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 memra_gguf::config::SwigluClamp;
17use std::sync::atomic::{AtomicU64, Ordering};
18
19/// Parse the documented `MEMRA_SPEC_REPLAY=1` rollback seam.
20///
21/// Keep this shared with serving admission so `=0` cannot select replay in one
22/// layer while another layer treats it as disabled.
23pub fn spec_replay_env_on(value: Option<&str>) -> bool {
24    value == Some("1")
25}
26
27pub fn spec_replay_env_enabled() -> bool {
28    let value = std::env::var("MEMRA_SPEC_REPLAY").ok();
29    spec_replay_env_on(value.as_deref())
30}
31
32/// step35 dcw draft-chain door (lane/step37-draft-graph-20260829). ON routes the step35 MTP
33/// block's draft attention through the WINDOWED device-counter family
34/// (`append_kv_quantized_dcw` + `fa_decode_dcw`, the step TP graph arc's kernels), which
35/// derives the SWA view entirely from device state (len_d, base_d, window): exactly the view
36/// offset the old capture refusal said `fa_decode_dc` could not express. BOTH draft modes
37/// switch together: eager and captured run the ONE launcher at the ONE bucket
38/// (min(cap, window)), so graph-vs-eager draft parity holds by construction (the
39/// `mtp_full_attn_dc` precedent).
40///
41/// DEFAULT ON since lane/step37-draft-graph-serving-20260830: the 20260829 lane shipped it
42/// OFF because it enabled nothing at the shipping head count (capture was structurally
43/// unreachable at heads=3); with the multi-head chain capture and the in-graph filtered
44/// sampler landed, this door is the kernel prerequisite for the captured chain on the
45/// QUALIFIED serving shape, and the exactness battery (greedy K=1..8 identity, per-K
46/// acceptance identity, seeded sampled twins) banks on the ON arm. Rollback seam:
47/// MEMRA_STEP35_DRAFT_DCW=0 restores the host-len eager arm (`mtp_step35_attn`) plus the
48/// named capture refusal, byte-for-byte the pre-lane serving; no state survives restart.
49fn step35_draft_dcw_on() -> bool {
50    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
51    *ON.get_or_init(|| std::env::var("MEMRA_STEP35_DRAFT_DCW").as_deref() != Ok("0"))
52}
53
54/// Multi-head MTP draft-chain capture door (lane/step37-draft-graph-serving-20260830,
55/// default ON — receipts in the lane RESULTS). ON lets the step-modulo prefix-replay chain
56/// (`mtp_extra` non-empty, the step37 3-head shipping shape) capture per-head single-row
57/// CUDA graphs and replay them in the exact eager launch order; the chain POLICY (head
58/// selection, prefix length, seed history) stays host-side, so graph-vs-eager drafts are
59/// bit-identical by construction. A failed capture degrades LOUDLY to the eager chain (the
60/// draft-graph WARN contract). OFF (=0) keeps the eager chain as the only multi-head path —
61/// the pre-lane serving byte-for-byte. Single-head capture is untouched by this door.
62fn mtp_chain_graph_on() -> bool {
63    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
64    *ON.get_or_init(|| std::env::var("MEMRA_MTP_CHAIN_GRAPH").as_deref() != Ok("0"))
65}
66
67/// In-graph FILTERED sampled draft door (lane/step37-draft-graph-serving-20260830, default
68/// ON — receipts in the lane RESULTS). ON widens the sampled draft-graph capture from the
69/// pure-temp regime to every truncation-filtered regime (top_k / top_p / min_p): the capture
70/// body runs `filter_stats` + `gumbel_perturb_filtered_ctr` IN-GRAPH, so the draft draws
71/// from the SAME filtered distribution the verify's accept test reconstructs (the
72/// graph-s-key exactness law, now satisfied inside the graph instead of by refusing it).
73/// Penalties stay eager either way (the history varies per round and cannot be baked).
74/// The pure-temp capture body is UNTOUCHED by this door (byte-identical to the pre-lane
75/// graph). OFF (=0) restores the pure-temp-only capture guard: filtered requests draft
76/// eager, byte-for-byte the pre-lane behavior.
77fn spec_graph_filtered_on() -> bool {
78    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
79    *ON.get_or_init(|| std::env::var("MEMRA_SPEC_GRAPH_FILTERED").as_deref() != Ok("0"))
80}
81
82fn parse_prime_trows_width(value: Option<&str>) -> Result<usize, String> {
83    let Some(raw) = value else {
84        return Ok(8);
85    };
86    let width = raw
87        .parse::<usize>()
88        .map_err(|_| format!("MEMRA_PRIME_TROWS_T must be an integer in 2..=8, got {raw:?}"))?;
89    if !(2..=8).contains(&width) {
90        return Err(format!("MEMRA_PRIME_TROWS_T must be in 2..=8, got {width}"));
91    }
92    Ok(width)
93}
94
95#[cfg(test)]
96mod prime_trows_width_tests {
97    #[test]
98    fn width_defaults_to_eight_and_refuses_invalid_operator_values() {
99        assert_eq!(super::parse_prime_trows_width(None), Ok(8));
100        assert_eq!(super::parse_prime_trows_width(Some("2")), Ok(2));
101        assert_eq!(super::parse_prime_trows_width(Some("8")), Ok(8));
102        for invalid in ["", "1", "9", "32", "wide"] {
103            let err = super::parse_prime_trows_width(Some(invalid)).unwrap_err();
104            assert!(err.contains("MEMRA_PRIME_TROWS_T"), "{err}");
105            assert!(err.contains("2..=8"), "{err}");
106        }
107    }
108}
109
110/// One compact, anchor-bounded DSpark supervision record. `tokens[0]` is the anchor at p and
111/// `hidden` is its predecessor carrier h[p-1], matching the live NextN/DSpark pairing. Target
112/// rows p..p+gamma-1 score tokens p+1..p+gamma. They are the full-target softmax's top-k
113/// entries; `target_tail_probs[j]` is the probability mass outside those rows. All flattened
114/// target arrays are `[gamma, top_k]` in row-major order.
115pub struct DsparkAnchorRecord {
116    pub position: usize,
117    pub hidden: Vec<f32>,
118    pub tokens: Vec<u32>,
119    pub target_top_ids: Vec<u32>,
120    pub target_top_logits: Vec<f32>,
121    pub target_top_probs: Vec<f32>,
122    pub target_tail_probs: Vec<f32>,
123}
124
125#[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
126fn dspark_sparse_softmax_topk(
127    logits: &[f32],
128    top_k: usize,
129    temperature: f32,
130) -> Result<(Vec<u32>, Vec<f32>, Vec<f32>, f32), Box<dyn std::error::Error>> {
131    if logits.is_empty() || top_k == 0 || top_k > logits.len() || temperature <= 0.0 {
132        return Err("invalid DSpark sparse-softmax shape or temperature".into());
133    }
134    if logits.iter().any(|value| !value.is_finite()) {
135        return Err("DSpark target logits contain a non-finite value".into());
136    }
137    let mut ranked: Vec<(u32, f32)> = logits
138        .iter()
139        .copied()
140        .enumerate()
141        .map(|(index, value)| (index as u32, value))
142        .collect();
143    let compare = |left: &(u32, f32), right: &(u32, f32)| {
144        right.1.total_cmp(&left.1).then(left.0.cmp(&right.0))
145    };
146    ranked.select_nth_unstable_by(top_k - 1, compare);
147    ranked[..top_k].sort_unstable_by(compare);
148
149    let max_logit = logits.iter().copied().fold(f32::NEG_INFINITY, f32::max);
150    let inv_temperature = 1.0f64 / temperature as f64;
151    let denominator: f64 = logits
152        .iter()
153        .map(|value| (((*value - max_logit) as f64) * inv_temperature).exp())
154        .sum();
155    let ids: Vec<u32> = ranked[..top_k].iter().map(|(index, _)| *index).collect();
156    let top_logits: Vec<f32> = ranked[..top_k].iter().map(|(_, value)| *value).collect();
157    let top_probs: Vec<f32> = top_logits
158        .iter()
159        .map(|value| ((((value - max_logit) as f64) * inv_temperature).exp() / denominator) as f32)
160        .collect();
161    let top_mass: f64 = top_probs.iter().map(|value| *value as f64).sum();
162    let tail = (1.0f64 - top_mass).clamp(0.0, 1.0) as f32;
163    Ok((ids, top_logits, top_probs, tail))
164}
165
166fn flatten_dspark_rows<T>(
167    rows: Vec<Option<Vec<T>>>,
168    position: usize,
169    label: &str,
170) -> Result<Vec<T>, Box<dyn std::error::Error>> {
171    let mut flattened = Vec::new();
172    for (slot, row) in rows.into_iter().enumerate() {
173        flattened.extend(
174            row.ok_or_else(|| format!("missing DSpark {label} at {position} slot {slot}"))?,
175        );
176    }
177    Ok(flattened)
178}
179
180/// H-SEED CONVENTION (MEMRA_SPEC_HPOST=1): feed the MTP head the POST-norm hidden — trunk rows
181/// hand over `output_norm(x)` and the draft chain recurrence hands over `shared_head_norm(h_nextn)`
182/// (= final_h) — matching the reference engines: llama.cpp #24025 ("qwen35: use post-norm hidden
183/// state for MTP", t_h_nextn is taken AFTER the final norm in both trunk and MTP graphs) and
184/// SGLang's qwen3_5_mtp (spec_info.hidden_states = the target model's post-norm output). memra's
185/// historical convention (default, MTP-PLAN §A) is PRE-norm x. Draft-quality-only: exactness is
186/// the verify's job either way; acceptance arbitrates. OnceLock: read once, hot-loop safe.
187/// `MEMRA_SPEC_HEAD_ROWS=1` — batch the verify tail's LM head over its t columns instead of running
188/// it at m=1 once per column. See the call site in `decode_step_t_core_stream` for why the batched
189/// form is the same per-row arithmetic (the bf16/q8 rows twins, not cuBLASLt) and what it costs
190/// today: the head is re-streamed t times per verify pass. Default off until the byte tape says so.
191pub(crate) fn head_rows_on() -> bool {
192    static ENV: std::sync::OnceLock<Option<bool>> = std::sync::OnceLock::new();
193    crate::step37_door(&ENV, "MEMRA_SPEC_HEAD_ROWS")
194}
195
196/// The serving walk's own doors, tri-stated the same way (owner flip 2026-08-27): env forces,
197/// unset takes the step37 family default. Call sites are the t-row verify walk itself.
198pub(crate) fn spec_verify_eager_on() -> bool {
199    static ENV: std::sync::OnceLock<Option<bool>> = std::sync::OnceLock::new();
200    crate::step37_door(&ENV, "MEMRA_SPEC_VERIFY_EAGER")
201}
202
203pub(crate) fn spec_verify_tcol_on() -> bool {
204    static ENV: std::sync::OnceLock<Option<bool>> = std::sync::OnceLock::new();
205    crate::step37_door(&ENV, "MEMRA_SPEC_VERIFY_TCOL")
206}
207
208/// NOT family-armed (2026-08-27): the walk's prime leaves its sub-32 TAIL chunk out of the
209/// DISTRIBUTED kv, so the server refuses before decode with "cache lengths diverged
210/// local=N distributed=floor(N/32)*32" for every prompt whose token count is not a multiple of
211/// 32 — i.e. nearly all real traffic. Isolated on the server route: defaults ERR (local=445
212/// distributed=416), MEMRA_PRIME_TROWS=0 OK. It was default-OFF before the 2026-08-27 flip and
213/// goes back to opt-in until the tail append is fixed and gated ON THE SERVER ROUTE, not just
214/// run-gen (run-gen calls decode_step_t on the whole prompt and never exercises this path — the
215/// reason a run-gen-only receipt could not see it). The GEMM prime supersedes it on this route.
216pub(crate) fn prime_trows_on() -> bool {
217    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
218    *ON.get_or_init(|| std::env::var("MEMRA_PRIME_TROWS").as_deref() == Ok("1"))
219}
220
221pub(crate) fn tcol_ffn_on() -> bool {
222    static ENV: std::sync::OnceLock<Option<bool>> = std::sync::OnceLock::new();
223    crate::step37_door(&ENV, "MEMRA_TCOL_FFN")
224}
225
226pub(crate) fn spec_hpost() -> bool {
227    static H: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
228    *H.get_or_init(|| {
229        std::env::var("MEMRA_SPEC_HPOST")
230            .map(|v| v != "0")
231            .unwrap_or(false)
232    })
233}
234
235/// LEAN VERIFY (default ON since 2026-07-08; MEMRA_SPEC_LEAN=0 reverts — close35 lane): the verify m-scaling
236/// probe + nsys diff showed the verify t-path pays ~1.0ms/call at m=1 over eager decode on the
237/// 35B, and the kernels are NOT the cause (dev-MoE identical, kernel-time delta only +179us).
238/// The overhead is (a) ~250 extra cuMemsetD8Async/call from `e.zeros()` on buffers every kernel
239/// fully overwrites (~0.9ms host issue + ~0.35ms GPU) and (b) the t=1 FA rows dispatch (rows_v2 +
240/// combine_rows, +50us vs the eager fa_decode pair). This flag switches (a) fully-overwritten
241/// verify buffers to `e.uninit` (identical bytes: every element is written before read) and
242/// (b) t==1 verify FA to the eager `fa_decode` entry (byte-identical: kernel-check pins the
243/// rows-vs-loop identity and the per-row loop at t=1 IS fa_decode on the same q). Gates arbitrate.
244pub(crate) fn spec_lean() -> bool {
245    static L: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
246    // DEFAULT ON since 2026-07-08 (MEMRA_SPEC_LEAN=0 reverts): bit-identical (buffers fully
247    // overwritten; gates green incl maxdiff-identical run-gen) and measured +2.4% e2e p3 /
248    // +1.5% p2 at the daily 35B config. m=1 verify now costs eager-decode parity.
249    *L.get_or_init(|| {
250        std::env::var("MEMRA_SPEC_LEAN")
251            .map(|v| v != "0")
252            .unwrap_or(true)
253    })
254}
255
256/// SMALL-M BATCHED VERIFY (default ON since 2026-07-09; MEMRA_SPEC_M2=0 reverts — lane/spec-m2): extend the
257/// batched linear-attn verify arm down to t=2 and batch the MoE dev token loop over a
258/// grid.z=token axis at every verify t. The close35 m-scaling probe put the m=2 verify tier at
259/// x1.54 of m=1 (llama x1.14); the per-column linear chain (t<3) and the serial MoE dev token
260/// loop are the two launch-structure causes. Both changes are LAUNCH-STRUCTURE ONLY:
261/// (a) the batched conv's t<pad ring update is pure copies (ssm_conv_ring_rebuild from a cloned
262///     ring — the ring stores raw input columns); every arithmetic kernel is the same one the
263///     t>=3 arm already runs (matmul_decode_exact bit-identical at m=2-4, gdn_scan's internal
264///     t-loop == chained T=1 steps);
265/// (b) the MoE dev-rows twins run the serial loop's per-token warp program with tok-offset
266///     pointers (same sel/w/aq/ad bytes, same dot order, same slot-ordered FMA chain).
267/// Gates arbitrate: run-spec K=1..8 self-consistency (35B+9B), kernel-check, run-gen argmax.
268pub(crate) fn spec_m2() -> bool {
269    static M: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
270    // DEFAULT ON since 2026-07-09 (MEMRA_SPEC_M2=0 reverts): launch-structure only — t=2
271    // batched linear arm (ring-roll copies, zero new FP order) + MoE dev-rows kernels
272    // (grid.z=token, 4 launches/layer at any verify t). Acceptance bit-identical at every K;
273    // 35B p2 +3.4% / p3 +3.6%; the profitable-K plateau widens (new optimum K=3 at 223).
274    *M.get_or_init(|| {
275        std::env::var("MEMRA_SPEC_M2")
276            .map(|v| v != "0")
277            .unwrap_or(true)
278    })
279}
280pub(crate) fn spec_stream() -> bool {
281    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
282    *ON.get_or_init(|| std::env::var("MEMRA_SPEC_STREAM").as_deref() == Ok("1"))
283}
284pub(crate) fn spec_stream_m() -> usize {
285    static M: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
286    *M.get_or_init(|| {
287        std::env::var("MEMRA_SPEC_STREAM_M")
288            .ok()
289            .and_then(|v| v.parse().ok())
290            .unwrap_or(4)
291    })
292}
293pub(crate) fn spec_devacc() -> bool {
294    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
295    *ON.get_or_init(|| std::env::var("MEMRA_SPEC_DEVACC").as_deref() == Ok("1"))
296}
297/// Engine-bundle slice 2 (DSF-ROUNDCOST-20260820 §1.1 host/device round trips + §2 rows 2-3),
298/// DEFAULT ON (`MEMRA_DSPARK_DEFER_READBACK=0` reverts): the dspark round's draft-chain DtoH
299/// is DEFERRED past verify dispatch and merged with the verify-argmax readback into ONE host
300/// sync (2 blocking DtoH/round -> 1). Verify embeds DEVICE tokens (`chain_d`) through the
301/// resident embed table — `embed_gather_u32_t`, bit-identical rows to the host gather by its
302/// own pinned contract. The host therefore dispatches snap + the whole verify while the DRAFT
303/// is still executing, instead of blocking ~1.7 ms on the chain and letting the device drain.
304/// Ladder arm only: the confidence policies size vt from a pre-verify head readback (their
305/// chain readback merges into that same sync instead). Exactness unchanged BY CONSTRUCTION —
306/// same tokens, same kernels, same order; E2E + accept-bank gates arbitrate.
307pub(crate) fn dspark_defer_readback_on() -> bool {
308    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
309    *ON.get_or_init(|| {
310        std::env::var("MEMRA_DSPARK_DEFER_READBACK")
311            .map(|v| v != "0")
312            .unwrap_or(true)
313    })
314}
315/// Engine-bundle slice 1 (DSF-ROUNDCOST-20260820 §1.1, lane/dspark-engine-bundle-20260820),
316/// DEFAULT ON (`MEMRA_STATE_COPY_BATCH=0` reverts): batch the dspark round's GDN state
317/// snapshot and partial-accept restore into single `copy_batch_uniform_f32` launches
318/// instead of ~2 memcpy dispatches (+2 alloc_zeros on the snap side) per linear layer per
319/// round — measured 0.67 ms/round snap + 0.25 ms/round commit of pure dispatch on the q38
320/// route. Launch-structure only: bytes, buffers and stream order are unchanged, so
321/// acceptance and streams stay bit-identical (E2E-gated on the B1 packs).
322pub(crate) fn state_copy_batch_on() -> bool {
323    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
324    *ON.get_or_init(|| {
325        std::env::var("MEMRA_STATE_COPY_BATCH")
326            .map(|v| v != "0")
327            .unwrap_or(true)
328    })
329}
330/// Engine-bundle slice 3 + fa-execupdate slice 4c (DSF-ROUNDCOST-20260820 §5 rank 1),
331/// DEFAULT OFF — `MEMRA_DSPARK_VERIFY_GRAPH=1` opts in: per-(segment, vt) CUDA graphs
332/// for the LINEAR-layer runs, plus the full-verify single graph per (vt, rung) when a
333/// round's rows all ride one seqs rung — see [`DsparkVerifyGraphs`]. Requires the
334/// slice-2 deferred path (device tokens); the eager walk is the byte-identical fallback.
335///
336/// MEASURED disposition (box6 card0, agentic pack, 2026-08-20, both slices): exactness
337/// holds everywhere (ALL EXACT, accept lines byte-match the banks, ckpt-gate oracle
338/// green over the graph + slab-commit paths). Slice-3's AUTO_FREE launch-scan limiter
339/// (25.6 us x 16 launches ≈ 0.41 ms/round) is FIXED — the captured bodies' alloc nodes
340/// are balanced by in-graph frees (census 84/84 per segment, 1776/1776 full) so graphs
341/// instantiate USE_NODE_PRIORITY and the scan is gone. What remains at gate scale:
342/// segment graphs +0.1 tok/s over the batched-rows default (114.4 vs 114.3 x5
343/// interleaved — the linear launch overhead was only ~0.1 ms); the FULL-verify graph is
344/// NET NEGATIVE at gate scale (110.6 vs 114.2: ~14-21 (vt, rung) captures/process at
345/// 2 full-walk executions + ~2.9k-node instantiate each eat far more than the ~0.2-0.3
346/// ms/round of remaining launch overhead). The orchestration ceiling of §1.3 is spent —
347/// the fa/append recovery landed DEFAULT-ON as the batched rows arm
348/// (`dspark_fa_rows_on`), not as a graph. The serve-lifetime cell (DSF-ROUNDCOST §9,
349/// nj-ws-solo) measured the amortization: crossover K≈33 requests, steady −0.246
350/// ms/round, −1.25% session wall over 240 requests — and the graphs-serve lane wired
351/// the door into the session arm (`dspark_spec_session_burst`) as a model-owned
352/// capture pool shared across sessions. Stays opt-in pending the owner's default-ON
353/// ratification on the serve-surface battery.
354pub(crate) fn dspark_verify_graph_on() -> bool {
355    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
356    *ON.get_or_init(|| std::env::var("MEMRA_DSPARK_VERIFY_GRAPH").as_deref() == Ok("1"))
357}
358/// MTP-ROUTE verify graphs, DEFAULT ON for the GDN+MoE family since 2026-08-23
359/// (`MEMRA_SPEC_VERIFY_GRAPH=0` is the kill switch, `=1` opts other families in).
360///
361/// The slice-4c capture already lived inside `qwen35_verify_tparallel` and said so in its own
362/// comment — "stream rides the qwen35moe burst, graphs ride the dspark route" — with no caller
363/// on this route. The MTP spec round is that caller.
364///
365/// WHY it is worth a default (receipts: `research/orndecode-20260822/VGRAPH.md`). With
366/// `MEMRA_SPEC_PHASE=1` this route's round reads verify-ISSUE 44-58% and verify-WAIT **0.0%**:
367/// the host is never waiting for the device, it is spending its own time launching the trunk.
368/// Replay collapses that into one graph launch and the phase all but disappears (55-62 ms ->
369/// 8-10 ms per burst).
370///
371/// MEASURED, two host generations, forced ON/OFF, balanced 4+4 boots in both orders:
372///   * current-generation host (9950X, the serving class): OFF 266.0-266.5, ON 318.8-319.5
373///     tok/s — **+19.7%**, no overlap, sub-1% spread per arm; per-round 6.9 -> 5.7 ms.
374///   * Zen 3 host: +3-9% (that rig's own clock drift is wider than the effect, so the ratio
375///     comes from per-round phase totals, which are internal to each boot).
376///     The ON arm lands at ~320 tok/s on BOTH hosts while OFF tracks host speed — the arm moves
377///     the round off the host and onto the device, which is the whole point.
378///
379/// EXACTNESS is structural (same kernels, same order) and gated anyway: a fixed-seed SAMPLED
380/// completion hashes identically ON vs OFF **and across both hosts** (`08941d5bb9762b21`),
381/// greedy seed-pinned likewise, `run-spec` K=1..8 PASS on both arms with identical acceptance
382/// at every K, kernel-check ALL GREEN.
383///
384/// SCOPE, deliberately narrow: default ON only where it was measured — the GatedDeltaNet +
385/// MoE family (`vgraph_family_default`). Qwen3.8-27B is GDN + DENSE mlp and would otherwise
386/// inherit this default unmeasured, which is the family-by-family law this repo keeps; it can
387/// opt in with `=1` once it has its own interleave. Also never armed together with
388/// ROUND-STREAM, and a round wider than the pool declines it for the eager walk.
389pub(crate) fn spec_verify_graph_env() -> Option<bool> {
390    static ON: std::sync::OnceLock<Option<bool>> = std::sync::OnceLock::new();
391    *ON.get_or_init(
392        || match std::env::var("MEMRA_SPEC_VERIFY_GRAPH").as_deref() {
393            Ok("1") => Some(true),
394            Ok("0") => Some(false),
395            _ => None,
396        },
397    )
398}
399/// SERVE-ROUTE twin of [`dspark_verify_graph_on`], DEFAULT ON — owner-ratified
400/// 2026-08-22 on the §10 serve-lifetime battery (DSF-ROUNDCOST-20260820 §10.3:
401/// crossover K=36–43, steady −0.357 ms/round, session wall −1.55..−1.65%, byte-exact
402/// 240/240 ×3 pairs, pool bounded at 8,852 MiB under `MEMRA_DSPARK_VG_MAX`). The env
403/// stays as the kill-switch: `MEMRA_DSPARK_VERIFY_GRAPH=0` restores the eager walk
404/// (byte-identical body); `MEMRA_DSPARK_VG_MAX=0` is the finer freeze valve. The BIN
405/// arm keeps its own opt-in default (`dspark_verify_graph_on`): at gate scale the
406/// capture toll is never repaid (§8 measured disposition — 14–21 captures over a
407/// 256-token run vs the serve session's thousands of rounds), and the two
408/// instruments must keep their own measured dispositions rather than share one flag.
409pub(crate) fn dspark_verify_graph_serve_on() -> bool {
410    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
411    *ON.get_or_init(|| std::env::var("MEMRA_DSPARK_VERIFY_GRAPH").as_deref() != Ok("0"))
412}
413/// Capture-count ceiling for the dspark verify-graph pool (graphs-serve lane) — the
414/// pool's memory policy STATED instead of silently unbounded. The keyspace is
415/// intrinsically finite — segment keys (run_start, vt) ≤ 16 runs x 7 windows, full
416/// keys (vt, rung, hi) ≤ 7 windows x the split-rung ladder (8 rungs at 32k ctx), ~168
417/// on the q38 export — so the default (256) never engages there; the knob is the
418/// safety valve for a future export with a wider ladder. At the ceiling the pool
419/// FREEZES: existing keys keep replaying, rounds needing a new capture run the eager
420/// walk byte-identically (round-atomic — a partial refusal would mix slab- and
421/// cols-stashed layers inside one commit). No eviction by design: destroying a live
422/// exec graph re-opens the stale-address class the indirect tables exist to close,
423/// and the bounded keyspace makes reclaim worthless.
424pub(crate) fn dspark_vg_cap() -> usize {
425    static CAP: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
426    *CAP.get_or_init(|| {
427        std::env::var("MEMRA_DSPARK_VG_MAX")
428            .ok()
429            .and_then(|v| v.parse().ok())
430            .unwrap_or(256)
431    })
432}
433
434/// PROJECTED REMAINING GROWTH of the verify-graph pool, in bytes (lane/hermes-perf-fixes,
435/// 2026-08-23 — the admission accounting the "pool dwarfs spec admission reserve" finding
436/// asks for). The pool was measured at 8,852 MiB at storm-complete on the q38 export while
437/// admission's transient floor (`SPEC_SHRINK_RESERVE`) is 1.5 GiB and never charged for it:
438/// sessions admitted while the pool is cold overcommit VRAM the pool WILL hold, because the
439/// pool grows monotonically (no eviction by design) and is model-owned across sessions.
440///
441/// SELF-MEASURING, no per-model constant (generic-model law — the 8,852 MiB is a q38 number
442/// and proves nothing about another export): the debt is remaining capture slots x the
443/// MARGINAL bytes a capture adds to this device's graph mem pool.
444///
445/// MARGINAL, NOT MEAN — measured correction (box9 on-box receipt, 2026-08-23). The first
446/// version of this used the mean (`reserved / captures`) and the live serve log showed why
447/// that is wrong: with the pool's reservation flat at ~33.6 MiB across captures 1..3, the
448/// mean-based debt printed **8,556 MB, then 4,261, then 2,830** — it extrapolated capture
449/// #1's ONE-TIME shared allocation (staging buffers, stash slabs, pointer tables: sized
450/// once per pool, shared by every key) across all 256 slots. An 8.5 GB phantom reserve at
451/// boot can refuse admissions that would have fit, which is a worse defect than the
452/// under-charge this accounting exists to remove. The marginal reading prices what an
453/// ADDITIONAL key actually costs: two observations `(captures, reserved)` give
454/// `(r1 - r0) / (c1 - c0)`, which is ~0 on an export whose pool does not grow per key and
455/// tracks real growth on one that does.
456///
457/// BOOTSTRAP (only one observation so far, so growth is unmeasurable): reserve one more
458/// pool's worth — `min(remaining x mean, reserved)`. "We have measured `reserved` bytes for
459/// `captures` keys; until growth is measurable, assume at most a doubling" is fail-safe in
460/// the same direction as the old rule without the 255x extrapolation.
461///
462/// Before the FIRST capture the debt is 0 (a single capture lands well inside the existing
463/// 1.5 GiB floor). `cap` is the intrinsic freeze ceiling (`MEMRA_DSPARK_VG_MAX`; =0 freeze
464/// valve => the pool cannot grow => debt 0); at or past the cap the pool FREEZES, so the
465/// debt is 0 there too.
466pub fn dspark_vg_debt_projection(
467    captures: usize,
468    cap: usize,
469    reserved_bytes: usize,
470    prev: Option<(usize, usize)>,
471) -> usize {
472    if captures == 0 || cap == 0 {
473        return 0;
474    }
475    let remaining = cap.saturating_sub(captures);
476    if remaining == 0 {
477        return 0;
478    }
479    match prev {
480        // marginal growth between two observations of the same pool
481        Some((c0, r0)) if captures > c0 => {
482            let marginal = reserved_bytes.saturating_sub(r0) / (captures - c0);
483            remaining.saturating_mul(marginal)
484        }
485        // bootstrap: at most one more pool's worth
486        _ => remaining
487            .saturating_mul(reserved_bytes / captures)
488            .min(reserved_bytes),
489    }
490}
491/// PRE-CAPTURE VRAM RESERVE CHECK door (lane/step37-vram-admission-20260830), DEFAULT ON.
492/// A draft-graph capture attempt on a tight card used to be try-and-fail: the 2 warmup
493/// forwards + instantiate grew the pool to the edge BEFORE the OOM surfaced, and the
494/// "eager fallback" then ran on a card the failed attempt had just exhausted (the owner's
495/// single-session second-prompt OOM: capture WARN followed by 28 step-OOM engine errors,
496/// device at 5 MiB free). With the gate ON, a capture is attempted only when the device's
497/// effective free (driver free + async-pool cached) covers the capture's expected appetite
498/// PLUS a post-capture safety floor — otherwise the session falls back to eager EARLY,
499/// with headroom intact, through the same LOUD once-per-flip WARN. `=0` restores
500/// try-and-fail (diagnostics door; the trim-on-OOM recovery below stays active either way).
501pub fn spec_capture_gate_on() -> bool {
502    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
503    *ON.get_or_init(|| std::env::var("MEMRA_SPEC_CAPTURE_GATE").as_deref() != Ok("0"))
504}
505
506/// Post-capture safety floor the reserve check keeps free ON TOP of the capture's own
507/// appetite: the same measured constant class as the admission transient floor
508/// (capture arenas + verify activations — the admit-oom control fit). A capture that
509/// would leave less than this behind is not worth its eager-coverage risk.
510pub(crate) const CAPTURE_HEADROOM_FLOOR: usize = 1536 << 20;
511
512/// Pure verdict half of the pre-capture reserve check (unit-testable): given the device's
513/// driver-free and pool-cached bytes and the capture's expected `need`, returns
514/// `Some((required, effective))` when the capture must be REFUSED, `None` when it fits.
515pub(crate) fn capture_headroom_verdict(
516    driver_free: usize,
517    pool_cached: usize,
518    need: usize,
519    floor: usize,
520) -> Option<(usize, usize)> {
521    let effective = driver_free.saturating_add(pool_cached);
522    let required = need.saturating_add(floor);
523    (effective < required).then_some((required, effective))
524}
525
526/// Expected device appetite of a draft-graph capture attempt when no measurement exists
527/// yet (bootstrap only — the model-owned high-water gauge takes over after the first
528/// observed capture). Deliberately conservative and shape-derived, never a per-family
529/// constant: per (head, mode) capture the two warmups + capture each walk one head
530/// forward whose dominant transients are a handful of `n_embd` rows and one `d_vocab`
531/// logits row, retained by the keeper; the sampled tail additionally parks
532/// `k` q-slots + perturb/q buffers of `d_vocab` each.
533pub(crate) fn draft_capture_bootstrap_estimate(
534    heads: usize,
535    k: usize,
536    d_vocab: usize,
537    n_embd: usize,
538) -> usize {
539    let per_capture = 3usize // 2 warmups + capture body, each retaining its transients
540        .saturating_mul(d_vocab.saturating_add(8 * n_embd))
541        .saturating_mul(4)
542        .max(32 << 20); // instantiate + driver-side graph backing per capture, floor
543    let captures = heads.max(1).saturating_mul(2); // interior + last per head
544    let sampled_slots = (k.saturating_add(2))
545        .saturating_mul(d_vocab)
546        .saturating_mul(4);
547    captures
548        .saturating_mul(per_capture)
549        .saturating_add(sampled_slots)
550        .max(64 << 20)
551}
552
553/// OOM predicate for capture-failure recovery (engine-side twin of the worker's
554/// `is_cuda_oom` — the same quoted-text contract).
555pub(crate) fn capture_err_is_oom(reason: &str) -> bool {
556    reason.contains("CUDA_ERROR_OUT_OF_MEMORY") || reason.contains("out of memory")
557}
558
559/// Impure half of the pre-capture reserve check: reads the device, trims the async pool
560/// when the driver alone is short but cached blocks would cover it (graph instantiate and
561/// cuBLAS workspaces allocate from the DRIVER, not from our pool — a pool sitting on freed
562/// blocks starves them), and returns the refusal reason line when the capture must not be
563/// attempted. `None` = go ahead.
564pub(crate) fn capture_headroom_refusal(e: &Engine, need: usize) -> Option<String> {
565    let Ok((driver_free, _total)) = e.ctx().mem_get_info() else {
566        return None; // unreadable device: keep the historical try-and-fail behavior
567    };
568    let pool_cached = e.pool_cached_bytes();
569    // A capture may take AT MOST HALF the discretionary headroom: required =
570    // 2x appetite + two floors (owner's contract: "fall back to eager EARLY with headroom
571    // intact"). Measured escalation on the owner-shape cells: one floor of slack let the
572    // capture walk the card to the edge and the burst step-OOM'd immediately; two floors
573    // still allowed a capture whose session then OOM'd on its own admission-charged work,
574    // because the capture had consumed the memory the charge was counting on. Requiring
575    // the appetite TWICE means the card retains a whole capture's worth of room after the
576    // capture lands - enough for the session's charged classes and its peers' bursts. The
577    // capture is an optimization worth ~2-3 ms of TTFT (draft-graph lane receipts); at the
578    // margin it is never worth an OOM incident.
579    let floor = CAPTURE_HEADROOM_FLOOR.saturating_mul(2);
580    let required_need = need.saturating_mul(2);
581    let required = required_need.saturating_add(floor);
582    match capture_headroom_verdict(driver_free, pool_cached, required_need, floor) {
583        Some((required, effective)) => Some(format!(
584            "insufficient VRAM headroom for capture: effective free {}MB (driver {}MB + pool-cached \
585             {}MB) < required {}MB (2x appetite {}MB + floor {}MB); capture skipped pre-attempt",
586            effective / (1 << 20),
587            driver_free / (1 << 20),
588            pool_cached / (1 << 20),
589            required / (1 << 20),
590            need / (1 << 20),
591            floor / (1 << 20),
592        )),
593        None => {
594            if driver_free < required && pool_cached > 0 {
595                let trimmed = e.pool_trim_to_zero();
596                if trimmed > 0 {
597                    eprintln!(
598                        "[spec] pre-capture pool trim: released {}MB cached back to the driver \
599                         (driver free {}MB < required {}MB; instantiate allocates from the driver)",
600                        trimmed / (1 << 20),
601                        driver_free / (1 << 20),
602                        required / (1 << 20),
603                    );
604                }
605            }
606            None
607        }
608    }
609}
610
611/// GRAPH-LAUNCH HEADROOM FLOOR (lane/step37-vram-admission-20260830, defect 3 root
612/// cause): `cuGraphLaunch` SEGFAULTS inside libcuda (offset +0x27c87f, a null internal
613/// dereference at address 0x60) when a captured graph is dispatched into a
614/// driver-exhausted card — reproduced on this lane's box with core dumps on BOTH the
615/// pre-lane and lane binaries (multi-active step-OOM squeeze; the crashing thread sits in
616/// `CudaGraph::launch` inside `generate_spec_inner2`). The eager arms fail RECOVERABLY on
617/// the same card (a quoted CUDA OOM the park path handles), so below this driver-free
618/// floor every graph arm yields to eager for the round. A named constant, not a knob: the
619/// winning value is the default and the guard exists to make a driver segfault
620/// unreachable, not to tune anything.
621pub(crate) const GRAPH_LAUNCH_MIN_FREE: usize = 256 << 20;
622
623/// Per-round guard for the floor above. Read failure keeps serving (never a false
624/// refusal from an unreadable device); one `mem_get_info` (~microseconds) per ~25ms round.
625pub(crate) fn graph_launch_headroom_ok(e: &Engine) -> bool {
626    match e.ctx().mem_get_info() {
627        Ok((free, _total)) => free >= GRAPH_LAUNCH_MIN_FREE,
628        Err(_) => true,
629    }
630}
631
632/// One grep-stable suspension line per ROUTE (each call site holds its own
633/// process-lifetime `Once`): every captured-graph launch route below the floor names
634/// itself in the tag while keeping the same `graph replay suspended:` key the step37
635/// admission lane's squeeze cell greps for. The spec-round guard keeps its original
636/// per-generation `[spec]` line; the sweep routes (graph-launch-guard-sweep lane,
637/// 2026-08-31) note once per process — presence is what the gates assert, and a
638/// suspended round is otherwise byte-identical to its eager twin.
639pub(crate) fn graph_replay_suspended_note(route: &str) {
640    eprintln!(
641        "[{route}] graph replay suspended: driver free below the {}MB launch floor \
642         (eager arms serve; cuGraphLaunch segfaults into an exhausted card)",
643        GRAPH_LAUNCH_MIN_FREE / (1 << 20)
644    );
645}
646
647/// Engine-bundle slice 4 (fa-execupdate lane, DSF-ROUNDCOST-20260820 §6 close: "the
648/// residual gap lives in the FULL-ATTENTION per-row section"), DEFAULT ON —
649/// `MEMRA_DSPARK_FA_ROWS=0` reverts to the per-row loop: when every row of a verify
650/// round takes the v4-seqs arm on ONE `fa_split_keys` rung (the straddle law, evaluated
651/// at the round's first and last t_kv — both eligibility gates are intervals in t_kv),
652/// the qwen35 t-parallel verify's per-row KV-append + fa-decode loop collapses into the
653/// z-batched serving twins: ONE `append_quantize_kv_q8_0_q5_1_seqs` + ONE
654/// `fa_decode_vec_q_seqs_v4` + ONE combine per full-attention layer, replacing
655/// T x (4 dtod row copies + append + 3 memsets + main + combine) launches. Bytes are
656/// pinned by the batched-tick increment-2 kernel-check (seqs-vs-per-seq-loop bit
657/// identity: per-row T_kv derives in-kernel from pos_seq[z]; splits >= ns_eff write the
658/// empty partial the combine never reads, so the shared n_splits_max stride changes no
659/// bytes) and re-gated e2e by this lane's battery.
660pub(crate) fn dspark_fa_rows_on() -> bool {
661    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
662    *ON.get_or_init(|| {
663        std::env::var("MEMRA_DSPARK_FA_ROWS")
664            .map(|v| v != "0")
665            .unwrap_or(true)
666    })
667}
668
669/// `t_pred0` for the `MEMRA_DEBUG_SPEC` per-round print, sampled-safe.
670///
671/// `generate_spec_inner2` fills its `preds` vector ONLY on the greedy path (`if !sampled`), and
672/// the per-round debug print was the sole consumer in the sampled arm: `t_pred(0)` survives round
673/// 0 (`base == 0` returns `last_pred`) and from round 1 (`base == 1`, a pending bonus) indexes an
674/// EMPTY vector — `index out of bounds: the len is 0 but the index is 0`, in the GPU worker
675/// thread, which then respawns and reloads weights while the request dies. So any sampled spec
676/// request longer than one round used to kill the worker whenever `MEMRA_DEBUG_SPEC` was set:
677/// the flag crashed precisely the regime it exists to investigate.
678///
679/// Fixed at the print site, not inside the closure, so the greedy accept walk keeps its strict
680/// indexing (an out-of-range pred there is a real bug and must still be loud).
681fn debug_t_pred0(sampled: bool, base: usize, last_pred: u32, preds: &[u32]) -> String {
682    if base == 0 {
683        return last_pred.to_string();
684    }
685    match preds.get(base - 1) {
686        Some(p) => p.to_string(),
687        // sampled: the greedy per-column argmax was never run for this round.
688        None => {
689            debug_assert!(
690                sampled,
691                "greedy spec: preds[{}] missing at base {base}",
692                base - 1
693            );
694            "n/a".to_string()
695        }
696    }
697}
698
699/// `MEMRA_SKEY_PROBE=1` — sampled-draft-graph key probe (lane/graph-s-key-exactness-20260819).
700///
701/// Reports, per burst and per round, which draft chain the sampled arm chose and under which
702/// filter regime, plus the ONE observable that separates a legal filtered draft from a stale
703/// pure-temp graph replayed under filters: an accept test whose gathered `q` is exactly 0.
704/// A draft token sampled from the FILTERED softmax can never gather q=0 (it was drawn from the
705/// kept set), so `q=0` in the verify means the draft came from a distribution the verify does
706/// not believe in — and `u * 0 < p` then accepts it unconditionally.
707///
708/// Its own env var, deliberately NOT `MEMRA_DEBUG_SPEC`: that flag panicked the GPU worker on
709/// any sampled spec request past round 0 until this lane fixed it (§2 of the bank note).
710pub(crate) fn skey_probe() -> bool {
711    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
712    *ON.get_or_init(|| std::env::var("MEMRA_SKEY_PROBE").as_deref() == Ok("1"))
713}
714
715/// GRAMMAR HOOK for constrained spec decode (lane/constrained-full, 2026-08-03). The engine
716/// stays llguidance-agnostic: the server adapts its per-session grammar state behind this
717/// trait. CONTRACT (the verify-side truncation rule — token-identical to constrained plain
718/// greedy decode): the exactness walk runs UNMASKED first; the hook then (a) truncates
719/// acceptance at the first grammar-illegal accepted token, and (b) when the truncation fired
720/// or the bonus is illegal, the engine recomputes that slot as the MASKED argmax of the
721/// target's own verify column (an unmasked argmax that is grammar-legal IS the masked argmax
722/// — masking only removes tokens — so the common case pays nothing). `consume` advances the
723/// state with each EMITTED token in order; EOS handling is the implementor's job (skip).
724pub trait SpecConstraint {
725    /// -inf the current state's banned ids on a HOST logits row (prompt-tail / init-feed
726    /// masked argmax).
727    fn mask_logits(&mut self, logits: &mut [f32]) -> Result<(), String>;
728    /// Packed 32-bit bitset words of the CURRENT state's allowed set (device-mask form).
729    fn mask_words(&mut self) -> Result<Vec<u32>, String>;
730    /// Is `tok` consumable in the CURRENT state?
731    fn is_allowed(&mut self, tok: u32) -> Result<bool, String>;
732    /// Advance the state with an emitted token.
733    fn consume(&mut self, tok: u32) -> Result<(), String>;
734
735    // --- DRAFT-SIDE MASKING (lane/draft-mask, 2026-08-04) ---
736    // The drafter proposed grammar-illegal tokens under tight schemas, so verify-side
737    // truncation cut nearly every round (measured acceptance 0.467-0.513 tight vs 0.62-0.82
738    // loose, research/constrained-full-20260803). These three methods let the engine mask the
739    // DRAFT model's own sampling with the grammar's legal set, so proposals are legal by
740    // construction. The state they walk is a SPECULATIVE CLONE of the session matcher — the
741    // real state is advanced only by `consume` (emitted tokens), so verify-side truncation
742    // stays the correctness backstop and the emitted stream is unchanged by construction
743    // (an accepted draft is the target's unmasked argmax AND grammar-legal, hence the masked
744    // argmax; a cut slot is recomputed as the masked argmax either way).
745    // Default impls = feature OFF (pre-lane behaviour: unmasked drafts).
746
747    /// Is draft-side masking available on this hook? Probed ONCE per burst, before the draft
748    /// graph is captured (the mask is an in-graph node — its presence is a capture-time shape).
749    fn draft_mask_enabled(&self) -> bool {
750        false
751    }
752    /// Start a draft chain: clone the CURRENT (committed) grammar state into the speculative
753    /// slot. Called once per spec round, before the first draft position.
754    fn draft_begin(&mut self) -> Result<(), String> {
755        Ok(())
756    }
757    /// Packed 32-bit bitset words of the SPECULATIVE state's allowed set (target-vocab ids),
758    /// for the draft position about to be sampled. `None` = draft masking off (no-op).
759    fn draft_mask_words(&mut self) -> Result<Option<Vec<u32>>, String> {
760        Ok(None)
761    }
762    /// Advance the SPECULATIVE state with a PROPOSED draft token. `false` = the chain cannot
763    /// continue (EOS proposed, or an unmasked position proposed something illegal) — the
764    /// engine stops drafting; the token already pushed still goes through verify.
765    fn draft_advance(&mut self, _tok: u32) -> Result<bool, String> {
766        Ok(false)
767    }
768}
769
770/// DRAFT-MASK UPLOAD (lane/draft-mask): pull the speculative state's allowed set (TARGET-id
771/// space) from the hook, project it into the DRAFT head's vocab space, and upload it into the
772/// stable device buffer the draft chain reads. Returns false when the chain must stop drafting:
773/// the hook handed out no mask, or NO draft-vocab row is grammar-legal at this position (a
774/// trimmed FR-Spec head genuinely cannot propose a legal token there — masking it would leave
775/// a fully-banned row whose argmax is meaningless, so the round drafts fewer tokens and the
776/// verify emits the masked argmax as usual).
777fn upload_draft_mask(
778    e: &Engine,
779    c: &mut dyn SpecConstraint,
780    dst: &mut CudaSlice<u32>,
781    d2t: Option<&Vec<u32>>,
782    d_vocab: usize,
783    words: usize,
784) -> Result<bool, Box<dyn std::error::Error>> {
785    let Some(tw) = c
786        .draft_mask_words()
787        .map_err(|e2| format!("constraint: {e2}"))?
788    else {
789        return Ok(false);
790    };
791    let bit = |t: usize| -> bool {
792        let w = t >> 5;
793        w < tw.len() && (tw[w] >> (t & 31)) & 1 == 1
794    };
795    let mut buf = vec![0u32; words];
796    match d2t {
797        // TRIMMED draft head: row i proposes target id d2t[i] — permute the mask accordingly.
798        Some(map) => {
799            for (i, &t) in map.iter().enumerate().take(d_vocab) {
800                if bit(t as usize) {
801                    buf[i >> 5] |= 1u32 << (i & 31);
802                }
803            }
804        }
805        // UNTRIMMED: draft ids ARE target ids; the packed words transfer verbatim (a short
806        // mask leaves the padded tail zeroed == banned, same rule as constrained::apply_mask).
807        None => {
808            let n = tw.len().min(words);
809            buf[..n].copy_from_slice(&tw[..n]);
810        }
811    }
812    if buf.iter().all(|w| *w == 0) {
813        return Ok(false);
814    }
815    e.htod_u32_into(dst, &buf)?;
816    Ok(true)
817}
818
819/// Keep the full token-embedding table in host memory and upload only the rows needed by each
820/// MTP/verify step. This is an exact memory-capacity seam for very large BF16 vocab tables: host
821/// gather expands the same source bits to f32, and only O(T*n_embd) bytes cross PCIe per step.
822/// CUDA-graph/round-stream draft paths require device token ids and therefore stay disabled.
823pub(crate) fn spec_host_embd() -> bool {
824    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
825    *ON.get_or_init(|| std::env::var("MEMRA_SPEC_HOST_EMBD").as_deref() == Ok("1"))
826}
827
828/// VERIFY-TIER TRUNK LAUNCH-FUSION (default ON since 2026-07-09; MEMRA_SPEC_FUSED_T=0 reverts — lane/close35b): extend
829/// the t=1 fused2/fused3 Q8_0 trunk launches to the batched verify tier (t=2-4, the K=1..3
830/// verify shapes). At t>1 the trunk pairs/triples (35B wqkv+wqkv_gate, wq/wk/wv,
831/// gate_shexp+up_shexp) each run a separate `matmul_decode_exact` — one q8_1 re-quantize of the
832/// SAME activation plus one _b2/_b4 launch per tensor. The fused twins share ONE quantize and
833/// ONE launch per group; per (tensor,token,row) the kernel body is q8_0_mmvq_batched verbatim
834/// with the identical row mapping -> BIT-IDENTICAL by construction (kernel-check pins it,
835/// run-spec K=1..8 + acceptance identity arbitrate e2e).
836pub(crate) fn spec_fused_t() -> bool {
837    static F: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
838    // DEFAULT ON since 2026-07-09 (MEMRA_SPEC_FUSED_T=0 reverts): verify t=2-4 trunk launch-fusion
839    // (fused2/fused3 Q8_0 batched twins, bit-identical by construction — m=1 block-offset split on
840    // the batched body). m=2 marginal token 2117->1762us; 35B daily: p3 +3.7% (crosses llama), p2 +5%.
841    *F.get_or_init(|| {
842        std::env::var("MEMRA_SPEC_FUSED_T")
843            .map(|v| v != "0")
844            .unwrap_or(true)
845    })
846}
847
848/// zeros/uninit switch for verify-path buffers that are FULLY OVERWRITTEN before any read.
849/// Only call this on such buffers — the lean contract is "identical bytes by construction".
850/// TOKEN-ID GUARD for every id that reaches an embed gather (#87 family).
851///
852/// A device argmax seeds its running index with 0x7FFFFFFF and replaces it only through
853/// comparisons, all of which are FALSE against NaN. An all-NaN logits row therefore returns
854/// the sentinel, and the next thing done with a token id is `embed_row(id)` — table +
855/// ~4.6 TB, never mapped, an MMU fault that kills the CUDA context for the whole process
856/// (research/pp2spec-crash-20260807). The draft chain and the GREEDY verify walk already
857/// trap this; the SAMPLED verify bonus, the boundary sampler and the replay arm's last_pred
858/// did not, which is why the recoverable fault on the greedy instrument is a TERMINAL one on
859/// the vendor-default sampled shape we actually serve.
860pub(crate) fn guard_vocab_token(
861    tok: u32,
862    n_vocab: usize,
863    what: &str,
864) -> Result<u32, Box<dyn std::error::Error>> {
865    if (tok as usize) >= n_vocab {
866        return Err(format!(
867            "{what}: token id 0x{tok:08x} >= n_vocab {n_vocab} — an all-NaN logits row left \
868             the device argmax's init sentinel in place; refusing to dereference the embed \
869             row (#87 trap)"
870        )
871        .into());
872    }
873    Ok(tok)
874}
875
876/// SPEC NaN-ORIGIN SCAN (`MEMRA_SPEC_NAN_SCAN=1`, DEFAULT OFF, diagnostic only).
877///
878/// The `#87` trap reports an all-NaN VERIFY logits column, which says the poison reached the
879/// head but not where it entered. With the scan armed the verify walk syncs and reads back
880/// every layer's output, so the FIRST layer whose residual carries a NaN names itself with the
881/// round's row and position. Off by default and never on a serving path: it costs one host
882/// sync + one `t*n_embd` D2H per layer, and the syncs change scheduling (so a run that stops
883/// reproducing under the scan is itself a datum, not an all-clear).
884///
885/// Rollback seam: unset `MEMRA_SPEC_NAN_SCAN` (or set it to 0). Every call site is behind
886/// `spec_nan_scan()`, so the default path keeps the exact launch sequence it had.
887pub(crate) fn spec_nan_scan() -> bool {
888    spec_nan_scan_level() > 0
889}
890
891/// `MEMRA_SPEC_NAN_SCAN` as a LEVEL, not a boolean. `1` scans each layer's residual, which
892/// names the layer. `2` also scans INSIDE the t-column layer body — the per-column attention
893/// output, the deferred-column o-proj/fa2 join, the post-attention norm and the routed-MoE
894/// output — because "layer 20 poisons row 0" does not say whether the attention or the routed
895/// MoE produced it, and those are different bugs with different fixes.
896pub(crate) fn spec_nan_scan_level() -> u8 {
897    static LVL: std::sync::OnceLock<u8> = std::sync::OnceLock::new();
898    *LVL.get_or_init(|| match std::env::var("MEMRA_SPEC_NAN_SCAN").as_deref() {
899        Ok("1") => 1,
900        Ok("2") => 2,
901        _ => 0,
902    })
903}
904
905/// Read back `[rows, cols]` and fail with the first NaN's coordinates. `what` names the
906/// producer (layer index, walk arm) so the error line is the localization.
907/// VERIFY-ARM RECEIPT (rides `MEMRA_SPEC_NAN_SCAN>=1`, bounded to 200 lines).
908///
909/// Names, per trunk layer, WHICH attention arm the t-column walk actually took. This exists
910/// because the level-1 residual scan below sat only on the non-fused tail: the fused
911/// rope+append+fa arm ends in `continue`, so every layer that fused was NEVER SCANNED and
912/// silently read as "clean". A poisoned residual therefore first reported at the next
913/// non-fused layer, which is how "layer 20 creates the poison" could be true of the scan and
914/// false of the engine. Also carries the row-table lookup counter, so "the fused path never
915/// ran" is distinguishable from "it ran and was innocent".
916/// KV-PLANE SCAN (`MEMRA_KV_PLANE_SCAN=1`, DEFAULT OFF, diagnostic only).
917///
918/// Reads back the STAGED rows of a layer's distributed K/V planes and reports the first row
919/// whose quantization scale is not finite. No kernel required: q8_0 blocks are
920/// `[half d][32 x i8]` and q5_1 blocks carry `half d` then `half m`, so the fp16 scale at the
921/// head of each block is host-checkable straight out of the byte plane.
922///
923/// It exists because the level-2 bad-row bitmap says EVERY verify row is non-finite at a
924/// global-attention layer's join, and row r attends a strict superset of row r-1's keys: that
925/// implicates the shared KV history those rows walk, not per-column staging. "The attention
926/// output is NaN" and "the KV history it attends is already NaN" are different bugs with
927/// different owners, and nothing measured so far separates them. A first-corrupt-row index
928/// also dates the corruption against the prime/decode boundary.
929///
930/// Bounded hard: only layers whose geometry has NO window (the global planes), only the first
931/// `MEMRA_KV_PLANE_SCAN_ROUNDS` verify rounds of a process (default 2), and it copies only
932/// `[0, staged_len)`, which is ~1.6 MB at the 1480-token repro rather than the 262144-row
933/// provision. It still syncs per layer, so it is never a serving or a measured-perf arm.
934pub(crate) fn kv_plane_scan_on() -> bool {
935    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
936    *ON.get_or_init(|| std::env::var("MEMRA_KV_PLANE_SCAN").as_deref() == Ok("1"))
937}
938
939fn kv_plane_scan_rounds() -> usize {
940    static R: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
941    *R.get_or_init(|| {
942        std::env::var("MEMRA_KV_PLANE_SCAN_ROUNDS")
943            .ok()
944            .and_then(|v| v.parse().ok())
945            .unwrap_or(2)
946    })
947}
948
949/// First non-finite fp16 block scale in `bytes`, as (block index, raw u16), scanning one
950/// scale every `stride` bytes. Returns None when every block scale is finite.
951fn first_bad_scale(bytes: &[u8], stride: usize) -> Option<(usize, u16)> {
952    if stride == 0 {
953        return None;
954    }
955    for (i, blk) in bytes.chunks_exact(stride).enumerate() {
956        let raw = u16::from_le_bytes([blk[0], blk[1]]);
957        if half_is_non_finite(raw) {
958            return Some((i, raw));
959        }
960    }
961    None
962}
963
964/// IEEE binary16: exponent all ones is Inf or NaN, whatever the mantissa says.
965fn half_is_non_finite(raw: u16) -> bool {
966    (raw & 0x7C00) == 0x7C00
967}
968
969/// Scan one layer's staged K/V planes for a non-finite quantization scale. Returns the
970/// receipt line, or None when the layer is out of scope or every scale is finite.
971pub(crate) fn scan_kv_plane(
972    e: &crate::Engine,
973    distributed: &memra_kv::ResidentTpKvCache,
974    il: usize,
975    pos0: usize,
976) -> Result<(), Box<dyn std::error::Error>> {
977    // One "round" is one pos0, not one layer: the walk visits 45 layers per verify. The
978    // default of 2 rounds is for a fault that shows up immediately; the step37 repro does not
979    // fire until rep 3 or later, i.e. round ~60 of the process, so that arm MUST raise
980    // MEMRA_KV_PLANE_SCAN_ROUNDS or it will scan only the two rounds that were never going to
981    // be poisoned and report a clean history it never looked at.
982    static ROUNDS: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
983    static LAST_POS: std::sync::atomic::AtomicUsize =
984        std::sync::atomic::AtomicUsize::new(usize::MAX);
985    if LAST_POS.swap(pos0, std::sync::atomic::Ordering::Relaxed) != pos0 {
986        ROUNDS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
987    }
988    if ROUNDS.load(std::sync::atomic::Ordering::Relaxed) > kv_plane_scan_rounds() {
989        return Ok(());
990    }
991    let staged = distributed.staged_len();
992    if staged == 0 {
993        return Ok(());
994    }
995    // ENGAGEMENT RECEIPT. This scan prints only on corruption, so `kvbad=0` in a cell would
996    // read the same whether the history was clean or the scan never ran once. Bounded so a
997    // 45-layer walk cannot flood the log.
998    static SEEN: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
999    let seen = SEEN.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
1000    let (ktb, vtb) = (distributed.k_tok_bytes(), distributed.v_tok_bytes());
1001    if seen < 4 {
1002        eprintln!(
1003            "[kv-plane] engaged #{seen} layer {il} pos0={pos0} staged={staged} \
1004             ktok={ktb} vtok={vtb} (scan armed; a corrupt plane prints its own line)"
1005        );
1006    }
1007    for rank in 0..distributed.ranks().len() {
1008        let Some(rc) = distributed.rank(rank) else {
1009            continue;
1010        };
1011        // q8_0 K blocks are [half d][32 x i8] = 34B; q5_1 V blocks lead with half d then half m.
1012        let kbytes = e.dtoh_u8_view(&rc.k().slice(0..staged * ktb))?;
1013        let vbytes = e.dtoh_u8_view(&rc.v().slice(0..staged * vtb))?;
1014        let kbad = first_bad_scale(&kbytes, 34);
1015        let vbad = first_bad_scale(&vbytes, 24);
1016        if kbad.is_some() || vbad.is_some() {
1017            let row = |b: Option<(usize, u16)>, tok: usize| {
1018                b.map(|(i, raw)| format!("blk {i} (row {}) raw={raw:#06x}", i * 34 / tok.max(1)))
1019                    .unwrap_or_else(|| "clean".into())
1020            };
1021            eprintln!(
1022                "[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",
1023                row(kbad, ktb),
1024                row(vbad, vtb)
1025            );
1026            return Ok(());
1027        }
1028    }
1029    Ok(())
1030}
1031
1032pub(crate) fn verify_arm_receipt(
1033    arm: &str,
1034    il: usize,
1035    pos0: usize,
1036    t: usize,
1037    staged: Option<usize>,
1038) {
1039    static N: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
1040    if N.fetch_add(1, std::sync::atomic::Ordering::Relaxed) >= 200 {
1041        return;
1042    }
1043    eprintln!(
1044        "[verify-arm] layer {il} arm={arm} pos0={pos0} t={t} staged_len={} rows_tab_lookups={}",
1045        staged.map(|v| v as i64).unwrap_or(-1),
1046        crate::tp::ROWS_TAB_ENGAGED.load(std::sync::atomic::Ordering::Relaxed)
1047    );
1048}
1049
1050pub(crate) fn nan_scan_rows(
1051    e: &Engine,
1052    buf: &CudaSlice<f32>,
1053    rows: usize,
1054    cols: usize,
1055    what: &str,
1056) -> Result<(), Box<dyn std::error::Error>> {
1057    // The readback is also the ATTRIBUTION point for an asynchronous fault: a
1058    // CUDA_ERROR_ILLEGAL_ADDRESS raised by any launch since the previous scan surfaces on this
1059    // sync, and the bare DriverError names nothing. Wrapping it with `what` turns "the process
1060    // died somewhere" into "it died at or before this layer, on this row, at this position".
1061    let host = e.dtoh(buf).map_err(|err| -> Box<dyn std::error::Error> {
1062        format!(
1063            "spec nan-scan: sync at {what} FAILED: {err} — the fault is at or before \
1064                     this point in the walk"
1065        )
1066        .into()
1067    })?;
1068    if host.len() < rows * cols {
1069        return Err(format!(
1070            "nan-scan {what}: buffer holds {} < {rows}x{cols}",
1071            host.len()
1072        )
1073        .into());
1074    }
1075    // SCAN EVERY ROW BEFORE REPORTING. A first-hit return says "row 0 is bad" and leaves the
1076    // other rows UNEXAMINED, which is exactly the bit that discriminates the two mechanisms: in
1077    // the t-column verify, row 0 attends keys [0..p+1) and row 1 attends [0..p+2), a strict
1078    // superset, so poison in the SHARED KV history must appear in BOTH rows, while poison in
1079    // per-column staging can appear in one. Report the whole map.
1080    let mut per_row: Vec<usize> = Vec::with_capacity(rows);
1081    let mut first_bad: Option<(usize, usize)> = None;
1082    for r in 0..rows {
1083        let row = &host[r * cols..(r + 1) * cols];
1084        let bad = row.iter().filter(|v| !v.is_finite()).count();
1085        per_row.push(bad);
1086        if bad > 0 && first_bad.is_none() {
1087            first_bad = Some((r, row.iter().position(|v| !v.is_finite()).unwrap_or(0)));
1088        }
1089    }
1090    if let Some((r0, c0)) = first_bad {
1091        let map: String = per_row
1092            .iter()
1093            .map(|&b| if b == 0 { '.' } else { 'X' })
1094            .collect();
1095        return Err(format!(
1096            "spec nan-scan: {what} produced non-finite values — rows[{rows}] map={map} \
1097             counts={per_row:?} of {cols} each; first at row {r0} element {c0}. Both rows bad \
1098             implicates shared state (the KV history this layer reads); one row bad implicates \
1099             per-column staging."
1100        )
1101        .into());
1102    }
1103    Ok(())
1104}
1105
1106fn vbuf(e: &Engine, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
1107    if spec_lean() { e.uninit(n) } else { e.zeros(n) }
1108}
1109
1110/// Scratch KV for the MTP block (one full-attn layer).
1111///
1112/// PERSISTENT MODE (default, 2026-07-03 — the acceptance lever): sized cap = max_ctx and kept in
1113/// sync with the COMMITTED sequence — slot p holds the MTP block's K/V for committed token p
1114/// (roped p+1, the chain's rope convention), so the draft chain's self-attention sees the FULL
1115/// committed history instead of only the current round's 1..K+1 chain tokens (the reference
1116/// engine's "mtp_update" design). Entries come from two sources:
1117///   - chain appends: accepted positions KEEP their chain-computed entries (embedding exact,
1118///     hidden chain-approximate — the reference engine accepts the same);
1119///   - `mtp_kv_fill` batches: prompt positions + the last-draft position on full accept, computed
1120///     from EXACT trunk hiddens (K/V-only MTP-block pass, no attention/FFN/lm_head).
1121///     Rejected drafts / p-min extras / pseudo-seed appends are all discarded by the round-start
1122///     `set_len` truncation (the KvLayer len mechanism — §C rollback for the draft side).
1123///     Multi-turn spec-decode session (2026-07-05): trunk Cache + persistent MTP draft scratch +
1124///     the committed token list, alive across generate_spec_session calls. Turn N+1 primes ONLY its
1125///     suffix (chunked continuation prime over the quantized past) and mtp_kv_fill's its suffix rows,
1126///     then the round loop runs unchanged. `last_h` carries the pre-output_norm hidden of the last
1127///     committed row across turns (the predecessor-pairing seed + fill anchor).
1128///     Per-request sampling config for the sampled-spec serve path.
1129#[derive(Clone, Copy, Debug)]
1130pub struct SpecSampling {
1131    pub temp: f32,
1132    pub seed: u64,
1133    pub top_k: i32,            // 0 = off
1134    pub top_p: f32,            // 1.0 = off
1135    pub min_p: f32,            // 0.0 = off
1136    pub penalty_last_n: usize, // 0 = penalties off
1137    pub penalty_repeat: f32,
1138    pub penalty_freq: f32,
1139    pub penalty_present: f32,
1140}
1141
1142impl SpecSampling {
1143    /// Non-identity penalties requested — THE `pen_on` predicate (one definition; the
1144    /// same group-off rule `SamplerIdentity::of` canonicalizes: a window with neutral
1145    /// coefficients is penalties-absent). Both spec routes and the dspark accept walk
1146    /// key their penalty arms off this.
1147    pub fn pen_on(&self) -> bool {
1148        self.penalty_last_n > 0
1149            && (self.penalty_repeat != 1.0
1150                || self.penalty_freq != 0.0
1151                || self.penalty_present != 0.0)
1152    }
1153}
1154
1155/// Which draft source a spec session is pinned to. The ENGINE-LEVEL half of
1156/// `DraftSourcePlan` (memra-gguf `model_plan.rs`, always general): the plan states what the
1157/// model DECLARES, this states what actually LOADED and therefore what the session runs.
1158/// Pinned at session creation for the session's lifetime.
1159///
1160/// Family-agnostic on purpose (lane/glm5-extract2, the DraftSource seam): glm5 is today's
1161/// consumer with NativeMtp | Dflash2; the hy3/qwen-next spec lanes select through the same
1162/// three-way law instead of re-deriving it. What each family still owns is the per-session
1163/// STATE behind the kind (see `dflash.rs`'s seam note for why that half is not a trait yet).
1164#[derive(Clone, Copy, Debug, PartialEq, Eq)]
1165pub enum DraftSourceKind {
1166    /// The model's own embedded NextN/MTP head.
1167    NativeMtp,
1168    /// A separately loaded DFlash2 block-diffusion drafter
1169    /// ([`crate::dflash::DflashDrafter`]).
1170    Dflash2,
1171}
1172
1173/// The uniform draft-source selection law. Pure — no env, no engine, no family types — so it
1174/// is CPU-gateable and so every spec family answers "which source" the same way.
1175///
1176/// THE LAW, in precedence order:
1177/// 1. A LOADED DFlash2 drafter IS the source. The operator asked for it by name (a set
1178///    drafter flag that cannot load is already a loud boot failure, never a silent
1179///    fallback), and the family's embedded head is deliberately NOT loaded for this source —
1180///    it is a full trunk layer of VRAM.
1181/// 2. Otherwise the embedded head, and only when the PLAN declares an embedded source: a
1182///    loaded head under a plan that does not declare `Embedded` is a load-path bug, not a
1183///    draft source, and it is refused by name rather than drafted from.
1184/// 3. Otherwise there is no draft source and speculative decode must refuse before drafting.
1185pub fn resolve_draft_source_kind(
1186    plan: memra_gguf::model_plan::DraftSourcePlan,
1187    embedded_head_loaded: bool,
1188    dflash_loaded: bool,
1189) -> Result<DraftSourceKind, String> {
1190    use memra_gguf::model_plan::DraftSourcePlan as P;
1191    if dflash_loaded {
1192        return Ok(DraftSourceKind::Dflash2);
1193    }
1194    if embedded_head_loaded {
1195        if plan != P::Embedded {
1196            return Err(format!(
1197                "an embedded draft head is loaded but the ModelPlan declares \
1198                 draft_source={plan:?} — refused rather than drafting from a head the plan \
1199                 does not claim"
1200            ));
1201        }
1202        return Ok(DraftSourceKind::NativeMtp);
1203    }
1204    Err(format!(
1205        "no draft source loaded (ModelPlan declares draft_source={plan:?}): speculative \
1206         decode has nothing to draft from"
1207    ))
1208}
1209
1210#[cfg(test)]
1211mod draft_source_kind_tests {
1212    use super::{DraftSourceKind, resolve_draft_source_kind};
1213    use memra_gguf::model_plan::DraftSourcePlan as P;
1214
1215    #[test]
1216    fn a_loaded_drafter_wins_over_a_co_loaded_embedded_head() {
1217        // The operator asked for the drafter BY NAME (a set drafter flag that cannot load is
1218        // already a loud boot failure), so it takes precedence under every plan value —
1219        // including ExternalArtifact, which is what a pack declares when the draft weights
1220        // are not in the model file.
1221        for plan in [P::Embedded, P::ExternalArtifact, P::None] {
1222            assert_eq!(
1223                resolve_draft_source_kind(plan, true, true).unwrap(),
1224                DraftSourceKind::Dflash2,
1225                "plan {plan:?}: a loaded drafter must win"
1226            );
1227            assert_eq!(
1228                resolve_draft_source_kind(plan, false, true).unwrap(),
1229                DraftSourceKind::Dflash2
1230            );
1231        }
1232    }
1233
1234    #[test]
1235    fn the_embedded_head_is_the_source_only_under_a_plan_that_claims_it() {
1236        assert_eq!(
1237            resolve_draft_source_kind(P::Embedded, true, false).unwrap(),
1238            DraftSourceKind::NativeMtp
1239        );
1240        // A head loaded under a plan that does not declare Embedded is a LOAD-PATH BUG, not a
1241        // draft source. Unreachable on glm5 today (its pack hardcodes Embedded and the head
1242        // only loads under it) — which is exactly why it is pinned here: an unreachable
1243        // refusal with no arm is an untested refusal, and the next family is the one that
1244        // makes it reachable.
1245        for plan in [P::ExternalArtifact, P::None] {
1246            let err = resolve_draft_source_kind(plan, true, false)
1247                .expect_err("a head under a non-Embedded plan must refuse");
1248            assert!(err.contains("does not claim"), "{err}");
1249            assert!(err.contains(&format!("{plan:?}")), "{err}");
1250        }
1251    }
1252
1253    #[test]
1254    fn nothing_loaded_refuses_before_drafting_and_names_the_plan() {
1255        for plan in [P::Embedded, P::ExternalArtifact, P::None] {
1256            let err =
1257                resolve_draft_source_kind(plan, false, false).expect_err("no source must refuse");
1258            assert!(err.contains("no draft source loaded"), "{err}");
1259            assert!(err.contains(&format!("{plan:?}")), "{err}");
1260        }
1261    }
1262}
1263
1264/// `MEMRA_SPEC_PMIN` break semantics over per-slot draft confidences (the chain break this
1265/// module's drafting loops apply inline: `p < p_min && (j > 0 || pmin0)`): keep the longest
1266/// prefix whose every slot clears `p_min`; slot 0 survives a miss unless PMIN0 arms
1267/// zero-draft rounds. Prefix truncation is forced by the accept rule anyway (a kept slot
1268/// after a dropped one could never commit — the dspark confidence-slot argument). Pure so
1269/// the rule is CPU-gateable; the SHARED K-policy surface every spec family consumes
1270/// (hoisted from the glm5 loop, lane/glm5-extract-general).
1271pub fn spec_conf_keep(q: &[f32], p_min: f32, pmin0: bool) -> usize {
1272    if p_min <= 0.0 {
1273        return q.len();
1274    }
1275    let mut kept = 0usize;
1276    for (j, &qj) in q.iter().enumerate() {
1277        if qj < p_min && (j > 0 || pmin0) {
1278            break;
1279        }
1280        kept += 1;
1281    }
1282    kept
1283}
1284
1285/// Host Philox4x32-10 uniform in (0,1) — mirrors spec_sample.cu's `philox4`/`u01` with the
1286/// ctr_lo tag 0xFFFF_FFFE, so the host accept-test stream never collides with any device
1287/// sampling event (device Gumbel uses (i>>2, stream_pos); device residual uses 0xFFFF_FFFD).
1288/// One value per (seed, ctr) EVENT; callers own the counter discipline. Extracted verbatim
1289/// from generate_spec_inner2's closure for the dspark sampled-admission walk (the two paths
1290/// MUST consume the identical stream construction — two ad-hoc Philox copies drifting apart
1291/// is a distributional bug, not a style problem).
1292pub(crate) fn host_u01(seed: u64, ctr: u32) -> f32 {
1293    let (m0, m1) = (0xD2511F53u32, 0xCD9E8D57u32);
1294    let (mut c0, mut c1, mut c2, mut c3) = (0xFFFF_FFFEu32, ctr, 0u32, 0u32);
1295    let (mut k0, mut k1) = ((seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
1296    for _ in 0..10 {
1297        let (h0, l0) = (((m0 as u64 * c0 as u64) >> 32) as u32, m0.wrapping_mul(c0));
1298        let (h1, l1) = (((m1 as u64 * c2 as u64) >> 32) as u32, m1.wrapping_mul(c2));
1299        let (n0, n1, n2, n3) = (h1 ^ c1 ^ k0, l1, h0 ^ c3 ^ k1, l0);
1300        c0 = n0;
1301        c1 = n1;
1302        c2 = n2;
1303        c3 = n3;
1304        k0 = k0.wrapping_add(0x9E3779B9);
1305        k1 = k1.wrapping_add(0xBB67AE85);
1306    }
1307    (c0 as f32 + 1.0) * (1.0 / 4294967296.0)
1308}
1309
1310/// Tracked draft positions for [`SpecTelemetry`] (serve K defaults to 3; the run-spec gate
1311/// sweeps K=1..8, and MEMRA_SPEC_CAPMAX defaults to 7 — 8 covers every tuned config).
1312pub const SPEC_TELEM_POS: usize = 8;
1313
1314/// Always-on per-draft-position acceptance telemetry (lane/accept-telemetry, 2026-08-05 —
1315/// the llama.cpp #26389 / vLLM spec-decode counter schema, upstream-sweeps 2026-08-05).
1316/// Lives on the [`SpecSession`] and accumulates across bursts; the serve worker diffs a
1317/// stashed copy per burst for its per-model /metrics aggregation and per-request usage.
1318/// Same normalization as the `[spec-stats]` line: p-min-discarded chain tokens are counted
1319/// in NEITHER drafted nor accepted.
1320#[derive(Clone, Copy, Default, Debug)]
1321pub struct SpecTelemetry {
1322    /// verify rounds completed (a round-stream burst counts each of its M rounds).
1323    pub rounds: u64,
1324    /// tokens drafted / accepted across all rounds.
1325    pub drafted: u64,
1326    pub accepted: u64,
1327    /// how often draft position j (0-based within a round's chain) was offered / accepted.
1328    /// Positions >= SPEC_TELEM_POS are untracked (totals still count them). The opt-in
1329    /// round-stream arm (MEMRA_SPEC_STREAM=1) reads back only totals, so under it these
1330    /// arrays cover the standard-path rounds only and their sums may undercount the totals.
1331    pub pos_drafted: [u64; SPEC_TELEM_POS],
1332    pub pos_accepted: [u64; SPEC_TELEM_POS],
1333}
1334
1335impl SpecTelemetry {
1336    /// Fieldwise `self - prev` — the worker's per-burst delta off a copy stashed before the
1337    /// burst call. Saturating: a caller diffing against the wrong snapshot gets zeros, not
1338    /// a wrapped counter.
1339    pub fn delta_since(&self, prev: &SpecTelemetry) -> SpecTelemetry {
1340        let mut d = SpecTelemetry {
1341            rounds: self.rounds.saturating_sub(prev.rounds),
1342            drafted: self.drafted.saturating_sub(prev.drafted),
1343            accepted: self.accepted.saturating_sub(prev.accepted),
1344            ..Default::default()
1345        };
1346        for j in 0..SPEC_TELEM_POS {
1347            d.pos_drafted[j] = self.pos_drafted[j].saturating_sub(prev.pos_drafted[j]);
1348            d.pos_accepted[j] = self.pos_accepted[j].saturating_sub(prev.pos_accepted[j]);
1349        }
1350        d
1351    }
1352    /// Fieldwise `self += d` — the worker's per-model aggregation.
1353    pub fn merge(&mut self, d: &SpecTelemetry) {
1354        self.rounds += d.rounds;
1355        self.drafted += d.drafted;
1356        self.accepted += d.accepted;
1357        for j in 0..SPEC_TELEM_POS {
1358            self.pos_drafted[j] += d.pos_drafted[j];
1359            self.pos_accepted[j] += d.pos_accepted[j];
1360        }
1361    }
1362
1363    /// Mean accepted draft-prefix length per verify round (tau).
1364    pub fn tau(&self) -> f64 {
1365        if self.rounds > 0 {
1366            self.accepted as f64 / self.rounds as f64
1367        } else {
1368            0.0
1369        }
1370    }
1371}
1372
1373/// Session-lifetime atomic acceptance counters. The verifier records only after the greedy or
1374/// rejection-sampling walk has resolved on the host, so these relaxed increments add no GPU
1375/// launch, synchronization, allocation, or ordering dependency to the numeric path.
1376struct SpecTelemetryCounters {
1377    rounds: AtomicU64,
1378    drafted: AtomicU64,
1379    accepted: AtomicU64,
1380    pos_drafted: [AtomicU64; SPEC_TELEM_POS],
1381    pos_accepted: [AtomicU64; SPEC_TELEM_POS],
1382}
1383
1384impl Default for SpecTelemetryCounters {
1385    fn default() -> Self {
1386        Self {
1387            rounds: AtomicU64::new(0),
1388            drafted: AtomicU64::new(0),
1389            accepted: AtomicU64::new(0),
1390            pos_drafted: std::array::from_fn(|_| AtomicU64::new(0)),
1391            pos_accepted: std::array::from_fn(|_| AtomicU64::new(0)),
1392        }
1393    }
1394}
1395
1396impl SpecTelemetryCounters {
1397    fn record_round(&self, drafted: usize, accepted: usize) {
1398        debug_assert!(accepted <= drafted);
1399        self.rounds.fetch_add(1, Ordering::Relaxed);
1400        self.drafted.fetch_add(drafted as u64, Ordering::Relaxed);
1401        self.accepted.fetch_add(accepted as u64, Ordering::Relaxed);
1402        for counter in self.pos_drafted.iter().take(drafted) {
1403            counter.fetch_add(1, Ordering::Relaxed);
1404        }
1405        for counter in self.pos_accepted.iter().take(accepted) {
1406            counter.fetch_add(1, Ordering::Relaxed);
1407        }
1408    }
1409
1410    /// Round-stream keeps each round's accept length on device; retain exact scalar totals while
1411    /// leaving the per-position arrays untouched, matching the pre-existing telemetry contract.
1412    fn record_totals(&self, rounds: usize, drafted: usize, accepted: usize) {
1413        self.rounds.fetch_add(rounds as u64, Ordering::Relaxed);
1414        self.drafted.fetch_add(drafted as u64, Ordering::Relaxed);
1415        self.accepted.fetch_add(accepted as u64, Ordering::Relaxed);
1416    }
1417
1418    fn snapshot(&self) -> SpecTelemetry {
1419        SpecTelemetry {
1420            rounds: self.rounds.load(Ordering::Relaxed),
1421            drafted: self.drafted.load(Ordering::Relaxed),
1422            accepted: self.accepted.load(Ordering::Relaxed),
1423            pos_drafted: std::array::from_fn(|j| self.pos_drafted[j].load(Ordering::Relaxed)),
1424            pos_accepted: std::array::from_fn(|j| self.pos_accepted[j].load(Ordering::Relaxed)),
1425        }
1426    }
1427}
1428
1429pub struct SpecSession {
1430    pub(crate) cache: Cache,
1431    pub(crate) scratch: MtpScratch,
1432    /// Every token whose state the caches hold, in order (prompt turns + generated), INCLUDING
1433    /// overshoot: spec commits accepted drafts past max_new; those rows are in the caches, so the
1434    /// session must count them. Callers render output from this, not from their own echo.
1435    pub committed: Vec<u32>,
1436    /// Pre-output_norm hidden of the LAST committed row (device). None before the first turn.
1437    pub(crate) last_h: Option<CudaSlice<f32>>,
1438    /// Greedy argmax predicting the token AFTER committed.last() (from the last turn's final
1439    /// logits). Fuels empty-suffix continuation bursts (serve): the next turn emits this token
1440    /// first, feeds it, and the round loop resumes without any prime. None before the first turn.
1441    pub next_pred: Option<u32>,
1442    /// SAMPLED-SPEC stream continuity across bursts: Philox event counters persist here so a
1443    /// session's randomness never repeats between generate_spec_session calls. (0,0) at admit.
1444    pub sctr: u32,
1445    pub uctr: u32,
1446    /// PERSISTENT DRAFT-GRAPH CONTEXT (2026-08-01, the serve-burst fixed-cost fix): the captured
1447    /// draft graph(s) + every device I/O buffer they bake, carried ACROSS generate_spec_session
1448    /// calls. Before this, every serve burst re-captured the draft graph (2 warmup forwards +
1449    /// instantiate) — measured ~16ms/burst on H100 q27 (MEMRA_SPEC_BURST sweep,
1450    /// research/spec-serving-20260801). None before the first turn; error paths drop it
1451    /// (next burst recaptures — serve retires errored sessions anyway).
1452    pub(crate) draft_ctx: Option<DraftGraphCtx>,
1453    /// PENDING-CARRY across bursts (2026-08-01, the serve burst-boundary fix): the bonus token
1454    /// emitted by the last round but NOT committed to the caches. The old tail committed it with
1455    /// a solo T=1 trunk pass (+ draft fill), and the next burst's setup fed the stashed next_pred
1456    /// with ANOTHER solo pass — 2x ~11.5ms/burst measured on H100 q27 ([spec-setup] trace).
1457    /// Carrying it lets the next empty-suffix greedy burst consume it as round-0 verify col 0,
1458    /// exactly like a mid-burst full-accept boundary (no solo passes). INVARIANT: when set,
1459    /// `committed` (== cache rows) EXCLUDES this token although it was already emitted in the
1460    /// last burst's output, and `last_h` holds the hidden of the last COMMITTED row (its
1461    /// predecessor — the chain-seed/fill anchor). `next_pred` is None (unknown without the
1462    /// commit pass). Non-empty-suffix or sampled turns must flush first (spec_flush_pending);
1463    /// generate_spec_session_sampled does this at entry, and serve parks only flushed sessions.
1464    pub pending_tok: Option<u32>,
1465    /// SESSION-AFFINITY TURN CHECKPOINT (lane/session-affinity, 2026-08-05): the state at this
1466    /// turn's PROMPT-END boundary, retained so a later turn can REWIND here. See
1467    /// [`SpecCheckpoint`]. Refreshed by every non-empty prime; None until the first one, and on
1468    /// a rig too tight to hold it (a failed capture is silent — resume just isn't available).
1469    pub(crate) turn_ckpt: Option<SpecCheckpoint>,
1470    /// Session-lifetime acceptance telemetry. Relaxed atomics update at the host-side round
1471    /// accounting the loop already does — no syncs, no allocation. NOTE a
1472    /// pool-resumed session carries the PREVIOUS requests' counts; per-request consumers
1473    /// diff with [`SpecTelemetry::delta_since`] around each burst.
1474    telem: SpecTelemetryCounters,
1475    /// PREFIX-CACHE publication request (lane/spec-prefix-cache): worker sets this to the
1476    /// miss-LCP boundary before a cold burst; the prime captures at exactly that split (it must
1477    /// coincide with the burst's `prime_split` or no capture happens). One-shot: consumed by the
1478    /// prime, result lands in `boundary_captures`.
1479    pub capture_at: Option<usize>,
1480    /// The captures the last prime produced (see [`SpecBoundaryCapture`]). Worker drains them
1481    /// post-burst to assemble prefix entries. A failed capture is silent, like `turn_ckpt` —
1482    /// publication just isn't available for that request. Plural since
1483    /// lane/frspec-multiturn-cache (2026-08-21): a cold burst can capture BOTH the miss-LCP
1484    /// split (the shared-prefix class) and the stable pre-generation boundary (the
1485    /// next-turn re-render class) — one entry per stop, exactly the boundary set the plain
1486    /// prefill tick publishes/checkpoints.
1487    pub boundary_captures: Vec<SpecBoundaryCapture>,
1488    /// STABLE-BOUNDARY TURN CHECKPOINT REQUEST (lane/frspec-multiturn-cache, 2026-08-21): the
1489    /// ABSOLUTE committed-length position the next non-empty prime should capture `turn_ckpt`
1490    /// at, instead of prompt-end. The worker sets it to the STABLE PRE-GENERATION boundary
1491    /// (`plain_checkpoint_boundary` — before the live generation header the client rewrites),
1492    /// porting the 2026-08-09 plain-tier fix: a prompt-end spec checkpoint includes the
1493    /// template's live assistant-generation header (`<|im_start|>assistant\n<think>\n`), which
1494    /// the NEXT turn's re-render replaces, so `affinity_match` diverged a couple tokens below
1495    /// the checkpoint and the spec pool declined 100% of multi-turn agent traffic (measured:
1496    /// `spec-affinity: declined (history diverged at 6811 of checkpoint 6813)`,
1497    /// research/multiturn-cache-20260821 B4). One-shot, `capture_at` convention; None = legacy
1498    /// prompt-end capture.
1499    pub ckpt_at: Option<usize>,
1500    /// FAIL-SAFE (lane/step37-vram-admission-20260830, external-review corroboration): set
1501    /// by the worker on a session serving a step-OOM park REPLAY. The burst entry pre-marks
1502    /// the draft-graph fallback so the replay never re-enters the capture path — the capture
1503    /// appetite is part of what drove the card to the OOM, and a replay that recaptures
1504    /// re-runs the incident. If the eager replay still cannot fit, the bounded retry budget
1505    /// exhausts into the honest recoverable Overloaded error instead of looping.
1506    pub capture_disabled: bool,
1507}
1508impl SpecSession {
1509    /// Context capacity of the session's caches (the server's ContextFull guard).
1510    pub fn cache_max_ctx(&self) -> usize {
1511        self.cache.max_ctx
1512    }
1513    /// Read access to the live trunk cache (lane/spec-prefix-cache): the worker slices
1514    /// full-attn KV rows `[0..capture.pos)` out of it when publishing a boundary capture —
1515    /// those rows are append-only for the session's lifetime (rollbacks never truncate below
1516    /// the prime boundary), so no copy was taken at prime time.
1517    pub fn cache_ref(&self) -> &Cache {
1518        &self.cache
1519    }
1520    /// Read access to the persistent draft-scratch plane (lane/spec-on-cache-hit): the
1521    /// worker slices rows `[0..capture.pos)` when publishing a boundary capture, exactly
1522    /// like the trunk KV — draft rows below the prompt end are append-only for the
1523    /// session's lifetime (the prime fill wrote them once; rollbacks reset `len_d` to the
1524    /// committed length, never below the prime boundary, and the true-hidden refresh
1525    /// rewrites generated positions only). Returns `(k, v, k_tok_bytes, v_tok_bytes)`.
1526    /// None when the scratch is ring-backed (Step35 SWA — physical rows are not
1527    /// prefix-addressable; the prefix cache already refuses that class end to end).
1528    pub fn draft_plane_ref(&self) -> Option<(&CudaSlice<u8>, &CudaSlice<u8>, usize, usize)> {
1529        if self.scratch.kv.ring.is_some() {
1530            return None;
1531        }
1532        Some((
1533            &self.scratch.kv.k,
1534            &self.scratch.kv.v,
1535            self.scratch.kv.k_tok_bytes,
1536            self.scratch.kv.v_tok_bytes,
1537        ))
1538    }
1539    /// Snapshot the session's process-local acceptance counters for per-burst diffing.
1540    pub fn telemetry(&self) -> SpecTelemetry {
1541        self.telem.snapshot()
1542    }
1543    /// Committed position this session can REWIND to (its retained prompt-end boundary), if any.
1544    /// A request whose prompt matches `committed[..pos]` exactly can resume from here — see
1545    /// `spec_rewind_to_checkpoint`.
1546    pub fn rewind_pos(&self) -> Option<usize> {
1547        self.turn_ckpt.as_ref().map(|c| c.pos)
1548    }
1549    /// Whether every ring-backed trunk/draft row needed by the retained checkpoint is resident.
1550    pub fn rewind_is_resident(&self) -> bool {
1551        self.turn_ckpt.as_ref().is_some_and(|ckpt| {
1552            self.cache.can_rollback(&ckpt.snap, 0) && self.scratch.can_rewind_to(ckpt.pos)
1553        })
1554    }
1555    /// Is this session in the DEMOTION-READY shape (see [`SpecSession::into_demoted`])?
1556    /// `false` means a carried pending must be flushed first (`spec_flush_pending`), or the
1557    /// session has never run a turn and has no prediction to hand over.
1558    pub fn demote_ready(&self) -> bool {
1559        self.pending_tok.is_none() && self.next_pred.is_some()
1560    }
1561    /// Does this session hold a carried pending bonus (flush required before a handoff/park)?
1562    pub fn has_pending(&self) -> bool {
1563        self.pending_tok.is_some()
1564    }
1565    /// Committed row count == cache rows (the session invariant), for the caller's own
1566    /// `fed`-length cross-check at a handoff boundary.
1567    pub fn committed_len(&self) -> usize {
1568        self.committed.len()
1569    }
1570    /// DEMOTION HANDOFF (lane/spec-gate, 2026-08-07): consume this session and hand its trunk
1571    /// cache + next-token prediction to the plain batched-decode path.
1572    ///
1573    /// WHY THIS IS EXACT (greedy). The invariant at a burst boundary is `cache.pos ==
1574    /// committed.len()`: every committed row has trunk KV + recurrent state, exactly as a plain
1575    /// tokenwise prime of the same `committed` sequence would have left it (that is the
1576    /// session-tail contract, and the same property `spec_rewind_to_checkpoint` and the reuse
1577    /// pool already rely on). `next_pred` is the argmax of the verify's logits for the LAST
1578    /// committed row — and verify-column logits are bit-identical to plain decode's logits at
1579    /// that position, because `matmul_decode_exact` bit-identity IS the basis of the greedy
1580    /// accept walk. So handing (cache, next_pred) to the batched path continues the stream from
1581    /// a state indistinguishable from one the batched path produced itself: the batched tick
1582    /// emits `next_pred`, feeds it into this same cache, and decodes on.
1583    ///
1584    /// `None` when the session is not in the handoff shape — a carried pending (its bonus row is
1585    /// NOT in the cache, so `spec_flush_pending` must commit it first) or no `next_pred` yet
1586    /// (never bursted). Callers must not force it: a half-committed cache handed to the batched
1587    /// path would silently skip a token.
1588    ///
1589    /// The MTP draft scratch, the persistent draft-graph context and the turn checkpoint are
1590    /// DROPPED here (freeing their VRAM): the batched path never drafts, and this handoff is
1591    /// one-way by design — there is no cheap symmetric re-promotion (rebuilding the draft KV
1592    /// would mean an `mtp_kv_fill` over the whole committed history).
1593    pub fn into_demoted(self) -> Option<(Cache, u32)> {
1594        if self.pending_tok.is_some() || self.cache.tainted {
1595            return None;
1596        }
1597        let np = self.next_pred?;
1598        debug_assert_eq!(
1599            self.cache.pos,
1600            self.committed.len(),
1601            "demotion handoff: cache rows != committed tokens"
1602        );
1603        Some((self.cache, np))
1604    }
1605    /// Pool-resume hook (audit Q2): clear the parked draft-graph failure memoization so a
1606    /// NEW request resuming this session gets one fresh capture chance — a transient-pressure
1607    /// capture failure must not persist for the pool's whole lifetime (the TRT #16072 class).
1608    /// Logs once iff a flag was actually set; a no-fallback resume is silent and free.
1609    pub fn reset_graph_fallback_on_resume(&mut self) {
1610        if let Some(line) = self
1611            .draft_ctx
1612            .as_mut()
1613            .and_then(|c| c.failed.reset_on_resume())
1614        {
1615            eprintln!("{line}");
1616        }
1617    }
1618}
1619
1620/// A session's PROMPT-END boundary state, the rewind target for session-affinity resume.
1621///
1622/// WHY THIS BOUNDARY, AND WHY IT IS THE ONLY ONE WORTH KEEPING. The rewrite class this lane
1623/// exists for (a client that strips `<think>` blocks out of prior assistant turns) mutates the
1624/// text the session GENERATED, never the prompt it was given. So turn N's prompt agrees with
1625/// turn N-1's committed tokens up to almost exactly where turn N-1's generation began — the
1626/// prompt-end boundary. Keeping a checkpoint there means the next turn re-primes only its own
1627/// delta (the rewritten answer + the new user turn) instead of the whole conversation.
1628///
1629/// WHAT IT MUST HOLD. Full-attn KV is append-only and position-addressed, so rewinding it is a
1630/// `len` truncation (no data). Linear-attn (GDN) conv/ssm state is mutated IN PLACE with no
1631/// position index, so it must be a real device COPY — that copy is the entire reason a spec
1632/// session could not previously rewind. The MTP draft scratch needs no copy either: its rows
1633/// below the boundary were written by this turn's fill and are never revisited (the per-round
1634/// true-hidden refresh only rewrites the CURRENT burst's committed positions), so rewinding it
1635/// is also just a `len` reset. `last_h` is the hidden of the last row below the boundary — the
1636/// predecessor-pairing anchor the next prime's fill reads for its first row.
1637///
1638/// COST: one `Cache::snapshot` per TURN, on a code path that already takes one per ROUND.
1639pub(crate) struct SpecCheckpoint {
1640    snap: crate::cache::CacheSnapshot,
1641    /// Committed length at the boundary (== cache.pos there, the session invariant).
1642    pos: usize,
1643    /// Pre-output_norm hidden of row `pos - 1`.
1644    last_h: CudaSlice<f32>,
1645}
1646
1647/// PREFIX-CACHE BOUNDARY CAPTURE (lane/spec-prefix-cache, 2026-08-14): the state a spec session
1648/// records at its cold-prime split so the WORKER can publish a cross-request prefix entry —
1649/// the commit-gated-publication port (research/cache-spec-design-20260814/PORT-PLAN.md item 1).
1650/// Only the pieces that are DESTROYED by continuing the prime need copies here: the in-place
1651/// GDN conv/ssm states (via `Cache::snapshot`, same mechanism as [`SpecCheckpoint`]) and the
1652/// boundary logits. Full-attn KV rows `[0..pos)` and draft-scratch rows `[0..pos)` are
1653/// append-only for the session's lifetime (rollbacks never truncate below the prime boundary),
1654/// so the worker slices those from the live caches post-burst instead of copying at prime time.
1655pub struct SpecBoundaryCapture {
1656    pub snap: crate::cache::CacheSnapshot,
1657    /// Token boundary (== cache.pos at capture; == the worker's miss-LCP split).
1658    pub pos: usize,
1659    /// Full-vocab logits after the prefix prime — the entry's boundary logits.
1660    pub logits: Vec<f32>,
1661    /// Pre-output_norm trunk hidden of row `pos - 1` (lane/spec-on-cache-hit): the
1662    /// predecessor-pairing anchor a RESTORED spec session's first suffix-fill row reads
1663    /// (the `SpecSession::last_h` convention). Empty = unavailable (capture stays valid;
1664    /// the fill's zeros row-0 fallback covers it at a bounded acceptance cost).
1665    pub last_h: Vec<f32>,
1666    /// Per-layer latent boundary tails (lane/glm5-prefix-latent2, 2026-09-01): the
1667    /// generation-destroyed slice of each MLA/DSA layer's boundary state, captured eagerly
1668    /// so the worker's DEFERRED publication can slice the append-only planes from the live
1669    /// cache (`LatentKvLayer::snapshot_plane_at`). EMPTY on every two-plane model — the
1670    /// pre-field captures are byte-identical; a latent-bearing cache with an EMPTY vec here
1671    /// keeps the publisher's loud refusal (the fail-closed door stays shut).
1672    pub latent_tails: Vec<Option<crate::cache::LatentTailCapture>>,
1673}
1674
1675/// D2H one hidden row out of a `[T, n_embd]` prime hidden stack — the boundary anchor a
1676/// spec boundary capture carries for later restored-session fills. Failure is silent
1677/// (`turn_ckpt` convention): the capture publishes without an anchor.
1678pub(crate) fn capture_boundary_hidden(
1679    e: &Engine,
1680    h_rows: &CudaSlice<f32>,
1681    pos: usize,
1682    n_embd: usize,
1683) -> Vec<f32> {
1684    if pos == 0 || h_rows.len() < pos * n_embd {
1685        return Vec::new();
1686    }
1687    let Ok(mut row) = e.uninit(n_embd) else {
1688        return Vec::new();
1689    };
1690    if e.copy_view_into(
1691        &mut row,
1692        0,
1693        &h_rows.slice((pos - 1) * n_embd..pos * n_embd),
1694        n_embd,
1695    )
1696    .is_err()
1697    {
1698        return Vec::new();
1699    }
1700    e.dtoh(&row).unwrap_or_default()
1701}
1702
1703/// ROLLBACK DOOR for sampled BOUNDARY tokens (lane/sampled-spec-quality, 2026-08-19).
1704/// Default ON: the token a burst emits at its own boundary is drawn from the request's
1705/// sampler. `MEMRA_SPEC_SAMPLED_BOUNDARY=0` restores the pre-lane posture (an ARGMAX at
1706/// every boundary) without touching greedy, which is byte-unaffected either way.
1707pub fn spec_sampled_boundary_on() -> bool {
1708    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1709    *ON.get_or_init(|| std::env::var("MEMRA_SPEC_SAMPLED_BOUNDARY").as_deref() != Ok("0"))
1710}
1711
1712/// ROLLBACK DOOR for SESSION-SPANNING penalty history (lane/sampled-spec-quality).
1713/// Default ON: `pen_hist` is seeded from the session's committed tail, so repetition /
1714/// frequency / presence penalties see the whole stream. `MEMRA_SPEC_PEN_SESSION=0`
1715/// restores the pre-lane posture (each burst restarts the window from its own prompt
1716/// slice, i.e. from NOTHING on a continuation burst) — and with the door shut the worker
1717/// must keep refusing penalized sampled prefix-cache restores, because the restored
1718/// session's continuation burst is handed no prompt slice at all.
1719pub fn spec_pen_session_on() -> bool {
1720    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1721    *ON.get_or_init(|| std::env::var("MEMRA_SPEC_PEN_SESSION").as_deref() != Ok("0"))
1722}
1723
1724/// ROLLBACK DOOR for extended-entry publication from a RESTORED session
1725/// (lane/sampled-spec-quality, Item 3). Default ON: a converted prefix-cache hit that fed a
1726/// suffix captures its own prompt-end boundary so the NEXT turn can hit a longer prefix.
1727/// `MEMRA_SPEC_RESTORE_REPUBLISH=0` restores the pre-lane posture (a namespace learns exactly
1728/// one boundary and never advances it). Whole-entry semantics only — the boundary is the
1729/// restored session's own prompt end, so `entry_pos != fed_len` still refuses on the way in.
1730pub fn spec_restore_republish_on() -> bool {
1731    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1732    *ON.get_or_init(|| std::env::var("MEMRA_SPEC_RESTORE_REPUBLISH").as_deref() != Ok("0"))
1733}
1734
1735/// Diagnostics: name every boundary token on stderr (`MEMRA_SPEC_BOUNDARY_TRACE=1`), with
1736/// the argmax the pre-lane code would have emitted from the same row. This is how the
1737/// lane MEASURES the boundary rate and the deviation rate instead of estimating them.
1738fn spec_boundary_trace() -> bool {
1739    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1740    *ON.get_or_init(|| std::env::var("MEMRA_SPEC_BOUNDARY_TRACE").as_deref() == Ok("1"))
1741}
1742
1743/// llama-parity floor for the penalty window when the request does not ask for a bigger
1744/// one (`repeat_last_n` default). The serve API arms `penalty_last_n = PEN_WINDOW_MAX` for any
1745/// non-identity penalty, so this floor only matters to explicit small windows and to the
1746/// CLI env path.
1747const PEN_WINDOW_FLOOR: usize = 64;
1748
1749/// CEILING on the penalty window, and it is a COST bound, not a semantic preference.
1750/// `penalize_logits_f32` (cu/spec_sample.cu) dedups on device by having thread `i` scan
1751/// `hist[0..i]`, so a pass is O(n_hist²) and it runs ~3x per verify round (the q rows, the
1752/// p column, the bonus column). The serve API uses this same bound for every non-identity
1753/// penalty so host/plain, sparse-device, and speculative sampling cannot change logits on
1754/// admission demotion. An uncapped 128k-token history would put ~1.7e10
1755/// comparisons per pass, tens of ms per round, i.e. penalties would silently destroy decode
1756/// throughput on exactly the long-context requests that most want them. 8192 keeps a pass
1757/// at ~7e7 comparisons (tens of microseconds) while still being **128x wider than the
1758/// pre-lane effective window** (64 prompt-tail tokens + whatever the current burst had
1759/// generated). A request that genuinely needs a window beyond this wants host-side dedup +
1760/// counts through a new kernel signature — a follow-up lane, named here rather than hidden.
1761/// `pub` since lane/dspark-penalized-sampled-20260821: the dspark route's accept walk and
1762/// the dspark_sample_gate binary trim their uploads with the SAME cap — a second constant
1763/// is a second thing to drift.
1764pub const PEN_WINDOW_MAX: usize = 8192;
1765
1766/// Seed a penalty window over the SESSION, not the burst (lane/sampled-spec-quality,
1767/// Item 2). The window is the last `max(penalty_last_n, 64)` tokens of
1768/// `session_committed ++ burst_prompt` — for a cold turn-1 burst (`session_committed`
1769/// empty, default `penalty_last_n`) that is byte-identically the pre-lane
1770/// `prompt.iter().rev().take(64).rev()`; for a continuation burst it is the stream the
1771/// client actually asked us to penalize, where the pre-lane code had NOTHING.
1772/// `pub` since lane/dspark-penalized-sampled-20260821: the dspark route seeds its session
1773/// window through the SAME function (one definition of "the window" across both spec
1774/// routes and the gate binary's trunk-only reference arm).
1775pub fn pen_window_seed(
1776    session_committed: &[u32],
1777    burst_prompt: &[u32],
1778    penalty_last_n: usize,
1779) -> Vec<u32> {
1780    let win = penalty_last_n.clamp(PEN_WINDOW_FLOOR, PEN_WINDOW_MAX);
1781    let take_prompt = burst_prompt.len().min(win);
1782    let take_sess = (win - take_prompt).min(session_committed.len());
1783    let mut hist = Vec::with_capacity(take_sess + take_prompt);
1784    hist.extend_from_slice(&session_committed[session_committed.len() - take_sess..]);
1785    hist.extend_from_slice(&burst_prompt[burst_prompt.len() - take_prompt..]);
1786    hist
1787}
1788
1789/// Draw a BOUNDARY token from the target distribution the request asked for
1790/// (lane/sampled-spec-quality, Item 1) — the fix for "sampled spec emits an ARGMAX token at
1791/// every burst boundary".
1792///
1793/// WHY THIS EXISTS. A spec burst's first emitted token is not produced by the accept walk:
1794/// it comes off a logits row that already exists (the prime's last row on a cold burst; the
1795/// row after the last committed token on a continuation burst; the prefix-cache entry's
1796/// boundary row on a restored one). Pre-lane that token was `argmax` in BOTH sampling
1797/// regimes, so a sampled stream took a greedy token once per burst — measured, not
1798/// estimated, in research/spec-cache-20260818/SAMPLED-QUALITY.md. At temperature > 0 the
1799/// customer asked for a sampled token, so this draws one.
1800///
1801/// THE PROGRAM IS THE FULL-ACCEPT BONUS'S PROGRAM, deliberately: penalize the row (over the
1802/// session's window), take this row's OWN filter stats (the sampfix-20260805 law — stats
1803/// from a neighbour row mis-scale every `e0` and can wipe the row to token 0), gumbel-perturb
1804/// with the session's Philox stream at `*sctr`, argmax the perturbed row. Reusing the bonus's
1805/// composition means `sample_check`'s distributional oracle covers this draw too, and the
1806/// boundary token is drawn from the same filtered/penalized `p` the accept walk targets.
1807///
1808/// THE STREAM IS THE SESSION'S, NOT A FRESH ONE. `sctr` is the caller's live counter and is
1809/// advanced by exactly one, so a boundary draw consumes the next value in the same Philox
1810/// stream the accept walk uses — never a second, independently seeded stream (which would be
1811/// a new distributional bug: two streams from one seed correlate wherever their counters
1812/// collide). That also makes a restored session's boundary draw at `sctr == 0` bit-identical
1813/// to the cold session's own first draw from the same logits row, which is what preserves the
1814/// sampled-hit lane's per-seed hit==cold byte identity.
1815#[allow(clippy::too_many_arguments)]
1816pub fn sample_boundary_token_dev(
1817    e: &Engine,
1818    logits: &CudaSlice<f32>,
1819    n_vocab: usize,
1820    sp: &SpecSampling,
1821    pen_hist: &[u32],
1822    sctr: &mut u32,
1823    site: &str,
1824) -> Result<u32, Box<dyn std::error::Error>> {
1825    debug_assert!(
1826        sp.temp > 0.0,
1827        "boundary sampling is the sampled regime only"
1828    );
1829    // Own copy: penalize_logits mutates in place and the caller's row is live state
1830    // (prime_logits back the constrained recompute; last_col_logits backs round 0's accept).
1831    let mut col = e.zeros(n_vocab)?;
1832    e.copy_into(&mut col, 0, logits, n_vocab)?;
1833    let pen_on = sp.penalty_last_n > 0
1834        && (sp.penalty_repeat != 1.0 || sp.penalty_freq != 0.0 || sp.penalty_present != 0.0);
1835    if pen_on && !pen_hist.is_empty() {
1836        // window trim mirrors the round loop's own upload (`pen_hist[w0..]`), cap included.
1837        let w0 = pen_hist
1838            .len()
1839            .saturating_sub(sp.penalty_last_n.min(PEN_WINDOW_MAX));
1840        let hist = &pen_hist[w0..];
1841        let hd = e.htod_u32_v(hist)?;
1842        e.penalize_logits(
1843            &mut col,
1844            &hd,
1845            hist.len(),
1846            sp.penalty_repeat,
1847            sp.penalty_freq,
1848            sp.penalty_present,
1849            n_vocab,
1850        )?;
1851    }
1852    let rows0 = e.htod_i32(&[0])?;
1853    let (mut th_d, mut z_d, mut mx_d) = (e.zeros(1)?, e.zeros(1)?, e.zeros(1)?);
1854    e.filter_stats(
1855        &col, n_vocab, &rows0, &mut th_d, &mut z_d, &mut mx_d, n_vocab, 1, sp.temp, sp.top_k,
1856        sp.top_p, sp.min_p,
1857    )?;
1858    let (th, mx) = (e.dtoh(&th_d)?[0], e.dtoh(&mx_d)?[0]);
1859    let mut perturb = e.zeros(n_vocab)?;
1860    e.gumbel_perturb_filtered(&col, &mut perturb, n_vocab, sp.seed, *sctr, sp.temp, mx, th)?;
1861    *sctr = sctr.wrapping_add(1);
1862    let td = e.argmax_token_device(&perturb, n_vocab)?;
1863    let tok = guard_vocab_token(
1864        e.dtoh_u32_one(&td)?,
1865        n_vocab,
1866        &format!("sampled boundary token (site={site})"),
1867    )?;
1868    if spec_boundary_trace() {
1869        // the pre-lane token, from the SAME row, so the deviation rate is measurable.
1870        let raw = e.argmax_token_device(logits, n_vocab)?;
1871        let greedy = e.dtoh_u32_one(&raw)?;
1872        eprintln!(
1873            "[spec-boundary] site={site} sampled={tok} argmax={greedy} \
1874             deviates={} temp={} sctr={}",
1875            (tok != greedy) as u8,
1876            sp.temp,
1877            sctr.wrapping_sub(1),
1878        );
1879    }
1880    Ok(tok)
1881}
1882
1883/// Host-row twin of [`sample_boundary_token_dev`] (the prime / feed / entry rows arrive as
1884/// host `Vec<f32>`).
1885#[allow(clippy::too_many_arguments)]
1886pub fn sample_boundary_token(
1887    e: &Engine,
1888    logits: &[f32],
1889    sp: &SpecSampling,
1890    pen_hist: &[u32],
1891    sctr: &mut u32,
1892    site: &str,
1893) -> Result<u32, Box<dyn std::error::Error>> {
1894    let n_vocab = logits.len();
1895    let d = e.htod(logits)?;
1896    sample_boundary_token_dev(e, &d, n_vocab, sp, pen_hist, sctr, site)
1897}
1898
1899struct SpecPipeTraceClock {
1900    pair: usize,
1901    started: std::time::Instant,
1902}
1903
1904#[derive(Clone)]
1905struct SpecPipeTraceCtx {
1906    clock: std::sync::Arc<SpecPipeTraceClock>,
1907    round: usize,
1908    lane: usize,
1909}
1910
1911struct SpecPipeTraceMarker {
1912    trace: SpecPipeTraceCtx,
1913    phase: &'static str,
1914    edge: &'static str,
1915    slot: Option<usize>,
1916}
1917
1918unsafe extern "C" fn spec_pipe_trace_marker(raw: *mut std::ffi::c_void) {
1919    let marker = unsafe { Box::from_raw(raw.cast::<SpecPipeTraceMarker>()) };
1920    let lane = if marker.trace.lane == 0 { "A" } else { "B" };
1921    let slot = marker
1922        .slot
1923        .map(|v| v.to_string())
1924        .unwrap_or_else(|| "-".into());
1925    let t_ms = marker.trace.clock.started.elapsed().as_secs_f64() * 1e3;
1926    use std::io::Write as _;
1927    let stderr = std::io::stderr();
1928    let mut stderr = stderr.lock();
1929    let _ = writeln!(
1930        stderr,
1931        "[spec-pipe-timeline] pair={} round={} lane={lane} phase={} edge={} \
1932         slot={slot} t_ms={t_ms:.3}",
1933        marker.trace.clock.pair, marker.trace.round, marker.phase, marker.edge,
1934    );
1935}
1936
1937fn enqueue_spec_pipe_trace_marker(
1938    stream: &cudarc::driver::CudaStream,
1939    trace: Option<&SpecPipeTraceCtx>,
1940    phase: &'static str,
1941    edge: &'static str,
1942    slot: Option<usize>,
1943) -> Result<(), Box<dyn std::error::Error>> {
1944    let Some(trace) = trace else {
1945        return Ok(());
1946    };
1947    let marker = Box::new(SpecPipeTraceMarker {
1948        trace: trace.clone(),
1949        phase,
1950        edge,
1951        slot,
1952    });
1953    let raw = Box::into_raw(marker);
1954    let result = unsafe {
1955        cudarc::driver::result::stream::launch_host_function(
1956            stream.cu_stream(),
1957            spec_pipe_trace_marker,
1958            raw.cast(),
1959        )
1960    };
1961    if let Err(err) = result {
1962        unsafe {
1963            drop(Box::from_raw(raw));
1964        }
1965        return Err(err.into());
1966    }
1967    Ok(())
1968}
1969
1970#[derive(Default)]
1971struct SpecPipeProgress {
1972    setup_done: [bool; 2],
1973    draft_done: [usize; 2],
1974    stage0_done: [usize; 2],
1975    verify_done: [usize; 2],
1976    accept_done: [usize; 2],
1977    finished: [bool; 2],
1978    aborted: bool,
1979}
1980
1981/// Host-side issue coordinator for the reduced two-session speculative pipeline. Each session
1982/// keeps its existing call stack and round locals; this object only orders phase entry. The
1983/// primary mutex spans whole draft/accept/tail issue regions so Engine's single-stream scratch
1984/// cannot be interleaved by the two host threads.
1985struct SpecPipeSync {
1986    progress: std::sync::Mutex<SpecPipeProgress>,
1987    changed: std::sync::Condvar,
1988    primary: std::sync::Mutex<()>,
1989    trace: Option<std::sync::Arc<SpecPipeTraceClock>>,
1990}
1991
1992impl SpecPipeSync {
1993    fn new() -> Self {
1994        static TRACE_PAIR: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
1995        let trace = (std::env::var("MEMRA_SPEC_PIPE_TRACE").as_deref() == Ok("1")).then(|| {
1996            std::sync::Arc::new(SpecPipeTraceClock {
1997                pair: TRACE_PAIR.fetch_add(1, std::sync::atomic::Ordering::Relaxed) + 1,
1998                started: std::time::Instant::now(),
1999            })
2000        });
2001        Self {
2002            progress: std::sync::Mutex::new(SpecPipeProgress::default()),
2003            changed: std::sync::Condvar::new(),
2004            primary: std::sync::Mutex::new(()),
2005            trace,
2006        }
2007    }
2008}
2009
2010#[derive(Clone)]
2011struct SpecPipeLane {
2012    sync: std::sync::Arc<SpecPipeSync>,
2013    lane: usize,
2014    rt: &'static crate::pp::PpNRt,
2015    walk_permit: crate::pp::PpWalkPermit,
2016}
2017
2018struct SpecPipePrimaryGuard<'a> {
2019    _primary: std::sync::MutexGuard<'a, ()>,
2020    _walk: crate::pp::PpWalkBorrowGuard,
2021}
2022
2023impl SpecPipeLane {
2024    fn peer(&self) -> usize {
2025        1 - self.lane
2026    }
2027
2028    fn aborted() -> Box<dyn std::error::Error> {
2029        "paired speculative peer aborted".into()
2030    }
2031
2032    fn trace(&self, round: usize) -> Option<SpecPipeTraceCtx> {
2033        self.sync.trace.as_ref().map(|clock| SpecPipeTraceCtx {
2034            clock: clock.clone(),
2035            round,
2036            lane: self.lane,
2037        })
2038    }
2039
2040    fn setup_begin(&self) -> Result<crate::pp::PpWalkBorrowGuard, Box<dyn std::error::Error>> {
2041        let mut p = self.sync.progress.lock().unwrap();
2042        while !p.aborted && self.lane == 1 && !p.setup_done[0] && !p.finished[0] {
2043            p = self.sync.changed.wait(p).unwrap();
2044        }
2045        if p.aborted {
2046            Err(Self::aborted())
2047        } else {
2048            drop(p);
2049            self.rt.borrow_walk(&self.walk_permit, "spec_pipe/setup")
2050        }
2051    }
2052
2053    fn setup_end(&self) {
2054        let mut p = self.sync.progress.lock().unwrap();
2055        p.setup_done[self.lane] = true;
2056        self.sync.changed.notify_all();
2057    }
2058
2059    fn draft_begin(
2060        &self,
2061        round: usize,
2062    ) -> Result<SpecPipePrimaryGuard<'_>, Box<dyn std::error::Error>> {
2063        let peer = self.peer();
2064        let mut p = self.sync.progress.lock().unwrap();
2065        loop {
2066            if p.aborted {
2067                return Err(Self::aborted());
2068            }
2069            let setup_ready =
2070                (p.setup_done[0] || p.finished[0]) && (p.setup_done[1] || p.finished[1]);
2071            let prior_ready = p.accept_done[self.lane] >= round
2072                && (p.accept_done[peer] >= round || p.finished[peer]);
2073            let turn_ready = if self.lane == 0 {
2074                true
2075            } else {
2076                p.draft_done[0] > round || p.finished[0]
2077            };
2078            if setup_ready && prior_ready && turn_ready {
2079                break;
2080            }
2081            p = self.sync.changed.wait(p).unwrap();
2082        }
2083        drop(p);
2084        let primary = self.sync.primary.lock().unwrap();
2085        let walk = self.rt.borrow_walk(&self.walk_permit, "spec_pipe/draft")?;
2086        Ok(SpecPipePrimaryGuard {
2087            _primary: primary,
2088            _walk: walk,
2089        })
2090    }
2091
2092    fn draft_end(&self, round: usize) {
2093        let mut p = self.sync.progress.lock().unwrap();
2094        p.draft_done[self.lane] = round + 1;
2095        self.sync.changed.notify_all();
2096    }
2097
2098    /// Admit stage 0 and return whether this lane owns the interval's one reverse fence.
2099    /// Lane B releases as soon as lane A has issued its boundary TX, not after A's full body.
2100    fn stage0_begin(&self, round: usize) -> Result<bool, Box<dyn std::error::Error>> {
2101        let peer = self.peer();
2102        let mut p = self.sync.progress.lock().unwrap();
2103        loop {
2104            if p.aborted {
2105                return Err(Self::aborted());
2106            }
2107            let ready = if self.lane == 0 {
2108                p.draft_done[0] > round && (p.draft_done[1] > round || p.finished[1])
2109            } else {
2110                p.draft_done[1] > round && (p.stage0_done[0] > round || p.finished[0])
2111            };
2112            if ready {
2113                return Ok(self.lane == 0 || p.finished[peer]);
2114            }
2115            p = self.sync.changed.wait(p).unwrap();
2116        }
2117    }
2118
2119    fn stage0_end(&self, round: usize) {
2120        let mut p = self.sync.progress.lock().unwrap();
2121        p.stage0_done[self.lane] = round + 1;
2122        self.sync.changed.notify_all();
2123    }
2124
2125    /// Stage 1 is single-owner per engine. A proceeds immediately after its own ticket; B waits
2126    /// for A's full stage1/head issue so only A.S1 and B.S0 can overlap.
2127    fn stage1_begin(&self, round: usize) -> Result<(), Box<dyn std::error::Error>> {
2128        let mut p = self.sync.progress.lock().unwrap();
2129        while !p.aborted
2130            && !(p.stage0_done[self.lane] > round
2131                && (self.lane == 0 || p.verify_done[0] > round || p.finished[0]))
2132        {
2133            p = self.sync.changed.wait(p).unwrap();
2134        }
2135        if p.aborted {
2136            Err(Self::aborted())
2137        } else {
2138            Ok(())
2139        }
2140    }
2141
2142    fn verify_end(&self, round: usize) {
2143        let mut p = self.sync.progress.lock().unwrap();
2144        p.verify_done[self.lane] = round + 1;
2145        self.sync.changed.notify_all();
2146    }
2147
2148    fn accept_begin(
2149        &self,
2150        round: usize,
2151    ) -> Result<SpecPipePrimaryGuard<'_>, Box<dyn std::error::Error>> {
2152        let mut p = self.sync.progress.lock().unwrap();
2153        loop {
2154            if p.aborted {
2155                return Err(Self::aborted());
2156            }
2157            let ready = if self.lane == 0 {
2158                p.verify_done[0] > round && (p.verify_done[1] > round || p.finished[1])
2159            } else {
2160                p.verify_done[1] > round && (p.accept_done[0] > round || p.finished[0])
2161            };
2162            if ready {
2163                break;
2164            }
2165            p = self.sync.changed.wait(p).unwrap();
2166        }
2167        drop(p);
2168        let primary = self.sync.primary.lock().unwrap();
2169        let walk = self.rt.borrow_walk(&self.walk_permit, "spec_pipe/accept")?;
2170        Ok(SpecPipePrimaryGuard {
2171            _primary: primary,
2172            _walk: walk,
2173        })
2174    }
2175
2176    fn accept_end(&self, round: usize) {
2177        let mut p = self.sync.progress.lock().unwrap();
2178        p.accept_done[self.lane] = round + 1;
2179        self.sync.changed.notify_all();
2180    }
2181
2182    fn primary(&self) -> Result<SpecPipePrimaryGuard<'_>, Box<dyn std::error::Error>> {
2183        let primary = self.sync.primary.lock().unwrap();
2184        let walk = self.rt.borrow_walk(&self.walk_permit, "spec_pipe/tail")?;
2185        Ok(SpecPipePrimaryGuard {
2186            _primary: primary,
2187            _walk: walk,
2188        })
2189    }
2190
2191    fn coordinated_walk(&self) -> Result<crate::pp::PpWalkBorrowGuard, Box<dyn std::error::Error>> {
2192        self.rt
2193            .borrow_walk(&self.walk_permit, "spec_pipe/coordinated_verify")
2194    }
2195
2196    fn finish(&self, failed: bool) {
2197        let mut p = self.sync.progress.lock().unwrap();
2198        p.finished[self.lane] = true;
2199        p.aborted |= failed;
2200        self.sync.changed.notify_all();
2201    }
2202}
2203
2204struct SpecPipeFinish<'a> {
2205    lane: &'a SpecPipeLane,
2206    closed: bool,
2207}
2208
2209impl<'a> SpecPipeFinish<'a> {
2210    fn new(lane: &'a SpecPipeLane) -> Self {
2211        Self {
2212            lane,
2213            closed: false,
2214        }
2215    }
2216
2217    fn close(&mut self, failed: bool) {
2218        self.lane.finish(failed);
2219        self.closed = true;
2220    }
2221}
2222
2223impl Drop for SpecPipeFinish<'_> {
2224    fn drop(&mut self) {
2225        if !self.closed {
2226            self.lane.finish(true);
2227        }
2228    }
2229}
2230
2231/// Scoped transfer of one exclusively-borrowed session to the second host issue thread.
2232/// `CudaGraph` is not marked Send by cudarc because its raw driver handles carry no automatic
2233/// trait. CUDA driver graph handles are context-scoped rather than OS-thread-affine; the caller
2234/// binds that context before touching the session, joins before returning, and never aliases the
2235/// pointer. Keep this exception local to the experimental pair call instead of marking the public
2236/// session type Send.
2237struct SpecPipeSessionPtr(*mut SpecSession);
2238
2239unsafe impl Send for SpecPipeSessionPtr {}
2240
2241impl SpecPipeSessionPtr {
2242    unsafe fn get_mut(&mut self) -> &mut SpecSession {
2243        unsafe { &mut *self.0 }
2244    }
2245}
2246
2247/// Per-session persistent draft-graph context: the captured CUDA graph(s) plus the device
2248/// buffers whose POINTERS the capture bakes. Reuse legality: the greedy capture bakes only
2249/// session-stable pointers (the session's own MtpScratch KV — allocated once, never realloc'd;
2250/// the model's resident embedding; the process-wide OnceLock p_min) and the g_* buffers held
2251/// HERE — so one capture serves the session's whole lifetime. The sampled capture additionally
2252/// bakes (seed, temp) as capture-time constants and needs k q-slots — keyed by `s_key`, dropped
2253/// and recaptured when a pool-resumed request changes them. `*_failed` memoizes a failed capture
2254/// so the eager fallback doesn't pay a doomed capture attempt every burst.
2255/// Capture identity of the parked SAMPLED draft graph (`DraftGraphCtx::graph_s`).
2256///
2257/// EXACTNESS, not perf (lane/graph-s-key-exactness-20260819; receipts
2258/// `research/spec-cache-20260818/GRAPH-S-KEY.md`). Two classes of field live here, both
2259/// load-bearing:
2260///
2261/// - **Baked constants.** `seed` and `temp` are capture-time constants INSIDE the graph and `k`
2262///   sizes the q slots its replays write. A resumed request changing any of them must recapture.
2263///   This is all the key used to carry.
2264/// - **Regime fields.** `top_k`/`top_p`/`min_p`/`pen_on` are not baked, but they decide whether
2265///   the captured graph is a legal draft chain AT ALL. The in-graph draw is one gumbel-max over
2266///   the RAW softmax (`gumbel_perturb_ctr`, unfiltered by construction), while the verify builds
2267///   the accept test's `q` from `filter_stats(q_slots, top_k, top_p, min_p)`. If those disagree
2268///   the accept test evaluates a distribution the draft was never sampled from: a draft token
2269///   below the filter threshold gathers `q = 0` (`softmax_gather_filtered_f32`,
2270///   `cu/spec_sample.cu`) and `u * 0 < p` accepts it UNCONDITIONALLY.
2271///
2272/// Omitting the regime fields was reachable — not through the prefix-cache spec restore (that
2273/// path is greedy-only, `memra-server` `spec_restore_convertible`), but through WHOLE-SESSION
2274/// spec reuse: a parked `SpecSession` carries this `DraftGraphCtx`, and the pool-resume probe
2275/// applies no sampler predicate at all. Turn 1 pure-temp parks a graph; turn 2 of the same
2276/// conversation, same explicit seed and temperature, adds `top_p`/`top_k` and inherits it.
2277#[derive(Clone, Copy, PartialEq, Eq, Debug)]
2278pub(crate) struct SampledGraphKey {
2279    seed: u64,
2280    temp_bits: u32,
2281    k: usize,
2282    top_k: i32,
2283    top_p_bits: u32,
2284    min_p_bits: u32,
2285    pen_on: bool,
2286}
2287
2288impl SampledGraphKey {
2289    pub(crate) fn new(
2290        seed: u64,
2291        temp: f32,
2292        k: usize,
2293        top_k: i32,
2294        top_p: f32,
2295        min_p: f32,
2296        pen_on: bool,
2297    ) -> Self {
2298        SampledGraphKey {
2299            seed,
2300            temp_bits: temp.to_bits(),
2301            k,
2302            top_k,
2303            top_p_bits: top_p.to_bits(),
2304            min_p_bits: min_p.to_bits(),
2305            pen_on,
2306        }
2307    }
2308
2309    /// The one regime the PURE-TEMP in-graph sampled chain may stand in for the eager one:
2310    /// nothing but temperature shapes `q`. Computed FROM THE KEY so the capture guard, the
2311    /// launch guard and the key can never drift apart (they were three separate expressions
2312    /// before this lane, and the launch site simply forgot to ask).
2313    pub(crate) fn pure_temp(&self) -> bool {
2314        self.top_k == 0
2315            && f32::from_bits(self.top_p_bits) >= 1.0
2316            && f32::from_bits(self.min_p_bits) <= 0.0
2317            && !self.pen_on
2318    }
2319
2320    /// Truncation filters active — the capture body needs the IN-GRAPH filter nodes
2321    /// (`filter_stats` + `gumbel_perturb_filtered_ctr`) so the draft draws from the same
2322    /// filtered distribution the accept test reconstructs. Meaningful only when
2323    /// `graph_capturable`; penalties never reach a capture body.
2324    pub(crate) fn filtered(&self) -> bool {
2325        !self.pure_temp()
2326    }
2327
2328    /// May the sampled draft graph be CAPTURED (and a parked one LAUNCHED) for this regime?
2329    /// Pure-temp always; filtered regimes when the filtered-capture door is on
2330    /// (lane/step37-draft-graph-serving-20260830); penalties never — the per-round history
2331    /// cannot be baked into a graph, and composing a raw-softmax (or stale-history) draw
2332    /// with a penalized accept test is the unconditional-accept exactness bug. Computed FROM
2333    /// THE KEY for the same no-drift reason as `pure_temp`.
2334    pub(crate) fn graph_capturable(&self) -> bool {
2335        !self.pen_on && (self.pure_temp() || spec_graph_filtered_on())
2336    }
2337}
2338
2339/// Per-head captured graphs for the MULTI-HEAD MTP draft chain (step-modulo prefix-replay,
2340/// lane/step37-draft-graph-serving-20260830). The chain POLICY — which head serves step j,
2341/// how long the replayed prefix is, which stored seed feeds row r — stays HOST-SIDE in the
2342/// launch loop, exactly `mtp_chain_forward_dev`'s order; the graphs capture ONE head-row
2343/// forward each, on the head's OWN scratch plane:
2344/// - `interior[i]`: head i, `with_head=false` — KV append + carrier only. Interior rows'
2345///   logits are dead in the eager chain too (`mtp_chain_forward_dev` keeps only the last
2346///   row), so skipping the head matmul changes no consumed byte and removes the eager
2347///   chain's per-replay-row full-vocab matmul.
2348/// - `last[i]`: head i, `with_head=true` + the mode's tail (greedy argmax, or the sampled
2349///   gumbel draw — filtered in-graph when the request carries filters).
2350///
2351/// One `DraftChainGraphs` per MODE (greedy vs sampled), owning its keeper: dropping the
2352/// sampled chain on an s_key change never invalidates the greedy one.
2353struct DraftChainGraphs {
2354    interior: Vec<cudarc::driver::CudaGraph>,
2355    last: Vec<cudarc::driver::CudaGraph>,
2356    /// Never read: exists to OWN the captured graphs' backing buffers for as long as the
2357    /// graphs replay (the capture-retain law; same class as `DsparkSegGraph::_keeper`).
2358    _keeper: Vec<Box<dyn std::any::Any + Send>>,
2359}
2360
2361/// Sampled-tail capture pack for `mtp_head_forward_cap`: the persistent buffers and baked
2362/// constants of the in-graph categorical draw. `filt: None` = the PURE-TEMP body (gumbel
2363/// over the raw softmax), byte-identical to the pre-lane capture; `Some` adds the in-graph
2364/// truncation filter (`filter_stats` + `gumbel_perturb_filtered_ctr`) so the draft draws
2365/// from the same filtered distribution the accept test reconstructs
2366/// (lane/step37-draft-graph-serving-20260830).
2367struct SampledCapArgs<'a> {
2368    ctr: &'a mut CudaSlice<u32>,
2369    perturb: &'a mut CudaSlice<f32>,
2370    q_out: &'a mut CudaSlice<f32>,
2371    seed: u64,
2372    temp: f32,
2373    filt: Option<SampledCapFilter<'a>>,
2374}
2375
2376/// In-graph truncation-filter nodes: the stat slots `filter_stats` fills and the perturb
2377/// reads, plus the filter constants baked into the capture (they live in `s_key`, so a
2378/// request whose filters differ drops the parked graph before this ever goes stale).
2379struct SampledCapFilter<'a> {
2380    rows0: &'a CudaSlice<i32>,
2381    th: &'a mut CudaSlice<f32>,
2382    z: &'a mut CudaSlice<f32>,
2383    mx: &'a mut CudaSlice<f32>,
2384    top_k: i32,
2385    top_p: f32,
2386    min_p: f32,
2387}
2388
2389pub(crate) struct DraftGraphCtx {
2390    g_tok: CudaSlice<u32>,
2391    g_pos: CudaSlice<i32>,
2392    g_seed: CudaSlice<f32>,
2393    g_p: CudaSlice<f32>,
2394    g_ctr: CudaSlice<u32>,
2395    g_q: CudaSlice<f32>,
2396    g_perturb: CudaSlice<f32>,
2397    /// IN-GRAPH filter-stat slots (filtered sampled capture): `filter_stats` writes
2398    /// (th, z, mx) here inside the graph; `gumbel_perturb_filtered_ctr` reads (mx, th) from
2399    /// the same slots. Persistent so the baked pointers survive replays. `g_rows0` is the
2400    /// constant row-index-0 the single-row `filter_stats` launch reads (a captured memcpy
2401    /// source must not be a host temporary).
2402    g_rows0: CudaSlice<i32>,
2403    g_th: CudaSlice<f32>,
2404    g_z: CudaSlice<f32>,
2405    g_mx: CudaSlice<f32>,
2406    q_slots: Vec<CudaSlice<f32>>,
2407    /// DRAFT-SIDE GRAMMAR MASK (lane/draft-mask): packed allowed-set words over the DRAFT
2408    /// head's vocab, at a STABLE address so the captured draft graph's mask node reads the
2409    /// per-position contents the host re-uploads before each replay (the graph-promote
2410    /// pattern from decode.rs). Empty unless the session drafts under a grammar.
2411    g_dmask: CudaSlice<u32>,
2412    /// was `graph` captured WITH the mask node? A parked graph of the wrong shape is dropped.
2413    /// Covers the multi-head `chain` too (single-head and chain are mutually exclusive for a
2414    /// given model, so one flag serves whichever is active).
2415    graph_masked: bool,
2416    graph: Option<cudarc::driver::CudaGraph>,
2417    graph_s: Option<cudarc::driver::CudaGraph>,
2418    /// Multi-head chain graphs (see [`DraftChainGraphs`]): greedy and sampled chains, the
2419    /// chain twins of `graph` / `graph_s`. `chain_s`'s capture identity is `s_key` (shared
2420    /// with `graph_s` — a session is either single-head or chain, never both), and it obeys
2421    /// the same drop rules (key mismatch, penalty regime, mask-shape change).
2422    chain: Option<DraftChainGraphs>,
2423    chain_s: Option<DraftChainGraphs>,
2424    /// Failed-capture memoization for both graphs — LOUD on flip, cleared on pool resume
2425    /// (audit Q2, the TRT #16072 silent-permanent-coverage-loss class).
2426    failed: DraftGraphFallback,
2427    /// Capture identity of `graph_s` — see [`SampledGraphKey`]. `None` iff no sampled graph is
2428    /// parked; a request whose key differs drops the parked graph (and its q slots/keeper).
2429    s_key: Option<SampledGraphKey>,
2430    /// CAPTURE-RETAIN keepers (#68 root cause, 2026-08-04): the warmup-run transients whose
2431    /// pool addresses the captured graph(s) bake. Without these, the transients return to the
2432    /// pool at capture-body exit and later work (burst-boundary prime/fill/commit passes, or a
2433    /// co-served session in the worker) reuses those addresses — the persisted graph's replay
2434    /// then reads/writes live unrelated buffers (exactness corruption, first seen as the ST
2435    /// serve-spec 4B graph-arm corruption; one-shot CLI calls never re-shuffled the pool, which
2436    /// is why run-spec K=1..8 passed on the same checkpoint). Same fix class as
2437    /// capture_graph_retained's gemma/decode.rs sites — hold as long as the graph replays.
2438    keeper: Vec<Box<dyn std::any::Any + Send>>,
2439    keeper_s: Vec<Box<dyn std::any::Any + Send>>,
2440}
2441
2442/// Failed-capture memoization for the two draft graphs (audit Q2, 2026-08-05 — the
2443/// TRT #16072 trap class: pressure-triggered, silent, long-lived coverage loss).
2444///
2445/// Three contracts:
2446/// - LOUD FLIP: `mark_*` returns the warn line exactly on the false→true transition
2447///   (returned, not printed, so the once-per-flip contract is unit-testable); the caller
2448///   `eprintln!`s it UNCONDITIONALLY — a dropped draft graph is never silent. Re-marking
2449///   an already-failed graph returns None (the per-burst memoization that keeps the eager
2450///   fallback from paying a doomed capture attempt every burst).
2451/// - RESET ON RESUME: `reset_on_resume` clears both flags — a parked session resumed by a
2452///   NEW request gets one fresh capture chance instead of carrying a transient-pressure
2453///   failure for the pool's whole lifetime. Returns the note line only when a flag was
2454///   actually set (quiet on the common clean-resume path).
2455/// - Shape-change clears (`clear_*`) stay silent, exactly as before: they precede a fresh
2456///   capture attempt whose own failure would re-flip loudly.
2457#[derive(Default)]
2458pub(crate) struct DraftGraphFallback {
2459    greedy: bool,
2460    sampled: bool,
2461}
2462impl DraftGraphFallback {
2463    fn mark_greedy(&mut self, reason: &str) -> Option<String> {
2464        if self.greedy {
2465            return None;
2466        }
2467        self.greedy = true;
2468        Some(format!(
2469            "[spec] WARN: draft-graph capture failed ({reason}); eager fallback until session resume"
2470        ))
2471    }
2472    fn mark_sampled(&mut self, reason: &str) -> Option<String> {
2473        if self.sampled {
2474            return None;
2475        }
2476        self.sampled = true;
2477        Some(format!(
2478            "[spec] WARN: sampled draft-graph capture failed ({reason}); eager fallback until session resume"
2479        ))
2480    }
2481    fn greedy_failed(&self) -> bool {
2482        self.greedy
2483    }
2484    fn sampled_failed(&self) -> bool {
2485        self.sampled
2486    }
2487    fn clear_greedy(&mut self) {
2488        self.greedy = false;
2489    }
2490    fn clear_sampled(&mut self) {
2491        self.sampled = false;
2492    }
2493    /// Pool-resume reset: both graphs get a fresh capture chance. Some(note) iff any flag
2494    /// was set (so clean resumes stay quiet).
2495    pub(crate) fn reset_on_resume(&mut self) -> Option<String> {
2496        if !self.greedy && !self.sampled {
2497            return None;
2498        }
2499        let which = match (self.greedy, self.sampled) {
2500            (true, true) => "greedy+sampled",
2501            (true, false) => "greedy",
2502            _ => "sampled",
2503        };
2504        self.greedy = false;
2505        self.sampled = false;
2506        Some(format!(
2507            "[spec] draft-graph fallback reset on session resume ({which}); recapture eligible"
2508        ))
2509    }
2510}
2511
2512impl DraftGraphCtx {
2513    fn new(e: &Engine, n_embd: usize, qlen: usize) -> Result<Self, Box<dyn std::error::Error>> {
2514        Ok(DraftGraphCtx {
2515            g_tok: e.alloc_u32_zeroed(1)?,
2516            g_pos: e.htod_i32(&[0])?,
2517            g_seed: e.zeros(n_embd)?,
2518            g_p: e.zeros(1)?,
2519            g_ctr: e.alloc_u32_zeroed(1)?,
2520            g_q: e.zeros(qlen)?,
2521            g_perturb: e.zeros(qlen)?,
2522            g_rows0: e.htod_i32(&[0])?,
2523            g_th: e.zeros(1)?,
2524            g_z: e.zeros(1)?,
2525            g_mx: e.zeros(1)?,
2526            q_slots: Vec::new(),
2527            g_dmask: e.alloc_u32_zeroed(1)?,
2528            graph_masked: false,
2529            graph: None,
2530            graph_s: None,
2531            chain: None,
2532            chain_s: None,
2533            failed: DraftGraphFallback::default(),
2534            s_key: None,
2535            keeper: Vec::new(),
2536            keeper_s: Vec::new(),
2537        })
2538    }
2539}
2540
2541pub(crate) struct MtpScratch {
2542    kv: KvLayer,
2543    /// Logical row capacity. On the graph/DC draft path it also doubles as fa_decode_dc's
2544    /// bucket_max: n_splits is sized from it ONCE, so the graph captured at round 0 stays valid
2545    /// for every later t_kv. Step35 refuses that path and may back this logical extent with the
2546    /// smaller host-indexed SWA ring instead.
2547    cap: usize,
2548    extra: Vec<MtpScratchPlane>,
2549}
2550
2551struct MtpScratchPlane {
2552    kv: KvLayer,
2553    cap: usize,
2554}
2555
2556fn mtp_scratch_layout(
2557    cfg: &memra_gguf::config::ModelConfig,
2558    geom: Option<&crate::hybrid::DraftGeom>,
2559) -> (usize, usize, usize, usize) {
2560    // Student draft heads carry fewer KV heads (head_dim unchanged) -> smaller scratch rows.
2561    let n_head_kv = geom.map(|g| g.n_head_kv).unwrap_or(cfg.n_head_kv as usize);
2562    let head_dim_k = cfg.head_dim_k as usize;
2563    let head_dim_v = cfg.head_dim_v as usize;
2564    assert!(
2565        head_dim_k.is_multiple_of(32) && head_dim_v.is_multiple_of(32),
2566        "KVQUANT requires head_dim%32==0 (MTP scratch)"
2567    );
2568    let kv_dim_k = head_dim_k * n_head_kv;
2569    let kv_dim_v = head_dim_v * n_head_kv;
2570    // The fp8-KV arm deliberately does not reach the draft scratch; keep the exact format
2571    // policy shared with `MtpScratch::new` so admission scales the same allocation.
2572    let (kbb, vbb) = crate::kv_blk_bytes();
2573    let k_tok_bytes = (kv_dim_k / 32) * kbb;
2574    let v_tok_bytes = (kv_dim_v / 32) * vbb;
2575    (kv_dim_k, kv_dim_v, k_tok_bytes, v_tok_bytes)
2576}
2577
2578fn mtp_chain_head_index(step: usize, head_count: usize) -> usize {
2579    assert!(head_count > 0, "MTP chain requires at least one head");
2580    step % head_count
2581}
2582
2583impl MtpScratch {
2584    fn alloc_plane(
2585        e: &Engine,
2586        cfg: &memra_gguf::config::ModelConfig,
2587        plan: &memra_gguf::model_plan::ModelPlan,
2588        cap: usize,
2589        geom: Option<&crate::hybrid::DraftGeom>,
2590    ) -> Result<MtpScratchPlane, Box<dyn std::error::Error>> {
2591        let (kv_dim_k, kv_dim_v, k_tok_bytes, v_tok_bytes) = mtp_scratch_layout(cfg, geom);
2592        let ring = if crate::cache::swa_ring_on()
2593            && crate::plan_backend::decode_batch_program(plan)
2594                == crate::plan_backend::DecodeBatchProgram::SlidingGatedMoe
2595        {
2596            let window = plan
2597                .layers
2598                .iter()
2599                .find_map(|layer| match layer.attention {
2600                    memra_gguf::model_plan::AttentionPlan::SlidingWindow { window, .. } => {
2601                        Some(window as usize)
2602                    }
2603                    _ => None,
2604                })
2605                .ok_or("sliding-gated-MoE draft scratch has no sliding-window layer")?;
2606            Some(crate::cache::KvRing::new(
2607                crate::cache::swa_ring_rows(window, cap),
2608                window,
2609            ))
2610        } else {
2611            None
2612        };
2613        let alloc_rows = ring.as_ref().map(crate::cache::KvRing::rows).unwrap_or(cap);
2614        // Ring-backed planes arm the device base mirror for the dcw draft arm (see
2615        // KvLayer::base_d): the captured chain derives its physical rows from
2616        // (len_d, base_d, window) with zero per-token node updates.
2617        let base_d = match ring.as_ref() {
2618            Some(_) => Some(e.htod_i32(&[0])?),
2619            None => None,
2620        };
2621        Ok(MtpScratchPlane {
2622            kv: KvLayer {
2623                k: e.alloc_u8(alloc_rows * k_tok_bytes)?,
2624                v: e.alloc_u8(alloc_rows * v_tok_bytes)?,
2625                kv_dim_k,
2626                kv_dim_v,
2627                k_tok_bytes,
2628                v_tok_bytes,
2629                len: 0,
2630                ring,
2631                len_d: e.htod_i32(&[0])?,
2632                base_d,
2633            },
2634            cap,
2635        })
2636    }
2637
2638    fn new(
2639        e: &Engine,
2640        cfg: &memra_gguf::config::ModelConfig,
2641        plan: &memra_gguf::model_plan::ModelPlan,
2642        cap: usize,
2643        geom: Option<&crate::hybrid::DraftGeom>,
2644    ) -> Result<Self, Box<dyn std::error::Error>> {
2645        // env-selected KV formats (default 34/24). The fp8-KV arm (MEMRA_KV_FP8) deliberately
2646        // does NOT reach the draft scratch: fp8 drafts drifted acceptance 69-88% -> 46%
2647        // (2026-07-12 A/B); the scratch is tiny, so it keeps baseline q8_0/q5_1 numerics
2648        // while the TRUNK cache carries the fp8 depth win. Scratch append/fa pass g=false.
2649        let primary = Self::alloc_plane(e, cfg, plan, cap, geom)?;
2650        Ok(MtpScratch {
2651            kv: primary.kv,
2652            cap: primary.cap,
2653            extra: Vec::new(),
2654        })
2655    }
2656
2657    fn push_plane(
2658        &mut self,
2659        e: &Engine,
2660        cfg: &memra_gguf::config::ModelConfig,
2661        plan: &memra_gguf::model_plan::ModelPlan,
2662        geom: Option<&crate::hybrid::DraftGeom>,
2663    ) -> Result<(), Box<dyn std::error::Error>> {
2664        self.extra
2665            .push(Self::alloc_plane(e, cfg, plan, self.cap, geom)?);
2666        Ok(())
2667    }
2668
2669    fn plane_count(&self) -> usize {
2670        1 + self.extra.len()
2671    }
2672
2673    fn plane(&self, index: usize) -> (&KvLayer, usize) {
2674        if index == 0 {
2675            (&self.kv, self.cap)
2676        } else {
2677            let plane = &self.extra[index - 1];
2678            (&plane.kv, plane.cap)
2679        }
2680    }
2681
2682    fn plane_mut(&mut self, index: usize) -> (&mut KvLayer, usize) {
2683        if index == 0 {
2684            (&mut self.kv, self.cap)
2685        } else {
2686            let plane = &mut self.extra[index - 1];
2687            (&mut plane.kv, plane.cap)
2688        }
2689    }
2690
2691    // #[track_caller]: set_len/set_plane_len have eight call sites (checkpoint restore, spec
2692    // rollback, session grow, seed replay ...) and the lap failure needs to say WHICH one, not
2693    // just that a rewind was refused.
2694    #[track_caller]
2695    fn set_plane_len(
2696        &mut self,
2697        e: &Engine,
2698        index: usize,
2699        n: usize,
2700    ) -> Result<(), Box<dyn std::error::Error>> {
2701        let caller = std::panic::Location::caller();
2702        let (kv, cap) = self.plane_mut(index);
2703        if let Some(ring) = kv.ring.as_ref()
2704            && !ring.can_rewind_to(n)
2705        {
2706            // NAME THE NUMBERS (2026-08-28). This error is a step37 serving blocker on the
2707            // vendor-default shape and it fires from more than one call path with more than
2708            // one trigger: a long generation walks the checkpoint out of the ring, but a
2709            // ~4.5k-token prompt also fails within 5 s of prime, which accumulation cannot
2710            // explain. A bare message forced two rounds of guessing; the operands make each
2711            // trigger name itself.
2712            let raw = n.saturating_sub(ring.window().saturating_sub(1));
2713            return Err(format!(
2714                    "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})",
2715                    ring.window(),
2716                    ring.base(),
2717                    ring.rows(),
2718                    raw & !31usize,
2719                )
2720                .into());
2721        }
2722        kv.len = n;
2723        e.set_i32_one(&mut kv.len_d, n as i32)
2724    }
2725
2726    /// Set BOTH length counters: the host mirror AND the device len_d the captured append/fa read
2727    /// (a 4-byte in-place htod — the counter pointer is baked into the graph, never realloc'd).
2728    /// This is the ONLY truncation/rollback mechanism the persistent draft KV needs.
2729    #[track_caller]
2730    fn set_len(&mut self, e: &Engine, n: usize) -> Result<(), Box<dyn std::error::Error>> {
2731        let caller = std::panic::Location::caller();
2732        if !self.can_rewind_to(n) {
2733            // set_plane_len re-checks and reports the operands; call it so the failure carries
2734            // which plane refused and why, instead of this bare aggregate.
2735            for index in 0..self.plane_count() {
2736                self.set_plane_len(e, index, n)?;
2737            }
2738            return Err(format!(
2739                "SWA ring MTP checkpoint has been lapped; full re-prime required (aggregate rewind_to={n}, no single plane reported, called from {caller})"
2740            )
2741            .into());
2742        }
2743        for index in 0..self.plane_count() {
2744            self.set_plane_len(e, index, n)?;
2745        }
2746        Ok(())
2747    }
2748
2749    fn can_rewind_to(&self, n: usize) -> bool {
2750        (0..self.plane_count()).all(|index| {
2751            self.plane(index)
2752                .0
2753                .ring
2754                .as_ref()
2755                .is_none_or(|ring| ring.can_rewind_to(n))
2756        })
2757    }
2758
2759    /// Pre-arm ring headroom for `rows` upcoming DEVICE-COUNTER appends (the dcw draft arm):
2760    /// a captured chain cannot rebase mid-replay, so any rebase the coming appends could need
2761    /// happens HERE, host-side, before the capture warmups or the round's replays (the rebase
2762    /// arm of `prepare_kv_append` also refreshes the plane's `base_d` device mirror). No-op on
2763    /// flat planes and when the ring already has room; `len` is untouched either way.
2764    fn ensure_dcw_headroom(
2765        &mut self,
2766        e: &Engine,
2767        rows: usize,
2768    ) -> Result<(), Box<dyn std::error::Error>> {
2769        for index in 0..self.plane_count() {
2770            let (kv, _) = self.plane_mut(index);
2771            let Some(ring) = kv.ring.as_ref() else {
2772                continue;
2773            };
2774            let retain = memra_kv::swa_retain_from(kv.len, ring.window(), ring.base());
2775            e.prepare_kv_append(kv, retain, rows)?;
2776        }
2777        Ok(())
2778    }
2779}
2780
2781/// Retained verify intermediates for the REPLAY-FREE partial accept (2026-07-03, the profiled
2782/// #1 spec cost at long ctx: the partial-accept replay was a DUPLICATE trunk pass — ~0.54 extra
2783/// full weight reads per round — recomputing columns the verify had already produced
2784/// bit-identically). Holds, per linear layer, everything needed to rebuild its recurrent state
2785/// to "after the first j verify columns" WITHOUT re-running the trunk:
2786/// - BATCHED-path layers (`gdn`): the exact token-major inputs the round's ONE gdn_scan
2787///   consumed. A prefix re-run of the SAME kernel (t=j) from the snapshot state is bit-identical
2788///   to the first j iterations of the verify's scan — the kernel's t-loop carries state in
2789///   registers and iteration t never depends on T. `qkv_mixed` (the conv input) feeds the
2790///   pure-copy ring rebuild.
2791/// - PER-COLUMN-path layers (`cols`): dtod clones of (conv_state, ssm_state) taken after each
2792///   column 0..t-2 — pure copies of the actual chain states (the last column is never a rebuild
2793///   target: j <= t-1).
2794///   Full-attn layers need nothing: their verify KV rows are bit-identical to eager's (the
2795///   decode-exact contract; verify-probe pins it), so rollback = len truncation.
2796struct GdnStash {
2797    qkv_mixed: CudaSlice<f32>, // [t, conv_dim] token-major (conv input)
2798    q_l2: CudaSlice<f32>,
2799    k_l2: CudaSlice<f32>,
2800    v_g: CudaSlice<f32>, // [t, num_v, d_state]
2801    g_log: CudaSlice<f32>,
2802    beta: CudaSlice<f32>, // [t, num_v]
2803}
2804pub(crate) struct VerifyCkpt {
2805    gdn: Vec<Option<GdnStash>>, // [n_layer], Some iff batched linear path ran
2806    #[allow(clippy::type_complexity)]
2807    // allow: one-shot composite type; naming it would hide the shape that matters at the call site
2808    cols: Vec<Option<Vec<(CudaSlice<f32>, CudaSlice<f32>)>>>, // [n_layer][col] = (conv, ssm) after col
2809}
2810/// Opaque handle for the dspark round (dflash.rs) — VerifyCkpt stays spec-private.
2811pub(crate) struct DsparkVerifyCkpt(VerifyCkpt);
2812
2813/// Engine-bundle slice 3 (DSF-ROUNDCOST-20260820 §2 row 4 / §5 rank 1): bucketed CUDA
2814/// graphs for the dspark verify's LINEAR-layer segments. The measured verify is ~2,800
2815/// eager launches whose residual cost is DEVICE-side per-launch overhead (slice 2 proved
2816/// host dispatch is not the binder: fully-deferred dispatch bought ~0 wall). The 48 GDN
2817/// layers between full-attention layers are shape-static given vt — no positions, no
2818/// t_kv, state addressed through pointer tables — so runs of them capture per
2819/// (segment, vt) and replay as ONE graph launch each. Full-attention layers stay eager
2820/// (their per-row append/fa arm picks are t_kv-driven — the exec-update extension).
2821///
2822/// Per round out-of-graph: one pointer-table refresh (gdn ping-pong moves the canonical
2823/// handles), one input-staging copy per segment, host parity bookkeeping. Captured via
2824/// `capture_graph_retained` (2 warmups + capture, keeper retains warmup transients so
2825/// pool addresses stay stable); the warmups EXECUTE, so segment conv/ssm state is saved
2826/// before and restored after — the graph's first real launch starts from the exact
2827/// pre-round state. The ckpt column stash rides persistent slabs (written inside the
2828/// graph as memcpy nodes); commit reads them via `dspark_commit_prefix_slab`.
2829/// `MEMRA_DSPARK_VERIFY_GRAPH=0` reverts to the eager walk (byte-identical body).
2830pub(crate) struct DsparkVerifyGraphs {
2831    /// Linear-attention layer indices ascending; `lin_pos[il]` = index into the vecs.
2832    lin: Vec<usize>,
2833    lin_pos: std::collections::HashMap<usize, usize>,
2834    /// [n_lin x 6] pointer table (conv, s0, s1, conv, s1, s0 per layer), refreshed per
2835    /// verify from the live handles; layer il's slice starts at lin_pos[il]*6.
2836    table_all: CudaSlice<u64>,
2837    host_table: Vec<u64>,
2838    /// Persistent per-layer ckpt stash slabs: row r of the verify at slab offset
2839    /// r*words. Shared by every (segment, vt) bucket — one verify runs at a time.
2840    stash_conv: Vec<CudaSlice<f32>>,
2841    stash_ssm: Vec<CudaSlice<f32>>,
2842    conv_words: usize,
2843    ssm_words: usize,
2844    /// Per-vt input/output staging (stable addresses the graphs bake).
2845    stage: std::collections::HashMap<usize, (CudaSlice<f32>, CudaSlice<f32>)>,
2846    /// Per-vt dflash tap-sink buffers — the captured segments bake the tap dst address,
2847    /// so the sink buffer must live (and persist) with the graphs, not with the round.
2848    pub(crate) tap_bufs: std::collections::HashMap<usize, CudaSlice<f32>>,
2849    graphs: std::collections::HashMap<(usize, usize), DsparkSegGraph>,
2850    /// Warmup-corruption guard scratch: pre-capture conv/ssm of every linear layer
2851    /// (sized n_lin — the slice-4c full-verify warmups execute the whole walk).
2852    save_conv: CudaSlice<f32>,
2853    save_ssm: CudaSlice<f32>,
2854    max_run: usize,
2855    n_embd: usize,
2856    /// Set by the verify walk: this round's linear ckpt lives in the slabs (the caller
2857    /// commits through `dspark_commit_prefix_slab` instead of the cols arm).
2858    pub(crate) round_slab: bool,
2859    // ---- slice 4c: full-verify single graph per (vt, rung) ----
2860    /// Full-attention layer indices ascending; `fa_pos[il]` = index into the vec.
2861    fa: Vec<usize>,
2862    fa_pos: std::collections::HashMap<usize, usize>,
2863    /// [n_fa x 2 x t_cap] interleaved (k,v) base-pointer pairs, refreshed per verify;
2864    /// layer il's slice starts at `fa_pos[il] * 2 * t_cap` (the seqs twins read pairs
2865    /// [2z], z < t <= t_cap, so one t_cap-sized table serves every vt).
2866    fa_table: CudaSlice<u64>,
2867    fa_host_table: Vec<u64>,
2868    t_cap: usize,
2869    /// Per-vt position staging for the captured bodies — contents refreshed per round
2870    /// (rope reads row r; the seqs twins derive append slot and T_kv per z from it).
2871    pos_stage: std::collections::HashMap<usize, CudaSlice<i32>>,
2872    /// Full-verify graphs keyed (vt, rung_end, hi).
2873    full: std::collections::HashMap<(usize, usize, usize), DsparkSegGraph>,
2874    /// Largest n with every layer in [0, n) linear or full-attention (walk coverage).
2875    covered: usize,
2876    /// Every layer in [0, n) is linear or full-attention (no MLA/unknown mixers) — the
2877    /// full-verify capture walks all of them.
2878    walk_uniform: bool,
2879    /// Last `(captures, device graph-mem reserved bytes)` reading taken by
2880    /// `HybridModel::dspark_vg_admission_debt` — the two-point base of the MARGINAL debt
2881    /// projection (see `dspark_vg_debt_projection`; a mean-based reading extrapolated the
2882    /// pool's one-time shared allocation and reserved 8.5 GB of phantom VRAM).
2883    debt_obs: Option<(usize, usize)>,
2884}
2885
2886struct DsparkSegGraph {
2887    graph: cudarc::driver::CudaGraph,
2888    _keeper: Vec<Box<dyn std::any::Any + Send>>,
2889}
2890
2891/// Per-call arguments of [`HybridModel::qwen35_tparallel_fa_layer`] — one struct so the
2892/// eager walk and the slice-4c captured full-verify graphs hand the SAME body its two
2893/// modes without a second copy of the math.
2894pub(crate) struct FaLayerArgs<'a> {
2895    /// [T] per-row positions (device): rope reads them row-indexed; the seqs twins read
2896    /// them per-z (append slot = pos, T_kv = pos + 1).
2897    pub pos_d: &'a CudaSlice<i32>,
2898    /// Verify-level lazy per-row 1-element position buffers — only the per-row fallback
2899    /// arm builds/uses them (graph mode refuses that arm).
2900    pub pos_rows: &'a mut Option<Vec<CudaSlice<i32>>>,
2901    pub pos0: usize,
2902    pub seqs_append: bool,
2903    pub batch_fa_on: bool,
2904    /// Some((kv pointer table, offset-in-u64s, rung_end)) = captured-graph mode.
2905    pub graph_cap: Option<(&'a CudaSlice<u64>, usize, usize)>,
2906    /// ROUND-STREAM (lane/draftcost-moe, v0.100 train merge): Some((token stream, device
2907    /// round counter)) routes the FA attend through the dc rows kernels and the Linear
2908    /// mixer through `linear_attn_verify_t` (the stream arms the old inline body carried).
2909    /// Never armed together with `graph_cap` (the verify-level merge guard refuses).
2910    pub stream: Option<(&'a CudaSlice<u32>, &'a CudaSlice<i32>)>,
2911    /// VerifyCkpt for the stream-Linear arm's GdnStash install; None in graph mode and
2912    /// for FA layers that never touch it.
2913    pub ckpt: Option<&'a mut VerifyCkpt>,
2914}
2915
2916// SAFETY: `CudaGraph` is not marked Send by cudarc because its raw driver handles carry
2917// no automatic trait; CUDA driver graph handles are context-scoped rather than
2918// OS-thread-affine (the SpecPipeSessionPtr precedent above). The ctx lives in
2919// `HybridModel::dspark_vgraphs` behind a Mutex and every touch happens on the engine's
2920// single decode-stream thread.
2921unsafe impl Send for DsparkVerifyGraphs {}
2922
2923impl DsparkVerifyGraphs {
2924    /// Live capture count (segment + full graphs) — the denominator of
2925    /// [`dspark_vg_debt_projection`]'s observed bytes/capture mean.
2926    pub(crate) fn captures(&self) -> usize {
2927        self.graphs.len() + self.full.len()
2928    }
2929
2930    /// Take the marginal-growth debt reading and record this observation for the next one.
2931    /// Called under the pool mutex by `HybridModel::dspark_vg_admission_debt`.
2932    pub(crate) fn admission_debt(&mut self, reserved_bytes: usize) -> usize {
2933        let captures = self.captures();
2934        let debt =
2935            dspark_vg_debt_projection(captures, dspark_vg_cap(), reserved_bytes, self.debt_obs);
2936        if captures > 0 {
2937            match self.debt_obs {
2938                Some((c0, _)) if captures <= c0 => {}
2939                _ => self.debt_obs = Some((captures, reserved_bytes)),
2940            }
2941        }
2942        debt
2943    }
2944
2945    /// Build for this cache's shape. None when there are no linear layers, sizes are
2946    /// non-uniform, or the trunk keeps a gemma4 config (never on the qwen35 family).
2947    pub(crate) fn new(
2948        e: &Engine,
2949        cache: &Cache,
2950        t_max: usize,
2951        n_embd: usize,
2952    ) -> Result<Option<Self>, Box<dyn std::error::Error>> {
2953        let lin: Vec<usize> = (0..cache.recur.len())
2954            .filter(|&il| cache.recur[il].is_some())
2955            .collect();
2956        if lin.is_empty() || t_max < 2 {
2957            return Ok(None);
2958        }
2959        let first = cache.recur[lin[0]].as_ref().unwrap();
2960        let (conv_words, ssm_words) = (first.conv_state.len(), first.ssm_state.len());
2961        for &il in &lin {
2962            let rl = cache.recur[il].as_ref().unwrap();
2963            if rl.conv_state.len() != conv_words || rl.ssm_state.len() != ssm_words {
2964                return Ok(None);
2965            }
2966        }
2967        let n = lin.len();
2968        let mut lin_pos = std::collections::HashMap::with_capacity(n);
2969        for (k, &il) in lin.iter().enumerate() {
2970            lin_pos.insert(il, k);
2971        }
2972        // longest run of consecutive linear layers (save-scratch sizing)
2973        let mut max_run = 1usize;
2974        let mut run = 1usize;
2975        for w in lin.windows(2) {
2976            if w[1] == w[0] + 1 {
2977                run += 1;
2978                max_run = max_run.max(run);
2979            } else {
2980                run = 1;
2981            }
2982        }
2983        let rows = t_max - 1;
2984        let mut stash_conv = Vec::with_capacity(n);
2985        let mut stash_ssm = Vec::with_capacity(n);
2986        for _ in 0..n {
2987            stash_conv.push(e.uninit(rows * conv_words)?);
2988            stash_ssm.push(e.uninit(rows * ssm_words)?);
2989        }
2990        let host_table = vec![0u64; n * 6];
2991        let table_all = e.htod_u64(&host_table)?;
2992        // slice 4c: full-attention census for the full-verify graphs.
2993        let fa: Vec<usize> = (0..cache.kv.len())
2994            .filter(|&il| cache.kv[il].is_some())
2995            .collect();
2996        let mut fa_pos = std::collections::HashMap::with_capacity(fa.len());
2997        for (k, &il) in fa.iter().enumerate() {
2998            fa_pos.insert(il, k);
2999        }
3000        let n_layers = cache.kv.len().max(cache.recur.len());
3001        // exactly one of (linear state, kv cache) per layer — no MLA/unknown mixers.
3002        let walk_uniform = (0..n_layers).all(|il| {
3003            cache.recur.get(il).is_some_and(|r| r.is_some())
3004                != cache.kv.get(il).is_some_and(|k| k.is_some())
3005        });
3006        // Contiguous covered prefix: the largest n such that every layer in [0, n) is
3007        // linear or full-attention. The TRUNK walk is [0, layers.len()) and the cache
3008        // vecs can carry EXTRA state slots past it (the q38 export keeps the MTP head
3009        // layer's kv at the tail — hi == lin+fa never held, the s4c battery's zero
3010        // 'full' captures). The full-graph guard is walk coverage, not slot arithmetic.
3011        let covered = (0..n_layers)
3012            .take_while(|il| lin_pos.contains_key(il) || fa_pos.contains_key(il))
3013            .count();
3014        let t_cap = t_max;
3015        let fa_host_table = vec![0u64; fa.len() * 2 * t_cap];
3016        let fa_table = e.htod_u64(&fa_host_table)?;
3017        Ok(Some(Self {
3018            lin,
3019            lin_pos,
3020            table_all,
3021            host_table,
3022            stash_conv,
3023            stash_ssm,
3024            conv_words,
3025            ssm_words,
3026            stage: std::collections::HashMap::new(),
3027            tap_bufs: std::collections::HashMap::new(),
3028            graphs: std::collections::HashMap::new(),
3029            save_conv: e.uninit(n * conv_words)?,
3030            save_ssm: e.uninit(n * ssm_words)?,
3031            max_run,
3032            n_embd,
3033            round_slab: false,
3034            fa,
3035            fa_pos,
3036            fa_table,
3037            fa_host_table,
3038            t_cap,
3039            pos_stage: std::collections::HashMap::new(),
3040            full: std::collections::HashMap::new(),
3041            covered,
3042            walk_uniform,
3043            debt_obs: None,
3044        }))
3045    }
3046
3047    /// Rebuild the pointer tables from the live handles (once per verify — the gdn
3048    /// ping-pong swaps the canonical/alt handles between rounds; a fresh generation's
3049    /// cache buffers land at new addresses; a stale table would read the wrong state).
3050    pub(crate) fn refresh_tables(
3051        &mut self,
3052        e: &Engine,
3053        cache: &Cache,
3054    ) -> Result<(), Box<dyn std::error::Error>> {
3055        use cudarc::driver::DevicePtr;
3056        {
3057            let s = &e.gpu.stream();
3058            for (k, &il) in self.lin.iter().enumerate() {
3059                let rl = cache.recur[il].as_ref().unwrap();
3060                let (pc, _g0) = rl.conv_state.device_ptr(s);
3061                let (p0, _g1) = rl.ssm_state.device_ptr(s);
3062                let (p1, _g2) = rl.ssm_state_alt.device_ptr(s);
3063                let o = k * 6;
3064                self.host_table[o] = pc;
3065                self.host_table[o + 1] = p0;
3066                self.host_table[o + 2] = p1;
3067                self.host_table[o + 3] = pc;
3068                self.host_table[o + 4] = p1;
3069                self.host_table[o + 5] = p0;
3070            }
3071            for (k, &il) in self.fa.iter().enumerate() {
3072                let kvl = cache.kv[il].as_ref().unwrap();
3073                let (pk, _g0) = kvl.k.device_ptr(s);
3074                let (pv, _g1) = kvl.v.device_ptr(s);
3075                let o = k * 2 * self.t_cap;
3076                for z in 0..self.t_cap {
3077                    self.fa_host_table[o + 2 * z] = pk;
3078                    self.fa_host_table[o + 2 * z + 1] = pv;
3079                }
3080            }
3081        }
3082        e.htod_u64_into(&self.host_table, &mut self.table_all)?;
3083        if !self.fa_host_table.is_empty() {
3084            e.htod_u64_into(&self.fa_host_table, &mut self.fa_table)?;
3085        }
3086        Ok(())
3087    }
3088
3089    /// Slice 4c eligibility: Some(rung_end) when this round can replay (or capture) a
3090    /// full-verify graph — the whole walk [lo, hi) is covered, every layer is linear or
3091    /// full-attention, and ALL of the round's per-row t_kv values take the v4-seqs arm
3092    /// on ONE `fa_split_keys` ladder step that the rung also sits on (the straddle law;
3093    /// both gates are t_kv intervals, so ends-inside means all-inside). The rung is the
3094    /// round's next power of two — grid/partial sizing only (`n_splits_max` is pure
3095    /// stride; splits >= ns_eff write the empty partial the combine never reads), so one
3096    /// captured graph is bit-identical for every round the rung covers.
3097    #[allow(clippy::too_many_arguments)]
3098    pub(crate) fn full_rung(
3099        &self,
3100        model: &crate::hybrid::HybridModel,
3101        cache: &Cache,
3102        lo: usize,
3103        hi: usize,
3104        t: usize,
3105        seqs_arms_on: bool,
3106    ) -> Option<usize> {
3107        if std::env::var("MEMRA_DSPARK_FULLG_DEBUG").as_deref() == Ok("1") {
3108            static ONCE: std::sync::Once = std::sync::Once::new();
3109            let len0 = self
3110                .fa
3111                .first()
3112                .and_then(|&il| cache.kv[il].as_ref())
3113                .map(|k| k.len);
3114            ONCE.call_once(|| {
3115                eprintln!(
3116                    "[fullg-debug] walk_uniform={} covered={} seqs_arms_on={} fa_rows_on={} t={} lo={} hi={} lin={} fa={} t_cap={} len0={:?}",
3117                    self.walk_uniform, self.covered, seqs_arms_on, dspark_fa_rows_on(), t, lo, hi,
3118                    self.lin.len(), self.fa.len(), self.t_cap, len0
3119                );
3120            });
3121        }
3122        if !self.walk_uniform
3123            || !seqs_arms_on
3124            || !dspark_fa_rows_on()
3125            || t < 2
3126            || lo != 0
3127            || hi > self.covered
3128            || t > self.t_cap
3129            || self.fa.is_empty()
3130        {
3131            return None;
3132        }
3133        let cfg = &model.cfg;
3134        let head_dim_global = cfg.head_dim_k as usize;
3135        let nkv = cfg.n_head_kv as usize;
3136        let kvl0 = cache.kv[self.fa[0]].as_ref().unwrap();
3137        // the z-batched twins read stacked rows at the cache's kv dims — must equal the
3138        // projection stride (the body's guard, hoisted so ineligible models fall back
3139        // instead of refusing mid-capture).
3140        let geom = cfg.full_attention_geometry_at(self.fa[0] as u32);
3141        let kv_dim = geom.n_head_kv as usize * geom.head_dim_k as usize;
3142        if kvl0.kv_dim_k != kv_dim || kvl0.kv_dim_v != kv_dim {
3143            return None;
3144        }
3145        let len0 = kvl0.len;
3146        let (t_kv_first, t_kv_last) = (len0 + 1, len0 + t);
3147        if !crate::fa_seqs_eligible(t_kv_first, head_dim_global)
3148            || !crate::fa_seqs_eligible(t_kv_last, head_dim_global)
3149            || crate::fa_split_keys(t_kv_first, nkv) != crate::fa_split_keys(t_kv_last, nkv)
3150        {
3151            return None;
3152        }
3153        let rung = t_kv_last.next_power_of_two().max(256);
3154        if crate::fa_split_keys(rung, nkv) != crate::fa_split_keys(t_kv_last, nkv) {
3155            return None;
3156        }
3157        Some(rung)
3158    }
3159
3160    /// Run the WHOLE verify walk [lo, hi) as one captured graph at (vt=t, rung): stage
3161    /// the residual + refresh the per-vt position staging, capture on first encounter
3162    /// (2 executing warmups bracketed by a full linear-state save/restore; KV warmup
3163    /// appends write the exact slots the replay writes — idempotent), launch, then apply
3164    /// the host bookkeeping the captured body skipped (per-linear-layer parity swap for
3165    /// odd t, per-fa-layer len bump). Returns the fresh residual.
3166    #[allow(clippy::too_many_arguments)]
3167    #[allow(clippy::map_entry)] // allow: the init body is fallible (`?`); Entry::or_insert_with cannot propagate errors
3168    pub(crate) fn run_full(
3169        &mut self,
3170        model: &crate::hybrid::HybridModel,
3171        e: &Engine,
3172        lo: usize,
3173        hi: usize,
3174        x: &CudaSlice<f32>,
3175        t: usize,
3176        pos0: usize,
3177        rung: usize,
3178        cache: &mut Cache,
3179    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3180        let n_embd = self.n_embd;
3181        if !self.stage.contains_key(&t) {
3182            let xin = e.uninit(t * n_embd)?;
3183            let xout = e.uninit(t * n_embd)?;
3184            self.stage.insert(t, (xin, xout));
3185        }
3186        if !self.pos_stage.contains_key(&t) {
3187            self.pos_stage.insert(t, e.htod_i32(&vec![0i32; t])?);
3188        }
3189        // Per-round refresh: position contents + input staging (both addresses are baked
3190        // by the captured bodies; only their CONTENTS change round to round).
3191        {
3192            let pos_host: Vec<i32> = (0..t).map(|r| (pos0 + r) as i32).collect();
3193            let pb = self.pos_stage.get_mut(&t).unwrap();
3194            e.htod_i32_into(pb, &pos_host)?;
3195            let (xin, _) = self.stage.get_mut(&t).unwrap();
3196            e.copy_into(xin, 0, x, t * n_embd)?;
3197        }
3198        let key = (t, rung, hi);
3199        if !self.full.contains_key(&key) {
3200            // The warmups EXECUTE the whole walk on live state — save every linear
3201            // layer's conv + canonical ssm first, restore after (KV needs no restore:
3202            // graph mode never bumps host lens and the appends write this round's own
3203            // slots).
3204            for (k, &il) in self.lin.iter().enumerate() {
3205                let rl = cache.recur[il].as_ref().unwrap();
3206                e.copy_into(
3207                    &mut self.save_conv,
3208                    k * self.conv_words,
3209                    &rl.conv_state,
3210                    self.conv_words,
3211                )?;
3212                e.copy_into(
3213                    &mut self.save_ssm,
3214                    k * self.ssm_words,
3215                    &rl.ssm_state,
3216                    self.ssm_words,
3217                )?;
3218            }
3219            let (graph, keeper) = {
3220                let table_all = &self.table_all;
3221                let lin_pos = &self.lin_pos;
3222                let fa_pos = &self.fa_pos;
3223                let fa_table = &self.fa_table;
3224                let t_cap = self.t_cap;
3225                let stash_conv = &mut self.stash_conv;
3226                let stash_ssm = &mut self.stash_ssm;
3227                let pos_d: &CudaSlice<i32> = &self.pos_stage[&t];
3228                let (xin, xout) = self
3229                    .stage
3230                    .get_mut(&t)
3231                    .map(|(a, b)| (&*a, b))
3232                    .expect("stage bucket created above");
3233                let cache_ref: &mut Cache = cache;
3234                let iflag = if std::env::var("MEMRA_DSPARK_VG_AUTOFREE").as_deref() == Ok("1") {
3235                    cudarc::driver::sys::CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH
3236                } else {
3237                    cudarc::driver::sys::CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_USE_NODE_PRIORITY
3238                };
3239                e.capture_graph_retained_flags(iflag, move |e| {
3240                    let mut xc: Option<CudaSlice<f32>> = None;
3241                    for il in lo..hi {
3242                        let xr: &CudaSlice<f32> = xc.as_ref().unwrap_or(xin);
3243                        let nx = if let Some(&k) = lin_pos.get(&il) {
3244                            model.qwen35_tparallel_linear_layer(
3245                                e,
3246                                il,
3247                                xr,
3248                                t,
3249                                cache_ref,
3250                                None,
3251                                Some((&mut stash_conv[k], &mut stash_ssm[k])),
3252                                Some((table_all, k * 6)),
3253                            )?
3254                        } else if let Some(&kf) = fa_pos.get(&il) {
3255                            let mut no_rows: Option<Vec<CudaSlice<i32>>> = None;
3256                            model.qwen35_tparallel_fa_layer(
3257                                e,
3258                                il,
3259                                xr,
3260                                t,
3261                                cache_ref,
3262                                FaLayerArgs {
3263                                    pos_d,
3264                                    pos_rows: &mut no_rows,
3265                                    pos0,
3266                                    seqs_append: true,
3267                                    batch_fa_on: true,
3268                                    graph_cap: Some((fa_table, kf * 2 * t_cap, rung)),
3269                                    stream: None,
3270                                    ckpt: None,
3271                                },
3272                            )?
3273                        } else {
3274                            return Err(format!(
3275                                "run_full: layer {il} is neither linear nor full-attention"
3276                            )
3277                            .into());
3278                        };
3279                        xc = Some(nx);
3280                    }
3281                    e.copy_into(xout, 0, xc.as_ref().unwrap(), t * n_embd)?;
3282                    Ok(())
3283                })?
3284            };
3285            // Undo the net host parity motion of the 3 body runs (each run swaps iff t
3286            // is odd -> 3 runs = net one swap), then restore the device state the
3287            // warmups consumed (walk scope only — layers past hi never executed). The
3288            // launch below then behaves exactly like one run.
3289            if t % 2 == 1 {
3290                for &il in &self.lin {
3291                    if il < lo || il >= hi {
3292                        continue;
3293                    }
3294                    let rl = cache.recur[il].as_mut().unwrap();
3295                    std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
3296                }
3297            }
3298            for (k, &il) in self.lin.iter().enumerate() {
3299                if il < lo || il >= hi {
3300                    continue;
3301                }
3302                let rl = cache.recur[il].as_mut().unwrap();
3303                let (cw, sw) = (self.conv_words, self.ssm_words);
3304                {
3305                    let sv = e.view(&self.save_conv, self.lin.len() * cw);
3306                    let win = sv.slice(k * cw..(k + 1) * cw);
3307                    e.copy_view_into(&mut rl.conv_state, 0, &win, cw)?;
3308                }
3309                {
3310                    let sv = e.view(&self.save_ssm, self.lin.len() * sw);
3311                    let win = sv.slice(k * sw..(k + 1) * sw);
3312                    e.copy_view_into(&mut rl.ssm_state, 0, &win, sw)?;
3313                }
3314            }
3315            if std::env::var("MEMRA_GRAPH_CENSUS").as_deref() == Ok("1")
3316                && let Ok(c) = crate::graph_update::node_census(&graph)
3317            {
3318                eprintln!("[dspark-vg-census] full vt={t} rung={rung} {c:?}");
3319            }
3320            self.full.insert(
3321                key,
3322                DsparkSegGraph {
3323                    graph,
3324                    _keeper: keeper,
3325                },
3326            );
3327        }
3328        self.full[&key].graph.launch()?;
3329        // Host bookkeeping for the replayed body (captured host code does not re-run):
3330        // gdn parity swap per linear layer (t odd), kv len bump per fa layer — scoped
3331        // to the WALK [lo, hi): the cache can carry extra state slots past it (the MTP
3332        // head layer's kv) that the walk never touches.
3333        if t % 2 == 1 {
3334            for &il in &self.lin {
3335                if il < lo || il >= hi {
3336                    continue;
3337                }
3338                let rl = cache.recur[il].as_mut().unwrap();
3339                std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
3340            }
3341        }
3342        for &il in &self.fa {
3343            if il < lo || il >= hi {
3344                continue;
3345            }
3346            cache.kv[il].as_mut().unwrap().len += t;
3347        }
3348        let (_, xout) = self.stage.get(&t).unwrap();
3349        let mut out = e.uninit(t * n_embd)?;
3350        e.copy_into(&mut out, 0, xout, t * n_embd)?;
3351        Ok(out)
3352    }
3353
3354    /// Run layers [start, end) (all linear) as one captured graph at this vt: stage the
3355    /// residual into the bucket's x_in, capture on first encounter (2 executing warmups
3356    /// bracketed by a segment state save/restore), launch, then apply the host parity
3357    /// bookkeeping the captured body would have done. Returns the fresh residual.
3358    #[allow(clippy::too_many_arguments)]
3359    #[allow(clippy::map_entry)] // allow: the init body is fallible (`?`); Entry::or_insert_with cannot propagate errors
3360    fn run_segment(
3361        &mut self,
3362        model: &crate::hybrid::HybridModel,
3363        e: &Engine,
3364        start: usize,
3365        end: usize,
3366        x: &CudaSlice<f32>,
3367        t: usize,
3368        cache: &mut Cache,
3369    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3370        let n_embd = self.n_embd;
3371        debug_assert!(end - start <= self.max_run);
3372        if !self.stage.contains_key(&t) {
3373            let xin = e.uninit(t * n_embd)?;
3374            let xout = e.uninit(t * n_embd)?;
3375            self.stage.insert(t, (xin, xout));
3376        }
3377        // Stage the residual at the bucket's baked input address.
3378        {
3379            let (xin, _) = self.stage.get_mut(&t).unwrap();
3380            e.copy_into(xin, 0, x, t * n_embd)?;
3381        }
3382        let key = (start, t);
3383        if !self.graphs.contains_key(&key) {
3384            // The 2 warmups EXECUTE the segment on live state — save conv + the canonical
3385            // ssm of every segment layer first, restore after, so the graph's first real
3386            // launch starts from the exact pre-round state (bytes gated e2e).
3387            for (k, il) in (start..end).enumerate() {
3388                let rl = cache.recur[il].as_ref().unwrap();
3389                e.copy_into(
3390                    &mut self.save_conv,
3391                    k * self.conv_words,
3392                    &rl.conv_state,
3393                    self.conv_words,
3394                )?;
3395                e.copy_into(
3396                    &mut self.save_ssm,
3397                    k * self.ssm_words,
3398                    &rl.ssm_state,
3399                    self.ssm_words,
3400                )?;
3401            }
3402            let (graph, keeper) = {
3403                let table_all = &self.table_all;
3404                let lin_pos = &self.lin_pos;
3405                let stash_conv = &mut self.stash_conv;
3406                let stash_ssm = &mut self.stash_ssm;
3407                let (xin, xout) = self
3408                    .stage
3409                    .get_mut(&t)
3410                    .map(|(a, b)| (&*a, b))
3411                    .expect("stage bucket created above");
3412                let cache_ref: &mut Cache = cache;
3413                // Slice 4 (fa-execupdate lane): USE_NODE_PRIORITY instead of
3414                // AUTO_FREE_ON_LAUNCH. The slice-3 measured limiter was AUTO_FREE's
3415                // launch-time mem-pool scan — 25.6 us per cuGraphLaunch x 16 segments
3416                // = ~0.41 ms/round, most of the eager-launch savings. The captured
3417                // body's cuMemAllocAsync transients are BALANCED by in-graph frees
3418                // (every transient drops inside the capture region — the generic
3419                // capture path's census precedent, 1589/1589), so AUTO_FREE has
3420                // nothing to reclaim and the graph is legal to instantiate without
3421                // it; PRIORITY is the flag the gemma slotted door ships for exactly
3422                // this reason (both alternatives drop the scan; UPLOAD via
3423                // cuGraphInstantiateWithFlags is WithParams-only and refused).
3424                // MEMRA_DSPARK_VG_AUTOFREE=1 reverts; MEMRA_GRAPH_CENSUS=1 prints
3425                // the node census at capture (the ALLOC==FREE receipt).
3426                let iflag = if std::env::var("MEMRA_DSPARK_VG_AUTOFREE").as_deref() == Ok("1") {
3427                    cudarc::driver::sys::CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH
3428                } else {
3429                    cudarc::driver::sys::CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_USE_NODE_PRIORITY
3430                };
3431                e.capture_graph_retained_flags(iflag, move |e| {
3432                    let mut xc: Option<CudaSlice<f32>> = None;
3433                    for il in start..end {
3434                        let k = lin_pos[&il];
3435                        let xr: &CudaSlice<f32> = xc.as_ref().unwrap_or(xin);
3436                        let nx = model.qwen35_tparallel_linear_layer(
3437                            e,
3438                            il,
3439                            xr,
3440                            t,
3441                            cache_ref,
3442                            None,
3443                            Some((&mut stash_conv[k], &mut stash_ssm[k])),
3444                            Some((table_all, k * 6)),
3445                        )?;
3446                        xc = Some(nx);
3447                    }
3448                    e.copy_into(xout, 0, xc.as_ref().unwrap(), t * n_embd)?;
3449                    Ok(())
3450                })?
3451            };
3452            // Undo the net host parity motion of the 3 body runs (each run swaps iff t
3453            // is odd -> 3 runs = net one swap), then restore the device state the
3454            // warmups consumed. The launch below then behaves exactly like one run.
3455            if t % 2 == 1 {
3456                for il in start..end {
3457                    let rl = cache.recur[il].as_mut().unwrap();
3458                    std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
3459                }
3460            }
3461            for (k, il) in (start..end).enumerate() {
3462                let rl = cache.recur[il].as_mut().unwrap();
3463                let (cw, sw) = (self.conv_words, self.ssm_words);
3464                {
3465                    let sv = e.view(&self.save_conv, self.lin.len() * cw);
3466                    let win = sv.slice(k * cw..(k + 1) * cw);
3467                    e.copy_view_into(&mut rl.conv_state, 0, &win, cw)?;
3468                }
3469                {
3470                    let sv = e.view(&self.save_ssm, self.lin.len() * sw);
3471                    let win = sv.slice(k * sw..(k + 1) * sw);
3472                    e.copy_view_into(&mut rl.ssm_state, 0, &win, sw)?;
3473                }
3474            }
3475            if std::env::var("MEMRA_GRAPH_CENSUS").as_deref() == Ok("1")
3476                && let Ok(c) = crate::graph_update::node_census(&graph)
3477            {
3478                eprintln!("[dspark-vg-census] seg={start}..{end} vt={t} {c:?}");
3479            }
3480            self.graphs.insert(
3481                key,
3482                DsparkSegGraph {
3483                    graph,
3484                    _keeper: keeper,
3485                },
3486            );
3487        }
3488        self.graphs[&key].graph.launch()?;
3489        // Host parity bookkeeping for the replayed body (the captured host swaps do not
3490        // re-run at replay).
3491        if t % 2 == 1 {
3492            for il in start..end {
3493                let rl = cache.recur[il].as_mut().unwrap();
3494                std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
3495            }
3496        }
3497        let (_, xout) = self.stage.get(&t).unwrap();
3498        let mut out = e.uninit(t * n_embd)?;
3499        e.copy_into(&mut out, 0, xout, t * n_embd)?;
3500        Ok(out)
3501    }
3502
3503    /// Pool freeze check (`dspark_vg_cap`): below the ceiling new keys may capture.
3504    fn can_capture(&self) -> bool {
3505        self.graphs.len() + self.full.len() < dspark_vg_cap()
3506    }
3507
3508    /// Round-atomic segment-door readiness: TRUE when this round's walk can ride the
3509    /// per-(segment, vt) graphs without a NEW capture past the pool ceiling — every
3510    /// linear run in [lo, hi) already has its (run_start, t) key, or capture is still
3511    /// allowed. FALSE sends the WHOLE round down the eager cols-ckpt walk: a partial
3512    /// refusal would stash some layers in the ctx slabs and others in the round's cols
3513    /// while one commit reads only one of them.
3514    pub(crate) fn segments_ready(
3515        &self,
3516        model: &crate::hybrid::HybridModel,
3517        lo: usize,
3518        hi: usize,
3519        t: usize,
3520    ) -> bool {
3521        if self.can_capture() {
3522            return true;
3523        }
3524        let mut il = lo;
3525        while il < hi {
3526            if matches!(model.layers[il].mixer, Mixer::Linear(_)) {
3527                let start = il;
3528                while il < hi && matches!(model.layers[il].mixer, Mixer::Linear(_)) {
3529                    il += 1;
3530                }
3531                if !self.graphs.contains_key(&(start, t)) {
3532                    return false;
3533                }
3534            } else {
3535                il += 1;
3536            }
3537        }
3538        true
3539    }
3540
3541    /// Widest verify window this pool was built for. A caller whose round exceeds it must
3542    /// take the eager walk: the stash slabs hold `t_capacity() - 1` column rows, and slicing
3543    /// past them is a panic rather than a refusal.
3544    pub(crate) fn t_capacity(&self) -> usize {
3545        self.t_cap
3546    }
3547
3548    /// Slab row (conv, ssm) device pointers + lengths for the commit restore of column
3549    /// `row` (0-based) of layer `il`. None for non-linear layers.
3550    pub(crate) fn slab_row(
3551        &self,
3552        e: &Engine,
3553        il: usize,
3554        row: usize,
3555    ) -> Option<(u64, u64, usize, usize)> {
3556        use cudarc::driver::DevicePtr;
3557        let k = *self.lin_pos.get(&il)?;
3558        let s = &e.gpu.stream();
3559        let (pc, _g0) = self.stash_conv[k].device_ptr(s);
3560        let (ps, _g1) = self.stash_ssm[k].device_ptr(s);
3561        Some((
3562            pc + (row * self.conv_words * 4) as u64,
3563            ps + (row * self.ssm_words * 4) as u64,
3564            self.conv_words,
3565            self.ssm_words,
3566        ))
3567    }
3568}
3569
3570impl VerifyCkpt {
3571    fn new(n_layer: usize) -> Self {
3572        VerifyCkpt {
3573            gdn: (0..n_layer).map(|_| None).collect(),
3574            cols: (0..n_layer).map(|_| None).collect(),
3575        }
3576    }
3577}
3578
3579/// The stage-0/TX half of one PP verify. The boundary slot is the ownership token: stage 1
3580/// consumes exactly the slot selected by `tx()` / `tx_pipelined()`, never a slot inferred from
3581/// a logical round number.
3582struct VerifyBoundaryTicket {
3583    rt: &'static crate::pp::PpNRt,
3584    caller_stream: std::sync::Arc<cudarc::driver::CudaStream>,
3585    slot: usize,
3586    pos0: usize,
3587    t: usize,
3588    payload: usize,
3589    n_st: usize,
3590    pipelined: bool,
3591    pp_anatomy: bool,
3592    pp_started: std::time::Instant,
3593    reverse_ms: f64,
3594    stage0_ms: f64,
3595    tx_ms: f64,
3596    trace: Option<SpecPipeTraceCtx>,
3597    _walk_owner: crate::pp::PpWalkLease,
3598}
3599
3600/// Explicit OPTIPIPE diagnostic control. Forced modes are set only by `optipipe-gate`; the
3601/// increment-2 controller can also be armed by the server's fresh-process research door.
3602#[derive(Clone, Copy, Debug, PartialEq, Eq)]
3603pub enum OptiForkGateMode {
3604    Disabled,
3605    Hit,
3606    Miss,
3607    Alternate,
3608    Abort,
3609    Controller,
3610}
3611
3612static OPTI_FORK_GATE_MODE: std::sync::atomic::AtomicU8 = std::sync::atomic::AtomicU8::new(0);
3613static OPTI_CONTROLLER_THRESHOLD: std::sync::atomic::AtomicU32 =
3614    std::sync::atomic::AtomicU32::new(0);
3615static OPTI_FORK_ATTEMPTS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
3616static OPTI_FORK_HITS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
3617static OPTI_FORK_MISSES: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
3618static OPTI_FORK_ABORT_DRAINS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
3619static OPTI_FORK_REFUSALS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
3620static OPTI_GATE_CHECKS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
3621static OPTI_GATE_ADMITS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
3622static OPTI_GATE_REJECTS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
3623static OPTI_RECONCILES: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
3624static OPTI_WASTED_DRAFT_TOKENS: std::sync::atomic::AtomicU64 =
3625    std::sync::atomic::AtomicU64::new(0);
3626static OPTI_SHADOW_DRAFT_TOKENS: std::sync::atomic::AtomicU64 =
3627    std::sync::atomic::AtomicU64::new(0);
3628static OPTI_BREAKER_TRIPS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
3629
3630impl OptiForkGateMode {
3631    fn code(self) -> u8 {
3632        match self {
3633            Self::Disabled => 0,
3634            Self::Hit => 1,
3635            Self::Miss => 2,
3636            Self::Alternate => 3,
3637            Self::Abort => 4,
3638            Self::Controller => 5,
3639        }
3640    }
3641
3642    fn configured() -> Self {
3643        match OPTI_FORK_GATE_MODE.load(std::sync::atomic::Ordering::Relaxed) {
3644            1 => Self::Hit,
3645            2 => Self::Miss,
3646            3 => Self::Alternate,
3647            4 => Self::Abort,
3648            5 => Self::Controller,
3649            _ => Self::Disabled,
3650        }
3651    }
3652
3653    fn action(self, generation: u64) -> OptiForkAction {
3654        match self {
3655            Self::Hit => OptiForkAction::Hit,
3656            Self::Miss => OptiForkAction::Miss,
3657            Self::Alternate if generation & 1 == 0 => OptiForkAction::Hit,
3658            Self::Alternate => OptiForkAction::Miss,
3659            Self::Abort => OptiForkAction::Abort,
3660            Self::Disabled | Self::Controller => {
3661                unreachable!("non-forced mode cannot choose a forced fork action")
3662            }
3663        }
3664    }
3665
3666    fn is_forced(self) -> bool {
3667        matches!(self, Self::Hit | Self::Miss | Self::Alternate | Self::Abort)
3668    }
3669}
3670
3671/// Arm or disarm the forced harness. Serving uses only `set_optipipe_controller_threshold`.
3672pub fn set_optipipe_gate_mode(mode: OptiForkGateMode) {
3673    OPTI_FORK_GATE_MODE.store(mode.code(), std::sync::atomic::Ordering::Relaxed);
3674}
3675
3676/// Arm the increment-2 diagnostic controller. The threshold applies to the uncalibrated
3677/// two-token draft-probability product. Serving can call this only through its explicit
3678/// fresh-process research door; the absent-door default remains byte-for-byte disabled.
3679pub fn set_optipipe_controller_threshold(threshold: f32) {
3680    assert!(threshold.is_finite() && (0.0..=1.0).contains(&threshold));
3681    OPTI_CONTROLLER_THRESHOLD.store(threshold.to_bits(), std::sync::atomic::Ordering::Relaxed);
3682    set_optipipe_gate_mode(OptiForkGateMode::Controller);
3683}
3684
3685#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
3686pub struct OptiForkGateStats {
3687    pub attempts: u64,
3688    pub hits: u64,
3689    pub misses: u64,
3690    pub abort_drains: u64,
3691    pub refusals: u64,
3692    pub gate_checks: u64,
3693    pub gate_admits: u64,
3694    pub gate_rejects: u64,
3695    pub reconciles: u64,
3696    pub wasted_draft_tokens: u64,
3697    pub shadow_draft_tokens: u64,
3698    pub breaker_trips: u64,
3699}
3700
3701#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
3702pub struct OptiForkStateIdentity {
3703    pub trunk_kv_bytes: usize,
3704    pub recurrent_bytes: usize,
3705    pub scratch_kv_bytes: usize,
3706    pub hidden_bytes: usize,
3707}
3708
3709pub fn reset_optipipe_gate_stats() {
3710    for counter in [
3711        &OPTI_FORK_ATTEMPTS,
3712        &OPTI_FORK_HITS,
3713        &OPTI_FORK_MISSES,
3714        &OPTI_FORK_ABORT_DRAINS,
3715        &OPTI_FORK_REFUSALS,
3716        &OPTI_GATE_CHECKS,
3717        &OPTI_GATE_ADMITS,
3718        &OPTI_GATE_REJECTS,
3719        &OPTI_RECONCILES,
3720        &OPTI_WASTED_DRAFT_TOKENS,
3721        &OPTI_SHADOW_DRAFT_TOKENS,
3722        &OPTI_BREAKER_TRIPS,
3723    ] {
3724        counter.store(0, std::sync::atomic::Ordering::Relaxed);
3725    }
3726}
3727
3728pub fn optipipe_gate_stats() -> OptiForkGateStats {
3729    let load = |v: &std::sync::atomic::AtomicU64| v.load(std::sync::atomic::Ordering::Relaxed);
3730    OptiForkGateStats {
3731        attempts: load(&OPTI_FORK_ATTEMPTS),
3732        hits: load(&OPTI_FORK_HITS),
3733        misses: load(&OPTI_FORK_MISSES),
3734        abort_drains: load(&OPTI_FORK_ABORT_DRAINS),
3735        refusals: load(&OPTI_FORK_REFUSALS),
3736        gate_checks: load(&OPTI_GATE_CHECKS),
3737        gate_admits: load(&OPTI_GATE_ADMITS),
3738        gate_rejects: load(&OPTI_GATE_REJECTS),
3739        reconciles: load(&OPTI_RECONCILES),
3740        wasted_draft_tokens: load(&OPTI_WASTED_DRAFT_TOKENS),
3741        shadow_draft_tokens: load(&OPTI_SHADOW_DRAFT_TOKENS),
3742        breaker_trips: load(&OPTI_BREAKER_TRIPS),
3743    }
3744}
3745
3746#[derive(Clone, Copy, Debug)]
3747struct OptiControllerPolicy {
3748    threshold: f32,
3749    consecutive_misses: u8,
3750    breaker_tripped: bool,
3751}
3752
3753impl OptiControllerPolicy {
3754    fn configured() -> Self {
3755        Self {
3756            threshold: f32::from_bits(
3757                OPTI_CONTROLLER_THRESHOLD.load(std::sync::atomic::Ordering::Relaxed),
3758            ),
3759            consecutive_misses: 0,
3760            breaker_tripped: false,
3761        }
3762    }
3763
3764    fn admit(&self, q_proxy: f32) -> bool {
3765        q_proxy.is_finite()
3766            && (0.0..=1.0).contains(&q_proxy)
3767            && (self.threshold == 0.0 || (!self.breaker_tripped && q_proxy >= self.threshold))
3768    }
3769
3770    /// Returns true exactly when this resolution newly trips the three-miss breaker.
3771    fn resolve(&mut self, hit: bool) -> bool {
3772        // q*=0 is the lane's explicit unconditional measurement arm. Its purpose is to price
3773        // every optimistic opportunity, so the safety breaker is measured separately and must
3774        // not silently turn this arm into "three attempts then serial".
3775        if self.threshold == 0.0 {
3776            self.consecutive_misses = 0;
3777            return false;
3778        }
3779        if hit {
3780            self.consecutive_misses = 0;
3781            return false;
3782        }
3783        self.consecutive_misses = self.consecutive_misses.saturating_add(1);
3784        if !self.breaker_tripped && self.consecutive_misses >= 3 {
3785            self.breaker_tripped = true;
3786            return true;
3787        }
3788        false
3789    }
3790}
3791
3792#[derive(Clone, Copy, Debug, PartialEq, Eq)]
3793enum OptiForkAction {
3794    Hit,
3795    Miss,
3796    Abort,
3797}
3798
3799#[derive(Clone, Copy, Debug, PartialEq, Eq)]
3800struct OptiForkGeneration {
3801    id: u64,
3802    slot: usize,
3803}
3804
3805#[derive(Default)]
3806struct OptiForkGenerationTracker {
3807    next: u64,
3808    live: [Option<u64>; 2],
3809}
3810
3811impl OptiForkGenerationTracker {
3812    fn reserve(&mut self) -> Result<OptiForkGeneration, Box<dyn std::error::Error>> {
3813        let generation = OptiForkGeneration {
3814            id: self.next,
3815            slot: (self.next & 1) as usize,
3816        };
3817        if let Some(live) = self.live[generation.slot] {
3818            return Err(format!(
3819                "optipipe snapshot slot {} still owns generation {live}; refusing to overwrite it",
3820                generation.slot,
3821            )
3822            .into());
3823        }
3824        self.next += 1;
3825        self.live[generation.slot] = Some(generation.id);
3826        Ok(generation)
3827    }
3828
3829    fn retire(&mut self, generation: OptiForkGeneration) -> Result<(), Box<dyn std::error::Error>> {
3830        match self.live[generation.slot] {
3831            Some(id) if id == generation.id => {
3832                self.live[generation.slot] = None;
3833                Ok(())
3834            }
3835            other => Err(format!(
3836                "optipipe generation teardown mismatch: ticket={} slot={} live={other:?}",
3837                generation.id, generation.slot,
3838            )
3839            .into()),
3840        }
3841    }
3842}
3843
3844struct OptiForkSeedGeneration {
3845    h_seed: CudaSlice<f32>,
3846    fill_prev: CudaSlice<f32>,
3847    scratch_len: usize,
3848}
3849
3850/// Allocate or refresh one full checkpoint through the engine that owns each PP stage. The
3851/// generic cache helper accepts one device and therefore cannot copy GDN state split across
3852/// devices. KV lengths and position stay host metadata; only recurrent buffers need stage-local
3853/// device ownership.
3854fn opti_snapshot_stage_owned(
3855    e: &Engine,
3856    cache: &Cache,
3857    rt: &'static crate::pp::PpNRt,
3858    fence: &[usize],
3859) -> Result<crate::cache::CacheSnapshot, Box<dyn std::error::Error>> {
3860    let n = cache.kv.len();
3861    let mut snapshot = crate::cache::CacheSnapshot {
3862        kv_len: vec![None; n],
3863        tp_kv_len: vec![None; n],
3864        conv: (0..n).map(|_| None).collect(),
3865        ssm: (0..n).map(|_| None).collect(),
3866        pos: cache.pos,
3867    };
3868    opti_snapshot_stage_owned_into(e, cache, rt, fence, &mut snapshot)?;
3869    Ok(snapshot)
3870}
3871
3872fn opti_snapshot_stage_owned_into(
3873    e: &Engine,
3874    cache: &Cache,
3875    rt: &'static crate::pp::PpNRt,
3876    fence: &[usize],
3877    snapshot: &mut crate::cache::CacheSnapshot,
3878) -> Result<(), Box<dyn std::error::Error>> {
3879    if fence.len() != rt.n_stages() + 1
3880        || snapshot.kv_len.len() != cache.kv.len()
3881        || snapshot.tp_kv_len.len() != cache.tp_kv.len()
3882    {
3883        return Err("optipipe stage-owned snapshot shape mismatch".into());
3884    }
3885    for stage in 0..rt.n_stages() {
3886        opti_snapshot_one_stage_owned_into(e, cache, rt, fence, stage, snapshot)?;
3887    }
3888    snapshot.pos = cache.pos;
3889    Ok(())
3890}
3891
3892/// Refresh one PP stage of a checkpoint. Increment 2 uses this split form so stage 0's
3893/// optimistic post-N state is captured before N+1 stage 0 is queued, while stage 1's matching
3894/// post-N state is captured only after N stage 1 is enqueued. Calling the all-stage helper at
3895/// either point would capture one side of the fork at the wrong generation.
3896fn opti_snapshot_one_stage_owned_into(
3897    e: &Engine,
3898    cache: &Cache,
3899    rt: &'static crate::pp::PpNRt,
3900    fence: &[usize],
3901    stage: usize,
3902    snapshot: &mut crate::cache::CacheSnapshot,
3903) -> Result<(), Box<dyn std::error::Error>> {
3904    if fence.len() != rt.n_stages() + 1
3905        || snapshot.kv_len.len() != cache.kv.len()
3906        || snapshot.tp_kv_len.len() != cache.tp_kv.len()
3907        || stage >= rt.n_stages()
3908    {
3909        return Err("optipipe single-stage snapshot shape mismatch".into());
3910    }
3911    let _scope = rt.enter(stage);
3912    let owner = rt.engine(stage, e);
3913    for il in fence[stage]..fence[stage + 1] {
3914        snapshot.kv_len[il] = cache.kv[il].as_ref().map(|kv| kv.len);
3915        snapshot.tp_kv_len[il] = cache.tp_kv[il]
3916            .as_ref()
3917            .map(crate::tp::ResidentTpKvCache::committed_len);
3918        match &cache.recur[il] {
3919            Some(recur) => {
3920                match snapshot.conv[il].as_mut() {
3921                    Some(dst) => {
3922                        owner.copy_into(dst, 0, &recur.conv_state, recur.conv_state.len())?
3923                    }
3924                    None => snapshot.conv[il] = Some(owner.clone_dtod(&recur.conv_state)?),
3925                }
3926                match snapshot.ssm[il].as_mut() {
3927                    Some(dst) => {
3928                        owner.copy_into(dst, 0, &recur.ssm_state, recur.ssm_state.len())?
3929                    }
3930                    None => snapshot.ssm[il] = Some(owner.clone_dtod(&recur.ssm_state)?),
3931                }
3932            }
3933            None if snapshot.conv[il].is_some() || snapshot.ssm[il].is_some() => {
3934                return Err(
3935                    format!("optipipe stage-owned snapshot layer {il} changed shape").into(),
3936                );
3937            }
3938            None => {}
3939        }
3940    }
3941    snapshot.pos = cache.pos;
3942    Ok(())
3943}
3944
3945/// Increment-1 persistent fork state. Exactly two snapshot/seed slots alternate; a live ticket
3946/// names its generation and keeps teardown fail-closed. Only stage 0 is allowed to mutate before
3947/// resolve, so the reconcile tables and conditional restores are stage-local.
3948struct OptiForkState {
3949    mode: OptiForkGateMode,
3950    controller: Option<OptiControllerPolicy>,
3951    generations: OptiForkGenerationTracker,
3952    active_snapshot_slot: usize,
3953    alternate_snapshot: crate::cache::CacheSnapshot,
3954    seeds: [OptiForkSeedGeneration; 2],
3955    rt: &'static crate::pp::PpNRt,
3956    fence: [usize; 3],
3957    split: usize,
3958    len_ptrs: CudaSlice<u64>,
3959    saved_lens: CudaSlice<i32>,
3960    forced_acc: CudaSlice<u32>,
3961    valid: CudaSlice<u32>,
3962    stage0_stream: std::sync::Arc<cudarc::driver::CudaStream>,
3963    logical_payload_bytes: [usize; 2],
3964}
3965
3966struct OptiForkTicket {
3967    generation: OptiForkGeneration,
3968    boundary: Option<VerifyBoundaryTicket>,
3969    drain: std::sync::Arc<cudarc::driver::CudaStream>,
3970    settled: bool,
3971}
3972
3973struct OptiControllerTicket {
3974    generation: OptiForkGeneration,
3975    boundary: Option<VerifyBoundaryTicket>,
3976    ckpt: Option<VerifyCkpt>,
3977    verify_tokens: [u32; 2],
3978    draft_prob: f32,
3979    eager_seed: Option<CudaSlice<f32>>,
3980    q_proxy: f32,
3981    scratch_len: usize,
3982    issued_at: std::time::Instant,
3983    drain: std::sync::Arc<cudarc::driver::CudaStream>,
3984    settled: bool,
3985}
3986
3987struct OptiControllerPrepared {
3988    verify_tokens: [u32; 2],
3989    draft_prob: f32,
3990    eager_seed: Option<CudaSlice<f32>>,
3991    q_proxy: f32,
3992    scratch_len: usize,
3993}
3994
3995impl OptiControllerTicket {
3996    fn take_boundary(&mut self) -> VerifyBoundaryTicket {
3997        self.boundary
3998            .take()
3999            .expect("controller boundary ticket already consumed")
4000    }
4001
4002    fn take_ckpt(&mut self) -> VerifyCkpt {
4003        self.ckpt
4004            .take()
4005            .expect("controller verify checkpoint already consumed")
4006    }
4007
4008    fn take_eager_seed(&mut self) -> Option<CudaSlice<f32>> {
4009        self.eager_seed.take()
4010    }
4011
4012    fn settle(&mut self) {
4013        self.settled = true;
4014    }
4015}
4016
4017impl Drop for OptiControllerTicket {
4018    fn drop(&mut self) {
4019        if !self.settled {
4020            let _ = self.drain.synchronize();
4021            OPTI_FORK_ABORT_DRAINS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
4022        }
4023    }
4024}
4025
4026impl OptiForkTicket {
4027    fn take_boundary(&mut self) -> VerifyBoundaryTicket {
4028        self.boundary
4029            .take()
4030            .expect("fork ticket boundary already consumed")
4031    }
4032
4033    fn settle(&mut self) {
4034        self.settled = true;
4035    }
4036}
4037
4038impl Drop for OptiForkTicket {
4039    fn drop(&mut self) {
4040        if !self.settled {
4041            let _ = self.drain.synchronize();
4042            OPTI_FORK_ABORT_DRAINS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
4043        }
4044    }
4045}
4046
4047impl OptiForkState {
4048    #[allow(clippy::too_many_arguments)]
4049    fn new(
4050        e: &Engine,
4051        cache: &Cache,
4052        mode: OptiForkGateMode,
4053        alternate_snapshot: crate::cache::CacheSnapshot,
4054        h_seed: &CudaSlice<f32>,
4055        fill_prev: &CudaSlice<f32>,
4056        rt: &'static crate::pp::PpNRt,
4057        split: usize,
4058        n_layer: usize,
4059    ) -> Result<Self, Box<dyn std::error::Error>> {
4060        let fence = [0, split, n_layer];
4061        let mut logical_payload_bytes = [0usize; 2];
4062        for stage in 0..2 {
4063            for il in fence[stage]..fence[stage + 1] {
4064                logical_payload_bytes[stage] += alternate_snapshot.conv[il]
4065                    .as_ref()
4066                    .map_or(0, |v| v.len() * std::mem::size_of::<f32>());
4067                logical_payload_bytes[stage] += alternate_snapshot.ssm[il]
4068                    .as_ref()
4069                    .map_or(0, |v| v.len() * std::mem::size_of::<f32>());
4070            }
4071        }
4072        let seeds = [
4073            OptiForkSeedGeneration {
4074                h_seed: e.clone_dtod(h_seed)?,
4075                fill_prev: e.clone_dtod(fill_prev)?,
4076                scratch_len: 0,
4077            },
4078            OptiForkSeedGeneration {
4079                h_seed: e.clone_dtod(h_seed)?,
4080                fill_prev: e.clone_dtod(fill_prev)?,
4081                scratch_len: 0,
4082            },
4083        ];
4084        let (len_ptrs, saved_lens, forced_acc, valid, stage0_stream) = {
4085            let _stage = rt.enter(0);
4086            let e0 = rt.engine(0, e);
4087            (
4088                crate::round_stream::kv_len_ptr_table_range(e0, cache, 0..split, None)?,
4089                e0.htod_i32(&vec![0; split])?,
4090                e0.alloc_u32_zeroed(2)?,
4091                e0.alloc_u32_zeroed(1)?,
4092                e0.stream(),
4093            )
4094        };
4095        logical_payload_bytes[0] += seeds
4096            .iter()
4097            .map(|seed| (seed.h_seed.len() + seed.fill_prev.len()) * std::mem::size_of::<f32>())
4098            .sum::<usize>();
4099        logical_payload_bytes[0] += len_ptrs.len() * std::mem::size_of::<u64>()
4100            + saved_lens.len() * std::mem::size_of::<i32>()
4101            + forced_acc.len() * std::mem::size_of::<u32>()
4102            + valid.len() * std::mem::size_of::<u32>();
4103        Ok(Self {
4104            mode,
4105            controller: (mode == OptiForkGateMode::Controller)
4106                .then(OptiControllerPolicy::configured),
4107            generations: OptiForkGenerationTracker::default(),
4108            active_snapshot_slot: 0,
4109            alternate_snapshot,
4110            seeds,
4111            rt,
4112            fence,
4113            split,
4114            len_ptrs,
4115            saved_lens,
4116            forced_acc,
4117            valid,
4118            stage0_stream,
4119            logical_payload_bytes,
4120        })
4121    }
4122
4123    fn reserve(
4124        &mut self,
4125        current_snapshot: &mut crate::cache::CacheSnapshot,
4126    ) -> Result<OptiForkGeneration, Box<dyn std::error::Error>> {
4127        let generation = self.generations.reserve()?;
4128        if generation.slot != self.active_snapshot_slot {
4129            std::mem::swap(current_snapshot, &mut self.alternate_snapshot);
4130            self.active_snapshot_slot = generation.slot;
4131        }
4132        Ok(generation)
4133    }
4134
4135    fn capture_seed(
4136        &mut self,
4137        e: &Engine,
4138        generation: OptiForkGeneration,
4139        h_seed: &CudaSlice<f32>,
4140        fill_prev: &CudaSlice<f32>,
4141        scratch_len: usize,
4142    ) -> Result<(), Box<dyn std::error::Error>> {
4143        let seed = &mut self.seeds[generation.slot];
4144        e.copy_into(&mut seed.h_seed, 0, h_seed, h_seed.len())?;
4145        e.copy_into(&mut seed.fill_prev, 0, fill_prev, fill_prev.len())?;
4146        seed.scratch_len = scratch_len;
4147        Ok(())
4148    }
4149
4150    fn ticket(
4151        &self,
4152        generation: OptiForkGeneration,
4153        boundary: VerifyBoundaryTicket,
4154    ) -> OptiForkTicket {
4155        OptiForkTicket {
4156            generation,
4157            boundary: Some(boundary),
4158            drain: self.stage0_stream.clone(),
4159            settled: false,
4160        }
4161    }
4162
4163    #[allow(clippy::too_many_arguments)]
4164    fn controller_ticket(
4165        &self,
4166        generation: OptiForkGeneration,
4167        boundary: VerifyBoundaryTicket,
4168        ckpt: VerifyCkpt,
4169        verify_tokens: [u32; 2],
4170        draft_prob: f32,
4171        eager_seed: Option<CudaSlice<f32>>,
4172        q_proxy: f32,
4173        scratch_len: usize,
4174    ) -> OptiControllerTicket {
4175        OptiControllerTicket {
4176            generation,
4177            boundary: Some(boundary),
4178            ckpt: Some(ckpt),
4179            verify_tokens,
4180            draft_prob,
4181            eager_seed,
4182            q_proxy,
4183            scratch_len,
4184            issued_at: std::time::Instant::now(),
4185            drain: self.stage0_stream.clone(),
4186            settled: false,
4187        }
4188    }
4189
4190    fn reserve_successor(&mut self) -> Result<OptiForkGeneration, Box<dyn std::error::Error>> {
4191        self.generations.reserve()
4192    }
4193
4194    fn successor_snapshot_mut(&mut self) -> &mut crate::cache::CacheSnapshot {
4195        &mut self.alternate_snapshot
4196    }
4197
4198    fn promote_successor_snapshot(
4199        &mut self,
4200        current_snapshot: &mut crate::cache::CacheSnapshot,
4201        generation: OptiForkGeneration,
4202    ) {
4203        std::mem::swap(current_snapshot, &mut self.alternate_snapshot);
4204        self.active_snapshot_slot = generation.slot;
4205    }
4206
4207    fn queue_actual_reconcile(
4208        &mut self,
4209        e: &Engine,
4210        snapshot: &crate::cache::CacheSnapshot,
4211        acc: &CudaSlice<u32>,
4212        optimistic_pending: u32,
4213        base: usize,
4214    ) -> Result<(), Box<dyn std::error::Error>> {
4215        let saved: Vec<i32> = (0..self.split)
4216            .map(|il| snapshot.kv_len[il].map(|v| v as i32).unwrap_or(0))
4217            .collect();
4218        // Serving keeps the caller/accept walk on the head (stage-1) device. Record the accept
4219        // decision point there and append a wait to stage 0 after its optimistic successor/TX;
4220        // the validity/reconcile kernels must never peer-read acc before it is written. The
4221        // increment-1 harness uses primary stage 0, where stream order already provides this.
4222        if self.rt.engine(0, e).ctx().ordinal() != e.ctx().ordinal() {
4223            self.rt.fence_stages_behind(&e.stream())?;
4224        }
4225        let _stage = self.rt.enter(0);
4226        let e0 = self.rt.engine(0, e);
4227        e0.htod_i32_into(&mut self.saved_lens, &saved)?;
4228        e0.spec_fork_valid(acc, optimistic_pending, &mut self.valid)?;
4229        e0.spec_fork_reconcile_kv(
4230            &self.len_ptrs,
4231            &self.saved_lens,
4232            acc,
4233            &self.valid,
4234            base,
4235            self.split,
4236        )
4237    }
4238
4239    fn finish_actual_reconcile(
4240        &mut self,
4241        e: &Engine,
4242        cache: &mut Cache,
4243        snapshot: &crate::cache::CacheSnapshot,
4244        n_acc: usize,
4245        base: usize,
4246        hit: bool,
4247    ) -> Result<(), Box<dyn std::error::Error>> {
4248        if hit {
4249            return Ok(());
4250        }
4251        let len_delta = base + n_acc;
4252        for il in 0..self.split {
4253            if let (Some(kv), Some(saved)) = (cache.kv[il].as_mut(), snapshot.kv_len[il]) {
4254                kv.len = saved + len_delta;
4255            }
4256        }
4257        {
4258            let _stage = self.rt.enter(1);
4259            let e1 = self.rt.engine(1, e);
4260            for il in self.split..self.fence[2] {
4261                if let (Some(kv), Some(saved)) = (cache.kv[il].as_mut(), snapshot.kv_len[il]) {
4262                    kv.len = saved + len_delta;
4263                    e1.set_i32_one(&mut kv.len_d, kv.len as i32)?;
4264                }
4265            }
4266        }
4267        self.rt.publish_to(0, &e.stream())?;
4268        Ok(())
4269    }
4270
4271    fn cancel_controller_ticket(
4272        &mut self,
4273        e: &Engine,
4274        cache: &mut Cache,
4275        scratch: &mut MtpScratch,
4276        snapshot: &crate::cache::CacheSnapshot,
4277        ticket: &mut OptiControllerTicket,
4278    ) -> Result<(), Box<dyn std::error::Error>> {
4279        {
4280            let _stage = self.rt.enter(0);
4281            let e0 = self.rt.engine(0, e);
4282            for il in 0..self.split {
4283                if let (Some(kv), Some(saved)) = (cache.kv[il].as_mut(), snapshot.kv_len[il]) {
4284                    kv.len = saved;
4285                    e0.set_i32_one(&mut kv.len_d, saved as i32)?;
4286                }
4287            }
4288        }
4289        scratch.set_len(e, snapshot.pos)?;
4290        ticket.settle();
4291        self.generations.retire(ticket.generation)?;
4292        OPTI_FORK_ABORT_DRAINS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
4293        OPTI_WASTED_DRAFT_TOKENS.fetch_add(2, std::sync::atomic::Ordering::Relaxed);
4294        eprintln!(
4295            "[opti-controller] tail-drain generation={} slot={}",
4296            ticket.generation.id, ticket.generation.slot,
4297        );
4298        Ok(())
4299    }
4300
4301    #[allow(clippy::too_many_arguments)]
4302    fn reconcile(
4303        &mut self,
4304        e: &Engine,
4305        cache: &mut Cache,
4306        scratch: &mut MtpScratch,
4307        snapshot: &crate::cache::CacheSnapshot,
4308        h_seed: &mut CudaSlice<f32>,
4309        fill_prev: &mut CudaSlice<f32>,
4310        generation: OptiForkGeneration,
4311        action: OptiForkAction,
4312        optimistic_pending: u32,
4313    ) -> Result<(), Box<dyn std::error::Error>> {
4314        debug_assert!(action != OptiForkAction::Abort);
4315        let miss_started = std::time::Instant::now();
4316        let keep = action == OptiForkAction::Hit;
4317        let saved: Vec<i32> = (0..self.split)
4318            .map(|il| snapshot.kv_len[il].map(|v| v as i32).unwrap_or(0))
4319            .collect();
4320        let seed = &self.seeds[generation.slot];
4321        {
4322            let _stage = self.rt.enter(0);
4323            let e0 = self.rt.engine(0, e);
4324            e0.htod_i32_into(&mut self.saved_lens, &saved)?;
4325            let forced = if keep {
4326                [1u32, optimistic_pending]
4327            } else {
4328                [0u32, optimistic_pending]
4329            };
4330            e0.htod_u32_into(&mut self.forced_acc, &forced)?;
4331            e0.spec_fork_valid(&self.forced_acc, optimistic_pending, &mut self.valid)?;
4332            e0.spec_fork_reconcile_kv(
4333                &self.len_ptrs,
4334                &self.saved_lens,
4335                &self.forced_acc,
4336                &self.valid,
4337                0,
4338                self.split,
4339            )?;
4340            for il in 0..self.split {
4341                if let Some(recur) = cache.recur[il].as_mut() {
4342                    let conv = snapshot.conv[il]
4343                        .as_ref()
4344                        .ok_or("optipipe stage0 snapshot missing conv state")?;
4345                    let ssm = snapshot.ssm[il]
4346                        .as_ref()
4347                        .ok_or("optipipe stage0 snapshot missing ssm state")?;
4348                    e0.spec_fork_restore_f32(conv, &mut recur.conv_state, &self.valid)?;
4349                    e0.spec_fork_restore_f32(ssm, &mut recur.ssm_state, &self.valid)?;
4350                }
4351            }
4352            e0.spec_fork_restore_f32(&seed.h_seed, h_seed, &self.valid)?;
4353            e0.spec_fork_restore_f32(&seed.fill_prev, fill_prev, &self.valid)?;
4354        }
4355
4356        if keep {
4357            OPTI_FORK_HITS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
4358            return Ok(());
4359        }
4360
4361        for il in 0..self.split {
4362            if let (Some(kv), Some(saved)) = (cache.kv[il].as_mut(), snapshot.kv_len[il]) {
4363                kv.len = saved;
4364            }
4365        }
4366        scratch.set_len(e, seed.scratch_len)?;
4367        // Targeted E_restart: publish only stage 0's reconcile to the caller, then bound the
4368        // forced diagnostic so the retained number is the actual miss cost, not enqueue time.
4369        let caller = e.stream();
4370        self.rt.publish_to(0, &caller)?;
4371        caller.synchronize()?;
4372        let miss_ms = miss_started.elapsed().as_secs_f64() * 1e3;
4373        eprintln!(
4374            "[opti-fork-reconcile] generation={} slot={} miss_ms={miss_ms:.3}",
4375            generation.id, generation.slot,
4376        );
4377        OPTI_FORK_MISSES.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
4378        Ok(())
4379    }
4380
4381    fn retire(&mut self, generation: OptiForkGeneration) -> Result<(), Box<dyn std::error::Error>> {
4382        self.generations.retire(generation)
4383    }
4384}
4385
4386/// MEMRA_SPEC_ROUND_PROF counters: whole-round wall, so the round can be weighed against the
4387/// draft-step ([spec-anatomy]) and verify-walk ([tcol-prof]) splits we already print.
4388static ROUND_PROF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
4389static ROUND_MS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
4390static ROUND_N: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
4391
4392fn validate_tp_kv_snapshot_shape(
4393    tp_kv: &[Option<crate::tp::ResidentTpKvCache>],
4394    saved_lens: &[Option<usize>],
4395) -> Result<(), Box<dyn std::error::Error>> {
4396    if tp_kv.len() != saved_lens.len() {
4397        return Err("spec TP KV snapshot shape mismatch".into());
4398    }
4399    for (layer, (cache, saved)) in tp_kv.iter().zip(saved_lens).enumerate() {
4400        if cache.is_some() != saved.is_some() {
4401            return Err(
4402                format!("spec TP KV layer {layer} changed shape since its snapshot").into(),
4403            );
4404        }
4405    }
4406    Ok(())
4407}
4408
4409impl HybridModel {
4410    /// memra#128: the canonical-to-rank byte copy that `5e0fffb97` added to
4411    /// `restore_step_tp_kv_verified_prefix`. OFF by default (written decision, docs/FLAGS.md
4412    /// `MEMRA_STEP_TP_KV_RESTORE`): on step37 NVFP4 TP2 the only shapes that pass the
4413    /// production acceptance gate are the engine before the copy (arm A) and the copy skipped
4414    /// on every step-TP layer (arm F, byte-identical answers to A); the copy on any layer
4415    /// spliced answers or shifted decode (darklanes research/memra128-bisect-20260903).
4416    /// `1` re-enables the copy for ordinary-commit layers; on-device-written layers
4417    /// (`rows_external`) are skipped either way, their canonical bytes are stale.
4418    fn step_tp_kv_restore_copy_on() -> bool {
4419        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
4420        *ON.get_or_init(|| std::env::var("MEMRA_STEP_TP_KV_RESTORE").ok().as_deref() == Some("1"))
4421    }
4422
4423    fn restore_step_tp_kv_verified_prefix(
4424        &self,
4425        e: &Engine,
4426        cache: &mut Cache,
4427        snap: &crate::cache::CacheSnapshot,
4428        accepted: usize,
4429        // memra#128: what an externally-written (dcw / fa2) layer needs from this call.
4430        // PARTIAL accept (commit_verified_prefix): 5e0fffb97 replaced the standalone
4431        // rewind_tp_kv_verified_prefix with this restore, so the length shrink to
4432        // saved+accepted must still happen here - without it E ran with distributed=259
4433        // against local=257. FULL accept: before 5e0fffb97 nothing touched the
4434        // distributed length there and it was right (arm A passed); rewinding to
4435        // saved+t_v shrinks it by one and the next verify's SWA ring view falls off the
4436        // end ("view [5148,5152) is outside resident [0,5151)", arm E2).
4437        rewind_external: bool,
4438    ) -> Result<(), Box<dyn std::error::Error>> {
4439        validate_tp_kv_snapshot_shape(&cache.tp_kv, &snap.tp_kv_len)?;
4440        e.stream().synchronize()?;
4441        {
4442            let (local_layers, distributed_layers) = (&cache.kv, &mut cache.tp_kv);
4443            let stream = e.gpu.stream();
4444            let mut runtime: Option<std::sync::Arc<crate::tp::TpE4m3HostBounce>> = None;
4445            let mut uniform_runtime = true;
4446            let mut batch = Vec::new();
4447            for (il, (distributed_slot, local_slot)) in distributed_layers
4448                .iter_mut()
4449                .zip(local_layers.iter())
4450                .enumerate()
4451            {
4452                let (Some(distributed), Some(saved)) =
4453                    (distributed_slot.as_mut(), snap.tp_kv_len[il])
4454                else {
4455                    continue;
4456                };
4457                // memra#128: on the dcw / fa2 verify path the rank rows for [saved, target)
4458                // were written on-device and the canonical cache holds only stale bytes for
4459                // them (the verify bumps `local.len` and writes nothing). Copying those over
4460                // the correct rank rows is exactly what spliced two requests' answers
4461                // together on step-3.7-flash. The rewind already kept the right rows; skip.
4462                let target = saved
4463                    .checked_add(accepted)
4464                    .ok_or("spec TP KV batch restore length overflow")?;
4465                if !Self::step_tp_kv_restore_copy_on() || distributed.rows_external() {
4466                    // Rank rows already right (written on-device). Length: see the
4467                    // `rewind_external` note on the signature.
4468                    if rewind_external {
4469                        distributed.rewind_to(target)?;
4470                    }
4471                    continue;
4472                }
4473                let local = local_slot
4474                    .as_ref()
4475                    .ok_or_else(|| format!("spec TP KV layer {il} lost its owning cache"))?;
4476                if local.len < target {
4477                    return Err(format!(
4478                        "spec TP KV layer {il} local length {} precedes restore target {target}",
4479                        local.len
4480                    )
4481                    .into());
4482                }
4483                let physical = local.physical_rows(saved, target)?;
4484                if physical.len() != accepted {
4485                    return Err(format!(
4486                        "spec TP KV layer {il} restore [{saved},{target}) is not contiguous"
4487                    )
4488                    .into());
4489                }
4490                let Mixer::Full(fa) = &self.layers[il].mixer else {
4491                    return Err(format!("spec TP KV layer {il} is not full attention").into());
4492                };
4493                let tp = fa
4494                    .step_tp_qkv
4495                    .as_ref()
4496                    .ok_or_else(|| format!("spec TP KV layer {il} lost its TP runtime"))?;
4497                if let Some(first) = runtime.as_ref() {
4498                    if !std::sync::Arc::ptr_eq(first, &tp.runtime) {
4499                        uniform_runtime = false;
4500                        break;
4501                    }
4502                } else {
4503                    runtime = Some(tp.runtime.clone());
4504                }
4505                use cudarc::driver::DevicePtr;
4506                let (k_base, _k_guard) = local.k.device_ptr(&stream);
4507                let (v_base, _v_guard) = local.v.device_ptr(&stream);
4508                batch.push(crate::tp::TpKvVerifiedLayer {
4509                    cache: distributed,
4510                    start: saved,
4511                    logical_len: target,
4512                    source_k_raw: k_base + (physical.start * local.k_tok_bytes) as u64,
4513                    source_v_raw: v_base + (physical.start * local.v_tok_bytes) as u64,
4514                    source_k_tok_bytes: local.k_tok_bytes,
4515                    source_v_tok_bytes: local.v_tok_bytes,
4516                });
4517            }
4518            if uniform_runtime
4519                && let Some(runtime) = runtime
4520                && runtime.restore_tp_kv_layers_from_device(&mut batch)?
4521            {
4522                return Ok(());
4523            }
4524        }
4525        for il in 0..self.layers.len() {
4526            let (Some(distributed), Some(saved)) = (cache.tp_kv[il].as_mut(), snap.tp_kv_len[il])
4527            else {
4528                continue;
4529            };
4530            let target = saved
4531                .checked_add(accepted)
4532                .ok_or("spec TP KV restore length overflow")?;
4533            if !Self::step_tp_kv_restore_copy_on() || distributed.rows_external() {
4534                if rewind_external {
4535                    distributed.rewind_to(target)?;
4536                }
4537                continue;
4538            }
4539            let local = cache.kv[il]
4540                .as_ref()
4541                .ok_or_else(|| format!("spec TP KV layer {il} lost its owning cache"))?;
4542            if local.len < target {
4543                return Err(format!(
4544                    "spec TP KV layer {il} local length {} precedes restore target {target}",
4545                    local.len
4546                )
4547                .into());
4548            }
4549            let physical = local.physical_rows(saved, target)?;
4550            if physical.len() != accepted {
4551                return Err(format!(
4552                    "spec TP KV layer {il} restore [{saved},{target}) is not contiguous"
4553                )
4554                .into());
4555            }
4556            use cudarc::driver::DevicePtr;
4557            let stream = e.gpu.stream();
4558            let (k_base, _k_guard) = local.k.device_ptr(&stream);
4559            let (v_base, _v_guard) = local.v.device_ptr(&stream);
4560            let k_raw = k_base + (physical.start * local.k_tok_bytes) as u64;
4561            let v_raw = v_base + (physical.start * local.v_tok_bytes) as u64;
4562            let Mixer::Full(fa) = &self.layers[il].mixer else {
4563                return Err(format!("spec TP KV layer {il} is not full attention").into());
4564            };
4565            let tp = fa
4566                .step_tp_qkv
4567                .as_ref()
4568                .ok_or_else(|| format!("spec TP KV layer {il} lost its TP runtime"))?;
4569            tp.runtime.restore_tp_kv_rows_from_device(
4570                distributed,
4571                saved,
4572                target,
4573                k_raw,
4574                v_raw,
4575                local.k_tok_bytes,
4576                local.v_tok_bytes,
4577            )?;
4578        }
4579        Ok(())
4580    }
4581
4582    fn mtp_head_count(&self) -> usize {
4583        usize::from(self.mtp.is_some()) + self.mtp_extra.len()
4584    }
4585
4586    fn mtp_head_at(&self, index: usize) -> &MtpHead {
4587        if index == 0 {
4588            self.mtp.as_ref().expect("MTP head 0 is unavailable")
4589        } else {
4590            &self.mtp_extra[index - 1]
4591        }
4592    }
4593
4594    fn new_mtp_scratch(
4595        &self,
4596        e: &Engine,
4597        cap: usize,
4598    ) -> Result<MtpScratch, Box<dyn std::error::Error>> {
4599        let mut scratch = MtpScratch::new(
4600            e,
4601            &self.cfg,
4602            &self.plan,
4603            cap,
4604            self.mtp.as_ref().and_then(|head| head.geom.as_ref()),
4605        )?;
4606        for head in &self.mtp_extra {
4607            scratch.push_plane(e, &self.cfg, &self.plan, head.geom.as_ref())?;
4608        }
4609        Ok(scratch)
4610    }
4611
4612    fn opti_graph_draft_step(
4613        &self,
4614        e: &Engine,
4615        mtp: &MtpHead,
4616        dctx: &mut DraftGraphCtx,
4617        scratch: &mut MtpScratch,
4618        d_vocab: usize,
4619    ) -> Result<(u32, f32), Box<dyn std::error::Error>> {
4620        // dcw door: one replay appends one device-counter row; pre-arm ring headroom
4621        // host-side before launching (no-op on flat planes).
4622        if step35_draft_dcw_on() {
4623            scratch.ensure_dcw_headroom(e, 2)?;
4624        }
4625        dctx.graph
4626            .as_ref()
4627            .ok_or("optipipe controller requires the greedy draft graph")?
4628            .launch()?;
4629        scratch.kv.len += 1;
4630        let idx = e.dtoh_u32_one(&dctx.g_tok)?;
4631        if (idx as usize) >= d_vocab {
4632            return Err(
4633                format!("optipipe draft argmax sentinel 0x{idx:08x} >= d_vocab {d_vocab}").into(),
4634            );
4635        }
4636        let probability = e.dtoh(&dctx.g_p)?[0];
4637        if !probability.is_finite() || !(0.0..=1.0).contains(&probability) {
4638            return Err(format!("optipipe draft probability is invalid: {probability}").into());
4639        }
4640        let token = match &mtp.d2t {
4641            Some(map) => map[idx as usize],
4642            None => idx,
4643        };
4644        if token != idx {
4645            e.set_u32_one(&mut dctx.g_tok, token)?;
4646        }
4647        Ok((token, probability))
4648    }
4649
4650    #[allow(clippy::too_many_arguments)]
4651    fn opti_controller_draft_step(
4652        &self,
4653        e: &Engine,
4654        mtp: &MtpHead,
4655        dctx: &mut DraftGraphCtx,
4656        scratch: &mut MtpScratch,
4657        d_vocab: usize,
4658        eager_state: &mut Option<(u32, CudaSlice<f32>)>,
4659        eager_pos: usize,
4660        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
4661        round_graph_ok: bool,
4662    ) -> Result<(u32, f32), Box<dyn std::error::Error>> {
4663        // GRAPH-LAUNCH HEADROOM GUARD (see GRAPH_LAUNCH_MIN_FREE): `round_graph_ok` is
4664        // the round's headroom snapshot. Below the floor the main draft arm already ran
4665        // eager (13651-class gate), which seeded `eager_state`, so the controller probe
4666        // rides its eager twin below instead of replaying the draft graph into an
4667        // exhausted card. The seed-unavailable Err beneath stays the recoverable
4668        // fail-closed for the shapes that never seed it.
4669        if dctx.graph.is_some() && round_graph_ok {
4670            return self.opti_graph_draft_step(e, mtp, dctx, scratch, d_vocab);
4671        }
4672        let (input_token, input_seed) = eager_state
4673            .take()
4674            .ok_or("optipipe eager continuation seed is unavailable")?;
4675        let (logits, next_seed) = self.mtp_head_forward_dev(
4676            e,
4677            mtp,
4678            input_token,
4679            &input_seed,
4680            scratch,
4681            eager_pos,
4682            embd_dev,
4683            None,
4684        )?;
4685        let token_d = e.argmax_token_device(&logits, d_vocab)?;
4686        let idx = e.dtoh_u32_one(&token_d)?;
4687        if (idx as usize) >= d_vocab {
4688            return Err(format!(
4689                "optipipe eager draft argmax sentinel 0x{idx:08x} >= d_vocab {d_vocab}"
4690            )
4691            .into());
4692        }
4693        let probability_d = e.prob_of_token_device(&logits, &token_d, d_vocab)?;
4694        let probability = e.dtoh(&probability_d)?[0];
4695        if !probability.is_finite() || !(0.0..=1.0).contains(&probability) {
4696            return Err(
4697                format!("optipipe eager draft probability is invalid: {probability}").into(),
4698            );
4699        }
4700        let token = match &mtp.d2t {
4701            Some(map) => map[idx as usize],
4702            None => idx,
4703        };
4704        *eager_state = Some((token, next_seed));
4705        Ok((token, probability))
4706    }
4707
4708    /// NextN head forward for ONE draft token (§A ops 1-13, T=1).
4709    /// Inputs: `e_tok` = the token to predict FROM (last committed / previous draft); `h_seed` =
4710    /// the trunk's pre-output_norm hidden of that token (§A op 2 input). `mtp_pos` = absolute
4711    /// position of the token being predicted from. Returns (draft_logits[n_vocab] host, h_nextn dev).
4712    /// `h_nextn` (§A op 10) becomes `h_seed` for the next autoregressive draft step.
4713    /// Device-resident: returns draft logits ON DEVICE (no [n_vocab] dtoh). The greedy draft
4714    /// loop only needs argmax — paired with `argmax_token_device` this cuts the ~600KB logits
4715    /// transfer + host argmax per draft token from the K-token draft chain.
4716    #[allow(clippy::too_many_arguments)]
4717    fn mtp_head_forward_dev(
4718        &self,
4719        e: &Engine,
4720        mtp: &MtpHead,
4721        e_tok: u32,
4722        h_seed: &CudaSlice<f32>,
4723        scratch: &mut MtpScratch,
4724        mtp_pos: usize,
4725        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
4726        mask: Option<(&CudaSlice<u32>, usize)>,
4727    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4728        self.mtp_head_forward_dev_at(e, mtp, e_tok, h_seed, scratch, 0, mtp_pos, embd_dev, mask)
4729    }
4730
4731    #[allow(clippy::too_many_arguments)]
4732    fn mtp_head_forward_dev_at(
4733        &self,
4734        e: &Engine,
4735        mtp: &MtpHead,
4736        e_tok: u32,
4737        h_seed: &CudaSlice<f32>,
4738        scratch: &mut MtpScratch,
4739        scratch_index: usize,
4740        mtp_pos: usize,
4741        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
4742        // DRAFT-SIDE GRAMMAR MASK (lane/draft-mask): (packed draft-vocab allowed set, words).
4743        // Applied to the head logits BEFORE they are returned, so every consumer (argmax,
4744        // gumbel draw, p-min prob) sees the grammar-legal row. None = unmasked (pre-lane).
4745        mask: Option<(&CudaSlice<u32>, usize)>,
4746    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4747        // MEMRA_SPEC_ANATOMY=1 — eager-step phase timers (diagnostic only). Phase boundaries
4748        // sync the stream, so absolute time inflates; the BREAKDOWN is the signal. Cumulative
4749        // summary on stderr every 128 steps: glue (embed..attn_norm), attn, ffn, head.
4750        use std::sync::atomic::{AtomicU64, Ordering::Relaxed};
4751        static ANAT_NS: [AtomicU64; 5] = [
4752            AtomicU64::new(0),
4753            AtomicU64::new(0),
4754            AtomicU64::new(0),
4755            AtomicU64::new(0),
4756            AtomicU64::new(0),
4757        ];
4758        static ANAT_STEPS: AtomicU64 = AtomicU64::new(0);
4759        let anat = {
4760            static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
4761            *ON.get_or_init(|| std::env::var("MEMRA_SPEC_ANATOMY").as_deref() == Ok("1"))
4762        };
4763        if anat {
4764            e.stream().synchronize()?; // drain prior queue so phase 0 starts clean
4765        }
4766        let t_all = std::time::Instant::now();
4767        let mut t_ph = std::time::Instant::now();
4768        let anat_mark = |i: usize,
4769                         e: &Engine,
4770                         t: &mut std::time::Instant|
4771         -> Result<(), Box<dyn std::error::Error>> {
4772            if anat {
4773                e.stream().synchronize()?;
4774                ANAT_NS[i].fetch_add(t.elapsed().as_nanos() as u64, Relaxed);
4775                *t = std::time::Instant::now();
4776            }
4777            Ok(())
4778        };
4779        let cfg = &self.cfg;
4780        let n_embd = cfg.n_embd as usize;
4781        // Distilled-student geometry: the block runs at the INNER width `di` (eh_proj out /
4782        // attn / ffn); the n_embd interface (embed, norms in, carrier out, head in) is unchanged.
4783        let di = mtp.geom.as_ref().map(|g| g.d_inner).unwrap_or(n_embd);
4784        let eps = cfg.rms_eps;
4785        let pos_d = e.htod_i32(&[mtp_pos as i32])?;
4786
4787        // op A: a resident table transfers one 4B token id. The exact host-row capacity path
4788        // expands this one row on CPU and transfers n_embd f32 values instead.
4789        let e_emb = match embd_dev {
4790            Some((g, qt, rb)) => e.embed_gather_device_t(g, &[e_tok], n_embd, qt, rb)?,
4791            None => e.htod(&self.embd.try_gather(n_embd, &[e_tok])?)?,
4792        };
4793
4794        // op 1/2: e_norm = RMSNorm(e, enorm); h_norm = RMSNorm(h_seed, hnorm)
4795        let mut e_norm = e.zeros(n_embd)?;
4796        e.rms_norm(&e_emb, mtp.enorm.float_data(), &mut e_norm, n_embd, 1, eps)?;
4797        let mut h_norm = e.zeros(n_embd)?;
4798        e.rms_norm(h_seed, mtp.hnorm.float_data(), &mut h_norm, n_embd, 1, eps)?;
4799
4800        // op 3: concat = [e_norm ; h_norm] -> [2*n_embd], e_norm in [0,n_embd), h_norm in [n_embd,2n_embd)
4801        let mut concat = e.zeros(2 * n_embd)?;
4802        e.copy_into(&mut concat, 0, &e_norm, n_embd)?;
4803        e.copy_into(&mut concat, n_embd, &h_norm, n_embd)?;
4804
4805        // op 4: inpSA = eh_proj @ concat  (eh_proj [2*n_embd, n_embd]) -> [n_embd]
4806        let inp_sa = e.matmul(&mtp.eh_proj, &concat, 1)?;
4807
4808        // op 5: a_norm = RMSNorm(inpSA, attn_norm)
4809        let mut a_norm = e.zeros(di)?;
4810        e.rms_norm(&inp_sa, mtp.attn_norm.float_data(), &mut a_norm, di, 1, eps)?;
4811        anat_mark(0, e, &mut t_ph)?;
4812
4813        // op 6: attention on the scratch KV. SAME dc launcher as the graph path (bucket_max =
4814        // scratch.cap, length from the device len_d) so eager drafts match graph drafts
4815        // bit-for-bit at any t_kv (the parity gate). Host len mirrored here (the dc append
4816        // advances only the device counter).
4817        let attn_out = match (&mtp.mixer, mtp.step35.as_ref()) {
4818            // step35 MTP block, dcw door armed: the SAME windowed device-counter launcher as
4819            // the captured chain (draft parity by construction). Per-step ring headroom runs
4820            // HERE (eager is host-len work, a rebase is legal); host len mirrored like the
4821            // plain dc arm below.
4822            (Mixer::Full(fa), Some(g))
4823                if self.step35_dcw_eligible(g, scratch.plane(scratch_index).1) =>
4824            {
4825                {
4826                    let (kv, _) = scratch.plane_mut(scratch_index);
4827                    let retain = match kv.ring.as_ref() {
4828                        Some(ring) => memra_kv::swa_retain_from(kv.len, ring.window(), ring.base()),
4829                        None => 0,
4830                    };
4831                    e.prepare_kv_append(kv, retain, 1)?;
4832                }
4833                let out =
4834                    self.mtp_step35_attn_dcw(e, fa, g, &a_norm, &pos_d, scratch, scratch_index)?;
4835                scratch.plane_mut(scratch_index).0.len += 1;
4836                out
4837            }
4838            // step35 MTP block, door off (MEMRA_STEP35_DRAFT_DCW=0 rollback) or class-
4839            // ineligible: PER-LAYER geometry + a separate head-wise gate + an SWA window,
4840            // none of which the plain dc launcher can express (see `mtp_step35_attn`).
4841            // Host-len arm. Advances BOTH the
4842            // host len and the device counter itself (unlike the dc arm, whose host-side
4843            // mirror the caller does).
4844            (Mixer::Full(fa), Some(g)) => {
4845                self.mtp_step35_attn(e, fa, g, &a_norm, &pos_d, scratch, scratch_index)?
4846            }
4847            (Mixer::Full(fa), None) => {
4848                let out = self.mtp_full_attn_dc(
4849                    e,
4850                    fa,
4851                    &a_norm,
4852                    &pos_d,
4853                    scratch,
4854                    scratch_index,
4855                    mtp.geom.as_ref(),
4856                )?;
4857                scratch.plane_mut(scratch_index).0.len += 1;
4858                out
4859            }
4860            (Mixer::Linear(_), _) => {
4861                panic!("MTP block is full-attn in qwen35; linear MTP not supported")
4862            }
4863            (Mixer::Mla(_), _) => crate::hybrid::mla_path_unimplemented("MTP head forward"),
4864            (Mixer::Kda(_), _) => crate::hybrid::kda_path_unimplemented("MTP head forward"),
4865        };
4866        anat_mark(1, e, &mut t_ph)?;
4867
4868        // op 7: x1 = inpSA + attn_out
4869        let mut x1 = e.zeros(di)?;
4870        e.add(&inp_sa, &attn_out, &mut x1, di)?;
4871
4872        // op 8: z = RMSNorm(x1, post_attn_norm)  (pre-FFN norm)
4873        let mut z = e.zeros(di)?;
4874        e.rms_norm(&x1, mtp.post_attn_norm.float_data(), &mut z, di, 1, eps)?;
4875
4876        // op 9: FFN (Dense or MoE) — same as the trunk decode FFN
4877        let ffn_out = match &mtp.ffn {
4878            crate::hybrid::Ffn::Dense {
4879                ffn_gate,
4880                ffn_up,
4881                ffn_down,
4882            } => {
4883                let n_ff = ffn_gate.out_features();
4884                let (gate, up) = if e.uses_q8_1_fast(ffn_gate) && e.uses_q8_1_fast(ffn_up) {
4885                    let (zq, zd) = e.quantize_q8_1(&z, 1, di)?;
4886                    (
4887                        e.matmul_pre(ffn_gate, &zq, &zd, &z, 1)?,
4888                        e.matmul_pre(ffn_up, &zq, &zd, &z, 1)?,
4889                    )
4890                } else {
4891                    (e.matmul(ffn_gate, &z, 1)?, e.matmul(ffn_up, &z, 1)?)
4892                };
4893                let mut act = e.zeros(n_ff)?;
4894                // step35: a DENSE FFN reads the per-layer SHEXP clamp (upstream's one `build_ffn`
4895                // serves the dense MLP and the shared expert off `swiglu_clamp_shexp` —
4896                // llama-graph.cpp:1751), resolved for the MTP block's OWN index. Every other arch
4897                // passes None, which is `ffn_act`'s dispatch verbatim.
4898                Self::ffn_act_lim(
4899                    e,
4900                    &self.cfg,
4901                    &gate,
4902                    &up,
4903                    1.0,
4904                    1.0,
4905                    mtp.step35
4906                        .as_ref()
4907                        .and_then(|s| s.clamp_shexp)
4908                        .map(SwigluClamp::Post),
4909                    &mut act,
4910                    n_ff,
4911                )?;
4912                e.matmul(ffn_down, &act, 1)?
4913            }
4914            // MTP head is a distinct block — key its experts under a separate layer index (u16::MAX)
4915            // so they never alias trunk layer 0's cache keys.
4916            crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il(e, m, &z, 1, u16::MAX)?,
4917        };
4918        anat_mark(2, e, &mut t_ph)?;
4919
4920        // op 10: h_nextn = x1 + ffn_out (at di)
4921        let mut h_inner = e.zeros(di)?;
4922        e.add(&x1, &ffn_out, &mut h_inner, di)?;
4923
4924        // op 10.5 (student): up-project the inner hidden back to n_embd — training semantics:
4925        // the chain carrier AND the head input are out_up(h_inner) (pre-final-norm).
4926        let h_nextn = match mtp.geom.as_ref() {
4927            Some(g) => e.matmul(&g.out_up, &h_inner, 1)?,
4928            None => h_inner,
4929        };
4930
4931        // op 11: final = RMSNorm(h_nextn, shared_head_norm OR output_norm)
4932        let final_norm = mtp.shared_head_norm.as_ref().unwrap_or(&self.output_norm);
4933        let mut final_h = e.zeros(n_embd)?;
4934        e.rms_norm(
4935            &h_nextn,
4936            final_norm.float_data(),
4937            &mut final_h,
4938            n_embd,
4939            1,
4940            eps,
4941        )?;
4942
4943        // op 12: draft_logits = (shared_head_head OR output) @ final — stays ON DEVICE.
4944        let head = mtp.shared_head_head.as_ref().unwrap_or(&self.output);
4945        let mut logits = e.matmul(head, &final_h, 1)?;
4946        // op 12b (lane/draft-mask): grammar mask over the DRAFT vocab, applied here so the
4947        // caller's argmax / gumbel draw / p-min prob all read the grammar-legal row.
4948        if let Some((mask_d, mw)) = mask {
4949            let d_vocab = head.out_features();
4950            e.mask_logits_col(&mut logits, mask_d, 0, d_vocab, mw)?;
4951        }
4952        anat_mark(3, e, &mut t_ph)?;
4953        if anat {
4954            ANAT_NS[4].fetch_add(t_all.elapsed().as_nanos() as u64, Relaxed);
4955            let n = ANAT_STEPS.fetch_add(1, Relaxed) + 1;
4956            if n.is_multiple_of(128) {
4957                let us = |i: usize| ANAT_NS[i].load(Relaxed) / n / 1000;
4958                eprintln!(
4959                    "[spec-anatomy] steps={n} avg us/step: glue={} attn={} ffn={} head={} total={}",
4960                    us(0),
4961                    us(1),
4962                    us(2),
4963                    us(3),
4964                    us(4)
4965                );
4966            }
4967        }
4968        // Chain recurrence hand-over: pre-norm h_nextn (default) or post-norm final_h
4969        // (MEMRA_SPEC_HPOST — llama.cpp #24025's t_h_nextn is taken AFTER the head norm).
4970        Ok((logits, if spec_hpost() { final_h } else { h_nextn }))
4971    }
4972
4973    /// One NextN/MTP draft step for an **MLA-mixer** MTP block (glm5_next class: MLA + own
4974    /// k-pool indexer + MoE, serial residual — the NextN layer carries no hc_* tensors), on
4975    /// the model `Cache`'s own MTP latent plane rather than the full-attn `MtpScratch` the
4976    /// qwen35/step35 chain uses. Gate: `glm5_mtp_head_gpu` (engine vs `memra_reference`
4977    /// `execute_mtp`, teacher-forced walk, eh_proj-transpose and h_seed-off-by-one red arms).
4978    ///
4979    /// The interface, stated precisely for the verify arc:
4980    /// - `h_seed`: `[n_embd]` f32 device — the trunk's COLLAPSED PRE-output_norm hidden of
4981    ///   the position whose next token is being drafted (MTP-PLAN §A; exactly what
4982    ///   `prime_cache`/`decode_step` return for hc models). `MEMRA_SPEC_HPOST` flips both
4983    ///   this producer and the returned carrier to the post-norm variant, same as the dev path.
4984    /// - `e_tok`: the token at the seeded position's SUCCESSOR — the token the trunk just
4985    ///   sampled/accepted (reference oracle pairing: `fused[i] = eh_proj([enorm(embed(ids[i]));
4986    ///   hnorm(trunk_hidden[i])])`, i.e. this call with `e_tok = ids[i]`, `h_seed = h[i]`,
4987    ///   `mtp_pos = i` reproduces the reference's row `i`).
4988    /// - `mtp_pos`: the absolute position this step appends to the MTP block's latent plane;
4989    ///   must equal that plane's current length (the plane advances by ONE row per call inside
4990    ///   `mla_attn_cached`; rollback on rejection = the verify arc's latent-plane len reset).
4991    /// - returns `(draft_logits [n_vocab], carrier [n_embd])` on device. glm5_next ships no
4992    ///   private MTP head, so the logits ride the trunk `lm_head` (full vocab, no d2t).
4993    pub fn mtp_head_forward_mla_cached(
4994        &self,
4995        e: &Engine,
4996        depth: usize,
4997        e_tok: u32,
4998        h_seed: &CudaSlice<f32>,
4999        cache: &mut Cache,
5000        mtp_pos: usize,
5001    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
5002        if depth >= self.mtp_head_count() {
5003            return Err(format!(
5004                "MTP depth {depth} out of range: {} embedded head(s) loaded \
5005                 (is MEMRA_GLM5_MTP=1 set for a glm5_next model?)",
5006                self.mtp_head_count()
5007            )
5008            .into());
5009        }
5010        let mtp = self.mtp_head_at(depth);
5011        let block = self
5012            .plan
5013            .mtp_blocks
5014            .get(depth)
5015            .ok_or_else(|| format!("ModelPlan declares no MTP block at depth {depth}"))?;
5016        let il = block.layer.index as usize;
5017        let Mixer::Mla(mla) = &mtp.mixer else {
5018            return Err(
5019                "mtp_head_forward_mla_cached serves MLA-mixer MTP blocks only; full-attn \
5020                 blocks take mtp_head_forward_dev's scratch path"
5021                    .into(),
5022            );
5023        };
5024        if matches!(mtp.ffn, crate::hybrid::Ffn::Dense { .. }) {
5025            return Err(
5026                "MLA-mixer MTP block with a Dense FFN has no gated arm yet (glm5_next and \
5027                 glm-dsa NextN blocks are MoE); refusing rather than running ungated math"
5028                    .into(),
5029            );
5030        }
5031        let plane_len = cache
5032            .latent
5033            .get(il)
5034            .and_then(|plane| plane.as_ref())
5035            .map(|plane| plane.len)
5036            .ok_or_else(|| {
5037                format!(
5038                    "MTP block layer {il} has no latent cache plane — the Cache must be \
5039                     built from a plan whose mtp_blocks declare StatePlan::LatentKvCache"
5040                )
5041            })?;
5042        if mtp_pos != plane_len {
5043            return Err(format!(
5044                "MTP draft position {mtp_pos} != the MTP latent plane's length {plane_len} — \
5045                 the plane advances one row per draft step and rolls back by len reset; a \
5046                 skipped or repeated position would attend the wrong horizon"
5047            )
5048            .into());
5049        }
5050
5051        let cfg = &self.cfg;
5052        let n_embd = cfg.n_embd as usize;
5053        let eps = cfg.rms_eps;
5054        let pos_d = e.htod_i32(&[mtp_pos as i32])?;
5055
5056        // Same op chain as `mtp_head_forward_dev_at` (ops 1-12), same kernels — only the
5057        // attention arm differs: `mla_attn_cached` on the plan's own MTP plane instead of
5058        // `mtp_full_attn_dc` on the MtpScratch.
5059        let e_emb = e.htod(&self.embd.try_gather(n_embd, &[e_tok])?)?;
5060        let mut e_norm = e.zeros(n_embd)?;
5061        e.rms_norm(&e_emb, mtp.enorm.float_data(), &mut e_norm, n_embd, 1, eps)?;
5062        let mut h_norm = e.zeros(n_embd)?;
5063        e.rms_norm(h_seed, mtp.hnorm.float_data(), &mut h_norm, n_embd, 1, eps)?;
5064
5065        let mut concat = e.zeros(2 * n_embd)?;
5066        e.copy_into(&mut concat, 0, &e_norm, n_embd)?;
5067        e.copy_into(&mut concat, n_embd, &h_norm, n_embd)?;
5068        let inp_sa = e.matmul(&mtp.eh_proj, &concat, 1)?;
5069
5070        let mut a_norm = e.zeros(n_embd)?;
5071        e.rms_norm(
5072            &inp_sa,
5073            mtp.attn_norm.float_data(),
5074            &mut a_norm,
5075            n_embd,
5076            1,
5077            eps,
5078        )?;
5079        let attn_out = self.mla_attn_cached(e, mla, &a_norm, &pos_d, 1, il, cache)?;
5080
5081        let mut x1 = e.zeros(n_embd)?;
5082        e.add(&inp_sa, &attn_out, &mut x1, n_embd)?;
5083        let mut z = e.zeros(n_embd)?;
5084        e.rms_norm(&x1, mtp.post_attn_norm.float_data(), &mut z, n_embd, 1, eps)?;
5085        let ffn_out = match &mtp.ffn {
5086            // Distinct block — key its experts off the trunk layers' cache keys (dev-path rule).
5087            crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il(e, m, &z, 1, u16::MAX)?,
5088            crate::hybrid::Ffn::Dense { .. } => unreachable!("refused above"),
5089        };
5090        let mut h_nextn = e.zeros(n_embd)?;
5091        e.add(&x1, &ffn_out, &mut h_nextn, n_embd)?;
5092
5093        let final_norm = mtp.shared_head_norm.as_ref().unwrap_or(&self.output_norm);
5094        let mut final_h = e.zeros(n_embd)?;
5095        e.rms_norm(
5096            &h_nextn,
5097            final_norm.float_data(),
5098            &mut final_h,
5099            n_embd,
5100            1,
5101            eps,
5102        )?;
5103        let head = mtp.shared_head_head.as_ref().unwrap_or(&self.output);
5104        let logits = e.matmul(head, &final_h, 1)?;
5105        Ok((logits, if spec_hpost() { final_h } else { h_nextn }))
5106    }
5107
5108    #[allow(clippy::too_many_arguments)]
5109    fn mtp_chain_forward_dev(
5110        &self,
5111        e: &Engine,
5112        tokens: &[u32],
5113        seeds: &[CudaSlice<f32>],
5114        scratch: &mut MtpScratch,
5115        committed_scratch_len: usize,
5116        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
5117        mask: Option<(&CudaSlice<u32>, usize)>,
5118    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
5119        if tokens.is_empty() || tokens.len() != seeds.len() {
5120            return Err("multi-head MTP prefix tokens/seeds are malformed".into());
5121        }
5122        let index = mtp_chain_head_index(tokens.len() - 1, self.mtp_head_count());
5123        let head = self.mtp_head_at(index);
5124        scratch.set_plane_len(e, index, committed_scratch_len)?;
5125
5126        let mut last = None;
5127        for row in 0..tokens.len() {
5128            let is_last = row + 1 == tokens.len();
5129            last = Some(self.mtp_head_forward_dev_at(
5130                e,
5131                head,
5132                tokens[row],
5133                &seeds[row],
5134                scratch,
5135                index,
5136                committed_scratch_len + row + 1,
5137                embd_dev,
5138                if is_last { mask } else { None },
5139            )?);
5140        }
5141        Ok(last.expect("non-empty MTP prefix produced no row"))
5142    }
5143
5144    /// step35 MTP-block attention, T=1, on the scratch KV — the EAGER-ONLY twin of
5145    /// `mtp_full_attn_dc`. Three things force a separate arm rather than a geometry parameter on
5146    /// the dc path, and all three are properties of this arch's MTP block:
5147    ///
5148    /// 1. **The SWA window.** Block 45 is an SWA-type block (`sliding_window_pattern[45]=true`,
5149    ///    window 512). Windowed decode in memra is a token-aligned VIEW OFFSET into the quantized
5150    ///    cache (the gemma4 R6 / `step35_decode_attn` pattern: keys carry absolute rope and the
5151    ///    mask is purely positional, so one query at `len-1` attending the last `win` rows IS the
5152    ///    windowed result). `fa_decode_dc` takes the key count from a DEVICE counter and always
5153    ///    starts at row 0 — it cannot express a nonzero offset. The windowed dc arm is
5154    ///    `mtp_step35_attn_dcw` (`fa_decode_dcw`, doored via MEMRA_STEP35_DRAFT_DCW —
5155    ///    default ON since lane/step37-draft-graph-serving-20260830); this host-len arm is
5156    ///    the =0 rollback and the class-ineligibility fallback.
5157    /// 2. **Per-layer head count.** 96 q heads over 8 KV (GQA 12) at this block, vs the trunk's 64
5158    ///    on its full-attn layers. The trunk cfg's `n_head` scalar is the MAX over layers, and the
5159    ///    trunk ARTIFACT's per-layer arrays stop at index 44 — so the count must come from the
5160    ///    resolved `Step35MtpGeom`, never from `cfg`.
5161    /// 3. **The separate head-wise gate.** `blk.45.attn_gate.weight [n_embd, 96]` produces one
5162    ///    sigmoid scalar per head (broadcast over head_dim) — `attn_head_gate`, not the qwen35
5163    ///    fused-into-wq `q_gate_split` form the dc arm handles.
5164    ///
5165    /// DOOR STATE: with MEMRA_STEP35_DRAFT_DCW=0 (or a sub-eligible kernel class),
5166    /// `mtp_head_forward_cap` refuses step35 heads explicitly (rather than silently capturing
5167    /// a window-less, wrong-past-`win` graph) and this eager chain IS the served path. With
5168    /// the door armed (the default), BOTH draft modes run the `mtp_step35_attn_dcw` twin
5169    /// instead of this arm.
5170    ///
5171    /// Unlike the dc arm this advances BOTH the host `kv.len` and the device counter, so the
5172    /// caller must not mirror.
5173    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
5174    fn mtp_step35_attn(
5175        &self,
5176        e: &Engine,
5177        fa: &FullAttnLayer,
5178        g: &crate::hybrid::Step35MtpGeom,
5179        h: &CudaSlice<f32>,
5180        pos_d: &CudaSlice<i32>,
5181        scratch: &mut MtpScratch,
5182        scratch_index: usize,
5183    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5184        let (nh, nkv, hd) = (g.n_head, g.n_head_kv, self.cfg.head_dim_k as usize);
5185        // MTP-GEOM RECEIPT, once per process, on the SERVED draft path. Slot-0 acceptance is
5186        // 0.725 here against 0.994 for vLLM MTP3 on the same checkpoint family and card class, and
5187        // the first three explanations for that gap were all wrong: head assignment (step-modulo
5188        // is index 0 at K=1, correct), MEMRA_SPEC_HPOST (identical 84/116 both arms), and this
5189        // block's geometry. Geometry was the one that could have failed SILENTLY — a wrong window
5190        // makes the draft attend the whole context instead of Step-3.7's 512, stays fluent, and
5191        // shows up only as acceptance — so it gets a standing receipt rather than another reading
5192        // of the source. Prints the resolved Step35MtpGeom the served path actually runs on;
5193        // `full_attention_geometry_at`'s missing-row fallback (window: None) does NOT reach here.
5194        {
5195            static ONCE: std::sync::OnceLock<()> = std::sync::OnceLock::new();
5196            ONCE.get_or_init(|| {
5197                eprintln!(
5198                    "[mtp-geom] arm=eager block={} swa={} window={} n_head={nh} n_head_kv={nkv} \
5199                     head_dim_k={hd} n_rot={} rope_base={} clamp_shexp={:?}",
5200                    g.il, g.swa, g.window, g.n_rot, g.rope_base, g.clamp_shexp,
5201                );
5202            });
5203        }
5204        let eps = self.cfg.rms_eps;
5205        let scale = 1.0 / (hd as f32).sqrt(); // step35.cpp:255 kq_scale
5206        let n_embd = self.cfg.n_embd as usize;
5207        let gw = fa
5208            .attn_gate
5209            .as_ref()
5210            .ok_or("step35 MTP block is missing attn_gate.weight (head-wise attention gate)")?;
5211
5212        let (q0, k0, v0, gt) = if e.uses_q8_1_fast(&fa.wq)
5213            && e.uses_q8_1_fast(&fa.wk)
5214            && e.uses_q8_1_fast(&fa.wv)
5215            && e.uses_q8_1_fast(gw)
5216        {
5217            let (hq, hdq) = e.quantize_q8_1(h, 1, n_embd)?;
5218            let (a, b, c) = match e.matmul_q8_fused3(&fa.wq, &fa.wk, &fa.wv, &hq, &hdq)? {
5219                Some(t3) => t3,
5220                None => (
5221                    e.matmul_pre(&fa.wq, &hq, &hdq, h, 1)?,
5222                    e.matmul_pre(&fa.wk, &hq, &hdq, h, 1)?,
5223                    e.matmul_pre(&fa.wv, &hq, &hdq, h, 1)?,
5224                ),
5225            };
5226            (a, b, c, e.matmul_pre(gw, &hq, &hdq, h, 1)?)
5227        } else {
5228            (
5229                e.matmul(&fa.wq, h, 1)?,
5230                e.matmul(&fa.wk, h, 1)?,
5231                e.matmul(&fa.wv, h, 1)?,
5232                e.matmul(gw, h, 1)?,
5233            )
5234        };
5235
5236        let mut q = e.uninit(nh * hd)?;
5237        e.rms_norm(&q0, fa.q_norm.float_data(), &mut q, hd, nh, eps)?;
5238        let mut k = e.uninit(nkv * hd)?;
5239        e.rms_norm(&k0, fa.k_norm.float_data(), &mut k, hd, nkv, eps)?;
5240        // `rope_freqs.weight` (llama3 factors) applies to the FULL-attn layers ONLY; SWA passes
5241        // null (llama-hparams / step35.cpp). Block 45 is SWA, so `ff` is None there — but read
5242        // the resolved flag, not the constant, so an all-full sibling stays correct.
5243        let ff = if g.swa {
5244            None
5245        } else {
5246            self.step35_aux.as_ref().and_then(|a| a.rope_freqs(e))
5247        };
5248        #[cfg(debug_assertions)]
5249        if let Some(ff) = ff {
5250            crate::debug_assert_tensor_stream_device(ff, &e.stream(), "mtp_step35_attn.rope_freqs");
5251        }
5252        e.rope_neox2(
5253            &mut q,
5254            &mut k,
5255            pos_d,
5256            hd,
5257            g.n_rot,
5258            nh,
5259            nkv,
5260            1,
5261            g.rope_base,
5262            1.0,
5263            ff,
5264        )?;
5265
5266        // Append at the HOST slot, then re-stamp the device counter: the eager chain has the
5267        // length on the host anyway, and the windowed view below needs it there to compute the
5268        // offset. The device counter is kept in lockstep so `mtp_kv_fill`'s `set_i32_one` and any
5269        // dc-family consumer of this scratch still agree.
5270        let (kv, scratch_cap) = scratch.plane_mut(scratch_index);
5271        assert!(
5272            kv.len < scratch_cap,
5273            "step35 MTP scratch overflow ({} >= {})",
5274            kv.len,
5275            scratch_cap
5276        );
5277        let next_len = kv.len + 1;
5278        let (off, t_kv) = if g.swa && next_len > g.window {
5279            (next_len - g.window, g.window)
5280        } else {
5281            (0, next_len)
5282        };
5283        // `off`/`t_kv` stay the ATTENTION view; the retain is a separate, lower bound so the
5284        // rewind that follows this append is still resident. THIS is the only site that rebases
5285        // this plane (MEMRA_KV_REBASE_TRACE, one run: 1 rebase, all from here), so it is the site
5286        // that decides `base` for everyone.
5287        let retain_from = match kv.ring.as_ref() {
5288            Some(ring) => memra_kv::swa_retain_from(kv.len, ring.window(), ring.base()),
5289            None => off & !31usize,
5290        };
5291        let write_row = e.prepare_kv_append(kv, retain_from, 1)?;
5292        e.append_kv_quantized(
5293            &k,
5294            &v0,
5295            &mut kv.k,
5296            &mut kv.v,
5297            write_row,
5298            kv.kv_dim_k,
5299            kv.kv_dim_v,
5300            kv.k_tok_bytes,
5301            kv.v_tok_bytes,
5302            false,
5303        )?;
5304        kv.len = next_len;
5305        e.set_i32_one(&mut kv.len_d, kv.len as i32)?;
5306        // SWA view offset (see note 1). The draft chain is short (k+2 rows), but the scratch is
5307        // PERSISTENT across rounds — `mtp_kv_fill` leaves one row per committed token behind, so
5308        // `kv.len` tracks absolute position and crosses 512 in any real generation. The window is
5309        // therefore live, not theoretical.
5310        let physical = kv.physical_rows(off, off + t_kv)?;
5311        let k_view = e.view_u8_range(
5312            &kv.k,
5313            physical.start * kv.k_tok_bytes,
5314            physical.end * kv.k_tok_bytes,
5315        );
5316        let v_view = e.view_u8_range(
5317            &kv.v,
5318            physical.start * kv.v_tok_bytes,
5319            physical.end * kv.v_tok_bytes,
5320        );
5321        let mut attn = e.uninit(nh * hd)?;
5322        e.fa_decode_kvmod(
5323            &q,
5324            &k_view,
5325            &v_view,
5326            &mut attn,
5327            hd,
5328            nh,
5329            nkv,
5330            t_kv,
5331            scale,
5332            kv.k_tok_bytes,
5333            kv.v_tok_bytes,
5334            false,
5335        )?;
5336
5337        let mut ag = e.uninit(nh * hd)?;
5338        e.attn_head_gate(&attn, &gt, &mut ag, None, hd, nh, 1)?;
5339        e.matmul(&fa.wo, &ag, 1)
5340    }
5341
5342    /// The dcw draft arm's kernel-class precondition, mirrored from `fa_decode_dcw`'s own
5343    /// refusal plus the v3 walk's format contract (`fa_v3_active`), so the DEV dispatch can
5344    /// never pick an arm the launcher would refuse mid-chain (the eager chain has no graceful
5345    /// fallback point) and the CAP site refuses with the named reason instead.
5346    ///
5347    /// `cap` = the SESSION's scratch-plane row capacity: the launcher's vec gate reads
5348    /// `bucket_max = min(window, cap)`, so a SMALL session (tiny prompt + tiny max_tokens,
5349    /// e.g. a max_tokens=8 probe: cap ~62 < the 96 vec floor) is OUTSIDE the dcw domain even
5350    /// though the WINDOW clears the floor. Mirroring the window alone shipped exactly that
5351    /// hole when the door default flipped ON (2026-08-30, vision-cell receipt: sampled
5352    /// capture WARN + `[engine-error] fa_decode_dcw supports the default v3-vec class only`
5353    /// hard-failing the burst — the eager dcw arm has no graceful fallback point). Sub-floor
5354    /// sessions now take the host-len kvmod arm, byte-for-byte the door-off serving.
5355    fn step35_dcw_eligible(&self, g: &crate::hybrid::Step35MtpGeom, cap: usize) -> bool {
5356        let hd = self.cfg.head_dim_k as usize;
5357        step35_draft_dcw_on()
5358            && g.swa
5359            && g.window.min(cap) >= crate::fa_vec_min_tkv()
5360            && std::env::var("MEMRA_NO_FA_VEC").is_err()
5361            && crate::fa_v3_active(hd)
5362            && hd <= 256
5363            && hd.is_multiple_of(32)
5364    }
5365
5366    /// step35 MTP-block attention, T=1, on the scratch KV: the WINDOWED DEVICE-COUNTER twin
5367    /// of `mtp_step35_attn`, serving BOTH draft paths when `step35_draft_dcw_on`. Write slot,
5368    /// key bound and SWA view offset all derive from device state (`len_d`, `base_d` written
5369    /// only at host-side rebases, and the block's `window`), so ONE captured graph serves the
5370    /// whole chain and replays see KV growth through the counter: the `mtp_full_attn_dc`
5371    /// contract plus the view offset the plain `_dc` kernel could not express (the old
5372    /// capture-refusal root cause). The three step35 properties stay per-geom exactly as in
5373    /// the eager twin: nh/nkv from `Step35MtpGeom`, the separate head-wise gate
5374    /// (`attn_head_gate`), per-layer rope width/base with SWA passing null freqs.
5375    ///
5376    /// bucket_max = min(cap, window): the windowed view never exceeds `window` rows, so the
5377    /// capture-time grid stays valid for every replayed len, and the kernel derives ns_eff
5378    /// from the LIVE T_kv at the fixed split_keys (one-partition law). Both arms call THIS
5379    /// launcher at THIS bucket, so eager and captured drafts are bit-identical by
5380    /// construction; vs the retired-by-flag `mtp_step35_attn` the only numeric-class deltas
5381    /// are the sub-vec-floor region (t_kv < 96: kvmod ran scalar, dcw stays vec) and any
5382    /// live-len split-ladder rung below the bucket's, both draft-side only (the verify
5383    /// arbitrates emitted bytes; acceptance is gated by the battery).
5384    ///
5385    /// Host len is NOT advanced here (graph contract); callers mirror. The EAGER caller runs
5386    /// `prepare_kv_append` per step (ring headroom, rebase legal there); the CAPTURED path
5387    /// pre-arms headroom at capture time and round start (`MtpScratch::ensure_dcw_headroom`)
5388    /// because a rebase is host work no captured chain may contain.
5389    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the capture/call contract; bundling into a struct is a refactor, not a lint fix
5390    fn mtp_step35_attn_dcw(
5391        &self,
5392        e: &Engine,
5393        fa: &FullAttnLayer,
5394        g: &crate::hybrid::Step35MtpGeom,
5395        h: &CudaSlice<f32>,
5396        pos_d: &CudaSlice<i32>,
5397        scratch: &mut MtpScratch,
5398        scratch_index: usize,
5399    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5400        let (nh, nkv, hd) = (g.n_head, g.n_head_kv, self.cfg.head_dim_k as usize);
5401        // MTP-GEOM RECEIPT (dcw twin of the `mtp_step35_attn` receipt): once per process,
5402        // naming the arm, so a serving log proves WHICH draft attention program ran (the
5403        // engagement receipt for the flag door, both directions).
5404        {
5405            static ONCE: std::sync::OnceLock<()> = std::sync::OnceLock::new();
5406            ONCE.get_or_init(|| {
5407                eprintln!(
5408                    "[mtp-geom] arm=dcw block={} swa={} window={} n_head={nh} n_head_kv={nkv} \
5409                     head_dim_k={hd} n_rot={} rope_base={} clamp_shexp={:?}",
5410                    g.il, g.swa, g.window, g.n_rot, g.rope_base, g.clamp_shexp,
5411                );
5412            });
5413        }
5414        let eps = self.cfg.rms_eps;
5415        let scale = 1.0 / (hd as f32).sqrt(); // step35.cpp:255 kq_scale
5416        let n_embd = self.cfg.n_embd as usize;
5417        let gw = fa
5418            .attn_gate
5419            .as_ref()
5420            .ok_or("step35 MTP block is missing attn_gate.weight (head-wise attention gate)")?;
5421
5422        let (q0, k0, v0, gt) = if e.uses_q8_1_fast(&fa.wq)
5423            && e.uses_q8_1_fast(&fa.wk)
5424            && e.uses_q8_1_fast(&fa.wv)
5425            && e.uses_q8_1_fast(gw)
5426        {
5427            let (hq, hdq) = e.quantize_q8_1(h, 1, n_embd)?;
5428            let (a, b, c) = match e.matmul_q8_fused3(&fa.wq, &fa.wk, &fa.wv, &hq, &hdq)? {
5429                Some(t3) => t3,
5430                None => (
5431                    e.matmul_pre(&fa.wq, &hq, &hdq, h, 1)?,
5432                    e.matmul_pre(&fa.wk, &hq, &hdq, h, 1)?,
5433                    e.matmul_pre(&fa.wv, &hq, &hdq, h, 1)?,
5434                ),
5435            };
5436            (a, b, c, e.matmul_pre(gw, &hq, &hdq, h, 1)?)
5437        } else {
5438            (
5439                e.matmul(&fa.wq, h, 1)?,
5440                e.matmul(&fa.wk, h, 1)?,
5441                e.matmul(&fa.wv, h, 1)?,
5442                e.matmul(gw, h, 1)?,
5443            )
5444        };
5445
5446        let mut q = e.zeros(nh * hd)?;
5447        e.rms_norm(&q0, fa.q_norm.float_data(), &mut q, hd, nh, eps)?;
5448        let mut k = e.zeros(nkv * hd)?;
5449        e.rms_norm(&k0, fa.k_norm.float_data(), &mut k, hd, nkv, eps)?;
5450        // rope_freqs (llama3 factors) apply to the FULL-attn layers ONLY; SWA passes null
5451        // (the eager twin's rule, resolved from the flag, not the constant).
5452        let ff = if g.swa {
5453            None
5454        } else {
5455            self.step35_aux.as_ref().and_then(|a| a.rope_freqs(e))
5456        };
5457        #[cfg(debug_assertions)]
5458        if let Some(ff) = ff {
5459            crate::debug_assert_tensor_stream_device(
5460                ff,
5461                &e.stream(),
5462                "mtp_step35_attn_dcw.rope_freqs",
5463            );
5464        }
5465        e.rope_neox2(
5466            &mut q,
5467            &mut k,
5468            pos_d,
5469            hd,
5470            g.n_rot,
5471            nh,
5472            nkv,
5473            1,
5474            g.rope_base,
5475            1.0,
5476            ff,
5477        )?;
5478
5479        let (kv, cap) = scratch.plane_mut(scratch_index);
5480        // Append at the DEVICE slot's PHYSICAL row (len_d - base_d), then advance the counter
5481        // in-graph. Physical room is the callers' headroom contract (see the fn doc).
5482        e.append_kv_quantized_dcw(
5483            &k,
5484            &v0,
5485            &mut kv.k,
5486            &mut kv.v,
5487            &kv.len_d,
5488            kv.base_d.as_ref(),
5489            kv.kv_dim_k,
5490            kv.kv_dim_v,
5491            kv.k_tok_bytes,
5492            kv.v_tok_bytes,
5493        )?;
5494        e.inc_seqlen(&mut kv.len_d)?;
5495        // Full-buffer views (any in-round physical row stays in range under the headroom
5496        // contract); the kernel bounds and offsets the key range from (len_d, base_d, window).
5497        let k_view = e.view_u8(&kv.k, kv.k.len());
5498        let v_view = e.view_u8(&kv.v, kv.v.len());
5499        let bucket = g.window.min(cap);
5500        let mut attn = e.zeros(nh * hd)?;
5501        e.fa_decode_dcw(
5502            &q,
5503            &k_view,
5504            &v_view,
5505            &mut attn,
5506            hd,
5507            nh,
5508            nkv,
5509            &kv.len_d,
5510            kv.base_d.as_ref(),
5511            if g.swa { g.window } else { 0 },
5512            bucket,
5513            scale,
5514            kv.k_tok_bytes,
5515            kv.v_tok_bytes,
5516            None,
5517        )?;
5518
5519        let mut ag = e.zeros(nh * hd)?;
5520        e.attn_head_gate(&attn, &gt, &mut ag, None, hd, nh, 1)?;
5521        e.matmul(&fa.wo, &ag, 1)
5522    }
5523
5524    /// MTP-block full attention, T=1, on the scratch KV (BOTH draft paths — eager and graph):
5525    /// the scratch write slot and the attention bound come from `scratch.kv.len_d` (device i32[1])
5526    /// so the launch args are FIXED across draft steps — ONE captured graph serves the whole
5527    /// chain, and replays keep seeing KV growth through the device counter (no recapture).
5528    /// Geometry contract: n_splits is sized from `scratch.cap` (the persistent capacity); splits
5529    /// whose key range lies beyond the device t_kv exit empty and the shared combine skips them
5530    /// (fa_decode_dc bit-correct-for-any-t_kv<=bucket_max contract). The eager path uses the SAME
5531    /// launcher with the SAME bucket_max -> identical dispatch -> bit-identical draft tokens (the
5532    /// graph-vs-eager parity gate). Host len is NOT advanced here (graph contract); callers mirror.
5533    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
5534    fn mtp_full_attn_dc(
5535        &self,
5536        e: &Engine,
5537        fa: &FullAttnLayer,
5538        h: &CudaSlice<f32>,
5539        pos_d: &CudaSlice<i32>,
5540        scratch: &mut MtpScratch,
5541        scratch_index: usize,
5542        geom: Option<&crate::hybrid::DraftGeom>,
5543    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5544        let cfg = &self.cfg;
5545        let mtp_il = cfg.n_layer.saturating_sub(cfg.nextn_predict_layers);
5546        let geometry = cfg.full_attention_geometry_at(mtp_il);
5547        let n_head = geom.map(|g| g.n_head).unwrap_or(geometry.n_head as usize);
5548        let n_head_kv = geom
5549            .map(|g| g.n_head_kv)
5550            .unwrap_or(geometry.n_head_kv as usize);
5551        let head_dim = geometry.head_dim_k as usize;
5552        let eps = cfg.rms_eps;
5553        let scale = geometry.attention_scale();
5554        let n_embd = geom.map(|g| g.d_inner).unwrap_or(cfg.n_embd as usize);
5555        let bucket_max = scratch.plane(scratch_index).1;
5556
5557        let (qf, mut k, v) =
5558            if e.uses_q8_1_fast(&fa.wq) && e.uses_q8_1_fast(&fa.wk) && e.uses_q8_1_fast(&fa.wv) {
5559                let (hq, hd) = e.quantize_q8_1(h, 1, n_embd)?;
5560                (
5561                    e.matmul_pre(&fa.wq, &hq, &hd, h, 1)?,
5562                    e.matmul_pre(&fa.wk, &hq, &hd, h, 1)?,
5563                    e.matmul_pre(&fa.wv, &hq, &hd, h, 1)?,
5564                )
5565            } else {
5566                (
5567                    e.matmul(&fa.wq, h, 1)?,
5568                    e.matmul(&fa.wk, h, 1)?,
5569                    e.matmul(&fa.wv, h, 1)?,
5570                )
5571            };
5572        // M3/Hy3 have no attention output gate — wq out is exactly q; skip the split.
5573        let gated = geometry.attention_gate == memra_gguf::config::AttentionGateKind::FusedQ;
5574        let (mut q, gate) = if gated {
5575            let mut q = e.zeros(n_head * head_dim)?;
5576            let mut gate = e.zeros(n_head * head_dim)?;
5577            e.q_gate_split(&qf, &mut q, &mut gate, head_dim, n_head, 1)?;
5578            (q, Some(gate))
5579        } else {
5580            (qf, None)
5581        };
5582
5583        let mut qn = e.zeros(n_head * head_dim)?;
5584        e.rms_norm(&q, fa.q_norm.float_data(), &mut qn, head_dim, n_head, eps)?;
5585        q = qn;
5586        let mut kn = e.zeros(n_head_kv * head_dim)?;
5587        e.rms_norm(
5588            &k,
5589            fa.k_norm.float_data(),
5590            &mut kn,
5591            head_dim,
5592            n_head_kv,
5593            eps,
5594        )?;
5595        k = kn;
5596        let rope_dims = geometry.n_rot as usize;
5597        e.rope_neox(
5598            &mut q,
5599            pos_d,
5600            head_dim,
5601            rope_dims,
5602            n_head,
5603            1,
5604            geometry.rope_base,
5605            1.0,
5606        )?;
5607        e.rope_neox(
5608            &mut k,
5609            pos_d,
5610            head_dim,
5611            rope_dims,
5612            n_head_kv,
5613            1,
5614            geometry.rope_base,
5615            1.0,
5616        )?;
5617
5618        let kv = scratch.plane_mut(scratch_index).0;
5619        // append at the DEVICE slot (kv.len_d == old len), then advance the counter in-graph.
5620        e.append_kv_quantized_dc(
5621            &k,
5622            &v,
5623            &mut kv.k,
5624            &mut kv.v,
5625            &kv.len_d,
5626            kv.kv_dim_k,
5627            kv.kv_dim_v,
5628            kv.k_tok_bytes,
5629            kv.v_tok_bytes,
5630            false,
5631        )?;
5632        e.inc_seqlen(&mut kv.len_d)?;
5633        // full-buffer views (any in-round t_kv stays in range on replay); the kernel bounds the
5634        // key range from the device counter.
5635        let k_view = e.view_u8(&kv.k, kv.k.len());
5636        let v_view = e.view_u8(&kv.v, kv.v.len());
5637        let (ktb, vtb) = (kv.k_tok_bytes, kv.v_tok_bytes);
5638        let mut attn = e.zeros(n_head * head_dim)?;
5639        e.fa_decode_dc(
5640            &q, &k_view, &v_view, &mut attn, head_dim, n_head, n_head_kv, &kv.len_d, bucket_max,
5641            scale, ktb, vtb, false,
5642        )?;
5643
5644        let attn_g = match &gate {
5645            Some(gate) => {
5646                let mut gsig = e.zeros(n_head * head_dim)?;
5647                e.sigmoid(gate, &mut gsig, n_head * head_dim)?;
5648                let mut ag = e.zeros(n_head * head_dim)?;
5649                e.mul(&attn, &gsig, &mut ag, n_head * head_dim)?;
5650                ag
5651            }
5652            None => attn,
5653        };
5654        e.matmul(&fa.wo, &attn_g, 1)
5655    }
5656
5657    /// PERSISTENT-DRAFT-KV fill (the reference engine's "mtp_update" analogue): compute the MTP
5658    /// block's K/V for `tokens` (committed tokens at positions pos0..pos0+T) from their EXACT
5659    /// trunk hiddens `h` ([T, n_embd] token-major, pre-output_norm) and append at slots pos0.. of
5660    /// the scratch KV. K/V-ONLY — ops A/1-5 plus the K-side of op 6 (wk/wv + k_norm + rope +
5661    /// quantized append); no wq/attention/FFN/lm_head, so per-token cost ~= eh_proj + wk/wv (a
5662    /// small fraction of one trunk layer), T-batched. Rope follows the chain convention
5663    /// rope(token@p) = p+1. Runs at round boundaries OUTSIDE the captured graph in BOTH draft
5664    /// modes -> draft parity by construction. Caller must have scratch.kv.len == pos0.
5665    #[allow(clippy::too_many_arguments)]
5666    fn mtp_kv_fill_at(
5667        &self,
5668        e: &Engine,
5669        mtp: &MtpHead,
5670        tokens: &[u32],
5671        h: &CudaSlice<f32>,
5672        pos0: usize,
5673        scratch: &mut MtpScratch,
5674        scratch_index: usize,
5675        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
5676    ) -> Result<(), Box<dyn std::error::Error>> {
5677        let cfg = &self.cfg;
5678        let n_embd = cfg.n_embd as usize;
5679        let eps = cfg.rms_eps;
5680        let t = tokens.len();
5681        let (scratch_kv, scratch_cap) = scratch.plane(scratch_index);
5682        assert_eq!(scratch_kv.len, pos0, "mtp_kv_fill: append slot mismatch");
5683        assert!(pos0 + t <= scratch_cap, "mtp_kv_fill: scratch overflow");
5684        let Mixer::Full(fa) = &mtp.mixer else {
5685            panic!("MTP block is full-attn in qwen35; linear MTP not supported")
5686        };
5687        let pos_vec: Vec<i32> = (0..t).map(|i| (pos0 + i + 1) as i32).collect();
5688        let pos_d = e.htod_i32(&pos_vec)?;
5689
5690        // ops A/1/2: embed + the two input norms, T-wide.
5691        let e_emb = match embd_dev {
5692            Some((g, qt, rb)) => e.embed_gather_device_t(g, tokens, n_embd, qt, rb)?,
5693            None => e.htod(&self.embd.try_gather(n_embd, tokens)?)?,
5694        };
5695        let mut e_norm = e.zeros(t * n_embd)?;
5696        e.rms_norm(&e_emb, mtp.enorm.float_data(), &mut e_norm, n_embd, t, eps)?;
5697        let mut h_norm = e.zeros(t * n_embd)?;
5698        e.rms_norm(h, mtp.hnorm.float_data(), &mut h_norm, n_embd, t, eps)?;
5699
5700        // op 3: per-row [e_norm ; h_norm] concat, token-major [T, 2*n_embd].
5701        let mut concat = e.zeros(t * 2 * n_embd)?;
5702        for i in 0..t {
5703            e.copy_view_into(
5704                &mut concat,
5705                i * 2 * n_embd,
5706                &e_norm.slice(i * n_embd..(i + 1) * n_embd),
5707                n_embd,
5708            )?;
5709            e.copy_view_into(
5710                &mut concat,
5711                i * 2 * n_embd + n_embd,
5712                &h_norm.slice(i * n_embd..(i + 1) * n_embd),
5713                n_embd,
5714            )?;
5715        }
5716
5717        // ops 4/5: eh_proj + attn_norm, T-wide (at the student inner width when geom is set).
5718        let di = mtp.geom.as_ref().map(|g| g.d_inner).unwrap_or(n_embd);
5719        let inp_sa = e.matmul(&mtp.eh_proj, &concat, t)?;
5720        let mut a_norm = e.zeros(t * di)?;
5721        e.rms_norm(&inp_sa, mtp.attn_norm.float_data(), &mut a_norm, di, t, eps)?;
5722
5723        // op 6 (K/V half): wk/wv + k_norm + rope + per-row quantized append. No wq/attention —
5724        // the fill only has to leave correct K/V rows behind for later chains to attend over.
5725        let n_head_kv = mtp
5726            .geom
5727            .as_ref()
5728            .map(|g| g.n_head_kv)
5729            .or(mtp.step35.as_ref().map(|s| s.n_head_kv))
5730            .unwrap_or_else(|| {
5731                let mtp_il = cfg.n_layer.saturating_sub(cfg.nextn_predict_layers);
5732                cfg.full_attention_geometry_at(mtp_il).n_head_kv as usize
5733            });
5734        let mtp_il = cfg.n_layer.saturating_sub(cfg.nextn_predict_layers);
5735        let geometry = cfg.full_attention_geometry_at(mtp_il);
5736        let head_dim = geometry.head_dim_k as usize;
5737        let mut k = e.matmul(&fa.wk, &a_norm, t)?;
5738        let v = e.matmul(&fa.wv, &a_norm, t)?;
5739        let mut kn = e.zeros(t * n_head_kv * head_dim)?;
5740        e.rms_norm(
5741            &k,
5742            fa.k_norm.float_data(),
5743            &mut kn,
5744            head_dim,
5745            n_head_kv * t,
5746            eps,
5747        )?;
5748        k = kn;
5749        // step35: rotary width AND base are per-layer, and the MTP block's values come from the
5750        // resolved `Step35MtpGeom` — NOT from `cfg.rope_dim_count`/`cfg.rope_freq_base`, which
5751        // carry the arch defaults (128 / 5e6, i.e. the FULL-attn layers' base). Getting this wrong
5752        // writes K rows the attention arm then re-derives at a different theta: correct-looking
5753        // output with dead acceptance, invisible to the exactness gates.
5754        let (rope_dims, rope_base, ff) = match mtp.step35.as_ref() {
5755            Some(s) => (
5756                s.n_rot,
5757                s.rope_base,
5758                if s.swa {
5759                    None
5760                } else {
5761                    self.step35_aux.as_ref().and_then(|a| a.rope_freqs(e))
5762                },
5763            ),
5764            None => (geometry.n_rot as usize, geometry.rope_base, None),
5765        };
5766        #[cfg(debug_assertions)]
5767        if let Some(ff) = ff {
5768            crate::debug_assert_tensor_stream_device(ff, &e.stream(), "mtp_kv_fill.rope_freqs");
5769        }
5770        match ff {
5771            Some(f) => e.rope_neox_ff(
5772                &mut k, &pos_d, head_dim, rope_dims, n_head_kv, t, rope_base, 1.0, f,
5773            )?,
5774            None => e.rope_neox(
5775                &mut k, &pos_d, head_dim, rope_dims, n_head_kv, t, rope_base, 1.0,
5776            )?,
5777        }
5778
5779        let kv = scratch.plane_mut(scratch_index).0;
5780        // Match the trunk prime contract: a chunk may need the aligned window immediately before
5781        // its first row, so preserve that prefix when the physical tail rebases at wrap.
5782        let retain_from = kv
5783            .ring
5784            .as_ref()
5785            .map(|ring| memra_kv::swa_retain_from(pos0, ring.window(), ring.base()))
5786            .unwrap_or(0);
5787        let write_row = e.prepare_kv_append(kv, retain_from, t)?;
5788        for i in 0..t {
5789            let k_row = k.slice(i * kv.kv_dim_k..(i + 1) * kv.kv_dim_k);
5790            let v_row = v.slice(i * kv.kv_dim_v..(i + 1) * kv.kv_dim_v);
5791            e.append_kv_quantized_view(
5792                &k_row,
5793                &v_row,
5794                &mut kv.k,
5795                &mut kv.v,
5796                write_row + i,
5797                kv.kv_dim_k,
5798                kv.kv_dim_v,
5799                kv.k_tok_bytes,
5800                kv.v_tok_bytes,
5801                false,
5802            )?;
5803        }
5804        kv.len = pos0 + t;
5805        e.set_i32_one(&mut kv.len_d, kv.len as i32)?;
5806        Ok(())
5807    }
5808
5809    #[allow(clippy::too_many_arguments)]
5810    fn mtp_kv_fill_all(
5811        &self,
5812        e: &Engine,
5813        tokens: &[u32],
5814        h: &CudaSlice<f32>,
5815        pos0: usize,
5816        scratch: &mut MtpScratch,
5817        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
5818    ) -> Result<(), Box<dyn std::error::Error>> {
5819        debug_assert_eq!(self.mtp_head_count(), scratch.plane_count());
5820        for index in 0..self.mtp_head_count() {
5821            self.mtp_kv_fill_at(
5822                e,
5823                self.mtp_head_at(index),
5824                tokens,
5825                h,
5826                pos0,
5827                scratch,
5828                index,
5829                embd_dev,
5830            )?;
5831        }
5832        Ok(())
5833    }
5834
5835    /// CAPTURE body for the GRAPH DRAFT (stage 2 of graph-grade spec): ONE MTP head forward with
5836    /// every varying input device-resident —
5837    ///   - token id from the persistent `tok_d` (the previous replay's in-graph argmax wrote it,
5838    ///     so the chain feeds itself; the host reads the same 4 bytes for the draft list),
5839    ///   - h_seed from the persistent `h_seed_d` (h_nextn is copied BACK into it at the end),
5840    ///   - rope pos from the persistent `pos_d` counter (inc'd in-graph),
5841    ///   - scratch KV slot/bound from `scratch.kv.len_d` (see mtp_full_attn_dc).
5842    ///     The p-min confidence lands in the persistent `p_d` iff `with_prob` (env is fixed per run).
5843    ///     Same kernels, same dispatch as the eager mtp_head_forward_dev chain -> same draft tokens
5844    ///     (exactness never depends on drafts — the verify arbitrates — but acceptance parity does).
5845    ///     `with_head=false` captures the HEAD-LESS twin for the pseudo-seed replay (2026-07-03):
5846    ///     the pseudo pass only needs h_nextn (op 10) + the scratch append — the lm_head read
5847    ///     (~1.06ms q6_K on the 9B), argmax and prob are dead weight there. h_nextn's inputs are
5848    ///     untouched, so the seed value is identical; round-start resets overwrite tok_d/p_d anyway.
5849    ///     `sampled_cap` = Some((ctr_d, perturb_d, q_out_d, seed, temp)) captures the SAMPLED twin
5850    ///     (step 3 of the sampled-spec arc): head logits are retained in the persistent `q_out_d`
5851    ///     (host D2Ds them to the round's q slot after each replay), the DEVICE event counter is
5852    ///     bumped in-graph, and the argmax reads GUMBEL-PERTURBED logits — one categorical draw per
5853    ///     replay, bit-identical to the eager arm's gumbel_perturb at the same (seed, sctr, temp).
5854    ///     seed/temp are capture-time constants (fixed per generate call, like p_min).
5855    #[allow(clippy::too_many_arguments)]
5856    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
5857    fn mtp_head_forward_cap(
5858        &self,
5859        e: &Engine,
5860        mtp: &MtpHead,
5861        tok_d: &mut CudaSlice<u32>,
5862        pos_d: &mut CudaSlice<i32>,
5863        h_seed_d: &mut CudaSlice<f32>,
5864        p_d: &mut CudaSlice<f32>,
5865        scratch: &mut MtpScratch,
5866        // Which scratch plane this head appends to / attends over: 0 for the single-head
5867        // chain (every pre-lane caller), the head's own plane index for the multi-head
5868        // chain graphs (each head owns one plane — `mtp_chain_forward_dev`'s contract).
5869        scratch_index: usize,
5870        with_prob: bool,
5871        with_head: bool,
5872        embd_gpu: &CudaSlice<u8>,
5873        embd_qt: i32,
5874        embd_rb: usize,
5875        d_vocab: usize,
5876        sampled_cap: Option<SampledCapArgs<'_>>,
5877        stream_pack: Option<(&mut CudaSlice<u32>, usize, Option<&CudaSlice<u32>>)>,
5878        // DRAFT-SIDE GRAMMAR MASK (lane/draft-mask): (packed draft-vocab allowed-set buffer,
5879        // word count). Captured as ONE mask_logits_f32 node between the head matmul and the
5880        // in-graph argmax; the buffer address is baked, its CONTENTS are re-uploaded by the
5881        // host before every replay (the decode.rs graph-mask pattern). All-ones contents = a
5882        // no-op ban, so a position the grammar cannot constrain costs one pass over the row.
5883        mask_cap: Option<(&CudaSlice<u32>, usize)>,
5884    ) -> Result<(), Box<dyn std::error::Error>> {
5885        let cfg = &self.cfg;
5886        let n_embd = cfg.n_embd as usize;
5887        // step35: capturable through the WINDOWED device-counter arm (`mtp_step35_attn_dcw`)
5888        // once the dcw door is armed and the v3-vec class is live. Without the door this stays
5889        // the deliberate, named refusal: the plain `_dc` attention's key bound always starts at
5890        // row 0, cannot express this block's SWA view offset, and a captured chain would
5891        // silently attend OUTSIDE the window once the persistent scratch passes 512 rows.
5892        // Returning Err (not a panic) is what the capture sites already handle by degrading to
5893        // the eager chain (`mtp_head_forward_dev` -> `mtp_step35_attn`).
5894        // ROUND-STREAM stays refused EITHER WAY: the stream VERIFY has no step35 twin (see the
5895        // step35_verify refusal), so a stream capture that succeeded here would only move the
5896        // failure from capture time (graceful stream-off) to serve time (a failed round).
5897        if let Some(g) = mtp.step35.as_ref() {
5898            if stream_pack.is_some() {
5899                return Err(
5900                    "step35 has no ROUND-STREAM draft arm (the stream verify has no step35 \
5901                     twin); stream off"
5902                        .into(),
5903                );
5904            }
5905            if !self.step35_dcw_eligible(g, scratch.plane(scratch_index).1) {
5906                return Err(format!(
5907                    "step35 has no captured draft chain (fa_decode_dc cannot express the MTP \
5908                        block's SWA view offset; the windowed dcw capture needs \
5909                        MEMRA_STEP35_DRAFT_DCW armed [default ON, =0 disarms] and the v3-vec \
5910                        class live at bucket=min(window {}, scratch cap {})) - the eager draft \
5911                        chain serves this shape",
5912                    g.window,
5913                    scratch.plane(scratch_index).1,
5914                )
5915                .into());
5916            }
5917        }
5918        // student inner width (see mtp_head_forward_dev) — interface dims stay n_embd.
5919        let di = mtp.geom.as_ref().map(|g| g.d_inner).unwrap_or(n_embd);
5920        let eps = cfg.rms_eps;
5921        let e_emb = e.embed_gather_device(embd_gpu, tok_d, n_embd, embd_qt, embd_rb)?;
5922        let mut e_norm = e.zeros(n_embd)?;
5923        e.rms_norm(&e_emb, mtp.enorm.float_data(), &mut e_norm, n_embd, 1, eps)?;
5924        let mut h_norm = e.zeros(n_embd)?;
5925        e.rms_norm(
5926            &*h_seed_d,
5927            mtp.hnorm.float_data(),
5928            &mut h_norm,
5929            n_embd,
5930            1,
5931            eps,
5932        )?;
5933        let mut concat = e.zeros(2 * n_embd)?;
5934        e.copy_into(&mut concat, 0, &e_norm, n_embd)?;
5935        e.copy_into(&mut concat, n_embd, &h_norm, n_embd)?;
5936        let inp_sa = e.matmul(&mtp.eh_proj, &concat, 1)?;
5937        let mut a_norm = e.zeros(di)?;
5938        e.rms_norm(&inp_sa, mtp.attn_norm.float_data(), &mut a_norm, di, 1, eps)?;
5939        let attn_out = match (&mtp.mixer, mtp.step35.as_ref()) {
5940            // step35 (eligibility already enforced by the refusal above): the windowed dcw
5941            // arm, the SAME launcher the eager dev arm runs when the door is armed. No host
5942            // work here (this is the capture body); headroom is the callers' pre-arm.
5943            (Mixer::Full(fa), Some(g)) => {
5944                self.mtp_step35_attn_dcw(e, fa, g, &a_norm, pos_d, scratch, scratch_index)?
5945            }
5946            (Mixer::Full(fa), None) => self.mtp_full_attn_dc(
5947                e,
5948                fa,
5949                &a_norm,
5950                pos_d,
5951                scratch,
5952                scratch_index,
5953                mtp.geom.as_ref(),
5954            )?,
5955            (Mixer::Linear(_), _) => {
5956                panic!("MTP block is full-attn in qwen35; linear MTP not supported")
5957            }
5958            (Mixer::Mla(_), _) => {
5959                crate::hybrid::mla_path_unimplemented("captured MTP head forward")
5960            }
5961            (Mixer::Kda(_), _) => {
5962                crate::hybrid::kda_path_unimplemented("captured MTP head forward")
5963            }
5964        };
5965        let mut x1 = e.zeros(di)?;
5966        e.add(&inp_sa, &attn_out, &mut x1, di)?;
5967        let mut z = e.zeros(di)?;
5968        e.rms_norm(&x1, mtp.post_attn_norm.float_data(), &mut z, di, 1, eps)?;
5969        let ffn_out = match &mtp.ffn {
5970            crate::hybrid::Ffn::Dense {
5971                ffn_gate,
5972                ffn_up,
5973                ffn_down,
5974            } => {
5975                let n_ff = ffn_gate.out_features();
5976                let (gate, up) = if e.uses_q8_1_fast(ffn_gate) && e.uses_q8_1_fast(ffn_up) {
5977                    let (zq, zd) = e.quantize_q8_1(&z, 1, di)?;
5978                    (
5979                        e.matmul_pre(ffn_gate, &zq, &zd, &z, 1)?,
5980                        e.matmul_pre(ffn_up, &zq, &zd, &z, 1)?,
5981                    )
5982                } else {
5983                    (e.matmul(ffn_gate, &z, 1)?, e.matmul(ffn_up, &z, 1)?)
5984                };
5985                let mut act = e.zeros(n_ff)?;
5986                // step35: the dense FFN reads the per-layer SHEXP clamp, resolved for the MTP
5987                // block's own index (the mtp_head_forward_dev rule; None for every other arch,
5988                // which is `ffn_act`'s dispatch verbatim). The eager and captured chains must
5989                // run the ONE activation program.
5990                Self::ffn_act_lim(
5991                    e,
5992                    &self.cfg,
5993                    &gate,
5994                    &up,
5995                    1.0,
5996                    1.0,
5997                    mtp.step35
5998                        .as_ref()
5999                        .and_then(|s| s.clamp_shexp)
6000                        .map(SwigluClamp::Post),
6001                    &mut act,
6002                    n_ff,
6003                )?;
6004                e.matmul(ffn_down, &act, 1)?
6005            }
6006            // ROUND-STREAM: a softmax-routed resident MoE takes the zero-D2H device router +
6007            // expert program and is capture-legal. Sigmoid-routed MoE (Hy3/M3/Step) still
6008            // selects through the host-visible sigmoid router; capturing that stream sync
6009            // invalidates CUDA capture, so it stays on the eager draft chain even when every
6010            // expert is resident. Non-resident (SLRU-lock) is likewise rejected.
6011            crate::hybrid::Ffn::Moe(m)
6012                if m.dev_exps.is_some() && self.cfg.sigmoid_router().is_none() =>
6013            {
6014                self.moe_ffn_il(e, m, &z, 1, u16::MAX)?
6015            }
6016            crate::hybrid::Ffn::Moe(_) => {
6017                return Err(
6018                    "graph draft requires a Dense or device-routed resident-MoE MTP FFN".into(),
6019                );
6020            }
6021        };
6022        let mut h_inner = e.zeros(di)?;
6023        e.add(&x1, &ffn_out, &mut h_inner, di)?;
6024        // student: up-project back to n_embd (carrier + head input; see mtp_head_forward_dev).
6025        let h_nextn = match mtp.geom.as_ref() {
6026            Some(g) => e.matmul(&g.out_up, &h_inner, 1)?,
6027            None => h_inner,
6028        };
6029        // MEMRA_SPEC_HPOST needs final_h even head-less (it IS the next seed under that convention).
6030        let final_h = if with_head || spec_hpost() {
6031            let final_norm = mtp.shared_head_norm.as_ref().unwrap_or(&self.output_norm);
6032            let mut fh = e.zeros(n_embd)?;
6033            e.rms_norm(&h_nextn, final_norm.float_data(), &mut fh, n_embd, 1, eps)?;
6034            Some(fh)
6035        } else {
6036            None
6037        };
6038        if with_head {
6039            let head = mtp.shared_head_head.as_ref().unwrap_or(&self.output);
6040            let mut logits = e.matmul(head, final_h.as_ref().unwrap(), 1)?;
6041            // DRAFT-SIDE GRAMMAR MASK: ban the grammar-illegal draft ids IN the captured chain,
6042            // before the argmax — proposals become legal by construction. Contents-only
6043            // per-replay upload keeps the capture valid.
6044            if let Some((mask_d, mw)) = mask_cap {
6045                e.mask_logits_col(&mut logits, mask_d, 0, d_vocab, mw)?;
6046            }
6047            if let Some(SampledCapArgs {
6048                ctr: ctr_d,
6049                perturb: perturb_d,
6050                q_out: q_out_d,
6051                seed,
6052                temp,
6053                filt,
6054            }) = sampled_cap
6055            {
6056                // SAMPLED chain: retain q (raw head logits -> persistent q_out_d; the matmul's
6057                // own buffer is pool-recycled after the capture body returns, so it can't be the
6058                // retention target), bump the device event counter, gumbel-perturb reading it,
6059                // and argmax the PERTURBED logits into tok_d — the in-graph categorical draw.
6060                e.copy_into(q_out_d, 0, &logits, d_vocab)?;
6061                e.sctr_inc(ctr_d)?;
6062                match filt {
6063                    // PURE-TEMP: gumbel over the raw softmax — byte-identical to the
6064                    // pre-lane capture body.
6065                    None => e.gumbel_perturb_ctr(&logits, perturb_d, d_vocab, seed, ctr_d, temp)?,
6066                    // FILTERED (lane/step37-draft-graph-serving-20260830): the SAME
6067                    // filter_stats program the eager arm and the accept path run (the
6068                    // wrapper's coop/plain choice is deployment-keyed, never per-call), then
6069                    // the device-stat/device-counter perturb twin — the draft draws from the
6070                    // exact filtered distribution the verify gathers `q` from. q was
6071                    // retained ABOVE, pre-perturb, so the accept path's post-replay stats
6072                    // recompute (same kernel, same bits) reconstructs these th/z exactly.
6073                    Some(f) => {
6074                        e.filter_stats(
6075                            &logits, d_vocab, f.rows0, f.th, f.z, f.mx, d_vocab, 1, temp, f.top_k,
6076                            f.top_p, f.min_p,
6077                        )?;
6078                        e.gumbel_perturb_filtered_ctr(
6079                            &logits, perturb_d, d_vocab, seed, ctr_d, temp, f.mx, f.th,
6080                        )?;
6081                    }
6082                }
6083                e.argmax_token_device_into(perturb_d, tok_d, d_vocab)?;
6084                // p-min prob = the head's RAW softmax confidence in the SAMPLED pick — same
6085                // semantics as the eager sampled arm's prob_of_token_device(dl_d, tok_d).
6086                if with_prob {
6087                    e.prob_of_token_device_into(&logits, tok_d, p_d, d_vocab)?;
6088                }
6089            } else {
6090                // draft token -> persistent tok_d (next replay's embed reads it; host reads the 4 bytes).
6091                e.argmax_token_device_into(&logits, tok_d, d_vocab)?;
6092                // p-min under a draft mask reads the MASKED row: confidence relative to the
6093                // grammar-LEGAL alternatives (illegal ids leave the softmax denominator), which
6094                // is the right semantics for "does the drafter know what comes next here" and
6095                // the same row the pick came from. Draft-quality only — verify arbitrates.
6096                if with_prob {
6097                    e.prob_of_token_device_into(&logits, tok_d, p_d, d_vocab)?;
6098                }
6099            }
6100        }
6101        // ROUND-STREAM K-chain: pack (tok, p) into slot j, then remap tok through d2t so the
6102        // NEXT chained body's embed reads the TARGET id — zero host involvement per step.
6103        if let Some((out, slot, d2t)) = stream_pack {
6104            e.pack_tok_p(tok_d, p_d, out, slot)?;
6105            if let Some(map) = d2t {
6106                e.tok_map_u32(tok_d, map)?;
6107            }
6108        }
6109        // Next draft step's h_seed: pre-norm h_nextn (default) or post-norm final_h (HPOST).
6110        if spec_hpost() {
6111            e.copy_into(h_seed_d, 0, final_h.as_ref().unwrap(), n_embd)?;
6112        } else {
6113            e.copy_into(h_seed_d, 0, &h_nextn, n_embd)?;
6114        }
6115        // advance the draft rope position in-graph.
6116        e.inc_seqlen(pos_d)?;
6117        Ok(())
6118    }
6119
6120    /// Batched target verify forward over `tokens` at positions `pos0..pos0+T` (§D.3, T=K+1).
6121    /// Returns ALL T logit columns (host f32, [T*n_vocab]); appends T cols to every full-attn KV
6122    /// and advances every linear-attn recur state by T steps (the recur steps are SEQUENTIAL T=1).
6123    /// Advances `cache.pos` by T.
6124    pub fn decode_step_t(
6125        &self,
6126        e: &Engine,
6127        tokens: &[u32],
6128        pos0: usize,
6129        cache: &mut Cache,
6130    ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
6131        if self.is_gemma4_e4b() {
6132            return Ok(self.gemma4_e4b_decode_step_t_h(e, tokens, pos0, cache)?.0);
6133        }
6134        if self.gemma_batch_program() {
6135            return self.gemma4_decode_step_t(e, tokens, pos0, cache);
6136        }
6137        Ok(self.decode_step_t_h(e, tokens, pos0, cache)?.0)
6138    }
6139
6140    /// Like `decode_step_t` but ALSO returns the LAST column's pre-output_norm hidden (h_seed for
6141    /// the next draft round). This lets partial-accept replay run as ONE batched T=(n_acc+1) forward
6142    /// (single weight read) instead of n_acc+1 separate T=1 decode_steps (n_acc+1 weight reads).
6143    /// At batch=1 decode is bandwidth-bound, so batching the replay is THE MTP profitability lever.
6144    pub fn decode_step_t_h(
6145        &self,
6146        e: &Engine,
6147        tokens: &[u32],
6148        pos0: usize,
6149        cache: &mut Cache,
6150    ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
6151        self.decode_step_t_h_emb(e, tokens, pos0, cache, None)
6152    }
6153
6154    /// Like `decode_step_t_h` with an optional RESIDENT embed table (spec hot loop): device
6155    /// gather instead of host dequant + [T, n_embd] f32 htod. Bit-identical rows.
6156    pub fn decode_step_t_h_emb(
6157        &self,
6158        e: &Engine,
6159        tokens: &[u32],
6160        pos0: usize,
6161        cache: &mut Cache,
6162        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
6163    ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
6164        let (logits_d, h_seed) = self.decode_step_t_h_emb_dev(e, tokens, pos0, cache, embd_dev)?;
6165        Ok((e.dtoh(&logits_d)?, h_seed))
6166    }
6167
6168    /// DEVICE-LOGITS verify forward (spec device-argmax lever): identical kernel chain to
6169    /// `decode_step_t_h_emb` but returns the [T, n_vocab] logits ON DEVICE — the accept walk
6170    /// argmaxes each column on-device and reads back ONE [T] u32 instead of dtoh'ing the full
6171    /// T x n_vocab f32 block (~1-4 MB + T host argmaxes, every round). Kernel dispatch is
6172    /// UNCHANGED (same decode-exact kernels); only the post-logits transfer moves.
6173    pub fn decode_step_t_h_emb_dev(
6174        &self,
6175        e: &Engine,
6176        tokens: &[u32],
6177        pos0: usize,
6178        cache: &mut Cache,
6179        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
6180    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
6181        cache.ensure_usable("decode_step_t")?;
6182        let n_embd = self.cfg.n_embd as usize;
6183        let t = tokens.len();
6184        let (logits, x) = self.decode_step_t_core(e, tokens, pos0, cache, embd_dev, None)?;
6185        // h_seed for the next round = LAST column's pre-output_norm hidden ([n_embd]).
6186        let mut hs = vbuf(e, n_embd)?; // fully written by copy_view_into below
6187        e.copy_view_into(&mut hs, 0, &x.slice((t - 1) * n_embd..t * n_embd), n_embd)?;
6188        Ok((logits, hs))
6189    }
6190
6191    /// CORE verify forward: the `decode_step_t_h_emb_dev` kernel chain, returning the FULL
6192    /// pre-output_norm hidden stack x ([T, n_embd], any column extractable) and optionally
6193    /// filling a `VerifyCkpt` (retained per-layer state-rebuild inputs) for the REPLAY-FREE
6194    /// partial accept. `ckpt: None` => byte-for-byte the old behavior (the ckpt writes are pure
6195    /// retains/copies — they never change what any kernel computes).
6196    fn decode_step_t_core(
6197        &self,
6198        e: &Engine,
6199        tokens: &[u32],
6200        pos0: usize,
6201        cache: &mut Cache,
6202        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
6203        mut ckpt: Option<&mut VerifyCkpt>,
6204    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
6205        self.decode_step_t_core_stream(
6206            e,
6207            tokens,
6208            pos0,
6209            cache,
6210            embd_dev,
6211            ckpt.take(),
6212            None,
6213            None,
6214            None,
6215            None,
6216        )
6217    }
6218
6219    /// [`Self::decode_step_t_core`] with the MTP route's verify-graph pool armed
6220    /// (`MEMRA_SPEC_VERIFY_GRAPH`). `graphs: None` reproduces `decode_step_t_core`
6221    /// argument-for-argument, so the eager walk stays the byte-identical fallback.
6222    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
6223    fn decode_step_t_core_vg(
6224        &self,
6225        e: &Engine,
6226        tokens: &[u32],
6227        pos0: usize,
6228        cache: &mut Cache,
6229        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
6230        mut ckpt: Option<&mut VerifyCkpt>,
6231        graphs: Option<&mut DsparkVerifyGraphs>,
6232    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
6233        self.decode_step_t_core_stream(
6234            e,
6235            tokens,
6236            pos0,
6237            cache,
6238            embd_dev,
6239            ckpt.take(),
6240            None,
6241            None,
6242            None,
6243            graphs,
6244        )
6245    }
6246
6247    /// Increment-0 two-session PP seam: release the peer after this lane's stage-0 boundary TX.
6248    /// The two independent sessions keep their own cache/checkpoint state; only issue order moves.
6249    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
6250    fn decode_step_t_core_pipelined(
6251        &self,
6252        e: &Engine,
6253        tokens: &[u32],
6254        pos0: usize,
6255        cache: &mut Cache,
6256        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
6257        mut ckpt: Option<&mut VerifyCkpt>,
6258        pipe: &SpecPipeLane,
6259        round: usize,
6260    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
6261        let fence = crate::pp::pp_cuts(self.layers.len())
6262            .ok_or("two-session speculative pipeline requires a PP stage cut")?;
6263        if crate::pp::pp2_streams_off() || !crate::pp::spec_pp_on() {
6264            return Err("two-session speculative pipeline requires the PP verify split".into());
6265        }
6266        let interval_fence = pipe.stage0_begin(round)?;
6267        let _walk = pipe.coordinated_walk()?;
6268        let ticket = self.verify_stage0_issue(
6269            e,
6270            tokens,
6271            pos0,
6272            cache,
6273            embd_dev,
6274            ckpt.as_deref_mut(),
6275            None,
6276            &fence,
6277            Some(interval_fence),
6278            pipe.trace(round),
6279        )?;
6280        pipe.stage0_end(round);
6281        pipe.stage1_begin(round)?;
6282        let result = self.verify_stage1_finish(e, ticket, cache, ckpt, None, &fence, true)?;
6283        pipe.verify_end(round);
6284        Ok(result)
6285    }
6286
6287    /// ROUND-STREAM stage (c) 4: `stream` = (device verify tokens [t], device pos counter) —
6288    /// when Some, rope positions come from pos_iota over the counter, the embed gathers the
6289    /// device tokens, and full_attn_verify routes appends/FA through the _dc twins reading the
6290    /// SAME counter (every layer's kvl.len == cache.pos, one counter drives all three). The
6291    /// host `tokens`/`pos0` args still size buffers (t is FIXED K+1 in stream mode).
6292    /// `vtok_dev` (engine-bundle slice 2): device verify tokens for the EMBED only —
6293    /// unlike `stream` mode it changes nothing else (host pos iota, host-len KV appends).
6294    /// `tokens` then only sizes buffers (the dummy-slice pattern the round-stream arm uses).
6295    #[allow(clippy::too_many_arguments)]
6296    fn decode_step_t_core_stream(
6297        &self,
6298        e: &Engine,
6299        tokens: &[u32],
6300        pos0: usize,
6301        cache: &mut Cache,
6302        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
6303        mut ckpt: Option<&mut VerifyCkpt>,
6304        stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
6305        pp_pipe: Option<bool>,
6306        vtok_dev: Option<&CudaSlice<u32>>,
6307        graphs: Option<&mut DsparkVerifyGraphs>,
6308    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
6309        // PP DOOR (lane/pp2-spec 2026-08-06): the verify trunk now takes its OWN stage split,
6310        // exactly as the eager and batched steps do. This is the single funnel every verify
6311        // forward reaches (decode_step_t / _h / _h_emb / _h_emb_dev / _core all land here), so
6312        // wiring it here wires the whole spec surface — the draft/accept/commit machinery above
6313        // is untouched.
6314        //
6315        // History: pp2-hardening (2026-08-06) made this funnel FAIL CLOSED, because its trunk
6316        // walk was unsplit on one stream and a sharded cross-device placement peer-read every
6317        // remote layer's weights on every spec round (measured 13.9-28x on the batched twin).
6318        // The refusal below survives to cover the residue — MEMRA_SPEC_PP=0, MEMRA_PP_STREAMS=0,
6319        // or a placement whose PpNRt fails to build — so a config that would still walk the
6320        // whole trunk on one stream refuses instead of regressing 28x.
6321        if let Some(fence) = crate::pp::pp_cuts(self.layers.len())
6322            && !crate::pp::pp2_streams_off()
6323            && crate::pp::spec_pp_on()
6324        {
6325            if vtok_dev.is_some() {
6326                return Err(
6327                    "device-token dspark verify (slice-2 deferred readback) has no PP \
6328                         stage-split arm; set MEMRA_DSPARK_DEFER_READBACK=0 or run the dspark \
6329                         route on one device"
6330                        .into(),
6331                );
6332            }
6333            return self.decode_step_t_core_ppn(
6334                e,
6335                tokens,
6336                pos0,
6337                cache,
6338                embd_dev,
6339                ckpt.take(),
6340                stream,
6341                &fence,
6342                pp_pipe,
6343            );
6344        }
6345        crate::pp::refuse_unsplit_if_remote(
6346            "decode_step_t (spec verify)",
6347            "drop MEMRA_SPEC_PP=0 / MEMRA_PP_STREAMS=0 so the verify trunk takes its OWN stage \
6348             split (decode_step_t_core_ppn); or run spec on one device",
6349        )?;
6350        let cfg = &self.cfg;
6351        let n_embd = cfg.n_embd as usize;
6352        let eps = cfg.rms_eps;
6353        let t = tokens.len();
6354        let pos_d = match stream {
6355            Some((_, ctr)) => {
6356                let mut p = e.alloc_uninit::<i32>(t)?;
6357                e.pos_iota(ctr, &mut p, t)?;
6358                p
6359            }
6360            None => {
6361                let pos_vec: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
6362                e.htod_i32(&pos_vec)?
6363            }
6364        };
6365
6366        // embed T tokens -> [T, n_embd] token-major (device gather on the spec hot loop)
6367        let x = match (stream, embd_dev) {
6368            (Some((vtok, _)), Some((g, qt, rb))) => {
6369                e.embed_gather_device_td(g, vtok, t, n_embd, qt, rb)?
6370            }
6371            (None, Some((g, qt, rb))) => match vtok_dev {
6372                // slice 2: device verify tokens, same embed_gather_u32_t kernel —
6373                // bit-identical rows to the host-token arm (same per-dtype deq).
6374                Some(vt_d) => e.embed_gather_device_td(g, vt_d, t, n_embd, qt, rb)?,
6375                None => e.embed_gather_device_t(g, tokens, n_embd, qt, rb)?,
6376            },
6377            _ => {
6378                assert!(
6379                    vtok_dev.is_none(),
6380                    "device-token verify requires the resident embed table (embd_dev)"
6381                );
6382                e.htod(&self.embd.try_gather(n_embd, tokens)?)?
6383            }
6384        };
6385
6386        // TRUNK WALK: layers [0, n_layers) through the SAME range-scoped subgraph the PP-N
6387        // stage split calls per stage (`verify_layers`) — one code path, so the split cannot
6388        // drift from the unsplit dispatch mirroring. lane/pp2-spec 2026-08-06.
6389        let x = self.verify_layers(
6390            e,
6391            x,
6392            0,
6393            self.layers.len(),
6394            &pos_d,
6395            pos0,
6396            t,
6397            cache,
6398            ckpt.take(),
6399            stream,
6400            graphs,
6401        )?;
6402        if spec_nan_scan() {
6403            nan_scan_rows(e, &x, t, n_embd, &format!("verify trunk exit pos0={pos0}"))?;
6404        }
6405
6406        let mut hn = vbuf(e, t * n_embd)?;
6407        // Stage-A door: with the serving-class row-outer verify walk, the TAIL must be the
6408        // t=1 decode program per row too (rms_norm t=1 + the single-row bf16 head — the
6409        // split head's concat is receipted bit-identical to it). The batched cuBLASLt head
6410        // is a different ULP class and flips near-tie argmaxes off the greedy tape.
6411        let eager_tail = self.sliding_gated_moe_batch_program() && spec_verify_eager_on();
6412        if eager_tail {
6413            let n_vocab = self.cfg.n_vocab as usize;
6414            // MEMRA_SPEC_HEAD_ROWS=1 — THE VERIFY TAIL'S REDUNDANT HEAD READ.
6415            //
6416            // The loop below runs the head at m=1 once PER COLUMN, so the LM head's weights are
6417            // streamed t times per verify pass. On step37 that head is ~0.49 GiB per card after the
6418            // rank split, ~1.07 ms of pure re-read at t=2 and worse at every wider t — which is a
6419            // large part of why the fixed K ladder LOSES (K=1 81.2 > K=2 73.1 > K=3 62.7 tok/s).
6420            //
6421            // The loop's justification is the comment above: the batched cuBLASLt head is a
6422            // different ULP class and flips near-tie argmaxes off the greedy tape. That is true of
6423            // cuBLASLt and it does NOT apply here, because a FloatBf16 head at 1..=32 rows never
6424            // reaches cuBLASLt: `matmul` routes it to `matvec_bf16_rows_into` (lib.rs:12248), whose
6425            // own doc says `matvec_bf16_f32acc_x4_rows` "runs the t=1 decode head program PER ROW
6426            // (identical dot + reduce), so decode/verify tiers keep the t=1 numeric class". Under
6427            // the W8 doors both widths route to the q8 mirror instead, and the t-column mirror is
6428            // documented "bit-identical to t single-row calls". So the batched form is the SAME
6429            // arithmetic per row on both paths, with one weight read instead of t.
6430            //
6431            // rms_norm is row-wise, so norm(t) is per-row identical to t x norm(1) by construction.
6432            //
6433            // DEFAULT OFF for exactly one turn of the crank: "bit-identical by two documented
6434            // claims" is still an argument. The greedy byte tape decides, and the door flips only
6435            // once the tape is a receipt.
6436            if head_rows_on() {
6437                e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
6438                let logits = e.matmul(&self.output, &hn, t)?;
6439                if stream.is_none() {
6440                    cache.pos += t;
6441                }
6442                return Ok((logits, if spec_hpost() { hn } else { x }));
6443            }
6444            let mut logits = vbuf(e, t * n_vocab)?;
6445            for r in 0..t {
6446                let mut row = e.uninit(n_embd)?;
6447                e.dtod_copy_view(&x.slice(r * n_embd..(r + 1) * n_embd), &mut row)?;
6448                let mut hr = e.uninit(n_embd)?;
6449                e.rms_norm(&row, self.output_norm.float_data(), &mut hr, n_embd, 1, eps)?;
6450                let lr = e.matmul(&self.output, &hr, 1)?;
6451                e.dtod_copy_into(&lr, &mut logits, r * n_vocab)?;
6452                e.dtod_copy_into(&hr, &mut hn, r * n_embd)?;
6453            }
6454            if stream.is_none() {
6455                cache.pos += t;
6456            }
6457            return Ok((logits, if spec_hpost() { hn } else { x }));
6458        }
6459        let serving_head =
6460            self.sliding_gated_moe_batch_program() || self.batched_serving_numeric_class();
6461        let logits = if serving_head {
6462            // Step35 and the qwen35 family (MoE 2026-08-14 AM, dense-hybrid same day PM — the
6463            // Q3.8 bring-up reproduced the identical near-tie class on dense: eager-class verify
6464            // vs batched-class live serving, ULP drift amplified through the GDN recurrence)
6465            // serve one batched numeric class at every live width, including B=1. Keep the
6466            // verify head in that same class; other generic families retain the decode-exact
6467            // head that their run-spec contract pins.
6468            e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
6469            e.matmul(&self.output, &hn, t)?
6470        } else {
6471            e.rms_norm_decode(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
6472            e.matmul_decode_exact(&self.output, &hn, t)?
6473        };
6474        // stream: the device pos counter owns position; host mirror reconciles at drain.
6475        if stream.is_none() {
6476            cache.pos += t;
6477        }
6478        // Hidden stack for seeds/refresh-fills: pre-norm x (default) or post-norm hn (HPOST).
6479        Ok((logits, if spec_hpost() { hn } else { x }))
6480    }
6481
6482    /// THE VERIFY TRUNK OVER PP-N (lane/pp2-spec 2026-08-06): `decode_step_t_core_stream`'s walk
6483    /// as N stage subgraphs, each on its own engine/stream (and, under `MEMRA_PP_DEVICES`, its own
6484    /// device), with a `[T, n_embd]` boundary transfer between them. T = K+1 (the verify batch),
6485    /// so this is the batched-boundary shape the pp2-batch lane's grow-only slots already handle
6486    /// (`tx(b, x, t*n_embd)`; the slot grows to the high-water T and the transport moves exactly
6487    /// the payload).
6488    ///
6489    /// Structure is `decode_step_batch_ppn`'s, which is `decode_step_h_ppn`'s. FOUR THINGS ARE
6490    /// PER-STAGE and each for a measured reason (see `decode_step_batch_ppn`'s header for the
6491    /// receipts):
6492    ///
6493    /// 1. THE ENGINE (`rt.engine(s, e)`) — `Engine` owns lazily-grown stable-pointer scratch
6494    ///    (`fa_part_pool`, `fa_vf16_scratch`, `argmax_partials`) that is single-stream-safe BY
6495    ///    DESIGN. Two stage streams through one Engine is the 2026-08-02 shared-scratch race
6496    ///    (35% flake, nondeterministic all-logits divergence). `PpNRt::build` gives every stage
6497    ///    s>0 its own Engine even on the primary device; honouring it here is what scopes the
6498    ///    pools. The verify path allocates MORE of that scratch than eager decode does (FA at
6499    ///    m=T, and the per-layer `GdnStash` retains), so this is load-bearing, not inherited.
6500    ///
6501    /// 2. `pos_d` — each stage uploads its OWN copy of the T rope positions on ITS stream, so the
6502    ///    buffer is allocated, consumed and freed on one stream. In `stream` mode that means each
6503    ///    stage runs its own `pos_iota` over the SHARED device counter (`pos_ctr`): the counter is
6504    ///    read-only during the forward (the round's `inc`/`copy_add` happen outside it), so every
6505    ///    stage derives the identical iota, and each stage's own output buffer is stream-local.
6506    ///
6507    /// 3. THE EMBED lives with stage 0 (`self.embd` / `embd_gpu` are host/primary-side; the
6508    ///    sharded loader leaves the table with stage 0 by construction).
6509    ///
6510    /// 4. THE HEAD (`output_norm` + `output`) runs on the LAST stage — the sharded loader uploaded
6511    ///    both through that stage's engine (`hybrid.rs`: `e_head = layer_engine(e, n_trunk,
6512    ///    n_trunk-1)`), so reading them anywhere else is a peer read of the biggest tensor in the
6513    ///    model, every round.
6514    ///
6515    /// WHAT STAYS ON THE PRIMARY, deliberately: the returned logits and hidden stack `x`. Both are
6516    /// last-stage-allocated device buffers, and every consumer (the device argmax walk, the accept
6517    /// kernels, `spec_seed_gather`, the ckpt rebuild in `commit_verified_prefix`) reads them
6518    /// through the primary context by UVA — the same read the batched serving epilogue's
6519    /// `last_logits_dev` park does. Those consumers are per-round O(T x n_vocab) and O(n_embd),
6520    /// not per-layer, so they are not the 28x class; splitting them is a separate lane.
6521    ///
6522    /// The MTP HEAD (draft side) is NOT split: it is one block, it lives wherever the loader put
6523    /// it (`load_mtp` uses the primary engine), and it is ~1-2 GB against the trunk's tens. Draft
6524    /// placement is measured, not assumed — see `research/pp2-spec-20260806`.
6525    ///
6526    /// EXACTNESS: PP-N adds ZERO deviation. Each stage runs the SAME kernels on the SAME bytes in
6527    /// the same order via the SAME `verify_layers` the unsplit body calls; the only change is
6528    /// where the residual is materialized, and the boundary is a straight f32 copy (dtod
6529    /// same-device / `cudaMemcpyPeerAsync` cross-device, no conversion). So the split MUST be
6530    /// BIT-IDENTICAL to the unsplit verify at the same T, in both placement orders. Gate:
6531    /// `decode-batch-gate --mode ppspec`. Acceptance counts are a DERIVED consequence — greedy
6532    /// accept argmaxes these logits, so bit-identical logits force identical accept walks; the
6533    /// `run-spec` K=1..8 arm checks that end-to-end rather than trusting the implication.
6534    #[allow(clippy::too_many_arguments)]
6535    fn decode_step_t_core_ppn(
6536        &self,
6537        e: &Engine,
6538        tokens: &[u32],
6539        pos0: usize,
6540        cache: &mut Cache,
6541        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
6542        mut ckpt: Option<&mut VerifyCkpt>,
6543        stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
6544        fence: &[usize],
6545        pp_pipe: Option<bool>,
6546    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
6547        let ticket = self.verify_stage0_issue(
6548            e,
6549            tokens,
6550            pos0,
6551            cache,
6552            embd_dev,
6553            ckpt.as_deref_mut(),
6554            stream,
6555            fence,
6556            pp_pipe,
6557            None,
6558        )?;
6559        self.verify_stage1_finish(e, ticket, cache, ckpt, stream, fence, true)
6560    }
6561
6562    /// Enqueue embed, stage 0, and the first boundary TX, then return the actual boundary slot.
6563    /// The ordinary PP verify wrapper calls `verify_stage1_finish` immediately after this return.
6564    #[allow(clippy::too_many_arguments)]
6565    fn verify_stage0_issue(
6566        &self,
6567        e: &Engine,
6568        tokens: &[u32],
6569        pos0: usize,
6570        cache: &mut Cache,
6571        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
6572        ckpt: Option<&mut VerifyCkpt>,
6573        stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
6574        fence: &[usize],
6575        pp_pipe: Option<bool>,
6576        trace: Option<SpecPipeTraceCtx>,
6577    ) -> Result<VerifyBoundaryTicket, Box<dyn std::error::Error>> {
6578        assert!(
6579            !self.is_gemma4_e4b() && !self.gemma_batch_program(),
6580            "decode_step_t_core_ppn covers the hybrid non-gemma4 verify trunk only \
6581             (the gemma4 arms have their own decode_step_t twins)"
6582        );
6583        if crate::pp::pp_host_bounce_active() && (stream.is_some() || embd_dev.is_some()) {
6584            return Err(
6585                "decode_step_t_core_ppn: refused with MEMRA_PP_HOST_BOUNCE=1 — the trunk \
6586                 boundary itself is host-staged, but device-resident verify still peer-reads \
6587                 primary-device token/position/embedding buffers from stage 0. Run plain PP \
6588                 serving on this host class; spec requires local per-stage inputs first."
6589                    .into(),
6590            );
6591        }
6592        let rt = crate::pp::PpNRt::get(e)?;
6593        // Pipelined callers do not bypass ownership: their explicit coordinator borrow makes
6594        // this acquire clone the same active generation. Ordinary callers acquire a fresh lease.
6595        let walk_owner = rt.acquire_walk("verify_stage0_issue")?;
6596        let n_st = fence.len() - 1;
6597        assert_eq!(
6598            rt.n_stages(),
6599            n_st,
6600            "PpNRt stage count {} != fence stages {n_st}",
6601            rt.n_stages()
6602        );
6603        let n_embd = self.cfg.n_embd as usize;
6604        let t = tokens.len();
6605        let payload = t * n_embd;
6606        if pp_pipe.is_some() {
6607            assert_eq!(n_st, 2, "spec pipeline requires exactly two PP stages");
6608        }
6609        // One-shot lane diagnostic: force natural PP-2 boundaries to completion so the server
6610        // log can price stage 0, the peer hop, the RX copy, and stage 1 + head separately without
6611        // nsys. The ordinary path keeps every enqueue asynchronous. N>2 is deliberately excluded:
6612        // the report below names exactly two stages and must never imply it measured middle ones.
6613        let pp_anatomy = n_st == 2 && std::env::var("MEMRA_SPEC_PP_ANATOMY").as_deref() == Ok("1");
6614        let pp_started = std::time::Instant::now();
6615        let (mut reverse_ms, mut stage0_ms, mut tx_ms) = (0.0f64, 0.0f64, 0.0f64);
6616        // The CALLER's ambient stream, captured BEFORE any `rt.enter()` pushes a stage stream:
6617        // this body returns DEVICE-RESIDENT buffers (the device-argmax accept walk's contract),
6618        // so the exit needs the same publication the boundaries get — see `PpNRt::publish_to`.
6619        // Taken here, not at the end, because inside the last-stage scope `e.stream()` IS the
6620        // stage stream and the wait would self-order into a no-op.
6621        let caller_stream = e.stream();
6622        // #87 ROOT-CAUSE FENCE (lane/pp2spec-crash): the PREVIOUS round's stage-allocated
6623        // outputs (logits/hidden/ckpt stashes) freed stream-ordered on the STAGE streams while
6624        // the primary stream still holds queued reads of them — with event tracking elided,
6625        // nothing stops the pool from reusing those blocks for THIS round's stage allocations,
6626        // whose writes then race the queued reads (measured: 13/4096-NaN random-bits garbage in
6627        // the spec round seed; the full anatomy is on `PpNRt::fence_stages_behind`). Order every
6628        // stage stream behind the caller before enqueueing new stage work.
6629        let reverse_started = std::time::Instant::now();
6630        if pp_pipe != Some(false) {
6631            rt.fence_stages_behind(&caller_stream)?;
6632        }
6633        if pp_pipe == Some(true) {
6634            // Both session verifies must alternate boundary slots even when the ordinary
6635            // decode overlap experiment is off. Prewarm before A's stage 0 so B cannot grow
6636            // slot 1 by synchronizing the RX stream while A's stage 1 is in flight.
6637            rt.prepare_overlap_slots(0, payload)?;
6638        }
6639        if pp_anatomy {
6640            // Drain the reverse-publication dependency before timing stage 0 itself. At c=1 this
6641            // prices any primary-stream rollback/refresh tail inherited from the prior round.
6642            for s in 0..n_st {
6643                let _st = rt.enter(s);
6644                rt.engine(s, e).stream().synchronize()?;
6645            }
6646            reverse_ms = reverse_started.elapsed().as_secs_f64() * 1e3;
6647        }
6648
6649        // Per-stage rope positions: in host mode the same [T] iota each stage uploads itself; in
6650        // stream mode each stage's own `pos_iota` over the shared read-only device counter.
6651        let stage_pos = |es: &Engine| -> Result<CudaSlice<i32>, Box<dyn std::error::Error>> {
6652            match stream {
6653                Some((_, ctr)) => {
6654                    let mut p = es.alloc_uninit::<i32>(t)?;
6655                    es.pos_iota(ctr, &mut p, t)?;
6656                    Ok(p)
6657                }
6658                None => {
6659                    let pos_vec: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
6660                    es.htod_i32(&pos_vec)
6661                }
6662            }
6663        };
6664
6665        // ---- STAGE 0: embed (the table lives with stage 0) + layers [0, fence[1]) + TX ----
6666        let slot = {
6667            let _st0 = rt.enter(0);
6668            let e0 = rt.engine(0, e);
6669            enqueue_spec_pipe_trace_marker(&e0.stream(), trace.as_ref(), "S0", "start", None)?;
6670            let stage0_started = std::time::Instant::now();
6671            let pos_d = stage_pos(e0)?;
6672            let x = match (stream, embd_dev) {
6673                (Some((vtok, _)), Some((g, qt, rb))) => {
6674                    e0.embed_gather_device_td(g, vtok, t, n_embd, qt, rb)?
6675                }
6676                (None, Some((g, qt, rb))) => e0.embed_gather_device_t(g, tokens, n_embd, qt, rb)?,
6677                _ => e0.htod(&self.embd.try_gather(n_embd, tokens)?)?,
6678            };
6679            let x = self.verify_layers(
6680                e0, x, fence[0], fence[1], &pos_d, pos0, t, cache, ckpt, stream, None,
6681            )?;
6682            if pp_anatomy {
6683                e0.stream().synchronize()?;
6684                stage0_ms = stage0_started.elapsed().as_secs_f64() * 1e3;
6685            }
6686            let tx_started = std::time::Instant::now();
6687            let slot = if pp_pipe.is_some() {
6688                rt.tx_pipelined(0, &x, payload)?
6689            } else {
6690                rt.tx(0, &x, payload)?
6691            };
6692            enqueue_spec_pipe_trace_marker(&e0.stream(), trace.as_ref(), "S0", "end", Some(slot))?;
6693            if pp_anatomy {
6694                e0.stream().synchronize()?;
6695                tx_ms = tx_started.elapsed().as_secs_f64() * 1e3;
6696            }
6697            slot
6698            // x + pos_d drop here: freed stream-ordered on stage-0's stream after use.
6699        };
6700
6701        Ok(VerifyBoundaryTicket {
6702            rt,
6703            caller_stream,
6704            slot,
6705            pos0,
6706            t,
6707            payload,
6708            n_st,
6709            pipelined: pp_pipe.is_some(),
6710            pp_anatomy,
6711            pp_started,
6712            reverse_ms,
6713            stage0_ms,
6714            tx_ms,
6715            trace,
6716            _walk_owner: walk_owner,
6717        })
6718    }
6719
6720    /// Consume a stage-0 boundary ticket and enqueue the remaining PP stages plus the head.
6721    /// On PP-2 this is exactly stage 1; PP-N keeps its pre-existing middle-stage walk here.
6722    #[allow(clippy::too_many_arguments)]
6723    fn verify_stage1_finish(
6724        &self,
6725        e: &Engine,
6726        ticket: VerifyBoundaryTicket,
6727        cache: &mut Cache,
6728        mut ckpt: Option<&mut VerifyCkpt>,
6729        stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
6730        fence: &[usize],
6731        publish_to_caller: bool,
6732    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
6733        let VerifyBoundaryTicket {
6734            rt,
6735            caller_stream,
6736            slot,
6737            pos0,
6738            t,
6739            payload,
6740            n_st,
6741            pipelined,
6742            pp_anatomy,
6743            pp_started,
6744            reverse_ms,
6745            stage0_ms,
6746            tx_ms,
6747            trace,
6748            _walk_owner,
6749        } = ticket;
6750        let n_embd = self.cfg.n_embd as usize;
6751        let eps = self.cfg.rms_eps;
6752        let mut slot = slot;
6753        let (mut rx_ms, mut stage1_ms) = (0.0f64, 0.0f64);
6754        let stage_pos = |es: &Engine| -> Result<CudaSlice<i32>, Box<dyn std::error::Error>> {
6755            match stream {
6756                Some((_, ctr)) => {
6757                    let mut p = es.alloc_uninit::<i32>(t)?;
6758                    es.pos_iota(ctr, &mut p, t)?;
6759                    Ok(p)
6760                }
6761                None => {
6762                    let pos_vec: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
6763                    es.htod_i32(&pos_vec)
6764                }
6765            }
6766        };
6767
6768        // ---- MIDDLE STAGES: RX boundary s-1 -> range -> TX boundary s ----
6769        for s in 1..n_st - 1 {
6770            let _st = rt.enter(s);
6771            let es = rt.engine(s, e);
6772            let pos_d = stage_pos(es)?;
6773            let x = rt.rx(s - 1, slot, payload)?;
6774            let x = self.verify_layers(
6775                es,
6776                x,
6777                fence[s],
6778                fence[s + 1],
6779                &pos_d,
6780                pos0,
6781                t,
6782                cache,
6783                ckpt.as_deref_mut(),
6784                stream,
6785                None,
6786            )?;
6787            slot = if pipelined {
6788                rt.tx_pipelined(s, &x, payload)?
6789            } else {
6790                rt.tx(s, &x, payload)?
6791            };
6792        }
6793
6794        // ---- LAST STAGE: RX + final range + output_norm + lm head ----
6795        let _stl = rt.enter(n_st - 1);
6796        let el = rt.engine(n_st - 1, e);
6797        let pos_d = stage_pos(el)?;
6798        let rx_started = std::time::Instant::now();
6799        let x = rt.rx(n_st - 2, slot, payload)?;
6800        if pp_anatomy {
6801            el.stream().synchronize()?;
6802            rx_ms = rx_started.elapsed().as_secs_f64() * 1e3;
6803        }
6804        enqueue_spec_pipe_trace_marker(&el.stream(), trace.as_ref(), "S1", "start", Some(slot))?;
6805        let stage1_started = std::time::Instant::now();
6806        let x = self.verify_layers(
6807            el,
6808            x,
6809            fence[n_st - 1],
6810            fence[n_st],
6811            &pos_d,
6812            pos0,
6813            t,
6814            cache,
6815            ckpt,
6816            stream,
6817            None,
6818        )?;
6819
6820        let mut hn = vbuf(el, payload)?;
6821        let logits = if self.sliding_gated_moe_batch_program() {
6822            // The PP Step35 serving path uses rms_norm + matmul for B=1 as well as B>1.
6823            // Verify must not switch numeric class merely because the same session speculates.
6824            el.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
6825            el.matmul(&self.output, &hn, t)?
6826        } else {
6827            el.rms_norm_decode(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
6828            el.matmul_decode_exact(&self.output, &hn, t)?
6829        };
6830        enqueue_spec_pipe_trace_marker(&el.stream(), trace.as_ref(), "S1", "end", Some(slot))?;
6831        if pp_anatomy {
6832            el.stream().synchronize()?;
6833            stage1_ms = stage1_started.elapsed().as_secs_f64() * 1e3;
6834        }
6835        // EXIT PUBLICATION: both returned buffers are still being produced on the last stage's
6836        // stream. Order the caller's stream behind that work before the buffers escape this
6837        // scope (the 2026-08-06 same-device ppspec find: without it the caller's primary-stream
6838        // consumer read unwritten logits — nondeterministic, one-device-only, and it poisoned
6839        // the following arm's KV in the same process).
6840        if publish_to_caller {
6841            rt.publish_to(n_st - 1, &caller_stream)?;
6842        }
6843        if pp_anatomy {
6844            if publish_to_caller {
6845                caller_stream.synchronize()?;
6846            }
6847            eprintln!(
6848                "[spec-pp-anatomy] t={t} reverse={reverse_ms:.3}ms stage0={stage0_ms:.3}ms \
6849                 tx={tx_ms:.3}ms rx={rx_ms:.3}ms stage1-head={stage1_ms:.3}ms total={:.3}ms",
6850                pp_started.elapsed().as_secs_f64() * 1e3,
6851            );
6852        }
6853        // stream: the device pos counter owns position; host mirror reconciles at drain.
6854        if stream.is_none() {
6855            cache.pos += t;
6856        }
6857        Ok((logits, if spec_hpost() { hn } else { x }))
6858    }
6859
6860    /// Step3.5/Step3.7 verify trunk in the serving batched numeric class.
6861    ///
6862    /// `step35_decode_batch_layers` is now authoritative at every live serving width, including
6863    /// B=1 (lane/cx-b1fix). The older verify walk deliberately mirrored the eager T=1 class:
6864    /// it replayed `step35_decode_attn` per row and used the eager/decode-exact FFN dispatch.
6865    /// Those classes are individually stable, but a near-tie prompt can choose different greedy
6866    /// bytes when a request moves from batched plain serving into speculative verify. Run the
6867    /// same authoritative B=1 stage subgraph for each verify row here. Rows still advance
6868    /// layer-by-layer, so every layer sees the preceding verify rows in its attention cache while
6869    /// every norm/projection/FFN uses exactly the live serving dispatch.
6870    #[allow(clippy::too_many_arguments)]
6871    /// PRIME-BY-T-ROWS (MEMRA_PRIME_TROWS=1): prefill the prompt through the same-session
6872    /// t-row walk in 32-row chunks — every row runs the t=1 decode program bit-for-bit
6873    /// (the TOKENWISE-prime ORACLE class), so this door is exact against the exactness
6874    /// reference while replacing the host-canonical per-token prime. Requires the walk
6875    /// doors (MEMRA_SPEC_VERIFY_EAGER/TCOL); returns the prime contract trio.
6876    #[allow(clippy::type_complexity)]
6877    pub(crate) fn step35_prime_trows(
6878        &self,
6879        e: &Engine,
6880        tokens: &[u32],
6881        cache: &mut Cache,
6882    ) -> Result<Option<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
6883    {
6884        let dbg = std::env::var("MEMRA_SPEC_FA2_DEBUG").as_deref() == Ok("1");
6885        if !prime_trows_on() {
6886            return Ok(None);
6887        }
6888        if !self.uses_sliding_gated_moe_program()
6889            || cache.pos != 0
6890            || cache.dflash_taps.is_some()
6891            || !spec_verify_eager_on()
6892            || !spec_verify_tcol_on()
6893        {
6894            if dbg {
6895                eprintln!(
6896                    "[prime-trows] refuse: program={} pos={} taps={} eager={:?} tcol={:?}",
6897                    self.uses_sliding_gated_moe_program(),
6898                    cache.pos,
6899                    cache.dflash_taps.is_some(),
6900                    std::env::var("MEMRA_SPEC_VERIFY_EAGER").ok(),
6901                    std::env::var("MEMRA_SPEC_VERIFY_TCOL").ok()
6902                );
6903            }
6904            return Ok(None);
6905        }
6906        let n_embd = self.cfg.n_embd as usize;
6907        let n_layers = self.layers.len();
6908        let t_total = tokens.len();
6909        let Some(embd_gpu) = self.embd_gpu_try(e) else {
6910            if dbg {
6911                eprintln!("[prime-trows] refuse: no device embed table");
6912            }
6913            return Ok(None);
6914        };
6915        let embd_qtype = match self.embd.ggml_type {
6916            memra_gguf::GgmlType::BF16 => crate::QT_BF16,
6917            memra_gguf::GgmlType::Q8_0 => crate::QT_Q8_0,
6918            other => {
6919                if dbg {
6920                    eprintln!("[prime-trows] refuse: embed dtype {other:?}");
6921                }
6922                return Ok(None);
6923            }
6924        };
6925        let embd_row_bytes = self.embd.raw.len() / self.cfg.n_vocab as usize;
6926        // Chunk plan: 32-row chunks; a 1-token tail folds into the previous chunk
6927        // (the walk floor is t >= 2).
6928        let mut bounds = Vec::new();
6929        let mut start = 0usize;
6930        while start < t_total {
6931            let mut end = (start + 32).min(t_total);
6932            if t_total - end == 1 {
6933                end -= 1;
6934            }
6935            bounds.push((start, end));
6936            start = end;
6937        }
6938        if bounds.iter().any(|(a, b)| b - a < 2) {
6939            return Ok(None); // degenerate short prompt keeps the ordinary prime
6940        }
6941        let mut hiddens = e.uninit(t_total * n_embd)?;
6942        let mut last: Option<CudaSlice<f32>> = None;
6943        for &(a, b) in &bounds {
6944            let tc = b - a;
6945            let tok_d = e.stream().clone_htod(&tokens[a..b])?;
6946            let x =
6947                e.embed_gather_device_td(embd_gpu, &tok_d, tc, n_embd, embd_qtype, embd_row_bytes)?;
6948            let out = self.step35_verify_batch_layers(e, x, 0, n_layers, a, tc, cache)?;
6949            e.copy_into(&mut hiddens, a * n_embd, &out, tc * n_embd)?;
6950            if b == t_total {
6951                let mut h = e.uninit(n_embd)?;
6952                e.dtod_copy_view(&out.slice((tc - 1) * n_embd..tc * n_embd), &mut h)?;
6953                last = Some(h);
6954            }
6955        }
6956        let h_seed = last.expect("last chunk produced the seed row");
6957        let mut hn = e.uninit(n_embd)?;
6958        e.rms_norm_decode(
6959            &h_seed,
6960            self.output_norm.float_data(),
6961            &mut hn,
6962            n_embd,
6963            1,
6964            self.cfg.rms_eps,
6965        )?;
6966        let logits_d = e.matmul_decode_exact(&self.output, &hn, 1)?;
6967        let logits = e.dtoh(&logits_d)?;
6968        cache.pos = t_total;
6969        Ok(Some((logits, h_seed, hiddens)))
6970    }
6971
6972    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
6973    fn step35_verify_batch_layers(
6974        &self,
6975        e: &Engine,
6976        mut x: CudaSlice<f32>,
6977        lo: usize,
6978        hi: usize,
6979        pos0: usize,
6980        t: usize,
6981        cache: &mut Cache,
6982    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6983        let n_embd = self.cfg.n_embd as usize;
6984        if !self.uses_sliding_gated_moe_program() {
6985            return Err(
6986                "serving-class verify requires sliding-gated-MoE canonical operations".into(),
6987            );
6988        }
6989        // SERVING-CLASS VERIFY (MEMRA_SPEC_VERIFY_EAGER=1, step37 MTP bring-up): each verify
6990        // column rides decode_layers_eager — the EXACT t=1 program live serving runs (all TP2
6991        // doors) — row-outer, so row r's appends land before row r+1 attends: bit-equal to
6992        // plain greedy by construction. Only the unsplit full-range walk qualifies; PP splits
6993        // and the tap path keep the batch-layer class.
6994        static VE: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
6995        let eager_verify =
6996            *VE.get_or_init(spec_verify_eager_on) && lo == 0 && hi == self.layers.len();
6997        if eager_verify {
6998            // T-COLUMN LAYER-OUTER WALK (MEMRA_SPEC_VERIFY_TCOL=1): per layer, one t-grid
6999            // attn norm + ONE weight-amortized QKV(+gate) over all T columns, then each
7000            // column runs the UNMODIFIED t=1 attention program via the col-select door and
7001            // the ordinary residual/FFN body. Values per column are bit-equal to the
7002            // row-outer walk: rms over the materialized residual == the fused add+norm
7003            // (kernel_check identity), the tcol kernel's per-column FP order == the t=1
7004            // kernel, and every downstream op IS the t=1 program.
7005            static TCOL: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
7006            let tcol = *TCOL.get_or_init(spec_verify_tcol_on);
7007            // T > 32 (prefill-class): run the SAME walk in 32-row chunks — each chunk's
7008            // rows are the t=1 program bit-for-bit and the rope pass advances the cache,
7009            // so a chunked call is value-identical to the row-outer loop it replaces.
7010            static TROWS_PREFILL: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
7011            // MEMRA_STEP_GEMM_PRIME outranks the walk: with the grouped GEMM prime armed, the
7012            // t-row walk defers so the batch path (GEMM trunk + grouped MoE) takes the prompt —
7013            // flag precedence between two existing doors, not a new flag. Without this, both
7014            // doors ON meant the walk still won and the GEMM prime needed PRIME_TROWS=0 by hand.
7015            let trows_prefill =
7016                *TROWS_PREFILL.get_or_init(|| prime_trows_on() && !crate::step_gemm_prime_on());
7017            // MEMRA_PRIME_TROWS_T=<w>: chunk width, default 8 = the REAL cap of this walk.
7018            // The workspace slabs go to 32 rows, but `matvec_bf16_qkvg_tcol_into` refuses
7019            // t > 8 (compile-time-T twins exist for 2/4/8 only; the runtime-t kernel spills
7020            // its accumulators to local memory), so a wider chunk fails the request with
7021            // "matvec_bf16_qkvg_tcol geometry" — which is exactly how the first server-path
7022            // TROWS arm died. Measured at 193 tokens: w=8 2.459 s, w=4 2.574 s.
7023            static TROWS_W: std::sync::OnceLock<Result<usize, String>> = std::sync::OnceLock::new();
7024            let trows_w = match TROWS_W.get_or_init(|| {
7025                let value = std::env::var("MEMRA_PRIME_TROWS_T").ok();
7026                parse_prime_trows_width(value.as_deref())
7027            }) {
7028                Ok(width) => *width,
7029                Err(err) => return Err(err.clone().into()),
7030            };
7031            if tcol && trows_prefill && t > trows_w {
7032                // One-time engagement receipt: without it a prefill gate cannot tell a
7033                // chunked walk from the row-outer fallback it is supposed to replace
7034                // (the first PRIME_TROWS gate passed vacuously on exactly that).
7035                static SEEN: std::sync::atomic::AtomicBool =
7036                    std::sync::atomic::AtomicBool::new(false);
7037                if !SEEN.swap(true, std::sync::atomic::Ordering::Relaxed) {
7038                    eprintln!(
7039                        "[prime-trows] ENGAGED t={t} width={trows_w} chunks={} layers={}..{}",
7040                        t.div_ceil(trows_w),
7041                        lo,
7042                        hi
7043                    );
7044                }
7045                let mut out = e.uninit(t * n_embd)?;
7046                let mut start = 0usize;
7047                while start < t {
7048                    let mut end = (start + trows_w).min(t);
7049                    if t - end == 1 {
7050                        end -= 1;
7051                    }
7052                    let tc = end - start;
7053                    let mut xc = e.uninit(tc * n_embd)?;
7054                    e.dtod_copy_view(&x.slice(start * n_embd..end * n_embd), &mut xc)?;
7055                    let oc =
7056                        self.step35_verify_batch_layers(e, xc, lo, hi, pos0 + start, tc, cache)?;
7057                    e.copy_into(&mut out, start * n_embd, &oc, tc * n_embd)?;
7058                    start = end;
7059                }
7060                return Ok(out);
7061            }
7062            if tcol && (2..=32).contains(&t) {
7063                // MEMRA_TCOL_PROF=1: synchronized per-segment wall profile of the walk
7064                // (norm+QKV precompute / per-col attention / per-col residual+FFN). The
7065                // syncs serialize the stream, so the split is for TARGETING amortization
7066                // work only — never a perf claim.
7067                static PROF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
7068                let prof =
7069                    *PROF.get_or_init(|| std::env::var("MEMRA_TCOL_PROF").as_deref() == Ok("1"));
7070                let mut prof_ms = [0f64; 3];
7071                let eps = self.cfg.rms_eps;
7072                let mut x_t = x;
7073                let mut h_t = e.uninit(t * n_embd)?;
7074                let mut h_row = e.uninit(n_embd)?; // real row: the non-dcw fallback reads it
7075                // Per-column pos buffers hoisted out of the layer loop (a per-col-per-layer
7076                // pageable htod was an in-stream engine turnaround x t x 45).
7077                let mut pos_rows = Vec::with_capacity(t);
7078                for r in 0..t {
7079                    pos_rows.push(e.htod_i32(&[(pos0 + r) as i32])?);
7080                }
7081                let mut ok = true;
7082                // MEMRA_TCOL_OPROJ=1: defer each column's o_proj — the finish seam
7083                // stashes `gated` instead of joining per column; one b4_tcol per rank +
7084                // one slab join produce every column's `mixed` after the attention pass.
7085                // Bit-exact per column (t=1 b4 program per column; elementwise join).
7086                // MEMRA_TCOL_FFN=1: today this only IMPLIES the o_proj defer above. Its
7087                // named feature, the two-column device-routed FFN sweep, rode the
7088                // slot-major v2 TP banks and was REMOVED with the MEMRA_NVFP4_BANK_V2 door
7089                // (2026-08-29, research/step37-bankv2-removal-20260829): the v2 layout
7090                // changed generated text in serving. The flag itself stays because it is
7091                // family-armed in the step37 serving defaults and killing it here would
7092                // silently drop the o_proj defer from the qualified serving shape.
7093                static FFN2: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
7094                let ffn_batch = *FFN2.get_or_init(tcol_ffn_on);
7095                let oproj_batch = crate::tp::tcol_oproj_on() || ffn_batch;
7096                // MEMRA_SPEC_FA2=1 (T=2 only): eligible layers defer BOTH columns' fa —
7097                // the per-column pass norms/ropes/appends and stashes q+gate, then one
7098                // shared-KV fa_decode_dcw2 per rank + the o_proj join produce the
7099                // [2, o_out] mixed slab. The precheck runs before arming (stashing is
7100                // unrecoverable); ineligible/boundary layers run the ordinary program.
7101                let fa2 = crate::tp::spec_fa2_on() && t <= 32;
7102                let mut mixed_row = e.uninit(n_embd)?;
7103                let mut pos_staged = false;
7104                for il in lo..hi {
7105                    let layer = &self.layers[il];
7106                    // BEFORE this layer touches its planes: is the history it is about to
7107                    // attend already poisoned? Global (non-ring) layers only, which are the
7108                    // ones the level-2 bitmap implicates.
7109                    if kv_plane_scan_on()
7110                        && self.step35_geom(il).window.is_none()
7111                        && let Some(distributed) = cache.tp_kv[il].as_ref()
7112                    {
7113                        scan_kv_plane(e, distributed, il, pos0)?;
7114                    }
7115                    let fa2_layer = fa2 && self.step35_fa_rows_precheck(cache, il, pos0, t)?;
7116                    let mut seg = std::time::Instant::now();
7117                    e.rms_norm(&x_t, layer.attn_norm.float_data(), &mut h_t, n_embd, t, eps)?;
7118                    if !self.step35_verify_qkv_precompute(e, il, &h_t, t)? {
7119                        ok = false;
7120                        break;
7121                    }
7122                    // FULL t-row attention pass (rope/append + fa + combine + o_proj in
7123                    // 3 launches/rank): same-session rows, slot = len-base+r, one len
7124                    // advance by t. Ring rebase happens during append BEFORE the fused
7125                    // kernel so device base_d and memory are already rebased for rows.
7126                    // Host cache bookkeeping mirrors the per-column tail.
7127                    let mut mixed_t_opt: Option<CudaSlice<f32>> = None;
7128                    if fa2_layer {
7129                        let crate::hybrid::Mixer::Full(fa) = &layer.mixer else {
7130                            return Err("verify rope pass expects full attention".into());
7131                        };
7132                        let tp = fa
7133                            .step_tp_qkv
7134                            .as_ref()
7135                            .ok_or("verify rope pass lost its TP state")?;
7136                        let empty: [CudaSlice<f32>; 0] = [];
7137                        let transaction = {
7138                            let tp_kv = cache.tp_kv[il]
7139                                .as_mut()
7140                                .expect("precheck verified the distributed cache");
7141                            let tx = tp_kv.begin_transaction()?;
7142                            tp.runtime.append_tp_kv_transaction_inner(
7143                                tp_kv, tx, &empty, &empty, t, true,
7144                            )?;
7145                            tx
7146                        };
7147                        match self.step35_verify_rope_fa_pass(e, il, cache, pos0, t, !pos_staged) {
7148                            Ok(Some(mixed_t)) => {
7149                                let tp_kv = cache.tp_kv[il]
7150                                    .as_mut()
7151                                    .expect("precheck verified the distributed cache");
7152                                tp.runtime.commit_tp_kv_transaction_external(
7153                                    tp_kv,
7154                                    transaction,
7155                                    t,
7156                                )?;
7157                                if let Some(local) = cache.kv[il].as_mut() {
7158                                    local.len = pos0 + t;
7159                                    if !crate::tp::len_mirror_lazy_on() {
7160                                        e.set_i32_one(&mut local.len_d, local.len as i32)?;
7161                                    }
7162                                    if let (Some(ring), Some(tp_base)) =
7163                                        (local.ring.as_mut(), tp_kv.ring_base())
7164                                        && ring.base() != tp_base
7165                                    {
7166                                        ring.apply_rebase(tp_base);
7167                                        if let Some(base_d) = local.base_d.as_mut() {
7168                                            e.set_i32_one(base_d, tp_base as i32)?;
7169                                        }
7170                                    }
7171                                }
7172                                pos_staged = true;
7173                                mixed_t_opt = Some(mixed_t);
7174                            }
7175                            Ok(None) => {
7176                                let tp_kv = cache.tp_kv[il]
7177                                    .as_mut()
7178                                    .expect("precheck verified the distributed cache");
7179                                tp.runtime.rollback_tp_kv_transaction(tp_kv, transaction)?;
7180                            }
7181                            Err(err) => {
7182                                if let Some(tp_kv) = cache.tp_kv[il].as_mut() {
7183                                    let _ =
7184                                        tp.runtime.rollback_tp_kv_transaction(tp_kv, transaction);
7185                                }
7186                                return Err(err);
7187                            }
7188                        }
7189                    }
7190                    if let Some(mixed_t) = mixed_t_opt {
7191                        if prof {
7192                            e.stream().synchronize()?;
7193                            prof_ms[1] += seg.elapsed().as_secs_f64() * 1e3;
7194                            seg = std::time::Instant::now();
7195                        }
7196                        let o_out = mixed_t.len() / t;
7197                        let mut next = e.uninit(t * n_embd)?;
7198                        {
7199                            for r in 0..t {
7200                                e.dtod_copy_view(
7201                                    &mixed_t.slice(r * o_out..(r + 1) * o_out),
7202                                    &mut mixed_row,
7203                                )?;
7204                                let mut x_row = e.uninit(n_embd)?;
7205                                e.dtod_copy_view(
7206                                    &x_t.slice(r * n_embd..(r + 1) * n_embd),
7207                                    &mut x_row,
7208                                )?;
7209                                let (x1, ffn_out) = self.residual_norm_ffn(
7210                                    e, layer, &x_row, &mixed_row, n_embd, il, eps,
7211                                )?;
7212                                let mut x2 = e.uninit(n_embd)?;
7213                                e.add(&x1, &ffn_out, &mut x2, n_embd)?;
7214                                e.dtod_copy_into(&x2, &mut next, r * n_embd)?;
7215                            }
7216                        }
7217                        if prof {
7218                            e.stream().synchronize()?;
7219                            prof_ms[2] += seg.elapsed().as_secs_f64() * 1e3;
7220                        }
7221                        x_t = next;
7222                        if spec_nan_scan() {
7223                            // The scan MUST sit on this arm too. It used to live only on
7224                            // the non-fused tail, so a fused layer's poison was first
7225                            // reported by the next non-fused layer.
7226                            verify_arm_receipt(
7227                                "fused",
7228                                il,
7229                                pos0,
7230                                t,
7231                                cache.tp_kv[il].as_ref().map(|d| d.staged_len()),
7232                            );
7233                            nan_scan_rows(
7234                                e,
7235                                &x_t,
7236                                t,
7237                                n_embd,
7238                                &format!("tcol layer {il} pos0={pos0} arm=fused"),
7239                            )?;
7240                        }
7241                        continue;
7242                    }
7243                    if prof {
7244                        e.stream().synchronize()?;
7245                        prof_ms[0] += seg.elapsed().as_secs_f64() * 1e3;
7246                        seg = std::time::Instant::now();
7247                    }
7248                    let mut next = e.uninit(t * n_embd)?;
7249                    // Columns whose o_proj was deferred (their FFN runs after the join).
7250                    // A NON-deferred column's FFN must run INSIDE the column loop: the
7251                    // oproj-tail handoff is a single cell that the same column's
7252                    // residual_norm_ffn consumes before the next column's finish.
7253                    let mut deferred: Vec<usize> = Vec::new();
7254                    let mut fa2_deferred: Vec<usize> = Vec::new();
7255                    let ffn_col = |r: usize,
7256                                   mixed: &CudaSlice<f32>,
7257                                   next: &mut CudaSlice<f32>|
7258                     -> Result<(), Box<dyn std::error::Error>> {
7259                        let mut x_row = e.uninit(n_embd)?;
7260                        e.dtod_copy_view(&x_t.slice(r * n_embd..(r + 1) * n_embd), &mut x_row)?;
7261                        let (x1, ffn_out) =
7262                            self.residual_norm_ffn(e, layer, &x_row, mixed, n_embd, il, eps)?;
7263                        if spec_nan_scan_level() >= 2 {
7264                            nan_scan_rows(
7265                                e,
7266                                &ffn_out,
7267                                1,
7268                                n_embd,
7269                                &format!("tcol layer {il} col {r} per-column FFN out"),
7270                            )?;
7271                        }
7272                        let mut x2 = e.uninit(n_embd)?;
7273                        e.add(&x1, &ffn_out, &mut x2, n_embd)?;
7274                        e.dtod_copy_into(&x2, next, r * n_embd)?;
7275                        Ok(())
7276                    };
7277                    #[allow(clippy::needless_range_loop)]
7278                    // allow: the explicit index loop keeps the offset arithmetic visible and aligned with the device-side indexing
7279                    for r in 0..t {
7280                        e.dtod_copy_view(&h_t.slice(r * n_embd..(r + 1) * n_embd), &mut h_row)?;
7281                        let row_pos = &pos_rows[r];
7282                        crate::tp::set_verify_tcol(Some(r));
7283                        if fa2_layer {
7284                            crate::tp::set_spec_fa2_defer(Some(r));
7285                        } else if oproj_batch {
7286                            crate::tp::set_tcol_oproj_defer(Some(r));
7287                        }
7288                        let mixed = match &layer.mixer {
7289                            crate::hybrid::Mixer::Full(fa) => {
7290                                self.full_attn_decode(e, fa, &h_row, row_pos, pos0 + r, cache, il)
7291                            }
7292                            _ => Err("step35 verify expects full attention".into()),
7293                        };
7294                        crate::tp::set_verify_tcol(None);
7295                        crate::tp::set_spec_fa2_defer(None);
7296                        crate::tp::set_tcol_oproj_defer(None);
7297                        let mixed = mixed?;
7298                        if fa2_layer && crate::tp::take_spec_fa2_stashed() {
7299                            fa2_deferred.push(r);
7300                        } else if oproj_batch && crate::tp::take_tcol_oproj_stashed() {
7301                            deferred.push(r);
7302                        } else {
7303                            if spec_nan_scan_level() >= 2 {
7304                                let cols = mixed.len();
7305                                nan_scan_rows(
7306                                    e,
7307                                    &mixed,
7308                                    1,
7309                                    cols,
7310                                    &format!("tcol layer {il} col {r} per-column ATTN out"),
7311                                )?;
7312                            }
7313                            ffn_col(r, &mixed, &mut next)?;
7314                        }
7315                    }
7316                    if !fa2_deferred.is_empty() && fa2_deferred.len() != t {
7317                        // The precheck guarantees both columns stash or neither; a strict
7318                        // subset means a column's output was never produced anywhere.
7319                        return Err("spec fa2 stash engaged for a subset of columns".into());
7320                    }
7321                    if prof {
7322                        e.stream().synchronize()?;
7323                        prof_ms[1] += seg.elapsed().as_secs_f64() * 1e3;
7324                        seg = std::time::Instant::now();
7325                    }
7326                    if !fa2_deferred.is_empty() {
7327                        deferred = fa2_deferred;
7328                    }
7329                    if !deferred.is_empty() {
7330                        let mixed_t = if fa2_layer {
7331                            self.step35_verify_fa_rows_join(e, il, cache, pos0, t)?
7332                        } else {
7333                            self.step35_verify_oproj_tcol(e, il, t)?
7334                        };
7335                        let o_out = mixed_t.len() / t;
7336                        if spec_nan_scan_level() >= 2 {
7337                            nan_scan_rows(
7338                                e,
7339                                &mixed_t,
7340                                t,
7341                                o_out,
7342                                &format!("tcol layer {il} JOINED attn over deferred cols"),
7343                            )?;
7344                        }
7345                        // Batched t=2 residual+MoE: one t-grid add_rms_norm (per-row
7346                        // program == t=1; bit-identical to the oproj-tail join per the
7347                        // M2 verbatim-program contract) feeding the two-column routed
7348                        // sweep. Ineligible layers (dense FFN, non-nvfp4) fall through
7349                        // to the per-column body.
7350                        {
7351                            for &r in &deferred {
7352                                e.dtod_copy_view(
7353                                    &mixed_t.slice(r * o_out..(r + 1) * o_out),
7354                                    &mut mixed_row,
7355                                )?;
7356                                ffn_col(r, &mixed_row, &mut next)?;
7357                            }
7358                        }
7359                    }
7360                    if prof {
7361                        e.stream().synchronize()?;
7362                        prof_ms[2] += seg.elapsed().as_secs_f64() * 1e3;
7363                    }
7364                    x_t = next;
7365                    if spec_nan_scan() {
7366                        verify_arm_receipt(
7367                            if fa2_layer { "join" } else { "percol" },
7368                            il,
7369                            pos0,
7370                            t,
7371                            cache.tp_kv[il].as_ref().map(|d| d.staged_len()),
7372                        );
7373                        nan_scan_rows(
7374                            e,
7375                            &x_t,
7376                            t,
7377                            n_embd,
7378                            &format!(
7379                                "tcol layer {il} pos0={pos0} arm={}",
7380                                if fa2_layer { "join" } else { "percol" }
7381                            ),
7382                        )?;
7383                    }
7384                }
7385                if prof {
7386                    eprintln!(
7387                        "[tcol-prof] t={t} norm+qkv={:.3}ms attn={:.3}ms ffn={:.3}ms",
7388                        prof_ms[0], prof_ms[1], prof_ms[2]
7389                    );
7390                }
7391                if ok {
7392                    return Ok(x_t);
7393                }
7394                // fall through to the row-outer walk on ineligible layers
7395                x = x_t;
7396            }
7397            let mut next = e.uninit(t * n_embd)?;
7398            let scan = spec_nan_scan();
7399            for r in 0..t {
7400                let mut row = e.uninit(n_embd)?;
7401                e.dtod_copy_view(&x.slice(r * n_embd..(r + 1) * n_embd), &mut row)?;
7402                let row_pos = e.htod_i32(&[(pos0 + r) as i32])?;
7403                let out = if scan {
7404                    // Diagnostic arm: the same range walked one layer at a time so the first
7405                    // poisoned layer names itself. `decode_layers_eager(lo, hi)` is range-scoped
7406                    // and executes its trailing residual add, so a per-layer chain is the same
7407                    // program with the cross-layer add+norm fusion unrolled.
7408                    nan_scan_rows(
7409                        e,
7410                        &row,
7411                        1,
7412                        n_embd,
7413                        &format!("embed row r={r} pos={}", pos0 + r),
7414                    )?;
7415                    let mut acc = row;
7416                    for il in lo..hi {
7417                        acc = self.decode_layers_eager(
7418                            e,
7419                            acc,
7420                            il,
7421                            il + 1,
7422                            &row_pos,
7423                            pos0 + r,
7424                            cache,
7425                        )?;
7426                        nan_scan_rows(
7427                            e,
7428                            &acc,
7429                            1,
7430                            n_embd,
7431                            &format!("row-outer layer {il} r={r} pos={}", pos0 + r),
7432                        )?;
7433                    }
7434                    acc
7435                } else {
7436                    self.decode_layers_eager(e, row, lo, hi, &row_pos, pos0 + r, cache)?
7437                };
7438                e.dtod_copy_into(&out, &mut next, r * n_embd)?;
7439            }
7440            // dflash taps are NOT produced on this arm (they need per-layer hiddens the
7441            // row-outer walk does not materialize); the door is a step37 MTP bring-up
7442            // surface where taps are unused.
7443            return Ok(next);
7444        }
7445        let mut ph_last = std::time::Instant::now();
7446        for il in lo..hi {
7447            let mut next = e.uninit(t * n_embd)?;
7448            for r in 0..t {
7449                let mut row = e.uninit(n_embd)?;
7450                e.dtod_copy_view(&x.slice(r * n_embd..(r + 1) * n_embd), &mut row)?;
7451                // The caller owns this verify's position. During controller overlap, cache.pos
7452                // still describes generation N while this stage-0 walk belongs to N+1.
7453                let row_pos = e.htod_i32(&[(pos0 + r) as i32])?;
7454                let mut one = [&mut *cache];
7455                let out = self.step35_decode_batch_layers(
7456                    e,
7457                    row,
7458                    &mut one,
7459                    &[(pos0 + r) as i32],
7460                    &row_pos,
7461                    il,
7462                    il + 1,
7463                    &mut ph_last,
7464                )?;
7465                e.dtod_copy_into(&out, &mut next, r * n_embd)?;
7466            }
7467            self.dflash_tap(e, cache, il, &next, t)?;
7468            x = next;
7469            if spec_nan_scan() {
7470                nan_scan_rows(e, &x, t, n_embd, &format!("batch-layer {il} pos0={pos0}"))?;
7471            }
7472        }
7473        Ok(x)
7474    }
7475
7476    /// DSpark drafter verify (lane/dspark-q38-recover): one t-row forward through the
7477    /// SERVING-CLASS verify funnel (`decode_step_t_core_stream` — the same numeric class
7478    /// MTP verify rides, GDN state advanced in place), returning per-row argmax tokens.
7479    /// Advances `cache.pos += t`; the caller owns snapshot/rollback (block acceptance is
7480    /// prefix-keep, not all-or-nothing).
7481    pub(crate) fn dspark_verify_t_am(
7482        &self,
7483        e: &Engine,
7484        tokens: &[u32],
7485        pos0: usize,
7486        cache: &mut Cache,
7487    ) -> Result<Vec<u32>, Box<dyn std::error::Error>> {
7488        let (logits, _hn) = self.decode_step_t_core_stream(
7489            e, tokens, pos0, cache, None, None, None, None, None, None,
7490        )?;
7491        let t = tokens.len();
7492        let v = self.output.out_features();
7493        let mut am_d = e.stream().alloc_zeros::<u32>(t)?;
7494        for r in 0..t {
7495            e.argmax_token_device_col(&logits, r, v, &mut am_d, r)?;
7496        }
7497        e.dtoh_u32(&am_d)
7498    }
7499
7500    /// DSpark verify returning the RAW verify logits [t, n_vocab] (device-resident) instead
7501    /// of per-row argmaxes — the sampled-admission arm's input (rejection-sampling accept
7502    /// gathers filtered p from these columns; lane/dspark-sampled-admission-20260820). Same
7503    /// forward as `dspark_verify_t_am`; the greedy arm keeps its argmax wrapper untouched.
7504    pub(crate) fn dspark_verify_t_logits(
7505        &self,
7506        e: &Engine,
7507        tokens: &[u32],
7508        pos0: usize,
7509        cache: &mut Cache,
7510    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7511        let (logits, _hn) = self.decode_step_t_core_stream(
7512            e, tokens, pos0, cache, None, None, None, None, None, None,
7513        )?;
7514        Ok(logits)
7515    }
7516
7517    /// DSpark verify with the MTP column-stash armed: identical forward to
7518    /// `dspark_verify_t_am`, but fills a `VerifyCkpt` so a partial accept can restore
7519    /// column state directly (`dspark_commit_prefix`) instead of snapshot-replay.
7520    /// The ckpt type is opaque outside spec.rs (newtype) — dflash.rs threads it through.
7521    pub(crate) fn dspark_verify_t_am_ckpt(
7522        &self,
7523        e: &Engine,
7524        tokens: &[u32],
7525        pos0: usize,
7526        cache: &mut Cache,
7527    ) -> Result<(Vec<u32>, DsparkVerifyCkpt), Box<dyn std::error::Error>> {
7528        let mut ck = VerifyCkpt::new(self.layers.len());
7529        let (logits, _hn) = self.decode_step_t_core_stream(
7530            e,
7531            tokens,
7532            pos0,
7533            cache,
7534            None,
7535            Some(&mut ck),
7536            None,
7537            None,
7538            None,
7539            None,
7540        )?;
7541        let t = tokens.len();
7542        let v = self.output.out_features();
7543        let mut am_d = e.stream().alloc_zeros::<u32>(t)?;
7544        for r in 0..t {
7545            e.argmax_token_device_col(&logits, r, v, &mut am_d, r)?;
7546        }
7547        Ok((e.dtoh_u32(&am_d)?, DsparkVerifyCkpt(ck)))
7548    }
7549
7550    /// Engine-bundle slice 2: `dspark_verify_t_am_ckpt` with DEVICE tokens and NO readback.
7551    /// The verify tokens are the round's `chain_d` (cand layout: [anchor, drafts...]); the
7552    /// embed gathers its first `t` entries on-device (`embed_gather_u32_t` — bit-identical
7553    /// rows to the host arm), so the host never blocks on the draft chain before dispatching
7554    /// verify. Returns the device per-row argmax buffer; the caller merges its readback with
7555    /// the chain's into ONE sync. Forward, ckpt fill and argmax walk are `_ckpt` verbatim.
7556    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
7557    pub(crate) fn dspark_verify_t_am_ckpt_dev(
7558        &self,
7559        e: &Engine,
7560        vtok: &CudaSlice<u32>,
7561        t: usize,
7562        pos0: usize,
7563        cache: &mut Cache,
7564        embd_dev: (&CudaSlice<u8>, i32, usize),
7565        graphs: Option<&mut DsparkVerifyGraphs>,
7566    ) -> Result<(CudaSlice<u32>, DsparkVerifyCkpt), Box<dyn std::error::Error>> {
7567        debug_assert!(
7568            vtok.len() >= t,
7569            "verify window exceeds the device token buffer"
7570        );
7571        // The slab flag is a per-round statement: clear it here so a verify that never
7572        // reaches the graphs door (rowwise env, a non-tparallel arm) cannot leave a
7573        // stale `true` steering the commit at slabs the round never wrote.
7574        let mut graphs = graphs;
7575        if let Some(g) = graphs.as_deref_mut() {
7576            g.round_slab = false;
7577        }
7578        let mut ck = VerifyCkpt::new(self.layers.len());
7579        // Dummy host tokens size the funnel; the embed reads `vtok` (the round-stream
7580        // arm's established pattern — spec.rs stream-mode verify does the same).
7581        let dummy = vec![0u32; t];
7582        let (logits, _hn) = self.decode_step_t_core_stream(
7583            e,
7584            &dummy,
7585            pos0,
7586            cache,
7587            Some(embd_dev),
7588            Some(&mut ck),
7589            None,
7590            None,
7591            Some(vtok),
7592            graphs,
7593        )?;
7594        let v = self.output.out_features();
7595        let mut am_d = e.stream().alloc_zeros::<u32>(t)?;
7596        for r in 0..t {
7597            e.argmax_token_device_col(&logits, r, v, &mut am_d, r)?;
7598        }
7599        Ok((am_d, DsparkVerifyCkpt(ck)))
7600    }
7601
7602    /// Ckpt-armed twin of [`Self::dspark_verify_t_logits`] (sampled-admission arm).
7603    pub(crate) fn dspark_verify_t_logits_ckpt(
7604        &self,
7605        e: &Engine,
7606        tokens: &[u32],
7607        pos0: usize,
7608        cache: &mut Cache,
7609    ) -> Result<(CudaSlice<f32>, DsparkVerifyCkpt), Box<dyn std::error::Error>> {
7610        let mut ck = VerifyCkpt::new(self.layers.len());
7611        let (logits, _hn) = self.decode_step_t_core_stream(
7612            e,
7613            tokens,
7614            pos0,
7615            cache,
7616            None,
7617            Some(&mut ck),
7618            None,
7619            None,
7620            None,
7621            None,
7622        )?;
7623        Ok((logits, DsparkVerifyCkpt(ck)))
7624    }
7625
7626    /// Restore the round to `keep` accepted columns from the verify stash: KV lens and
7627    /// pos from the pre-verify snapshot + keep, GDN conv/ssm from the stashed column
7628    /// state — no replay forward. The exact `commit_verified_prefix` the MTP path ships.
7629    pub(crate) fn dspark_commit_prefix(
7630        &self,
7631        e: &Engine,
7632        cache: &mut Cache,
7633        snap: &crate::cache::CacheSnapshot,
7634        ckpt: &DsparkVerifyCkpt,
7635        keep: usize,
7636    ) -> Result<(), Box<dyn std::error::Error>> {
7637        self.commit_verified_prefix(e, cache, snap, &ckpt.0, keep, false, None)
7638    }
7639
7640    /// Slice-3 commit twin: restore to `keep` accepted columns when the round's linear
7641    /// column stash lives in the graphs ctx's persistent slabs (`DsparkVerifyGraphs`) —
7642    /// the cols arm's exact semantics (KV lens + pos from the snapshot, GDN conv/ssm
7643    /// from the stash of column keep-1), slab-addressed and batched into two copy
7644    /// launches. `MEMRA_STATE_COPY_BATCH=0` falls back to per-layer view copies.
7645    pub(crate) fn dspark_commit_prefix_slab(
7646        &self,
7647        e: &Engine,
7648        cache: &mut Cache,
7649        snap: &crate::cache::CacheSnapshot,
7650        ctx: &DsparkVerifyGraphs,
7651        keep: usize,
7652    ) -> Result<(), Box<dyn std::error::Error>> {
7653        use cudarc::driver::DevicePtr;
7654        debug_assert!(keep >= 1, "keep==0 rounds take the legacy rollback");
7655        let mut conv_src: Vec<u64> = Vec::new();
7656        let mut ssm_src: Vec<u64> = Vec::new();
7657        let mut conv_dst: Vec<u64> = Vec::new();
7658        let mut ssm_dst: Vec<u64> = Vec::new();
7659        for il in 0..self.layers.len() {
7660            if let (Some(kvl), Some(saved)) = (cache.kv[il].as_mut(), snap.kv_len[il]) {
7661                kvl.len = saved + keep;
7662                e.set_i32_one(&mut kvl.len_d, kvl.len as i32)?;
7663            }
7664            if let Some(rl) = cache.recur[il].as_ref() {
7665                let (pc, ps, _cw, _sw) = ctx
7666                    .slab_row(e, il, keep - 1)
7667                    .ok_or("slab commit: linear layer missing from the graphs ctx")?;
7668                conv_src.push(pc);
7669                ssm_src.push(ps);
7670                let st = &e.gpu.stream();
7671                let (dc, _g0) = rl.conv_state.device_ptr(st);
7672                let (ds, _g1) = rl.ssm_state.device_ptr(st);
7673                conv_dst.push(dc);
7674                ssm_dst.push(ds);
7675            }
7676        }
7677        let n = conv_src.len();
7678        if n > 0 {
7679            if state_copy_batch_on() {
7680                let mut tt = vec![0u64; 2 * n];
7681                tt[..n].copy_from_slice(&conv_src);
7682                tt[n..].copy_from_slice(&conv_dst);
7683                let ct = e.htod_u64(&tt)?;
7684                tt[..n].copy_from_slice(&ssm_src);
7685                tt[n..].copy_from_slice(&ssm_dst);
7686                let st = e.htod_u64(&tt)?;
7687                e.copy_batch_uniform_f32(&ct, n, ctx.conv_words)?;
7688                e.copy_batch_uniform_f32(&st, n, ctx.ssm_words)?;
7689            } else {
7690                let (cw, sw) = (ctx.conv_words, ctx.ssm_words);
7691                let row = keep - 1;
7692                for il in 0..self.layers.len() {
7693                    let Some(rl) = cache.recur[il].as_mut() else {
7694                        continue;
7695                    };
7696                    let k = ctx.lin_pos[&il];
7697                    {
7698                        let sv = e.view(&ctx.stash_conv[k], (row + 1) * cw);
7699                        let win = sv.slice(row * cw..(row + 1) * cw);
7700                        e.copy_view_into(&mut rl.conv_state, 0, &win, cw)?;
7701                    }
7702                    {
7703                        let sv = e.view(&ctx.stash_ssm[k], (row + 1) * sw);
7704                        let win = sv.slice(row * sw..(row + 1) * sw);
7705                        e.copy_view_into(&mut rl.ssm_state, 0, &win, sw)?;
7706                    }
7707                }
7708            }
7709        }
7710        cache.pos = snap.pos + keep;
7711        Ok(())
7712    }
7713
7714    /// Qwen35-family verify trunk in the live serving numeric class.
7715    ///
7716    /// Serving intentionally keeps this architecture in the generic batched program even at
7717    /// B=1. The older verify walk used its own mirrored dispatch and can flip near-tie argmaxes.
7718    ///
7719    /// Two arms, one numeric class:
7720    /// - DENSE GDN (`DenseMlp`, t<=16): `qwen35_verify_tparallel` — the weight ops (norms,
7721    ///   projections, FFN) hoist to m=T through the exact-tier batched kernels whose per-row
7722    ///   program IS the m=1 program (`matmul_pre == fused2 per (tensor,row); _bN mmvq per-row
7723    ///   == m=1` — decode_batch.rs v2 note), while the state ops (conv ring, gdn scan, KV
7724    ///   append, fa decode) stay a per-row loop running the b_n=1 serving kernels with each
7725    ///   row's own t_kv-driven arm pick (the straddle law: every row executes the exact
7726    ///   program its isolated serving step would). One weight read per layer per round
7727    ///   instead of T — this is what makes MTP profitable in the exact class (the per-row
7728    ///   walk measured verify(K+1) ~= (K+1) plain steps: 69 -> 44 tok/s served, 2026-08-15).
7729    /// - MoE / t>16 / `MEMRA_SPEC_VERIFY_ROWWISE=1`: the per-row replay of the authoritative
7730    ///   serving layer body, preserving single-session autoregressive cache order (the
7731    ///   correctness reference; also the rollback seam for the t-parallel arm).
7732    ///
7733    /// Bit-identity of the t-parallel arm vs the rowwise arm is gated by spec-serve-gate
7734    /// (zero differing logits at T=1..4, K arms) + the 8-prompt ON/OFF canary before ship.
7735    #[allow(clippy::too_many_arguments)]
7736    fn qwen35_verify_batch_layers(
7737        &self,
7738        e: &Engine,
7739        x: CudaSlice<f32>,
7740        lo: usize,
7741        hi: usize,
7742        pos0: usize,
7743        t: usize,
7744        cache: &mut Cache,
7745        ckpt: Option<&mut VerifyCkpt>,
7746        stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
7747        graphs: Option<&mut DsparkVerifyGraphs>,
7748    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7749        // Qwen35Moe admitted 2026-08-20 (lane/draftcost-moe): the t-parallel arm already
7750        // carries the MoE FFN (`moe_ffn_il_zq8` at m=T) and the GDN per-row state loop; the
7751        // arch fence was a qualification gate, not a mechanism gap. Measured disease on the
7752        // 35B-A3B class: rowwise verify ~= 5.6 ms per drafted token (one full trunk step
7753        // each) — the same (K+1)-plain-steps wall the dense admission fixed on 2026-08-15.
7754        // Rollback seam unchanged: MEMRA_SPEC_VERIFY_ROWWISE=1.
7755        let rowwise = std::env::var("MEMRA_SPEC_VERIFY_ROWWISE").as_deref() == Ok("1")
7756            || !self.batched_serving_numeric_class()
7757            || t > 16;
7758        if rowwise {
7759            if stream.is_some() {
7760                // rowwise replays per row with host cache.pos — irreconcilable with a
7761                // device position counter. Burst callers must keep t <= 16 and the
7762                // ROWWISE env unset; refusing beats silently mispositioned rows.
7763                return Err("qwen35 rowwise verify has no ROUND-STREAM arm \
7764                            (t > 16 or MEMRA_SPEC_VERIFY_ROWWISE=1)"
7765                    .into());
7766            }
7767            self.qwen35_verify_rowwise(e, x, lo, hi, pos0, t, cache, ckpt)
7768        } else {
7769            self.qwen35_verify_tparallel(e, x, lo, hi, pos0, t, cache, ckpt, stream, graphs)
7770        }
7771    }
7772
7773    /// The per-row correctness reference: replay each verify row through the authoritative
7774    /// serving layer body (`decode_batch_layers` at b_n=1). T full weight reads per layer.
7775    #[allow(clippy::too_many_arguments)]
7776    fn qwen35_verify_rowwise(
7777        &self,
7778        e: &Engine,
7779        mut x: CudaSlice<f32>,
7780        lo: usize,
7781        hi: usize,
7782        pos0: usize,
7783        t: usize,
7784        cache: &mut Cache,
7785        mut ckpt: Option<&mut VerifyCkpt>,
7786    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7787        let n_embd = self.cfg.n_embd as usize;
7788        let saved_pos = cache.pos;
7789        let mut ph_last = std::time::Instant::now();
7790        for il in lo..hi {
7791            let mut next = e.uninit(t * n_embd)?;
7792            let mut col_states: Option<Vec<(CudaSlice<f32>, CudaSlice<f32>)>> =
7793                if ckpt.is_some() && t >= 2 && matches!(self.layers[il].mixer, Mixer::Linear(_)) {
7794                    Some(Vec::with_capacity(t - 1))
7795                } else {
7796                    None
7797                };
7798            for r in 0..t {
7799                cache.pos = pos0 + r;
7800                let mut row = e.uninit(n_embd)?;
7801                e.dtod_copy_view(&x.slice(r * n_embd..(r + 1) * n_embd), &mut row)?;
7802                let row_pos = e.htod_i32(&[(pos0 + r) as i32])?;
7803                let mut one = [&mut *cache];
7804                let ctx = self.batch_layer_ctx(e, &one, il, il + 1)?;
7805                let out = match self.decode_batch_layers(
7806                    e,
7807                    row,
7808                    &mut one,
7809                    &ctx,
7810                    &row_pos,
7811                    &mut ph_last,
7812                ) {
7813                    Ok(out) => out,
7814                    Err(error) => {
7815                        cache.pos = saved_pos;
7816                        return Err(error);
7817                    }
7818                };
7819                e.dtod_copy_into(&out, &mut next, r * n_embd)?;
7820                if r + 1 < t
7821                    && let Some(states) = col_states.as_mut()
7822                {
7823                    let recur = cache.recur[il]
7824                        .as_ref()
7825                        .ok_or("Qwen35-MoE linear verify layer has no recurrent state")?;
7826                    states.push((
7827                        e.clone_dtod(&recur.conv_state)?,
7828                        e.clone_dtod(&recur.ssm_state)?,
7829                    ));
7830                }
7831            }
7832            if let (Some(checkpoint), Some(states)) = (ckpt.as_deref_mut(), col_states) {
7833                checkpoint.cols[il] = Some(states);
7834            }
7835            x = next;
7836        }
7837        cache.pos = saved_pos;
7838        Ok(x)
7839    }
7840
7841    /// T-PARALLEL VERIFY IN THE SERVING NUMERIC CLASS (lane/tparallel-verify, 2026-08-15).
7842    ///
7843    /// The weight ops run ONCE per layer at m=T; the state ops run per row through the same
7844    /// b_n=1 serving kernels the rowwise replay uses. Per-row bit-identity rests on the two
7845    /// pins the serving batch tier already carries:
7846    ///   * `matmul_pre` / `_bN` mmvq: per-row program == m=1 program (decode_batch.rs v2 note,
7847    ///     kernel-check pinned) — so a [T, n_embd] projection row equals the row projected
7848    ///     alone;
7849    ///   * row-indexed norms/elementwise (`rms_norm`, `quantize_q8_1`, `add_rms_norm`,
7850    ///     `gated_rmsnorm[_q8_1]`, `silu_mul`, `rope_neox` with per-row positions): the T-row
7851    ///     launch is the per-row program (same pin the generic verify's fused norms rely on).
7852    ///     The sequential dependencies keep their exact serving order: the conv ring / gdn scan
7853    ///     chain state row -> row through the `_b` kernels at b_n=1 (ping-pong via a 6-entry
7854    ///     alternating pointer table, host handles swapped per row so VerifyCkpt clones the
7855    ///     canonical state exactly as the rowwise arm does), and each row's KV append + fa decode
7856    ///     picks its arm from ITS OWN t_kv (append: format-only; fa: `fa_seqs_eligible` + its own
7857    ///     `fa_split_keys` rung at b_n=1) — the straddle law per row, so every row executes the
7858    ///     program its isolated B=1 serving step would.
7859    ///
7860    /// Cost: 1 weight read per layer per round + T state micro-launches, vs the rowwise arm's
7861    /// T weight reads. Gated bit-identical vs the rowwise arm by spec-serve-gate + canary.
7862    #[allow(clippy::too_many_arguments)]
7863    fn qwen35_verify_tparallel(
7864        &self,
7865        e: &Engine,
7866        mut x: CudaSlice<f32>,
7867        lo: usize,
7868        hi: usize,
7869        pos0: usize,
7870        t: usize,
7871        cache: &mut Cache,
7872        mut ckpt: Option<&mut VerifyCkpt>,
7873        stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
7874        mut graphs: Option<&mut DsparkVerifyGraphs>,
7875    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7876        let seqs_append =
7877            std::env::var("MEMRA_BATCH_APPEND").as_deref() != Ok("0") && !Engine::kv_fp8_on();
7878        let batch_fa_on = std::env::var("MEMRA_BATCH_FA").as_deref() != Ok("0");
7879
7880        // Merge guard (v0.98 train, re-affirmed on the v0.100 train over slice 4c): the
7881        // ROUND-STREAM arm (lane/draftcost-moe, device position counter) and the dspark
7882        // verify graphs (engine-bundle slice 3 / trunk slice 4c) have no common caller —
7883        // stream rides the qwen35moe burst, graphs ride the dspark route. If a future
7884        // caller arms both, refuse loudly instead of silently dropping the graphs ctx
7885        // (the stream linear arm takes linear_attn_verify_t, not the graphed segment or
7886        // full-verify bodies).
7887        if stream.is_some() && graphs.is_some() {
7888            return Err(
7889                "qwen35 tparallel verify: ROUND-STREAM and dspark verify graphs \
7890                        cannot arm together"
7891                    .into(),
7892            );
7893        }
7894        // Engine-bundle slice 3 + slice 4c: with a graphs ctx armed, pointer tables are
7895        // refreshed once per verify (the gdn ping-pong moves handles; a fresh generation
7896        // moves the kv caches). Then:
7897        //  - slice 4c: when the WHOLE round rides one seqs rung (every row batchable, one
7898        //    split-ladder step, rung covers the round), the ENTIRE walk replays as ONE
7899        //    full-verify graph per (vt, rung) — linear layers through the shared
7900        //    `qwen35_tparallel_linear_layer` body, full-attention layers through the
7901        //    shared `qwen35_tparallel_fa_layer` body in graph mode.
7902        //  - fallback (straddle rounds, below the vec floor, partial walks): runs of
7903        //    consecutive LINEAR layers replay the slice-3 per-(segment, vt) graphs and
7904        //    the full-attention layers run eager (batched rows when eligible).
7905        //
7906        // GRAPH-LAUNCH HEADROOM GUARD (see GRAPH_LAUNCH_MIN_FREE): the dspark verify
7907        // graphs replay through this walk from THREE callers — the MTP spec round's vg
7908        // door (already dropped per round by `graph_round_ok` before it gets here), the
7909        // dspark one-shot, and the dspark SERVE round (default ON since v0.108). Below
7910        // the driver-free floor the WHOLE round takes the byte-identical eager
7911        // cols-ckpt walk — the same drop-the-ctx fallback the pool ceiling already
7912        // takes — instead of feeding cuGraphLaunch a card it segfaults on.
7913        if let Some(g) = graphs.as_deref_mut()
7914            && !graph_launch_headroom_ok(e)
7915        {
7916            g.round_slab = false;
7917            graphs = None;
7918            static NOTED: std::sync::Once = std::sync::Once::new();
7919            NOTED.call_once(|| graph_replay_suspended_note("dspark-vg"));
7920        }
7921        if let Some(g) = graphs.as_deref_mut() {
7922            g.refresh_tables(e, cache)?;
7923            g.round_slab = false;
7924            if let Some(rung) = g.full_rung(self, cache, lo, hi, t, seqs_append && batch_fa_on) {
7925                // Pool ceiling (dspark_vg_cap): an existing key always replays; a NEW
7926                // full capture past the ceiling falls through to the segment/eager arms.
7927                if g.full.contains_key(&(t, rung, hi)) || g.can_capture() {
7928                    let out = g.run_full(self, e, lo, hi, &x, t, pos0, rung, cache)?;
7929                    g.round_slab = true;
7930                    return Ok(out);
7931                }
7932            }
7933            // Round-atomic ceiling check for the segment door: if any linear run in this
7934            // walk would need a NEW capture past the ceiling, the whole round runs the
7935            // eager cols-ckpt walk (mixing slab- and cols-stashed layers in one round
7936            // would corrupt the commit).
7937            if !g.segments_ready(self, lo, hi, t) {
7938                graphs = None;
7939            }
7940        }
7941        // STREAM (2b, lane/draftcost-moe): positions come from the device round counter
7942        // (pos_iota / i32_copy_add) so a burst round needs no host position knowledge.
7943        let pos_d = match stream {
7944            Some((_, ctr)) => {
7945                let mut p = e.alloc_uninit::<i32>(t)?;
7946                e.pos_iota(ctr, &mut p, t)?;
7947                p
7948            }
7949            None => {
7950                let pos_host: Vec<i32> = (0..t).map(|r| (pos0 + r) as i32).collect();
7951                e.htod_i32(&pos_host)?
7952            }
7953        };
7954        // Per-row 1-element position buffers, built ONCE per verify (the append/fa wrappers
7955        // take owned pos slices; building these inside the layer x row loops cost 16xT H2Ds).
7956        // LAZY since slice 4: the batched fa/append arm never touches them — they are built
7957        // on the first per-row fallback layer only (stream-aware there; the stream FA arm
7958        // rides the dc rows kernels and never reaches the fallback).
7959        let mut pos_rows: Option<Vec<CudaSlice<i32>>> = None;
7960        let mut il = lo;
7961        while il < hi {
7962            if graphs.is_some() && matches!(self.layers[il].mixer, Mixer::Linear(_)) {
7963                let mut end = il;
7964                while end < hi && matches!(self.layers[end].mixer, Mixer::Linear(_)) {
7965                    end += 1;
7966                }
7967                let g = graphs.as_deref_mut().expect("checked above");
7968                x = g.run_segment(self, e, il, end, &x, t, cache)?;
7969                g.round_slab = true;
7970                il = end;
7971                continue;
7972            }
7973            let layer = &self.layers[il];
7974            if stream.is_none() && matches!(layer.mixer, Mixer::Linear(_)) {
7975                // Eager linear layer (no graphs ctx): the shared body, legacy cols-ckpt arm.
7976                // Under ROUND-STREAM the linear layers ride the fa-body match's stream arm
7977                // below (linear_attn_verify_t — the stream COMMIT needs its GdnStash).
7978                x = self.qwen35_tparallel_linear_layer(
7979                    e,
7980                    il,
7981                    &x,
7982                    t,
7983                    cache,
7984                    ckpt.as_deref_mut(),
7985                    None,
7986                    None,
7987                )?;
7988                il += 1;
7989                continue;
7990            }
7991            // Full-attention (or stream-Linear, or MLA-refusing) layer: the extracted
7992            // shared body — eager arm (fresh per-verify pos/table, exact t_kv sizing,
7993            // in-body len bump). The slice-4c captured full-verify graphs run the SAME
7994            // body in graph mode; under ROUND-STREAM the body's dc-rows / GDN stream arms
7995            // run (lane/draftcost-moe).
7996            x = self.qwen35_tparallel_fa_layer(
7997                e,
7998                il,
7999                &x,
8000                t,
8001                cache,
8002                FaLayerArgs {
8003                    pos_d: &pos_d,
8004                    pos_rows: &mut pos_rows,
8005                    pos0,
8006                    seqs_append,
8007                    batch_fa_on,
8008                    graph_cap: None,
8009                    stream,
8010                    ckpt: ckpt.as_deref_mut(),
8011                },
8012            )?;
8013            il += 1;
8014        }
8015        Ok(x)
8016    }
8017
8018    /// SHARED dense-FFN body for the qwen35 t-parallel layers (trunk-kernels slice B) —
8019    /// ONE copy for the fa and linear layer bodies (the verify_layers extraction lesson).
8020    /// Dual arm (MEMRA_TK_FFN_DUAL, default on): gate+up in ONE dual launch from the
8021    /// pre-quantized activation with macro-scales DEFERRED into the fused SwiGLU+q8_1
8022    /// epilogue, then ffn_down from the fused (aq, ad) — the q27 verify chain verbatim.
8023    /// Every door is the bit-identical proven one: `matmul_decode_exact_dual_pre` (per
8024    /// (tensor,token,row) == the two singles), `silu_mul_scaled_q8_1` (y*s inline == the
8025    /// scale_inplace store, value-exact; fused quantize == quantize_q8_1 bytes),
8026    /// `matmul_decode_exact_pre` (dispatch mirror of the singles' q8_1-fast tail).
8027    /// Dual-refused (t outside 2..=7, non-NVFP4, layout mismatch) or seam off -> the
8028    /// original singles chain, byte-for-byte.
8029    #[allow(clippy::too_many_arguments)]
8030    fn qwen35_tparallel_dense_ffn(
8031        &self,
8032        e: &Engine,
8033        ffn_gate: &crate::model::GpuTensor,
8034        ffn_up: &crate::model::GpuTensor,
8035        ffn_down: &crate::model::GpuTensor,
8036        zn: &CudaSlice<f32>,
8037        t: usize,
8038        n_embd: usize,
8039    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8040        let n_ff = ffn_gate.out_features();
8041        let (zq, zd) = e.quantize_q8_1(zn, t, n_embd)?;
8042        if Engine::tk_ffn_dual_on()
8043            && let Some(((g, gs), (u, us))) =
8044                e.matmul_decode_exact_dual_pre(ffn_gate, ffn_up, &zq, &zd, t)?
8045        {
8046            if e.uses_q8_1_fast(ffn_down) {
8047                let (aq, ad) = e.silu_mul_scaled_q8_1(&g, &u, gs, us, t * n_ff)?;
8048                return e.matmul_decode_exact_pre(ffn_down, &aq, &ad, t);
8049            }
8050            let mut act = e.uninit(t * n_ff)?;
8051            e.silu_mul_scaled(&g, &u, gs, us, &mut act, t * n_ff)?;
8052            let (aq, ad) = e.quantize_q8_1(&act, t, n_ff)?;
8053            return e.matmul_pre(ffn_down, &aq, &ad, &act, t);
8054        }
8055        // v1 singles chain (seam off or dual-refused) — the pre-slice-B body verbatim.
8056        let g = e.matmul_pre(ffn_gate, &zq, &zd, zn, t)?;
8057        let u = e.matmul_pre(ffn_up, &zq, &zd, zn, t)?;
8058        let mut act = e.uninit(t * n_ff)?;
8059        e.silu_mul(&g, &u, &mut act, t * n_ff)?;
8060        let (aq, ad) = e.quantize_q8_1(&act, t, n_ff)?;
8061        e.matmul_pre(ffn_down, &aq, &ad, &act, t)
8062    }
8063
8064    /// ONE t-parallel FULL-ATTENTION layer (attn_norm + fa mixer + post_attn_norm + FFN +
8065    /// tap) — extracted from the walk exactly like `qwen35_tparallel_linear_layer` so the
8066    /// eager walk and the slice-4c captured full-verify graphs execute the SAME body (a
8067    /// second copy is how dispatch mirrors drift — the verify_layers extraction lesson).
8068    ///
8069    /// `args.graph_cap = Some((table, off, rung_end))` is the captured-graph mode:
8070    /// - kv base-pointer pairs come from the ctx-owned persistent table at `off` (a fresh
8071    ///   generation's cache lands at new addresses that only the per-verify table refresh
8072    ///   knows — the slice-3 baked-address lesson);
8073    /// - the seqs twins size partials/grid at `rung_end` and pin `split_keys` to the
8074    ///   rung's ladder value: `n_splits_max` is pure stride, splits >= ns_eff write the
8075    ///   EMPTY partial the combine never reads, and every per-row T_kv derives in-kernel
8076    ///   from `pos_seq[z]` — so one captured launch replays bit-identically for every
8077    ///   round whose rows all sit inside the rung;
8078    /// - the host len bump moves to the replay caller (captured host code does not
8079    ///   re-run at replay).
8080    ///   Graph mode REFUSES any round the batched arm cannot take: the per-row fallback
8081    ///   host-branches on t_kv and must never be captured.
8082    #[allow(clippy::too_many_arguments)]
8083    fn qwen35_tparallel_fa_layer(
8084        &self,
8085        e: &Engine,
8086        il: usize,
8087        x: &CudaSlice<f32>,
8088        t: usize,
8089        cache: &mut Cache,
8090        args: FaLayerArgs<'_>,
8091    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8092        use cudarc::driver::DevicePtr;
8093        let cfg = &self.cfg;
8094        let n_embd = cfg.n_embd as usize;
8095        let eps = cfg.rms_eps;
8096        let head_dim_global = cfg.head_dim_k as usize;
8097        let layer = &self.layers[il];
8098        let FaLayerArgs {
8099            pos_d,
8100            pos_rows,
8101            pos0,
8102            seqs_append,
8103            batch_fa_on,
8104            graph_cap,
8105            stream,
8106            ckpt,
8107        } = args;
8108
8109        // ---- attn_norm + q8_1 quantize at m=T (row-indexed == per-row) ----
8110        let anorm = layer.attn_norm.float_data();
8111        let mut xn = e.uninit(t * n_embd)?;
8112        e.rms_norm(x, anorm, &mut xn, n_embd, t, eps)?;
8113        let (hq, hd) = e.quantize_q8_1(&xn, t, n_embd)?;
8114
8115        let mixed: CudaSlice<f32> = match &layer.mixer {
8116            Mixer::Mla(_) => crate::hybrid::mla_path_unimplemented("tensor-parallel attention"),
8117            Mixer::Kda(_) => crate::hybrid::kda_path_unimplemented("T-parallel attention"),
8118            // STREAM ARM (2b, lane/draftcost-moe): under a device position counter the
8119            // per-row serving-kernel chain cannot run (host state swaps keyed on host
8120            // row index are fine, but the stream COMMIT needs the GdnStash for its _dc
8121            // rebuild — the per-row chain only produces per-column clones). GDN rides
8122            // `linear_attn_verify_t`: batched q8_1-class projections, stash-producing,
8123            // and its one-scan recurrence is pinned bit-identical to T chained T=1
8124            // steps (its header + kernel-check). Position-independent, so no counter
8125            // plumbing is needed. Guards mirror the generic call site exactly.
8126            Mixer::Linear(la) if stream.is_some() => {
8127                if !(t >= 3 || (t == 2 && spec_m2()))
8128                    || !self.mixer_in_q8_1_fast(e, &layer.mixer)
8129                    || !e.uses_q8_1_fast(&la.ssm_out)
8130                {
8131                    return Err("qwen35 stream verify: GDN batched arm requires t>=3 \
8132                                (or MEMRA_SPEC_M2 at t=2) and q8_1-fast projections"
8133                        .into());
8134                }
8135                let want = ckpt.is_some();
8136                let (out, stash) =
8137                    self.linear_attn_verify_t(e, la, &xn, Some((&hq, &hd)), t, cache, il, want)?;
8138                if let (Some(ck), Some(st)) = (ckpt, stash) {
8139                    ck.gdn[il] = Some(st);
8140                }
8141                out
8142            }
8143            Mixer::Linear(_) => {
8144                unreachable!("linear layers ride qwen35_tparallel_linear_layer")
8145            }
8146            Mixer::Full(fa) => {
8147                let geometry = cfg.full_attention_geometry_at(il as u32);
8148                let n_head = geometry.n_head as usize;
8149                let n_head_kv = geometry.n_head_kv as usize;
8150                let head_dim = geometry.head_dim_k as usize;
8151                let rope_dims = geometry.n_rot as usize;
8152                let rope_base = geometry.rope_base;
8153                let scale = geometry.attention_scale();
8154                // Batched projections: one weight read serves all T rows.
8155                // GROUP-3 twin (trunk-kernels slice D): q/k/v in ONE launch — the group4
8156                // kernel with n3=0, bit-identical per (tensor, token, row) to the three
8157                // singles; refused or MEMRA_TK_FA_GROUP=0 -> singles byte-for-byte.
8158                let (qf, mut k, v) = match e.matmul_decode_exact_group3_pre(
8159                    [&fa.wq, &fa.wk, &fa.wv],
8160                    &hq,
8161                    &hd,
8162                    t,
8163                )? {
8164                    Some(mut g3) => {
8165                        let v = g3.pop().unwrap();
8166                        let k = g3.pop().unwrap();
8167                        let qf = g3.pop().unwrap();
8168                        (qf, k, v)
8169                    }
8170                    None => (
8171                        e.matmul_pre(&fa.wq, &hq, &hd, &xn, t)?,
8172                        e.matmul_pre(&fa.wk, &hq, &hd, &xn, t)?,
8173                        e.matmul_pre(&fa.wv, &hq, &hd, &xn, t)?,
8174                    ),
8175                };
8176                let gated =
8177                    geometry.attention_gate == memra_gguf::config::AttentionGateKind::FusedQ;
8178                let (mut q, gate) = if gated {
8179                    let mut qs = e.uninit(t * n_head * head_dim)?;
8180                    let mut gs = e.uninit(t * n_head * head_dim)?;
8181                    e.q_gate_split(&qf, &mut qs, &mut gs, head_dim, n_head, t)?;
8182                    (qs, Some(gs))
8183                } else {
8184                    (qf, None)
8185                };
8186                let mut qn = e.uninit(t * n_head * head_dim)?;
8187                e.rms_norm(
8188                    &q,
8189                    fa.q_norm.float_data(),
8190                    &mut qn,
8191                    head_dim,
8192                    t * n_head,
8193                    eps,
8194                )?;
8195                q = qn;
8196                let mut kn = e.uninit(t * n_head_kv * head_dim)?;
8197                e.rms_norm(
8198                    &k,
8199                    fa.k_norm.float_data(),
8200                    &mut kn,
8201                    head_dim,
8202                    t * n_head_kv,
8203                    eps,
8204                )?;
8205                k = kn;
8206                e.rope_neox(
8207                    &mut q, pos_d, head_dim, rope_dims, n_head, t, rope_base, 1.0,
8208                )?;
8209                e.rope_neox(
8210                    &mut k, pos_d, head_dim, rope_dims, n_head_kv, t, rope_base, 1.0,
8211                )?;
8212
8213                // Per-row append + attend: row r sees rows 0..r in KV (causal within the
8214                // draft), each through the b_n=1 serving kernels at its own t_kv.
8215                let q_dim = n_head * head_dim;
8216                let kv_dim = n_head_kv * head_dim;
8217                let mut attn = e.uninit(t * q_dim)?;
8218                let (kdk, kdv, ktb, vtb, len0, kv_local) = {
8219                    let kvl = cache.kv[il].as_ref().unwrap();
8220                    // [2T] interleaved k,v base pointers: entry pair z serves row z of
8221                    // the batched twins; the per-row fallback reads pair 0 (same cache
8222                    // for every row of one layer). Graph mode reads the ctx table.
8223                    let local: Option<CudaSlice<u64>> = match graph_cap {
8224                        Some(_) => None,
8225                        None => {
8226                            let s = &e.gpu.stream();
8227                            let (pk, _g) = kvl.k.device_ptr(s);
8228                            let (pv, _g2) = kvl.v.device_ptr(s);
8229                            let mut tbl = Vec::with_capacity(2 * t);
8230                            for _ in 0..t {
8231                                tbl.push(pk);
8232                                tbl.push(pv);
8233                            }
8234                            Some(e.htod_u64(&tbl)?)
8235                        }
8236                    };
8237                    (
8238                        kvl.kv_dim_k,
8239                        kvl.kv_dim_v,
8240                        kvl.k_tok_bytes,
8241                        kvl.v_tok_bytes,
8242                        kvl.len,
8243                        local,
8244                    )
8245                };
8246                let (kv_tbl, kv_off): (&CudaSlice<u64>, usize) = match graph_cap {
8247                    Some((tb, off, _)) => (tb, off),
8248                    None => (kv_local.as_ref().expect("built above"), 0),
8249                };
8250                // Slice 4 (fa/append rows — see dspark_fa_rows_on): the whole per-row
8251                // section batches into the z-batched serving twins when every row of
8252                // this round takes the v4-seqs arm on ONE fa_split_keys rung. Both
8253                // guards are evaluated at the round's FIRST and LAST t_kv — the
8254                // eligibility window (vec floor .. v4 max) and each split-ladder rung
8255                // are intervals in t_kv, so ends-inside means all-inside (the straddle
8256                // law). Appending all T rows before any attend is read-equivalent to
8257                // the interleaved order: row r's walk reads keys 0..len0+r only, and
8258                // rows > r land at slots it never touches; every written cache row is
8259                // the per-token appender's exact warp program (kernel-check pinned).
8260                let t_kv_first = len0 + 1;
8261                let t_kv_last = len0 + t;
8262                let rows_batched = t >= 2
8263                    && seqs_append
8264                    && batch_fa_on
8265                    && dspark_fa_rows_on()
8266                    // the z-batched twins read stacked rows at the CACHE's kv dims;
8267                    // the projection stack is [T, n_head_kv*head_dim] — they must be
8268                    // the same stride or row z misaligns (true for this family; the
8269                    // guard keeps any asymmetric-kv model on the per-row loop).
8270                    && kdk == kv_dim
8271                    && kdv == kv_dim
8272                    && crate::fa_seqs_eligible(t_kv_first, head_dim_global)
8273                    && crate::fa_seqs_eligible(t_kv_last, head_dim_global)
8274                    && crate::fa_split_keys(t_kv_first, cfg.n_head_kv as usize)
8275                        == crate::fa_split_keys(t_kv_last, cfg.n_head_kv as usize);
8276                // Sizing: eager = exact round bound; graph mode = the rung end (stride +
8277                // grid only — bytes proven equal above). Capture-time invariants refuse
8278                // loudly rather than bake a divergent body.
8279                let (size_kv_max, sp) = match graph_cap {
8280                    Some((_, _, rung)) => {
8281                        if !rows_batched {
8282                            return Err(format!(
8283                                "fa graph capture: layer {il} round is not batchable \
8284                                 (t_kv {t_kv_first}..{t_kv_last}) — the per-row fallback \
8285                                 must never be captured"
8286                            )
8287                            .into());
8288                        }
8289                        let sp_r = crate::fa_split_keys(rung, cfg.n_head_kv as usize);
8290                        if t_kv_last > rung
8291                            || sp_r != crate::fa_split_keys(t_kv_last, cfg.n_head_kv as usize)
8292                        {
8293                            return Err(format!(
8294                                "fa graph capture: rung {rung} does not cover round \
8295                                 t_kv {t_kv_first}..{t_kv_last} on one split ladder step"
8296                            )
8297                            .into());
8298                        }
8299                        (rung, sp_r)
8300                    }
8301                    None => (
8302                        t_kv_last,
8303                        crate::fa_split_keys(t_kv_last, cfg.n_head_kv as usize),
8304                    ),
8305                };
8306                if let Some((_, ctr)) = stream {
8307                    // STREAM ARM (2b): one batched dc append + the multi-row dc attention
8308                    // — the generic stream arm's exact shape (rows kernels are pinned
8309                    // byte-identical to the per-row programs by kernel-check). Host len
8310                    // stays a stale lower bound; the burst drain reconciles it.
8311                    let kvl = cache.kv[il].as_mut().unwrap();
8312                    e.append_kv_quantized_rows_dc(
8313                        &k,
8314                        &v,
8315                        &mut kvl.k,
8316                        &mut kvl.v,
8317                        ctr,
8318                        t,
8319                        kdk,
8320                        kdv,
8321                        ktb,
8322                        vtb,
8323                        Engine::kv_fp8_on(),
8324                    )?;
8325                    let upper = (kvl.len + t + 64).min(cache.max_ctx);
8326                    let k_view = e.view_u8(&kvl.k, upper * ktb);
8327                    let v_view = e.view_u8(&kvl.v, upper * vtb);
8328                    e.fa_decode_rows_dc(
8329                        &q, &k_view, &v_view, &mut attn, head_dim, n_head, n_head_kv, ctr, upper,
8330                        t, scale, ktb, vtb, 0, false,
8331                    )?;
8332                } else if rows_batched {
8333                    e.append_kv_quantized_seqs(
8334                        &k,
8335                        &v,
8336                        &kv_tbl.slice(kv_off..kv_off + 2 * t),
8337                        pos_d,
8338                        t,
8339                        kdk,
8340                        kdv,
8341                        ktb,
8342                        vtb,
8343                    )?;
8344                    if graph_cap.is_none() {
8345                        cache.kv[il].as_mut().unwrap().len += t;
8346                    }
8347                    e.fa_decode_batch_seqs_v4(
8348                        &q,
8349                        &kv_tbl.slice(kv_off..kv_off + 2 * t),
8350                        pos_d,
8351                        &mut attn,
8352                        head_dim,
8353                        n_head,
8354                        n_head_kv,
8355                        t,
8356                        size_kv_max,
8357                        scale,
8358                        sp,
8359                        ktb,
8360                        vtb,
8361                    )?;
8362                } else {
8363                    if pos_rows.is_none() {
8364                        // Stream-aware for symmetry with pos_d (the stream FA arm rides
8365                        // the dc rows kernels above and never reaches this fallback).
8366                        *pos_rows = Some(match stream {
8367                            Some((_, ctr)) => (0..t)
8368                                .map(|r| {
8369                                    let mut b = e.alloc_uninit::<i32>(1)?;
8370                                    e.i32_copy_add(ctr, &mut b, r as i32)?;
8371                                    Ok(b)
8372                                })
8373                                .collect::<Result<_, Box<dyn std::error::Error>>>()?,
8374                            None => (0..t)
8375                                .map(|r| e.htod_i32(&[(pos0 + r) as i32]))
8376                                .collect::<Result<_, _>>()?,
8377                        });
8378                    }
8379                    let pos_rows = pos_rows.as_ref().unwrap();
8380                    #[allow(clippy::needless_range_loop)]
8381                    // allow: the explicit index loop keeps the offset arithmetic visible and aligned with the device-side indexing
8382                    for r in 0..t {
8383                        // Owned per-row scratch: the b_n=1 kernels take packed batch buffers
8384                        // whose row 0 is this row (arithmetic-free materialization copies,
8385                        // same as decode's per-seq fallback arm).
8386                        let mut k_row = e.uninit(kv_dim)?;
8387                        e.dtod_copy_view(&k.slice(r * kv_dim..(r + 1) * kv_dim), &mut k_row)?;
8388                        let mut v_row = e.uninit(kv_dim)?;
8389                        e.dtod_copy_view(&v.slice(r * kv_dim..(r + 1) * kv_dim), &mut v_row)?;
8390                        let pos_row = &pos_rows[r];
8391                        let kvl = cache.kv[il].as_mut().unwrap();
8392                        if seqs_append {
8393                            e.append_kv_quantized_seqs(
8394                                &k_row,
8395                                &v_row,
8396                                &kv_tbl.slice(kv_off..kv_off + 2),
8397                                pos_row,
8398                                1,
8399                                kdk,
8400                                kdv,
8401                                ktb,
8402                                vtb,
8403                            )?;
8404                            kvl.len += 1;
8405                        } else {
8406                            e.append_kv_quantized_view(
8407                                &k_row.slice(0..kv_dim),
8408                                &v_row.slice(0..kv_dim),
8409                                &mut kvl.k,
8410                                &mut kvl.v,
8411                                kvl.len,
8412                                kvl.kv_dim_k,
8413                                kvl.kv_dim_v,
8414                                kvl.k_tok_bytes,
8415                                kvl.v_tok_bytes,
8416                                Engine::kv_fp8_on(),
8417                            )?;
8418                            kvl.len += 1;
8419                        }
8420                        let t_kv = kvl.len;
8421                        let mut q_row = e.uninit(q_dim)?;
8422                        e.dtod_copy_view(&q.slice(r * q_dim..(r + 1) * q_dim), &mut q_row)?;
8423                        let mut a_row = e.uninit(q_dim)?;
8424                        if batch_fa_on && crate::fa_seqs_eligible(t_kv, head_dim_global) {
8425                            let sp0_r = crate::fa_split_keys(t_kv, cfg.n_head_kv as usize);
8426                            e.fa_decode_batch_seqs_v4(
8427                                &q_row,
8428                                &kv_tbl.slice(kv_off..kv_off + 2),
8429                                pos_row,
8430                                &mut a_row,
8431                                head_dim,
8432                                n_head,
8433                                n_head_kv,
8434                                1,
8435                                t_kv,
8436                                scale,
8437                                sp0_r,
8438                                ktb,
8439                                vtb,
8440                            )?;
8441                        } else {
8442                            let k_view = e.view_u8(&kvl.k, t_kv * kvl.k_tok_bytes);
8443                            let v_view = e.view_u8(&kvl.v, t_kv * kvl.v_tok_bytes);
8444                            let mut a_view = a_row.slice_mut(0..q_dim);
8445                            e.fa_decode_kvmod_view(
8446                                &q_row.slice(0..q_dim),
8447                                &k_view,
8448                                &v_view,
8449                                &mut a_view,
8450                                head_dim,
8451                                n_head,
8452                                n_head_kv,
8453                                t_kv,
8454                                scale,
8455                                kvl.k_tok_bytes,
8456                                kvl.v_tok_bytes,
8457                                Engine::kv_fp8_on(),
8458                            )?;
8459                        }
8460                        e.dtod_copy_into(&a_row, &mut attn, r * q_dim)?;
8461                    }
8462                }
8463
8464                // Output gate (element-wise) + o-proj at m=T.
8465                let attn_g = match &gate {
8466                    Some(g) => {
8467                        let n = t * q_dim;
8468                        let mut gsig = e.uninit(n)?;
8469                        e.sigmoid(g, &mut gsig, n)?;
8470                        let mut ag = e.uninit(n)?;
8471                        e.mul(&attn, &gsig, &mut ag, n)?;
8472                        ag
8473                    }
8474                    None => attn,
8475                };
8476                e.matmul(&fa.wo, &attn_g, t)?
8477            }
8478        };
8479
8480        // ---- residual add + post_attn_norm + FFN at m=T (serving dispatch verbatim) ----
8481        let pnorm = layer.post_attn_norm.float_data();
8482        let mut x1 = e.uninit(t * n_embd)?;
8483        let mut zn = e.uninit(t * n_embd)?;
8484        e.add_rms_norm(x, &mixed, pnorm, &mut x1, &mut zn, n_embd, t, eps)?;
8485        let ffn_out = match &layer.ffn {
8486            crate::hybrid::Ffn::Dense {
8487                ffn_gate,
8488                ffn_up,
8489                ffn_down,
8490            } => {
8491                assert!(
8492                    self.cfg.m3.is_none(),
8493                    "qwen35 t-parallel verify: M3 swigluoai FFN not yet batched"
8494                );
8495                self.qwen35_tparallel_dense_ffn(e, ffn_gate, ffn_up, ffn_down, &zn, t, n_embd)?
8496            }
8497            crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il_zq8(e, m, &zn, None, t, il as u16)?,
8498        };
8499        let mut x2 = e.uninit(t * n_embd)?;
8500        e.add(&x1, &ffn_out, &mut x2, t * n_embd)?;
8501        // dspark drafter tap (no-op when no sink armed): post-layer residual verify rows
8502        self.dflash_tap(e, cache, il, &x2, t)?;
8503        Ok(x2)
8504    }
8505
8506    /// ONE t-parallel LINEAR layer (attn_norm + gdn mixer + post_attn_norm + FFN + tap) —
8507    /// the exact body the old in-loop Linear arm ran, extracted so the eager walk and the
8508    /// slice-3 captured segments execute the SAME code (a second copy is how dispatch
8509    /// mirrors drift — the verify_layers extraction lesson). Two deliberate changes, both
8510    /// bit-identical by construction:
8511    /// - the gdn ping-pong host swap moves from per-row to ONE end-of-body swap (t odd):
8512    ///   the device sequence is driven entirely by the 6-entry pointer table, which
8513    ///   already encodes both parities; the ckpt stash reads name row r's out buffer
8514    ///   directly (r even -> alt handle, odd -> canonical) — the same physical bytes the
8515    ///   legacy post-swap clone read.
8516    /// - `stash` (slice-3 ctx): persistent per-layer slabs written by copy_into instead of
8517    ///   per-row clone_dtod allocs — same bytes, capture-legal (no per-round host objects).
8518    ///   `table_src` = (persistent pointer table, offset) when the ctx owns the tables;
8519    ///   None builds the per-verify table exactly as before.
8520    #[allow(clippy::too_many_arguments)]
8521    fn qwen35_tparallel_linear_layer(
8522        &self,
8523        e: &Engine,
8524        il: usize,
8525        x: &CudaSlice<f32>,
8526        t: usize,
8527        cache: &mut Cache,
8528        ckpt: Option<&mut VerifyCkpt>,
8529        stash: Option<(&mut CudaSlice<f32>, &mut CudaSlice<f32>)>,
8530        table_src: Option<(&CudaSlice<u64>, usize)>,
8531    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8532        use cudarc::driver::DevicePtr;
8533        let cfg = &self.cfg;
8534        let n_embd = cfg.n_embd as usize;
8535        let eps = cfg.rms_eps;
8536        let layer = &self.layers[il];
8537        let Mixer::Linear(la) = &layer.mixer else {
8538            return Err("qwen35_tparallel_linear_layer on a non-linear layer".into());
8539        };
8540        // ---- attn_norm + q8_1 quantize at m=T (row-indexed == per-row) ----
8541        let anorm = layer.attn_norm.float_data();
8542        let mut xn = e.uninit(t * n_embd)?;
8543        e.rms_norm(x, anorm, &mut xn, n_embd, t, eps)?;
8544        let (hq, hd) = e.quantize_q8_1(&xn, t, n_embd)?;
8545
8546        let geometry = la.geometry;
8547        let d_state = geometry.key_head_dim as usize;
8548        let num_k = geometry.key_heads as usize;
8549        let num_v = geometry.value_heads as usize;
8550        let d_conv = geometry.conv_kernel as usize;
8551        let key_dim = d_state * num_k;
8552        let value_dim = geometry.value_head_dim as usize * num_v;
8553        let conv_dim = key_dim * 2 + value_dim;
8554        let gdn_scale = 1.0 / (d_state as f32).sqrt();
8555
8556        // ---- batched projections: one weight read for all T rows ----
8557        // GROUP-4 twin (trunk-kernels slice C): the whole 4-tuple in ONE launch, bit-identical
8558        // per (tensor, token, row) to the four singles; refused (layout/tier) or
8559        // MEMRA_TK_GDN_GROUP=0 -> the singles chain byte-for-byte.
8560        let (qkv_mixed, z, beta_raw, alpha) = match e.matmul_decode_exact_group4_pre(
8561            [&la.wqkv, &la.wqkv_gate, &la.ssm_beta, &la.ssm_alpha],
8562            &hq,
8563            &hd,
8564            t,
8565        )? {
8566            Some(mut g4) => {
8567                let alpha = g4.pop().unwrap();
8568                let beta_raw = g4.pop().unwrap();
8569                let z = g4.pop().unwrap();
8570                let qkv_mixed = g4.pop().unwrap();
8571                (qkv_mixed, z, beta_raw, alpha)
8572            }
8573            None => (
8574                e.matmul_pre(&la.wqkv, &hq, &hd, &xn, t)?,
8575                e.matmul_pre(&la.wqkv_gate, &hq, &hd, &xn, t)?,
8576                e.matmul_pre(&la.ssm_beta, &hq, &hd, &xn, t)?,
8577                e.matmul_pre(&la.ssm_alpha, &hq, &hd, &xn, t)?,
8578            ),
8579        };
8580        let beta_w = la.ssm_beta.out_features();
8581        let alpha_w = la.ssm_alpha.out_features();
8582        let qkv_w = la.wqkv.out_features();
8583
8584        // ---- per-row state chain through the b_n=1 serving kernels ----
8585        // 6-entry alternating pointer table expresses the ping-pong without a rebuild per
8586        // row: even rows scan s0 -> s1, odd rows s1 -> s0.
8587        let table_local: Option<CudaSlice<u64>> = match table_src {
8588            Some(_) => None,
8589            None => {
8590                let rl = cache.recur[il].as_ref().unwrap();
8591                let s = &e.gpu.stream();
8592                let (pc, _g0) = rl.conv_state.device_ptr(s);
8593                let (p0, _g1) = rl.ssm_state.device_ptr(s);
8594                let (p1, _g2) = rl.ssm_state_alt.device_ptr(s);
8595                Some(e.htod_u64(&[pc, p0, p1, pc, p1, p0])?)
8596            }
8597        };
8598        let (table, toff): (&CudaSlice<u64>, usize) = match table_src {
8599            Some((tb, off)) => (tb, off),
8600            None => (table_local.as_ref().unwrap(), 0),
8601        };
8602        let mut o_all = e.uninit(t * value_dim)?;
8603        let mut col_states: Option<Vec<(CudaSlice<f32>, CudaSlice<f32>)>> =
8604            if ckpt.is_some() && stash.is_none() && t >= 2 {
8605                Some(Vec::with_capacity(t - 1))
8606            } else {
8607                None
8608            };
8609        let mut stash = stash;
8610        // Per-row scratch reused across rows (uninit is cheap but not free at
8611        // 48 layers x T rows); row inputs/outputs pass as VIEWS into the packed
8612        // [T, ...] buffers — zero arithmetic-free copies in this loop.
8613        let mut conv_out = e.uninit(conv_dim)?;
8614        let mut q_l2 = e.uninit(value_dim)?;
8615        let mut k_l2 = e.uninit(value_dim)?;
8616        let mut v_gd = e.uninit(value_dim)?;
8617        let mut beta_b = e.uninit(num_v)?;
8618        let mut g_log = e.uninit(num_v)?;
8619        for r in 0..t {
8620            let base = toff + if r % 2 == 0 { 0 } else { 3 };
8621            let conv_view = table.slice(base..base + 1);
8622            let in_view = table.slice(base + 1..base + 2);
8623            let out_view = table.slice(base + 2..base + 3);
8624            e.ssm_conv1d_fused_decode_b_view(
8625                &qkv_mixed.slice(r * qkv_w..(r + 1) * qkv_w),
8626                &conv_view,
8627                la.ssm_conv1d.float_data(),
8628                &mut conv_out,
8629                conv_dim,
8630                d_conv,
8631                1,
8632            )?;
8633            e.gdn_prep_decode_b_view(
8634                &conv_out,
8635                &beta_raw.slice(r * beta_w..(r + 1) * beta_w),
8636                &alpha.slice(r * alpha_w..(r + 1) * alpha_w),
8637                la.ssm_dt.float_data(),
8638                la.ssm_a.float_data(),
8639                &mut q_l2,
8640                &mut k_l2,
8641                &mut v_gd,
8642                &mut beta_b,
8643                &mut g_log,
8644                d_state,
8645                num_v,
8646                num_k,
8647                key_dim,
8648                eps,
8649                conv_dim,
8650                1,
8651            )?;
8652            let mut o_row = o_all.slice_mut(r * value_dim..(r + 1) * value_dim);
8653            e.gdn_scan_s128_batched_view(
8654                &q_l2, &k_l2, &v_gd, &g_log, &beta_b, &in_view, &out_view, &mut o_row, num_v, 1,
8655                gdn_scale,
8656            )?;
8657            if r + 1 < t {
8658                // Row r's out buffer: even rows write s1 (the alt handle — no swaps ran),
8659                // odd rows write s0 — the same physical state the legacy post-swap
8660                // canonical clone read.
8661                let rl = cache.recur[il]
8662                    .as_ref()
8663                    .ok_or("qwen35 linear verify layer has no recurrent state")?;
8664                let ssm_src = if r % 2 == 0 {
8665                    &rl.ssm_state_alt
8666                } else {
8667                    &rl.ssm_state
8668                };
8669                match stash.as_mut() {
8670                    Some((conv_slab, ssm_slab)) => {
8671                        // BOTH stash reads go through the pointer table at run time: the
8672                        // ssm handles ping-pong between rounds, and the ctx (with its
8673                        // captured graphs) outlives the Cache — a fresh generation's
8674                        // conv/ssm buffers land at new addresses that only the per-round
8675                        // table refresh knows. A baked direct copy would read freed
8676                        // memory (parity was the slice-3 smoke divergence; cache
8677                        // lifetime is the cross-generation twin).
8678                        e.copy_indirect_src_f32(
8679                            &conv_view,
8680                            conv_slab,
8681                            r * conv_dim * (d_conv - 1),
8682                            conv_dim * (d_conv - 1),
8683                        )?;
8684                        // The ssm handles PING-PONG between rounds: a captured direct
8685                        // copy would bake the capture-time physical buffer and read the
8686                        // wrong parity after any odd-vt round (the slice-3 smoke
8687                        // divergence). Read the src address from row r's OUT table
8688                        // entry at run time — the same entry the scan just wrote.
8689                        e.copy_indirect_src_f32(
8690                            &out_view,
8691                            ssm_slab,
8692                            r * d_state * d_state * num_v,
8693                            d_state * d_state * num_v,
8694                        )?;
8695                    }
8696                    None => {
8697                        if let Some(states) = col_states.as_mut() {
8698                            states.push((e.clone_dtod(&rl.conv_state)?, e.clone_dtod(ssm_src)?));
8699                        }
8700                    }
8701                }
8702            }
8703        }
8704        // ONE end-of-body parity swap (t odd) — the legacy loop swapped per row; the net
8705        // handle motion is identical and the device sequence never read the handles.
8706        if t % 2 == 1 {
8707            let rl = cache.recur[il].as_mut().unwrap();
8708            std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
8709        }
8710        if let (Some(checkpoint), Some(states)) = (ckpt, col_states) {
8711            checkpoint.cols[il] = Some(states);
8712        }
8713
8714        // ---- batched gated norm + out-projection at m=T ----
8715        let mixed = if e.uses_q8_1_fast(&la.ssm_out) {
8716            let (gq, gd) = e.gated_rmsnorm_q8_1(
8717                &o_all,
8718                la.ssm_norm.float_data(),
8719                &z,
8720                d_state,
8721                t * num_v,
8722                eps,
8723            )?;
8724            let g0 = e.zeros(0)?;
8725            e.matmul_pre(&la.ssm_out, &gq, &gd, &g0, t)?
8726        } else {
8727            let mut gn = e.uninit(t * value_dim)?;
8728            e.gated_rmsnorm(
8729                &o_all,
8730                la.ssm_norm.float_data(),
8731                &z,
8732                &mut gn,
8733                d_state,
8734                t * num_v,
8735                eps,
8736            )?;
8737            e.matmul(&la.ssm_out, &gn, t)?
8738        };
8739
8740        // ---- residual add + post_attn_norm + FFN at m=T (serving dispatch verbatim) ----
8741        let pnorm = layer.post_attn_norm.float_data();
8742        let mut x1 = e.uninit(t * n_embd)?;
8743        let mut zn = e.uninit(t * n_embd)?;
8744        e.add_rms_norm(x, &mixed, pnorm, &mut x1, &mut zn, n_embd, t, eps)?;
8745        let ffn_out = match &layer.ffn {
8746            crate::hybrid::Ffn::Dense {
8747                ffn_gate,
8748                ffn_up,
8749                ffn_down,
8750            } => {
8751                assert!(
8752                    self.cfg.m3.is_none(),
8753                    "qwen35 t-parallel verify: M3 swigluoai FFN not yet batched"
8754                );
8755                self.qwen35_tparallel_dense_ffn(e, ffn_gate, ffn_up, ffn_down, &zn, t, n_embd)?
8756            }
8757            crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il_zq8(e, m, &zn, None, t, il as u16)?,
8758        };
8759        let mut x2 = e.uninit(t * n_embd)?;
8760        e.add(&x1, &ffn_out, &mut x2, t * n_embd)?;
8761        // dspark drafter tap (no-op when no sink armed): post-layer residual verify rows
8762        self.dflash_tap(e, cache, il, &x2, t)?;
8763        Ok(x2)
8764    }
8765
8766    /// PP-N STAGE SUBGRAPH of the verify trunk: layers `[lo, hi)` of `decode_step_t_core_stream`'s
8767    /// walk, verbatim. Enters with a MATERIALIZED `[T, n_embd]` residual (no pending fusion pair
8768    /// carried in from outside the range) and exits with the range's final residual materialized
8769    /// (the trailing add executed) — exactly the `decode_layers_eager(lo, hi)` contract, T rows
8770    /// instead of one.
8771    ///
8772    /// EXTRACTED (lane/pp2-spec 2026-08-06) rather than duplicated: `decode_step_t_core_stream` IS
8773    /// the single funnel every verify forward reaches, and its per-layer dispatch MIRRORING (norm
8774    /// fusion per layer, the t>=3/spec_m2 batched-linear window, the fused-q8 FFN chain, the
8775    /// decode-exact projections) is what makes verify bit-identical to eager decode. A second copy
8776    /// for the split arm is how those mirrors drift apart on the next lever. The unsplit body now
8777    /// calls this with `(0, n_layers)`, so the whole-trunk path and every stage range run the SAME
8778    /// code — there is no "split version" of the verify math.
8779    ///
8780    /// Bit-identity of a cut rests on the same kernel-check-pinned identity the eager arm's cut
8781    /// does — `add_rms_norm_q8_1 == add then rms_norm_q8_1` at nrows=T — because the ONLY thing a
8782    /// fence changes is that the cross-layer fusion carry breaks at `hi-1` and is re-materialized
8783    /// as an explicit `add`. `decode-batch-gate --mode ppspec` verifies end-to-end on real weights.
8784    #[allow(clippy::too_many_arguments)]
8785    fn verify_layers(
8786        &self,
8787        e: &Engine,
8788        mut x: CudaSlice<f32>,
8789        lo: usize,
8790        hi: usize,
8791        pos_d: &CudaSlice<i32>,
8792        pos0: usize,
8793        t: usize,
8794        cache: &mut Cache,
8795        mut ckpt: Option<&mut VerifyCkpt>,
8796        stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
8797        graphs: Option<&mut DsparkVerifyGraphs>,
8798    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8799        if self.sliding_gated_moe_batch_program() {
8800            if stream.is_some() {
8801                return Err(
8802                    "step35 has no ROUND-STREAM verify arm (the device-counter _dc twins \
8803                            cannot express the SWA offset KV view)"
8804                        .into(),
8805                );
8806            }
8807            return self.step35_verify_batch_layers(e, x, lo, hi, pos0, t, cache);
8808        }
8809        if self.batched_serving_numeric_class() {
8810            return self.qwen35_verify_batch_layers(
8811                e,
8812                x,
8813                lo,
8814                hi,
8815                pos0,
8816                t,
8817                cache,
8818                ckpt.take(),
8819                stream,
8820                graphs,
8821            );
8822        }
8823        let n_embd = self.cfg.n_embd as usize;
8824        let eps = self.cfg.rms_eps;
8825        // CROSS-LAYER ADD+NORM FUSION (lane/vt-fixes fix 2, mirroring decode_step_h's
8826        // launch-arc form): layer il's post-FFN residual add (x2 = x1 + ffn_out) and layer
8827        // il+1's attn_norm(+quantize) are consecutive row-wise ops — ONE add_rms_norm_q8_1
8828        // launch at nrows=t does all three (bit-identity pinned by the T-row kernel-check
8829        // arms). Carry the un-added (x1, ffn_out) pair; the fused launch materializes x2 (the
8830        // residual the next layer needs) as its `res` output. Falls back to the separate add
8831        // when the next layer is off the fused-q8 path.
8832        let mut pending: Option<(CudaSlice<f32>, CudaSlice<f32>)> = None;
8833        for il in lo..hi {
8834            let layer = &self.layers[il];
8835            // DISPATCH-MIRRORED attn-input RMSNorm (FP-order lesson #8): eager decode fuses the
8836            // 1024-thread rms_norm_q8_1 ONLY when every mixer projection is q8_1-fast; layers with
8837            // Float projections (ssm_beta/ssm_alpha on layers 1/2/4 of the 9B NVFP4 GGUF) take the
8838            // UNFUSED 256-thread rms_norm. The verify norm must mirror that PER-LAYER choice —
8839            // blockDim changes the sum-of-squares reduce order, and the ULP shift amplifies through
8840            // the GDN recurrence into argmax flips (measured: 9B text prompt, 1 ULP at layer 2 ->
8841            // 2.3e-1 logit maxdiff at the head -> K=1..8 divergence at a 0.03-margin token).
8842            let mixer_fast = self.mixer_in_q8_1_fast(e, &layer.mixer);
8843            let norm_fused = std::env::var("MEMRA_NO_FUSE_NORMQ").is_err() && mixer_fast;
8844            // BATCHED EPILOGUE RE-FUSE (lane/vt-fixes fix 2, 2026-08-03): when the norm is
8845            // dispatch-fused AND every consumer of `h` reads only its q8_1 form (Full mixer:
8846            // projections only; Linear mixer: the batched arm — the per-column fallback needs
8847            // f32 h), emit the attn-input norm DIRECTLY as q8_1 via `rms_norm_q8_1` at nrows=t
8848            // (row-indexed kernel — the T-row launch is the per-row m=1 program, kernel-check
8849            // pins bit-identity vs rms_norm_decode -> quantize_q8_1). Kills the standalone
8850            // quantize launch(es) + the f32 h HBM round-trip that decode never pays.
8851            // step35 (Full mixer) is the third case that needs f32 `h`: its verify arm is a
8852            // per-ROW replay of the eager decode mixer, whose `pre_q` contract is a single row —
8853            // a T-row q8_1 pair cannot be handed to it, and re-deriving per-row q8_1 from the f32
8854            // rows is exactly the dispatch being mirrored. Keep step35 on the unfused arm.
8855            let lin_q8_only = match &layer.mixer {
8856                Mixer::Linear(la) => {
8857                    (t >= 3 || (t == 2 && spec_m2())) && e.uses_q8_1_fast(&la.ssm_out)
8858                }
8859                Mixer::Full(_) if self.sliding_gated_moe_batch_program() => false,
8860                _ => true,
8861            };
8862            // NOTE decode.rs's take()-first lesson: take the pending pair BEFORE branching so
8863            // a non-fused layer still performs the residual add.
8864            let taken = pending.take();
8865            let (h, h_q8) = if norm_fused && lin_q8_only {
8866                let pair = match taken {
8867                    // fused add + attn_norm + q8_1: ONE launch resolves the carried residual
8868                    // AND emits this layer's mixer input pre-quantized. res -> x2 (= new x).
8869                    Some((x1p, f1p)) => {
8870                        let mut x2 = vbuf(e, t * n_embd)?; // fully written (res output)
8871                        let p = e.add_rms_norm_q8_1(
8872                            &x1p,
8873                            &f1p,
8874                            layer.attn_norm.float_data(),
8875                            &mut x2,
8876                            n_embd,
8877                            t,
8878                            eps,
8879                        )?;
8880                        x = x2;
8881                        p
8882                    }
8883                    None => e.rms_norm_q8_1(&x, layer.attn_norm.float_data(), n_embd, t, eps)?,
8884                };
8885                (e.zeros(0)?, Some(pair)) // h unused on this path (q8-only consumers)
8886            } else {
8887                if let Some((x1p, f1p)) = taken {
8888                    let mut x2 = vbuf(e, t * n_embd)?; // fully written by add
8889                    e.add(&x1p, &f1p, &mut x2, t * n_embd)?;
8890                    x = x2;
8891                }
8892                let mut h = vbuf(e, t * n_embd)?; // fully written by either rms_norm arm
8893                if norm_fused {
8894                    e.rms_norm_decode(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
8895                } else {
8896                    e.rms_norm(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
8897                }
8898                (h, None)
8899            };
8900            let h_q8_ref = h_q8.as_ref().map(|(q, d)| (q, d));
8901
8902            let mixed = match &layer.mixer {
8903                Mixer::Full(fa) => self.full_attn_verify(
8904                    e,
8905                    fa,
8906                    &h,
8907                    h_q8_ref,
8908                    pos_d,
8909                    t,
8910                    cache,
8911                    il,
8912                    stream.map(|(_, c)| c),
8913                )?,
8914                Mixer::Mla(_) => crate::hybrid::mla_path_unimplemented("speculative verify"),
8915                Mixer::Kda(_) => crate::hybrid::kda_path_unimplemented("speculative verify"),
8916                Mixer::Linear(la) => {
8917                    // BATCHED linear verify (2026-07-03, the MTP-profit lever): one T-token pass —
8918                    // batched projections (weight read ONCE, hits the m=2-4 weight-resident matvec),
8919                    // carried-state conv (ssm_conv1d_tm_state), GDN prep on the prefill kernels, and
8920                    // ONE gdn_scan whose internal sequential t-loop is the SAME recurrence as T
8921                    // chained T=1 steps (bit-identical). Falls back to the sequential per-column
8922                    // chain when T < d_conv-1 (conv ring update needs T >= pad) — or when ANY
8923                    // projection is off the q8_1 fast path: matmul_decode_exact would route a Float
8924                    // tensor to cuBLAS at m=t (different FP accumulation than eager's per-token
8925                    // GEMV), so mixed-dtype layers stay on the eager-identical per-column chain.
8926                    // MEMRA_SPEC_M2 (lane/spec-m2): the t==2 batch rides the same arm — the conv
8927                    // wrapper handles t<pad with a pure-copy ring rebuild; see spec_m2() header.
8928                    if (t >= 3 || (t == 2 && spec_m2()))
8929                        && mixer_fast
8930                        && e.uses_q8_1_fast(&la.ssm_out)
8931                    {
8932                        let want = ckpt.is_some();
8933                        let (out, stash) =
8934                            self.linear_attn_verify_t(e, la, &h, h_q8_ref, t, cache, il, want)?;
8935                        if let (Some(ck), Some(st)) = (ckpt.as_deref_mut(), stash) {
8936                            ck.gdn[il] = Some(st);
8937                        }
8938                        out
8939                    } else {
8940                        let mut out = vbuf(e, t * n_embd)?; // every col written by copy_into
8941                        let mut col_states: Option<Vec<(CudaSlice<f32>, CudaSlice<f32>)>> =
8942                            if ckpt.is_some() && t >= 2 {
8943                                Some(Vec::with_capacity(t - 1))
8944                            } else {
8945                                None
8946                            };
8947                        for col in 0..t {
8948                            let mut h_col = vbuf(e, n_embd)?; // fully written by copy_view_into
8949                            let src = h.slice(col * n_embd..(col + 1) * n_embd);
8950                            e.copy_view_into(&mut h_col, 0, &src, n_embd)?;
8951                            let m_col = self.linear_attn_decode(e, la, &h_col, cache, il)?;
8952                            e.copy_into(&mut out, col * n_embd, &m_col, n_embd)?;
8953                            // REPLAY-FREE ckpt: clone the chain's ACTUAL state after this column
8954                            // (pure dtod — cannot change any computed value). Last column skipped:
8955                            // rebuild targets are j <= t-1 columns.
8956                            if let Some(cs) = col_states.as_mut()
8957                                && col + 1 < t
8958                            {
8959                                let rl = cache.recur[il].as_ref().unwrap();
8960                                cs.push((
8961                                    e.clone_dtod(&rl.conv_state)?,
8962                                    e.clone_dtod(&rl.ssm_state)?,
8963                                ));
8964                            }
8965                        }
8966                        if let (Some(ck), Some(cs)) = (ckpt.as_deref_mut(), col_states) {
8967                            // ReplaySSM-assessment instrumentation (2026-07-30): the
8968                            // per-column clones are the only true state snapshots left in
8969                            // the verify (the batched path stashes INPUTS and replays).
8970                            if std::env::var("MEMRA_SPEC_STATS").as_deref() == Ok("1") {
8971                                static ONCE: std::sync::Once = std::sync::Once::new();
8972                                let bytes: usize =
8973                                    cs.iter().map(|(c, s)| (c.len() + s.len()) * 4).sum();
8974                                ONCE.call_once(|| eprintln!(
8975                                    "[verify-ckpt] per-column layer il={il}: {} clones, {:.2} MB/layer/round",
8976                                    cs.len(), bytes as f64 / 1e6));
8977                            }
8978                            ck.cols[il] = Some(cs);
8979                        }
8980                        out
8981                    }
8982                }
8983            };
8984            if spec_nan_scan_level() >= 2 {
8985                let mixed_width = mixed.len() / t;
8986                nan_scan_rows(
8987                    e,
8988                    &mixed,
8989                    t,
8990                    mixed_width,
8991                    &format!("verify layer {il} batched ATTN out pos0={pos0}"),
8992                )?;
8993            }
8994
8995            // DISPATCH-MIRRORED post-attn norm: eager residual_norm_ffn fuses add+norm+quant
8996            // (1024-thread add_rms_norm_q8_1) only for Dense FFNs whose gate+up are q8_1-fast;
8997            // otherwise (and for MoE) it runs the 256-thread fused add_rms_norm. Mirror per layer.
8998            let ffn_fuse = match &layer.ffn {
8999                crate::hybrid::Ffn::Dense {
9000                    ffn_gate, ffn_up, ..
9001                } => {
9002                    std::env::var("MEMRA_NO_FUSE_NORMQ").is_err()
9003                        && e.uses_q8_1_fast(ffn_gate)
9004                        && e.uses_q8_1_fast(ffn_up)
9005                }
9006                crate::hybrid::Ffn::Moe(_) => false,
9007            };
9008            // BATCHED EPILOGUE RE-FUSE (lane/vt-fixes fix 2): on the ffn_fuse path (Dense,
9009            // gate+up q8_1-fast, non-M3) the FFN input is emitted DIRECTLY as q8_1 by ONE
9010            // add_rms_norm_q8_1 launch at nrows=t (row-indexed kernel: the T-row launch is the
9011            // per-row m=1 program; kernel-check pins bit-identity vs the unfused
9012            // add_f32 -> rms_norm_decode -> quantize_q8_1 chain at T=2/4/5/8) — replacing the
9013            // add + rms_norm_decode launches AND the dual/singles' internal re-quantize.
9014            // M3's swigluoai must keep the f32 chain (the fused SwiGLU epilogue encodes plain
9015            // SiLU), mirroring residual_norm_ffn's m3 guard on the decode path.
9016            // step35: same guard per LAYER. A dense FFN's clamp is the SHEXP array (upstream's
9017            // one build_ffn serves dense + shared expert, llama-graph.cpp:1751), and verify MUST
9018            // mirror decode's dispatch or spec self-consistency fails.
9019            let dense_lim = self.cfg.clamp_shexp_at(il as u32);
9020            let fuse_q8 = ffn_fuse && self.cfg.m3.is_none() && dense_lim.is_none();
9021            let mut x1 = vbuf(e, t * n_embd)?; // fully written by add / add_rms_norm*
9022            let mut z = e.zeros(0)?; // replaced below on the unfused arms
9023            let z_q8 = if fuse_q8 {
9024                Some(e.add_rms_norm_q8_1(
9025                    &x,
9026                    &mixed,
9027                    layer.post_attn_norm.float_data(),
9028                    &mut x1,
9029                    n_embd,
9030                    t,
9031                    eps,
9032                )?)
9033            } else {
9034                let mut zf = vbuf(e, t * n_embd)?; // fully written by rms_norm_decode / add_rms_norm
9035                if ffn_fuse {
9036                    e.add(&x, &mixed, &mut x1, t * n_embd)?;
9037                    e.rms_norm_decode(
9038                        &x1,
9039                        layer.post_attn_norm.float_data(),
9040                        &mut zf,
9041                        n_embd,
9042                        t,
9043                        eps,
9044                    )?;
9045                } else {
9046                    e.add_rms_norm(
9047                        &x,
9048                        &mixed,
9049                        layer.post_attn_norm.float_data(),
9050                        &mut x1,
9051                        &mut zf,
9052                        n_embd,
9053                        t,
9054                        eps,
9055                    )?;
9056                }
9057                z = zf;
9058                None
9059            };
9060            if spec_nan_scan_level() >= 2 && !z.is_empty() {
9061                nan_scan_rows(
9062                    e,
9063                    &z,
9064                    t,
9065                    n_embd,
9066                    &format!("verify layer {il} post-attn norm z pos0={pos0}"),
9067                )?;
9068            }
9069            // DECODE-EXACT FFN projections: force MMVQ for gate/up/down at any T to match the
9070            // T=1 decode FP accumulation order. At T>=5 the generic matmul/matmul_pre falls to dp4a
9071            // (128-thread, different FP sum order). At T=2-4 the batched MMVQ is already bit-identical.
9072            let ffn_out = match &layer.ffn {
9073                crate::hybrid::Ffn::Dense {
9074                    ffn_gate,
9075                    ffn_up,
9076                    ffn_down,
9077                } => {
9078                    let n_ff = ffn_gate.out_features();
9079                    if let Some((zq, zd)) = z_q8.as_ref() {
9080                        // FUSED CHAIN (fix 2): pre-quantized z feeds the projections; the SwiGLU
9081                        // epilogue emits act pre-quantized for ffn_down (silu_mul_scaled_q8_1,
9082                        // bit-identical to silu_mul + quantize — kernel-check-pinned) with the
9083                        // NVFP4 macro-scales folded (deferred-scale dual: y*s inline == the
9084                        // scale_inplace store, value-exact) — the exact m=1 decode epilogue
9085                        // structure at nrows=t.
9086                        let pair = e
9087                            .matmul_decode_exact_dual_pre(ffn_gate, ffn_up, zq, zd, t)?
9088                            .map(|((g, gs), (u, us))| (g, gs, u, us));
9089                        let (gate, gs, up, us) = match pair {
9090                            Some(x4) => x4,
9091                            None => (
9092                                e.matmul_decode_exact_pre(ffn_gate, zq, zd, t)?,
9093                                1.0, // scale already applied inside _pre
9094                                e.matmul_decode_exact_pre(ffn_up, zq, zd, t)?,
9095                                1.0,
9096                            ),
9097                        };
9098                        if e.uses_q8_1_fast(ffn_down) {
9099                            let (aq, ad) = e.silu_mul_scaled_q8_1(&gate, &up, gs, us, t * n_ff)?;
9100                            e.matmul_decode_exact_pre(ffn_down, &aq, &ad, t)?
9101                        } else {
9102                            let mut act = vbuf(e, t * n_ff)?;
9103                            e.silu_mul_scaled(&gate, &up, gs, us, &mut act, t * n_ff)?;
9104                            e.matmul_decode_exact(ffn_down, &act, t)?
9105                        }
9106                    } else {
9107                        // UNFUSED (pre-fix) chain — MoE-adjacent/M3/off-fast layers, unchanged.
9108                        // DUAL gate+up batched twin (lane/verify-economics, 2026-08-02): one launch
9109                        // for the pair at t=2..8 — bit-identical per (tensor,token,row) to the two
9110                        // singles (kernel-check pins bitwise; MEMRA_SPEC_DUAL_T=0 reverts). None
9111                        // (non-NVFP4 / t outside the tier / seam off) -> the two singles, unchanged.
9112                        let (gate, up) =
9113                            match e.matmul_decode_exact_dual(ffn_gate, ffn_up, &z, t)? {
9114                                Some(pair) => pair,
9115                                None => (
9116                                    e.matmul_decode_exact(ffn_gate, &z, t)?,
9117                                    e.matmul_decode_exact(ffn_up, &z, t)?,
9118                                ),
9119                            };
9120                        let mut act = vbuf(e, t * n_ff)?; // fully written by ffn_act_lim
9121                        Self::ffn_act_lim(
9122                            e,
9123                            &self.cfg,
9124                            &gate,
9125                            &up,
9126                            1.0,
9127                            1.0,
9128                            dense_lim,
9129                            &mut act,
9130                            t * n_ff,
9131                        )?;
9132                        e.matmul_decode_exact(ffn_down, &act, t)?
9133                    }
9134                }
9135                crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il(e, m, &z, t, il as u16)?,
9136            };
9137            if spec_nan_scan_level() >= 2 {
9138                nan_scan_rows(
9139                    e,
9140                    &ffn_out,
9141                    t,
9142                    n_embd,
9143                    &format!("verify layer {il} batched FFN out pos0={pos0}"),
9144                )?;
9145            }
9146            if spec_nan_scan() {
9147                let mut residual = vbuf(e, t * n_embd)?;
9148                e.add(&x1, &ffn_out, &mut residual, t * n_embd)?;
9149                nan_scan_rows(
9150                    e,
9151                    &residual,
9152                    t,
9153                    n_embd,
9154                    &format!("verify layer {il} residual pos0={pos0}"),
9155                )?;
9156            }
9157            // CROSS-LAYER fusion: defer this layer's post-FFN residual add — the next layer's
9158            // fused-q8 attn norm folds it in (add_rms_norm_q8_1 == add; rms_norm; quantize,
9159            // kernel-check-pinned at nrows=T). Non-fused next layers add explicitly above.
9160            pending = Some((x1, ffn_out));
9161        }
9162        // RANGE's final add (no next norm INSIDE the range to fuse with; for the
9163        // whole-trunk call that is the last layer, whose next norm is output_norm — f32-out).
9164        if let Some((x1p, f1p)) = pending.take() {
9165            let mut x2 = vbuf(e, t * n_embd)?; // fully written by add
9166            e.add(&x1p, &f1p, &mut x2, t * n_embd)?;
9167            x = x2;
9168        }
9169        Ok(x)
9170    }
9171    /// BATCHED linear-attn verify (T=K+1): the whole layer in ~10 launches instead of T x the
9172    /// T=1 decode chain (T x ~12 launches + T weight reads of the four projections). The GDN
9173    /// recurrence itself is inherently sequential — gdn_scan_s128 runs its internal t-loop with
9174    /// the SAME per-token math as chained T=1 calls (bit-identical state evolution); everything
9175    /// around it (projections, conv, prep, gated norm, out-proj) batches. Advances conv ring +
9176    /// ssm state exactly like T sequential decode steps.
9177    /// `want_stash`: additionally RETAIN the gdn-scan inputs (pure buffer keep-alives, zero extra
9178    /// kernels) so a partial accept can rebuild the state after any column prefix (REPLAY-FREE).
9179    #[allow(clippy::too_many_arguments)]
9180    fn linear_attn_verify_t(
9181        &self,
9182        e: &Engine,
9183        la: &LinearAttnLayer,
9184        h: &CudaSlice<f32>,
9185        h_q8: Option<(&CudaSlice<i8>, &CudaSlice<f32>)>,
9186        t: usize,
9187        cache: &mut Cache,
9188        il: usize,
9189        want_stash: bool,
9190    ) -> Result<(CudaSlice<f32>, Option<GdnStash>), Box<dyn std::error::Error>> {
9191        let cfg = &self.cfg;
9192        let geometry = la.geometry;
9193        let d_state = geometry.key_head_dim as usize;
9194        let num_k = geometry.key_heads as usize;
9195        let num_v = geometry.value_heads as usize;
9196        let d_conv = geometry.conv_kernel as usize;
9197        let key_dim = d_state * num_k;
9198        let conv_dim = key_dim * 2 + geometry.value_head_dim as usize * num_v;
9199        let eps = cfg.rms_eps;
9200        let scale = 1.0 / (d_state as f32).sqrt();
9201
9202        // DECODE-EXACT projections: matmul_decode_exact forces the MMVQ (warp-per-row, 32-thread)
9203        // accumulation order for EVERY m, matching the T=1 decode path bit-for-bit. The generic
9204        // `matmul` at m>=5 falls to dp4a (128-thread, two-level reduce) which has a different FP
9205        // sum order — ULP differences propagate through gdn_scan and flip argmax on the 27B.
9206        // Q8 TRUNK-FUSION at T=1 (35B: wqkv+wqkv_gate both Q8_0): one fused2 launch, bit-identical
9207        // per (tensor,row) to the two m=1 MMVQ dispatches below — decode-exact contract holds.
9208        // VERIFY-TIER TRUNK FUSION (MEMRA_SPEC_FUSED_T, t=2-4): quantize h ONCE for every
9209        // fused-eligible same-input Q8_0 pair of this layer (35B wqkv+wqkv_gate; 9B
9210        // ssm_beta+ssm_alpha) — each fused2 batched launch then replaces two decode-exact
9211        // calls (each of which re-quantizes the same h + runs its own _b2/_b4 launch).
9212        // Bit-identical per (tensor,token,row) — see spec_fused_t().
9213        // BATCHED EPILOGUE RE-FUSE (lane/vt-fixes fix 2): `h_q8` = the attn-input norm emitted
9214        // directly as q8_1 by the caller's fused rms_norm_q8_1 (bit-identical to the unfused
9215        // chain, kernel-check-pinned). When present it REPLACES the standalone quantize below
9216        // and feeds every projection; the caller guaranteed all four input projections are
9217        // q8_1-fast. When absent, the old shared-quantize (fused-t window) stands.
9218        let h_q8_t = if h_q8.is_none()
9219            && spec_fused_t()
9220            && (2..=4).contains(&t)
9221            && ((e.uses_q8_1_fast(&la.wqkv) && e.uses_q8_1_fast(&la.wqkv_gate))
9222                || (e.uses_q8_1_fast(&la.ssm_beta) && e.uses_q8_1_fast(&la.ssm_alpha)))
9223        {
9224            Some(e.quantize_q8_1(h, t, cfg.n_embd as usize)?)
9225        } else {
9226            None
9227        };
9228        // one view: the caller's fused-norm q8 or this fn's own shared quantize.
9229        let hq8_any: Option<(&CudaSlice<i8>, &CudaSlice<f32>)> =
9230            h_q8.or(h_q8_t.as_ref().map(|(q, d)| (q, d)));
9231        let (qkv_mixed, z) = {
9232            let mut fused = None;
9233            if t == 1 && e.uses_q8_1_fast(&la.wqkv) && e.uses_q8_1_fast(&la.wqkv_gate) {
9234                let (hq, hd) = e.quantize_q8_1(h, 1, cfg.n_embd as usize)?;
9235                fused = e.matmul_q8_fused2(&la.wqkv, &la.wqkv_gate, &hq, &hd)?;
9236            } else if let Some((hq, hd)) = hq8_any
9237                && spec_fused_t()
9238                && (2..=4).contains(&t)
9239            {
9240                fused = e.matmul_q8_fused2_t(&la.wqkv, &la.wqkv_gate, hq, hd, t)?;
9241            }
9242            match (fused, hq8_any) {
9243                (Some(pair), _) => pair,
9244                (None, Some((hq, hd))) if h_q8.is_some() => (
9245                    e.matmul_decode_exact_pre(&la.wqkv, hq, hd, t)?,
9246                    e.matmul_decode_exact_pre(&la.wqkv_gate, hq, hd, t)?,
9247                ),
9248                (None, _) => (
9249                    e.matmul_decode_exact(&la.wqkv, h, t)?,
9250                    e.matmul_decode_exact(&la.wqkv_gate, h, t)?,
9251                ),
9252            }
9253        };
9254        // beta+alpha DUAL at T=1 (75% of p3 rounds run T=1 verify — p-min chain cuts): the dual
9255        // mr2 kernel is bit-identical per element to the m=1 MMVQ matmul_decode_exact dispatches
9256        // (same warp-per-row body, blockIdx.y picks the weight), so the decode-exact contract
9257        // holds; the run-spec battery is the arbiter. T>1 keeps the per-tensor decode-exact path.
9258        let (beta_raw, alpha) = if t == 1 {
9259            let (hq, hd) = e.quantize_q8_1(h, 1, cfg.n_embd as usize)?;
9260            match e.matmul_pre_dual_noscale(&la.ssm_beta, &la.ssm_alpha, &hq, &hd, 1)? {
9261                Some(((mut b, bs), (mut a, as_))) => {
9262                    if bs != 1.0 {
9263                        e.scale_inplace(&mut b, bs, la.ssm_beta.out_features())?;
9264                    }
9265                    if as_ != 1.0 {
9266                        e.scale_inplace(&mut a, as_, la.ssm_alpha.out_features())?;
9267                    }
9268                    (b, a)
9269                }
9270                // Q8_0 fused2 twin (9B stores beta/alpha as Q8_0): DISPATCH-MIRRORS the eager
9271                // decode's beta_alpha closure — the fused body is qmatvec_q8_0_mmvq verbatim,
9272                // bit-identical per row (kernel-check rel=0.00e0 gate), so decode==verify holds.
9273                None => match e.matmul_q8_fused2(&la.ssm_beta, &la.ssm_alpha, &hq, &hd)? {
9274                    Some((b, a)) => (b, a),
9275                    None => (
9276                        e.matmul_decode_exact(&la.ssm_beta, h, 1)?,
9277                        e.matmul_decode_exact(&la.ssm_alpha, h, 1)?,
9278                    ),
9279                },
9280            }
9281        } else {
9282            // fused-t twin (9B stores beta/alpha as Q8_0): same shared-quantize + one launch
9283            // contract as the wqkv pair above; 35B beta/alpha are Float -> None -> fallback.
9284            let mut nvfp4_fused = None;
9285            let mut q8_fused = None;
9286            if let Some((hq, hd)) = hq8_any {
9287                if t == 3 && std::env::var("MEMRA_NVFP4_AUX_DUAL").as_deref() != Ok("0") {
9288                    nvfp4_fused =
9289                        e.matmul_decode_exact_dual_pre(&la.ssm_beta, &la.ssm_alpha, hq, hd, t)?;
9290                    if nvfp4_fused.is_some() && std::env::var("MEMRA_DEBUG").is_ok() {
9291                        static ONCE: std::sync::Once = std::sync::Once::new();
9292                        ONCE.call_once(|| {
9293                            eprintln!("[memra] NVFP4 beta+alpha batched aux dual ENGAGED (t={t})")
9294                        });
9295                    }
9296                }
9297                if nvfp4_fused.is_none() && spec_fused_t() && (2..=4).contains(&t) {
9298                    q8_fused = e.matmul_q8_fused2_t(&la.ssm_beta, &la.ssm_alpha, hq, hd, t)?;
9299                }
9300            }
9301            if let Some(((mut b, bs), (mut a, as_))) = nvfp4_fused {
9302                if bs != 1.0 {
9303                    e.scale_inplace(&mut b, bs, t * la.ssm_beta.out_features())?;
9304                }
9305                if as_ != 1.0 {
9306                    e.scale_inplace(&mut a, as_, t * la.ssm_alpha.out_features())?;
9307                }
9308                (b, a)
9309            } else if let Some(pair) = q8_fused {
9310                pair
9311            } else {
9312                match hq8_any {
9313                    Some((hq, hd)) if h_q8.is_some() => (
9314                        e.matmul_decode_exact_pre(&la.ssm_beta, hq, hd, t)?,
9315                        e.matmul_decode_exact_pre(&la.ssm_alpha, hq, hd, t)?,
9316                    ),
9317                    _ => (
9318                        e.matmul_decode_exact(&la.ssm_beta, h, t)?,
9319                        e.matmul_decode_exact(&la.ssm_alpha, h, t)?,
9320                    ),
9321                }
9322            }
9323        };
9324
9325        // conv with CARRIED state + ring roll (T >= pad rides the input-column update kernel;
9326        // T < pad — the MEMRA_SPEC_M2 t=2 arm — rolls via the pure-copy ring rebuild).
9327        let rl = cache.recur[il].as_mut().unwrap();
9328        let mut conv_out = e.uninit(conv_dim * t)?;
9329        e.ssm_conv1d_tm_state(
9330            &qkv_mixed,
9331            &mut rl.conv_state,
9332            la.ssm_conv1d.float_data(),
9333            &mut conv_out,
9334            conv_dim,
9335            t,
9336            d_conv,
9337        )?;
9338
9339        // GDN prep via the prefill kernels (repack + L2 + sigmoid + glog), T-wide.
9340        let mut q_g = e.uninit(d_state * num_v * t)?;
9341        let mut k_g = e.uninit(d_state * num_v * t)?;
9342        let mut v_g = e.uninit(d_state * num_v * t)?;
9343        e.qkv_to_gdn_repack(
9344            &conv_out, &mut q_g, &mut k_g, &mut v_g, d_state, num_v, num_k, key_dim, t,
9345        )?;
9346        let mut q_l2 = e.uninit(d_state * num_v * t)?;
9347        e.l2_norm_decode(&q_g, &mut q_l2, d_state, num_v * t, eps)?;
9348        let mut k_l2 = e.uninit(d_state * num_v * t)?;
9349        e.l2_norm_decode(&k_g, &mut k_l2, d_state, num_v * t, eps)?;
9350        let mut beta = e.uninit(t * num_v)?;
9351        e.sigmoid(&beta_raw, &mut beta, t * num_v)?;
9352        let mut g_log = e.uninit(t * num_v)?;
9353        e.gdn_glog(
9354            &alpha,
9355            la.ssm_dt.float_data(),
9356            la.ssm_a.float_data(),
9357            &mut g_log,
9358            num_v,
9359            t,
9360        )?;
9361
9362        // ONE gdn_scan over T tokens from the carried state (internal sequential loop ==
9363        // T chained T=1 steps). Ping-pong the resident buffers like eager decode.
9364        let mut o = e.uninit(d_state * num_v * t)?;
9365        {
9366            let crate::cache::RecurLayer {
9367                ssm_state,
9368                ssm_state_alt,
9369                ..
9370            } = rl;
9371            e.gdn_scan_s128(
9372                &q_l2,
9373                &k_l2,
9374                &v_g,
9375                &g_log,
9376                &beta,
9377                ssm_state,
9378                ssm_state_alt,
9379                &mut o,
9380                num_v,
9381                t,
9382                scale,
9383            )?;
9384        }
9385        std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
9386
9387        // gated RMSNorm + out projection, T-wide. FUSED-QUANTIZE ARM (lane/vt-fixes fix 2,
9388        // mirroring the T=1 decode's launch-arc form): when ssm_out rides the q8_1 fast path,
9389        // emit q8_1 straight from the gated norm at nrows=num_v*t (row-indexed kernel, the
9390        // T-wide launch is the per-row program; kernel-check pins bit-identity vs
9391        // gated_rmsnorm -> quantize_q8_1 at T=1 and T=5) and feed the decode-exact dispatch
9392        // pre-quantized — one launch replaces norm + quantize. Fallback = the f32 chain.
9393        let out = if e.uses_q8_1_fast(&la.ssm_out) {
9394            let (gq, gd) =
9395                e.gated_rmsnorm_q8_1(&o, la.ssm_norm.float_data(), &z, d_state, num_v * t, eps)?;
9396            e.matmul_decode_exact_pre(&la.ssm_out, &gq, &gd, t)?
9397        } else {
9398            let mut gn = e.uninit(d_state * num_v * t)?;
9399            e.gated_rmsnorm(
9400                &o,
9401                la.ssm_norm.float_data(),
9402                &z,
9403                &mut gn,
9404                d_state,
9405                num_v * t,
9406                eps,
9407            )?;
9408            // DECODE-EXACT out-projection: same MMVQ path as the T=1 decode (ssm_out at m>=5
9409            // would fall to dp4a with a different FP reduction order — same class of bug as
9410            // the input projs).
9411            e.matmul_decode_exact(&la.ssm_out, &gn, t)?
9412        };
9413        let stash = if want_stash {
9414            Some(GdnStash {
9415                qkv_mixed,
9416                q_l2,
9417                k_l2,
9418                v_g,
9419                g_log,
9420                beta,
9421            })
9422        } else {
9423            None
9424        };
9425        Ok((out, stash))
9426    }
9427
9428    /// REPLAY-FREE partial-accept commit (2026-07-03): make the cache state == "committed through
9429    /// the first `j` verify columns" WITHOUT the legacy rollback + duplicate trunk replay.
9430    /// - Full-attn KV: truncate both the owning-stage shadow and every TP rank to snapshot + j.
9431    ///   The verify's appended rows for those columns are bit-identical to what an eager T=1
9432    ///   chain writes (the decode-exact contract the verify-probe gates), so keeping them ==
9433    ///   replaying them.
9434    /// - Linear layers, batched path: rebuild the conv ring by PURE COPIES (ring holds raw input
9435    ///   columns) and the ssm state by a prefix re-run of the SAME gdn_scan kernel (t=j) from the
9436    ///   snapshot state over the stash's identical inputs — the kernel's t-loop carries state in
9437    ///   registers and writes it once at the end, so iterations 0..j-1 are independent of T:
9438    ///   bit-identical to the verify's own state after j tokens == the eager chain state.
9439    /// - Linear layers, per-column path: restore the cloned actual state after column j-1.
9440    ///   Caller guarantees 1 <= j <= t-1 (j==0 rounds take the legacy rollback; j==t is full accept).
9441    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
9442    fn commit_verified_prefix(
9443        &self,
9444        e: &Engine,
9445        cache: &mut Cache,
9446        snap: &crate::cache::CacheSnapshot,
9447        ckpt: &VerifyCkpt,
9448        j: usize,
9449        kv_lens_done: bool,
9450        dev_j: Option<(&CudaSlice<u32>, usize, usize)>,
9451    ) -> Result<(), Box<dyn std::error::Error>> {
9452        // GDN geometry derives lazily inside recurrent-layer arms. Full-attention plans carry no
9453        // recurrent state and must never be forced through a synthetic SSM geometry.
9454        // Engine-bundle slice 1 (DSF-ROUNDCOST-20260820 §1.1): the per-column-arm restores
9455        // are 2 tiny D2D copies per linear layer (~96 dispatches/partial round on the q38
9456        // route). When every cols-arm layer shares uniform state sizes (single ssm cfg —
9457        // always true today), batch them into two `copy_batch_uniform_f32` launches. Bytes,
9458        // buffers and stream order are identical to the per-layer memcpy sequence; the
9459        // kernel-rebuild (gdn-stash) arm below is untouched. MEMRA_STATE_COPY_BATCH=0 reverts.
9460        let mut batched_cols = false;
9461        if state_copy_batch_on() && dev_j.is_none() {
9462            use cudarc::driver::DevicePtr;
9463            let s = &e.gpu.stream();
9464            let mut conv_pairs: Vec<(u64, u64)> = Vec::new();
9465            let mut ssm_pairs: Vec<(u64, u64)> = Vec::new();
9466            let (mut conv_words, mut ssm_words) = (0usize, 0usize);
9467            let mut uniform = true;
9468            for il in 0..self.layers.len() {
9469                let Some(rl) = cache.recur[il].as_ref() else {
9470                    continue;
9471                };
9472                if ckpt.gdn[il].is_some() {
9473                    continue; // kernel-rebuild arm restores below, per layer
9474                }
9475                let Some(cols) = &ckpt.cols[il] else {
9476                    continue; // missing-ckpt error surfaces in the main loop
9477                };
9478                let (c, st) = &cols[j - 1];
9479                if conv_pairs.is_empty() {
9480                    conv_words = c.len();
9481                    ssm_words = st.len();
9482                } else if c.len() != conv_words || st.len() != ssm_words {
9483                    uniform = false;
9484                    break;
9485                }
9486                let (pc, _g0) = c.device_ptr(s);
9487                let (dc, _g1) = rl.conv_state.device_ptr(s);
9488                let (ps, _g2) = st.device_ptr(s);
9489                let (ds, _g3) = rl.ssm_state.device_ptr(s);
9490                conv_pairs.push((pc, dc));
9491                ssm_pairs.push((ps, ds));
9492            }
9493            if uniform && !conv_pairs.is_empty() {
9494                let n = conv_pairs.len();
9495                let mut t = vec![0u64; 2 * n];
9496                for (k, &(src, dst)) in conv_pairs.iter().enumerate() {
9497                    t[k] = src;
9498                    t[n + k] = dst;
9499                }
9500                let conv_t = e.htod_u64(&t)?;
9501                for (k, &(src, dst)) in ssm_pairs.iter().enumerate() {
9502                    t[k] = src;
9503                    t[n + k] = dst;
9504                }
9505                let ssm_t = e.htod_u64(&t)?;
9506                e.copy_batch_uniform_f32(&conv_t, n, conv_words)?;
9507                e.copy_batch_uniform_f32(&ssm_t, n, ssm_words)?;
9508                batched_cols = true;
9509            }
9510        }
9511        for il in 0..self.layers.len() {
9512            if let (Some(kvl), Some(saved)) = (cache.kv[il].as_mut(), snap.kv_len[il]) {
9513                kvl.len = saved + j;
9514                // devacc 3a: spec_rollback_kv already wrote len_d on-device (same value).
9515                if !kv_lens_done {
9516                    e.set_i32_one(&mut kvl.len_d, kvl.len as i32)?;
9517                }
9518            }
9519            if let Some(rl) = cache.recur[il].as_mut() {
9520                let Mixer::Linear(linear) = &self.layers[il].mixer else {
9521                    return Err(format!("recurrent cache layer {il} has no GDN plan").into());
9522                };
9523                let geometry = linear.geometry;
9524                let d_state = geometry.key_head_dim as usize;
9525                let num_k = geometry.key_heads as usize;
9526                let num_v = geometry.value_heads as usize;
9527                let d_conv = geometry.conv_kernel as usize;
9528                let conv_dim = d_state * num_k * 2 + geometry.value_head_dim as usize * num_v;
9529                let scale = 1.0 / (d_state as f32).sqrt();
9530                if let Some(st) = &ckpt.gdn[il] {
9531                    let ring_old = snap.conv[il].as_ref().expect("snapshot missing conv");
9532                    let state_in = snap.ssm[il].as_ref().expect("snapshot missing ssm");
9533                    if let Some((acc, base, t_v)) = dev_j {
9534                        // 3b: j read on-device (_dc twins, same bodies; full accept early-exits).
9535                        e.ssm_conv_ring_rebuild_dc(
9536                            &st.qkv_mixed,
9537                            ring_old,
9538                            &mut rl.conv_state,
9539                            conv_dim,
9540                            acc,
9541                            base,
9542                            t_v,
9543                            d_conv,
9544                        )?;
9545                        let mut o = e.uninit(d_state * num_v * j.max(1))?;
9546                        e.gdn_scan_s128_dc(
9547                            &st.q_l2,
9548                            &st.k_l2,
9549                            &st.v_g,
9550                            &st.g_log,
9551                            &st.beta,
9552                            state_in,
9553                            &mut rl.ssm_state,
9554                            &mut o,
9555                            num_v,
9556                            acc,
9557                            base,
9558                            t_v,
9559                            scale,
9560                        )?;
9561                    } else {
9562                        e.ssm_conv_ring_rebuild(
9563                            &st.qkv_mixed,
9564                            ring_old,
9565                            &mut rl.conv_state,
9566                            conv_dim,
9567                            j,
9568                            d_conv,
9569                        )?;
9570                        let mut o = e.uninit(d_state * num_v * j)?; // scan output, discarded
9571                        e.gdn_scan_s128(
9572                            &st.q_l2,
9573                            &st.k_l2,
9574                            &st.v_g,
9575                            &st.g_log,
9576                            &st.beta,
9577                            state_in,
9578                            &mut rl.ssm_state,
9579                            &mut o,
9580                            num_v,
9581                            j,
9582                            scale,
9583                        )?;
9584                    }
9585                } else if let Some(cols) = &ckpt.cols[il] {
9586                    if !batched_cols {
9587                        let (c, s) = &cols[j - 1];
9588                        e.copy_into(&mut rl.conv_state, 0, c, c.len())?;
9589                        e.copy_into(&mut rl.ssm_state, 0, s, s.len())?;
9590                    }
9591                } else {
9592                    return Err(
9593                        "commit_verified_prefix: verify ckpt missing for linear layer".into(),
9594                    );
9595                }
9596            }
9597        }
9598        self.restore_step_tp_kv_verified_prefix(e, cache, snap, j, true)?;
9599        cache.pos = snap.pos + j;
9600        Ok(())
9601    }
9602
9603    /// ROUND-STREAM: recur restore with device-j (the _dc twins; full accept early-exits
9604    /// in-kernel). Requires the batched-linear stash on every linear layer (stream gate).
9605    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
9606    fn commit_verified_prefix_stream(
9607        &self,
9608        e: &Engine,
9609        cache: &mut Cache,
9610        snap: &crate::cache::CacheSnapshot,
9611        ckpt: &VerifyCkpt,
9612        acc: &CudaSlice<u32>,
9613        base: usize,
9614        t_v: usize,
9615    ) -> Result<(), Box<dyn std::error::Error>> {
9616        for il in 0..self.layers.len() {
9617            if let Some(rl) = cache.recur[il].as_mut() {
9618                let Mixer::Linear(linear) = &self.layers[il].mixer else {
9619                    return Err(format!("recurrent cache layer {il} has no GDN plan").into());
9620                };
9621                let geometry = linear.geometry;
9622                let d_state = geometry.key_head_dim as usize;
9623                let num_k = geometry.key_heads as usize;
9624                let num_v = geometry.value_heads as usize;
9625                let d_conv = geometry.conv_kernel as usize;
9626                let conv_dim = d_state * num_k * 2 + geometry.value_head_dim as usize * num_v;
9627                let scale = 1.0 / (d_state as f32).sqrt();
9628                let st = ckpt.gdn[il]
9629                    .as_ref()
9630                    .ok_or("stream restore: batched-linear stash missing")?;
9631                let ring_old = snap.conv[il].as_ref().expect("snapshot missing conv");
9632                let state_in = snap.ssm[il].as_ref().expect("snapshot missing ssm");
9633                e.ssm_conv_ring_rebuild_dc(
9634                    &st.qkv_mixed,
9635                    ring_old,
9636                    &mut rl.conv_state,
9637                    conv_dim,
9638                    acc,
9639                    base,
9640                    t_v,
9641                    d_conv,
9642                )?;
9643                let mut o = e.uninit(d_state * num_v * t_v)?;
9644                e.gdn_scan_s128_dc(
9645                    &st.q_l2,
9646                    &st.k_l2,
9647                    &st.v_g,
9648                    &st.g_log,
9649                    &st.beta,
9650                    state_in,
9651                    &mut rl.ssm_state,
9652                    &mut o,
9653                    num_v,
9654                    acc,
9655                    base,
9656                    t_v,
9657                    scale,
9658                )?;
9659            }
9660        }
9661        Ok(())
9662    }
9663
9664    /// EAGLE3 aux-capturing verify forward over `tokens` (T) — mirrors `decode_step_t_h` exactly
9665    /// (same KV append, same causal verify, same recur advance) but ALSO clones the aux residual-
9666    /// stream hiddens (blocks in `aux_layers`) for TWO columns: the LAST column (always) and the
9667    /// optional `pred_col` (the EAGLE seed = bonus's predecessor). Returns
9668    /// (all_T_logits host, last_col_aux, pred_col_aux?). Used by the EAGLE3 orchestrator's commit.
9669    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
9670    pub fn decode_step_t_aux2(
9671        &self,
9672        e: &Engine,
9673        tokens: &[u32],
9674        pos0: usize,
9675        cache: &mut Cache,
9676        aux_layers: &[usize],
9677        pred_col: Option<usize>,
9678    ) -> Result<
9679        (Vec<f32>, Vec<CudaSlice<f32>>, Option<Vec<CudaSlice<f32>>>),
9680        Box<dyn std::error::Error>,
9681    > {
9682        cache.ensure_usable("decode_step_t_aux2")?;
9683        let cfg = &self.cfg;
9684        let n_embd = cfg.n_embd as usize;
9685        let eps = cfg.rms_eps;
9686        let t = tokens.len();
9687        let pos_vec: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
9688        let pos_d = e.htod_i32(&pos_vec)?;
9689        let mut x = e.htod(&self.embd.try_gather(n_embd, tokens)?)?;
9690        let mut aux_last: Vec<CudaSlice<f32>> = Vec::with_capacity(aux_layers.len());
9691        let mut aux_pred: Vec<CudaSlice<f32>> = Vec::new();
9692        let want_pred = pred_col.is_some();
9693
9694        for (il, layer) in self.layers.iter().enumerate() {
9695            // DISPATCH-MIRRORED norms (FP-order lesson #8) — see decode_step_t_h_emb.
9696            let mixer_fast = self.mixer_in_q8_1_fast(e, &layer.mixer);
9697            let norm_fused = std::env::var("MEMRA_NO_FUSE_NORMQ").is_err() && mixer_fast;
9698            let mut h = vbuf(e, t * n_embd)?; // fully written by either rms_norm arm
9699            if norm_fused {
9700                e.rms_norm_decode(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
9701            } else {
9702                e.rms_norm(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
9703            }
9704            let mixed = match &layer.mixer {
9705                Mixer::Full(fa) => {
9706                    self.full_attn_verify(e, fa, &h, None, &pos_d, t, cache, il, None)?
9707                }
9708                Mixer::Mla(_) => {
9709                    crate::hybrid::mla_path_unimplemented("auxiliary T-parallel decode")
9710                }
9711                Mixer::Kda(_) => crate::hybrid::kda_path_unimplemented("aux decode step"),
9712                Mixer::Linear(la) => {
9713                    let mut out = e.zeros(t * n_embd)?;
9714                    for col in 0..t {
9715                        let mut h_col = e.zeros(n_embd)?;
9716                        let src = h.slice(col * n_embd..(col + 1) * n_embd);
9717                        e.copy_view_into(&mut h_col, 0, &src, n_embd)?;
9718                        let m_col = self.linear_attn_decode(e, la, &h_col, cache, il)?;
9719                        e.copy_into(&mut out, col * n_embd, &m_col, n_embd)?;
9720                    }
9721                    out
9722                }
9723            };
9724            let ffn_fuse = match &layer.ffn {
9725                crate::hybrid::Ffn::Dense {
9726                    ffn_gate, ffn_up, ..
9727                } => {
9728                    std::env::var("MEMRA_NO_FUSE_NORMQ").is_err()
9729                        && e.uses_q8_1_fast(ffn_gate)
9730                        && e.uses_q8_1_fast(ffn_up)
9731                }
9732                crate::hybrid::Ffn::Moe(_) => false,
9733            };
9734            let mut x1 = vbuf(e, t * n_embd)?; // fully written by add / add_rms_norm
9735            let mut z = vbuf(e, t * n_embd)?; // fully written by rms_norm_decode / add_rms_norm
9736            if ffn_fuse {
9737                e.add(&x, &mixed, &mut x1, t * n_embd)?;
9738                e.rms_norm_decode(
9739                    &x1,
9740                    layer.post_attn_norm.float_data(),
9741                    &mut z,
9742                    n_embd,
9743                    t,
9744                    eps,
9745                )?;
9746            } else {
9747                e.add_rms_norm(
9748                    &x,
9749                    &mixed,
9750                    layer.post_attn_norm.float_data(),
9751                    &mut x1,
9752                    &mut z,
9753                    n_embd,
9754                    t,
9755                    eps,
9756                )?;
9757            }
9758            let ffn_out = match &layer.ffn {
9759                crate::hybrid::Ffn::Dense {
9760                    ffn_gate,
9761                    ffn_up,
9762                    ffn_down,
9763                } => {
9764                    let n_ff = ffn_gate.out_features();
9765                    let gate = e.matmul_decode_exact(ffn_gate, &z, t)?;
9766                    let up = e.matmul_decode_exact(ffn_up, &z, t)?;
9767                    let mut act = vbuf(e, t * n_ff)?; // fully written by ffn_act_lim
9768                    // dense FFN clamp = the SHEXP array (upstream build_ffn serves both).
9769                    Self::ffn_act_lim(
9770                        e,
9771                        &self.cfg,
9772                        &gate,
9773                        &up,
9774                        1.0,
9775                        1.0,
9776                        self.cfg.clamp_shexp_at(il as u32),
9777                        &mut act,
9778                        t * n_ff,
9779                    )?;
9780                    e.matmul_decode_exact(ffn_down, &act, t)?
9781                }
9782                crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il(e, m, &z, t, il as u16)?,
9783            };
9784            let mut x2 = vbuf(e, t * n_embd)?; // fully written by add
9785            e.add(&x1, &ffn_out, &mut x2, t * n_embd)?;
9786            if aux_layers.contains(&il) {
9787                let mut a = e.zeros(n_embd)?;
9788                e.copy_view_into(&mut a, 0, &x2.slice((t - 1) * n_embd..t * n_embd), n_embd)?;
9789                aux_last.push(a);
9790                if let Some(pc) = pred_col {
9791                    let mut ap = e.zeros(n_embd)?;
9792                    e.copy_view_into(
9793                        &mut ap,
9794                        0,
9795                        &x2.slice(pc * n_embd..(pc + 1) * n_embd),
9796                        n_embd,
9797                    )?;
9798                    aux_pred.push(ap);
9799                }
9800            }
9801            x = x2;
9802        }
9803        let mut hn = vbuf(e, t * n_embd)?; // fully written by rms_norm_decode
9804        e.rms_norm_decode(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
9805        let logits = e.matmul_decode_exact(&self.output, &hn, t)?;
9806        let host = e.dtoh(&logits)?;
9807        cache.pos += t;
9808        Ok((
9809            host,
9810            aux_last,
9811            if want_pred { Some(aux_pred) } else { None },
9812        ))
9813    }
9814
9815    /// step35 SPEC-VERIFY attention over T query tokens — a per-row REPLAY of the eager
9816    /// `step35_decode_attn`.
9817    ///
9818    /// WHY A REPLAY AND NOT A BATCHED TWIN. The verify's whole job is to be bit-identical to what
9819    /// the eager decode would have computed for the same tokens; that is what makes greedy spec
9820    /// decode exact (run-spec asserts token identity for K=1..8). Every other verify arm in this
9821    /// file earns that identity by carefully mirroring dispatch (`matmul_decode_exact` to force
9822    /// MMVQ at any m, per-layer `ffn_fuse` mirroring, per-row `fa_decode` key bounds). step35
9823    /// stacks FOUR more per-layer degrees of freedom on top of that — per-layer `n_head`
9824    /// (64 full / 96 SWA), per-layer rotary width (64 full / 128 SWA), per-layer rope base, and a
9825    /// SEPARATE `attn_gate` tensor whose projection shares the attn-normed input — and its SWA
9826    /// layers attend through a token-OFFSET view whose offset is a function of the ABSOLUTE
9827    /// position of each query row. A batched twin would have to reproduce all of that AND the
9828    /// per-row offset in one launch; the offset alone rules out the existing rows kernels (they
9829    /// take one `base_len`, not a per-row offset).
9830    ///
9831    /// So this arm calls the eager path itself, once per row, on the same cache. Identity is then
9832    /// true BY CONSTRUCTION rather than by mirroring: row r runs exactly the kernel sequence that
9833    /// eager decode step r runs (same projections, same q8_1 fusion decision, same append, same
9834    /// view arithmetic, same `fa_decode_kvmod`, same gate), because it IS that code. Cost: T x the
9835    /// eager decode mixer instead of one batched pass — the same trade the generic arm's `else`
9836    /// per-row loop already accepts when `fa_rows_eligible` says no. Correctness first; a batched
9837    /// step35 twin is a perf lane's job and must be gated against this arm.
9838    ///
9839    /// The `h_q8` pre-quantized pair from the caller's fused norm is NOT forwarded: it is a
9840    /// T-row buffer and `step35_decode_attn`'s `pre_q` contract is one row. Instead each row's
9841    /// f32 `h` slice is handed over and the callee re-derives its own q8_1 exactly as eager decode
9842    /// does (`quantize_q8_1(h, 1, n_embd)`) — which is the dispatch being mirrored. Callers that
9843    /// took the fused arm therefore MUST still pass a live `h`; `step35_verify` asserts that.
9844    #[allow(clippy::too_many_arguments)]
9845    fn step35_verify(
9846        &self,
9847        e: &Engine,
9848        fa: &FullAttnLayer,
9849        h: &CudaSlice<f32>,
9850        h_q8: Option<(&CudaSlice<i8>, &CudaSlice<f32>)>,
9851        t: usize,
9852        cache: &mut Cache,
9853        il: usize,
9854    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9855        let n_embd = self.cfg.n_embd as usize;
9856        // The fused (h-less) attn-norm arm hands `h` as a zero-length placeholder. This arm needs
9857        // the f32 rows, so the caller must not take that lever for step35 — enforced at the call
9858        // site by the sliding-gated-MoE `Mixer::Full(_) => false` arm of
9859        // `lin_q8_only` in `decode_step_t_core_stream`, and asserted here so a future caller
9860        // cannot regress it into silently reading an empty buffer.
9861        assert_eq!(
9862            h.len(),
9863            t * n_embd,
9864            "step35_verify needs the f32 attn-normed rows ([t*n_embd]); the caller took the \
9865             fused q8-only norm arm (h_q8={}) — step35 must stay on the unfused arm",
9866            h_q8.is_some()
9867        );
9868        // ROW WIDTH IS n_embd, NOT n_head*head_dim: `step35_decode_attn` returns the mixer output
9869        // AFTER `wo`, so a row is [n_embd] — the same contract the generic arm's
9870        // `matmul_decode_exact(&fa.wo, &attn_g, t)` return has. Sizing this buffer from the
9871        // per-layer head geometry (8192 on full-attn, 12288 on SWA) instead overran the row on the
9872        // FIRST copy and panicked inside `copy_into`'s `CudaView::slice` unwrap
9873        // (raw/mtp-bt-20260806T212127Z.log frames 12-13).
9874        let mut out = vbuf(e, t * n_embd)?; // each row fully written by the copy below
9875        for r in 0..t {
9876            // Absolute position of this query row. `cache.pos` is the committed length at round
9877            // start and every row before r has already been appended by this loop, so the r-th
9878            // verify token sits at cache.pos + r — the same position eager decode would give it.
9879            let pos_d = e.htod_i32(&[(cache.pos + r) as i32])?;
9880            let mut h_row = vbuf(e, n_embd)?; // fully written by copy_view_into
9881            e.copy_view_into(
9882                &mut h_row,
9883                0,
9884                &h.slice(r * n_embd..(r + 1) * n_embd),
9885                n_embd,
9886            )?;
9887            // THE eager decode mixer: appends this row's K/V at kvl.len, advances it, then
9888            // attends over the (SWA-offset) view. Post-`wo`, same contract as this fn returns.
9889            let o = self.step35_decode_attn(e, fa, il, &h_row, None, &pos_d, cache)?;
9890            debug_assert_eq!(
9891                o.len(),
9892                n_embd,
9893                "step35_decode_attn returns post-wo [n_embd]"
9894            );
9895            e.copy_into(&mut out, r * n_embd, &o, n_embd)?;
9896        }
9897        Ok(out)
9898    }
9899
9900    /// Full-attention mixer over T query tokens with a GROWING resident KV (verify path, §D.3).
9901    /// Appends the T new K/V columns to cache.kv[il] then attends causally over [0..len) via
9902    /// fa_prefill. Token-major [T, kv_dim] projection layout == cache row layout (single copy).
9903    #[allow(clippy::too_many_arguments)]
9904    fn full_attn_verify(
9905        &self,
9906        e: &Engine,
9907        fa: &FullAttnLayer,
9908        h: &CudaSlice<f32>,
9909        h_q8: Option<(&CudaSlice<i8>, &CudaSlice<f32>)>,
9910        pos_d: &CudaSlice<i32>,
9911        t: usize,
9912        cache: &mut Cache,
9913        il: usize,
9914        stream_ctr: Option<&CudaSlice<i32>>,
9915    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9916        // step35: the generic geometry below is wrong for this arch (per-layer n_head, partial
9917        // per-layer rope, the SWA offset view, and a SEPARATE head-wise gate tensor), so it takes
9918        // its own arm. A verify that silently computes different attention than decode defeats the
9919        // whole self-consistency gate, so the arm is a per-row REPLAY of `step35_decode_attn`
9920        // rather than a batched twin — see `step35_verify` for why that is the exactness-correct
9921        // shape and not laziness.
9922        if self.sliding_gated_moe_batch_program() {
9923            if stream_ctr.is_some() {
9924                return Err(
9925                    "step35 has no ROUND-STREAM verify arm (the device-counter _dc twins \
9926                            cannot express the SWA offset KV view; same root cause as the dc \
9927                            decode refusal) — run spec without the stream arm"
9928                        .into(),
9929                );
9930            }
9931            return self.step35_verify(e, fa, h, h_q8, t, cache, il);
9932        }
9933        let cfg = &self.cfg;
9934        let geometry = cfg.full_attention_geometry_at(il as u32);
9935        let n_head = geometry.n_head as usize;
9936        let n_head_kv = geometry.n_head_kv as usize;
9937        let head_dim = geometry.head_dim_k as usize;
9938        let eps = cfg.rms_eps;
9939        let scale = geometry.attention_scale();
9940        let n_embd = cfg.n_embd as usize;
9941
9942        // DECODE-EXACT Q/K/V projections: matmul_decode_exact forces the MMVQ (warp-per-row) path
9943        // for every m, matching the T=1 decode's FP accumulation order. matmul_pre at m>=5 would
9944        // fall to dp4a (128-thread, two-level reduce) with a different FP sum order.
9945        // Q8 TRUNK-FUSION at T=1: DISPATCH-MIRRORS the eager decode's fused3 (bit-identical body).
9946        // BATCHED EPILOGUE RE-FUSE (lane/vt-fixes fix 2): `h_q8` = the attn-input norm's q8_1
9947        // form emitted by the fused rms_norm_q8_1 (bit-identical to rms_norm_decode ->
9948        // quantize_q8_1, kernel-check-pinned). When present (caller checked mixer q8_1-fast),
9949        // every projection consumes it — `h` may be a zero-len placeholder and must not be read.
9950        let (qf, mut k, v) = if let Some(mut qkv) = self.full_attn_tp_qkv(e, fa, h, t)? {
9951            let v = qkv.pop().ok_or("full-attention TP verify QKV omitted V")?;
9952            let k = qkv.pop().ok_or("full-attention TP verify QKV omitted K")?;
9953            let q = qkv.pop().ok_or("full-attention TP verify QKV omitted Q")?;
9954            if !qkv.is_empty() {
9955                return Err("full-attention TP verify QKV returned extra projections".into());
9956            }
9957            (q, k, v)
9958        } else {
9959            let mut fused = None;
9960            let qkv_fast =
9961                e.uses_q8_1_fast(&fa.wq) && e.uses_q8_1_fast(&fa.wk) && e.uses_q8_1_fast(&fa.wv);
9962            if t == 1 && qkv_fast {
9963                let (hq_o, hd_o);
9964                let (hq, hd): (&CudaSlice<i8>, &CudaSlice<f32>) = match h_q8 {
9965                    Some(p) => p,
9966                    None => {
9967                        (hq_o, hd_o) = e.quantize_q8_1(h, 1, n_embd)?;
9968                        (&hq_o, &hd_o)
9969                    }
9970                };
9971                fused = e.matmul_q8_fused3(&fa.wq, &fa.wk, &fa.wv, hq, hd)?;
9972            } else if spec_fused_t() && (2..=4).contains(&t) && qkv_fast {
9973                // VERIFY-TIER TRUNK FUSION (MEMRA_SPEC_FUSED_T): one shared quantize + one
9974                // fused3 batched launch replaces three decode-exact calls (3 re-quantizes of
9975                // the same h + 3 _b2/_b4 launches). Bit-identical per (tensor,token,row).
9976                let (hq_o, hd_o);
9977                let (hq, hd): (&CudaSlice<i8>, &CudaSlice<f32>) = match h_q8 {
9978                    Some(p) => p,
9979                    None => {
9980                        (hq_o, hd_o) = e.quantize_q8_1(h, t, n_embd)?;
9981                        (&hq_o, &hd_o)
9982                    }
9983                };
9984                fused = e.matmul_q8_fused3_t(&fa.wq, &fa.wk, &fa.wv, hq, hd, t)?;
9985            }
9986            match (fused, h_q8) {
9987                (Some(triple), _) => triple,
9988                // shared pre-quantized activation (q8_1-fast guaranteed by the caller): the
9989                // decode-exact dispatch consumes (hq, hd) instead of re-quantizing 3x.
9990                (None, Some((hq, hd))) if qkv_fast => (
9991                    e.matmul_decode_exact_pre(&fa.wq, hq, hd, t)?,
9992                    e.matmul_decode_exact_pre(&fa.wk, hq, hd, t)?,
9993                    e.matmul_decode_exact_pre(&fa.wv, hq, hd, t)?,
9994                ),
9995                (None, _) => (
9996                    e.matmul_decode_exact(&fa.wq, h, t)?,
9997                    e.matmul_decode_exact(&fa.wk, h, t)?,
9998                    e.matmul_decode_exact(&fa.wv, h, t)?,
9999                ),
10000            }
10001        };
10002        // M3/Hy3 have no attention output gate — wq out is exactly q; skip the split.
10003        let gated = geometry.attention_gate == memra_gguf::config::AttentionGateKind::FusedQ;
10004        let (mut q, gate) = if gated {
10005            let mut q = vbuf(e, t * n_head * head_dim)?; // fully written by q_gate_split
10006            let mut gate = vbuf(e, t * n_head * head_dim)?; // fully written by q_gate_split
10007            e.q_gate_split(&qf, &mut q, &mut gate, head_dim, n_head, t)?;
10008            (q, Some(gate))
10009        } else {
10010            (qf, None)
10011        };
10012
10013        let mut qn = vbuf(e, t * n_head * head_dim)?; // fully written by rms_norm
10014        e.rms_norm(
10015            &q,
10016            fa.q_norm.float_data(),
10017            &mut qn,
10018            head_dim,
10019            n_head * t,
10020            eps,
10021        )?;
10022        q = qn;
10023        let mut kn = vbuf(e, t * n_head_kv * head_dim)?; // fully written by rms_norm
10024        e.rms_norm(
10025            &k,
10026            fa.k_norm.float_data(),
10027            &mut kn,
10028            head_dim,
10029            n_head_kv * t,
10030            eps,
10031        )?;
10032        k = kn;
10033        let rope_dims = geometry.n_rot as usize;
10034        e.rope_neox(
10035            &mut q,
10036            pos_d,
10037            head_dim,
10038            rope_dims,
10039            n_head,
10040            t,
10041            geometry.rope_base,
10042            1.0,
10043        )?;
10044        e.rope_neox(
10045            &mut k,
10046            pos_d,
10047            head_dim,
10048            rope_dims,
10049            n_head_kv,
10050            t,
10051            geometry.rope_base,
10052            1.0,
10053        )?;
10054
10055        // append T new K/V columns to the resident QUANTIZED cache. k/v are token-major [T, kv_dim]
10056        // f32; append-quantize each of the T token rows into the byte cache (q8_0 K / q5_1 V).
10057        let kvl = cache.kv[il].as_mut().unwrap();
10058        let (kv_dim_k, kv_dim_v, ktb, vtb) =
10059            (kvl.kv_dim_k, kvl.kv_dim_v, kvl.k_tok_bytes, kvl.v_tok_bytes);
10060        if let Some(ctr) = stream_ctr {
10061            // stream: ONE batched append at the device counter (rows kernel = the per-view warp
10062            // math on a (block, token) grid, documented byte-identical); host len is a stale
10063            // LOWER BOUND under pre-issue (drain reconciles it).
10064            e.append_kv_quantized_rows_dc(
10065                &k,
10066                &v,
10067                &mut kvl.k,
10068                &mut kvl.v,
10069                ctr,
10070                t,
10071                kv_dim_k,
10072                kv_dim_v,
10073                ktb,
10074                vtb,
10075                crate::Engine::kv_fp8_on(),
10076            )?;
10077        } else {
10078            for i in 0..t {
10079                let k_row = k.slice(i * kv_dim_k..(i + 1) * kv_dim_k);
10080                let v_row = v.slice(i * kv_dim_v..(i + 1) * kv_dim_v);
10081                e.append_kv_quantized_view(
10082                    &k_row,
10083                    &v_row,
10084                    &mut kvl.k,
10085                    &mut kvl.v,
10086                    kvl.len + i,
10087                    kv_dim_k,
10088                    kv_dim_v,
10089                    ktb,
10090                    vtb,
10091                    crate::Engine::kv_fp8_on(),
10092                )?;
10093            }
10094            kvl.len += t;
10095        }
10096
10097        // BIT-IDENTICAL VERIFY ATTENTION (spec-exactness fix): the FP accumulation order must be
10098        // byte-for-byte identical to the eager decode path. fa_prefill uses a different tile size
10099        // (BLOCK_Q=64, BK=32) and online-softmax structure than fa_decode's split-K + combine,
10100        // which changes FP summation order and can flip argmax at tight logit margins. Query row r
10101        // attends to keys [0..base_len+r+1) — each successive row sees one more key (the causal
10102        // property). This matches eager: decode appends k at len, then fa_decode sees t_kv = len+1
10103        // keys. The verify appends all T tokens first but bounds the key range per row.
10104        //
10105        // MULTI-ROW FUSED PATH (the long-ctx spec fix, 2026-07-03): when every row takes the vec
10106        // kernel (base_len+1 >= FA_VEC_MIN_TKV), ONE fa_decode_rows launch executes the exact
10107        // per-row program for all T rows (grid.z = row, per-row n_splits from the same
10108        // fa_split_keys formula) — replacing T x (2 launches + 2 dtod copies + 5 partial allocs)
10109        // and multiplying resident CTAs by T on a latency-bound kernel. Bit-identical per row by
10110        // construction; kernel-check pins rows-vs-loop byte identity, run-spec is the end gate.
10111        // Short ctx (any row below the vec crossover) and MEMRA_NO_FA_VEC/MEMRA_FA_ROWS_OFF keep the
10112        // per-row loop (whose fa_decode picks scalar/vec per row exactly like eager decode).
10113        let mut attn = vbuf(e, t * n_head * head_dim)?; // fully written by every FA arm below
10114        let base_len = kvl.len - t; // KV len BEFORE this round's T tokens were appended
10115        // T=1 INCLUDED (2026-07-05): p-min cuts the draft to 1 in ~75% of rounds on hard
10116        // (agentic) content — the old t>1 gate sent those rounds to the per-row loop (262us/row
10117        // + q-row copy + per-row allocs vs 93us/row through the fused kernel at grid.z=1, same
10118        // program). nsys accounting: 1088 of 1456 verify FA launches were T=1 escapees.
10119        // LEAN T=1 ARM (MEMRA_SPEC_LEAN, close35): at t==1, q IS one row and fa_decode on it is
10120        // the EXACT eager decode dispatch (vec_q_v2 + combine_f32; the rows pair measured +50us
10121        // at m=1). Byte-identical: kernel-check pins rows-vs-loop identity, and the per-row loop
10122        // at t=1 is fa_decode on the same q with zero-offset copies. Gates arbitrate.
10123        if let Some(ctr) = stream_ctr {
10124            // STREAM ARM: causal base from the device counter; host kvl.len is a stale lower
10125            // bound used only for the split-sizing upper bound (+64 slack covers M pre-issued
10126            // rounds at K<=8). Views span the bound; per-row limits derive in-kernel.
10127            let upper = kvl.len + t + 64;
10128            let k_view = e.view_u8(&kvl.k, (upper.min(cache.max_ctx)) * ktb);
10129            let v_view = e.view_u8(&kvl.v, (upper.min(cache.max_ctx)) * vtb);
10130            e.fa_decode_rows_dc(
10131                &q,
10132                &k_view,
10133                &v_view,
10134                &mut attn,
10135                head_dim,
10136                n_head,
10137                n_head_kv,
10138                ctr,
10139                upper.min(cache.max_ctx),
10140                t,
10141                scale,
10142                ktb,
10143                vtb,
10144                0,
10145                false,
10146            )?;
10147        } else if spec_lean() && t == 1 {
10148            let t_kv = base_len + 1;
10149            let k_view = e.view_u8(&kvl.k, t_kv * ktb);
10150            let v_view = e.view_u8(&kvl.v, t_kv * vtb);
10151            e.fa_decode_kvmod(
10152                &q,
10153                &k_view,
10154                &v_view,
10155                &mut attn,
10156                head_dim,
10157                n_head,
10158                n_head_kv,
10159                t_kv,
10160                scale,
10161                ktb,
10162                vtb,
10163                crate::Engine::kv_fp8_on(),
10164            )?;
10165        } else if e.fa_rows_eligible(base_len, head_dim) {
10166            let k_view = e.view_u8(&kvl.k, (base_len + t) * ktb);
10167            let v_view = e.view_u8(&kvl.v, (base_len + t) * vtb);
10168            e.fa_decode_rows(
10169                &q,
10170                &k_view,
10171                &v_view,
10172                &mut attn,
10173                head_dim,
10174                n_head,
10175                n_head_kv,
10176                base_len,
10177                t,
10178                scale,
10179                ktb,
10180                vtb,
10181                None,
10182                false,
10183                crate::Engine::kv_fp8_on(),
10184                None,
10185            )?;
10186        } else {
10187            for r in 0..t {
10188                let t_kv_r = base_len + r + 1; // this row sees keys [0..t_kv_r)
10189                let k_view_r = e.view_u8(&kvl.k, t_kv_r * ktb);
10190                let v_view_r = e.view_u8(&kvl.v, t_kv_r * vtb);
10191                // copy q row into an owned buffer (fa_decode takes &CudaSlice, not CudaView)
10192                let mut q_row = vbuf(e, n_head * head_dim)?; // fully written by copy_view_into
10193                let q_src = q.slice(r * n_head * head_dim..(r + 1) * n_head * head_dim);
10194                e.copy_view_into(&mut q_row, 0, &q_src, n_head * head_dim)?;
10195                let mut attn_row = vbuf(e, n_head * head_dim)?; // fully written by fa_decode
10196                e.fa_decode_kvmod(
10197                    &q_row,
10198                    &k_view_r,
10199                    &v_view_r,
10200                    &mut attn_row,
10201                    head_dim,
10202                    n_head,
10203                    n_head_kv,
10204                    t_kv_r,
10205                    scale,
10206                    ktb,
10207                    vtb,
10208                    crate::Engine::kv_fp8_on(),
10209                )?;
10210                e.copy_into(
10211                    &mut attn,
10212                    r * n_head * head_dim,
10213                    &attn_row,
10214                    n_head * head_dim,
10215                )?;
10216            }
10217        }
10218
10219        let attn_g = match &gate {
10220            Some(gate) => {
10221                let mut gsig = vbuf(e, t * n_head * head_dim)?; // fully written by sigmoid
10222                e.sigmoid(gate, &mut gsig, t * n_head * head_dim)?;
10223                let mut ag = vbuf(e, t * n_head * head_dim)?; // fully written by mul
10224                e.mul(&attn, &gsig, &mut ag, t * n_head * head_dim)?;
10225                ag
10226            }
10227            None => attn,
10228        };
10229        // DECODE-EXACT wo projection: at m>=5 (K=4+ with pending) the generic matmul would use dp4a
10230        // (128-thread, different FP sum order than MMVQ). Force MMVQ for bit-identity with decode.
10231        match self.full_attn_tp_o(e, fa, &attn_g, t)? {
10232            Some(output) => Ok(output),
10233            None => Ok(e.matmul_decode_exact(&fa.wo, &attn_g, t)?),
10234        }
10235    }
10236
10237    /// Context-linear bytes for a plain serving session's trunk cache.
10238    pub fn plain_session_kv_bytes_per_token(&self) -> usize {
10239        crate::cache::cache_bytes_per_token_for_plan(
10240            &self.cfg,
10241            &self.plan,
10242            0,
10243            self.plan.layers.len(),
10244        )
10245    }
10246
10247    /// `(logical bytes/token, ring-capped bytes/token, ring row cap)` for exact admission.
10248    pub fn plain_session_kv_shape(&self) -> (usize, usize, usize) {
10249        (
10250            self.plain_session_kv_bytes_per_token(),
10251            crate::cache::cache_ring_bytes_per_token_for_plan(
10252                &self.cfg,
10253                &self.plan,
10254                0,
10255                self.plan.layers.len(),
10256            ),
10257            crate::cache::cache_ring_row_cap_for_plan(&self.plan),
10258        )
10259    }
10260
10261    /// Context-linear bytes for a speculative serving session: trunk cache plus persistent MTP
10262    /// scratch. With no MTP head this equals the plain coefficient.
10263    pub fn spec_session_kv_bytes_per_token(&self) -> usize {
10264        let scratch = self
10265            .mtp
10266            .iter()
10267            .chain(self.mtp_extra.iter())
10268            .map(|mtp| {
10269                let (_, _, k, v) = mtp_scratch_layout(&self.cfg, mtp.geom.as_ref());
10270                k + v
10271            })
10272            .sum::<usize>();
10273        self.plain_session_kv_bytes_per_token()
10274            .saturating_add(scratch)
10275    }
10276
10277    /// Spec twin of [`HybridModel::plain_session_kv_shape`]; Step35's persistent MTP scratch is
10278    /// capped by the same SWA ring rows as the trunk.
10279    pub fn spec_session_kv_shape(&self) -> (usize, usize, usize) {
10280        let total = self.spec_session_kv_bytes_per_token();
10281        let (_, mut ring, rows) = self.plain_session_kv_shape();
10282        if rows > 0 {
10283            ring = ring.saturating_add(
10284                self.mtp
10285                    .iter()
10286                    .chain(self.mtp_extra.iter())
10287                    .map(|mtp| {
10288                        let (_, _, k, v) = mtp_scratch_layout(&self.cfg, mtp.geom.as_ref());
10289                        k + v
10290                    })
10291                    .sum::<usize>(),
10292            );
10293        }
10294        (total, ring, rows)
10295    }
10296
10297    /// Greedy MTP speculative decode (§B). Token-identical to `generate(prompt, max_new)` but uses
10298    /// the NextN head to draft K tokens then verifies them in one batched target forward.
10299    /// Returns (generated tokens, total_drafted, total_accepted) so the caller can report
10300    /// acceptance rate. `k` = draft length per round.
10301    ///
10302    /// GRAPH DRAFT (stage 2 of graph-grade spec): when the model is all-Dense and the MTP head is
10303    /// Dense (no MoE host readbacks), the fixed-shape T=1 MTP forward is CUDA-graph-captured ONCE
10304    /// and replayed per draft step — the ~40 eager launches per drafted token collapse into one
10305    /// graph dispatch; only the 4-byte token id (and 4-byte p-min confidence) round-trip per step.
10306    /// Event tracking is disabled for the whole call (generate_graph pattern) so every buffer the
10307    /// captured graph references is event-free; the spec loop is strictly single-stream.
10308    /// MEMRA_SPEC_NOGRAPH=1 forces the eager draft chain.
10309    /// SAMPLED mode (MEMRA_SPEC_TEMP>0) has its OWN capture (gumbel-perturbed in-graph argmax,
10310    /// device Philox event counter, persistent q retention) — graph-vs-eager sampled streams are
10311    /// bit-identical for the same (seed, prompt, K, temp); see the sampled-graph setup in
10312    /// generate_spec_inner2.
10313    /// Multi-turn session: trunk cache + MTP draft scratch persist across generate calls, so
10314    /// turn N+1 primes ONLY its new suffix (the 124k-conversation daily pattern — re-priming a
10315    /// 32k history costs ~54s; a suffix prime costs seconds). APPEND-ONLY by construction: the
10316    /// hybrid linear-attn states are in-place (no position index), so a session can extend but
10317    /// never rewind — `committed` is the exact token list whose state the caches hold (includes
10318    /// any overshoot tokens past max_new; the caller renders from `committed`, not its own echo).
10319    pub fn new_session(
10320        &self,
10321        e: &Engine,
10322        max_ctx: usize,
10323    ) -> Result<SpecSession, Box<dyn std::error::Error>> {
10324        Ok(SpecSession {
10325            // STAGE-OWNED KV (lane/pp2-spec 2026-08-06): `pp::new_cache`, not `Cache::new`. This
10326            // is the SERVING spec-session path, and with the ppN door open across two cards a
10327            // primary-homed cache makes every remote stage peer-read its OWN KV on every verify
10328            // round — the wrong-card class already fixed on the two batched serving paths
10329            // (worker.rs 2483 / 2837). With the door shut `new_cache` IS `Cache::new` (same
10330            // branch, same allocations), so single-device behavior is byte-unchanged.
10331            cache: crate::pp::new_cache_planned(e, &self.cfg, &self.plan, max_ctx)?,
10332            scratch: self.new_mtp_scratch(e, max_ctx)?,
10333            committed: Vec::new(),
10334            last_h: None,
10335            next_pred: None,
10336            sctr: 0,
10337            uctr: 0,
10338            draft_ctx: None,
10339            pending_tok: None,
10340            turn_ckpt: None,
10341            telem: SpecTelemetryCounters::default(),
10342            capture_at: None,
10343            boundary_captures: Vec::new(),
10344            ckpt_at: None,
10345            capture_disabled: false,
10346        })
10347    }
10348
10349    /// SPEC-ON-CACHE-HIT restore (lane/spec-on-cache-hit, 2026-08-18 — PORT-PLAN item 3,
10350    /// research/cache-spec-design-20260814, scoped to WHOLE-ENTRY restores only): build a
10351    /// SpecSession around a trunk cache the worker already restored from a prefix-cache
10352    /// entry, re-installing the entry's published draft plane as the MTP scratch rows
10353    /// `[0..prefix.len())` and the entry's boundary hidden as `last_h`, then feeding the
10354    /// prompt SUFFIX here — through EXACTLY the plain path's program selection — so the
10355    /// worker always receives a fully-warm continuation session (committed = whole
10356    /// prompt, `next_pred` + `last_h` set; caller sets `next_pred` from the entry's
10357    /// boundary logits on the empty-suffix shape).
10358    ///
10359    /// PROGRAM LAW (the splitiso two-programs class, learned AGAIN in this lane's own
10360    /// gate): the identity target for a converted hit is the PLAIN hit serving the same
10361    /// request, and plain feeds a carried suffix via eager `decode_step` below
10362    /// PRIME_MIN_T and via `prime_cache` at/above it (prefill_tick's arms). The generate
10363    /// path's tokenwise arm routes qwen35-class through the BATCHED T=1 program
10364    /// (`spec_target_step_h`) instead — ULP-different suffix rows, and the gate measured
10365    /// the near-tie flip at generated token ~8 (research/spec-cache-20260818, qwen r3).
10366    /// So the suffix is fed HERE, mirroring prefill_tick arm-for-arm, not handed to the
10367    /// burst prime.
10368    ///
10369    /// SEED RULE (both sampling regimes; lane/sampled-hit-spec 2026-08-19, sampled draw
10370    /// added by lane/sampled-spec-quality 2026-08-19). The boundary token is produced by
10371    /// EXACTLY the rule the cold burst entry applies to its own first token from the same
10372    /// logits row: `argmax` when greedy, and a `sample_boundary_token` draw at Philox
10373    /// counter 0 when sampled. Both shapes are covered — the entry's boundary logits on a
10374    /// full-cover (empty-suffix) hit, this feed's own boundary logits on a suffix hit.
10375    /// That is what keeps a restored session seed-identical to a cold one PER SEED: the
10376    /// cold session draws from the identical row at counter 0 and then runs its rounds from
10377    /// counter 1, so the restored session admits with `sctr = 1` after its own draw.
10378    /// The WORKER owns the one refusal this constructor cannot see — a constrained request.
10379    /// (The penalized-sampled refusal was LIFTED once the burst's penalty window learned to
10380    /// span the session: `committed` here is the WHOLE prompt, so the restored session's
10381    /// window is the cold session's window. It comes back if `MEMRA_SPEC_PEN_SESSION=0`.)
10382    ///
10383    /// NOT the rolled-back partial-restore hazard: the caller restores at exactly the
10384    /// entry's captured endpoint (`e.pos`) through the shipping whole-entry path;
10385    /// mid-entry (`at < e.pos`) trunk restores stay behind MEMRA_PREFIX_PARTIAL_RESTORE
10386    /// and are never routed here.
10387    ///
10388    /// Failure contract: `Err((Some(cache), why))` before any trunk mutation — the
10389    /// worker rebuilds the plain carrier and the hit serves plain, byte-unchanged.
10390    /// `Err((None, why))` after the suffix feed began — the carrier is part-fed and
10391    /// UNUSABLE; the worker serves the request cold-plain (correct, slower) and the
10392    /// entry stays published for the next request.
10393    #[allow(clippy::too_many_arguments)]
10394    #[allow(clippy::result_large_err)] // allow: the fat error type is the diagnostic contract here; boxing it would change the error surface
10395    pub fn spec_session_from_restored(
10396        &self,
10397        e: &Engine,
10398        mut cache: Cache,
10399        prefix: Vec<u32>,
10400        suffix: &[u32],
10401        draft_k: &CudaSlice<u8>,
10402        draft_v: &CudaSlice<u8>,
10403        draft_k_tok_bytes: usize,
10404        draft_v_tok_bytes: usize,
10405        draft_len: usize,
10406        last_h: &[f32],
10407        // The ENTRY's boundary logits row (the full-cover shape's seed source). May be empty
10408        // when a suffix follows — the feed's own logits are the boundary then.
10409        boundary_logits: &[f32],
10410        // The request's sampler, or None for greedy. Owned here so the seed rule lives in
10411        // ONE place instead of being half-applied by the worker.
10412        sampling: Option<SpecSampling>,
10413        require_anchor: bool,
10414        max_ctx: usize,
10415        // STABLE-BOUNDARY REPUBLICATION (lane/frspec-multiturn-cache, 2026-08-21): ABSOLUTE
10416        // prompt position to split the suffix feed at and capture the extended-entry
10417        // publication + this session's `turn_ckpt` — the worker's stable pre-generation
10418        // boundary (`plain_checkpoint_boundary`). None = legacy prompt-end republication.
10419        // WHY: the prompt-end capture below includes the template's live generation header
10420        // (`<|im_start|>assistant\n<think>\n`), which the next turn's re-render replaces, so
10421        // for a hybrid (whole-entry restores only) every extended entry's last ~2 tokens
10422        // diverged from every future prompt and the hit boundary FROZE at the first
10423        // lcp-split entry forever (measured: cached 6811 of 38228 by turn 8, B4).
10424        republish_at: Option<usize>,
10425    ) -> Result<SpecSession, (Option<Cache>, String)> {
10426        let pos = prefix.len();
10427        let fail = |cache: Cache, msg: String| -> Result<SpecSession, (Option<Cache>, String)> {
10428            Err((Some(cache), msg))
10429        };
10430        if let Err(error) = cache.ensure_usable("spec_session_from_restored") {
10431            drop(cache);
10432            return Err((None, error.to_string()));
10433        }
10434        if self.mtp.is_none() {
10435            return fail(cache, "no MTP head attached (nothing to draft with)".into());
10436        }
10437        if pos == 0 {
10438            return fail(cache, "empty committed prefix".into());
10439        }
10440        if cache.pos != pos {
10441            let msg = format!(
10442                "restored cache pos {} != restored prefix len {pos}",
10443                cache.pos
10444            );
10445            return fail(cache, msg);
10446        }
10447        if draft_len != pos {
10448            return fail(
10449                cache,
10450                format!("draft plane len {draft_len} != restored prefix len {pos}"),
10451            );
10452        }
10453        if pos + suffix.len() >= max_ctx {
10454            return fail(
10455                cache,
10456                format!(
10457                    "prompt {} + suffix would not leave generation room in ctx {max_ctx}",
10458                    pos + suffix.len(),
10459                ),
10460            );
10461        }
10462        let mut scratch = match MtpScratch::new(
10463            e,
10464            &self.cfg,
10465            &self.plan,
10466            max_ctx,
10467            self.mtp.as_ref().and_then(|m| m.geom.as_ref()),
10468        ) {
10469            Ok(s) => s,
10470            Err(err) => return fail(cache, format!("draft scratch alloc failed: {err}")),
10471        };
10472        if scratch.kv.ring.is_some() {
10473            return fail(
10474                cache,
10475                "ring-backed draft scratch (Step35 SWA) cannot take a flat prefix restore".into(),
10476            );
10477        }
10478        if scratch.kv.k_tok_bytes != draft_k_tok_bytes
10479            || scratch.kv.v_tok_bytes != draft_v_tok_bytes
10480        {
10481            return fail(
10482                cache,
10483                format!(
10484                    "draft plane layout {draft_k_tok_bytes}/{draft_v_tok_bytes} != scratch \
10485                     {}/{} bytes/token (stale entry across a format change)",
10486                    scratch.kv.k_tok_bytes, scratch.kv.v_tok_bytes,
10487                ),
10488            );
10489        }
10490        if pos > scratch.cap {
10491            return fail(
10492                cache,
10493                format!(
10494                    "draft plane rows {pos} exceed scratch capacity {}",
10495                    scratch.cap
10496                ),
10497            );
10498        }
10499        let kb = pos * draft_k_tok_bytes;
10500        let vb = pos * draft_v_tok_bytes;
10501        if draft_k.len() < kb || draft_v.len() < vb {
10502            return fail(
10503                cache,
10504                format!(
10505                    "truncated draft plane: K {} < {kb} or V {} < {vb} bytes",
10506                    draft_k.len(),
10507                    draft_v.len(),
10508                ),
10509            );
10510        }
10511        if kb > 0
10512            && let Err(err) = e.copy_u8_into(&mut scratch.kv.k, 0, draft_k, kb)
10513        {
10514            return fail(cache, format!("draft K restore copy failed: {err}"));
10515        }
10516        if vb > 0
10517            && let Err(err) = e.copy_u8_into(&mut scratch.kv.v, 0, draft_v, vb)
10518        {
10519            return fail(cache, format!("draft V restore copy failed: {err}"));
10520        }
10521        if let Err(err) = scratch.set_len(e, pos) {
10522            return fail(cache, format!("draft scratch len set failed: {err}"));
10523        }
10524        let mut last_h_dev = if last_h.len() == self.cfg.n_embd as usize {
10525            // anchor upload failure is acceptance-only when a suffix feed follows (fill
10526            // row-0 falls back to zeros) but FATAL for an empty-suffix continuation (the
10527            // burst entry asserts committed + last_h + next_pred) — the caller says which.
10528            e.htod(last_h).ok()
10529        } else {
10530            None
10531        };
10532        if require_anchor && last_h_dev.is_none() {
10533            return fail(
10534                cache,
10535                "empty-suffix continuation requires the entry's boundary hidden anchor".into(),
10536            );
10537        }
10538        let mut committed = prefix;
10539        // Set on BOTH shapes below (suffix-fed and full-cover) — never left None, which is
10540        // what the empty-suffix continuation assert in the burst entry requires.
10541        let next_pred;
10542        // Philox: (0,0) at admit exactly like a fresh session; a sampled boundary draw below
10543        // consumes counter 0 and leaves 1, which is the state a cold session reaches after
10544        // drawing its own first token from the same row.
10545        let mut sctr = 0u32;
10546        let sampled = sampling.is_some_and(|s| s.temp > 0.0) && spec_sampled_boundary_on();
10547        // Penalty window for the boundary draw: the last `penalty_last_n` tokens of the WHOLE
10548        // prompt, which is what the cold session's own burst sees (Item 2's window). Built
10549        // after the suffix joins `committed` below.
10550        let mut boundary_captures: Vec<SpecBoundaryCapture> = Vec::new();
10551        let mut restored_turn_ckpt: Option<SpecCheckpoint> = None;
10552        if !suffix.is_empty() {
10553            // ---- SUFFIX FEED, mirroring prefill_tick's program selection exactly ----
10554            // From here on the trunk cache mutates: failures return Err((None, _)) and
10555            // the worker serves the request cold-plain instead of reusing the carrier.
10556            let dirty =
10557                |msg: String| -> Result<SpecSession, (Option<Cache>, String)> { Err((None, msg)) };
10558            let n_embd = self.cfg.n_embd as usize;
10559            let t = suffix.len();
10560            let mut h_rows = match e.uninit(t * n_embd) {
10561                Ok(b) => b,
10562                Err(err) => return fail(cache, format!("suffix hidden buffer alloc: {err}")),
10563            };
10564            // STABLE-BOUNDARY split (see `republish_at`): feed stops at the boundary so the
10565            // in-place GDN conv/ssm state can be snapshotted there — the only moment it
10566            // exists (the cold prime-split law). suffix-relative; None = one-segment legacy.
10567            let b_rel = republish_at
10568                .and_then(|abs| abs.checked_sub(pos))
10569                .filter(|&r| r > 0 && r < t);
10570            let mut feed_logits = Vec::new();
10571            let tokenwise_env = std::env::var("MEMRA_PRIME_TOKENWISE").is_ok()
10572                || e.frozen_cpu_experts_prefer_tokenwise_prime();
10573            let mut fed = 0usize;
10574            for seg_end in [b_rel, Some(t)].into_iter().flatten() {
10575                if seg_end <= fed {
10576                    continue;
10577                }
10578                let seg = &suffix[fed..seg_end];
10579                let batched = seg.len() >= crate::hybrid_forward::PRIME_MIN_T && !tokenwise_env;
10580                if batched {
10581                    // prefill_tick's prime arm: request-level prime_cache call; tokens still
10582                    // queued after this segment ride `queued_after` so Step35 arm selection
10583                    // stays keyed to the request's end (tick-seg law).
10584                    match self.prime_cache(e, seg, &mut cache, t - seg_end) {
10585                        Ok((l, _h_seed, hiddens)) => {
10586                            if let Err(err) =
10587                                e.copy_into(&mut h_rows, fed * n_embd, &hiddens, seg.len() * n_embd)
10588                            {
10589                                return dirty(format!("suffix hidden copy: {err}"));
10590                            }
10591                            feed_logits = l;
10592                        }
10593                        Err(err) => return dirty(format!("suffix prime failed: {err}")),
10594                    }
10595                } else {
10596                    // prefill_tick's tokenwise arm: eager decode_step, one token at a time.
10597                    for (i, &tok) in seg.iter().enumerate() {
10598                        match self.decode_step_h(e, tok, &mut cache) {
10599                            Ok((l, h)) => {
10600                                if let Err(err) =
10601                                    e.copy_into(&mut h_rows, (fed + i) * n_embd, &h, n_embd)
10602                                {
10603                                    return dirty(format!("suffix hidden copy: {err}"));
10604                                }
10605                                feed_logits = l;
10606                            }
10607                            Err(err) => return dirty(format!("suffix decode_step failed: {err}")),
10608                        }
10609                    }
10610                }
10611                fed = seg_end;
10612                if Some(seg_end) == b_rel {
10613                    // The stable pre-generation boundary: capture the extended-entry
10614                    // publication AND this session's own turn checkpoint here instead of at
10615                    // prompt-end (both would otherwise carry the volatile live-header tail
10616                    // the next re-render replaces). Failure silent, turn_ckpt convention.
10617                    debug_assert_eq!(
10618                        cache.pos,
10619                        pos + seg_end,
10620                        "stable-boundary capture off the feed split"
10621                    );
10622                    if spec_restore_republish_on()
10623                        && let Ok(snap) = cache.snapshot(e)
10624                    {
10625                        boundary_captures.push(SpecBoundaryCapture {
10626                            snap,
10627                            pos: pos + seg_end,
10628                            logits: feed_logits.clone(),
10629                            last_h: capture_boundary_hidden(e, &h_rows, seg_end, n_embd),
10630                            latent_tails: Vec::new(),
10631                        });
10632                    }
10633                    let anchor: Result<CudaSlice<f32>, Box<dyn std::error::Error>> =
10634                        e.uninit(n_embd).and_then(|mut a| {
10635                            e.copy_view_into(
10636                                &mut a,
10637                                0,
10638                                &h_rows.slice((seg_end - 1) * n_embd..seg_end * n_embd),
10639                                n_embd,
10640                            )?;
10641                            Ok(a)
10642                        });
10643                    if let (Ok(snap), Ok(last_h)) = (cache.snapshot(e), anchor) {
10644                        restored_turn_ckpt = Some(SpecCheckpoint {
10645                            snap,
10646                            pos: pos + seg_end,
10647                            last_h,
10648                        });
10649                    }
10650                }
10651            }
10652            // Draft-scratch fill for the suffix rows, predecessor-paired: row `pos` reads
10653            // the entry's boundary anchor (zeros fallback — acceptance-only), row `pos+i`
10654            // reads h_rows[i-1]. Chunked like the generate path's fill (transients scale
10655            // with T). Fill failures are acceptance-only — truncate to the restored rows
10656            // and continue; the burst's own set_len keeps the invariant.
10657            let _mtp = self.mtp.as_ref().expect("mtp checked above"); // invariant check only; the fill below re-reads self.mtp
10658            let (embd_qt, embd_rb) = self.embd.qt_and_row_bytes(n_embd);
10659            let embd_gpu = if spec_host_embd() {
10660                None
10661            } else {
10662                Some(
10663                    self.embd_gpu
10664                        .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload")),
10665                )
10666            };
10667            let embd_dev = embd_gpu.map(|g| (g, embd_qt, embd_rb));
10668            let fill_chunk = 4096usize;
10669            let mut filled = true;
10670            let mut start = 0usize;
10671            'fill: while start < t {
10672                let end = (start + fill_chunk).min(t);
10673                let tc = end - start;
10674                let Ok(mut phs) = e.zeros(tc * n_embd) else {
10675                    filled = false;
10676                    break 'fill;
10677                };
10678                let (src_lo, dst_off, n_copy) = if start == 0 {
10679                    (0, n_embd, (tc - 1) * n_embd)
10680                } else {
10681                    ((start - 1) * n_embd, 0, tc * n_embd)
10682                };
10683                if start == 0
10684                    && let Some(lh) = last_h_dev.as_ref()
10685                    && e.copy_into(&mut phs, 0, lh, n_embd).is_err()
10686                {
10687                    filled = false;
10688                    break 'fill;
10689                }
10690                if n_copy > 0
10691                    && e.copy_view_into(
10692                        &mut phs,
10693                        dst_off,
10694                        &h_rows.slice(src_lo..src_lo + n_copy),
10695                        n_copy,
10696                    )
10697                    .is_err()
10698                {
10699                    filled = false;
10700                    break 'fill;
10701                }
10702                if self
10703                    .mtp_kv_fill_all(
10704                        e,
10705                        &suffix[start..end],
10706                        &phs,
10707                        pos + start,
10708                        &mut scratch,
10709                        embd_dev,
10710                    )
10711                    .is_err()
10712                {
10713                    filled = false;
10714                    break 'fill;
10715                }
10716                start = end;
10717            }
10718            if !filled {
10719                // acceptance-only: drafts over missing suffix rows are cheap and wrong,
10720                // so keep only the restored rows resident and let verify arbitrate.
10721                if let Err(err) = scratch.set_len(e, pos) {
10722                    return dirty(format!("scratch truncation after failed fill: {err}"));
10723                }
10724            }
10725            // EXTENDED-ENTRY PUBLICATION (lane/sampled-spec-quality, Item 3 — the fix for
10726            // "a restored spec session never publishes an extended entry", SAMPLED-HIT.md
10727            // finding (d)). Pre-lane, publication was armed only for COLD sessions
10728            // (`spec_resumed == 0` in the worker) and both engine capture sites require a
10729            // non-continuation burst — but a converted hit's first burst IS a continuation,
10730            // so a growing conversation learned exactly ONE boundary and turn 3 could never
10731            // hit a longer prefix than turn 2 did.
10732            //
10733            // WHERE, and why it is safe here: `cache.pos == prefix + suffix` at this exact
10734            // line — the trunk is primed over the whole prompt, nothing is generated, and the
10735            // draft plane rows [0..prompt) are filled just above. That is a complete
10736            // whole-entry boundary (`pos == fed_len`), the same shape the cold seed capture
10737            // publishes; the worker's existing publication sweep picks it up because it is
10738            // keyed on non-empty `boundary_captures` and is sampler- and resume-independent.
10739            // NOT the partial-restore hazard: the boundary is this session's own prompt END,
10740            // never mid-entry, so `entry_pos != fed_len` still refuses on the way back in.
10741            // Failure is SILENT by design (the turn_ckpt / boundary-capture convention):
10742            // publication is an optimization, never a correctness dependency.
10743            //
10744            // SUPERSEDED WHEN `republish_at` FIRED (lane/frspec-multiturn-cache): a prompt-end
10745            // entry's tail is the live generation header the next re-render replaces, so on a
10746            // hybrid (whole-entry restores) it can never serve the conversation's next turn —
10747            // the stable-boundary capture above IS this publication, minus the poisoned tail.
10748            if spec_restore_republish_on() && boundary_captures.is_empty() {
10749                debug_assert_eq!(
10750                    cache.pos,
10751                    pos + t,
10752                    "extended-entry capture must sit at the restored session's prompt end",
10753                );
10754                if let Ok(snap) = cache.snapshot(e) {
10755                    boundary_captures.push(SpecBoundaryCapture {
10756                        snap,
10757                        pos: pos + t,
10758                        logits: feed_logits.clone(),
10759                        last_h: capture_boundary_hidden(e, &h_rows, t, n_embd),
10760                        latent_tails: Vec::new(),
10761                    });
10762                }
10763            }
10764            // continuation seed: the feed's boundary logits ARE the plain path's boundary
10765            // logits (same program), so greedy's argmax here is plain's first emitted token,
10766            // and the sampled draw is the cold sampled session's own first token.
10767            next_pred = Some(if sampled {
10768                let sp = sampling.expect("sampled implies a sampler");
10769                // `committed` is still the restored prefix here; the suffix joins it below —
10770                // so this is the last-N window over the WHOLE prompt, exactly the cold
10771                // session's own window at its first token.
10772                let hist = pen_window_seed(&committed, suffix, sp.penalty_last_n);
10773                match sample_boundary_token(
10774                    e,
10775                    &feed_logits,
10776                    &sp,
10777                    &hist,
10778                    &mut sctr,
10779                    "restore-suffix-feed",
10780                ) {
10781                    Ok(t) => t,
10782                    // the trunk is already fed: hand nothing back, the worker serves the
10783                    // request cold-plain. Never fall back to an argmax — that would put a
10784                    // greedy token in a sampled stream to save a slow path.
10785                    Err(err) => {
10786                        return dirty(format!("boundary token draw failed: {err}"));
10787                    }
10788                }
10789            } else {
10790                argmax(&feed_logits) as u32
10791            });
10792            let mut lh = match e.uninit(n_embd) {
10793                Ok(b) => b,
10794                Err(err) => return dirty(format!("boundary hidden alloc: {err}")),
10795            };
10796            if let Err(err) = e.copy_view_into(
10797                &mut lh,
10798                0,
10799                &h_rows.slice((t - 1) * n_embd..t * n_embd),
10800                n_embd,
10801            ) {
10802                return dirty(format!("boundary hidden copy: {err}"));
10803            }
10804            last_h_dev = Some(lh);
10805            committed.extend_from_slice(suffix);
10806        } else {
10807            // FULL-COVER shape (empty suffix — the identical-repeat / agent-loop shape): the
10808            // ENTRY's boundary logits are the boundary row, and this is the token the cold
10809            // session emits from that same row. Owned here rather than in the worker so the
10810            // sampled draw cannot be half-applied on one shape (the worker used to argmax it).
10811            if boundary_logits.is_empty() {
10812                return fail(
10813                    cache,
10814                    "full-cover restore without the entry's boundary logits".into(),
10815                );
10816            }
10817            next_pred = Some(if sampled {
10818                let sp = sampling.expect("sampled implies a sampler");
10819                let hist = pen_window_seed(&committed, &[], sp.penalty_last_n);
10820                match sample_boundary_token(
10821                    e,
10822                    boundary_logits,
10823                    &sp,
10824                    &hist,
10825                    &mut sctr,
10826                    "restore-full-cover",
10827                ) {
10828                    Ok(t) => t,
10829                    // nothing has been mutated on this shape — hand the carrier back and let
10830                    // the hit serve PLAIN (the banked pre-lane path).
10831                    Err(err) => {
10832                        return fail(cache, format!("boundary token draw failed: {err}"));
10833                    }
10834                }
10835            } else {
10836                argmax(boundary_logits) as u32
10837            });
10838        }
10839        Ok(SpecSession {
10840            cache,
10841            scratch,
10842            committed,
10843            last_h: last_h_dev,
10844            next_pred,
10845            sctr,
10846            uctr: 0,
10847            draft_ctx: None,
10848            pending_tok: None,
10849            // Stable-boundary capture from the split feed above (None on the legacy shape):
10850            // a restored session previously parked WITHOUT a checkpoint, so the next turn's
10851            // affinity probe declined ("no turn checkpoint retained") and the conversation
10852            // fell back to the frozen prefix entry forever.
10853            turn_ckpt: restored_turn_ckpt,
10854            telem: SpecTelemetryCounters::default(),
10855            capture_at: None,
10856            boundary_captures,
10857            ckpt_at: None,
10858            capture_disabled: false,
10859        })
10860    }
10861
10862    /// Forced-gate exact state comparison. This intentionally reads the real live prefixes from
10863    /// their owning PP devices: matching emitted ids alone would miss a stale `len_d`, recurrent
10864    /// snapshot, or draft-KV row that only corrupts the following round.
10865    pub fn optipipe_compare_session_state(
10866        &self,
10867        e: &Engine,
10868        reference: &SpecSession,
10869        candidate: &SpecSession,
10870    ) -> Result<OptiForkStateIdentity, Box<dyn std::error::Error>> {
10871        fn fail(what: &str) -> Box<dyn std::error::Error> {
10872            format!("optipipe state mismatch: {what}").into()
10873        }
10874        fn same_f32(a: &[f32], b: &[f32]) -> bool {
10875            a.len() == b.len() && a.iter().zip(b).all(|(x, y)| x.to_bits() == y.to_bits())
10876        }
10877        fn compare_layers(
10878            es: &Engine,
10879            range: std::ops::Range<usize>,
10880            reference: &SpecSession,
10881            candidate: &SpecSession,
10882            report: &mut OptiForkStateIdentity,
10883        ) -> Result<(), Box<dyn std::error::Error>> {
10884            for il in range {
10885                match (&reference.cache.kv[il], &candidate.cache.kv[il]) {
10886                    (Some(a), Some(b)) => {
10887                        if a.len != b.len {
10888                            return Err(fail(&format!(
10889                                "layer {il} host KV len {} != {}",
10890                                a.len, b.len
10891                            )));
10892                        }
10893                        let ad = es.dtoh_i32(&a.len_d)?;
10894                        let bd = es.dtoh_i32(&b.len_d)?;
10895                        if ad != bd || ad.first().copied() != Some(a.len as i32) {
10896                            return Err(fail(&format!(
10897                                "layer {il} device KV len {ad:?} != {bd:?} (host={})",
10898                                a.len,
10899                            )));
10900                        }
10901                        let kb = a.len * a.k_tok_bytes;
10902                        let vb = a.len * a.v_tok_bytes;
10903                        if kb > 0 {
10904                            let ak = es.dtoh_u8_view(&a.k.slice(0..kb))?;
10905                            let bk = es.dtoh_u8_view(&b.k.slice(0..kb))?;
10906                            if ak != bk {
10907                                let at = ak.iter().zip(&bk).position(|(x, y)| x != y).unwrap();
10908                                return Err(fail(&format!(
10909                                    "layer {il} K bytes at byte {at} row {} offset {}: {} != {}",
10910                                    at / a.k_tok_bytes,
10911                                    at % a.k_tok_bytes,
10912                                    ak[at],
10913                                    bk[at],
10914                                )));
10915                            }
10916                        }
10917                        if vb > 0 {
10918                            let av = es.dtoh_u8_view(&a.v.slice(0..vb))?;
10919                            let bv = es.dtoh_u8_view(&b.v.slice(0..vb))?;
10920                            if av != bv {
10921                                let at = av.iter().zip(&bv).position(|(x, y)| x != y).unwrap();
10922                                return Err(fail(&format!(
10923                                    "layer {il} V bytes at byte {at} row {} offset {}: {} != {}",
10924                                    at / a.v_tok_bytes,
10925                                    at % a.v_tok_bytes,
10926                                    av[at],
10927                                    bv[at],
10928                                )));
10929                            }
10930                        }
10931                        report.trunk_kv_bytes += kb + vb;
10932                    }
10933                    (None, None) => {}
10934                    _ => return Err(fail(&format!("layer {il} KV presence"))),
10935                }
10936                match (&reference.cache.recur[il], &candidate.cache.recur[il]) {
10937                    (Some(a), Some(b)) => {
10938                        let ac = es.dtoh(&a.conv_state)?;
10939                        let bc = es.dtoh(&b.conv_state)?;
10940                        if !same_f32(&ac, &bc) {
10941                            return Err(fail(&format!("layer {il} conv state")));
10942                        }
10943                        let as_ = es.dtoh(&a.ssm_state)?;
10944                        let bs = es.dtoh(&b.ssm_state)?;
10945                        if !same_f32(&as_, &bs) {
10946                            return Err(fail(&format!("layer {il} SSM state")));
10947                        }
10948                        report.recurrent_bytes += (ac.len() + as_.len()) * 4;
10949                    }
10950                    (None, None) => {}
10951                    _ => return Err(fail(&format!("layer {il} recurrent presence"))),
10952                }
10953            }
10954            Ok(())
10955        }
10956
10957        if reference.committed != candidate.committed {
10958            return Err(fail("committed token ids"));
10959        }
10960        if reference.cache.pos != candidate.cache.pos
10961            || reference.cache.max_ctx != candidate.cache.max_ctx
10962        {
10963            return Err(fail("cache pos/capacity"));
10964        }
10965        if reference.pending_tok != candidate.pending_tok
10966            || reference.next_pred != candidate.next_pred
10967            || reference.sctr != candidate.sctr
10968            || reference.uctr != candidate.uctr
10969        {
10970            return Err(fail("pending/prediction/counter tail"));
10971        }
10972
10973        let mut report = OptiForkStateIdentity::default();
10974        if let Some(fence) = crate::pp::pp_cuts(self.layers.len()) {
10975            let rt = crate::pp::PpNRt::get(e)?;
10976            for stage in 0..rt.n_stages() {
10977                let _scope = rt.enter(stage);
10978                compare_layers(
10979                    rt.engine(stage, e),
10980                    fence[stage]..fence[stage + 1],
10981                    reference,
10982                    candidate,
10983                    &mut report,
10984                )?;
10985            }
10986        } else {
10987            compare_layers(e, 0..self.layers.len(), reference, candidate, &mut report)?;
10988        }
10989
10990        if reference.scratch.plane_count() != candidate.scratch.plane_count() {
10991            return Err(fail("draft scratch plane count"));
10992        }
10993        for index in 0..reference.scratch.plane_count() {
10994            let (a, _) = reference.scratch.plane(index);
10995            let (b, _) = candidate.scratch.plane(index);
10996            if a.len != b.len
10997                || a.kv_dim_k != b.kv_dim_k
10998                || a.kv_dim_v != b.kv_dim_v
10999                || a.k_tok_bytes != b.k_tok_bytes
11000                || a.v_tok_bytes != b.v_tok_bytes
11001                || e.dtoh_i32(&a.len_d)? != e.dtoh_i32(&b.len_d)?
11002            {
11003                return Err(fail(&format!("draft scratch plane {index} length/layout")));
11004            }
11005            let kb = a.len * a.k_tok_bytes;
11006            let vb = a.len * a.v_tok_bytes;
11007            if kb > 0 && e.dtoh_u8_view(&a.k.slice(0..kb))? != e.dtoh_u8_view(&b.k.slice(0..kb))? {
11008                return Err(fail(&format!("draft scratch plane {index} K bytes")));
11009            }
11010            if vb > 0 && e.dtoh_u8_view(&a.v.slice(0..vb))? != e.dtoh_u8_view(&b.v.slice(0..vb))? {
11011                return Err(fail(&format!("draft scratch plane {index} V bytes")));
11012            }
11013            report.scratch_kv_bytes += kb + vb;
11014        }
11015
11016        match (&reference.last_h, &candidate.last_h) {
11017            (Some(a), Some(b)) => {
11018                let ah = e.dtoh(a)?;
11019                let bh = e.dtoh(b)?;
11020                if !same_f32(&ah, &bh) {
11021                    return Err(fail("last hidden/seed bytes"));
11022                }
11023                report.hidden_bytes = ah.len() * 4;
11024            }
11025            (None, None) => {}
11026            _ => return Err(fail("last hidden/seed presence")),
11027        }
11028        Ok(report)
11029    }
11030
11031    /// SESSION-AFFINITY REWIND (lane/session-affinity, 2026-08-05): roll `sess` back to its
11032    /// retained prompt-end checkpoint, so a request whose prompt matches
11033    /// `committed[..rewind_pos()]` exactly can resume there and prime only its own delta.
11034    ///
11035    /// EXACTNESS. After this returns, the session is byte-for-byte the state it was in AT that
11036    /// boundary: full-attn KV truncated to it (append-only, position-addressed), GDN conv/ssm
11037    /// restored from the device copy taken there, draft scratch length reset, `committed`
11038    /// truncated, `last_h` = the boundary's predecessor anchor. That is precisely the state a
11039    /// fresh prime of `committed[..pos]` would have produced, so the following suffix prime and
11040    /// every burst after it are identical to a cold run of the same token stream — the
11041    /// committed-tokens-authoritative contract.
11042    ///
11043    /// `next_pred` and `pending_tok` are CLEARED: both describe generation past the boundary,
11044    /// which the rewind discards. The caller therefore must supply a non-empty suffix (a
11045    /// rewound session cannot serve an empty-suffix continuation burst — there is nothing to
11046    /// continue). The persistent draft graph survives: it bakes only session-stable pointers
11047    /// (the scratch KV, the resident embedding), none of which the rewind moves.
11048    ///
11049    /// The checkpoint is CONSUMED (`turn_ckpt` taken): its snapshot buffers are freed here, and
11050    /// this turn's own prime installs a fresh one at the new prompt end. Returns the position
11051    /// rewound to, or `None` when the session holds no checkpoint (caller: full re-prime).
11052    pub fn spec_rewind_to_checkpoint(
11053        &self,
11054        e: &Engine,
11055        sess: &mut SpecSession,
11056    ) -> Result<Option<usize>, Box<dyn std::error::Error>> {
11057        if sess.turn_ckpt.as_ref().is_some_and(|ckpt| {
11058            !sess.cache.can_rollback(&ckpt.snap, 0) || !sess.scratch.can_rewind_to(ckpt.pos)
11059        }) {
11060            return Err(
11061                "SWA ring rewind checkpoint has been lapped; full re-prime required".into(),
11062            );
11063        }
11064        let Some(ckpt) = sess.turn_ckpt.take() else {
11065            return Ok(None);
11066        };
11067        assert!(
11068            ckpt.pos <= sess.committed.len(),
11069            "checkpoint past committed ({} > {})",
11070            ckpt.pos,
11071            sess.committed.len()
11072        );
11073        // Restore through each layer's owning engine. A single primary-engine rollback is not
11074        // sufficient when the serving cache is stage-owned under cross-device PP.
11075        crate::pp::restore_cache_checkpoint(e, self, None, &mut sess.cache, &ckpt.snap)?;
11076        debug_assert_eq!(
11077            sess.cache.pos, ckpt.pos,
11078            "rollback landed off the checkpoint"
11079        );
11080        sess.scratch.set_len(e, ckpt.pos)?;
11081        sess.committed.truncate(ckpt.pos);
11082        sess.last_h = Some(ckpt.last_h);
11083        sess.next_pred = None;
11084        sess.pending_tok = None;
11085        Ok(Some(ckpt.pos))
11086    }
11087
11088    /// Grow a parked speculative session to `target_cap` and rewind it to its retained turn
11089    /// checkpoint without re-priming the checkpoint prefix.
11090    ///
11091    /// The trunk cache is restored exactly like a plain grown cache: append-only full-attention
11092    /// KV rows come from the parked cache, while recurrent state comes from the checkpoint's
11093    /// owned snapshot. The MTP scratch is also context-linear and its rows below the checkpoint
11094    /// remain authoritative, so they are copied into a fresh larger scratch before its length is
11095    /// truncated. Pointer-baking draft graphs are dropped and recaptured on the next burst.
11096    ///
11097    /// All fallible work completes before `sess` is mutated. A failed allocation or copy leaves
11098    /// the parked session intact, allowing the caller one reclaim-and-retry attempt.
11099    pub fn spec_grow_and_rewind_to_checkpoint(
11100        &self,
11101        e: &Engine,
11102        sess: &mut SpecSession,
11103        target_cap: usize,
11104    ) -> Result<Option<usize>, Box<dyn std::error::Error>> {
11105        if target_cap <= sess.cache.max_ctx {
11106            return self.spec_rewind_to_checkpoint(e, sess);
11107        }
11108        let Some(ckpt) = sess.turn_ckpt.as_ref() else {
11109            return Ok(None);
11110        };
11111        if ckpt.pos == 0 || ckpt.pos > sess.committed.len() {
11112            return Err(format!(
11113                "checkpoint pos {} outside committed length {}",
11114                ckpt.pos,
11115                sess.committed.len(),
11116            )
11117            .into());
11118        }
11119        if ckpt.pos > target_cap {
11120            return Err(format!(
11121                "checkpoint pos {} exceeds grown capacity {target_cap}",
11122                ckpt.pos,
11123            )
11124            .into());
11125        }
11126
11127        let mut grown_cache = crate::pp::new_cache_planned(e, &self.cfg, &self.plan, target_cap)?;
11128        let mut grown_scratch = self.new_mtp_scratch(e, target_cap)?;
11129        crate::pp::restore_cache_checkpoint(
11130            e,
11131            self,
11132            Some(&sess.cache),
11133            &mut grown_cache,
11134            &ckpt.snap,
11135        )?;
11136
11137        if sess.scratch.plane_count() != grown_scratch.plane_count() {
11138            return Err("checkpoint draft plane count mismatch".into());
11139        }
11140        for index in 0..sess.scratch.plane_count() {
11141            let (src, _) = sess.scratch.plane(index);
11142            let (dst, _) = grown_scratch.plane_mut(index);
11143            if ckpt.pos > src.len
11144                || src.kv_dim_k != dst.kv_dim_k
11145                || src.kv_dim_v != dst.kv_dim_v
11146                || src.k_tok_bytes != dst.k_tok_bytes
11147                || src.v_tok_bytes != dst.v_tok_bytes
11148            {
11149                return Err(format!(
11150                    "checkpoint draft plane {index} layout mismatch (pos {}, source len {})",
11151                    ckpt.pos, src.len,
11152                )
11153                .into());
11154            }
11155            match (&src.ring, dst.ring.as_ref()) {
11156                (Some(sring), Some(_)) => {
11157                    // Ring-backed draft plane (step35): `ckpt.pos` is absolute and exceeds the
11158                    // physical rows once lapped — same class as the trunk-KV restore panic
11159                    // (2026-08-29 warm-turn-at-40k). Copy the aligned live window, rebase.
11160                    let (new_base, phys) = sring.restore_plan(ckpt.pos).map_err(|err| {
11161                        format!("checkpoint draft plane {index} SWA restore refused: {err}")
11162                    })?;
11163                    let rows = phys.len();
11164                    let kb = rows * src.k_tok_bytes;
11165                    let vb = rows * src.v_tok_bytes;
11166                    if kb > 0 {
11167                        e.copy_u8_range_into(
11168                            &mut dst.k,
11169                            0,
11170                            &src.k,
11171                            phys.start * src.k_tok_bytes,
11172                            kb,
11173                        )?;
11174                    }
11175                    if vb > 0 {
11176                        e.copy_u8_range_into(
11177                            &mut dst.v,
11178                            0,
11179                            &src.v,
11180                            phys.start * src.v_tok_bytes,
11181                            vb,
11182                        )?;
11183                    }
11184                    dst.ring
11185                        .as_mut()
11186                        .expect("ring presence checked above")
11187                        .apply_rebase(new_base);
11188                    if let Some(base_d) = dst.base_d.as_mut() {
11189                        e.set_i32_one(base_d, new_base as i32)?;
11190                    }
11191                }
11192                (None, None) => {
11193                    let kb = ckpt.pos * src.k_tok_bytes;
11194                    let vb = ckpt.pos * src.v_tok_bytes;
11195                    if kb > 0 {
11196                        e.copy_u8_into(&mut dst.k, 0, &src.k, kb)?;
11197                    }
11198                    if vb > 0 {
11199                        e.copy_u8_into(&mut dst.v, 0, &src.v, vb)?;
11200                    }
11201                }
11202                _ => {
11203                    return Err(format!("checkpoint draft plane {index} ring/flat mismatch").into());
11204                }
11205            }
11206        }
11207        grown_scratch.set_len(e, ckpt.pos)?;
11208        // The old scratch is dropped immediately after publication below. Bound its D2D reads
11209        // first; growth happens once per rewritten turn, outside the decode hot loop.
11210        e.stream().synchronize()?;
11211
11212        let ckpt = sess
11213            .turn_ckpt
11214            .take()
11215            .expect("checkpoint remained present through transactional grow");
11216        let pos = ckpt.pos;
11217        sess.cache = grown_cache;
11218        sess.scratch = grown_scratch;
11219        sess.committed.truncate(pos);
11220        sess.last_h = Some(ckpt.last_h);
11221        sess.next_pred = None;
11222        sess.pending_tok = None;
11223        sess.draft_ctx = None;
11224        debug_assert_eq!(sess.cache.pos, pos, "grown rewind landed off checkpoint");
11225        debug_assert!(
11226            (0..sess.scratch.plane_count()).all(|index| sess.scratch.plane(index).0.len == pos),
11227            "grown draft rewind landed off checkpoint"
11228        );
11229        Ok(Some(pos))
11230    }
11231
11232    /// Commit a carried pending bonus (see SpecSession::pending_tok): one T=1 trunk pass
11233    /// (its logits' argmax becomes next_pred) + the draft-KV fill at the carried anchor —
11234    /// byte-identical to the pre-carry session tail. Required before a non-empty-suffix
11235    /// prime, a sampled turn, or parking a session for pool reuse. No-op without a pending.
11236    /// `sampling` is the sampler of the request that will CONSUME the resulting `next_pred`
11237    /// (lane/sampled-spec-quality): this is a boundary site like any other, so a sampled
11238    /// consumer must get a DRAWN token, not an argmax. Pass `None` from the park/demote
11239    /// callers — a pending only ever exists on the GREEDY tail, and the consumer of a
11240    /// park-time flush is a future request whose sampler is not knowable here (residual
11241    /// named at the pool-resume probe in worker.rs and in SAMPLED-QUALITY.md).
11242    pub fn spec_flush_pending(
11243        &self,
11244        e: &Engine,
11245        sess: &mut SpecSession,
11246        sampling: Option<SpecSampling>,
11247    ) -> Result<(), Box<dyn std::error::Error>> {
11248        sess.cache.ensure_usable("spec_flush_pending")?;
11249        let Some(b) = sess.pending_tok.take() else {
11250            return Ok(());
11251        };
11252        if self.mtp.is_none() {
11253            return Err("pending carry requires an MTP head".into());
11254        }
11255        let n_embd = self.cfg.n_embd as usize;
11256        let (embd_qt, embd_rb) = self.embd.qt_and_row_bytes(n_embd);
11257        let embd_gpu = if spec_host_embd() {
11258            None
11259        } else {
11260            Some(
11261                self.embd_gpu
11262                    .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload")),
11263            )
11264        };
11265        let embd_dev = embd_gpu.map(|g| (g, embd_qt, embd_rb));
11266        let pos_b = sess.cache.pos;
11267        sess.scratch.set_len(e, pos_b)?;
11268        let (lg_b, hb) = self.spec_target_step_h(e, b, &mut sess.cache)?;
11269        sess.next_pred = Some(match sampling {
11270            Some(sp) if sp.temp > 0.0 && spec_sampled_boundary_on() => {
11271                // window includes `b` itself: it is committed by this pass, and the pre-lane
11272                // code never counted a boundary token in the penalty history at all.
11273                let hist = pen_window_seed(&sess.committed, &[b], sp.penalty_last_n);
11274                sample_boundary_token(e, &lg_b, &sp, &hist, &mut sess.sctr, "flush-pending")?
11275            }
11276            _ => argmax(&lg_b) as u32,
11277        });
11278        let anchor = sess
11279            .last_h
11280            .as_ref()
11281            .expect("pending carry requires last_h (the predecessor-row anchor)");
11282        self.mtp_kv_fill_all(e, &[b], anchor, pos_b, &mut sess.scratch, embd_dev)?;
11283        sess.last_h = Some(hb);
11284        sess.committed.push(b);
11285        Ok(())
11286    }
11287
11288    /// Solo target feed used only at speculative round boundaries. Step35 serving made its
11289    /// staged batched B=1 graph authoritative, so a speculative session must enter and leave
11290    /// rounds through that same graph. Other model families keep their eager T=1 contract.
11291    fn spec_target_step_h(
11292        &self,
11293        e: &Engine,
11294        token: u32,
11295        cache: &mut Cache,
11296    ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
11297        cache.ensure_usable("spec_target_step_h")?;
11298        if !self.sliding_gated_moe_batch_program() && !self.batched_serving_numeric_class() {
11299            return self.decode_step_h(e, token, cache);
11300        }
11301        let pos0 = cache.pos;
11302        let (logits, hidden) = self.decode_step_t_core(e, &[token], pos0, cache, None, None)?;
11303        Ok((e.dtoh(&logits)?, hidden))
11304    }
11305
11306    /// The archs whose LIVE B=1 serving runs the generic BATCHED numeric class (decode_step_batch
11307    /// walk + batched head), so their spec verify must run the SAME class. MoE learned this
11308    /// 2026-08-14 AM (4b777ccc5); the dense hybrid reproduced the identical near-tie flip class
11309    /// the same day on Qwen3.8-27B — eager-class verify logits drift from batched-class serving
11310    /// logits ("1 ULP at layer 2 → 2.3e-1 logit maxdiff at the head"), and the GDN recurrence
11311    /// carries the drift until a near-tie flips deep in generation. One predicate so the five
11312    /// dispatch sites cannot drift apart again.
11313    /// Draft-graph head admissibility (lane/draftcost-moe, 2026-08-20): the capture body
11314    /// (`mtp_head_forward_cap`) supports Dense heads and SOFTMAX device-routed resident-MoE
11315    /// heads. Residency alone is insufficient: Hy3/M3/Step sigmoid routing returns selected
11316    /// experts through a host synchronization, which is capture-illegal. Those heads use the
11317    /// exact eager draft chain until a device-only sigmoid expert program lands. Trunk FFN class
11318    /// is irrelevant — the graph body is the HEAD forward only. One predicate for all three
11319    /// eligibility sites so they cannot drift (the serving numeric-class lesson).
11320    fn mtp_graph_capturable(&self) -> bool {
11321        let sigmoid_router = self.cfg.sigmoid_router().is_some();
11322        for head in self.mtp.iter().chain(self.mtp_extra.iter()) {
11323            let reason = match &head.ffn {
11324                crate::hybrid::Ffn::Dense { .. } => None,
11325                crate::hybrid::Ffn::Moe(mo) if mo.dev_exps.is_none() => {
11326                    Some("non-resident MoE MTP head")
11327                }
11328                crate::hybrid::Ffn::Moe(_) if sigmoid_router => {
11329                    Some("sigmoid-router MoE MTP head requires host-visible routing")
11330                }
11331                crate::hybrid::Ffn::Moe(_) => None,
11332            };
11333            if let Some(reason) = reason {
11334                static NOTICE: std::sync::Once = std::sync::Once::new();
11335                NOTICE.call_once(|| {
11336                    eprintln!(
11337                        "[spec] draft graph unavailable: {reason}; eager draft chain engaged"
11338                    );
11339                });
11340                return false;
11341            }
11342        }
11343        self.mtp.is_some()
11344    }
11345
11346    fn batched_serving_numeric_class(&self) -> bool {
11347        self.plan
11348            .trunk_operations()
11349            .contains(&memra_gguf::model_plan::OperationKind::GatedDeltaNet)
11350    }
11351
11352    /// The family the MTP verify-graph default was measured on: GatedDeltaNet state layers
11353    /// (a `recur` mixer) together with a routed-MoE FFN — Ornith-1.5-35B-A3B and its kin. The
11354    /// server-side twin of this test is `model_forces_spec_replay` (GatedDeltaNet + MoeMlp);
11355    /// keeping the engine's own version structural rather than name-based means a new
11356    /// checkpoint of the same shape inherits the default, and a different shape does not.
11357    /// pub(crate) since lane/graph-launch-guard-sweep-20260831: `dspark_vg_admission_debt`
11358    /// consults it so the MTP-route pool stops escaping the admission charge.
11359    pub(crate) fn vgraph_family_default(&self) -> bool {
11360        let has_linear = self
11361            .layers
11362            .iter()
11363            .any(|l| matches!(l.mixer, Mixer::Linear(_)));
11364        let has_moe = self
11365            .layers
11366            .iter()
11367            .any(|l| matches!(l.ffn, crate::hybrid::Ffn::Moe(_)));
11368        has_linear && has_moe
11369    }
11370
11371    fn sliding_gated_moe_batch_program(&self) -> bool {
11372        self.uses_sliding_gated_moe_program()
11373    }
11374
11375    fn gemma_batch_program(&self) -> bool {
11376        self.uses_gemma_program()
11377    }
11378
11379    /// Reduced-matrix admission for increment 1. This deliberately does not change the PP-2
11380    /// serving policy: the worker calls it only after `MEMRA_SPEC_PIPE=1` and an explicit spec
11381    /// session already exist.
11382    pub fn spec_pipe_available(&self, e: &Engine) -> bool {
11383        if std::env::var("MEMRA_SPEC_PIPE").as_deref() != Ok("1")
11384            || !spec_devacc()
11385            || spec_replay_env_enabled()
11386            || spec_stream()
11387            || std::env::var("MEMRA_SPEC_ADAPT").as_deref() == Ok("1")
11388            || std::env::var("MEMRA_SPEC_PMIN0").as_deref() == Ok("1")
11389            || std::env::var("MEMRA_SPEC_PP_ANATOMY").as_deref() == Ok("1")
11390            || std::env::var("MEMRA_SPEC_PMIN")
11391                .ok()
11392                .and_then(|v| v.parse::<f32>().ok())
11393                .unwrap_or(0.0)
11394                > 0.0
11395            || self.is_gemma4_e4b()
11396            || self.gemma_batch_program()
11397            || self.mtp.is_none()
11398            || !self.mtp_extra.is_empty()
11399            // Both paired lanes would otherwise hold the model-global verify-graph mutex across
11400            // setup and wait for each other. Independent graph pools are future work; the pair
11401            // requires the explicit eager-verify arm today.
11402            || crate::spec::spec_verify_graph_env()
11403                .unwrap_or_else(|| self.vgraph_family_default())
11404        {
11405            return false;
11406        }
11407        let Some(cuts) = crate::pp::pp_cuts(self.layers.len()) else {
11408            return false;
11409        };
11410        if cuts.len() != 3 || crate::pp::pp2_streams_off() || !crate::pp::spec_pp_on() {
11411            return false;
11412        }
11413        crate::pp::PpNRt::get(e)
11414            .map(|rt| rt.n_stages() == 2 && rt.cross_device())
11415            .unwrap_or(false)
11416    }
11417
11418    /// Two warm greedy continuation bursts over one PP-2 interval coordinator. The two existing
11419    /// `generate_spec_inner2` call stacks own all per-session round locals; only phase issue order
11420    /// changes. No callback is accepted in increment 1 — the worker publishes each completed burst.
11421    #[allow(clippy::too_many_arguments)]
11422    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
11423    pub fn generate_spec_session_pair(
11424        &self,
11425        e: &Engine,
11426        sess_a: &mut SpecSession,
11427        max_new_a: usize,
11428        k_a: usize,
11429        sess_b: &mut SpecSession,
11430        max_new_b: usize,
11431        k_b: usize,
11432    ) -> Result<((Vec<u32>, usize, usize), (Vec<u32>, usize, usize)), Box<dyn std::error::Error>>
11433    {
11434        self.refuse_hyper("generate_spec_session_pair")?;
11435        if !self.spec_pipe_available(e) {
11436            return Err("two-session speculative pipeline is outside its reduced matrix".into());
11437        }
11438        let rt = crate::pp::PpNRt::get(e)?;
11439        let pp_walk = rt.acquire_walk("generate_spec_session_pair")?;
11440        let pp_permit = rt.walk_permit(&pp_walk, "generate_spec_session_pair")?;
11441        if max_new_a == 0 || max_new_b == 0 || k_a == 0 || k_b == 0 {
11442            return Err(
11443                "two-session speculative pipeline requires non-empty positive-K bursts".into(),
11444            );
11445        }
11446        for sess in [&*sess_a, &*sess_b] {
11447            if sess.committed.is_empty()
11448                || sess.last_h.is_none()
11449                || (sess.next_pred.is_none() && sess.pending_tok.is_none())
11450            {
11451                return Err("two-session speculative pipeline requires warm continuations".into());
11452            }
11453        }
11454
11455        let graph_ok = std::env::var("MEMRA_SPEC_NOGRAPH").is_err()
11456            && !spec_host_embd()
11457            && self.mtp_graph_capturable()
11458            && self.mtp_extra.is_empty()
11459            && !crate::model::full_prec_enabled();
11460        let graph_a = graph_ok && k_a + 2 < 96;
11461        let graph_b = graph_ok && k_b + 2 < 96;
11462        let was_tracking = e.ctx().is_event_tracking();
11463        if (graph_a || graph_b) && was_tracking {
11464            unsafe {
11465                e.ctx().disable_event_tracking();
11466            }
11467        }
11468
11469        static LOGGED: std::sync::Once = std::sync::Once::new();
11470        LOGGED.call_once(|| {
11471            eprintln!("[spec-pipe] two-session PP-2 continuation pipeline engaged");
11472        });
11473        let sync = std::sync::Arc::new(SpecPipeSync::new());
11474        let lane_a = SpecPipeLane {
11475            sync: sync.clone(),
11476            lane: 0,
11477            rt,
11478            walk_permit: pp_permit.clone(),
11479        };
11480        let lane_b = SpecPipeLane {
11481            sync,
11482            lane: 1,
11483            rt,
11484            walk_permit: pp_permit,
11485        };
11486        let mut sess_b_ptr = SpecPipeSessionPtr(sess_b as *mut SpecSession);
11487        let (result_a, result_b) = std::thread::scope(|scope| {
11488            let b = scope.spawn(move || {
11489                let mut finish = SpecPipeFinish::new(&lane_b);
11490                let sess_b = unsafe { sess_b_ptr.get_mut() };
11491                let result = (|| -> Result<_, String> {
11492                    e.ctx().bind_to_thread().map_err(|err| err.to_string())?;
11493                    self.generate_spec_inner2(
11494                        e,
11495                        &[],
11496                        max_new_b,
11497                        k_b,
11498                        graph_b,
11499                        Some(sess_b),
11500                        None,
11501                        None,
11502                        None,
11503                        None,
11504                        Some(&lane_b),
11505                    )
11506                    .map_err(|err| err.to_string())
11507                })();
11508                finish.close(result.is_err());
11509                result
11510            });
11511            let mut finish = SpecPipeFinish::new(&lane_a);
11512            let result_a = self.generate_spec_inner2(
11513                e,
11514                &[],
11515                max_new_a,
11516                k_a,
11517                graph_a,
11518                Some(sess_a),
11519                None,
11520                None,
11521                None,
11522                None,
11523                Some(&lane_a),
11524            );
11525            finish.close(result_a.is_err());
11526            let result_b = b
11527                .join()
11528                .map_err(|_| "paired speculative session B panicked".to_string())
11529                .and_then(|r| r);
11530            (result_a, result_b)
11531        });
11532
11533        if (graph_a || graph_b) && was_tracking {
11534            unsafe {
11535                e.ctx().enable_event_tracking();
11536            }
11537        }
11538        let result_a = result_a?;
11539        let result_b = result_b.map_err(|err| -> Box<dyn std::error::Error> { err.into() })?;
11540        Ok((result_a, result_b))
11541    }
11542
11543    /// One spec-decode turn on a live session. `suffix` = the NEW tokens only (turn N+1's user
11544    /// message rendered through the chat template continuation). Returns (new tokens emitted,
11545    /// drafted, accepted); session.committed grows by suffix + emitted.
11546    pub fn generate_spec_session(
11547        &self,
11548        e: &Engine,
11549        sess: &mut SpecSession,
11550        suffix: &[u32],
11551        max_new: usize,
11552        k: usize,
11553    ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
11554        self.generate_spec_session_sampled(e, sess, suffix, max_new, k, None, None)
11555    }
11556
11557    /// Serve-path sampled spec: routes the burst through the rejection-sampling verify with
11558    /// per-SESSION Philox continuity (sess.sctr/uctr). None = env-driven (CLI) or greedy.
11559    /// Filters (top-k/p/min-p) apply SYMMETRICALLY to draft q and verify p — distribution-exact
11560    /// for the filtered target (feat/filtered-spec).
11561    ///
11562    /// `on_commit` (sse-cadence, 2026-08-05): called with each newly-emitted slice of the
11563    /// output — once right after the prime's first token, then once per round commit — so a
11564    /// streaming caller can flush text at round cadence instead of once per burst. The slices
11565    /// are disjoint, in order, and concatenate to exactly the returned token vec. Emission-
11566    /// timing only: token bytes, session state, and exactness are untouched.
11567    ///
11568    /// The returned bool is a CONTINUE-VERDICT (admission yield, 2026-08-06): `false` ends
11569    /// the burst at the current round boundary, exactly as if `max_new` had been reached —
11570    /// the caller's scheduler regains control without waiting the burst out. Burst size is
11571    /// content-neutral (spec-levers battery), so an early exit moves WHEN the burst returns,
11572    /// never what tokens say. The slice may be EMPTY (a poll-only boundary — round-stream
11573    /// drains and the defensive tail flush can land with nothing new committed).
11574    #[allow(clippy::too_many_arguments)]
11575    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
11576    pub fn generate_spec_session_sampled(
11577        &self,
11578        e: &Engine,
11579        sess: &mut SpecSession,
11580        suffix: &[u32],
11581        max_new: usize,
11582        k: usize,
11583        sampling: Option<SpecSampling>,
11584        on_commit: Option<&mut dyn FnMut(&[u32]) -> bool>,
11585    ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
11586        self.generate_spec_session_sampled_prime_split(
11587            e, sess, suffix, max_new, k, sampling, None, on_commit,
11588        )
11589    }
11590
11591    /// Serve-only cold-prime segmentation twin. `prime_split` is the same stable boundary the
11592    /// plain worker would honor before entering its sub-floor tokenwise tail; warm continuations
11593    /// pass `None` and stay on the existing zero-prime path.
11594    #[allow(clippy::too_many_arguments)]
11595    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
11596    pub fn generate_spec_session_sampled_prime_split(
11597        &self,
11598        e: &Engine,
11599        sess: &mut SpecSession,
11600        suffix: &[u32],
11601        max_new: usize,
11602        k: usize,
11603        sampling: Option<SpecSampling>,
11604        prime_split: Option<usize>,
11605        on_commit: Option<&mut dyn FnMut(&[u32]) -> bool>,
11606    ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
11607        self.generate_spec_session_constrained_prime_split(
11608            e,
11609            sess,
11610            suffix,
11611            max_new,
11612            k,
11613            sampling,
11614            None,
11615            prime_split,
11616            on_commit,
11617        )
11618    }
11619
11620    /// `generate_spec_session_sampled` + GRAMMAR (constrained decoding, 2026-08-03): the
11621    /// hook truncates acceptance at the first grammar-illegal token AFTER the exactness
11622    /// verify (grammar is an extra rejection rule, ordering like the batched-verify twins)
11623    /// and replaces an illegal bonus with the MASKED argmax of the target's own verify
11624    /// column — token-identical to constrained plain greedy decode. GREEDY only (the
11625    /// worker routes sampled constrained to plain decode). Acceptance under tight grammars
11626    /// may drop (drafter is unconstrained); that is measured, not hidden.
11627    #[allow(clippy::too_many_arguments)]
11628    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
11629    pub fn generate_spec_session_constrained(
11630        &self,
11631        e: &Engine,
11632        sess: &mut SpecSession,
11633        suffix: &[u32],
11634        max_new: usize,
11635        k: usize,
11636        sampling: Option<SpecSampling>,
11637        constraint: Option<&mut dyn SpecConstraint>,
11638        on_commit: Option<&mut dyn FnMut(&[u32]) -> bool>,
11639    ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
11640        self.generate_spec_session_constrained_prime_split(
11641            e, sess, suffix, max_new, k, sampling, constraint, None, on_commit,
11642        )
11643    }
11644
11645    #[allow(clippy::too_many_arguments)]
11646    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
11647    pub fn generate_spec_session_constrained_prime_split(
11648        &self,
11649        e: &Engine,
11650        sess: &mut SpecSession,
11651        suffix: &[u32],
11652        max_new: usize,
11653        k: usize,
11654        sampling: Option<SpecSampling>,
11655        constraint: Option<&mut dyn SpecConstraint>,
11656        prime_split: Option<usize>,
11657        on_commit: Option<&mut dyn FnMut(&[u32]) -> bool>,
11658    ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
11659        if constraint.is_some() && sampling.is_some_and(|s| s.temp > 0.0) {
11660            return Err(
11661                "constrained spec decode is greedy-only (worker routes sampled \
11662                        constrained to plain decode)"
11663                    .into(),
11664            );
11665        }
11666        // PENDING-CARRY entry flush: a carried bonus precedes any new suffix in the sequence,
11667        // so it must commit BEFORE the suffix primes; the sampled path doesn't carry (its
11668        // round-0 accept needs the commit pass's logits). Empty-suffix greedy bursts — the
11669        // serve continuation case — consume the carry in-loop with zero solo passes.
11670        if sess.pending_tok.is_some()
11671            && (!suffix.is_empty() || sampling.is_some_and(|s| s.temp > 0.0))
11672        {
11673            self.spec_flush_pending(e, sess, sampling)?;
11674        }
11675
11676        // FULL_PREC forces the EAGER draft: the graph capture would enclose cuBLASLt f32 GEMV
11677        // (the FloatBf16 else-branches) and a bf16_to_f32 dequant alloc — neither is stream-capture
11678        // safe. Eager rides matmul/matmul_decode_exact, which dequant FloatBf16 on use. (§item 2.)
11679        // Multi-head MTP (mtp_extra non-empty) no longer disqualifies: the chain captures
11680        // per-head graphs (lane/step37-draft-graph-serving-20260830, MEMRA_MTP_CHAIN_GRAPH).
11681        let graph_draft = std::env::var("MEMRA_SPEC_NOGRAPH").is_err()
11682            && !spec_host_embd()
11683            && self.mtp_graph_capturable()
11684            && k + 2 < 96
11685            && !crate::model::full_prec_enabled();
11686        let was_tracking = e.ctx().is_event_tracking();
11687        if graph_draft && was_tracking {
11688            unsafe {
11689                e.ctx().disable_event_tracking();
11690            }
11691        }
11692        let r = self.generate_spec_inner2(
11693            e,
11694            suffix,
11695            max_new,
11696            k,
11697            graph_draft,
11698            Some(sess),
11699            sampling,
11700            constraint,
11701            on_commit,
11702            prime_split,
11703            None,
11704        );
11705        if graph_draft && was_tracking {
11706            unsafe {
11707                e.ctx().enable_event_tracking();
11708            }
11709        }
11710        let (out, d, a) = r?;
11711        Ok((out, d, a))
11712    }
11713
11714    pub fn generate_spec(
11715        &self,
11716        e: &Engine,
11717        prompt: &[u32],
11718        max_new: usize,
11719        k: usize,
11720    ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
11721        // glm5 T-parallel verify door (lane/glm5-tparallel-verify): an hc trunk with a
11722        // loaded DRAFT SOURCE — the embedded MTP head OR the DFlash2 drafter
11723        // (lane/glm5-dflash-draft-src) — routes to the glm5 draft->verify->rollback loop —
11724        // MEMRA_GLM5_SPEC=1 only (default OFF; flag row in FLAGS.md). Unset/0 falls
11725        // through to the standing named refusal below, byte-identical to the pre-lane
11726        // binary. Same fail-closed manifest stance as the generic path: an unqualified
11727        // MtpSpec rewrite refuses before any drafting.
11728        if self.hyper.is_some()
11729            && crate::glm_spec::glm5_spec_on()
11730            && (self.mtp.is_some() || self.glm5_dflash.is_some())
11731        {
11732            if !self.rewrite_allowed(memra_gguf::execution_manifest::RewriteSurface::MtpSpec) {
11733                return Err("speculative rewrite is not qualified for this ModelPlan".into());
11734            }
11735            return self.generate_spec_glm5(e, prompt, max_new, k);
11736        }
11737        self.refuse_hyper("generate_spec")?;
11738        if crate::pp::pp_cuts(self.layers.len()).is_some()
11739            && !self.rewrite_allowed(memra_gguf::execution_manifest::RewriteSurface::Pipeline)
11740        {
11741            return Err("pipeline rewrite is not qualified for speculative decode".into());
11742        }
11743        if !self.rewrite_allowed(memra_gguf::execution_manifest::RewriteSurface::MtpSpec) {
11744            return Err("speculative rewrite is not qualified for this ModelPlan".into());
11745        }
11746        // FULL_PREC forces eager (see generate_spec_session note): CUDA graph capture cannot
11747        // enclose cuBLASLt f32 GEMV or the bf16_to_f32 dequant alloc the FloatBf16 path needs.
11748        // Multi-head MTP no longer disqualifies (chain graphs; see generate_spec_session).
11749        let graph_draft = std::env::var("MEMRA_SPEC_NOGRAPH").is_err()
11750            && !spec_host_embd()
11751            && self.mtp_graph_capturable()
11752            && k + 2 < 96
11753            && !crate::model::full_prec_enabled();
11754        if !graph_draft {
11755            return self.generate_spec_inner2(
11756                e, prompt, max_new, k, false, None, None, None, None, None, None,
11757            );
11758        }
11759        let was_tracking = e.ctx().is_event_tracking();
11760        if was_tracking {
11761            unsafe {
11762                e.ctx().disable_event_tracking();
11763            }
11764        }
11765        let r = self.generate_spec_inner2(
11766            e, prompt, max_new, k, true, None, None, None, None, None, None,
11767        );
11768        if was_tracking {
11769            unsafe {
11770                e.ctx().enable_event_tracking();
11771            }
11772        }
11773        r
11774    }
11775
11776    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
11777    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
11778    fn generate_spec_inner2(
11779        &self,
11780        e: &Engine,
11781        prompt: &[u32],
11782        max_new: usize,
11783        k: usize,
11784        graph_draft: bool,
11785        mut sess: Option<&mut SpecSession>,
11786        sampling: Option<SpecSampling>,
11787        mut constraint: Option<&mut dyn SpecConstraint>,
11788        mut on_commit: Option<&mut dyn FnMut(&[u32]) -> bool>,
11789        prime_split: Option<usize>,
11790        pipe: Option<&SpecPipeLane>,
11791    ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
11792        assert!(k >= 1, "k must be >= 1");
11793        let pipe_setup_walk = match pipe {
11794            Some(p) => Some(p.setup_begin()?),
11795            None => None,
11796        };
11797        // sse-cadence flush cursor: everything in out[..flushed] has been handed to on_commit.
11798        let mut flushed = 0usize;
11799        // admission yield (2026-08-06): on_commit's continue-verdict; false = end the burst
11800        // at the next round boundary (same exit as max_new reached — the session tail runs).
11801        // Initialized by the unconditional post-prime flush below.
11802        let mut keep_going;
11803        let mtp = self
11804            .mtp
11805            .as_ref()
11806            .expect("generate_spec requires an MTP head (nextn_predict_layers>0)");
11807        let n_vocab = self.output.out_features();
11808        // FR-Spec: the draft head may be TRIMMED (fewer rows than n_vocab); the draft argmax runs
11809        // over the draft vocab and the winning index maps through d2t to a TARGET token id.
11810        // Everything downstream (verify/accept/commit) sees target ids only — exactness unchanged.
11811        let d_vocab = mtp
11812            .shared_head_head
11813            .as_ref()
11814            .unwrap_or(&self.output)
11815            .out_features();
11816        if !self.mtp_extra.is_empty() {
11817            if self.plan.draft_source != memra_gguf::model_plan::DraftSourcePlan::Embedded
11818                || self.plan.mtp_blocks.len() != self.mtp_head_count()
11819            {
11820                return Err(
11821                    "multi-head MTP requires one embedded canonical block per loaded head".into(),
11822                );
11823            }
11824            // TRIMMED chains (2026-08-27): every head must carry the SAME d2t — the ranking is
11825            // token-frequency and head-independent, and every downstream remap (per-step argmax,
11826            // stream pack, sampled d2t_dev) reads head 0's map, so equality is what makes that
11827            // single map correct for the whole chain. Mixed trimmed/untrimmed is refused.
11828            for (offset, head) in self.mtp_extra.iter().enumerate() {
11829                if head.d2t != mtp.d2t
11830                    || head
11831                        .shared_head_head
11832                        .as_ref()
11833                        .unwrap_or(&self.output)
11834                        .out_features()
11835                        != d_vocab
11836                {
11837                    return Err(format!(
11838                        "embedded MTP head {} has incompatible draft vocabulary",
11839                        offset + 1
11840                    )
11841                    .into());
11842                }
11843            }
11844            eprintln!(
11845                "[mtp-chain] heads={} policy=step-modulo prefix-replay kv=per-head",
11846                self.mtp_head_count()
11847            );
11848        }
11849        let n_embd = self.cfg.n_embd as usize;
11850        // SESSION MODE: reuse the live cache/scratch, prime only the suffix. `base` = tokens
11851        // already committed (their state is in the caches); 0 = fresh single-shot call.
11852        let session_mode = sess.is_some();
11853        let max_ctx = match sess.as_ref() {
11854            Some(s) => s.cache.max_ctx,
11855            None => prompt.len() + max_new + k + 8,
11856        };
11857        let mut own_cache;
11858        let mut own_scratch;
11859        // PREFIX-CACHE capture request threaded out of the session (lane/spec-prefix-cache):
11860        // (requested split, destination list). Single-shot per burst; fresh calls have none.
11861        let mut sess_capture: Option<(Option<usize>, &mut Vec<SpecBoundaryCapture>)> = None;
11862        // STABLE-BOUNDARY turn-checkpoint request (lane/frspec-multiturn-cache): ABSOLUTE
11863        // committed-length position; consumed one-shot like `capture_at`. None = legacy
11864        // prompt-end capture below.
11865        let mut ckpt_req: Option<usize> = None;
11866        // FAIL-SAFE bit threaded out of the session (see `SpecSession::capture_disabled`).
11867        let mut sess_capture_disabled = false;
11868        let (
11869            cache,
11870            scratch,
11871            mut sess_tail,
11872            mut sess_draft_slot,
11873            mut sess_pending_slot,
11874            sess_ckpt_slot,
11875            sess_telem,
11876        ): (
11877            &mut Cache,
11878            &mut MtpScratch,
11879            Option<(
11880                &mut Vec<u32>,
11881                &mut Option<CudaSlice<f32>>,
11882                &mut Option<u32>,
11883                &mut u32,
11884                &mut u32,
11885            )>,
11886            Option<&mut Option<DraftGraphCtx>>,
11887            Option<&mut Option<u32>>,
11888            Option<&mut Option<SpecCheckpoint>>,
11889            Option<&SpecTelemetryCounters>,
11890        ) = match sess.take() {
11891            Some(sr) => {
11892                let SpecSession {
11893                    cache,
11894                    scratch,
11895                    committed,
11896                    last_h,
11897                    next_pred,
11898                    sctr: s_sctr,
11899                    uctr: s_uctr,
11900                    draft_ctx,
11901                    pending_tok,
11902                    turn_ckpt,
11903                    telem,
11904                    capture_at,
11905                    boundary_captures,
11906                    ckpt_at,
11907                    capture_disabled,
11908                } = sr;
11909                sess_capture_disabled = *capture_disabled;
11910                sess_capture = Some((capture_at.take(), boundary_captures));
11911                ckpt_req = ckpt_at.take();
11912                (
11913                    cache,
11914                    scratch,
11915                    Some((committed, last_h, next_pred, s_sctr, s_uctr)),
11916                    Some(draft_ctx),
11917                    Some(pending_tok),
11918                    Some(turn_ckpt),
11919                    Some(telem),
11920                )
11921            }
11922            None => {
11923                // STAGE-OWNED KV (lane/pp2-spec 2026-08-06) — see `new_session`. Door shut =
11924                // `Cache::new` verbatim.
11925                own_cache = crate::pp::new_cache_planned(e, &self.cfg, &self.plan, max_ctx)?;
11926                // Persistent scratch = max_ctx rows (~2KB/token quantized).
11927                own_scratch = self.new_mtp_scratch(e, max_ctx)?;
11928                (
11929                    &mut own_cache,
11930                    &mut own_scratch,
11931                    None,
11932                    None,
11933                    None,
11934                    None,
11935                    None,
11936                )
11937            }
11938        };
11939        cache.ensure_usable("generate_spec")?;
11940        if scratch.plane_count() != self.mtp_head_count() {
11941            return Err(format!(
11942                "MTP scratch/head count mismatch ({}/{})",
11943                scratch.plane_count(),
11944                self.mtp_head_count()
11945            )
11946            .into());
11947        }
11948        let base = cache.pos;
11949        // PENDING-CARRY consume (2026-08-01): a carried bonus reaches here only on the
11950        // empty-suffix GREEDY continuation path (generate_spec_session_sampled flushed every
11951        // other case). It enters the round loop as round-0's pending — verify col 0 — exactly
11952        // like a mid-burst full-accept boundary: no init feed, no tail commit pass.
11953        let carried_pending: Option<u32> = sess_pending_slot.as_mut().and_then(|s| s.take());
11954        // PERSISTENT DRAFT KV (the only mode since 2026-07-08 — the legacy round-local scratch,
11955        // MEMRA_SPEC_KVLOCAL, measured -35 acceptance pts on the 27B p3 sweep and was removed;
11956        // acceptance-only — exactness is verify's job either way).
11957        // HIDDEN-PAIRING CONVENTION (DEFAULT = predecessor-row, 2026-07-04 — the 27B acceptance
11958        // unlock, +16pts): the MTP head is TRAINED on rows pairing token x_p with the trunk
11959        // hidden of its PREDECESSOR h_{p-1} (the reference engine's mtp_update shifts the target
11960        // hiddens right by one; its draft step 0 feeds (id_last, TRUE hidden of the row id_last
11961        // was sampled from)). memra's historical convention paired SAME-ROW (x_p, h_p) in the fill
11962        // and seeded chain step 0 through an extra MTP pass on a duplicated token (the
11963        // pseudo-seed) — measured 27B p2 K=3 acceptance 0.569 vs 0.731, p3 0.445 vs 0.63+, and
11964        // the chain steps j>=1 were already predecessor-shaped, so ONLY the fill + step-0 seed
11965        // move. The fill shifts by one and the chain seeds from the predecessor's true hidden
11966        // DIRECTLY (vh_seed / vx[j-1]) — the pseudo pass disappears (one MTP-block pass saved
11967        // per round on top of the acceptance win). Draft-quality-only: exactness stays the
11968        // verify's job either way. (The legacy same-row pairing seam, MEMRA_SPEC_HSAME, and its
11969        // pseudo-seed passes were removed 2026-07-08 — predecessor pairing won by +16 acc pts;
11970        // the legacy round-local scratch, MEMRA_SPEC_KVLOCAL, went with it.)
11971        // REPLAY-FREE PARTIAL ACCEPT (default, 2026-07-03): partial rounds keep the verify's own
11972        // bit-identical committed-prefix state (KV truncate + recur rebuild from the VerifyCkpt)
11973        // and leave the bonus PENDING — no duplicate trunk pass (profiled ~0.54 extra full weight
11974        // reads/round at long ctx). MEMRA_SPEC_REPLAY=1 restores the legacy rollback+replay (A/B
11975        // + fallback seam).
11976        // Qwen35-MoE replay pin LIFTED (lane/draftcost-moe, 2026-08-20). The pin's stated
11977        // bar — the retained verify-state commit proven equivalent to sequential serving —
11978        // was waiting on this arch running the serving batched verify class, which the
11979        // t-parallel admission (this lane, increment 1) provided: the VerifyCkpt the
11980        // replay-free commit consumes is now produced by the SAME serving-class verify that
11981        // qualified dense qwen35 on 2026-08-15 (where the per-round duplicate replay
11982        // measured 69 -> 30 tok/s). Qualification receipts (run-spec K=1..8 both arms,
11983        // 8-prompt replay-vs-replay-free canary, long-prompt cell):
11984        // research/draftcost-moe-20260820/RECEIPTS.md. MEMRA_SPEC_REPLAY=1 stays the
11985        // rollback + A/B seam.
11986        let spec_replay = spec_replay_env_enabled();
11987        if constraint.is_some() && spec_replay {
11988            return Err(
11989                "constrained spec decode does not support MEMRA_SPEC_REPLAY=1 \
11990                        (legacy replay commits an unmasked bonus)"
11991                    .into(),
11992            );
11993        }
11994        // TRUE-HIDDEN REFRESH (default in persistent-draft-KV mode): every round overwrites the
11995        // committed positions' scratch entries from the verify's exact hiddens (mtp_kv_fill batch)
11996        // instead of keeping chain-approximate entries. MEMRA_SPEC_NOREFRESH=1 = legacy (A/B seam).
11997        let refresh = std::env::var("MEMRA_SPEC_NOREFRESH").is_err();
11998        if !refresh && !self.mtp_extra.is_empty() {
11999            return Err("multi-head MTP requires exact accepted-prefix refresh".into());
12000        }
12001
12002        // prime: BATCHED cache prime (prime_cache — the measured #1 e2e gap: tokenwise primed at
12003        // ~102/38 tok/s vs the engine's ~2000-5900 tok/s batched prefill). prime_cache returns the
12004        // full pre-output_norm hidden stack [T, n_embd], which IS prompt_h (the persistent-draft-KV
12005        // mtp_kv_fill input) — no per-token collection needed. Prompts below PRIME_MIN_T, and
12006        // MEMRA_PRIME_TOKENWISE=1, and frozen Hy3 CPU/GPU expert splits take the tokenwise
12007        // decode_step_h loop. The latter avoids transient GPU staging of the spilled expert bank.
12008        // EMPTY-SUFFIX CONTINUATION (serve bursts): a session turn with NO new tokens resumes
12009        // generation exactly where the last turn stopped — no prime at all. The stashed
12010        // `next_pred` plays prime_logits' role: it is the token produced from the logits after
12011        // committed.last() by the same rule this entry applies to a cold prime's last row —
12012        // an argmax when greedy, a `sample_boundary_token` draw when sampled (the burst tail,
12013        // or `spec_session_from_restored` for a converted prefix-cache hit, did the drawing
12014        // where the sampler and the session's Philox counters were live). `last_h` seeds the
12015        // predecessor pairing below. Fresh calls and non-empty suffixes take the normal path.
12016        let continuation = prompt.is_empty();
12017        if continuation {
12018            assert!(session_mode, "empty prompt requires a session");
12019            assert!(
12020                sess_tail
12021                    .as_ref()
12022                    .is_some_and(|(c, lh, np, _, _)| !c.is_empty()
12023                        && lh.is_some()
12024                        && (np.is_some() || carried_pending.is_some())),
12025                "empty-suffix continuation needs a primed session (committed + last_h + next_pred|pending)"
12026            );
12027        }
12028        let mut prime_logits;
12029        let mut prompt_h: Option<CudaSlice<f32>> = None;
12030        let t_prime = std::time::Instant::now();
12031        let batched_prime = !continuation
12032            && prompt.len() >= crate::hybrid_forward::PRIME_MIN_T
12033            && std::env::var("MEMRA_PRIME_TOKENWISE").is_err()
12034            && !e.frozen_cpu_experts_prefer_tokenwise_prime();
12035        let prime_split = prime_split.filter(|&split| split > 0 && split < prompt.len());
12036        if prime_split.is_some() && continuation {
12037            return Err("spec prime split requires a non-empty prime".into());
12038        }
12039        // STABLE-BOUNDARY TURN CHECKPOINT stop (lane/frspec-multiturn-cache, 2026-08-21):
12040        // the worker's `ckpt_at` request, ABSOLUTE -> prompt-relative. On WARM bursts
12041        // (base != 0, an affinity-rewound or pool-resumed session priming its own delta)
12042        // this is the only stop; on COLD bursts it usually coincides with `prime_split`
12043        // (both are the plain tier's stable pre-generation boundary). A boundary the prime
12044        // cannot honor (outside this prime's range) silently drops the capture — the
12045        // turn_ckpt convention: the next turn re-primes in full, never a wrong resume.
12046        let ckpt_rel = if continuation {
12047            None
12048        } else {
12049            ckpt_req
12050                .and_then(|abs| abs.checked_sub(base))
12051                .filter(|&r| r > 0 && r < prompt.len())
12052        };
12053        // Prime stops, ordered: each is a boundary the prime halts at so the in-place GDN
12054        // conv/ssm state can be snapshotted there (the only moment it exists). One stop =
12055        // the legacy single-split program, byte-for-byte.
12056        let mut stops: Vec<usize> = Vec::new();
12057        for b in [prime_split, ckpt_rel].into_iter().flatten() {
12058            if !stops.contains(&b) {
12059                stops.push(b);
12060            }
12061        }
12062        stops.sort_unstable();
12063        // Captured at the ckpt stop, installed into the session slot post-prime (replacing
12064        // the legacy prompt-end capture). Some(None) = capture attempted and failed -> the
12065        // slot is cleared (a stale checkpoint would rewind to the WRONG boundary).
12066        let mut ckpt_early: Option<Option<SpecCheckpoint>> = None;
12067        if continuation {
12068            prime_logits = Vec::new();
12069        } else if !stops.is_empty() {
12070            if let Some(&first) = stops.first()
12071                && prime_split == Some(first)
12072                && first < crate::hybrid_forward::PRIME_MIN_T
12073            {
12074                return Err(format!(
12075                    "spec prime split {first} is below PRIME_MIN_T {}",
12076                    crate::hybrid_forward::PRIME_MIN_T,
12077                )
12078                .into());
12079            }
12080            // Mirror the plain worker's boundary stops exactly. Each segment is a
12081            // request-level prime (`queued_after` keeps Step35 arm selection independent of
12082            // the stops — tick-seg law); a segment below PRIME_MIN_T (and the final tail
12083            // under MEMRA_PRIME_TOKENWISE) takes the same eager tokenwise continuation as
12084            // prefill_tick. Retain every hidden row so the draft scratch fill remains one
12085            // coherent prompt.
12086            let mut h_all = e.uninit(prompt.len() * n_embd)?;
12087            prime_logits = Vec::new();
12088            let mut prev = 0usize;
12089            for seg_end in stops.iter().copied().chain(std::iter::once(prompt.len())) {
12090                if seg_end <= prev {
12091                    continue;
12092                }
12093                let seg = &prompt[prev..seg_end];
12094                let is_final = seg_end == prompt.len();
12095                let batched_seg = seg.len() >= crate::hybrid_forward::PRIME_MIN_T
12096                    && (!is_final
12097                        || (std::env::var("MEMRA_PRIME_TOKENWISE").is_err()
12098                            && !e.frozen_cpu_experts_prefer_tokenwise_prime()));
12099                if batched_seg {
12100                    let (l, _, h_seg) =
12101                        self.prime_cache(e, seg, &mut *cache, prompt.len() - seg_end)?;
12102                    e.copy_into(&mut h_all, prev * n_embd, &h_seg, seg.len() * n_embd)?;
12103                    prime_logits = l;
12104                } else {
12105                    for (i, &tok) in seg.iter().enumerate() {
12106                        let (l, h) = self.decode_step_h(e, tok, &mut *cache)?;
12107                        e.copy_into(&mut h_all, (prev + i) * n_embd, &h, n_embd)?;
12108                        prime_logits = l;
12109                    }
12110                }
12111                prev = seg_end;
12112                if is_final {
12113                    break;
12114                }
12115                debug_assert_eq!(cache.pos, base + seg_end, "prime stop landed off boundary");
12116                // PREFIX-CACHE BOUNDARY CAPTURE (lane/spec-prefix-cache): the GDN conv/ssm
12117                // states are about to be advanced in place by the next segment, so this is
12118                // the ONLY moment the boundary's recurrent state exists. Capture iff the
12119                // worker requested exactly this stop (cold sessions only — `capture_at` is
12120                // never armed warm). A failed snapshot is silent (turn_ckpt convention) —
12121                // publication is an optimization, never a correctness dependency.
12122                if base == 0
12123                    && let Some((requested, slot)) = sess_capture.as_mut()
12124                {
12125                    // Publish at the requested miss-LCP stop (the shared-prefix class)
12126                    // AND at the stable-boundary stop (the next-turn re-render class,
12127                    // lane/frspec-multiturn-cache) — the same boundary set the plain
12128                    // prefill tick learns. Without the second entry, the turn after a
12129                    // cold re-park could only hit the OLDER lcp entry (the measured
12130                    // one-turn transient: t3 restored 607 of 24122 while the plain arm
12131                    // rewound to 15222). Dedupe is the worker sweep's has_key.
12132                    if (*requested == Some(seg_end) || ckpt_rel == Some(seg_end))
12133                        && let Ok(snap) = cache.snapshot(e)
12134                    {
12135                        slot.push(SpecBoundaryCapture {
12136                            snap,
12137                            pos: seg_end,
12138                            logits: prime_logits.clone(),
12139                            // rows [0..seg_end) of h_all are primed — the following
12140                            // segments append, never overwrite.
12141                            last_h: capture_boundary_hidden(e, &h_all, seg_end, n_embd),
12142                            latent_tails: Vec::new(),
12143                        });
12144                    }
12145                }
12146                // SESSION-AFFINITY TURN CHECKPOINT at the STABLE boundary (see `ckpt_at`):
12147                // same snapshot mechanics, installed post-prime in place of the prompt-end
12148                // capture the re-render class always diverged below.
12149                if ckpt_rel == Some(seg_end) {
12150                    let anchor: Result<CudaSlice<f32>, Box<dyn std::error::Error>> =
12151                        e.uninit(n_embd).and_then(|mut a| {
12152                            e.copy_view_into(
12153                                &mut a,
12154                                0,
12155                                &h_all.slice((seg_end - 1) * n_embd..seg_end * n_embd),
12156                                n_embd,
12157                            )?;
12158                            Ok(a)
12159                        });
12160                    ckpt_early = Some(match (cache.snapshot(e), anchor) {
12161                        (Ok(snap), Ok(last_h)) => Some(SpecCheckpoint {
12162                            snap,
12163                            pos: base + seg_end,
12164                            last_h,
12165                        }),
12166                        _ => None,
12167                    });
12168                }
12169            }
12170            if std::env::var("MEMRA_SPEC_STATS").as_deref() == Ok("1") {
12171                eprintln!(
12172                    "[spec-prime] stops={stops:?} tail={}",
12173                    prompt.len() - stops.last().copied().unwrap_or(0)
12174                );
12175            }
12176            prompt_h = Some(h_all);
12177        } else if batched_prime {
12178            let (l, _h_seed, hiddens) = self.prime_cache(e, prompt, &mut *cache, 0)?;
12179            prime_logits = l;
12180            prompt_h = Some(hiddens);
12181        } else {
12182            prime_logits = Vec::new();
12183            prompt_h = Some(e.uninit(prompt.len() * n_embd)?);
12184            for (i, &tok) in prompt.iter().enumerate() {
12185                let (l, h) = self.spec_target_step_h(e, tok, &mut *cache)?;
12186                if let Some(ph) = prompt_h.as_mut() {
12187                    e.copy_into(ph, i * n_embd, &h, n_embd)?;
12188                }
12189                prime_logits = l;
12190            }
12191        }
12192        e.stream().synchronize()?;
12193        // PREFIX-CACHE SEED CAPTURE (lane/spec-prefix-cache): boundary == prompt end (the seed
12194        // case — no shared-prefix split, publish the whole prompt). The prime just finished, so
12195        // cache.pos == base + prompt.len() and the recurrent state IS the boundary state;
12196        // prime_logits are the boundary logits. Cold sessions only (base == 0) — same law as
12197        // prime_split. The mid-prompt capture above already consumed the request if it matched.
12198        if !continuation
12199            && base == 0
12200            && let Some((requested, slot)) = sess_capture.as_mut()
12201            && *requested == Some(prompt.len())
12202            && slot.is_empty()
12203        {
12204            debug_assert_eq!(cache.pos, prompt.len(), "seed capture off prompt end");
12205            if let Ok(snap) = cache.snapshot(e) {
12206                slot.push(SpecBoundaryCapture {
12207                    snap,
12208                    pos: prompt.len(),
12209                    logits: prime_logits.clone(),
12210                    last_h: prompt_h
12211                        .as_ref()
12212                        .map(|ph| capture_boundary_hidden(e, ph, prompt.len(), n_embd))
12213                        .unwrap_or_default(),
12214                    latent_tails: Vec::new(),
12215                });
12216            }
12217        }
12218        // Harness timing contract (see crate::PRIME_NANOS): gen-only throughput without the
12219        // prime-subtraction hack.
12220        crate::PRIME_NANOS.store(
12221            t_prime.elapsed().as_nanos() as u64,
12222            std::sync::atomic::Ordering::Relaxed,
12223        );
12224
12225        let (embd_qt, embd_rb) = self.embd.qt_and_row_bytes(n_embd);
12226        // Resident table is fastest when it fits. Large spill deployments can preserve that HBM
12227        // for expert-cache slots and gather only the exact rows needed by MTP/verify from host.
12228        let host_embd = spec_host_embd();
12229        let embd_gpu = if host_embd {
12230            None
12231        } else {
12232            Some(
12233                self.embd_gpu
12234                    .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload")),
12235            )
12236        };
12237        let embd_dev = embd_gpu.map(|g| (g, embd_qt, embd_rb));
12238        if host_embd {
12239            eprintln!(
12240                "[spec] host-row embedding: {} bytes kept off HBM",
12241                self.embd.raw.len()
12242            );
12243        }
12244        let mut out: Vec<u32> = Vec::with_capacity(max_new);
12245        let mut total_drafted = 0usize;
12246        let mut total_accepted = 0usize;
12247
12248        // --- SAMPLER FIRST (lane/sampled-spec-quality, 2026-08-19) ---
12249        // The sampler config, the session's Philox counters and the penalty window are parsed
12250        // HERE, above the boundary-token selection, because the boundary token must be drawn
12251        // from the sampler the request asked for. Pre-lane this block sat ~50 lines BELOW the
12252        // selection, which is the whole mechanical reason the boundary token was an argmax:
12253        // the sampler state was not in scope yet. Nothing here depends on the round loop, so
12254        // moving it up is a pure reordering for greedy (`sampled == false` ⇒ every branch
12255        // below takes the argmax path it always took).
12256        // --- SAMPLED SPEC (MEMRA_SPEC_TEMP>0, research/sampled-spec-impl-map.md): rejection-
12257        // sampling verify (Leviathan/Chen) — accept draft x at u < p(x)/q(x), resample from
12258        // norm(max(0,p-q)) on reject, bonus sampled from p on full accept. Counter-based Philox
12259        // everywhere (seed, event) -> reproducible. temp==0/unset = the greedy path, untouched.
12260        let sp = sampling.unwrap_or_else(|| SpecSampling {
12261            temp: std::env::var("MEMRA_SPEC_TEMP")
12262                .ok()
12263                .and_then(|v| v.parse().ok())
12264                .unwrap_or(0.0),
12265            seed: std::env::var("MEMRA_SEED")
12266                .ok()
12267                .and_then(|v| v.parse().ok())
12268                .unwrap_or(42),
12269            top_k: std::env::var("MEMRA_TOP_K")
12270                .ok()
12271                .and_then(|v| v.parse().ok())
12272                .unwrap_or(0),
12273            top_p: std::env::var("MEMRA_TOP_P")
12274                .ok()
12275                .and_then(|v| v.parse().ok())
12276                .unwrap_or(1.0),
12277            min_p: std::env::var("MEMRA_MIN_P")
12278                .ok()
12279                .and_then(|v| v.parse().ok())
12280                .unwrap_or(0.0),
12281            penalty_last_n: std::env::var("MEMRA_PENALTY_LAST_N")
12282                .ok()
12283                .and_then(|v| v.parse().ok())
12284                .unwrap_or(0),
12285            penalty_repeat: std::env::var("MEMRA_PENALTY_REPEAT")
12286                .ok()
12287                .and_then(|v| v.parse().ok())
12288                .unwrap_or(1.0),
12289            penalty_freq: std::env::var("MEMRA_PENALTY_FREQ")
12290                .ok()
12291                .and_then(|v| v.parse().ok())
12292                .unwrap_or(0.0),
12293            penalty_present: std::env::var("MEMRA_PENALTY_PRESENT")
12294                .ok()
12295                .and_then(|v| v.parse().ok())
12296                .unwrap_or(0.0),
12297        });
12298        let (sp_temp, sp_seed) = (sp.temp, sp.seed);
12299        let sampled = sp_temp > 0.0;
12300        // Counters resume from the session (burst continuity: randomness must never repeat
12301        // across generate_spec_session calls); one-shot callers start at (0,0). Read through
12302        // sess_tail — `sess` was take()n into it above, so sess.as_ref() here is always None.
12303        let mut sctr: u32 = sess_tail.as_ref().map(|(_, _, _, s, _)| **s).unwrap_or(0);
12304        let mut uctr: u32 = sess_tail.as_ref().map(|(_, _, _, _, u)| **u).unwrap_or(0);
12305        // Penalties (v2.1): applied to COPIES of q rows and p columns symmetrically (exactness
12306        // for the penalized+filtered target). History = generated tokens, host-tracked window.
12307        let pen_on = sampled
12308            && sp.penalty_last_n > 0
12309            && (sp.penalty_repeat != 1.0 || sp.penalty_freq != 0.0 || sp.penalty_present != 0.0);
12310        // SESSION-SPANNING PENALTY WINDOW (Item 2). Pre-lane this was
12311        // `prompt.iter().rev().take(64).rev()` — the BURST's suffix slice — so a continuation
12312        // burst (the majority of a stream's tokens, and ALL of a converted cache hit's) started
12313        // with an EMPTY penalty history and the client's repetition/frequency/presence penalties
12314        // silently reset at every burst boundary. The window now spans `committed ++ prompt`,
12315        // which is what the API contract says and what the plain sampler's own `history` does.
12316        // Byte-identical to the pre-lane seed for a cold turn-1 burst at the default window.
12317        let mut pen_hist: Vec<u32> = if pen_on {
12318            let sess_hist: &[u32] = if spec_pen_session_on() {
12319                sess_tail
12320                    .as_ref()
12321                    .map(|(c, ..)| c.as_slice())
12322                    .unwrap_or(&[])
12323            } else {
12324                &[] // MEMRA_SPEC_PEN_SESSION=0: pre-lane burst-local window
12325            };
12326            pen_window_seed(sess_hist, prompt, sp.penalty_last_n)
12327        } else {
12328            Vec::new()
12329        };
12330        // First generated token = the BOUNDARY token: greedy takes the argmax of the prompt's
12331        // last logits (== greedy's first token, byte-contract); SAMPLED draws it from the
12332        // request's own filtered/penalized target through the session's Philox stream
12333        // (`sample_boundary_token`, lane/sampled-spec-quality Item 1 — pre-lane this was an
12334        // argmax in both regimes, so ~1 token per burst of a sampled stream was greedy).
12335        // Emit it, then FEED it to establish the loop invariant below.
12336        // PENDING-CARRY: the carried bonus was already emitted by the LAST burst — it becomes
12337        // last_token WITHOUT re-emission, and round 0 consumes it as pending (no init feed).
12338        // CONSTRAINED entry rules: the first emitted token is the MASKED argmax of the
12339        // prompt's last logits (plain constrained-greedy identity); a continuation without
12340        // a carried pending would emit an UNMASKED stashed next_pred — refused loudly (the
12341        // worker never resumes constrained sessions from the pool, so this cannot fire).
12342        if let Some(c) = constraint.as_deref_mut() {
12343            if continuation && carried_pending.is_none() {
12344                return Err("constrained spec continuation requires a carried pending \
12345                            (pool resume is unconstrained-only)"
12346                    .into());
12347            }
12348            if !continuation {
12349                c.mask_logits(&mut prime_logits)
12350                    .map_err(|e2| format!("constraint: {e2}"))?;
12351            }
12352        }
12353        let mut last_token = if let Some(b) = carried_pending {
12354            b
12355        } else if continuation {
12356            // A continuation's boundary token was DRAWN by the burst that stashed it (the
12357            // session tail below), or by `spec_session_from_restored` for a converted
12358            // prefix-cache hit — in both cases from the correct logits row with this same
12359            // session's Philox stream, which is why it can be consumed here as-is.
12360            sess_tail.as_ref().unwrap().2.unwrap()
12361        } else if sampled && constraint.is_none() && spec_sampled_boundary_on() {
12362            sample_boundary_token(e, &prime_logits, &sp, &pen_hist, &mut sctr, "cold-prime")?
12363        } else {
12364            // greedy (byte contract), the rollback door, or constrained (masked-argmax
12365            // identity — the worker routes sampled+constrained to the plain path, and this
12366            // function refuses the combination outright above).
12367            argmax(&prime_logits) as u32
12368        };
12369        if pen_on {
12370            // The boundary token is a GENERATED token: the plain sampler `accept()`s every
12371            // emitted token into its penalty history, and pre-lane the burst's first token
12372            // was invisible to penalties forever (never pushed, and never in `committed`
12373            // until this burst's tail). Covers the carry/continuation seeds too — neither is
12374            // in `committed` yet.
12375            pen_hist.push(last_token);
12376        }
12377        if carried_pending.is_none() {
12378            out.push(last_token);
12379            // grammar advances with every emitted token (carried pendings were consumed
12380            // by the burst that emitted them).
12381            if let Some(c) = constraint.as_deref_mut() {
12382                c.consume(last_token)
12383                    .map_err(|e2| format!("constraint: {e2}"))?;
12384            }
12385        }
12386        if continuation {
12387            // draft-KV invariant: entries [0..base) are the session's exact fills; truncate any
12388            // overhang so the chain's first append lands at slot base (== committed.len()).
12389            scratch.set_len(e, base)?;
12390        }
12391        // sse-cadence: hand the caller every not-yet-flushed token (disjoint in-order slices
12392        // concatenating to the full `out`). Called after the prime's first token and after each
12393        // round commit — emission timing only, token bytes untouched. The slice may be EMPTY
12394        // (poll-only boundary: zero-round folds commit nothing new); returns the caller's
12395        // continue-verdict (admission yield, 2026-08-06) — false ends the burst at this round.
12396        #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
12397        fn flush_commit(
12398            cb: &mut Option<&mut dyn FnMut(&[u32]) -> bool>,
12399            out: &[u32],
12400            flushed: &mut usize,
12401        ) -> bool {
12402            if let Some(f) = cb.as_mut() {
12403                let keep = f(&out[*flushed..]);
12404                *flushed = out.len();
12405                keep
12406            } else {
12407                true
12408            }
12409        }
12410        keep_going = flush_commit(&mut on_commit, &out, &mut flushed);
12411        // INVARIANT at loop top: `last_token` is the most-recently-committed/emitted token, its
12412        // KV+recur state IS in `cache` (cache.pos = position right AFTER last_token), `last_pred`
12413        // is the greedy ARGMAX of the logits that predict the token FOLLOWING last_token, and
12414        // `h_seed` = last_token's pre-output_norm hidden. Establish it by feeding last_token once
12415        // (mirrors plain greedy). DEVICE-ARGMAX lever: the accept walk only ever consumes the
12416        // argmax of those logits — never the full vector — so a host u32 replaces the Vec<f32>.
12417        // Trimmed heads: q lives on the trimmed vocab; accept gathers use the TRIMMED index and
12418        // the residual scatters q into target-id space (q=-inf off-trim — the head cannot propose
12419        // those, so their residual mass is p(x), correct by construction).
12420        let d2t_dev: Option<CudaSlice<u32>> = if sampled || crate::spec::spec_stream() {
12421            match &mtp.d2t {
12422                Some(map) => Some(e.htod_u32_v(map)?),
12423                None => None,
12424            }
12425        } else {
12426            None
12427        };
12428        let mut q_full_buf: Option<CudaSlice<f32>> = None;
12429        // host Philox4x32-10 accept-test uniforms: module fn `host_u01` (shared with the
12430        // dspark sampled-admission walk); byte-identical to the closure it replaces.
12431        let mut draft_logits: Vec<CudaSlice<f32>> = Vec::new(); // retained head logits (q), per slot
12432        let mut draft_stats: Vec<(f32, f32, f32)> = Vec::new(); // (row_max, th_e, z_e) per slot
12433        let mut perturb_buf: Option<CudaSlice<f32>> = None; // gumbel scratch (max(n_vocab,d_vocab))
12434        let mut sample_tok = e.alloc_u32_zeroed(1)?; // residual/bonus sample out
12435        let mut col_buf: Option<CudaSlice<f32>> = None; // materialized verify column
12436        let mut pen_hist_d: Option<CudaSlice<u32>> = None;
12437        let mut pcol_buf: Option<CudaSlice<f32>> = None; // penalized p-column scratch
12438        // MEMRA_SPEC_SETUP_TRACE=1 (diagnostics): per-call wall decomposition of the burst
12439        // SETUP + TAIL segments (the round loop's internals are MEMRA_SPEC_PHASE's job) —
12440        // built to pin the serve per-burst fixed cost (research/spec-serving-20260801).
12441        let setup_trace = std::env::var("MEMRA_SPEC_SETUP_TRACE").as_deref() == Ok("1");
12442        let t_ent = std::time::Instant::now();
12443
12444        // SESSION-AFFINITY TURN CHECKPOINT (lane/session-affinity, 2026-08-05): capture the
12445        // PROMPT-END boundary state so a LATER turn can rewind here and re-prime only its own
12446        // delta instead of the whole conversation. See `SpecCheckpoint` for why this boundary is
12447        // the one that matters (a history-rewriting client mutates what the session GENERATED,
12448        // so the next turn's prompt agrees with this one up to exactly here).
12449        //
12450        // WHERE — AND WHY THIS EXACT LINE. Right after the trunk prime, BEFORE the init feed
12451        // (`decode_step_h(last_token)`) and before round 0: the last instant at which the caches
12452        // hold exactly `base + prompt.len()` rows and nothing generated.
12453        //
12454        // This was WRONG in the first cut of this lane: the capture sat after the draft-KV fill,
12455        // which is also after the init feed, so `cache.pos` was `base + prompt.len() + 1` — the
12456        // boundary included the FIRST GENERATED TOKEN. That token is the first thing inside the
12457        // `<think>` block the client strips, so every later turn's diff diverged exactly one
12458        // token below the checkpoint and affinity declined 100% of the time. Measured on the
12459        // owner regime: "history diverged at 12233 of checkpoint 12234". The off-by-one made the
12460        // whole mechanism inert while looking, from the outside, like a working
12461        // correctness-declines-safely path — hence the decline log carries the offsets.
12462        //
12463        // The full-attn planes are `len`-truncatable so the snapshot copies only the GDN conv/ssm
12464        // state (the reason a spec session could not rewind before). The draft scratch needs no
12465        // copy: rows below the boundary are rewritten by the next turn's own fill.
12466        //
12467        // WHEN: non-empty prime only. An empty-suffix continuation burst adds no prompt boundary
12468        // (its "prompt end" IS the previous checkpoint's, already held), so it keeps the existing
12469        // checkpoint rather than replacing it with a strictly worse one.
12470        //
12471        // FAILURE IS SILENT BY DESIGN: on a VRAM-tight rig the snapshot alloc can fail. That
12472        // costs the NEXT turn its rewind (it re-primes fully, today's behavior) and must never
12473        // fail the burst that is already running — so the error is swallowed, loud only under
12474        // MEMRA_DEBUG_SPEC.
12475        //
12476        // STABLE-BOUNDARY OVERRIDE (lane/frspec-multiturn-cache, 2026-08-21): the prompt-end
12477        // posture above was DISPROVED for the think-posture template class — the prompt's own
12478        // tail is the live generation header (`<|im_start|>assistant\n<think>\n`) that the
12479        // next turn's re-render replaces, so the diff diverged a couple tokens BELOW the
12480        // checkpoint and affinity declined 100% of multi-turn agent traffic (the same class
12481        // the plain tier fixed on 2026-08-09 via `plain_checkpoint_boundary`; the port to the
12482        // spec tier is this lane). When the worker armed `ckpt_at`, the capture happened at
12483        // that stop inside the prime above (`ckpt_early`) and is installed here instead;
12484        // capture-attempted-but-failed clears the slot exactly like the legacy arm.
12485        if let Some(slot) = sess_ckpt_slot {
12486            if let Some(early) = ckpt_early {
12487                if early.is_none() && std::env::var("MEMRA_DEBUG_SPEC").is_ok() {
12488                    eprintln!(
12489                        "[spec] stable-boundary turn checkpoint skipped; \
12490                               next turn re-primes in full"
12491                    );
12492                }
12493                *slot = early;
12494            } else if !continuation {
12495                let pos = cache.pos;
12496                debug_assert_eq!(
12497                    pos,
12498                    base + prompt.len(),
12499                    "turn checkpoint must sit at the prompt end, before the init feed"
12500                );
12501                let anchor: Result<CudaSlice<f32>, Box<dyn std::error::Error>> =
12502                    if let Some(ph) = &prompt_h {
12503                        // hidden of the LAST primed row = the predecessor anchor at this
12504                        // boundary (exactly what a fresh prime of committed[..pos] leaves in
12505                        // last_h, and what the next prime's fill reads for its first row).
12506                        let np = prompt.len();
12507                        e.uninit(n_embd).and_then(|mut a| {
12508                            e.copy_view_into(
12509                                &mut a,
12510                                0,
12511                                &ph.slice((np - 1) * n_embd..np * n_embd),
12512                                n_embd,
12513                            )?;
12514                            Ok(a)
12515                        })
12516                    } else {
12517                        Err("no prompt hiddens".into())
12518                    };
12519                match (cache.snapshot(e), anchor) {
12520                    (Ok(snap), Ok(last_h)) => {
12521                        *slot = Some(SpecCheckpoint { snap, pos, last_h });
12522                    }
12523                    (s, a) => {
12524                        *slot = None; // a stale checkpoint would rewind to the WRONG boundary
12525                        if std::env::var("MEMRA_DEBUG_SPEC").is_ok() {
12526                            let err = s
12527                                .err()
12528                                .map(|e| e.to_string())
12529                                .or_else(|| a.err().map(|e| e.to_string()))
12530                                .unwrap_or_default();
12531                            eprintln!(
12532                                "[spec] turn checkpoint skipped ({err}); \
12533                                       next turn re-primes in full"
12534                            );
12535                        }
12536                    }
12537                }
12538            }
12539        }
12540        // INIT FEED — skipped on a pending carry: last_token (the carried bonus) is NOT in the
12541        // caches and must NOT be fed solo; round 0's batched verify commits it as col 0. Its
12542        // seed/anchor hidden is the carried last_h (copied below); last_pred is dead in the
12543        // pending path (t_pred reads verify col 0 — the accept walk overwrites it).
12544        let mut last_pred = 0u32;
12545        let mut last_col_logits: Option<CudaSlice<f32>> = None;
12546        // CONSTRAINED: the init feed's logits back the (n_acc==0, base==0) masked-argmax
12547        // recompute in the grammar-truncation walk — retained host-side, round 0 only.
12548        let mut init_logits_host: Option<Vec<f32>> = None;
12549        let h_seed0: CudaSlice<f32> = if carried_pending.is_none() {
12550            let (init_logits, h) = self.spec_target_step_h(e, last_token, &mut *cache)?;
12551            last_pred = argmax(&init_logits) as u32;
12552            if constraint.is_some() {
12553                init_logits_host = Some(init_logits.clone());
12554            }
12555            // sampled mode: p-distribution after last_token, for the j==0/base==0 accept test.
12556            if sampled {
12557                last_col_logits = Some(e.htod(&init_logits)?);
12558            }
12559            h
12560        } else {
12561            // predecessor-row anchor: hidden of the last COMMITTED row (the carry contract).
12562            let lh = sess_tail
12563                .as_ref()
12564                .unwrap()
12565                .1
12566                .as_ref()
12567                .expect("pending carry requires last_h");
12568            e.clone_dtod(lh)?
12569        };
12570        let t_init = t_ent.elapsed();
12571        let mut last_col_stats: Option<(f32, f32, f32)> = None;
12572        // PERSISTENT h_seed buffer (allocated BEFORE any graph capture so no captured scratch can
12573        // alias it): every path that updates the round seed copies INTO it — no per-round allocs,
12574        // stable pointer for the graph-draft round-start copy.
12575        let mut h_seed_buf = e.clone_dtod(&h_seed0)?;
12576        // Predecessor-pairing trackers: `fill_prev` = trunk hidden AT the last COMMITTED row (the
12577        // predecessor of the next verify's col 0 — the reference's carried pending-h analogue;
12578        // also the predecessor-row hidden for the round-0 legacy-replay seed). At round 0 that
12579        // row is last_token's own (h_seed0). The chain step-0 seed under the pairing default =
12580        // hidden of the row BEFORE last_token = the prompt's last row at round 0 (h_seed_buf
12581        // overwritten below).
12582        let mut fill_prev = e.clone_dtod(&h_seed0)?;
12583        {
12584            if let Some(ph) = &prompt_h {
12585                let np = prompt.len();
12586                e.copy_view_into(
12587                    &mut h_seed_buf,
12588                    0,
12589                    &ph.slice((np - 1) * n_embd..np * n_embd),
12590                    n_embd,
12591                )?;
12592            } else if continuation
12593                && let Some((_, lh, _, _, _)) = sess_tail.as_ref()
12594                && let Some(lh) = lh.as_ref()
12595            {
12596                e.copy_into(&mut h_seed_buf, 0, lh, n_embd)?;
12597            }
12598        }
12599        // Persistent device prediction slots for the accept walk (max k+1 verify columns).
12600        let mut preds_d = e.alloc_u32_zeroed(k + 2)?;
12601
12602        let debug_spec = std::env::var("MEMRA_DEBUG_SPEC").is_ok();
12603        let fork_mode = OptiForkGateMode::configured();
12604        // MEMRA_SPEC_STATS=1: per-slot accept histogram + draft-length histogram, printed once at
12605        // the end. Metric normalization vs the reference engine: BOTH engines count
12606        // accepted/drafted where the chain stopped at p-min and the sub-threshold token is
12607        // discarded uncounted — per-slot decay + chain-length mix are the extra dimensions.
12608        let spec_stats = std::env::var("MEMRA_SPEC_STATS").is_ok();
12609        let mut st_drafted = vec![0usize; k];
12610        let mut st_accepted = vec![0usize; k];
12611        let mut st_len_hist = vec![0usize; k + 1];
12612        let mut st_full = 0usize;
12613        // P-MIN CONFIDENCE GATE (MEMRA_SPEC_PMIN, the serve script's --spec-draft-p-min mechanism):
12614        // stop the draft chain early when the head's softmax confidence in its own pick drops
12615        // below p_min. Hoisted above the loop: the graph capture bakes the prob kernels iff on.
12616        static PMIN: std::sync::OnceLock<f32> = std::sync::OnceLock::new();
12617        let p_min = *PMIN.get_or_init(|| {
12618            std::env::var("MEMRA_SPEC_PMIN")
12619                .ok()
12620                .and_then(|v| v.parse().ok())
12621                .unwrap_or(0.0)
12622        });
12623        // ZERO-DRAFT ROUNDS (MEMRA_SPEC_PMIN0=1, vendored from llama.cpp's draft gating): let the
12624        // p-min gate apply at j==0 too, so a low-confidence round drafts NOTHING and the verify
12625        // batch is just the pending bonus (m=1 = a plain decode step). llama's 35B win rides
12626        // exactly this — draft acceptance 76% at mean len 2.5 because unpredictable stretches
12627        // never pay draft+verify overhead. Only legal when a pending bonus exists (an empty
12628        // verify batch is not); the j==0 exemption stays for pending-less rounds.
12629        let pmin0 = std::env::var("MEMRA_SPEC_PMIN0")
12630            .map(|v| v == "1")
12631            .unwrap_or(false);
12632
12633        // --- GRAPH DRAFT setup: persistent I/O buffers + ONE capture (2 warmups inside). The
12634        // warmups mutate scratch len_d / pos / tok / seed — all reset at every round start, so the
12635        // only restore needed is the scratch counter. Capture failure (e.g. a non-capturable
12636        // cuBLAS path in an exotic head) falls back to the eager draft chain.
12637        // PER-SESSION PERSISTENCE (2026-08-01): session calls reuse the DraftGraphCtx parked on
12638        // the SpecSession — the capture (2 warmup head forwards + instantiate) ran ONCE at the
12639        // session's first burst, not per burst (measured ~16ms/burst fixed cost on H100 q27,
12640        // research/spec-serving-20260801). Reuse is pointer-exact: the graph bakes the session's
12641        // own scratch KV (never realloc'd), the model's resident embedding, the OnceLock p_min,
12642        // and the g_* buffers carried in the ctx — replay dispatch is identical to a fresh
12643        // capture, so draft tokens are bit-identical (drafts never decide exactness anyway; the
12644        // verify arbitrates). Single-shot calls (sess=None) build a fresh ctx and drop it.
12645        let mut dctx: DraftGraphCtx = match sess_draft_slot.as_mut().and_then(|s| s.take()) {
12646            Some(c) => c,
12647            None => DraftGraphCtx::new(e, n_embd, if sampled { d_vocab } else { 1 })?,
12648        };
12649        // FAIL-SAFE (step-OOM park replay): pre-mark both fallback flags so no capture arm
12650        // below can fire — LOUD once per replayed session through the standard WARN line.
12651        if sess_capture_disabled {
12652            let reason =
12653                "session replayed after a step-OOM park; draft capture disabled (fail-safe)";
12654            let flip = dctx.failed.mark_greedy(reason);
12655            let flip_s = dctx.failed.mark_sampled(reason);
12656            if let Some(line) = flip.or(flip_s) {
12657                eprintln!("{line}");
12658            }
12659        }
12660        // A session that ran greedy bursts first sized g_q/g_perturb at 1; a sampled resume
12661        // needs d_vocab. Realloc is legal exactly while graph_s is None (nothing baked them).
12662        if sampled && dctx.g_q.len() < d_vocab {
12663            dctx.g_q = e.zeros(d_vocab)?;
12664            dctx.g_perturb = e.zeros(d_vocab)?;
12665        }
12666        // DRAFT-SIDE GRAMMAR MASK (lane/draft-mask, 2026-08-04): the drafter samples the
12667        // grammar's legal set, so proposals are legal BY CONSTRUCTION and the verify-side
12668        // truncation (the correctness backstop) stops cutting every tight-schema round.
12669        // The mask is one node inside the captured draft chain — presence is a CAPTURE-TIME
12670        // shape, so a parked graph of the other shape is dropped and recaptured.
12671        let dmask_on = constraint
12672            .as_deref()
12673            .is_some_and(|c| c.draft_mask_enabled());
12674        let dmask_words = if dmask_on { d_vocab.div_ceil(32) } else { 0 };
12675        if dmask_on && dctx.g_dmask.len() < dmask_words {
12676            dctx.g_dmask = e.alloc_u32_zeroed(dmask_words)?;
12677            dctx.graph = None; // the old capture baked the old (or no) mask pointer
12678            dctx.chain = None; // chain last-row graphs bake the same pointer
12679            dctx.failed.clear_greedy();
12680            dctx.keeper.clear();
12681        }
12682        if (dctx.graph.is_some() || dctx.chain.is_some()) && dctx.graph_masked != dmask_on {
12683            dctx.graph = None;
12684            dctx.chain = None;
12685            dctx.failed.clear_greedy();
12686            dctx.keeper.clear();
12687        }
12688        // MULTI-HEAD CHAIN mode (mtp_extra non-empty — step37's 3-head shipping shape): the
12689        // step-modulo prefix-replay chain captures PER-HEAD single-row graphs
12690        // (`DraftChainGraphs`) instead of the one self-feeding graph below; the single-head
12691        // capture arms are untouched and unreachable in this mode (the launch arms branch the
12692        // same way). This removes the historical `mtp_extra.is_empty()` capture exclusion —
12693        // and with it the silent no-attempt hole: a chain capture that FAILS now trips the
12694        // same LOUD draft-graph WARN as a single-head failure.
12695        let chain_mode = !self.mtp_extra.is_empty();
12696        // ---- PRE-CAPTURE VRAM RESERVE CHECK + PER-SESSION DRAFT-STATE MEASUREMENT ----
12697        // (lane/step37-vram-admission-20260830). `cap_eff0` opens the measurement bracket:
12698        // when any capture succeeds in THIS call, the effective-free delta across the whole
12699        // capture section is recorded as the model's per-session draft-state high-water
12700        // (admission charges it per spec-capable session — this state was charged at ZERO
12701        // before the lane). The reserve check runs BEFORE any capture arm can allocate: a
12702        // refused capture trips the same LOUD once-per-flip WARN class as a failed one, but
12703        // with the card's headroom still intact (the owner's single-session OOM was a capture
12704        // attempt walking the card to the edge and stranding the eager fallback at 5 MiB free).
12705        let cap_eff0 = e
12706            .ctx()
12707            .mem_get_info()
12708            .ok()
12709            .map(|(f, _)| f.saturating_add(e.pool_cached_bytes()));
12710        // Peak instrument for the same bracket: the CAPTURE-TIME peak (warmup transients +
12711        // instantiate scratch, alive together) dwarfs the parked delta — measured on the
12712        // owner shape: a capture whose PARKED state reads ~2.6GB walked a ~7GB-free card to
12713        // OOM mid-capture. Reset the pool watermark here; read it at bracket end.
12714        let _ = e.pool_high_water_reset();
12715        let cap_used0 = e.pool_reserved_used().1;
12716        let mut captured_now = false;
12717        let mut capture_oom_entry_eff: Option<usize> = None;
12718        let capture_need = {
12719            let observed = self.draft_session_admission_bytes();
12720            if observed > 0 {
12721                observed
12722            } else {
12723                draft_capture_bootstrap_estimate(
12724                    if chain_mode { self.mtp_head_count() } else { 1 },
12725                    k,
12726                    d_vocab,
12727                    n_embd,
12728                )
12729            }
12730        };
12731        if spec_capture_gate_on()
12732            && graph_draft
12733            && !sampled
12734            && !dctx.failed.greedy_failed()
12735            && ((chain_mode && dctx.chain.is_none() && mtp_chain_graph_on())
12736                || (!chain_mode && dctx.graph.is_none()))
12737            && let Some(reason) = capture_headroom_refusal(e, capture_need)
12738            && let Some(line) = dctx.failed.mark_greedy(&reason)
12739        {
12740            eprintln!("{line}");
12741        }
12742        if graph_draft
12743            && !sampled
12744            && chain_mode
12745            && dctx.chain.is_none()
12746            && !dctx.failed.greedy_failed()
12747        {
12748            if mtp_chain_graph_on() {
12749                let heads_n = self.mtp_head_count();
12750                let DraftGraphCtx {
12751                    g_tok,
12752                    g_pos,
12753                    g_seed,
12754                    g_p,
12755                    g_dmask,
12756                    ..
12757                } = &mut dctx;
12758                if dmask_on {
12759                    e.htod_u32_into(g_dmask, &vec![u32::MAX; dmask_words])?;
12760                }
12761                let g_dmask_ro: &CudaSlice<u32> = &*g_dmask;
12762                let with_prob = p_min > 0.0;
12763                // CAPTURE-RETAIN (#68 fix): one keeper for the whole chain — every graph's
12764                // warmup transients stay pinned as long as any of them replays.
12765                let cap_res = (|| -> Result<DraftChainGraphs, Box<dyn std::error::Error>> {
12766                    // dcw door: same warmup headroom pre-arm as the single-head capture
12767                    // below — every plane, because each head's capture warmups append on
12768                    // its OWN plane. INSIDE the fallible closure (vram-admission lane): an
12769                    // OOM here used to `?` out of the whole burst as a step error; now it
12770                    // is a capture failure — LOUD WARN, eager chain serves.
12771                    if step35_draft_dcw_on() {
12772                        scratch.ensure_dcw_headroom(e, k + 2)?;
12773                    }
12774                    let mut interior = Vec::with_capacity(heads_n);
12775                    let mut last = Vec::with_capacity(heads_n);
12776                    let mut keeper: Vec<Box<dyn std::any::Any + Send>> = Vec::new();
12777                    for hi in 0..heads_n {
12778                        let head = self.mtp_head_at(hi);
12779                        // interior row: KV append + carrier only (`with_head=false` — the
12780                        // eager chain discards interior logits too, so this is the same
12781                        // consumed-byte program minus the dead full-vocab head matmul).
12782                        let (g, keep) = e.capture_graph_retained(|e| {
12783                            self.mtp_head_forward_cap(
12784                                e,
12785                                head,
12786                                g_tok,
12787                                g_pos,
12788                                g_seed,
12789                                g_p,
12790                                &mut *scratch,
12791                                hi,
12792                                false,
12793                                false,
12794                                embd_gpu.expect("graph draft requires resident embedding"),
12795                                embd_qt,
12796                                embd_rb,
12797                                d_vocab,
12798                                None,
12799                                None,
12800                                None,
12801                            )
12802                        })?;
12803                        // the warmups appended rows on plane hi; rewind before the next
12804                        // capture so successive warmups never outrun the pre-armed headroom.
12805                        scratch.set_plane_len(e, hi, base)?;
12806                        interior.push(g);
12807                        keeper.extend(keep);
12808                        // last row: head matmul + greedy argmax tail (+ p when the policy
12809                        // reads it, + the grammar-mask node when constrained).
12810                        let (g2, keep2) = e.capture_graph_retained(|e| {
12811                            self.mtp_head_forward_cap(
12812                                e,
12813                                head,
12814                                g_tok,
12815                                g_pos,
12816                                g_seed,
12817                                g_p,
12818                                &mut *scratch,
12819                                hi,
12820                                with_prob,
12821                                true,
12822                                embd_gpu.expect("graph draft requires resident embedding"),
12823                                embd_qt,
12824                                embd_rb,
12825                                d_vocab,
12826                                None,
12827                                None,
12828                                if dmask_on {
12829                                    Some((g_dmask_ro, dmask_words))
12830                                } else {
12831                                    None
12832                                },
12833                            )
12834                        })?;
12835                        scratch.set_plane_len(e, hi, base)?;
12836                        last.push(g2);
12837                        keeper.extend(keep2);
12838                    }
12839                    Ok(DraftChainGraphs {
12840                        interior,
12841                        last,
12842                        _keeper: keeper,
12843                    })
12844                })();
12845                match cap_res {
12846                    Ok(cg) => {
12847                        scratch.set_len(e, base)?;
12848                        // POSITIVE engagement receipt (the 3a lesson: a WARN-free boot is
12849                        // NOT evidence of capture — the captured state must name itself).
12850                        eprintln!(
12851                            "[mtp-chain-graph] captured mode=greedy heads={heads_n} \
12852                             interior={heads_n} last={heads_n} masked={}",
12853                            dmask_on as u8
12854                        );
12855                        dctx.chain = Some(cg);
12856                        dctx.graph_masked = dmask_on;
12857                        captured_now = true;
12858                    }
12859                    Err(err) => {
12860                        scratch.set_len(e, base)?;
12861                        // LOUD flip (audit Q2): a dropped draft graph is a coverage loss,
12862                        // never silent — now including the multi-head shipping shape.
12863                        // OOM RECOVERY (vram-admission lane): a failed attempt's freed
12864                        // transients sit CACHED in the async pool where the driver cannot
12865                        // see them; trim them back so the eager fallback (and any driver-
12866                        // side allocation) actually has the headroom the free suggests.
12867                        let mut reason = err.to_string();
12868                        if capture_err_is_oom(&reason) {
12869                            capture_oom_entry_eff = capture_oom_entry_eff.max(cap_eff0);
12870                            let trimmed = e.pool_trim_to_zero();
12871                            if trimmed > 0 {
12872                                reason.push_str(&format!(
12873                                    "; pool trimmed {}MB back to the driver",
12874                                    trimmed / (1 << 20)
12875                                ));
12876                            }
12877                        }
12878                        if let Some(line) = dctx.failed.mark_greedy(&reason) {
12879                            eprintln!("{line}");
12880                        }
12881                    }
12882                }
12883            } else {
12884                // Disarmed by MEMRA_MTP_CHAIN_GRAPH=0: say so once per process — the OFF arm
12885                // must be attributable in a boot log, never inferable from silence.
12886                static NOTE: std::sync::Once = std::sync::Once::new();
12887                NOTE.call_once(|| {
12888                    eprintln!(
12889                        "[spec] multi-head draft-chain capture disarmed \
12890                         (MEMRA_MTP_CHAIN_GRAPH=0); eager chain serves this shape"
12891                    );
12892                });
12893            }
12894        }
12895        if graph_draft
12896            && !sampled
12897            && !chain_mode
12898            && dctx.graph.is_none()
12899            && !dctx.failed.greedy_failed()
12900        {
12901            let DraftGraphCtx {
12902                g_tok,
12903                g_pos,
12904                g_seed,
12905                g_p,
12906                g_dmask,
12907                ..
12908            } = &mut dctx;
12909            // capture-time contents: ALL-ONES (ban nothing). A replay only ever runs after the
12910            // host uploads the position's real words, so the warmups stay grammar-free.
12911            if dmask_on {
12912                e.htod_u32_into(g_dmask, &vec![u32::MAX; dmask_words])?;
12913            }
12914            let g_dmask_ro: &CudaSlice<u32> = &*g_dmask;
12915            // CAPTURE-RETAIN (#68 fix): the warmup transients' pool addresses are baked into the
12916            // captured graph; the keeper pins them for the graph's lifetime. capture_graph (non-
12917            // retained) freed them at exit — safe for one-shot generate_spec (nothing else touches
12918            // the pool between replays) but WRONG for sessions: burst-boundary prime/fill/commit
12919            // passes (and, in serve, other sessions) recycle those addresses and the replay then
12920            // clobbers live buffers — the ST serve-spec corruption (research/serve-st-20260803).
12921            let cap_res = (|| {
12922                // dcw door: the capture warmups append device-counter rows the capture body
12923                // cannot rebase for; pre-arm ring headroom host-side (no-op on flat planes /
12924                // room-enough rings, and the door-off path is untouched). INSIDE the fallible
12925                // closure (vram-admission lane): an OOM here is a capture failure, not a
12926                // burst-killing step error.
12927                if step35_draft_dcw_on() {
12928                    scratch.ensure_dcw_headroom(e, k + 2)?;
12929                }
12930                e.capture_graph_retained(|e| {
12931                    self.mtp_head_forward_cap(
12932                        e,
12933                        mtp,
12934                        g_tok,
12935                        g_pos,
12936                        g_seed,
12937                        g_p,
12938                        &mut *scratch,
12939                        0,
12940                        p_min > 0.0 || fork_mode == OptiForkGateMode::Controller,
12941                        true,
12942                        embd_gpu.expect("graph draft requires resident embedding"),
12943                        embd_qt,
12944                        embd_rb,
12945                        d_vocab,
12946                        None,
12947                        None,
12948                        if dmask_on {
12949                            Some((g_dmask_ro, dmask_words))
12950                        } else {
12951                            None
12952                        },
12953                    )
12954                })
12955            })();
12956            match cap_res {
12957                Ok((g, keep)) => {
12958                    scratch.set_len(e, base)?;
12959                    dctx.graph = Some(g);
12960                    dctx.graph_masked = dmask_on;
12961                    dctx.keeper = keep;
12962                    captured_now = true;
12963                }
12964                Err(err) => {
12965                    scratch.set_len(e, base)?;
12966                    // LOUD flip (audit Q2): a dropped draft graph is a coverage loss, never
12967                    // silent. Once per flip — mark returns None on an already-failed ctx.
12968                    let mut reason = err.to_string();
12969                    if capture_err_is_oom(&reason) {
12970                        capture_oom_entry_eff = capture_oom_entry_eff.max(cap_eff0);
12971                        let trimmed = e.pool_trim_to_zero();
12972                        if trimmed > 0 {
12973                            reason.push_str(&format!(
12974                                "; pool trimmed {}MB back to the driver",
12975                                trimmed / (1 << 20)
12976                            ));
12977                        }
12978                    }
12979                    if let Some(line) = dctx.failed.mark_greedy(&reason) {
12980                        eprintln!("{line}");
12981                    }
12982                }
12983            }
12984        }
12985        // --- SAMPLED GRAPH DRAFT setup (step 3 of the sampled-spec arc): a SECOND capture, own
12986        // graph object, built only when sampled && graph-eligible — the greedy capture above is
12987        // untouched (and skipped when sampled: its graph would never be launched). Same head
12988        // forward, but the in-graph argmax reads GUMBEL-PERTURBED logits; the Philox event
12989        // counter lives in the persistent device g_ctr (bumped in-graph, host-seeded from sctr
12990        // once per round); the raw head logits land in the persistent g_q for the host's
12991        // per-replay async D2D into the round's q slot (q_slots, K x d_vocab, allocated once).
12992        // seed/temp are capture-time constants — baked into graph_s, so a pool-resumed request
12993        // with a different (seed, temp, k) drops the parked sampled graph and recaptures.
12994        // COST OF THE FRESH-SEED SERVE DEFAULT (dogfood F4, 2026-08-04): omitting `seed` on a
12995        // serve request now draws fresh per-request entropy (it used to default to a pinned 0),
12996        // so a seed-omitting request that RESUMES a parked spec session finds an s_key baked
12997        // with the PREVIOUS request's seed and pays one recapture. Bounded, and it does not
12998        // reopen the ~16ms/burst regression the persistent ctx exists to fix: a session's seed
12999        // is fixed for its whole lifetime (worker.rs reads s.sampler.seed() per burst), so
13000        // this compare misses at most ONCE per resumed request — the first burst recaptures
13001        // and every later burst in that request replays. A client that wants the parked graph
13002        // AND reproducibility supplies an explicit `seed`, honored exactly, which keeps s_key
13003        // stable across its whole conversation.
13004        // COMPOSITION RULE (fspec x gsd merge): the in-graph chain samples from the RAW
13005        // softmax — it can hold neither per-row filter stats nor the varying penalty history.
13006        // The sampled graph therefore engages only in the PURE-TEMP regime; filters/penalties
13007        // force the eager draft (which computes stats/penalties per row).
13008        // KEY THE WHOLE REGIME, not just the baked constants (lane/graph-s-key-exactness-
13009        // 20260819). `s_key` used to be `(seed, temp, k)`; the filters and penalties were left
13010        // out, so a filtered request resuming a session that parked a PURE-TEMP graph kept it —
13011        // and the launch site never re-asked `pure_temp`. See [`SampledGraphKey`] for what that
13012        // costs (an unconditional accept of out-of-head draft tokens, i.e. an exactness bug on
13013        // the request shape the vendor-default flip makes the majority).
13014        let s_key = SampledGraphKey::new(sp_seed, sp_temp, k, sp.top_k, sp.top_p, sp.min_p, pen_on);
13015        let pure_temp = s_key.pure_temp();
13016        // The regime the sampled graph may be captured/launched in: pure-temp always;
13017        // truncation-filtered when the filtered-capture door is on (the filter runs
13018        // IN-GRAPH — lane/step37-draft-graph-serving-20260830); penalties never.
13019        let s_capturable = s_key.graph_capturable();
13020        if sampled && dctx.s_key.is_some_and(|old| old != s_key) {
13021            dctx.graph_s = None;
13022            dctx.chain_s = None;
13023            dctx.failed.clear_sampled();
13024            dctx.s_key = None;
13025            dctx.q_slots.clear();
13026            dctx.keeper_s.clear();
13027        }
13028        // PRE-CAPTURE VRAM RESERVE CHECK, sampled arms (vram-admission lane): same contract
13029        // as the greedy check above — refuse BEFORE allocating, LOUD once, eager serves.
13030        if spec_capture_gate_on()
13031            && graph_draft
13032            && sampled
13033            && s_capturable
13034            && !dctx.failed.sampled_failed()
13035            && ((chain_mode && dctx.chain_s.is_none() && mtp_chain_graph_on())
13036                || (!chain_mode && dctx.graph_s.is_none()))
13037            && let Some(reason) = capture_headroom_refusal(e, capture_need)
13038            && let Some(line) = dctx.failed.mark_sampled(&reason)
13039        {
13040            eprintln!("{line}");
13041        }
13042        // FILTERED capture nodes need q slots sized d_vocab AND the stat slots; the pure-temp
13043        // body leaves g_th/g_z/g_mx untouched (they exist from ctx creation either way).
13044        if graph_draft
13045            && sampled
13046            && s_capturable
13047            && chain_mode
13048            && dctx.chain_s.is_none()
13049            && !dctx.failed.sampled_failed()
13050        {
13051            if mtp_chain_graph_on() {
13052                let heads_n = self.mtp_head_count();
13053                let filtered = s_key.filtered();
13054                let DraftGraphCtx {
13055                    g_tok,
13056                    g_pos,
13057                    g_seed,
13058                    g_p,
13059                    g_ctr,
13060                    g_perturb,
13061                    g_q,
13062                    g_rows0,
13063                    g_th,
13064                    g_z,
13065                    g_mx,
13066                    ..
13067                } = &mut dctx;
13068                let with_prob = p_min > 0.0;
13069                let cap_res = (|| -> Result<DraftChainGraphs, Box<dyn std::error::Error>> {
13070                    // dcw pre-arm INSIDE the fallible closure (vram-admission lane): an OOM
13071                    // here is a capture failure with the LOUD WARN, never a step error.
13072                    if step35_draft_dcw_on() {
13073                        scratch.ensure_dcw_headroom(e, k + 2)?;
13074                    }
13075                    let mut interior = Vec::with_capacity(heads_n);
13076                    let mut last = Vec::with_capacity(heads_n);
13077                    let mut keeper: Vec<Box<dyn std::any::Any + Send>> = Vec::new();
13078                    for hi in 0..heads_n {
13079                        let head = self.mtp_head_at(hi);
13080                        // interior row: no head, no draw — shared shape with the greedy
13081                        // chain's interior, captured per mode for keeper-lifetime hygiene.
13082                        let (g, keep) = e.capture_graph_retained(|e| {
13083                            self.mtp_head_forward_cap(
13084                                e,
13085                                head,
13086                                g_tok,
13087                                g_pos,
13088                                g_seed,
13089                                g_p,
13090                                &mut *scratch,
13091                                hi,
13092                                false,
13093                                false,
13094                                embd_gpu.expect("graph draft requires resident embedding"),
13095                                embd_qt,
13096                                embd_rb,
13097                                d_vocab,
13098                                None,
13099                                None,
13100                                None,
13101                            )
13102                        })?;
13103                        scratch.set_plane_len(e, hi, base)?;
13104                        interior.push(g);
13105                        keeper.extend(keep);
13106                        // last row: head matmul + the in-graph categorical draw (filtered
13107                        // nodes when the request carries filters).
13108                        let (g2, keep2) = e.capture_graph_retained(|e| {
13109                            self.mtp_head_forward_cap(
13110                                e,
13111                                head,
13112                                g_tok,
13113                                g_pos,
13114                                g_seed,
13115                                g_p,
13116                                &mut *scratch,
13117                                hi,
13118                                with_prob,
13119                                true,
13120                                embd_gpu.expect("graph draft requires resident embedding"),
13121                                embd_qt,
13122                                embd_rb,
13123                                d_vocab,
13124                                Some(SampledCapArgs {
13125                                    ctr: &mut *g_ctr,
13126                                    perturb: &mut *g_perturb,
13127                                    q_out: &mut *g_q,
13128                                    seed: sp_seed,
13129                                    temp: sp_temp,
13130                                    filt: if filtered {
13131                                        Some(SampledCapFilter {
13132                                            rows0: &*g_rows0,
13133                                            th: &mut *g_th,
13134                                            z: &mut *g_z,
13135                                            mx: &mut *g_mx,
13136                                            top_k: sp.top_k,
13137                                            top_p: sp.top_p,
13138                                            min_p: sp.min_p,
13139                                        })
13140                                    } else {
13141                                        None
13142                                    },
13143                                }),
13144                                None,
13145                                None, // constrained spec is greedy-only
13146                            )
13147                        })?;
13148                        scratch.set_plane_len(e, hi, base)?;
13149                        last.push(g2);
13150                        keeper.extend(keep2);
13151                    }
13152                    Ok(DraftChainGraphs {
13153                        interior,
13154                        last,
13155                        _keeper: keeper,
13156                    })
13157                })();
13158                match cap_res {
13159                    Ok(cg) => {
13160                        scratch.set_len(e, base)?;
13161                        // NO STRANDED PARTIAL STATE (vram-admission lane): the q-slot allocs
13162                        // after a successful capture are themselves fallible on a tight card.
13163                        // A mid-loop failure used to `?` out as a step error, leaving orphan
13164                        // slots parked on the ctx (wrong count, stale contents) for the next
13165                        // capture attempt to stack onto. Allocate all-or-nothing: on failure
13166                        // drop the fresh graphs AND the partial slots, mark the LOUD fallback.
13167                        dctx.q_slots.clear();
13168                        let slots = (0..k)
13169                            .map(|_| e.zeros(d_vocab))
13170                            .collect::<Result<Vec<_>, _>>();
13171                        match slots {
13172                            Ok(slots) => {
13173                                dctx.q_slots = slots;
13174                                eprintln!(
13175                                    "[mtp-chain-graph] captured mode=sampled heads={heads_n} \
13176                                     interior={heads_n} last={heads_n} filtered={} key={s_key:?}",
13177                                    s_key.filtered() as u8
13178                                );
13179                                dctx.chain_s = Some(cg);
13180                                dctx.s_key = Some(s_key);
13181                                captured_now = true;
13182                            }
13183                            Err(err) => {
13184                                drop(cg);
13185                                dctx.q_slots.clear();
13186                                let mut reason = format!("q-slot alloc failed: {err}");
13187                                if capture_err_is_oom(&reason) {
13188                                    capture_oom_entry_eff = capture_oom_entry_eff.max(cap_eff0);
13189                                    let trimmed = e.pool_trim_to_zero();
13190                                    if trimmed > 0 {
13191                                        reason.push_str(&format!(
13192                                            "; pool trimmed {}MB back to the driver",
13193                                            trimmed / (1 << 20)
13194                                        ));
13195                                    }
13196                                }
13197                                if let Some(line) = dctx.failed.mark_sampled(&reason) {
13198                                    eprintln!("{line}");
13199                                }
13200                            }
13201                        }
13202                    }
13203                    Err(err) => {
13204                        scratch.set_len(e, base)?;
13205                        let mut reason = err.to_string();
13206                        if capture_err_is_oom(&reason) {
13207                            capture_oom_entry_eff = capture_oom_entry_eff.max(cap_eff0);
13208                            let trimmed = e.pool_trim_to_zero();
13209                            if trimmed > 0 {
13210                                reason.push_str(&format!(
13211                                    "; pool trimmed {}MB back to the driver",
13212                                    trimmed / (1 << 20)
13213                                ));
13214                            }
13215                        }
13216                        if let Some(line) = dctx.failed.mark_sampled(&reason) {
13217                            eprintln!("{line}");
13218                        }
13219                    }
13220                }
13221            } else {
13222                static NOTE_S: std::sync::Once = std::sync::Once::new();
13223                NOTE_S.call_once(|| {
13224                    eprintln!(
13225                        "[spec] multi-head draft-chain capture disarmed \
13226                         (MEMRA_MTP_CHAIN_GRAPH=0); eager chain serves this shape"
13227                    );
13228                });
13229            }
13230        }
13231        if graph_draft
13232            && sampled
13233            && s_capturable
13234            && !chain_mode
13235            && dctx.graph_s.is_none()
13236            && !dctx.failed.sampled_failed()
13237        {
13238            let filtered = s_key.filtered();
13239            let DraftGraphCtx {
13240                g_tok,
13241                g_pos,
13242                g_seed,
13243                g_p,
13244                g_ctr,
13245                g_perturb,
13246                g_q,
13247                g_rows0,
13248                g_th,
13249                g_z,
13250                g_mx,
13251                ..
13252            } = &mut dctx;
13253            // CAPTURE-RETAIN (#68 fix): same keeper contract as the greedy capture above.
13254            let cap_res = (|| {
13255                // dcw pre-arm INSIDE the fallible closure (vram-admission lane): an OOM
13256                // here is a capture failure with the LOUD WARN, never a step error.
13257                if step35_draft_dcw_on() {
13258                    scratch.ensure_dcw_headroom(e, k + 2)?;
13259                }
13260                e.capture_graph_retained(|e| {
13261                    self.mtp_head_forward_cap(
13262                        e,
13263                        mtp,
13264                        g_tok,
13265                        g_pos,
13266                        g_seed,
13267                        g_p,
13268                        &mut *scratch,
13269                        0,
13270                        p_min > 0.0,
13271                        true,
13272                        embd_gpu.expect("graph draft requires resident embedding"),
13273                        embd_qt,
13274                        embd_rb,
13275                        d_vocab,
13276                        Some(SampledCapArgs {
13277                            ctr: &mut *g_ctr,
13278                            perturb: &mut *g_perturb,
13279                            q_out: &mut *g_q,
13280                            seed: sp_seed,
13281                            temp: sp_temp,
13282                            filt: if filtered {
13283                                Some(SampledCapFilter {
13284                                    rows0: &*g_rows0,
13285                                    th: &mut *g_th,
13286                                    z: &mut *g_z,
13287                                    mx: &mut *g_mx,
13288                                    top_k: sp.top_k,
13289                                    top_p: sp.top_p,
13290                                    min_p: sp.min_p,
13291                                })
13292                            } else {
13293                                None
13294                            },
13295                        }),
13296                        None,
13297                        None, // constrained spec is greedy-only — sampled never carries a hook
13298                    )
13299                })
13300            })();
13301            match cap_res {
13302                Ok((g, keep)) => {
13303                    scratch.set_len(e, base)?;
13304                    // NO STRANDED PARTIAL STATE: all-or-nothing q slots, same contract as
13305                    // the chain arm above.
13306                    dctx.q_slots.clear();
13307                    let slots = (0..k)
13308                        .map(|_| e.zeros(d_vocab))
13309                        .collect::<Result<Vec<_>, _>>();
13310                    match slots {
13311                        Ok(slots) => {
13312                            dctx.q_slots = slots;
13313                            dctx.graph_s = Some(g);
13314                            dctx.s_key = Some(s_key);
13315                            dctx.keeper_s = keep;
13316                            captured_now = true;
13317                        }
13318                        Err(err) => {
13319                            drop(g);
13320                            drop(keep);
13321                            dctx.q_slots.clear();
13322                            let mut reason = format!("q-slot alloc failed: {err}");
13323                            if capture_err_is_oom(&reason) {
13324                                capture_oom_entry_eff = capture_oom_entry_eff.max(cap_eff0);
13325                                let trimmed = e.pool_trim_to_zero();
13326                                if trimmed > 0 {
13327                                    reason.push_str(&format!(
13328                                        "; pool trimmed {}MB back to the driver",
13329                                        trimmed / (1 << 20)
13330                                    ));
13331                                }
13332                            }
13333                            if let Some(line) = dctx.failed.mark_sampled(&reason) {
13334                                eprintln!("{line}");
13335                            }
13336                        }
13337                    }
13338                }
13339                Err(err) => {
13340                    scratch.set_len(e, base)?;
13341                    // LOUD flip (audit Q2): same contract as the greedy capture above.
13342                    let mut reason = err.to_string();
13343                    if capture_err_is_oom(&reason) {
13344                        capture_oom_entry_eff = capture_oom_entry_eff.max(cap_eff0);
13345                        let trimmed = e.pool_trim_to_zero();
13346                        if trimmed > 0 {
13347                            reason.push_str(&format!(
13348                                "; pool trimmed {}MB back to the driver",
13349                                trimmed / (1 << 20)
13350                            ));
13351                        }
13352                    }
13353                    if let Some(line) = dctx.failed.mark_sampled(&reason) {
13354                        eprintln!("{line}");
13355                    }
13356                }
13357            }
13358        }
13359        // ---- PER-SESSION DRAFT-STATE MEASUREMENT bracket end (vram-admission lane): when a
13360        // capture landed in THIS call, the effective-free delta across the capture section is
13361        // this session's parked draft-graph state (keepers + q slots + instantiated graphs'
13362        // backing). Recorded as a model-owned high-water; admission charges it per
13363        // spec-capable session (see `draft_session_admission_bytes`).
13364        if captured_now
13365            && let Some(eff0) = cap_eff0
13366            && let Ok((f1, _)) = e.ctx().mem_get_info()
13367        {
13368            let eff1 = f1.saturating_add(e.pool_cached_bytes());
13369            let parked_delta = eff0.saturating_sub(eff1);
13370            let (_res_high, used_high) = e.pool_high_water_reset();
13371            let peak_delta = used_high.saturating_sub(cap_used0);
13372            let observed = parked_delta.max(peak_delta);
13373            if observed > 0
13374                && let Some(hw) = self.record_draft_state_bytes(observed)
13375            {
13376                eprintln!(
13377                    "[spec] draft-session state high-water: {}MB (max of parked delta {}MB \
13378                     and capture-time pool peak {}MB; charged per spec admission and gating \
13379                     future captures)",
13380                    hw / (1 << 20),
13381                    parked_delta / (1 << 20),
13382                    peak_delta / (1 << 20),
13383                );
13384            }
13385        }
13386        // FAILURE IS AN OBSERVATION TOO: a capture that OOM'd at entry-effective E proved
13387        // the capture-time peak exceeds E. Feed E into the gauge so every future gate
13388        // refuses at or below the headroom that just failed (self-healing even when the
13389        // boot probe is disarmed and the bootstrap estimate was blind).
13390        if let Some(entry_eff) = capture_oom_entry_eff
13391            && let Some(hw) = self.record_draft_state_bytes(entry_eff)
13392        {
13393            eprintln!(
13394                "[spec] draft-session capture appetite floor raised to {}MB: a capture \
13395                 attempt OOM'd with that much effective free (failure-observed bound)",
13396                hw / (1 << 20)
13397            );
13398        }
13399        // ---- EXACTNESS GUARD, the enforceable half (lane/graph-s-key-exactness-20260819,
13400        // widened by lane/step37-draft-graph-serving-20260830) ----
13401        // With the filters and penalties in `s_key`, a graph that SURVIVED the drop above was
13402        // captured under THIS request's exact regime, and capture requires `graph_capturable`
13403        // (pure-temp, or filtered with the in-graph filter nodes; never penalties) — so a
13404        // parked graph implies both. That implication is the whole exactness argument for the
13405        // graph arm, so it is asserted here rather than assumed: a future change that widens
13406        // the capture condition, narrows the key, or copies a `DraftGraphCtx` across regimes
13407        // fails LOUDLY at this line instead of silently drafting from a distribution the
13408        // verify never reconstructs. Release builds refuse the graph (drop it, draft eager)
13409        // rather than launching it; the launch site re-tests the regime independently.
13410        if sampled
13411            && (dctx.graph_s.is_some() || dctx.chain_s.is_some())
13412            && (!s_capturable || dctx.s_key != Some(s_key))
13413        {
13414            debug_assert!(
13415                false,
13416                "sampled draft graph parked under {:?} survived into a request outside its \
13417                 capture regime (top_k={} top_p={} min_p={} pen_on={} capturable={}): the \
13418                 in-graph draw and the verify's accept test would see different distributions",
13419                dctx.s_key, sp.top_k, sp.top_p, sp.min_p, pen_on, s_capturable,
13420            );
13421            eprintln!(
13422                "[spec] BUG: dropping a parked sampled draft graph that outlived its capture \
13423                 regime (s_key={:?}, request top_k={} top_p={} min_p={} pen_on={} \
13424                 capturable={}); drafting EAGER — the key must carry every field that shapes q",
13425                dctx.s_key, sp.top_k, sp.top_p, sp.min_p, pen_on, s_capturable,
13426            );
13427            dctx.graph_s = None;
13428            dctx.chain_s = None;
13429            dctx.s_key = None;
13430            dctx.q_slots.clear();
13431            dctx.keeper_s.clear();
13432        }
13433        // SKEY PROBE (MEMRA_SKEY_PROBE=1): the burst-entry facts the reachability question turns
13434        // on — is this request sampled, is it in a regime the sampled graph is legal in, and is
13435        // a graph PARKED from an earlier request of the same session? The launch arms below
13436        // print which chain actually ran, so the probe never restates the condition.
13437        if skey_probe() {
13438            eprintln!(
13439                "[skey] burst sampled={} pure_temp={} capturable={} temp={} top_k={} top_p={} \
13440                 min_p={} pen_on={} k={} graph_draft={} graph_s_parked={} chain_s_parked={} \
13441                 s_key_parked={:?}",
13442                sampled as u8,
13443                pure_temp as u8,
13444                s_capturable as u8,
13445                sp_temp,
13446                sp.top_k,
13447                sp.top_p,
13448                sp.min_p,
13449                pen_on as u8,
13450                k,
13451                graph_draft as u8,
13452                dctx.graph_s.is_some() as u8,
13453                dctx.chain_s.is_some() as u8,
13454                dctx.s_key,
13455            );
13456        }
13457        let t_cap = t_ent.elapsed();
13458        // PERSISTENT DRAFT KV: fill the MTP block's K/V for every prompt position from the exact
13459        // trunk hiddens collected during prime — ONE batched K/V-only pass (overwrites any
13460        // capture-warmup garbage; capture left len at 0). last_token (the init feed) needs no
13461        // fill: the first chain step processes it and appends its entry at slot prompt.len().
13462        if let Some(ph) = &prompt_h {
13463            // SESSION: rows [0..base) are the previous turns' exact fills (refresh overwrote them
13464            // with true verify hiddens) — truncate any draft overhang, fill ONLY the suffix at
13465            // global positions [base..base+tp). Fresh call: base==0, identical to before.
13466            scratch.set_len(e, base)?;
13467            // CHUNKED FILL (long-ctx OOM fix, 2026-07-05): mtp_kv_fill's transients scale with its
13468            // T (concat = T*2*n_embd*4B — 1.5GB at 40k) and its concat loop is 2*T launches. The
13469            // fill is a pure sequential append, so chunking is exact: each chunk appends its rows
13470            // at pos0=base+start with the identical per-row math. Same knob as the trunk prime.
13471            let tp = prompt.len();
13472            let fill_chunk: usize = if crate::cache::swa_ring_on() {
13473                crate::hybrid_forward::prime_chunk_tokens(tp, self.layers.len())
13474            } else {
13475                // Preserve the flag-OFF schedule byte-for-byte, including the legacy zero value
13476                // meaning one monolithic fill.
13477                std::env::var("MEMRA_PRIME_CHUNK")
13478                    .ok()
13479                    .and_then(|v| v.parse().ok())
13480                    .unwrap_or(4096)
13481            };
13482            let fill_chunk = if fill_chunk == 0 { tp } else { fill_chunk };
13483            // CUDA launch wall (same class as the trunk prime's PRIME_CHUNK_LAUNCH_CAP):
13484            // a fill call's matmuls can land on the grid.y=m dp4a family, and grid.y caps
13485            // at 65,535. This loop has no tail fold, so the raw limit is exact:
13486            // tp <= 65,535 keeps the legacy schedule (monolithic included) byte-for-byte,
13487            // and larger fills — unreachable before the trunk prime's own cap fix — chunk.
13488            let fill_chunk = fill_chunk.min(crate::hybrid_forward::CUDA_GRID_YZ_MAX);
13489            let mut start = 0usize;
13490            while start < tp {
13491                let end = (start + fill_chunk).min(tp);
13492                let tc = end - start;
13493                {
13494                    // PREDECESSOR pairing: row i gets h[i-1]; global row 0 a zeros row (the
13495                    // reference engine's initial pending-h is zeroed too); a session turn's row 0
13496                    // gets the PREVIOUS turn's last committed hidden (sess.last_h). Per chunk:
13497                    // rows start..end read h[start-1..end-1] — one dtod into a chunk buffer.
13498                    let mut phs = e.zeros(tc * n_embd)?;
13499                    let (src_lo, dst_off) = if start == 0 {
13500                        (0, n_embd)
13501                    } else {
13502                        ((start - 1) * n_embd, 0)
13503                    };
13504                    let n_copy = if start == 0 {
13505                        (tc - 1) * n_embd
13506                    } else {
13507                        tc * n_embd
13508                    };
13509                    if start == 0
13510                        && let Some((_, lh, _, _, _)) = sess_tail.as_ref()
13511                        && let Some(lh) = lh.as_ref()
13512                    {
13513                        e.copy_into(&mut phs, 0, lh, n_embd)?;
13514                    }
13515                    if n_copy > 0 {
13516                        e.copy_view_into(
13517                            &mut phs,
13518                            dst_off,
13519                            &ph.slice(src_lo..src_lo + n_copy),
13520                            n_copy,
13521                        )?;
13522                    }
13523                    self.mtp_kv_fill_all(
13524                        e,
13525                        &prompt[start..end],
13526                        &phs,
13527                        base + start,
13528                        &mut *scratch,
13529                        embd_dev,
13530                    )?;
13531                }
13532                start = end;
13533            }
13534        }
13535        // MEMRA_PROFILE_SPEC=2: profiler capture starts HERE — after the prime, so an
13536        // `nsys -c cudaProfilerApi` capture contains ONLY the round loop (draft/verify/commit).
13537        // (=1 brackets the whole call in run_spec.rs, prime included.)
13538        if std::env::var("MEMRA_PROFILE_SPEC").as_deref() == Ok("2") {
13539            unsafe extern "C" {
13540                fn cudaProfilerStart() -> i32;
13541            }
13542            unsafe {
13543                cudaProfilerStart();
13544            }
13545        }
13546        // ROUND-STREAM stage (c) 4 (MEMRA_SPEC_STREAM=1, experimental): pre-issued M-round
13547        // bursts with ZERO per-round host readbacks — the accept/seed/rollback/ring kernels
13548        // consume each other's device outputs; the host drains the ring every M rounds. v1
13549        // constraints: greedy, !spec_replay, single-shot, batched-linear layers, no refresh
13550        // fills (acceptance effect A/B-arbitrated), enters from round 1 (pending guaranteed).
13551        // NOTE: not gated on the caller's graph_draft (its trunk_dense conjunct turns the 35B
13552        // MoE off) — the stream capture encloses ONLY the dense MTP head; the head-dense /
13553        // full-prec / k gates are re-derived here and a failed capture degrades to stream-off.
13554        let stream_on = crate::spec::spec_stream()
13555            && !sampled
13556            && !spec_replay
13557            && self.mtp_extra.is_empty()
13558            && constraint.is_none()
13559            && !session_mode
13560            && embd_gpu.is_some()
13561            && !crate::model::full_prec_enabled()
13562            && k + 2 < 96;
13563        let mut stream_graph: Option<cudarc::driver::CudaGraph> = None;
13564        let mut g_tokp2k = e.alloc_u32_zeroed(2 * k.max(1))?;
13565        if stream_on {
13566            let cap = e.capture_graph(|e| {
13567                for j in 0..k.max(1) {
13568                    self.mtp_head_forward_cap(
13569                        e,
13570                        mtp,
13571                        &mut dctx.g_tok,
13572                        &mut dctx.g_pos,
13573                        &mut dctx.g_seed,
13574                        &mut dctx.g_p,
13575                        &mut *scratch,
13576                        0,
13577                        true,
13578                        true,
13579                        embd_gpu.expect("round stream requires resident embedding"),
13580                        embd_qt,
13581                        embd_rb,
13582                        d_vocab,
13583                        None,
13584                        Some((&mut g_tokp2k, j, d2t_dev.as_ref())),
13585                        None, // round-stream requires constraint.is_none() (see stream_on)
13586                    )?;
13587                }
13588                Ok(())
13589            });
13590            match cap {
13591                Ok(g) => {
13592                    scratch.set_len(e, 0)?;
13593                    stream_graph = Some(g);
13594                }
13595                Err(err) => {
13596                    scratch.set_len(e, 0)?;
13597                    if debug_spec {
13598                        eprintln!("[spec] stream-graph capture failed ({err}); stream off");
13599                    }
13600                }
13601            }
13602        }
13603        let stream_active = stream_on && stream_graph.is_some();
13604        if debug_spec {
13605            eprintln!(
13606                "[spec] stream_on={stream_on} env={} samp={sampled} dg={} captured={} active={stream_active} session={session_mode} replay={spec_replay}",
13607                crate::spec::spec_stream(),
13608                dctx.graph.is_some(),
13609                stream_graph.is_some()
13610            );
13611        }
13612        let t_v_s = k + 1;
13613        // ROUND-STREAM buffers + ptr tables now live in the model-generic round_stream
13614        // module (extracted 2026-07-12; the gemma burst reuses them).
13615        let sb = crate::round_stream::StreamBufs::new(e, k, crate::spec::spec_stream_m())?;
13616        let crate::round_stream::StreamBufs {
13617            mut vtok_d,
13618            mut brk_d,
13619            mut pend_d,
13620            last_pred_d,
13621            mut pos_ctr,
13622            mut pos_start_d,
13623            mut ring_d,
13624            acc_d: mut stream_acc,
13625            m_rounds,
13626            k: _,
13627        } = sb;
13628        let stream_ptrs: Option<CudaSlice<u64>> = if stream_active {
13629            Some(crate::round_stream::kv_len_ptr_table(
13630                e,
13631                cache,
13632                Some(&pos_ctr),
13633            )?)
13634        } else {
13635            None
13636        };
13637
13638        let t_fill = t_ent.elapsed();
13639        let mut round = 0usize;
13640        // ADAPTIVE DRAFT LENGTH (MEMRA_SPEC_ADAPT=1, opt-in — the gemma_spec accepted-run law,
13641        // ported 2026-08-01): next round's draft depth = last round's accepted run + 1, clamped
13642        // to [floor(pos), k_cap] — a miss shrinks the next draft to the miss point + 1,
13643        // full-accept streaks re-deepen one step per round. NOT the 2026-07-07 acceptance-EMA
13644        // (that arm measured an HONEST LOSS to static per-class optima — 115.0/85.8/73.4 vs
13645        // 121.6/92.7/75.6, EMA lag — and was removed 2026-07-08; rig5090.jsonl has the record).
13646        // The gemma law has no lag class: it reacts within one round, and was worth +7-20% on
13647        // the gemma cells at unchanged exactness (2026-07-10 flip; floor sweep 2026-07-25;
13648        // position key 2026-07-26). Signal = n_acc from the round's EXISTING accept readback —
13649        // zero new syncs; the draft graph is a SINGLE-STEP capture replayed per drafted token,
13650        // so a per-round depth needs no re-capture (unlike gemma's whole-chain graphs). qwen's
13651        // in-round p-min cut already shortens chains mid-round, so gemma's one-round-late p-min
13652        // fold into kc is unnecessary here — the accepted-run law sees the cut via n_acc.
13653        // Exactness is the verify's job at ANY depth (same contract as p-min variable rounds).
13654        // DEFAULT OFF on the qwen path until its cells gate a flip (gemma's is default-on).
13655        // MEASURED 2026-08-01 (H100 GPU-3, interleaved x3, NGEN=256, same-invocation plain
13656        // denominators; research/qwen-adaptive-k-20260801/): REFUTED on the tuned qwen configs.
13657        // q27 K=3+HPOST+PMIN=0.3: short +0.8% (noise; law ~idles, len_hist identical), board
13658        // -1.9%, agentic -0.5%; board PMIN=0 -2.1% (not p-min shadowing — the law itself);
13659        // floor=1 -2.8% (gemma's floor-collapse, reproduced). q35 K=2 board: -6.4% (52/136
13660        // rounds shrink to depth 1; no depth to reclaim at K=2). The gemma direction DOES
13661        // appear at untuned depth-K — q27 K=6 floor=4 +1.5% over fixed K=6 — but stays -3.7%
13662        // below fixed K=3: same verdict class as the retired EMA arm (honest loss to static
13663        // per-class optima). Acceptance-rate rises under the law while tokens/round falls —
13664        // it buys accept-% by adding rounds, and a round's fixed draft+verify cost wins.
13665        // K=1..8 self-consistency PASS both models with the law ON (exactness held).
13666        let adapt = std::env::var("MEMRA_SPEC_ADAPT").as_deref() == Ok("1");
13667        // floor: per-model default keyed on n_embd (gemma's tiering — models with an expensive
13668        // verify keep deep drafts after a miss); MEMRA_SPEC_ADAPT_FLOOR pins it everywhere.
13669        let adapt_floor_env: Option<usize> = std::env::var("MEMRA_SPEC_ADAPT_FLOOR")
13670            .ok()
13671            .and_then(|v| v.parse().ok());
13672        let adapt_floor_default: usize = if self.cfg.n_embd as usize >= 3500 {
13673            4
13674        } else if self.cfg.n_embd as usize >= 2500 {
13675            2
13676        } else {
13677            1
13678        };
13679        let adapt_floor: usize = adapt_floor_env.unwrap_or(adapt_floor_default);
13680        // position key: past floor_ctx a HIGH floor (>=4) relaxes to 1 — forced-deep drafts
13681        // turn net-negative at depth (gemma 31B d1736 evidence); MEMRA_SPEC_FLOOR_CTX moves
13682        // the boundary, an explicit MEMRA_SPEC_ADAPT_FLOOR pins the floor everywhere.
13683        let floor_ctx: usize = std::env::var("MEMRA_SPEC_FLOOR_CTX")
13684            .ok()
13685            .and_then(|v| v.parse().ok())
13686            .unwrap_or(1024);
13687        let floor_at = |pos: usize| -> usize {
13688            if adapt_floor_env.is_some() || pos < floor_ctx {
13689                adapt_floor
13690            } else if adapt_floor >= 4 {
13691                1
13692            } else {
13693                adapt_floor
13694            }
13695        };
13696        // cap: MEMRA_SPEC_CAPMAX (gemma semantics, default 7). Binds only under adapt — the
13697        // fixed-K default path is untouched by this whole block.
13698        let cap_max: usize = std::env::var("MEMRA_SPEC_CAPMAX")
13699            .ok()
13700            .and_then(|v| v.parse().ok())
13701            .unwrap_or(7);
13702        let k_cap = k.min(cap_max).max(1);
13703        let mut kc = k_cap;
13704        let mut opti_fork: Option<OptiForkState> = None;
13705        let mut _opti_walk: Option<crate::pp::PpWalkLease> = None;
13706        let mut _opti_walk_borrow: Option<crate::pp::PpWalkBorrowGuard> = None;
13707        let mut fork_snapshot: Option<crate::cache::CacheSnapshot> = None;
13708        if fork_mode != OptiForkGateMode::Disabled {
13709            let fence = crate::pp::pp_cuts(self.layers.len());
13710            let refusal = if !session_mode {
13711                Some("not-session")
13712            } else if k != 1 || adapt {
13713                Some("requires-fixed-k1")
13714            } else if sampled || constraint.is_some() || spec_replay {
13715                Some("sampled-constrained-or-replay")
13716            } else if pipe.is_some() {
13717                Some("two-session-pipeline")
13718            } else if !spec_devacc() {
13719                Some("requires-device-accept")
13720            } else if stream_active || crate::spec::spec_stream() {
13721                Some("round-stream")
13722            } else if !self.mtp_extra.is_empty() {
13723                Some("multi-head-mtp")
13724            } else if crate::cache::swa_ring_on() || cache.has_swa_ring() {
13725                Some("swa-ring")
13726            } else if crate::pp::pp_host_bounce_active() {
13727                Some("host-bounce")
13728            } else if fork_mode == OptiForkGateMode::Controller
13729                && cache.recur.iter().any(Option::is_some)
13730            {
13731                Some("controller-requires-zero-recurrent-state")
13732            } else if fence.as_ref().is_none_or(|f| f.len() != 3) {
13733                Some("requires-pp2")
13734            } else {
13735                None
13736            };
13737            if let Some(reason) = refusal {
13738                OPTI_FORK_REFUSALS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
13739                eprintln!("[opti-fork] refused reason={reason}");
13740            } else {
13741                let fence = fence.expect("validated PP-2 fence");
13742                let rt = crate::pp::PpNRt::get(e)?;
13743                let primary_stage0 = rt.engine(0, e).ctx().ordinal() == e.ctx().ordinal();
13744                let primary_stage1 = rt.engine(1, e).ctx().ordinal() == e.ctx().ordinal();
13745                let primary_supported =
13746                    primary_stage0 || (fork_mode == OptiForkGateMode::Controller && primary_stage1);
13747                if !rt.cross_device() || !primary_supported {
13748                    OPTI_FORK_REFUSALS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
13749                    eprintln!("[opti-fork] refused reason=requires-supported-primary-cross-device");
13750                } else {
13751                    // The optimistic controller can keep two boundary tickets in flight. Give
13752                    // every nested verify an explicit borrow of one whole-walk generation; no
13753                    // `pp_pipe` boolean is allowed to bypass ownership on its own.
13754                    let walk = rt.acquire_walk("opti_fork_coordinator")?;
13755                    let permit = rt.walk_permit(&walk, "opti_fork_coordinator")?;
13756                    let borrow = rt.borrow_walk(&permit, "opti_fork_coordinator")?;
13757                    // Both recurrent snapshots and both seed generations are allocated before
13758                    // the first fork, each through its owning PP stage. Allocation failure
13759                    // therefore happens before any optimistic state mutation can occur.
13760                    let current_snapshot = opti_snapshot_stage_owned(e, cache, rt, &fence)?;
13761                    let alternate_snapshot = opti_snapshot_stage_owned(e, cache, rt, &fence)?;
13762                    let fork = OptiForkState::new(
13763                        e,
13764                        cache,
13765                        fork_mode,
13766                        alternate_snapshot,
13767                        &h_seed_buf,
13768                        &fill_prev,
13769                        rt,
13770                        fence[1],
13771                        self.layers.len(),
13772                    )?;
13773                    eprintln!(
13774                        "[opti-fork] armed mode={fork_mode:?} snapshots=2 seeds=2 split={} \
13775                         payload_dev0={} payload_dev1={} q_threshold={:.3}",
13776                        fence[1],
13777                        fork.logical_payload_bytes[0],
13778                        fork.logical_payload_bytes[1],
13779                        fork.controller.map_or(0.0, |policy| policy.threshold),
13780                    );
13781                    fork_snapshot = Some(current_snapshot);
13782                    opti_fork = Some(fork);
13783                    _opti_walk = Some(walk);
13784                    _opti_walk_borrow = Some(borrow);
13785                }
13786            }
13787        }
13788        // Persistent snapshot buffers are allocated once and refreshed in place. The fork arm
13789        // uses stage-owned snapshots; refused/disabled arms retain the existing generic helper.
13790        let mut snap = match fork_snapshot {
13791            Some(snapshot) => snapshot,
13792            None => cache.snapshot(e)?,
13793        };
13794        let mut carried_opti: Option<OptiControllerTicket> = None;
13795        // ROUND-STREAM stage (b) 3a: device table of per-layer kvl.len_d pointers (stable — the
13796        // cache never reallocates len_d; see cache.rs "stable pointer" note). 0 = no KV layer.
13797        let kv_len_ptrs: Option<CudaSlice<u64>> = if spec_devacc() && !spec_replay {
13798            Some(crate::round_stream::kv_len_ptr_table(e, cache, None)?)
13799        } else {
13800            None
13801        };
13802        // BONUS FOLD (2026-07-04): after a FULL accept the bonus token is NOT committed with a
13803        // separate T=1 trunk pass (a full weight read per round). It stays PENDING and rides as
13804        // column 0 of the NEXT round's verify batch. Under predecessor pairing the next chain
13805        // seeds from the bonus's predecessor's TRUE verify hidden (free — no extra
13806        // pass of any kind). Verify still
13807        // checks every emitted token against the target -> exactness holds by construction; only
13808        // DRAFT QUALITY can shift, which the acceptance numbers arbitrate.
13809        // bonus emitted but not yet committed to cache. A carried pending (see SpecSession::
13810        // pending_tok) enters round 0 directly — the burst boundary becomes a plain round edge.
13811        let mut pending: Option<u32> = carried_pending;
13812        // MEMRA_SPEC_PHASE=1: per-round wall decomposition (draft / verify / accept+commit) —
13813        // no tracing, no extra syncs (each phase is naturally sync-bounded: draft readbacks,
13814        // the verify accept readback). Printed once at loop end via spec-stats.
13815        let anatomy_on = std::env::var("MEMRA_SPEC_PP_ANATOMY").as_deref() == Ok("1");
13816        let phase_on = anatomy_on || std::env::var("MEMRA_SPEC_PHASE").as_deref() == Ok("1");
13817        // MEMRA_SPEC_PHASE_SYNC=1 — reads the phase split correctly, and proves it. `ph_mark` is a
13818        // bare Instant, so `verify-issue` is the host QUEUEING the walk (the GPU is already running
13819        // under it) and `verify-wait` is only the residual drain at the accept readback: one
13820        // overlapped interval cut at the first blocking call, NOT "GPU time" beside "host time".
13821        // Syncing right after the walk is issued moves the whole GPU wall into `verify-issue`. If
13822        // the walk's GPU total is really issue+wait, then with this on verify-issue jumps to that
13823        // sum, verify-wait collapses to the readback alone, and the ROUND WALL DOES NOT MOVE —
13824        // which is what says the queueing time was hidden and is not a target. Diagnostic only.
13825        let phase_sync = std::env::var("MEMRA_SPEC_PHASE_SYNC").as_deref() == Ok("1");
13826        // DRAFT-MASK receipt (lane/draft-mask): speculative-clone wall + rounds, printed with
13827        // spec-stats. The clone is the one cost the design adds per round — measured, not assumed.
13828        let (mut dm_clone_ns, mut dm_rounds) = (0u128, 0usize);
13829        // grammar-truncation counters: how many rounds the verify-side cut fired and how many
13830        // already-verified tokens it threw away. THIS is the quantity draft masking targets.
13831        let (mut dm_cuts, mut dm_cut_tokens) = (0usize, 0usize);
13832        let (mut ph_draft, mut ph_verify, mut ph_rest) = (0f64, 0f64, 0f64);
13833        let mut ph_wait = 0f64;
13834        let mut ph_commit = 0f64;
13835        let mut ph_t = std::time::Instant::now();
13836        let mut ph_mark = |acc: &mut f64, on: bool| {
13837            if on {
13838                let now = std::time::Instant::now();
13839                *acc += (now - ph_t).as_secs_f64();
13840                ph_t = now;
13841            }
13842        };
13843        // MTP-ROUTE VERIFY GRAPHS (`MEMRA_SPEC_VERIFY_GRAPH`, see the flag doc): the
13844        // model-owned capture pool, locked for the whole burst exactly as the dspark serve
13845        // arm holds it — the slab stash is live verify -> commit inside a round, and the
13846        // worker drives rounds from one scheduler thread. PERSISTENT across generations on
13847        // the model (rebuilding per call re-captures the pool per prompt, which is the
13848        // measured way to lose more than the launches cost); the captured bodies are
13849        // cache-independent, every state read going through per-round refreshed pointer
13850        // tables. None = the eager walk, byte-identical.
13851        //
13852        // Never armed together with ROUND-STREAM: the tparallel verify refuses that pair
13853        // loudly, and `stream_active` owns the burst arm above, so the door stays shut
13854        // whenever the stream is live rather than relying on that refusal.
13855        // The lock is taken ONLY when the door is armed: with the flag off this whole block
13856        // is inert, so the default path cannot serialize two spec generations behind a mutex
13857        // it never reads.
13858        let vg_armed =
13859            crate::spec::spec_verify_graph_env().unwrap_or_else(|| self.vgraph_family_default());
13860        let mut vg_guard = if vg_armed && !stream_active {
13861            let mut g = self.dspark_vgraphs.lock().unwrap();
13862            if g.is_none() {
13863                // Size by the WIDEST verify this run can present, which is k+1 and NOT
13864                // k_cap+1: the sampled arm's own window is `t_v_s = k + 1`, so a pool built
13865                // from a smaller adaptive cap gets sliced past its stash rows (a `slice_mut`
13866                // panic in the sampled ON arm, measured before this line said k+1).
13867                let vt_cap = (k.max(k_cap) + 1).max(2);
13868                *g = DsparkVerifyGraphs::new(e, cache, vt_cap, n_embd)?;
13869                if g.is_some() {
13870                    // Engagement receipt (the dead-arm lesson): prove the door is LIVE rather
13871                    // than trusting that a flag set means a pool built.
13872                    eprintln!("[spec-vg] MTP verify-graph pool ENGAGED (vt_cap={vt_cap})");
13873                } else {
13874                    eprintln!(
13875                        "[spec-vg] MTP verify-graph pool declined (no linear layers, \
13876                         non-uniform state, or vt_cap < 2) — eager walk"
13877                    );
13878                }
13879            }
13880            Some(g)
13881        } else {
13882            None
13883        };
13884        // Capacity fail-safe: a round wider than the pool was built for must take the eager
13885        // walk, not slice the stash past its rows. The sizing above already covers every
13886        // round this run can present; this keeps a future caller (or a k that grows behind
13887        // the pool's back) on the byte-identical fallback instead of a panic.
13888        let vg_t_cap = vg_guard
13889            .as_ref()
13890            .and_then(|g| g.as_ref())
13891            .map(|g| g.t_capacity())
13892            .unwrap_or(0);
13893        if let Some(p) = pipe {
13894            p.setup_end();
13895        }
13896        drop(pipe_setup_walk);
13897        let mut graph_guard_noted = false;
13898        while keep_going && out.len() < max_new {
13899            // GRAPH-LAUNCH HEADROOM GUARD (see GRAPH_LAUNCH_MIN_FREE): below the floor,
13900            // every captured-graph arm in this round yields to its byte-identical eager
13901            // twin instead of feeding cuGraphLaunch a card it segfaults on.
13902            let graph_round_ok = graph_launch_headroom_ok(e);
13903            if !graph_round_ok && !graph_guard_noted {
13904                graph_guard_noted = true;
13905                eprintln!(
13906                    "[spec] graph replay suspended: driver free below the {}MB launch floor \
13907                     (eager arms serve; cuGraphLaunch segfaults into an exhausted card)",
13908                    GRAPH_LAUNCH_MIN_FREE / (1 << 20)
13909                );
13910            }
13911            // MEMRA_SPEC_ROUND_PROF=1: wall of the WHOLE round against the pieces we already
13912            // instrument. Needed because the parts do not add up: the draft step measures 1.27 ms
13913            // ([spec-anatomy] glue 92 / attn 280 / ffn 222 / head 670 us) and the t=2 verify walk
13914            // 25.6 ms ([tcol-prof] attn 10.1 + ffn 15.3), yet a K=1 round takes 177 ms on the
13915            // step37 TP2 stack. This prints where the other ~150 ms lives.
13916            let round_prof = ROUND_PROF
13917                .get_or_init(|| std::env::var("MEMRA_SPEC_ROUND_PROF").as_deref() == Ok("1"));
13918            let round_t0 = round_prof.then(std::time::Instant::now);
13919            // ROUND-STREAM BURST: from round 1 (pending guaranteed by every non-replay arm),
13920            // issue M rounds with zero readbacks, then drain the ring + reconcile mirrors.
13921            if let (true, Some(sg), Some(ptrs)) = (
13922                stream_active && round >= 1 && pending.is_some() && graph_round_ok,
13923                &stream_graph,
13924                &stream_ptrs,
13925            ) {
13926                if debug_spec {
13927                    static ONCE: std::sync::Once = std::sync::Once::new();
13928                    ONCE.call_once(|| {
13929                        eprintln!("[memra] ROUND-STREAM burst engaged (M={m_rounds} k={k})")
13930                    });
13931                }
13932                e.set_i32_one(&mut pos_ctr, cache.pos as i32)?;
13933                e.set_u32_one(&mut pend_d, pending.unwrap())?;
13934                e.set_u32_one(&mut ring_d, 0)?; // ring count = 0 (writes element 0)
13935                for _mi in 0..m_rounds {
13936                    e.i32_copy_add(&pos_ctr, &mut pos_start_d, 0)?;
13937                    cache.snapshot_into(e, &mut snap)?; // device D2Ds, stream-ordered
13938                    e.i32_copy_add(&pos_ctr, &mut scratch.kv.len_d, 0)?; // draft-KV rollback
13939                    e.i32_copy_add(&pos_ctr, &mut dctx.g_pos, 1)?; // rope pos = pos + base
13940                    e.u32_copy(&pend_d, &mut dctx.g_tok)?;
13941                    e.copy_into(&mut dctx.g_seed, 0, &h_seed_buf, n_embd)?;
13942                    sg.launch()?;
13943                    e.spec_assemble_verify(
13944                        &g_tokp2k,
13945                        &pend_d,
13946                        d2t_dev.as_ref(),
13947                        &mut vtok_d,
13948                        &mut brk_d,
13949                        p_min,
13950                        k,
13951                        pmin0,
13952                    )?;
13953                    let mut ck = VerifyCkpt::new(self.layers.len());
13954                    let dummy = vec![0u32; t_v_s];
13955                    let (tl_d, vx) = self.decode_step_t_core_stream(
13956                        e,
13957                        &dummy,
13958                        0,
13959                        &mut *cache,
13960                        embd_dev,
13961                        Some(&mut ck),
13962                        Some((&vtok_d, &pos_ctr)),
13963                        None,
13964                        None,
13965                        None,
13966                    )?;
13967                    for j in 0..t_v_s {
13968                        e.argmax_token_device_col(&tl_d, j, n_vocab, &mut preds_d, j)?;
13969                    }
13970                    e.spec_accept_greedy_dc(
13971                        &preds_d,
13972                        &vtok_d,
13973                        &last_pred_d,
13974                        &brk_d,
13975                        &mut stream_acc,
13976                    )?;
13977                    e.spec_seed_gather(&vx, &fill_prev, &stream_acc, &mut h_seed_buf, 1, n_embd)?;
13978                    e.copy_into(&mut fill_prev, 0, &h_seed_buf, n_embd)?;
13979                    self.commit_verified_prefix_stream(
13980                        e,
13981                        &mut *cache,
13982                        &snap,
13983                        &ck,
13984                        &stream_acc,
13985                        1,
13986                        t_v_s,
13987                    )?;
13988                    e.spec_rollback_stream(
13989                        ptrs,
13990                        &pos_start_d,
13991                        &stream_acc,
13992                        1,
13993                        self.layers.len() + 1,
13994                    )?;
13995                    e.spec_ring_commit(&vtok_d, &stream_acc, &brk_d, &mut ring_d, &mut pend_d)?;
13996                }
13997                e.stream().synchronize()?;
13998                let ring_h = e.dtoh_u32(&ring_d)?;
13999                let cnt = ring_h[0] as usize;
14000                for i in 0..cnt {
14001                    if out.len() < max_new {
14002                        out.push(ring_h[1 + i]);
14003                    }
14004                }
14005                let pos_h = e.dtoh_i32(&pos_ctr)?[0] as usize;
14006                for il in 0..self.layers.len() {
14007                    if let Some(kvl) = cache.kv[il].as_mut() {
14008                        kvl.len = pos_h;
14009                    }
14010                }
14011                cache.pos = pos_h;
14012                scratch.kv.len = pos_h;
14013                pending = Some(ring_h[cnt]); // last drained token = the live bonus
14014                last_token = ring_h[cnt];
14015                total_drafted += k * m_rounds; // upper bound (p-min breaks uncounted)
14016                total_accepted += cnt.saturating_sub(m_rounds);
14017                if let Some(t) = sess_telem {
14018                    // totals only — the burst's per-round accept counts stayed on device
14019                    // (that is the point of the round-stream arm). pos_* untouched.
14020                    t.record_totals(m_rounds, k * m_rounds, cnt.saturating_sub(m_rounds));
14021                }
14022                round += m_rounds;
14023                // sse-cadence: the drained ring is committed — flush it at burst-drain cadence.
14024                keep_going = flush_commit(&mut on_commit, &out, &mut flushed);
14025                continue;
14026            }
14027            let pipe_draft = match pipe {
14028                Some(p) => Some(p.draft_begin(round)?),
14029                None => None,
14030            };
14031            let pos = cache.pos; // #tokens committed (EXCLUDES a pending bonus)
14032            let mut current_opti = carried_opti.take();
14033            let mut fork_generation = if current_opti.is_none() && pending.is_some() {
14034                match opti_fork.as_mut() {
14035                    Some(fork) if fork.mode.is_forced() => Some(fork.reserve(&mut snap)?),
14036                    None => None,
14037                    Some(_) => None,
14038                }
14039            } else {
14040                None
14041            };
14042            if current_opti.is_none() {
14043                if let Some(fork) = opti_fork.as_ref() {
14044                    opti_snapshot_stage_owned_into(e, cache, fork.rt, &fork.fence, &mut snap)?;
14045                } else {
14046                    cache.snapshot_into(e, &mut snap)?;
14047                }
14048            } else if snap.pos != pos {
14049                return Err(format!(
14050                    "optipipe carried snapshot pos {} != current pos {pos}",
14051                    snap.pos
14052                )
14053                .into());
14054            } // §C: snapshot BEFORE draft+verify (already retained for a carried successor)
14055            ph_mark(&mut ph_rest, phase_on);
14056
14057            // --- 1. DRAFT k tokens with the NextN head (autoregressive, T=1 each) ---
14058            // p-min semantics (both paths): stop the chain early when the head's confidence in
14059            // its own pick drops below p_min — the just-drafted token is DISCARDED, but its
14060            // scratch append stands (identical to the eager chain's ordering). j==0 always drafts.
14061            let base0 = if pending.is_some() { 1usize } else { 0usize };
14062            // fixed draft length by default; MEMRA_SPEC_ADAPT=1 drafts at last round's
14063            // accepted run + 1 (the gemma law — see the setup block above the loop).
14064            let k_this = if adapt { kc } else { k };
14065            let mut draft: Vec<u32> = Vec::with_capacity(k);
14066            let mut draft_idx: Vec<u32> = Vec::with_capacity(k); // trimmed-vocab ids (== draft when untrimmed)
14067            let mut controller_draft_prob: Option<f32> = None;
14068            let mut controller_eager_state: Option<(u32, CudaSlice<f32>)> = None;
14069            if let Some(ticket) = current_opti.as_mut() {
14070                let carried_pending = pending.ok_or("optipipe carried successor lost pending")?;
14071                if ticket.verify_tokens[0] != carried_pending {
14072                    return Err(format!(
14073                        "optipipe carried pending mismatch: ticket={} live={carried_pending}",
14074                        ticket.verify_tokens[0],
14075                    )
14076                    .into());
14077                }
14078                draft.push(ticket.verify_tokens[1]);
14079                controller_draft_prob = Some(ticket.draft_prob);
14080                controller_eager_state = ticket
14081                    .take_eager_seed()
14082                    .map(|seed| (ticket.verify_tokens[1], seed));
14083            } else {
14084                // Round-start draft-KV sync (BOTH paths). Persistent: truncate/align to the committed
14085                // history — slots 0..P hold entries for the tokens before last_token@P (P = pos +
14086                // base0 - 1); this single set_len IS the draft-side rollback (drops last round's
14087                // rejected drafts and p-min extras via the len mechanism).
14088                scratch.set_len(e, pos + base0 - 1)?;
14089                // dcw door: a captured chain appends k_this device-counter rows (plus the
14090                // pseudo-seed replay) with no host intervention; any ring rebase those appends
14091                // could need happens HERE, host-side, before the replays. The eager arm keeps
14092                // its own per-step prepare, so this is graph-path-only work.
14093                if step35_draft_dcw_on()
14094                    && (dctx.graph.is_some()
14095                        || dctx.graph_s.is_some()
14096                        || dctx.chain.is_some()
14097                        || dctx.chain_s.is_some())
14098                {
14099                    scratch.ensure_dcw_headroom(e, k_this + 2)?;
14100                }
14101                if pen_on {
14102                    // PEN_WINDOW_MAX also bounds the per-round upload and the O(n_hist^2)
14103                    // device dedup: the serve window is already PEN_WINDOW_MAX, and this
14104                    // defensive min also bounds non-server callers.
14105                    let win = sp.penalty_last_n.min(PEN_WINDOW_MAX);
14106                    let w0 = pen_hist.len().saturating_sub(win);
14107                    pen_hist_d = Some(e.htod_u32_v(&pen_hist[w0..])?);
14108                }
14109                if sampled {
14110                    draft_logits.clear();
14111                    draft_stats.clear();
14112                }
14113                // DRAFT-SIDE GRAMMAR MASK: clone the committed grammar state ONCE per round; each
14114                // position's mask is computed on that clone and advanced by the PROPOSED token. The
14115                // real state moves only on emission (verify's job), so the emitted stream is
14116                // unchanged — the mask only removes tokens the verify would have truncated anyway.
14117                let mut dmask_live = dmask_on;
14118                if dmask_live {
14119                    let t_c = std::time::Instant::now();
14120                    constraint
14121                        .as_deref_mut()
14122                        .unwrap()
14123                        .draft_begin()
14124                        .map_err(|e2| format!("constraint: {e2}"))?;
14125                    dm_clone_ns += t_c.elapsed().as_nanos();
14126                    dm_rounds += 1;
14127                }
14128                if let (false, Some(cg)) = (sampled || pen_on || !graph_round_ok, &dctx.chain) {
14129                    // GREEDY CHAIN GRAPH (lane/step37-draft-graph-serving-20260830): the
14130                    // eager multi-head chain's EXACT launch order — step j rewinds head
14131                    // (j % heads)'s plane to the committed length and replays rows 0..=j —
14132                    // with each row's whole head-forward as ONE graph launch. The chain
14133                    // POLICY (head choice, prefix length, stored-seed feed) is host-side,
14134                    // identical to `mtp_chain_forward_dev`, so graph-vs-eager drafts are
14135                    // bit-identical by construction (same launcher, same bucket — the dcw
14136                    // parity contract). Interior rows launch the head-less graph: their
14137                    // logits are dead in the eager chain too, so the consumed bytes match.
14138                    let heads_n = self.mtp_head_count();
14139                    let committed = pos + base0 - 1;
14140                    let mut chain_tokens: Vec<u32> = vec![last_token];
14141                    let mut chain_seed_bufs: Vec<CudaSlice<f32>> = vec![e.clone_dtod(&h_seed_buf)?];
14142                    for j in 0..k_this {
14143                        let index = mtp_chain_head_index(j, heads_n);
14144                        if debug_spec {
14145                            eprintln!(
14146                                "[mtp-chain-step] round={round} j={j} head={index} \
14147                                 replay_rows={} arm=graph",
14148                                chain_tokens.len(),
14149                            );
14150                        }
14151                        scratch.set_plane_len(e, index, committed)?;
14152                        e.set_i32_one(&mut dctx.g_pos, (committed + 1) as i32)?;
14153                        for row in 0..=j {
14154                            e.set_u32_one(&mut dctx.g_tok, chain_tokens[row])?;
14155                            e.copy_into(&mut dctx.g_seed, 0, &chain_seed_bufs[row], n_embd)?;
14156                            if row < j {
14157                                cg.interior[index].launch()?;
14158                            } else {
14159                                // per-position mask upload before the LAST row only — the
14160                                // eager chain applies the mask on is_last exactly the same.
14161                                if dmask_live
14162                                    && !upload_draft_mask(
14163                                        e,
14164                                        constraint.as_deref_mut().unwrap(),
14165                                        &mut dctx.g_dmask,
14166                                        mtp.d2t.as_ref(),
14167                                        d_vocab,
14168                                        dmask_words,
14169                                    )?
14170                                {
14171                                    e.htod_u32_into(
14172                                        &mut dctx.g_dmask,
14173                                        &vec![u32::MAX; dmask_words],
14174                                    )?;
14175                                    dmask_live = false;
14176                                }
14177                                cg.last[index].launch()?;
14178                            }
14179                            // host mirror (len_d advanced in-graph by the dcw append)
14180                            scratch.plane_mut(index).0.len += 1;
14181                        }
14182                        let idx = e.dtoh_u32_one(&dctx.g_tok)?;
14183                        // #87 SENTINEL TRAP (see the single-head graph arm below).
14184                        if (idx as usize) >= d_vocab {
14185                            let seed_h = e.dtoh(&dctx.g_seed)?;
14186                            let seed_nan = seed_h.iter().filter(|v| v.is_nan()).count();
14187                            return Err(format!(
14188                                "draft(chain-graph) argmax sentinel 0x{idx:08x} >= d_vocab \
14189                             {d_vocab} at round {round} j={j} head={index} pos={pos}: \
14190                             head-out NaN {seed_nan}/{n_embd} — refusing to dereference \
14191                             the embed row (#87 trap)"
14192                            )
14193                            .into());
14194                        }
14195                        // multi-head MTP forbids a trimmed head (validated at entry), so the
14196                        // draft index IS the target id; keep the map for uniformity.
14197                        let d = match &mtp.d2t {
14198                            Some(map) => map[idx as usize],
14199                            None => idx,
14200                        };
14201                        let draft_p = if p_min > 0.0 {
14202                            Some(e.dtoh(&dctx.g_p)?[0])
14203                        } else {
14204                            None
14205                        };
14206                        if j == 0 {
14207                            controller_draft_prob = draft_p;
14208                        }
14209                        if let Some(p) = draft_p.filter(|_| p_min > 0.0)
14210                            && p < p_min
14211                            && (j > 0 || (pmin0 && base0 == 1))
14212                        {
14213                            break;
14214                        }
14215                        draft.push(d);
14216                        chain_tokens.push(d);
14217                        // step j's h_nextn: the last-row graph self-fed it into g_seed —
14218                        // snapshot it as the chain history seed for row j+1 (stream-ordered
14219                        // after the launch, exactly the eager chain's chain_seeds push).
14220                        chain_seed_bufs.push(e.clone_dtod(&dctx.g_seed)?);
14221                        // speculative grammar advance (see the single-head graph arm).
14222                        if dmask_live
14223                            && !constraint
14224                                .as_deref_mut()
14225                                .unwrap()
14226                                .draft_advance(d)
14227                                .map_err(|e2| format!("constraint: {e2}"))?
14228                        {
14229                            e.htod_u32_into(&mut dctx.g_dmask, &vec![u32::MAX; dmask_words])?;
14230                            break;
14231                        }
14232                    }
14233                } else if let (true, Some(cg)) = (
14234                    sampled && s_capturable && dctx.s_key == Some(s_key) && graph_round_ok,
14235                    &dctx.chain_s,
14236                ) {
14237                    if skey_probe() {
14238                        eprintln!(
14239                            "[skey] chain=graph_chain_s round={round} capturable={} top_k={} \
14240                             top_p={} min_p={} s_key_parked={:?}",
14241                            s_capturable as u8, sp.top_k, sp.top_p, sp.min_p, dctx.s_key,
14242                        );
14243                    }
14244                    // SAMPLED CHAIN GRAPH: the greedy chain arm's launch order with the
14245                    // sampled last-row graphs — in-graph counter bump + (filtered) gumbel
14246                    // draw + argmax; q retained per step into q_slots exactly like the
14247                    // single-head sampled graph arm. Counter continuity: g_ctr host-seeded
14248                    // to sctr-1 once per ROUND; each step's last-row graph bumps it BEFORE
14249                    // the perturb, so step j consumes counter sctr+j — the eager Philox
14250                    // stream (interior rows never draw, never bump).
14251                    let heads_n = self.mtp_head_count();
14252                    let committed = pos + base0 - 1;
14253                    let filtered_stats_in_graph = s_key.filtered();
14254                    let mut chain_tokens: Vec<u32> = vec![last_token];
14255                    let mut chain_seed_bufs: Vec<CudaSlice<f32>> = vec![e.clone_dtod(&h_seed_buf)?];
14256                    e.set_u32_one(&mut dctx.g_ctr, sctr.wrapping_sub(1))?;
14257                    for j in 0..k_this {
14258                        let index = mtp_chain_head_index(j, heads_n);
14259                        if debug_spec {
14260                            eprintln!(
14261                                "[mtp-chain-step] round={round} j={j} head={index} \
14262                                 replay_rows={} arm=graph_s",
14263                                chain_tokens.len(),
14264                            );
14265                        }
14266                        scratch.set_plane_len(e, index, committed)?;
14267                        e.set_i32_one(&mut dctx.g_pos, (committed + 1) as i32)?;
14268                        for row in 0..=j {
14269                            e.set_u32_one(&mut dctx.g_tok, chain_tokens[row])?;
14270                            e.copy_into(&mut dctx.g_seed, 0, &chain_seed_bufs[row], n_embd)?;
14271                            if row < j {
14272                                cg.interior[index].launch()?;
14273                            } else {
14274                                cg.last[index].launch()?;
14275                            }
14276                            scratch.plane_mut(index).0.len += 1;
14277                        }
14278                        sctr += 1; // mirrors the in-graph g_ctr bump (eager parity:
14279                        // counts the p-min-discarded token too)
14280                        // q retention: ONE async D2D of the persistent head-logits buffer
14281                        // into this round's slot j (stream-ordered after the replay).
14282                        e.copy_into(&mut dctx.q_slots[j], 0, &dctx.g_q, d_vocab)?;
14283                        // FILTERED capture: read the in-graph filter_stats scalars back per
14284                        // replay instead of a second full-vocab filter_stats per slot post-
14285                        // chain — bit-exact (the values the in-graph perturb consumed) and
14286                        // measured worth ~5% of vendor-default serving tok/s at K=3. Before
14287                        // the p-min break so the discarded slot's stats land too.
14288                        if filtered_stats_in_graph {
14289                            draft_stats.push((
14290                                e.dtoh(&dctx.g_mx)?[0],
14291                                e.dtoh(&dctx.g_th)?[0],
14292                                e.dtoh(&dctx.g_z)?[0],
14293                            ));
14294                        }
14295                        let idx = e.dtoh_u32_one(&dctx.g_tok)?;
14296                        // #87 SENTINEL TRAP (see the single-head graph arms).
14297                        if (idx as usize) >= d_vocab {
14298                            let seed_h = e.dtoh(&dctx.g_seed)?;
14299                            let seed_nan = seed_h.iter().filter(|v| v.is_nan()).count();
14300                            return Err(format!(
14301                                "draft(chain-graph-sampled) argmax sentinel 0x{idx:08x} >= \
14302                             d_vocab {d_vocab} at round {round} j={j} head={index} pos={pos}: \
14303                             head-out NaN {seed_nan}/{n_embd} — refusing to dereference the \
14304                             embed row (#87 trap)"
14305                            )
14306                            .into());
14307                        }
14308                        let d = match &mtp.d2t {
14309                            Some(map) => map[idx as usize],
14310                            None => idx,
14311                        };
14312                        draft_idx.push(idx);
14313                        if p_min > 0.0 {
14314                            let p = e.dtoh(&dctx.g_p)?[0];
14315                            if p < p_min && (j > 0 || (pmin0 && base0 == 1)) {
14316                                break;
14317                            }
14318                        }
14319                        draft.push(d);
14320                        chain_tokens.push(d);
14321                        chain_seed_bufs.push(e.clone_dtod(&dctx.g_seed)?);
14322                    }
14323                    // PURE-TEMP accept path: stats per used slot recomputed from the RETAINED
14324                    // q with the SAME filter_stats program the eager arm runs (deployment-
14325                    // keyed coop/plain choice, same input bits). The FILTERED graph read its
14326                    // stats back per replay above.
14327                    if !filtered_stats_in_graph {
14328                        for j in 0..draft.len().max(draft_idx.len()) {
14329                            let rows0 = e.htod_i32(&[0])?;
14330                            let (mut th_d, mut z_d, mut mx_d) =
14331                                (e.zeros(1)?, e.zeros(1)?, e.zeros(1)?);
14332                            e.filter_stats(
14333                                &dctx.q_slots[j],
14334                                d_vocab,
14335                                &rows0,
14336                                &mut th_d,
14337                                &mut z_d,
14338                                &mut mx_d,
14339                                d_vocab,
14340                                1,
14341                                sp_temp,
14342                                sp.top_k,
14343                                sp.top_p,
14344                                sp.min_p,
14345                            )?;
14346                            draft_stats.push((
14347                                e.dtoh(&mx_d)?[0],
14348                                e.dtoh(&th_d)?[0],
14349                                e.dtoh(&z_d)?[0],
14350                            ));
14351                        }
14352                    }
14353                } else if let (false, Some(gr)) =
14354                    (sampled || pen_on || !graph_round_ok, &dctx.graph)
14355                {
14356                    // GRAPH DRAFT: one dispatch per drafted token. The chain feeds itself on-device
14357                    // (in-graph argmax -> tok_d -> next replay's embed; h_nextn -> h_seed_d; pos_d
14358                    // inc'd in-graph); the host only reads 4B token (+4B p) and decides the break.
14359                    e.set_i32_one(&mut dctx.g_pos, (pos + base0) as i32)?;
14360                    e.set_u32_one(&mut dctx.g_tok, last_token)?;
14361                    e.copy_into(&mut dctx.g_seed, 0, &h_seed_buf, n_embd)?;
14362                    for j in 0..k_this {
14363                        // per-position mask upload (contents only — the graph's baked pointer is
14364                        // dctx.g_dmask). All-ones once masking goes dead mid-chain, so the captured
14365                        // mask node degrades to a no-op ban instead of needing a second graph.
14366                        if dmask_live
14367                            && !upload_draft_mask(
14368                                e,
14369                                constraint.as_deref_mut().unwrap(),
14370                                &mut dctx.g_dmask,
14371                                mtp.d2t.as_ref(),
14372                                d_vocab,
14373                                dmask_words,
14374                            )?
14375                        {
14376                            // no draft-vocab row is grammar-legal here (a trimmed FR-Spec head can
14377                            // genuinely miss the legal set): neutralize the captured mask node and
14378                            // finish the chain UNMASKED — exactly pre-lane behaviour, never worse.
14379                            e.htod_u32_into(&mut dctx.g_dmask, &vec![u32::MAX; dmask_words])?;
14380                            dmask_live = false;
14381                        }
14382                        gr.launch()?;
14383                        scratch.kv.len += 1; // host mirror (len_d advanced in-graph)
14384                        let idx = e.dtoh_u32_one(&dctx.g_tok)?;
14385                        // #87 SENTINEL TRAP: an all-NaN head-logits row leaves the device argmax's
14386                        // init sentinel (0x7FFFFFFF) in g_tok — feeding it onward dereferences
14387                        // embed_row(sentinel) = table + ~4.6TB (never mapped) inside the NEXT graph
14388                        // replay's embed node, and the MMU fault kills the CUDA context for the
14389                        // whole process (research/pp2spec-crash-20260807: 3 coredumps, byte-exact
14390                        // VA arithmetic). Refuse loudly instead; the diagnostics name the first-NaN
14391                        // buffer (g_seed = the verify-side handoff vs head-side compute).
14392                        if (idx as usize) >= d_vocab {
14393                            // g_seed is SELF-FED (the replay writes h_nextn back into it), so it
14394                            // reads as the head's OUTPUT at j; h_seed_buf is the round's INPUT
14395                            // seed, untouched since the round-start copy — the pair discriminates
14396                            // "seed arrived poisoned" from "head forward produced NaN".
14397                            let seed_h = e.dtoh(&dctx.g_seed)?;
14398                            let seed_nan = seed_h.iter().filter(|v| v.is_nan()).count();
14399                            let in_h = e.dtoh(&h_seed_buf)?;
14400                            let in_nan = in_h.iter().filter(|v| v.is_nan()).count();
14401                            return Err(format!(
14402                                "draft(graph) argmax sentinel 0x{idx:08x} >= d_vocab {d_vocab} at \
14403                             round {round} j={j} pos={pos}: head-out NaN {seed_nan}/{n_embd}, \
14404                             round-input-seed NaN {in_nan}/{n_embd} — refusing to dereference \
14405                             the embed row (#87 trap)"
14406                            )
14407                            .into());
14408                        }
14409                        // trimmed draft vocab -> target token id (identity when no d2t map)
14410                        let d = match &mtp.d2t {
14411                            Some(map) => map[idx as usize],
14412                            None => idx,
14413                        };
14414                        let draft_p = if p_min > 0.0
14415                            || opti_fork
14416                                .as_ref()
14417                                .is_some_and(|fork| fork.controller.is_some())
14418                        {
14419                            Some(e.dtoh(&dctx.g_p)?[0])
14420                        } else {
14421                            None
14422                        };
14423                        if j == 0 {
14424                            controller_draft_prob = draft_p;
14425                        }
14426                        if let Some(p) = draft_p.filter(|_| p_min > 0.0)
14427                            && p < p_min
14428                            && (j > 0 || (pmin0 && base0 == 1))
14429                        {
14430                            break;
14431                        }
14432                        draft.push(d);
14433                        // with a trimmed head the NEXT embed must read the TARGET id, not the draft
14434                        // index the argmax wrote — patch the persistent token buffer (4B htod).
14435                        if d != idx {
14436                            e.set_u32_one(&mut dctx.g_tok, d)?;
14437                        }
14438                        // advance the SPECULATIVE state with the proposal; a dead chain drops to
14439                        // unmasked drafting for the remaining positions (verify still arbitrates).
14440                        // speculative advance; a chain the grammar can no longer follow (EOS
14441                        // proposed) ends here. The captured mask node always runs, so a dead chain
14442                        // leaves the buffer NEUTRAL (all-ones = ban nothing) before it exits.
14443                        if dmask_live
14444                            && !constraint
14445                                .as_deref_mut()
14446                                .unwrap()
14447                                .draft_advance(d)
14448                                .map_err(|e2| format!("constraint: {e2}"))?
14449                        {
14450                            e.htod_u32_into(&mut dctx.g_dmask, &vec![u32::MAX; dmask_words])?;
14451                            break;
14452                        }
14453                    }
14454                // REGIME RE-TEST (lane/graph-s-key-exactness-20260819, widened by
14455                // lane/step37-draft-graph-serving-20260830): the sampled graph is legal ONLY
14456                // in the regime it was captured in. The condition used to read
14457                // `(sampled, &dctx.graph_s)` and trusted `s_key` to have dropped anything
14458                // else — which it could not, because the key omitted the filters. Both
14459                // halves are enforced: the key drops a stale graph, and this site refuses to
14460                // launch one whose key differs or whose regime is uncapturable (penalties).
14461                } else if let (true, Some(gr)) = (
14462                    sampled && s_capturable && dctx.s_key == Some(s_key) && graph_round_ok,
14463                    &dctx.graph_s,
14464                ) {
14465                    if skey_probe() {
14466                        eprintln!(
14467                            "[skey] chain=graph_s round={round} pure_temp={} capturable={} \
14468                             top_k={} top_p={} min_p={} s_key_parked={:?}",
14469                            pure_temp as u8,
14470                            s_capturable as u8,
14471                            sp.top_k,
14472                            sp.top_p,
14473                            sp.min_p,
14474                            dctx.s_key,
14475                        );
14476                    }
14477                    // SAMPLED GRAPH DRAFT: one replay per drafted token — head forward + gumbel +
14478                    // argmax in ONE dispatch; the host reads 4B token (+4B p), D2Ds q into slot j,
14479                    // and decides the break. Event-counter continuity: g_ctr is host-seeded to
14480                    // sctr-1 ONCE per round (outside the graph); the in-graph bump runs BEFORE the
14481                    // perturb, so replay j consumes counter sctr+j — exactly the eager arm's Philox
14482                    // stream. Host sctr advances in lockstep (computed, no readback needed).
14483                    e.set_i32_one(&mut dctx.g_pos, (pos + base0) as i32)?;
14484                    e.set_u32_one(&mut dctx.g_tok, last_token)?;
14485                    e.copy_into(&mut dctx.g_seed, 0, &h_seed_buf, n_embd)?;
14486                    e.set_u32_one(&mut dctx.g_ctr, sctr.wrapping_sub(1))?;
14487                    let filtered_stats_in_graph = s_key.filtered();
14488                    for j in 0..k_this {
14489                        gr.launch()?;
14490                        scratch.kv.len += 1; // host mirror (len_d advanced in-graph)
14491                        sctr += 1; // mirrors the in-graph g_ctr bump (eager parity:
14492                        // counts the p-min-discarded token too)
14493                        // q retention: ONE async D2D of the persistent head-logits buffer into this
14494                        // round's slot j (stream-ordered after the replay, before the next one).
14495                        e.copy_into(&mut dctx.q_slots[j], 0, &dctx.g_q, d_vocab)?;
14496                        // FILTERED capture: the replay's own filter_stats node already computed
14497                        // (th, z, mx) — read the three scalars back instead of paying a SECOND
14498                        // full-vocab filter_stats per slot post-chain (measured ~5% of vendor-
14499                        // default serving tok/s at K=3). Bit-exact by construction: these are
14500                        // the very values the in-graph perturb consumed. Read BEFORE the p-min
14501                        // break so the discarded slot's stats land too (accept-path indexing).
14502                        if filtered_stats_in_graph {
14503                            draft_stats.push((
14504                                e.dtoh(&dctx.g_mx)?[0],
14505                                e.dtoh(&dctx.g_th)?[0],
14506                                e.dtoh(&dctx.g_z)?[0],
14507                            ));
14508                        }
14509                        let idx = e.dtoh_u32_one(&dctx.g_tok)?;
14510                        // #87 SENTINEL TRAP (see the greedy graph arm above).
14511                        if (idx as usize) >= d_vocab {
14512                            let seed_h = e.dtoh(&dctx.g_seed)?;
14513                            let seed_nan = seed_h.iter().filter(|v| v.is_nan()).count();
14514                            return Err(format!(
14515                                "draft(graph-sampled) argmax sentinel 0x{idx:08x} >= d_vocab \
14516                             {d_vocab} at round {round} j={j} pos={pos}: round-seed NaN \
14517                             {seed_nan}/{n_embd} — refusing to dereference the embed row \
14518                             (#87 trap)"
14519                            )
14520                            .into());
14521                        }
14522                        let d = match &mtp.d2t {
14523                            Some(map) => map[idx as usize],
14524                            None => idx,
14525                        };
14526                        draft_idx.push(idx);
14527                        if p_min > 0.0 {
14528                            let p = e.dtoh(&dctx.g_p)?[0];
14529                            if p < p_min && (j > 0 || (pmin0 && base0 == 1)) {
14530                                break;
14531                            }
14532                        }
14533                        draft.push(d);
14534                        // trimmed head: the NEXT embed must read the TARGET id (see the greedy arm).
14535                        if d != idx {
14536                            e.set_u32_one(&mut dctx.g_tok, d)?;
14537                        }
14538                    }
14539                    // PURE-TEMP accept path: fill draft_stats per used slot post-chain (the
14540                    // stats degenerate to th=0 / full-Z; one filter_stats launch per slot).
14541                    // The FILTERED graph read its stats back per replay above.
14542                    if !filtered_stats_in_graph {
14543                        for j in 0..draft.len().max(draft_idx.len()) {
14544                            let rows0 = e.htod_i32(&[0])?;
14545                            let (mut th_d, mut z_d, mut mx_d) =
14546                                (e.zeros(1)?, e.zeros(1)?, e.zeros(1)?);
14547                            e.filter_stats(
14548                                &dctx.q_slots[j],
14549                                d_vocab,
14550                                &rows0,
14551                                &mut th_d,
14552                                &mut z_d,
14553                                &mut mx_d,
14554                                d_vocab,
14555                                1,
14556                                sp_temp,
14557                                sp.top_k,
14558                                sp.top_p,
14559                                sp.min_p,
14560                            )?;
14561                            draft_stats.push((
14562                                e.dtoh(&mx_d)?[0],
14563                                e.dtoh(&th_d)?[0],
14564                                e.dtoh(&z_d)?[0],
14565                            ));
14566                        }
14567                    }
14568                } else {
14569                    if skey_probe() && sampled {
14570                        eprintln!(
14571                            "[skey] chain=eager round={round} pure_temp={} top_k={} \
14572                             top_p={} min_p={} s_key_parked={:?}",
14573                            pure_temp as u8, sp.top_k, sp.top_p, sp.min_p, dctx.s_key,
14574                        );
14575                    }
14576                    // EAGER DRAFT (fallback: MoE head/trunk, huge k, MEMRA_SPEC_NOGRAPH, capture fail).
14577                    let chain_heads = !self.mtp_extra.is_empty();
14578                    let mut e_tok = last_token;
14579                    let mut d_seed = e.clone_dtod(&h_seed_buf)?;
14580                    let mut chain_tokens = if chain_heads {
14581                        vec![last_token]
14582                    } else {
14583                        Vec::new()
14584                    };
14585                    let mut chain_seeds = if chain_heads {
14586                        vec![e.clone_dtod(&h_seed_buf)?]
14587                    } else {
14588                        Vec::new()
14589                    };
14590                    for j in 0..k_this {
14591                        // GPU-ARGMAX DRAFT (2026-07-03): device logits + device argmax + 4-byte token
14592                        // read instead of the ~600KB full-vocab dtoh + host argmax per draft token.
14593                        let mtp_pos = pos + base0 + j;
14594                        // draft-side grammar mask (eager twin of the graph arm's in-graph node).
14595                        // A position with no legal draft-vocab row drops to unmasked drafting for
14596                        // the rest of the chain (pre-lane behaviour; verify still arbitrates).
14597                        if dmask_live {
14598                            dmask_live = upload_draft_mask(
14599                                e,
14600                                constraint.as_deref_mut().unwrap(),
14601                                &mut dctx.g_dmask,
14602                                mtp.d2t.as_ref(),
14603                                d_vocab,
14604                                dmask_words,
14605                            )?;
14606                        }
14607                        let mask = if dmask_live {
14608                            Some((&dctx.g_dmask, dmask_words))
14609                        } else {
14610                            None
14611                        };
14612                        let (dl_d, h_nextn) = if chain_heads {
14613                            if debug_spec {
14614                                eprintln!(
14615                                    "[mtp-chain-step] round={round} j={j} head={} replay_rows={}",
14616                                    mtp_chain_head_index(j, self.mtp_head_count()),
14617                                    chain_tokens.len(),
14618                                );
14619                            }
14620                            self.mtp_chain_forward_dev(
14621                                e,
14622                                &chain_tokens,
14623                                &chain_seeds,
14624                                &mut *scratch,
14625                                pos + base0 - 1,
14626                                embd_dev,
14627                                mask,
14628                            )?
14629                        } else {
14630                            self.mtp_head_forward_dev(
14631                                e,
14632                                mtp,
14633                                e_tok,
14634                                &d_seed,
14635                                &mut *scratch,
14636                                mtp_pos,
14637                                embd_dev,
14638                                mask,
14639                            )?
14640                        };
14641                        let tok_d = if sampled {
14642                            // FILTERED Gumbel-max: stats -> masked perturb -> argmax = one draw from
14643                            // the filtered softmax (filters off => th=0, exact v1 semantics).
14644                            if perturb_buf.is_none() {
14645                                perturb_buf = Some(e.zeros(d_vocab.max(n_vocab))?);
14646                            }
14647                            let mut q_row = e.clone_dtod(&dl_d)?; // retained q (penalized when on)
14648                            if pen_on {
14649                                let h = pen_hist_d.as_ref().unwrap();
14650                                let nh = h.len();
14651                                e.penalize_logits(
14652                                    &mut q_row,
14653                                    h,
14654                                    nh,
14655                                    sp.penalty_repeat,
14656                                    sp.penalty_freq,
14657                                    sp.penalty_present,
14658                                    d_vocab,
14659                                )?;
14660                            }
14661                            let rows0 = e.htod_i32(&[0])?;
14662                            let (mut th_d, mut z_d, mut mx_d) =
14663                                (e.zeros(1)?, e.zeros(1)?, e.zeros(1)?);
14664                            e.filter_stats(
14665                                &q_row, d_vocab, &rows0, &mut th_d, &mut z_d, &mut mx_d, d_vocab,
14666                                1, sp_temp, sp.top_k, sp.top_p, sp.min_p,
14667                            )?;
14668                            let (th, z, mx) =
14669                                (e.dtoh(&th_d)?[0], e.dtoh(&z_d)?[0], e.dtoh(&mx_d)?[0]);
14670                            let pb = perturb_buf.as_mut().unwrap();
14671                            e.gumbel_perturb_filtered(
14672                                &q_row, pb, d_vocab, sp_seed, sctr, sp_temp, mx, th,
14673                            )?;
14674                            sctr += 1;
14675                            draft_logits.push(q_row);
14676                            draft_stats.push((mx, th, z));
14677                            e.argmax_token_device(pb, d_vocab)?
14678                        } else {
14679                            e.argmax_token_device(&dl_d, d_vocab)?
14680                        };
14681                        let idx = e.dtoh_u32_one(&tok_d)?;
14682                        // #87 SENTINEL TRAP (eager twin — see the graph arm). Extra diagnostics
14683                        // here because the eager chain's operands are all readable: dl_d (the head
14684                        // logits row) and d_seed (this step's h_seed) name the first-NaN buffer.
14685                        if (idx as usize) >= d_vocab {
14686                            let dl_h = e.dtoh(&dl_d)?;
14687                            let dl_nan = dl_h.iter().filter(|v| v.is_nan()).count();
14688                            let seed_h = if chain_heads {
14689                                e.dtoh(chain_seeds.last().unwrap())?
14690                            } else {
14691                                e.dtoh(&d_seed)?
14692                            };
14693                            let seed_nan = seed_h.iter().filter(|v| v.is_nan()).count();
14694                            return Err(format!(
14695                                "draft(eager) argmax sentinel 0x{idx:08x} >= d_vocab {d_vocab} at \
14696                             round {round} j={j} pos={pos}: head-logits NaN {dl_nan}/{d_vocab}, \
14697                             step-seed NaN {seed_nan}/{n_embd} — refusing to dereference the \
14698                             embed row (#87 trap)"
14699                            )
14700                            .into());
14701                        }
14702                        let d = match &mtp.d2t {
14703                            Some(map) => map[idx as usize],
14704                            None => idx,
14705                        };
14706                        if sampled {
14707                            draft_idx.push(idx);
14708                        }
14709                        let draft_p = if p_min > 0.0
14710                            || opti_fork
14711                                .as_ref()
14712                                .is_some_and(|fork| fork.controller.is_some())
14713                        {
14714                            let p_d = e.prob_of_token_device(&dl_d, &tok_d, d_vocab)?;
14715                            Some(e.dtoh(&p_d)?[0])
14716                        } else {
14717                            None
14718                        };
14719                        if j == 0 {
14720                            controller_draft_prob = draft_p;
14721                        }
14722                        if let Some(p) = draft_p.filter(|_| p_min > 0.0)
14723                            && p < p_min
14724                            && (j > 0 || (pmin0 && base0 == 1))
14725                        {
14726                            break;
14727                        }
14728                        draft.push(d);
14729                        if chain_heads {
14730                            chain_tokens.push(d);
14731                            chain_seeds.push(h_nextn);
14732                        } else {
14733                            e_tok = d;
14734                            d_seed = h_nextn;
14735                        }
14736                        // speculative advance; a chain the grammar can no longer follow (EOS
14737                        // proposed) ends here — the prefix already proposed still rides verify.
14738                        if dmask_live
14739                            && !constraint
14740                                .as_deref_mut()
14741                                .unwrap()
14742                                .draft_advance(d)
14743                                .map_err(|e2| format!("constraint: {e2}"))?
14744                        {
14745                            break;
14746                        }
14747                    }
14748                    if !chain_heads
14749                        && opti_fork
14750                            .as_ref()
14751                            .is_some_and(|fork| fork.controller.is_some())
14752                    {
14753                        controller_eager_state = Some((e_tok, d_seed));
14754                    }
14755                }
14756            }
14757            let k_round = draft.len();
14758            if let Some(p) = pipe {
14759                p.draft_end(round);
14760            }
14761            drop(pipe_draft);
14762
14763            ph_mark(&mut ph_draft, phase_on);
14764            // --- 2. VERIFY: one batched target forward. With a pending bonus, it rides as col 0
14765            //         (committing its KV/recur inside the SAME weight read); drafts follow. ---
14766            let verify_tokens: Vec<u32> = match pending {
14767                Some(b) => {
14768                    let mut v = Vec::with_capacity(k_round + 1);
14769                    v.push(b);
14770                    v.extend_from_slice(&draft);
14771                    v
14772                }
14773                None => draft.clone(),
14774            };
14775            let base = if pending.is_some() { 1 } else { 0 };
14776            // ckpt (REPLAY-FREE partial accept): retain per-layer state-rebuild inputs alongside
14777            // the verify. Pure buffer keep-alives + dtod clones — kernel work is unchanged.
14778            let mut ckpt = if let Some(ticket) = current_opti.as_mut() {
14779                Some(ticket.take_ckpt())
14780            } else if spec_replay {
14781                None
14782            } else {
14783                Some(VerifyCkpt::new(self.layers.len()))
14784            };
14785            let controller_can_probe = base == 1
14786                && k_round == 1
14787                && out.len().saturating_add(2) < max_new
14788                && controller_draft_prob.is_some()
14789                && opti_fork
14790                    .as_ref()
14791                    .and_then(|fork| fork.controller.as_ref())
14792                    .is_some_and(|policy| !policy.breaker_tripped);
14793            let mut successor_attempt: Option<OptiControllerTicket> = None;
14794            let mut rejected_probe: Option<(f32, u32)> = None;
14795            let mut controller_prepared: Option<OptiControllerPrepared> = None;
14796            if controller_can_probe {
14797                // Prepare d2/q and, on admission, d3 before either current verify half is
14798                // issued. N stage 0 can then be followed immediately by N+1 stage 0; once N's
14799                // boundary fires, those dev0 launches overlap N stage 1 on dev1. Preparing on
14800                // the primary stream after N stage 1 would serialize the supposed pipeline.
14801                let eager_pos = scratch.kv.len + 1;
14802                let (optimistic_pending, pending_probability) = self.opti_controller_draft_step(
14803                    e,
14804                    mtp,
14805                    &mut dctx,
14806                    &mut *scratch,
14807                    d_vocab,
14808                    &mut controller_eager_state,
14809                    eager_pos,
14810                    embd_dev,
14811                    graph_round_ok,
14812                )?;
14813                let first_probability = controller_draft_prob
14814                    .ok_or("optipipe controller probe lost first-token probability")?;
14815                let q_proxy = first_probability * pending_probability;
14816                OPTI_GATE_CHECKS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
14817                OPTI_SHADOW_DRAFT_TOKENS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
14818                let admitted = opti_fork
14819                    .as_ref()
14820                    .and_then(|fork| fork.controller.as_ref())
14821                    .ok_or("optipipe controller policy disappeared")?
14822                    .admit(q_proxy);
14823                if admitted {
14824                    OPTI_GATE_ADMITS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
14825                    let eager_pos = scratch.kv.len + 1;
14826                    let (optimistic_draft, optimistic_draft_probability) = self
14827                        .opti_controller_draft_step(
14828                            e,
14829                            mtp,
14830                            &mut dctx,
14831                            &mut *scratch,
14832                            d_vocab,
14833                            &mut controller_eager_state,
14834                            eager_pos,
14835                            embd_dev,
14836                            graph_round_ok,
14837                        )?;
14838                    OPTI_SHADOW_DRAFT_TOKENS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
14839                    let eager_seed = controller_eager_state.take().map(|(token, seed)| {
14840                        debug_assert_eq!(token, optimistic_draft);
14841                        seed
14842                    });
14843                    controller_prepared = Some(OptiControllerPrepared {
14844                        verify_tokens: [optimistic_pending, optimistic_draft],
14845                        draft_prob: optimistic_draft_probability,
14846                        eager_seed,
14847                        q_proxy,
14848                        scratch_len: scratch.kv.len,
14849                    });
14850                } else {
14851                    OPTI_GATE_REJECTS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
14852                    OPTI_WASTED_DRAFT_TOKENS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
14853                    rejected_probe = Some((q_proxy, optimistic_pending));
14854                    eprintln!(
14855                        "[opti-controller] reject q={q_proxy:.6} threshold={:.3}",
14856                        opti_fork
14857                            .as_ref()
14858                            .and_then(|fork| fork.controller.as_ref())
14859                            .expect("controller policy")
14860                            .threshold,
14861                    );
14862                }
14863            }
14864            let fork_attempt = match fork_generation.take() {
14865                Some(generation) if base == 1 && k_round == 1 => Some(generation),
14866                Some(generation) => {
14867                    opti_fork
14868                        .as_mut()
14869                        .expect("fork generation without fork state")
14870                        .retire(generation)?;
14871                    None
14872                }
14873                None => None,
14874            };
14875            let (tlogits_d, vx) = if let Some(p) = pipe {
14876                self.decode_step_t_core_pipelined(
14877                    e,
14878                    &verify_tokens,
14879                    pos,
14880                    &mut *cache,
14881                    embd_dev,
14882                    ckpt.as_mut(),
14883                    p,
14884                    round,
14885                )?
14886            } else if controller_can_probe {
14887                let fence = opti_fork
14888                    .as_ref()
14889                    .ok_or("optipipe controller probe lost fork state")?
14890                    .fence;
14891                let boundary = match current_opti.as_mut() {
14892                    Some(ticket) => ticket.take_boundary(),
14893                    None => self.verify_stage0_issue(
14894                        e,
14895                        &verify_tokens,
14896                        pos,
14897                        &mut *cache,
14898                        embd_dev,
14899                        ckpt.as_mut(),
14900                        None,
14901                        &fence,
14902                        Some(true),
14903                        None,
14904                    )?,
14905                };
14906                if let Some(prepared) = controller_prepared.take() {
14907                    let generation = {
14908                        let fork = opti_fork
14909                            .as_mut()
14910                            .ok_or("optipipe controller admission lost fork state")?;
14911                        let generation = fork.reserve_successor()?;
14912                        let rt = fork.rt;
14913                        let snapshot_fence = fork.fence;
14914                        opti_snapshot_one_stage_owned_into(
14915                            e,
14916                            cache,
14917                            rt,
14918                            &snapshot_fence,
14919                            0,
14920                            fork.successor_snapshot_mut(),
14921                        )?;
14922                        generation
14923                    };
14924                    let mut successor_ckpt = VerifyCkpt::new(self.layers.len());
14925                    let successor_boundary = self.verify_stage0_issue(
14926                        e,
14927                        &prepared.verify_tokens,
14928                        pos + verify_tokens.len(),
14929                        &mut *cache,
14930                        embd_dev,
14931                        Some(&mut successor_ckpt),
14932                        None,
14933                        &fence,
14934                        Some(false),
14935                        None,
14936                    )?;
14937                    OPTI_FORK_ATTEMPTS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
14938                    let fork = opti_fork
14939                        .as_ref()
14940                        .ok_or("optipipe controller ticket lost fork state")?;
14941                    successor_attempt = Some(fork.controller_ticket(
14942                        generation,
14943                        successor_boundary,
14944                        successor_ckpt,
14945                        prepared.verify_tokens,
14946                        prepared.draft_prob,
14947                        prepared.eager_seed,
14948                        prepared.q_proxy,
14949                        prepared.scratch_len,
14950                    ));
14951                    eprintln!(
14952                        "[opti-controller] issue generation={} q={:.6} threshold={:.3} \
14953                         verify={:?}",
14954                        generation.id,
14955                        prepared.q_proxy,
14956                        fork.controller.expect("controller policy").threshold,
14957                        prepared.verify_tokens,
14958                    );
14959                }
14960                let result = self.verify_stage1_finish(
14961                    e,
14962                    boundary,
14963                    &mut *cache,
14964                    ckpt.as_mut(),
14965                    None,
14966                    &fence,
14967                    successor_attempt.is_none(),
14968                )?;
14969                if let Some(ticket) = current_opti.as_mut() {
14970                    ticket.settle();
14971                }
14972                if successor_attempt.is_some() {
14973                    let fork = opti_fork
14974                        .as_mut()
14975                        .ok_or("optipipe successor snapshot lost fork state")?;
14976                    let rt = fork.rt;
14977                    let snapshot_fence = fork.fence;
14978                    opti_snapshot_one_stage_owned_into(
14979                        e,
14980                        cache,
14981                        rt,
14982                        &snapshot_fence,
14983                        1,
14984                        fork.successor_snapshot_mut(),
14985                    )?;
14986                    // Publish N only after both independent successor-state queues are complete.
14987                    fork.rt.publish_to(1, &e.stream())?;
14988                }
14989                result
14990            } else if let Some(ticket) = current_opti.as_mut() {
14991                let fork = opti_fork
14992                    .as_mut()
14993                    .ok_or("optipipe carried controller ticket lost fork state")?;
14994                let boundary = ticket.take_boundary();
14995                let result = self.verify_stage1_finish(
14996                    e,
14997                    boundary,
14998                    &mut *cache,
14999                    ckpt.as_mut(),
15000                    None,
15001                    &fork.fence,
15002                    true,
15003                )?;
15004                ticket.settle();
15005                result
15006            } else if let Some(generation) = fork_attempt {
15007                let fork = opti_fork
15008                    .as_mut()
15009                    .expect("fork generation without fork state");
15010                fork.capture_seed(e, generation, &h_seed_buf, &fill_prev, scratch.kv.len)?;
15011                let action = fork.mode.action(generation.id);
15012                let boundary = self.verify_stage0_issue(
15013                    e,
15014                    &verify_tokens,
15015                    pos,
15016                    &mut *cache,
15017                    embd_dev,
15018                    ckpt.as_mut(),
15019                    None,
15020                    &fork.fence,
15021                    Some(true),
15022                    None,
15023                )?;
15024                OPTI_FORK_ATTEMPTS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
15025                let mut ticket = fork.ticket(generation, boundary);
15026                if action == OptiForkAction::Abort {
15027                    return Err(format!(
15028                        "optipipe forced abort with generation {} stage0 in flight",
15029                        generation.id,
15030                    )
15031                    .into());
15032                }
15033                fork.reconcile(
15034                    e,
15035                    &mut *cache,
15036                    &mut *scratch,
15037                    &snap,
15038                    &mut h_seed_buf,
15039                    &mut fill_prev,
15040                    generation,
15041                    action,
15042                    verify_tokens[0],
15043                )?;
15044                let result = if action == OptiForkAction::Hit {
15045                    let boundary = ticket.take_boundary();
15046                    self.verify_stage1_finish(
15047                        e,
15048                        boundary,
15049                        &mut *cache,
15050                        ckpt.as_mut(),
15051                        None,
15052                        &fork.fence,
15053                        true,
15054                    )?
15055                } else {
15056                    // The optimistic boundary slot has no reader. Re-run the unchanged serial
15057                    // verify only after E_restart published the restored stage-0 state.
15058                    self.decode_step_t_core(
15059                        e,
15060                        &verify_tokens,
15061                        pos,
15062                        &mut *cache,
15063                        embd_dev,
15064                        ckpt.as_mut(),
15065                    )?
15066                };
15067                ticket.settle();
15068                debug_assert_eq!(ticket.generation, generation);
15069                fork.retire(generation)?;
15070                result
15071            } else {
15072                // The serial verify every non-fork round takes — the MTP route's
15073                // verify-graph door. The pool is None unless MEMRA_SPEC_VERIFY_GRAPH armed
15074                // a pool above, and then the walk replays the captured trunk instead of
15075                // re-issuing it launch by launch. `graph_round_ok` is the round's
15076                // headroom snapshot (see GRAPH_LAUNCH_MIN_FREE): below the floor the
15077                // round declines the pool exactly like an over-cap round and rides the
15078                // byte-identical eager walk — the `[spec]` suspension line above
15079                // already named the round.
15080                let vg_round = if verify_tokens.len() <= vg_t_cap && graph_round_ok {
15081                    vg_guard.as_mut().and_then(|g| g.as_mut())
15082                } else {
15083                    if let Some(g) = vg_guard.as_mut().and_then(|g| g.as_mut()) {
15084                        // The commit reads this flag to pick its arm; a round that declines
15085                        // the pool must not inherit a stale `true` from the round before it.
15086                        g.round_slab = false;
15087                    }
15088                    None
15089                };
15090                self.decode_step_t_core_vg(
15091                    e,
15092                    &verify_tokens,
15093                    pos,
15094                    &mut *cache,
15095                    embd_dev,
15096                    ckpt.as_mut(),
15097                    vg_round,
15098                )?
15099            };
15100            let pipe_accept = match pipe {
15101                Some(p) => Some(p.accept_begin(round)?),
15102                None => None,
15103            };
15104
15105            if phase_sync {
15106                e.stream().synchronize()?;
15107            }
15108            ph_mark(&mut ph_verify, phase_on);
15109            // --- 3. GREEDY ACCEPT (walk prefix, stop at first mismatch) ---
15110            // DEVICE-ARGMAX ACCEPT: argmax every verify column ON DEVICE (same 2-pass kernels +
15111            // smallest-index tie-break as host argmax, argmax_gate-validated) and read back ONE
15112            // [T] u32 — replaces the T x n_vocab f32 dtoh + T host argmaxes per round.
15113            // t_pred[j] = target's greedy prediction for the slot after draft[j-1] (j>=1) or after
15114            // last_token (j==0). With a pending bonus, col 0 IS the prediction after last_token
15115            // (== the bonus), so every index shifts by `base` and last_pred is unused.
15116            let t_v = verify_tokens.len();
15117            let mut preds: Vec<u32> = Vec::new();
15118            if !sampled {
15119                for j in 0..t_v {
15120                    e.argmax_token_device_col(&tlogits_d, j, n_vocab, &mut preds_d, j)?;
15121                }
15122                preds = e.dtoh_u32(&preds_d)?; // <- the verify-GPU wait lands here
15123                // #87 SENTINEL TRAP, verify side: a sentinel pred becomes the round's bonus =
15124                // next round's last_token = the next chain's embed lookup. Catch it at the
15125                // source with the column named — an all-NaN VERIFY column implicates the
15126                // stage-split trunk (decode_step_t_core_ppn), not the draft head.
15127                if let Some(bad) = preds[..t_v].iter().position(|&p| (p as usize) >= n_vocab) {
15128                    let col = &tlogits_d.slice(bad * n_vocab..(bad + 1) * n_vocab);
15129                    let mut probe = e.zeros(n_vocab)?;
15130                    e.copy_view_into(&mut probe, 0, col, n_vocab)?;
15131                    let col_h = e.dtoh(&probe)?;
15132                    let col_nan = col_h.iter().filter(|v| v.is_nan()).count();
15133                    return Err(format!(
15134                        "verify argmax sentinel 0x{:08x} >= n_vocab {n_vocab} at round {round} \
15135                         col {bad}/{t_v} pos={pos}: verify-logits col NaN {col_nan}/{n_vocab} \
15136                         — the verify TRUNK produced a poisoned column (#87 trap). Run \
15137                         MEMRA_SPEC_NAN_SCAN=1 to name the layer that creates it (=2 to split \
15138                         that layer into attention and routed MoE). NOT the draft head, and NOT \
15139                         the PP stage split this message used to name: pp_cuts() returns None \
15140                         without MEMRA_PP_STAGES, so decode_step_t_core_ppn never runs unless \
15141                         that variable is set.",
15142                        preds[bad]
15143                    )
15144                    .into());
15145                }
15146            }
15147            ph_mark(&mut ph_wait, phase_on);
15148            let t_pred = |j: usize| -> u32 {
15149                if j == 0 && base == 0 {
15150                    last_pred
15151                } else {
15152                    // GREEDY-ONLY: `preds` is filled under `if !sampled` above. The debug print
15153                    // used to call this from the sampled arm and panicked the worker; it now goes
15154                    // through `debug_t_pred0`. Keep the strict index here — in the greedy walk an
15155                    // out-of-range pred is a real bug, not something to paper over.
15156                    debug_assert!(
15157                        !sampled,
15158                        "t_pred is greedy-only: `preds` is empty in the sampled arm"
15159                    );
15160                    preds[base + j - 1]
15161                }
15162            };
15163            let mut devacc_seeded = false;
15164            let mut devacc_acc: Option<CudaSlice<u32>> = None;
15165            let (n_acc, bonus) = if !sampled {
15166                // ROUND-STREAM stage (a) (MEMRA_SPEC_DEVACC=1 opt-in): the walk runs ON DEVICE
15167                // (spec_accept_greedy, verbatim rule) and the host reads back 8B (n_acc, bonus)
15168                // instead of the [T] preds. Same sync count — machinery for stages (b)/(c),
15169                // gated on token identity vs the host walk (the arms below are bit-equal rules).
15170                if crate::spec::spec_devacc() && k_round > 0 && !spec_replay && constraint.is_none()
15171                {
15172                    let draft_d = e.htod_u32_v(&draft)?;
15173                    let mut acc_out = e.alloc_u32_zeroed(2)?;
15174                    e.spec_accept_greedy(
15175                        &preds_d,
15176                        &draft_d,
15177                        last_pred,
15178                        base,
15179                        k_round,
15180                        &mut acc_out,
15181                    )?;
15182                    devacc_acc = Some(acc_out.clone());
15183                    // stage (b): next-round seed gathered ON DEVICE from acc_out before the host
15184                    // ever reads n_acc (j=base+n_acc -> vx col j-1; j==0 -> fill_prev). The three
15185                    // non-replay commit arms skip their host-offset seed copies (guarded below);
15186                    // the legacy spec_replay arm keeps its own rx-based seeding (excluded here).
15187                    // NOTE: fill_prev is NOT updated here — the commit arms' TRUE-HIDDEN
15188                    // REFRESH reads the OLD fill_prev (predecessor of this round's verify batch);
15189                    // the update lands after the arms (devacc_seeded guard below).
15190                    e.spec_seed_gather(&vx, &fill_prev, &acc_out, &mut h_seed_buf, base, n_embd)?;
15191                    // 3a: KV lens roll back on device (len = saved + base + n_acc, all arms'
15192                    // unified rule; full accept rewrites the verify-left value). Host mirrors
15193                    // update after the readback; commit_verified_prefix skips its len_d writes.
15194                    if let Some(successor) = successor_attempt.as_ref() {
15195                        opti_fork
15196                            .as_mut()
15197                            .ok_or("optipipe successor reconcile lost fork state")?
15198                            .queue_actual_reconcile(
15199                                e,
15200                                &snap,
15201                                &acc_out,
15202                                successor.verify_tokens[0],
15203                                base,
15204                            )?;
15205                    } else if let Some(ptrs) = &kv_len_ptrs {
15206                        let saved: Vec<i32> = (0..self.layers.len())
15207                            .map(|il| snap.kv_len[il].map(|v| v as i32).unwrap_or(0))
15208                            .collect();
15209                        let saved_d = e.htod_i32(&saved)?;
15210                        e.spec_rollback_kv(ptrs, &saved_d, &acc_out, base, self.layers.len())?;
15211                    }
15212                    devacc_seeded = true;
15213                    let ab = e.dtoh_u32(&acc_out)?;
15214                    (ab[0] as usize, ab[1])
15215                } else {
15216                    let mut n_acc = 0usize;
15217                    #[allow(clippy::needless_range_loop)]
15218                    // allow: the explicit index loop keeps the offset arithmetic visible and aligned with the device-side indexing
15219                    for j in 0..k_round {
15220                        if t_pred(j) == draft[j] {
15221                            n_acc += 1;
15222                        } else {
15223                            break;
15224                        }
15225                    }
15226                    // bonus = target's own token at the first non-accepted slot. n_acc in 0..=k; t_pred
15227                    // is defined for j in 0..=k (j==0 -> last_logits, j>=1 -> col j-1, last col = k-1).
15228                    (n_acc, t_pred(n_acc))
15229                }
15230            } else {
15231                // --- SAMPLED ACCEPT (rejection sampling): u_j < p_j(x_j)/q_j(x_j) walk ---
15232                if col_buf.is_none() {
15233                    col_buf = Some(e.zeros(n_vocab)?);
15234                }
15235                // FILTERED p_j: per-verify-col stats (one batched filter_stats call), then the
15236                // filtered gather. j==0&&base==0 reads last_col (its own stats row appended).
15237                let mut pj = vec![0f32; k_round.max(1)];
15238                let mut col_stats: Vec<(f32, f32, f32)> = Vec::new(); // (max, th, z) per verify col used
15239                if k_round > 0 {
15240                    let mut ids: Vec<u32> = Vec::new();
15241                    let mut rows: Vec<i32> = Vec::new();
15242                    #[allow(clippy::needless_range_loop)]
15243                    // allow: the explicit index loop keeps the offset arithmetic visible and aligned with the device-side indexing
15244                    for j in 0..k_round {
15245                        if j > 0 || base == 1 {
15246                            ids.push(draft[j]);
15247                            rows.push((base + j) as i32 - 1);
15248                        }
15249                    }
15250                    if !ids.is_empty() {
15251                        let nr = rows.len();
15252                        // penalties: materialize the used columns into one contiguous penalized
15253                        // buffer (rows remapped 0..nr) so stats+gathers see the penalized p.
15254                        // penalties: materialize used columns contiguously, penalize all rows in
15255                        // one launch, and point stats+gathers at the penalized buffer (rows 0..nr).
15256                        let p_rows: Vec<i32> = if pen_on {
15257                            (0..nr as i32).collect()
15258                        } else {
15259                            rows.clone()
15260                        };
15261                        if pen_on {
15262                            if pcol_buf.as_ref().map(|b| b.len()).unwrap_or(0) < nr * n_vocab {
15263                                pcol_buf = Some(e.zeros(nr * n_vocab)?);
15264                            }
15265                            let pc = pcol_buf.as_mut().unwrap();
15266                            for (i2, &r) in rows.iter().enumerate() {
15267                                let c = r as usize;
15268                                e.copy_view_into(
15269                                    pc,
15270                                    i2 * n_vocab,
15271                                    &tlogits_d.slice(c * n_vocab..(c + 1) * n_vocab),
15272                                    n_vocab,
15273                                )?;
15274                            }
15275                            let h = pen_hist_d.as_ref().unwrap();
15276                            let nh = h.len();
15277                            e.penalize_logits_rows(
15278                                pc,
15279                                h,
15280                                nh,
15281                                sp.penalty_repeat,
15282                                sp.penalty_freq,
15283                                sp.penalty_present,
15284                                n_vocab,
15285                                nr,
15286                            )?;
15287                        }
15288                        let p_src: &CudaSlice<f32> = if pen_on {
15289                            pcol_buf.as_ref().unwrap()
15290                        } else {
15291                            &tlogits_d
15292                        };
15293                        let rowsd = e.htod_i32(&p_rows)?;
15294                        let (mut th_d, mut z_d, mut mx_d) =
15295                            (e.zeros(nr)?, e.zeros(nr)?, e.zeros(nr)?);
15296                        e.filter_stats(
15297                            p_src, n_vocab, &rowsd, &mut th_d, &mut z_d, &mut mx_d, n_vocab, nr,
15298                            sp_temp, sp.top_k, sp.top_p, sp.min_p,
15299                        )?;
15300                        let idsd = e.htod_u32_v(&ids)?;
15301                        let mut outd = e.zeros(nr)?;
15302                        e.softmax_gather_filtered(
15303                            p_src, n_vocab, &idsd, &rowsd, &th_d, &z_d, &mut outd, n_vocab, nr,
15304                            sp_temp,
15305                        )?;
15306                        let outv = e.dtoh(&outd)?;
15307                        let (thv, zv, mxv) = (e.dtoh(&th_d)?, e.dtoh(&z_d)?, e.dtoh(&mx_d)?);
15308                        let mut oi = 0usize;
15309                        #[allow(clippy::needless_range_loop)]
15310                        // allow: the explicit index loop keeps the offset arithmetic visible and aligned with the device-side indexing
15311                        for j in 0..k_round {
15312                            if j > 0 || base == 1 {
15313                                pj[j] = outv[oi];
15314                                oi += 1;
15315                            }
15316                        }
15317                        col_stats = (0..nr).map(|i| (mxv[i], thv[i], zv[i])).collect();
15318                    }
15319                    if base == 0 {
15320                        let lc: &CudaSlice<f32> = if pen_on {
15321                            if col_buf.is_none() {
15322                                col_buf = Some(e.zeros(n_vocab)?);
15323                            }
15324                            let cb = col_buf.as_mut().unwrap();
15325                            e.copy_into(
15326                                cb,
15327                                0,
15328                                last_col_logits
15329                                    .as_ref()
15330                                    .expect("sampled: last_col_logits unset"),
15331                                n_vocab,
15332                            )?;
15333                            let h = pen_hist_d.as_ref().unwrap();
15334                            let nh = h.len();
15335                            e.penalize_logits(
15336                                cb,
15337                                h,
15338                                nh,
15339                                sp.penalty_repeat,
15340                                sp.penalty_freq,
15341                                sp.penalty_present,
15342                                n_vocab,
15343                            )?;
15344                            col_buf.as_ref().unwrap()
15345                        } else {
15346                            last_col_logits
15347                                .as_ref()
15348                                .expect("sampled: last_col_logits unset")
15349                        };
15350                        let rows0 = e.htod_i32(&[0])?;
15351                        let (mut th_d, mut z_d, mut mx_d) = (e.zeros(1)?, e.zeros(1)?, e.zeros(1)?);
15352                        e.filter_stats(
15353                            lc, n_vocab, &rows0, &mut th_d, &mut z_d, &mut mx_d, n_vocab, 1,
15354                            sp_temp, sp.top_k, sp.top_p, sp.min_p,
15355                        )?;
15356                        let idsd = e.htod_u32_v(&[draft[0]])?;
15357                        let mut outd = e.zeros(1)?;
15358                        e.softmax_gather_filtered(
15359                            lc, n_vocab, &idsd, &rows0, &th_d, &z_d, &mut outd, n_vocab, 1, sp_temp,
15360                        )?;
15361                        pj[0] = e.dtoh(&outd)?[0];
15362                        last_col_stats =
15363                            Some((e.dtoh(&mx_d)?[0], e.dtoh(&th_d)?[0], e.dtoh(&z_d)?[0]));
15364                    }
15365                }
15366                // q source: the graph arms (single-head AND chain) retained the head logits
15367                // in the persistent q_slots; the eager arm in per-round draft_logits clones.
15368                // Same raw-logit values either way. FILTERED q_j: stats from draft_stats
15369                // (eager pushes in-chain; the graph arms compute them post-replay from the
15370                // retained q with the same filter_stats program — bit-identical to the
15371                // in-graph stats that shaped the draw, keeping ONE accept path).
15372                let q_bufs: &[CudaSlice<f32>] = if dctx.graph_s.is_some() || dctx.chain_s.is_some()
15373                {
15374                    &dctx.q_slots
15375                } else {
15376                    &draft_logits
15377                };
15378                let mut n_acc = 0usize;
15379                for j in 0..k_round {
15380                    let (qmx, qth, qz) = draft_stats[j];
15381                    let idsd = e.htod_u32_v(&[draft_idx[j]])?;
15382                    let rowsd = e.htod_i32(&[0])?;
15383                    let thd = e.htod(&[qth])?;
15384                    let zd = e.htod(&[qz])?;
15385                    let _ = qmx;
15386                    let mut outd = e.zeros(1)?;
15387                    e.softmax_gather_filtered(
15388                        &q_bufs[j], d_vocab, &idsd, &rowsd, &thd, &zd, &mut outd, d_vocab, 1,
15389                        sp_temp,
15390                    )?;
15391                    let qj = e.dtoh(&outd)?[0];
15392                    let u = host_u01(sp_seed, uctr);
15393                    uctr += 1;
15394                    let accept = (u as f64) * (qj as f64) < pj[j] as f64;
15395                    // SKEY PROBE: q == 0 for the token the draft actually proposed is the
15396                    // exactness signature (see `skey_probe`). Impossible when the draft was
15397                    // drawn from the same filtered distribution the verify reconstructs here;
15398                    // `u * 0 < p` makes it an UNCONDITIONAL accept whenever p > 0.
15399                    if skey_probe() && qj == 0.0 {
15400                        eprintln!(
15401                            "[skey] EXACTNESS q=0 round={round} j={j} draft_tok={} \
15402                             draft_idx={} p={:e} u={u} accepted={} th_z={:?}",
15403                            draft[j], draft_idx[j], pj[j], accept as u8, draft_stats[j],
15404                        );
15405                    }
15406                    if accept {
15407                        n_acc += 1;
15408                    } else {
15409                        break;
15410                    }
15411                }
15412                let bonus = if n_acc == k_round {
15413                    // FULL ACCEPT: bonus ~ FILTERED softmax at the last verify column.
15414                    let col = base + k_round - 1;
15415                    let cb = col_buf.as_mut().unwrap();
15416                    e.copy_view_into(
15417                        cb,
15418                        0,
15419                        &tlogits_d.slice(col * n_vocab..(col + 1) * n_vocab),
15420                        n_vocab,
15421                    )?;
15422                    if pen_on {
15423                        let h = pen_hist_d.as_ref().unwrap();
15424                        let nh = h.len();
15425                        e.penalize_logits(
15426                            cb,
15427                            h,
15428                            nh,
15429                            sp.penalty_repeat,
15430                            sp.penalty_freq,
15431                            sp.penalty_present,
15432                            n_vocab,
15433                        )?;
15434                    }
15435                    if perturb_buf.is_none() {
15436                        perturb_buf = Some(e.zeros(d_vocab.max(n_vocab))?);
15437                    }
15438                    // STATS MUST COME FROM THIS COLUMN (bug fix 2026-08-05, lane/sampler-
15439                    // truncation-fix; receipts research/sampfix-20260805/). The old code reused
15440                    // `col_stats.last()` here, which is ALWAYS the wrong row: the gathered set
15441                    // covers verify columns 0..=(base+k_round-2) (rows pushed as base+j-1), while
15442                    // the full-accept bonus samples column base+k_round-1 — exactly ONE PAST the
15443                    // last gathered column, in both base arms. `th` is a threshold in e-units of
15444                    // its OWN row's max, so feeding a neighbour's (row_max, th) into
15445                    // gumbel_perturb_filtered mis-scales every e0 = exp((x-row_max)/T). When the
15446                    // donor column's peak is higher by more than T*ln(1/th), EVERY id fails
15447                    // `e0 >= th`, the whole perturbed row becomes -3.4e38, and the 2-pass argmax
15448                    // falls through to its smallest-index tie-break => token id 0 ("!") spliced
15449                    // mid-word. Fragility is ordered by how large th is: min_p pins th = min_p
15450                    // (0.05 => trigger at delta > 2.4 at T=0.8, fires constantly), top_p's
15451                    // mass-boundary th is smaller, top_k's k-th-largest th smaller still — which
15452                    // is why the head-to-head matrix saw min_p and top_p corrupt while top_k-only
15453                    // stayed clean. The pure-temp default regime is immune (th == 0 masks nothing,
15454                    // and row_max is unused once nothing is masked), so this fix is a byte-level
15455                    // no-op for the untruncated serve default. One extra one-block filter_stats
15456                    // per full-accept round is the whole cost.
15457                    let (mx, th) = {
15458                        let rows0 = e.htod_i32(&[0])?;
15459                        let (mut th_d, mut z_d, mut mx_d) = (e.zeros(1)?, e.zeros(1)?, e.zeros(1)?);
15460                        let cb0 = col_buf.as_ref().unwrap();
15461                        e.filter_stats(
15462                            cb0, n_vocab, &rows0, &mut th_d, &mut z_d, &mut mx_d, n_vocab, 1,
15463                            sp_temp, sp.top_k, sp.top_p, sp.min_p,
15464                        )?;
15465                        (e.dtoh(&mx_d)?[0], e.dtoh(&th_d)?[0])
15466                    };
15467                    let pb = perturb_buf.as_mut().unwrap();
15468                    let cb2 = col_buf.as_ref().unwrap();
15469                    e.gumbel_perturb_filtered(cb2, pb, n_vocab, sp_seed, sctr, sp_temp, mx, th)?;
15470                    sctr += 1;
15471                    let td = e.argmax_token_device(pb, n_vocab)?;
15472                    e.dtoh_u32_one(&td)?
15473                } else {
15474                    // REJECT at n_acc: bonus ~ norm(max(0, softmax_T(p) - softmax_T(q))).
15475                    let cb = col_buf.as_mut().unwrap();
15476                    if n_acc > 0 || base == 1 {
15477                        let col = base + n_acc - 1;
15478                        e.copy_view_into(
15479                            cb,
15480                            0,
15481                            &tlogits_d.slice(col * n_vocab..(col + 1) * n_vocab),
15482                            n_vocab,
15483                        )?;
15484                    } else {
15485                        let lc = last_col_logits.as_ref().unwrap();
15486                        e.copy_into(cb, 0, lc, n_vocab)?;
15487                    }
15488                    if pen_on {
15489                        let h = pen_hist_d.as_ref().unwrap();
15490                        let nh = h.len();
15491                        e.penalize_logits(
15492                            cb,
15493                            h,
15494                            nh,
15495                            sp.penalty_repeat,
15496                            sp.penalty_freq,
15497                            sp.penalty_present,
15498                            n_vocab,
15499                        )?;
15500                    }
15501                    let cb2 = col_buf.as_ref().unwrap();
15502                    let sc = sctr;
15503                    sctr += 1;
15504                    // p-stats for the reject column: from col_stats when the col was gathered,
15505                    // else (j==0&&base==0) from last_col_stats.
15506                    let p_stats = if n_acc > 0 || base == 1 {
15507                        // col index within the gathered set == number of gathered cols before n_acc
15508                        let gi = if base == 1 { n_acc } else { n_acc - 1 };
15509                        col_stats.get(gi).copied().unwrap_or({
15510                            (0.0, 0.0, 1.0) // unreachable: gathered cols always cover the reject slot
15511                        })
15512                    } else {
15513                        last_col_stats.expect("sampled: last_col_stats unset at reject")
15514                    };
15515                    let q_stats = draft_stats[n_acc];
15516                    if let Some(map) = &d2t_dev {
15517                        if q_full_buf.is_none() {
15518                            q_full_buf = Some(e.zeros(n_vocab)?);
15519                        }
15520                        let qf = q_full_buf.as_mut().unwrap();
15521                        e.scatter_trim_logits(&q_bufs[n_acc], map, qf, d_vocab, n_vocab)?;
15522                        let qf2 = q_full_buf.as_ref().unwrap();
15523                        e.residual_sample_filtered(
15524                            cb2,
15525                            Some(qf2),
15526                            n_vocab,
15527                            sp_temp,
15528                            sp_seed,
15529                            sc,
15530                            p_stats,
15531                            q_stats,
15532                            &mut sample_tok,
15533                        )?;
15534                    } else {
15535                        e.residual_sample_filtered(
15536                            cb2,
15537                            Some(&q_bufs[n_acc]),
15538                            n_vocab,
15539                            sp_temp,
15540                            sp_seed,
15541                            sc,
15542                            p_stats,
15543                            q_stats,
15544                            &mut sample_tok,
15545                        )?;
15546                    }
15547                    e.dtoh_u32(&sample_tok)?[0]
15548                };
15549                (
15550                    n_acc,
15551                    guard_vocab_token(
15552                        bonus,
15553                        n_vocab,
15554                        &format!("sampled verify bonus at round {round} pos={pos} n_acc={n_acc}"),
15555                    )?,
15556                )
15557            };
15558            // --- 3b. GRAMMAR TRUNCATION (constrained spec, 2026-08-03): the grammar is
15559            // an extra rejection rule AFTER the exactness verify (the batched-verify-twins
15560            // ordering). Walk the accepted drafts through the grammar in commit order; the
15561            // first illegal token truncates acceptance at its slot, and that slot's emission
15562            // is recomputed as the MASKED argmax of the target's own verify column — token-
15563            // identical to constrained plain greedy decode (an unmasked argmax that is
15564            // grammar-legal IS the masked argmax: masking only removes competitors). The
15565            // column D2H (~1MB) is paid only when a cut fires — the tight-grammar cost,
15566            // measured in acceptance numbers, never hidden.
15567            let (n_acc, bonus) = match constraint.as_deref_mut() {
15568                None => (n_acc, bonus),
15569                Some(c) => {
15570                    fn ce(e2: String) -> Box<dyn std::error::Error> {
15571                        format!("constraint: {e2}").into()
15572                    }
15573                    let mut na = n_acc;
15574                    let mut cut = false;
15575                    for (j, &d) in draft.iter().enumerate().take(n_acc) {
15576                        if c.is_allowed(d).map_err(ce)? {
15577                            c.consume(d).map_err(ce)?;
15578                        } else {
15579                            na = j;
15580                            cut = true;
15581                            dm_cut_tokens += n_acc - j;
15582                            break;
15583                        }
15584                    }
15585                    if cut {
15586                        dm_cuts += 1;
15587                    }
15588                    let mut bo = bonus;
15589                    if cut || !c.is_allowed(bo).map_err(ce)? {
15590                        let mut row = if na == 0 && base == 0 {
15591                            init_logits_host
15592                                .clone()
15593                                .ok_or("constraint: init logits missing (round-0 cut)")?
15594                        } else {
15595                            e.dtoh_view(
15596                                &tlogits_d.slice((base + na - 1) * n_vocab..(base + na) * n_vocab),
15597                            )?
15598                        };
15599                        c.mask_logits(&mut row).map_err(ce)?;
15600                        bo = argmax(&row) as u32;
15601                    }
15602                    c.consume(bo).map_err(ce)?;
15603                    (na, bo)
15604                }
15605            };
15606            let mut successor_valid = false;
15607            if let Some((q_proxy, expected_d2)) = rejected_probe {
15608                let v_n = n_acc == 1 && bonus == expected_d2;
15609                eprintln!(
15610                    "[opti-controller] shadow q={q_proxy:.6} admitted=false v_n={v_n} \
15611                     expected_d2={expected_d2} n_acc={n_acc} bonus={bonus}",
15612                );
15613            }
15614            if let Some(successor) = successor_attempt.as_ref() {
15615                successor_valid = n_acc == 1 && bonus == successor.verify_tokens[0];
15616                let generation = successor.generation;
15617                let q_proxy = successor.q_proxy;
15618                let expected_pending = successor.verify_tokens[0];
15619                let resolution_ms = successor.issued_at.elapsed().as_secs_f64() * 1e3;
15620                let fork = opti_fork
15621                    .as_mut()
15622                    .ok_or("optipipe successor resolution lost fork state")?;
15623                fork.finish_actual_reconcile(e, &mut *cache, &snap, n_acc, base, successor_valid)?;
15624                if successor_valid {
15625                    OPTI_FORK_HITS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
15626                } else {
15627                    OPTI_FORK_MISSES.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
15628                    OPTI_RECONCILES.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
15629                    OPTI_WASTED_DRAFT_TOKENS.fetch_add(2, std::sync::atomic::Ordering::Relaxed);
15630                }
15631                let breaker_tripped = fork
15632                    .controller
15633                    .as_mut()
15634                    .expect("controller policy")
15635                    .resolve(successor_valid);
15636                if breaker_tripped {
15637                    OPTI_BREAKER_TRIPS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
15638                }
15639                eprintln!(
15640                    "[opti-controller] resolve generation={} hit={} q={q_proxy:.6} \
15641                     expected_pending={expected_pending} n_acc={n_acc} bonus={bonus} \
15642                     resolution_ms={resolution_ms:.3} reconcile={} breaker={}",
15643                    generation.id, successor_valid, !successor_valid, breaker_tripped,
15644                );
15645                if !successor_valid {
15646                    let mut successor = successor_attempt
15647                        .take()
15648                        .expect("controller successor disappeared on miss");
15649                    successor.settle();
15650                    fork.retire(generation)?;
15651                }
15652            }
15653            total_drafted += k_round;
15654            total_accepted += n_acc;
15655            if let Some(t) = sess_telem {
15656                // Greedy, rejection-sampling, and grammar truncation all converge here after
15657                // the accept decision is already on host. Fixed-size relaxed atomics only.
15658                t.record_round(k_round, n_acc);
15659            }
15660            if spec_stats {
15661                st_len_hist[k_round] += 1;
15662                #[allow(clippy::needless_range_loop)]
15663                // allow: the explicit index loop keeps the offset arithmetic visible and aligned with the device-side indexing
15664                for j in 0..k_round {
15665                    st_drafted[j] += 1;
15666                }
15667                #[allow(clippy::needless_range_loop)]
15668                // allow: the explicit index loop keeps the offset arithmetic visible and aligned with the device-side indexing
15669                for j in 0..n_acc {
15670                    st_accepted[j] += 1;
15671                }
15672                if n_acc == k_round {
15673                    st_full += 1;
15674                }
15675            }
15676
15677            if debug_spec {
15678                eprintln!(
15679                    "[R{round}] pos={pos} out_len={} last_tok={last_token} draft={draft:?} n_acc={n_acc} bonus={bonus} t_pred0={}",
15680                    out.len(),
15681                    // NOT `t_pred(0)`: `preds` is filled only under `if !sampled` above, so on a
15682                    // sampled request round >= 1 (base == 1) indexed an EMPTY vector and PANICKED
15683                    // the GPU worker thread — a debug flag that killed the exact regime you would
15684                    // set it to investigate. See `debug_t_pred0`.
15685                    debug_t_pred0(sampled, base, last_pred, &preds)
15686                );
15687            }
15688
15689            // --- 4. COMMIT: draft[0..n_acc] then bonus (n_acc + 1 tokens) ---
15690            let commit_started = std::time::Instant::now();
15691            // SESSION MODE: every accepted column is already in the CACHE — `out` must carry all
15692            // of them (overshoot past max_new included) or `committed` under-counts the cache rows
15693            // and the next turn's continuation seeds one token off (gate-caught 2026-07-05). The
15694            // single-shot path keeps the cap (its caller truncates + drops the cache anyway).
15695            #[allow(clippy::needless_range_loop)]
15696            // allow: the explicit index loop keeps the offset arithmetic visible and aligned with the device-side indexing
15697            for j in 0..n_acc {
15698                if !session_mode && out.len() >= max_new {
15699                    break;
15700                }
15701                out.push(draft[j]);
15702            }
15703            if pen_on {
15704                pen_hist.extend_from_slice(&draft[0..n_acc]);
15705                pen_hist.push(bonus);
15706            }
15707            let bonus_emitted = session_mode || out.len() < max_new;
15708            if bonus_emitted {
15709                out.push(bonus);
15710            }
15711            last_token = bonus;
15712
15713            // --- 5. ROLLBACK + advance (§C) ---
15714            if n_acc == k_round && !spec_replay {
15715                // FULL ACCEPT, BONUS FOLD: all verify columns (pending? + drafts) are committed in
15716                // cache; the NEW bonus stays PENDING for the next round's verify batch — NO extra
15717                // T=1 trunk pass. The next draft chain seeds from the MTP block's h_nextn at the
15718                // bonus position: one MTP-block pass (~1/33 trunk cost) replaces the trunk read.
15719                // last_pred is dead in the pending path (t_pred reads verify col 0).
15720                //
15721                // PERSISTENT DRAFT KV, full-accept fill: the chain covered last_token +
15722                // draft[0..k_round-2] as INPUTS (slots P..P'-2); draft[k_round-1] (slot P'-1) was
15723                // only ever an output, so its entry is MISSING. Fill it from vh_seed — its EXACT
15724                // trunk hidden (the last verify column). set_len first: a p-min break may have
15725                // left one extra chain append at that slot. Partial accepts need NO fill (the
15726                // chain already covered every accepted position; round-start set_len truncates).
15727                self.restore_step_tp_kv_verified_prefix(e, &mut *cache, &snap, t_v, false)?;
15728                let mut vh_seed = e.zeros(n_embd)?;
15729                e.copy_view_into(
15730                    &mut vh_seed,
15731                    0,
15732                    &vx.slice((t_v - 1) * n_embd..t_v * n_embd),
15733                    n_embd,
15734                )?;
15735                if refresh {
15736                    // TRUE-HIDDEN REFRESH (2026-07-03, the HANDOVER-listed acceptance lever):
15737                    // overwrite ALL committed positions' scratch entries with K/V from their EXACT
15738                    // verify hiddens — the reference engine's mtp_update fills from true hiddens;
15739                    // the full stack (vx) is already resident from the verify. Replaces both the
15740                    // chain-approximate entries AND the old last-token-only fill. Acceptance-only
15741                    // (draft attention quality); exactness stays the verify's job.
15742                    scratch.set_len(e, pos)?;
15743                    // PREDECESSOR pairing: row i gets vx[i-1]; row 0 the carried fill_prev
15744                    // (hidden of the last committed row before this verify batch).
15745                    let mut vxs = e.zeros(t_v * n_embd)?;
15746                    e.copy_into(&mut vxs, 0, &fill_prev, n_embd)?;
15747                    if t_v > 1 {
15748                        e.copy_view_into(
15749                            &mut vxs,
15750                            n_embd,
15751                            &vx.slice(0..(t_v - 1) * n_embd),
15752                            (t_v - 1) * n_embd,
15753                        )?;
15754                    }
15755                    self.mtp_kv_fill_all(e, &verify_tokens, &vxs, pos, &mut *scratch, embd_dev)?;
15756                } else {
15757                    scratch.set_len(e, pos + base + k_round - 1)?;
15758                    // predecessor of the last draft = verify col t_v-2 (or fill_prev at t_v==1)
15759                    let mut hp = e.zeros(n_embd)?;
15760                    if t_v >= 2 {
15761                        e.copy_view_into(
15762                            &mut hp,
15763                            0,
15764                            &vx.slice((t_v - 2) * n_embd..(t_v - 1) * n_embd),
15765                            n_embd,
15766                        )?;
15767                    } else {
15768                        e.copy_into(&mut hp, 0, &fill_prev, n_embd)?;
15769                    }
15770                    self.mtp_kv_fill_all(
15771                        e,
15772                        &[draft[k_round - 1]],
15773                        &hp,
15774                        pos + base + k_round - 1,
15775                        &mut *scratch,
15776                        embd_dev,
15777                    )?;
15778                }
15779                // REFERENCE SEEDING: no pseudo pass — the next chain's step 0 IS the
15780                // reference's (id_last, h_prev) draft row; it appends the bonus's scratch
15781                // entry itself. Seed = TRUE hidden of the bonus's predecessor (last verify
15782                // col). Saves one MTP-block pass per round on top of the pairing fix.
15783                if !devacc_seeded {
15784                    e.copy_into(&mut h_seed_buf, 0, &vh_seed, n_embd)?;
15785                    e.copy_into(&mut fill_prev, 0, &vh_seed, n_embd)?;
15786                }
15787                pending = Some(bonus);
15788                if debug_spec {
15789                    eprintln!("  -> FULL ACCEPT (bonus pending, prev-h seed)");
15790                }
15791            } else if !spec_replay && base + n_acc >= 1 {
15792                // PARTIAL ACCEPT, REPLAY-FREE (2026-07-03 — the profiled #1 long-ctx spec cost):
15793                // the verify's first j = base+n_acc columns ARE the committed sequence, computed
15794                // bit-identically to eager (decode-exact contract) — so KEEP them: KV truncates to
15795                // pos+j, recurrent state rebuilds from the VerifyCkpt (same-kernel gdn prefix
15796                // re-run / pure state-clone restore), and the bonus stays PENDING exactly like the
15797                // full-accept path — the legacy duplicate trunk replay is gone. The next chain
15798                // seeds from the MTP pseudo-hidden of the bonus, whose seed = the TRUE verify
15799                // hidden of its predecessor (col j-1) — same one-hop pseudo structure as full
15800                // accept (never compounds: the next verify recomputes true hiddens for all
15801                // committed columns).
15802                let j = base + n_acc;
15803                // VERIFY-GRAPH SLAB COMMIT: when the captured trunk ran, the linear layers'
15804                // column stash was written into the graphs ctx's persistent slabs as in-graph
15805                // memcpy nodes, NOT into the per-column VerifyCkpt the cols arm reads — so the
15806                // commit must take the slab twin (same semantics, slab-addressed sources). The
15807                // ctx states which of the two this round produced via `round_slab`; trusting the
15808                // flag rather than the env keeps a round that fell back to the eager walk (a
15809                // capture that declined, a t the pool never captured) on the cols arm.
15810                let slab_commit = vg_guard
15811                    .as_ref()
15812                    .and_then(|g| g.as_ref())
15813                    .map(|g| g.round_slab)
15814                    .unwrap_or(false);
15815                if slab_commit {
15816                    self.dspark_commit_prefix_slab(
15817                        e,
15818                        &mut *cache,
15819                        &snap,
15820                        vg_guard
15821                            .as_ref()
15822                            .and_then(|g| g.as_ref())
15823                            .expect("slab_commit implies a graphs ctx"),
15824                        j,
15825                    )?;
15826                } else {
15827                    self.commit_verified_prefix(
15828                        e,
15829                        &mut *cache,
15830                        &snap,
15831                        ckpt.as_ref().unwrap(),
15832                        j,
15833                        devacc_seeded,
15834                        if devacc_seeded {
15835                            devacc_acc.as_ref().map(|a| (a, base, t_v))
15836                        } else {
15837                            None
15838                        },
15839                    )?;
15840                }
15841                let mut seed = e.zeros(n_embd)?;
15842                e.copy_view_into(
15843                    &mut seed,
15844                    0,
15845                    &vx.slice((j - 1) * n_embd..j * n_embd),
15846                    n_embd,
15847                )?;
15848                // Draft scratch: TRUE-HIDDEN REFRESH of the committed prefix (see the full-accept
15849                // branch); without it the chain entries stand and only the tail truncates. Either
15850                // way len ends at pos+j so the pseudo append lands at the bonus's slot pos+j
15851                // (persistent mode), rope pos+j+1 (chain convention).
15852                if refresh {
15853                    scratch.set_len(e, pos)?;
15854                    let mut vxs = e.zeros(j * n_embd)?;
15855                    e.copy_into(&mut vxs, 0, &fill_prev, n_embd)?;
15856                    if j > 1 {
15857                        e.copy_view_into(
15858                            &mut vxs,
15859                            n_embd,
15860                            &vx.slice(0..(j - 1) * n_embd),
15861                            (j - 1) * n_embd,
15862                        )?;
15863                    }
15864                    self.mtp_kv_fill_all(
15865                        e,
15866                        &verify_tokens[0..j],
15867                        &vxs,
15868                        pos,
15869                        &mut *scratch,
15870                        embd_dev,
15871                    )?;
15872                } else {
15873                    scratch.set_len(e, pos + j)?;
15874                }
15875                // REFERENCE SEEDING (see the full-accept branch): seed = TRUE hidden of the
15876                // bonus's predecessor (verify col j-1); no pseudo pass.
15877                if !devacc_seeded {
15878                    e.copy_into(&mut h_seed_buf, 0, &seed, n_embd)?;
15879                    e.copy_into(&mut fill_prev, 0, &seed, n_embd)?;
15880                }
15881                pending = Some(bonus);
15882                if debug_spec {
15883                    eprintln!("  -> PARTIAL(replay-free j={j}, bonus pending, prev-h seed)");
15884                }
15885            } else if !spec_replay {
15886                // ZERO ROUND FOLD (2026-07-10, verify-cost target #3): base+n_acc == 0 — a
15887                // pending-less round where nothing was accepted (PMIN0 zero-draft chains after a
15888                // replay/commit, or plain 0-accept rounds at round 0). The old path replayed
15889                // [bonus] through a FULL m=1 trunk+head forward (the 489us full-vocab head pass
15890                // measured at ~0.75/round on PMIN0 configs). Instead: restore the pre-round
15891                // snapshot and let the bonus ride the NEXT round's verify as col 0 — the existing
15892                // base=1 pending machinery, bit-identical by the decode-exact verify contract.
15893                // Seed: the bonus's predecessor is the last COMMITTED token, whose hidden
15894                // fill_prev already carries (same seeding as the 1-token-replay case it replaces).
15895                cache.rollback(e, &snap, 0)?;
15896                scratch.set_len(e, pos)?;
15897                e.copy_into(&mut h_seed_buf, 0, &fill_prev, n_embd)?;
15898                pending = Some(bonus);
15899                if debug_spec {
15900                    eprintln!("  -> ZERO-ROUND FOLD (bonus pending, fill_prev seed)");
15901                }
15902            } else {
15903                // PARTIAL ACCEPT, LEGACY REPLAY (seam MEMRA_SPEC_REPLAY=1 — or j==0: nothing of
15904                // this round survives, only possible before the first pending exists, ~round 0):
15905                // restore EVERYTHING to the pre-round snapshot (KV truncate to pos + recur
15906                // restore), then replay the committed prefix pending? ++ draft[0..n_acc] ++
15907                // [bonus] as ONE batched T forward — single weight read, bit-identical to greedy
15908                // (the verify-all-columns path is the same math). Commits the bonus with a TRUE
15909                // trunk hidden.
15910                cache.rollback(e, &snap, 0)?; // accept_len=0: KV len = pos, recur = snapshot
15911                let mut replay: Vec<u32> = Vec::with_capacity(base + n_acc + 1);
15912                if let Some(b) = pending.take() {
15913                    replay.push(b);
15914                }
15915                replay.extend_from_slice(&draft[0..n_acc]);
15916                replay.push(bonus);
15917                // Full-stack forward (decode_step_t_core = decode_step_t_h_emb_dev's body):
15918                // Predecessor pairing seeds from the PREDECESSOR row (col len-2) — the same-row path takes the
15919                // last col exactly as before (byte-identical to the old _h_emb_dev call).
15920                let (rl_d, rx) = if self.batched_serving_numeric_class() {
15921                    let mut logits = Vec::with_capacity(replay.len() * n_vocab);
15922                    let mut hidden = e.uninit(replay.len() * n_embd)?;
15923                    for (row, &token) in replay.iter().enumerate() {
15924                        let (row_logits, row_hidden) =
15925                            self.spec_target_step_h(e, token, &mut *cache)?;
15926                        logits.extend_from_slice(&row_logits);
15927                        e.dtod_copy_into(&row_hidden, &mut hidden, row * n_embd)?;
15928                    }
15929                    (e.htod(&logits)?, hidden)
15930                } else {
15931                    self.decode_step_t_core(e, &replay, pos, &mut *cache, embd_dev, None)?
15932                };
15933                // last_pred = argmax of the LAST column's logits (predicts the token after `bonus`)
15934                // — device argmax + one 4-byte read instead of the full-vocab column dtoh.
15935                e.argmax_token_device_col(&rl_d, replay.len() - 1, n_vocab, &mut preds_d, 0)?;
15936                last_pred = guard_vocab_token(
15937                    e.dtoh_u32(&preds_d)?[0],
15938                    n_vocab,
15939                    &format!("replay last_pred at round {round} pos={pos}"),
15940                )?;
15941                if sampled {
15942                    let lr0 = replay.len();
15943                    let lc = last_col_logits
15944                        .as_mut()
15945                        .expect("sampled: last_col_logits unset");
15946                    e.copy_view_into(
15947                        lc,
15948                        0,
15949                        &rl_d.slice((lr0 - 1) * n_vocab..lr0 * n_vocab),
15950                        n_vocab,
15951                    )?;
15952                }
15953                let lr = replay.len();
15954                if lr >= 2 {
15955                    e.copy_view_into(
15956                        &mut h_seed_buf,
15957                        0,
15958                        &rx.slice((lr - 2) * n_embd..(lr - 1) * n_embd),
15959                        n_embd,
15960                    )?;
15961                } else {
15962                    // 1-token replay (round-0 miss): the bonus's predecessor is the OLD
15963                    // last_token, whose own-row hidden fill_prev still holds.
15964                    e.copy_into(&mut h_seed_buf, 0, &fill_prev, n_embd)?;
15965                }
15966                // the bonus is COMMITTED here — it becomes the last committed row.
15967                let mut rh_last = e.zeros(n_embd)?;
15968                e.copy_view_into(
15969                    &mut rh_last,
15970                    0,
15971                    &rx.slice((lr - 1) * n_embd..lr * n_embd),
15972                    n_embd,
15973                )?;
15974                e.copy_into(&mut fill_prev, 0, &rh_last, n_embd)?;
15975                if debug_spec {
15976                    eprintln!("  -> PARTIAL(replay={replay:?}), next_pred={last_pred}");
15977                }
15978            }
15979            if devacc_seeded {
15980                // stage (b) epilogue: fill_prev takes the gathered seed AFTER the refresh fills
15981                // consumed the old value (both slots carry the same value in every non-replay arm).
15982                e.copy_into(&mut fill_prev, 0, &h_seed_buf, n_embd)?;
15983            }
15984            if successor_valid {
15985                let optimistic_scratch_len = successor_attempt
15986                    .as_ref()
15987                    .expect("valid controller successor disappeared")
15988                    .scratch_len;
15989                // The normal current-round commit refreshed/truncated the logical scratch tail.
15990                // Its optimistic successor row was already written physically, so restoring only
15991                // the retained logical length makes that row live for the carried round.
15992                scratch.set_len(e, optimistic_scratch_len)?;
15993            }
15994            if let Some(current) = current_opti.take() {
15995                opti_fork
15996                    .as_mut()
15997                    .ok_or("optipipe current retirement lost fork state")?
15998                    .retire(current.generation)?;
15999            }
16000            if successor_valid {
16001                let successor = successor_attempt
16002                    .take()
16003                    .expect("valid controller successor disappeared before promotion");
16004                let generation = successor.generation;
16005                opti_fork
16006                    .as_mut()
16007                    .ok_or("optipipe successor promotion lost fork state")?
16008                    .promote_successor_snapshot(&mut snap, generation);
16009                carried_opti = Some(successor);
16010            }
16011            if anatomy_on {
16012                // Commit/rollback is normally asynchronous on the primary/head stream. Bound it
16013                // only for this diagnostic so it does not disappear into the following draft's
16014                // first token readback.
16015                e.stream().synchronize()?;
16016                ph_commit += commit_started.elapsed().as_secs_f64();
16017            }
16018            // adaptive-K update (host math, zero syncs): next round drafts accepted-run + 1,
16019            // clamped to [floor(pos), k_cap]. cache.pos is post-rollback here (the round's
16020            // final position — the floor's position key reads the committed depth). Burst
16021            // rounds (`continue` above) draft the captured fixed depth and skip this, exactly
16022            // like gemma's burst arm.
16023            if adapt {
16024                let fl_now = floor_at(cache.pos);
16025                kc = (n_acc + 1).clamp(fl_now.min(k_cap), k_cap);
16026            }
16027            ph_mark(&mut ph_rest, phase_on);
16028            if let Some(p) = pipe {
16029                p.accept_end(round);
16030            }
16031            drop(pipe_accept);
16032            if let Some(t0) = round_t0 {
16033                let ms = t0.elapsed().as_secs_f64() * 1e3;
16034                ROUND_MS.fetch_add((ms * 1e3) as u64, std::sync::atomic::Ordering::Relaxed);
16035                let n = ROUND_N.fetch_add(1, std::sync::atomic::Ordering::Relaxed) + 1;
16036                if n.is_multiple_of(32) {
16037                    eprintln!(
16038                        "[spec-round] rounds={n} avg round wall={:.2} ms (emitted={} drafted so far)",
16039                        ROUND_MS.load(std::sync::atomic::Ordering::Relaxed) as f64 / 1e3 / n as f64,
16040                        out.len()
16041                    );
16042                }
16043            }
16044            round += 1;
16045            // sse-cadence: this round's accepted drafts + bonus are committed (out is
16046            // append-only past step 4) — flush at round cadence.
16047            keep_going = flush_commit(&mut on_commit, &out, &mut flushed);
16048        }
16049        if let Some(mut ticket) = carried_opti.take() {
16050            opti_fork
16051                .as_mut()
16052                .ok_or("optipipe tail drain lost fork state")?
16053                .cancel_controller_ticket(e, &mut *cache, &mut *scratch, &snap, &mut ticket)?;
16054        }
16055        // sse-cadence: nothing below appends to `out`; flush any remainder (defensive).
16056        // (verdict ignored — the burst is over either way; the session tail runs unchanged.)
16057        let _ = flush_commit(&mut on_commit, &out, &mut flushed);
16058
16059        if spec_stats {
16060            let per_slot: Vec<String> = (0..k)
16061                .map(|j| {
16062                    if st_drafted[j] > 0 {
16063                        format!(
16064                            "{}/{}={:.3}",
16065                            st_accepted[j],
16066                            st_drafted[j],
16067                            st_accepted[j] as f64 / st_drafted[j] as f64
16068                        )
16069                    } else {
16070                        "0/0".into()
16071                    }
16072                })
16073                .collect();
16074            let acc = if total_drafted > 0 {
16075                total_accepted as f64 / total_drafted as f64
16076            } else {
16077                0.0
16078            };
16079            eprintln!(
16080                "[spec-stats] rounds={round} full_accept={st_full} len_hist={st_len_hist:?} \
16081                       per_slot=[{}] total={total_accepted}/{total_drafted}={acc:.3} \
16082                       tok_per_round={:.3}",
16083                per_slot.join(" "),
16084                (total_accepted + round) as f64 / round.max(1) as f64
16085            );
16086        }
16087        if constraint.is_some() {
16088            eprintln!(
16089                "[draft-mask] mask_rounds={dm_rounds} clone_total={:.3}ms \
16090                 clone_per_round={:.4}ms gram_cuts={dm_cuts}/{round} cut_tokens={dm_cut_tokens}",
16091                dm_clone_ns as f64 / 1e6,
16092                dm_clone_ns as f64 / 1e6 / dm_rounds.max(1) as f64
16093            );
16094        }
16095        if phase_on {
16096            let tot = ph_draft + ph_verify + ph_wait + ph_rest;
16097            eprintln!(
16098                "[spec-phase] draft={:.1}ms ({:.1}%) verify-issue={:.1}ms ({:.1}%) verify-wait={:.1}ms ({:.1}%) commit-host={:.1}ms ({:.1}%) rounds={round}",
16099                ph_draft * 1e3,
16100                ph_draft / tot * 100.0,
16101                ph_verify * 1e3,
16102                ph_verify / tot * 100.0,
16103                ph_wait * 1e3,
16104                ph_wait / tot * 100.0,
16105                ph_rest * 1e3,
16106                ph_rest / tot * 100.0
16107            );
16108        }
16109        if anatomy_on {
16110            let rounds_f = round.max(1) as f64;
16111            let other = (ph_rest - ph_commit).max(0.0);
16112            eprintln!(
16113                "[spec-anatomy] per-round draft={:.3}ms pp-verify={:.3}ms \
16114                 verify-accept={:.3}ms commit-rollback={:.3}ms other={:.3}ms rounds={round}",
16115                ph_draft * 1e3 / rounds_f,
16116                ph_verify * 1e3 / rounds_f,
16117                ph_wait * 1e3 / rounds_f,
16118                ph_commit * 1e3 / rounds_f,
16119                other * 1e3 / rounds_f,
16120            );
16121        }
16122        let _pipe_tail = pipe.map(|p| p.primary()).transpose()?;
16123        // SESSION TAIL: leave the session in the exact invariant the next turn's suffix prime
16124        // expects — every row in `committed` has trunk KV/recur state AND an exact draft-KV row.
16125        // Park the draft-graph ctx back on the session (the serve-burst fixed-cost fix): the next
16126        // burst replays instead of recapturing. Error paths (`?` above) drop it — recaptured then.
16127        if let Some(slot) = sess_draft_slot.take() {
16128            *slot = Some(dctx);
16129        }
16130        let t_rounds = t_ent.elapsed();
16131        if let Some((committed, last_h, next_pred_slot, sctr_slot, uctr_slot)) = sess_tail.take() {
16132            // NEXT BURST'S BOUNDARY TOKEN (lane/sampled-spec-quality, Item 1). Greedy stashes
16133            // the argmax `last_pred` exactly as before (byte contract). SAMPLED draws the token
16134            // HERE, where the sampler, the session Philox counters and the penalty window are
16135            // all live and the boundary logits row still exists — that is the "make the state
16136            // available" half of the fix; the consuming burst then just emits it. `sctr` is
16137            // written to the session BELOW the draws so the advance is never lost.
16138            *next_pred_slot = Some(last_pred);
16139            let sample_boundary = sampled && constraint.is_none() && spec_sampled_boundary_on();
16140            let mut stashed_pending = false;
16141            if let Some(b) = pending.take() {
16142                if !sampled {
16143                    // PENDING-CARRY (2026-08-01): stash the bonus on the session instead of
16144                    // committing it with a solo T=1 pass — the next empty-suffix greedy burst
16145                    // consumes it as round-0 verify col 0 (a plain round edge; the old tail
16146                    // commit + next burst's init feed were 11.6+11.5ms solo trunk passes per
16147                    // burst on H100 q27, [spec-setup] trace). b stays in `out` (emitted) but
16148                    // OUT of `committed` (cache rows == committed); the consuming call
16149                    // prepends it once its verify commits the row. next_pred is unknowable
16150                    // without the commit pass — None; callers gate on pending_tok too.
16151                    debug_assert_eq!(out.last(), Some(&b), "pending must be the last emitted");
16152                    if let Some(slot) = sess_pending_slot.take() {
16153                        *slot = Some(b);
16154                    }
16155                    *next_pred_slot = None;
16156                    // fill_prev = hidden of the last COMMITTED row (b's predecessor) — the
16157                    // exact chain-seed/fill anchor the consuming burst (or a flush) needs.
16158                    *last_h = Some(e.clone_dtod(&fill_prev)?);
16159                    stashed_pending = true;
16160                } else {
16161                    // SAMPLED tail (unchanged): commit the bonus (one T=1 pass) + draft fill —
16162                    // the sampled round-0 accept needs this pass's logits (last_col_logits).
16163                    let pos_b = cache.pos;
16164                    scratch.set_len(e, pos_b)?;
16165                    let (lg_b, hb) = self.spec_target_step_h(e, b, &mut *cache)?;
16166                    // after a FULL-accept exit `last_pred` is STALE (it predicted the bonus
16167                    // itself — the prediction AFTER the bonus never materialized; it would have
16168                    // been the next round's verify col 0). The commit's logits ARE that
16169                    // prediction — so they are also the row the next burst's boundary token
16170                    // comes off, and (lane/sampled-spec-quality) it is DRAWN from them here.
16171                    *next_pred_slot = Some(if sample_boundary {
16172                        sample_boundary_token(
16173                            e,
16174                            &lg_b,
16175                            &sp,
16176                            &pen_hist,
16177                            &mut sctr,
16178                            "burst-tail-commit",
16179                        )?
16180                    } else {
16181                        argmax(&lg_b) as u32
16182                    });
16183                    self.mtp_kv_fill_all(e, &[b], &fill_prev, pos_b, &mut *scratch, embd_dev)?;
16184                    *last_h = Some(hb);
16185                }
16186            } else {
16187                // fill_prev tracks the hidden of the last COMMITTED row throughout the loop.
16188                *last_h = Some(e.clone_dtod(&fill_prev)?);
16189                if sample_boundary {
16190                    // No pending to commit, so the boundary row is the one `last_pred` was
16191                    // argmaxed from and the sampled path keeps it on device: the init feed's
16192                    // logits when the burst ran zero rounds, else the legacy-replay path's
16193                    // last verify column (both predict the token AFTER the last committed
16194                    // row). It is retained precisely because round 0's accept test needs it,
16195                    // so the draw costs no extra D2H of the [n_vocab] row.
16196                    match last_col_logits.as_ref() {
16197                        Some(lc) => {
16198                            *next_pred_slot = Some(sample_boundary_token_dev(
16199                                e,
16200                                lc,
16201                                n_vocab,
16202                                &sp,
16203                                &pen_hist,
16204                                &mut sctr,
16205                                "burst-tail-nopending",
16206                            )?);
16207                        }
16208                        // NAME THE FALLBACK (house standard): unreachable today — a sampled
16209                        // burst always feeds or replays, so the row exists — but if it ever
16210                        // is, the stream takes a greedy token and SAYS so rather than
16211                        // silently regressing to the pre-lane behaviour.
16212                        None => eprintln!(
16213                            "[spec-boundary] sampled tail kept the ARGMAX boundary token \
16214                             (reason: no retained boundary logits row)"
16215                        ),
16216                    }
16217                }
16218            }
16219            *sctr_slot = sctr;
16220            *uctr_slot = uctr;
16221            committed.extend_from_slice(prompt);
16222            if let Some(cb) = carried_pending {
16223                // the consumed carry's cache row landed in round 0's verify (every pending
16224                // round commits col 0) — it joins `committed` here, in sequence order.
16225                committed.push(cb);
16226            }
16227            if stashed_pending {
16228                // ZERO-EMIT BURST (2026-08-06 c=8 serve panic, pre-existing since b4aea184):
16229                // `out.len() - 1` underflowed on an EMPTY `out` — "range end index
16230                // 18446744073709551615 out of range for slice of length 0", killing the
16231                // memra-gpu-worker and failing 31 of 32 concurrent requests with "worker closed
16232                // stream". Reachable because `pending` starts as `carried_pending` (a bonus
16233                // stashed by the PREVIOUS burst) while `out` starts empty, and the carry is
16234                // deliberately NOT pushed to `out` (line ~3239: the burst that emitted it already
16235                // did). So a burst that stashes a pending without emitting anything of its own —
16236                // the round loop exits before a push, e.g. the ring drain's `out.len() < max_new`
16237                // guard skipping every token under a tight budget — arrives here with
16238                // out.len() == 0 and stashed_pending == true.
16239                //
16240                // The invariant is unchanged: `committed` gets every emitted token EXCEPT the
16241                // stashed bonus. With nothing emitted, that is nothing — and the carry pushed
16242                // just above is already accounted. Saturating, not a min/assert: an empty `out`
16243                // here is a legitimate burst shape, not a corrupt state.
16244                let emitted = out.len().saturating_sub(1);
16245                committed.extend_from_slice(&out[..emitted]);
16246            } else {
16247                committed.extend_from_slice(&out); // FULL out incl. overshoot — all committed
16248            }
16249            debug_assert_eq!(
16250                cache.pos,
16251                committed.len(),
16252                "session invariant: cache rows == committed tokens"
16253            );
16254            if setup_trace {
16255                e.stream().synchronize()?; // bound the async tail fill in the trace
16256                let t_tail = t_ent.elapsed();
16257                eprintln!(
16258                    "[spec-setup] init={:.2}ms cap={:.2}ms fill={:.2}ms rounds={:.2}ms tail={:.2}ms total={:.2}ms out={} cont={}",
16259                    t_init.as_secs_f64() * 1e3,
16260                    (t_cap - t_init).as_secs_f64() * 1e3,
16261                    (t_fill - t_cap).as_secs_f64() * 1e3,
16262                    (t_rounds - t_fill).as_secs_f64() * 1e3,
16263                    (t_tail - t_rounds).as_secs_f64() * 1e3,
16264                    t_tail.as_secs_f64() * 1e3,
16265                    out.len(),
16266                    continuation
16267                );
16268            }
16269            return Ok((out, total_drafted, total_accepted));
16270        }
16271        out.truncate(max_new);
16272        Ok((out, total_drafted, total_accepted))
16273    }
16274
16275    /// Anchor-bounded DSpark target extraction. The trunk sees the exact generated token tape;
16276    /// only requested hidden rows and target-logit rows cross PCIe. An anchor token at p pairs
16277    /// with the pre-output-norm h[p-1] carrier, exactly as the existing replay/NextN path does.
16278    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
16279    pub fn extract_dspark_anchors(
16280        &self,
16281        e: &Engine,
16282        tokens: &[u32],
16283        anchor_positions: &[usize],
16284        gamma: usize,
16285        top_k: usize,
16286        chunk: usize,
16287        temperature: f32,
16288    ) -> Result<Vec<DsparkAnchorRecord>, Box<dyn std::error::Error>> {
16289        if tokens.len() < gamma + 2 || gamma == 0 || chunk < 2 {
16290            return Err("DSpark extraction token tape/gamma/chunk is invalid".into());
16291        }
16292        if anchor_positions.windows(2).any(|pair| pair[0] >= pair[1]) {
16293            return Err("DSpark anchor positions must be sorted and unique".into());
16294        }
16295        for &position in anchor_positions {
16296            if position == 0 || position + gamma >= tokens.len() {
16297                return Err(format!(
16298                    "DSpark anchor {position} has no predecessor or cannot cover gamma={gamma} in {} tokens",
16299                    tokens.len()
16300                )
16301                .into());
16302            }
16303        }
16304
16305        let n_vocab = self.output.out_features();
16306        let n_embd = self.cfg.n_embd as usize;
16307        let mut cache =
16308            crate::pp::new_cache_planned(e, &self.cfg, &self.plan, tokens.len() + gamma + 8)?;
16309        let (embd_qt, embd_rb) = self.embd.qt_and_row_bytes(n_embd);
16310        let embd_gpu = if spec_host_embd() {
16311            None
16312        } else {
16313            Some(
16314                self.embd_gpu
16315                    .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload")),
16316            )
16317        };
16318        let embd_dev = embd_gpu.map(|gpu| (gpu, embd_qt, embd_rb));
16319
16320        struct PendingRecord {
16321            position: usize,
16322            hidden: Option<Vec<f32>>,
16323            tokens: Vec<u32>,
16324            target_top_ids: Vec<Option<Vec<u32>>>,
16325            target_top_logits: Vec<Option<Vec<f32>>>,
16326            target_top_probs: Vec<Option<Vec<f32>>>,
16327            target_tail_probs: Vec<Option<f32>>,
16328        }
16329
16330        let mut pending: Vec<PendingRecord> = anchor_positions
16331            .iter()
16332            .map(|&position| PendingRecord {
16333                position,
16334                hidden: None,
16335                tokens: tokens[position..=position + gamma].to_vec(),
16336                target_top_ids: vec![None; gamma],
16337                target_top_logits: vec![None; gamma],
16338                target_top_probs: vec![None; gamma],
16339                target_tail_probs: vec![None; gamma],
16340            })
16341            .collect();
16342
16343        let mut start = 0usize;
16344        while start < tokens.len() {
16345            let end = (start + chunk).min(tokens.len());
16346            let chunk_tokens = &tokens[start..end];
16347            let (target_logits, hidden_rows) =
16348                self.decode_step_t_core(e, chunk_tokens, start, &mut cache, embd_dev, None)?;
16349            for record in &mut pending {
16350                let hidden_position = record.position - 1;
16351                if hidden_position >= start && hidden_position < end {
16352                    let local = hidden_position - start;
16353                    record.hidden = Some(
16354                        e.dtoh_view(&hidden_rows.slice(local * n_embd..(local + 1) * n_embd))?,
16355                    );
16356                }
16357                for slot in 0..gamma {
16358                    let target_row = record.position + slot;
16359                    if target_row < start || target_row >= end {
16360                        continue;
16361                    }
16362                    let local = target_row - start;
16363                    let logits =
16364                        e.dtoh_view(&target_logits.slice(local * n_vocab..(local + 1) * n_vocab))?;
16365                    let (ids, top_logits, probs, tail) =
16366                        dspark_sparse_softmax_topk(&logits, top_k, temperature)?;
16367                    record.target_top_ids[slot] = Some(ids);
16368                    record.target_top_logits[slot] = Some(top_logits);
16369                    record.target_top_probs[slot] = Some(probs);
16370                    record.target_tail_probs[slot] = Some(tail);
16371                }
16372            }
16373            start = end;
16374        }
16375
16376        pending
16377            .into_iter()
16378            .map(|record| {
16379                let hidden = record
16380                    .hidden
16381                    .ok_or_else(|| format!("missing DSpark hidden at {}", record.position))?;
16382                let target_top_ids =
16383                    flatten_dspark_rows(record.target_top_ids, record.position, "target ids")?;
16384                let target_top_logits = flatten_dspark_rows(
16385                    record.target_top_logits,
16386                    record.position,
16387                    "target logits",
16388                )?;
16389                let target_top_probs =
16390                    flatten_dspark_rows(record.target_top_probs, record.position, "target probs")?;
16391                let target_tail_probs = record
16392                    .target_tail_probs
16393                    .into_iter()
16394                    .enumerate()
16395                    .map(|(slot, value)| {
16396                        value.ok_or_else(|| {
16397                            format!("missing DSpark tail at {} slot {slot}", record.position)
16398                        })
16399                    })
16400                    .collect::<Result<Vec<_>, _>>()?;
16401                Ok(DsparkAnchorRecord {
16402                    position: record.position,
16403                    hidden,
16404                    tokens: record.tokens,
16405                    target_top_ids,
16406                    target_top_logits,
16407                    target_top_probs,
16408                    target_tail_probs,
16409                })
16410            })
16411            .collect()
16412    }
16413
16414    /// TEACHER-FORCED REPLAY ACCEPTANCE (hqmtp MTP-heal protocol): walk a FIXED token
16415    /// sequence and, at sampled positions, compare the MTP head's K-token draft chain against
16416    /// the trunk's own teacher-forced greedy predictions. Nothing is generated — the context is
16417    /// the corpus text itself, so (a) degenerate self-generated loops cannot inflate acceptance
16418    /// and (b) two arms (bf16 ceiling vs NVFP4) score on IDENTICAL contexts, isolating the
16419    /// quant-induced head/hidden-state mismatch from text drift.
16420    ///
16421    /// Per eval position p (context = tokens[0..=p], predecessor pairing as in spec decode):
16422    ///   draft_j  = chain token j from (tokens[p], h_{p-1}), then its own drafts — the exact
16423    ///              eager spec-decode chain (same mtp_head_forward_dev, same rope positions).
16424    ///   target_j = teacher-forced greedy pick for position p+1+j (argmax of the trunk logits
16425    ///              at forced context tokens[0..p+j]). For j==0 this equals live spec
16426    ///              acceptance; for j>=1 live verify would condition on the drafts, here it
16427    ///              conditions on the corpus — deterministic and arm-comparable by design.
16428    ///
16429    /// Returns (rows, bg): one (p, drafts[k], targets[k]) row per eval position (ascending p),
16430    /// plus the full teacher-forced greedy track bg (bg[i] = greedy pick for position i, i>=1)
16431    /// so harnesses can cross-check runs (e.g. different chunk sizes must give identical bg).
16432    ///
16433    /// `hdump`: when Some, every position's pre-output_norm trunk hidden (the exact rows the
16434    /// draft-KV fill pairs from) streams to the file as little-endian f32 [t_total, n_embd] —
16435    /// the head-distillation extraction (hqmtp): the ENGINE is the source of truth for trunk
16436    /// hiddens (HF torch reproductions of the hybrid trunk measured only ~0.5 greedy
16437    /// agreement vs this path — not usable as a training-data source).
16438    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
16439    pub fn replay_acceptance(
16440        &self,
16441        e: &Engine,
16442        tokens: &[u32],
16443        k: usize,
16444        stride: usize,
16445        chunk: usize,
16446        mut hdump: Option<&mut std::fs::File>,
16447    ) -> Result<(Vec<(usize, Vec<u32>, Vec<u32>)>, Vec<u32>), Box<dyn std::error::Error>> {
16448        assert!(k >= 1 && stride >= 1 && chunk >= 2);
16449        let mtp = self
16450            .mtp
16451            .as_ref()
16452            .expect("replay_acceptance requires an MTP head");
16453        let n_vocab = self.output.out_features();
16454        let d_vocab = mtp
16455            .shared_head_head
16456            .as_ref()
16457            .unwrap_or(&self.output)
16458            .out_features();
16459        let n_embd = self.cfg.n_embd as usize;
16460        let t_total = tokens.len();
16461        assert!(t_total >= 8, "corpus too short ({t_total} tokens)");
16462        // STAGE-OWNED KV (lane/pp2-spec 2026-08-06) — see `new_session`. Door shut = `Cache::new`.
16463        let mut cache = crate::pp::new_cache_planned(e, &self.cfg, &self.plan, t_total + k + 8)?;
16464        let mut scratch = self.new_mtp_scratch(e, t_total + k + 8)?;
16465        let (embd_qt, embd_rb) = self.embd.qt_and_row_bytes(n_embd);
16466        let embd_gpu = if spec_host_embd() {
16467            None
16468        } else {
16469            Some(
16470                self.embd_gpu
16471                    .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload")),
16472            )
16473        };
16474        let embd_dev = embd_gpu.map(|g| (g, embd_qt, embd_rb));
16475
16476        // bg[i] = the trunk's greedy pick for position i under the forced context (i >= 1).
16477        let mut bg: Vec<u32> = vec![0; t_total + 1];
16478        let mut rows: Vec<(usize, Vec<u32>, Vec<u32>)> = Vec::new();
16479        let mut prev_last_h = e.zeros(n_embd)?; // predecessor hidden entering the chunk
16480        let mut seed_buf = e.zeros(n_embd)?;
16481        let mut preds_d = e.alloc_u32_zeroed(chunk)?;
16482        let nll_on = std::env::var("MEMRA_REPLAY_NLL").as_deref() == Ok("1");
16483        let (mut nll_sum, mut nll_cnt) = (0f64, 0u64);
16484        let mut s = 0usize;
16485        while s < t_total {
16486            let cend = (s + chunk).min(t_total);
16487            let tc = cend - s;
16488            let ch = &tokens[s..cend];
16489            // 1. forced trunk pass — verify path (decode-exact contract): all-column logits +
16490            //    the chunk's true hiddens.
16491            let (tl_d, vx) = self.decode_step_t_core(e, ch, s, &mut cache, embd_dev, None)?;
16492            for j in 0..tc {
16493                e.argmax_token_device_col(&tl_d, j, n_vocab, &mut preds_d, j)?;
16494            }
16495            let preds = e.dtoh_u32(&preds_d)?;
16496            for j in 0..tc {
16497                bg[s + j + 1] = preds[j];
16498            }
16499            // MEMRA_REPLAY_NLL=1: teacher-forced NLL/perplexity over the same forced pass — the
16500            // checkpoint-quality metric (position j's logits score the GOLD next token).
16501            if nll_on {
16502                let jmax = if cend < t_total { tc } else { tc - 1 }; // last pos has no gold next
16503                if jmax > 0 {
16504                    let ids: Vec<u32> = (0..jmax).map(|j| tokens[s + j + 1]).collect();
16505                    let rows: Vec<i32> = (0..jmax as i32).collect();
16506                    let idsd = e.htod_u32_v(&ids)?;
16507                    let rowsd = e.htod_i32(&rows)?;
16508                    let mut outd = e.zeros(jmax)?;
16509                    e.softmax_gather(&tl_d, n_vocab, &idsd, &rowsd, &mut outd, n_vocab, jmax, 1.0)?;
16510                    for pr in e.dtoh(&outd)? {
16511                        nll_sum += -((pr.max(1e-30)) as f64).ln();
16512                        nll_cnt += 1;
16513                    }
16514                }
16515            }
16516            if let Some(f) = hdump.as_deref_mut() {
16517                use std::io::Write;
16518                let host: Vec<f32> = e.dtoh(&vx)?;
16519                // bf16 round-to-nearest-even — f32 doubled the disk bill at bulk
16520                // extraction scale (20M tokens x 4096 = 320GB f32 vs 160GB bf16).
16521                let mut bytes = Vec::with_capacity(tc * n_embd * 2);
16522                for v in &host[..tc * n_embd] {
16523                    let b = v.to_bits();
16524                    let r = b.wrapping_add(0x7FFF + ((b >> 16) & 1));
16525                    bytes.extend_from_slice(&((r >> 16) as u16).to_le_bytes());
16526                }
16527                f.write_all(&bytes)?;
16528            }
16529            // CHAINLESS extraction (stride > corpus, the bulk-hdump mode): no chunk ever
16530            // drafts, so the draft-KV fills are pure waste — skip them (2 MTP-block passes
16531            // per token saved; the forced trunk pass + hdump is all the mode needs).
16532            let chainless = stride > t_total;
16533            if chainless {
16534                e.copy_view_into(
16535                    &mut prev_last_h,
16536                    0,
16537                    &vx.slice((tc - 1) * n_embd..tc * n_embd),
16538                    n_embd,
16539                )?;
16540                s = cend;
16541                continue;
16542            }
16543            // 2. TRUE predecessor-paired draft-KV fill for the chunk (row i carries h_{i-1};
16544            //    row s reads the previous chunk's last true hidden, zeros at corpus start).
16545            let mut vxs = e.zeros(tc * n_embd)?;
16546            e.copy_into(&mut vxs, 0, &prev_last_h, n_embd)?;
16547            if tc > 1 {
16548                e.copy_view_into(
16549                    &mut vxs,
16550                    n_embd,
16551                    &vx.slice(0..(tc - 1) * n_embd),
16552                    (tc - 1) * n_embd,
16553                )?;
16554            }
16555            scratch.set_len(e, s)?;
16556            self.mtp_kv_fill_all(e, ch, &vxs, s, &mut scratch, embd_dev)?;
16557            // 3. draft chains at sampled positions, DESCENDING: a chain reads only slots
16558            //    [0..p) (true fills) and appends at >= p; the next (smaller-p) chain's set_len
16559            //    truncates those approximate appends before they can ever be read.
16560            let ps: Vec<usize> = (s..cend)
16561                .filter(|p| *p >= 1 && *p % stride == 0 && *p + k <= t_total)
16562                .collect();
16563            for &p in ps.iter().rev() {
16564                scratch.set_len(e, p)?;
16565                if p == s {
16566                    e.copy_into(&mut seed_buf, 0, &prev_last_h, n_embd)?;
16567                } else {
16568                    e.copy_view_into(
16569                        &mut seed_buf,
16570                        0,
16571                        &vx.slice((p - 1 - s) * n_embd..(p - s) * n_embd),
16572                        n_embd,
16573                    )?;
16574                }
16575                let mut e_tok = tokens[p];
16576                let mut d_seed = e.clone_dtod(&seed_buf)?;
16577                let chain_heads = !self.mtp_extra.is_empty();
16578                let mut chain_tokens = if chain_heads {
16579                    vec![tokens[p]]
16580                } else {
16581                    Vec::new()
16582                };
16583                let mut chain_seeds = if chain_heads {
16584                    vec![e.clone_dtod(&seed_buf)?]
16585                } else {
16586                    Vec::new()
16587                };
16588                let mut drafts: Vec<u32> = Vec::with_capacity(k);
16589                for j in 0..k {
16590                    let (dl_d, h_nextn) = if chain_heads {
16591                        self.mtp_chain_forward_dev(
16592                            e,
16593                            &chain_tokens,
16594                            &chain_seeds,
16595                            &mut scratch,
16596                            p,
16597                            embd_dev,
16598                            None,
16599                        )?
16600                    } else {
16601                        self.mtp_head_forward_dev(
16602                            e,
16603                            mtp,
16604                            e_tok,
16605                            &d_seed,
16606                            &mut scratch,
16607                            p + 1 + j,
16608                            embd_dev,
16609                            None,
16610                        )?
16611                    };
16612                    let tok_d = e.argmax_token_device(&dl_d, d_vocab)?;
16613                    let idx = e.dtoh_u32_one(&tok_d)?;
16614                    let d = match &mtp.d2t {
16615                        Some(map) => map[idx as usize],
16616                        None => idx,
16617                    };
16618                    drafts.push(d);
16619                    if chain_heads {
16620                        chain_tokens.push(d);
16621                        chain_seeds.push(h_nextn);
16622                    } else {
16623                        e_tok = d;
16624                        d_seed = h_nextn;
16625                    }
16626                }
16627                // targets may live in a LATER chunk's bg — resolved after the walk.
16628                rows.push((p, drafts, Vec::new()));
16629            }
16630            // 4. restore TRUE entries for the whole chunk (the next chunk's chains and fills
16631            //    expect scratch.len == cend with exact rows).
16632            scratch.set_len(e, s)?;
16633            self.mtp_kv_fill_all(e, ch, &vxs, s, &mut scratch, embd_dev)?;
16634            e.copy_view_into(
16635                &mut prev_last_h,
16636                0,
16637                &vx.slice((tc - 1) * n_embd..tc * n_embd),
16638                n_embd,
16639            )?;
16640            s = cend;
16641        }
16642        for (p, drafts, targets) in rows.iter_mut() {
16643            for j in 0..drafts.len() {
16644                targets.push(bg[*p + 1 + j]);
16645            }
16646        }
16647        rows.sort_by_key(|r| r.0);
16648        if nll_cnt > 0 {
16649            let mean = nll_sum / nll_cnt as f64;
16650            println!(
16651                "[replay-nll] tokens={nll_cnt} nll/token={mean:.5} ppl={:.4}",
16652                mean.exp()
16653            );
16654        }
16655        Ok((rows, bg))
16656    }
16657}
16658
16659#[cfg(test)]
16660mod vg_debt_tests {
16661    use super::dspark_vg_debt_projection;
16662
16663    /// TOOTH for the verify-graph admission accounting: the pool's projected remaining
16664    /// growth must be charged (pre-fix, admission charged 0 for a pool measured at
16665    /// 8,852 MiB), the projection must price the MARGINAL cost of one more key rather than
16666    /// extrapolating the pool's one-time shared allocation, and the doors that make growth
16667    /// impossible must zero the debt.
16668    #[test]
16669    fn vg_debt_projects_remaining_growth_and_respects_the_freeze_valves() {
16670        const MIB: usize = 1 << 20;
16671        let d = dspark_vg_debt_projection;
16672        // cold pool: nothing observed, one capture fits inside SPEC_SHRINK_RESERVE.
16673        assert_eq!(d(0, 256, 0, None), 0);
16674        // freeze valve MEMRA_DSPARK_VG_MAX=0: the pool cannot grow.
16675        assert_eq!(d(10, 0, 500 * MIB, None), 0);
16676        // saturated pool: at/past the cap the pool FREEZES, nothing left to reserve.
16677        assert_eq!(d(256, 256, 8852 * MIB, None), 0);
16678        assert_eq!(d(300, 256, 8852 * MIB, None), 0);
16679
16680        // BOOTSTRAP (one observation, growth unmeasurable): at most one more pool's worth.
16681        // The pre-fix mean rule extrapolated 255x here — the measured 8.5 GB phantom.
16682        assert_eq!(d(1, 256, 33 * MIB, None), 33 * MIB);
16683
16684        // MARGINAL, flat pool (the box9 receipt: reserved stayed ~33.6 MiB across captures
16685        // 1..3, so an additional key costs ~nothing and the debt must collapse to ~0 —
16686        // NOT the 8,556/4,261/2,830 MB the mean rule printed).
16687        assert_eq!(d(3, 256, 33 * MIB, Some((1, 33 * MIB))), 0);
16688
16689        // MARGINAL, genuinely growing pool: 40 MiB per new key over 2 keys, 250 slots left.
16690        let debt = d(6, 256, 273 * MIB, Some((4, 193 * MIB)));
16691        assert_eq!(debt, 250 * (40 * MIB));
16692        assert!(
16693            debt > 3 * (1536 * MIB),
16694            "real growth must dwarf SPEC_SHRINK_RESERVE"
16695        );
16696
16697        // a shrinking/recycled reading never becomes a negative charge.
16698        assert_eq!(d(6, 256, 10 * MIB, Some((4, 99 * MIB))), 0);
16699        // a stale observation at the same capture count falls back to bootstrap.
16700        assert_eq!(d(4, 256, 80 * MIB, Some((4, 80 * MIB))), 80 * MIB);
16701    }
16702}
16703
16704#[cfg(test)]
16705mod capture_headroom_tests {
16706    use super::{
16707        CAPTURE_HEADROOM_FLOOR, capture_err_is_oom, capture_headroom_verdict,
16708        draft_capture_bootstrap_estimate,
16709    };
16710
16711    /// TOOTH for the pre-capture reserve check (lane/step37-vram-admission-20260830): a
16712    /// capture attempt must be refused BEFORE it allocates when the device cannot cover its
16713    /// appetite plus the post-capture floor — and pool-cached bytes count as headroom
16714    /// (driver `free` alone under-counts, the wrong direction for a gate that drops
16715    /// coverage).
16716    #[test]
16717    fn capture_reserve_check_refuses_short_devices_and_counts_pool_cache() {
16718        const MIB: usize = 1 << 20;
16719        let need = 900 * MIB;
16720        // Plenty of room: no refusal.
16721        assert_eq!(
16722            capture_headroom_verdict(8_000 * MIB, 0, need, CAPTURE_HEADROOM_FLOOR),
16723            None
16724        );
16725        // The owner's shape: capture appetite would walk the card to the edge — refused,
16726        // with the arithmetic surfaced for the WARN line.
16727        let (required, effective) =
16728            capture_headroom_verdict(1_200 * MIB, 0, need, CAPTURE_HEADROOM_FLOOR)
16729                .expect("short device must refuse");
16730        assert_eq!(required, need + CAPTURE_HEADROOM_FLOOR);
16731        assert_eq!(effective, 1_200 * MIB);
16732        // Pool-cached bytes are real headroom (the trim path makes them driver-visible).
16733        assert_eq!(
16734            capture_headroom_verdict(1_200 * MIB, 7_000 * MIB, need, CAPTURE_HEADROOM_FLOOR),
16735            None
16736        );
16737        // Boundary: exactly enough is enough (>=, never a fencepost refusal).
16738        assert_eq!(
16739            capture_headroom_verdict(
16740                need + CAPTURE_HEADROOM_FLOOR,
16741                0,
16742                need,
16743                CAPTURE_HEADROOM_FLOOR
16744            ),
16745            None
16746        );
16747        // POLICY at the call site (owner-shape receipts, escalated twice on-box): the
16748        // refusal fn is handed 2x the appetite plus TWO floors — a capture may take at
16749        // most half the discretionary headroom, so the card retains a whole capture's
16750        // worth of room after it lands. One floor of slack above one appetite (the shape
16751        // that step-OOM'd on the owner cell) must therefore REFUSE under the call-site
16752        // requirement.
16753        assert!(
16754            capture_headroom_verdict(
16755                need + CAPTURE_HEADROOM_FLOOR + (100 << 20),
16756                0,
16757                2 * need,
16758                CAPTURE_HEADROOM_FLOOR * 2
16759            )
16760            .is_some()
16761        );
16762    }
16763
16764    #[test]
16765    fn bootstrap_estimate_scales_with_heads_and_never_underflows() {
16766        // 3-head chain on a step37-shaped vocab must expect strictly more than one head.
16767        let one = draft_capture_bootstrap_estimate(1, 3, 128_896, 4_096);
16768        let three = draft_capture_bootstrap_estimate(3, 3, 128_896, 4_096);
16769        assert!(three > one);
16770        // Degenerate shapes keep a sane minimum (the estimate feeds a refusal gate; a
16771        // zero-need gate refuses nothing).
16772        assert!(draft_capture_bootstrap_estimate(0, 0, 0, 0) >= 64 << 20);
16773    }
16774
16775    #[test]
16776    fn capture_oom_predicate_matches_the_quoted_driver_text() {
16777        assert!(capture_err_is_oom(
16778            "DriverError(CUDA_ERROR_OUT_OF_MEMORY, \"out of memory\")"
16779        ));
16780        assert!(capture_err_is_oom("allocation failed: out of memory"));
16781        assert!(!capture_err_is_oom("capture produced no graph"));
16782    }
16783}
16784
16785#[cfg(test)]
16786mod mtp_chain_tests {
16787    use super::mtp_chain_head_index;
16788
16789    #[test]
16790    fn embedded_step_heads_cycle_in_declared_order() {
16791        let actual: Vec<usize> = (0..8).map(|step| mtp_chain_head_index(step, 3)).collect();
16792        assert_eq!(actual, [0, 1, 2, 0, 1, 2, 0, 1]);
16793    }
16794
16795    #[test]
16796    fn standalone_draft_remains_single_head() {
16797        assert!((0..8).all(|step| mtp_chain_head_index(step, 1) == 0));
16798    }
16799}
16800
16801#[cfg(test)]
16802mod tp_verified_prefix_tests {
16803    use super::validate_tp_kv_snapshot_shape;
16804    use crate::tp::ResidentTpKvCache;
16805
16806    #[test]
16807    fn snapshot_shape_accepts_matching_tp_presence() {
16808        let layers = vec![
16809            Some(ResidentTpKvCache::new(Vec::new(), 1, 1, 1, 1, 8)),
16810            None,
16811        ];
16812        validate_tp_kv_snapshot_shape(&layers, &[Some(2), None]).unwrap();
16813    }
16814
16815    #[test]
16816    fn snapshot_shape_rejects_changed_tp_presence() {
16817        let layers = vec![Some(ResidentTpKvCache::new(Vec::new(), 1, 1, 1, 1, 8))];
16818        let error = validate_tp_kv_snapshot_shape(&layers, &[None])
16819            .unwrap_err()
16820            .to_string();
16821        assert!(error.contains("changed shape"), "unexpected error: {error}");
16822    }
16823
16824    #[test]
16825    fn step37_dcw_rebase_at_5151_boundary() {
16826        use crate::cache::KvRingAppend;
16827        let mut cache = ResidentTpKvCache::new_swa(Vec::new(), 128, 128, 136, 96, 262_144, 512);
16828        assert_eq!(cache.physical_capacity(), 5151);
16829        cache.publish_hydration(5150, 0).unwrap();
16830        assert_eq!(cache.ring_base(), Some(0));
16831
16832        // Before rebase, attempting to view rows past capacity fails with the exact bug error
16833        let err = cache.physical_range(5150, 5153).unwrap_err();
16834        assert_eq!(
16835            err,
16836            "SWA ring view [5150,5153) is outside resident [0,5151)"
16837        );
16838
16839        let (write_row, would_rebase) = cache.peek_append_ring(3).unwrap();
16840        assert!(would_rebase);
16841        assert_eq!(write_row, 542);
16842
16843        let tx = cache.begin_transaction().unwrap();
16844        let plan = cache.prepare_append(tx, 3).unwrap();
16845        assert_eq!(plan.target(), 5153);
16846        assert_eq!(plan.write_row(), 542);
16847        assert_eq!(
16848            plan.ring_append(),
16849            Some(KvRingAppend::Rebase {
16850                src_row: 4608,
16851                keep_rows: 542,
16852                new_base: 4608,
16853                write_row: 542,
16854            })
16855        );
16856        cache.publish_append_rebase(plan).unwrap();
16857        cache.publish_append_plan(plan).unwrap();
16858        assert_eq!(cache.ring_base(), Some(4608));
16859
16860        // After rebase, physical range is within bounds
16861        let range = cache.physical_range(5150, 5153).unwrap();
16862        assert_eq!(range, 542..545);
16863
16864        let target = cache.commit_target(tx, 3).unwrap();
16865        cache.publish_finalize(tx, target).unwrap();
16866        assert_eq!((cache.committed_len(), cache.staged_len()), (5153, 5153));
16867    }
16868
16869    #[test]
16870    fn step37_dcw_rebase_rollback_preserves_view_and_base() {
16871        let mut cache = ResidentTpKvCache::new_swa(Vec::new(), 128, 128, 136, 96, 262_144, 512);
16872        cache.publish_hydration(5150, 0).unwrap();
16873
16874        let tx = cache.begin_transaction().unwrap();
16875        let plan = cache.prepare_append(tx, 3).unwrap();
16876        cache.publish_append_rebase(plan).unwrap();
16877        cache.publish_append_plan(plan).unwrap();
16878        assert_eq!(cache.ring_base(), Some(4608));
16879
16880        // Rollback 0 rows accepted (pass declined)
16881        let rollback = cache.commit_target(tx, 0).unwrap();
16882        cache.publish_finalize(tx, rollback).unwrap();
16883        assert_eq!((cache.committed_len(), cache.staged_len()), (5150, 5150));
16884        assert_eq!(cache.ring_base(), Some(4608));
16885
16886        // View for base_len (5150) is still valid in resident [4608, 4608 + 5151)
16887        let range = cache.physical_range(4608, 5150).unwrap();
16888        assert_eq!(range, 0..542);
16889    }
16890}
16891
16892#[cfg(test)]
16893mod dspark_sparse_tests {
16894    use super::dspark_sparse_softmax_topk;
16895
16896    #[test]
16897    fn topk_keeps_full_softmax_mass_and_stable_ties() {
16898        let logits = [1.0f32, 3.0, 3.0, -2.0];
16899        let (ids, top_logits, probs, tail) = dspark_sparse_softmax_topk(&logits, 2, 1.0).unwrap();
16900        assert_eq!(ids, vec![1, 2]);
16901        assert_eq!(top_logits, vec![3.0, 3.0]);
16902        let denominator = logits.iter().map(|value| (value - 3.0).exp()).sum::<f32>();
16903        let expected = 1.0 / denominator;
16904        assert!((probs[0] - expected).abs() < 1.0e-6);
16905        assert!((probs[1] - expected).abs() < 1.0e-6);
16906        assert!((tail - (1.0 - 2.0 * expected)).abs() < 1.0e-6);
16907        assert!((probs.iter().sum::<f32>() + tail - 1.0).abs() < 1.0e-6);
16908    }
16909}
16910
16911#[cfg(test)]
16912mod spec_replay_env_tests {
16913    use super::spec_replay_env_on;
16914
16915    #[test]
16916    fn replay_requires_literal_one() {
16917        assert!(!spec_replay_env_on(None));
16918        assert!(!spec_replay_env_on(Some("")));
16919        assert!(!spec_replay_env_on(Some("0")));
16920        assert!(!spec_replay_env_on(Some("true")));
16921        assert!(!spec_replay_env_on(Some("2")));
16922        assert!(spec_replay_env_on(Some("1")));
16923    }
16924}
16925
16926#[cfg(test)]
16927mod telem_tests {
16928    use super::{SPEC_TELEM_POS, SpecTelemetry, SpecTelemetryCounters};
16929
16930    #[test]
16931    fn synthetic_accept_masks_produce_tau_and_position_histogram() {
16932        let counters = SpecTelemetryCounters::default();
16933        for mask in [
16934            [true, true, true],
16935            [true, true, false],
16936            [true, false, false],
16937            [false, false, false],
16938        ] {
16939            let accepted = mask.iter().take_while(|&&value| value).count();
16940            counters.record_round(mask.len(), accepted);
16941        }
16942
16943        let snapshot = counters.snapshot();
16944        assert_eq!(
16945            (snapshot.rounds, snapshot.drafted, snapshot.accepted),
16946            (4, 12, 6)
16947        );
16948        assert_eq!(&snapshot.pos_drafted[..3], &[4, 4, 4]);
16949        assert_eq!(&snapshot.pos_accepted[..3], &[3, 2, 1]);
16950        assert_eq!(snapshot.tau(), 1.5);
16951        assert_eq!(snapshot.pos_drafted[3..], [0; SPEC_TELEM_POS - 3]);
16952        assert_eq!(snapshot.pos_accepted[3..], [0; SPEC_TELEM_POS - 3]);
16953    }
16954
16955    /// The worker's per-burst pattern: stash, accumulate, diff — the delta must isolate
16956    /// exactly the burst's contribution (pool-resumed sessions carry prior requests' counts).
16957    #[test]
16958    fn delta_isolates_burst_contribution() {
16959        let mut t = SpecTelemetry::default();
16960        // "previous request": 2 rounds of k=3, accepts 3 then 1.
16961        for (kr, na) in [(3usize, 3usize), (3, 1)] {
16962            t.rounds += 1;
16963            t.drafted += kr as u64;
16964            t.accepted += na as u64;
16965            for j in 0..kr {
16966                t.pos_drafted[j] += 1;
16967            }
16968            for j in 0..na {
16969                t.pos_accepted[j] += 1;
16970            }
16971        }
16972        let before = t;
16973        // "this burst": 1 round k=3, accepts 2.
16974        t.rounds += 1;
16975        t.drafted += 3;
16976        t.accepted += 2;
16977        for j in 0..3 {
16978            t.pos_drafted[j] += 1;
16979        }
16980        for j in 0..2 {
16981            t.pos_accepted[j] += 1;
16982        }
16983        let d = t.delta_since(&before);
16984        assert_eq!((d.rounds, d.drafted, d.accepted), (1, 3, 2));
16985        assert_eq!(&d.pos_drafted[..3], &[1, 1, 1]);
16986        assert_eq!(&d.pos_accepted[..3], &[1, 1, 0]);
16987        assert_eq!(d.pos_drafted[3..], [0; SPEC_TELEM_POS - 3]);
16988    }
16989
16990    /// merge(delta) then merge(delta2) equals accumulating both — the per-model /metrics
16991    /// aggregation invariant.
16992    #[test]
16993    fn merge_accumulates_fieldwise() {
16994        let mut agg = SpecTelemetry::default();
16995        let mut d1 = SpecTelemetry {
16996            rounds: 2,
16997            drafted: 6,
16998            accepted: 4,
16999            ..Default::default()
17000        };
17001        d1.pos_drafted[0] = 2;
17002        d1.pos_accepted[0] = 2;
17003        let mut d2 = SpecTelemetry {
17004            rounds: 1,
17005            drafted: 3,
17006            accepted: 1,
17007            ..Default::default()
17008        };
17009        d2.pos_drafted[0] = 1;
17010        d2.pos_accepted[0] = 1;
17011        d2.pos_drafted[1] = 1;
17012        agg.merge(&d1);
17013        agg.merge(&d2);
17014        assert_eq!((agg.rounds, agg.drafted, agg.accepted), (3, 9, 5));
17015        assert_eq!(agg.pos_drafted[0], 3);
17016        assert_eq!(agg.pos_accepted[0], 3);
17017        assert_eq!(agg.pos_drafted[1], 1);
17018        assert_eq!(agg.pos_accepted[1], 0);
17019    }
17020
17021    /// Wrong-snapshot diff saturates to zero instead of wrapping — the counters feed a
17022    /// public metrics surface and must never publish a u64-wrapped garbage value.
17023    #[test]
17024    fn delta_saturates_never_wraps() {
17025        let small = SpecTelemetry {
17026            rounds: 1,
17027            drafted: 2,
17028            accepted: 1,
17029            ..Default::default()
17030        };
17031        let big = SpecTelemetry {
17032            rounds: 5,
17033            drafted: 15,
17034            accepted: 9,
17035            ..Default::default()
17036        };
17037        let d = small.delta_since(&big);
17038        assert_eq!((d.rounds, d.drafted, d.accepted), (0, 0, 0));
17039    }
17040}
17041
17042#[cfg(test)]
17043mod opti_fork_tests {
17044    use super::{
17045        OptiControllerPolicy, OptiForkAction, OptiForkGateMode, OptiForkGenerationTracker,
17046    };
17047
17048    #[test]
17049    fn controller_threshold_and_three_miss_breaker_are_exact() {
17050        let mut policy = OptiControllerPolicy {
17051            threshold: 0.7,
17052            consecutive_misses: 0,
17053            breaker_tripped: false,
17054        };
17055        assert!(!policy.admit(0.699_999));
17056        assert!(policy.admit(0.7));
17057        assert!(!policy.resolve(false));
17058        assert!(!policy.resolve(false));
17059        assert!(policy.resolve(false));
17060        assert!(policy.breaker_tripped);
17061        assert!(!policy.admit(1.0));
17062        assert!(
17063            !policy.resolve(true),
17064            "a resolved hit cannot re-arm a tripped request"
17065        );
17066        assert!(policy.breaker_tripped);
17067    }
17068
17069    #[test]
17070    fn zero_threshold_is_the_true_unconditional_measurement_arm() {
17071        let mut policy = OptiControllerPolicy {
17072            threshold: 0.0,
17073            consecutive_misses: 0,
17074            breaker_tripped: false,
17075        };
17076        for _ in 0..16 {
17077            assert!(policy.admit(0.0));
17078            assert!(!policy.resolve(false));
17079        }
17080        for invalid in [f32::NAN, f32::INFINITY, -0.01, 1.01] {
17081            assert!(
17082                !policy.admit(invalid),
17083                "invalid q proxy must fail closed: {invalid}"
17084            );
17085        }
17086        assert!(!policy.breaker_tripped);
17087        assert_eq!(policy.consecutive_misses, 0);
17088    }
17089
17090    #[test]
17091    fn alternating_mode_flips_by_generation_not_round_parity() {
17092        assert_eq!(OptiForkGateMode::Alternate.action(0), OptiForkAction::Hit);
17093        assert_eq!(OptiForkGateMode::Alternate.action(1), OptiForkAction::Miss);
17094        assert_eq!(OptiForkGateMode::Alternate.action(8), OptiForkAction::Hit);
17095        assert_eq!(OptiForkGateMode::Alternate.action(9), OptiForkAction::Miss);
17096    }
17097
17098    #[test]
17099    fn live_generation_cannot_be_overwritten() {
17100        let mut tracker = OptiForkGenerationTracker::default();
17101        let g0 = tracker.reserve().unwrap();
17102        let g1 = tracker.reserve().unwrap();
17103        let err = tracker.reserve().unwrap_err().to_string();
17104        assert!(
17105            err.contains("still owns generation 0"),
17106            "unexpected error: {err}"
17107        );
17108        tracker.retire(g0).unwrap();
17109        let g2 = tracker.reserve().unwrap();
17110        assert_eq!((g2.id, g2.slot), (2, 0));
17111        tracker.retire(g1).unwrap();
17112        tracker.retire(g2).unwrap();
17113    }
17114
17115    #[test]
17116    fn teardown_rejects_a_stale_generation_tag() {
17117        let mut tracker = OptiForkGenerationTracker::default();
17118        let g0 = tracker.reserve().unwrap();
17119        tracker.retire(g0).unwrap();
17120        let err = tracker.retire(g0).unwrap_err().to_string();
17121        assert!(err.contains("teardown mismatch"), "unexpected error: {err}");
17122    }
17123}
17124
17125#[cfg(test)]
17126mod draft_graph_fallback_tests {
17127    use super::DraftGraphFallback;
17128
17129    /// The Q2 contract, part (a): a fallback flip is LOUD — exactly once per flip.
17130    #[test]
17131    fn flip_is_loud_once_and_memoized_after() {
17132        let mut f = DraftGraphFallback::default();
17133        let line = f
17134            .mark_greedy("out of memory")
17135            .expect("first flip must return the warn line");
17136        assert!(
17137            line.contains("WARN"),
17138            "flip line must be warn-level: {line}"
17139        );
17140        assert!(
17141            line.contains("out of memory"),
17142            "flip line must carry the reason: {line}"
17143        );
17144        assert!(f.greedy_failed());
17145        // re-marking an already-failed graph is the memoization: quiet, still failed.
17146        assert!(f.mark_greedy("out of memory").is_none());
17147        assert!(f.greedy_failed());
17148        // the two graphs' flags are independent (greedy flip leaves sampled capturable).
17149        assert!(!f.sampled_failed());
17150        let line_s = f
17151            .mark_sampled("capture unsupported")
17152            .expect("sampled flip is its own flip");
17153        assert!(
17154            line_s.contains("sampled"),
17155            "sampled flip names itself: {line_s}"
17156        );
17157        assert!(f.mark_sampled("capture unsupported").is_none());
17158    }
17159
17160    /// The Q2 contract, part (b): resume-from-pool RESETS both flags (fresh capture chance),
17161    /// and says so exactly when there was something to reset.
17162    #[test]
17163    fn reset_on_resume_clears_flags_and_logs_once() {
17164        let mut f = DraftGraphFallback::default();
17165        // clean session: resume is silent, nothing to reset.
17166        assert!(f.reset_on_resume().is_none());
17167        f.mark_greedy("oom").unwrap();
17168        f.mark_sampled("oom").unwrap();
17169        let note = f
17170            .reset_on_resume()
17171            .expect("a set flag must produce the reset note");
17172        assert!(
17173            note.contains("greedy+sampled"),
17174            "note names what was reset: {note}"
17175        );
17176        assert!(
17177            !f.greedy_failed() && !f.sampled_failed(),
17178            "both flags cleared"
17179        );
17180        // and the NEXT failure after a reset is a fresh flip — loud again.
17181        assert!(f.mark_greedy("oom again").is_some());
17182        let note2 = f.reset_on_resume().expect("greedy-only reset");
17183        assert!(note2.contains("(greedy)"), "single-flag note: {note2}");
17184    }
17185
17186    /// Shape-change clears (dmask realloc / mask-shape mismatch / s_key change) stay silent —
17187    /// they precede a fresh capture attempt whose own failure re-flips loudly.
17188    #[test]
17189    fn shape_change_clears_are_silent() {
17190        let mut f = DraftGraphFallback::default();
17191        f.mark_greedy("oom").unwrap();
17192        f.clear_greedy();
17193        assert!(!f.greedy_failed());
17194        f.mark_sampled("oom").unwrap();
17195        f.clear_sampled();
17196        assert!(!f.sampled_failed());
17197        // after a silent clear there is nothing left for resume to report.
17198        assert!(f.reset_on_resume().is_none());
17199    }
17200}
17201
17202/// SAMPLED DRAFT-GRAPH KEY (lane/graph-s-key-exactness-20260819).
17203///
17204/// These are the CPU teeth for an exactness bug whose live reproduction needs a GPU, a trunk, a
17205/// drafter and a two-turn session: the key itself. Every test below fails against the pre-fix key
17206/// `(seed, temp.to_bits(), k)` — `legacy_key` restates it so the collision is explicit rather
17207/// than remembered.
17208#[cfg(test)]
17209mod sampled_graph_key_tests {
17210    use super::{SampledGraphKey, debug_t_pred0};
17211
17212    /// The pre-fix key, verbatim: `let s_key = (sp_seed, sp_temp.to_bits(), k);`
17213    fn legacy_key(k: &SampledGraphKey) -> (u64, u32, usize) {
17214        (k.seed, k.temp_bits, k.k)
17215    }
17216
17217    fn pure_temp_key() -> SampledGraphKey {
17218        // temperature 1.0, filters off — today's serve default, the shape that parks a graph.
17219        SampledGraphKey::new(12345, 1.0, 3, 0, 1.0, 0.0, false)
17220    }
17221
17222    /// THE COLLISION. Two requests that differ ONLY in the truncation filters shared one key, so
17223    /// a parked pure-temp graph survived into a filtered request and the launch site launched it.
17224    #[test]
17225    fn vendor_filters_change_the_key() {
17226        let parked = pure_temp_key();
17227        // qwen3.8 generation_config.json — what the vendor-default flip makes the default shape.
17228        let vendor = SampledGraphKey::new(12345, 1.0, 3, 20, 0.95, 0.0, false);
17229        assert_eq!(
17230            legacy_key(&parked),
17231            legacy_key(&vendor),
17232            "pre-fix key collided: this is the bug, and the reason a test asserts on it",
17233        );
17234        assert_ne!(parked, vendor, "post-fix key must separate the two regimes");
17235        assert!(parked.pure_temp());
17236        assert!(!vendor.pure_temp());
17237    }
17238
17239    /// Each distribution-shaping field alone is enough to drop the parked graph.
17240    #[test]
17241    fn every_filter_field_is_keyed() {
17242        let base = pure_temp_key();
17243        for (what, other) in [
17244            (
17245                "top_k",
17246                SampledGraphKey::new(12345, 1.0, 3, 20, 1.0, 0.0, false),
17247            ),
17248            (
17249                "top_p",
17250                SampledGraphKey::new(12345, 1.0, 3, 0, 0.95, 0.0, false),
17251            ),
17252            (
17253                "min_p",
17254                SampledGraphKey::new(12345, 1.0, 3, 0, 1.0, 0.05, false),
17255            ),
17256            (
17257                "penalties",
17258                SampledGraphKey::new(12345, 1.0, 3, 0, 1.0, 0.0, true),
17259            ),
17260        ] {
17261            assert_ne!(base, other, "{what} must be part of the key");
17262            assert!(!other.pure_temp(), "{what} leaves the pure-temp regime");
17263            assert_eq!(
17264                legacy_key(&base),
17265                legacy_key(&other),
17266                "{what} was invisible to the pre-fix key",
17267            );
17268        }
17269    }
17270
17271    /// The baked constants stay keyed (this half was always right — regression cover for it).
17272    #[test]
17273    fn baked_constants_stay_keyed() {
17274        let base = pure_temp_key();
17275        assert_ne!(
17276            base,
17277            SampledGraphKey::new(999, 1.0, 3, 0, 1.0, 0.0, false),
17278            "seed"
17279        );
17280        assert_ne!(
17281            base,
17282            SampledGraphKey::new(12345, 0.7, 3, 0, 1.0, 0.0, false),
17283            "temp"
17284        );
17285        assert_ne!(
17286            base,
17287            SampledGraphKey::new(12345, 1.0, 4, 0, 1.0, 0.0, false),
17288            "k"
17289        );
17290        // bitwise on temperature: 0.7f32 vs the same value re-derived must NOT differ.
17291        assert_eq!(
17292            SampledGraphKey::new(1, 0.7, 3, 0, 1.0, 0.0, false),
17293            SampledGraphKey::new(1, 7.0 / 10.0, 3, 0, 1.0, 0.0, false),
17294        );
17295    }
17296
17297    /// THE LOAD-BEARING HALF OF THE SEED DECISION (lane/session-resume-sampler-predicate-
17298    /// 20260820). The whole-session resume predicate deliberately does NOT compare `seed`: an
17299    /// omitted serve `seed` draws fresh per-request entropy, so comparing it would refuse every
17300    /// seed-omitting sampled conversation. That is only sound because the one piece of parked state
17301    /// that BAKES the seed — this graph — is re-keyed on it, so a seed change drops and recaptures.
17302    ///
17303    /// This test is the other end of that argument, asserted here rather than remembered in a
17304    /// comment: if a future change dropped `seed` from the key, the resume predicate's exclusion
17305    /// would silently become the unsound thing it is documented not to be.
17306    /// (Paired with `seed_alone_does_not_refuse` in `memra-sampling`.)
17307    #[test]
17308    fn seed_alone_still_rekeys_the_draft_graph() {
17309        let parked = pure_temp_key();
17310        let reseeded = SampledGraphKey::new(999, 1.0, 3, 0, 1.0, 0.0, false);
17311        assert_ne!(
17312            parked, reseeded,
17313            "a seed-only change MUST drop the parked sampled graph — the resume predicate's \
17314             decision not to compare seed rests on exactly this",
17315        );
17316        // Same regime on both sides: the drop is a recapture, not a fall to the eager chain
17317        // because of a filter difference.
17318        assert!(parked.pure_temp() && reseeded.pure_temp());
17319    }
17320
17321    /// `pure_temp()` is the capture guard's predicate, computed from the key so the two cannot
17322    /// drift. The equality below is the invariant the launch-site guard asserts: identical keys
17323    /// agree on the regime, so a graph that survives the drop is legal to launch.
17324    #[test]
17325    fn equal_keys_agree_on_the_regime() {
17326        let a = SampledGraphKey::new(7, 0.8, 3, 20, 0.95, 0.0, false);
17327        let b = SampledGraphKey::new(7, 0.8, 3, 20, 0.95, 0.0, false);
17328        assert_eq!(a, b);
17329        assert_eq!(a.pure_temp(), b.pure_temp());
17330        // top_p slightly above 1.0 (a client sending 1.0 exactly, or an operator default) is
17331        // still the unfiltered regime, matching the original `sp.top_p >= 1.0` test.
17332        assert!(SampledGraphKey::new(7, 0.8, 3, 0, 1.0, 0.0, false).pure_temp());
17333        assert!(SampledGraphKey::new(7, 0.8, 3, 0, 1.5, -1.0, false).pure_temp());
17334    }
17335
17336    /// The WIDENED capture regime (lane/step37-draft-graph-serving-20260830): truncation-
17337    /// filtered shapes are capturable — the filter runs IN-GRAPH (`filter_stats` +
17338    /// `gumbel_perturb_filtered_ctr`), so the draft draws from the same filtered
17339    /// distribution the accept test reconstructs. Penalties never are: the per-round
17340    /// history cannot be baked. The step37 vendor-default shape (temp 0.5 / top_p 0.9) is
17341    /// exactly the previously-excluded regime this lane exists to capture.
17342    #[test]
17343    fn filtered_regimes_are_capturable_penalties_never() {
17344        let vendor = SampledGraphKey::new(12345, 0.5, 3, 0, 0.9, 0.0, false);
17345        assert!(!vendor.pure_temp());
17346        assert!(vendor.filtered());
17347        assert!(
17348            vendor.graph_capturable(),
17349            "the vendor-default filtered shape must be capturable (default door state)",
17350        );
17351        assert!(pure_temp_key().graph_capturable());
17352        assert!(
17353            !pure_temp_key().filtered(),
17354            "pure-temp takes the legacy (filterless) capture body",
17355        );
17356        let pen = SampledGraphKey::new(12345, 0.5, 3, 0, 0.9, 0.0, true);
17357        assert!(
17358            !pen.graph_capturable(),
17359            "penalty history varies per round and can never be baked into a graph",
17360        );
17361    }
17362
17363    /// MEMRA_DEBUG_SPEC on a SAMPLED spec request past round 0: the print must render without
17364    /// indexing the empty greedy `preds` vector (it panicked the GPU worker before this lane).
17365    #[test]
17366    fn debug_print_survives_the_sampled_arm() {
17367        // round >= 1 with a pending bonus == base 1, sampled == `preds` empty.
17368        assert_eq!(debug_t_pred0(true, 1, 4242, &[]), "n/a");
17369        assert_eq!(debug_t_pred0(true, 2, 4242, &[]), "n/a");
17370        // round 0 without a pending bonus still reports last_pred, in both arms.
17371        assert_eq!(debug_t_pred0(true, 0, 4242, &[]), "4242");
17372        assert_eq!(debug_t_pred0(false, 0, 4242, &[7, 8]), "4242");
17373        // greedy keeps the real prediction it always printed.
17374        assert_eq!(debug_t_pred0(false, 1, 4242, &[7, 8]), "7");
17375        assert_eq!(debug_t_pred0(false, 2, 4242, &[7, 8]), "8");
17376    }
17377}