Skip to main content

memra_engine/
spec.rs

1//! Qwen3.5 MTP (NextN) greedy speculative decode (research/mtp/MTP-PLAN.md §A/§B/§C/§D).
2//!
3//! Greedy spec decode is MATHEMATICALLY EXACT: the accepted+bonus token stream is token-for-token
4//! identical to plain greedy `generate`. This module provides:
5//!   - `mtp_head_forward`  (§A, T=1): one NextN draft-token forward.
6//!   - `decode_step_t`     (§D.3, T=K+1): batched target verify forward, all-column logits.
7//!   - `generate_spec`     (§B): the draft/verify/accept/rollback orchestrator.
8//! Cache snapshot/rollback lives in cache.rs (§D.4). The MTP head uses its OWN scratch KV (§D.6),
9//! PERSISTENT over the committed sequence (see `MtpScratch`).
10
11use crate::Engine;
12use crate::cache::{Cache, KvLayer};
13use crate::forward::argmax;
14use crate::hybrid::{FullAttnLayer, HybridModel, LinearAttnLayer, Mixer, MtpHead};
15use cudarc::driver::CudaSlice;
16use std::sync::atomic::{AtomicU64, Ordering};
17
18/// Parse the documented `MEMRA_SPEC_REPLAY=1` rollback seam.
19///
20/// Keep this shared with serving admission so `=0` cannot select replay in one
21/// layer while another layer treats it as disabled.
22pub fn spec_replay_env_on(value: Option<&str>) -> bool {
23    value == Some("1")
24}
25
26pub fn spec_replay_env_enabled() -> bool {
27    let value = std::env::var("MEMRA_SPEC_REPLAY").ok();
28    spec_replay_env_on(value.as_deref())
29}
30
31/// step35 dcw draft-chain door (lane/step37-draft-graph-20260829). ON routes the step35 MTP
32/// block's draft attention through the WINDOWED device-counter family
33/// (`append_kv_quantized_dcw` + `fa_decode_dcw`, the step TP graph arc's kernels), which
34/// derives the SWA view entirely from device state (len_d, base_d, window): exactly the view
35/// offset the old capture refusal said `fa_decode_dc` could not express. BOTH draft modes
36/// switch together: eager and captured run the ONE launcher at the ONE bucket
37/// (min(cap, window)), so graph-vs-eager draft parity holds by construction (the
38/// `mtp_full_attn_dc` precedent).
39///
40/// DEFAULT ON since lane/step37-draft-graph-serving-20260830: the 20260829 lane shipped it
41/// OFF because it enabled nothing at the shipping head count (capture was structurally
42/// unreachable at heads=3); with the multi-head chain capture and the in-graph filtered
43/// sampler landed, this door is the kernel prerequisite for the captured chain on the
44/// QUALIFIED serving shape, and the exactness battery (greedy K=1..8 identity, per-K
45/// acceptance identity, seeded sampled twins) banks on the ON arm. Rollback seam:
46/// MEMRA_STEP35_DRAFT_DCW=0 restores the host-len eager arm (`mtp_step35_attn`) plus the
47/// named capture refusal, byte-for-byte the pre-lane serving; no state survives restart.
48fn step35_draft_dcw_on() -> bool {
49    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
50    *ON.get_or_init(|| std::env::var("MEMRA_STEP35_DRAFT_DCW").as_deref() != Ok("0"))
51}
52
53/// Multi-head MTP draft-chain capture door (lane/step37-draft-graph-serving-20260830,
54/// default ON — receipts in the lane RESULTS). ON lets the step-modulo prefix-replay chain
55/// (`mtp_extra` non-empty, the step37 3-head shipping shape) capture per-head single-row
56/// CUDA graphs and replay them in the exact eager launch order; the chain POLICY (head
57/// selection, prefix length, seed history) stays host-side, so graph-vs-eager drafts are
58/// bit-identical by construction. A failed capture degrades LOUDLY to the eager chain (the
59/// draft-graph WARN contract). OFF (=0) keeps the eager chain as the only multi-head path —
60/// the pre-lane serving byte-for-byte. Single-head capture is untouched by this door.
61fn mtp_chain_graph_on() -> bool {
62    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
63    *ON.get_or_init(|| std::env::var("MEMRA_MTP_CHAIN_GRAPH").as_deref() != Ok("0"))
64}
65
66/// In-graph FILTERED sampled draft door (lane/step37-draft-graph-serving-20260830, default
67/// ON — receipts in the lane RESULTS). ON widens the sampled draft-graph capture from the
68/// pure-temp regime to every truncation-filtered regime (top_k / top_p / min_p): the capture
69/// body runs `filter_stats` + `gumbel_perturb_filtered_ctr` IN-GRAPH, so the draft draws
70/// from the SAME filtered distribution the verify's accept test reconstructs (the
71/// graph-s-key exactness law, now satisfied inside the graph instead of by refusing it).
72/// Penalties stay eager either way (the history varies per round and cannot be baked).
73/// The pure-temp capture body is UNTOUCHED by this door (byte-identical to the pre-lane
74/// graph). OFF (=0) restores the pure-temp-only capture guard: filtered requests draft
75/// eager, byte-for-byte the pre-lane behavior.
76fn spec_graph_filtered_on() -> bool {
77    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
78    *ON.get_or_init(|| std::env::var("MEMRA_SPEC_GRAPH_FILTERED").as_deref() != Ok("0"))
79}
80
81fn parse_prime_trows_width(value: Option<&str>) -> Result<usize, String> {
82    let Some(raw) = value else {
83        return Ok(8);
84    };
85    let width = raw
86        .parse::<usize>()
87        .map_err(|_| format!("MEMRA_PRIME_TROWS_T must be an integer in 2..=8, got {raw:?}"))?;
88    if !(2..=8).contains(&width) {
89        return Err(format!("MEMRA_PRIME_TROWS_T must be in 2..=8, got {width}"));
90    }
91    Ok(width)
92}
93
94#[cfg(test)]
95mod prime_trows_width_tests {
96    #[test]
97    fn width_defaults_to_eight_and_refuses_invalid_operator_values() {
98        assert_eq!(super::parse_prime_trows_width(None), Ok(8));
99        assert_eq!(super::parse_prime_trows_width(Some("2")), Ok(2));
100        assert_eq!(super::parse_prime_trows_width(Some("8")), Ok(8));
101        for invalid in ["", "1", "9", "32", "wide"] {
102            let err = super::parse_prime_trows_width(Some(invalid)).unwrap_err();
103            assert!(err.contains("MEMRA_PRIME_TROWS_T"), "{err}");
104            assert!(err.contains("2..=8"), "{err}");
105        }
106    }
107}
108
109/// One compact, anchor-bounded DSpark supervision record. `tokens[0]` is the anchor at p and
110/// `hidden` is its predecessor carrier h[p-1], matching the live NextN/DSpark pairing. Target
111/// rows p..p+gamma-1 score tokens p+1..p+gamma. They are the full-target softmax's top-k
112/// entries; `target_tail_probs[j]` is the probability mass outside those rows. All flattened
113/// target arrays are `[gamma, top_k]` in row-major order.
114pub struct DsparkAnchorRecord {
115    pub position: usize,
116    pub hidden: Vec<f32>,
117    pub tokens: Vec<u32>,
118    pub target_top_ids: Vec<u32>,
119    pub target_top_logits: Vec<f32>,
120    pub target_top_probs: Vec<f32>,
121    pub target_tail_probs: Vec<f32>,
122}
123
124fn dspark_sparse_softmax_topk(
125    logits: &[f32],
126    top_k: usize,
127    temperature: f32,
128) -> Result<(Vec<u32>, Vec<f32>, Vec<f32>, f32), Box<dyn std::error::Error>> {
129    if logits.is_empty() || top_k == 0 || top_k > logits.len() || temperature <= 0.0 {
130        return Err("invalid DSpark sparse-softmax shape or temperature".into());
131    }
132    if logits.iter().any(|value| !value.is_finite()) {
133        return Err("DSpark target logits contain a non-finite value".into());
134    }
135    let mut ranked: Vec<(u32, f32)> = logits
136        .iter()
137        .copied()
138        .enumerate()
139        .map(|(index, value)| (index as u32, value))
140        .collect();
141    let compare = |left: &(u32, f32), right: &(u32, f32)| {
142        right.1.total_cmp(&left.1).then(left.0.cmp(&right.0))
143    };
144    ranked.select_nth_unstable_by(top_k - 1, compare);
145    ranked[..top_k].sort_unstable_by(compare);
146
147    let max_logit = logits.iter().copied().fold(f32::NEG_INFINITY, f32::max);
148    let inv_temperature = 1.0f64 / temperature as f64;
149    let denominator: f64 = logits
150        .iter()
151        .map(|value| (((*value - max_logit) as f64) * inv_temperature).exp())
152        .sum();
153    let ids: Vec<u32> = ranked[..top_k].iter().map(|(index, _)| *index).collect();
154    let top_logits: Vec<f32> = ranked[..top_k].iter().map(|(_, value)| *value).collect();
155    let top_probs: Vec<f32> = top_logits
156        .iter()
157        .map(|value| ((((value - max_logit) as f64) * inv_temperature).exp() / denominator) as f32)
158        .collect();
159    let top_mass: f64 = top_probs.iter().map(|value| *value as f64).sum();
160    let tail = (1.0f64 - top_mass).clamp(0.0, 1.0) as f32;
161    Ok((ids, top_logits, top_probs, tail))
162}
163
164fn flatten_dspark_rows<T>(
165    rows: Vec<Option<Vec<T>>>,
166    position: usize,
167    label: &str,
168) -> Result<Vec<T>, Box<dyn std::error::Error>> {
169    let mut flattened = Vec::new();
170    for (slot, row) in rows.into_iter().enumerate() {
171        flattened.extend(
172            row.ok_or_else(|| format!("missing DSpark {label} at {position} slot {slot}"))?,
173        );
174    }
175    Ok(flattened)
176}
177
178/// H-SEED CONVENTION (MEMRA_SPEC_HPOST=1): feed the MTP head the POST-norm hidden — trunk rows
179/// hand over `output_norm(x)` and the draft chain recurrence hands over `shared_head_norm(h_nextn)`
180/// (= final_h) — matching the reference engines: llama.cpp #24025 ("qwen35: use post-norm hidden
181/// state for MTP", t_h_nextn is taken AFTER the final norm in both trunk and MTP graphs) and
182/// SGLang's qwen3_5_mtp (spec_info.hidden_states = the target model's post-norm output). memra's
183/// historical convention (default, MTP-PLAN §A) is PRE-norm x. Draft-quality-only: exactness is
184/// the verify's job either way; acceptance arbitrates. OnceLock: read once, hot-loop safe.
185/// `MEMRA_SPEC_HEAD_ROWS=1` — batch the verify tail's LM head over its t columns instead of running
186/// it at m=1 once per column. See the call site in `decode_step_t_core_stream` for why the batched
187/// form is the same per-row arithmetic (the bf16/q8 rows twins, not cuBLASLt) and what it costs
188/// today: the head is re-streamed t times per verify pass. Default off until the byte tape says so.
189pub(crate) fn head_rows_on() -> bool {
190    static ENV: std::sync::OnceLock<Option<bool>> = std::sync::OnceLock::new();
191    crate::step37_door(&ENV, "MEMRA_SPEC_HEAD_ROWS")
192}
193
194/// The serving walk's own doors, tri-stated the same way (owner flip 2026-08-27): env forces,
195/// unset takes the step37 family default. Call sites are the t-row verify walk itself.
196pub(crate) fn spec_verify_eager_on() -> bool {
197    static ENV: std::sync::OnceLock<Option<bool>> = std::sync::OnceLock::new();
198    crate::step37_door(&ENV, "MEMRA_SPEC_VERIFY_EAGER")
199}
200
201pub(crate) fn spec_verify_tcol_on() -> bool {
202    static ENV: std::sync::OnceLock<Option<bool>> = std::sync::OnceLock::new();
203    crate::step37_door(&ENV, "MEMRA_SPEC_VERIFY_TCOL")
204}
205
206/// NOT family-armed (2026-08-27): the walk's prime leaves its sub-32 TAIL chunk out of the
207/// DISTRIBUTED kv, so the server refuses before decode with "cache lengths diverged
208/// local=N distributed=floor(N/32)*32" for every prompt whose token count is not a multiple of
209/// 32 — i.e. nearly all real traffic. Isolated on the server route: defaults ERR (local=445
210/// distributed=416), MEMRA_PRIME_TROWS=0 OK. It was default-OFF before the 2026-08-27 flip and
211/// goes back to opt-in until the tail append is fixed and gated ON THE SERVER ROUTE, not just
212/// run-gen (run-gen calls decode_step_t on the whole prompt and never exercises this path — the
213/// reason a run-gen-only receipt could not see it). The GEMM prime supersedes it on this route.
214pub(crate) fn prime_trows_on() -> bool {
215    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
216    *ON.get_or_init(|| std::env::var("MEMRA_PRIME_TROWS").as_deref() == Ok("1"))
217}
218
219pub(crate) fn tcol_ffn_on() -> bool {
220    static ENV: std::sync::OnceLock<Option<bool>> = std::sync::OnceLock::new();
221    crate::step37_door(&ENV, "MEMRA_TCOL_FFN")
222}
223
224pub(crate) fn spec_hpost() -> bool {
225    static H: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
226    *H.get_or_init(|| {
227        std::env::var("MEMRA_SPEC_HPOST")
228            .map(|v| v != "0")
229            .unwrap_or(false)
230    })
231}
232
233/// LEAN VERIFY (default ON since 2026-07-08; MEMRA_SPEC_LEAN=0 reverts — close35 lane): the verify m-scaling
234/// probe + nsys diff showed the verify t-path pays ~1.0ms/call at m=1 over eager decode on the
235/// 35B, and the kernels are NOT the cause (dev-MoE identical, kernel-time delta only +179us).
236/// The overhead is (a) ~250 extra cuMemsetD8Async/call from `e.zeros()` on buffers every kernel
237/// fully overwrites (~0.9ms host issue + ~0.35ms GPU) and (b) the t=1 FA rows dispatch (rows_v2 +
238/// combine_rows, +50us vs the eager fa_decode pair). This flag switches (a) fully-overwritten
239/// verify buffers to `e.uninit` (identical bytes: every element is written before read) and
240/// (b) t==1 verify FA to the eager `fa_decode` entry (byte-identical: kernel-check pins the
241/// rows-vs-loop identity and the per-row loop at t=1 IS fa_decode on the same q). Gates arbitrate.
242pub(crate) fn spec_lean() -> bool {
243    static L: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
244    // DEFAULT ON since 2026-07-08 (MEMRA_SPEC_LEAN=0 reverts): bit-identical (buffers fully
245    // overwritten; gates green incl maxdiff-identical run-gen) and measured +2.4% e2e p3 /
246    // +1.5% p2 at the daily 35B config. m=1 verify now costs eager-decode parity.
247    *L.get_or_init(|| {
248        std::env::var("MEMRA_SPEC_LEAN")
249            .map(|v| v != "0")
250            .unwrap_or(true)
251    })
252}
253
254/// SMALL-M BATCHED VERIFY (default ON since 2026-07-09; MEMRA_SPEC_M2=0 reverts — lane/spec-m2): extend the
255/// batched linear-attn verify arm down to t=2 and batch the MoE dev token loop over a
256/// grid.z=token axis at every verify t. The close35 m-scaling probe put the m=2 verify tier at
257/// x1.54 of m=1 (llama x1.14); the per-column linear chain (t<3) and the serial MoE dev token
258/// loop are the two launch-structure causes. Both changes are LAUNCH-STRUCTURE ONLY:
259/// (a) the batched conv's t<pad ring update is pure copies (ssm_conv_ring_rebuild from a cloned
260///     ring — the ring stores raw input columns); every arithmetic kernel is the same one the
261///     t>=3 arm already runs (matmul_decode_exact bit-identical at m=2-4, gdn_scan's internal
262///     t-loop == chained T=1 steps);
263/// (b) the MoE dev-rows twins run the serial loop's per-token warp program with tok-offset
264///     pointers (same sel/w/aq/ad bytes, same dot order, same slot-ordered FMA chain).
265/// Gates arbitrate: run-spec K=1..8 self-consistency (35B+9B), kernel-check, run-gen argmax.
266pub(crate) fn spec_m2() -> bool {
267    static M: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
268    // DEFAULT ON since 2026-07-09 (MEMRA_SPEC_M2=0 reverts): launch-structure only — t=2
269    // batched linear arm (ring-roll copies, zero new FP order) + MoE dev-rows kernels
270    // (grid.z=token, 4 launches/layer at any verify t). Acceptance bit-identical at every K;
271    // 35B p2 +3.4% / p3 +3.6%; the profitable-K plateau widens (new optimum K=3 at 223).
272    *M.get_or_init(|| {
273        std::env::var("MEMRA_SPEC_M2")
274            .map(|v| v != "0")
275            .unwrap_or(true)
276    })
277}
278pub(crate) fn spec_stream() -> bool {
279    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
280    *ON.get_or_init(|| std::env::var("MEMRA_SPEC_STREAM").as_deref() == Ok("1"))
281}
282pub(crate) fn spec_stream_m() -> usize {
283    static M: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
284    *M.get_or_init(|| {
285        std::env::var("MEMRA_SPEC_STREAM_M")
286            .ok()
287            .and_then(|v| v.parse().ok())
288            .unwrap_or(4)
289    })
290}
291pub(crate) fn spec_devacc() -> bool {
292    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
293    *ON.get_or_init(|| std::env::var("MEMRA_SPEC_DEVACC").as_deref() == Ok("1"))
294}
295/// Engine-bundle slice 2 (DSF-ROUNDCOST-20260820 §1.1 host/device round trips + §2 rows 2-3),
296/// DEFAULT ON (`MEMRA_DSPARK_DEFER_READBACK=0` reverts): the dspark round's draft-chain DtoH
297/// is DEFERRED past verify dispatch and merged with the verify-argmax readback into ONE host
298/// sync (2 blocking DtoH/round -> 1). Verify embeds DEVICE tokens (`chain_d`) through the
299/// resident embed table — `embed_gather_u32_t`, bit-identical rows to the host gather by its
300/// own pinned contract. The host therefore dispatches snap + the whole verify while the DRAFT
301/// is still executing, instead of blocking ~1.7 ms on the chain and letting the device drain.
302/// Ladder arm only: the confidence policies size vt from a pre-verify head readback (their
303/// chain readback merges into that same sync instead). Exactness unchanged BY CONSTRUCTION —
304/// same tokens, same kernels, same order; E2E + accept-bank gates arbitrate.
305pub(crate) fn dspark_defer_readback_on() -> bool {
306    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
307    *ON.get_or_init(|| {
308        std::env::var("MEMRA_DSPARK_DEFER_READBACK")
309            .map(|v| v != "0")
310            .unwrap_or(true)
311    })
312}
313/// Engine-bundle slice 1 (DSF-ROUNDCOST-20260820 §1.1, lane/dspark-engine-bundle-20260820),
314/// DEFAULT ON (`MEMRA_STATE_COPY_BATCH=0` reverts): batch the dspark round's GDN state
315/// snapshot and partial-accept restore into single `copy_batch_uniform_f32` launches
316/// instead of ~2 memcpy dispatches (+2 alloc_zeros on the snap side) per linear layer per
317/// round — measured 0.67 ms/round snap + 0.25 ms/round commit of pure dispatch on the q38
318/// route. Launch-structure only: bytes, buffers and stream order are unchanged, so
319/// acceptance and streams stay bit-identical (E2E-gated on the B1 packs).
320pub(crate) fn state_copy_batch_on() -> bool {
321    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
322    *ON.get_or_init(|| {
323        std::env::var("MEMRA_STATE_COPY_BATCH")
324            .map(|v| v != "0")
325            .unwrap_or(true)
326    })
327}
328/// Engine-bundle slice 3 + fa-execupdate slice 4c (DSF-ROUNDCOST-20260820 §5 rank 1),
329/// DEFAULT OFF — `MEMRA_DSPARK_VERIFY_GRAPH=1` opts in: per-(segment, vt) CUDA graphs
330/// for the LINEAR-layer runs, plus the full-verify single graph per (vt, rung) when a
331/// round's rows all ride one seqs rung — see [`DsparkVerifyGraphs`]. Requires the
332/// slice-2 deferred path (device tokens); the eager walk is the byte-identical fallback.
333///
334/// MEASURED disposition (box6 card0, agentic pack, 2026-08-20, both slices): exactness
335/// holds everywhere (ALL EXACT, accept lines byte-match the banks, ckpt-gate oracle
336/// green over the graph + slab-commit paths). Slice-3's AUTO_FREE launch-scan limiter
337/// (25.6 us x 16 launches ≈ 0.41 ms/round) is FIXED — the captured bodies' alloc nodes
338/// are balanced by in-graph frees (census 84/84 per segment, 1776/1776 full) so graphs
339/// instantiate USE_NODE_PRIORITY and the scan is gone. What remains at gate scale:
340/// segment graphs +0.1 tok/s over the batched-rows default (114.4 vs 114.3 x5
341/// interleaved — the linear launch overhead was only ~0.1 ms); the FULL-verify graph is
342/// NET NEGATIVE at gate scale (110.6 vs 114.2: ~14-21 (vt, rung) captures/process at
343/// 2 full-walk executions + ~2.9k-node instantiate each eat far more than the ~0.2-0.3
344/// ms/round of remaining launch overhead). The orchestration ceiling of §1.3 is spent —
345/// the fa/append recovery landed DEFAULT-ON as the batched rows arm
346/// (`dspark_fa_rows_on`), not as a graph. The serve-lifetime cell (DSF-ROUNDCOST §9,
347/// nj-ws-solo) measured the amortization: crossover K≈33 requests, steady −0.246
348/// ms/round, −1.25% session wall over 240 requests — and the graphs-serve lane wired
349/// the door into the session arm (`dspark_spec_session_burst`) as a model-owned
350/// capture pool shared across sessions. Stays opt-in pending the owner's default-ON
351/// ratification on the serve-surface battery.
352pub(crate) fn dspark_verify_graph_on() -> bool {
353    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
354    *ON.get_or_init(|| std::env::var("MEMRA_DSPARK_VERIFY_GRAPH").as_deref() == Ok("1"))
355}
356/// MTP-ROUTE verify graphs, DEFAULT ON for the GDN+MoE family since 2026-08-23
357/// (`MEMRA_SPEC_VERIFY_GRAPH=0` is the kill switch, `=1` opts other families in).
358///
359/// The slice-4c capture already lived inside `qwen35_verify_tparallel` and said so in its own
360/// comment — "stream rides the qwen35moe burst, graphs ride the dspark route" — with no caller
361/// on this route. The MTP spec round is that caller.
362///
363/// WHY it is worth a default (receipts: `research/orndecode-20260822/VGRAPH.md`). With
364/// `MEMRA_SPEC_PHASE=1` this route's round reads verify-ISSUE 44-58% and verify-WAIT **0.0%**:
365/// the host is never waiting for the device, it is spending its own time launching the trunk.
366/// Replay collapses that into one graph launch and the phase all but disappears (55-62 ms ->
367/// 8-10 ms per burst).
368///
369/// MEASURED, two host generations, forced ON/OFF, balanced 4+4 boots in both orders:
370///   * current-generation host (9950X, the serving class): OFF 266.0-266.5, ON 318.8-319.5
371///     tok/s — **+19.7%**, no overlap, sub-1% spread per arm; per-round 6.9 -> 5.7 ms.
372///   * Zen 3 host: +3-9% (that rig's own clock drift is wider than the effect, so the ratio
373///     comes from per-round phase totals, which are internal to each boot).
374/// The ON arm lands at ~320 tok/s on BOTH hosts while OFF tracks host speed — the arm moves
375/// the round off the host and onto the device, which is the whole point.
376///
377/// EXACTNESS is structural (same kernels, same order) and gated anyway: a fixed-seed SAMPLED
378/// completion hashes identically ON vs OFF **and across both hosts** (`08941d5bb9762b21`),
379/// greedy seed-pinned likewise, `run-spec` K=1..8 PASS on both arms with identical acceptance
380/// at every K, kernel-check ALL GREEN.
381///
382/// SCOPE, deliberately narrow: default ON only where it was measured — the GatedDeltaNet +
383/// MoE family (`vgraph_family_default`). Qwen3.8-27B is GDN + DENSE mlp and would otherwise
384/// inherit this default unmeasured, which is the family-by-family law this repo keeps; it can
385/// opt in with `=1` once it has its own interleave. Also never armed together with
386/// ROUND-STREAM, and a round wider than the pool declines it for the eager walk.
387pub(crate) fn spec_verify_graph_env() -> Option<bool> {
388    static ON: std::sync::OnceLock<Option<bool>> = std::sync::OnceLock::new();
389    *ON.get_or_init(
390        || match std::env::var("MEMRA_SPEC_VERIFY_GRAPH").as_deref() {
391            Ok("1") => Some(true),
392            Ok("0") => Some(false),
393            _ => None,
394        },
395    )
396}
397/// SERVE-ROUTE twin of [`dspark_verify_graph_on`], DEFAULT ON — owner-ratified
398/// 2026-08-22 on the §10 serve-lifetime battery (DSF-ROUNDCOST-20260820 §10.3:
399/// crossover K=36–43, steady −0.357 ms/round, session wall −1.55..−1.65%, byte-exact
400/// 240/240 ×3 pairs, pool bounded at 8,852 MiB under `MEMRA_DSPARK_VG_MAX`). The env
401/// stays as the kill-switch: `MEMRA_DSPARK_VERIFY_GRAPH=0` restores the eager walk
402/// (byte-identical body); `MEMRA_DSPARK_VG_MAX=0` is the finer freeze valve. The BIN
403/// arm keeps its own opt-in default (`dspark_verify_graph_on`): at gate scale the
404/// capture toll is never repaid (§8 measured disposition — 14–21 captures over a
405/// 256-token run vs the serve session's thousands of rounds), and the two
406/// instruments must keep their own measured dispositions rather than share one flag.
407pub(crate) fn dspark_verify_graph_serve_on() -> bool {
408    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
409    *ON.get_or_init(|| std::env::var("MEMRA_DSPARK_VERIFY_GRAPH").as_deref() != Ok("0"))
410}
411/// Capture-count ceiling for the dspark verify-graph pool (graphs-serve lane) — the
412/// pool's memory policy STATED instead of silently unbounded. The keyspace is
413/// intrinsically finite — segment keys (run_start, vt) ≤ 16 runs x 7 windows, full
414/// keys (vt, rung, hi) ≤ 7 windows x the split-rung ladder (8 rungs at 32k ctx), ~168
415/// on the q38 export — so the default (256) never engages there; the knob is the
416/// safety valve for a future export with a wider ladder. At the ceiling the pool
417/// FREEZES: existing keys keep replaying, rounds needing a new capture run the eager
418/// walk byte-identically (round-atomic — a partial refusal would mix slab- and
419/// cols-stashed layers inside one commit). No eviction by design: destroying a live
420/// exec graph re-opens the stale-address class the indirect tables exist to close,
421/// and the bounded keyspace makes reclaim worthless.
422pub(crate) fn dspark_vg_cap() -> usize {
423    static CAP: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
424    *CAP.get_or_init(|| {
425        std::env::var("MEMRA_DSPARK_VG_MAX")
426            .ok()
427            .and_then(|v| v.parse().ok())
428            .unwrap_or(256)
429    })
430}
431
432/// PROJECTED REMAINING GROWTH of the verify-graph pool, in bytes (lane/hermes-perf-fixes,
433/// 2026-08-23 — the admission accounting the "pool dwarfs spec admission reserve" finding
434/// asks for). The pool was measured at 8,852 MiB at storm-complete on the q38 export while
435/// admission's transient floor (`SPEC_SHRINK_RESERVE`) is 1.5 GiB and never charged for it:
436/// sessions admitted while the pool is cold overcommit VRAM the pool WILL hold, because the
437/// pool grows monotonically (no eviction by design) and is model-owned across sessions.
438///
439/// SELF-MEASURING, no per-model constant (generic-model law — the 8,852 MiB is a q38 number
440/// and proves nothing about another export): the debt is remaining capture slots x the
441/// MARGINAL bytes a capture adds to this device's graph mem pool.
442///
443/// MARGINAL, NOT MEAN — measured correction (box9 on-box receipt, 2026-08-23). The first
444/// version of this used the mean (`reserved / captures`) and the live serve log showed why
445/// that is wrong: with the pool's reservation flat at ~33.6 MiB across captures 1..3, the
446/// mean-based debt printed **8,556 MB, then 4,261, then 2,830** — it extrapolated capture
447/// #1's ONE-TIME shared allocation (staging buffers, stash slabs, pointer tables: sized
448/// once per pool, shared by every key) across all 256 slots. An 8.5 GB phantom reserve at
449/// boot can refuse admissions that would have fit, which is a worse defect than the
450/// under-charge this accounting exists to remove. The marginal reading prices what an
451/// ADDITIONAL key actually costs: two observations `(captures, reserved)` give
452/// `(r1 - r0) / (c1 - c0)`, which is ~0 on an export whose pool does not grow per key and
453/// tracks real growth on one that does.
454///
455/// BOOTSTRAP (only one observation so far, so growth is unmeasurable): reserve one more
456/// pool's worth — `min(remaining x mean, reserved)`. "We have measured `reserved` bytes for
457/// `captures` keys; until growth is measurable, assume at most a doubling" is fail-safe in
458/// the same direction as the old rule without the 255x extrapolation.
459///
460/// Before the FIRST capture the debt is 0 (a single capture lands well inside the existing
461/// 1.5 GiB floor). `cap` is the intrinsic freeze ceiling (`MEMRA_DSPARK_VG_MAX`; =0 freeze
462/// valve => the pool cannot grow => debt 0); at or past the cap the pool FREEZES, so the
463/// debt is 0 there too.
464pub fn dspark_vg_debt_projection(
465    captures: usize,
466    cap: usize,
467    reserved_bytes: usize,
468    prev: Option<(usize, usize)>,
469) -> usize {
470    if captures == 0 || cap == 0 {
471        return 0;
472    }
473    let remaining = cap.saturating_sub(captures);
474    if remaining == 0 {
475        return 0;
476    }
477    match prev {
478        // marginal growth between two observations of the same pool
479        Some((c0, r0)) if captures > c0 => {
480            let marginal = reserved_bytes.saturating_sub(r0) / (captures - c0);
481            remaining.saturating_mul(marginal)
482        }
483        // bootstrap: at most one more pool's worth
484        _ => remaining
485            .saturating_mul(reserved_bytes / captures)
486            .min(reserved_bytes),
487    }
488}
489/// PRE-CAPTURE VRAM RESERVE CHECK door (lane/step37-vram-admission-20260830), DEFAULT ON.
490/// A draft-graph capture attempt on a tight card used to be try-and-fail: the 2 warmup
491/// forwards + instantiate grew the pool to the edge BEFORE the OOM surfaced, and the
492/// "eager fallback" then ran on a card the failed attempt had just exhausted (the owner's
493/// single-session second-prompt OOM: capture WARN followed by 28 step-OOM engine errors,
494/// device at 5 MiB free). With the gate ON, a capture is attempted only when the device's
495/// effective free (driver free + async-pool cached) covers the capture's expected appetite
496/// PLUS a post-capture safety floor — otherwise the session falls back to eager EARLY,
497/// with headroom intact, through the same LOUD once-per-flip WARN. `=0` restores
498/// try-and-fail (diagnostics door; the trim-on-OOM recovery below stays active either way).
499pub fn spec_capture_gate_on() -> bool {
500    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
501    *ON.get_or_init(|| std::env::var("MEMRA_SPEC_CAPTURE_GATE").as_deref() != Ok("0"))
502}
503
504/// Post-capture safety floor the reserve check keeps free ON TOP of the capture's own
505/// appetite: the same measured constant class as the admission transient floor
506/// (capture arenas + verify activations — the admit-oom control fit). A capture that
507/// would leave less than this behind is not worth its eager-coverage risk.
508pub(crate) const CAPTURE_HEADROOM_FLOOR: usize = 1536 << 20;
509
510/// Pure verdict half of the pre-capture reserve check (unit-testable): given the device's
511/// driver-free and pool-cached bytes and the capture's expected `need`, returns
512/// `Some((required, effective))` when the capture must be REFUSED, `None` when it fits.
513pub(crate) fn capture_headroom_verdict(
514    driver_free: usize,
515    pool_cached: usize,
516    need: usize,
517    floor: usize,
518) -> Option<(usize, usize)> {
519    let effective = driver_free.saturating_add(pool_cached);
520    let required = need.saturating_add(floor);
521    (effective < required).then_some((required, effective))
522}
523
524/// Expected device appetite of a draft-graph capture attempt when no measurement exists
525/// yet (bootstrap only — the model-owned high-water gauge takes over after the first
526/// observed capture). Deliberately conservative and shape-derived, never a per-family
527/// constant: per (head, mode) capture the two warmups + capture each walk one head
528/// forward whose dominant transients are a handful of `n_embd` rows and one `d_vocab`
529/// logits row, retained by the keeper; the sampled tail additionally parks
530/// `k` q-slots + perturb/q buffers of `d_vocab` each.
531pub(crate) fn draft_capture_bootstrap_estimate(
532    heads: usize,
533    k: usize,
534    d_vocab: usize,
535    n_embd: usize,
536) -> usize {
537    let per_capture = 3usize // 2 warmups + capture body, each retaining its transients
538        .saturating_mul(d_vocab.saturating_add(8 * n_embd))
539        .saturating_mul(4)
540        .max(32 << 20); // instantiate + driver-side graph backing per capture, floor
541    let captures = heads.max(1).saturating_mul(2); // interior + last per head
542    let sampled_slots = (k.saturating_add(2))
543        .saturating_mul(d_vocab)
544        .saturating_mul(4);
545    captures
546        .saturating_mul(per_capture)
547        .saturating_add(sampled_slots)
548        .max(64 << 20)
549}
550
551/// OOM predicate for capture-failure recovery (engine-side twin of the worker's
552/// `is_cuda_oom` — the same quoted-text contract).
553pub(crate) fn capture_err_is_oom(reason: &str) -> bool {
554    reason.contains("CUDA_ERROR_OUT_OF_MEMORY") || reason.contains("out of memory")
555}
556
557/// Impure half of the pre-capture reserve check: reads the device, trims the async pool
558/// when the driver alone is short but cached blocks would cover it (graph instantiate and
559/// cuBLAS workspaces allocate from the DRIVER, not from our pool — a pool sitting on freed
560/// blocks starves them), and returns the refusal reason line when the capture must not be
561/// attempted. `None` = go ahead.
562pub(crate) fn capture_headroom_refusal(e: &Engine, need: usize) -> Option<String> {
563    let Ok((driver_free, _total)) = e.ctx().mem_get_info() else {
564        return None; // unreadable device: keep the historical try-and-fail behavior
565    };
566    let pool_cached = e.pool_cached_bytes();
567    // A capture may take AT MOST HALF the discretionary headroom: required =
568    // 2x appetite + two floors (owner's contract: "fall back to eager EARLY with headroom
569    // intact"). Measured escalation on the owner-shape cells: one floor of slack let the
570    // capture walk the card to the edge and the burst step-OOM'd immediately; two floors
571    // still allowed a capture whose session then OOM'd on its own admission-charged work,
572    // because the capture had consumed the memory the charge was counting on. Requiring
573    // the appetite TWICE means the card retains a whole capture's worth of room after the
574    // capture lands - enough for the session's charged classes and its peers' bursts. The
575    // capture is an optimization worth ~2-3 ms of TTFT (draft-graph lane receipts); at the
576    // margin it is never worth an OOM incident.
577    let floor = CAPTURE_HEADROOM_FLOOR.saturating_mul(2);
578    let required_need = need.saturating_mul(2);
579    let required = required_need.saturating_add(floor);
580    match capture_headroom_verdict(driver_free, pool_cached, required_need, floor) {
581        Some((required, effective)) => Some(format!(
582            "insufficient VRAM headroom for capture: effective free {}MB (driver {}MB + pool-cached \
583             {}MB) < required {}MB (2x appetite {}MB + floor {}MB); capture skipped pre-attempt",
584            effective / (1 << 20),
585            driver_free / (1 << 20),
586            pool_cached / (1 << 20),
587            required / (1 << 20),
588            need / (1 << 20),
589            floor / (1 << 20),
590        )),
591        None => {
592            if driver_free < required && pool_cached > 0 {
593                let trimmed = e.pool_trim_to_zero();
594                if trimmed > 0 {
595                    eprintln!(
596                        "[spec] pre-capture pool trim: released {}MB cached back to the driver \
597                         (driver free {}MB < required {}MB; instantiate allocates from the driver)",
598                        trimmed / (1 << 20),
599                        driver_free / (1 << 20),
600                        required / (1 << 20),
601                    );
602                }
603            }
604            None
605        }
606    }
607}
608
609/// GRAPH-LAUNCH HEADROOM FLOOR (lane/step37-vram-admission-20260830, defect 3 root
610/// cause): `cuGraphLaunch` SEGFAULTS inside libcuda (offset +0x27c87f, a null internal
611/// dereference at address 0x60) when a captured graph is dispatched into a
612/// driver-exhausted card — reproduced on this lane's box with core dumps on BOTH the
613/// pre-lane and lane binaries (multi-active step-OOM squeeze; the crashing thread sits in
614/// `CudaGraph::launch` inside `generate_spec_inner2`). The eager arms fail RECOVERABLY on
615/// the same card (a quoted CUDA OOM the park path handles), so below this driver-free
616/// floor every graph arm yields to eager for the round. A named constant, not a knob: the
617/// winning value is the default and the guard exists to make a driver segfault
618/// unreachable, not to tune anything.
619pub(crate) const GRAPH_LAUNCH_MIN_FREE: usize = 256 << 20;
620
621/// Per-round guard for the floor above. Read failure keeps serving (never a false
622/// refusal from an unreadable device); one `mem_get_info` (~microseconds) per ~25ms round.
623pub(crate) fn graph_launch_headroom_ok(e: &Engine) -> bool {
624    match e.ctx().mem_get_info() {
625        Ok((free, _total)) => free >= GRAPH_LAUNCH_MIN_FREE,
626        Err(_) => true,
627    }
628}
629
630/// One grep-stable suspension line per ROUTE (each call site holds its own
631/// process-lifetime `Once`): every captured-graph launch route below the floor names
632/// itself in the tag while keeping the same `graph replay suspended:` key the step37
633/// admission lane's squeeze cell greps for. The spec-round guard keeps its original
634/// per-generation `[spec]` line; the sweep routes (graph-launch-guard-sweep lane,
635/// 2026-08-31) note once per process — presence is what the gates assert, and a
636/// suspended round is otherwise byte-identical to its eager twin.
637pub(crate) fn graph_replay_suspended_note(route: &str) {
638    eprintln!(
639        "[{route}] graph replay suspended: driver free below the {}MB launch floor \
640         (eager arms serve; cuGraphLaunch segfaults into an exhausted card)",
641        GRAPH_LAUNCH_MIN_FREE / (1 << 20)
642    );
643}
644
645/// Engine-bundle slice 4 (fa-execupdate lane, DSF-ROUNDCOST-20260820 §6 close: "the
646/// residual gap lives in the FULL-ATTENTION per-row section"), DEFAULT ON —
647/// `MEMRA_DSPARK_FA_ROWS=0` reverts to the per-row loop: when every row of a verify
648/// round takes the v4-seqs arm on ONE `fa_split_keys` rung (the straddle law, evaluated
649/// at the round's first and last t_kv — both eligibility gates are intervals in t_kv),
650/// the qwen35 t-parallel verify's per-row KV-append + fa-decode loop collapses into the
651/// z-batched serving twins: ONE `append_quantize_kv_q8_0_q5_1_seqs` + ONE
652/// `fa_decode_vec_q_seqs_v4` + ONE combine per full-attention layer, replacing
653/// T x (4 dtod row copies + append + 3 memsets + main + combine) launches. Bytes are
654/// pinned by the batched-tick increment-2 kernel-check (seqs-vs-per-seq-loop bit
655/// identity: per-row T_kv derives in-kernel from pos_seq[z]; splits >= ns_eff write the
656/// empty partial the combine never reads, so the shared n_splits_max stride changes no
657/// bytes) and re-gated e2e by this lane's battery.
658pub(crate) fn dspark_fa_rows_on() -> bool {
659    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
660    *ON.get_or_init(|| {
661        std::env::var("MEMRA_DSPARK_FA_ROWS")
662            .map(|v| v != "0")
663            .unwrap_or(true)
664    })
665}
666
667/// `t_pred0` for the `MEMRA_DEBUG_SPEC` per-round print, sampled-safe.
668///
669/// `generate_spec_inner2` fills its `preds` vector ONLY on the greedy path (`if !sampled`), and
670/// the per-round debug print was the sole consumer in the sampled arm: `t_pred(0)` survives round
671/// 0 (`base == 0` returns `last_pred`) and from round 1 (`base == 1`, a pending bonus) indexes an
672/// EMPTY vector — `index out of bounds: the len is 0 but the index is 0`, in the GPU worker
673/// thread, which then respawns and reloads weights while the request dies. So any sampled spec
674/// request longer than one round used to kill the worker whenever `MEMRA_DEBUG_SPEC` was set:
675/// the flag crashed precisely the regime it exists to investigate.
676///
677/// Fixed at the print site, not inside the closure, so the greedy accept walk keeps its strict
678/// indexing (an out-of-range pred there is a real bug and must still be loud).
679fn debug_t_pred0(sampled: bool, base: usize, last_pred: u32, preds: &[u32]) -> String {
680    if base == 0 {
681        return last_pred.to_string();
682    }
683    match preds.get(base - 1) {
684        Some(p) => p.to_string(),
685        // sampled: the greedy per-column argmax was never run for this round.
686        None => {
687            debug_assert!(
688                sampled,
689                "greedy spec: preds[{}] missing at base {base}",
690                base - 1
691            );
692            "n/a".to_string()
693        }
694    }
695}
696
697/// `MEMRA_SKEY_PROBE=1` — sampled-draft-graph key probe (lane/graph-s-key-exactness-20260819).
698///
699/// Reports, per burst and per round, which draft chain the sampled arm chose and under which
700/// filter regime, plus the ONE observable that separates a legal filtered draft from a stale
701/// pure-temp graph replayed under filters: an accept test whose gathered `q` is exactly 0.
702/// A draft token sampled from the FILTERED softmax can never gather q=0 (it was drawn from the
703/// kept set), so `q=0` in the verify means the draft came from a distribution the verify does
704/// not believe in — and `u * 0 < p` then accepts it unconditionally.
705///
706/// Its own env var, deliberately NOT `MEMRA_DEBUG_SPEC`: that flag panicked the GPU worker on
707/// any sampled spec request past round 0 until this lane fixed it (§2 of the bank note).
708pub(crate) fn skey_probe() -> bool {
709    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
710    *ON.get_or_init(|| std::env::var("MEMRA_SKEY_PROBE").as_deref() == Ok("1"))
711}
712
713/// GRAMMAR HOOK for constrained spec decode (lane/constrained-full, 2026-08-03). The engine
714/// stays llguidance-agnostic: the server adapts its per-session grammar state behind this
715/// trait. CONTRACT (the verify-side truncation rule — token-identical to constrained plain
716/// greedy decode): the exactness walk runs UNMASKED first; the hook then (a) truncates
717/// acceptance at the first grammar-illegal accepted token, and (b) when the truncation fired
718/// or the bonus is illegal, the engine recomputes that slot as the MASKED argmax of the
719/// target's own verify column (an unmasked argmax that is grammar-legal IS the masked argmax
720/// — masking only removes tokens — so the common case pays nothing). `consume` advances the
721/// state with each EMITTED token in order; EOS handling is the implementor's job (skip).
722pub trait SpecConstraint {
723    /// -inf the current state's banned ids on a HOST logits row (prompt-tail / init-feed
724    /// masked argmax).
725    fn mask_logits(&mut self, logits: &mut [f32]) -> Result<(), String>;
726    /// Packed 32-bit bitset words of the CURRENT state's allowed set (device-mask form).
727    fn mask_words(&mut self) -> Result<Vec<u32>, String>;
728    /// Is `tok` consumable in the CURRENT state?
729    fn is_allowed(&mut self, tok: u32) -> Result<bool, String>;
730    /// Advance the state with an emitted token.
731    fn consume(&mut self, tok: u32) -> Result<(), String>;
732
733    // --- DRAFT-SIDE MASKING (lane/draft-mask, 2026-08-04) ---
734    // The drafter proposed grammar-illegal tokens under tight schemas, so verify-side
735    // truncation cut nearly every round (measured acceptance 0.467-0.513 tight vs 0.62-0.82
736    // loose, research/constrained-full-20260803). These three methods let the engine mask the
737    // DRAFT model's own sampling with the grammar's legal set, so proposals are legal by
738    // construction. The state they walk is a SPECULATIVE CLONE of the session matcher — the
739    // real state is advanced only by `consume` (emitted tokens), so verify-side truncation
740    // stays the correctness backstop and the emitted stream is unchanged by construction
741    // (an accepted draft is the target's unmasked argmax AND grammar-legal, hence the masked
742    // argmax; a cut slot is recomputed as the masked argmax either way).
743    // Default impls = feature OFF (pre-lane behaviour: unmasked drafts).
744
745    /// Is draft-side masking available on this hook? Probed ONCE per burst, before the draft
746    /// graph is captured (the mask is an in-graph node — its presence is a capture-time shape).
747    fn draft_mask_enabled(&self) -> bool {
748        false
749    }
750    /// Start a draft chain: clone the CURRENT (committed) grammar state into the speculative
751    /// slot. Called once per spec round, before the first draft position.
752    fn draft_begin(&mut self) -> Result<(), String> {
753        Ok(())
754    }
755    /// Packed 32-bit bitset words of the SPECULATIVE state's allowed set (target-vocab ids),
756    /// for the draft position about to be sampled. `None` = draft masking off (no-op).
757    fn draft_mask_words(&mut self) -> Result<Option<Vec<u32>>, String> {
758        Ok(None)
759    }
760    /// Advance the SPECULATIVE state with a PROPOSED draft token. `false` = the chain cannot
761    /// continue (EOS proposed, or an unmasked position proposed something illegal) — the
762    /// engine stops drafting; the token already pushed still goes through verify.
763    fn draft_advance(&mut self, _tok: u32) -> Result<bool, String> {
764        Ok(false)
765    }
766}
767
768/// DRAFT-MASK UPLOAD (lane/draft-mask): pull the speculative state's allowed set (TARGET-id
769/// space) from the hook, project it into the DRAFT head's vocab space, and upload it into the
770/// stable device buffer the draft chain reads. Returns false when the chain must stop drafting:
771/// the hook handed out no mask, or NO draft-vocab row is grammar-legal at this position (a
772/// trimmed FR-Spec head genuinely cannot propose a legal token there — masking it would leave
773/// a fully-banned row whose argmax is meaningless, so the round drafts fewer tokens and the
774/// verify emits the masked argmax as usual).
775fn upload_draft_mask(
776    e: &Engine,
777    c: &mut dyn SpecConstraint,
778    dst: &mut CudaSlice<u32>,
779    d2t: Option<&Vec<u32>>,
780    d_vocab: usize,
781    words: usize,
782) -> Result<bool, Box<dyn std::error::Error>> {
783    let Some(tw) = c
784        .draft_mask_words()
785        .map_err(|e2| format!("constraint: {e2}"))?
786    else {
787        return Ok(false);
788    };
789    let bit = |t: usize| -> bool {
790        let w = t >> 5;
791        w < tw.len() && (tw[w] >> (t & 31)) & 1 == 1
792    };
793    let mut buf = vec![0u32; words];
794    match d2t {
795        // TRIMMED draft head: row i proposes target id d2t[i] — permute the mask accordingly.
796        Some(map) => {
797            for (i, &t) in map.iter().enumerate().take(d_vocab) {
798                if bit(t as usize) {
799                    buf[i >> 5] |= 1u32 << (i & 31);
800                }
801            }
802        }
803        // UNTRIMMED: draft ids ARE target ids; the packed words transfer verbatim (a short
804        // mask leaves the padded tail zeroed == banned, same rule as constrained::apply_mask).
805        None => {
806            let n = tw.len().min(words);
807            buf[..n].copy_from_slice(&tw[..n]);
808        }
809    }
810    if buf.iter().all(|w| *w == 0) {
811        return Ok(false);
812    }
813    e.htod_u32_into(dst, &buf)?;
814    Ok(true)
815}
816
817/// Keep the full token-embedding table in host memory and upload only the rows needed by each
818/// MTP/verify step. This is an exact memory-capacity seam for very large BF16 vocab tables: host
819/// gather expands the same source bits to f32, and only O(T*n_embd) bytes cross PCIe per step.
820/// CUDA-graph/round-stream draft paths require device token ids and therefore stay disabled.
821pub(crate) fn spec_host_embd() -> bool {
822    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
823    *ON.get_or_init(|| std::env::var("MEMRA_SPEC_HOST_EMBD").as_deref() == Ok("1"))
824}
825
826/// VERIFY-TIER TRUNK LAUNCH-FUSION (default ON since 2026-07-09; MEMRA_SPEC_FUSED_T=0 reverts — lane/close35b): extend
827/// the t=1 fused2/fused3 Q8_0 trunk launches to the batched verify tier (t=2-4, the K=1..3
828/// verify shapes). At t>1 the trunk pairs/triples (35B wqkv+wqkv_gate, wq/wk/wv,
829/// gate_shexp+up_shexp) each run a separate `matmul_decode_exact` — one q8_1 re-quantize of the
830/// SAME activation plus one _b2/_b4 launch per tensor. The fused twins share ONE quantize and
831/// ONE launch per group; per (tensor,token,row) the kernel body is q8_0_mmvq_batched verbatim
832/// with the identical row mapping -> BIT-IDENTICAL by construction (kernel-check pins it,
833/// run-spec K=1..8 + acceptance identity arbitrate e2e).
834pub(crate) fn spec_fused_t() -> bool {
835    static F: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
836    // DEFAULT ON since 2026-07-09 (MEMRA_SPEC_FUSED_T=0 reverts): verify t=2-4 trunk launch-fusion
837    // (fused2/fused3 Q8_0 batched twins, bit-identical by construction — m=1 block-offset split on
838    // the batched body). m=2 marginal token 2117->1762us; 35B daily: p3 +3.7% (crosses llama), p2 +5%.
839    *F.get_or_init(|| {
840        std::env::var("MEMRA_SPEC_FUSED_T")
841            .map(|v| v != "0")
842            .unwrap_or(true)
843    })
844}
845
846/// zeros/uninit switch for verify-path buffers that are FULLY OVERWRITTEN before any read.
847/// Only call this on such buffers — the lean contract is "identical bytes by construction".
848/// TOKEN-ID GUARD for every id that reaches an embed gather (#87 family).
849///
850/// A device argmax seeds its running index with 0x7FFFFFFF and replaces it only through
851/// comparisons, all of which are FALSE against NaN. An all-NaN logits row therefore returns
852/// the sentinel, and the next thing done with a token id is `embed_row(id)` — table +
853/// ~4.6 TB, never mapped, an MMU fault that kills the CUDA context for the whole process
854/// (research/pp2spec-crash-20260807). The draft chain and the GREEDY verify walk already
855/// trap this; the SAMPLED verify bonus, the boundary sampler and the replay arm's last_pred
856/// did not, which is why the recoverable fault on the greedy instrument is a TERMINAL one on
857/// the vendor-default sampled shape we actually serve.
858pub(crate) fn guard_vocab_token(
859    tok: u32,
860    n_vocab: usize,
861    what: &str,
862) -> Result<u32, Box<dyn std::error::Error>> {
863    if (tok as usize) >= n_vocab {
864        return Err(format!(
865            "{what}: token id 0x{tok:08x} >= n_vocab {n_vocab} — an all-NaN logits row left \
866             the device argmax's init sentinel in place; refusing to dereference the embed \
867             row (#87 trap)"
868        )
869        .into());
870    }
871    Ok(tok)
872}
873
874/// SPEC NaN-ORIGIN SCAN (`MEMRA_SPEC_NAN_SCAN=1`, DEFAULT OFF, diagnostic only).
875///
876/// The `#87` trap reports an all-NaN VERIFY logits column, which says the poison reached the
877/// head but not where it entered. With the scan armed the verify walk syncs and reads back
878/// every layer's output, so the FIRST layer whose residual carries a NaN names itself with the
879/// round's row and position. Off by default and never on a serving path: it costs one host
880/// sync + one `t*n_embd` D2H per layer, and the syncs change scheduling (so a run that stops
881/// reproducing under the scan is itself a datum, not an all-clear).
882///
883/// Rollback seam: unset `MEMRA_SPEC_NAN_SCAN` (or set it to 0). Every call site is behind
884/// `spec_nan_scan()`, so the default path keeps the exact launch sequence it had.
885pub(crate) fn spec_nan_scan() -> bool {
886    spec_nan_scan_level() > 0
887}
888
889/// `MEMRA_SPEC_NAN_SCAN` as a LEVEL, not a boolean. `1` scans each layer's residual, which
890/// names the layer. `2` also scans INSIDE the t-column layer body — the per-column attention
891/// output, the deferred-column o-proj/fa2 join, the post-attention norm and the routed-MoE
892/// output — because "layer 20 poisons row 0" does not say whether the attention or the routed
893/// MoE produced it, and those are different bugs with different fixes.
894pub(crate) fn spec_nan_scan_level() -> u8 {
895    static LVL: std::sync::OnceLock<u8> = std::sync::OnceLock::new();
896    *LVL.get_or_init(|| match std::env::var("MEMRA_SPEC_NAN_SCAN").as_deref() {
897        Ok("1") => 1,
898        Ok("2") => 2,
899        _ => 0,
900    })
901}
902
903/// Read back `[rows, cols]` and fail with the first NaN's coordinates. `what` names the
904/// producer (layer index, walk arm) so the error line is the localization.
905/// VERIFY-ARM RECEIPT (rides `MEMRA_SPEC_NAN_SCAN>=1`, bounded to 200 lines).
906///
907/// Names, per trunk layer, WHICH attention arm the t-column walk actually took. This exists
908/// because the level-1 residual scan below sat only on the non-fused tail: the fused
909/// rope+append+fa arm ends in `continue`, so every layer that fused was NEVER SCANNED and
910/// silently read as "clean". A poisoned residual therefore first reported at the next
911/// non-fused layer, which is how "layer 20 creates the poison" could be true of the scan and
912/// false of the engine. Also carries the row-table lookup counter, so "the fused path never
913/// ran" is distinguishable from "it ran and was innocent".
914/// KV-PLANE SCAN (`MEMRA_KV_PLANE_SCAN=1`, DEFAULT OFF, diagnostic only).
915///
916/// Reads back the STAGED rows of a layer's distributed K/V planes and reports the first row
917/// whose quantization scale is not finite. No kernel required: q8_0 blocks are
918/// `[half d][32 x i8]` and q5_1 blocks carry `half d` then `half m`, so the fp16 scale at the
919/// head of each block is host-checkable straight out of the byte plane.
920///
921/// It exists because the level-2 bad-row bitmap says EVERY verify row is non-finite at a
922/// global-attention layer's join, and row r attends a strict superset of row r-1's keys: that
923/// implicates the shared KV history those rows walk, not per-column staging. "The attention
924/// output is NaN" and "the KV history it attends is already NaN" are different bugs with
925/// different owners, and nothing measured so far separates them. A first-corrupt-row index
926/// also dates the corruption against the prime/decode boundary.
927///
928/// Bounded hard: only layers whose geometry has NO window (the global planes), only the first
929/// `MEMRA_KV_PLANE_SCAN_ROUNDS` verify rounds of a process (default 2), and it copies only
930/// `[0, staged_len)`, which is ~1.6 MB at the 1480-token repro rather than the 262144-row
931/// provision. It still syncs per layer, so it is never a serving or a measured-perf arm.
932pub(crate) fn kv_plane_scan_on() -> bool {
933    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
934    *ON.get_or_init(|| std::env::var("MEMRA_KV_PLANE_SCAN").as_deref() == Ok("1"))
935}
936
937fn kv_plane_scan_rounds() -> usize {
938    static R: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
939    *R.get_or_init(|| {
940        std::env::var("MEMRA_KV_PLANE_SCAN_ROUNDS")
941            .ok()
942            .and_then(|v| v.parse().ok())
943            .unwrap_or(2)
944    })
945}
946
947/// First non-finite fp16 block scale in `bytes`, as (block index, raw u16), scanning one
948/// scale every `stride` bytes. Returns None when every block scale is finite.
949fn first_bad_scale(bytes: &[u8], stride: usize) -> Option<(usize, u16)> {
950    if stride == 0 {
951        return None;
952    }
953    for (i, blk) in bytes.chunks_exact(stride).enumerate() {
954        let raw = u16::from_le_bytes([blk[0], blk[1]]);
955        if half_is_non_finite(raw) {
956            return Some((i, raw));
957        }
958    }
959    None
960}
961
962/// IEEE binary16: exponent all ones is Inf or NaN, whatever the mantissa says.
963fn half_is_non_finite(raw: u16) -> bool {
964    (raw & 0x7C00) == 0x7C00
965}
966
967/// Scan one layer's staged K/V planes for a non-finite quantization scale. Returns the
968/// receipt line, or None when the layer is out of scope or every scale is finite.
969pub(crate) fn scan_kv_plane(
970    e: &crate::Engine,
971    distributed: &memra_kv::ResidentTpKvCache,
972    il: usize,
973    pos0: usize,
974) -> Result<(), Box<dyn std::error::Error>> {
975    // One "round" is one pos0, not one layer: the walk visits 45 layers per verify. The
976    // default of 2 rounds is for a fault that shows up immediately; the step37 repro does not
977    // fire until rep 3 or later, i.e. round ~60 of the process, so that arm MUST raise
978    // MEMRA_KV_PLANE_SCAN_ROUNDS or it will scan only the two rounds that were never going to
979    // be poisoned and report a clean history it never looked at.
980    static ROUNDS: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
981    static LAST_POS: std::sync::atomic::AtomicUsize =
982        std::sync::atomic::AtomicUsize::new(usize::MAX);
983    if LAST_POS.swap(pos0, std::sync::atomic::Ordering::Relaxed) != pos0 {
984        ROUNDS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
985    }
986    if ROUNDS.load(std::sync::atomic::Ordering::Relaxed) > kv_plane_scan_rounds() {
987        return Ok(());
988    }
989    let staged = distributed.staged_len();
990    if staged == 0 {
991        return Ok(());
992    }
993    // ENGAGEMENT RECEIPT. This scan prints only on corruption, so `kvbad=0` in a cell would
994    // read the same whether the history was clean or the scan never ran once. Bounded so a
995    // 45-layer walk cannot flood the log.
996    static SEEN: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
997    let seen = SEEN.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
998    let (ktb, vtb) = (distributed.k_tok_bytes(), distributed.v_tok_bytes());
999    if seen < 4 {
1000        eprintln!(
1001            "[kv-plane] engaged #{seen} layer {il} pos0={pos0} staged={staged} \
1002             ktok={ktb} vtok={vtb} (scan armed; a corrupt plane prints its own line)"
1003        );
1004    }
1005    for rank in 0..distributed.ranks().len() {
1006        let Some(rc) = distributed.rank(rank) else {
1007            continue;
1008        };
1009        // q8_0 K blocks are [half d][32 x i8] = 34B; q5_1 V blocks lead with half d then half m.
1010        let kbytes = e.dtoh_u8_view(&rc.k().slice(0..staged * ktb))?;
1011        let vbytes = e.dtoh_u8_view(&rc.v().slice(0..staged * vtb))?;
1012        let kbad = first_bad_scale(&kbytes, 34);
1013        let vbad = first_bad_scale(&vbytes, 24);
1014        if kbad.is_some() || vbad.is_some() {
1015            let row = |b: Option<(usize, u16)>, tok: usize| {
1016                b.map(|(i, raw)| format!("blk {i} (row {}) raw={raw:#06x}", i * 34 / tok.max(1)))
1017                    .unwrap_or_else(|| "clean".into())
1018            };
1019            eprintln!(
1020                "[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",
1021                row(kbad, ktb),
1022                row(vbad, vtb)
1023            );
1024            return Ok(());
1025        }
1026    }
1027    Ok(())
1028}
1029
1030pub(crate) fn verify_arm_receipt(
1031    arm: &str,
1032    il: usize,
1033    pos0: usize,
1034    t: usize,
1035    staged: Option<usize>,
1036) {
1037    static N: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
1038    if N.fetch_add(1, std::sync::atomic::Ordering::Relaxed) >= 200 {
1039        return;
1040    }
1041    eprintln!(
1042        "[verify-arm] layer {il} arm={arm} pos0={pos0} t={t} staged_len={} rows_tab_lookups={}",
1043        staged.map(|v| v as i64).unwrap_or(-1),
1044        crate::tp::ROWS_TAB_ENGAGED.load(std::sync::atomic::Ordering::Relaxed)
1045    );
1046}
1047
1048pub(crate) fn nan_scan_rows(
1049    e: &Engine,
1050    buf: &CudaSlice<f32>,
1051    rows: usize,
1052    cols: usize,
1053    what: &str,
1054) -> Result<(), Box<dyn std::error::Error>> {
1055    // The readback is also the ATTRIBUTION point for an asynchronous fault: a
1056    // CUDA_ERROR_ILLEGAL_ADDRESS raised by any launch since the previous scan surfaces on this
1057    // sync, and the bare DriverError names nothing. Wrapping it with `what` turns "the process
1058    // died somewhere" into "it died at or before this layer, on this row, at this position".
1059    let host = e.dtoh(buf).map_err(|err| -> Box<dyn std::error::Error> {
1060        format!(
1061            "spec nan-scan: sync at {what} FAILED: {err} — the fault is at or before \
1062                     this point in the walk"
1063        )
1064        .into()
1065    })?;
1066    if host.len() < rows * cols {
1067        return Err(format!(
1068            "nan-scan {what}: buffer holds {} < {rows}x{cols}",
1069            host.len()
1070        )
1071        .into());
1072    }
1073    // SCAN EVERY ROW BEFORE REPORTING. A first-hit return says "row 0 is bad" and leaves the
1074    // other rows UNEXAMINED, which is exactly the bit that discriminates the two mechanisms: in
1075    // the t-column verify, row 0 attends keys [0..p+1) and row 1 attends [0..p+2), a strict
1076    // superset, so poison in the SHARED KV history must appear in BOTH rows, while poison in
1077    // per-column staging can appear in one. Report the whole map.
1078    let mut per_row: Vec<usize> = Vec::with_capacity(rows);
1079    let mut first_bad: Option<(usize, usize)> = None;
1080    for r in 0..rows {
1081        let row = &host[r * cols..(r + 1) * cols];
1082        let bad = row.iter().filter(|v| !v.is_finite()).count();
1083        per_row.push(bad);
1084        if bad > 0 && first_bad.is_none() {
1085            first_bad = Some((r, row.iter().position(|v| !v.is_finite()).unwrap_or(0)));
1086        }
1087    }
1088    if let Some((r0, c0)) = first_bad {
1089        let map: String = per_row
1090            .iter()
1091            .map(|&b| if b == 0 { '.' } else { 'X' })
1092            .collect();
1093        return Err(format!(
1094            "spec nan-scan: {what} produced non-finite values — rows[{rows}] map={map} \
1095             counts={per_row:?} of {cols} each; first at row {r0} element {c0}. Both rows bad \
1096             implicates shared state (the KV history this layer reads); one row bad implicates \
1097             per-column staging."
1098        )
1099        .into());
1100    }
1101    Ok(())
1102}
1103
1104fn vbuf(e: &Engine, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
1105    if spec_lean() { e.uninit(n) } else { e.zeros(n) }
1106}
1107
1108/// Scratch KV for the MTP block (one full-attn layer).
1109///
1110/// PERSISTENT MODE (default, 2026-07-03 — the acceptance lever): sized cap = max_ctx and kept in
1111/// sync with the COMMITTED sequence — slot p holds the MTP block's K/V for committed token p
1112/// (roped p+1, the chain's rope convention), so the draft chain's self-attention sees the FULL
1113/// committed history instead of only the current round's 1..K+1 chain tokens (the reference
1114/// engine's "mtp_update" design). Entries come from two sources:
1115///   - chain appends: accepted positions KEEP their chain-computed entries (embedding exact,
1116///     hidden chain-approximate — the reference engine accepts the same);
1117///   - `mtp_kv_fill` batches: prompt positions + the last-draft position on full accept, computed
1118///     from EXACT trunk hiddens (K/V-only MTP-block pass, no attention/FFN/lm_head).
1119/// Rejected drafts / p-min extras / pseudo-seed appends are all discarded by the round-start
1120/// `set_len` truncation (the KvLayer len mechanism — §C rollback for the draft side).
1121/// Multi-turn spec-decode session (2026-07-05): trunk Cache + persistent MTP draft scratch +
1122/// the committed token list, alive across generate_spec_session calls. Turn N+1 primes ONLY its
1123/// suffix (chunked continuation prime over the quantized past) and mtp_kv_fill's its suffix rows,
1124/// then the round loop runs unchanged. `last_h` carries the pre-output_norm hidden of the last
1125/// committed row across turns (the predecessor-pairing seed + fill anchor).
1126/// Per-request sampling config for the sampled-spec serve path.
1127#[derive(Clone, Copy, Debug)]
1128pub struct SpecSampling {
1129    pub temp: f32,
1130    pub seed: u64,
1131    pub top_k: i32,            // 0 = off
1132    pub top_p: f32,            // 1.0 = off
1133    pub min_p: f32,            // 0.0 = off
1134    pub penalty_last_n: usize, // 0 = penalties off
1135    pub penalty_repeat: f32,
1136    pub penalty_freq: f32,
1137    pub penalty_present: f32,
1138}
1139
1140impl SpecSampling {
1141    /// Non-identity penalties requested — THE `pen_on` predicate (one definition; the
1142    /// same group-off rule `SamplerIdentity::of` canonicalizes: a window with neutral
1143    /// coefficients is penalties-absent). Both spec routes and the dspark accept walk
1144    /// key their penalty arms off this.
1145    pub fn pen_on(&self) -> bool {
1146        self.penalty_last_n > 0
1147            && (self.penalty_repeat != 1.0
1148                || self.penalty_freq != 0.0
1149                || self.penalty_present != 0.0)
1150    }
1151}
1152
1153/// Host Philox4x32-10 uniform in (0,1) — mirrors spec_sample.cu's `philox4`/`u01` with the
1154/// ctr_lo tag 0xFFFF_FFFE, so the host accept-test stream never collides with any device
1155/// sampling event (device Gumbel uses (i>>2, stream_pos); device residual uses 0xFFFF_FFFD).
1156/// One value per (seed, ctr) EVENT; callers own the counter discipline. Extracted verbatim
1157/// from generate_spec_inner2's closure for the dspark sampled-admission walk (the two paths
1158/// MUST consume the identical stream construction — two ad-hoc Philox copies drifting apart
1159/// is a distributional bug, not a style problem).
1160pub(crate) fn host_u01(seed: u64, ctr: u32) -> f32 {
1161    let (m0, m1) = (0xD2511F53u32, 0xCD9E8D57u32);
1162    let (mut c0, mut c1, mut c2, mut c3) = (0xFFFF_FFFEu32, ctr, 0u32, 0u32);
1163    let (mut k0, mut k1) = ((seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
1164    for _ in 0..10 {
1165        let (h0, l0) = (((m0 as u64 * c0 as u64) >> 32) as u32, m0.wrapping_mul(c0));
1166        let (h1, l1) = (((m1 as u64 * c2 as u64) >> 32) as u32, m1.wrapping_mul(c2));
1167        let (n0, n1, n2, n3) = (h1 ^ c1 ^ k0, l1, h0 ^ c3 ^ k1, l0);
1168        c0 = n0;
1169        c1 = n1;
1170        c2 = n2;
1171        c3 = n3;
1172        k0 = k0.wrapping_add(0x9E3779B9);
1173        k1 = k1.wrapping_add(0xBB67AE85);
1174    }
1175    (c0 as f32 + 1.0) * (1.0 / 4294967296.0)
1176}
1177
1178/// Tracked draft positions for [`SpecTelemetry`] (serve K defaults to 3; the run-spec gate
1179/// sweeps K=1..8, and MEMRA_SPEC_CAPMAX defaults to 7 — 8 covers every tuned config).
1180pub const SPEC_TELEM_POS: usize = 8;
1181
1182/// Always-on per-draft-position acceptance telemetry (lane/accept-telemetry, 2026-08-05 —
1183/// the llama.cpp #26389 / vLLM spec-decode counter schema, upstream-sweeps 2026-08-05).
1184/// Lives on the [`SpecSession`] and accumulates across bursts; the serve worker diffs a
1185/// stashed copy per burst for its per-model /metrics aggregation and per-request usage.
1186/// Same normalization as the `[spec-stats]` line: p-min-discarded chain tokens are counted
1187/// in NEITHER drafted nor accepted.
1188#[derive(Clone, Copy, Default, Debug)]
1189pub struct SpecTelemetry {
1190    /// verify rounds completed (a round-stream burst counts each of its M rounds).
1191    pub rounds: u64,
1192    /// tokens drafted / accepted across all rounds.
1193    pub drafted: u64,
1194    pub accepted: u64,
1195    /// how often draft position j (0-based within a round's chain) was offered / accepted.
1196    /// Positions >= SPEC_TELEM_POS are untracked (totals still count them). The opt-in
1197    /// round-stream arm (MEMRA_SPEC_STREAM=1) reads back only totals, so under it these
1198    /// arrays cover the standard-path rounds only and their sums may undercount the totals.
1199    pub pos_drafted: [u64; SPEC_TELEM_POS],
1200    pub pos_accepted: [u64; SPEC_TELEM_POS],
1201}
1202
1203impl SpecTelemetry {
1204    /// Fieldwise `self - prev` — the worker's per-burst delta off a copy stashed before the
1205    /// burst call. Saturating: a caller diffing against the wrong snapshot gets zeros, not
1206    /// a wrapped counter.
1207    pub fn delta_since(&self, prev: &SpecTelemetry) -> SpecTelemetry {
1208        let mut d = SpecTelemetry {
1209            rounds: self.rounds.saturating_sub(prev.rounds),
1210            drafted: self.drafted.saturating_sub(prev.drafted),
1211            accepted: self.accepted.saturating_sub(prev.accepted),
1212            ..Default::default()
1213        };
1214        for j in 0..SPEC_TELEM_POS {
1215            d.pos_drafted[j] = self.pos_drafted[j].saturating_sub(prev.pos_drafted[j]);
1216            d.pos_accepted[j] = self.pos_accepted[j].saturating_sub(prev.pos_accepted[j]);
1217        }
1218        d
1219    }
1220    /// Fieldwise `self += d` — the worker's per-model aggregation.
1221    pub fn merge(&mut self, d: &SpecTelemetry) {
1222        self.rounds += d.rounds;
1223        self.drafted += d.drafted;
1224        self.accepted += d.accepted;
1225        for j in 0..SPEC_TELEM_POS {
1226            self.pos_drafted[j] += d.pos_drafted[j];
1227            self.pos_accepted[j] += d.pos_accepted[j];
1228        }
1229    }
1230
1231    /// Mean accepted draft-prefix length per verify round (tau).
1232    pub fn tau(&self) -> f64 {
1233        if self.rounds > 0 {
1234            self.accepted as f64 / self.rounds as f64
1235        } else {
1236            0.0
1237        }
1238    }
1239}
1240
1241/// Session-lifetime atomic acceptance counters. The verifier records only after the greedy or
1242/// rejection-sampling walk has resolved on the host, so these relaxed increments add no GPU
1243/// launch, synchronization, allocation, or ordering dependency to the numeric path.
1244struct SpecTelemetryCounters {
1245    rounds: AtomicU64,
1246    drafted: AtomicU64,
1247    accepted: AtomicU64,
1248    pos_drafted: [AtomicU64; SPEC_TELEM_POS],
1249    pos_accepted: [AtomicU64; SPEC_TELEM_POS],
1250}
1251
1252impl Default for SpecTelemetryCounters {
1253    fn default() -> Self {
1254        Self {
1255            rounds: AtomicU64::new(0),
1256            drafted: AtomicU64::new(0),
1257            accepted: AtomicU64::new(0),
1258            pos_drafted: std::array::from_fn(|_| AtomicU64::new(0)),
1259            pos_accepted: std::array::from_fn(|_| AtomicU64::new(0)),
1260        }
1261    }
1262}
1263
1264impl SpecTelemetryCounters {
1265    fn record_round(&self, drafted: usize, accepted: usize) {
1266        debug_assert!(accepted <= drafted);
1267        self.rounds.fetch_add(1, Ordering::Relaxed);
1268        self.drafted.fetch_add(drafted as u64, Ordering::Relaxed);
1269        self.accepted.fetch_add(accepted as u64, Ordering::Relaxed);
1270        for counter in self.pos_drafted.iter().take(drafted) {
1271            counter.fetch_add(1, Ordering::Relaxed);
1272        }
1273        for counter in self.pos_accepted.iter().take(accepted) {
1274            counter.fetch_add(1, Ordering::Relaxed);
1275        }
1276    }
1277
1278    /// Round-stream keeps each round's accept length on device; retain exact scalar totals while
1279    /// leaving the per-position arrays untouched, matching the pre-existing telemetry contract.
1280    fn record_totals(&self, rounds: usize, drafted: usize, accepted: usize) {
1281        self.rounds.fetch_add(rounds as u64, Ordering::Relaxed);
1282        self.drafted.fetch_add(drafted as u64, Ordering::Relaxed);
1283        self.accepted.fetch_add(accepted as u64, Ordering::Relaxed);
1284    }
1285
1286    fn snapshot(&self) -> SpecTelemetry {
1287        SpecTelemetry {
1288            rounds: self.rounds.load(Ordering::Relaxed),
1289            drafted: self.drafted.load(Ordering::Relaxed),
1290            accepted: self.accepted.load(Ordering::Relaxed),
1291            pos_drafted: std::array::from_fn(|j| self.pos_drafted[j].load(Ordering::Relaxed)),
1292            pos_accepted: std::array::from_fn(|j| self.pos_accepted[j].load(Ordering::Relaxed)),
1293        }
1294    }
1295}
1296
1297pub struct SpecSession {
1298    pub(crate) cache: Cache,
1299    pub(crate) scratch: MtpScratch,
1300    /// Every token whose state the caches hold, in order (prompt turns + generated), INCLUDING
1301    /// overshoot: spec commits accepted drafts past max_new; those rows are in the caches, so the
1302    /// session must count them. Callers render output from this, not from their own echo.
1303    pub committed: Vec<u32>,
1304    /// Pre-output_norm hidden of the LAST committed row (device). None before the first turn.
1305    pub(crate) last_h: Option<CudaSlice<f32>>,
1306    /// Greedy argmax predicting the token AFTER committed.last() (from the last turn's final
1307    /// logits). Fuels empty-suffix continuation bursts (serve): the next turn emits this token
1308    /// first, feeds it, and the round loop resumes without any prime. None before the first turn.
1309    pub next_pred: Option<u32>,
1310    /// SAMPLED-SPEC stream continuity across bursts: Philox event counters persist here so a
1311    /// session's randomness never repeats between generate_spec_session calls. (0,0) at admit.
1312    pub sctr: u32,
1313    pub uctr: u32,
1314    /// PERSISTENT DRAFT-GRAPH CONTEXT (2026-08-01, the serve-burst fixed-cost fix): the captured
1315    /// draft graph(s) + every device I/O buffer they bake, carried ACROSS generate_spec_session
1316    /// calls. Before this, every serve burst re-captured the draft graph (2 warmup forwards +
1317    /// instantiate) — measured ~16ms/burst on H100 q27 (MEMRA_SPEC_BURST sweep,
1318    /// research/spec-serving-20260801). None before the first turn; error paths drop it
1319    /// (next burst recaptures — serve retires errored sessions anyway).
1320    pub(crate) draft_ctx: Option<DraftGraphCtx>,
1321    /// PENDING-CARRY across bursts (2026-08-01, the serve burst-boundary fix): the bonus token
1322    /// emitted by the last round but NOT committed to the caches. The old tail committed it with
1323    /// a solo T=1 trunk pass (+ draft fill), and the next burst's setup fed the stashed next_pred
1324    /// with ANOTHER solo pass — 2x ~11.5ms/burst measured on H100 q27 ([spec-setup] trace).
1325    /// Carrying it lets the next empty-suffix greedy burst consume it as round-0 verify col 0,
1326    /// exactly like a mid-burst full-accept boundary (no solo passes). INVARIANT: when set,
1327    /// `committed` (== cache rows) EXCLUDES this token although it was already emitted in the
1328    /// last burst's output, and `last_h` holds the hidden of the last COMMITTED row (its
1329    /// predecessor — the chain-seed/fill anchor). `next_pred` is None (unknown without the
1330    /// commit pass). Non-empty-suffix or sampled turns must flush first (spec_flush_pending);
1331    /// generate_spec_session_sampled does this at entry, and serve parks only flushed sessions.
1332    pub pending_tok: Option<u32>,
1333    /// SESSION-AFFINITY TURN CHECKPOINT (lane/session-affinity, 2026-08-05): the state at this
1334    /// turn's PROMPT-END boundary, retained so a later turn can REWIND here. See
1335    /// [`SpecCheckpoint`]. Refreshed by every non-empty prime; None until the first one, and on
1336    /// a rig too tight to hold it (a failed capture is silent — resume just isn't available).
1337    pub(crate) turn_ckpt: Option<SpecCheckpoint>,
1338    /// Session-lifetime acceptance telemetry. Relaxed atomics update at the host-side round
1339    /// accounting the loop already does — no syncs, no allocation. NOTE a
1340    /// pool-resumed session carries the PREVIOUS requests' counts; per-request consumers
1341    /// diff with [`SpecTelemetry::delta_since`] around each burst.
1342    telem: SpecTelemetryCounters,
1343    /// PREFIX-CACHE publication request (lane/spec-prefix-cache): worker sets this to the
1344    /// miss-LCP boundary before a cold burst; the prime captures at exactly that split (it must
1345    /// coincide with the burst's `prime_split` or no capture happens). One-shot: consumed by the
1346    /// prime, result lands in `boundary_captures`.
1347    pub capture_at: Option<usize>,
1348    /// The captures the last prime produced (see [`SpecBoundaryCapture`]). Worker drains them
1349    /// post-burst to assemble prefix entries. A failed capture is silent, like `turn_ckpt` —
1350    /// publication just isn't available for that request. Plural since
1351    /// lane/frspec-multiturn-cache (2026-08-21): a cold burst can capture BOTH the miss-LCP
1352    /// split (the shared-prefix class) and the stable pre-generation boundary (the
1353    /// next-turn re-render class) — one entry per stop, exactly the boundary set the plain
1354    /// prefill tick publishes/checkpoints.
1355    pub boundary_captures: Vec<SpecBoundaryCapture>,
1356    /// STABLE-BOUNDARY TURN CHECKPOINT REQUEST (lane/frspec-multiturn-cache, 2026-08-21): the
1357    /// ABSOLUTE committed-length position the next non-empty prime should capture `turn_ckpt`
1358    /// at, instead of prompt-end. The worker sets it to the STABLE PRE-GENERATION boundary
1359    /// (`plain_checkpoint_boundary` — before the live generation header the client rewrites),
1360    /// porting the 2026-08-09 plain-tier fix: a prompt-end spec checkpoint includes the
1361    /// template's live assistant-generation header (`<|im_start|>assistant\n<think>\n`), which
1362    /// the NEXT turn's re-render replaces, so `affinity_match` diverged a couple tokens below
1363    /// the checkpoint and the spec pool declined 100% of multi-turn agent traffic (measured:
1364    /// `spec-affinity: declined (history diverged at 6811 of checkpoint 6813)`,
1365    /// research/multiturn-cache-20260821 B4). One-shot, `capture_at` convention; None = legacy
1366    /// prompt-end capture.
1367    pub ckpt_at: Option<usize>,
1368    /// FAIL-SAFE (lane/step37-vram-admission-20260830, external-review corroboration): set
1369    /// by the worker on a session serving a step-OOM park REPLAY. The burst entry pre-marks
1370    /// the draft-graph fallback so the replay never re-enters the capture path — the capture
1371    /// appetite is part of what drove the card to the OOM, and a replay that recaptures
1372    /// re-runs the incident. If the eager replay still cannot fit, the bounded retry budget
1373    /// exhausts into the honest recoverable Overloaded error instead of looping.
1374    pub capture_disabled: bool,
1375}
1376impl SpecSession {
1377    /// Context capacity of the session's caches (the server's ContextFull guard).
1378    pub fn cache_max_ctx(&self) -> usize {
1379        self.cache.max_ctx
1380    }
1381    /// Read access to the live trunk cache (lane/spec-prefix-cache): the worker slices
1382    /// full-attn KV rows `[0..capture.pos)` out of it when publishing a boundary capture —
1383    /// those rows are append-only for the session's lifetime (rollbacks never truncate below
1384    /// the prime boundary), so no copy was taken at prime time.
1385    pub fn cache_ref(&self) -> &Cache {
1386        &self.cache
1387    }
1388    /// Read access to the persistent draft-scratch plane (lane/spec-on-cache-hit): the
1389    /// worker slices rows `[0..capture.pos)` when publishing a boundary capture, exactly
1390    /// like the trunk KV — draft rows below the prompt end are append-only for the
1391    /// session's lifetime (the prime fill wrote them once; rollbacks reset `len_d` to the
1392    /// committed length, never below the prime boundary, and the true-hidden refresh
1393    /// rewrites generated positions only). Returns `(k, v, k_tok_bytes, v_tok_bytes)`.
1394    /// None when the scratch is ring-backed (Step35 SWA — physical rows are not
1395    /// prefix-addressable; the prefix cache already refuses that class end to end).
1396    pub fn draft_plane_ref(&self) -> Option<(&CudaSlice<u8>, &CudaSlice<u8>, usize, usize)> {
1397        if self.scratch.kv.ring.is_some() {
1398            return None;
1399        }
1400        Some((
1401            &self.scratch.kv.k,
1402            &self.scratch.kv.v,
1403            self.scratch.kv.k_tok_bytes,
1404            self.scratch.kv.v_tok_bytes,
1405        ))
1406    }
1407    /// Snapshot the session's process-local acceptance counters for per-burst diffing.
1408    pub fn telemetry(&self) -> SpecTelemetry {
1409        self.telem.snapshot()
1410    }
1411    /// Committed position this session can REWIND to (its retained prompt-end boundary), if any.
1412    /// A request whose prompt matches `committed[..pos]` exactly can resume from here — see
1413    /// `spec_rewind_to_checkpoint`.
1414    pub fn rewind_pos(&self) -> Option<usize> {
1415        self.turn_ckpt.as_ref().map(|c| c.pos)
1416    }
1417    /// Whether every ring-backed trunk/draft row needed by the retained checkpoint is resident.
1418    pub fn rewind_is_resident(&self) -> bool {
1419        self.turn_ckpt.as_ref().is_some_and(|ckpt| {
1420            self.cache.can_rollback(&ckpt.snap, 0) && self.scratch.can_rewind_to(ckpt.pos)
1421        })
1422    }
1423    /// Is this session in the DEMOTION-READY shape (see [`SpecSession::into_demoted`])?
1424    /// `false` means a carried pending must be flushed first (`spec_flush_pending`), or the
1425    /// session has never run a turn and has no prediction to hand over.
1426    pub fn demote_ready(&self) -> bool {
1427        self.pending_tok.is_none() && self.next_pred.is_some()
1428    }
1429    /// Does this session hold a carried pending bonus (flush required before a handoff/park)?
1430    pub fn has_pending(&self) -> bool {
1431        self.pending_tok.is_some()
1432    }
1433    /// Committed row count == cache rows (the session invariant), for the caller's own
1434    /// `fed`-length cross-check at a handoff boundary.
1435    pub fn committed_len(&self) -> usize {
1436        self.committed.len()
1437    }
1438    /// DEMOTION HANDOFF (lane/spec-gate, 2026-08-07): consume this session and hand its trunk
1439    /// cache + next-token prediction to the plain batched-decode path.
1440    ///
1441    /// WHY THIS IS EXACT (greedy). The invariant at a burst boundary is `cache.pos ==
1442    /// committed.len()`: every committed row has trunk KV + recurrent state, exactly as a plain
1443    /// tokenwise prime of the same `committed` sequence would have left it (that is the
1444    /// session-tail contract, and the same property `spec_rewind_to_checkpoint` and the reuse
1445    /// pool already rely on). `next_pred` is the argmax of the verify's logits for the LAST
1446    /// committed row — and verify-column logits are bit-identical to plain decode's logits at
1447    /// that position, because `matmul_decode_exact` bit-identity IS the basis of the greedy
1448    /// accept walk. So handing (cache, next_pred) to the batched path continues the stream from
1449    /// a state indistinguishable from one the batched path produced itself: the batched tick
1450    /// emits `next_pred`, feeds it into this same cache, and decodes on.
1451    ///
1452    /// `None` when the session is not in the handoff shape — a carried pending (its bonus row is
1453    /// NOT in the cache, so `spec_flush_pending` must commit it first) or no `next_pred` yet
1454    /// (never bursted). Callers must not force it: a half-committed cache handed to the batched
1455    /// path would silently skip a token.
1456    ///
1457    /// The MTP draft scratch, the persistent draft-graph context and the turn checkpoint are
1458    /// DROPPED here (freeing their VRAM): the batched path never drafts, and this handoff is
1459    /// one-way by design — there is no cheap symmetric re-promotion (rebuilding the draft KV
1460    /// would mean an `mtp_kv_fill` over the whole committed history).
1461    pub fn into_demoted(self) -> Option<(Cache, u32)> {
1462        if self.pending_tok.is_some() {
1463            return None;
1464        }
1465        let np = self.next_pred?;
1466        debug_assert_eq!(
1467            self.cache.pos,
1468            self.committed.len(),
1469            "demotion handoff: cache rows != committed tokens"
1470        );
1471        Some((self.cache, np))
1472    }
1473    /// Pool-resume hook (audit Q2): clear the parked draft-graph failure memoization so a
1474    /// NEW request resuming this session gets one fresh capture chance — a transient-pressure
1475    /// capture failure must not persist for the pool's whole lifetime (the TRT #16072 class).
1476    /// Logs once iff a flag was actually set; a no-fallback resume is silent and free.
1477    pub fn reset_graph_fallback_on_resume(&mut self) {
1478        if let Some(line) = self
1479            .draft_ctx
1480            .as_mut()
1481            .and_then(|c| c.failed.reset_on_resume())
1482        {
1483            eprintln!("{line}");
1484        }
1485    }
1486}
1487
1488/// A session's PROMPT-END boundary state, the rewind target for session-affinity resume.
1489///
1490/// WHY THIS BOUNDARY, AND WHY IT IS THE ONLY ONE WORTH KEEPING. The rewrite class this lane
1491/// exists for (a client that strips `<think>` blocks out of prior assistant turns) mutates the
1492/// text the session GENERATED, never the prompt it was given. So turn N's prompt agrees with
1493/// turn N-1's committed tokens up to almost exactly where turn N-1's generation began — the
1494/// prompt-end boundary. Keeping a checkpoint there means the next turn re-primes only its own
1495/// delta (the rewritten answer + the new user turn) instead of the whole conversation.
1496///
1497/// WHAT IT MUST HOLD. Full-attn KV is append-only and position-addressed, so rewinding it is a
1498/// `len` truncation (no data). Linear-attn (GDN) conv/ssm state is mutated IN PLACE with no
1499/// position index, so it must be a real device COPY — that copy is the entire reason a spec
1500/// session could not previously rewind. The MTP draft scratch needs no copy either: its rows
1501/// below the boundary were written by this turn's fill and are never revisited (the per-round
1502/// true-hidden refresh only rewrites the CURRENT burst's committed positions), so rewinding it
1503/// is also just a `len` reset. `last_h` is the hidden of the last row below the boundary — the
1504/// predecessor-pairing anchor the next prime's fill reads for its first row.
1505///
1506/// COST: one `Cache::snapshot` per TURN, on a code path that already takes one per ROUND.
1507pub(crate) struct SpecCheckpoint {
1508    snap: crate::cache::CacheSnapshot,
1509    /// Committed length at the boundary (== cache.pos there, the session invariant).
1510    pos: usize,
1511    /// Pre-output_norm hidden of row `pos - 1`.
1512    last_h: CudaSlice<f32>,
1513}
1514
1515/// PREFIX-CACHE BOUNDARY CAPTURE (lane/spec-prefix-cache, 2026-08-14): the state a spec session
1516/// records at its cold-prime split so the WORKER can publish a cross-request prefix entry —
1517/// the commit-gated-publication port (research/cache-spec-design-20260814/PORT-PLAN.md item 1).
1518/// Only the pieces that are DESTROYED by continuing the prime need copies here: the in-place
1519/// GDN conv/ssm states (via `Cache::snapshot`, same mechanism as [`SpecCheckpoint`]) and the
1520/// boundary logits. Full-attn KV rows `[0..pos)` and draft-scratch rows `[0..pos)` are
1521/// append-only for the session's lifetime (rollbacks never truncate below the prime boundary),
1522/// so the worker slices those from the live caches post-burst instead of copying at prime time.
1523pub struct SpecBoundaryCapture {
1524    pub snap: crate::cache::CacheSnapshot,
1525    /// Token boundary (== cache.pos at capture; == the worker's miss-LCP split).
1526    pub pos: usize,
1527    /// Full-vocab logits after the prefix prime — the entry's boundary logits.
1528    pub logits: Vec<f32>,
1529    /// Pre-output_norm trunk hidden of row `pos - 1` (lane/spec-on-cache-hit): the
1530    /// predecessor-pairing anchor a RESTORED spec session's first suffix-fill row reads
1531    /// (the `SpecSession::last_h` convention). Empty = unavailable (capture stays valid;
1532    /// the fill's zeros row-0 fallback covers it at a bounded acceptance cost).
1533    pub last_h: Vec<f32>,
1534}
1535
1536/// D2H one hidden row out of a `[T, n_embd]` prime hidden stack — the boundary anchor a
1537/// spec boundary capture carries for later restored-session fills. Failure is silent
1538/// (`turn_ckpt` convention): the capture publishes without an anchor.
1539fn capture_boundary_hidden(
1540    e: &Engine,
1541    h_rows: &CudaSlice<f32>,
1542    pos: usize,
1543    n_embd: usize,
1544) -> Vec<f32> {
1545    if pos == 0 || h_rows.len() < pos * n_embd {
1546        return Vec::new();
1547    }
1548    let Ok(mut row) = e.uninit(n_embd) else {
1549        return Vec::new();
1550    };
1551    if e.copy_view_into(
1552        &mut row,
1553        0,
1554        &h_rows.slice((pos - 1) * n_embd..pos * n_embd),
1555        n_embd,
1556    )
1557    .is_err()
1558    {
1559        return Vec::new();
1560    }
1561    e.dtoh(&row).unwrap_or_default()
1562}
1563
1564/// ROLLBACK DOOR for sampled BOUNDARY tokens (lane/sampled-spec-quality, 2026-08-19).
1565/// Default ON: the token a burst emits at its own boundary is drawn from the request's
1566/// sampler. `MEMRA_SPEC_SAMPLED_BOUNDARY=0` restores the pre-lane posture (an ARGMAX at
1567/// every boundary) without touching greedy, which is byte-unaffected either way.
1568pub fn spec_sampled_boundary_on() -> bool {
1569    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1570    *ON.get_or_init(|| std::env::var("MEMRA_SPEC_SAMPLED_BOUNDARY").as_deref() != Ok("0"))
1571}
1572
1573/// ROLLBACK DOOR for SESSION-SPANNING penalty history (lane/sampled-spec-quality).
1574/// Default ON: `pen_hist` is seeded from the session's committed tail, so repetition /
1575/// frequency / presence penalties see the whole stream. `MEMRA_SPEC_PEN_SESSION=0`
1576/// restores the pre-lane posture (each burst restarts the window from its own prompt
1577/// slice, i.e. from NOTHING on a continuation burst) — and with the door shut the worker
1578/// must keep refusing penalized sampled prefix-cache restores, because the restored
1579/// session's continuation burst is handed no prompt slice at all.
1580pub fn spec_pen_session_on() -> bool {
1581    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1582    *ON.get_or_init(|| std::env::var("MEMRA_SPEC_PEN_SESSION").as_deref() != Ok("0"))
1583}
1584
1585/// ROLLBACK DOOR for extended-entry publication from a RESTORED session
1586/// (lane/sampled-spec-quality, Item 3). Default ON: a converted prefix-cache hit that fed a
1587/// suffix captures its own prompt-end boundary so the NEXT turn can hit a longer prefix.
1588/// `MEMRA_SPEC_RESTORE_REPUBLISH=0` restores the pre-lane posture (a namespace learns exactly
1589/// one boundary and never advances it). Whole-entry semantics only — the boundary is the
1590/// restored session's own prompt end, so `entry_pos != fed_len` still refuses on the way in.
1591pub fn spec_restore_republish_on() -> bool {
1592    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1593    *ON.get_or_init(|| std::env::var("MEMRA_SPEC_RESTORE_REPUBLISH").as_deref() != Ok("0"))
1594}
1595
1596/// Diagnostics: name every boundary token on stderr (`MEMRA_SPEC_BOUNDARY_TRACE=1`), with
1597/// the argmax the pre-lane code would have emitted from the same row. This is how the
1598/// lane MEASURES the boundary rate and the deviation rate instead of estimating them.
1599fn spec_boundary_trace() -> bool {
1600    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1601    *ON.get_or_init(|| std::env::var("MEMRA_SPEC_BOUNDARY_TRACE").as_deref() == Ok("1"))
1602}
1603
1604/// llama-parity floor for the penalty window when the request does not ask for a bigger
1605/// one (`repeat_last_n` default). The serve API arms `penalty_last_n = PEN_WINDOW_MAX` for any
1606/// non-identity penalty, so this floor only matters to explicit small windows and to the
1607/// CLI env path.
1608const PEN_WINDOW_FLOOR: usize = 64;
1609
1610/// CEILING on the penalty window, and it is a COST bound, not a semantic preference.
1611/// `penalize_logits_f32` (cu/spec_sample.cu) dedups on device by having thread `i` scan
1612/// `hist[0..i]`, so a pass is O(n_hist²) and it runs ~3x per verify round (the q rows, the
1613/// p column, the bonus column). The serve API uses this same bound for every non-identity
1614/// penalty so host/plain, sparse-device, and speculative sampling cannot change logits on
1615/// admission demotion. An uncapped 128k-token history would put ~1.7e10
1616/// comparisons per pass, tens of ms per round, i.e. penalties would silently destroy decode
1617/// throughput on exactly the long-context requests that most want them. 8192 keeps a pass
1618/// at ~7e7 comparisons (tens of microseconds) while still being **128x wider than the
1619/// pre-lane effective window** (64 prompt-tail tokens + whatever the current burst had
1620/// generated). A request that genuinely needs a window beyond this wants host-side dedup +
1621/// counts through a new kernel signature — a follow-up lane, named here rather than hidden.
1622/// `pub` since lane/dspark-penalized-sampled-20260821: the dspark route's accept walk and
1623/// the dspark_sample_gate binary trim their uploads with the SAME cap — a second constant
1624/// is a second thing to drift.
1625pub const PEN_WINDOW_MAX: usize = 8192;
1626
1627/// Seed a penalty window over the SESSION, not the burst (lane/sampled-spec-quality,
1628/// Item 2). The window is the last `max(penalty_last_n, 64)` tokens of
1629/// `session_committed ++ burst_prompt` — for a cold turn-1 burst (`session_committed`
1630/// empty, default `penalty_last_n`) that is byte-identically the pre-lane
1631/// `prompt.iter().rev().take(64).rev()`; for a continuation burst it is the stream the
1632/// client actually asked us to penalize, where the pre-lane code had NOTHING.
1633/// `pub` since lane/dspark-penalized-sampled-20260821: the dspark route seeds its session
1634/// window through the SAME function (one definition of "the window" across both spec
1635/// routes and the gate binary's trunk-only reference arm).
1636pub fn pen_window_seed(
1637    session_committed: &[u32],
1638    burst_prompt: &[u32],
1639    penalty_last_n: usize,
1640) -> Vec<u32> {
1641    let win = penalty_last_n.clamp(PEN_WINDOW_FLOOR, PEN_WINDOW_MAX);
1642    let take_prompt = burst_prompt.len().min(win);
1643    let take_sess = (win - take_prompt).min(session_committed.len());
1644    let mut hist = Vec::with_capacity(take_sess + take_prompt);
1645    hist.extend_from_slice(&session_committed[session_committed.len() - take_sess..]);
1646    hist.extend_from_slice(&burst_prompt[burst_prompt.len() - take_prompt..]);
1647    hist
1648}
1649
1650/// Draw a BOUNDARY token from the target distribution the request asked for
1651/// (lane/sampled-spec-quality, Item 1) — the fix for "sampled spec emits an ARGMAX token at
1652/// every burst boundary".
1653///
1654/// WHY THIS EXISTS. A spec burst's first emitted token is not produced by the accept walk:
1655/// it comes off a logits row that already exists (the prime's last row on a cold burst; the
1656/// row after the last committed token on a continuation burst; the prefix-cache entry's
1657/// boundary row on a restored one). Pre-lane that token was `argmax` in BOTH sampling
1658/// regimes, so a sampled stream took a greedy token once per burst — measured, not
1659/// estimated, in research/spec-cache-20260818/SAMPLED-QUALITY.md. At temperature > 0 the
1660/// customer asked for a sampled token, so this draws one.
1661///
1662/// THE PROGRAM IS THE FULL-ACCEPT BONUS'S PROGRAM, deliberately: penalize the row (over the
1663/// session's window), take this row's OWN filter stats (the sampfix-20260805 law — stats
1664/// from a neighbour row mis-scale every `e0` and can wipe the row to token 0), gumbel-perturb
1665/// with the session's Philox stream at `*sctr`, argmax the perturbed row. Reusing the bonus's
1666/// composition means `sample_check`'s distributional oracle covers this draw too, and the
1667/// boundary token is drawn from the same filtered/penalized `p` the accept walk targets.
1668///
1669/// THE STREAM IS THE SESSION'S, NOT A FRESH ONE. `sctr` is the caller's live counter and is
1670/// advanced by exactly one, so a boundary draw consumes the next value in the same Philox
1671/// stream the accept walk uses — never a second, independently seeded stream (which would be
1672/// a new distributional bug: two streams from one seed correlate wherever their counters
1673/// collide). That also makes a restored session's boundary draw at `sctr == 0` bit-identical
1674/// to the cold session's own first draw from the same logits row, which is what preserves the
1675/// sampled-hit lane's per-seed hit==cold byte identity.
1676#[allow(clippy::too_many_arguments)]
1677pub fn sample_boundary_token_dev(
1678    e: &Engine,
1679    logits: &CudaSlice<f32>,
1680    n_vocab: usize,
1681    sp: &SpecSampling,
1682    pen_hist: &[u32],
1683    sctr: &mut u32,
1684    site: &str,
1685) -> Result<u32, Box<dyn std::error::Error>> {
1686    debug_assert!(
1687        sp.temp > 0.0,
1688        "boundary sampling is the sampled regime only"
1689    );
1690    // Own copy: penalize_logits mutates in place and the caller's row is live state
1691    // (prime_logits back the constrained recompute; last_col_logits backs round 0's accept).
1692    let mut col = e.zeros(n_vocab)?;
1693    e.copy_into(&mut col, 0, logits, n_vocab)?;
1694    let pen_on = sp.penalty_last_n > 0
1695        && (sp.penalty_repeat != 1.0 || sp.penalty_freq != 0.0 || sp.penalty_present != 0.0);
1696    if pen_on && !pen_hist.is_empty() {
1697        // window trim mirrors the round loop's own upload (`pen_hist[w0..]`), cap included.
1698        let w0 = pen_hist
1699            .len()
1700            .saturating_sub(sp.penalty_last_n.min(PEN_WINDOW_MAX));
1701        let hist = &pen_hist[w0..];
1702        let hd = e.htod_u32_v(hist)?;
1703        e.penalize_logits(
1704            &mut col,
1705            &hd,
1706            hist.len(),
1707            sp.penalty_repeat,
1708            sp.penalty_freq,
1709            sp.penalty_present,
1710            n_vocab,
1711        )?;
1712    }
1713    let rows0 = e.htod_i32(&[0])?;
1714    let (mut th_d, mut z_d, mut mx_d) = (e.zeros(1)?, e.zeros(1)?, e.zeros(1)?);
1715    e.filter_stats(
1716        &col, n_vocab, &rows0, &mut th_d, &mut z_d, &mut mx_d, n_vocab, 1, sp.temp, sp.top_k,
1717        sp.top_p, sp.min_p,
1718    )?;
1719    let (th, mx) = (e.dtoh(&th_d)?[0], e.dtoh(&mx_d)?[0]);
1720    let mut perturb = e.zeros(n_vocab)?;
1721    e.gumbel_perturb_filtered(&col, &mut perturb, n_vocab, sp.seed, *sctr, sp.temp, mx, th)?;
1722    *sctr = sctr.wrapping_add(1);
1723    let td = e.argmax_token_device(&perturb, n_vocab)?;
1724    let tok = guard_vocab_token(
1725        e.dtoh_u32_one(&td)?,
1726        n_vocab,
1727        &format!("sampled boundary token (site={site})"),
1728    )?;
1729    if spec_boundary_trace() {
1730        // the pre-lane token, from the SAME row, so the deviation rate is measurable.
1731        let raw = e.argmax_token_device(logits, n_vocab)?;
1732        let greedy = e.dtoh_u32_one(&raw)?;
1733        eprintln!(
1734            "[spec-boundary] site={site} sampled={tok} argmax={greedy} \
1735             deviates={} temp={} sctr={}",
1736            (tok != greedy) as u8,
1737            sp.temp,
1738            sctr.wrapping_sub(1),
1739        );
1740    }
1741    Ok(tok)
1742}
1743
1744/// Host-row twin of [`sample_boundary_token_dev`] (the prime / feed / entry rows arrive as
1745/// host `Vec<f32>`).
1746#[allow(clippy::too_many_arguments)]
1747pub fn sample_boundary_token(
1748    e: &Engine,
1749    logits: &[f32],
1750    sp: &SpecSampling,
1751    pen_hist: &[u32],
1752    sctr: &mut u32,
1753    site: &str,
1754) -> Result<u32, Box<dyn std::error::Error>> {
1755    let n_vocab = logits.len();
1756    let d = e.htod(logits)?;
1757    sample_boundary_token_dev(e, &d, n_vocab, sp, pen_hist, sctr, site)
1758}
1759
1760struct SpecPipeTraceClock {
1761    pair: usize,
1762    started: std::time::Instant,
1763}
1764
1765#[derive(Clone)]
1766struct SpecPipeTraceCtx {
1767    clock: std::sync::Arc<SpecPipeTraceClock>,
1768    round: usize,
1769    lane: usize,
1770}
1771
1772struct SpecPipeTraceMarker {
1773    trace: SpecPipeTraceCtx,
1774    phase: &'static str,
1775    edge: &'static str,
1776    slot: Option<usize>,
1777}
1778
1779unsafe extern "C" fn spec_pipe_trace_marker(raw: *mut std::ffi::c_void) {
1780    let marker = unsafe { Box::from_raw(raw.cast::<SpecPipeTraceMarker>()) };
1781    let lane = if marker.trace.lane == 0 { "A" } else { "B" };
1782    let slot = marker
1783        .slot
1784        .map(|v| v.to_string())
1785        .unwrap_or_else(|| "-".into());
1786    let t_ms = marker.trace.clock.started.elapsed().as_secs_f64() * 1e3;
1787    use std::io::Write as _;
1788    let stderr = std::io::stderr();
1789    let mut stderr = stderr.lock();
1790    let _ = writeln!(
1791        stderr,
1792        "[spec-pipe-timeline] pair={} round={} lane={lane} phase={} edge={} \
1793         slot={slot} t_ms={t_ms:.3}",
1794        marker.trace.clock.pair, marker.trace.round, marker.phase, marker.edge,
1795    );
1796}
1797
1798fn enqueue_spec_pipe_trace_marker(
1799    stream: &cudarc::driver::CudaStream,
1800    trace: Option<&SpecPipeTraceCtx>,
1801    phase: &'static str,
1802    edge: &'static str,
1803    slot: Option<usize>,
1804) -> Result<(), Box<dyn std::error::Error>> {
1805    let Some(trace) = trace else {
1806        return Ok(());
1807    };
1808    let marker = Box::new(SpecPipeTraceMarker {
1809        trace: trace.clone(),
1810        phase,
1811        edge,
1812        slot,
1813    });
1814    let raw = Box::into_raw(marker);
1815    let result = unsafe {
1816        cudarc::driver::result::stream::launch_host_function(
1817            stream.cu_stream(),
1818            spec_pipe_trace_marker,
1819            raw.cast(),
1820        )
1821    };
1822    if let Err(err) = result {
1823        unsafe {
1824            drop(Box::from_raw(raw));
1825        }
1826        return Err(err.into());
1827    }
1828    Ok(())
1829}
1830
1831#[derive(Default)]
1832struct SpecPipeProgress {
1833    setup_done: [bool; 2],
1834    draft_done: [usize; 2],
1835    stage0_done: [usize; 2],
1836    verify_done: [usize; 2],
1837    accept_done: [usize; 2],
1838    finished: [bool; 2],
1839    aborted: bool,
1840}
1841
1842/// Host-side issue coordinator for the reduced two-session speculative pipeline. Each session
1843/// keeps its existing call stack and round locals; this object only orders phase entry. The
1844/// primary mutex spans whole draft/accept/tail issue regions so Engine's single-stream scratch
1845/// cannot be interleaved by the two host threads.
1846struct SpecPipeSync {
1847    progress: std::sync::Mutex<SpecPipeProgress>,
1848    changed: std::sync::Condvar,
1849    primary: std::sync::Mutex<()>,
1850    trace: Option<std::sync::Arc<SpecPipeTraceClock>>,
1851}
1852
1853impl SpecPipeSync {
1854    fn new() -> Self {
1855        static TRACE_PAIR: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
1856        let trace = (std::env::var("MEMRA_SPEC_PIPE_TRACE").as_deref() == Ok("1")).then(|| {
1857            std::sync::Arc::new(SpecPipeTraceClock {
1858                pair: TRACE_PAIR.fetch_add(1, std::sync::atomic::Ordering::Relaxed) + 1,
1859                started: std::time::Instant::now(),
1860            })
1861        });
1862        Self {
1863            progress: std::sync::Mutex::new(SpecPipeProgress::default()),
1864            changed: std::sync::Condvar::new(),
1865            primary: std::sync::Mutex::new(()),
1866            trace,
1867        }
1868    }
1869}
1870
1871#[derive(Clone)]
1872struct SpecPipeLane {
1873    sync: std::sync::Arc<SpecPipeSync>,
1874    lane: usize,
1875}
1876
1877impl SpecPipeLane {
1878    fn peer(&self) -> usize {
1879        1 - self.lane
1880    }
1881
1882    fn aborted() -> Box<dyn std::error::Error> {
1883        "paired speculative peer aborted".into()
1884    }
1885
1886    fn trace(&self, round: usize) -> Option<SpecPipeTraceCtx> {
1887        self.sync.trace.as_ref().map(|clock| SpecPipeTraceCtx {
1888            clock: clock.clone(),
1889            round,
1890            lane: self.lane,
1891        })
1892    }
1893
1894    fn setup_begin(&self) -> Result<(), Box<dyn std::error::Error>> {
1895        let mut p = self.sync.progress.lock().unwrap();
1896        while !p.aborted && self.lane == 1 && !p.setup_done[0] && !p.finished[0] {
1897            p = self.sync.changed.wait(p).unwrap();
1898        }
1899        if p.aborted {
1900            Err(Self::aborted())
1901        } else {
1902            Ok(())
1903        }
1904    }
1905
1906    fn setup_end(&self) {
1907        let mut p = self.sync.progress.lock().unwrap();
1908        p.setup_done[self.lane] = true;
1909        self.sync.changed.notify_all();
1910    }
1911
1912    fn draft_begin(
1913        &self,
1914        round: usize,
1915    ) -> Result<std::sync::MutexGuard<'_, ()>, Box<dyn std::error::Error>> {
1916        let peer = self.peer();
1917        let mut p = self.sync.progress.lock().unwrap();
1918        loop {
1919            if p.aborted {
1920                return Err(Self::aborted());
1921            }
1922            let setup_ready =
1923                (p.setup_done[0] || p.finished[0]) && (p.setup_done[1] || p.finished[1]);
1924            let prior_ready = p.accept_done[self.lane] >= round
1925                && (p.accept_done[peer] >= round || p.finished[peer]);
1926            let turn_ready = if self.lane == 0 {
1927                true
1928            } else {
1929                p.draft_done[0] > round || p.finished[0]
1930            };
1931            if setup_ready && prior_ready && turn_ready {
1932                break;
1933            }
1934            p = self.sync.changed.wait(p).unwrap();
1935        }
1936        drop(p);
1937        Ok(self.sync.primary.lock().unwrap())
1938    }
1939
1940    fn draft_end(&self, round: usize) {
1941        let mut p = self.sync.progress.lock().unwrap();
1942        p.draft_done[self.lane] = round + 1;
1943        self.sync.changed.notify_all();
1944    }
1945
1946    /// Admit stage 0 and return whether this lane owns the interval's one reverse fence.
1947    /// Lane B releases as soon as lane A has issued its boundary TX, not after A's full body.
1948    fn stage0_begin(&self, round: usize) -> Result<bool, Box<dyn std::error::Error>> {
1949        let peer = self.peer();
1950        let mut p = self.sync.progress.lock().unwrap();
1951        loop {
1952            if p.aborted {
1953                return Err(Self::aborted());
1954            }
1955            let ready = if self.lane == 0 {
1956                p.draft_done[0] > round && (p.draft_done[1] > round || p.finished[1])
1957            } else {
1958                p.draft_done[1] > round && (p.stage0_done[0] > round || p.finished[0])
1959            };
1960            if ready {
1961                return Ok(self.lane == 0 || p.finished[peer]);
1962            }
1963            p = self.sync.changed.wait(p).unwrap();
1964        }
1965    }
1966
1967    fn stage0_end(&self, round: usize) {
1968        let mut p = self.sync.progress.lock().unwrap();
1969        p.stage0_done[self.lane] = round + 1;
1970        self.sync.changed.notify_all();
1971    }
1972
1973    /// Stage 1 is single-owner per engine. A proceeds immediately after its own ticket; B waits
1974    /// for A's full stage1/head issue so only A.S1 and B.S0 can overlap.
1975    fn stage1_begin(&self, round: usize) -> Result<(), Box<dyn std::error::Error>> {
1976        let mut p = self.sync.progress.lock().unwrap();
1977        while !p.aborted
1978            && !(p.stage0_done[self.lane] > round
1979                && (self.lane == 0 || p.verify_done[0] > round || p.finished[0]))
1980        {
1981            p = self.sync.changed.wait(p).unwrap();
1982        }
1983        if p.aborted {
1984            Err(Self::aborted())
1985        } else {
1986            Ok(())
1987        }
1988    }
1989
1990    fn verify_end(&self, round: usize) {
1991        let mut p = self.sync.progress.lock().unwrap();
1992        p.verify_done[self.lane] = round + 1;
1993        self.sync.changed.notify_all();
1994    }
1995
1996    fn accept_begin(
1997        &self,
1998        round: usize,
1999    ) -> Result<std::sync::MutexGuard<'_, ()>, Box<dyn std::error::Error>> {
2000        let mut p = self.sync.progress.lock().unwrap();
2001        loop {
2002            if p.aborted {
2003                return Err(Self::aborted());
2004            }
2005            let ready = if self.lane == 0 {
2006                p.verify_done[0] > round && (p.verify_done[1] > round || p.finished[1])
2007            } else {
2008                p.verify_done[1] > round && (p.accept_done[0] > round || p.finished[0])
2009            };
2010            if ready {
2011                break;
2012            }
2013            p = self.sync.changed.wait(p).unwrap();
2014        }
2015        drop(p);
2016        Ok(self.sync.primary.lock().unwrap())
2017    }
2018
2019    fn accept_end(&self, round: usize) {
2020        let mut p = self.sync.progress.lock().unwrap();
2021        p.accept_done[self.lane] = round + 1;
2022        self.sync.changed.notify_all();
2023    }
2024
2025    fn primary(&self) -> std::sync::MutexGuard<'_, ()> {
2026        self.sync.primary.lock().unwrap()
2027    }
2028
2029    fn finish(&self, failed: bool) {
2030        let mut p = self.sync.progress.lock().unwrap();
2031        p.finished[self.lane] = true;
2032        p.aborted |= failed;
2033        self.sync.changed.notify_all();
2034    }
2035}
2036
2037struct SpecPipeFinish<'a> {
2038    lane: &'a SpecPipeLane,
2039    closed: bool,
2040}
2041
2042impl<'a> SpecPipeFinish<'a> {
2043    fn new(lane: &'a SpecPipeLane) -> Self {
2044        Self {
2045            lane,
2046            closed: false,
2047        }
2048    }
2049
2050    fn close(&mut self, failed: bool) {
2051        self.lane.finish(failed);
2052        self.closed = true;
2053    }
2054}
2055
2056impl Drop for SpecPipeFinish<'_> {
2057    fn drop(&mut self) {
2058        if !self.closed {
2059            self.lane.finish(true);
2060        }
2061    }
2062}
2063
2064/// Scoped transfer of one exclusively-borrowed session to the second host issue thread.
2065/// `CudaGraph` is not marked Send by cudarc because its raw driver handles carry no automatic
2066/// trait. CUDA driver graph handles are context-scoped rather than OS-thread-affine; the caller
2067/// binds that context before touching the session, joins before returning, and never aliases the
2068/// pointer. Keep this exception local to the experimental pair call instead of marking the public
2069/// session type Send.
2070struct SpecPipeSessionPtr(*mut SpecSession);
2071
2072unsafe impl Send for SpecPipeSessionPtr {}
2073
2074impl SpecPipeSessionPtr {
2075    unsafe fn get_mut(&mut self) -> &mut SpecSession {
2076        unsafe { &mut *self.0 }
2077    }
2078}
2079
2080/// Per-session persistent draft-graph context: the captured CUDA graph(s) plus the device
2081/// buffers whose POINTERS the capture bakes. Reuse legality: the greedy capture bakes only
2082/// session-stable pointers (the session's own MtpScratch KV — allocated once, never realloc'd;
2083/// the model's resident embedding; the process-wide OnceLock p_min) and the g_* buffers held
2084/// HERE — so one capture serves the session's whole lifetime. The sampled capture additionally
2085/// bakes (seed, temp) as capture-time constants and needs k q-slots — keyed by `s_key`, dropped
2086/// and recaptured when a pool-resumed request changes them. `*_failed` memoizes a failed capture
2087/// so the eager fallback doesn't pay a doomed capture attempt every burst.
2088/// Capture identity of the parked SAMPLED draft graph (`DraftGraphCtx::graph_s`).
2089///
2090/// EXACTNESS, not perf (lane/graph-s-key-exactness-20260819; receipts
2091/// `research/spec-cache-20260818/GRAPH-S-KEY.md`). Two classes of field live here, both
2092/// load-bearing:
2093///
2094/// - **Baked constants.** `seed` and `temp` are capture-time constants INSIDE the graph and `k`
2095///   sizes the q slots its replays write. A resumed request changing any of them must recapture.
2096///   This is all the key used to carry.
2097/// - **Regime fields.** `top_k`/`top_p`/`min_p`/`pen_on` are not baked, but they decide whether
2098///   the captured graph is a legal draft chain AT ALL. The in-graph draw is one gumbel-max over
2099///   the RAW softmax (`gumbel_perturb_ctr`, unfiltered by construction), while the verify builds
2100///   the accept test's `q` from `filter_stats(q_slots, top_k, top_p, min_p)`. If those disagree
2101///   the accept test evaluates a distribution the draft was never sampled from: a draft token
2102///   below the filter threshold gathers `q = 0` (`softmax_gather_filtered_f32`,
2103///   `cu/spec_sample.cu`) and `u * 0 < p` accepts it UNCONDITIONALLY.
2104///
2105/// Omitting the regime fields was reachable — not through the prefix-cache spec restore (that
2106/// path is greedy-only, `memra-server` `spec_restore_convertible`), but through WHOLE-SESSION
2107/// spec reuse: a parked `SpecSession` carries this `DraftGraphCtx`, and the pool-resume probe
2108/// applies no sampler predicate at all. Turn 1 pure-temp parks a graph; turn 2 of the same
2109/// conversation, same explicit seed and temperature, adds `top_p`/`top_k` and inherits it.
2110#[derive(Clone, Copy, PartialEq, Eq, Debug)]
2111pub(crate) struct SampledGraphKey {
2112    seed: u64,
2113    temp_bits: u32,
2114    k: usize,
2115    top_k: i32,
2116    top_p_bits: u32,
2117    min_p_bits: u32,
2118    pen_on: bool,
2119}
2120
2121impl SampledGraphKey {
2122    pub(crate) fn new(
2123        seed: u64,
2124        temp: f32,
2125        k: usize,
2126        top_k: i32,
2127        top_p: f32,
2128        min_p: f32,
2129        pen_on: bool,
2130    ) -> Self {
2131        SampledGraphKey {
2132            seed,
2133            temp_bits: temp.to_bits(),
2134            k,
2135            top_k,
2136            top_p_bits: top_p.to_bits(),
2137            min_p_bits: min_p.to_bits(),
2138            pen_on,
2139        }
2140    }
2141
2142    /// The one regime the PURE-TEMP in-graph sampled chain may stand in for the eager one:
2143    /// nothing but temperature shapes `q`. Computed FROM THE KEY so the capture guard, the
2144    /// launch guard and the key can never drift apart (they were three separate expressions
2145    /// before this lane, and the launch site simply forgot to ask).
2146    pub(crate) fn pure_temp(&self) -> bool {
2147        self.top_k == 0
2148            && f32::from_bits(self.top_p_bits) >= 1.0
2149            && f32::from_bits(self.min_p_bits) <= 0.0
2150            && !self.pen_on
2151    }
2152
2153    /// Truncation filters active — the capture body needs the IN-GRAPH filter nodes
2154    /// (`filter_stats` + `gumbel_perturb_filtered_ctr`) so the draft draws from the same
2155    /// filtered distribution the accept test reconstructs. Meaningful only when
2156    /// `graph_capturable`; penalties never reach a capture body.
2157    pub(crate) fn filtered(&self) -> bool {
2158        !self.pure_temp()
2159    }
2160
2161    /// May the sampled draft graph be CAPTURED (and a parked one LAUNCHED) for this regime?
2162    /// Pure-temp always; filtered regimes when the filtered-capture door is on
2163    /// (lane/step37-draft-graph-serving-20260830); penalties never — the per-round history
2164    /// cannot be baked into a graph, and composing a raw-softmax (or stale-history) draw
2165    /// with a penalized accept test is the unconditional-accept exactness bug. Computed FROM
2166    /// THE KEY for the same no-drift reason as `pure_temp`.
2167    pub(crate) fn graph_capturable(&self) -> bool {
2168        !self.pen_on && (self.pure_temp() || spec_graph_filtered_on())
2169    }
2170}
2171
2172/// Per-head captured graphs for the MULTI-HEAD MTP draft chain (step-modulo prefix-replay,
2173/// lane/step37-draft-graph-serving-20260830). The chain POLICY — which head serves step j,
2174/// how long the replayed prefix is, which stored seed feeds row r — stays HOST-SIDE in the
2175/// launch loop, exactly `mtp_chain_forward_dev`'s order; the graphs capture ONE head-row
2176/// forward each, on the head's OWN scratch plane:
2177/// - `interior[i]`: head i, `with_head=false` — KV append + carrier only. Interior rows'
2178///   logits are dead in the eager chain too (`mtp_chain_forward_dev` keeps only the last
2179///   row), so skipping the head matmul changes no consumed byte and removes the eager
2180///   chain's per-replay-row full-vocab matmul.
2181/// - `last[i]`: head i, `with_head=true` + the mode's tail (greedy argmax, or the sampled
2182///   gumbel draw — filtered in-graph when the request carries filters).
2183/// One `DraftChainGraphs` per MODE (greedy vs sampled), owning its keeper: dropping the
2184/// sampled chain on an s_key change never invalidates the greedy one.
2185struct DraftChainGraphs {
2186    interior: Vec<cudarc::driver::CudaGraph>,
2187    last: Vec<cudarc::driver::CudaGraph>,
2188    keeper: Vec<Box<dyn std::any::Any + Send>>,
2189}
2190
2191/// Sampled-tail capture pack for `mtp_head_forward_cap`: the persistent buffers and baked
2192/// constants of the in-graph categorical draw. `filt: None` = the PURE-TEMP body (gumbel
2193/// over the raw softmax), byte-identical to the pre-lane capture; `Some` adds the in-graph
2194/// truncation filter (`filter_stats` + `gumbel_perturb_filtered_ctr`) so the draft draws
2195/// from the same filtered distribution the accept test reconstructs
2196/// (lane/step37-draft-graph-serving-20260830).
2197struct SampledCapArgs<'a> {
2198    ctr: &'a mut CudaSlice<u32>,
2199    perturb: &'a mut CudaSlice<f32>,
2200    q_out: &'a mut CudaSlice<f32>,
2201    seed: u64,
2202    temp: f32,
2203    filt: Option<SampledCapFilter<'a>>,
2204}
2205
2206/// In-graph truncation-filter nodes: the stat slots `filter_stats` fills and the perturb
2207/// reads, plus the filter constants baked into the capture (they live in `s_key`, so a
2208/// request whose filters differ drops the parked graph before this ever goes stale).
2209struct SampledCapFilter<'a> {
2210    rows0: &'a CudaSlice<i32>,
2211    th: &'a mut CudaSlice<f32>,
2212    z: &'a mut CudaSlice<f32>,
2213    mx: &'a mut CudaSlice<f32>,
2214    top_k: i32,
2215    top_p: f32,
2216    min_p: f32,
2217}
2218
2219pub(crate) struct DraftGraphCtx {
2220    g_tok: CudaSlice<u32>,
2221    g_pos: CudaSlice<i32>,
2222    g_seed: CudaSlice<f32>,
2223    g_p: CudaSlice<f32>,
2224    g_ctr: CudaSlice<u32>,
2225    g_q: CudaSlice<f32>,
2226    g_perturb: CudaSlice<f32>,
2227    /// IN-GRAPH filter-stat slots (filtered sampled capture): `filter_stats` writes
2228    /// (th, z, mx) here inside the graph; `gumbel_perturb_filtered_ctr` reads (mx, th) from
2229    /// the same slots. Persistent so the baked pointers survive replays. `g_rows0` is the
2230    /// constant row-index-0 the single-row `filter_stats` launch reads (a captured memcpy
2231    /// source must not be a host temporary).
2232    g_rows0: CudaSlice<i32>,
2233    g_th: CudaSlice<f32>,
2234    g_z: CudaSlice<f32>,
2235    g_mx: CudaSlice<f32>,
2236    q_slots: Vec<CudaSlice<f32>>,
2237    /// DRAFT-SIDE GRAMMAR MASK (lane/draft-mask): packed allowed-set words over the DRAFT
2238    /// head's vocab, at a STABLE address so the captured draft graph's mask node reads the
2239    /// per-position contents the host re-uploads before each replay (the graph-promote
2240    /// pattern from decode.rs). Empty unless the session drafts under a grammar.
2241    g_dmask: CudaSlice<u32>,
2242    /// was `graph` captured WITH the mask node? A parked graph of the wrong shape is dropped.
2243    /// Covers the multi-head `chain` too (single-head and chain are mutually exclusive for a
2244    /// given model, so one flag serves whichever is active).
2245    graph_masked: bool,
2246    graph: Option<cudarc::driver::CudaGraph>,
2247    graph_s: Option<cudarc::driver::CudaGraph>,
2248    /// Multi-head chain graphs (see [`DraftChainGraphs`]): greedy and sampled chains, the
2249    /// chain twins of `graph` / `graph_s`. `chain_s`'s capture identity is `s_key` (shared
2250    /// with `graph_s` — a session is either single-head or chain, never both), and it obeys
2251    /// the same drop rules (key mismatch, penalty regime, mask-shape change).
2252    chain: Option<DraftChainGraphs>,
2253    chain_s: Option<DraftChainGraphs>,
2254    /// Failed-capture memoization for both graphs — LOUD on flip, cleared on pool resume
2255    /// (audit Q2, the TRT #16072 silent-permanent-coverage-loss class).
2256    failed: DraftGraphFallback,
2257    /// Capture identity of `graph_s` — see [`SampledGraphKey`]. `None` iff no sampled graph is
2258    /// parked; a request whose key differs drops the parked graph (and its q slots/keeper).
2259    s_key: Option<SampledGraphKey>,
2260    /// CAPTURE-RETAIN keepers (#68 root cause, 2026-08-04): the warmup-run transients whose
2261    /// pool addresses the captured graph(s) bake. Without these, the transients return to the
2262    /// pool at capture-body exit and later work (burst-boundary prime/fill/commit passes, or a
2263    /// co-served session in the worker) reuses those addresses — the persisted graph's replay
2264    /// then reads/writes live unrelated buffers (exactness corruption, first seen as the ST
2265    /// serve-spec 4B graph-arm corruption; one-shot CLI calls never re-shuffled the pool, which
2266    /// is why run-spec K=1..8 passed on the same checkpoint). Same fix class as
2267    /// capture_graph_retained's gemma/decode.rs sites — hold as long as the graph replays.
2268    keeper: Vec<Box<dyn std::any::Any + Send>>,
2269    keeper_s: Vec<Box<dyn std::any::Any + Send>>,
2270}
2271
2272/// Failed-capture memoization for the two draft graphs (audit Q2, 2026-08-05 — the
2273/// TRT #16072 trap class: pressure-triggered, silent, long-lived coverage loss).
2274///
2275/// Three contracts:
2276/// - LOUD FLIP: `mark_*` returns the warn line exactly on the false→true transition
2277///   (returned, not printed, so the once-per-flip contract is unit-testable); the caller
2278///   `eprintln!`s it UNCONDITIONALLY — a dropped draft graph is never silent. Re-marking
2279///   an already-failed graph returns None (the per-burst memoization that keeps the eager
2280///   fallback from paying a doomed capture attempt every burst).
2281/// - RESET ON RESUME: `reset_on_resume` clears both flags — a parked session resumed by a
2282///   NEW request gets one fresh capture chance instead of carrying a transient-pressure
2283///   failure for the pool's whole lifetime. Returns the note line only when a flag was
2284///   actually set (quiet on the common clean-resume path).
2285/// - Shape-change clears (`clear_*`) stay silent, exactly as before: they precede a fresh
2286///   capture attempt whose own failure would re-flip loudly.
2287#[derive(Default)]
2288pub(crate) struct DraftGraphFallback {
2289    greedy: bool,
2290    sampled: bool,
2291}
2292impl DraftGraphFallback {
2293    fn mark_greedy(&mut self, reason: &str) -> Option<String> {
2294        if self.greedy {
2295            return None;
2296        }
2297        self.greedy = true;
2298        Some(format!(
2299            "[spec] WARN: draft-graph capture failed ({reason}); eager fallback until session resume"
2300        ))
2301    }
2302    fn mark_sampled(&mut self, reason: &str) -> Option<String> {
2303        if self.sampled {
2304            return None;
2305        }
2306        self.sampled = true;
2307        Some(format!(
2308            "[spec] WARN: sampled draft-graph capture failed ({reason}); eager fallback until session resume"
2309        ))
2310    }
2311    fn greedy_failed(&self) -> bool {
2312        self.greedy
2313    }
2314    fn sampled_failed(&self) -> bool {
2315        self.sampled
2316    }
2317    fn clear_greedy(&mut self) {
2318        self.greedy = false;
2319    }
2320    fn clear_sampled(&mut self) {
2321        self.sampled = false;
2322    }
2323    /// Pool-resume reset: both graphs get a fresh capture chance. Some(note) iff any flag
2324    /// was set (so clean resumes stay quiet).
2325    pub(crate) fn reset_on_resume(&mut self) -> Option<String> {
2326        if !self.greedy && !self.sampled {
2327            return None;
2328        }
2329        let which = match (self.greedy, self.sampled) {
2330            (true, true) => "greedy+sampled",
2331            (true, false) => "greedy",
2332            _ => "sampled",
2333        };
2334        self.greedy = false;
2335        self.sampled = false;
2336        Some(format!(
2337            "[spec] draft-graph fallback reset on session resume ({which}); recapture eligible"
2338        ))
2339    }
2340}
2341
2342impl DraftGraphCtx {
2343    fn new(e: &Engine, n_embd: usize, qlen: usize) -> Result<Self, Box<dyn std::error::Error>> {
2344        Ok(DraftGraphCtx {
2345            g_tok: e.alloc_u32_zeroed(1)?,
2346            g_pos: e.htod_i32(&[0])?,
2347            g_seed: e.zeros(n_embd)?,
2348            g_p: e.zeros(1)?,
2349            g_ctr: e.alloc_u32_zeroed(1)?,
2350            g_q: e.zeros(qlen)?,
2351            g_perturb: e.zeros(qlen)?,
2352            g_rows0: e.htod_i32(&[0])?,
2353            g_th: e.zeros(1)?,
2354            g_z: e.zeros(1)?,
2355            g_mx: e.zeros(1)?,
2356            q_slots: Vec::new(),
2357            g_dmask: e.alloc_u32_zeroed(1)?,
2358            graph_masked: false,
2359            graph: None,
2360            graph_s: None,
2361            chain: None,
2362            chain_s: None,
2363            failed: DraftGraphFallback::default(),
2364            s_key: None,
2365            keeper: Vec::new(),
2366            keeper_s: Vec::new(),
2367        })
2368    }
2369}
2370
2371pub(crate) struct MtpScratch {
2372    kv: KvLayer,
2373    /// Logical row capacity. On the graph/DC draft path it also doubles as fa_decode_dc's
2374    /// bucket_max: n_splits is sized from it ONCE, so the graph captured at round 0 stays valid
2375    /// for every later t_kv. Step35 refuses that path and may back this logical extent with the
2376    /// smaller host-indexed SWA ring instead.
2377    cap: usize,
2378    extra: Vec<MtpScratchPlane>,
2379}
2380
2381struct MtpScratchPlane {
2382    kv: KvLayer,
2383    cap: usize,
2384}
2385
2386fn mtp_scratch_layout(
2387    cfg: &memra_gguf::config::ModelConfig,
2388    geom: Option<&crate::hybrid::DraftGeom>,
2389) -> (usize, usize, usize, usize) {
2390    // Student draft heads carry fewer KV heads (head_dim unchanged) -> smaller scratch rows.
2391    let n_head_kv = geom.map(|g| g.n_head_kv).unwrap_or(cfg.n_head_kv as usize);
2392    let head_dim_k = cfg.head_dim_k as usize;
2393    let head_dim_v = cfg.head_dim_v as usize;
2394    assert!(
2395        head_dim_k % 32 == 0 && head_dim_v % 32 == 0,
2396        "KVQUANT requires head_dim%32==0 (MTP scratch)"
2397    );
2398    let kv_dim_k = head_dim_k * n_head_kv;
2399    let kv_dim_v = head_dim_v * n_head_kv;
2400    // The fp8-KV arm deliberately does not reach the draft scratch; keep the exact format
2401    // policy shared with `MtpScratch::new` so admission scales the same allocation.
2402    let (kbb, vbb) = crate::kv_blk_bytes();
2403    let k_tok_bytes = (kv_dim_k / 32) * kbb;
2404    let v_tok_bytes = (kv_dim_v / 32) * vbb;
2405    (kv_dim_k, kv_dim_v, k_tok_bytes, v_tok_bytes)
2406}
2407
2408fn mtp_chain_head_index(step: usize, head_count: usize) -> usize {
2409    assert!(head_count > 0, "MTP chain requires at least one head");
2410    step % head_count
2411}
2412
2413impl MtpScratch {
2414    fn alloc_plane(
2415        e: &Engine,
2416        cfg: &memra_gguf::config::ModelConfig,
2417        plan: &memra_gguf::model_plan::ModelPlan,
2418        cap: usize,
2419        geom: Option<&crate::hybrid::DraftGeom>,
2420    ) -> Result<MtpScratchPlane, Box<dyn std::error::Error>> {
2421        let (kv_dim_k, kv_dim_v, k_tok_bytes, v_tok_bytes) = mtp_scratch_layout(cfg, geom);
2422        let ring = if crate::cache::swa_ring_on()
2423            && crate::plan_backend::decode_batch_program(plan)
2424                == crate::plan_backend::DecodeBatchProgram::SlidingGatedMoe
2425        {
2426            let window = plan
2427                .layers
2428                .iter()
2429                .find_map(|layer| match layer.attention {
2430                    memra_gguf::model_plan::AttentionPlan::SlidingWindow { window, .. } => {
2431                        Some(window as usize)
2432                    }
2433                    _ => None,
2434                })
2435                .ok_or("sliding-gated-MoE draft scratch has no sliding-window layer")?;
2436            Some(crate::cache::KvRing::new(
2437                crate::cache::swa_ring_rows(window, cap),
2438                window,
2439            ))
2440        } else {
2441            None
2442        };
2443        let alloc_rows = ring.as_ref().map(crate::cache::KvRing::rows).unwrap_or(cap);
2444        // Ring-backed planes arm the device base mirror for the dcw draft arm (see
2445        // KvLayer::base_d): the captured chain derives its physical rows from
2446        // (len_d, base_d, window) with zero per-token node updates.
2447        let base_d = match ring.as_ref() {
2448            Some(_) => Some(e.htod_i32(&[0])?),
2449            None => None,
2450        };
2451        Ok(MtpScratchPlane {
2452            kv: KvLayer {
2453                k: e.alloc_u8(alloc_rows * k_tok_bytes)?,
2454                v: e.alloc_u8(alloc_rows * v_tok_bytes)?,
2455                kv_dim_k,
2456                kv_dim_v,
2457                k_tok_bytes,
2458                v_tok_bytes,
2459                len: 0,
2460                ring,
2461                len_d: e.htod_i32(&[0])?,
2462                base_d,
2463            },
2464            cap,
2465        })
2466    }
2467
2468    fn new(
2469        e: &Engine,
2470        cfg: &memra_gguf::config::ModelConfig,
2471        plan: &memra_gguf::model_plan::ModelPlan,
2472        cap: usize,
2473        geom: Option<&crate::hybrid::DraftGeom>,
2474    ) -> Result<Self, Box<dyn std::error::Error>> {
2475        // env-selected KV formats (default 34/24). The fp8-KV arm (MEMRA_KV_FP8) deliberately
2476        // does NOT reach the draft scratch: fp8 drafts drifted acceptance 69-88% -> 46%
2477        // (2026-07-12 A/B); the scratch is tiny, so it keeps baseline q8_0/q5_1 numerics
2478        // while the TRUNK cache carries the fp8 depth win. Scratch append/fa pass g=false.
2479        let primary = Self::alloc_plane(e, cfg, plan, cap, geom)?;
2480        Ok(MtpScratch {
2481            kv: primary.kv,
2482            cap: primary.cap,
2483            extra: Vec::new(),
2484        })
2485    }
2486
2487    fn push_plane(
2488        &mut self,
2489        e: &Engine,
2490        cfg: &memra_gguf::config::ModelConfig,
2491        plan: &memra_gguf::model_plan::ModelPlan,
2492        geom: Option<&crate::hybrid::DraftGeom>,
2493    ) -> Result<(), Box<dyn std::error::Error>> {
2494        self.extra
2495            .push(Self::alloc_plane(e, cfg, plan, self.cap, geom)?);
2496        Ok(())
2497    }
2498
2499    fn plane_count(&self) -> usize {
2500        1 + self.extra.len()
2501    }
2502
2503    fn plane(&self, index: usize) -> (&KvLayer, usize) {
2504        if index == 0 {
2505            (&self.kv, self.cap)
2506        } else {
2507            let plane = &self.extra[index - 1];
2508            (&plane.kv, plane.cap)
2509        }
2510    }
2511
2512    fn plane_mut(&mut self, index: usize) -> (&mut KvLayer, usize) {
2513        if index == 0 {
2514            (&mut self.kv, self.cap)
2515        } else {
2516            let plane = &mut self.extra[index - 1];
2517            (&mut plane.kv, plane.cap)
2518        }
2519    }
2520
2521    // #[track_caller]: set_len/set_plane_len have eight call sites (checkpoint restore, spec
2522    // rollback, session grow, seed replay ...) and the lap failure needs to say WHICH one, not
2523    // just that a rewind was refused.
2524    #[track_caller]
2525    fn set_plane_len(
2526        &mut self,
2527        e: &Engine,
2528        index: usize,
2529        n: usize,
2530    ) -> Result<(), Box<dyn std::error::Error>> {
2531        let caller = std::panic::Location::caller();
2532        let (kv, cap) = self.plane_mut(index);
2533        if let Some(ring) = kv.ring.as_ref() {
2534            if !ring.can_rewind_to(n) {
2535                // NAME THE NUMBERS (2026-08-28). This error is a step37 serving blocker on the
2536                // vendor-default shape and it fires from more than one call path with more than
2537                // one trigger: a long generation walks the checkpoint out of the ring, but a
2538                // ~4.5k-token prompt also fails within 5 s of prime, which accumulation cannot
2539                // explain. A bare message forced two rounds of guessing; the operands make each
2540                // trigger name itself.
2541                let raw = n.saturating_sub(ring.window().saturating_sub(1));
2542                return Err(format!(
2543                    "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})",
2544                    ring.window(),
2545                    ring.base(),
2546                    ring.rows(),
2547                    raw & !31usize,
2548                )
2549                .into());
2550            }
2551        }
2552        kv.len = n;
2553        e.set_i32_one(&mut kv.len_d, n as i32)
2554    }
2555
2556    /// Set BOTH length counters: the host mirror AND the device len_d the captured append/fa read
2557    /// (a 4-byte in-place htod — the counter pointer is baked into the graph, never realloc'd).
2558    /// This is the ONLY truncation/rollback mechanism the persistent draft KV needs.
2559    #[track_caller]
2560    fn set_len(&mut self, e: &Engine, n: usize) -> Result<(), Box<dyn std::error::Error>> {
2561        let caller = std::panic::Location::caller();
2562        if !self.can_rewind_to(n) {
2563            // set_plane_len re-checks and reports the operands; call it so the failure carries
2564            // which plane refused and why, instead of this bare aggregate.
2565            for index in 0..self.plane_count() {
2566                self.set_plane_len(e, index, n)?;
2567            }
2568            return Err(format!(
2569                "SWA ring MTP checkpoint has been lapped; full re-prime required (aggregate rewind_to={n}, no single plane reported, called from {caller})"
2570            )
2571            .into());
2572        }
2573        for index in 0..self.plane_count() {
2574            self.set_plane_len(e, index, n)?;
2575        }
2576        Ok(())
2577    }
2578
2579    fn can_rewind_to(&self, n: usize) -> bool {
2580        (0..self.plane_count()).all(|index| {
2581            self.plane(index)
2582                .0
2583                .ring
2584                .as_ref()
2585                .is_none_or(|ring| ring.can_rewind_to(n))
2586        })
2587    }
2588
2589    /// Pre-arm ring headroom for `rows` upcoming DEVICE-COUNTER appends (the dcw draft arm):
2590    /// a captured chain cannot rebase mid-replay, so any rebase the coming appends could need
2591    /// happens HERE, host-side, before the capture warmups or the round's replays (the rebase
2592    /// arm of `prepare_kv_append` also refreshes the plane's `base_d` device mirror). No-op on
2593    /// flat planes and when the ring already has room; `len` is untouched either way.
2594    fn ensure_dcw_headroom(
2595        &mut self,
2596        e: &Engine,
2597        rows: usize,
2598    ) -> Result<(), Box<dyn std::error::Error>> {
2599        for index in 0..self.plane_count() {
2600            let (kv, _) = self.plane_mut(index);
2601            let Some(ring) = kv.ring.as_ref() else {
2602                continue;
2603            };
2604            let retain = memra_kv::swa_retain_from(kv.len, ring.window(), ring.base());
2605            e.prepare_kv_append(kv, retain, rows)?;
2606        }
2607        Ok(())
2608    }
2609}
2610
2611/// Retained verify intermediates for the REPLAY-FREE partial accept (2026-07-03, the profiled
2612/// #1 spec cost at long ctx: the partial-accept replay was a DUPLICATE trunk pass — ~0.54 extra
2613/// full weight reads per round — recomputing columns the verify had already produced
2614/// bit-identically). Holds, per linear layer, everything needed to rebuild its recurrent state
2615/// to "after the first j verify columns" WITHOUT re-running the trunk:
2616/// - BATCHED-path layers (`gdn`): the exact token-major inputs the round's ONE gdn_scan
2617///   consumed. A prefix re-run of the SAME kernel (t=j) from the snapshot state is bit-identical
2618///   to the first j iterations of the verify's scan — the kernel's t-loop carries state in
2619///   registers and iteration t never depends on T. `qkv_mixed` (the conv input) feeds the
2620///   pure-copy ring rebuild.
2621/// - PER-COLUMN-path layers (`cols`): dtod clones of (conv_state, ssm_state) taken after each
2622///   column 0..t-2 — pure copies of the actual chain states (the last column is never a rebuild
2623///   target: j <= t-1).
2624/// Full-attn layers need nothing: their verify KV rows are bit-identical to eager's (the
2625/// decode-exact contract; verify-probe pins it), so rollback = len truncation.
2626struct GdnStash {
2627    qkv_mixed: CudaSlice<f32>, // [t, conv_dim] token-major (conv input)
2628    q_l2: CudaSlice<f32>,
2629    k_l2: CudaSlice<f32>,
2630    v_g: CudaSlice<f32>, // [t, num_v, d_state]
2631    g_log: CudaSlice<f32>,
2632    beta: CudaSlice<f32>, // [t, num_v]
2633}
2634pub(crate) struct VerifyCkpt {
2635    gdn: Vec<Option<GdnStash>>, // [n_layer], Some iff batched linear path ran
2636    cols: Vec<Option<Vec<(CudaSlice<f32>, CudaSlice<f32>)>>>, // [n_layer][col] = (conv, ssm) after col
2637}
2638/// Opaque handle for the dspark round (dflash.rs) — VerifyCkpt stays spec-private.
2639pub(crate) struct DsparkVerifyCkpt(VerifyCkpt);
2640
2641/// Engine-bundle slice 3 (DSF-ROUNDCOST-20260820 §2 row 4 / §5 rank 1): bucketed CUDA
2642/// graphs for the dspark verify's LINEAR-layer segments. The measured verify is ~2,800
2643/// eager launches whose residual cost is DEVICE-side per-launch overhead (slice 2 proved
2644/// host dispatch is not the binder: fully-deferred dispatch bought ~0 wall). The 48 GDN
2645/// layers between full-attention layers are shape-static given vt — no positions, no
2646/// t_kv, state addressed through pointer tables — so runs of them capture per
2647/// (segment, vt) and replay as ONE graph launch each. Full-attention layers stay eager
2648/// (their per-row append/fa arm picks are t_kv-driven — the exec-update extension).
2649///
2650/// Per round out-of-graph: one pointer-table refresh (gdn ping-pong moves the canonical
2651/// handles), one input-staging copy per segment, host parity bookkeeping. Captured via
2652/// `capture_graph_retained` (2 warmups + capture, keeper retains warmup transients so
2653/// pool addresses stay stable); the warmups EXECUTE, so segment conv/ssm state is saved
2654/// before and restored after — the graph's first real launch starts from the exact
2655/// pre-round state. The ckpt column stash rides persistent slabs (written inside the
2656/// graph as memcpy nodes); commit reads them via `dspark_commit_prefix_slab`.
2657/// `MEMRA_DSPARK_VERIFY_GRAPH=0` reverts to the eager walk (byte-identical body).
2658pub(crate) struct DsparkVerifyGraphs {
2659    /// Linear-attention layer indices ascending; `lin_pos[il]` = index into the vecs.
2660    lin: Vec<usize>,
2661    lin_pos: std::collections::HashMap<usize, usize>,
2662    /// [n_lin x 6] pointer table (conv, s0, s1, conv, s1, s0 per layer), refreshed per
2663    /// verify from the live handles; layer il's slice starts at lin_pos[il]*6.
2664    table_all: CudaSlice<u64>,
2665    host_table: Vec<u64>,
2666    /// Persistent per-layer ckpt stash slabs: row r of the verify at slab offset
2667    /// r*words. Shared by every (segment, vt) bucket — one verify runs at a time.
2668    stash_conv: Vec<CudaSlice<f32>>,
2669    stash_ssm: Vec<CudaSlice<f32>>,
2670    conv_words: usize,
2671    ssm_words: usize,
2672    /// Per-vt input/output staging (stable addresses the graphs bake).
2673    stage: std::collections::HashMap<usize, (CudaSlice<f32>, CudaSlice<f32>)>,
2674    /// Per-vt dflash tap-sink buffers — the captured segments bake the tap dst address,
2675    /// so the sink buffer must live (and persist) with the graphs, not with the round.
2676    pub(crate) tap_bufs: std::collections::HashMap<usize, CudaSlice<f32>>,
2677    graphs: std::collections::HashMap<(usize, usize), DsparkSegGraph>,
2678    /// Warmup-corruption guard scratch: pre-capture conv/ssm of every linear layer
2679    /// (sized n_lin — the slice-4c full-verify warmups execute the whole walk).
2680    save_conv: CudaSlice<f32>,
2681    save_ssm: CudaSlice<f32>,
2682    max_run: usize,
2683    n_embd: usize,
2684    /// Set by the verify walk: this round's linear ckpt lives in the slabs (the caller
2685    /// commits through `dspark_commit_prefix_slab` instead of the cols arm).
2686    pub(crate) round_slab: bool,
2687    // ---- slice 4c: full-verify single graph per (vt, rung) ----
2688    /// Full-attention layer indices ascending; `fa_pos[il]` = index into the vec.
2689    fa: Vec<usize>,
2690    fa_pos: std::collections::HashMap<usize, usize>,
2691    /// [n_fa x 2 x t_cap] interleaved (k,v) base-pointer pairs, refreshed per verify;
2692    /// layer il's slice starts at `fa_pos[il] * 2 * t_cap` (the seqs twins read pairs
2693    /// [2z], z < t <= t_cap, so one t_cap-sized table serves every vt).
2694    fa_table: CudaSlice<u64>,
2695    fa_host_table: Vec<u64>,
2696    t_cap: usize,
2697    /// Per-vt position staging for the captured bodies — contents refreshed per round
2698    /// (rope reads row r; the seqs twins derive append slot and T_kv per z from it).
2699    pos_stage: std::collections::HashMap<usize, CudaSlice<i32>>,
2700    /// Full-verify graphs keyed (vt, rung_end, hi).
2701    full: std::collections::HashMap<(usize, usize, usize), DsparkSegGraph>,
2702    /// Largest n with every layer in [0, n) linear or full-attention (walk coverage).
2703    covered: usize,
2704    /// Every layer in [0, n) is linear or full-attention (no MLA/unknown mixers) — the
2705    /// full-verify capture walks all of them.
2706    walk_uniform: bool,
2707    /// Last `(captures, device graph-mem reserved bytes)` reading taken by
2708    /// `HybridModel::dspark_vg_admission_debt` — the two-point base of the MARGINAL debt
2709    /// projection (see `dspark_vg_debt_projection`; a mean-based reading extrapolated the
2710    /// pool's one-time shared allocation and reserved 8.5 GB of phantom VRAM).
2711    debt_obs: Option<(usize, usize)>,
2712}
2713
2714struct DsparkSegGraph {
2715    graph: cudarc::driver::CudaGraph,
2716    _keeper: Vec<Box<dyn std::any::Any + Send>>,
2717}
2718
2719/// Per-call arguments of [`HybridModel::qwen35_tparallel_fa_layer`] — one struct so the
2720/// eager walk and the slice-4c captured full-verify graphs hand the SAME body its two
2721/// modes without a second copy of the math.
2722pub(crate) struct FaLayerArgs<'a> {
2723    /// [T] per-row positions (device): rope reads them row-indexed; the seqs twins read
2724    /// them per-z (append slot = pos, T_kv = pos + 1).
2725    pub pos_d: &'a CudaSlice<i32>,
2726    /// Verify-level lazy per-row 1-element position buffers — only the per-row fallback
2727    /// arm builds/uses them (graph mode refuses that arm).
2728    pub pos_rows: &'a mut Option<Vec<CudaSlice<i32>>>,
2729    pub pos0: usize,
2730    pub seqs_append: bool,
2731    pub batch_fa_on: bool,
2732    /// Some((kv pointer table, offset-in-u64s, rung_end)) = captured-graph mode.
2733    pub graph_cap: Option<(&'a CudaSlice<u64>, usize, usize)>,
2734    /// ROUND-STREAM (lane/draftcost-moe, v0.100 train merge): Some((token stream, device
2735    /// round counter)) routes the FA attend through the dc rows kernels and the Linear
2736    /// mixer through `linear_attn_verify_t` (the stream arms the old inline body carried).
2737    /// Never armed together with `graph_cap` (the verify-level merge guard refuses).
2738    pub stream: Option<(&'a CudaSlice<u32>, &'a CudaSlice<i32>)>,
2739    /// VerifyCkpt for the stream-Linear arm's GdnStash install; None in graph mode and
2740    /// for FA layers that never touch it.
2741    pub ckpt: Option<&'a mut VerifyCkpt>,
2742}
2743
2744// SAFETY: `CudaGraph` is not marked Send by cudarc because its raw driver handles carry
2745// no automatic trait; CUDA driver graph handles are context-scoped rather than
2746// OS-thread-affine (the SpecPipeSessionPtr precedent above). The ctx lives in
2747// `HybridModel::dspark_vgraphs` behind a Mutex and every touch happens on the engine's
2748// single decode-stream thread.
2749unsafe impl Send for DsparkVerifyGraphs {}
2750
2751impl DsparkVerifyGraphs {
2752    /// Live capture count (segment + full graphs) — the denominator of
2753    /// [`dspark_vg_debt_projection`]'s observed bytes/capture mean.
2754    pub(crate) fn captures(&self) -> usize {
2755        self.graphs.len() + self.full.len()
2756    }
2757
2758    /// Take the marginal-growth debt reading and record this observation for the next one.
2759    /// Called under the pool mutex by `HybridModel::dspark_vg_admission_debt`.
2760    pub(crate) fn admission_debt(&mut self, reserved_bytes: usize) -> usize {
2761        let captures = self.captures();
2762        let debt =
2763            dspark_vg_debt_projection(captures, dspark_vg_cap(), reserved_bytes, self.debt_obs);
2764        if captures > 0 {
2765            match self.debt_obs {
2766                Some((c0, _)) if captures <= c0 => {}
2767                _ => self.debt_obs = Some((captures, reserved_bytes)),
2768            }
2769        }
2770        debt
2771    }
2772
2773    /// Build for this cache's shape. None when there are no linear layers, sizes are
2774    /// non-uniform, or the trunk keeps a gemma4 config (never on the qwen35 family).
2775    pub(crate) fn new(
2776        e: &Engine,
2777        cache: &Cache,
2778        t_max: usize,
2779        n_embd: usize,
2780    ) -> Result<Option<Self>, Box<dyn std::error::Error>> {
2781        let lin: Vec<usize> = (0..cache.recur.len())
2782            .filter(|&il| cache.recur[il].is_some())
2783            .collect();
2784        if lin.is_empty() || t_max < 2 {
2785            return Ok(None);
2786        }
2787        let first = cache.recur[lin[0]].as_ref().unwrap();
2788        let (conv_words, ssm_words) = (first.conv_state.len(), first.ssm_state.len());
2789        for &il in &lin {
2790            let rl = cache.recur[il].as_ref().unwrap();
2791            if rl.conv_state.len() != conv_words || rl.ssm_state.len() != ssm_words {
2792                return Ok(None);
2793            }
2794        }
2795        let n = lin.len();
2796        let mut lin_pos = std::collections::HashMap::with_capacity(n);
2797        for (k, &il) in lin.iter().enumerate() {
2798            lin_pos.insert(il, k);
2799        }
2800        // longest run of consecutive linear layers (save-scratch sizing)
2801        let mut max_run = 1usize;
2802        let mut run = 1usize;
2803        for w in lin.windows(2) {
2804            if w[1] == w[0] + 1 {
2805                run += 1;
2806                max_run = max_run.max(run);
2807            } else {
2808                run = 1;
2809            }
2810        }
2811        let rows = t_max - 1;
2812        let mut stash_conv = Vec::with_capacity(n);
2813        let mut stash_ssm = Vec::with_capacity(n);
2814        for _ in 0..n {
2815            stash_conv.push(e.uninit(rows * conv_words)?);
2816            stash_ssm.push(e.uninit(rows * ssm_words)?);
2817        }
2818        let host_table = vec![0u64; n * 6];
2819        let table_all = e.htod_u64(&host_table)?;
2820        // slice 4c: full-attention census for the full-verify graphs.
2821        let fa: Vec<usize> = (0..cache.kv.len())
2822            .filter(|&il| cache.kv[il].is_some())
2823            .collect();
2824        let mut fa_pos = std::collections::HashMap::with_capacity(fa.len());
2825        for (k, &il) in fa.iter().enumerate() {
2826            fa_pos.insert(il, k);
2827        }
2828        let n_layers = cache.kv.len().max(cache.recur.len());
2829        // exactly one of (linear state, kv cache) per layer — no MLA/unknown mixers.
2830        let walk_uniform = (0..n_layers).all(|il| {
2831            cache.recur.get(il).is_some_and(|r| r.is_some())
2832                != cache.kv.get(il).is_some_and(|k| k.is_some())
2833        });
2834        // Contiguous covered prefix: the largest n such that every layer in [0, n) is
2835        // linear or full-attention. The TRUNK walk is [0, layers.len()) and the cache
2836        // vecs can carry EXTRA state slots past it (the q38 export keeps the MTP head
2837        // layer's kv at the tail — hi == lin+fa never held, the s4c battery's zero
2838        // 'full' captures). The full-graph guard is walk coverage, not slot arithmetic.
2839        let covered = (0..n_layers)
2840            .take_while(|il| lin_pos.contains_key(il) || fa_pos.contains_key(il))
2841            .count();
2842        let t_cap = t_max;
2843        let fa_host_table = vec![0u64; fa.len() * 2 * t_cap];
2844        let fa_table = e.htod_u64(&fa_host_table)?;
2845        Ok(Some(Self {
2846            lin,
2847            lin_pos,
2848            table_all,
2849            host_table,
2850            stash_conv,
2851            stash_ssm,
2852            conv_words,
2853            ssm_words,
2854            stage: std::collections::HashMap::new(),
2855            tap_bufs: std::collections::HashMap::new(),
2856            graphs: std::collections::HashMap::new(),
2857            save_conv: e.uninit(n * conv_words)?,
2858            save_ssm: e.uninit(n * ssm_words)?,
2859            max_run,
2860            n_embd,
2861            round_slab: false,
2862            fa,
2863            fa_pos,
2864            fa_table,
2865            fa_host_table,
2866            t_cap,
2867            pos_stage: std::collections::HashMap::new(),
2868            full: std::collections::HashMap::new(),
2869            covered,
2870            walk_uniform,
2871            debt_obs: None,
2872        }))
2873    }
2874
2875    /// Rebuild the pointer tables from the live handles (once per verify — the gdn
2876    /// ping-pong swaps the canonical/alt handles between rounds; a fresh generation's
2877    /// cache buffers land at new addresses; a stale table would read the wrong state).
2878    pub(crate) fn refresh_tables(
2879        &mut self,
2880        e: &Engine,
2881        cache: &Cache,
2882    ) -> Result<(), Box<dyn std::error::Error>> {
2883        use cudarc::driver::DevicePtr;
2884        {
2885            let s = &e.gpu.stream();
2886            for (k, &il) in self.lin.iter().enumerate() {
2887                let rl = cache.recur[il].as_ref().unwrap();
2888                let (pc, _g0) = rl.conv_state.device_ptr(s);
2889                let (p0, _g1) = rl.ssm_state.device_ptr(s);
2890                let (p1, _g2) = rl.ssm_state_alt.device_ptr(s);
2891                let o = k * 6;
2892                self.host_table[o] = pc as u64;
2893                self.host_table[o + 1] = p0 as u64;
2894                self.host_table[o + 2] = p1 as u64;
2895                self.host_table[o + 3] = pc as u64;
2896                self.host_table[o + 4] = p1 as u64;
2897                self.host_table[o + 5] = p0 as u64;
2898            }
2899            for (k, &il) in self.fa.iter().enumerate() {
2900                let kvl = cache.kv[il].as_ref().unwrap();
2901                let (pk, _g0) = kvl.k.device_ptr(s);
2902                let (pv, _g1) = kvl.v.device_ptr(s);
2903                let o = k * 2 * self.t_cap;
2904                for z in 0..self.t_cap {
2905                    self.fa_host_table[o + 2 * z] = pk as u64;
2906                    self.fa_host_table[o + 2 * z + 1] = pv as u64;
2907                }
2908            }
2909        }
2910        e.htod_u64_into(&self.host_table, &mut self.table_all)?;
2911        if !self.fa_host_table.is_empty() {
2912            e.htod_u64_into(&self.fa_host_table, &mut self.fa_table)?;
2913        }
2914        Ok(())
2915    }
2916
2917    /// Slice 4c eligibility: Some(rung_end) when this round can replay (or capture) a
2918    /// full-verify graph — the whole walk [lo, hi) is covered, every layer is linear or
2919    /// full-attention, and ALL of the round's per-row t_kv values take the v4-seqs arm
2920    /// on ONE `fa_split_keys` ladder step that the rung also sits on (the straddle law;
2921    /// both gates are t_kv intervals, so ends-inside means all-inside). The rung is the
2922    /// round's next power of two — grid/partial sizing only (`n_splits_max` is pure
2923    /// stride; splits >= ns_eff write the empty partial the combine never reads), so one
2924    /// captured graph is bit-identical for every round the rung covers.
2925    #[allow(clippy::too_many_arguments)]
2926    pub(crate) fn full_rung(
2927        &self,
2928        model: &crate::hybrid::HybridModel,
2929        cache: &Cache,
2930        lo: usize,
2931        hi: usize,
2932        t: usize,
2933        seqs_arms_on: bool,
2934    ) -> Option<usize> {
2935        if std::env::var("MEMRA_DSPARK_FULLG_DEBUG").as_deref() == Ok("1") {
2936            static ONCE: std::sync::Once = std::sync::Once::new();
2937            let len0 = self
2938                .fa
2939                .first()
2940                .and_then(|&il| cache.kv[il].as_ref())
2941                .map(|k| k.len);
2942            ONCE.call_once(|| {
2943                eprintln!(
2944                    "[fullg-debug] walk_uniform={} covered={} seqs_arms_on={} fa_rows_on={} t={} lo={} hi={} lin={} fa={} t_cap={} len0={:?}",
2945                    self.walk_uniform, self.covered, seqs_arms_on, dspark_fa_rows_on(), t, lo, hi,
2946                    self.lin.len(), self.fa.len(), self.t_cap, len0
2947                );
2948            });
2949        }
2950        if !self.walk_uniform
2951            || !seqs_arms_on
2952            || !dspark_fa_rows_on()
2953            || t < 2
2954            || lo != 0
2955            || hi > self.covered
2956            || t > self.t_cap
2957            || self.fa.is_empty()
2958        {
2959            return None;
2960        }
2961        let cfg = &model.cfg;
2962        let head_dim_global = cfg.head_dim_k as usize;
2963        let nkv = cfg.n_head_kv as usize;
2964        let kvl0 = cache.kv[self.fa[0]].as_ref().unwrap();
2965        // the z-batched twins read stacked rows at the cache's kv dims — must equal the
2966        // projection stride (the body's guard, hoisted so ineligible models fall back
2967        // instead of refusing mid-capture).
2968        let geom = cfg.full_attention_geometry_at(self.fa[0] as u32);
2969        let kv_dim = geom.n_head_kv as usize * geom.head_dim_k as usize;
2970        if kvl0.kv_dim_k != kv_dim || kvl0.kv_dim_v != kv_dim {
2971            return None;
2972        }
2973        let len0 = kvl0.len;
2974        let (t_kv_first, t_kv_last) = (len0 + 1, len0 + t);
2975        if !crate::fa_seqs_eligible(t_kv_first, head_dim_global)
2976            || !crate::fa_seqs_eligible(t_kv_last, head_dim_global)
2977            || crate::fa_split_keys(t_kv_first, nkv) != crate::fa_split_keys(t_kv_last, nkv)
2978        {
2979            return None;
2980        }
2981        let rung = t_kv_last.next_power_of_two().max(256);
2982        if crate::fa_split_keys(rung, nkv) != crate::fa_split_keys(t_kv_last, nkv) {
2983            return None;
2984        }
2985        Some(rung)
2986    }
2987
2988    /// Run the WHOLE verify walk [lo, hi) as one captured graph at (vt=t, rung): stage
2989    /// the residual + refresh the per-vt position staging, capture on first encounter
2990    /// (2 executing warmups bracketed by a full linear-state save/restore; KV warmup
2991    /// appends write the exact slots the replay writes — idempotent), launch, then apply
2992    /// the host bookkeeping the captured body skipped (per-linear-layer parity swap for
2993    /// odd t, per-fa-layer len bump). Returns the fresh residual.
2994    #[allow(clippy::too_many_arguments)]
2995    pub(crate) fn run_full(
2996        &mut self,
2997        model: &crate::hybrid::HybridModel,
2998        e: &Engine,
2999        lo: usize,
3000        hi: usize,
3001        x: &CudaSlice<f32>,
3002        t: usize,
3003        pos0: usize,
3004        rung: usize,
3005        cache: &mut Cache,
3006    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3007        let n_embd = self.n_embd;
3008        if !self.stage.contains_key(&t) {
3009            let xin = e.uninit(t * n_embd)?;
3010            let xout = e.uninit(t * n_embd)?;
3011            self.stage.insert(t, (xin, xout));
3012        }
3013        if !self.pos_stage.contains_key(&t) {
3014            self.pos_stage.insert(t, e.htod_i32(&vec![0i32; t])?);
3015        }
3016        // Per-round refresh: position contents + input staging (both addresses are baked
3017        // by the captured bodies; only their CONTENTS change round to round).
3018        {
3019            let pos_host: Vec<i32> = (0..t).map(|r| (pos0 + r) as i32).collect();
3020            let pb = self.pos_stage.get_mut(&t).unwrap();
3021            e.htod_i32_into(pb, &pos_host)?;
3022            let (xin, _) = self.stage.get_mut(&t).unwrap();
3023            e.copy_into(xin, 0, x, t * n_embd)?;
3024        }
3025        let key = (t, rung, hi);
3026        if !self.full.contains_key(&key) {
3027            // The warmups EXECUTE the whole walk on live state — save every linear
3028            // layer's conv + canonical ssm first, restore after (KV needs no restore:
3029            // graph mode never bumps host lens and the appends write this round's own
3030            // slots).
3031            for (k, &il) in self.lin.iter().enumerate() {
3032                let rl = cache.recur[il].as_ref().unwrap();
3033                e.copy_into(
3034                    &mut self.save_conv,
3035                    k * self.conv_words,
3036                    &rl.conv_state,
3037                    self.conv_words,
3038                )?;
3039                e.copy_into(
3040                    &mut self.save_ssm,
3041                    k * self.ssm_words,
3042                    &rl.ssm_state,
3043                    self.ssm_words,
3044                )?;
3045            }
3046            let (graph, keeper) = {
3047                let table_all = &self.table_all;
3048                let lin_pos = &self.lin_pos;
3049                let fa_pos = &self.fa_pos;
3050                let fa_table = &self.fa_table;
3051                let t_cap = self.t_cap;
3052                let stash_conv = &mut self.stash_conv;
3053                let stash_ssm = &mut self.stash_ssm;
3054                let pos_d: &CudaSlice<i32> = &self.pos_stage[&t];
3055                let (xin, xout) = self
3056                    .stage
3057                    .get_mut(&t)
3058                    .map(|(a, b)| (&*a, b))
3059                    .expect("stage bucket created above");
3060                let cache_ref: &mut Cache = cache;
3061                let iflag = if std::env::var("MEMRA_DSPARK_VG_AUTOFREE").as_deref() == Ok("1") {
3062                    cudarc::driver::sys::CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH
3063                } else {
3064                    cudarc::driver::sys::CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_USE_NODE_PRIORITY
3065                };
3066                e.capture_graph_retained_flags(iflag, move |e| {
3067                    let mut xc: Option<CudaSlice<f32>> = None;
3068                    for il in lo..hi {
3069                        let xr: &CudaSlice<f32> = xc.as_ref().unwrap_or(xin);
3070                        let nx = if let Some(&k) = lin_pos.get(&il) {
3071                            model.qwen35_tparallel_linear_layer(
3072                                e,
3073                                il,
3074                                xr,
3075                                t,
3076                                cache_ref,
3077                                None,
3078                                Some((&mut stash_conv[k], &mut stash_ssm[k])),
3079                                Some((table_all, k * 6)),
3080                            )?
3081                        } else if let Some(&kf) = fa_pos.get(&il) {
3082                            let mut no_rows: Option<Vec<CudaSlice<i32>>> = None;
3083                            model.qwen35_tparallel_fa_layer(
3084                                e,
3085                                il,
3086                                xr,
3087                                t,
3088                                cache_ref,
3089                                FaLayerArgs {
3090                                    pos_d,
3091                                    pos_rows: &mut no_rows,
3092                                    pos0,
3093                                    seqs_append: true,
3094                                    batch_fa_on: true,
3095                                    graph_cap: Some((fa_table, kf * 2 * t_cap, rung)),
3096                                    stream: None,
3097                                    ckpt: None,
3098                                },
3099                            )?
3100                        } else {
3101                            return Err(format!(
3102                                "run_full: layer {il} is neither linear nor full-attention"
3103                            )
3104                            .into());
3105                        };
3106                        xc = Some(nx);
3107                    }
3108                    e.copy_into(xout, 0, xc.as_ref().unwrap(), t * n_embd)?;
3109                    Ok(())
3110                })?
3111            };
3112            // Undo the net host parity motion of the 3 body runs (each run swaps iff t
3113            // is odd -> 3 runs = net one swap), then restore the device state the
3114            // warmups consumed (walk scope only — layers past hi never executed). The
3115            // launch below then behaves exactly like one run.
3116            if t % 2 == 1 {
3117                for &il in &self.lin {
3118                    if il < lo || il >= hi {
3119                        continue;
3120                    }
3121                    let rl = cache.recur[il].as_mut().unwrap();
3122                    std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
3123                }
3124            }
3125            for (k, &il) in self.lin.iter().enumerate() {
3126                if il < lo || il >= hi {
3127                    continue;
3128                }
3129                let rl = cache.recur[il].as_mut().unwrap();
3130                let (cw, sw) = (self.conv_words, self.ssm_words);
3131                {
3132                    let sv = e.view(&self.save_conv, self.lin.len() * cw);
3133                    let win = sv.slice(k * cw..(k + 1) * cw);
3134                    e.copy_view_into(&mut rl.conv_state, 0, &win, cw)?;
3135                }
3136                {
3137                    let sv = e.view(&self.save_ssm, self.lin.len() * sw);
3138                    let win = sv.slice(k * sw..(k + 1) * sw);
3139                    e.copy_view_into(&mut rl.ssm_state, 0, &win, sw)?;
3140                }
3141            }
3142            if std::env::var("MEMRA_GRAPH_CENSUS").as_deref() == Ok("1") {
3143                if let Ok(c) = crate::graph_update::node_census(&graph) {
3144                    eprintln!("[dspark-vg-census] full vt={t} rung={rung} {c:?}");
3145                }
3146            }
3147            self.full.insert(
3148                key,
3149                DsparkSegGraph {
3150                    graph,
3151                    _keeper: keeper,
3152                },
3153            );
3154        }
3155        self.full[&key].graph.launch()?;
3156        // Host bookkeeping for the replayed body (captured host code does not re-run):
3157        // gdn parity swap per linear layer (t odd), kv len bump per fa layer — scoped
3158        // to the WALK [lo, hi): the cache can carry extra state slots past it (the MTP
3159        // head layer's kv) that the walk never touches.
3160        if t % 2 == 1 {
3161            for &il in &self.lin {
3162                if il < lo || il >= hi {
3163                    continue;
3164                }
3165                let rl = cache.recur[il].as_mut().unwrap();
3166                std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
3167            }
3168        }
3169        for &il in &self.fa {
3170            if il < lo || il >= hi {
3171                continue;
3172            }
3173            cache.kv[il].as_mut().unwrap().len += t;
3174        }
3175        let (_, xout) = self.stage.get(&t).unwrap();
3176        let mut out = e.uninit(t * n_embd)?;
3177        e.copy_into(&mut out, 0, xout, t * n_embd)?;
3178        Ok(out)
3179    }
3180
3181    /// Run layers [start, end) (all linear) as one captured graph at this vt: stage the
3182    /// residual into the bucket's x_in, capture on first encounter (2 executing warmups
3183    /// bracketed by a segment state save/restore), launch, then apply the host parity
3184    /// bookkeeping the captured body would have done. Returns the fresh residual.
3185    #[allow(clippy::too_many_arguments)]
3186    fn run_segment(
3187        &mut self,
3188        model: &crate::hybrid::HybridModel,
3189        e: &Engine,
3190        start: usize,
3191        end: usize,
3192        x: &CudaSlice<f32>,
3193        t: usize,
3194        cache: &mut Cache,
3195    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3196        let n_embd = self.n_embd;
3197        debug_assert!(end - start <= self.max_run);
3198        if !self.stage.contains_key(&t) {
3199            let xin = e.uninit(t * n_embd)?;
3200            let xout = e.uninit(t * n_embd)?;
3201            self.stage.insert(t, (xin, xout));
3202        }
3203        // Stage the residual at the bucket's baked input address.
3204        {
3205            let (xin, _) = self.stage.get_mut(&t).unwrap();
3206            e.copy_into(xin, 0, x, t * n_embd)?;
3207        }
3208        let key = (start, t);
3209        if !self.graphs.contains_key(&key) {
3210            // The 2 warmups EXECUTE the segment on live state — save conv + the canonical
3211            // ssm of every segment layer first, restore after, so the graph's first real
3212            // launch starts from the exact pre-round state (bytes gated e2e).
3213            for (k, il) in (start..end).enumerate() {
3214                let rl = cache.recur[il].as_ref().unwrap();
3215                e.copy_into(
3216                    &mut self.save_conv,
3217                    k * self.conv_words,
3218                    &rl.conv_state,
3219                    self.conv_words,
3220                )?;
3221                e.copy_into(
3222                    &mut self.save_ssm,
3223                    k * self.ssm_words,
3224                    &rl.ssm_state,
3225                    self.ssm_words,
3226                )?;
3227            }
3228            let (graph, keeper) = {
3229                let table_all = &self.table_all;
3230                let lin_pos = &self.lin_pos;
3231                let stash_conv = &mut self.stash_conv;
3232                let stash_ssm = &mut self.stash_ssm;
3233                let (xin, xout) = self
3234                    .stage
3235                    .get_mut(&t)
3236                    .map(|(a, b)| (&*a, b))
3237                    .expect("stage bucket created above");
3238                let cache_ref: &mut Cache = cache;
3239                // Slice 4 (fa-execupdate lane): USE_NODE_PRIORITY instead of
3240                // AUTO_FREE_ON_LAUNCH. The slice-3 measured limiter was AUTO_FREE's
3241                // launch-time mem-pool scan — 25.6 us per cuGraphLaunch x 16 segments
3242                // = ~0.41 ms/round, most of the eager-launch savings. The captured
3243                // body's cuMemAllocAsync transients are BALANCED by in-graph frees
3244                // (every transient drops inside the capture region — the generic
3245                // capture path's census precedent, 1589/1589), so AUTO_FREE has
3246                // nothing to reclaim and the graph is legal to instantiate without
3247                // it; PRIORITY is the flag the gemma slotted door ships for exactly
3248                // this reason (both alternatives drop the scan; UPLOAD via
3249                // cuGraphInstantiateWithFlags is WithParams-only and refused).
3250                // MEMRA_DSPARK_VG_AUTOFREE=1 reverts; MEMRA_GRAPH_CENSUS=1 prints
3251                // the node census at capture (the ALLOC==FREE receipt).
3252                let iflag = if std::env::var("MEMRA_DSPARK_VG_AUTOFREE").as_deref() == Ok("1") {
3253                    cudarc::driver::sys::CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH
3254                } else {
3255                    cudarc::driver::sys::CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_USE_NODE_PRIORITY
3256                };
3257                e.capture_graph_retained_flags(iflag, move |e| {
3258                    let mut xc: Option<CudaSlice<f32>> = None;
3259                    for il in start..end {
3260                        let k = lin_pos[&il];
3261                        let xr: &CudaSlice<f32> = xc.as_ref().unwrap_or(xin);
3262                        let nx = model.qwen35_tparallel_linear_layer(
3263                            e,
3264                            il,
3265                            xr,
3266                            t,
3267                            cache_ref,
3268                            None,
3269                            Some((&mut stash_conv[k], &mut stash_ssm[k])),
3270                            Some((table_all, k * 6)),
3271                        )?;
3272                        xc = Some(nx);
3273                    }
3274                    e.copy_into(xout, 0, xc.as_ref().unwrap(), t * n_embd)?;
3275                    Ok(())
3276                })?
3277            };
3278            // Undo the net host parity motion of the 3 body runs (each run swaps iff t
3279            // is odd -> 3 runs = net one swap), then restore the device state the
3280            // warmups consumed. The launch below then behaves exactly like one run.
3281            if t % 2 == 1 {
3282                for il in start..end {
3283                    let rl = cache.recur[il].as_mut().unwrap();
3284                    std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
3285                }
3286            }
3287            for (k, il) in (start..end).enumerate() {
3288                let rl = cache.recur[il].as_mut().unwrap();
3289                let (cw, sw) = (self.conv_words, self.ssm_words);
3290                {
3291                    let sv = e.view(&self.save_conv, self.lin.len() * cw);
3292                    let win = sv.slice(k * cw..(k + 1) * cw);
3293                    e.copy_view_into(&mut rl.conv_state, 0, &win, cw)?;
3294                }
3295                {
3296                    let sv = e.view(&self.save_ssm, self.lin.len() * sw);
3297                    let win = sv.slice(k * sw..(k + 1) * sw);
3298                    e.copy_view_into(&mut rl.ssm_state, 0, &win, sw)?;
3299                }
3300            }
3301            if std::env::var("MEMRA_GRAPH_CENSUS").as_deref() == Ok("1") {
3302                if let Ok(c) = crate::graph_update::node_census(&graph) {
3303                    eprintln!("[dspark-vg-census] seg={start}..{end} vt={t} {c:?}");
3304                }
3305            }
3306            self.graphs.insert(
3307                key,
3308                DsparkSegGraph {
3309                    graph,
3310                    _keeper: keeper,
3311                },
3312            );
3313        }
3314        self.graphs[&key].graph.launch()?;
3315        // Host parity bookkeeping for the replayed body (the captured host swaps do not
3316        // re-run at replay).
3317        if t % 2 == 1 {
3318            for il in start..end {
3319                let rl = cache.recur[il].as_mut().unwrap();
3320                std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
3321            }
3322        }
3323        let (_, xout) = self.stage.get(&t).unwrap();
3324        let mut out = e.uninit(t * n_embd)?;
3325        e.copy_into(&mut out, 0, xout, t * n_embd)?;
3326        Ok(out)
3327    }
3328
3329    /// Pool freeze check (`dspark_vg_cap`): below the ceiling new keys may capture.
3330    fn can_capture(&self) -> bool {
3331        self.graphs.len() + self.full.len() < dspark_vg_cap()
3332    }
3333
3334    /// Round-atomic segment-door readiness: TRUE when this round's walk can ride the
3335    /// per-(segment, vt) graphs without a NEW capture past the pool ceiling — every
3336    /// linear run in [lo, hi) already has its (run_start, t) key, or capture is still
3337    /// allowed. FALSE sends the WHOLE round down the eager cols-ckpt walk: a partial
3338    /// refusal would stash some layers in the ctx slabs and others in the round's cols
3339    /// while one commit reads only one of them.
3340    pub(crate) fn segments_ready(
3341        &self,
3342        model: &crate::hybrid::HybridModel,
3343        lo: usize,
3344        hi: usize,
3345        t: usize,
3346    ) -> bool {
3347        if self.can_capture() {
3348            return true;
3349        }
3350        let mut il = lo;
3351        while il < hi {
3352            if matches!(model.layers[il].mixer, Mixer::Linear(_)) {
3353                let start = il;
3354                while il < hi && matches!(model.layers[il].mixer, Mixer::Linear(_)) {
3355                    il += 1;
3356                }
3357                if !self.graphs.contains_key(&(start, t)) {
3358                    return false;
3359                }
3360            } else {
3361                il += 1;
3362            }
3363        }
3364        true
3365    }
3366
3367    /// Widest verify window this pool was built for. A caller whose round exceeds it must
3368    /// take the eager walk: the stash slabs hold `t_capacity() - 1` column rows, and slicing
3369    /// past them is a panic rather than a refusal.
3370    pub(crate) fn t_capacity(&self) -> usize {
3371        self.t_cap
3372    }
3373
3374    /// Slab row (conv, ssm) device pointers + lengths for the commit restore of column
3375    /// `row` (0-based) of layer `il`. None for non-linear layers.
3376    pub(crate) fn slab_row(
3377        &self,
3378        e: &Engine,
3379        il: usize,
3380        row: usize,
3381    ) -> Option<(u64, u64, usize, usize)> {
3382        use cudarc::driver::DevicePtr;
3383        let k = *self.lin_pos.get(&il)?;
3384        let s = &e.gpu.stream();
3385        let (pc, _g0) = self.stash_conv[k].device_ptr(s);
3386        let (ps, _g1) = self.stash_ssm[k].device_ptr(s);
3387        Some((
3388            pc as u64 + (row * self.conv_words * 4) as u64,
3389            ps as u64 + (row * self.ssm_words * 4) as u64,
3390            self.conv_words,
3391            self.ssm_words,
3392        ))
3393    }
3394}
3395
3396impl VerifyCkpt {
3397    fn new(n_layer: usize) -> Self {
3398        VerifyCkpt {
3399            gdn: (0..n_layer).map(|_| None).collect(),
3400            cols: (0..n_layer).map(|_| None).collect(),
3401        }
3402    }
3403}
3404
3405/// The stage-0/TX half of one PP verify. The boundary slot is the ownership token: stage 1
3406/// consumes exactly the slot selected by `tx()` / `tx_pipelined()`, never a slot inferred from
3407/// a logical round number.
3408struct VerifyBoundaryTicket {
3409    rt: &'static crate::pp::PpNRt,
3410    caller_stream: std::sync::Arc<cudarc::driver::CudaStream>,
3411    slot: usize,
3412    pos0: usize,
3413    t: usize,
3414    payload: usize,
3415    n_st: usize,
3416    pipelined: bool,
3417    pp_anatomy: bool,
3418    pp_started: std::time::Instant,
3419    reverse_ms: f64,
3420    stage0_ms: f64,
3421    tx_ms: f64,
3422    trace: Option<SpecPipeTraceCtx>,
3423}
3424
3425/// Explicit OPTIPIPE diagnostic control. Forced modes are set only by `optipipe-gate`; the
3426/// increment-2 controller can also be armed by the server's fresh-process research door.
3427#[derive(Clone, Copy, Debug, PartialEq, Eq)]
3428pub enum OptiForkGateMode {
3429    Disabled,
3430    Hit,
3431    Miss,
3432    Alternate,
3433    Abort,
3434    Controller,
3435}
3436
3437static OPTI_FORK_GATE_MODE: std::sync::atomic::AtomicU8 = std::sync::atomic::AtomicU8::new(0);
3438static OPTI_CONTROLLER_THRESHOLD: std::sync::atomic::AtomicU32 =
3439    std::sync::atomic::AtomicU32::new(0);
3440static OPTI_FORK_ATTEMPTS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
3441static OPTI_FORK_HITS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
3442static OPTI_FORK_MISSES: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
3443static OPTI_FORK_ABORT_DRAINS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
3444static OPTI_FORK_REFUSALS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
3445static OPTI_GATE_CHECKS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
3446static OPTI_GATE_ADMITS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
3447static OPTI_GATE_REJECTS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
3448static OPTI_RECONCILES: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
3449static OPTI_WASTED_DRAFT_TOKENS: std::sync::atomic::AtomicU64 =
3450    std::sync::atomic::AtomicU64::new(0);
3451static OPTI_SHADOW_DRAFT_TOKENS: std::sync::atomic::AtomicU64 =
3452    std::sync::atomic::AtomicU64::new(0);
3453static OPTI_BREAKER_TRIPS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
3454
3455impl OptiForkGateMode {
3456    fn code(self) -> u8 {
3457        match self {
3458            Self::Disabled => 0,
3459            Self::Hit => 1,
3460            Self::Miss => 2,
3461            Self::Alternate => 3,
3462            Self::Abort => 4,
3463            Self::Controller => 5,
3464        }
3465    }
3466
3467    fn configured() -> Self {
3468        match OPTI_FORK_GATE_MODE.load(std::sync::atomic::Ordering::Relaxed) {
3469            1 => Self::Hit,
3470            2 => Self::Miss,
3471            3 => Self::Alternate,
3472            4 => Self::Abort,
3473            5 => Self::Controller,
3474            _ => Self::Disabled,
3475        }
3476    }
3477
3478    fn action(self, generation: u64) -> OptiForkAction {
3479        match self {
3480            Self::Hit => OptiForkAction::Hit,
3481            Self::Miss => OptiForkAction::Miss,
3482            Self::Alternate if generation & 1 == 0 => OptiForkAction::Hit,
3483            Self::Alternate => OptiForkAction::Miss,
3484            Self::Abort => OptiForkAction::Abort,
3485            Self::Disabled | Self::Controller => {
3486                unreachable!("non-forced mode cannot choose a forced fork action")
3487            }
3488        }
3489    }
3490
3491    fn is_forced(self) -> bool {
3492        matches!(self, Self::Hit | Self::Miss | Self::Alternate | Self::Abort)
3493    }
3494}
3495
3496/// Arm or disarm the forced harness. Serving uses only `set_optipipe_controller_threshold`.
3497pub fn set_optipipe_gate_mode(mode: OptiForkGateMode) {
3498    OPTI_FORK_GATE_MODE.store(mode.code(), std::sync::atomic::Ordering::Relaxed);
3499}
3500
3501/// Arm the increment-2 diagnostic controller. The threshold applies to the uncalibrated
3502/// two-token draft-probability product. Serving can call this only through its explicit
3503/// fresh-process research door; the absent-door default remains byte-for-byte disabled.
3504pub fn set_optipipe_controller_threshold(threshold: f32) {
3505    assert!(threshold.is_finite() && (0.0..=1.0).contains(&threshold));
3506    OPTI_CONTROLLER_THRESHOLD.store(threshold.to_bits(), std::sync::atomic::Ordering::Relaxed);
3507    set_optipipe_gate_mode(OptiForkGateMode::Controller);
3508}
3509
3510#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
3511pub struct OptiForkGateStats {
3512    pub attempts: u64,
3513    pub hits: u64,
3514    pub misses: u64,
3515    pub abort_drains: u64,
3516    pub refusals: u64,
3517    pub gate_checks: u64,
3518    pub gate_admits: u64,
3519    pub gate_rejects: u64,
3520    pub reconciles: u64,
3521    pub wasted_draft_tokens: u64,
3522    pub shadow_draft_tokens: u64,
3523    pub breaker_trips: u64,
3524}
3525
3526#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
3527pub struct OptiForkStateIdentity {
3528    pub trunk_kv_bytes: usize,
3529    pub recurrent_bytes: usize,
3530    pub scratch_kv_bytes: usize,
3531    pub hidden_bytes: usize,
3532}
3533
3534pub fn reset_optipipe_gate_stats() {
3535    for counter in [
3536        &OPTI_FORK_ATTEMPTS,
3537        &OPTI_FORK_HITS,
3538        &OPTI_FORK_MISSES,
3539        &OPTI_FORK_ABORT_DRAINS,
3540        &OPTI_FORK_REFUSALS,
3541        &OPTI_GATE_CHECKS,
3542        &OPTI_GATE_ADMITS,
3543        &OPTI_GATE_REJECTS,
3544        &OPTI_RECONCILES,
3545        &OPTI_WASTED_DRAFT_TOKENS,
3546        &OPTI_SHADOW_DRAFT_TOKENS,
3547        &OPTI_BREAKER_TRIPS,
3548    ] {
3549        counter.store(0, std::sync::atomic::Ordering::Relaxed);
3550    }
3551}
3552
3553pub fn optipipe_gate_stats() -> OptiForkGateStats {
3554    let load = |v: &std::sync::atomic::AtomicU64| v.load(std::sync::atomic::Ordering::Relaxed);
3555    OptiForkGateStats {
3556        attempts: load(&OPTI_FORK_ATTEMPTS),
3557        hits: load(&OPTI_FORK_HITS),
3558        misses: load(&OPTI_FORK_MISSES),
3559        abort_drains: load(&OPTI_FORK_ABORT_DRAINS),
3560        refusals: load(&OPTI_FORK_REFUSALS),
3561        gate_checks: load(&OPTI_GATE_CHECKS),
3562        gate_admits: load(&OPTI_GATE_ADMITS),
3563        gate_rejects: load(&OPTI_GATE_REJECTS),
3564        reconciles: load(&OPTI_RECONCILES),
3565        wasted_draft_tokens: load(&OPTI_WASTED_DRAFT_TOKENS),
3566        shadow_draft_tokens: load(&OPTI_SHADOW_DRAFT_TOKENS),
3567        breaker_trips: load(&OPTI_BREAKER_TRIPS),
3568    }
3569}
3570
3571#[derive(Clone, Copy, Debug)]
3572struct OptiControllerPolicy {
3573    threshold: f32,
3574    consecutive_misses: u8,
3575    breaker_tripped: bool,
3576}
3577
3578impl OptiControllerPolicy {
3579    fn configured() -> Self {
3580        Self {
3581            threshold: f32::from_bits(
3582                OPTI_CONTROLLER_THRESHOLD.load(std::sync::atomic::Ordering::Relaxed),
3583            ),
3584            consecutive_misses: 0,
3585            breaker_tripped: false,
3586        }
3587    }
3588
3589    fn admit(&self, q_proxy: f32) -> bool {
3590        q_proxy.is_finite()
3591            && (0.0..=1.0).contains(&q_proxy)
3592            && (self.threshold == 0.0 || (!self.breaker_tripped && q_proxy >= self.threshold))
3593    }
3594
3595    /// Returns true exactly when this resolution newly trips the three-miss breaker.
3596    fn resolve(&mut self, hit: bool) -> bool {
3597        // q*=0 is the lane's explicit unconditional measurement arm. Its purpose is to price
3598        // every optimistic opportunity, so the safety breaker is measured separately and must
3599        // not silently turn this arm into "three attempts then serial".
3600        if self.threshold == 0.0 {
3601            self.consecutive_misses = 0;
3602            return false;
3603        }
3604        if hit {
3605            self.consecutive_misses = 0;
3606            return false;
3607        }
3608        self.consecutive_misses = self.consecutive_misses.saturating_add(1);
3609        if !self.breaker_tripped && self.consecutive_misses >= 3 {
3610            self.breaker_tripped = true;
3611            return true;
3612        }
3613        false
3614    }
3615}
3616
3617#[derive(Clone, Copy, Debug, PartialEq, Eq)]
3618enum OptiForkAction {
3619    Hit,
3620    Miss,
3621    Abort,
3622}
3623
3624#[derive(Clone, Copy, Debug, PartialEq, Eq)]
3625struct OptiForkGeneration {
3626    id: u64,
3627    slot: usize,
3628}
3629
3630#[derive(Default)]
3631struct OptiForkGenerationTracker {
3632    next: u64,
3633    live: [Option<u64>; 2],
3634}
3635
3636impl OptiForkGenerationTracker {
3637    fn reserve(&mut self) -> Result<OptiForkGeneration, Box<dyn std::error::Error>> {
3638        let generation = OptiForkGeneration {
3639            id: self.next,
3640            slot: (self.next & 1) as usize,
3641        };
3642        if let Some(live) = self.live[generation.slot] {
3643            return Err(format!(
3644                "optipipe snapshot slot {} still owns generation {live}; refusing to overwrite it",
3645                generation.slot,
3646            )
3647            .into());
3648        }
3649        self.next += 1;
3650        self.live[generation.slot] = Some(generation.id);
3651        Ok(generation)
3652    }
3653
3654    fn retire(&mut self, generation: OptiForkGeneration) -> Result<(), Box<dyn std::error::Error>> {
3655        match self.live[generation.slot] {
3656            Some(id) if id == generation.id => {
3657                self.live[generation.slot] = None;
3658                Ok(())
3659            }
3660            other => Err(format!(
3661                "optipipe generation teardown mismatch: ticket={} slot={} live={other:?}",
3662                generation.id, generation.slot,
3663            )
3664            .into()),
3665        }
3666    }
3667}
3668
3669struct OptiForkSeedGeneration {
3670    h_seed: CudaSlice<f32>,
3671    fill_prev: CudaSlice<f32>,
3672    scratch_len: usize,
3673}
3674
3675/// Allocate or refresh one full checkpoint through the engine that owns each PP stage. The
3676/// generic cache helper accepts one device and therefore cannot copy GDN state split across
3677/// devices. KV lengths and position stay host metadata; only recurrent buffers need stage-local
3678/// device ownership.
3679fn opti_snapshot_stage_owned(
3680    e: &Engine,
3681    cache: &Cache,
3682    rt: &'static crate::pp::PpNRt,
3683    fence: &[usize],
3684) -> Result<crate::cache::CacheSnapshot, Box<dyn std::error::Error>> {
3685    let n = cache.kv.len();
3686    let mut snapshot = crate::cache::CacheSnapshot {
3687        kv_len: vec![None; n],
3688        tp_kv_len: vec![None; n],
3689        conv: (0..n).map(|_| None).collect(),
3690        ssm: (0..n).map(|_| None).collect(),
3691        pos: cache.pos,
3692    };
3693    opti_snapshot_stage_owned_into(e, cache, rt, fence, &mut snapshot)?;
3694    Ok(snapshot)
3695}
3696
3697fn opti_snapshot_stage_owned_into(
3698    e: &Engine,
3699    cache: &Cache,
3700    rt: &'static crate::pp::PpNRt,
3701    fence: &[usize],
3702    snapshot: &mut crate::cache::CacheSnapshot,
3703) -> Result<(), Box<dyn std::error::Error>> {
3704    if fence.len() != rt.n_stages() + 1
3705        || snapshot.kv_len.len() != cache.kv.len()
3706        || snapshot.tp_kv_len.len() != cache.tp_kv.len()
3707    {
3708        return Err("optipipe stage-owned snapshot shape mismatch".into());
3709    }
3710    for stage in 0..rt.n_stages() {
3711        opti_snapshot_one_stage_owned_into(e, cache, rt, fence, stage, snapshot)?;
3712    }
3713    snapshot.pos = cache.pos;
3714    Ok(())
3715}
3716
3717/// Refresh one PP stage of a checkpoint. Increment 2 uses this split form so stage 0's
3718/// optimistic post-N state is captured before N+1 stage 0 is queued, while stage 1's matching
3719/// post-N state is captured only after N stage 1 is enqueued. Calling the all-stage helper at
3720/// either point would capture one side of the fork at the wrong generation.
3721fn opti_snapshot_one_stage_owned_into(
3722    e: &Engine,
3723    cache: &Cache,
3724    rt: &'static crate::pp::PpNRt,
3725    fence: &[usize],
3726    stage: usize,
3727    snapshot: &mut crate::cache::CacheSnapshot,
3728) -> Result<(), Box<dyn std::error::Error>> {
3729    if fence.len() != rt.n_stages() + 1
3730        || snapshot.kv_len.len() != cache.kv.len()
3731        || snapshot.tp_kv_len.len() != cache.tp_kv.len()
3732        || stage >= rt.n_stages()
3733    {
3734        return Err("optipipe single-stage snapshot shape mismatch".into());
3735    }
3736    let _scope = rt.enter(stage);
3737    let owner = rt.engine(stage, e);
3738    for il in fence[stage]..fence[stage + 1] {
3739        snapshot.kv_len[il] = cache.kv[il].as_ref().map(|kv| kv.len);
3740        snapshot.tp_kv_len[il] = cache.tp_kv[il]
3741            .as_ref()
3742            .map(crate::tp::ResidentTpKvCache::committed_len);
3743        match &cache.recur[il] {
3744            Some(recur) => {
3745                match snapshot.conv[il].as_mut() {
3746                    Some(dst) => {
3747                        owner.copy_into(dst, 0, &recur.conv_state, recur.conv_state.len())?
3748                    }
3749                    None => snapshot.conv[il] = Some(owner.clone_dtod(&recur.conv_state)?),
3750                }
3751                match snapshot.ssm[il].as_mut() {
3752                    Some(dst) => {
3753                        owner.copy_into(dst, 0, &recur.ssm_state, recur.ssm_state.len())?
3754                    }
3755                    None => snapshot.ssm[il] = Some(owner.clone_dtod(&recur.ssm_state)?),
3756                }
3757            }
3758            None if snapshot.conv[il].is_some() || snapshot.ssm[il].is_some() => {
3759                return Err(
3760                    format!("optipipe stage-owned snapshot layer {il} changed shape").into(),
3761                );
3762            }
3763            None => {}
3764        }
3765    }
3766    snapshot.pos = cache.pos;
3767    Ok(())
3768}
3769
3770/// Increment-1 persistent fork state. Exactly two snapshot/seed slots alternate; a live ticket
3771/// names its generation and keeps teardown fail-closed. Only stage 0 is allowed to mutate before
3772/// resolve, so the reconcile tables and conditional restores are stage-local.
3773struct OptiForkState {
3774    mode: OptiForkGateMode,
3775    controller: Option<OptiControllerPolicy>,
3776    generations: OptiForkGenerationTracker,
3777    active_snapshot_slot: usize,
3778    alternate_snapshot: crate::cache::CacheSnapshot,
3779    seeds: [OptiForkSeedGeneration; 2],
3780    rt: &'static crate::pp::PpNRt,
3781    fence: [usize; 3],
3782    split: usize,
3783    len_ptrs: CudaSlice<u64>,
3784    saved_lens: CudaSlice<i32>,
3785    forced_acc: CudaSlice<u32>,
3786    valid: CudaSlice<u32>,
3787    stage0_stream: std::sync::Arc<cudarc::driver::CudaStream>,
3788    logical_payload_bytes: [usize; 2],
3789}
3790
3791struct OptiForkTicket {
3792    generation: OptiForkGeneration,
3793    boundary: Option<VerifyBoundaryTicket>,
3794    drain: std::sync::Arc<cudarc::driver::CudaStream>,
3795    settled: bool,
3796}
3797
3798struct OptiControllerTicket {
3799    generation: OptiForkGeneration,
3800    boundary: Option<VerifyBoundaryTicket>,
3801    ckpt: Option<VerifyCkpt>,
3802    verify_tokens: [u32; 2],
3803    draft_prob: f32,
3804    eager_seed: Option<CudaSlice<f32>>,
3805    q_proxy: f32,
3806    scratch_len: usize,
3807    issued_at: std::time::Instant,
3808    drain: std::sync::Arc<cudarc::driver::CudaStream>,
3809    settled: bool,
3810}
3811
3812struct OptiControllerPrepared {
3813    verify_tokens: [u32; 2],
3814    draft_prob: f32,
3815    eager_seed: Option<CudaSlice<f32>>,
3816    q_proxy: f32,
3817    scratch_len: usize,
3818}
3819
3820impl OptiControllerTicket {
3821    fn take_boundary(&mut self) -> VerifyBoundaryTicket {
3822        self.boundary
3823            .take()
3824            .expect("controller boundary ticket already consumed")
3825    }
3826
3827    fn take_ckpt(&mut self) -> VerifyCkpt {
3828        self.ckpt
3829            .take()
3830            .expect("controller verify checkpoint already consumed")
3831    }
3832
3833    fn take_eager_seed(&mut self) -> Option<CudaSlice<f32>> {
3834        self.eager_seed.take()
3835    }
3836
3837    fn settle(&mut self) {
3838        self.settled = true;
3839    }
3840}
3841
3842impl Drop for OptiControllerTicket {
3843    fn drop(&mut self) {
3844        if !self.settled {
3845            let _ = self.drain.synchronize();
3846            OPTI_FORK_ABORT_DRAINS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
3847        }
3848    }
3849}
3850
3851impl OptiForkTicket {
3852    fn take_boundary(&mut self) -> VerifyBoundaryTicket {
3853        self.boundary
3854            .take()
3855            .expect("fork ticket boundary already consumed")
3856    }
3857
3858    fn settle(&mut self) {
3859        self.settled = true;
3860    }
3861}
3862
3863impl Drop for OptiForkTicket {
3864    fn drop(&mut self) {
3865        if !self.settled {
3866            let _ = self.drain.synchronize();
3867            OPTI_FORK_ABORT_DRAINS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
3868        }
3869    }
3870}
3871
3872impl OptiForkState {
3873    #[allow(clippy::too_many_arguments)]
3874    fn new(
3875        e: &Engine,
3876        cache: &Cache,
3877        mode: OptiForkGateMode,
3878        alternate_snapshot: crate::cache::CacheSnapshot,
3879        h_seed: &CudaSlice<f32>,
3880        fill_prev: &CudaSlice<f32>,
3881        rt: &'static crate::pp::PpNRt,
3882        split: usize,
3883        n_layer: usize,
3884    ) -> Result<Self, Box<dyn std::error::Error>> {
3885        let fence = [0, split, n_layer];
3886        let mut logical_payload_bytes = [0usize; 2];
3887        for stage in 0..2 {
3888            for il in fence[stage]..fence[stage + 1] {
3889                logical_payload_bytes[stage] += alternate_snapshot.conv[il]
3890                    .as_ref()
3891                    .map_or(0, |v| v.len() * std::mem::size_of::<f32>());
3892                logical_payload_bytes[stage] += alternate_snapshot.ssm[il]
3893                    .as_ref()
3894                    .map_or(0, |v| v.len() * std::mem::size_of::<f32>());
3895            }
3896        }
3897        let seeds = [
3898            OptiForkSeedGeneration {
3899                h_seed: e.clone_dtod(h_seed)?,
3900                fill_prev: e.clone_dtod(fill_prev)?,
3901                scratch_len: 0,
3902            },
3903            OptiForkSeedGeneration {
3904                h_seed: e.clone_dtod(h_seed)?,
3905                fill_prev: e.clone_dtod(fill_prev)?,
3906                scratch_len: 0,
3907            },
3908        ];
3909        let (len_ptrs, saved_lens, forced_acc, valid, stage0_stream) = {
3910            let _stage = rt.enter(0);
3911            let e0 = rt.engine(0, e);
3912            (
3913                crate::round_stream::kv_len_ptr_table_range(e0, cache, 0..split, None)?,
3914                e0.htod_i32(&vec![0; split])?,
3915                e0.alloc_u32_zeroed(2)?,
3916                e0.alloc_u32_zeroed(1)?,
3917                e0.stream(),
3918            )
3919        };
3920        logical_payload_bytes[0] += seeds
3921            .iter()
3922            .map(|seed| (seed.h_seed.len() + seed.fill_prev.len()) * std::mem::size_of::<f32>())
3923            .sum::<usize>();
3924        logical_payload_bytes[0] += len_ptrs.len() * std::mem::size_of::<u64>()
3925            + saved_lens.len() * std::mem::size_of::<i32>()
3926            + forced_acc.len() * std::mem::size_of::<u32>()
3927            + valid.len() * std::mem::size_of::<u32>();
3928        Ok(Self {
3929            mode,
3930            controller: (mode == OptiForkGateMode::Controller)
3931                .then(OptiControllerPolicy::configured),
3932            generations: OptiForkGenerationTracker::default(),
3933            active_snapshot_slot: 0,
3934            alternate_snapshot,
3935            seeds,
3936            rt,
3937            fence,
3938            split,
3939            len_ptrs,
3940            saved_lens,
3941            forced_acc,
3942            valid,
3943            stage0_stream,
3944            logical_payload_bytes,
3945        })
3946    }
3947
3948    fn reserve(
3949        &mut self,
3950        current_snapshot: &mut crate::cache::CacheSnapshot,
3951    ) -> Result<OptiForkGeneration, Box<dyn std::error::Error>> {
3952        let generation = self.generations.reserve()?;
3953        if generation.slot != self.active_snapshot_slot {
3954            std::mem::swap(current_snapshot, &mut self.alternate_snapshot);
3955            self.active_snapshot_slot = generation.slot;
3956        }
3957        Ok(generation)
3958    }
3959
3960    fn capture_seed(
3961        &mut self,
3962        e: &Engine,
3963        generation: OptiForkGeneration,
3964        h_seed: &CudaSlice<f32>,
3965        fill_prev: &CudaSlice<f32>,
3966        scratch_len: usize,
3967    ) -> Result<(), Box<dyn std::error::Error>> {
3968        let seed = &mut self.seeds[generation.slot];
3969        e.copy_into(&mut seed.h_seed, 0, h_seed, h_seed.len())?;
3970        e.copy_into(&mut seed.fill_prev, 0, fill_prev, fill_prev.len())?;
3971        seed.scratch_len = scratch_len;
3972        Ok(())
3973    }
3974
3975    fn ticket(
3976        &self,
3977        generation: OptiForkGeneration,
3978        boundary: VerifyBoundaryTicket,
3979    ) -> OptiForkTicket {
3980        OptiForkTicket {
3981            generation,
3982            boundary: Some(boundary),
3983            drain: self.stage0_stream.clone(),
3984            settled: false,
3985        }
3986    }
3987
3988    #[allow(clippy::too_many_arguments)]
3989    fn controller_ticket(
3990        &self,
3991        generation: OptiForkGeneration,
3992        boundary: VerifyBoundaryTicket,
3993        ckpt: VerifyCkpt,
3994        verify_tokens: [u32; 2],
3995        draft_prob: f32,
3996        eager_seed: Option<CudaSlice<f32>>,
3997        q_proxy: f32,
3998        scratch_len: usize,
3999    ) -> OptiControllerTicket {
4000        OptiControllerTicket {
4001            generation,
4002            boundary: Some(boundary),
4003            ckpt: Some(ckpt),
4004            verify_tokens,
4005            draft_prob,
4006            eager_seed,
4007            q_proxy,
4008            scratch_len,
4009            issued_at: std::time::Instant::now(),
4010            drain: self.stage0_stream.clone(),
4011            settled: false,
4012        }
4013    }
4014
4015    fn reserve_successor(&mut self) -> Result<OptiForkGeneration, Box<dyn std::error::Error>> {
4016        self.generations.reserve()
4017    }
4018
4019    fn successor_snapshot_mut(&mut self) -> &mut crate::cache::CacheSnapshot {
4020        &mut self.alternate_snapshot
4021    }
4022
4023    fn promote_successor_snapshot(
4024        &mut self,
4025        current_snapshot: &mut crate::cache::CacheSnapshot,
4026        generation: OptiForkGeneration,
4027    ) {
4028        std::mem::swap(current_snapshot, &mut self.alternate_snapshot);
4029        self.active_snapshot_slot = generation.slot;
4030    }
4031
4032    fn queue_actual_reconcile(
4033        &mut self,
4034        e: &Engine,
4035        snapshot: &crate::cache::CacheSnapshot,
4036        acc: &CudaSlice<u32>,
4037        optimistic_pending: u32,
4038        base: usize,
4039    ) -> Result<(), Box<dyn std::error::Error>> {
4040        let saved: Vec<i32> = (0..self.split)
4041            .map(|il| snapshot.kv_len[il].map(|v| v as i32).unwrap_or(0))
4042            .collect();
4043        // Serving keeps the caller/accept walk on the head (stage-1) device. Record the accept
4044        // decision point there and append a wait to stage 0 after its optimistic successor/TX;
4045        // the validity/reconcile kernels must never peer-read acc before it is written. The
4046        // increment-1 harness uses primary stage 0, where stream order already provides this.
4047        if self.rt.engine(0, e).ctx().ordinal() != e.ctx().ordinal() {
4048            self.rt.fence_stages_behind(&e.stream())?;
4049        }
4050        let _stage = self.rt.enter(0);
4051        let e0 = self.rt.engine(0, e);
4052        e0.htod_i32_into(&mut self.saved_lens, &saved)?;
4053        e0.spec_fork_valid(acc, optimistic_pending, &mut self.valid)?;
4054        e0.spec_fork_reconcile_kv(
4055            &self.len_ptrs,
4056            &self.saved_lens,
4057            acc,
4058            &self.valid,
4059            base,
4060            self.split,
4061        )
4062    }
4063
4064    fn finish_actual_reconcile(
4065        &mut self,
4066        e: &Engine,
4067        cache: &mut Cache,
4068        snapshot: &crate::cache::CacheSnapshot,
4069        n_acc: usize,
4070        base: usize,
4071        hit: bool,
4072    ) -> Result<(), Box<dyn std::error::Error>> {
4073        if hit {
4074            return Ok(());
4075        }
4076        let len_delta = base + n_acc;
4077        for il in 0..self.split {
4078            if let (Some(kv), Some(saved)) = (cache.kv[il].as_mut(), snapshot.kv_len[il]) {
4079                kv.len = saved + len_delta;
4080            }
4081        }
4082        {
4083            let _stage = self.rt.enter(1);
4084            let e1 = self.rt.engine(1, e);
4085            for il in self.split..self.fence[2] {
4086                if let (Some(kv), Some(saved)) = (cache.kv[il].as_mut(), snapshot.kv_len[il]) {
4087                    kv.len = saved + len_delta;
4088                    e1.set_i32_one(&mut kv.len_d, kv.len as i32)?;
4089                }
4090            }
4091        }
4092        self.rt.publish_to(0, &e.stream())?;
4093        Ok(())
4094    }
4095
4096    fn cancel_controller_ticket(
4097        &mut self,
4098        e: &Engine,
4099        cache: &mut Cache,
4100        scratch: &mut MtpScratch,
4101        snapshot: &crate::cache::CacheSnapshot,
4102        ticket: &mut OptiControllerTicket,
4103    ) -> Result<(), Box<dyn std::error::Error>> {
4104        {
4105            let _stage = self.rt.enter(0);
4106            let e0 = self.rt.engine(0, e);
4107            for il in 0..self.split {
4108                if let (Some(kv), Some(saved)) = (cache.kv[il].as_mut(), snapshot.kv_len[il]) {
4109                    kv.len = saved;
4110                    e0.set_i32_one(&mut kv.len_d, saved as i32)?;
4111                }
4112            }
4113        }
4114        scratch.set_len(e, snapshot.pos)?;
4115        ticket.settle();
4116        self.generations.retire(ticket.generation)?;
4117        OPTI_FORK_ABORT_DRAINS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
4118        OPTI_WASTED_DRAFT_TOKENS.fetch_add(2, std::sync::atomic::Ordering::Relaxed);
4119        eprintln!(
4120            "[opti-controller] tail-drain generation={} slot={}",
4121            ticket.generation.id, ticket.generation.slot,
4122        );
4123        Ok(())
4124    }
4125
4126    #[allow(clippy::too_many_arguments)]
4127    fn reconcile(
4128        &mut self,
4129        e: &Engine,
4130        cache: &mut Cache,
4131        scratch: &mut MtpScratch,
4132        snapshot: &crate::cache::CacheSnapshot,
4133        h_seed: &mut CudaSlice<f32>,
4134        fill_prev: &mut CudaSlice<f32>,
4135        generation: OptiForkGeneration,
4136        action: OptiForkAction,
4137        optimistic_pending: u32,
4138    ) -> Result<(), Box<dyn std::error::Error>> {
4139        debug_assert!(action != OptiForkAction::Abort);
4140        let miss_started = std::time::Instant::now();
4141        let keep = action == OptiForkAction::Hit;
4142        let saved: Vec<i32> = (0..self.split)
4143            .map(|il| snapshot.kv_len[il].map(|v| v as i32).unwrap_or(0))
4144            .collect();
4145        let seed = &self.seeds[generation.slot];
4146        {
4147            let _stage = self.rt.enter(0);
4148            let e0 = self.rt.engine(0, e);
4149            e0.htod_i32_into(&mut self.saved_lens, &saved)?;
4150            let forced = if keep {
4151                [1u32, optimistic_pending]
4152            } else {
4153                [0u32, optimistic_pending]
4154            };
4155            e0.htod_u32_into(&mut self.forced_acc, &forced)?;
4156            e0.spec_fork_valid(&self.forced_acc, optimistic_pending, &mut self.valid)?;
4157            e0.spec_fork_reconcile_kv(
4158                &self.len_ptrs,
4159                &self.saved_lens,
4160                &self.forced_acc,
4161                &self.valid,
4162                0,
4163                self.split,
4164            )?;
4165            for il in 0..self.split {
4166                if let Some(recur) = cache.recur[il].as_mut() {
4167                    let conv = snapshot.conv[il]
4168                        .as_ref()
4169                        .ok_or("optipipe stage0 snapshot missing conv state")?;
4170                    let ssm = snapshot.ssm[il]
4171                        .as_ref()
4172                        .ok_or("optipipe stage0 snapshot missing ssm state")?;
4173                    e0.spec_fork_restore_f32(conv, &mut recur.conv_state, &self.valid)?;
4174                    e0.spec_fork_restore_f32(ssm, &mut recur.ssm_state, &self.valid)?;
4175                }
4176            }
4177            e0.spec_fork_restore_f32(&seed.h_seed, h_seed, &self.valid)?;
4178            e0.spec_fork_restore_f32(&seed.fill_prev, fill_prev, &self.valid)?;
4179        }
4180
4181        if keep {
4182            OPTI_FORK_HITS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
4183            return Ok(());
4184        }
4185
4186        for il in 0..self.split {
4187            if let (Some(kv), Some(saved)) = (cache.kv[il].as_mut(), snapshot.kv_len[il]) {
4188                kv.len = saved;
4189            }
4190        }
4191        scratch.set_len(e, seed.scratch_len)?;
4192        // Targeted E_restart: publish only stage 0's reconcile to the caller, then bound the
4193        // forced diagnostic so the retained number is the actual miss cost, not enqueue time.
4194        let caller = e.stream();
4195        self.rt.publish_to(0, &caller)?;
4196        caller.synchronize()?;
4197        let miss_ms = miss_started.elapsed().as_secs_f64() * 1e3;
4198        eprintln!(
4199            "[opti-fork-reconcile] generation={} slot={} miss_ms={miss_ms:.3}",
4200            generation.id, generation.slot,
4201        );
4202        OPTI_FORK_MISSES.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
4203        Ok(())
4204    }
4205
4206    fn retire(&mut self, generation: OptiForkGeneration) -> Result<(), Box<dyn std::error::Error>> {
4207        self.generations.retire(generation)
4208    }
4209}
4210
4211fn rewind_tp_kv_verified_prefix(
4212    tp_kv: &mut [Option<crate::tp::ResidentTpKvCache>],
4213    saved_lens: &[Option<usize>],
4214    accepted: usize,
4215) -> Result<(), Box<dyn std::error::Error>> {
4216    if tp_kv.len() != saved_lens.len() {
4217        return Err("spec TP KV snapshot shape mismatch".into());
4218    }
4219    for (layer, (cache, saved)) in tp_kv.iter_mut().zip(saved_lens).enumerate() {
4220        match (cache.as_mut(), *saved) {
4221            (Some(cache), Some(saved)) => {
4222                let target = saved
4223                    .checked_add(accepted)
4224                    .ok_or("spec TP KV committed length overflow")?;
4225                cache.rewind_to(target)?;
4226            }
4227            (None, None) => {}
4228            _ => {
4229                return Err(
4230                    format!("spec TP KV layer {layer} changed shape since its snapshot").into(),
4231                );
4232            }
4233        }
4234    }
4235    Ok(())
4236}
4237
4238/// MEMRA_SPEC_ROUND_PROF counters: whole-round wall, so the round can be weighed against the
4239/// draft-step ([spec-anatomy]) and verify-walk ([tcol-prof]) splits we already print.
4240static ROUND_PROF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
4241static ROUND_MS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
4242static ROUND_N: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
4243
4244impl HybridModel {
4245    fn mtp_head_count(&self) -> usize {
4246        usize::from(self.mtp.is_some()) + self.mtp_extra.len()
4247    }
4248
4249    fn mtp_head_at(&self, index: usize) -> &MtpHead {
4250        if index == 0 {
4251            self.mtp.as_ref().expect("MTP head 0 is unavailable")
4252        } else {
4253            &self.mtp_extra[index - 1]
4254        }
4255    }
4256
4257    fn new_mtp_scratch(
4258        &self,
4259        e: &Engine,
4260        cap: usize,
4261    ) -> Result<MtpScratch, Box<dyn std::error::Error>> {
4262        let mut scratch = MtpScratch::new(
4263            e,
4264            &self.cfg,
4265            &self.plan,
4266            cap,
4267            self.mtp.as_ref().and_then(|head| head.geom.as_ref()),
4268        )?;
4269        for head in &self.mtp_extra {
4270            scratch.push_plane(e, &self.cfg, &self.plan, head.geom.as_ref())?;
4271        }
4272        Ok(scratch)
4273    }
4274
4275    fn opti_graph_draft_step(
4276        &self,
4277        e: &Engine,
4278        mtp: &MtpHead,
4279        dctx: &mut DraftGraphCtx,
4280        scratch: &mut MtpScratch,
4281        d_vocab: usize,
4282    ) -> Result<(u32, f32), Box<dyn std::error::Error>> {
4283        // dcw door: one replay appends one device-counter row; pre-arm ring headroom
4284        // host-side before launching (no-op on flat planes).
4285        if step35_draft_dcw_on() {
4286            scratch.ensure_dcw_headroom(e, 2)?;
4287        }
4288        dctx.graph
4289            .as_ref()
4290            .ok_or("optipipe controller requires the greedy draft graph")?
4291            .launch()?;
4292        scratch.kv.len += 1;
4293        let idx = e.dtoh_u32_one(&dctx.g_tok)?;
4294        if (idx as usize) >= d_vocab {
4295            return Err(
4296                format!("optipipe draft argmax sentinel 0x{idx:08x} >= d_vocab {d_vocab}").into(),
4297            );
4298        }
4299        let probability = e.dtoh(&dctx.g_p)?[0];
4300        if !probability.is_finite() || !(0.0..=1.0).contains(&probability) {
4301            return Err(format!("optipipe draft probability is invalid: {probability}").into());
4302        }
4303        let token = match &mtp.d2t {
4304            Some(map) => map[idx as usize],
4305            None => idx,
4306        };
4307        if token != idx {
4308            e.set_u32_one(&mut dctx.g_tok, token)?;
4309        }
4310        Ok((token, probability))
4311    }
4312
4313    #[allow(clippy::too_many_arguments)]
4314    fn opti_controller_draft_step(
4315        &self,
4316        e: &Engine,
4317        mtp: &MtpHead,
4318        dctx: &mut DraftGraphCtx,
4319        scratch: &mut MtpScratch,
4320        d_vocab: usize,
4321        eager_state: &mut Option<(u32, CudaSlice<f32>)>,
4322        eager_pos: usize,
4323        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
4324        round_graph_ok: bool,
4325    ) -> Result<(u32, f32), Box<dyn std::error::Error>> {
4326        // GRAPH-LAUNCH HEADROOM GUARD (see GRAPH_LAUNCH_MIN_FREE): `round_graph_ok` is
4327        // the round's headroom snapshot. Below the floor the main draft arm already ran
4328        // eager (13651-class gate), which seeded `eager_state`, so the controller probe
4329        // rides its eager twin below instead of replaying the draft graph into an
4330        // exhausted card. The seed-unavailable Err beneath stays the recoverable
4331        // fail-closed for the shapes that never seed it.
4332        if dctx.graph.is_some() && round_graph_ok {
4333            return self.opti_graph_draft_step(e, mtp, dctx, scratch, d_vocab);
4334        }
4335        let (input_token, input_seed) = eager_state
4336            .take()
4337            .ok_or("optipipe eager continuation seed is unavailable")?;
4338        let (logits, next_seed) = self.mtp_head_forward_dev(
4339            e,
4340            mtp,
4341            input_token,
4342            &input_seed,
4343            scratch,
4344            eager_pos,
4345            embd_dev,
4346            None,
4347        )?;
4348        let token_d = e.argmax_token_device(&logits, d_vocab)?;
4349        let idx = e.dtoh_u32_one(&token_d)?;
4350        if (idx as usize) >= d_vocab {
4351            return Err(format!(
4352                "optipipe eager draft argmax sentinel 0x{idx:08x} >= d_vocab {d_vocab}"
4353            )
4354            .into());
4355        }
4356        let probability_d = e.prob_of_token_device(&logits, &token_d, d_vocab)?;
4357        let probability = e.dtoh(&probability_d)?[0];
4358        if !probability.is_finite() || !(0.0..=1.0).contains(&probability) {
4359            return Err(
4360                format!("optipipe eager draft probability is invalid: {probability}").into(),
4361            );
4362        }
4363        let token = match &mtp.d2t {
4364            Some(map) => map[idx as usize],
4365            None => idx,
4366        };
4367        *eager_state = Some((token, next_seed));
4368        Ok((token, probability))
4369    }
4370
4371    /// NextN head forward for ONE draft token (§A ops 1-13, T=1).
4372    /// Inputs: `e_tok` = the token to predict FROM (last committed / previous draft); `h_seed` =
4373    /// the trunk's pre-output_norm hidden of that token (§A op 2 input). `mtp_pos` = absolute
4374    /// position of the token being predicted from. Returns (draft_logits[n_vocab] host, h_nextn dev).
4375    /// `h_nextn` (§A op 10) becomes `h_seed` for the next autoregressive draft step.
4376    /// Device-resident: returns draft logits ON DEVICE (no [n_vocab] dtoh). The greedy draft
4377    /// loop only needs argmax — paired with `argmax_token_device` this cuts the ~600KB logits
4378    /// transfer + host argmax per draft token from the K-token draft chain.
4379    #[allow(clippy::too_many_arguments)]
4380    fn mtp_head_forward_dev(
4381        &self,
4382        e: &Engine,
4383        mtp: &MtpHead,
4384        e_tok: u32,
4385        h_seed: &CudaSlice<f32>,
4386        scratch: &mut MtpScratch,
4387        mtp_pos: usize,
4388        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
4389        mask: Option<(&CudaSlice<u32>, usize)>,
4390    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4391        self.mtp_head_forward_dev_at(e, mtp, e_tok, h_seed, scratch, 0, mtp_pos, embd_dev, mask)
4392    }
4393
4394    #[allow(clippy::too_many_arguments)]
4395    fn mtp_head_forward_dev_at(
4396        &self,
4397        e: &Engine,
4398        mtp: &MtpHead,
4399        e_tok: u32,
4400        h_seed: &CudaSlice<f32>,
4401        scratch: &mut MtpScratch,
4402        scratch_index: usize,
4403        mtp_pos: usize,
4404        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
4405        // DRAFT-SIDE GRAMMAR MASK (lane/draft-mask): (packed draft-vocab allowed set, words).
4406        // Applied to the head logits BEFORE they are returned, so every consumer (argmax,
4407        // gumbel draw, p-min prob) sees the grammar-legal row. None = unmasked (pre-lane).
4408        mask: Option<(&CudaSlice<u32>, usize)>,
4409    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4410        // MEMRA_SPEC_ANATOMY=1 — eager-step phase timers (diagnostic only). Phase boundaries
4411        // sync the stream, so absolute time inflates; the BREAKDOWN is the signal. Cumulative
4412        // summary on stderr every 128 steps: glue (embed..attn_norm), attn, ffn, head.
4413        use std::sync::atomic::{AtomicU64, Ordering::Relaxed};
4414        static ANAT_NS: [AtomicU64; 5] = [
4415            AtomicU64::new(0),
4416            AtomicU64::new(0),
4417            AtomicU64::new(0),
4418            AtomicU64::new(0),
4419            AtomicU64::new(0),
4420        ];
4421        static ANAT_STEPS: AtomicU64 = AtomicU64::new(0);
4422        let anat = {
4423            static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
4424            *ON.get_or_init(|| std::env::var("MEMRA_SPEC_ANATOMY").as_deref() == Ok("1"))
4425        };
4426        if anat {
4427            e.stream().synchronize()?; // drain prior queue so phase 0 starts clean
4428        }
4429        let t_all = std::time::Instant::now();
4430        let mut t_ph = std::time::Instant::now();
4431        let mut anat_mark = |i: usize,
4432                             e: &Engine,
4433                             t: &mut std::time::Instant|
4434         -> Result<(), Box<dyn std::error::Error>> {
4435            if anat {
4436                e.stream().synchronize()?;
4437                ANAT_NS[i].fetch_add(t.elapsed().as_nanos() as u64, Relaxed);
4438                *t = std::time::Instant::now();
4439            }
4440            Ok(())
4441        };
4442        let cfg = &self.cfg;
4443        let n_embd = cfg.n_embd as usize;
4444        // Distilled-student geometry: the block runs at the INNER width `di` (eh_proj out /
4445        // attn / ffn); the n_embd interface (embed, norms in, carrier out, head in) is unchanged.
4446        let di = mtp.geom.as_ref().map(|g| g.d_inner).unwrap_or(n_embd);
4447        let eps = cfg.rms_eps;
4448        let pos_d = e.htod_i32(&[mtp_pos as i32])?;
4449
4450        // op A: a resident table transfers one 4B token id. The exact host-row capacity path
4451        // expands this one row on CPU and transfers n_embd f32 values instead.
4452        let e_emb = match embd_dev {
4453            Some((g, qt, rb)) => e.embed_gather_device_t(g, &[e_tok], n_embd, qt, rb)?,
4454            None => e.htod(&self.embd.gather(n_embd, &[e_tok]))?,
4455        };
4456
4457        // op 1/2: e_norm = RMSNorm(e, enorm); h_norm = RMSNorm(h_seed, hnorm)
4458        let mut e_norm = e.zeros(n_embd)?;
4459        e.rms_norm(&e_emb, mtp.enorm.float_data(), &mut e_norm, n_embd, 1, eps)?;
4460        let mut h_norm = e.zeros(n_embd)?;
4461        e.rms_norm(h_seed, mtp.hnorm.float_data(), &mut h_norm, n_embd, 1, eps)?;
4462
4463        // op 3: concat = [e_norm ; h_norm] -> [2*n_embd], e_norm in [0,n_embd), h_norm in [n_embd,2n_embd)
4464        let mut concat = e.zeros(2 * n_embd)?;
4465        e.copy_into(&mut concat, 0, &e_norm, n_embd)?;
4466        e.copy_into(&mut concat, n_embd, &h_norm, n_embd)?;
4467
4468        // op 4: inpSA = eh_proj @ concat  (eh_proj [2*n_embd, n_embd]) -> [n_embd]
4469        let inp_sa = e.matmul(&mtp.eh_proj, &concat, 1)?;
4470
4471        // op 5: a_norm = RMSNorm(inpSA, attn_norm)
4472        let mut a_norm = e.zeros(di)?;
4473        e.rms_norm(&inp_sa, mtp.attn_norm.float_data(), &mut a_norm, di, 1, eps)?;
4474        anat_mark(0, e, &mut t_ph)?;
4475
4476        // op 6: attention on the scratch KV. SAME dc launcher as the graph path (bucket_max =
4477        // scratch.cap, length from the device len_d) so eager drafts match graph drafts
4478        // bit-for-bit at any t_kv (the parity gate). Host len mirrored here (the dc append
4479        // advances only the device counter).
4480        let attn_out = match (&mtp.mixer, mtp.step35.as_ref()) {
4481            // step35 MTP block, dcw door armed: the SAME windowed device-counter launcher as
4482            // the captured chain (draft parity by construction). Per-step ring headroom runs
4483            // HERE (eager is host-len work, a rebase is legal); host len mirrored like the
4484            // plain dc arm below.
4485            (Mixer::Full(fa), Some(g))
4486                if self.step35_dcw_eligible(g, scratch.plane(scratch_index).1) =>
4487            {
4488                {
4489                    let (kv, _) = scratch.plane_mut(scratch_index);
4490                    let retain = match kv.ring.as_ref() {
4491                        Some(ring) => memra_kv::swa_retain_from(kv.len, ring.window(), ring.base()),
4492                        None => 0,
4493                    };
4494                    e.prepare_kv_append(kv, retain, 1)?;
4495                }
4496                let out =
4497                    self.mtp_step35_attn_dcw(e, fa, g, &a_norm, &pos_d, scratch, scratch_index)?;
4498                scratch.plane_mut(scratch_index).0.len += 1;
4499                out
4500            }
4501            // step35 MTP block, door off (MEMRA_STEP35_DRAFT_DCW=0 rollback) or class-
4502            // ineligible: PER-LAYER geometry + a separate head-wise gate + an SWA window,
4503            // none of which the plain dc launcher can express (see `mtp_step35_attn`).
4504            // Host-len arm. Advances BOTH the
4505            // host len and the device counter itself (unlike the dc arm, whose host-side
4506            // mirror the caller does).
4507            (Mixer::Full(fa), Some(g)) => {
4508                self.mtp_step35_attn(e, fa, g, &a_norm, &pos_d, scratch, scratch_index)?
4509            }
4510            (Mixer::Full(fa), None) => {
4511                let out = self.mtp_full_attn_dc(
4512                    e,
4513                    fa,
4514                    &a_norm,
4515                    &pos_d,
4516                    scratch,
4517                    scratch_index,
4518                    mtp.geom.as_ref(),
4519                )?;
4520                scratch.plane_mut(scratch_index).0.len += 1;
4521                out
4522            }
4523            (Mixer::Linear(_), _) => {
4524                panic!("MTP block is full-attn in qwen35; linear MTP not supported")
4525            }
4526            (Mixer::Mla(_), _) => crate::hybrid::mla_forward_unimplemented(),
4527        };
4528        anat_mark(1, e, &mut t_ph)?;
4529
4530        // op 7: x1 = inpSA + attn_out
4531        let mut x1 = e.zeros(di)?;
4532        e.add(&inp_sa, &attn_out, &mut x1, di)?;
4533
4534        // op 8: z = RMSNorm(x1, post_attn_norm)  (pre-FFN norm)
4535        let mut z = e.zeros(di)?;
4536        e.rms_norm(&x1, mtp.post_attn_norm.float_data(), &mut z, di, 1, eps)?;
4537
4538        // op 9: FFN (Dense or MoE) — same as the trunk decode FFN
4539        let ffn_out = match &mtp.ffn {
4540            crate::hybrid::Ffn::Dense {
4541                ffn_gate,
4542                ffn_up,
4543                ffn_down,
4544            } => {
4545                let n_ff = ffn_gate.out_features();
4546                let (gate, up) = if e.uses_q8_1_fast(ffn_gate) && e.uses_q8_1_fast(ffn_up) {
4547                    let (zq, zd) = e.quantize_q8_1(&z, 1, di)?;
4548                    (
4549                        e.matmul_pre(ffn_gate, &zq, &zd, &z, 1)?,
4550                        e.matmul_pre(ffn_up, &zq, &zd, &z, 1)?,
4551                    )
4552                } else {
4553                    (e.matmul(ffn_gate, &z, 1)?, e.matmul(ffn_up, &z, 1)?)
4554                };
4555                let mut act = e.zeros(n_ff)?;
4556                // step35: a DENSE FFN reads the per-layer SHEXP clamp (upstream's one `build_ffn`
4557                // serves the dense MLP and the shared expert off `swiglu_clamp_shexp` —
4558                // llama-graph.cpp:1751), resolved for the MTP block's OWN index. Every other arch
4559                // passes None, which is `ffn_act`'s dispatch verbatim.
4560                Self::ffn_act_lim(
4561                    e,
4562                    &self.cfg,
4563                    &gate,
4564                    &up,
4565                    1.0,
4566                    1.0,
4567                    mtp.step35.as_ref().and_then(|s| s.clamp_shexp),
4568                    &mut act,
4569                    n_ff,
4570                )?;
4571                e.matmul(ffn_down, &act, 1)?
4572            }
4573            // MTP head is a distinct block — key its experts under a separate layer index (u16::MAX)
4574            // so they never alias trunk layer 0's cache keys.
4575            crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il(e, m, &z, 1, u16::MAX)?,
4576        };
4577        anat_mark(2, e, &mut t_ph)?;
4578
4579        // op 10: h_nextn = x1 + ffn_out (at di)
4580        let mut h_inner = e.zeros(di)?;
4581        e.add(&x1, &ffn_out, &mut h_inner, di)?;
4582
4583        // op 10.5 (student): up-project the inner hidden back to n_embd — training semantics:
4584        // the chain carrier AND the head input are out_up(h_inner) (pre-final-norm).
4585        let h_nextn = match mtp.geom.as_ref() {
4586            Some(g) => e.matmul(&g.out_up, &h_inner, 1)?,
4587            None => h_inner,
4588        };
4589
4590        // op 11: final = RMSNorm(h_nextn, shared_head_norm OR output_norm)
4591        let final_norm = mtp.shared_head_norm.as_ref().unwrap_or(&self.output_norm);
4592        let mut final_h = e.zeros(n_embd)?;
4593        e.rms_norm(
4594            &h_nextn,
4595            final_norm.float_data(),
4596            &mut final_h,
4597            n_embd,
4598            1,
4599            eps,
4600        )?;
4601
4602        // op 12: draft_logits = (shared_head_head OR output) @ final — stays ON DEVICE.
4603        let head = mtp.shared_head_head.as_ref().unwrap_or(&self.output);
4604        let mut logits = e.matmul(head, &final_h, 1)?;
4605        // op 12b (lane/draft-mask): grammar mask over the DRAFT vocab, applied here so the
4606        // caller's argmax / gumbel draw / p-min prob all read the grammar-legal row.
4607        if let Some((mask_d, mw)) = mask {
4608            let d_vocab = head.out_features();
4609            e.mask_logits_col(&mut logits, mask_d, 0, d_vocab, mw)?;
4610        }
4611        anat_mark(3, e, &mut t_ph)?;
4612        if anat {
4613            ANAT_NS[4].fetch_add(t_all.elapsed().as_nanos() as u64, Relaxed);
4614            let n = ANAT_STEPS.fetch_add(1, Relaxed) + 1;
4615            if n % 128 == 0 {
4616                let us = |i: usize| ANAT_NS[i].load(Relaxed) / n / 1000;
4617                eprintln!(
4618                    "[spec-anatomy] steps={n} avg us/step: glue={} attn={} ffn={} head={} total={}",
4619                    us(0),
4620                    us(1),
4621                    us(2),
4622                    us(3),
4623                    us(4)
4624                );
4625            }
4626        }
4627        // Chain recurrence hand-over: pre-norm h_nextn (default) or post-norm final_h
4628        // (MEMRA_SPEC_HPOST — llama.cpp #24025's t_h_nextn is taken AFTER the head norm).
4629        Ok((logits, if spec_hpost() { final_h } else { h_nextn }))
4630    }
4631
4632    #[allow(clippy::too_many_arguments)]
4633    fn mtp_chain_forward_dev(
4634        &self,
4635        e: &Engine,
4636        tokens: &[u32],
4637        seeds: &[CudaSlice<f32>],
4638        scratch: &mut MtpScratch,
4639        committed_scratch_len: usize,
4640        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
4641        mask: Option<(&CudaSlice<u32>, usize)>,
4642    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4643        if tokens.is_empty() || tokens.len() != seeds.len() {
4644            return Err("multi-head MTP prefix tokens/seeds are malformed".into());
4645        }
4646        let index = mtp_chain_head_index(tokens.len() - 1, self.mtp_head_count());
4647        let head = self.mtp_head_at(index);
4648        scratch.set_plane_len(e, index, committed_scratch_len)?;
4649
4650        let mut last = None;
4651        for row in 0..tokens.len() {
4652            let is_last = row + 1 == tokens.len();
4653            last = Some(self.mtp_head_forward_dev_at(
4654                e,
4655                head,
4656                tokens[row],
4657                &seeds[row],
4658                scratch,
4659                index,
4660                committed_scratch_len + row + 1,
4661                embd_dev,
4662                if is_last { mask } else { None },
4663            )?);
4664        }
4665        Ok(last.expect("non-empty MTP prefix produced no row"))
4666    }
4667
4668    /// step35 MTP-block attention, T=1, on the scratch KV — the EAGER-ONLY twin of
4669    /// `mtp_full_attn_dc`. Three things force a separate arm rather than a geometry parameter on
4670    /// the dc path, and all three are properties of this arch's MTP block:
4671    ///
4672    /// 1. **The SWA window.** Block 45 is an SWA-type block (`sliding_window_pattern[45]=true`,
4673    ///    window 512). Windowed decode in memra is a token-aligned VIEW OFFSET into the quantized
4674    ///    cache (the gemma4 R6 / `step35_decode_attn` pattern: keys carry absolute rope and the
4675    ///    mask is purely positional, so one query at `len-1` attending the last `win` rows IS the
4676    ///    windowed result). `fa_decode_dc` takes the key count from a DEVICE counter and always
4677    ///    starts at row 0 — it cannot express a nonzero offset. The windowed dc arm is
4678    ///    `mtp_step35_attn_dcw` (`fa_decode_dcw`, doored via MEMRA_STEP35_DRAFT_DCW —
4679    ///    default ON since lane/step37-draft-graph-serving-20260830); this host-len arm is
4680    ///    the =0 rollback and the class-ineligibility fallback.
4681    /// 2. **Per-layer head count.** 96 q heads over 8 KV (GQA 12) at this block, vs the trunk's 64
4682    ///    on its full-attn layers. The trunk cfg's `n_head` scalar is the MAX over layers, and the
4683    ///    trunk ARTIFACT's per-layer arrays stop at index 44 — so the count must come from the
4684    ///    resolved `Step35MtpGeom`, never from `cfg`.
4685    /// 3. **The separate head-wise gate.** `blk.45.attn_gate.weight [n_embd, 96]` produces one
4686    ///    sigmoid scalar per head (broadcast over head_dim) — `attn_head_gate`, not the qwen35
4687    ///    fused-into-wq `q_gate_split` form the dc arm handles.
4688    ///
4689    /// DOOR STATE: with MEMRA_STEP35_DRAFT_DCW=0 (or a sub-eligible kernel class),
4690    /// `mtp_head_forward_cap` refuses step35 heads explicitly (rather than silently capturing
4691    /// a window-less, wrong-past-`win` graph) and this eager chain IS the served path. With
4692    /// the door armed (the default), BOTH draft modes run the `mtp_step35_attn_dcw` twin
4693    /// instead of this arm.
4694    ///
4695    /// Unlike the dc arm this advances BOTH the host `kv.len` and the device counter, so the
4696    /// caller must not mirror.
4697    fn mtp_step35_attn(
4698        &self,
4699        e: &Engine,
4700        fa: &FullAttnLayer,
4701        g: &crate::hybrid::Step35MtpGeom,
4702        h: &CudaSlice<f32>,
4703        pos_d: &CudaSlice<i32>,
4704        scratch: &mut MtpScratch,
4705        scratch_index: usize,
4706    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4707        let (nh, nkv, hd) = (g.n_head, g.n_head_kv, self.cfg.head_dim_k as usize);
4708        // MTP-GEOM RECEIPT, once per process, on the SERVED draft path. Slot-0 acceptance is
4709        // 0.725 here against 0.994 for vLLM MTP3 on the same checkpoint family and card class, and
4710        // the first three explanations for that gap were all wrong: head assignment (step-modulo
4711        // is index 0 at K=1, correct), MEMRA_SPEC_HPOST (identical 84/116 both arms), and this
4712        // block's geometry. Geometry was the one that could have failed SILENTLY — a wrong window
4713        // makes the draft attend the whole context instead of Step-3.7's 512, stays fluent, and
4714        // shows up only as acceptance — so it gets a standing receipt rather than another reading
4715        // of the source. Prints the resolved Step35MtpGeom the served path actually runs on;
4716        // `full_attention_geometry_at`'s missing-row fallback (window: None) does NOT reach here.
4717        {
4718            static ONCE: std::sync::OnceLock<()> = std::sync::OnceLock::new();
4719            ONCE.get_or_init(|| {
4720                eprintln!(
4721                    "[mtp-geom] arm=eager block={} swa={} window={} n_head={nh} n_head_kv={nkv} \
4722                     head_dim_k={hd} n_rot={} rope_base={} clamp_shexp={:?}",
4723                    g.il, g.swa, g.window, g.n_rot, g.rope_base, g.clamp_shexp,
4724                );
4725            });
4726        }
4727        let eps = self.cfg.rms_eps;
4728        let scale = 1.0 / (hd as f32).sqrt(); // step35.cpp:255 kq_scale
4729        let n_embd = self.cfg.n_embd as usize;
4730        let gw = fa
4731            .attn_gate
4732            .as_ref()
4733            .ok_or("step35 MTP block is missing attn_gate.weight (head-wise attention gate)")?;
4734
4735        let (q0, k0, v0, gt) = if e.uses_q8_1_fast(&fa.wq)
4736            && e.uses_q8_1_fast(&fa.wk)
4737            && e.uses_q8_1_fast(&fa.wv)
4738            && e.uses_q8_1_fast(gw)
4739        {
4740            let (hq, hdq) = e.quantize_q8_1(h, 1, n_embd)?;
4741            let (a, b, c) = match e.matmul_q8_fused3(&fa.wq, &fa.wk, &fa.wv, &hq, &hdq)? {
4742                Some(t3) => t3,
4743                None => (
4744                    e.matmul_pre(&fa.wq, &hq, &hdq, h, 1)?,
4745                    e.matmul_pre(&fa.wk, &hq, &hdq, h, 1)?,
4746                    e.matmul_pre(&fa.wv, &hq, &hdq, h, 1)?,
4747                ),
4748            };
4749            (a, b, c, e.matmul_pre(gw, &hq, &hdq, h, 1)?)
4750        } else {
4751            (
4752                e.matmul(&fa.wq, h, 1)?,
4753                e.matmul(&fa.wk, h, 1)?,
4754                e.matmul(&fa.wv, h, 1)?,
4755                e.matmul(gw, h, 1)?,
4756            )
4757        };
4758
4759        let mut q = e.uninit(nh * hd)?;
4760        e.rms_norm(&q0, fa.q_norm.float_data(), &mut q, hd, nh, eps)?;
4761        let mut k = e.uninit(nkv * hd)?;
4762        e.rms_norm(&k0, fa.k_norm.float_data(), &mut k, hd, nkv, eps)?;
4763        // `rope_freqs.weight` (llama3 factors) applies to the FULL-attn layers ONLY; SWA passes
4764        // null (llama-hparams / step35.cpp). Block 45 is SWA, so `ff` is None there — but read
4765        // the resolved flag, not the constant, so an all-full sibling stays correct.
4766        let ff = if g.swa {
4767            None
4768        } else {
4769            self.step35_aux.as_ref().and_then(|a| a.rope_freqs(e))
4770        };
4771        #[cfg(debug_assertions)]
4772        if let Some(ff) = ff {
4773            crate::debug_assert_tensor_stream_device(ff, &e.stream(), "mtp_step35_attn.rope_freqs");
4774        }
4775        e.rope_neox2(
4776            &mut q,
4777            &mut k,
4778            pos_d,
4779            hd,
4780            g.n_rot,
4781            nh,
4782            nkv,
4783            1,
4784            g.rope_base,
4785            1.0,
4786            ff,
4787        )?;
4788
4789        // Append at the HOST slot, then re-stamp the device counter: the eager chain has the
4790        // length on the host anyway, and the windowed view below needs it there to compute the
4791        // offset. The device counter is kept in lockstep so `mtp_kv_fill`'s `set_i32_one` and any
4792        // dc-family consumer of this scratch still agree.
4793        let (kv, scratch_cap) = scratch.plane_mut(scratch_index);
4794        assert!(
4795            kv.len < scratch_cap,
4796            "step35 MTP scratch overflow ({} >= {})",
4797            kv.len,
4798            scratch_cap
4799        );
4800        let next_len = kv.len + 1;
4801        let (off, t_kv) = if g.swa && next_len > g.window {
4802            (next_len - g.window, g.window)
4803        } else {
4804            (0, next_len)
4805        };
4806        // `off`/`t_kv` stay the ATTENTION view; the retain is a separate, lower bound so the
4807        // rewind that follows this append is still resident. THIS is the only site that rebases
4808        // this plane (MEMRA_KV_REBASE_TRACE, one run: 1 rebase, all from here), so it is the site
4809        // that decides `base` for everyone.
4810        let retain_from = match kv.ring.as_ref() {
4811            Some(ring) => memra_kv::swa_retain_from(kv.len, ring.window(), ring.base()),
4812            None => off & !31usize,
4813        };
4814        let write_row = e.prepare_kv_append(kv, retain_from, 1)?;
4815        e.append_kv_quantized(
4816            &k,
4817            &v0,
4818            &mut kv.k,
4819            &mut kv.v,
4820            write_row,
4821            kv.kv_dim_k,
4822            kv.kv_dim_v,
4823            kv.k_tok_bytes,
4824            kv.v_tok_bytes,
4825            false,
4826        )?;
4827        kv.len = next_len;
4828        e.set_i32_one(&mut kv.len_d, kv.len as i32)?;
4829        // SWA view offset (see note 1). The draft chain is short (k+2 rows), but the scratch is
4830        // PERSISTENT across rounds — `mtp_kv_fill` leaves one row per committed token behind, so
4831        // `kv.len` tracks absolute position and crosses 512 in any real generation. The window is
4832        // therefore live, not theoretical.
4833        let physical = kv.physical_rows(off, off + t_kv)?;
4834        let k_view = e.view_u8_range(
4835            &kv.k,
4836            physical.start * kv.k_tok_bytes,
4837            physical.end * kv.k_tok_bytes,
4838        );
4839        let v_view = e.view_u8_range(
4840            &kv.v,
4841            physical.start * kv.v_tok_bytes,
4842            physical.end * kv.v_tok_bytes,
4843        );
4844        let mut attn = e.uninit(nh * hd)?;
4845        e.fa_decode_kvmod(
4846            &q,
4847            &k_view,
4848            &v_view,
4849            &mut attn,
4850            hd,
4851            nh,
4852            nkv,
4853            t_kv,
4854            scale,
4855            kv.k_tok_bytes,
4856            kv.v_tok_bytes,
4857            false,
4858        )?;
4859
4860        let mut ag = e.uninit(nh * hd)?;
4861        e.attn_head_gate(&attn, &gt, &mut ag, None, hd, nh, 1)?;
4862        Ok(e.matmul(&fa.wo, &ag, 1)?)
4863    }
4864
4865    /// The dcw draft arm's kernel-class precondition, mirrored from `fa_decode_dcw`'s own
4866    /// refusal plus the v3 walk's format contract (`fa_v3_active`), so the DEV dispatch can
4867    /// never pick an arm the launcher would refuse mid-chain (the eager chain has no graceful
4868    /// fallback point) and the CAP site refuses with the named reason instead.
4869    ///
4870    /// `cap` = the SESSION's scratch-plane row capacity: the launcher's vec gate reads
4871    /// `bucket_max = min(window, cap)`, so a SMALL session (tiny prompt + tiny max_tokens,
4872    /// e.g. a max_tokens=8 probe: cap ~62 < the 96 vec floor) is OUTSIDE the dcw domain even
4873    /// though the WINDOW clears the floor. Mirroring the window alone shipped exactly that
4874    /// hole when the door default flipped ON (2026-08-30, vision-cell receipt: sampled
4875    /// capture WARN + `[engine-error] fa_decode_dcw supports the default v3-vec class only`
4876    /// hard-failing the burst — the eager dcw arm has no graceful fallback point). Sub-floor
4877    /// sessions now take the host-len kvmod arm, byte-for-byte the door-off serving.
4878    fn step35_dcw_eligible(&self, g: &crate::hybrid::Step35MtpGeom, cap: usize) -> bool {
4879        let hd = self.cfg.head_dim_k as usize;
4880        step35_draft_dcw_on()
4881            && g.swa
4882            && g.window.min(cap) >= crate::fa_vec_min_tkv()
4883            && std::env::var("MEMRA_NO_FA_VEC").is_err()
4884            && crate::fa_v3_active(hd)
4885            && hd <= 256
4886            && hd % 32 == 0
4887    }
4888
4889    /// step35 MTP-block attention, T=1, on the scratch KV: the WINDOWED DEVICE-COUNTER twin
4890    /// of `mtp_step35_attn`, serving BOTH draft paths when `step35_draft_dcw_on`. Write slot,
4891    /// key bound and SWA view offset all derive from device state (`len_d`, `base_d` written
4892    /// only at host-side rebases, and the block's `window`), so ONE captured graph serves the
4893    /// whole chain and replays see KV growth through the counter: the `mtp_full_attn_dc`
4894    /// contract plus the view offset the plain `_dc` kernel could not express (the old
4895    /// capture-refusal root cause). The three step35 properties stay per-geom exactly as in
4896    /// the eager twin: nh/nkv from `Step35MtpGeom`, the separate head-wise gate
4897    /// (`attn_head_gate`), per-layer rope width/base with SWA passing null freqs.
4898    ///
4899    /// bucket_max = min(cap, window): the windowed view never exceeds `window` rows, so the
4900    /// capture-time grid stays valid for every replayed len, and the kernel derives ns_eff
4901    /// from the LIVE T_kv at the fixed split_keys (one-partition law). Both arms call THIS
4902    /// launcher at THIS bucket, so eager and captured drafts are bit-identical by
4903    /// construction; vs the retired-by-flag `mtp_step35_attn` the only numeric-class deltas
4904    /// are the sub-vec-floor region (t_kv < 96: kvmod ran scalar, dcw stays vec) and any
4905    /// live-len split-ladder rung below the bucket's, both draft-side only (the verify
4906    /// arbitrates emitted bytes; acceptance is gated by the battery).
4907    ///
4908    /// Host len is NOT advanced here (graph contract); callers mirror. The EAGER caller runs
4909    /// `prepare_kv_append` per step (ring headroom, rebase legal there); the CAPTURED path
4910    /// pre-arms headroom at capture time and round start (`MtpScratch::ensure_dcw_headroom`)
4911    /// because a rebase is host work no captured chain may contain.
4912    fn mtp_step35_attn_dcw(
4913        &self,
4914        e: &Engine,
4915        fa: &FullAttnLayer,
4916        g: &crate::hybrid::Step35MtpGeom,
4917        h: &CudaSlice<f32>,
4918        pos_d: &CudaSlice<i32>,
4919        scratch: &mut MtpScratch,
4920        scratch_index: usize,
4921    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4922        let (nh, nkv, hd) = (g.n_head, g.n_head_kv, self.cfg.head_dim_k as usize);
4923        // MTP-GEOM RECEIPT (dcw twin of the `mtp_step35_attn` receipt): once per process,
4924        // naming the arm, so a serving log proves WHICH draft attention program ran (the
4925        // engagement receipt for the flag door, both directions).
4926        {
4927            static ONCE: std::sync::OnceLock<()> = std::sync::OnceLock::new();
4928            ONCE.get_or_init(|| {
4929                eprintln!(
4930                    "[mtp-geom] arm=dcw block={} swa={} window={} n_head={nh} n_head_kv={nkv} \
4931                     head_dim_k={hd} n_rot={} rope_base={} clamp_shexp={:?}",
4932                    g.il, g.swa, g.window, g.n_rot, g.rope_base, g.clamp_shexp,
4933                );
4934            });
4935        }
4936        let eps = self.cfg.rms_eps;
4937        let scale = 1.0 / (hd as f32).sqrt(); // step35.cpp:255 kq_scale
4938        let n_embd = self.cfg.n_embd as usize;
4939        let gw = fa
4940            .attn_gate
4941            .as_ref()
4942            .ok_or("step35 MTP block is missing attn_gate.weight (head-wise attention gate)")?;
4943
4944        let (q0, k0, v0, gt) = if e.uses_q8_1_fast(&fa.wq)
4945            && e.uses_q8_1_fast(&fa.wk)
4946            && e.uses_q8_1_fast(&fa.wv)
4947            && e.uses_q8_1_fast(gw)
4948        {
4949            let (hq, hdq) = e.quantize_q8_1(h, 1, n_embd)?;
4950            let (a, b, c) = match e.matmul_q8_fused3(&fa.wq, &fa.wk, &fa.wv, &hq, &hdq)? {
4951                Some(t3) => t3,
4952                None => (
4953                    e.matmul_pre(&fa.wq, &hq, &hdq, h, 1)?,
4954                    e.matmul_pre(&fa.wk, &hq, &hdq, h, 1)?,
4955                    e.matmul_pre(&fa.wv, &hq, &hdq, h, 1)?,
4956                ),
4957            };
4958            (a, b, c, e.matmul_pre(gw, &hq, &hdq, h, 1)?)
4959        } else {
4960            (
4961                e.matmul(&fa.wq, h, 1)?,
4962                e.matmul(&fa.wk, h, 1)?,
4963                e.matmul(&fa.wv, h, 1)?,
4964                e.matmul(gw, h, 1)?,
4965            )
4966        };
4967
4968        let mut q = e.zeros(nh * hd)?;
4969        e.rms_norm(&q0, fa.q_norm.float_data(), &mut q, hd, nh, eps)?;
4970        let mut k = e.zeros(nkv * hd)?;
4971        e.rms_norm(&k0, fa.k_norm.float_data(), &mut k, hd, nkv, eps)?;
4972        // rope_freqs (llama3 factors) apply to the FULL-attn layers ONLY; SWA passes null
4973        // (the eager twin's rule, resolved from the flag, not the constant).
4974        let ff = if g.swa {
4975            None
4976        } else {
4977            self.step35_aux.as_ref().and_then(|a| a.rope_freqs(e))
4978        };
4979        #[cfg(debug_assertions)]
4980        if let Some(ff) = ff {
4981            crate::debug_assert_tensor_stream_device(
4982                ff,
4983                &e.stream(),
4984                "mtp_step35_attn_dcw.rope_freqs",
4985            );
4986        }
4987        e.rope_neox2(
4988            &mut q,
4989            &mut k,
4990            pos_d,
4991            hd,
4992            g.n_rot,
4993            nh,
4994            nkv,
4995            1,
4996            g.rope_base,
4997            1.0,
4998            ff,
4999        )?;
5000
5001        let (kv, cap) = scratch.plane_mut(scratch_index);
5002        // Append at the DEVICE slot's PHYSICAL row (len_d - base_d), then advance the counter
5003        // in-graph. Physical room is the callers' headroom contract (see the fn doc).
5004        e.append_kv_quantized_dcw(
5005            &k,
5006            &v0,
5007            &mut kv.k,
5008            &mut kv.v,
5009            &kv.len_d,
5010            kv.base_d.as_ref(),
5011            kv.kv_dim_k,
5012            kv.kv_dim_v,
5013            kv.k_tok_bytes,
5014            kv.v_tok_bytes,
5015        )?;
5016        e.inc_seqlen(&mut kv.len_d)?;
5017        // Full-buffer views (any in-round physical row stays in range under the headroom
5018        // contract); the kernel bounds and offsets the key range from (len_d, base_d, window).
5019        let k_view = e.view_u8(&kv.k, kv.k.len());
5020        let v_view = e.view_u8(&kv.v, kv.v.len());
5021        let bucket = g.window.min(cap);
5022        let mut attn = e.zeros(nh * hd)?;
5023        e.fa_decode_dcw(
5024            &q,
5025            &k_view,
5026            &v_view,
5027            &mut attn,
5028            hd,
5029            nh,
5030            nkv,
5031            &kv.len_d,
5032            kv.base_d.as_ref(),
5033            if g.swa { g.window } else { 0 },
5034            bucket,
5035            scale,
5036            kv.k_tok_bytes,
5037            kv.v_tok_bytes,
5038            None,
5039        )?;
5040
5041        let mut ag = e.zeros(nh * hd)?;
5042        e.attn_head_gate(&attn, &gt, &mut ag, None, hd, nh, 1)?;
5043        Ok(e.matmul(&fa.wo, &ag, 1)?)
5044    }
5045
5046    /// MTP-block full attention, T=1, on the scratch KV (BOTH draft paths — eager and graph):
5047    /// the scratch write slot and the attention bound come from `scratch.kv.len_d` (device i32[1])
5048    /// so the launch args are FIXED across draft steps — ONE captured graph serves the whole
5049    /// chain, and replays keep seeing KV growth through the device counter (no recapture).
5050    /// Geometry contract: n_splits is sized from `scratch.cap` (the persistent capacity); splits
5051    /// whose key range lies beyond the device t_kv exit empty and the shared combine skips them
5052    /// (fa_decode_dc bit-correct-for-any-t_kv<=bucket_max contract). The eager path uses the SAME
5053    /// launcher with the SAME bucket_max -> identical dispatch -> bit-identical draft tokens (the
5054    /// graph-vs-eager parity gate). Host len is NOT advanced here (graph contract); callers mirror.
5055    fn mtp_full_attn_dc(
5056        &self,
5057        e: &Engine,
5058        fa: &FullAttnLayer,
5059        h: &CudaSlice<f32>,
5060        pos_d: &CudaSlice<i32>,
5061        scratch: &mut MtpScratch,
5062        scratch_index: usize,
5063        geom: Option<&crate::hybrid::DraftGeom>,
5064    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5065        let cfg = &self.cfg;
5066        let mtp_il = cfg.n_layer.saturating_sub(cfg.nextn_predict_layers);
5067        let geometry = cfg.full_attention_geometry_at(mtp_il);
5068        let n_head = geom.map(|g| g.n_head).unwrap_or(geometry.n_head as usize);
5069        let n_head_kv = geom
5070            .map(|g| g.n_head_kv)
5071            .unwrap_or(geometry.n_head_kv as usize);
5072        let head_dim = geometry.head_dim_k as usize;
5073        let eps = cfg.rms_eps;
5074        let scale = geometry.attention_scale();
5075        let n_embd = geom.map(|g| g.d_inner).unwrap_or(cfg.n_embd as usize);
5076        let bucket_max = scratch.plane(scratch_index).1;
5077
5078        let (qf, mut k, v) =
5079            if e.uses_q8_1_fast(&fa.wq) && e.uses_q8_1_fast(&fa.wk) && e.uses_q8_1_fast(&fa.wv) {
5080                let (hq, hd) = e.quantize_q8_1(h, 1, n_embd)?;
5081                (
5082                    e.matmul_pre(&fa.wq, &hq, &hd, h, 1)?,
5083                    e.matmul_pre(&fa.wk, &hq, &hd, h, 1)?,
5084                    e.matmul_pre(&fa.wv, &hq, &hd, h, 1)?,
5085                )
5086            } else {
5087                (
5088                    e.matmul(&fa.wq, h, 1)?,
5089                    e.matmul(&fa.wk, h, 1)?,
5090                    e.matmul(&fa.wv, h, 1)?,
5091                )
5092            };
5093        // M3/Hy3 have no attention output gate — wq out is exactly q; skip the split.
5094        let gated = geometry.attention_gate == memra_gguf::config::AttentionGateKind::FusedQ;
5095        let (mut q, gate) = if gated {
5096            let mut q = e.zeros(n_head * head_dim)?;
5097            let mut gate = e.zeros(n_head * head_dim)?;
5098            e.q_gate_split(&qf, &mut q, &mut gate, head_dim, n_head, 1)?;
5099            (q, Some(gate))
5100        } else {
5101            (qf, None)
5102        };
5103
5104        let mut qn = e.zeros(n_head * head_dim)?;
5105        e.rms_norm(&q, fa.q_norm.float_data(), &mut qn, head_dim, n_head, eps)?;
5106        q = qn;
5107        let mut kn = e.zeros(n_head_kv * head_dim)?;
5108        e.rms_norm(
5109            &k,
5110            fa.k_norm.float_data(),
5111            &mut kn,
5112            head_dim,
5113            n_head_kv,
5114            eps,
5115        )?;
5116        k = kn;
5117        let rope_dims = geometry.n_rot as usize;
5118        e.rope_neox(
5119            &mut q,
5120            pos_d,
5121            head_dim,
5122            rope_dims,
5123            n_head,
5124            1,
5125            geometry.rope_base,
5126            1.0,
5127        )?;
5128        e.rope_neox(
5129            &mut k,
5130            pos_d,
5131            head_dim,
5132            rope_dims,
5133            n_head_kv,
5134            1,
5135            geometry.rope_base,
5136            1.0,
5137        )?;
5138
5139        let kv = scratch.plane_mut(scratch_index).0;
5140        // append at the DEVICE slot (kv.len_d == old len), then advance the counter in-graph.
5141        e.append_kv_quantized_dc(
5142            &k,
5143            &v,
5144            &mut kv.k,
5145            &mut kv.v,
5146            &kv.len_d,
5147            kv.kv_dim_k,
5148            kv.kv_dim_v,
5149            kv.k_tok_bytes,
5150            kv.v_tok_bytes,
5151            false,
5152        )?;
5153        e.inc_seqlen(&mut kv.len_d)?;
5154        // full-buffer views (any in-round t_kv stays in range on replay); the kernel bounds the
5155        // key range from the device counter.
5156        let k_view = e.view_u8(&kv.k, kv.k.len());
5157        let v_view = e.view_u8(&kv.v, kv.v.len());
5158        let (ktb, vtb) = (kv.k_tok_bytes, kv.v_tok_bytes);
5159        let mut attn = e.zeros(n_head * head_dim)?;
5160        e.fa_decode_dc(
5161            &q, &k_view, &v_view, &mut attn, head_dim, n_head, n_head_kv, &kv.len_d, bucket_max,
5162            scale, ktb, vtb, false,
5163        )?;
5164
5165        let attn_g = match &gate {
5166            Some(gate) => {
5167                let mut gsig = e.zeros(n_head * head_dim)?;
5168                e.sigmoid(gate, &mut gsig, n_head * head_dim)?;
5169                let mut ag = e.zeros(n_head * head_dim)?;
5170                e.mul(&attn, &gsig, &mut ag, n_head * head_dim)?;
5171                ag
5172            }
5173            None => attn,
5174        };
5175        Ok(e.matmul(&fa.wo, &attn_g, 1)?)
5176    }
5177
5178    /// PERSISTENT-DRAFT-KV fill (the reference engine's "mtp_update" analogue): compute the MTP
5179    /// block's K/V for `tokens` (committed tokens at positions pos0..pos0+T) from their EXACT
5180    /// trunk hiddens `h` ([T, n_embd] token-major, pre-output_norm) and append at slots pos0.. of
5181    /// the scratch KV. K/V-ONLY — ops A/1-5 plus the K-side of op 6 (wk/wv + k_norm + rope +
5182    /// quantized append); no wq/attention/FFN/lm_head, so per-token cost ~= eh_proj + wk/wv (a
5183    /// small fraction of one trunk layer), T-batched. Rope follows the chain convention
5184    /// rope(token@p) = p+1. Runs at round boundaries OUTSIDE the captured graph in BOTH draft
5185    /// modes -> draft parity by construction. Caller must have scratch.kv.len == pos0.
5186    #[allow(clippy::too_many_arguments)]
5187    fn mtp_kv_fill_at(
5188        &self,
5189        e: &Engine,
5190        mtp: &MtpHead,
5191        tokens: &[u32],
5192        h: &CudaSlice<f32>,
5193        pos0: usize,
5194        scratch: &mut MtpScratch,
5195        scratch_index: usize,
5196        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
5197    ) -> Result<(), Box<dyn std::error::Error>> {
5198        let cfg = &self.cfg;
5199        let n_embd = cfg.n_embd as usize;
5200        let eps = cfg.rms_eps;
5201        let t = tokens.len();
5202        let (scratch_kv, scratch_cap) = scratch.plane(scratch_index);
5203        assert_eq!(scratch_kv.len, pos0, "mtp_kv_fill: append slot mismatch");
5204        assert!(pos0 + t <= scratch_cap, "mtp_kv_fill: scratch overflow");
5205        let Mixer::Full(fa) = &mtp.mixer else {
5206            panic!("MTP block is full-attn in qwen35; linear MTP not supported")
5207        };
5208        let pos_vec: Vec<i32> = (0..t).map(|i| (pos0 + i + 1) as i32).collect();
5209        let pos_d = e.htod_i32(&pos_vec)?;
5210
5211        // ops A/1/2: embed + the two input norms, T-wide.
5212        let e_emb = match embd_dev {
5213            Some((g, qt, rb)) => e.embed_gather_device_t(g, tokens, n_embd, qt, rb)?,
5214            None => e.htod(&self.embd.gather(n_embd, tokens))?,
5215        };
5216        let mut e_norm = e.zeros(t * n_embd)?;
5217        e.rms_norm(&e_emb, mtp.enorm.float_data(), &mut e_norm, n_embd, t, eps)?;
5218        let mut h_norm = e.zeros(t * n_embd)?;
5219        e.rms_norm(h, mtp.hnorm.float_data(), &mut h_norm, n_embd, t, eps)?;
5220
5221        // op 3: per-row [e_norm ; h_norm] concat, token-major [T, 2*n_embd].
5222        let mut concat = e.zeros(t * 2 * n_embd)?;
5223        for i in 0..t {
5224            e.copy_view_into(
5225                &mut concat,
5226                i * 2 * n_embd,
5227                &e_norm.slice(i * n_embd..(i + 1) * n_embd),
5228                n_embd,
5229            )?;
5230            e.copy_view_into(
5231                &mut concat,
5232                i * 2 * n_embd + n_embd,
5233                &h_norm.slice(i * n_embd..(i + 1) * n_embd),
5234                n_embd,
5235            )?;
5236        }
5237
5238        // ops 4/5: eh_proj + attn_norm, T-wide (at the student inner width when geom is set).
5239        let di = mtp.geom.as_ref().map(|g| g.d_inner).unwrap_or(n_embd);
5240        let inp_sa = e.matmul(&mtp.eh_proj, &concat, t)?;
5241        let mut a_norm = e.zeros(t * di)?;
5242        e.rms_norm(&inp_sa, mtp.attn_norm.float_data(), &mut a_norm, di, t, eps)?;
5243
5244        // op 6 (K/V half): wk/wv + k_norm + rope + per-row quantized append. No wq/attention —
5245        // the fill only has to leave correct K/V rows behind for later chains to attend over.
5246        let n_head_kv = mtp
5247            .geom
5248            .as_ref()
5249            .map(|g| g.n_head_kv)
5250            .or(mtp.step35.as_ref().map(|s| s.n_head_kv))
5251            .unwrap_or_else(|| {
5252                let mtp_il = cfg.n_layer.saturating_sub(cfg.nextn_predict_layers);
5253                cfg.full_attention_geometry_at(mtp_il).n_head_kv as usize
5254            });
5255        let mtp_il = cfg.n_layer.saturating_sub(cfg.nextn_predict_layers);
5256        let geometry = cfg.full_attention_geometry_at(mtp_il);
5257        let head_dim = geometry.head_dim_k as usize;
5258        let mut k = e.matmul(&fa.wk, &a_norm, t)?;
5259        let v = e.matmul(&fa.wv, &a_norm, t)?;
5260        let mut kn = e.zeros(t * n_head_kv * head_dim)?;
5261        e.rms_norm(
5262            &k,
5263            fa.k_norm.float_data(),
5264            &mut kn,
5265            head_dim,
5266            n_head_kv * t,
5267            eps,
5268        )?;
5269        k = kn;
5270        // step35: rotary width AND base are per-layer, and the MTP block's values come from the
5271        // resolved `Step35MtpGeom` — NOT from `cfg.rope_dim_count`/`cfg.rope_freq_base`, which
5272        // carry the arch defaults (128 / 5e6, i.e. the FULL-attn layers' base). Getting this wrong
5273        // writes K rows the attention arm then re-derives at a different theta: correct-looking
5274        // output with dead acceptance, invisible to the exactness gates.
5275        let (rope_dims, rope_base, ff) = match mtp.step35.as_ref() {
5276            Some(s) => (
5277                s.n_rot,
5278                s.rope_base,
5279                if s.swa {
5280                    None
5281                } else {
5282                    self.step35_aux.as_ref().and_then(|a| a.rope_freqs(e))
5283                },
5284            ),
5285            None => (geometry.n_rot as usize, geometry.rope_base, None),
5286        };
5287        #[cfg(debug_assertions)]
5288        if let Some(ff) = ff {
5289            crate::debug_assert_tensor_stream_device(ff, &e.stream(), "mtp_kv_fill.rope_freqs");
5290        }
5291        match ff {
5292            Some(f) => e.rope_neox_ff(
5293                &mut k, &pos_d, head_dim, rope_dims, n_head_kv, t, rope_base, 1.0, f,
5294            )?,
5295            None => e.rope_neox(
5296                &mut k, &pos_d, head_dim, rope_dims, n_head_kv, t, rope_base, 1.0,
5297            )?,
5298        }
5299
5300        let kv = scratch.plane_mut(scratch_index).0;
5301        // Match the trunk prime contract: a chunk may need the aligned window immediately before
5302        // its first row, so preserve that prefix when the physical tail rebases at wrap.
5303        let retain_from = kv
5304            .ring
5305            .as_ref()
5306            .map(|ring| memra_kv::swa_retain_from(pos0, ring.window(), ring.base()))
5307            .unwrap_or(0);
5308        let write_row = e.prepare_kv_append(kv, retain_from, t)?;
5309        for i in 0..t {
5310            let k_row = k.slice(i * kv.kv_dim_k..(i + 1) * kv.kv_dim_k);
5311            let v_row = v.slice(i * kv.kv_dim_v..(i + 1) * kv.kv_dim_v);
5312            e.append_kv_quantized_view(
5313                &k_row,
5314                &v_row,
5315                &mut kv.k,
5316                &mut kv.v,
5317                write_row + i,
5318                kv.kv_dim_k,
5319                kv.kv_dim_v,
5320                kv.k_tok_bytes,
5321                kv.v_tok_bytes,
5322                false,
5323            )?;
5324        }
5325        kv.len = pos0 + t;
5326        e.set_i32_one(&mut kv.len_d, kv.len as i32)?;
5327        Ok(())
5328    }
5329
5330    #[allow(clippy::too_many_arguments)]
5331    fn mtp_kv_fill_all(
5332        &self,
5333        e: &Engine,
5334        tokens: &[u32],
5335        h: &CudaSlice<f32>,
5336        pos0: usize,
5337        scratch: &mut MtpScratch,
5338        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
5339    ) -> Result<(), Box<dyn std::error::Error>> {
5340        debug_assert_eq!(self.mtp_head_count(), scratch.plane_count());
5341        for index in 0..self.mtp_head_count() {
5342            self.mtp_kv_fill_at(
5343                e,
5344                self.mtp_head_at(index),
5345                tokens,
5346                h,
5347                pos0,
5348                scratch,
5349                index,
5350                embd_dev,
5351            )?;
5352        }
5353        Ok(())
5354    }
5355
5356    /// CAPTURE body for the GRAPH DRAFT (stage 2 of graph-grade spec): ONE MTP head forward with
5357    /// every varying input device-resident —
5358    ///   - token id from the persistent `tok_d` (the previous replay's in-graph argmax wrote it,
5359    ///     so the chain feeds itself; the host reads the same 4 bytes for the draft list),
5360    ///   - h_seed from the persistent `h_seed_d` (h_nextn is copied BACK into it at the end),
5361    ///   - rope pos from the persistent `pos_d` counter (inc'd in-graph),
5362    ///   - scratch KV slot/bound from `scratch.kv.len_d` (see mtp_full_attn_dc).
5363    /// The p-min confidence lands in the persistent `p_d` iff `with_prob` (env is fixed per run).
5364    /// Same kernels, same dispatch as the eager mtp_head_forward_dev chain -> same draft tokens
5365    /// (exactness never depends on drafts — the verify arbitrates — but acceptance parity does).
5366    /// `with_head=false` captures the HEAD-LESS twin for the pseudo-seed replay (2026-07-03):
5367    /// the pseudo pass only needs h_nextn (op 10) + the scratch append — the lm_head read
5368    /// (~1.06ms q6_K on the 9B), argmax and prob are dead weight there. h_nextn's inputs are
5369    /// untouched, so the seed value is identical; round-start resets overwrite tok_d/p_d anyway.
5370    /// `sampled_cap` = Some((ctr_d, perturb_d, q_out_d, seed, temp)) captures the SAMPLED twin
5371    /// (step 3 of the sampled-spec arc): head logits are retained in the persistent `q_out_d`
5372    /// (host D2Ds them to the round's q slot after each replay), the DEVICE event counter is
5373    /// bumped in-graph, and the argmax reads GUMBEL-PERTURBED logits — one categorical draw per
5374    /// replay, bit-identical to the eager arm's gumbel_perturb at the same (seed, sctr, temp).
5375    /// seed/temp are capture-time constants (fixed per generate call, like p_min).
5376    #[allow(clippy::too_many_arguments)]
5377    fn mtp_head_forward_cap(
5378        &self,
5379        e: &Engine,
5380        mtp: &MtpHead,
5381        tok_d: &mut CudaSlice<u32>,
5382        pos_d: &mut CudaSlice<i32>,
5383        h_seed_d: &mut CudaSlice<f32>,
5384        p_d: &mut CudaSlice<f32>,
5385        scratch: &mut MtpScratch,
5386        // Which scratch plane this head appends to / attends over: 0 for the single-head
5387        // chain (every pre-lane caller), the head's own plane index for the multi-head
5388        // chain graphs (each head owns one plane — `mtp_chain_forward_dev`'s contract).
5389        scratch_index: usize,
5390        with_prob: bool,
5391        with_head: bool,
5392        embd_gpu: &CudaSlice<u8>,
5393        embd_qt: i32,
5394        embd_rb: usize,
5395        d_vocab: usize,
5396        sampled_cap: Option<SampledCapArgs<'_>>,
5397        stream_pack: Option<(&mut CudaSlice<u32>, usize, Option<&CudaSlice<u32>>)>,
5398        // DRAFT-SIDE GRAMMAR MASK (lane/draft-mask): (packed draft-vocab allowed-set buffer,
5399        // word count). Captured as ONE mask_logits_f32 node between the head matmul and the
5400        // in-graph argmax; the buffer address is baked, its CONTENTS are re-uploaded by the
5401        // host before every replay (the decode.rs graph-mask pattern). All-ones contents = a
5402        // no-op ban, so a position the grammar cannot constrain costs one pass over the row.
5403        mask_cap: Option<(&CudaSlice<u32>, usize)>,
5404    ) -> Result<(), Box<dyn std::error::Error>> {
5405        let cfg = &self.cfg;
5406        let n_embd = cfg.n_embd as usize;
5407        // step35: capturable through the WINDOWED device-counter arm (`mtp_step35_attn_dcw`)
5408        // once the dcw door is armed and the v3-vec class is live. Without the door this stays
5409        // the deliberate, named refusal: the plain `_dc` attention's key bound always starts at
5410        // row 0, cannot express this block's SWA view offset, and a captured chain would
5411        // silently attend OUTSIDE the window once the persistent scratch passes 512 rows.
5412        // Returning Err (not a panic) is what the capture sites already handle by degrading to
5413        // the eager chain (`mtp_head_forward_dev` -> `mtp_step35_attn`).
5414        // ROUND-STREAM stays refused EITHER WAY: the stream VERIFY has no step35 twin (see the
5415        // step35_verify refusal), so a stream capture that succeeded here would only move the
5416        // failure from capture time (graceful stream-off) to serve time (a failed round).
5417        if let Some(g) = mtp.step35.as_ref() {
5418            if stream_pack.is_some() {
5419                return Err(
5420                    "step35 has no ROUND-STREAM draft arm (the stream verify has no step35 \
5421                     twin); stream off"
5422                        .into(),
5423                );
5424            }
5425            if !self.step35_dcw_eligible(g, scratch.plane(scratch_index).1) {
5426                return Err(format!(
5427                    "step35 has no captured draft chain (fa_decode_dc cannot express the MTP \
5428                        block's SWA view offset; the windowed dcw capture needs \
5429                        MEMRA_STEP35_DRAFT_DCW armed [default ON, =0 disarms] and the v3-vec \
5430                        class live at bucket=min(window {}, scratch cap {})) - the eager draft \
5431                        chain serves this shape",
5432                    g.window,
5433                    scratch.plane(scratch_index).1,
5434                )
5435                .into());
5436            }
5437        }
5438        // student inner width (see mtp_head_forward_dev) — interface dims stay n_embd.
5439        let di = mtp.geom.as_ref().map(|g| g.d_inner).unwrap_or(n_embd);
5440        let eps = cfg.rms_eps;
5441        let e_emb = e.embed_gather_device(embd_gpu, tok_d, n_embd, embd_qt, embd_rb)?;
5442        let mut e_norm = e.zeros(n_embd)?;
5443        e.rms_norm(&e_emb, mtp.enorm.float_data(), &mut e_norm, n_embd, 1, eps)?;
5444        let mut h_norm = e.zeros(n_embd)?;
5445        e.rms_norm(
5446            &*h_seed_d,
5447            mtp.hnorm.float_data(),
5448            &mut h_norm,
5449            n_embd,
5450            1,
5451            eps,
5452        )?;
5453        let mut concat = e.zeros(2 * n_embd)?;
5454        e.copy_into(&mut concat, 0, &e_norm, n_embd)?;
5455        e.copy_into(&mut concat, n_embd, &h_norm, n_embd)?;
5456        let inp_sa = e.matmul(&mtp.eh_proj, &concat, 1)?;
5457        let mut a_norm = e.zeros(di)?;
5458        e.rms_norm(&inp_sa, mtp.attn_norm.float_data(), &mut a_norm, di, 1, eps)?;
5459        let attn_out = match (&mtp.mixer, mtp.step35.as_ref()) {
5460            // step35 (eligibility already enforced by the refusal above): the windowed dcw
5461            // arm, the SAME launcher the eager dev arm runs when the door is armed. No host
5462            // work here (this is the capture body); headroom is the callers' pre-arm.
5463            (Mixer::Full(fa), Some(g)) => {
5464                self.mtp_step35_attn_dcw(e, fa, g, &a_norm, pos_d, scratch, scratch_index)?
5465            }
5466            (Mixer::Full(fa), None) => self.mtp_full_attn_dc(
5467                e,
5468                fa,
5469                &a_norm,
5470                pos_d,
5471                scratch,
5472                scratch_index,
5473                mtp.geom.as_ref(),
5474            )?,
5475            (Mixer::Linear(_), _) => {
5476                panic!("MTP block is full-attn in qwen35; linear MTP not supported")
5477            }
5478            (Mixer::Mla(_), _) => crate::hybrid::mla_forward_unimplemented(),
5479        };
5480        let mut x1 = e.zeros(di)?;
5481        e.add(&inp_sa, &attn_out, &mut x1, di)?;
5482        let mut z = e.zeros(di)?;
5483        e.rms_norm(&x1, mtp.post_attn_norm.float_data(), &mut z, di, 1, eps)?;
5484        let ffn_out = match &mtp.ffn {
5485            crate::hybrid::Ffn::Dense {
5486                ffn_gate,
5487                ffn_up,
5488                ffn_down,
5489            } => {
5490                let n_ff = ffn_gate.out_features();
5491                let (gate, up) = if e.uses_q8_1_fast(ffn_gate) && e.uses_q8_1_fast(ffn_up) {
5492                    let (zq, zd) = e.quantize_q8_1(&z, 1, di)?;
5493                    (
5494                        e.matmul_pre(ffn_gate, &zq, &zd, &z, 1)?,
5495                        e.matmul_pre(ffn_up, &zq, &zd, &z, 1)?,
5496                    )
5497                } else {
5498                    (e.matmul(ffn_gate, &z, 1)?, e.matmul(ffn_up, &z, 1)?)
5499                };
5500                let mut act = e.zeros(n_ff)?;
5501                // step35: the dense FFN reads the per-layer SHEXP clamp, resolved for the MTP
5502                // block's own index (the mtp_head_forward_dev rule; None for every other arch,
5503                // which is `ffn_act`'s dispatch verbatim). The eager and captured chains must
5504                // run the ONE activation program.
5505                Self::ffn_act_lim(
5506                    e,
5507                    &self.cfg,
5508                    &gate,
5509                    &up,
5510                    1.0,
5511                    1.0,
5512                    mtp.step35.as_ref().and_then(|s| s.clamp_shexp),
5513                    &mut act,
5514                    n_ff,
5515                )?;
5516                e.matmul(ffn_down, &act, 1)?
5517            }
5518            // ROUND-STREAM: the 35B NextN block carries a MoE FFN. With RESIDENT experts the
5519            // dev path is pure device launches (device top-k + rows kernels, ZERO-DtoH by
5520            // design) — capture-legal. Non-resident (SLRU-lock) stays rejected: the capture
5521            // error arm degrades the caller to eager/stream-off.
5522            crate::hybrid::Ffn::Moe(m) if m.dev_exps.is_some() => {
5523                self.moe_ffn_il(e, m, &z, 1, u16::MAX)?
5524            }
5525            crate::hybrid::Ffn::Moe(_) => {
5526                return Err("graph draft requires a Dense (or resident-MoE) MTP FFN".into());
5527            }
5528        };
5529        let mut h_inner = e.zeros(di)?;
5530        e.add(&x1, &ffn_out, &mut h_inner, di)?;
5531        // student: up-project back to n_embd (carrier + head input; see mtp_head_forward_dev).
5532        let h_nextn = match mtp.geom.as_ref() {
5533            Some(g) => e.matmul(&g.out_up, &h_inner, 1)?,
5534            None => h_inner,
5535        };
5536        // MEMRA_SPEC_HPOST needs final_h even head-less (it IS the next seed under that convention).
5537        let final_h = if with_head || spec_hpost() {
5538            let final_norm = mtp.shared_head_norm.as_ref().unwrap_or(&self.output_norm);
5539            let mut fh = e.zeros(n_embd)?;
5540            e.rms_norm(&h_nextn, final_norm.float_data(), &mut fh, n_embd, 1, eps)?;
5541            Some(fh)
5542        } else {
5543            None
5544        };
5545        if with_head {
5546            let head = mtp.shared_head_head.as_ref().unwrap_or(&self.output);
5547            let mut logits = e.matmul(head, final_h.as_ref().unwrap(), 1)?;
5548            // DRAFT-SIDE GRAMMAR MASK: ban the grammar-illegal draft ids IN the captured chain,
5549            // before the argmax — proposals become legal by construction. Contents-only
5550            // per-replay upload keeps the capture valid.
5551            if let Some((mask_d, mw)) = mask_cap {
5552                e.mask_logits_col(&mut logits, mask_d, 0, d_vocab, mw)?;
5553            }
5554            if let Some(SampledCapArgs {
5555                ctr: ctr_d,
5556                perturb: perturb_d,
5557                q_out: q_out_d,
5558                seed,
5559                temp,
5560                filt,
5561            }) = sampled_cap
5562            {
5563                // SAMPLED chain: retain q (raw head logits -> persistent q_out_d; the matmul's
5564                // own buffer is pool-recycled after the capture body returns, so it can't be the
5565                // retention target), bump the device event counter, gumbel-perturb reading it,
5566                // and argmax the PERTURBED logits into tok_d — the in-graph categorical draw.
5567                e.copy_into(q_out_d, 0, &logits, d_vocab)?;
5568                e.sctr_inc(ctr_d)?;
5569                match filt {
5570                    // PURE-TEMP: gumbel over the raw softmax — byte-identical to the
5571                    // pre-lane capture body.
5572                    None => e.gumbel_perturb_ctr(&logits, perturb_d, d_vocab, seed, ctr_d, temp)?,
5573                    // FILTERED (lane/step37-draft-graph-serving-20260830): the SAME
5574                    // filter_stats program the eager arm and the accept path run (the
5575                    // wrapper's coop/plain choice is deployment-keyed, never per-call), then
5576                    // the device-stat/device-counter perturb twin — the draft draws from the
5577                    // exact filtered distribution the verify gathers `q` from. q was
5578                    // retained ABOVE, pre-perturb, so the accept path's post-replay stats
5579                    // recompute (same kernel, same bits) reconstructs these th/z exactly.
5580                    Some(f) => {
5581                        e.filter_stats(
5582                            &logits, d_vocab, f.rows0, f.th, f.z, f.mx, d_vocab, 1, temp, f.top_k,
5583                            f.top_p, f.min_p,
5584                        )?;
5585                        e.gumbel_perturb_filtered_ctr(
5586                            &logits, perturb_d, d_vocab, seed, ctr_d, temp, f.mx, f.th,
5587                        )?;
5588                    }
5589                }
5590                e.argmax_token_device_into(perturb_d, tok_d, d_vocab)?;
5591                // p-min prob = the head's RAW softmax confidence in the SAMPLED pick — same
5592                // semantics as the eager sampled arm's prob_of_token_device(dl_d, tok_d).
5593                if with_prob {
5594                    e.prob_of_token_device_into(&logits, tok_d, p_d, d_vocab)?;
5595                }
5596            } else {
5597                // draft token -> persistent tok_d (next replay's embed reads it; host reads the 4 bytes).
5598                e.argmax_token_device_into(&logits, tok_d, d_vocab)?;
5599                // p-min under a draft mask reads the MASKED row: confidence relative to the
5600                // grammar-LEGAL alternatives (illegal ids leave the softmax denominator), which
5601                // is the right semantics for "does the drafter know what comes next here" and
5602                // the same row the pick came from. Draft-quality only — verify arbitrates.
5603                if with_prob {
5604                    e.prob_of_token_device_into(&logits, tok_d, p_d, d_vocab)?;
5605                }
5606            }
5607        }
5608        // ROUND-STREAM K-chain: pack (tok, p) into slot j, then remap tok through d2t so the
5609        // NEXT chained body's embed reads the TARGET id — zero host involvement per step.
5610        if let Some((out, slot, d2t)) = stream_pack {
5611            e.pack_tok_p(tok_d, p_d, out, slot)?;
5612            if let Some(map) = d2t {
5613                e.tok_map_u32(tok_d, map)?;
5614            }
5615        }
5616        // Next draft step's h_seed: pre-norm h_nextn (default) or post-norm final_h (HPOST).
5617        if spec_hpost() {
5618            e.copy_into(h_seed_d, 0, final_h.as_ref().unwrap(), n_embd)?;
5619        } else {
5620            e.copy_into(h_seed_d, 0, &h_nextn, n_embd)?;
5621        }
5622        // advance the draft rope position in-graph.
5623        e.inc_seqlen(pos_d)?;
5624        Ok(())
5625    }
5626
5627    /// Batched target verify forward over `tokens` at positions `pos0..pos0+T` (§D.3, T=K+1).
5628    /// Returns ALL T logit columns (host f32, [T*n_vocab]); appends T cols to every full-attn KV
5629    /// and advances every linear-attn recur state by T steps (the recur steps are SEQUENTIAL T=1).
5630    /// Advances `cache.pos` by T.
5631    pub fn decode_step_t(
5632        &self,
5633        e: &Engine,
5634        tokens: &[u32],
5635        pos0: usize,
5636        cache: &mut Cache,
5637    ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
5638        if self.is_gemma4_e4b() {
5639            return Ok(self.gemma4_e4b_decode_step_t_h(e, tokens, pos0, cache)?.0);
5640        }
5641        if self.gemma_batch_program() {
5642            return self.gemma4_decode_step_t(e, tokens, pos0, cache);
5643        }
5644        Ok(self.decode_step_t_h(e, tokens, pos0, cache)?.0)
5645    }
5646
5647    /// Like `decode_step_t` but ALSO returns the LAST column's pre-output_norm hidden (h_seed for
5648    /// the next draft round). This lets partial-accept replay run as ONE batched T=(n_acc+1) forward
5649    /// (single weight read) instead of n_acc+1 separate T=1 decode_steps (n_acc+1 weight reads).
5650    /// At batch=1 decode is bandwidth-bound, so batching the replay is THE MTP profitability lever.
5651    pub fn decode_step_t_h(
5652        &self,
5653        e: &Engine,
5654        tokens: &[u32],
5655        pos0: usize,
5656        cache: &mut Cache,
5657    ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
5658        self.decode_step_t_h_emb(e, tokens, pos0, cache, None)
5659    }
5660
5661    /// Like `decode_step_t_h` with an optional RESIDENT embed table (spec hot loop): device
5662    /// gather instead of host dequant + [T, n_embd] f32 htod. Bit-identical rows.
5663    pub fn decode_step_t_h_emb(
5664        &self,
5665        e: &Engine,
5666        tokens: &[u32],
5667        pos0: usize,
5668        cache: &mut Cache,
5669        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
5670    ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
5671        let (logits_d, h_seed) = self.decode_step_t_h_emb_dev(e, tokens, pos0, cache, embd_dev)?;
5672        Ok((e.dtoh(&logits_d)?, h_seed))
5673    }
5674
5675    /// DEVICE-LOGITS verify forward (spec device-argmax lever): identical kernel chain to
5676    /// `decode_step_t_h_emb` but returns the [T, n_vocab] logits ON DEVICE — the accept walk
5677    /// argmaxes each column on-device and reads back ONE [T] u32 instead of dtoh'ing the full
5678    /// T x n_vocab f32 block (~1-4 MB + T host argmaxes, every round). Kernel dispatch is
5679    /// UNCHANGED (same decode-exact kernels); only the post-logits transfer moves.
5680    pub fn decode_step_t_h_emb_dev(
5681        &self,
5682        e: &Engine,
5683        tokens: &[u32],
5684        pos0: usize,
5685        cache: &mut Cache,
5686        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
5687    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
5688        let n_embd = self.cfg.n_embd as usize;
5689        let t = tokens.len();
5690        let (logits, x) = self.decode_step_t_core(e, tokens, pos0, cache, embd_dev, None)?;
5691        // h_seed for the next round = LAST column's pre-output_norm hidden ([n_embd]).
5692        let mut hs = vbuf(e, n_embd)?; // fully written by copy_view_into below
5693        e.copy_view_into(&mut hs, 0, &x.slice((t - 1) * n_embd..t * n_embd), n_embd)?;
5694        Ok((logits, hs))
5695    }
5696
5697    /// CORE verify forward: the `decode_step_t_h_emb_dev` kernel chain, returning the FULL
5698    /// pre-output_norm hidden stack x ([T, n_embd], any column extractable) and optionally
5699    /// filling a `VerifyCkpt` (retained per-layer state-rebuild inputs) for the REPLAY-FREE
5700    /// partial accept. `ckpt: None` => byte-for-byte the old behavior (the ckpt writes are pure
5701    /// retains/copies — they never change what any kernel computes).
5702    fn decode_step_t_core(
5703        &self,
5704        e: &Engine,
5705        tokens: &[u32],
5706        pos0: usize,
5707        cache: &mut Cache,
5708        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
5709        mut ckpt: Option<&mut VerifyCkpt>,
5710    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
5711        self.decode_step_t_core_stream(
5712            e,
5713            tokens,
5714            pos0,
5715            cache,
5716            embd_dev,
5717            ckpt.take(),
5718            None,
5719            None,
5720            None,
5721            None,
5722        )
5723    }
5724
5725    /// [`Self::decode_step_t_core`] with the MTP route's verify-graph pool armed
5726    /// (`MEMRA_SPEC_VERIFY_GRAPH`). `graphs: None` reproduces `decode_step_t_core`
5727    /// argument-for-argument, so the eager walk stays the byte-identical fallback.
5728    fn decode_step_t_core_vg(
5729        &self,
5730        e: &Engine,
5731        tokens: &[u32],
5732        pos0: usize,
5733        cache: &mut Cache,
5734        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
5735        mut ckpt: Option<&mut VerifyCkpt>,
5736        graphs: Option<&mut DsparkVerifyGraphs>,
5737    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
5738        self.decode_step_t_core_stream(
5739            e,
5740            tokens,
5741            pos0,
5742            cache,
5743            embd_dev,
5744            ckpt.take(),
5745            None,
5746            None,
5747            None,
5748            graphs,
5749        )
5750    }
5751
5752    /// Increment-0 two-session PP seam: release the peer after this lane's stage-0 boundary TX.
5753    /// The two independent sessions keep their own cache/checkpoint state; only issue order moves.
5754    fn decode_step_t_core_pipelined(
5755        &self,
5756        e: &Engine,
5757        tokens: &[u32],
5758        pos0: usize,
5759        cache: &mut Cache,
5760        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
5761        mut ckpt: Option<&mut VerifyCkpt>,
5762        pipe: &SpecPipeLane,
5763        round: usize,
5764    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
5765        let fence = crate::pp::pp_cuts(self.layers.len())
5766            .ok_or("two-session speculative pipeline requires a PP stage cut")?;
5767        if crate::pp::pp2_streams_off() || !crate::pp::spec_pp_on() {
5768            return Err("two-session speculative pipeline requires the PP verify split".into());
5769        }
5770        let interval_fence = pipe.stage0_begin(round)?;
5771        let ticket = self.verify_stage0_issue(
5772            e,
5773            tokens,
5774            pos0,
5775            cache,
5776            embd_dev,
5777            ckpt.as_deref_mut(),
5778            None,
5779            &fence,
5780            Some(interval_fence),
5781            pipe.trace(round),
5782        )?;
5783        pipe.stage0_end(round);
5784        pipe.stage1_begin(round)?;
5785        let result = self.verify_stage1_finish(e, ticket, cache, ckpt, None, &fence, true)?;
5786        pipe.verify_end(round);
5787        Ok(result)
5788    }
5789
5790    /// ROUND-STREAM stage (c) 4: `stream` = (device verify tokens [t], device pos counter) —
5791    /// when Some, rope positions come from pos_iota over the counter, the embed gathers the
5792    /// device tokens, and full_attn_verify routes appends/FA through the _dc twins reading the
5793    /// SAME counter (every layer's kvl.len == cache.pos, one counter drives all three). The
5794    /// host `tokens`/`pos0` args still size buffers (t is FIXED K+1 in stream mode).
5795    /// `vtok_dev` (engine-bundle slice 2): device verify tokens for the EMBED only —
5796    /// unlike `stream` mode it changes nothing else (host pos iota, host-len KV appends).
5797    /// `tokens` then only sizes buffers (the dummy-slice pattern the round-stream arm uses).
5798    #[allow(clippy::too_many_arguments)]
5799    fn decode_step_t_core_stream(
5800        &self,
5801        e: &Engine,
5802        tokens: &[u32],
5803        pos0: usize,
5804        cache: &mut Cache,
5805        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
5806        mut ckpt: Option<&mut VerifyCkpt>,
5807        stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
5808        pp_pipe: Option<bool>,
5809        vtok_dev: Option<&CudaSlice<u32>>,
5810        graphs: Option<&mut DsparkVerifyGraphs>,
5811    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
5812        // PP DOOR (lane/pp2-spec 2026-08-06): the verify trunk now takes its OWN stage split,
5813        // exactly as the eager and batched steps do. This is the single funnel every verify
5814        // forward reaches (decode_step_t / _h / _h_emb / _h_emb_dev / _core all land here), so
5815        // wiring it here wires the whole spec surface — the draft/accept/commit machinery above
5816        // is untouched.
5817        //
5818        // History: pp2-hardening (2026-08-06) made this funnel FAIL CLOSED, because its trunk
5819        // walk was unsplit on one stream and a sharded cross-device placement peer-read every
5820        // remote layer's weights on every spec round (measured 13.9-28x on the batched twin).
5821        // The refusal below survives to cover the residue — MEMRA_SPEC_PP=0, MEMRA_PP_STREAMS=0,
5822        // or a placement whose PpNRt fails to build — so a config that would still walk the
5823        // whole trunk on one stream refuses instead of regressing 28x.
5824        if let Some(fence) = crate::pp::pp_cuts(self.layers.len()) {
5825            if !crate::pp::pp2_streams_off() && crate::pp::spec_pp_on() {
5826                if vtok_dev.is_some() {
5827                    return Err(
5828                        "device-token dspark verify (slice-2 deferred readback) has no PP \
5829                         stage-split arm; set MEMRA_DSPARK_DEFER_READBACK=0 or run the dspark \
5830                         route on one device"
5831                            .into(),
5832                    );
5833                }
5834                return self.decode_step_t_core_ppn(
5835                    e,
5836                    tokens,
5837                    pos0,
5838                    cache,
5839                    embd_dev,
5840                    ckpt.take(),
5841                    stream,
5842                    &fence,
5843                    pp_pipe,
5844                );
5845            }
5846        }
5847        crate::pp::refuse_unsplit_if_remote(
5848            "decode_step_t (spec verify)",
5849            "drop MEMRA_SPEC_PP=0 / MEMRA_PP_STREAMS=0 so the verify trunk takes its OWN stage \
5850             split (decode_step_t_core_ppn); or run spec on one device",
5851        )?;
5852        let cfg = &self.cfg;
5853        let n_embd = cfg.n_embd as usize;
5854        let eps = cfg.rms_eps;
5855        let t = tokens.len();
5856        let pos_d = match stream {
5857            Some((_, ctr)) => {
5858                let mut p = e.alloc_uninit::<i32>(t)?;
5859                e.pos_iota(ctr, &mut p, t)?;
5860                p
5861            }
5862            None => {
5863                let pos_vec: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
5864                e.htod_i32(&pos_vec)?
5865            }
5866        };
5867
5868        // embed T tokens -> [T, n_embd] token-major (device gather on the spec hot loop)
5869        let x = match (stream, embd_dev) {
5870            (Some((vtok, _)), Some((g, qt, rb))) => {
5871                e.embed_gather_device_td(g, vtok, t, n_embd, qt, rb)?
5872            }
5873            (None, Some((g, qt, rb))) => match vtok_dev {
5874                // slice 2: device verify tokens, same embed_gather_u32_t kernel —
5875                // bit-identical rows to the host-token arm (same per-dtype deq).
5876                Some(vt_d) => e.embed_gather_device_td(g, vt_d, t, n_embd, qt, rb)?,
5877                None => e.embed_gather_device_t(g, tokens, n_embd, qt, rb)?,
5878            },
5879            _ => {
5880                assert!(
5881                    vtok_dev.is_none(),
5882                    "device-token verify requires the resident embed table (embd_dev)"
5883                );
5884                e.htod(&self.embd.gather(n_embd, tokens))?
5885            }
5886        };
5887
5888        // TRUNK WALK: layers [0, n_layers) through the SAME range-scoped subgraph the PP-N
5889        // stage split calls per stage (`verify_layers`) — one code path, so the split cannot
5890        // drift from the unsplit dispatch mirroring. lane/pp2-spec 2026-08-06.
5891        let x = self.verify_layers(
5892            e,
5893            x,
5894            0,
5895            self.layers.len(),
5896            &pos_d,
5897            pos0,
5898            t,
5899            cache,
5900            ckpt.take(),
5901            stream,
5902            graphs,
5903        )?;
5904        if spec_nan_scan() {
5905            nan_scan_rows(e, &x, t, n_embd, &format!("verify trunk exit pos0={pos0}"))?;
5906        }
5907
5908        let mut hn = vbuf(e, t * n_embd)?;
5909        // Stage-A door: with the serving-class row-outer verify walk, the TAIL must be the
5910        // t=1 decode program per row too (rms_norm t=1 + the single-row bf16 head — the
5911        // split head's concat is receipted bit-identical to it). The batched cuBLASLt head
5912        // is a different ULP class and flips near-tie argmaxes off the greedy tape.
5913        let eager_tail = self.sliding_gated_moe_batch_program() && spec_verify_eager_on();
5914        if eager_tail {
5915            let n_vocab = self.cfg.n_vocab as usize;
5916            // MEMRA_SPEC_HEAD_ROWS=1 — THE VERIFY TAIL'S REDUNDANT HEAD READ.
5917            //
5918            // The loop below runs the head at m=1 once PER COLUMN, so the LM head's weights are
5919            // streamed t times per verify pass. On step37 that head is ~0.49 GiB per card after the
5920            // rank split, ~1.07 ms of pure re-read at t=2 and worse at every wider t — which is a
5921            // large part of why the fixed K ladder LOSES (K=1 81.2 > K=2 73.1 > K=3 62.7 tok/s).
5922            //
5923            // The loop's justification is the comment above: the batched cuBLASLt head is a
5924            // different ULP class and flips near-tie argmaxes off the greedy tape. That is true of
5925            // cuBLASLt and it does NOT apply here, because a FloatBf16 head at 1..=32 rows never
5926            // reaches cuBLASLt: `matmul` routes it to `matvec_bf16_rows_into` (lib.rs:12248), whose
5927            // own doc says `matvec_bf16_f32acc_x4_rows` "runs the t=1 decode head program PER ROW
5928            // (identical dot + reduce), so decode/verify tiers keep the t=1 numeric class". Under
5929            // the W8 doors both widths route to the q8 mirror instead, and the t-column mirror is
5930            // documented "bit-identical to t single-row calls". So the batched form is the SAME
5931            // arithmetic per row on both paths, with one weight read instead of t.
5932            //
5933            // rms_norm is row-wise, so norm(t) is per-row identical to t x norm(1) by construction.
5934            //
5935            // DEFAULT OFF for exactly one turn of the crank: "bit-identical by two documented
5936            // claims" is still an argument. The greedy byte tape decides, and the door flips only
5937            // once the tape is a receipt.
5938            if head_rows_on() {
5939                e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
5940                let logits = e.matmul(&self.output, &hn, t)?;
5941                if stream.is_none() {
5942                    cache.pos += t;
5943                }
5944                return Ok((logits, if spec_hpost() { hn } else { x }));
5945            }
5946            let mut logits = vbuf(e, t * n_vocab)?;
5947            for r in 0..t {
5948                let mut row = e.uninit(n_embd)?;
5949                e.dtod_copy_view(&x.slice(r * n_embd..(r + 1) * n_embd), &mut row)?;
5950                let mut hr = e.uninit(n_embd)?;
5951                e.rms_norm(&row, self.output_norm.float_data(), &mut hr, n_embd, 1, eps)?;
5952                let lr = e.matmul(&self.output, &hr, 1)?;
5953                e.dtod_copy_into(&lr, &mut logits, r * n_vocab)?;
5954                e.dtod_copy_into(&hr, &mut hn, r * n_embd)?;
5955            }
5956            if stream.is_none() {
5957                cache.pos += t;
5958            }
5959            return Ok((logits, if spec_hpost() { hn } else { x }));
5960        }
5961        let serving_head =
5962            self.sliding_gated_moe_batch_program() || self.batched_serving_numeric_class();
5963        let logits = if serving_head {
5964            // Step35 and the qwen35 family (MoE 2026-08-14 AM, dense-hybrid same day PM — the
5965            // Q3.8 bring-up reproduced the identical near-tie class on dense: eager-class verify
5966            // vs batched-class live serving, ULP drift amplified through the GDN recurrence)
5967            // serve one batched numeric class at every live width, including B=1. Keep the
5968            // verify head in that same class; other generic families retain the decode-exact
5969            // head that their run-spec contract pins.
5970            e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
5971            e.matmul(&self.output, &hn, t)?
5972        } else {
5973            e.rms_norm_decode(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
5974            e.matmul_decode_exact(&self.output, &hn, t)?
5975        };
5976        // stream: the device pos counter owns position; host mirror reconciles at drain.
5977        if stream.is_none() {
5978            cache.pos += t;
5979        }
5980        // Hidden stack for seeds/refresh-fills: pre-norm x (default) or post-norm hn (HPOST).
5981        Ok((logits, if spec_hpost() { hn } else { x }))
5982    }
5983
5984    /// THE VERIFY TRUNK OVER PP-N (lane/pp2-spec 2026-08-06): `decode_step_t_core_stream`'s walk
5985    /// as N stage subgraphs, each on its own engine/stream (and, under `MEMRA_PP_DEVICES`, its own
5986    /// device), with a `[T, n_embd]` boundary transfer between them. T = K+1 (the verify batch),
5987    /// so this is the batched-boundary shape the pp2-batch lane's grow-only slots already handle
5988    /// (`tx(b, x, t*n_embd)`; the slot grows to the high-water T and the transport moves exactly
5989    /// the payload).
5990    ///
5991    /// Structure is `decode_step_batch_ppn`'s, which is `decode_step_h_ppn`'s. FOUR THINGS ARE
5992    /// PER-STAGE and each for a measured reason (see `decode_step_batch_ppn`'s header for the
5993    /// receipts):
5994    ///
5995    /// 1. THE ENGINE (`rt.engine(s, e)`) — `Engine` owns lazily-grown stable-pointer scratch
5996    ///    (`fa_part_pool`, `fa_vf16_scratch`, `argmax_partials`) that is single-stream-safe BY
5997    ///    DESIGN. Two stage streams through one Engine is the 2026-08-02 shared-scratch race
5998    ///    (35% flake, nondeterministic all-logits divergence). `PpNRt::build` gives every stage
5999    ///    s>0 its own Engine even on the primary device; honouring it here is what scopes the
6000    ///    pools. The verify path allocates MORE of that scratch than eager decode does (FA at
6001    ///    m=T, and the per-layer `GdnStash` retains), so this is load-bearing, not inherited.
6002    ///
6003    /// 2. `pos_d` — each stage uploads its OWN copy of the T rope positions on ITS stream, so the
6004    ///    buffer is allocated, consumed and freed on one stream. In `stream` mode that means each
6005    ///    stage runs its own `pos_iota` over the SHARED device counter (`pos_ctr`): the counter is
6006    ///    read-only during the forward (the round's `inc`/`copy_add` happen outside it), so every
6007    ///    stage derives the identical iota, and each stage's own output buffer is stream-local.
6008    ///
6009    /// 3. THE EMBED lives with stage 0 (`self.embd` / `embd_gpu` are host/primary-side; the
6010    ///    sharded loader leaves the table with stage 0 by construction).
6011    ///
6012    /// 4. THE HEAD (`output_norm` + `output`) runs on the LAST stage — the sharded loader uploaded
6013    ///    both through that stage's engine (`hybrid.rs`: `e_head = layer_engine(e, n_trunk,
6014    ///    n_trunk-1)`), so reading them anywhere else is a peer read of the biggest tensor in the
6015    ///    model, every round.
6016    ///
6017    /// WHAT STAYS ON THE PRIMARY, deliberately: the returned logits and hidden stack `x`. Both are
6018    /// last-stage-allocated device buffers, and every consumer (the device argmax walk, the accept
6019    /// kernels, `spec_seed_gather`, the ckpt rebuild in `commit_verified_prefix`) reads them
6020    /// through the primary context by UVA — the same read the batched serving epilogue's
6021    /// `last_logits_dev` park does. Those consumers are per-round O(T x n_vocab) and O(n_embd),
6022    /// not per-layer, so they are not the 28x class; splitting them is a separate lane.
6023    ///
6024    /// The MTP HEAD (draft side) is NOT split: it is one block, it lives wherever the loader put
6025    /// it (`load_mtp` uses the primary engine), and it is ~1-2 GB against the trunk's tens. Draft
6026    /// placement is measured, not assumed — see `research/pp2-spec-20260806`.
6027    ///
6028    /// EXACTNESS: PP-N adds ZERO deviation. Each stage runs the SAME kernels on the SAME bytes in
6029    /// the same order via the SAME `verify_layers` the unsplit body calls; the only change is
6030    /// where the residual is materialized, and the boundary is a straight f32 copy (dtod
6031    /// same-device / `cudaMemcpyPeerAsync` cross-device, no conversion). So the split MUST be
6032    /// BIT-IDENTICAL to the unsplit verify at the same T, in both placement orders. Gate:
6033    /// `decode-batch-gate --mode ppspec`. Acceptance counts are a DERIVED consequence — greedy
6034    /// accept argmaxes these logits, so bit-identical logits force identical accept walks; the
6035    /// `run-spec` K=1..8 arm checks that end-to-end rather than trusting the implication.
6036    #[allow(clippy::too_many_arguments)]
6037    fn decode_step_t_core_ppn(
6038        &self,
6039        e: &Engine,
6040        tokens: &[u32],
6041        pos0: usize,
6042        cache: &mut Cache,
6043        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
6044        mut ckpt: Option<&mut VerifyCkpt>,
6045        stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
6046        fence: &[usize],
6047        pp_pipe: Option<bool>,
6048    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
6049        let ticket = self.verify_stage0_issue(
6050            e,
6051            tokens,
6052            pos0,
6053            cache,
6054            embd_dev,
6055            ckpt.as_deref_mut(),
6056            stream,
6057            fence,
6058            pp_pipe,
6059            None,
6060        )?;
6061        self.verify_stage1_finish(e, ticket, cache, ckpt, stream, fence, true)
6062    }
6063
6064    /// Enqueue embed, stage 0, and the first boundary TX, then return the actual boundary slot.
6065    /// The ordinary PP verify wrapper calls `verify_stage1_finish` immediately after this return.
6066    #[allow(clippy::too_many_arguments)]
6067    fn verify_stage0_issue(
6068        &self,
6069        e: &Engine,
6070        tokens: &[u32],
6071        pos0: usize,
6072        cache: &mut Cache,
6073        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
6074        mut ckpt: Option<&mut VerifyCkpt>,
6075        stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
6076        fence: &[usize],
6077        pp_pipe: Option<bool>,
6078        trace: Option<SpecPipeTraceCtx>,
6079    ) -> Result<VerifyBoundaryTicket, Box<dyn std::error::Error>> {
6080        assert!(
6081            !self.is_gemma4_e4b() && !self.gemma_batch_program(),
6082            "decode_step_t_core_ppn covers the hybrid non-gemma4 verify trunk only \
6083             (the gemma4 arms have their own decode_step_t twins)"
6084        );
6085        if crate::pp::pp_host_bounce_active() && (stream.is_some() || embd_dev.is_some()) {
6086            return Err(
6087                "decode_step_t_core_ppn: refused with MEMRA_PP_HOST_BOUNCE=1 — the trunk \
6088                 boundary itself is host-staged, but device-resident verify still peer-reads \
6089                 primary-device token/position/embedding buffers from stage 0. Run plain PP \
6090                 serving on this host class; spec requires local per-stage inputs first."
6091                    .into(),
6092            );
6093        }
6094        let rt = crate::pp::PpNRt::get(e)?;
6095        let n_st = fence.len() - 1;
6096        assert_eq!(
6097            rt.n_stages(),
6098            n_st,
6099            "PpNRt stage count {} != fence stages {n_st}",
6100            rt.n_stages()
6101        );
6102        let n_embd = self.cfg.n_embd as usize;
6103        let t = tokens.len();
6104        let payload = t * n_embd;
6105        if pp_pipe.is_some() {
6106            assert_eq!(n_st, 2, "spec pipeline requires exactly two PP stages");
6107        }
6108        // One-shot lane diagnostic: force natural PP-2 boundaries to completion so the server
6109        // log can price stage 0, the peer hop, the RX copy, and stage 1 + head separately without
6110        // nsys. The ordinary path keeps every enqueue asynchronous. N>2 is deliberately excluded:
6111        // the report below names exactly two stages and must never imply it measured middle ones.
6112        let pp_anatomy = n_st == 2 && std::env::var("MEMRA_SPEC_PP_ANATOMY").as_deref() == Ok("1");
6113        let pp_started = std::time::Instant::now();
6114        let (mut reverse_ms, mut stage0_ms, mut tx_ms) = (0.0f64, 0.0f64, 0.0f64);
6115        // The CALLER's ambient stream, captured BEFORE any `rt.enter()` pushes a stage stream:
6116        // this body returns DEVICE-RESIDENT buffers (the device-argmax accept walk's contract),
6117        // so the exit needs the same publication the boundaries get — see `PpNRt::publish_to`.
6118        // Taken here, not at the end, because inside the last-stage scope `e.stream()` IS the
6119        // stage stream and the wait would self-order into a no-op.
6120        let caller_stream = e.stream();
6121        // #87 ROOT-CAUSE FENCE (lane/pp2spec-crash): the PREVIOUS round's stage-allocated
6122        // outputs (logits/hidden/ckpt stashes) freed stream-ordered on the STAGE streams while
6123        // the primary stream still holds queued reads of them — with event tracking elided,
6124        // nothing stops the pool from reusing those blocks for THIS round's stage allocations,
6125        // whose writes then race the queued reads (measured: 13/4096-NaN random-bits garbage in
6126        // the spec round seed; the full anatomy is on `PpNRt::fence_stages_behind`). Order every
6127        // stage stream behind the caller before enqueueing new stage work.
6128        let reverse_started = std::time::Instant::now();
6129        if pp_pipe != Some(false) {
6130            rt.fence_stages_behind(&caller_stream)?;
6131        }
6132        if pp_pipe == Some(true) {
6133            // Both session verifies must alternate boundary slots even when the ordinary
6134            // decode overlap experiment is off. Prewarm before A's stage 0 so B cannot grow
6135            // slot 1 by synchronizing the RX stream while A's stage 1 is in flight.
6136            rt.prepare_overlap_slots(0, payload)?;
6137        }
6138        if pp_anatomy {
6139            // Drain the reverse-publication dependency before timing stage 0 itself. At c=1 this
6140            // prices any primary-stream rollback/refresh tail inherited from the prior round.
6141            for s in 0..n_st {
6142                let _st = rt.enter(s);
6143                rt.engine(s, e).stream().synchronize()?;
6144            }
6145            reverse_ms = reverse_started.elapsed().as_secs_f64() * 1e3;
6146        }
6147
6148        // Per-stage rope positions: in host mode the same [T] iota each stage uploads itself; in
6149        // stream mode each stage's own `pos_iota` over the shared read-only device counter.
6150        let stage_pos = |es: &Engine| -> Result<CudaSlice<i32>, Box<dyn std::error::Error>> {
6151            match stream {
6152                Some((_, ctr)) => {
6153                    let mut p = es.alloc_uninit::<i32>(t)?;
6154                    es.pos_iota(ctr, &mut p, t)?;
6155                    Ok(p)
6156                }
6157                None => {
6158                    let pos_vec: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
6159                    es.htod_i32(&pos_vec)
6160                }
6161            }
6162        };
6163
6164        // ---- STAGE 0: embed (the table lives with stage 0) + layers [0, fence[1]) + TX ----
6165        let slot = {
6166            let _st0 = rt.enter(0);
6167            let e0 = rt.engine(0, e);
6168            enqueue_spec_pipe_trace_marker(&e0.stream(), trace.as_ref(), "S0", "start", None)?;
6169            let stage0_started = std::time::Instant::now();
6170            let pos_d = stage_pos(e0)?;
6171            let x = match (stream, embd_dev) {
6172                (Some((vtok, _)), Some((g, qt, rb))) => {
6173                    e0.embed_gather_device_td(g, vtok, t, n_embd, qt, rb)?
6174                }
6175                (None, Some((g, qt, rb))) => e0.embed_gather_device_t(g, tokens, n_embd, qt, rb)?,
6176                _ => e0.htod(&self.embd.gather(n_embd, tokens))?,
6177            };
6178            let x = self.verify_layers(
6179                e0,
6180                x,
6181                fence[0],
6182                fence[1],
6183                &pos_d,
6184                pos0,
6185                t,
6186                cache,
6187                ckpt.as_deref_mut(),
6188                stream,
6189                None,
6190            )?;
6191            if pp_anatomy {
6192                e0.stream().synchronize()?;
6193                stage0_ms = stage0_started.elapsed().as_secs_f64() * 1e3;
6194            }
6195            let tx_started = std::time::Instant::now();
6196            let slot = if pp_pipe.is_some() {
6197                rt.tx_pipelined(0, &x, payload)?
6198            } else {
6199                rt.tx(0, &x, payload)?
6200            };
6201            enqueue_spec_pipe_trace_marker(&e0.stream(), trace.as_ref(), "S0", "end", Some(slot))?;
6202            if pp_anatomy {
6203                e0.stream().synchronize()?;
6204                tx_ms = tx_started.elapsed().as_secs_f64() * 1e3;
6205            }
6206            slot
6207            // x + pos_d drop here: freed stream-ordered on stage-0's stream after use.
6208        };
6209
6210        Ok(VerifyBoundaryTicket {
6211            rt,
6212            caller_stream,
6213            slot,
6214            pos0,
6215            t,
6216            payload,
6217            n_st,
6218            pipelined: pp_pipe.is_some(),
6219            pp_anatomy,
6220            pp_started,
6221            reverse_ms,
6222            stage0_ms,
6223            tx_ms,
6224            trace,
6225        })
6226    }
6227
6228    /// Consume a stage-0 boundary ticket and enqueue the remaining PP stages plus the head.
6229    /// On PP-2 this is exactly stage 1; PP-N keeps its pre-existing middle-stage walk here.
6230    #[allow(clippy::too_many_arguments)]
6231    fn verify_stage1_finish(
6232        &self,
6233        e: &Engine,
6234        ticket: VerifyBoundaryTicket,
6235        cache: &mut Cache,
6236        mut ckpt: Option<&mut VerifyCkpt>,
6237        stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
6238        fence: &[usize],
6239        publish_to_caller: bool,
6240    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
6241        let VerifyBoundaryTicket {
6242            rt,
6243            caller_stream,
6244            slot,
6245            pos0,
6246            t,
6247            payload,
6248            n_st,
6249            pipelined,
6250            pp_anatomy,
6251            pp_started,
6252            reverse_ms,
6253            stage0_ms,
6254            tx_ms,
6255            trace,
6256        } = ticket;
6257        let n_embd = self.cfg.n_embd as usize;
6258        let eps = self.cfg.rms_eps;
6259        let mut slot = slot;
6260        let (mut rx_ms, mut stage1_ms) = (0.0f64, 0.0f64);
6261        let stage_pos = |es: &Engine| -> Result<CudaSlice<i32>, Box<dyn std::error::Error>> {
6262            match stream {
6263                Some((_, ctr)) => {
6264                    let mut p = es.alloc_uninit::<i32>(t)?;
6265                    es.pos_iota(ctr, &mut p, t)?;
6266                    Ok(p)
6267                }
6268                None => {
6269                    let pos_vec: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
6270                    es.htod_i32(&pos_vec)
6271                }
6272            }
6273        };
6274
6275        // ---- MIDDLE STAGES: RX boundary s-1 -> range -> TX boundary s ----
6276        for s in 1..n_st - 1 {
6277            let _st = rt.enter(s);
6278            let es = rt.engine(s, e);
6279            let pos_d = stage_pos(es)?;
6280            let x = rt.rx(s - 1, slot, payload)?;
6281            let x = self.verify_layers(
6282                es,
6283                x,
6284                fence[s],
6285                fence[s + 1],
6286                &pos_d,
6287                pos0,
6288                t,
6289                cache,
6290                ckpt.as_deref_mut(),
6291                stream,
6292                None,
6293            )?;
6294            slot = if pipelined {
6295                rt.tx_pipelined(s, &x, payload)?
6296            } else {
6297                rt.tx(s, &x, payload)?
6298            };
6299        }
6300
6301        // ---- LAST STAGE: RX + final range + output_norm + lm head ----
6302        let _stl = rt.enter(n_st - 1);
6303        let el = rt.engine(n_st - 1, e);
6304        let pos_d = stage_pos(el)?;
6305        let rx_started = std::time::Instant::now();
6306        let x = rt.rx(n_st - 2, slot, payload)?;
6307        if pp_anatomy {
6308            el.stream().synchronize()?;
6309            rx_ms = rx_started.elapsed().as_secs_f64() * 1e3;
6310        }
6311        enqueue_spec_pipe_trace_marker(&el.stream(), trace.as_ref(), "S1", "start", Some(slot))?;
6312        let stage1_started = std::time::Instant::now();
6313        let x = self.verify_layers(
6314            el,
6315            x,
6316            fence[n_st - 1],
6317            fence[n_st],
6318            &pos_d,
6319            pos0,
6320            t,
6321            cache,
6322            ckpt.as_deref_mut(),
6323            stream,
6324            None,
6325        )?;
6326
6327        let mut hn = vbuf(el, payload)?;
6328        let logits = if self.sliding_gated_moe_batch_program() {
6329            // The PP Step35 serving path uses rms_norm + matmul for B=1 as well as B>1.
6330            // Verify must not switch numeric class merely because the same session speculates.
6331            el.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
6332            el.matmul(&self.output, &hn, t)?
6333        } else {
6334            el.rms_norm_decode(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
6335            el.matmul_decode_exact(&self.output, &hn, t)?
6336        };
6337        enqueue_spec_pipe_trace_marker(&el.stream(), trace.as_ref(), "S1", "end", Some(slot))?;
6338        if pp_anatomy {
6339            el.stream().synchronize()?;
6340            stage1_ms = stage1_started.elapsed().as_secs_f64() * 1e3;
6341        }
6342        // EXIT PUBLICATION: both returned buffers are still being produced on the last stage's
6343        // stream. Order the caller's stream behind that work before the buffers escape this
6344        // scope (the 2026-08-06 same-device ppspec find: without it the caller's primary-stream
6345        // consumer read unwritten logits — nondeterministic, one-device-only, and it poisoned
6346        // the following arm's KV in the same process).
6347        if publish_to_caller {
6348            rt.publish_to(n_st - 1, &caller_stream)?;
6349        }
6350        if pp_anatomy {
6351            if publish_to_caller {
6352                caller_stream.synchronize()?;
6353            }
6354            eprintln!(
6355                "[spec-pp-anatomy] t={t} reverse={reverse_ms:.3}ms stage0={stage0_ms:.3}ms \
6356                 tx={tx_ms:.3}ms rx={rx_ms:.3}ms stage1-head={stage1_ms:.3}ms total={:.3}ms",
6357                pp_started.elapsed().as_secs_f64() * 1e3,
6358            );
6359        }
6360        // stream: the device pos counter owns position; host mirror reconciles at drain.
6361        if stream.is_none() {
6362            cache.pos += t;
6363        }
6364        Ok((logits, if spec_hpost() { hn } else { x }))
6365    }
6366
6367    /// Step3.5/Step3.7 verify trunk in the serving batched numeric class.
6368    ///
6369    /// `step35_decode_batch_layers` is now authoritative at every live serving width, including
6370    /// B=1 (lane/cx-b1fix). The older verify walk deliberately mirrored the eager T=1 class:
6371    /// it replayed `step35_decode_attn` per row and used the eager/decode-exact FFN dispatch.
6372    /// Those classes are individually stable, but a near-tie prompt can choose different greedy
6373    /// bytes when a request moves from batched plain serving into speculative verify. Run the
6374    /// same authoritative B=1 stage subgraph for each verify row here. Rows still advance
6375    /// layer-by-layer, so every layer sees the preceding verify rows in its attention cache while
6376    /// every norm/projection/FFN uses exactly the live serving dispatch.
6377    #[allow(clippy::too_many_arguments)]
6378    /// PRIME-BY-T-ROWS (MEMRA_PRIME_TROWS=1): prefill the prompt through the same-session
6379    /// t-row walk in 32-row chunks — every row runs the t=1 decode program bit-for-bit
6380    /// (the TOKENWISE-prime ORACLE class), so this door is exact against the exactness
6381    /// reference while replacing the host-canonical per-token prime. Requires the walk
6382    /// doors (MEMRA_SPEC_VERIFY_EAGER/TCOL); returns the prime contract trio.
6383    #[allow(clippy::type_complexity)]
6384    pub(crate) fn step35_prime_trows(
6385        &self,
6386        e: &Engine,
6387        tokens: &[u32],
6388        cache: &mut Cache,
6389    ) -> Result<Option<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
6390    {
6391        let dbg = std::env::var("MEMRA_SPEC_FA2_DEBUG").as_deref() == Ok("1");
6392        if !prime_trows_on() {
6393            return Ok(None);
6394        }
6395        if !self.uses_sliding_gated_moe_program()
6396            || cache.pos != 0
6397            || cache.dflash_taps.is_some()
6398            || !spec_verify_eager_on()
6399            || !spec_verify_tcol_on()
6400        {
6401            if dbg {
6402                eprintln!(
6403                    "[prime-trows] refuse: program={} pos={} taps={} eager={:?} tcol={:?}",
6404                    self.uses_sliding_gated_moe_program(),
6405                    cache.pos,
6406                    cache.dflash_taps.is_some(),
6407                    std::env::var("MEMRA_SPEC_VERIFY_EAGER").ok(),
6408                    std::env::var("MEMRA_SPEC_VERIFY_TCOL").ok()
6409                );
6410            }
6411            return Ok(None);
6412        }
6413        let n_embd = self.cfg.n_embd as usize;
6414        let n_layers = self.layers.len();
6415        let t_total = tokens.len();
6416        let Some(embd_gpu) = self.embd_gpu_try(e) else {
6417            if dbg {
6418                eprintln!("[prime-trows] refuse: no device embed table");
6419            }
6420            return Ok(None);
6421        };
6422        let embd_qtype = match self.embd.ggml_type {
6423            memra_gguf::GgmlType::BF16 => crate::QT_BF16,
6424            memra_gguf::GgmlType::Q8_0 => crate::QT_Q8_0,
6425            other => {
6426                if dbg {
6427                    eprintln!("[prime-trows] refuse: embed dtype {other:?}");
6428                }
6429                return Ok(None);
6430            }
6431        };
6432        let embd_row_bytes = self.embd.raw.len() / self.cfg.n_vocab as usize;
6433        // Chunk plan: 32-row chunks; a 1-token tail folds into the previous chunk
6434        // (the walk floor is t >= 2).
6435        let mut bounds = Vec::new();
6436        let mut start = 0usize;
6437        while start < t_total {
6438            let mut end = (start + 32).min(t_total);
6439            if t_total - end == 1 {
6440                end -= 1;
6441            }
6442            bounds.push((start, end));
6443            start = end;
6444        }
6445        if bounds.iter().any(|(a, b)| b - a < 2) {
6446            return Ok(None); // degenerate short prompt keeps the ordinary prime
6447        }
6448        let mut hiddens = e.uninit(t_total * n_embd)?;
6449        let mut last: Option<CudaSlice<f32>> = None;
6450        for &(a, b) in &bounds {
6451            let tc = b - a;
6452            let tok_d = e.stream().clone_htod(&tokens[a..b])?;
6453            let x =
6454                e.embed_gather_device_td(embd_gpu, &tok_d, tc, n_embd, embd_qtype, embd_row_bytes)?;
6455            let out = self.step35_verify_batch_layers(e, x, 0, n_layers, a, tc, cache)?;
6456            e.copy_into(&mut hiddens, a * n_embd, &out, tc * n_embd)?;
6457            if b == t_total {
6458                let mut h = e.uninit(n_embd)?;
6459                e.dtod_copy_view(&out.slice((tc - 1) * n_embd..tc * n_embd), &mut h)?;
6460                last = Some(h);
6461            }
6462        }
6463        let h_seed = last.expect("last chunk produced the seed row");
6464        let mut hn = e.uninit(n_embd)?;
6465        e.rms_norm_decode(
6466            &h_seed,
6467            self.output_norm.float_data(),
6468            &mut hn,
6469            n_embd,
6470            1,
6471            self.cfg.rms_eps,
6472        )?;
6473        let logits_d = e.matmul_decode_exact(&self.output, &hn, 1)?;
6474        let logits = e.dtoh(&logits_d)?;
6475        cache.pos = t_total;
6476        Ok(Some((logits, h_seed, hiddens)))
6477    }
6478
6479    fn step35_verify_batch_layers(
6480        &self,
6481        e: &Engine,
6482        mut x: CudaSlice<f32>,
6483        lo: usize,
6484        hi: usize,
6485        pos0: usize,
6486        t: usize,
6487        cache: &mut Cache,
6488    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6489        let n_embd = self.cfg.n_embd as usize;
6490        if !self.uses_sliding_gated_moe_program() {
6491            return Err(
6492                "serving-class verify requires sliding-gated-MoE canonical operations".into(),
6493            );
6494        }
6495        // SERVING-CLASS VERIFY (MEMRA_SPEC_VERIFY_EAGER=1, step37 MTP bring-up): each verify
6496        // column rides decode_layers_eager — the EXACT t=1 program live serving runs (all TP2
6497        // doors) — row-outer, so row r's appends land before row r+1 attends: bit-equal to
6498        // plain greedy by construction. Only the unsplit full-range walk qualifies; PP splits
6499        // and the tap path keep the batch-layer class.
6500        static VE: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
6501        let eager_verify =
6502            *VE.get_or_init(spec_verify_eager_on) && lo == 0 && hi == self.layers.len();
6503        if eager_verify {
6504            // T-COLUMN LAYER-OUTER WALK (MEMRA_SPEC_VERIFY_TCOL=1): per layer, one t-grid
6505            // attn norm + ONE weight-amortized QKV(+gate) over all T columns, then each
6506            // column runs the UNMODIFIED t=1 attention program via the col-select door and
6507            // the ordinary residual/FFN body. Values per column are bit-equal to the
6508            // row-outer walk: rms over the materialized residual == the fused add+norm
6509            // (kernel_check identity), the tcol kernel's per-column FP order == the t=1
6510            // kernel, and every downstream op IS the t=1 program.
6511            static TCOL: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
6512            let tcol = *TCOL.get_or_init(spec_verify_tcol_on);
6513            // T > 32 (prefill-class): run the SAME walk in 32-row chunks — each chunk's
6514            // rows are the t=1 program bit-for-bit and the rope pass advances the cache,
6515            // so a chunked call is value-identical to the row-outer loop it replaces.
6516            static TROWS_PREFILL: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
6517            // MEMRA_STEP_GEMM_PRIME outranks the walk: with the grouped GEMM prime armed, the
6518            // t-row walk defers so the batch path (GEMM trunk + grouped MoE) takes the prompt —
6519            // flag precedence between two existing doors, not a new flag. Without this, both
6520            // doors ON meant the walk still won and the GEMM prime needed PRIME_TROWS=0 by hand.
6521            let trows_prefill =
6522                *TROWS_PREFILL.get_or_init(|| prime_trows_on() && !crate::step_gemm_prime_on());
6523            // MEMRA_PRIME_TROWS_T=<w>: chunk width, default 8 = the REAL cap of this walk.
6524            // The workspace slabs go to 32 rows, but `matvec_bf16_qkvg_tcol_into` refuses
6525            // t > 8 (compile-time-T twins exist for 2/4/8 only; the runtime-t kernel spills
6526            // its accumulators to local memory), so a wider chunk fails the request with
6527            // "matvec_bf16_qkvg_tcol geometry" — which is exactly how the first server-path
6528            // TROWS arm died. Measured at 193 tokens: w=8 2.459 s, w=4 2.574 s.
6529            static TROWS_W: std::sync::OnceLock<Result<usize, String>> = std::sync::OnceLock::new();
6530            let trows_w = match TROWS_W.get_or_init(|| {
6531                let value = std::env::var("MEMRA_PRIME_TROWS_T").ok();
6532                parse_prime_trows_width(value.as_deref())
6533            }) {
6534                Ok(width) => *width,
6535                Err(err) => return Err(err.clone().into()),
6536            };
6537            if tcol && trows_prefill && t > trows_w {
6538                // One-time engagement receipt: without it a prefill gate cannot tell a
6539                // chunked walk from the row-outer fallback it is supposed to replace
6540                // (the first PRIME_TROWS gate passed vacuously on exactly that).
6541                static SEEN: std::sync::atomic::AtomicBool =
6542                    std::sync::atomic::AtomicBool::new(false);
6543                if !SEEN.swap(true, std::sync::atomic::Ordering::Relaxed) {
6544                    eprintln!(
6545                        "[prime-trows] ENGAGED t={t} width={trows_w} chunks={} layers={}..{}",
6546                        t.div_ceil(trows_w),
6547                        lo,
6548                        hi
6549                    );
6550                }
6551                let mut out = e.uninit(t * n_embd)?;
6552                let mut start = 0usize;
6553                while start < t {
6554                    let mut end = (start + trows_w).min(t);
6555                    if t - end == 1 {
6556                        end -= 1;
6557                    }
6558                    let tc = end - start;
6559                    let mut xc = e.uninit(tc * n_embd)?;
6560                    e.dtod_copy_view(&x.slice(start * n_embd..end * n_embd), &mut xc)?;
6561                    let oc =
6562                        self.step35_verify_batch_layers(e, xc, lo, hi, pos0 + start, tc, cache)?;
6563                    e.copy_into(&mut out, start * n_embd, &oc, tc * n_embd)?;
6564                    start = end;
6565                }
6566                return Ok(out);
6567            }
6568            if tcol && t >= 2 && t <= 32 {
6569                // MEMRA_TCOL_PROF=1: synchronized per-segment wall profile of the walk
6570                // (norm+QKV precompute / per-col attention / per-col residual+FFN). The
6571                // syncs serialize the stream, so the split is for TARGETING amortization
6572                // work only — never a perf claim.
6573                static PROF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
6574                let prof =
6575                    *PROF.get_or_init(|| std::env::var("MEMRA_TCOL_PROF").as_deref() == Ok("1"));
6576                let mut prof_ms = [0f64; 3];
6577                let eps = self.cfg.rms_eps;
6578                let mut x_t = x;
6579                let mut h_t = e.uninit(t * n_embd)?;
6580                let mut h_row = e.uninit(n_embd)?; // real row: the non-dcw fallback reads it
6581                // Per-column pos buffers hoisted out of the layer loop (a per-col-per-layer
6582                // pageable htod was an in-stream engine turnaround x t x 45).
6583                let mut pos_rows = Vec::with_capacity(t);
6584                for r in 0..t {
6585                    pos_rows.push(e.htod_i32(&[(pos0 + r) as i32])?);
6586                }
6587                let mut ok = true;
6588                // MEMRA_TCOL_OPROJ=1: defer each column's o_proj — the finish seam
6589                // stashes `gated` instead of joining per column; one b4_tcol per rank +
6590                // one slab join produce every column's `mixed` after the attention pass.
6591                // Bit-exact per column (t=1 b4 program per column; elementwise join).
6592                // MEMRA_TCOL_FFN=1: today this only IMPLIES the o_proj defer above. Its
6593                // named feature, the two-column device-routed FFN sweep, rode the
6594                // slot-major v2 TP banks and was REMOVED with the MEMRA_NVFP4_BANK_V2 door
6595                // (2026-08-29, research/step37-bankv2-removal-20260829): the v2 layout
6596                // changed generated text in serving. The flag itself stays because it is
6597                // family-armed in the step37 serving defaults and killing it here would
6598                // silently drop the o_proj defer from the qualified serving shape.
6599                static FFN2: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
6600                let ffn_batch = *FFN2.get_or_init(tcol_ffn_on);
6601                let oproj_batch = crate::tp::tcol_oproj_on() || ffn_batch;
6602                // MEMRA_SPEC_FA2=1 (T=2 only): eligible layers defer BOTH columns' fa —
6603                // the per-column pass norms/ropes/appends and stashes q+gate, then one
6604                // shared-KV fa_decode_dcw2 per rank + the o_proj join produce the
6605                // [2, o_out] mixed slab. The precheck runs before arming (stashing is
6606                // unrecoverable); ineligible/boundary layers run the ordinary program.
6607                let fa2 = crate::tp::spec_fa2_on() && t <= 32;
6608                let mut mixed_row = e.uninit(n_embd)?;
6609                let mut pos_staged = false;
6610                for il in lo..hi {
6611                    let layer = &self.layers[il];
6612                    // BEFORE this layer touches its planes: is the history it is about to
6613                    // attend already poisoned? Global (non-ring) layers only, which are the
6614                    // ones the level-2 bitmap implicates.
6615                    if kv_plane_scan_on() && self.step35_geom(il).window.is_none() {
6616                        if let Some(distributed) = cache.tp_kv[il].as_ref() {
6617                            scan_kv_plane(e, distributed, il, pos0)?;
6618                        }
6619                    }
6620                    let fa2_layer = fa2 && self.step35_fa_rows_precheck(cache, il, pos0, t)?;
6621                    let mut seg = std::time::Instant::now();
6622                    e.rms_norm(&x_t, layer.attn_norm.float_data(), &mut h_t, n_embd, t, eps)?;
6623                    if !self.step35_verify_qkv_precompute(e, il, &h_t, t)? {
6624                        ok = false;
6625                        break;
6626                    }
6627                    // FULL t-row attention pass (rope/append + fa + combine + o_proj in
6628                    // 3 launches/rank): same-session rows, slot = len-base+r, one len
6629                    // advance by t. Host cache bookkeeping mirrors the per-column tail.
6630                    if fa2_layer {
6631                        if let Some(mixed_t) =
6632                            self.step35_verify_rope_fa_pass(e, il, cache, pos0, t, !pos_staged)?
6633                        {
6634                            pos_staged = true;
6635                            {
6636                                let tp_kv = cache.tp_kv[il]
6637                                    .as_mut()
6638                                    .expect("precheck verified the distributed cache");
6639                                let transaction = tp_kv.begin_transaction()?;
6640                                let crate::hybrid::Mixer::Full(fa) = &layer.mixer else {
6641                                    return Err("verify rope pass expects full attention".into());
6642                                };
6643                                let tp = fa
6644                                    .step_tp_qkv
6645                                    .as_ref()
6646                                    .ok_or("verify rope pass lost its TP state")?;
6647                                let empty: [CudaSlice<f32>; 0] = [];
6648                                tp.runtime.append_tp_kv_transaction_inner(
6649                                    tp_kv,
6650                                    transaction,
6651                                    &empty,
6652                                    &empty,
6653                                    t,
6654                                    true,
6655                                )?;
6656                                tp.runtime.commit_tp_kv_transaction_external(
6657                                    tp_kv,
6658                                    transaction,
6659                                    t,
6660                                )?;
6661                                if let Some(local) = cache.kv[il].as_mut() {
6662                                    local.len = pos0 + t;
6663                                    if !crate::tp::len_mirror_lazy_on() {
6664                                        e.set_i32_one(&mut local.len_d, local.len as i32)?;
6665                                    }
6666                                }
6667                            }
6668                            if prof {
6669                                e.stream().synchronize()?;
6670                                prof_ms[1] += seg.elapsed().as_secs_f64() * 1e3;
6671                                seg = std::time::Instant::now();
6672                            }
6673                            let o_out = mixed_t.len() / t;
6674                            let mut next = e.uninit(t * n_embd)?;
6675                            {
6676                                for r in 0..t {
6677                                    e.dtod_copy_view(
6678                                        &mixed_t.slice(r * o_out..(r + 1) * o_out),
6679                                        &mut mixed_row,
6680                                    )?;
6681                                    let mut x_row = e.uninit(n_embd)?;
6682                                    e.dtod_copy_view(
6683                                        &x_t.slice(r * n_embd..(r + 1) * n_embd),
6684                                        &mut x_row,
6685                                    )?;
6686                                    let (x1, ffn_out) = self.residual_norm_ffn(
6687                                        e, layer, &x_row, &mixed_row, n_embd, il, eps,
6688                                    )?;
6689                                    let mut x2 = e.uninit(n_embd)?;
6690                                    e.add(&x1, &ffn_out, &mut x2, n_embd)?;
6691                                    e.dtod_copy_into(&x2, &mut next, r * n_embd)?;
6692                                }
6693                            }
6694                            if prof {
6695                                e.stream().synchronize()?;
6696                                prof_ms[2] += seg.elapsed().as_secs_f64() * 1e3;
6697                            }
6698                            x_t = next;
6699                            if spec_nan_scan() {
6700                                // The scan MUST sit on this arm too. It used to live only on
6701                                // the non-fused tail, so a fused layer's poison was first
6702                                // reported by the next non-fused layer.
6703                                verify_arm_receipt(
6704                                    "fused",
6705                                    il,
6706                                    pos0,
6707                                    t,
6708                                    cache.tp_kv[il].as_ref().map(|d| d.staged_len()),
6709                                );
6710                                nan_scan_rows(
6711                                    e,
6712                                    &x_t,
6713                                    t,
6714                                    n_embd,
6715                                    &format!("tcol layer {il} pos0={pos0} arm=fused"),
6716                                )?;
6717                            }
6718                            continue;
6719                        }
6720                    }
6721                    if prof {
6722                        e.stream().synchronize()?;
6723                        prof_ms[0] += seg.elapsed().as_secs_f64() * 1e3;
6724                        seg = std::time::Instant::now();
6725                    }
6726                    let mut next = e.uninit(t * n_embd)?;
6727                    // Columns whose o_proj was deferred (their FFN runs after the join).
6728                    // A NON-deferred column's FFN must run INSIDE the column loop: the
6729                    // oproj-tail handoff is a single cell that the same column's
6730                    // residual_norm_ffn consumes before the next column's finish.
6731                    let mut deferred: Vec<usize> = Vec::new();
6732                    let mut fa2_deferred: Vec<usize> = Vec::new();
6733                    let mut ffn_col =
6734                        |r: usize,
6735                         mixed: &CudaSlice<f32>,
6736                         next: &mut CudaSlice<f32>|
6737                         -> Result<(), Box<dyn std::error::Error>> {
6738                            let mut x_row = e.uninit(n_embd)?;
6739                            e.dtod_copy_view(&x_t.slice(r * n_embd..(r + 1) * n_embd), &mut x_row)?;
6740                            let (x1, ffn_out) =
6741                                self.residual_norm_ffn(e, layer, &x_row, mixed, n_embd, il, eps)?;
6742                            if spec_nan_scan_level() >= 2 {
6743                                nan_scan_rows(
6744                                    e,
6745                                    &ffn_out,
6746                                    1,
6747                                    n_embd,
6748                                    &format!("tcol layer {il} col {r} per-column FFN out"),
6749                                )?;
6750                            }
6751                            let mut x2 = e.uninit(n_embd)?;
6752                            e.add(&x1, &ffn_out, &mut x2, n_embd)?;
6753                            e.dtod_copy_into(&x2, next, r * n_embd)?;
6754                            Ok(())
6755                        };
6756                    for r in 0..t {
6757                        e.dtod_copy_view(&h_t.slice(r * n_embd..(r + 1) * n_embd), &mut h_row)?;
6758                        let row_pos = &pos_rows[r];
6759                        crate::tp::set_verify_tcol(Some(r));
6760                        if fa2_layer {
6761                            crate::tp::set_spec_fa2_defer(Some(r));
6762                        } else if oproj_batch {
6763                            crate::tp::set_tcol_oproj_defer(Some(r));
6764                        }
6765                        let mixed = match &layer.mixer {
6766                            crate::hybrid::Mixer::Full(fa) => {
6767                                self.full_attn_decode(e, fa, &h_row, row_pos, pos0 + r, cache, il)
6768                            }
6769                            _ => Err("step35 verify expects full attention".into()),
6770                        };
6771                        crate::tp::set_verify_tcol(None);
6772                        crate::tp::set_spec_fa2_defer(None);
6773                        crate::tp::set_tcol_oproj_defer(None);
6774                        let mixed = mixed?;
6775                        if fa2_layer && crate::tp::take_spec_fa2_stashed() {
6776                            fa2_deferred.push(r);
6777                        } else if oproj_batch && crate::tp::take_tcol_oproj_stashed() {
6778                            deferred.push(r);
6779                        } else {
6780                            if spec_nan_scan_level() >= 2 {
6781                                let cols = mixed.len();
6782                                nan_scan_rows(
6783                                    e,
6784                                    &mixed,
6785                                    1,
6786                                    cols,
6787                                    &format!("tcol layer {il} col {r} per-column ATTN out"),
6788                                )?;
6789                            }
6790                            ffn_col(r, &mixed, &mut next)?;
6791                        }
6792                    }
6793                    if !fa2_deferred.is_empty() && fa2_deferred.len() != t {
6794                        // The precheck guarantees both columns stash or neither; a strict
6795                        // subset means a column's output was never produced anywhere.
6796                        return Err("spec fa2 stash engaged for a subset of columns".into());
6797                    }
6798                    if prof {
6799                        e.stream().synchronize()?;
6800                        prof_ms[1] += seg.elapsed().as_secs_f64() * 1e3;
6801                        seg = std::time::Instant::now();
6802                    }
6803                    if !fa2_deferred.is_empty() {
6804                        deferred = fa2_deferred;
6805                    }
6806                    if !deferred.is_empty() {
6807                        let mixed_t = if fa2_layer {
6808                            self.step35_verify_fa_rows_join(e, il, cache, pos0, t)?
6809                        } else {
6810                            self.step35_verify_oproj_tcol(e, il, t)?
6811                        };
6812                        let o_out = mixed_t.len() / t;
6813                        if spec_nan_scan_level() >= 2 {
6814                            nan_scan_rows(
6815                                e,
6816                                &mixed_t,
6817                                t,
6818                                o_out,
6819                                &format!("tcol layer {il} JOINED attn over deferred cols"),
6820                            )?;
6821                        }
6822                        // Batched t=2 residual+MoE: one t-grid add_rms_norm (per-row
6823                        // program == t=1; bit-identical to the oproj-tail join per the
6824                        // M2 verbatim-program contract) feeding the two-column routed
6825                        // sweep. Ineligible layers (dense FFN, non-nvfp4) fall through
6826                        // to the per-column body.
6827                        {
6828                            for &r in &deferred {
6829                                e.dtod_copy_view(
6830                                    &mixed_t.slice(r * o_out..(r + 1) * o_out),
6831                                    &mut mixed_row,
6832                                )?;
6833                                ffn_col(r, &mixed_row, &mut next)?;
6834                            }
6835                        }
6836                    }
6837                    if prof {
6838                        e.stream().synchronize()?;
6839                        prof_ms[2] += seg.elapsed().as_secs_f64() * 1e3;
6840                    }
6841                    drop(ffn_col);
6842                    x_t = next;
6843                    if spec_nan_scan() {
6844                        verify_arm_receipt(
6845                            if fa2_layer { "join" } else { "percol" },
6846                            il,
6847                            pos0,
6848                            t,
6849                            cache.tp_kv[il].as_ref().map(|d| d.staged_len()),
6850                        );
6851                        nan_scan_rows(
6852                            e,
6853                            &x_t,
6854                            t,
6855                            n_embd,
6856                            &format!(
6857                                "tcol layer {il} pos0={pos0} arm={}",
6858                                if fa2_layer { "join" } else { "percol" }
6859                            ),
6860                        )?;
6861                    }
6862                }
6863                if prof {
6864                    eprintln!(
6865                        "[tcol-prof] t={t} norm+qkv={:.3}ms attn={:.3}ms ffn={:.3}ms",
6866                        prof_ms[0], prof_ms[1], prof_ms[2]
6867                    );
6868                }
6869                if ok {
6870                    return Ok(x_t);
6871                }
6872                // fall through to the row-outer walk on ineligible layers
6873                x = x_t;
6874            }
6875            let mut next = e.uninit(t * n_embd)?;
6876            let scan = spec_nan_scan();
6877            for r in 0..t {
6878                let mut row = e.uninit(n_embd)?;
6879                e.dtod_copy_view(&x.slice(r * n_embd..(r + 1) * n_embd), &mut row)?;
6880                let row_pos = e.htod_i32(&[(pos0 + r) as i32])?;
6881                let out = if scan {
6882                    // Diagnostic arm: the same range walked one layer at a time so the first
6883                    // poisoned layer names itself. `decode_layers_eager(lo, hi)` is range-scoped
6884                    // and executes its trailing residual add, so a per-layer chain is the same
6885                    // program with the cross-layer add+norm fusion unrolled.
6886                    nan_scan_rows(
6887                        e,
6888                        &row,
6889                        1,
6890                        n_embd,
6891                        &format!("embed row r={r} pos={}", pos0 + r),
6892                    )?;
6893                    let mut acc = row;
6894                    for il in lo..hi {
6895                        acc = self.decode_layers_eager(
6896                            e,
6897                            acc,
6898                            il,
6899                            il + 1,
6900                            &row_pos,
6901                            pos0 + r,
6902                            cache,
6903                        )?;
6904                        nan_scan_rows(
6905                            e,
6906                            &acc,
6907                            1,
6908                            n_embd,
6909                            &format!("row-outer layer {il} r={r} pos={}", pos0 + r),
6910                        )?;
6911                    }
6912                    acc
6913                } else {
6914                    self.decode_layers_eager(e, row, lo, hi, &row_pos, pos0 + r, cache)?
6915                };
6916                e.dtod_copy_into(&out, &mut next, r * n_embd)?;
6917            }
6918            // dflash taps are NOT produced on this arm (they need per-layer hiddens the
6919            // row-outer walk does not materialize); the door is a step37 MTP bring-up
6920            // surface where taps are unused.
6921            return Ok(next);
6922        }
6923        let mut ph_last = std::time::Instant::now();
6924        for il in lo..hi {
6925            let mut next = e.uninit(t * n_embd)?;
6926            for r in 0..t {
6927                let mut row = e.uninit(n_embd)?;
6928                e.dtod_copy_view(&x.slice(r * n_embd..(r + 1) * n_embd), &mut row)?;
6929                // The caller owns this verify's position. During controller overlap, cache.pos
6930                // still describes generation N while this stage-0 walk belongs to N+1.
6931                let row_pos = e.htod_i32(&[(pos0 + r) as i32])?;
6932                let mut one = [&mut *cache];
6933                let out = self.step35_decode_batch_layers(
6934                    e,
6935                    row,
6936                    &mut one,
6937                    &[(pos0 + r) as i32],
6938                    &row_pos,
6939                    il,
6940                    il + 1,
6941                    &mut ph_last,
6942                )?;
6943                e.dtod_copy_into(&out, &mut next, r * n_embd)?;
6944            }
6945            self.dflash_tap(e, cache, il, &next, t)?;
6946            x = next;
6947            if spec_nan_scan() {
6948                nan_scan_rows(e, &x, t, n_embd, &format!("batch-layer {il} pos0={pos0}"))?;
6949            }
6950        }
6951        Ok(x)
6952    }
6953
6954    /// DSpark drafter verify (lane/dspark-q38-recover): one t-row forward through the
6955    /// SERVING-CLASS verify funnel (`decode_step_t_core_stream` — the same numeric class
6956    /// MTP verify rides, GDN state advanced in place), returning per-row argmax tokens.
6957    /// Advances `cache.pos += t`; the caller owns snapshot/rollback (block acceptance is
6958    /// prefix-keep, not all-or-nothing).
6959    pub(crate) fn dspark_verify_t_am(
6960        &self,
6961        e: &Engine,
6962        tokens: &[u32],
6963        pos0: usize,
6964        cache: &mut Cache,
6965    ) -> Result<Vec<u32>, Box<dyn std::error::Error>> {
6966        let (logits, _hn) = self.decode_step_t_core_stream(
6967            e, tokens, pos0, cache, None, None, None, None, None, None,
6968        )?;
6969        let t = tokens.len();
6970        let v = self.output.out_features();
6971        let mut am_d = e.stream().alloc_zeros::<u32>(t)?;
6972        for r in 0..t {
6973            e.argmax_token_device_col(&logits, r, v, &mut am_d, r)?;
6974        }
6975        Ok(e.dtoh_u32(&am_d)?)
6976    }
6977
6978    /// DSpark verify returning the RAW verify logits [t, n_vocab] (device-resident) instead
6979    /// of per-row argmaxes — the sampled-admission arm's input (rejection-sampling accept
6980    /// gathers filtered p from these columns; lane/dspark-sampled-admission-20260820). Same
6981    /// forward as `dspark_verify_t_am`; the greedy arm keeps its argmax wrapper untouched.
6982    pub(crate) fn dspark_verify_t_logits(
6983        &self,
6984        e: &Engine,
6985        tokens: &[u32],
6986        pos0: usize,
6987        cache: &mut Cache,
6988    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6989        let (logits, _hn) = self.decode_step_t_core_stream(
6990            e, tokens, pos0, cache, None, None, None, None, None, None,
6991        )?;
6992        Ok(logits)
6993    }
6994
6995    /// DSpark verify with the MTP column-stash armed: identical forward to
6996    /// `dspark_verify_t_am`, but fills a `VerifyCkpt` so a partial accept can restore
6997    /// column state directly (`dspark_commit_prefix`) instead of snapshot-replay.
6998    /// The ckpt type is opaque outside spec.rs (newtype) — dflash.rs threads it through.
6999    pub(crate) fn dspark_verify_t_am_ckpt(
7000        &self,
7001        e: &Engine,
7002        tokens: &[u32],
7003        pos0: usize,
7004        cache: &mut Cache,
7005    ) -> Result<(Vec<u32>, DsparkVerifyCkpt), Box<dyn std::error::Error>> {
7006        let mut ck = VerifyCkpt::new(self.layers.len());
7007        let (logits, _hn) = self.decode_step_t_core_stream(
7008            e,
7009            tokens,
7010            pos0,
7011            cache,
7012            None,
7013            Some(&mut ck),
7014            None,
7015            None,
7016            None,
7017            None,
7018        )?;
7019        let t = tokens.len();
7020        let v = self.output.out_features();
7021        let mut am_d = e.stream().alloc_zeros::<u32>(t)?;
7022        for r in 0..t {
7023            e.argmax_token_device_col(&logits, r, v, &mut am_d, r)?;
7024        }
7025        Ok((e.dtoh_u32(&am_d)?, DsparkVerifyCkpt(ck)))
7026    }
7027
7028    /// Engine-bundle slice 2: `dspark_verify_t_am_ckpt` with DEVICE tokens and NO readback.
7029    /// The verify tokens are the round's `chain_d` (cand layout: [anchor, drafts...]); the
7030    /// embed gathers its first `t` entries on-device (`embed_gather_u32_t` — bit-identical
7031    /// rows to the host arm), so the host never blocks on the draft chain before dispatching
7032    /// verify. Returns the device per-row argmax buffer; the caller merges its readback with
7033    /// the chain's into ONE sync. Forward, ckpt fill and argmax walk are `_ckpt` verbatim.
7034    pub(crate) fn dspark_verify_t_am_ckpt_dev(
7035        &self,
7036        e: &Engine,
7037        vtok: &CudaSlice<u32>,
7038        t: usize,
7039        pos0: usize,
7040        cache: &mut Cache,
7041        embd_dev: (&CudaSlice<u8>, i32, usize),
7042        graphs: Option<&mut DsparkVerifyGraphs>,
7043    ) -> Result<(CudaSlice<u32>, DsparkVerifyCkpt), Box<dyn std::error::Error>> {
7044        debug_assert!(
7045            vtok.len() >= t,
7046            "verify window exceeds the device token buffer"
7047        );
7048        // The slab flag is a per-round statement: clear it here so a verify that never
7049        // reaches the graphs door (rowwise env, a non-tparallel arm) cannot leave a
7050        // stale `true` steering the commit at slabs the round never wrote.
7051        let mut graphs = graphs;
7052        if let Some(g) = graphs.as_deref_mut() {
7053            g.round_slab = false;
7054        }
7055        let mut ck = VerifyCkpt::new(self.layers.len());
7056        // Dummy host tokens size the funnel; the embed reads `vtok` (the round-stream
7057        // arm's established pattern — spec.rs stream-mode verify does the same).
7058        let dummy = vec![0u32; t];
7059        let (logits, _hn) = self.decode_step_t_core_stream(
7060            e,
7061            &dummy,
7062            pos0,
7063            cache,
7064            Some(embd_dev),
7065            Some(&mut ck),
7066            None,
7067            None,
7068            Some(vtok),
7069            graphs,
7070        )?;
7071        let v = self.output.out_features();
7072        let mut am_d = e.stream().alloc_zeros::<u32>(t)?;
7073        for r in 0..t {
7074            e.argmax_token_device_col(&logits, r, v, &mut am_d, r)?;
7075        }
7076        Ok((am_d, DsparkVerifyCkpt(ck)))
7077    }
7078
7079    /// Ckpt-armed twin of [`Self::dspark_verify_t_logits`] (sampled-admission arm).
7080    pub(crate) fn dspark_verify_t_logits_ckpt(
7081        &self,
7082        e: &Engine,
7083        tokens: &[u32],
7084        pos0: usize,
7085        cache: &mut Cache,
7086    ) -> Result<(CudaSlice<f32>, DsparkVerifyCkpt), Box<dyn std::error::Error>> {
7087        let mut ck = VerifyCkpt::new(self.layers.len());
7088        let (logits, _hn) = self.decode_step_t_core_stream(
7089            e,
7090            tokens,
7091            pos0,
7092            cache,
7093            None,
7094            Some(&mut ck),
7095            None,
7096            None,
7097            None,
7098            None,
7099        )?;
7100        Ok((logits, DsparkVerifyCkpt(ck)))
7101    }
7102
7103    /// Restore the round to `keep` accepted columns from the verify stash: KV lens and
7104    /// pos from the pre-verify snapshot + keep, GDN conv/ssm from the stashed column
7105    /// state — no replay forward. The exact `commit_verified_prefix` the MTP path ships.
7106    pub(crate) fn dspark_commit_prefix(
7107        &self,
7108        e: &Engine,
7109        cache: &mut Cache,
7110        snap: &crate::cache::CacheSnapshot,
7111        ckpt: &DsparkVerifyCkpt,
7112        keep: usize,
7113    ) -> Result<(), Box<dyn std::error::Error>> {
7114        self.commit_verified_prefix(e, cache, snap, &ckpt.0, keep, false, None)
7115    }
7116
7117    /// Slice-3 commit twin: restore to `keep` accepted columns when the round's linear
7118    /// column stash lives in the graphs ctx's persistent slabs (`DsparkVerifyGraphs`) —
7119    /// the cols arm's exact semantics (KV lens + pos from the snapshot, GDN conv/ssm
7120    /// from the stash of column keep-1), slab-addressed and batched into two copy
7121    /// launches. `MEMRA_STATE_COPY_BATCH=0` falls back to per-layer view copies.
7122    pub(crate) fn dspark_commit_prefix_slab(
7123        &self,
7124        e: &Engine,
7125        cache: &mut Cache,
7126        snap: &crate::cache::CacheSnapshot,
7127        ctx: &DsparkVerifyGraphs,
7128        keep: usize,
7129    ) -> Result<(), Box<dyn std::error::Error>> {
7130        use cudarc::driver::DevicePtr;
7131        debug_assert!(keep >= 1, "keep==0 rounds take the legacy rollback");
7132        let mut conv_src: Vec<u64> = Vec::new();
7133        let mut ssm_src: Vec<u64> = Vec::new();
7134        let mut conv_dst: Vec<u64> = Vec::new();
7135        let mut ssm_dst: Vec<u64> = Vec::new();
7136        for il in 0..self.layers.len() {
7137            if let (Some(kvl), Some(saved)) = (cache.kv[il].as_mut(), snap.kv_len[il]) {
7138                kvl.len = saved + keep;
7139                e.set_i32_one(&mut kvl.len_d, kvl.len as i32)?;
7140            }
7141            if let Some(rl) = cache.recur[il].as_ref() {
7142                let (pc, ps, _cw, _sw) = ctx
7143                    .slab_row(e, il, keep - 1)
7144                    .ok_or("slab commit: linear layer missing from the graphs ctx")?;
7145                conv_src.push(pc);
7146                ssm_src.push(ps);
7147                let st = &e.gpu.stream();
7148                let (dc, _g0) = rl.conv_state.device_ptr(st);
7149                let (ds, _g1) = rl.ssm_state.device_ptr(st);
7150                conv_dst.push(dc as u64);
7151                ssm_dst.push(ds as u64);
7152            }
7153        }
7154        let n = conv_src.len();
7155        if n > 0 {
7156            if state_copy_batch_on() {
7157                let mut tt = vec![0u64; 2 * n];
7158                tt[..n].copy_from_slice(&conv_src);
7159                tt[n..].copy_from_slice(&conv_dst);
7160                let ct = e.htod_u64(&tt)?;
7161                tt[..n].copy_from_slice(&ssm_src);
7162                tt[n..].copy_from_slice(&ssm_dst);
7163                let st = e.htod_u64(&tt)?;
7164                e.copy_batch_uniform_f32(&ct, n, ctx.conv_words)?;
7165                e.copy_batch_uniform_f32(&st, n, ctx.ssm_words)?;
7166            } else {
7167                let (cw, sw) = (ctx.conv_words, ctx.ssm_words);
7168                let row = keep - 1;
7169                for il in 0..self.layers.len() {
7170                    let Some(rl) = cache.recur[il].as_mut() else {
7171                        continue;
7172                    };
7173                    let k = ctx.lin_pos[&il];
7174                    {
7175                        let sv = e.view(&ctx.stash_conv[k], (row + 1) * cw);
7176                        let win = sv.slice(row * cw..(row + 1) * cw);
7177                        e.copy_view_into(&mut rl.conv_state, 0, &win, cw)?;
7178                    }
7179                    {
7180                        let sv = e.view(&ctx.stash_ssm[k], (row + 1) * sw);
7181                        let win = sv.slice(row * sw..(row + 1) * sw);
7182                        e.copy_view_into(&mut rl.ssm_state, 0, &win, sw)?;
7183                    }
7184                }
7185            }
7186        }
7187        cache.pos = snap.pos + keep;
7188        Ok(())
7189    }
7190
7191    /// Qwen35-family verify trunk in the live serving numeric class.
7192    ///
7193    /// Serving intentionally keeps this architecture in the generic batched program even at
7194    /// B=1. The older verify walk used its own mirrored dispatch and can flip near-tie argmaxes.
7195    ///
7196    /// Two arms, one numeric class:
7197    /// - DENSE GDN (`DenseMlp`, t<=16): `qwen35_verify_tparallel` — the weight ops (norms,
7198    ///   projections, FFN) hoist to m=T through the exact-tier batched kernels whose per-row
7199    ///   program IS the m=1 program (`matmul_pre == fused2 per (tensor,row); _bN mmvq per-row
7200    ///   == m=1` — decode_batch.rs v2 note), while the state ops (conv ring, gdn scan, KV
7201    ///   append, fa decode) stay a per-row loop running the b_n=1 serving kernels with each
7202    ///   row's own t_kv-driven arm pick (the straddle law: every row executes the exact
7203    ///   program its isolated serving step would). One weight read per layer per round
7204    ///   instead of T — this is what makes MTP profitable in the exact class (the per-row
7205    ///   walk measured verify(K+1) ~= (K+1) plain steps: 69 -> 44 tok/s served, 2026-08-15).
7206    /// - MoE / t>16 / `MEMRA_SPEC_VERIFY_ROWWISE=1`: the per-row replay of the authoritative
7207    ///   serving layer body, preserving single-session autoregressive cache order (the
7208    ///   correctness reference; also the rollback seam for the t-parallel arm).
7209    ///
7210    /// Bit-identity of the t-parallel arm vs the rowwise arm is gated by spec-serve-gate
7211    /// (zero differing logits at T=1..4, K arms) + the 8-prompt ON/OFF canary before ship.
7212    #[allow(clippy::too_many_arguments)]
7213    fn qwen35_verify_batch_layers(
7214        &self,
7215        e: &Engine,
7216        x: CudaSlice<f32>,
7217        lo: usize,
7218        hi: usize,
7219        pos0: usize,
7220        t: usize,
7221        cache: &mut Cache,
7222        ckpt: Option<&mut VerifyCkpt>,
7223        stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
7224        graphs: Option<&mut DsparkVerifyGraphs>,
7225    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7226        // Qwen35Moe admitted 2026-08-20 (lane/draftcost-moe): the t-parallel arm already
7227        // carries the MoE FFN (`moe_ffn_il_zq8` at m=T) and the GDN per-row state loop; the
7228        // arch fence was a qualification gate, not a mechanism gap. Measured disease on the
7229        // 35B-A3B class: rowwise verify ~= 5.6 ms per drafted token (one full trunk step
7230        // each) — the same (K+1)-plain-steps wall the dense admission fixed on 2026-08-15.
7231        // Rollback seam unchanged: MEMRA_SPEC_VERIFY_ROWWISE=1.
7232        let rowwise = std::env::var("MEMRA_SPEC_VERIFY_ROWWISE").as_deref() == Ok("1")
7233            || !self.batched_serving_numeric_class()
7234            || t > 16;
7235        if rowwise {
7236            if stream.is_some() {
7237                // rowwise replays per row with host cache.pos — irreconcilable with a
7238                // device position counter. Burst callers must keep t <= 16 and the
7239                // ROWWISE env unset; refusing beats silently mispositioned rows.
7240                return Err("qwen35 rowwise verify has no ROUND-STREAM arm \
7241                            (t > 16 or MEMRA_SPEC_VERIFY_ROWWISE=1)"
7242                    .into());
7243            }
7244            self.qwen35_verify_rowwise(e, x, lo, hi, pos0, t, cache, ckpt)
7245        } else {
7246            self.qwen35_verify_tparallel(e, x, lo, hi, pos0, t, cache, ckpt, stream, graphs)
7247        }
7248    }
7249
7250    /// The per-row correctness reference: replay each verify row through the authoritative
7251    /// serving layer body (`decode_batch_layers` at b_n=1). T full weight reads per layer.
7252    #[allow(clippy::too_many_arguments)]
7253    fn qwen35_verify_rowwise(
7254        &self,
7255        e: &Engine,
7256        mut x: CudaSlice<f32>,
7257        lo: usize,
7258        hi: usize,
7259        pos0: usize,
7260        t: usize,
7261        cache: &mut Cache,
7262        mut ckpt: Option<&mut VerifyCkpt>,
7263    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7264        let n_embd = self.cfg.n_embd as usize;
7265        let saved_pos = cache.pos;
7266        let mut ph_last = std::time::Instant::now();
7267        for il in lo..hi {
7268            let mut next = e.uninit(t * n_embd)?;
7269            let mut col_states: Option<Vec<(CudaSlice<f32>, CudaSlice<f32>)>> =
7270                if ckpt.is_some() && t >= 2 && matches!(self.layers[il].mixer, Mixer::Linear(_)) {
7271                    Some(Vec::with_capacity(t - 1))
7272                } else {
7273                    None
7274                };
7275            for r in 0..t {
7276                cache.pos = pos0 + r;
7277                let mut row = e.uninit(n_embd)?;
7278                e.dtod_copy_view(&x.slice(r * n_embd..(r + 1) * n_embd), &mut row)?;
7279                let row_pos = e.htod_i32(&[(pos0 + r) as i32])?;
7280                let mut one = [&mut *cache];
7281                let ctx = self.batch_layer_ctx(e, &one, il, il + 1)?;
7282                let out = match self.decode_batch_layers(
7283                    e,
7284                    row,
7285                    &mut one,
7286                    &ctx,
7287                    &row_pos,
7288                    &mut ph_last,
7289                ) {
7290                    Ok(out) => out,
7291                    Err(error) => {
7292                        cache.pos = saved_pos;
7293                        return Err(error);
7294                    }
7295                };
7296                e.dtod_copy_into(&out, &mut next, r * n_embd)?;
7297                if r + 1 < t {
7298                    if let Some(states) = col_states.as_mut() {
7299                        let recur = cache.recur[il]
7300                            .as_ref()
7301                            .ok_or("Qwen35-MoE linear verify layer has no recurrent state")?;
7302                        states.push((
7303                            e.clone_dtod(&recur.conv_state)?,
7304                            e.clone_dtod(&recur.ssm_state)?,
7305                        ));
7306                    }
7307                }
7308            }
7309            if let (Some(checkpoint), Some(states)) = (ckpt.as_deref_mut(), col_states) {
7310                checkpoint.cols[il] = Some(states);
7311            }
7312            x = next;
7313        }
7314        cache.pos = saved_pos;
7315        Ok(x)
7316    }
7317
7318    /// T-PARALLEL VERIFY IN THE SERVING NUMERIC CLASS (lane/tparallel-verify, 2026-08-15).
7319    ///
7320    /// The weight ops run ONCE per layer at m=T; the state ops run per row through the same
7321    /// b_n=1 serving kernels the rowwise replay uses. Per-row bit-identity rests on the two
7322    /// pins the serving batch tier already carries:
7323    ///   * `matmul_pre` / `_bN` mmvq: per-row program == m=1 program (decode_batch.rs v2 note,
7324    ///     kernel-check pinned) — so a [T, n_embd] projection row equals the row projected
7325    ///     alone;
7326    ///   * row-indexed norms/elementwise (`rms_norm`, `quantize_q8_1`, `add_rms_norm`,
7327    ///     `gated_rmsnorm[_q8_1]`, `silu_mul`, `rope_neox` with per-row positions): the T-row
7328    ///     launch is the per-row program (same pin the generic verify's fused norms rely on).
7329    /// The sequential dependencies keep their exact serving order: the conv ring / gdn scan
7330    /// chain state row -> row through the `_b` kernels at b_n=1 (ping-pong via a 6-entry
7331    /// alternating pointer table, host handles swapped per row so VerifyCkpt clones the
7332    /// canonical state exactly as the rowwise arm does), and each row's KV append + fa decode
7333    /// picks its arm from ITS OWN t_kv (append: format-only; fa: `fa_seqs_eligible` + its own
7334    /// `fa_split_keys` rung at b_n=1) — the straddle law per row, so every row executes the
7335    /// program its isolated B=1 serving step would.
7336    ///
7337    /// Cost: 1 weight read per layer per round + T state micro-launches, vs the rowwise arm's
7338    /// T weight reads. Gated bit-identical vs the rowwise arm by spec-serve-gate + canary.
7339    #[allow(clippy::too_many_arguments)]
7340    fn qwen35_verify_tparallel(
7341        &self,
7342        e: &Engine,
7343        mut x: CudaSlice<f32>,
7344        lo: usize,
7345        hi: usize,
7346        pos0: usize,
7347        t: usize,
7348        cache: &mut Cache,
7349        mut ckpt: Option<&mut VerifyCkpt>,
7350        stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
7351        mut graphs: Option<&mut DsparkVerifyGraphs>,
7352    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7353        let seqs_append =
7354            std::env::var("MEMRA_BATCH_APPEND").as_deref() != Ok("0") && !Engine::kv_fp8_on();
7355        let batch_fa_on = std::env::var("MEMRA_BATCH_FA").as_deref() != Ok("0");
7356
7357        // Merge guard (v0.98 train, re-affirmed on the v0.100 train over slice 4c): the
7358        // ROUND-STREAM arm (lane/draftcost-moe, device position counter) and the dspark
7359        // verify graphs (engine-bundle slice 3 / trunk slice 4c) have no common caller —
7360        // stream rides the qwen35moe burst, graphs ride the dspark route. If a future
7361        // caller arms both, refuse loudly instead of silently dropping the graphs ctx
7362        // (the stream linear arm takes linear_attn_verify_t, not the graphed segment or
7363        // full-verify bodies).
7364        if stream.is_some() && graphs.is_some() {
7365            return Err(
7366                "qwen35 tparallel verify: ROUND-STREAM and dspark verify graphs \
7367                        cannot arm together"
7368                    .into(),
7369            );
7370        }
7371        // Engine-bundle slice 3 + slice 4c: with a graphs ctx armed, pointer tables are
7372        // refreshed once per verify (the gdn ping-pong moves handles; a fresh generation
7373        // moves the kv caches). Then:
7374        //  - slice 4c: when the WHOLE round rides one seqs rung (every row batchable, one
7375        //    split-ladder step, rung covers the round), the ENTIRE walk replays as ONE
7376        //    full-verify graph per (vt, rung) — linear layers through the shared
7377        //    `qwen35_tparallel_linear_layer` body, full-attention layers through the
7378        //    shared `qwen35_tparallel_fa_layer` body in graph mode.
7379        //  - fallback (straddle rounds, below the vec floor, partial walks): runs of
7380        //    consecutive LINEAR layers replay the slice-3 per-(segment, vt) graphs and
7381        //    the full-attention layers run eager (batched rows when eligible).
7382        //
7383        // GRAPH-LAUNCH HEADROOM GUARD (see GRAPH_LAUNCH_MIN_FREE): the dspark verify
7384        // graphs replay through this walk from THREE callers — the MTP spec round's vg
7385        // door (already dropped per round by `graph_round_ok` before it gets here), the
7386        // dspark one-shot, and the dspark SERVE round (default ON since v0.108). Below
7387        // the driver-free floor the WHOLE round takes the byte-identical eager
7388        // cols-ckpt walk — the same drop-the-ctx fallback the pool ceiling already
7389        // takes — instead of feeding cuGraphLaunch a card it segfaults on.
7390        if let Some(g) = graphs.as_deref_mut() {
7391            if !graph_launch_headroom_ok(e) {
7392                g.round_slab = false;
7393                graphs = None;
7394                static NOTED: std::sync::Once = std::sync::Once::new();
7395                NOTED.call_once(|| graph_replay_suspended_note("dspark-vg"));
7396            }
7397        }
7398        if let Some(g) = graphs.as_deref_mut() {
7399            g.refresh_tables(e, cache)?;
7400            g.round_slab = false;
7401            if let Some(rung) = g.full_rung(self, cache, lo, hi, t, seqs_append && batch_fa_on) {
7402                // Pool ceiling (dspark_vg_cap): an existing key always replays; a NEW
7403                // full capture past the ceiling falls through to the segment/eager arms.
7404                if g.full.contains_key(&(t, rung, hi)) || g.can_capture() {
7405                    let out = g.run_full(self, e, lo, hi, &x, t, pos0, rung, cache)?;
7406                    g.round_slab = true;
7407                    return Ok(out);
7408                }
7409            }
7410            // Round-atomic ceiling check for the segment door: if any linear run in this
7411            // walk would need a NEW capture past the ceiling, the whole round runs the
7412            // eager cols-ckpt walk (mixing slab- and cols-stashed layers in one round
7413            // would corrupt the commit).
7414            if !g.segments_ready(self, lo, hi, t) {
7415                graphs = None;
7416            }
7417        }
7418        // STREAM (2b, lane/draftcost-moe): positions come from the device round counter
7419        // (pos_iota / i32_copy_add) so a burst round needs no host position knowledge.
7420        let pos_d = match stream {
7421            Some((_, ctr)) => {
7422                let mut p = e.alloc_uninit::<i32>(t)?;
7423                e.pos_iota(ctr, &mut p, t)?;
7424                p
7425            }
7426            None => {
7427                let pos_host: Vec<i32> = (0..t).map(|r| (pos0 + r) as i32).collect();
7428                e.htod_i32(&pos_host)?
7429            }
7430        };
7431        // Per-row 1-element position buffers, built ONCE per verify (the append/fa wrappers
7432        // take owned pos slices; building these inside the layer x row loops cost 16xT H2Ds).
7433        // LAZY since slice 4: the batched fa/append arm never touches them — they are built
7434        // on the first per-row fallback layer only (stream-aware there; the stream FA arm
7435        // rides the dc rows kernels and never reaches the fallback).
7436        let mut pos_rows: Option<Vec<CudaSlice<i32>>> = None;
7437        let mut il = lo;
7438        while il < hi {
7439            if graphs.is_some() && matches!(self.layers[il].mixer, Mixer::Linear(_)) {
7440                let mut end = il;
7441                while end < hi && matches!(self.layers[end].mixer, Mixer::Linear(_)) {
7442                    end += 1;
7443                }
7444                let g = graphs.as_deref_mut().expect("checked above");
7445                x = g.run_segment(self, e, il, end, &x, t, cache)?;
7446                g.round_slab = true;
7447                il = end;
7448                continue;
7449            }
7450            let layer = &self.layers[il];
7451            if stream.is_none() && matches!(layer.mixer, Mixer::Linear(_)) {
7452                // Eager linear layer (no graphs ctx): the shared body, legacy cols-ckpt arm.
7453                // Under ROUND-STREAM the linear layers ride the fa-body match's stream arm
7454                // below (linear_attn_verify_t — the stream COMMIT needs its GdnStash).
7455                x = self.qwen35_tparallel_linear_layer(
7456                    e,
7457                    il,
7458                    &x,
7459                    t,
7460                    cache,
7461                    ckpt.as_deref_mut(),
7462                    None,
7463                    None,
7464                )?;
7465                il += 1;
7466                continue;
7467            }
7468            // Full-attention (or stream-Linear, or MLA-refusing) layer: the extracted
7469            // shared body — eager arm (fresh per-verify pos/table, exact t_kv sizing,
7470            // in-body len bump). The slice-4c captured full-verify graphs run the SAME
7471            // body in graph mode; under ROUND-STREAM the body's dc-rows / GDN stream arms
7472            // run (lane/draftcost-moe).
7473            x = self.qwen35_tparallel_fa_layer(
7474                e,
7475                il,
7476                &x,
7477                t,
7478                cache,
7479                FaLayerArgs {
7480                    pos_d: &pos_d,
7481                    pos_rows: &mut pos_rows,
7482                    pos0,
7483                    seqs_append,
7484                    batch_fa_on,
7485                    graph_cap: None,
7486                    stream,
7487                    ckpt: ckpt.as_deref_mut(),
7488                },
7489            )?;
7490            il += 1;
7491        }
7492        Ok(x)
7493    }
7494
7495    /// SHARED dense-FFN body for the qwen35 t-parallel layers (trunk-kernels slice B) —
7496    /// ONE copy for the fa and linear layer bodies (the verify_layers extraction lesson).
7497    /// Dual arm (MEMRA_TK_FFN_DUAL, default on): gate+up in ONE dual launch from the
7498    /// pre-quantized activation with macro-scales DEFERRED into the fused SwiGLU+q8_1
7499    /// epilogue, then ffn_down from the fused (aq, ad) — the q27 verify chain verbatim.
7500    /// Every door is the bit-identical proven one: `matmul_decode_exact_dual_pre` (per
7501    /// (tensor,token,row) == the two singles), `silu_mul_scaled_q8_1` (y*s inline == the
7502    /// scale_inplace store, value-exact; fused quantize == quantize_q8_1 bytes),
7503    /// `matmul_decode_exact_pre` (dispatch mirror of the singles' q8_1-fast tail).
7504    /// Dual-refused (t outside 2..=7, non-NVFP4, layout mismatch) or seam off -> the
7505    /// original singles chain, byte-for-byte.
7506    #[allow(clippy::too_many_arguments)]
7507    fn qwen35_tparallel_dense_ffn(
7508        &self,
7509        e: &Engine,
7510        ffn_gate: &crate::model::GpuTensor,
7511        ffn_up: &crate::model::GpuTensor,
7512        ffn_down: &crate::model::GpuTensor,
7513        zn: &CudaSlice<f32>,
7514        t: usize,
7515        n_embd: usize,
7516    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7517        let n_ff = ffn_gate.out_features();
7518        let (zq, zd) = e.quantize_q8_1(zn, t, n_embd)?;
7519        if Engine::tk_ffn_dual_on() {
7520            if let Some(((g, gs), (u, us))) =
7521                e.matmul_decode_exact_dual_pre(ffn_gate, ffn_up, &zq, &zd, t)?
7522            {
7523                if e.uses_q8_1_fast(ffn_down) {
7524                    let (aq, ad) = e.silu_mul_scaled_q8_1(&g, &u, gs, us, t * n_ff)?;
7525                    return e.matmul_decode_exact_pre(ffn_down, &aq, &ad, t);
7526                }
7527                let mut act = e.uninit(t * n_ff)?;
7528                e.silu_mul_scaled(&g, &u, gs, us, &mut act, t * n_ff)?;
7529                let (aq, ad) = e.quantize_q8_1(&act, t, n_ff)?;
7530                return e.matmul_pre(ffn_down, &aq, &ad, &act, t);
7531            }
7532        }
7533        // v1 singles chain (seam off or dual-refused) — the pre-slice-B body verbatim.
7534        let g = e.matmul_pre(ffn_gate, &zq, &zd, zn, t)?;
7535        let u = e.matmul_pre(ffn_up, &zq, &zd, zn, t)?;
7536        let mut act = e.uninit(t * n_ff)?;
7537        e.silu_mul(&g, &u, &mut act, t * n_ff)?;
7538        let (aq, ad) = e.quantize_q8_1(&act, t, n_ff)?;
7539        e.matmul_pre(ffn_down, &aq, &ad, &act, t)
7540    }
7541
7542    /// ONE t-parallel FULL-ATTENTION layer (attn_norm + fa mixer + post_attn_norm + FFN +
7543    /// tap) — extracted from the walk exactly like `qwen35_tparallel_linear_layer` so the
7544    /// eager walk and the slice-4c captured full-verify graphs execute the SAME body (a
7545    /// second copy is how dispatch mirrors drift — the verify_layers extraction lesson).
7546    ///
7547    /// `args.graph_cap = Some((table, off, rung_end))` is the captured-graph mode:
7548    /// - kv base-pointer pairs come from the ctx-owned persistent table at `off` (a fresh
7549    ///   generation's cache lands at new addresses that only the per-verify table refresh
7550    ///   knows — the slice-3 baked-address lesson);
7551    /// - the seqs twins size partials/grid at `rung_end` and pin `split_keys` to the
7552    ///   rung's ladder value: `n_splits_max` is pure stride, splits >= ns_eff write the
7553    ///   EMPTY partial the combine never reads, and every per-row T_kv derives in-kernel
7554    ///   from `pos_seq[z]` — so one captured launch replays bit-identically for every
7555    ///   round whose rows all sit inside the rung;
7556    /// - the host len bump moves to the replay caller (captured host code does not
7557    ///   re-run at replay).
7558    /// Graph mode REFUSES any round the batched arm cannot take: the per-row fallback
7559    /// host-branches on t_kv and must never be captured.
7560    #[allow(clippy::too_many_arguments)]
7561    fn qwen35_tparallel_fa_layer(
7562        &self,
7563        e: &Engine,
7564        il: usize,
7565        x: &CudaSlice<f32>,
7566        t: usize,
7567        cache: &mut Cache,
7568        args: FaLayerArgs<'_>,
7569    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7570        use cudarc::driver::DevicePtr;
7571        let cfg = &self.cfg;
7572        let n_embd = cfg.n_embd as usize;
7573        let eps = cfg.rms_eps;
7574        let head_dim_global = cfg.head_dim_k as usize;
7575        let layer = &self.layers[il];
7576        let FaLayerArgs {
7577            pos_d,
7578            pos_rows,
7579            pos0,
7580            seqs_append,
7581            batch_fa_on,
7582            graph_cap,
7583            stream,
7584            mut ckpt,
7585        } = args;
7586
7587        // ---- attn_norm + q8_1 quantize at m=T (row-indexed == per-row) ----
7588        let anorm = layer.attn_norm.float_data();
7589        let mut xn = e.uninit(t * n_embd)?;
7590        e.rms_norm(x, anorm, &mut xn, n_embd, t, eps)?;
7591        let (hq, hd) = e.quantize_q8_1(&xn, t, n_embd)?;
7592
7593        let mixed: CudaSlice<f32> = match &layer.mixer {
7594            Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
7595            // STREAM ARM (2b, lane/draftcost-moe): under a device position counter the
7596            // per-row serving-kernel chain cannot run (host state swaps keyed on host
7597            // row index are fine, but the stream COMMIT needs the GdnStash for its _dc
7598            // rebuild — the per-row chain only produces per-column clones). GDN rides
7599            // `linear_attn_verify_t`: batched q8_1-class projections, stash-producing,
7600            // and its one-scan recurrence is pinned bit-identical to T chained T=1
7601            // steps (its header + kernel-check). Position-independent, so no counter
7602            // plumbing is needed. Guards mirror the generic call site exactly.
7603            Mixer::Linear(la) if stream.is_some() => {
7604                if !(t >= 3 || (t == 2 && spec_m2()))
7605                    || !self.mixer_in_q8_1_fast(e, &layer.mixer)
7606                    || !e.uses_q8_1_fast(&la.ssm_out)
7607                {
7608                    return Err("qwen35 stream verify: GDN batched arm requires t>=3 \
7609                                (or MEMRA_SPEC_M2 at t=2) and q8_1-fast projections"
7610                        .into());
7611                }
7612                let want = ckpt.is_some();
7613                let (out, stash) =
7614                    self.linear_attn_verify_t(e, la, &xn, Some((&hq, &hd)), t, cache, il, want)?;
7615                if let (Some(ck), Some(st)) = (ckpt.as_deref_mut(), stash) {
7616                    ck.gdn[il] = Some(st);
7617                }
7618                out
7619            }
7620            Mixer::Linear(_) => {
7621                unreachable!("linear layers ride qwen35_tparallel_linear_layer")
7622            }
7623            Mixer::Full(fa) => {
7624                let geometry = cfg.full_attention_geometry_at(il as u32);
7625                let n_head = geometry.n_head as usize;
7626                let n_head_kv = geometry.n_head_kv as usize;
7627                let head_dim = geometry.head_dim_k as usize;
7628                let rope_dims = geometry.n_rot as usize;
7629                let rope_base = geometry.rope_base;
7630                let scale = geometry.attention_scale();
7631                // Batched projections: one weight read serves all T rows.
7632                // GROUP-3 twin (trunk-kernels slice D): q/k/v in ONE launch — the group4
7633                // kernel with n3=0, bit-identical per (tensor, token, row) to the three
7634                // singles; refused or MEMRA_TK_FA_GROUP=0 -> singles byte-for-byte.
7635                let (qf, mut k, v) = match e.matmul_decode_exact_group3_pre(
7636                    [&fa.wq, &fa.wk, &fa.wv],
7637                    &hq,
7638                    &hd,
7639                    t,
7640                )? {
7641                    Some(mut g3) => {
7642                        let v = g3.pop().unwrap();
7643                        let k = g3.pop().unwrap();
7644                        let qf = g3.pop().unwrap();
7645                        (qf, k, v)
7646                    }
7647                    None => (
7648                        e.matmul_pre(&fa.wq, &hq, &hd, &xn, t)?,
7649                        e.matmul_pre(&fa.wk, &hq, &hd, &xn, t)?,
7650                        e.matmul_pre(&fa.wv, &hq, &hd, &xn, t)?,
7651                    ),
7652                };
7653                let gated =
7654                    geometry.attention_gate == memra_gguf::config::AttentionGateKind::FusedQ;
7655                let (mut q, gate) = if gated {
7656                    let mut qs = e.uninit(t * n_head * head_dim)?;
7657                    let mut gs = e.uninit(t * n_head * head_dim)?;
7658                    e.q_gate_split(&qf, &mut qs, &mut gs, head_dim, n_head, t)?;
7659                    (qs, Some(gs))
7660                } else {
7661                    (qf, None)
7662                };
7663                let mut qn = e.uninit(t * n_head * head_dim)?;
7664                e.rms_norm(
7665                    &q,
7666                    fa.q_norm.float_data(),
7667                    &mut qn,
7668                    head_dim,
7669                    t * n_head,
7670                    eps,
7671                )?;
7672                q = qn;
7673                let mut kn = e.uninit(t * n_head_kv * head_dim)?;
7674                e.rms_norm(
7675                    &k,
7676                    fa.k_norm.float_data(),
7677                    &mut kn,
7678                    head_dim,
7679                    t * n_head_kv,
7680                    eps,
7681                )?;
7682                k = kn;
7683                e.rope_neox(
7684                    &mut q, pos_d, head_dim, rope_dims, n_head, t, rope_base, 1.0,
7685                )?;
7686                e.rope_neox(
7687                    &mut k, pos_d, head_dim, rope_dims, n_head_kv, t, rope_base, 1.0,
7688                )?;
7689
7690                // Per-row append + attend: row r sees rows 0..r in KV (causal within the
7691                // draft), each through the b_n=1 serving kernels at its own t_kv.
7692                let q_dim = n_head * head_dim;
7693                let kv_dim = n_head_kv * head_dim;
7694                let mut attn = e.uninit(t * q_dim)?;
7695                let (kdk, kdv, ktb, vtb, len0, kv_local) = {
7696                    let kvl = cache.kv[il].as_ref().unwrap();
7697                    // [2T] interleaved k,v base pointers: entry pair z serves row z of
7698                    // the batched twins; the per-row fallback reads pair 0 (same cache
7699                    // for every row of one layer). Graph mode reads the ctx table.
7700                    let local: Option<CudaSlice<u64>> = match graph_cap {
7701                        Some(_) => None,
7702                        None => {
7703                            let s = &e.gpu.stream();
7704                            let (pk, _g) = kvl.k.device_ptr(s);
7705                            let (pv, _g2) = kvl.v.device_ptr(s);
7706                            let mut tbl = Vec::with_capacity(2 * t);
7707                            for _ in 0..t {
7708                                tbl.push(pk as u64);
7709                                tbl.push(pv as u64);
7710                            }
7711                            Some(e.htod_u64(&tbl)?)
7712                        }
7713                    };
7714                    (
7715                        kvl.kv_dim_k,
7716                        kvl.kv_dim_v,
7717                        kvl.k_tok_bytes,
7718                        kvl.v_tok_bytes,
7719                        kvl.len,
7720                        local,
7721                    )
7722                };
7723                let (kv_tbl, kv_off): (&CudaSlice<u64>, usize) = match graph_cap {
7724                    Some((tb, off, _)) => (tb, off),
7725                    None => (kv_local.as_ref().expect("built above"), 0),
7726                };
7727                // Slice 4 (fa/append rows — see dspark_fa_rows_on): the whole per-row
7728                // section batches into the z-batched serving twins when every row of
7729                // this round takes the v4-seqs arm on ONE fa_split_keys rung. Both
7730                // guards are evaluated at the round's FIRST and LAST t_kv — the
7731                // eligibility window (vec floor .. v4 max) and each split-ladder rung
7732                // are intervals in t_kv, so ends-inside means all-inside (the straddle
7733                // law). Appending all T rows before any attend is read-equivalent to
7734                // the interleaved order: row r's walk reads keys 0..len0+r only, and
7735                // rows > r land at slots it never touches; every written cache row is
7736                // the per-token appender's exact warp program (kernel-check pinned).
7737                let t_kv_first = len0 + 1;
7738                let t_kv_last = len0 + t;
7739                let rows_batched = t >= 2
7740                    && seqs_append
7741                    && batch_fa_on
7742                    && dspark_fa_rows_on()
7743                    // the z-batched twins read stacked rows at the CACHE's kv dims;
7744                    // the projection stack is [T, n_head_kv*head_dim] — they must be
7745                    // the same stride or row z misaligns (true for this family; the
7746                    // guard keeps any asymmetric-kv model on the per-row loop).
7747                    && kdk == kv_dim
7748                    && kdv == kv_dim
7749                    && crate::fa_seqs_eligible(t_kv_first, head_dim_global)
7750                    && crate::fa_seqs_eligible(t_kv_last, head_dim_global)
7751                    && crate::fa_split_keys(t_kv_first, cfg.n_head_kv as usize)
7752                        == crate::fa_split_keys(t_kv_last, cfg.n_head_kv as usize);
7753                // Sizing: eager = exact round bound; graph mode = the rung end (stride +
7754                // grid only — bytes proven equal above). Capture-time invariants refuse
7755                // loudly rather than bake a divergent body.
7756                let (size_kv_max, sp) = match graph_cap {
7757                    Some((_, _, rung)) => {
7758                        if !rows_batched {
7759                            return Err(format!(
7760                                "fa graph capture: layer {il} round is not batchable \
7761                                 (t_kv {t_kv_first}..{t_kv_last}) — the per-row fallback \
7762                                 must never be captured"
7763                            )
7764                            .into());
7765                        }
7766                        let sp_r = crate::fa_split_keys(rung, cfg.n_head_kv as usize);
7767                        if t_kv_last > rung
7768                            || sp_r != crate::fa_split_keys(t_kv_last, cfg.n_head_kv as usize)
7769                        {
7770                            return Err(format!(
7771                                "fa graph capture: rung {rung} does not cover round \
7772                                 t_kv {t_kv_first}..{t_kv_last} on one split ladder step"
7773                            )
7774                            .into());
7775                        }
7776                        (rung, sp_r)
7777                    }
7778                    None => (
7779                        t_kv_last,
7780                        crate::fa_split_keys(t_kv_last, cfg.n_head_kv as usize),
7781                    ),
7782                };
7783                if let Some((_, ctr)) = stream {
7784                    // STREAM ARM (2b): one batched dc append + the multi-row dc attention
7785                    // — the generic stream arm's exact shape (rows kernels are pinned
7786                    // byte-identical to the per-row programs by kernel-check). Host len
7787                    // stays a stale lower bound; the burst drain reconciles it.
7788                    let kvl = cache.kv[il].as_mut().unwrap();
7789                    e.append_kv_quantized_rows_dc(
7790                        &k,
7791                        &v,
7792                        &mut kvl.k,
7793                        &mut kvl.v,
7794                        ctr,
7795                        t,
7796                        kdk,
7797                        kdv,
7798                        ktb,
7799                        vtb,
7800                        Engine::kv_fp8_on(),
7801                    )?;
7802                    let upper = (kvl.len + t + 64).min(cache.max_ctx);
7803                    let k_view = e.view_u8(&kvl.k, upper * ktb);
7804                    let v_view = e.view_u8(&kvl.v, upper * vtb);
7805                    e.fa_decode_rows_dc(
7806                        &q, &k_view, &v_view, &mut attn, head_dim, n_head, n_head_kv, ctr, upper,
7807                        t, scale, ktb, vtb, 0, false,
7808                    )?;
7809                } else if rows_batched {
7810                    e.append_kv_quantized_seqs(
7811                        &k,
7812                        &v,
7813                        &kv_tbl.slice(kv_off..kv_off + 2 * t),
7814                        pos_d,
7815                        t,
7816                        kdk,
7817                        kdv,
7818                        ktb,
7819                        vtb,
7820                    )?;
7821                    if graph_cap.is_none() {
7822                        cache.kv[il].as_mut().unwrap().len += t;
7823                    }
7824                    e.fa_decode_batch_seqs_v4(
7825                        &q,
7826                        &kv_tbl.slice(kv_off..kv_off + 2 * t),
7827                        pos_d,
7828                        &mut attn,
7829                        head_dim,
7830                        n_head,
7831                        n_head_kv,
7832                        t,
7833                        size_kv_max,
7834                        scale,
7835                        sp,
7836                        ktb,
7837                        vtb,
7838                    )?;
7839                } else {
7840                    if pos_rows.is_none() {
7841                        // Stream-aware for symmetry with pos_d (the stream FA arm rides
7842                        // the dc rows kernels above and never reaches this fallback).
7843                        *pos_rows = Some(match stream {
7844                            Some((_, ctr)) => (0..t)
7845                                .map(|r| {
7846                                    let mut b = e.alloc_uninit::<i32>(1)?;
7847                                    e.i32_copy_add(ctr, &mut b, r as i32)?;
7848                                    Ok(b)
7849                                })
7850                                .collect::<Result<_, Box<dyn std::error::Error>>>()?,
7851                            None => (0..t)
7852                                .map(|r| e.htod_i32(&[(pos0 + r) as i32]))
7853                                .collect::<Result<_, _>>()?,
7854                        });
7855                    }
7856                    let pos_rows = pos_rows.as_ref().unwrap();
7857                    for r in 0..t {
7858                        // Owned per-row scratch: the b_n=1 kernels take packed batch buffers
7859                        // whose row 0 is this row (arithmetic-free materialization copies,
7860                        // same as decode's per-seq fallback arm).
7861                        let mut k_row = e.uninit(kv_dim)?;
7862                        e.dtod_copy_view(&k.slice(r * kv_dim..(r + 1) * kv_dim), &mut k_row)?;
7863                        let mut v_row = e.uninit(kv_dim)?;
7864                        e.dtod_copy_view(&v.slice(r * kv_dim..(r + 1) * kv_dim), &mut v_row)?;
7865                        let pos_row = &pos_rows[r];
7866                        let kvl = cache.kv[il].as_mut().unwrap();
7867                        if seqs_append {
7868                            e.append_kv_quantized_seqs(
7869                                &k_row,
7870                                &v_row,
7871                                &kv_tbl.slice(kv_off..kv_off + 2),
7872                                pos_row,
7873                                1,
7874                                kdk,
7875                                kdv,
7876                                ktb,
7877                                vtb,
7878                            )?;
7879                            kvl.len += 1;
7880                        } else {
7881                            e.append_kv_quantized_view(
7882                                &k_row.slice(0..kv_dim),
7883                                &v_row.slice(0..kv_dim),
7884                                &mut kvl.k,
7885                                &mut kvl.v,
7886                                kvl.len,
7887                                kvl.kv_dim_k,
7888                                kvl.kv_dim_v,
7889                                kvl.k_tok_bytes,
7890                                kvl.v_tok_bytes,
7891                                Engine::kv_fp8_on(),
7892                            )?;
7893                            kvl.len += 1;
7894                        }
7895                        let t_kv = kvl.len;
7896                        let mut q_row = e.uninit(q_dim)?;
7897                        e.dtod_copy_view(&q.slice(r * q_dim..(r + 1) * q_dim), &mut q_row)?;
7898                        let mut a_row = e.uninit(q_dim)?;
7899                        if batch_fa_on && crate::fa_seqs_eligible(t_kv, head_dim_global) {
7900                            let sp0_r = crate::fa_split_keys(t_kv, cfg.n_head_kv as usize);
7901                            e.fa_decode_batch_seqs_v4(
7902                                &q_row,
7903                                &kv_tbl.slice(kv_off..kv_off + 2),
7904                                pos_row,
7905                                &mut a_row,
7906                                head_dim,
7907                                n_head,
7908                                n_head_kv,
7909                                1,
7910                                t_kv,
7911                                scale,
7912                                sp0_r,
7913                                ktb,
7914                                vtb,
7915                            )?;
7916                        } else {
7917                            let k_view = e.view_u8(&kvl.k, t_kv * kvl.k_tok_bytes);
7918                            let v_view = e.view_u8(&kvl.v, t_kv * kvl.v_tok_bytes);
7919                            let mut a_view = a_row.slice_mut(0..q_dim);
7920                            e.fa_decode_kvmod_view(
7921                                &q_row.slice(0..q_dim),
7922                                &k_view,
7923                                &v_view,
7924                                &mut a_view,
7925                                head_dim,
7926                                n_head,
7927                                n_head_kv,
7928                                t_kv,
7929                                scale,
7930                                kvl.k_tok_bytes,
7931                                kvl.v_tok_bytes,
7932                                Engine::kv_fp8_on(),
7933                            )?;
7934                        }
7935                        e.dtod_copy_into(&a_row, &mut attn, r * q_dim)?;
7936                    }
7937                }
7938
7939                // Output gate (element-wise) + o-proj at m=T.
7940                let attn_g = match &gate {
7941                    Some(g) => {
7942                        let n = t * q_dim;
7943                        let mut gsig = e.uninit(n)?;
7944                        e.sigmoid(g, &mut gsig, n)?;
7945                        let mut ag = e.uninit(n)?;
7946                        e.mul(&attn, &gsig, &mut ag, n)?;
7947                        ag
7948                    }
7949                    None => attn,
7950                };
7951                e.matmul(&fa.wo, &attn_g, t)?
7952            }
7953        };
7954
7955        // ---- residual add + post_attn_norm + FFN at m=T (serving dispatch verbatim) ----
7956        let pnorm = layer.post_attn_norm.float_data();
7957        let mut x1 = e.uninit(t * n_embd)?;
7958        let mut zn = e.uninit(t * n_embd)?;
7959        e.add_rms_norm(x, &mixed, pnorm, &mut x1, &mut zn, n_embd, t, eps)?;
7960        let ffn_out = match &layer.ffn {
7961            crate::hybrid::Ffn::Dense {
7962                ffn_gate,
7963                ffn_up,
7964                ffn_down,
7965            } => {
7966                assert!(
7967                    self.cfg.m3.is_none(),
7968                    "qwen35 t-parallel verify: M3 swigluoai FFN not yet batched"
7969                );
7970                self.qwen35_tparallel_dense_ffn(e, ffn_gate, ffn_up, ffn_down, &zn, t, n_embd)?
7971            }
7972            crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il_zq8(e, m, &zn, None, t, il as u16)?,
7973        };
7974        let mut x2 = e.uninit(t * n_embd)?;
7975        e.add(&x1, &ffn_out, &mut x2, t * n_embd)?;
7976        // dspark drafter tap (no-op when no sink armed): post-layer residual verify rows
7977        self.dflash_tap(e, cache, il, &x2, t)?;
7978        Ok(x2)
7979    }
7980
7981    /// ONE t-parallel LINEAR layer (attn_norm + gdn mixer + post_attn_norm + FFN + tap) —
7982    /// the exact body the old in-loop Linear arm ran, extracted so the eager walk and the
7983    /// slice-3 captured segments execute the SAME code (a second copy is how dispatch
7984    /// mirrors drift — the verify_layers extraction lesson). Two deliberate changes, both
7985    /// bit-identical by construction:
7986    /// - the gdn ping-pong host swap moves from per-row to ONE end-of-body swap (t odd):
7987    ///   the device sequence is driven entirely by the 6-entry pointer table, which
7988    ///   already encodes both parities; the ckpt stash reads name row r's out buffer
7989    ///   directly (r even -> alt handle, odd -> canonical) — the same physical bytes the
7990    ///   legacy post-swap clone read.
7991    /// - `stash` (slice-3 ctx): persistent per-layer slabs written by copy_into instead of
7992    ///   per-row clone_dtod allocs — same bytes, capture-legal (no per-round host objects).
7993    /// `table_src` = (persistent pointer table, offset) when the ctx owns the tables;
7994    /// None builds the per-verify table exactly as before.
7995    #[allow(clippy::too_many_arguments)]
7996    fn qwen35_tparallel_linear_layer(
7997        &self,
7998        e: &Engine,
7999        il: usize,
8000        x: &CudaSlice<f32>,
8001        t: usize,
8002        cache: &mut Cache,
8003        mut ckpt: Option<&mut VerifyCkpt>,
8004        stash: Option<(&mut CudaSlice<f32>, &mut CudaSlice<f32>)>,
8005        table_src: Option<(&CudaSlice<u64>, usize)>,
8006    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8007        use cudarc::driver::DevicePtr;
8008        let cfg = &self.cfg;
8009        let n_embd = cfg.n_embd as usize;
8010        let eps = cfg.rms_eps;
8011        let layer = &self.layers[il];
8012        let Mixer::Linear(la) = &layer.mixer else {
8013            return Err("qwen35_tparallel_linear_layer on a non-linear layer".into());
8014        };
8015        // ---- attn_norm + q8_1 quantize at m=T (row-indexed == per-row) ----
8016        let anorm = layer.attn_norm.float_data();
8017        let mut xn = e.uninit(t * n_embd)?;
8018        e.rms_norm(x, anorm, &mut xn, n_embd, t, eps)?;
8019        let (hq, hd) = e.quantize_q8_1(&xn, t, n_embd)?;
8020
8021        let geometry = la.geometry;
8022        let d_state = geometry.key_head_dim as usize;
8023        let num_k = geometry.key_heads as usize;
8024        let num_v = geometry.value_heads as usize;
8025        let d_conv = geometry.conv_kernel as usize;
8026        let key_dim = d_state * num_k;
8027        let value_dim = geometry.value_head_dim as usize * num_v;
8028        let conv_dim = key_dim * 2 + value_dim;
8029        let gdn_scale = 1.0 / (d_state as f32).sqrt();
8030
8031        // ---- batched projections: one weight read for all T rows ----
8032        // GROUP-4 twin (trunk-kernels slice C): the whole 4-tuple in ONE launch, bit-identical
8033        // per (tensor, token, row) to the four singles; refused (layout/tier) or
8034        // MEMRA_TK_GDN_GROUP=0 -> the singles chain byte-for-byte.
8035        let (qkv_mixed, z, beta_raw, alpha) = match e.matmul_decode_exact_group4_pre(
8036            [&la.wqkv, &la.wqkv_gate, &la.ssm_beta, &la.ssm_alpha],
8037            &hq,
8038            &hd,
8039            t,
8040        )? {
8041            Some(mut g4) => {
8042                let alpha = g4.pop().unwrap();
8043                let beta_raw = g4.pop().unwrap();
8044                let z = g4.pop().unwrap();
8045                let qkv_mixed = g4.pop().unwrap();
8046                (qkv_mixed, z, beta_raw, alpha)
8047            }
8048            None => (
8049                e.matmul_pre(&la.wqkv, &hq, &hd, &xn, t)?,
8050                e.matmul_pre(&la.wqkv_gate, &hq, &hd, &xn, t)?,
8051                e.matmul_pre(&la.ssm_beta, &hq, &hd, &xn, t)?,
8052                e.matmul_pre(&la.ssm_alpha, &hq, &hd, &xn, t)?,
8053            ),
8054        };
8055        let beta_w = la.ssm_beta.out_features();
8056        let alpha_w = la.ssm_alpha.out_features();
8057        let qkv_w = la.wqkv.out_features();
8058
8059        // ---- per-row state chain through the b_n=1 serving kernels ----
8060        // 6-entry alternating pointer table expresses the ping-pong without a rebuild per
8061        // row: even rows scan s0 -> s1, odd rows s1 -> s0.
8062        let table_local: Option<CudaSlice<u64>> = match table_src {
8063            Some(_) => None,
8064            None => {
8065                let rl = cache.recur[il].as_ref().unwrap();
8066                let s = &e.gpu.stream();
8067                let (pc, _g0) = rl.conv_state.device_ptr(s);
8068                let (p0, _g1) = rl.ssm_state.device_ptr(s);
8069                let (p1, _g2) = rl.ssm_state_alt.device_ptr(s);
8070                Some(e.htod_u64(&[
8071                    pc as u64, p0 as u64, p1 as u64, pc as u64, p1 as u64, p0 as u64,
8072                ])?)
8073            }
8074        };
8075        let (table, toff): (&CudaSlice<u64>, usize) = match table_src {
8076            Some((tb, off)) => (tb, off),
8077            None => (table_local.as_ref().unwrap(), 0),
8078        };
8079        let mut o_all = e.uninit(t * value_dim)?;
8080        let mut col_states: Option<Vec<(CudaSlice<f32>, CudaSlice<f32>)>> =
8081            if ckpt.is_some() && stash.is_none() && t >= 2 {
8082                Some(Vec::with_capacity(t - 1))
8083            } else {
8084                None
8085            };
8086        let mut stash = stash;
8087        // Per-row scratch reused across rows (uninit is cheap but not free at
8088        // 48 layers x T rows); row inputs/outputs pass as VIEWS into the packed
8089        // [T, ...] buffers — zero arithmetic-free copies in this loop.
8090        let mut conv_out = e.uninit(conv_dim)?;
8091        let mut q_l2 = e.uninit(value_dim)?;
8092        let mut k_l2 = e.uninit(value_dim)?;
8093        let mut v_gd = e.uninit(value_dim)?;
8094        let mut beta_b = e.uninit(num_v)?;
8095        let mut g_log = e.uninit(num_v)?;
8096        for r in 0..t {
8097            let base = toff + if r % 2 == 0 { 0 } else { 3 };
8098            let conv_view = table.slice(base..base + 1);
8099            let in_view = table.slice(base + 1..base + 2);
8100            let out_view = table.slice(base + 2..base + 3);
8101            e.ssm_conv1d_fused_decode_b_view(
8102                &qkv_mixed.slice(r * qkv_w..(r + 1) * qkv_w),
8103                &conv_view,
8104                la.ssm_conv1d.float_data(),
8105                &mut conv_out,
8106                conv_dim,
8107                d_conv,
8108                1,
8109            )?;
8110            e.gdn_prep_decode_b_view(
8111                &conv_out,
8112                &beta_raw.slice(r * beta_w..(r + 1) * beta_w),
8113                &alpha.slice(r * alpha_w..(r + 1) * alpha_w),
8114                la.ssm_dt.float_data(),
8115                la.ssm_a.float_data(),
8116                &mut q_l2,
8117                &mut k_l2,
8118                &mut v_gd,
8119                &mut beta_b,
8120                &mut g_log,
8121                d_state,
8122                num_v,
8123                num_k,
8124                key_dim,
8125                eps,
8126                conv_dim,
8127                1,
8128            )?;
8129            let mut o_row = o_all.slice_mut(r * value_dim..(r + 1) * value_dim);
8130            e.gdn_scan_s128_batched_view(
8131                &q_l2, &k_l2, &v_gd, &g_log, &beta_b, &in_view, &out_view, &mut o_row, num_v, 1,
8132                gdn_scale,
8133            )?;
8134            if r + 1 < t {
8135                // Row r's out buffer: even rows write s1 (the alt handle — no swaps ran),
8136                // odd rows write s0 — the same physical state the legacy post-swap
8137                // canonical clone read.
8138                let rl = cache.recur[il]
8139                    .as_ref()
8140                    .ok_or("qwen35 linear verify layer has no recurrent state")?;
8141                let ssm_src = if r % 2 == 0 {
8142                    &rl.ssm_state_alt
8143                } else {
8144                    &rl.ssm_state
8145                };
8146                match stash.as_mut() {
8147                    Some((conv_slab, ssm_slab)) => {
8148                        // BOTH stash reads go through the pointer table at run time: the
8149                        // ssm handles ping-pong between rounds, and the ctx (with its
8150                        // captured graphs) outlives the Cache — a fresh generation's
8151                        // conv/ssm buffers land at new addresses that only the per-round
8152                        // table refresh knows. A baked direct copy would read freed
8153                        // memory (parity was the slice-3 smoke divergence; cache
8154                        // lifetime is the cross-generation twin).
8155                        e.copy_indirect_src_f32(
8156                            &conv_view,
8157                            conv_slab,
8158                            r * conv_dim * (d_conv - 1),
8159                            conv_dim * (d_conv - 1),
8160                        )?;
8161                        // The ssm handles PING-PONG between rounds: a captured direct
8162                        // copy would bake the capture-time physical buffer and read the
8163                        // wrong parity after any odd-vt round (the slice-3 smoke
8164                        // divergence). Read the src address from row r's OUT table
8165                        // entry at run time — the same entry the scan just wrote.
8166                        e.copy_indirect_src_f32(
8167                            &out_view,
8168                            ssm_slab,
8169                            r * d_state * d_state * num_v,
8170                            d_state * d_state * num_v,
8171                        )?;
8172                    }
8173                    None => {
8174                        if let Some(states) = col_states.as_mut() {
8175                            states.push((e.clone_dtod(&rl.conv_state)?, e.clone_dtod(ssm_src)?));
8176                        }
8177                    }
8178                }
8179            }
8180        }
8181        // ONE end-of-body parity swap (t odd) — the legacy loop swapped per row; the net
8182        // handle motion is identical and the device sequence never read the handles.
8183        if t % 2 == 1 {
8184            let rl = cache.recur[il].as_mut().unwrap();
8185            std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
8186        }
8187        if let (Some(checkpoint), Some(states)) = (ckpt.as_deref_mut(), col_states) {
8188            checkpoint.cols[il] = Some(states);
8189        }
8190
8191        // ---- batched gated norm + out-projection at m=T ----
8192        let mixed = if e.uses_q8_1_fast(&la.ssm_out) {
8193            let (gq, gd) = e.gated_rmsnorm_q8_1(
8194                &o_all,
8195                la.ssm_norm.float_data(),
8196                &z,
8197                d_state,
8198                t * num_v,
8199                eps,
8200            )?;
8201            let g0 = e.zeros(0)?;
8202            e.matmul_pre(&la.ssm_out, &gq, &gd, &g0, t)?
8203        } else {
8204            let mut gn = e.uninit(t * value_dim)?;
8205            e.gated_rmsnorm(
8206                &o_all,
8207                la.ssm_norm.float_data(),
8208                &z,
8209                &mut gn,
8210                d_state,
8211                t * num_v,
8212                eps,
8213            )?;
8214            e.matmul(&la.ssm_out, &gn, t)?
8215        };
8216
8217        // ---- residual add + post_attn_norm + FFN at m=T (serving dispatch verbatim) ----
8218        let pnorm = layer.post_attn_norm.float_data();
8219        let mut x1 = e.uninit(t * n_embd)?;
8220        let mut zn = e.uninit(t * n_embd)?;
8221        e.add_rms_norm(x, &mixed, pnorm, &mut x1, &mut zn, n_embd, t, eps)?;
8222        let ffn_out = match &layer.ffn {
8223            crate::hybrid::Ffn::Dense {
8224                ffn_gate,
8225                ffn_up,
8226                ffn_down,
8227            } => {
8228                assert!(
8229                    self.cfg.m3.is_none(),
8230                    "qwen35 t-parallel verify: M3 swigluoai FFN not yet batched"
8231                );
8232                self.qwen35_tparallel_dense_ffn(e, ffn_gate, ffn_up, ffn_down, &zn, t, n_embd)?
8233            }
8234            crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il_zq8(e, m, &zn, None, t, il as u16)?,
8235        };
8236        let mut x2 = e.uninit(t * n_embd)?;
8237        e.add(&x1, &ffn_out, &mut x2, t * n_embd)?;
8238        // dspark drafter tap (no-op when no sink armed): post-layer residual verify rows
8239        self.dflash_tap(e, cache, il, &x2, t)?;
8240        Ok(x2)
8241    }
8242
8243    /// PP-N STAGE SUBGRAPH of the verify trunk: layers `[lo, hi)` of `decode_step_t_core_stream`'s
8244    /// walk, verbatim. Enters with a MATERIALIZED `[T, n_embd]` residual (no pending fusion pair
8245    /// carried in from outside the range) and exits with the range's final residual materialized
8246    /// (the trailing add executed) — exactly the `decode_layers_eager(lo, hi)` contract, T rows
8247    /// instead of one.
8248    ///
8249    /// EXTRACTED (lane/pp2-spec 2026-08-06) rather than duplicated: `decode_step_t_core_stream` IS
8250    /// the single funnel every verify forward reaches, and its per-layer dispatch MIRRORING (norm
8251    /// fusion per layer, the t>=3/spec_m2 batched-linear window, the fused-q8 FFN chain, the
8252    /// decode-exact projections) is what makes verify bit-identical to eager decode. A second copy
8253    /// for the split arm is how those mirrors drift apart on the next lever. The unsplit body now
8254    /// calls this with `(0, n_layers)`, so the whole-trunk path and every stage range run the SAME
8255    /// code — there is no "split version" of the verify math.
8256    ///
8257    /// Bit-identity of a cut rests on the same kernel-check-pinned identity the eager arm's cut
8258    /// does — `add_rms_norm_q8_1 == add then rms_norm_q8_1` at nrows=T — because the ONLY thing a
8259    /// fence changes is that the cross-layer fusion carry breaks at `hi-1` and is re-materialized
8260    /// as an explicit `add`. `decode-batch-gate --mode ppspec` verifies end-to-end on real weights.
8261    #[allow(clippy::too_many_arguments)]
8262    fn verify_layers(
8263        &self,
8264        e: &Engine,
8265        mut x: CudaSlice<f32>,
8266        lo: usize,
8267        hi: usize,
8268        pos_d: &CudaSlice<i32>,
8269        pos0: usize,
8270        t: usize,
8271        cache: &mut Cache,
8272        mut ckpt: Option<&mut VerifyCkpt>,
8273        stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
8274        graphs: Option<&mut DsparkVerifyGraphs>,
8275    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8276        if self.sliding_gated_moe_batch_program() {
8277            if stream.is_some() {
8278                return Err(
8279                    "step35 has no ROUND-STREAM verify arm (the device-counter _dc twins \
8280                            cannot express the SWA offset KV view)"
8281                        .into(),
8282                );
8283            }
8284            return self.step35_verify_batch_layers(e, x, lo, hi, pos0, t, cache);
8285        }
8286        if self.batched_serving_numeric_class() {
8287            return self.qwen35_verify_batch_layers(
8288                e,
8289                x,
8290                lo,
8291                hi,
8292                pos0,
8293                t,
8294                cache,
8295                ckpt.take(),
8296                stream,
8297                graphs,
8298            );
8299        }
8300        let n_embd = self.cfg.n_embd as usize;
8301        let eps = self.cfg.rms_eps;
8302        // CROSS-LAYER ADD+NORM FUSION (lane/vt-fixes fix 2, mirroring decode_step_h's
8303        // launch-arc form): layer il's post-FFN residual add (x2 = x1 + ffn_out) and layer
8304        // il+1's attn_norm(+quantize) are consecutive row-wise ops — ONE add_rms_norm_q8_1
8305        // launch at nrows=t does all three (bit-identity pinned by the T-row kernel-check
8306        // arms). Carry the un-added (x1, ffn_out) pair; the fused launch materializes x2 (the
8307        // residual the next layer needs) as its `res` output. Falls back to the separate add
8308        // when the next layer is off the fused-q8 path.
8309        let mut pending: Option<(CudaSlice<f32>, CudaSlice<f32>)> = None;
8310        for il in lo..hi {
8311            let layer = &self.layers[il];
8312            // DISPATCH-MIRRORED attn-input RMSNorm (FP-order lesson #8): eager decode fuses the
8313            // 1024-thread rms_norm_q8_1 ONLY when every mixer projection is q8_1-fast; layers with
8314            // Float projections (ssm_beta/ssm_alpha on layers 1/2/4 of the 9B NVFP4 GGUF) take the
8315            // UNFUSED 256-thread rms_norm. The verify norm must mirror that PER-LAYER choice —
8316            // blockDim changes the sum-of-squares reduce order, and the ULP shift amplifies through
8317            // the GDN recurrence into argmax flips (measured: 9B text prompt, 1 ULP at layer 2 ->
8318            // 2.3e-1 logit maxdiff at the head -> K=1..8 divergence at a 0.03-margin token).
8319            let mixer_fast = self.mixer_in_q8_1_fast(e, &layer.mixer);
8320            let norm_fused = std::env::var("MEMRA_NO_FUSE_NORMQ").is_err() && mixer_fast;
8321            // BATCHED EPILOGUE RE-FUSE (lane/vt-fixes fix 2, 2026-08-03): when the norm is
8322            // dispatch-fused AND every consumer of `h` reads only its q8_1 form (Full mixer:
8323            // projections only; Linear mixer: the batched arm — the per-column fallback needs
8324            // f32 h), emit the attn-input norm DIRECTLY as q8_1 via `rms_norm_q8_1` at nrows=t
8325            // (row-indexed kernel — the T-row launch is the per-row m=1 program, kernel-check
8326            // pins bit-identity vs rms_norm_decode -> quantize_q8_1). Kills the standalone
8327            // quantize launch(es) + the f32 h HBM round-trip that decode never pays.
8328            // step35 (Full mixer) is the third case that needs f32 `h`: its verify arm is a
8329            // per-ROW replay of the eager decode mixer, whose `pre_q` contract is a single row —
8330            // a T-row q8_1 pair cannot be handed to it, and re-deriving per-row q8_1 from the f32
8331            // rows is exactly the dispatch being mirrored. Keep step35 on the unfused arm.
8332            let lin_q8_only = match &layer.mixer {
8333                Mixer::Linear(la) => {
8334                    (t >= 3 || (t == 2 && spec_m2())) && e.uses_q8_1_fast(&la.ssm_out)
8335                }
8336                Mixer::Full(_) if self.sliding_gated_moe_batch_program() => false,
8337                _ => true,
8338            };
8339            // NOTE decode.rs's take()-first lesson: take the pending pair BEFORE branching so
8340            // a non-fused layer still performs the residual add.
8341            let taken = pending.take();
8342            let (h, h_q8) = if norm_fused && lin_q8_only {
8343                let pair = match taken {
8344                    // fused add + attn_norm + q8_1: ONE launch resolves the carried residual
8345                    // AND emits this layer's mixer input pre-quantized. res -> x2 (= new x).
8346                    Some((x1p, f1p)) => {
8347                        let mut x2 = vbuf(e, t * n_embd)?; // fully written (res output)
8348                        let p = e.add_rms_norm_q8_1(
8349                            &x1p,
8350                            &f1p,
8351                            layer.attn_norm.float_data(),
8352                            &mut x2,
8353                            n_embd,
8354                            t,
8355                            eps,
8356                        )?;
8357                        x = x2;
8358                        p
8359                    }
8360                    None => e.rms_norm_q8_1(&x, layer.attn_norm.float_data(), n_embd, t, eps)?,
8361                };
8362                (e.zeros(0)?, Some(pair)) // h unused on this path (q8-only consumers)
8363            } else {
8364                if let Some((x1p, f1p)) = taken {
8365                    let mut x2 = vbuf(e, t * n_embd)?; // fully written by add
8366                    e.add(&x1p, &f1p, &mut x2, t * n_embd)?;
8367                    x = x2;
8368                }
8369                let mut h = vbuf(e, t * n_embd)?; // fully written by either rms_norm arm
8370                if norm_fused {
8371                    e.rms_norm_decode(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
8372                } else {
8373                    e.rms_norm(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
8374                }
8375                (h, None)
8376            };
8377            let h_q8_ref = h_q8.as_ref().map(|(q, d)| (q, d));
8378
8379            let mixed = match &layer.mixer {
8380                Mixer::Full(fa) => self.full_attn_verify(
8381                    e,
8382                    fa,
8383                    &h,
8384                    h_q8_ref,
8385                    pos_d,
8386                    t,
8387                    cache,
8388                    il,
8389                    stream.map(|(_, c)| c),
8390                )?,
8391                Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
8392                Mixer::Linear(la) => {
8393                    // BATCHED linear verify (2026-07-03, the MTP-profit lever): one T-token pass —
8394                    // batched projections (weight read ONCE, hits the m=2-4 weight-resident matvec),
8395                    // carried-state conv (ssm_conv1d_tm_state), GDN prep on the prefill kernels, and
8396                    // ONE gdn_scan whose internal sequential t-loop is the SAME recurrence as T
8397                    // chained T=1 steps (bit-identical). Falls back to the sequential per-column
8398                    // chain when T < d_conv-1 (conv ring update needs T >= pad) — or when ANY
8399                    // projection is off the q8_1 fast path: matmul_decode_exact would route a Float
8400                    // tensor to cuBLAS at m=t (different FP accumulation than eager's per-token
8401                    // GEMV), so mixed-dtype layers stay on the eager-identical per-column chain.
8402                    // MEMRA_SPEC_M2 (lane/spec-m2): the t==2 batch rides the same arm — the conv
8403                    // wrapper handles t<pad with a pure-copy ring rebuild; see spec_m2() header.
8404                    if (t >= 3 || (t == 2 && spec_m2()))
8405                        && mixer_fast
8406                        && e.uses_q8_1_fast(&la.ssm_out)
8407                    {
8408                        let want = ckpt.is_some();
8409                        let (out, stash) =
8410                            self.linear_attn_verify_t(e, la, &h, h_q8_ref, t, cache, il, want)?;
8411                        if let (Some(ck), Some(st)) = (ckpt.as_deref_mut(), stash) {
8412                            ck.gdn[il] = Some(st);
8413                        }
8414                        out
8415                    } else {
8416                        let mut out = vbuf(e, t * n_embd)?; // every col written by copy_into
8417                        let mut col_states: Option<Vec<(CudaSlice<f32>, CudaSlice<f32>)>> =
8418                            if ckpt.is_some() && t >= 2 {
8419                                Some(Vec::with_capacity(t - 1))
8420                            } else {
8421                                None
8422                            };
8423                        for col in 0..t {
8424                            let mut h_col = vbuf(e, n_embd)?; // fully written by copy_view_into
8425                            let src = h.slice(col * n_embd..(col + 1) * n_embd);
8426                            e.copy_view_into(&mut h_col, 0, &src, n_embd)?;
8427                            let m_col = self.linear_attn_decode(e, la, &h_col, cache, il)?;
8428                            e.copy_into(&mut out, col * n_embd, &m_col, n_embd)?;
8429                            // REPLAY-FREE ckpt: clone the chain's ACTUAL state after this column
8430                            // (pure dtod — cannot change any computed value). Last column skipped:
8431                            // rebuild targets are j <= t-1 columns.
8432                            if let Some(cs) = col_states.as_mut() {
8433                                if col + 1 < t {
8434                                    let rl = cache.recur[il].as_ref().unwrap();
8435                                    cs.push((
8436                                        e.clone_dtod(&rl.conv_state)?,
8437                                        e.clone_dtod(&rl.ssm_state)?,
8438                                    ));
8439                                }
8440                            }
8441                        }
8442                        if let (Some(ck), Some(cs)) = (ckpt.as_deref_mut(), col_states) {
8443                            // ReplaySSM-assessment instrumentation (2026-07-30): the
8444                            // per-column clones are the only true state snapshots left in
8445                            // the verify (the batched path stashes INPUTS and replays).
8446                            if std::env::var("MEMRA_SPEC_STATS").as_deref() == Ok("1") {
8447                                static ONCE: std::sync::Once = std::sync::Once::new();
8448                                let bytes: usize =
8449                                    cs.iter().map(|(c, s)| (c.len() + s.len()) * 4).sum();
8450                                ONCE.call_once(|| eprintln!(
8451                                    "[verify-ckpt] per-column layer il={il}: {} clones, {:.2} MB/layer/round",
8452                                    cs.len(), bytes as f64 / 1e6));
8453                            }
8454                            ck.cols[il] = Some(cs);
8455                        }
8456                        out
8457                    }
8458                }
8459            };
8460
8461            // DISPATCH-MIRRORED post-attn norm: eager residual_norm_ffn fuses add+norm+quant
8462            // (1024-thread add_rms_norm_q8_1) only for Dense FFNs whose gate+up are q8_1-fast;
8463            // otherwise (and for MoE) it runs the 256-thread fused add_rms_norm. Mirror per layer.
8464            let ffn_fuse = match &layer.ffn {
8465                crate::hybrid::Ffn::Dense {
8466                    ffn_gate, ffn_up, ..
8467                } => {
8468                    std::env::var("MEMRA_NO_FUSE_NORMQ").is_err()
8469                        && e.uses_q8_1_fast(ffn_gate)
8470                        && e.uses_q8_1_fast(ffn_up)
8471                }
8472                crate::hybrid::Ffn::Moe(_) => false,
8473            };
8474            // BATCHED EPILOGUE RE-FUSE (lane/vt-fixes fix 2): on the ffn_fuse path (Dense,
8475            // gate+up q8_1-fast, non-M3) the FFN input is emitted DIRECTLY as q8_1 by ONE
8476            // add_rms_norm_q8_1 launch at nrows=t (row-indexed kernel: the T-row launch is the
8477            // per-row m=1 program; kernel-check pins bit-identity vs the unfused
8478            // add_f32 -> rms_norm_decode -> quantize_q8_1 chain at T=2/4/5/8) — replacing the
8479            // add + rms_norm_decode launches AND the dual/singles' internal re-quantize.
8480            // M3's swigluoai must keep the f32 chain (the fused SwiGLU epilogue encodes plain
8481            // SiLU), mirroring residual_norm_ffn's m3 guard on the decode path.
8482            // step35: same guard per LAYER. A dense FFN's clamp is the SHEXP array (upstream's
8483            // one build_ffn serves dense + shared expert, llama-graph.cpp:1751), and verify MUST
8484            // mirror decode's dispatch or spec self-consistency fails.
8485            let dense_lim = self.cfg.clamp_shexp_at(il as u32);
8486            let fuse_q8 = ffn_fuse && self.cfg.m3.is_none() && dense_lim.is_none();
8487            let mut x1 = vbuf(e, t * n_embd)?; // fully written by add / add_rms_norm*
8488            let mut z = e.zeros(0)?; // replaced below on the unfused arms
8489            let z_q8 = if fuse_q8 {
8490                Some(e.add_rms_norm_q8_1(
8491                    &x,
8492                    &mixed,
8493                    layer.post_attn_norm.float_data(),
8494                    &mut x1,
8495                    n_embd,
8496                    t,
8497                    eps,
8498                )?)
8499            } else {
8500                let mut zf = vbuf(e, t * n_embd)?; // fully written by rms_norm_decode / add_rms_norm
8501                if ffn_fuse {
8502                    e.add(&x, &mixed, &mut x1, t * n_embd)?;
8503                    e.rms_norm_decode(
8504                        &x1,
8505                        layer.post_attn_norm.float_data(),
8506                        &mut zf,
8507                        n_embd,
8508                        t,
8509                        eps,
8510                    )?;
8511                } else {
8512                    e.add_rms_norm(
8513                        &x,
8514                        &mixed,
8515                        layer.post_attn_norm.float_data(),
8516                        &mut x1,
8517                        &mut zf,
8518                        n_embd,
8519                        t,
8520                        eps,
8521                    )?;
8522                }
8523                z = zf;
8524                None
8525            };
8526            // DECODE-EXACT FFN projections: force MMVQ for gate/up/down at any T to match the
8527            // T=1 decode FP accumulation order. At T>=5 the generic matmul/matmul_pre falls to dp4a
8528            // (128-thread, different FP sum order). At T=2-4 the batched MMVQ is already bit-identical.
8529            let ffn_out = match &layer.ffn {
8530                crate::hybrid::Ffn::Dense {
8531                    ffn_gate,
8532                    ffn_up,
8533                    ffn_down,
8534                } => {
8535                    let n_ff = ffn_gate.out_features();
8536                    if let Some((zq, zd)) = z_q8.as_ref() {
8537                        // FUSED CHAIN (fix 2): pre-quantized z feeds the projections; the SwiGLU
8538                        // epilogue emits act pre-quantized for ffn_down (silu_mul_scaled_q8_1,
8539                        // bit-identical to silu_mul + quantize — kernel-check-pinned) with the
8540                        // NVFP4 macro-scales folded (deferred-scale dual: y*s inline == the
8541                        // scale_inplace store, value-exact) — the exact m=1 decode epilogue
8542                        // structure at nrows=t.
8543                        let pair =
8544                            match e.matmul_decode_exact_dual_pre(ffn_gate, ffn_up, zq, zd, t)? {
8545                                Some(((g, gs), (u, us))) => Some((g, gs, u, us)),
8546                                None => None,
8547                            };
8548                        let (gate, gs, up, us) = match pair {
8549                            Some(x4) => x4,
8550                            None => (
8551                                e.matmul_decode_exact_pre(ffn_gate, zq, zd, t)?,
8552                                1.0, // scale already applied inside _pre
8553                                e.matmul_decode_exact_pre(ffn_up, zq, zd, t)?,
8554                                1.0,
8555                            ),
8556                        };
8557                        if e.uses_q8_1_fast(ffn_down) {
8558                            let (aq, ad) = e.silu_mul_scaled_q8_1(&gate, &up, gs, us, t * n_ff)?;
8559                            e.matmul_decode_exact_pre(ffn_down, &aq, &ad, t)?
8560                        } else {
8561                            let mut act = vbuf(e, t * n_ff)?;
8562                            e.silu_mul_scaled(&gate, &up, gs, us, &mut act, t * n_ff)?;
8563                            e.matmul_decode_exact(ffn_down, &act, t)?
8564                        }
8565                    } else {
8566                        // UNFUSED (pre-fix) chain — MoE-adjacent/M3/off-fast layers, unchanged.
8567                        // DUAL gate+up batched twin (lane/verify-economics, 2026-08-02): one launch
8568                        // for the pair at t=2..8 — bit-identical per (tensor,token,row) to the two
8569                        // singles (kernel-check pins bitwise; MEMRA_SPEC_DUAL_T=0 reverts). None
8570                        // (non-NVFP4 / t outside the tier / seam off) -> the two singles, unchanged.
8571                        let (gate, up) =
8572                            match e.matmul_decode_exact_dual(ffn_gate, ffn_up, &z, t)? {
8573                                Some(pair) => pair,
8574                                None => (
8575                                    e.matmul_decode_exact(ffn_gate, &z, t)?,
8576                                    e.matmul_decode_exact(ffn_up, &z, t)?,
8577                                ),
8578                            };
8579                        let mut act = vbuf(e, t * n_ff)?; // fully written by ffn_act_lim
8580                        Self::ffn_act_lim(
8581                            e,
8582                            &self.cfg,
8583                            &gate,
8584                            &up,
8585                            1.0,
8586                            1.0,
8587                            dense_lim,
8588                            &mut act,
8589                            t * n_ff,
8590                        )?;
8591                        e.matmul_decode_exact(ffn_down, &act, t)?
8592                    }
8593                }
8594                crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il(e, m, &z, t, il as u16)?,
8595            };
8596            // CROSS-LAYER fusion: defer this layer's post-FFN residual add — the next layer's
8597            // fused-q8 attn norm folds it in (add_rms_norm_q8_1 == add; rms_norm; quantize,
8598            // kernel-check-pinned at nrows=T). Non-fused next layers add explicitly above.
8599            pending = Some((x1, ffn_out));
8600        }
8601        // RANGE's final add (no next norm INSIDE the range to fuse with; for the
8602        // whole-trunk call that is the last layer, whose next norm is output_norm — f32-out).
8603        if let Some((x1p, f1p)) = pending.take() {
8604            let mut x2 = vbuf(e, t * n_embd)?; // fully written by add
8605            e.add(&x1p, &f1p, &mut x2, t * n_embd)?;
8606            x = x2;
8607        }
8608        Ok(x)
8609    }
8610    /// BATCHED linear-attn verify (T=K+1): the whole layer in ~10 launches instead of T x the
8611    /// T=1 decode chain (T x ~12 launches + T weight reads of the four projections). The GDN
8612    /// recurrence itself is inherently sequential — gdn_scan_s128 runs its internal t-loop with
8613    /// the SAME per-token math as chained T=1 calls (bit-identical state evolution); everything
8614    /// around it (projections, conv, prep, gated norm, out-proj) batches. Advances conv ring +
8615    /// ssm state exactly like T sequential decode steps.
8616    /// `want_stash`: additionally RETAIN the gdn-scan inputs (pure buffer keep-alives, zero extra
8617    /// kernels) so a partial accept can rebuild the state after any column prefix (REPLAY-FREE).
8618    #[allow(clippy::too_many_arguments)]
8619    fn linear_attn_verify_t(
8620        &self,
8621        e: &Engine,
8622        la: &LinearAttnLayer,
8623        h: &CudaSlice<f32>,
8624        h_q8: Option<(&CudaSlice<i8>, &CudaSlice<f32>)>,
8625        t: usize,
8626        cache: &mut Cache,
8627        il: usize,
8628        want_stash: bool,
8629    ) -> Result<(CudaSlice<f32>, Option<GdnStash>), Box<dyn std::error::Error>> {
8630        let cfg = &self.cfg;
8631        let geometry = la.geometry;
8632        let d_state = geometry.key_head_dim as usize;
8633        let num_k = geometry.key_heads as usize;
8634        let num_v = geometry.value_heads as usize;
8635        let d_conv = geometry.conv_kernel as usize;
8636        let key_dim = d_state * num_k;
8637        let conv_dim = key_dim * 2 + geometry.value_head_dim as usize * num_v;
8638        let eps = cfg.rms_eps;
8639        let scale = 1.0 / (d_state as f32).sqrt();
8640
8641        // DECODE-EXACT projections: matmul_decode_exact forces the MMVQ (warp-per-row, 32-thread)
8642        // accumulation order for EVERY m, matching the T=1 decode path bit-for-bit. The generic
8643        // `matmul` at m>=5 falls to dp4a (128-thread, two-level reduce) which has a different FP
8644        // sum order — ULP differences propagate through gdn_scan and flip argmax on the 27B.
8645        // Q8 TRUNK-FUSION at T=1 (35B: wqkv+wqkv_gate both Q8_0): one fused2 launch, bit-identical
8646        // per (tensor,row) to the two m=1 MMVQ dispatches below — decode-exact contract holds.
8647        // VERIFY-TIER TRUNK FUSION (MEMRA_SPEC_FUSED_T, t=2-4): quantize h ONCE for every
8648        // fused-eligible same-input Q8_0 pair of this layer (35B wqkv+wqkv_gate; 9B
8649        // ssm_beta+ssm_alpha) — each fused2 batched launch then replaces two decode-exact
8650        // calls (each of which re-quantizes the same h + runs its own _b2/_b4 launch).
8651        // Bit-identical per (tensor,token,row) — see spec_fused_t().
8652        // BATCHED EPILOGUE RE-FUSE (lane/vt-fixes fix 2): `h_q8` = the attn-input norm emitted
8653        // directly as q8_1 by the caller's fused rms_norm_q8_1 (bit-identical to the unfused
8654        // chain, kernel-check-pinned). When present it REPLACES the standalone quantize below
8655        // and feeds every projection; the caller guaranteed all four input projections are
8656        // q8_1-fast. When absent, the old shared-quantize (fused-t window) stands.
8657        let h_q8_t = if h_q8.is_none()
8658            && spec_fused_t()
8659            && (2..=4).contains(&t)
8660            && ((e.uses_q8_1_fast(&la.wqkv) && e.uses_q8_1_fast(&la.wqkv_gate))
8661                || (e.uses_q8_1_fast(&la.ssm_beta) && e.uses_q8_1_fast(&la.ssm_alpha)))
8662        {
8663            Some(e.quantize_q8_1(h, t, cfg.n_embd as usize)?)
8664        } else {
8665            None
8666        };
8667        // one view: the caller's fused-norm q8 or this fn's own shared quantize.
8668        let hq8_any: Option<(&CudaSlice<i8>, &CudaSlice<f32>)> =
8669            h_q8.or(h_q8_t.as_ref().map(|(q, d)| (q, d)));
8670        let (qkv_mixed, z) = {
8671            let mut fused = None;
8672            if t == 1 && e.uses_q8_1_fast(&la.wqkv) && e.uses_q8_1_fast(&la.wqkv_gate) {
8673                let (hq, hd) = e.quantize_q8_1(h, 1, cfg.n_embd as usize)?;
8674                fused = e.matmul_q8_fused2(&la.wqkv, &la.wqkv_gate, &hq, &hd)?;
8675            } else if let Some((hq, hd)) = hq8_any {
8676                if spec_fused_t() && (2..=4).contains(&t) {
8677                    fused = e.matmul_q8_fused2_t(&la.wqkv, &la.wqkv_gate, hq, hd, t)?;
8678                }
8679            }
8680            match (fused, hq8_any) {
8681                (Some(pair), _) => pair,
8682                (None, Some((hq, hd))) if h_q8.is_some() => (
8683                    e.matmul_decode_exact_pre(&la.wqkv, hq, hd, t)?,
8684                    e.matmul_decode_exact_pre(&la.wqkv_gate, hq, hd, t)?,
8685                ),
8686                (None, _) => (
8687                    e.matmul_decode_exact(&la.wqkv, h, t)?,
8688                    e.matmul_decode_exact(&la.wqkv_gate, h, t)?,
8689                ),
8690            }
8691        };
8692        // beta+alpha DUAL at T=1 (75% of p3 rounds run T=1 verify — p-min chain cuts): the dual
8693        // mr2 kernel is bit-identical per element to the m=1 MMVQ matmul_decode_exact dispatches
8694        // (same warp-per-row body, blockIdx.y picks the weight), so the decode-exact contract
8695        // holds; the run-spec battery is the arbiter. T>1 keeps the per-tensor decode-exact path.
8696        let (beta_raw, alpha) = if t == 1 {
8697            let (hq, hd) = e.quantize_q8_1(h, 1, cfg.n_embd as usize)?;
8698            match e.matmul_pre_dual_noscale(&la.ssm_beta, &la.ssm_alpha, &hq, &hd, 1)? {
8699                Some(((mut b, bs), (mut a, as_))) => {
8700                    if bs != 1.0 {
8701                        e.scale_inplace(&mut b, bs, la.ssm_beta.out_features())?;
8702                    }
8703                    if as_ != 1.0 {
8704                        e.scale_inplace(&mut a, as_, la.ssm_alpha.out_features())?;
8705                    }
8706                    (b, a)
8707                }
8708                // Q8_0 fused2 twin (9B stores beta/alpha as Q8_0): DISPATCH-MIRRORS the eager
8709                // decode's beta_alpha closure — the fused body is qmatvec_q8_0_mmvq verbatim,
8710                // bit-identical per row (kernel-check rel=0.00e0 gate), so decode==verify holds.
8711                None => match e.matmul_q8_fused2(&la.ssm_beta, &la.ssm_alpha, &hq, &hd)? {
8712                    Some((b, a)) => (b, a),
8713                    None => (
8714                        e.matmul_decode_exact(&la.ssm_beta, h, 1)?,
8715                        e.matmul_decode_exact(&la.ssm_alpha, h, 1)?,
8716                    ),
8717                },
8718            }
8719        } else {
8720            // fused-t twin (9B stores beta/alpha as Q8_0): same shared-quantize + one launch
8721            // contract as the wqkv pair above; 35B beta/alpha are Float -> None -> fallback.
8722            let mut nvfp4_fused = None;
8723            let mut q8_fused = None;
8724            if let Some((hq, hd)) = hq8_any {
8725                if t == 3 && std::env::var("MEMRA_NVFP4_AUX_DUAL").as_deref() != Ok("0") {
8726                    nvfp4_fused =
8727                        e.matmul_decode_exact_dual_pre(&la.ssm_beta, &la.ssm_alpha, hq, hd, t)?;
8728                    if nvfp4_fused.is_some() && std::env::var("MEMRA_DEBUG").is_ok() {
8729                        static ONCE: std::sync::Once = std::sync::Once::new();
8730                        ONCE.call_once(|| {
8731                            eprintln!("[memra] NVFP4 beta+alpha batched aux dual ENGAGED (t={t})")
8732                        });
8733                    }
8734                }
8735                if nvfp4_fused.is_none() && spec_fused_t() && (2..=4).contains(&t) {
8736                    q8_fused = e.matmul_q8_fused2_t(&la.ssm_beta, &la.ssm_alpha, hq, hd, t)?;
8737                }
8738            }
8739            if let Some(((mut b, bs), (mut a, as_))) = nvfp4_fused {
8740                if bs != 1.0 {
8741                    e.scale_inplace(&mut b, bs, t * la.ssm_beta.out_features())?;
8742                }
8743                if as_ != 1.0 {
8744                    e.scale_inplace(&mut a, as_, t * la.ssm_alpha.out_features())?;
8745                }
8746                (b, a)
8747            } else if let Some(pair) = q8_fused {
8748                pair
8749            } else {
8750                match hq8_any {
8751                    Some((hq, hd)) if h_q8.is_some() => (
8752                        e.matmul_decode_exact_pre(&la.ssm_beta, hq, hd, t)?,
8753                        e.matmul_decode_exact_pre(&la.ssm_alpha, hq, hd, t)?,
8754                    ),
8755                    _ => (
8756                        e.matmul_decode_exact(&la.ssm_beta, h, t)?,
8757                        e.matmul_decode_exact(&la.ssm_alpha, h, t)?,
8758                    ),
8759                }
8760            }
8761        };
8762
8763        // conv with CARRIED state + ring roll (T >= pad rides the input-column update kernel;
8764        // T < pad — the MEMRA_SPEC_M2 t=2 arm — rolls via the pure-copy ring rebuild).
8765        let rl = cache.recur[il].as_mut().unwrap();
8766        let mut conv_out = e.uninit(conv_dim * t)?;
8767        e.ssm_conv1d_tm_state(
8768            &qkv_mixed,
8769            &mut rl.conv_state,
8770            la.ssm_conv1d.float_data(),
8771            &mut conv_out,
8772            conv_dim,
8773            t,
8774            d_conv,
8775        )?;
8776
8777        // GDN prep via the prefill kernels (repack + L2 + sigmoid + glog), T-wide.
8778        let mut q_g = e.uninit(d_state * num_v * t)?;
8779        let mut k_g = e.uninit(d_state * num_v * t)?;
8780        let mut v_g = e.uninit(d_state * num_v * t)?;
8781        e.qkv_to_gdn_repack(
8782            &conv_out, &mut q_g, &mut k_g, &mut v_g, d_state, num_v, num_k, key_dim, t,
8783        )?;
8784        let mut q_l2 = e.uninit(d_state * num_v * t)?;
8785        e.l2_norm_decode(&q_g, &mut q_l2, d_state, num_v * t, eps)?;
8786        let mut k_l2 = e.uninit(d_state * num_v * t)?;
8787        e.l2_norm_decode(&k_g, &mut k_l2, d_state, num_v * t, eps)?;
8788        let mut beta = e.uninit(t * num_v)?;
8789        e.sigmoid(&beta_raw, &mut beta, t * num_v)?;
8790        let mut g_log = e.uninit(t * num_v)?;
8791        e.gdn_glog(
8792            &alpha,
8793            la.ssm_dt.float_data(),
8794            la.ssm_a.float_data(),
8795            &mut g_log,
8796            num_v,
8797            t,
8798        )?;
8799
8800        // ONE gdn_scan over T tokens from the carried state (internal sequential loop ==
8801        // T chained T=1 steps). Ping-pong the resident buffers like eager decode.
8802        let mut o = e.uninit(d_state * num_v * t)?;
8803        {
8804            let crate::cache::RecurLayer {
8805                ssm_state,
8806                ssm_state_alt,
8807                ..
8808            } = rl;
8809            e.gdn_scan_s128(
8810                &q_l2,
8811                &k_l2,
8812                &v_g,
8813                &g_log,
8814                &beta,
8815                ssm_state,
8816                ssm_state_alt,
8817                &mut o,
8818                num_v,
8819                t,
8820                scale,
8821            )?;
8822        }
8823        std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
8824
8825        // gated RMSNorm + out projection, T-wide. FUSED-QUANTIZE ARM (lane/vt-fixes fix 2,
8826        // mirroring the T=1 decode's launch-arc form): when ssm_out rides the q8_1 fast path,
8827        // emit q8_1 straight from the gated norm at nrows=num_v*t (row-indexed kernel, the
8828        // T-wide launch is the per-row program; kernel-check pins bit-identity vs
8829        // gated_rmsnorm -> quantize_q8_1 at T=1 and T=5) and feed the decode-exact dispatch
8830        // pre-quantized — one launch replaces norm + quantize. Fallback = the f32 chain.
8831        let out = if e.uses_q8_1_fast(&la.ssm_out) {
8832            let (gq, gd) =
8833                e.gated_rmsnorm_q8_1(&o, la.ssm_norm.float_data(), &z, d_state, num_v * t, eps)?;
8834            e.matmul_decode_exact_pre(&la.ssm_out, &gq, &gd, t)?
8835        } else {
8836            let mut gn = e.uninit(d_state * num_v * t)?;
8837            e.gated_rmsnorm(
8838                &o,
8839                la.ssm_norm.float_data(),
8840                &z,
8841                &mut gn,
8842                d_state,
8843                num_v * t,
8844                eps,
8845            )?;
8846            // DECODE-EXACT out-projection: same MMVQ path as the T=1 decode (ssm_out at m>=5
8847            // would fall to dp4a with a different FP reduction order — same class of bug as
8848            // the input projs).
8849            e.matmul_decode_exact(&la.ssm_out, &gn, t)?
8850        };
8851        let stash = if want_stash {
8852            Some(GdnStash {
8853                qkv_mixed,
8854                q_l2,
8855                k_l2,
8856                v_g,
8857                g_log,
8858                beta,
8859            })
8860        } else {
8861            None
8862        };
8863        Ok((out, stash))
8864    }
8865
8866    /// REPLAY-FREE partial-accept commit (2026-07-03): make the cache state == "committed through
8867    /// the first `j` verify columns" WITHOUT the legacy rollback + duplicate trunk replay.
8868    /// - Full-attn KV: truncate both the owning-stage shadow and every TP rank to snapshot + j.
8869    ///   The verify's appended rows for those columns are bit-identical to what an eager T=1
8870    ///   chain writes (the decode-exact contract the verify-probe gates), so keeping them ==
8871    ///   replaying them.
8872    /// - Linear layers, batched path: rebuild the conv ring by PURE COPIES (ring holds raw input
8873    ///   columns) and the ssm state by a prefix re-run of the SAME gdn_scan kernel (t=j) from the
8874    ///   snapshot state over the stash's identical inputs — the kernel's t-loop carries state in
8875    ///   registers and writes it once at the end, so iterations 0..j-1 are independent of T:
8876    ///   bit-identical to the verify's own state after j tokens == the eager chain state.
8877    /// - Linear layers, per-column path: restore the cloned actual state after column j-1.
8878    /// Caller guarantees 1 <= j <= t-1 (j==0 rounds take the legacy rollback; j==t is full accept).
8879    fn commit_verified_prefix(
8880        &self,
8881        e: &Engine,
8882        cache: &mut Cache,
8883        snap: &crate::cache::CacheSnapshot,
8884        ckpt: &VerifyCkpt,
8885        j: usize,
8886        kv_lens_done: bool,
8887        dev_j: Option<(&CudaSlice<u32>, usize, usize)>,
8888    ) -> Result<(), Box<dyn std::error::Error>> {
8889        // GDN geometry derives lazily inside recurrent-layer arms. Full-attention plans carry no
8890        // recurrent state and must never be forced through a synthetic SSM geometry.
8891        // Engine-bundle slice 1 (DSF-ROUNDCOST-20260820 §1.1): the per-column-arm restores
8892        // are 2 tiny D2D copies per linear layer (~96 dispatches/partial round on the q38
8893        // route). When every cols-arm layer shares uniform state sizes (single ssm cfg —
8894        // always true today), batch them into two `copy_batch_uniform_f32` launches. Bytes,
8895        // buffers and stream order are identical to the per-layer memcpy sequence; the
8896        // kernel-rebuild (gdn-stash) arm below is untouched. MEMRA_STATE_COPY_BATCH=0 reverts.
8897        let mut batched_cols = false;
8898        if state_copy_batch_on() && dev_j.is_none() {
8899            use cudarc::driver::DevicePtr;
8900            let s = &e.gpu.stream();
8901            let mut conv_pairs: Vec<(u64, u64)> = Vec::new();
8902            let mut ssm_pairs: Vec<(u64, u64)> = Vec::new();
8903            let (mut conv_words, mut ssm_words) = (0usize, 0usize);
8904            let mut uniform = true;
8905            for il in 0..self.layers.len() {
8906                let Some(rl) = cache.recur[il].as_ref() else {
8907                    continue;
8908                };
8909                if ckpt.gdn[il].is_some() {
8910                    continue; // kernel-rebuild arm restores below, per layer
8911                }
8912                let Some(cols) = &ckpt.cols[il] else {
8913                    continue; // missing-ckpt error surfaces in the main loop
8914                };
8915                let (c, st) = &cols[j - 1];
8916                if conv_pairs.is_empty() {
8917                    conv_words = c.len();
8918                    ssm_words = st.len();
8919                } else if c.len() != conv_words || st.len() != ssm_words {
8920                    uniform = false;
8921                    break;
8922                }
8923                let (pc, _g0) = c.device_ptr(s);
8924                let (dc, _g1) = rl.conv_state.device_ptr(s);
8925                let (ps, _g2) = st.device_ptr(s);
8926                let (ds, _g3) = rl.ssm_state.device_ptr(s);
8927                conv_pairs.push((pc as u64, dc as u64));
8928                ssm_pairs.push((ps as u64, ds as u64));
8929            }
8930            if uniform && !conv_pairs.is_empty() {
8931                let n = conv_pairs.len();
8932                let mut t = vec![0u64; 2 * n];
8933                for (k, &(src, dst)) in conv_pairs.iter().enumerate() {
8934                    t[k] = src;
8935                    t[n + k] = dst;
8936                }
8937                let conv_t = e.htod_u64(&t)?;
8938                for (k, &(src, dst)) in ssm_pairs.iter().enumerate() {
8939                    t[k] = src;
8940                    t[n + k] = dst;
8941                }
8942                let ssm_t = e.htod_u64(&t)?;
8943                e.copy_batch_uniform_f32(&conv_t, n, conv_words)?;
8944                e.copy_batch_uniform_f32(&ssm_t, n, ssm_words)?;
8945                batched_cols = true;
8946            }
8947        }
8948        rewind_tp_kv_verified_prefix(&mut cache.tp_kv, &snap.tp_kv_len, j)?;
8949        for il in 0..self.layers.len() {
8950            if let (Some(kvl), Some(saved)) = (cache.kv[il].as_mut(), snap.kv_len[il]) {
8951                kvl.len = saved + j;
8952                // devacc 3a: spec_rollback_kv already wrote len_d on-device (same value).
8953                if !kv_lens_done {
8954                    e.set_i32_one(&mut kvl.len_d, kvl.len as i32)?;
8955                }
8956            }
8957            if let Some(rl) = cache.recur[il].as_mut() {
8958                let Mixer::Linear(linear) = &self.layers[il].mixer else {
8959                    return Err(format!("recurrent cache layer {il} has no GDN plan").into());
8960                };
8961                let geometry = linear.geometry;
8962                let d_state = geometry.key_head_dim as usize;
8963                let num_k = geometry.key_heads as usize;
8964                let num_v = geometry.value_heads as usize;
8965                let d_conv = geometry.conv_kernel as usize;
8966                let conv_dim = d_state * num_k * 2 + geometry.value_head_dim as usize * num_v;
8967                let scale = 1.0 / (d_state as f32).sqrt();
8968                if let Some(st) = &ckpt.gdn[il] {
8969                    let ring_old = snap.conv[il].as_ref().expect("snapshot missing conv");
8970                    let state_in = snap.ssm[il].as_ref().expect("snapshot missing ssm");
8971                    if let Some((acc, base, t_v)) = dev_j {
8972                        // 3b: j read on-device (_dc twins, same bodies; full accept early-exits).
8973                        e.ssm_conv_ring_rebuild_dc(
8974                            &st.qkv_mixed,
8975                            ring_old,
8976                            &mut rl.conv_state,
8977                            conv_dim,
8978                            acc,
8979                            base,
8980                            t_v,
8981                            d_conv,
8982                        )?;
8983                        let mut o = e.uninit(d_state * num_v * j.max(1))?;
8984                        e.gdn_scan_s128_dc(
8985                            &st.q_l2,
8986                            &st.k_l2,
8987                            &st.v_g,
8988                            &st.g_log,
8989                            &st.beta,
8990                            state_in,
8991                            &mut rl.ssm_state,
8992                            &mut o,
8993                            num_v,
8994                            acc,
8995                            base,
8996                            t_v,
8997                            scale,
8998                        )?;
8999                    } else {
9000                        e.ssm_conv_ring_rebuild(
9001                            &st.qkv_mixed,
9002                            ring_old,
9003                            &mut rl.conv_state,
9004                            conv_dim,
9005                            j,
9006                            d_conv,
9007                        )?;
9008                        let mut o = e.uninit(d_state * num_v * j)?; // scan output, discarded
9009                        e.gdn_scan_s128(
9010                            &st.q_l2,
9011                            &st.k_l2,
9012                            &st.v_g,
9013                            &st.g_log,
9014                            &st.beta,
9015                            state_in,
9016                            &mut rl.ssm_state,
9017                            &mut o,
9018                            num_v,
9019                            j,
9020                            scale,
9021                        )?;
9022                    }
9023                } else if let Some(cols) = &ckpt.cols[il] {
9024                    if !batched_cols {
9025                        let (c, s) = &cols[j - 1];
9026                        e.copy_into(&mut rl.conv_state, 0, c, c.len())?;
9027                        e.copy_into(&mut rl.ssm_state, 0, s, s.len())?;
9028                    }
9029                } else {
9030                    return Err(
9031                        "commit_verified_prefix: verify ckpt missing for linear layer".into(),
9032                    );
9033                }
9034            }
9035        }
9036        cache.pos = snap.pos + j;
9037        Ok(())
9038    }
9039
9040    /// ROUND-STREAM: recur restore with device-j (the _dc twins; full accept early-exits
9041    /// in-kernel). Requires the batched-linear stash on every linear layer (stream gate).
9042    fn commit_verified_prefix_stream(
9043        &self,
9044        e: &Engine,
9045        cache: &mut Cache,
9046        snap: &crate::cache::CacheSnapshot,
9047        ckpt: &VerifyCkpt,
9048        acc: &CudaSlice<u32>,
9049        base: usize,
9050        t_v: usize,
9051    ) -> Result<(), Box<dyn std::error::Error>> {
9052        for il in 0..self.layers.len() {
9053            if let Some(rl) = cache.recur[il].as_mut() {
9054                let Mixer::Linear(linear) = &self.layers[il].mixer else {
9055                    return Err(format!("recurrent cache layer {il} has no GDN plan").into());
9056                };
9057                let geometry = linear.geometry;
9058                let d_state = geometry.key_head_dim as usize;
9059                let num_k = geometry.key_heads as usize;
9060                let num_v = geometry.value_heads as usize;
9061                let d_conv = geometry.conv_kernel as usize;
9062                let conv_dim = d_state * num_k * 2 + geometry.value_head_dim as usize * num_v;
9063                let scale = 1.0 / (d_state as f32).sqrt();
9064                let st = ckpt.gdn[il]
9065                    .as_ref()
9066                    .ok_or("stream restore: batched-linear stash missing")?;
9067                let ring_old = snap.conv[il].as_ref().expect("snapshot missing conv");
9068                let state_in = snap.ssm[il].as_ref().expect("snapshot missing ssm");
9069                e.ssm_conv_ring_rebuild_dc(
9070                    &st.qkv_mixed,
9071                    ring_old,
9072                    &mut rl.conv_state,
9073                    conv_dim,
9074                    acc,
9075                    base,
9076                    t_v,
9077                    d_conv,
9078                )?;
9079                let mut o = e.uninit(d_state * num_v * t_v)?;
9080                e.gdn_scan_s128_dc(
9081                    &st.q_l2,
9082                    &st.k_l2,
9083                    &st.v_g,
9084                    &st.g_log,
9085                    &st.beta,
9086                    state_in,
9087                    &mut rl.ssm_state,
9088                    &mut o,
9089                    num_v,
9090                    acc,
9091                    base,
9092                    t_v,
9093                    scale,
9094                )?;
9095            }
9096        }
9097        Ok(())
9098    }
9099
9100    /// EAGLE3 aux-capturing verify forward over `tokens` (T) — mirrors `decode_step_t_h` exactly
9101    /// (same KV append, same causal verify, same recur advance) but ALSO clones the aux residual-
9102    /// stream hiddens (blocks in `aux_layers`) for TWO columns: the LAST column (always) and the
9103    /// optional `pred_col` (the EAGLE seed = bonus's predecessor). Returns
9104    /// (all_T_logits host, last_col_aux, pred_col_aux?). Used by the EAGLE3 orchestrator's commit.
9105    pub fn decode_step_t_aux2(
9106        &self,
9107        e: &Engine,
9108        tokens: &[u32],
9109        pos0: usize,
9110        cache: &mut Cache,
9111        aux_layers: &[usize],
9112        pred_col: Option<usize>,
9113    ) -> Result<
9114        (Vec<f32>, Vec<CudaSlice<f32>>, Option<Vec<CudaSlice<f32>>>),
9115        Box<dyn std::error::Error>,
9116    > {
9117        let cfg = &self.cfg;
9118        let n_embd = cfg.n_embd as usize;
9119        let eps = cfg.rms_eps;
9120        let t = tokens.len();
9121        let pos_vec: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
9122        let pos_d = e.htod_i32(&pos_vec)?;
9123        let mut x = e.htod(&self.embd.gather(n_embd, tokens))?;
9124        let mut aux_last: Vec<CudaSlice<f32>> = Vec::with_capacity(aux_layers.len());
9125        let mut aux_pred: Vec<CudaSlice<f32>> = Vec::new();
9126        let want_pred = pred_col.is_some();
9127
9128        for (il, layer) in self.layers.iter().enumerate() {
9129            // DISPATCH-MIRRORED norms (FP-order lesson #8) — see decode_step_t_h_emb.
9130            let mixer_fast = self.mixer_in_q8_1_fast(e, &layer.mixer);
9131            let norm_fused = std::env::var("MEMRA_NO_FUSE_NORMQ").is_err() && mixer_fast;
9132            let mut h = vbuf(e, t * n_embd)?; // fully written by either rms_norm arm
9133            if norm_fused {
9134                e.rms_norm_decode(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
9135            } else {
9136                e.rms_norm(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
9137            }
9138            let mixed = match &layer.mixer {
9139                Mixer::Full(fa) => {
9140                    self.full_attn_verify(e, fa, &h, None, &pos_d, t, cache, il, None)?
9141                }
9142                Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
9143                Mixer::Linear(la) => {
9144                    let mut out = e.zeros(t * n_embd)?;
9145                    for col in 0..t {
9146                        let mut h_col = e.zeros(n_embd)?;
9147                        let src = h.slice(col * n_embd..(col + 1) * n_embd);
9148                        e.copy_view_into(&mut h_col, 0, &src, n_embd)?;
9149                        let m_col = self.linear_attn_decode(e, la, &h_col, cache, il)?;
9150                        e.copy_into(&mut out, col * n_embd, &m_col, n_embd)?;
9151                    }
9152                    out
9153                }
9154            };
9155            let ffn_fuse = match &layer.ffn {
9156                crate::hybrid::Ffn::Dense {
9157                    ffn_gate, ffn_up, ..
9158                } => {
9159                    std::env::var("MEMRA_NO_FUSE_NORMQ").is_err()
9160                        && e.uses_q8_1_fast(ffn_gate)
9161                        && e.uses_q8_1_fast(ffn_up)
9162                }
9163                crate::hybrid::Ffn::Moe(_) => false,
9164            };
9165            let mut x1 = vbuf(e, t * n_embd)?; // fully written by add / add_rms_norm
9166            let mut z = vbuf(e, t * n_embd)?; // fully written by rms_norm_decode / add_rms_norm
9167            if ffn_fuse {
9168                e.add(&x, &mixed, &mut x1, t * n_embd)?;
9169                e.rms_norm_decode(
9170                    &x1,
9171                    layer.post_attn_norm.float_data(),
9172                    &mut z,
9173                    n_embd,
9174                    t,
9175                    eps,
9176                )?;
9177            } else {
9178                e.add_rms_norm(
9179                    &x,
9180                    &mixed,
9181                    layer.post_attn_norm.float_data(),
9182                    &mut x1,
9183                    &mut z,
9184                    n_embd,
9185                    t,
9186                    eps,
9187                )?;
9188            }
9189            let ffn_out = match &layer.ffn {
9190                crate::hybrid::Ffn::Dense {
9191                    ffn_gate,
9192                    ffn_up,
9193                    ffn_down,
9194                } => {
9195                    let n_ff = ffn_gate.out_features();
9196                    let gate = e.matmul_decode_exact(ffn_gate, &z, t)?;
9197                    let up = e.matmul_decode_exact(ffn_up, &z, t)?;
9198                    let mut act = vbuf(e, t * n_ff)?; // fully written by ffn_act_lim
9199                    // dense FFN clamp = the SHEXP array (upstream build_ffn serves both).
9200                    Self::ffn_act_lim(
9201                        e,
9202                        &self.cfg,
9203                        &gate,
9204                        &up,
9205                        1.0,
9206                        1.0,
9207                        self.cfg.clamp_shexp_at(il as u32),
9208                        &mut act,
9209                        t * n_ff,
9210                    )?;
9211                    e.matmul_decode_exact(ffn_down, &act, t)?
9212                }
9213                crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il(e, m, &z, t, il as u16)?,
9214            };
9215            let mut x2 = vbuf(e, t * n_embd)?; // fully written by add
9216            e.add(&x1, &ffn_out, &mut x2, t * n_embd)?;
9217            if aux_layers.contains(&il) {
9218                let mut a = e.zeros(n_embd)?;
9219                e.copy_view_into(&mut a, 0, &x2.slice((t - 1) * n_embd..t * n_embd), n_embd)?;
9220                aux_last.push(a);
9221                if let Some(pc) = pred_col {
9222                    let mut ap = e.zeros(n_embd)?;
9223                    e.copy_view_into(
9224                        &mut ap,
9225                        0,
9226                        &x2.slice(pc * n_embd..(pc + 1) * n_embd),
9227                        n_embd,
9228                    )?;
9229                    aux_pred.push(ap);
9230                }
9231            }
9232            x = x2;
9233        }
9234        let mut hn = vbuf(e, t * n_embd)?; // fully written by rms_norm_decode
9235        e.rms_norm_decode(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
9236        let logits = e.matmul_decode_exact(&self.output, &hn, t)?;
9237        let host = e.dtoh(&logits)?;
9238        cache.pos += t;
9239        Ok((
9240            host,
9241            aux_last,
9242            if want_pred { Some(aux_pred) } else { None },
9243        ))
9244    }
9245
9246    /// step35 SPEC-VERIFY attention over T query tokens — a per-row REPLAY of the eager
9247    /// `step35_decode_attn`.
9248    ///
9249    /// WHY A REPLAY AND NOT A BATCHED TWIN. The verify's whole job is to be bit-identical to what
9250    /// the eager decode would have computed for the same tokens; that is what makes greedy spec
9251    /// decode exact (run-spec asserts token identity for K=1..8). Every other verify arm in this
9252    /// file earns that identity by carefully mirroring dispatch (`matmul_decode_exact` to force
9253    /// MMVQ at any m, per-layer `ffn_fuse` mirroring, per-row `fa_decode` key bounds). step35
9254    /// stacks FOUR more per-layer degrees of freedom on top of that — per-layer `n_head`
9255    /// (64 full / 96 SWA), per-layer rotary width (64 full / 128 SWA), per-layer rope base, and a
9256    /// SEPARATE `attn_gate` tensor whose projection shares the attn-normed input — and its SWA
9257    /// layers attend through a token-OFFSET view whose offset is a function of the ABSOLUTE
9258    /// position of each query row. A batched twin would have to reproduce all of that AND the
9259    /// per-row offset in one launch; the offset alone rules out the existing rows kernels (they
9260    /// take one `base_len`, not a per-row offset).
9261    ///
9262    /// So this arm calls the eager path itself, once per row, on the same cache. Identity is then
9263    /// true BY CONSTRUCTION rather than by mirroring: row r runs exactly the kernel sequence that
9264    /// eager decode step r runs (same projections, same q8_1 fusion decision, same append, same
9265    /// view arithmetic, same `fa_decode_kvmod`, same gate), because it IS that code. Cost: T x the
9266    /// eager decode mixer instead of one batched pass — the same trade the generic arm's `else`
9267    /// per-row loop already accepts when `fa_rows_eligible` says no. Correctness first; a batched
9268    /// step35 twin is a perf lane's job and must be gated against this arm.
9269    ///
9270    /// The `h_q8` pre-quantized pair from the caller's fused norm is NOT forwarded: it is a
9271    /// T-row buffer and `step35_decode_attn`'s `pre_q` contract is one row. Instead each row's
9272    /// f32 `h` slice is handed over and the callee re-derives its own q8_1 exactly as eager decode
9273    /// does (`quantize_q8_1(h, 1, n_embd)`) — which is the dispatch being mirrored. Callers that
9274    /// took the fused arm therefore MUST still pass a live `h`; `step35_verify` asserts that.
9275    #[allow(clippy::too_many_arguments)]
9276    fn step35_verify(
9277        &self,
9278        e: &Engine,
9279        fa: &FullAttnLayer,
9280        h: &CudaSlice<f32>,
9281        h_q8: Option<(&CudaSlice<i8>, &CudaSlice<f32>)>,
9282        t: usize,
9283        cache: &mut Cache,
9284        il: usize,
9285    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9286        let n_embd = self.cfg.n_embd as usize;
9287        // The fused (h-less) attn-norm arm hands `h` as a zero-length placeholder. This arm needs
9288        // the f32 rows, so the caller must not take that lever for step35 — enforced at the call
9289        // site by the sliding-gated-MoE `Mixer::Full(_) => false` arm of
9290        // `lin_q8_only` in `decode_step_t_core_stream`, and asserted here so a future caller
9291        // cannot regress it into silently reading an empty buffer.
9292        assert_eq!(
9293            h.len(),
9294            t * n_embd,
9295            "step35_verify needs the f32 attn-normed rows ([t*n_embd]); the caller took the \
9296             fused q8-only norm arm (h_q8={}) — step35 must stay on the unfused arm",
9297            h_q8.is_some()
9298        );
9299        // ROW WIDTH IS n_embd, NOT n_head*head_dim: `step35_decode_attn` returns the mixer output
9300        // AFTER `wo`, so a row is [n_embd] — the same contract the generic arm's
9301        // `matmul_decode_exact(&fa.wo, &attn_g, t)` return has. Sizing this buffer from the
9302        // per-layer head geometry (8192 on full-attn, 12288 on SWA) instead overran the row on the
9303        // FIRST copy and panicked inside `copy_into`'s `CudaView::slice` unwrap
9304        // (raw/mtp-bt-20260806T212127Z.log frames 12-13).
9305        let mut out = vbuf(e, t * n_embd)?; // each row fully written by the copy below
9306        for r in 0..t {
9307            // Absolute position of this query row. `cache.pos` is the committed length at round
9308            // start and every row before r has already been appended by this loop, so the r-th
9309            // verify token sits at cache.pos + r — the same position eager decode would give it.
9310            let pos_d = e.htod_i32(&[(cache.pos + r) as i32])?;
9311            let mut h_row = vbuf(e, n_embd)?; // fully written by copy_view_into
9312            e.copy_view_into(
9313                &mut h_row,
9314                0,
9315                &h.slice(r * n_embd..(r + 1) * n_embd),
9316                n_embd,
9317            )?;
9318            // THE eager decode mixer: appends this row's K/V at kvl.len, advances it, then
9319            // attends over the (SWA-offset) view. Post-`wo`, same contract as this fn returns.
9320            let o = self.step35_decode_attn(e, fa, il, &h_row, None, &pos_d, cache)?;
9321            debug_assert_eq!(
9322                o.len(),
9323                n_embd,
9324                "step35_decode_attn returns post-wo [n_embd]"
9325            );
9326            e.copy_into(&mut out, r * n_embd, &o, n_embd)?;
9327        }
9328        Ok(out)
9329    }
9330
9331    /// Full-attention mixer over T query tokens with a GROWING resident KV (verify path, §D.3).
9332    /// Appends the T new K/V columns to cache.kv[il] then attends causally over [0..len) via
9333    /// fa_prefill. Token-major [T, kv_dim] projection layout == cache row layout (single copy).
9334    #[allow(clippy::too_many_arguments)]
9335    fn full_attn_verify(
9336        &self,
9337        e: &Engine,
9338        fa: &FullAttnLayer,
9339        h: &CudaSlice<f32>,
9340        h_q8: Option<(&CudaSlice<i8>, &CudaSlice<f32>)>,
9341        pos_d: &CudaSlice<i32>,
9342        t: usize,
9343        cache: &mut Cache,
9344        il: usize,
9345        stream_ctr: Option<&CudaSlice<i32>>,
9346    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9347        // step35: the generic geometry below is wrong for this arch (per-layer n_head, partial
9348        // per-layer rope, the SWA offset view, and a SEPARATE head-wise gate tensor), so it takes
9349        // its own arm. A verify that silently computes different attention than decode defeats the
9350        // whole self-consistency gate, so the arm is a per-row REPLAY of `step35_decode_attn`
9351        // rather than a batched twin — see `step35_verify` for why that is the exactness-correct
9352        // shape and not laziness.
9353        if self.sliding_gated_moe_batch_program() {
9354            if stream_ctr.is_some() {
9355                return Err(
9356                    "step35 has no ROUND-STREAM verify arm (the device-counter _dc twins \
9357                            cannot express the SWA offset KV view; same root cause as the dc \
9358                            decode refusal) — run spec without the stream arm"
9359                        .into(),
9360                );
9361            }
9362            return self.step35_verify(e, fa, h, h_q8, t, cache, il);
9363        }
9364        let cfg = &self.cfg;
9365        let geometry = cfg.full_attention_geometry_at(il as u32);
9366        let n_head = geometry.n_head as usize;
9367        let n_head_kv = geometry.n_head_kv as usize;
9368        let head_dim = geometry.head_dim_k as usize;
9369        let eps = cfg.rms_eps;
9370        let scale = geometry.attention_scale();
9371        let n_embd = cfg.n_embd as usize;
9372
9373        // DECODE-EXACT Q/K/V projections: matmul_decode_exact forces the MMVQ (warp-per-row) path
9374        // for every m, matching the T=1 decode's FP accumulation order. matmul_pre at m>=5 would
9375        // fall to dp4a (128-thread, two-level reduce) with a different FP sum order.
9376        // Q8 TRUNK-FUSION at T=1: DISPATCH-MIRRORS the eager decode's fused3 (bit-identical body).
9377        // BATCHED EPILOGUE RE-FUSE (lane/vt-fixes fix 2): `h_q8` = the attn-input norm's q8_1
9378        // form emitted by the fused rms_norm_q8_1 (bit-identical to rms_norm_decode ->
9379        // quantize_q8_1, kernel-check-pinned). When present (caller checked mixer q8_1-fast),
9380        // every projection consumes it — `h` may be a zero-len placeholder and must not be read.
9381        let (qf, mut k, v) = {
9382            let mut fused = None;
9383            let qkv_fast =
9384                e.uses_q8_1_fast(&fa.wq) && e.uses_q8_1_fast(&fa.wk) && e.uses_q8_1_fast(&fa.wv);
9385            if t == 1 && qkv_fast {
9386                let (hq_o, hd_o);
9387                let (hq, hd): (&CudaSlice<i8>, &CudaSlice<f32>) = match h_q8 {
9388                    Some(p) => p,
9389                    None => {
9390                        (hq_o, hd_o) = e.quantize_q8_1(h, 1, n_embd)?;
9391                        (&hq_o, &hd_o)
9392                    }
9393                };
9394                fused = e.matmul_q8_fused3(&fa.wq, &fa.wk, &fa.wv, hq, hd)?;
9395            } else if spec_fused_t() && (2..=4).contains(&t) && qkv_fast {
9396                // VERIFY-TIER TRUNK FUSION (MEMRA_SPEC_FUSED_T): one shared quantize + one
9397                // fused3 batched launch replaces three decode-exact calls (3 re-quantizes of
9398                // the same h + 3 _b2/_b4 launches). Bit-identical per (tensor,token,row).
9399                let (hq_o, hd_o);
9400                let (hq, hd): (&CudaSlice<i8>, &CudaSlice<f32>) = match h_q8 {
9401                    Some(p) => p,
9402                    None => {
9403                        (hq_o, hd_o) = e.quantize_q8_1(h, t, n_embd)?;
9404                        (&hq_o, &hd_o)
9405                    }
9406                };
9407                fused = e.matmul_q8_fused3_t(&fa.wq, &fa.wk, &fa.wv, hq, hd, t)?;
9408            }
9409            match (fused, h_q8) {
9410                (Some(triple), _) => triple,
9411                // shared pre-quantized activation (q8_1-fast guaranteed by the caller): the
9412                // decode-exact dispatch consumes (hq, hd) instead of re-quantizing 3x.
9413                (None, Some((hq, hd))) if qkv_fast => (
9414                    e.matmul_decode_exact_pre(&fa.wq, hq, hd, t)?,
9415                    e.matmul_decode_exact_pre(&fa.wk, hq, hd, t)?,
9416                    e.matmul_decode_exact_pre(&fa.wv, hq, hd, t)?,
9417                ),
9418                (None, _) => (
9419                    e.matmul_decode_exact(&fa.wq, h, t)?,
9420                    e.matmul_decode_exact(&fa.wk, h, t)?,
9421                    e.matmul_decode_exact(&fa.wv, h, t)?,
9422                ),
9423            }
9424        };
9425        // M3/Hy3 have no attention output gate — wq out is exactly q; skip the split.
9426        let gated = geometry.attention_gate == memra_gguf::config::AttentionGateKind::FusedQ;
9427        let (mut q, gate) = if gated {
9428            let mut q = vbuf(e, t * n_head * head_dim)?; // fully written by q_gate_split
9429            let mut gate = vbuf(e, t * n_head * head_dim)?; // fully written by q_gate_split
9430            e.q_gate_split(&qf, &mut q, &mut gate, head_dim, n_head, t)?;
9431            (q, Some(gate))
9432        } else {
9433            (qf, None)
9434        };
9435
9436        let mut qn = vbuf(e, t * n_head * head_dim)?; // fully written by rms_norm
9437        e.rms_norm(
9438            &q,
9439            fa.q_norm.float_data(),
9440            &mut qn,
9441            head_dim,
9442            n_head * t,
9443            eps,
9444        )?;
9445        q = qn;
9446        let mut kn = vbuf(e, t * n_head_kv * head_dim)?; // fully written by rms_norm
9447        e.rms_norm(
9448            &k,
9449            fa.k_norm.float_data(),
9450            &mut kn,
9451            head_dim,
9452            n_head_kv * t,
9453            eps,
9454        )?;
9455        k = kn;
9456        let rope_dims = geometry.n_rot as usize;
9457        e.rope_neox(
9458            &mut q,
9459            pos_d,
9460            head_dim,
9461            rope_dims,
9462            n_head,
9463            t,
9464            geometry.rope_base,
9465            1.0,
9466        )?;
9467        e.rope_neox(
9468            &mut k,
9469            pos_d,
9470            head_dim,
9471            rope_dims,
9472            n_head_kv,
9473            t,
9474            geometry.rope_base,
9475            1.0,
9476        )?;
9477
9478        // append T new K/V columns to the resident QUANTIZED cache. k/v are token-major [T, kv_dim]
9479        // f32; append-quantize each of the T token rows into the byte cache (q8_0 K / q5_1 V).
9480        let kvl = cache.kv[il].as_mut().unwrap();
9481        let (kv_dim_k, kv_dim_v, ktb, vtb) =
9482            (kvl.kv_dim_k, kvl.kv_dim_v, kvl.k_tok_bytes, kvl.v_tok_bytes);
9483        if let Some(ctr) = stream_ctr {
9484            // stream: ONE batched append at the device counter (rows kernel = the per-view warp
9485            // math on a (block, token) grid, documented byte-identical); host len is a stale
9486            // LOWER BOUND under pre-issue (drain reconciles it).
9487            e.append_kv_quantized_rows_dc(
9488                &k,
9489                &v,
9490                &mut kvl.k,
9491                &mut kvl.v,
9492                ctr,
9493                t,
9494                kv_dim_k,
9495                kv_dim_v,
9496                ktb,
9497                vtb,
9498                crate::Engine::kv_fp8_on(),
9499            )?;
9500        } else {
9501            for i in 0..t {
9502                let k_row = k.slice(i * kv_dim_k..(i + 1) * kv_dim_k);
9503                let v_row = v.slice(i * kv_dim_v..(i + 1) * kv_dim_v);
9504                e.append_kv_quantized_view(
9505                    &k_row,
9506                    &v_row,
9507                    &mut kvl.k,
9508                    &mut kvl.v,
9509                    kvl.len + i,
9510                    kv_dim_k,
9511                    kv_dim_v,
9512                    ktb,
9513                    vtb,
9514                    crate::Engine::kv_fp8_on(),
9515                )?;
9516            }
9517            kvl.len += t;
9518        }
9519
9520        // BIT-IDENTICAL VERIFY ATTENTION (spec-exactness fix): the FP accumulation order must be
9521        // byte-for-byte identical to the eager decode path. fa_prefill uses a different tile size
9522        // (BLOCK_Q=64, BK=32) and online-softmax structure than fa_decode's split-K + combine,
9523        // which changes FP summation order and can flip argmax at tight logit margins. Query row r
9524        // attends to keys [0..base_len+r+1) — each successive row sees one more key (the causal
9525        // property). This matches eager: decode appends k at len, then fa_decode sees t_kv = len+1
9526        // keys. The verify appends all T tokens first but bounds the key range per row.
9527        //
9528        // MULTI-ROW FUSED PATH (the long-ctx spec fix, 2026-07-03): when every row takes the vec
9529        // kernel (base_len+1 >= FA_VEC_MIN_TKV), ONE fa_decode_rows launch executes the exact
9530        // per-row program for all T rows (grid.z = row, per-row n_splits from the same
9531        // fa_split_keys formula) — replacing T x (2 launches + 2 dtod copies + 5 partial allocs)
9532        // and multiplying resident CTAs by T on a latency-bound kernel. Bit-identical per row by
9533        // construction; kernel-check pins rows-vs-loop byte identity, run-spec is the end gate.
9534        // Short ctx (any row below the vec crossover) and MEMRA_NO_FA_VEC/MEMRA_FA_ROWS_OFF keep the
9535        // per-row loop (whose fa_decode picks scalar/vec per row exactly like eager decode).
9536        let mut attn = vbuf(e, t * n_head * head_dim)?; // fully written by every FA arm below
9537        let base_len = kvl.len - t; // KV len BEFORE this round's T tokens were appended
9538        // T=1 INCLUDED (2026-07-05): p-min cuts the draft to 1 in ~75% of rounds on hard
9539        // (agentic) content — the old t>1 gate sent those rounds to the per-row loop (262us/row
9540        // + q-row copy + per-row allocs vs 93us/row through the fused kernel at grid.z=1, same
9541        // program). nsys accounting: 1088 of 1456 verify FA launches were T=1 escapees.
9542        // LEAN T=1 ARM (MEMRA_SPEC_LEAN, close35): at t==1, q IS one row and fa_decode on it is
9543        // the EXACT eager decode dispatch (vec_q_v2 + combine_f32; the rows pair measured +50us
9544        // at m=1). Byte-identical: kernel-check pins rows-vs-loop identity, and the per-row loop
9545        // at t=1 is fa_decode on the same q with zero-offset copies. Gates arbitrate.
9546        if let Some(ctr) = stream_ctr {
9547            // STREAM ARM: causal base from the device counter; host kvl.len is a stale lower
9548            // bound used only for the split-sizing upper bound (+64 slack covers M pre-issued
9549            // rounds at K<=8). Views span the bound; per-row limits derive in-kernel.
9550            let upper = kvl.len + t + 64;
9551            let k_view = e.view_u8(&kvl.k, (upper.min(cache.max_ctx)) * ktb);
9552            let v_view = e.view_u8(&kvl.v, (upper.min(cache.max_ctx)) * vtb);
9553            e.fa_decode_rows_dc(
9554                &q,
9555                &k_view,
9556                &v_view,
9557                &mut attn,
9558                head_dim,
9559                n_head,
9560                n_head_kv,
9561                ctr,
9562                upper.min(cache.max_ctx),
9563                t,
9564                scale,
9565                ktb,
9566                vtb,
9567                0,
9568                false,
9569            )?;
9570        } else if spec_lean() && t == 1 {
9571            let t_kv = base_len + 1;
9572            let k_view = e.view_u8(&kvl.k, t_kv * ktb);
9573            let v_view = e.view_u8(&kvl.v, t_kv * vtb);
9574            e.fa_decode_kvmod(
9575                &q,
9576                &k_view,
9577                &v_view,
9578                &mut attn,
9579                head_dim,
9580                n_head,
9581                n_head_kv,
9582                t_kv,
9583                scale,
9584                ktb,
9585                vtb,
9586                crate::Engine::kv_fp8_on(),
9587            )?;
9588        } else if e.fa_rows_eligible(base_len, head_dim) {
9589            let k_view = e.view_u8(&kvl.k, (base_len + t) * ktb);
9590            let v_view = e.view_u8(&kvl.v, (base_len + t) * vtb);
9591            e.fa_decode_rows(
9592                &q,
9593                &k_view,
9594                &v_view,
9595                &mut attn,
9596                head_dim,
9597                n_head,
9598                n_head_kv,
9599                base_len,
9600                t,
9601                scale,
9602                ktb,
9603                vtb,
9604                None,
9605                false,
9606                crate::Engine::kv_fp8_on(),
9607                None,
9608            )?;
9609        } else {
9610            for r in 0..t {
9611                let t_kv_r = base_len + r + 1; // this row sees keys [0..t_kv_r)
9612                let k_view_r = e.view_u8(&kvl.k, t_kv_r * ktb);
9613                let v_view_r = e.view_u8(&kvl.v, t_kv_r * vtb);
9614                // copy q row into an owned buffer (fa_decode takes &CudaSlice, not CudaView)
9615                let mut q_row = vbuf(e, n_head * head_dim)?; // fully written by copy_view_into
9616                let q_src = q.slice(r * n_head * head_dim..(r + 1) * n_head * head_dim);
9617                e.copy_view_into(&mut q_row, 0, &q_src, n_head * head_dim)?;
9618                let mut attn_row = vbuf(e, n_head * head_dim)?; // fully written by fa_decode
9619                e.fa_decode_kvmod(
9620                    &q_row,
9621                    &k_view_r,
9622                    &v_view_r,
9623                    &mut attn_row,
9624                    head_dim,
9625                    n_head,
9626                    n_head_kv,
9627                    t_kv_r,
9628                    scale,
9629                    ktb,
9630                    vtb,
9631                    crate::Engine::kv_fp8_on(),
9632                )?;
9633                e.copy_into(
9634                    &mut attn,
9635                    r * n_head * head_dim,
9636                    &attn_row,
9637                    n_head * head_dim,
9638                )?;
9639            }
9640        }
9641
9642        let attn_g = match &gate {
9643            Some(gate) => {
9644                let mut gsig = vbuf(e, t * n_head * head_dim)?; // fully written by sigmoid
9645                e.sigmoid(gate, &mut gsig, t * n_head * head_dim)?;
9646                let mut ag = vbuf(e, t * n_head * head_dim)?; // fully written by mul
9647                e.mul(&attn, &gsig, &mut ag, t * n_head * head_dim)?;
9648                ag
9649            }
9650            None => attn,
9651        };
9652        // DECODE-EXACT wo projection: at m>=5 (K=4+ with pending) the generic matmul would use dp4a
9653        // (128-thread, different FP sum order than MMVQ). Force MMVQ for bit-identity with decode.
9654        Ok(e.matmul_decode_exact(&fa.wo, &attn_g, t)?)
9655    }
9656
9657    /// Context-linear bytes for a plain serving session's trunk cache.
9658    pub fn plain_session_kv_bytes_per_token(&self) -> usize {
9659        crate::cache::cache_bytes_per_token_for_plan(
9660            &self.cfg,
9661            &self.plan,
9662            0,
9663            self.plan.layers.len(),
9664        )
9665    }
9666
9667    /// `(logical bytes/token, ring-capped bytes/token, ring row cap)` for exact admission.
9668    pub fn plain_session_kv_shape(&self) -> (usize, usize, usize) {
9669        (
9670            self.plain_session_kv_bytes_per_token(),
9671            crate::cache::cache_ring_bytes_per_token_for_plan(
9672                &self.cfg,
9673                &self.plan,
9674                0,
9675                self.plan.layers.len(),
9676            ),
9677            crate::cache::cache_ring_row_cap_for_plan(&self.plan),
9678        )
9679    }
9680
9681    /// Context-linear bytes for a speculative serving session: trunk cache plus persistent MTP
9682    /// scratch. With no MTP head this equals the plain coefficient.
9683    pub fn spec_session_kv_bytes_per_token(&self) -> usize {
9684        let scratch = self
9685            .mtp
9686            .iter()
9687            .chain(self.mtp_extra.iter())
9688            .map(|mtp| {
9689                let (_, _, k, v) = mtp_scratch_layout(&self.cfg, mtp.geom.as_ref());
9690                k + v
9691            })
9692            .sum::<usize>();
9693        self.plain_session_kv_bytes_per_token()
9694            .saturating_add(scratch)
9695    }
9696
9697    /// Spec twin of [`HybridModel::plain_session_kv_shape`]; Step35's persistent MTP scratch is
9698    /// capped by the same SWA ring rows as the trunk.
9699    pub fn spec_session_kv_shape(&self) -> (usize, usize, usize) {
9700        let total = self.spec_session_kv_bytes_per_token();
9701        let (_, mut ring, rows) = self.plain_session_kv_shape();
9702        if rows > 0 {
9703            ring = ring.saturating_add(
9704                self.mtp
9705                    .iter()
9706                    .chain(self.mtp_extra.iter())
9707                    .map(|mtp| {
9708                        let (_, _, k, v) = mtp_scratch_layout(&self.cfg, mtp.geom.as_ref());
9709                        k + v
9710                    })
9711                    .sum::<usize>(),
9712            );
9713        }
9714        (total, ring, rows)
9715    }
9716
9717    /// Greedy MTP speculative decode (§B). Token-identical to `generate(prompt, max_new)` but uses
9718    /// the NextN head to draft K tokens then verifies them in one batched target forward.
9719    /// Returns (generated tokens, total_drafted, total_accepted) so the caller can report
9720    /// acceptance rate. `k` = draft length per round.
9721    ///
9722    /// GRAPH DRAFT (stage 2 of graph-grade spec): when the model is all-Dense and the MTP head is
9723    /// Dense (no MoE host readbacks), the fixed-shape T=1 MTP forward is CUDA-graph-captured ONCE
9724    /// and replayed per draft step — the ~40 eager launches per drafted token collapse into one
9725    /// graph dispatch; only the 4-byte token id (and 4-byte p-min confidence) round-trip per step.
9726    /// Event tracking is disabled for the whole call (generate_graph pattern) so every buffer the
9727    /// captured graph references is event-free; the spec loop is strictly single-stream.
9728    /// MEMRA_SPEC_NOGRAPH=1 forces the eager draft chain.
9729    /// SAMPLED mode (MEMRA_SPEC_TEMP>0) has its OWN capture (gumbel-perturbed in-graph argmax,
9730    /// device Philox event counter, persistent q retention) — graph-vs-eager sampled streams are
9731    /// bit-identical for the same (seed, prompt, K, temp); see the sampled-graph setup in
9732    /// generate_spec_inner2.
9733    /// Multi-turn session: trunk cache + MTP draft scratch persist across generate calls, so
9734    /// turn N+1 primes ONLY its new suffix (the 124k-conversation daily pattern — re-priming a
9735    /// 32k history costs ~54s; a suffix prime costs seconds). APPEND-ONLY by construction: the
9736    /// hybrid linear-attn states are in-place (no position index), so a session can extend but
9737    /// never rewind — `committed` is the exact token list whose state the caches hold (includes
9738    /// any overshoot tokens past max_new; the caller renders from `committed`, not its own echo).
9739    pub fn new_session(
9740        &self,
9741        e: &Engine,
9742        max_ctx: usize,
9743    ) -> Result<SpecSession, Box<dyn std::error::Error>> {
9744        Ok(SpecSession {
9745            // STAGE-OWNED KV (lane/pp2-spec 2026-08-06): `pp::new_cache`, not `Cache::new`. This
9746            // is the SERVING spec-session path, and with the ppN door open across two cards a
9747            // primary-homed cache makes every remote stage peer-read its OWN KV on every verify
9748            // round — the wrong-card class already fixed on the two batched serving paths
9749            // (worker.rs 2483 / 2837). With the door shut `new_cache` IS `Cache::new` (same
9750            // branch, same allocations), so single-device behavior is byte-unchanged.
9751            cache: crate::pp::new_cache_planned(e, &self.cfg, &self.plan, max_ctx)?,
9752            scratch: self.new_mtp_scratch(e, max_ctx)?,
9753            committed: Vec::new(),
9754            last_h: None,
9755            next_pred: None,
9756            sctr: 0,
9757            uctr: 0,
9758            draft_ctx: None,
9759            pending_tok: None,
9760            turn_ckpt: None,
9761            telem: SpecTelemetryCounters::default(),
9762            capture_at: None,
9763            boundary_captures: Vec::new(),
9764            ckpt_at: None,
9765            capture_disabled: false,
9766        })
9767    }
9768
9769    /// SPEC-ON-CACHE-HIT restore (lane/spec-on-cache-hit, 2026-08-18 — PORT-PLAN item 3,
9770    /// research/cache-spec-design-20260814, scoped to WHOLE-ENTRY restores only): build a
9771    /// SpecSession around a trunk cache the worker already restored from a prefix-cache
9772    /// entry, re-installing the entry's published draft plane as the MTP scratch rows
9773    /// `[0..prefix.len())` and the entry's boundary hidden as `last_h`, then feeding the
9774    /// prompt SUFFIX here — through EXACTLY the plain path's program selection — so the
9775    /// worker always receives a fully-warm continuation session (committed = whole
9776    /// prompt, `next_pred` + `last_h` set; caller sets `next_pred` from the entry's
9777    /// boundary logits on the empty-suffix shape).
9778    ///
9779    /// PROGRAM LAW (the splitiso two-programs class, learned AGAIN in this lane's own
9780    /// gate): the identity target for a converted hit is the PLAIN hit serving the same
9781    /// request, and plain feeds a carried suffix via eager `decode_step` below
9782    /// PRIME_MIN_T and via `prime_cache` at/above it (prefill_tick's arms). The generate
9783    /// path's tokenwise arm routes qwen35-class through the BATCHED T=1 program
9784    /// (`spec_target_step_h`) instead — ULP-different suffix rows, and the gate measured
9785    /// the near-tie flip at generated token ~8 (research/spec-cache-20260818, qwen r3).
9786    /// So the suffix is fed HERE, mirroring prefill_tick arm-for-arm, not handed to the
9787    /// burst prime.
9788    ///
9789    /// SEED RULE (both sampling regimes; lane/sampled-hit-spec 2026-08-19, sampled draw
9790    /// added by lane/sampled-spec-quality 2026-08-19). The boundary token is produced by
9791    /// EXACTLY the rule the cold burst entry applies to its own first token from the same
9792    /// logits row: `argmax` when greedy, and a `sample_boundary_token` draw at Philox
9793    /// counter 0 when sampled. Both shapes are covered — the entry's boundary logits on a
9794    /// full-cover (empty-suffix) hit, this feed's own boundary logits on a suffix hit.
9795    /// That is what keeps a restored session seed-identical to a cold one PER SEED: the
9796    /// cold session draws from the identical row at counter 0 and then runs its rounds from
9797    /// counter 1, so the restored session admits with `sctr = 1` after its own draw.
9798    /// The WORKER owns the one refusal this constructor cannot see — a constrained request.
9799    /// (The penalized-sampled refusal was LIFTED once the burst's penalty window learned to
9800    /// span the session: `committed` here is the WHOLE prompt, so the restored session's
9801    /// window is the cold session's window. It comes back if `MEMRA_SPEC_PEN_SESSION=0`.)
9802    ///
9803    /// NOT the rolled-back partial-restore hazard: the caller restores at exactly the
9804    /// entry's captured endpoint (`e.pos`) through the shipping whole-entry path;
9805    /// mid-entry (`at < e.pos`) trunk restores stay behind MEMRA_PREFIX_PARTIAL_RESTORE
9806    /// and are never routed here.
9807    ///
9808    /// Failure contract: `Err((Some(cache), why))` before any trunk mutation — the
9809    /// worker rebuilds the plain carrier and the hit serves plain, byte-unchanged.
9810    /// `Err((None, why))` after the suffix feed began — the carrier is part-fed and
9811    /// UNUSABLE; the worker serves the request cold-plain (correct, slower) and the
9812    /// entry stays published for the next request.
9813    #[allow(clippy::too_many_arguments)]
9814    pub fn spec_session_from_restored(
9815        &self,
9816        e: &Engine,
9817        mut cache: Cache,
9818        prefix: Vec<u32>,
9819        suffix: &[u32],
9820        draft_k: &CudaSlice<u8>,
9821        draft_v: &CudaSlice<u8>,
9822        draft_k_tok_bytes: usize,
9823        draft_v_tok_bytes: usize,
9824        draft_len: usize,
9825        last_h: &[f32],
9826        // The ENTRY's boundary logits row (the full-cover shape's seed source). May be empty
9827        // when a suffix follows — the feed's own logits are the boundary then.
9828        boundary_logits: &[f32],
9829        // The request's sampler, or None for greedy. Owned here so the seed rule lives in
9830        // ONE place instead of being half-applied by the worker.
9831        sampling: Option<SpecSampling>,
9832        require_anchor: bool,
9833        max_ctx: usize,
9834        // STABLE-BOUNDARY REPUBLICATION (lane/frspec-multiturn-cache, 2026-08-21): ABSOLUTE
9835        // prompt position to split the suffix feed at and capture the extended-entry
9836        // publication + this session's `turn_ckpt` — the worker's stable pre-generation
9837        // boundary (`plain_checkpoint_boundary`). None = legacy prompt-end republication.
9838        // WHY: the prompt-end capture below includes the template's live generation header
9839        // (`<|im_start|>assistant\n<think>\n`), which the next turn's re-render replaces, so
9840        // for a hybrid (whole-entry restores only) every extended entry's last ~2 tokens
9841        // diverged from every future prompt and the hit boundary FROZE at the first
9842        // lcp-split entry forever (measured: cached 6811 of 38228 by turn 8, B4).
9843        republish_at: Option<usize>,
9844    ) -> Result<SpecSession, (Option<Cache>, String)> {
9845        let pos = prefix.len();
9846        let fail = |cache: Cache, msg: String| -> Result<SpecSession, (Option<Cache>, String)> {
9847            Err((Some(cache), msg))
9848        };
9849        if self.mtp.is_none() {
9850            return fail(cache, "no MTP head attached (nothing to draft with)".into());
9851        }
9852        if pos == 0 {
9853            return fail(cache, "empty committed prefix".into());
9854        }
9855        if cache.pos != pos {
9856            let msg = format!(
9857                "restored cache pos {} != restored prefix len {pos}",
9858                cache.pos
9859            );
9860            return fail(cache, msg);
9861        }
9862        if draft_len != pos {
9863            return fail(
9864                cache,
9865                format!("draft plane len {draft_len} != restored prefix len {pos}"),
9866            );
9867        }
9868        if pos + suffix.len() >= max_ctx {
9869            return fail(
9870                cache,
9871                format!(
9872                    "prompt {} + suffix would not leave generation room in ctx {max_ctx}",
9873                    pos + suffix.len(),
9874                ),
9875            );
9876        }
9877        let mut scratch = match MtpScratch::new(
9878            e,
9879            &self.cfg,
9880            &self.plan,
9881            max_ctx,
9882            self.mtp.as_ref().and_then(|m| m.geom.as_ref()),
9883        ) {
9884            Ok(s) => s,
9885            Err(err) => return fail(cache, format!("draft scratch alloc failed: {err}")),
9886        };
9887        if scratch.kv.ring.is_some() {
9888            return fail(
9889                cache,
9890                "ring-backed draft scratch (Step35 SWA) cannot take a flat prefix restore".into(),
9891            );
9892        }
9893        if scratch.kv.k_tok_bytes != draft_k_tok_bytes
9894            || scratch.kv.v_tok_bytes != draft_v_tok_bytes
9895        {
9896            return fail(
9897                cache,
9898                format!(
9899                    "draft plane layout {draft_k_tok_bytes}/{draft_v_tok_bytes} != scratch \
9900                     {}/{} bytes/token (stale entry across a format change)",
9901                    scratch.kv.k_tok_bytes, scratch.kv.v_tok_bytes,
9902                ),
9903            );
9904        }
9905        if pos > scratch.cap {
9906            return fail(
9907                cache,
9908                format!(
9909                    "draft plane rows {pos} exceed scratch capacity {}",
9910                    scratch.cap
9911                ),
9912            );
9913        }
9914        let kb = pos * draft_k_tok_bytes;
9915        let vb = pos * draft_v_tok_bytes;
9916        if draft_k.len() < kb || draft_v.len() < vb {
9917            return fail(
9918                cache,
9919                format!(
9920                    "truncated draft plane: K {} < {kb} or V {} < {vb} bytes",
9921                    draft_k.len(),
9922                    draft_v.len(),
9923                ),
9924            );
9925        }
9926        if kb > 0 {
9927            if let Err(err) = e.copy_u8_into(&mut scratch.kv.k, 0, draft_k, kb) {
9928                return fail(cache, format!("draft K restore copy failed: {err}"));
9929            }
9930        }
9931        if vb > 0 {
9932            if let Err(err) = e.copy_u8_into(&mut scratch.kv.v, 0, draft_v, vb) {
9933                return fail(cache, format!("draft V restore copy failed: {err}"));
9934            }
9935        }
9936        if let Err(err) = scratch.set_len(e, pos) {
9937            return fail(cache, format!("draft scratch len set failed: {err}"));
9938        }
9939        let mut last_h_dev = if last_h.len() == self.cfg.n_embd as usize {
9940            // anchor upload failure is acceptance-only when a suffix feed follows (fill
9941            // row-0 falls back to zeros) but FATAL for an empty-suffix continuation (the
9942            // burst entry asserts committed + last_h + next_pred) — the caller says which.
9943            e.htod(last_h).ok()
9944        } else {
9945            None
9946        };
9947        if require_anchor && last_h_dev.is_none() {
9948            return fail(
9949                cache,
9950                "empty-suffix continuation requires the entry's boundary hidden anchor".into(),
9951            );
9952        }
9953        let mut committed = prefix;
9954        // Set on BOTH shapes below (suffix-fed and full-cover) — never left None, which is
9955        // what the empty-suffix continuation assert in the burst entry requires.
9956        let next_pred;
9957        // Philox: (0,0) at admit exactly like a fresh session; a sampled boundary draw below
9958        // consumes counter 0 and leaves 1, which is the state a cold session reaches after
9959        // drawing its own first token from the same row.
9960        let mut sctr = 0u32;
9961        let sampled = sampling.is_some_and(|s| s.temp > 0.0) && spec_sampled_boundary_on();
9962        // Penalty window for the boundary draw: the last `penalty_last_n` tokens of the WHOLE
9963        // prompt, which is what the cold session's own burst sees (Item 2's window). Built
9964        // after the suffix joins `committed` below.
9965        let mut boundary_captures: Vec<SpecBoundaryCapture> = Vec::new();
9966        let mut restored_turn_ckpt: Option<SpecCheckpoint> = None;
9967        if !suffix.is_empty() {
9968            // ---- SUFFIX FEED, mirroring prefill_tick's program selection exactly ----
9969            // From here on the trunk cache mutates: failures return Err((None, _)) and
9970            // the worker serves the request cold-plain instead of reusing the carrier.
9971            let dirty =
9972                |msg: String| -> Result<SpecSession, (Option<Cache>, String)> { Err((None, msg)) };
9973            let n_embd = self.cfg.n_embd as usize;
9974            let t = suffix.len();
9975            let mut h_rows = match e.uninit(t * n_embd) {
9976                Ok(b) => b,
9977                Err(err) => return fail(cache, format!("suffix hidden buffer alloc: {err}")),
9978            };
9979            // STABLE-BOUNDARY split (see `republish_at`): feed stops at the boundary so the
9980            // in-place GDN conv/ssm state can be snapshotted there — the only moment it
9981            // exists (the cold prime-split law). suffix-relative; None = one-segment legacy.
9982            let b_rel = republish_at
9983                .and_then(|abs| abs.checked_sub(pos))
9984                .filter(|&r| r > 0 && r < t);
9985            let mut feed_logits = Vec::new();
9986            let tokenwise_env = std::env::var("MEMRA_PRIME_TOKENWISE").is_ok()
9987                || e.frozen_cpu_experts_prefer_tokenwise_prime();
9988            let mut fed = 0usize;
9989            for seg_end in [b_rel, Some(t)].into_iter().flatten() {
9990                if seg_end <= fed {
9991                    continue;
9992                }
9993                let seg = &suffix[fed..seg_end];
9994                let batched = seg.len() >= crate::hybrid_forward::PRIME_MIN_T && !tokenwise_env;
9995                if batched {
9996                    // prefill_tick's prime arm: request-level prime_cache call; tokens still
9997                    // queued after this segment ride `queued_after` so Step35 arm selection
9998                    // stays keyed to the request's end (tick-seg law).
9999                    match self.prime_cache(e, seg, &mut cache, t - seg_end) {
10000                        Ok((l, _h_seed, hiddens)) => {
10001                            if let Err(err) =
10002                                e.copy_into(&mut h_rows, fed * n_embd, &hiddens, seg.len() * n_embd)
10003                            {
10004                                return dirty(format!("suffix hidden copy: {err}"));
10005                            }
10006                            feed_logits = l;
10007                        }
10008                        Err(err) => return dirty(format!("suffix prime failed: {err}")),
10009                    }
10010                } else {
10011                    // prefill_tick's tokenwise arm: eager decode_step, one token at a time.
10012                    for (i, &tok) in seg.iter().enumerate() {
10013                        match self.decode_step_h(e, tok, &mut cache) {
10014                            Ok((l, h)) => {
10015                                if let Err(err) =
10016                                    e.copy_into(&mut h_rows, (fed + i) * n_embd, &h, n_embd)
10017                                {
10018                                    return dirty(format!("suffix hidden copy: {err}"));
10019                                }
10020                                feed_logits = l;
10021                            }
10022                            Err(err) => return dirty(format!("suffix decode_step failed: {err}")),
10023                        }
10024                    }
10025                }
10026                fed = seg_end;
10027                if Some(seg_end) == b_rel {
10028                    // The stable pre-generation boundary: capture the extended-entry
10029                    // publication AND this session's own turn checkpoint here instead of at
10030                    // prompt-end (both would otherwise carry the volatile live-header tail
10031                    // the next re-render replaces). Failure silent, turn_ckpt convention.
10032                    debug_assert_eq!(
10033                        cache.pos,
10034                        pos + seg_end,
10035                        "stable-boundary capture off the feed split"
10036                    );
10037                    if spec_restore_republish_on() {
10038                        if let Ok(snap) = cache.snapshot(e) {
10039                            boundary_captures.push(SpecBoundaryCapture {
10040                                snap,
10041                                pos: pos + seg_end,
10042                                logits: feed_logits.clone(),
10043                                last_h: capture_boundary_hidden(e, &h_rows, seg_end, n_embd),
10044                            });
10045                        }
10046                    }
10047                    let anchor: Result<CudaSlice<f32>, Box<dyn std::error::Error>> =
10048                        e.uninit(n_embd).and_then(|mut a| {
10049                            e.copy_view_into(
10050                                &mut a,
10051                                0,
10052                                &h_rows.slice((seg_end - 1) * n_embd..seg_end * n_embd),
10053                                n_embd,
10054                            )?;
10055                            Ok(a)
10056                        });
10057                    if let (Ok(snap), Ok(last_h)) = (cache.snapshot(e), anchor) {
10058                        restored_turn_ckpt = Some(SpecCheckpoint {
10059                            snap,
10060                            pos: pos + seg_end,
10061                            last_h,
10062                        });
10063                    }
10064                }
10065            }
10066            // Draft-scratch fill for the suffix rows, predecessor-paired: row `pos` reads
10067            // the entry's boundary anchor (zeros fallback — acceptance-only), row `pos+i`
10068            // reads h_rows[i-1]. Chunked like the generate path's fill (transients scale
10069            // with T). Fill failures are acceptance-only — truncate to the restored rows
10070            // and continue; the burst's own set_len keeps the invariant.
10071            let mtp = self.mtp.as_ref().expect("mtp checked above");
10072            let (embd_qt, embd_rb) = self.embd.qt_and_row_bytes(n_embd);
10073            let embd_gpu = if spec_host_embd() {
10074                None
10075            } else {
10076                Some(
10077                    self.embd_gpu
10078                        .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload")),
10079                )
10080            };
10081            let embd_dev = embd_gpu.map(|g| (g, embd_qt, embd_rb));
10082            let fill_chunk = 4096usize;
10083            let mut filled = true;
10084            let mut start = 0usize;
10085            'fill: while start < t {
10086                let end = (start + fill_chunk).min(t);
10087                let tc = end - start;
10088                let Ok(mut phs) = e.zeros(tc * n_embd) else {
10089                    filled = false;
10090                    break 'fill;
10091                };
10092                let (src_lo, dst_off, n_copy) = if start == 0 {
10093                    (0, n_embd, (tc - 1) * n_embd)
10094                } else {
10095                    ((start - 1) * n_embd, 0, tc * n_embd)
10096                };
10097                if start == 0 {
10098                    if let Some(lh) = last_h_dev.as_ref() {
10099                        if e.copy_into(&mut phs, 0, lh, n_embd).is_err() {
10100                            filled = false;
10101                            break 'fill;
10102                        }
10103                    }
10104                }
10105                if n_copy > 0
10106                    && e.copy_view_into(
10107                        &mut phs,
10108                        dst_off,
10109                        &h_rows.slice(src_lo..src_lo + n_copy),
10110                        n_copy,
10111                    )
10112                    .is_err()
10113                {
10114                    filled = false;
10115                    break 'fill;
10116                }
10117                if self
10118                    .mtp_kv_fill_all(
10119                        e,
10120                        &suffix[start..end],
10121                        &phs,
10122                        pos + start,
10123                        &mut scratch,
10124                        embd_dev,
10125                    )
10126                    .is_err()
10127                {
10128                    filled = false;
10129                    break 'fill;
10130                }
10131                start = end;
10132            }
10133            if !filled {
10134                // acceptance-only: drafts over missing suffix rows are cheap and wrong,
10135                // so keep only the restored rows resident and let verify arbitrate.
10136                if let Err(err) = scratch.set_len(e, pos) {
10137                    return dirty(format!("scratch truncation after failed fill: {err}"));
10138                }
10139            }
10140            // EXTENDED-ENTRY PUBLICATION (lane/sampled-spec-quality, Item 3 — the fix for
10141            // "a restored spec session never publishes an extended entry", SAMPLED-HIT.md
10142            // finding (d)). Pre-lane, publication was armed only for COLD sessions
10143            // (`spec_resumed == 0` in the worker) and both engine capture sites require a
10144            // non-continuation burst — but a converted hit's first burst IS a continuation,
10145            // so a growing conversation learned exactly ONE boundary and turn 3 could never
10146            // hit a longer prefix than turn 2 did.
10147            //
10148            // WHERE, and why it is safe here: `cache.pos == prefix + suffix` at this exact
10149            // line — the trunk is primed over the whole prompt, nothing is generated, and the
10150            // draft plane rows [0..prompt) are filled just above. That is a complete
10151            // whole-entry boundary (`pos == fed_len`), the same shape the cold seed capture
10152            // publishes; the worker's existing publication sweep picks it up because it is
10153            // keyed on non-empty `boundary_captures` and is sampler- and resume-independent.
10154            // NOT the partial-restore hazard: the boundary is this session's own prompt END,
10155            // never mid-entry, so `entry_pos != fed_len` still refuses on the way back in.
10156            // Failure is SILENT by design (the turn_ckpt / boundary-capture convention):
10157            // publication is an optimization, never a correctness dependency.
10158            //
10159            // SUPERSEDED WHEN `republish_at` FIRED (lane/frspec-multiturn-cache): a prompt-end
10160            // entry's tail is the live generation header the next re-render replaces, so on a
10161            // hybrid (whole-entry restores) it can never serve the conversation's next turn —
10162            // the stable-boundary capture above IS this publication, minus the poisoned tail.
10163            if spec_restore_republish_on() && boundary_captures.is_empty() {
10164                debug_assert_eq!(
10165                    cache.pos,
10166                    pos + t,
10167                    "extended-entry capture must sit at the restored session's prompt end",
10168                );
10169                if let Ok(snap) = cache.snapshot(e) {
10170                    boundary_captures.push(SpecBoundaryCapture {
10171                        snap,
10172                        pos: pos + t,
10173                        logits: feed_logits.clone(),
10174                        last_h: capture_boundary_hidden(e, &h_rows, t, n_embd),
10175                    });
10176                }
10177            }
10178            // continuation seed: the feed's boundary logits ARE the plain path's boundary
10179            // logits (same program), so greedy's argmax here is plain's first emitted token,
10180            // and the sampled draw is the cold sampled session's own first token.
10181            next_pred = Some(if sampled {
10182                let sp = sampling.expect("sampled implies a sampler");
10183                // `committed` is still the restored prefix here; the suffix joins it below —
10184                // so this is the last-N window over the WHOLE prompt, exactly the cold
10185                // session's own window at its first token.
10186                let hist = pen_window_seed(&committed, suffix, sp.penalty_last_n);
10187                match sample_boundary_token(
10188                    e,
10189                    &feed_logits,
10190                    &sp,
10191                    &hist,
10192                    &mut sctr,
10193                    "restore-suffix-feed",
10194                ) {
10195                    Ok(t) => t,
10196                    // the trunk is already fed: hand nothing back, the worker serves the
10197                    // request cold-plain. Never fall back to an argmax — that would put a
10198                    // greedy token in a sampled stream to save a slow path.
10199                    Err(err) => {
10200                        return dirty(format!("boundary token draw failed: {err}"));
10201                    }
10202                }
10203            } else {
10204                argmax(&feed_logits) as u32
10205            });
10206            let mut lh = match e.uninit(n_embd) {
10207                Ok(b) => b,
10208                Err(err) => return dirty(format!("boundary hidden alloc: {err}")),
10209            };
10210            if let Err(err) = e.copy_view_into(
10211                &mut lh,
10212                0,
10213                &h_rows.slice((t - 1) * n_embd..t * n_embd),
10214                n_embd,
10215            ) {
10216                return dirty(format!("boundary hidden copy: {err}"));
10217            }
10218            last_h_dev = Some(lh);
10219            committed.extend_from_slice(suffix);
10220        } else {
10221            // FULL-COVER shape (empty suffix — the identical-repeat / agent-loop shape): the
10222            // ENTRY's boundary logits are the boundary row, and this is the token the cold
10223            // session emits from that same row. Owned here rather than in the worker so the
10224            // sampled draw cannot be half-applied on one shape (the worker used to argmax it).
10225            if boundary_logits.is_empty() {
10226                return fail(
10227                    cache,
10228                    "full-cover restore without the entry's boundary logits".into(),
10229                );
10230            }
10231            next_pred = Some(if sampled {
10232                let sp = sampling.expect("sampled implies a sampler");
10233                let hist = pen_window_seed(&committed, &[], sp.penalty_last_n);
10234                match sample_boundary_token(
10235                    e,
10236                    boundary_logits,
10237                    &sp,
10238                    &hist,
10239                    &mut sctr,
10240                    "restore-full-cover",
10241                ) {
10242                    Ok(t) => t,
10243                    // nothing has been mutated on this shape — hand the carrier back and let
10244                    // the hit serve PLAIN (the banked pre-lane path).
10245                    Err(err) => {
10246                        return fail(cache, format!("boundary token draw failed: {err}"));
10247                    }
10248                }
10249            } else {
10250                argmax(boundary_logits) as u32
10251            });
10252        }
10253        Ok(SpecSession {
10254            cache,
10255            scratch,
10256            committed,
10257            last_h: last_h_dev,
10258            next_pred,
10259            sctr,
10260            uctr: 0,
10261            draft_ctx: None,
10262            pending_tok: None,
10263            // Stable-boundary capture from the split feed above (None on the legacy shape):
10264            // a restored session previously parked WITHOUT a checkpoint, so the next turn's
10265            // affinity probe declined ("no turn checkpoint retained") and the conversation
10266            // fell back to the frozen prefix entry forever.
10267            turn_ckpt: restored_turn_ckpt,
10268            telem: SpecTelemetryCounters::default(),
10269            capture_at: None,
10270            boundary_captures,
10271            ckpt_at: None,
10272            capture_disabled: false,
10273        })
10274    }
10275
10276    /// Forced-gate exact state comparison. This intentionally reads the real live prefixes from
10277    /// their owning PP devices: matching emitted ids alone would miss a stale `len_d`, recurrent
10278    /// snapshot, or draft-KV row that only corrupts the following round.
10279    pub fn optipipe_compare_session_state(
10280        &self,
10281        e: &Engine,
10282        reference: &SpecSession,
10283        candidate: &SpecSession,
10284    ) -> Result<OptiForkStateIdentity, Box<dyn std::error::Error>> {
10285        fn fail(what: &str) -> Box<dyn std::error::Error> {
10286            format!("optipipe state mismatch: {what}").into()
10287        }
10288        fn same_f32(a: &[f32], b: &[f32]) -> bool {
10289            a.len() == b.len() && a.iter().zip(b).all(|(x, y)| x.to_bits() == y.to_bits())
10290        }
10291        fn compare_layers(
10292            es: &Engine,
10293            range: std::ops::Range<usize>,
10294            reference: &SpecSession,
10295            candidate: &SpecSession,
10296            report: &mut OptiForkStateIdentity,
10297        ) -> Result<(), Box<dyn std::error::Error>> {
10298            for il in range {
10299                match (&reference.cache.kv[il], &candidate.cache.kv[il]) {
10300                    (Some(a), Some(b)) => {
10301                        if a.len != b.len {
10302                            return Err(fail(&format!(
10303                                "layer {il} host KV len {} != {}",
10304                                a.len, b.len
10305                            )));
10306                        }
10307                        let ad = es.dtoh_i32(&a.len_d)?;
10308                        let bd = es.dtoh_i32(&b.len_d)?;
10309                        if ad != bd || ad.first().copied() != Some(a.len as i32) {
10310                            return Err(fail(&format!(
10311                                "layer {il} device KV len {ad:?} != {bd:?} (host={})",
10312                                a.len,
10313                            )));
10314                        }
10315                        let kb = a.len * a.k_tok_bytes;
10316                        let vb = a.len * a.v_tok_bytes;
10317                        if kb > 0 {
10318                            let ak = es.dtoh_u8_view(&a.k.slice(0..kb))?;
10319                            let bk = es.dtoh_u8_view(&b.k.slice(0..kb))?;
10320                            if ak != bk {
10321                                let at = ak.iter().zip(&bk).position(|(x, y)| x != y).unwrap();
10322                                return Err(fail(&format!(
10323                                    "layer {il} K bytes at byte {at} row {} offset {}: {} != {}",
10324                                    at / a.k_tok_bytes,
10325                                    at % a.k_tok_bytes,
10326                                    ak[at],
10327                                    bk[at],
10328                                )));
10329                            }
10330                        }
10331                        if vb > 0 {
10332                            let av = es.dtoh_u8_view(&a.v.slice(0..vb))?;
10333                            let bv = es.dtoh_u8_view(&b.v.slice(0..vb))?;
10334                            if av != bv {
10335                                let at = av.iter().zip(&bv).position(|(x, y)| x != y).unwrap();
10336                                return Err(fail(&format!(
10337                                    "layer {il} V bytes at byte {at} row {} offset {}: {} != {}",
10338                                    at / a.v_tok_bytes,
10339                                    at % a.v_tok_bytes,
10340                                    av[at],
10341                                    bv[at],
10342                                )));
10343                            }
10344                        }
10345                        report.trunk_kv_bytes += kb + vb;
10346                    }
10347                    (None, None) => {}
10348                    _ => return Err(fail(&format!("layer {il} KV presence"))),
10349                }
10350                match (&reference.cache.recur[il], &candidate.cache.recur[il]) {
10351                    (Some(a), Some(b)) => {
10352                        let ac = es.dtoh(&a.conv_state)?;
10353                        let bc = es.dtoh(&b.conv_state)?;
10354                        if !same_f32(&ac, &bc) {
10355                            return Err(fail(&format!("layer {il} conv state")));
10356                        }
10357                        let as_ = es.dtoh(&a.ssm_state)?;
10358                        let bs = es.dtoh(&b.ssm_state)?;
10359                        if !same_f32(&as_, &bs) {
10360                            return Err(fail(&format!("layer {il} SSM state")));
10361                        }
10362                        report.recurrent_bytes += (ac.len() + as_.len()) * 4;
10363                    }
10364                    (None, None) => {}
10365                    _ => return Err(fail(&format!("layer {il} recurrent presence"))),
10366                }
10367            }
10368            Ok(())
10369        }
10370
10371        if reference.committed != candidate.committed {
10372            return Err(fail("committed token ids"));
10373        }
10374        if reference.cache.pos != candidate.cache.pos
10375            || reference.cache.max_ctx != candidate.cache.max_ctx
10376        {
10377            return Err(fail("cache pos/capacity"));
10378        }
10379        if reference.pending_tok != candidate.pending_tok
10380            || reference.next_pred != candidate.next_pred
10381            || reference.sctr != candidate.sctr
10382            || reference.uctr != candidate.uctr
10383        {
10384            return Err(fail("pending/prediction/counter tail"));
10385        }
10386
10387        let mut report = OptiForkStateIdentity::default();
10388        if let Some(fence) = crate::pp::pp_cuts(self.layers.len()) {
10389            let rt = crate::pp::PpNRt::get(e)?;
10390            for stage in 0..rt.n_stages() {
10391                let _scope = rt.enter(stage);
10392                compare_layers(
10393                    rt.engine(stage, e),
10394                    fence[stage]..fence[stage + 1],
10395                    reference,
10396                    candidate,
10397                    &mut report,
10398                )?;
10399            }
10400        } else {
10401            compare_layers(e, 0..self.layers.len(), reference, candidate, &mut report)?;
10402        }
10403
10404        if reference.scratch.plane_count() != candidate.scratch.plane_count() {
10405            return Err(fail("draft scratch plane count"));
10406        }
10407        for index in 0..reference.scratch.plane_count() {
10408            let (a, _) = reference.scratch.plane(index);
10409            let (b, _) = candidate.scratch.plane(index);
10410            if a.len != b.len
10411                || a.kv_dim_k != b.kv_dim_k
10412                || a.kv_dim_v != b.kv_dim_v
10413                || a.k_tok_bytes != b.k_tok_bytes
10414                || a.v_tok_bytes != b.v_tok_bytes
10415                || e.dtoh_i32(&a.len_d)? != e.dtoh_i32(&b.len_d)?
10416            {
10417                return Err(fail(&format!("draft scratch plane {index} length/layout")));
10418            }
10419            let kb = a.len * a.k_tok_bytes;
10420            let vb = a.len * a.v_tok_bytes;
10421            if kb > 0 && e.dtoh_u8_view(&a.k.slice(0..kb))? != e.dtoh_u8_view(&b.k.slice(0..kb))? {
10422                return Err(fail(&format!("draft scratch plane {index} K bytes")));
10423            }
10424            if vb > 0 && e.dtoh_u8_view(&a.v.slice(0..vb))? != e.dtoh_u8_view(&b.v.slice(0..vb))? {
10425                return Err(fail(&format!("draft scratch plane {index} V bytes")));
10426            }
10427            report.scratch_kv_bytes += kb + vb;
10428        }
10429
10430        match (&reference.last_h, &candidate.last_h) {
10431            (Some(a), Some(b)) => {
10432                let ah = e.dtoh(a)?;
10433                let bh = e.dtoh(b)?;
10434                if !same_f32(&ah, &bh) {
10435                    return Err(fail("last hidden/seed bytes"));
10436                }
10437                report.hidden_bytes = ah.len() * 4;
10438            }
10439            (None, None) => {}
10440            _ => return Err(fail("last hidden/seed presence")),
10441        }
10442        Ok(report)
10443    }
10444
10445    /// SESSION-AFFINITY REWIND (lane/session-affinity, 2026-08-05): roll `sess` back to its
10446    /// retained prompt-end checkpoint, so a request whose prompt matches
10447    /// `committed[..rewind_pos()]` exactly can resume there and prime only its own delta.
10448    ///
10449    /// EXACTNESS. After this returns, the session is byte-for-byte the state it was in AT that
10450    /// boundary: full-attn KV truncated to it (append-only, position-addressed), GDN conv/ssm
10451    /// restored from the device copy taken there, draft scratch length reset, `committed`
10452    /// truncated, `last_h` = the boundary's predecessor anchor. That is precisely the state a
10453    /// fresh prime of `committed[..pos]` would have produced, so the following suffix prime and
10454    /// every burst after it are identical to a cold run of the same token stream — the
10455    /// committed-tokens-authoritative contract.
10456    ///
10457    /// `next_pred` and `pending_tok` are CLEARED: both describe generation past the boundary,
10458    /// which the rewind discards. The caller therefore must supply a non-empty suffix (a
10459    /// rewound session cannot serve an empty-suffix continuation burst — there is nothing to
10460    /// continue). The persistent draft graph survives: it bakes only session-stable pointers
10461    /// (the scratch KV, the resident embedding), none of which the rewind moves.
10462    ///
10463    /// The checkpoint is CONSUMED (`turn_ckpt` taken): its snapshot buffers are freed here, and
10464    /// this turn's own prime installs a fresh one at the new prompt end. Returns the position
10465    /// rewound to, or `None` when the session holds no checkpoint (caller: full re-prime).
10466    pub fn spec_rewind_to_checkpoint(
10467        &self,
10468        e: &Engine,
10469        sess: &mut SpecSession,
10470    ) -> Result<Option<usize>, Box<dyn std::error::Error>> {
10471        if sess.turn_ckpt.as_ref().is_some_and(|ckpt| {
10472            !sess.cache.can_rollback(&ckpt.snap, 0) || !sess.scratch.can_rewind_to(ckpt.pos)
10473        }) {
10474            return Err(
10475                "SWA ring rewind checkpoint has been lapped; full re-prime required".into(),
10476            );
10477        }
10478        let Some(ckpt) = sess.turn_ckpt.take() else {
10479            return Ok(None);
10480        };
10481        assert!(
10482            ckpt.pos <= sess.committed.len(),
10483            "checkpoint past committed ({} > {})",
10484            ckpt.pos,
10485            sess.committed.len()
10486        );
10487        // Restore through each layer's owning engine. A single primary-engine rollback is not
10488        // sufficient when the serving cache is stage-owned under cross-device PP.
10489        crate::pp::restore_cache_checkpoint(e, self, None, &mut sess.cache, &ckpt.snap)?;
10490        debug_assert_eq!(
10491            sess.cache.pos, ckpt.pos,
10492            "rollback landed off the checkpoint"
10493        );
10494        sess.scratch.set_len(e, ckpt.pos)?;
10495        sess.committed.truncate(ckpt.pos);
10496        sess.last_h = Some(ckpt.last_h);
10497        sess.next_pred = None;
10498        sess.pending_tok = None;
10499        Ok(Some(ckpt.pos))
10500    }
10501
10502    /// Grow a parked speculative session to `target_cap` and rewind it to its retained turn
10503    /// checkpoint without re-priming the checkpoint prefix.
10504    ///
10505    /// The trunk cache is restored exactly like a plain grown cache: append-only full-attention
10506    /// KV rows come from the parked cache, while recurrent state comes from the checkpoint's
10507    /// owned snapshot. The MTP scratch is also context-linear and its rows below the checkpoint
10508    /// remain authoritative, so they are copied into a fresh larger scratch before its length is
10509    /// truncated. Pointer-baking draft graphs are dropped and recaptured on the next burst.
10510    ///
10511    /// All fallible work completes before `sess` is mutated. A failed allocation or copy leaves
10512    /// the parked session intact, allowing the caller one reclaim-and-retry attempt.
10513    pub fn spec_grow_and_rewind_to_checkpoint(
10514        &self,
10515        e: &Engine,
10516        sess: &mut SpecSession,
10517        target_cap: usize,
10518    ) -> Result<Option<usize>, Box<dyn std::error::Error>> {
10519        if target_cap <= sess.cache.max_ctx {
10520            return self.spec_rewind_to_checkpoint(e, sess);
10521        }
10522        let Some(ckpt) = sess.turn_ckpt.as_ref() else {
10523            return Ok(None);
10524        };
10525        if ckpt.pos == 0 || ckpt.pos > sess.committed.len() {
10526            return Err(format!(
10527                "checkpoint pos {} outside committed length {}",
10528                ckpt.pos,
10529                sess.committed.len(),
10530            )
10531            .into());
10532        }
10533        if ckpt.pos > target_cap {
10534            return Err(format!(
10535                "checkpoint pos {} exceeds grown capacity {target_cap}",
10536                ckpt.pos,
10537            )
10538            .into());
10539        }
10540
10541        let mut grown_cache = crate::pp::new_cache_planned(e, &self.cfg, &self.plan, target_cap)?;
10542        let mut grown_scratch = self.new_mtp_scratch(e, target_cap)?;
10543        crate::pp::restore_cache_checkpoint(
10544            e,
10545            self,
10546            Some(&sess.cache),
10547            &mut grown_cache,
10548            &ckpt.snap,
10549        )?;
10550
10551        if sess.scratch.plane_count() != grown_scratch.plane_count() {
10552            return Err("checkpoint draft plane count mismatch".into());
10553        }
10554        for index in 0..sess.scratch.plane_count() {
10555            let (src, _) = sess.scratch.plane(index);
10556            let (dst, _) = grown_scratch.plane_mut(index);
10557            if ckpt.pos > src.len
10558                || src.kv_dim_k != dst.kv_dim_k
10559                || src.kv_dim_v != dst.kv_dim_v
10560                || src.k_tok_bytes != dst.k_tok_bytes
10561                || src.v_tok_bytes != dst.v_tok_bytes
10562            {
10563                return Err(format!(
10564                    "checkpoint draft plane {index} layout mismatch (pos {}, source len {})",
10565                    ckpt.pos, src.len,
10566                )
10567                .into());
10568            }
10569            match (&src.ring, dst.ring.as_ref()) {
10570                (Some(sring), Some(_)) => {
10571                    // Ring-backed draft plane (step35): `ckpt.pos` is absolute and exceeds the
10572                    // physical rows once lapped — same class as the trunk-KV restore panic
10573                    // (2026-08-29 warm-turn-at-40k). Copy the aligned live window, rebase.
10574                    let (new_base, phys) = sring.restore_plan(ckpt.pos).map_err(|err| {
10575                        format!("checkpoint draft plane {index} SWA restore refused: {err}")
10576                    })?;
10577                    let rows = phys.len();
10578                    let kb = rows * src.k_tok_bytes;
10579                    let vb = rows * src.v_tok_bytes;
10580                    if kb > 0 {
10581                        e.copy_u8_range_into(
10582                            &mut dst.k,
10583                            0,
10584                            &src.k,
10585                            phys.start * src.k_tok_bytes,
10586                            kb,
10587                        )?;
10588                    }
10589                    if vb > 0 {
10590                        e.copy_u8_range_into(
10591                            &mut dst.v,
10592                            0,
10593                            &src.v,
10594                            phys.start * src.v_tok_bytes,
10595                            vb,
10596                        )?;
10597                    }
10598                    dst.ring
10599                        .as_mut()
10600                        .expect("ring presence checked above")
10601                        .apply_rebase(new_base);
10602                    if let Some(base_d) = dst.base_d.as_mut() {
10603                        e.set_i32_one(base_d, new_base as i32)?;
10604                    }
10605                }
10606                (None, None) => {
10607                    let kb = ckpt.pos * src.k_tok_bytes;
10608                    let vb = ckpt.pos * src.v_tok_bytes;
10609                    if kb > 0 {
10610                        e.copy_u8_into(&mut dst.k, 0, &src.k, kb)?;
10611                    }
10612                    if vb > 0 {
10613                        e.copy_u8_into(&mut dst.v, 0, &src.v, vb)?;
10614                    }
10615                }
10616                _ => {
10617                    return Err(format!("checkpoint draft plane {index} ring/flat mismatch").into());
10618                }
10619            }
10620        }
10621        grown_scratch.set_len(e, ckpt.pos)?;
10622        // The old scratch is dropped immediately after publication below. Bound its D2D reads
10623        // first; growth happens once per rewritten turn, outside the decode hot loop.
10624        e.stream().synchronize()?;
10625
10626        let ckpt = sess
10627            .turn_ckpt
10628            .take()
10629            .expect("checkpoint remained present through transactional grow");
10630        let pos = ckpt.pos;
10631        sess.cache = grown_cache;
10632        sess.scratch = grown_scratch;
10633        sess.committed.truncate(pos);
10634        sess.last_h = Some(ckpt.last_h);
10635        sess.next_pred = None;
10636        sess.pending_tok = None;
10637        sess.draft_ctx = None;
10638        debug_assert_eq!(sess.cache.pos, pos, "grown rewind landed off checkpoint");
10639        debug_assert!(
10640            (0..sess.scratch.plane_count()).all(|index| sess.scratch.plane(index).0.len == pos),
10641            "grown draft rewind landed off checkpoint"
10642        );
10643        Ok(Some(pos))
10644    }
10645
10646    /// Commit a carried pending bonus (see SpecSession::pending_tok): one T=1 trunk pass
10647    /// (its logits' argmax becomes next_pred) + the draft-KV fill at the carried anchor —
10648    /// byte-identical to the pre-carry session tail. Required before a non-empty-suffix
10649    /// prime, a sampled turn, or parking a session for pool reuse. No-op without a pending.
10650    /// `sampling` is the sampler of the request that will CONSUME the resulting `next_pred`
10651    /// (lane/sampled-spec-quality): this is a boundary site like any other, so a sampled
10652    /// consumer must get a DRAWN token, not an argmax. Pass `None` from the park/demote
10653    /// callers — a pending only ever exists on the GREEDY tail, and the consumer of a
10654    /// park-time flush is a future request whose sampler is not knowable here (residual
10655    /// named at the pool-resume probe in worker.rs and in SAMPLED-QUALITY.md).
10656    pub fn spec_flush_pending(
10657        &self,
10658        e: &Engine,
10659        sess: &mut SpecSession,
10660        sampling: Option<SpecSampling>,
10661    ) -> Result<(), Box<dyn std::error::Error>> {
10662        let Some(b) = sess.pending_tok.take() else {
10663            return Ok(());
10664        };
10665        if self.mtp.is_none() {
10666            return Err("pending carry requires an MTP head".into());
10667        }
10668        let n_embd = self.cfg.n_embd as usize;
10669        let (embd_qt, embd_rb) = self.embd.qt_and_row_bytes(n_embd);
10670        let embd_gpu = if spec_host_embd() {
10671            None
10672        } else {
10673            Some(
10674                self.embd_gpu
10675                    .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload")),
10676            )
10677        };
10678        let embd_dev = embd_gpu.map(|g| (g, embd_qt, embd_rb));
10679        let pos_b = sess.cache.pos;
10680        sess.scratch.set_len(e, pos_b)?;
10681        let (lg_b, hb) = self.spec_target_step_h(e, b, &mut sess.cache)?;
10682        sess.next_pred = Some(match sampling {
10683            Some(sp) if sp.temp > 0.0 && spec_sampled_boundary_on() => {
10684                // window includes `b` itself: it is committed by this pass, and the pre-lane
10685                // code never counted a boundary token in the penalty history at all.
10686                let hist = pen_window_seed(&sess.committed, &[b], sp.penalty_last_n);
10687                sample_boundary_token(e, &lg_b, &sp, &hist, &mut sess.sctr, "flush-pending")?
10688            }
10689            _ => argmax(&lg_b) as u32,
10690        });
10691        let anchor = sess
10692            .last_h
10693            .as_ref()
10694            .expect("pending carry requires last_h (the predecessor-row anchor)");
10695        self.mtp_kv_fill_all(e, &[b], anchor, pos_b, &mut sess.scratch, embd_dev)?;
10696        sess.last_h = Some(hb);
10697        sess.committed.push(b);
10698        Ok(())
10699    }
10700
10701    /// Solo target feed used only at speculative round boundaries. Step35 serving made its
10702    /// staged batched B=1 graph authoritative, so a speculative session must enter and leave
10703    /// rounds through that same graph. Other model families keep their eager T=1 contract.
10704    fn spec_target_step_h(
10705        &self,
10706        e: &Engine,
10707        token: u32,
10708        cache: &mut Cache,
10709    ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
10710        if !self.sliding_gated_moe_batch_program() && !self.batched_serving_numeric_class() {
10711            return self.decode_step_h(e, token, cache);
10712        }
10713        let pos0 = cache.pos;
10714        let (logits, hidden) = self.decode_step_t_core(e, &[token], pos0, cache, None, None)?;
10715        Ok((e.dtoh(&logits)?, hidden))
10716    }
10717
10718    /// The archs whose LIVE B=1 serving runs the generic BATCHED numeric class (decode_step_batch
10719    /// walk + batched head), so their spec verify must run the SAME class. MoE learned this
10720    /// 2026-08-14 AM (4b777ccc5); the dense hybrid reproduced the identical near-tie flip class
10721    /// the same day on Qwen3.8-27B — eager-class verify logits drift from batched-class serving
10722    /// logits ("1 ULP at layer 2 → 2.3e-1 logit maxdiff at the head"), and the GDN recurrence
10723    /// carries the drift until a near-tie flips deep in generation. One predicate so the five
10724    /// dispatch sites cannot drift apart again.
10725    /// Draft-graph head admissibility (lane/draftcost-moe, 2026-08-20): the capture body
10726    /// (`mtp_head_forward_cap`) supports Dense heads AND resident-MoE heads
10727    /// (`Ffn::Moe(m) if m.dev_exps.is_some()`); non-resident MoE still refuses inside the
10728    /// capture and the caller falls back to the eager chain by design. Trunk FFN class is
10729    /// irrelevant — the graph body is the HEAD forward only. One predicate for all three
10730    /// eligibility sites so they cannot drift (the serving numeric-class lesson).
10731    fn mtp_graph_capturable(&self) -> bool {
10732        // EVERY loaded head must be capturable: the multi-head chain graphs capture each
10733        // head's forward (lane/step37-draft-graph-serving-20260830), so one SLRU-locked MoE
10734        // head anywhere in the chain refuses capture for the whole chain (loudly, via the
10735        // capture-site WARN) rather than capturing a subset the launch order cannot honor.
10736        let head_ok = |m: &MtpHead| match &m.ffn {
10737            crate::hybrid::Ffn::Dense { .. } => true,
10738            crate::hybrid::Ffn::Moe(mo) => mo.dev_exps.is_some(),
10739        };
10740        self.mtp.as_ref().map(&head_ok).unwrap_or(false) && self.mtp_extra.iter().all(head_ok)
10741    }
10742
10743    fn batched_serving_numeric_class(&self) -> bool {
10744        self.plan
10745            .trunk_operations()
10746            .contains(&memra_gguf::model_plan::OperationKind::GatedDeltaNet)
10747    }
10748
10749    /// The family the MTP verify-graph default was measured on: GatedDeltaNet state layers
10750    /// (a `recur` mixer) together with a routed-MoE FFN — Ornith-1.5-35B-A3B and its kin. The
10751    /// server-side twin of this test is `model_forces_spec_replay` (GatedDeltaNet + MoeMlp);
10752    /// keeping the engine's own version structural rather than name-based means a new
10753    /// checkpoint of the same shape inherits the default, and a different shape does not.
10754    /// pub(crate) since lane/graph-launch-guard-sweep-20260831: `dspark_vg_admission_debt`
10755    /// consults it so the MTP-route pool stops escaping the admission charge.
10756    pub(crate) fn vgraph_family_default(&self) -> bool {
10757        let has_linear = self
10758            .layers
10759            .iter()
10760            .any(|l| matches!(l.mixer, Mixer::Linear(_)));
10761        let has_moe = self
10762            .layers
10763            .iter()
10764            .any(|l| matches!(l.ffn, crate::hybrid::Ffn::Moe(_)));
10765        has_linear && has_moe
10766    }
10767
10768    fn sliding_gated_moe_batch_program(&self) -> bool {
10769        self.uses_sliding_gated_moe_program()
10770    }
10771
10772    fn gemma_batch_program(&self) -> bool {
10773        self.uses_gemma_program()
10774    }
10775
10776    /// Reduced-matrix admission for increment 1. This deliberately does not change the PP-2
10777    /// serving policy: the worker calls it only after `MEMRA_SPEC_PIPE=1` and an explicit spec
10778    /// session already exist.
10779    pub fn spec_pipe_available(&self, e: &Engine) -> bool {
10780        if std::env::var("MEMRA_SPEC_PIPE").as_deref() != Ok("1")
10781            || !spec_devacc()
10782            || spec_replay_env_enabled()
10783            || spec_stream()
10784            || std::env::var("MEMRA_SPEC_ADAPT").as_deref() == Ok("1")
10785            || std::env::var("MEMRA_SPEC_PMIN0").as_deref() == Ok("1")
10786            || std::env::var("MEMRA_SPEC_PP_ANATOMY").as_deref() == Ok("1")
10787            || std::env::var("MEMRA_SPEC_PMIN")
10788                .ok()
10789                .and_then(|v| v.parse::<f32>().ok())
10790                .unwrap_or(0.0)
10791                > 0.0
10792            || self.is_gemma4_e4b()
10793            || self.gemma_batch_program()
10794            || self.mtp.is_none()
10795            || !self.mtp_extra.is_empty()
10796        {
10797            return false;
10798        }
10799        let Some(cuts) = crate::pp::pp_cuts(self.layers.len()) else {
10800            return false;
10801        };
10802        if cuts.len() != 3 || crate::pp::pp2_streams_off() || !crate::pp::spec_pp_on() {
10803            return false;
10804        }
10805        crate::pp::PpNRt::get(e)
10806            .map(|rt| rt.n_stages() == 2 && rt.cross_device())
10807            .unwrap_or(false)
10808    }
10809
10810    /// Two warm greedy continuation bursts over one PP-2 interval coordinator. The two existing
10811    /// `generate_spec_inner2` call stacks own all per-session round locals; only phase issue order
10812    /// changes. No callback is accepted in increment 1 — the worker publishes each completed burst.
10813    #[allow(clippy::too_many_arguments)]
10814    pub fn generate_spec_session_pair(
10815        &self,
10816        e: &Engine,
10817        sess_a: &mut SpecSession,
10818        max_new_a: usize,
10819        k_a: usize,
10820        sess_b: &mut SpecSession,
10821        max_new_b: usize,
10822        k_b: usize,
10823    ) -> Result<((Vec<u32>, usize, usize), (Vec<u32>, usize, usize)), Box<dyn std::error::Error>>
10824    {
10825        if !self.spec_pipe_available(e) {
10826            return Err("two-session speculative pipeline is outside its reduced matrix".into());
10827        }
10828        if max_new_a == 0 || max_new_b == 0 || k_a == 0 || k_b == 0 {
10829            return Err(
10830                "two-session speculative pipeline requires non-empty positive-K bursts".into(),
10831            );
10832        }
10833        for sess in [&*sess_a, &*sess_b] {
10834            if sess.committed.is_empty()
10835                || sess.last_h.is_none()
10836                || (sess.next_pred.is_none() && sess.pending_tok.is_none())
10837            {
10838                return Err("two-session speculative pipeline requires warm continuations".into());
10839            }
10840        }
10841
10842        let graph_ok = std::env::var("MEMRA_SPEC_NOGRAPH").is_err()
10843            && !spec_host_embd()
10844            && self.mtp_graph_capturable()
10845            && self.mtp_extra.is_empty()
10846            && !crate::model::full_prec_enabled();
10847        let graph_a = graph_ok && k_a + 2 < 96;
10848        let graph_b = graph_ok && k_b + 2 < 96;
10849        let was_tracking = e.ctx().is_event_tracking();
10850        if (graph_a || graph_b) && was_tracking {
10851            unsafe {
10852                e.ctx().disable_event_tracking();
10853            }
10854        }
10855
10856        static LOGGED: std::sync::Once = std::sync::Once::new();
10857        LOGGED.call_once(|| {
10858            eprintln!("[spec-pipe] two-session PP-2 continuation pipeline engaged");
10859        });
10860        let sync = std::sync::Arc::new(SpecPipeSync::new());
10861        let lane_a = SpecPipeLane {
10862            sync: sync.clone(),
10863            lane: 0,
10864        };
10865        let lane_b = SpecPipeLane { sync, lane: 1 };
10866        let mut sess_b_ptr = SpecPipeSessionPtr(sess_b as *mut SpecSession);
10867        let (result_a, result_b) = std::thread::scope(|scope| {
10868            let b = scope.spawn(move || {
10869                let mut finish = SpecPipeFinish::new(&lane_b);
10870                let sess_b = unsafe { sess_b_ptr.get_mut() };
10871                let result = e
10872                    .ctx()
10873                    .bind_to_thread()
10874                    .map_err(|err| err.to_string())
10875                    .and_then(|_| {
10876                        self.generate_spec_inner2(
10877                            e,
10878                            &[],
10879                            max_new_b,
10880                            k_b,
10881                            graph_b,
10882                            Some(sess_b),
10883                            None,
10884                            None,
10885                            None,
10886                            None,
10887                            Some(&lane_b),
10888                        )
10889                        .map_err(|err| err.to_string())
10890                    });
10891                finish.close(result.is_err());
10892                result
10893            });
10894            let mut finish = SpecPipeFinish::new(&lane_a);
10895            let result_a = self.generate_spec_inner2(
10896                e,
10897                &[],
10898                max_new_a,
10899                k_a,
10900                graph_a,
10901                Some(sess_a),
10902                None,
10903                None,
10904                None,
10905                None,
10906                Some(&lane_a),
10907            );
10908            finish.close(result_a.is_err());
10909            let result_b = b
10910                .join()
10911                .map_err(|_| "paired speculative session B panicked".to_string())
10912                .and_then(|r| r);
10913            (result_a, result_b)
10914        });
10915
10916        if (graph_a || graph_b) && was_tracking {
10917            unsafe {
10918                e.ctx().enable_event_tracking();
10919            }
10920        }
10921        let result_a = result_a?;
10922        let result_b = result_b.map_err(|err| -> Box<dyn std::error::Error> { err.into() })?;
10923        Ok((result_a, result_b))
10924    }
10925
10926    /// One spec-decode turn on a live session. `suffix` = the NEW tokens only (turn N+1's user
10927    /// message rendered through the chat template continuation). Returns (new tokens emitted,
10928    /// drafted, accepted); session.committed grows by suffix + emitted.
10929    pub fn generate_spec_session(
10930        &self,
10931        e: &Engine,
10932        sess: &mut SpecSession,
10933        suffix: &[u32],
10934        max_new: usize,
10935        k: usize,
10936    ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
10937        self.generate_spec_session_sampled(e, sess, suffix, max_new, k, None, None)
10938    }
10939
10940    /// Serve-path sampled spec: routes the burst through the rejection-sampling verify with
10941    /// per-SESSION Philox continuity (sess.sctr/uctr). None = env-driven (CLI) or greedy.
10942    /// Filters (top-k/p/min-p) apply SYMMETRICALLY to draft q and verify p — distribution-exact
10943    /// for the filtered target (feat/filtered-spec).
10944    ///
10945    /// `on_commit` (sse-cadence, 2026-08-05): called with each newly-emitted slice of the
10946    /// output — once right after the prime's first token, then once per round commit — so a
10947    /// streaming caller can flush text at round cadence instead of once per burst. The slices
10948    /// are disjoint, in order, and concatenate to exactly the returned token vec. Emission-
10949    /// timing only: token bytes, session state, and exactness are untouched.
10950    ///
10951    /// The returned bool is a CONTINUE-VERDICT (admission yield, 2026-08-06): `false` ends
10952    /// the burst at the current round boundary, exactly as if `max_new` had been reached —
10953    /// the caller's scheduler regains control without waiting the burst out. Burst size is
10954    /// content-neutral (spec-levers battery), so an early exit moves WHEN the burst returns,
10955    /// never what tokens say. The slice may be EMPTY (a poll-only boundary — round-stream
10956    /// drains and the defensive tail flush can land with nothing new committed).
10957    #[allow(clippy::too_many_arguments)]
10958    pub fn generate_spec_session_sampled(
10959        &self,
10960        e: &Engine,
10961        sess: &mut SpecSession,
10962        suffix: &[u32],
10963        max_new: usize,
10964        k: usize,
10965        sampling: Option<SpecSampling>,
10966        on_commit: Option<&mut dyn FnMut(&[u32]) -> bool>,
10967    ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
10968        self.generate_spec_session_sampled_prime_split(
10969            e, sess, suffix, max_new, k, sampling, None, on_commit,
10970        )
10971    }
10972
10973    /// Serve-only cold-prime segmentation twin. `prime_split` is the same stable boundary the
10974    /// plain worker would honor before entering its sub-floor tokenwise tail; warm continuations
10975    /// pass `None` and stay on the existing zero-prime path.
10976    #[allow(clippy::too_many_arguments)]
10977    pub fn generate_spec_session_sampled_prime_split(
10978        &self,
10979        e: &Engine,
10980        sess: &mut SpecSession,
10981        suffix: &[u32],
10982        max_new: usize,
10983        k: usize,
10984        sampling: Option<SpecSampling>,
10985        prime_split: Option<usize>,
10986        on_commit: Option<&mut dyn FnMut(&[u32]) -> bool>,
10987    ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
10988        self.generate_spec_session_constrained_prime_split(
10989            e,
10990            sess,
10991            suffix,
10992            max_new,
10993            k,
10994            sampling,
10995            None,
10996            prime_split,
10997            on_commit,
10998        )
10999    }
11000
11001    /// `generate_spec_session_sampled` + GRAMMAR (constrained decoding, 2026-08-03): the
11002    /// hook truncates acceptance at the first grammar-illegal token AFTER the exactness
11003    /// verify (grammar is an extra rejection rule, ordering like the batched-verify twins)
11004    /// and replaces an illegal bonus with the MASKED argmax of the target's own verify
11005    /// column — token-identical to constrained plain greedy decode. GREEDY only (the
11006    /// worker routes sampled constrained to plain decode). Acceptance under tight grammars
11007    /// may drop (drafter is unconstrained); that is measured, not hidden.
11008    #[allow(clippy::too_many_arguments)]
11009    pub fn generate_spec_session_constrained(
11010        &self,
11011        e: &Engine,
11012        sess: &mut SpecSession,
11013        suffix: &[u32],
11014        max_new: usize,
11015        k: usize,
11016        sampling: Option<SpecSampling>,
11017        constraint: Option<&mut dyn SpecConstraint>,
11018        on_commit: Option<&mut dyn FnMut(&[u32]) -> bool>,
11019    ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
11020        self.generate_spec_session_constrained_prime_split(
11021            e, sess, suffix, max_new, k, sampling, constraint, None, on_commit,
11022        )
11023    }
11024
11025    #[allow(clippy::too_many_arguments)]
11026    pub fn generate_spec_session_constrained_prime_split(
11027        &self,
11028        e: &Engine,
11029        sess: &mut SpecSession,
11030        suffix: &[u32],
11031        max_new: usize,
11032        k: usize,
11033        sampling: Option<SpecSampling>,
11034        constraint: Option<&mut dyn SpecConstraint>,
11035        prime_split: Option<usize>,
11036        on_commit: Option<&mut dyn FnMut(&[u32]) -> bool>,
11037    ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
11038        if constraint.is_some() && sampling.is_some_and(|s| s.temp > 0.0) {
11039            return Err(
11040                "constrained spec decode is greedy-only (worker routes sampled \
11041                        constrained to plain decode)"
11042                    .into(),
11043            );
11044        }
11045        // PENDING-CARRY entry flush: a carried bonus precedes any new suffix in the sequence,
11046        // so it must commit BEFORE the suffix primes; the sampled path doesn't carry (its
11047        // round-0 accept needs the commit pass's logits). Empty-suffix greedy bursts — the
11048        // serve continuation case — consume the carry in-loop with zero solo passes.
11049        if sess.pending_tok.is_some()
11050            && (!suffix.is_empty() || sampling.map_or(false, |s| s.temp > 0.0))
11051        {
11052            self.spec_flush_pending(e, sess, sampling)?;
11053        }
11054
11055        // FULL_PREC forces the EAGER draft: the graph capture would enclose cuBLASLt f32 GEMV
11056        // (the FloatBf16 else-branches) and a bf16_to_f32 dequant alloc — neither is stream-capture
11057        // safe. Eager rides matmul/matmul_decode_exact, which dequant FloatBf16 on use. (§item 2.)
11058        // Multi-head MTP (mtp_extra non-empty) no longer disqualifies: the chain captures
11059        // per-head graphs (lane/step37-draft-graph-serving-20260830, MEMRA_MTP_CHAIN_GRAPH).
11060        let graph_draft = std::env::var("MEMRA_SPEC_NOGRAPH").is_err()
11061            && !spec_host_embd()
11062            && self.mtp_graph_capturable()
11063            && k + 2 < 96
11064            && !crate::model::full_prec_enabled();
11065        let was_tracking = e.ctx().is_event_tracking();
11066        if graph_draft && was_tracking {
11067            unsafe {
11068                e.ctx().disable_event_tracking();
11069            }
11070        }
11071        let r = self.generate_spec_inner2(
11072            e,
11073            suffix,
11074            max_new,
11075            k,
11076            graph_draft,
11077            Some(sess),
11078            sampling,
11079            constraint,
11080            on_commit,
11081            prime_split,
11082            None,
11083        );
11084        if graph_draft && was_tracking {
11085            unsafe {
11086                e.ctx().enable_event_tracking();
11087            }
11088        }
11089        let (out, d, a) = r?;
11090        Ok((out, d, a))
11091    }
11092
11093    pub fn generate_spec(
11094        &self,
11095        e: &Engine,
11096        prompt: &[u32],
11097        max_new: usize,
11098        k: usize,
11099    ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
11100        if crate::pp::pp_cuts(self.layers.len()).is_some()
11101            && !self.rewrite_allowed(memra_gguf::execution_manifest::RewriteSurface::Pipeline)
11102        {
11103            return Err("pipeline rewrite is not qualified for speculative decode".into());
11104        }
11105        if !self.rewrite_allowed(memra_gguf::execution_manifest::RewriteSurface::MtpSpec) {
11106            return Err("speculative rewrite is not qualified for this ModelPlan".into());
11107        }
11108        // FULL_PREC forces eager (see generate_spec_session note): CUDA graph capture cannot
11109        // enclose cuBLASLt f32 GEMV or the bf16_to_f32 dequant alloc the FloatBf16 path needs.
11110        // Multi-head MTP no longer disqualifies (chain graphs; see generate_spec_session).
11111        let graph_draft = std::env::var("MEMRA_SPEC_NOGRAPH").is_err()
11112            && !spec_host_embd()
11113            && self.mtp_graph_capturable()
11114            && k + 2 < 96
11115            && !crate::model::full_prec_enabled();
11116        if !graph_draft {
11117            return self.generate_spec_inner2(
11118                e, prompt, max_new, k, false, None, None, None, None, None, None,
11119            );
11120        }
11121        let was_tracking = e.ctx().is_event_tracking();
11122        if was_tracking {
11123            unsafe {
11124                e.ctx().disable_event_tracking();
11125            }
11126        }
11127        let r = self.generate_spec_inner2(
11128            e, prompt, max_new, k, true, None, None, None, None, None, None,
11129        );
11130        if was_tracking {
11131            unsafe {
11132                e.ctx().enable_event_tracking();
11133            }
11134        }
11135        r
11136    }
11137
11138    fn generate_spec_inner2(
11139        &self,
11140        e: &Engine,
11141        prompt: &[u32],
11142        max_new: usize,
11143        k: usize,
11144        graph_draft: bool,
11145        mut sess: Option<&mut SpecSession>,
11146        sampling: Option<SpecSampling>,
11147        mut constraint: Option<&mut dyn SpecConstraint>,
11148        mut on_commit: Option<&mut dyn FnMut(&[u32]) -> bool>,
11149        prime_split: Option<usize>,
11150        pipe: Option<&SpecPipeLane>,
11151    ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
11152        assert!(k >= 1, "k must be >= 1");
11153        if let Some(p) = pipe {
11154            p.setup_begin()?;
11155        }
11156        // sse-cadence flush cursor: everything in out[..flushed] has been handed to on_commit.
11157        let mut flushed = 0usize;
11158        // admission yield (2026-08-06): on_commit's continue-verdict; false = end the burst
11159        // at the next round boundary (same exit as max_new reached — the session tail runs).
11160        // Initialized by the unconditional post-prime flush below.
11161        let mut keep_going;
11162        let mtp = self
11163            .mtp
11164            .as_ref()
11165            .expect("generate_spec requires an MTP head (nextn_predict_layers>0)");
11166        let n_vocab = self.output.out_features();
11167        // FR-Spec: the draft head may be TRIMMED (fewer rows than n_vocab); the draft argmax runs
11168        // over the draft vocab and the winning index maps through d2t to a TARGET token id.
11169        // Everything downstream (verify/accept/commit) sees target ids only — exactness unchanged.
11170        let d_vocab = mtp
11171            .shared_head_head
11172            .as_ref()
11173            .unwrap_or(&self.output)
11174            .out_features();
11175        if !self.mtp_extra.is_empty() {
11176            if self.plan.draft_source != memra_gguf::model_plan::DraftSourcePlan::Embedded
11177                || self.plan.mtp_blocks.len() != self.mtp_head_count()
11178            {
11179                return Err(
11180                    "multi-head MTP requires one embedded canonical block per loaded head".into(),
11181                );
11182            }
11183            // TRIMMED chains (2026-08-27): every head must carry the SAME d2t — the ranking is
11184            // token-frequency and head-independent, and every downstream remap (per-step argmax,
11185            // stream pack, sampled d2t_dev) reads head 0's map, so equality is what makes that
11186            // single map correct for the whole chain. Mixed trimmed/untrimmed is refused.
11187            for (offset, head) in self.mtp_extra.iter().enumerate() {
11188                if head.d2t != mtp.d2t
11189                    || head
11190                        .shared_head_head
11191                        .as_ref()
11192                        .unwrap_or(&self.output)
11193                        .out_features()
11194                        != d_vocab
11195                {
11196                    return Err(format!(
11197                        "embedded MTP head {} has incompatible draft vocabulary",
11198                        offset + 1
11199                    )
11200                    .into());
11201                }
11202            }
11203            eprintln!(
11204                "[mtp-chain] heads={} policy=step-modulo prefix-replay kv=per-head",
11205                self.mtp_head_count()
11206            );
11207        }
11208        let n_embd = self.cfg.n_embd as usize;
11209        // SESSION MODE: reuse the live cache/scratch, prime only the suffix. `base` = tokens
11210        // already committed (their state is in the caches); 0 = fresh single-shot call.
11211        let session_mode = sess.is_some();
11212        let max_ctx = match sess.as_ref() {
11213            Some(s) => s.cache.max_ctx,
11214            None => prompt.len() + max_new + k + 8,
11215        };
11216        let mut own_cache;
11217        let mut own_scratch;
11218        // PREFIX-CACHE capture request threaded out of the session (lane/spec-prefix-cache):
11219        // (requested split, destination list). Single-shot per burst; fresh calls have none.
11220        let mut sess_capture: Option<(Option<usize>, &mut Vec<SpecBoundaryCapture>)> = None;
11221        // STABLE-BOUNDARY turn-checkpoint request (lane/frspec-multiturn-cache): ABSOLUTE
11222        // committed-length position; consumed one-shot like `capture_at`. None = legacy
11223        // prompt-end capture below.
11224        let mut ckpt_req: Option<usize> = None;
11225        // FAIL-SAFE bit threaded out of the session (see `SpecSession::capture_disabled`).
11226        let mut sess_capture_disabled = false;
11227        let (
11228            cache,
11229            scratch,
11230            mut sess_tail,
11231            mut sess_draft_slot,
11232            mut sess_pending_slot,
11233            sess_ckpt_slot,
11234            sess_telem,
11235        ): (
11236            &mut Cache,
11237            &mut MtpScratch,
11238            Option<(
11239                &mut Vec<u32>,
11240                &mut Option<CudaSlice<f32>>,
11241                &mut Option<u32>,
11242                &mut u32,
11243                &mut u32,
11244            )>,
11245            Option<&mut Option<DraftGraphCtx>>,
11246            Option<&mut Option<u32>>,
11247            Option<&mut Option<SpecCheckpoint>>,
11248            Option<&SpecTelemetryCounters>,
11249        ) = match sess.take() {
11250            Some(sr) => {
11251                let SpecSession {
11252                    cache,
11253                    scratch,
11254                    committed,
11255                    last_h,
11256                    next_pred,
11257                    sctr: s_sctr,
11258                    uctr: s_uctr,
11259                    draft_ctx,
11260                    pending_tok,
11261                    turn_ckpt,
11262                    telem,
11263                    capture_at,
11264                    boundary_captures,
11265                    ckpt_at,
11266                    capture_disabled,
11267                } = sr;
11268                sess_capture_disabled = *capture_disabled;
11269                sess_capture = Some((capture_at.take(), boundary_captures));
11270                ckpt_req = ckpt_at.take();
11271                (
11272                    cache,
11273                    scratch,
11274                    Some((committed, last_h, next_pred, s_sctr, s_uctr)),
11275                    Some(draft_ctx),
11276                    Some(pending_tok),
11277                    Some(turn_ckpt),
11278                    Some(telem),
11279                )
11280            }
11281            None => {
11282                // STAGE-OWNED KV (lane/pp2-spec 2026-08-06) — see `new_session`. Door shut =
11283                // `Cache::new` verbatim.
11284                own_cache = crate::pp::new_cache_planned(e, &self.cfg, &self.plan, max_ctx)?;
11285                // Persistent scratch = max_ctx rows (~2KB/token quantized).
11286                own_scratch = self.new_mtp_scratch(e, max_ctx)?;
11287                (
11288                    &mut own_cache,
11289                    &mut own_scratch,
11290                    None,
11291                    None,
11292                    None,
11293                    None,
11294                    None,
11295                )
11296            }
11297        };
11298        if scratch.plane_count() != self.mtp_head_count() {
11299            return Err(format!(
11300                "MTP scratch/head count mismatch ({}/{})",
11301                scratch.plane_count(),
11302                self.mtp_head_count()
11303            )
11304            .into());
11305        }
11306        let base = cache.pos;
11307        // PENDING-CARRY consume (2026-08-01): a carried bonus reaches here only on the
11308        // empty-suffix GREEDY continuation path (generate_spec_session_sampled flushed every
11309        // other case). It enters the round loop as round-0's pending — verify col 0 — exactly
11310        // like a mid-burst full-accept boundary: no init feed, no tail commit pass.
11311        let carried_pending: Option<u32> = sess_pending_slot.as_mut().and_then(|s| s.take());
11312        // PERSISTENT DRAFT KV (the only mode since 2026-07-08 — the legacy round-local scratch,
11313        // MEMRA_SPEC_KVLOCAL, measured -35 acceptance pts on the 27B p3 sweep and was removed;
11314        // acceptance-only — exactness is verify's job either way).
11315        // HIDDEN-PAIRING CONVENTION (DEFAULT = predecessor-row, 2026-07-04 — the 27B acceptance
11316        // unlock, +16pts): the MTP head is TRAINED on rows pairing token x_p with the trunk
11317        // hidden of its PREDECESSOR h_{p-1} (the reference engine's mtp_update shifts the target
11318        // hiddens right by one; its draft step 0 feeds (id_last, TRUE hidden of the row id_last
11319        // was sampled from)). memra's historical convention paired SAME-ROW (x_p, h_p) in the fill
11320        // and seeded chain step 0 through an extra MTP pass on a duplicated token (the
11321        // pseudo-seed) — measured 27B p2 K=3 acceptance 0.569 vs 0.731, p3 0.445 vs 0.63+, and
11322        // the chain steps j>=1 were already predecessor-shaped, so ONLY the fill + step-0 seed
11323        // move. The fill shifts by one and the chain seeds from the predecessor's true hidden
11324        // DIRECTLY (vh_seed / vx[j-1]) — the pseudo pass disappears (one MTP-block pass saved
11325        // per round on top of the acceptance win). Draft-quality-only: exactness stays the
11326        // verify's job either way. (The legacy same-row pairing seam, MEMRA_SPEC_HSAME, and its
11327        // pseudo-seed passes were removed 2026-07-08 — predecessor pairing won by +16 acc pts;
11328        // the legacy round-local scratch, MEMRA_SPEC_KVLOCAL, went with it.)
11329        // REPLAY-FREE PARTIAL ACCEPT (default, 2026-07-03): partial rounds keep the verify's own
11330        // bit-identical committed-prefix state (KV truncate + recur rebuild from the VerifyCkpt)
11331        // and leave the bonus PENDING — no duplicate trunk pass (profiled ~0.54 extra full weight
11332        // reads/round at long ctx). MEMRA_SPEC_REPLAY=1 restores the legacy rollback+replay (A/B
11333        // + fallback seam).
11334        // Qwen35-MoE replay pin LIFTED (lane/draftcost-moe, 2026-08-20). The pin's stated
11335        // bar — the retained verify-state commit proven equivalent to sequential serving —
11336        // was waiting on this arch running the serving batched verify class, which the
11337        // t-parallel admission (this lane, increment 1) provided: the VerifyCkpt the
11338        // replay-free commit consumes is now produced by the SAME serving-class verify that
11339        // qualified dense qwen35 on 2026-08-15 (where the per-round duplicate replay
11340        // measured 69 -> 30 tok/s). Qualification receipts (run-spec K=1..8 both arms,
11341        // 8-prompt replay-vs-replay-free canary, long-prompt cell):
11342        // research/draftcost-moe-20260820/RECEIPTS.md. MEMRA_SPEC_REPLAY=1 stays the
11343        // rollback + A/B seam.
11344        let spec_replay = spec_replay_env_enabled();
11345        if constraint.is_some() && spec_replay {
11346            return Err(
11347                "constrained spec decode does not support MEMRA_SPEC_REPLAY=1 \
11348                        (legacy replay commits an unmasked bonus)"
11349                    .into(),
11350            );
11351        }
11352        // TRUE-HIDDEN REFRESH (default in persistent-draft-KV mode): every round overwrites the
11353        // committed positions' scratch entries from the verify's exact hiddens (mtp_kv_fill batch)
11354        // instead of keeping chain-approximate entries. MEMRA_SPEC_NOREFRESH=1 = legacy (A/B seam).
11355        let refresh = std::env::var("MEMRA_SPEC_NOREFRESH").is_err();
11356        if !refresh && !self.mtp_extra.is_empty() {
11357            return Err("multi-head MTP requires exact accepted-prefix refresh".into());
11358        }
11359
11360        // prime: BATCHED cache prime (prime_cache — the measured #1 e2e gap: tokenwise primed at
11361        // ~102/38 tok/s vs the engine's ~2000-5900 tok/s batched prefill). prime_cache returns the
11362        // full pre-output_norm hidden stack [T, n_embd], which IS prompt_h (the persistent-draft-KV
11363        // mtp_kv_fill input) — no per-token collection needed. Prompts below PRIME_MIN_T, and
11364        // MEMRA_PRIME_TOKENWISE=1, and frozen Hy3 CPU/GPU expert splits take the tokenwise
11365        // decode_step_h loop. The latter avoids transient GPU staging of the spilled expert bank.
11366        // EMPTY-SUFFIX CONTINUATION (serve bursts): a session turn with NO new tokens resumes
11367        // generation exactly where the last turn stopped — no prime at all. The stashed
11368        // `next_pred` plays prime_logits' role: it is the token produced from the logits after
11369        // committed.last() by the same rule this entry applies to a cold prime's last row —
11370        // an argmax when greedy, a `sample_boundary_token` draw when sampled (the burst tail,
11371        // or `spec_session_from_restored` for a converted prefix-cache hit, did the drawing
11372        // where the sampler and the session's Philox counters were live). `last_h` seeds the
11373        // predecessor pairing below. Fresh calls and non-empty suffixes take the normal path.
11374        let continuation = prompt.is_empty();
11375        if continuation {
11376            assert!(session_mode, "empty prompt requires a session");
11377            assert!(
11378                sess_tail
11379                    .as_ref()
11380                    .map_or(false, |(c, lh, np, _, _)| !c.is_empty()
11381                        && lh.is_some()
11382                        && (np.is_some() || carried_pending.is_some())),
11383                "empty-suffix continuation needs a primed session (committed + last_h + next_pred|pending)"
11384            );
11385        }
11386        let mut prime_logits;
11387        let mut prompt_h: Option<CudaSlice<f32>> = None;
11388        let t_prime = std::time::Instant::now();
11389        let batched_prime = !continuation
11390            && prompt.len() >= crate::hybrid_forward::PRIME_MIN_T
11391            && std::env::var("MEMRA_PRIME_TOKENWISE").is_err()
11392            && !e.frozen_cpu_experts_prefer_tokenwise_prime();
11393        let prime_split = prime_split.filter(|&split| split > 0 && split < prompt.len());
11394        if prime_split.is_some() && continuation {
11395            return Err("spec prime split requires a non-empty prime".into());
11396        }
11397        // STABLE-BOUNDARY TURN CHECKPOINT stop (lane/frspec-multiturn-cache, 2026-08-21):
11398        // the worker's `ckpt_at` request, ABSOLUTE -> prompt-relative. On WARM bursts
11399        // (base != 0, an affinity-rewound or pool-resumed session priming its own delta)
11400        // this is the only stop; on COLD bursts it usually coincides with `prime_split`
11401        // (both are the plain tier's stable pre-generation boundary). A boundary the prime
11402        // cannot honor (outside this prime's range) silently drops the capture — the
11403        // turn_ckpt convention: the next turn re-primes in full, never a wrong resume.
11404        let ckpt_rel = if continuation {
11405            None
11406        } else {
11407            ckpt_req
11408                .and_then(|abs| abs.checked_sub(base))
11409                .filter(|&r| r > 0 && r < prompt.len())
11410        };
11411        // Prime stops, ordered: each is a boundary the prime halts at so the in-place GDN
11412        // conv/ssm state can be snapshotted there (the only moment it exists). One stop =
11413        // the legacy single-split program, byte-for-byte.
11414        let mut stops: Vec<usize> = Vec::new();
11415        for b in [prime_split, ckpt_rel].into_iter().flatten() {
11416            if !stops.contains(&b) {
11417                stops.push(b);
11418            }
11419        }
11420        stops.sort_unstable();
11421        // Captured at the ckpt stop, installed into the session slot post-prime (replacing
11422        // the legacy prompt-end capture). Some(None) = capture attempted and failed -> the
11423        // slot is cleared (a stale checkpoint would rewind to the WRONG boundary).
11424        let mut ckpt_early: Option<Option<SpecCheckpoint>> = None;
11425        if continuation {
11426            prime_logits = Vec::new();
11427        } else if !stops.is_empty() {
11428            if let Some(&first) = stops.first() {
11429                if prime_split == Some(first) && first < crate::hybrid_forward::PRIME_MIN_T {
11430                    return Err(format!(
11431                        "spec prime split {first} is below PRIME_MIN_T {}",
11432                        crate::hybrid_forward::PRIME_MIN_T,
11433                    )
11434                    .into());
11435                }
11436            }
11437            // Mirror the plain worker's boundary stops exactly. Each segment is a
11438            // request-level prime (`queued_after` keeps Step35 arm selection independent of
11439            // the stops — tick-seg law); a segment below PRIME_MIN_T (and the final tail
11440            // under MEMRA_PRIME_TOKENWISE) takes the same eager tokenwise continuation as
11441            // prefill_tick. Retain every hidden row so the draft scratch fill remains one
11442            // coherent prompt.
11443            let mut h_all = e.uninit(prompt.len() * n_embd)?;
11444            prime_logits = Vec::new();
11445            let mut prev = 0usize;
11446            for seg_end in stops.iter().copied().chain(std::iter::once(prompt.len())) {
11447                if seg_end <= prev {
11448                    continue;
11449                }
11450                let seg = &prompt[prev..seg_end];
11451                let is_final = seg_end == prompt.len();
11452                let batched_seg = seg.len() >= crate::hybrid_forward::PRIME_MIN_T
11453                    && (!is_final
11454                        || (std::env::var("MEMRA_PRIME_TOKENWISE").is_err()
11455                            && !e.frozen_cpu_experts_prefer_tokenwise_prime()));
11456                if batched_seg {
11457                    let (l, _, h_seg) =
11458                        self.prime_cache(e, seg, &mut *cache, prompt.len() - seg_end)?;
11459                    e.copy_into(&mut h_all, prev * n_embd, &h_seg, seg.len() * n_embd)?;
11460                    prime_logits = l;
11461                } else {
11462                    for (i, &tok) in seg.iter().enumerate() {
11463                        let (l, h) = self.decode_step_h(e, tok, &mut *cache)?;
11464                        e.copy_into(&mut h_all, (prev + i) * n_embd, &h, n_embd)?;
11465                        prime_logits = l;
11466                    }
11467                }
11468                prev = seg_end;
11469                if is_final {
11470                    break;
11471                }
11472                debug_assert_eq!(cache.pos, base + seg_end, "prime stop landed off boundary");
11473                // PREFIX-CACHE BOUNDARY CAPTURE (lane/spec-prefix-cache): the GDN conv/ssm
11474                // states are about to be advanced in place by the next segment, so this is
11475                // the ONLY moment the boundary's recurrent state exists. Capture iff the
11476                // worker requested exactly this stop (cold sessions only — `capture_at` is
11477                // never armed warm). A failed snapshot is silent (turn_ckpt convention) —
11478                // publication is an optimization, never a correctness dependency.
11479                if base == 0 {
11480                    if let Some((requested, slot)) = sess_capture.as_mut() {
11481                        // Publish at the requested miss-LCP stop (the shared-prefix class)
11482                        // AND at the stable-boundary stop (the next-turn re-render class,
11483                        // lane/frspec-multiturn-cache) — the same boundary set the plain
11484                        // prefill tick learns. Without the second entry, the turn after a
11485                        // cold re-park could only hit the OLDER lcp entry (the measured
11486                        // one-turn transient: t3 restored 607 of 24122 while the plain arm
11487                        // rewound to 15222). Dedupe is the worker sweep's has_key.
11488                        if *requested == Some(seg_end) || ckpt_rel == Some(seg_end) {
11489                            if let Ok(snap) = cache.snapshot(e) {
11490                                slot.push(SpecBoundaryCapture {
11491                                    snap,
11492                                    pos: seg_end,
11493                                    logits: prime_logits.clone(),
11494                                    // rows [0..seg_end) of h_all are primed — the following
11495                                    // segments append, never overwrite.
11496                                    last_h: capture_boundary_hidden(e, &h_all, seg_end, n_embd),
11497                                });
11498                            }
11499                        }
11500                    }
11501                }
11502                // SESSION-AFFINITY TURN CHECKPOINT at the STABLE boundary (see `ckpt_at`):
11503                // same snapshot mechanics, installed post-prime in place of the prompt-end
11504                // capture the re-render class always diverged below.
11505                if ckpt_rel == Some(seg_end) {
11506                    let anchor: Result<CudaSlice<f32>, Box<dyn std::error::Error>> =
11507                        e.uninit(n_embd).and_then(|mut a| {
11508                            e.copy_view_into(
11509                                &mut a,
11510                                0,
11511                                &h_all.slice((seg_end - 1) * n_embd..seg_end * n_embd),
11512                                n_embd,
11513                            )?;
11514                            Ok(a)
11515                        });
11516                    ckpt_early = Some(match (cache.snapshot(e), anchor) {
11517                        (Ok(snap), Ok(last_h)) => Some(SpecCheckpoint {
11518                            snap,
11519                            pos: base + seg_end,
11520                            last_h,
11521                        }),
11522                        _ => None,
11523                    });
11524                }
11525            }
11526            if std::env::var("MEMRA_SPEC_STATS").as_deref() == Ok("1") {
11527                eprintln!(
11528                    "[spec-prime] stops={stops:?} tail={}",
11529                    prompt.len() - stops.last().copied().unwrap_or(0)
11530                );
11531            }
11532            prompt_h = Some(h_all);
11533        } else if batched_prime {
11534            let (l, _h_seed, hiddens) = self.prime_cache(e, prompt, &mut *cache, 0)?;
11535            prime_logits = l;
11536            prompt_h = Some(hiddens);
11537        } else {
11538            prime_logits = Vec::new();
11539            prompt_h = Some(e.uninit(prompt.len() * n_embd)?);
11540            for (i, &tok) in prompt.iter().enumerate() {
11541                let (l, h) = self.spec_target_step_h(e, tok, &mut *cache)?;
11542                if let Some(ph) = prompt_h.as_mut() {
11543                    e.copy_into(ph, i * n_embd, &h, n_embd)?;
11544                }
11545                prime_logits = l;
11546            }
11547        }
11548        e.stream().synchronize()?;
11549        // PREFIX-CACHE SEED CAPTURE (lane/spec-prefix-cache): boundary == prompt end (the seed
11550        // case — no shared-prefix split, publish the whole prompt). The prime just finished, so
11551        // cache.pos == base + prompt.len() and the recurrent state IS the boundary state;
11552        // prime_logits are the boundary logits. Cold sessions only (base == 0) — same law as
11553        // prime_split. The mid-prompt capture above already consumed the request if it matched.
11554        if !continuation && base == 0 {
11555            if let Some((requested, slot)) = sess_capture.as_mut() {
11556                if *requested == Some(prompt.len()) && slot.is_empty() {
11557                    debug_assert_eq!(cache.pos, prompt.len(), "seed capture off prompt end");
11558                    if let Ok(snap) = cache.snapshot(e) {
11559                        slot.push(SpecBoundaryCapture {
11560                            snap,
11561                            pos: prompt.len(),
11562                            logits: prime_logits.clone(),
11563                            last_h: prompt_h
11564                                .as_ref()
11565                                .map(|ph| capture_boundary_hidden(e, ph, prompt.len(), n_embd))
11566                                .unwrap_or_default(),
11567                        });
11568                    }
11569                }
11570            }
11571        }
11572        // Harness timing contract (see crate::PRIME_NANOS): gen-only throughput without the
11573        // prime-subtraction hack.
11574        crate::PRIME_NANOS.store(
11575            t_prime.elapsed().as_nanos() as u64,
11576            std::sync::atomic::Ordering::Relaxed,
11577        );
11578
11579        let (embd_qt, embd_rb) = self.embd.qt_and_row_bytes(n_embd);
11580        // Resident table is fastest when it fits. Large spill deployments can preserve that HBM
11581        // for expert-cache slots and gather only the exact rows needed by MTP/verify from host.
11582        let host_embd = spec_host_embd();
11583        let embd_gpu = if host_embd {
11584            None
11585        } else {
11586            Some(
11587                self.embd_gpu
11588                    .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload")),
11589            )
11590        };
11591        let embd_dev = embd_gpu.map(|g| (g, embd_qt, embd_rb));
11592        if host_embd {
11593            eprintln!(
11594                "[spec] host-row embedding: {} bytes kept off HBM",
11595                self.embd.raw.len()
11596            );
11597        }
11598        let mut out: Vec<u32> = Vec::with_capacity(max_new);
11599        let mut total_drafted = 0usize;
11600        let mut total_accepted = 0usize;
11601
11602        // --- SAMPLER FIRST (lane/sampled-spec-quality, 2026-08-19) ---
11603        // The sampler config, the session's Philox counters and the penalty window are parsed
11604        // HERE, above the boundary-token selection, because the boundary token must be drawn
11605        // from the sampler the request asked for. Pre-lane this block sat ~50 lines BELOW the
11606        // selection, which is the whole mechanical reason the boundary token was an argmax:
11607        // the sampler state was not in scope yet. Nothing here depends on the round loop, so
11608        // moving it up is a pure reordering for greedy (`sampled == false` ⇒ every branch
11609        // below takes the argmax path it always took).
11610        // --- SAMPLED SPEC (MEMRA_SPEC_TEMP>0, research/sampled-spec-impl-map.md): rejection-
11611        // sampling verify (Leviathan/Chen) — accept draft x at u < p(x)/q(x), resample from
11612        // norm(max(0,p-q)) on reject, bonus sampled from p on full accept. Counter-based Philox
11613        // everywhere (seed, event) -> reproducible. temp==0/unset = the greedy path, untouched.
11614        let sp = sampling.unwrap_or_else(|| SpecSampling {
11615            temp: std::env::var("MEMRA_SPEC_TEMP")
11616                .ok()
11617                .and_then(|v| v.parse().ok())
11618                .unwrap_or(0.0),
11619            seed: std::env::var("MEMRA_SEED")
11620                .ok()
11621                .and_then(|v| v.parse().ok())
11622                .unwrap_or(42),
11623            top_k: std::env::var("MEMRA_TOP_K")
11624                .ok()
11625                .and_then(|v| v.parse().ok())
11626                .unwrap_or(0),
11627            top_p: std::env::var("MEMRA_TOP_P")
11628                .ok()
11629                .and_then(|v| v.parse().ok())
11630                .unwrap_or(1.0),
11631            min_p: std::env::var("MEMRA_MIN_P")
11632                .ok()
11633                .and_then(|v| v.parse().ok())
11634                .unwrap_or(0.0),
11635            penalty_last_n: std::env::var("MEMRA_PENALTY_LAST_N")
11636                .ok()
11637                .and_then(|v| v.parse().ok())
11638                .unwrap_or(0),
11639            penalty_repeat: std::env::var("MEMRA_PENALTY_REPEAT")
11640                .ok()
11641                .and_then(|v| v.parse().ok())
11642                .unwrap_or(1.0),
11643            penalty_freq: std::env::var("MEMRA_PENALTY_FREQ")
11644                .ok()
11645                .and_then(|v| v.parse().ok())
11646                .unwrap_or(0.0),
11647            penalty_present: std::env::var("MEMRA_PENALTY_PRESENT")
11648                .ok()
11649                .and_then(|v| v.parse().ok())
11650                .unwrap_or(0.0),
11651        });
11652        let (sp_temp, sp_seed) = (sp.temp, sp.seed);
11653        let sampled = sp_temp > 0.0;
11654        // Counters resume from the session (burst continuity: randomness must never repeat
11655        // across generate_spec_session calls); one-shot callers start at (0,0). Read through
11656        // sess_tail — `sess` was take()n into it above, so sess.as_ref() here is always None.
11657        let mut sctr: u32 = sess_tail.as_ref().map(|(_, _, _, s, _)| **s).unwrap_or(0);
11658        let mut uctr: u32 = sess_tail.as_ref().map(|(_, _, _, _, u)| **u).unwrap_or(0);
11659        // Penalties (v2.1): applied to COPIES of q rows and p columns symmetrically (exactness
11660        // for the penalized+filtered target). History = generated tokens, host-tracked window.
11661        let pen_on = sampled
11662            && sp.penalty_last_n > 0
11663            && (sp.penalty_repeat != 1.0 || sp.penalty_freq != 0.0 || sp.penalty_present != 0.0);
11664        // SESSION-SPANNING PENALTY WINDOW (Item 2). Pre-lane this was
11665        // `prompt.iter().rev().take(64).rev()` — the BURST's suffix slice — so a continuation
11666        // burst (the majority of a stream's tokens, and ALL of a converted cache hit's) started
11667        // with an EMPTY penalty history and the client's repetition/frequency/presence penalties
11668        // silently reset at every burst boundary. The window now spans `committed ++ prompt`,
11669        // which is what the API contract says and what the plain sampler's own `history` does.
11670        // Byte-identical to the pre-lane seed for a cold turn-1 burst at the default window.
11671        let mut pen_hist: Vec<u32> = if pen_on {
11672            let sess_hist: &[u32] = if spec_pen_session_on() {
11673                sess_tail
11674                    .as_ref()
11675                    .map(|(c, ..)| c.as_slice())
11676                    .unwrap_or(&[])
11677            } else {
11678                &[] // MEMRA_SPEC_PEN_SESSION=0: pre-lane burst-local window
11679            };
11680            pen_window_seed(sess_hist, prompt, sp.penalty_last_n)
11681        } else {
11682            Vec::new()
11683        };
11684        // First generated token = the BOUNDARY token: greedy takes the argmax of the prompt's
11685        // last logits (== greedy's first token, byte-contract); SAMPLED draws it from the
11686        // request's own filtered/penalized target through the session's Philox stream
11687        // (`sample_boundary_token`, lane/sampled-spec-quality Item 1 — pre-lane this was an
11688        // argmax in both regimes, so ~1 token per burst of a sampled stream was greedy).
11689        // Emit it, then FEED it to establish the loop invariant below.
11690        // PENDING-CARRY: the carried bonus was already emitted by the LAST burst — it becomes
11691        // last_token WITHOUT re-emission, and round 0 consumes it as pending (no init feed).
11692        // CONSTRAINED entry rules: the first emitted token is the MASKED argmax of the
11693        // prompt's last logits (plain constrained-greedy identity); a continuation without
11694        // a carried pending would emit an UNMASKED stashed next_pred — refused loudly (the
11695        // worker never resumes constrained sessions from the pool, so this cannot fire).
11696        if let Some(c) = constraint.as_deref_mut() {
11697            if continuation && carried_pending.is_none() {
11698                return Err("constrained spec continuation requires a carried pending \
11699                            (pool resume is unconstrained-only)"
11700                    .into());
11701            }
11702            if !continuation {
11703                c.mask_logits(&mut prime_logits)
11704                    .map_err(|e2| format!("constraint: {e2}"))?;
11705            }
11706        }
11707        let mut last_token = if let Some(b) = carried_pending {
11708            b
11709        } else if continuation {
11710            // A continuation's boundary token was DRAWN by the burst that stashed it (the
11711            // session tail below), or by `spec_session_from_restored` for a converted
11712            // prefix-cache hit — in both cases from the correct logits row with this same
11713            // session's Philox stream, which is why it can be consumed here as-is.
11714            sess_tail.as_ref().unwrap().2.unwrap()
11715        } else if sampled && constraint.is_none() && spec_sampled_boundary_on() {
11716            sample_boundary_token(e, &prime_logits, &sp, &pen_hist, &mut sctr, "cold-prime")?
11717        } else {
11718            // greedy (byte contract), the rollback door, or constrained (masked-argmax
11719            // identity — the worker routes sampled+constrained to the plain path, and this
11720            // function refuses the combination outright above).
11721            argmax(&prime_logits) as u32
11722        };
11723        if pen_on {
11724            // The boundary token is a GENERATED token: the plain sampler `accept()`s every
11725            // emitted token into its penalty history, and pre-lane the burst's first token
11726            // was invisible to penalties forever (never pushed, and never in `committed`
11727            // until this burst's tail). Covers the carry/continuation seeds too — neither is
11728            // in `committed` yet.
11729            pen_hist.push(last_token);
11730        }
11731        if carried_pending.is_none() {
11732            out.push(last_token);
11733            // grammar advances with every emitted token (carried pendings were consumed
11734            // by the burst that emitted them).
11735            if let Some(c) = constraint.as_deref_mut() {
11736                c.consume(last_token)
11737                    .map_err(|e2| format!("constraint: {e2}"))?;
11738            }
11739        }
11740        if continuation {
11741            // draft-KV invariant: entries [0..base) are the session's exact fills; truncate any
11742            // overhang so the chain's first append lands at slot base (== committed.len()).
11743            scratch.set_len(e, base)?;
11744        }
11745        // sse-cadence: hand the caller every not-yet-flushed token (disjoint in-order slices
11746        // concatenating to the full `out`). Called after the prime's first token and after each
11747        // round commit — emission timing only, token bytes untouched. The slice may be EMPTY
11748        // (poll-only boundary: zero-round folds commit nothing new); returns the caller's
11749        // continue-verdict (admission yield, 2026-08-06) — false ends the burst at this round.
11750        fn flush_commit(
11751            cb: &mut Option<&mut dyn FnMut(&[u32]) -> bool>,
11752            out: &[u32],
11753            flushed: &mut usize,
11754        ) -> bool {
11755            if let Some(f) = cb.as_mut() {
11756                let keep = f(&out[*flushed..]);
11757                *flushed = out.len();
11758                keep
11759            } else {
11760                true
11761            }
11762        }
11763        keep_going = flush_commit(&mut on_commit, &out, &mut flushed);
11764        // INVARIANT at loop top: `last_token` is the most-recently-committed/emitted token, its
11765        // KV+recur state IS in `cache` (cache.pos = position right AFTER last_token), `last_pred`
11766        // is the greedy ARGMAX of the logits that predict the token FOLLOWING last_token, and
11767        // `h_seed` = last_token's pre-output_norm hidden. Establish it by feeding last_token once
11768        // (mirrors plain greedy). DEVICE-ARGMAX lever: the accept walk only ever consumes the
11769        // argmax of those logits — never the full vector — so a host u32 replaces the Vec<f32>.
11770        // Trimmed heads: q lives on the trimmed vocab; accept gathers use the TRIMMED index and
11771        // the residual scatters q into target-id space (q=-inf off-trim — the head cannot propose
11772        // those, so their residual mass is p(x), correct by construction).
11773        let d2t_dev: Option<CudaSlice<u32>> = if sampled || crate::spec::spec_stream() {
11774            match &mtp.d2t {
11775                Some(map) => Some(e.htod_u32_v(map)?),
11776                None => None,
11777            }
11778        } else {
11779            None
11780        };
11781        let mut q_full_buf: Option<CudaSlice<f32>> = None;
11782        // host Philox4x32-10 accept-test uniforms: module fn `host_u01` (shared with the
11783        // dspark sampled-admission walk); byte-identical to the closure it replaces.
11784        let mut draft_logits: Vec<CudaSlice<f32>> = Vec::new(); // retained head logits (q), per slot
11785        let mut draft_stats: Vec<(f32, f32, f32)> = Vec::new(); // (row_max, th_e, z_e) per slot
11786        let mut perturb_buf: Option<CudaSlice<f32>> = None; // gumbel scratch (max(n_vocab,d_vocab))
11787        let mut sample_tok = e.alloc_u32_zeroed(1)?; // residual/bonus sample out
11788        let mut col_buf: Option<CudaSlice<f32>> = None; // materialized verify column
11789        let mut pen_hist_d: Option<CudaSlice<u32>> = None;
11790        let mut pcol_buf: Option<CudaSlice<f32>> = None; // penalized p-column scratch
11791        // MEMRA_SPEC_SETUP_TRACE=1 (diagnostics): per-call wall decomposition of the burst
11792        // SETUP + TAIL segments (the round loop's internals are MEMRA_SPEC_PHASE's job) —
11793        // built to pin the serve per-burst fixed cost (research/spec-serving-20260801).
11794        let setup_trace = std::env::var("MEMRA_SPEC_SETUP_TRACE").as_deref() == Ok("1");
11795        let t_ent = std::time::Instant::now();
11796
11797        // SESSION-AFFINITY TURN CHECKPOINT (lane/session-affinity, 2026-08-05): capture the
11798        // PROMPT-END boundary state so a LATER turn can rewind here and re-prime only its own
11799        // delta instead of the whole conversation. See `SpecCheckpoint` for why this boundary is
11800        // the one that matters (a history-rewriting client mutates what the session GENERATED,
11801        // so the next turn's prompt agrees with this one up to exactly here).
11802        //
11803        // WHERE — AND WHY THIS EXACT LINE. Right after the trunk prime, BEFORE the init feed
11804        // (`decode_step_h(last_token)`) and before round 0: the last instant at which the caches
11805        // hold exactly `base + prompt.len()` rows and nothing generated.
11806        //
11807        // This was WRONG in the first cut of this lane: the capture sat after the draft-KV fill,
11808        // which is also after the init feed, so `cache.pos` was `base + prompt.len() + 1` — the
11809        // boundary included the FIRST GENERATED TOKEN. That token is the first thing inside the
11810        // `<think>` block the client strips, so every later turn's diff diverged exactly one
11811        // token below the checkpoint and affinity declined 100% of the time. Measured on the
11812        // owner regime: "history diverged at 12233 of checkpoint 12234". The off-by-one made the
11813        // whole mechanism inert while looking, from the outside, like a working
11814        // correctness-declines-safely path — hence the decline log carries the offsets.
11815        //
11816        // The full-attn planes are `len`-truncatable so the snapshot copies only the GDN conv/ssm
11817        // state (the reason a spec session could not rewind before). The draft scratch needs no
11818        // copy: rows below the boundary are rewritten by the next turn's own fill.
11819        //
11820        // WHEN: non-empty prime only. An empty-suffix continuation burst adds no prompt boundary
11821        // (its "prompt end" IS the previous checkpoint's, already held), so it keeps the existing
11822        // checkpoint rather than replacing it with a strictly worse one.
11823        //
11824        // FAILURE IS SILENT BY DESIGN: on a VRAM-tight rig the snapshot alloc can fail. That
11825        // costs the NEXT turn its rewind (it re-primes fully, today's behavior) and must never
11826        // fail the burst that is already running — so the error is swallowed, loud only under
11827        // MEMRA_DEBUG_SPEC.
11828        //
11829        // STABLE-BOUNDARY OVERRIDE (lane/frspec-multiturn-cache, 2026-08-21): the prompt-end
11830        // posture above was DISPROVED for the think-posture template class — the prompt's own
11831        // tail is the live generation header (`<|im_start|>assistant\n<think>\n`) that the
11832        // next turn's re-render replaces, so the diff diverged a couple tokens BELOW the
11833        // checkpoint and affinity declined 100% of multi-turn agent traffic (the same class
11834        // the plain tier fixed on 2026-08-09 via `plain_checkpoint_boundary`; the port to the
11835        // spec tier is this lane). When the worker armed `ckpt_at`, the capture happened at
11836        // that stop inside the prime above (`ckpt_early`) and is installed here instead;
11837        // capture-attempted-but-failed clears the slot exactly like the legacy arm.
11838        if let Some(slot) = sess_ckpt_slot {
11839            if let Some(early) = ckpt_early {
11840                if early.is_none() && std::env::var("MEMRA_DEBUG_SPEC").is_ok() {
11841                    eprintln!(
11842                        "[spec] stable-boundary turn checkpoint skipped; \
11843                               next turn re-primes in full"
11844                    );
11845                }
11846                *slot = early;
11847            } else if !continuation {
11848                let pos = cache.pos;
11849                debug_assert_eq!(
11850                    pos,
11851                    base + prompt.len(),
11852                    "turn checkpoint must sit at the prompt end, before the init feed"
11853                );
11854                let anchor: Result<CudaSlice<f32>, Box<dyn std::error::Error>> =
11855                    if let Some(ph) = &prompt_h {
11856                        // hidden of the LAST primed row = the predecessor anchor at this
11857                        // boundary (exactly what a fresh prime of committed[..pos] leaves in
11858                        // last_h, and what the next prime's fill reads for its first row).
11859                        let np = prompt.len();
11860                        e.uninit(n_embd).and_then(|mut a| {
11861                            e.copy_view_into(
11862                                &mut a,
11863                                0,
11864                                &ph.slice((np - 1) * n_embd..np * n_embd),
11865                                n_embd,
11866                            )?;
11867                            Ok(a)
11868                        })
11869                    } else {
11870                        Err("no prompt hiddens".into())
11871                    };
11872                match (cache.snapshot(e), anchor) {
11873                    (Ok(snap), Ok(last_h)) => {
11874                        *slot = Some(SpecCheckpoint { snap, pos, last_h });
11875                    }
11876                    (s, a) => {
11877                        *slot = None; // a stale checkpoint would rewind to the WRONG boundary
11878                        if std::env::var("MEMRA_DEBUG_SPEC").is_ok() {
11879                            let err = s
11880                                .err()
11881                                .map(|e| e.to_string())
11882                                .or_else(|| a.err().map(|e| e.to_string()))
11883                                .unwrap_or_default();
11884                            eprintln!(
11885                                "[spec] turn checkpoint skipped ({err}); \
11886                                       next turn re-primes in full"
11887                            );
11888                        }
11889                    }
11890                }
11891            }
11892        }
11893        // INIT FEED — skipped on a pending carry: last_token (the carried bonus) is NOT in the
11894        // caches and must NOT be fed solo; round 0's batched verify commits it as col 0. Its
11895        // seed/anchor hidden is the carried last_h (copied below); last_pred is dead in the
11896        // pending path (t_pred reads verify col 0 — the accept walk overwrites it).
11897        let mut last_pred = 0u32;
11898        let mut last_col_logits: Option<CudaSlice<f32>> = None;
11899        // CONSTRAINED: the init feed's logits back the (n_acc==0, base==0) masked-argmax
11900        // recompute in the grammar-truncation walk — retained host-side, round 0 only.
11901        let mut init_logits_host: Option<Vec<f32>> = None;
11902        let h_seed0: CudaSlice<f32> = if carried_pending.is_none() {
11903            let (init_logits, h) = self.spec_target_step_h(e, last_token, &mut *cache)?;
11904            last_pred = argmax(&init_logits) as u32;
11905            if constraint.is_some() {
11906                init_logits_host = Some(init_logits.clone());
11907            }
11908            // sampled mode: p-distribution after last_token, for the j==0/base==0 accept test.
11909            if sampled {
11910                last_col_logits = Some(e.htod(&init_logits)?);
11911            }
11912            h
11913        } else {
11914            // predecessor-row anchor: hidden of the last COMMITTED row (the carry contract).
11915            let lh = sess_tail
11916                .as_ref()
11917                .unwrap()
11918                .1
11919                .as_ref()
11920                .expect("pending carry requires last_h");
11921            e.clone_dtod(lh)?
11922        };
11923        let t_init = t_ent.elapsed();
11924        let mut last_col_stats: Option<(f32, f32, f32)> = None;
11925        // PERSISTENT h_seed buffer (allocated BEFORE any graph capture so no captured scratch can
11926        // alias it): every path that updates the round seed copies INTO it — no per-round allocs,
11927        // stable pointer for the graph-draft round-start copy.
11928        let mut h_seed_buf = e.clone_dtod(&h_seed0)?;
11929        // Predecessor-pairing trackers: `fill_prev` = trunk hidden AT the last COMMITTED row (the
11930        // predecessor of the next verify's col 0 — the reference's carried pending-h analogue;
11931        // also the predecessor-row hidden for the round-0 legacy-replay seed). At round 0 that
11932        // row is last_token's own (h_seed0). The chain step-0 seed under the pairing default =
11933        // hidden of the row BEFORE last_token = the prompt's last row at round 0 (h_seed_buf
11934        // overwritten below).
11935        let mut fill_prev = e.clone_dtod(&h_seed0)?;
11936        {
11937            if let Some(ph) = &prompt_h {
11938                let np = prompt.len();
11939                e.copy_view_into(
11940                    &mut h_seed_buf,
11941                    0,
11942                    &ph.slice((np - 1) * n_embd..np * n_embd),
11943                    n_embd,
11944                )?;
11945            } else if continuation {
11946                if let Some((_, lh, _, _, _)) = sess_tail.as_ref() {
11947                    if let Some(lh) = lh.as_ref() {
11948                        e.copy_into(&mut h_seed_buf, 0, lh, n_embd)?;
11949                    }
11950                }
11951            }
11952        }
11953        // Persistent device prediction slots for the accept walk (max k+1 verify columns).
11954        let mut preds_d = e.alloc_u32_zeroed(k + 2)?;
11955
11956        let debug_spec = std::env::var("MEMRA_DEBUG_SPEC").is_ok();
11957        let fork_mode = OptiForkGateMode::configured();
11958        // MEMRA_SPEC_STATS=1: per-slot accept histogram + draft-length histogram, printed once at
11959        // the end. Metric normalization vs the reference engine: BOTH engines count
11960        // accepted/drafted where the chain stopped at p-min and the sub-threshold token is
11961        // discarded uncounted — per-slot decay + chain-length mix are the extra dimensions.
11962        let spec_stats = std::env::var("MEMRA_SPEC_STATS").is_ok();
11963        let mut st_drafted = vec![0usize; k];
11964        let mut st_accepted = vec![0usize; k];
11965        let mut st_len_hist = vec![0usize; k + 1];
11966        let mut st_full = 0usize;
11967        // P-MIN CONFIDENCE GATE (MEMRA_SPEC_PMIN, the serve script's --spec-draft-p-min mechanism):
11968        // stop the draft chain early when the head's softmax confidence in its own pick drops
11969        // below p_min. Hoisted above the loop: the graph capture bakes the prob kernels iff on.
11970        static PMIN: std::sync::OnceLock<f32> = std::sync::OnceLock::new();
11971        let p_min = *PMIN.get_or_init(|| {
11972            std::env::var("MEMRA_SPEC_PMIN")
11973                .ok()
11974                .and_then(|v| v.parse().ok())
11975                .unwrap_or(0.0)
11976        });
11977        // ZERO-DRAFT ROUNDS (MEMRA_SPEC_PMIN0=1, vendored from llama.cpp's draft gating): let the
11978        // p-min gate apply at j==0 too, so a low-confidence round drafts NOTHING and the verify
11979        // batch is just the pending bonus (m=1 = a plain decode step). llama's 35B win rides
11980        // exactly this — draft acceptance 76% at mean len 2.5 because unpredictable stretches
11981        // never pay draft+verify overhead. Only legal when a pending bonus exists (an empty
11982        // verify batch is not); the j==0 exemption stays for pending-less rounds.
11983        let pmin0 = std::env::var("MEMRA_SPEC_PMIN0")
11984            .map(|v| v == "1")
11985            .unwrap_or(false);
11986
11987        // --- GRAPH DRAFT setup: persistent I/O buffers + ONE capture (2 warmups inside). The
11988        // warmups mutate scratch len_d / pos / tok / seed — all reset at every round start, so the
11989        // only restore needed is the scratch counter. Capture failure (e.g. a non-capturable
11990        // cuBLAS path in an exotic head) falls back to the eager draft chain.
11991        // PER-SESSION PERSISTENCE (2026-08-01): session calls reuse the DraftGraphCtx parked on
11992        // the SpecSession — the capture (2 warmup head forwards + instantiate) ran ONCE at the
11993        // session's first burst, not per burst (measured ~16ms/burst fixed cost on H100 q27,
11994        // research/spec-serving-20260801). Reuse is pointer-exact: the graph bakes the session's
11995        // own scratch KV (never realloc'd), the model's resident embedding, the OnceLock p_min,
11996        // and the g_* buffers carried in the ctx — replay dispatch is identical to a fresh
11997        // capture, so draft tokens are bit-identical (drafts never decide exactness anyway; the
11998        // verify arbitrates). Single-shot calls (sess=None) build a fresh ctx and drop it.
11999        let mut dctx: DraftGraphCtx = match sess_draft_slot.as_mut().and_then(|s| s.take()) {
12000            Some(c) => c,
12001            None => DraftGraphCtx::new(e, n_embd, if sampled { d_vocab } else { 1 })?,
12002        };
12003        // FAIL-SAFE (step-OOM park replay): pre-mark both fallback flags so no capture arm
12004        // below can fire — LOUD once per replayed session through the standard WARN line.
12005        if sess_capture_disabled {
12006            let reason =
12007                "session replayed after a step-OOM park; draft capture disabled (fail-safe)";
12008            let flip = dctx.failed.mark_greedy(reason);
12009            let flip_s = dctx.failed.mark_sampled(reason);
12010            if let Some(line) = flip.or(flip_s) {
12011                eprintln!("{line}");
12012            }
12013        }
12014        // A session that ran greedy bursts first sized g_q/g_perturb at 1; a sampled resume
12015        // needs d_vocab. Realloc is legal exactly while graph_s is None (nothing baked them).
12016        if sampled && dctx.g_q.len() < d_vocab {
12017            dctx.g_q = e.zeros(d_vocab)?;
12018            dctx.g_perturb = e.zeros(d_vocab)?;
12019        }
12020        // DRAFT-SIDE GRAMMAR MASK (lane/draft-mask, 2026-08-04): the drafter samples the
12021        // grammar's legal set, so proposals are legal BY CONSTRUCTION and the verify-side
12022        // truncation (the correctness backstop) stops cutting every tight-schema round.
12023        // The mask is one node inside the captured draft chain — presence is a CAPTURE-TIME
12024        // shape, so a parked graph of the other shape is dropped and recaptured.
12025        let dmask_on = constraint
12026            .as_deref()
12027            .is_some_and(|c| c.draft_mask_enabled());
12028        let dmask_words = if dmask_on { d_vocab.div_ceil(32) } else { 0 };
12029        if dmask_on && dctx.g_dmask.len() < dmask_words {
12030            dctx.g_dmask = e.alloc_u32_zeroed(dmask_words)?;
12031            dctx.graph = None; // the old capture baked the old (or no) mask pointer
12032            dctx.chain = None; // chain last-row graphs bake the same pointer
12033            dctx.failed.clear_greedy();
12034            dctx.keeper.clear();
12035        }
12036        if (dctx.graph.is_some() || dctx.chain.is_some()) && dctx.graph_masked != dmask_on {
12037            dctx.graph = None;
12038            dctx.chain = None;
12039            dctx.failed.clear_greedy();
12040            dctx.keeper.clear();
12041        }
12042        // MULTI-HEAD CHAIN mode (mtp_extra non-empty — step37's 3-head shipping shape): the
12043        // step-modulo prefix-replay chain captures PER-HEAD single-row graphs
12044        // (`DraftChainGraphs`) instead of the one self-feeding graph below; the single-head
12045        // capture arms are untouched and unreachable in this mode (the launch arms branch the
12046        // same way). This removes the historical `mtp_extra.is_empty()` capture exclusion —
12047        // and with it the silent no-attempt hole: a chain capture that FAILS now trips the
12048        // same LOUD draft-graph WARN as a single-head failure.
12049        let chain_mode = !self.mtp_extra.is_empty();
12050        // ---- PRE-CAPTURE VRAM RESERVE CHECK + PER-SESSION DRAFT-STATE MEASUREMENT ----
12051        // (lane/step37-vram-admission-20260830). `cap_eff0` opens the measurement bracket:
12052        // when any capture succeeds in THIS call, the effective-free delta across the whole
12053        // capture section is recorded as the model's per-session draft-state high-water
12054        // (admission charges it per spec-capable session — this state was charged at ZERO
12055        // before the lane). The reserve check runs BEFORE any capture arm can allocate: a
12056        // refused capture trips the same LOUD once-per-flip WARN class as a failed one, but
12057        // with the card's headroom still intact (the owner's single-session OOM was a capture
12058        // attempt walking the card to the edge and stranding the eager fallback at 5 MiB free).
12059        let cap_eff0 = e
12060            .ctx()
12061            .mem_get_info()
12062            .ok()
12063            .map(|(f, _)| f.saturating_add(e.pool_cached_bytes()));
12064        // Peak instrument for the same bracket: the CAPTURE-TIME peak (warmup transients +
12065        // instantiate scratch, alive together) dwarfs the parked delta — measured on the
12066        // owner shape: a capture whose PARKED state reads ~2.6GB walked a ~7GB-free card to
12067        // OOM mid-capture. Reset the pool watermark here; read it at bracket end.
12068        let _ = e.pool_high_water_reset();
12069        let cap_used0 = e.pool_reserved_used().1;
12070        let mut captured_now = false;
12071        let mut capture_oom_entry_eff: Option<usize> = None;
12072        let capture_need = {
12073            let observed = self.draft_session_admission_bytes();
12074            if observed > 0 {
12075                observed
12076            } else {
12077                draft_capture_bootstrap_estimate(
12078                    if chain_mode { self.mtp_head_count() } else { 1 },
12079                    k,
12080                    d_vocab,
12081                    n_embd,
12082                )
12083            }
12084        };
12085        if spec_capture_gate_on()
12086            && graph_draft
12087            && !sampled
12088            && !dctx.failed.greedy_failed()
12089            && ((chain_mode && dctx.chain.is_none() && mtp_chain_graph_on())
12090                || (!chain_mode && dctx.graph.is_none()))
12091            && let Some(reason) = capture_headroom_refusal(e, capture_need)
12092            && let Some(line) = dctx.failed.mark_greedy(&reason)
12093        {
12094            eprintln!("{line}");
12095        }
12096        if graph_draft
12097            && !sampled
12098            && chain_mode
12099            && dctx.chain.is_none()
12100            && !dctx.failed.greedy_failed()
12101        {
12102            if mtp_chain_graph_on() {
12103                let heads_n = self.mtp_head_count();
12104                let DraftGraphCtx {
12105                    g_tok,
12106                    g_pos,
12107                    g_seed,
12108                    g_p,
12109                    g_dmask,
12110                    ..
12111                } = &mut dctx;
12112                if dmask_on {
12113                    e.htod_u32_into(g_dmask, &vec![u32::MAX; dmask_words])?;
12114                }
12115                let g_dmask_ro: &CudaSlice<u32> = &*g_dmask;
12116                let with_prob = p_min > 0.0;
12117                // CAPTURE-RETAIN (#68 fix): one keeper for the whole chain — every graph's
12118                // warmup transients stay pinned as long as any of them replays.
12119                let cap_res = (|| -> Result<DraftChainGraphs, Box<dyn std::error::Error>> {
12120                    // dcw door: same warmup headroom pre-arm as the single-head capture
12121                    // below — every plane, because each head's capture warmups append on
12122                    // its OWN plane. INSIDE the fallible closure (vram-admission lane): an
12123                    // OOM here used to `?` out of the whole burst as a step error; now it
12124                    // is a capture failure — LOUD WARN, eager chain serves.
12125                    if step35_draft_dcw_on() {
12126                        scratch.ensure_dcw_headroom(e, k + 2)?;
12127                    }
12128                    let mut interior = Vec::with_capacity(heads_n);
12129                    let mut last = Vec::with_capacity(heads_n);
12130                    let mut keeper: Vec<Box<dyn std::any::Any + Send>> = Vec::new();
12131                    for hi in 0..heads_n {
12132                        let head = self.mtp_head_at(hi);
12133                        // interior row: KV append + carrier only (`with_head=false` — the
12134                        // eager chain discards interior logits too, so this is the same
12135                        // consumed-byte program minus the dead full-vocab head matmul).
12136                        let (g, keep) = e.capture_graph_retained(|e| {
12137                            self.mtp_head_forward_cap(
12138                                e,
12139                                head,
12140                                g_tok,
12141                                g_pos,
12142                                g_seed,
12143                                g_p,
12144                                &mut *scratch,
12145                                hi,
12146                                false,
12147                                false,
12148                                embd_gpu.expect("graph draft requires resident embedding"),
12149                                embd_qt,
12150                                embd_rb,
12151                                d_vocab,
12152                                None,
12153                                None,
12154                                None,
12155                            )
12156                        })?;
12157                        // the warmups appended rows on plane hi; rewind before the next
12158                        // capture so successive warmups never outrun the pre-armed headroom.
12159                        scratch.set_plane_len(e, hi, base)?;
12160                        interior.push(g);
12161                        keeper.extend(keep);
12162                        // last row: head matmul + greedy argmax tail (+ p when the policy
12163                        // reads it, + the grammar-mask node when constrained).
12164                        let (g2, keep2) = e.capture_graph_retained(|e| {
12165                            self.mtp_head_forward_cap(
12166                                e,
12167                                head,
12168                                g_tok,
12169                                g_pos,
12170                                g_seed,
12171                                g_p,
12172                                &mut *scratch,
12173                                hi,
12174                                with_prob,
12175                                true,
12176                                embd_gpu.expect("graph draft requires resident embedding"),
12177                                embd_qt,
12178                                embd_rb,
12179                                d_vocab,
12180                                None,
12181                                None,
12182                                if dmask_on {
12183                                    Some((g_dmask_ro, dmask_words))
12184                                } else {
12185                                    None
12186                                },
12187                            )
12188                        })?;
12189                        scratch.set_plane_len(e, hi, base)?;
12190                        last.push(g2);
12191                        keeper.extend(keep2);
12192                    }
12193                    Ok(DraftChainGraphs {
12194                        interior,
12195                        last,
12196                        keeper,
12197                    })
12198                })();
12199                match cap_res {
12200                    Ok(cg) => {
12201                        scratch.set_len(e, base)?;
12202                        // POSITIVE engagement receipt (the 3a lesson: a WARN-free boot is
12203                        // NOT evidence of capture — the captured state must name itself).
12204                        eprintln!(
12205                            "[mtp-chain-graph] captured mode=greedy heads={heads_n} \
12206                             interior={heads_n} last={heads_n} masked={}",
12207                            dmask_on as u8
12208                        );
12209                        dctx.chain = Some(cg);
12210                        dctx.graph_masked = dmask_on;
12211                        captured_now = true;
12212                    }
12213                    Err(err) => {
12214                        scratch.set_len(e, base)?;
12215                        // LOUD flip (audit Q2): a dropped draft graph is a coverage loss,
12216                        // never silent — now including the multi-head shipping shape.
12217                        // OOM RECOVERY (vram-admission lane): a failed attempt's freed
12218                        // transients sit CACHED in the async pool where the driver cannot
12219                        // see them; trim them back so the eager fallback (and any driver-
12220                        // side allocation) actually has the headroom the free suggests.
12221                        let mut reason = err.to_string();
12222                        if capture_err_is_oom(&reason) {
12223                            capture_oom_entry_eff = capture_oom_entry_eff.max(cap_eff0);
12224                            let trimmed = e.pool_trim_to_zero();
12225                            if trimmed > 0 {
12226                                reason.push_str(&format!(
12227                                    "; pool trimmed {}MB back to the driver",
12228                                    trimmed / (1 << 20)
12229                                ));
12230                            }
12231                        }
12232                        if let Some(line) = dctx.failed.mark_greedy(&reason) {
12233                            eprintln!("{line}");
12234                        }
12235                    }
12236                }
12237            } else {
12238                // Disarmed by MEMRA_MTP_CHAIN_GRAPH=0: say so once per process — the OFF arm
12239                // must be attributable in a boot log, never inferable from silence.
12240                static NOTE: std::sync::Once = std::sync::Once::new();
12241                NOTE.call_once(|| {
12242                    eprintln!(
12243                        "[spec] multi-head draft-chain capture disarmed \
12244                         (MEMRA_MTP_CHAIN_GRAPH=0); eager chain serves this shape"
12245                    );
12246                });
12247            }
12248        }
12249        if graph_draft
12250            && !sampled
12251            && !chain_mode
12252            && dctx.graph.is_none()
12253            && !dctx.failed.greedy_failed()
12254        {
12255            let DraftGraphCtx {
12256                g_tok,
12257                g_pos,
12258                g_seed,
12259                g_p,
12260                g_dmask,
12261                ..
12262            } = &mut dctx;
12263            // capture-time contents: ALL-ONES (ban nothing). A replay only ever runs after the
12264            // host uploads the position's real words, so the warmups stay grammar-free.
12265            if dmask_on {
12266                e.htod_u32_into(g_dmask, &vec![u32::MAX; dmask_words])?;
12267            }
12268            let g_dmask_ro: &CudaSlice<u32> = &*g_dmask;
12269            // CAPTURE-RETAIN (#68 fix): the warmup transients' pool addresses are baked into the
12270            // captured graph; the keeper pins them for the graph's lifetime. capture_graph (non-
12271            // retained) freed them at exit — safe for one-shot generate_spec (nothing else touches
12272            // the pool between replays) but WRONG for sessions: burst-boundary prime/fill/commit
12273            // passes (and, in serve, other sessions) recycle those addresses and the replay then
12274            // clobbers live buffers — the ST serve-spec corruption (research/serve-st-20260803).
12275            let cap_res = (|| {
12276                // dcw door: the capture warmups append device-counter rows the capture body
12277                // cannot rebase for; pre-arm ring headroom host-side (no-op on flat planes /
12278                // room-enough rings, and the door-off path is untouched). INSIDE the fallible
12279                // closure (vram-admission lane): an OOM here is a capture failure, not a
12280                // burst-killing step error.
12281                if step35_draft_dcw_on() {
12282                    scratch.ensure_dcw_headroom(e, k + 2)?;
12283                }
12284                e.capture_graph_retained(|e| {
12285                    self.mtp_head_forward_cap(
12286                        e,
12287                        mtp,
12288                        g_tok,
12289                        g_pos,
12290                        g_seed,
12291                        g_p,
12292                        &mut *scratch,
12293                        0,
12294                        p_min > 0.0 || fork_mode == OptiForkGateMode::Controller,
12295                        true,
12296                        embd_gpu.expect("graph draft requires resident embedding"),
12297                        embd_qt,
12298                        embd_rb,
12299                        d_vocab,
12300                        None,
12301                        None,
12302                        if dmask_on {
12303                            Some((g_dmask_ro, dmask_words))
12304                        } else {
12305                            None
12306                        },
12307                    )
12308                })
12309            })();
12310            match cap_res {
12311                Ok((g, keep)) => {
12312                    scratch.set_len(e, base)?;
12313                    dctx.graph = Some(g);
12314                    dctx.graph_masked = dmask_on;
12315                    dctx.keeper = keep;
12316                    captured_now = true;
12317                }
12318                Err(err) => {
12319                    scratch.set_len(e, base)?;
12320                    // LOUD flip (audit Q2): a dropped draft graph is a coverage loss, never
12321                    // silent. Once per flip — mark returns None on an already-failed ctx.
12322                    let mut reason = err.to_string();
12323                    if capture_err_is_oom(&reason) {
12324                        capture_oom_entry_eff = capture_oom_entry_eff.max(cap_eff0);
12325                        let trimmed = e.pool_trim_to_zero();
12326                        if trimmed > 0 {
12327                            reason.push_str(&format!(
12328                                "; pool trimmed {}MB back to the driver",
12329                                trimmed / (1 << 20)
12330                            ));
12331                        }
12332                    }
12333                    if let Some(line) = dctx.failed.mark_greedy(&reason) {
12334                        eprintln!("{line}");
12335                    }
12336                }
12337            }
12338        }
12339        // --- SAMPLED GRAPH DRAFT setup (step 3 of the sampled-spec arc): a SECOND capture, own
12340        // graph object, built only when sampled && graph-eligible — the greedy capture above is
12341        // untouched (and skipped when sampled: its graph would never be launched). Same head
12342        // forward, but the in-graph argmax reads GUMBEL-PERTURBED logits; the Philox event
12343        // counter lives in the persistent device g_ctr (bumped in-graph, host-seeded from sctr
12344        // once per round); the raw head logits land in the persistent g_q for the host's
12345        // per-replay async D2D into the round's q slot (q_slots, K x d_vocab, allocated once).
12346        // seed/temp are capture-time constants — baked into graph_s, so a pool-resumed request
12347        // with a different (seed, temp, k) drops the parked sampled graph and recaptures.
12348        // COST OF THE FRESH-SEED SERVE DEFAULT (dogfood F4, 2026-08-04): omitting `seed` on a
12349        // serve request now draws fresh per-request entropy (it used to default to a pinned 0),
12350        // so a seed-omitting request that RESUMES a parked spec session finds an s_key baked
12351        // with the PREVIOUS request's seed and pays one recapture. Bounded, and it does not
12352        // reopen the ~16ms/burst regression the persistent ctx exists to fix: a session's seed
12353        // is fixed for its whole lifetime (worker.rs reads s.sampler.seed() per burst), so
12354        // this compare misses at most ONCE per resumed request — the first burst recaptures
12355        // and every later burst in that request replays. A client that wants the parked graph
12356        // AND reproducibility supplies an explicit `seed`, honored exactly, which keeps s_key
12357        // stable across its whole conversation.
12358        // COMPOSITION RULE (fspec x gsd merge): the in-graph chain samples from the RAW
12359        // softmax — it can hold neither per-row filter stats nor the varying penalty history.
12360        // The sampled graph therefore engages only in the PURE-TEMP regime; filters/penalties
12361        // force the eager draft (which computes stats/penalties per row).
12362        // KEY THE WHOLE REGIME, not just the baked constants (lane/graph-s-key-exactness-
12363        // 20260819). `s_key` used to be `(seed, temp, k)`; the filters and penalties were left
12364        // out, so a filtered request resuming a session that parked a PURE-TEMP graph kept it —
12365        // and the launch site never re-asked `pure_temp`. See [`SampledGraphKey`] for what that
12366        // costs (an unconditional accept of out-of-head draft tokens, i.e. an exactness bug on
12367        // the request shape the vendor-default flip makes the majority).
12368        let s_key = SampledGraphKey::new(sp_seed, sp_temp, k, sp.top_k, sp.top_p, sp.min_p, pen_on);
12369        let pure_temp = s_key.pure_temp();
12370        // The regime the sampled graph may be captured/launched in: pure-temp always;
12371        // truncation-filtered when the filtered-capture door is on (the filter runs
12372        // IN-GRAPH — lane/step37-draft-graph-serving-20260830); penalties never.
12373        let s_capturable = s_key.graph_capturable();
12374        if sampled && dctx.s_key.is_some_and(|old| old != s_key) {
12375            dctx.graph_s = None;
12376            dctx.chain_s = None;
12377            dctx.failed.clear_sampled();
12378            dctx.s_key = None;
12379            dctx.q_slots.clear();
12380            dctx.keeper_s.clear();
12381        }
12382        // PRE-CAPTURE VRAM RESERVE CHECK, sampled arms (vram-admission lane): same contract
12383        // as the greedy check above — refuse BEFORE allocating, LOUD once, eager serves.
12384        if spec_capture_gate_on()
12385            && graph_draft
12386            && sampled
12387            && s_capturable
12388            && !dctx.failed.sampled_failed()
12389            && ((chain_mode && dctx.chain_s.is_none() && mtp_chain_graph_on())
12390                || (!chain_mode && dctx.graph_s.is_none()))
12391            && let Some(reason) = capture_headroom_refusal(e, capture_need)
12392            && let Some(line) = dctx.failed.mark_sampled(&reason)
12393        {
12394            eprintln!("{line}");
12395        }
12396        // FILTERED capture nodes need q slots sized d_vocab AND the stat slots; the pure-temp
12397        // body leaves g_th/g_z/g_mx untouched (they exist from ctx creation either way).
12398        if graph_draft
12399            && sampled
12400            && s_capturable
12401            && chain_mode
12402            && dctx.chain_s.is_none()
12403            && !dctx.failed.sampled_failed()
12404        {
12405            if mtp_chain_graph_on() {
12406                let heads_n = self.mtp_head_count();
12407                let filtered = s_key.filtered();
12408                let DraftGraphCtx {
12409                    g_tok,
12410                    g_pos,
12411                    g_seed,
12412                    g_p,
12413                    g_ctr,
12414                    g_perturb,
12415                    g_q,
12416                    g_rows0,
12417                    g_th,
12418                    g_z,
12419                    g_mx,
12420                    ..
12421                } = &mut dctx;
12422                let with_prob = p_min > 0.0;
12423                let cap_res = (|| -> Result<DraftChainGraphs, Box<dyn std::error::Error>> {
12424                    // dcw pre-arm INSIDE the fallible closure (vram-admission lane): an OOM
12425                    // here is a capture failure with the LOUD WARN, never a step error.
12426                    if step35_draft_dcw_on() {
12427                        scratch.ensure_dcw_headroom(e, k + 2)?;
12428                    }
12429                    let mut interior = Vec::with_capacity(heads_n);
12430                    let mut last = Vec::with_capacity(heads_n);
12431                    let mut keeper: Vec<Box<dyn std::any::Any + Send>> = Vec::new();
12432                    for hi in 0..heads_n {
12433                        let head = self.mtp_head_at(hi);
12434                        // interior row: no head, no draw — shared shape with the greedy
12435                        // chain's interior, captured per mode for keeper-lifetime hygiene.
12436                        let (g, keep) = e.capture_graph_retained(|e| {
12437                            self.mtp_head_forward_cap(
12438                                e,
12439                                head,
12440                                g_tok,
12441                                g_pos,
12442                                g_seed,
12443                                g_p,
12444                                &mut *scratch,
12445                                hi,
12446                                false,
12447                                false,
12448                                embd_gpu.expect("graph draft requires resident embedding"),
12449                                embd_qt,
12450                                embd_rb,
12451                                d_vocab,
12452                                None,
12453                                None,
12454                                None,
12455                            )
12456                        })?;
12457                        scratch.set_plane_len(e, hi, base)?;
12458                        interior.push(g);
12459                        keeper.extend(keep);
12460                        // last row: head matmul + the in-graph categorical draw (filtered
12461                        // nodes when the request carries filters).
12462                        let (g2, keep2) = e.capture_graph_retained(|e| {
12463                            self.mtp_head_forward_cap(
12464                                e,
12465                                head,
12466                                g_tok,
12467                                g_pos,
12468                                g_seed,
12469                                g_p,
12470                                &mut *scratch,
12471                                hi,
12472                                with_prob,
12473                                true,
12474                                embd_gpu.expect("graph draft requires resident embedding"),
12475                                embd_qt,
12476                                embd_rb,
12477                                d_vocab,
12478                                Some(SampledCapArgs {
12479                                    ctr: &mut *g_ctr,
12480                                    perturb: &mut *g_perturb,
12481                                    q_out: &mut *g_q,
12482                                    seed: sp_seed,
12483                                    temp: sp_temp,
12484                                    filt: if filtered {
12485                                        Some(SampledCapFilter {
12486                                            rows0: &*g_rows0,
12487                                            th: &mut *g_th,
12488                                            z: &mut *g_z,
12489                                            mx: &mut *g_mx,
12490                                            top_k: sp.top_k,
12491                                            top_p: sp.top_p,
12492                                            min_p: sp.min_p,
12493                                        })
12494                                    } else {
12495                                        None
12496                                    },
12497                                }),
12498                                None,
12499                                None, // constrained spec is greedy-only
12500                            )
12501                        })?;
12502                        scratch.set_plane_len(e, hi, base)?;
12503                        last.push(g2);
12504                        keeper.extend(keep2);
12505                    }
12506                    Ok(DraftChainGraphs {
12507                        interior,
12508                        last,
12509                        keeper,
12510                    })
12511                })();
12512                match cap_res {
12513                    Ok(cg) => {
12514                        scratch.set_len(e, base)?;
12515                        // NO STRANDED PARTIAL STATE (vram-admission lane): the q-slot allocs
12516                        // after a successful capture are themselves fallible on a tight card.
12517                        // A mid-loop failure used to `?` out as a step error, leaving orphan
12518                        // slots parked on the ctx (wrong count, stale contents) for the next
12519                        // capture attempt to stack onto. Allocate all-or-nothing: on failure
12520                        // drop the fresh graphs AND the partial slots, mark the LOUD fallback.
12521                        dctx.q_slots.clear();
12522                        let slots = (0..k)
12523                            .map(|_| e.zeros(d_vocab))
12524                            .collect::<Result<Vec<_>, _>>();
12525                        match slots {
12526                            Ok(slots) => {
12527                                dctx.q_slots = slots;
12528                                eprintln!(
12529                                    "[mtp-chain-graph] captured mode=sampled heads={heads_n} \
12530                                     interior={heads_n} last={heads_n} filtered={} key={s_key:?}",
12531                                    s_key.filtered() as u8
12532                                );
12533                                dctx.chain_s = Some(cg);
12534                                dctx.s_key = Some(s_key);
12535                                captured_now = true;
12536                            }
12537                            Err(err) => {
12538                                drop(cg);
12539                                dctx.q_slots.clear();
12540                                let mut reason = format!("q-slot alloc failed: {err}");
12541                                if capture_err_is_oom(&reason) {
12542                                    capture_oom_entry_eff = capture_oom_entry_eff.max(cap_eff0);
12543                                    let trimmed = e.pool_trim_to_zero();
12544                                    if trimmed > 0 {
12545                                        reason.push_str(&format!(
12546                                            "; pool trimmed {}MB back to the driver",
12547                                            trimmed / (1 << 20)
12548                                        ));
12549                                    }
12550                                }
12551                                if let Some(line) = dctx.failed.mark_sampled(&reason) {
12552                                    eprintln!("{line}");
12553                                }
12554                            }
12555                        }
12556                    }
12557                    Err(err) => {
12558                        scratch.set_len(e, base)?;
12559                        let mut reason = err.to_string();
12560                        if capture_err_is_oom(&reason) {
12561                            capture_oom_entry_eff = capture_oom_entry_eff.max(cap_eff0);
12562                            let trimmed = e.pool_trim_to_zero();
12563                            if trimmed > 0 {
12564                                reason.push_str(&format!(
12565                                    "; pool trimmed {}MB back to the driver",
12566                                    trimmed / (1 << 20)
12567                                ));
12568                            }
12569                        }
12570                        if let Some(line) = dctx.failed.mark_sampled(&reason) {
12571                            eprintln!("{line}");
12572                        }
12573                    }
12574                }
12575            } else {
12576                static NOTE_S: std::sync::Once = std::sync::Once::new();
12577                NOTE_S.call_once(|| {
12578                    eprintln!(
12579                        "[spec] multi-head draft-chain capture disarmed \
12580                         (MEMRA_MTP_CHAIN_GRAPH=0); eager chain serves this shape"
12581                    );
12582                });
12583            }
12584        }
12585        if graph_draft
12586            && sampled
12587            && s_capturable
12588            && !chain_mode
12589            && dctx.graph_s.is_none()
12590            && !dctx.failed.sampled_failed()
12591        {
12592            let filtered = s_key.filtered();
12593            let DraftGraphCtx {
12594                g_tok,
12595                g_pos,
12596                g_seed,
12597                g_p,
12598                g_ctr,
12599                g_perturb,
12600                g_q,
12601                g_rows0,
12602                g_th,
12603                g_z,
12604                g_mx,
12605                ..
12606            } = &mut dctx;
12607            // CAPTURE-RETAIN (#68 fix): same keeper contract as the greedy capture above.
12608            let cap_res = (|| {
12609                // dcw pre-arm INSIDE the fallible closure (vram-admission lane): an OOM
12610                // here is a capture failure with the LOUD WARN, never a step error.
12611                if step35_draft_dcw_on() {
12612                    scratch.ensure_dcw_headroom(e, k + 2)?;
12613                }
12614                e.capture_graph_retained(|e| {
12615                    self.mtp_head_forward_cap(
12616                        e,
12617                        mtp,
12618                        g_tok,
12619                        g_pos,
12620                        g_seed,
12621                        g_p,
12622                        &mut *scratch,
12623                        0,
12624                        p_min > 0.0,
12625                        true,
12626                        embd_gpu.expect("graph draft requires resident embedding"),
12627                        embd_qt,
12628                        embd_rb,
12629                        d_vocab,
12630                        Some(SampledCapArgs {
12631                            ctr: &mut *g_ctr,
12632                            perturb: &mut *g_perturb,
12633                            q_out: &mut *g_q,
12634                            seed: sp_seed,
12635                            temp: sp_temp,
12636                            filt: if filtered {
12637                                Some(SampledCapFilter {
12638                                    rows0: &*g_rows0,
12639                                    th: &mut *g_th,
12640                                    z: &mut *g_z,
12641                                    mx: &mut *g_mx,
12642                                    top_k: sp.top_k,
12643                                    top_p: sp.top_p,
12644                                    min_p: sp.min_p,
12645                                })
12646                            } else {
12647                                None
12648                            },
12649                        }),
12650                        None,
12651                        None, // constrained spec is greedy-only — sampled never carries a hook
12652                    )
12653                })
12654            })();
12655            match cap_res {
12656                Ok((g, keep)) => {
12657                    scratch.set_len(e, base)?;
12658                    // NO STRANDED PARTIAL STATE: all-or-nothing q slots, same contract as
12659                    // the chain arm above.
12660                    dctx.q_slots.clear();
12661                    let slots = (0..k)
12662                        .map(|_| e.zeros(d_vocab))
12663                        .collect::<Result<Vec<_>, _>>();
12664                    match slots {
12665                        Ok(slots) => {
12666                            dctx.q_slots = slots;
12667                            dctx.graph_s = Some(g);
12668                            dctx.s_key = Some(s_key);
12669                            dctx.keeper_s = keep;
12670                            captured_now = true;
12671                        }
12672                        Err(err) => {
12673                            drop(g);
12674                            drop(keep);
12675                            dctx.q_slots.clear();
12676                            let mut reason = format!("q-slot alloc failed: {err}");
12677                            if capture_err_is_oom(&reason) {
12678                                capture_oom_entry_eff = capture_oom_entry_eff.max(cap_eff0);
12679                                let trimmed = e.pool_trim_to_zero();
12680                                if trimmed > 0 {
12681                                    reason.push_str(&format!(
12682                                        "; pool trimmed {}MB back to the driver",
12683                                        trimmed / (1 << 20)
12684                                    ));
12685                                }
12686                            }
12687                            if let Some(line) = dctx.failed.mark_sampled(&reason) {
12688                                eprintln!("{line}");
12689                            }
12690                        }
12691                    }
12692                }
12693                Err(err) => {
12694                    scratch.set_len(e, base)?;
12695                    // LOUD flip (audit Q2): same contract as the greedy capture above.
12696                    let mut reason = err.to_string();
12697                    if capture_err_is_oom(&reason) {
12698                        capture_oom_entry_eff = capture_oom_entry_eff.max(cap_eff0);
12699                        let trimmed = e.pool_trim_to_zero();
12700                        if trimmed > 0 {
12701                            reason.push_str(&format!(
12702                                "; pool trimmed {}MB back to the driver",
12703                                trimmed / (1 << 20)
12704                            ));
12705                        }
12706                    }
12707                    if let Some(line) = dctx.failed.mark_sampled(&reason) {
12708                        eprintln!("{line}");
12709                    }
12710                }
12711            }
12712        }
12713        // ---- PER-SESSION DRAFT-STATE MEASUREMENT bracket end (vram-admission lane): when a
12714        // capture landed in THIS call, the effective-free delta across the capture section is
12715        // this session's parked draft-graph state (keepers + q slots + instantiated graphs'
12716        // backing). Recorded as a model-owned high-water; admission charges it per
12717        // spec-capable session (see `draft_session_admission_bytes`).
12718        if captured_now
12719            && let Some(eff0) = cap_eff0
12720            && let Ok((f1, _)) = e.ctx().mem_get_info()
12721        {
12722            let eff1 = f1.saturating_add(e.pool_cached_bytes());
12723            let parked_delta = eff0.saturating_sub(eff1);
12724            let (_res_high, used_high) = e.pool_high_water_reset();
12725            let peak_delta = used_high.saturating_sub(cap_used0);
12726            let observed = parked_delta.max(peak_delta);
12727            if observed > 0
12728                && let Some(hw) = self.record_draft_state_bytes(observed)
12729            {
12730                eprintln!(
12731                    "[spec] draft-session state high-water: {}MB (max of parked delta {}MB \
12732                     and capture-time pool peak {}MB; charged per spec admission and gating \
12733                     future captures)",
12734                    hw / (1 << 20),
12735                    parked_delta / (1 << 20),
12736                    peak_delta / (1 << 20),
12737                );
12738            }
12739        }
12740        // FAILURE IS AN OBSERVATION TOO: a capture that OOM'd at entry-effective E proved
12741        // the capture-time peak exceeds E. Feed E into the gauge so every future gate
12742        // refuses at or below the headroom that just failed (self-healing even when the
12743        // boot probe is disarmed and the bootstrap estimate was blind).
12744        if let Some(entry_eff) = capture_oom_entry_eff
12745            && let Some(hw) = self.record_draft_state_bytes(entry_eff)
12746        {
12747            eprintln!(
12748                "[spec] draft-session capture appetite floor raised to {}MB: a capture \
12749                 attempt OOM'd with that much effective free (failure-observed bound)",
12750                hw / (1 << 20)
12751            );
12752        }
12753        // ---- EXACTNESS GUARD, the enforceable half (lane/graph-s-key-exactness-20260819,
12754        // widened by lane/step37-draft-graph-serving-20260830) ----
12755        // With the filters and penalties in `s_key`, a graph that SURVIVED the drop above was
12756        // captured under THIS request's exact regime, and capture requires `graph_capturable`
12757        // (pure-temp, or filtered with the in-graph filter nodes; never penalties) — so a
12758        // parked graph implies both. That implication is the whole exactness argument for the
12759        // graph arm, so it is asserted here rather than assumed: a future change that widens
12760        // the capture condition, narrows the key, or copies a `DraftGraphCtx` across regimes
12761        // fails LOUDLY at this line instead of silently drafting from a distribution the
12762        // verify never reconstructs. Release builds refuse the graph (drop it, draft eager)
12763        // rather than launching it; the launch site re-tests the regime independently.
12764        if sampled
12765            && (dctx.graph_s.is_some() || dctx.chain_s.is_some())
12766            && (!s_capturable || dctx.s_key != Some(s_key))
12767        {
12768            debug_assert!(
12769                false,
12770                "sampled draft graph parked under {:?} survived into a request outside its \
12771                 capture regime (top_k={} top_p={} min_p={} pen_on={} capturable={}): the \
12772                 in-graph draw and the verify's accept test would see different distributions",
12773                dctx.s_key, sp.top_k, sp.top_p, sp.min_p, pen_on, s_capturable,
12774            );
12775            eprintln!(
12776                "[spec] BUG: dropping a parked sampled draft graph that outlived its capture \
12777                 regime (s_key={:?}, request top_k={} top_p={} min_p={} pen_on={} \
12778                 capturable={}); drafting EAGER — the key must carry every field that shapes q",
12779                dctx.s_key, sp.top_k, sp.top_p, sp.min_p, pen_on, s_capturable,
12780            );
12781            dctx.graph_s = None;
12782            dctx.chain_s = None;
12783            dctx.s_key = None;
12784            dctx.q_slots.clear();
12785            dctx.keeper_s.clear();
12786        }
12787        // SKEY PROBE (MEMRA_SKEY_PROBE=1): the burst-entry facts the reachability question turns
12788        // on — is this request sampled, is it in a regime the sampled graph is legal in, and is
12789        // a graph PARKED from an earlier request of the same session? The launch arms below
12790        // print which chain actually ran, so the probe never restates the condition.
12791        if skey_probe() {
12792            eprintln!(
12793                "[skey] burst sampled={} pure_temp={} capturable={} temp={} top_k={} top_p={} \
12794                 min_p={} pen_on={} k={} graph_draft={} graph_s_parked={} chain_s_parked={} \
12795                 s_key_parked={:?}",
12796                sampled as u8,
12797                pure_temp as u8,
12798                s_capturable as u8,
12799                sp_temp,
12800                sp.top_k,
12801                sp.top_p,
12802                sp.min_p,
12803                pen_on as u8,
12804                k,
12805                graph_draft as u8,
12806                dctx.graph_s.is_some() as u8,
12807                dctx.chain_s.is_some() as u8,
12808                dctx.s_key,
12809            );
12810        }
12811        let t_cap = t_ent.elapsed();
12812        // PERSISTENT DRAFT KV: fill the MTP block's K/V for every prompt position from the exact
12813        // trunk hiddens collected during prime — ONE batched K/V-only pass (overwrites any
12814        // capture-warmup garbage; capture left len at 0). last_token (the init feed) needs no
12815        // fill: the first chain step processes it and appends its entry at slot prompt.len().
12816        if let Some(ph) = &prompt_h {
12817            // SESSION: rows [0..base) are the previous turns' exact fills (refresh overwrote them
12818            // with true verify hiddens) — truncate any draft overhang, fill ONLY the suffix at
12819            // global positions [base..base+tp). Fresh call: base==0, identical to before.
12820            scratch.set_len(e, base)?;
12821            // CHUNKED FILL (long-ctx OOM fix, 2026-07-05): mtp_kv_fill's transients scale with its
12822            // T (concat = T*2*n_embd*4B — 1.5GB at 40k) and its concat loop is 2*T launches. The
12823            // fill is a pure sequential append, so chunking is exact: each chunk appends its rows
12824            // at pos0=base+start with the identical per-row math. Same knob as the trunk prime.
12825            let tp = prompt.len();
12826            let fill_chunk: usize = if crate::cache::swa_ring_on() {
12827                crate::hybrid_forward::prime_chunk_tokens(tp, self.layers.len())
12828            } else {
12829                // Preserve the flag-OFF schedule byte-for-byte, including the legacy zero value
12830                // meaning one monolithic fill.
12831                std::env::var("MEMRA_PRIME_CHUNK")
12832                    .ok()
12833                    .and_then(|v| v.parse().ok())
12834                    .unwrap_or(4096)
12835            };
12836            let fill_chunk = if fill_chunk == 0 { tp } else { fill_chunk };
12837            let mut start = 0usize;
12838            while start < tp {
12839                let end = (start + fill_chunk).min(tp);
12840                let tc = end - start;
12841                {
12842                    // PREDECESSOR pairing: row i gets h[i-1]; global row 0 a zeros row (the
12843                    // reference engine's initial pending-h is zeroed too); a session turn's row 0
12844                    // gets the PREVIOUS turn's last committed hidden (sess.last_h). Per chunk:
12845                    // rows start..end read h[start-1..end-1] — one dtod into a chunk buffer.
12846                    let mut phs = e.zeros(tc * n_embd)?;
12847                    let (src_lo, dst_off) = if start == 0 {
12848                        (0, n_embd)
12849                    } else {
12850                        ((start - 1) * n_embd, 0)
12851                    };
12852                    let n_copy = if start == 0 {
12853                        (tc - 1) * n_embd
12854                    } else {
12855                        tc * n_embd
12856                    };
12857                    if start == 0 {
12858                        if let Some((_, lh, _, _, _)) = sess_tail.as_ref() {
12859                            if let Some(lh) = lh.as_ref() {
12860                                e.copy_into(&mut phs, 0, lh, n_embd)?;
12861                            }
12862                        }
12863                    }
12864                    if n_copy > 0 {
12865                        e.copy_view_into(
12866                            &mut phs,
12867                            dst_off,
12868                            &ph.slice(src_lo..src_lo + n_copy),
12869                            n_copy,
12870                        )?;
12871                    }
12872                    self.mtp_kv_fill_all(
12873                        e,
12874                        &prompt[start..end],
12875                        &phs,
12876                        base + start,
12877                        &mut *scratch,
12878                        embd_dev,
12879                    )?;
12880                }
12881                start = end;
12882            }
12883        }
12884        // MEMRA_PROFILE_SPEC=2: profiler capture starts HERE — after the prime, so an
12885        // `nsys -c cudaProfilerApi` capture contains ONLY the round loop (draft/verify/commit).
12886        // (=1 brackets the whole call in run_spec.rs, prime included.)
12887        if std::env::var("MEMRA_PROFILE_SPEC").as_deref() == Ok("2") {
12888            unsafe extern "C" {
12889                fn cudaProfilerStart() -> i32;
12890            }
12891            unsafe {
12892                cudaProfilerStart();
12893            }
12894        }
12895        // ROUND-STREAM stage (c) 4 (MEMRA_SPEC_STREAM=1, experimental): pre-issued M-round
12896        // bursts with ZERO per-round host readbacks — the accept/seed/rollback/ring kernels
12897        // consume each other's device outputs; the host drains the ring every M rounds. v1
12898        // constraints: greedy, !spec_replay, single-shot, batched-linear layers, no refresh
12899        // fills (acceptance effect A/B-arbitrated), enters from round 1 (pending guaranteed).
12900        // NOTE: not gated on the caller's graph_draft (its trunk_dense conjunct turns the 35B
12901        // MoE off) — the stream capture encloses ONLY the dense MTP head; the head-dense /
12902        // full-prec / k gates are re-derived here and a failed capture degrades to stream-off.
12903        let stream_on = crate::spec::spec_stream()
12904            && !sampled
12905            && !spec_replay
12906            && self.mtp_extra.is_empty()
12907            && constraint.is_none()
12908            && !session_mode
12909            && embd_gpu.is_some()
12910            && !crate::model::full_prec_enabled()
12911            && k + 2 < 96;
12912        let mut stream_graph: Option<cudarc::driver::CudaGraph> = None;
12913        let mut g_tokp2k = e.alloc_u32_zeroed(2 * k.max(1))?;
12914        if stream_on {
12915            let cap = e.capture_graph(|e| {
12916                for j in 0..k.max(1) {
12917                    self.mtp_head_forward_cap(
12918                        e,
12919                        mtp,
12920                        &mut dctx.g_tok,
12921                        &mut dctx.g_pos,
12922                        &mut dctx.g_seed,
12923                        &mut dctx.g_p,
12924                        &mut *scratch,
12925                        0,
12926                        true,
12927                        true,
12928                        embd_gpu.expect("round stream requires resident embedding"),
12929                        embd_qt,
12930                        embd_rb,
12931                        d_vocab,
12932                        None,
12933                        Some((&mut g_tokp2k, j, d2t_dev.as_ref())),
12934                        None, // round-stream requires constraint.is_none() (see stream_on)
12935                    )?;
12936                }
12937                Ok(())
12938            });
12939            match cap {
12940                Ok(g) => {
12941                    scratch.set_len(e, 0)?;
12942                    stream_graph = Some(g);
12943                }
12944                Err(err) => {
12945                    scratch.set_len(e, 0)?;
12946                    if debug_spec {
12947                        eprintln!("[spec] stream-graph capture failed ({err}); stream off");
12948                    }
12949                }
12950            }
12951        }
12952        let stream_active = stream_on && stream_graph.is_some();
12953        if debug_spec {
12954            eprintln!(
12955                "[spec] stream_on={stream_on} env={} samp={sampled} dg={} captured={} active={stream_active} session={session_mode} replay={spec_replay}",
12956                crate::spec::spec_stream(),
12957                dctx.graph.is_some(),
12958                stream_graph.is_some()
12959            );
12960        }
12961        let t_v_s = k + 1;
12962        // ROUND-STREAM buffers + ptr tables now live in the model-generic round_stream
12963        // module (extracted 2026-07-12; the gemma burst reuses them).
12964        let sb = crate::round_stream::StreamBufs::new(e, k, crate::spec::spec_stream_m())?;
12965        let crate::round_stream::StreamBufs {
12966            mut vtok_d,
12967            mut brk_d,
12968            mut pend_d,
12969            last_pred_d,
12970            mut pos_ctr,
12971            mut pos_start_d,
12972            mut ring_d,
12973            acc_d: mut stream_acc,
12974            m_rounds,
12975            k: _,
12976        } = sb;
12977        let stream_ptrs: Option<CudaSlice<u64>> = if stream_active {
12978            Some(crate::round_stream::kv_len_ptr_table(
12979                e,
12980                cache,
12981                Some(&pos_ctr),
12982            )?)
12983        } else {
12984            None
12985        };
12986
12987        let t_fill = t_ent.elapsed();
12988        let mut round = 0usize;
12989        // ADAPTIVE DRAFT LENGTH (MEMRA_SPEC_ADAPT=1, opt-in — the gemma_spec accepted-run law,
12990        // ported 2026-08-01): next round's draft depth = last round's accepted run + 1, clamped
12991        // to [floor(pos), k_cap] — a miss shrinks the next draft to the miss point + 1,
12992        // full-accept streaks re-deepen one step per round. NOT the 2026-07-07 acceptance-EMA
12993        // (that arm measured an HONEST LOSS to static per-class optima — 115.0/85.8/73.4 vs
12994        // 121.6/92.7/75.6, EMA lag — and was removed 2026-07-08; rig5090.jsonl has the record).
12995        // The gemma law has no lag class: it reacts within one round, and was worth +7-20% on
12996        // the gemma cells at unchanged exactness (2026-07-10 flip; floor sweep 2026-07-25;
12997        // position key 2026-07-26). Signal = n_acc from the round's EXISTING accept readback —
12998        // zero new syncs; the draft graph is a SINGLE-STEP capture replayed per drafted token,
12999        // so a per-round depth needs no re-capture (unlike gemma's whole-chain graphs). qwen's
13000        // in-round p-min cut already shortens chains mid-round, so gemma's one-round-late p-min
13001        // fold into kc is unnecessary here — the accepted-run law sees the cut via n_acc.
13002        // Exactness is the verify's job at ANY depth (same contract as p-min variable rounds).
13003        // DEFAULT OFF on the qwen path until its cells gate a flip (gemma's is default-on).
13004        // MEASURED 2026-08-01 (H100 GPU-3, interleaved x3, NGEN=256, same-invocation plain
13005        // denominators; research/qwen-adaptive-k-20260801/): REFUTED on the tuned qwen configs.
13006        // q27 K=3+HPOST+PMIN=0.3: short +0.8% (noise; law ~idles, len_hist identical), board
13007        // -1.9%, agentic -0.5%; board PMIN=0 -2.1% (not p-min shadowing — the law itself);
13008        // floor=1 -2.8% (gemma's floor-collapse, reproduced). q35 K=2 board: -6.4% (52/136
13009        // rounds shrink to depth 1; no depth to reclaim at K=2). The gemma direction DOES
13010        // appear at untuned depth-K — q27 K=6 floor=4 +1.5% over fixed K=6 — but stays -3.7%
13011        // below fixed K=3: same verdict class as the retired EMA arm (honest loss to static
13012        // per-class optima). Acceptance-rate rises under the law while tokens/round falls —
13013        // it buys accept-% by adding rounds, and a round's fixed draft+verify cost wins.
13014        // K=1..8 self-consistency PASS both models with the law ON (exactness held).
13015        let adapt = std::env::var("MEMRA_SPEC_ADAPT").as_deref() == Ok("1");
13016        // floor: per-model default keyed on n_embd (gemma's tiering — models with an expensive
13017        // verify keep deep drafts after a miss); MEMRA_SPEC_ADAPT_FLOOR pins it everywhere.
13018        let adapt_floor_env: Option<usize> = std::env::var("MEMRA_SPEC_ADAPT_FLOOR")
13019            .ok()
13020            .and_then(|v| v.parse().ok());
13021        let adapt_floor_default: usize = if self.cfg.n_embd as usize >= 3500 {
13022            4
13023        } else if self.cfg.n_embd as usize >= 2500 {
13024            2
13025        } else {
13026            1
13027        };
13028        let adapt_floor: usize = adapt_floor_env.unwrap_or(adapt_floor_default);
13029        // position key: past floor_ctx a HIGH floor (>=4) relaxes to 1 — forced-deep drafts
13030        // turn net-negative at depth (gemma 31B d1736 evidence); MEMRA_SPEC_FLOOR_CTX moves
13031        // the boundary, an explicit MEMRA_SPEC_ADAPT_FLOOR pins the floor everywhere.
13032        let floor_ctx: usize = std::env::var("MEMRA_SPEC_FLOOR_CTX")
13033            .ok()
13034            .and_then(|v| v.parse().ok())
13035            .unwrap_or(1024);
13036        let floor_at = |pos: usize| -> usize {
13037            if adapt_floor_env.is_some() || pos < floor_ctx {
13038                adapt_floor
13039            } else if adapt_floor >= 4 {
13040                1
13041            } else {
13042                adapt_floor
13043            }
13044        };
13045        // cap: MEMRA_SPEC_CAPMAX (gemma semantics, default 7). Binds only under adapt — the
13046        // fixed-K default path is untouched by this whole block.
13047        let cap_max: usize = std::env::var("MEMRA_SPEC_CAPMAX")
13048            .ok()
13049            .and_then(|v| v.parse().ok())
13050            .unwrap_or(7);
13051        let k_cap = k.min(cap_max).max(1);
13052        let mut kc = k_cap;
13053        let mut opti_fork: Option<OptiForkState> = None;
13054        let mut fork_snapshot: Option<crate::cache::CacheSnapshot> = None;
13055        if fork_mode != OptiForkGateMode::Disabled {
13056            let fence = crate::pp::pp_cuts(self.layers.len());
13057            let refusal = if !session_mode {
13058                Some("not-session")
13059            } else if k != 1 || adapt {
13060                Some("requires-fixed-k1")
13061            } else if sampled || constraint.is_some() || spec_replay {
13062                Some("sampled-constrained-or-replay")
13063            } else if pipe.is_some() {
13064                Some("two-session-pipeline")
13065            } else if !spec_devacc() {
13066                Some("requires-device-accept")
13067            } else if stream_active || crate::spec::spec_stream() {
13068                Some("round-stream")
13069            } else if !self.mtp_extra.is_empty() {
13070                Some("multi-head-mtp")
13071            } else if crate::cache::swa_ring_on() || cache.has_swa_ring() {
13072                Some("swa-ring")
13073            } else if crate::pp::pp_host_bounce_active() {
13074                Some("host-bounce")
13075            } else if fork_mode == OptiForkGateMode::Controller
13076                && cache.recur.iter().any(Option::is_some)
13077            {
13078                Some("controller-requires-zero-recurrent-state")
13079            } else if fence.as_ref().is_none_or(|f| f.len() != 3) {
13080                Some("requires-pp2")
13081            } else {
13082                None
13083            };
13084            if let Some(reason) = refusal {
13085                OPTI_FORK_REFUSALS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
13086                eprintln!("[opti-fork] refused reason={reason}");
13087            } else {
13088                let fence = fence.expect("validated PP-2 fence");
13089                let rt = crate::pp::PpNRt::get(e)?;
13090                let primary_stage0 = rt.engine(0, e).ctx().ordinal() == e.ctx().ordinal();
13091                let primary_stage1 = rt.engine(1, e).ctx().ordinal() == e.ctx().ordinal();
13092                let primary_supported =
13093                    primary_stage0 || (fork_mode == OptiForkGateMode::Controller && primary_stage1);
13094                if !rt.cross_device() || !primary_supported {
13095                    OPTI_FORK_REFUSALS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
13096                    eprintln!("[opti-fork] refused reason=requires-supported-primary-cross-device");
13097                } else {
13098                    // Both recurrent snapshots and both seed generations are allocated before
13099                    // the first fork, each through its owning PP stage. Allocation failure
13100                    // therefore happens before any optimistic state mutation can occur.
13101                    let current_snapshot = opti_snapshot_stage_owned(e, cache, rt, &fence)?;
13102                    let alternate_snapshot = opti_snapshot_stage_owned(e, cache, rt, &fence)?;
13103                    let fork = OptiForkState::new(
13104                        e,
13105                        cache,
13106                        fork_mode,
13107                        alternate_snapshot,
13108                        &h_seed_buf,
13109                        &fill_prev,
13110                        rt,
13111                        fence[1],
13112                        self.layers.len(),
13113                    )?;
13114                    eprintln!(
13115                        "[opti-fork] armed mode={fork_mode:?} snapshots=2 seeds=2 split={} \
13116                         payload_dev0={} payload_dev1={} q_threshold={:.3}",
13117                        fence[1],
13118                        fork.logical_payload_bytes[0],
13119                        fork.logical_payload_bytes[1],
13120                        fork.controller.map_or(0.0, |policy| policy.threshold),
13121                    );
13122                    fork_snapshot = Some(current_snapshot);
13123                    opti_fork = Some(fork);
13124                }
13125            }
13126        }
13127        // Persistent snapshot buffers are allocated once and refreshed in place. The fork arm
13128        // uses stage-owned snapshots; refused/disabled arms retain the existing generic helper.
13129        let mut snap = match fork_snapshot {
13130            Some(snapshot) => snapshot,
13131            None => cache.snapshot(e)?,
13132        };
13133        let mut carried_opti: Option<OptiControllerTicket> = None;
13134        // ROUND-STREAM stage (b) 3a: device table of per-layer kvl.len_d pointers (stable — the
13135        // cache never reallocates len_d; see cache.rs "stable pointer" note). 0 = no KV layer.
13136        let kv_len_ptrs: Option<CudaSlice<u64>> = if spec_devacc() && !spec_replay {
13137            Some(crate::round_stream::kv_len_ptr_table(e, cache, None)?)
13138        } else {
13139            None
13140        };
13141        // BONUS FOLD (2026-07-04): after a FULL accept the bonus token is NOT committed with a
13142        // separate T=1 trunk pass (a full weight read per round). It stays PENDING and rides as
13143        // column 0 of the NEXT round's verify batch. Under predecessor pairing the next chain
13144        // seeds from the bonus's predecessor's TRUE verify hidden (free — no extra
13145        // pass of any kind). Verify still
13146        // checks every emitted token against the target -> exactness holds by construction; only
13147        // DRAFT QUALITY can shift, which the acceptance numbers arbitrate.
13148        // bonus emitted but not yet committed to cache. A carried pending (see SpecSession::
13149        // pending_tok) enters round 0 directly — the burst boundary becomes a plain round edge.
13150        let mut pending: Option<u32> = carried_pending;
13151        // MEMRA_SPEC_PHASE=1: per-round wall decomposition (draft / verify / accept+commit) —
13152        // no tracing, no extra syncs (each phase is naturally sync-bounded: draft readbacks,
13153        // the verify accept readback). Printed once at loop end via spec-stats.
13154        let anatomy_on = std::env::var("MEMRA_SPEC_PP_ANATOMY").as_deref() == Ok("1");
13155        let phase_on = anatomy_on || std::env::var("MEMRA_SPEC_PHASE").as_deref() == Ok("1");
13156        // MEMRA_SPEC_PHASE_SYNC=1 — reads the phase split correctly, and proves it. `ph_mark` is a
13157        // bare Instant, so `verify-issue` is the host QUEUEING the walk (the GPU is already running
13158        // under it) and `verify-wait` is only the residual drain at the accept readback: one
13159        // overlapped interval cut at the first blocking call, NOT "GPU time" beside "host time".
13160        // Syncing right after the walk is issued moves the whole GPU wall into `verify-issue`. If
13161        // the walk's GPU total is really issue+wait, then with this on verify-issue jumps to that
13162        // sum, verify-wait collapses to the readback alone, and the ROUND WALL DOES NOT MOVE —
13163        // which is what says the queueing time was hidden and is not a target. Diagnostic only.
13164        let phase_sync = std::env::var("MEMRA_SPEC_PHASE_SYNC").as_deref() == Ok("1");
13165        // DRAFT-MASK receipt (lane/draft-mask): speculative-clone wall + rounds, printed with
13166        // spec-stats. The clone is the one cost the design adds per round — measured, not assumed.
13167        let (mut dm_clone_ns, mut dm_rounds) = (0u128, 0usize);
13168        // grammar-truncation counters: how many rounds the verify-side cut fired and how many
13169        // already-verified tokens it threw away. THIS is the quantity draft masking targets.
13170        let (mut dm_cuts, mut dm_cut_tokens) = (0usize, 0usize);
13171        let (mut ph_draft, mut ph_verify, mut ph_rest) = (0f64, 0f64, 0f64);
13172        let mut ph_wait = 0f64;
13173        let mut ph_commit = 0f64;
13174        let mut ph_t = std::time::Instant::now();
13175        let mut ph_mark = |acc: &mut f64, on: bool| {
13176            if on {
13177                let now = std::time::Instant::now();
13178                *acc += (now - ph_t).as_secs_f64();
13179                ph_t = now;
13180            }
13181        };
13182        // MTP-ROUTE VERIFY GRAPHS (`MEMRA_SPEC_VERIFY_GRAPH`, see the flag doc): the
13183        // model-owned capture pool, locked for the whole burst exactly as the dspark serve
13184        // arm holds it — the slab stash is live verify -> commit inside a round, and the
13185        // worker drives rounds from one scheduler thread. PERSISTENT across generations on
13186        // the model (rebuilding per call re-captures the pool per prompt, which is the
13187        // measured way to lose more than the launches cost); the captured bodies are
13188        // cache-independent, every state read going through per-round refreshed pointer
13189        // tables. None = the eager walk, byte-identical.
13190        //
13191        // Never armed together with ROUND-STREAM: the tparallel verify refuses that pair
13192        // loudly, and `stream_active` owns the burst arm above, so the door stays shut
13193        // whenever the stream is live rather than relying on that refusal.
13194        // The lock is taken ONLY when the door is armed: with the flag off this whole block
13195        // is inert, so the default path cannot serialize two spec generations behind a mutex
13196        // it never reads.
13197        let vg_armed =
13198            crate::spec::spec_verify_graph_env().unwrap_or_else(|| self.vgraph_family_default());
13199        let mut vg_guard = if vg_armed && !stream_active {
13200            let mut g = self.dspark_vgraphs.lock().unwrap();
13201            if g.is_none() {
13202                // Size by the WIDEST verify this run can present, which is k+1 and NOT
13203                // k_cap+1: the sampled arm's own window is `t_v_s = k + 1`, so a pool built
13204                // from a smaller adaptive cap gets sliced past its stash rows (a `slice_mut`
13205                // panic in the sampled ON arm, measured before this line said k+1).
13206                let vt_cap = (k.max(k_cap) + 1).max(2);
13207                *g = DsparkVerifyGraphs::new(e, cache, vt_cap, n_embd)?;
13208                if g.is_some() {
13209                    // Engagement receipt (the dead-arm lesson): prove the door is LIVE rather
13210                    // than trusting that a flag set means a pool built.
13211                    eprintln!("[spec-vg] MTP verify-graph pool ENGAGED (vt_cap={vt_cap})");
13212                } else {
13213                    eprintln!(
13214                        "[spec-vg] MTP verify-graph pool declined (no linear layers, \
13215                         non-uniform state, or vt_cap < 2) — eager walk"
13216                    );
13217                }
13218            }
13219            Some(g)
13220        } else {
13221            None
13222        };
13223        // Capacity fail-safe: a round wider than the pool was built for must take the eager
13224        // walk, not slice the stash past its rows. The sizing above already covers every
13225        // round this run can present; this keeps a future caller (or a k that grows behind
13226        // the pool's back) on the byte-identical fallback instead of a panic.
13227        let vg_t_cap = vg_guard
13228            .as_ref()
13229            .and_then(|g| g.as_ref())
13230            .map(|g| g.t_capacity())
13231            .unwrap_or(0);
13232        if let Some(p) = pipe {
13233            p.setup_end();
13234        }
13235        let mut graph_guard_noted = false;
13236        while keep_going && out.len() < max_new {
13237            // GRAPH-LAUNCH HEADROOM GUARD (see GRAPH_LAUNCH_MIN_FREE): below the floor,
13238            // every captured-graph arm in this round yields to its byte-identical eager
13239            // twin instead of feeding cuGraphLaunch a card it segfaults on.
13240            let graph_round_ok = graph_launch_headroom_ok(e);
13241            if !graph_round_ok && !graph_guard_noted {
13242                graph_guard_noted = true;
13243                eprintln!(
13244                    "[spec] graph replay suspended: driver free below the {}MB launch floor \
13245                     (eager arms serve; cuGraphLaunch segfaults into an exhausted card)",
13246                    GRAPH_LAUNCH_MIN_FREE / (1 << 20)
13247                );
13248            }
13249            // MEMRA_SPEC_ROUND_PROF=1: wall of the WHOLE round against the pieces we already
13250            // instrument. Needed because the parts do not add up: the draft step measures 1.27 ms
13251            // ([spec-anatomy] glue 92 / attn 280 / ffn 222 / head 670 us) and the t=2 verify walk
13252            // 25.6 ms ([tcol-prof] attn 10.1 + ffn 15.3), yet a K=1 round takes 177 ms on the
13253            // step37 TP2 stack. This prints where the other ~150 ms lives.
13254            let round_prof = ROUND_PROF
13255                .get_or_init(|| std::env::var("MEMRA_SPEC_ROUND_PROF").as_deref() == Ok("1"));
13256            let round_t0 = round_prof.then(std::time::Instant::now);
13257            // ROUND-STREAM BURST: from round 1 (pending guaranteed by every non-replay arm),
13258            // issue M rounds with zero readbacks, then drain the ring + reconcile mirrors.
13259            if let (true, Some(sg), Some(ptrs)) = (
13260                stream_active && round >= 1 && pending.is_some() && graph_round_ok,
13261                &stream_graph,
13262                &stream_ptrs,
13263            ) {
13264                if debug_spec {
13265                    static ONCE: std::sync::Once = std::sync::Once::new();
13266                    ONCE.call_once(|| {
13267                        eprintln!("[memra] ROUND-STREAM burst engaged (M={m_rounds} k={k})")
13268                    });
13269                }
13270                e.set_i32_one(&mut pos_ctr, cache.pos as i32)?;
13271                e.set_u32_one(&mut pend_d, pending.unwrap())?;
13272                e.set_u32_one(&mut ring_d, 0)?; // ring count = 0 (writes element 0)
13273                for _mi in 0..m_rounds {
13274                    e.i32_copy_add(&pos_ctr, &mut pos_start_d, 0)?;
13275                    cache.snapshot_into(e, &mut snap)?; // device D2Ds, stream-ordered
13276                    e.i32_copy_add(&pos_ctr, &mut scratch.kv.len_d, 0)?; // draft-KV rollback
13277                    e.i32_copy_add(&pos_ctr, &mut dctx.g_pos, 1)?; // rope pos = pos + base
13278                    e.u32_copy(&pend_d, &mut dctx.g_tok)?;
13279                    e.copy_into(&mut dctx.g_seed, 0, &h_seed_buf, n_embd)?;
13280                    sg.launch()?;
13281                    e.spec_assemble_verify(
13282                        &g_tokp2k,
13283                        &pend_d,
13284                        d2t_dev.as_ref(),
13285                        &mut vtok_d,
13286                        &mut brk_d,
13287                        p_min,
13288                        k,
13289                        pmin0,
13290                    )?;
13291                    let mut ck = VerifyCkpt::new(self.layers.len());
13292                    let dummy = vec![0u32; t_v_s];
13293                    let (tl_d, vx) = self.decode_step_t_core_stream(
13294                        e,
13295                        &dummy,
13296                        0,
13297                        &mut *cache,
13298                        embd_dev,
13299                        Some(&mut ck),
13300                        Some((&vtok_d, &pos_ctr)),
13301                        None,
13302                        None,
13303                        None,
13304                    )?;
13305                    for j in 0..t_v_s {
13306                        e.argmax_token_device_col(&tl_d, j, n_vocab, &mut preds_d, j)?;
13307                    }
13308                    e.spec_accept_greedy_dc(
13309                        &preds_d,
13310                        &vtok_d,
13311                        &last_pred_d,
13312                        &brk_d,
13313                        &mut stream_acc,
13314                    )?;
13315                    e.spec_seed_gather(&vx, &fill_prev, &stream_acc, &mut h_seed_buf, 1, n_embd)?;
13316                    e.copy_into(&mut fill_prev, 0, &h_seed_buf, n_embd)?;
13317                    self.commit_verified_prefix_stream(
13318                        e,
13319                        &mut *cache,
13320                        &snap,
13321                        &ck,
13322                        &stream_acc,
13323                        1,
13324                        t_v_s,
13325                    )?;
13326                    e.spec_rollback_stream(
13327                        ptrs,
13328                        &pos_start_d,
13329                        &stream_acc,
13330                        1,
13331                        self.layers.len() + 1,
13332                    )?;
13333                    e.spec_ring_commit(&vtok_d, &stream_acc, &brk_d, &mut ring_d, &mut pend_d)?;
13334                }
13335                e.stream().synchronize()?;
13336                let ring_h = e.dtoh_u32(&ring_d)?;
13337                let cnt = ring_h[0] as usize;
13338                for i in 0..cnt {
13339                    if out.len() < max_new {
13340                        out.push(ring_h[1 + i]);
13341                    }
13342                }
13343                let pos_h = e.dtoh_i32(&pos_ctr)?[0] as usize;
13344                for il in 0..self.layers.len() {
13345                    if let Some(kvl) = cache.kv[il].as_mut() {
13346                        kvl.len = pos_h;
13347                    }
13348                }
13349                cache.pos = pos_h;
13350                scratch.kv.len = pos_h;
13351                pending = Some(ring_h[cnt]); // last drained token = the live bonus
13352                last_token = ring_h[cnt];
13353                total_drafted += k * m_rounds; // upper bound (p-min breaks uncounted)
13354                total_accepted += cnt.saturating_sub(m_rounds);
13355                if let Some(t) = sess_telem {
13356                    // totals only — the burst's per-round accept counts stayed on device
13357                    // (that is the point of the round-stream arm). pos_* untouched.
13358                    t.record_totals(m_rounds, k * m_rounds, cnt.saturating_sub(m_rounds));
13359                }
13360                round += m_rounds;
13361                // sse-cadence: the drained ring is committed — flush it at burst-drain cadence.
13362                keep_going = flush_commit(&mut on_commit, &out, &mut flushed);
13363                continue;
13364            }
13365            let pipe_draft = match pipe {
13366                Some(p) => Some(p.draft_begin(round)?),
13367                None => None,
13368            };
13369            let pos = cache.pos; // #tokens committed (EXCLUDES a pending bonus)
13370            let mut current_opti = carried_opti.take();
13371            let mut fork_generation = if current_opti.is_none() && pending.is_some() {
13372                match opti_fork.as_mut() {
13373                    Some(fork) if fork.mode.is_forced() => Some(fork.reserve(&mut snap)?),
13374                    None => None,
13375                    Some(_) => None,
13376                }
13377            } else {
13378                None
13379            };
13380            if current_opti.is_none() {
13381                if let Some(fork) = opti_fork.as_ref() {
13382                    opti_snapshot_stage_owned_into(e, cache, fork.rt, &fork.fence, &mut snap)?;
13383                } else {
13384                    cache.snapshot_into(e, &mut snap)?;
13385                }
13386            } else if snap.pos != pos {
13387                return Err(format!(
13388                    "optipipe carried snapshot pos {} != current pos {pos}",
13389                    snap.pos
13390                )
13391                .into());
13392            } // §C: snapshot BEFORE draft+verify (already retained for a carried successor)
13393            ph_mark(&mut ph_rest, phase_on);
13394
13395            // --- 1. DRAFT k tokens with the NextN head (autoregressive, T=1 each) ---
13396            // p-min semantics (both paths): stop the chain early when the head's confidence in
13397            // its own pick drops below p_min — the just-drafted token is DISCARDED, but its
13398            // scratch append stands (identical to the eager chain's ordering). j==0 always drafts.
13399            let base0 = if pending.is_some() { 1usize } else { 0usize };
13400            // fixed draft length by default; MEMRA_SPEC_ADAPT=1 drafts at last round's
13401            // accepted run + 1 (the gemma law — see the setup block above the loop).
13402            let k_this = if adapt { kc } else { k };
13403            let mut draft: Vec<u32> = Vec::with_capacity(k);
13404            let mut draft_idx: Vec<u32> = Vec::with_capacity(k); // trimmed-vocab ids (== draft when untrimmed)
13405            let mut controller_draft_prob: Option<f32> = None;
13406            let mut controller_eager_state: Option<(u32, CudaSlice<f32>)> = None;
13407            if let Some(ticket) = current_opti.as_mut() {
13408                let carried_pending = pending.ok_or("optipipe carried successor lost pending")?;
13409                if ticket.verify_tokens[0] != carried_pending {
13410                    return Err(format!(
13411                        "optipipe carried pending mismatch: ticket={} live={carried_pending}",
13412                        ticket.verify_tokens[0],
13413                    )
13414                    .into());
13415                }
13416                draft.push(ticket.verify_tokens[1]);
13417                controller_draft_prob = Some(ticket.draft_prob);
13418                controller_eager_state = ticket
13419                    .take_eager_seed()
13420                    .map(|seed| (ticket.verify_tokens[1], seed));
13421            } else {
13422                // Round-start draft-KV sync (BOTH paths). Persistent: truncate/align to the committed
13423                // history — slots 0..P hold entries for the tokens before last_token@P (P = pos +
13424                // base0 - 1); this single set_len IS the draft-side rollback (drops last round's
13425                // rejected drafts and p-min extras via the len mechanism).
13426                scratch.set_len(e, pos + base0 - 1)?;
13427                // dcw door: a captured chain appends k_this device-counter rows (plus the
13428                // pseudo-seed replay) with no host intervention; any ring rebase those appends
13429                // could need happens HERE, host-side, before the replays. The eager arm keeps
13430                // its own per-step prepare, so this is graph-path-only work.
13431                if step35_draft_dcw_on()
13432                    && (dctx.graph.is_some()
13433                        || dctx.graph_s.is_some()
13434                        || dctx.chain.is_some()
13435                        || dctx.chain_s.is_some())
13436                {
13437                    scratch.ensure_dcw_headroom(e, k_this + 2)?;
13438                }
13439                if pen_on {
13440                    // PEN_WINDOW_MAX also bounds the per-round upload and the O(n_hist^2)
13441                    // device dedup: the serve window is already PEN_WINDOW_MAX, and this
13442                    // defensive min also bounds non-server callers.
13443                    let win = sp.penalty_last_n.min(PEN_WINDOW_MAX);
13444                    let w0 = pen_hist.len().saturating_sub(win);
13445                    pen_hist_d = Some(e.htod_u32_v(&pen_hist[w0..])?);
13446                }
13447                if sampled {
13448                    draft_logits.clear();
13449                    draft_stats.clear();
13450                }
13451                // DRAFT-SIDE GRAMMAR MASK: clone the committed grammar state ONCE per round; each
13452                // position's mask is computed on that clone and advanced by the PROPOSED token. The
13453                // real state moves only on emission (verify's job), so the emitted stream is
13454                // unchanged — the mask only removes tokens the verify would have truncated anyway.
13455                let mut dmask_live = dmask_on;
13456                if dmask_live {
13457                    let t_c = std::time::Instant::now();
13458                    constraint
13459                        .as_deref_mut()
13460                        .unwrap()
13461                        .draft_begin()
13462                        .map_err(|e2| format!("constraint: {e2}"))?;
13463                    dm_clone_ns += t_c.elapsed().as_nanos();
13464                    dm_rounds += 1;
13465                }
13466                if let (false, Some(cg)) = (sampled || pen_on || !graph_round_ok, &dctx.chain) {
13467                    // GREEDY CHAIN GRAPH (lane/step37-draft-graph-serving-20260830): the
13468                    // eager multi-head chain's EXACT launch order — step j rewinds head
13469                    // (j % heads)'s plane to the committed length and replays rows 0..=j —
13470                    // with each row's whole head-forward as ONE graph launch. The chain
13471                    // POLICY (head choice, prefix length, stored-seed feed) is host-side,
13472                    // identical to `mtp_chain_forward_dev`, so graph-vs-eager drafts are
13473                    // bit-identical by construction (same launcher, same bucket — the dcw
13474                    // parity contract). Interior rows launch the head-less graph: their
13475                    // logits are dead in the eager chain too, so the consumed bytes match.
13476                    let heads_n = self.mtp_head_count();
13477                    let committed = pos + base0 - 1;
13478                    let mut chain_tokens: Vec<u32> = vec![last_token];
13479                    let mut chain_seed_bufs: Vec<CudaSlice<f32>> = vec![e.clone_dtod(&h_seed_buf)?];
13480                    for j in 0..k_this {
13481                        let index = mtp_chain_head_index(j, heads_n);
13482                        if debug_spec {
13483                            eprintln!(
13484                                "[mtp-chain-step] round={round} j={j} head={index} \
13485                                 replay_rows={} arm=graph",
13486                                chain_tokens.len(),
13487                            );
13488                        }
13489                        scratch.set_plane_len(e, index, committed)?;
13490                        e.set_i32_one(&mut dctx.g_pos, (committed + 1) as i32)?;
13491                        for row in 0..=j {
13492                            e.set_u32_one(&mut dctx.g_tok, chain_tokens[row])?;
13493                            e.copy_into(&mut dctx.g_seed, 0, &chain_seed_bufs[row], n_embd)?;
13494                            if row < j {
13495                                cg.interior[index].launch()?;
13496                            } else {
13497                                // per-position mask upload before the LAST row only — the
13498                                // eager chain applies the mask on is_last exactly the same.
13499                                if dmask_live
13500                                    && !upload_draft_mask(
13501                                        e,
13502                                        constraint.as_deref_mut().unwrap(),
13503                                        &mut dctx.g_dmask,
13504                                        mtp.d2t.as_ref(),
13505                                        d_vocab,
13506                                        dmask_words,
13507                                    )?
13508                                {
13509                                    e.htod_u32_into(
13510                                        &mut dctx.g_dmask,
13511                                        &vec![u32::MAX; dmask_words],
13512                                    )?;
13513                                    dmask_live = false;
13514                                }
13515                                cg.last[index].launch()?;
13516                            }
13517                            // host mirror (len_d advanced in-graph by the dcw append)
13518                            scratch.plane_mut(index).0.len += 1;
13519                        }
13520                        let idx = e.dtoh_u32_one(&dctx.g_tok)?;
13521                        // #87 SENTINEL TRAP (see the single-head graph arm below).
13522                        if (idx as usize) >= d_vocab {
13523                            let seed_h = e.dtoh(&dctx.g_seed)?;
13524                            let seed_nan = seed_h.iter().filter(|v| v.is_nan()).count();
13525                            return Err(format!(
13526                                "draft(chain-graph) argmax sentinel 0x{idx:08x} >= d_vocab \
13527                             {d_vocab} at round {round} j={j} head={index} pos={pos}: \
13528                             head-out NaN {seed_nan}/{n_embd} — refusing to dereference \
13529                             the embed row (#87 trap)"
13530                            )
13531                            .into());
13532                        }
13533                        // multi-head MTP forbids a trimmed head (validated at entry), so the
13534                        // draft index IS the target id; keep the map for uniformity.
13535                        let d = match &mtp.d2t {
13536                            Some(map) => map[idx as usize],
13537                            None => idx,
13538                        };
13539                        let draft_p = if p_min > 0.0 {
13540                            Some(e.dtoh(&dctx.g_p)?[0])
13541                        } else {
13542                            None
13543                        };
13544                        if j == 0 {
13545                            controller_draft_prob = draft_p;
13546                        }
13547                        if let Some(p) = draft_p.filter(|_| p_min > 0.0) {
13548                            if p < p_min && (j > 0 || (pmin0 && base0 == 1)) {
13549                                break;
13550                            }
13551                        }
13552                        draft.push(d);
13553                        chain_tokens.push(d);
13554                        // step j's h_nextn: the last-row graph self-fed it into g_seed —
13555                        // snapshot it as the chain history seed for row j+1 (stream-ordered
13556                        // after the launch, exactly the eager chain's chain_seeds push).
13557                        chain_seed_bufs.push(e.clone_dtod(&dctx.g_seed)?);
13558                        // speculative grammar advance (see the single-head graph arm).
13559                        if dmask_live
13560                            && !constraint
13561                                .as_deref_mut()
13562                                .unwrap()
13563                                .draft_advance(d)
13564                                .map_err(|e2| format!("constraint: {e2}"))?
13565                        {
13566                            e.htod_u32_into(&mut dctx.g_dmask, &vec![u32::MAX; dmask_words])?;
13567                            break;
13568                        }
13569                    }
13570                } else if let (true, Some(cg)) = (
13571                    sampled && s_capturable && dctx.s_key == Some(s_key) && graph_round_ok,
13572                    &dctx.chain_s,
13573                ) {
13574                    if skey_probe() {
13575                        eprintln!(
13576                            "[skey] chain=graph_chain_s round={round} capturable={} top_k={} \
13577                             top_p={} min_p={} s_key_parked={:?}",
13578                            s_capturable as u8, sp.top_k, sp.top_p, sp.min_p, dctx.s_key,
13579                        );
13580                    }
13581                    // SAMPLED CHAIN GRAPH: the greedy chain arm's launch order with the
13582                    // sampled last-row graphs — in-graph counter bump + (filtered) gumbel
13583                    // draw + argmax; q retained per step into q_slots exactly like the
13584                    // single-head sampled graph arm. Counter continuity: g_ctr host-seeded
13585                    // to sctr-1 once per ROUND; each step's last-row graph bumps it BEFORE
13586                    // the perturb, so step j consumes counter sctr+j — the eager Philox
13587                    // stream (interior rows never draw, never bump).
13588                    let heads_n = self.mtp_head_count();
13589                    let committed = pos + base0 - 1;
13590                    let filtered_stats_in_graph = s_key.filtered();
13591                    let mut chain_tokens: Vec<u32> = vec![last_token];
13592                    let mut chain_seed_bufs: Vec<CudaSlice<f32>> = vec![e.clone_dtod(&h_seed_buf)?];
13593                    e.set_u32_one(&mut dctx.g_ctr, sctr.wrapping_sub(1))?;
13594                    for j in 0..k_this {
13595                        let index = mtp_chain_head_index(j, heads_n);
13596                        if debug_spec {
13597                            eprintln!(
13598                                "[mtp-chain-step] round={round} j={j} head={index} \
13599                                 replay_rows={} arm=graph_s",
13600                                chain_tokens.len(),
13601                            );
13602                        }
13603                        scratch.set_plane_len(e, index, committed)?;
13604                        e.set_i32_one(&mut dctx.g_pos, (committed + 1) as i32)?;
13605                        for row in 0..=j {
13606                            e.set_u32_one(&mut dctx.g_tok, chain_tokens[row])?;
13607                            e.copy_into(&mut dctx.g_seed, 0, &chain_seed_bufs[row], n_embd)?;
13608                            if row < j {
13609                                cg.interior[index].launch()?;
13610                            } else {
13611                                cg.last[index].launch()?;
13612                            }
13613                            scratch.plane_mut(index).0.len += 1;
13614                        }
13615                        sctr += 1; // mirrors the in-graph g_ctr bump (eager parity:
13616                        // counts the p-min-discarded token too)
13617                        // q retention: ONE async D2D of the persistent head-logits buffer
13618                        // into this round's slot j (stream-ordered after the replay).
13619                        e.copy_into(&mut dctx.q_slots[j], 0, &dctx.g_q, d_vocab)?;
13620                        // FILTERED capture: read the in-graph filter_stats scalars back per
13621                        // replay instead of a second full-vocab filter_stats per slot post-
13622                        // chain — bit-exact (the values the in-graph perturb consumed) and
13623                        // measured worth ~5% of vendor-default serving tok/s at K=3. Before
13624                        // the p-min break so the discarded slot's stats land too.
13625                        if filtered_stats_in_graph {
13626                            draft_stats.push((
13627                                e.dtoh(&dctx.g_mx)?[0],
13628                                e.dtoh(&dctx.g_th)?[0],
13629                                e.dtoh(&dctx.g_z)?[0],
13630                            ));
13631                        }
13632                        let idx = e.dtoh_u32_one(&dctx.g_tok)?;
13633                        // #87 SENTINEL TRAP (see the single-head graph arms).
13634                        if (idx as usize) >= d_vocab {
13635                            let seed_h = e.dtoh(&dctx.g_seed)?;
13636                            let seed_nan = seed_h.iter().filter(|v| v.is_nan()).count();
13637                            return Err(format!(
13638                                "draft(chain-graph-sampled) argmax sentinel 0x{idx:08x} >= \
13639                             d_vocab {d_vocab} at round {round} j={j} head={index} pos={pos}: \
13640                             head-out NaN {seed_nan}/{n_embd} — refusing to dereference the \
13641                             embed row (#87 trap)"
13642                            )
13643                            .into());
13644                        }
13645                        let d = match &mtp.d2t {
13646                            Some(map) => map[idx as usize],
13647                            None => idx,
13648                        };
13649                        draft_idx.push(idx);
13650                        if p_min > 0.0 {
13651                            let p = e.dtoh(&dctx.g_p)?[0];
13652                            if p < p_min && (j > 0 || (pmin0 && base0 == 1)) {
13653                                break;
13654                            }
13655                        }
13656                        draft.push(d);
13657                        chain_tokens.push(d);
13658                        chain_seed_bufs.push(e.clone_dtod(&dctx.g_seed)?);
13659                    }
13660                    // PURE-TEMP accept path: stats per used slot recomputed from the RETAINED
13661                    // q with the SAME filter_stats program the eager arm runs (deployment-
13662                    // keyed coop/plain choice, same input bits). The FILTERED graph read its
13663                    // stats back per replay above.
13664                    if !filtered_stats_in_graph {
13665                        for j in 0..draft.len().max(draft_idx.len()) {
13666                            let rows0 = e.htod_i32(&[0])?;
13667                            let (mut th_d, mut z_d, mut mx_d) =
13668                                (e.zeros(1)?, e.zeros(1)?, e.zeros(1)?);
13669                            e.filter_stats(
13670                                &dctx.q_slots[j],
13671                                d_vocab,
13672                                &rows0,
13673                                &mut th_d,
13674                                &mut z_d,
13675                                &mut mx_d,
13676                                d_vocab,
13677                                1,
13678                                sp_temp,
13679                                sp.top_k,
13680                                sp.top_p,
13681                                sp.min_p,
13682                            )?;
13683                            draft_stats.push((
13684                                e.dtoh(&mx_d)?[0],
13685                                e.dtoh(&th_d)?[0],
13686                                e.dtoh(&z_d)?[0],
13687                            ));
13688                        }
13689                    }
13690                } else if let (false, Some(gr)) =
13691                    (sampled || pen_on || !graph_round_ok, &dctx.graph)
13692                {
13693                    // GRAPH DRAFT: one dispatch per drafted token. The chain feeds itself on-device
13694                    // (in-graph argmax -> tok_d -> next replay's embed; h_nextn -> h_seed_d; pos_d
13695                    // inc'd in-graph); the host only reads 4B token (+4B p) and decides the break.
13696                    e.set_i32_one(&mut dctx.g_pos, (pos + base0) as i32)?;
13697                    e.set_u32_one(&mut dctx.g_tok, last_token)?;
13698                    e.copy_into(&mut dctx.g_seed, 0, &h_seed_buf, n_embd)?;
13699                    for j in 0..k_this {
13700                        // per-position mask upload (contents only — the graph's baked pointer is
13701                        // dctx.g_dmask). All-ones once masking goes dead mid-chain, so the captured
13702                        // mask node degrades to a no-op ban instead of needing a second graph.
13703                        if dmask_live
13704                            && !upload_draft_mask(
13705                                e,
13706                                constraint.as_deref_mut().unwrap(),
13707                                &mut dctx.g_dmask,
13708                                mtp.d2t.as_ref(),
13709                                d_vocab,
13710                                dmask_words,
13711                            )?
13712                        {
13713                            // no draft-vocab row is grammar-legal here (a trimmed FR-Spec head can
13714                            // genuinely miss the legal set): neutralize the captured mask node and
13715                            // finish the chain UNMASKED — exactly pre-lane behaviour, never worse.
13716                            e.htod_u32_into(&mut dctx.g_dmask, &vec![u32::MAX; dmask_words])?;
13717                            dmask_live = false;
13718                        }
13719                        gr.launch()?;
13720                        scratch.kv.len += 1; // host mirror (len_d advanced in-graph)
13721                        let idx = e.dtoh_u32_one(&dctx.g_tok)?;
13722                        // #87 SENTINEL TRAP: an all-NaN head-logits row leaves the device argmax's
13723                        // init sentinel (0x7FFFFFFF) in g_tok — feeding it onward dereferences
13724                        // embed_row(sentinel) = table + ~4.6TB (never mapped) inside the NEXT graph
13725                        // replay's embed node, and the MMU fault kills the CUDA context for the
13726                        // whole process (research/pp2spec-crash-20260807: 3 coredumps, byte-exact
13727                        // VA arithmetic). Refuse loudly instead; the diagnostics name the first-NaN
13728                        // buffer (g_seed = the verify-side handoff vs head-side compute).
13729                        if (idx as usize) >= d_vocab {
13730                            // g_seed is SELF-FED (the replay writes h_nextn back into it), so it
13731                            // reads as the head's OUTPUT at j; h_seed_buf is the round's INPUT
13732                            // seed, untouched since the round-start copy — the pair discriminates
13733                            // "seed arrived poisoned" from "head forward produced NaN".
13734                            let seed_h = e.dtoh(&dctx.g_seed)?;
13735                            let seed_nan = seed_h.iter().filter(|v| v.is_nan()).count();
13736                            let in_h = e.dtoh(&h_seed_buf)?;
13737                            let in_nan = in_h.iter().filter(|v| v.is_nan()).count();
13738                            return Err(format!(
13739                                "draft(graph) argmax sentinel 0x{idx:08x} >= d_vocab {d_vocab} at \
13740                             round {round} j={j} pos={pos}: head-out NaN {seed_nan}/{n_embd}, \
13741                             round-input-seed NaN {in_nan}/{n_embd} — refusing to dereference \
13742                             the embed row (#87 trap)"
13743                            )
13744                            .into());
13745                        }
13746                        // trimmed draft vocab -> target token id (identity when no d2t map)
13747                        let d = match &mtp.d2t {
13748                            Some(map) => map[idx as usize],
13749                            None => idx,
13750                        };
13751                        let draft_p = if p_min > 0.0
13752                            || opti_fork
13753                                .as_ref()
13754                                .is_some_and(|fork| fork.controller.is_some())
13755                        {
13756                            Some(e.dtoh(&dctx.g_p)?[0])
13757                        } else {
13758                            None
13759                        };
13760                        if j == 0 {
13761                            controller_draft_prob = draft_p;
13762                        }
13763                        if let Some(p) = draft_p.filter(|_| p_min > 0.0) {
13764                            if p < p_min && (j > 0 || (pmin0 && base0 == 1)) {
13765                                break;
13766                            }
13767                        }
13768                        draft.push(d);
13769                        // with a trimmed head the NEXT embed must read the TARGET id, not the draft
13770                        // index the argmax wrote — patch the persistent token buffer (4B htod).
13771                        if d != idx {
13772                            e.set_u32_one(&mut dctx.g_tok, d)?;
13773                        }
13774                        // advance the SPECULATIVE state with the proposal; a dead chain drops to
13775                        // unmasked drafting for the remaining positions (verify still arbitrates).
13776                        // speculative advance; a chain the grammar can no longer follow (EOS
13777                        // proposed) ends here. The captured mask node always runs, so a dead chain
13778                        // leaves the buffer NEUTRAL (all-ones = ban nothing) before it exits.
13779                        if dmask_live
13780                            && !constraint
13781                                .as_deref_mut()
13782                                .unwrap()
13783                                .draft_advance(d)
13784                                .map_err(|e2| format!("constraint: {e2}"))?
13785                        {
13786                            e.htod_u32_into(&mut dctx.g_dmask, &vec![u32::MAX; dmask_words])?;
13787                            break;
13788                        }
13789                    }
13790                // REGIME RE-TEST (lane/graph-s-key-exactness-20260819, widened by
13791                // lane/step37-draft-graph-serving-20260830): the sampled graph is legal ONLY
13792                // in the regime it was captured in. The condition used to read
13793                // `(sampled, &dctx.graph_s)` and trusted `s_key` to have dropped anything
13794                // else — which it could not, because the key omitted the filters. Both
13795                // halves are enforced: the key drops a stale graph, and this site refuses to
13796                // launch one whose key differs or whose regime is uncapturable (penalties).
13797                } else if let (true, Some(gr)) = (
13798                    sampled && s_capturable && dctx.s_key == Some(s_key) && graph_round_ok,
13799                    &dctx.graph_s,
13800                ) {
13801                    if skey_probe() {
13802                        eprintln!(
13803                            "[skey] chain=graph_s round={round} pure_temp={} capturable={} \
13804                             top_k={} top_p={} min_p={} s_key_parked={:?}",
13805                            pure_temp as u8,
13806                            s_capturable as u8,
13807                            sp.top_k,
13808                            sp.top_p,
13809                            sp.min_p,
13810                            dctx.s_key,
13811                        );
13812                    }
13813                    // SAMPLED GRAPH DRAFT: one replay per drafted token — head forward + gumbel +
13814                    // argmax in ONE dispatch; the host reads 4B token (+4B p), D2Ds q into slot j,
13815                    // and decides the break. Event-counter continuity: g_ctr is host-seeded to
13816                    // sctr-1 ONCE per round (outside the graph); the in-graph bump runs BEFORE the
13817                    // perturb, so replay j consumes counter sctr+j — exactly the eager arm's Philox
13818                    // stream. Host sctr advances in lockstep (computed, no readback needed).
13819                    e.set_i32_one(&mut dctx.g_pos, (pos + base0) as i32)?;
13820                    e.set_u32_one(&mut dctx.g_tok, last_token)?;
13821                    e.copy_into(&mut dctx.g_seed, 0, &h_seed_buf, n_embd)?;
13822                    e.set_u32_one(&mut dctx.g_ctr, sctr.wrapping_sub(1))?;
13823                    let filtered_stats_in_graph = s_key.filtered();
13824                    for j in 0..k_this {
13825                        gr.launch()?;
13826                        scratch.kv.len += 1; // host mirror (len_d advanced in-graph)
13827                        sctr += 1; // mirrors the in-graph g_ctr bump (eager parity:
13828                        // counts the p-min-discarded token too)
13829                        // q retention: ONE async D2D of the persistent head-logits buffer into this
13830                        // round's slot j (stream-ordered after the replay, before the next one).
13831                        e.copy_into(&mut dctx.q_slots[j], 0, &dctx.g_q, d_vocab)?;
13832                        // FILTERED capture: the replay's own filter_stats node already computed
13833                        // (th, z, mx) — read the three scalars back instead of paying a SECOND
13834                        // full-vocab filter_stats per slot post-chain (measured ~5% of vendor-
13835                        // default serving tok/s at K=3). Bit-exact by construction: these are
13836                        // the very values the in-graph perturb consumed. Read BEFORE the p-min
13837                        // break so the discarded slot's stats land too (accept-path indexing).
13838                        if filtered_stats_in_graph {
13839                            draft_stats.push((
13840                                e.dtoh(&dctx.g_mx)?[0],
13841                                e.dtoh(&dctx.g_th)?[0],
13842                                e.dtoh(&dctx.g_z)?[0],
13843                            ));
13844                        }
13845                        let idx = e.dtoh_u32_one(&dctx.g_tok)?;
13846                        // #87 SENTINEL TRAP (see the greedy graph arm above).
13847                        if (idx as usize) >= d_vocab {
13848                            let seed_h = e.dtoh(&dctx.g_seed)?;
13849                            let seed_nan = seed_h.iter().filter(|v| v.is_nan()).count();
13850                            return Err(format!(
13851                                "draft(graph-sampled) argmax sentinel 0x{idx:08x} >= d_vocab \
13852                             {d_vocab} at round {round} j={j} pos={pos}: round-seed NaN \
13853                             {seed_nan}/{n_embd} — refusing to dereference the embed row \
13854                             (#87 trap)"
13855                            )
13856                            .into());
13857                        }
13858                        let d = match &mtp.d2t {
13859                            Some(map) => map[idx as usize],
13860                            None => idx,
13861                        };
13862                        draft_idx.push(idx);
13863                        if p_min > 0.0 {
13864                            let p = e.dtoh(&dctx.g_p)?[0];
13865                            if p < p_min && (j > 0 || (pmin0 && base0 == 1)) {
13866                                break;
13867                            }
13868                        }
13869                        draft.push(d);
13870                        // trimmed head: the NEXT embed must read the TARGET id (see the greedy arm).
13871                        if d != idx {
13872                            e.set_u32_one(&mut dctx.g_tok, d)?;
13873                        }
13874                    }
13875                    // PURE-TEMP accept path: fill draft_stats per used slot post-chain (the
13876                    // stats degenerate to th=0 / full-Z; one filter_stats launch per slot).
13877                    // The FILTERED graph read its stats back per replay above.
13878                    if !filtered_stats_in_graph {
13879                        for j in 0..draft.len().max(draft_idx.len()) {
13880                            let rows0 = e.htod_i32(&[0])?;
13881                            let (mut th_d, mut z_d, mut mx_d) =
13882                                (e.zeros(1)?, e.zeros(1)?, e.zeros(1)?);
13883                            e.filter_stats(
13884                                &dctx.q_slots[j],
13885                                d_vocab,
13886                                &rows0,
13887                                &mut th_d,
13888                                &mut z_d,
13889                                &mut mx_d,
13890                                d_vocab,
13891                                1,
13892                                sp_temp,
13893                                sp.top_k,
13894                                sp.top_p,
13895                                sp.min_p,
13896                            )?;
13897                            draft_stats.push((
13898                                e.dtoh(&mx_d)?[0],
13899                                e.dtoh(&th_d)?[0],
13900                                e.dtoh(&z_d)?[0],
13901                            ));
13902                        }
13903                    }
13904                } else {
13905                    if skey_probe() && sampled {
13906                        eprintln!(
13907                            "[skey] chain=eager round={round} pure_temp={} top_k={} \
13908                             top_p={} min_p={} s_key_parked={:?}",
13909                            pure_temp as u8, sp.top_k, sp.top_p, sp.min_p, dctx.s_key,
13910                        );
13911                    }
13912                    // EAGER DRAFT (fallback: MoE head/trunk, huge k, MEMRA_SPEC_NOGRAPH, capture fail).
13913                    let chain_heads = !self.mtp_extra.is_empty();
13914                    let mut e_tok = last_token;
13915                    let mut d_seed = e.clone_dtod(&h_seed_buf)?;
13916                    let mut chain_tokens = if chain_heads {
13917                        vec![last_token]
13918                    } else {
13919                        Vec::new()
13920                    };
13921                    let mut chain_seeds = if chain_heads {
13922                        vec![e.clone_dtod(&h_seed_buf)?]
13923                    } else {
13924                        Vec::new()
13925                    };
13926                    for j in 0..k_this {
13927                        // GPU-ARGMAX DRAFT (2026-07-03): device logits + device argmax + 4-byte token
13928                        // read instead of the ~600KB full-vocab dtoh + host argmax per draft token.
13929                        let mtp_pos = pos + base0 + j;
13930                        // draft-side grammar mask (eager twin of the graph arm's in-graph node).
13931                        // A position with no legal draft-vocab row drops to unmasked drafting for
13932                        // the rest of the chain (pre-lane behaviour; verify still arbitrates).
13933                        if dmask_live {
13934                            dmask_live = upload_draft_mask(
13935                                e,
13936                                constraint.as_deref_mut().unwrap(),
13937                                &mut dctx.g_dmask,
13938                                mtp.d2t.as_ref(),
13939                                d_vocab,
13940                                dmask_words,
13941                            )?;
13942                        }
13943                        let mask = if dmask_live {
13944                            Some((&dctx.g_dmask, dmask_words))
13945                        } else {
13946                            None
13947                        };
13948                        let (dl_d, h_nextn) = if chain_heads {
13949                            if debug_spec {
13950                                eprintln!(
13951                                    "[mtp-chain-step] round={round} j={j} head={} replay_rows={}",
13952                                    mtp_chain_head_index(j, self.mtp_head_count()),
13953                                    chain_tokens.len(),
13954                                );
13955                            }
13956                            self.mtp_chain_forward_dev(
13957                                e,
13958                                &chain_tokens,
13959                                &chain_seeds,
13960                                &mut *scratch,
13961                                pos + base0 - 1,
13962                                embd_dev,
13963                                mask,
13964                            )?
13965                        } else {
13966                            self.mtp_head_forward_dev(
13967                                e,
13968                                mtp,
13969                                e_tok,
13970                                &d_seed,
13971                                &mut *scratch,
13972                                mtp_pos,
13973                                embd_dev,
13974                                mask,
13975                            )?
13976                        };
13977                        let tok_d = if sampled {
13978                            // FILTERED Gumbel-max: stats -> masked perturb -> argmax = one draw from
13979                            // the filtered softmax (filters off => th=0, exact v1 semantics).
13980                            if perturb_buf.is_none() {
13981                                perturb_buf = Some(e.zeros(d_vocab.max(n_vocab))?);
13982                            }
13983                            let mut q_row = e.clone_dtod(&dl_d)?; // retained q (penalized when on)
13984                            if pen_on {
13985                                let h = pen_hist_d.as_ref().unwrap();
13986                                let nh = h.len();
13987                                e.penalize_logits(
13988                                    &mut q_row,
13989                                    h,
13990                                    nh,
13991                                    sp.penalty_repeat,
13992                                    sp.penalty_freq,
13993                                    sp.penalty_present,
13994                                    d_vocab,
13995                                )?;
13996                            }
13997                            let rows0 = e.htod_i32(&[0])?;
13998                            let (mut th_d, mut z_d, mut mx_d) =
13999                                (e.zeros(1)?, e.zeros(1)?, e.zeros(1)?);
14000                            e.filter_stats(
14001                                &q_row, d_vocab, &rows0, &mut th_d, &mut z_d, &mut mx_d, d_vocab,
14002                                1, sp_temp, sp.top_k, sp.top_p, sp.min_p,
14003                            )?;
14004                            let (th, z, mx) =
14005                                (e.dtoh(&th_d)?[0], e.dtoh(&z_d)?[0], e.dtoh(&mx_d)?[0]);
14006                            let pb = perturb_buf.as_mut().unwrap();
14007                            e.gumbel_perturb_filtered(
14008                                &q_row, pb, d_vocab, sp_seed, sctr, sp_temp, mx, th,
14009                            )?;
14010                            sctr += 1;
14011                            draft_logits.push(q_row);
14012                            draft_stats.push((mx, th, z));
14013                            e.argmax_token_device(pb, d_vocab)?
14014                        } else {
14015                            e.argmax_token_device(&dl_d, d_vocab)?
14016                        };
14017                        let idx = e.dtoh_u32_one(&tok_d)?;
14018                        // #87 SENTINEL TRAP (eager twin — see the graph arm). Extra diagnostics
14019                        // here because the eager chain's operands are all readable: dl_d (the head
14020                        // logits row) and d_seed (this step's h_seed) name the first-NaN buffer.
14021                        if (idx as usize) >= d_vocab {
14022                            let dl_h = e.dtoh(&dl_d)?;
14023                            let dl_nan = dl_h.iter().filter(|v| v.is_nan()).count();
14024                            let seed_h = if chain_heads {
14025                                e.dtoh(chain_seeds.last().unwrap())?
14026                            } else {
14027                                e.dtoh(&d_seed)?
14028                            };
14029                            let seed_nan = seed_h.iter().filter(|v| v.is_nan()).count();
14030                            return Err(format!(
14031                                "draft(eager) argmax sentinel 0x{idx:08x} >= d_vocab {d_vocab} at \
14032                             round {round} j={j} pos={pos}: head-logits NaN {dl_nan}/{d_vocab}, \
14033                             step-seed NaN {seed_nan}/{n_embd} — refusing to dereference the \
14034                             embed row (#87 trap)"
14035                            )
14036                            .into());
14037                        }
14038                        let d = match &mtp.d2t {
14039                            Some(map) => map[idx as usize],
14040                            None => idx,
14041                        };
14042                        if sampled {
14043                            draft_idx.push(idx);
14044                        }
14045                        let draft_p = if p_min > 0.0
14046                            || opti_fork
14047                                .as_ref()
14048                                .is_some_and(|fork| fork.controller.is_some())
14049                        {
14050                            let p_d = e.prob_of_token_device(&dl_d, &tok_d, d_vocab)?;
14051                            Some(e.dtoh(&p_d)?[0])
14052                        } else {
14053                            None
14054                        };
14055                        if j == 0 {
14056                            controller_draft_prob = draft_p;
14057                        }
14058                        if let Some(p) = draft_p.filter(|_| p_min > 0.0) {
14059                            if p < p_min && (j > 0 || (pmin0 && base0 == 1)) {
14060                                break;
14061                            }
14062                        }
14063                        draft.push(d);
14064                        if chain_heads {
14065                            chain_tokens.push(d);
14066                            chain_seeds.push(h_nextn);
14067                        } else {
14068                            e_tok = d;
14069                            d_seed = h_nextn;
14070                        }
14071                        // speculative advance; a chain the grammar can no longer follow (EOS
14072                        // proposed) ends here — the prefix already proposed still rides verify.
14073                        if dmask_live
14074                            && !constraint
14075                                .as_deref_mut()
14076                                .unwrap()
14077                                .draft_advance(d)
14078                                .map_err(|e2| format!("constraint: {e2}"))?
14079                        {
14080                            break;
14081                        }
14082                    }
14083                    if !chain_heads
14084                        && opti_fork
14085                            .as_ref()
14086                            .is_some_and(|fork| fork.controller.is_some())
14087                    {
14088                        controller_eager_state = Some((e_tok, d_seed));
14089                    }
14090                }
14091            }
14092            let k_round = draft.len();
14093            if let Some(p) = pipe {
14094                p.draft_end(round);
14095            }
14096            drop(pipe_draft);
14097
14098            ph_mark(&mut ph_draft, phase_on);
14099            // --- 2. VERIFY: one batched target forward. With a pending bonus, it rides as col 0
14100            //         (committing its KV/recur inside the SAME weight read); drafts follow. ---
14101            let verify_tokens: Vec<u32> = match pending {
14102                Some(b) => {
14103                    let mut v = Vec::with_capacity(k_round + 1);
14104                    v.push(b);
14105                    v.extend_from_slice(&draft);
14106                    v
14107                }
14108                None => draft.clone(),
14109            };
14110            let base = if pending.is_some() { 1 } else { 0 };
14111            // ckpt (REPLAY-FREE partial accept): retain per-layer state-rebuild inputs alongside
14112            // the verify. Pure buffer keep-alives + dtod clones — kernel work is unchanged.
14113            let mut ckpt = if let Some(ticket) = current_opti.as_mut() {
14114                Some(ticket.take_ckpt())
14115            } else if spec_replay {
14116                None
14117            } else {
14118                Some(VerifyCkpt::new(self.layers.len()))
14119            };
14120            let controller_can_probe = base == 1
14121                && k_round == 1
14122                && out.len().saturating_add(2) < max_new
14123                && controller_draft_prob.is_some()
14124                && opti_fork
14125                    .as_ref()
14126                    .and_then(|fork| fork.controller.as_ref())
14127                    .is_some_and(|policy| !policy.breaker_tripped);
14128            let mut successor_attempt: Option<OptiControllerTicket> = None;
14129            let mut rejected_probe: Option<(f32, u32)> = None;
14130            let mut controller_prepared: Option<OptiControllerPrepared> = None;
14131            if controller_can_probe {
14132                // Prepare d2/q and, on admission, d3 before either current verify half is
14133                // issued. N stage 0 can then be followed immediately by N+1 stage 0; once N's
14134                // boundary fires, those dev0 launches overlap N stage 1 on dev1. Preparing on
14135                // the primary stream after N stage 1 would serialize the supposed pipeline.
14136                let eager_pos = scratch.kv.len + 1;
14137                let (optimistic_pending, pending_probability) = self.opti_controller_draft_step(
14138                    e,
14139                    mtp,
14140                    &mut dctx,
14141                    &mut *scratch,
14142                    d_vocab,
14143                    &mut controller_eager_state,
14144                    eager_pos,
14145                    embd_dev,
14146                    graph_round_ok,
14147                )?;
14148                let first_probability = controller_draft_prob
14149                    .ok_or("optipipe controller probe lost first-token probability")?;
14150                let q_proxy = first_probability * pending_probability;
14151                OPTI_GATE_CHECKS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
14152                OPTI_SHADOW_DRAFT_TOKENS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
14153                let admitted = opti_fork
14154                    .as_ref()
14155                    .and_then(|fork| fork.controller.as_ref())
14156                    .ok_or("optipipe controller policy disappeared")?
14157                    .admit(q_proxy);
14158                if admitted {
14159                    OPTI_GATE_ADMITS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
14160                    let eager_pos = scratch.kv.len + 1;
14161                    let (optimistic_draft, optimistic_draft_probability) = self
14162                        .opti_controller_draft_step(
14163                            e,
14164                            mtp,
14165                            &mut dctx,
14166                            &mut *scratch,
14167                            d_vocab,
14168                            &mut controller_eager_state,
14169                            eager_pos,
14170                            embd_dev,
14171                            graph_round_ok,
14172                        )?;
14173                    OPTI_SHADOW_DRAFT_TOKENS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
14174                    let eager_seed = controller_eager_state.take().map(|(token, seed)| {
14175                        debug_assert_eq!(token, optimistic_draft);
14176                        seed
14177                    });
14178                    controller_prepared = Some(OptiControllerPrepared {
14179                        verify_tokens: [optimistic_pending, optimistic_draft],
14180                        draft_prob: optimistic_draft_probability,
14181                        eager_seed,
14182                        q_proxy,
14183                        scratch_len: scratch.kv.len,
14184                    });
14185                } else {
14186                    OPTI_GATE_REJECTS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
14187                    OPTI_WASTED_DRAFT_TOKENS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
14188                    rejected_probe = Some((q_proxy, optimistic_pending));
14189                    eprintln!(
14190                        "[opti-controller] reject q={q_proxy:.6} threshold={:.3}",
14191                        opti_fork
14192                            .as_ref()
14193                            .and_then(|fork| fork.controller.as_ref())
14194                            .expect("controller policy")
14195                            .threshold,
14196                    );
14197                }
14198            }
14199            let fork_attempt = match fork_generation.take() {
14200                Some(generation) if base == 1 && k_round == 1 => Some(generation),
14201                Some(generation) => {
14202                    opti_fork
14203                        .as_mut()
14204                        .expect("fork generation without fork state")
14205                        .retire(generation)?;
14206                    None
14207                }
14208                None => None,
14209            };
14210            let (tlogits_d, vx) = if let Some(p) = pipe {
14211                self.decode_step_t_core_pipelined(
14212                    e,
14213                    &verify_tokens,
14214                    pos,
14215                    &mut *cache,
14216                    embd_dev,
14217                    ckpt.as_mut(),
14218                    p,
14219                    round,
14220                )?
14221            } else if controller_can_probe {
14222                let fence = opti_fork
14223                    .as_ref()
14224                    .ok_or("optipipe controller probe lost fork state")?
14225                    .fence;
14226                let boundary = match current_opti.as_mut() {
14227                    Some(ticket) => ticket.take_boundary(),
14228                    None => self.verify_stage0_issue(
14229                        e,
14230                        &verify_tokens,
14231                        pos,
14232                        &mut *cache,
14233                        embd_dev,
14234                        ckpt.as_mut(),
14235                        None,
14236                        &fence,
14237                        Some(true),
14238                        None,
14239                    )?,
14240                };
14241                if let Some(prepared) = controller_prepared.take() {
14242                    let generation = {
14243                        let fork = opti_fork
14244                            .as_mut()
14245                            .ok_or("optipipe controller admission lost fork state")?;
14246                        let generation = fork.reserve_successor()?;
14247                        let rt = fork.rt;
14248                        let snapshot_fence = fork.fence;
14249                        opti_snapshot_one_stage_owned_into(
14250                            e,
14251                            cache,
14252                            rt,
14253                            &snapshot_fence,
14254                            0,
14255                            fork.successor_snapshot_mut(),
14256                        )?;
14257                        generation
14258                    };
14259                    let mut successor_ckpt = VerifyCkpt::new(self.layers.len());
14260                    let successor_boundary = self.verify_stage0_issue(
14261                        e,
14262                        &prepared.verify_tokens,
14263                        pos + verify_tokens.len(),
14264                        &mut *cache,
14265                        embd_dev,
14266                        Some(&mut successor_ckpt),
14267                        None,
14268                        &fence,
14269                        Some(false),
14270                        None,
14271                    )?;
14272                    OPTI_FORK_ATTEMPTS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
14273                    let fork = opti_fork
14274                        .as_ref()
14275                        .ok_or("optipipe controller ticket lost fork state")?;
14276                    successor_attempt = Some(fork.controller_ticket(
14277                        generation,
14278                        successor_boundary,
14279                        successor_ckpt,
14280                        prepared.verify_tokens,
14281                        prepared.draft_prob,
14282                        prepared.eager_seed,
14283                        prepared.q_proxy,
14284                        prepared.scratch_len,
14285                    ));
14286                    eprintln!(
14287                        "[opti-controller] issue generation={} q={:.6} threshold={:.3} \
14288                         verify={:?}",
14289                        generation.id,
14290                        prepared.q_proxy,
14291                        fork.controller.expect("controller policy").threshold,
14292                        prepared.verify_tokens,
14293                    );
14294                }
14295                let result = self.verify_stage1_finish(
14296                    e,
14297                    boundary,
14298                    &mut *cache,
14299                    ckpt.as_mut(),
14300                    None,
14301                    &fence,
14302                    successor_attempt.is_none(),
14303                )?;
14304                if let Some(ticket) = current_opti.as_mut() {
14305                    ticket.settle();
14306                }
14307                if successor_attempt.is_some() {
14308                    let fork = opti_fork
14309                        .as_mut()
14310                        .ok_or("optipipe successor snapshot lost fork state")?;
14311                    let rt = fork.rt;
14312                    let snapshot_fence = fork.fence;
14313                    opti_snapshot_one_stage_owned_into(
14314                        e,
14315                        cache,
14316                        rt,
14317                        &snapshot_fence,
14318                        1,
14319                        fork.successor_snapshot_mut(),
14320                    )?;
14321                    // Publish N only after both independent successor-state queues are complete.
14322                    fork.rt.publish_to(1, &e.stream())?;
14323                }
14324                result
14325            } else if let Some(ticket) = current_opti.as_mut() {
14326                let fork = opti_fork
14327                    .as_mut()
14328                    .ok_or("optipipe carried controller ticket lost fork state")?;
14329                let boundary = ticket.take_boundary();
14330                let result = self.verify_stage1_finish(
14331                    e,
14332                    boundary,
14333                    &mut *cache,
14334                    ckpt.as_mut(),
14335                    None,
14336                    &fork.fence,
14337                    true,
14338                )?;
14339                ticket.settle();
14340                result
14341            } else if let Some(generation) = fork_attempt {
14342                let fork = opti_fork
14343                    .as_mut()
14344                    .expect("fork generation without fork state");
14345                fork.capture_seed(e, generation, &h_seed_buf, &fill_prev, scratch.kv.len)?;
14346                let action = fork.mode.action(generation.id);
14347                let boundary = self.verify_stage0_issue(
14348                    e,
14349                    &verify_tokens,
14350                    pos,
14351                    &mut *cache,
14352                    embd_dev,
14353                    ckpt.as_mut(),
14354                    None,
14355                    &fork.fence,
14356                    Some(true),
14357                    None,
14358                )?;
14359                OPTI_FORK_ATTEMPTS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
14360                let mut ticket = fork.ticket(generation, boundary);
14361                if action == OptiForkAction::Abort {
14362                    return Err(format!(
14363                        "optipipe forced abort with generation {} stage0 in flight",
14364                        generation.id,
14365                    )
14366                    .into());
14367                }
14368                fork.reconcile(
14369                    e,
14370                    &mut *cache,
14371                    &mut *scratch,
14372                    &snap,
14373                    &mut h_seed_buf,
14374                    &mut fill_prev,
14375                    generation,
14376                    action,
14377                    verify_tokens[0],
14378                )?;
14379                let result = if action == OptiForkAction::Hit {
14380                    let boundary = ticket.take_boundary();
14381                    self.verify_stage1_finish(
14382                        e,
14383                        boundary,
14384                        &mut *cache,
14385                        ckpt.as_mut(),
14386                        None,
14387                        &fork.fence,
14388                        true,
14389                    )?
14390                } else {
14391                    // The optimistic boundary slot has no reader. Re-run the unchanged serial
14392                    // verify only after E_restart published the restored stage-0 state.
14393                    self.decode_step_t_core(
14394                        e,
14395                        &verify_tokens,
14396                        pos,
14397                        &mut *cache,
14398                        embd_dev,
14399                        ckpt.as_mut(),
14400                    )?
14401                };
14402                ticket.settle();
14403                debug_assert_eq!(ticket.generation, generation);
14404                fork.retire(generation)?;
14405                result
14406            } else {
14407                // The serial verify every non-fork round takes — the MTP route's
14408                // verify-graph door. The pool is None unless MEMRA_SPEC_VERIFY_GRAPH armed
14409                // a pool above, and then the walk replays the captured trunk instead of
14410                // re-issuing it launch by launch. `graph_round_ok` is the round's
14411                // headroom snapshot (see GRAPH_LAUNCH_MIN_FREE): below the floor the
14412                // round declines the pool exactly like an over-cap round and rides the
14413                // byte-identical eager walk — the `[spec]` suspension line above
14414                // already named the round.
14415                let vg_round = if verify_tokens.len() <= vg_t_cap && graph_round_ok {
14416                    vg_guard.as_mut().and_then(|g| g.as_mut())
14417                } else {
14418                    if let Some(g) = vg_guard.as_mut().and_then(|g| g.as_mut()) {
14419                        // The commit reads this flag to pick its arm; a round that declines
14420                        // the pool must not inherit a stale `true` from the round before it.
14421                        g.round_slab = false;
14422                    }
14423                    None
14424                };
14425                self.decode_step_t_core_vg(
14426                    e,
14427                    &verify_tokens,
14428                    pos,
14429                    &mut *cache,
14430                    embd_dev,
14431                    ckpt.as_mut(),
14432                    vg_round,
14433                )?
14434            };
14435            let pipe_accept = match pipe {
14436                Some(p) => Some(p.accept_begin(round)?),
14437                None => None,
14438            };
14439
14440            if phase_sync {
14441                e.stream().synchronize()?;
14442            }
14443            ph_mark(&mut ph_verify, phase_on);
14444            // --- 3. GREEDY ACCEPT (walk prefix, stop at first mismatch) ---
14445            // DEVICE-ARGMAX ACCEPT: argmax every verify column ON DEVICE (same 2-pass kernels +
14446            // smallest-index tie-break as host argmax, argmax_gate-validated) and read back ONE
14447            // [T] u32 — replaces the T x n_vocab f32 dtoh + T host argmaxes per round.
14448            // t_pred[j] = target's greedy prediction for the slot after draft[j-1] (j>=1) or after
14449            // last_token (j==0). With a pending bonus, col 0 IS the prediction after last_token
14450            // (== the bonus), so every index shifts by `base` and last_pred is unused.
14451            let t_v = verify_tokens.len();
14452            let mut preds: Vec<u32> = Vec::new();
14453            if !sampled {
14454                for j in 0..t_v {
14455                    e.argmax_token_device_col(&tlogits_d, j, n_vocab, &mut preds_d, j)?;
14456                }
14457                preds = e.dtoh_u32(&preds_d)?; // <- the verify-GPU wait lands here
14458                // #87 SENTINEL TRAP, verify side: a sentinel pred becomes the round's bonus =
14459                // next round's last_token = the next chain's embed lookup. Catch it at the
14460                // source with the column named — an all-NaN VERIFY column implicates the
14461                // stage-split trunk (decode_step_t_core_ppn), not the draft head.
14462                if let Some(bad) = preds[..t_v].iter().position(|&p| (p as usize) >= n_vocab) {
14463                    let col = &tlogits_d.slice(bad * n_vocab..(bad + 1) * n_vocab);
14464                    let mut probe = e.zeros(n_vocab)?;
14465                    e.copy_view_into(&mut probe, 0, col, n_vocab)?;
14466                    let col_h = e.dtoh(&probe)?;
14467                    let col_nan = col_h.iter().filter(|v| v.is_nan()).count();
14468                    return Err(format!(
14469                        "verify argmax sentinel 0x{:08x} >= n_vocab {n_vocab} at round {round} \
14470                         col {bad}/{t_v} pos={pos}: verify-logits col NaN {col_nan}/{n_vocab} \
14471                         — the verify TRUNK produced a poisoned column (#87 trap). Run \
14472                         MEMRA_SPEC_NAN_SCAN=1 to name the layer that creates it (=2 to split \
14473                         that layer into attention and routed MoE). NOT the draft head, and NOT \
14474                         the PP stage split this message used to name: pp_cuts() returns None \
14475                         without MEMRA_PP_STAGES, so decode_step_t_core_ppn never runs unless \
14476                         that variable is set.",
14477                        preds[bad]
14478                    )
14479                    .into());
14480                }
14481            }
14482            ph_mark(&mut ph_wait, phase_on);
14483            let t_pred = |j: usize| -> u32 {
14484                if j == 0 && base == 0 {
14485                    last_pred
14486                } else {
14487                    // GREEDY-ONLY: `preds` is filled under `if !sampled` above. The debug print
14488                    // used to call this from the sampled arm and panicked the worker; it now goes
14489                    // through `debug_t_pred0`. Keep the strict index here — in the greedy walk an
14490                    // out-of-range pred is a real bug, not something to paper over.
14491                    debug_assert!(
14492                        !sampled,
14493                        "t_pred is greedy-only: `preds` is empty in the sampled arm"
14494                    );
14495                    preds[base + j - 1]
14496                }
14497            };
14498            let mut devacc_seeded = false;
14499            let mut devacc_acc: Option<CudaSlice<u32>> = None;
14500            let (n_acc, bonus) = if !sampled {
14501                // ROUND-STREAM stage (a) (MEMRA_SPEC_DEVACC=1 opt-in): the walk runs ON DEVICE
14502                // (spec_accept_greedy, verbatim rule) and the host reads back 8B (n_acc, bonus)
14503                // instead of the [T] preds. Same sync count — machinery for stages (b)/(c),
14504                // gated on token identity vs the host walk (the arms below are bit-equal rules).
14505                if crate::spec::spec_devacc() && k_round > 0 && !spec_replay && constraint.is_none()
14506                {
14507                    let draft_d = e.htod_u32_v(&draft)?;
14508                    let mut acc_out = e.alloc_u32_zeroed(2)?;
14509                    e.spec_accept_greedy(
14510                        &preds_d,
14511                        &draft_d,
14512                        last_pred,
14513                        base,
14514                        k_round,
14515                        &mut acc_out,
14516                    )?;
14517                    devacc_acc = Some(acc_out.clone());
14518                    // stage (b): next-round seed gathered ON DEVICE from acc_out before the host
14519                    // ever reads n_acc (j=base+n_acc -> vx col j-1; j==0 -> fill_prev). The three
14520                    // non-replay commit arms skip their host-offset seed copies (guarded below);
14521                    // the legacy spec_replay arm keeps its own rx-based seeding (excluded here).
14522                    // NOTE: fill_prev is NOT updated here — the commit arms' TRUE-HIDDEN
14523                    // REFRESH reads the OLD fill_prev (predecessor of this round's verify batch);
14524                    // the update lands after the arms (devacc_seeded guard below).
14525                    e.spec_seed_gather(&vx, &fill_prev, &acc_out, &mut h_seed_buf, base, n_embd)?;
14526                    // 3a: KV lens roll back on device (len = saved + base + n_acc, all arms'
14527                    // unified rule; full accept rewrites the verify-left value). Host mirrors
14528                    // update after the readback; commit_verified_prefix skips its len_d writes.
14529                    if let Some(successor) = successor_attempt.as_ref() {
14530                        opti_fork
14531                            .as_mut()
14532                            .ok_or("optipipe successor reconcile lost fork state")?
14533                            .queue_actual_reconcile(
14534                                e,
14535                                &snap,
14536                                &acc_out,
14537                                successor.verify_tokens[0],
14538                                base,
14539                            )?;
14540                    } else if let Some(ptrs) = &kv_len_ptrs {
14541                        let saved: Vec<i32> = (0..self.layers.len())
14542                            .map(|il| snap.kv_len[il].map(|v| v as i32).unwrap_or(0))
14543                            .collect();
14544                        let saved_d = e.htod_i32(&saved)?;
14545                        e.spec_rollback_kv(ptrs, &saved_d, &acc_out, base, self.layers.len())?;
14546                    }
14547                    devacc_seeded = true;
14548                    let ab = e.dtoh_u32(&acc_out)?;
14549                    (ab[0] as usize, ab[1])
14550                } else {
14551                    let mut n_acc = 0usize;
14552                    for j in 0..k_round {
14553                        if t_pred(j) == draft[j] {
14554                            n_acc += 1;
14555                        } else {
14556                            break;
14557                        }
14558                    }
14559                    // bonus = target's own token at the first non-accepted slot. n_acc in 0..=k; t_pred
14560                    // is defined for j in 0..=k (j==0 -> last_logits, j>=1 -> col j-1, last col = k-1).
14561                    (n_acc, t_pred(n_acc))
14562                }
14563            } else {
14564                // --- SAMPLED ACCEPT (rejection sampling): u_j < p_j(x_j)/q_j(x_j) walk ---
14565                if col_buf.is_none() {
14566                    col_buf = Some(e.zeros(n_vocab)?);
14567                }
14568                // FILTERED p_j: per-verify-col stats (one batched filter_stats call), then the
14569                // filtered gather. j==0&&base==0 reads last_col (its own stats row appended).
14570                let mut pj = vec![0f32; k_round.max(1)];
14571                let mut col_stats: Vec<(f32, f32, f32)> = Vec::new(); // (max, th, z) per verify col used
14572                if k_round > 0 {
14573                    let mut ids: Vec<u32> = Vec::new();
14574                    let mut rows: Vec<i32> = Vec::new();
14575                    for j in 0..k_round {
14576                        if j > 0 || base == 1 {
14577                            ids.push(draft[j]);
14578                            rows.push((base + j) as i32 - 1);
14579                        }
14580                    }
14581                    if !ids.is_empty() {
14582                        let nr = rows.len();
14583                        // penalties: materialize the used columns into one contiguous penalized
14584                        // buffer (rows remapped 0..nr) so stats+gathers see the penalized p.
14585                        // penalties: materialize used columns contiguously, penalize all rows in
14586                        // one launch, and point stats+gathers at the penalized buffer (rows 0..nr).
14587                        let p_rows: Vec<i32> = if pen_on {
14588                            (0..nr as i32).collect()
14589                        } else {
14590                            rows.clone()
14591                        };
14592                        if pen_on {
14593                            if pcol_buf.as_ref().map(|b| b.len()).unwrap_or(0) < nr * n_vocab {
14594                                pcol_buf = Some(e.zeros(nr * n_vocab)?);
14595                            }
14596                            let pc = pcol_buf.as_mut().unwrap();
14597                            for (i2, &r) in rows.iter().enumerate() {
14598                                let c = r as usize;
14599                                e.copy_view_into(
14600                                    pc,
14601                                    i2 * n_vocab,
14602                                    &tlogits_d.slice(c * n_vocab..(c + 1) * n_vocab),
14603                                    n_vocab,
14604                                )?;
14605                            }
14606                            let h = pen_hist_d.as_ref().unwrap();
14607                            let nh = h.len();
14608                            e.penalize_logits_rows(
14609                                pc,
14610                                h,
14611                                nh,
14612                                sp.penalty_repeat,
14613                                sp.penalty_freq,
14614                                sp.penalty_present,
14615                                n_vocab,
14616                                nr,
14617                            )?;
14618                        }
14619                        let p_src: &CudaSlice<f32> = if pen_on {
14620                            pcol_buf.as_ref().unwrap()
14621                        } else {
14622                            &tlogits_d
14623                        };
14624                        let rowsd = e.htod_i32(&p_rows)?;
14625                        let (mut th_d, mut z_d, mut mx_d) =
14626                            (e.zeros(nr)?, e.zeros(nr)?, e.zeros(nr)?);
14627                        e.filter_stats(
14628                            p_src, n_vocab, &rowsd, &mut th_d, &mut z_d, &mut mx_d, n_vocab, nr,
14629                            sp_temp, sp.top_k, sp.top_p, sp.min_p,
14630                        )?;
14631                        let idsd = e.htod_u32_v(&ids)?;
14632                        let mut outd = e.zeros(nr)?;
14633                        e.softmax_gather_filtered(
14634                            p_src, n_vocab, &idsd, &rowsd, &th_d, &z_d, &mut outd, n_vocab, nr,
14635                            sp_temp,
14636                        )?;
14637                        let outv = e.dtoh(&outd)?;
14638                        let (thv, zv, mxv) = (e.dtoh(&th_d)?, e.dtoh(&z_d)?, e.dtoh(&mx_d)?);
14639                        let mut oi = 0usize;
14640                        for j in 0..k_round {
14641                            if j > 0 || base == 1 {
14642                                pj[j] = outv[oi];
14643                                oi += 1;
14644                            }
14645                        }
14646                        col_stats = (0..nr).map(|i| (mxv[i], thv[i], zv[i])).collect();
14647                    }
14648                    if base == 0 {
14649                        let lc: &CudaSlice<f32> = if pen_on {
14650                            if col_buf.is_none() {
14651                                col_buf = Some(e.zeros(n_vocab)?);
14652                            }
14653                            let cb = col_buf.as_mut().unwrap();
14654                            e.copy_into(
14655                                cb,
14656                                0,
14657                                last_col_logits
14658                                    .as_ref()
14659                                    .expect("sampled: last_col_logits unset"),
14660                                n_vocab,
14661                            )?;
14662                            let h = pen_hist_d.as_ref().unwrap();
14663                            let nh = h.len();
14664                            e.penalize_logits(
14665                                cb,
14666                                h,
14667                                nh,
14668                                sp.penalty_repeat,
14669                                sp.penalty_freq,
14670                                sp.penalty_present,
14671                                n_vocab,
14672                            )?;
14673                            col_buf.as_ref().unwrap()
14674                        } else {
14675                            last_col_logits
14676                                .as_ref()
14677                                .expect("sampled: last_col_logits unset")
14678                        };
14679                        let rows0 = e.htod_i32(&[0])?;
14680                        let (mut th_d, mut z_d, mut mx_d) = (e.zeros(1)?, e.zeros(1)?, e.zeros(1)?);
14681                        e.filter_stats(
14682                            lc, n_vocab, &rows0, &mut th_d, &mut z_d, &mut mx_d, n_vocab, 1,
14683                            sp_temp, sp.top_k, sp.top_p, sp.min_p,
14684                        )?;
14685                        let idsd = e.htod_u32_v(&[draft[0]])?;
14686                        let mut outd = e.zeros(1)?;
14687                        e.softmax_gather_filtered(
14688                            lc, n_vocab, &idsd, &rows0, &th_d, &z_d, &mut outd, n_vocab, 1, sp_temp,
14689                        )?;
14690                        pj[0] = e.dtoh(&outd)?[0];
14691                        last_col_stats =
14692                            Some((e.dtoh(&mx_d)?[0], e.dtoh(&th_d)?[0], e.dtoh(&z_d)?[0]));
14693                    }
14694                }
14695                // q source: the graph arms (single-head AND chain) retained the head logits
14696                // in the persistent q_slots; the eager arm in per-round draft_logits clones.
14697                // Same raw-logit values either way. FILTERED q_j: stats from draft_stats
14698                // (eager pushes in-chain; the graph arms compute them post-replay from the
14699                // retained q with the same filter_stats program — bit-identical to the
14700                // in-graph stats that shaped the draw, keeping ONE accept path).
14701                let q_bufs: &[CudaSlice<f32>] = if dctx.graph_s.is_some() || dctx.chain_s.is_some()
14702                {
14703                    &dctx.q_slots
14704                } else {
14705                    &draft_logits
14706                };
14707                let mut n_acc = 0usize;
14708                for j in 0..k_round {
14709                    let (qmx, qth, qz) = draft_stats[j];
14710                    let idsd = e.htod_u32_v(&[draft_idx[j]])?;
14711                    let rowsd = e.htod_i32(&[0])?;
14712                    let thd = e.htod(&[qth])?;
14713                    let zd = e.htod(&[qz])?;
14714                    let _ = qmx;
14715                    let mut outd = e.zeros(1)?;
14716                    e.softmax_gather_filtered(
14717                        &q_bufs[j], d_vocab, &idsd, &rowsd, &thd, &zd, &mut outd, d_vocab, 1,
14718                        sp_temp,
14719                    )?;
14720                    let qj = e.dtoh(&outd)?[0];
14721                    let u = host_u01(sp_seed, uctr);
14722                    uctr += 1;
14723                    let accept = (u as f64) * (qj as f64) < pj[j] as f64;
14724                    // SKEY PROBE: q == 0 for the token the draft actually proposed is the
14725                    // exactness signature (see `skey_probe`). Impossible when the draft was
14726                    // drawn from the same filtered distribution the verify reconstructs here;
14727                    // `u * 0 < p` makes it an UNCONDITIONAL accept whenever p > 0.
14728                    if skey_probe() && qj == 0.0 {
14729                        eprintln!(
14730                            "[skey] EXACTNESS q=0 round={round} j={j} draft_tok={} \
14731                             draft_idx={} p={:e} u={u} accepted={} th_z={:?}",
14732                            draft[j], draft_idx[j], pj[j], accept as u8, draft_stats[j],
14733                        );
14734                    }
14735                    if accept {
14736                        n_acc += 1;
14737                    } else {
14738                        break;
14739                    }
14740                }
14741                let bonus = if n_acc == k_round {
14742                    // FULL ACCEPT: bonus ~ FILTERED softmax at the last verify column.
14743                    let col = base + k_round - 1;
14744                    let cb = col_buf.as_mut().unwrap();
14745                    e.copy_view_into(
14746                        cb,
14747                        0,
14748                        &tlogits_d.slice(col * n_vocab..(col + 1) * n_vocab),
14749                        n_vocab,
14750                    )?;
14751                    if pen_on {
14752                        let h = pen_hist_d.as_ref().unwrap();
14753                        let nh = h.len();
14754                        e.penalize_logits(
14755                            cb,
14756                            h,
14757                            nh,
14758                            sp.penalty_repeat,
14759                            sp.penalty_freq,
14760                            sp.penalty_present,
14761                            n_vocab,
14762                        )?;
14763                    }
14764                    if perturb_buf.is_none() {
14765                        perturb_buf = Some(e.zeros(d_vocab.max(n_vocab))?);
14766                    }
14767                    // STATS MUST COME FROM THIS COLUMN (bug fix 2026-08-05, lane/sampler-
14768                    // truncation-fix; receipts research/sampfix-20260805/). The old code reused
14769                    // `col_stats.last()` here, which is ALWAYS the wrong row: the gathered set
14770                    // covers verify columns 0..=(base+k_round-2) (rows pushed as base+j-1), while
14771                    // the full-accept bonus samples column base+k_round-1 — exactly ONE PAST the
14772                    // last gathered column, in both base arms. `th` is a threshold in e-units of
14773                    // its OWN row's max, so feeding a neighbour's (row_max, th) into
14774                    // gumbel_perturb_filtered mis-scales every e0 = exp((x-row_max)/T). When the
14775                    // donor column's peak is higher by more than T*ln(1/th), EVERY id fails
14776                    // `e0 >= th`, the whole perturbed row becomes -3.4e38, and the 2-pass argmax
14777                    // falls through to its smallest-index tie-break => token id 0 ("!") spliced
14778                    // mid-word. Fragility is ordered by how large th is: min_p pins th = min_p
14779                    // (0.05 => trigger at delta > 2.4 at T=0.8, fires constantly), top_p's
14780                    // mass-boundary th is smaller, top_k's k-th-largest th smaller still — which
14781                    // is why the head-to-head matrix saw min_p and top_p corrupt while top_k-only
14782                    // stayed clean. The pure-temp default regime is immune (th == 0 masks nothing,
14783                    // and row_max is unused once nothing is masked), so this fix is a byte-level
14784                    // no-op for the untruncated serve default. One extra one-block filter_stats
14785                    // per full-accept round is the whole cost.
14786                    let (mx, th) = {
14787                        let rows0 = e.htod_i32(&[0])?;
14788                        let (mut th_d, mut z_d, mut mx_d) = (e.zeros(1)?, e.zeros(1)?, e.zeros(1)?);
14789                        let cb0 = col_buf.as_ref().unwrap();
14790                        e.filter_stats(
14791                            cb0, n_vocab, &rows0, &mut th_d, &mut z_d, &mut mx_d, n_vocab, 1,
14792                            sp_temp, sp.top_k, sp.top_p, sp.min_p,
14793                        )?;
14794                        (e.dtoh(&mx_d)?[0], e.dtoh(&th_d)?[0])
14795                    };
14796                    let pb = perturb_buf.as_mut().unwrap();
14797                    let cb2 = col_buf.as_ref().unwrap();
14798                    e.gumbel_perturb_filtered(cb2, pb, n_vocab, sp_seed, sctr, sp_temp, mx, th)?;
14799                    sctr += 1;
14800                    let td = e.argmax_token_device(pb, n_vocab)?;
14801                    e.dtoh_u32_one(&td)?
14802                } else {
14803                    // REJECT at n_acc: bonus ~ norm(max(0, softmax_T(p) - softmax_T(q))).
14804                    let cb = col_buf.as_mut().unwrap();
14805                    if n_acc > 0 || base == 1 {
14806                        let col = base + n_acc - 1;
14807                        e.copy_view_into(
14808                            cb,
14809                            0,
14810                            &tlogits_d.slice(col * n_vocab..(col + 1) * n_vocab),
14811                            n_vocab,
14812                        )?;
14813                    } else {
14814                        let lc = last_col_logits.as_ref().unwrap();
14815                        e.copy_into(cb, 0, lc, n_vocab)?;
14816                    }
14817                    if pen_on {
14818                        let h = pen_hist_d.as_ref().unwrap();
14819                        let nh = h.len();
14820                        e.penalize_logits(
14821                            cb,
14822                            h,
14823                            nh,
14824                            sp.penalty_repeat,
14825                            sp.penalty_freq,
14826                            sp.penalty_present,
14827                            n_vocab,
14828                        )?;
14829                    }
14830                    let cb2 = col_buf.as_ref().unwrap();
14831                    let sc = sctr;
14832                    sctr += 1;
14833                    // p-stats for the reject column: from col_stats when the col was gathered,
14834                    // else (j==0&&base==0) from last_col_stats.
14835                    let p_stats = if n_acc > 0 || base == 1 {
14836                        // col index within the gathered set == number of gathered cols before n_acc
14837                        let gi = if base == 1 { n_acc } else { n_acc - 1 };
14838                        col_stats.get(gi).copied().unwrap_or_else(|| {
14839                            (0.0, 0.0, 1.0) // unreachable: gathered cols always cover the reject slot
14840                        })
14841                    } else {
14842                        last_col_stats.expect("sampled: last_col_stats unset at reject")
14843                    };
14844                    let q_stats = draft_stats[n_acc];
14845                    if let Some(map) = &d2t_dev {
14846                        if q_full_buf.is_none() {
14847                            q_full_buf = Some(e.zeros(n_vocab)?);
14848                        }
14849                        let qf = q_full_buf.as_mut().unwrap();
14850                        e.scatter_trim_logits(&q_bufs[n_acc], map, qf, d_vocab, n_vocab)?;
14851                        let qf2 = q_full_buf.as_ref().unwrap();
14852                        e.residual_sample_filtered(
14853                            cb2,
14854                            Some(qf2),
14855                            n_vocab,
14856                            sp_temp,
14857                            sp_seed,
14858                            sc,
14859                            p_stats,
14860                            q_stats,
14861                            &mut sample_tok,
14862                        )?;
14863                    } else {
14864                        e.residual_sample_filtered(
14865                            cb2,
14866                            Some(&q_bufs[n_acc]),
14867                            n_vocab,
14868                            sp_temp,
14869                            sp_seed,
14870                            sc,
14871                            p_stats,
14872                            q_stats,
14873                            &mut sample_tok,
14874                        )?;
14875                    }
14876                    e.dtoh_u32(&sample_tok)?[0]
14877                };
14878                (
14879                    n_acc,
14880                    guard_vocab_token(
14881                        bonus,
14882                        n_vocab,
14883                        &format!("sampled verify bonus at round {round} pos={pos} n_acc={n_acc}"),
14884                    )?,
14885                )
14886            };
14887            // --- 3b. GRAMMAR TRUNCATION (constrained spec, 2026-08-03): the grammar is
14888            // an extra rejection rule AFTER the exactness verify (the batched-verify-twins
14889            // ordering). Walk the accepted drafts through the grammar in commit order; the
14890            // first illegal token truncates acceptance at its slot, and that slot's emission
14891            // is recomputed as the MASKED argmax of the target's own verify column — token-
14892            // identical to constrained plain greedy decode (an unmasked argmax that is
14893            // grammar-legal IS the masked argmax: masking only removes competitors). The
14894            // column D2H (~1MB) is paid only when a cut fires — the tight-grammar cost,
14895            // measured in acceptance numbers, never hidden.
14896            let (n_acc, bonus) = match constraint.as_deref_mut() {
14897                None => (n_acc, bonus),
14898                Some(c) => {
14899                    fn ce(e2: String) -> Box<dyn std::error::Error> {
14900                        format!("constraint: {e2}").into()
14901                    }
14902                    let mut na = n_acc;
14903                    let mut cut = false;
14904                    for (j, &d) in draft.iter().enumerate().take(n_acc) {
14905                        if c.is_allowed(d).map_err(ce)? {
14906                            c.consume(d).map_err(ce)?;
14907                        } else {
14908                            na = j;
14909                            cut = true;
14910                            dm_cut_tokens += n_acc - j;
14911                            break;
14912                        }
14913                    }
14914                    if cut {
14915                        dm_cuts += 1;
14916                    }
14917                    let mut bo = bonus;
14918                    if cut || !c.is_allowed(bo).map_err(ce)? {
14919                        let mut row = if na == 0 && base == 0 {
14920                            init_logits_host
14921                                .clone()
14922                                .ok_or("constraint: init logits missing (round-0 cut)")?
14923                        } else {
14924                            e.dtoh_view(
14925                                &tlogits_d.slice((base + na - 1) * n_vocab..(base + na) * n_vocab),
14926                            )?
14927                        };
14928                        c.mask_logits(&mut row).map_err(ce)?;
14929                        bo = argmax(&row) as u32;
14930                    }
14931                    c.consume(bo).map_err(ce)?;
14932                    (na, bo)
14933                }
14934            };
14935            let mut successor_valid = false;
14936            if let Some((q_proxy, expected_d2)) = rejected_probe {
14937                let v_n = n_acc == 1 && bonus == expected_d2;
14938                eprintln!(
14939                    "[opti-controller] shadow q={q_proxy:.6} admitted=false v_n={v_n} \
14940                     expected_d2={expected_d2} n_acc={n_acc} bonus={bonus}",
14941                );
14942            }
14943            if let Some(successor) = successor_attempt.as_ref() {
14944                successor_valid = n_acc == 1 && bonus == successor.verify_tokens[0];
14945                let generation = successor.generation;
14946                let q_proxy = successor.q_proxy;
14947                let expected_pending = successor.verify_tokens[0];
14948                let resolution_ms = successor.issued_at.elapsed().as_secs_f64() * 1e3;
14949                let fork = opti_fork
14950                    .as_mut()
14951                    .ok_or("optipipe successor resolution lost fork state")?;
14952                fork.finish_actual_reconcile(e, &mut *cache, &snap, n_acc, base, successor_valid)?;
14953                if successor_valid {
14954                    OPTI_FORK_HITS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
14955                } else {
14956                    OPTI_FORK_MISSES.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
14957                    OPTI_RECONCILES.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
14958                    OPTI_WASTED_DRAFT_TOKENS.fetch_add(2, std::sync::atomic::Ordering::Relaxed);
14959                }
14960                let breaker_tripped = fork
14961                    .controller
14962                    .as_mut()
14963                    .expect("controller policy")
14964                    .resolve(successor_valid);
14965                if breaker_tripped {
14966                    OPTI_BREAKER_TRIPS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
14967                }
14968                eprintln!(
14969                    "[opti-controller] resolve generation={} hit={} q={q_proxy:.6} \
14970                     expected_pending={expected_pending} n_acc={n_acc} bonus={bonus} \
14971                     resolution_ms={resolution_ms:.3} reconcile={} breaker={}",
14972                    generation.id, successor_valid, !successor_valid, breaker_tripped,
14973                );
14974                if !successor_valid {
14975                    let mut successor = successor_attempt
14976                        .take()
14977                        .expect("controller successor disappeared on miss");
14978                    successor.settle();
14979                    fork.retire(generation)?;
14980                }
14981            }
14982            total_drafted += k_round;
14983            total_accepted += n_acc;
14984            if let Some(t) = sess_telem {
14985                // Greedy, rejection-sampling, and grammar truncation all converge here after
14986                // the accept decision is already on host. Fixed-size relaxed atomics only.
14987                t.record_round(k_round, n_acc);
14988            }
14989            if spec_stats {
14990                st_len_hist[k_round] += 1;
14991                for j in 0..k_round {
14992                    st_drafted[j] += 1;
14993                }
14994                for j in 0..n_acc {
14995                    st_accepted[j] += 1;
14996                }
14997                if n_acc == k_round {
14998                    st_full += 1;
14999                }
15000            }
15001
15002            if debug_spec {
15003                eprintln!(
15004                    "[R{round}] pos={pos} out_len={} last_tok={last_token} draft={draft:?} n_acc={n_acc} bonus={bonus} t_pred0={}",
15005                    out.len(),
15006                    // NOT `t_pred(0)`: `preds` is filled only under `if !sampled` above, so on a
15007                    // sampled request round >= 1 (base == 1) indexed an EMPTY vector and PANICKED
15008                    // the GPU worker thread — a debug flag that killed the exact regime you would
15009                    // set it to investigate. See `debug_t_pred0`.
15010                    debug_t_pred0(sampled, base, last_pred, &preds)
15011                );
15012            }
15013
15014            // --- 4. COMMIT: draft[0..n_acc] then bonus (n_acc + 1 tokens) ---
15015            let commit_started = std::time::Instant::now();
15016            // SESSION MODE: every accepted column is already in the CACHE — `out` must carry all
15017            // of them (overshoot past max_new included) or `committed` under-counts the cache rows
15018            // and the next turn's continuation seeds one token off (gate-caught 2026-07-05). The
15019            // single-shot path keeps the cap (its caller truncates + drops the cache anyway).
15020            for j in 0..n_acc {
15021                if !session_mode && out.len() >= max_new {
15022                    break;
15023                }
15024                out.push(draft[j]);
15025            }
15026            if pen_on {
15027                pen_hist.extend_from_slice(&draft[0..n_acc]);
15028                pen_hist.push(bonus);
15029            }
15030            let bonus_emitted = session_mode || out.len() < max_new;
15031            if bonus_emitted {
15032                out.push(bonus);
15033            }
15034            last_token = bonus;
15035
15036            // --- 5. ROLLBACK + advance (§C) ---
15037            if n_acc == k_round && !spec_replay {
15038                // FULL ACCEPT, BONUS FOLD: all verify columns (pending? + drafts) are committed in
15039                // cache; the NEW bonus stays PENDING for the next round's verify batch — NO extra
15040                // T=1 trunk pass. The next draft chain seeds from the MTP block's h_nextn at the
15041                // bonus position: one MTP-block pass (~1/33 trunk cost) replaces the trunk read.
15042                // last_pred is dead in the pending path (t_pred reads verify col 0).
15043                //
15044                // PERSISTENT DRAFT KV, full-accept fill: the chain covered last_token +
15045                // draft[0..k_round-2] as INPUTS (slots P..P'-2); draft[k_round-1] (slot P'-1) was
15046                // only ever an output, so its entry is MISSING. Fill it from vh_seed — its EXACT
15047                // trunk hidden (the last verify column). set_len first: a p-min break may have
15048                // left one extra chain append at that slot. Partial accepts need NO fill (the
15049                // chain already covered every accepted position; round-start set_len truncates).
15050                let mut vh_seed = e.zeros(n_embd)?;
15051                e.copy_view_into(
15052                    &mut vh_seed,
15053                    0,
15054                    &vx.slice((t_v - 1) * n_embd..t_v * n_embd),
15055                    n_embd,
15056                )?;
15057                if refresh {
15058                    // TRUE-HIDDEN REFRESH (2026-07-03, the HANDOVER-listed acceptance lever):
15059                    // overwrite ALL committed positions' scratch entries with K/V from their EXACT
15060                    // verify hiddens — the reference engine's mtp_update fills from true hiddens;
15061                    // the full stack (vx) is already resident from the verify. Replaces both the
15062                    // chain-approximate entries AND the old last-token-only fill. Acceptance-only
15063                    // (draft attention quality); exactness stays the verify's job.
15064                    scratch.set_len(e, pos)?;
15065                    // PREDECESSOR pairing: row i gets vx[i-1]; row 0 the carried fill_prev
15066                    // (hidden of the last committed row before this verify batch).
15067                    let mut vxs = e.zeros(t_v * n_embd)?;
15068                    e.copy_into(&mut vxs, 0, &fill_prev, n_embd)?;
15069                    if t_v > 1 {
15070                        e.copy_view_into(
15071                            &mut vxs,
15072                            n_embd,
15073                            &vx.slice(0..(t_v - 1) * n_embd),
15074                            (t_v - 1) * n_embd,
15075                        )?;
15076                    }
15077                    self.mtp_kv_fill_all(e, &verify_tokens, &vxs, pos, &mut *scratch, embd_dev)?;
15078                } else {
15079                    scratch.set_len(e, pos + base + k_round - 1)?;
15080                    // predecessor of the last draft = verify col t_v-2 (or fill_prev at t_v==1)
15081                    let mut hp = e.zeros(n_embd)?;
15082                    if t_v >= 2 {
15083                        e.copy_view_into(
15084                            &mut hp,
15085                            0,
15086                            &vx.slice((t_v - 2) * n_embd..(t_v - 1) * n_embd),
15087                            n_embd,
15088                        )?;
15089                    } else {
15090                        e.copy_into(&mut hp, 0, &fill_prev, n_embd)?;
15091                    }
15092                    self.mtp_kv_fill_all(
15093                        e,
15094                        &[draft[k_round - 1]],
15095                        &hp,
15096                        pos + base + k_round - 1,
15097                        &mut *scratch,
15098                        embd_dev,
15099                    )?;
15100                }
15101                // REFERENCE SEEDING: no pseudo pass — the next chain's step 0 IS the
15102                // reference's (id_last, h_prev) draft row; it appends the bonus's scratch
15103                // entry itself. Seed = TRUE hidden of the bonus's predecessor (last verify
15104                // col). Saves one MTP-block pass per round on top of the pairing fix.
15105                if !devacc_seeded {
15106                    e.copy_into(&mut h_seed_buf, 0, &vh_seed, n_embd)?;
15107                    e.copy_into(&mut fill_prev, 0, &vh_seed, n_embd)?;
15108                }
15109                pending = Some(bonus);
15110                if debug_spec {
15111                    eprintln!("  -> FULL ACCEPT (bonus pending, prev-h seed)");
15112                }
15113            } else if !spec_replay && base + n_acc >= 1 {
15114                // PARTIAL ACCEPT, REPLAY-FREE (2026-07-03 — the profiled #1 long-ctx spec cost):
15115                // the verify's first j = base+n_acc columns ARE the committed sequence, computed
15116                // bit-identically to eager (decode-exact contract) — so KEEP them: KV truncates to
15117                // pos+j, recurrent state rebuilds from the VerifyCkpt (same-kernel gdn prefix
15118                // re-run / pure state-clone restore), and the bonus stays PENDING exactly like the
15119                // full-accept path — the legacy duplicate trunk replay is gone. The next chain
15120                // seeds from the MTP pseudo-hidden of the bonus, whose seed = the TRUE verify
15121                // hidden of its predecessor (col j-1) — same one-hop pseudo structure as full
15122                // accept (never compounds: the next verify recomputes true hiddens for all
15123                // committed columns).
15124                let j = base + n_acc;
15125                // VERIFY-GRAPH SLAB COMMIT: when the captured trunk ran, the linear layers'
15126                // column stash was written into the graphs ctx's persistent slabs as in-graph
15127                // memcpy nodes, NOT into the per-column VerifyCkpt the cols arm reads — so the
15128                // commit must take the slab twin (same semantics, slab-addressed sources). The
15129                // ctx states which of the two this round produced via `round_slab`; trusting the
15130                // flag rather than the env keeps a round that fell back to the eager walk (a
15131                // capture that declined, a t the pool never captured) on the cols arm.
15132                let slab_commit = vg_guard
15133                    .as_ref()
15134                    .and_then(|g| g.as_ref())
15135                    .map(|g| g.round_slab)
15136                    .unwrap_or(false);
15137                if slab_commit {
15138                    self.dspark_commit_prefix_slab(
15139                        e,
15140                        &mut *cache,
15141                        &snap,
15142                        vg_guard
15143                            .as_ref()
15144                            .and_then(|g| g.as_ref())
15145                            .expect("slab_commit implies a graphs ctx"),
15146                        j,
15147                    )?;
15148                } else {
15149                    self.commit_verified_prefix(
15150                        e,
15151                        &mut *cache,
15152                        &snap,
15153                        ckpt.as_ref().unwrap(),
15154                        j,
15155                        devacc_seeded,
15156                        if devacc_seeded {
15157                            devacc_acc.as_ref().map(|a| (a, base, t_v))
15158                        } else {
15159                            None
15160                        },
15161                    )?;
15162                }
15163                let mut seed = e.zeros(n_embd)?;
15164                e.copy_view_into(
15165                    &mut seed,
15166                    0,
15167                    &vx.slice((j - 1) * n_embd..j * n_embd),
15168                    n_embd,
15169                )?;
15170                // Draft scratch: TRUE-HIDDEN REFRESH of the committed prefix (see the full-accept
15171                // branch); without it the chain entries stand and only the tail truncates. Either
15172                // way len ends at pos+j so the pseudo append lands at the bonus's slot pos+j
15173                // (persistent mode), rope pos+j+1 (chain convention).
15174                if refresh {
15175                    scratch.set_len(e, pos)?;
15176                    let mut vxs = e.zeros(j * n_embd)?;
15177                    e.copy_into(&mut vxs, 0, &fill_prev, n_embd)?;
15178                    if j > 1 {
15179                        e.copy_view_into(
15180                            &mut vxs,
15181                            n_embd,
15182                            &vx.slice(0..(j - 1) * n_embd),
15183                            (j - 1) * n_embd,
15184                        )?;
15185                    }
15186                    self.mtp_kv_fill_all(
15187                        e,
15188                        &verify_tokens[0..j],
15189                        &vxs,
15190                        pos,
15191                        &mut *scratch,
15192                        embd_dev,
15193                    )?;
15194                } else {
15195                    scratch.set_len(e, pos + j)?;
15196                }
15197                // REFERENCE SEEDING (see the full-accept branch): seed = TRUE hidden of the
15198                // bonus's predecessor (verify col j-1); no pseudo pass.
15199                if !devacc_seeded {
15200                    e.copy_into(&mut h_seed_buf, 0, &seed, n_embd)?;
15201                    e.copy_into(&mut fill_prev, 0, &seed, n_embd)?;
15202                }
15203                pending = Some(bonus);
15204                if debug_spec {
15205                    eprintln!("  -> PARTIAL(replay-free j={j}, bonus pending, prev-h seed)");
15206                }
15207            } else if !spec_replay {
15208                // ZERO ROUND FOLD (2026-07-10, verify-cost target #3): base+n_acc == 0 — a
15209                // pending-less round where nothing was accepted (PMIN0 zero-draft chains after a
15210                // replay/commit, or plain 0-accept rounds at round 0). The old path replayed
15211                // [bonus] through a FULL m=1 trunk+head forward (the 489us full-vocab head pass
15212                // measured at ~0.75/round on PMIN0 configs). Instead: restore the pre-round
15213                // snapshot and let the bonus ride the NEXT round's verify as col 0 — the existing
15214                // base=1 pending machinery, bit-identical by the decode-exact verify contract.
15215                // Seed: the bonus's predecessor is the last COMMITTED token, whose hidden
15216                // fill_prev already carries (same seeding as the 1-token-replay case it replaces).
15217                cache.rollback(e, &snap, 0)?;
15218                scratch.set_len(e, pos)?;
15219                e.copy_into(&mut h_seed_buf, 0, &fill_prev, n_embd)?;
15220                pending = Some(bonus);
15221                if debug_spec {
15222                    eprintln!("  -> ZERO-ROUND FOLD (bonus pending, fill_prev seed)");
15223                }
15224            } else {
15225                // PARTIAL ACCEPT, LEGACY REPLAY (seam MEMRA_SPEC_REPLAY=1 — or j==0: nothing of
15226                // this round survives, only possible before the first pending exists, ~round 0):
15227                // restore EVERYTHING to the pre-round snapshot (KV truncate to pos + recur
15228                // restore), then replay the committed prefix pending? ++ draft[0..n_acc] ++
15229                // [bonus] as ONE batched T forward — single weight read, bit-identical to greedy
15230                // (the verify-all-columns path is the same math). Commits the bonus with a TRUE
15231                // trunk hidden.
15232                cache.rollback(e, &snap, 0)?; // accept_len=0: KV len = pos, recur = snapshot
15233                let mut replay: Vec<u32> = Vec::with_capacity(base + n_acc + 1);
15234                if let Some(b) = pending.take() {
15235                    replay.push(b);
15236                }
15237                replay.extend_from_slice(&draft[0..n_acc]);
15238                replay.push(bonus);
15239                // Full-stack forward (decode_step_t_core = decode_step_t_h_emb_dev's body):
15240                // Predecessor pairing seeds from the PREDECESSOR row (col len-2) — the same-row path takes the
15241                // last col exactly as before (byte-identical to the old _h_emb_dev call).
15242                let (rl_d, rx) = if self.batched_serving_numeric_class() {
15243                    let mut logits = Vec::with_capacity(replay.len() * n_vocab);
15244                    let mut hidden = e.uninit(replay.len() * n_embd)?;
15245                    for (row, &token) in replay.iter().enumerate() {
15246                        let (row_logits, row_hidden) =
15247                            self.spec_target_step_h(e, token, &mut *cache)?;
15248                        logits.extend_from_slice(&row_logits);
15249                        e.dtod_copy_into(&row_hidden, &mut hidden, row * n_embd)?;
15250                    }
15251                    (e.htod(&logits)?, hidden)
15252                } else {
15253                    self.decode_step_t_core(e, &replay, pos, &mut *cache, embd_dev, None)?
15254                };
15255                // last_pred = argmax of the LAST column's logits (predicts the token after `bonus`)
15256                // — device argmax + one 4-byte read instead of the full-vocab column dtoh.
15257                e.argmax_token_device_col(&rl_d, replay.len() - 1, n_vocab, &mut preds_d, 0)?;
15258                last_pred = guard_vocab_token(
15259                    e.dtoh_u32(&preds_d)?[0],
15260                    n_vocab,
15261                    &format!("replay last_pred at round {round} pos={pos}"),
15262                )?;
15263                if sampled {
15264                    let lr0 = replay.len();
15265                    let lc = last_col_logits
15266                        .as_mut()
15267                        .expect("sampled: last_col_logits unset");
15268                    e.copy_view_into(
15269                        lc,
15270                        0,
15271                        &rl_d.slice((lr0 - 1) * n_vocab..lr0 * n_vocab),
15272                        n_vocab,
15273                    )?;
15274                }
15275                let lr = replay.len();
15276                if lr >= 2 {
15277                    e.copy_view_into(
15278                        &mut h_seed_buf,
15279                        0,
15280                        &rx.slice((lr - 2) * n_embd..(lr - 1) * n_embd),
15281                        n_embd,
15282                    )?;
15283                } else {
15284                    // 1-token replay (round-0 miss): the bonus's predecessor is the OLD
15285                    // last_token, whose own-row hidden fill_prev still holds.
15286                    e.copy_into(&mut h_seed_buf, 0, &fill_prev, n_embd)?;
15287                }
15288                // the bonus is COMMITTED here — it becomes the last committed row.
15289                let mut rh_last = e.zeros(n_embd)?;
15290                e.copy_view_into(
15291                    &mut rh_last,
15292                    0,
15293                    &rx.slice((lr - 1) * n_embd..lr * n_embd),
15294                    n_embd,
15295                )?;
15296                e.copy_into(&mut fill_prev, 0, &rh_last, n_embd)?;
15297                if debug_spec {
15298                    eprintln!("  -> PARTIAL(replay={replay:?}), next_pred={last_pred}");
15299                }
15300            }
15301            if devacc_seeded {
15302                // stage (b) epilogue: fill_prev takes the gathered seed AFTER the refresh fills
15303                // consumed the old value (both slots carry the same value in every non-replay arm).
15304                e.copy_into(&mut fill_prev, 0, &h_seed_buf, n_embd)?;
15305            }
15306            if successor_valid {
15307                let optimistic_scratch_len = successor_attempt
15308                    .as_ref()
15309                    .expect("valid controller successor disappeared")
15310                    .scratch_len;
15311                // The normal current-round commit refreshed/truncated the logical scratch tail.
15312                // Its optimistic successor row was already written physically, so restoring only
15313                // the retained logical length makes that row live for the carried round.
15314                scratch.set_len(e, optimistic_scratch_len)?;
15315            }
15316            if let Some(current) = current_opti.take() {
15317                opti_fork
15318                    .as_mut()
15319                    .ok_or("optipipe current retirement lost fork state")?
15320                    .retire(current.generation)?;
15321            }
15322            if successor_valid {
15323                let successor = successor_attempt
15324                    .take()
15325                    .expect("valid controller successor disappeared before promotion");
15326                let generation = successor.generation;
15327                opti_fork
15328                    .as_mut()
15329                    .ok_or("optipipe successor promotion lost fork state")?
15330                    .promote_successor_snapshot(&mut snap, generation);
15331                carried_opti = Some(successor);
15332            }
15333            if anatomy_on {
15334                // Commit/rollback is normally asynchronous on the primary/head stream. Bound it
15335                // only for this diagnostic so it does not disappear into the following draft's
15336                // first token readback.
15337                e.stream().synchronize()?;
15338                ph_commit += commit_started.elapsed().as_secs_f64();
15339            }
15340            // adaptive-K update (host math, zero syncs): next round drafts accepted-run + 1,
15341            // clamped to [floor(pos), k_cap]. cache.pos is post-rollback here (the round's
15342            // final position — the floor's position key reads the committed depth). Burst
15343            // rounds (`continue` above) draft the captured fixed depth and skip this, exactly
15344            // like gemma's burst arm.
15345            if adapt {
15346                let fl_now = floor_at(cache.pos);
15347                kc = (n_acc + 1).clamp(fl_now.min(k_cap), k_cap);
15348            }
15349            ph_mark(&mut ph_rest, phase_on);
15350            if let Some(p) = pipe {
15351                p.accept_end(round);
15352            }
15353            drop(pipe_accept);
15354            if let Some(t0) = round_t0 {
15355                let ms = t0.elapsed().as_secs_f64() * 1e3;
15356                ROUND_MS.fetch_add((ms * 1e3) as u64, std::sync::atomic::Ordering::Relaxed);
15357                let n = ROUND_N.fetch_add(1, std::sync::atomic::Ordering::Relaxed) + 1;
15358                if n % 32 == 0 {
15359                    eprintln!(
15360                        "[spec-round] rounds={n} avg round wall={:.2} ms (emitted={} drafted so far)",
15361                        ROUND_MS.load(std::sync::atomic::Ordering::Relaxed) as f64 / 1e3 / n as f64,
15362                        out.len()
15363                    );
15364                }
15365            }
15366            round += 1;
15367            // sse-cadence: this round's accepted drafts + bonus are committed (out is
15368            // append-only past step 4) — flush at round cadence.
15369            keep_going = flush_commit(&mut on_commit, &out, &mut flushed);
15370        }
15371        if let Some(mut ticket) = carried_opti.take() {
15372            opti_fork
15373                .as_mut()
15374                .ok_or("optipipe tail drain lost fork state")?
15375                .cancel_controller_ticket(e, &mut *cache, &mut *scratch, &snap, &mut ticket)?;
15376        }
15377        // sse-cadence: nothing below appends to `out`; flush any remainder (defensive).
15378        // (verdict ignored — the burst is over either way; the session tail runs unchanged.)
15379        let _ = flush_commit(&mut on_commit, &out, &mut flushed);
15380
15381        if spec_stats {
15382            let per_slot: Vec<String> = (0..k)
15383                .map(|j| {
15384                    if st_drafted[j] > 0 {
15385                        format!(
15386                            "{}/{}={:.3}",
15387                            st_accepted[j],
15388                            st_drafted[j],
15389                            st_accepted[j] as f64 / st_drafted[j] as f64
15390                        )
15391                    } else {
15392                        "0/0".into()
15393                    }
15394                })
15395                .collect();
15396            let acc = if total_drafted > 0 {
15397                total_accepted as f64 / total_drafted as f64
15398            } else {
15399                0.0
15400            };
15401            eprintln!(
15402                "[spec-stats] rounds={round} full_accept={st_full} len_hist={st_len_hist:?} \
15403                       per_slot=[{}] total={total_accepted}/{total_drafted}={acc:.3} \
15404                       tok_per_round={:.3}",
15405                per_slot.join(" "),
15406                (total_accepted + round) as f64 / round.max(1) as f64
15407            );
15408        }
15409        if constraint.is_some() {
15410            eprintln!(
15411                "[draft-mask] mask_rounds={dm_rounds} clone_total={:.3}ms \
15412                 clone_per_round={:.4}ms gram_cuts={dm_cuts}/{round} cut_tokens={dm_cut_tokens}",
15413                dm_clone_ns as f64 / 1e6,
15414                dm_clone_ns as f64 / 1e6 / dm_rounds.max(1) as f64
15415            );
15416        }
15417        if phase_on {
15418            let tot = ph_draft + ph_verify + ph_wait + ph_rest;
15419            eprintln!(
15420                "[spec-phase] draft={:.1}ms ({:.1}%) verify-issue={:.1}ms ({:.1}%) verify-wait={:.1}ms ({:.1}%) commit-host={:.1}ms ({:.1}%) rounds={round}",
15421                ph_draft * 1e3,
15422                ph_draft / tot * 100.0,
15423                ph_verify * 1e3,
15424                ph_verify / tot * 100.0,
15425                ph_wait * 1e3,
15426                ph_wait / tot * 100.0,
15427                ph_rest * 1e3,
15428                ph_rest / tot * 100.0
15429            );
15430        }
15431        if anatomy_on {
15432            let rounds_f = round.max(1) as f64;
15433            let other = (ph_rest - ph_commit).max(0.0);
15434            eprintln!(
15435                "[spec-anatomy] per-round draft={:.3}ms pp-verify={:.3}ms \
15436                 verify-accept={:.3}ms commit-rollback={:.3}ms other={:.3}ms rounds={round}",
15437                ph_draft * 1e3 / rounds_f,
15438                ph_verify * 1e3 / rounds_f,
15439                ph_wait * 1e3 / rounds_f,
15440                ph_commit * 1e3 / rounds_f,
15441                other * 1e3 / rounds_f,
15442            );
15443        }
15444        let _pipe_tail = pipe.map(|p| p.primary());
15445        // SESSION TAIL: leave the session in the exact invariant the next turn's suffix prime
15446        // expects — every row in `committed` has trunk KV/recur state AND an exact draft-KV row.
15447        // Park the draft-graph ctx back on the session (the serve-burst fixed-cost fix): the next
15448        // burst replays instead of recapturing. Error paths (`?` above) drop it — recaptured then.
15449        if let Some(slot) = sess_draft_slot.take() {
15450            *slot = Some(dctx);
15451        }
15452        let t_rounds = t_ent.elapsed();
15453        if let Some((committed, last_h, next_pred_slot, sctr_slot, uctr_slot)) = sess_tail.take() {
15454            // NEXT BURST'S BOUNDARY TOKEN (lane/sampled-spec-quality, Item 1). Greedy stashes
15455            // the argmax `last_pred` exactly as before (byte contract). SAMPLED draws the token
15456            // HERE, where the sampler, the session Philox counters and the penalty window are
15457            // all live and the boundary logits row still exists — that is the "make the state
15458            // available" half of the fix; the consuming burst then just emits it. `sctr` is
15459            // written to the session BELOW the draws so the advance is never lost.
15460            *next_pred_slot = Some(last_pred);
15461            let sample_boundary = sampled && constraint.is_none() && spec_sampled_boundary_on();
15462            let mut stashed_pending = false;
15463            if let Some(b) = pending.take() {
15464                if !sampled {
15465                    // PENDING-CARRY (2026-08-01): stash the bonus on the session instead of
15466                    // committing it with a solo T=1 pass — the next empty-suffix greedy burst
15467                    // consumes it as round-0 verify col 0 (a plain round edge; the old tail
15468                    // commit + next burst's init feed were 11.6+11.5ms solo trunk passes per
15469                    // burst on H100 q27, [spec-setup] trace). b stays in `out` (emitted) but
15470                    // OUT of `committed` (cache rows == committed); the consuming call
15471                    // prepends it once its verify commits the row. next_pred is unknowable
15472                    // without the commit pass — None; callers gate on pending_tok too.
15473                    debug_assert_eq!(out.last(), Some(&b), "pending must be the last emitted");
15474                    if let Some(slot) = sess_pending_slot.take() {
15475                        *slot = Some(b);
15476                    }
15477                    *next_pred_slot = None;
15478                    // fill_prev = hidden of the last COMMITTED row (b's predecessor) — the
15479                    // exact chain-seed/fill anchor the consuming burst (or a flush) needs.
15480                    *last_h = Some(e.clone_dtod(&fill_prev)?);
15481                    stashed_pending = true;
15482                } else {
15483                    // SAMPLED tail (unchanged): commit the bonus (one T=1 pass) + draft fill —
15484                    // the sampled round-0 accept needs this pass's logits (last_col_logits).
15485                    let pos_b = cache.pos;
15486                    scratch.set_len(e, pos_b)?;
15487                    let (lg_b, hb) = self.spec_target_step_h(e, b, &mut *cache)?;
15488                    // after a FULL-accept exit `last_pred` is STALE (it predicted the bonus
15489                    // itself — the prediction AFTER the bonus never materialized; it would have
15490                    // been the next round's verify col 0). The commit's logits ARE that
15491                    // prediction — so they are also the row the next burst's boundary token
15492                    // comes off, and (lane/sampled-spec-quality) it is DRAWN from them here.
15493                    *next_pred_slot = Some(if sample_boundary {
15494                        sample_boundary_token(
15495                            e,
15496                            &lg_b,
15497                            &sp,
15498                            &pen_hist,
15499                            &mut sctr,
15500                            "burst-tail-commit",
15501                        )?
15502                    } else {
15503                        argmax(&lg_b) as u32
15504                    });
15505                    self.mtp_kv_fill_all(e, &[b], &fill_prev, pos_b, &mut *scratch, embd_dev)?;
15506                    *last_h = Some(hb);
15507                }
15508            } else {
15509                // fill_prev tracks the hidden of the last COMMITTED row throughout the loop.
15510                *last_h = Some(e.clone_dtod(&fill_prev)?);
15511                if sample_boundary {
15512                    // No pending to commit, so the boundary row is the one `last_pred` was
15513                    // argmaxed from and the sampled path keeps it on device: the init feed's
15514                    // logits when the burst ran zero rounds, else the legacy-replay path's
15515                    // last verify column (both predict the token AFTER the last committed
15516                    // row). It is retained precisely because round 0's accept test needs it,
15517                    // so the draw costs no extra D2H of the [n_vocab] row.
15518                    match last_col_logits.as_ref() {
15519                        Some(lc) => {
15520                            *next_pred_slot = Some(sample_boundary_token_dev(
15521                                e,
15522                                lc,
15523                                n_vocab,
15524                                &sp,
15525                                &pen_hist,
15526                                &mut sctr,
15527                                "burst-tail-nopending",
15528                            )?);
15529                        }
15530                        // NAME THE FALLBACK (house standard): unreachable today — a sampled
15531                        // burst always feeds or replays, so the row exists — but if it ever
15532                        // is, the stream takes a greedy token and SAYS so rather than
15533                        // silently regressing to the pre-lane behaviour.
15534                        None => eprintln!(
15535                            "[spec-boundary] sampled tail kept the ARGMAX boundary token \
15536                             (reason: no retained boundary logits row)"
15537                        ),
15538                    }
15539                }
15540            }
15541            *sctr_slot = sctr;
15542            *uctr_slot = uctr;
15543            committed.extend_from_slice(prompt);
15544            if let Some(cb) = carried_pending {
15545                // the consumed carry's cache row landed in round 0's verify (every pending
15546                // round commits col 0) — it joins `committed` here, in sequence order.
15547                committed.push(cb);
15548            }
15549            if stashed_pending {
15550                // ZERO-EMIT BURST (2026-08-06 c=8 serve panic, pre-existing since b4aea184):
15551                // `out.len() - 1` underflowed on an EMPTY `out` — "range end index
15552                // 18446744073709551615 out of range for slice of length 0", killing the
15553                // memra-gpu-worker and failing 31 of 32 concurrent requests with "worker closed
15554                // stream". Reachable because `pending` starts as `carried_pending` (a bonus
15555                // stashed by the PREVIOUS burst) while `out` starts empty, and the carry is
15556                // deliberately NOT pushed to `out` (line ~3239: the burst that emitted it already
15557                // did). So a burst that stashes a pending without emitting anything of its own —
15558                // the round loop exits before a push, e.g. the ring drain's `out.len() < max_new`
15559                // guard skipping every token under a tight budget — arrives here with
15560                // out.len() == 0 and stashed_pending == true.
15561                //
15562                // The invariant is unchanged: `committed` gets every emitted token EXCEPT the
15563                // stashed bonus. With nothing emitted, that is nothing — and the carry pushed
15564                // just above is already accounted. Saturating, not a min/assert: an empty `out`
15565                // here is a legitimate burst shape, not a corrupt state.
15566                let emitted = out.len().saturating_sub(1);
15567                committed.extend_from_slice(&out[..emitted]);
15568            } else {
15569                committed.extend_from_slice(&out); // FULL out incl. overshoot — all committed
15570            }
15571            debug_assert_eq!(
15572                cache.pos,
15573                committed.len(),
15574                "session invariant: cache rows == committed tokens"
15575            );
15576            if setup_trace {
15577                e.stream().synchronize()?; // bound the async tail fill in the trace
15578                let t_tail = t_ent.elapsed();
15579                eprintln!(
15580                    "[spec-setup] init={:.2}ms cap={:.2}ms fill={:.2}ms rounds={:.2}ms tail={:.2}ms total={:.2}ms out={} cont={}",
15581                    t_init.as_secs_f64() * 1e3,
15582                    (t_cap - t_init).as_secs_f64() * 1e3,
15583                    (t_fill - t_cap).as_secs_f64() * 1e3,
15584                    (t_rounds - t_fill).as_secs_f64() * 1e3,
15585                    (t_tail - t_rounds).as_secs_f64() * 1e3,
15586                    t_tail.as_secs_f64() * 1e3,
15587                    out.len(),
15588                    continuation
15589                );
15590            }
15591            return Ok((out, total_drafted, total_accepted));
15592        }
15593        out.truncate(max_new);
15594        Ok((out, total_drafted, total_accepted))
15595    }
15596
15597    /// Anchor-bounded DSpark target extraction. The trunk sees the exact generated token tape;
15598    /// only requested hidden rows and target-logit rows cross PCIe. An anchor token at p pairs
15599    /// with the pre-output-norm h[p-1] carrier, exactly as the existing replay/NextN path does.
15600    pub fn extract_dspark_anchors(
15601        &self,
15602        e: &Engine,
15603        tokens: &[u32],
15604        anchor_positions: &[usize],
15605        gamma: usize,
15606        top_k: usize,
15607        chunk: usize,
15608        temperature: f32,
15609    ) -> Result<Vec<DsparkAnchorRecord>, Box<dyn std::error::Error>> {
15610        if tokens.len() < gamma + 2 || gamma == 0 || chunk < 2 {
15611            return Err("DSpark extraction token tape/gamma/chunk is invalid".into());
15612        }
15613        if anchor_positions.windows(2).any(|pair| pair[0] >= pair[1]) {
15614            return Err("DSpark anchor positions must be sorted and unique".into());
15615        }
15616        for &position in anchor_positions {
15617            if position == 0 || position + gamma >= tokens.len() {
15618                return Err(format!(
15619                    "DSpark anchor {position} has no predecessor or cannot cover gamma={gamma} in {} tokens",
15620                    tokens.len()
15621                )
15622                .into());
15623            }
15624        }
15625
15626        let n_vocab = self.output.out_features();
15627        let n_embd = self.cfg.n_embd as usize;
15628        let mut cache =
15629            crate::pp::new_cache_planned(e, &self.cfg, &self.plan, tokens.len() + gamma + 8)?;
15630        let (embd_qt, embd_rb) = self.embd.qt_and_row_bytes(n_embd);
15631        let embd_gpu = if spec_host_embd() {
15632            None
15633        } else {
15634            Some(
15635                self.embd_gpu
15636                    .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload")),
15637            )
15638        };
15639        let embd_dev = embd_gpu.map(|gpu| (gpu, embd_qt, embd_rb));
15640
15641        struct PendingRecord {
15642            position: usize,
15643            hidden: Option<Vec<f32>>,
15644            tokens: Vec<u32>,
15645            target_top_ids: Vec<Option<Vec<u32>>>,
15646            target_top_logits: Vec<Option<Vec<f32>>>,
15647            target_top_probs: Vec<Option<Vec<f32>>>,
15648            target_tail_probs: Vec<Option<f32>>,
15649        }
15650
15651        let mut pending: Vec<PendingRecord> = anchor_positions
15652            .iter()
15653            .map(|&position| PendingRecord {
15654                position,
15655                hidden: None,
15656                tokens: tokens[position..=position + gamma].to_vec(),
15657                target_top_ids: vec![None; gamma],
15658                target_top_logits: vec![None; gamma],
15659                target_top_probs: vec![None; gamma],
15660                target_tail_probs: vec![None; gamma],
15661            })
15662            .collect();
15663
15664        let mut start = 0usize;
15665        while start < tokens.len() {
15666            let end = (start + chunk).min(tokens.len());
15667            let chunk_tokens = &tokens[start..end];
15668            let (target_logits, hidden_rows) =
15669                self.decode_step_t_core(e, chunk_tokens, start, &mut cache, embd_dev, None)?;
15670            for record in &mut pending {
15671                let hidden_position = record.position - 1;
15672                if hidden_position >= start && hidden_position < end {
15673                    let local = hidden_position - start;
15674                    record.hidden = Some(
15675                        e.dtoh_view(&hidden_rows.slice(local * n_embd..(local + 1) * n_embd))?,
15676                    );
15677                }
15678                for slot in 0..gamma {
15679                    let target_row = record.position + slot;
15680                    if target_row < start || target_row >= end {
15681                        continue;
15682                    }
15683                    let local = target_row - start;
15684                    let logits =
15685                        e.dtoh_view(&target_logits.slice(local * n_vocab..(local + 1) * n_vocab))?;
15686                    let (ids, top_logits, probs, tail) =
15687                        dspark_sparse_softmax_topk(&logits, top_k, temperature)?;
15688                    record.target_top_ids[slot] = Some(ids);
15689                    record.target_top_logits[slot] = Some(top_logits);
15690                    record.target_top_probs[slot] = Some(probs);
15691                    record.target_tail_probs[slot] = Some(tail);
15692                }
15693            }
15694            start = end;
15695        }
15696
15697        pending
15698            .into_iter()
15699            .map(|record| {
15700                let hidden = record
15701                    .hidden
15702                    .ok_or_else(|| format!("missing DSpark hidden at {}", record.position))?;
15703                let target_top_ids =
15704                    flatten_dspark_rows(record.target_top_ids, record.position, "target ids")?;
15705                let target_top_logits = flatten_dspark_rows(
15706                    record.target_top_logits,
15707                    record.position,
15708                    "target logits",
15709                )?;
15710                let target_top_probs =
15711                    flatten_dspark_rows(record.target_top_probs, record.position, "target probs")?;
15712                let target_tail_probs = record
15713                    .target_tail_probs
15714                    .into_iter()
15715                    .enumerate()
15716                    .map(|(slot, value)| {
15717                        value.ok_or_else(|| {
15718                            format!("missing DSpark tail at {} slot {slot}", record.position)
15719                        })
15720                    })
15721                    .collect::<Result<Vec<_>, _>>()?;
15722                Ok(DsparkAnchorRecord {
15723                    position: record.position,
15724                    hidden,
15725                    tokens: record.tokens,
15726                    target_top_ids,
15727                    target_top_logits,
15728                    target_top_probs,
15729                    target_tail_probs,
15730                })
15731            })
15732            .collect()
15733    }
15734
15735    /// TEACHER-FORCED REPLAY ACCEPTANCE (hqmtp MTP-heal protocol): walk a FIXED token
15736    /// sequence and, at sampled positions, compare the MTP head's K-token draft chain against
15737    /// the trunk's own teacher-forced greedy predictions. Nothing is generated — the context is
15738    /// the corpus text itself, so (a) degenerate self-generated loops cannot inflate acceptance
15739    /// and (b) two arms (bf16 ceiling vs NVFP4) score on IDENTICAL contexts, isolating the
15740    /// quant-induced head/hidden-state mismatch from text drift.
15741    ///
15742    /// Per eval position p (context = tokens[0..=p], predecessor pairing as in spec decode):
15743    ///   draft_j  = chain token j from (tokens[p], h_{p-1}), then its own drafts — the exact
15744    ///              eager spec-decode chain (same mtp_head_forward_dev, same rope positions).
15745    ///   target_j = teacher-forced greedy pick for position p+1+j (argmax of the trunk logits
15746    ///              at forced context tokens[0..p+j]). For j==0 this equals live spec
15747    ///              acceptance; for j>=1 live verify would condition on the drafts, here it
15748    ///              conditions on the corpus — deterministic and arm-comparable by design.
15749    ///
15750    /// Returns (rows, bg): one (p, drafts[k], targets[k]) row per eval position (ascending p),
15751    /// plus the full teacher-forced greedy track bg (bg[i] = greedy pick for position i, i>=1)
15752    /// so harnesses can cross-check runs (e.g. different chunk sizes must give identical bg).
15753    ///
15754    /// `hdump`: when Some, every position's pre-output_norm trunk hidden (the exact rows the
15755    /// draft-KV fill pairs from) streams to the file as little-endian f32 [t_total, n_embd] —
15756    /// the head-distillation extraction (hqmtp): the ENGINE is the source of truth for trunk
15757    /// hiddens (HF torch reproductions of the hybrid trunk measured only ~0.5 greedy
15758    /// agreement vs this path — not usable as a training-data source).
15759    pub fn replay_acceptance(
15760        &self,
15761        e: &Engine,
15762        tokens: &[u32],
15763        k: usize,
15764        stride: usize,
15765        chunk: usize,
15766        mut hdump: Option<&mut std::fs::File>,
15767    ) -> Result<(Vec<(usize, Vec<u32>, Vec<u32>)>, Vec<u32>), Box<dyn std::error::Error>> {
15768        assert!(k >= 1 && stride >= 1 && chunk >= 2);
15769        let mtp = self
15770            .mtp
15771            .as_ref()
15772            .expect("replay_acceptance requires an MTP head");
15773        let n_vocab = self.output.out_features();
15774        let d_vocab = mtp
15775            .shared_head_head
15776            .as_ref()
15777            .unwrap_or(&self.output)
15778            .out_features();
15779        let n_embd = self.cfg.n_embd as usize;
15780        let t_total = tokens.len();
15781        assert!(t_total >= 8, "corpus too short ({t_total} tokens)");
15782        // STAGE-OWNED KV (lane/pp2-spec 2026-08-06) — see `new_session`. Door shut = `Cache::new`.
15783        let mut cache = crate::pp::new_cache_planned(e, &self.cfg, &self.plan, t_total + k + 8)?;
15784        let mut scratch = self.new_mtp_scratch(e, t_total + k + 8)?;
15785        let (embd_qt, embd_rb) = self.embd.qt_and_row_bytes(n_embd);
15786        let embd_gpu = if spec_host_embd() {
15787            None
15788        } else {
15789            Some(
15790                self.embd_gpu
15791                    .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload")),
15792            )
15793        };
15794        let embd_dev = embd_gpu.map(|g| (g, embd_qt, embd_rb));
15795
15796        // bg[i] = the trunk's greedy pick for position i under the forced context (i >= 1).
15797        let mut bg: Vec<u32> = vec![0; t_total + 1];
15798        let mut rows: Vec<(usize, Vec<u32>, Vec<u32>)> = Vec::new();
15799        let mut prev_last_h = e.zeros(n_embd)?; // predecessor hidden entering the chunk
15800        let mut seed_buf = e.zeros(n_embd)?;
15801        let mut preds_d = e.alloc_u32_zeroed(chunk)?;
15802        let nll_on = std::env::var("MEMRA_REPLAY_NLL").as_deref() == Ok("1");
15803        let (mut nll_sum, mut nll_cnt) = (0f64, 0u64);
15804        let mut s = 0usize;
15805        while s < t_total {
15806            let cend = (s + chunk).min(t_total);
15807            let tc = cend - s;
15808            let ch = &tokens[s..cend];
15809            // 1. forced trunk pass — verify path (decode-exact contract): all-column logits +
15810            //    the chunk's true hiddens.
15811            let (tl_d, vx) = self.decode_step_t_core(e, ch, s, &mut cache, embd_dev, None)?;
15812            for j in 0..tc {
15813                e.argmax_token_device_col(&tl_d, j, n_vocab, &mut preds_d, j)?;
15814            }
15815            let preds = e.dtoh_u32(&preds_d)?;
15816            for j in 0..tc {
15817                bg[s + j + 1] = preds[j];
15818            }
15819            // MEMRA_REPLAY_NLL=1: teacher-forced NLL/perplexity over the same forced pass — the
15820            // checkpoint-quality metric (position j's logits score the GOLD next token).
15821            if nll_on {
15822                let jmax = if cend < t_total { tc } else { tc - 1 }; // last pos has no gold next
15823                if jmax > 0 {
15824                    let ids: Vec<u32> = (0..jmax).map(|j| tokens[s + j + 1]).collect();
15825                    let rows: Vec<i32> = (0..jmax as i32).collect();
15826                    let idsd = e.htod_u32_v(&ids)?;
15827                    let rowsd = e.htod_i32(&rows)?;
15828                    let mut outd = e.zeros(jmax)?;
15829                    e.softmax_gather(&tl_d, n_vocab, &idsd, &rowsd, &mut outd, n_vocab, jmax, 1.0)?;
15830                    for pr in e.dtoh(&outd)? {
15831                        nll_sum += -((pr.max(1e-30)) as f64).ln();
15832                        nll_cnt += 1;
15833                    }
15834                }
15835            }
15836            if let Some(f) = hdump.as_deref_mut() {
15837                use std::io::Write;
15838                let host: Vec<f32> = e.dtoh(&vx)?;
15839                // bf16 round-to-nearest-even — f32 doubled the disk bill at bulk
15840                // extraction scale (20M tokens x 4096 = 320GB f32 vs 160GB bf16).
15841                let mut bytes = Vec::with_capacity(tc * n_embd * 2);
15842                for v in &host[..tc * n_embd] {
15843                    let b = v.to_bits();
15844                    let r = b.wrapping_add(0x7FFF + ((b >> 16) & 1));
15845                    bytes.extend_from_slice(&((r >> 16) as u16).to_le_bytes());
15846                }
15847                f.write_all(&bytes)?;
15848            }
15849            // CHAINLESS extraction (stride > corpus, the bulk-hdump mode): no chunk ever
15850            // drafts, so the draft-KV fills are pure waste — skip them (2 MTP-block passes
15851            // per token saved; the forced trunk pass + hdump is all the mode needs).
15852            let chainless = stride > t_total;
15853            if chainless {
15854                e.copy_view_into(
15855                    &mut prev_last_h,
15856                    0,
15857                    &vx.slice((tc - 1) * n_embd..tc * n_embd),
15858                    n_embd,
15859                )?;
15860                s = cend;
15861                continue;
15862            }
15863            // 2. TRUE predecessor-paired draft-KV fill for the chunk (row i carries h_{i-1};
15864            //    row s reads the previous chunk's last true hidden, zeros at corpus start).
15865            let mut vxs = e.zeros(tc * n_embd)?;
15866            e.copy_into(&mut vxs, 0, &prev_last_h, n_embd)?;
15867            if tc > 1 {
15868                e.copy_view_into(
15869                    &mut vxs,
15870                    n_embd,
15871                    &vx.slice(0..(tc - 1) * n_embd),
15872                    (tc - 1) * n_embd,
15873                )?;
15874            }
15875            scratch.set_len(e, s)?;
15876            self.mtp_kv_fill_all(e, ch, &vxs, s, &mut scratch, embd_dev)?;
15877            // 3. draft chains at sampled positions, DESCENDING: a chain reads only slots
15878            //    [0..p) (true fills) and appends at >= p; the next (smaller-p) chain's set_len
15879            //    truncates those approximate appends before they can ever be read.
15880            let ps: Vec<usize> = (s..cend)
15881                .filter(|p| *p >= 1 && *p % stride == 0 && *p + k <= t_total)
15882                .collect();
15883            for &p in ps.iter().rev() {
15884                scratch.set_len(e, p)?;
15885                if p == s {
15886                    e.copy_into(&mut seed_buf, 0, &prev_last_h, n_embd)?;
15887                } else {
15888                    e.copy_view_into(
15889                        &mut seed_buf,
15890                        0,
15891                        &vx.slice((p - 1 - s) * n_embd..(p - s) * n_embd),
15892                        n_embd,
15893                    )?;
15894                }
15895                let mut e_tok = tokens[p];
15896                let mut d_seed = e.clone_dtod(&seed_buf)?;
15897                let chain_heads = !self.mtp_extra.is_empty();
15898                let mut chain_tokens = if chain_heads {
15899                    vec![tokens[p]]
15900                } else {
15901                    Vec::new()
15902                };
15903                let mut chain_seeds = if chain_heads {
15904                    vec![e.clone_dtod(&seed_buf)?]
15905                } else {
15906                    Vec::new()
15907                };
15908                let mut drafts: Vec<u32> = Vec::with_capacity(k);
15909                for j in 0..k {
15910                    let (dl_d, h_nextn) = if chain_heads {
15911                        self.mtp_chain_forward_dev(
15912                            e,
15913                            &chain_tokens,
15914                            &chain_seeds,
15915                            &mut scratch,
15916                            p,
15917                            embd_dev,
15918                            None,
15919                        )?
15920                    } else {
15921                        self.mtp_head_forward_dev(
15922                            e,
15923                            mtp,
15924                            e_tok,
15925                            &d_seed,
15926                            &mut scratch,
15927                            p + 1 + j,
15928                            embd_dev,
15929                            None,
15930                        )?
15931                    };
15932                    let tok_d = e.argmax_token_device(&dl_d, d_vocab)?;
15933                    let idx = e.dtoh_u32_one(&tok_d)?;
15934                    let d = match &mtp.d2t {
15935                        Some(map) => map[idx as usize],
15936                        None => idx,
15937                    };
15938                    drafts.push(d);
15939                    if chain_heads {
15940                        chain_tokens.push(d);
15941                        chain_seeds.push(h_nextn);
15942                    } else {
15943                        e_tok = d;
15944                        d_seed = h_nextn;
15945                    }
15946                }
15947                // targets may live in a LATER chunk's bg — resolved after the walk.
15948                rows.push((p, drafts, Vec::new()));
15949            }
15950            // 4. restore TRUE entries for the whole chunk (the next chunk's chains and fills
15951            //    expect scratch.len == cend with exact rows).
15952            scratch.set_len(e, s)?;
15953            self.mtp_kv_fill_all(e, ch, &vxs, s, &mut scratch, embd_dev)?;
15954            e.copy_view_into(
15955                &mut prev_last_h,
15956                0,
15957                &vx.slice((tc - 1) * n_embd..tc * n_embd),
15958                n_embd,
15959            )?;
15960            s = cend;
15961        }
15962        for (p, drafts, targets) in rows.iter_mut() {
15963            for j in 0..drafts.len() {
15964                targets.push(bg[*p + 1 + j]);
15965            }
15966        }
15967        rows.sort_by_key(|r| r.0);
15968        if nll_cnt > 0 {
15969            let mean = nll_sum / nll_cnt as f64;
15970            println!(
15971                "[replay-nll] tokens={nll_cnt} nll/token={mean:.5} ppl={:.4}",
15972                mean.exp()
15973            );
15974        }
15975        Ok((rows, bg))
15976    }
15977}
15978
15979#[cfg(test)]
15980mod vg_debt_tests {
15981    use super::dspark_vg_debt_projection;
15982
15983    /// TOOTH for the verify-graph admission accounting: the pool's projected remaining
15984    /// growth must be charged (pre-fix, admission charged 0 for a pool measured at
15985    /// 8,852 MiB), the projection must price the MARGINAL cost of one more key rather than
15986    /// extrapolating the pool's one-time shared allocation, and the doors that make growth
15987    /// impossible must zero the debt.
15988    #[test]
15989    fn vg_debt_projects_remaining_growth_and_respects_the_freeze_valves() {
15990        const MIB: usize = 1 << 20;
15991        let d = dspark_vg_debt_projection;
15992        // cold pool: nothing observed, one capture fits inside SPEC_SHRINK_RESERVE.
15993        assert_eq!(d(0, 256, 0, None), 0);
15994        // freeze valve MEMRA_DSPARK_VG_MAX=0: the pool cannot grow.
15995        assert_eq!(d(10, 0, 500 * MIB, None), 0);
15996        // saturated pool: at/past the cap the pool FREEZES, nothing left to reserve.
15997        assert_eq!(d(256, 256, 8852 * MIB, None), 0);
15998        assert_eq!(d(300, 256, 8852 * MIB, None), 0);
15999
16000        // BOOTSTRAP (one observation, growth unmeasurable): at most one more pool's worth.
16001        // The pre-fix mean rule extrapolated 255x here — the measured 8.5 GB phantom.
16002        assert_eq!(d(1, 256, 33 * MIB, None), 33 * MIB);
16003
16004        // MARGINAL, flat pool (the box9 receipt: reserved stayed ~33.6 MiB across captures
16005        // 1..3, so an additional key costs ~nothing and the debt must collapse to ~0 —
16006        // NOT the 8,556/4,261/2,830 MB the mean rule printed).
16007        assert_eq!(d(3, 256, 33 * MIB, Some((1, 33 * MIB))), 0);
16008
16009        // MARGINAL, genuinely growing pool: 40 MiB per new key over 2 keys, 250 slots left.
16010        let debt = d(6, 256, 273 * MIB, Some((4, 193 * MIB)));
16011        assert_eq!(debt, 250 * (40 * MIB));
16012        assert!(
16013            debt > 3 * (1536 * MIB),
16014            "real growth must dwarf SPEC_SHRINK_RESERVE"
16015        );
16016
16017        // a shrinking/recycled reading never becomes a negative charge.
16018        assert_eq!(d(6, 256, 10 * MIB, Some((4, 99 * MIB))), 0);
16019        // a stale observation at the same capture count falls back to bootstrap.
16020        assert_eq!(d(4, 256, 80 * MIB, Some((4, 80 * MIB))), 80 * MIB);
16021    }
16022}
16023
16024#[cfg(test)]
16025mod capture_headroom_tests {
16026    use super::{
16027        CAPTURE_HEADROOM_FLOOR, capture_err_is_oom, capture_headroom_verdict,
16028        draft_capture_bootstrap_estimate,
16029    };
16030
16031    /// TOOTH for the pre-capture reserve check (lane/step37-vram-admission-20260830): a
16032    /// capture attempt must be refused BEFORE it allocates when the device cannot cover its
16033    /// appetite plus the post-capture floor — and pool-cached bytes count as headroom
16034    /// (driver `free` alone under-counts, the wrong direction for a gate that drops
16035    /// coverage).
16036    #[test]
16037    fn capture_reserve_check_refuses_short_devices_and_counts_pool_cache() {
16038        const MIB: usize = 1 << 20;
16039        let need = 900 * MIB;
16040        // Plenty of room: no refusal.
16041        assert_eq!(
16042            capture_headroom_verdict(8_000 * MIB, 0, need, CAPTURE_HEADROOM_FLOOR),
16043            None
16044        );
16045        // The owner's shape: capture appetite would walk the card to the edge — refused,
16046        // with the arithmetic surfaced for the WARN line.
16047        let (required, effective) =
16048            capture_headroom_verdict(1_200 * MIB, 0, need, CAPTURE_HEADROOM_FLOOR)
16049                .expect("short device must refuse");
16050        assert_eq!(required, need + CAPTURE_HEADROOM_FLOOR);
16051        assert_eq!(effective, 1_200 * MIB);
16052        // Pool-cached bytes are real headroom (the trim path makes them driver-visible).
16053        assert_eq!(
16054            capture_headroom_verdict(1_200 * MIB, 7_000 * MIB, need, CAPTURE_HEADROOM_FLOOR),
16055            None
16056        );
16057        // Boundary: exactly enough is enough (>=, never a fencepost refusal).
16058        assert_eq!(
16059            capture_headroom_verdict(
16060                need + CAPTURE_HEADROOM_FLOOR,
16061                0,
16062                need,
16063                CAPTURE_HEADROOM_FLOOR
16064            ),
16065            None
16066        );
16067        // POLICY at the call site (owner-shape receipts, escalated twice on-box): the
16068        // refusal fn is handed 2x the appetite plus TWO floors — a capture may take at
16069        // most half the discretionary headroom, so the card retains a whole capture's
16070        // worth of room after it lands. One floor of slack above one appetite (the shape
16071        // that step-OOM'd on the owner cell) must therefore REFUSE under the call-site
16072        // requirement.
16073        assert!(
16074            capture_headroom_verdict(
16075                need + CAPTURE_HEADROOM_FLOOR + (100 << 20),
16076                0,
16077                2 * need,
16078                CAPTURE_HEADROOM_FLOOR * 2
16079            )
16080            .is_some()
16081        );
16082    }
16083
16084    #[test]
16085    fn bootstrap_estimate_scales_with_heads_and_never_underflows() {
16086        // 3-head chain on a step37-shaped vocab must expect strictly more than one head.
16087        let one = draft_capture_bootstrap_estimate(1, 3, 128_896, 4_096);
16088        let three = draft_capture_bootstrap_estimate(3, 3, 128_896, 4_096);
16089        assert!(three > one);
16090        // Degenerate shapes keep a sane minimum (the estimate feeds a refusal gate; a
16091        // zero-need gate refuses nothing).
16092        assert!(draft_capture_bootstrap_estimate(0, 0, 0, 0) >= 64 << 20);
16093    }
16094
16095    #[test]
16096    fn capture_oom_predicate_matches_the_quoted_driver_text() {
16097        assert!(capture_err_is_oom(
16098            "DriverError(CUDA_ERROR_OUT_OF_MEMORY, \"out of memory\")"
16099        ));
16100        assert!(capture_err_is_oom("allocation failed: out of memory"));
16101        assert!(!capture_err_is_oom("capture produced no graph"));
16102    }
16103}
16104
16105#[cfg(test)]
16106mod mtp_chain_tests {
16107    use super::mtp_chain_head_index;
16108
16109    #[test]
16110    fn embedded_step_heads_cycle_in_declared_order() {
16111        let actual: Vec<usize> = (0..8).map(|step| mtp_chain_head_index(step, 3)).collect();
16112        assert_eq!(actual, [0, 1, 2, 0, 1, 2, 0, 1]);
16113    }
16114
16115    #[test]
16116    fn standalone_draft_remains_single_head() {
16117        assert!((0..8).all(|step| mtp_chain_head_index(step, 1) == 0));
16118    }
16119}
16120
16121#[cfg(test)]
16122mod tp_verified_prefix_tests {
16123    use super::rewind_tp_kv_verified_prefix;
16124    use crate::tp::ResidentTpKvCache;
16125
16126    fn cache_with_committed_len(committed: usize) -> ResidentTpKvCache {
16127        let mut cache = ResidentTpKvCache::new(Vec::new(), 1, 1, 1, 1, 8);
16128        let transaction = cache.begin_transaction().unwrap();
16129        let target = cache.append_target(transaction, committed).unwrap();
16130        cache.publish_append(transaction, target).unwrap();
16131        let target = cache.commit_target(transaction, committed).unwrap();
16132        cache.publish_finalize(transaction, target).unwrap();
16133        cache
16134    }
16135
16136    #[test]
16137    fn replay_free_prefix_rewinds_tp_visibility_to_snapshot_plus_accepts() {
16138        let mut layers = vec![Some(cache_with_committed_len(5)), None];
16139        rewind_tp_kv_verified_prefix(&mut layers, &[Some(2), None], 1).unwrap();
16140        let cache = layers[0].as_ref().unwrap();
16141        assert_eq!(cache.committed_len(), 3);
16142        assert_eq!(cache.staged_len(), 3);
16143    }
16144
16145    #[test]
16146    fn replay_free_prefix_rejects_a_changed_tp_cache_shape() {
16147        let mut layers = vec![Some(cache_with_committed_len(1))];
16148        let error = rewind_tp_kv_verified_prefix(&mut layers, &[None], 1)
16149            .unwrap_err()
16150            .to_string();
16151        assert!(error.contains("changed shape"), "unexpected error: {error}");
16152    }
16153}
16154
16155#[cfg(test)]
16156mod dspark_sparse_tests {
16157    use super::dspark_sparse_softmax_topk;
16158
16159    #[test]
16160    fn topk_keeps_full_softmax_mass_and_stable_ties() {
16161        let logits = [1.0f32, 3.0, 3.0, -2.0];
16162        let (ids, top_logits, probs, tail) = dspark_sparse_softmax_topk(&logits, 2, 1.0).unwrap();
16163        assert_eq!(ids, vec![1, 2]);
16164        assert_eq!(top_logits, vec![3.0, 3.0]);
16165        let denominator = logits.iter().map(|value| (value - 3.0).exp()).sum::<f32>();
16166        let expected = 1.0 / denominator;
16167        assert!((probs[0] - expected).abs() < 1.0e-6);
16168        assert!((probs[1] - expected).abs() < 1.0e-6);
16169        assert!((tail - (1.0 - 2.0 * expected)).abs() < 1.0e-6);
16170        assert!((probs.iter().sum::<f32>() + tail - 1.0).abs() < 1.0e-6);
16171    }
16172}
16173
16174#[cfg(test)]
16175mod spec_replay_env_tests {
16176    use super::spec_replay_env_on;
16177
16178    #[test]
16179    fn replay_requires_literal_one() {
16180        assert!(!spec_replay_env_on(None));
16181        assert!(!spec_replay_env_on(Some("")));
16182        assert!(!spec_replay_env_on(Some("0")));
16183        assert!(!spec_replay_env_on(Some("true")));
16184        assert!(!spec_replay_env_on(Some("2")));
16185        assert!(spec_replay_env_on(Some("1")));
16186    }
16187}
16188
16189#[cfg(test)]
16190mod telem_tests {
16191    use super::{SPEC_TELEM_POS, SpecTelemetry, SpecTelemetryCounters};
16192
16193    #[test]
16194    fn synthetic_accept_masks_produce_tau_and_position_histogram() {
16195        let counters = SpecTelemetryCounters::default();
16196        for mask in [
16197            [true, true, true],
16198            [true, true, false],
16199            [true, false, false],
16200            [false, false, false],
16201        ] {
16202            let accepted = mask.iter().take_while(|&&value| value).count();
16203            counters.record_round(mask.len(), accepted);
16204        }
16205
16206        let snapshot = counters.snapshot();
16207        assert_eq!(
16208            (snapshot.rounds, snapshot.drafted, snapshot.accepted),
16209            (4, 12, 6)
16210        );
16211        assert_eq!(&snapshot.pos_drafted[..3], &[4, 4, 4]);
16212        assert_eq!(&snapshot.pos_accepted[..3], &[3, 2, 1]);
16213        assert_eq!(snapshot.tau(), 1.5);
16214        assert_eq!(snapshot.pos_drafted[3..], [0; SPEC_TELEM_POS - 3]);
16215        assert_eq!(snapshot.pos_accepted[3..], [0; SPEC_TELEM_POS - 3]);
16216    }
16217
16218    /// The worker's per-burst pattern: stash, accumulate, diff — the delta must isolate
16219    /// exactly the burst's contribution (pool-resumed sessions carry prior requests' counts).
16220    #[test]
16221    fn delta_isolates_burst_contribution() {
16222        let mut t = SpecTelemetry::default();
16223        // "previous request": 2 rounds of k=3, accepts 3 then 1.
16224        for (kr, na) in [(3usize, 3usize), (3, 1)] {
16225            t.rounds += 1;
16226            t.drafted += kr as u64;
16227            t.accepted += na as u64;
16228            for j in 0..kr {
16229                t.pos_drafted[j] += 1;
16230            }
16231            for j in 0..na {
16232                t.pos_accepted[j] += 1;
16233            }
16234        }
16235        let before = t;
16236        // "this burst": 1 round k=3, accepts 2.
16237        t.rounds += 1;
16238        t.drafted += 3;
16239        t.accepted += 2;
16240        for j in 0..3 {
16241            t.pos_drafted[j] += 1;
16242        }
16243        for j in 0..2 {
16244            t.pos_accepted[j] += 1;
16245        }
16246        let d = t.delta_since(&before);
16247        assert_eq!((d.rounds, d.drafted, d.accepted), (1, 3, 2));
16248        assert_eq!(&d.pos_drafted[..3], &[1, 1, 1]);
16249        assert_eq!(&d.pos_accepted[..3], &[1, 1, 0]);
16250        assert_eq!(d.pos_drafted[3..], [0; SPEC_TELEM_POS - 3]);
16251    }
16252
16253    /// merge(delta) then merge(delta2) equals accumulating both — the per-model /metrics
16254    /// aggregation invariant.
16255    #[test]
16256    fn merge_accumulates_fieldwise() {
16257        let mut agg = SpecTelemetry::default();
16258        let mut d1 = SpecTelemetry {
16259            rounds: 2,
16260            drafted: 6,
16261            accepted: 4,
16262            ..Default::default()
16263        };
16264        d1.pos_drafted[0] = 2;
16265        d1.pos_accepted[0] = 2;
16266        let mut d2 = SpecTelemetry {
16267            rounds: 1,
16268            drafted: 3,
16269            accepted: 1,
16270            ..Default::default()
16271        };
16272        d2.pos_drafted[0] = 1;
16273        d2.pos_accepted[0] = 1;
16274        d2.pos_drafted[1] = 1;
16275        agg.merge(&d1);
16276        agg.merge(&d2);
16277        assert_eq!((agg.rounds, agg.drafted, agg.accepted), (3, 9, 5));
16278        assert_eq!(agg.pos_drafted[0], 3);
16279        assert_eq!(agg.pos_accepted[0], 3);
16280        assert_eq!(agg.pos_drafted[1], 1);
16281        assert_eq!(agg.pos_accepted[1], 0);
16282    }
16283
16284    /// Wrong-snapshot diff saturates to zero instead of wrapping — the counters feed a
16285    /// public metrics surface and must never publish a u64-wrapped garbage value.
16286    #[test]
16287    fn delta_saturates_never_wraps() {
16288        let small = SpecTelemetry {
16289            rounds: 1,
16290            drafted: 2,
16291            accepted: 1,
16292            ..Default::default()
16293        };
16294        let big = SpecTelemetry {
16295            rounds: 5,
16296            drafted: 15,
16297            accepted: 9,
16298            ..Default::default()
16299        };
16300        let d = small.delta_since(&big);
16301        assert_eq!((d.rounds, d.drafted, d.accepted), (0, 0, 0));
16302    }
16303}
16304
16305#[cfg(test)]
16306mod opti_fork_tests {
16307    use super::{
16308        OptiControllerPolicy, OptiForkAction, OptiForkGateMode, OptiForkGenerationTracker,
16309    };
16310
16311    #[test]
16312    fn controller_threshold_and_three_miss_breaker_are_exact() {
16313        let mut policy = OptiControllerPolicy {
16314            threshold: 0.7,
16315            consecutive_misses: 0,
16316            breaker_tripped: false,
16317        };
16318        assert!(!policy.admit(0.699_999));
16319        assert!(policy.admit(0.7));
16320        assert!(!policy.resolve(false));
16321        assert!(!policy.resolve(false));
16322        assert!(policy.resolve(false));
16323        assert!(policy.breaker_tripped);
16324        assert!(!policy.admit(1.0));
16325        assert!(
16326            !policy.resolve(true),
16327            "a resolved hit cannot re-arm a tripped request"
16328        );
16329        assert!(policy.breaker_tripped);
16330    }
16331
16332    #[test]
16333    fn zero_threshold_is_the_true_unconditional_measurement_arm() {
16334        let mut policy = OptiControllerPolicy {
16335            threshold: 0.0,
16336            consecutive_misses: 0,
16337            breaker_tripped: false,
16338        };
16339        for _ in 0..16 {
16340            assert!(policy.admit(0.0));
16341            assert!(!policy.resolve(false));
16342        }
16343        for invalid in [f32::NAN, f32::INFINITY, -0.01, 1.01] {
16344            assert!(
16345                !policy.admit(invalid),
16346                "invalid q proxy must fail closed: {invalid}"
16347            );
16348        }
16349        assert!(!policy.breaker_tripped);
16350        assert_eq!(policy.consecutive_misses, 0);
16351    }
16352
16353    #[test]
16354    fn alternating_mode_flips_by_generation_not_round_parity() {
16355        assert_eq!(OptiForkGateMode::Alternate.action(0), OptiForkAction::Hit);
16356        assert_eq!(OptiForkGateMode::Alternate.action(1), OptiForkAction::Miss);
16357        assert_eq!(OptiForkGateMode::Alternate.action(8), OptiForkAction::Hit);
16358        assert_eq!(OptiForkGateMode::Alternate.action(9), OptiForkAction::Miss);
16359    }
16360
16361    #[test]
16362    fn live_generation_cannot_be_overwritten() {
16363        let mut tracker = OptiForkGenerationTracker::default();
16364        let g0 = tracker.reserve().unwrap();
16365        let g1 = tracker.reserve().unwrap();
16366        let err = tracker.reserve().unwrap_err().to_string();
16367        assert!(
16368            err.contains("still owns generation 0"),
16369            "unexpected error: {err}"
16370        );
16371        tracker.retire(g0).unwrap();
16372        let g2 = tracker.reserve().unwrap();
16373        assert_eq!((g2.id, g2.slot), (2, 0));
16374        tracker.retire(g1).unwrap();
16375        tracker.retire(g2).unwrap();
16376    }
16377
16378    #[test]
16379    fn teardown_rejects_a_stale_generation_tag() {
16380        let mut tracker = OptiForkGenerationTracker::default();
16381        let g0 = tracker.reserve().unwrap();
16382        tracker.retire(g0).unwrap();
16383        let err = tracker.retire(g0).unwrap_err().to_string();
16384        assert!(err.contains("teardown mismatch"), "unexpected error: {err}");
16385    }
16386}
16387
16388#[cfg(test)]
16389mod draft_graph_fallback_tests {
16390    use super::DraftGraphFallback;
16391
16392    /// The Q2 contract, part (a): a fallback flip is LOUD — exactly once per flip.
16393    #[test]
16394    fn flip_is_loud_once_and_memoized_after() {
16395        let mut f = DraftGraphFallback::default();
16396        let line = f
16397            .mark_greedy("out of memory")
16398            .expect("first flip must return the warn line");
16399        assert!(
16400            line.contains("WARN"),
16401            "flip line must be warn-level: {line}"
16402        );
16403        assert!(
16404            line.contains("out of memory"),
16405            "flip line must carry the reason: {line}"
16406        );
16407        assert!(f.greedy_failed());
16408        // re-marking an already-failed graph is the memoization: quiet, still failed.
16409        assert!(f.mark_greedy("out of memory").is_none());
16410        assert!(f.greedy_failed());
16411        // the two graphs' flags are independent (greedy flip leaves sampled capturable).
16412        assert!(!f.sampled_failed());
16413        let line_s = f
16414            .mark_sampled("capture unsupported")
16415            .expect("sampled flip is its own flip");
16416        assert!(
16417            line_s.contains("sampled"),
16418            "sampled flip names itself: {line_s}"
16419        );
16420        assert!(f.mark_sampled("capture unsupported").is_none());
16421    }
16422
16423    /// The Q2 contract, part (b): resume-from-pool RESETS both flags (fresh capture chance),
16424    /// and says so exactly when there was something to reset.
16425    #[test]
16426    fn reset_on_resume_clears_flags_and_logs_once() {
16427        let mut f = DraftGraphFallback::default();
16428        // clean session: resume is silent, nothing to reset.
16429        assert!(f.reset_on_resume().is_none());
16430        f.mark_greedy("oom").unwrap();
16431        f.mark_sampled("oom").unwrap();
16432        let note = f
16433            .reset_on_resume()
16434            .expect("a set flag must produce the reset note");
16435        assert!(
16436            note.contains("greedy+sampled"),
16437            "note names what was reset: {note}"
16438        );
16439        assert!(
16440            !f.greedy_failed() && !f.sampled_failed(),
16441            "both flags cleared"
16442        );
16443        // and the NEXT failure after a reset is a fresh flip — loud again.
16444        assert!(f.mark_greedy("oom again").is_some());
16445        let note2 = f.reset_on_resume().expect("greedy-only reset");
16446        assert!(note2.contains("(greedy)"), "single-flag note: {note2}");
16447    }
16448
16449    /// Shape-change clears (dmask realloc / mask-shape mismatch / s_key change) stay silent —
16450    /// they precede a fresh capture attempt whose own failure re-flips loudly.
16451    #[test]
16452    fn shape_change_clears_are_silent() {
16453        let mut f = DraftGraphFallback::default();
16454        f.mark_greedy("oom").unwrap();
16455        f.clear_greedy();
16456        assert!(!f.greedy_failed());
16457        f.mark_sampled("oom").unwrap();
16458        f.clear_sampled();
16459        assert!(!f.sampled_failed());
16460        // after a silent clear there is nothing left for resume to report.
16461        assert!(f.reset_on_resume().is_none());
16462    }
16463}
16464
16465/// SAMPLED DRAFT-GRAPH KEY (lane/graph-s-key-exactness-20260819).
16466///
16467/// These are the CPU teeth for an exactness bug whose live reproduction needs a GPU, a trunk, a
16468/// drafter and a two-turn session: the key itself. Every test below fails against the pre-fix key
16469/// `(seed, temp.to_bits(), k)` — `legacy_key` restates it so the collision is explicit rather
16470/// than remembered.
16471#[cfg(test)]
16472mod sampled_graph_key_tests {
16473    use super::{SampledGraphKey, debug_t_pred0};
16474
16475    /// The pre-fix key, verbatim: `let s_key = (sp_seed, sp_temp.to_bits(), k);`
16476    fn legacy_key(k: &SampledGraphKey) -> (u64, u32, usize) {
16477        (k.seed, k.temp_bits, k.k)
16478    }
16479
16480    fn pure_temp_key() -> SampledGraphKey {
16481        // temperature 1.0, filters off — today's serve default, the shape that parks a graph.
16482        SampledGraphKey::new(12345, 1.0, 3, 0, 1.0, 0.0, false)
16483    }
16484
16485    /// THE COLLISION. Two requests that differ ONLY in the truncation filters shared one key, so
16486    /// a parked pure-temp graph survived into a filtered request and the launch site launched it.
16487    #[test]
16488    fn vendor_filters_change_the_key() {
16489        let parked = pure_temp_key();
16490        // qwen3.8 generation_config.json — what the vendor-default flip makes the default shape.
16491        let vendor = SampledGraphKey::new(12345, 1.0, 3, 20, 0.95, 0.0, false);
16492        assert_eq!(
16493            legacy_key(&parked),
16494            legacy_key(&vendor),
16495            "pre-fix key collided: this is the bug, and the reason a test asserts on it",
16496        );
16497        assert_ne!(parked, vendor, "post-fix key must separate the two regimes");
16498        assert!(parked.pure_temp());
16499        assert!(!vendor.pure_temp());
16500    }
16501
16502    /// Each distribution-shaping field alone is enough to drop the parked graph.
16503    #[test]
16504    fn every_filter_field_is_keyed() {
16505        let base = pure_temp_key();
16506        for (what, other) in [
16507            (
16508                "top_k",
16509                SampledGraphKey::new(12345, 1.0, 3, 20, 1.0, 0.0, false),
16510            ),
16511            (
16512                "top_p",
16513                SampledGraphKey::new(12345, 1.0, 3, 0, 0.95, 0.0, false),
16514            ),
16515            (
16516                "min_p",
16517                SampledGraphKey::new(12345, 1.0, 3, 0, 1.0, 0.05, false),
16518            ),
16519            (
16520                "penalties",
16521                SampledGraphKey::new(12345, 1.0, 3, 0, 1.0, 0.0, true),
16522            ),
16523        ] {
16524            assert_ne!(base, other, "{what} must be part of the key");
16525            assert!(!other.pure_temp(), "{what} leaves the pure-temp regime");
16526            assert_eq!(
16527                legacy_key(&base),
16528                legacy_key(&other),
16529                "{what} was invisible to the pre-fix key",
16530            );
16531        }
16532    }
16533
16534    /// The baked constants stay keyed (this half was always right — regression cover for it).
16535    #[test]
16536    fn baked_constants_stay_keyed() {
16537        let base = pure_temp_key();
16538        assert_ne!(
16539            base,
16540            SampledGraphKey::new(999, 1.0, 3, 0, 1.0, 0.0, false),
16541            "seed"
16542        );
16543        assert_ne!(
16544            base,
16545            SampledGraphKey::new(12345, 0.7, 3, 0, 1.0, 0.0, false),
16546            "temp"
16547        );
16548        assert_ne!(
16549            base,
16550            SampledGraphKey::new(12345, 1.0, 4, 0, 1.0, 0.0, false),
16551            "k"
16552        );
16553        // bitwise on temperature: 0.7f32 vs the same value re-derived must NOT differ.
16554        assert_eq!(
16555            SampledGraphKey::new(1, 0.7, 3, 0, 1.0, 0.0, false),
16556            SampledGraphKey::new(1, 7.0 / 10.0, 3, 0, 1.0, 0.0, false),
16557        );
16558    }
16559
16560    /// THE LOAD-BEARING HALF OF THE SEED DECISION (lane/session-resume-sampler-predicate-
16561    /// 20260820). The whole-session resume predicate deliberately does NOT compare `seed`: an
16562    /// omitted serve `seed` draws fresh per-request entropy, so comparing it would refuse every
16563    /// seed-omitting sampled conversation. That is only sound because the one piece of parked state
16564    /// that BAKES the seed — this graph — is re-keyed on it, so a seed change drops and recaptures.
16565    ///
16566    /// This test is the other end of that argument, asserted here rather than remembered in a
16567    /// comment: if a future change dropped `seed` from the key, the resume predicate's exclusion
16568    /// would silently become the unsound thing it is documented not to be.
16569    /// (Paired with `seed_alone_does_not_refuse` in `memra-sampling`.)
16570    #[test]
16571    fn seed_alone_still_rekeys_the_draft_graph() {
16572        let parked = pure_temp_key();
16573        let reseeded = SampledGraphKey::new(999, 1.0, 3, 0, 1.0, 0.0, false);
16574        assert_ne!(
16575            parked, reseeded,
16576            "a seed-only change MUST drop the parked sampled graph — the resume predicate's \
16577             decision not to compare seed rests on exactly this",
16578        );
16579        // Same regime on both sides: the drop is a recapture, not a fall to the eager chain
16580        // because of a filter difference.
16581        assert!(parked.pure_temp() && reseeded.pure_temp());
16582    }
16583
16584    /// `pure_temp()` is the capture guard's predicate, computed from the key so the two cannot
16585    /// drift. The equality below is the invariant the launch-site guard asserts: identical keys
16586    /// agree on the regime, so a graph that survives the drop is legal to launch.
16587    #[test]
16588    fn equal_keys_agree_on_the_regime() {
16589        let a = SampledGraphKey::new(7, 0.8, 3, 20, 0.95, 0.0, false);
16590        let b = SampledGraphKey::new(7, 0.8, 3, 20, 0.95, 0.0, false);
16591        assert_eq!(a, b);
16592        assert_eq!(a.pure_temp(), b.pure_temp());
16593        // top_p slightly above 1.0 (a client sending 1.0 exactly, or an operator default) is
16594        // still the unfiltered regime, matching the original `sp.top_p >= 1.0` test.
16595        assert!(SampledGraphKey::new(7, 0.8, 3, 0, 1.0, 0.0, false).pure_temp());
16596        assert!(SampledGraphKey::new(7, 0.8, 3, 0, 1.5, -1.0, false).pure_temp());
16597    }
16598
16599    /// The WIDENED capture regime (lane/step37-draft-graph-serving-20260830): truncation-
16600    /// filtered shapes are capturable — the filter runs IN-GRAPH (`filter_stats` +
16601    /// `gumbel_perturb_filtered_ctr`), so the draft draws from the same filtered
16602    /// distribution the accept test reconstructs. Penalties never are: the per-round
16603    /// history cannot be baked. The step37 vendor-default shape (temp 0.5 / top_p 0.9) is
16604    /// exactly the previously-excluded regime this lane exists to capture.
16605    #[test]
16606    fn filtered_regimes_are_capturable_penalties_never() {
16607        let vendor = SampledGraphKey::new(12345, 0.5, 3, 0, 0.9, 0.0, false);
16608        assert!(!vendor.pure_temp());
16609        assert!(vendor.filtered());
16610        assert!(
16611            vendor.graph_capturable(),
16612            "the vendor-default filtered shape must be capturable (default door state)",
16613        );
16614        assert!(pure_temp_key().graph_capturable());
16615        assert!(
16616            !pure_temp_key().filtered(),
16617            "pure-temp takes the legacy (filterless) capture body",
16618        );
16619        let pen = SampledGraphKey::new(12345, 0.5, 3, 0, 0.9, 0.0, true);
16620        assert!(
16621            !pen.graph_capturable(),
16622            "penalty history varies per round and can never be baked into a graph",
16623        );
16624    }
16625
16626    /// MEMRA_DEBUG_SPEC on a SAMPLED spec request past round 0: the print must render without
16627    /// indexing the empty greedy `preds` vector (it panicked the GPU worker before this lane).
16628    #[test]
16629    fn debug_print_survives_the_sampled_arm() {
16630        // round >= 1 with a pending bonus == base 1, sampled == `preds` empty.
16631        assert_eq!(debug_t_pred0(true, 1, 4242, &[]), "n/a");
16632        assert_eq!(debug_t_pred0(true, 2, 4242, &[]), "n/a");
16633        // round 0 without a pending bonus still reports last_pred, in both arms.
16634        assert_eq!(debug_t_pred0(true, 0, 4242, &[]), "4242");
16635        assert_eq!(debug_t_pred0(false, 0, 4242, &[7, 8]), "4242");
16636        // greedy keeps the real prediction it always printed.
16637        assert_eq!(debug_t_pred0(false, 1, 4242, &[7, 8]), "7");
16638        assert_eq!(debug_t_pred0(false, 2, 4242, &[7, 8]), "8");
16639    }
16640}