Skip to main content

memra_engine/
spec.rs

1//! Qwen3.5 MTP (NextN) greedy speculative decode (research/mtp/MTP-PLAN.md §A/§B/§C/§D).
2//!
3//! Greedy spec decode is MATHEMATICALLY EXACT: the accepted+bonus token stream is token-for-token
4//! identical to plain greedy `generate`. This module provides:
5//!   - `mtp_head_forward`  (§A, T=1): one NextN draft-token forward.
6//!   - `decode_step_t`     (§D.3, T=K+1): batched target verify forward, all-column logits.
7//!   - `generate_spec`     (§B): the draft/verify/accept/rollback orchestrator.
8//!     Cache snapshot/rollback lives in cache.rs (§D.4). The MTP head uses its OWN scratch KV (§D.6),
9//!     PERSISTENT over the committed sequence (see `MtpScratch`).
10
11use crate::Engine;
12use crate::cache::{Cache, KvLayer};
13use crate::forward::argmax;
14use crate::hybrid::{FullAttnLayer, HybridModel, LinearAttnLayer, Mixer, MtpHead};
15use cudarc::driver::CudaSlice;
16use memra_gguf::config::SwigluClamp;
17use std::sync::atomic::{AtomicU64, Ordering};
18
19/// Parse the documented `MEMRA_SPEC_REPLAY=1` rollback seam.
20///
21/// Keep this shared with serving admission so `=0` cannot select replay in one
22/// layer while another layer treats it as disabled.
23pub fn spec_replay_env_on(value: Option<&str>) -> bool {
24    value == Some("1")
25}
26
27pub fn spec_replay_env_enabled() -> bool {
28    let value = std::env::var("MEMRA_SPEC_REPLAY").ok();
29    spec_replay_env_on(value.as_deref())
30}
31
32/// step35 dcw draft-chain door (lane/step37-draft-graph-20260829). ON routes the step35 MTP
33/// block's draft attention through the WINDOWED device-counter family
34/// (`append_kv_quantized_dcw` + `fa_decode_dcw`, the step TP graph arc's kernels), which
35/// derives the SWA view entirely from device state (len_d, base_d, window): exactly the view
36/// offset the old capture refusal said `fa_decode_dc` could not express. BOTH draft modes
37/// switch together: eager and captured run the ONE launcher at the ONE bucket
38/// (min(cap, window)), so graph-vs-eager draft parity holds by construction (the
39/// `mtp_full_attn_dc` precedent).
40///
41/// DEFAULT ON since lane/step37-draft-graph-serving-20260830: the 20260829 lane shipped it
42/// OFF because it enabled nothing at the shipping head count (capture was structurally
43/// unreachable at heads=3); with the multi-head chain capture and the in-graph filtered
44/// sampler landed, this door is the kernel prerequisite for the captured chain on the
45/// QUALIFIED serving shape, and the exactness battery (greedy K=1..8 identity, per-K
46/// acceptance identity, seeded sampled twins) banks on the ON arm. Rollback seam:
47/// MEMRA_STEP35_DRAFT_DCW=0 restores the host-len eager arm (`mtp_step35_attn`) plus the
48/// named capture refusal, byte-for-byte the pre-lane serving; no state survives restart.
49fn step35_draft_dcw_on() -> bool {
50    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
51    *ON.get_or_init(|| std::env::var("MEMRA_STEP35_DRAFT_DCW").as_deref() != Ok("0"))
52}
53
54/// Multi-head MTP draft-chain capture door (lane/step37-draft-graph-serving-20260830,
55/// default ON — receipts in the lane RESULTS). ON lets the step-modulo prefix-replay chain
56/// (`mtp_extra` non-empty, the step37 3-head shipping shape) capture per-head single-row
57/// CUDA graphs and replay them in the exact eager launch order; the chain POLICY (head
58/// selection, prefix length, seed history) stays host-side, so graph-vs-eager drafts are
59/// bit-identical by construction. A failed capture degrades LOUDLY to the eager chain (the
60/// draft-graph WARN contract). OFF (=0) keeps the eager chain as the only multi-head path —
61/// the pre-lane serving byte-for-byte. Single-head capture is untouched by this door.
62fn mtp_chain_graph_on() -> bool {
63    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
64    *ON.get_or_init(|| std::env::var("MEMRA_MTP_CHAIN_GRAPH").as_deref() != Ok("0"))
65}
66
67/// In-graph FILTERED sampled draft door (lane/step37-draft-graph-serving-20260830, default
68/// ON — receipts in the lane RESULTS). ON widens the sampled draft-graph capture from the
69/// pure-temp regime to every truncation-filtered regime (top_k / top_p / min_p): the capture
70/// body runs `filter_stats` + `gumbel_perturb_filtered_ctr` IN-GRAPH, so the draft draws
71/// from the SAME filtered distribution the verify's accept test reconstructs (the
72/// graph-s-key exactness law, now satisfied inside the graph instead of by refusing it).
73/// Penalties stay eager either way (the history varies per round and cannot be baked).
74/// The pure-temp capture body is UNTOUCHED by this door (byte-identical to the pre-lane
75/// graph). OFF (=0) restores the pure-temp-only capture guard: filtered requests draft
76/// eager, byte-for-byte the pre-lane behavior.
77fn spec_graph_filtered_on() -> bool {
78    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
79    *ON.get_or_init(|| std::env::var("MEMRA_SPEC_GRAPH_FILTERED").as_deref() != Ok("0"))
80}
81
82fn parse_prime_trows_width(value: Option<&str>) -> Result<usize, String> {
83    let Some(raw) = value else {
84        return Ok(8);
85    };
86    let width = raw
87        .parse::<usize>()
88        .map_err(|_| format!("MEMRA_PRIME_TROWS_T must be an integer in 2..=8, got {raw:?}"))?;
89    if !(2..=8).contains(&width) {
90        return Err(format!("MEMRA_PRIME_TROWS_T must be in 2..=8, got {width}"));
91    }
92    Ok(width)
93}
94
95#[cfg(test)]
96mod prime_trows_width_tests {
97    #[test]
98    fn width_defaults_to_eight_and_refuses_invalid_operator_values() {
99        assert_eq!(super::parse_prime_trows_width(None), Ok(8));
100        assert_eq!(super::parse_prime_trows_width(Some("2")), Ok(2));
101        assert_eq!(super::parse_prime_trows_width(Some("8")), Ok(8));
102        for invalid in ["", "1", "9", "32", "wide"] {
103            let err = super::parse_prime_trows_width(Some(invalid)).unwrap_err();
104            assert!(err.contains("MEMRA_PRIME_TROWS_T"), "{err}");
105            assert!(err.contains("2..=8"), "{err}");
106        }
107    }
108}
109
110/// One compact, anchor-bounded DSpark supervision record. `tokens[0]` is the anchor at p and
111/// `hidden` is its predecessor carrier h[p-1], matching the live NextN/DSpark pairing. Target
112/// rows p..p+gamma-1 score tokens p+1..p+gamma. They are the full-target softmax's top-k
113/// entries; `target_tail_probs[j]` is the probability mass outside those rows. All flattened
114/// target arrays are `[gamma, top_k]` in row-major order.
115pub struct DsparkAnchorRecord {
116    pub position: usize,
117    pub hidden: Vec<f32>,
118    pub tokens: Vec<u32>,
119    pub target_top_ids: Vec<u32>,
120    pub target_top_logits: Vec<f32>,
121    pub target_top_probs: Vec<f32>,
122    pub target_tail_probs: Vec<f32>,
123}
124
125#[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
126fn dspark_sparse_softmax_topk(
127    logits: &[f32],
128    top_k: usize,
129    temperature: f32,
130) -> Result<(Vec<u32>, Vec<f32>, Vec<f32>, f32), Box<dyn std::error::Error>> {
131    if logits.is_empty() || top_k == 0 || top_k > logits.len() || temperature <= 0.0 {
132        return Err("invalid DSpark sparse-softmax shape or temperature".into());
133    }
134    if logits.iter().any(|value| !value.is_finite()) {
135        return Err("DSpark target logits contain a non-finite value".into());
136    }
137    let mut ranked: Vec<(u32, f32)> = logits
138        .iter()
139        .copied()
140        .enumerate()
141        .map(|(index, value)| (index as u32, value))
142        .collect();
143    let compare = |left: &(u32, f32), right: &(u32, f32)| {
144        right.1.total_cmp(&left.1).then(left.0.cmp(&right.0))
145    };
146    ranked.select_nth_unstable_by(top_k - 1, compare);
147    ranked[..top_k].sort_unstable_by(compare);
148
149    let max_logit = logits.iter().copied().fold(f32::NEG_INFINITY, f32::max);
150    let inv_temperature = 1.0f64 / temperature as f64;
151    let denominator: f64 = logits
152        .iter()
153        .map(|value| (((*value - max_logit) as f64) * inv_temperature).exp())
154        .sum();
155    let ids: Vec<u32> = ranked[..top_k].iter().map(|(index, _)| *index).collect();
156    let top_logits: Vec<f32> = ranked[..top_k].iter().map(|(_, value)| *value).collect();
157    let top_probs: Vec<f32> = top_logits
158        .iter()
159        .map(|value| ((((value - max_logit) as f64) * inv_temperature).exp() / denominator) as f32)
160        .collect();
161    let top_mass: f64 = top_probs.iter().map(|value| *value as f64).sum();
162    let tail = (1.0f64 - top_mass).clamp(0.0, 1.0) as f32;
163    Ok((ids, top_logits, top_probs, tail))
164}
165
166fn flatten_dspark_rows<T>(
167    rows: Vec<Option<Vec<T>>>,
168    position: usize,
169    label: &str,
170) -> Result<Vec<T>, Box<dyn std::error::Error>> {
171    let mut flattened = Vec::new();
172    for (slot, row) in rows.into_iter().enumerate() {
173        flattened.extend(
174            row.ok_or_else(|| format!("missing DSpark {label} at {position} slot {slot}"))?,
175        );
176    }
177    Ok(flattened)
178}
179
180/// H-SEED CONVENTION (MEMRA_SPEC_HPOST=1): feed the MTP head the POST-norm hidden — trunk rows
181/// hand over `output_norm(x)` and the draft chain recurrence hands over `shared_head_norm(h_nextn)`
182/// (= final_h) — matching the reference engines: llama.cpp #24025 ("qwen35: use post-norm hidden
183/// state for MTP", t_h_nextn is taken AFTER the final norm in both trunk and MTP graphs) and
184/// SGLang's qwen3_5_mtp (spec_info.hidden_states = the target model's post-norm output). memra's
185/// historical convention (default, MTP-PLAN §A) is PRE-norm x. Draft-quality-only: exactness is
186/// the verify's job either way; acceptance arbitrates. OnceLock: read once, hot-loop safe.
187/// `MEMRA_SPEC_HEAD_ROWS=1` — batch the verify tail's LM head over its t columns instead of running
188/// it at m=1 once per column. See the call site in `decode_step_t_core_stream` for why the batched
189/// form is the same per-row arithmetic (the bf16/q8 rows twins, not cuBLASLt) and what it costs
190/// today: the head is re-streamed t times per verify pass. Default off until the byte tape says so.
191pub(crate) fn head_rows_on() -> bool {
192    static ENV: std::sync::OnceLock<Option<bool>> = std::sync::OnceLock::new();
193    crate::step37_door(&ENV, "MEMRA_SPEC_HEAD_ROWS")
194}
195
196/// The serving walk's own doors, tri-stated the same way (owner flip 2026-08-27): env forces,
197/// unset takes the step37 family default. Call sites are the t-row verify walk itself.
198pub(crate) fn spec_verify_eager_on() -> bool {
199    static ENV: std::sync::OnceLock<Option<bool>> = std::sync::OnceLock::new();
200    crate::step37_door(&ENV, "MEMRA_SPEC_VERIFY_EAGER")
201}
202
203pub(crate) fn spec_verify_tcol_on() -> bool {
204    static ENV: std::sync::OnceLock<Option<bool>> = std::sync::OnceLock::new();
205    crate::step37_door(&ENV, "MEMRA_SPEC_VERIFY_TCOL")
206}
207
208/// NOT family-armed (2026-08-27): the walk's prime leaves its sub-32 TAIL chunk out of the
209/// DISTRIBUTED kv, so the server refuses before decode with "cache lengths diverged
210/// local=N distributed=floor(N/32)*32" for every prompt whose token count is not a multiple of
211/// 32 — i.e. nearly all real traffic. Isolated on the server route: defaults ERR (local=445
212/// distributed=416), MEMRA_PRIME_TROWS=0 OK. It was default-OFF before the 2026-08-27 flip and
213/// goes back to opt-in until the tail append is fixed and gated ON THE SERVER ROUTE, not just
214/// run-gen (run-gen calls decode_step_t on the whole prompt and never exercises this path — the
215/// reason a run-gen-only receipt could not see it). The GEMM prime supersedes it on this route.
216pub(crate) fn prime_trows_on() -> bool {
217    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
218    *ON.get_or_init(|| std::env::var("MEMRA_PRIME_TROWS").as_deref() == Ok("1"))
219}
220
221pub(crate) fn tcol_ffn_on() -> bool {
222    static ENV: std::sync::OnceLock<Option<bool>> = std::sync::OnceLock::new();
223    crate::step37_door(&ENV, "MEMRA_TCOL_FFN")
224}
225
226pub(crate) fn spec_hpost() -> bool {
227    static H: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
228    *H.get_or_init(|| {
229        std::env::var("MEMRA_SPEC_HPOST")
230            .map(|v| v != "0")
231            .unwrap_or(false)
232    })
233}
234
235/// LEAN VERIFY (default ON since 2026-07-08; MEMRA_SPEC_LEAN=0 reverts — close35 lane): the verify m-scaling
236/// probe + nsys diff showed the verify t-path pays ~1.0ms/call at m=1 over eager decode on the
237/// 35B, and the kernels are NOT the cause (dev-MoE identical, kernel-time delta only +179us).
238/// The overhead is (a) ~250 extra cuMemsetD8Async/call from `e.zeros()` on buffers every kernel
239/// fully overwrites (~0.9ms host issue + ~0.35ms GPU) and (b) the t=1 FA rows dispatch (rows_v2 +
240/// combine_rows, +50us vs the eager fa_decode pair). This flag switches (a) fully-overwritten
241/// verify buffers to `e.uninit` (identical bytes: every element is written before read) and
242/// (b) t==1 verify FA to the eager `fa_decode` entry (byte-identical: kernel-check pins the
243/// rows-vs-loop identity and the per-row loop at t=1 IS fa_decode on the same q). Gates arbitrate.
244pub(crate) fn spec_lean() -> bool {
245    static L: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
246    // DEFAULT ON since 2026-07-08 (MEMRA_SPEC_LEAN=0 reverts): bit-identical (buffers fully
247    // overwritten; gates green incl maxdiff-identical run-gen) and measured +2.4% e2e p3 /
248    // +1.5% p2 at the daily 35B config. m=1 verify now costs eager-decode parity.
249    *L.get_or_init(|| {
250        std::env::var("MEMRA_SPEC_LEAN")
251            .map(|v| v != "0")
252            .unwrap_or(true)
253    })
254}
255
256/// SMALL-M BATCHED VERIFY (default ON since 2026-07-09; MEMRA_SPEC_M2=0 reverts — lane/spec-m2): extend the
257/// batched linear-attn verify arm down to t=2 and batch the MoE dev token loop over a
258/// grid.z=token axis at every verify t. The close35 m-scaling probe put the m=2 verify tier at
259/// x1.54 of m=1 (llama x1.14); the per-column linear chain (t<3) and the serial MoE dev token
260/// loop are the two launch-structure causes. Both changes are LAUNCH-STRUCTURE ONLY:
261/// (a) the batched conv's t<pad ring update is pure copies (ssm_conv_ring_rebuild from a cloned
262///     ring — the ring stores raw input columns); every arithmetic kernel is the same one the
263///     t>=3 arm already runs (matmul_decode_exact bit-identical at m=2-4, gdn_scan's internal
264///     t-loop == chained T=1 steps);
265/// (b) the MoE dev-rows twins run the serial loop's per-token warp program with tok-offset
266///     pointers (same sel/w/aq/ad bytes, same dot order, same slot-ordered FMA chain).
267/// Gates arbitrate: run-spec K=1..8 self-consistency (35B+9B), kernel-check, run-gen argmax.
268pub(crate) fn spec_m2() -> bool {
269    static M: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
270    // DEFAULT ON since 2026-07-09 (MEMRA_SPEC_M2=0 reverts): launch-structure only — t=2
271    // batched linear arm (ring-roll copies, zero new FP order) + MoE dev-rows kernels
272    // (grid.z=token, 4 launches/layer at any verify t). Acceptance bit-identical at every K;
273    // 35B p2 +3.4% / p3 +3.6%; the profitable-K plateau widens (new optimum K=3 at 223).
274    *M.get_or_init(|| {
275        std::env::var("MEMRA_SPEC_M2")
276            .map(|v| v != "0")
277            .unwrap_or(true)
278    })
279}
280pub(crate) fn spec_stream() -> bool {
281    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
282    *ON.get_or_init(|| std::env::var("MEMRA_SPEC_STREAM").as_deref() == Ok("1"))
283}
284pub(crate) fn spec_stream_m() -> usize {
285    static M: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
286    *M.get_or_init(|| {
287        std::env::var("MEMRA_SPEC_STREAM_M")
288            .ok()
289            .and_then(|v| v.parse().ok())
290            .unwrap_or(4)
291    })
292}
293/// Engine-bundle slice 3 + fa-execupdate slice 4c (DSF-ROUNDCOST-20260820 §5 rank 1),
294/// DEFAULT OFF — `MEMRA_DSPARK_VERIFY_GRAPH=1` opts in: per-(segment, vt) CUDA graphs
295/// for the LINEAR-layer runs, plus the full-verify single graph per (vt, rung) when a
296/// round's rows all ride one seqs rung — see [`DsparkVerifyGraphs`]. Requires the
297/// slice-2 deferred path (device tokens); the eager walk is the byte-identical fallback.
298///
299/// MEASURED disposition (box6 card0, agentic pack, 2026-08-20, both slices): exactness
300/// holds everywhere (ALL EXACT, accept lines byte-match the banks, ckpt-gate oracle
301/// green over the graph + slab-commit paths). Slice-3's AUTO_FREE launch-scan limiter
302/// (25.6 us x 16 launches ≈ 0.41 ms/round) is FIXED — the captured bodies' alloc nodes
303/// are balanced by in-graph frees (census 84/84 per segment, 1776/1776 full) so graphs
304/// instantiate USE_NODE_PRIORITY and the scan is gone. What remains at gate scale:
305/// segment graphs +0.1 tok/s over the batched-rows default (114.4 vs 114.3 x5
306/// interleaved — the linear launch overhead was only ~0.1 ms); the FULL-verify graph is
307/// NET NEGATIVE at gate scale (110.6 vs 114.2: ~14-21 (vt, rung) captures/process at
308/// 2 full-walk executions + ~2.9k-node instantiate each eat far more than the ~0.2-0.3
309/// ms/round of remaining launch overhead). The orchestration ceiling of §1.3 is spent —
310/// the fa/append recovery landed DEFAULT-ON as the batched rows arm
311/// (`dspark_fa_rows_on`), not as a graph. The serve-lifetime cell (DSF-ROUNDCOST §9,
312/// nj-ws-solo) measured the amortization: crossover K≈33 requests, steady −0.246
313/// ms/round, −1.25% session wall over 240 requests — and the graphs-serve lane wired
314/// the door into the session arm (`dspark_spec_session_burst`) as a model-owned
315/// capture pool shared across sessions. Stays opt-in pending the owner's default-ON
316/// ratification on the serve-surface battery.
317pub(crate) fn dspark_verify_graph_on() -> bool {
318    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
319    *ON.get_or_init(|| std::env::var("MEMRA_DSPARK_VERIFY_GRAPH").as_deref() == Ok("1"))
320}
321/// MTP-ROUTE verify graphs, DEFAULT ON for the GDN+MoE family since 2026-08-23
322/// (`MEMRA_SPEC_VERIFY_GRAPH=0` is the kill switch, `=1` opts other families in).
323///
324/// The slice-4c capture already lived inside `qwen35_verify_tparallel` and said so in its own
325/// comment — "stream rides the qwen35moe burst, graphs ride the dspark route" — with no caller
326/// on this route. The MTP spec round is that caller.
327///
328/// WHY it is worth a default (receipts: `research/orndecode-20260822/VGRAPH.md`). With
329/// `MEMRA_SPEC_PHASE=1` this route's round reads verify-ISSUE 44-58% and verify-WAIT **0.0%**:
330/// the host is never waiting for the device, it is spending its own time launching the trunk.
331/// Replay collapses that into one graph launch and the phase all but disappears (55-62 ms ->
332/// 8-10 ms per burst).
333///
334/// MEASURED, two host generations, forced ON/OFF, balanced 4+4 boots in both orders:
335///   * current-generation host (9950X, the serving class): OFF 266.0-266.5, ON 318.8-319.5
336///     tok/s — **+19.7%**, no overlap, sub-1% spread per arm; per-round 6.9 -> 5.7 ms.
337///   * Zen 3 host: +3-9% (that rig's own clock drift is wider than the effect, so the ratio
338///     comes from per-round phase totals, which are internal to each boot).
339///     The ON arm lands at ~320 tok/s on BOTH hosts while OFF tracks host speed — the arm moves
340///     the round off the host and onto the device, which is the whole point.
341///
342/// EXACTNESS is structural (same kernels, same order) and gated anyway: a fixed-seed SAMPLED
343/// completion hashes identically ON vs OFF **and across both hosts** (`08941d5bb9762b21`),
344/// greedy seed-pinned likewise, `run-spec` K=1..8 PASS on both arms with identical acceptance
345/// at every K, kernel-check ALL GREEN.
346///
347/// SCOPE, deliberately narrow: default ON only where it was measured — the GatedDeltaNet +
348/// MoE family (`vgraph_family_default`). Qwen3.8-27B is GDN + DENSE mlp and would otherwise
349/// inherit this default unmeasured, which is the family-by-family law this repo keeps; it can
350/// opt in with `=1` once it has its own interleave. Also never armed together with
351/// ROUND-STREAM, and a round wider than the pool declines it for the eager walk.
352pub(crate) fn spec_verify_graph_env() -> Option<bool> {
353    static ON: std::sync::OnceLock<Option<bool>> = std::sync::OnceLock::new();
354    *ON.get_or_init(
355        || match std::env::var("MEMRA_SPEC_VERIFY_GRAPH").as_deref() {
356            Ok("1") => Some(true),
357            Ok("0") => Some(false),
358            _ => None,
359        },
360    )
361}
362/// SERVE-ROUTE twin of [`dspark_verify_graph_on`], DEFAULT ON — owner-ratified
363/// 2026-08-22 on the §10 serve-lifetime battery (DSF-ROUNDCOST-20260820 §10.3:
364/// crossover K=36–43, steady −0.357 ms/round, session wall −1.55..−1.65%, byte-exact
365/// 240/240 ×3 pairs, pool bounded at 8,852 MiB under `MEMRA_DSPARK_VG_MAX`). The env
366/// stays as the kill-switch: `MEMRA_DSPARK_VERIFY_GRAPH=0` restores the eager walk
367/// (byte-identical body); `MEMRA_DSPARK_VG_MAX=0` is the finer freeze valve. The BIN
368/// arm keeps its own opt-in default (`dspark_verify_graph_on`): at gate scale the
369/// capture toll is never repaid (§8 measured disposition — 14–21 captures over a
370/// 256-token run vs the serve session's thousands of rounds), and the two
371/// instruments must keep their own measured dispositions rather than share one flag.
372pub(crate) fn dspark_verify_graph_serve_on() -> bool {
373    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
374    *ON.get_or_init(|| std::env::var("MEMRA_DSPARK_VERIFY_GRAPH").as_deref() != Ok("0"))
375}
376/// Capture-count ceiling for the dspark verify-graph pool (graphs-serve lane) — the
377/// pool's memory policy STATED instead of silently unbounded. The keyspace is
378/// intrinsically finite — segment keys (run_start, vt) ≤ 16 runs x 7 windows, full
379/// keys (vt, rung, hi) ≤ 7 windows x the split-rung ladder (8 rungs at 32k ctx), ~168
380/// on the q38 export — so the default (256) never engages there; the knob is the
381/// safety valve for a future export with a wider ladder. At the ceiling the pool
382/// FREEZES: existing keys keep replaying, rounds needing a new capture run the eager
383/// walk byte-identically (round-atomic — a partial refusal would mix slab- and
384/// cols-stashed layers inside one commit). No eviction by design: destroying a live
385/// exec graph re-opens the stale-address class the indirect tables exist to close,
386/// and the bounded keyspace makes reclaim worthless.
387pub(crate) fn dspark_vg_cap() -> usize {
388    static CAP: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
389    *CAP.get_or_init(|| {
390        std::env::var("MEMRA_DSPARK_VG_MAX")
391            .ok()
392            .and_then(|v| v.parse().ok())
393            .unwrap_or(256)
394    })
395}
396
397/// PROJECTED REMAINING GROWTH of the verify-graph pool, in bytes (lane/hermes-perf-fixes,
398/// 2026-08-23 — the admission accounting the "pool dwarfs spec admission reserve" finding
399/// asks for). The pool was measured at 8,852 MiB at storm-complete on the q38 export while
400/// admission's transient floor (`SPEC_SHRINK_RESERVE`) is 1.5 GiB and never charged for it:
401/// sessions admitted while the pool is cold overcommit VRAM the pool WILL hold, because the
402/// pool grows monotonically (no eviction by design) and is model-owned across sessions.
403///
404/// SELF-MEASURING, no per-model constant (generic-model law — the 8,852 MiB is a q38 number
405/// and proves nothing about another export): the debt is remaining capture slots x the
406/// MARGINAL bytes a capture adds to this device's graph mem pool.
407///
408/// MARGINAL, NOT MEAN — measured correction (box9 on-box receipt, 2026-08-23). The first
409/// version of this used the mean (`reserved / captures`) and the live serve log showed why
410/// that is wrong: with the pool's reservation flat at ~33.6 MiB across captures 1..3, the
411/// mean-based debt printed **8,556 MB, then 4,261, then 2,830** — it extrapolated capture
412/// #1's ONE-TIME shared allocation (staging buffers, stash slabs, pointer tables: sized
413/// once per pool, shared by every key) across all 256 slots. An 8.5 GB phantom reserve at
414/// boot can refuse admissions that would have fit, which is a worse defect than the
415/// under-charge this accounting exists to remove. The marginal reading prices what an
416/// ADDITIONAL key actually costs: two observations `(captures, reserved)` give
417/// `(r1 - r0) / (c1 - c0)`, which is ~0 on an export whose pool does not grow per key and
418/// tracks real growth on one that does.
419///
420/// BOOTSTRAP (only one observation so far, so growth is unmeasurable): reserve one more
421/// pool's worth — `min(remaining x mean, reserved)`. "We have measured `reserved` bytes for
422/// `captures` keys; until growth is measurable, assume at most a doubling" is fail-safe in
423/// the same direction as the old rule without the 255x extrapolation.
424///
425/// Before the FIRST capture the debt is 0 (a single capture lands well inside the existing
426/// 1.5 GiB floor). `cap` is the intrinsic freeze ceiling (`MEMRA_DSPARK_VG_MAX`; =0 freeze
427/// valve => the pool cannot grow => debt 0); at or past the cap the pool FREEZES, so the
428/// debt is 0 there too.
429pub fn dspark_vg_debt_projection(
430    captures: usize,
431    cap: usize,
432    reserved_bytes: usize,
433    prev: Option<(usize, usize)>,
434) -> usize {
435    if captures == 0 || cap == 0 {
436        return 0;
437    }
438    let remaining = cap.saturating_sub(captures);
439    if remaining == 0 {
440        return 0;
441    }
442    match prev {
443        // marginal growth between two observations of the same pool
444        Some((c0, r0)) if captures > c0 => {
445            let marginal = reserved_bytes.saturating_sub(r0) / (captures - c0);
446            remaining.saturating_mul(marginal)
447        }
448        // bootstrap: at most one more pool's worth
449        _ => remaining
450            .saturating_mul(reserved_bytes / captures)
451            .min(reserved_bytes),
452    }
453}
454/// PRE-CAPTURE VRAM RESERVE CHECK door (lane/step37-vram-admission-20260830), DEFAULT ON.
455/// A draft-graph capture attempt on a tight card used to be try-and-fail: the 2 warmup
456/// forwards + instantiate grew the pool to the edge BEFORE the OOM surfaced, and the
457/// "eager fallback" then ran on a card the failed attempt had just exhausted (the owner's
458/// single-session second-prompt OOM: capture WARN followed by 28 step-OOM engine errors,
459/// device at 5 MiB free). With the gate ON, a capture is attempted only when the device's
460/// effective free (driver free + async-pool cached) covers the capture's expected appetite
461/// PLUS a post-capture safety floor — otherwise the session falls back to eager EARLY,
462/// with headroom intact, through the same LOUD once-per-flip WARN. `=0` restores
463/// try-and-fail (diagnostics door; the trim-on-OOM recovery below stays active either way).
464pub fn spec_capture_gate_on() -> bool {
465    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
466    *ON.get_or_init(|| std::env::var("MEMRA_SPEC_CAPTURE_GATE").as_deref() != Ok("0"))
467}
468
469/// Post-capture safety floor the reserve check keeps free ON TOP of the capture's own
470/// appetite: the same measured constant class as the admission transient floor
471/// (capture arenas + verify activations — the admit-oom control fit). A capture that
472/// would leave less than this behind is not worth its eager-coverage risk.
473pub(crate) const CAPTURE_HEADROOM_FLOOR: usize = 1536 << 20;
474
475/// Pure verdict half of the pre-capture reserve check (unit-testable): given the device's
476/// driver-free and pool-cached bytes and the capture's expected `need`, returns
477/// `Some((required, effective))` when the capture must be REFUSED, `None` when it fits.
478pub(crate) fn capture_headroom_verdict(
479    driver_free: usize,
480    pool_cached: usize,
481    need: usize,
482    floor: usize,
483) -> Option<(usize, usize)> {
484    let effective = driver_free.saturating_add(pool_cached);
485    let required = need.saturating_add(floor);
486    (effective < required).then_some((required, effective))
487}
488
489/// Expected device appetite of a draft-graph capture attempt when no measurement exists
490/// yet (bootstrap only — the model-owned high-water gauge takes over after the first
491/// observed capture). Deliberately conservative and shape-derived, never a per-family
492/// constant: per (head, mode) capture the two warmups + capture each walk one head
493/// forward whose dominant transients are a handful of `n_embd` rows and one `d_vocab`
494/// logits row, retained by the keeper; the sampled tail additionally parks
495/// `k` q-slots + perturb/q buffers of `d_vocab` each.
496pub(crate) fn draft_capture_bootstrap_estimate(
497    heads: usize,
498    k: usize,
499    d_vocab: usize,
500    n_embd: usize,
501) -> usize {
502    let per_capture = 3usize // 2 warmups + capture body, each retaining its transients
503        .saturating_mul(d_vocab.saturating_add(8 * n_embd))
504        .saturating_mul(4)
505        .max(32 << 20); // instantiate + driver-side graph backing per capture, floor
506    let captures = heads.max(1).saturating_mul(2); // interior + last per head
507    let sampled_slots = (k.saturating_add(2))
508        .saturating_mul(d_vocab)
509        .saturating_mul(4);
510    captures
511        .saturating_mul(per_capture)
512        .saturating_add(sampled_slots)
513        .max(64 << 20)
514}
515
516/// OOM predicate for capture-failure recovery (engine-side twin of the worker's
517/// `is_cuda_oom` — the same quoted-text contract).
518pub(crate) fn capture_err_is_oom(reason: &str) -> bool {
519    reason.contains("CUDA_ERROR_OUT_OF_MEMORY") || reason.contains("out of memory")
520}
521
522/// Impure half of the pre-capture reserve check: reads the device, trims the async pool
523/// when the driver alone is short but cached blocks would cover it (graph instantiate and
524/// cuBLAS workspaces allocate from the DRIVER, not from our pool — a pool sitting on freed
525/// blocks starves them), and returns the refusal reason line when the capture must not be
526/// attempted. `None` = go ahead.
527pub(crate) fn capture_headroom_refusal(e: &Engine, need: usize) -> Option<String> {
528    let Ok((driver_free, _total)) = e.ctx().mem_get_info() else {
529        return None; // unreadable device: keep the historical try-and-fail behavior
530    };
531    let pool_cached = e.pool_cached_bytes();
532    // A capture may take AT MOST HALF the discretionary headroom: required =
533    // 2x appetite + two floors (owner's contract: "fall back to eager EARLY with headroom
534    // intact"). Measured escalation on the owner-shape cells: one floor of slack let the
535    // capture walk the card to the edge and the burst step-OOM'd immediately; two floors
536    // still allowed a capture whose session then OOM'd on its own admission-charged work,
537    // because the capture had consumed the memory the charge was counting on. Requiring
538    // the appetite TWICE means the card retains a whole capture's worth of room after the
539    // capture lands - enough for the session's charged classes and its peers' bursts. The
540    // capture is an optimization worth ~2-3 ms of TTFT (draft-graph lane receipts); at the
541    // margin it is never worth an OOM incident.
542    let floor = CAPTURE_HEADROOM_FLOOR.saturating_mul(2);
543    let required_need = need.saturating_mul(2);
544    let required = required_need.saturating_add(floor);
545    match capture_headroom_verdict(driver_free, pool_cached, required_need, floor) {
546        Some((required, effective)) => Some(format!(
547            "insufficient VRAM headroom for capture: effective free {}MB (driver {}MB + pool-cached \
548             {}MB) < required {}MB (2x appetite {}MB + floor {}MB); capture skipped pre-attempt",
549            effective / (1 << 20),
550            driver_free / (1 << 20),
551            pool_cached / (1 << 20),
552            required / (1 << 20),
553            need / (1 << 20),
554            floor / (1 << 20),
555        )),
556        None => {
557            if driver_free < required && pool_cached > 0 {
558                let trimmed = e.pool_trim_to_zero();
559                if trimmed > 0 {
560                    eprintln!(
561                        "[spec] pre-capture pool trim: released {}MB cached back to the driver \
562                         (driver free {}MB < required {}MB; instantiate allocates from the driver)",
563                        trimmed / (1 << 20),
564                        driver_free / (1 << 20),
565                        required / (1 << 20),
566                    );
567                }
568            }
569            None
570        }
571    }
572}
573
574/// GRAPH-LAUNCH HEADROOM FLOOR (lane/step37-vram-admission-20260830, defect 3 root
575/// cause): `cuGraphLaunch` SEGFAULTS inside libcuda (offset +0x27c87f, a null internal
576/// dereference at address 0x60) when a captured graph is dispatched into a
577/// driver-exhausted card — reproduced on this lane's box with core dumps on BOTH the
578/// pre-lane and lane binaries (multi-active step-OOM squeeze; the crashing thread sits in
579/// `CudaGraph::launch` inside `generate_spec_inner2`). The eager arms fail RECOVERABLY on
580/// the same card (a quoted CUDA OOM the park path handles), so below this driver-free
581/// floor every graph arm yields to eager for the round. A named constant, not a knob: the
582/// winning value is the default and the guard exists to make a driver segfault
583/// unreachable, not to tune anything.
584pub(crate) const GRAPH_LAUNCH_MIN_FREE: usize = 256 << 20;
585
586/// Per-round guard for the floor above. Read failure keeps serving (never a false
587/// refusal from an unreadable device); one `mem_get_info` (~microseconds) per ~25ms round.
588pub(crate) fn graph_launch_headroom_ok(e: &Engine) -> bool {
589    match e.ctx().mem_get_info() {
590        Ok((free, _total)) => free >= GRAPH_LAUNCH_MIN_FREE,
591        Err(_) => true,
592    }
593}
594
595/// One grep-stable suspension line per ROUTE (each call site holds its own
596/// process-lifetime `Once`): every captured-graph launch route below the floor names
597/// itself in the tag while keeping the same `graph replay suspended:` key the step37
598/// admission lane's squeeze cell greps for. The spec-round guard keeps its original
599/// per-generation `[spec]` line; the sweep routes (graph-launch-guard-sweep lane,
600/// 2026-08-31) note once per process — presence is what the gates assert, and a
601/// suspended round is otherwise byte-identical to its eager twin.
602pub(crate) fn graph_replay_suspended_note(route: &str) {
603    eprintln!(
604        "[{route}] graph replay suspended: driver free below the {}MB launch floor \
605         (eager arms serve; cuGraphLaunch segfaults into an exhausted card)",
606        GRAPH_LAUNCH_MIN_FREE / (1 << 20)
607    );
608}
609
610/// Engine-bundle slice 4 (fa-execupdate lane, DSF-ROUNDCOST-20260820 §6 close: "the
611/// residual gap lives in the FULL-ATTENTION per-row section"), DEFAULT ON —
612/// `MEMRA_DSPARK_FA_ROWS=0` reverts to the per-row loop: when every row of a verify
613/// round takes the v4-seqs arm on ONE `fa_split_keys` rung (the straddle law, evaluated
614/// at the round's first and last t_kv — both eligibility gates are intervals in t_kv),
615/// the qwen35 t-parallel verify's per-row KV-append + fa-decode loop collapses into the
616/// z-batched serving twins: ONE `append_quantize_kv_q8_0_q5_1_seqs` + ONE
617/// `fa_decode_vec_q_seqs_v4` + ONE combine per full-attention layer, replacing
618/// T x (4 dtod row copies + append + 3 memsets + main + combine) launches. Bytes are
619/// pinned by the batched-tick increment-2 kernel-check (seqs-vs-per-seq-loop bit
620/// identity: per-row T_kv derives in-kernel from pos_seq[z]; splits >= ns_eff write the
621/// empty partial the combine never reads, so the shared n_splits_max stride changes no
622/// bytes) and re-gated e2e by this lane's battery.
623pub(crate) fn dspark_fa_rows_on() -> bool {
624    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
625    *ON.get_or_init(|| {
626        std::env::var("MEMRA_DSPARK_FA_ROWS")
627            .map(|v| v != "0")
628            .unwrap_or(true)
629    })
630}
631
632/// `t_pred0` for the `MEMRA_DEBUG_SPEC` per-round print, sampled-safe.
633///
634/// `generate_spec_inner2` fills its `preds` vector ONLY on the greedy path (`if !sampled`), and
635/// the per-round debug print was the sole consumer in the sampled arm: `t_pred(0)` survives round
636/// 0 (`base == 0` returns `last_pred`) and from round 1 (`base == 1`, a pending bonus) indexes an
637/// EMPTY vector — `index out of bounds: the len is 0 but the index is 0`, in the GPU worker
638/// thread, which then respawns and reloads weights while the request dies. So any sampled spec
639/// request longer than one round used to kill the worker whenever `MEMRA_DEBUG_SPEC` was set:
640/// the flag crashed precisely the regime it exists to investigate.
641///
642/// Fixed at the print site, not inside the closure, so the greedy accept walk keeps its strict
643/// indexing (an out-of-range pred there is a real bug and must still be loud).
644fn debug_t_pred0(sampled: bool, base: usize, last_pred: u32, preds: &[u32]) -> String {
645    if base == 0 {
646        return last_pred.to_string();
647    }
648    match preds.get(base - 1) {
649        Some(p) => p.to_string(),
650        // sampled: the greedy per-column argmax was never run for this round.
651        None => {
652            debug_assert!(
653                sampled,
654                "greedy spec: preds[{}] missing at base {base}",
655                base - 1
656            );
657            "n/a".to_string()
658        }
659    }
660}
661
662/// `MEMRA_SKEY_PROBE=1` — sampled-draft-graph key probe (lane/graph-s-key-exactness-20260819).
663///
664/// Reports, per burst and per round, which draft chain the sampled arm chose and under which
665/// filter regime, plus the ONE observable that separates a legal filtered draft from a stale
666/// pure-temp graph replayed under filters: an accept test whose gathered `q` is exactly 0.
667/// A draft token sampled from the FILTERED softmax can never gather q=0 (it was drawn from the
668/// kept set), so `q=0` in the verify means the draft came from a distribution the verify does
669/// not believe in — and `u * 0 < p` then accepts it unconditionally.
670///
671/// Its own env var, deliberately NOT `MEMRA_DEBUG_SPEC`: that flag panicked the GPU worker on
672/// any sampled spec request past round 0 until this lane fixed it (§2 of the bank note).
673pub(crate) fn skey_probe() -> bool {
674    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
675    *ON.get_or_init(|| std::env::var("MEMRA_SKEY_PROBE").as_deref() == Ok("1"))
676}
677
678/// GRAMMAR HOOK for constrained spec decode (lane/constrained-full, 2026-08-03). The engine
679/// stays llguidance-agnostic: the server adapts its per-session grammar state behind this
680/// trait. CONTRACT (the verify-side truncation rule — token-identical to constrained plain
681/// greedy decode): the exactness walk runs UNMASKED first; the hook then (a) truncates
682/// acceptance at the first grammar-illegal accepted token, and (b) when the truncation fired
683/// or the bonus is illegal, the engine recomputes that slot as the MASKED argmax of the
684/// target's own verify column (an unmasked argmax that is grammar-legal IS the masked argmax
685/// — masking only removes tokens — so the common case pays nothing). `consume` advances the
686/// state with each EMITTED token in order; EOS handling is the implementor's job (skip).
687pub trait SpecConstraint {
688    /// -inf the current state's banned ids on a HOST logits row (prompt-tail / init-feed
689    /// masked argmax).
690    fn mask_logits(&mut self, logits: &mut [f32]) -> Result<(), String>;
691    /// Packed 32-bit bitset words of the CURRENT state's allowed set (device-mask form).
692    fn mask_words(&mut self) -> Result<Vec<u32>, String>;
693    /// Is `tok` consumable in the CURRENT state?
694    fn is_allowed(&mut self, tok: u32) -> Result<bool, String>;
695    /// Advance the state with an emitted token.
696    fn consume(&mut self, tok: u32) -> Result<(), String>;
697
698    // --- DRAFT-SIDE MASKING (lane/draft-mask, 2026-08-04) ---
699    // The drafter proposed grammar-illegal tokens under tight schemas, so verify-side
700    // truncation cut nearly every round (measured acceptance 0.467-0.513 tight vs 0.62-0.82
701    // loose, research/constrained-full-20260803). These three methods let the engine mask the
702    // DRAFT model's own sampling with the grammar's legal set, so proposals are legal by
703    // construction. The state they walk is a SPECULATIVE CLONE of the session matcher — the
704    // real state is advanced only by `consume` (emitted tokens), so verify-side truncation
705    // stays the correctness backstop and the emitted stream is unchanged by construction
706    // (an accepted draft is the target's unmasked argmax AND grammar-legal, hence the masked
707    // argmax; a cut slot is recomputed as the masked argmax either way).
708    // Default impls = feature OFF (pre-lane behaviour: unmasked drafts).
709
710    /// Is draft-side masking available on this hook? Probed ONCE per burst, before the draft
711    /// graph is captured (the mask is an in-graph node — its presence is a capture-time shape).
712    fn draft_mask_enabled(&self) -> bool {
713        false
714    }
715    /// Start a draft chain: clone the CURRENT (committed) grammar state into the speculative
716    /// slot. Called once per spec round, before the first draft position.
717    fn draft_begin(&mut self) -> Result<(), String> {
718        Ok(())
719    }
720    /// Packed 32-bit bitset words of the SPECULATIVE state's allowed set (target-vocab ids),
721    /// for the draft position about to be sampled. `None` = draft masking off (no-op).
722    fn draft_mask_words(&mut self) -> Result<Option<Vec<u32>>, String> {
723        Ok(None)
724    }
725    /// Advance the SPECULATIVE state with a PROPOSED draft token. `false` = the chain cannot
726    /// continue (EOS proposed, or an unmasked position proposed something illegal) — the
727    /// engine stops drafting; the token already pushed still goes through verify.
728    fn draft_advance(&mut self, _tok: u32) -> Result<bool, String> {
729        Ok(false)
730    }
731}
732
733/// DRAFT-MASK UPLOAD (lane/draft-mask): pull the speculative state's allowed set (TARGET-id
734/// space) from the hook, project it into the DRAFT head's vocab space, and upload it into the
735/// stable device buffer the draft chain reads. Returns false when the chain must stop drafting:
736/// the hook handed out no mask, or NO draft-vocab row is grammar-legal at this position (a
737/// trimmed FR-Spec head genuinely cannot propose a legal token there — masking it would leave
738/// a fully-banned row whose argmax is meaningless, so the round drafts fewer tokens and the
739/// verify emits the masked argmax as usual).
740fn upload_draft_mask(
741    e: &Engine,
742    c: &mut dyn SpecConstraint,
743    dst: &mut CudaSlice<u32>,
744    d2t: Option<&Vec<u32>>,
745    d_vocab: usize,
746    words: usize,
747) -> Result<bool, Box<dyn std::error::Error>> {
748    let Some(tw) = c
749        .draft_mask_words()
750        .map_err(|e2| format!("constraint: {e2}"))?
751    else {
752        return Ok(false);
753    };
754    let bit = |t: usize| -> bool {
755        let w = t >> 5;
756        w < tw.len() && (tw[w] >> (t & 31)) & 1 == 1
757    };
758    let mut buf = vec![0u32; words];
759    match d2t {
760        // TRIMMED draft head: row i proposes target id d2t[i] — permute the mask accordingly.
761        Some(map) => {
762            for (i, &t) in map.iter().enumerate().take(d_vocab) {
763                if bit(t as usize) {
764                    buf[i >> 5] |= 1u32 << (i & 31);
765                }
766            }
767        }
768        // UNTRIMMED: draft ids ARE target ids; the packed words transfer verbatim (a short
769        // mask leaves the padded tail zeroed == banned, same rule as constrained::apply_mask).
770        None => {
771            let n = tw.len().min(words);
772            buf[..n].copy_from_slice(&tw[..n]);
773        }
774    }
775    if buf.iter().all(|w| *w == 0) {
776        return Ok(false);
777    }
778    e.htod_u32_into(dst, &buf)?;
779    Ok(true)
780}
781
782/// Keep the full token-embedding table in host memory and upload only the rows needed by each
783/// MTP/verify step. This is an exact memory-capacity seam for very large BF16 vocab tables: host
784/// gather expands the same source bits to f32, and only O(T*n_embd) bytes cross PCIe per step.
785/// CUDA-graph/round-stream draft paths require device token ids and therefore stay disabled.
786pub(crate) fn spec_host_embd() -> bool {
787    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
788    *ON.get_or_init(|| std::env::var("MEMRA_SPEC_HOST_EMBD").as_deref() == Ok("1"))
789}
790
791/// VERIFY-TIER TRUNK LAUNCH-FUSION (default ON since 2026-07-09; MEMRA_SPEC_FUSED_T=0 reverts — lane/close35b): extend
792/// the t=1 fused2/fused3 Q8_0 trunk launches to the batched verify tier (t=2-4, the K=1..3
793/// verify shapes). At t>1 the trunk pairs/triples (35B wqkv+wqkv_gate, wq/wk/wv,
794/// gate_shexp+up_shexp) each run a separate `matmul_decode_exact` — one q8_1 re-quantize of the
795/// SAME activation plus one _b2/_b4 launch per tensor. The fused twins share ONE quantize and
796/// ONE launch per group; per (tensor,token,row) the kernel body is q8_0_mmvq_batched verbatim
797/// with the identical row mapping -> BIT-IDENTICAL by construction (kernel-check pins it,
798/// run-spec K=1..8 + acceptance identity arbitrate e2e).
799pub(crate) fn spec_fused_t() -> bool {
800    static F: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
801    // DEFAULT ON since 2026-07-09 (MEMRA_SPEC_FUSED_T=0 reverts): verify t=2-4 trunk launch-fusion
802    // (fused2/fused3 Q8_0 batched twins, bit-identical by construction — m=1 block-offset split on
803    // the batched body). m=2 marginal token 2117->1762us; 35B daily: p3 +3.7% (crosses llama), p2 +5%.
804    *F.get_or_init(|| {
805        std::env::var("MEMRA_SPEC_FUSED_T")
806            .map(|v| v != "0")
807            .unwrap_or(true)
808    })
809}
810
811/// zeros/uninit switch for verify-path buffers that are FULLY OVERWRITTEN before any read.
812/// Only call this on such buffers — the lean contract is "identical bytes by construction".
813/// TOKEN-ID GUARD for every id that reaches an embed gather (#87 family).
814///
815/// A device argmax seeds its running index with 0x7FFFFFFF and replaces it only through
816/// comparisons, all of which are FALSE against NaN. An all-NaN logits row therefore returns
817/// the sentinel, and the next thing done with a token id is `embed_row(id)` — table +
818/// ~4.6 TB, never mapped, an MMU fault that kills the CUDA context for the whole process
819/// (research/pp2spec-crash-20260807). The draft chain and the GREEDY verify walk already
820/// trap this; the SAMPLED verify bonus, the boundary sampler and the replay arm's last_pred
821/// did not, which is why the recoverable fault on the greedy instrument is a TERMINAL one on
822/// the vendor-default sampled shape we actually serve.
823pub(crate) fn guard_vocab_token(
824    tok: u32,
825    n_vocab: usize,
826    what: &str,
827) -> Result<u32, Box<dyn std::error::Error>> {
828    if (tok as usize) >= n_vocab {
829        return Err(format!(
830            "{what}: token id 0x{tok:08x} >= n_vocab {n_vocab} — an all-NaN logits row left \
831             the device argmax's init sentinel in place; refusing to dereference the embed \
832             row (#87 trap)"
833        )
834        .into());
835    }
836    Ok(tok)
837}
838
839/// SPEC NaN-ORIGIN SCAN (`MEMRA_SPEC_NAN_SCAN=1`, DEFAULT OFF, diagnostic only).
840///
841/// The `#87` trap reports an all-NaN VERIFY logits column, which says the poison reached the
842/// head but not where it entered. With the scan armed the verify walk syncs and reads back
843/// every layer's output, so the FIRST layer whose residual carries a NaN names itself with the
844/// round's row and position. Off by default and never on a serving path: it costs one host
845/// sync + one `t*n_embd` D2H per layer, and the syncs change scheduling (so a run that stops
846/// reproducing under the scan is itself a datum, not an all-clear).
847///
848/// Rollback seam: unset `MEMRA_SPEC_NAN_SCAN` (or set it to 0). Every call site is behind
849/// `spec_nan_scan()`, so the default path keeps the exact launch sequence it had.
850pub(crate) fn spec_nan_scan() -> bool {
851    spec_nan_scan_level() > 0
852}
853
854/// `MEMRA_SPEC_NAN_SCAN` as a LEVEL, not a boolean. `1` scans each layer's residual, which
855/// names the layer. `2` also scans INSIDE the t-column layer body — the per-column attention
856/// output, the deferred-column o-proj/fa2 join, the post-attention norm and the routed-MoE
857/// output — because "layer 20 poisons row 0" does not say whether the attention or the routed
858/// MoE produced it, and those are different bugs with different fixes.
859pub(crate) fn spec_nan_scan_level() -> u8 {
860    static LVL: std::sync::OnceLock<u8> = std::sync::OnceLock::new();
861    *LVL.get_or_init(|| match std::env::var("MEMRA_SPEC_NAN_SCAN").as_deref() {
862        Ok("1") => 1,
863        Ok("2") => 2,
864        _ => 0,
865    })
866}
867
868/// Read back `[rows, cols]` and fail with the first NaN's coordinates. `what` names the
869/// producer (layer index, walk arm) so the error line is the localization.
870/// VERIFY-ARM RECEIPT (rides `MEMRA_SPEC_NAN_SCAN>=1`, bounded to 200 lines).
871///
872/// Names, per trunk layer, WHICH attention arm the t-column walk actually took. This exists
873/// because the level-1 residual scan below sat only on the non-fused tail: the fused
874/// rope+append+fa arm ends in `continue`, so every layer that fused was NEVER SCANNED and
875/// silently read as "clean". A poisoned residual therefore first reported at the next
876/// non-fused layer, which is how "layer 20 creates the poison" could be true of the scan and
877/// false of the engine. Also carries the row-table lookup counter, so "the fused path never
878/// ran" is distinguishable from "it ran and was innocent".
879/// KV-PLANE SCAN (`MEMRA_KV_PLANE_SCAN=1`, DEFAULT OFF, diagnostic only).
880///
881/// Reads back the STAGED rows of a layer's distributed K/V planes and reports the first row
882/// whose quantization scale is not finite. No kernel required: q8_0 blocks are
883/// `[half d][32 x i8]` and q5_1 blocks carry `half d` then `half m`, so the fp16 scale at the
884/// head of each block is host-checkable straight out of the byte plane.
885///
886/// It exists because the level-2 bad-row bitmap says EVERY verify row is non-finite at a
887/// global-attention layer's join, and row r attends a strict superset of row r-1's keys: that
888/// implicates the shared KV history those rows walk, not per-column staging. "The attention
889/// output is NaN" and "the KV history it attends is already NaN" are different bugs with
890/// different owners, and nothing measured so far separates them. A first-corrupt-row index
891/// also dates the corruption against the prime/decode boundary.
892///
893/// Bounded hard: only layers whose geometry has NO window (the global planes), only the first
894/// `MEMRA_KV_PLANE_SCAN_ROUNDS` verify rounds of a process (default 2), and it copies only
895/// `[0, staged_len)`, which is ~1.6 MB at the 1480-token repro rather than the 262144-row
896/// provision. It still syncs per layer, so it is never a serving or a measured-perf arm.
897pub(crate) fn kv_plane_scan_on() -> bool {
898    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
899    *ON.get_or_init(|| std::env::var("MEMRA_KV_PLANE_SCAN").as_deref() == Ok("1"))
900}
901
902fn kv_plane_scan_rounds() -> usize {
903    static R: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
904    *R.get_or_init(|| {
905        std::env::var("MEMRA_KV_PLANE_SCAN_ROUNDS")
906            .ok()
907            .and_then(|v| v.parse().ok())
908            .unwrap_or(2)
909    })
910}
911
912/// First non-finite fp16 block scale in `bytes`, as (block index, raw u16), scanning one
913/// scale every `stride` bytes. Returns None when every block scale is finite.
914fn first_bad_scale(bytes: &[u8], stride: usize) -> Option<(usize, u16)> {
915    if stride == 0 {
916        return None;
917    }
918    for (i, blk) in bytes.chunks_exact(stride).enumerate() {
919        let raw = u16::from_le_bytes([blk[0], blk[1]]);
920        if half_is_non_finite(raw) {
921            return Some((i, raw));
922        }
923    }
924    None
925}
926
927/// IEEE binary16: exponent all ones is Inf or NaN, whatever the mantissa says.
928fn half_is_non_finite(raw: u16) -> bool {
929    (raw & 0x7C00) == 0x7C00
930}
931
932/// Scan one layer's staged K/V planes for a non-finite quantization scale. Returns the
933/// receipt line, or None when the layer is out of scope or every scale is finite.
934pub(crate) fn scan_kv_plane(
935    e: &crate::Engine,
936    distributed: &memra_kv::ResidentTpKvCache,
937    il: usize,
938    pos0: usize,
939) -> Result<(), Box<dyn std::error::Error>> {
940    // One "round" is one pos0, not one layer: the walk visits 45 layers per verify. The
941    // default of 2 rounds is for a fault that shows up immediately; the step37 repro does not
942    // fire until rep 3 or later, i.e. round ~60 of the process, so that arm MUST raise
943    // MEMRA_KV_PLANE_SCAN_ROUNDS or it will scan only the two rounds that were never going to
944    // be poisoned and report a clean history it never looked at.
945    static ROUNDS: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
946    static LAST_POS: std::sync::atomic::AtomicUsize =
947        std::sync::atomic::AtomicUsize::new(usize::MAX);
948    if LAST_POS.swap(pos0, std::sync::atomic::Ordering::Relaxed) != pos0 {
949        ROUNDS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
950    }
951    if ROUNDS.load(std::sync::atomic::Ordering::Relaxed) > kv_plane_scan_rounds() {
952        return Ok(());
953    }
954    let staged = distributed.staged_len();
955    if staged == 0 {
956        return Ok(());
957    }
958    // ENGAGEMENT RECEIPT. This scan prints only on corruption, so `kvbad=0` in a cell would
959    // read the same whether the history was clean or the scan never ran once. Bounded so a
960    // 45-layer walk cannot flood the log.
961    static SEEN: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
962    let seen = SEEN.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
963    let (ktb, vtb) = (distributed.k_tok_bytes(), distributed.v_tok_bytes());
964    if seen < 4 {
965        eprintln!(
966            "[kv-plane] engaged #{seen} layer {il} pos0={pos0} staged={staged} \
967             ktok={ktb} vtok={vtb} (scan armed; a corrupt plane prints its own line)"
968        );
969    }
970    for rank in 0..distributed.ranks().len() {
971        let Some(rc) = distributed.rank(rank) else {
972            continue;
973        };
974        // q8_0 K blocks are [half d][32 x i8] = 34B; q5_1 V blocks lead with half d then half m.
975        let kbytes = e.dtoh_u8_view(&rc.k().slice(0..staged * ktb))?;
976        let vbytes = e.dtoh_u8_view(&rc.v().slice(0..staged * vtb))?;
977        let kbad = first_bad_scale(&kbytes, 34);
978        let vbad = first_bad_scale(&vbytes, 24);
979        if kbad.is_some() || vbad.is_some() {
980            let row = |b: Option<(usize, u16)>, tok: usize| {
981                b.map(|(i, raw)| format!("blk {i} (row {}) raw={raw:#06x}", i * 34 / tok.max(1)))
982                    .unwrap_or_else(|| "clean".into())
983            };
984            eprintln!(
985                "[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",
986                row(kbad, ktb),
987                row(vbad, vtb)
988            );
989            return Ok(());
990        }
991    }
992    Ok(())
993}
994
995pub(crate) fn verify_arm_receipt(
996    arm: &str,
997    il: usize,
998    pos0: usize,
999    t: usize,
1000    staged: Option<usize>,
1001) {
1002    static N: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
1003    if N.fetch_add(1, std::sync::atomic::Ordering::Relaxed) >= 200 {
1004        return;
1005    }
1006    eprintln!(
1007        "[verify-arm] layer {il} arm={arm} pos0={pos0} t={t} staged_len={} rows_tab_lookups={}",
1008        staged.map(|v| v as i64).unwrap_or(-1),
1009        crate::tp::ROWS_TAB_ENGAGED.load(std::sync::atomic::Ordering::Relaxed)
1010    );
1011}
1012
1013pub(crate) fn nan_scan_rows(
1014    e: &Engine,
1015    buf: &CudaSlice<f32>,
1016    rows: usize,
1017    cols: usize,
1018    what: &str,
1019) -> Result<(), Box<dyn std::error::Error>> {
1020    // The readback is also the ATTRIBUTION point for an asynchronous fault: a
1021    // CUDA_ERROR_ILLEGAL_ADDRESS raised by any launch since the previous scan surfaces on this
1022    // sync, and the bare DriverError names nothing. Wrapping it with `what` turns "the process
1023    // died somewhere" into "it died at or before this layer, on this row, at this position".
1024    let host = e.dtoh(buf).map_err(|err| -> Box<dyn std::error::Error> {
1025        format!(
1026            "spec nan-scan: sync at {what} FAILED: {err} — the fault is at or before \
1027                     this point in the walk"
1028        )
1029        .into()
1030    })?;
1031    if host.len() < rows * cols {
1032        return Err(format!(
1033            "nan-scan {what}: buffer holds {} < {rows}x{cols}",
1034            host.len()
1035        )
1036        .into());
1037    }
1038    // SCAN EVERY ROW BEFORE REPORTING. A first-hit return says "row 0 is bad" and leaves the
1039    // other rows UNEXAMINED, which is exactly the bit that discriminates the two mechanisms: in
1040    // the t-column verify, row 0 attends keys [0..p+1) and row 1 attends [0..p+2), a strict
1041    // superset, so poison in the SHARED KV history must appear in BOTH rows, while poison in
1042    // per-column staging can appear in one. Report the whole map.
1043    let mut per_row: Vec<usize> = Vec::with_capacity(rows);
1044    let mut first_bad: Option<(usize, usize)> = None;
1045    for r in 0..rows {
1046        let row = &host[r * cols..(r + 1) * cols];
1047        let bad = row.iter().filter(|v| !v.is_finite()).count();
1048        per_row.push(bad);
1049        if bad > 0 && first_bad.is_none() {
1050            first_bad = Some((r, row.iter().position(|v| !v.is_finite()).unwrap_or(0)));
1051        }
1052    }
1053    if let Some((r0, c0)) = first_bad {
1054        let map: String = per_row
1055            .iter()
1056            .map(|&b| if b == 0 { '.' } else { 'X' })
1057            .collect();
1058        return Err(format!(
1059            "spec nan-scan: {what} produced non-finite values — rows[{rows}] map={map} \
1060             counts={per_row:?} of {cols} each; first at row {r0} element {c0}. Both rows bad \
1061             implicates shared state (the KV history this layer reads); one row bad implicates \
1062             per-column staging."
1063        )
1064        .into());
1065    }
1066    Ok(())
1067}
1068
1069fn vbuf(e: &Engine, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
1070    if spec_lean() { e.uninit(n) } else { e.zeros(n) }
1071}
1072
1073/// Scratch KV for the MTP block (one full-attn layer).
1074///
1075/// PERSISTENT MODE (default, 2026-07-03 — the acceptance lever): sized cap = max_ctx and kept in
1076/// sync with the COMMITTED sequence — slot p holds the MTP block's K/V for committed token p
1077/// (roped p+1, the chain's rope convention), so the draft chain's self-attention sees the FULL
1078/// committed history instead of only the current round's 1..K+1 chain tokens (the reference
1079/// engine's "mtp_update" design). Entries come from two sources:
1080///   - chain appends: accepted positions KEEP their chain-computed entries (embedding exact,
1081///     hidden chain-approximate — the reference engine accepts the same);
1082///   - `mtp_kv_fill` batches: prompt positions + the last-draft position on full accept, computed
1083///     from EXACT trunk hiddens (K/V-only MTP-block pass, no attention/FFN/lm_head).
1084///     Rejected drafts / p-min extras / pseudo-seed appends are all discarded by the round-start
1085///     `set_len` truncation (the KvLayer len mechanism — §C rollback for the draft side).
1086///     Multi-turn spec-decode session (2026-07-05): trunk Cache + persistent MTP draft scratch +
1087///     the committed token list, alive across generate_spec_session calls. Turn N+1 primes ONLY its
1088///     suffix (chunked continuation prime over the quantized past) and mtp_kv_fill's its suffix rows,
1089///     then the round loop runs unchanged. `last_h` carries the pre-output_norm hidden of the last
1090///     committed row across turns (the predecessor-pairing seed + fill anchor).
1091///     Per-request sampling config for the sampled-spec serve path.
1092#[derive(Clone, Copy, Debug)]
1093pub struct SpecSampling {
1094    pub temp: f32,
1095    pub seed: u64,
1096    pub top_k: i32,            // 0 = off
1097    pub top_p: f32,            // 1.0 = off
1098    pub min_p: f32,            // 0.0 = off
1099    pub penalty_last_n: usize, // 0 = penalties off
1100    pub penalty_repeat: f32,
1101    pub penalty_freq: f32,
1102    pub penalty_present: f32,
1103}
1104
1105impl SpecSampling {
1106    /// Non-identity penalties requested — THE `pen_on` predicate (one definition; the
1107    /// same group-off rule `SamplerIdentity::of` canonicalizes: a window with neutral
1108    /// coefficients is penalties-absent). Both spec routes and the dspark accept walk
1109    /// key their penalty arms off this.
1110    pub fn pen_on(&self) -> bool {
1111        self.penalty_last_n > 0
1112            && (self.penalty_repeat != 1.0
1113                || self.penalty_freq != 0.0
1114                || self.penalty_present != 0.0)
1115    }
1116}
1117
1118/// Which draft source a spec session is pinned to. The ENGINE-LEVEL half of
1119/// `DraftSourcePlan` (memra-gguf `model_plan.rs`, always general): the plan states what the
1120/// model DECLARES, this states what actually LOADED and therefore what the session runs.
1121/// Pinned at session creation for the session's lifetime.
1122///
1123/// Family-agnostic on purpose (lane/glm5-extract2, the DraftSource seam): glm5 is today's
1124/// consumer with NativeMtp | Dflash2; the hy3/qwen-next spec lanes select through the same
1125/// three-way law instead of re-deriving it. What each family still owns is the per-session
1126/// STATE behind the kind (see `dflash.rs`'s seam note for why that half is not a trait yet).
1127#[derive(Clone, Copy, Debug, PartialEq, Eq)]
1128pub enum DraftSourceKind {
1129    /// The model's own embedded NextN/MTP head.
1130    NativeMtp,
1131    /// A separately loaded DFlash2 block-diffusion drafter
1132    /// ([`crate::dflash::DflashDrafter`]).
1133    Dflash2,
1134}
1135
1136/// The uniform draft-source selection law. Pure — no env, no engine, no family types — so it
1137/// is CPU-gateable and so every spec family answers "which source" the same way.
1138///
1139/// THE LAW, in precedence order:
1140/// 1. A LOADED DFlash2 drafter IS the source. The operator asked for it by name (a set
1141///    drafter flag that cannot load is already a loud boot failure, never a silent
1142///    fallback), and the family's embedded head is deliberately NOT loaded for this source —
1143///    it is a full trunk layer of VRAM.
1144/// 2. Otherwise the embedded head, and only when the PLAN declares an embedded source: a
1145///    loaded head under a plan that does not declare `Embedded` is a load-path bug, not a
1146///    draft source, and it is refused by name rather than drafted from.
1147/// 3. Otherwise there is no draft source and speculative decode must refuse before drafting.
1148pub fn resolve_draft_source_kind(
1149    plan: memra_gguf::model_plan::DraftSourcePlan,
1150    embedded_head_loaded: bool,
1151    dflash_loaded: bool,
1152) -> Result<DraftSourceKind, String> {
1153    use memra_gguf::model_plan::DraftSourcePlan as P;
1154    if dflash_loaded {
1155        return Ok(DraftSourceKind::Dflash2);
1156    }
1157    if embedded_head_loaded {
1158        if plan != P::Embedded {
1159            return Err(format!(
1160                "an embedded draft head is loaded but the ModelPlan declares \
1161                 draft_source={plan:?} — refused rather than drafting from a head the plan \
1162                 does not claim"
1163            ));
1164        }
1165        return Ok(DraftSourceKind::NativeMtp);
1166    }
1167    Err(format!(
1168        "no draft source loaded (ModelPlan declares draft_source={plan:?}): speculative \
1169         decode has nothing to draft from"
1170    ))
1171}
1172
1173#[cfg(test)]
1174mod draft_source_kind_tests {
1175    use super::{DraftSourceKind, resolve_draft_source_kind};
1176    use memra_gguf::model_plan::DraftSourcePlan as P;
1177
1178    #[test]
1179    fn a_loaded_drafter_wins_over_a_co_loaded_embedded_head() {
1180        // The operator asked for the drafter BY NAME (a set drafter flag that cannot load is
1181        // already a loud boot failure), so it takes precedence under every plan value —
1182        // including ExternalArtifact, which is what a pack declares when the draft weights
1183        // are not in the model file.
1184        for plan in [P::Embedded, P::ExternalArtifact, P::None] {
1185            assert_eq!(
1186                resolve_draft_source_kind(plan, true, true).unwrap(),
1187                DraftSourceKind::Dflash2,
1188                "plan {plan:?}: a loaded drafter must win"
1189            );
1190            assert_eq!(
1191                resolve_draft_source_kind(plan, false, true).unwrap(),
1192                DraftSourceKind::Dflash2
1193            );
1194        }
1195    }
1196
1197    #[test]
1198    fn the_embedded_head_is_the_source_only_under_a_plan_that_claims_it() {
1199        assert_eq!(
1200            resolve_draft_source_kind(P::Embedded, true, false).unwrap(),
1201            DraftSourceKind::NativeMtp
1202        );
1203        // A head loaded under a plan that does not declare Embedded is a LOAD-PATH BUG, not a
1204        // draft source. Unreachable on glm5 today (its pack hardcodes Embedded and the head
1205        // only loads under it) — which is exactly why it is pinned here: an unreachable
1206        // refusal with no arm is an untested refusal, and the next family is the one that
1207        // makes it reachable.
1208        for plan in [P::ExternalArtifact, P::None] {
1209            let err = resolve_draft_source_kind(plan, true, false)
1210                .expect_err("a head under a non-Embedded plan must refuse");
1211            assert!(err.contains("does not claim"), "{err}");
1212            assert!(err.contains(&format!("{plan:?}")), "{err}");
1213        }
1214    }
1215
1216    #[test]
1217    fn nothing_loaded_refuses_before_drafting_and_names_the_plan() {
1218        for plan in [P::Embedded, P::ExternalArtifact, P::None] {
1219            let err =
1220                resolve_draft_source_kind(plan, false, false).expect_err("no source must refuse");
1221            assert!(err.contains("no draft source loaded"), "{err}");
1222            assert!(err.contains(&format!("{plan:?}")), "{err}");
1223        }
1224    }
1225}
1226
1227/// `MEMRA_SPEC_PMIN` break semantics over per-slot draft confidences (the chain break this
1228/// module's drafting loops apply inline: `p < p_min && (j > 0 || pmin0)`): keep the longest
1229/// prefix whose every slot clears `p_min`; slot 0 survives a miss unless PMIN0 arms
1230/// zero-draft rounds. Prefix truncation is forced by the accept rule anyway (a kept slot
1231/// after a dropped one could never commit — the dspark confidence-slot argument). Pure so
1232/// the rule is CPU-gateable; the SHARED K-policy surface every spec family consumes
1233/// (hoisted from the glm5 loop, lane/glm5-extract-general).
1234pub fn spec_conf_keep(q: &[f32], p_min: f32, pmin0: bool) -> usize {
1235    if p_min <= 0.0 {
1236        return q.len();
1237    }
1238    let mut kept = 0usize;
1239    for (j, &qj) in q.iter().enumerate() {
1240        if qj < p_min && (j > 0 || pmin0) {
1241            break;
1242        }
1243        kept += 1;
1244    }
1245    kept
1246}
1247
1248/// Host Philox4x32-10 uniform in (0,1) — mirrors spec_sample.cu's `philox4`/`u01` with the
1249/// ctr_lo tag 0xFFFF_FFFE, so the host accept-test stream never collides with any device
1250/// sampling event (device Gumbel uses (i>>2, stream_pos); device residual uses 0xFFFF_FFFD).
1251/// One value per (seed, ctr) EVENT; callers own the counter discipline. Extracted verbatim
1252/// from generate_spec_inner2's closure for the dspark sampled-admission walk (the two paths
1253/// MUST consume the identical stream construction — two ad-hoc Philox copies drifting apart
1254/// is a distributional bug, not a style problem).
1255pub(crate) fn host_u01(seed: u64, ctr: u32) -> f32 {
1256    let (m0, m1) = (0xD2511F53u32, 0xCD9E8D57u32);
1257    let (mut c0, mut c1, mut c2, mut c3) = (0xFFFF_FFFEu32, ctr, 0u32, 0u32);
1258    let (mut k0, mut k1) = ((seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
1259    for _ in 0..10 {
1260        let (h0, l0) = (((m0 as u64 * c0 as u64) >> 32) as u32, m0.wrapping_mul(c0));
1261        let (h1, l1) = (((m1 as u64 * c2 as u64) >> 32) as u32, m1.wrapping_mul(c2));
1262        let (n0, n1, n2, n3) = (h1 ^ c1 ^ k0, l1, h0 ^ c3 ^ k1, l0);
1263        c0 = n0;
1264        c1 = n1;
1265        c2 = n2;
1266        c3 = n3;
1267        k0 = k0.wrapping_add(0x9E3779B9);
1268        k1 = k1.wrapping_add(0xBB67AE85);
1269    }
1270    (c0 as f32 + 1.0) * (1.0 / 4294967296.0)
1271}
1272
1273/// Tracked draft positions for [`SpecTelemetry`] (serve K defaults to 3; the run-spec gate
1274/// sweeps K=1..8, and MEMRA_SPEC_CAPMAX defaults to 7 — 8 covers every tuned config).
1275pub const SPEC_TELEM_POS: usize = 8;
1276
1277/// Always-on per-draft-position acceptance telemetry (lane/accept-telemetry, 2026-08-05 —
1278/// the llama.cpp #26389 / vLLM spec-decode counter schema, upstream-sweeps 2026-08-05).
1279/// Lives on the [`SpecSession`] and accumulates across bursts; the serve worker diffs a
1280/// stashed copy per burst for its per-model /metrics aggregation and per-request usage.
1281/// Same normalization as the `[spec-stats]` line: p-min-discarded chain tokens are counted
1282/// in NEITHER drafted nor accepted.
1283#[derive(Clone, Copy, Default, Debug)]
1284pub struct SpecTelemetry {
1285    /// verify rounds completed (a round-stream burst counts each of its M rounds).
1286    pub rounds: u64,
1287    /// tokens drafted / accepted across all rounds.
1288    pub drafted: u64,
1289    pub accepted: u64,
1290    /// how often draft position j (0-based within a round's chain) was offered / accepted.
1291    /// Positions >= SPEC_TELEM_POS are untracked (totals still count them). The opt-in
1292    /// round-stream arm (MEMRA_SPEC_STREAM=1) reads back only totals, so under it these
1293    /// arrays cover the standard-path rounds only and their sums may undercount the totals.
1294    pub pos_drafted: [u64; SPEC_TELEM_POS],
1295    pub pos_accepted: [u64; SPEC_TELEM_POS],
1296}
1297
1298impl SpecTelemetry {
1299    /// Fieldwise `self - prev` — the worker's per-burst delta off a copy stashed before the
1300    /// burst call. Saturating: a caller diffing against the wrong snapshot gets zeros, not
1301    /// a wrapped counter.
1302    pub fn delta_since(&self, prev: &SpecTelemetry) -> SpecTelemetry {
1303        let mut d = SpecTelemetry {
1304            rounds: self.rounds.saturating_sub(prev.rounds),
1305            drafted: self.drafted.saturating_sub(prev.drafted),
1306            accepted: self.accepted.saturating_sub(prev.accepted),
1307            ..Default::default()
1308        };
1309        for j in 0..SPEC_TELEM_POS {
1310            d.pos_drafted[j] = self.pos_drafted[j].saturating_sub(prev.pos_drafted[j]);
1311            d.pos_accepted[j] = self.pos_accepted[j].saturating_sub(prev.pos_accepted[j]);
1312        }
1313        d
1314    }
1315    /// Fieldwise `self += d` — the worker's per-model aggregation.
1316    pub fn merge(&mut self, d: &SpecTelemetry) {
1317        self.rounds += d.rounds;
1318        self.drafted += d.drafted;
1319        self.accepted += d.accepted;
1320        for j in 0..SPEC_TELEM_POS {
1321            self.pos_drafted[j] += d.pos_drafted[j];
1322            self.pos_accepted[j] += d.pos_accepted[j];
1323        }
1324    }
1325
1326    /// Mean accepted draft-prefix length per verify round (tau).
1327    pub fn tau(&self) -> f64 {
1328        if self.rounds > 0 {
1329            self.accepted as f64 / self.rounds as f64
1330        } else {
1331            0.0
1332        }
1333    }
1334}
1335
1336/// Session-lifetime atomic acceptance counters. The verifier records only after the greedy or
1337/// rejection-sampling walk has resolved on the host, so these relaxed increments add no GPU
1338/// launch, synchronization, allocation, or ordering dependency to the numeric path.
1339struct SpecTelemetryCounters {
1340    rounds: AtomicU64,
1341    drafted: AtomicU64,
1342    accepted: AtomicU64,
1343    pos_drafted: [AtomicU64; SPEC_TELEM_POS],
1344    pos_accepted: [AtomicU64; SPEC_TELEM_POS],
1345}
1346
1347impl Default for SpecTelemetryCounters {
1348    fn default() -> Self {
1349        Self {
1350            rounds: AtomicU64::new(0),
1351            drafted: AtomicU64::new(0),
1352            accepted: AtomicU64::new(0),
1353            pos_drafted: std::array::from_fn(|_| AtomicU64::new(0)),
1354            pos_accepted: std::array::from_fn(|_| AtomicU64::new(0)),
1355        }
1356    }
1357}
1358
1359impl SpecTelemetryCounters {
1360    fn record_round(&self, drafted: usize, accepted: usize) {
1361        debug_assert!(accepted <= drafted);
1362        self.rounds.fetch_add(1, Ordering::Relaxed);
1363        self.drafted.fetch_add(drafted as u64, Ordering::Relaxed);
1364        self.accepted.fetch_add(accepted as u64, Ordering::Relaxed);
1365        for counter in self.pos_drafted.iter().take(drafted) {
1366            counter.fetch_add(1, Ordering::Relaxed);
1367        }
1368        for counter in self.pos_accepted.iter().take(accepted) {
1369            counter.fetch_add(1, Ordering::Relaxed);
1370        }
1371    }
1372
1373    /// Round-stream keeps each round's accept length on device; retain exact scalar totals while
1374    /// leaving the per-position arrays untouched, matching the pre-existing telemetry contract.
1375    fn record_totals(&self, rounds: usize, drafted: usize, accepted: usize) {
1376        self.rounds.fetch_add(rounds as u64, Ordering::Relaxed);
1377        self.drafted.fetch_add(drafted as u64, Ordering::Relaxed);
1378        self.accepted.fetch_add(accepted as u64, Ordering::Relaxed);
1379    }
1380
1381    fn snapshot(&self) -> SpecTelemetry {
1382        SpecTelemetry {
1383            rounds: self.rounds.load(Ordering::Relaxed),
1384            drafted: self.drafted.load(Ordering::Relaxed),
1385            accepted: self.accepted.load(Ordering::Relaxed),
1386            pos_drafted: std::array::from_fn(|j| self.pos_drafted[j].load(Ordering::Relaxed)),
1387            pos_accepted: std::array::from_fn(|j| self.pos_accepted[j].load(Ordering::Relaxed)),
1388        }
1389    }
1390}
1391
1392pub struct SpecSession {
1393    pub(crate) cache: Cache,
1394    pub(crate) scratch: MtpScratch,
1395    /// Every token whose state the caches hold, in order (prompt turns + generated), INCLUDING
1396    /// overshoot: spec commits accepted drafts past max_new; those rows are in the caches, so the
1397    /// session must count them. Callers render output from this, not from their own echo.
1398    pub committed: Vec<u32>,
1399    /// Pre-output_norm hidden of the LAST committed row (device). None before the first turn.
1400    pub(crate) last_h: Option<CudaSlice<f32>>,
1401    /// Greedy argmax predicting the token AFTER committed.last() (from the last turn's final
1402    /// logits). Fuels empty-suffix continuation bursts (serve): the next turn emits this token
1403    /// first, feeds it, and the round loop resumes without any prime. None before the first turn.
1404    pub next_pred: Option<u32>,
1405    /// SAMPLED-SPEC stream continuity across bursts: Philox event counters persist here so a
1406    /// session's randomness never repeats between generate_spec_session calls. (0,0) at admit.
1407    pub sctr: u32,
1408    pub uctr: u32,
1409    /// PERSISTENT DRAFT-GRAPH CONTEXT (2026-08-01, the serve-burst fixed-cost fix): the captured
1410    /// draft graph(s) + every device I/O buffer they bake, carried ACROSS generate_spec_session
1411    /// calls. Before this, every serve burst re-captured the draft graph (2 warmup forwards +
1412    /// instantiate) — measured ~16ms/burst on H100 q27 (MEMRA_SPEC_BURST sweep,
1413    /// research/spec-serving-20260801). None before the first turn; error paths drop it
1414    /// (next burst recaptures — serve retires errored sessions anyway).
1415    pub(crate) draft_ctx: Option<DraftGraphCtx>,
1416    /// PENDING-CARRY across bursts (2026-08-01, the serve burst-boundary fix): the bonus token
1417    /// emitted by the last round but NOT committed to the caches. The old tail committed it with
1418    /// a solo T=1 trunk pass (+ draft fill), and the next burst's setup fed the stashed next_pred
1419    /// with ANOTHER solo pass — 2x ~11.5ms/burst measured on H100 q27 ([spec-setup] trace).
1420    /// Carrying it lets the next empty-suffix greedy burst consume it as round-0 verify col 0,
1421    /// exactly like a mid-burst full-accept boundary (no solo passes). INVARIANT: when set,
1422    /// `committed` (== cache rows) EXCLUDES this token although it was already emitted in the
1423    /// last burst's output, and `last_h` holds the hidden of the last COMMITTED row (its
1424    /// predecessor — the chain-seed/fill anchor). `next_pred` is None (unknown without the
1425    /// commit pass). Non-empty-suffix or sampled turns must flush first (spec_flush_pending);
1426    /// generate_spec_session_sampled does this at entry, and serve parks only flushed sessions.
1427    pub pending_tok: Option<u32>,
1428    /// SESSION-AFFINITY TURN CHECKPOINT (lane/session-affinity, 2026-08-05): the state at this
1429    /// turn's PROMPT-END boundary, retained so a later turn can REWIND here. See
1430    /// [`SpecCheckpoint`]. Refreshed by every non-empty prime; None until the first one, and on
1431    /// a rig too tight to hold it (a failed capture is silent — resume just isn't available).
1432    pub(crate) turn_ckpt: Option<SpecCheckpoint>,
1433    /// Session-lifetime acceptance telemetry. Relaxed atomics update at the host-side round
1434    /// accounting the loop already does — no syncs, no allocation. NOTE a
1435    /// pool-resumed session carries the PREVIOUS requests' counts; per-request consumers
1436    /// diff with [`SpecTelemetry::delta_since`] around each burst.
1437    telem: SpecTelemetryCounters,
1438    /// PREFIX-CACHE publication request (lane/spec-prefix-cache): worker sets this to the
1439    /// miss-LCP boundary before a cold burst; the prime captures at exactly that split (it must
1440    /// coincide with the burst's `prime_split` or no capture happens). One-shot: consumed by the
1441    /// prime, result lands in `boundary_captures`.
1442    pub capture_at: Option<usize>,
1443    /// The captures the last prime produced (see [`SpecBoundaryCapture`]). Worker drains them
1444    /// post-burst to assemble prefix entries. A failed capture is silent, like `turn_ckpt` —
1445    /// publication just isn't available for that request. Plural since
1446    /// lane/frspec-multiturn-cache (2026-08-21): a cold burst can capture BOTH the miss-LCP
1447    /// split (the shared-prefix class) and the stable pre-generation boundary (the
1448    /// next-turn re-render class) — one entry per stop, exactly the boundary set the plain
1449    /// prefill tick publishes/checkpoints.
1450    pub boundary_captures: Vec<SpecBoundaryCapture>,
1451    /// STABLE-BOUNDARY TURN CHECKPOINT REQUEST (lane/frspec-multiturn-cache, 2026-08-21): the
1452    /// ABSOLUTE committed-length position the next non-empty prime should capture `turn_ckpt`
1453    /// at, instead of prompt-end. The worker sets it to the STABLE PRE-GENERATION boundary
1454    /// (`plain_checkpoint_boundary` — before the live generation header the client rewrites),
1455    /// porting the 2026-08-09 plain-tier fix: a prompt-end spec checkpoint includes the
1456    /// template's live assistant-generation header (`<|im_start|>assistant\n<think>\n`), which
1457    /// the NEXT turn's re-render replaces, so `affinity_match` diverged a couple tokens below
1458    /// the checkpoint and the spec pool declined 100% of multi-turn agent traffic (measured:
1459    /// `spec-affinity: declined (history diverged at 6811 of checkpoint 6813)`,
1460    /// research/multiturn-cache-20260821 B4). One-shot, `capture_at` convention; None = legacy
1461    /// prompt-end capture.
1462    pub ckpt_at: Option<usize>,
1463    /// FAIL-SAFE (lane/step37-vram-admission-20260830, external-review corroboration): set
1464    /// by the worker on a session serving a step-OOM park REPLAY. The burst entry pre-marks
1465    /// the draft-graph fallback so the replay never re-enters the capture path — the capture
1466    /// appetite is part of what drove the card to the OOM, and a replay that recaptures
1467    /// re-runs the incident. If the eager replay still cannot fit, the bounded retry budget
1468    /// exhausts into the honest recoverable Overloaded error instead of looping.
1469    pub capture_disabled: bool,
1470}
1471impl SpecSession {
1472    /// Context capacity of the session's caches (the server's ContextFull guard).
1473    pub fn cache_max_ctx(&self) -> usize {
1474        self.cache.max_ctx
1475    }
1476    /// Read access to the live trunk cache (lane/spec-prefix-cache): the worker slices
1477    /// full-attn KV rows `[0..capture.pos)` out of it when publishing a boundary capture —
1478    /// those rows are append-only for the session's lifetime (rollbacks never truncate below
1479    /// the prime boundary), so no copy was taken at prime time.
1480    pub fn cache_ref(&self) -> &Cache {
1481        &self.cache
1482    }
1483    /// Read access to the persistent draft-scratch plane (lane/spec-on-cache-hit): the
1484    /// worker slices rows `[0..capture.pos)` when publishing a boundary capture, exactly
1485    /// like the trunk KV — draft rows below the prompt end are append-only for the
1486    /// session's lifetime (the prime fill wrote them once; rollbacks reset `len_d` to the
1487    /// committed length, never below the prime boundary, and the true-hidden refresh
1488    /// rewrites generated positions only). Returns `(k, v, k_tok_bytes, v_tok_bytes)`.
1489    /// None when the scratch is ring-backed (Step35 SWA — physical rows are not
1490    /// prefix-addressable; the prefix cache already refuses that class end to end).
1491    pub fn draft_plane_ref(&self) -> Option<(&CudaSlice<u8>, &CudaSlice<u8>, usize, usize)> {
1492        if self.scratch.kv.ring.is_some() {
1493            return None;
1494        }
1495        Some((
1496            &self.scratch.kv.k,
1497            &self.scratch.kv.v,
1498            self.scratch.kv.k_tok_bytes,
1499            self.scratch.kv.v_tok_bytes,
1500        ))
1501    }
1502    /// Snapshot the session's process-local acceptance counters for per-burst diffing.
1503    pub fn telemetry(&self) -> SpecTelemetry {
1504        self.telem.snapshot()
1505    }
1506    /// Committed position this session can REWIND to (its retained prompt-end boundary), if any.
1507    /// A request whose prompt matches `committed[..pos]` exactly can resume from here — see
1508    /// `spec_rewind_to_checkpoint`.
1509    pub fn rewind_pos(&self) -> Option<usize> {
1510        self.turn_ckpt.as_ref().map(|c| c.pos)
1511    }
1512    /// Whether every ring-backed trunk/draft row needed by the retained checkpoint is resident.
1513    pub fn rewind_is_resident(&self) -> bool {
1514        self.turn_ckpt.as_ref().is_some_and(|ckpt| {
1515            self.cache.can_rollback(&ckpt.snap, 0) && self.scratch.can_rewind_to(ckpt.pos)
1516        })
1517    }
1518    /// Is this session in the DEMOTION-READY shape (see [`SpecSession::into_demoted`])?
1519    /// `false` means a carried pending must be flushed first (`spec_flush_pending`), or the
1520    /// session has never run a turn and has no prediction to hand over.
1521    pub fn demote_ready(&self) -> bool {
1522        self.pending_tok.is_none() && self.next_pred.is_some()
1523    }
1524    /// Does this session hold a carried pending bonus (flush required before a handoff/park)?
1525    pub fn has_pending(&self) -> bool {
1526        self.pending_tok.is_some()
1527    }
1528    /// Committed row count == cache rows (the session invariant), for the caller's own
1529    /// `fed`-length cross-check at a handoff boundary.
1530    pub fn committed_len(&self) -> usize {
1531        self.committed.len()
1532    }
1533    /// DEMOTION HANDOFF (lane/spec-gate, 2026-08-07): consume this session and hand its trunk
1534    /// cache + next-token prediction to the plain batched-decode path.
1535    ///
1536    /// WHY THIS IS EXACT (greedy). The invariant at a burst boundary is `cache.pos ==
1537    /// committed.len()`: every committed row has trunk KV + recurrent state, exactly as a plain
1538    /// tokenwise prime of the same `committed` sequence would have left it (that is the
1539    /// session-tail contract, and the same property `spec_rewind_to_checkpoint` and the reuse
1540    /// pool already rely on). `next_pred` is the argmax of the verify's logits for the LAST
1541    /// committed row — and verify-column logits are bit-identical to plain decode's logits at
1542    /// that position, because `matmul_decode_exact` bit-identity IS the basis of the greedy
1543    /// accept walk. So handing (cache, next_pred) to the batched path continues the stream from
1544    /// a state indistinguishable from one the batched path produced itself: the batched tick
1545    /// emits `next_pred`, feeds it into this same cache, and decodes on.
1546    ///
1547    /// `None` when the session is not in the handoff shape — a carried pending (its bonus row is
1548    /// NOT in the cache, so `spec_flush_pending` must commit it first) or no `next_pred` yet
1549    /// (never bursted). Callers must not force it: a half-committed cache handed to the batched
1550    /// path would silently skip a token.
1551    ///
1552    /// The MTP draft scratch, the persistent draft-graph context and the turn checkpoint are
1553    /// DROPPED here (freeing their VRAM): the batched path never drafts, and this handoff is
1554    /// one-way by design — there is no cheap symmetric re-promotion (rebuilding the draft KV
1555    /// would mean an `mtp_kv_fill` over the whole committed history).
1556    pub fn into_demoted(self) -> Option<(Cache, u32)> {
1557        if self.pending_tok.is_some() || self.cache.tainted {
1558            return None;
1559        }
1560        let np = self.next_pred?;
1561        debug_assert_eq!(
1562            self.cache.pos,
1563            self.committed.len(),
1564            "demotion handoff: cache rows != committed tokens"
1565        );
1566        Some((self.cache, np))
1567    }
1568    /// Pool-resume hook (audit Q2): clear the parked draft-graph failure memoization so a
1569    /// NEW request resuming this session gets one fresh capture chance — a transient-pressure
1570    /// capture failure must not persist for the pool's whole lifetime (the TRT #16072 class).
1571    /// Logs once iff a flag was actually set; a no-fallback resume is silent and free.
1572    pub fn reset_graph_fallback_on_resume(&mut self) {
1573        if let Some(line) = self
1574            .draft_ctx
1575            .as_mut()
1576            .and_then(|c| c.failed.reset_on_resume())
1577        {
1578            eprintln!("{line}");
1579        }
1580    }
1581}
1582
1583/// A session's PROMPT-END boundary state, the rewind target for session-affinity resume.
1584///
1585/// WHY THIS BOUNDARY, AND WHY IT IS THE ONLY ONE WORTH KEEPING. The rewrite class this lane
1586/// exists for (a client that strips `<think>` blocks out of prior assistant turns) mutates the
1587/// text the session GENERATED, never the prompt it was given. So turn N's prompt agrees with
1588/// turn N-1's committed tokens up to almost exactly where turn N-1's generation began — the
1589/// prompt-end boundary. Keeping a checkpoint there means the next turn re-primes only its own
1590/// delta (the rewritten answer + the new user turn) instead of the whole conversation.
1591///
1592/// WHAT IT MUST HOLD. Full-attn KV is append-only and position-addressed, so rewinding it is a
1593/// `len` truncation (no data). Linear-attn (GDN) conv/ssm state is mutated IN PLACE with no
1594/// position index, so it must be a real device COPY — that copy is the entire reason a spec
1595/// session could not previously rewind. The MTP draft scratch needs no copy either: its rows
1596/// below the boundary were written by this turn's fill and are never revisited (the per-round
1597/// true-hidden refresh only rewrites the CURRENT burst's committed positions), so rewinding it
1598/// is also just a `len` reset. `last_h` is the hidden of the last row below the boundary — the
1599/// predecessor-pairing anchor the next prime's fill reads for its first row.
1600///
1601/// COST: one `Cache::snapshot` per TURN, on a code path that already takes one per ROUND.
1602pub(crate) struct SpecCheckpoint {
1603    snap: crate::cache::CacheSnapshot,
1604    /// Committed length at the boundary (== cache.pos there, the session invariant).
1605    pos: usize,
1606    /// Pre-output_norm hidden of row `pos - 1`.
1607    last_h: CudaSlice<f32>,
1608}
1609
1610/// PREFIX-CACHE BOUNDARY CAPTURE (lane/spec-prefix-cache, 2026-08-14): the state a spec session
1611/// records at its cold-prime split so the WORKER can publish a cross-request prefix entry —
1612/// the commit-gated-publication port (research/cache-spec-design-20260814/PORT-PLAN.md item 1).
1613/// Only the pieces that are DESTROYED by continuing the prime need copies here: the in-place
1614/// GDN conv/ssm states (via `Cache::snapshot`, same mechanism as [`SpecCheckpoint`]) and the
1615/// boundary logits. Full-attn KV rows `[0..pos)` and draft-scratch rows `[0..pos)` are
1616/// append-only for the session's lifetime (rollbacks never truncate below the prime boundary),
1617/// so the worker slices those from the live caches post-burst instead of copying at prime time.
1618pub struct SpecBoundaryCapture {
1619    pub snap: crate::cache::CacheSnapshot,
1620    /// Token boundary (== cache.pos at capture; == the worker's miss-LCP split).
1621    pub pos: usize,
1622    /// Full-vocab logits after the prefix prime — the entry's boundary logits.
1623    pub logits: Vec<f32>,
1624    /// Pre-output_norm trunk hidden of row `pos - 1` (lane/spec-on-cache-hit): the
1625    /// predecessor-pairing anchor a RESTORED spec session's first suffix-fill row reads
1626    /// (the `SpecSession::last_h` convention). Empty = unavailable (capture stays valid;
1627    /// the fill's zeros row-0 fallback covers it at a bounded acceptance cost).
1628    pub last_h: Vec<f32>,
1629    /// Per-layer latent boundary tails (lane/glm5-prefix-latent2, 2026-09-01): the
1630    /// generation-destroyed slice of each MLA/DSA layer's boundary state, captured eagerly
1631    /// so the worker's DEFERRED publication can slice the append-only planes from the live
1632    /// cache (`LatentKvLayer::snapshot_plane_at`). EMPTY on every two-plane model — the
1633    /// pre-field captures are byte-identical; a latent-bearing cache with an EMPTY vec here
1634    /// keeps the publisher's loud refusal (the fail-closed door stays shut).
1635    pub latent_tails: Vec<Option<crate::cache::LatentTailCapture>>,
1636}
1637
1638/// D2H one hidden row out of a `[T, n_embd]` prime hidden stack — the boundary anchor a
1639/// spec boundary capture carries for later restored-session fills. Failure is silent
1640/// (`turn_ckpt` convention): the capture publishes without an anchor.
1641pub(crate) fn capture_boundary_hidden(
1642    e: &Engine,
1643    h_rows: &CudaSlice<f32>,
1644    pos: usize,
1645    n_embd: usize,
1646) -> Vec<f32> {
1647    if pos == 0 || h_rows.len() < pos * n_embd {
1648        return Vec::new();
1649    }
1650    let Ok(mut row) = e.uninit(n_embd) else {
1651        return Vec::new();
1652    };
1653    if e.copy_view_into(
1654        &mut row,
1655        0,
1656        &h_rows.slice((pos - 1) * n_embd..pos * n_embd),
1657        n_embd,
1658    )
1659    .is_err()
1660    {
1661        return Vec::new();
1662    }
1663    e.dtoh(&row).unwrap_or_default()
1664}
1665
1666/// ROLLBACK DOOR for sampled BOUNDARY tokens (lane/sampled-spec-quality, 2026-08-19).
1667/// Default ON: the token a burst emits at its own boundary is drawn from the request's
1668/// sampler. `MEMRA_SPEC_SAMPLED_BOUNDARY=0` restores the pre-lane posture (an ARGMAX at
1669/// every boundary) without touching greedy, which is byte-unaffected either way.
1670pub fn spec_sampled_boundary_on() -> bool {
1671    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1672    *ON.get_or_init(|| std::env::var("MEMRA_SPEC_SAMPLED_BOUNDARY").as_deref() != Ok("0"))
1673}
1674
1675/// ROLLBACK DOOR for SESSION-SPANNING penalty history (lane/sampled-spec-quality).
1676/// Default ON: `pen_hist` is seeded from the session's committed tail, so repetition /
1677/// frequency / presence penalties see the whole stream. `MEMRA_SPEC_PEN_SESSION=0`
1678/// restores the pre-lane posture (each burst restarts the window from its own prompt
1679/// slice, i.e. from NOTHING on a continuation burst) — and with the door shut the worker
1680/// must keep refusing penalized sampled prefix-cache restores, because the restored
1681/// session's continuation burst is handed no prompt slice at all.
1682pub fn spec_pen_session_on() -> bool {
1683    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1684    *ON.get_or_init(|| std::env::var("MEMRA_SPEC_PEN_SESSION").as_deref() != Ok("0"))
1685}
1686
1687/// ROLLBACK DOOR for extended-entry publication from a RESTORED session
1688/// (lane/sampled-spec-quality, Item 3). Default ON: a converted prefix-cache hit that fed a
1689/// suffix captures its own prompt-end boundary so the NEXT turn can hit a longer prefix.
1690/// `MEMRA_SPEC_RESTORE_REPUBLISH=0` restores the pre-lane posture (a namespace learns exactly
1691/// one boundary and never advances it). Whole-entry semantics only — the boundary is the
1692/// restored session's own prompt end, so `entry_pos != fed_len` still refuses on the way in.
1693pub fn spec_restore_republish_on() -> bool {
1694    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1695    *ON.get_or_init(|| std::env::var("MEMRA_SPEC_RESTORE_REPUBLISH").as_deref() != Ok("0"))
1696}
1697
1698/// Diagnostics: name every boundary token on stderr (`MEMRA_SPEC_BOUNDARY_TRACE=1`), with
1699/// the argmax the pre-lane code would have emitted from the same row. This is how the
1700/// lane MEASURES the boundary rate and the deviation rate instead of estimating them.
1701fn spec_boundary_trace() -> bool {
1702    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1703    *ON.get_or_init(|| std::env::var("MEMRA_SPEC_BOUNDARY_TRACE").as_deref() == Ok("1"))
1704}
1705
1706/// llama-parity floor for the penalty window when the request does not ask for a bigger
1707/// one (`repeat_last_n` default). The serve API arms `penalty_last_n = PEN_WINDOW_MAX` for any
1708/// non-identity penalty, so this floor only matters to explicit small windows and to the
1709/// CLI env path.
1710const PEN_WINDOW_FLOOR: usize = 64;
1711
1712/// CEILING on the penalty window, and it is a COST bound, not a semantic preference.
1713/// `penalize_logits_f32` (cu/spec_sample.cu) dedups on device by having thread `i` scan
1714/// `hist[0..i]`, so a pass is O(n_hist²) and it runs ~3x per verify round (the q rows, the
1715/// p column, the bonus column). The serve API uses this same bound for every non-identity
1716/// penalty so host/plain, sparse-device, and speculative sampling cannot change logits on
1717/// admission demotion. An uncapped 128k-token history would put ~1.7e10
1718/// comparisons per pass, tens of ms per round, i.e. penalties would silently destroy decode
1719/// throughput on exactly the long-context requests that most want them. 8192 keeps a pass
1720/// at ~7e7 comparisons (tens of microseconds) while still being **128x wider than the
1721/// pre-lane effective window** (64 prompt-tail tokens + whatever the current burst had
1722/// generated). A request that genuinely needs a window beyond this wants host-side dedup +
1723/// counts through a new kernel signature — a follow-up lane, named here rather than hidden.
1724/// `pub` since lane/dspark-penalized-sampled-20260821: the dspark route's accept walk and
1725/// the dspark_sample_gate binary trim their uploads with the SAME cap — a second constant
1726/// is a second thing to drift.
1727pub const PEN_WINDOW_MAX: usize = 8192;
1728
1729/// Seed a penalty window over the SESSION, not the burst (lane/sampled-spec-quality,
1730/// Item 2). The window is the last `max(penalty_last_n, 64)` tokens of
1731/// `session_committed ++ burst_prompt` — for a cold turn-1 burst (`session_committed`
1732/// empty, default `penalty_last_n`) that is byte-identically the pre-lane
1733/// `prompt.iter().rev().take(64).rev()`; for a continuation burst it is the stream the
1734/// client actually asked us to penalize, where the pre-lane code had NOTHING.
1735/// `pub` since lane/dspark-penalized-sampled-20260821: the dspark route seeds its session
1736/// window through the SAME function (one definition of "the window" across both spec
1737/// routes and the gate binary's trunk-only reference arm).
1738pub fn pen_window_seed(
1739    session_committed: &[u32],
1740    burst_prompt: &[u32],
1741    penalty_last_n: usize,
1742) -> Vec<u32> {
1743    let win = penalty_last_n.clamp(PEN_WINDOW_FLOOR, PEN_WINDOW_MAX);
1744    let take_prompt = burst_prompt.len().min(win);
1745    let take_sess = (win - take_prompt).min(session_committed.len());
1746    let mut hist = Vec::with_capacity(take_sess + take_prompt);
1747    hist.extend_from_slice(&session_committed[session_committed.len() - take_sess..]);
1748    hist.extend_from_slice(&burst_prompt[burst_prompt.len() - take_prompt..]);
1749    hist
1750}
1751
1752/// Draw a BOUNDARY token from the target distribution the request asked for
1753/// (lane/sampled-spec-quality, Item 1) — the fix for "sampled spec emits an ARGMAX token at
1754/// every burst boundary".
1755///
1756/// WHY THIS EXISTS. A spec burst's first emitted token is not produced by the accept walk:
1757/// it comes off a logits row that already exists (the prime's last row on a cold burst; the
1758/// row after the last committed token on a continuation burst; the prefix-cache entry's
1759/// boundary row on a restored one). Pre-lane that token was `argmax` in BOTH sampling
1760/// regimes, so a sampled stream took a greedy token once per burst — measured, not
1761/// estimated, in research/spec-cache-20260818/SAMPLED-QUALITY.md. At temperature > 0 the
1762/// customer asked for a sampled token, so this draws one.
1763///
1764/// THE PROGRAM IS THE FULL-ACCEPT BONUS'S PROGRAM, deliberately: penalize the row (over the
1765/// session's window), take this row's OWN filter stats (the sampfix-20260805 law — stats
1766/// from a neighbour row mis-scale every `e0` and can wipe the row to token 0), gumbel-perturb
1767/// with the session's Philox stream at `*sctr`, argmax the perturbed row. Reusing the bonus's
1768/// composition means `sample_check`'s distributional oracle covers this draw too, and the
1769/// boundary token is drawn from the same filtered/penalized `p` the accept walk targets.
1770///
1771/// THE STREAM IS THE SESSION'S, NOT A FRESH ONE. `sctr` is the caller's live counter and is
1772/// advanced by exactly one, so a boundary draw consumes the next value in the same Philox
1773/// stream the accept walk uses — never a second, independently seeded stream (which would be
1774/// a new distributional bug: two streams from one seed correlate wherever their counters
1775/// collide). That also makes a restored session's boundary draw at `sctr == 0` bit-identical
1776/// to the cold session's own first draw from the same logits row, which is what preserves the
1777/// sampled-hit lane's per-seed hit==cold byte identity.
1778#[allow(clippy::too_many_arguments)]
1779pub fn sample_boundary_token_dev(
1780    e: &Engine,
1781    logits: &CudaSlice<f32>,
1782    n_vocab: usize,
1783    sp: &SpecSampling,
1784    pen_hist: &[u32],
1785    sctr: &mut u32,
1786    site: &str,
1787) -> Result<u32, Box<dyn std::error::Error>> {
1788    debug_assert!(
1789        sp.temp > 0.0,
1790        "boundary sampling is the sampled regime only"
1791    );
1792    // Own copy: penalize_logits mutates in place and the caller's row is live state
1793    // (prime_logits back the constrained recompute; last_col_logits backs round 0's accept).
1794    let mut col = e.zeros(n_vocab)?;
1795    e.copy_into(&mut col, 0, logits, n_vocab)?;
1796    let pen_on = sp.penalty_last_n > 0
1797        && (sp.penalty_repeat != 1.0 || sp.penalty_freq != 0.0 || sp.penalty_present != 0.0);
1798    if pen_on && !pen_hist.is_empty() {
1799        // window trim mirrors the round loop's own upload (`pen_hist[w0..]`), cap included.
1800        let w0 = pen_hist
1801            .len()
1802            .saturating_sub(sp.penalty_last_n.min(PEN_WINDOW_MAX));
1803        let hist = &pen_hist[w0..];
1804        let hd = e.htod_u32_v(hist)?;
1805        e.penalize_logits(
1806            &mut col,
1807            &hd,
1808            hist.len(),
1809            sp.penalty_repeat,
1810            sp.penalty_freq,
1811            sp.penalty_present,
1812            n_vocab,
1813        )?;
1814    }
1815    let rows0 = e.htod_i32(&[0])?;
1816    let (mut th_d, mut z_d, mut mx_d) = (e.zeros(1)?, e.zeros(1)?, e.zeros(1)?);
1817    e.filter_stats(
1818        &col, n_vocab, &rows0, &mut th_d, &mut z_d, &mut mx_d, n_vocab, 1, sp.temp, sp.top_k,
1819        sp.top_p, sp.min_p,
1820    )?;
1821    let (th, mx) = (e.dtoh(&th_d)?[0], e.dtoh(&mx_d)?[0]);
1822    let mut perturb = e.zeros(n_vocab)?;
1823    e.gumbel_perturb_filtered(&col, &mut perturb, n_vocab, sp.seed, *sctr, sp.temp, mx, th)?;
1824    *sctr = sctr.wrapping_add(1);
1825    let td = e.argmax_token_device(&perturb, n_vocab)?;
1826    let tok = guard_vocab_token(
1827        e.dtoh_u32_one(&td)?,
1828        n_vocab,
1829        &format!("sampled boundary token (site={site})"),
1830    )?;
1831    if spec_boundary_trace() {
1832        // the pre-lane token, from the SAME row, so the deviation rate is measurable.
1833        let raw = e.argmax_token_device(logits, n_vocab)?;
1834        let greedy = e.dtoh_u32_one(&raw)?;
1835        eprintln!(
1836            "[spec-boundary] site={site} sampled={tok} argmax={greedy} \
1837             deviates={} temp={} sctr={}",
1838            (tok != greedy) as u8,
1839            sp.temp,
1840            sctr.wrapping_sub(1),
1841        );
1842    }
1843    Ok(tok)
1844}
1845
1846/// Host-row twin of [`sample_boundary_token_dev`] (the prime / feed / entry rows arrive as
1847/// host `Vec<f32>`).
1848#[allow(clippy::too_many_arguments)]
1849pub fn sample_boundary_token(
1850    e: &Engine,
1851    logits: &[f32],
1852    sp: &SpecSampling,
1853    pen_hist: &[u32],
1854    sctr: &mut u32,
1855    site: &str,
1856) -> Result<u32, Box<dyn std::error::Error>> {
1857    let n_vocab = logits.len();
1858    let d = e.htod(logits)?;
1859    sample_boundary_token_dev(e, &d, n_vocab, sp, pen_hist, sctr, site)
1860}
1861
1862#[derive(Clone, Copy, PartialEq, Eq, Debug)]
1863pub(crate) struct SampledGraphKey {
1864    seed: u64,
1865    temp_bits: u32,
1866    k: usize,
1867    top_k: i32,
1868    top_p_bits: u32,
1869    min_p_bits: u32,
1870    pen_on: bool,
1871}
1872
1873impl SampledGraphKey {
1874    pub(crate) fn new(
1875        seed: u64,
1876        temp: f32,
1877        k: usize,
1878        top_k: i32,
1879        top_p: f32,
1880        min_p: f32,
1881        pen_on: bool,
1882    ) -> Self {
1883        SampledGraphKey {
1884            seed,
1885            temp_bits: temp.to_bits(),
1886            k,
1887            top_k,
1888            top_p_bits: top_p.to_bits(),
1889            min_p_bits: min_p.to_bits(),
1890            pen_on,
1891        }
1892    }
1893
1894    /// The one regime the PURE-TEMP in-graph sampled chain may stand in for the eager one:
1895    /// nothing but temperature shapes `q`. Computed FROM THE KEY so the capture guard, the
1896    /// launch guard and the key can never drift apart (they were three separate expressions
1897    /// before this lane, and the launch site simply forgot to ask).
1898    pub(crate) fn pure_temp(&self) -> bool {
1899        self.top_k == 0
1900            && f32::from_bits(self.top_p_bits) >= 1.0
1901            && f32::from_bits(self.min_p_bits) <= 0.0
1902            && !self.pen_on
1903    }
1904
1905    /// Truncation filters active — the capture body needs the IN-GRAPH filter nodes
1906    /// (`filter_stats` + `gumbel_perturb_filtered_ctr`) so the draft draws from the same
1907    /// filtered distribution the accept test reconstructs. Meaningful only when
1908    /// `graph_capturable`; penalties never reach a capture body.
1909    pub(crate) fn filtered(&self) -> bool {
1910        !self.pure_temp()
1911    }
1912
1913    /// May the sampled draft graph be CAPTURED (and a parked one LAUNCHED) for this regime?
1914    /// Pure-temp always; filtered regimes when the filtered-capture door is on
1915    /// (lane/step37-draft-graph-serving-20260830); penalties never — the per-round history
1916    /// cannot be baked into a graph, and composing a raw-softmax (or stale-history) draw
1917    /// with a penalized accept test is the unconditional-accept exactness bug. Computed FROM
1918    /// THE KEY for the same no-drift reason as `pure_temp`.
1919    pub(crate) fn graph_capturable(&self) -> bool {
1920        !self.pen_on && (self.pure_temp() || spec_graph_filtered_on())
1921    }
1922}
1923
1924/// Per-head captured graphs for the MULTI-HEAD MTP draft chain (step-modulo prefix-replay,
1925/// lane/step37-draft-graph-serving-20260830). The chain POLICY — which head serves step j,
1926/// how long the replayed prefix is, which stored seed feeds row r — stays HOST-SIDE in the
1927/// launch loop, exactly `mtp_chain_forward_dev`'s order; the graphs capture ONE head-row
1928/// forward each, on the head's OWN scratch plane:
1929/// - `interior[i]`: head i, `with_head=false` — KV append + carrier only. Interior rows'
1930///   logits are dead in the eager chain too (`mtp_chain_forward_dev` keeps only the last
1931///   row), so skipping the head matmul changes no consumed byte and removes the eager
1932///   chain's per-replay-row full-vocab matmul.
1933/// - `last[i]`: head i, `with_head=true` + the mode's tail (greedy argmax, or the sampled
1934///   gumbel draw — filtered in-graph when the request carries filters).
1935///
1936/// One `DraftChainGraphs` per MODE (greedy vs sampled), owning its keeper: dropping the
1937/// sampled chain on an s_key change never invalidates the greedy one.
1938struct DraftChainGraphs {
1939    interior: Vec<cudarc::driver::CudaGraph>,
1940    last: Vec<cudarc::driver::CudaGraph>,
1941    /// Never read: exists to OWN the captured graphs' backing buffers for as long as the
1942    /// graphs replay (the capture-retain law; same class as `DsparkSegGraph::_keeper`).
1943    _keeper: Vec<Box<dyn std::any::Any + Send>>,
1944}
1945
1946/// Sampled-tail capture pack for `mtp_head_forward_cap`: the persistent buffers and baked
1947/// constants of the in-graph categorical draw. `filt: None` = the PURE-TEMP body (gumbel
1948/// over the raw softmax), byte-identical to the pre-lane capture; `Some` adds the in-graph
1949/// truncation filter (`filter_stats` + `gumbel_perturb_filtered_ctr`) so the draft draws
1950/// from the same filtered distribution the accept test reconstructs
1951/// (lane/step37-draft-graph-serving-20260830).
1952struct SampledCapArgs<'a> {
1953    ctr: &'a mut CudaSlice<u32>,
1954    perturb: &'a mut CudaSlice<f32>,
1955    q_out: &'a mut CudaSlice<f32>,
1956    seed: u64,
1957    temp: f32,
1958    filt: Option<SampledCapFilter<'a>>,
1959}
1960
1961/// In-graph truncation-filter nodes: the stat slots `filter_stats` fills and the perturb
1962/// reads, plus the filter constants baked into the capture (they live in `s_key`, so a
1963/// request whose filters differ drops the parked graph before this ever goes stale).
1964struct SampledCapFilter<'a> {
1965    rows0: &'a CudaSlice<i32>,
1966    th: &'a mut CudaSlice<f32>,
1967    z: &'a mut CudaSlice<f32>,
1968    mx: &'a mut CudaSlice<f32>,
1969    top_k: i32,
1970    top_p: f32,
1971    min_p: f32,
1972}
1973
1974pub(crate) struct DraftGraphCtx {
1975    g_tok: CudaSlice<u32>,
1976    g_pos: CudaSlice<i32>,
1977    g_seed: CudaSlice<f32>,
1978    g_p: CudaSlice<f32>,
1979    g_ctr: CudaSlice<u32>,
1980    g_q: CudaSlice<f32>,
1981    g_perturb: CudaSlice<f32>,
1982    /// IN-GRAPH filter-stat slots (filtered sampled capture): `filter_stats` writes
1983    /// (th, z, mx) here inside the graph; `gumbel_perturb_filtered_ctr` reads (mx, th) from
1984    /// the same slots. Persistent so the baked pointers survive replays. `g_rows0` is the
1985    /// constant row-index-0 the single-row `filter_stats` launch reads (a captured memcpy
1986    /// source must not be a host temporary).
1987    g_rows0: CudaSlice<i32>,
1988    g_th: CudaSlice<f32>,
1989    g_z: CudaSlice<f32>,
1990    g_mx: CudaSlice<f32>,
1991    q_slots: Vec<CudaSlice<f32>>,
1992    /// DRAFT-SIDE GRAMMAR MASK (lane/draft-mask): packed allowed-set words over the DRAFT
1993    /// head's vocab, at a STABLE address so the captured draft graph's mask node reads the
1994    /// per-position contents the host re-uploads before each replay (the graph-promote
1995    /// pattern from decode.rs). Empty unless the session drafts under a grammar.
1996    g_dmask: CudaSlice<u32>,
1997    /// was `graph` captured WITH the mask node? A parked graph of the wrong shape is dropped.
1998    /// Covers the multi-head `chain` too (single-head and chain are mutually exclusive for a
1999    /// given model, so one flag serves whichever is active).
2000    graph_masked: bool,
2001    graph: Option<cudarc::driver::CudaGraph>,
2002    graph_s: Option<cudarc::driver::CudaGraph>,
2003    /// Multi-head chain graphs (see [`DraftChainGraphs`]): greedy and sampled chains, the
2004    /// chain twins of `graph` / `graph_s`. `chain_s`'s capture identity is `s_key` (shared
2005    /// with `graph_s` — a session is either single-head or chain, never both), and it obeys
2006    /// the same drop rules (key mismatch, penalty regime, mask-shape change).
2007    chain: Option<DraftChainGraphs>,
2008    chain_s: Option<DraftChainGraphs>,
2009    /// Failed-capture memoization for both graphs — LOUD on flip, cleared on pool resume
2010    /// (audit Q2, the TRT #16072 silent-permanent-coverage-loss class).
2011    failed: DraftGraphFallback,
2012    /// Capture identity of `graph_s` — see [`SampledGraphKey`]. `None` iff no sampled graph is
2013    /// parked; a request whose key differs drops the parked graph (and its q slots/keeper).
2014    s_key: Option<SampledGraphKey>,
2015    /// CAPTURE-RETAIN keepers (#68 root cause, 2026-08-04): the warmup-run transients whose
2016    /// pool addresses the captured graph(s) bake. Without these, the transients return to the
2017    /// pool at capture-body exit and later work (burst-boundary prime/fill/commit passes, or a
2018    /// co-served session in the worker) reuses those addresses — the persisted graph's replay
2019    /// then reads/writes live unrelated buffers (exactness corruption, first seen as the ST
2020    /// serve-spec 4B graph-arm corruption; one-shot CLI calls never re-shuffled the pool, which
2021    /// is why run-spec K=1..8 passed on the same checkpoint). Same fix class as
2022    /// capture_graph_retained's gemma/decode.rs sites — hold as long as the graph replays.
2023    keeper: Vec<Box<dyn std::any::Any + Send>>,
2024    keeper_s: Vec<Box<dyn std::any::Any + Send>>,
2025}
2026
2027/// Failed-capture memoization for the two draft graphs (audit Q2, 2026-08-05 — the
2028/// TRT #16072 trap class: pressure-triggered, silent, long-lived coverage loss).
2029///
2030/// Three contracts:
2031/// - LOUD FLIP: `mark_*` returns the warn line exactly on the false→true transition
2032///   (returned, not printed, so the once-per-flip contract is unit-testable); the caller
2033///   `eprintln!`s it UNCONDITIONALLY — a dropped draft graph is never silent. Re-marking
2034///   an already-failed graph returns None (the per-burst memoization that keeps the eager
2035///   fallback from paying a doomed capture attempt every burst).
2036/// - RESET ON RESUME: `reset_on_resume` clears both flags — a parked session resumed by a
2037///   NEW request gets one fresh capture chance instead of carrying a transient-pressure
2038///   failure for the pool's whole lifetime. Returns the note line only when a flag was
2039///   actually set (quiet on the common clean-resume path).
2040/// - Shape-change clears (`clear_*`) stay silent, exactly as before: they precede a fresh
2041///   capture attempt whose own failure would re-flip loudly.
2042#[derive(Default)]
2043pub(crate) struct DraftGraphFallback {
2044    greedy: bool,
2045    sampled: bool,
2046}
2047impl DraftGraphFallback {
2048    fn mark_greedy(&mut self, reason: &str) -> Option<String> {
2049        if self.greedy {
2050            return None;
2051        }
2052        self.greedy = true;
2053        Some(format!(
2054            "[spec] WARN: draft-graph capture failed ({reason}); eager fallback until session resume"
2055        ))
2056    }
2057    fn mark_sampled(&mut self, reason: &str) -> Option<String> {
2058        if self.sampled {
2059            return None;
2060        }
2061        self.sampled = true;
2062        Some(format!(
2063            "[spec] WARN: sampled draft-graph capture failed ({reason}); eager fallback until session resume"
2064        ))
2065    }
2066    fn greedy_failed(&self) -> bool {
2067        self.greedy
2068    }
2069    fn sampled_failed(&self) -> bool {
2070        self.sampled
2071    }
2072    fn clear_greedy(&mut self) {
2073        self.greedy = false;
2074    }
2075    fn clear_sampled(&mut self) {
2076        self.sampled = false;
2077    }
2078    /// Pool-resume reset: both graphs get a fresh capture chance. Some(note) iff any flag
2079    /// was set (so clean resumes stay quiet).
2080    pub(crate) fn reset_on_resume(&mut self) -> Option<String> {
2081        if !self.greedy && !self.sampled {
2082            return None;
2083        }
2084        let which = match (self.greedy, self.sampled) {
2085            (true, true) => "greedy+sampled",
2086            (true, false) => "greedy",
2087            _ => "sampled",
2088        };
2089        self.greedy = false;
2090        self.sampled = false;
2091        Some(format!(
2092            "[spec] draft-graph fallback reset on session resume ({which}); recapture eligible"
2093        ))
2094    }
2095}
2096
2097impl DraftGraphCtx {
2098    fn new(e: &Engine, n_embd: usize, qlen: usize) -> Result<Self, Box<dyn std::error::Error>> {
2099        Ok(DraftGraphCtx {
2100            g_tok: e.alloc_u32_zeroed(1)?,
2101            g_pos: e.htod_i32(&[0])?,
2102            g_seed: e.zeros(n_embd)?,
2103            g_p: e.zeros(1)?,
2104            g_ctr: e.alloc_u32_zeroed(1)?,
2105            g_q: e.zeros(qlen)?,
2106            g_perturb: e.zeros(qlen)?,
2107            g_rows0: e.htod_i32(&[0])?,
2108            g_th: e.zeros(1)?,
2109            g_z: e.zeros(1)?,
2110            g_mx: e.zeros(1)?,
2111            q_slots: Vec::new(),
2112            g_dmask: e.alloc_u32_zeroed(1)?,
2113            graph_masked: false,
2114            graph: None,
2115            graph_s: None,
2116            chain: None,
2117            chain_s: None,
2118            failed: DraftGraphFallback::default(),
2119            s_key: None,
2120            keeper: Vec::new(),
2121            keeper_s: Vec::new(),
2122        })
2123    }
2124}
2125
2126pub(crate) struct MtpScratch {
2127    kv: KvLayer,
2128    /// Logical row capacity. On the graph/DC draft path it also doubles as fa_decode_dc's
2129    /// bucket_max: n_splits is sized from it ONCE, so the graph captured at round 0 stays valid
2130    /// for every later t_kv. Step35 refuses that path and may back this logical extent with the
2131    /// smaller host-indexed SWA ring instead.
2132    cap: usize,
2133    extra: Vec<MtpScratchPlane>,
2134}
2135
2136struct MtpScratchPlane {
2137    kv: KvLayer,
2138    cap: usize,
2139}
2140
2141fn mtp_scratch_layout(
2142    cfg: &memra_gguf::config::ModelConfig,
2143    geom: Option<&crate::hybrid::DraftGeom>,
2144) -> (usize, usize, usize, usize) {
2145    // Student draft heads carry fewer KV heads (head_dim unchanged) -> smaller scratch rows.
2146    let n_head_kv = geom.map(|g| g.n_head_kv).unwrap_or(cfg.n_head_kv as usize);
2147    let head_dim_k = cfg.head_dim_k as usize;
2148    let head_dim_v = cfg.head_dim_v as usize;
2149    assert!(
2150        head_dim_k.is_multiple_of(32) && head_dim_v.is_multiple_of(32),
2151        "KVQUANT requires head_dim%32==0 (MTP scratch)"
2152    );
2153    let kv_dim_k = head_dim_k * n_head_kv;
2154    let kv_dim_v = head_dim_v * n_head_kv;
2155    // The fp8-KV arm deliberately does not reach the draft scratch; keep the exact format
2156    // policy shared with `MtpScratch::new` so admission scales the same allocation.
2157    let (kbb, vbb) = crate::kv_blk_bytes();
2158    let k_tok_bytes = (kv_dim_k / 32) * kbb;
2159    let v_tok_bytes = (kv_dim_v / 32) * vbb;
2160    (kv_dim_k, kv_dim_v, k_tok_bytes, v_tok_bytes)
2161}
2162
2163fn mtp_chain_head_index(step: usize, head_count: usize) -> usize {
2164    assert!(head_count > 0, "MTP chain requires at least one head");
2165    step % head_count
2166}
2167
2168impl MtpScratch {
2169    fn alloc_plane(
2170        e: &Engine,
2171        cfg: &memra_gguf::config::ModelConfig,
2172        plan: &memra_gguf::model_plan::ModelPlan,
2173        cap: usize,
2174        geom: Option<&crate::hybrid::DraftGeom>,
2175    ) -> Result<MtpScratchPlane, Box<dyn std::error::Error>> {
2176        let (kv_dim_k, kv_dim_v, k_tok_bytes, v_tok_bytes) = mtp_scratch_layout(cfg, geom);
2177        let ring = if crate::cache::swa_ring_on()
2178            && crate::plan_backend::decode_batch_program(plan)
2179                == crate::plan_backend::DecodeBatchProgram::SlidingGatedMoe
2180        {
2181            let window = plan
2182                .layers
2183                .iter()
2184                .find_map(|layer| match layer.attention {
2185                    memra_gguf::model_plan::AttentionPlan::SlidingWindow { window, .. } => {
2186                        Some(window as usize)
2187                    }
2188                    _ => None,
2189                })
2190                .ok_or("sliding-gated-MoE draft scratch has no sliding-window layer")?;
2191            Some(crate::cache::KvRing::new(
2192                crate::cache::swa_ring_rows(window, cap),
2193                window,
2194            ))
2195        } else {
2196            None
2197        };
2198        let alloc_rows = ring.as_ref().map(crate::cache::KvRing::rows).unwrap_or(cap);
2199        // Ring-backed planes arm the device base mirror for the dcw draft arm (see
2200        // KvLayer::base_d): the captured chain derives its physical rows from
2201        // (len_d, base_d, window) with zero per-token node updates.
2202        let base_d = match ring.as_ref() {
2203            Some(_) => Some(e.htod_i32(&[0])?),
2204            None => None,
2205        };
2206        Ok(MtpScratchPlane {
2207            kv: KvLayer {
2208                k: e.alloc_u8(alloc_rows * k_tok_bytes)?,
2209                v: e.alloc_u8(alloc_rows * v_tok_bytes)?,
2210                kv_dim_k,
2211                kv_dim_v,
2212                k_tok_bytes,
2213                v_tok_bytes,
2214                len: 0,
2215                ring,
2216                len_d: e.htod_i32(&[0])?,
2217                base_d,
2218            },
2219            cap,
2220        })
2221    }
2222
2223    fn new(
2224        e: &Engine,
2225        cfg: &memra_gguf::config::ModelConfig,
2226        plan: &memra_gguf::model_plan::ModelPlan,
2227        cap: usize,
2228        geom: Option<&crate::hybrid::DraftGeom>,
2229    ) -> Result<Self, Box<dyn std::error::Error>> {
2230        // Trunk KV block bytes (34/24, q8_0/q5_1). The scratch keeps the baseline numerics and
2231        // its append/fa pass g=false (the e4m3 arms drifted draft acceptance 69-88% -> 46% in the
2232        // 2026-07-12 A/B, and the trunk fp8 door itself was removed 2026-09-05).
2233        let primary = Self::alloc_plane(e, cfg, plan, cap, geom)?;
2234        Ok(MtpScratch {
2235            kv: primary.kv,
2236            cap: primary.cap,
2237            extra: Vec::new(),
2238        })
2239    }
2240
2241    fn push_plane(
2242        &mut self,
2243        e: &Engine,
2244        cfg: &memra_gguf::config::ModelConfig,
2245        plan: &memra_gguf::model_plan::ModelPlan,
2246        geom: Option<&crate::hybrid::DraftGeom>,
2247    ) -> Result<(), Box<dyn std::error::Error>> {
2248        self.extra
2249            .push(Self::alloc_plane(e, cfg, plan, self.cap, geom)?);
2250        Ok(())
2251    }
2252
2253    fn plane_count(&self) -> usize {
2254        1 + self.extra.len()
2255    }
2256
2257    fn plane(&self, index: usize) -> (&KvLayer, usize) {
2258        if index == 0 {
2259            (&self.kv, self.cap)
2260        } else {
2261            let plane = &self.extra[index - 1];
2262            (&plane.kv, plane.cap)
2263        }
2264    }
2265
2266    fn plane_mut(&mut self, index: usize) -> (&mut KvLayer, usize) {
2267        if index == 0 {
2268            (&mut self.kv, self.cap)
2269        } else {
2270            let plane = &mut self.extra[index - 1];
2271            (&mut plane.kv, plane.cap)
2272        }
2273    }
2274
2275    // #[track_caller]: set_len/set_plane_len have eight call sites (checkpoint restore, spec
2276    // rollback, session grow, seed replay ...) and the lap failure needs to say WHICH one, not
2277    // just that a rewind was refused.
2278    #[track_caller]
2279    fn set_plane_len(
2280        &mut self,
2281        e: &Engine,
2282        index: usize,
2283        n: usize,
2284    ) -> Result<(), Box<dyn std::error::Error>> {
2285        let caller = std::panic::Location::caller();
2286        let (kv, cap) = self.plane_mut(index);
2287        if let Some(ring) = kv.ring.as_ref()
2288            && !ring.can_rewind_to(n)
2289        {
2290            // NAME THE NUMBERS (2026-08-28). This error is a step37 serving blocker on the
2291            // vendor-default shape and it fires from more than one call path with more than
2292            // one trigger: a long generation walks the checkpoint out of the ring, but a
2293            // ~4.5k-token prompt also fails within 5 s of prime, which accumulation cannot
2294            // explain. A bare message forced two rounds of guessing; the operands make each
2295            // trigger name itself.
2296            let raw = n.saturating_sub(ring.window().saturating_sub(1));
2297            return Err(format!(
2298                    "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})",
2299                    ring.window(),
2300                    ring.base(),
2301                    ring.rows(),
2302                    raw & !31usize,
2303                )
2304                .into());
2305        }
2306        kv.len = n;
2307        e.set_i32_one(&mut kv.len_d, n as i32)
2308    }
2309
2310    /// Set BOTH length counters: the host mirror AND the device len_d the captured append/fa read
2311    /// (a 4-byte in-place htod — the counter pointer is baked into the graph, never realloc'd).
2312    /// This is the ONLY truncation/rollback mechanism the persistent draft KV needs.
2313    #[track_caller]
2314    fn set_len(&mut self, e: &Engine, n: usize) -> Result<(), Box<dyn std::error::Error>> {
2315        let caller = std::panic::Location::caller();
2316        if !self.can_rewind_to(n) {
2317            // set_plane_len re-checks and reports the operands; call it so the failure carries
2318            // which plane refused and why, instead of this bare aggregate.
2319            for index in 0..self.plane_count() {
2320                self.set_plane_len(e, index, n)?;
2321            }
2322            return Err(format!(
2323                "SWA ring MTP checkpoint has been lapped; full re-prime required (aggregate rewind_to={n}, no single plane reported, called from {caller})"
2324            )
2325            .into());
2326        }
2327        for index in 0..self.plane_count() {
2328            self.set_plane_len(e, index, n)?;
2329        }
2330        Ok(())
2331    }
2332
2333    fn can_rewind_to(&self, n: usize) -> bool {
2334        (0..self.plane_count()).all(|index| {
2335            self.plane(index)
2336                .0
2337                .ring
2338                .as_ref()
2339                .is_none_or(|ring| ring.can_rewind_to(n))
2340        })
2341    }
2342
2343    /// Pre-arm ring headroom for `rows` upcoming DEVICE-COUNTER appends (the dcw draft arm):
2344    /// a captured chain cannot rebase mid-replay, so any rebase the coming appends could need
2345    /// happens HERE, host-side, before the capture warmups or the round's replays (the rebase
2346    /// arm of `prepare_kv_append` also refreshes the plane's `base_d` device mirror). No-op on
2347    /// flat planes and when the ring already has room; `len` is untouched either way.
2348    fn ensure_dcw_headroom(
2349        &mut self,
2350        e: &Engine,
2351        rows: usize,
2352    ) -> Result<(), Box<dyn std::error::Error>> {
2353        for index in 0..self.plane_count() {
2354            let (kv, _) = self.plane_mut(index);
2355            let Some(ring) = kv.ring.as_ref() else {
2356                continue;
2357            };
2358            let retain = memra_kv::swa_retain_from(kv.len, ring.window(), ring.base());
2359            e.prepare_kv_append(kv, retain, rows)?;
2360        }
2361        Ok(())
2362    }
2363}
2364
2365/// Retained verify intermediates for the REPLAY-FREE partial accept (2026-07-03, the profiled
2366/// #1 spec cost at long ctx: the partial-accept replay was a DUPLICATE trunk pass — ~0.54 extra
2367/// full weight reads per round — recomputing columns the verify had already produced
2368/// bit-identically). Holds, per linear layer, everything needed to rebuild its recurrent state
2369/// to "after the first j verify columns" WITHOUT re-running the trunk:
2370/// - BATCHED-path layers (`gdn`): the exact token-major inputs the round's ONE gdn_scan
2371///   consumed. A prefix re-run of the SAME kernel (t=j) from the snapshot state is bit-identical
2372///   to the first j iterations of the verify's scan — the kernel's t-loop carries state in
2373///   registers and iteration t never depends on T. `qkv_mixed` (the conv input) feeds the
2374///   pure-copy ring rebuild.
2375/// - PER-COLUMN-path layers (`cols`): dtod clones of (conv_state, ssm_state) taken after each
2376///   column 0..t-2 — pure copies of the actual chain states (the last column is never a rebuild
2377///   target: j <= t-1).
2378///   Full-attn layers need nothing: their verify KV rows are bit-identical to eager's (the
2379///   decode-exact contract; verify-probe pins it), so rollback = len truncation.
2380struct GdnStash {
2381    qkv_mixed: CudaSlice<f32>, // [t, conv_dim] token-major (conv input)
2382    q_l2: CudaSlice<f32>,
2383    k_l2: CudaSlice<f32>,
2384    v_g: CudaSlice<f32>, // [t, num_v, d_state]
2385    g_log: CudaSlice<f32>,
2386    beta: CudaSlice<f32>, // [t, num_v]
2387}
2388pub(crate) struct VerifyCkpt {
2389    gdn: Vec<Option<GdnStash>>, // [n_layer], Some iff batched linear path ran
2390    #[allow(clippy::type_complexity)]
2391    // allow: one-shot composite type; naming it would hide the shape that matters at the call site
2392    cols: Vec<Option<Vec<(CudaSlice<f32>, CudaSlice<f32>)>>>, // [n_layer][col] = (conv, ssm) after col
2393}
2394/// Opaque handle for the dspark round (dflash.rs) — VerifyCkpt stays spec-private.
2395pub(crate) struct DsparkVerifyCkpt(VerifyCkpt);
2396
2397/// Engine-bundle slice 3 (DSF-ROUNDCOST-20260820 §2 row 4 / §5 rank 1): bucketed CUDA
2398/// graphs for the dspark verify's LINEAR-layer segments. The measured verify is ~2,800
2399/// eager launches whose residual cost is DEVICE-side per-launch overhead (slice 2 proved
2400/// host dispatch is not the binder: fully-deferred dispatch bought ~0 wall). The 48 GDN
2401/// layers between full-attention layers are shape-static given vt — no positions, no
2402/// t_kv, state addressed through pointer tables — so runs of them capture per
2403/// (segment, vt) and replay as ONE graph launch each. Full-attention layers stay eager
2404/// (their per-row append/fa arm picks are t_kv-driven — the exec-update extension).
2405///
2406/// Per round out-of-graph: one pointer-table refresh (gdn ping-pong moves the canonical
2407/// handles), one input-staging copy per segment, host parity bookkeeping. Captured via
2408/// `capture_graph_retained` (2 warmups + capture, keeper retains warmup transients so
2409/// pool addresses stay stable); the warmups EXECUTE, so segment conv/ssm state is saved
2410/// before and restored after — the graph's first real launch starts from the exact
2411/// pre-round state. The ckpt column stash rides persistent slabs (written inside the
2412/// graph as memcpy nodes); commit reads them via `dspark_commit_prefix_slab`.
2413/// `MEMRA_DSPARK_VERIFY_GRAPH=0` reverts to the eager walk (byte-identical body).
2414pub(crate) struct DsparkVerifyGraphs {
2415    /// Linear-attention layer indices ascending; `lin_pos[il]` = index into the vecs.
2416    lin: Vec<usize>,
2417    lin_pos: std::collections::HashMap<usize, usize>,
2418    /// [n_lin x 6] pointer table (conv, s0, s1, conv, s1, s0 per layer), refreshed per
2419    /// verify from the live handles; layer il's slice starts at lin_pos[il]*6.
2420    table_all: CudaSlice<u64>,
2421    host_table: Vec<u64>,
2422    /// Persistent per-layer ckpt stash slabs: row r of the verify at slab offset
2423    /// r*words. Shared by every (segment, vt) bucket — one verify runs at a time.
2424    stash_conv: Vec<CudaSlice<f32>>,
2425    stash_ssm: Vec<CudaSlice<f32>>,
2426    conv_words: usize,
2427    ssm_words: usize,
2428    /// Per-vt input/output staging (stable addresses the graphs bake).
2429    stage: std::collections::HashMap<usize, (CudaSlice<f32>, CudaSlice<f32>)>,
2430    /// Per-vt dflash tap-sink buffers — the captured segments bake the tap dst address,
2431    /// so the sink buffer must live (and persist) with the graphs, not with the round.
2432    pub(crate) tap_bufs: std::collections::HashMap<usize, CudaSlice<f32>>,
2433    graphs: std::collections::HashMap<(usize, usize), DsparkSegGraph>,
2434    /// Warmup-corruption guard scratch: pre-capture conv/ssm of every linear layer
2435    /// (sized n_lin — the slice-4c full-verify warmups execute the whole walk).
2436    save_conv: CudaSlice<f32>,
2437    save_ssm: CudaSlice<f32>,
2438    max_run: usize,
2439    n_embd: usize,
2440    /// Set by the verify walk: this round's linear ckpt lives in the slabs (the caller
2441    /// commits through `dspark_commit_prefix_slab` instead of the cols arm).
2442    pub(crate) round_slab: bool,
2443    // ---- slice 4c: full-verify single graph per (vt, rung) ----
2444    /// Full-attention layer indices ascending; `fa_pos[il]` = index into the vec.
2445    fa: Vec<usize>,
2446    fa_pos: std::collections::HashMap<usize, usize>,
2447    /// [n_fa x 2 x t_cap] interleaved (k,v) base-pointer pairs, refreshed per verify;
2448    /// layer il's slice starts at `fa_pos[il] * 2 * t_cap` (the seqs twins read pairs
2449    /// [2z], z < t <= t_cap, so one t_cap-sized table serves every vt).
2450    fa_table: CudaSlice<u64>,
2451    fa_host_table: Vec<u64>,
2452    t_cap: usize,
2453    /// Per-vt position staging for the captured bodies — contents refreshed per round
2454    /// (rope reads row r; the seqs twins derive append slot and T_kv per z from it).
2455    pos_stage: std::collections::HashMap<usize, CudaSlice<i32>>,
2456    /// Full-verify graphs keyed (vt, rung_end, hi).
2457    full: std::collections::HashMap<(usize, usize, usize), DsparkSegGraph>,
2458    /// Largest n with every layer in [0, n) linear or full-attention (walk coverage).
2459    covered: usize,
2460    /// Every layer in [0, n) is linear or full-attention (no MLA/unknown mixers) — the
2461    /// full-verify capture walks all of them.
2462    walk_uniform: bool,
2463    /// Last `(captures, device graph-mem reserved bytes)` reading taken by
2464    /// `HybridModel::dspark_vg_admission_debt` — the two-point base of the MARGINAL debt
2465    /// projection (see `dspark_vg_debt_projection`; a mean-based reading extrapolated the
2466    /// pool's one-time shared allocation and reserved 8.5 GB of phantom VRAM).
2467    debt_obs: Option<(usize, usize)>,
2468}
2469
2470struct DsparkSegGraph {
2471    graph: cudarc::driver::CudaGraph,
2472    _keeper: Vec<Box<dyn std::any::Any + Send>>,
2473}
2474
2475/// Per-call arguments of [`HybridModel::qwen35_tparallel_fa_layer`] — one struct so the
2476/// eager walk and the slice-4c captured full-verify graphs hand the SAME body its two
2477/// modes without a second copy of the math.
2478pub(crate) struct FaLayerArgs<'a> {
2479    /// [T] per-row positions (device): rope reads them row-indexed; the seqs twins read
2480    /// them per-z (append slot = pos, T_kv = pos + 1).
2481    pub pos_d: &'a CudaSlice<i32>,
2482    /// Verify-level lazy per-row 1-element position buffers — only the per-row fallback
2483    /// arm builds/uses them (graph mode refuses that arm).
2484    pub pos_rows: &'a mut Option<Vec<CudaSlice<i32>>>,
2485    pub pos0: usize,
2486    /// Some((kv pointer table, offset-in-u64s, rung_end)) = captured-graph mode.
2487    pub graph_cap: Option<(&'a CudaSlice<u64>, usize, usize)>,
2488    /// ROUND-STREAM (lane/draftcost-moe, v0.100 train merge): Some((token stream, device
2489    /// round counter)) routes the FA attend through the dc rows kernels and the Linear
2490    /// mixer through `linear_attn_verify_t` (the stream arms the old inline body carried).
2491    /// Never armed together with `graph_cap` (the verify-level merge guard refuses).
2492    pub stream: Option<(&'a CudaSlice<u32>, &'a CudaSlice<i32>)>,
2493    /// VerifyCkpt for the stream-Linear arm's GdnStash install; None in graph mode and
2494    /// for FA layers that never touch it.
2495    pub ckpt: Option<&'a mut VerifyCkpt>,
2496}
2497
2498// SAFETY: `CudaGraph` is not marked Send by cudarc because its raw driver handles carry
2499// no automatic trait; CUDA driver graph handles are context-scoped rather than
2500// OS-thread-affine. The ctx lives in
2501// `HybridModel::dspark_vgraphs` behind a Mutex and every touch happens on the engine's
2502// single decode-stream thread.
2503unsafe impl Send for DsparkVerifyGraphs {}
2504
2505impl DsparkVerifyGraphs {
2506    /// Live capture count (segment + full graphs) — the denominator of
2507    /// [`dspark_vg_debt_projection`]'s observed bytes/capture mean.
2508    pub(crate) fn captures(&self) -> usize {
2509        self.graphs.len() + self.full.len()
2510    }
2511
2512    /// Take the marginal-growth debt reading and record this observation for the next one.
2513    /// Called under the pool mutex by `HybridModel::dspark_vg_admission_debt`.
2514    pub(crate) fn admission_debt(&mut self, reserved_bytes: usize) -> usize {
2515        let captures = self.captures();
2516        let debt =
2517            dspark_vg_debt_projection(captures, dspark_vg_cap(), reserved_bytes, self.debt_obs);
2518        if captures > 0 {
2519            match self.debt_obs {
2520                Some((c0, _)) if captures <= c0 => {}
2521                _ => self.debt_obs = Some((captures, reserved_bytes)),
2522            }
2523        }
2524        debt
2525    }
2526
2527    /// Build for this cache's shape. None when there are no linear layers, sizes are
2528    /// non-uniform, or the trunk keeps a gemma4 config (never on the qwen35 family).
2529    pub(crate) fn new(
2530        e: &Engine,
2531        cache: &Cache,
2532        t_max: usize,
2533        n_embd: usize,
2534    ) -> Result<Option<Self>, Box<dyn std::error::Error>> {
2535        let lin: Vec<usize> = (0..cache.recur.len())
2536            .filter(|&il| cache.recur[il].is_some())
2537            .collect();
2538        if lin.is_empty() || t_max < 2 {
2539            return Ok(None);
2540        }
2541        let first = cache.recur[lin[0]].as_ref().unwrap();
2542        let (conv_words, ssm_words) = (first.conv_state.len(), first.ssm_state.len());
2543        for &il in &lin {
2544            let rl = cache.recur[il].as_ref().unwrap();
2545            if rl.conv_state.len() != conv_words || rl.ssm_state.len() != ssm_words {
2546                return Ok(None);
2547            }
2548        }
2549        let n = lin.len();
2550        let mut lin_pos = std::collections::HashMap::with_capacity(n);
2551        for (k, &il) in lin.iter().enumerate() {
2552            lin_pos.insert(il, k);
2553        }
2554        // longest run of consecutive linear layers (save-scratch sizing)
2555        let mut max_run = 1usize;
2556        let mut run = 1usize;
2557        for w in lin.windows(2) {
2558            if w[1] == w[0] + 1 {
2559                run += 1;
2560                max_run = max_run.max(run);
2561            } else {
2562                run = 1;
2563            }
2564        }
2565        let rows = t_max - 1;
2566        let mut stash_conv = Vec::with_capacity(n);
2567        let mut stash_ssm = Vec::with_capacity(n);
2568        for _ in 0..n {
2569            stash_conv.push(e.uninit(rows * conv_words)?);
2570            stash_ssm.push(e.uninit(rows * ssm_words)?);
2571        }
2572        let host_table = vec![0u64; n * 6];
2573        let table_all = e.htod_u64(&host_table)?;
2574        // slice 4c: full-attention census for the full-verify graphs.
2575        let fa: Vec<usize> = (0..cache.kv.len())
2576            .filter(|&il| cache.kv[il].is_some())
2577            .collect();
2578        let mut fa_pos = std::collections::HashMap::with_capacity(fa.len());
2579        for (k, &il) in fa.iter().enumerate() {
2580            fa_pos.insert(il, k);
2581        }
2582        let n_layers = cache.kv.len().max(cache.recur.len());
2583        // exactly one of (linear state, kv cache) per layer — no MLA/unknown mixers.
2584        let walk_uniform = (0..n_layers).all(|il| {
2585            cache.recur.get(il).is_some_and(|r| r.is_some())
2586                != cache.kv.get(il).is_some_and(|k| k.is_some())
2587        });
2588        // Contiguous covered prefix: the largest n such that every layer in [0, n) is
2589        // linear or full-attention. The TRUNK walk is [0, layers.len()) and the cache
2590        // vecs can carry EXTRA state slots past it (the q38 export keeps the MTP head
2591        // layer's kv at the tail — hi == lin+fa never held, the s4c battery's zero
2592        // 'full' captures). The full-graph guard is walk coverage, not slot arithmetic.
2593        let covered = (0..n_layers)
2594            .take_while(|il| lin_pos.contains_key(il) || fa_pos.contains_key(il))
2595            .count();
2596        let t_cap = t_max;
2597        let fa_host_table = vec![0u64; fa.len() * 2 * t_cap];
2598        let fa_table = e.htod_u64(&fa_host_table)?;
2599        Ok(Some(Self {
2600            lin,
2601            lin_pos,
2602            table_all,
2603            host_table,
2604            stash_conv,
2605            stash_ssm,
2606            conv_words,
2607            ssm_words,
2608            stage: std::collections::HashMap::new(),
2609            tap_bufs: std::collections::HashMap::new(),
2610            graphs: std::collections::HashMap::new(),
2611            save_conv: e.uninit(n * conv_words)?,
2612            save_ssm: e.uninit(n * ssm_words)?,
2613            max_run,
2614            n_embd,
2615            round_slab: false,
2616            fa,
2617            fa_pos,
2618            fa_table,
2619            fa_host_table,
2620            t_cap,
2621            pos_stage: std::collections::HashMap::new(),
2622            full: std::collections::HashMap::new(),
2623            covered,
2624            walk_uniform,
2625            debt_obs: None,
2626        }))
2627    }
2628
2629    /// Rebuild the pointer tables from the live handles (once per verify — the gdn
2630    /// ping-pong swaps the canonical/alt handles between rounds; a fresh generation's
2631    /// cache buffers land at new addresses; a stale table would read the wrong state).
2632    pub(crate) fn refresh_tables(
2633        &mut self,
2634        e: &Engine,
2635        cache: &Cache,
2636    ) -> Result<(), Box<dyn std::error::Error>> {
2637        use cudarc::driver::DevicePtr;
2638        {
2639            let s = &e.gpu.stream();
2640            for (k, &il) in self.lin.iter().enumerate() {
2641                let rl = cache.recur[il].as_ref().unwrap();
2642                let (pc, _g0) = rl.conv_state.device_ptr(s);
2643                let (p0, _g1) = rl.ssm_state.device_ptr(s);
2644                let (p1, _g2) = rl.ssm_state_alt.device_ptr(s);
2645                let o = k * 6;
2646                self.host_table[o] = pc;
2647                self.host_table[o + 1] = p0;
2648                self.host_table[o + 2] = p1;
2649                self.host_table[o + 3] = pc;
2650                self.host_table[o + 4] = p1;
2651                self.host_table[o + 5] = p0;
2652            }
2653            for (k, &il) in self.fa.iter().enumerate() {
2654                let kvl = cache.kv[il].as_ref().unwrap();
2655                let (pk, _g0) = kvl.k.device_ptr(s);
2656                let (pv, _g1) = kvl.v.device_ptr(s);
2657                let o = k * 2 * self.t_cap;
2658                for z in 0..self.t_cap {
2659                    self.fa_host_table[o + 2 * z] = pk;
2660                    self.fa_host_table[o + 2 * z + 1] = pv;
2661                }
2662            }
2663        }
2664        e.htod_u64_into(&self.host_table, &mut self.table_all)?;
2665        if !self.fa_host_table.is_empty() {
2666            e.htod_u64_into(&self.fa_host_table, &mut self.fa_table)?;
2667        }
2668        Ok(())
2669    }
2670
2671    /// Slice 4c eligibility: Some(rung_end) when this round can replay (or capture) a
2672    /// full-verify graph — the whole walk [lo, hi) is covered, every layer is linear or
2673    /// full-attention, and ALL of the round's per-row t_kv values take the v4-seqs arm
2674    /// on ONE `fa_split_keys` ladder step that the rung also sits on (the straddle law;
2675    /// both gates are t_kv intervals, so ends-inside means all-inside). The rung is the
2676    /// round's next power of two — grid/partial sizing only (`n_splits_max` is pure
2677    /// stride; splits >= ns_eff write the empty partial the combine never reads), so one
2678    /// captured graph is bit-identical for every round the rung covers.
2679    #[allow(clippy::too_many_arguments)]
2680    pub(crate) fn full_rung(
2681        &self,
2682        model: &crate::hybrid::HybridModel,
2683        cache: &Cache,
2684        lo: usize,
2685        hi: usize,
2686        t: usize,
2687    ) -> Option<usize> {
2688        if std::env::var("MEMRA_DSPARK_FULLG_DEBUG").as_deref() == Ok("1") {
2689            static ONCE: std::sync::Once = std::sync::Once::new();
2690            let len0 = self
2691                .fa
2692                .first()
2693                .and_then(|&il| cache.kv[il].as_ref())
2694                .map(|k| k.len);
2695            ONCE.call_once(|| {
2696                eprintln!(
2697                    "[fullg-debug] walk_uniform={} covered={} fa_rows_on={} t={} lo={} hi={} lin={} fa={} t_cap={} len0={:?}",
2698                    self.walk_uniform, self.covered, dspark_fa_rows_on(), t, lo, hi,
2699                    self.lin.len(), self.fa.len(), self.t_cap, len0
2700                );
2701            });
2702        }
2703        if !self.walk_uniform
2704            || !dspark_fa_rows_on()
2705            || t < 2
2706            || lo != 0
2707            || hi > self.covered
2708            || t > self.t_cap
2709            || self.fa.is_empty()
2710        {
2711            return None;
2712        }
2713        let cfg = &model.cfg;
2714        let head_dim_global = cfg.head_dim_k as usize;
2715        let nkv = cfg.n_head_kv as usize;
2716        let kvl0 = cache.kv[self.fa[0]].as_ref().unwrap();
2717        // the z-batched twins read stacked rows at the cache's kv dims — must equal the
2718        // projection stride (the body's guard, hoisted so ineligible models fall back
2719        // instead of refusing mid-capture).
2720        let geom = cfg.full_attention_geometry_at(self.fa[0] as u32);
2721        let kv_dim = geom.n_head_kv as usize * geom.head_dim_k as usize;
2722        if kvl0.kv_dim_k != kv_dim || kvl0.kv_dim_v != kv_dim {
2723            return None;
2724        }
2725        let len0 = kvl0.len;
2726        let (t_kv_first, t_kv_last) = (len0 + 1, len0 + t);
2727        if !crate::fa_seqs_eligible(t_kv_first, head_dim_global)
2728            || !crate::fa_seqs_eligible(t_kv_last, head_dim_global)
2729            || crate::fa_split_keys(t_kv_first, nkv) != crate::fa_split_keys(t_kv_last, nkv)
2730        {
2731            return None;
2732        }
2733        let rung = t_kv_last.next_power_of_two().max(256);
2734        if crate::fa_split_keys(rung, nkv) != crate::fa_split_keys(t_kv_last, nkv) {
2735            return None;
2736        }
2737        Some(rung)
2738    }
2739
2740    /// Run the WHOLE verify walk [lo, hi) as one captured graph at (vt=t, rung): stage
2741    /// the residual + refresh the per-vt position staging, capture on first encounter
2742    /// (2 executing warmups bracketed by a full linear-state save/restore; KV warmup
2743    /// appends write the exact slots the replay writes — idempotent), launch, then apply
2744    /// the host bookkeeping the captured body skipped (per-linear-layer parity swap for
2745    /// odd t, per-fa-layer len bump). Returns the fresh residual.
2746    #[allow(clippy::too_many_arguments)]
2747    #[allow(clippy::map_entry)] // allow: the init body is fallible (`?`); Entry::or_insert_with cannot propagate errors
2748    pub(crate) fn run_full(
2749        &mut self,
2750        model: &crate::hybrid::HybridModel,
2751        e: &Engine,
2752        lo: usize,
2753        hi: usize,
2754        x: &CudaSlice<f32>,
2755        t: usize,
2756        pos0: usize,
2757        rung: usize,
2758        cache: &mut Cache,
2759    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2760        let n_embd = self.n_embd;
2761        if !self.stage.contains_key(&t) {
2762            let xin = e.uninit(t * n_embd)?;
2763            let xout = e.uninit(t * n_embd)?;
2764            self.stage.insert(t, (xin, xout));
2765        }
2766        if !self.pos_stage.contains_key(&t) {
2767            self.pos_stage.insert(t, e.htod_i32(&vec![0i32; t])?);
2768        }
2769        // Per-round refresh: position contents + input staging (both addresses are baked
2770        // by the captured bodies; only their CONTENTS change round to round).
2771        {
2772            let pos_host: Vec<i32> = (0..t).map(|r| (pos0 + r) as i32).collect();
2773            let pb = self.pos_stage.get_mut(&t).unwrap();
2774            e.htod_i32_into(pb, &pos_host)?;
2775            let (xin, _) = self.stage.get_mut(&t).unwrap();
2776            e.copy_into(xin, 0, x, t * n_embd)?;
2777        }
2778        let key = (t, rung, hi);
2779        if !self.full.contains_key(&key) {
2780            // The warmups EXECUTE the whole walk on live state — save every linear
2781            // layer's conv + canonical ssm first, restore after (KV needs no restore:
2782            // graph mode never bumps host lens and the appends write this round's own
2783            // slots).
2784            for (k, &il) in self.lin.iter().enumerate() {
2785                let rl = cache.recur[il].as_ref().unwrap();
2786                e.copy_into(
2787                    &mut self.save_conv,
2788                    k * self.conv_words,
2789                    &rl.conv_state,
2790                    self.conv_words,
2791                )?;
2792                e.copy_into(
2793                    &mut self.save_ssm,
2794                    k * self.ssm_words,
2795                    &rl.ssm_state,
2796                    self.ssm_words,
2797                )?;
2798            }
2799            let (graph, keeper) = {
2800                let table_all = &self.table_all;
2801                let lin_pos = &self.lin_pos;
2802                let fa_pos = &self.fa_pos;
2803                let fa_table = &self.fa_table;
2804                let t_cap = self.t_cap;
2805                let stash_conv = &mut self.stash_conv;
2806                let stash_ssm = &mut self.stash_ssm;
2807                let pos_d: &CudaSlice<i32> = &self.pos_stage[&t];
2808                let (xin, xout) = self
2809                    .stage
2810                    .get_mut(&t)
2811                    .map(|(a, b)| (&*a, b))
2812                    .expect("stage bucket created above");
2813                let cache_ref: &mut Cache = cache;
2814                let iflag = if std::env::var("MEMRA_DSPARK_VG_AUTOFREE").as_deref() == Ok("1") {
2815                    cudarc::driver::sys::CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH
2816                } else {
2817                    cudarc::driver::sys::CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_USE_NODE_PRIORITY
2818                };
2819                e.capture_graph_retained_flags(iflag, move |e| {
2820                    let mut xc: Option<CudaSlice<f32>> = None;
2821                    for il in lo..hi {
2822                        let xr: &CudaSlice<f32> = xc.as_ref().unwrap_or(xin);
2823                        let nx = if let Some(&k) = lin_pos.get(&il) {
2824                            model.qwen35_tparallel_linear_layer(
2825                                e,
2826                                il,
2827                                xr,
2828                                t,
2829                                cache_ref,
2830                                None,
2831                                Some((&mut stash_conv[k], &mut stash_ssm[k])),
2832                                Some((table_all, k * 6)),
2833                            )?
2834                        } else if let Some(&kf) = fa_pos.get(&il) {
2835                            let mut no_rows: Option<Vec<CudaSlice<i32>>> = None;
2836                            model.qwen35_tparallel_fa_layer(
2837                                e,
2838                                il,
2839                                xr,
2840                                t,
2841                                cache_ref,
2842                                FaLayerArgs {
2843                                    pos_d,
2844                                    pos_rows: &mut no_rows,
2845                                    pos0,
2846                                    graph_cap: Some((fa_table, kf * 2 * t_cap, rung)),
2847                                    stream: None,
2848                                    ckpt: None,
2849                                },
2850                            )?
2851                        } else {
2852                            return Err(format!(
2853                                "run_full: layer {il} is neither linear nor full-attention"
2854                            )
2855                            .into());
2856                        };
2857                        xc = Some(nx);
2858                    }
2859                    e.copy_into(xout, 0, xc.as_ref().unwrap(), t * n_embd)?;
2860                    Ok(())
2861                })?
2862            };
2863            // Undo the net host parity motion of the 3 body runs (each run swaps iff t
2864            // is odd -> 3 runs = net one swap), then restore the device state the
2865            // warmups consumed (walk scope only — layers past hi never executed). The
2866            // launch below then behaves exactly like one run.
2867            if t % 2 == 1 {
2868                for &il in &self.lin {
2869                    if il < lo || il >= hi {
2870                        continue;
2871                    }
2872                    let rl = cache.recur[il].as_mut().unwrap();
2873                    std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
2874                }
2875            }
2876            for (k, &il) in self.lin.iter().enumerate() {
2877                if il < lo || il >= hi {
2878                    continue;
2879                }
2880                let rl = cache.recur[il].as_mut().unwrap();
2881                let (cw, sw) = (self.conv_words, self.ssm_words);
2882                {
2883                    let sv = e.view(&self.save_conv, self.lin.len() * cw);
2884                    let win = sv.slice(k * cw..(k + 1) * cw);
2885                    e.copy_view_into(&mut rl.conv_state, 0, &win, cw)?;
2886                }
2887                {
2888                    let sv = e.view(&self.save_ssm, self.lin.len() * sw);
2889                    let win = sv.slice(k * sw..(k + 1) * sw);
2890                    e.copy_view_into(&mut rl.ssm_state, 0, &win, sw)?;
2891                }
2892            }
2893            if std::env::var("MEMRA_GRAPH_CENSUS").as_deref() == Ok("1")
2894                && let Ok(c) = crate::graph_update::node_census(&graph)
2895            {
2896                eprintln!("[dspark-vg-census] full vt={t} rung={rung} {c:?}");
2897            }
2898            self.full.insert(
2899                key,
2900                DsparkSegGraph {
2901                    graph,
2902                    _keeper: keeper,
2903                },
2904            );
2905        }
2906        self.full[&key].graph.launch()?;
2907        // Host bookkeeping for the replayed body (captured host code does not re-run):
2908        // gdn parity swap per linear layer (t odd), kv len bump per fa layer — scoped
2909        // to the WALK [lo, hi): the cache can carry extra state slots past it (the MTP
2910        // head layer's kv) that the walk never touches.
2911        if t % 2 == 1 {
2912            for &il in &self.lin {
2913                if il < lo || il >= hi {
2914                    continue;
2915                }
2916                let rl = cache.recur[il].as_mut().unwrap();
2917                std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
2918            }
2919        }
2920        for &il in &self.fa {
2921            if il < lo || il >= hi {
2922                continue;
2923            }
2924            cache.kv[il].as_mut().unwrap().len += t;
2925        }
2926        let (_, xout) = self.stage.get(&t).unwrap();
2927        let mut out = e.uninit(t * n_embd)?;
2928        e.copy_into(&mut out, 0, xout, t * n_embd)?;
2929        Ok(out)
2930    }
2931
2932    /// Run layers [start, end) (all linear) as one captured graph at this vt: stage the
2933    /// residual into the bucket's x_in, capture on first encounter (2 executing warmups
2934    /// bracketed by a segment state save/restore), launch, then apply the host parity
2935    /// bookkeeping the captured body would have done. Returns the fresh residual.
2936    #[allow(clippy::too_many_arguments)]
2937    #[allow(clippy::map_entry)] // allow: the init body is fallible (`?`); Entry::or_insert_with cannot propagate errors
2938    fn run_segment(
2939        &mut self,
2940        model: &crate::hybrid::HybridModel,
2941        e: &Engine,
2942        start: usize,
2943        end: usize,
2944        x: &CudaSlice<f32>,
2945        t: usize,
2946        cache: &mut Cache,
2947    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2948        let n_embd = self.n_embd;
2949        debug_assert!(end - start <= self.max_run);
2950        if !self.stage.contains_key(&t) {
2951            let xin = e.uninit(t * n_embd)?;
2952            let xout = e.uninit(t * n_embd)?;
2953            self.stage.insert(t, (xin, xout));
2954        }
2955        // Stage the residual at the bucket's baked input address.
2956        {
2957            let (xin, _) = self.stage.get_mut(&t).unwrap();
2958            e.copy_into(xin, 0, x, t * n_embd)?;
2959        }
2960        let key = (start, t);
2961        if !self.graphs.contains_key(&key) {
2962            // The 2 warmups EXECUTE the segment on live state — save conv + the canonical
2963            // ssm of every segment layer first, restore after, so the graph's first real
2964            // launch starts from the exact pre-round state (bytes gated e2e).
2965            for (k, il) in (start..end).enumerate() {
2966                let rl = cache.recur[il].as_ref().unwrap();
2967                e.copy_into(
2968                    &mut self.save_conv,
2969                    k * self.conv_words,
2970                    &rl.conv_state,
2971                    self.conv_words,
2972                )?;
2973                e.copy_into(
2974                    &mut self.save_ssm,
2975                    k * self.ssm_words,
2976                    &rl.ssm_state,
2977                    self.ssm_words,
2978                )?;
2979            }
2980            let (graph, keeper) = {
2981                let table_all = &self.table_all;
2982                let lin_pos = &self.lin_pos;
2983                let stash_conv = &mut self.stash_conv;
2984                let stash_ssm = &mut self.stash_ssm;
2985                let (xin, xout) = self
2986                    .stage
2987                    .get_mut(&t)
2988                    .map(|(a, b)| (&*a, b))
2989                    .expect("stage bucket created above");
2990                let cache_ref: &mut Cache = cache;
2991                // Slice 4 (fa-execupdate lane): USE_NODE_PRIORITY instead of
2992                // AUTO_FREE_ON_LAUNCH. The slice-3 measured limiter was AUTO_FREE's
2993                // launch-time mem-pool scan — 25.6 us per cuGraphLaunch x 16 segments
2994                // = ~0.41 ms/round, most of the eager-launch savings. The captured
2995                // body's cuMemAllocAsync transients are BALANCED by in-graph frees
2996                // (every transient drops inside the capture region — the generic
2997                // capture path's census precedent, 1589/1589), so AUTO_FREE has
2998                // nothing to reclaim and the graph is legal to instantiate without
2999                // it; PRIORITY is the flag the gemma slotted door ships for exactly
3000                // this reason (both alternatives drop the scan; UPLOAD via
3001                // cuGraphInstantiateWithFlags is WithParams-only and refused).
3002                // MEMRA_DSPARK_VG_AUTOFREE=1 reverts; MEMRA_GRAPH_CENSUS=1 prints
3003                // the node census at capture (the ALLOC==FREE receipt).
3004                let iflag = if std::env::var("MEMRA_DSPARK_VG_AUTOFREE").as_deref() == Ok("1") {
3005                    cudarc::driver::sys::CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH
3006                } else {
3007                    cudarc::driver::sys::CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_USE_NODE_PRIORITY
3008                };
3009                e.capture_graph_retained_flags(iflag, move |e| {
3010                    let mut xc: Option<CudaSlice<f32>> = None;
3011                    for il in start..end {
3012                        let k = lin_pos[&il];
3013                        let xr: &CudaSlice<f32> = xc.as_ref().unwrap_or(xin);
3014                        let nx = model.qwen35_tparallel_linear_layer(
3015                            e,
3016                            il,
3017                            xr,
3018                            t,
3019                            cache_ref,
3020                            None,
3021                            Some((&mut stash_conv[k], &mut stash_ssm[k])),
3022                            Some((table_all, k * 6)),
3023                        )?;
3024                        xc = Some(nx);
3025                    }
3026                    e.copy_into(xout, 0, xc.as_ref().unwrap(), t * n_embd)?;
3027                    Ok(())
3028                })?
3029            };
3030            // Undo the net host parity motion of the 3 body runs (each run swaps iff t
3031            // is odd -> 3 runs = net one swap), then restore the device state the
3032            // warmups consumed. The launch below then behaves exactly like one run.
3033            if t % 2 == 1 {
3034                for il in start..end {
3035                    let rl = cache.recur[il].as_mut().unwrap();
3036                    std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
3037                }
3038            }
3039            for (k, il) in (start..end).enumerate() {
3040                let rl = cache.recur[il].as_mut().unwrap();
3041                let (cw, sw) = (self.conv_words, self.ssm_words);
3042                {
3043                    let sv = e.view(&self.save_conv, self.lin.len() * cw);
3044                    let win = sv.slice(k * cw..(k + 1) * cw);
3045                    e.copy_view_into(&mut rl.conv_state, 0, &win, cw)?;
3046                }
3047                {
3048                    let sv = e.view(&self.save_ssm, self.lin.len() * sw);
3049                    let win = sv.slice(k * sw..(k + 1) * sw);
3050                    e.copy_view_into(&mut rl.ssm_state, 0, &win, sw)?;
3051                }
3052            }
3053            if std::env::var("MEMRA_GRAPH_CENSUS").as_deref() == Ok("1")
3054                && let Ok(c) = crate::graph_update::node_census(&graph)
3055            {
3056                eprintln!("[dspark-vg-census] seg={start}..{end} vt={t} {c:?}");
3057            }
3058            self.graphs.insert(
3059                key,
3060                DsparkSegGraph {
3061                    graph,
3062                    _keeper: keeper,
3063                },
3064            );
3065        }
3066        self.graphs[&key].graph.launch()?;
3067        // Host parity bookkeeping for the replayed body (the captured host swaps do not
3068        // re-run at replay).
3069        if t % 2 == 1 {
3070            for il in start..end {
3071                let rl = cache.recur[il].as_mut().unwrap();
3072                std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
3073            }
3074        }
3075        let (_, xout) = self.stage.get(&t).unwrap();
3076        let mut out = e.uninit(t * n_embd)?;
3077        e.copy_into(&mut out, 0, xout, t * n_embd)?;
3078        Ok(out)
3079    }
3080
3081    /// Pool freeze check (`dspark_vg_cap`): below the ceiling new keys may capture.
3082    fn can_capture(&self) -> bool {
3083        self.graphs.len() + self.full.len() < dspark_vg_cap()
3084    }
3085
3086    /// Round-atomic segment-door readiness: TRUE when this round's walk can ride the
3087    /// per-(segment, vt) graphs without a NEW capture past the pool ceiling — every
3088    /// linear run in [lo, hi) already has its (run_start, t) key, or capture is still
3089    /// allowed. FALSE sends the WHOLE round down the eager cols-ckpt walk: a partial
3090    /// refusal would stash some layers in the ctx slabs and others in the round's cols
3091    /// while one commit reads only one of them.
3092    pub(crate) fn segments_ready(
3093        &self,
3094        model: &crate::hybrid::HybridModel,
3095        lo: usize,
3096        hi: usize,
3097        t: usize,
3098    ) -> bool {
3099        if self.can_capture() {
3100            return true;
3101        }
3102        let mut il = lo;
3103        while il < hi {
3104            if matches!(model.layers[il].mixer, Mixer::Linear(_)) {
3105                let start = il;
3106                while il < hi && matches!(model.layers[il].mixer, Mixer::Linear(_)) {
3107                    il += 1;
3108                }
3109                if !self.graphs.contains_key(&(start, t)) {
3110                    return false;
3111                }
3112            } else {
3113                il += 1;
3114            }
3115        }
3116        true
3117    }
3118
3119    /// Widest verify window this pool was built for. A caller whose round exceeds it must
3120    /// take the eager walk: the stash slabs hold `t_capacity() - 1` column rows, and slicing
3121    /// past them is a panic rather than a refusal.
3122    pub(crate) fn t_capacity(&self) -> usize {
3123        self.t_cap
3124    }
3125
3126    /// Slab row (conv, ssm) device pointers + lengths for the commit restore of column
3127    /// `row` (0-based) of layer `il`. None for non-linear layers.
3128    pub(crate) fn slab_row(
3129        &self,
3130        e: &Engine,
3131        il: usize,
3132        row: usize,
3133    ) -> Option<(u64, u64, usize, usize)> {
3134        use cudarc::driver::DevicePtr;
3135        let k = *self.lin_pos.get(&il)?;
3136        let s = &e.gpu.stream();
3137        let (pc, _g0) = self.stash_conv[k].device_ptr(s);
3138        let (ps, _g1) = self.stash_ssm[k].device_ptr(s);
3139        Some((
3140            pc + (row * self.conv_words * 4) as u64,
3141            ps + (row * self.ssm_words * 4) as u64,
3142            self.conv_words,
3143            self.ssm_words,
3144        ))
3145    }
3146}
3147
3148impl VerifyCkpt {
3149    fn new(n_layer: usize) -> Self {
3150        VerifyCkpt {
3151            gdn: (0..n_layer).map(|_| None).collect(),
3152            cols: (0..n_layer).map(|_| None).collect(),
3153        }
3154    }
3155}
3156
3157/// The stage-0/TX half of one PP verify. The boundary slot is the ownership token: stage 1
3158/// consumes exactly the slot selected by `tx()`, never a slot inferred from
3159/// a logical round number.
3160struct VerifyBoundaryTicket {
3161    rt: &'static crate::pp::PpNRt,
3162    caller_stream: std::sync::Arc<cudarc::driver::CudaStream>,
3163    slot: usize,
3164    pos0: usize,
3165    t: usize,
3166    payload: usize,
3167    n_st: usize,
3168    pp_anatomy: bool,
3169    pp_started: std::time::Instant,
3170    reverse_ms: f64,
3171    stage0_ms: f64,
3172    tx_ms: f64,
3173    _walk_owner: crate::pp::PpWalkLease,
3174}
3175
3176/// MEMRA_SPEC_ROUND_PROF counters: whole-round wall, so the round can be weighed against the
3177/// draft-step ([spec-anatomy]) and verify-walk ([tcol-prof]) splits we already print.
3178static ROUND_PROF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
3179static ROUND_MS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
3180static ROUND_N: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
3181
3182fn validate_tp_kv_snapshot_shape(
3183    tp_kv: &[Option<crate::tp::ResidentTpKvCache>],
3184    saved_lens: &[Option<usize>],
3185) -> Result<(), Box<dyn std::error::Error>> {
3186    if tp_kv.len() != saved_lens.len() {
3187        return Err("spec TP KV snapshot shape mismatch".into());
3188    }
3189    for (layer, (cache, saved)) in tp_kv.iter().zip(saved_lens).enumerate() {
3190        if cache.is_some() != saved.is_some() {
3191            return Err(
3192                format!("spec TP KV layer {layer} changed shape since its snapshot").into(),
3193            );
3194        }
3195    }
3196    Ok(())
3197}
3198
3199impl HybridModel {
3200    /// memra#128: the canonical-to-rank byte copy that `5e0fffb97` added to
3201    /// `restore_step_tp_kv_verified_prefix`. OFF by default (written decision, docs/FLAGS.md
3202    /// `MEMRA_STEP_TP_KV_RESTORE`): on step37 NVFP4 TP2 the only shapes that pass the
3203    /// production acceptance gate are the engine before the copy (arm A) and the copy skipped
3204    /// on every step-TP layer (arm F, byte-identical answers to A); the copy on any layer
3205    /// spliced answers or shifted decode (darklanes research/memra128-bisect-20260903).
3206    /// `1` re-enables the copy for ordinary-commit layers; on-device-written layers
3207    /// (`rows_external`) are skipped either way, their canonical bytes are stale.
3208    fn step_tp_kv_restore_copy_on() -> bool {
3209        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
3210        *ON.get_or_init(|| std::env::var("MEMRA_STEP_TP_KV_RESTORE").ok().as_deref() == Some("1"))
3211    }
3212
3213    fn restore_step_tp_kv_verified_prefix(
3214        &self,
3215        e: &Engine,
3216        cache: &mut Cache,
3217        snap: &crate::cache::CacheSnapshot,
3218        accepted: usize,
3219        // memra#128: what an externally-written (dcw / fa2) layer needs from this call.
3220        // PARTIAL accept (commit_verified_prefix): 5e0fffb97 replaced the standalone
3221        // rewind_tp_kv_verified_prefix with this restore, so the length shrink to
3222        // saved+accepted must still happen here - without it E ran with distributed=259
3223        // against local=257. FULL accept: before 5e0fffb97 nothing touched the
3224        // distributed length there and it was right (arm A passed); rewinding to
3225        // saved+t_v shrinks it by one and the next verify's SWA ring view falls off the
3226        // end ("view [5148,5152) is outside resident [0,5151)", arm E2).
3227        rewind_external: bool,
3228    ) -> Result<(), Box<dyn std::error::Error>> {
3229        validate_tp_kv_snapshot_shape(&cache.tp_kv, &snap.tp_kv_len)?;
3230        e.stream().synchronize()?;
3231        {
3232            let (local_layers, distributed_layers) = (&cache.kv, &mut cache.tp_kv);
3233            let stream = e.gpu.stream();
3234            let mut runtime: Option<std::sync::Arc<crate::tp::TpE4m3HostBounce>> = None;
3235            let mut uniform_runtime = true;
3236            let mut batch = Vec::new();
3237            for (il, (distributed_slot, local_slot)) in distributed_layers
3238                .iter_mut()
3239                .zip(local_layers.iter())
3240                .enumerate()
3241            {
3242                let (Some(distributed), Some(saved)) =
3243                    (distributed_slot.as_mut(), snap.tp_kv_len[il])
3244                else {
3245                    continue;
3246                };
3247                // memra#128: on the dcw / fa2 verify path the rank rows for [saved, target)
3248                // were written on-device and the canonical cache holds only stale bytes for
3249                // them (the verify bumps `local.len` and writes nothing). Copying those over
3250                // the correct rank rows is exactly what spliced two requests' answers
3251                // together on step-3.7-flash. The rewind already kept the right rows; skip.
3252                let target = saved
3253                    .checked_add(accepted)
3254                    .ok_or("spec TP KV batch restore length overflow")?;
3255                if !Self::step_tp_kv_restore_copy_on() || distributed.rows_external() {
3256                    // Rank rows already right (written on-device). Length: see the
3257                    // `rewind_external` note on the signature.
3258                    if rewind_external {
3259                        distributed.rewind_to(target)?;
3260                    }
3261                    continue;
3262                }
3263                let local = local_slot
3264                    .as_ref()
3265                    .ok_or_else(|| format!("spec TP KV layer {il} lost its owning cache"))?;
3266                if local.len < target {
3267                    return Err(format!(
3268                        "spec TP KV layer {il} local length {} precedes restore target {target}",
3269                        local.len
3270                    )
3271                    .into());
3272                }
3273                let physical = local.physical_rows(saved, target)?;
3274                if physical.len() != accepted {
3275                    return Err(format!(
3276                        "spec TP KV layer {il} restore [{saved},{target}) is not contiguous"
3277                    )
3278                    .into());
3279                }
3280                let Mixer::Full(fa) = &self.layers[il].mixer else {
3281                    return Err(format!("spec TP KV layer {il} is not full attention").into());
3282                };
3283                let tp = fa
3284                    .step_tp_qkv
3285                    .as_ref()
3286                    .ok_or_else(|| format!("spec TP KV layer {il} lost its TP runtime"))?;
3287                if let Some(first) = runtime.as_ref() {
3288                    if !std::sync::Arc::ptr_eq(first, &tp.runtime) {
3289                        uniform_runtime = false;
3290                        break;
3291                    }
3292                } else {
3293                    runtime = Some(tp.runtime.clone());
3294                }
3295                use cudarc::driver::DevicePtr;
3296                let (k_base, _k_guard) = local.k.device_ptr(&stream);
3297                let (v_base, _v_guard) = local.v.device_ptr(&stream);
3298                batch.push(crate::tp::TpKvVerifiedLayer {
3299                    cache: distributed,
3300                    start: saved,
3301                    logical_len: target,
3302                    source_k_raw: k_base + (physical.start * local.k_tok_bytes) as u64,
3303                    source_v_raw: v_base + (physical.start * local.v_tok_bytes) as u64,
3304                    source_k_tok_bytes: local.k_tok_bytes,
3305                    source_v_tok_bytes: local.v_tok_bytes,
3306                });
3307            }
3308            if uniform_runtime
3309                && let Some(runtime) = runtime
3310                && runtime.restore_tp_kv_layers_from_device(&mut batch)?
3311            {
3312                return Ok(());
3313            }
3314        }
3315        for il in 0..self.layers.len() {
3316            let (Some(distributed), Some(saved)) = (cache.tp_kv[il].as_mut(), snap.tp_kv_len[il])
3317            else {
3318                continue;
3319            };
3320            let target = saved
3321                .checked_add(accepted)
3322                .ok_or("spec TP KV restore length overflow")?;
3323            if !Self::step_tp_kv_restore_copy_on() || distributed.rows_external() {
3324                if rewind_external {
3325                    distributed.rewind_to(target)?;
3326                }
3327                continue;
3328            }
3329            let local = cache.kv[il]
3330                .as_ref()
3331                .ok_or_else(|| format!("spec TP KV layer {il} lost its owning cache"))?;
3332            if local.len < target {
3333                return Err(format!(
3334                    "spec TP KV layer {il} local length {} precedes restore target {target}",
3335                    local.len
3336                )
3337                .into());
3338            }
3339            let physical = local.physical_rows(saved, target)?;
3340            if physical.len() != accepted {
3341                return Err(format!(
3342                    "spec TP KV layer {il} restore [{saved},{target}) is not contiguous"
3343                )
3344                .into());
3345            }
3346            use cudarc::driver::DevicePtr;
3347            let stream = e.gpu.stream();
3348            let (k_base, _k_guard) = local.k.device_ptr(&stream);
3349            let (v_base, _v_guard) = local.v.device_ptr(&stream);
3350            let k_raw = k_base + (physical.start * local.k_tok_bytes) as u64;
3351            let v_raw = v_base + (physical.start * local.v_tok_bytes) as u64;
3352            let Mixer::Full(fa) = &self.layers[il].mixer else {
3353                return Err(format!("spec TP KV layer {il} is not full attention").into());
3354            };
3355            let tp = fa
3356                .step_tp_qkv
3357                .as_ref()
3358                .ok_or_else(|| format!("spec TP KV layer {il} lost its TP runtime"))?;
3359            tp.runtime.restore_tp_kv_rows_from_device(
3360                distributed,
3361                saved,
3362                target,
3363                k_raw,
3364                v_raw,
3365                local.k_tok_bytes,
3366                local.v_tok_bytes,
3367            )?;
3368        }
3369        Ok(())
3370    }
3371
3372    fn mtp_head_count(&self) -> usize {
3373        usize::from(self.mtp.is_some()) + self.mtp_extra.len()
3374    }
3375
3376    fn mtp_head_at(&self, index: usize) -> &MtpHead {
3377        if index == 0 {
3378            self.mtp.as_ref().expect("MTP head 0 is unavailable")
3379        } else {
3380            &self.mtp_extra[index - 1]
3381        }
3382    }
3383
3384    fn new_mtp_scratch(
3385        &self,
3386        e: &Engine,
3387        cap: usize,
3388    ) -> Result<MtpScratch, Box<dyn std::error::Error>> {
3389        let mut scratch = MtpScratch::new(
3390            e,
3391            &self.cfg,
3392            &self.plan,
3393            cap,
3394            self.mtp.as_ref().and_then(|head| head.geom.as_ref()),
3395        )?;
3396        for head in &self.mtp_extra {
3397            scratch.push_plane(e, &self.cfg, &self.plan, head.geom.as_ref())?;
3398        }
3399        Ok(scratch)
3400    }
3401
3402    /// NextN head forward for ONE draft token (§A ops 1-13, T=1).
3403    /// Inputs: `e_tok` = the token to predict FROM (last committed / previous draft); `h_seed` =
3404    /// the trunk's pre-output_norm hidden of that token (§A op 2 input). `mtp_pos` = absolute
3405    /// position of the token being predicted from. Returns (draft_logits[n_vocab] host, h_nextn dev).
3406    /// `h_nextn` (§A op 10) becomes `h_seed` for the next autoregressive draft step.
3407    /// Device-resident: returns draft logits ON DEVICE (no [n_vocab] dtoh). The greedy draft
3408    /// loop only needs argmax — paired with `argmax_token_device` this cuts the ~600KB logits
3409    /// transfer + host argmax per draft token from the K-token draft chain.
3410    #[allow(clippy::too_many_arguments)]
3411    fn mtp_head_forward_dev(
3412        &self,
3413        e: &Engine,
3414        mtp: &MtpHead,
3415        e_tok: u32,
3416        h_seed: &CudaSlice<f32>,
3417        scratch: &mut MtpScratch,
3418        mtp_pos: usize,
3419        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
3420        mask: Option<(&CudaSlice<u32>, usize)>,
3421    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
3422        self.mtp_head_forward_dev_at(e, mtp, e_tok, h_seed, scratch, 0, mtp_pos, embd_dev, mask)
3423    }
3424
3425    #[allow(clippy::too_many_arguments)]
3426    fn mtp_head_forward_dev_at(
3427        &self,
3428        e: &Engine,
3429        mtp: &MtpHead,
3430        e_tok: u32,
3431        h_seed: &CudaSlice<f32>,
3432        scratch: &mut MtpScratch,
3433        scratch_index: usize,
3434        mtp_pos: usize,
3435        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
3436        // DRAFT-SIDE GRAMMAR MASK (lane/draft-mask): (packed draft-vocab allowed set, words).
3437        // Applied to the head logits BEFORE they are returned, so every consumer (argmax,
3438        // gumbel draw, p-min prob) sees the grammar-legal row. None = unmasked (pre-lane).
3439        mask: Option<(&CudaSlice<u32>, usize)>,
3440    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
3441        // MEMRA_SPEC_ANATOMY=1 — eager-step phase timers (diagnostic only). Phase boundaries
3442        // sync the stream, so absolute time inflates; the BREAKDOWN is the signal. Cumulative
3443        // summary on stderr every 128 steps: glue (embed..attn_norm), attn, ffn, head.
3444        use std::sync::atomic::{AtomicU64, Ordering::Relaxed};
3445        static ANAT_NS: [AtomicU64; 5] = [
3446            AtomicU64::new(0),
3447            AtomicU64::new(0),
3448            AtomicU64::new(0),
3449            AtomicU64::new(0),
3450            AtomicU64::new(0),
3451        ];
3452        static ANAT_STEPS: AtomicU64 = AtomicU64::new(0);
3453        let anat = {
3454            static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
3455            *ON.get_or_init(|| std::env::var("MEMRA_SPEC_ANATOMY").as_deref() == Ok("1"))
3456        };
3457        if anat {
3458            e.stream().synchronize()?; // drain prior queue so phase 0 starts clean
3459        }
3460        let t_all = std::time::Instant::now();
3461        let mut t_ph = std::time::Instant::now();
3462        let anat_mark = |i: usize,
3463                         e: &Engine,
3464                         t: &mut std::time::Instant|
3465         -> Result<(), Box<dyn std::error::Error>> {
3466            if anat {
3467                e.stream().synchronize()?;
3468                ANAT_NS[i].fetch_add(t.elapsed().as_nanos() as u64, Relaxed);
3469                *t = std::time::Instant::now();
3470            }
3471            Ok(())
3472        };
3473        let cfg = &self.cfg;
3474        let n_embd = cfg.n_embd as usize;
3475        // Distilled-student geometry: the block runs at the INNER width `di` (eh_proj out /
3476        // attn / ffn); the n_embd interface (embed, norms in, carrier out, head in) is unchanged.
3477        let di = mtp.geom.as_ref().map(|g| g.d_inner).unwrap_or(n_embd);
3478        let eps = cfg.rms_eps;
3479        let pos_d = e.htod_i32(&[mtp_pos as i32])?;
3480
3481        // op A: a resident table transfers one 4B token id. The exact host-row capacity path
3482        // expands this one row on CPU and transfers n_embd f32 values instead.
3483        let e_emb = match embd_dev {
3484            Some((g, qt, rb)) => e.embed_gather_device_t(g, &[e_tok], n_embd, qt, rb)?,
3485            None => e.htod(&self.embd.try_gather(n_embd, &[e_tok])?)?,
3486        };
3487
3488        // op 1/2: e_norm = RMSNorm(e, enorm); h_norm = RMSNorm(h_seed, hnorm)
3489        let mut e_norm = e.zeros(n_embd)?;
3490        e.rms_norm(&e_emb, mtp.enorm.float_data(), &mut e_norm, n_embd, 1, eps)?;
3491        let mut h_norm = e.zeros(n_embd)?;
3492        e.rms_norm(h_seed, mtp.hnorm.float_data(), &mut h_norm, n_embd, 1, eps)?;
3493
3494        // op 3: concat = [e_norm ; h_norm] -> [2*n_embd], e_norm in [0,n_embd), h_norm in [n_embd,2n_embd)
3495        let mut concat = e.zeros(2 * n_embd)?;
3496        e.copy_into(&mut concat, 0, &e_norm, n_embd)?;
3497        e.copy_into(&mut concat, n_embd, &h_norm, n_embd)?;
3498
3499        // op 4: inpSA = eh_proj @ concat  (eh_proj [2*n_embd, n_embd]) -> [n_embd]
3500        let inp_sa = e.matmul(&mtp.eh_proj, &concat, 1)?;
3501
3502        // op 5: a_norm = RMSNorm(inpSA, attn_norm)
3503        let mut a_norm = e.zeros(di)?;
3504        e.rms_norm(&inp_sa, mtp.attn_norm.float_data(), &mut a_norm, di, 1, eps)?;
3505        anat_mark(0, e, &mut t_ph)?;
3506
3507        // op 6: attention on the scratch KV. SAME dc launcher as the graph path (bucket_max =
3508        // scratch.cap, length from the device len_d) so eager drafts match graph drafts
3509        // bit-for-bit at any t_kv (the parity gate). Host len mirrored here (the dc append
3510        // advances only the device counter).
3511        let attn_out = match (&mtp.mixer, mtp.step35.as_ref()) {
3512            // step35 MTP block, dcw door armed: the SAME windowed device-counter launcher as
3513            // the captured chain (draft parity by construction). Per-step ring headroom runs
3514            // HERE (eager is host-len work, a rebase is legal); host len mirrored like the
3515            // plain dc arm below.
3516            (Mixer::Full(fa), Some(g))
3517                if self.step35_dcw_eligible(g, scratch.plane(scratch_index).1) =>
3518            {
3519                {
3520                    let (kv, _) = scratch.plane_mut(scratch_index);
3521                    let retain = match kv.ring.as_ref() {
3522                        Some(ring) => memra_kv::swa_retain_from(kv.len, ring.window(), ring.base()),
3523                        None => 0,
3524                    };
3525                    e.prepare_kv_append(kv, retain, 1)?;
3526                }
3527                let out =
3528                    self.mtp_step35_attn_dcw(e, fa, g, &a_norm, &pos_d, scratch, scratch_index)?;
3529                scratch.plane_mut(scratch_index).0.len += 1;
3530                out
3531            }
3532            // step35 MTP block, door off (MEMRA_STEP35_DRAFT_DCW=0 rollback) or class-
3533            // ineligible: PER-LAYER geometry + a separate head-wise gate + an SWA window,
3534            // none of which the plain dc launcher can express (see `mtp_step35_attn`).
3535            // Host-len arm. Advances BOTH the
3536            // host len and the device counter itself (unlike the dc arm, whose host-side
3537            // mirror the caller does).
3538            (Mixer::Full(fa), Some(g)) => {
3539                self.mtp_step35_attn(e, fa, g, &a_norm, &pos_d, scratch, scratch_index)?
3540            }
3541            (Mixer::Full(fa), None) => {
3542                let out = self.mtp_full_attn_dc(
3543                    e,
3544                    fa,
3545                    &a_norm,
3546                    &pos_d,
3547                    scratch,
3548                    scratch_index,
3549                    mtp.geom.as_ref(),
3550                )?;
3551                scratch.plane_mut(scratch_index).0.len += 1;
3552                out
3553            }
3554            (Mixer::Linear(_), _) => {
3555                panic!("MTP block is full-attn in qwen35; linear MTP not supported")
3556            }
3557            (Mixer::Mla(_), _) => crate::hybrid::mla_path_unimplemented("MTP head forward"),
3558            (Mixer::Kda(_), _) => crate::hybrid::kda_path_unimplemented("MTP head forward"),
3559        };
3560        anat_mark(1, e, &mut t_ph)?;
3561
3562        // op 7: x1 = inpSA + attn_out
3563        let mut x1 = e.zeros(di)?;
3564        e.add(&inp_sa, &attn_out, &mut x1, di)?;
3565
3566        // op 8: z = RMSNorm(x1, post_attn_norm)  (pre-FFN norm)
3567        let mut z = e.zeros(di)?;
3568        e.rms_norm(&x1, mtp.post_attn_norm.float_data(), &mut z, di, 1, eps)?;
3569
3570        // op 9: FFN (Dense or MoE) — same as the trunk decode FFN
3571        let ffn_out = match &mtp.ffn {
3572            crate::hybrid::Ffn::Dense {
3573                ffn_gate,
3574                ffn_up,
3575                ffn_down,
3576            } => {
3577                let n_ff = ffn_gate.out_features();
3578                let (gate, up) = if e.uses_q8_1_fast(ffn_gate) && e.uses_q8_1_fast(ffn_up) {
3579                    let (zq, zd) = e.quantize_q8_1(&z, 1, di)?;
3580                    (
3581                        e.matmul_pre(ffn_gate, &zq, &zd, &z, 1)?,
3582                        e.matmul_pre(ffn_up, &zq, &zd, &z, 1)?,
3583                    )
3584                } else {
3585                    (e.matmul(ffn_gate, &z, 1)?, e.matmul(ffn_up, &z, 1)?)
3586                };
3587                let mut act = e.zeros(n_ff)?;
3588                // step35: a DENSE FFN reads the per-layer SHEXP clamp (upstream's one `build_ffn`
3589                // serves the dense MLP and the shared expert off `swiglu_clamp_shexp` —
3590                // llama-graph.cpp:1751), resolved for the MTP block's OWN index. Every other arch
3591                // passes None, which is `ffn_act`'s dispatch verbatim.
3592                Self::ffn_act_lim(
3593                    e,
3594                    &self.cfg,
3595                    &gate,
3596                    &up,
3597                    1.0,
3598                    1.0,
3599                    mtp.step35
3600                        .as_ref()
3601                        .and_then(|s| s.clamp_shexp)
3602                        .map(SwigluClamp::Post),
3603                    &mut act,
3604                    n_ff,
3605                )?;
3606                e.matmul(ffn_down, &act, 1)?
3607            }
3608            // MTP head is a distinct block — key its experts under a separate layer index (u16::MAX)
3609            // so they never alias trunk layer 0's cache keys.
3610            crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il(e, m, &z, 1, u16::MAX)?,
3611        };
3612        anat_mark(2, e, &mut t_ph)?;
3613
3614        // op 10: h_nextn = x1 + ffn_out (at di)
3615        let mut h_inner = e.zeros(di)?;
3616        e.add(&x1, &ffn_out, &mut h_inner, di)?;
3617
3618        // op 10.5 (student): up-project the inner hidden back to n_embd — training semantics:
3619        // the chain carrier AND the head input are out_up(h_inner) (pre-final-norm).
3620        let h_nextn = match mtp.geom.as_ref() {
3621            Some(g) => e.matmul(&g.out_up, &h_inner, 1)?,
3622            None => h_inner,
3623        };
3624
3625        // op 11: final = RMSNorm(h_nextn, shared_head_norm OR output_norm)
3626        let final_norm = mtp.shared_head_norm.as_ref().unwrap_or(&self.output_norm);
3627        let mut final_h = e.zeros(n_embd)?;
3628        e.rms_norm(
3629            &h_nextn,
3630            final_norm.float_data(),
3631            &mut final_h,
3632            n_embd,
3633            1,
3634            eps,
3635        )?;
3636
3637        // op 12: draft_logits = (shared_head_head OR output) @ final — stays ON DEVICE.
3638        let head = mtp.shared_head_head.as_ref().unwrap_or(&self.output);
3639        let mut logits = e.matmul(head, &final_h, 1)?;
3640        // op 12b (lane/draft-mask): grammar mask over the DRAFT vocab, applied here so the
3641        // caller's argmax / gumbel draw / p-min prob all read the grammar-legal row.
3642        if let Some((mask_d, mw)) = mask {
3643            let d_vocab = head.out_features();
3644            e.mask_logits_col(&mut logits, mask_d, 0, d_vocab, mw)?;
3645        }
3646        anat_mark(3, e, &mut t_ph)?;
3647        if anat {
3648            ANAT_NS[4].fetch_add(t_all.elapsed().as_nanos() as u64, Relaxed);
3649            let n = ANAT_STEPS.fetch_add(1, Relaxed) + 1;
3650            if n.is_multiple_of(128) {
3651                let us = |i: usize| ANAT_NS[i].load(Relaxed) / n / 1000;
3652                eprintln!(
3653                    "[spec-anatomy] steps={n} avg us/step: glue={} attn={} ffn={} head={} total={}",
3654                    us(0),
3655                    us(1),
3656                    us(2),
3657                    us(3),
3658                    us(4)
3659                );
3660            }
3661        }
3662        // Chain recurrence hand-over: pre-norm h_nextn (default) or post-norm final_h
3663        // (MEMRA_SPEC_HPOST — llama.cpp #24025's t_h_nextn is taken AFTER the head norm).
3664        Ok((logits, if spec_hpost() { final_h } else { h_nextn }))
3665    }
3666
3667    /// One NextN/MTP draft step for an **MLA-mixer** MTP block (glm5_next class: MLA + own
3668    /// k-pool indexer + MoE, serial residual — the NextN layer carries no hc_* tensors), on
3669    /// the model `Cache`'s own MTP latent plane rather than the full-attn `MtpScratch` the
3670    /// qwen35/step35 chain uses. Gate: `glm5_mtp_head_gpu` (engine vs `memra_reference`
3671    /// `execute_mtp`, teacher-forced walk, eh_proj-transpose and h_seed-off-by-one red arms).
3672    ///
3673    /// The interface, stated precisely for the verify arc:
3674    /// - `h_seed`: `[n_embd]` f32 device — the trunk's COLLAPSED PRE-output_norm hidden of
3675    ///   the position whose next token is being drafted (MTP-PLAN §A; exactly what
3676    ///   `prime_cache`/`decode_step` return for hc models). `MEMRA_SPEC_HPOST` flips both
3677    ///   this producer and the returned carrier to the post-norm variant, same as the dev path.
3678    /// - `e_tok`: the token at the seeded position's SUCCESSOR — the token the trunk just
3679    ///   sampled/accepted (reference oracle pairing: `fused[i] = eh_proj([enorm(embed(ids[i]));
3680    ///   hnorm(trunk_hidden[i])])`, i.e. this call with `e_tok = ids[i]`, `h_seed = h[i]`,
3681    ///   `mtp_pos = i` reproduces the reference's row `i`).
3682    /// - `mtp_pos`: the absolute position this step appends to the MTP block's latent plane;
3683    ///   must equal that plane's current length (the plane advances by ONE row per call inside
3684    ///   `mla_attn_cached`; rollback on rejection = the verify arc's latent-plane len reset).
3685    /// - returns `(draft_logits [n_vocab], carrier [n_embd])` on device. glm5_next ships no
3686    ///   private MTP head, so the logits ride the trunk `lm_head` (full vocab, no d2t).
3687    pub fn mtp_head_forward_mla_cached(
3688        &self,
3689        e: &Engine,
3690        depth: usize,
3691        e_tok: u32,
3692        h_seed: &CudaSlice<f32>,
3693        cache: &mut Cache,
3694        mtp_pos: usize,
3695    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
3696        if depth >= self.mtp_head_count() {
3697            return Err(format!(
3698                "MTP depth {depth} out of range: {} embedded head(s) loaded \
3699                 (is MEMRA_GLM5_MTP=1 set for a glm5_next model?)",
3700                self.mtp_head_count()
3701            )
3702            .into());
3703        }
3704        let mtp = self.mtp_head_at(depth);
3705        let block = self
3706            .plan
3707            .mtp_blocks
3708            .get(depth)
3709            .ok_or_else(|| format!("ModelPlan declares no MTP block at depth {depth}"))?;
3710        let il = block.layer.index as usize;
3711        let Mixer::Mla(mla) = &mtp.mixer else {
3712            return Err(
3713                "mtp_head_forward_mla_cached serves MLA-mixer MTP blocks only; full-attn \
3714                 blocks take mtp_head_forward_dev's scratch path"
3715                    .into(),
3716            );
3717        };
3718        if matches!(mtp.ffn, crate::hybrid::Ffn::Dense { .. }) {
3719            return Err(
3720                "MLA-mixer MTP block with a Dense FFN has no gated arm yet (glm5_next and \
3721                 glm-dsa NextN blocks are MoE); refusing rather than running ungated math"
3722                    .into(),
3723            );
3724        }
3725        let plane_len = cache
3726            .latent
3727            .get(il)
3728            .and_then(|plane| plane.as_ref())
3729            .map(|plane| plane.len)
3730            .ok_or_else(|| {
3731                format!(
3732                    "MTP block layer {il} has no latent cache plane — the Cache must be \
3733                     built from a plan whose mtp_blocks declare StatePlan::LatentKvCache"
3734                )
3735            })?;
3736        if mtp_pos != plane_len {
3737            return Err(format!(
3738                "MTP draft position {mtp_pos} != the MTP latent plane's length {plane_len} — \
3739                 the plane advances one row per draft step and rolls back by len reset; a \
3740                 skipped or repeated position would attend the wrong horizon"
3741            )
3742            .into());
3743        }
3744
3745        let cfg = &self.cfg;
3746        let n_embd = cfg.n_embd as usize;
3747        let eps = cfg.rms_eps;
3748        let pos_d = e.htod_i32(&[mtp_pos as i32])?;
3749
3750        // Same op chain as `mtp_head_forward_dev_at` (ops 1-12), same kernels — only the
3751        // attention arm differs: `mla_attn_cached` on the plan's own MTP plane instead of
3752        // `mtp_full_attn_dc` on the MtpScratch.
3753        let e_emb = e.htod(&self.embd.try_gather(n_embd, &[e_tok])?)?;
3754        let mut e_norm = e.zeros(n_embd)?;
3755        e.rms_norm(&e_emb, mtp.enorm.float_data(), &mut e_norm, n_embd, 1, eps)?;
3756        let mut h_norm = e.zeros(n_embd)?;
3757        e.rms_norm(h_seed, mtp.hnorm.float_data(), &mut h_norm, n_embd, 1, eps)?;
3758
3759        let mut concat = e.zeros(2 * n_embd)?;
3760        e.copy_into(&mut concat, 0, &e_norm, n_embd)?;
3761        e.copy_into(&mut concat, n_embd, &h_norm, n_embd)?;
3762        let inp_sa = e.matmul(&mtp.eh_proj, &concat, 1)?;
3763
3764        let mut a_norm = e.zeros(n_embd)?;
3765        e.rms_norm(
3766            &inp_sa,
3767            mtp.attn_norm.float_data(),
3768            &mut a_norm,
3769            n_embd,
3770            1,
3771            eps,
3772        )?;
3773        let attn_out = self.mla_attn_cached(e, mla, &a_norm, &pos_d, 1, il, cache)?;
3774
3775        let mut x1 = e.zeros(n_embd)?;
3776        e.add(&inp_sa, &attn_out, &mut x1, n_embd)?;
3777        let mut z = e.zeros(n_embd)?;
3778        e.rms_norm(&x1, mtp.post_attn_norm.float_data(), &mut z, n_embd, 1, eps)?;
3779        let ffn_out = match &mtp.ffn {
3780            // Distinct block — key its experts off the trunk layers' cache keys (dev-path rule).
3781            crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il(e, m, &z, 1, u16::MAX)?,
3782            crate::hybrid::Ffn::Dense { .. } => unreachable!("refused above"),
3783        };
3784        let mut h_nextn = e.zeros(n_embd)?;
3785        e.add(&x1, &ffn_out, &mut h_nextn, n_embd)?;
3786
3787        let final_norm = mtp.shared_head_norm.as_ref().unwrap_or(&self.output_norm);
3788        let mut final_h = e.zeros(n_embd)?;
3789        e.rms_norm(
3790            &h_nextn,
3791            final_norm.float_data(),
3792            &mut final_h,
3793            n_embd,
3794            1,
3795            eps,
3796        )?;
3797        let head = mtp.shared_head_head.as_ref().unwrap_or(&self.output);
3798        let logits = e.matmul(head, &final_h, 1)?;
3799        Ok((logits, if spec_hpost() { final_h } else { h_nextn }))
3800    }
3801
3802    #[allow(clippy::too_many_arguments)]
3803    fn mtp_chain_forward_dev(
3804        &self,
3805        e: &Engine,
3806        tokens: &[u32],
3807        seeds: &[CudaSlice<f32>],
3808        scratch: &mut MtpScratch,
3809        committed_scratch_len: usize,
3810        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
3811        mask: Option<(&CudaSlice<u32>, usize)>,
3812    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
3813        if tokens.is_empty() || tokens.len() != seeds.len() {
3814            return Err("multi-head MTP prefix tokens/seeds are malformed".into());
3815        }
3816        let index = mtp_chain_head_index(tokens.len() - 1, self.mtp_head_count());
3817        let head = self.mtp_head_at(index);
3818        scratch.set_plane_len(e, index, committed_scratch_len)?;
3819
3820        let mut last = None;
3821        for row in 0..tokens.len() {
3822            let is_last = row + 1 == tokens.len();
3823            last = Some(self.mtp_head_forward_dev_at(
3824                e,
3825                head,
3826                tokens[row],
3827                &seeds[row],
3828                scratch,
3829                index,
3830                committed_scratch_len + row + 1,
3831                embd_dev,
3832                if is_last { mask } else { None },
3833            )?);
3834        }
3835        Ok(last.expect("non-empty MTP prefix produced no row"))
3836    }
3837
3838    /// step35 MTP-block attention, T=1, on the scratch KV — the EAGER-ONLY twin of
3839    /// `mtp_full_attn_dc`. Three things force a separate arm rather than a geometry parameter on
3840    /// the dc path, and all three are properties of this arch's MTP block:
3841    ///
3842    /// 1. **The SWA window.** Block 45 is an SWA-type block (`sliding_window_pattern[45]=true`,
3843    ///    window 512). Windowed decode in memra is a token-aligned VIEW OFFSET into the quantized
3844    ///    cache (the gemma4 R6 / `step35_decode_attn` pattern: keys carry absolute rope and the
3845    ///    mask is purely positional, so one query at `len-1` attending the last `win` rows IS the
3846    ///    windowed result). `fa_decode_dc` takes the key count from a DEVICE counter and always
3847    ///    starts at row 0 — it cannot express a nonzero offset. The windowed dc arm is
3848    ///    `mtp_step35_attn_dcw` (`fa_decode_dcw`, doored via MEMRA_STEP35_DRAFT_DCW —
3849    ///    default ON since lane/step37-draft-graph-serving-20260830); this host-len arm is
3850    ///    the =0 rollback and the class-ineligibility fallback.
3851    /// 2. **Per-layer head count.** 96 q heads over 8 KV (GQA 12) at this block, vs the trunk's 64
3852    ///    on its full-attn layers. The trunk cfg's `n_head` scalar is the MAX over layers, and the
3853    ///    trunk ARTIFACT's per-layer arrays stop at index 44 — so the count must come from the
3854    ///    resolved `Step35MtpGeom`, never from `cfg`.
3855    /// 3. **The separate head-wise gate.** `blk.45.attn_gate.weight [n_embd, 96]` produces one
3856    ///    sigmoid scalar per head (broadcast over head_dim) — `attn_head_gate`, not the qwen35
3857    ///    fused-into-wq `q_gate_split` form the dc arm handles.
3858    ///
3859    /// DOOR STATE: with MEMRA_STEP35_DRAFT_DCW=0 (or a sub-eligible kernel class),
3860    /// `mtp_head_forward_cap` refuses step35 heads explicitly (rather than silently capturing
3861    /// a window-less, wrong-past-`win` graph) and this eager chain IS the served path. With
3862    /// the door armed (the default), BOTH draft modes run the `mtp_step35_attn_dcw` twin
3863    /// instead of this arm.
3864    ///
3865    /// Unlike the dc arm this advances BOTH the host `kv.len` and the device counter, so the
3866    /// caller must not mirror.
3867    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
3868    fn mtp_step35_attn(
3869        &self,
3870        e: &Engine,
3871        fa: &FullAttnLayer,
3872        g: &crate::hybrid::Step35MtpGeom,
3873        h: &CudaSlice<f32>,
3874        pos_d: &CudaSlice<i32>,
3875        scratch: &mut MtpScratch,
3876        scratch_index: usize,
3877    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3878        let (nh, nkv, hd) = (g.n_head, g.n_head_kv, self.cfg.head_dim_k as usize);
3879        // MTP-GEOM RECEIPT, once per process, on the SERVED draft path. Slot-0 acceptance is
3880        // 0.725 here against 0.994 for vLLM MTP3 on the same checkpoint family and card class, and
3881        // the first three explanations for that gap were all wrong: head assignment (step-modulo
3882        // is index 0 at K=1, correct), MEMRA_SPEC_HPOST (identical 84/116 both arms), and this
3883        // block's geometry. Geometry was the one that could have failed SILENTLY — a wrong window
3884        // makes the draft attend the whole context instead of Step-3.7's 512, stays fluent, and
3885        // shows up only as acceptance — so it gets a standing receipt rather than another reading
3886        // of the source. Prints the resolved Step35MtpGeom the served path actually runs on;
3887        // `full_attention_geometry_at`'s missing-row fallback (window: None) does NOT reach here.
3888        {
3889            static ONCE: std::sync::OnceLock<()> = std::sync::OnceLock::new();
3890            ONCE.get_or_init(|| {
3891                eprintln!(
3892                    "[mtp-geom] arm=eager block={} swa={} window={} n_head={nh} n_head_kv={nkv} \
3893                     head_dim_k={hd} n_rot={} rope_base={} clamp_shexp={:?}",
3894                    g.il, g.swa, g.window, g.n_rot, g.rope_base, g.clamp_shexp,
3895                );
3896            });
3897        }
3898        let eps = self.cfg.rms_eps;
3899        let scale = 1.0 / (hd as f32).sqrt(); // step35.cpp:255 kq_scale
3900        let n_embd = self.cfg.n_embd as usize;
3901        let gw = fa
3902            .attn_gate
3903            .as_ref()
3904            .ok_or("step35 MTP block is missing attn_gate.weight (head-wise attention gate)")?;
3905
3906        let (q0, k0, v0, gt) = if e.uses_q8_1_fast(&fa.wq)
3907            && e.uses_q8_1_fast(&fa.wk)
3908            && e.uses_q8_1_fast(&fa.wv)
3909            && e.uses_q8_1_fast(gw)
3910        {
3911            let (hq, hdq) = e.quantize_q8_1(h, 1, n_embd)?;
3912            let (a, b, c) = match e.matmul_q8_fused3(&fa.wq, &fa.wk, &fa.wv, &hq, &hdq)? {
3913                Some(t3) => t3,
3914                None => (
3915                    e.matmul_pre(&fa.wq, &hq, &hdq, h, 1)?,
3916                    e.matmul_pre(&fa.wk, &hq, &hdq, h, 1)?,
3917                    e.matmul_pre(&fa.wv, &hq, &hdq, h, 1)?,
3918                ),
3919            };
3920            (a, b, c, e.matmul_pre(gw, &hq, &hdq, h, 1)?)
3921        } else {
3922            (
3923                e.matmul(&fa.wq, h, 1)?,
3924                e.matmul(&fa.wk, h, 1)?,
3925                e.matmul(&fa.wv, h, 1)?,
3926                e.matmul(gw, h, 1)?,
3927            )
3928        };
3929
3930        let mut q = e.uninit(nh * hd)?;
3931        e.rms_norm(&q0, fa.q_norm.float_data(), &mut q, hd, nh, eps)?;
3932        let mut k = e.uninit(nkv * hd)?;
3933        e.rms_norm(&k0, fa.k_norm.float_data(), &mut k, hd, nkv, eps)?;
3934        // `rope_freqs.weight` (llama3 factors) applies to the FULL-attn layers ONLY; SWA passes
3935        // null (llama-hparams / step35.cpp). Block 45 is SWA, so `ff` is None there — but read
3936        // the resolved flag, not the constant, so an all-full sibling stays correct.
3937        let ff = if g.swa {
3938            None
3939        } else {
3940            self.step35_aux.as_ref().and_then(|a| a.rope_freqs(e))
3941        };
3942        #[cfg(debug_assertions)]
3943        if let Some(ff) = ff {
3944            crate::debug_assert_tensor_stream_device(ff, &e.stream(), "mtp_step35_attn.rope_freqs");
3945        }
3946        e.rope_neox2(
3947            &mut q,
3948            &mut k,
3949            pos_d,
3950            hd,
3951            g.n_rot,
3952            nh,
3953            nkv,
3954            1,
3955            g.rope_base,
3956            1.0,
3957            ff,
3958        )?;
3959
3960        // Append at the HOST slot, then re-stamp the device counter: the eager chain has the
3961        // length on the host anyway, and the windowed view below needs it there to compute the
3962        // offset. The device counter is kept in lockstep so `mtp_kv_fill`'s `set_i32_one` and any
3963        // dc-family consumer of this scratch still agree.
3964        let (kv, scratch_cap) = scratch.plane_mut(scratch_index);
3965        assert!(
3966            kv.len < scratch_cap,
3967            "step35 MTP scratch overflow ({} >= {})",
3968            kv.len,
3969            scratch_cap
3970        );
3971        let next_len = kv.len + 1;
3972        let (off, t_kv) = if g.swa && next_len > g.window {
3973            (next_len - g.window, g.window)
3974        } else {
3975            (0, next_len)
3976        };
3977        // `off`/`t_kv` stay the ATTENTION view; the retain is a separate, lower bound so the
3978        // rewind that follows this append is still resident. THIS is the only site that rebases
3979        // this plane (MEMRA_KV_REBASE_TRACE, one run: 1 rebase, all from here), so it is the site
3980        // that decides `base` for everyone.
3981        let retain_from = match kv.ring.as_ref() {
3982            Some(ring) => memra_kv::swa_retain_from(kv.len, ring.window(), ring.base()),
3983            None => off & !31usize,
3984        };
3985        let write_row = e.prepare_kv_append(kv, retain_from, 1)?;
3986        e.append_kv_quantized(
3987            &k,
3988            &v0,
3989            &mut kv.k,
3990            &mut kv.v,
3991            write_row,
3992            kv.kv_dim_k,
3993            kv.kv_dim_v,
3994            kv.k_tok_bytes,
3995            kv.v_tok_bytes,
3996            false,
3997        )?;
3998        kv.len = next_len;
3999        e.set_i32_one(&mut kv.len_d, kv.len as i32)?;
4000        // SWA view offset (see note 1). The draft chain is short (k+2 rows), but the scratch is
4001        // PERSISTENT across rounds — `mtp_kv_fill` leaves one row per committed token behind, so
4002        // `kv.len` tracks absolute position and crosses 512 in any real generation. The window is
4003        // therefore live, not theoretical.
4004        let physical = kv.physical_rows(off, off + t_kv)?;
4005        let k_view = e.view_u8_range(
4006            &kv.k,
4007            physical.start * kv.k_tok_bytes,
4008            physical.end * kv.k_tok_bytes,
4009        );
4010        let v_view = e.view_u8_range(
4011            &kv.v,
4012            physical.start * kv.v_tok_bytes,
4013            physical.end * kv.v_tok_bytes,
4014        );
4015        let mut attn = e.uninit(nh * hd)?;
4016        e.fa_decode_kvmod(
4017            &q,
4018            &k_view,
4019            &v_view,
4020            &mut attn,
4021            hd,
4022            nh,
4023            nkv,
4024            t_kv,
4025            scale,
4026            kv.k_tok_bytes,
4027            kv.v_tok_bytes,
4028            false,
4029        )?;
4030
4031        let mut ag = e.uninit(nh * hd)?;
4032        e.attn_head_gate(&attn, &gt, &mut ag, None, hd, nh, 1)?;
4033        e.matmul(&fa.wo, &ag, 1)
4034    }
4035
4036    /// The dcw draft arm's kernel-class precondition, mirrored from `fa_decode_dcw`'s own
4037    /// refusal plus the v3 walk's format contract (`fa_v3_active`), so the DEV dispatch can
4038    /// never pick an arm the launcher would refuse mid-chain (the eager chain has no graceful
4039    /// fallback point) and the CAP site refuses with the named reason instead.
4040    ///
4041    /// `cap` = the SESSION's scratch-plane row capacity: the launcher's vec gate reads
4042    /// `bucket_max = min(window, cap)`, so a SMALL session (tiny prompt + tiny max_tokens,
4043    /// e.g. a max_tokens=8 probe: cap ~62 < the 96 vec floor) is OUTSIDE the dcw domain even
4044    /// though the WINDOW clears the floor. Mirroring the window alone shipped exactly that
4045    /// hole when the door default flipped ON (2026-08-30, vision-cell receipt: sampled
4046    /// capture WARN + `[engine-error] fa_decode_dcw supports the default v3-vec class only`
4047    /// hard-failing the burst — the eager dcw arm has no graceful fallback point). Sub-floor
4048    /// sessions now take the host-len kvmod arm, byte-for-byte the door-off serving.
4049    fn step35_dcw_eligible(&self, g: &crate::hybrid::Step35MtpGeom, cap: usize) -> bool {
4050        let hd = self.cfg.head_dim_k as usize;
4051        step35_draft_dcw_on()
4052            && g.swa
4053            && g.window.min(cap) >= crate::fa_vec_min_tkv()
4054            && std::env::var("MEMRA_NO_FA_VEC").is_err()
4055            && crate::fa_v3_active(hd)
4056            && hd <= 256
4057            && hd.is_multiple_of(32)
4058    }
4059
4060    /// step35 MTP-block attention, T=1, on the scratch KV: the WINDOWED DEVICE-COUNTER twin
4061    /// of `mtp_step35_attn`, serving BOTH draft paths when `step35_draft_dcw_on`. Write slot,
4062    /// key bound and SWA view offset all derive from device state (`len_d`, `base_d` written
4063    /// only at host-side rebases, and the block's `window`), so ONE captured graph serves the
4064    /// whole chain and replays see KV growth through the counter: the `mtp_full_attn_dc`
4065    /// contract plus the view offset the plain `_dc` kernel could not express (the old
4066    /// capture-refusal root cause). The three step35 properties stay per-geom exactly as in
4067    /// the eager twin: nh/nkv from `Step35MtpGeom`, the separate head-wise gate
4068    /// (`attn_head_gate`), per-layer rope width/base with SWA passing null freqs.
4069    ///
4070    /// bucket_max = min(cap, window): the windowed view never exceeds `window` rows, so the
4071    /// capture-time grid stays valid for every replayed len, and the kernel derives ns_eff
4072    /// from the LIVE T_kv at the fixed split_keys (one-partition law). Both arms call THIS
4073    /// launcher at THIS bucket, so eager and captured drafts are bit-identical by
4074    /// construction; vs the retired-by-flag `mtp_step35_attn` the only numeric-class deltas
4075    /// are the sub-vec-floor region (t_kv < 96: kvmod ran scalar, dcw stays vec) and any
4076    /// live-len split-ladder rung below the bucket's, both draft-side only (the verify
4077    /// arbitrates emitted bytes; acceptance is gated by the battery).
4078    ///
4079    /// Host len is NOT advanced here (graph contract); callers mirror. The EAGER caller runs
4080    /// `prepare_kv_append` per step (ring headroom, rebase legal there); the CAPTURED path
4081    /// pre-arms headroom at capture time and round start (`MtpScratch::ensure_dcw_headroom`)
4082    /// because a rebase is host work no captured chain may contain.
4083    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the capture/call contract; bundling into a struct is a refactor, not a lint fix
4084    fn mtp_step35_attn_dcw(
4085        &self,
4086        e: &Engine,
4087        fa: &FullAttnLayer,
4088        g: &crate::hybrid::Step35MtpGeom,
4089        h: &CudaSlice<f32>,
4090        pos_d: &CudaSlice<i32>,
4091        scratch: &mut MtpScratch,
4092        scratch_index: usize,
4093    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4094        let (nh, nkv, hd) = (g.n_head, g.n_head_kv, self.cfg.head_dim_k as usize);
4095        // MTP-GEOM RECEIPT (dcw twin of the `mtp_step35_attn` receipt): once per process,
4096        // naming the arm, so a serving log proves WHICH draft attention program ran (the
4097        // engagement receipt for the flag door, both directions).
4098        {
4099            static ONCE: std::sync::OnceLock<()> = std::sync::OnceLock::new();
4100            ONCE.get_or_init(|| {
4101                eprintln!(
4102                    "[mtp-geom] arm=dcw block={} swa={} window={} n_head={nh} n_head_kv={nkv} \
4103                     head_dim_k={hd} n_rot={} rope_base={} clamp_shexp={:?}",
4104                    g.il, g.swa, g.window, g.n_rot, g.rope_base, g.clamp_shexp,
4105                );
4106            });
4107        }
4108        let eps = self.cfg.rms_eps;
4109        let scale = 1.0 / (hd as f32).sqrt(); // step35.cpp:255 kq_scale
4110        let n_embd = self.cfg.n_embd as usize;
4111        let gw = fa
4112            .attn_gate
4113            .as_ref()
4114            .ok_or("step35 MTP block is missing attn_gate.weight (head-wise attention gate)")?;
4115
4116        let (q0, k0, v0, gt) = if e.uses_q8_1_fast(&fa.wq)
4117            && e.uses_q8_1_fast(&fa.wk)
4118            && e.uses_q8_1_fast(&fa.wv)
4119            && e.uses_q8_1_fast(gw)
4120        {
4121            let (hq, hdq) = e.quantize_q8_1(h, 1, n_embd)?;
4122            let (a, b, c) = match e.matmul_q8_fused3(&fa.wq, &fa.wk, &fa.wv, &hq, &hdq)? {
4123                Some(t3) => t3,
4124                None => (
4125                    e.matmul_pre(&fa.wq, &hq, &hdq, h, 1)?,
4126                    e.matmul_pre(&fa.wk, &hq, &hdq, h, 1)?,
4127                    e.matmul_pre(&fa.wv, &hq, &hdq, h, 1)?,
4128                ),
4129            };
4130            (a, b, c, e.matmul_pre(gw, &hq, &hdq, h, 1)?)
4131        } else {
4132            (
4133                e.matmul(&fa.wq, h, 1)?,
4134                e.matmul(&fa.wk, h, 1)?,
4135                e.matmul(&fa.wv, h, 1)?,
4136                e.matmul(gw, h, 1)?,
4137            )
4138        };
4139
4140        let mut q = e.zeros(nh * hd)?;
4141        e.rms_norm(&q0, fa.q_norm.float_data(), &mut q, hd, nh, eps)?;
4142        let mut k = e.zeros(nkv * hd)?;
4143        e.rms_norm(&k0, fa.k_norm.float_data(), &mut k, hd, nkv, eps)?;
4144        // rope_freqs (llama3 factors) apply to the FULL-attn layers ONLY; SWA passes null
4145        // (the eager twin's rule, resolved from the flag, not the constant).
4146        let ff = if g.swa {
4147            None
4148        } else {
4149            self.step35_aux.as_ref().and_then(|a| a.rope_freqs(e))
4150        };
4151        #[cfg(debug_assertions)]
4152        if let Some(ff) = ff {
4153            crate::debug_assert_tensor_stream_device(
4154                ff,
4155                &e.stream(),
4156                "mtp_step35_attn_dcw.rope_freqs",
4157            );
4158        }
4159        e.rope_neox2(
4160            &mut q,
4161            &mut k,
4162            pos_d,
4163            hd,
4164            g.n_rot,
4165            nh,
4166            nkv,
4167            1,
4168            g.rope_base,
4169            1.0,
4170            ff,
4171        )?;
4172
4173        let (kv, cap) = scratch.plane_mut(scratch_index);
4174        // Append at the DEVICE slot's PHYSICAL row (len_d - base_d), then advance the counter
4175        // in-graph. Physical room is the callers' headroom contract (see the fn doc).
4176        e.append_kv_quantized_dcw(
4177            &k,
4178            &v0,
4179            &mut kv.k,
4180            &mut kv.v,
4181            &kv.len_d,
4182            kv.base_d.as_ref(),
4183            kv.kv_dim_k,
4184            kv.kv_dim_v,
4185            kv.k_tok_bytes,
4186            kv.v_tok_bytes,
4187        )?;
4188        e.inc_seqlen(&mut kv.len_d)?;
4189        // Full-buffer views (any in-round physical row stays in range under the headroom
4190        // contract); the kernel bounds and offsets the key range from (len_d, base_d, window).
4191        let k_view = e.view_u8(&kv.k, kv.k.len());
4192        let v_view = e.view_u8(&kv.v, kv.v.len());
4193        let bucket = g.window.min(cap);
4194        let mut attn = e.zeros(nh * hd)?;
4195        e.fa_decode_dcw(
4196            &q,
4197            &k_view,
4198            &v_view,
4199            &mut attn,
4200            hd,
4201            nh,
4202            nkv,
4203            &kv.len_d,
4204            kv.base_d.as_ref(),
4205            if g.swa { g.window } else { 0 },
4206            bucket,
4207            scale,
4208            kv.k_tok_bytes,
4209            kv.v_tok_bytes,
4210            None,
4211        )?;
4212
4213        let mut ag = e.zeros(nh * hd)?;
4214        e.attn_head_gate(&attn, &gt, &mut ag, None, hd, nh, 1)?;
4215        e.matmul(&fa.wo, &ag, 1)
4216    }
4217
4218    /// MTP-block full attention, T=1, on the scratch KV (BOTH draft paths — eager and graph):
4219    /// the scratch write slot and the attention bound come from `scratch.kv.len_d` (device i32[1])
4220    /// so the launch args are FIXED across draft steps — ONE captured graph serves the whole
4221    /// chain, and replays keep seeing KV growth through the device counter (no recapture).
4222    /// Geometry contract: n_splits is sized from `scratch.cap` (the persistent capacity); splits
4223    /// whose key range lies beyond the device t_kv exit empty and the shared combine skips them
4224    /// (fa_decode_dc bit-correct-for-any-t_kv<=bucket_max contract). The eager path uses the SAME
4225    /// launcher with the SAME bucket_max -> identical dispatch -> bit-identical draft tokens (the
4226    /// graph-vs-eager parity gate). Host len is NOT advanced here (graph contract); callers mirror.
4227    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
4228    fn mtp_full_attn_dc(
4229        &self,
4230        e: &Engine,
4231        fa: &FullAttnLayer,
4232        h: &CudaSlice<f32>,
4233        pos_d: &CudaSlice<i32>,
4234        scratch: &mut MtpScratch,
4235        scratch_index: usize,
4236        geom: Option<&crate::hybrid::DraftGeom>,
4237    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4238        let cfg = &self.cfg;
4239        let mtp_il = cfg.n_layer.saturating_sub(cfg.nextn_predict_layers);
4240        let geometry = cfg.full_attention_geometry_at(mtp_il);
4241        let n_head = geom.map(|g| g.n_head).unwrap_or(geometry.n_head as usize);
4242        let n_head_kv = geom
4243            .map(|g| g.n_head_kv)
4244            .unwrap_or(geometry.n_head_kv as usize);
4245        let head_dim = geometry.head_dim_k as usize;
4246        let eps = cfg.rms_eps;
4247        let scale = geometry.attention_scale();
4248        let n_embd = geom.map(|g| g.d_inner).unwrap_or(cfg.n_embd as usize);
4249        let bucket_max = scratch.plane(scratch_index).1;
4250
4251        let (qf, mut k, v) =
4252            if e.uses_q8_1_fast(&fa.wq) && e.uses_q8_1_fast(&fa.wk) && e.uses_q8_1_fast(&fa.wv) {
4253                let (hq, hd) = e.quantize_q8_1(h, 1, n_embd)?;
4254                (
4255                    e.matmul_pre(&fa.wq, &hq, &hd, h, 1)?,
4256                    e.matmul_pre(&fa.wk, &hq, &hd, h, 1)?,
4257                    e.matmul_pre(&fa.wv, &hq, &hd, h, 1)?,
4258                )
4259            } else {
4260                (
4261                    e.matmul(&fa.wq, h, 1)?,
4262                    e.matmul(&fa.wk, h, 1)?,
4263                    e.matmul(&fa.wv, h, 1)?,
4264                )
4265            };
4266        // M3/Hy3 have no attention output gate — wq out is exactly q; skip the split.
4267        let gated = geometry.attention_gate == memra_gguf::config::AttentionGateKind::FusedQ;
4268        let (mut q, gate) = if gated {
4269            let mut q = e.zeros(n_head * head_dim)?;
4270            let mut gate = e.zeros(n_head * head_dim)?;
4271            e.q_gate_split(&qf, &mut q, &mut gate, head_dim, n_head, 1)?;
4272            (q, Some(gate))
4273        } else {
4274            (qf, None)
4275        };
4276
4277        let mut qn = e.zeros(n_head * head_dim)?;
4278        e.rms_norm(&q, fa.q_norm.float_data(), &mut qn, head_dim, n_head, eps)?;
4279        q = qn;
4280        let mut kn = e.zeros(n_head_kv * head_dim)?;
4281        e.rms_norm(
4282            &k,
4283            fa.k_norm.float_data(),
4284            &mut kn,
4285            head_dim,
4286            n_head_kv,
4287            eps,
4288        )?;
4289        k = kn;
4290        let rope_dims = geometry.n_rot as usize;
4291        e.rope_neox(
4292            &mut q,
4293            pos_d,
4294            head_dim,
4295            rope_dims,
4296            n_head,
4297            1,
4298            geometry.rope_base,
4299            1.0,
4300        )?;
4301        e.rope_neox(
4302            &mut k,
4303            pos_d,
4304            head_dim,
4305            rope_dims,
4306            n_head_kv,
4307            1,
4308            geometry.rope_base,
4309            1.0,
4310        )?;
4311
4312        let kv = scratch.plane_mut(scratch_index).0;
4313        // append at the DEVICE slot (kv.len_d == old len), then advance the counter in-graph.
4314        e.append_kv_quantized_dc(
4315            &k,
4316            &v,
4317            &mut kv.k,
4318            &mut kv.v,
4319            &kv.len_d,
4320            kv.kv_dim_k,
4321            kv.kv_dim_v,
4322            kv.k_tok_bytes,
4323            kv.v_tok_bytes,
4324            false,
4325        )?;
4326        e.inc_seqlen(&mut kv.len_d)?;
4327        // full-buffer views (any in-round t_kv stays in range on replay); the kernel bounds the
4328        // key range from the device counter.
4329        let k_view = e.view_u8(&kv.k, kv.k.len());
4330        let v_view = e.view_u8(&kv.v, kv.v.len());
4331        let (ktb, vtb) = (kv.k_tok_bytes, kv.v_tok_bytes);
4332        let mut attn = e.zeros(n_head * head_dim)?;
4333        e.fa_decode_dc(
4334            &q, &k_view, &v_view, &mut attn, head_dim, n_head, n_head_kv, &kv.len_d, bucket_max,
4335            scale, ktb, vtb, false,
4336        )?;
4337
4338        let attn_g = match &gate {
4339            Some(gate) => {
4340                let mut gsig = e.zeros(n_head * head_dim)?;
4341                e.sigmoid(gate, &mut gsig, n_head * head_dim)?;
4342                let mut ag = e.zeros(n_head * head_dim)?;
4343                e.mul(&attn, &gsig, &mut ag, n_head * head_dim)?;
4344                ag
4345            }
4346            None => attn,
4347        };
4348        e.matmul(&fa.wo, &attn_g, 1)
4349    }
4350
4351    /// PERSISTENT-DRAFT-KV fill (the reference engine's "mtp_update" analogue): compute the MTP
4352    /// block's K/V for `tokens` (committed tokens at positions pos0..pos0+T) from their EXACT
4353    /// trunk hiddens `h` ([T, n_embd] token-major, pre-output_norm) and append at slots pos0.. of
4354    /// the scratch KV. K/V-ONLY — ops A/1-5 plus the K-side of op 6 (wk/wv + k_norm + rope +
4355    /// quantized append); no wq/attention/FFN/lm_head, so per-token cost ~= eh_proj + wk/wv (a
4356    /// small fraction of one trunk layer), T-batched. Rope follows the chain convention
4357    /// rope(token@p) = p+1. Runs at round boundaries OUTSIDE the captured graph in BOTH draft
4358    /// modes -> draft parity by construction. Caller must have scratch.kv.len == pos0.
4359    #[allow(clippy::too_many_arguments)]
4360    fn mtp_kv_fill_at(
4361        &self,
4362        e: &Engine,
4363        mtp: &MtpHead,
4364        tokens: &[u32],
4365        h: &CudaSlice<f32>,
4366        pos0: usize,
4367        scratch: &mut MtpScratch,
4368        scratch_index: usize,
4369        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
4370    ) -> Result<(), Box<dyn std::error::Error>> {
4371        let cfg = &self.cfg;
4372        let n_embd = cfg.n_embd as usize;
4373        let eps = cfg.rms_eps;
4374        let t = tokens.len();
4375        let (scratch_kv, scratch_cap) = scratch.plane(scratch_index);
4376        assert_eq!(scratch_kv.len, pos0, "mtp_kv_fill: append slot mismatch");
4377        assert!(pos0 + t <= scratch_cap, "mtp_kv_fill: scratch overflow");
4378        let Mixer::Full(fa) = &mtp.mixer else {
4379            panic!("MTP block is full-attn in qwen35; linear MTP not supported")
4380        };
4381        let pos_vec: Vec<i32> = (0..t).map(|i| (pos0 + i + 1) as i32).collect();
4382        let pos_d = e.htod_i32(&pos_vec)?;
4383
4384        // ops A/1/2: embed + the two input norms, T-wide.
4385        let e_emb = match embd_dev {
4386            Some((g, qt, rb)) => e.embed_gather_device_t(g, tokens, n_embd, qt, rb)?,
4387            None => e.htod(&self.embd.try_gather(n_embd, tokens)?)?,
4388        };
4389        let mut e_norm = e.zeros(t * n_embd)?;
4390        e.rms_norm(&e_emb, mtp.enorm.float_data(), &mut e_norm, n_embd, t, eps)?;
4391        let mut h_norm = e.zeros(t * n_embd)?;
4392        e.rms_norm(h, mtp.hnorm.float_data(), &mut h_norm, n_embd, t, eps)?;
4393
4394        // op 3: per-row [e_norm ; h_norm] concat, token-major [T, 2*n_embd].
4395        let mut concat = e.zeros(t * 2 * n_embd)?;
4396        for i in 0..t {
4397            e.copy_view_into(
4398                &mut concat,
4399                i * 2 * n_embd,
4400                &e_norm.slice(i * n_embd..(i + 1) * n_embd),
4401                n_embd,
4402            )?;
4403            e.copy_view_into(
4404                &mut concat,
4405                i * 2 * n_embd + n_embd,
4406                &h_norm.slice(i * n_embd..(i + 1) * n_embd),
4407                n_embd,
4408            )?;
4409        }
4410
4411        // ops 4/5: eh_proj + attn_norm, T-wide (at the student inner width when geom is set).
4412        let di = mtp.geom.as_ref().map(|g| g.d_inner).unwrap_or(n_embd);
4413        let inp_sa = e.matmul(&mtp.eh_proj, &concat, t)?;
4414        let mut a_norm = e.zeros(t * di)?;
4415        e.rms_norm(&inp_sa, mtp.attn_norm.float_data(), &mut a_norm, di, t, eps)?;
4416
4417        // op 6 (K/V half): wk/wv + k_norm + rope + per-row quantized append. No wq/attention —
4418        // the fill only has to leave correct K/V rows behind for later chains to attend over.
4419        let n_head_kv = mtp
4420            .geom
4421            .as_ref()
4422            .map(|g| g.n_head_kv)
4423            .or(mtp.step35.as_ref().map(|s| s.n_head_kv))
4424            .unwrap_or_else(|| {
4425                let mtp_il = cfg.n_layer.saturating_sub(cfg.nextn_predict_layers);
4426                cfg.full_attention_geometry_at(mtp_il).n_head_kv as usize
4427            });
4428        let mtp_il = cfg.n_layer.saturating_sub(cfg.nextn_predict_layers);
4429        let geometry = cfg.full_attention_geometry_at(mtp_il);
4430        let head_dim = geometry.head_dim_k as usize;
4431        let mut k = e.matmul(&fa.wk, &a_norm, t)?;
4432        let v = e.matmul(&fa.wv, &a_norm, t)?;
4433        let mut kn = e.zeros(t * n_head_kv * head_dim)?;
4434        e.rms_norm(
4435            &k,
4436            fa.k_norm.float_data(),
4437            &mut kn,
4438            head_dim,
4439            n_head_kv * t,
4440            eps,
4441        )?;
4442        k = kn;
4443        // step35: rotary width AND base are per-layer, and the MTP block's values come from the
4444        // resolved `Step35MtpGeom` — NOT from `cfg.rope_dim_count`/`cfg.rope_freq_base`, which
4445        // carry the arch defaults (128 / 5e6, i.e. the FULL-attn layers' base). Getting this wrong
4446        // writes K rows the attention arm then re-derives at a different theta: correct-looking
4447        // output with dead acceptance, invisible to the exactness gates.
4448        let (rope_dims, rope_base, ff) = match mtp.step35.as_ref() {
4449            Some(s) => (
4450                s.n_rot,
4451                s.rope_base,
4452                if s.swa {
4453                    None
4454                } else {
4455                    self.step35_aux.as_ref().and_then(|a| a.rope_freqs(e))
4456                },
4457            ),
4458            None => (geometry.n_rot as usize, geometry.rope_base, None),
4459        };
4460        #[cfg(debug_assertions)]
4461        if let Some(ff) = ff {
4462            crate::debug_assert_tensor_stream_device(ff, &e.stream(), "mtp_kv_fill.rope_freqs");
4463        }
4464        match ff {
4465            Some(f) => e.rope_neox_ff(
4466                &mut k, &pos_d, head_dim, rope_dims, n_head_kv, t, rope_base, 1.0, f,
4467            )?,
4468            None => e.rope_neox(
4469                &mut k, &pos_d, head_dim, rope_dims, n_head_kv, t, rope_base, 1.0,
4470            )?,
4471        }
4472
4473        let kv = scratch.plane_mut(scratch_index).0;
4474        // Match the trunk prime contract: a chunk may need the aligned window immediately before
4475        // its first row, so preserve that prefix when the physical tail rebases at wrap.
4476        let retain_from = kv
4477            .ring
4478            .as_ref()
4479            .map(|ring| memra_kv::swa_retain_from(pos0, ring.window(), ring.base()))
4480            .unwrap_or(0);
4481        let write_row = e.prepare_kv_append(kv, retain_from, t)?;
4482        for i in 0..t {
4483            let k_row = k.slice(i * kv.kv_dim_k..(i + 1) * kv.kv_dim_k);
4484            let v_row = v.slice(i * kv.kv_dim_v..(i + 1) * kv.kv_dim_v);
4485            e.append_kv_quantized_view(
4486                &k_row,
4487                &v_row,
4488                &mut kv.k,
4489                &mut kv.v,
4490                write_row + i,
4491                kv.kv_dim_k,
4492                kv.kv_dim_v,
4493                kv.k_tok_bytes,
4494                kv.v_tok_bytes,
4495                false,
4496            )?;
4497        }
4498        kv.len = pos0 + t;
4499        e.set_i32_one(&mut kv.len_d, kv.len as i32)?;
4500        Ok(())
4501    }
4502
4503    #[allow(clippy::too_many_arguments)]
4504    fn mtp_kv_fill_all(
4505        &self,
4506        e: &Engine,
4507        tokens: &[u32],
4508        h: &CudaSlice<f32>,
4509        pos0: usize,
4510        scratch: &mut MtpScratch,
4511        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
4512    ) -> Result<(), Box<dyn std::error::Error>> {
4513        debug_assert_eq!(self.mtp_head_count(), scratch.plane_count());
4514        for index in 0..self.mtp_head_count() {
4515            self.mtp_kv_fill_at(
4516                e,
4517                self.mtp_head_at(index),
4518                tokens,
4519                h,
4520                pos0,
4521                scratch,
4522                index,
4523                embd_dev,
4524            )?;
4525        }
4526        Ok(())
4527    }
4528
4529    /// CAPTURE body for the GRAPH DRAFT (stage 2 of graph-grade spec): ONE MTP head forward with
4530    /// every varying input device-resident —
4531    ///   - token id from the persistent `tok_d` (the previous replay's in-graph argmax wrote it,
4532    ///     so the chain feeds itself; the host reads the same 4 bytes for the draft list),
4533    ///   - h_seed from the persistent `h_seed_d` (h_nextn is copied BACK into it at the end),
4534    ///   - rope pos from the persistent `pos_d` counter (inc'd in-graph),
4535    ///   - scratch KV slot/bound from `scratch.kv.len_d` (see mtp_full_attn_dc).
4536    ///     The p-min confidence lands in the persistent `p_d` iff `with_prob` (env is fixed per run).
4537    ///     Same kernels, same dispatch as the eager mtp_head_forward_dev chain -> same draft tokens
4538    ///     (exactness never depends on drafts — the verify arbitrates — but acceptance parity does).
4539    ///     `with_head=false` captures the HEAD-LESS twin for the pseudo-seed replay (2026-07-03):
4540    ///     the pseudo pass only needs h_nextn (op 10) + the scratch append — the lm_head read
4541    ///     (~1.06ms q6_K on the 9B), argmax and prob are dead weight there. h_nextn's inputs are
4542    ///     untouched, so the seed value is identical; round-start resets overwrite tok_d/p_d anyway.
4543    ///     `sampled_cap` = Some((ctr_d, perturb_d, q_out_d, seed, temp)) captures the SAMPLED twin
4544    ///     (step 3 of the sampled-spec arc): head logits are retained in the persistent `q_out_d`
4545    ///     (host D2Ds them to the round's q slot after each replay), the DEVICE event counter is
4546    ///     bumped in-graph, and the argmax reads GUMBEL-PERTURBED logits — one categorical draw per
4547    ///     replay, bit-identical to the eager arm's gumbel_perturb at the same (seed, sctr, temp).
4548    ///     seed/temp are capture-time constants (fixed per generate call, like p_min).
4549    #[allow(clippy::too_many_arguments)]
4550    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
4551    fn mtp_head_forward_cap(
4552        &self,
4553        e: &Engine,
4554        mtp: &MtpHead,
4555        tok_d: &mut CudaSlice<u32>,
4556        pos_d: &mut CudaSlice<i32>,
4557        h_seed_d: &mut CudaSlice<f32>,
4558        p_d: &mut CudaSlice<f32>,
4559        scratch: &mut MtpScratch,
4560        // Which scratch plane this head appends to / attends over: 0 for the single-head
4561        // chain (every pre-lane caller), the head's own plane index for the multi-head
4562        // chain graphs (each head owns one plane — `mtp_chain_forward_dev`'s contract).
4563        scratch_index: usize,
4564        with_prob: bool,
4565        with_head: bool,
4566        embd_gpu: &CudaSlice<u8>,
4567        embd_qt: i32,
4568        embd_rb: usize,
4569        d_vocab: usize,
4570        sampled_cap: Option<SampledCapArgs<'_>>,
4571        stream_pack: Option<(&mut CudaSlice<u32>, usize, Option<&CudaSlice<u32>>)>,
4572        // DRAFT-SIDE GRAMMAR MASK (lane/draft-mask): (packed draft-vocab allowed-set buffer,
4573        // word count). Captured as ONE mask_logits_f32 node between the head matmul and the
4574        // in-graph argmax; the buffer address is baked, its CONTENTS are re-uploaded by the
4575        // host before every replay (the decode.rs graph-mask pattern). All-ones contents = a
4576        // no-op ban, so a position the grammar cannot constrain costs one pass over the row.
4577        mask_cap: Option<(&CudaSlice<u32>, usize)>,
4578    ) -> Result<(), Box<dyn std::error::Error>> {
4579        let cfg = &self.cfg;
4580        let n_embd = cfg.n_embd as usize;
4581        // step35: capturable through the WINDOWED device-counter arm (`mtp_step35_attn_dcw`)
4582        // once the dcw door is armed and the v3-vec class is live. Without the door this stays
4583        // the deliberate, named refusal: the plain `_dc` attention's key bound always starts at
4584        // row 0, cannot express this block's SWA view offset, and a captured chain would
4585        // silently attend OUTSIDE the window once the persistent scratch passes 512 rows.
4586        // Returning Err (not a panic) is what the capture sites already handle by degrading to
4587        // the eager chain (`mtp_head_forward_dev` -> `mtp_step35_attn`).
4588        // ROUND-STREAM stays refused EITHER WAY: the stream VERIFY has no step35 twin (see the
4589        // step35_verify refusal), so a stream capture that succeeded here would only move the
4590        // failure from capture time (graceful stream-off) to serve time (a failed round).
4591        if let Some(g) = mtp.step35.as_ref() {
4592            if stream_pack.is_some() {
4593                return Err(
4594                    "step35 has no ROUND-STREAM draft arm (the stream verify has no step35 \
4595                     twin); stream off"
4596                        .into(),
4597                );
4598            }
4599            if !self.step35_dcw_eligible(g, scratch.plane(scratch_index).1) {
4600                return Err(format!(
4601                    "step35 has no captured draft chain (fa_decode_dc cannot express the MTP \
4602                        block's SWA view offset; the windowed dcw capture needs \
4603                        MEMRA_STEP35_DRAFT_DCW armed [default ON, =0 disarms] and the v3-vec \
4604                        class live at bucket=min(window {}, scratch cap {})) - the eager draft \
4605                        chain serves this shape",
4606                    g.window,
4607                    scratch.plane(scratch_index).1,
4608                )
4609                .into());
4610            }
4611        }
4612        // student inner width (see mtp_head_forward_dev) — interface dims stay n_embd.
4613        let di = mtp.geom.as_ref().map(|g| g.d_inner).unwrap_or(n_embd);
4614        let eps = cfg.rms_eps;
4615        let e_emb = e.embed_gather_device(embd_gpu, tok_d, n_embd, embd_qt, embd_rb)?;
4616        let mut e_norm = e.zeros(n_embd)?;
4617        e.rms_norm(&e_emb, mtp.enorm.float_data(), &mut e_norm, n_embd, 1, eps)?;
4618        let mut h_norm = e.zeros(n_embd)?;
4619        e.rms_norm(
4620            &*h_seed_d,
4621            mtp.hnorm.float_data(),
4622            &mut h_norm,
4623            n_embd,
4624            1,
4625            eps,
4626        )?;
4627        let mut concat = e.zeros(2 * n_embd)?;
4628        e.copy_into(&mut concat, 0, &e_norm, n_embd)?;
4629        e.copy_into(&mut concat, n_embd, &h_norm, n_embd)?;
4630        let inp_sa = e.matmul(&mtp.eh_proj, &concat, 1)?;
4631        let mut a_norm = e.zeros(di)?;
4632        e.rms_norm(&inp_sa, mtp.attn_norm.float_data(), &mut a_norm, di, 1, eps)?;
4633        let attn_out = match (&mtp.mixer, mtp.step35.as_ref()) {
4634            // step35 (eligibility already enforced by the refusal above): the windowed dcw
4635            // arm, the SAME launcher the eager dev arm runs when the door is armed. No host
4636            // work here (this is the capture body); headroom is the callers' pre-arm.
4637            (Mixer::Full(fa), Some(g)) => {
4638                self.mtp_step35_attn_dcw(e, fa, g, &a_norm, pos_d, scratch, scratch_index)?
4639            }
4640            (Mixer::Full(fa), None) => self.mtp_full_attn_dc(
4641                e,
4642                fa,
4643                &a_norm,
4644                pos_d,
4645                scratch,
4646                scratch_index,
4647                mtp.geom.as_ref(),
4648            )?,
4649            (Mixer::Linear(_), _) => {
4650                panic!("MTP block is full-attn in qwen35; linear MTP not supported")
4651            }
4652            (Mixer::Mla(_), _) => {
4653                crate::hybrid::mla_path_unimplemented("captured MTP head forward")
4654            }
4655            (Mixer::Kda(_), _) => {
4656                crate::hybrid::kda_path_unimplemented("captured MTP head forward")
4657            }
4658        };
4659        let mut x1 = e.zeros(di)?;
4660        e.add(&inp_sa, &attn_out, &mut x1, di)?;
4661        let mut z = e.zeros(di)?;
4662        e.rms_norm(&x1, mtp.post_attn_norm.float_data(), &mut z, di, 1, eps)?;
4663        let ffn_out = match &mtp.ffn {
4664            crate::hybrid::Ffn::Dense {
4665                ffn_gate,
4666                ffn_up,
4667                ffn_down,
4668            } => {
4669                let n_ff = ffn_gate.out_features();
4670                let (gate, up) = if e.uses_q8_1_fast(ffn_gate) && e.uses_q8_1_fast(ffn_up) {
4671                    let (zq, zd) = e.quantize_q8_1(&z, 1, di)?;
4672                    (
4673                        e.matmul_pre(ffn_gate, &zq, &zd, &z, 1)?,
4674                        e.matmul_pre(ffn_up, &zq, &zd, &z, 1)?,
4675                    )
4676                } else {
4677                    (e.matmul(ffn_gate, &z, 1)?, e.matmul(ffn_up, &z, 1)?)
4678                };
4679                let mut act = e.zeros(n_ff)?;
4680                // step35: the dense FFN reads the per-layer SHEXP clamp, resolved for the MTP
4681                // block's own index (the mtp_head_forward_dev rule; None for every other arch,
4682                // which is `ffn_act`'s dispatch verbatim). The eager and captured chains must
4683                // run the ONE activation program.
4684                Self::ffn_act_lim(
4685                    e,
4686                    &self.cfg,
4687                    &gate,
4688                    &up,
4689                    1.0,
4690                    1.0,
4691                    mtp.step35
4692                        .as_ref()
4693                        .and_then(|s| s.clamp_shexp)
4694                        .map(SwigluClamp::Post),
4695                    &mut act,
4696                    n_ff,
4697                )?;
4698                e.matmul(ffn_down, &act, 1)?
4699            }
4700            // ROUND-STREAM: a softmax-routed resident MoE takes the zero-D2H device router +
4701            // expert program and is capture-legal. Sigmoid-routed MoE (Hy3/M3/Step) still
4702            // selects through the host-visible sigmoid router; capturing that stream sync
4703            // invalidates CUDA capture, so it stays on the eager draft chain even when every
4704            // expert is resident. Non-resident (SLRU-lock) is likewise rejected.
4705            crate::hybrid::Ffn::Moe(m)
4706                if m.dev_exps.is_some() && self.cfg.sigmoid_router().is_none() =>
4707            {
4708                self.moe_ffn_il(e, m, &z, 1, u16::MAX)?
4709            }
4710            crate::hybrid::Ffn::Moe(_) => {
4711                return Err(
4712                    "graph draft requires a Dense or device-routed resident-MoE MTP FFN".into(),
4713                );
4714            }
4715        };
4716        let mut h_inner = e.zeros(di)?;
4717        e.add(&x1, &ffn_out, &mut h_inner, di)?;
4718        // student: up-project back to n_embd (carrier + head input; see mtp_head_forward_dev).
4719        let h_nextn = match mtp.geom.as_ref() {
4720            Some(g) => e.matmul(&g.out_up, &h_inner, 1)?,
4721            None => h_inner,
4722        };
4723        // MEMRA_SPEC_HPOST needs final_h even head-less (it IS the next seed under that convention).
4724        let final_h = if with_head || spec_hpost() {
4725            let final_norm = mtp.shared_head_norm.as_ref().unwrap_or(&self.output_norm);
4726            let mut fh = e.zeros(n_embd)?;
4727            e.rms_norm(&h_nextn, final_norm.float_data(), &mut fh, n_embd, 1, eps)?;
4728            Some(fh)
4729        } else {
4730            None
4731        };
4732        if with_head {
4733            let head = mtp.shared_head_head.as_ref().unwrap_or(&self.output);
4734            let mut logits = e.matmul(head, final_h.as_ref().unwrap(), 1)?;
4735            // DRAFT-SIDE GRAMMAR MASK: ban the grammar-illegal draft ids IN the captured chain,
4736            // before the argmax — proposals become legal by construction. Contents-only
4737            // per-replay upload keeps the capture valid.
4738            if let Some((mask_d, mw)) = mask_cap {
4739                e.mask_logits_col(&mut logits, mask_d, 0, d_vocab, mw)?;
4740            }
4741            if let Some(SampledCapArgs {
4742                ctr: ctr_d,
4743                perturb: perturb_d,
4744                q_out: q_out_d,
4745                seed,
4746                temp,
4747                filt,
4748            }) = sampled_cap
4749            {
4750                // SAMPLED chain: retain q (raw head logits -> persistent q_out_d; the matmul's
4751                // own buffer is pool-recycled after the capture body returns, so it can't be the
4752                // retention target), bump the device event counter, gumbel-perturb reading it,
4753                // and argmax the PERTURBED logits into tok_d — the in-graph categorical draw.
4754                e.copy_into(q_out_d, 0, &logits, d_vocab)?;
4755                e.sctr_inc(ctr_d)?;
4756                match filt {
4757                    // PURE-TEMP: gumbel over the raw softmax — byte-identical to the
4758                    // pre-lane capture body.
4759                    None => e.gumbel_perturb_ctr(&logits, perturb_d, d_vocab, seed, ctr_d, temp)?,
4760                    // FILTERED (lane/step37-draft-graph-serving-20260830): the SAME
4761                    // filter_stats program the eager arm and the accept path run (the
4762                    // wrapper's coop/plain choice is deployment-keyed, never per-call), then
4763                    // the device-stat/device-counter perturb twin — the draft draws from the
4764                    // exact filtered distribution the verify gathers `q` from. q was
4765                    // retained ABOVE, pre-perturb, so the accept path's post-replay stats
4766                    // recompute (same kernel, same bits) reconstructs these th/z exactly.
4767                    Some(f) => {
4768                        e.filter_stats(
4769                            &logits, d_vocab, f.rows0, f.th, f.z, f.mx, d_vocab, 1, temp, f.top_k,
4770                            f.top_p, f.min_p,
4771                        )?;
4772                        e.gumbel_perturb_filtered_ctr(
4773                            &logits, perturb_d, d_vocab, seed, ctr_d, temp, f.mx, f.th,
4774                        )?;
4775                    }
4776                }
4777                e.argmax_token_device_into(perturb_d, tok_d, d_vocab)?;
4778                // p-min prob = the head's RAW softmax confidence in the SAMPLED pick — same
4779                // semantics as the eager sampled arm's prob_of_token_device(dl_d, tok_d).
4780                if with_prob {
4781                    e.prob_of_token_device_into(&logits, tok_d, p_d, d_vocab)?;
4782                }
4783            } else {
4784                // draft token -> persistent tok_d (next replay's embed reads it; host reads the 4 bytes).
4785                e.argmax_token_device_into(&logits, tok_d, d_vocab)?;
4786                // p-min under a draft mask reads the MASKED row: confidence relative to the
4787                // grammar-LEGAL alternatives (illegal ids leave the softmax denominator), which
4788                // is the right semantics for "does the drafter know what comes next here" and
4789                // the same row the pick came from. Draft-quality only — verify arbitrates.
4790                if with_prob {
4791                    e.prob_of_token_device_into(&logits, tok_d, p_d, d_vocab)?;
4792                }
4793            }
4794        }
4795        // ROUND-STREAM K-chain: pack (tok, p) into slot j, then remap tok through d2t so the
4796        // NEXT chained body's embed reads the TARGET id — zero host involvement per step.
4797        if let Some((out, slot, d2t)) = stream_pack {
4798            e.pack_tok_p(tok_d, p_d, out, slot)?;
4799            if let Some(map) = d2t {
4800                e.tok_map_u32(tok_d, map)?;
4801            }
4802        }
4803        // Next draft step's h_seed: pre-norm h_nextn (default) or post-norm final_h (HPOST).
4804        if spec_hpost() {
4805            e.copy_into(h_seed_d, 0, final_h.as_ref().unwrap(), n_embd)?;
4806        } else {
4807            e.copy_into(h_seed_d, 0, &h_nextn, n_embd)?;
4808        }
4809        // advance the draft rope position in-graph.
4810        e.inc_seqlen(pos_d)?;
4811        Ok(())
4812    }
4813
4814    /// Batched target verify forward over `tokens` at positions `pos0..pos0+T` (§D.3, T=K+1).
4815    /// Returns ALL T logit columns (host f32, [T*n_vocab]); appends T cols to every full-attn KV
4816    /// and advances every linear-attn recur state by T steps (the recur steps are SEQUENTIAL T=1).
4817    /// Advances `cache.pos` by T.
4818    pub fn decode_step_t(
4819        &self,
4820        e: &Engine,
4821        tokens: &[u32],
4822        pos0: usize,
4823        cache: &mut Cache,
4824    ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
4825        if self.is_gemma4_e4b() {
4826            return Ok(self.gemma4_e4b_decode_step_t_h(e, tokens, pos0, cache)?.0);
4827        }
4828        if self.gemma_batch_program() {
4829            return self.gemma4_decode_step_t(e, tokens, pos0, cache);
4830        }
4831        Ok(self.decode_step_t_h(e, tokens, pos0, cache)?.0)
4832    }
4833
4834    /// Like `decode_step_t` but ALSO returns the LAST column's pre-output_norm hidden (h_seed for
4835    /// the next draft round). This lets partial-accept replay run as ONE batched T=(n_acc+1) forward
4836    /// (single weight read) instead of n_acc+1 separate T=1 decode_steps (n_acc+1 weight reads).
4837    /// At batch=1 decode is bandwidth-bound, so batching the replay is THE MTP profitability lever.
4838    pub fn decode_step_t_h(
4839        &self,
4840        e: &Engine,
4841        tokens: &[u32],
4842        pos0: usize,
4843        cache: &mut Cache,
4844    ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4845        self.decode_step_t_h_emb(e, tokens, pos0, cache, None)
4846    }
4847
4848    /// Like `decode_step_t_h` with an optional RESIDENT embed table (spec hot loop): device
4849    /// gather instead of host dequant + [T, n_embd] f32 htod. Bit-identical rows.
4850    pub fn decode_step_t_h_emb(
4851        &self,
4852        e: &Engine,
4853        tokens: &[u32],
4854        pos0: usize,
4855        cache: &mut Cache,
4856        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
4857    ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4858        let (logits_d, h_seed) = self.decode_step_t_h_emb_dev(e, tokens, pos0, cache, embd_dev)?;
4859        Ok((e.dtoh(&logits_d)?, h_seed))
4860    }
4861
4862    /// DEVICE-LOGITS verify forward (spec device-argmax lever): identical kernel chain to
4863    /// `decode_step_t_h_emb` but returns the [T, n_vocab] logits ON DEVICE — the accept walk
4864    /// argmaxes each column on-device and reads back ONE [T] u32 instead of dtoh'ing the full
4865    /// T x n_vocab f32 block (~1-4 MB + T host argmaxes, every round). Kernel dispatch is
4866    /// UNCHANGED (same decode-exact kernels); only the post-logits transfer moves.
4867    pub fn decode_step_t_h_emb_dev(
4868        &self,
4869        e: &Engine,
4870        tokens: &[u32],
4871        pos0: usize,
4872        cache: &mut Cache,
4873        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
4874    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4875        cache.ensure_usable("decode_step_t")?;
4876        let n_embd = self.cfg.n_embd as usize;
4877        let t = tokens.len();
4878        let (logits, x) = self.decode_step_t_core(e, tokens, pos0, cache, embd_dev, None)?;
4879        // h_seed for the next round = LAST column's pre-output_norm hidden ([n_embd]).
4880        let mut hs = vbuf(e, n_embd)?; // fully written by copy_view_into below
4881        e.copy_view_into(&mut hs, 0, &x.slice((t - 1) * n_embd..t * n_embd), n_embd)?;
4882        Ok((logits, hs))
4883    }
4884
4885    /// CORE verify forward: the `decode_step_t_h_emb_dev` kernel chain, returning the FULL
4886    /// pre-output_norm hidden stack x ([T, n_embd], any column extractable) and optionally
4887    /// filling a `VerifyCkpt` (retained per-layer state-rebuild inputs) for the REPLAY-FREE
4888    /// partial accept. `ckpt: None` => byte-for-byte the old behavior (the ckpt writes are pure
4889    /// retains/copies — they never change what any kernel computes).
4890    fn decode_step_t_core(
4891        &self,
4892        e: &Engine,
4893        tokens: &[u32],
4894        pos0: usize,
4895        cache: &mut Cache,
4896        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
4897        mut ckpt: Option<&mut VerifyCkpt>,
4898    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4899        self.decode_step_t_core_stream(
4900            e,
4901            tokens,
4902            pos0,
4903            cache,
4904            embd_dev,
4905            ckpt.take(),
4906            None,
4907            None,
4908            None,
4909        )
4910    }
4911
4912    /// [`Self::decode_step_t_core`] with the MTP route's verify-graph pool armed
4913    /// (`MEMRA_SPEC_VERIFY_GRAPH`). `graphs: None` reproduces `decode_step_t_core`
4914    /// argument-for-argument, so the eager walk stays the byte-identical fallback.
4915    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
4916    fn decode_step_t_core_vg(
4917        &self,
4918        e: &Engine,
4919        tokens: &[u32],
4920        pos0: usize,
4921        cache: &mut Cache,
4922        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
4923        mut ckpt: Option<&mut VerifyCkpt>,
4924        graphs: Option<&mut DsparkVerifyGraphs>,
4925    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4926        self.decode_step_t_core_stream(
4927            e,
4928            tokens,
4929            pos0,
4930            cache,
4931            embd_dev,
4932            ckpt.take(),
4933            None,
4934            None,
4935            graphs,
4936        )
4937    }
4938
4939    /// ROUND-STREAM stage (c) 4: `stream` = (device verify tokens [t], device pos counter) —
4940    /// when Some, rope positions come from pos_iota over the counter, the embed gathers the
4941    /// device tokens, and full_attn_verify routes appends/FA through the _dc twins reading the
4942    /// SAME counter (every layer's kvl.len == cache.pos, one counter drives all three). The
4943    /// host `tokens`/`pos0` args still size buffers (t is FIXED K+1 in stream mode).
4944    /// `vtok_dev` (engine-bundle slice 2): device verify tokens for the EMBED only —
4945    /// unlike `stream` mode it changes nothing else (host pos iota, host-len KV appends).
4946    /// `tokens` then only sizes buffers (the dummy-slice pattern the round-stream arm uses).
4947    #[allow(clippy::too_many_arguments)]
4948    fn decode_step_t_core_stream(
4949        &self,
4950        e: &Engine,
4951        tokens: &[u32],
4952        pos0: usize,
4953        cache: &mut Cache,
4954        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
4955        mut ckpt: Option<&mut VerifyCkpt>,
4956        stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
4957        vtok_dev: Option<&CudaSlice<u32>>,
4958        graphs: Option<&mut DsparkVerifyGraphs>,
4959    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4960        // PP DOOR (lane/pp2-spec 2026-08-06): the verify trunk now takes its OWN stage split,
4961        // exactly as the eager and batched steps do. This is the single funnel every verify
4962        // forward reaches (decode_step_t / _h / _h_emb / _h_emb_dev / _core all land here), so
4963        // wiring it here wires the whole spec surface — the draft/accept/commit machinery above
4964        // is untouched.
4965        //
4966        // History: pp2-hardening (2026-08-06) made this funnel FAIL CLOSED, because its trunk
4967        // walk was unsplit on one stream and a sharded cross-device placement peer-read every
4968        // remote layer's weights on every spec round (measured 13.9-28x on the batched twin).
4969        // The refusal below survives to cover the residue — MEMRA_SPEC_PP=0, MEMRA_PP_STREAMS=0,
4970        // or a placement whose PpNRt fails to build — so a config that would still walk the
4971        // whole trunk on one stream refuses instead of regressing 28x.
4972        if let Some(fence) = crate::pp::pp_cuts(self.layers.len())
4973            && !crate::pp::pp2_streams_off()
4974            && crate::pp::spec_pp_on()
4975        {
4976            if vtok_dev.is_some() {
4977                return Err(
4978                    "device-token dspark verify (slice-2 deferred readback) has no PP \
4979                         stage-split arm; run the dspark route on one device"
4980                        .into(),
4981                );
4982            }
4983            return self.decode_step_t_core_ppn(
4984                e,
4985                tokens,
4986                pos0,
4987                cache,
4988                embd_dev,
4989                ckpt.take(),
4990                stream,
4991                &fence,
4992            );
4993        }
4994        crate::pp::refuse_unsplit_if_remote(
4995            "decode_step_t (spec verify)",
4996            "drop MEMRA_SPEC_PP=0 / MEMRA_PP_STREAMS=0 so the verify trunk takes its OWN stage \
4997             split (decode_step_t_core_ppn); or run spec on one device",
4998        )?;
4999        let cfg = &self.cfg;
5000        let n_embd = cfg.n_embd as usize;
5001        let eps = cfg.rms_eps;
5002        let t = tokens.len();
5003        let pos_d = match stream {
5004            Some((_, ctr)) => {
5005                let mut p = e.alloc_uninit::<i32>(t)?;
5006                e.pos_iota(ctr, &mut p, t)?;
5007                p
5008            }
5009            None => {
5010                let pos_vec: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
5011                e.htod_i32(&pos_vec)?
5012            }
5013        };
5014
5015        // embed T tokens -> [T, n_embd] token-major (device gather on the spec hot loop)
5016        let x = match (stream, embd_dev) {
5017            (Some((vtok, _)), Some((g, qt, rb))) => {
5018                e.embed_gather_device_td(g, vtok, t, n_embd, qt, rb)?
5019            }
5020            (None, Some((g, qt, rb))) => match vtok_dev {
5021                // slice 2: device verify tokens, same embed_gather_u32_t kernel —
5022                // bit-identical rows to the host-token arm (same per-dtype deq).
5023                Some(vt_d) => e.embed_gather_device_td(g, vt_d, t, n_embd, qt, rb)?,
5024                None => e.embed_gather_device_t(g, tokens, n_embd, qt, rb)?,
5025            },
5026            _ => {
5027                assert!(
5028                    vtok_dev.is_none(),
5029                    "device-token verify requires the resident embed table (embd_dev)"
5030                );
5031                e.htod(&self.embd.try_gather(n_embd, tokens)?)?
5032            }
5033        };
5034
5035        // TRUNK WALK: layers [0, n_layers) through the SAME range-scoped subgraph the PP-N
5036        // stage split calls per stage (`verify_layers`) — one code path, so the split cannot
5037        // drift from the unsplit dispatch mirroring. lane/pp2-spec 2026-08-06.
5038        let x = self.verify_layers(
5039            e,
5040            x,
5041            0,
5042            self.layers.len(),
5043            &pos_d,
5044            pos0,
5045            t,
5046            cache,
5047            ckpt.take(),
5048            stream,
5049            graphs,
5050        )?;
5051        if spec_nan_scan() {
5052            nan_scan_rows(e, &x, t, n_embd, &format!("verify trunk exit pos0={pos0}"))?;
5053        }
5054
5055        let mut hn = vbuf(e, t * n_embd)?;
5056        // Stage-A door: with the serving-class row-outer verify walk, the TAIL must be the
5057        // t=1 decode program per row too (rms_norm t=1 + the single-row bf16 head — the
5058        // split head's concat is receipted bit-identical to it). The batched cuBLASLt head
5059        // is a different ULP class and flips near-tie argmaxes off the greedy tape.
5060        let eager_tail = self.sliding_gated_moe_batch_program() && spec_verify_eager_on();
5061        if eager_tail {
5062            let n_vocab = self.cfg.n_vocab as usize;
5063            // MEMRA_SPEC_HEAD_ROWS=1 — THE VERIFY TAIL'S REDUNDANT HEAD READ.
5064            //
5065            // The loop below runs the head at m=1 once PER COLUMN, so the LM head's weights are
5066            // streamed t times per verify pass. On step37 that head is ~0.49 GiB per card after the
5067            // rank split, ~1.07 ms of pure re-read at t=2 and worse at every wider t — which is a
5068            // large part of why the fixed K ladder LOSES (K=1 81.2 > K=2 73.1 > K=3 62.7 tok/s).
5069            //
5070            // The loop's justification is the comment above: the batched cuBLASLt head is a
5071            // different ULP class and flips near-tie argmaxes off the greedy tape. That is true of
5072            // cuBLASLt and it does NOT apply here, because a FloatBf16 head at 1..=32 rows never
5073            // reaches cuBLASLt: `matmul` routes it to `matvec_bf16_rows_into` (lib.rs:12248), whose
5074            // own doc says `matvec_bf16_f32acc_x4_rows` "runs the t=1 decode head program PER ROW
5075            // (identical dot + reduce), so decode/verify tiers keep the t=1 numeric class". Under
5076            // the W8 doors both widths route to the q8 mirror instead, and the t-column mirror is
5077            // documented "bit-identical to t single-row calls". So the batched form is the SAME
5078            // arithmetic per row on both paths, with one weight read instead of t.
5079            //
5080            // rms_norm is row-wise, so norm(t) is per-row identical to t x norm(1) by construction.
5081            //
5082            // DEFAULT OFF for exactly one turn of the crank: "bit-identical by two documented
5083            // claims" is still an argument. The greedy byte tape decides, and the door flips only
5084            // once the tape is a receipt.
5085            if head_rows_on() {
5086                e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
5087                let logits = e.matmul(&self.output, &hn, t)?;
5088                if stream.is_none() {
5089                    cache.pos += t;
5090                }
5091                return Ok((logits, if spec_hpost() { hn } else { x }));
5092            }
5093            let mut logits = vbuf(e, t * n_vocab)?;
5094            for r in 0..t {
5095                let mut row = e.uninit(n_embd)?;
5096                e.dtod_copy_view(&x.slice(r * n_embd..(r + 1) * n_embd), &mut row)?;
5097                let mut hr = e.uninit(n_embd)?;
5098                e.rms_norm(&row, self.output_norm.float_data(), &mut hr, n_embd, 1, eps)?;
5099                let lr = e.matmul(&self.output, &hr, 1)?;
5100                e.dtod_copy_into(&lr, &mut logits, r * n_vocab)?;
5101                e.dtod_copy_into(&hr, &mut hn, r * n_embd)?;
5102            }
5103            if stream.is_none() {
5104                cache.pos += t;
5105            }
5106            return Ok((logits, if spec_hpost() { hn } else { x }));
5107        }
5108        let serving_head =
5109            self.sliding_gated_moe_batch_program() || self.batched_serving_numeric_class();
5110        let logits = if serving_head {
5111            // Step35 and the qwen35 family (MoE 2026-08-14 AM, dense-hybrid same day PM — the
5112            // Q3.8 bring-up reproduced the identical near-tie class on dense: eager-class verify
5113            // vs batched-class live serving, ULP drift amplified through the GDN recurrence)
5114            // serve one batched numeric class at every live width, including B=1. Keep the
5115            // verify head in that same class; other generic families retain the decode-exact
5116            // head that their run-spec contract pins.
5117            e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
5118            e.matmul(&self.output, &hn, t)?
5119        } else {
5120            e.rms_norm_decode(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
5121            e.matmul_decode_exact(&self.output, &hn, t)?
5122        };
5123        // stream: the device pos counter owns position; host mirror reconciles at drain.
5124        if stream.is_none() {
5125            cache.pos += t;
5126        }
5127        // Hidden stack for seeds/refresh-fills: pre-norm x (default) or post-norm hn (HPOST).
5128        Ok((logits, if spec_hpost() { hn } else { x }))
5129    }
5130
5131    /// THE VERIFY TRUNK OVER PP-N (lane/pp2-spec 2026-08-06): `decode_step_t_core_stream`'s walk
5132    /// as N stage subgraphs, each on its own engine/stream (and, under `MEMRA_PP_DEVICES`, its own
5133    /// device), with a `[T, n_embd]` boundary transfer between them. T = K+1 (the verify batch),
5134    /// so this is the batched-boundary shape the pp2-batch lane's grow-only slots already handle
5135    /// (`tx(b, x, t*n_embd)`; the slot grows to the high-water T and the transport moves exactly
5136    /// the payload).
5137    ///
5138    /// Structure is `decode_step_batch_ppn`'s, which is `decode_step_h_ppn`'s. FOUR THINGS ARE
5139    /// PER-STAGE and each for a measured reason (see `decode_step_batch_ppn`'s header for the
5140    /// receipts):
5141    ///
5142    /// 1. THE ENGINE (`rt.engine(s, e)`) — `Engine` owns lazily-grown stable-pointer scratch
5143    ///    (`fa_part_pool`, `fa_vf16_scratch`, `argmax_partials`) that is single-stream-safe BY
5144    ///    DESIGN. Two stage streams through one Engine is the 2026-08-02 shared-scratch race
5145    ///    (35% flake, nondeterministic all-logits divergence). `PpNRt::build` gives every stage
5146    ///    s>0 its own Engine even on the primary device; honouring it here is what scopes the
5147    ///    pools. The verify path allocates MORE of that scratch than eager decode does (FA at
5148    ///    m=T, and the per-layer `GdnStash` retains), so this is load-bearing, not inherited.
5149    ///
5150    /// 2. `pos_d` — each stage uploads its OWN copy of the T rope positions on ITS stream, so the
5151    ///    buffer is allocated, consumed and freed on one stream. In `stream` mode that means each
5152    ///    stage runs its own `pos_iota` over the SHARED device counter (`pos_ctr`): the counter is
5153    ///    read-only during the forward (the round's `inc`/`copy_add` happen outside it), so every
5154    ///    stage derives the identical iota, and each stage's own output buffer is stream-local.
5155    ///
5156    /// 3. THE EMBED lives with stage 0 (`self.embd` / `embd_gpu` are host/primary-side; the
5157    ///    sharded loader leaves the table with stage 0 by construction).
5158    ///
5159    /// 4. THE HEAD (`output_norm` + `output`) runs on the LAST stage — the sharded loader uploaded
5160    ///    both through that stage's engine (`hybrid.rs`: `e_head = layer_engine(e, n_trunk,
5161    ///    n_trunk-1)`), so reading them anywhere else is a peer read of the biggest tensor in the
5162    ///    model, every round.
5163    ///
5164    /// WHAT STAYS ON THE PRIMARY, deliberately: the returned logits and hidden stack `x`. Both are
5165    /// last-stage-allocated device buffers, and every consumer (the device argmax walk, the accept
5166    /// kernels, `spec_seed_gather`, the ckpt rebuild in `commit_verified_prefix`) reads them
5167    /// through the primary context by UVA — the same read the batched serving epilogue's
5168    /// `last_logits_dev` park does. Those consumers are per-round O(T x n_vocab) and O(n_embd),
5169    /// not per-layer, so they are not the 28x class; splitting them is a separate lane.
5170    ///
5171    /// The MTP HEAD (draft side) is NOT split: it is one block, it lives wherever the loader put
5172    /// it (`load_mtp` uses the primary engine), and it is ~1-2 GB against the trunk's tens. Draft
5173    /// placement is measured, not assumed — see `research/pp2-spec-20260806`.
5174    ///
5175    /// EXACTNESS: PP-N adds ZERO deviation. Each stage runs the SAME kernels on the SAME bytes in
5176    /// the same order via the SAME `verify_layers` the unsplit body calls; the only change is
5177    /// where the residual is materialized, and the boundary is a straight f32 copy (dtod
5178    /// same-device / `cudaMemcpyPeerAsync` cross-device, no conversion). So the split MUST be
5179    /// BIT-IDENTICAL to the unsplit verify at the same T, in both placement orders. Gate:
5180    /// `decode-batch-gate --mode ppspec`. Acceptance counts are a DERIVED consequence — greedy
5181    /// accept argmaxes these logits, so bit-identical logits force identical accept walks; the
5182    /// `run-spec` K=1..8 arm checks that end-to-end rather than trusting the implication.
5183    #[allow(clippy::too_many_arguments)]
5184    fn decode_step_t_core_ppn(
5185        &self,
5186        e: &Engine,
5187        tokens: &[u32],
5188        pos0: usize,
5189        cache: &mut Cache,
5190        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
5191        mut ckpt: Option<&mut VerifyCkpt>,
5192        stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
5193        fence: &[usize],
5194    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
5195        let ticket = self.verify_stage0_issue(
5196            e,
5197            tokens,
5198            pos0,
5199            cache,
5200            embd_dev,
5201            ckpt.as_deref_mut(),
5202            stream,
5203            fence,
5204        )?;
5205        self.verify_stage1_finish(e, ticket, cache, ckpt, stream, fence, true)
5206    }
5207
5208    /// Enqueue embed, stage 0, and the first boundary TX, then return the actual boundary slot.
5209    /// The ordinary PP verify wrapper calls `verify_stage1_finish` immediately after this return.
5210    #[allow(clippy::too_many_arguments)]
5211    fn verify_stage0_issue(
5212        &self,
5213        e: &Engine,
5214        tokens: &[u32],
5215        pos0: usize,
5216        cache: &mut Cache,
5217        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
5218        ckpt: Option<&mut VerifyCkpt>,
5219        stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
5220        fence: &[usize],
5221    ) -> Result<VerifyBoundaryTicket, Box<dyn std::error::Error>> {
5222        assert!(
5223            !self.is_gemma4_e4b() && !self.gemma_batch_program(),
5224            "decode_step_t_core_ppn covers the hybrid non-gemma4 verify trunk only \
5225             (the gemma4 arms have their own decode_step_t twins)"
5226        );
5227        if crate::pp::pp_host_bounce_active() && (stream.is_some() || embd_dev.is_some()) {
5228            return Err(
5229                "decode_step_t_core_ppn: refused with MEMRA_PP_HOST_BOUNCE=1 — the trunk \
5230                 boundary itself is host-staged, but device-resident verify still peer-reads \
5231                 primary-device token/position/embedding buffers from stage 0. Run plain PP \
5232                 serving on this host class; spec requires local per-stage inputs first."
5233                    .into(),
5234            );
5235        }
5236        let rt = crate::pp::PpNRt::get(e)?;
5237        // Pipelined callers do not bypass ownership: their explicit coordinator borrow makes
5238        // this acquire clone the same active generation. Ordinary callers acquire a fresh lease.
5239        let walk_owner = rt.acquire_walk("verify_stage0_issue")?;
5240        let n_st = fence.len() - 1;
5241        assert_eq!(
5242            rt.n_stages(),
5243            n_st,
5244            "PpNRt stage count {} != fence stages {n_st}",
5245            rt.n_stages()
5246        );
5247        let n_embd = self.cfg.n_embd as usize;
5248        let t = tokens.len();
5249        let payload = t * n_embd;
5250        // One-shot lane diagnostic: force natural PP-2 boundaries to completion so the server
5251        // log can price stage 0, the peer hop, the RX copy, and stage 1 + head separately without
5252        // nsys. The ordinary path keeps every enqueue asynchronous. N>2 is deliberately excluded:
5253        // the report below names exactly two stages and must never imply it measured middle ones.
5254        let pp_anatomy = n_st == 2 && std::env::var("MEMRA_SPEC_PP_ANATOMY").as_deref() == Ok("1");
5255        let pp_started = std::time::Instant::now();
5256        let (mut reverse_ms, mut stage0_ms, mut tx_ms) = (0.0f64, 0.0f64, 0.0f64);
5257        // The CALLER's ambient stream, captured BEFORE any `rt.enter()` pushes a stage stream:
5258        // this body returns DEVICE-RESIDENT buffers (the device-argmax accept walk's contract),
5259        // so the exit needs the same publication the boundaries get — see `PpNRt::publish_to`.
5260        // Taken here, not at the end, because inside the last-stage scope `e.stream()` IS the
5261        // stage stream and the wait would self-order into a no-op.
5262        let caller_stream = e.stream();
5263        // #87 ROOT-CAUSE FENCE (lane/pp2spec-crash): the PREVIOUS round's stage-allocated
5264        // outputs (logits/hidden/ckpt stashes) freed stream-ordered on the STAGE streams while
5265        // the primary stream still holds queued reads of them — with event tracking elided,
5266        // nothing stops the pool from reusing those blocks for THIS round's stage allocations,
5267        // whose writes then race the queued reads (measured: 13/4096-NaN random-bits garbage in
5268        // the spec round seed; the full anatomy is on `PpNRt::fence_stages_behind`). Order every
5269        // stage stream behind the caller before enqueueing new stage work.
5270        let reverse_started = std::time::Instant::now();
5271        rt.fence_stages_behind(&caller_stream)?;
5272        if pp_anatomy {
5273            // Drain the reverse-publication dependency before timing stage 0 itself. At c=1 this
5274            // prices any primary-stream rollback/refresh tail inherited from the prior round.
5275            for s in 0..n_st {
5276                let _st = rt.enter(s);
5277                rt.engine(s, e).stream().synchronize()?;
5278            }
5279            reverse_ms = reverse_started.elapsed().as_secs_f64() * 1e3;
5280        }
5281
5282        // Per-stage rope positions: in host mode the same [T] iota each stage uploads itself; in
5283        // stream mode each stage's own `pos_iota` over the shared read-only device counter.
5284        let stage_pos = |es: &Engine| -> Result<CudaSlice<i32>, Box<dyn std::error::Error>> {
5285            match stream {
5286                Some((_, ctr)) => {
5287                    let mut p = es.alloc_uninit::<i32>(t)?;
5288                    es.pos_iota(ctr, &mut p, t)?;
5289                    Ok(p)
5290                }
5291                None => {
5292                    let pos_vec: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
5293                    es.htod_i32(&pos_vec)
5294                }
5295            }
5296        };
5297
5298        // ---- STAGE 0: embed (the table lives with stage 0) + layers [0, fence[1]) + TX ----
5299        let slot = {
5300            let _st0 = rt.enter(0);
5301            let e0 = rt.engine(0, e);
5302            let stage0_started = std::time::Instant::now();
5303            let pos_d = stage_pos(e0)?;
5304            let x = match (stream, embd_dev) {
5305                (Some((vtok, _)), Some((g, qt, rb))) => {
5306                    e0.embed_gather_device_td(g, vtok, t, n_embd, qt, rb)?
5307                }
5308                (None, Some((g, qt, rb))) => e0.embed_gather_device_t(g, tokens, n_embd, qt, rb)?,
5309                _ => e0.htod(&self.embd.try_gather(n_embd, tokens)?)?,
5310            };
5311            let x = self.verify_layers(
5312                e0, x, fence[0], fence[1], &pos_d, pos0, t, cache, ckpt, stream, None,
5313            )?;
5314            if pp_anatomy {
5315                e0.stream().synchronize()?;
5316                stage0_ms = stage0_started.elapsed().as_secs_f64() * 1e3;
5317            }
5318            let tx_started = std::time::Instant::now();
5319            let slot = rt.tx(0, &x, payload)?;
5320            if pp_anatomy {
5321                e0.stream().synchronize()?;
5322                tx_ms = tx_started.elapsed().as_secs_f64() * 1e3;
5323            }
5324            slot
5325            // x + pos_d drop here: freed stream-ordered on stage-0's stream after use.
5326        };
5327
5328        Ok(VerifyBoundaryTicket {
5329            rt,
5330            caller_stream,
5331            slot,
5332            pos0,
5333            t,
5334            payload,
5335            n_st,
5336            pp_anatomy,
5337            pp_started,
5338            reverse_ms,
5339            stage0_ms,
5340            tx_ms,
5341            _walk_owner: walk_owner,
5342        })
5343    }
5344
5345    /// Consume a stage-0 boundary ticket and enqueue the remaining PP stages plus the head.
5346    /// On PP-2 this is exactly stage 1; PP-N keeps its pre-existing middle-stage walk here.
5347    #[allow(clippy::too_many_arguments)]
5348    fn verify_stage1_finish(
5349        &self,
5350        e: &Engine,
5351        ticket: VerifyBoundaryTicket,
5352        cache: &mut Cache,
5353        mut ckpt: Option<&mut VerifyCkpt>,
5354        stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
5355        fence: &[usize],
5356        publish_to_caller: bool,
5357    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
5358        let VerifyBoundaryTicket {
5359            rt,
5360            caller_stream,
5361            slot,
5362            pos0,
5363            t,
5364            payload,
5365            n_st,
5366            pp_anatomy,
5367            pp_started,
5368            reverse_ms,
5369            stage0_ms,
5370            tx_ms,
5371            _walk_owner,
5372        } = ticket;
5373        let n_embd = self.cfg.n_embd as usize;
5374        let eps = self.cfg.rms_eps;
5375        let mut slot = slot;
5376        let (mut rx_ms, mut stage1_ms) = (0.0f64, 0.0f64);
5377        let stage_pos = |es: &Engine| -> Result<CudaSlice<i32>, Box<dyn std::error::Error>> {
5378            match stream {
5379                Some((_, ctr)) => {
5380                    let mut p = es.alloc_uninit::<i32>(t)?;
5381                    es.pos_iota(ctr, &mut p, t)?;
5382                    Ok(p)
5383                }
5384                None => {
5385                    let pos_vec: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
5386                    es.htod_i32(&pos_vec)
5387                }
5388            }
5389        };
5390
5391        // ---- MIDDLE STAGES: RX boundary s-1 -> range -> TX boundary s ----
5392        for s in 1..n_st - 1 {
5393            let _st = rt.enter(s);
5394            let es = rt.engine(s, e);
5395            let pos_d = stage_pos(es)?;
5396            let x = rt.rx(s - 1, slot, payload)?;
5397            let x = self.verify_layers(
5398                es,
5399                x,
5400                fence[s],
5401                fence[s + 1],
5402                &pos_d,
5403                pos0,
5404                t,
5405                cache,
5406                ckpt.as_deref_mut(),
5407                stream,
5408                None,
5409            )?;
5410            slot = rt.tx(s, &x, payload)?;
5411        }
5412
5413        // ---- LAST STAGE: RX + final range + output_norm + lm head ----
5414        let _stl = rt.enter(n_st - 1);
5415        let el = rt.engine(n_st - 1, e);
5416        let pos_d = stage_pos(el)?;
5417        let rx_started = std::time::Instant::now();
5418        let x = rt.rx(n_st - 2, slot, payload)?;
5419        if pp_anatomy {
5420            el.stream().synchronize()?;
5421            rx_ms = rx_started.elapsed().as_secs_f64() * 1e3;
5422        }
5423        let stage1_started = std::time::Instant::now();
5424        let x = self.verify_layers(
5425            el,
5426            x,
5427            fence[n_st - 1],
5428            fence[n_st],
5429            &pos_d,
5430            pos0,
5431            t,
5432            cache,
5433            ckpt,
5434            stream,
5435            None,
5436        )?;
5437
5438        let mut hn = vbuf(el, payload)?;
5439        let logits = if self.sliding_gated_moe_batch_program() {
5440            // The PP Step35 serving path uses rms_norm + matmul for B=1 as well as B>1.
5441            // Verify must not switch numeric class merely because the same session speculates.
5442            el.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
5443            el.matmul(&self.output, &hn, t)?
5444        } else {
5445            el.rms_norm_decode(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
5446            el.matmul_decode_exact(&self.output, &hn, t)?
5447        };
5448        if pp_anatomy {
5449            el.stream().synchronize()?;
5450            stage1_ms = stage1_started.elapsed().as_secs_f64() * 1e3;
5451        }
5452        // EXIT PUBLICATION: both returned buffers are still being produced on the last stage's
5453        // stream. Order the caller's stream behind that work before the buffers escape this
5454        // scope (the 2026-08-06 same-device ppspec find: without it the caller's primary-stream
5455        // consumer read unwritten logits — nondeterministic, one-device-only, and it poisoned
5456        // the following arm's KV in the same process).
5457        if publish_to_caller {
5458            rt.publish_to(n_st - 1, &caller_stream)?;
5459        }
5460        if pp_anatomy {
5461            if publish_to_caller {
5462                caller_stream.synchronize()?;
5463            }
5464            eprintln!(
5465                "[spec-pp-anatomy] t={t} reverse={reverse_ms:.3}ms stage0={stage0_ms:.3}ms \
5466                 tx={tx_ms:.3}ms rx={rx_ms:.3}ms stage1-head={stage1_ms:.3}ms total={:.3}ms",
5467                pp_started.elapsed().as_secs_f64() * 1e3,
5468            );
5469        }
5470        // stream: the device pos counter owns position; host mirror reconciles at drain.
5471        if stream.is_none() {
5472            cache.pos += t;
5473        }
5474        Ok((logits, if spec_hpost() { hn } else { x }))
5475    }
5476
5477    /// Step3.5/Step3.7 verify trunk in the serving batched numeric class.
5478    ///
5479    /// `step35_decode_batch_layers` is now authoritative at every live serving width, including
5480    /// B=1 (lane/cx-b1fix). The older verify walk deliberately mirrored the eager T=1 class:
5481    /// it replayed `step35_decode_attn` per row and used the eager/decode-exact FFN dispatch.
5482    /// Those classes are individually stable, but a near-tie prompt can choose different greedy
5483    /// bytes when a request moves from batched plain serving into speculative verify. Run the
5484    /// same authoritative B=1 stage subgraph for each verify row here. Rows still advance
5485    /// layer-by-layer, so every layer sees the preceding verify rows in its attention cache while
5486    /// every norm/projection/FFN uses exactly the live serving dispatch.
5487    #[allow(clippy::too_many_arguments)]
5488    /// PRIME-BY-T-ROWS (MEMRA_PRIME_TROWS=1): prefill the prompt through the same-session
5489    /// t-row walk in 32-row chunks — every row runs the t=1 decode program bit-for-bit
5490    /// (the TOKENWISE-prime ORACLE class), so this door is exact against the exactness
5491    /// reference while replacing the host-canonical per-token prime. Requires the walk
5492    /// doors (MEMRA_SPEC_VERIFY_EAGER/TCOL); returns the prime contract trio.
5493    #[allow(clippy::type_complexity)]
5494    pub(crate) fn step35_prime_trows(
5495        &self,
5496        e: &Engine,
5497        tokens: &[u32],
5498        cache: &mut Cache,
5499    ) -> Result<Option<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
5500    {
5501        let dbg = std::env::var("MEMRA_SPEC_FA2_DEBUG").as_deref() == Ok("1");
5502        if !prime_trows_on() {
5503            return Ok(None);
5504        }
5505        if !self.uses_sliding_gated_moe_program()
5506            || cache.pos != 0
5507            || cache.dflash_taps.is_some()
5508            || !spec_verify_eager_on()
5509            || !spec_verify_tcol_on()
5510        {
5511            if dbg {
5512                eprintln!(
5513                    "[prime-trows] refuse: program={} pos={} taps={} eager={:?} tcol={:?}",
5514                    self.uses_sliding_gated_moe_program(),
5515                    cache.pos,
5516                    cache.dflash_taps.is_some(),
5517                    std::env::var("MEMRA_SPEC_VERIFY_EAGER").ok(),
5518                    std::env::var("MEMRA_SPEC_VERIFY_TCOL").ok()
5519                );
5520            }
5521            return Ok(None);
5522        }
5523        let n_embd = self.cfg.n_embd as usize;
5524        let n_layers = self.layers.len();
5525        let t_total = tokens.len();
5526        let Some(embd_gpu) = self.embd_gpu_try(e) else {
5527            if dbg {
5528                eprintln!("[prime-trows] refuse: no device embed table");
5529            }
5530            return Ok(None);
5531        };
5532        let embd_qtype = match self.embd.ggml_type {
5533            memra_gguf::GgmlType::BF16 => crate::QT_BF16,
5534            memra_gguf::GgmlType::Q8_0 => crate::QT_Q8_0,
5535            other => {
5536                if dbg {
5537                    eprintln!("[prime-trows] refuse: embed dtype {other:?}");
5538                }
5539                return Ok(None);
5540            }
5541        };
5542        let embd_row_bytes = self.embd.raw.len() / self.cfg.n_vocab as usize;
5543        // Chunk plan: 32-row chunks; a 1-token tail folds into the previous chunk
5544        // (the walk floor is t >= 2).
5545        let mut bounds = Vec::new();
5546        let mut start = 0usize;
5547        while start < t_total {
5548            let mut end = (start + 32).min(t_total);
5549            if t_total - end == 1 {
5550                end -= 1;
5551            }
5552            bounds.push((start, end));
5553            start = end;
5554        }
5555        if bounds.iter().any(|(a, b)| b - a < 2) {
5556            return Ok(None); // degenerate short prompt keeps the ordinary prime
5557        }
5558        let mut hiddens = e.uninit(t_total * n_embd)?;
5559        let mut last: Option<CudaSlice<f32>> = None;
5560        for &(a, b) in &bounds {
5561            let tc = b - a;
5562            let tok_d = e.stream().clone_htod(&tokens[a..b])?;
5563            let x =
5564                e.embed_gather_device_td(embd_gpu, &tok_d, tc, n_embd, embd_qtype, embd_row_bytes)?;
5565            let out = self.step35_verify_batch_layers(e, x, 0, n_layers, a, tc, cache)?;
5566            e.copy_into(&mut hiddens, a * n_embd, &out, tc * n_embd)?;
5567            if b == t_total {
5568                let mut h = e.uninit(n_embd)?;
5569                e.dtod_copy_view(&out.slice((tc - 1) * n_embd..tc * n_embd), &mut h)?;
5570                last = Some(h);
5571            }
5572        }
5573        let h_seed = last.expect("last chunk produced the seed row");
5574        let mut hn = e.uninit(n_embd)?;
5575        e.rms_norm_decode(
5576            &h_seed,
5577            self.output_norm.float_data(),
5578            &mut hn,
5579            n_embd,
5580            1,
5581            self.cfg.rms_eps,
5582        )?;
5583        let logits_d = e.matmul_decode_exact(&self.output, &hn, 1)?;
5584        let logits = e.dtoh(&logits_d)?;
5585        cache.pos = t_total;
5586        Ok(Some((logits, h_seed, hiddens)))
5587    }
5588
5589    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
5590    fn step35_verify_batch_layers(
5591        &self,
5592        e: &Engine,
5593        mut x: CudaSlice<f32>,
5594        lo: usize,
5595        hi: usize,
5596        pos0: usize,
5597        t: usize,
5598        cache: &mut Cache,
5599    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5600        let n_embd = self.cfg.n_embd as usize;
5601        if !self.uses_sliding_gated_moe_program() {
5602            return Err(
5603                "serving-class verify requires sliding-gated-MoE canonical operations".into(),
5604            );
5605        }
5606        // SERVING-CLASS VERIFY (MEMRA_SPEC_VERIFY_EAGER=1, step37 MTP bring-up): each verify
5607        // column rides decode_layers_eager — the EXACT t=1 program live serving runs (all TP2
5608        // doors) — row-outer, so row r's appends land before row r+1 attends: bit-equal to
5609        // plain greedy by construction. Only the unsplit full-range walk qualifies; PP splits
5610        // and the tap path keep the batch-layer class.
5611        static VE: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
5612        let eager_verify =
5613            *VE.get_or_init(spec_verify_eager_on) && lo == 0 && hi == self.layers.len();
5614        if eager_verify {
5615            // T-COLUMN LAYER-OUTER WALK (MEMRA_SPEC_VERIFY_TCOL=1): per layer, one t-grid
5616            // attn norm + ONE weight-amortized QKV(+gate) over all T columns, then each
5617            // column runs the UNMODIFIED t=1 attention program via the col-select door and
5618            // the ordinary residual/FFN body. Values per column are bit-equal to the
5619            // row-outer walk: rms over the materialized residual == the fused add+norm
5620            // (kernel_check identity), the tcol kernel's per-column FP order == the t=1
5621            // kernel, and every downstream op IS the t=1 program.
5622            static TCOL: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
5623            let tcol = *TCOL.get_or_init(spec_verify_tcol_on);
5624            // T > 32 (prefill-class): run the SAME walk in 32-row chunks — each chunk's
5625            // rows are the t=1 program bit-for-bit and the rope pass advances the cache,
5626            // so a chunked call is value-identical to the row-outer loop it replaces.
5627            static TROWS_PREFILL: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
5628            // MEMRA_STEP_GEMM_PRIME outranks the walk: with the grouped GEMM prime armed, the
5629            // t-row walk defers so the batch path (GEMM trunk + grouped MoE) takes the prompt —
5630            // flag precedence between two existing doors, not a new flag. Without this, both
5631            // doors ON meant the walk still won and the GEMM prime needed PRIME_TROWS=0 by hand.
5632            let trows_prefill =
5633                *TROWS_PREFILL.get_or_init(|| prime_trows_on() && !crate::step_gemm_prime_on());
5634            // MEMRA_PRIME_TROWS_T=<w>: chunk width, default 8 = the REAL cap of this walk.
5635            // The workspace slabs go to 32 rows, but `matvec_bf16_qkvg_tcol_into` refuses
5636            // t > 8 (compile-time-T twins exist for 2/4/8 only; the runtime-t kernel spills
5637            // its accumulators to local memory), so a wider chunk fails the request with
5638            // "matvec_bf16_qkvg_tcol geometry" — which is exactly how the first server-path
5639            // TROWS arm died. Measured at 193 tokens: w=8 2.459 s, w=4 2.574 s.
5640            static TROWS_W: std::sync::OnceLock<Result<usize, String>> = std::sync::OnceLock::new();
5641            let trows_w = match TROWS_W.get_or_init(|| {
5642                let value = std::env::var("MEMRA_PRIME_TROWS_T").ok();
5643                parse_prime_trows_width(value.as_deref())
5644            }) {
5645                Ok(width) => *width,
5646                Err(err) => return Err(err.clone().into()),
5647            };
5648            if tcol && trows_prefill && t > trows_w {
5649                // One-time engagement receipt: without it a prefill gate cannot tell a
5650                // chunked walk from the row-outer fallback it is supposed to replace
5651                // (the first PRIME_TROWS gate passed vacuously on exactly that).
5652                static SEEN: std::sync::atomic::AtomicBool =
5653                    std::sync::atomic::AtomicBool::new(false);
5654                if !SEEN.swap(true, std::sync::atomic::Ordering::Relaxed) {
5655                    eprintln!(
5656                        "[prime-trows] ENGAGED t={t} width={trows_w} chunks={} layers={}..{}",
5657                        t.div_ceil(trows_w),
5658                        lo,
5659                        hi
5660                    );
5661                }
5662                let mut out = e.uninit(t * n_embd)?;
5663                let mut start = 0usize;
5664                while start < t {
5665                    let mut end = (start + trows_w).min(t);
5666                    if t - end == 1 {
5667                        end -= 1;
5668                    }
5669                    let tc = end - start;
5670                    let mut xc = e.uninit(tc * n_embd)?;
5671                    e.dtod_copy_view(&x.slice(start * n_embd..end * n_embd), &mut xc)?;
5672                    let oc =
5673                        self.step35_verify_batch_layers(e, xc, lo, hi, pos0 + start, tc, cache)?;
5674                    e.copy_into(&mut out, start * n_embd, &oc, tc * n_embd)?;
5675                    start = end;
5676                }
5677                return Ok(out);
5678            }
5679            if tcol && (2..=32).contains(&t) {
5680                // MEMRA_TCOL_PROF=1: synchronized per-segment wall profile of the walk
5681                // (norm+QKV precompute / per-col attention / per-col residual+FFN). The
5682                // syncs serialize the stream, so the split is for TARGETING amortization
5683                // work only — never a perf claim.
5684                static PROF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
5685                let prof =
5686                    *PROF.get_or_init(|| std::env::var("MEMRA_TCOL_PROF").as_deref() == Ok("1"));
5687                let mut prof_ms = [0f64; 3];
5688                let eps = self.cfg.rms_eps;
5689                let mut x_t = x;
5690                let mut h_t = e.uninit(t * n_embd)?;
5691                let mut h_row = e.uninit(n_embd)?; // real row: the non-dcw fallback reads it
5692                // Per-column pos buffers hoisted out of the layer loop (a per-col-per-layer
5693                // pageable htod was an in-stream engine turnaround x t x 45).
5694                let mut pos_rows = Vec::with_capacity(t);
5695                for r in 0..t {
5696                    pos_rows.push(e.htod_i32(&[(pos0 + r) as i32])?);
5697                }
5698                let mut ok = true;
5699                // MEMRA_TCOL_OPROJ=1: defer each column's o_proj — the finish seam
5700                // stashes `gated` instead of joining per column; one b4_tcol per rank +
5701                // one slab join produce every column's `mixed` after the attention pass.
5702                // Bit-exact per column (t=1 b4 program per column; elementwise join).
5703                // MEMRA_TCOL_FFN=1: today this only IMPLIES the o_proj defer above. Its
5704                // named feature, the two-column device-routed FFN sweep, rode the
5705                // slot-major v2 TP banks and was REMOVED with the MEMRA_NVFP4_BANK_V2 door
5706                // (2026-08-29, research/step37-bankv2-removal-20260829): the v2 layout
5707                // changed generated text in serving. The flag itself stays because it is
5708                // family-armed in the step37 serving defaults and killing it here would
5709                // silently drop the o_proj defer from the qualified serving shape.
5710                static FFN2: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
5711                let ffn_batch = *FFN2.get_or_init(tcol_ffn_on);
5712                let oproj_batch = crate::tp::tcol_oproj_on() || ffn_batch;
5713                // MEMRA_SPEC_FA2=1: eligible layers defer every verify column's fa —
5714                // the per-column pass norms/ropes/appends and stashes q+gate, then one
5715                // fa_decode_vec_q_v3_dcw_rows per rank attends every stashed row and the
5716                // o_proj join produces the mixed slab. Ineligible/boundary layers run the
5717                // ordinary program.
5718                let fa2 = crate::tp::spec_fa2_on() && t <= 32;
5719                let mut mixed_row = e.uninit(n_embd)?;
5720                let mut pos_staged = false;
5721                for il in lo..hi {
5722                    let layer = &self.layers[il];
5723                    // BEFORE this layer touches its planes: is the history it is about to
5724                    // attend already poisoned? Global (non-ring) layers only, which are the
5725                    // ones the level-2 bitmap implicates.
5726                    if kv_plane_scan_on()
5727                        && self.step35_geom(il).window.is_none()
5728                        && let Some(distributed) = cache.tp_kv[il].as_ref()
5729                    {
5730                        scan_kv_plane(e, distributed, il, pos0)?;
5731                    }
5732                    let fa2_layer = fa2 && self.step35_fa_rows_precheck(cache, il, pos0, t)?;
5733                    let mut seg = std::time::Instant::now();
5734                    e.rms_norm(&x_t, layer.attn_norm.float_data(), &mut h_t, n_embd, t, eps)?;
5735                    if !self.step35_verify_qkv_precompute(e, il, &h_t, t)? {
5736                        ok = false;
5737                        break;
5738                    }
5739                    // FULL t-row attention pass (rope/append + fa + combine + o_proj in
5740                    // 3 launches/rank): same-session rows, slot = len-base+r, one len
5741                    // advance by t. Ring rebase happens during append BEFORE the fused
5742                    // kernel so device base_d and memory are already rebased for rows.
5743                    // Host cache bookkeeping mirrors the per-column tail.
5744                    let mut mixed_t_opt: Option<CudaSlice<f32>> = None;
5745                    if fa2_layer {
5746                        let crate::hybrid::Mixer::Full(fa) = &layer.mixer else {
5747                            return Err("verify rope pass expects full attention".into());
5748                        };
5749                        let tp = fa
5750                            .step_tp_qkv
5751                            .as_ref()
5752                            .ok_or("verify rope pass lost its TP state")?;
5753                        let empty: [CudaSlice<f32>; 0] = [];
5754                        let transaction = {
5755                            let tp_kv = cache.tp_kv[il]
5756                                .as_mut()
5757                                .expect("precheck verified the distributed cache");
5758                            let tx = tp_kv.begin_transaction()?;
5759                            tp.runtime.append_tp_kv_transaction_inner(
5760                                tp_kv, tx, &empty, &empty, t, true,
5761                            )?;
5762                            tx
5763                        };
5764                        match self.step35_verify_rope_fa_pass(e, il, cache, pos0, t, !pos_staged) {
5765                            Ok(Some(mixed_t)) => {
5766                                let tp_kv = cache.tp_kv[il]
5767                                    .as_mut()
5768                                    .expect("precheck verified the distributed cache");
5769                                tp.runtime.commit_tp_kv_transaction_external(
5770                                    tp_kv,
5771                                    transaction,
5772                                    t,
5773                                )?;
5774                                if let Some(local) = cache.kv[il].as_mut() {
5775                                    local.len = pos0 + t;
5776                                    if !crate::tp::len_mirror_lazy_on() {
5777                                        e.set_i32_one(&mut local.len_d, local.len as i32)?;
5778                                    }
5779                                    if let (Some(ring), Some(tp_base)) =
5780                                        (local.ring.as_mut(), tp_kv.ring_base())
5781                                        && ring.base() != tp_base
5782                                    {
5783                                        ring.apply_rebase(tp_base);
5784                                        if let Some(base_d) = local.base_d.as_mut() {
5785                                            e.set_i32_one(base_d, tp_base as i32)?;
5786                                        }
5787                                    }
5788                                }
5789                                pos_staged = true;
5790                                mixed_t_opt = Some(mixed_t);
5791                            }
5792                            Ok(None) => {
5793                                let tp_kv = cache.tp_kv[il]
5794                                    .as_mut()
5795                                    .expect("precheck verified the distributed cache");
5796                                tp.runtime.rollback_tp_kv_transaction(tp_kv, transaction)?;
5797                            }
5798                            Err(err) => {
5799                                if let Some(tp_kv) = cache.tp_kv[il].as_mut() {
5800                                    let _ =
5801                                        tp.runtime.rollback_tp_kv_transaction(tp_kv, transaction);
5802                                }
5803                                return Err(err);
5804                            }
5805                        }
5806                    }
5807                    if let Some(mixed_t) = mixed_t_opt {
5808                        if prof {
5809                            e.stream().synchronize()?;
5810                            prof_ms[1] += seg.elapsed().as_secs_f64() * 1e3;
5811                            seg = std::time::Instant::now();
5812                        }
5813                        let o_out = mixed_t.len() / t;
5814                        let mut next = e.uninit(t * n_embd)?;
5815                        {
5816                            for r in 0..t {
5817                                e.dtod_copy_view(
5818                                    &mixed_t.slice(r * o_out..(r + 1) * o_out),
5819                                    &mut mixed_row,
5820                                )?;
5821                                let mut x_row = e.uninit(n_embd)?;
5822                                e.dtod_copy_view(
5823                                    &x_t.slice(r * n_embd..(r + 1) * n_embd),
5824                                    &mut x_row,
5825                                )?;
5826                                let (x1, ffn_out) = self.residual_norm_ffn(
5827                                    e, layer, &x_row, &mixed_row, n_embd, il, eps,
5828                                )?;
5829                                let mut x2 = e.uninit(n_embd)?;
5830                                e.add(&x1, &ffn_out, &mut x2, n_embd)?;
5831                                e.dtod_copy_into(&x2, &mut next, r * n_embd)?;
5832                            }
5833                        }
5834                        if prof {
5835                            e.stream().synchronize()?;
5836                            prof_ms[2] += seg.elapsed().as_secs_f64() * 1e3;
5837                        }
5838                        x_t = next;
5839                        if spec_nan_scan() {
5840                            // The scan MUST sit on this arm too. It used to live only on
5841                            // the non-fused tail, so a fused layer's poison was first
5842                            // reported by the next non-fused layer.
5843                            verify_arm_receipt(
5844                                "fused",
5845                                il,
5846                                pos0,
5847                                t,
5848                                cache.tp_kv[il].as_ref().map(|d| d.staged_len()),
5849                            );
5850                            nan_scan_rows(
5851                                e,
5852                                &x_t,
5853                                t,
5854                                n_embd,
5855                                &format!("tcol layer {il} pos0={pos0} arm=fused"),
5856                            )?;
5857                        }
5858                        continue;
5859                    }
5860                    if prof {
5861                        e.stream().synchronize()?;
5862                        prof_ms[0] += seg.elapsed().as_secs_f64() * 1e3;
5863                        seg = std::time::Instant::now();
5864                    }
5865                    let mut next = e.uninit(t * n_embd)?;
5866                    // Columns whose o_proj was deferred (their FFN runs after the join).
5867                    // A NON-deferred column's FFN must run INSIDE the column loop: the
5868                    // oproj-tail handoff is a single cell that the same column's
5869                    // residual_norm_ffn consumes before the next column's finish.
5870                    let mut deferred: Vec<usize> = Vec::new();
5871                    let mut fa2_deferred: Vec<usize> = Vec::new();
5872                    let ffn_col = |r: usize,
5873                                   mixed: &CudaSlice<f32>,
5874                                   next: &mut CudaSlice<f32>|
5875                     -> Result<(), Box<dyn std::error::Error>> {
5876                        let mut x_row = e.uninit(n_embd)?;
5877                        e.dtod_copy_view(&x_t.slice(r * n_embd..(r + 1) * n_embd), &mut x_row)?;
5878                        let (x1, ffn_out) =
5879                            self.residual_norm_ffn(e, layer, &x_row, mixed, n_embd, il, eps)?;
5880                        if spec_nan_scan_level() >= 2 {
5881                            nan_scan_rows(
5882                                e,
5883                                &ffn_out,
5884                                1,
5885                                n_embd,
5886                                &format!("tcol layer {il} col {r} per-column FFN out"),
5887                            )?;
5888                        }
5889                        let mut x2 = e.uninit(n_embd)?;
5890                        e.add(&x1, &ffn_out, &mut x2, n_embd)?;
5891                        e.dtod_copy_into(&x2, next, r * n_embd)?;
5892                        Ok(())
5893                    };
5894                    #[allow(clippy::needless_range_loop)]
5895                    // allow: the explicit index loop keeps the offset arithmetic visible and aligned with the device-side indexing
5896                    for r in 0..t {
5897                        e.dtod_copy_view(&h_t.slice(r * n_embd..(r + 1) * n_embd), &mut h_row)?;
5898                        let row_pos = &pos_rows[r];
5899                        crate::tp::set_verify_tcol(Some(r));
5900                        if fa2_layer {
5901                            crate::tp::set_spec_fa2_defer(Some(r));
5902                        } else if oproj_batch {
5903                            crate::tp::set_tcol_oproj_defer(Some(r));
5904                        }
5905                        let mixed = match &layer.mixer {
5906                            crate::hybrid::Mixer::Full(fa) => {
5907                                self.full_attn_decode(e, fa, &h_row, row_pos, pos0 + r, cache, il)
5908                            }
5909                            _ => Err("step35 verify expects full attention".into()),
5910                        };
5911                        crate::tp::set_verify_tcol(None);
5912                        crate::tp::set_spec_fa2_defer(None);
5913                        crate::tp::set_tcol_oproj_defer(None);
5914                        let mixed = mixed?;
5915                        if fa2_layer && crate::tp::take_spec_fa2_stashed() {
5916                            fa2_deferred.push(r);
5917                        } else if oproj_batch && crate::tp::take_tcol_oproj_stashed() {
5918                            deferred.push(r);
5919                        } else {
5920                            if spec_nan_scan_level() >= 2 {
5921                                let cols = mixed.len();
5922                                nan_scan_rows(
5923                                    e,
5924                                    &mixed,
5925                                    1,
5926                                    cols,
5927                                    &format!("tcol layer {il} col {r} per-column ATTN out"),
5928                                )?;
5929                            }
5930                            ffn_col(r, &mixed, &mut next)?;
5931                        }
5932                    }
5933                    if !fa2_deferred.is_empty() && fa2_deferred.len() != t {
5934                        // The precheck guarantees both columns stash or neither; a strict
5935                        // subset means a column's output was never produced anywhere.
5936                        return Err("spec fa2 stash engaged for a subset of columns".into());
5937                    }
5938                    if prof {
5939                        e.stream().synchronize()?;
5940                        prof_ms[1] += seg.elapsed().as_secs_f64() * 1e3;
5941                        seg = std::time::Instant::now();
5942                    }
5943                    if !fa2_deferred.is_empty() {
5944                        deferred = fa2_deferred;
5945                    }
5946                    if !deferred.is_empty() {
5947                        let mixed_t = if fa2_layer {
5948                            self.step35_verify_fa_rows_join(e, il, cache, pos0, t)?
5949                        } else {
5950                            self.step35_verify_oproj_tcol(e, il, t)?
5951                        };
5952                        let o_out = mixed_t.len() / t;
5953                        if spec_nan_scan_level() >= 2 {
5954                            nan_scan_rows(
5955                                e,
5956                                &mixed_t,
5957                                t,
5958                                o_out,
5959                                &format!("tcol layer {il} JOINED attn over deferred cols"),
5960                            )?;
5961                        }
5962                        // Batched t=2 residual+MoE: one t-grid add_rms_norm (per-row
5963                        // program == t=1; bit-identical to the oproj-tail join per the
5964                        // M2 verbatim-program contract) feeding the two-column routed
5965                        // sweep. Ineligible layers (dense FFN, non-nvfp4) fall through
5966                        // to the per-column body.
5967                        {
5968                            for &r in &deferred {
5969                                e.dtod_copy_view(
5970                                    &mixed_t.slice(r * o_out..(r + 1) * o_out),
5971                                    &mut mixed_row,
5972                                )?;
5973                                ffn_col(r, &mixed_row, &mut next)?;
5974                            }
5975                        }
5976                    }
5977                    if prof {
5978                        e.stream().synchronize()?;
5979                        prof_ms[2] += seg.elapsed().as_secs_f64() * 1e3;
5980                    }
5981                    x_t = next;
5982                    if spec_nan_scan() {
5983                        verify_arm_receipt(
5984                            if fa2_layer { "join" } else { "percol" },
5985                            il,
5986                            pos0,
5987                            t,
5988                            cache.tp_kv[il].as_ref().map(|d| d.staged_len()),
5989                        );
5990                        nan_scan_rows(
5991                            e,
5992                            &x_t,
5993                            t,
5994                            n_embd,
5995                            &format!(
5996                                "tcol layer {il} pos0={pos0} arm={}",
5997                                if fa2_layer { "join" } else { "percol" }
5998                            ),
5999                        )?;
6000                    }
6001                }
6002                if prof {
6003                    eprintln!(
6004                        "[tcol-prof] t={t} norm+qkv={:.3}ms attn={:.3}ms ffn={:.3}ms",
6005                        prof_ms[0], prof_ms[1], prof_ms[2]
6006                    );
6007                }
6008                if ok {
6009                    return Ok(x_t);
6010                }
6011                // fall through to the row-outer walk on ineligible layers
6012                x = x_t;
6013            }
6014            let mut next = e.uninit(t * n_embd)?;
6015            let scan = spec_nan_scan();
6016            for r in 0..t {
6017                let mut row = e.uninit(n_embd)?;
6018                e.dtod_copy_view(&x.slice(r * n_embd..(r + 1) * n_embd), &mut row)?;
6019                let row_pos = e.htod_i32(&[(pos0 + r) as i32])?;
6020                let out = if scan {
6021                    // Diagnostic arm: the same range walked one layer at a time so the first
6022                    // poisoned layer names itself. `decode_layers_eager(lo, hi)` is range-scoped
6023                    // and executes its trailing residual add, so a per-layer chain is the same
6024                    // program with the cross-layer add+norm fusion unrolled.
6025                    nan_scan_rows(
6026                        e,
6027                        &row,
6028                        1,
6029                        n_embd,
6030                        &format!("embed row r={r} pos={}", pos0 + r),
6031                    )?;
6032                    let mut acc = row;
6033                    for il in lo..hi {
6034                        acc = self.decode_layers_eager(
6035                            e,
6036                            acc,
6037                            il,
6038                            il + 1,
6039                            &row_pos,
6040                            pos0 + r,
6041                            cache,
6042                        )?;
6043                        nan_scan_rows(
6044                            e,
6045                            &acc,
6046                            1,
6047                            n_embd,
6048                            &format!("row-outer layer {il} r={r} pos={}", pos0 + r),
6049                        )?;
6050                    }
6051                    acc
6052                } else {
6053                    self.decode_layers_eager(e, row, lo, hi, &row_pos, pos0 + r, cache)?
6054                };
6055                e.dtod_copy_into(&out, &mut next, r * n_embd)?;
6056            }
6057            // dflash taps are NOT produced on this arm (they need per-layer hiddens the
6058            // row-outer walk does not materialize); the door is a step37 MTP bring-up
6059            // surface where taps are unused.
6060            return Ok(next);
6061        }
6062        let mut ph_last = std::time::Instant::now();
6063        for il in lo..hi {
6064            let mut next = e.uninit(t * n_embd)?;
6065            for r in 0..t {
6066                let mut row = e.uninit(n_embd)?;
6067                e.dtod_copy_view(&x.slice(r * n_embd..(r + 1) * n_embd), &mut row)?;
6068                // The caller owns this verify's position. During controller overlap, cache.pos
6069                // still describes generation N while this stage-0 walk belongs to N+1.
6070                let row_pos = e.htod_i32(&[(pos0 + r) as i32])?;
6071                let mut one = [&mut *cache];
6072                let out = self.step35_decode_batch_layers(
6073                    e,
6074                    row,
6075                    &mut one,
6076                    &[(pos0 + r) as i32],
6077                    &row_pos,
6078                    il,
6079                    il + 1,
6080                    &mut ph_last,
6081                )?;
6082                e.dtod_copy_into(&out, &mut next, r * n_embd)?;
6083            }
6084            self.dflash_tap(e, cache, il, &next, t)?;
6085            x = next;
6086            if spec_nan_scan() {
6087                nan_scan_rows(e, &x, t, n_embd, &format!("batch-layer {il} pos0={pos0}"))?;
6088            }
6089        }
6090        Ok(x)
6091    }
6092
6093    /// DSpark drafter verify (lane/dspark-q38-recover): one t-row forward through the
6094    /// SERVING-CLASS verify funnel (`decode_step_t_core_stream` — the same numeric class
6095    /// MTP verify rides, GDN state advanced in place), returning per-row argmax tokens.
6096    /// Advances `cache.pos += t`; the caller owns snapshot/rollback (block acceptance is
6097    /// prefix-keep, not all-or-nothing).
6098    pub(crate) fn dspark_verify_t_am(
6099        &self,
6100        e: &Engine,
6101        tokens: &[u32],
6102        pos0: usize,
6103        cache: &mut Cache,
6104    ) -> Result<Vec<u32>, Box<dyn std::error::Error>> {
6105        let (logits, _hn) =
6106            self.decode_step_t_core_stream(e, tokens, pos0, cache, None, None, None, None, None)?;
6107        let t = tokens.len();
6108        let v = self.output.out_features();
6109        let mut am_d = e.stream().alloc_zeros::<u32>(t)?;
6110        for r in 0..t {
6111            e.argmax_token_device_col(&logits, r, v, &mut am_d, r)?;
6112        }
6113        e.dtoh_u32(&am_d)
6114    }
6115
6116    /// DSpark verify returning the RAW verify logits [t, n_vocab] (device-resident) instead
6117    /// of per-row argmaxes — the sampled-admission arm's input (rejection-sampling accept
6118    /// gathers filtered p from these columns; lane/dspark-sampled-admission-20260820). Same
6119    /// forward as `dspark_verify_t_am`; the greedy arm keeps its argmax wrapper untouched.
6120    pub(crate) fn dspark_verify_t_logits(
6121        &self,
6122        e: &Engine,
6123        tokens: &[u32],
6124        pos0: usize,
6125        cache: &mut Cache,
6126    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6127        let (logits, _hn) =
6128            self.decode_step_t_core_stream(e, tokens, pos0, cache, None, None, None, None, None)?;
6129        Ok(logits)
6130    }
6131
6132    /// DSpark verify with the MTP column-stash armed: identical forward to
6133    /// `dspark_verify_t_am`, but fills a `VerifyCkpt` so a partial accept can restore
6134    /// column state directly (`dspark_commit_prefix`) instead of snapshot-replay.
6135    /// The ckpt type is opaque outside spec.rs (newtype) — dflash.rs threads it through.
6136    pub(crate) fn dspark_verify_t_am_ckpt(
6137        &self,
6138        e: &Engine,
6139        tokens: &[u32],
6140        pos0: usize,
6141        cache: &mut Cache,
6142    ) -> Result<(Vec<u32>, DsparkVerifyCkpt), Box<dyn std::error::Error>> {
6143        let mut ck = VerifyCkpt::new(self.layers.len());
6144        let (logits, _hn) = self.decode_step_t_core_stream(
6145            e,
6146            tokens,
6147            pos0,
6148            cache,
6149            None,
6150            Some(&mut ck),
6151            None,
6152            None,
6153            None,
6154        )?;
6155        let t = tokens.len();
6156        let v = self.output.out_features();
6157        let mut am_d = e.stream().alloc_zeros::<u32>(t)?;
6158        for r in 0..t {
6159            e.argmax_token_device_col(&logits, r, v, &mut am_d, r)?;
6160        }
6161        Ok((e.dtoh_u32(&am_d)?, DsparkVerifyCkpt(ck)))
6162    }
6163
6164    /// Engine-bundle slice 2: `dspark_verify_t_am_ckpt` with DEVICE tokens and NO readback.
6165    /// The verify tokens are the round's `chain_d` (cand layout: [anchor, drafts...]); the
6166    /// embed gathers its first `t` entries on-device (`embed_gather_u32_t` — bit-identical
6167    /// rows to the host arm), so the host never blocks on the draft chain before dispatching
6168    /// verify. Returns the device per-row argmax buffer; the caller merges its readback with
6169    /// the chain's into ONE sync. Forward, ckpt fill and argmax walk are `_ckpt` verbatim.
6170    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
6171    pub(crate) fn dspark_verify_t_am_ckpt_dev(
6172        &self,
6173        e: &Engine,
6174        vtok: &CudaSlice<u32>,
6175        t: usize,
6176        pos0: usize,
6177        cache: &mut Cache,
6178        embd_dev: (&CudaSlice<u8>, i32, usize),
6179        graphs: Option<&mut DsparkVerifyGraphs>,
6180    ) -> Result<(CudaSlice<u32>, DsparkVerifyCkpt), Box<dyn std::error::Error>> {
6181        debug_assert!(
6182            vtok.len() >= t,
6183            "verify window exceeds the device token buffer"
6184        );
6185        // The slab flag is a per-round statement: clear it here so a verify that never
6186        // reaches the graphs door (rowwise env, a non-tparallel arm) cannot leave a
6187        // stale `true` steering the commit at slabs the round never wrote.
6188        let mut graphs = graphs;
6189        if let Some(g) = graphs.as_deref_mut() {
6190            g.round_slab = false;
6191        }
6192        let mut ck = VerifyCkpt::new(self.layers.len());
6193        // Dummy host tokens size the funnel; the embed reads `vtok` (the round-stream
6194        // arm's established pattern — spec.rs stream-mode verify does the same).
6195        let dummy = vec![0u32; t];
6196        let (logits, _hn) = self.decode_step_t_core_stream(
6197            e,
6198            &dummy,
6199            pos0,
6200            cache,
6201            Some(embd_dev),
6202            Some(&mut ck),
6203            None,
6204            Some(vtok),
6205            graphs,
6206        )?;
6207        let v = self.output.out_features();
6208        let mut am_d = e.stream().alloc_zeros::<u32>(t)?;
6209        for r in 0..t {
6210            e.argmax_token_device_col(&logits, r, v, &mut am_d, r)?;
6211        }
6212        Ok((am_d, DsparkVerifyCkpt(ck)))
6213    }
6214
6215    /// Ckpt-armed twin of [`Self::dspark_verify_t_logits`] (sampled-admission arm).
6216    pub(crate) fn dspark_verify_t_logits_ckpt(
6217        &self,
6218        e: &Engine,
6219        tokens: &[u32],
6220        pos0: usize,
6221        cache: &mut Cache,
6222    ) -> Result<(CudaSlice<f32>, DsparkVerifyCkpt), Box<dyn std::error::Error>> {
6223        let mut ck = VerifyCkpt::new(self.layers.len());
6224        let (logits, _hn) = self.decode_step_t_core_stream(
6225            e,
6226            tokens,
6227            pos0,
6228            cache,
6229            None,
6230            Some(&mut ck),
6231            None,
6232            None,
6233            None,
6234        )?;
6235        Ok((logits, DsparkVerifyCkpt(ck)))
6236    }
6237
6238    /// Restore the round to `keep` accepted columns from the verify stash: KV lens and
6239    /// pos from the pre-verify snapshot + keep, GDN conv/ssm from the stashed column
6240    /// state — no replay forward. The exact `commit_verified_prefix` the MTP path ships.
6241    pub(crate) fn dspark_commit_prefix(
6242        &self,
6243        e: &Engine,
6244        cache: &mut Cache,
6245        snap: &crate::cache::CacheSnapshot,
6246        ckpt: &DsparkVerifyCkpt,
6247        keep: usize,
6248    ) -> Result<(), Box<dyn std::error::Error>> {
6249        self.commit_verified_prefix(e, cache, snap, &ckpt.0, keep)
6250    }
6251
6252    /// Slice-3 commit twin: restore to `keep` accepted columns when the round's linear
6253    /// column stash lives in the graphs ctx's persistent slabs (`DsparkVerifyGraphs`) —
6254    /// the cols arm's exact semantics (KV lens + pos from the snapshot, GDN conv/ssm
6255    /// from the stash of column keep-1), slab-addressed and batched into two copy
6256    /// launches.
6257    pub(crate) fn dspark_commit_prefix_slab(
6258        &self,
6259        e: &Engine,
6260        cache: &mut Cache,
6261        snap: &crate::cache::CacheSnapshot,
6262        ctx: &DsparkVerifyGraphs,
6263        keep: usize,
6264    ) -> Result<(), Box<dyn std::error::Error>> {
6265        use cudarc::driver::DevicePtr;
6266        debug_assert!(keep >= 1, "keep==0 rounds take the legacy rollback");
6267        let mut conv_src: Vec<u64> = Vec::new();
6268        let mut ssm_src: Vec<u64> = Vec::new();
6269        let mut conv_dst: Vec<u64> = Vec::new();
6270        let mut ssm_dst: Vec<u64> = Vec::new();
6271        for il in 0..self.layers.len() {
6272            if let (Some(kvl), Some(saved)) = (cache.kv[il].as_mut(), snap.kv_len[il]) {
6273                kvl.len = saved + keep;
6274                e.set_i32_one(&mut kvl.len_d, kvl.len as i32)?;
6275            }
6276            if let Some(rl) = cache.recur[il].as_ref() {
6277                let (pc, ps, _cw, _sw) = ctx
6278                    .slab_row(e, il, keep - 1)
6279                    .ok_or("slab commit: linear layer missing from the graphs ctx")?;
6280                conv_src.push(pc);
6281                ssm_src.push(ps);
6282                let st = &e.gpu.stream();
6283                let (dc, _g0) = rl.conv_state.device_ptr(st);
6284                let (ds, _g1) = rl.ssm_state.device_ptr(st);
6285                conv_dst.push(dc);
6286                ssm_dst.push(ds);
6287            }
6288        }
6289        let n = conv_src.len();
6290        if n > 0 {
6291            let mut tt = vec![0u64; 2 * n];
6292            tt[..n].copy_from_slice(&conv_src);
6293            tt[n..].copy_from_slice(&conv_dst);
6294            let ct = e.htod_u64(&tt)?;
6295            tt[..n].copy_from_slice(&ssm_src);
6296            tt[n..].copy_from_slice(&ssm_dst);
6297            let st = e.htod_u64(&tt)?;
6298            e.copy_batch_uniform_f32(&ct, n, ctx.conv_words)?;
6299            e.copy_batch_uniform_f32(&st, n, ctx.ssm_words)?;
6300        }
6301        cache.pos = snap.pos + keep;
6302        Ok(())
6303    }
6304
6305    /// Qwen35-family verify trunk in the live serving numeric class.
6306    ///
6307    /// Serving intentionally keeps this architecture in the generic batched program even at
6308    /// B=1. The older verify walk used its own mirrored dispatch and can flip near-tie argmaxes.
6309    ///
6310    /// Two arms, one numeric class:
6311    /// - DENSE GDN (`DenseMlp`, t<=16): `qwen35_verify_tparallel` — the weight ops (norms,
6312    ///   projections, FFN) hoist to m=T through the exact-tier batched kernels whose per-row
6313    ///   program IS the m=1 program (`matmul_pre == fused2 per (tensor,row); _bN mmvq per-row
6314    ///   == m=1` — decode_batch.rs v2 note), while the state ops (conv ring, gdn scan, KV
6315    ///   append, fa decode) stay a per-row loop running the b_n=1 serving kernels with each
6316    ///   row's own t_kv-driven arm pick (the straddle law: every row executes the exact
6317    ///   program its isolated serving step would). One weight read per layer per round
6318    ///   instead of T — this is what makes MTP profitable in the exact class (the per-row
6319    ///   walk measured verify(K+1) ~= (K+1) plain steps: 69 -> 44 tok/s served, 2026-08-15).
6320    /// - MoE / t>16 / `MEMRA_SPEC_VERIFY_ROWWISE=1`: the per-row replay of the authoritative
6321    ///   serving layer body, preserving single-session autoregressive cache order (the
6322    ///   correctness reference; also the rollback seam for the t-parallel arm).
6323    ///
6324    /// Bit-identity of the t-parallel arm vs the rowwise arm is gated by spec-serve-gate
6325    /// (zero differing logits at T=1..4, K arms) + the 8-prompt ON/OFF canary before ship.
6326    #[allow(clippy::too_many_arguments)]
6327    fn qwen35_verify_batch_layers(
6328        &self,
6329        e: &Engine,
6330        x: CudaSlice<f32>,
6331        lo: usize,
6332        hi: usize,
6333        pos0: usize,
6334        t: usize,
6335        cache: &mut Cache,
6336        ckpt: Option<&mut VerifyCkpt>,
6337        stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
6338        graphs: Option<&mut DsparkVerifyGraphs>,
6339    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6340        // Qwen35Moe admitted 2026-08-20 (lane/draftcost-moe): the t-parallel arm already
6341        // carries the MoE FFN (`moe_ffn_il_zq8` at m=T) and the GDN per-row state loop; the
6342        // arch fence was a qualification gate, not a mechanism gap. Measured disease on the
6343        // 35B-A3B class: rowwise verify ~= 5.6 ms per drafted token (one full trunk step
6344        // each) — the same (K+1)-plain-steps wall the dense admission fixed on 2026-08-15.
6345        // Rollback seam unchanged: MEMRA_SPEC_VERIFY_ROWWISE=1.
6346        let rowwise = std::env::var("MEMRA_SPEC_VERIFY_ROWWISE").as_deref() == Ok("1")
6347            || !self.batched_serving_numeric_class()
6348            || t > 16;
6349        if rowwise {
6350            if stream.is_some() {
6351                // rowwise replays per row with host cache.pos — irreconcilable with a
6352                // device position counter. Burst callers must keep t <= 16 and the
6353                // ROWWISE env unset; refusing beats silently mispositioned rows.
6354                return Err("qwen35 rowwise verify has no ROUND-STREAM arm \
6355                            (t > 16 or MEMRA_SPEC_VERIFY_ROWWISE=1)"
6356                    .into());
6357            }
6358            self.qwen35_verify_rowwise(e, x, lo, hi, pos0, t, cache, ckpt)
6359        } else {
6360            self.qwen35_verify_tparallel(e, x, lo, hi, pos0, t, cache, ckpt, stream, graphs)
6361        }
6362    }
6363
6364    /// The per-row correctness reference: replay each verify row through the authoritative
6365    /// serving layer body (`decode_batch_layers` at b_n=1). T full weight reads per layer.
6366    #[allow(clippy::too_many_arguments)]
6367    fn qwen35_verify_rowwise(
6368        &self,
6369        e: &Engine,
6370        mut x: CudaSlice<f32>,
6371        lo: usize,
6372        hi: usize,
6373        pos0: usize,
6374        t: usize,
6375        cache: &mut Cache,
6376        mut ckpt: Option<&mut VerifyCkpt>,
6377    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6378        let n_embd = self.cfg.n_embd as usize;
6379        let saved_pos = cache.pos;
6380        let mut ph_last = std::time::Instant::now();
6381        for il in lo..hi {
6382            let mut next = e.uninit(t * n_embd)?;
6383            let mut col_states: Option<Vec<(CudaSlice<f32>, CudaSlice<f32>)>> =
6384                if ckpt.is_some() && t >= 2 && matches!(self.layers[il].mixer, Mixer::Linear(_)) {
6385                    Some(Vec::with_capacity(t - 1))
6386                } else {
6387                    None
6388                };
6389            for r in 0..t {
6390                cache.pos = pos0 + r;
6391                let mut row = e.uninit(n_embd)?;
6392                e.dtod_copy_view(&x.slice(r * n_embd..(r + 1) * n_embd), &mut row)?;
6393                let row_pos = e.htod_i32(&[(pos0 + r) as i32])?;
6394                let mut one = [&mut *cache];
6395                let ctx = self.batch_layer_ctx(e, &one, il, il + 1)?;
6396                let out = match self.decode_batch_layers(
6397                    e,
6398                    row,
6399                    &mut one,
6400                    &ctx,
6401                    &row_pos,
6402                    &mut ph_last,
6403                ) {
6404                    Ok(out) => out,
6405                    Err(error) => {
6406                        cache.pos = saved_pos;
6407                        return Err(error);
6408                    }
6409                };
6410                e.dtod_copy_into(&out, &mut next, r * n_embd)?;
6411                if r + 1 < t
6412                    && let Some(states) = col_states.as_mut()
6413                {
6414                    let recur = cache.recur[il]
6415                        .as_ref()
6416                        .ok_or("Qwen35-MoE linear verify layer has no recurrent state")?;
6417                    states.push((
6418                        e.clone_dtod(&recur.conv_state)?,
6419                        e.clone_dtod(&recur.ssm_state)?,
6420                    ));
6421                }
6422            }
6423            if let (Some(checkpoint), Some(states)) = (ckpt.as_deref_mut(), col_states) {
6424                checkpoint.cols[il] = Some(states);
6425            }
6426            x = next;
6427        }
6428        cache.pos = saved_pos;
6429        Ok(x)
6430    }
6431
6432    /// T-PARALLEL VERIFY IN THE SERVING NUMERIC CLASS (lane/tparallel-verify, 2026-08-15).
6433    ///
6434    /// The weight ops run ONCE per layer at m=T; the state ops run per row through the same
6435    /// b_n=1 serving kernels the rowwise replay uses. Per-row bit-identity rests on the two
6436    /// pins the serving batch tier already carries:
6437    ///   * `matmul_pre` / `_bN` mmvq: per-row program == m=1 program (decode_batch.rs v2 note,
6438    ///     kernel-check pinned) — so a [T, n_embd] projection row equals the row projected
6439    ///     alone;
6440    ///   * row-indexed norms/elementwise (`rms_norm`, `quantize_q8_1`, `add_rms_norm`,
6441    ///     `gated_rmsnorm[_q8_1]`, `silu_mul`, `rope_neox` with per-row positions): the T-row
6442    ///     launch is the per-row program (same pin the generic verify's fused norms rely on).
6443    ///     The sequential dependencies keep their exact serving order: the conv ring / gdn scan
6444    ///     chain state row -> row through the `_b` kernels at b_n=1 (ping-pong via a 6-entry
6445    ///     alternating pointer table, host handles swapped per row so VerifyCkpt clones the
6446    ///     canonical state exactly as the rowwise arm does), and each row's KV append + fa decode
6447    ///     picks its arm from ITS OWN t_kv (append: format-only; fa: `fa_seqs_eligible` + its own
6448    ///     `fa_split_keys` rung at b_n=1) — the straddle law per row, so every row executes the
6449    ///     program its isolated B=1 serving step would.
6450    ///
6451    /// Cost: 1 weight read per layer per round + T state micro-launches, vs the rowwise arm's
6452    /// T weight reads. Gated bit-identical vs the rowwise arm by spec-serve-gate + canary.
6453    #[allow(clippy::too_many_arguments)]
6454    fn qwen35_verify_tparallel(
6455        &self,
6456        e: &Engine,
6457        mut x: CudaSlice<f32>,
6458        lo: usize,
6459        hi: usize,
6460        pos0: usize,
6461        t: usize,
6462        cache: &mut Cache,
6463        mut ckpt: Option<&mut VerifyCkpt>,
6464        stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
6465        mut graphs: Option<&mut DsparkVerifyGraphs>,
6466    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6467        // Merge guard (v0.98 train, re-affirmed on the v0.100 train over slice 4c): the
6468        // ROUND-STREAM arm (lane/draftcost-moe, device position counter) and the dspark
6469        // verify graphs (engine-bundle slice 3 / trunk slice 4c) have no common caller —
6470        // stream rides the qwen35moe burst, graphs ride the dspark route. If a future
6471        // caller arms both, refuse loudly instead of silently dropping the graphs ctx
6472        // (the stream linear arm takes linear_attn_verify_t, not the graphed segment or
6473        // full-verify bodies).
6474        if stream.is_some() && graphs.is_some() {
6475            return Err(
6476                "qwen35 tparallel verify: ROUND-STREAM and dspark verify graphs \
6477                        cannot arm together"
6478                    .into(),
6479            );
6480        }
6481        // Engine-bundle slice 3 + slice 4c: with a graphs ctx armed, pointer tables are
6482        // refreshed once per verify (the gdn ping-pong moves handles; a fresh generation
6483        // moves the kv caches). Then:
6484        //  - slice 4c: when the WHOLE round rides one seqs rung (every row batchable, one
6485        //    split-ladder step, rung covers the round), the ENTIRE walk replays as ONE
6486        //    full-verify graph per (vt, rung) — linear layers through the shared
6487        //    `qwen35_tparallel_linear_layer` body, full-attention layers through the
6488        //    shared `qwen35_tparallel_fa_layer` body in graph mode.
6489        //  - fallback (straddle rounds, below the vec floor, partial walks): runs of
6490        //    consecutive LINEAR layers replay the slice-3 per-(segment, vt) graphs and
6491        //    the full-attention layers run eager (batched rows when eligible).
6492        //
6493        // GRAPH-LAUNCH HEADROOM GUARD (see GRAPH_LAUNCH_MIN_FREE): the dspark verify
6494        // graphs replay through this walk from THREE callers — the MTP spec round's vg
6495        // door (already dropped per round by `graph_round_ok` before it gets here), the
6496        // dspark one-shot, and the dspark SERVE round (default ON since v0.108). Below
6497        // the driver-free floor the WHOLE round takes the byte-identical eager
6498        // cols-ckpt walk — the same drop-the-ctx fallback the pool ceiling already
6499        // takes — instead of feeding cuGraphLaunch a card it segfaults on.
6500        if let Some(g) = graphs.as_deref_mut()
6501            && !graph_launch_headroom_ok(e)
6502        {
6503            g.round_slab = false;
6504            graphs = None;
6505            static NOTED: std::sync::Once = std::sync::Once::new();
6506            NOTED.call_once(|| graph_replay_suspended_note("dspark-vg"));
6507        }
6508        if let Some(g) = graphs.as_deref_mut() {
6509            g.refresh_tables(e, cache)?;
6510            g.round_slab = false;
6511            if let Some(rung) = g.full_rung(self, cache, lo, hi, t) {
6512                // Pool ceiling (dspark_vg_cap): an existing key always replays; a NEW
6513                // full capture past the ceiling falls through to the segment/eager arms.
6514                if g.full.contains_key(&(t, rung, hi)) || g.can_capture() {
6515                    let out = g.run_full(self, e, lo, hi, &x, t, pos0, rung, cache)?;
6516                    g.round_slab = true;
6517                    return Ok(out);
6518                }
6519            }
6520            // Round-atomic ceiling check for the segment door: if any linear run in this
6521            // walk would need a NEW capture past the ceiling, the whole round runs the
6522            // eager cols-ckpt walk (mixing slab- and cols-stashed layers in one round
6523            // would corrupt the commit).
6524            if !g.segments_ready(self, lo, hi, t) {
6525                graphs = None;
6526            }
6527        }
6528        // STREAM (2b, lane/draftcost-moe): positions come from the device round counter
6529        // (pos_iota / i32_copy_add) so a burst round needs no host position knowledge.
6530        let pos_d = match stream {
6531            Some((_, ctr)) => {
6532                let mut p = e.alloc_uninit::<i32>(t)?;
6533                e.pos_iota(ctr, &mut p, t)?;
6534                p
6535            }
6536            None => {
6537                let pos_host: Vec<i32> = (0..t).map(|r| (pos0 + r) as i32).collect();
6538                e.htod_i32(&pos_host)?
6539            }
6540        };
6541        // Per-row 1-element position buffers, built ONCE per verify (the append/fa wrappers
6542        // take owned pos slices; building these inside the layer x row loops cost 16xT H2Ds).
6543        // LAZY since slice 4: the batched fa/append arm never touches them — they are built
6544        // on the first per-row fallback layer only (stream-aware there; the stream FA arm
6545        // rides the dc rows kernels and never reaches the fallback).
6546        let mut pos_rows: Option<Vec<CudaSlice<i32>>> = None;
6547        let mut il = lo;
6548        while il < hi {
6549            if graphs.is_some() && matches!(self.layers[il].mixer, Mixer::Linear(_)) {
6550                let mut end = il;
6551                while end < hi && matches!(self.layers[end].mixer, Mixer::Linear(_)) {
6552                    end += 1;
6553                }
6554                let g = graphs.as_deref_mut().expect("checked above");
6555                x = g.run_segment(self, e, il, end, &x, t, cache)?;
6556                g.round_slab = true;
6557                il = end;
6558                continue;
6559            }
6560            let layer = &self.layers[il];
6561            if stream.is_none() && matches!(layer.mixer, Mixer::Linear(_)) {
6562                // Eager linear layer (no graphs ctx): the shared body, legacy cols-ckpt arm.
6563                // Under ROUND-STREAM the linear layers ride the fa-body match's stream arm
6564                // below (linear_attn_verify_t — the stream COMMIT needs its GdnStash).
6565                x = self.qwen35_tparallel_linear_layer(
6566                    e,
6567                    il,
6568                    &x,
6569                    t,
6570                    cache,
6571                    ckpt.as_deref_mut(),
6572                    None,
6573                    None,
6574                )?;
6575                il += 1;
6576                continue;
6577            }
6578            // Full-attention (or stream-Linear, or MLA-refusing) layer: the extracted
6579            // shared body — eager arm (fresh per-verify pos/table, exact t_kv sizing,
6580            // in-body len bump). The slice-4c captured full-verify graphs run the SAME
6581            // body in graph mode; under ROUND-STREAM the body's dc-rows / GDN stream arms
6582            // run (lane/draftcost-moe).
6583            x = self.qwen35_tparallel_fa_layer(
6584                e,
6585                il,
6586                &x,
6587                t,
6588                cache,
6589                FaLayerArgs {
6590                    pos_d: &pos_d,
6591                    pos_rows: &mut pos_rows,
6592                    pos0,
6593                    graph_cap: None,
6594                    stream,
6595                    ckpt: ckpt.as_deref_mut(),
6596                },
6597            )?;
6598            il += 1;
6599        }
6600        Ok(x)
6601    }
6602
6603    /// SHARED dense-FFN body for the qwen35 t-parallel layers (trunk-kernels slice B) —
6604    /// ONE copy for the fa and linear layer bodies (the verify_layers extraction lesson).
6605    /// Dual arm (MEMRA_TK_FFN_DUAL, default on): gate+up in ONE dual launch from the
6606    /// pre-quantized activation with macro-scales DEFERRED into the fused SwiGLU+q8_1
6607    /// epilogue, then ffn_down from the fused (aq, ad) — the q27 verify chain verbatim.
6608    /// Every door is the bit-identical proven one: `matmul_decode_exact_dual_pre` (per
6609    /// (tensor,token,row) == the two singles), `silu_mul_scaled_q8_1` (y*s inline == the
6610    /// scale_inplace store, value-exact; fused quantize == quantize_q8_1 bytes),
6611    /// `matmul_decode_exact_pre` (dispatch mirror of the singles' q8_1-fast tail).
6612    /// Dual-refused (t outside 2..=7, non-NVFP4, layout mismatch) or seam off -> the
6613    /// original singles chain, byte-for-byte.
6614    #[allow(clippy::too_many_arguments)]
6615    fn qwen35_tparallel_dense_ffn(
6616        &self,
6617        e: &Engine,
6618        ffn_gate: &crate::model::GpuTensor,
6619        ffn_up: &crate::model::GpuTensor,
6620        ffn_down: &crate::model::GpuTensor,
6621        zn: &CudaSlice<f32>,
6622        t: usize,
6623        n_embd: usize,
6624    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6625        let n_ff = ffn_gate.out_features();
6626        let (zq, zd) = e.quantize_q8_1(zn, t, n_embd)?;
6627        if Engine::tk_ffn_dual_on()
6628            && let Some(((g, gs), (u, us))) =
6629                e.matmul_decode_exact_dual_pre(ffn_gate, ffn_up, &zq, &zd, t)?
6630        {
6631            if e.uses_q8_1_fast(ffn_down) {
6632                let (aq, ad) = e.silu_mul_scaled_q8_1(&g, &u, gs, us, t * n_ff)?;
6633                return e.matmul_decode_exact_pre(ffn_down, &aq, &ad, t);
6634            }
6635            let mut act = e.uninit(t * n_ff)?;
6636            e.silu_mul_scaled(&g, &u, gs, us, &mut act, t * n_ff)?;
6637            let (aq, ad) = e.quantize_q8_1(&act, t, n_ff)?;
6638            return e.matmul_pre(ffn_down, &aq, &ad, &act, t);
6639        }
6640        // v1 singles chain (seam off or dual-refused) — the pre-slice-B body verbatim.
6641        let g = e.matmul_pre(ffn_gate, &zq, &zd, zn, t)?;
6642        let u = e.matmul_pre(ffn_up, &zq, &zd, zn, t)?;
6643        let mut act = e.uninit(t * n_ff)?;
6644        e.silu_mul(&g, &u, &mut act, t * n_ff)?;
6645        let (aq, ad) = e.quantize_q8_1(&act, t, n_ff)?;
6646        e.matmul_pre(ffn_down, &aq, &ad, &act, t)
6647    }
6648
6649    /// ONE t-parallel FULL-ATTENTION layer (attn_norm + fa mixer + post_attn_norm + FFN +
6650    /// tap) — extracted from the walk exactly like `qwen35_tparallel_linear_layer` so the
6651    /// eager walk and the slice-4c captured full-verify graphs execute the SAME body (a
6652    /// second copy is how dispatch mirrors drift — the verify_layers extraction lesson).
6653    ///
6654    /// `args.graph_cap = Some((table, off, rung_end))` is the captured-graph mode:
6655    /// - kv base-pointer pairs come from the ctx-owned persistent table at `off` (a fresh
6656    ///   generation's cache lands at new addresses that only the per-verify table refresh
6657    ///   knows — the slice-3 baked-address lesson);
6658    /// - the seqs twins size partials/grid at `rung_end` and pin `split_keys` to the
6659    ///   rung's ladder value: `n_splits_max` is pure stride, splits >= ns_eff write the
6660    ///   EMPTY partial the combine never reads, and every per-row T_kv derives in-kernel
6661    ///   from `pos_seq[z]` — so one captured launch replays bit-identically for every
6662    ///   round whose rows all sit inside the rung;
6663    /// - the host len bump moves to the replay caller (captured host code does not
6664    ///   re-run at replay).
6665    ///   Graph mode REFUSES any round the batched arm cannot take: the per-row fallback
6666    ///   host-branches on t_kv and must never be captured.
6667    #[allow(clippy::too_many_arguments)]
6668    fn qwen35_tparallel_fa_layer(
6669        &self,
6670        e: &Engine,
6671        il: usize,
6672        x: &CudaSlice<f32>,
6673        t: usize,
6674        cache: &mut Cache,
6675        args: FaLayerArgs<'_>,
6676    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6677        use cudarc::driver::DevicePtr;
6678        let cfg = &self.cfg;
6679        let n_embd = cfg.n_embd as usize;
6680        let eps = cfg.rms_eps;
6681        let head_dim_global = cfg.head_dim_k as usize;
6682        let layer = &self.layers[il];
6683        let FaLayerArgs {
6684            pos_d,
6685            pos_rows,
6686            pos0,
6687            graph_cap,
6688            stream,
6689            ckpt,
6690        } = args;
6691
6692        // ---- attn_norm + q8_1 quantize at m=T (row-indexed == per-row) ----
6693        let anorm = layer.attn_norm.float_data();
6694        let mut xn = e.uninit(t * n_embd)?;
6695        e.rms_norm(x, anorm, &mut xn, n_embd, t, eps)?;
6696        let (hq, hd) = e.quantize_q8_1(&xn, t, n_embd)?;
6697
6698        let mixed: CudaSlice<f32> = match &layer.mixer {
6699            Mixer::Mla(_) => crate::hybrid::mla_path_unimplemented("tensor-parallel attention"),
6700            Mixer::Kda(_) => crate::hybrid::kda_path_unimplemented("T-parallel attention"),
6701            // STREAM ARM (2b, lane/draftcost-moe): under a device position counter the
6702            // per-row serving-kernel chain cannot run (host state swaps keyed on host
6703            // row index are fine, but the stream COMMIT needs the GdnStash for its _dc
6704            // rebuild — the per-row chain only produces per-column clones). GDN rides
6705            // `linear_attn_verify_t`: batched q8_1-class projections, stash-producing,
6706            // and its one-scan recurrence is pinned bit-identical to T chained T=1
6707            // steps (its header + kernel-check). Position-independent, so no counter
6708            // plumbing is needed. Guards mirror the generic call site exactly.
6709            Mixer::Linear(la) if stream.is_some() => {
6710                if !(t >= 3 || (t == 2 && spec_m2()))
6711                    || !self.mixer_in_q8_1_fast(e, &layer.mixer)
6712                    || !e.uses_q8_1_fast(&la.ssm_out)
6713                {
6714                    return Err("qwen35 stream verify: GDN batched arm requires t>=3 \
6715                                (or MEMRA_SPEC_M2 at t=2) and q8_1-fast projections"
6716                        .into());
6717                }
6718                let want = ckpt.is_some();
6719                let (out, stash) =
6720                    self.linear_attn_verify_t(e, la, &xn, Some((&hq, &hd)), t, cache, il, want)?;
6721                if let (Some(ck), Some(st)) = (ckpt, stash) {
6722                    ck.gdn[il] = Some(st);
6723                }
6724                out
6725            }
6726            Mixer::Linear(_) => {
6727                unreachable!("linear layers ride qwen35_tparallel_linear_layer")
6728            }
6729            Mixer::Full(fa) => {
6730                let geometry = cfg.full_attention_geometry_at(il as u32);
6731                let n_head = geometry.n_head as usize;
6732                let n_head_kv = geometry.n_head_kv as usize;
6733                let head_dim = geometry.head_dim_k as usize;
6734                let rope_dims = geometry.n_rot as usize;
6735                let rope_base = geometry.rope_base;
6736                let scale = geometry.attention_scale();
6737                // Batched projections: one weight read serves all T rows.
6738                // GROUP-3 twin (trunk-kernels slice D): q/k/v in ONE launch — the group4
6739                // kernel with n3=0, bit-identical per (tensor, token, row) to the three
6740                // singles; refused or MEMRA_TK_FA_GROUP=0 -> singles byte-for-byte.
6741                let (qf, mut k, v) = match e.matmul_decode_exact_group3_pre(
6742                    [&fa.wq, &fa.wk, &fa.wv],
6743                    &hq,
6744                    &hd,
6745                    t,
6746                )? {
6747                    Some(mut g3) => {
6748                        let v = g3.pop().unwrap();
6749                        let k = g3.pop().unwrap();
6750                        let qf = g3.pop().unwrap();
6751                        (qf, k, v)
6752                    }
6753                    None => (
6754                        e.matmul_pre(&fa.wq, &hq, &hd, &xn, t)?,
6755                        e.matmul_pre(&fa.wk, &hq, &hd, &xn, t)?,
6756                        e.matmul_pre(&fa.wv, &hq, &hd, &xn, t)?,
6757                    ),
6758                };
6759                let gated =
6760                    geometry.attention_gate == memra_gguf::config::AttentionGateKind::FusedQ;
6761                let (mut q, gate) = if gated {
6762                    let mut qs = e.uninit(t * n_head * head_dim)?;
6763                    let mut gs = e.uninit(t * n_head * head_dim)?;
6764                    e.q_gate_split(&qf, &mut qs, &mut gs, head_dim, n_head, t)?;
6765                    (qs, Some(gs))
6766                } else {
6767                    (qf, None)
6768                };
6769                let mut qn = e.uninit(t * n_head * head_dim)?;
6770                e.rms_norm(
6771                    &q,
6772                    fa.q_norm.float_data(),
6773                    &mut qn,
6774                    head_dim,
6775                    t * n_head,
6776                    eps,
6777                )?;
6778                q = qn;
6779                let mut kn = e.uninit(t * n_head_kv * head_dim)?;
6780                e.rms_norm(
6781                    &k,
6782                    fa.k_norm.float_data(),
6783                    &mut kn,
6784                    head_dim,
6785                    t * n_head_kv,
6786                    eps,
6787                )?;
6788                k = kn;
6789                e.rope_neox(
6790                    &mut q, pos_d, head_dim, rope_dims, n_head, t, rope_base, 1.0,
6791                )?;
6792                e.rope_neox(
6793                    &mut k, pos_d, head_dim, rope_dims, n_head_kv, t, rope_base, 1.0,
6794                )?;
6795
6796                // Per-row append + attend: row r sees rows 0..r in KV (causal within the
6797                // draft), each through the b_n=1 serving kernels at its own t_kv.
6798                let q_dim = n_head * head_dim;
6799                let kv_dim = n_head_kv * head_dim;
6800                let mut attn = e.uninit(t * q_dim)?;
6801                let (kdk, kdv, ktb, vtb, len0, kv_local) = {
6802                    let kvl = cache.kv[il].as_ref().unwrap();
6803                    // [2T] interleaved k,v base pointers: entry pair z serves row z of
6804                    // the batched twins; the per-row fallback reads pair 0 (same cache
6805                    // for every row of one layer). Graph mode reads the ctx table.
6806                    let local: Option<CudaSlice<u64>> = match graph_cap {
6807                        Some(_) => None,
6808                        None => {
6809                            let s = &e.gpu.stream();
6810                            let (pk, _g) = kvl.k.device_ptr(s);
6811                            let (pv, _g2) = kvl.v.device_ptr(s);
6812                            let mut tbl = Vec::with_capacity(2 * t);
6813                            for _ in 0..t {
6814                                tbl.push(pk);
6815                                tbl.push(pv);
6816                            }
6817                            Some(e.htod_u64(&tbl)?)
6818                        }
6819                    };
6820                    (
6821                        kvl.kv_dim_k,
6822                        kvl.kv_dim_v,
6823                        kvl.k_tok_bytes,
6824                        kvl.v_tok_bytes,
6825                        kvl.len,
6826                        local,
6827                    )
6828                };
6829                let (kv_tbl, kv_off): (&CudaSlice<u64>, usize) = match graph_cap {
6830                    Some((tb, off, _)) => (tb, off),
6831                    None => (kv_local.as_ref().expect("built above"), 0),
6832                };
6833                // Slice 4 (fa/append rows — see dspark_fa_rows_on): the whole per-row
6834                // section batches into the z-batched serving twins when every row of
6835                // this round takes the v4-seqs arm on ONE fa_split_keys rung. Both
6836                // guards are evaluated at the round's FIRST and LAST t_kv — the
6837                // eligibility window (vec floor .. v4 max) and each split-ladder rung
6838                // are intervals in t_kv, so ends-inside means all-inside (the straddle
6839                // law). Appending all T rows before any attend is read-equivalent to
6840                // the interleaved order: row r's walk reads keys 0..len0+r only, and
6841                // rows > r land at slots it never touches; every written cache row is
6842                // the per-token appender's exact warp program (kernel-check pinned).
6843                let t_kv_first = len0 + 1;
6844                let t_kv_last = len0 + t;
6845                let rows_batched = t >= 2
6846                    && dspark_fa_rows_on()
6847                    // the z-batched twins read stacked rows at the CACHE's kv dims;
6848                    // the projection stack is [T, n_head_kv*head_dim] — they must be
6849                    // the same stride or row z misaligns (true for this family; the
6850                    // guard keeps any asymmetric-kv model on the per-row loop).
6851                    && kdk == kv_dim
6852                    && kdv == kv_dim
6853                    && crate::fa_seqs_eligible(t_kv_first, head_dim_global)
6854                    && crate::fa_seqs_eligible(t_kv_last, head_dim_global)
6855                    && crate::fa_split_keys(t_kv_first, cfg.n_head_kv as usize)
6856                        == crate::fa_split_keys(t_kv_last, cfg.n_head_kv as usize);
6857                // Sizing: eager = exact round bound; graph mode = the rung end (stride +
6858                // grid only — bytes proven equal above). Capture-time invariants refuse
6859                // loudly rather than bake a divergent body.
6860                let (size_kv_max, sp) = match graph_cap {
6861                    Some((_, _, rung)) => {
6862                        if !rows_batched {
6863                            return Err(format!(
6864                                "fa graph capture: layer {il} round is not batchable \
6865                                 (t_kv {t_kv_first}..{t_kv_last}) — the per-row fallback \
6866                                 must never be captured"
6867                            )
6868                            .into());
6869                        }
6870                        let sp_r = crate::fa_split_keys(rung, cfg.n_head_kv as usize);
6871                        if t_kv_last > rung
6872                            || sp_r != crate::fa_split_keys(t_kv_last, cfg.n_head_kv as usize)
6873                        {
6874                            return Err(format!(
6875                                "fa graph capture: rung {rung} does not cover round \
6876                                 t_kv {t_kv_first}..{t_kv_last} on one split ladder step"
6877                            )
6878                            .into());
6879                        }
6880                        (rung, sp_r)
6881                    }
6882                    None => (
6883                        t_kv_last,
6884                        crate::fa_split_keys(t_kv_last, cfg.n_head_kv as usize),
6885                    ),
6886                };
6887                if let Some((_, ctr)) = stream {
6888                    // STREAM ARM (2b): one batched dc append + the multi-row dc attention
6889                    // — the generic stream arm's exact shape (rows kernels are pinned
6890                    // byte-identical to the per-row programs by kernel-check). Host len
6891                    // stays a stale lower bound; the burst drain reconciles it.
6892                    let kvl = cache.kv[il].as_mut().unwrap();
6893                    e.append_kv_quantized_rows_dc(
6894                        &k, &v, &mut kvl.k, &mut kvl.v, ctr, t, kdk, kdv, ktb, vtb, false,
6895                    )?;
6896                    let upper = (kvl.len + t + 64).min(cache.max_ctx);
6897                    let k_view = e.view_u8(&kvl.k, upper * ktb);
6898                    let v_view = e.view_u8(&kvl.v, upper * vtb);
6899                    e.fa_decode_rows_dc(
6900                        &q, &k_view, &v_view, &mut attn, head_dim, n_head, n_head_kv, ctr, upper,
6901                        t, scale, ktb, vtb, 0, false,
6902                    )?;
6903                } else if rows_batched {
6904                    e.append_kv_quantized_seqs(
6905                        &k,
6906                        &v,
6907                        &kv_tbl.slice(kv_off..kv_off + 2 * t),
6908                        pos_d,
6909                        t,
6910                        kdk,
6911                        kdv,
6912                        ktb,
6913                        vtb,
6914                    )?;
6915                    if graph_cap.is_none() {
6916                        cache.kv[il].as_mut().unwrap().len += t;
6917                    }
6918                    e.fa_decode_batch_seqs_v4(
6919                        &q,
6920                        &kv_tbl.slice(kv_off..kv_off + 2 * t),
6921                        pos_d,
6922                        &mut attn,
6923                        head_dim,
6924                        n_head,
6925                        n_head_kv,
6926                        t,
6927                        size_kv_max,
6928                        scale,
6929                        sp,
6930                        ktb,
6931                        vtb,
6932                    )?;
6933                } else {
6934                    if pos_rows.is_none() {
6935                        // Stream-aware for symmetry with pos_d (the stream FA arm rides
6936                        // the dc rows kernels above and never reaches this fallback).
6937                        *pos_rows = Some(match stream {
6938                            Some((_, ctr)) => (0..t)
6939                                .map(|r| {
6940                                    let mut b = e.alloc_uninit::<i32>(1)?;
6941                                    e.i32_copy_add(ctr, &mut b, r as i32)?;
6942                                    Ok(b)
6943                                })
6944                                .collect::<Result<_, Box<dyn std::error::Error>>>()?,
6945                            None => (0..t)
6946                                .map(|r| e.htod_i32(&[(pos0 + r) as i32]))
6947                                .collect::<Result<_, _>>()?,
6948                        });
6949                    }
6950                    let pos_rows = pos_rows.as_ref().unwrap();
6951                    #[allow(clippy::needless_range_loop)]
6952                    // allow: the explicit index loop keeps the offset arithmetic visible and aligned with the device-side indexing
6953                    for r in 0..t {
6954                        // Owned per-row scratch: the b_n=1 kernels take packed batch buffers
6955                        // whose row 0 is this row (arithmetic-free materialization copies,
6956                        // same as decode's per-seq fallback arm).
6957                        let mut k_row = e.uninit(kv_dim)?;
6958                        e.dtod_copy_view(&k.slice(r * kv_dim..(r + 1) * kv_dim), &mut k_row)?;
6959                        let mut v_row = e.uninit(kv_dim)?;
6960                        e.dtod_copy_view(&v.slice(r * kv_dim..(r + 1) * kv_dim), &mut v_row)?;
6961                        let pos_row = &pos_rows[r];
6962                        let kvl = cache.kv[il].as_mut().unwrap();
6963                        e.append_kv_quantized_seqs(
6964                            &k_row,
6965                            &v_row,
6966                            &kv_tbl.slice(kv_off..kv_off + 2),
6967                            pos_row,
6968                            1,
6969                            kdk,
6970                            kdv,
6971                            ktb,
6972                            vtb,
6973                        )?;
6974                        kvl.len += 1;
6975                        let t_kv = kvl.len;
6976                        let mut q_row = e.uninit(q_dim)?;
6977                        e.dtod_copy_view(&q.slice(r * q_dim..(r + 1) * q_dim), &mut q_row)?;
6978                        let mut a_row = e.uninit(q_dim)?;
6979                        if crate::fa_seqs_eligible(t_kv, head_dim_global) {
6980                            let sp0_r = crate::fa_split_keys(t_kv, cfg.n_head_kv as usize);
6981                            e.fa_decode_batch_seqs_v4(
6982                                &q_row,
6983                                &kv_tbl.slice(kv_off..kv_off + 2),
6984                                pos_row,
6985                                &mut a_row,
6986                                head_dim,
6987                                n_head,
6988                                n_head_kv,
6989                                1,
6990                                t_kv,
6991                                scale,
6992                                sp0_r,
6993                                ktb,
6994                                vtb,
6995                            )?;
6996                        } else {
6997                            let k_view = e.view_u8(&kvl.k, t_kv * kvl.k_tok_bytes);
6998                            let v_view = e.view_u8(&kvl.v, t_kv * kvl.v_tok_bytes);
6999                            let mut a_view = a_row.slice_mut(0..q_dim);
7000                            e.fa_decode_kvmod_view(
7001                                &q_row.slice(0..q_dim),
7002                                &k_view,
7003                                &v_view,
7004                                &mut a_view,
7005                                head_dim,
7006                                n_head,
7007                                n_head_kv,
7008                                t_kv,
7009                                scale,
7010                                kvl.k_tok_bytes,
7011                                kvl.v_tok_bytes,
7012                                false,
7013                            )?;
7014                        }
7015                        e.dtod_copy_into(&a_row, &mut attn, r * q_dim)?;
7016                    }
7017                }
7018
7019                // Output gate (element-wise) + o-proj at m=T.
7020                let attn_g = match &gate {
7021                    Some(g) => {
7022                        let n = t * q_dim;
7023                        let mut gsig = e.uninit(n)?;
7024                        e.sigmoid(g, &mut gsig, n)?;
7025                        let mut ag = e.uninit(n)?;
7026                        e.mul(&attn, &gsig, &mut ag, n)?;
7027                        ag
7028                    }
7029                    None => attn,
7030                };
7031                e.matmul(&fa.wo, &attn_g, t)?
7032            }
7033        };
7034
7035        // ---- residual add + post_attn_norm + FFN at m=T (serving dispatch verbatim) ----
7036        let pnorm = layer.post_attn_norm.float_data();
7037        let mut x1 = e.uninit(t * n_embd)?;
7038        let mut zn = e.uninit(t * n_embd)?;
7039        e.add_rms_norm(x, &mixed, pnorm, &mut x1, &mut zn, n_embd, t, eps)?;
7040        let ffn_out = match &layer.ffn {
7041            crate::hybrid::Ffn::Dense {
7042                ffn_gate,
7043                ffn_up,
7044                ffn_down,
7045            } => {
7046                assert!(
7047                    self.cfg.m3.is_none(),
7048                    "qwen35 t-parallel verify: M3 swigluoai FFN not yet batched"
7049                );
7050                self.qwen35_tparallel_dense_ffn(e, ffn_gate, ffn_up, ffn_down, &zn, t, n_embd)?
7051            }
7052            crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il_zq8(e, m, &zn, None, t, il as u16)?,
7053        };
7054        let mut x2 = e.uninit(t * n_embd)?;
7055        e.add(&x1, &ffn_out, &mut x2, t * n_embd)?;
7056        // dspark drafter tap (no-op when no sink armed): post-layer residual verify rows
7057        self.dflash_tap(e, cache, il, &x2, t)?;
7058        Ok(x2)
7059    }
7060
7061    /// ONE t-parallel LINEAR layer (attn_norm + gdn mixer + post_attn_norm + FFN + tap) —
7062    /// the exact body the old in-loop Linear arm ran, extracted so the eager walk and the
7063    /// slice-3 captured segments execute the SAME code (a second copy is how dispatch
7064    /// mirrors drift — the verify_layers extraction lesson). Two deliberate changes, both
7065    /// bit-identical by construction:
7066    /// - the gdn ping-pong host swap moves from per-row to ONE end-of-body swap (t odd):
7067    ///   the device sequence is driven entirely by the 6-entry pointer table, which
7068    ///   already encodes both parities; the ckpt stash reads name row r's out buffer
7069    ///   directly (r even -> alt handle, odd -> canonical) — the same physical bytes the
7070    ///   legacy post-swap clone read.
7071    /// - `stash` (slice-3 ctx): persistent per-layer slabs written by copy_into instead of
7072    ///   per-row clone_dtod allocs — same bytes, capture-legal (no per-round host objects).
7073    ///   `table_src` = (persistent pointer table, offset) when the ctx owns the tables;
7074    ///   None builds the per-verify table exactly as before.
7075    #[allow(clippy::too_many_arguments)]
7076    fn qwen35_tparallel_linear_layer(
7077        &self,
7078        e: &Engine,
7079        il: usize,
7080        x: &CudaSlice<f32>,
7081        t: usize,
7082        cache: &mut Cache,
7083        ckpt: Option<&mut VerifyCkpt>,
7084        stash: Option<(&mut CudaSlice<f32>, &mut CudaSlice<f32>)>,
7085        table_src: Option<(&CudaSlice<u64>, usize)>,
7086    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7087        use cudarc::driver::DevicePtr;
7088        let cfg = &self.cfg;
7089        let n_embd = cfg.n_embd as usize;
7090        let eps = cfg.rms_eps;
7091        let layer = &self.layers[il];
7092        let Mixer::Linear(la) = &layer.mixer else {
7093            return Err("qwen35_tparallel_linear_layer on a non-linear layer".into());
7094        };
7095        // ---- attn_norm + q8_1 quantize at m=T (row-indexed == per-row) ----
7096        let anorm = layer.attn_norm.float_data();
7097        let mut xn = e.uninit(t * n_embd)?;
7098        e.rms_norm(x, anorm, &mut xn, n_embd, t, eps)?;
7099        let (hq, hd) = e.quantize_q8_1(&xn, t, n_embd)?;
7100
7101        let geometry = la.geometry;
7102        let d_state = geometry.key_head_dim as usize;
7103        let num_k = geometry.key_heads as usize;
7104        let num_v = geometry.value_heads as usize;
7105        let d_conv = geometry.conv_kernel as usize;
7106        let key_dim = d_state * num_k;
7107        let value_dim = geometry.value_head_dim as usize * num_v;
7108        let conv_dim = key_dim * 2 + value_dim;
7109        let gdn_scale = 1.0 / (d_state as f32).sqrt();
7110
7111        // ---- batched projections: one weight read for all T rows ----
7112        // GROUP-4 twin (trunk-kernels slice C): the whole 4-tuple in ONE launch, bit-identical
7113        // per (tensor, token, row) to the four singles; refused (layout/tier) or
7114        // MEMRA_TK_GDN_GROUP=0 -> the singles chain byte-for-byte.
7115        let (qkv_mixed, z, beta_raw, alpha) = match e.matmul_decode_exact_group4_pre(
7116            [&la.wqkv, &la.wqkv_gate, &la.ssm_beta, &la.ssm_alpha],
7117            &hq,
7118            &hd,
7119            t,
7120        )? {
7121            Some(mut g4) => {
7122                let alpha = g4.pop().unwrap();
7123                let beta_raw = g4.pop().unwrap();
7124                let z = g4.pop().unwrap();
7125                let qkv_mixed = g4.pop().unwrap();
7126                (qkv_mixed, z, beta_raw, alpha)
7127            }
7128            None => (
7129                e.matmul_pre(&la.wqkv, &hq, &hd, &xn, t)?,
7130                e.matmul_pre(&la.wqkv_gate, &hq, &hd, &xn, t)?,
7131                e.matmul_pre(&la.ssm_beta, &hq, &hd, &xn, t)?,
7132                e.matmul_pre(&la.ssm_alpha, &hq, &hd, &xn, t)?,
7133            ),
7134        };
7135        let beta_w = la.ssm_beta.out_features();
7136        let alpha_w = la.ssm_alpha.out_features();
7137        let qkv_w = la.wqkv.out_features();
7138
7139        // ---- per-row state chain through the b_n=1 serving kernels ----
7140        // 6-entry alternating pointer table expresses the ping-pong without a rebuild per
7141        // row: even rows scan s0 -> s1, odd rows s1 -> s0.
7142        let table_local: Option<CudaSlice<u64>> = match table_src {
7143            Some(_) => None,
7144            None => {
7145                let rl = cache.recur[il].as_ref().unwrap();
7146                let s = &e.gpu.stream();
7147                let (pc, _g0) = rl.conv_state.device_ptr(s);
7148                let (p0, _g1) = rl.ssm_state.device_ptr(s);
7149                let (p1, _g2) = rl.ssm_state_alt.device_ptr(s);
7150                Some(e.htod_u64(&[pc, p0, p1, pc, p1, p0])?)
7151            }
7152        };
7153        let (table, toff): (&CudaSlice<u64>, usize) = match table_src {
7154            Some((tb, off)) => (tb, off),
7155            None => (table_local.as_ref().unwrap(), 0),
7156        };
7157        let mut o_all = e.uninit(t * value_dim)?;
7158        let mut col_states: Option<Vec<(CudaSlice<f32>, CudaSlice<f32>)>> =
7159            if ckpt.is_some() && stash.is_none() && t >= 2 {
7160                Some(Vec::with_capacity(t - 1))
7161            } else {
7162                None
7163            };
7164        let mut stash = stash;
7165        // Per-row scratch reused across rows (uninit is cheap but not free at
7166        // 48 layers x T rows); row inputs/outputs pass as VIEWS into the packed
7167        // [T, ...] buffers — zero arithmetic-free copies in this loop.
7168        let mut conv_out = e.uninit(conv_dim)?;
7169        let mut q_l2 = e.uninit(value_dim)?;
7170        let mut k_l2 = e.uninit(value_dim)?;
7171        let mut v_gd = e.uninit(value_dim)?;
7172        let mut beta_b = e.uninit(num_v)?;
7173        let mut g_log = e.uninit(num_v)?;
7174        for r in 0..t {
7175            let base = toff + if r % 2 == 0 { 0 } else { 3 };
7176            let conv_view = table.slice(base..base + 1);
7177            let in_view = table.slice(base + 1..base + 2);
7178            let out_view = table.slice(base + 2..base + 3);
7179            e.ssm_conv1d_fused_decode_b_view(
7180                &qkv_mixed.slice(r * qkv_w..(r + 1) * qkv_w),
7181                &conv_view,
7182                la.ssm_conv1d.float_data(),
7183                &mut conv_out,
7184                conv_dim,
7185                d_conv,
7186                1,
7187            )?;
7188            e.gdn_prep_decode_b_view(
7189                &conv_out,
7190                &beta_raw.slice(r * beta_w..(r + 1) * beta_w),
7191                &alpha.slice(r * alpha_w..(r + 1) * alpha_w),
7192                la.ssm_dt.float_data(),
7193                la.ssm_a.float_data(),
7194                &mut q_l2,
7195                &mut k_l2,
7196                &mut v_gd,
7197                &mut beta_b,
7198                &mut g_log,
7199                d_state,
7200                num_v,
7201                num_k,
7202                key_dim,
7203                eps,
7204                conv_dim,
7205                1,
7206            )?;
7207            let mut o_row = o_all.slice_mut(r * value_dim..(r + 1) * value_dim);
7208            e.gdn_scan_s128_batched_view(
7209                &q_l2, &k_l2, &v_gd, &g_log, &beta_b, &in_view, &out_view, &mut o_row, num_v, 1,
7210                gdn_scale,
7211            )?;
7212            if r + 1 < t {
7213                // Row r's out buffer: even rows write s1 (the alt handle — no swaps ran),
7214                // odd rows write s0 — the same physical state the legacy post-swap
7215                // canonical clone read.
7216                let rl = cache.recur[il]
7217                    .as_ref()
7218                    .ok_or("qwen35 linear verify layer has no recurrent state")?;
7219                let ssm_src = if r % 2 == 0 {
7220                    &rl.ssm_state_alt
7221                } else {
7222                    &rl.ssm_state
7223                };
7224                match stash.as_mut() {
7225                    Some((conv_slab, ssm_slab)) => {
7226                        // BOTH stash reads go through the pointer table at run time: the
7227                        // ssm handles ping-pong between rounds, and the ctx (with its
7228                        // captured graphs) outlives the Cache — a fresh generation's
7229                        // conv/ssm buffers land at new addresses that only the per-round
7230                        // table refresh knows. A baked direct copy would read freed
7231                        // memory (parity was the slice-3 smoke divergence; cache
7232                        // lifetime is the cross-generation twin).
7233                        e.copy_indirect_src_f32(
7234                            &conv_view,
7235                            conv_slab,
7236                            r * conv_dim * (d_conv - 1),
7237                            conv_dim * (d_conv - 1),
7238                        )?;
7239                        // The ssm handles PING-PONG between rounds: a captured direct
7240                        // copy would bake the capture-time physical buffer and read the
7241                        // wrong parity after any odd-vt round (the slice-3 smoke
7242                        // divergence). Read the src address from row r's OUT table
7243                        // entry at run time — the same entry the scan just wrote.
7244                        e.copy_indirect_src_f32(
7245                            &out_view,
7246                            ssm_slab,
7247                            r * d_state * d_state * num_v,
7248                            d_state * d_state * num_v,
7249                        )?;
7250                    }
7251                    None => {
7252                        if let Some(states) = col_states.as_mut() {
7253                            states.push((e.clone_dtod(&rl.conv_state)?, e.clone_dtod(ssm_src)?));
7254                        }
7255                    }
7256                }
7257            }
7258        }
7259        // ONE end-of-body parity swap (t odd) — the legacy loop swapped per row; the net
7260        // handle motion is identical and the device sequence never read the handles.
7261        if t % 2 == 1 {
7262            let rl = cache.recur[il].as_mut().unwrap();
7263            std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
7264        }
7265        if let (Some(checkpoint), Some(states)) = (ckpt, col_states) {
7266            checkpoint.cols[il] = Some(states);
7267        }
7268
7269        // ---- batched gated norm + out-projection at m=T ----
7270        let mixed = if e.uses_q8_1_fast(&la.ssm_out) {
7271            let (gq, gd) = e.gated_rmsnorm_q8_1(
7272                &o_all,
7273                la.ssm_norm.float_data(),
7274                &z,
7275                d_state,
7276                t * num_v,
7277                eps,
7278            )?;
7279            let g0 = e.zeros(0)?;
7280            e.matmul_pre(&la.ssm_out, &gq, &gd, &g0, t)?
7281        } else {
7282            let mut gn = e.uninit(t * value_dim)?;
7283            e.gated_rmsnorm(
7284                &o_all,
7285                la.ssm_norm.float_data(),
7286                &z,
7287                &mut gn,
7288                d_state,
7289                t * num_v,
7290                eps,
7291            )?;
7292            e.matmul(&la.ssm_out, &gn, t)?
7293        };
7294
7295        // ---- residual add + post_attn_norm + FFN at m=T (serving dispatch verbatim) ----
7296        let pnorm = layer.post_attn_norm.float_data();
7297        let mut x1 = e.uninit(t * n_embd)?;
7298        let mut zn = e.uninit(t * n_embd)?;
7299        e.add_rms_norm(x, &mixed, pnorm, &mut x1, &mut zn, n_embd, t, eps)?;
7300        let ffn_out = match &layer.ffn {
7301            crate::hybrid::Ffn::Dense {
7302                ffn_gate,
7303                ffn_up,
7304                ffn_down,
7305            } => {
7306                assert!(
7307                    self.cfg.m3.is_none(),
7308                    "qwen35 t-parallel verify: M3 swigluoai FFN not yet batched"
7309                );
7310                self.qwen35_tparallel_dense_ffn(e, ffn_gate, ffn_up, ffn_down, &zn, t, n_embd)?
7311            }
7312            crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il_zq8(e, m, &zn, None, t, il as u16)?,
7313        };
7314        let mut x2 = e.uninit(t * n_embd)?;
7315        e.add(&x1, &ffn_out, &mut x2, t * n_embd)?;
7316        // dspark drafter tap (no-op when no sink armed): post-layer residual verify rows
7317        self.dflash_tap(e, cache, il, &x2, t)?;
7318        Ok(x2)
7319    }
7320
7321    /// PP-N STAGE SUBGRAPH of the verify trunk: layers `[lo, hi)` of `decode_step_t_core_stream`'s
7322    /// walk, verbatim. Enters with a MATERIALIZED `[T, n_embd]` residual (no pending fusion pair
7323    /// carried in from outside the range) and exits with the range's final residual materialized
7324    /// (the trailing add executed) — exactly the `decode_layers_eager(lo, hi)` contract, T rows
7325    /// instead of one.
7326    ///
7327    /// EXTRACTED (lane/pp2-spec 2026-08-06) rather than duplicated: `decode_step_t_core_stream` IS
7328    /// the single funnel every verify forward reaches, and its per-layer dispatch MIRRORING (norm
7329    /// fusion per layer, the t>=3/spec_m2 batched-linear window, the fused-q8 FFN chain, the
7330    /// decode-exact projections) is what makes verify bit-identical to eager decode. A second copy
7331    /// for the split arm is how those mirrors drift apart on the next lever. The unsplit body now
7332    /// calls this with `(0, n_layers)`, so the whole-trunk path and every stage range run the SAME
7333    /// code — there is no "split version" of the verify math.
7334    ///
7335    /// Bit-identity of a cut rests on the same kernel-check-pinned identity the eager arm's cut
7336    /// does — `add_rms_norm_q8_1 == add then rms_norm_q8_1` at nrows=T — because the ONLY thing a
7337    /// fence changes is that the cross-layer fusion carry breaks at `hi-1` and is re-materialized
7338    /// as an explicit `add`. `decode-batch-gate --mode ppspec` verifies end-to-end on real weights.
7339    #[allow(clippy::too_many_arguments)]
7340    fn verify_layers(
7341        &self,
7342        e: &Engine,
7343        mut x: CudaSlice<f32>,
7344        lo: usize,
7345        hi: usize,
7346        pos_d: &CudaSlice<i32>,
7347        pos0: usize,
7348        t: usize,
7349        cache: &mut Cache,
7350        mut ckpt: Option<&mut VerifyCkpt>,
7351        stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
7352        graphs: Option<&mut DsparkVerifyGraphs>,
7353    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7354        if self.sliding_gated_moe_batch_program() {
7355            if stream.is_some() {
7356                return Err(
7357                    "step35 has no ROUND-STREAM verify arm (the device-counter _dc twins \
7358                            cannot express the SWA offset KV view)"
7359                        .into(),
7360                );
7361            }
7362            return self.step35_verify_batch_layers(e, x, lo, hi, pos0, t, cache);
7363        }
7364        if self.batched_serving_numeric_class() {
7365            return self.qwen35_verify_batch_layers(
7366                e,
7367                x,
7368                lo,
7369                hi,
7370                pos0,
7371                t,
7372                cache,
7373                ckpt.take(),
7374                stream,
7375                graphs,
7376            );
7377        }
7378        let n_embd = self.cfg.n_embd as usize;
7379        let eps = self.cfg.rms_eps;
7380        // CROSS-LAYER ADD+NORM FUSION (lane/vt-fixes fix 2, mirroring decode_step_h's
7381        // launch-arc form): layer il's post-FFN residual add (x2 = x1 + ffn_out) and layer
7382        // il+1's attn_norm(+quantize) are consecutive row-wise ops — ONE add_rms_norm_q8_1
7383        // launch at nrows=t does all three (bit-identity pinned by the T-row kernel-check
7384        // arms). Carry the un-added (x1, ffn_out) pair; the fused launch materializes x2 (the
7385        // residual the next layer needs) as its `res` output. Falls back to the separate add
7386        // when the next layer is off the fused-q8 path.
7387        let mut pending: Option<(CudaSlice<f32>, CudaSlice<f32>)> = None;
7388        for il in lo..hi {
7389            let layer = &self.layers[il];
7390            // DISPATCH-MIRRORED attn-input RMSNorm (FP-order lesson #8): eager decode fuses the
7391            // 1024-thread rms_norm_q8_1 ONLY when every mixer projection is q8_1-fast; layers with
7392            // Float projections (ssm_beta/ssm_alpha on layers 1/2/4 of the 9B NVFP4 GGUF) take the
7393            // UNFUSED 256-thread rms_norm. The verify norm must mirror that PER-LAYER choice —
7394            // blockDim changes the sum-of-squares reduce order, and the ULP shift amplifies through
7395            // the GDN recurrence into argmax flips (measured: 9B text prompt, 1 ULP at layer 2 ->
7396            // 2.3e-1 logit maxdiff at the head -> K=1..8 divergence at a 0.03-margin token).
7397            let mixer_fast = self.mixer_in_q8_1_fast(e, &layer.mixer);
7398            let norm_fused = std::env::var("MEMRA_NO_FUSE_NORMQ").is_err() && mixer_fast;
7399            // BATCHED EPILOGUE RE-FUSE (lane/vt-fixes fix 2, 2026-08-03): when the norm is
7400            // dispatch-fused AND every consumer of `h` reads only its q8_1 form (Full mixer:
7401            // projections only; Linear mixer: the batched arm — the per-column fallback needs
7402            // f32 h), emit the attn-input norm DIRECTLY as q8_1 via `rms_norm_q8_1` at nrows=t
7403            // (row-indexed kernel — the T-row launch is the per-row m=1 program, kernel-check
7404            // pins bit-identity vs rms_norm_decode -> quantize_q8_1). Kills the standalone
7405            // quantize launch(es) + the f32 h HBM round-trip that decode never pays.
7406            // step35 (Full mixer) is the third case that needs f32 `h`: its verify arm is a
7407            // per-ROW replay of the eager decode mixer, whose `pre_q` contract is a single row —
7408            // a T-row q8_1 pair cannot be handed to it, and re-deriving per-row q8_1 from the f32
7409            // rows is exactly the dispatch being mirrored. Keep step35 on the unfused arm.
7410            let lin_q8_only = match &layer.mixer {
7411                Mixer::Linear(la) => {
7412                    (t >= 3 || (t == 2 && spec_m2())) && e.uses_q8_1_fast(&la.ssm_out)
7413                }
7414                Mixer::Full(_) if self.sliding_gated_moe_batch_program() => false,
7415                _ => true,
7416            };
7417            // NOTE decode.rs's take()-first lesson: take the pending pair BEFORE branching so
7418            // a non-fused layer still performs the residual add.
7419            let taken = pending.take();
7420            let (h, h_q8) = if norm_fused && lin_q8_only {
7421                let pair = match taken {
7422                    // fused add + attn_norm + q8_1: ONE launch resolves the carried residual
7423                    // AND emits this layer's mixer input pre-quantized. res -> x2 (= new x).
7424                    Some((x1p, f1p)) => {
7425                        let mut x2 = vbuf(e, t * n_embd)?; // fully written (res output)
7426                        let p = e.add_rms_norm_q8_1(
7427                            &x1p,
7428                            &f1p,
7429                            layer.attn_norm.float_data(),
7430                            &mut x2,
7431                            n_embd,
7432                            t,
7433                            eps,
7434                        )?;
7435                        x = x2;
7436                        p
7437                    }
7438                    None => e.rms_norm_q8_1(&x, layer.attn_norm.float_data(), n_embd, t, eps)?,
7439                };
7440                (e.zeros(0)?, Some(pair)) // h unused on this path (q8-only consumers)
7441            } else {
7442                if let Some((x1p, f1p)) = taken {
7443                    let mut x2 = vbuf(e, t * n_embd)?; // fully written by add
7444                    e.add(&x1p, &f1p, &mut x2, t * n_embd)?;
7445                    x = x2;
7446                }
7447                let mut h = vbuf(e, t * n_embd)?; // fully written by either rms_norm arm
7448                if norm_fused {
7449                    e.rms_norm_decode(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
7450                } else {
7451                    e.rms_norm(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
7452                }
7453                (h, None)
7454            };
7455            let h_q8_ref = h_q8.as_ref().map(|(q, d)| (q, d));
7456
7457            let mixed = match &layer.mixer {
7458                Mixer::Full(fa) => self.full_attn_verify(
7459                    e,
7460                    fa,
7461                    &h,
7462                    h_q8_ref,
7463                    pos_d,
7464                    t,
7465                    cache,
7466                    il,
7467                    stream.map(|(_, c)| c),
7468                )?,
7469                Mixer::Mla(_) => crate::hybrid::mla_path_unimplemented("speculative verify"),
7470                Mixer::Kda(_) => crate::hybrid::kda_path_unimplemented("speculative verify"),
7471                Mixer::Linear(la) => {
7472                    // BATCHED linear verify (2026-07-03, the MTP-profit lever): one T-token pass —
7473                    // batched projections (weight read ONCE, hits the m=2-4 weight-resident matvec),
7474                    // carried-state conv (ssm_conv1d_tm_state), GDN prep on the prefill kernels, and
7475                    // ONE gdn_scan whose internal sequential t-loop is the SAME recurrence as T
7476                    // chained T=1 steps (bit-identical). Falls back to the sequential per-column
7477                    // chain when T < d_conv-1 (conv ring update needs T >= pad) — or when ANY
7478                    // projection is off the q8_1 fast path: matmul_decode_exact would route a Float
7479                    // tensor to cuBLAS at m=t (different FP accumulation than eager's per-token
7480                    // GEMV), so mixed-dtype layers stay on the eager-identical per-column chain.
7481                    // MEMRA_SPEC_M2 (lane/spec-m2): the t==2 batch rides the same arm — the conv
7482                    // wrapper handles t<pad with a pure-copy ring rebuild; see spec_m2() header.
7483                    if (t >= 3 || (t == 2 && spec_m2()))
7484                        && mixer_fast
7485                        && e.uses_q8_1_fast(&la.ssm_out)
7486                    {
7487                        let want = ckpt.is_some();
7488                        let (out, stash) =
7489                            self.linear_attn_verify_t(e, la, &h, h_q8_ref, t, cache, il, want)?;
7490                        if let (Some(ck), Some(st)) = (ckpt.as_deref_mut(), stash) {
7491                            ck.gdn[il] = Some(st);
7492                        }
7493                        out
7494                    } else {
7495                        let mut out = vbuf(e, t * n_embd)?; // every col written by copy_into
7496                        let mut col_states: Option<Vec<(CudaSlice<f32>, CudaSlice<f32>)>> =
7497                            if ckpt.is_some() && t >= 2 {
7498                                Some(Vec::with_capacity(t - 1))
7499                            } else {
7500                                None
7501                            };
7502                        for col in 0..t {
7503                            let mut h_col = vbuf(e, n_embd)?; // fully written by copy_view_into
7504                            let src = h.slice(col * n_embd..(col + 1) * n_embd);
7505                            e.copy_view_into(&mut h_col, 0, &src, n_embd)?;
7506                            let m_col = self.linear_attn_decode(e, la, &h_col, cache, il)?;
7507                            e.copy_into(&mut out, col * n_embd, &m_col, n_embd)?;
7508                            // REPLAY-FREE ckpt: clone the chain's ACTUAL state after this column
7509                            // (pure dtod — cannot change any computed value). Last column skipped:
7510                            // rebuild targets are j <= t-1 columns.
7511                            if let Some(cs) = col_states.as_mut()
7512                                && col + 1 < t
7513                            {
7514                                let rl = cache.recur[il].as_ref().unwrap();
7515                                cs.push((
7516                                    e.clone_dtod(&rl.conv_state)?,
7517                                    e.clone_dtod(&rl.ssm_state)?,
7518                                ));
7519                            }
7520                        }
7521                        if let (Some(ck), Some(cs)) = (ckpt.as_deref_mut(), col_states) {
7522                            // ReplaySSM-assessment instrumentation (2026-07-30): the
7523                            // per-column clones are the only true state snapshots left in
7524                            // the verify (the batched path stashes INPUTS and replays).
7525                            if std::env::var("MEMRA_SPEC_STATS").as_deref() == Ok("1") {
7526                                static ONCE: std::sync::Once = std::sync::Once::new();
7527                                let bytes: usize =
7528                                    cs.iter().map(|(c, s)| (c.len() + s.len()) * 4).sum();
7529                                ONCE.call_once(|| eprintln!(
7530                                    "[verify-ckpt] per-column layer il={il}: {} clones, {:.2} MB/layer/round",
7531                                    cs.len(), bytes as f64 / 1e6));
7532                            }
7533                            ck.cols[il] = Some(cs);
7534                        }
7535                        out
7536                    }
7537                }
7538            };
7539            if spec_nan_scan_level() >= 2 {
7540                let mixed_width = mixed.len() / t;
7541                nan_scan_rows(
7542                    e,
7543                    &mixed,
7544                    t,
7545                    mixed_width,
7546                    &format!("verify layer {il} batched ATTN out pos0={pos0}"),
7547                )?;
7548            }
7549
7550            // DISPATCH-MIRRORED post-attn norm: eager residual_norm_ffn fuses add+norm+quant
7551            // (1024-thread add_rms_norm_q8_1) only for Dense FFNs whose gate+up are q8_1-fast;
7552            // otherwise (and for MoE) it runs the 256-thread fused add_rms_norm. Mirror per layer.
7553            let ffn_fuse = match &layer.ffn {
7554                crate::hybrid::Ffn::Dense {
7555                    ffn_gate, ffn_up, ..
7556                } => {
7557                    std::env::var("MEMRA_NO_FUSE_NORMQ").is_err()
7558                        && e.uses_q8_1_fast(ffn_gate)
7559                        && e.uses_q8_1_fast(ffn_up)
7560                }
7561                crate::hybrid::Ffn::Moe(_) => false,
7562            };
7563            // BATCHED EPILOGUE RE-FUSE (lane/vt-fixes fix 2): on the ffn_fuse path (Dense,
7564            // gate+up q8_1-fast, non-M3) the FFN input is emitted DIRECTLY as q8_1 by ONE
7565            // add_rms_norm_q8_1 launch at nrows=t (row-indexed kernel: the T-row launch is the
7566            // per-row m=1 program; kernel-check pins bit-identity vs the unfused
7567            // add_f32 -> rms_norm_decode -> quantize_q8_1 chain at T=2/4/5/8) — replacing the
7568            // add + rms_norm_decode launches AND the dual/singles' internal re-quantize.
7569            // M3's swigluoai must keep the f32 chain (the fused SwiGLU epilogue encodes plain
7570            // SiLU), mirroring residual_norm_ffn's m3 guard on the decode path.
7571            // step35: same guard per LAYER. A dense FFN's clamp is the SHEXP array (upstream's
7572            // one build_ffn serves dense + shared expert, llama-graph.cpp:1751), and verify MUST
7573            // mirror decode's dispatch or spec self-consistency fails.
7574            let dense_lim = self.cfg.clamp_shexp_at(il as u32);
7575            let fuse_q8 = ffn_fuse && self.cfg.m3.is_none() && dense_lim.is_none();
7576            let mut x1 = vbuf(e, t * n_embd)?; // fully written by add / add_rms_norm*
7577            let mut z = e.zeros(0)?; // replaced below on the unfused arms
7578            let z_q8 = if fuse_q8 {
7579                Some(e.add_rms_norm_q8_1(
7580                    &x,
7581                    &mixed,
7582                    layer.post_attn_norm.float_data(),
7583                    &mut x1,
7584                    n_embd,
7585                    t,
7586                    eps,
7587                )?)
7588            } else {
7589                let mut zf = vbuf(e, t * n_embd)?; // fully written by rms_norm_decode / add_rms_norm
7590                if ffn_fuse {
7591                    e.add(&x, &mixed, &mut x1, t * n_embd)?;
7592                    e.rms_norm_decode(
7593                        &x1,
7594                        layer.post_attn_norm.float_data(),
7595                        &mut zf,
7596                        n_embd,
7597                        t,
7598                        eps,
7599                    )?;
7600                } else {
7601                    e.add_rms_norm(
7602                        &x,
7603                        &mixed,
7604                        layer.post_attn_norm.float_data(),
7605                        &mut x1,
7606                        &mut zf,
7607                        n_embd,
7608                        t,
7609                        eps,
7610                    )?;
7611                }
7612                z = zf;
7613                None
7614            };
7615            if spec_nan_scan_level() >= 2 && !z.is_empty() {
7616                nan_scan_rows(
7617                    e,
7618                    &z,
7619                    t,
7620                    n_embd,
7621                    &format!("verify layer {il} post-attn norm z pos0={pos0}"),
7622                )?;
7623            }
7624            // DECODE-EXACT FFN projections: force MMVQ for gate/up/down at any T to match the
7625            // T=1 decode FP accumulation order. At T>=5 the generic matmul/matmul_pre falls to dp4a
7626            // (128-thread, different FP sum order). At T=2-4 the batched MMVQ is already bit-identical.
7627            let ffn_out = match &layer.ffn {
7628                crate::hybrid::Ffn::Dense {
7629                    ffn_gate,
7630                    ffn_up,
7631                    ffn_down,
7632                } => {
7633                    let n_ff = ffn_gate.out_features();
7634                    if let Some((zq, zd)) = z_q8.as_ref() {
7635                        // FUSED CHAIN (fix 2): pre-quantized z feeds the projections; the SwiGLU
7636                        // epilogue emits act pre-quantized for ffn_down (silu_mul_scaled_q8_1,
7637                        // bit-identical to silu_mul + quantize — kernel-check-pinned) with the
7638                        // NVFP4 macro-scales folded (deferred-scale dual: y*s inline == the
7639                        // scale_inplace store, value-exact) — the exact m=1 decode epilogue
7640                        // structure at nrows=t.
7641                        let pair = e
7642                            .matmul_decode_exact_dual_pre(ffn_gate, ffn_up, zq, zd, t)?
7643                            .map(|((g, gs), (u, us))| (g, gs, u, us));
7644                        let (gate, gs, up, us) = match pair {
7645                            Some(x4) => x4,
7646                            None => (
7647                                e.matmul_decode_exact_pre(ffn_gate, zq, zd, t)?,
7648                                1.0, // scale already applied inside _pre
7649                                e.matmul_decode_exact_pre(ffn_up, zq, zd, t)?,
7650                                1.0,
7651                            ),
7652                        };
7653                        if e.uses_q8_1_fast(ffn_down) {
7654                            let (aq, ad) = e.silu_mul_scaled_q8_1(&gate, &up, gs, us, t * n_ff)?;
7655                            e.matmul_decode_exact_pre(ffn_down, &aq, &ad, t)?
7656                        } else {
7657                            let mut act = vbuf(e, t * n_ff)?;
7658                            e.silu_mul_scaled(&gate, &up, gs, us, &mut act, t * n_ff)?;
7659                            e.matmul_decode_exact(ffn_down, &act, t)?
7660                        }
7661                    } else {
7662                        // UNFUSED (pre-fix) chain — MoE-adjacent/M3/off-fast layers, unchanged.
7663                        // DUAL gate+up batched twin (lane/verify-economics, 2026-08-02): one launch
7664                        // for the pair at t=2..8 — bit-identical per (tensor,token,row) to the two
7665                        // singles (kernel-check pins bitwise; MEMRA_SPEC_DUAL_T=0 reverts). None
7666                        // (non-NVFP4 / t outside the tier / seam off) -> the two singles, unchanged.
7667                        let (gate, up) =
7668                            match e.matmul_decode_exact_dual(ffn_gate, ffn_up, &z, t)? {
7669                                Some(pair) => pair,
7670                                None => (
7671                                    e.matmul_decode_exact(ffn_gate, &z, t)?,
7672                                    e.matmul_decode_exact(ffn_up, &z, t)?,
7673                                ),
7674                            };
7675                        let mut act = vbuf(e, t * n_ff)?; // fully written by ffn_act_lim
7676                        Self::ffn_act_lim(
7677                            e,
7678                            &self.cfg,
7679                            &gate,
7680                            &up,
7681                            1.0,
7682                            1.0,
7683                            dense_lim,
7684                            &mut act,
7685                            t * n_ff,
7686                        )?;
7687                        e.matmul_decode_exact(ffn_down, &act, t)?
7688                    }
7689                }
7690                crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il(e, m, &z, t, il as u16)?,
7691            };
7692            if spec_nan_scan_level() >= 2 {
7693                nan_scan_rows(
7694                    e,
7695                    &ffn_out,
7696                    t,
7697                    n_embd,
7698                    &format!("verify layer {il} batched FFN out pos0={pos0}"),
7699                )?;
7700            }
7701            if spec_nan_scan() {
7702                let mut residual = vbuf(e, t * n_embd)?;
7703                e.add(&x1, &ffn_out, &mut residual, t * n_embd)?;
7704                nan_scan_rows(
7705                    e,
7706                    &residual,
7707                    t,
7708                    n_embd,
7709                    &format!("verify layer {il} residual pos0={pos0}"),
7710                )?;
7711            }
7712            // CROSS-LAYER fusion: defer this layer's post-FFN residual add — the next layer's
7713            // fused-q8 attn norm folds it in (add_rms_norm_q8_1 == add; rms_norm; quantize,
7714            // kernel-check-pinned at nrows=T). Non-fused next layers add explicitly above.
7715            pending = Some((x1, ffn_out));
7716        }
7717        // RANGE's final add (no next norm INSIDE the range to fuse with; for the
7718        // whole-trunk call that is the last layer, whose next norm is output_norm — f32-out).
7719        if let Some((x1p, f1p)) = pending.take() {
7720            let mut x2 = vbuf(e, t * n_embd)?; // fully written by add
7721            e.add(&x1p, &f1p, &mut x2, t * n_embd)?;
7722            x = x2;
7723        }
7724        Ok(x)
7725    }
7726    /// BATCHED linear-attn verify (T=K+1): the whole layer in ~10 launches instead of T x the
7727    /// T=1 decode chain (T x ~12 launches + T weight reads of the four projections). The GDN
7728    /// recurrence itself is inherently sequential — gdn_scan_s128 runs its internal t-loop with
7729    /// the SAME per-token math as chained T=1 calls (bit-identical state evolution); everything
7730    /// around it (projections, conv, prep, gated norm, out-proj) batches. Advances conv ring +
7731    /// ssm state exactly like T sequential decode steps.
7732    /// `want_stash`: additionally RETAIN the gdn-scan inputs (pure buffer keep-alives, zero extra
7733    /// kernels) so a partial accept can rebuild the state after any column prefix (REPLAY-FREE).
7734    #[allow(clippy::too_many_arguments)]
7735    fn linear_attn_verify_t(
7736        &self,
7737        e: &Engine,
7738        la: &LinearAttnLayer,
7739        h: &CudaSlice<f32>,
7740        h_q8: Option<(&CudaSlice<i8>, &CudaSlice<f32>)>,
7741        t: usize,
7742        cache: &mut Cache,
7743        il: usize,
7744        want_stash: bool,
7745    ) -> Result<(CudaSlice<f32>, Option<GdnStash>), Box<dyn std::error::Error>> {
7746        let cfg = &self.cfg;
7747        let geometry = la.geometry;
7748        let d_state = geometry.key_head_dim as usize;
7749        let num_k = geometry.key_heads as usize;
7750        let num_v = geometry.value_heads as usize;
7751        let d_conv = geometry.conv_kernel as usize;
7752        let key_dim = d_state * num_k;
7753        let conv_dim = key_dim * 2 + geometry.value_head_dim as usize * num_v;
7754        let eps = cfg.rms_eps;
7755        let scale = 1.0 / (d_state as f32).sqrt();
7756
7757        // DECODE-EXACT projections: matmul_decode_exact forces the MMVQ (warp-per-row, 32-thread)
7758        // accumulation order for EVERY m, matching the T=1 decode path bit-for-bit. The generic
7759        // `matmul` at m>=5 falls to dp4a (128-thread, two-level reduce) which has a different FP
7760        // sum order — ULP differences propagate through gdn_scan and flip argmax on the 27B.
7761        // Q8 TRUNK-FUSION at T=1 (35B: wqkv+wqkv_gate both Q8_0): one fused2 launch, bit-identical
7762        // per (tensor,row) to the two m=1 MMVQ dispatches below — decode-exact contract holds.
7763        // VERIFY-TIER TRUNK FUSION (MEMRA_SPEC_FUSED_T, t=2-4): quantize h ONCE for every
7764        // fused-eligible same-input Q8_0 pair of this layer (35B wqkv+wqkv_gate; 9B
7765        // ssm_beta+ssm_alpha) — each fused2 batched launch then replaces two decode-exact
7766        // calls (each of which re-quantizes the same h + runs its own _b2/_b4 launch).
7767        // Bit-identical per (tensor,token,row) — see spec_fused_t().
7768        // BATCHED EPILOGUE RE-FUSE (lane/vt-fixes fix 2): `h_q8` = the attn-input norm emitted
7769        // directly as q8_1 by the caller's fused rms_norm_q8_1 (bit-identical to the unfused
7770        // chain, kernel-check-pinned). When present it REPLACES the standalone quantize below
7771        // and feeds every projection; the caller guaranteed all four input projections are
7772        // q8_1-fast. When absent, the old shared-quantize (fused-t window) stands.
7773        let h_q8_t = if h_q8.is_none()
7774            && spec_fused_t()
7775            && (2..=4).contains(&t)
7776            && ((e.uses_q8_1_fast(&la.wqkv) && e.uses_q8_1_fast(&la.wqkv_gate))
7777                || (e.uses_q8_1_fast(&la.ssm_beta) && e.uses_q8_1_fast(&la.ssm_alpha)))
7778        {
7779            Some(e.quantize_q8_1(h, t, cfg.n_embd as usize)?)
7780        } else {
7781            None
7782        };
7783        // one view: the caller's fused-norm q8 or this fn's own shared quantize.
7784        let hq8_any: Option<(&CudaSlice<i8>, &CudaSlice<f32>)> =
7785            h_q8.or(h_q8_t.as_ref().map(|(q, d)| (q, d)));
7786        let (qkv_mixed, z) = {
7787            let mut fused = None;
7788            if t == 1 && e.uses_q8_1_fast(&la.wqkv) && e.uses_q8_1_fast(&la.wqkv_gate) {
7789                let (hq, hd) = e.quantize_q8_1(h, 1, cfg.n_embd as usize)?;
7790                fused = e.matmul_q8_fused2(&la.wqkv, &la.wqkv_gate, &hq, &hd)?;
7791            } else if let Some((hq, hd)) = hq8_any
7792                && spec_fused_t()
7793                && (2..=4).contains(&t)
7794            {
7795                fused = e.matmul_q8_fused2_t(&la.wqkv, &la.wqkv_gate, hq, hd, t)?;
7796            }
7797            match (fused, hq8_any) {
7798                (Some(pair), _) => pair,
7799                (None, Some((hq, hd))) if h_q8.is_some() => (
7800                    e.matmul_decode_exact_pre(&la.wqkv, hq, hd, t)?,
7801                    e.matmul_decode_exact_pre(&la.wqkv_gate, hq, hd, t)?,
7802                ),
7803                (None, _) => (
7804                    e.matmul_decode_exact(&la.wqkv, h, t)?,
7805                    e.matmul_decode_exact(&la.wqkv_gate, h, t)?,
7806                ),
7807            }
7808        };
7809        // beta+alpha DUAL at T=1 (75% of p3 rounds run T=1 verify — p-min chain cuts): the dual
7810        // mr2 kernel is bit-identical per element to the m=1 MMVQ matmul_decode_exact dispatches
7811        // (same warp-per-row body, blockIdx.y picks the weight), so the decode-exact contract
7812        // holds; the run-spec battery is the arbiter. T>1 keeps the per-tensor decode-exact path.
7813        let (beta_raw, alpha) = if t == 1 {
7814            let (hq, hd) = e.quantize_q8_1(h, 1, cfg.n_embd as usize)?;
7815            match e.matmul_pre_dual_noscale(&la.ssm_beta, &la.ssm_alpha, &hq, &hd, 1)? {
7816                Some(((mut b, bs), (mut a, as_))) => {
7817                    if bs != 1.0 {
7818                        e.scale_inplace(&mut b, bs, la.ssm_beta.out_features())?;
7819                    }
7820                    if as_ != 1.0 {
7821                        e.scale_inplace(&mut a, as_, la.ssm_alpha.out_features())?;
7822                    }
7823                    (b, a)
7824                }
7825                // Q8_0 fused2 twin (9B stores beta/alpha as Q8_0): DISPATCH-MIRRORS the eager
7826                // decode's beta_alpha closure — the fused body is qmatvec_q8_0_mmvq verbatim,
7827                // bit-identical per row (kernel-check rel=0.00e0 gate), so decode==verify holds.
7828                None => match e.matmul_q8_fused2(&la.ssm_beta, &la.ssm_alpha, &hq, &hd)? {
7829                    Some((b, a)) => (b, a),
7830                    None => (
7831                        e.matmul_decode_exact(&la.ssm_beta, h, 1)?,
7832                        e.matmul_decode_exact(&la.ssm_alpha, h, 1)?,
7833                    ),
7834                },
7835            }
7836        } else {
7837            // fused-t twin (9B stores beta/alpha as Q8_0): same shared-quantize + one launch
7838            // contract as the wqkv pair above; 35B beta/alpha are Float -> None -> fallback.
7839            let mut nvfp4_fused = None;
7840            let mut q8_fused = None;
7841            if let Some((hq, hd)) = hq8_any {
7842                if t == 3 && std::env::var("MEMRA_NVFP4_AUX_DUAL").as_deref() != Ok("0") {
7843                    nvfp4_fused =
7844                        e.matmul_decode_exact_dual_pre(&la.ssm_beta, &la.ssm_alpha, hq, hd, t)?;
7845                    if nvfp4_fused.is_some() && std::env::var("MEMRA_DEBUG").is_ok() {
7846                        static ONCE: std::sync::Once = std::sync::Once::new();
7847                        ONCE.call_once(|| {
7848                            eprintln!("[memra] NVFP4 beta+alpha batched aux dual ENGAGED (t={t})")
7849                        });
7850                    }
7851                }
7852                if nvfp4_fused.is_none() && spec_fused_t() && (2..=4).contains(&t) {
7853                    q8_fused = e.matmul_q8_fused2_t(&la.ssm_beta, &la.ssm_alpha, hq, hd, t)?;
7854                }
7855            }
7856            if let Some(((mut b, bs), (mut a, as_))) = nvfp4_fused {
7857                if bs != 1.0 {
7858                    e.scale_inplace(&mut b, bs, t * la.ssm_beta.out_features())?;
7859                }
7860                if as_ != 1.0 {
7861                    e.scale_inplace(&mut a, as_, t * la.ssm_alpha.out_features())?;
7862                }
7863                (b, a)
7864            } else if let Some(pair) = q8_fused {
7865                pair
7866            } else {
7867                match hq8_any {
7868                    Some((hq, hd)) if h_q8.is_some() => (
7869                        e.matmul_decode_exact_pre(&la.ssm_beta, hq, hd, t)?,
7870                        e.matmul_decode_exact_pre(&la.ssm_alpha, hq, hd, t)?,
7871                    ),
7872                    _ => (
7873                        e.matmul_decode_exact(&la.ssm_beta, h, t)?,
7874                        e.matmul_decode_exact(&la.ssm_alpha, h, t)?,
7875                    ),
7876                }
7877            }
7878        };
7879
7880        // conv with CARRIED state + ring roll (T >= pad rides the input-column update kernel;
7881        // T < pad — the MEMRA_SPEC_M2 t=2 arm — rolls via the pure-copy ring rebuild).
7882        let rl = cache.recur[il].as_mut().unwrap();
7883        let mut conv_out = e.uninit(conv_dim * t)?;
7884        e.ssm_conv1d_tm_state(
7885            &qkv_mixed,
7886            &mut rl.conv_state,
7887            la.ssm_conv1d.float_data(),
7888            &mut conv_out,
7889            conv_dim,
7890            t,
7891            d_conv,
7892        )?;
7893
7894        // GDN prep via the prefill kernels (repack + L2 + sigmoid + glog), T-wide.
7895        let mut q_g = e.uninit(d_state * num_v * t)?;
7896        let mut k_g = e.uninit(d_state * num_v * t)?;
7897        let mut v_g = e.uninit(d_state * num_v * t)?;
7898        e.qkv_to_gdn_repack(
7899            &conv_out, &mut q_g, &mut k_g, &mut v_g, d_state, num_v, num_k, key_dim, t,
7900        )?;
7901        let mut q_l2 = e.uninit(d_state * num_v * t)?;
7902        e.l2_norm_decode(&q_g, &mut q_l2, d_state, num_v * t, eps)?;
7903        let mut k_l2 = e.uninit(d_state * num_v * t)?;
7904        e.l2_norm_decode(&k_g, &mut k_l2, d_state, num_v * t, eps)?;
7905        let mut beta = e.uninit(t * num_v)?;
7906        e.sigmoid(&beta_raw, &mut beta, t * num_v)?;
7907        let mut g_log = e.uninit(t * num_v)?;
7908        e.gdn_glog(
7909            &alpha,
7910            la.ssm_dt.float_data(),
7911            la.ssm_a.float_data(),
7912            &mut g_log,
7913            num_v,
7914            t,
7915        )?;
7916
7917        // ONE gdn_scan over T tokens from the carried state (internal sequential loop ==
7918        // T chained T=1 steps). Ping-pong the resident buffers like eager decode.
7919        let mut o = e.uninit(d_state * num_v * t)?;
7920        {
7921            let crate::cache::RecurLayer {
7922                ssm_state,
7923                ssm_state_alt,
7924                ..
7925            } = rl;
7926            e.gdn_scan_s128(
7927                &q_l2,
7928                &k_l2,
7929                &v_g,
7930                &g_log,
7931                &beta,
7932                ssm_state,
7933                ssm_state_alt,
7934                &mut o,
7935                num_v,
7936                t,
7937                scale,
7938            )?;
7939        }
7940        std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
7941
7942        // gated RMSNorm + out projection, T-wide. FUSED-QUANTIZE ARM (lane/vt-fixes fix 2,
7943        // mirroring the T=1 decode's launch-arc form): when ssm_out rides the q8_1 fast path,
7944        // emit q8_1 straight from the gated norm at nrows=num_v*t (row-indexed kernel, the
7945        // T-wide launch is the per-row program; kernel-check pins bit-identity vs
7946        // gated_rmsnorm -> quantize_q8_1 at T=1 and T=5) and feed the decode-exact dispatch
7947        // pre-quantized — one launch replaces norm + quantize. Fallback = the f32 chain.
7948        let out = if e.uses_q8_1_fast(&la.ssm_out) {
7949            let (gq, gd) =
7950                e.gated_rmsnorm_q8_1(&o, la.ssm_norm.float_data(), &z, d_state, num_v * t, eps)?;
7951            e.matmul_decode_exact_pre(&la.ssm_out, &gq, &gd, t)?
7952        } else {
7953            let mut gn = e.uninit(d_state * num_v * t)?;
7954            e.gated_rmsnorm(
7955                &o,
7956                la.ssm_norm.float_data(),
7957                &z,
7958                &mut gn,
7959                d_state,
7960                num_v * t,
7961                eps,
7962            )?;
7963            // DECODE-EXACT out-projection: same MMVQ path as the T=1 decode (ssm_out at m>=5
7964            // would fall to dp4a with a different FP reduction order — same class of bug as
7965            // the input projs).
7966            e.matmul_decode_exact(&la.ssm_out, &gn, t)?
7967        };
7968        let stash = if want_stash {
7969            Some(GdnStash {
7970                qkv_mixed,
7971                q_l2,
7972                k_l2,
7973                v_g,
7974                g_log,
7975                beta,
7976            })
7977        } else {
7978            None
7979        };
7980        Ok((out, stash))
7981    }
7982
7983    /// REPLAY-FREE partial-accept commit (2026-07-03): make the cache state == "committed through
7984    /// the first `j` verify columns" WITHOUT the legacy rollback + duplicate trunk replay.
7985    /// - Full-attn KV: truncate both the owning-stage shadow and every TP rank to snapshot + j.
7986    ///   The verify's appended rows for those columns are bit-identical to what an eager T=1
7987    ///   chain writes (the decode-exact contract the verify-probe gates), so keeping them ==
7988    ///   replaying them.
7989    /// - Linear layers, batched path: rebuild the conv ring by PURE COPIES (ring holds raw input
7990    ///   columns) and the ssm state by a prefix re-run of the SAME gdn_scan kernel (t=j) from the
7991    ///   snapshot state over the stash's identical inputs — the kernel's t-loop carries state in
7992    ///   registers and writes it once at the end, so iterations 0..j-1 are independent of T:
7993    ///   bit-identical to the verify's own state after j tokens == the eager chain state.
7994    /// - Linear layers, per-column path: restore the cloned actual state after column j-1.
7995    ///   Caller guarantees 1 <= j <= t-1 (j==0 rounds take the legacy rollback; j==t is full accept).
7996    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
7997    fn commit_verified_prefix(
7998        &self,
7999        e: &Engine,
8000        cache: &mut Cache,
8001        snap: &crate::cache::CacheSnapshot,
8002        ckpt: &VerifyCkpt,
8003        j: usize,
8004    ) -> Result<(), Box<dyn std::error::Error>> {
8005        // GDN geometry derives lazily inside recurrent-layer arms. Full-attention plans carry no
8006        // recurrent state and must never be forced through a synthetic SSM geometry.
8007        // Engine-bundle slice 1 (DSF-ROUNDCOST-20260820 §1.1): the per-column-arm restores
8008        // are 2 tiny D2D copies per linear layer (~96 dispatches/partial round on the q38
8009        // route). When every cols-arm layer shares uniform state sizes (single ssm cfg —
8010        // always true today), batch them into two `copy_batch_uniform_f32` launches. Bytes,
8011        // buffers and stream order are identical to the per-layer memcpy sequence; the
8012        // kernel-rebuild (gdn-stash) arm below is untouched.
8013        let mut batched_cols = false;
8014        {
8015            use cudarc::driver::DevicePtr;
8016            let s = &e.gpu.stream();
8017            let mut conv_pairs: Vec<(u64, u64)> = Vec::new();
8018            let mut ssm_pairs: Vec<(u64, u64)> = Vec::new();
8019            let (mut conv_words, mut ssm_words) = (0usize, 0usize);
8020            let mut uniform = true;
8021            for il in 0..self.layers.len() {
8022                let Some(rl) = cache.recur[il].as_ref() else {
8023                    continue;
8024                };
8025                if ckpt.gdn[il].is_some() {
8026                    continue; // kernel-rebuild arm restores below, per layer
8027                }
8028                let Some(cols) = &ckpt.cols[il] else {
8029                    continue; // missing-ckpt error surfaces in the main loop
8030                };
8031                let (c, st) = &cols[j - 1];
8032                if conv_pairs.is_empty() {
8033                    conv_words = c.len();
8034                    ssm_words = st.len();
8035                } else if c.len() != conv_words || st.len() != ssm_words {
8036                    uniform = false;
8037                    break;
8038                }
8039                let (pc, _g0) = c.device_ptr(s);
8040                let (dc, _g1) = rl.conv_state.device_ptr(s);
8041                let (ps, _g2) = st.device_ptr(s);
8042                let (ds, _g3) = rl.ssm_state.device_ptr(s);
8043                conv_pairs.push((pc, dc));
8044                ssm_pairs.push((ps, ds));
8045            }
8046            if uniform && !conv_pairs.is_empty() {
8047                let n = conv_pairs.len();
8048                let mut t = vec![0u64; 2 * n];
8049                for (k, &(src, dst)) in conv_pairs.iter().enumerate() {
8050                    t[k] = src;
8051                    t[n + k] = dst;
8052                }
8053                let conv_t = e.htod_u64(&t)?;
8054                for (k, &(src, dst)) in ssm_pairs.iter().enumerate() {
8055                    t[k] = src;
8056                    t[n + k] = dst;
8057                }
8058                let ssm_t = e.htod_u64(&t)?;
8059                e.copy_batch_uniform_f32(&conv_t, n, conv_words)?;
8060                e.copy_batch_uniform_f32(&ssm_t, n, ssm_words)?;
8061                batched_cols = true;
8062            }
8063        }
8064        for il in 0..self.layers.len() {
8065            if let (Some(kvl), Some(saved)) = (cache.kv[il].as_mut(), snap.kv_len[il]) {
8066                kvl.len = saved + j;
8067                e.set_i32_one(&mut kvl.len_d, kvl.len as i32)?;
8068            }
8069            if let Some(rl) = cache.recur[il].as_mut() {
8070                let Mixer::Linear(linear) = &self.layers[il].mixer else {
8071                    return Err(format!("recurrent cache layer {il} has no GDN plan").into());
8072                };
8073                let geometry = linear.geometry;
8074                let d_state = geometry.key_head_dim as usize;
8075                let num_k = geometry.key_heads as usize;
8076                let num_v = geometry.value_heads as usize;
8077                let d_conv = geometry.conv_kernel as usize;
8078                let conv_dim = d_state * num_k * 2 + geometry.value_head_dim as usize * num_v;
8079                let scale = 1.0 / (d_state as f32).sqrt();
8080                if let Some(st) = &ckpt.gdn[il] {
8081                    let ring_old = snap.conv[il].as_ref().expect("snapshot missing conv");
8082                    let state_in = snap.ssm[il].as_ref().expect("snapshot missing ssm");
8083                    e.ssm_conv_ring_rebuild(
8084                        &st.qkv_mixed,
8085                        ring_old,
8086                        &mut rl.conv_state,
8087                        conv_dim,
8088                        j,
8089                        d_conv,
8090                    )?;
8091                    let mut o = e.uninit(d_state * num_v * j)?; // scan output, discarded
8092                    e.gdn_scan_s128(
8093                        &st.q_l2,
8094                        &st.k_l2,
8095                        &st.v_g,
8096                        &st.g_log,
8097                        &st.beta,
8098                        state_in,
8099                        &mut rl.ssm_state,
8100                        &mut o,
8101                        num_v,
8102                        j,
8103                        scale,
8104                    )?;
8105                } else if let Some(cols) = &ckpt.cols[il] {
8106                    if !batched_cols {
8107                        let (c, s) = &cols[j - 1];
8108                        e.copy_into(&mut rl.conv_state, 0, c, c.len())?;
8109                        e.copy_into(&mut rl.ssm_state, 0, s, s.len())?;
8110                    }
8111                } else {
8112                    return Err(
8113                        "commit_verified_prefix: verify ckpt missing for linear layer".into(),
8114                    );
8115                }
8116            }
8117        }
8118        self.restore_step_tp_kv_verified_prefix(e, cache, snap, j, true)?;
8119        cache.pos = snap.pos + j;
8120        Ok(())
8121    }
8122
8123    /// ROUND-STREAM: recur restore with device-j (the _dc twins; full accept early-exits
8124    /// in-kernel). Requires the batched-linear stash on every linear layer (stream gate).
8125    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
8126    fn commit_verified_prefix_stream(
8127        &self,
8128        e: &Engine,
8129        cache: &mut Cache,
8130        snap: &crate::cache::CacheSnapshot,
8131        ckpt: &VerifyCkpt,
8132        acc: &CudaSlice<u32>,
8133        base: usize,
8134        t_v: usize,
8135    ) -> Result<(), Box<dyn std::error::Error>> {
8136        for il in 0..self.layers.len() {
8137            if let Some(rl) = cache.recur[il].as_mut() {
8138                let Mixer::Linear(linear) = &self.layers[il].mixer else {
8139                    return Err(format!("recurrent cache layer {il} has no GDN plan").into());
8140                };
8141                let geometry = linear.geometry;
8142                let d_state = geometry.key_head_dim as usize;
8143                let num_k = geometry.key_heads as usize;
8144                let num_v = geometry.value_heads as usize;
8145                let d_conv = geometry.conv_kernel as usize;
8146                let conv_dim = d_state * num_k * 2 + geometry.value_head_dim as usize * num_v;
8147                let scale = 1.0 / (d_state as f32).sqrt();
8148                let st = ckpt.gdn[il]
8149                    .as_ref()
8150                    .ok_or("stream restore: batched-linear stash missing")?;
8151                let ring_old = snap.conv[il].as_ref().expect("snapshot missing conv");
8152                let state_in = snap.ssm[il].as_ref().expect("snapshot missing ssm");
8153                e.ssm_conv_ring_rebuild_dc(
8154                    &st.qkv_mixed,
8155                    ring_old,
8156                    &mut rl.conv_state,
8157                    conv_dim,
8158                    acc,
8159                    base,
8160                    t_v,
8161                    d_conv,
8162                )?;
8163                let mut o = e.uninit(d_state * num_v * t_v)?;
8164                e.gdn_scan_s128_dc(
8165                    &st.q_l2,
8166                    &st.k_l2,
8167                    &st.v_g,
8168                    &st.g_log,
8169                    &st.beta,
8170                    state_in,
8171                    &mut rl.ssm_state,
8172                    &mut o,
8173                    num_v,
8174                    acc,
8175                    base,
8176                    t_v,
8177                    scale,
8178                )?;
8179            }
8180        }
8181        Ok(())
8182    }
8183
8184    /// EAGLE3 aux-capturing verify forward over `tokens` (T) — mirrors `decode_step_t_h` exactly
8185    /// (same KV append, same causal verify, same recur advance) but ALSO clones the aux residual-
8186    /// stream hiddens (blocks in `aux_layers`) for TWO columns: the LAST column (always) and the
8187    /// optional `pred_col` (the EAGLE seed = bonus's predecessor). Returns
8188    /// (all_T_logits host, last_col_aux, pred_col_aux?). Used by the EAGLE3 orchestrator's commit.
8189    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
8190    pub fn decode_step_t_aux2(
8191        &self,
8192        e: &Engine,
8193        tokens: &[u32],
8194        pos0: usize,
8195        cache: &mut Cache,
8196        aux_layers: &[usize],
8197        pred_col: Option<usize>,
8198    ) -> Result<
8199        (Vec<f32>, Vec<CudaSlice<f32>>, Option<Vec<CudaSlice<f32>>>),
8200        Box<dyn std::error::Error>,
8201    > {
8202        cache.ensure_usable("decode_step_t_aux2")?;
8203        let cfg = &self.cfg;
8204        let n_embd = cfg.n_embd as usize;
8205        let eps = cfg.rms_eps;
8206        let t = tokens.len();
8207        let pos_vec: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
8208        let pos_d = e.htod_i32(&pos_vec)?;
8209        let mut x = e.htod(&self.embd.try_gather(n_embd, tokens)?)?;
8210        let mut aux_last: Vec<CudaSlice<f32>> = Vec::with_capacity(aux_layers.len());
8211        let mut aux_pred: Vec<CudaSlice<f32>> = Vec::new();
8212        let want_pred = pred_col.is_some();
8213
8214        for (il, layer) in self.layers.iter().enumerate() {
8215            // DISPATCH-MIRRORED norms (FP-order lesson #8) — see decode_step_t_h_emb.
8216            let mixer_fast = self.mixer_in_q8_1_fast(e, &layer.mixer);
8217            let norm_fused = std::env::var("MEMRA_NO_FUSE_NORMQ").is_err() && mixer_fast;
8218            let mut h = vbuf(e, t * n_embd)?; // fully written by either rms_norm arm
8219            if norm_fused {
8220                e.rms_norm_decode(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
8221            } else {
8222                e.rms_norm(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
8223            }
8224            let mixed = match &layer.mixer {
8225                Mixer::Full(fa) => {
8226                    self.full_attn_verify(e, fa, &h, None, &pos_d, t, cache, il, None)?
8227                }
8228                Mixer::Mla(_) => {
8229                    crate::hybrid::mla_path_unimplemented("auxiliary T-parallel decode")
8230                }
8231                Mixer::Kda(_) => crate::hybrid::kda_path_unimplemented("aux decode step"),
8232                Mixer::Linear(la) => {
8233                    let mut out = e.zeros(t * n_embd)?;
8234                    for col in 0..t {
8235                        let mut h_col = e.zeros(n_embd)?;
8236                        let src = h.slice(col * n_embd..(col + 1) * n_embd);
8237                        e.copy_view_into(&mut h_col, 0, &src, n_embd)?;
8238                        let m_col = self.linear_attn_decode(e, la, &h_col, cache, il)?;
8239                        e.copy_into(&mut out, col * n_embd, &m_col, n_embd)?;
8240                    }
8241                    out
8242                }
8243            };
8244            let ffn_fuse = match &layer.ffn {
8245                crate::hybrid::Ffn::Dense {
8246                    ffn_gate, ffn_up, ..
8247                } => {
8248                    std::env::var("MEMRA_NO_FUSE_NORMQ").is_err()
8249                        && e.uses_q8_1_fast(ffn_gate)
8250                        && e.uses_q8_1_fast(ffn_up)
8251                }
8252                crate::hybrid::Ffn::Moe(_) => false,
8253            };
8254            let mut x1 = vbuf(e, t * n_embd)?; // fully written by add / add_rms_norm
8255            let mut z = vbuf(e, t * n_embd)?; // fully written by rms_norm_decode / add_rms_norm
8256            if ffn_fuse {
8257                e.add(&x, &mixed, &mut x1, t * n_embd)?;
8258                e.rms_norm_decode(
8259                    &x1,
8260                    layer.post_attn_norm.float_data(),
8261                    &mut z,
8262                    n_embd,
8263                    t,
8264                    eps,
8265                )?;
8266            } else {
8267                e.add_rms_norm(
8268                    &x,
8269                    &mixed,
8270                    layer.post_attn_norm.float_data(),
8271                    &mut x1,
8272                    &mut z,
8273                    n_embd,
8274                    t,
8275                    eps,
8276                )?;
8277            }
8278            let ffn_out = match &layer.ffn {
8279                crate::hybrid::Ffn::Dense {
8280                    ffn_gate,
8281                    ffn_up,
8282                    ffn_down,
8283                } => {
8284                    let n_ff = ffn_gate.out_features();
8285                    let gate = e.matmul_decode_exact(ffn_gate, &z, t)?;
8286                    let up = e.matmul_decode_exact(ffn_up, &z, t)?;
8287                    let mut act = vbuf(e, t * n_ff)?; // fully written by ffn_act_lim
8288                    // dense FFN clamp = the SHEXP array (upstream build_ffn serves both).
8289                    Self::ffn_act_lim(
8290                        e,
8291                        &self.cfg,
8292                        &gate,
8293                        &up,
8294                        1.0,
8295                        1.0,
8296                        self.cfg.clamp_shexp_at(il as u32),
8297                        &mut act,
8298                        t * n_ff,
8299                    )?;
8300                    e.matmul_decode_exact(ffn_down, &act, t)?
8301                }
8302                crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il(e, m, &z, t, il as u16)?,
8303            };
8304            let mut x2 = vbuf(e, t * n_embd)?; // fully written by add
8305            e.add(&x1, &ffn_out, &mut x2, t * n_embd)?;
8306            if aux_layers.contains(&il) {
8307                let mut a = e.zeros(n_embd)?;
8308                e.copy_view_into(&mut a, 0, &x2.slice((t - 1) * n_embd..t * n_embd), n_embd)?;
8309                aux_last.push(a);
8310                if let Some(pc) = pred_col {
8311                    let mut ap = e.zeros(n_embd)?;
8312                    e.copy_view_into(
8313                        &mut ap,
8314                        0,
8315                        &x2.slice(pc * n_embd..(pc + 1) * n_embd),
8316                        n_embd,
8317                    )?;
8318                    aux_pred.push(ap);
8319                }
8320            }
8321            x = x2;
8322        }
8323        let mut hn = vbuf(e, t * n_embd)?; // fully written by rms_norm_decode
8324        e.rms_norm_decode(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
8325        let logits = e.matmul_decode_exact(&self.output, &hn, t)?;
8326        let host = e.dtoh(&logits)?;
8327        cache.pos += t;
8328        Ok((
8329            host,
8330            aux_last,
8331            if want_pred { Some(aux_pred) } else { None },
8332        ))
8333    }
8334
8335    /// step35 SPEC-VERIFY attention over T query tokens — a per-row REPLAY of the eager
8336    /// `step35_decode_attn`.
8337    ///
8338    /// WHY A REPLAY AND NOT A BATCHED TWIN. The verify's whole job is to be bit-identical to what
8339    /// the eager decode would have computed for the same tokens; that is what makes greedy spec
8340    /// decode exact (run-spec asserts token identity for K=1..8). Every other verify arm in this
8341    /// file earns that identity by carefully mirroring dispatch (`matmul_decode_exact` to force
8342    /// MMVQ at any m, per-layer `ffn_fuse` mirroring, per-row `fa_decode` key bounds). step35
8343    /// stacks FOUR more per-layer degrees of freedom on top of that — per-layer `n_head`
8344    /// (64 full / 96 SWA), per-layer rotary width (64 full / 128 SWA), per-layer rope base, and a
8345    /// SEPARATE `attn_gate` tensor whose projection shares the attn-normed input — and its SWA
8346    /// layers attend through a token-OFFSET view whose offset is a function of the ABSOLUTE
8347    /// position of each query row. A batched twin would have to reproduce all of that AND the
8348    /// per-row offset in one launch; the offset alone rules out the existing rows kernels (they
8349    /// take one `base_len`, not a per-row offset).
8350    ///
8351    /// So this arm calls the eager path itself, once per row, on the same cache. Identity is then
8352    /// true BY CONSTRUCTION rather than by mirroring: row r runs exactly the kernel sequence that
8353    /// eager decode step r runs (same projections, same q8_1 fusion decision, same append, same
8354    /// view arithmetic, same `fa_decode_kvmod`, same gate), because it IS that code. Cost: T x the
8355    /// eager decode mixer instead of one batched pass — the same trade the generic arm's `else`
8356    /// per-row loop already accepts when `fa_rows_eligible` says no. Correctness first; a batched
8357    /// step35 twin is a perf lane's job and must be gated against this arm.
8358    ///
8359    /// The `h_q8` pre-quantized pair from the caller's fused norm is NOT forwarded: it is a
8360    /// T-row buffer and `step35_decode_attn`'s `pre_q` contract is one row. Instead each row's
8361    /// f32 `h` slice is handed over and the callee re-derives its own q8_1 exactly as eager decode
8362    /// does (`quantize_q8_1(h, 1, n_embd)`) — which is the dispatch being mirrored. Callers that
8363    /// took the fused arm therefore MUST still pass a live `h`; `step35_verify` asserts that.
8364    #[allow(clippy::too_many_arguments)]
8365    fn step35_verify(
8366        &self,
8367        e: &Engine,
8368        fa: &FullAttnLayer,
8369        h: &CudaSlice<f32>,
8370        h_q8: Option<(&CudaSlice<i8>, &CudaSlice<f32>)>,
8371        t: usize,
8372        cache: &mut Cache,
8373        il: usize,
8374    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8375        let n_embd = self.cfg.n_embd as usize;
8376        // The fused (h-less) attn-norm arm hands `h` as a zero-length placeholder. This arm needs
8377        // the f32 rows, so the caller must not take that lever for step35 — enforced at the call
8378        // site by the sliding-gated-MoE `Mixer::Full(_) => false` arm of
8379        // `lin_q8_only` in `decode_step_t_core_stream`, and asserted here so a future caller
8380        // cannot regress it into silently reading an empty buffer.
8381        assert_eq!(
8382            h.len(),
8383            t * n_embd,
8384            "step35_verify needs the f32 attn-normed rows ([t*n_embd]); the caller took the \
8385             fused q8-only norm arm (h_q8={}) — step35 must stay on the unfused arm",
8386            h_q8.is_some()
8387        );
8388        // ROW WIDTH IS n_embd, NOT n_head*head_dim: `step35_decode_attn` returns the mixer output
8389        // AFTER `wo`, so a row is [n_embd] — the same contract the generic arm's
8390        // `matmul_decode_exact(&fa.wo, &attn_g, t)` return has. Sizing this buffer from the
8391        // per-layer head geometry (8192 on full-attn, 12288 on SWA) instead overran the row on the
8392        // FIRST copy and panicked inside `copy_into`'s `CudaView::slice` unwrap
8393        // (raw/mtp-bt-20260806T212127Z.log frames 12-13).
8394        let mut out = vbuf(e, t * n_embd)?; // each row fully written by the copy below
8395        for r in 0..t {
8396            // Absolute position of this query row. `cache.pos` is the committed length at round
8397            // start and every row before r has already been appended by this loop, so the r-th
8398            // verify token sits at cache.pos + r — the same position eager decode would give it.
8399            let pos_d = e.htod_i32(&[(cache.pos + r) as i32])?;
8400            let mut h_row = vbuf(e, n_embd)?; // fully written by copy_view_into
8401            e.copy_view_into(
8402                &mut h_row,
8403                0,
8404                &h.slice(r * n_embd..(r + 1) * n_embd),
8405                n_embd,
8406            )?;
8407            // THE eager decode mixer: appends this row's K/V at kvl.len, advances it, then
8408            // attends over the (SWA-offset) view. Post-`wo`, same contract as this fn returns.
8409            let o = self.step35_decode_attn(e, fa, il, &h_row, None, &pos_d, cache)?;
8410            debug_assert_eq!(
8411                o.len(),
8412                n_embd,
8413                "step35_decode_attn returns post-wo [n_embd]"
8414            );
8415            e.copy_into(&mut out, r * n_embd, &o, n_embd)?;
8416        }
8417        Ok(out)
8418    }
8419
8420    /// Full-attention mixer over T query tokens with a GROWING resident KV (verify path, §D.3).
8421    /// Appends the T new K/V columns to cache.kv[il] then attends causally over [0..len) via
8422    /// fa_prefill. Token-major [T, kv_dim] projection layout == cache row layout (single copy).
8423    #[allow(clippy::too_many_arguments)]
8424    fn full_attn_verify(
8425        &self,
8426        e: &Engine,
8427        fa: &FullAttnLayer,
8428        h: &CudaSlice<f32>,
8429        h_q8: Option<(&CudaSlice<i8>, &CudaSlice<f32>)>,
8430        pos_d: &CudaSlice<i32>,
8431        t: usize,
8432        cache: &mut Cache,
8433        il: usize,
8434        stream_ctr: Option<&CudaSlice<i32>>,
8435    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8436        // step35: the generic geometry below is wrong for this arch (per-layer n_head, partial
8437        // per-layer rope, the SWA offset view, and a SEPARATE head-wise gate tensor), so it takes
8438        // its own arm. A verify that silently computes different attention than decode defeats the
8439        // whole self-consistency gate, so the arm is a per-row REPLAY of `step35_decode_attn`
8440        // rather than a batched twin — see `step35_verify` for why that is the exactness-correct
8441        // shape and not laziness.
8442        if self.sliding_gated_moe_batch_program() {
8443            if stream_ctr.is_some() {
8444                return Err(
8445                    "step35 has no ROUND-STREAM verify arm (the device-counter _dc twins \
8446                            cannot express the SWA offset KV view; same root cause as the dc \
8447                            decode refusal) — run spec without the stream arm"
8448                        .into(),
8449                );
8450            }
8451            return self.step35_verify(e, fa, h, h_q8, t, cache, il);
8452        }
8453        let cfg = &self.cfg;
8454        let geometry = cfg.full_attention_geometry_at(il as u32);
8455        let n_head = geometry.n_head as usize;
8456        let n_head_kv = geometry.n_head_kv as usize;
8457        let head_dim = geometry.head_dim_k as usize;
8458        let eps = cfg.rms_eps;
8459        let scale = geometry.attention_scale();
8460        let n_embd = cfg.n_embd as usize;
8461
8462        // DECODE-EXACT Q/K/V projections: matmul_decode_exact forces the MMVQ (warp-per-row) path
8463        // for every m, matching the T=1 decode's FP accumulation order. matmul_pre at m>=5 would
8464        // fall to dp4a (128-thread, two-level reduce) with a different FP sum order.
8465        // Q8 TRUNK-FUSION at T=1: DISPATCH-MIRRORS the eager decode's fused3 (bit-identical body).
8466        // BATCHED EPILOGUE RE-FUSE (lane/vt-fixes fix 2): `h_q8` = the attn-input norm's q8_1
8467        // form emitted by the fused rms_norm_q8_1 (bit-identical to rms_norm_decode ->
8468        // quantize_q8_1, kernel-check-pinned). When present (caller checked mixer q8_1-fast),
8469        // every projection consumes it — `h` may be a zero-len placeholder and must not be read.
8470        let (qf, mut k, v) = if let Some(mut qkv) = self.full_attn_tp_qkv(e, fa, h, t)? {
8471            let v = qkv.pop().ok_or("full-attention TP verify QKV omitted V")?;
8472            let k = qkv.pop().ok_or("full-attention TP verify QKV omitted K")?;
8473            let q = qkv.pop().ok_or("full-attention TP verify QKV omitted Q")?;
8474            if !qkv.is_empty() {
8475                return Err("full-attention TP verify QKV returned extra projections".into());
8476            }
8477            (q, k, v)
8478        } else {
8479            let mut fused = None;
8480            let qkv_fast =
8481                e.uses_q8_1_fast(&fa.wq) && e.uses_q8_1_fast(&fa.wk) && e.uses_q8_1_fast(&fa.wv);
8482            if t == 1 && qkv_fast {
8483                let (hq_o, hd_o);
8484                let (hq, hd): (&CudaSlice<i8>, &CudaSlice<f32>) = match h_q8 {
8485                    Some(p) => p,
8486                    None => {
8487                        (hq_o, hd_o) = e.quantize_q8_1(h, 1, n_embd)?;
8488                        (&hq_o, &hd_o)
8489                    }
8490                };
8491                fused = e.matmul_q8_fused3(&fa.wq, &fa.wk, &fa.wv, hq, hd)?;
8492            } else if spec_fused_t() && (2..=4).contains(&t) && qkv_fast {
8493                // VERIFY-TIER TRUNK FUSION (MEMRA_SPEC_FUSED_T): one shared quantize + one
8494                // fused3 batched launch replaces three decode-exact calls (3 re-quantizes of
8495                // the same h + 3 _b2/_b4 launches). Bit-identical per (tensor,token,row).
8496                let (hq_o, hd_o);
8497                let (hq, hd): (&CudaSlice<i8>, &CudaSlice<f32>) = match h_q8 {
8498                    Some(p) => p,
8499                    None => {
8500                        (hq_o, hd_o) = e.quantize_q8_1(h, t, n_embd)?;
8501                        (&hq_o, &hd_o)
8502                    }
8503                };
8504                fused = e.matmul_q8_fused3_t(&fa.wq, &fa.wk, &fa.wv, hq, hd, t)?;
8505            }
8506            match (fused, h_q8) {
8507                (Some(triple), _) => triple,
8508                // shared pre-quantized activation (q8_1-fast guaranteed by the caller): the
8509                // decode-exact dispatch consumes (hq, hd) instead of re-quantizing 3x.
8510                (None, Some((hq, hd))) if qkv_fast => (
8511                    e.matmul_decode_exact_pre(&fa.wq, hq, hd, t)?,
8512                    e.matmul_decode_exact_pre(&fa.wk, hq, hd, t)?,
8513                    e.matmul_decode_exact_pre(&fa.wv, hq, hd, t)?,
8514                ),
8515                (None, _) => (
8516                    e.matmul_decode_exact(&fa.wq, h, t)?,
8517                    e.matmul_decode_exact(&fa.wk, h, t)?,
8518                    e.matmul_decode_exact(&fa.wv, h, t)?,
8519                ),
8520            }
8521        };
8522        // M3/Hy3 have no attention output gate — wq out is exactly q; skip the split.
8523        let gated = geometry.attention_gate == memra_gguf::config::AttentionGateKind::FusedQ;
8524        let (mut q, gate) = if gated {
8525            let mut q = vbuf(e, t * n_head * head_dim)?; // fully written by q_gate_split
8526            let mut gate = vbuf(e, t * n_head * head_dim)?; // fully written by q_gate_split
8527            e.q_gate_split(&qf, &mut q, &mut gate, head_dim, n_head, t)?;
8528            (q, Some(gate))
8529        } else {
8530            (qf, None)
8531        };
8532
8533        let mut qn = vbuf(e, t * n_head * head_dim)?; // fully written by rms_norm
8534        e.rms_norm(
8535            &q,
8536            fa.q_norm.float_data(),
8537            &mut qn,
8538            head_dim,
8539            n_head * t,
8540            eps,
8541        )?;
8542        q = qn;
8543        let mut kn = vbuf(e, t * n_head_kv * head_dim)?; // fully written by rms_norm
8544        e.rms_norm(
8545            &k,
8546            fa.k_norm.float_data(),
8547            &mut kn,
8548            head_dim,
8549            n_head_kv * t,
8550            eps,
8551        )?;
8552        k = kn;
8553        let rope_dims = geometry.n_rot as usize;
8554        e.rope_neox(
8555            &mut q,
8556            pos_d,
8557            head_dim,
8558            rope_dims,
8559            n_head,
8560            t,
8561            geometry.rope_base,
8562            1.0,
8563        )?;
8564        e.rope_neox(
8565            &mut k,
8566            pos_d,
8567            head_dim,
8568            rope_dims,
8569            n_head_kv,
8570            t,
8571            geometry.rope_base,
8572            1.0,
8573        )?;
8574
8575        // append T new K/V columns to the resident QUANTIZED cache. k/v are token-major [T, kv_dim]
8576        // f32; append-quantize each of the T token rows into the byte cache (q8_0 K / q5_1 V).
8577        let kvl = cache.kv[il].as_mut().unwrap();
8578        let (kv_dim_k, kv_dim_v, ktb, vtb) =
8579            (kvl.kv_dim_k, kvl.kv_dim_v, kvl.k_tok_bytes, kvl.v_tok_bytes);
8580        if let Some(ctr) = stream_ctr {
8581            // stream: ONE batched append at the device counter (rows kernel = the per-view warp
8582            // math on a (block, token) grid, documented byte-identical); host len is a stale
8583            // LOWER BOUND under pre-issue (drain reconciles it).
8584            e.append_kv_quantized_rows_dc(
8585                &k, &v, &mut kvl.k, &mut kvl.v, ctr, t, kv_dim_k, kv_dim_v, ktb, vtb, false,
8586            )?;
8587        } else {
8588            for i in 0..t {
8589                let k_row = k.slice(i * kv_dim_k..(i + 1) * kv_dim_k);
8590                let v_row = v.slice(i * kv_dim_v..(i + 1) * kv_dim_v);
8591                e.append_kv_quantized_view(
8592                    &k_row,
8593                    &v_row,
8594                    &mut kvl.k,
8595                    &mut kvl.v,
8596                    kvl.len + i,
8597                    kv_dim_k,
8598                    kv_dim_v,
8599                    ktb,
8600                    vtb,
8601                    false,
8602                )?;
8603            }
8604            kvl.len += t;
8605        }
8606
8607        // BIT-IDENTICAL VERIFY ATTENTION (spec-exactness fix): the FP accumulation order must be
8608        // byte-for-byte identical to the eager decode path. fa_prefill uses a different tile size
8609        // (BLOCK_Q=64, BK=32) and online-softmax structure than fa_decode's split-K + combine,
8610        // which changes FP summation order and can flip argmax at tight logit margins. Query row r
8611        // attends to keys [0..base_len+r+1) — each successive row sees one more key (the causal
8612        // property). This matches eager: decode appends k at len, then fa_decode sees t_kv = len+1
8613        // keys. The verify appends all T tokens first but bounds the key range per row.
8614        //
8615        // MULTI-ROW FUSED PATH (the long-ctx spec fix, 2026-07-03): when every row takes the vec
8616        // kernel (base_len+1 >= FA_VEC_MIN_TKV), ONE fa_decode_rows launch executes the exact
8617        // per-row program for all T rows (grid.z = row, per-row n_splits from the same
8618        // fa_split_keys formula) — replacing T x (2 launches + 2 dtod copies + 5 partial allocs)
8619        // and multiplying resident CTAs by T on a latency-bound kernel. Bit-identical per row by
8620        // construction; kernel-check pins rows-vs-loop byte identity, run-spec is the end gate.
8621        // Short ctx (any row below the vec crossover) and MEMRA_NO_FA_VEC/MEMRA_FA_ROWS_OFF keep the
8622        // per-row loop (whose fa_decode picks scalar/vec per row exactly like eager decode).
8623        let mut attn = vbuf(e, t * n_head * head_dim)?; // fully written by every FA arm below
8624        let base_len = kvl.len - t; // KV len BEFORE this round's T tokens were appended
8625        // T=1 INCLUDED (2026-07-05): p-min cuts the draft to 1 in ~75% of rounds on hard
8626        // (agentic) content — the old t>1 gate sent those rounds to the per-row loop (262us/row
8627        // + q-row copy + per-row allocs vs 93us/row through the fused kernel at grid.z=1, same
8628        // program). nsys accounting: 1088 of 1456 verify FA launches were T=1 escapees.
8629        // LEAN T=1 ARM (MEMRA_SPEC_LEAN, close35): at t==1, q IS one row and fa_decode on it is
8630        // the EXACT eager decode dispatch (vec_q_v2 + combine_f32; the rows pair measured +50us
8631        // at m=1). Byte-identical: kernel-check pins rows-vs-loop identity, and the per-row loop
8632        // at t=1 is fa_decode on the same q with zero-offset copies. Gates arbitrate.
8633        if let Some(ctr) = stream_ctr {
8634            // STREAM ARM: causal base from the device counter; host kvl.len is a stale lower
8635            // bound used only for the split-sizing upper bound (+64 slack covers M pre-issued
8636            // rounds at K<=8). Views span the bound; per-row limits derive in-kernel.
8637            let upper = kvl.len + t + 64;
8638            let k_view = e.view_u8(&kvl.k, (upper.min(cache.max_ctx)) * ktb);
8639            let v_view = e.view_u8(&kvl.v, (upper.min(cache.max_ctx)) * vtb);
8640            e.fa_decode_rows_dc(
8641                &q,
8642                &k_view,
8643                &v_view,
8644                &mut attn,
8645                head_dim,
8646                n_head,
8647                n_head_kv,
8648                ctr,
8649                upper.min(cache.max_ctx),
8650                t,
8651                scale,
8652                ktb,
8653                vtb,
8654                0,
8655                false,
8656            )?;
8657        } else if spec_lean() && t == 1 {
8658            let t_kv = base_len + 1;
8659            let k_view = e.view_u8(&kvl.k, t_kv * ktb);
8660            let v_view = e.view_u8(&kvl.v, t_kv * vtb);
8661            e.fa_decode_kvmod(
8662                &q, &k_view, &v_view, &mut attn, head_dim, n_head, n_head_kv, t_kv, scale, ktb,
8663                vtb, false,
8664            )?;
8665        } else if e.fa_rows_eligible(base_len, head_dim) {
8666            let k_view = e.view_u8(&kvl.k, (base_len + t) * ktb);
8667            let v_view = e.view_u8(&kvl.v, (base_len + t) * vtb);
8668            e.fa_decode_rows(
8669                &q, &k_view, &v_view, &mut attn, head_dim, n_head, n_head_kv, base_len, t, scale,
8670                ktb, vtb, None, false, false, None,
8671            )?;
8672        } else {
8673            for r in 0..t {
8674                let t_kv_r = base_len + r + 1; // this row sees keys [0..t_kv_r)
8675                let k_view_r = e.view_u8(&kvl.k, t_kv_r * ktb);
8676                let v_view_r = e.view_u8(&kvl.v, t_kv_r * vtb);
8677                // copy q row into an owned buffer (fa_decode takes &CudaSlice, not CudaView)
8678                let mut q_row = vbuf(e, n_head * head_dim)?; // fully written by copy_view_into
8679                let q_src = q.slice(r * n_head * head_dim..(r + 1) * n_head * head_dim);
8680                e.copy_view_into(&mut q_row, 0, &q_src, n_head * head_dim)?;
8681                let mut attn_row = vbuf(e, n_head * head_dim)?; // fully written by fa_decode
8682                e.fa_decode_kvmod(
8683                    &q_row,
8684                    &k_view_r,
8685                    &v_view_r,
8686                    &mut attn_row,
8687                    head_dim,
8688                    n_head,
8689                    n_head_kv,
8690                    t_kv_r,
8691                    scale,
8692                    ktb,
8693                    vtb,
8694                    false,
8695                )?;
8696                e.copy_into(
8697                    &mut attn,
8698                    r * n_head * head_dim,
8699                    &attn_row,
8700                    n_head * head_dim,
8701                )?;
8702            }
8703        }
8704
8705        let attn_g = match &gate {
8706            Some(gate) => {
8707                let mut gsig = vbuf(e, t * n_head * head_dim)?; // fully written by sigmoid
8708                e.sigmoid(gate, &mut gsig, t * n_head * head_dim)?;
8709                let mut ag = vbuf(e, t * n_head * head_dim)?; // fully written by mul
8710                e.mul(&attn, &gsig, &mut ag, t * n_head * head_dim)?;
8711                ag
8712            }
8713            None => attn,
8714        };
8715        // DECODE-EXACT wo projection: at m>=5 (K=4+ with pending) the generic matmul would use dp4a
8716        // (128-thread, different FP sum order than MMVQ). Force MMVQ for bit-identity with decode.
8717        match self.full_attn_tp_o(e, fa, &attn_g, t)? {
8718            Some(output) => Ok(output),
8719            None => Ok(e.matmul_decode_exact(&fa.wo, &attn_g, t)?),
8720        }
8721    }
8722
8723    /// Context-linear bytes for a plain serving session's trunk cache.
8724    pub fn plain_session_kv_bytes_per_token(&self) -> usize {
8725        crate::cache::cache_bytes_per_token_for_plan(
8726            &self.cfg,
8727            &self.plan,
8728            0,
8729            self.plan.layers.len(),
8730        )
8731    }
8732
8733    /// `(logical bytes/token, ring-capped bytes/token, ring row cap)` for exact admission.
8734    pub fn plain_session_kv_shape(&self) -> (usize, usize, usize) {
8735        (
8736            self.plain_session_kv_bytes_per_token(),
8737            crate::cache::cache_ring_bytes_per_token_for_plan(
8738                &self.cfg,
8739                &self.plan,
8740                0,
8741                self.plan.layers.len(),
8742            ),
8743            crate::cache::cache_ring_row_cap_for_plan(&self.plan),
8744        )
8745    }
8746
8747    /// Context-linear bytes for a speculative serving session: trunk cache plus persistent MTP
8748    /// scratch. With no MTP head this equals the plain coefficient.
8749    pub fn spec_session_kv_bytes_per_token(&self) -> usize {
8750        let scratch = self
8751            .mtp
8752            .iter()
8753            .chain(self.mtp_extra.iter())
8754            .map(|mtp| {
8755                let (_, _, k, v) = mtp_scratch_layout(&self.cfg, mtp.geom.as_ref());
8756                k + v
8757            })
8758            .sum::<usize>();
8759        self.plain_session_kv_bytes_per_token()
8760            .saturating_add(scratch)
8761    }
8762
8763    /// Spec twin of [`HybridModel::plain_session_kv_shape`]; Step35's persistent MTP scratch is
8764    /// capped by the same SWA ring rows as the trunk.
8765    pub fn spec_session_kv_shape(&self) -> (usize, usize, usize) {
8766        let total = self.spec_session_kv_bytes_per_token();
8767        let (_, mut ring, rows) = self.plain_session_kv_shape();
8768        if rows > 0 {
8769            ring = ring.saturating_add(
8770                self.mtp
8771                    .iter()
8772                    .chain(self.mtp_extra.iter())
8773                    .map(|mtp| {
8774                        let (_, _, k, v) = mtp_scratch_layout(&self.cfg, mtp.geom.as_ref());
8775                        k + v
8776                    })
8777                    .sum::<usize>(),
8778            );
8779        }
8780        (total, ring, rows)
8781    }
8782
8783    /// Greedy MTP speculative decode (§B). Token-identical to `generate(prompt, max_new)` but uses
8784    /// the NextN head to draft K tokens then verifies them in one batched target forward.
8785    /// Returns (generated tokens, total_drafted, total_accepted) so the caller can report
8786    /// acceptance rate. `k` = draft length per round.
8787    ///
8788    /// GRAPH DRAFT (stage 2 of graph-grade spec): when the model is all-Dense and the MTP head is
8789    /// Dense (no MoE host readbacks), the fixed-shape T=1 MTP forward is CUDA-graph-captured ONCE
8790    /// and replayed per draft step — the ~40 eager launches per drafted token collapse into one
8791    /// graph dispatch; only the 4-byte token id (and 4-byte p-min confidence) round-trip per step.
8792    /// Event tracking is disabled for the whole call (generate_graph pattern) so every buffer the
8793    /// captured graph references is event-free; the spec loop is strictly single-stream.
8794    /// MEMRA_SPEC_NOGRAPH=1 forces the eager draft chain.
8795    /// SAMPLED mode (MEMRA_SPEC_TEMP>0) has its OWN capture (gumbel-perturbed in-graph argmax,
8796    /// device Philox event counter, persistent q retention) — graph-vs-eager sampled streams are
8797    /// bit-identical for the same (seed, prompt, K, temp); see the sampled-graph setup in
8798    /// generate_spec_inner2.
8799    /// Multi-turn session: trunk cache + MTP draft scratch persist across generate calls, so
8800    /// turn N+1 primes ONLY its new suffix (the 124k-conversation daily pattern — re-priming a
8801    /// 32k history costs ~54s; a suffix prime costs seconds). APPEND-ONLY by construction: the
8802    /// hybrid linear-attn states are in-place (no position index), so a session can extend but
8803    /// never rewind — `committed` is the exact token list whose state the caches hold (includes
8804    /// any overshoot tokens past max_new; the caller renders from `committed`, not its own echo).
8805    pub fn new_session(
8806        &self,
8807        e: &Engine,
8808        max_ctx: usize,
8809    ) -> Result<SpecSession, Box<dyn std::error::Error>> {
8810        Ok(SpecSession {
8811            // STAGE-OWNED KV (lane/pp2-spec 2026-08-06): `pp::new_cache`, not `Cache::new`. This
8812            // is the SERVING spec-session path, and with the ppN door open across two cards a
8813            // primary-homed cache makes every remote stage peer-read its OWN KV on every verify
8814            // round — the wrong-card class already fixed on the two batched serving paths
8815            // (worker.rs 2483 / 2837). With the door shut `new_cache` IS `Cache::new` (same
8816            // branch, same allocations), so single-device behavior is byte-unchanged.
8817            cache: crate::pp::new_cache_planned(e, &self.cfg, &self.plan, max_ctx)?,
8818            scratch: self.new_mtp_scratch(e, max_ctx)?,
8819            committed: Vec::new(),
8820            last_h: None,
8821            next_pred: None,
8822            sctr: 0,
8823            uctr: 0,
8824            draft_ctx: None,
8825            pending_tok: None,
8826            turn_ckpt: None,
8827            telem: SpecTelemetryCounters::default(),
8828            capture_at: None,
8829            boundary_captures: Vec::new(),
8830            ckpt_at: None,
8831            capture_disabled: false,
8832        })
8833    }
8834
8835    /// SPEC-ON-CACHE-HIT restore (lane/spec-on-cache-hit, 2026-08-18 — PORT-PLAN item 3,
8836    /// research/cache-spec-design-20260814, scoped to WHOLE-ENTRY restores only): build a
8837    /// SpecSession around a trunk cache the worker already restored from a prefix-cache
8838    /// entry, re-installing the entry's published draft plane as the MTP scratch rows
8839    /// `[0..prefix.len())` and the entry's boundary hidden as `last_h`, then feeding the
8840    /// prompt SUFFIX here — through EXACTLY the plain path's program selection — so the
8841    /// worker always receives a fully-warm continuation session (committed = whole
8842    /// prompt, `next_pred` + `last_h` set; caller sets `next_pred` from the entry's
8843    /// boundary logits on the empty-suffix shape).
8844    ///
8845    /// PROGRAM LAW (the splitiso two-programs class, learned AGAIN in this lane's own
8846    /// gate): the identity target for a converted hit is the PLAIN hit serving the same
8847    /// request, and plain feeds a carried suffix via eager `decode_step` below
8848    /// PRIME_MIN_T and via `prime_cache` at/above it (prefill_tick's arms). The generate
8849    /// path's tokenwise arm routes qwen35-class through the BATCHED T=1 program
8850    /// (`spec_target_step_h`) instead — ULP-different suffix rows, and the gate measured
8851    /// the near-tie flip at generated token ~8 (research/spec-cache-20260818, qwen r3).
8852    /// So the suffix is fed HERE, mirroring prefill_tick arm-for-arm, not handed to the
8853    /// burst prime.
8854    ///
8855    /// SEED RULE (both sampling regimes; lane/sampled-hit-spec 2026-08-19, sampled draw
8856    /// added by lane/sampled-spec-quality 2026-08-19). The boundary token is produced by
8857    /// EXACTLY the rule the cold burst entry applies to its own first token from the same
8858    /// logits row: `argmax` when greedy, and a `sample_boundary_token` draw at Philox
8859    /// counter 0 when sampled. Both shapes are covered — the entry's boundary logits on a
8860    /// full-cover (empty-suffix) hit, this feed's own boundary logits on a suffix hit.
8861    /// That is what keeps a restored session seed-identical to a cold one PER SEED: the
8862    /// cold session draws from the identical row at counter 0 and then runs its rounds from
8863    /// counter 1, so the restored session admits with `sctr = 1` after its own draw.
8864    /// The WORKER owns the one refusal this constructor cannot see — a constrained request.
8865    /// (The penalized-sampled refusal was LIFTED once the burst's penalty window learned to
8866    /// span the session: `committed` here is the WHOLE prompt, so the restored session's
8867    /// window is the cold session's window. It comes back if `MEMRA_SPEC_PEN_SESSION=0`.)
8868    ///
8869    /// NOT the rolled-back partial-restore hazard: the caller restores at exactly the
8870    /// entry's captured endpoint (`e.pos`) through the shipping whole-entry path;
8871    /// mid-entry (`at < e.pos`) trunk restores stay behind MEMRA_PREFIX_PARTIAL_RESTORE
8872    /// and are never routed here.
8873    ///
8874    /// Failure contract: `Err((Some(cache), why))` before any trunk mutation — the
8875    /// worker rebuilds the plain carrier and the hit serves plain, byte-unchanged.
8876    /// `Err((None, why))` after the suffix feed began — the carrier is part-fed and
8877    /// UNUSABLE; the worker serves the request cold-plain (correct, slower) and the
8878    /// entry stays published for the next request.
8879    #[allow(clippy::too_many_arguments)]
8880    #[allow(clippy::result_large_err)] // allow: the fat error type is the diagnostic contract here; boxing it would change the error surface
8881    pub fn spec_session_from_restored(
8882        &self,
8883        e: &Engine,
8884        mut cache: Cache,
8885        prefix: Vec<u32>,
8886        suffix: &[u32],
8887        draft_k: &CudaSlice<u8>,
8888        draft_v: &CudaSlice<u8>,
8889        draft_k_tok_bytes: usize,
8890        draft_v_tok_bytes: usize,
8891        draft_len: usize,
8892        last_h: &[f32],
8893        // The ENTRY's boundary logits row (the full-cover shape's seed source). May be empty
8894        // when a suffix follows — the feed's own logits are the boundary then.
8895        boundary_logits: &[f32],
8896        // The request's sampler, or None for greedy. Owned here so the seed rule lives in
8897        // ONE place instead of being half-applied by the worker.
8898        sampling: Option<SpecSampling>,
8899        require_anchor: bool,
8900        max_ctx: usize,
8901        // STABLE-BOUNDARY REPUBLICATION (lane/frspec-multiturn-cache, 2026-08-21): ABSOLUTE
8902        // prompt position to split the suffix feed at and capture the extended-entry
8903        // publication + this session's `turn_ckpt` — the worker's stable pre-generation
8904        // boundary (`plain_checkpoint_boundary`). None = legacy prompt-end republication.
8905        // WHY: the prompt-end capture below includes the template's live generation header
8906        // (`<|im_start|>assistant\n<think>\n`), which the next turn's re-render replaces, so
8907        // for a hybrid (whole-entry restores only) every extended entry's last ~2 tokens
8908        // diverged from every future prompt and the hit boundary FROZE at the first
8909        // lcp-split entry forever (measured: cached 6811 of 38228 by turn 8, B4).
8910        republish_at: Option<usize>,
8911    ) -> Result<SpecSession, (Option<Cache>, String)> {
8912        let pos = prefix.len();
8913        let fail = |cache: Cache, msg: String| -> Result<SpecSession, (Option<Cache>, String)> {
8914            Err((Some(cache), msg))
8915        };
8916        if let Err(error) = cache.ensure_usable("spec_session_from_restored") {
8917            drop(cache);
8918            return Err((None, error.to_string()));
8919        }
8920        if self.mtp.is_none() {
8921            return fail(cache, "no MTP head attached (nothing to draft with)".into());
8922        }
8923        if pos == 0 {
8924            return fail(cache, "empty committed prefix".into());
8925        }
8926        if cache.pos != pos {
8927            let msg = format!(
8928                "restored cache pos {} != restored prefix len {pos}",
8929                cache.pos
8930            );
8931            return fail(cache, msg);
8932        }
8933        if draft_len != pos {
8934            return fail(
8935                cache,
8936                format!("draft plane len {draft_len} != restored prefix len {pos}"),
8937            );
8938        }
8939        if pos + suffix.len() >= max_ctx {
8940            return fail(
8941                cache,
8942                format!(
8943                    "prompt {} + suffix would not leave generation room in ctx {max_ctx}",
8944                    pos + suffix.len(),
8945                ),
8946            );
8947        }
8948        let mut scratch = match MtpScratch::new(
8949            e,
8950            &self.cfg,
8951            &self.plan,
8952            max_ctx,
8953            self.mtp.as_ref().and_then(|m| m.geom.as_ref()),
8954        ) {
8955            Ok(s) => s,
8956            Err(err) => return fail(cache, format!("draft scratch alloc failed: {err}")),
8957        };
8958        if scratch.kv.ring.is_some() {
8959            return fail(
8960                cache,
8961                "ring-backed draft scratch (Step35 SWA) cannot take a flat prefix restore".into(),
8962            );
8963        }
8964        if scratch.kv.k_tok_bytes != draft_k_tok_bytes
8965            || scratch.kv.v_tok_bytes != draft_v_tok_bytes
8966        {
8967            return fail(
8968                cache,
8969                format!(
8970                    "draft plane layout {draft_k_tok_bytes}/{draft_v_tok_bytes} != scratch \
8971                     {}/{} bytes/token (stale entry across a format change)",
8972                    scratch.kv.k_tok_bytes, scratch.kv.v_tok_bytes,
8973                ),
8974            );
8975        }
8976        if pos > scratch.cap {
8977            return fail(
8978                cache,
8979                format!(
8980                    "draft plane rows {pos} exceed scratch capacity {}",
8981                    scratch.cap
8982                ),
8983            );
8984        }
8985        let kb = pos * draft_k_tok_bytes;
8986        let vb = pos * draft_v_tok_bytes;
8987        if draft_k.len() < kb || draft_v.len() < vb {
8988            return fail(
8989                cache,
8990                format!(
8991                    "truncated draft plane: K {} < {kb} or V {} < {vb} bytes",
8992                    draft_k.len(),
8993                    draft_v.len(),
8994                ),
8995            );
8996        }
8997        if kb > 0
8998            && let Err(err) = e.copy_u8_into(&mut scratch.kv.k, 0, draft_k, kb)
8999        {
9000            return fail(cache, format!("draft K restore copy failed: {err}"));
9001        }
9002        if vb > 0
9003            && let Err(err) = e.copy_u8_into(&mut scratch.kv.v, 0, draft_v, vb)
9004        {
9005            return fail(cache, format!("draft V restore copy failed: {err}"));
9006        }
9007        if let Err(err) = scratch.set_len(e, pos) {
9008            return fail(cache, format!("draft scratch len set failed: {err}"));
9009        }
9010        let mut last_h_dev = if last_h.len() == self.cfg.n_embd as usize {
9011            // anchor upload failure is acceptance-only when a suffix feed follows (fill
9012            // row-0 falls back to zeros) but FATAL for an empty-suffix continuation (the
9013            // burst entry asserts committed + last_h + next_pred) — the caller says which.
9014            e.htod(last_h).ok()
9015        } else {
9016            None
9017        };
9018        if require_anchor && last_h_dev.is_none() {
9019            return fail(
9020                cache,
9021                "empty-suffix continuation requires the entry's boundary hidden anchor".into(),
9022            );
9023        }
9024        let mut committed = prefix;
9025        // Set on BOTH shapes below (suffix-fed and full-cover) — never left None, which is
9026        // what the empty-suffix continuation assert in the burst entry requires.
9027        let next_pred;
9028        // Philox: (0,0) at admit exactly like a fresh session; a sampled boundary draw below
9029        // consumes counter 0 and leaves 1, which is the state a cold session reaches after
9030        // drawing its own first token from the same row.
9031        let mut sctr = 0u32;
9032        let sampled = sampling.is_some_and(|s| s.temp > 0.0) && spec_sampled_boundary_on();
9033        // Penalty window for the boundary draw: the last `penalty_last_n` tokens of the WHOLE
9034        // prompt, which is what the cold session's own burst sees (Item 2's window). Built
9035        // after the suffix joins `committed` below.
9036        let mut boundary_captures: Vec<SpecBoundaryCapture> = Vec::new();
9037        let mut restored_turn_ckpt: Option<SpecCheckpoint> = None;
9038        if !suffix.is_empty() {
9039            // ---- SUFFIX FEED, mirroring prefill_tick's program selection exactly ----
9040            // From here on the trunk cache mutates: failures return Err((None, _)) and
9041            // the worker serves the request cold-plain instead of reusing the carrier.
9042            let dirty =
9043                |msg: String| -> Result<SpecSession, (Option<Cache>, String)> { Err((None, msg)) };
9044            let n_embd = self.cfg.n_embd as usize;
9045            let t = suffix.len();
9046            let mut h_rows = match e.uninit(t * n_embd) {
9047                Ok(b) => b,
9048                Err(err) => return fail(cache, format!("suffix hidden buffer alloc: {err}")),
9049            };
9050            // STABLE-BOUNDARY split (see `republish_at`): feed stops at the boundary so the
9051            // in-place GDN conv/ssm state can be snapshotted there — the only moment it
9052            // exists (the cold prime-split law). suffix-relative; None = one-segment legacy.
9053            let b_rel = republish_at
9054                .and_then(|abs| abs.checked_sub(pos))
9055                .filter(|&r| r > 0 && r < t);
9056            let mut feed_logits = Vec::new();
9057            let tokenwise_env = std::env::var("MEMRA_PRIME_TOKENWISE").is_ok()
9058                || e.frozen_cpu_experts_prefer_tokenwise_prime();
9059            let mut fed = 0usize;
9060            for seg_end in [b_rel, Some(t)].into_iter().flatten() {
9061                if seg_end <= fed {
9062                    continue;
9063                }
9064                let seg = &suffix[fed..seg_end];
9065                let batched = seg.len() >= crate::hybrid_forward::PRIME_MIN_T && !tokenwise_env;
9066                if batched {
9067                    // prefill_tick's prime arm: request-level prime_cache call; tokens still
9068                    // queued after this segment ride `queued_after` so Step35 arm selection
9069                    // stays keyed to the request's end (tick-seg law).
9070                    match self.prime_cache(e, seg, &mut cache, t - seg_end) {
9071                        Ok((l, _h_seed, hiddens)) => {
9072                            if let Err(err) =
9073                                e.copy_into(&mut h_rows, fed * n_embd, &hiddens, seg.len() * n_embd)
9074                            {
9075                                return dirty(format!("suffix hidden copy: {err}"));
9076                            }
9077                            feed_logits = l;
9078                        }
9079                        Err(err) => return dirty(format!("suffix prime failed: {err}")),
9080                    }
9081                } else {
9082                    // prefill_tick's tokenwise arm: eager decode_step, one token at a time.
9083                    for (i, &tok) in seg.iter().enumerate() {
9084                        match self.decode_step_h(e, tok, &mut cache) {
9085                            Ok((l, h)) => {
9086                                if let Err(err) =
9087                                    e.copy_into(&mut h_rows, (fed + i) * n_embd, &h, n_embd)
9088                                {
9089                                    return dirty(format!("suffix hidden copy: {err}"));
9090                                }
9091                                feed_logits = l;
9092                            }
9093                            Err(err) => return dirty(format!("suffix decode_step failed: {err}")),
9094                        }
9095                    }
9096                }
9097                fed = seg_end;
9098                if Some(seg_end) == b_rel {
9099                    // The stable pre-generation boundary: capture the extended-entry
9100                    // publication AND this session's own turn checkpoint here instead of at
9101                    // prompt-end (both would otherwise carry the volatile live-header tail
9102                    // the next re-render replaces). Failure silent, turn_ckpt convention.
9103                    debug_assert_eq!(
9104                        cache.pos,
9105                        pos + seg_end,
9106                        "stable-boundary capture off the feed split"
9107                    );
9108                    if spec_restore_republish_on()
9109                        && let Ok(snap) = cache.snapshot(e)
9110                    {
9111                        boundary_captures.push(SpecBoundaryCapture {
9112                            snap,
9113                            pos: pos + seg_end,
9114                            logits: feed_logits.clone(),
9115                            last_h: capture_boundary_hidden(e, &h_rows, seg_end, n_embd),
9116                            latent_tails: Vec::new(),
9117                        });
9118                    }
9119                    let anchor: Result<CudaSlice<f32>, Box<dyn std::error::Error>> =
9120                        e.uninit(n_embd).and_then(|mut a| {
9121                            e.copy_view_into(
9122                                &mut a,
9123                                0,
9124                                &h_rows.slice((seg_end - 1) * n_embd..seg_end * n_embd),
9125                                n_embd,
9126                            )?;
9127                            Ok(a)
9128                        });
9129                    if let (Ok(snap), Ok(last_h)) = (cache.snapshot(e), anchor) {
9130                        restored_turn_ckpt = Some(SpecCheckpoint {
9131                            snap,
9132                            pos: pos + seg_end,
9133                            last_h,
9134                        });
9135                    }
9136                }
9137            }
9138            // Draft-scratch fill for the suffix rows, predecessor-paired: row `pos` reads
9139            // the entry's boundary anchor (zeros fallback — acceptance-only), row `pos+i`
9140            // reads h_rows[i-1]. Chunked like the generate path's fill (transients scale
9141            // with T). Fill failures are acceptance-only — truncate to the restored rows
9142            // and continue; the burst's own set_len keeps the invariant.
9143            let _mtp = self.mtp.as_ref().expect("mtp checked above"); // invariant check only; the fill below re-reads self.mtp
9144            let (embd_qt, embd_rb) = self.embd.qt_and_row_bytes(n_embd);
9145            let embd_gpu = if spec_host_embd() {
9146                None
9147            } else {
9148                Some(
9149                    self.embd_gpu
9150                        .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload")),
9151                )
9152            };
9153            let embd_dev = embd_gpu.map(|g| (g, embd_qt, embd_rb));
9154            let fill_chunk = 4096usize;
9155            let mut filled = true;
9156            let mut start = 0usize;
9157            'fill: while start < t {
9158                let end = (start + fill_chunk).min(t);
9159                let tc = end - start;
9160                let Ok(mut phs) = e.zeros(tc * n_embd) else {
9161                    filled = false;
9162                    break 'fill;
9163                };
9164                let (src_lo, dst_off, n_copy) = if start == 0 {
9165                    (0, n_embd, (tc - 1) * n_embd)
9166                } else {
9167                    ((start - 1) * n_embd, 0, tc * n_embd)
9168                };
9169                if start == 0
9170                    && let Some(lh) = last_h_dev.as_ref()
9171                    && e.copy_into(&mut phs, 0, lh, n_embd).is_err()
9172                {
9173                    filled = false;
9174                    break 'fill;
9175                }
9176                if n_copy > 0
9177                    && e.copy_view_into(
9178                        &mut phs,
9179                        dst_off,
9180                        &h_rows.slice(src_lo..src_lo + n_copy),
9181                        n_copy,
9182                    )
9183                    .is_err()
9184                {
9185                    filled = false;
9186                    break 'fill;
9187                }
9188                if self
9189                    .mtp_kv_fill_all(
9190                        e,
9191                        &suffix[start..end],
9192                        &phs,
9193                        pos + start,
9194                        &mut scratch,
9195                        embd_dev,
9196                    )
9197                    .is_err()
9198                {
9199                    filled = false;
9200                    break 'fill;
9201                }
9202                start = end;
9203            }
9204            if !filled {
9205                // acceptance-only: drafts over missing suffix rows are cheap and wrong,
9206                // so keep only the restored rows resident and let verify arbitrate.
9207                if let Err(err) = scratch.set_len(e, pos) {
9208                    return dirty(format!("scratch truncation after failed fill: {err}"));
9209                }
9210            }
9211            // EXTENDED-ENTRY PUBLICATION (lane/sampled-spec-quality, Item 3 — the fix for
9212            // "a restored spec session never publishes an extended entry", SAMPLED-HIT.md
9213            // finding (d)). Pre-lane, publication was armed only for COLD sessions
9214            // (`spec_resumed == 0` in the worker) and both engine capture sites require a
9215            // non-continuation burst — but a converted hit's first burst IS a continuation,
9216            // so a growing conversation learned exactly ONE boundary and turn 3 could never
9217            // hit a longer prefix than turn 2 did.
9218            //
9219            // WHERE, and why it is safe here: `cache.pos == prefix + suffix` at this exact
9220            // line — the trunk is primed over the whole prompt, nothing is generated, and the
9221            // draft plane rows [0..prompt) are filled just above. That is a complete
9222            // whole-entry boundary (`pos == fed_len`), the same shape the cold seed capture
9223            // publishes; the worker's existing publication sweep picks it up because it is
9224            // keyed on non-empty `boundary_captures` and is sampler- and resume-independent.
9225            // NOT the partial-restore hazard: the boundary is this session's own prompt END,
9226            // never mid-entry, so `entry_pos != fed_len` still refuses on the way back in.
9227            // Failure is SILENT by design (the turn_ckpt / boundary-capture convention):
9228            // publication is an optimization, never a correctness dependency.
9229            //
9230            // SUPERSEDED WHEN `republish_at` FIRED (lane/frspec-multiturn-cache): a prompt-end
9231            // entry's tail is the live generation header the next re-render replaces, so on a
9232            // hybrid (whole-entry restores) it can never serve the conversation's next turn —
9233            // the stable-boundary capture above IS this publication, minus the poisoned tail.
9234            if spec_restore_republish_on() && boundary_captures.is_empty() {
9235                debug_assert_eq!(
9236                    cache.pos,
9237                    pos + t,
9238                    "extended-entry capture must sit at the restored session's prompt end",
9239                );
9240                if let Ok(snap) = cache.snapshot(e) {
9241                    boundary_captures.push(SpecBoundaryCapture {
9242                        snap,
9243                        pos: pos + t,
9244                        logits: feed_logits.clone(),
9245                        last_h: capture_boundary_hidden(e, &h_rows, t, n_embd),
9246                        latent_tails: Vec::new(),
9247                    });
9248                }
9249            }
9250            // continuation seed: the feed's boundary logits ARE the plain path's boundary
9251            // logits (same program), so greedy's argmax here is plain's first emitted token,
9252            // and the sampled draw is the cold sampled session's own first token.
9253            next_pred = Some(if sampled {
9254                let sp = sampling.expect("sampled implies a sampler");
9255                // `committed` is still the restored prefix here; the suffix joins it below —
9256                // so this is the last-N window over the WHOLE prompt, exactly the cold
9257                // session's own window at its first token.
9258                let hist = pen_window_seed(&committed, suffix, sp.penalty_last_n);
9259                match sample_boundary_token(
9260                    e,
9261                    &feed_logits,
9262                    &sp,
9263                    &hist,
9264                    &mut sctr,
9265                    "restore-suffix-feed",
9266                ) {
9267                    Ok(t) => t,
9268                    // the trunk is already fed: hand nothing back, the worker serves the
9269                    // request cold-plain. Never fall back to an argmax — that would put a
9270                    // greedy token in a sampled stream to save a slow path.
9271                    Err(err) => {
9272                        return dirty(format!("boundary token draw failed: {err}"));
9273                    }
9274                }
9275            } else {
9276                argmax(&feed_logits) as u32
9277            });
9278            let mut lh = match e.uninit(n_embd) {
9279                Ok(b) => b,
9280                Err(err) => return dirty(format!("boundary hidden alloc: {err}")),
9281            };
9282            if let Err(err) = e.copy_view_into(
9283                &mut lh,
9284                0,
9285                &h_rows.slice((t - 1) * n_embd..t * n_embd),
9286                n_embd,
9287            ) {
9288                return dirty(format!("boundary hidden copy: {err}"));
9289            }
9290            last_h_dev = Some(lh);
9291            committed.extend_from_slice(suffix);
9292        } else {
9293            // FULL-COVER shape (empty suffix — the identical-repeat / agent-loop shape): the
9294            // ENTRY's boundary logits are the boundary row, and this is the token the cold
9295            // session emits from that same row. Owned here rather than in the worker so the
9296            // sampled draw cannot be half-applied on one shape (the worker used to argmax it).
9297            if boundary_logits.is_empty() {
9298                return fail(
9299                    cache,
9300                    "full-cover restore without the entry's boundary logits".into(),
9301                );
9302            }
9303            next_pred = Some(if sampled {
9304                let sp = sampling.expect("sampled implies a sampler");
9305                let hist = pen_window_seed(&committed, &[], sp.penalty_last_n);
9306                match sample_boundary_token(
9307                    e,
9308                    boundary_logits,
9309                    &sp,
9310                    &hist,
9311                    &mut sctr,
9312                    "restore-full-cover",
9313                ) {
9314                    Ok(t) => t,
9315                    // nothing has been mutated on this shape — hand the carrier back and let
9316                    // the hit serve PLAIN (the banked pre-lane path).
9317                    Err(err) => {
9318                        return fail(cache, format!("boundary token draw failed: {err}"));
9319                    }
9320                }
9321            } else {
9322                argmax(boundary_logits) as u32
9323            });
9324        }
9325        Ok(SpecSession {
9326            cache,
9327            scratch,
9328            committed,
9329            last_h: last_h_dev,
9330            next_pred,
9331            sctr,
9332            uctr: 0,
9333            draft_ctx: None,
9334            pending_tok: None,
9335            // Stable-boundary capture from the split feed above (None on the legacy shape):
9336            // a restored session previously parked WITHOUT a checkpoint, so the next turn's
9337            // affinity probe declined ("no turn checkpoint retained") and the conversation
9338            // fell back to the frozen prefix entry forever.
9339            turn_ckpt: restored_turn_ckpt,
9340            telem: SpecTelemetryCounters::default(),
9341            capture_at: None,
9342            boundary_captures,
9343            ckpt_at: None,
9344            capture_disabled: false,
9345        })
9346    }
9347
9348    /// SESSION-AFFINITY REWIND (lane/session-affinity, 2026-08-05): roll `sess` back to its
9349    /// retained prompt-end checkpoint, so a request whose prompt matches
9350    /// `committed[..rewind_pos()]` exactly can resume there and prime only its own delta.
9351    ///
9352    /// EXACTNESS. After this returns, the session is byte-for-byte the state it was in AT that
9353    /// boundary: full-attn KV truncated to it (append-only, position-addressed), GDN conv/ssm
9354    /// restored from the device copy taken there, draft scratch length reset, `committed`
9355    /// truncated, `last_h` = the boundary's predecessor anchor. That is precisely the state a
9356    /// fresh prime of `committed[..pos]` would have produced, so the following suffix prime and
9357    /// every burst after it are identical to a cold run of the same token stream — the
9358    /// committed-tokens-authoritative contract.
9359    ///
9360    /// `next_pred` and `pending_tok` are CLEARED: both describe generation past the boundary,
9361    /// which the rewind discards. The caller therefore must supply a non-empty suffix (a
9362    /// rewound session cannot serve an empty-suffix continuation burst — there is nothing to
9363    /// continue). The persistent draft graph survives: it bakes only session-stable pointers
9364    /// (the scratch KV, the resident embedding), none of which the rewind moves.
9365    ///
9366    /// The checkpoint is CONSUMED (`turn_ckpt` taken): its snapshot buffers are freed here, and
9367    /// this turn's own prime installs a fresh one at the new prompt end. Returns the position
9368    /// rewound to, or `None` when the session holds no checkpoint (caller: full re-prime).
9369    pub fn spec_rewind_to_checkpoint(
9370        &self,
9371        e: &Engine,
9372        sess: &mut SpecSession,
9373    ) -> Result<Option<usize>, Box<dyn std::error::Error>> {
9374        if sess.turn_ckpt.as_ref().is_some_and(|ckpt| {
9375            !sess.cache.can_rollback(&ckpt.snap, 0) || !sess.scratch.can_rewind_to(ckpt.pos)
9376        }) {
9377            return Err(
9378                "SWA ring rewind checkpoint has been lapped; full re-prime required".into(),
9379            );
9380        }
9381        let Some(ckpt) = sess.turn_ckpt.take() else {
9382            return Ok(None);
9383        };
9384        assert!(
9385            ckpt.pos <= sess.committed.len(),
9386            "checkpoint past committed ({} > {})",
9387            ckpt.pos,
9388            sess.committed.len()
9389        );
9390        // Restore through each layer's owning engine. A single primary-engine rollback is not
9391        // sufficient when the serving cache is stage-owned under cross-device PP.
9392        crate::pp::restore_cache_checkpoint(e, self, None, &mut sess.cache, &ckpt.snap)?;
9393        debug_assert_eq!(
9394            sess.cache.pos, ckpt.pos,
9395            "rollback landed off the checkpoint"
9396        );
9397        sess.scratch.set_len(e, ckpt.pos)?;
9398        sess.committed.truncate(ckpt.pos);
9399        sess.last_h = Some(ckpt.last_h);
9400        sess.next_pred = None;
9401        sess.pending_tok = None;
9402        Ok(Some(ckpt.pos))
9403    }
9404
9405    /// Grow a parked speculative session to `target_cap` and rewind it to its retained turn
9406    /// checkpoint without re-priming the checkpoint prefix.
9407    ///
9408    /// The trunk cache is restored exactly like a plain grown cache: append-only full-attention
9409    /// KV rows come from the parked cache, while recurrent state comes from the checkpoint's
9410    /// owned snapshot. The MTP scratch is also context-linear and its rows below the checkpoint
9411    /// remain authoritative, so they are copied into a fresh larger scratch before its length is
9412    /// truncated. Pointer-baking draft graphs are dropped and recaptured on the next burst.
9413    ///
9414    /// All fallible work completes before `sess` is mutated. A failed allocation or copy leaves
9415    /// the parked session intact, allowing the caller one reclaim-and-retry attempt.
9416    pub fn spec_grow_and_rewind_to_checkpoint(
9417        &self,
9418        e: &Engine,
9419        sess: &mut SpecSession,
9420        target_cap: usize,
9421    ) -> Result<Option<usize>, Box<dyn std::error::Error>> {
9422        if target_cap <= sess.cache.max_ctx {
9423            return self.spec_rewind_to_checkpoint(e, sess);
9424        }
9425        let Some(ckpt) = sess.turn_ckpt.as_ref() else {
9426            return Ok(None);
9427        };
9428        if ckpt.pos == 0 || ckpt.pos > sess.committed.len() {
9429            return Err(format!(
9430                "checkpoint pos {} outside committed length {}",
9431                ckpt.pos,
9432                sess.committed.len(),
9433            )
9434            .into());
9435        }
9436        if ckpt.pos > target_cap {
9437            return Err(format!(
9438                "checkpoint pos {} exceeds grown capacity {target_cap}",
9439                ckpt.pos,
9440            )
9441            .into());
9442        }
9443
9444        let mut grown_cache = crate::pp::new_cache_planned(e, &self.cfg, &self.plan, target_cap)?;
9445        let mut grown_scratch = self.new_mtp_scratch(e, target_cap)?;
9446        crate::pp::restore_cache_checkpoint(
9447            e,
9448            self,
9449            Some(&sess.cache),
9450            &mut grown_cache,
9451            &ckpt.snap,
9452        )?;
9453
9454        if sess.scratch.plane_count() != grown_scratch.plane_count() {
9455            return Err("checkpoint draft plane count mismatch".into());
9456        }
9457        for index in 0..sess.scratch.plane_count() {
9458            let (src, _) = sess.scratch.plane(index);
9459            let (dst, _) = grown_scratch.plane_mut(index);
9460            if ckpt.pos > src.len
9461                || src.kv_dim_k != dst.kv_dim_k
9462                || src.kv_dim_v != dst.kv_dim_v
9463                || src.k_tok_bytes != dst.k_tok_bytes
9464                || src.v_tok_bytes != dst.v_tok_bytes
9465            {
9466                return Err(format!(
9467                    "checkpoint draft plane {index} layout mismatch (pos {}, source len {})",
9468                    ckpt.pos, src.len,
9469                )
9470                .into());
9471            }
9472            match (&src.ring, dst.ring.as_ref()) {
9473                (Some(sring), Some(_)) => {
9474                    // Ring-backed draft plane (step35): `ckpt.pos` is absolute and exceeds the
9475                    // physical rows once lapped — same class as the trunk-KV restore panic
9476                    // (2026-08-29 warm-turn-at-40k). Copy the aligned live window, rebase.
9477                    let (new_base, phys) = sring.restore_plan(ckpt.pos).map_err(|err| {
9478                        format!("checkpoint draft plane {index} SWA restore refused: {err}")
9479                    })?;
9480                    let rows = phys.len();
9481                    let kb = rows * src.k_tok_bytes;
9482                    let vb = rows * src.v_tok_bytes;
9483                    if kb > 0 {
9484                        e.copy_u8_range_into(
9485                            &mut dst.k,
9486                            0,
9487                            &src.k,
9488                            phys.start * src.k_tok_bytes,
9489                            kb,
9490                        )?;
9491                    }
9492                    if vb > 0 {
9493                        e.copy_u8_range_into(
9494                            &mut dst.v,
9495                            0,
9496                            &src.v,
9497                            phys.start * src.v_tok_bytes,
9498                            vb,
9499                        )?;
9500                    }
9501                    dst.ring
9502                        .as_mut()
9503                        .expect("ring presence checked above")
9504                        .apply_rebase(new_base);
9505                    if let Some(base_d) = dst.base_d.as_mut() {
9506                        e.set_i32_one(base_d, new_base as i32)?;
9507                    }
9508                }
9509                (None, None) => {
9510                    let kb = ckpt.pos * src.k_tok_bytes;
9511                    let vb = ckpt.pos * src.v_tok_bytes;
9512                    if kb > 0 {
9513                        e.copy_u8_into(&mut dst.k, 0, &src.k, kb)?;
9514                    }
9515                    if vb > 0 {
9516                        e.copy_u8_into(&mut dst.v, 0, &src.v, vb)?;
9517                    }
9518                }
9519                _ => {
9520                    return Err(format!("checkpoint draft plane {index} ring/flat mismatch").into());
9521                }
9522            }
9523        }
9524        grown_scratch.set_len(e, ckpt.pos)?;
9525        // The old scratch is dropped immediately after publication below. Bound its D2D reads
9526        // first; growth happens once per rewritten turn, outside the decode hot loop.
9527        e.stream().synchronize()?;
9528
9529        let ckpt = sess
9530            .turn_ckpt
9531            .take()
9532            .expect("checkpoint remained present through transactional grow");
9533        let pos = ckpt.pos;
9534        sess.cache = grown_cache;
9535        sess.scratch = grown_scratch;
9536        sess.committed.truncate(pos);
9537        sess.last_h = Some(ckpt.last_h);
9538        sess.next_pred = None;
9539        sess.pending_tok = None;
9540        sess.draft_ctx = None;
9541        debug_assert_eq!(sess.cache.pos, pos, "grown rewind landed off checkpoint");
9542        debug_assert!(
9543            (0..sess.scratch.plane_count()).all(|index| sess.scratch.plane(index).0.len == pos),
9544            "grown draft rewind landed off checkpoint"
9545        );
9546        Ok(Some(pos))
9547    }
9548
9549    /// Commit a carried pending bonus (see SpecSession::pending_tok): one T=1 trunk pass
9550    /// (its logits' argmax becomes next_pred) + the draft-KV fill at the carried anchor —
9551    /// byte-identical to the pre-carry session tail. Required before a non-empty-suffix
9552    /// prime, a sampled turn, or parking a session for pool reuse. No-op without a pending.
9553    /// `sampling` is the sampler of the request that will CONSUME the resulting `next_pred`
9554    /// (lane/sampled-spec-quality): this is a boundary site like any other, so a sampled
9555    /// consumer must get a DRAWN token, not an argmax. Pass `None` from the park/demote
9556    /// callers — a pending only ever exists on the GREEDY tail, and the consumer of a
9557    /// park-time flush is a future request whose sampler is not knowable here (residual
9558    /// named at the pool-resume probe in worker.rs and in SAMPLED-QUALITY.md).
9559    pub fn spec_flush_pending(
9560        &self,
9561        e: &Engine,
9562        sess: &mut SpecSession,
9563        sampling: Option<SpecSampling>,
9564    ) -> Result<(), Box<dyn std::error::Error>> {
9565        sess.cache.ensure_usable("spec_flush_pending")?;
9566        let Some(b) = sess.pending_tok.take() else {
9567            return Ok(());
9568        };
9569        if self.mtp.is_none() {
9570            return Err("pending carry requires an MTP head".into());
9571        }
9572        let n_embd = self.cfg.n_embd as usize;
9573        let (embd_qt, embd_rb) = self.embd.qt_and_row_bytes(n_embd);
9574        let embd_gpu = if spec_host_embd() {
9575            None
9576        } else {
9577            Some(
9578                self.embd_gpu
9579                    .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload")),
9580            )
9581        };
9582        let embd_dev = embd_gpu.map(|g| (g, embd_qt, embd_rb));
9583        let pos_b = sess.cache.pos;
9584        sess.scratch.set_len(e, pos_b)?;
9585        let (lg_b, hb) = self.spec_target_step_h(e, b, &mut sess.cache)?;
9586        sess.next_pred = Some(match sampling {
9587            Some(sp) if sp.temp > 0.0 && spec_sampled_boundary_on() => {
9588                // window includes `b` itself: it is committed by this pass, and the pre-lane
9589                // code never counted a boundary token in the penalty history at all.
9590                let hist = pen_window_seed(&sess.committed, &[b], sp.penalty_last_n);
9591                sample_boundary_token(e, &lg_b, &sp, &hist, &mut sess.sctr, "flush-pending")?
9592            }
9593            _ => argmax(&lg_b) as u32,
9594        });
9595        let anchor = sess
9596            .last_h
9597            .as_ref()
9598            .expect("pending carry requires last_h (the predecessor-row anchor)");
9599        self.mtp_kv_fill_all(e, &[b], anchor, pos_b, &mut sess.scratch, embd_dev)?;
9600        sess.last_h = Some(hb);
9601        sess.committed.push(b);
9602        Ok(())
9603    }
9604
9605    /// Solo target feed used only at speculative round boundaries. Step35 serving made its
9606    /// staged batched B=1 graph authoritative, so a speculative session must enter and leave
9607    /// rounds through that same graph. Other model families keep their eager T=1 contract.
9608    fn spec_target_step_h(
9609        &self,
9610        e: &Engine,
9611        token: u32,
9612        cache: &mut Cache,
9613    ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9614        cache.ensure_usable("spec_target_step_h")?;
9615        if !self.sliding_gated_moe_batch_program() && !self.batched_serving_numeric_class() {
9616            return self.decode_step_h(e, token, cache);
9617        }
9618        let pos0 = cache.pos;
9619        let (logits, hidden) = self.decode_step_t_core(e, &[token], pos0, cache, None, None)?;
9620        Ok((e.dtoh(&logits)?, hidden))
9621    }
9622
9623    /// The archs whose LIVE B=1 serving runs the generic BATCHED numeric class (decode_step_batch
9624    /// walk + batched head), so their spec verify must run the SAME class. MoE learned this
9625    /// 2026-08-14 AM (4b777ccc5); the dense hybrid reproduced the identical near-tie flip class
9626    /// the same day on Qwen3.8-27B — eager-class verify logits drift from batched-class serving
9627    /// logits ("1 ULP at layer 2 → 2.3e-1 logit maxdiff at the head"), and the GDN recurrence
9628    /// carries the drift until a near-tie flips deep in generation. One predicate so the five
9629    /// dispatch sites cannot drift apart again.
9630    /// Draft-graph head admissibility (lane/draftcost-moe, 2026-08-20): the capture body
9631    /// (`mtp_head_forward_cap`) supports Dense heads and SOFTMAX device-routed resident-MoE
9632    /// heads. Residency alone is insufficient: Hy3/M3/Step sigmoid routing returns selected
9633    /// experts through a host synchronization, which is capture-illegal. Those heads use the
9634    /// exact eager draft chain until a device-only sigmoid expert program lands. Trunk FFN class
9635    /// is irrelevant — the graph body is the HEAD forward only. One predicate for all three
9636    /// eligibility sites so they cannot drift (the serving numeric-class lesson).
9637    fn mtp_graph_capturable(&self) -> bool {
9638        let sigmoid_router = self.cfg.sigmoid_router().is_some();
9639        for head in self.mtp.iter().chain(self.mtp_extra.iter()) {
9640            let reason = match &head.ffn {
9641                crate::hybrid::Ffn::Dense { .. } => None,
9642                crate::hybrid::Ffn::Moe(mo) if mo.dev_exps.is_none() => {
9643                    Some("non-resident MoE MTP head")
9644                }
9645                crate::hybrid::Ffn::Moe(_) if sigmoid_router => {
9646                    Some("sigmoid-router MoE MTP head requires host-visible routing")
9647                }
9648                crate::hybrid::Ffn::Moe(_) => None,
9649            };
9650            if let Some(reason) = reason {
9651                static NOTICE: std::sync::Once = std::sync::Once::new();
9652                NOTICE.call_once(|| {
9653                    eprintln!(
9654                        "[spec] draft graph unavailable: {reason}; eager draft chain engaged"
9655                    );
9656                });
9657                return false;
9658            }
9659        }
9660        self.mtp.is_some()
9661    }
9662
9663    fn batched_serving_numeric_class(&self) -> bool {
9664        self.plan
9665            .trunk_operations()
9666            .contains(&memra_gguf::model_plan::OperationKind::GatedDeltaNet)
9667    }
9668
9669    /// The family the MTP verify-graph default was measured on: GatedDeltaNet state layers
9670    /// (a `recur` mixer) together with a routed-MoE FFN — Ornith-1.5-35B-A3B and its kin. The
9671    /// server-side twin of this test is `model_forces_spec_replay` (GatedDeltaNet + MoeMlp);
9672    /// keeping the engine's own version structural rather than name-based means a new
9673    /// checkpoint of the same shape inherits the default, and a different shape does not.
9674    /// pub(crate) since lane/graph-launch-guard-sweep-20260831: `dspark_vg_admission_debt`
9675    /// consults it so the MTP-route pool stops escaping the admission charge.
9676    pub(crate) fn vgraph_family_default(&self) -> bool {
9677        let has_linear = self
9678            .layers
9679            .iter()
9680            .any(|l| matches!(l.mixer, Mixer::Linear(_)));
9681        let has_moe = self
9682            .layers
9683            .iter()
9684            .any(|l| matches!(l.ffn, crate::hybrid::Ffn::Moe(_)));
9685        has_linear && has_moe
9686    }
9687
9688    fn sliding_gated_moe_batch_program(&self) -> bool {
9689        self.uses_sliding_gated_moe_program()
9690    }
9691
9692    fn gemma_batch_program(&self) -> bool {
9693        self.uses_gemma_program()
9694    }
9695
9696    /// One spec-decode turn on a live session. `suffix` = the NEW tokens only (turn N+1's user
9697    /// message rendered through the chat template continuation). Returns (new tokens emitted,
9698    /// drafted, accepted); session.committed grows by suffix + emitted.
9699    pub fn generate_spec_session(
9700        &self,
9701        e: &Engine,
9702        sess: &mut SpecSession,
9703        suffix: &[u32],
9704        max_new: usize,
9705        k: usize,
9706    ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
9707        self.generate_spec_session_sampled(e, sess, suffix, max_new, k, None, None)
9708    }
9709
9710    /// Serve-path sampled spec: routes the burst through the rejection-sampling verify with
9711    /// per-SESSION Philox continuity (sess.sctr/uctr). None = env-driven (CLI) or greedy.
9712    /// Filters (top-k/p/min-p) apply SYMMETRICALLY to draft q and verify p — distribution-exact
9713    /// for the filtered target (feat/filtered-spec).
9714    ///
9715    /// `on_commit` (sse-cadence, 2026-08-05): called with each newly-emitted slice of the
9716    /// output — once right after the prime's first token, then once per round commit — so a
9717    /// streaming caller can flush text at round cadence instead of once per burst. The slices
9718    /// are disjoint, in order, and concatenate to exactly the returned token vec. Emission-
9719    /// timing only: token bytes, session state, and exactness are untouched.
9720    ///
9721    /// The returned bool is a CONTINUE-VERDICT (admission yield, 2026-08-06): `false` ends
9722    /// the burst at the current round boundary, exactly as if `max_new` had been reached —
9723    /// the caller's scheduler regains control without waiting the burst out. Burst size is
9724    /// content-neutral (spec-levers battery), so an early exit moves WHEN the burst returns,
9725    /// never what tokens say. The slice may be EMPTY (a poll-only boundary — round-stream
9726    /// drains and the defensive tail flush can land with nothing new committed).
9727    #[allow(clippy::too_many_arguments)]
9728    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
9729    pub fn generate_spec_session_sampled(
9730        &self,
9731        e: &Engine,
9732        sess: &mut SpecSession,
9733        suffix: &[u32],
9734        max_new: usize,
9735        k: usize,
9736        sampling: Option<SpecSampling>,
9737        on_commit: Option<&mut dyn FnMut(&[u32]) -> bool>,
9738    ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
9739        self.generate_spec_session_sampled_prime_split(
9740            e, sess, suffix, max_new, k, sampling, None, on_commit,
9741        )
9742    }
9743
9744    /// Serve-only cold-prime segmentation twin. `prime_split` is the same stable boundary the
9745    /// plain worker would honor before entering its sub-floor tokenwise tail; warm continuations
9746    /// pass `None` and stay on the existing zero-prime path.
9747    #[allow(clippy::too_many_arguments)]
9748    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
9749    pub fn generate_spec_session_sampled_prime_split(
9750        &self,
9751        e: &Engine,
9752        sess: &mut SpecSession,
9753        suffix: &[u32],
9754        max_new: usize,
9755        k: usize,
9756        sampling: Option<SpecSampling>,
9757        prime_split: Option<usize>,
9758        on_commit: Option<&mut dyn FnMut(&[u32]) -> bool>,
9759    ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
9760        self.generate_spec_session_constrained_prime_split(
9761            e,
9762            sess,
9763            suffix,
9764            max_new,
9765            k,
9766            sampling,
9767            None,
9768            prime_split,
9769            on_commit,
9770        )
9771    }
9772
9773    /// `generate_spec_session_sampled` + GRAMMAR (constrained decoding, 2026-08-03): the
9774    /// hook truncates acceptance at the first grammar-illegal token AFTER the exactness
9775    /// verify (grammar is an extra rejection rule, ordering like the batched-verify twins)
9776    /// and replaces an illegal bonus with the MASKED argmax of the target's own verify
9777    /// column — token-identical to constrained plain greedy decode. GREEDY only (the
9778    /// worker routes sampled constrained to plain decode). Acceptance under tight grammars
9779    /// may drop (drafter is unconstrained); that is measured, not hidden.
9780    #[allow(clippy::too_many_arguments)]
9781    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
9782    pub fn generate_spec_session_constrained(
9783        &self,
9784        e: &Engine,
9785        sess: &mut SpecSession,
9786        suffix: &[u32],
9787        max_new: usize,
9788        k: usize,
9789        sampling: Option<SpecSampling>,
9790        constraint: Option<&mut dyn SpecConstraint>,
9791        on_commit: Option<&mut dyn FnMut(&[u32]) -> bool>,
9792    ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
9793        self.generate_spec_session_constrained_prime_split(
9794            e, sess, suffix, max_new, k, sampling, constraint, None, on_commit,
9795        )
9796    }
9797
9798    #[allow(clippy::too_many_arguments)]
9799    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
9800    pub fn generate_spec_session_constrained_prime_split(
9801        &self,
9802        e: &Engine,
9803        sess: &mut SpecSession,
9804        suffix: &[u32],
9805        max_new: usize,
9806        k: usize,
9807        sampling: Option<SpecSampling>,
9808        constraint: Option<&mut dyn SpecConstraint>,
9809        prime_split: Option<usize>,
9810        on_commit: Option<&mut dyn FnMut(&[u32]) -> bool>,
9811    ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
9812        if constraint.is_some() && sampling.is_some_and(|s| s.temp > 0.0) {
9813            return Err(
9814                "constrained spec decode is greedy-only (worker routes sampled \
9815                        constrained to plain decode)"
9816                    .into(),
9817            );
9818        }
9819        // PENDING-CARRY entry flush: a carried bonus precedes any new suffix in the sequence,
9820        // so it must commit BEFORE the suffix primes; the sampled path doesn't carry (its
9821        // round-0 accept needs the commit pass's logits). Empty-suffix greedy bursts — the
9822        // serve continuation case — consume the carry in-loop with zero solo passes.
9823        if sess.pending_tok.is_some()
9824            && (!suffix.is_empty() || sampling.is_some_and(|s| s.temp > 0.0))
9825        {
9826            self.spec_flush_pending(e, sess, sampling)?;
9827        }
9828
9829        // FULL_PREC forces the EAGER draft: the graph capture would enclose cuBLASLt f32 GEMV
9830        // (the FloatBf16 else-branches) and a bf16_to_f32 dequant alloc — neither is stream-capture
9831        // safe. Eager rides matmul/matmul_decode_exact, which dequant FloatBf16 on use. (§item 2.)
9832        // Multi-head MTP (mtp_extra non-empty) no longer disqualifies: the chain captures
9833        // per-head graphs (lane/step37-draft-graph-serving-20260830, MEMRA_MTP_CHAIN_GRAPH).
9834        let graph_draft = std::env::var("MEMRA_SPEC_NOGRAPH").is_err()
9835            && !spec_host_embd()
9836            && self.mtp_graph_capturable()
9837            && k + 2 < 96
9838            && !crate::model::full_prec_enabled();
9839        let was_tracking = e.ctx().is_event_tracking();
9840        if graph_draft && was_tracking {
9841            unsafe {
9842                e.ctx().disable_event_tracking();
9843            }
9844        }
9845        let r = self.generate_spec_inner2(
9846            e,
9847            suffix,
9848            max_new,
9849            k,
9850            graph_draft,
9851            Some(sess),
9852            sampling,
9853            constraint,
9854            on_commit,
9855            prime_split,
9856        );
9857        if graph_draft && was_tracking {
9858            unsafe {
9859                e.ctx().enable_event_tracking();
9860            }
9861        }
9862        let (out, d, a) = r?;
9863        Ok((out, d, a))
9864    }
9865
9866    pub fn generate_spec(
9867        &self,
9868        e: &Engine,
9869        prompt: &[u32],
9870        max_new: usize,
9871        k: usize,
9872    ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
9873        // glm5 T-parallel verify door (lane/glm5-tparallel-verify): an hc trunk with a
9874        // loaded DRAFT SOURCE — the embedded MTP head OR the DFlash2 drafter
9875        // (lane/glm5-dflash-draft-src) — routes to the glm5 draft->verify->rollback loop —
9876        // MEMRA_GLM5_SPEC=1 only (default OFF; flag row in FLAGS.md). Unset/0 falls
9877        // through to the standing named refusal below, byte-identical to the pre-lane
9878        // binary. Same fail-closed manifest stance as the generic path: an unqualified
9879        // MtpSpec rewrite refuses before any drafting.
9880        if self.hyper.is_some()
9881            && crate::glm_spec::glm5_spec_on()
9882            && (self.mtp.is_some() || self.glm5_dflash.is_some())
9883        {
9884            if !self.rewrite_allowed(memra_gguf::execution_manifest::RewriteSurface::MtpSpec) {
9885                return Err("speculative rewrite is not qualified for this ModelPlan".into());
9886            }
9887            return self.generate_spec_glm5(e, prompt, max_new, k);
9888        }
9889        self.refuse_hyper("generate_spec")?;
9890        if crate::pp::pp_cuts(self.layers.len()).is_some()
9891            && !self.rewrite_allowed(memra_gguf::execution_manifest::RewriteSurface::Pipeline)
9892        {
9893            return Err("pipeline rewrite is not qualified for speculative decode".into());
9894        }
9895        if !self.rewrite_allowed(memra_gguf::execution_manifest::RewriteSurface::MtpSpec) {
9896            return Err("speculative rewrite is not qualified for this ModelPlan".into());
9897        }
9898        // FULL_PREC forces eager (see generate_spec_session note): CUDA graph capture cannot
9899        // enclose cuBLASLt f32 GEMV or the bf16_to_f32 dequant alloc the FloatBf16 path needs.
9900        // Multi-head MTP no longer disqualifies (chain graphs; see generate_spec_session).
9901        let graph_draft = std::env::var("MEMRA_SPEC_NOGRAPH").is_err()
9902            && !spec_host_embd()
9903            && self.mtp_graph_capturable()
9904            && k + 2 < 96
9905            && !crate::model::full_prec_enabled();
9906        if !graph_draft {
9907            return self
9908                .generate_spec_inner2(e, prompt, max_new, k, false, None, None, None, None, None);
9909        }
9910        let was_tracking = e.ctx().is_event_tracking();
9911        if was_tracking {
9912            unsafe {
9913                e.ctx().disable_event_tracking();
9914            }
9915        }
9916        let r =
9917            self.generate_spec_inner2(e, prompt, max_new, k, true, None, None, None, None, None);
9918        if was_tracking {
9919            unsafe {
9920                e.ctx().enable_event_tracking();
9921            }
9922        }
9923        r
9924    }
9925
9926    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
9927    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
9928    fn generate_spec_inner2(
9929        &self,
9930        e: &Engine,
9931        prompt: &[u32],
9932        max_new: usize,
9933        k: usize,
9934        graph_draft: bool,
9935        mut sess: Option<&mut SpecSession>,
9936        sampling: Option<SpecSampling>,
9937        mut constraint: Option<&mut dyn SpecConstraint>,
9938        mut on_commit: Option<&mut dyn FnMut(&[u32]) -> bool>,
9939        prime_split: Option<usize>,
9940    ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
9941        assert!(k >= 1, "k must be >= 1");
9942        // sse-cadence flush cursor: everything in out[..flushed] has been handed to on_commit.
9943        let mut flushed = 0usize;
9944        // admission yield (2026-08-06): on_commit's continue-verdict; false = end the burst
9945        // at the next round boundary (same exit as max_new reached — the session tail runs).
9946        // Initialized by the unconditional post-prime flush below.
9947        let mut keep_going;
9948        let mtp = self
9949            .mtp
9950            .as_ref()
9951            .expect("generate_spec requires an MTP head (nextn_predict_layers>0)");
9952        let n_vocab = self.output.out_features();
9953        // FR-Spec: the draft head may be TRIMMED (fewer rows than n_vocab); the draft argmax runs
9954        // over the draft vocab and the winning index maps through d2t to a TARGET token id.
9955        // Everything downstream (verify/accept/commit) sees target ids only — exactness unchanged.
9956        let d_vocab = mtp
9957            .shared_head_head
9958            .as_ref()
9959            .unwrap_or(&self.output)
9960            .out_features();
9961        if !self.mtp_extra.is_empty() {
9962            if self.plan.draft_source != memra_gguf::model_plan::DraftSourcePlan::Embedded
9963                || self.plan.mtp_blocks.len() != self.mtp_head_count()
9964            {
9965                return Err(
9966                    "multi-head MTP requires one embedded canonical block per loaded head".into(),
9967                );
9968            }
9969            // TRIMMED chains (2026-08-27): every head must carry the SAME d2t — the ranking is
9970            // token-frequency and head-independent, and every downstream remap (per-step argmax,
9971            // stream pack, sampled d2t_dev) reads head 0's map, so equality is what makes that
9972            // single map correct for the whole chain. Mixed trimmed/untrimmed is refused.
9973            for (offset, head) in self.mtp_extra.iter().enumerate() {
9974                if head.d2t != mtp.d2t
9975                    || head
9976                        .shared_head_head
9977                        .as_ref()
9978                        .unwrap_or(&self.output)
9979                        .out_features()
9980                        != d_vocab
9981                {
9982                    return Err(format!(
9983                        "embedded MTP head {} has incompatible draft vocabulary",
9984                        offset + 1
9985                    )
9986                    .into());
9987                }
9988            }
9989            eprintln!(
9990                "[mtp-chain] heads={} policy=step-modulo prefix-replay kv=per-head",
9991                self.mtp_head_count()
9992            );
9993        }
9994        let n_embd = self.cfg.n_embd as usize;
9995        // SESSION MODE: reuse the live cache/scratch, prime only the suffix. `base` = tokens
9996        // already committed (their state is in the caches); 0 = fresh single-shot call.
9997        let session_mode = sess.is_some();
9998        let max_ctx = match sess.as_ref() {
9999            Some(s) => s.cache.max_ctx,
10000            None => prompt.len() + max_new + k + 8,
10001        };
10002        let mut own_cache;
10003        let mut own_scratch;
10004        // PREFIX-CACHE capture request threaded out of the session (lane/spec-prefix-cache):
10005        // (requested split, destination list). Single-shot per burst; fresh calls have none.
10006        let mut sess_capture: Option<(Option<usize>, &mut Vec<SpecBoundaryCapture>)> = None;
10007        // STABLE-BOUNDARY turn-checkpoint request (lane/frspec-multiturn-cache): ABSOLUTE
10008        // committed-length position; consumed one-shot like `capture_at`. None = legacy
10009        // prompt-end capture below.
10010        let mut ckpt_req: Option<usize> = None;
10011        // FAIL-SAFE bit threaded out of the session (see `SpecSession::capture_disabled`).
10012        let mut sess_capture_disabled = false;
10013        let (
10014            cache,
10015            scratch,
10016            mut sess_tail,
10017            mut sess_draft_slot,
10018            mut sess_pending_slot,
10019            sess_ckpt_slot,
10020            sess_telem,
10021        ): (
10022            &mut Cache,
10023            &mut MtpScratch,
10024            Option<(
10025                &mut Vec<u32>,
10026                &mut Option<CudaSlice<f32>>,
10027                &mut Option<u32>,
10028                &mut u32,
10029                &mut u32,
10030            )>,
10031            Option<&mut Option<DraftGraphCtx>>,
10032            Option<&mut Option<u32>>,
10033            Option<&mut Option<SpecCheckpoint>>,
10034            Option<&SpecTelemetryCounters>,
10035        ) = match sess.take() {
10036            Some(sr) => {
10037                let SpecSession {
10038                    cache,
10039                    scratch,
10040                    committed,
10041                    last_h,
10042                    next_pred,
10043                    sctr: s_sctr,
10044                    uctr: s_uctr,
10045                    draft_ctx,
10046                    pending_tok,
10047                    turn_ckpt,
10048                    telem,
10049                    capture_at,
10050                    boundary_captures,
10051                    ckpt_at,
10052                    capture_disabled,
10053                } = sr;
10054                sess_capture_disabled = *capture_disabled;
10055                sess_capture = Some((capture_at.take(), boundary_captures));
10056                ckpt_req = ckpt_at.take();
10057                (
10058                    cache,
10059                    scratch,
10060                    Some((committed, last_h, next_pred, s_sctr, s_uctr)),
10061                    Some(draft_ctx),
10062                    Some(pending_tok),
10063                    Some(turn_ckpt),
10064                    Some(telem),
10065                )
10066            }
10067            None => {
10068                // STAGE-OWNED KV (lane/pp2-spec 2026-08-06) — see `new_session`. Door shut =
10069                // `Cache::new` verbatim.
10070                own_cache = crate::pp::new_cache_planned(e, &self.cfg, &self.plan, max_ctx)?;
10071                // Persistent scratch = max_ctx rows (~2KB/token quantized).
10072                own_scratch = self.new_mtp_scratch(e, max_ctx)?;
10073                (
10074                    &mut own_cache,
10075                    &mut own_scratch,
10076                    None,
10077                    None,
10078                    None,
10079                    None,
10080                    None,
10081                )
10082            }
10083        };
10084        cache.ensure_usable("generate_spec")?;
10085        if scratch.plane_count() != self.mtp_head_count() {
10086            return Err(format!(
10087                "MTP scratch/head count mismatch ({}/{})",
10088                scratch.plane_count(),
10089                self.mtp_head_count()
10090            )
10091            .into());
10092        }
10093        let base = cache.pos;
10094        // PENDING-CARRY consume (2026-08-01): a carried bonus reaches here only on the
10095        // empty-suffix GREEDY continuation path (generate_spec_session_sampled flushed every
10096        // other case). It enters the round loop as round-0's pending — verify col 0 — exactly
10097        // like a mid-burst full-accept boundary: no init feed, no tail commit pass.
10098        let carried_pending: Option<u32> = sess_pending_slot.as_mut().and_then(|s| s.take());
10099        // PERSISTENT DRAFT KV (the only mode since 2026-07-08 — the legacy round-local scratch,
10100        // MEMRA_SPEC_KVLOCAL, measured -35 acceptance pts on the 27B p3 sweep and was removed;
10101        // acceptance-only — exactness is verify's job either way).
10102        // HIDDEN-PAIRING CONVENTION (DEFAULT = predecessor-row, 2026-07-04 — the 27B acceptance
10103        // unlock, +16pts): the MTP head is TRAINED on rows pairing token x_p with the trunk
10104        // hidden of its PREDECESSOR h_{p-1} (the reference engine's mtp_update shifts the target
10105        // hiddens right by one; its draft step 0 feeds (id_last, TRUE hidden of the row id_last
10106        // was sampled from)). memra's historical convention paired SAME-ROW (x_p, h_p) in the fill
10107        // and seeded chain step 0 through an extra MTP pass on a duplicated token (the
10108        // pseudo-seed) — measured 27B p2 K=3 acceptance 0.569 vs 0.731, p3 0.445 vs 0.63+, and
10109        // the chain steps j>=1 were already predecessor-shaped, so ONLY the fill + step-0 seed
10110        // move. The fill shifts by one and the chain seeds from the predecessor's true hidden
10111        // DIRECTLY (vh_seed / vx[j-1]) — the pseudo pass disappears (one MTP-block pass saved
10112        // per round on top of the acceptance win). Draft-quality-only: exactness stays the
10113        // verify's job either way. (The legacy same-row pairing seam, MEMRA_SPEC_HSAME, and its
10114        // pseudo-seed passes were removed 2026-07-08 — predecessor pairing won by +16 acc pts;
10115        // the legacy round-local scratch, MEMRA_SPEC_KVLOCAL, went with it.)
10116        // REPLAY-FREE PARTIAL ACCEPT (default, 2026-07-03): partial rounds keep the verify's own
10117        // bit-identical committed-prefix state (KV truncate + recur rebuild from the VerifyCkpt)
10118        // and leave the bonus PENDING — no duplicate trunk pass (profiled ~0.54 extra full weight
10119        // reads/round at long ctx). MEMRA_SPEC_REPLAY=1 restores the legacy rollback+replay (A/B
10120        // + fallback seam).
10121        // Qwen35-MoE replay pin LIFTED (lane/draftcost-moe, 2026-08-20). The pin's stated
10122        // bar — the retained verify-state commit proven equivalent to sequential serving —
10123        // was waiting on this arch running the serving batched verify class, which the
10124        // t-parallel admission (this lane, increment 1) provided: the VerifyCkpt the
10125        // replay-free commit consumes is now produced by the SAME serving-class verify that
10126        // qualified dense qwen35 on 2026-08-15 (where the per-round duplicate replay
10127        // measured 69 -> 30 tok/s). Qualification receipts (run-spec K=1..8 both arms,
10128        // 8-prompt replay-vs-replay-free canary, long-prompt cell):
10129        // research/draftcost-moe-20260820/RECEIPTS.md. MEMRA_SPEC_REPLAY=1 stays the
10130        // rollback + A/B seam.
10131        let spec_replay = spec_replay_env_enabled();
10132        if constraint.is_some() && spec_replay {
10133            return Err(
10134                "constrained spec decode does not support MEMRA_SPEC_REPLAY=1 \
10135                        (legacy replay commits an unmasked bonus)"
10136                    .into(),
10137            );
10138        }
10139        // TRUE-HIDDEN REFRESH (default in persistent-draft-KV mode): every round overwrites the
10140        // committed positions' scratch entries from the verify's exact hiddens (mtp_kv_fill batch)
10141        // instead of keeping chain-approximate entries. MEMRA_SPEC_NOREFRESH=1 = legacy (A/B seam).
10142        let refresh = std::env::var("MEMRA_SPEC_NOREFRESH").is_err();
10143        if !refresh && !self.mtp_extra.is_empty() {
10144            return Err("multi-head MTP requires exact accepted-prefix refresh".into());
10145        }
10146
10147        // prime: BATCHED cache prime (prime_cache — the measured #1 e2e gap: tokenwise primed at
10148        // ~102/38 tok/s vs the engine's ~2000-5900 tok/s batched prefill). prime_cache returns the
10149        // full pre-output_norm hidden stack [T, n_embd], which IS prompt_h (the persistent-draft-KV
10150        // mtp_kv_fill input) — no per-token collection needed. Prompts below PRIME_MIN_T, and
10151        // MEMRA_PRIME_TOKENWISE=1, and frozen Hy3 CPU/GPU expert splits take the tokenwise
10152        // decode_step_h loop. The latter avoids transient GPU staging of the spilled expert bank.
10153        // EMPTY-SUFFIX CONTINUATION (serve bursts): a session turn with NO new tokens resumes
10154        // generation exactly where the last turn stopped — no prime at all. The stashed
10155        // `next_pred` plays prime_logits' role: it is the token produced from the logits after
10156        // committed.last() by the same rule this entry applies to a cold prime's last row —
10157        // an argmax when greedy, a `sample_boundary_token` draw when sampled (the burst tail,
10158        // or `spec_session_from_restored` for a converted prefix-cache hit, did the drawing
10159        // where the sampler and the session's Philox counters were live). `last_h` seeds the
10160        // predecessor pairing below. Fresh calls and non-empty suffixes take the normal path.
10161        let continuation = prompt.is_empty();
10162        if continuation {
10163            assert!(session_mode, "empty prompt requires a session");
10164            assert!(
10165                sess_tail
10166                    .as_ref()
10167                    .is_some_and(|(c, lh, np, _, _)| !c.is_empty()
10168                        && lh.is_some()
10169                        && (np.is_some() || carried_pending.is_some())),
10170                "empty-suffix continuation needs a primed session (committed + last_h + next_pred|pending)"
10171            );
10172        }
10173        let mut prime_logits;
10174        let mut prompt_h: Option<CudaSlice<f32>> = None;
10175        let t_prime = std::time::Instant::now();
10176        let batched_prime = !continuation
10177            && prompt.len() >= crate::hybrid_forward::PRIME_MIN_T
10178            && std::env::var("MEMRA_PRIME_TOKENWISE").is_err()
10179            && !e.frozen_cpu_experts_prefer_tokenwise_prime();
10180        let prime_split = prime_split.filter(|&split| split > 0 && split < prompt.len());
10181        if prime_split.is_some() && continuation {
10182            return Err("spec prime split requires a non-empty prime".into());
10183        }
10184        // STABLE-BOUNDARY TURN CHECKPOINT stop (lane/frspec-multiturn-cache, 2026-08-21):
10185        // the worker's `ckpt_at` request, ABSOLUTE -> prompt-relative. On WARM bursts
10186        // (base != 0, an affinity-rewound or pool-resumed session priming its own delta)
10187        // this is the only stop; on COLD bursts it usually coincides with `prime_split`
10188        // (both are the plain tier's stable pre-generation boundary). A boundary the prime
10189        // cannot honor (outside this prime's range) silently drops the capture — the
10190        // turn_ckpt convention: the next turn re-primes in full, never a wrong resume.
10191        let ckpt_rel = if continuation {
10192            None
10193        } else {
10194            ckpt_req
10195                .and_then(|abs| abs.checked_sub(base))
10196                .filter(|&r| r > 0 && r < prompt.len())
10197        };
10198        // Prime stops, ordered: each is a boundary the prime halts at so the in-place GDN
10199        // conv/ssm state can be snapshotted there (the only moment it exists). One stop =
10200        // the legacy single-split program, byte-for-byte.
10201        let mut stops: Vec<usize> = Vec::new();
10202        for b in [prime_split, ckpt_rel].into_iter().flatten() {
10203            if !stops.contains(&b) {
10204                stops.push(b);
10205            }
10206        }
10207        stops.sort_unstable();
10208        // Captured at the ckpt stop, installed into the session slot post-prime (replacing
10209        // the legacy prompt-end capture). Some(None) = capture attempted and failed -> the
10210        // slot is cleared (a stale checkpoint would rewind to the WRONG boundary).
10211        let mut ckpt_early: Option<Option<SpecCheckpoint>> = None;
10212        if continuation {
10213            prime_logits = Vec::new();
10214        } else if !stops.is_empty() {
10215            if let Some(&first) = stops.first()
10216                && prime_split == Some(first)
10217                && first < crate::hybrid_forward::PRIME_MIN_T
10218            {
10219                return Err(format!(
10220                    "spec prime split {first} is below PRIME_MIN_T {}",
10221                    crate::hybrid_forward::PRIME_MIN_T,
10222                )
10223                .into());
10224            }
10225            // Mirror the plain worker's boundary stops exactly. Each segment is a
10226            // request-level prime (`queued_after` keeps Step35 arm selection independent of
10227            // the stops — tick-seg law); a segment below PRIME_MIN_T (and the final tail
10228            // under MEMRA_PRIME_TOKENWISE) takes the same eager tokenwise continuation as
10229            // prefill_tick. Retain every hidden row so the draft scratch fill remains one
10230            // coherent prompt.
10231            let mut h_all = e.uninit(prompt.len() * n_embd)?;
10232            prime_logits = Vec::new();
10233            let mut prev = 0usize;
10234            for seg_end in stops.iter().copied().chain(std::iter::once(prompt.len())) {
10235                if seg_end <= prev {
10236                    continue;
10237                }
10238                let seg = &prompt[prev..seg_end];
10239                let is_final = seg_end == prompt.len();
10240                let batched_seg = seg.len() >= crate::hybrid_forward::PRIME_MIN_T
10241                    && (!is_final
10242                        || (std::env::var("MEMRA_PRIME_TOKENWISE").is_err()
10243                            && !e.frozen_cpu_experts_prefer_tokenwise_prime()));
10244                if batched_seg {
10245                    let (l, _, h_seg) =
10246                        self.prime_cache(e, seg, &mut *cache, prompt.len() - seg_end)?;
10247                    e.copy_into(&mut h_all, prev * n_embd, &h_seg, seg.len() * n_embd)?;
10248                    prime_logits = l;
10249                } else {
10250                    for (i, &tok) in seg.iter().enumerate() {
10251                        let (l, h) = self.decode_step_h(e, tok, &mut *cache)?;
10252                        e.copy_into(&mut h_all, (prev + i) * n_embd, &h, n_embd)?;
10253                        prime_logits = l;
10254                    }
10255                }
10256                prev = seg_end;
10257                if is_final {
10258                    break;
10259                }
10260                debug_assert_eq!(cache.pos, base + seg_end, "prime stop landed off boundary");
10261                // PREFIX-CACHE BOUNDARY CAPTURE (lane/spec-prefix-cache): the GDN conv/ssm
10262                // states are about to be advanced in place by the next segment, so this is
10263                // the ONLY moment the boundary's recurrent state exists. Capture iff the
10264                // worker requested exactly this stop (cold sessions only — `capture_at` is
10265                // never armed warm). A failed snapshot is silent (turn_ckpt convention) —
10266                // publication is an optimization, never a correctness dependency.
10267                if base == 0
10268                    && let Some((requested, slot)) = sess_capture.as_mut()
10269                {
10270                    // Publish at the requested miss-LCP stop (the shared-prefix class)
10271                    // AND at the stable-boundary stop (the next-turn re-render class,
10272                    // lane/frspec-multiturn-cache) — the same boundary set the plain
10273                    // prefill tick learns. Without the second entry, the turn after a
10274                    // cold re-park could only hit the OLDER lcp entry (the measured
10275                    // one-turn transient: t3 restored 607 of 24122 while the plain arm
10276                    // rewound to 15222). Dedupe is the worker sweep's has_key.
10277                    if (*requested == Some(seg_end) || ckpt_rel == Some(seg_end))
10278                        && let Ok(snap) = cache.snapshot(e)
10279                    {
10280                        slot.push(SpecBoundaryCapture {
10281                            snap,
10282                            pos: seg_end,
10283                            logits: prime_logits.clone(),
10284                            // rows [0..seg_end) of h_all are primed — the following
10285                            // segments append, never overwrite.
10286                            last_h: capture_boundary_hidden(e, &h_all, seg_end, n_embd),
10287                            latent_tails: Vec::new(),
10288                        });
10289                    }
10290                }
10291                // SESSION-AFFINITY TURN CHECKPOINT at the STABLE boundary (see `ckpt_at`):
10292                // same snapshot mechanics, installed post-prime in place of the prompt-end
10293                // capture the re-render class always diverged below.
10294                if ckpt_rel == Some(seg_end) {
10295                    let anchor: Result<CudaSlice<f32>, Box<dyn std::error::Error>> =
10296                        e.uninit(n_embd).and_then(|mut a| {
10297                            e.copy_view_into(
10298                                &mut a,
10299                                0,
10300                                &h_all.slice((seg_end - 1) * n_embd..seg_end * n_embd),
10301                                n_embd,
10302                            )?;
10303                            Ok(a)
10304                        });
10305                    ckpt_early = Some(match (cache.snapshot(e), anchor) {
10306                        (Ok(snap), Ok(last_h)) => Some(SpecCheckpoint {
10307                            snap,
10308                            pos: base + seg_end,
10309                            last_h,
10310                        }),
10311                        _ => None,
10312                    });
10313                }
10314            }
10315            if std::env::var("MEMRA_SPEC_STATS").as_deref() == Ok("1") {
10316                eprintln!(
10317                    "[spec-prime] stops={stops:?} tail={}",
10318                    prompt.len() - stops.last().copied().unwrap_or(0)
10319                );
10320            }
10321            prompt_h = Some(h_all);
10322        } else if batched_prime {
10323            let (l, _h_seed, hiddens) = self.prime_cache(e, prompt, &mut *cache, 0)?;
10324            prime_logits = l;
10325            prompt_h = Some(hiddens);
10326        } else {
10327            prime_logits = Vec::new();
10328            prompt_h = Some(e.uninit(prompt.len() * n_embd)?);
10329            for (i, &tok) in prompt.iter().enumerate() {
10330                let (l, h) = self.spec_target_step_h(e, tok, &mut *cache)?;
10331                if let Some(ph) = prompt_h.as_mut() {
10332                    e.copy_into(ph, i * n_embd, &h, n_embd)?;
10333                }
10334                prime_logits = l;
10335            }
10336        }
10337        e.stream().synchronize()?;
10338        // PREFIX-CACHE SEED CAPTURE (lane/spec-prefix-cache): boundary == prompt end (the seed
10339        // case — no shared-prefix split, publish the whole prompt). The prime just finished, so
10340        // cache.pos == base + prompt.len() and the recurrent state IS the boundary state;
10341        // prime_logits are the boundary logits. Cold sessions only (base == 0) — same law as
10342        // prime_split. The mid-prompt capture above already consumed the request if it matched.
10343        if !continuation
10344            && base == 0
10345            && let Some((requested, slot)) = sess_capture.as_mut()
10346            && *requested == Some(prompt.len())
10347            && slot.is_empty()
10348        {
10349            debug_assert_eq!(cache.pos, prompt.len(), "seed capture off prompt end");
10350            if let Ok(snap) = cache.snapshot(e) {
10351                slot.push(SpecBoundaryCapture {
10352                    snap,
10353                    pos: prompt.len(),
10354                    logits: prime_logits.clone(),
10355                    last_h: prompt_h
10356                        .as_ref()
10357                        .map(|ph| capture_boundary_hidden(e, ph, prompt.len(), n_embd))
10358                        .unwrap_or_default(),
10359                    latent_tails: Vec::new(),
10360                });
10361            }
10362        }
10363        // Harness timing contract (see crate::PRIME_NANOS): gen-only throughput without the
10364        // prime-subtraction hack.
10365        crate::PRIME_NANOS.store(
10366            t_prime.elapsed().as_nanos() as u64,
10367            std::sync::atomic::Ordering::Relaxed,
10368        );
10369
10370        let (embd_qt, embd_rb) = self.embd.qt_and_row_bytes(n_embd);
10371        // Resident table is fastest when it fits. Large spill deployments can preserve that HBM
10372        // for expert-cache slots and gather only the exact rows needed by MTP/verify from host.
10373        let host_embd = spec_host_embd();
10374        let embd_gpu = if host_embd {
10375            None
10376        } else {
10377            Some(
10378                self.embd_gpu
10379                    .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload")),
10380            )
10381        };
10382        let embd_dev = embd_gpu.map(|g| (g, embd_qt, embd_rb));
10383        if host_embd {
10384            eprintln!(
10385                "[spec] host-row embedding: {} bytes kept off HBM",
10386                self.embd.raw.len()
10387            );
10388        }
10389        let mut out: Vec<u32> = Vec::with_capacity(max_new);
10390        let mut total_drafted = 0usize;
10391        let mut total_accepted = 0usize;
10392
10393        // --- SAMPLER FIRST (lane/sampled-spec-quality, 2026-08-19) ---
10394        // The sampler config, the session's Philox counters and the penalty window are parsed
10395        // HERE, above the boundary-token selection, because the boundary token must be drawn
10396        // from the sampler the request asked for. Pre-lane this block sat ~50 lines BELOW the
10397        // selection, which is the whole mechanical reason the boundary token was an argmax:
10398        // the sampler state was not in scope yet. Nothing here depends on the round loop, so
10399        // moving it up is a pure reordering for greedy (`sampled == false` ⇒ every branch
10400        // below takes the argmax path it always took).
10401        // --- SAMPLED SPEC (MEMRA_SPEC_TEMP>0, research/sampled-spec-impl-map.md): rejection-
10402        // sampling verify (Leviathan/Chen) — accept draft x at u < p(x)/q(x), resample from
10403        // norm(max(0,p-q)) on reject, bonus sampled from p on full accept. Counter-based Philox
10404        // everywhere (seed, event) -> reproducible. temp==0/unset = the greedy path, untouched.
10405        let sp = sampling.unwrap_or_else(|| SpecSampling {
10406            temp: std::env::var("MEMRA_SPEC_TEMP")
10407                .ok()
10408                .and_then(|v| v.parse().ok())
10409                .unwrap_or(0.0),
10410            seed: std::env::var("MEMRA_SEED")
10411                .ok()
10412                .and_then(|v| v.parse().ok())
10413                .unwrap_or(42),
10414            top_k: std::env::var("MEMRA_TOP_K")
10415                .ok()
10416                .and_then(|v| v.parse().ok())
10417                .unwrap_or(0),
10418            top_p: std::env::var("MEMRA_TOP_P")
10419                .ok()
10420                .and_then(|v| v.parse().ok())
10421                .unwrap_or(1.0),
10422            min_p: std::env::var("MEMRA_MIN_P")
10423                .ok()
10424                .and_then(|v| v.parse().ok())
10425                .unwrap_or(0.0),
10426            penalty_last_n: std::env::var("MEMRA_PENALTY_LAST_N")
10427                .ok()
10428                .and_then(|v| v.parse().ok())
10429                .unwrap_or(0),
10430            penalty_repeat: std::env::var("MEMRA_PENALTY_REPEAT")
10431                .ok()
10432                .and_then(|v| v.parse().ok())
10433                .unwrap_or(1.0),
10434            penalty_freq: std::env::var("MEMRA_PENALTY_FREQ")
10435                .ok()
10436                .and_then(|v| v.parse().ok())
10437                .unwrap_or(0.0),
10438            penalty_present: std::env::var("MEMRA_PENALTY_PRESENT")
10439                .ok()
10440                .and_then(|v| v.parse().ok())
10441                .unwrap_or(0.0),
10442        });
10443        let (sp_temp, sp_seed) = (sp.temp, sp.seed);
10444        let sampled = sp_temp > 0.0;
10445        // Counters resume from the session (burst continuity: randomness must never repeat
10446        // across generate_spec_session calls); one-shot callers start at (0,0). Read through
10447        // sess_tail — `sess` was take()n into it above, so sess.as_ref() here is always None.
10448        let mut sctr: u32 = sess_tail.as_ref().map(|(_, _, _, s, _)| **s).unwrap_or(0);
10449        let mut uctr: u32 = sess_tail.as_ref().map(|(_, _, _, _, u)| **u).unwrap_or(0);
10450        // Penalties (v2.1): applied to COPIES of q rows and p columns symmetrically (exactness
10451        // for the penalized+filtered target). History = generated tokens, host-tracked window.
10452        let pen_on = sampled
10453            && sp.penalty_last_n > 0
10454            && (sp.penalty_repeat != 1.0 || sp.penalty_freq != 0.0 || sp.penalty_present != 0.0);
10455        // SESSION-SPANNING PENALTY WINDOW (Item 2). Pre-lane this was
10456        // `prompt.iter().rev().take(64).rev()` — the BURST's suffix slice — so a continuation
10457        // burst (the majority of a stream's tokens, and ALL of a converted cache hit's) started
10458        // with an EMPTY penalty history and the client's repetition/frequency/presence penalties
10459        // silently reset at every burst boundary. The window now spans `committed ++ prompt`,
10460        // which is what the API contract says and what the plain sampler's own `history` does.
10461        // Byte-identical to the pre-lane seed for a cold turn-1 burst at the default window.
10462        let mut pen_hist: Vec<u32> = if pen_on {
10463            let sess_hist: &[u32] = if spec_pen_session_on() {
10464                sess_tail
10465                    .as_ref()
10466                    .map(|(c, ..)| c.as_slice())
10467                    .unwrap_or(&[])
10468            } else {
10469                &[] // MEMRA_SPEC_PEN_SESSION=0: pre-lane burst-local window
10470            };
10471            pen_window_seed(sess_hist, prompt, sp.penalty_last_n)
10472        } else {
10473            Vec::new()
10474        };
10475        // First generated token = the BOUNDARY token: greedy takes the argmax of the prompt's
10476        // last logits (== greedy's first token, byte-contract); SAMPLED draws it from the
10477        // request's own filtered/penalized target through the session's Philox stream
10478        // (`sample_boundary_token`, lane/sampled-spec-quality Item 1 — pre-lane this was an
10479        // argmax in both regimes, so ~1 token per burst of a sampled stream was greedy).
10480        // Emit it, then FEED it to establish the loop invariant below.
10481        // PENDING-CARRY: the carried bonus was already emitted by the LAST burst — it becomes
10482        // last_token WITHOUT re-emission, and round 0 consumes it as pending (no init feed).
10483        // CONSTRAINED entry rules: the first emitted token is the MASKED argmax of the
10484        // prompt's last logits (plain constrained-greedy identity); a continuation without
10485        // a carried pending would emit an UNMASKED stashed next_pred — refused loudly (the
10486        // worker never resumes constrained sessions from the pool, so this cannot fire).
10487        if let Some(c) = constraint.as_deref_mut() {
10488            if continuation && carried_pending.is_none() {
10489                return Err("constrained spec continuation requires a carried pending \
10490                            (pool resume is unconstrained-only)"
10491                    .into());
10492            }
10493            if !continuation {
10494                c.mask_logits(&mut prime_logits)
10495                    .map_err(|e2| format!("constraint: {e2}"))?;
10496            }
10497        }
10498        let mut last_token = if let Some(b) = carried_pending {
10499            b
10500        } else if continuation {
10501            // A continuation's boundary token was DRAWN by the burst that stashed it (the
10502            // session tail below), or by `spec_session_from_restored` for a converted
10503            // prefix-cache hit — in both cases from the correct logits row with this same
10504            // session's Philox stream, which is why it can be consumed here as-is.
10505            sess_tail.as_ref().unwrap().2.unwrap()
10506        } else if sampled && constraint.is_none() && spec_sampled_boundary_on() {
10507            sample_boundary_token(e, &prime_logits, &sp, &pen_hist, &mut sctr, "cold-prime")?
10508        } else {
10509            // greedy (byte contract), the rollback door, or constrained (masked-argmax
10510            // identity — the worker routes sampled+constrained to the plain path, and this
10511            // function refuses the combination outright above).
10512            argmax(&prime_logits) as u32
10513        };
10514        if pen_on {
10515            // The boundary token is a GENERATED token: the plain sampler `accept()`s every
10516            // emitted token into its penalty history, and pre-lane the burst's first token
10517            // was invisible to penalties forever (never pushed, and never in `committed`
10518            // until this burst's tail). Covers the carry/continuation seeds too — neither is
10519            // in `committed` yet.
10520            pen_hist.push(last_token);
10521        }
10522        if carried_pending.is_none() {
10523            out.push(last_token);
10524            // grammar advances with every emitted token (carried pendings were consumed
10525            // by the burst that emitted them).
10526            if let Some(c) = constraint.as_deref_mut() {
10527                c.consume(last_token)
10528                    .map_err(|e2| format!("constraint: {e2}"))?;
10529            }
10530        }
10531        if continuation {
10532            // draft-KV invariant: entries [0..base) are the session's exact fills; truncate any
10533            // overhang so the chain's first append lands at slot base (== committed.len()).
10534            scratch.set_len(e, base)?;
10535        }
10536        // sse-cadence: hand the caller every not-yet-flushed token (disjoint in-order slices
10537        // concatenating to the full `out`). Called after the prime's first token and after each
10538        // round commit — emission timing only, token bytes untouched. The slice may be EMPTY
10539        // (poll-only boundary: zero-round folds commit nothing new); returns the caller's
10540        // continue-verdict (admission yield, 2026-08-06) — false ends the burst at this round.
10541        #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
10542        fn flush_commit(
10543            cb: &mut Option<&mut dyn FnMut(&[u32]) -> bool>,
10544            out: &[u32],
10545            flushed: &mut usize,
10546        ) -> bool {
10547            if let Some(f) = cb.as_mut() {
10548                let keep = f(&out[*flushed..]);
10549                *flushed = out.len();
10550                keep
10551            } else {
10552                true
10553            }
10554        }
10555        keep_going = flush_commit(&mut on_commit, &out, &mut flushed);
10556        // INVARIANT at loop top: `last_token` is the most-recently-committed/emitted token, its
10557        // KV+recur state IS in `cache` (cache.pos = position right AFTER last_token), `last_pred`
10558        // is the greedy ARGMAX of the logits that predict the token FOLLOWING last_token, and
10559        // `h_seed` = last_token's pre-output_norm hidden. Establish it by feeding last_token once
10560        // (mirrors plain greedy). DEVICE-ARGMAX lever: the accept walk only ever consumes the
10561        // argmax of those logits — never the full vector — so a host u32 replaces the Vec<f32>.
10562        // Trimmed heads: q lives on the trimmed vocab; accept gathers use the TRIMMED index and
10563        // the residual scatters q into target-id space (q=-inf off-trim — the head cannot propose
10564        // those, so their residual mass is p(x), correct by construction).
10565        let d2t_dev: Option<CudaSlice<u32>> = if sampled || crate::spec::spec_stream() {
10566            match &mtp.d2t {
10567                Some(map) => Some(e.htod_u32_v(map)?),
10568                None => None,
10569            }
10570        } else {
10571            None
10572        };
10573        let mut q_full_buf: Option<CudaSlice<f32>> = None;
10574        // host Philox4x32-10 accept-test uniforms: module fn `host_u01` (shared with the
10575        // dspark sampled-admission walk); byte-identical to the closure it replaces.
10576        let mut draft_logits: Vec<CudaSlice<f32>> = Vec::new(); // retained head logits (q), per slot
10577        let mut draft_stats: Vec<(f32, f32, f32)> = Vec::new(); // (row_max, th_e, z_e) per slot
10578        let mut perturb_buf: Option<CudaSlice<f32>> = None; // gumbel scratch (max(n_vocab,d_vocab))
10579        let mut sample_tok = e.alloc_u32_zeroed(1)?; // residual/bonus sample out
10580        let mut col_buf: Option<CudaSlice<f32>> = None; // materialized verify column
10581        let mut pen_hist_d: Option<CudaSlice<u32>> = None;
10582        let mut pcol_buf: Option<CudaSlice<f32>> = None; // penalized p-column scratch
10583        // MEMRA_SPEC_SETUP_TRACE=1 (diagnostics): per-call wall decomposition of the burst
10584        // SETUP + TAIL segments (the round loop's internals are MEMRA_SPEC_PHASE's job) —
10585        // built to pin the serve per-burst fixed cost (research/spec-serving-20260801).
10586        let setup_trace = std::env::var("MEMRA_SPEC_SETUP_TRACE").as_deref() == Ok("1");
10587        let t_ent = std::time::Instant::now();
10588
10589        // SESSION-AFFINITY TURN CHECKPOINT (lane/session-affinity, 2026-08-05): capture the
10590        // PROMPT-END boundary state so a LATER turn can rewind here and re-prime only its own
10591        // delta instead of the whole conversation. See `SpecCheckpoint` for why this boundary is
10592        // the one that matters (a history-rewriting client mutates what the session GENERATED,
10593        // so the next turn's prompt agrees with this one up to exactly here).
10594        //
10595        // WHERE — AND WHY THIS EXACT LINE. Right after the trunk prime, BEFORE the init feed
10596        // (`decode_step_h(last_token)`) and before round 0: the last instant at which the caches
10597        // hold exactly `base + prompt.len()` rows and nothing generated.
10598        //
10599        // This was WRONG in the first cut of this lane: the capture sat after the draft-KV fill,
10600        // which is also after the init feed, so `cache.pos` was `base + prompt.len() + 1` — the
10601        // boundary included the FIRST GENERATED TOKEN. That token is the first thing inside the
10602        // `<think>` block the client strips, so every later turn's diff diverged exactly one
10603        // token below the checkpoint and affinity declined 100% of the time. Measured on the
10604        // owner regime: "history diverged at 12233 of checkpoint 12234". The off-by-one made the
10605        // whole mechanism inert while looking, from the outside, like a working
10606        // correctness-declines-safely path — hence the decline log carries the offsets.
10607        //
10608        // The full-attn planes are `len`-truncatable so the snapshot copies only the GDN conv/ssm
10609        // state (the reason a spec session could not rewind before). The draft scratch needs no
10610        // copy: rows below the boundary are rewritten by the next turn's own fill.
10611        //
10612        // WHEN: non-empty prime only. An empty-suffix continuation burst adds no prompt boundary
10613        // (its "prompt end" IS the previous checkpoint's, already held), so it keeps the existing
10614        // checkpoint rather than replacing it with a strictly worse one.
10615        //
10616        // FAILURE IS SILENT BY DESIGN: on a VRAM-tight rig the snapshot alloc can fail. That
10617        // costs the NEXT turn its rewind (it re-primes fully, today's behavior) and must never
10618        // fail the burst that is already running — so the error is swallowed, loud only under
10619        // MEMRA_DEBUG_SPEC.
10620        //
10621        // STABLE-BOUNDARY OVERRIDE (lane/frspec-multiturn-cache, 2026-08-21): the prompt-end
10622        // posture above was DISPROVED for the think-posture template class — the prompt's own
10623        // tail is the live generation header (`<|im_start|>assistant\n<think>\n`) that the
10624        // next turn's re-render replaces, so the diff diverged a couple tokens BELOW the
10625        // checkpoint and affinity declined 100% of multi-turn agent traffic (the same class
10626        // the plain tier fixed on 2026-08-09 via `plain_checkpoint_boundary`; the port to the
10627        // spec tier is this lane). When the worker armed `ckpt_at`, the capture happened at
10628        // that stop inside the prime above (`ckpt_early`) and is installed here instead;
10629        // capture-attempted-but-failed clears the slot exactly like the legacy arm.
10630        if let Some(slot) = sess_ckpt_slot {
10631            if let Some(early) = ckpt_early {
10632                if early.is_none() && std::env::var("MEMRA_DEBUG_SPEC").is_ok() {
10633                    eprintln!(
10634                        "[spec] stable-boundary turn checkpoint skipped; \
10635                               next turn re-primes in full"
10636                    );
10637                }
10638                *slot = early;
10639            } else if !continuation {
10640                let pos = cache.pos;
10641                debug_assert_eq!(
10642                    pos,
10643                    base + prompt.len(),
10644                    "turn checkpoint must sit at the prompt end, before the init feed"
10645                );
10646                let anchor: Result<CudaSlice<f32>, Box<dyn std::error::Error>> =
10647                    if let Some(ph) = &prompt_h {
10648                        // hidden of the LAST primed row = the predecessor anchor at this
10649                        // boundary (exactly what a fresh prime of committed[..pos] leaves in
10650                        // last_h, and what the next prime's fill reads for its first row).
10651                        let np = prompt.len();
10652                        e.uninit(n_embd).and_then(|mut a| {
10653                            e.copy_view_into(
10654                                &mut a,
10655                                0,
10656                                &ph.slice((np - 1) * n_embd..np * n_embd),
10657                                n_embd,
10658                            )?;
10659                            Ok(a)
10660                        })
10661                    } else {
10662                        Err("no prompt hiddens".into())
10663                    };
10664                match (cache.snapshot(e), anchor) {
10665                    (Ok(snap), Ok(last_h)) => {
10666                        *slot = Some(SpecCheckpoint { snap, pos, last_h });
10667                    }
10668                    (s, a) => {
10669                        *slot = None; // a stale checkpoint would rewind to the WRONG boundary
10670                        if std::env::var("MEMRA_DEBUG_SPEC").is_ok() {
10671                            let err = s
10672                                .err()
10673                                .map(|e| e.to_string())
10674                                .or_else(|| a.err().map(|e| e.to_string()))
10675                                .unwrap_or_default();
10676                            eprintln!(
10677                                "[spec] turn checkpoint skipped ({err}); \
10678                                       next turn re-primes in full"
10679                            );
10680                        }
10681                    }
10682                }
10683            }
10684        }
10685        // INIT FEED — skipped on a pending carry: last_token (the carried bonus) is NOT in the
10686        // caches and must NOT be fed solo; round 0's batched verify commits it as col 0. Its
10687        // seed/anchor hidden is the carried last_h (copied below); last_pred is dead in the
10688        // pending path (t_pred reads verify col 0 — the accept walk overwrites it).
10689        let mut last_pred = 0u32;
10690        let mut last_col_logits: Option<CudaSlice<f32>> = None;
10691        // CONSTRAINED: the init feed's logits back the (n_acc==0, base==0) masked-argmax
10692        // recompute in the grammar-truncation walk — retained host-side, round 0 only.
10693        let mut init_logits_host: Option<Vec<f32>> = None;
10694        let h_seed0: CudaSlice<f32> = if carried_pending.is_none() {
10695            let (init_logits, h) = self.spec_target_step_h(e, last_token, &mut *cache)?;
10696            last_pred = argmax(&init_logits) as u32;
10697            if constraint.is_some() {
10698                init_logits_host = Some(init_logits.clone());
10699            }
10700            // sampled mode: p-distribution after last_token, for the j==0/base==0 accept test.
10701            if sampled {
10702                last_col_logits = Some(e.htod(&init_logits)?);
10703            }
10704            h
10705        } else {
10706            // predecessor-row anchor: hidden of the last COMMITTED row (the carry contract).
10707            let lh = sess_tail
10708                .as_ref()
10709                .unwrap()
10710                .1
10711                .as_ref()
10712                .expect("pending carry requires last_h");
10713            e.clone_dtod(lh)?
10714        };
10715        let t_init = t_ent.elapsed();
10716        let mut last_col_stats: Option<(f32, f32, f32)> = None;
10717        // PERSISTENT h_seed buffer (allocated BEFORE any graph capture so no captured scratch can
10718        // alias it): every path that updates the round seed copies INTO it — no per-round allocs,
10719        // stable pointer for the graph-draft round-start copy.
10720        let mut h_seed_buf = e.clone_dtod(&h_seed0)?;
10721        // Predecessor-pairing trackers: `fill_prev` = trunk hidden AT the last COMMITTED row (the
10722        // predecessor of the next verify's col 0 — the reference's carried pending-h analogue;
10723        // also the predecessor-row hidden for the round-0 legacy-replay seed). At round 0 that
10724        // row is last_token's own (h_seed0). The chain step-0 seed under the pairing default =
10725        // hidden of the row BEFORE last_token = the prompt's last row at round 0 (h_seed_buf
10726        // overwritten below).
10727        let mut fill_prev = e.clone_dtod(&h_seed0)?;
10728        {
10729            if let Some(ph) = &prompt_h {
10730                let np = prompt.len();
10731                e.copy_view_into(
10732                    &mut h_seed_buf,
10733                    0,
10734                    &ph.slice((np - 1) * n_embd..np * n_embd),
10735                    n_embd,
10736                )?;
10737            } else if continuation
10738                && let Some((_, lh, _, _, _)) = sess_tail.as_ref()
10739                && let Some(lh) = lh.as_ref()
10740            {
10741                e.copy_into(&mut h_seed_buf, 0, lh, n_embd)?;
10742            }
10743        }
10744        // Persistent device prediction slots for the accept walk (max k+1 verify columns).
10745        let mut preds_d = e.alloc_u32_zeroed(k + 2)?;
10746
10747        let debug_spec = std::env::var("MEMRA_DEBUG_SPEC").is_ok();
10748        // MEMRA_SPEC_STATS=1: per-slot accept histogram + draft-length histogram, printed once at
10749        // the end. Metric normalization vs the reference engine: BOTH engines count
10750        // accepted/drafted where the chain stopped at p-min and the sub-threshold token is
10751        // discarded uncounted — per-slot decay + chain-length mix are the extra dimensions.
10752        let spec_stats = std::env::var("MEMRA_SPEC_STATS").is_ok();
10753        let mut st_drafted = vec![0usize; k];
10754        let mut st_accepted = vec![0usize; k];
10755        let mut st_len_hist = vec![0usize; k + 1];
10756        let mut st_full = 0usize;
10757        // P-MIN CONFIDENCE GATE (MEMRA_SPEC_PMIN, the serve script's --spec-draft-p-min mechanism):
10758        // stop the draft chain early when the head's softmax confidence in its own pick drops
10759        // below p_min. Hoisted above the loop: the graph capture bakes the prob kernels iff on.
10760        static PMIN: std::sync::OnceLock<f32> = std::sync::OnceLock::new();
10761        let p_min = *PMIN.get_or_init(|| {
10762            std::env::var("MEMRA_SPEC_PMIN")
10763                .ok()
10764                .and_then(|v| v.parse().ok())
10765                .unwrap_or(0.0)
10766        });
10767        // ZERO-DRAFT ROUNDS (MEMRA_SPEC_PMIN0=1, vendored from llama.cpp's draft gating): let the
10768        // p-min gate apply at j==0 too, so a low-confidence round drafts NOTHING and the verify
10769        // batch is just the pending bonus (m=1 = a plain decode step). llama's 35B win rides
10770        // exactly this — draft acceptance 76% at mean len 2.5 because unpredictable stretches
10771        // never pay draft+verify overhead. Only legal when a pending bonus exists (an empty
10772        // verify batch is not); the j==0 exemption stays for pending-less rounds.
10773        let pmin0 = std::env::var("MEMRA_SPEC_PMIN0")
10774            .map(|v| v == "1")
10775            .unwrap_or(false);
10776
10777        // --- GRAPH DRAFT setup: persistent I/O buffers + ONE capture (2 warmups inside). The
10778        // warmups mutate scratch len_d / pos / tok / seed — all reset at every round start, so the
10779        // only restore needed is the scratch counter. Capture failure (e.g. a non-capturable
10780        // cuBLAS path in an exotic head) falls back to the eager draft chain.
10781        // PER-SESSION PERSISTENCE (2026-08-01): session calls reuse the DraftGraphCtx parked on
10782        // the SpecSession — the capture (2 warmup head forwards + instantiate) ran ONCE at the
10783        // session's first burst, not per burst (measured ~16ms/burst fixed cost on H100 q27,
10784        // research/spec-serving-20260801). Reuse is pointer-exact: the graph bakes the session's
10785        // own scratch KV (never realloc'd), the model's resident embedding, the OnceLock p_min,
10786        // and the g_* buffers carried in the ctx — replay dispatch is identical to a fresh
10787        // capture, so draft tokens are bit-identical (drafts never decide exactness anyway; the
10788        // verify arbitrates). Single-shot calls (sess=None) build a fresh ctx and drop it.
10789        let mut dctx: DraftGraphCtx = match sess_draft_slot.as_mut().and_then(|s| s.take()) {
10790            Some(c) => c,
10791            None => DraftGraphCtx::new(e, n_embd, if sampled { d_vocab } else { 1 })?,
10792        };
10793        // FAIL-SAFE (step-OOM park replay): pre-mark both fallback flags so no capture arm
10794        // below can fire — LOUD once per replayed session through the standard WARN line.
10795        if sess_capture_disabled {
10796            let reason =
10797                "session replayed after a step-OOM park; draft capture disabled (fail-safe)";
10798            let flip = dctx.failed.mark_greedy(reason);
10799            let flip_s = dctx.failed.mark_sampled(reason);
10800            if let Some(line) = flip.or(flip_s) {
10801                eprintln!("{line}");
10802            }
10803        }
10804        // A session that ran greedy bursts first sized g_q/g_perturb at 1; a sampled resume
10805        // needs d_vocab. Realloc is legal exactly while graph_s is None (nothing baked them).
10806        if sampled && dctx.g_q.len() < d_vocab {
10807            dctx.g_q = e.zeros(d_vocab)?;
10808            dctx.g_perturb = e.zeros(d_vocab)?;
10809        }
10810        // DRAFT-SIDE GRAMMAR MASK (lane/draft-mask, 2026-08-04): the drafter samples the
10811        // grammar's legal set, so proposals are legal BY CONSTRUCTION and the verify-side
10812        // truncation (the correctness backstop) stops cutting every tight-schema round.
10813        // The mask is one node inside the captured draft chain — presence is a CAPTURE-TIME
10814        // shape, so a parked graph of the other shape is dropped and recaptured.
10815        let dmask_on = constraint
10816            .as_deref()
10817            .is_some_and(|c| c.draft_mask_enabled());
10818        let dmask_words = if dmask_on { d_vocab.div_ceil(32) } else { 0 };
10819        if dmask_on && dctx.g_dmask.len() < dmask_words {
10820            dctx.g_dmask = e.alloc_u32_zeroed(dmask_words)?;
10821            dctx.graph = None; // the old capture baked the old (or no) mask pointer
10822            dctx.chain = None; // chain last-row graphs bake the same pointer
10823            dctx.failed.clear_greedy();
10824            dctx.keeper.clear();
10825        }
10826        if (dctx.graph.is_some() || dctx.chain.is_some()) && dctx.graph_masked != dmask_on {
10827            dctx.graph = None;
10828            dctx.chain = None;
10829            dctx.failed.clear_greedy();
10830            dctx.keeper.clear();
10831        }
10832        // MULTI-HEAD CHAIN mode (mtp_extra non-empty — step37's 3-head shipping shape): the
10833        // step-modulo prefix-replay chain captures PER-HEAD single-row graphs
10834        // (`DraftChainGraphs`) instead of the one self-feeding graph below; the single-head
10835        // capture arms are untouched and unreachable in this mode (the launch arms branch the
10836        // same way). This removes the historical `mtp_extra.is_empty()` capture exclusion —
10837        // and with it the silent no-attempt hole: a chain capture that FAILS now trips the
10838        // same LOUD draft-graph WARN as a single-head failure.
10839        let chain_mode = !self.mtp_extra.is_empty();
10840        // ---- PRE-CAPTURE VRAM RESERVE CHECK + PER-SESSION DRAFT-STATE MEASUREMENT ----
10841        // (lane/step37-vram-admission-20260830). `cap_eff0` opens the measurement bracket:
10842        // when any capture succeeds in THIS call, the effective-free delta across the whole
10843        // capture section is recorded as the model's per-session draft-state high-water
10844        // (admission charges it per spec-capable session — this state was charged at ZERO
10845        // before the lane). The reserve check runs BEFORE any capture arm can allocate: a
10846        // refused capture trips the same LOUD once-per-flip WARN class as a failed one, but
10847        // with the card's headroom still intact (the owner's single-session OOM was a capture
10848        // attempt walking the card to the edge and stranding the eager fallback at 5 MiB free).
10849        let cap_eff0 = e
10850            .ctx()
10851            .mem_get_info()
10852            .ok()
10853            .map(|(f, _)| f.saturating_add(e.pool_cached_bytes()));
10854        // Peak instrument for the same bracket: the CAPTURE-TIME peak (warmup transients +
10855        // instantiate scratch, alive together) dwarfs the parked delta — measured on the
10856        // owner shape: a capture whose PARKED state reads ~2.6GB walked a ~7GB-free card to
10857        // OOM mid-capture. Reset the pool watermark here; read it at bracket end.
10858        let _ = e.pool_high_water_reset();
10859        let cap_used0 = e.pool_reserved_used().1;
10860        let mut captured_now = false;
10861        let mut capture_oom_entry_eff: Option<usize> = None;
10862        let capture_need = {
10863            let observed = self.draft_session_admission_bytes();
10864            if observed > 0 {
10865                observed
10866            } else {
10867                draft_capture_bootstrap_estimate(
10868                    if chain_mode { self.mtp_head_count() } else { 1 },
10869                    k,
10870                    d_vocab,
10871                    n_embd,
10872                )
10873            }
10874        };
10875        if spec_capture_gate_on()
10876            && graph_draft
10877            && !sampled
10878            && !dctx.failed.greedy_failed()
10879            && ((chain_mode && dctx.chain.is_none() && mtp_chain_graph_on())
10880                || (!chain_mode && dctx.graph.is_none()))
10881            && let Some(reason) = capture_headroom_refusal(e, capture_need)
10882            && let Some(line) = dctx.failed.mark_greedy(&reason)
10883        {
10884            eprintln!("{line}");
10885        }
10886        if graph_draft
10887            && !sampled
10888            && chain_mode
10889            && dctx.chain.is_none()
10890            && !dctx.failed.greedy_failed()
10891        {
10892            if mtp_chain_graph_on() {
10893                let heads_n = self.mtp_head_count();
10894                let DraftGraphCtx {
10895                    g_tok,
10896                    g_pos,
10897                    g_seed,
10898                    g_p,
10899                    g_dmask,
10900                    ..
10901                } = &mut dctx;
10902                if dmask_on {
10903                    e.htod_u32_into(g_dmask, &vec![u32::MAX; dmask_words])?;
10904                }
10905                let g_dmask_ro: &CudaSlice<u32> = &*g_dmask;
10906                let with_prob = p_min > 0.0;
10907                // CAPTURE-RETAIN (#68 fix): one keeper for the whole chain — every graph's
10908                // warmup transients stay pinned as long as any of them replays.
10909                let cap_res = (|| -> Result<DraftChainGraphs, Box<dyn std::error::Error>> {
10910                    // dcw door: same warmup headroom pre-arm as the single-head capture
10911                    // below — every plane, because each head's capture warmups append on
10912                    // its OWN plane. INSIDE the fallible closure (vram-admission lane): an
10913                    // OOM here used to `?` out of the whole burst as a step error; now it
10914                    // is a capture failure — LOUD WARN, eager chain serves.
10915                    if step35_draft_dcw_on() {
10916                        scratch.ensure_dcw_headroom(e, k + 2)?;
10917                    }
10918                    let mut interior = Vec::with_capacity(heads_n);
10919                    let mut last = Vec::with_capacity(heads_n);
10920                    let mut keeper: Vec<Box<dyn std::any::Any + Send>> = Vec::new();
10921                    for hi in 0..heads_n {
10922                        let head = self.mtp_head_at(hi);
10923                        // interior row: KV append + carrier only (`with_head=false` — the
10924                        // eager chain discards interior logits too, so this is the same
10925                        // consumed-byte program minus the dead full-vocab head matmul).
10926                        let (g, keep) = e.capture_graph_retained(|e| {
10927                            self.mtp_head_forward_cap(
10928                                e,
10929                                head,
10930                                g_tok,
10931                                g_pos,
10932                                g_seed,
10933                                g_p,
10934                                &mut *scratch,
10935                                hi,
10936                                false,
10937                                false,
10938                                embd_gpu.expect("graph draft requires resident embedding"),
10939                                embd_qt,
10940                                embd_rb,
10941                                d_vocab,
10942                                None,
10943                                None,
10944                                None,
10945                            )
10946                        })?;
10947                        // the warmups appended rows on plane hi; rewind before the next
10948                        // capture so successive warmups never outrun the pre-armed headroom.
10949                        scratch.set_plane_len(e, hi, base)?;
10950                        interior.push(g);
10951                        keeper.extend(keep);
10952                        // last row: head matmul + greedy argmax tail (+ p when the policy
10953                        // reads it, + the grammar-mask node when constrained).
10954                        let (g2, keep2) = e.capture_graph_retained(|e| {
10955                            self.mtp_head_forward_cap(
10956                                e,
10957                                head,
10958                                g_tok,
10959                                g_pos,
10960                                g_seed,
10961                                g_p,
10962                                &mut *scratch,
10963                                hi,
10964                                with_prob,
10965                                true,
10966                                embd_gpu.expect("graph draft requires resident embedding"),
10967                                embd_qt,
10968                                embd_rb,
10969                                d_vocab,
10970                                None,
10971                                None,
10972                                if dmask_on {
10973                                    Some((g_dmask_ro, dmask_words))
10974                                } else {
10975                                    None
10976                                },
10977                            )
10978                        })?;
10979                        scratch.set_plane_len(e, hi, base)?;
10980                        last.push(g2);
10981                        keeper.extend(keep2);
10982                    }
10983                    Ok(DraftChainGraphs {
10984                        interior,
10985                        last,
10986                        _keeper: keeper,
10987                    })
10988                })();
10989                match cap_res {
10990                    Ok(cg) => {
10991                        scratch.set_len(e, base)?;
10992                        // POSITIVE engagement receipt (the 3a lesson: a WARN-free boot is
10993                        // NOT evidence of capture — the captured state must name itself).
10994                        eprintln!(
10995                            "[mtp-chain-graph] captured mode=greedy heads={heads_n} \
10996                             interior={heads_n} last={heads_n} masked={}",
10997                            dmask_on as u8
10998                        );
10999                        dctx.chain = Some(cg);
11000                        dctx.graph_masked = dmask_on;
11001                        captured_now = true;
11002                    }
11003                    Err(err) => {
11004                        scratch.set_len(e, base)?;
11005                        // LOUD flip (audit Q2): a dropped draft graph is a coverage loss,
11006                        // never silent — now including the multi-head shipping shape.
11007                        // OOM RECOVERY (vram-admission lane): a failed attempt's freed
11008                        // transients sit CACHED in the async pool where the driver cannot
11009                        // see them; trim them back so the eager fallback (and any driver-
11010                        // side allocation) actually has the headroom the free suggests.
11011                        let mut reason = err.to_string();
11012                        if capture_err_is_oom(&reason) {
11013                            capture_oom_entry_eff = capture_oom_entry_eff.max(cap_eff0);
11014                            let trimmed = e.pool_trim_to_zero();
11015                            if trimmed > 0 {
11016                                reason.push_str(&format!(
11017                                    "; pool trimmed {}MB back to the driver",
11018                                    trimmed / (1 << 20)
11019                                ));
11020                            }
11021                        }
11022                        if let Some(line) = dctx.failed.mark_greedy(&reason) {
11023                            eprintln!("{line}");
11024                        }
11025                    }
11026                }
11027            } else {
11028                // Disarmed by MEMRA_MTP_CHAIN_GRAPH=0: say so once per process — the OFF arm
11029                // must be attributable in a boot log, never inferable from silence.
11030                static NOTE: std::sync::Once = std::sync::Once::new();
11031                NOTE.call_once(|| {
11032                    eprintln!(
11033                        "[spec] multi-head draft-chain capture disarmed \
11034                         (MEMRA_MTP_CHAIN_GRAPH=0); eager chain serves this shape"
11035                    );
11036                });
11037            }
11038        }
11039        if graph_draft
11040            && !sampled
11041            && !chain_mode
11042            && dctx.graph.is_none()
11043            && !dctx.failed.greedy_failed()
11044        {
11045            let DraftGraphCtx {
11046                g_tok,
11047                g_pos,
11048                g_seed,
11049                g_p,
11050                g_dmask,
11051                ..
11052            } = &mut dctx;
11053            // capture-time contents: ALL-ONES (ban nothing). A replay only ever runs after the
11054            // host uploads the position's real words, so the warmups stay grammar-free.
11055            if dmask_on {
11056                e.htod_u32_into(g_dmask, &vec![u32::MAX; dmask_words])?;
11057            }
11058            let g_dmask_ro: &CudaSlice<u32> = &*g_dmask;
11059            // CAPTURE-RETAIN (#68 fix): the warmup transients' pool addresses are baked into the
11060            // captured graph; the keeper pins them for the graph's lifetime. capture_graph (non-
11061            // retained) freed them at exit — safe for one-shot generate_spec (nothing else touches
11062            // the pool between replays) but WRONG for sessions: burst-boundary prime/fill/commit
11063            // passes (and, in serve, other sessions) recycle those addresses and the replay then
11064            // clobbers live buffers — the ST serve-spec corruption (research/serve-st-20260803).
11065            let cap_res = (|| {
11066                // dcw door: the capture warmups append device-counter rows the capture body
11067                // cannot rebase for; pre-arm ring headroom host-side (no-op on flat planes /
11068                // room-enough rings, and the door-off path is untouched). INSIDE the fallible
11069                // closure (vram-admission lane): an OOM here is a capture failure, not a
11070                // burst-killing step error.
11071                if step35_draft_dcw_on() {
11072                    scratch.ensure_dcw_headroom(e, k + 2)?;
11073                }
11074                e.capture_graph_retained(|e| {
11075                    self.mtp_head_forward_cap(
11076                        e,
11077                        mtp,
11078                        g_tok,
11079                        g_pos,
11080                        g_seed,
11081                        g_p,
11082                        &mut *scratch,
11083                        0,
11084                        p_min > 0.0,
11085                        true,
11086                        embd_gpu.expect("graph draft requires resident embedding"),
11087                        embd_qt,
11088                        embd_rb,
11089                        d_vocab,
11090                        None,
11091                        None,
11092                        if dmask_on {
11093                            Some((g_dmask_ro, dmask_words))
11094                        } else {
11095                            None
11096                        },
11097                    )
11098                })
11099            })();
11100            match cap_res {
11101                Ok((g, keep)) => {
11102                    scratch.set_len(e, base)?;
11103                    dctx.graph = Some(g);
11104                    dctx.graph_masked = dmask_on;
11105                    dctx.keeper = keep;
11106                    captured_now = true;
11107                }
11108                Err(err) => {
11109                    scratch.set_len(e, base)?;
11110                    // LOUD flip (audit Q2): a dropped draft graph is a coverage loss, never
11111                    // silent. Once per flip — mark returns None on an already-failed ctx.
11112                    let mut reason = err.to_string();
11113                    if capture_err_is_oom(&reason) {
11114                        capture_oom_entry_eff = capture_oom_entry_eff.max(cap_eff0);
11115                        let trimmed = e.pool_trim_to_zero();
11116                        if trimmed > 0 {
11117                            reason.push_str(&format!(
11118                                "; pool trimmed {}MB back to the driver",
11119                                trimmed / (1 << 20)
11120                            ));
11121                        }
11122                    }
11123                    if let Some(line) = dctx.failed.mark_greedy(&reason) {
11124                        eprintln!("{line}");
11125                    }
11126                }
11127            }
11128        }
11129        // --- SAMPLED GRAPH DRAFT setup (step 3 of the sampled-spec arc): a SECOND capture, own
11130        // graph object, built only when sampled && graph-eligible — the greedy capture above is
11131        // untouched (and skipped when sampled: its graph would never be launched). Same head
11132        // forward, but the in-graph argmax reads GUMBEL-PERTURBED logits; the Philox event
11133        // counter lives in the persistent device g_ctr (bumped in-graph, host-seeded from sctr
11134        // once per round); the raw head logits land in the persistent g_q for the host's
11135        // per-replay async D2D into the round's q slot (q_slots, K x d_vocab, allocated once).
11136        // seed/temp are capture-time constants — baked into graph_s, so a pool-resumed request
11137        // with a different (seed, temp, k) drops the parked sampled graph and recaptures.
11138        // COST OF THE FRESH-SEED SERVE DEFAULT (dogfood F4, 2026-08-04): omitting `seed` on a
11139        // serve request now draws fresh per-request entropy (it used to default to a pinned 0),
11140        // so a seed-omitting request that RESUMES a parked spec session finds an s_key baked
11141        // with the PREVIOUS request's seed and pays one recapture. Bounded, and it does not
11142        // reopen the ~16ms/burst regression the persistent ctx exists to fix: a session's seed
11143        // is fixed for its whole lifetime (worker.rs reads s.sampler.seed() per burst), so
11144        // this compare misses at most ONCE per resumed request — the first burst recaptures
11145        // and every later burst in that request replays. A client that wants the parked graph
11146        // AND reproducibility supplies an explicit `seed`, honored exactly, which keeps s_key
11147        // stable across its whole conversation.
11148        // COMPOSITION RULE (fspec x gsd merge): the in-graph chain samples from the RAW
11149        // softmax — it can hold neither per-row filter stats nor the varying penalty history.
11150        // The sampled graph therefore engages only in the PURE-TEMP regime; filters/penalties
11151        // force the eager draft (which computes stats/penalties per row).
11152        // KEY THE WHOLE REGIME, not just the baked constants (lane/graph-s-key-exactness-
11153        // 20260819). `s_key` used to be `(seed, temp, k)`; the filters and penalties were left
11154        // out, so a filtered request resuming a session that parked a PURE-TEMP graph kept it —
11155        // and the launch site never re-asked `pure_temp`. See [`SampledGraphKey`] for what that
11156        // costs (an unconditional accept of out-of-head draft tokens, i.e. an exactness bug on
11157        // the request shape the vendor-default flip makes the majority).
11158        let s_key = SampledGraphKey::new(sp_seed, sp_temp, k, sp.top_k, sp.top_p, sp.min_p, pen_on);
11159        let pure_temp = s_key.pure_temp();
11160        // The regime the sampled graph may be captured/launched in: pure-temp always;
11161        // truncation-filtered when the filtered-capture door is on (the filter runs
11162        // IN-GRAPH — lane/step37-draft-graph-serving-20260830); penalties never.
11163        let s_capturable = s_key.graph_capturable();
11164        if sampled && dctx.s_key.is_some_and(|old| old != s_key) {
11165            dctx.graph_s = None;
11166            dctx.chain_s = None;
11167            dctx.failed.clear_sampled();
11168            dctx.s_key = None;
11169            dctx.q_slots.clear();
11170            dctx.keeper_s.clear();
11171        }
11172        // PRE-CAPTURE VRAM RESERVE CHECK, sampled arms (vram-admission lane): same contract
11173        // as the greedy check above — refuse BEFORE allocating, LOUD once, eager serves.
11174        if spec_capture_gate_on()
11175            && graph_draft
11176            && sampled
11177            && s_capturable
11178            && !dctx.failed.sampled_failed()
11179            && ((chain_mode && dctx.chain_s.is_none() && mtp_chain_graph_on())
11180                || (!chain_mode && dctx.graph_s.is_none()))
11181            && let Some(reason) = capture_headroom_refusal(e, capture_need)
11182            && let Some(line) = dctx.failed.mark_sampled(&reason)
11183        {
11184            eprintln!("{line}");
11185        }
11186        // FILTERED capture nodes need q slots sized d_vocab AND the stat slots; the pure-temp
11187        // body leaves g_th/g_z/g_mx untouched (they exist from ctx creation either way).
11188        if graph_draft
11189            && sampled
11190            && s_capturable
11191            && chain_mode
11192            && dctx.chain_s.is_none()
11193            && !dctx.failed.sampled_failed()
11194        {
11195            if mtp_chain_graph_on() {
11196                let heads_n = self.mtp_head_count();
11197                let filtered = s_key.filtered();
11198                let DraftGraphCtx {
11199                    g_tok,
11200                    g_pos,
11201                    g_seed,
11202                    g_p,
11203                    g_ctr,
11204                    g_perturb,
11205                    g_q,
11206                    g_rows0,
11207                    g_th,
11208                    g_z,
11209                    g_mx,
11210                    ..
11211                } = &mut dctx;
11212                let with_prob = p_min > 0.0;
11213                let cap_res = (|| -> Result<DraftChainGraphs, Box<dyn std::error::Error>> {
11214                    // dcw pre-arm INSIDE the fallible closure (vram-admission lane): an OOM
11215                    // here is a capture failure with the LOUD WARN, never a step error.
11216                    if step35_draft_dcw_on() {
11217                        scratch.ensure_dcw_headroom(e, k + 2)?;
11218                    }
11219                    let mut interior = Vec::with_capacity(heads_n);
11220                    let mut last = Vec::with_capacity(heads_n);
11221                    let mut keeper: Vec<Box<dyn std::any::Any + Send>> = Vec::new();
11222                    for hi in 0..heads_n {
11223                        let head = self.mtp_head_at(hi);
11224                        // interior row: no head, no draw — shared shape with the greedy
11225                        // chain's interior, captured per mode for keeper-lifetime hygiene.
11226                        let (g, keep) = e.capture_graph_retained(|e| {
11227                            self.mtp_head_forward_cap(
11228                                e,
11229                                head,
11230                                g_tok,
11231                                g_pos,
11232                                g_seed,
11233                                g_p,
11234                                &mut *scratch,
11235                                hi,
11236                                false,
11237                                false,
11238                                embd_gpu.expect("graph draft requires resident embedding"),
11239                                embd_qt,
11240                                embd_rb,
11241                                d_vocab,
11242                                None,
11243                                None,
11244                                None,
11245                            )
11246                        })?;
11247                        scratch.set_plane_len(e, hi, base)?;
11248                        interior.push(g);
11249                        keeper.extend(keep);
11250                        // last row: head matmul + the in-graph categorical draw (filtered
11251                        // nodes when the request carries filters).
11252                        let (g2, keep2) = e.capture_graph_retained(|e| {
11253                            self.mtp_head_forward_cap(
11254                                e,
11255                                head,
11256                                g_tok,
11257                                g_pos,
11258                                g_seed,
11259                                g_p,
11260                                &mut *scratch,
11261                                hi,
11262                                with_prob,
11263                                true,
11264                                embd_gpu.expect("graph draft requires resident embedding"),
11265                                embd_qt,
11266                                embd_rb,
11267                                d_vocab,
11268                                Some(SampledCapArgs {
11269                                    ctr: &mut *g_ctr,
11270                                    perturb: &mut *g_perturb,
11271                                    q_out: &mut *g_q,
11272                                    seed: sp_seed,
11273                                    temp: sp_temp,
11274                                    filt: if filtered {
11275                                        Some(SampledCapFilter {
11276                                            rows0: &*g_rows0,
11277                                            th: &mut *g_th,
11278                                            z: &mut *g_z,
11279                                            mx: &mut *g_mx,
11280                                            top_k: sp.top_k,
11281                                            top_p: sp.top_p,
11282                                            min_p: sp.min_p,
11283                                        })
11284                                    } else {
11285                                        None
11286                                    },
11287                                }),
11288                                None,
11289                                None, // constrained spec is greedy-only
11290                            )
11291                        })?;
11292                        scratch.set_plane_len(e, hi, base)?;
11293                        last.push(g2);
11294                        keeper.extend(keep2);
11295                    }
11296                    Ok(DraftChainGraphs {
11297                        interior,
11298                        last,
11299                        _keeper: keeper,
11300                    })
11301                })();
11302                match cap_res {
11303                    Ok(cg) => {
11304                        scratch.set_len(e, base)?;
11305                        // NO STRANDED PARTIAL STATE (vram-admission lane): the q-slot allocs
11306                        // after a successful capture are themselves fallible on a tight card.
11307                        // A mid-loop failure used to `?` out as a step error, leaving orphan
11308                        // slots parked on the ctx (wrong count, stale contents) for the next
11309                        // capture attempt to stack onto. Allocate all-or-nothing: on failure
11310                        // drop the fresh graphs AND the partial slots, mark the LOUD fallback.
11311                        dctx.q_slots.clear();
11312                        let slots = (0..k)
11313                            .map(|_| e.zeros(d_vocab))
11314                            .collect::<Result<Vec<_>, _>>();
11315                        match slots {
11316                            Ok(slots) => {
11317                                dctx.q_slots = slots;
11318                                eprintln!(
11319                                    "[mtp-chain-graph] captured mode=sampled heads={heads_n} \
11320                                     interior={heads_n} last={heads_n} filtered={} key={s_key:?}",
11321                                    s_key.filtered() as u8
11322                                );
11323                                dctx.chain_s = Some(cg);
11324                                dctx.s_key = Some(s_key);
11325                                captured_now = true;
11326                            }
11327                            Err(err) => {
11328                                drop(cg);
11329                                dctx.q_slots.clear();
11330                                let mut reason = format!("q-slot alloc failed: {err}");
11331                                if capture_err_is_oom(&reason) {
11332                                    capture_oom_entry_eff = capture_oom_entry_eff.max(cap_eff0);
11333                                    let trimmed = e.pool_trim_to_zero();
11334                                    if trimmed > 0 {
11335                                        reason.push_str(&format!(
11336                                            "; pool trimmed {}MB back to the driver",
11337                                            trimmed / (1 << 20)
11338                                        ));
11339                                    }
11340                                }
11341                                if let Some(line) = dctx.failed.mark_sampled(&reason) {
11342                                    eprintln!("{line}");
11343                                }
11344                            }
11345                        }
11346                    }
11347                    Err(err) => {
11348                        scratch.set_len(e, base)?;
11349                        let mut reason = err.to_string();
11350                        if capture_err_is_oom(&reason) {
11351                            capture_oom_entry_eff = capture_oom_entry_eff.max(cap_eff0);
11352                            let trimmed = e.pool_trim_to_zero();
11353                            if trimmed > 0 {
11354                                reason.push_str(&format!(
11355                                    "; pool trimmed {}MB back to the driver",
11356                                    trimmed / (1 << 20)
11357                                ));
11358                            }
11359                        }
11360                        if let Some(line) = dctx.failed.mark_sampled(&reason) {
11361                            eprintln!("{line}");
11362                        }
11363                    }
11364                }
11365            } else {
11366                static NOTE_S: std::sync::Once = std::sync::Once::new();
11367                NOTE_S.call_once(|| {
11368                    eprintln!(
11369                        "[spec] multi-head draft-chain capture disarmed \
11370                         (MEMRA_MTP_CHAIN_GRAPH=0); eager chain serves this shape"
11371                    );
11372                });
11373            }
11374        }
11375        if graph_draft
11376            && sampled
11377            && s_capturable
11378            && !chain_mode
11379            && dctx.graph_s.is_none()
11380            && !dctx.failed.sampled_failed()
11381        {
11382            let filtered = s_key.filtered();
11383            let DraftGraphCtx {
11384                g_tok,
11385                g_pos,
11386                g_seed,
11387                g_p,
11388                g_ctr,
11389                g_perturb,
11390                g_q,
11391                g_rows0,
11392                g_th,
11393                g_z,
11394                g_mx,
11395                ..
11396            } = &mut dctx;
11397            // CAPTURE-RETAIN (#68 fix): same keeper contract as the greedy capture above.
11398            let cap_res = (|| {
11399                // dcw pre-arm INSIDE the fallible closure (vram-admission lane): an OOM
11400                // here is a capture failure with the LOUD WARN, never a step error.
11401                if step35_draft_dcw_on() {
11402                    scratch.ensure_dcw_headroom(e, k + 2)?;
11403                }
11404                e.capture_graph_retained(|e| {
11405                    self.mtp_head_forward_cap(
11406                        e,
11407                        mtp,
11408                        g_tok,
11409                        g_pos,
11410                        g_seed,
11411                        g_p,
11412                        &mut *scratch,
11413                        0,
11414                        p_min > 0.0,
11415                        true,
11416                        embd_gpu.expect("graph draft requires resident embedding"),
11417                        embd_qt,
11418                        embd_rb,
11419                        d_vocab,
11420                        Some(SampledCapArgs {
11421                            ctr: &mut *g_ctr,
11422                            perturb: &mut *g_perturb,
11423                            q_out: &mut *g_q,
11424                            seed: sp_seed,
11425                            temp: sp_temp,
11426                            filt: if filtered {
11427                                Some(SampledCapFilter {
11428                                    rows0: &*g_rows0,
11429                                    th: &mut *g_th,
11430                                    z: &mut *g_z,
11431                                    mx: &mut *g_mx,
11432                                    top_k: sp.top_k,
11433                                    top_p: sp.top_p,
11434                                    min_p: sp.min_p,
11435                                })
11436                            } else {
11437                                None
11438                            },
11439                        }),
11440                        None,
11441                        None, // constrained spec is greedy-only — sampled never carries a hook
11442                    )
11443                })
11444            })();
11445            match cap_res {
11446                Ok((g, keep)) => {
11447                    scratch.set_len(e, base)?;
11448                    // NO STRANDED PARTIAL STATE: all-or-nothing q slots, same contract as
11449                    // the chain arm above.
11450                    dctx.q_slots.clear();
11451                    let slots = (0..k)
11452                        .map(|_| e.zeros(d_vocab))
11453                        .collect::<Result<Vec<_>, _>>();
11454                    match slots {
11455                        Ok(slots) => {
11456                            dctx.q_slots = slots;
11457                            dctx.graph_s = Some(g);
11458                            dctx.s_key = Some(s_key);
11459                            dctx.keeper_s = keep;
11460                            captured_now = true;
11461                        }
11462                        Err(err) => {
11463                            drop(g);
11464                            drop(keep);
11465                            dctx.q_slots.clear();
11466                            let mut reason = format!("q-slot alloc failed: {err}");
11467                            if capture_err_is_oom(&reason) {
11468                                capture_oom_entry_eff = capture_oom_entry_eff.max(cap_eff0);
11469                                let trimmed = e.pool_trim_to_zero();
11470                                if trimmed > 0 {
11471                                    reason.push_str(&format!(
11472                                        "; pool trimmed {}MB back to the driver",
11473                                        trimmed / (1 << 20)
11474                                    ));
11475                                }
11476                            }
11477                            if let Some(line) = dctx.failed.mark_sampled(&reason) {
11478                                eprintln!("{line}");
11479                            }
11480                        }
11481                    }
11482                }
11483                Err(err) => {
11484                    scratch.set_len(e, base)?;
11485                    // LOUD flip (audit Q2): same contract as the greedy capture above.
11486                    let mut reason = err.to_string();
11487                    if capture_err_is_oom(&reason) {
11488                        capture_oom_entry_eff = capture_oom_entry_eff.max(cap_eff0);
11489                        let trimmed = e.pool_trim_to_zero();
11490                        if trimmed > 0 {
11491                            reason.push_str(&format!(
11492                                "; pool trimmed {}MB back to the driver",
11493                                trimmed / (1 << 20)
11494                            ));
11495                        }
11496                    }
11497                    if let Some(line) = dctx.failed.mark_sampled(&reason) {
11498                        eprintln!("{line}");
11499                    }
11500                }
11501            }
11502        }
11503        // ---- PER-SESSION DRAFT-STATE MEASUREMENT bracket end (vram-admission lane): when a
11504        // capture landed in THIS call, the effective-free delta across the capture section is
11505        // this session's parked draft-graph state (keepers + q slots + instantiated graphs'
11506        // backing). Recorded as a model-owned high-water; admission charges it per
11507        // spec-capable session (see `draft_session_admission_bytes`).
11508        if captured_now
11509            && let Some(eff0) = cap_eff0
11510            && let Ok((f1, _)) = e.ctx().mem_get_info()
11511        {
11512            let eff1 = f1.saturating_add(e.pool_cached_bytes());
11513            let parked_delta = eff0.saturating_sub(eff1);
11514            let (_res_high, used_high) = e.pool_high_water_reset();
11515            let peak_delta = used_high.saturating_sub(cap_used0);
11516            let observed = parked_delta.max(peak_delta);
11517            if observed > 0
11518                && let Some(hw) = self.record_draft_state_bytes(observed)
11519            {
11520                eprintln!(
11521                    "[spec] draft-session state high-water: {}MB (max of parked delta {}MB \
11522                     and capture-time pool peak {}MB; charged per spec admission and gating \
11523                     future captures)",
11524                    hw / (1 << 20),
11525                    parked_delta / (1 << 20),
11526                    peak_delta / (1 << 20),
11527                );
11528            }
11529        }
11530        // FAILURE IS AN OBSERVATION TOO: a capture that OOM'd at entry-effective E proved
11531        // the capture-time peak exceeds E. Feed E into the gauge so every future gate
11532        // refuses at or below the headroom that just failed (self-healing even when the
11533        // boot probe is disarmed and the bootstrap estimate was blind).
11534        if let Some(entry_eff) = capture_oom_entry_eff
11535            && let Some(hw) = self.record_draft_state_bytes(entry_eff)
11536        {
11537            eprintln!(
11538                "[spec] draft-session capture appetite floor raised to {}MB: a capture \
11539                 attempt OOM'd with that much effective free (failure-observed bound)",
11540                hw / (1 << 20)
11541            );
11542        }
11543        // ---- EXACTNESS GUARD, the enforceable half (lane/graph-s-key-exactness-20260819,
11544        // widened by lane/step37-draft-graph-serving-20260830) ----
11545        // With the filters and penalties in `s_key`, a graph that SURVIVED the drop above was
11546        // captured under THIS request's exact regime, and capture requires `graph_capturable`
11547        // (pure-temp, or filtered with the in-graph filter nodes; never penalties) — so a
11548        // parked graph implies both. That implication is the whole exactness argument for the
11549        // graph arm, so it is asserted here rather than assumed: a future change that widens
11550        // the capture condition, narrows the key, or copies a `DraftGraphCtx` across regimes
11551        // fails LOUDLY at this line instead of silently drafting from a distribution the
11552        // verify never reconstructs. Release builds refuse the graph (drop it, draft eager)
11553        // rather than launching it; the launch site re-tests the regime independently.
11554        if sampled
11555            && (dctx.graph_s.is_some() || dctx.chain_s.is_some())
11556            && (!s_capturable || dctx.s_key != Some(s_key))
11557        {
11558            debug_assert!(
11559                false,
11560                "sampled draft graph parked under {:?} survived into a request outside its \
11561                 capture regime (top_k={} top_p={} min_p={} pen_on={} capturable={}): the \
11562                 in-graph draw and the verify's accept test would see different distributions",
11563                dctx.s_key, sp.top_k, sp.top_p, sp.min_p, pen_on, s_capturable,
11564            );
11565            eprintln!(
11566                "[spec] BUG: dropping a parked sampled draft graph that outlived its capture \
11567                 regime (s_key={:?}, request top_k={} top_p={} min_p={} pen_on={} \
11568                 capturable={}); drafting EAGER — the key must carry every field that shapes q",
11569                dctx.s_key, sp.top_k, sp.top_p, sp.min_p, pen_on, s_capturable,
11570            );
11571            dctx.graph_s = None;
11572            dctx.chain_s = None;
11573            dctx.s_key = None;
11574            dctx.q_slots.clear();
11575            dctx.keeper_s.clear();
11576        }
11577        // SKEY PROBE (MEMRA_SKEY_PROBE=1): the burst-entry facts the reachability question turns
11578        // on — is this request sampled, is it in a regime the sampled graph is legal in, and is
11579        // a graph PARKED from an earlier request of the same session? The launch arms below
11580        // print which chain actually ran, so the probe never restates the condition.
11581        if skey_probe() {
11582            eprintln!(
11583                "[skey] burst sampled={} pure_temp={} capturable={} temp={} top_k={} top_p={} \
11584                 min_p={} pen_on={} k={} graph_draft={} graph_s_parked={} chain_s_parked={} \
11585                 s_key_parked={:?}",
11586                sampled as u8,
11587                pure_temp as u8,
11588                s_capturable as u8,
11589                sp_temp,
11590                sp.top_k,
11591                sp.top_p,
11592                sp.min_p,
11593                pen_on as u8,
11594                k,
11595                graph_draft as u8,
11596                dctx.graph_s.is_some() as u8,
11597                dctx.chain_s.is_some() as u8,
11598                dctx.s_key,
11599            );
11600        }
11601        let t_cap = t_ent.elapsed();
11602        // PERSISTENT DRAFT KV: fill the MTP block's K/V for every prompt position from the exact
11603        // trunk hiddens collected during prime — ONE batched K/V-only pass (overwrites any
11604        // capture-warmup garbage; capture left len at 0). last_token (the init feed) needs no
11605        // fill: the first chain step processes it and appends its entry at slot prompt.len().
11606        if let Some(ph) = &prompt_h {
11607            // SESSION: rows [0..base) are the previous turns' exact fills (refresh overwrote them
11608            // with true verify hiddens) — truncate any draft overhang, fill ONLY the suffix at
11609            // global positions [base..base+tp). Fresh call: base==0, identical to before.
11610            scratch.set_len(e, base)?;
11611            // CHUNKED FILL (long-ctx OOM fix, 2026-07-05): mtp_kv_fill's transients scale with its
11612            // T (concat = T*2*n_embd*4B — 1.5GB at 40k) and its concat loop is 2*T launches. The
11613            // fill is a pure sequential append, so chunking is exact: each chunk appends its rows
11614            // at pos0=base+start with the identical per-row math. Same knob as the trunk prime.
11615            let tp = prompt.len();
11616            let fill_chunk: usize = if crate::cache::swa_ring_on() {
11617                crate::hybrid_forward::prime_chunk_tokens(tp, self.layers.len())
11618            } else {
11619                // Preserve the flag-OFF schedule byte-for-byte, including the legacy zero value
11620                // meaning one monolithic fill.
11621                std::env::var("MEMRA_PRIME_CHUNK")
11622                    .ok()
11623                    .and_then(|v| v.parse().ok())
11624                    .unwrap_or(4096)
11625            };
11626            let fill_chunk = if fill_chunk == 0 { tp } else { fill_chunk };
11627            // CUDA launch wall (same class as the trunk prime's PRIME_CHUNK_LAUNCH_CAP):
11628            // a fill call's matmuls can land on the grid.y=m dp4a family, and grid.y caps
11629            // at 65,535. This loop has no tail fold, so the raw limit is exact:
11630            // tp <= 65,535 keeps the legacy schedule (monolithic included) byte-for-byte,
11631            // and larger fills — unreachable before the trunk prime's own cap fix — chunk.
11632            let fill_chunk = fill_chunk.min(crate::hybrid_forward::CUDA_GRID_YZ_MAX);
11633            let mut start = 0usize;
11634            while start < tp {
11635                let end = (start + fill_chunk).min(tp);
11636                let tc = end - start;
11637                {
11638                    // PREDECESSOR pairing: row i gets h[i-1]; global row 0 a zeros row (the
11639                    // reference engine's initial pending-h is zeroed too); a session turn's row 0
11640                    // gets the PREVIOUS turn's last committed hidden (sess.last_h). Per chunk:
11641                    // rows start..end read h[start-1..end-1] — one dtod into a chunk buffer.
11642                    let mut phs = e.zeros(tc * n_embd)?;
11643                    let (src_lo, dst_off) = if start == 0 {
11644                        (0, n_embd)
11645                    } else {
11646                        ((start - 1) * n_embd, 0)
11647                    };
11648                    let n_copy = if start == 0 {
11649                        (tc - 1) * n_embd
11650                    } else {
11651                        tc * n_embd
11652                    };
11653                    if start == 0
11654                        && let Some((_, lh, _, _, _)) = sess_tail.as_ref()
11655                        && let Some(lh) = lh.as_ref()
11656                    {
11657                        e.copy_into(&mut phs, 0, lh, n_embd)?;
11658                    }
11659                    if n_copy > 0 {
11660                        e.copy_view_into(
11661                            &mut phs,
11662                            dst_off,
11663                            &ph.slice(src_lo..src_lo + n_copy),
11664                            n_copy,
11665                        )?;
11666                    }
11667                    self.mtp_kv_fill_all(
11668                        e,
11669                        &prompt[start..end],
11670                        &phs,
11671                        base + start,
11672                        &mut *scratch,
11673                        embd_dev,
11674                    )?;
11675                }
11676                start = end;
11677            }
11678        }
11679        // MEMRA_PROFILE_SPEC=2: profiler capture starts HERE — after the prime, so an
11680        // `nsys -c cudaProfilerApi` capture contains ONLY the round loop (draft/verify/commit).
11681        // (=1 brackets the whole call in run_spec.rs, prime included.)
11682        if std::env::var("MEMRA_PROFILE_SPEC").as_deref() == Ok("2") {
11683            unsafe extern "C" {
11684                fn cudaProfilerStart() -> i32;
11685            }
11686            unsafe {
11687                cudaProfilerStart();
11688            }
11689        }
11690        // ROUND-STREAM stage (c) 4 (MEMRA_SPEC_STREAM=1, experimental): pre-issued M-round
11691        // bursts with ZERO per-round host readbacks — the accept/seed/rollback/ring kernels
11692        // consume each other's device outputs; the host drains the ring every M rounds. v1
11693        // constraints: greedy, !spec_replay, single-shot, batched-linear layers, no refresh
11694        // fills (acceptance effect A/B-arbitrated), enters from round 1 (pending guaranteed).
11695        // NOTE: not gated on the caller's graph_draft (its trunk_dense conjunct turns the 35B
11696        // MoE off) — the stream capture encloses ONLY the dense MTP head; the head-dense /
11697        // full-prec / k gates are re-derived here and a failed capture degrades to stream-off.
11698        let stream_on = crate::spec::spec_stream()
11699            && !sampled
11700            && !spec_replay
11701            && self.mtp_extra.is_empty()
11702            && constraint.is_none()
11703            && !session_mode
11704            && embd_gpu.is_some()
11705            && !crate::model::full_prec_enabled()
11706            && k + 2 < 96;
11707        let mut stream_graph: Option<cudarc::driver::CudaGraph> = None;
11708        let mut g_tokp2k = e.alloc_u32_zeroed(2 * k.max(1))?;
11709        if stream_on {
11710            let cap = e.capture_graph(|e| {
11711                for j in 0..k.max(1) {
11712                    self.mtp_head_forward_cap(
11713                        e,
11714                        mtp,
11715                        &mut dctx.g_tok,
11716                        &mut dctx.g_pos,
11717                        &mut dctx.g_seed,
11718                        &mut dctx.g_p,
11719                        &mut *scratch,
11720                        0,
11721                        true,
11722                        true,
11723                        embd_gpu.expect("round stream requires resident embedding"),
11724                        embd_qt,
11725                        embd_rb,
11726                        d_vocab,
11727                        None,
11728                        Some((&mut g_tokp2k, j, d2t_dev.as_ref())),
11729                        None, // round-stream requires constraint.is_none() (see stream_on)
11730                    )?;
11731                }
11732                Ok(())
11733            });
11734            match cap {
11735                Ok(g) => {
11736                    scratch.set_len(e, 0)?;
11737                    stream_graph = Some(g);
11738                }
11739                Err(err) => {
11740                    scratch.set_len(e, 0)?;
11741                    if debug_spec {
11742                        eprintln!("[spec] stream-graph capture failed ({err}); stream off");
11743                    }
11744                }
11745            }
11746        }
11747        let stream_active = stream_on && stream_graph.is_some();
11748        if debug_spec {
11749            eprintln!(
11750                "[spec] stream_on={stream_on} env={} samp={sampled} dg={} captured={} active={stream_active} session={session_mode} replay={spec_replay}",
11751                crate::spec::spec_stream(),
11752                dctx.graph.is_some(),
11753                stream_graph.is_some()
11754            );
11755        }
11756        let t_v_s = k + 1;
11757        // ROUND-STREAM buffers + ptr tables now live in the model-generic round_stream
11758        // module (extracted 2026-07-12; the gemma burst reuses them).
11759        let sb = crate::round_stream::StreamBufs::new(e, k, crate::spec::spec_stream_m())?;
11760        let crate::round_stream::StreamBufs {
11761            mut vtok_d,
11762            mut brk_d,
11763            mut pend_d,
11764            last_pred_d,
11765            mut pos_ctr,
11766            mut pos_start_d,
11767            mut ring_d,
11768            acc_d: mut stream_acc,
11769            m_rounds,
11770            k: _,
11771        } = sb;
11772        let stream_ptrs: Option<CudaSlice<u64>> = if stream_active {
11773            Some(crate::round_stream::kv_len_ptr_table(
11774                e,
11775                cache,
11776                Some(&pos_ctr),
11777            )?)
11778        } else {
11779            None
11780        };
11781
11782        let t_fill = t_ent.elapsed();
11783        let mut round = 0usize;
11784        // ADAPTIVE DRAFT LENGTH (MEMRA_SPEC_ADAPT=1, opt-in — the gemma_spec accepted-run law,
11785        // ported 2026-08-01): next round's draft depth = last round's accepted run + 1, clamped
11786        // to [floor(pos), k_cap] — a miss shrinks the next draft to the miss point + 1,
11787        // full-accept streaks re-deepen one step per round. NOT the 2026-07-07 acceptance-EMA
11788        // (that arm measured an HONEST LOSS to static per-class optima — 115.0/85.8/73.4 vs
11789        // 121.6/92.7/75.6, EMA lag — and was removed 2026-07-08; rig5090.jsonl has the record).
11790        // The gemma law has no lag class: it reacts within one round, and was worth +7-20% on
11791        // the gemma cells at unchanged exactness (2026-07-10 flip; floor sweep 2026-07-25;
11792        // position key 2026-07-26). Signal = n_acc from the round's EXISTING accept readback —
11793        // zero new syncs; the draft graph is a SINGLE-STEP capture replayed per drafted token,
11794        // so a per-round depth needs no re-capture (unlike gemma's whole-chain graphs). qwen's
11795        // in-round p-min cut already shortens chains mid-round, so gemma's one-round-late p-min
11796        // fold into kc is unnecessary here — the accepted-run law sees the cut via n_acc.
11797        // Exactness is the verify's job at ANY depth (same contract as p-min variable rounds).
11798        // DEFAULT OFF on the qwen path until its cells gate a flip (gemma's is default-on).
11799        // MEASURED 2026-08-01 (H100 GPU-3, interleaved x3, NGEN=256, same-invocation plain
11800        // denominators; research/qwen-adaptive-k-20260801/): REFUTED on the tuned qwen configs.
11801        // q27 K=3+HPOST+PMIN=0.3: short +0.8% (noise; law ~idles, len_hist identical), board
11802        // -1.9%, agentic -0.5%; board PMIN=0 -2.1% (not p-min shadowing — the law itself);
11803        // floor=1 -2.8% (gemma's floor-collapse, reproduced). q35 K=2 board: -6.4% (52/136
11804        // rounds shrink to depth 1; no depth to reclaim at K=2). The gemma direction DOES
11805        // appear at untuned depth-K — q27 K=6 floor=4 +1.5% over fixed K=6 — but stays -3.7%
11806        // below fixed K=3: same verdict class as the retired EMA arm (honest loss to static
11807        // per-class optima). Acceptance-rate rises under the law while tokens/round falls —
11808        // it buys accept-% by adding rounds, and a round's fixed draft+verify cost wins.
11809        // K=1..8 self-consistency PASS both models with the law ON (exactness held).
11810        let adapt = std::env::var("MEMRA_SPEC_ADAPT").as_deref() == Ok("1");
11811        // floor: per-model default keyed on n_embd (gemma's tiering — models with an expensive
11812        // verify keep deep drafts after a miss); MEMRA_SPEC_ADAPT_FLOOR pins it everywhere.
11813        let adapt_floor_env: Option<usize> = std::env::var("MEMRA_SPEC_ADAPT_FLOOR")
11814            .ok()
11815            .and_then(|v| v.parse().ok());
11816        let adapt_floor_default: usize = if self.cfg.n_embd as usize >= 3500 {
11817            4
11818        } else if self.cfg.n_embd as usize >= 2500 {
11819            2
11820        } else {
11821            1
11822        };
11823        let adapt_floor: usize = adapt_floor_env.unwrap_or(adapt_floor_default);
11824        // position key: past floor_ctx a HIGH floor (>=4) relaxes to 1 — forced-deep drafts
11825        // turn net-negative at depth (gemma 31B d1736 evidence); MEMRA_SPEC_FLOOR_CTX moves
11826        // the boundary, an explicit MEMRA_SPEC_ADAPT_FLOOR pins the floor everywhere.
11827        let floor_ctx: usize = std::env::var("MEMRA_SPEC_FLOOR_CTX")
11828            .ok()
11829            .and_then(|v| v.parse().ok())
11830            .unwrap_or(1024);
11831        let floor_at = |pos: usize| -> usize {
11832            if adapt_floor_env.is_some() || pos < floor_ctx {
11833                adapt_floor
11834            } else if adapt_floor >= 4 {
11835                1
11836            } else {
11837                adapt_floor
11838            }
11839        };
11840        // cap: MEMRA_SPEC_CAPMAX (gemma semantics, default 7). Binds only under adapt — the
11841        // fixed-K default path is untouched by this whole block.
11842        let cap_max: usize = std::env::var("MEMRA_SPEC_CAPMAX")
11843            .ok()
11844            .and_then(|v| v.parse().ok())
11845            .unwrap_or(7);
11846        let k_cap = k.min(cap_max).max(1);
11847        let mut kc = k_cap;
11848        // Persistent snapshot buffers are allocated once and refreshed in place.
11849        let mut snap = cache.snapshot(e)?;
11850        // BONUS FOLD (2026-07-04): after a FULL accept the bonus token is NOT committed with a
11851        // separate T=1 trunk pass (a full weight read per round). It stays PENDING and rides as
11852        // column 0 of the NEXT round's verify batch. Under predecessor pairing the next chain
11853        // seeds from the bonus's predecessor's TRUE verify hidden (free — no extra
11854        // pass of any kind). Verify still
11855        // checks every emitted token against the target -> exactness holds by construction; only
11856        // DRAFT QUALITY can shift, which the acceptance numbers arbitrate.
11857        // bonus emitted but not yet committed to cache. A carried pending (see SpecSession::
11858        // pending_tok) enters round 0 directly — the burst boundary becomes a plain round edge.
11859        let mut pending: Option<u32> = carried_pending;
11860        // MEMRA_SPEC_PHASE=1: per-round wall decomposition (draft / verify / accept+commit) —
11861        // no tracing, no extra syncs (each phase is naturally sync-bounded: draft readbacks,
11862        // the verify accept readback). Printed once at loop end via spec-stats.
11863        let anatomy_on = std::env::var("MEMRA_SPEC_PP_ANATOMY").as_deref() == Ok("1");
11864        let phase_on = anatomy_on || std::env::var("MEMRA_SPEC_PHASE").as_deref() == Ok("1");
11865        // MEMRA_SPEC_PHASE_SYNC=1 — reads the phase split correctly, and proves it. `ph_mark` is a
11866        // bare Instant, so `verify-issue` is the host QUEUEING the walk (the GPU is already running
11867        // under it) and `verify-wait` is only the residual drain at the accept readback: one
11868        // overlapped interval cut at the first blocking call, NOT "GPU time" beside "host time".
11869        // Syncing right after the walk is issued moves the whole GPU wall into `verify-issue`. If
11870        // the walk's GPU total is really issue+wait, then with this on verify-issue jumps to that
11871        // sum, verify-wait collapses to the readback alone, and the ROUND WALL DOES NOT MOVE —
11872        // which is what says the queueing time was hidden and is not a target. Diagnostic only.
11873        let phase_sync = std::env::var("MEMRA_SPEC_PHASE_SYNC").as_deref() == Ok("1");
11874        // DRAFT-MASK receipt (lane/draft-mask): speculative-clone wall + rounds, printed with
11875        // spec-stats. The clone is the one cost the design adds per round — measured, not assumed.
11876        let (mut dm_clone_ns, mut dm_rounds) = (0u128, 0usize);
11877        // grammar-truncation counters: how many rounds the verify-side cut fired and how many
11878        // already-verified tokens it threw away. THIS is the quantity draft masking targets.
11879        let (mut dm_cuts, mut dm_cut_tokens) = (0usize, 0usize);
11880        let (mut ph_draft, mut ph_verify, mut ph_rest) = (0f64, 0f64, 0f64);
11881        let mut ph_wait = 0f64;
11882        let mut ph_commit = 0f64;
11883        let mut ph_t = std::time::Instant::now();
11884        let mut ph_mark = |acc: &mut f64, on: bool| {
11885            if on {
11886                let now = std::time::Instant::now();
11887                *acc += (now - ph_t).as_secs_f64();
11888                ph_t = now;
11889            }
11890        };
11891        // MTP-ROUTE VERIFY GRAPHS (`MEMRA_SPEC_VERIFY_GRAPH`, see the flag doc): the
11892        // model-owned capture pool, locked for the whole burst exactly as the dspark serve
11893        // arm holds it — the slab stash is live verify -> commit inside a round, and the
11894        // worker drives rounds from one scheduler thread. PERSISTENT across generations on
11895        // the model (rebuilding per call re-captures the pool per prompt, which is the
11896        // measured way to lose more than the launches cost); the captured bodies are
11897        // cache-independent, every state read going through per-round refreshed pointer
11898        // tables. None = the eager walk, byte-identical.
11899        //
11900        // Never armed together with ROUND-STREAM: the tparallel verify refuses that pair
11901        // loudly, and `stream_active` owns the burst arm above, so the door stays shut
11902        // whenever the stream is live rather than relying on that refusal.
11903        // The lock is taken ONLY when the door is armed: with the flag off this whole block
11904        // is inert, so the default path cannot serialize two spec generations behind a mutex
11905        // it never reads.
11906        let vg_armed =
11907            crate::spec::spec_verify_graph_env().unwrap_or_else(|| self.vgraph_family_default());
11908        let mut vg_guard = if vg_armed && !stream_active {
11909            let mut g = self.dspark_vgraphs.lock().unwrap();
11910            if g.is_none() {
11911                // Size by the WIDEST verify this run can present, which is k+1 and NOT
11912                // k_cap+1: the sampled arm's own window is `t_v_s = k + 1`, so a pool built
11913                // from a smaller adaptive cap gets sliced past its stash rows (a `slice_mut`
11914                // panic in the sampled ON arm, measured before this line said k+1).
11915                let vt_cap = (k.max(k_cap) + 1).max(2);
11916                *g = DsparkVerifyGraphs::new(e, cache, vt_cap, n_embd)?;
11917                if g.is_some() {
11918                    // Engagement receipt (the dead-arm lesson): prove the door is LIVE rather
11919                    // than trusting that a flag set means a pool built.
11920                    eprintln!("[spec-vg] MTP verify-graph pool ENGAGED (vt_cap={vt_cap})");
11921                } else {
11922                    eprintln!(
11923                        "[spec-vg] MTP verify-graph pool declined (no linear layers, \
11924                         non-uniform state, or vt_cap < 2) — eager walk"
11925                    );
11926                }
11927            }
11928            Some(g)
11929        } else {
11930            None
11931        };
11932        // Capacity fail-safe: a round wider than the pool was built for must take the eager
11933        // walk, not slice the stash past its rows. The sizing above already covers every
11934        // round this run can present; this keeps a future caller (or a k that grows behind
11935        // the pool's back) on the byte-identical fallback instead of a panic.
11936        let vg_t_cap = vg_guard
11937            .as_ref()
11938            .and_then(|g| g.as_ref())
11939            .map(|g| g.t_capacity())
11940            .unwrap_or(0);
11941        let mut graph_guard_noted = false;
11942        while keep_going && out.len() < max_new {
11943            // GRAPH-LAUNCH HEADROOM GUARD (see GRAPH_LAUNCH_MIN_FREE): below the floor,
11944            // every captured-graph arm in this round yields to its byte-identical eager
11945            // twin instead of feeding cuGraphLaunch a card it segfaults on.
11946            let graph_round_ok = graph_launch_headroom_ok(e);
11947            if !graph_round_ok && !graph_guard_noted {
11948                graph_guard_noted = true;
11949                eprintln!(
11950                    "[spec] graph replay suspended: driver free below the {}MB launch floor \
11951                     (eager arms serve; cuGraphLaunch segfaults into an exhausted card)",
11952                    GRAPH_LAUNCH_MIN_FREE / (1 << 20)
11953                );
11954            }
11955            // MEMRA_SPEC_ROUND_PROF=1: wall of the WHOLE round against the pieces we already
11956            // instrument. Needed because the parts do not add up: the draft step measures 1.27 ms
11957            // ([spec-anatomy] glue 92 / attn 280 / ffn 222 / head 670 us) and the t=2 verify walk
11958            // 25.6 ms ([tcol-prof] attn 10.1 + ffn 15.3), yet a K=1 round takes 177 ms on the
11959            // step37 TP2 stack. This prints where the other ~150 ms lives.
11960            let round_prof = ROUND_PROF
11961                .get_or_init(|| std::env::var("MEMRA_SPEC_ROUND_PROF").as_deref() == Ok("1"));
11962            let round_t0 = round_prof.then(std::time::Instant::now);
11963            // ROUND-STREAM BURST: from round 1 (pending guaranteed by every non-replay arm),
11964            // issue M rounds with zero readbacks, then drain the ring + reconcile mirrors.
11965            if let (true, Some(sg), Some(ptrs)) = (
11966                stream_active && round >= 1 && pending.is_some() && graph_round_ok,
11967                &stream_graph,
11968                &stream_ptrs,
11969            ) {
11970                if debug_spec {
11971                    static ONCE: std::sync::Once = std::sync::Once::new();
11972                    ONCE.call_once(|| {
11973                        eprintln!("[memra] ROUND-STREAM burst engaged (M={m_rounds} k={k})")
11974                    });
11975                }
11976                e.set_i32_one(&mut pos_ctr, cache.pos as i32)?;
11977                e.set_u32_one(&mut pend_d, pending.unwrap())?;
11978                e.set_u32_one(&mut ring_d, 0)?; // ring count = 0 (writes element 0)
11979                for _mi in 0..m_rounds {
11980                    e.i32_copy_add(&pos_ctr, &mut pos_start_d, 0)?;
11981                    cache.snapshot_into(e, &mut snap)?; // device D2Ds, stream-ordered
11982                    e.i32_copy_add(&pos_ctr, &mut scratch.kv.len_d, 0)?; // draft-KV rollback
11983                    e.i32_copy_add(&pos_ctr, &mut dctx.g_pos, 1)?; // rope pos = pos + base
11984                    e.u32_copy(&pend_d, &mut dctx.g_tok)?;
11985                    e.copy_into(&mut dctx.g_seed, 0, &h_seed_buf, n_embd)?;
11986                    sg.launch()?;
11987                    e.spec_assemble_verify(
11988                        &g_tokp2k,
11989                        &pend_d,
11990                        d2t_dev.as_ref(),
11991                        &mut vtok_d,
11992                        &mut brk_d,
11993                        p_min,
11994                        k,
11995                        pmin0,
11996                    )?;
11997                    let mut ck = VerifyCkpt::new(self.layers.len());
11998                    let dummy = vec![0u32; t_v_s];
11999                    let (tl_d, vx) = self.decode_step_t_core_stream(
12000                        e,
12001                        &dummy,
12002                        0,
12003                        &mut *cache,
12004                        embd_dev,
12005                        Some(&mut ck),
12006                        Some((&vtok_d, &pos_ctr)),
12007                        None,
12008                        None,
12009                    )?;
12010                    for j in 0..t_v_s {
12011                        e.argmax_token_device_col(&tl_d, j, n_vocab, &mut preds_d, j)?;
12012                    }
12013                    e.spec_accept_greedy_dc(
12014                        &preds_d,
12015                        &vtok_d,
12016                        &last_pred_d,
12017                        &brk_d,
12018                        &mut stream_acc,
12019                    )?;
12020                    e.spec_seed_gather(&vx, &fill_prev, &stream_acc, &mut h_seed_buf, 1, n_embd)?;
12021                    e.copy_into(&mut fill_prev, 0, &h_seed_buf, n_embd)?;
12022                    self.commit_verified_prefix_stream(
12023                        e,
12024                        &mut *cache,
12025                        &snap,
12026                        &ck,
12027                        &stream_acc,
12028                        1,
12029                        t_v_s,
12030                    )?;
12031                    e.spec_rollback_stream(
12032                        ptrs,
12033                        &pos_start_d,
12034                        &stream_acc,
12035                        1,
12036                        self.layers.len() + 1,
12037                    )?;
12038                    e.spec_ring_commit(&vtok_d, &stream_acc, &brk_d, &mut ring_d, &mut pend_d)?;
12039                }
12040                e.stream().synchronize()?;
12041                let ring_h = e.dtoh_u32(&ring_d)?;
12042                let cnt = ring_h[0] as usize;
12043                for i in 0..cnt {
12044                    if out.len() < max_new {
12045                        out.push(ring_h[1 + i]);
12046                    }
12047                }
12048                let pos_h = e.dtoh_i32(&pos_ctr)?[0] as usize;
12049                for il in 0..self.layers.len() {
12050                    if let Some(kvl) = cache.kv[il].as_mut() {
12051                        kvl.len = pos_h;
12052                    }
12053                }
12054                cache.pos = pos_h;
12055                scratch.kv.len = pos_h;
12056                pending = Some(ring_h[cnt]); // last drained token = the live bonus
12057                last_token = ring_h[cnt];
12058                total_drafted += k * m_rounds; // upper bound (p-min breaks uncounted)
12059                total_accepted += cnt.saturating_sub(m_rounds);
12060                if let Some(t) = sess_telem {
12061                    // totals only — the burst's per-round accept counts stayed on device
12062                    // (that is the point of the round-stream arm). pos_* untouched.
12063                    t.record_totals(m_rounds, k * m_rounds, cnt.saturating_sub(m_rounds));
12064                }
12065                round += m_rounds;
12066                // sse-cadence: the drained ring is committed — flush it at burst-drain cadence.
12067                keep_going = flush_commit(&mut on_commit, &out, &mut flushed);
12068                continue;
12069            }
12070            let pos = cache.pos; // #tokens committed (EXCLUDES a pending bonus)
12071            cache.snapshot_into(e, &mut snap)?; // §C: snapshot BEFORE draft+verify
12072            ph_mark(&mut ph_rest, phase_on);
12073
12074            // --- 1. DRAFT k tokens with the NextN head (autoregressive, T=1 each) ---
12075            // p-min semantics (both paths): stop the chain early when the head's confidence in
12076            // its own pick drops below p_min — the just-drafted token is DISCARDED, but its
12077            // scratch append stands (identical to the eager chain's ordering). j==0 always drafts.
12078            let base0 = if pending.is_some() { 1usize } else { 0usize };
12079            // fixed draft length by default; MEMRA_SPEC_ADAPT=1 drafts at last round's
12080            // accepted run + 1 (the gemma law — see the setup block above the loop).
12081            let k_this = if adapt { kc } else { k };
12082            let mut draft: Vec<u32> = Vec::with_capacity(k);
12083            let mut draft_idx: Vec<u32> = Vec::with_capacity(k); // trimmed-vocab ids (== draft when untrimmed)
12084            {
12085                // Round-start draft-KV sync (BOTH paths). Persistent: truncate/align to the committed
12086                // history — slots 0..P hold entries for the tokens before last_token@P (P = pos +
12087                // base0 - 1); this single set_len IS the draft-side rollback (drops last round's
12088                // rejected drafts and p-min extras via the len mechanism).
12089                scratch.set_len(e, pos + base0 - 1)?;
12090                // dcw door: a captured chain appends k_this device-counter rows (plus the
12091                // pseudo-seed replay) with no host intervention; any ring rebase those appends
12092                // could need happens HERE, host-side, before the replays. The eager arm keeps
12093                // its own per-step prepare, so this is graph-path-only work.
12094                if step35_draft_dcw_on()
12095                    && (dctx.graph.is_some()
12096                        || dctx.graph_s.is_some()
12097                        || dctx.chain.is_some()
12098                        || dctx.chain_s.is_some())
12099                {
12100                    scratch.ensure_dcw_headroom(e, k_this + 2)?;
12101                }
12102                if pen_on {
12103                    // PEN_WINDOW_MAX also bounds the per-round upload and the O(n_hist^2)
12104                    // device dedup: the serve window is already PEN_WINDOW_MAX, and this
12105                    // defensive min also bounds non-server callers.
12106                    let win = sp.penalty_last_n.min(PEN_WINDOW_MAX);
12107                    let w0 = pen_hist.len().saturating_sub(win);
12108                    pen_hist_d = Some(e.htod_u32_v(&pen_hist[w0..])?);
12109                }
12110                if sampled {
12111                    draft_logits.clear();
12112                    draft_stats.clear();
12113                }
12114                // DRAFT-SIDE GRAMMAR MASK: clone the committed grammar state ONCE per round; each
12115                // position's mask is computed on that clone and advanced by the PROPOSED token. The
12116                // real state moves only on emission (verify's job), so the emitted stream is
12117                // unchanged — the mask only removes tokens the verify would have truncated anyway.
12118                let mut dmask_live = dmask_on;
12119                if dmask_live {
12120                    let t_c = std::time::Instant::now();
12121                    constraint
12122                        .as_deref_mut()
12123                        .unwrap()
12124                        .draft_begin()
12125                        .map_err(|e2| format!("constraint: {e2}"))?;
12126                    dm_clone_ns += t_c.elapsed().as_nanos();
12127                    dm_rounds += 1;
12128                }
12129                if let (false, Some(cg)) = (sampled || pen_on || !graph_round_ok, &dctx.chain) {
12130                    // GREEDY CHAIN GRAPH (lane/step37-draft-graph-serving-20260830): the
12131                    // eager multi-head chain's EXACT launch order — step j rewinds head
12132                    // (j % heads)'s plane to the committed length and replays rows 0..=j —
12133                    // with each row's whole head-forward as ONE graph launch. The chain
12134                    // POLICY (head choice, prefix length, stored-seed feed) is host-side,
12135                    // identical to `mtp_chain_forward_dev`, so graph-vs-eager drafts are
12136                    // bit-identical by construction (same launcher, same bucket — the dcw
12137                    // parity contract). Interior rows launch the head-less graph: their
12138                    // logits are dead in the eager chain too, so the consumed bytes match.
12139                    let heads_n = self.mtp_head_count();
12140                    let committed = pos + base0 - 1;
12141                    let mut chain_tokens: Vec<u32> = vec![last_token];
12142                    let mut chain_seed_bufs: Vec<CudaSlice<f32>> = vec![e.clone_dtod(&h_seed_buf)?];
12143                    for j in 0..k_this {
12144                        let index = mtp_chain_head_index(j, heads_n);
12145                        if debug_spec {
12146                            eprintln!(
12147                                "[mtp-chain-step] round={round} j={j} head={index} \
12148                                 replay_rows={} arm=graph",
12149                                chain_tokens.len(),
12150                            );
12151                        }
12152                        scratch.set_plane_len(e, index, committed)?;
12153                        e.set_i32_one(&mut dctx.g_pos, (committed + 1) as i32)?;
12154                        for row in 0..=j {
12155                            e.set_u32_one(&mut dctx.g_tok, chain_tokens[row])?;
12156                            e.copy_into(&mut dctx.g_seed, 0, &chain_seed_bufs[row], n_embd)?;
12157                            if row < j {
12158                                cg.interior[index].launch()?;
12159                            } else {
12160                                // per-position mask upload before the LAST row only — the
12161                                // eager chain applies the mask on is_last exactly the same.
12162                                if dmask_live
12163                                    && !upload_draft_mask(
12164                                        e,
12165                                        constraint.as_deref_mut().unwrap(),
12166                                        &mut dctx.g_dmask,
12167                                        mtp.d2t.as_ref(),
12168                                        d_vocab,
12169                                        dmask_words,
12170                                    )?
12171                                {
12172                                    e.htod_u32_into(
12173                                        &mut dctx.g_dmask,
12174                                        &vec![u32::MAX; dmask_words],
12175                                    )?;
12176                                    dmask_live = false;
12177                                }
12178                                cg.last[index].launch()?;
12179                            }
12180                            // host mirror (len_d advanced in-graph by the dcw append)
12181                            scratch.plane_mut(index).0.len += 1;
12182                        }
12183                        let idx = e.dtoh_u32_one(&dctx.g_tok)?;
12184                        // #87 SENTINEL TRAP (see the single-head graph arm below).
12185                        if (idx as usize) >= d_vocab {
12186                            let seed_h = e.dtoh(&dctx.g_seed)?;
12187                            let seed_nan = seed_h.iter().filter(|v| v.is_nan()).count();
12188                            return Err(format!(
12189                                "draft(chain-graph) argmax sentinel 0x{idx:08x} >= d_vocab \
12190                             {d_vocab} at round {round} j={j} head={index} pos={pos}: \
12191                             head-out NaN {seed_nan}/{n_embd} — refusing to dereference \
12192                             the embed row (#87 trap)"
12193                            )
12194                            .into());
12195                        }
12196                        // multi-head MTP forbids a trimmed head (validated at entry), so the
12197                        // draft index IS the target id; keep the map for uniformity.
12198                        let d = match &mtp.d2t {
12199                            Some(map) => map[idx as usize],
12200                            None => idx,
12201                        };
12202                        let draft_p = if p_min > 0.0 {
12203                            Some(e.dtoh(&dctx.g_p)?[0])
12204                        } else {
12205                            None
12206                        };
12207                        if let Some(p) = draft_p.filter(|_| p_min > 0.0)
12208                            && p < p_min
12209                            && (j > 0 || (pmin0 && base0 == 1))
12210                        {
12211                            break;
12212                        }
12213                        draft.push(d);
12214                        chain_tokens.push(d);
12215                        // step j's h_nextn: the last-row graph self-fed it into g_seed —
12216                        // snapshot it as the chain history seed for row j+1 (stream-ordered
12217                        // after the launch, exactly the eager chain's chain_seeds push).
12218                        chain_seed_bufs.push(e.clone_dtod(&dctx.g_seed)?);
12219                        // speculative grammar advance (see the single-head graph arm).
12220                        if dmask_live
12221                            && !constraint
12222                                .as_deref_mut()
12223                                .unwrap()
12224                                .draft_advance(d)
12225                                .map_err(|e2| format!("constraint: {e2}"))?
12226                        {
12227                            e.htod_u32_into(&mut dctx.g_dmask, &vec![u32::MAX; dmask_words])?;
12228                            break;
12229                        }
12230                    }
12231                } else if let (true, Some(cg)) = (
12232                    sampled && s_capturable && dctx.s_key == Some(s_key) && graph_round_ok,
12233                    &dctx.chain_s,
12234                ) {
12235                    if skey_probe() {
12236                        eprintln!(
12237                            "[skey] chain=graph_chain_s round={round} capturable={} top_k={} \
12238                             top_p={} min_p={} s_key_parked={:?}",
12239                            s_capturable as u8, sp.top_k, sp.top_p, sp.min_p, dctx.s_key,
12240                        );
12241                    }
12242                    // SAMPLED CHAIN GRAPH: the greedy chain arm's launch order with the
12243                    // sampled last-row graphs — in-graph counter bump + (filtered) gumbel
12244                    // draw + argmax; q retained per step into q_slots exactly like the
12245                    // single-head sampled graph arm. Counter continuity: g_ctr host-seeded
12246                    // to sctr-1 once per ROUND; each step's last-row graph bumps it BEFORE
12247                    // the perturb, so step j consumes counter sctr+j — the eager Philox
12248                    // stream (interior rows never draw, never bump).
12249                    let heads_n = self.mtp_head_count();
12250                    let committed = pos + base0 - 1;
12251                    let filtered_stats_in_graph = s_key.filtered();
12252                    let mut chain_tokens: Vec<u32> = vec![last_token];
12253                    let mut chain_seed_bufs: Vec<CudaSlice<f32>> = vec![e.clone_dtod(&h_seed_buf)?];
12254                    e.set_u32_one(&mut dctx.g_ctr, sctr.wrapping_sub(1))?;
12255                    for j in 0..k_this {
12256                        let index = mtp_chain_head_index(j, heads_n);
12257                        if debug_spec {
12258                            eprintln!(
12259                                "[mtp-chain-step] round={round} j={j} head={index} \
12260                                 replay_rows={} arm=graph_s",
12261                                chain_tokens.len(),
12262                            );
12263                        }
12264                        scratch.set_plane_len(e, index, committed)?;
12265                        e.set_i32_one(&mut dctx.g_pos, (committed + 1) as i32)?;
12266                        for row in 0..=j {
12267                            e.set_u32_one(&mut dctx.g_tok, chain_tokens[row])?;
12268                            e.copy_into(&mut dctx.g_seed, 0, &chain_seed_bufs[row], n_embd)?;
12269                            if row < j {
12270                                cg.interior[index].launch()?;
12271                            } else {
12272                                cg.last[index].launch()?;
12273                            }
12274                            scratch.plane_mut(index).0.len += 1;
12275                        }
12276                        sctr += 1; // mirrors the in-graph g_ctr bump (eager parity:
12277                        // counts the p-min-discarded token too)
12278                        // q retention: ONE async D2D of the persistent head-logits buffer
12279                        // into this round's slot j (stream-ordered after the replay).
12280                        e.copy_into(&mut dctx.q_slots[j], 0, &dctx.g_q, d_vocab)?;
12281                        // FILTERED capture: read the in-graph filter_stats scalars back per
12282                        // replay instead of a second full-vocab filter_stats per slot post-
12283                        // chain — bit-exact (the values the in-graph perturb consumed) and
12284                        // measured worth ~5% of vendor-default serving tok/s at K=3. Before
12285                        // the p-min break so the discarded slot's stats land too.
12286                        if filtered_stats_in_graph {
12287                            draft_stats.push((
12288                                e.dtoh(&dctx.g_mx)?[0],
12289                                e.dtoh(&dctx.g_th)?[0],
12290                                e.dtoh(&dctx.g_z)?[0],
12291                            ));
12292                        }
12293                        let idx = e.dtoh_u32_one(&dctx.g_tok)?;
12294                        // #87 SENTINEL TRAP (see the single-head graph arms).
12295                        if (idx as usize) >= d_vocab {
12296                            let seed_h = e.dtoh(&dctx.g_seed)?;
12297                            let seed_nan = seed_h.iter().filter(|v| v.is_nan()).count();
12298                            return Err(format!(
12299                                "draft(chain-graph-sampled) argmax sentinel 0x{idx:08x} >= \
12300                             d_vocab {d_vocab} at round {round} j={j} head={index} pos={pos}: \
12301                             head-out NaN {seed_nan}/{n_embd} — refusing to dereference the \
12302                             embed row (#87 trap)"
12303                            )
12304                            .into());
12305                        }
12306                        let d = match &mtp.d2t {
12307                            Some(map) => map[idx as usize],
12308                            None => idx,
12309                        };
12310                        draft_idx.push(idx);
12311                        if p_min > 0.0 {
12312                            let p = e.dtoh(&dctx.g_p)?[0];
12313                            if p < p_min && (j > 0 || (pmin0 && base0 == 1)) {
12314                                break;
12315                            }
12316                        }
12317                        draft.push(d);
12318                        chain_tokens.push(d);
12319                        chain_seed_bufs.push(e.clone_dtod(&dctx.g_seed)?);
12320                    }
12321                    // PURE-TEMP accept path: stats per used slot recomputed from the RETAINED
12322                    // q with the SAME filter_stats program the eager arm runs (deployment-
12323                    // keyed coop/plain choice, same input bits). The FILTERED graph read its
12324                    // stats back per replay above.
12325                    if !filtered_stats_in_graph {
12326                        for j in 0..draft.len().max(draft_idx.len()) {
12327                            let rows0 = e.htod_i32(&[0])?;
12328                            let (mut th_d, mut z_d, mut mx_d) =
12329                                (e.zeros(1)?, e.zeros(1)?, e.zeros(1)?);
12330                            e.filter_stats(
12331                                &dctx.q_slots[j],
12332                                d_vocab,
12333                                &rows0,
12334                                &mut th_d,
12335                                &mut z_d,
12336                                &mut mx_d,
12337                                d_vocab,
12338                                1,
12339                                sp_temp,
12340                                sp.top_k,
12341                                sp.top_p,
12342                                sp.min_p,
12343                            )?;
12344                            draft_stats.push((
12345                                e.dtoh(&mx_d)?[0],
12346                                e.dtoh(&th_d)?[0],
12347                                e.dtoh(&z_d)?[0],
12348                            ));
12349                        }
12350                    }
12351                } else if let (false, Some(gr)) =
12352                    (sampled || pen_on || !graph_round_ok, &dctx.graph)
12353                {
12354                    // GRAPH DRAFT: one dispatch per drafted token. The chain feeds itself on-device
12355                    // (in-graph argmax -> tok_d -> next replay's embed; h_nextn -> h_seed_d; pos_d
12356                    // inc'd in-graph); the host only reads 4B token (+4B p) and decides the break.
12357                    e.set_i32_one(&mut dctx.g_pos, (pos + base0) as i32)?;
12358                    e.set_u32_one(&mut dctx.g_tok, last_token)?;
12359                    e.copy_into(&mut dctx.g_seed, 0, &h_seed_buf, n_embd)?;
12360                    for j in 0..k_this {
12361                        // per-position mask upload (contents only — the graph's baked pointer is
12362                        // dctx.g_dmask). All-ones once masking goes dead mid-chain, so the captured
12363                        // mask node degrades to a no-op ban instead of needing a second graph.
12364                        if dmask_live
12365                            && !upload_draft_mask(
12366                                e,
12367                                constraint.as_deref_mut().unwrap(),
12368                                &mut dctx.g_dmask,
12369                                mtp.d2t.as_ref(),
12370                                d_vocab,
12371                                dmask_words,
12372                            )?
12373                        {
12374                            // no draft-vocab row is grammar-legal here (a trimmed FR-Spec head can
12375                            // genuinely miss the legal set): neutralize the captured mask node and
12376                            // finish the chain UNMASKED — exactly pre-lane behaviour, never worse.
12377                            e.htod_u32_into(&mut dctx.g_dmask, &vec![u32::MAX; dmask_words])?;
12378                            dmask_live = false;
12379                        }
12380                        gr.launch()?;
12381                        scratch.kv.len += 1; // host mirror (len_d advanced in-graph)
12382                        let idx = e.dtoh_u32_one(&dctx.g_tok)?;
12383                        // #87 SENTINEL TRAP: an all-NaN head-logits row leaves the device argmax's
12384                        // init sentinel (0x7FFFFFFF) in g_tok — feeding it onward dereferences
12385                        // embed_row(sentinel) = table + ~4.6TB (never mapped) inside the NEXT graph
12386                        // replay's embed node, and the MMU fault kills the CUDA context for the
12387                        // whole process (research/pp2spec-crash-20260807: 3 coredumps, byte-exact
12388                        // VA arithmetic). Refuse loudly instead; the diagnostics name the first-NaN
12389                        // buffer (g_seed = the verify-side handoff vs head-side compute).
12390                        if (idx as usize) >= d_vocab {
12391                            // g_seed is SELF-FED (the replay writes h_nextn back into it), so it
12392                            // reads as the head's OUTPUT at j; h_seed_buf is the round's INPUT
12393                            // seed, untouched since the round-start copy — the pair discriminates
12394                            // "seed arrived poisoned" from "head forward produced NaN".
12395                            let seed_h = e.dtoh(&dctx.g_seed)?;
12396                            let seed_nan = seed_h.iter().filter(|v| v.is_nan()).count();
12397                            let in_h = e.dtoh(&h_seed_buf)?;
12398                            let in_nan = in_h.iter().filter(|v| v.is_nan()).count();
12399                            return Err(format!(
12400                                "draft(graph) argmax sentinel 0x{idx:08x} >= d_vocab {d_vocab} at \
12401                             round {round} j={j} pos={pos}: head-out NaN {seed_nan}/{n_embd}, \
12402                             round-input-seed NaN {in_nan}/{n_embd} — refusing to dereference \
12403                             the embed row (#87 trap)"
12404                            )
12405                            .into());
12406                        }
12407                        // trimmed draft vocab -> target token id (identity when no d2t map)
12408                        let d = match &mtp.d2t {
12409                            Some(map) => map[idx as usize],
12410                            None => idx,
12411                        };
12412                        let draft_p = if p_min > 0.0 {
12413                            Some(e.dtoh(&dctx.g_p)?[0])
12414                        } else {
12415                            None
12416                        };
12417                        if let Some(p) = draft_p.filter(|_| p_min > 0.0)
12418                            && p < p_min
12419                            && (j > 0 || (pmin0 && base0 == 1))
12420                        {
12421                            break;
12422                        }
12423                        draft.push(d);
12424                        // with a trimmed head the NEXT embed must read the TARGET id, not the draft
12425                        // index the argmax wrote — patch the persistent token buffer (4B htod).
12426                        if d != idx {
12427                            e.set_u32_one(&mut dctx.g_tok, d)?;
12428                        }
12429                        // advance the SPECULATIVE state with the proposal; a dead chain drops to
12430                        // unmasked drafting for the remaining positions (verify still arbitrates).
12431                        // speculative advance; a chain the grammar can no longer follow (EOS
12432                        // proposed) ends here. The captured mask node always runs, so a dead chain
12433                        // leaves the buffer NEUTRAL (all-ones = ban nothing) before it exits.
12434                        if dmask_live
12435                            && !constraint
12436                                .as_deref_mut()
12437                                .unwrap()
12438                                .draft_advance(d)
12439                                .map_err(|e2| format!("constraint: {e2}"))?
12440                        {
12441                            e.htod_u32_into(&mut dctx.g_dmask, &vec![u32::MAX; dmask_words])?;
12442                            break;
12443                        }
12444                    }
12445                // REGIME RE-TEST (lane/graph-s-key-exactness-20260819, widened by
12446                // lane/step37-draft-graph-serving-20260830): the sampled graph is legal ONLY
12447                // in the regime it was captured in. The condition used to read
12448                // `(sampled, &dctx.graph_s)` and trusted `s_key` to have dropped anything
12449                // else — which it could not, because the key omitted the filters. Both
12450                // halves are enforced: the key drops a stale graph, and this site refuses to
12451                // launch one whose key differs or whose regime is uncapturable (penalties).
12452                } else if let (true, Some(gr)) = (
12453                    sampled && s_capturable && dctx.s_key == Some(s_key) && graph_round_ok,
12454                    &dctx.graph_s,
12455                ) {
12456                    if skey_probe() {
12457                        eprintln!(
12458                            "[skey] chain=graph_s round={round} pure_temp={} capturable={} \
12459                             top_k={} top_p={} min_p={} s_key_parked={:?}",
12460                            pure_temp as u8,
12461                            s_capturable as u8,
12462                            sp.top_k,
12463                            sp.top_p,
12464                            sp.min_p,
12465                            dctx.s_key,
12466                        );
12467                    }
12468                    // SAMPLED GRAPH DRAFT: one replay per drafted token — head forward + gumbel +
12469                    // argmax in ONE dispatch; the host reads 4B token (+4B p), D2Ds q into slot j,
12470                    // and decides the break. Event-counter continuity: g_ctr is host-seeded to
12471                    // sctr-1 ONCE per round (outside the graph); the in-graph bump runs BEFORE the
12472                    // perturb, so replay j consumes counter sctr+j — exactly the eager arm's Philox
12473                    // stream. Host sctr advances in lockstep (computed, no readback needed).
12474                    e.set_i32_one(&mut dctx.g_pos, (pos + base0) as i32)?;
12475                    e.set_u32_one(&mut dctx.g_tok, last_token)?;
12476                    e.copy_into(&mut dctx.g_seed, 0, &h_seed_buf, n_embd)?;
12477                    e.set_u32_one(&mut dctx.g_ctr, sctr.wrapping_sub(1))?;
12478                    let filtered_stats_in_graph = s_key.filtered();
12479                    for j in 0..k_this {
12480                        gr.launch()?;
12481                        scratch.kv.len += 1; // host mirror (len_d advanced in-graph)
12482                        sctr += 1; // mirrors the in-graph g_ctr bump (eager parity:
12483                        // counts the p-min-discarded token too)
12484                        // q retention: ONE async D2D of the persistent head-logits buffer into this
12485                        // round's slot j (stream-ordered after the replay, before the next one).
12486                        e.copy_into(&mut dctx.q_slots[j], 0, &dctx.g_q, d_vocab)?;
12487                        // FILTERED capture: the replay's own filter_stats node already computed
12488                        // (th, z, mx) — read the three scalars back instead of paying a SECOND
12489                        // full-vocab filter_stats per slot post-chain (measured ~5% of vendor-
12490                        // default serving tok/s at K=3). Bit-exact by construction: these are
12491                        // the very values the in-graph perturb consumed. Read BEFORE the p-min
12492                        // break so the discarded slot's stats land too (accept-path indexing).
12493                        if filtered_stats_in_graph {
12494                            draft_stats.push((
12495                                e.dtoh(&dctx.g_mx)?[0],
12496                                e.dtoh(&dctx.g_th)?[0],
12497                                e.dtoh(&dctx.g_z)?[0],
12498                            ));
12499                        }
12500                        let idx = e.dtoh_u32_one(&dctx.g_tok)?;
12501                        // #87 SENTINEL TRAP (see the greedy graph arm above).
12502                        if (idx as usize) >= d_vocab {
12503                            let seed_h = e.dtoh(&dctx.g_seed)?;
12504                            let seed_nan = seed_h.iter().filter(|v| v.is_nan()).count();
12505                            return Err(format!(
12506                                "draft(graph-sampled) argmax sentinel 0x{idx:08x} >= d_vocab \
12507                             {d_vocab} at round {round} j={j} pos={pos}: round-seed NaN \
12508                             {seed_nan}/{n_embd} — refusing to dereference the embed row \
12509                             (#87 trap)"
12510                            )
12511                            .into());
12512                        }
12513                        let d = match &mtp.d2t {
12514                            Some(map) => map[idx as usize],
12515                            None => idx,
12516                        };
12517                        draft_idx.push(idx);
12518                        if p_min > 0.0 {
12519                            let p = e.dtoh(&dctx.g_p)?[0];
12520                            if p < p_min && (j > 0 || (pmin0 && base0 == 1)) {
12521                                break;
12522                            }
12523                        }
12524                        draft.push(d);
12525                        // trimmed head: the NEXT embed must read the TARGET id (see the greedy arm).
12526                        if d != idx {
12527                            e.set_u32_one(&mut dctx.g_tok, d)?;
12528                        }
12529                    }
12530                    // PURE-TEMP accept path: fill draft_stats per used slot post-chain (the
12531                    // stats degenerate to th=0 / full-Z; one filter_stats launch per slot).
12532                    // The FILTERED graph read its stats back per replay above.
12533                    if !filtered_stats_in_graph {
12534                        for j in 0..draft.len().max(draft_idx.len()) {
12535                            let rows0 = e.htod_i32(&[0])?;
12536                            let (mut th_d, mut z_d, mut mx_d) =
12537                                (e.zeros(1)?, e.zeros(1)?, e.zeros(1)?);
12538                            e.filter_stats(
12539                                &dctx.q_slots[j],
12540                                d_vocab,
12541                                &rows0,
12542                                &mut th_d,
12543                                &mut z_d,
12544                                &mut mx_d,
12545                                d_vocab,
12546                                1,
12547                                sp_temp,
12548                                sp.top_k,
12549                                sp.top_p,
12550                                sp.min_p,
12551                            )?;
12552                            draft_stats.push((
12553                                e.dtoh(&mx_d)?[0],
12554                                e.dtoh(&th_d)?[0],
12555                                e.dtoh(&z_d)?[0],
12556                            ));
12557                        }
12558                    }
12559                } else {
12560                    if skey_probe() && sampled {
12561                        eprintln!(
12562                            "[skey] chain=eager round={round} pure_temp={} top_k={} \
12563                             top_p={} min_p={} s_key_parked={:?}",
12564                            pure_temp as u8, sp.top_k, sp.top_p, sp.min_p, dctx.s_key,
12565                        );
12566                    }
12567                    // EAGER DRAFT (fallback: MoE head/trunk, huge k, MEMRA_SPEC_NOGRAPH, capture fail).
12568                    let chain_heads = !self.mtp_extra.is_empty();
12569                    let mut e_tok = last_token;
12570                    let mut d_seed = e.clone_dtod(&h_seed_buf)?;
12571                    let mut chain_tokens = if chain_heads {
12572                        vec![last_token]
12573                    } else {
12574                        Vec::new()
12575                    };
12576                    let mut chain_seeds = if chain_heads {
12577                        vec![e.clone_dtod(&h_seed_buf)?]
12578                    } else {
12579                        Vec::new()
12580                    };
12581                    for j in 0..k_this {
12582                        // GPU-ARGMAX DRAFT (2026-07-03): device logits + device argmax + 4-byte token
12583                        // read instead of the ~600KB full-vocab dtoh + host argmax per draft token.
12584                        let mtp_pos = pos + base0 + j;
12585                        // draft-side grammar mask (eager twin of the graph arm's in-graph node).
12586                        // A position with no legal draft-vocab row drops to unmasked drafting for
12587                        // the rest of the chain (pre-lane behaviour; verify still arbitrates).
12588                        if dmask_live {
12589                            dmask_live = upload_draft_mask(
12590                                e,
12591                                constraint.as_deref_mut().unwrap(),
12592                                &mut dctx.g_dmask,
12593                                mtp.d2t.as_ref(),
12594                                d_vocab,
12595                                dmask_words,
12596                            )?;
12597                        }
12598                        let mask = if dmask_live {
12599                            Some((&dctx.g_dmask, dmask_words))
12600                        } else {
12601                            None
12602                        };
12603                        let (dl_d, h_nextn) = if chain_heads {
12604                            if debug_spec {
12605                                eprintln!(
12606                                    "[mtp-chain-step] round={round} j={j} head={} replay_rows={}",
12607                                    mtp_chain_head_index(j, self.mtp_head_count()),
12608                                    chain_tokens.len(),
12609                                );
12610                            }
12611                            self.mtp_chain_forward_dev(
12612                                e,
12613                                &chain_tokens,
12614                                &chain_seeds,
12615                                &mut *scratch,
12616                                pos + base0 - 1,
12617                                embd_dev,
12618                                mask,
12619                            )?
12620                        } else {
12621                            self.mtp_head_forward_dev(
12622                                e,
12623                                mtp,
12624                                e_tok,
12625                                &d_seed,
12626                                &mut *scratch,
12627                                mtp_pos,
12628                                embd_dev,
12629                                mask,
12630                            )?
12631                        };
12632                        let tok_d = if sampled {
12633                            // FILTERED Gumbel-max: stats -> masked perturb -> argmax = one draw from
12634                            // the filtered softmax (filters off => th=0, exact v1 semantics).
12635                            if perturb_buf.is_none() {
12636                                perturb_buf = Some(e.zeros(d_vocab.max(n_vocab))?);
12637                            }
12638                            let mut q_row = e.clone_dtod(&dl_d)?; // retained q (penalized when on)
12639                            if pen_on {
12640                                let h = pen_hist_d.as_ref().unwrap();
12641                                let nh = h.len();
12642                                e.penalize_logits(
12643                                    &mut q_row,
12644                                    h,
12645                                    nh,
12646                                    sp.penalty_repeat,
12647                                    sp.penalty_freq,
12648                                    sp.penalty_present,
12649                                    d_vocab,
12650                                )?;
12651                            }
12652                            let rows0 = e.htod_i32(&[0])?;
12653                            let (mut th_d, mut z_d, mut mx_d) =
12654                                (e.zeros(1)?, e.zeros(1)?, e.zeros(1)?);
12655                            e.filter_stats(
12656                                &q_row, d_vocab, &rows0, &mut th_d, &mut z_d, &mut mx_d, d_vocab,
12657                                1, sp_temp, sp.top_k, sp.top_p, sp.min_p,
12658                            )?;
12659                            let (th, z, mx) =
12660                                (e.dtoh(&th_d)?[0], e.dtoh(&z_d)?[0], e.dtoh(&mx_d)?[0]);
12661                            let pb = perturb_buf.as_mut().unwrap();
12662                            e.gumbel_perturb_filtered(
12663                                &q_row, pb, d_vocab, sp_seed, sctr, sp_temp, mx, th,
12664                            )?;
12665                            sctr += 1;
12666                            draft_logits.push(q_row);
12667                            draft_stats.push((mx, th, z));
12668                            e.argmax_token_device(pb, d_vocab)?
12669                        } else {
12670                            e.argmax_token_device(&dl_d, d_vocab)?
12671                        };
12672                        let idx = e.dtoh_u32_one(&tok_d)?;
12673                        // #87 SENTINEL TRAP (eager twin — see the graph arm). Extra diagnostics
12674                        // here because the eager chain's operands are all readable: dl_d (the head
12675                        // logits row) and d_seed (this step's h_seed) name the first-NaN buffer.
12676                        if (idx as usize) >= d_vocab {
12677                            let dl_h = e.dtoh(&dl_d)?;
12678                            let dl_nan = dl_h.iter().filter(|v| v.is_nan()).count();
12679                            let seed_h = if chain_heads {
12680                                e.dtoh(chain_seeds.last().unwrap())?
12681                            } else {
12682                                e.dtoh(&d_seed)?
12683                            };
12684                            let seed_nan = seed_h.iter().filter(|v| v.is_nan()).count();
12685                            return Err(format!(
12686                                "draft(eager) argmax sentinel 0x{idx:08x} >= d_vocab {d_vocab} at \
12687                             round {round} j={j} pos={pos}: head-logits NaN {dl_nan}/{d_vocab}, \
12688                             step-seed NaN {seed_nan}/{n_embd} — refusing to dereference the \
12689                             embed row (#87 trap)"
12690                            )
12691                            .into());
12692                        }
12693                        let d = match &mtp.d2t {
12694                            Some(map) => map[idx as usize],
12695                            None => idx,
12696                        };
12697                        if sampled {
12698                            draft_idx.push(idx);
12699                        }
12700                        let draft_p = if p_min > 0.0 {
12701                            let p_d = e.prob_of_token_device(&dl_d, &tok_d, d_vocab)?;
12702                            Some(e.dtoh(&p_d)?[0])
12703                        } else {
12704                            None
12705                        };
12706                        if let Some(p) = draft_p.filter(|_| p_min > 0.0)
12707                            && p < p_min
12708                            && (j > 0 || (pmin0 && base0 == 1))
12709                        {
12710                            break;
12711                        }
12712                        draft.push(d);
12713                        if chain_heads {
12714                            chain_tokens.push(d);
12715                            chain_seeds.push(h_nextn);
12716                        } else {
12717                            e_tok = d;
12718                            d_seed = h_nextn;
12719                        }
12720                        // speculative advance; a chain the grammar can no longer follow (EOS
12721                        // proposed) ends here — the prefix already proposed still rides verify.
12722                        if dmask_live
12723                            && !constraint
12724                                .as_deref_mut()
12725                                .unwrap()
12726                                .draft_advance(d)
12727                                .map_err(|e2| format!("constraint: {e2}"))?
12728                        {
12729                            break;
12730                        }
12731                    }
12732                }
12733            }
12734            let k_round = draft.len();
12735
12736            ph_mark(&mut ph_draft, phase_on);
12737            // --- 2. VERIFY: one batched target forward. With a pending bonus, it rides as col 0
12738            //         (committing its KV/recur inside the SAME weight read); drafts follow. ---
12739            let verify_tokens: Vec<u32> = match pending {
12740                Some(b) => {
12741                    let mut v = Vec::with_capacity(k_round + 1);
12742                    v.push(b);
12743                    v.extend_from_slice(&draft);
12744                    v
12745                }
12746                None => draft.clone(),
12747            };
12748            let base = if pending.is_some() { 1 } else { 0 };
12749            // ckpt (REPLAY-FREE partial accept): retain per-layer state-rebuild inputs alongside
12750            // the verify. Pure buffer keep-alives + dtod clones — kernel work is unchanged.
12751            let mut ckpt = if spec_replay {
12752                None
12753            } else {
12754                Some(VerifyCkpt::new(self.layers.len()))
12755            };
12756            let (tlogits_d, vx) = {
12757                // The serial verify every non-fork round takes — the MTP route's
12758                // verify-graph door. The pool is None unless MEMRA_SPEC_VERIFY_GRAPH armed
12759                // a pool above, and then the walk replays the captured trunk instead of
12760                // re-issuing it launch by launch. `graph_round_ok` is the round's
12761                // headroom snapshot (see GRAPH_LAUNCH_MIN_FREE): below the floor the
12762                // round declines the pool exactly like an over-cap round and rides the
12763                // byte-identical eager walk — the `[spec]` suspension line above
12764                // already named the round.
12765                let vg_round = if verify_tokens.len() <= vg_t_cap && graph_round_ok {
12766                    vg_guard.as_mut().and_then(|g| g.as_mut())
12767                } else {
12768                    if let Some(g) = vg_guard.as_mut().and_then(|g| g.as_mut()) {
12769                        // The commit reads this flag to pick its arm; a round that declines
12770                        // the pool must not inherit a stale `true` from the round before it.
12771                        g.round_slab = false;
12772                    }
12773                    None
12774                };
12775                self.decode_step_t_core_vg(
12776                    e,
12777                    &verify_tokens,
12778                    pos,
12779                    &mut *cache,
12780                    embd_dev,
12781                    ckpt.as_mut(),
12782                    vg_round,
12783                )?
12784            };
12785
12786            if phase_sync {
12787                e.stream().synchronize()?;
12788            }
12789            ph_mark(&mut ph_verify, phase_on);
12790            // --- 3. GREEDY ACCEPT (walk prefix, stop at first mismatch) ---
12791            // DEVICE-ARGMAX ACCEPT: argmax every verify column ON DEVICE (same 2-pass kernels +
12792            // smallest-index tie-break as host argmax, argmax_gate-validated) and read back ONE
12793            // [T] u32 — replaces the T x n_vocab f32 dtoh + T host argmaxes per round.
12794            // t_pred[j] = target's greedy prediction for the slot after draft[j-1] (j>=1) or after
12795            // last_token (j==0). With a pending bonus, col 0 IS the prediction after last_token
12796            // (== the bonus), so every index shifts by `base` and last_pred is unused.
12797            let t_v = verify_tokens.len();
12798            let mut preds: Vec<u32> = Vec::new();
12799            if !sampled {
12800                for j in 0..t_v {
12801                    e.argmax_token_device_col(&tlogits_d, j, n_vocab, &mut preds_d, j)?;
12802                }
12803                preds = e.dtoh_u32(&preds_d)?; // <- the verify-GPU wait lands here
12804                // #87 SENTINEL TRAP, verify side: a sentinel pred becomes the round's bonus =
12805                // next round's last_token = the next chain's embed lookup. Catch it at the
12806                // source with the column named — an all-NaN VERIFY column implicates the
12807                // stage-split trunk (decode_step_t_core_ppn), not the draft head.
12808                if let Some(bad) = preds[..t_v].iter().position(|&p| (p as usize) >= n_vocab) {
12809                    let col = &tlogits_d.slice(bad * n_vocab..(bad + 1) * n_vocab);
12810                    let mut probe = e.zeros(n_vocab)?;
12811                    e.copy_view_into(&mut probe, 0, col, n_vocab)?;
12812                    let col_h = e.dtoh(&probe)?;
12813                    let col_nan = col_h.iter().filter(|v| v.is_nan()).count();
12814                    return Err(format!(
12815                        "verify argmax sentinel 0x{:08x} >= n_vocab {n_vocab} at round {round} \
12816                         col {bad}/{t_v} pos={pos}: verify-logits col NaN {col_nan}/{n_vocab} \
12817                         — the verify TRUNK produced a poisoned column (#87 trap). Run \
12818                         MEMRA_SPEC_NAN_SCAN=1 to name the layer that creates it (=2 to split \
12819                         that layer into attention and routed MoE). NOT the draft head, and NOT \
12820                         the PP stage split this message used to name: pp_cuts() returns None \
12821                         without MEMRA_PP_STAGES, so decode_step_t_core_ppn never runs unless \
12822                         that variable is set.",
12823                        preds[bad]
12824                    )
12825                    .into());
12826                }
12827            }
12828            ph_mark(&mut ph_wait, phase_on);
12829            let t_pred = |j: usize| -> u32 {
12830                if j == 0 && base == 0 {
12831                    last_pred
12832                } else {
12833                    // GREEDY-ONLY: `preds` is filled under `if !sampled` above. The debug print
12834                    // used to call this from the sampled arm and panicked the worker; it now goes
12835                    // through `debug_t_pred0`. Keep the strict index here — in the greedy walk an
12836                    // out-of-range pred is a real bug, not something to paper over.
12837                    debug_assert!(
12838                        !sampled,
12839                        "t_pred is greedy-only: `preds` is empty in the sampled arm"
12840                    );
12841                    preds[base + j - 1]
12842                }
12843            };
12844            let (n_acc, bonus) = if !sampled {
12845                let mut n_acc = 0usize;
12846                #[allow(clippy::needless_range_loop)]
12847                // allow: the explicit index loop keeps the offset arithmetic visible and aligned with the device-side indexing
12848                for j in 0..k_round {
12849                    if t_pred(j) == draft[j] {
12850                        n_acc += 1;
12851                    } else {
12852                        break;
12853                    }
12854                }
12855                // bonus = target's own token at the first non-accepted slot. n_acc in 0..=k; t_pred
12856                // is defined for j in 0..=k (j==0 -> last_logits, j>=1 -> col j-1, last col = k-1).
12857                (n_acc, t_pred(n_acc))
12858            } else {
12859                // --- SAMPLED ACCEPT (rejection sampling): u_j < p_j(x_j)/q_j(x_j) walk ---
12860                if col_buf.is_none() {
12861                    col_buf = Some(e.zeros(n_vocab)?);
12862                }
12863                // FILTERED p_j: per-verify-col stats (one batched filter_stats call), then the
12864                // filtered gather. j==0&&base==0 reads last_col (its own stats row appended).
12865                let mut pj = vec![0f32; k_round.max(1)];
12866                let mut col_stats: Vec<(f32, f32, f32)> = Vec::new(); // (max, th, z) per verify col used
12867                if k_round > 0 {
12868                    let mut ids: Vec<u32> = Vec::new();
12869                    let mut rows: Vec<i32> = Vec::new();
12870                    #[allow(clippy::needless_range_loop)]
12871                    // allow: the explicit index loop keeps the offset arithmetic visible and aligned with the device-side indexing
12872                    for j in 0..k_round {
12873                        if j > 0 || base == 1 {
12874                            ids.push(draft[j]);
12875                            rows.push((base + j) as i32 - 1);
12876                        }
12877                    }
12878                    if !ids.is_empty() {
12879                        let nr = rows.len();
12880                        // penalties: materialize the used columns into one contiguous penalized
12881                        // buffer (rows remapped 0..nr) so stats+gathers see the penalized p.
12882                        // penalties: materialize used columns contiguously, penalize all rows in
12883                        // one launch, and point stats+gathers at the penalized buffer (rows 0..nr).
12884                        let p_rows: Vec<i32> = if pen_on {
12885                            (0..nr as i32).collect()
12886                        } else {
12887                            rows.clone()
12888                        };
12889                        if pen_on {
12890                            if pcol_buf.as_ref().map(|b| b.len()).unwrap_or(0) < nr * n_vocab {
12891                                pcol_buf = Some(e.zeros(nr * n_vocab)?);
12892                            }
12893                            let pc = pcol_buf.as_mut().unwrap();
12894                            for (i2, &r) in rows.iter().enumerate() {
12895                                let c = r as usize;
12896                                e.copy_view_into(
12897                                    pc,
12898                                    i2 * n_vocab,
12899                                    &tlogits_d.slice(c * n_vocab..(c + 1) * n_vocab),
12900                                    n_vocab,
12901                                )?;
12902                            }
12903                            let h = pen_hist_d.as_ref().unwrap();
12904                            let nh = h.len();
12905                            e.penalize_logits_rows(
12906                                pc,
12907                                h,
12908                                nh,
12909                                sp.penalty_repeat,
12910                                sp.penalty_freq,
12911                                sp.penalty_present,
12912                                n_vocab,
12913                                nr,
12914                            )?;
12915                        }
12916                        let p_src: &CudaSlice<f32> = if pen_on {
12917                            pcol_buf.as_ref().unwrap()
12918                        } else {
12919                            &tlogits_d
12920                        };
12921                        let rowsd = e.htod_i32(&p_rows)?;
12922                        let (mut th_d, mut z_d, mut mx_d) =
12923                            (e.zeros(nr)?, e.zeros(nr)?, e.zeros(nr)?);
12924                        e.filter_stats(
12925                            p_src, n_vocab, &rowsd, &mut th_d, &mut z_d, &mut mx_d, n_vocab, nr,
12926                            sp_temp, sp.top_k, sp.top_p, sp.min_p,
12927                        )?;
12928                        let idsd = e.htod_u32_v(&ids)?;
12929                        let mut outd = e.zeros(nr)?;
12930                        e.softmax_gather_filtered(
12931                            p_src, n_vocab, &idsd, &rowsd, &th_d, &z_d, &mut outd, n_vocab, nr,
12932                            sp_temp,
12933                        )?;
12934                        let outv = e.dtoh(&outd)?;
12935                        let (thv, zv, mxv) = (e.dtoh(&th_d)?, e.dtoh(&z_d)?, e.dtoh(&mx_d)?);
12936                        let mut oi = 0usize;
12937                        #[allow(clippy::needless_range_loop)]
12938                        // allow: the explicit index loop keeps the offset arithmetic visible and aligned with the device-side indexing
12939                        for j in 0..k_round {
12940                            if j > 0 || base == 1 {
12941                                pj[j] = outv[oi];
12942                                oi += 1;
12943                            }
12944                        }
12945                        col_stats = (0..nr).map(|i| (mxv[i], thv[i], zv[i])).collect();
12946                    }
12947                    if base == 0 {
12948                        let lc: &CudaSlice<f32> = if pen_on {
12949                            if col_buf.is_none() {
12950                                col_buf = Some(e.zeros(n_vocab)?);
12951                            }
12952                            let cb = col_buf.as_mut().unwrap();
12953                            e.copy_into(
12954                                cb,
12955                                0,
12956                                last_col_logits
12957                                    .as_ref()
12958                                    .expect("sampled: last_col_logits unset"),
12959                                n_vocab,
12960                            )?;
12961                            let h = pen_hist_d.as_ref().unwrap();
12962                            let nh = h.len();
12963                            e.penalize_logits(
12964                                cb,
12965                                h,
12966                                nh,
12967                                sp.penalty_repeat,
12968                                sp.penalty_freq,
12969                                sp.penalty_present,
12970                                n_vocab,
12971                            )?;
12972                            col_buf.as_ref().unwrap()
12973                        } else {
12974                            last_col_logits
12975                                .as_ref()
12976                                .expect("sampled: last_col_logits unset")
12977                        };
12978                        let rows0 = e.htod_i32(&[0])?;
12979                        let (mut th_d, mut z_d, mut mx_d) = (e.zeros(1)?, e.zeros(1)?, e.zeros(1)?);
12980                        e.filter_stats(
12981                            lc, n_vocab, &rows0, &mut th_d, &mut z_d, &mut mx_d, n_vocab, 1,
12982                            sp_temp, sp.top_k, sp.top_p, sp.min_p,
12983                        )?;
12984                        let idsd = e.htod_u32_v(&[draft[0]])?;
12985                        let mut outd = e.zeros(1)?;
12986                        e.softmax_gather_filtered(
12987                            lc, n_vocab, &idsd, &rows0, &th_d, &z_d, &mut outd, n_vocab, 1, sp_temp,
12988                        )?;
12989                        pj[0] = e.dtoh(&outd)?[0];
12990                        last_col_stats =
12991                            Some((e.dtoh(&mx_d)?[0], e.dtoh(&th_d)?[0], e.dtoh(&z_d)?[0]));
12992                    }
12993                }
12994                // q source: the graph arms (single-head AND chain) retained the head logits
12995                // in the persistent q_slots; the eager arm in per-round draft_logits clones.
12996                // Same raw-logit values either way. FILTERED q_j: stats from draft_stats
12997                // (eager pushes in-chain; the graph arms compute them post-replay from the
12998                // retained q with the same filter_stats program — bit-identical to the
12999                // in-graph stats that shaped the draw, keeping ONE accept path).
13000                let q_bufs: &[CudaSlice<f32>] = if dctx.graph_s.is_some() || dctx.chain_s.is_some()
13001                {
13002                    &dctx.q_slots
13003                } else {
13004                    &draft_logits
13005                };
13006                let mut n_acc = 0usize;
13007                for j in 0..k_round {
13008                    let (qmx, qth, qz) = draft_stats[j];
13009                    let idsd = e.htod_u32_v(&[draft_idx[j]])?;
13010                    let rowsd = e.htod_i32(&[0])?;
13011                    let thd = e.htod(&[qth])?;
13012                    let zd = e.htod(&[qz])?;
13013                    let _ = qmx;
13014                    let mut outd = e.zeros(1)?;
13015                    e.softmax_gather_filtered(
13016                        &q_bufs[j], d_vocab, &idsd, &rowsd, &thd, &zd, &mut outd, d_vocab, 1,
13017                        sp_temp,
13018                    )?;
13019                    let qj = e.dtoh(&outd)?[0];
13020                    let u = host_u01(sp_seed, uctr);
13021                    uctr += 1;
13022                    let accept = (u as f64) * (qj as f64) < pj[j] as f64;
13023                    // SKEY PROBE: q == 0 for the token the draft actually proposed is the
13024                    // exactness signature (see `skey_probe`). Impossible when the draft was
13025                    // drawn from the same filtered distribution the verify reconstructs here;
13026                    // `u * 0 < p` makes it an UNCONDITIONAL accept whenever p > 0.
13027                    if skey_probe() && qj == 0.0 {
13028                        eprintln!(
13029                            "[skey] EXACTNESS q=0 round={round} j={j} draft_tok={} \
13030                             draft_idx={} p={:e} u={u} accepted={} th_z={:?}",
13031                            draft[j], draft_idx[j], pj[j], accept as u8, draft_stats[j],
13032                        );
13033                    }
13034                    if accept {
13035                        n_acc += 1;
13036                    } else {
13037                        break;
13038                    }
13039                }
13040                let bonus = if n_acc == k_round {
13041                    // FULL ACCEPT: bonus ~ FILTERED softmax at the last verify column.
13042                    let col = base + k_round - 1;
13043                    let cb = col_buf.as_mut().unwrap();
13044                    e.copy_view_into(
13045                        cb,
13046                        0,
13047                        &tlogits_d.slice(col * n_vocab..(col + 1) * n_vocab),
13048                        n_vocab,
13049                    )?;
13050                    if pen_on {
13051                        let h = pen_hist_d.as_ref().unwrap();
13052                        let nh = h.len();
13053                        e.penalize_logits(
13054                            cb,
13055                            h,
13056                            nh,
13057                            sp.penalty_repeat,
13058                            sp.penalty_freq,
13059                            sp.penalty_present,
13060                            n_vocab,
13061                        )?;
13062                    }
13063                    if perturb_buf.is_none() {
13064                        perturb_buf = Some(e.zeros(d_vocab.max(n_vocab))?);
13065                    }
13066                    // STATS MUST COME FROM THIS COLUMN (bug fix 2026-08-05, lane/sampler-
13067                    // truncation-fix; receipts research/sampfix-20260805/). The old code reused
13068                    // `col_stats.last()` here, which is ALWAYS the wrong row: the gathered set
13069                    // covers verify columns 0..=(base+k_round-2) (rows pushed as base+j-1), while
13070                    // the full-accept bonus samples column base+k_round-1 — exactly ONE PAST the
13071                    // last gathered column, in both base arms. `th` is a threshold in e-units of
13072                    // its OWN row's max, so feeding a neighbour's (row_max, th) into
13073                    // gumbel_perturb_filtered mis-scales every e0 = exp((x-row_max)/T). When the
13074                    // donor column's peak is higher by more than T*ln(1/th), EVERY id fails
13075                    // `e0 >= th`, the whole perturbed row becomes -3.4e38, and the 2-pass argmax
13076                    // falls through to its smallest-index tie-break => token id 0 ("!") spliced
13077                    // mid-word. Fragility is ordered by how large th is: min_p pins th = min_p
13078                    // (0.05 => trigger at delta > 2.4 at T=0.8, fires constantly), top_p's
13079                    // mass-boundary th is smaller, top_k's k-th-largest th smaller still — which
13080                    // is why the head-to-head matrix saw min_p and top_p corrupt while top_k-only
13081                    // stayed clean. The pure-temp default regime is immune (th == 0 masks nothing,
13082                    // and row_max is unused once nothing is masked), so this fix is a byte-level
13083                    // no-op for the untruncated serve default. One extra one-block filter_stats
13084                    // per full-accept round is the whole cost.
13085                    let (mx, th) = {
13086                        let rows0 = e.htod_i32(&[0])?;
13087                        let (mut th_d, mut z_d, mut mx_d) = (e.zeros(1)?, e.zeros(1)?, e.zeros(1)?);
13088                        let cb0 = col_buf.as_ref().unwrap();
13089                        e.filter_stats(
13090                            cb0, n_vocab, &rows0, &mut th_d, &mut z_d, &mut mx_d, n_vocab, 1,
13091                            sp_temp, sp.top_k, sp.top_p, sp.min_p,
13092                        )?;
13093                        (e.dtoh(&mx_d)?[0], e.dtoh(&th_d)?[0])
13094                    };
13095                    let pb = perturb_buf.as_mut().unwrap();
13096                    let cb2 = col_buf.as_ref().unwrap();
13097                    e.gumbel_perturb_filtered(cb2, pb, n_vocab, sp_seed, sctr, sp_temp, mx, th)?;
13098                    sctr += 1;
13099                    let td = e.argmax_token_device(pb, n_vocab)?;
13100                    e.dtoh_u32_one(&td)?
13101                } else {
13102                    // REJECT at n_acc: bonus ~ norm(max(0, softmax_T(p) - softmax_T(q))).
13103                    let cb = col_buf.as_mut().unwrap();
13104                    if n_acc > 0 || base == 1 {
13105                        let col = base + n_acc - 1;
13106                        e.copy_view_into(
13107                            cb,
13108                            0,
13109                            &tlogits_d.slice(col * n_vocab..(col + 1) * n_vocab),
13110                            n_vocab,
13111                        )?;
13112                    } else {
13113                        let lc = last_col_logits.as_ref().unwrap();
13114                        e.copy_into(cb, 0, lc, n_vocab)?;
13115                    }
13116                    if pen_on {
13117                        let h = pen_hist_d.as_ref().unwrap();
13118                        let nh = h.len();
13119                        e.penalize_logits(
13120                            cb,
13121                            h,
13122                            nh,
13123                            sp.penalty_repeat,
13124                            sp.penalty_freq,
13125                            sp.penalty_present,
13126                            n_vocab,
13127                        )?;
13128                    }
13129                    let cb2 = col_buf.as_ref().unwrap();
13130                    let sc = sctr;
13131                    sctr += 1;
13132                    // p-stats for the reject column: from col_stats when the col was gathered,
13133                    // else (j==0&&base==0) from last_col_stats.
13134                    let p_stats = if n_acc > 0 || base == 1 {
13135                        // col index within the gathered set == number of gathered cols before n_acc
13136                        let gi = if base == 1 { n_acc } else { n_acc - 1 };
13137                        col_stats.get(gi).copied().unwrap_or({
13138                            (0.0, 0.0, 1.0) // unreachable: gathered cols always cover the reject slot
13139                        })
13140                    } else {
13141                        last_col_stats.expect("sampled: last_col_stats unset at reject")
13142                    };
13143                    let q_stats = draft_stats[n_acc];
13144                    if let Some(map) = &d2t_dev {
13145                        if q_full_buf.is_none() {
13146                            q_full_buf = Some(e.zeros(n_vocab)?);
13147                        }
13148                        let qf = q_full_buf.as_mut().unwrap();
13149                        e.scatter_trim_logits(&q_bufs[n_acc], map, qf, d_vocab, n_vocab)?;
13150                        let qf2 = q_full_buf.as_ref().unwrap();
13151                        e.residual_sample_filtered(
13152                            cb2,
13153                            Some(qf2),
13154                            n_vocab,
13155                            sp_temp,
13156                            sp_seed,
13157                            sc,
13158                            p_stats,
13159                            q_stats,
13160                            &mut sample_tok,
13161                        )?;
13162                    } else {
13163                        e.residual_sample_filtered(
13164                            cb2,
13165                            Some(&q_bufs[n_acc]),
13166                            n_vocab,
13167                            sp_temp,
13168                            sp_seed,
13169                            sc,
13170                            p_stats,
13171                            q_stats,
13172                            &mut sample_tok,
13173                        )?;
13174                    }
13175                    e.dtoh_u32(&sample_tok)?[0]
13176                };
13177                (
13178                    n_acc,
13179                    guard_vocab_token(
13180                        bonus,
13181                        n_vocab,
13182                        &format!("sampled verify bonus at round {round} pos={pos} n_acc={n_acc}"),
13183                    )?,
13184                )
13185            };
13186            // --- 3b. GRAMMAR TRUNCATION (constrained spec, 2026-08-03): the grammar is
13187            // an extra rejection rule AFTER the exactness verify (the batched-verify-twins
13188            // ordering). Walk the accepted drafts through the grammar in commit order; the
13189            // first illegal token truncates acceptance at its slot, and that slot's emission
13190            // is recomputed as the MASKED argmax of the target's own verify column — token-
13191            // identical to constrained plain greedy decode (an unmasked argmax that is
13192            // grammar-legal IS the masked argmax: masking only removes competitors). The
13193            // column D2H (~1MB) is paid only when a cut fires — the tight-grammar cost,
13194            // measured in acceptance numbers, never hidden.
13195            let (n_acc, bonus) = match constraint.as_deref_mut() {
13196                None => (n_acc, bonus),
13197                Some(c) => {
13198                    fn ce(e2: String) -> Box<dyn std::error::Error> {
13199                        format!("constraint: {e2}").into()
13200                    }
13201                    let mut na = n_acc;
13202                    let mut cut = false;
13203                    for (j, &d) in draft.iter().enumerate().take(n_acc) {
13204                        if c.is_allowed(d).map_err(ce)? {
13205                            c.consume(d).map_err(ce)?;
13206                        } else {
13207                            na = j;
13208                            cut = true;
13209                            dm_cut_tokens += n_acc - j;
13210                            break;
13211                        }
13212                    }
13213                    if cut {
13214                        dm_cuts += 1;
13215                    }
13216                    let mut bo = bonus;
13217                    if cut || !c.is_allowed(bo).map_err(ce)? {
13218                        let mut row = if na == 0 && base == 0 {
13219                            init_logits_host
13220                                .clone()
13221                                .ok_or("constraint: init logits missing (round-0 cut)")?
13222                        } else {
13223                            e.dtoh_view(
13224                                &tlogits_d.slice((base + na - 1) * n_vocab..(base + na) * n_vocab),
13225                            )?
13226                        };
13227                        c.mask_logits(&mut row).map_err(ce)?;
13228                        bo = argmax(&row) as u32;
13229                    }
13230                    c.consume(bo).map_err(ce)?;
13231                    (na, bo)
13232                }
13233            };
13234            total_drafted += k_round;
13235            total_accepted += n_acc;
13236            if let Some(t) = sess_telem {
13237                // Greedy, rejection-sampling, and grammar truncation all converge here after
13238                // the accept decision is already on host. Fixed-size relaxed atomics only.
13239                t.record_round(k_round, n_acc);
13240            }
13241            if spec_stats {
13242                st_len_hist[k_round] += 1;
13243                #[allow(clippy::needless_range_loop)]
13244                // allow: the explicit index loop keeps the offset arithmetic visible and aligned with the device-side indexing
13245                for j in 0..k_round {
13246                    st_drafted[j] += 1;
13247                }
13248                #[allow(clippy::needless_range_loop)]
13249                // allow: the explicit index loop keeps the offset arithmetic visible and aligned with the device-side indexing
13250                for j in 0..n_acc {
13251                    st_accepted[j] += 1;
13252                }
13253                if n_acc == k_round {
13254                    st_full += 1;
13255                }
13256            }
13257
13258            if debug_spec {
13259                eprintln!(
13260                    "[R{round}] pos={pos} out_len={} last_tok={last_token} draft={draft:?} n_acc={n_acc} bonus={bonus} t_pred0={}",
13261                    out.len(),
13262                    // NOT `t_pred(0)`: `preds` is filled only under `if !sampled` above, so on a
13263                    // sampled request round >= 1 (base == 1) indexed an EMPTY vector and PANICKED
13264                    // the GPU worker thread — a debug flag that killed the exact regime you would
13265                    // set it to investigate. See `debug_t_pred0`.
13266                    debug_t_pred0(sampled, base, last_pred, &preds)
13267                );
13268            }
13269
13270            // --- 4. COMMIT: draft[0..n_acc] then bonus (n_acc + 1 tokens) ---
13271            let commit_started = std::time::Instant::now();
13272            // SESSION MODE: every accepted column is already in the CACHE — `out` must carry all
13273            // of them (overshoot past max_new included) or `committed` under-counts the cache rows
13274            // and the next turn's continuation seeds one token off (gate-caught 2026-07-05). The
13275            // single-shot path keeps the cap (its caller truncates + drops the cache anyway).
13276            #[allow(clippy::needless_range_loop)]
13277            // allow: the explicit index loop keeps the offset arithmetic visible and aligned with the device-side indexing
13278            for j in 0..n_acc {
13279                if !session_mode && out.len() >= max_new {
13280                    break;
13281                }
13282                out.push(draft[j]);
13283            }
13284            if pen_on {
13285                pen_hist.extend_from_slice(&draft[0..n_acc]);
13286                pen_hist.push(bonus);
13287            }
13288            let bonus_emitted = session_mode || out.len() < max_new;
13289            if bonus_emitted {
13290                out.push(bonus);
13291            }
13292            last_token = bonus;
13293
13294            // --- 5. ROLLBACK + advance (§C) ---
13295            if n_acc == k_round && !spec_replay {
13296                // FULL ACCEPT, BONUS FOLD: all verify columns (pending? + drafts) are committed in
13297                // cache; the NEW bonus stays PENDING for the next round's verify batch — NO extra
13298                // T=1 trunk pass. The next draft chain seeds from the MTP block's h_nextn at the
13299                // bonus position: one MTP-block pass (~1/33 trunk cost) replaces the trunk read.
13300                // last_pred is dead in the pending path (t_pred reads verify col 0).
13301                //
13302                // PERSISTENT DRAFT KV, full-accept fill: the chain covered last_token +
13303                // draft[0..k_round-2] as INPUTS (slots P..P'-2); draft[k_round-1] (slot P'-1) was
13304                // only ever an output, so its entry is MISSING. Fill it from vh_seed — its EXACT
13305                // trunk hidden (the last verify column). set_len first: a p-min break may have
13306                // left one extra chain append at that slot. Partial accepts need NO fill (the
13307                // chain already covered every accepted position; round-start set_len truncates).
13308                self.restore_step_tp_kv_verified_prefix(e, &mut *cache, &snap, t_v, false)?;
13309                let mut vh_seed = e.zeros(n_embd)?;
13310                e.copy_view_into(
13311                    &mut vh_seed,
13312                    0,
13313                    &vx.slice((t_v - 1) * n_embd..t_v * n_embd),
13314                    n_embd,
13315                )?;
13316                if refresh {
13317                    // TRUE-HIDDEN REFRESH (2026-07-03, the HANDOVER-listed acceptance lever):
13318                    // overwrite ALL committed positions' scratch entries with K/V from their EXACT
13319                    // verify hiddens — the reference engine's mtp_update fills from true hiddens;
13320                    // the full stack (vx) is already resident from the verify. Replaces both the
13321                    // chain-approximate entries AND the old last-token-only fill. Acceptance-only
13322                    // (draft attention quality); exactness stays the verify's job.
13323                    scratch.set_len(e, pos)?;
13324                    // PREDECESSOR pairing: row i gets vx[i-1]; row 0 the carried fill_prev
13325                    // (hidden of the last committed row before this verify batch).
13326                    let mut vxs = e.zeros(t_v * n_embd)?;
13327                    e.copy_into(&mut vxs, 0, &fill_prev, n_embd)?;
13328                    if t_v > 1 {
13329                        e.copy_view_into(
13330                            &mut vxs,
13331                            n_embd,
13332                            &vx.slice(0..(t_v - 1) * n_embd),
13333                            (t_v - 1) * n_embd,
13334                        )?;
13335                    }
13336                    self.mtp_kv_fill_all(e, &verify_tokens, &vxs, pos, &mut *scratch, embd_dev)?;
13337                } else {
13338                    scratch.set_len(e, pos + base + k_round - 1)?;
13339                    // predecessor of the last draft = verify col t_v-2 (or fill_prev at t_v==1)
13340                    let mut hp = e.zeros(n_embd)?;
13341                    if t_v >= 2 {
13342                        e.copy_view_into(
13343                            &mut hp,
13344                            0,
13345                            &vx.slice((t_v - 2) * n_embd..(t_v - 1) * n_embd),
13346                            n_embd,
13347                        )?;
13348                    } else {
13349                        e.copy_into(&mut hp, 0, &fill_prev, n_embd)?;
13350                    }
13351                    self.mtp_kv_fill_all(
13352                        e,
13353                        &[draft[k_round - 1]],
13354                        &hp,
13355                        pos + base + k_round - 1,
13356                        &mut *scratch,
13357                        embd_dev,
13358                    )?;
13359                }
13360                // REFERENCE SEEDING: no pseudo pass — the next chain's step 0 IS the
13361                // reference's (id_last, h_prev) draft row; it appends the bonus's scratch
13362                // entry itself. Seed = TRUE hidden of the bonus's predecessor (last verify
13363                // col). Saves one MTP-block pass per round on top of the pairing fix.
13364                e.copy_into(&mut h_seed_buf, 0, &vh_seed, n_embd)?;
13365                e.copy_into(&mut fill_prev, 0, &vh_seed, n_embd)?;
13366                pending = Some(bonus);
13367                if debug_spec {
13368                    eprintln!("  -> FULL ACCEPT (bonus pending, prev-h seed)");
13369                }
13370            } else if !spec_replay && base + n_acc >= 1 {
13371                // PARTIAL ACCEPT, REPLAY-FREE (2026-07-03 — the profiled #1 long-ctx spec cost):
13372                // the verify's first j = base+n_acc columns ARE the committed sequence, computed
13373                // bit-identically to eager (decode-exact contract) — so KEEP them: KV truncates to
13374                // pos+j, recurrent state rebuilds from the VerifyCkpt (same-kernel gdn prefix
13375                // re-run / pure state-clone restore), and the bonus stays PENDING exactly like the
13376                // full-accept path — the legacy duplicate trunk replay is gone. The next chain
13377                // seeds from the MTP pseudo-hidden of the bonus, whose seed = the TRUE verify
13378                // hidden of its predecessor (col j-1) — same one-hop pseudo structure as full
13379                // accept (never compounds: the next verify recomputes true hiddens for all
13380                // committed columns).
13381                let j = base + n_acc;
13382                // VERIFY-GRAPH SLAB COMMIT: when the captured trunk ran, the linear layers'
13383                // column stash was written into the graphs ctx's persistent slabs as in-graph
13384                // memcpy nodes, NOT into the per-column VerifyCkpt the cols arm reads — so the
13385                // commit must take the slab twin (same semantics, slab-addressed sources). The
13386                // ctx states which of the two this round produced via `round_slab`; trusting the
13387                // flag rather than the env keeps a round that fell back to the eager walk (a
13388                // capture that declined, a t the pool never captured) on the cols arm.
13389                let slab_commit = vg_guard
13390                    .as_ref()
13391                    .and_then(|g| g.as_ref())
13392                    .map(|g| g.round_slab)
13393                    .unwrap_or(false);
13394                if slab_commit {
13395                    self.dspark_commit_prefix_slab(
13396                        e,
13397                        &mut *cache,
13398                        &snap,
13399                        vg_guard
13400                            .as_ref()
13401                            .and_then(|g| g.as_ref())
13402                            .expect("slab_commit implies a graphs ctx"),
13403                        j,
13404                    )?;
13405                } else {
13406                    self.commit_verified_prefix(e, &mut *cache, &snap, ckpt.as_ref().unwrap(), j)?;
13407                }
13408                let mut seed = e.zeros(n_embd)?;
13409                e.copy_view_into(
13410                    &mut seed,
13411                    0,
13412                    &vx.slice((j - 1) * n_embd..j * n_embd),
13413                    n_embd,
13414                )?;
13415                // Draft scratch: TRUE-HIDDEN REFRESH of the committed prefix (see the full-accept
13416                // branch); without it the chain entries stand and only the tail truncates. Either
13417                // way len ends at pos+j so the pseudo append lands at the bonus's slot pos+j
13418                // (persistent mode), rope pos+j+1 (chain convention).
13419                if refresh {
13420                    scratch.set_len(e, pos)?;
13421                    let mut vxs = e.zeros(j * n_embd)?;
13422                    e.copy_into(&mut vxs, 0, &fill_prev, n_embd)?;
13423                    if j > 1 {
13424                        e.copy_view_into(
13425                            &mut vxs,
13426                            n_embd,
13427                            &vx.slice(0..(j - 1) * n_embd),
13428                            (j - 1) * n_embd,
13429                        )?;
13430                    }
13431                    self.mtp_kv_fill_all(
13432                        e,
13433                        &verify_tokens[0..j],
13434                        &vxs,
13435                        pos,
13436                        &mut *scratch,
13437                        embd_dev,
13438                    )?;
13439                } else {
13440                    scratch.set_len(e, pos + j)?;
13441                }
13442                // REFERENCE SEEDING (see the full-accept branch): seed = TRUE hidden of the
13443                // bonus's predecessor (verify col j-1); no pseudo pass.
13444                e.copy_into(&mut h_seed_buf, 0, &seed, n_embd)?;
13445                e.copy_into(&mut fill_prev, 0, &seed, n_embd)?;
13446                pending = Some(bonus);
13447                if debug_spec {
13448                    eprintln!("  -> PARTIAL(replay-free j={j}, bonus pending, prev-h seed)");
13449                }
13450            } else if !spec_replay {
13451                // ZERO ROUND FOLD (2026-07-10, verify-cost target #3): base+n_acc == 0 — a
13452                // pending-less round where nothing was accepted (PMIN0 zero-draft chains after a
13453                // replay/commit, or plain 0-accept rounds at round 0). The old path replayed
13454                // [bonus] through a FULL m=1 trunk+head forward (the 489us full-vocab head pass
13455                // measured at ~0.75/round on PMIN0 configs). Instead: restore the pre-round
13456                // snapshot and let the bonus ride the NEXT round's verify as col 0 — the existing
13457                // base=1 pending machinery, bit-identical by the decode-exact verify contract.
13458                // Seed: the bonus's predecessor is the last COMMITTED token, whose hidden
13459                // fill_prev already carries (same seeding as the 1-token-replay case it replaces).
13460                cache.rollback(e, &snap, 0)?;
13461                scratch.set_len(e, pos)?;
13462                e.copy_into(&mut h_seed_buf, 0, &fill_prev, n_embd)?;
13463                pending = Some(bonus);
13464                if debug_spec {
13465                    eprintln!("  -> ZERO-ROUND FOLD (bonus pending, fill_prev seed)");
13466                }
13467            } else {
13468                // PARTIAL ACCEPT, LEGACY REPLAY (seam MEMRA_SPEC_REPLAY=1 — or j==0: nothing of
13469                // this round survives, only possible before the first pending exists, ~round 0):
13470                // restore EVERYTHING to the pre-round snapshot (KV truncate to pos + recur
13471                // restore), then replay the committed prefix pending? ++ draft[0..n_acc] ++
13472                // [bonus] as ONE batched T forward — single weight read, bit-identical to greedy
13473                // (the verify-all-columns path is the same math). Commits the bonus with a TRUE
13474                // trunk hidden.
13475                cache.rollback(e, &snap, 0)?; // accept_len=0: KV len = pos, recur = snapshot
13476                let mut replay: Vec<u32> = Vec::with_capacity(base + n_acc + 1);
13477                if let Some(b) = pending.take() {
13478                    replay.push(b);
13479                }
13480                replay.extend_from_slice(&draft[0..n_acc]);
13481                replay.push(bonus);
13482                // Full-stack forward (decode_step_t_core = decode_step_t_h_emb_dev's body):
13483                // Predecessor pairing seeds from the PREDECESSOR row (col len-2) — the same-row path takes the
13484                // last col exactly as before (byte-identical to the old _h_emb_dev call).
13485                let (rl_d, rx) = if self.batched_serving_numeric_class() {
13486                    let mut logits = Vec::with_capacity(replay.len() * n_vocab);
13487                    let mut hidden = e.uninit(replay.len() * n_embd)?;
13488                    for (row, &token) in replay.iter().enumerate() {
13489                        let (row_logits, row_hidden) =
13490                            self.spec_target_step_h(e, token, &mut *cache)?;
13491                        logits.extend_from_slice(&row_logits);
13492                        e.dtod_copy_into(&row_hidden, &mut hidden, row * n_embd)?;
13493                    }
13494                    (e.htod(&logits)?, hidden)
13495                } else {
13496                    self.decode_step_t_core(e, &replay, pos, &mut *cache, embd_dev, None)?
13497                };
13498                // last_pred = argmax of the LAST column's logits (predicts the token after `bonus`)
13499                // — device argmax + one 4-byte read instead of the full-vocab column dtoh.
13500                e.argmax_token_device_col(&rl_d, replay.len() - 1, n_vocab, &mut preds_d, 0)?;
13501                last_pred = guard_vocab_token(
13502                    e.dtoh_u32(&preds_d)?[0],
13503                    n_vocab,
13504                    &format!("replay last_pred at round {round} pos={pos}"),
13505                )?;
13506                if sampled {
13507                    let lr0 = replay.len();
13508                    let lc = last_col_logits
13509                        .as_mut()
13510                        .expect("sampled: last_col_logits unset");
13511                    e.copy_view_into(
13512                        lc,
13513                        0,
13514                        &rl_d.slice((lr0 - 1) * n_vocab..lr0 * n_vocab),
13515                        n_vocab,
13516                    )?;
13517                }
13518                let lr = replay.len();
13519                if lr >= 2 {
13520                    e.copy_view_into(
13521                        &mut h_seed_buf,
13522                        0,
13523                        &rx.slice((lr - 2) * n_embd..(lr - 1) * n_embd),
13524                        n_embd,
13525                    )?;
13526                } else {
13527                    // 1-token replay (round-0 miss): the bonus's predecessor is the OLD
13528                    // last_token, whose own-row hidden fill_prev still holds.
13529                    e.copy_into(&mut h_seed_buf, 0, &fill_prev, n_embd)?;
13530                }
13531                // the bonus is COMMITTED here — it becomes the last committed row.
13532                let mut rh_last = e.zeros(n_embd)?;
13533                e.copy_view_into(
13534                    &mut rh_last,
13535                    0,
13536                    &rx.slice((lr - 1) * n_embd..lr * n_embd),
13537                    n_embd,
13538                )?;
13539                e.copy_into(&mut fill_prev, 0, &rh_last, n_embd)?;
13540                if debug_spec {
13541                    eprintln!("  -> PARTIAL(replay={replay:?}), next_pred={last_pred}");
13542                }
13543            }
13544            if anatomy_on {
13545                // Commit/rollback is normally asynchronous on the primary/head stream. Bound it
13546                // only for this diagnostic so it does not disappear into the following draft's
13547                // first token readback.
13548                e.stream().synchronize()?;
13549                ph_commit += commit_started.elapsed().as_secs_f64();
13550            }
13551            // adaptive-K update (host math, zero syncs): next round drafts accepted-run + 1,
13552            // clamped to [floor(pos), k_cap]. cache.pos is post-rollback here (the round's
13553            // final position — the floor's position key reads the committed depth). Burst
13554            // rounds (`continue` above) draft the captured fixed depth and skip this, exactly
13555            // like gemma's burst arm.
13556            if adapt {
13557                let fl_now = floor_at(cache.pos);
13558                kc = (n_acc + 1).clamp(fl_now.min(k_cap), k_cap);
13559            }
13560            ph_mark(&mut ph_rest, phase_on);
13561            if let Some(t0) = round_t0 {
13562                let ms = t0.elapsed().as_secs_f64() * 1e3;
13563                ROUND_MS.fetch_add((ms * 1e3) as u64, std::sync::atomic::Ordering::Relaxed);
13564                let n = ROUND_N.fetch_add(1, std::sync::atomic::Ordering::Relaxed) + 1;
13565                if n.is_multiple_of(32) {
13566                    eprintln!(
13567                        "[spec-round] rounds={n} avg round wall={:.2} ms (emitted={} drafted so far)",
13568                        ROUND_MS.load(std::sync::atomic::Ordering::Relaxed) as f64 / 1e3 / n as f64,
13569                        out.len()
13570                    );
13571                }
13572            }
13573            round += 1;
13574            // sse-cadence: this round's accepted drafts + bonus are committed (out is
13575            // append-only past step 4) — flush at round cadence.
13576            keep_going = flush_commit(&mut on_commit, &out, &mut flushed);
13577        }
13578        // sse-cadence: nothing below appends to `out`; flush any remainder (defensive).
13579        // (verdict ignored — the burst is over either way; the session tail runs unchanged.)
13580        let _ = flush_commit(&mut on_commit, &out, &mut flushed);
13581
13582        if spec_stats {
13583            let per_slot: Vec<String> = (0..k)
13584                .map(|j| {
13585                    if st_drafted[j] > 0 {
13586                        format!(
13587                            "{}/{}={:.3}",
13588                            st_accepted[j],
13589                            st_drafted[j],
13590                            st_accepted[j] as f64 / st_drafted[j] as f64
13591                        )
13592                    } else {
13593                        "0/0".into()
13594                    }
13595                })
13596                .collect();
13597            let acc = if total_drafted > 0 {
13598                total_accepted as f64 / total_drafted as f64
13599            } else {
13600                0.0
13601            };
13602            eprintln!(
13603                "[spec-stats] rounds={round} full_accept={st_full} len_hist={st_len_hist:?} \
13604                       per_slot=[{}] total={total_accepted}/{total_drafted}={acc:.3} \
13605                       tok_per_round={:.3}",
13606                per_slot.join(" "),
13607                (total_accepted + round) as f64 / round.max(1) as f64
13608            );
13609        }
13610        if constraint.is_some() {
13611            eprintln!(
13612                "[draft-mask] mask_rounds={dm_rounds} clone_total={:.3}ms \
13613                 clone_per_round={:.4}ms gram_cuts={dm_cuts}/{round} cut_tokens={dm_cut_tokens}",
13614                dm_clone_ns as f64 / 1e6,
13615                dm_clone_ns as f64 / 1e6 / dm_rounds.max(1) as f64
13616            );
13617        }
13618        if phase_on {
13619            let tot = ph_draft + ph_verify + ph_wait + ph_rest;
13620            eprintln!(
13621                "[spec-phase] draft={:.1}ms ({:.1}%) verify-issue={:.1}ms ({:.1}%) verify-wait={:.1}ms ({:.1}%) commit-host={:.1}ms ({:.1}%) rounds={round}",
13622                ph_draft * 1e3,
13623                ph_draft / tot * 100.0,
13624                ph_verify * 1e3,
13625                ph_verify / tot * 100.0,
13626                ph_wait * 1e3,
13627                ph_wait / tot * 100.0,
13628                ph_rest * 1e3,
13629                ph_rest / tot * 100.0
13630            );
13631        }
13632        if anatomy_on {
13633            let rounds_f = round.max(1) as f64;
13634            let other = (ph_rest - ph_commit).max(0.0);
13635            eprintln!(
13636                "[spec-anatomy] per-round draft={:.3}ms pp-verify={:.3}ms \
13637                 verify-accept={:.3}ms commit-rollback={:.3}ms other={:.3}ms rounds={round}",
13638                ph_draft * 1e3 / rounds_f,
13639                ph_verify * 1e3 / rounds_f,
13640                ph_wait * 1e3 / rounds_f,
13641                ph_commit * 1e3 / rounds_f,
13642                other * 1e3 / rounds_f,
13643            );
13644        }
13645        // SESSION TAIL: leave the session in the exact invariant the next turn's suffix prime
13646        // expects — every row in `committed` has trunk KV/recur state AND an exact draft-KV row.
13647        // Park the draft-graph ctx back on the session (the serve-burst fixed-cost fix): the next
13648        // burst replays instead of recapturing. Error paths (`?` above) drop it — recaptured then.
13649        if let Some(slot) = sess_draft_slot.take() {
13650            *slot = Some(dctx);
13651        }
13652        let t_rounds = t_ent.elapsed();
13653        if let Some((committed, last_h, next_pred_slot, sctr_slot, uctr_slot)) = sess_tail.take() {
13654            // NEXT BURST'S BOUNDARY TOKEN (lane/sampled-spec-quality, Item 1). Greedy stashes
13655            // the argmax `last_pred` exactly as before (byte contract). SAMPLED draws the token
13656            // HERE, where the sampler, the session Philox counters and the penalty window are
13657            // all live and the boundary logits row still exists — that is the "make the state
13658            // available" half of the fix; the consuming burst then just emits it. `sctr` is
13659            // written to the session BELOW the draws so the advance is never lost.
13660            *next_pred_slot = Some(last_pred);
13661            let sample_boundary = sampled && constraint.is_none() && spec_sampled_boundary_on();
13662            let mut stashed_pending = false;
13663            if let Some(b) = pending.take() {
13664                if !sampled {
13665                    // PENDING-CARRY (2026-08-01): stash the bonus on the session instead of
13666                    // committing it with a solo T=1 pass — the next empty-suffix greedy burst
13667                    // consumes it as round-0 verify col 0 (a plain round edge; the old tail
13668                    // commit + next burst's init feed were 11.6+11.5ms solo trunk passes per
13669                    // burst on H100 q27, [spec-setup] trace). b stays in `out` (emitted) but
13670                    // OUT of `committed` (cache rows == committed); the consuming call
13671                    // prepends it once its verify commits the row. next_pred is unknowable
13672                    // without the commit pass — None; callers gate on pending_tok too.
13673                    debug_assert_eq!(out.last(), Some(&b), "pending must be the last emitted");
13674                    if let Some(slot) = sess_pending_slot.take() {
13675                        *slot = Some(b);
13676                    }
13677                    *next_pred_slot = None;
13678                    // fill_prev = hidden of the last COMMITTED row (b's predecessor) — the
13679                    // exact chain-seed/fill anchor the consuming burst (or a flush) needs.
13680                    *last_h = Some(e.clone_dtod(&fill_prev)?);
13681                    stashed_pending = true;
13682                } else {
13683                    // SAMPLED tail (unchanged): commit the bonus (one T=1 pass) + draft fill —
13684                    // the sampled round-0 accept needs this pass's logits (last_col_logits).
13685                    let pos_b = cache.pos;
13686                    scratch.set_len(e, pos_b)?;
13687                    let (lg_b, hb) = self.spec_target_step_h(e, b, &mut *cache)?;
13688                    // after a FULL-accept exit `last_pred` is STALE (it predicted the bonus
13689                    // itself — the prediction AFTER the bonus never materialized; it would have
13690                    // been the next round's verify col 0). The commit's logits ARE that
13691                    // prediction — so they are also the row the next burst's boundary token
13692                    // comes off, and (lane/sampled-spec-quality) it is DRAWN from them here.
13693                    *next_pred_slot = Some(if sample_boundary {
13694                        sample_boundary_token(
13695                            e,
13696                            &lg_b,
13697                            &sp,
13698                            &pen_hist,
13699                            &mut sctr,
13700                            "burst-tail-commit",
13701                        )?
13702                    } else {
13703                        argmax(&lg_b) as u32
13704                    });
13705                    self.mtp_kv_fill_all(e, &[b], &fill_prev, pos_b, &mut *scratch, embd_dev)?;
13706                    *last_h = Some(hb);
13707                }
13708            } else {
13709                // fill_prev tracks the hidden of the last COMMITTED row throughout the loop.
13710                *last_h = Some(e.clone_dtod(&fill_prev)?);
13711                if sample_boundary {
13712                    // No pending to commit, so the boundary row is the one `last_pred` was
13713                    // argmaxed from and the sampled path keeps it on device: the init feed's
13714                    // logits when the burst ran zero rounds, else the legacy-replay path's
13715                    // last verify column (both predict the token AFTER the last committed
13716                    // row). It is retained precisely because round 0's accept test needs it,
13717                    // so the draw costs no extra D2H of the [n_vocab] row.
13718                    match last_col_logits.as_ref() {
13719                        Some(lc) => {
13720                            *next_pred_slot = Some(sample_boundary_token_dev(
13721                                e,
13722                                lc,
13723                                n_vocab,
13724                                &sp,
13725                                &pen_hist,
13726                                &mut sctr,
13727                                "burst-tail-nopending",
13728                            )?);
13729                        }
13730                        // NAME THE FALLBACK (house standard): unreachable today — a sampled
13731                        // burst always feeds or replays, so the row exists — but if it ever
13732                        // is, the stream takes a greedy token and SAYS so rather than
13733                        // silently regressing to the pre-lane behaviour.
13734                        None => eprintln!(
13735                            "[spec-boundary] sampled tail kept the ARGMAX boundary token \
13736                             (reason: no retained boundary logits row)"
13737                        ),
13738                    }
13739                }
13740            }
13741            *sctr_slot = sctr;
13742            *uctr_slot = uctr;
13743            committed.extend_from_slice(prompt);
13744            if let Some(cb) = carried_pending {
13745                // the consumed carry's cache row landed in round 0's verify (every pending
13746                // round commits col 0) — it joins `committed` here, in sequence order.
13747                committed.push(cb);
13748            }
13749            if stashed_pending {
13750                // ZERO-EMIT BURST (2026-08-06 c=8 serve panic, pre-existing since b4aea184):
13751                // `out.len() - 1` underflowed on an EMPTY `out` — "range end index
13752                // 18446744073709551615 out of range for slice of length 0", killing the
13753                // memra-gpu-worker and failing 31 of 32 concurrent requests with "worker closed
13754                // stream". Reachable because `pending` starts as `carried_pending` (a bonus
13755                // stashed by the PREVIOUS burst) while `out` starts empty, and the carry is
13756                // deliberately NOT pushed to `out` (line ~3239: the burst that emitted it already
13757                // did). So a burst that stashes a pending without emitting anything of its own —
13758                // the round loop exits before a push, e.g. the ring drain's `out.len() < max_new`
13759                // guard skipping every token under a tight budget — arrives here with
13760                // out.len() == 0 and stashed_pending == true.
13761                //
13762                // The invariant is unchanged: `committed` gets every emitted token EXCEPT the
13763                // stashed bonus. With nothing emitted, that is nothing — and the carry pushed
13764                // just above is already accounted. Saturating, not a min/assert: an empty `out`
13765                // here is a legitimate burst shape, not a corrupt state.
13766                let emitted = out.len().saturating_sub(1);
13767                committed.extend_from_slice(&out[..emitted]);
13768            } else {
13769                committed.extend_from_slice(&out); // FULL out incl. overshoot — all committed
13770            }
13771            debug_assert_eq!(
13772                cache.pos,
13773                committed.len(),
13774                "session invariant: cache rows == committed tokens"
13775            );
13776            if setup_trace {
13777                e.stream().synchronize()?; // bound the async tail fill in the trace
13778                let t_tail = t_ent.elapsed();
13779                eprintln!(
13780                    "[spec-setup] init={:.2}ms cap={:.2}ms fill={:.2}ms rounds={:.2}ms tail={:.2}ms total={:.2}ms out={} cont={}",
13781                    t_init.as_secs_f64() * 1e3,
13782                    (t_cap - t_init).as_secs_f64() * 1e3,
13783                    (t_fill - t_cap).as_secs_f64() * 1e3,
13784                    (t_rounds - t_fill).as_secs_f64() * 1e3,
13785                    (t_tail - t_rounds).as_secs_f64() * 1e3,
13786                    t_tail.as_secs_f64() * 1e3,
13787                    out.len(),
13788                    continuation
13789                );
13790            }
13791            return Ok((out, total_drafted, total_accepted));
13792        }
13793        out.truncate(max_new);
13794        Ok((out, total_drafted, total_accepted))
13795    }
13796
13797    /// Anchor-bounded DSpark target extraction. The trunk sees the exact generated token tape;
13798    /// only requested hidden rows and target-logit rows cross PCIe. An anchor token at p pairs
13799    /// with the pre-output-norm h[p-1] carrier, exactly as the existing replay/NextN path does.
13800    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
13801    pub fn extract_dspark_anchors(
13802        &self,
13803        e: &Engine,
13804        tokens: &[u32],
13805        anchor_positions: &[usize],
13806        gamma: usize,
13807        top_k: usize,
13808        chunk: usize,
13809        temperature: f32,
13810    ) -> Result<Vec<DsparkAnchorRecord>, Box<dyn std::error::Error>> {
13811        if tokens.len() < gamma + 2 || gamma == 0 || chunk < 2 {
13812            return Err("DSpark extraction token tape/gamma/chunk is invalid".into());
13813        }
13814        if anchor_positions.windows(2).any(|pair| pair[0] >= pair[1]) {
13815            return Err("DSpark anchor positions must be sorted and unique".into());
13816        }
13817        for &position in anchor_positions {
13818            if position == 0 || position + gamma >= tokens.len() {
13819                return Err(format!(
13820                    "DSpark anchor {position} has no predecessor or cannot cover gamma={gamma} in {} tokens",
13821                    tokens.len()
13822                )
13823                .into());
13824            }
13825        }
13826
13827        let n_vocab = self.output.out_features();
13828        let n_embd = self.cfg.n_embd as usize;
13829        let mut cache =
13830            crate::pp::new_cache_planned(e, &self.cfg, &self.plan, tokens.len() + gamma + 8)?;
13831        let (embd_qt, embd_rb) = self.embd.qt_and_row_bytes(n_embd);
13832        let embd_gpu = if spec_host_embd() {
13833            None
13834        } else {
13835            Some(
13836                self.embd_gpu
13837                    .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload")),
13838            )
13839        };
13840        let embd_dev = embd_gpu.map(|gpu| (gpu, embd_qt, embd_rb));
13841
13842        struct PendingRecord {
13843            position: usize,
13844            hidden: Option<Vec<f32>>,
13845            tokens: Vec<u32>,
13846            target_top_ids: Vec<Option<Vec<u32>>>,
13847            target_top_logits: Vec<Option<Vec<f32>>>,
13848            target_top_probs: Vec<Option<Vec<f32>>>,
13849            target_tail_probs: Vec<Option<f32>>,
13850        }
13851
13852        let mut pending: Vec<PendingRecord> = anchor_positions
13853            .iter()
13854            .map(|&position| PendingRecord {
13855                position,
13856                hidden: None,
13857                tokens: tokens[position..=position + gamma].to_vec(),
13858                target_top_ids: vec![None; gamma],
13859                target_top_logits: vec![None; gamma],
13860                target_top_probs: vec![None; gamma],
13861                target_tail_probs: vec![None; gamma],
13862            })
13863            .collect();
13864
13865        let mut start = 0usize;
13866        while start < tokens.len() {
13867            let end = (start + chunk).min(tokens.len());
13868            let chunk_tokens = &tokens[start..end];
13869            let (target_logits, hidden_rows) =
13870                self.decode_step_t_core(e, chunk_tokens, start, &mut cache, embd_dev, None)?;
13871            for record in &mut pending {
13872                let hidden_position = record.position - 1;
13873                if hidden_position >= start && hidden_position < end {
13874                    let local = hidden_position - start;
13875                    record.hidden = Some(
13876                        e.dtoh_view(&hidden_rows.slice(local * n_embd..(local + 1) * n_embd))?,
13877                    );
13878                }
13879                for slot in 0..gamma {
13880                    let target_row = record.position + slot;
13881                    if target_row < start || target_row >= end {
13882                        continue;
13883                    }
13884                    let local = target_row - start;
13885                    let logits =
13886                        e.dtoh_view(&target_logits.slice(local * n_vocab..(local + 1) * n_vocab))?;
13887                    let (ids, top_logits, probs, tail) =
13888                        dspark_sparse_softmax_topk(&logits, top_k, temperature)?;
13889                    record.target_top_ids[slot] = Some(ids);
13890                    record.target_top_logits[slot] = Some(top_logits);
13891                    record.target_top_probs[slot] = Some(probs);
13892                    record.target_tail_probs[slot] = Some(tail);
13893                }
13894            }
13895            start = end;
13896        }
13897
13898        pending
13899            .into_iter()
13900            .map(|record| {
13901                let hidden = record
13902                    .hidden
13903                    .ok_or_else(|| format!("missing DSpark hidden at {}", record.position))?;
13904                let target_top_ids =
13905                    flatten_dspark_rows(record.target_top_ids, record.position, "target ids")?;
13906                let target_top_logits = flatten_dspark_rows(
13907                    record.target_top_logits,
13908                    record.position,
13909                    "target logits",
13910                )?;
13911                let target_top_probs =
13912                    flatten_dspark_rows(record.target_top_probs, record.position, "target probs")?;
13913                let target_tail_probs = record
13914                    .target_tail_probs
13915                    .into_iter()
13916                    .enumerate()
13917                    .map(|(slot, value)| {
13918                        value.ok_or_else(|| {
13919                            format!("missing DSpark tail at {} slot {slot}", record.position)
13920                        })
13921                    })
13922                    .collect::<Result<Vec<_>, _>>()?;
13923                Ok(DsparkAnchorRecord {
13924                    position: record.position,
13925                    hidden,
13926                    tokens: record.tokens,
13927                    target_top_ids,
13928                    target_top_logits,
13929                    target_top_probs,
13930                    target_tail_probs,
13931                })
13932            })
13933            .collect()
13934    }
13935
13936    /// TEACHER-FORCED REPLAY ACCEPTANCE (hqmtp MTP-heal protocol): walk a FIXED token
13937    /// sequence and, at sampled positions, compare the MTP head's K-token draft chain against
13938    /// the trunk's own teacher-forced greedy predictions. Nothing is generated — the context is
13939    /// the corpus text itself, so (a) degenerate self-generated loops cannot inflate acceptance
13940    /// and (b) two arms (bf16 ceiling vs NVFP4) score on IDENTICAL contexts, isolating the
13941    /// quant-induced head/hidden-state mismatch from text drift.
13942    ///
13943    /// Per eval position p (context = tokens[0..=p], predecessor pairing as in spec decode):
13944    ///   draft_j  = chain token j from (tokens[p], h_{p-1}), then its own drafts — the exact
13945    ///              eager spec-decode chain (same mtp_head_forward_dev, same rope positions).
13946    ///   target_j = teacher-forced greedy pick for position p+1+j (argmax of the trunk logits
13947    ///              at forced context tokens[0..p+j]). For j==0 this equals live spec
13948    ///              acceptance; for j>=1 live verify would condition on the drafts, here it
13949    ///              conditions on the corpus — deterministic and arm-comparable by design.
13950    ///
13951    /// Returns (rows, bg): one (p, drafts[k], targets[k]) row per eval position (ascending p),
13952    /// plus the full teacher-forced greedy track bg (bg[i] = greedy pick for position i, i>=1)
13953    /// so harnesses can cross-check runs (e.g. different chunk sizes must give identical bg).
13954    ///
13955    /// `hdump`: when Some, every position's pre-output_norm trunk hidden (the exact rows the
13956    /// draft-KV fill pairs from) streams to the file as little-endian f32 [t_total, n_embd] —
13957    /// the head-distillation extraction (hqmtp): the ENGINE is the source of truth for trunk
13958    /// hiddens (HF torch reproductions of the hybrid trunk measured only ~0.5 greedy
13959    /// agreement vs this path — not usable as a training-data source).
13960    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
13961    pub fn replay_acceptance(
13962        &self,
13963        e: &Engine,
13964        tokens: &[u32],
13965        k: usize,
13966        stride: usize,
13967        chunk: usize,
13968        mut hdump: Option<&mut std::fs::File>,
13969    ) -> Result<(Vec<(usize, Vec<u32>, Vec<u32>)>, Vec<u32>), Box<dyn std::error::Error>> {
13970        assert!(k >= 1 && stride >= 1 && chunk >= 2);
13971        let mtp = self
13972            .mtp
13973            .as_ref()
13974            .expect("replay_acceptance requires an MTP head");
13975        let n_vocab = self.output.out_features();
13976        let d_vocab = mtp
13977            .shared_head_head
13978            .as_ref()
13979            .unwrap_or(&self.output)
13980            .out_features();
13981        let n_embd = self.cfg.n_embd as usize;
13982        let t_total = tokens.len();
13983        assert!(t_total >= 8, "corpus too short ({t_total} tokens)");
13984        // STAGE-OWNED KV (lane/pp2-spec 2026-08-06) — see `new_session`. Door shut = `Cache::new`.
13985        let mut cache = crate::pp::new_cache_planned(e, &self.cfg, &self.plan, t_total + k + 8)?;
13986        let mut scratch = self.new_mtp_scratch(e, t_total + k + 8)?;
13987        let (embd_qt, embd_rb) = self.embd.qt_and_row_bytes(n_embd);
13988        let embd_gpu = if spec_host_embd() {
13989            None
13990        } else {
13991            Some(
13992                self.embd_gpu
13993                    .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload")),
13994            )
13995        };
13996        let embd_dev = embd_gpu.map(|g| (g, embd_qt, embd_rb));
13997
13998        // bg[i] = the trunk's greedy pick for position i under the forced context (i >= 1).
13999        let mut bg: Vec<u32> = vec![0; t_total + 1];
14000        let mut rows: Vec<(usize, Vec<u32>, Vec<u32>)> = Vec::new();
14001        let mut prev_last_h = e.zeros(n_embd)?; // predecessor hidden entering the chunk
14002        let mut seed_buf = e.zeros(n_embd)?;
14003        let mut preds_d = e.alloc_u32_zeroed(chunk)?;
14004        let nll_on = std::env::var("MEMRA_REPLAY_NLL").as_deref() == Ok("1");
14005        let (mut nll_sum, mut nll_cnt) = (0f64, 0u64);
14006        let mut s = 0usize;
14007        while s < t_total {
14008            let cend = (s + chunk).min(t_total);
14009            let tc = cend - s;
14010            let ch = &tokens[s..cend];
14011            // 1. forced trunk pass — verify path (decode-exact contract): all-column logits +
14012            //    the chunk's true hiddens.
14013            let (tl_d, vx) = self.decode_step_t_core(e, ch, s, &mut cache, embd_dev, None)?;
14014            for j in 0..tc {
14015                e.argmax_token_device_col(&tl_d, j, n_vocab, &mut preds_d, j)?;
14016            }
14017            let preds = e.dtoh_u32(&preds_d)?;
14018            for j in 0..tc {
14019                bg[s + j + 1] = preds[j];
14020            }
14021            // MEMRA_REPLAY_NLL=1: teacher-forced NLL/perplexity over the same forced pass — the
14022            // checkpoint-quality metric (position j's logits score the GOLD next token).
14023            if nll_on {
14024                let jmax = if cend < t_total { tc } else { tc - 1 }; // last pos has no gold next
14025                if jmax > 0 {
14026                    let ids: Vec<u32> = (0..jmax).map(|j| tokens[s + j + 1]).collect();
14027                    let rows: Vec<i32> = (0..jmax as i32).collect();
14028                    let idsd = e.htod_u32_v(&ids)?;
14029                    let rowsd = e.htod_i32(&rows)?;
14030                    let mut outd = e.zeros(jmax)?;
14031                    e.softmax_gather(&tl_d, n_vocab, &idsd, &rowsd, &mut outd, n_vocab, jmax, 1.0)?;
14032                    for pr in e.dtoh(&outd)? {
14033                        nll_sum += -((pr.max(1e-30)) as f64).ln();
14034                        nll_cnt += 1;
14035                    }
14036                }
14037            }
14038            if let Some(f) = hdump.as_deref_mut() {
14039                use std::io::Write;
14040                let host: Vec<f32> = e.dtoh(&vx)?;
14041                // bf16 round-to-nearest-even — f32 doubled the disk bill at bulk
14042                // extraction scale (20M tokens x 4096 = 320GB f32 vs 160GB bf16).
14043                let mut bytes = Vec::with_capacity(tc * n_embd * 2);
14044                for v in &host[..tc * n_embd] {
14045                    let b = v.to_bits();
14046                    let r = b.wrapping_add(0x7FFF + ((b >> 16) & 1));
14047                    bytes.extend_from_slice(&((r >> 16) as u16).to_le_bytes());
14048                }
14049                f.write_all(&bytes)?;
14050            }
14051            // CHAINLESS extraction (stride > corpus, the bulk-hdump mode): no chunk ever
14052            // drafts, so the draft-KV fills are pure waste — skip them (2 MTP-block passes
14053            // per token saved; the forced trunk pass + hdump is all the mode needs).
14054            let chainless = stride > t_total;
14055            if chainless {
14056                e.copy_view_into(
14057                    &mut prev_last_h,
14058                    0,
14059                    &vx.slice((tc - 1) * n_embd..tc * n_embd),
14060                    n_embd,
14061                )?;
14062                s = cend;
14063                continue;
14064            }
14065            // 2. TRUE predecessor-paired draft-KV fill for the chunk (row i carries h_{i-1};
14066            //    row s reads the previous chunk's last true hidden, zeros at corpus start).
14067            let mut vxs = e.zeros(tc * n_embd)?;
14068            e.copy_into(&mut vxs, 0, &prev_last_h, n_embd)?;
14069            if tc > 1 {
14070                e.copy_view_into(
14071                    &mut vxs,
14072                    n_embd,
14073                    &vx.slice(0..(tc - 1) * n_embd),
14074                    (tc - 1) * n_embd,
14075                )?;
14076            }
14077            scratch.set_len(e, s)?;
14078            self.mtp_kv_fill_all(e, ch, &vxs, s, &mut scratch, embd_dev)?;
14079            // 3. draft chains at sampled positions, DESCENDING: a chain reads only slots
14080            //    [0..p) (true fills) and appends at >= p; the next (smaller-p) chain's set_len
14081            //    truncates those approximate appends before they can ever be read.
14082            let ps: Vec<usize> = (s..cend)
14083                .filter(|p| *p >= 1 && *p % stride == 0 && *p + k <= t_total)
14084                .collect();
14085            for &p in ps.iter().rev() {
14086                scratch.set_len(e, p)?;
14087                if p == s {
14088                    e.copy_into(&mut seed_buf, 0, &prev_last_h, n_embd)?;
14089                } else {
14090                    e.copy_view_into(
14091                        &mut seed_buf,
14092                        0,
14093                        &vx.slice((p - 1 - s) * n_embd..(p - s) * n_embd),
14094                        n_embd,
14095                    )?;
14096                }
14097                let mut e_tok = tokens[p];
14098                let mut d_seed = e.clone_dtod(&seed_buf)?;
14099                let chain_heads = !self.mtp_extra.is_empty();
14100                let mut chain_tokens = if chain_heads {
14101                    vec![tokens[p]]
14102                } else {
14103                    Vec::new()
14104                };
14105                let mut chain_seeds = if chain_heads {
14106                    vec![e.clone_dtod(&seed_buf)?]
14107                } else {
14108                    Vec::new()
14109                };
14110                let mut drafts: Vec<u32> = Vec::with_capacity(k);
14111                for j in 0..k {
14112                    let (dl_d, h_nextn) = if chain_heads {
14113                        self.mtp_chain_forward_dev(
14114                            e,
14115                            &chain_tokens,
14116                            &chain_seeds,
14117                            &mut scratch,
14118                            p,
14119                            embd_dev,
14120                            None,
14121                        )?
14122                    } else {
14123                        self.mtp_head_forward_dev(
14124                            e,
14125                            mtp,
14126                            e_tok,
14127                            &d_seed,
14128                            &mut scratch,
14129                            p + 1 + j,
14130                            embd_dev,
14131                            None,
14132                        )?
14133                    };
14134                    let tok_d = e.argmax_token_device(&dl_d, d_vocab)?;
14135                    let idx = e.dtoh_u32_one(&tok_d)?;
14136                    let d = match &mtp.d2t {
14137                        Some(map) => map[idx as usize],
14138                        None => idx,
14139                    };
14140                    drafts.push(d);
14141                    if chain_heads {
14142                        chain_tokens.push(d);
14143                        chain_seeds.push(h_nextn);
14144                    } else {
14145                        e_tok = d;
14146                        d_seed = h_nextn;
14147                    }
14148                }
14149                // targets may live in a LATER chunk's bg — resolved after the walk.
14150                rows.push((p, drafts, Vec::new()));
14151            }
14152            // 4. restore TRUE entries for the whole chunk (the next chunk's chains and fills
14153            //    expect scratch.len == cend with exact rows).
14154            scratch.set_len(e, s)?;
14155            self.mtp_kv_fill_all(e, ch, &vxs, s, &mut scratch, embd_dev)?;
14156            e.copy_view_into(
14157                &mut prev_last_h,
14158                0,
14159                &vx.slice((tc - 1) * n_embd..tc * n_embd),
14160                n_embd,
14161            )?;
14162            s = cend;
14163        }
14164        for (p, drafts, targets) in rows.iter_mut() {
14165            for j in 0..drafts.len() {
14166                targets.push(bg[*p + 1 + j]);
14167            }
14168        }
14169        rows.sort_by_key(|r| r.0);
14170        if nll_cnt > 0 {
14171            let mean = nll_sum / nll_cnt as f64;
14172            println!(
14173                "[replay-nll] tokens={nll_cnt} nll/token={mean:.5} ppl={:.4}",
14174                mean.exp()
14175            );
14176        }
14177        Ok((rows, bg))
14178    }
14179}
14180
14181#[cfg(test)]
14182mod vg_debt_tests {
14183    use super::dspark_vg_debt_projection;
14184
14185    /// TOOTH for the verify-graph admission accounting: the pool's projected remaining
14186    /// growth must be charged (pre-fix, admission charged 0 for a pool measured at
14187    /// 8,852 MiB), the projection must price the MARGINAL cost of one more key rather than
14188    /// extrapolating the pool's one-time shared allocation, and the doors that make growth
14189    /// impossible must zero the debt.
14190    #[test]
14191    fn vg_debt_projects_remaining_growth_and_respects_the_freeze_valves() {
14192        const MIB: usize = 1 << 20;
14193        let d = dspark_vg_debt_projection;
14194        // cold pool: nothing observed, one capture fits inside SPEC_SHRINK_RESERVE.
14195        assert_eq!(d(0, 256, 0, None), 0);
14196        // freeze valve MEMRA_DSPARK_VG_MAX=0: the pool cannot grow.
14197        assert_eq!(d(10, 0, 500 * MIB, None), 0);
14198        // saturated pool: at/past the cap the pool FREEZES, nothing left to reserve.
14199        assert_eq!(d(256, 256, 8852 * MIB, None), 0);
14200        assert_eq!(d(300, 256, 8852 * MIB, None), 0);
14201
14202        // BOOTSTRAP (one observation, growth unmeasurable): at most one more pool's worth.
14203        // The pre-fix mean rule extrapolated 255x here — the measured 8.5 GB phantom.
14204        assert_eq!(d(1, 256, 33 * MIB, None), 33 * MIB);
14205
14206        // MARGINAL, flat pool (the box9 receipt: reserved stayed ~33.6 MiB across captures
14207        // 1..3, so an additional key costs ~nothing and the debt must collapse to ~0 —
14208        // NOT the 8,556/4,261/2,830 MB the mean rule printed).
14209        assert_eq!(d(3, 256, 33 * MIB, Some((1, 33 * MIB))), 0);
14210
14211        // MARGINAL, genuinely growing pool: 40 MiB per new key over 2 keys, 250 slots left.
14212        let debt = d(6, 256, 273 * MIB, Some((4, 193 * MIB)));
14213        assert_eq!(debt, 250 * (40 * MIB));
14214        assert!(
14215            debt > 3 * (1536 * MIB),
14216            "real growth must dwarf SPEC_SHRINK_RESERVE"
14217        );
14218
14219        // a shrinking/recycled reading never becomes a negative charge.
14220        assert_eq!(d(6, 256, 10 * MIB, Some((4, 99 * MIB))), 0);
14221        // a stale observation at the same capture count falls back to bootstrap.
14222        assert_eq!(d(4, 256, 80 * MIB, Some((4, 80 * MIB))), 80 * MIB);
14223    }
14224}
14225
14226#[cfg(test)]
14227mod capture_headroom_tests {
14228    use super::{
14229        CAPTURE_HEADROOM_FLOOR, capture_err_is_oom, capture_headroom_verdict,
14230        draft_capture_bootstrap_estimate,
14231    };
14232
14233    /// TOOTH for the pre-capture reserve check (lane/step37-vram-admission-20260830): a
14234    /// capture attempt must be refused BEFORE it allocates when the device cannot cover its
14235    /// appetite plus the post-capture floor — and pool-cached bytes count as headroom
14236    /// (driver `free` alone under-counts, the wrong direction for a gate that drops
14237    /// coverage).
14238    #[test]
14239    fn capture_reserve_check_refuses_short_devices_and_counts_pool_cache() {
14240        const MIB: usize = 1 << 20;
14241        let need = 900 * MIB;
14242        // Plenty of room: no refusal.
14243        assert_eq!(
14244            capture_headroom_verdict(8_000 * MIB, 0, need, CAPTURE_HEADROOM_FLOOR),
14245            None
14246        );
14247        // The owner's shape: capture appetite would walk the card to the edge — refused,
14248        // with the arithmetic surfaced for the WARN line.
14249        let (required, effective) =
14250            capture_headroom_verdict(1_200 * MIB, 0, need, CAPTURE_HEADROOM_FLOOR)
14251                .expect("short device must refuse");
14252        assert_eq!(required, need + CAPTURE_HEADROOM_FLOOR);
14253        assert_eq!(effective, 1_200 * MIB);
14254        // Pool-cached bytes are real headroom (the trim path makes them driver-visible).
14255        assert_eq!(
14256            capture_headroom_verdict(1_200 * MIB, 7_000 * MIB, need, CAPTURE_HEADROOM_FLOOR),
14257            None
14258        );
14259        // Boundary: exactly enough is enough (>=, never a fencepost refusal).
14260        assert_eq!(
14261            capture_headroom_verdict(
14262                need + CAPTURE_HEADROOM_FLOOR,
14263                0,
14264                need,
14265                CAPTURE_HEADROOM_FLOOR
14266            ),
14267            None
14268        );
14269        // POLICY at the call site (owner-shape receipts, escalated twice on-box): the
14270        // refusal fn is handed 2x the appetite plus TWO floors — a capture may take at
14271        // most half the discretionary headroom, so the card retains a whole capture's
14272        // worth of room after it lands. One floor of slack above one appetite (the shape
14273        // that step-OOM'd on the owner cell) must therefore REFUSE under the call-site
14274        // requirement.
14275        assert!(
14276            capture_headroom_verdict(
14277                need + CAPTURE_HEADROOM_FLOOR + (100 << 20),
14278                0,
14279                2 * need,
14280                CAPTURE_HEADROOM_FLOOR * 2
14281            )
14282            .is_some()
14283        );
14284    }
14285
14286    #[test]
14287    fn bootstrap_estimate_scales_with_heads_and_never_underflows() {
14288        // 3-head chain on a step37-shaped vocab must expect strictly more than one head.
14289        let one = draft_capture_bootstrap_estimate(1, 3, 128_896, 4_096);
14290        let three = draft_capture_bootstrap_estimate(3, 3, 128_896, 4_096);
14291        assert!(three > one);
14292        // Degenerate shapes keep a sane minimum (the estimate feeds a refusal gate; a
14293        // zero-need gate refuses nothing).
14294        assert!(draft_capture_bootstrap_estimate(0, 0, 0, 0) >= 64 << 20);
14295    }
14296
14297    #[test]
14298    fn capture_oom_predicate_matches_the_quoted_driver_text() {
14299        assert!(capture_err_is_oom(
14300            "DriverError(CUDA_ERROR_OUT_OF_MEMORY, \"out of memory\")"
14301        ));
14302        assert!(capture_err_is_oom("allocation failed: out of memory"));
14303        assert!(!capture_err_is_oom("capture produced no graph"));
14304    }
14305}
14306
14307#[cfg(test)]
14308mod mtp_chain_tests {
14309    use super::mtp_chain_head_index;
14310
14311    #[test]
14312    fn embedded_step_heads_cycle_in_declared_order() {
14313        let actual: Vec<usize> = (0..8).map(|step| mtp_chain_head_index(step, 3)).collect();
14314        assert_eq!(actual, [0, 1, 2, 0, 1, 2, 0, 1]);
14315    }
14316
14317    #[test]
14318    fn standalone_draft_remains_single_head() {
14319        assert!((0..8).all(|step| mtp_chain_head_index(step, 1) == 0));
14320    }
14321}
14322
14323#[cfg(test)]
14324mod tp_verified_prefix_tests {
14325    use super::validate_tp_kv_snapshot_shape;
14326    use crate::tp::ResidentTpKvCache;
14327
14328    #[test]
14329    fn snapshot_shape_accepts_matching_tp_presence() {
14330        let layers = vec![
14331            Some(ResidentTpKvCache::new(Vec::new(), 1, 1, 1, 1, 8)),
14332            None,
14333        ];
14334        validate_tp_kv_snapshot_shape(&layers, &[Some(2), None]).unwrap();
14335    }
14336
14337    #[test]
14338    fn snapshot_shape_rejects_changed_tp_presence() {
14339        let layers = vec![Some(ResidentTpKvCache::new(Vec::new(), 1, 1, 1, 1, 8))];
14340        let error = validate_tp_kv_snapshot_shape(&layers, &[None])
14341            .unwrap_err()
14342            .to_string();
14343        assert!(error.contains("changed shape"), "unexpected error: {error}");
14344    }
14345
14346    #[test]
14347    fn step37_dcw_rebase_at_5151_boundary() {
14348        use crate::cache::KvRingAppend;
14349        let mut cache = ResidentTpKvCache::new_swa(Vec::new(), 128, 128, 136, 96, 262_144, 512);
14350        assert_eq!(cache.physical_capacity(), 5151);
14351        cache.publish_hydration(5150, 0).unwrap();
14352        assert_eq!(cache.ring_base(), Some(0));
14353
14354        // Before rebase, attempting to view rows past capacity fails with the exact bug error
14355        let err = cache.physical_range(5150, 5153).unwrap_err();
14356        assert_eq!(
14357            err,
14358            "SWA ring view [5150,5153) is outside resident [0,5151)"
14359        );
14360
14361        let (write_row, would_rebase) = cache.peek_append_ring(3).unwrap();
14362        assert!(would_rebase);
14363        assert_eq!(write_row, 542);
14364
14365        let tx = cache.begin_transaction().unwrap();
14366        let plan = cache.prepare_append(tx, 3).unwrap();
14367        assert_eq!(plan.target(), 5153);
14368        assert_eq!(plan.write_row(), 542);
14369        assert_eq!(
14370            plan.ring_append(),
14371            Some(KvRingAppend::Rebase {
14372                src_row: 4608,
14373                keep_rows: 542,
14374                new_base: 4608,
14375                write_row: 542,
14376            })
14377        );
14378        cache.publish_append_rebase(plan).unwrap();
14379        cache.publish_append_plan(plan).unwrap();
14380        assert_eq!(cache.ring_base(), Some(4608));
14381
14382        // After rebase, physical range is within bounds
14383        let range = cache.physical_range(5150, 5153).unwrap();
14384        assert_eq!(range, 542..545);
14385
14386        let target = cache.commit_target(tx, 3).unwrap();
14387        cache.publish_finalize(tx, target).unwrap();
14388        assert_eq!((cache.committed_len(), cache.staged_len()), (5153, 5153));
14389    }
14390
14391    #[test]
14392    fn step37_dcw_rebase_rollback_preserves_view_and_base() {
14393        let mut cache = ResidentTpKvCache::new_swa(Vec::new(), 128, 128, 136, 96, 262_144, 512);
14394        cache.publish_hydration(5150, 0).unwrap();
14395
14396        let tx = cache.begin_transaction().unwrap();
14397        let plan = cache.prepare_append(tx, 3).unwrap();
14398        cache.publish_append_rebase(plan).unwrap();
14399        cache.publish_append_plan(plan).unwrap();
14400        assert_eq!(cache.ring_base(), Some(4608));
14401
14402        // Rollback 0 rows accepted (pass declined)
14403        let rollback = cache.commit_target(tx, 0).unwrap();
14404        cache.publish_finalize(tx, rollback).unwrap();
14405        assert_eq!((cache.committed_len(), cache.staged_len()), (5150, 5150));
14406        assert_eq!(cache.ring_base(), Some(4608));
14407
14408        // View for base_len (5150) is still valid in resident [4608, 4608 + 5151)
14409        let range = cache.physical_range(4608, 5150).unwrap();
14410        assert_eq!(range, 0..542);
14411    }
14412}
14413
14414#[cfg(test)]
14415mod dspark_sparse_tests {
14416    use super::dspark_sparse_softmax_topk;
14417
14418    #[test]
14419    fn topk_keeps_full_softmax_mass_and_stable_ties() {
14420        let logits = [1.0f32, 3.0, 3.0, -2.0];
14421        let (ids, top_logits, probs, tail) = dspark_sparse_softmax_topk(&logits, 2, 1.0).unwrap();
14422        assert_eq!(ids, vec![1, 2]);
14423        assert_eq!(top_logits, vec![3.0, 3.0]);
14424        let denominator = logits.iter().map(|value| (value - 3.0).exp()).sum::<f32>();
14425        let expected = 1.0 / denominator;
14426        assert!((probs[0] - expected).abs() < 1.0e-6);
14427        assert!((probs[1] - expected).abs() < 1.0e-6);
14428        assert!((tail - (1.0 - 2.0 * expected)).abs() < 1.0e-6);
14429        assert!((probs.iter().sum::<f32>() + tail - 1.0).abs() < 1.0e-6);
14430    }
14431}
14432
14433#[cfg(test)]
14434mod spec_replay_env_tests {
14435    use super::spec_replay_env_on;
14436
14437    #[test]
14438    fn replay_requires_literal_one() {
14439        assert!(!spec_replay_env_on(None));
14440        assert!(!spec_replay_env_on(Some("")));
14441        assert!(!spec_replay_env_on(Some("0")));
14442        assert!(!spec_replay_env_on(Some("true")));
14443        assert!(!spec_replay_env_on(Some("2")));
14444        assert!(spec_replay_env_on(Some("1")));
14445    }
14446}
14447
14448#[cfg(test)]
14449mod telem_tests {
14450    use super::{SPEC_TELEM_POS, SpecTelemetry, SpecTelemetryCounters};
14451
14452    #[test]
14453    fn synthetic_accept_masks_produce_tau_and_position_histogram() {
14454        let counters = SpecTelemetryCounters::default();
14455        for mask in [
14456            [true, true, true],
14457            [true, true, false],
14458            [true, false, false],
14459            [false, false, false],
14460        ] {
14461            let accepted = mask.iter().take_while(|&&value| value).count();
14462            counters.record_round(mask.len(), accepted);
14463        }
14464
14465        let snapshot = counters.snapshot();
14466        assert_eq!(
14467            (snapshot.rounds, snapshot.drafted, snapshot.accepted),
14468            (4, 12, 6)
14469        );
14470        assert_eq!(&snapshot.pos_drafted[..3], &[4, 4, 4]);
14471        assert_eq!(&snapshot.pos_accepted[..3], &[3, 2, 1]);
14472        assert_eq!(snapshot.tau(), 1.5);
14473        assert_eq!(snapshot.pos_drafted[3..], [0; SPEC_TELEM_POS - 3]);
14474        assert_eq!(snapshot.pos_accepted[3..], [0; SPEC_TELEM_POS - 3]);
14475    }
14476
14477    /// The worker's per-burst pattern: stash, accumulate, diff — the delta must isolate
14478    /// exactly the burst's contribution (pool-resumed sessions carry prior requests' counts).
14479    #[test]
14480    fn delta_isolates_burst_contribution() {
14481        let mut t = SpecTelemetry::default();
14482        // "previous request": 2 rounds of k=3, accepts 3 then 1.
14483        for (kr, na) in [(3usize, 3usize), (3, 1)] {
14484            t.rounds += 1;
14485            t.drafted += kr as u64;
14486            t.accepted += na as u64;
14487            for j in 0..kr {
14488                t.pos_drafted[j] += 1;
14489            }
14490            for j in 0..na {
14491                t.pos_accepted[j] += 1;
14492            }
14493        }
14494        let before = t;
14495        // "this burst": 1 round k=3, accepts 2.
14496        t.rounds += 1;
14497        t.drafted += 3;
14498        t.accepted += 2;
14499        for j in 0..3 {
14500            t.pos_drafted[j] += 1;
14501        }
14502        for j in 0..2 {
14503            t.pos_accepted[j] += 1;
14504        }
14505        let d = t.delta_since(&before);
14506        assert_eq!((d.rounds, d.drafted, d.accepted), (1, 3, 2));
14507        assert_eq!(&d.pos_drafted[..3], &[1, 1, 1]);
14508        assert_eq!(&d.pos_accepted[..3], &[1, 1, 0]);
14509        assert_eq!(d.pos_drafted[3..], [0; SPEC_TELEM_POS - 3]);
14510    }
14511
14512    /// merge(delta) then merge(delta2) equals accumulating both — the per-model /metrics
14513    /// aggregation invariant.
14514    #[test]
14515    fn merge_accumulates_fieldwise() {
14516        let mut agg = SpecTelemetry::default();
14517        let mut d1 = SpecTelemetry {
14518            rounds: 2,
14519            drafted: 6,
14520            accepted: 4,
14521            ..Default::default()
14522        };
14523        d1.pos_drafted[0] = 2;
14524        d1.pos_accepted[0] = 2;
14525        let mut d2 = SpecTelemetry {
14526            rounds: 1,
14527            drafted: 3,
14528            accepted: 1,
14529            ..Default::default()
14530        };
14531        d2.pos_drafted[0] = 1;
14532        d2.pos_accepted[0] = 1;
14533        d2.pos_drafted[1] = 1;
14534        agg.merge(&d1);
14535        agg.merge(&d2);
14536        assert_eq!((agg.rounds, agg.drafted, agg.accepted), (3, 9, 5));
14537        assert_eq!(agg.pos_drafted[0], 3);
14538        assert_eq!(agg.pos_accepted[0], 3);
14539        assert_eq!(agg.pos_drafted[1], 1);
14540        assert_eq!(agg.pos_accepted[1], 0);
14541    }
14542
14543    /// Wrong-snapshot diff saturates to zero instead of wrapping — the counters feed a
14544    /// public metrics surface and must never publish a u64-wrapped garbage value.
14545    #[test]
14546    fn delta_saturates_never_wraps() {
14547        let small = SpecTelemetry {
14548            rounds: 1,
14549            drafted: 2,
14550            accepted: 1,
14551            ..Default::default()
14552        };
14553        let big = SpecTelemetry {
14554            rounds: 5,
14555            drafted: 15,
14556            accepted: 9,
14557            ..Default::default()
14558        };
14559        let d = small.delta_since(&big);
14560        assert_eq!((d.rounds, d.drafted, d.accepted), (0, 0, 0));
14561    }
14562}
14563
14564#[cfg(test)]
14565mod draft_graph_fallback_tests {
14566    use super::DraftGraphFallback;
14567
14568    /// The Q2 contract, part (a): a fallback flip is LOUD — exactly once per flip.
14569    #[test]
14570    fn flip_is_loud_once_and_memoized_after() {
14571        let mut f = DraftGraphFallback::default();
14572        let line = f
14573            .mark_greedy("out of memory")
14574            .expect("first flip must return the warn line");
14575        assert!(
14576            line.contains("WARN"),
14577            "flip line must be warn-level: {line}"
14578        );
14579        assert!(
14580            line.contains("out of memory"),
14581            "flip line must carry the reason: {line}"
14582        );
14583        assert!(f.greedy_failed());
14584        // re-marking an already-failed graph is the memoization: quiet, still failed.
14585        assert!(f.mark_greedy("out of memory").is_none());
14586        assert!(f.greedy_failed());
14587        // the two graphs' flags are independent (greedy flip leaves sampled capturable).
14588        assert!(!f.sampled_failed());
14589        let line_s = f
14590            .mark_sampled("capture unsupported")
14591            .expect("sampled flip is its own flip");
14592        assert!(
14593            line_s.contains("sampled"),
14594            "sampled flip names itself: {line_s}"
14595        );
14596        assert!(f.mark_sampled("capture unsupported").is_none());
14597    }
14598
14599    /// The Q2 contract, part (b): resume-from-pool RESETS both flags (fresh capture chance),
14600    /// and says so exactly when there was something to reset.
14601    #[test]
14602    fn reset_on_resume_clears_flags_and_logs_once() {
14603        let mut f = DraftGraphFallback::default();
14604        // clean session: resume is silent, nothing to reset.
14605        assert!(f.reset_on_resume().is_none());
14606        f.mark_greedy("oom").unwrap();
14607        f.mark_sampled("oom").unwrap();
14608        let note = f
14609            .reset_on_resume()
14610            .expect("a set flag must produce the reset note");
14611        assert!(
14612            note.contains("greedy+sampled"),
14613            "note names what was reset: {note}"
14614        );
14615        assert!(
14616            !f.greedy_failed() && !f.sampled_failed(),
14617            "both flags cleared"
14618        );
14619        // and the NEXT failure after a reset is a fresh flip — loud again.
14620        assert!(f.mark_greedy("oom again").is_some());
14621        let note2 = f.reset_on_resume().expect("greedy-only reset");
14622        assert!(note2.contains("(greedy)"), "single-flag note: {note2}");
14623    }
14624
14625    /// Shape-change clears (dmask realloc / mask-shape mismatch / s_key change) stay silent —
14626    /// they precede a fresh capture attempt whose own failure re-flips loudly.
14627    #[test]
14628    fn shape_change_clears_are_silent() {
14629        let mut f = DraftGraphFallback::default();
14630        f.mark_greedy("oom").unwrap();
14631        f.clear_greedy();
14632        assert!(!f.greedy_failed());
14633        f.mark_sampled("oom").unwrap();
14634        f.clear_sampled();
14635        assert!(!f.sampled_failed());
14636        // after a silent clear there is nothing left for resume to report.
14637        assert!(f.reset_on_resume().is_none());
14638    }
14639}
14640
14641/// SAMPLED DRAFT-GRAPH KEY (lane/graph-s-key-exactness-20260819).
14642///
14643/// These are the CPU teeth for an exactness bug whose live reproduction needs a GPU, a trunk, a
14644/// drafter and a two-turn session: the key itself. Every test below fails against the pre-fix key
14645/// `(seed, temp.to_bits(), k)` — `legacy_key` restates it so the collision is explicit rather
14646/// than remembered.
14647#[cfg(test)]
14648mod sampled_graph_key_tests {
14649    use super::{SampledGraphKey, debug_t_pred0};
14650
14651    /// The pre-fix key, verbatim: `let s_key = (sp_seed, sp_temp.to_bits(), k);`
14652    fn legacy_key(k: &SampledGraphKey) -> (u64, u32, usize) {
14653        (k.seed, k.temp_bits, k.k)
14654    }
14655
14656    fn pure_temp_key() -> SampledGraphKey {
14657        // temperature 1.0, filters off — today's serve default, the shape that parks a graph.
14658        SampledGraphKey::new(12345, 1.0, 3, 0, 1.0, 0.0, false)
14659    }
14660
14661    /// THE COLLISION. Two requests that differ ONLY in the truncation filters shared one key, so
14662    /// a parked pure-temp graph survived into a filtered request and the launch site launched it.
14663    #[test]
14664    fn vendor_filters_change_the_key() {
14665        let parked = pure_temp_key();
14666        // qwen3.8 generation_config.json — what the vendor-default flip makes the default shape.
14667        let vendor = SampledGraphKey::new(12345, 1.0, 3, 20, 0.95, 0.0, false);
14668        assert_eq!(
14669            legacy_key(&parked),
14670            legacy_key(&vendor),
14671            "pre-fix key collided: this is the bug, and the reason a test asserts on it",
14672        );
14673        assert_ne!(parked, vendor, "post-fix key must separate the two regimes");
14674        assert!(parked.pure_temp());
14675        assert!(!vendor.pure_temp());
14676    }
14677
14678    /// Each distribution-shaping field alone is enough to drop the parked graph.
14679    #[test]
14680    fn every_filter_field_is_keyed() {
14681        let base = pure_temp_key();
14682        for (what, other) in [
14683            (
14684                "top_k",
14685                SampledGraphKey::new(12345, 1.0, 3, 20, 1.0, 0.0, false),
14686            ),
14687            (
14688                "top_p",
14689                SampledGraphKey::new(12345, 1.0, 3, 0, 0.95, 0.0, false),
14690            ),
14691            (
14692                "min_p",
14693                SampledGraphKey::new(12345, 1.0, 3, 0, 1.0, 0.05, false),
14694            ),
14695            (
14696                "penalties",
14697                SampledGraphKey::new(12345, 1.0, 3, 0, 1.0, 0.0, true),
14698            ),
14699        ] {
14700            assert_ne!(base, other, "{what} must be part of the key");
14701            assert!(!other.pure_temp(), "{what} leaves the pure-temp regime");
14702            assert_eq!(
14703                legacy_key(&base),
14704                legacy_key(&other),
14705                "{what} was invisible to the pre-fix key",
14706            );
14707        }
14708    }
14709
14710    /// The baked constants stay keyed (this half was always right — regression cover for it).
14711    #[test]
14712    fn baked_constants_stay_keyed() {
14713        let base = pure_temp_key();
14714        assert_ne!(
14715            base,
14716            SampledGraphKey::new(999, 1.0, 3, 0, 1.0, 0.0, false),
14717            "seed"
14718        );
14719        assert_ne!(
14720            base,
14721            SampledGraphKey::new(12345, 0.7, 3, 0, 1.0, 0.0, false),
14722            "temp"
14723        );
14724        assert_ne!(
14725            base,
14726            SampledGraphKey::new(12345, 1.0, 4, 0, 1.0, 0.0, false),
14727            "k"
14728        );
14729        // bitwise on temperature: 0.7f32 vs the same value re-derived must NOT differ.
14730        assert_eq!(
14731            SampledGraphKey::new(1, 0.7, 3, 0, 1.0, 0.0, false),
14732            SampledGraphKey::new(1, 7.0 / 10.0, 3, 0, 1.0, 0.0, false),
14733        );
14734    }
14735
14736    /// THE LOAD-BEARING HALF OF THE SEED DECISION (lane/session-resume-sampler-predicate-
14737    /// 20260820). The whole-session resume predicate deliberately does NOT compare `seed`: an
14738    /// omitted serve `seed` draws fresh per-request entropy, so comparing it would refuse every
14739    /// seed-omitting sampled conversation. That is only sound because the one piece of parked state
14740    /// that BAKES the seed — this graph — is re-keyed on it, so a seed change drops and recaptures.
14741    ///
14742    /// This test is the other end of that argument, asserted here rather than remembered in a
14743    /// comment: if a future change dropped `seed` from the key, the resume predicate's exclusion
14744    /// would silently become the unsound thing it is documented not to be.
14745    /// (Paired with `seed_alone_does_not_refuse` in `memra-sampling`.)
14746    #[test]
14747    fn seed_alone_still_rekeys_the_draft_graph() {
14748        let parked = pure_temp_key();
14749        let reseeded = SampledGraphKey::new(999, 1.0, 3, 0, 1.0, 0.0, false);
14750        assert_ne!(
14751            parked, reseeded,
14752            "a seed-only change MUST drop the parked sampled graph — the resume predicate's \
14753             decision not to compare seed rests on exactly this",
14754        );
14755        // Same regime on both sides: the drop is a recapture, not a fall to the eager chain
14756        // because of a filter difference.
14757        assert!(parked.pure_temp() && reseeded.pure_temp());
14758    }
14759
14760    /// `pure_temp()` is the capture guard's predicate, computed from the key so the two cannot
14761    /// drift. The equality below is the invariant the launch-site guard asserts: identical keys
14762    /// agree on the regime, so a graph that survives the drop is legal to launch.
14763    #[test]
14764    fn equal_keys_agree_on_the_regime() {
14765        let a = SampledGraphKey::new(7, 0.8, 3, 20, 0.95, 0.0, false);
14766        let b = SampledGraphKey::new(7, 0.8, 3, 20, 0.95, 0.0, false);
14767        assert_eq!(a, b);
14768        assert_eq!(a.pure_temp(), b.pure_temp());
14769        // top_p slightly above 1.0 (a client sending 1.0 exactly, or an operator default) is
14770        // still the unfiltered regime, matching the original `sp.top_p >= 1.0` test.
14771        assert!(SampledGraphKey::new(7, 0.8, 3, 0, 1.0, 0.0, false).pure_temp());
14772        assert!(SampledGraphKey::new(7, 0.8, 3, 0, 1.5, -1.0, false).pure_temp());
14773    }
14774
14775    /// The WIDENED capture regime (lane/step37-draft-graph-serving-20260830): truncation-
14776    /// filtered shapes are capturable — the filter runs IN-GRAPH (`filter_stats` +
14777    /// `gumbel_perturb_filtered_ctr`), so the draft draws from the same filtered
14778    /// distribution the accept test reconstructs. Penalties never are: the per-round
14779    /// history cannot be baked. The step37 vendor-default shape (temp 0.5 / top_p 0.9) is
14780    /// exactly the previously-excluded regime this lane exists to capture.
14781    #[test]
14782    fn filtered_regimes_are_capturable_penalties_never() {
14783        let vendor = SampledGraphKey::new(12345, 0.5, 3, 0, 0.9, 0.0, false);
14784        assert!(!vendor.pure_temp());
14785        assert!(vendor.filtered());
14786        assert!(
14787            vendor.graph_capturable(),
14788            "the vendor-default filtered shape must be capturable (default door state)",
14789        );
14790        assert!(pure_temp_key().graph_capturable());
14791        assert!(
14792            !pure_temp_key().filtered(),
14793            "pure-temp takes the legacy (filterless) capture body",
14794        );
14795        let pen = SampledGraphKey::new(12345, 0.5, 3, 0, 0.9, 0.0, true);
14796        assert!(
14797            !pen.graph_capturable(),
14798            "penalty history varies per round and can never be baked into a graph",
14799        );
14800    }
14801
14802    /// MEMRA_DEBUG_SPEC on a SAMPLED spec request past round 0: the print must render without
14803    /// indexing the empty greedy `preds` vector (it panicked the GPU worker before this lane).
14804    #[test]
14805    fn debug_print_survives_the_sampled_arm() {
14806        // round >= 1 with a pending bonus == base 1, sampled == `preds` empty.
14807        assert_eq!(debug_t_pred0(true, 1, 4242, &[]), "n/a");
14808        assert_eq!(debug_t_pred0(true, 2, 4242, &[]), "n/a");
14809        // round 0 without a pending bonus still reports last_pred, in both arms.
14810        assert_eq!(debug_t_pred0(true, 0, 4242, &[]), "4242");
14811        assert_eq!(debug_t_pred0(false, 0, 4242, &[7, 8]), "4242");
14812        // greedy keeps the real prediction it always printed.
14813        assert_eq!(debug_t_pred0(false, 1, 4242, &[7, 8]), "7");
14814        assert_eq!(debug_t_pred0(false, 2, 4242, &[7, 8]), "8");
14815    }
14816}