Skip to main content

memra_engine/
spec.rs

1//! Qwen3.5 MTP (NextN) greedy speculative decode (research/mtp/MTP-PLAN.md §A/§B/§C/§D).
2//!
3//! Greedy spec decode is MATHEMATICALLY EXACT: the accepted+bonus token stream is token-for-token
4//! identical to plain greedy `generate`. This module provides:
5//!   - `mtp_head_forward`  (§A, T=1): one NextN draft-token forward.
6//!   - `decode_step_t`     (§D.3, T=K+1): batched target verify forward, all-column logits.
7//!   - `generate_spec`     (§B): the draft/verify/accept/rollback orchestrator.
8//! Cache snapshot/rollback lives in cache.rs (§D.4). The MTP head uses its OWN scratch KV (§D.6),
9//! PERSISTENT over the committed sequence (see `MtpScratch`).
10
11use crate::Engine;
12use crate::cache::{Cache, KvLayer};
13use crate::forward::argmax;
14use crate::hybrid::{FullAttnLayer, HybridModel, LinearAttnLayer, Mixer, MtpHead};
15use cudarc::driver::CudaSlice;
16use std::sync::atomic::{AtomicU64, Ordering};
17
18/// Parse the documented `MEMRA_SPEC_REPLAY=1` rollback seam.
19///
20/// Keep this shared with serving admission so `=0` cannot select replay in one
21/// layer while another layer treats it as disabled.
22pub fn spec_replay_env_on(value: Option<&str>) -> bool {
23    value == Some("1")
24}
25
26pub fn spec_replay_env_enabled() -> bool {
27    let value = std::env::var("MEMRA_SPEC_REPLAY").ok();
28    spec_replay_env_on(value.as_deref())
29}
30
31/// One compact, anchor-bounded DSpark supervision record. `tokens[0]` is the anchor at p and
32/// `hidden` is its predecessor carrier h[p-1], matching the live NextN/DSpark pairing. Target
33/// rows p..p+gamma-1 score tokens p+1..p+gamma. They are the full-target softmax's top-k
34/// entries; `target_tail_probs[j]` is the probability mass outside those rows. All flattened
35/// target arrays are `[gamma, top_k]` in row-major order.
36pub struct DsparkAnchorRecord {
37    pub position: usize,
38    pub hidden: Vec<f32>,
39    pub tokens: Vec<u32>,
40    pub target_top_ids: Vec<u32>,
41    pub target_top_logits: Vec<f32>,
42    pub target_top_probs: Vec<f32>,
43    pub target_tail_probs: Vec<f32>,
44}
45
46fn dspark_sparse_softmax_topk(
47    logits: &[f32],
48    top_k: usize,
49    temperature: f32,
50) -> Result<(Vec<u32>, Vec<f32>, Vec<f32>, f32), Box<dyn std::error::Error>> {
51    if logits.is_empty() || top_k == 0 || top_k > logits.len() || temperature <= 0.0 {
52        return Err("invalid DSpark sparse-softmax shape or temperature".into());
53    }
54    if logits.iter().any(|value| !value.is_finite()) {
55        return Err("DSpark target logits contain a non-finite value".into());
56    }
57    let mut ranked: Vec<(u32, f32)> = logits
58        .iter()
59        .copied()
60        .enumerate()
61        .map(|(index, value)| (index as u32, value))
62        .collect();
63    let compare = |left: &(u32, f32), right: &(u32, f32)| {
64        right.1.total_cmp(&left.1).then(left.0.cmp(&right.0))
65    };
66    ranked.select_nth_unstable_by(top_k - 1, compare);
67    ranked[..top_k].sort_unstable_by(compare);
68
69    let max_logit = logits.iter().copied().fold(f32::NEG_INFINITY, f32::max);
70    let inv_temperature = 1.0f64 / temperature as f64;
71    let denominator: f64 = logits
72        .iter()
73        .map(|value| (((*value - max_logit) as f64) * inv_temperature).exp())
74        .sum();
75    let ids: Vec<u32> = ranked[..top_k].iter().map(|(index, _)| *index).collect();
76    let top_logits: Vec<f32> = ranked[..top_k].iter().map(|(_, value)| *value).collect();
77    let top_probs: Vec<f32> = top_logits
78        .iter()
79        .map(|value| ((((value - max_logit) as f64) * inv_temperature).exp() / denominator) as f32)
80        .collect();
81    let top_mass: f64 = top_probs.iter().map(|value| *value as f64).sum();
82    let tail = (1.0f64 - top_mass).clamp(0.0, 1.0) as f32;
83    Ok((ids, top_logits, top_probs, tail))
84}
85
86fn flatten_dspark_rows<T>(
87    rows: Vec<Option<Vec<T>>>,
88    position: usize,
89    label: &str,
90) -> Result<Vec<T>, Box<dyn std::error::Error>> {
91    let mut flattened = Vec::new();
92    for (slot, row) in rows.into_iter().enumerate() {
93        flattened.extend(
94            row.ok_or_else(|| format!("missing DSpark {label} at {position} slot {slot}"))?,
95        );
96    }
97    Ok(flattened)
98}
99
100/// H-SEED CONVENTION (MEMRA_SPEC_HPOST=1): feed the MTP head the POST-norm hidden — trunk rows
101/// hand over `output_norm(x)` and the draft chain recurrence hands over `shared_head_norm(h_nextn)`
102/// (= final_h) — matching the reference engines: llama.cpp #24025 ("qwen35: use post-norm hidden
103/// state for MTP", t_h_nextn is taken AFTER the final norm in both trunk and MTP graphs) and
104/// SGLang's qwen3_5_mtp (spec_info.hidden_states = the target model's post-norm output). memra's
105/// historical convention (default, MTP-PLAN §A) is PRE-norm x. Draft-quality-only: exactness is
106/// the verify's job either way; acceptance arbitrates. OnceLock: read once, hot-loop safe.
107pub(crate) fn spec_hpost() -> bool {
108    static H: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
109    *H.get_or_init(|| {
110        std::env::var("MEMRA_SPEC_HPOST")
111            .map(|v| v != "0")
112            .unwrap_or(false)
113    })
114}
115
116/// LEAN VERIFY (default ON since 2026-07-08; MEMRA_SPEC_LEAN=0 reverts — close35 lane): the verify m-scaling
117/// probe + nsys diff showed the verify t-path pays ~1.0ms/call at m=1 over eager decode on the
118/// 35B, and the kernels are NOT the cause (dev-MoE identical, kernel-time delta only +179us).
119/// The overhead is (a) ~250 extra cuMemsetD8Async/call from `e.zeros()` on buffers every kernel
120/// fully overwrites (~0.9ms host issue + ~0.35ms GPU) and (b) the t=1 FA rows dispatch (rows_v2 +
121/// combine_rows, +50us vs the eager fa_decode pair). This flag switches (a) fully-overwritten
122/// verify buffers to `e.uninit` (identical bytes: every element is written before read) and
123/// (b) t==1 verify FA to the eager `fa_decode` entry (byte-identical: kernel-check pins the
124/// rows-vs-loop identity and the per-row loop at t=1 IS fa_decode on the same q). Gates arbitrate.
125pub(crate) fn spec_lean() -> bool {
126    static L: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
127    // DEFAULT ON since 2026-07-08 (MEMRA_SPEC_LEAN=0 reverts): bit-identical (buffers fully
128    // overwritten; gates green incl maxdiff-identical run-gen) and measured +2.4% e2e p3 /
129    // +1.5% p2 at the daily 35B config. m=1 verify now costs eager-decode parity.
130    *L.get_or_init(|| {
131        std::env::var("MEMRA_SPEC_LEAN")
132            .map(|v| v != "0")
133            .unwrap_or(true)
134    })
135}
136
137/// SMALL-M BATCHED VERIFY (default ON since 2026-07-09; MEMRA_SPEC_M2=0 reverts — lane/spec-m2): extend the
138/// batched linear-attn verify arm down to t=2 and batch the MoE dev token loop over a
139/// grid.z=token axis at every verify t. The close35 m-scaling probe put the m=2 verify tier at
140/// x1.54 of m=1 (llama x1.14); the per-column linear chain (t<3) and the serial MoE dev token
141/// loop are the two launch-structure causes. Both changes are LAUNCH-STRUCTURE ONLY:
142/// (a) the batched conv's t<pad ring update is pure copies (ssm_conv_ring_rebuild from a cloned
143///     ring — the ring stores raw input columns); every arithmetic kernel is the same one the
144///     t>=3 arm already runs (matmul_decode_exact bit-identical at m=2-4, gdn_scan's internal
145///     t-loop == chained T=1 steps);
146/// (b) the MoE dev-rows twins run the serial loop's per-token warp program with tok-offset
147///     pointers (same sel/w/aq/ad bytes, same dot order, same slot-ordered FMA chain).
148/// Gates arbitrate: run-spec K=1..8 self-consistency (35B+9B), kernel-check, run-gen argmax.
149pub(crate) fn spec_m2() -> bool {
150    static M: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
151    // DEFAULT ON since 2026-07-09 (MEMRA_SPEC_M2=0 reverts): launch-structure only — t=2
152    // batched linear arm (ring-roll copies, zero new FP order) + MoE dev-rows kernels
153    // (grid.z=token, 4 launches/layer at any verify t). Acceptance bit-identical at every K;
154    // 35B p2 +3.4% / p3 +3.6%; the profitable-K plateau widens (new optimum K=3 at 223).
155    *M.get_or_init(|| {
156        std::env::var("MEMRA_SPEC_M2")
157            .map(|v| v != "0")
158            .unwrap_or(true)
159    })
160}
161pub(crate) fn spec_stream() -> bool {
162    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
163    *ON.get_or_init(|| std::env::var("MEMRA_SPEC_STREAM").as_deref() == Ok("1"))
164}
165pub(crate) fn spec_stream_m() -> usize {
166    static M: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
167    *M.get_or_init(|| {
168        std::env::var("MEMRA_SPEC_STREAM_M")
169            .ok()
170            .and_then(|v| v.parse().ok())
171            .unwrap_or(4)
172    })
173}
174pub(crate) fn spec_devacc() -> bool {
175    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
176    *ON.get_or_init(|| std::env::var("MEMRA_SPEC_DEVACC").as_deref() == Ok("1"))
177}
178
179/// GRAMMAR HOOK for constrained spec decode (lane/constrained-full, 2026-08-03). The engine
180/// stays llguidance-agnostic: the server adapts its per-session grammar state behind this
181/// trait. CONTRACT (the verify-side truncation rule — token-identical to constrained plain
182/// greedy decode): the exactness walk runs UNMASKED first; the hook then (a) truncates
183/// acceptance at the first grammar-illegal accepted token, and (b) when the truncation fired
184/// or the bonus is illegal, the engine recomputes that slot as the MASKED argmax of the
185/// target's own verify column (an unmasked argmax that is grammar-legal IS the masked argmax
186/// — masking only removes tokens — so the common case pays nothing). `consume` advances the
187/// state with each EMITTED token in order; EOS handling is the implementor's job (skip).
188pub trait SpecConstraint {
189    /// -inf the current state's banned ids on a HOST logits row (prompt-tail / init-feed
190    /// masked argmax).
191    fn mask_logits(&mut self, logits: &mut [f32]) -> Result<(), String>;
192    /// Packed 32-bit bitset words of the CURRENT state's allowed set (device-mask form).
193    fn mask_words(&mut self) -> Result<Vec<u32>, String>;
194    /// Is `tok` consumable in the CURRENT state?
195    fn is_allowed(&mut self, tok: u32) -> Result<bool, String>;
196    /// Advance the state with an emitted token.
197    fn consume(&mut self, tok: u32) -> Result<(), String>;
198
199    // --- DRAFT-SIDE MASKING (lane/draft-mask, 2026-08-04) ---
200    // The drafter proposed grammar-illegal tokens under tight schemas, so verify-side
201    // truncation cut nearly every round (measured acceptance 0.467-0.513 tight vs 0.62-0.82
202    // loose, research/constrained-full-20260803). These three methods let the engine mask the
203    // DRAFT model's own sampling with the grammar's legal set, so proposals are legal by
204    // construction. The state they walk is a SPECULATIVE CLONE of the session matcher — the
205    // real state is advanced only by `consume` (emitted tokens), so verify-side truncation
206    // stays the correctness backstop and the emitted stream is unchanged by construction
207    // (an accepted draft is the target's unmasked argmax AND grammar-legal, hence the masked
208    // argmax; a cut slot is recomputed as the masked argmax either way).
209    // Default impls = feature OFF (pre-lane behaviour: unmasked drafts).
210
211    /// Is draft-side masking available on this hook? Probed ONCE per burst, before the draft
212    /// graph is captured (the mask is an in-graph node — its presence is a capture-time shape).
213    fn draft_mask_enabled(&self) -> bool {
214        false
215    }
216    /// Start a draft chain: clone the CURRENT (committed) grammar state into the speculative
217    /// slot. Called once per spec round, before the first draft position.
218    fn draft_begin(&mut self) -> Result<(), String> {
219        Ok(())
220    }
221    /// Packed 32-bit bitset words of the SPECULATIVE state's allowed set (target-vocab ids),
222    /// for the draft position about to be sampled. `None` = draft masking off (no-op).
223    fn draft_mask_words(&mut self) -> Result<Option<Vec<u32>>, String> {
224        Ok(None)
225    }
226    /// Advance the SPECULATIVE state with a PROPOSED draft token. `false` = the chain cannot
227    /// continue (EOS proposed, or an unmasked position proposed something illegal) — the
228    /// engine stops drafting; the token already pushed still goes through verify.
229    fn draft_advance(&mut self, _tok: u32) -> Result<bool, String> {
230        Ok(false)
231    }
232}
233
234/// DRAFT-MASK UPLOAD (lane/draft-mask): pull the speculative state's allowed set (TARGET-id
235/// space) from the hook, project it into the DRAFT head's vocab space, and upload it into the
236/// stable device buffer the draft chain reads. Returns false when the chain must stop drafting:
237/// the hook handed out no mask, or NO draft-vocab row is grammar-legal at this position (a
238/// trimmed FR-Spec head genuinely cannot propose a legal token there — masking it would leave
239/// a fully-banned row whose argmax is meaningless, so the round drafts fewer tokens and the
240/// verify emits the masked argmax as usual).
241fn upload_draft_mask(
242    e: &Engine,
243    c: &mut dyn SpecConstraint,
244    dst: &mut CudaSlice<u32>,
245    d2t: Option<&Vec<u32>>,
246    d_vocab: usize,
247    words: usize,
248) -> Result<bool, Box<dyn std::error::Error>> {
249    let Some(tw) = c
250        .draft_mask_words()
251        .map_err(|e2| format!("constraint: {e2}"))?
252    else {
253        return Ok(false);
254    };
255    let bit = |t: usize| -> bool {
256        let w = t >> 5;
257        w < tw.len() && (tw[w] >> (t & 31)) & 1 == 1
258    };
259    let mut buf = vec![0u32; words];
260    match d2t {
261        // TRIMMED draft head: row i proposes target id d2t[i] — permute the mask accordingly.
262        Some(map) => {
263            for (i, &t) in map.iter().enumerate().take(d_vocab) {
264                if bit(t as usize) {
265                    buf[i >> 5] |= 1u32 << (i & 31);
266                }
267            }
268        }
269        // UNTRIMMED: draft ids ARE target ids; the packed words transfer verbatim (a short
270        // mask leaves the padded tail zeroed == banned, same rule as constrained::apply_mask).
271        None => {
272            let n = tw.len().min(words);
273            buf[..n].copy_from_slice(&tw[..n]);
274        }
275    }
276    if buf.iter().all(|w| *w == 0) {
277        return Ok(false);
278    }
279    e.htod_u32_into(dst, &buf)?;
280    Ok(true)
281}
282
283/// Keep the full token-embedding table in host memory and upload only the rows needed by each
284/// MTP/verify step. This is an exact memory-capacity seam for very large BF16 vocab tables: host
285/// gather expands the same source bits to f32, and only O(T*n_embd) bytes cross PCIe per step.
286/// CUDA-graph/round-stream draft paths require device token ids and therefore stay disabled.
287pub(crate) fn spec_host_embd() -> bool {
288    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
289    *ON.get_or_init(|| std::env::var("MEMRA_SPEC_HOST_EMBD").as_deref() == Ok("1"))
290}
291
292/// VERIFY-TIER TRUNK LAUNCH-FUSION (default ON since 2026-07-09; MEMRA_SPEC_FUSED_T=0 reverts — lane/close35b): extend
293/// the t=1 fused2/fused3 Q8_0 trunk launches to the batched verify tier (t=2-4, the K=1..3
294/// verify shapes). At t>1 the trunk pairs/triples (35B wqkv+wqkv_gate, wq/wk/wv,
295/// gate_shexp+up_shexp) each run a separate `matmul_decode_exact` — one q8_1 re-quantize of the
296/// SAME activation plus one _b2/_b4 launch per tensor. The fused twins share ONE quantize and
297/// ONE launch per group; per (tensor,token,row) the kernel body is q8_0_mmvq_batched verbatim
298/// with the identical row mapping -> BIT-IDENTICAL by construction (kernel-check pins it,
299/// run-spec K=1..8 + acceptance identity arbitrate e2e).
300pub(crate) fn spec_fused_t() -> bool {
301    static F: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
302    // DEFAULT ON since 2026-07-09 (MEMRA_SPEC_FUSED_T=0 reverts): verify t=2-4 trunk launch-fusion
303    // (fused2/fused3 Q8_0 batched twins, bit-identical by construction — m=1 block-offset split on
304    // the batched body). m=2 marginal token 2117->1762us; 35B daily: p3 +3.7% (crosses llama), p2 +5%.
305    *F.get_or_init(|| {
306        std::env::var("MEMRA_SPEC_FUSED_T")
307            .map(|v| v != "0")
308            .unwrap_or(true)
309    })
310}
311
312/// zeros/uninit switch for verify-path buffers that are FULLY OVERWRITTEN before any read.
313/// Only call this on such buffers — the lean contract is "identical bytes by construction".
314fn vbuf(e: &Engine, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
315    if spec_lean() { e.uninit(n) } else { e.zeros(n) }
316}
317
318/// Scratch KV for the MTP block (one full-attn layer).
319///
320/// PERSISTENT MODE (default, 2026-07-03 — the acceptance lever): sized cap = max_ctx and kept in
321/// sync with the COMMITTED sequence — slot p holds the MTP block's K/V for committed token p
322/// (roped p+1, the chain's rope convention), so the draft chain's self-attention sees the FULL
323/// committed history instead of only the current round's 1..K+1 chain tokens (the reference
324/// engine's "mtp_update" design). Entries come from two sources:
325///   - chain appends: accepted positions KEEP their chain-computed entries (embedding exact,
326///     hidden chain-approximate — the reference engine accepts the same);
327///   - `mtp_kv_fill` batches: prompt positions + the last-draft position on full accept, computed
328///     from EXACT trunk hiddens (K/V-only MTP-block pass, no attention/FFN/lm_head).
329/// Rejected drafts / p-min extras / pseudo-seed appends are all discarded by the round-start
330/// `set_len` truncation (the KvLayer len mechanism — §C rollback for the draft side).
331/// Multi-turn spec-decode session (2026-07-05): trunk Cache + persistent MTP draft scratch +
332/// the committed token list, alive across generate_spec_session calls. Turn N+1 primes ONLY its
333/// suffix (chunked continuation prime over the quantized past) and mtp_kv_fill's its suffix rows,
334/// then the round loop runs unchanged. `last_h` carries the pre-output_norm hidden of the last
335/// committed row across turns (the predecessor-pairing seed + fill anchor).
336/// Per-request sampling config for the sampled-spec serve path.
337#[derive(Clone, Copy, Debug)]
338pub struct SpecSampling {
339    pub temp: f32,
340    pub seed: u64,
341    pub top_k: i32,            // 0 = off
342    pub top_p: f32,            // 1.0 = off
343    pub min_p: f32,            // 0.0 = off
344    pub penalty_last_n: usize, // 0 = penalties off
345    pub penalty_repeat: f32,
346    pub penalty_freq: f32,
347    pub penalty_present: f32,
348}
349
350/// Tracked draft positions for [`SpecTelemetry`] (serve K defaults to 3; the run-spec gate
351/// sweeps K=1..8, and MEMRA_SPEC_CAPMAX defaults to 7 — 8 covers every tuned config).
352pub const SPEC_TELEM_POS: usize = 8;
353
354/// Always-on per-draft-position acceptance telemetry (lane/accept-telemetry, 2026-08-05 —
355/// the llama.cpp #26389 / vLLM spec-decode counter schema, upstream-sweeps 2026-08-05).
356/// Lives on the [`SpecSession`] and accumulates across bursts; the serve worker diffs a
357/// stashed copy per burst for its per-model /metrics aggregation and per-request usage.
358/// Same normalization as the `[spec-stats]` line: p-min-discarded chain tokens are counted
359/// in NEITHER drafted nor accepted.
360#[derive(Clone, Copy, Default, Debug)]
361pub struct SpecTelemetry {
362    /// verify rounds completed (a round-stream burst counts each of its M rounds).
363    pub rounds: u64,
364    /// tokens drafted / accepted across all rounds.
365    pub drafted: u64,
366    pub accepted: u64,
367    /// how often draft position j (0-based within a round's chain) was offered / accepted.
368    /// Positions >= SPEC_TELEM_POS are untracked (totals still count them). The opt-in
369    /// round-stream arm (MEMRA_SPEC_STREAM=1) reads back only totals, so under it these
370    /// arrays cover the standard-path rounds only and their sums may undercount the totals.
371    pub pos_drafted: [u64; SPEC_TELEM_POS],
372    pub pos_accepted: [u64; SPEC_TELEM_POS],
373}
374
375impl SpecTelemetry {
376    /// Fieldwise `self - prev` — the worker's per-burst delta off a copy stashed before the
377    /// burst call. Saturating: a caller diffing against the wrong snapshot gets zeros, not
378    /// a wrapped counter.
379    pub fn delta_since(&self, prev: &SpecTelemetry) -> SpecTelemetry {
380        let mut d = SpecTelemetry {
381            rounds: self.rounds.saturating_sub(prev.rounds),
382            drafted: self.drafted.saturating_sub(prev.drafted),
383            accepted: self.accepted.saturating_sub(prev.accepted),
384            ..Default::default()
385        };
386        for j in 0..SPEC_TELEM_POS {
387            d.pos_drafted[j] = self.pos_drafted[j].saturating_sub(prev.pos_drafted[j]);
388            d.pos_accepted[j] = self.pos_accepted[j].saturating_sub(prev.pos_accepted[j]);
389        }
390        d
391    }
392    /// Fieldwise `self += d` — the worker's per-model aggregation.
393    pub fn merge(&mut self, d: &SpecTelemetry) {
394        self.rounds += d.rounds;
395        self.drafted += d.drafted;
396        self.accepted += d.accepted;
397        for j in 0..SPEC_TELEM_POS {
398            self.pos_drafted[j] += d.pos_drafted[j];
399            self.pos_accepted[j] += d.pos_accepted[j];
400        }
401    }
402
403    /// Mean accepted draft-prefix length per verify round (tau).
404    pub fn tau(&self) -> f64 {
405        if self.rounds > 0 {
406            self.accepted as f64 / self.rounds as f64
407        } else {
408            0.0
409        }
410    }
411}
412
413/// Session-lifetime atomic acceptance counters. The verifier records only after the greedy or
414/// rejection-sampling walk has resolved on the host, so these relaxed increments add no GPU
415/// launch, synchronization, allocation, or ordering dependency to the numeric path.
416struct SpecTelemetryCounters {
417    rounds: AtomicU64,
418    drafted: AtomicU64,
419    accepted: AtomicU64,
420    pos_drafted: [AtomicU64; SPEC_TELEM_POS],
421    pos_accepted: [AtomicU64; SPEC_TELEM_POS],
422}
423
424impl Default for SpecTelemetryCounters {
425    fn default() -> Self {
426        Self {
427            rounds: AtomicU64::new(0),
428            drafted: AtomicU64::new(0),
429            accepted: AtomicU64::new(0),
430            pos_drafted: std::array::from_fn(|_| AtomicU64::new(0)),
431            pos_accepted: std::array::from_fn(|_| AtomicU64::new(0)),
432        }
433    }
434}
435
436impl SpecTelemetryCounters {
437    fn record_round(&self, drafted: usize, accepted: usize) {
438        debug_assert!(accepted <= drafted);
439        self.rounds.fetch_add(1, Ordering::Relaxed);
440        self.drafted.fetch_add(drafted as u64, Ordering::Relaxed);
441        self.accepted.fetch_add(accepted as u64, Ordering::Relaxed);
442        for counter in self.pos_drafted.iter().take(drafted) {
443            counter.fetch_add(1, Ordering::Relaxed);
444        }
445        for counter in self.pos_accepted.iter().take(accepted) {
446            counter.fetch_add(1, Ordering::Relaxed);
447        }
448    }
449
450    /// Round-stream keeps each round's accept length on device; retain exact scalar totals while
451    /// leaving the per-position arrays untouched, matching the pre-existing telemetry contract.
452    fn record_totals(&self, rounds: usize, drafted: usize, accepted: usize) {
453        self.rounds.fetch_add(rounds as u64, Ordering::Relaxed);
454        self.drafted.fetch_add(drafted as u64, Ordering::Relaxed);
455        self.accepted.fetch_add(accepted as u64, Ordering::Relaxed);
456    }
457
458    fn snapshot(&self) -> SpecTelemetry {
459        SpecTelemetry {
460            rounds: self.rounds.load(Ordering::Relaxed),
461            drafted: self.drafted.load(Ordering::Relaxed),
462            accepted: self.accepted.load(Ordering::Relaxed),
463            pos_drafted: std::array::from_fn(|j| self.pos_drafted[j].load(Ordering::Relaxed)),
464            pos_accepted: std::array::from_fn(|j| self.pos_accepted[j].load(Ordering::Relaxed)),
465        }
466    }
467}
468
469pub struct SpecSession {
470    pub(crate) cache: Cache,
471    pub(crate) scratch: MtpScratch,
472    /// Every token whose state the caches hold, in order (prompt turns + generated), INCLUDING
473    /// overshoot: spec commits accepted drafts past max_new; those rows are in the caches, so the
474    /// session must count them. Callers render output from this, not from their own echo.
475    pub committed: Vec<u32>,
476    /// Pre-output_norm hidden of the LAST committed row (device). None before the first turn.
477    pub(crate) last_h: Option<CudaSlice<f32>>,
478    /// Greedy argmax predicting the token AFTER committed.last() (from the last turn's final
479    /// logits). Fuels empty-suffix continuation bursts (serve): the next turn emits this token
480    /// first, feeds it, and the round loop resumes without any prime. None before the first turn.
481    pub next_pred: Option<u32>,
482    /// SAMPLED-SPEC stream continuity across bursts: Philox event counters persist here so a
483    /// session's randomness never repeats between generate_spec_session calls. (0,0) at admit.
484    pub sctr: u32,
485    pub uctr: u32,
486    /// PERSISTENT DRAFT-GRAPH CONTEXT (2026-08-01, the serve-burst fixed-cost fix): the captured
487    /// draft graph(s) + every device I/O buffer they bake, carried ACROSS generate_spec_session
488    /// calls. Before this, every serve burst re-captured the draft graph (2 warmup forwards +
489    /// instantiate) — measured ~16ms/burst on H100 q27 (MEMRA_SPEC_BURST sweep,
490    /// research/spec-serving-20260801). None before the first turn; error paths drop it
491    /// (next burst recaptures — serve retires errored sessions anyway).
492    pub(crate) draft_ctx: Option<DraftGraphCtx>,
493    /// PENDING-CARRY across bursts (2026-08-01, the serve burst-boundary fix): the bonus token
494    /// emitted by the last round but NOT committed to the caches. The old tail committed it with
495    /// a solo T=1 trunk pass (+ draft fill), and the next burst's setup fed the stashed next_pred
496    /// with ANOTHER solo pass — 2x ~11.5ms/burst measured on H100 q27 ([spec-setup] trace).
497    /// Carrying it lets the next empty-suffix greedy burst consume it as round-0 verify col 0,
498    /// exactly like a mid-burst full-accept boundary (no solo passes). INVARIANT: when set,
499    /// `committed` (== cache rows) EXCLUDES this token although it was already emitted in the
500    /// last burst's output, and `last_h` holds the hidden of the last COMMITTED row (its
501    /// predecessor — the chain-seed/fill anchor). `next_pred` is None (unknown without the
502    /// commit pass). Non-empty-suffix or sampled turns must flush first (spec_flush_pending);
503    /// generate_spec_session_sampled does this at entry, and serve parks only flushed sessions.
504    pub pending_tok: Option<u32>,
505    /// SESSION-AFFINITY TURN CHECKPOINT (lane/session-affinity, 2026-08-05): the state at this
506    /// turn's PROMPT-END boundary, retained so a later turn can REWIND here. See
507    /// [`SpecCheckpoint`]. Refreshed by every non-empty prime; None until the first one, and on
508    /// a rig too tight to hold it (a failed capture is silent — resume just isn't available).
509    pub(crate) turn_ckpt: Option<SpecCheckpoint>,
510    /// Session-lifetime acceptance telemetry. Relaxed atomics update at the host-side round
511    /// accounting the loop already does — no syncs, no allocation. NOTE a
512    /// pool-resumed session carries the PREVIOUS requests' counts; per-request consumers
513    /// diff with [`SpecTelemetry::delta_since`] around each burst.
514    telem: SpecTelemetryCounters,
515    /// PREFIX-CACHE publication request (lane/spec-prefix-cache): worker sets this to the
516    /// miss-LCP boundary before a cold burst; the prime captures at exactly that split (it must
517    /// coincide with the burst's `prime_split` or no capture happens). One-shot: consumed by the
518    /// prime, result lands in `boundary_capture`.
519    pub capture_at: Option<usize>,
520    /// The capture the last cold prime produced (see [`SpecBoundaryCapture`]). Worker takes it
521    /// post-burst to assemble the prefix entry. A failed capture is silent, like `turn_ckpt` —
522    /// publication just isn't available for that request.
523    pub boundary_capture: Option<SpecBoundaryCapture>,
524}
525impl SpecSession {
526    /// Context capacity of the session's caches (the server's ContextFull guard).
527    pub fn cache_max_ctx(&self) -> usize {
528        self.cache.max_ctx
529    }
530    /// Read access to the live trunk cache (lane/spec-prefix-cache): the worker slices
531    /// full-attn KV rows `[0..capture.pos)` out of it when publishing a boundary capture —
532    /// those rows are append-only for the session's lifetime (rollbacks never truncate below
533    /// the prime boundary), so no copy was taken at prime time.
534    pub fn cache_ref(&self) -> &Cache {
535        &self.cache
536    }
537    /// Snapshot the session's process-local acceptance counters for per-burst diffing.
538    pub fn telemetry(&self) -> SpecTelemetry {
539        self.telem.snapshot()
540    }
541    /// Committed position this session can REWIND to (its retained prompt-end boundary), if any.
542    /// A request whose prompt matches `committed[..pos]` exactly can resume from here — see
543    /// `spec_rewind_to_checkpoint`.
544    pub fn rewind_pos(&self) -> Option<usize> {
545        self.turn_ckpt.as_ref().map(|c| c.pos)
546    }
547    /// Whether every ring-backed trunk/draft row needed by the retained checkpoint is resident.
548    pub fn rewind_is_resident(&self) -> bool {
549        self.turn_ckpt.as_ref().is_some_and(|ckpt| {
550            self.cache.can_rollback(&ckpt.snap, 0) && self.scratch.can_rewind_to(ckpt.pos)
551        })
552    }
553    /// Is this session in the DEMOTION-READY shape (see [`SpecSession::into_demoted`])?
554    /// `false` means a carried pending must be flushed first (`spec_flush_pending`), or the
555    /// session has never run a turn and has no prediction to hand over.
556    pub fn demote_ready(&self) -> bool {
557        self.pending_tok.is_none() && self.next_pred.is_some()
558    }
559    /// Does this session hold a carried pending bonus (flush required before a handoff/park)?
560    pub fn has_pending(&self) -> bool {
561        self.pending_tok.is_some()
562    }
563    /// Committed row count == cache rows (the session invariant), for the caller's own
564    /// `fed`-length cross-check at a handoff boundary.
565    pub fn committed_len(&self) -> usize {
566        self.committed.len()
567    }
568    /// DEMOTION HANDOFF (lane/spec-gate, 2026-08-07): consume this session and hand its trunk
569    /// cache + next-token prediction to the plain batched-decode path.
570    ///
571    /// WHY THIS IS EXACT (greedy). The invariant at a burst boundary is `cache.pos ==
572    /// committed.len()`: every committed row has trunk KV + recurrent state, exactly as a plain
573    /// tokenwise prime of the same `committed` sequence would have left it (that is the
574    /// session-tail contract, and the same property `spec_rewind_to_checkpoint` and the reuse
575    /// pool already rely on). `next_pred` is the argmax of the verify's logits for the LAST
576    /// committed row — and verify-column logits are bit-identical to plain decode's logits at
577    /// that position, because `matmul_decode_exact` bit-identity IS the basis of the greedy
578    /// accept walk. So handing (cache, next_pred) to the batched path continues the stream from
579    /// a state indistinguishable from one the batched path produced itself: the batched tick
580    /// emits `next_pred`, feeds it into this same cache, and decodes on.
581    ///
582    /// `None` when the session is not in the handoff shape — a carried pending (its bonus row is
583    /// NOT in the cache, so `spec_flush_pending` must commit it first) or no `next_pred` yet
584    /// (never bursted). Callers must not force it: a half-committed cache handed to the batched
585    /// path would silently skip a token.
586    ///
587    /// The MTP draft scratch, the persistent draft-graph context and the turn checkpoint are
588    /// DROPPED here (freeing their VRAM): the batched path never drafts, and this handoff is
589    /// one-way by design — there is no cheap symmetric re-promotion (rebuilding the draft KV
590    /// would mean an `mtp_kv_fill` over the whole committed history).
591    pub fn into_demoted(self) -> Option<(Cache, u32)> {
592        if self.pending_tok.is_some() {
593            return None;
594        }
595        let np = self.next_pred?;
596        debug_assert_eq!(
597            self.cache.pos,
598            self.committed.len(),
599            "demotion handoff: cache rows != committed tokens"
600        );
601        Some((self.cache, np))
602    }
603    /// Pool-resume hook (audit Q2): clear the parked draft-graph failure memoization so a
604    /// NEW request resuming this session gets one fresh capture chance — a transient-pressure
605    /// capture failure must not persist for the pool's whole lifetime (the TRT #16072 class).
606    /// Logs once iff a flag was actually set; a no-fallback resume is silent and free.
607    pub fn reset_graph_fallback_on_resume(&mut self) {
608        if let Some(line) = self
609            .draft_ctx
610            .as_mut()
611            .and_then(|c| c.failed.reset_on_resume())
612        {
613            eprintln!("{line}");
614        }
615    }
616}
617
618/// A session's PROMPT-END boundary state, the rewind target for session-affinity resume.
619///
620/// WHY THIS BOUNDARY, AND WHY IT IS THE ONLY ONE WORTH KEEPING. The rewrite class this lane
621/// exists for (a client that strips `<think>` blocks out of prior assistant turns) mutates the
622/// text the session GENERATED, never the prompt it was given. So turn N's prompt agrees with
623/// turn N-1's committed tokens up to almost exactly where turn N-1's generation began — the
624/// prompt-end boundary. Keeping a checkpoint there means the next turn re-primes only its own
625/// delta (the rewritten answer + the new user turn) instead of the whole conversation.
626///
627/// WHAT IT MUST HOLD. Full-attn KV is append-only and position-addressed, so rewinding it is a
628/// `len` truncation (no data). Linear-attn (GDN) conv/ssm state is mutated IN PLACE with no
629/// position index, so it must be a real device COPY — that copy is the entire reason a spec
630/// session could not previously rewind. The MTP draft scratch needs no copy either: its rows
631/// below the boundary were written by this turn's fill and are never revisited (the per-round
632/// true-hidden refresh only rewrites the CURRENT burst's committed positions), so rewinding it
633/// is also just a `len` reset. `last_h` is the hidden of the last row below the boundary — the
634/// predecessor-pairing anchor the next prime's fill reads for its first row.
635///
636/// COST: one `Cache::snapshot` per TURN, on a code path that already takes one per ROUND.
637pub(crate) struct SpecCheckpoint {
638    snap: crate::cache::CacheSnapshot,
639    /// Committed length at the boundary (== cache.pos there, the session invariant).
640    pos: usize,
641    /// Pre-output_norm hidden of row `pos - 1`.
642    last_h: CudaSlice<f32>,
643}
644
645/// PREFIX-CACHE BOUNDARY CAPTURE (lane/spec-prefix-cache, 2026-08-14): the state a spec session
646/// records at its cold-prime split so the WORKER can publish a cross-request prefix entry —
647/// the commit-gated-publication port (research/cache-spec-design-20260814/PORT-PLAN.md item 1).
648/// Only the pieces that are DESTROYED by continuing the prime need copies here: the in-place
649/// GDN conv/ssm states (via `Cache::snapshot`, same mechanism as [`SpecCheckpoint`]) and the
650/// boundary logits. Full-attn KV rows `[0..pos)` and draft-scratch rows `[0..pos)` are
651/// append-only for the session's lifetime (rollbacks never truncate below the prime boundary),
652/// so the worker slices those from the live caches post-burst instead of copying at prime time.
653pub struct SpecBoundaryCapture {
654    pub snap: crate::cache::CacheSnapshot,
655    /// Token boundary (== cache.pos at capture; == the worker's miss-LCP split).
656    pub pos: usize,
657    /// Full-vocab logits after the prefix prime — the entry's boundary logits.
658    pub logits: Vec<f32>,
659}
660
661struct SpecPipeTraceClock {
662    pair: usize,
663    started: std::time::Instant,
664}
665
666#[derive(Clone)]
667struct SpecPipeTraceCtx {
668    clock: std::sync::Arc<SpecPipeTraceClock>,
669    round: usize,
670    lane: usize,
671}
672
673struct SpecPipeTraceMarker {
674    trace: SpecPipeTraceCtx,
675    phase: &'static str,
676    edge: &'static str,
677    slot: Option<usize>,
678}
679
680unsafe extern "C" fn spec_pipe_trace_marker(raw: *mut std::ffi::c_void) {
681    let marker = unsafe { Box::from_raw(raw.cast::<SpecPipeTraceMarker>()) };
682    let lane = if marker.trace.lane == 0 { "A" } else { "B" };
683    let slot = marker
684        .slot
685        .map(|v| v.to_string())
686        .unwrap_or_else(|| "-".into());
687    let t_ms = marker.trace.clock.started.elapsed().as_secs_f64() * 1e3;
688    use std::io::Write as _;
689    let stderr = std::io::stderr();
690    let mut stderr = stderr.lock();
691    let _ = writeln!(
692        stderr,
693        "[spec-pipe-timeline] pair={} round={} lane={lane} phase={} edge={} \
694         slot={slot} t_ms={t_ms:.3}",
695        marker.trace.clock.pair, marker.trace.round, marker.phase, marker.edge,
696    );
697}
698
699fn enqueue_spec_pipe_trace_marker(
700    stream: &cudarc::driver::CudaStream,
701    trace: Option<&SpecPipeTraceCtx>,
702    phase: &'static str,
703    edge: &'static str,
704    slot: Option<usize>,
705) -> Result<(), Box<dyn std::error::Error>> {
706    let Some(trace) = trace else {
707        return Ok(());
708    };
709    let marker = Box::new(SpecPipeTraceMarker {
710        trace: trace.clone(),
711        phase,
712        edge,
713        slot,
714    });
715    let raw = Box::into_raw(marker);
716    let result = unsafe {
717        cudarc::driver::result::stream::launch_host_function(
718            stream.cu_stream(),
719            spec_pipe_trace_marker,
720            raw.cast(),
721        )
722    };
723    if let Err(err) = result {
724        unsafe {
725            drop(Box::from_raw(raw));
726        }
727        return Err(err.into());
728    }
729    Ok(())
730}
731
732#[derive(Default)]
733struct SpecPipeProgress {
734    setup_done: [bool; 2],
735    draft_done: [usize; 2],
736    stage0_done: [usize; 2],
737    verify_done: [usize; 2],
738    accept_done: [usize; 2],
739    finished: [bool; 2],
740    aborted: bool,
741}
742
743/// Host-side issue coordinator for the reduced two-session speculative pipeline. Each session
744/// keeps its existing call stack and round locals; this object only orders phase entry. The
745/// primary mutex spans whole draft/accept/tail issue regions so Engine's single-stream scratch
746/// cannot be interleaved by the two host threads.
747struct SpecPipeSync {
748    progress: std::sync::Mutex<SpecPipeProgress>,
749    changed: std::sync::Condvar,
750    primary: std::sync::Mutex<()>,
751    trace: Option<std::sync::Arc<SpecPipeTraceClock>>,
752}
753
754impl SpecPipeSync {
755    fn new() -> Self {
756        static TRACE_PAIR: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
757        let trace = (std::env::var("MEMRA_SPEC_PIPE_TRACE").as_deref() == Ok("1")).then(|| {
758            std::sync::Arc::new(SpecPipeTraceClock {
759                pair: TRACE_PAIR.fetch_add(1, std::sync::atomic::Ordering::Relaxed) + 1,
760                started: std::time::Instant::now(),
761            })
762        });
763        Self {
764            progress: std::sync::Mutex::new(SpecPipeProgress::default()),
765            changed: std::sync::Condvar::new(),
766            primary: std::sync::Mutex::new(()),
767            trace,
768        }
769    }
770}
771
772#[derive(Clone)]
773struct SpecPipeLane {
774    sync: std::sync::Arc<SpecPipeSync>,
775    lane: usize,
776}
777
778impl SpecPipeLane {
779    fn peer(&self) -> usize {
780        1 - self.lane
781    }
782
783    fn aborted() -> Box<dyn std::error::Error> {
784        "paired speculative peer aborted".into()
785    }
786
787    fn trace(&self, round: usize) -> Option<SpecPipeTraceCtx> {
788        self.sync.trace.as_ref().map(|clock| SpecPipeTraceCtx {
789            clock: clock.clone(),
790            round,
791            lane: self.lane,
792        })
793    }
794
795    fn setup_begin(&self) -> Result<(), Box<dyn std::error::Error>> {
796        let mut p = self.sync.progress.lock().unwrap();
797        while !p.aborted && self.lane == 1 && !p.setup_done[0] && !p.finished[0] {
798            p = self.sync.changed.wait(p).unwrap();
799        }
800        if p.aborted {
801            Err(Self::aborted())
802        } else {
803            Ok(())
804        }
805    }
806
807    fn setup_end(&self) {
808        let mut p = self.sync.progress.lock().unwrap();
809        p.setup_done[self.lane] = true;
810        self.sync.changed.notify_all();
811    }
812
813    fn draft_begin(
814        &self,
815        round: usize,
816    ) -> Result<std::sync::MutexGuard<'_, ()>, Box<dyn std::error::Error>> {
817        let peer = self.peer();
818        let mut p = self.sync.progress.lock().unwrap();
819        loop {
820            if p.aborted {
821                return Err(Self::aborted());
822            }
823            let setup_ready =
824                (p.setup_done[0] || p.finished[0]) && (p.setup_done[1] || p.finished[1]);
825            let prior_ready = p.accept_done[self.lane] >= round
826                && (p.accept_done[peer] >= round || p.finished[peer]);
827            let turn_ready = if self.lane == 0 {
828                true
829            } else {
830                p.draft_done[0] > round || p.finished[0]
831            };
832            if setup_ready && prior_ready && turn_ready {
833                break;
834            }
835            p = self.sync.changed.wait(p).unwrap();
836        }
837        drop(p);
838        Ok(self.sync.primary.lock().unwrap())
839    }
840
841    fn draft_end(&self, round: usize) {
842        let mut p = self.sync.progress.lock().unwrap();
843        p.draft_done[self.lane] = round + 1;
844        self.sync.changed.notify_all();
845    }
846
847    /// Admit stage 0 and return whether this lane owns the interval's one reverse fence.
848    /// Lane B releases as soon as lane A has issued its boundary TX, not after A's full body.
849    fn stage0_begin(&self, round: usize) -> Result<bool, Box<dyn std::error::Error>> {
850        let peer = self.peer();
851        let mut p = self.sync.progress.lock().unwrap();
852        loop {
853            if p.aborted {
854                return Err(Self::aborted());
855            }
856            let ready = if self.lane == 0 {
857                p.draft_done[0] > round && (p.draft_done[1] > round || p.finished[1])
858            } else {
859                p.draft_done[1] > round && (p.stage0_done[0] > round || p.finished[0])
860            };
861            if ready {
862                return Ok(self.lane == 0 || p.finished[peer]);
863            }
864            p = self.sync.changed.wait(p).unwrap();
865        }
866    }
867
868    fn stage0_end(&self, round: usize) {
869        let mut p = self.sync.progress.lock().unwrap();
870        p.stage0_done[self.lane] = round + 1;
871        self.sync.changed.notify_all();
872    }
873
874    /// Stage 1 is single-owner per engine. A proceeds immediately after its own ticket; B waits
875    /// for A's full stage1/head issue so only A.S1 and B.S0 can overlap.
876    fn stage1_begin(&self, round: usize) -> Result<(), Box<dyn std::error::Error>> {
877        let mut p = self.sync.progress.lock().unwrap();
878        while !p.aborted
879            && !(p.stage0_done[self.lane] > round
880                && (self.lane == 0 || p.verify_done[0] > round || p.finished[0]))
881        {
882            p = self.sync.changed.wait(p).unwrap();
883        }
884        if p.aborted {
885            Err(Self::aborted())
886        } else {
887            Ok(())
888        }
889    }
890
891    fn verify_end(&self, round: usize) {
892        let mut p = self.sync.progress.lock().unwrap();
893        p.verify_done[self.lane] = round + 1;
894        self.sync.changed.notify_all();
895    }
896
897    fn accept_begin(
898        &self,
899        round: usize,
900    ) -> Result<std::sync::MutexGuard<'_, ()>, Box<dyn std::error::Error>> {
901        let mut p = self.sync.progress.lock().unwrap();
902        loop {
903            if p.aborted {
904                return Err(Self::aborted());
905            }
906            let ready = if self.lane == 0 {
907                p.verify_done[0] > round && (p.verify_done[1] > round || p.finished[1])
908            } else {
909                p.verify_done[1] > round && (p.accept_done[0] > round || p.finished[0])
910            };
911            if ready {
912                break;
913            }
914            p = self.sync.changed.wait(p).unwrap();
915        }
916        drop(p);
917        Ok(self.sync.primary.lock().unwrap())
918    }
919
920    fn accept_end(&self, round: usize) {
921        let mut p = self.sync.progress.lock().unwrap();
922        p.accept_done[self.lane] = round + 1;
923        self.sync.changed.notify_all();
924    }
925
926    fn primary(&self) -> std::sync::MutexGuard<'_, ()> {
927        self.sync.primary.lock().unwrap()
928    }
929
930    fn finish(&self, failed: bool) {
931        let mut p = self.sync.progress.lock().unwrap();
932        p.finished[self.lane] = true;
933        p.aborted |= failed;
934        self.sync.changed.notify_all();
935    }
936}
937
938struct SpecPipeFinish<'a> {
939    lane: &'a SpecPipeLane,
940    closed: bool,
941}
942
943impl<'a> SpecPipeFinish<'a> {
944    fn new(lane: &'a SpecPipeLane) -> Self {
945        Self {
946            lane,
947            closed: false,
948        }
949    }
950
951    fn close(&mut self, failed: bool) {
952        self.lane.finish(failed);
953        self.closed = true;
954    }
955}
956
957impl Drop for SpecPipeFinish<'_> {
958    fn drop(&mut self) {
959        if !self.closed {
960            self.lane.finish(true);
961        }
962    }
963}
964
965/// Scoped transfer of one exclusively-borrowed session to the second host issue thread.
966/// `CudaGraph` is not marked Send by cudarc because its raw driver handles carry no automatic
967/// trait. CUDA driver graph handles are context-scoped rather than OS-thread-affine; the caller
968/// binds that context before touching the session, joins before returning, and never aliases the
969/// pointer. Keep this exception local to the experimental pair call instead of marking the public
970/// session type Send.
971struct SpecPipeSessionPtr(*mut SpecSession);
972
973unsafe impl Send for SpecPipeSessionPtr {}
974
975impl SpecPipeSessionPtr {
976    unsafe fn get_mut(&mut self) -> &mut SpecSession {
977        unsafe { &mut *self.0 }
978    }
979}
980
981/// Per-session persistent draft-graph context: the captured CUDA graph(s) plus the device
982/// buffers whose POINTERS the capture bakes. Reuse legality: the greedy capture bakes only
983/// session-stable pointers (the session's own MtpScratch KV — allocated once, never realloc'd;
984/// the model's resident embedding; the process-wide OnceLock p_min) and the g_* buffers held
985/// HERE — so one capture serves the session's whole lifetime. The sampled capture additionally
986/// bakes (seed, temp) as capture-time constants and needs k q-slots — keyed by `s_key`, dropped
987/// and recaptured when a pool-resumed request changes them. `*_failed` memoizes a failed capture
988/// so the eager fallback doesn't pay a doomed capture attempt every burst.
989pub(crate) struct DraftGraphCtx {
990    g_tok: CudaSlice<u32>,
991    g_pos: CudaSlice<i32>,
992    g_seed: CudaSlice<f32>,
993    g_p: CudaSlice<f32>,
994    g_ctr: CudaSlice<u32>,
995    g_q: CudaSlice<f32>,
996    g_perturb: CudaSlice<f32>,
997    q_slots: Vec<CudaSlice<f32>>,
998    /// DRAFT-SIDE GRAMMAR MASK (lane/draft-mask): packed allowed-set words over the DRAFT
999    /// head's vocab, at a STABLE address so the captured draft graph's mask node reads the
1000    /// per-position contents the host re-uploads before each replay (the graph-promote
1001    /// pattern from decode.rs). Empty unless the session drafts under a grammar.
1002    g_dmask: CudaSlice<u32>,
1003    /// was `graph` captured WITH the mask node? A parked graph of the wrong shape is dropped.
1004    graph_masked: bool,
1005    graph: Option<cudarc::driver::CudaGraph>,
1006    graph_s: Option<cudarc::driver::CudaGraph>,
1007    /// Failed-capture memoization for both graphs — LOUD on flip, cleared on pool resume
1008    /// (audit Q2, the TRT #16072 silent-permanent-coverage-loss class).
1009    failed: DraftGraphFallback,
1010    /// (seed, temp.to_bits(), k) baked into graph_s at its capture.
1011    s_key: Option<(u64, u32, usize)>,
1012    /// CAPTURE-RETAIN keepers (#68 root cause, 2026-08-04): the warmup-run transients whose
1013    /// pool addresses the captured graph(s) bake. Without these, the transients return to the
1014    /// pool at capture-body exit and later work (burst-boundary prime/fill/commit passes, or a
1015    /// co-served session in the worker) reuses those addresses — the persisted graph's replay
1016    /// then reads/writes live unrelated buffers (exactness corruption, first seen as the ST
1017    /// serve-spec 4B graph-arm corruption; one-shot CLI calls never re-shuffled the pool, which
1018    /// is why run-spec K=1..8 passed on the same checkpoint). Same fix class as
1019    /// capture_graph_retained's gemma/decode.rs sites — hold as long as the graph replays.
1020    keeper: Vec<Box<dyn std::any::Any + Send>>,
1021    keeper_s: Vec<Box<dyn std::any::Any + Send>>,
1022}
1023
1024/// Failed-capture memoization for the two draft graphs (audit Q2, 2026-08-05 — the
1025/// TRT #16072 trap class: pressure-triggered, silent, long-lived coverage loss).
1026///
1027/// Three contracts:
1028/// - LOUD FLIP: `mark_*` returns the warn line exactly on the false→true transition
1029///   (returned, not printed, so the once-per-flip contract is unit-testable); the caller
1030///   `eprintln!`s it UNCONDITIONALLY — a dropped draft graph is never silent. Re-marking
1031///   an already-failed graph returns None (the per-burst memoization that keeps the eager
1032///   fallback from paying a doomed capture attempt every burst).
1033/// - RESET ON RESUME: `reset_on_resume` clears both flags — a parked session resumed by a
1034///   NEW request gets one fresh capture chance instead of carrying a transient-pressure
1035///   failure for the pool's whole lifetime. Returns the note line only when a flag was
1036///   actually set (quiet on the common clean-resume path).
1037/// - Shape-change clears (`clear_*`) stay silent, exactly as before: they precede a fresh
1038///   capture attempt whose own failure would re-flip loudly.
1039#[derive(Default)]
1040pub(crate) struct DraftGraphFallback {
1041    greedy: bool,
1042    sampled: bool,
1043}
1044impl DraftGraphFallback {
1045    fn mark_greedy(&mut self, reason: &str) -> Option<String> {
1046        if self.greedy {
1047            return None;
1048        }
1049        self.greedy = true;
1050        Some(format!(
1051            "[spec] WARN: draft-graph capture failed ({reason}); eager fallback until session resume"
1052        ))
1053    }
1054    fn mark_sampled(&mut self, reason: &str) -> Option<String> {
1055        if self.sampled {
1056            return None;
1057        }
1058        self.sampled = true;
1059        Some(format!(
1060            "[spec] WARN: sampled draft-graph capture failed ({reason}); eager fallback until session resume"
1061        ))
1062    }
1063    fn greedy_failed(&self) -> bool {
1064        self.greedy
1065    }
1066    fn sampled_failed(&self) -> bool {
1067        self.sampled
1068    }
1069    fn clear_greedy(&mut self) {
1070        self.greedy = false;
1071    }
1072    fn clear_sampled(&mut self) {
1073        self.sampled = false;
1074    }
1075    /// Pool-resume reset: both graphs get a fresh capture chance. Some(note) iff any flag
1076    /// was set (so clean resumes stay quiet).
1077    pub(crate) fn reset_on_resume(&mut self) -> Option<String> {
1078        if !self.greedy && !self.sampled {
1079            return None;
1080        }
1081        let which = match (self.greedy, self.sampled) {
1082            (true, true) => "greedy+sampled",
1083            (true, false) => "greedy",
1084            _ => "sampled",
1085        };
1086        self.greedy = false;
1087        self.sampled = false;
1088        Some(format!(
1089            "[spec] draft-graph fallback reset on session resume ({which}); recapture eligible"
1090        ))
1091    }
1092}
1093
1094impl DraftGraphCtx {
1095    fn new(e: &Engine, n_embd: usize, qlen: usize) -> Result<Self, Box<dyn std::error::Error>> {
1096        Ok(DraftGraphCtx {
1097            g_tok: e.alloc_u32_zeroed(1)?,
1098            g_pos: e.htod_i32(&[0])?,
1099            g_seed: e.zeros(n_embd)?,
1100            g_p: e.zeros(1)?,
1101            g_ctr: e.alloc_u32_zeroed(1)?,
1102            g_q: e.zeros(qlen)?,
1103            g_perturb: e.zeros(qlen)?,
1104            q_slots: Vec::new(),
1105            g_dmask: e.alloc_u32_zeroed(1)?,
1106            graph_masked: false,
1107            graph: None,
1108            graph_s: None,
1109            failed: DraftGraphFallback::default(),
1110            s_key: None,
1111            keeper: Vec::new(),
1112            keeper_s: Vec::new(),
1113        })
1114    }
1115}
1116
1117pub(crate) struct MtpScratch {
1118    kv: KvLayer,
1119    /// Logical row capacity. On the graph/DC draft path it also doubles as fa_decode_dc's
1120    /// bucket_max: n_splits is sized from it ONCE, so the graph captured at round 0 stays valid
1121    /// for every later t_kv. Step35 refuses that path and may back this logical extent with the
1122    /// smaller host-indexed SWA ring instead.
1123    cap: usize,
1124}
1125
1126fn mtp_scratch_layout(
1127    cfg: &memra_gguf::config::ModelConfig,
1128    geom: Option<&crate::hybrid::DraftGeom>,
1129) -> (usize, usize, usize, usize) {
1130    // Student draft heads carry fewer KV heads (head_dim unchanged) -> smaller scratch rows.
1131    let n_head_kv = geom.map(|g| g.n_head_kv).unwrap_or(cfg.n_head_kv as usize);
1132    let head_dim_k = cfg.head_dim_k as usize;
1133    let head_dim_v = cfg.head_dim_v as usize;
1134    assert!(
1135        head_dim_k % 32 == 0 && head_dim_v % 32 == 0,
1136        "KVQUANT requires head_dim%32==0 (MTP scratch)"
1137    );
1138    let kv_dim_k = head_dim_k * n_head_kv;
1139    let kv_dim_v = head_dim_v * n_head_kv;
1140    // The fp8-KV arm deliberately does not reach the draft scratch; keep the exact format
1141    // policy shared with `MtpScratch::new` so admission scales the same allocation.
1142    let (kbb, vbb) = crate::kv_blk_bytes();
1143    let k_tok_bytes = (kv_dim_k / 32) * kbb;
1144    let v_tok_bytes = (kv_dim_v / 32) * vbb;
1145    (kv_dim_k, kv_dim_v, k_tok_bytes, v_tok_bytes)
1146}
1147
1148impl MtpScratch {
1149    fn new(
1150        e: &Engine,
1151        cfg: &memra_gguf::config::ModelConfig,
1152        cap: usize,
1153        geom: Option<&crate::hybrid::DraftGeom>,
1154    ) -> Result<Self, Box<dyn std::error::Error>> {
1155        // env-selected KV formats (default 34/24). The fp8-KV arm (MEMRA_KV_FP8) deliberately
1156        // does NOT reach the draft scratch: fp8 drafts drifted acceptance 69-88% -> 46%
1157        // (2026-07-12 A/B); the scratch is tiny, so it keeps baseline q8_0/q5_1 numerics
1158        // while the TRUNK cache carries the fp8 depth win. Scratch append/fa pass g=false.
1159        let (kv_dim_k, kv_dim_v, k_tok_bytes, v_tok_bytes) = mtp_scratch_layout(cfg, geom);
1160        let ring = if crate::cache::swa_ring_on() && cfg.arch.is_step35() {
1161            let window = cfg.step35.as_ref().unwrap().sliding_window as usize;
1162            Some(crate::cache::KvRing::new(
1163                crate::cache::swa_ring_rows(window, cap),
1164                window,
1165            ))
1166        } else {
1167            None
1168        };
1169        let alloc_rows = ring.as_ref().map(crate::cache::KvRing::rows).unwrap_or(cap);
1170        Ok(MtpScratch {
1171            kv: KvLayer {
1172                k: e.alloc_u8(alloc_rows * k_tok_bytes)?,
1173                v: e.alloc_u8(alloc_rows * v_tok_bytes)?,
1174                kv_dim_k,
1175                kv_dim_v,
1176                k_tok_bytes,
1177                v_tok_bytes,
1178                len: 0,
1179                ring,
1180                len_d: e.htod_i32(&[0])?,
1181            },
1182            cap,
1183        })
1184    }
1185    /// Set BOTH length counters: the host mirror AND the device len_d the captured append/fa read
1186    /// (a 4-byte in-place htod — the counter pointer is baked into the graph, never realloc'd).
1187    /// This is the ONLY truncation/rollback mechanism the persistent draft KV needs.
1188    fn set_len(&mut self, e: &Engine, n: usize) -> Result<(), Box<dyn std::error::Error>> {
1189        if self
1190            .kv
1191            .ring
1192            .as_ref()
1193            .is_some_and(|ring| !ring.can_rewind_to(n))
1194        {
1195            return Err("SWA ring MTP checkpoint has been lapped; full re-prime required".into());
1196        }
1197        self.kv.len = n;
1198        e.set_i32_one(&mut self.kv.len_d, n as i32)
1199    }
1200
1201    fn can_rewind_to(&self, n: usize) -> bool {
1202        self.kv
1203            .ring
1204            .as_ref()
1205            .is_none_or(|ring| ring.can_rewind_to(n))
1206    }
1207}
1208
1209/// Retained verify intermediates for the REPLAY-FREE partial accept (2026-07-03, the profiled
1210/// #1 spec cost at long ctx: the partial-accept replay was a DUPLICATE trunk pass — ~0.54 extra
1211/// full weight reads per round — recomputing columns the verify had already produced
1212/// bit-identically). Holds, per linear layer, everything needed to rebuild its recurrent state
1213/// to "after the first j verify columns" WITHOUT re-running the trunk:
1214/// - BATCHED-path layers (`gdn`): the exact token-major inputs the round's ONE gdn_scan
1215///   consumed. A prefix re-run of the SAME kernel (t=j) from the snapshot state is bit-identical
1216///   to the first j iterations of the verify's scan — the kernel's t-loop carries state in
1217///   registers and iteration t never depends on T. `qkv_mixed` (the conv input) feeds the
1218///   pure-copy ring rebuild.
1219/// - PER-COLUMN-path layers (`cols`): dtod clones of (conv_state, ssm_state) taken after each
1220///   column 0..t-2 — pure copies of the actual chain states (the last column is never a rebuild
1221///   target: j <= t-1).
1222/// Full-attn layers need nothing: their verify KV rows are bit-identical to eager's (the
1223/// decode-exact contract; verify-probe pins it), so rollback = len truncation.
1224struct GdnStash {
1225    qkv_mixed: CudaSlice<f32>, // [t, conv_dim] token-major (conv input)
1226    q_l2: CudaSlice<f32>,
1227    k_l2: CudaSlice<f32>,
1228    v_g: CudaSlice<f32>, // [t, num_v, d_state]
1229    g_log: CudaSlice<f32>,
1230    beta: CudaSlice<f32>, // [t, num_v]
1231}
1232struct VerifyCkpt {
1233    gdn: Vec<Option<GdnStash>>, // [n_layer], Some iff batched linear path ran
1234    cols: Vec<Option<Vec<(CudaSlice<f32>, CudaSlice<f32>)>>>, // [n_layer][col] = (conv, ssm) after col
1235}
1236impl VerifyCkpt {
1237    fn new(n_layer: usize) -> Self {
1238        VerifyCkpt {
1239            gdn: (0..n_layer).map(|_| None).collect(),
1240            cols: (0..n_layer).map(|_| None).collect(),
1241        }
1242    }
1243}
1244
1245/// The stage-0/TX half of one PP verify. The boundary slot is the ownership token: stage 1
1246/// consumes exactly the slot selected by `tx()` / `tx_pipelined()`, never a slot inferred from
1247/// a logical round number.
1248struct VerifyBoundaryTicket {
1249    rt: &'static crate::pp::PpNRt,
1250    caller_stream: std::sync::Arc<cudarc::driver::CudaStream>,
1251    slot: usize,
1252    pos0: usize,
1253    t: usize,
1254    payload: usize,
1255    n_st: usize,
1256    pipelined: bool,
1257    pp_anatomy: bool,
1258    pp_started: std::time::Instant,
1259    reverse_ms: f64,
1260    stage0_ms: f64,
1261    tx_ms: f64,
1262    trace: Option<SpecPipeTraceCtx>,
1263}
1264
1265/// Explicit OPTIPIPE diagnostic control. Forced modes are set only by `optipipe-gate`; the
1266/// increment-2 controller can also be armed by the server's fresh-process research door.
1267#[derive(Clone, Copy, Debug, PartialEq, Eq)]
1268pub enum OptiForkGateMode {
1269    Disabled,
1270    Hit,
1271    Miss,
1272    Alternate,
1273    Abort,
1274    Controller,
1275}
1276
1277static OPTI_FORK_GATE_MODE: std::sync::atomic::AtomicU8 = std::sync::atomic::AtomicU8::new(0);
1278static OPTI_CONTROLLER_THRESHOLD: std::sync::atomic::AtomicU32 =
1279    std::sync::atomic::AtomicU32::new(0);
1280static OPTI_FORK_ATTEMPTS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
1281static OPTI_FORK_HITS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
1282static OPTI_FORK_MISSES: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
1283static OPTI_FORK_ABORT_DRAINS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
1284static OPTI_FORK_REFUSALS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
1285static OPTI_GATE_CHECKS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
1286static OPTI_GATE_ADMITS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
1287static OPTI_GATE_REJECTS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
1288static OPTI_RECONCILES: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
1289static OPTI_WASTED_DRAFT_TOKENS: std::sync::atomic::AtomicU64 =
1290    std::sync::atomic::AtomicU64::new(0);
1291static OPTI_SHADOW_DRAFT_TOKENS: std::sync::atomic::AtomicU64 =
1292    std::sync::atomic::AtomicU64::new(0);
1293static OPTI_BREAKER_TRIPS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
1294
1295impl OptiForkGateMode {
1296    fn code(self) -> u8 {
1297        match self {
1298            Self::Disabled => 0,
1299            Self::Hit => 1,
1300            Self::Miss => 2,
1301            Self::Alternate => 3,
1302            Self::Abort => 4,
1303            Self::Controller => 5,
1304        }
1305    }
1306
1307    fn configured() -> Self {
1308        match OPTI_FORK_GATE_MODE.load(std::sync::atomic::Ordering::Relaxed) {
1309            1 => Self::Hit,
1310            2 => Self::Miss,
1311            3 => Self::Alternate,
1312            4 => Self::Abort,
1313            5 => Self::Controller,
1314            _ => Self::Disabled,
1315        }
1316    }
1317
1318    fn action(self, generation: u64) -> OptiForkAction {
1319        match self {
1320            Self::Hit => OptiForkAction::Hit,
1321            Self::Miss => OptiForkAction::Miss,
1322            Self::Alternate if generation & 1 == 0 => OptiForkAction::Hit,
1323            Self::Alternate => OptiForkAction::Miss,
1324            Self::Abort => OptiForkAction::Abort,
1325            Self::Disabled | Self::Controller => {
1326                unreachable!("non-forced mode cannot choose a forced fork action")
1327            }
1328        }
1329    }
1330
1331    fn is_forced(self) -> bool {
1332        matches!(self, Self::Hit | Self::Miss | Self::Alternate | Self::Abort)
1333    }
1334}
1335
1336/// Arm or disarm the forced harness. Serving uses only `set_optipipe_controller_threshold`.
1337pub fn set_optipipe_gate_mode(mode: OptiForkGateMode) {
1338    OPTI_FORK_GATE_MODE.store(mode.code(), std::sync::atomic::Ordering::Relaxed);
1339}
1340
1341/// Arm the increment-2 diagnostic controller. The threshold applies to the uncalibrated
1342/// two-token draft-probability product. Serving can call this only through its explicit
1343/// fresh-process research door; the absent-door default remains byte-for-byte disabled.
1344pub fn set_optipipe_controller_threshold(threshold: f32) {
1345    assert!(threshold.is_finite() && (0.0..=1.0).contains(&threshold));
1346    OPTI_CONTROLLER_THRESHOLD.store(threshold.to_bits(), std::sync::atomic::Ordering::Relaxed);
1347    set_optipipe_gate_mode(OptiForkGateMode::Controller);
1348}
1349
1350#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
1351pub struct OptiForkGateStats {
1352    pub attempts: u64,
1353    pub hits: u64,
1354    pub misses: u64,
1355    pub abort_drains: u64,
1356    pub refusals: u64,
1357    pub gate_checks: u64,
1358    pub gate_admits: u64,
1359    pub gate_rejects: u64,
1360    pub reconciles: u64,
1361    pub wasted_draft_tokens: u64,
1362    pub shadow_draft_tokens: u64,
1363    pub breaker_trips: u64,
1364}
1365
1366#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
1367pub struct OptiForkStateIdentity {
1368    pub trunk_kv_bytes: usize,
1369    pub recurrent_bytes: usize,
1370    pub scratch_kv_bytes: usize,
1371    pub hidden_bytes: usize,
1372}
1373
1374pub fn reset_optipipe_gate_stats() {
1375    for counter in [
1376        &OPTI_FORK_ATTEMPTS,
1377        &OPTI_FORK_HITS,
1378        &OPTI_FORK_MISSES,
1379        &OPTI_FORK_ABORT_DRAINS,
1380        &OPTI_FORK_REFUSALS,
1381        &OPTI_GATE_CHECKS,
1382        &OPTI_GATE_ADMITS,
1383        &OPTI_GATE_REJECTS,
1384        &OPTI_RECONCILES,
1385        &OPTI_WASTED_DRAFT_TOKENS,
1386        &OPTI_SHADOW_DRAFT_TOKENS,
1387        &OPTI_BREAKER_TRIPS,
1388    ] {
1389        counter.store(0, std::sync::atomic::Ordering::Relaxed);
1390    }
1391}
1392
1393pub fn optipipe_gate_stats() -> OptiForkGateStats {
1394    let load = |v: &std::sync::atomic::AtomicU64| v.load(std::sync::atomic::Ordering::Relaxed);
1395    OptiForkGateStats {
1396        attempts: load(&OPTI_FORK_ATTEMPTS),
1397        hits: load(&OPTI_FORK_HITS),
1398        misses: load(&OPTI_FORK_MISSES),
1399        abort_drains: load(&OPTI_FORK_ABORT_DRAINS),
1400        refusals: load(&OPTI_FORK_REFUSALS),
1401        gate_checks: load(&OPTI_GATE_CHECKS),
1402        gate_admits: load(&OPTI_GATE_ADMITS),
1403        gate_rejects: load(&OPTI_GATE_REJECTS),
1404        reconciles: load(&OPTI_RECONCILES),
1405        wasted_draft_tokens: load(&OPTI_WASTED_DRAFT_TOKENS),
1406        shadow_draft_tokens: load(&OPTI_SHADOW_DRAFT_TOKENS),
1407        breaker_trips: load(&OPTI_BREAKER_TRIPS),
1408    }
1409}
1410
1411#[derive(Clone, Copy, Debug)]
1412struct OptiControllerPolicy {
1413    threshold: f32,
1414    consecutive_misses: u8,
1415    breaker_tripped: bool,
1416}
1417
1418impl OptiControllerPolicy {
1419    fn configured() -> Self {
1420        Self {
1421            threshold: f32::from_bits(
1422                OPTI_CONTROLLER_THRESHOLD.load(std::sync::atomic::Ordering::Relaxed),
1423            ),
1424            consecutive_misses: 0,
1425            breaker_tripped: false,
1426        }
1427    }
1428
1429    fn admit(&self, q_proxy: f32) -> bool {
1430        q_proxy.is_finite()
1431            && (0.0..=1.0).contains(&q_proxy)
1432            && (self.threshold == 0.0 || (!self.breaker_tripped && q_proxy >= self.threshold))
1433    }
1434
1435    /// Returns true exactly when this resolution newly trips the three-miss breaker.
1436    fn resolve(&mut self, hit: bool) -> bool {
1437        // q*=0 is the lane's explicit unconditional measurement arm. Its purpose is to price
1438        // every optimistic opportunity, so the safety breaker is measured separately and must
1439        // not silently turn this arm into "three attempts then serial".
1440        if self.threshold == 0.0 {
1441            self.consecutive_misses = 0;
1442            return false;
1443        }
1444        if hit {
1445            self.consecutive_misses = 0;
1446            return false;
1447        }
1448        self.consecutive_misses = self.consecutive_misses.saturating_add(1);
1449        if !self.breaker_tripped && self.consecutive_misses >= 3 {
1450            self.breaker_tripped = true;
1451            return true;
1452        }
1453        false
1454    }
1455}
1456
1457#[derive(Clone, Copy, Debug, PartialEq, Eq)]
1458enum OptiForkAction {
1459    Hit,
1460    Miss,
1461    Abort,
1462}
1463
1464#[derive(Clone, Copy, Debug, PartialEq, Eq)]
1465struct OptiForkGeneration {
1466    id: u64,
1467    slot: usize,
1468}
1469
1470#[derive(Default)]
1471struct OptiForkGenerationTracker {
1472    next: u64,
1473    live: [Option<u64>; 2],
1474}
1475
1476impl OptiForkGenerationTracker {
1477    fn reserve(&mut self) -> Result<OptiForkGeneration, Box<dyn std::error::Error>> {
1478        let generation = OptiForkGeneration {
1479            id: self.next,
1480            slot: (self.next & 1) as usize,
1481        };
1482        if let Some(live) = self.live[generation.slot] {
1483            return Err(format!(
1484                "optipipe snapshot slot {} still owns generation {live}; refusing to overwrite it",
1485                generation.slot,
1486            )
1487            .into());
1488        }
1489        self.next += 1;
1490        self.live[generation.slot] = Some(generation.id);
1491        Ok(generation)
1492    }
1493
1494    fn retire(&mut self, generation: OptiForkGeneration) -> Result<(), Box<dyn std::error::Error>> {
1495        match self.live[generation.slot] {
1496            Some(id) if id == generation.id => {
1497                self.live[generation.slot] = None;
1498                Ok(())
1499            }
1500            other => Err(format!(
1501                "optipipe generation teardown mismatch: ticket={} slot={} live={other:?}",
1502                generation.id, generation.slot,
1503            )
1504            .into()),
1505        }
1506    }
1507}
1508
1509struct OptiForkSeedGeneration {
1510    h_seed: CudaSlice<f32>,
1511    fill_prev: CudaSlice<f32>,
1512    scratch_len: usize,
1513}
1514
1515/// Allocate or refresh one full checkpoint through the engine that owns each PP stage. The
1516/// generic cache helper accepts one device and therefore cannot copy GDN state split across
1517/// devices. KV lengths and position stay host metadata; only recurrent buffers need stage-local
1518/// device ownership.
1519fn opti_snapshot_stage_owned(
1520    e: &Engine,
1521    cache: &Cache,
1522    rt: &'static crate::pp::PpNRt,
1523    fence: &[usize],
1524) -> Result<crate::cache::CacheSnapshot, Box<dyn std::error::Error>> {
1525    let n = cache.kv.len();
1526    let mut snapshot = crate::cache::CacheSnapshot {
1527        kv_len: vec![None; n],
1528        conv: (0..n).map(|_| None).collect(),
1529        ssm: (0..n).map(|_| None).collect(),
1530        pos: cache.pos,
1531    };
1532    opti_snapshot_stage_owned_into(e, cache, rt, fence, &mut snapshot)?;
1533    Ok(snapshot)
1534}
1535
1536fn opti_snapshot_stage_owned_into(
1537    e: &Engine,
1538    cache: &Cache,
1539    rt: &'static crate::pp::PpNRt,
1540    fence: &[usize],
1541    snapshot: &mut crate::cache::CacheSnapshot,
1542) -> Result<(), Box<dyn std::error::Error>> {
1543    if fence.len() != rt.n_stages() + 1 || snapshot.kv_len.len() != cache.kv.len() {
1544        return Err("optipipe stage-owned snapshot shape mismatch".into());
1545    }
1546    for stage in 0..rt.n_stages() {
1547        opti_snapshot_one_stage_owned_into(e, cache, rt, fence, stage, snapshot)?;
1548    }
1549    snapshot.pos = cache.pos;
1550    Ok(())
1551}
1552
1553/// Refresh one PP stage of a checkpoint. Increment 2 uses this split form so stage 0's
1554/// optimistic post-N state is captured before N+1 stage 0 is queued, while stage 1's matching
1555/// post-N state is captured only after N stage 1 is enqueued. Calling the all-stage helper at
1556/// either point would capture one side of the fork at the wrong generation.
1557fn opti_snapshot_one_stage_owned_into(
1558    e: &Engine,
1559    cache: &Cache,
1560    rt: &'static crate::pp::PpNRt,
1561    fence: &[usize],
1562    stage: usize,
1563    snapshot: &mut crate::cache::CacheSnapshot,
1564) -> Result<(), Box<dyn std::error::Error>> {
1565    if fence.len() != rt.n_stages() + 1
1566        || snapshot.kv_len.len() != cache.kv.len()
1567        || stage >= rt.n_stages()
1568    {
1569        return Err("optipipe single-stage snapshot shape mismatch".into());
1570    }
1571    let _scope = rt.enter(stage);
1572    let owner = rt.engine(stage, e);
1573    for il in fence[stage]..fence[stage + 1] {
1574        snapshot.kv_len[il] = cache.kv[il].as_ref().map(|kv| kv.len);
1575        match &cache.recur[il] {
1576            Some(recur) => {
1577                match snapshot.conv[il].as_mut() {
1578                    Some(dst) => {
1579                        owner.copy_into(dst, 0, &recur.conv_state, recur.conv_state.len())?
1580                    }
1581                    None => snapshot.conv[il] = Some(owner.clone_dtod(&recur.conv_state)?),
1582                }
1583                match snapshot.ssm[il].as_mut() {
1584                    Some(dst) => {
1585                        owner.copy_into(dst, 0, &recur.ssm_state, recur.ssm_state.len())?
1586                    }
1587                    None => snapshot.ssm[il] = Some(owner.clone_dtod(&recur.ssm_state)?),
1588                }
1589            }
1590            None if snapshot.conv[il].is_some() || snapshot.ssm[il].is_some() => {
1591                return Err(
1592                    format!("optipipe stage-owned snapshot layer {il} changed shape").into(),
1593                );
1594            }
1595            None => {}
1596        }
1597    }
1598    snapshot.pos = cache.pos;
1599    Ok(())
1600}
1601
1602/// Increment-1 persistent fork state. Exactly two snapshot/seed slots alternate; a live ticket
1603/// names its generation and keeps teardown fail-closed. Only stage 0 is allowed to mutate before
1604/// resolve, so the reconcile tables and conditional restores are stage-local.
1605struct OptiForkState {
1606    mode: OptiForkGateMode,
1607    controller: Option<OptiControllerPolicy>,
1608    generations: OptiForkGenerationTracker,
1609    active_snapshot_slot: usize,
1610    alternate_snapshot: crate::cache::CacheSnapshot,
1611    seeds: [OptiForkSeedGeneration; 2],
1612    rt: &'static crate::pp::PpNRt,
1613    fence: [usize; 3],
1614    split: usize,
1615    len_ptrs: CudaSlice<u64>,
1616    saved_lens: CudaSlice<i32>,
1617    forced_acc: CudaSlice<u32>,
1618    valid: CudaSlice<u32>,
1619    stage0_stream: std::sync::Arc<cudarc::driver::CudaStream>,
1620    logical_payload_bytes: [usize; 2],
1621}
1622
1623struct OptiForkTicket {
1624    generation: OptiForkGeneration,
1625    boundary: Option<VerifyBoundaryTicket>,
1626    drain: std::sync::Arc<cudarc::driver::CudaStream>,
1627    settled: bool,
1628}
1629
1630struct OptiControllerTicket {
1631    generation: OptiForkGeneration,
1632    boundary: Option<VerifyBoundaryTicket>,
1633    ckpt: Option<VerifyCkpt>,
1634    verify_tokens: [u32; 2],
1635    draft_prob: f32,
1636    eager_seed: Option<CudaSlice<f32>>,
1637    q_proxy: f32,
1638    scratch_len: usize,
1639    issued_at: std::time::Instant,
1640    drain: std::sync::Arc<cudarc::driver::CudaStream>,
1641    settled: bool,
1642}
1643
1644struct OptiControllerPrepared {
1645    verify_tokens: [u32; 2],
1646    draft_prob: f32,
1647    eager_seed: Option<CudaSlice<f32>>,
1648    q_proxy: f32,
1649    scratch_len: usize,
1650}
1651
1652impl OptiControllerTicket {
1653    fn take_boundary(&mut self) -> VerifyBoundaryTicket {
1654        self.boundary
1655            .take()
1656            .expect("controller boundary ticket already consumed")
1657    }
1658
1659    fn take_ckpt(&mut self) -> VerifyCkpt {
1660        self.ckpt
1661            .take()
1662            .expect("controller verify checkpoint already consumed")
1663    }
1664
1665    fn take_eager_seed(&mut self) -> Option<CudaSlice<f32>> {
1666        self.eager_seed.take()
1667    }
1668
1669    fn settle(&mut self) {
1670        self.settled = true;
1671    }
1672}
1673
1674impl Drop for OptiControllerTicket {
1675    fn drop(&mut self) {
1676        if !self.settled {
1677            let _ = self.drain.synchronize();
1678            OPTI_FORK_ABORT_DRAINS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
1679        }
1680    }
1681}
1682
1683impl OptiForkTicket {
1684    fn take_boundary(&mut self) -> VerifyBoundaryTicket {
1685        self.boundary
1686            .take()
1687            .expect("fork ticket boundary already consumed")
1688    }
1689
1690    fn settle(&mut self) {
1691        self.settled = true;
1692    }
1693}
1694
1695impl Drop for OptiForkTicket {
1696    fn drop(&mut self) {
1697        if !self.settled {
1698            let _ = self.drain.synchronize();
1699            OPTI_FORK_ABORT_DRAINS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
1700        }
1701    }
1702}
1703
1704impl OptiForkState {
1705    #[allow(clippy::too_many_arguments)]
1706    fn new(
1707        e: &Engine,
1708        cache: &Cache,
1709        mode: OptiForkGateMode,
1710        alternate_snapshot: crate::cache::CacheSnapshot,
1711        h_seed: &CudaSlice<f32>,
1712        fill_prev: &CudaSlice<f32>,
1713        rt: &'static crate::pp::PpNRt,
1714        split: usize,
1715        n_layer: usize,
1716    ) -> Result<Self, Box<dyn std::error::Error>> {
1717        let fence = [0, split, n_layer];
1718        let mut logical_payload_bytes = [0usize; 2];
1719        for stage in 0..2 {
1720            for il in fence[stage]..fence[stage + 1] {
1721                logical_payload_bytes[stage] += alternate_snapshot.conv[il]
1722                    .as_ref()
1723                    .map_or(0, |v| v.len() * std::mem::size_of::<f32>());
1724                logical_payload_bytes[stage] += alternate_snapshot.ssm[il]
1725                    .as_ref()
1726                    .map_or(0, |v| v.len() * std::mem::size_of::<f32>());
1727            }
1728        }
1729        let seeds = [
1730            OptiForkSeedGeneration {
1731                h_seed: e.clone_dtod(h_seed)?,
1732                fill_prev: e.clone_dtod(fill_prev)?,
1733                scratch_len: 0,
1734            },
1735            OptiForkSeedGeneration {
1736                h_seed: e.clone_dtod(h_seed)?,
1737                fill_prev: e.clone_dtod(fill_prev)?,
1738                scratch_len: 0,
1739            },
1740        ];
1741        let (len_ptrs, saved_lens, forced_acc, valid, stage0_stream) = {
1742            let _stage = rt.enter(0);
1743            let e0 = rt.engine(0, e);
1744            (
1745                crate::round_stream::kv_len_ptr_table_range(e0, cache, 0..split, None)?,
1746                e0.htod_i32(&vec![0; split])?,
1747                e0.alloc_u32_zeroed(2)?,
1748                e0.alloc_u32_zeroed(1)?,
1749                e0.stream(),
1750            )
1751        };
1752        logical_payload_bytes[0] += seeds
1753            .iter()
1754            .map(|seed| (seed.h_seed.len() + seed.fill_prev.len()) * std::mem::size_of::<f32>())
1755            .sum::<usize>();
1756        logical_payload_bytes[0] += len_ptrs.len() * std::mem::size_of::<u64>()
1757            + saved_lens.len() * std::mem::size_of::<i32>()
1758            + forced_acc.len() * std::mem::size_of::<u32>()
1759            + valid.len() * std::mem::size_of::<u32>();
1760        Ok(Self {
1761            mode,
1762            controller: (mode == OptiForkGateMode::Controller)
1763                .then(OptiControllerPolicy::configured),
1764            generations: OptiForkGenerationTracker::default(),
1765            active_snapshot_slot: 0,
1766            alternate_snapshot,
1767            seeds,
1768            rt,
1769            fence,
1770            split,
1771            len_ptrs,
1772            saved_lens,
1773            forced_acc,
1774            valid,
1775            stage0_stream,
1776            logical_payload_bytes,
1777        })
1778    }
1779
1780    fn reserve(
1781        &mut self,
1782        current_snapshot: &mut crate::cache::CacheSnapshot,
1783    ) -> Result<OptiForkGeneration, Box<dyn std::error::Error>> {
1784        let generation = self.generations.reserve()?;
1785        if generation.slot != self.active_snapshot_slot {
1786            std::mem::swap(current_snapshot, &mut self.alternate_snapshot);
1787            self.active_snapshot_slot = generation.slot;
1788        }
1789        Ok(generation)
1790    }
1791
1792    fn capture_seed(
1793        &mut self,
1794        e: &Engine,
1795        generation: OptiForkGeneration,
1796        h_seed: &CudaSlice<f32>,
1797        fill_prev: &CudaSlice<f32>,
1798        scratch_len: usize,
1799    ) -> Result<(), Box<dyn std::error::Error>> {
1800        let seed = &mut self.seeds[generation.slot];
1801        e.copy_into(&mut seed.h_seed, 0, h_seed, h_seed.len())?;
1802        e.copy_into(&mut seed.fill_prev, 0, fill_prev, fill_prev.len())?;
1803        seed.scratch_len = scratch_len;
1804        Ok(())
1805    }
1806
1807    fn ticket(
1808        &self,
1809        generation: OptiForkGeneration,
1810        boundary: VerifyBoundaryTicket,
1811    ) -> OptiForkTicket {
1812        OptiForkTicket {
1813            generation,
1814            boundary: Some(boundary),
1815            drain: self.stage0_stream.clone(),
1816            settled: false,
1817        }
1818    }
1819
1820    #[allow(clippy::too_many_arguments)]
1821    fn controller_ticket(
1822        &self,
1823        generation: OptiForkGeneration,
1824        boundary: VerifyBoundaryTicket,
1825        ckpt: VerifyCkpt,
1826        verify_tokens: [u32; 2],
1827        draft_prob: f32,
1828        eager_seed: Option<CudaSlice<f32>>,
1829        q_proxy: f32,
1830        scratch_len: usize,
1831    ) -> OptiControllerTicket {
1832        OptiControllerTicket {
1833            generation,
1834            boundary: Some(boundary),
1835            ckpt: Some(ckpt),
1836            verify_tokens,
1837            draft_prob,
1838            eager_seed,
1839            q_proxy,
1840            scratch_len,
1841            issued_at: std::time::Instant::now(),
1842            drain: self.stage0_stream.clone(),
1843            settled: false,
1844        }
1845    }
1846
1847    fn reserve_successor(&mut self) -> Result<OptiForkGeneration, Box<dyn std::error::Error>> {
1848        self.generations.reserve()
1849    }
1850
1851    fn successor_snapshot_mut(&mut self) -> &mut crate::cache::CacheSnapshot {
1852        &mut self.alternate_snapshot
1853    }
1854
1855    fn promote_successor_snapshot(
1856        &mut self,
1857        current_snapshot: &mut crate::cache::CacheSnapshot,
1858        generation: OptiForkGeneration,
1859    ) {
1860        std::mem::swap(current_snapshot, &mut self.alternate_snapshot);
1861        self.active_snapshot_slot = generation.slot;
1862    }
1863
1864    fn queue_actual_reconcile(
1865        &mut self,
1866        e: &Engine,
1867        snapshot: &crate::cache::CacheSnapshot,
1868        acc: &CudaSlice<u32>,
1869        optimistic_pending: u32,
1870        base: usize,
1871    ) -> Result<(), Box<dyn std::error::Error>> {
1872        let saved: Vec<i32> = (0..self.split)
1873            .map(|il| snapshot.kv_len[il].map(|v| v as i32).unwrap_or(0))
1874            .collect();
1875        // Serving keeps the caller/accept walk on the head (stage-1) device. Record the accept
1876        // decision point there and append a wait to stage 0 after its optimistic successor/TX;
1877        // the validity/reconcile kernels must never peer-read acc before it is written. The
1878        // increment-1 harness uses primary stage 0, where stream order already provides this.
1879        if self.rt.engine(0, e).ctx().ordinal() != e.ctx().ordinal() {
1880            self.rt.fence_stages_behind(&e.stream())?;
1881        }
1882        let _stage = self.rt.enter(0);
1883        let e0 = self.rt.engine(0, e);
1884        e0.htod_i32_into(&mut self.saved_lens, &saved)?;
1885        e0.spec_fork_valid(acc, optimistic_pending, &mut self.valid)?;
1886        e0.spec_fork_reconcile_kv(
1887            &self.len_ptrs,
1888            &self.saved_lens,
1889            acc,
1890            &self.valid,
1891            base,
1892            self.split,
1893        )
1894    }
1895
1896    fn finish_actual_reconcile(
1897        &mut self,
1898        e: &Engine,
1899        cache: &mut Cache,
1900        snapshot: &crate::cache::CacheSnapshot,
1901        n_acc: usize,
1902        base: usize,
1903        hit: bool,
1904    ) -> Result<(), Box<dyn std::error::Error>> {
1905        if hit {
1906            return Ok(());
1907        }
1908        let len_delta = base + n_acc;
1909        for il in 0..self.split {
1910            if let (Some(kv), Some(saved)) = (cache.kv[il].as_mut(), snapshot.kv_len[il]) {
1911                kv.len = saved + len_delta;
1912            }
1913        }
1914        {
1915            let _stage = self.rt.enter(1);
1916            let e1 = self.rt.engine(1, e);
1917            for il in self.split..self.fence[2] {
1918                if let (Some(kv), Some(saved)) = (cache.kv[il].as_mut(), snapshot.kv_len[il]) {
1919                    kv.len = saved + len_delta;
1920                    e1.set_i32_one(&mut kv.len_d, kv.len as i32)?;
1921                }
1922            }
1923        }
1924        self.rt.publish_to(0, &e.stream())?;
1925        Ok(())
1926    }
1927
1928    fn cancel_controller_ticket(
1929        &mut self,
1930        e: &Engine,
1931        cache: &mut Cache,
1932        scratch: &mut MtpScratch,
1933        snapshot: &crate::cache::CacheSnapshot,
1934        ticket: &mut OptiControllerTicket,
1935    ) -> Result<(), Box<dyn std::error::Error>> {
1936        {
1937            let _stage = self.rt.enter(0);
1938            let e0 = self.rt.engine(0, e);
1939            for il in 0..self.split {
1940                if let (Some(kv), Some(saved)) = (cache.kv[il].as_mut(), snapshot.kv_len[il]) {
1941                    kv.len = saved;
1942                    e0.set_i32_one(&mut kv.len_d, saved as i32)?;
1943                }
1944            }
1945        }
1946        scratch.set_len(e, snapshot.pos)?;
1947        ticket.settle();
1948        self.generations.retire(ticket.generation)?;
1949        OPTI_FORK_ABORT_DRAINS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
1950        OPTI_WASTED_DRAFT_TOKENS.fetch_add(2, std::sync::atomic::Ordering::Relaxed);
1951        eprintln!(
1952            "[opti-controller] tail-drain generation={} slot={}",
1953            ticket.generation.id, ticket.generation.slot,
1954        );
1955        Ok(())
1956    }
1957
1958    #[allow(clippy::too_many_arguments)]
1959    fn reconcile(
1960        &mut self,
1961        e: &Engine,
1962        cache: &mut Cache,
1963        scratch: &mut MtpScratch,
1964        snapshot: &crate::cache::CacheSnapshot,
1965        h_seed: &mut CudaSlice<f32>,
1966        fill_prev: &mut CudaSlice<f32>,
1967        generation: OptiForkGeneration,
1968        action: OptiForkAction,
1969        optimistic_pending: u32,
1970    ) -> Result<(), Box<dyn std::error::Error>> {
1971        debug_assert!(action != OptiForkAction::Abort);
1972        let miss_started = std::time::Instant::now();
1973        let keep = action == OptiForkAction::Hit;
1974        let saved: Vec<i32> = (0..self.split)
1975            .map(|il| snapshot.kv_len[il].map(|v| v as i32).unwrap_or(0))
1976            .collect();
1977        let seed = &self.seeds[generation.slot];
1978        {
1979            let _stage = self.rt.enter(0);
1980            let e0 = self.rt.engine(0, e);
1981            e0.htod_i32_into(&mut self.saved_lens, &saved)?;
1982            let forced = if keep {
1983                [1u32, optimistic_pending]
1984            } else {
1985                [0u32, optimistic_pending]
1986            };
1987            e0.htod_u32_into(&mut self.forced_acc, &forced)?;
1988            e0.spec_fork_valid(&self.forced_acc, optimistic_pending, &mut self.valid)?;
1989            e0.spec_fork_reconcile_kv(
1990                &self.len_ptrs,
1991                &self.saved_lens,
1992                &self.forced_acc,
1993                &self.valid,
1994                0,
1995                self.split,
1996            )?;
1997            for il in 0..self.split {
1998                if let Some(recur) = cache.recur[il].as_mut() {
1999                    let conv = snapshot.conv[il]
2000                        .as_ref()
2001                        .ok_or("optipipe stage0 snapshot missing conv state")?;
2002                    let ssm = snapshot.ssm[il]
2003                        .as_ref()
2004                        .ok_or("optipipe stage0 snapshot missing ssm state")?;
2005                    e0.spec_fork_restore_f32(conv, &mut recur.conv_state, &self.valid)?;
2006                    e0.spec_fork_restore_f32(ssm, &mut recur.ssm_state, &self.valid)?;
2007                }
2008            }
2009            e0.spec_fork_restore_f32(&seed.h_seed, h_seed, &self.valid)?;
2010            e0.spec_fork_restore_f32(&seed.fill_prev, fill_prev, &self.valid)?;
2011        }
2012
2013        if keep {
2014            OPTI_FORK_HITS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2015            return Ok(());
2016        }
2017
2018        for il in 0..self.split {
2019            if let (Some(kv), Some(saved)) = (cache.kv[il].as_mut(), snapshot.kv_len[il]) {
2020                kv.len = saved;
2021            }
2022        }
2023        scratch.set_len(e, seed.scratch_len)?;
2024        // Targeted E_restart: publish only stage 0's reconcile to the caller, then bound the
2025        // forced diagnostic so the retained number is the actual miss cost, not enqueue time.
2026        let caller = e.stream();
2027        self.rt.publish_to(0, &caller)?;
2028        caller.synchronize()?;
2029        let miss_ms = miss_started.elapsed().as_secs_f64() * 1e3;
2030        eprintln!(
2031            "[opti-fork-reconcile] generation={} slot={} miss_ms={miss_ms:.3}",
2032            generation.id, generation.slot,
2033        );
2034        OPTI_FORK_MISSES.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2035        Ok(())
2036    }
2037
2038    fn retire(&mut self, generation: OptiForkGeneration) -> Result<(), Box<dyn std::error::Error>> {
2039        self.generations.retire(generation)
2040    }
2041}
2042
2043impl HybridModel {
2044    fn opti_graph_draft_step(
2045        &self,
2046        e: &Engine,
2047        mtp: &MtpHead,
2048        dctx: &mut DraftGraphCtx,
2049        scratch: &mut MtpScratch,
2050        d_vocab: usize,
2051    ) -> Result<(u32, f32), Box<dyn std::error::Error>> {
2052        dctx.graph
2053            .as_ref()
2054            .ok_or("optipipe controller requires the greedy draft graph")?
2055            .launch()?;
2056        scratch.kv.len += 1;
2057        let idx = e.dtoh_u32_one(&dctx.g_tok)?;
2058        if (idx as usize) >= d_vocab {
2059            return Err(
2060                format!("optipipe draft argmax sentinel 0x{idx:08x} >= d_vocab {d_vocab}").into(),
2061            );
2062        }
2063        let probability = e.dtoh(&dctx.g_p)?[0];
2064        if !probability.is_finite() || !(0.0..=1.0).contains(&probability) {
2065            return Err(format!("optipipe draft probability is invalid: {probability}").into());
2066        }
2067        let token = match &mtp.d2t {
2068            Some(map) => map[idx as usize],
2069            None => idx,
2070        };
2071        if token != idx {
2072            e.set_u32_one(&mut dctx.g_tok, token)?;
2073        }
2074        Ok((token, probability))
2075    }
2076
2077    #[allow(clippy::too_many_arguments)]
2078    fn opti_controller_draft_step(
2079        &self,
2080        e: &Engine,
2081        mtp: &MtpHead,
2082        dctx: &mut DraftGraphCtx,
2083        scratch: &mut MtpScratch,
2084        d_vocab: usize,
2085        eager_state: &mut Option<(u32, CudaSlice<f32>)>,
2086        eager_pos: usize,
2087        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
2088    ) -> Result<(u32, f32), Box<dyn std::error::Error>> {
2089        if dctx.graph.is_some() {
2090            return self.opti_graph_draft_step(e, mtp, dctx, scratch, d_vocab);
2091        }
2092        let (input_token, input_seed) = eager_state
2093            .take()
2094            .ok_or("optipipe eager continuation seed is unavailable")?;
2095        let (logits, next_seed) = self.mtp_head_forward_dev(
2096            e,
2097            mtp,
2098            input_token,
2099            &input_seed,
2100            scratch,
2101            eager_pos,
2102            embd_dev,
2103            None,
2104        )?;
2105        let token_d = e.argmax_token_device(&logits, d_vocab)?;
2106        let idx = e.dtoh_u32_one(&token_d)?;
2107        if (idx as usize) >= d_vocab {
2108            return Err(format!(
2109                "optipipe eager draft argmax sentinel 0x{idx:08x} >= d_vocab {d_vocab}"
2110            )
2111            .into());
2112        }
2113        let probability_d = e.prob_of_token_device(&logits, &token_d, d_vocab)?;
2114        let probability = e.dtoh(&probability_d)?[0];
2115        if !probability.is_finite() || !(0.0..=1.0).contains(&probability) {
2116            return Err(
2117                format!("optipipe eager draft probability is invalid: {probability}").into(),
2118            );
2119        }
2120        let token = match &mtp.d2t {
2121            Some(map) => map[idx as usize],
2122            None => idx,
2123        };
2124        *eager_state = Some((token, next_seed));
2125        Ok((token, probability))
2126    }
2127
2128    /// NextN head forward for ONE draft token (§A ops 1-13, T=1).
2129    /// Inputs: `e_tok` = the token to predict FROM (last committed / previous draft); `h_seed` =
2130    /// the trunk's pre-output_norm hidden of that token (§A op 2 input). `mtp_pos` = absolute
2131    /// position of the token being predicted from. Returns (draft_logits[n_vocab] host, h_nextn dev).
2132    /// `h_nextn` (§A op 10) becomes `h_seed` for the next autoregressive draft step.
2133    /// Device-resident: returns draft logits ON DEVICE (no [n_vocab] dtoh). The greedy draft
2134    /// loop only needs argmax — paired with `argmax_token_device` this cuts the ~600KB logits
2135    /// transfer + host argmax per draft token from the K-token draft chain.
2136    #[allow(clippy::too_many_arguments)]
2137    fn mtp_head_forward_dev(
2138        &self,
2139        e: &Engine,
2140        mtp: &MtpHead,
2141        e_tok: u32,
2142        h_seed: &CudaSlice<f32>,
2143        scratch: &mut MtpScratch,
2144        mtp_pos: usize,
2145        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
2146        // DRAFT-SIDE GRAMMAR MASK (lane/draft-mask): (packed draft-vocab allowed set, words).
2147        // Applied to the head logits BEFORE they are returned, so every consumer (argmax,
2148        // gumbel draw, p-min prob) sees the grammar-legal row. None = unmasked (pre-lane).
2149        mask: Option<(&CudaSlice<u32>, usize)>,
2150    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
2151        let cfg = &self.cfg;
2152        let n_embd = cfg.n_embd as usize;
2153        // Distilled-student geometry: the block runs at the INNER width `di` (eh_proj out /
2154        // attn / ffn); the n_embd interface (embed, norms in, carrier out, head in) is unchanged.
2155        let di = mtp.geom.as_ref().map(|g| g.d_inner).unwrap_or(n_embd);
2156        let eps = cfg.rms_eps;
2157        let pos_d = e.htod_i32(&[mtp_pos as i32])?;
2158
2159        // op A: a resident table transfers one 4B token id. The exact host-row capacity path
2160        // expands this one row on CPU and transfers n_embd f32 values instead.
2161        let e_emb = match embd_dev {
2162            Some((g, qt, rb)) => e.embed_gather_device_t(g, &[e_tok], n_embd, qt, rb)?,
2163            None => e.htod(&self.embd.gather(n_embd, &[e_tok]))?,
2164        };
2165
2166        // op 1/2: e_norm = RMSNorm(e, enorm); h_norm = RMSNorm(h_seed, hnorm)
2167        let mut e_norm = e.zeros(n_embd)?;
2168        e.rms_norm(&e_emb, mtp.enorm.float_data(), &mut e_norm, n_embd, 1, eps)?;
2169        let mut h_norm = e.zeros(n_embd)?;
2170        e.rms_norm(h_seed, mtp.hnorm.float_data(), &mut h_norm, n_embd, 1, eps)?;
2171
2172        // op 3: concat = [e_norm ; h_norm] -> [2*n_embd], e_norm in [0,n_embd), h_norm in [n_embd,2n_embd)
2173        let mut concat = e.zeros(2 * n_embd)?;
2174        e.copy_into(&mut concat, 0, &e_norm, n_embd)?;
2175        e.copy_into(&mut concat, n_embd, &h_norm, n_embd)?;
2176
2177        // op 4: inpSA = eh_proj @ concat  (eh_proj [2*n_embd, n_embd]) -> [n_embd]
2178        let inp_sa = e.matmul(&mtp.eh_proj, &concat, 1)?;
2179
2180        // op 5: a_norm = RMSNorm(inpSA, attn_norm)
2181        let mut a_norm = e.zeros(di)?;
2182        e.rms_norm(&inp_sa, mtp.attn_norm.float_data(), &mut a_norm, di, 1, eps)?;
2183
2184        // op 6: attention on the scratch KV. SAME dc launcher as the graph path (bucket_max =
2185        // scratch.cap, length from the device len_d) so eager drafts match graph drafts
2186        // bit-for-bit at any t_kv (the parity gate). Host len mirrored here (the dc append
2187        // advances only the device counter).
2188        let attn_out = match (&mtp.mixer, mtp.step35.as_ref()) {
2189            // step35 MTP block: PER-LAYER geometry + a separate head-wise gate + an SWA window,
2190            // none of which the dc launcher can express (see `mtp_step35_attn`). Host-len arm.
2191            // Advances BOTH the host len and the device counter itself (unlike the dc arm,
2192            // whose host-side mirror the caller does).
2193            (Mixer::Full(fa), Some(g)) => {
2194                self.mtp_step35_attn(e, fa, g, &a_norm, &pos_d, scratch)?
2195            }
2196            (Mixer::Full(fa), None) => {
2197                let out =
2198                    self.mtp_full_attn_dc(e, fa, &a_norm, &pos_d, scratch, mtp.geom.as_ref())?;
2199                scratch.kv.len += 1;
2200                out
2201            }
2202            (Mixer::Linear(_), _) => {
2203                panic!("MTP block is full-attn in qwen35; linear MTP not supported")
2204            }
2205            (Mixer::Mla(_), _) => crate::hybrid::mla_forward_unimplemented(),
2206        };
2207
2208        // op 7: x1 = inpSA + attn_out
2209        let mut x1 = e.zeros(di)?;
2210        e.add(&inp_sa, &attn_out, &mut x1, di)?;
2211
2212        // op 8: z = RMSNorm(x1, post_attn_norm)  (pre-FFN norm)
2213        let mut z = e.zeros(di)?;
2214        e.rms_norm(&x1, mtp.post_attn_norm.float_data(), &mut z, di, 1, eps)?;
2215
2216        // op 9: FFN (Dense or MoE) — same as the trunk decode FFN
2217        let ffn_out = match &mtp.ffn {
2218            crate::hybrid::Ffn::Dense {
2219                ffn_gate,
2220                ffn_up,
2221                ffn_down,
2222            } => {
2223                let n_ff = ffn_gate.out_features();
2224                let (gate, up) = if e.uses_q8_1_fast(ffn_gate) && e.uses_q8_1_fast(ffn_up) {
2225                    let (zq, zd) = e.quantize_q8_1(&z, 1, di)?;
2226                    (
2227                        e.matmul_pre(ffn_gate, &zq, &zd, &z, 1)?,
2228                        e.matmul_pre(ffn_up, &zq, &zd, &z, 1)?,
2229                    )
2230                } else {
2231                    (e.matmul(ffn_gate, &z, 1)?, e.matmul(ffn_up, &z, 1)?)
2232                };
2233                let mut act = e.zeros(n_ff)?;
2234                // step35: a DENSE FFN reads the per-layer SHEXP clamp (upstream's one `build_ffn`
2235                // serves the dense MLP and the shared expert off `swiglu_clamp_shexp` —
2236                // llama-graph.cpp:1751), resolved for the MTP block's OWN index. Every other arch
2237                // passes None, which is `ffn_act`'s dispatch verbatim.
2238                Self::ffn_act_lim(
2239                    e,
2240                    &self.cfg,
2241                    &gate,
2242                    &up,
2243                    1.0,
2244                    1.0,
2245                    mtp.step35.as_ref().and_then(|s| s.clamp_shexp),
2246                    &mut act,
2247                    n_ff,
2248                )?;
2249                e.matmul(ffn_down, &act, 1)?
2250            }
2251            // MTP head is a distinct block — key its experts under a separate layer index (u16::MAX)
2252            // so they never alias trunk layer 0's cache keys.
2253            crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il(e, m, &z, 1, u16::MAX)?,
2254        };
2255
2256        // op 10: h_nextn = x1 + ffn_out (at di)
2257        let mut h_inner = e.zeros(di)?;
2258        e.add(&x1, &ffn_out, &mut h_inner, di)?;
2259
2260        // op 10.5 (student): up-project the inner hidden back to n_embd — training semantics:
2261        // the chain carrier AND the head input are out_up(h_inner) (pre-final-norm).
2262        let h_nextn = match mtp.geom.as_ref() {
2263            Some(g) => e.matmul(&g.out_up, &h_inner, 1)?,
2264            None => h_inner,
2265        };
2266
2267        // op 11: final = RMSNorm(h_nextn, shared_head_norm OR output_norm)
2268        let final_norm = mtp.shared_head_norm.as_ref().unwrap_or(&self.output_norm);
2269        let mut final_h = e.zeros(n_embd)?;
2270        e.rms_norm(
2271            &h_nextn,
2272            final_norm.float_data(),
2273            &mut final_h,
2274            n_embd,
2275            1,
2276            eps,
2277        )?;
2278
2279        // op 12: draft_logits = (shared_head_head OR output) @ final — stays ON DEVICE.
2280        let head = mtp.shared_head_head.as_ref().unwrap_or(&self.output);
2281        let mut logits = e.matmul(head, &final_h, 1)?;
2282        // op 12b (lane/draft-mask): grammar mask over the DRAFT vocab, applied here so the
2283        // caller's argmax / gumbel draw / p-min prob all read the grammar-legal row.
2284        if let Some((mask_d, mw)) = mask {
2285            let d_vocab = head.out_features();
2286            e.mask_logits_col(&mut logits, mask_d, 0, d_vocab, mw)?;
2287        }
2288        // Chain recurrence hand-over: pre-norm h_nextn (default) or post-norm final_h
2289        // (MEMRA_SPEC_HPOST — llama.cpp #24025's t_h_nextn is taken AFTER the head norm).
2290        Ok((logits, if spec_hpost() { final_h } else { h_nextn }))
2291    }
2292
2293    /// step35 MTP-block attention, T=1, on the scratch KV — the EAGER-ONLY twin of
2294    /// `mtp_full_attn_dc`. Three things force a separate arm rather than a geometry parameter on
2295    /// the dc path, and all three are properties of this arch's MTP block:
2296    ///
2297    /// 1. **The SWA window.** Block 45 is an SWA-type block (`sliding_window_pattern[45]=true`,
2298    ///    window 512). Windowed decode in memra is a token-aligned VIEW OFFSET into the quantized
2299    ///    cache (the gemma4 R6 / `step35_decode_attn` pattern: keys carry absolute rope and the
2300    ///    mask is purely positional, so one query at `len-1` attending the last `win` rows IS the
2301    ///    windowed result). `fa_decode_dc` takes the key count from a DEVICE counter and always
2302    ///    starts at row 0 — it cannot express a nonzero offset, so a windowed dc arm would need a
2303    ///    new kernel. That is deliberately not built here: see the CUDA-graph note below.
2304    /// 2. **Per-layer head count.** 96 q heads over 8 KV (GQA 12) at this block, vs the trunk's 64
2305    ///    on its full-attn layers. The trunk cfg's `n_head` scalar is the MAX over layers, and the
2306    ///    trunk ARTIFACT's per-layer arrays stop at index 44 — so the count must come from the
2307    ///    resolved `Step35MtpGeom`, never from `cfg`.
2308    /// 3. **The separate head-wise gate.** `blk.45.attn_gate.weight [n_embd, 96]` produces one
2309    ///    sigmoid scalar per head (broadcast over head_dim) — `attn_head_gate`, not the qwen35
2310    ///    fused-into-wq `q_gate_split` form the dc arm handles.
2311    ///
2312    /// WHY EAGER-ONLY IS NOT A GAP TODAY: `graph_draft` requires `trunk_dense`, and this SKU's
2313    /// trunk is a 288-expert MoE, so the graph draft is already off for every step35 model — the
2314    /// eager chain IS the served path. `mtp_head_forward_cap` refuses step35 explicitly rather
2315    /// than silently capturing a window-less (wrong past `win` draft rows) graph.
2316    ///
2317    /// Unlike the dc arm this advances BOTH the host `kv.len` and the device counter, so the
2318    /// caller must not mirror.
2319    fn mtp_step35_attn(
2320        &self,
2321        e: &Engine,
2322        fa: &FullAttnLayer,
2323        g: &crate::hybrid::Step35MtpGeom,
2324        h: &CudaSlice<f32>,
2325        pos_d: &CudaSlice<i32>,
2326        scratch: &mut MtpScratch,
2327    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2328        let (nh, nkv, hd) = (g.n_head, g.n_head_kv, self.cfg.head_dim_k as usize);
2329        let eps = self.cfg.rms_eps;
2330        let scale = 1.0 / (hd as f32).sqrt(); // step35.cpp:255 kq_scale
2331        let n_embd = self.cfg.n_embd as usize;
2332        let gw = fa
2333            .attn_gate
2334            .as_ref()
2335            .ok_or("step35 MTP block is missing attn_gate.weight (head-wise attention gate)")?;
2336
2337        let (q0, k0, v0, gt) = if e.uses_q8_1_fast(&fa.wq)
2338            && e.uses_q8_1_fast(&fa.wk)
2339            && e.uses_q8_1_fast(&fa.wv)
2340            && e.uses_q8_1_fast(gw)
2341        {
2342            let (hq, hdq) = e.quantize_q8_1(h, 1, n_embd)?;
2343            let (a, b, c) = match e.matmul_q8_fused3(&fa.wq, &fa.wk, &fa.wv, &hq, &hdq)? {
2344                Some(t3) => t3,
2345                None => (
2346                    e.matmul_pre(&fa.wq, &hq, &hdq, h, 1)?,
2347                    e.matmul_pre(&fa.wk, &hq, &hdq, h, 1)?,
2348                    e.matmul_pre(&fa.wv, &hq, &hdq, h, 1)?,
2349                ),
2350            };
2351            (a, b, c, e.matmul_pre(gw, &hq, &hdq, h, 1)?)
2352        } else {
2353            (
2354                e.matmul(&fa.wq, h, 1)?,
2355                e.matmul(&fa.wk, h, 1)?,
2356                e.matmul(&fa.wv, h, 1)?,
2357                e.matmul(gw, h, 1)?,
2358            )
2359        };
2360
2361        let mut q = e.uninit(nh * hd)?;
2362        e.rms_norm(&q0, fa.q_norm.float_data(), &mut q, hd, nh, eps)?;
2363        let mut k = e.uninit(nkv * hd)?;
2364        e.rms_norm(&k0, fa.k_norm.float_data(), &mut k, hd, nkv, eps)?;
2365        // `rope_freqs.weight` (llama3 factors) applies to the FULL-attn layers ONLY; SWA passes
2366        // null (llama-hparams / step35.cpp). Block 45 is SWA, so `ff` is None there — but read
2367        // the resolved flag, not the constant, so an all-full sibling stays correct.
2368        let ff = if g.swa {
2369            None
2370        } else {
2371            self.step35_aux.as_ref().and_then(|a| a.rope_freqs(e))
2372        };
2373        #[cfg(debug_assertions)]
2374        if let Some(ff) = ff {
2375            crate::debug_assert_tensor_stream_device(ff, &e.stream(), "mtp_step35_attn.rope_freqs");
2376        }
2377        e.rope_neox2(
2378            &mut q,
2379            &mut k,
2380            pos_d,
2381            hd,
2382            g.n_rot,
2383            nh,
2384            nkv,
2385            1,
2386            g.rope_base,
2387            1.0,
2388            ff,
2389        )?;
2390
2391        // Append at the HOST slot, then re-stamp the device counter: the eager chain has the
2392        // length on the host anyway, and the windowed view below needs it there to compute the
2393        // offset. The device counter is kept in lockstep so `mtp_kv_fill`'s `set_i32_one` and any
2394        // dc-family consumer of this scratch still agree.
2395        let kv = &mut scratch.kv;
2396        assert!(
2397            kv.len < scratch.cap,
2398            "step35 MTP scratch overflow ({} >= {})",
2399            kv.len,
2400            scratch.cap
2401        );
2402        let next_len = kv.len + 1;
2403        let (off, t_kv) = if g.swa && next_len > g.window {
2404            (next_len - g.window, g.window)
2405        } else {
2406            (0, next_len)
2407        };
2408        let write_row = e.prepare_kv_append(kv, off & !31usize, 1)?;
2409        e.append_kv_quantized(
2410            &k,
2411            &v0,
2412            &mut kv.k,
2413            &mut kv.v,
2414            write_row,
2415            kv.kv_dim_k,
2416            kv.kv_dim_v,
2417            kv.k_tok_bytes,
2418            kv.v_tok_bytes,
2419            false,
2420        )?;
2421        kv.len = next_len;
2422        e.set_i32_one(&mut kv.len_d, kv.len as i32)?;
2423        // SWA view offset (see note 1). The draft chain is short (k+2 rows), but the scratch is
2424        // PERSISTENT across rounds — `mtp_kv_fill` leaves one row per committed token behind, so
2425        // `kv.len` tracks absolute position and crosses 512 in any real generation. The window is
2426        // therefore live, not theoretical.
2427        let physical = kv.physical_rows(off, off + t_kv)?;
2428        let k_view = e.view_u8_range(
2429            &kv.k,
2430            physical.start * kv.k_tok_bytes,
2431            physical.end * kv.k_tok_bytes,
2432        );
2433        let v_view = e.view_u8_range(
2434            &kv.v,
2435            physical.start * kv.v_tok_bytes,
2436            physical.end * kv.v_tok_bytes,
2437        );
2438        let mut attn = e.uninit(nh * hd)?;
2439        e.fa_decode_kvmod(
2440            &q,
2441            &k_view,
2442            &v_view,
2443            &mut attn,
2444            hd,
2445            nh,
2446            nkv,
2447            t_kv,
2448            scale,
2449            kv.k_tok_bytes,
2450            kv.v_tok_bytes,
2451            false,
2452        )?;
2453
2454        let mut ag = e.uninit(nh * hd)?;
2455        e.attn_head_gate(&attn, &gt, &mut ag, None, hd, nh, 1)?;
2456        Ok(e.matmul(&fa.wo, &ag, 1)?)
2457    }
2458
2459    /// MTP-block full attention, T=1, on the scratch KV (BOTH draft paths — eager and graph):
2460    /// the scratch write slot and the attention bound come from `scratch.kv.len_d` (device i32[1])
2461    /// so the launch args are FIXED across draft steps — ONE captured graph serves the whole
2462    /// chain, and replays keep seeing KV growth through the device counter (no recapture).
2463    /// Geometry contract: n_splits is sized from `scratch.cap` (the persistent capacity); splits
2464    /// whose key range lies beyond the device t_kv exit empty and the shared combine skips them
2465    /// (fa_decode_dc bit-correct-for-any-t_kv<=bucket_max contract). The eager path uses the SAME
2466    /// launcher with the SAME bucket_max -> identical dispatch -> bit-identical draft tokens (the
2467    /// graph-vs-eager parity gate). Host len is NOT advanced here (graph contract); callers mirror.
2468    fn mtp_full_attn_dc(
2469        &self,
2470        e: &Engine,
2471        fa: &FullAttnLayer,
2472        h: &CudaSlice<f32>,
2473        pos_d: &CudaSlice<i32>,
2474        scratch: &mut MtpScratch,
2475        geom: Option<&crate::hybrid::DraftGeom>,
2476    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2477        let cfg = &self.cfg;
2478        let mtp_il = cfg.n_layer.saturating_sub(cfg.nextn_predict_layers);
2479        let geometry = cfg.full_attention_geometry_at(mtp_il);
2480        let n_head = geom.map(|g| g.n_head).unwrap_or(geometry.n_head as usize);
2481        let n_head_kv = geom
2482            .map(|g| g.n_head_kv)
2483            .unwrap_or(geometry.n_head_kv as usize);
2484        let head_dim = geometry.head_dim_k as usize;
2485        let eps = cfg.rms_eps;
2486        let scale = geometry.attention_scale();
2487        let n_embd = geom.map(|g| g.d_inner).unwrap_or(cfg.n_embd as usize);
2488        let bucket_max = scratch.cap; // < 96 guaranteed by the graph_draft eligibility gate
2489
2490        let (qf, mut k, v) =
2491            if e.uses_q8_1_fast(&fa.wq) && e.uses_q8_1_fast(&fa.wk) && e.uses_q8_1_fast(&fa.wv) {
2492                let (hq, hd) = e.quantize_q8_1(h, 1, n_embd)?;
2493                (
2494                    e.matmul_pre(&fa.wq, &hq, &hd, h, 1)?,
2495                    e.matmul_pre(&fa.wk, &hq, &hd, h, 1)?,
2496                    e.matmul_pre(&fa.wv, &hq, &hd, h, 1)?,
2497                )
2498            } else {
2499                (
2500                    e.matmul(&fa.wq, h, 1)?,
2501                    e.matmul(&fa.wk, h, 1)?,
2502                    e.matmul(&fa.wv, h, 1)?,
2503                )
2504            };
2505        // M3/Hy3 have no attention output gate — wq out is exactly q; skip the split.
2506        let gated = geometry.attention_gate == memra_gguf::config::AttentionGateKind::FusedQ;
2507        let (mut q, gate) = if gated {
2508            let mut q = e.zeros(n_head * head_dim)?;
2509            let mut gate = e.zeros(n_head * head_dim)?;
2510            e.q_gate_split(&qf, &mut q, &mut gate, head_dim, n_head, 1)?;
2511            (q, Some(gate))
2512        } else {
2513            (qf, None)
2514        };
2515
2516        let mut qn = e.zeros(n_head * head_dim)?;
2517        e.rms_norm(&q, fa.q_norm.float_data(), &mut qn, head_dim, n_head, eps)?;
2518        q = qn;
2519        let mut kn = e.zeros(n_head_kv * head_dim)?;
2520        e.rms_norm(
2521            &k,
2522            fa.k_norm.float_data(),
2523            &mut kn,
2524            head_dim,
2525            n_head_kv,
2526            eps,
2527        )?;
2528        k = kn;
2529        let rope_dims = geometry.n_rot as usize;
2530        e.rope_neox(
2531            &mut q,
2532            pos_d,
2533            head_dim,
2534            rope_dims,
2535            n_head,
2536            1,
2537            geometry.rope_base,
2538            1.0,
2539        )?;
2540        e.rope_neox(
2541            &mut k,
2542            pos_d,
2543            head_dim,
2544            rope_dims,
2545            n_head_kv,
2546            1,
2547            geometry.rope_base,
2548            1.0,
2549        )?;
2550
2551        let kv = &mut scratch.kv;
2552        // append at the DEVICE slot (kv.len_d == old len), then advance the counter in-graph.
2553        e.append_kv_quantized_dc(
2554            &k,
2555            &v,
2556            &mut kv.k,
2557            &mut kv.v,
2558            &kv.len_d,
2559            kv.kv_dim_k,
2560            kv.kv_dim_v,
2561            kv.k_tok_bytes,
2562            kv.v_tok_bytes,
2563            false,
2564        )?;
2565        e.inc_seqlen(&mut kv.len_d)?;
2566        // full-buffer views (any in-round t_kv stays in range on replay); the kernel bounds the
2567        // key range from the device counter.
2568        let k_view = e.view_u8(&kv.k, kv.k.len());
2569        let v_view = e.view_u8(&kv.v, kv.v.len());
2570        let (ktb, vtb) = (kv.k_tok_bytes, kv.v_tok_bytes);
2571        let mut attn = e.zeros(n_head * head_dim)?;
2572        e.fa_decode_dc(
2573            &q, &k_view, &v_view, &mut attn, head_dim, n_head, n_head_kv, &kv.len_d, bucket_max,
2574            scale, ktb, vtb, false,
2575        )?;
2576
2577        let attn_g = match &gate {
2578            Some(gate) => {
2579                let mut gsig = e.zeros(n_head * head_dim)?;
2580                e.sigmoid(gate, &mut gsig, n_head * head_dim)?;
2581                let mut ag = e.zeros(n_head * head_dim)?;
2582                e.mul(&attn, &gsig, &mut ag, n_head * head_dim)?;
2583                ag
2584            }
2585            None => attn,
2586        };
2587        Ok(e.matmul(&fa.wo, &attn_g, 1)?)
2588    }
2589
2590    /// PERSISTENT-DRAFT-KV fill (the reference engine's "mtp_update" analogue): compute the MTP
2591    /// block's K/V for `tokens` (committed tokens at positions pos0..pos0+T) from their EXACT
2592    /// trunk hiddens `h` ([T, n_embd] token-major, pre-output_norm) and append at slots pos0.. of
2593    /// the scratch KV. K/V-ONLY — ops A/1-5 plus the K-side of op 6 (wk/wv + k_norm + rope +
2594    /// quantized append); no wq/attention/FFN/lm_head, so per-token cost ~= eh_proj + wk/wv (a
2595    /// small fraction of one trunk layer), T-batched. Rope follows the chain convention
2596    /// rope(token@p) = p+1. Runs at round boundaries OUTSIDE the captured graph in BOTH draft
2597    /// modes -> draft parity by construction. Caller must have scratch.kv.len == pos0.
2598    #[allow(clippy::too_many_arguments)]
2599    fn mtp_kv_fill(
2600        &self,
2601        e: &Engine,
2602        mtp: &MtpHead,
2603        tokens: &[u32],
2604        h: &CudaSlice<f32>,
2605        pos0: usize,
2606        scratch: &mut MtpScratch,
2607        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
2608    ) -> Result<(), Box<dyn std::error::Error>> {
2609        let cfg = &self.cfg;
2610        let n_embd = cfg.n_embd as usize;
2611        let eps = cfg.rms_eps;
2612        let t = tokens.len();
2613        assert_eq!(scratch.kv.len, pos0, "mtp_kv_fill: append slot mismatch");
2614        assert!(pos0 + t <= scratch.cap, "mtp_kv_fill: scratch overflow");
2615        let Mixer::Full(fa) = &mtp.mixer else {
2616            panic!("MTP block is full-attn in qwen35; linear MTP not supported")
2617        };
2618        let pos_vec: Vec<i32> = (0..t).map(|i| (pos0 + i + 1) as i32).collect();
2619        let pos_d = e.htod_i32(&pos_vec)?;
2620
2621        // ops A/1/2: embed + the two input norms, T-wide.
2622        let e_emb = match embd_dev {
2623            Some((g, qt, rb)) => e.embed_gather_device_t(g, tokens, n_embd, qt, rb)?,
2624            None => e.htod(&self.embd.gather(n_embd, tokens))?,
2625        };
2626        let mut e_norm = e.zeros(t * n_embd)?;
2627        e.rms_norm(&e_emb, mtp.enorm.float_data(), &mut e_norm, n_embd, t, eps)?;
2628        let mut h_norm = e.zeros(t * n_embd)?;
2629        e.rms_norm(h, mtp.hnorm.float_data(), &mut h_norm, n_embd, t, eps)?;
2630
2631        // op 3: per-row [e_norm ; h_norm] concat, token-major [T, 2*n_embd].
2632        let mut concat = e.zeros(t * 2 * n_embd)?;
2633        for i in 0..t {
2634            e.copy_view_into(
2635                &mut concat,
2636                i * 2 * n_embd,
2637                &e_norm.slice(i * n_embd..(i + 1) * n_embd),
2638                n_embd,
2639            )?;
2640            e.copy_view_into(
2641                &mut concat,
2642                i * 2 * n_embd + n_embd,
2643                &h_norm.slice(i * n_embd..(i + 1) * n_embd),
2644                n_embd,
2645            )?;
2646        }
2647
2648        // ops 4/5: eh_proj + attn_norm, T-wide (at the student inner width when geom is set).
2649        let di = mtp.geom.as_ref().map(|g| g.d_inner).unwrap_or(n_embd);
2650        let inp_sa = e.matmul(&mtp.eh_proj, &concat, t)?;
2651        let mut a_norm = e.zeros(t * di)?;
2652        e.rms_norm(&inp_sa, mtp.attn_norm.float_data(), &mut a_norm, di, t, eps)?;
2653
2654        // op 6 (K/V half): wk/wv + k_norm + rope + per-row quantized append. No wq/attention —
2655        // the fill only has to leave correct K/V rows behind for later chains to attend over.
2656        let n_head_kv = mtp
2657            .geom
2658            .as_ref()
2659            .map(|g| g.n_head_kv)
2660            .or(mtp.step35.as_ref().map(|s| s.n_head_kv))
2661            .unwrap_or_else(|| {
2662                let mtp_il = cfg.n_layer.saturating_sub(cfg.nextn_predict_layers);
2663                cfg.full_attention_geometry_at(mtp_il).n_head_kv as usize
2664            });
2665        let mtp_il = cfg.n_layer.saturating_sub(cfg.nextn_predict_layers);
2666        let geometry = cfg.full_attention_geometry_at(mtp_il);
2667        let head_dim = geometry.head_dim_k as usize;
2668        let mut k = e.matmul(&fa.wk, &a_norm, t)?;
2669        let v = e.matmul(&fa.wv, &a_norm, t)?;
2670        let mut kn = e.zeros(t * n_head_kv * head_dim)?;
2671        e.rms_norm(
2672            &k,
2673            fa.k_norm.float_data(),
2674            &mut kn,
2675            head_dim,
2676            n_head_kv * t,
2677            eps,
2678        )?;
2679        k = kn;
2680        // step35: rotary width AND base are per-layer, and the MTP block's values come from the
2681        // resolved `Step35MtpGeom` — NOT from `cfg.rope_dim_count`/`cfg.rope_freq_base`, which
2682        // carry the arch defaults (128 / 5e6, i.e. the FULL-attn layers' base). Getting this wrong
2683        // writes K rows the attention arm then re-derives at a different theta: correct-looking
2684        // output with dead acceptance, invisible to the exactness gates.
2685        let (rope_dims, rope_base, ff) = match mtp.step35.as_ref() {
2686            Some(s) => (
2687                s.n_rot,
2688                s.rope_base,
2689                if s.swa {
2690                    None
2691                } else {
2692                    self.step35_aux.as_ref().and_then(|a| a.rope_freqs(e))
2693                },
2694            ),
2695            None => (geometry.n_rot as usize, geometry.rope_base, None),
2696        };
2697        #[cfg(debug_assertions)]
2698        if let Some(ff) = ff {
2699            crate::debug_assert_tensor_stream_device(ff, &e.stream(), "mtp_kv_fill.rope_freqs");
2700        }
2701        match ff {
2702            Some(f) => e.rope_neox_ff(
2703                &mut k, &pos_d, head_dim, rope_dims, n_head_kv, t, rope_base, 1.0, f,
2704            )?,
2705            None => e.rope_neox(
2706                &mut k, &pos_d, head_dim, rope_dims, n_head_kv, t, rope_base, 1.0,
2707            )?,
2708        }
2709
2710        let kv = &mut scratch.kv;
2711        // Match the trunk prime contract: a chunk may need the aligned window immediately before
2712        // its first row, so preserve that prefix when the physical tail rebases at wrap.
2713        let retain_from = kv
2714            .ring
2715            .as_ref()
2716            .map(|ring| pos0.saturating_sub(ring.window() - 1) & !31usize)
2717            .unwrap_or(0);
2718        let write_row = e.prepare_kv_append(kv, retain_from, t)?;
2719        for i in 0..t {
2720            let k_row = k.slice(i * kv.kv_dim_k..(i + 1) * kv.kv_dim_k);
2721            let v_row = v.slice(i * kv.kv_dim_v..(i + 1) * kv.kv_dim_v);
2722            e.append_kv_quantized_view(
2723                &k_row,
2724                &v_row,
2725                &mut kv.k,
2726                &mut kv.v,
2727                write_row + i,
2728                kv.kv_dim_k,
2729                kv.kv_dim_v,
2730                kv.k_tok_bytes,
2731                kv.v_tok_bytes,
2732                false,
2733            )?;
2734        }
2735        kv.len = pos0 + t;
2736        e.set_i32_one(&mut kv.len_d, kv.len as i32)?;
2737        Ok(())
2738    }
2739
2740    /// CAPTURE body for the GRAPH DRAFT (stage 2 of graph-grade spec): ONE MTP head forward with
2741    /// every varying input device-resident —
2742    ///   - token id from the persistent `tok_d` (the previous replay's in-graph argmax wrote it,
2743    ///     so the chain feeds itself; the host reads the same 4 bytes for the draft list),
2744    ///   - h_seed from the persistent `h_seed_d` (h_nextn is copied BACK into it at the end),
2745    ///   - rope pos from the persistent `pos_d` counter (inc'd in-graph),
2746    ///   - scratch KV slot/bound from `scratch.kv.len_d` (see mtp_full_attn_dc).
2747    /// The p-min confidence lands in the persistent `p_d` iff `with_prob` (env is fixed per run).
2748    /// Same kernels, same dispatch as the eager mtp_head_forward_dev chain -> same draft tokens
2749    /// (exactness never depends on drafts — the verify arbitrates — but acceptance parity does).
2750    /// `with_head=false` captures the HEAD-LESS twin for the pseudo-seed replay (2026-07-03):
2751    /// the pseudo pass only needs h_nextn (op 10) + the scratch append — the lm_head read
2752    /// (~1.06ms q6_K on the 9B), argmax and prob are dead weight there. h_nextn's inputs are
2753    /// untouched, so the seed value is identical; round-start resets overwrite tok_d/p_d anyway.
2754    /// `sampled_cap` = Some((ctr_d, perturb_d, q_out_d, seed, temp)) captures the SAMPLED twin
2755    /// (step 3 of the sampled-spec arc): head logits are retained in the persistent `q_out_d`
2756    /// (host D2Ds them to the round's q slot after each replay), the DEVICE event counter is
2757    /// bumped in-graph, and the argmax reads GUMBEL-PERTURBED logits — one categorical draw per
2758    /// replay, bit-identical to the eager arm's gumbel_perturb at the same (seed, sctr, temp).
2759    /// seed/temp are capture-time constants (fixed per generate call, like p_min).
2760    #[allow(clippy::too_many_arguments)]
2761    fn mtp_head_forward_cap(
2762        &self,
2763        e: &Engine,
2764        mtp: &MtpHead,
2765        tok_d: &mut CudaSlice<u32>,
2766        pos_d: &mut CudaSlice<i32>,
2767        h_seed_d: &mut CudaSlice<f32>,
2768        p_d: &mut CudaSlice<f32>,
2769        scratch: &mut MtpScratch,
2770        with_prob: bool,
2771        with_head: bool,
2772        embd_gpu: &CudaSlice<u8>,
2773        embd_qt: i32,
2774        embd_rb: usize,
2775        d_vocab: usize,
2776        sampled_cap: Option<(
2777            &mut CudaSlice<u32>,
2778            &mut CudaSlice<f32>,
2779            &mut CudaSlice<f32>,
2780            u64,
2781            f32,
2782        )>,
2783        stream_pack: Option<(&mut CudaSlice<u32>, usize, Option<&CudaSlice<u32>>)>,
2784        // DRAFT-SIDE GRAMMAR MASK (lane/draft-mask): (packed draft-vocab allowed-set buffer,
2785        // word count). Captured as ONE mask_logits_f32 node between the head matmul and the
2786        // in-graph argmax; the buffer address is baked, its CONTENTS are re-uploaded by the
2787        // host before every replay (the decode.rs graph-mask pattern). All-ones contents = a
2788        // no-op ban, so a position the grammar cannot constrain costs one pass over the row.
2789        mask_cap: Option<(&CudaSlice<u32>, usize)>,
2790    ) -> Result<(), Box<dyn std::error::Error>> {
2791        let cfg = &self.cfg;
2792        let n_embd = cfg.n_embd as usize;
2793        // step35 REFUSAL (deliberate, named): the graph body's attention is `mtp_full_attn_dc`,
2794        // whose device-counter key bound always starts at row 0 — it cannot express this block's
2795        // SWA view offset, so a captured chain would silently attend OUTSIDE the window once the
2796        // persistent scratch passes 512 rows. Nothing is lost today: `graph_draft` also requires
2797        // `trunk_dense`, and step35's trunk is a 288-expert MoE, so the eager chain
2798        // (`mtp_head_forward_dev` -> `mtp_step35_attn`) is the served path. Returning Err (not a
2799        // panic) is what the two capture sites and the round-stream capture already handle by
2800        // degrading to eager / stream-off.
2801        if mtp.step35.is_some() {
2802            return Err(
2803                "step35 has no captured draft chain (fa_decode_dc cannot express the MTP \
2804                        block's SWA view offset; same root cause as the dc decode refusal) — the \
2805                        eager draft chain serves this arch"
2806                    .into(),
2807            );
2808        }
2809        // student inner width (see mtp_head_forward_dev) — interface dims stay n_embd.
2810        let di = mtp.geom.as_ref().map(|g| g.d_inner).unwrap_or(n_embd);
2811        let eps = cfg.rms_eps;
2812        let e_emb = e.embed_gather_device(embd_gpu, tok_d, n_embd, embd_qt, embd_rb)?;
2813        let mut e_norm = e.zeros(n_embd)?;
2814        e.rms_norm(&e_emb, mtp.enorm.float_data(), &mut e_norm, n_embd, 1, eps)?;
2815        let mut h_norm = e.zeros(n_embd)?;
2816        e.rms_norm(
2817            &*h_seed_d,
2818            mtp.hnorm.float_data(),
2819            &mut h_norm,
2820            n_embd,
2821            1,
2822            eps,
2823        )?;
2824        let mut concat = e.zeros(2 * n_embd)?;
2825        e.copy_into(&mut concat, 0, &e_norm, n_embd)?;
2826        e.copy_into(&mut concat, n_embd, &h_norm, n_embd)?;
2827        let inp_sa = e.matmul(&mtp.eh_proj, &concat, 1)?;
2828        let mut a_norm = e.zeros(di)?;
2829        e.rms_norm(&inp_sa, mtp.attn_norm.float_data(), &mut a_norm, di, 1, eps)?;
2830        let attn_out = match &mtp.mixer {
2831            Mixer::Full(fa) => {
2832                self.mtp_full_attn_dc(e, fa, &a_norm, pos_d, scratch, mtp.geom.as_ref())?
2833            }
2834            Mixer::Linear(_) => {
2835                panic!("MTP block is full-attn in qwen35; linear MTP not supported")
2836            }
2837            Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
2838        };
2839        let mut x1 = e.zeros(di)?;
2840        e.add(&inp_sa, &attn_out, &mut x1, di)?;
2841        let mut z = e.zeros(di)?;
2842        e.rms_norm(&x1, mtp.post_attn_norm.float_data(), &mut z, di, 1, eps)?;
2843        let ffn_out = match &mtp.ffn {
2844            crate::hybrid::Ffn::Dense {
2845                ffn_gate,
2846                ffn_up,
2847                ffn_down,
2848            } => {
2849                let n_ff = ffn_gate.out_features();
2850                let (gate, up) = if e.uses_q8_1_fast(ffn_gate) && e.uses_q8_1_fast(ffn_up) {
2851                    let (zq, zd) = e.quantize_q8_1(&z, 1, di)?;
2852                    (
2853                        e.matmul_pre(ffn_gate, &zq, &zd, &z, 1)?,
2854                        e.matmul_pre(ffn_up, &zq, &zd, &z, 1)?,
2855                    )
2856                } else {
2857                    (e.matmul(ffn_gate, &z, 1)?, e.matmul(ffn_up, &z, 1)?)
2858                };
2859                let mut act = e.zeros(n_ff)?;
2860                Self::ffn_act(e, &self.cfg, &gate, &up, &mut act, n_ff)?;
2861                e.matmul(ffn_down, &act, 1)?
2862            }
2863            // ROUND-STREAM: the 35B NextN block carries a MoE FFN. With RESIDENT experts the
2864            // dev path is pure device launches (device top-k + rows kernels, ZERO-DtoH by
2865            // design) — capture-legal. Non-resident (SLRU-lock) stays rejected: the capture
2866            // error arm degrades the caller to eager/stream-off.
2867            crate::hybrid::Ffn::Moe(m) if m.dev_exps.is_some() => {
2868                self.moe_ffn_il(e, m, &z, 1, u16::MAX)?
2869            }
2870            crate::hybrid::Ffn::Moe(_) => {
2871                return Err("graph draft requires a Dense (or resident-MoE) MTP FFN".into());
2872            }
2873        };
2874        let mut h_inner = e.zeros(di)?;
2875        e.add(&x1, &ffn_out, &mut h_inner, di)?;
2876        // student: up-project back to n_embd (carrier + head input; see mtp_head_forward_dev).
2877        let h_nextn = match mtp.geom.as_ref() {
2878            Some(g) => e.matmul(&g.out_up, &h_inner, 1)?,
2879            None => h_inner,
2880        };
2881        // MEMRA_SPEC_HPOST needs final_h even head-less (it IS the next seed under that convention).
2882        let final_h = if with_head || spec_hpost() {
2883            let final_norm = mtp.shared_head_norm.as_ref().unwrap_or(&self.output_norm);
2884            let mut fh = e.zeros(n_embd)?;
2885            e.rms_norm(&h_nextn, final_norm.float_data(), &mut fh, n_embd, 1, eps)?;
2886            Some(fh)
2887        } else {
2888            None
2889        };
2890        if with_head {
2891            let head = mtp.shared_head_head.as_ref().unwrap_or(&self.output);
2892            let mut logits = e.matmul(head, final_h.as_ref().unwrap(), 1)?;
2893            // DRAFT-SIDE GRAMMAR MASK: ban the grammar-illegal draft ids IN the captured chain,
2894            // before the argmax — proposals become legal by construction. Contents-only
2895            // per-replay upload keeps the capture valid.
2896            if let Some((mask_d, mw)) = mask_cap {
2897                e.mask_logits_col(&mut logits, mask_d, 0, d_vocab, mw)?;
2898            }
2899            if let Some((ctr_d, perturb_d, q_out_d, seed, temp)) = sampled_cap {
2900                // SAMPLED chain: retain q (raw head logits -> persistent q_out_d; the matmul's
2901                // own buffer is pool-recycled after the capture body returns, so it can't be the
2902                // retention target), bump the device event counter, gumbel-perturb reading it,
2903                // and argmax the PERTURBED logits into tok_d — the in-graph categorical draw.
2904                e.copy_into(q_out_d, 0, &logits, d_vocab)?;
2905                e.sctr_inc(ctr_d)?;
2906                e.gumbel_perturb_ctr(&logits, perturb_d, d_vocab, seed, ctr_d, temp)?;
2907                e.argmax_token_device_into(perturb_d, tok_d, d_vocab)?;
2908                // p-min prob = the head's RAW softmax confidence in the SAMPLED pick — same
2909                // semantics as the eager sampled arm's prob_of_token_device(dl_d, tok_d).
2910                if with_prob {
2911                    e.prob_of_token_device_into(&logits, tok_d, p_d, d_vocab)?;
2912                }
2913            } else {
2914                // draft token -> persistent tok_d (next replay's embed reads it; host reads the 4 bytes).
2915                e.argmax_token_device_into(&logits, tok_d, d_vocab)?;
2916                // p-min under a draft mask reads the MASKED row: confidence relative to the
2917                // grammar-LEGAL alternatives (illegal ids leave the softmax denominator), which
2918                // is the right semantics for "does the drafter know what comes next here" and
2919                // the same row the pick came from. Draft-quality only — verify arbitrates.
2920                if with_prob {
2921                    e.prob_of_token_device_into(&logits, tok_d, p_d, d_vocab)?;
2922                }
2923            }
2924        }
2925        // ROUND-STREAM K-chain: pack (tok, p) into slot j, then remap tok through d2t so the
2926        // NEXT chained body's embed reads the TARGET id — zero host involvement per step.
2927        if let Some((out, slot, d2t)) = stream_pack {
2928            e.pack_tok_p(tok_d, p_d, out, slot)?;
2929            if let Some(map) = d2t {
2930                e.tok_map_u32(tok_d, map)?;
2931            }
2932        }
2933        // Next draft step's h_seed: pre-norm h_nextn (default) or post-norm final_h (HPOST).
2934        if spec_hpost() {
2935            e.copy_into(h_seed_d, 0, final_h.as_ref().unwrap(), n_embd)?;
2936        } else {
2937            e.copy_into(h_seed_d, 0, &h_nextn, n_embd)?;
2938        }
2939        // advance the draft rope position in-graph.
2940        e.inc_seqlen(pos_d)?;
2941        Ok(())
2942    }
2943
2944    /// Batched target verify forward over `tokens` at positions `pos0..pos0+T` (§D.3, T=K+1).
2945    /// Returns ALL T logit columns (host f32, [T*n_vocab]); appends T cols to every full-attn KV
2946    /// and advances every linear-attn recur state by T steps (the recur steps are SEQUENTIAL T=1).
2947    /// Advances `cache.pos` by T.
2948    pub fn decode_step_t(
2949        &self,
2950        e: &Engine,
2951        tokens: &[u32],
2952        pos0: usize,
2953        cache: &mut Cache,
2954    ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
2955        if self.is_gemma4_e4b() {
2956            return Ok(self.gemma4_e4b_decode_step_t_h(e, tokens, pos0, cache)?.0);
2957        }
2958        if self.cfg.gemma4.is_some() {
2959            return self.gemma4_decode_step_t(e, tokens, pos0, cache);
2960        }
2961        Ok(self.decode_step_t_h(e, tokens, pos0, cache)?.0)
2962    }
2963
2964    /// Like `decode_step_t` but ALSO returns the LAST column's pre-output_norm hidden (h_seed for
2965    /// the next draft round). This lets partial-accept replay run as ONE batched T=(n_acc+1) forward
2966    /// (single weight read) instead of n_acc+1 separate T=1 decode_steps (n_acc+1 weight reads).
2967    /// At batch=1 decode is bandwidth-bound, so batching the replay is THE MTP profitability lever.
2968    pub fn decode_step_t_h(
2969        &self,
2970        e: &Engine,
2971        tokens: &[u32],
2972        pos0: usize,
2973        cache: &mut Cache,
2974    ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
2975        self.decode_step_t_h_emb(e, tokens, pos0, cache, None)
2976    }
2977
2978    /// Like `decode_step_t_h` with an optional RESIDENT embed table (spec hot loop): device
2979    /// gather instead of host dequant + [T, n_embd] f32 htod. Bit-identical rows.
2980    pub fn decode_step_t_h_emb(
2981        &self,
2982        e: &Engine,
2983        tokens: &[u32],
2984        pos0: usize,
2985        cache: &mut Cache,
2986        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
2987    ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
2988        let (logits_d, h_seed) = self.decode_step_t_h_emb_dev(e, tokens, pos0, cache, embd_dev)?;
2989        Ok((e.dtoh(&logits_d)?, h_seed))
2990    }
2991
2992    /// DEVICE-LOGITS verify forward (spec device-argmax lever): identical kernel chain to
2993    /// `decode_step_t_h_emb` but returns the [T, n_vocab] logits ON DEVICE — the accept walk
2994    /// argmaxes each column on-device and reads back ONE [T] u32 instead of dtoh'ing the full
2995    /// T x n_vocab f32 block (~1-4 MB + T host argmaxes, every round). Kernel dispatch is
2996    /// UNCHANGED (same decode-exact kernels); only the post-logits transfer moves.
2997    pub fn decode_step_t_h_emb_dev(
2998        &self,
2999        e: &Engine,
3000        tokens: &[u32],
3001        pos0: usize,
3002        cache: &mut Cache,
3003        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
3004    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
3005        let n_embd = self.cfg.n_embd as usize;
3006        let t = tokens.len();
3007        let (logits, x) = self.decode_step_t_core(e, tokens, pos0, cache, embd_dev, None)?;
3008        // h_seed for the next round = LAST column's pre-output_norm hidden ([n_embd]).
3009        let mut hs = vbuf(e, n_embd)?; // fully written by copy_view_into below
3010        e.copy_view_into(&mut hs, 0, &x.slice((t - 1) * n_embd..t * n_embd), n_embd)?;
3011        Ok((logits, hs))
3012    }
3013
3014    /// CORE verify forward: the `decode_step_t_h_emb_dev` kernel chain, returning the FULL
3015    /// pre-output_norm hidden stack x ([T, n_embd], any column extractable) and optionally
3016    /// filling a `VerifyCkpt` (retained per-layer state-rebuild inputs) for the REPLAY-FREE
3017    /// partial accept. `ckpt: None` => byte-for-byte the old behavior (the ckpt writes are pure
3018    /// retains/copies — they never change what any kernel computes).
3019    fn decode_step_t_core(
3020        &self,
3021        e: &Engine,
3022        tokens: &[u32],
3023        pos0: usize,
3024        cache: &mut Cache,
3025        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
3026        mut ckpt: Option<&mut VerifyCkpt>,
3027    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
3028        self.decode_step_t_core_stream(e, tokens, pos0, cache, embd_dev, ckpt.take(), None, None)
3029    }
3030
3031    /// Increment-0 two-session PP seam: release the peer after this lane's stage-0 boundary TX.
3032    /// The two independent sessions keep their own cache/checkpoint state; only issue order moves.
3033    fn decode_step_t_core_pipelined(
3034        &self,
3035        e: &Engine,
3036        tokens: &[u32],
3037        pos0: usize,
3038        cache: &mut Cache,
3039        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
3040        mut ckpt: Option<&mut VerifyCkpt>,
3041        pipe: &SpecPipeLane,
3042        round: usize,
3043    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
3044        let fence = crate::pp::pp_cuts(self.layers.len())
3045            .ok_or("two-session speculative pipeline requires a PP stage cut")?;
3046        if crate::pp::pp2_streams_off() || !crate::pp::spec_pp_on() {
3047            return Err("two-session speculative pipeline requires the PP verify split".into());
3048        }
3049        let interval_fence = pipe.stage0_begin(round)?;
3050        let ticket = self.verify_stage0_issue(
3051            e,
3052            tokens,
3053            pos0,
3054            cache,
3055            embd_dev,
3056            ckpt.as_deref_mut(),
3057            None,
3058            &fence,
3059            Some(interval_fence),
3060            pipe.trace(round),
3061        )?;
3062        pipe.stage0_end(round);
3063        pipe.stage1_begin(round)?;
3064        let result = self.verify_stage1_finish(e, ticket, cache, ckpt, None, &fence, true)?;
3065        pipe.verify_end(round);
3066        Ok(result)
3067    }
3068
3069    /// ROUND-STREAM stage (c) 4: `stream` = (device verify tokens [t], device pos counter) —
3070    /// when Some, rope positions come from pos_iota over the counter, the embed gathers the
3071    /// device tokens, and full_attn_verify routes appends/FA through the _dc twins reading the
3072    /// SAME counter (every layer's kvl.len == cache.pos, one counter drives all three). The
3073    /// host `tokens`/`pos0` args still size buffers (t is FIXED K+1 in stream mode).
3074    #[allow(clippy::too_many_arguments)]
3075    fn decode_step_t_core_stream(
3076        &self,
3077        e: &Engine,
3078        tokens: &[u32],
3079        pos0: usize,
3080        cache: &mut Cache,
3081        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
3082        mut ckpt: Option<&mut VerifyCkpt>,
3083        stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
3084        pp_pipe: Option<bool>,
3085    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
3086        // PP DOOR (lane/pp2-spec 2026-08-06): the verify trunk now takes its OWN stage split,
3087        // exactly as the eager and batched steps do. This is the single funnel every verify
3088        // forward reaches (decode_step_t / _h / _h_emb / _h_emb_dev / _core all land here), so
3089        // wiring it here wires the whole spec surface — the draft/accept/commit machinery above
3090        // is untouched.
3091        //
3092        // History: pp2-hardening (2026-08-06) made this funnel FAIL CLOSED, because its trunk
3093        // walk was unsplit on one stream and a sharded cross-device placement peer-read every
3094        // remote layer's weights on every spec round (measured 13.9-28x on the batched twin).
3095        // The refusal below survives to cover the residue — MEMRA_SPEC_PP=0, MEMRA_PP_STREAMS=0,
3096        // or a placement whose PpNRt fails to build — so a config that would still walk the
3097        // whole trunk on one stream refuses instead of regressing 28x.
3098        if let Some(fence) = crate::pp::pp_cuts(self.layers.len()) {
3099            if !crate::pp::pp2_streams_off() && crate::pp::spec_pp_on() {
3100                return self.decode_step_t_core_ppn(
3101                    e,
3102                    tokens,
3103                    pos0,
3104                    cache,
3105                    embd_dev,
3106                    ckpt.take(),
3107                    stream,
3108                    &fence,
3109                    pp_pipe,
3110                );
3111            }
3112        }
3113        crate::pp::refuse_unsplit_if_remote(
3114            "decode_step_t (spec verify)",
3115            "drop MEMRA_SPEC_PP=0 / MEMRA_PP_STREAMS=0 so the verify trunk takes its OWN stage \
3116             split (decode_step_t_core_ppn); or run spec on one device",
3117        )?;
3118        let cfg = &self.cfg;
3119        let n_embd = cfg.n_embd as usize;
3120        let eps = cfg.rms_eps;
3121        let t = tokens.len();
3122        let pos_d = match stream {
3123            Some((_, ctr)) => {
3124                let mut p = e.alloc_uninit::<i32>(t)?;
3125                e.pos_iota(ctr, &mut p, t)?;
3126                p
3127            }
3128            None => {
3129                let pos_vec: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
3130                e.htod_i32(&pos_vec)?
3131            }
3132        };
3133
3134        // embed T tokens -> [T, n_embd] token-major (device gather on the spec hot loop)
3135        let x = match (stream, embd_dev) {
3136            (Some((vtok, _)), Some((g, qt, rb))) => {
3137                e.embed_gather_device_td(g, vtok, t, n_embd, qt, rb)?
3138            }
3139            (None, Some((g, qt, rb))) => e.embed_gather_device_t(g, tokens, n_embd, qt, rb)?,
3140            _ => e.htod(&self.embd.gather(n_embd, tokens))?,
3141        };
3142
3143        // TRUNK WALK: layers [0, n_layers) through the SAME range-scoped subgraph the PP-N
3144        // stage split calls per stage (`verify_layers`) — one code path, so the split cannot
3145        // drift from the unsplit dispatch mirroring. lane/pp2-spec 2026-08-06.
3146        let x = self.verify_layers(
3147            e,
3148            x,
3149            0,
3150            self.layers.len(),
3151            &pos_d,
3152            pos0,
3153            t,
3154            cache,
3155            ckpt.take(),
3156            stream,
3157        )?;
3158
3159        let mut hn = vbuf(e, t * n_embd)?;
3160        let serving_head = self.cfg.step35.is_some() || self.qwen35_serving_class();
3161        let logits = if serving_head {
3162            // Step35 and the qwen35 family (MoE 2026-08-14 AM, dense-hybrid same day PM — the
3163            // Q3.8 bring-up reproduced the identical near-tie class on dense: eager-class verify
3164            // vs batched-class live serving, ULP drift amplified through the GDN recurrence)
3165            // serve one batched numeric class at every live width, including B=1. Keep the
3166            // verify head in that same class; other generic families retain the decode-exact
3167            // head that their run-spec contract pins.
3168            e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
3169            e.matmul(&self.output, &hn, t)?
3170        } else {
3171            e.rms_norm_decode(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
3172            e.matmul_decode_exact(&self.output, &hn, t)?
3173        };
3174        // stream: the device pos counter owns position; host mirror reconciles at drain.
3175        if stream.is_none() {
3176            cache.pos += t;
3177        }
3178        // Hidden stack for seeds/refresh-fills: pre-norm x (default) or post-norm hn (HPOST).
3179        Ok((logits, if spec_hpost() { hn } else { x }))
3180    }
3181
3182    /// THE VERIFY TRUNK OVER PP-N (lane/pp2-spec 2026-08-06): `decode_step_t_core_stream`'s walk
3183    /// as N stage subgraphs, each on its own engine/stream (and, under `MEMRA_PP_DEVICES`, its own
3184    /// device), with a `[T, n_embd]` boundary transfer between them. T = K+1 (the verify batch),
3185    /// so this is the batched-boundary shape the pp2-batch lane's grow-only slots already handle
3186    /// (`tx(b, x, t*n_embd)`; the slot grows to the high-water T and the transport moves exactly
3187    /// the payload).
3188    ///
3189    /// Structure is `decode_step_batch_ppn`'s, which is `decode_step_h_ppn`'s. FOUR THINGS ARE
3190    /// PER-STAGE and each for a measured reason (see `decode_step_batch_ppn`'s header for the
3191    /// receipts):
3192    ///
3193    /// 1. THE ENGINE (`rt.engine(s, e)`) — `Engine` owns lazily-grown stable-pointer scratch
3194    ///    (`fa_part_pool`, `fa_vf16_scratch`, `argmax_partials`) that is single-stream-safe BY
3195    ///    DESIGN. Two stage streams through one Engine is the 2026-08-02 shared-scratch race
3196    ///    (35% flake, nondeterministic all-logits divergence). `PpNRt::build` gives every stage
3197    ///    s>0 its own Engine even on the primary device; honouring it here is what scopes the
3198    ///    pools. The verify path allocates MORE of that scratch than eager decode does (FA at
3199    ///    m=T, and the per-layer `GdnStash` retains), so this is load-bearing, not inherited.
3200    ///
3201    /// 2. `pos_d` — each stage uploads its OWN copy of the T rope positions on ITS stream, so the
3202    ///    buffer is allocated, consumed and freed on one stream. In `stream` mode that means each
3203    ///    stage runs its own `pos_iota` over the SHARED device counter (`pos_ctr`): the counter is
3204    ///    read-only during the forward (the round's `inc`/`copy_add` happen outside it), so every
3205    ///    stage derives the identical iota, and each stage's own output buffer is stream-local.
3206    ///
3207    /// 3. THE EMBED lives with stage 0 (`self.embd` / `embd_gpu` are host/primary-side; the
3208    ///    sharded loader leaves the table with stage 0 by construction).
3209    ///
3210    /// 4. THE HEAD (`output_norm` + `output`) runs on the LAST stage — the sharded loader uploaded
3211    ///    both through that stage's engine (`hybrid.rs`: `e_head = layer_engine(e, n_trunk,
3212    ///    n_trunk-1)`), so reading them anywhere else is a peer read of the biggest tensor in the
3213    ///    model, every round.
3214    ///
3215    /// WHAT STAYS ON THE PRIMARY, deliberately: the returned logits and hidden stack `x`. Both are
3216    /// last-stage-allocated device buffers, and every consumer (the device argmax walk, the accept
3217    /// kernels, `spec_seed_gather`, the ckpt rebuild in `commit_verified_prefix`) reads them
3218    /// through the primary context by UVA — the same read the batched serving epilogue's
3219    /// `last_logits_dev` park does. Those consumers are per-round O(T x n_vocab) and O(n_embd),
3220    /// not per-layer, so they are not the 28x class; splitting them is a separate lane.
3221    ///
3222    /// The MTP HEAD (draft side) is NOT split: it is one block, it lives wherever the loader put
3223    /// it (`load_mtp` uses the primary engine), and it is ~1-2 GB against the trunk's tens. Draft
3224    /// placement is measured, not assumed — see `research/pp2-spec-20260806`.
3225    ///
3226    /// EXACTNESS: PP-N adds ZERO deviation. Each stage runs the SAME kernels on the SAME bytes in
3227    /// the same order via the SAME `verify_layers` the unsplit body calls; the only change is
3228    /// where the residual is materialized, and the boundary is a straight f32 copy (dtod
3229    /// same-device / `cudaMemcpyPeerAsync` cross-device, no conversion). So the split MUST be
3230    /// BIT-IDENTICAL to the unsplit verify at the same T, in both placement orders. Gate:
3231    /// `decode-batch-gate --mode ppspec`. Acceptance counts are a DERIVED consequence — greedy
3232    /// accept argmaxes these logits, so bit-identical logits force identical accept walks; the
3233    /// `run-spec` K=1..8 arm checks that end-to-end rather than trusting the implication.
3234    #[allow(clippy::too_many_arguments)]
3235    fn decode_step_t_core_ppn(
3236        &self,
3237        e: &Engine,
3238        tokens: &[u32],
3239        pos0: usize,
3240        cache: &mut Cache,
3241        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
3242        mut ckpt: Option<&mut VerifyCkpt>,
3243        stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
3244        fence: &[usize],
3245        pp_pipe: Option<bool>,
3246    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
3247        let ticket = self.verify_stage0_issue(
3248            e,
3249            tokens,
3250            pos0,
3251            cache,
3252            embd_dev,
3253            ckpt.as_deref_mut(),
3254            stream,
3255            fence,
3256            pp_pipe,
3257            None,
3258        )?;
3259        self.verify_stage1_finish(e, ticket, cache, ckpt, stream, fence, true)
3260    }
3261
3262    /// Enqueue embed, stage 0, and the first boundary TX, then return the actual boundary slot.
3263    /// The ordinary PP verify wrapper calls `verify_stage1_finish` immediately after this return.
3264    #[allow(clippy::too_many_arguments)]
3265    fn verify_stage0_issue(
3266        &self,
3267        e: &Engine,
3268        tokens: &[u32],
3269        pos0: usize,
3270        cache: &mut Cache,
3271        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
3272        mut ckpt: Option<&mut VerifyCkpt>,
3273        stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
3274        fence: &[usize],
3275        pp_pipe: Option<bool>,
3276        trace: Option<SpecPipeTraceCtx>,
3277    ) -> Result<VerifyBoundaryTicket, Box<dyn std::error::Error>> {
3278        assert!(
3279            !self.is_gemma4_e4b() && self.cfg.gemma4.is_none(),
3280            "decode_step_t_core_ppn covers the hybrid non-gemma4 verify trunk only \
3281             (the gemma4 arms have their own decode_step_t twins)"
3282        );
3283        if crate::pp::pp_host_bounce_active() && (stream.is_some() || embd_dev.is_some()) {
3284            return Err(
3285                "decode_step_t_core_ppn: refused with MEMRA_PP_HOST_BOUNCE=1 — the trunk \
3286                 boundary itself is host-staged, but device-resident verify still peer-reads \
3287                 primary-device token/position/embedding buffers from stage 0. Run plain PP \
3288                 serving on this host class; spec requires local per-stage inputs first."
3289                    .into(),
3290            );
3291        }
3292        let rt = crate::pp::PpNRt::get(e)?;
3293        let n_st = fence.len() - 1;
3294        assert_eq!(
3295            rt.n_stages(),
3296            n_st,
3297            "PpNRt stage count {} != fence stages {n_st}",
3298            rt.n_stages()
3299        );
3300        let n_embd = self.cfg.n_embd as usize;
3301        let t = tokens.len();
3302        let payload = t * n_embd;
3303        if pp_pipe.is_some() {
3304            assert_eq!(n_st, 2, "spec pipeline requires exactly two PP stages");
3305        }
3306        // One-shot lane diagnostic: force natural PP-2 boundaries to completion so the server
3307        // log can price stage 0, the peer hop, the RX copy, and stage 1 + head separately without
3308        // nsys. The ordinary path keeps every enqueue asynchronous. N>2 is deliberately excluded:
3309        // the report below names exactly two stages and must never imply it measured middle ones.
3310        let pp_anatomy = n_st == 2 && std::env::var("MEMRA_SPEC_PP_ANATOMY").as_deref() == Ok("1");
3311        let pp_started = std::time::Instant::now();
3312        let (mut reverse_ms, mut stage0_ms, mut tx_ms) = (0.0f64, 0.0f64, 0.0f64);
3313        // The CALLER's ambient stream, captured BEFORE any `rt.enter()` pushes a stage stream:
3314        // this body returns DEVICE-RESIDENT buffers (the device-argmax accept walk's contract),
3315        // so the exit needs the same publication the boundaries get — see `PpNRt::publish_to`.
3316        // Taken here, not at the end, because inside the last-stage scope `e.stream()` IS the
3317        // stage stream and the wait would self-order into a no-op.
3318        let caller_stream = e.stream();
3319        // #87 ROOT-CAUSE FENCE (lane/pp2spec-crash): the PREVIOUS round's stage-allocated
3320        // outputs (logits/hidden/ckpt stashes) freed stream-ordered on the STAGE streams while
3321        // the primary stream still holds queued reads of them — with event tracking elided,
3322        // nothing stops the pool from reusing those blocks for THIS round's stage allocations,
3323        // whose writes then race the queued reads (measured: 13/4096-NaN random-bits garbage in
3324        // the spec round seed; the full anatomy is on `PpNRt::fence_stages_behind`). Order every
3325        // stage stream behind the caller before enqueueing new stage work.
3326        let reverse_started = std::time::Instant::now();
3327        if pp_pipe != Some(false) {
3328            rt.fence_stages_behind(&caller_stream)?;
3329        }
3330        if pp_pipe == Some(true) {
3331            // Both session verifies must alternate boundary slots even when the ordinary
3332            // decode overlap experiment is off. Prewarm before A's stage 0 so B cannot grow
3333            // slot 1 by synchronizing the RX stream while A's stage 1 is in flight.
3334            rt.prepare_overlap_slots(0, payload)?;
3335        }
3336        if pp_anatomy {
3337            // Drain the reverse-publication dependency before timing stage 0 itself. At c=1 this
3338            // prices any primary-stream rollback/refresh tail inherited from the prior round.
3339            for s in 0..n_st {
3340                let _st = rt.enter(s);
3341                rt.engine(s, e).stream().synchronize()?;
3342            }
3343            reverse_ms = reverse_started.elapsed().as_secs_f64() * 1e3;
3344        }
3345
3346        // Per-stage rope positions: in host mode the same [T] iota each stage uploads itself; in
3347        // stream mode each stage's own `pos_iota` over the shared read-only device counter.
3348        let stage_pos = |es: &Engine| -> Result<CudaSlice<i32>, Box<dyn std::error::Error>> {
3349            match stream {
3350                Some((_, ctr)) => {
3351                    let mut p = es.alloc_uninit::<i32>(t)?;
3352                    es.pos_iota(ctr, &mut p, t)?;
3353                    Ok(p)
3354                }
3355                None => {
3356                    let pos_vec: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
3357                    es.htod_i32(&pos_vec)
3358                }
3359            }
3360        };
3361
3362        // ---- STAGE 0: embed (the table lives with stage 0) + layers [0, fence[1]) + TX ----
3363        let slot = {
3364            let _st0 = rt.enter(0);
3365            let e0 = rt.engine(0, e);
3366            enqueue_spec_pipe_trace_marker(&e0.stream(), trace.as_ref(), "S0", "start", None)?;
3367            let stage0_started = std::time::Instant::now();
3368            let pos_d = stage_pos(e0)?;
3369            let x = match (stream, embd_dev) {
3370                (Some((vtok, _)), Some((g, qt, rb))) => {
3371                    e0.embed_gather_device_td(g, vtok, t, n_embd, qt, rb)?
3372                }
3373                (None, Some((g, qt, rb))) => e0.embed_gather_device_t(g, tokens, n_embd, qt, rb)?,
3374                _ => e0.htod(&self.embd.gather(n_embd, tokens))?,
3375            };
3376            let x = self.verify_layers(
3377                e0,
3378                x,
3379                fence[0],
3380                fence[1],
3381                &pos_d,
3382                pos0,
3383                t,
3384                cache,
3385                ckpt.as_deref_mut(),
3386                stream,
3387            )?;
3388            if pp_anatomy {
3389                e0.stream().synchronize()?;
3390                stage0_ms = stage0_started.elapsed().as_secs_f64() * 1e3;
3391            }
3392            let tx_started = std::time::Instant::now();
3393            let slot = if pp_pipe.is_some() {
3394                rt.tx_pipelined(0, &x, payload)?
3395            } else {
3396                rt.tx(0, &x, payload)?
3397            };
3398            enqueue_spec_pipe_trace_marker(&e0.stream(), trace.as_ref(), "S0", "end", Some(slot))?;
3399            if pp_anatomy {
3400                e0.stream().synchronize()?;
3401                tx_ms = tx_started.elapsed().as_secs_f64() * 1e3;
3402            }
3403            slot
3404            // x + pos_d drop here: freed stream-ordered on stage-0's stream after use.
3405        };
3406
3407        Ok(VerifyBoundaryTicket {
3408            rt,
3409            caller_stream,
3410            slot,
3411            pos0,
3412            t,
3413            payload,
3414            n_st,
3415            pipelined: pp_pipe.is_some(),
3416            pp_anatomy,
3417            pp_started,
3418            reverse_ms,
3419            stage0_ms,
3420            tx_ms,
3421            trace,
3422        })
3423    }
3424
3425    /// Consume a stage-0 boundary ticket and enqueue the remaining PP stages plus the head.
3426    /// On PP-2 this is exactly stage 1; PP-N keeps its pre-existing middle-stage walk here.
3427    #[allow(clippy::too_many_arguments)]
3428    fn verify_stage1_finish(
3429        &self,
3430        e: &Engine,
3431        ticket: VerifyBoundaryTicket,
3432        cache: &mut Cache,
3433        mut ckpt: Option<&mut VerifyCkpt>,
3434        stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
3435        fence: &[usize],
3436        publish_to_caller: bool,
3437    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
3438        let VerifyBoundaryTicket {
3439            rt,
3440            caller_stream,
3441            slot,
3442            pos0,
3443            t,
3444            payload,
3445            n_st,
3446            pipelined,
3447            pp_anatomy,
3448            pp_started,
3449            reverse_ms,
3450            stage0_ms,
3451            tx_ms,
3452            trace,
3453        } = ticket;
3454        let n_embd = self.cfg.n_embd as usize;
3455        let eps = self.cfg.rms_eps;
3456        let mut slot = slot;
3457        let (mut rx_ms, mut stage1_ms) = (0.0f64, 0.0f64);
3458        let stage_pos = |es: &Engine| -> Result<CudaSlice<i32>, Box<dyn std::error::Error>> {
3459            match stream {
3460                Some((_, ctr)) => {
3461                    let mut p = es.alloc_uninit::<i32>(t)?;
3462                    es.pos_iota(ctr, &mut p, t)?;
3463                    Ok(p)
3464                }
3465                None => {
3466                    let pos_vec: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
3467                    es.htod_i32(&pos_vec)
3468                }
3469            }
3470        };
3471
3472        // ---- MIDDLE STAGES: RX boundary s-1 -> range -> TX boundary s ----
3473        for s in 1..n_st - 1 {
3474            let _st = rt.enter(s);
3475            let es = rt.engine(s, e);
3476            let pos_d = stage_pos(es)?;
3477            let x = rt.rx(s - 1, slot, payload)?;
3478            let x = self.verify_layers(
3479                es,
3480                x,
3481                fence[s],
3482                fence[s + 1],
3483                &pos_d,
3484                pos0,
3485                t,
3486                cache,
3487                ckpt.as_deref_mut(),
3488                stream,
3489            )?;
3490            slot = if pipelined {
3491                rt.tx_pipelined(s, &x, payload)?
3492            } else {
3493                rt.tx(s, &x, payload)?
3494            };
3495        }
3496
3497        // ---- LAST STAGE: RX + final range + output_norm + lm head ----
3498        let _stl = rt.enter(n_st - 1);
3499        let el = rt.engine(n_st - 1, e);
3500        let pos_d = stage_pos(el)?;
3501        let rx_started = std::time::Instant::now();
3502        let x = rt.rx(n_st - 2, slot, payload)?;
3503        if pp_anatomy {
3504            el.stream().synchronize()?;
3505            rx_ms = rx_started.elapsed().as_secs_f64() * 1e3;
3506        }
3507        enqueue_spec_pipe_trace_marker(&el.stream(), trace.as_ref(), "S1", "start", Some(slot))?;
3508        let stage1_started = std::time::Instant::now();
3509        let x = self.verify_layers(
3510            el,
3511            x,
3512            fence[n_st - 1],
3513            fence[n_st],
3514            &pos_d,
3515            pos0,
3516            t,
3517            cache,
3518            ckpt.as_deref_mut(),
3519            stream,
3520        )?;
3521
3522        let mut hn = vbuf(el, payload)?;
3523        let logits = if self.cfg.step35.is_some() {
3524            // The PP Step35 serving path uses rms_norm + matmul for B=1 as well as B>1.
3525            // Verify must not switch numeric class merely because the same session speculates.
3526            el.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
3527            el.matmul(&self.output, &hn, t)?
3528        } else {
3529            el.rms_norm_decode(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
3530            el.matmul_decode_exact(&self.output, &hn, t)?
3531        };
3532        enqueue_spec_pipe_trace_marker(&el.stream(), trace.as_ref(), "S1", "end", Some(slot))?;
3533        if pp_anatomy {
3534            el.stream().synchronize()?;
3535            stage1_ms = stage1_started.elapsed().as_secs_f64() * 1e3;
3536        }
3537        // EXIT PUBLICATION: both returned buffers are still being produced on the last stage's
3538        // stream. Order the caller's stream behind that work before the buffers escape this
3539        // scope (the 2026-08-06 same-device ppspec find: without it the caller's primary-stream
3540        // consumer read unwritten logits — nondeterministic, one-device-only, and it poisoned
3541        // the following arm's KV in the same process).
3542        if publish_to_caller {
3543            rt.publish_to(n_st - 1, &caller_stream)?;
3544        }
3545        if pp_anatomy {
3546            if publish_to_caller {
3547                caller_stream.synchronize()?;
3548            }
3549            eprintln!(
3550                "[spec-pp-anatomy] t={t} reverse={reverse_ms:.3}ms stage0={stage0_ms:.3}ms \
3551                 tx={tx_ms:.3}ms rx={rx_ms:.3}ms stage1-head={stage1_ms:.3}ms total={:.3}ms",
3552                pp_started.elapsed().as_secs_f64() * 1e3,
3553            );
3554        }
3555        // stream: the device pos counter owns position; host mirror reconciles at drain.
3556        if stream.is_none() {
3557            cache.pos += t;
3558        }
3559        Ok((logits, if spec_hpost() { hn } else { x }))
3560    }
3561
3562    /// Step3.5/Step3.7 verify trunk in the serving batched numeric class.
3563    ///
3564    /// `step35_decode_batch_layers` is now authoritative at every live serving width, including
3565    /// B=1 (lane/cx-b1fix). The older verify walk deliberately mirrored the eager T=1 class:
3566    /// it replayed `step35_decode_attn` per row and used the eager/decode-exact FFN dispatch.
3567    /// Those classes are individually stable, but a near-tie prompt can choose different greedy
3568    /// bytes when a request moves from batched plain serving into speculative verify. Run the
3569    /// same authoritative B=1 stage subgraph for each verify row here. Rows still advance
3570    /// layer-by-layer, so every layer sees the preceding verify rows in its attention cache while
3571    /// every norm/projection/FFN uses exactly the live serving dispatch.
3572    #[allow(clippy::too_many_arguments)]
3573    fn step35_verify_batch_layers(
3574        &self,
3575        e: &Engine,
3576        mut x: CudaSlice<f32>,
3577        lo: usize,
3578        hi: usize,
3579        pos0: usize,
3580        t: usize,
3581        cache: &mut Cache,
3582    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3583        let n_embd = self.cfg.n_embd as usize;
3584        self.cfg
3585            .step35
3586            .as_ref()
3587            .ok_or("step35 verify batch requires step35 cfg")?;
3588        let mut ph_last = std::time::Instant::now();
3589        for il in lo..hi {
3590            let mut next = e.uninit(t * n_embd)?;
3591            for r in 0..t {
3592                let mut row = e.uninit(n_embd)?;
3593                e.dtod_copy_view(&x.slice(r * n_embd..(r + 1) * n_embd), &mut row)?;
3594                // The caller owns this verify's position. During controller overlap, cache.pos
3595                // still describes generation N while this stage-0 walk belongs to N+1.
3596                let row_pos = e.htod_i32(&[(pos0 + r) as i32])?;
3597                let mut one = [&mut *cache];
3598                let out = self.step35_decode_batch_layers(
3599                    e,
3600                    row,
3601                    &mut one,
3602                    &row_pos,
3603                    il,
3604                    il + 1,
3605                    &mut ph_last,
3606                )?;
3607                e.dtod_copy_into(&out, &mut next, r * n_embd)?;
3608            }
3609            x = next;
3610        }
3611        Ok(x)
3612    }
3613
3614    /// Qwen35-family verify trunk in the live serving numeric class.
3615    ///
3616    /// Serving intentionally keeps this architecture in the generic batched program even at
3617    /// B=1. The older verify walk used its own mirrored dispatch and can flip near-tie argmaxes.
3618    ///
3619    /// Two arms, one numeric class:
3620    /// - DENSE (`Arch::Qwen35`, t<=16): `qwen35_verify_tparallel` — the weight ops (norms,
3621    ///   projections, FFN) hoist to m=T through the exact-tier batched kernels whose per-row
3622    ///   program IS the m=1 program (`matmul_pre == fused2 per (tensor,row); _bN mmvq per-row
3623    ///   == m=1` — decode_batch.rs v2 note), while the state ops (conv ring, gdn scan, KV
3624    ///   append, fa decode) stay a per-row loop running the b_n=1 serving kernels with each
3625    ///   row's own t_kv-driven arm pick (the straddle law: every row executes the exact
3626    ///   program its isolated serving step would). One weight read per layer per round
3627    ///   instead of T — this is what makes MTP profitable in the exact class (the per-row
3628    ///   walk measured verify(K+1) ~= (K+1) plain steps: 69 -> 44 tok/s served, 2026-08-15).
3629    /// - MoE / t>16 / `MEMRA_SPEC_VERIFY_ROWWISE=1`: the per-row replay of the authoritative
3630    ///   serving layer body, preserving single-session autoregressive cache order (the
3631    ///   correctness reference; also the rollback seam for the t-parallel arm).
3632    ///
3633    /// Bit-identity of the t-parallel arm vs the rowwise arm is gated by spec-serve-gate
3634    /// (zero differing logits at T=1..4, K arms) + the 8-prompt ON/OFF canary before ship.
3635    #[allow(clippy::too_many_arguments)]
3636    fn qwen35_verify_batch_layers(
3637        &self,
3638        e: &Engine,
3639        x: CudaSlice<f32>,
3640        lo: usize,
3641        hi: usize,
3642        pos0: usize,
3643        t: usize,
3644        cache: &mut Cache,
3645        ckpt: Option<&mut VerifyCkpt>,
3646    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3647        let rowwise = std::env::var("MEMRA_SPEC_VERIFY_ROWWISE").as_deref() == Ok("1")
3648            || !matches!(self.cfg.arch, memra_gguf::config::Arch::Qwen35)
3649            || t > 16;
3650        if rowwise {
3651            self.qwen35_verify_rowwise(e, x, lo, hi, pos0, t, cache, ckpt)
3652        } else {
3653            self.qwen35_verify_tparallel(e, x, lo, hi, pos0, t, cache, ckpt)
3654        }
3655    }
3656
3657    /// The per-row correctness reference: replay each verify row through the authoritative
3658    /// serving layer body (`decode_batch_layers` at b_n=1). T full weight reads per layer.
3659    #[allow(clippy::too_many_arguments)]
3660    fn qwen35_verify_rowwise(
3661        &self,
3662        e: &Engine,
3663        mut x: CudaSlice<f32>,
3664        lo: usize,
3665        hi: usize,
3666        pos0: usize,
3667        t: usize,
3668        cache: &mut Cache,
3669        mut ckpt: Option<&mut VerifyCkpt>,
3670    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3671        let n_embd = self.cfg.n_embd as usize;
3672        let saved_pos = cache.pos;
3673        let mut ph_last = std::time::Instant::now();
3674        for il in lo..hi {
3675            let mut next = e.uninit(t * n_embd)?;
3676            let mut col_states: Option<Vec<(CudaSlice<f32>, CudaSlice<f32>)>> =
3677                if ckpt.is_some() && t >= 2 && matches!(self.layers[il].mixer, Mixer::Linear(_)) {
3678                    Some(Vec::with_capacity(t - 1))
3679                } else {
3680                    None
3681                };
3682            for r in 0..t {
3683                cache.pos = pos0 + r;
3684                let mut row = e.uninit(n_embd)?;
3685                e.dtod_copy_view(&x.slice(r * n_embd..(r + 1) * n_embd), &mut row)?;
3686                let row_pos = e.htod_i32(&[(pos0 + r) as i32])?;
3687                let mut one = [&mut *cache];
3688                let ctx = self.batch_layer_ctx(e, &one, il, il + 1)?;
3689                let out = match self.decode_batch_layers(
3690                    e,
3691                    row,
3692                    &mut one,
3693                    &ctx,
3694                    &row_pos,
3695                    &mut ph_last,
3696                ) {
3697                    Ok(out) => out,
3698                    Err(error) => {
3699                        cache.pos = saved_pos;
3700                        return Err(error);
3701                    }
3702                };
3703                e.dtod_copy_into(&out, &mut next, r * n_embd)?;
3704                if r + 1 < t {
3705                    if let Some(states) = col_states.as_mut() {
3706                        let recur = cache.recur[il]
3707                            .as_ref()
3708                            .ok_or("Qwen35-MoE linear verify layer has no recurrent state")?;
3709                        states.push((
3710                            e.clone_dtod(&recur.conv_state)?,
3711                            e.clone_dtod(&recur.ssm_state)?,
3712                        ));
3713                    }
3714                }
3715            }
3716            if let (Some(checkpoint), Some(states)) = (ckpt.as_deref_mut(), col_states) {
3717                checkpoint.cols[il] = Some(states);
3718            }
3719            x = next;
3720        }
3721        cache.pos = saved_pos;
3722        Ok(x)
3723    }
3724
3725    /// T-PARALLEL VERIFY IN THE SERVING NUMERIC CLASS (lane/tparallel-verify, 2026-08-15).
3726    ///
3727    /// The weight ops run ONCE per layer at m=T; the state ops run per row through the same
3728    /// b_n=1 serving kernels the rowwise replay uses. Per-row bit-identity rests on the two
3729    /// pins the serving batch tier already carries:
3730    ///   * `matmul_pre` / `_bN` mmvq: per-row program == m=1 program (decode_batch.rs v2 note,
3731    ///     kernel-check pinned) — so a [T, n_embd] projection row equals the row projected
3732    ///     alone;
3733    ///   * row-indexed norms/elementwise (`rms_norm`, `quantize_q8_1`, `add_rms_norm`,
3734    ///     `gated_rmsnorm[_q8_1]`, `silu_mul`, `rope_neox` with per-row positions): the T-row
3735    ///     launch is the per-row program (same pin the generic verify's fused norms rely on).
3736    /// The sequential dependencies keep their exact serving order: the conv ring / gdn scan
3737    /// chain state row -> row through the `_b` kernels at b_n=1 (ping-pong via a 6-entry
3738    /// alternating pointer table, host handles swapped per row so VerifyCkpt clones the
3739    /// canonical state exactly as the rowwise arm does), and each row's KV append + fa decode
3740    /// picks its arm from ITS OWN t_kv (append: format-only; fa: `fa_seqs_eligible` + its own
3741    /// `fa_split_keys` rung at b_n=1) — the straddle law per row, so every row executes the
3742    /// program its isolated B=1 serving step would.
3743    ///
3744    /// Cost: 1 weight read per layer per round + T state micro-launches, vs the rowwise arm's
3745    /// T weight reads. Gated bit-identical vs the rowwise arm by spec-serve-gate + canary.
3746    #[allow(clippy::too_many_arguments)]
3747    fn qwen35_verify_tparallel(
3748        &self,
3749        e: &Engine,
3750        mut x: CudaSlice<f32>,
3751        lo: usize,
3752        hi: usize,
3753        pos0: usize,
3754        t: usize,
3755        cache: &mut Cache,
3756        mut ckpt: Option<&mut VerifyCkpt>,
3757    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3758        use cudarc::driver::DevicePtr;
3759        let cfg = &self.cfg;
3760        let n_embd = cfg.n_embd as usize;
3761        let eps = cfg.rms_eps;
3762        let head_dim_global = cfg.head_dim_k as usize;
3763        let pos_host: Vec<i32> = (0..t).map(|r| (pos0 + r) as i32).collect();
3764        let pos_d = e.htod_i32(&pos_host)?;
3765        // Per-row 1-element position buffers, built ONCE per verify (the append/fa wrappers
3766        // take owned pos slices; building these inside the layer x row loops cost 16xT H2Ds).
3767        let pos_rows: Vec<CudaSlice<i32>> = (0..t)
3768            .map(|r| e.htod_i32(&[(pos0 + r) as i32]))
3769            .collect::<Result<_, _>>()?;
3770        let seqs_append =
3771            std::env::var("MEMRA_BATCH_APPEND").as_deref() != Ok("0") && !Engine::kv_fp8_on();
3772        let batch_fa_on = std::env::var("MEMRA_BATCH_FA").as_deref() != Ok("0");
3773
3774        for il in lo..hi {
3775            let layer = &self.layers[il];
3776            // ---- attn_norm + q8_1 quantize at m=T (row-indexed == per-row) ----
3777            let anorm = layer.attn_norm.float_data();
3778            let mut xn = e.uninit(t * n_embd)?;
3779            e.rms_norm(&x, anorm, &mut xn, n_embd, t, eps)?;
3780            let (hq, hd) = e.quantize_q8_1(&xn, t, n_embd)?;
3781
3782            let mixed: CudaSlice<f32> = match &layer.mixer {
3783                Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
3784                Mixer::Full(fa) => {
3785                    let geometry = cfg.full_attention_geometry_at(il as u32);
3786                    let n_head = geometry.n_head as usize;
3787                    let n_head_kv = geometry.n_head_kv as usize;
3788                    let head_dim = geometry.head_dim_k as usize;
3789                    let rope_dims = geometry.n_rot as usize;
3790                    let rope_base = geometry.rope_base;
3791                    let scale = geometry.attention_scale();
3792                    // Batched projections: one weight read serves all T rows.
3793                    let qf = e.matmul_pre(&fa.wq, &hq, &hd, &xn, t)?;
3794                    let mut k = e.matmul_pre(&fa.wk, &hq, &hd, &xn, t)?;
3795                    let v = e.matmul_pre(&fa.wv, &hq, &hd, &xn, t)?;
3796                    let gated =
3797                        geometry.attention_gate == memra_gguf::config::AttentionGateKind::FusedQ;
3798                    let (mut q, gate) = if gated {
3799                        let mut qs = e.uninit(t * n_head * head_dim)?;
3800                        let mut gs = e.uninit(t * n_head * head_dim)?;
3801                        e.q_gate_split(&qf, &mut qs, &mut gs, head_dim, n_head, t)?;
3802                        (qs, Some(gs))
3803                    } else {
3804                        (qf, None)
3805                    };
3806                    let mut qn = e.uninit(t * n_head * head_dim)?;
3807                    e.rms_norm(
3808                        &q,
3809                        fa.q_norm.float_data(),
3810                        &mut qn,
3811                        head_dim,
3812                        t * n_head,
3813                        eps,
3814                    )?;
3815                    q = qn;
3816                    let mut kn = e.uninit(t * n_head_kv * head_dim)?;
3817                    e.rms_norm(
3818                        &k,
3819                        fa.k_norm.float_data(),
3820                        &mut kn,
3821                        head_dim,
3822                        t * n_head_kv,
3823                        eps,
3824                    )?;
3825                    k = kn;
3826                    e.rope_neox(
3827                        &mut q, &pos_d, head_dim, rope_dims, n_head, t, rope_base, 1.0,
3828                    )?;
3829                    e.rope_neox(
3830                        &mut k, &pos_d, head_dim, rope_dims, n_head_kv, t, rope_base, 1.0,
3831                    )?;
3832
3833                    // Per-row append + attend: row r sees rows 0..r in KV (causal within the
3834                    // draft), each through the b_n=1 serving kernels at its own t_kv.
3835                    let q_dim = n_head * head_dim;
3836                    let kv_dim = n_head_kv * head_dim;
3837                    let mut attn = e.uninit(t * q_dim)?;
3838                    let (kdk, kdv, ktb, vtb, kv_view) = {
3839                        let kvl = cache.kv[il].as_ref().unwrap();
3840                        let s = &e.gpu.stream();
3841                        let (pk, _g) = kvl.k.device_ptr(s);
3842                        let (pv, _g2) = kvl.v.device_ptr(s);
3843                        (
3844                            kvl.kv_dim_k,
3845                            kvl.kv_dim_v,
3846                            kvl.k_tok_bytes,
3847                            kvl.v_tok_bytes,
3848                            e.htod_u64(&[pk as u64, pv as u64])?,
3849                        )
3850                    };
3851                    for r in 0..t {
3852                        // Owned per-row scratch: the b_n=1 kernels take packed batch buffers
3853                        // whose row 0 is this row (arithmetic-free materialization copies,
3854                        // same as decode's per-seq fallback arm).
3855                        let mut k_row = e.uninit(kv_dim)?;
3856                        e.dtod_copy_view(&k.slice(r * kv_dim..(r + 1) * kv_dim), &mut k_row)?;
3857                        let mut v_row = e.uninit(kv_dim)?;
3858                        e.dtod_copy_view(&v.slice(r * kv_dim..(r + 1) * kv_dim), &mut v_row)?;
3859                        let pos_row = &pos_rows[r];
3860                        let kvl = cache.kv[il].as_mut().unwrap();
3861                        if seqs_append {
3862                            e.append_kv_quantized_seqs(
3863                                &k_row,
3864                                &v_row,
3865                                &kv_view.slice(0..2),
3866                                pos_row,
3867                                1,
3868                                kdk,
3869                                kdv,
3870                                ktb,
3871                                vtb,
3872                            )?;
3873                            kvl.len += 1;
3874                        } else {
3875                            e.append_kv_quantized_view(
3876                                &k_row.slice(0..kv_dim),
3877                                &v_row.slice(0..kv_dim),
3878                                &mut kvl.k,
3879                                &mut kvl.v,
3880                                kvl.len,
3881                                kvl.kv_dim_k,
3882                                kvl.kv_dim_v,
3883                                kvl.k_tok_bytes,
3884                                kvl.v_tok_bytes,
3885                                Engine::kv_fp8_on(),
3886                            )?;
3887                            kvl.len += 1;
3888                        }
3889                        let t_kv = kvl.len;
3890                        let mut q_row = e.uninit(q_dim)?;
3891                        e.dtod_copy_view(&q.slice(r * q_dim..(r + 1) * q_dim), &mut q_row)?;
3892                        let mut a_row = e.uninit(q_dim)?;
3893                        if batch_fa_on && crate::fa_seqs_eligible(t_kv, head_dim_global) {
3894                            let sp0_r = crate::fa_split_keys(t_kv, cfg.n_head_kv as usize);
3895                            e.fa_decode_batch_seqs_v4(
3896                                &q_row,
3897                                &kv_view.slice(0..2),
3898                                pos_row,
3899                                &mut a_row,
3900                                head_dim,
3901                                n_head,
3902                                n_head_kv,
3903                                1,
3904                                t_kv,
3905                                scale,
3906                                sp0_r,
3907                                ktb,
3908                                vtb,
3909                            )?;
3910                        } else {
3911                            let k_view = e.view_u8(&kvl.k, t_kv * kvl.k_tok_bytes);
3912                            let v_view = e.view_u8(&kvl.v, t_kv * kvl.v_tok_bytes);
3913                            let mut a_view = a_row.slice_mut(0..q_dim);
3914                            e.fa_decode_kvmod_view(
3915                                &q_row.slice(0..q_dim),
3916                                &k_view,
3917                                &v_view,
3918                                &mut a_view,
3919                                head_dim,
3920                                n_head,
3921                                n_head_kv,
3922                                t_kv,
3923                                scale,
3924                                kvl.k_tok_bytes,
3925                                kvl.v_tok_bytes,
3926                                Engine::kv_fp8_on(),
3927                            )?;
3928                        }
3929                        e.dtod_copy_into(&a_row, &mut attn, r * q_dim)?;
3930                    }
3931
3932                    // Output gate (element-wise) + o-proj at m=T.
3933                    let attn_g = match &gate {
3934                        Some(g) => {
3935                            let n = t * q_dim;
3936                            let mut gsig = e.uninit(n)?;
3937                            e.sigmoid(g, &mut gsig, n)?;
3938                            let mut ag = e.uninit(n)?;
3939                            e.mul(&attn, &gsig, &mut ag, n)?;
3940                            ag
3941                        }
3942                        None => attn,
3943                    };
3944                    e.matmul(&fa.wo, &attn_g, t)?
3945                }
3946                Mixer::Linear(la) => {
3947                    let ssm = cfg.ssm.as_ref().expect("linear mixer requires ssm cfg");
3948                    let d_state = ssm.state_size as usize;
3949                    let num_k = ssm.group_count as usize;
3950                    let num_v = ssm.time_step_rank as usize;
3951                    let d_conv = ssm.conv_kernel as usize;
3952                    let key_dim = d_state * num_k;
3953                    let value_dim = d_state * num_v;
3954                    let conv_dim = key_dim * 2 + value_dim;
3955                    let gdn_scale = 1.0 / (d_state as f32).sqrt();
3956
3957                    // ---- batched projections: one weight read for all T rows ----
3958                    let qkv_mixed = e.matmul_pre(&la.wqkv, &hq, &hd, &xn, t)?;
3959                    let z = e.matmul_pre(&la.wqkv_gate, &hq, &hd, &xn, t)?;
3960                    let beta_raw = e.matmul_pre(&la.ssm_beta, &hq, &hd, &xn, t)?;
3961                    let alpha = e.matmul_pre(&la.ssm_alpha, &hq, &hd, &xn, t)?;
3962                    let beta_w = la.ssm_beta.out_features();
3963                    let alpha_w = la.ssm_alpha.out_features();
3964                    let qkv_w = la.wqkv.out_features();
3965
3966                    // ---- per-row state chain through the b_n=1 serving kernels ----
3967                    // 6-entry alternating pointer table expresses the ping-pong without a
3968                    // rebuild per row: even rows scan s0 -> s1, odd rows s1 -> s0. Host
3969                    // handles swap per row so ckpt clones the canonical state (and the
3970                    // post-verify canonical handle matches the last write), exactly as the
3971                    // rowwise arm leaves them.
3972                    let table = {
3973                        let rl = cache.recur[il].as_ref().unwrap();
3974                        let s = &e.gpu.stream();
3975                        let (pc, _g0) = rl.conv_state.device_ptr(s);
3976                        let (p0, _g1) = rl.ssm_state.device_ptr(s);
3977                        let (p1, _g2) = rl.ssm_state_alt.device_ptr(s);
3978                        e.htod_u64(&[
3979                            pc as u64, p0 as u64, p1 as u64, pc as u64, p1 as u64, p0 as u64,
3980                        ])?
3981                    };
3982                    let mut o_all = e.uninit(t * value_dim)?;
3983                    let mut col_states: Option<Vec<(CudaSlice<f32>, CudaSlice<f32>)>> =
3984                        if ckpt.is_some() && t >= 2 {
3985                            Some(Vec::with_capacity(t - 1))
3986                        } else {
3987                            None
3988                        };
3989                    // Per-row scratch reused across rows (uninit is cheap but not free at
3990                    // 48 layers x T rows); row inputs/outputs pass as VIEWS into the packed
3991                    // [T, ...] buffers — zero arithmetic-free copies in this loop.
3992                    let mut conv_out = e.uninit(conv_dim)?;
3993                    let mut q_l2 = e.uninit(value_dim)?;
3994                    let mut k_l2 = e.uninit(value_dim)?;
3995                    let mut v_gd = e.uninit(value_dim)?;
3996                    let mut beta_b = e.uninit(num_v)?;
3997                    let mut g_log = e.uninit(num_v)?;
3998                    for r in 0..t {
3999                        let base = if r % 2 == 0 { 0 } else { 3 };
4000                        let conv_view = table.slice(base..base + 1);
4001                        let in_view = table.slice(base + 1..base + 2);
4002                        let out_view = table.slice(base + 2..base + 3);
4003                        e.ssm_conv1d_fused_decode_b_view(
4004                            &qkv_mixed.slice(r * qkv_w..(r + 1) * qkv_w),
4005                            &conv_view,
4006                            la.ssm_conv1d.float_data(),
4007                            &mut conv_out,
4008                            conv_dim,
4009                            d_conv,
4010                            1,
4011                        )?;
4012                        e.gdn_prep_decode_b_view(
4013                            &conv_out,
4014                            &beta_raw.slice(r * beta_w..(r + 1) * beta_w),
4015                            &alpha.slice(r * alpha_w..(r + 1) * alpha_w),
4016                            la.ssm_dt.float_data(),
4017                            la.ssm_a.float_data(),
4018                            &mut q_l2,
4019                            &mut k_l2,
4020                            &mut v_gd,
4021                            &mut beta_b,
4022                            &mut g_log,
4023                            d_state,
4024                            num_v,
4025                            num_k,
4026                            key_dim,
4027                            eps,
4028                            conv_dim,
4029                            1,
4030                        )?;
4031                        let mut o_row = o_all.slice_mut(r * value_dim..(r + 1) * value_dim);
4032                        e.gdn_scan_s128_batched_view(
4033                            &q_l2, &k_l2, &v_gd, &g_log, &beta_b, &in_view, &out_view, &mut o_row,
4034                            num_v, 1, gdn_scale,
4035                        )?;
4036                        {
4037                            let rl = cache.recur[il].as_mut().unwrap();
4038                            std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
4039                        }
4040                        if r + 1 < t {
4041                            if let Some(states) = col_states.as_mut() {
4042                                let recur = cache.recur[il]
4043                                    .as_ref()
4044                                    .ok_or("qwen35 linear verify layer has no recurrent state")?;
4045                                states.push((
4046                                    e.clone_dtod(&recur.conv_state)?,
4047                                    e.clone_dtod(&recur.ssm_state)?,
4048                                ));
4049                            }
4050                        }
4051                    }
4052                    if let (Some(checkpoint), Some(states)) = (ckpt.as_deref_mut(), col_states) {
4053                        checkpoint.cols[il] = Some(states);
4054                    }
4055
4056                    // ---- batched gated norm + out-projection at m=T ----
4057                    if e.uses_q8_1_fast(&la.ssm_out) {
4058                        let (gq, gd) = e.gated_rmsnorm_q8_1(
4059                            &o_all,
4060                            la.ssm_norm.float_data(),
4061                            &z,
4062                            d_state,
4063                            t * num_v,
4064                            eps,
4065                        )?;
4066                        let g0 = e.zeros(0)?;
4067                        e.matmul_pre(&la.ssm_out, &gq, &gd, &g0, t)?
4068                    } else {
4069                        let mut gn = e.uninit(t * value_dim)?;
4070                        e.gated_rmsnorm(
4071                            &o_all,
4072                            la.ssm_norm.float_data(),
4073                            &z,
4074                            &mut gn,
4075                            d_state,
4076                            t * num_v,
4077                            eps,
4078                        )?;
4079                        e.matmul(&la.ssm_out, &gn, t)?
4080                    }
4081                }
4082            };
4083
4084            // ---- residual add + post_attn_norm + FFN at m=T (serving dispatch verbatim) ----
4085            let pnorm = layer.post_attn_norm.float_data();
4086            let mut x1 = e.uninit(t * n_embd)?;
4087            let mut zn = e.uninit(t * n_embd)?;
4088            e.add_rms_norm(&x, &mixed, pnorm, &mut x1, &mut zn, n_embd, t, eps)?;
4089            let ffn_out = match &layer.ffn {
4090                crate::hybrid::Ffn::Dense {
4091                    ffn_gate,
4092                    ffn_up,
4093                    ffn_down,
4094                } => {
4095                    assert!(
4096                        self.cfg.m3.is_none(),
4097                        "qwen35 t-parallel verify: M3 swigluoai FFN not yet batched"
4098                    );
4099                    let n_ff = ffn_gate.out_features();
4100                    let (zq, zd) = e.quantize_q8_1(&zn, t, n_embd)?;
4101                    let g = e.matmul_pre(ffn_gate, &zq, &zd, &zn, t)?;
4102                    let u = e.matmul_pre(ffn_up, &zq, &zd, &zn, t)?;
4103                    let mut act = e.uninit(t * n_ff)?;
4104                    e.silu_mul(&g, &u, &mut act, t * n_ff)?;
4105                    let (aq, ad) = e.quantize_q8_1(&act, t, n_ff)?;
4106                    e.matmul_pre(ffn_down, &aq, &ad, &act, t)?
4107                }
4108                crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il_zq8(e, m, &zn, None, t, il as u16)?,
4109            };
4110            let mut x2 = e.uninit(t * n_embd)?;
4111            e.add(&x1, &ffn_out, &mut x2, t * n_embd)?;
4112            x = x2;
4113        }
4114        Ok(x)
4115    }
4116
4117    /// PP-N STAGE SUBGRAPH of the verify trunk: layers `[lo, hi)` of `decode_step_t_core_stream`'s
4118    /// walk, verbatim. Enters with a MATERIALIZED `[T, n_embd]` residual (no pending fusion pair
4119    /// carried in from outside the range) and exits with the range's final residual materialized
4120    /// (the trailing add executed) — exactly the `decode_layers_eager(lo, hi)` contract, T rows
4121    /// instead of one.
4122    ///
4123    /// EXTRACTED (lane/pp2-spec 2026-08-06) rather than duplicated: `decode_step_t_core_stream` IS
4124    /// the single funnel every verify forward reaches, and its per-layer dispatch MIRRORING (norm
4125    /// fusion per layer, the t>=3/spec_m2 batched-linear window, the fused-q8 FFN chain, the
4126    /// decode-exact projections) is what makes verify bit-identical to eager decode. A second copy
4127    /// for the split arm is how those mirrors drift apart on the next lever. The unsplit body now
4128    /// calls this with `(0, n_layers)`, so the whole-trunk path and every stage range run the SAME
4129    /// code — there is no "split version" of the verify math.
4130    ///
4131    /// Bit-identity of a cut rests on the same kernel-check-pinned identity the eager arm's cut
4132    /// does — `add_rms_norm_q8_1 == add then rms_norm_q8_1` at nrows=T — because the ONLY thing a
4133    /// fence changes is that the cross-layer fusion carry breaks at `hi-1` and is re-materialized
4134    /// as an explicit `add`. `decode-batch-gate --mode ppspec` verifies end-to-end on real weights.
4135    #[allow(clippy::too_many_arguments)]
4136    fn verify_layers(
4137        &self,
4138        e: &Engine,
4139        mut x: CudaSlice<f32>,
4140        lo: usize,
4141        hi: usize,
4142        pos_d: &CudaSlice<i32>,
4143        pos0: usize,
4144        t: usize,
4145        cache: &mut Cache,
4146        mut ckpt: Option<&mut VerifyCkpt>,
4147        stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
4148    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4149        if self.cfg.step35.is_some() {
4150            if stream.is_some() {
4151                return Err(
4152                    "step35 has no ROUND-STREAM verify arm (the device-counter _dc twins \
4153                            cannot express the SWA offset KV view)"
4154                        .into(),
4155                );
4156            }
4157            return self.step35_verify_batch_layers(e, x, lo, hi, pos0, t, cache);
4158        }
4159        if self.qwen35_serving_class() {
4160            if stream.is_some() {
4161                return Err("qwen35-family serving-class verify has no ROUND-STREAM arm".into());
4162            }
4163            return self.qwen35_verify_batch_layers(e, x, lo, hi, pos0, t, cache, ckpt.take());
4164        }
4165        let n_embd = self.cfg.n_embd as usize;
4166        let eps = self.cfg.rms_eps;
4167        // CROSS-LAYER ADD+NORM FUSION (lane/vt-fixes fix 2, mirroring decode_step_h's
4168        // launch-arc form): layer il's post-FFN residual add (x2 = x1 + ffn_out) and layer
4169        // il+1's attn_norm(+quantize) are consecutive row-wise ops — ONE add_rms_norm_q8_1
4170        // launch at nrows=t does all three (bit-identity pinned by the T-row kernel-check
4171        // arms). Carry the un-added (x1, ffn_out) pair; the fused launch materializes x2 (the
4172        // residual the next layer needs) as its `res` output. Falls back to the separate add
4173        // when the next layer is off the fused-q8 path.
4174        let mut pending: Option<(CudaSlice<f32>, CudaSlice<f32>)> = None;
4175        for il in lo..hi {
4176            let layer = &self.layers[il];
4177            // DISPATCH-MIRRORED attn-input RMSNorm (FP-order lesson #8): eager decode fuses the
4178            // 1024-thread rms_norm_q8_1 ONLY when every mixer projection is q8_1-fast; layers with
4179            // Float projections (ssm_beta/ssm_alpha on layers 1/2/4 of the 9B NVFP4 GGUF) take the
4180            // UNFUSED 256-thread rms_norm. The verify norm must mirror that PER-LAYER choice —
4181            // blockDim changes the sum-of-squares reduce order, and the ULP shift amplifies through
4182            // the GDN recurrence into argmax flips (measured: 9B text prompt, 1 ULP at layer 2 ->
4183            // 2.3e-1 logit maxdiff at the head -> K=1..8 divergence at a 0.03-margin token).
4184            let mixer_fast = self.mixer_in_q8_1_fast(e, &layer.mixer);
4185            let norm_fused = std::env::var("MEMRA_NO_FUSE_NORMQ").is_err() && mixer_fast;
4186            // BATCHED EPILOGUE RE-FUSE (lane/vt-fixes fix 2, 2026-08-03): when the norm is
4187            // dispatch-fused AND every consumer of `h` reads only its q8_1 form (Full mixer:
4188            // projections only; Linear mixer: the batched arm — the per-column fallback needs
4189            // f32 h), emit the attn-input norm DIRECTLY as q8_1 via `rms_norm_q8_1` at nrows=t
4190            // (row-indexed kernel — the T-row launch is the per-row m=1 program, kernel-check
4191            // pins bit-identity vs rms_norm_decode -> quantize_q8_1). Kills the standalone
4192            // quantize launch(es) + the f32 h HBM round-trip that decode never pays.
4193            // step35 (Full mixer) is the third case that needs f32 `h`: its verify arm is a
4194            // per-ROW replay of the eager decode mixer, whose `pre_q` contract is a single row —
4195            // a T-row q8_1 pair cannot be handed to it, and re-deriving per-row q8_1 from the f32
4196            // rows is exactly the dispatch being mirrored. Keep step35 on the unfused arm.
4197            let lin_q8_only = match &layer.mixer {
4198                Mixer::Linear(la) => {
4199                    (t >= 3 || (t == 2 && spec_m2())) && e.uses_q8_1_fast(&la.ssm_out)
4200                }
4201                Mixer::Full(_) if self.cfg.step35.is_some() => false,
4202                _ => true,
4203            };
4204            // NOTE decode.rs's take()-first lesson: take the pending pair BEFORE branching so
4205            // a non-fused layer still performs the residual add.
4206            let taken = pending.take();
4207            let (h, h_q8) = if norm_fused && lin_q8_only {
4208                let pair = match taken {
4209                    // fused add + attn_norm + q8_1: ONE launch resolves the carried residual
4210                    // AND emits this layer's mixer input pre-quantized. res -> x2 (= new x).
4211                    Some((x1p, f1p)) => {
4212                        let mut x2 = vbuf(e, t * n_embd)?; // fully written (res output)
4213                        let p = e.add_rms_norm_q8_1(
4214                            &x1p,
4215                            &f1p,
4216                            layer.attn_norm.float_data(),
4217                            &mut x2,
4218                            n_embd,
4219                            t,
4220                            eps,
4221                        )?;
4222                        x = x2;
4223                        p
4224                    }
4225                    None => e.rms_norm_q8_1(&x, layer.attn_norm.float_data(), n_embd, t, eps)?,
4226                };
4227                (e.zeros(0)?, Some(pair)) // h unused on this path (q8-only consumers)
4228            } else {
4229                if let Some((x1p, f1p)) = taken {
4230                    let mut x2 = vbuf(e, t * n_embd)?; // fully written by add
4231                    e.add(&x1p, &f1p, &mut x2, t * n_embd)?;
4232                    x = x2;
4233                }
4234                let mut h = vbuf(e, t * n_embd)?; // fully written by either rms_norm arm
4235                if norm_fused {
4236                    e.rms_norm_decode(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
4237                } else {
4238                    e.rms_norm(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
4239                }
4240                (h, None)
4241            };
4242            let h_q8_ref = h_q8.as_ref().map(|(q, d)| (q, d));
4243
4244            let mixed = match &layer.mixer {
4245                Mixer::Full(fa) => self.full_attn_verify(
4246                    e,
4247                    fa,
4248                    &h,
4249                    h_q8_ref,
4250                    pos_d,
4251                    t,
4252                    cache,
4253                    il,
4254                    stream.map(|(_, c)| c),
4255                )?,
4256                Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
4257                Mixer::Linear(la) => {
4258                    // BATCHED linear verify (2026-07-03, the MTP-profit lever): one T-token pass —
4259                    // batched projections (weight read ONCE, hits the m=2-4 weight-resident matvec),
4260                    // carried-state conv (ssm_conv1d_tm_state), GDN prep on the prefill kernels, and
4261                    // ONE gdn_scan whose internal sequential t-loop is the SAME recurrence as T
4262                    // chained T=1 steps (bit-identical). Falls back to the sequential per-column
4263                    // chain when T < d_conv-1 (conv ring update needs T >= pad) — or when ANY
4264                    // projection is off the q8_1 fast path: matmul_decode_exact would route a Float
4265                    // tensor to cuBLAS at m=t (different FP accumulation than eager's per-token
4266                    // GEMV), so mixed-dtype layers stay on the eager-identical per-column chain.
4267                    // MEMRA_SPEC_M2 (lane/spec-m2): the t==2 batch rides the same arm — the conv
4268                    // wrapper handles t<pad with a pure-copy ring rebuild; see spec_m2() header.
4269                    if (t >= 3 || (t == 2 && spec_m2()))
4270                        && mixer_fast
4271                        && e.uses_q8_1_fast(&la.ssm_out)
4272                    {
4273                        let want = ckpt.is_some();
4274                        let (out, stash) =
4275                            self.linear_attn_verify_t(e, la, &h, h_q8_ref, t, cache, il, want)?;
4276                        if let (Some(ck), Some(st)) = (ckpt.as_deref_mut(), stash) {
4277                            ck.gdn[il] = Some(st);
4278                        }
4279                        out
4280                    } else {
4281                        let mut out = vbuf(e, t * n_embd)?; // every col written by copy_into
4282                        let mut col_states: Option<Vec<(CudaSlice<f32>, CudaSlice<f32>)>> =
4283                            if ckpt.is_some() && t >= 2 {
4284                                Some(Vec::with_capacity(t - 1))
4285                            } else {
4286                                None
4287                            };
4288                        for col in 0..t {
4289                            let mut h_col = vbuf(e, n_embd)?; // fully written by copy_view_into
4290                            let src = h.slice(col * n_embd..(col + 1) * n_embd);
4291                            e.copy_view_into(&mut h_col, 0, &src, n_embd)?;
4292                            let m_col = self.linear_attn_decode(e, la, &h_col, cache, il)?;
4293                            e.copy_into(&mut out, col * n_embd, &m_col, n_embd)?;
4294                            // REPLAY-FREE ckpt: clone the chain's ACTUAL state after this column
4295                            // (pure dtod — cannot change any computed value). Last column skipped:
4296                            // rebuild targets are j <= t-1 columns.
4297                            if let Some(cs) = col_states.as_mut() {
4298                                if col + 1 < t {
4299                                    let rl = cache.recur[il].as_ref().unwrap();
4300                                    cs.push((
4301                                        e.clone_dtod(&rl.conv_state)?,
4302                                        e.clone_dtod(&rl.ssm_state)?,
4303                                    ));
4304                                }
4305                            }
4306                        }
4307                        if let (Some(ck), Some(cs)) = (ckpt.as_deref_mut(), col_states) {
4308                            // ReplaySSM-assessment instrumentation (2026-07-30): the
4309                            // per-column clones are the only true state snapshots left in
4310                            // the verify (the batched path stashes INPUTS and replays).
4311                            if std::env::var("MEMRA_SPEC_STATS").as_deref() == Ok("1") {
4312                                static ONCE: std::sync::Once = std::sync::Once::new();
4313                                let bytes: usize =
4314                                    cs.iter().map(|(c, s)| (c.len() + s.len()) * 4).sum();
4315                                ONCE.call_once(|| eprintln!(
4316                                    "[verify-ckpt] per-column layer il={il}: {} clones, {:.2} MB/layer/round",
4317                                    cs.len(), bytes as f64 / 1e6));
4318                            }
4319                            ck.cols[il] = Some(cs);
4320                        }
4321                        out
4322                    }
4323                }
4324            };
4325
4326            // DISPATCH-MIRRORED post-attn norm: eager residual_norm_ffn fuses add+norm+quant
4327            // (1024-thread add_rms_norm_q8_1) only for Dense FFNs whose gate+up are q8_1-fast;
4328            // otherwise (and for MoE) it runs the 256-thread fused add_rms_norm. Mirror per layer.
4329            let ffn_fuse = match &layer.ffn {
4330                crate::hybrid::Ffn::Dense {
4331                    ffn_gate, ffn_up, ..
4332                } => {
4333                    std::env::var("MEMRA_NO_FUSE_NORMQ").is_err()
4334                        && e.uses_q8_1_fast(ffn_gate)
4335                        && e.uses_q8_1_fast(ffn_up)
4336                }
4337                crate::hybrid::Ffn::Moe(_) => false,
4338            };
4339            // BATCHED EPILOGUE RE-FUSE (lane/vt-fixes fix 2): on the ffn_fuse path (Dense,
4340            // gate+up q8_1-fast, non-M3) the FFN input is emitted DIRECTLY as q8_1 by ONE
4341            // add_rms_norm_q8_1 launch at nrows=t (row-indexed kernel: the T-row launch is the
4342            // per-row m=1 program; kernel-check pins bit-identity vs the unfused
4343            // add_f32 -> rms_norm_decode -> quantize_q8_1 chain at T=2/4/5/8) — replacing the
4344            // add + rms_norm_decode launches AND the dual/singles' internal re-quantize.
4345            // M3's swigluoai must keep the f32 chain (the fused SwiGLU epilogue encodes plain
4346            // SiLU), mirroring residual_norm_ffn's m3 guard on the decode path.
4347            // step35: same guard per LAYER. A dense FFN's clamp is the SHEXP array (upstream's
4348            // one build_ffn serves dense + shared expert, llama-graph.cpp:1751), and verify MUST
4349            // mirror decode's dispatch or spec self-consistency fails.
4350            let dense_lim = self.cfg.clamp_shexp_at(il as u32);
4351            let fuse_q8 = ffn_fuse && self.cfg.m3.is_none() && dense_lim.is_none();
4352            let mut x1 = vbuf(e, t * n_embd)?; // fully written by add / add_rms_norm*
4353            let mut z = e.zeros(0)?; // replaced below on the unfused arms
4354            let z_q8 = if fuse_q8 {
4355                Some(e.add_rms_norm_q8_1(
4356                    &x,
4357                    &mixed,
4358                    layer.post_attn_norm.float_data(),
4359                    &mut x1,
4360                    n_embd,
4361                    t,
4362                    eps,
4363                )?)
4364            } else {
4365                let mut zf = vbuf(e, t * n_embd)?; // fully written by rms_norm_decode / add_rms_norm
4366                if ffn_fuse {
4367                    e.add(&x, &mixed, &mut x1, t * n_embd)?;
4368                    e.rms_norm_decode(
4369                        &x1,
4370                        layer.post_attn_norm.float_data(),
4371                        &mut zf,
4372                        n_embd,
4373                        t,
4374                        eps,
4375                    )?;
4376                } else {
4377                    e.add_rms_norm(
4378                        &x,
4379                        &mixed,
4380                        layer.post_attn_norm.float_data(),
4381                        &mut x1,
4382                        &mut zf,
4383                        n_embd,
4384                        t,
4385                        eps,
4386                    )?;
4387                }
4388                z = zf;
4389                None
4390            };
4391            // DECODE-EXACT FFN projections: force MMVQ for gate/up/down at any T to match the
4392            // T=1 decode FP accumulation order. At T>=5 the generic matmul/matmul_pre falls to dp4a
4393            // (128-thread, different FP sum order). At T=2-4 the batched MMVQ is already bit-identical.
4394            let ffn_out = match &layer.ffn {
4395                crate::hybrid::Ffn::Dense {
4396                    ffn_gate,
4397                    ffn_up,
4398                    ffn_down,
4399                } => {
4400                    let n_ff = ffn_gate.out_features();
4401                    if let Some((zq, zd)) = z_q8.as_ref() {
4402                        // FUSED CHAIN (fix 2): pre-quantized z feeds the projections; the SwiGLU
4403                        // epilogue emits act pre-quantized for ffn_down (silu_mul_scaled_q8_1,
4404                        // bit-identical to silu_mul + quantize — kernel-check-pinned) with the
4405                        // NVFP4 macro-scales folded (deferred-scale dual: y*s inline == the
4406                        // scale_inplace store, value-exact) — the exact m=1 decode epilogue
4407                        // structure at nrows=t.
4408                        let pair =
4409                            match e.matmul_decode_exact_dual_pre(ffn_gate, ffn_up, zq, zd, t)? {
4410                                Some(((g, gs), (u, us))) => Some((g, gs, u, us)),
4411                                None => None,
4412                            };
4413                        let (gate, gs, up, us) = match pair {
4414                            Some(x4) => x4,
4415                            None => (
4416                                e.matmul_decode_exact_pre(ffn_gate, zq, zd, t)?,
4417                                1.0, // scale already applied inside _pre
4418                                e.matmul_decode_exact_pre(ffn_up, zq, zd, t)?,
4419                                1.0,
4420                            ),
4421                        };
4422                        if e.uses_q8_1_fast(ffn_down) {
4423                            let (aq, ad) = e.silu_mul_scaled_q8_1(&gate, &up, gs, us, t * n_ff)?;
4424                            e.matmul_decode_exact_pre(ffn_down, &aq, &ad, t)?
4425                        } else {
4426                            let mut act = vbuf(e, t * n_ff)?;
4427                            e.silu_mul_scaled(&gate, &up, gs, us, &mut act, t * n_ff)?;
4428                            e.matmul_decode_exact(ffn_down, &act, t)?
4429                        }
4430                    } else {
4431                        // UNFUSED (pre-fix) chain — MoE-adjacent/M3/off-fast layers, unchanged.
4432                        // DUAL gate+up batched twin (lane/verify-economics, 2026-08-02): one launch
4433                        // for the pair at t=2..8 — bit-identical per (tensor,token,row) to the two
4434                        // singles (kernel-check pins bitwise; MEMRA_SPEC_DUAL_T=0 reverts). None
4435                        // (non-NVFP4 / t outside the tier / seam off) -> the two singles, unchanged.
4436                        let (gate, up) =
4437                            match e.matmul_decode_exact_dual(ffn_gate, ffn_up, &z, t)? {
4438                                Some(pair) => pair,
4439                                None => (
4440                                    e.matmul_decode_exact(ffn_gate, &z, t)?,
4441                                    e.matmul_decode_exact(ffn_up, &z, t)?,
4442                                ),
4443                            };
4444                        let mut act = vbuf(e, t * n_ff)?; // fully written by ffn_act_lim
4445                        Self::ffn_act_lim(
4446                            e,
4447                            &self.cfg,
4448                            &gate,
4449                            &up,
4450                            1.0,
4451                            1.0,
4452                            dense_lim,
4453                            &mut act,
4454                            t * n_ff,
4455                        )?;
4456                        e.matmul_decode_exact(ffn_down, &act, t)?
4457                    }
4458                }
4459                crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il(e, m, &z, t, il as u16)?,
4460            };
4461            // CROSS-LAYER fusion: defer this layer's post-FFN residual add — the next layer's
4462            // fused-q8 attn norm folds it in (add_rms_norm_q8_1 == add; rms_norm; quantize,
4463            // kernel-check-pinned at nrows=T). Non-fused next layers add explicitly above.
4464            pending = Some((x1, ffn_out));
4465        }
4466        // RANGE's final add (no next norm INSIDE the range to fuse with; for the
4467        // whole-trunk call that is the last layer, whose next norm is output_norm — f32-out).
4468        if let Some((x1p, f1p)) = pending.take() {
4469            let mut x2 = vbuf(e, t * n_embd)?; // fully written by add
4470            e.add(&x1p, &f1p, &mut x2, t * n_embd)?;
4471            x = x2;
4472        }
4473        Ok(x)
4474    }
4475    /// BATCHED linear-attn verify (T=K+1): the whole layer in ~10 launches instead of T x the
4476    /// T=1 decode chain (T x ~12 launches + T weight reads of the four projections). The GDN
4477    /// recurrence itself is inherently sequential — gdn_scan_s128 runs its internal t-loop with
4478    /// the SAME per-token math as chained T=1 calls (bit-identical state evolution); everything
4479    /// around it (projections, conv, prep, gated norm, out-proj) batches. Advances conv ring +
4480    /// ssm state exactly like T sequential decode steps.
4481    /// `want_stash`: additionally RETAIN the gdn-scan inputs (pure buffer keep-alives, zero extra
4482    /// kernels) so a partial accept can rebuild the state after any column prefix (REPLAY-FREE).
4483    #[allow(clippy::too_many_arguments)]
4484    fn linear_attn_verify_t(
4485        &self,
4486        e: &Engine,
4487        la: &LinearAttnLayer,
4488        h: &CudaSlice<f32>,
4489        h_q8: Option<(&CudaSlice<i8>, &CudaSlice<f32>)>,
4490        t: usize,
4491        cache: &mut Cache,
4492        il: usize,
4493        want_stash: bool,
4494    ) -> Result<(CudaSlice<f32>, Option<GdnStash>), Box<dyn std::error::Error>> {
4495        let cfg = &self.cfg;
4496        let ssm = cfg.ssm.as_ref().unwrap();
4497        let d_state = ssm.state_size as usize;
4498        let num_k = ssm.group_count as usize;
4499        let num_v = ssm.time_step_rank as usize;
4500        let d_conv = ssm.conv_kernel as usize;
4501        let key_dim = d_state * num_k;
4502        let conv_dim = key_dim * 2 + d_state * num_v;
4503        let eps = cfg.rms_eps;
4504        let scale = 1.0 / (d_state as f32).sqrt();
4505
4506        // DECODE-EXACT projections: matmul_decode_exact forces the MMVQ (warp-per-row, 32-thread)
4507        // accumulation order for EVERY m, matching the T=1 decode path bit-for-bit. The generic
4508        // `matmul` at m>=5 falls to dp4a (128-thread, two-level reduce) which has a different FP
4509        // sum order — ULP differences propagate through gdn_scan and flip argmax on the 27B.
4510        // Q8 TRUNK-FUSION at T=1 (35B: wqkv+wqkv_gate both Q8_0): one fused2 launch, bit-identical
4511        // per (tensor,row) to the two m=1 MMVQ dispatches below — decode-exact contract holds.
4512        // VERIFY-TIER TRUNK FUSION (MEMRA_SPEC_FUSED_T, t=2-4): quantize h ONCE for every
4513        // fused-eligible same-input Q8_0 pair of this layer (35B wqkv+wqkv_gate; 9B
4514        // ssm_beta+ssm_alpha) — each fused2 batched launch then replaces two decode-exact
4515        // calls (each of which re-quantizes the same h + runs its own _b2/_b4 launch).
4516        // Bit-identical per (tensor,token,row) — see spec_fused_t().
4517        // BATCHED EPILOGUE RE-FUSE (lane/vt-fixes fix 2): `h_q8` = the attn-input norm emitted
4518        // directly as q8_1 by the caller's fused rms_norm_q8_1 (bit-identical to the unfused
4519        // chain, kernel-check-pinned). When present it REPLACES the standalone quantize below
4520        // and feeds every projection; the caller guaranteed all four input projections are
4521        // q8_1-fast. When absent, the old shared-quantize (fused-t window) stands.
4522        let h_q8_t = if h_q8.is_none()
4523            && spec_fused_t()
4524            && (2..=4).contains(&t)
4525            && ((e.uses_q8_1_fast(&la.wqkv) && e.uses_q8_1_fast(&la.wqkv_gate))
4526                || (e.uses_q8_1_fast(&la.ssm_beta) && e.uses_q8_1_fast(&la.ssm_alpha)))
4527        {
4528            Some(e.quantize_q8_1(h, t, cfg.n_embd as usize)?)
4529        } else {
4530            None
4531        };
4532        // one view: the caller's fused-norm q8 or this fn's own shared quantize.
4533        let hq8_any: Option<(&CudaSlice<i8>, &CudaSlice<f32>)> =
4534            h_q8.or(h_q8_t.as_ref().map(|(q, d)| (q, d)));
4535        let (qkv_mixed, z) = {
4536            let mut fused = None;
4537            if t == 1 && e.uses_q8_1_fast(&la.wqkv) && e.uses_q8_1_fast(&la.wqkv_gate) {
4538                let (hq, hd) = e.quantize_q8_1(h, 1, cfg.n_embd as usize)?;
4539                fused = e.matmul_q8_fused2(&la.wqkv, &la.wqkv_gate, &hq, &hd)?;
4540            } else if let Some((hq, hd)) = hq8_any {
4541                if spec_fused_t() && (2..=4).contains(&t) {
4542                    fused = e.matmul_q8_fused2_t(&la.wqkv, &la.wqkv_gate, hq, hd, t)?;
4543                }
4544            }
4545            match (fused, hq8_any) {
4546                (Some(pair), _) => pair,
4547                (None, Some((hq, hd))) if h_q8.is_some() => (
4548                    e.matmul_decode_exact_pre(&la.wqkv, hq, hd, t)?,
4549                    e.matmul_decode_exact_pre(&la.wqkv_gate, hq, hd, t)?,
4550                ),
4551                (None, _) => (
4552                    e.matmul_decode_exact(&la.wqkv, h, t)?,
4553                    e.matmul_decode_exact(&la.wqkv_gate, h, t)?,
4554                ),
4555            }
4556        };
4557        // beta+alpha DUAL at T=1 (75% of p3 rounds run T=1 verify — p-min chain cuts): the dual
4558        // mr2 kernel is bit-identical per element to the m=1 MMVQ matmul_decode_exact dispatches
4559        // (same warp-per-row body, blockIdx.y picks the weight), so the decode-exact contract
4560        // holds; the run-spec battery is the arbiter. T>1 keeps the per-tensor decode-exact path.
4561        let (beta_raw, alpha) = if t == 1 {
4562            let (hq, hd) = e.quantize_q8_1(h, 1, cfg.n_embd as usize)?;
4563            match e.matmul_pre_dual_noscale(&la.ssm_beta, &la.ssm_alpha, &hq, &hd, 1)? {
4564                Some(((mut b, bs), (mut a, as_))) => {
4565                    if bs != 1.0 {
4566                        e.scale_inplace(&mut b, bs, la.ssm_beta.out_features())?;
4567                    }
4568                    if as_ != 1.0 {
4569                        e.scale_inplace(&mut a, as_, la.ssm_alpha.out_features())?;
4570                    }
4571                    (b, a)
4572                }
4573                // Q8_0 fused2 twin (9B stores beta/alpha as Q8_0): DISPATCH-MIRRORS the eager
4574                // decode's beta_alpha closure — the fused body is qmatvec_q8_0_mmvq verbatim,
4575                // bit-identical per row (kernel-check rel=0.00e0 gate), so decode==verify holds.
4576                None => match e.matmul_q8_fused2(&la.ssm_beta, &la.ssm_alpha, &hq, &hd)? {
4577                    Some((b, a)) => (b, a),
4578                    None => (
4579                        e.matmul_decode_exact(&la.ssm_beta, h, 1)?,
4580                        e.matmul_decode_exact(&la.ssm_alpha, h, 1)?,
4581                    ),
4582                },
4583            }
4584        } else {
4585            // fused-t twin (9B stores beta/alpha as Q8_0): same shared-quantize + one launch
4586            // contract as the wqkv pair above; 35B beta/alpha are Float -> None -> fallback.
4587            let mut nvfp4_fused = None;
4588            let mut q8_fused = None;
4589            if let Some((hq, hd)) = hq8_any {
4590                if t == 3 && std::env::var("MEMRA_NVFP4_AUX_DUAL").as_deref() != Ok("0") {
4591                    nvfp4_fused =
4592                        e.matmul_decode_exact_dual_pre(&la.ssm_beta, &la.ssm_alpha, hq, hd, t)?;
4593                    if nvfp4_fused.is_some() && std::env::var("MEMRA_DEBUG").is_ok() {
4594                        static ONCE: std::sync::Once = std::sync::Once::new();
4595                        ONCE.call_once(|| {
4596                            eprintln!("[memra] NVFP4 beta+alpha batched aux dual ENGAGED (t={t})")
4597                        });
4598                    }
4599                }
4600                if nvfp4_fused.is_none() && spec_fused_t() && (2..=4).contains(&t) {
4601                    q8_fused = e.matmul_q8_fused2_t(&la.ssm_beta, &la.ssm_alpha, hq, hd, t)?;
4602                }
4603            }
4604            if let Some(((mut b, bs), (mut a, as_))) = nvfp4_fused {
4605                if bs != 1.0 {
4606                    e.scale_inplace(&mut b, bs, t * la.ssm_beta.out_features())?;
4607                }
4608                if as_ != 1.0 {
4609                    e.scale_inplace(&mut a, as_, t * la.ssm_alpha.out_features())?;
4610                }
4611                (b, a)
4612            } else if let Some(pair) = q8_fused {
4613                pair
4614            } else {
4615                match hq8_any {
4616                    Some((hq, hd)) if h_q8.is_some() => (
4617                        e.matmul_decode_exact_pre(&la.ssm_beta, hq, hd, t)?,
4618                        e.matmul_decode_exact_pre(&la.ssm_alpha, hq, hd, t)?,
4619                    ),
4620                    _ => (
4621                        e.matmul_decode_exact(&la.ssm_beta, h, t)?,
4622                        e.matmul_decode_exact(&la.ssm_alpha, h, t)?,
4623                    ),
4624                }
4625            }
4626        };
4627
4628        // conv with CARRIED state + ring roll (T >= pad rides the input-column update kernel;
4629        // T < pad — the MEMRA_SPEC_M2 t=2 arm — rolls via the pure-copy ring rebuild).
4630        let rl = cache.recur[il].as_mut().unwrap();
4631        let mut conv_out = e.uninit(conv_dim * t)?;
4632        e.ssm_conv1d_tm_state(
4633            &qkv_mixed,
4634            &mut rl.conv_state,
4635            la.ssm_conv1d.float_data(),
4636            &mut conv_out,
4637            conv_dim,
4638            t,
4639            d_conv,
4640        )?;
4641
4642        // GDN prep via the prefill kernels (repack + L2 + sigmoid + glog), T-wide.
4643        let mut q_g = e.uninit(d_state * num_v * t)?;
4644        let mut k_g = e.uninit(d_state * num_v * t)?;
4645        let mut v_g = e.uninit(d_state * num_v * t)?;
4646        e.qkv_to_gdn_repack(
4647            &conv_out, &mut q_g, &mut k_g, &mut v_g, d_state, num_v, num_k, key_dim, t,
4648        )?;
4649        let mut q_l2 = e.uninit(d_state * num_v * t)?;
4650        e.l2_norm_decode(&q_g, &mut q_l2, d_state, num_v * t, eps)?;
4651        let mut k_l2 = e.uninit(d_state * num_v * t)?;
4652        e.l2_norm_decode(&k_g, &mut k_l2, d_state, num_v * t, eps)?;
4653        let mut beta = e.uninit(t * num_v)?;
4654        e.sigmoid(&beta_raw, &mut beta, t * num_v)?;
4655        let mut g_log = e.uninit(t * num_v)?;
4656        e.gdn_glog(
4657            &alpha,
4658            la.ssm_dt.float_data(),
4659            la.ssm_a.float_data(),
4660            &mut g_log,
4661            num_v,
4662            t,
4663        )?;
4664
4665        // ONE gdn_scan over T tokens from the carried state (internal sequential loop ==
4666        // T chained T=1 steps). Ping-pong the resident buffers like eager decode.
4667        let mut o = e.uninit(d_state * num_v * t)?;
4668        {
4669            let crate::cache::RecurLayer {
4670                ssm_state,
4671                ssm_state_alt,
4672                ..
4673            } = rl;
4674            e.gdn_scan_s128(
4675                &q_l2,
4676                &k_l2,
4677                &v_g,
4678                &g_log,
4679                &beta,
4680                ssm_state,
4681                ssm_state_alt,
4682                &mut o,
4683                num_v,
4684                t,
4685                scale,
4686            )?;
4687        }
4688        std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
4689
4690        // gated RMSNorm + out projection, T-wide. FUSED-QUANTIZE ARM (lane/vt-fixes fix 2,
4691        // mirroring the T=1 decode's launch-arc form): when ssm_out rides the q8_1 fast path,
4692        // emit q8_1 straight from the gated norm at nrows=num_v*t (row-indexed kernel, the
4693        // T-wide launch is the per-row program; kernel-check pins bit-identity vs
4694        // gated_rmsnorm -> quantize_q8_1 at T=1 and T=5) and feed the decode-exact dispatch
4695        // pre-quantized — one launch replaces norm + quantize. Fallback = the f32 chain.
4696        let out = if e.uses_q8_1_fast(&la.ssm_out) {
4697            let (gq, gd) =
4698                e.gated_rmsnorm_q8_1(&o, la.ssm_norm.float_data(), &z, d_state, num_v * t, eps)?;
4699            e.matmul_decode_exact_pre(&la.ssm_out, &gq, &gd, t)?
4700        } else {
4701            let mut gn = e.uninit(d_state * num_v * t)?;
4702            e.gated_rmsnorm(
4703                &o,
4704                la.ssm_norm.float_data(),
4705                &z,
4706                &mut gn,
4707                d_state,
4708                num_v * t,
4709                eps,
4710            )?;
4711            // DECODE-EXACT out-projection: same MMVQ path as the T=1 decode (ssm_out at m>=5
4712            // would fall to dp4a with a different FP reduction order — same class of bug as
4713            // the input projs).
4714            e.matmul_decode_exact(&la.ssm_out, &gn, t)?
4715        };
4716        let stash = if want_stash {
4717            Some(GdnStash {
4718                qkv_mixed,
4719                q_l2,
4720                k_l2,
4721                v_g,
4722                g_log,
4723                beta,
4724            })
4725        } else {
4726            None
4727        };
4728        Ok((out, stash))
4729    }
4730
4731    /// REPLAY-FREE partial-accept commit (2026-07-03): make the cache state == "committed through
4732    /// the first `j` verify columns" WITHOUT the legacy rollback + duplicate trunk replay.
4733    /// - Full-attn KV: truncate len to snapshot + j. The verify's appended rows for those columns
4734    ///   are bit-identical to what an eager T=1 chain writes (the decode-exact contract the
4735    ///   verify-probe gates), so keeping them == replaying them.
4736    /// - Linear layers, batched path: rebuild the conv ring by PURE COPIES (ring holds raw input
4737    ///   columns) and the ssm state by a prefix re-run of the SAME gdn_scan kernel (t=j) from the
4738    ///   snapshot state over the stash's identical inputs — the kernel's t-loop carries state in
4739    ///   registers and writes it once at the end, so iterations 0..j-1 are independent of T:
4740    ///   bit-identical to the verify's own state after j tokens == the eager chain state.
4741    /// - Linear layers, per-column path: restore the cloned actual state after column j-1.
4742    /// Caller guarantees 1 <= j <= t-1 (j==0 rounds take the legacy rollback; j==t is full accept).
4743    fn commit_verified_prefix(
4744        &self,
4745        e: &Engine,
4746        cache: &mut Cache,
4747        snap: &crate::cache::CacheSnapshot,
4748        ckpt: &VerifyCkpt,
4749        j: usize,
4750        kv_lens_done: bool,
4751        dev_j: Option<(&CudaSlice<u32>, usize, usize)>,
4752    ) -> Result<(), Box<dyn std::error::Error>> {
4753        let cfg = &self.cfg;
4754        let ssm = cfg.ssm.as_ref().unwrap();
4755        let d_state = ssm.state_size as usize;
4756        let num_k = ssm.group_count as usize;
4757        let num_v = ssm.time_step_rank as usize;
4758        let d_conv = ssm.conv_kernel as usize;
4759        let conv_dim = d_state * num_k * 2 + d_state * num_v;
4760        let scale = 1.0 / (d_state as f32).sqrt();
4761        for il in 0..self.layers.len() {
4762            if let (Some(kvl), Some(saved)) = (cache.kv[il].as_mut(), snap.kv_len[il]) {
4763                kvl.len = saved + j;
4764                // devacc 3a: spec_rollback_kv already wrote len_d on-device (same value).
4765                if !kv_lens_done {
4766                    e.set_i32_one(&mut kvl.len_d, kvl.len as i32)?;
4767                }
4768            }
4769            if let Some(rl) = cache.recur[il].as_mut() {
4770                if let Some(st) = &ckpt.gdn[il] {
4771                    let ring_old = snap.conv[il].as_ref().expect("snapshot missing conv");
4772                    let state_in = snap.ssm[il].as_ref().expect("snapshot missing ssm");
4773                    if let Some((acc, base, t_v)) = dev_j {
4774                        // 3b: j read on-device (_dc twins, same bodies; full accept early-exits).
4775                        e.ssm_conv_ring_rebuild_dc(
4776                            &st.qkv_mixed,
4777                            ring_old,
4778                            &mut rl.conv_state,
4779                            conv_dim,
4780                            acc,
4781                            base,
4782                            t_v,
4783                            d_conv,
4784                        )?;
4785                        let mut o = e.uninit(d_state * num_v * j.max(1))?;
4786                        e.gdn_scan_s128_dc(
4787                            &st.q_l2,
4788                            &st.k_l2,
4789                            &st.v_g,
4790                            &st.g_log,
4791                            &st.beta,
4792                            state_in,
4793                            &mut rl.ssm_state,
4794                            &mut o,
4795                            num_v,
4796                            acc,
4797                            base,
4798                            t_v,
4799                            scale,
4800                        )?;
4801                    } else {
4802                        e.ssm_conv_ring_rebuild(
4803                            &st.qkv_mixed,
4804                            ring_old,
4805                            &mut rl.conv_state,
4806                            conv_dim,
4807                            j,
4808                            d_conv,
4809                        )?;
4810                        let mut o = e.uninit(d_state * num_v * j)?; // scan output, discarded
4811                        e.gdn_scan_s128(
4812                            &st.q_l2,
4813                            &st.k_l2,
4814                            &st.v_g,
4815                            &st.g_log,
4816                            &st.beta,
4817                            state_in,
4818                            &mut rl.ssm_state,
4819                            &mut o,
4820                            num_v,
4821                            j,
4822                            scale,
4823                        )?;
4824                    }
4825                } else if let Some(cols) = &ckpt.cols[il] {
4826                    let (c, s) = &cols[j - 1];
4827                    e.copy_into(&mut rl.conv_state, 0, c, c.len())?;
4828                    e.copy_into(&mut rl.ssm_state, 0, s, s.len())?;
4829                } else {
4830                    return Err(
4831                        "commit_verified_prefix: verify ckpt missing for linear layer".into(),
4832                    );
4833                }
4834            }
4835        }
4836        cache.pos = snap.pos + j;
4837        Ok(())
4838    }
4839
4840    /// ROUND-STREAM: recur restore with device-j (the _dc twins; full accept early-exits
4841    /// in-kernel). Requires the batched-linear stash on every linear layer (stream gate).
4842    fn commit_verified_prefix_stream(
4843        &self,
4844        e: &Engine,
4845        cache: &mut Cache,
4846        snap: &crate::cache::CacheSnapshot,
4847        ckpt: &VerifyCkpt,
4848        acc: &CudaSlice<u32>,
4849        base: usize,
4850        t_v: usize,
4851    ) -> Result<(), Box<dyn std::error::Error>> {
4852        let cfg = &self.cfg;
4853        let ssm = cfg.ssm.as_ref().unwrap();
4854        let d_state = ssm.state_size as usize;
4855        let num_k = ssm.group_count as usize;
4856        let num_v = ssm.time_step_rank as usize;
4857        let d_conv = ssm.conv_kernel as usize;
4858        let conv_dim = d_state * num_k * 2 + d_state * num_v;
4859        let scale = 1.0 / (d_state as f32).sqrt();
4860        for il in 0..self.layers.len() {
4861            if let Some(rl) = cache.recur[il].as_mut() {
4862                let st = ckpt.gdn[il]
4863                    .as_ref()
4864                    .ok_or("stream restore: batched-linear stash missing")?;
4865                let ring_old = snap.conv[il].as_ref().expect("snapshot missing conv");
4866                let state_in = snap.ssm[il].as_ref().expect("snapshot missing ssm");
4867                e.ssm_conv_ring_rebuild_dc(
4868                    &st.qkv_mixed,
4869                    ring_old,
4870                    &mut rl.conv_state,
4871                    conv_dim,
4872                    acc,
4873                    base,
4874                    t_v,
4875                    d_conv,
4876                )?;
4877                let mut o = e.uninit(d_state * num_v * t_v)?;
4878                e.gdn_scan_s128_dc(
4879                    &st.q_l2,
4880                    &st.k_l2,
4881                    &st.v_g,
4882                    &st.g_log,
4883                    &st.beta,
4884                    state_in,
4885                    &mut rl.ssm_state,
4886                    &mut o,
4887                    num_v,
4888                    acc,
4889                    base,
4890                    t_v,
4891                    scale,
4892                )?;
4893            }
4894        }
4895        Ok(())
4896    }
4897
4898    /// EAGLE3 aux-capturing verify forward over `tokens` (T) — mirrors `decode_step_t_h` exactly
4899    /// (same KV append, same causal verify, same recur advance) but ALSO clones the aux residual-
4900    /// stream hiddens (blocks in `aux_layers`) for TWO columns: the LAST column (always) and the
4901    /// optional `pred_col` (the EAGLE seed = bonus's predecessor). Returns
4902    /// (all_T_logits host, last_col_aux, pred_col_aux?). Used by the EAGLE3 orchestrator's commit.
4903    pub fn decode_step_t_aux2(
4904        &self,
4905        e: &Engine,
4906        tokens: &[u32],
4907        pos0: usize,
4908        cache: &mut Cache,
4909        aux_layers: &[usize],
4910        pred_col: Option<usize>,
4911    ) -> Result<
4912        (Vec<f32>, Vec<CudaSlice<f32>>, Option<Vec<CudaSlice<f32>>>),
4913        Box<dyn std::error::Error>,
4914    > {
4915        let cfg = &self.cfg;
4916        let n_embd = cfg.n_embd as usize;
4917        let eps = cfg.rms_eps;
4918        let t = tokens.len();
4919        let pos_vec: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
4920        let pos_d = e.htod_i32(&pos_vec)?;
4921        let mut x = e.htod(&self.embd.gather(n_embd, tokens))?;
4922        let mut aux_last: Vec<CudaSlice<f32>> = Vec::with_capacity(aux_layers.len());
4923        let mut aux_pred: Vec<CudaSlice<f32>> = Vec::new();
4924        let want_pred = pred_col.is_some();
4925
4926        for (il, layer) in self.layers.iter().enumerate() {
4927            // DISPATCH-MIRRORED norms (FP-order lesson #8) — see decode_step_t_h_emb.
4928            let mixer_fast = self.mixer_in_q8_1_fast(e, &layer.mixer);
4929            let norm_fused = std::env::var("MEMRA_NO_FUSE_NORMQ").is_err() && mixer_fast;
4930            let mut h = vbuf(e, t * n_embd)?; // fully written by either rms_norm arm
4931            if norm_fused {
4932                e.rms_norm_decode(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
4933            } else {
4934                e.rms_norm(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
4935            }
4936            let mixed = match &layer.mixer {
4937                Mixer::Full(fa) => {
4938                    self.full_attn_verify(e, fa, &h, None, &pos_d, t, cache, il, None)?
4939                }
4940                Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
4941                Mixer::Linear(la) => {
4942                    let mut out = e.zeros(t * n_embd)?;
4943                    for col in 0..t {
4944                        let mut h_col = e.zeros(n_embd)?;
4945                        let src = h.slice(col * n_embd..(col + 1) * n_embd);
4946                        e.copy_view_into(&mut h_col, 0, &src, n_embd)?;
4947                        let m_col = self.linear_attn_decode(e, la, &h_col, cache, il)?;
4948                        e.copy_into(&mut out, col * n_embd, &m_col, n_embd)?;
4949                    }
4950                    out
4951                }
4952            };
4953            let ffn_fuse = match &layer.ffn {
4954                crate::hybrid::Ffn::Dense {
4955                    ffn_gate, ffn_up, ..
4956                } => {
4957                    std::env::var("MEMRA_NO_FUSE_NORMQ").is_err()
4958                        && e.uses_q8_1_fast(ffn_gate)
4959                        && e.uses_q8_1_fast(ffn_up)
4960                }
4961                crate::hybrid::Ffn::Moe(_) => false,
4962            };
4963            let mut x1 = vbuf(e, t * n_embd)?; // fully written by add / add_rms_norm
4964            let mut z = vbuf(e, t * n_embd)?; // fully written by rms_norm_decode / add_rms_norm
4965            if ffn_fuse {
4966                e.add(&x, &mixed, &mut x1, t * n_embd)?;
4967                e.rms_norm_decode(
4968                    &x1,
4969                    layer.post_attn_norm.float_data(),
4970                    &mut z,
4971                    n_embd,
4972                    t,
4973                    eps,
4974                )?;
4975            } else {
4976                e.add_rms_norm(
4977                    &x,
4978                    &mixed,
4979                    layer.post_attn_norm.float_data(),
4980                    &mut x1,
4981                    &mut z,
4982                    n_embd,
4983                    t,
4984                    eps,
4985                )?;
4986            }
4987            let ffn_out = match &layer.ffn {
4988                crate::hybrid::Ffn::Dense {
4989                    ffn_gate,
4990                    ffn_up,
4991                    ffn_down,
4992                } => {
4993                    let n_ff = ffn_gate.out_features();
4994                    let gate = e.matmul_decode_exact(ffn_gate, &z, t)?;
4995                    let up = e.matmul_decode_exact(ffn_up, &z, t)?;
4996                    let mut act = vbuf(e, t * n_ff)?; // fully written by ffn_act_lim
4997                    // dense FFN clamp = the SHEXP array (upstream build_ffn serves both).
4998                    Self::ffn_act_lim(
4999                        e,
5000                        &self.cfg,
5001                        &gate,
5002                        &up,
5003                        1.0,
5004                        1.0,
5005                        self.cfg.clamp_shexp_at(il as u32),
5006                        &mut act,
5007                        t * n_ff,
5008                    )?;
5009                    e.matmul_decode_exact(ffn_down, &act, t)?
5010                }
5011                crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il(e, m, &z, t, il as u16)?,
5012            };
5013            let mut x2 = vbuf(e, t * n_embd)?; // fully written by add
5014            e.add(&x1, &ffn_out, &mut x2, t * n_embd)?;
5015            if aux_layers.contains(&il) {
5016                let mut a = e.zeros(n_embd)?;
5017                e.copy_view_into(&mut a, 0, &x2.slice((t - 1) * n_embd..t * n_embd), n_embd)?;
5018                aux_last.push(a);
5019                if let Some(pc) = pred_col {
5020                    let mut ap = e.zeros(n_embd)?;
5021                    e.copy_view_into(
5022                        &mut ap,
5023                        0,
5024                        &x2.slice(pc * n_embd..(pc + 1) * n_embd),
5025                        n_embd,
5026                    )?;
5027                    aux_pred.push(ap);
5028                }
5029            }
5030            x = x2;
5031        }
5032        let mut hn = vbuf(e, t * n_embd)?; // fully written by rms_norm_decode
5033        e.rms_norm_decode(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
5034        let logits = e.matmul_decode_exact(&self.output, &hn, t)?;
5035        let host = e.dtoh(&logits)?;
5036        cache.pos += t;
5037        Ok((
5038            host,
5039            aux_last,
5040            if want_pred { Some(aux_pred) } else { None },
5041        ))
5042    }
5043
5044    /// step35 SPEC-VERIFY attention over T query tokens — a per-row REPLAY of the eager
5045    /// `step35_decode_attn`.
5046    ///
5047    /// WHY A REPLAY AND NOT A BATCHED TWIN. The verify's whole job is to be bit-identical to what
5048    /// the eager decode would have computed for the same tokens; that is what makes greedy spec
5049    /// decode exact (run-spec asserts token identity for K=1..8). Every other verify arm in this
5050    /// file earns that identity by carefully mirroring dispatch (`matmul_decode_exact` to force
5051    /// MMVQ at any m, per-layer `ffn_fuse` mirroring, per-row `fa_decode` key bounds). step35
5052    /// stacks FOUR more per-layer degrees of freedom on top of that — per-layer `n_head`
5053    /// (64 full / 96 SWA), per-layer rotary width (64 full / 128 SWA), per-layer rope base, and a
5054    /// SEPARATE `attn_gate` tensor whose projection shares the attn-normed input — and its SWA
5055    /// layers attend through a token-OFFSET view whose offset is a function of the ABSOLUTE
5056    /// position of each query row. A batched twin would have to reproduce all of that AND the
5057    /// per-row offset in one launch; the offset alone rules out the existing rows kernels (they
5058    /// take one `base_len`, not a per-row offset).
5059    ///
5060    /// So this arm calls the eager path itself, once per row, on the same cache. Identity is then
5061    /// true BY CONSTRUCTION rather than by mirroring: row r runs exactly the kernel sequence that
5062    /// eager decode step r runs (same projections, same q8_1 fusion decision, same append, same
5063    /// view arithmetic, same `fa_decode_kvmod`, same gate), because it IS that code. Cost: T x the
5064    /// eager decode mixer instead of one batched pass — the same trade the generic arm's `else`
5065    /// per-row loop already accepts when `fa_rows_eligible` says no. Correctness first; a batched
5066    /// step35 twin is a perf lane's job and must be gated against this arm.
5067    ///
5068    /// The `h_q8` pre-quantized pair from the caller's fused norm is NOT forwarded: it is a
5069    /// T-row buffer and `step35_decode_attn`'s `pre_q` contract is one row. Instead each row's
5070    /// f32 `h` slice is handed over and the callee re-derives its own q8_1 exactly as eager decode
5071    /// does (`quantize_q8_1(h, 1, n_embd)`) — which is the dispatch being mirrored. Callers that
5072    /// took the fused arm therefore MUST still pass a live `h`; `step35_verify` asserts that.
5073    #[allow(clippy::too_many_arguments)]
5074    fn step35_verify(
5075        &self,
5076        e: &Engine,
5077        fa: &FullAttnLayer,
5078        h: &CudaSlice<f32>,
5079        h_q8: Option<(&CudaSlice<i8>, &CudaSlice<f32>)>,
5080        t: usize,
5081        cache: &mut Cache,
5082        il: usize,
5083    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5084        let n_embd = self.cfg.n_embd as usize;
5085        // The fused (h-less) attn-norm arm hands `h` as a zero-length placeholder. This arm needs
5086        // the f32 rows, so the caller must not take that lever for step35 — enforced at the call
5087        // site by the `Mixer::Full(_) if self.cfg.step35.is_some() => false` arm of
5088        // `lin_q8_only` in `decode_step_t_core_stream`, and asserted here so a future caller
5089        // cannot regress it into silently reading an empty buffer.
5090        assert_eq!(
5091            h.len(),
5092            t * n_embd,
5093            "step35_verify needs the f32 attn-normed rows ([t*n_embd]); the caller took the \
5094             fused q8-only norm arm (h_q8={}) — step35 must stay on the unfused arm",
5095            h_q8.is_some()
5096        );
5097        // ROW WIDTH IS n_embd, NOT n_head*head_dim: `step35_decode_attn` returns the mixer output
5098        // AFTER `wo`, so a row is [n_embd] — the same contract the generic arm's
5099        // `matmul_decode_exact(&fa.wo, &attn_g, t)` return has. Sizing this buffer from the
5100        // per-layer head geometry (8192 on full-attn, 12288 on SWA) instead overran the row on the
5101        // FIRST copy and panicked inside `copy_into`'s `CudaView::slice` unwrap
5102        // (raw/mtp-bt-20260806T212127Z.log frames 12-13).
5103        let mut out = vbuf(e, t * n_embd)?; // each row fully written by the copy below
5104        for r in 0..t {
5105            // Absolute position of this query row. `cache.pos` is the committed length at round
5106            // start and every row before r has already been appended by this loop, so the r-th
5107            // verify token sits at cache.pos + r — the same position eager decode would give it.
5108            let pos_d = e.htod_i32(&[(cache.pos + r) as i32])?;
5109            let mut h_row = vbuf(e, n_embd)?; // fully written by copy_view_into
5110            e.copy_view_into(
5111                &mut h_row,
5112                0,
5113                &h.slice(r * n_embd..(r + 1) * n_embd),
5114                n_embd,
5115            )?;
5116            // THE eager decode mixer: appends this row's K/V at kvl.len, advances it, then
5117            // attends over the (SWA-offset) view. Post-`wo`, same contract as this fn returns.
5118            let o = self.step35_decode_attn(e, fa, il, &h_row, None, &pos_d, cache)?;
5119            debug_assert_eq!(
5120                o.len(),
5121                n_embd,
5122                "step35_decode_attn returns post-wo [n_embd]"
5123            );
5124            e.copy_into(&mut out, r * n_embd, &o, n_embd)?;
5125        }
5126        Ok(out)
5127    }
5128
5129    /// Full-attention mixer over T query tokens with a GROWING resident KV (verify path, §D.3).
5130    /// Appends the T new K/V columns to cache.kv[il] then attends causally over [0..len) via
5131    /// fa_prefill. Token-major [T, kv_dim] projection layout == cache row layout (single copy).
5132    #[allow(clippy::too_many_arguments)]
5133    fn full_attn_verify(
5134        &self,
5135        e: &Engine,
5136        fa: &FullAttnLayer,
5137        h: &CudaSlice<f32>,
5138        h_q8: Option<(&CudaSlice<i8>, &CudaSlice<f32>)>,
5139        pos_d: &CudaSlice<i32>,
5140        t: usize,
5141        cache: &mut Cache,
5142        il: usize,
5143        stream_ctr: Option<&CudaSlice<i32>>,
5144    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5145        // step35: the generic geometry below is wrong for this arch (per-layer n_head, partial
5146        // per-layer rope, the SWA offset view, and a SEPARATE head-wise gate tensor), so it takes
5147        // its own arm. A verify that silently computes different attention than decode defeats the
5148        // whole self-consistency gate, so the arm is a per-row REPLAY of `step35_decode_attn`
5149        // rather than a batched twin — see `step35_verify` for why that is the exactness-correct
5150        // shape and not laziness.
5151        if self.cfg.step35.is_some() {
5152            if stream_ctr.is_some() {
5153                return Err(
5154                    "step35 has no ROUND-STREAM verify arm (the device-counter _dc twins \
5155                            cannot express the SWA offset KV view; same root cause as the dc \
5156                            decode refusal) — run spec without the stream arm"
5157                        .into(),
5158                );
5159            }
5160            return self.step35_verify(e, fa, h, h_q8, t, cache, il);
5161        }
5162        let cfg = &self.cfg;
5163        let geometry = cfg.full_attention_geometry_at(il as u32);
5164        let n_head = geometry.n_head as usize;
5165        let n_head_kv = geometry.n_head_kv as usize;
5166        let head_dim = geometry.head_dim_k as usize;
5167        let eps = cfg.rms_eps;
5168        let scale = geometry.attention_scale();
5169        let n_embd = cfg.n_embd as usize;
5170
5171        // DECODE-EXACT Q/K/V projections: matmul_decode_exact forces the MMVQ (warp-per-row) path
5172        // for every m, matching the T=1 decode's FP accumulation order. matmul_pre at m>=5 would
5173        // fall to dp4a (128-thread, two-level reduce) with a different FP sum order.
5174        // Q8 TRUNK-FUSION at T=1: DISPATCH-MIRRORS the eager decode's fused3 (bit-identical body).
5175        // BATCHED EPILOGUE RE-FUSE (lane/vt-fixes fix 2): `h_q8` = the attn-input norm's q8_1
5176        // form emitted by the fused rms_norm_q8_1 (bit-identical to rms_norm_decode ->
5177        // quantize_q8_1, kernel-check-pinned). When present (caller checked mixer q8_1-fast),
5178        // every projection consumes it — `h` may be a zero-len placeholder and must not be read.
5179        let (qf, mut k, v) = {
5180            let mut fused = None;
5181            let qkv_fast =
5182                e.uses_q8_1_fast(&fa.wq) && e.uses_q8_1_fast(&fa.wk) && e.uses_q8_1_fast(&fa.wv);
5183            if t == 1 && qkv_fast {
5184                let (hq_o, hd_o);
5185                let (hq, hd): (&CudaSlice<i8>, &CudaSlice<f32>) = match h_q8 {
5186                    Some(p) => p,
5187                    None => {
5188                        (hq_o, hd_o) = e.quantize_q8_1(h, 1, n_embd)?;
5189                        (&hq_o, &hd_o)
5190                    }
5191                };
5192                fused = e.matmul_q8_fused3(&fa.wq, &fa.wk, &fa.wv, hq, hd)?;
5193            } else if spec_fused_t() && (2..=4).contains(&t) && qkv_fast {
5194                // VERIFY-TIER TRUNK FUSION (MEMRA_SPEC_FUSED_T): one shared quantize + one
5195                // fused3 batched launch replaces three decode-exact calls (3 re-quantizes of
5196                // the same h + 3 _b2/_b4 launches). Bit-identical per (tensor,token,row).
5197                let (hq_o, hd_o);
5198                let (hq, hd): (&CudaSlice<i8>, &CudaSlice<f32>) = match h_q8 {
5199                    Some(p) => p,
5200                    None => {
5201                        (hq_o, hd_o) = e.quantize_q8_1(h, t, n_embd)?;
5202                        (&hq_o, &hd_o)
5203                    }
5204                };
5205                fused = e.matmul_q8_fused3_t(&fa.wq, &fa.wk, &fa.wv, hq, hd, t)?;
5206            }
5207            match (fused, h_q8) {
5208                (Some(triple), _) => triple,
5209                // shared pre-quantized activation (q8_1-fast guaranteed by the caller): the
5210                // decode-exact dispatch consumes (hq, hd) instead of re-quantizing 3x.
5211                (None, Some((hq, hd))) if qkv_fast => (
5212                    e.matmul_decode_exact_pre(&fa.wq, hq, hd, t)?,
5213                    e.matmul_decode_exact_pre(&fa.wk, hq, hd, t)?,
5214                    e.matmul_decode_exact_pre(&fa.wv, hq, hd, t)?,
5215                ),
5216                (None, _) => (
5217                    e.matmul_decode_exact(&fa.wq, h, t)?,
5218                    e.matmul_decode_exact(&fa.wk, h, t)?,
5219                    e.matmul_decode_exact(&fa.wv, h, t)?,
5220                ),
5221            }
5222        };
5223        // M3/Hy3 have no attention output gate — wq out is exactly q; skip the split.
5224        let gated = geometry.attention_gate == memra_gguf::config::AttentionGateKind::FusedQ;
5225        let (mut q, gate) = if gated {
5226            let mut q = vbuf(e, t * n_head * head_dim)?; // fully written by q_gate_split
5227            let mut gate = vbuf(e, t * n_head * head_dim)?; // fully written by q_gate_split
5228            e.q_gate_split(&qf, &mut q, &mut gate, head_dim, n_head, t)?;
5229            (q, Some(gate))
5230        } else {
5231            (qf, None)
5232        };
5233
5234        let mut qn = vbuf(e, t * n_head * head_dim)?; // fully written by rms_norm
5235        e.rms_norm(
5236            &q,
5237            fa.q_norm.float_data(),
5238            &mut qn,
5239            head_dim,
5240            n_head * t,
5241            eps,
5242        )?;
5243        q = qn;
5244        let mut kn = vbuf(e, t * n_head_kv * head_dim)?; // fully written by rms_norm
5245        e.rms_norm(
5246            &k,
5247            fa.k_norm.float_data(),
5248            &mut kn,
5249            head_dim,
5250            n_head_kv * t,
5251            eps,
5252        )?;
5253        k = kn;
5254        let rope_dims = geometry.n_rot as usize;
5255        e.rope_neox(
5256            &mut q,
5257            pos_d,
5258            head_dim,
5259            rope_dims,
5260            n_head,
5261            t,
5262            geometry.rope_base,
5263            1.0,
5264        )?;
5265        e.rope_neox(
5266            &mut k,
5267            pos_d,
5268            head_dim,
5269            rope_dims,
5270            n_head_kv,
5271            t,
5272            geometry.rope_base,
5273            1.0,
5274        )?;
5275
5276        // append T new K/V columns to the resident QUANTIZED cache. k/v are token-major [T, kv_dim]
5277        // f32; append-quantize each of the T token rows into the byte cache (q8_0 K / q5_1 V).
5278        let kvl = cache.kv[il].as_mut().unwrap();
5279        let (kv_dim_k, kv_dim_v, ktb, vtb) =
5280            (kvl.kv_dim_k, kvl.kv_dim_v, kvl.k_tok_bytes, kvl.v_tok_bytes);
5281        if let Some(ctr) = stream_ctr {
5282            // stream: ONE batched append at the device counter (rows kernel = the per-view warp
5283            // math on a (block, token) grid, documented byte-identical); host len is a stale
5284            // LOWER BOUND under pre-issue (drain reconciles it).
5285            e.append_kv_quantized_rows_dc(
5286                &k,
5287                &v,
5288                &mut kvl.k,
5289                &mut kvl.v,
5290                ctr,
5291                t,
5292                kv_dim_k,
5293                kv_dim_v,
5294                ktb,
5295                vtb,
5296                crate::Engine::kv_fp8_on(),
5297            )?;
5298        } else {
5299            for i in 0..t {
5300                let k_row = k.slice(i * kv_dim_k..(i + 1) * kv_dim_k);
5301                let v_row = v.slice(i * kv_dim_v..(i + 1) * kv_dim_v);
5302                e.append_kv_quantized_view(
5303                    &k_row,
5304                    &v_row,
5305                    &mut kvl.k,
5306                    &mut kvl.v,
5307                    kvl.len + i,
5308                    kv_dim_k,
5309                    kv_dim_v,
5310                    ktb,
5311                    vtb,
5312                    crate::Engine::kv_fp8_on(),
5313                )?;
5314            }
5315            kvl.len += t;
5316        }
5317
5318        // BIT-IDENTICAL VERIFY ATTENTION (spec-exactness fix): the FP accumulation order must be
5319        // byte-for-byte identical to the eager decode path. fa_prefill uses a different tile size
5320        // (BLOCK_Q=64, BK=32) and online-softmax structure than fa_decode's split-K + combine,
5321        // which changes FP summation order and can flip argmax at tight logit margins. Query row r
5322        // attends to keys [0..base_len+r+1) — each successive row sees one more key (the causal
5323        // property). This matches eager: decode appends k at len, then fa_decode sees t_kv = len+1
5324        // keys. The verify appends all T tokens first but bounds the key range per row.
5325        //
5326        // MULTI-ROW FUSED PATH (the long-ctx spec fix, 2026-07-03): when every row takes the vec
5327        // kernel (base_len+1 >= FA_VEC_MIN_TKV), ONE fa_decode_rows launch executes the exact
5328        // per-row program for all T rows (grid.z = row, per-row n_splits from the same
5329        // fa_split_keys formula) — replacing T x (2 launches + 2 dtod copies + 5 partial allocs)
5330        // and multiplying resident CTAs by T on a latency-bound kernel. Bit-identical per row by
5331        // construction; kernel-check pins rows-vs-loop byte identity, run-spec is the end gate.
5332        // Short ctx (any row below the vec crossover) and MEMRA_NO_FA_VEC/MEMRA_FA_ROWS_OFF keep the
5333        // per-row loop (whose fa_decode picks scalar/vec per row exactly like eager decode).
5334        let mut attn = vbuf(e, t * n_head * head_dim)?; // fully written by every FA arm below
5335        let base_len = kvl.len - t; // KV len BEFORE this round's T tokens were appended
5336        // T=1 INCLUDED (2026-07-05): p-min cuts the draft to 1 in ~75% of rounds on hard
5337        // (agentic) content — the old t>1 gate sent those rounds to the per-row loop (262us/row
5338        // + q-row copy + per-row allocs vs 93us/row through the fused kernel at grid.z=1, same
5339        // program). nsys accounting: 1088 of 1456 verify FA launches were T=1 escapees.
5340        // LEAN T=1 ARM (MEMRA_SPEC_LEAN, close35): at t==1, q IS one row and fa_decode on it is
5341        // the EXACT eager decode dispatch (vec_q_v2 + combine_f32; the rows pair measured +50us
5342        // at m=1). Byte-identical: kernel-check pins rows-vs-loop identity, and the per-row loop
5343        // at t=1 is fa_decode on the same q with zero-offset copies. Gates arbitrate.
5344        if let Some(ctr) = stream_ctr {
5345            // STREAM ARM: causal base from the device counter; host kvl.len is a stale lower
5346            // bound used only for the split-sizing upper bound (+64 slack covers M pre-issued
5347            // rounds at K<=8). Views span the bound; per-row limits derive in-kernel.
5348            let upper = kvl.len + t + 64;
5349            let k_view = e.view_u8(&kvl.k, (upper.min(cache.max_ctx)) * ktb);
5350            let v_view = e.view_u8(&kvl.v, (upper.min(cache.max_ctx)) * vtb);
5351            e.fa_decode_rows_dc(
5352                &q,
5353                &k_view,
5354                &v_view,
5355                &mut attn,
5356                head_dim,
5357                n_head,
5358                n_head_kv,
5359                ctr,
5360                upper.min(cache.max_ctx),
5361                t,
5362                scale,
5363                ktb,
5364                vtb,
5365                0,
5366                false,
5367            )?;
5368        } else if spec_lean() && t == 1 {
5369            let t_kv = base_len + 1;
5370            let k_view = e.view_u8(&kvl.k, t_kv * ktb);
5371            let v_view = e.view_u8(&kvl.v, t_kv * vtb);
5372            e.fa_decode_kvmod(
5373                &q,
5374                &k_view,
5375                &v_view,
5376                &mut attn,
5377                head_dim,
5378                n_head,
5379                n_head_kv,
5380                t_kv,
5381                scale,
5382                ktb,
5383                vtb,
5384                crate::Engine::kv_fp8_on(),
5385            )?;
5386        } else if e.fa_rows_eligible(base_len, head_dim) {
5387            let k_view = e.view_u8(&kvl.k, (base_len + t) * ktb);
5388            let v_view = e.view_u8(&kvl.v, (base_len + t) * vtb);
5389            e.fa_decode_rows(
5390                &q,
5391                &k_view,
5392                &v_view,
5393                &mut attn,
5394                head_dim,
5395                n_head,
5396                n_head_kv,
5397                base_len,
5398                t,
5399                scale,
5400                ktb,
5401                vtb,
5402                None,
5403                false,
5404                crate::Engine::kv_fp8_on(),
5405                None,
5406            )?;
5407        } else {
5408            for r in 0..t {
5409                let t_kv_r = base_len + r + 1; // this row sees keys [0..t_kv_r)
5410                let k_view_r = e.view_u8(&kvl.k, t_kv_r * ktb);
5411                let v_view_r = e.view_u8(&kvl.v, t_kv_r * vtb);
5412                // copy q row into an owned buffer (fa_decode takes &CudaSlice, not CudaView)
5413                let mut q_row = vbuf(e, n_head * head_dim)?; // fully written by copy_view_into
5414                let q_src = q.slice(r * n_head * head_dim..(r + 1) * n_head * head_dim);
5415                e.copy_view_into(&mut q_row, 0, &q_src, n_head * head_dim)?;
5416                let mut attn_row = vbuf(e, n_head * head_dim)?; // fully written by fa_decode
5417                e.fa_decode_kvmod(
5418                    &q_row,
5419                    &k_view_r,
5420                    &v_view_r,
5421                    &mut attn_row,
5422                    head_dim,
5423                    n_head,
5424                    n_head_kv,
5425                    t_kv_r,
5426                    scale,
5427                    ktb,
5428                    vtb,
5429                    crate::Engine::kv_fp8_on(),
5430                )?;
5431                e.copy_into(
5432                    &mut attn,
5433                    r * n_head * head_dim,
5434                    &attn_row,
5435                    n_head * head_dim,
5436                )?;
5437            }
5438        }
5439
5440        let attn_g = match &gate {
5441            Some(gate) => {
5442                let mut gsig = vbuf(e, t * n_head * head_dim)?; // fully written by sigmoid
5443                e.sigmoid(gate, &mut gsig, t * n_head * head_dim)?;
5444                let mut ag = vbuf(e, t * n_head * head_dim)?; // fully written by mul
5445                e.mul(&attn, &gsig, &mut ag, t * n_head * head_dim)?;
5446                ag
5447            }
5448            None => attn,
5449        };
5450        // DECODE-EXACT wo projection: at m>=5 (K=4+ with pending) the generic matmul would use dp4a
5451        // (128-thread, different FP sum order than MMVQ). Force MMVQ for bit-identity with decode.
5452        Ok(e.matmul_decode_exact(&fa.wo, &attn_g, t)?)
5453    }
5454
5455    /// Context-linear bytes for a plain serving session's trunk cache.
5456    pub fn plain_session_kv_bytes_per_token(&self) -> usize {
5457        crate::cache::cache_bytes_per_token(&self.cfg)
5458    }
5459
5460    /// `(logical bytes/token, ring-capped bytes/token, ring row cap)` for exact admission.
5461    pub fn plain_session_kv_shape(&self) -> (usize, usize, usize) {
5462        (
5463            self.plain_session_kv_bytes_per_token(),
5464            crate::cache::cache_ring_bytes_per_token(&self.cfg),
5465            crate::cache::cache_ring_row_cap(&self.cfg),
5466        )
5467    }
5468
5469    /// Context-linear bytes for a speculative serving session: trunk cache plus persistent MTP
5470    /// scratch. With no MTP head this equals the plain coefficient.
5471    pub fn spec_session_kv_bytes_per_token(&self) -> usize {
5472        let scratch = self
5473            .mtp
5474            .as_ref()
5475            .map(|mtp| {
5476                let (_, _, k, v) = mtp_scratch_layout(&self.cfg, mtp.geom.as_ref());
5477                k + v
5478            })
5479            .unwrap_or(0);
5480        self.plain_session_kv_bytes_per_token()
5481            .saturating_add(scratch)
5482    }
5483
5484    /// Spec twin of [`HybridModel::plain_session_kv_shape`]; Step35's persistent MTP scratch is
5485    /// capped by the same SWA ring rows as the trunk.
5486    pub fn spec_session_kv_shape(&self) -> (usize, usize, usize) {
5487        let total = self.spec_session_kv_bytes_per_token();
5488        let (_, mut ring, rows) = self.plain_session_kv_shape();
5489        if rows > 0 {
5490            ring = ring.saturating_add(
5491                self.mtp
5492                    .as_ref()
5493                    .map(|mtp| {
5494                        let (_, _, k, v) = mtp_scratch_layout(&self.cfg, mtp.geom.as_ref());
5495                        k + v
5496                    })
5497                    .unwrap_or(0),
5498            );
5499        }
5500        (total, ring, rows)
5501    }
5502
5503    /// Greedy MTP speculative decode (§B). Token-identical to `generate(prompt, max_new)` but uses
5504    /// the NextN head to draft K tokens then verifies them in one batched target forward.
5505    /// Returns (generated tokens, total_drafted, total_accepted) so the caller can report
5506    /// acceptance rate. `k` = draft length per round.
5507    ///
5508    /// GRAPH DRAFT (stage 2 of graph-grade spec): when the model is all-Dense and the MTP head is
5509    /// Dense (no MoE host readbacks), the fixed-shape T=1 MTP forward is CUDA-graph-captured ONCE
5510    /// and replayed per draft step — the ~40 eager launches per drafted token collapse into one
5511    /// graph dispatch; only the 4-byte token id (and 4-byte p-min confidence) round-trip per step.
5512    /// Event tracking is disabled for the whole call (generate_graph pattern) so every buffer the
5513    /// captured graph references is event-free; the spec loop is strictly single-stream.
5514    /// MEMRA_SPEC_NOGRAPH=1 forces the eager draft chain.
5515    /// SAMPLED mode (MEMRA_SPEC_TEMP>0) has its OWN capture (gumbel-perturbed in-graph argmax,
5516    /// device Philox event counter, persistent q retention) — graph-vs-eager sampled streams are
5517    /// bit-identical for the same (seed, prompt, K, temp); see the sampled-graph setup in
5518    /// generate_spec_inner2.
5519    /// Multi-turn session: trunk cache + MTP draft scratch persist across generate calls, so
5520    /// turn N+1 primes ONLY its new suffix (the 124k-conversation daily pattern — re-priming a
5521    /// 32k history costs ~54s; a suffix prime costs seconds). APPEND-ONLY by construction: the
5522    /// hybrid linear-attn states are in-place (no position index), so a session can extend but
5523    /// never rewind — `committed` is the exact token list whose state the caches hold (includes
5524    /// any overshoot tokens past max_new; the caller renders from `committed`, not its own echo).
5525    pub fn new_session(
5526        &self,
5527        e: &Engine,
5528        max_ctx: usize,
5529    ) -> Result<SpecSession, Box<dyn std::error::Error>> {
5530        Ok(SpecSession {
5531            // STAGE-OWNED KV (lane/pp2-spec 2026-08-06): `pp::new_cache`, not `Cache::new`. This
5532            // is the SERVING spec-session path, and with the ppN door open across two cards a
5533            // primary-homed cache makes every remote stage peer-read its OWN KV on every verify
5534            // round — the wrong-card class already fixed on the two batched serving paths
5535            // (worker.rs 2483 / 2837). With the door shut `new_cache` IS `Cache::new` (same
5536            // branch, same allocations), so single-device behavior is byte-unchanged.
5537            cache: crate::pp::new_cache(e, &self.cfg, max_ctx)?,
5538            scratch: MtpScratch::new(
5539                e,
5540                &self.cfg,
5541                max_ctx,
5542                self.mtp.as_ref().and_then(|m| m.geom.as_ref()),
5543            )?,
5544            committed: Vec::new(),
5545            last_h: None,
5546            next_pred: None,
5547            sctr: 0,
5548            uctr: 0,
5549            draft_ctx: None,
5550            pending_tok: None,
5551            turn_ckpt: None,
5552            telem: SpecTelemetryCounters::default(),
5553            capture_at: None,
5554            boundary_capture: None,
5555        })
5556    }
5557
5558    /// Forced-gate exact state comparison. This intentionally reads the real live prefixes from
5559    /// their owning PP devices: matching emitted ids alone would miss a stale `len_d`, recurrent
5560    /// snapshot, or draft-KV row that only corrupts the following round.
5561    pub fn optipipe_compare_session_state(
5562        &self,
5563        e: &Engine,
5564        reference: &SpecSession,
5565        candidate: &SpecSession,
5566    ) -> Result<OptiForkStateIdentity, Box<dyn std::error::Error>> {
5567        fn fail(what: &str) -> Box<dyn std::error::Error> {
5568            format!("optipipe state mismatch: {what}").into()
5569        }
5570        fn same_f32(a: &[f32], b: &[f32]) -> bool {
5571            a.len() == b.len() && a.iter().zip(b).all(|(x, y)| x.to_bits() == y.to_bits())
5572        }
5573        fn compare_layers(
5574            es: &Engine,
5575            range: std::ops::Range<usize>,
5576            reference: &SpecSession,
5577            candidate: &SpecSession,
5578            report: &mut OptiForkStateIdentity,
5579        ) -> Result<(), Box<dyn std::error::Error>> {
5580            for il in range {
5581                match (&reference.cache.kv[il], &candidate.cache.kv[il]) {
5582                    (Some(a), Some(b)) => {
5583                        if a.len != b.len {
5584                            return Err(fail(&format!(
5585                                "layer {il} host KV len {} != {}",
5586                                a.len, b.len
5587                            )));
5588                        }
5589                        let ad = es.dtoh_i32(&a.len_d)?;
5590                        let bd = es.dtoh_i32(&b.len_d)?;
5591                        if ad != bd || ad.first().copied() != Some(a.len as i32) {
5592                            return Err(fail(&format!(
5593                                "layer {il} device KV len {ad:?} != {bd:?} (host={})",
5594                                a.len,
5595                            )));
5596                        }
5597                        let kb = a.len * a.k_tok_bytes;
5598                        let vb = a.len * a.v_tok_bytes;
5599                        if kb > 0 {
5600                            let ak = es.dtoh_u8_view(&a.k.slice(0..kb))?;
5601                            let bk = es.dtoh_u8_view(&b.k.slice(0..kb))?;
5602                            if ak != bk {
5603                                let at = ak.iter().zip(&bk).position(|(x, y)| x != y).unwrap();
5604                                return Err(fail(&format!(
5605                                    "layer {il} K bytes at byte {at} row {} offset {}: {} != {}",
5606                                    at / a.k_tok_bytes,
5607                                    at % a.k_tok_bytes,
5608                                    ak[at],
5609                                    bk[at],
5610                                )));
5611                            }
5612                        }
5613                        if vb > 0 {
5614                            let av = es.dtoh_u8_view(&a.v.slice(0..vb))?;
5615                            let bv = es.dtoh_u8_view(&b.v.slice(0..vb))?;
5616                            if av != bv {
5617                                let at = av.iter().zip(&bv).position(|(x, y)| x != y).unwrap();
5618                                return Err(fail(&format!(
5619                                    "layer {il} V bytes at byte {at} row {} offset {}: {} != {}",
5620                                    at / a.v_tok_bytes,
5621                                    at % a.v_tok_bytes,
5622                                    av[at],
5623                                    bv[at],
5624                                )));
5625                            }
5626                        }
5627                        report.trunk_kv_bytes += kb + vb;
5628                    }
5629                    (None, None) => {}
5630                    _ => return Err(fail(&format!("layer {il} KV presence"))),
5631                }
5632                match (&reference.cache.recur[il], &candidate.cache.recur[il]) {
5633                    (Some(a), Some(b)) => {
5634                        let ac = es.dtoh(&a.conv_state)?;
5635                        let bc = es.dtoh(&b.conv_state)?;
5636                        if !same_f32(&ac, &bc) {
5637                            return Err(fail(&format!("layer {il} conv state")));
5638                        }
5639                        let as_ = es.dtoh(&a.ssm_state)?;
5640                        let bs = es.dtoh(&b.ssm_state)?;
5641                        if !same_f32(&as_, &bs) {
5642                            return Err(fail(&format!("layer {il} SSM state")));
5643                        }
5644                        report.recurrent_bytes += (ac.len() + as_.len()) * 4;
5645                    }
5646                    (None, None) => {}
5647                    _ => return Err(fail(&format!("layer {il} recurrent presence"))),
5648                }
5649            }
5650            Ok(())
5651        }
5652
5653        if reference.committed != candidate.committed {
5654            return Err(fail("committed token ids"));
5655        }
5656        if reference.cache.pos != candidate.cache.pos
5657            || reference.cache.max_ctx != candidate.cache.max_ctx
5658        {
5659            return Err(fail("cache pos/capacity"));
5660        }
5661        if reference.pending_tok != candidate.pending_tok
5662            || reference.next_pred != candidate.next_pred
5663            || reference.sctr != candidate.sctr
5664            || reference.uctr != candidate.uctr
5665        {
5666            return Err(fail("pending/prediction/counter tail"));
5667        }
5668
5669        let mut report = OptiForkStateIdentity::default();
5670        if let Some(fence) = crate::pp::pp_cuts(self.layers.len()) {
5671            let rt = crate::pp::PpNRt::get(e)?;
5672            for stage in 0..rt.n_stages() {
5673                let _scope = rt.enter(stage);
5674                compare_layers(
5675                    rt.engine(stage, e),
5676                    fence[stage]..fence[stage + 1],
5677                    reference,
5678                    candidate,
5679                    &mut report,
5680                )?;
5681            }
5682        } else {
5683            compare_layers(e, 0..self.layers.len(), reference, candidate, &mut report)?;
5684        }
5685
5686        let (a, b) = (&reference.scratch.kv, &candidate.scratch.kv);
5687        if a.len != b.len || e.dtoh_i32(&a.len_d)? != e.dtoh_i32(&b.len_d)? {
5688            return Err(fail("draft scratch length"));
5689        }
5690        let kb = a.len * a.k_tok_bytes;
5691        let vb = a.len * a.v_tok_bytes;
5692        if kb > 0 && e.dtoh_u8_view(&a.k.slice(0..kb))? != e.dtoh_u8_view(&b.k.slice(0..kb))? {
5693            return Err(fail("draft scratch K bytes"));
5694        }
5695        if vb > 0 && e.dtoh_u8_view(&a.v.slice(0..vb))? != e.dtoh_u8_view(&b.v.slice(0..vb))? {
5696            return Err(fail("draft scratch V bytes"));
5697        }
5698        report.scratch_kv_bytes = kb + vb;
5699
5700        match (&reference.last_h, &candidate.last_h) {
5701            (Some(a), Some(b)) => {
5702                let ah = e.dtoh(a)?;
5703                let bh = e.dtoh(b)?;
5704                if !same_f32(&ah, &bh) {
5705                    return Err(fail("last hidden/seed bytes"));
5706                }
5707                report.hidden_bytes = ah.len() * 4;
5708            }
5709            (None, None) => {}
5710            _ => return Err(fail("last hidden/seed presence")),
5711        }
5712        Ok(report)
5713    }
5714
5715    /// SESSION-AFFINITY REWIND (lane/session-affinity, 2026-08-05): roll `sess` back to its
5716    /// retained prompt-end checkpoint, so a request whose prompt matches
5717    /// `committed[..rewind_pos()]` exactly can resume there and prime only its own delta.
5718    ///
5719    /// EXACTNESS. After this returns, the session is byte-for-byte the state it was in AT that
5720    /// boundary: full-attn KV truncated to it (append-only, position-addressed), GDN conv/ssm
5721    /// restored from the device copy taken there, draft scratch length reset, `committed`
5722    /// truncated, `last_h` = the boundary's predecessor anchor. That is precisely the state a
5723    /// fresh prime of `committed[..pos]` would have produced, so the following suffix prime and
5724    /// every burst after it are identical to a cold run of the same token stream — the
5725    /// committed-tokens-authoritative contract.
5726    ///
5727    /// `next_pred` and `pending_tok` are CLEARED: both describe generation past the boundary,
5728    /// which the rewind discards. The caller therefore must supply a non-empty suffix (a
5729    /// rewound session cannot serve an empty-suffix continuation burst — there is nothing to
5730    /// continue). The persistent draft graph survives: it bakes only session-stable pointers
5731    /// (the scratch KV, the resident embedding), none of which the rewind moves.
5732    ///
5733    /// The checkpoint is CONSUMED (`turn_ckpt` taken): its snapshot buffers are freed here, and
5734    /// this turn's own prime installs a fresh one at the new prompt end. Returns the position
5735    /// rewound to, or `None` when the session holds no checkpoint (caller: full re-prime).
5736    pub fn spec_rewind_to_checkpoint(
5737        &self,
5738        e: &Engine,
5739        sess: &mut SpecSession,
5740    ) -> Result<Option<usize>, Box<dyn std::error::Error>> {
5741        if sess.turn_ckpt.as_ref().is_some_and(|ckpt| {
5742            !sess.cache.can_rollback(&ckpt.snap, 0) || !sess.scratch.can_rewind_to(ckpt.pos)
5743        }) {
5744            return Err(
5745                "SWA ring rewind checkpoint has been lapped; full re-prime required".into(),
5746            );
5747        }
5748        let Some(ckpt) = sess.turn_ckpt.take() else {
5749            return Ok(None);
5750        };
5751        assert!(
5752            ckpt.pos <= sess.committed.len(),
5753            "checkpoint past committed ({} > {})",
5754            ckpt.pos,
5755            sess.committed.len()
5756        );
5757        // Restore through each layer's owning engine. A single primary-engine rollback is not
5758        // sufficient when the serving cache is stage-owned under cross-device PP.
5759        crate::pp::restore_cache_checkpoint(e, &self.cfg, None, &mut sess.cache, &ckpt.snap)?;
5760        debug_assert_eq!(
5761            sess.cache.pos, ckpt.pos,
5762            "rollback landed off the checkpoint"
5763        );
5764        sess.scratch.set_len(e, ckpt.pos)?;
5765        sess.committed.truncate(ckpt.pos);
5766        sess.last_h = Some(ckpt.last_h);
5767        sess.next_pred = None;
5768        sess.pending_tok = None;
5769        Ok(Some(ckpt.pos))
5770    }
5771
5772    /// Grow a parked speculative session to `target_cap` and rewind it to its retained turn
5773    /// checkpoint without re-priming the checkpoint prefix.
5774    ///
5775    /// The trunk cache is restored exactly like a plain grown cache: append-only full-attention
5776    /// KV rows come from the parked cache, while recurrent state comes from the checkpoint's
5777    /// owned snapshot. The MTP scratch is also context-linear and its rows below the checkpoint
5778    /// remain authoritative, so they are copied into a fresh larger scratch before its length is
5779    /// truncated. Pointer-baking draft graphs are dropped and recaptured on the next burst.
5780    ///
5781    /// All fallible work completes before `sess` is mutated. A failed allocation or copy leaves
5782    /// the parked session intact, allowing the caller one reclaim-and-retry attempt.
5783    pub fn spec_grow_and_rewind_to_checkpoint(
5784        &self,
5785        e: &Engine,
5786        sess: &mut SpecSession,
5787        target_cap: usize,
5788    ) -> Result<Option<usize>, Box<dyn std::error::Error>> {
5789        if target_cap <= sess.cache.max_ctx {
5790            return self.spec_rewind_to_checkpoint(e, sess);
5791        }
5792        let Some(ckpt) = sess.turn_ckpt.as_ref() else {
5793            return Ok(None);
5794        };
5795        if ckpt.pos == 0 || ckpt.pos > sess.committed.len() {
5796            return Err(format!(
5797                "checkpoint pos {} outside committed length {}",
5798                ckpt.pos,
5799                sess.committed.len(),
5800            )
5801            .into());
5802        }
5803        if ckpt.pos > target_cap {
5804            return Err(format!(
5805                "checkpoint pos {} exceeds grown capacity {target_cap}",
5806                ckpt.pos,
5807            )
5808            .into());
5809        }
5810
5811        let mut grown_cache = crate::pp::new_cache(e, &self.cfg, target_cap)?;
5812        let mut grown_scratch = MtpScratch::new(
5813            e,
5814            &self.cfg,
5815            target_cap,
5816            self.mtp.as_ref().and_then(|m| m.geom.as_ref()),
5817        )?;
5818        crate::pp::restore_cache_checkpoint(
5819            e,
5820            &self.cfg,
5821            Some(&sess.cache),
5822            &mut grown_cache,
5823            &ckpt.snap,
5824        )?;
5825
5826        let src = &sess.scratch.kv;
5827        let dst = &mut grown_scratch.kv;
5828        if ckpt.pos > src.len
5829            || src.kv_dim_k != dst.kv_dim_k
5830            || src.kv_dim_v != dst.kv_dim_v
5831            || src.k_tok_bytes != dst.k_tok_bytes
5832            || src.v_tok_bytes != dst.v_tok_bytes
5833        {
5834            return Err(format!(
5835                "checkpoint draft layout mismatch (pos {}, source len {})",
5836                ckpt.pos, src.len,
5837            )
5838            .into());
5839        }
5840        let kb = ckpt.pos * src.k_tok_bytes;
5841        let vb = ckpt.pos * src.v_tok_bytes;
5842        if kb > 0 {
5843            e.copy_u8_into(&mut dst.k, 0, &src.k, kb)?;
5844        }
5845        if vb > 0 {
5846            e.copy_u8_into(&mut dst.v, 0, &src.v, vb)?;
5847        }
5848        grown_scratch.set_len(e, ckpt.pos)?;
5849        // The old scratch is dropped immediately after publication below. Bound its D2D reads
5850        // first; growth happens once per rewritten turn, outside the decode hot loop.
5851        e.stream().synchronize()?;
5852
5853        let ckpt = sess
5854            .turn_ckpt
5855            .take()
5856            .expect("checkpoint remained present through transactional grow");
5857        let pos = ckpt.pos;
5858        sess.cache = grown_cache;
5859        sess.scratch = grown_scratch;
5860        sess.committed.truncate(pos);
5861        sess.last_h = Some(ckpt.last_h);
5862        sess.next_pred = None;
5863        sess.pending_tok = None;
5864        sess.draft_ctx = None;
5865        debug_assert_eq!(sess.cache.pos, pos, "grown rewind landed off checkpoint");
5866        debug_assert_eq!(
5867            sess.scratch.kv.len, pos,
5868            "grown draft rewind landed off checkpoint"
5869        );
5870        Ok(Some(pos))
5871    }
5872
5873    /// Commit a carried pending bonus (see SpecSession::pending_tok): one T=1 trunk pass
5874    /// (its logits' argmax becomes next_pred) + the draft-KV fill at the carried anchor —
5875    /// byte-identical to the pre-carry session tail. Required before a non-empty-suffix
5876    /// prime, a sampled turn, or parking a session for pool reuse. No-op without a pending.
5877    pub fn spec_flush_pending(
5878        &self,
5879        e: &Engine,
5880        sess: &mut SpecSession,
5881    ) -> Result<(), Box<dyn std::error::Error>> {
5882        let Some(b) = sess.pending_tok.take() else {
5883            return Ok(());
5884        };
5885        let mtp = self
5886            .mtp
5887            .as_ref()
5888            .expect("pending carry requires an MTP head");
5889        let n_embd = self.cfg.n_embd as usize;
5890        let (embd_qt, embd_rb) = self.embd.qt_and_row_bytes(n_embd);
5891        let embd_gpu = if spec_host_embd() {
5892            None
5893        } else {
5894            Some(
5895                self.embd_gpu
5896                    .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload")),
5897            )
5898        };
5899        let embd_dev = embd_gpu.map(|g| (g, embd_qt, embd_rb));
5900        let pos_b = sess.cache.pos;
5901        sess.scratch.set_len(e, pos_b)?;
5902        let (lg_b, hb) = self.spec_target_step_h(e, b, &mut sess.cache)?;
5903        sess.next_pred = Some(argmax(&lg_b) as u32);
5904        let anchor = sess
5905            .last_h
5906            .as_ref()
5907            .expect("pending carry requires last_h (the predecessor-row anchor)");
5908        self.mtp_kv_fill(e, mtp, &[b], anchor, pos_b, &mut sess.scratch, embd_dev)?;
5909        sess.last_h = Some(hb);
5910        sess.committed.push(b);
5911        Ok(())
5912    }
5913
5914    /// Solo target feed used only at speculative round boundaries. Step35 serving made its
5915    /// staged batched B=1 graph authoritative, so a speculative session must enter and leave
5916    /// rounds through that same graph. Other model families keep their eager T=1 contract.
5917    fn spec_target_step_h(
5918        &self,
5919        e: &Engine,
5920        token: u32,
5921        cache: &mut Cache,
5922    ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
5923        if self.cfg.step35.is_none() && !self.qwen35_serving_class() {
5924            return self.decode_step_h(e, token, cache);
5925        }
5926        let pos0 = cache.pos;
5927        let (logits, hidden) = self.decode_step_t_core(e, &[token], pos0, cache, None, None)?;
5928        Ok((e.dtoh(&logits)?, hidden))
5929    }
5930
5931    /// The archs whose LIVE B=1 serving runs the generic BATCHED numeric class (decode_step_batch
5932    /// walk + batched head), so their spec verify must run the SAME class. MoE learned this
5933    /// 2026-08-14 AM (4b777ccc5); the dense hybrid reproduced the identical near-tie flip class
5934    /// the same day on Qwen3.8-27B — eager-class verify logits drift from batched-class serving
5935    /// logits ("1 ULP at layer 2 → 2.3e-1 logit maxdiff at the head"), and the GDN recurrence
5936    /// carries the drift until a near-tie flips deep in generation. One predicate so the five
5937    /// dispatch sites cannot drift apart again.
5938    fn qwen35_serving_class(&self) -> bool {
5939        matches!(
5940            self.cfg.arch,
5941            memra_gguf::config::Arch::Qwen35 | memra_gguf::config::Arch::Qwen35Moe
5942        )
5943    }
5944
5945    /// Reduced-matrix admission for increment 1. This deliberately does not change the PP-2
5946    /// serving policy: the worker calls it only after `MEMRA_SPEC_PIPE=1` and an explicit spec
5947    /// session already exist.
5948    pub fn spec_pipe_available(&self, e: &Engine) -> bool {
5949        if std::env::var("MEMRA_SPEC_PIPE").as_deref() != Ok("1")
5950            || !spec_devacc()
5951            || spec_replay_env_enabled()
5952            || spec_stream()
5953            || std::env::var("MEMRA_SPEC_ADAPT").as_deref() == Ok("1")
5954            || std::env::var("MEMRA_SPEC_PMIN0").as_deref() == Ok("1")
5955            || std::env::var("MEMRA_SPEC_PP_ANATOMY").as_deref() == Ok("1")
5956            || std::env::var("MEMRA_SPEC_PMIN")
5957                .ok()
5958                .and_then(|v| v.parse::<f32>().ok())
5959                .unwrap_or(0.0)
5960                > 0.0
5961            || self.is_gemma4_e4b()
5962            || self.cfg.gemma4.is_some()
5963            || self.mtp.is_none()
5964        {
5965            return false;
5966        }
5967        let Some(cuts) = crate::pp::pp_cuts(self.layers.len()) else {
5968            return false;
5969        };
5970        if cuts.len() != 3 || crate::pp::pp2_streams_off() || !crate::pp::spec_pp_on() {
5971            return false;
5972        }
5973        crate::pp::PpNRt::get(e)
5974            .map(|rt| rt.n_stages() == 2 && rt.cross_device())
5975            .unwrap_or(false)
5976    }
5977
5978    /// Two warm greedy continuation bursts over one PP-2 interval coordinator. The two existing
5979    /// `generate_spec_inner2` call stacks own all per-session round locals; only phase issue order
5980    /// changes. No callback is accepted in increment 1 — the worker publishes each completed burst.
5981    #[allow(clippy::too_many_arguments)]
5982    pub fn generate_spec_session_pair(
5983        &self,
5984        e: &Engine,
5985        sess_a: &mut SpecSession,
5986        max_new_a: usize,
5987        k_a: usize,
5988        sess_b: &mut SpecSession,
5989        max_new_b: usize,
5990        k_b: usize,
5991    ) -> Result<((Vec<u32>, usize, usize), (Vec<u32>, usize, usize)), Box<dyn std::error::Error>>
5992    {
5993        if !self.spec_pipe_available(e) {
5994            return Err("two-session speculative pipeline is outside its reduced matrix".into());
5995        }
5996        if max_new_a == 0 || max_new_b == 0 || k_a == 0 || k_b == 0 {
5997            return Err(
5998                "two-session speculative pipeline requires non-empty positive-K bursts".into(),
5999            );
6000        }
6001        for sess in [&*sess_a, &*sess_b] {
6002            if sess.committed.is_empty()
6003                || sess.last_h.is_none()
6004                || (sess.next_pred.is_none() && sess.pending_tok.is_none())
6005            {
6006                return Err("two-session speculative pipeline requires warm continuations".into());
6007            }
6008        }
6009
6010        let mtp_dense = self
6011            .mtp
6012            .as_ref()
6013            .map(|m| matches!(m.ffn, crate::hybrid::Ffn::Dense { .. }))
6014            .unwrap_or(false);
6015        let trunk_dense = self
6016            .layers
6017            .iter()
6018            .all(|l| matches!(l.ffn, crate::hybrid::Ffn::Dense { .. }));
6019        let graph_ok = std::env::var("MEMRA_SPEC_NOGRAPH").is_err()
6020            && !spec_host_embd()
6021            && mtp_dense
6022            && trunk_dense
6023            && !crate::model::full_prec_enabled();
6024        let graph_a = graph_ok && k_a + 2 < 96;
6025        let graph_b = graph_ok && k_b + 2 < 96;
6026        let was_tracking = e.ctx().is_event_tracking();
6027        if (graph_a || graph_b) && was_tracking {
6028            unsafe {
6029                e.ctx().disable_event_tracking();
6030            }
6031        }
6032
6033        static LOGGED: std::sync::Once = std::sync::Once::new();
6034        LOGGED.call_once(|| {
6035            eprintln!("[spec-pipe] two-session PP-2 continuation pipeline engaged");
6036        });
6037        let sync = std::sync::Arc::new(SpecPipeSync::new());
6038        let lane_a = SpecPipeLane {
6039            sync: sync.clone(),
6040            lane: 0,
6041        };
6042        let lane_b = SpecPipeLane { sync, lane: 1 };
6043        let mut sess_b_ptr = SpecPipeSessionPtr(sess_b as *mut SpecSession);
6044        let (result_a, result_b) = std::thread::scope(|scope| {
6045            let b = scope.spawn(move || {
6046                let mut finish = SpecPipeFinish::new(&lane_b);
6047                let sess_b = unsafe { sess_b_ptr.get_mut() };
6048                let result = e
6049                    .ctx()
6050                    .bind_to_thread()
6051                    .map_err(|err| err.to_string())
6052                    .and_then(|_| {
6053                        self.generate_spec_inner2(
6054                            e,
6055                            &[],
6056                            max_new_b,
6057                            k_b,
6058                            graph_b,
6059                            Some(sess_b),
6060                            None,
6061                            None,
6062                            None,
6063                            None,
6064                            Some(&lane_b),
6065                        )
6066                        .map_err(|err| err.to_string())
6067                    });
6068                finish.close(result.is_err());
6069                result
6070            });
6071            let mut finish = SpecPipeFinish::new(&lane_a);
6072            let result_a = self.generate_spec_inner2(
6073                e,
6074                &[],
6075                max_new_a,
6076                k_a,
6077                graph_a,
6078                Some(sess_a),
6079                None,
6080                None,
6081                None,
6082                None,
6083                Some(&lane_a),
6084            );
6085            finish.close(result_a.is_err());
6086            let result_b = b
6087                .join()
6088                .map_err(|_| "paired speculative session B panicked".to_string())
6089                .and_then(|r| r);
6090            (result_a, result_b)
6091        });
6092
6093        if (graph_a || graph_b) && was_tracking {
6094            unsafe {
6095                e.ctx().enable_event_tracking();
6096            }
6097        }
6098        let result_a = result_a?;
6099        let result_b = result_b.map_err(|err| -> Box<dyn std::error::Error> { err.into() })?;
6100        Ok((result_a, result_b))
6101    }
6102
6103    /// One spec-decode turn on a live session. `suffix` = the NEW tokens only (turn N+1's user
6104    /// message rendered through the chat template continuation). Returns (new tokens emitted,
6105    /// drafted, accepted); session.committed grows by suffix + emitted.
6106    pub fn generate_spec_session(
6107        &self,
6108        e: &Engine,
6109        sess: &mut SpecSession,
6110        suffix: &[u32],
6111        max_new: usize,
6112        k: usize,
6113    ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
6114        self.generate_spec_session_sampled(e, sess, suffix, max_new, k, None, None)
6115    }
6116
6117    /// Serve-path sampled spec: routes the burst through the rejection-sampling verify with
6118    /// per-SESSION Philox continuity (sess.sctr/uctr). None = env-driven (CLI) or greedy.
6119    /// Filters (top-k/p/min-p) apply SYMMETRICALLY to draft q and verify p — distribution-exact
6120    /// for the filtered target (feat/filtered-spec).
6121    ///
6122    /// `on_commit` (sse-cadence, 2026-08-05): called with each newly-emitted slice of the
6123    /// output — once right after the prime's first token, then once per round commit — so a
6124    /// streaming caller can flush text at round cadence instead of once per burst. The slices
6125    /// are disjoint, in order, and concatenate to exactly the returned token vec. Emission-
6126    /// timing only: token bytes, session state, and exactness are untouched.
6127    ///
6128    /// The returned bool is a CONTINUE-VERDICT (admission yield, 2026-08-06): `false` ends
6129    /// the burst at the current round boundary, exactly as if `max_new` had been reached —
6130    /// the caller's scheduler regains control without waiting the burst out. Burst size is
6131    /// content-neutral (spec-levers battery), so an early exit moves WHEN the burst returns,
6132    /// never what tokens say. The slice may be EMPTY (a poll-only boundary — round-stream
6133    /// drains and the defensive tail flush can land with nothing new committed).
6134    #[allow(clippy::too_many_arguments)]
6135    pub fn generate_spec_session_sampled(
6136        &self,
6137        e: &Engine,
6138        sess: &mut SpecSession,
6139        suffix: &[u32],
6140        max_new: usize,
6141        k: usize,
6142        sampling: Option<SpecSampling>,
6143        on_commit: Option<&mut dyn FnMut(&[u32]) -> bool>,
6144    ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
6145        self.generate_spec_session_sampled_prime_split(
6146            e, sess, suffix, max_new, k, sampling, None, on_commit,
6147        )
6148    }
6149
6150    /// Serve-only cold-prime segmentation twin. `prime_split` is the same stable boundary the
6151    /// plain worker would honor before entering its sub-floor tokenwise tail; warm continuations
6152    /// pass `None` and stay on the existing zero-prime path.
6153    #[allow(clippy::too_many_arguments)]
6154    pub fn generate_spec_session_sampled_prime_split(
6155        &self,
6156        e: &Engine,
6157        sess: &mut SpecSession,
6158        suffix: &[u32],
6159        max_new: usize,
6160        k: usize,
6161        sampling: Option<SpecSampling>,
6162        prime_split: Option<usize>,
6163        on_commit: Option<&mut dyn FnMut(&[u32]) -> bool>,
6164    ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
6165        self.generate_spec_session_constrained_prime_split(
6166            e,
6167            sess,
6168            suffix,
6169            max_new,
6170            k,
6171            sampling,
6172            None,
6173            prime_split,
6174            on_commit,
6175        )
6176    }
6177
6178    /// `generate_spec_session_sampled` + GRAMMAR (constrained decoding, 2026-08-03): the
6179    /// hook truncates acceptance at the first grammar-illegal token AFTER the exactness
6180    /// verify (grammar is an extra rejection rule, ordering like the batched-verify twins)
6181    /// and replaces an illegal bonus with the MASKED argmax of the target's own verify
6182    /// column — token-identical to constrained plain greedy decode. GREEDY only (the
6183    /// worker routes sampled constrained to plain decode). Acceptance under tight grammars
6184    /// may drop (drafter is unconstrained); that is measured, not hidden.
6185    #[allow(clippy::too_many_arguments)]
6186    pub fn generate_spec_session_constrained(
6187        &self,
6188        e: &Engine,
6189        sess: &mut SpecSession,
6190        suffix: &[u32],
6191        max_new: usize,
6192        k: usize,
6193        sampling: Option<SpecSampling>,
6194        constraint: Option<&mut dyn SpecConstraint>,
6195        on_commit: Option<&mut dyn FnMut(&[u32]) -> bool>,
6196    ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
6197        self.generate_spec_session_constrained_prime_split(
6198            e, sess, suffix, max_new, k, sampling, constraint, None, on_commit,
6199        )
6200    }
6201
6202    #[allow(clippy::too_many_arguments)]
6203    pub fn generate_spec_session_constrained_prime_split(
6204        &self,
6205        e: &Engine,
6206        sess: &mut SpecSession,
6207        suffix: &[u32],
6208        max_new: usize,
6209        k: usize,
6210        sampling: Option<SpecSampling>,
6211        constraint: Option<&mut dyn SpecConstraint>,
6212        prime_split: Option<usize>,
6213        on_commit: Option<&mut dyn FnMut(&[u32]) -> bool>,
6214    ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
6215        if constraint.is_some() && sampling.is_some_and(|s| s.temp > 0.0) {
6216            return Err(
6217                "constrained spec decode is greedy-only (worker routes sampled \
6218                        constrained to plain decode)"
6219                    .into(),
6220            );
6221        }
6222        // PENDING-CARRY entry flush: a carried bonus precedes any new suffix in the sequence,
6223        // so it must commit BEFORE the suffix primes; the sampled path doesn't carry (its
6224        // round-0 accept needs the commit pass's logits). Empty-suffix greedy bursts — the
6225        // serve continuation case — consume the carry in-loop with zero solo passes.
6226        if sess.pending_tok.is_some()
6227            && (!suffix.is_empty() || sampling.map_or(false, |s| s.temp > 0.0))
6228        {
6229            self.spec_flush_pending(e, sess)?;
6230        }
6231        let mtp_dense = self
6232            .mtp
6233            .as_ref()
6234            .map(|m| matches!(m.ffn, crate::hybrid::Ffn::Dense { .. }))
6235            .unwrap_or(false);
6236        let trunk_dense = self
6237            .layers
6238            .iter()
6239            .all(|l| matches!(l.ffn, crate::hybrid::Ffn::Dense { .. }));
6240        // FULL_PREC forces the EAGER draft: the graph capture would enclose cuBLASLt f32 GEMV
6241        // (the FloatBf16 else-branches) and a bf16_to_f32 dequant alloc — neither is stream-capture
6242        // safe. Eager rides matmul/matmul_decode_exact, which dequant FloatBf16 on use. (§item 2.)
6243        let graph_draft = std::env::var("MEMRA_SPEC_NOGRAPH").is_err()
6244            && !spec_host_embd()
6245            && mtp_dense
6246            && trunk_dense
6247            && k + 2 < 96
6248            && !crate::model::full_prec_enabled();
6249        let was_tracking = e.ctx().is_event_tracking();
6250        if graph_draft && was_tracking {
6251            unsafe {
6252                e.ctx().disable_event_tracking();
6253            }
6254        }
6255        let r = self.generate_spec_inner2(
6256            e,
6257            suffix,
6258            max_new,
6259            k,
6260            graph_draft,
6261            Some(sess),
6262            sampling,
6263            constraint,
6264            on_commit,
6265            prime_split,
6266            None,
6267        );
6268        if graph_draft && was_tracking {
6269            unsafe {
6270                e.ctx().enable_event_tracking();
6271            }
6272        }
6273        let (out, d, a) = r?;
6274        Ok((out, d, a))
6275    }
6276
6277    pub fn generate_spec(
6278        &self,
6279        e: &Engine,
6280        prompt: &[u32],
6281        max_new: usize,
6282        k: usize,
6283    ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
6284        let mtp_dense = self
6285            .mtp
6286            .as_ref()
6287            .map(|m| matches!(m.ffn, crate::hybrid::Ffn::Dense { .. }))
6288            .unwrap_or(false);
6289        let trunk_dense = self
6290            .layers
6291            .iter()
6292            .all(|l| matches!(l.ffn, crate::hybrid::Ffn::Dense { .. }));
6293        // FULL_PREC forces eager (see generate_spec_session note): CUDA graph capture cannot
6294        // enclose cuBLASLt f32 GEMV or the bf16_to_f32 dequant alloc the FloatBf16 path needs.
6295        let graph_draft = std::env::var("MEMRA_SPEC_NOGRAPH").is_err()
6296            && !spec_host_embd()
6297            && mtp_dense
6298            && trunk_dense
6299            && k + 2 < 96
6300            && !crate::model::full_prec_enabled();
6301        if !graph_draft {
6302            return self.generate_spec_inner2(
6303                e, prompt, max_new, k, false, None, None, None, None, None, None,
6304            );
6305        }
6306        let was_tracking = e.ctx().is_event_tracking();
6307        if was_tracking {
6308            unsafe {
6309                e.ctx().disable_event_tracking();
6310            }
6311        }
6312        let r = self.generate_spec_inner2(
6313            e, prompt, max_new, k, true, None, None, None, None, None, None,
6314        );
6315        if was_tracking {
6316            unsafe {
6317                e.ctx().enable_event_tracking();
6318            }
6319        }
6320        r
6321    }
6322
6323    fn generate_spec_inner2(
6324        &self,
6325        e: &Engine,
6326        prompt: &[u32],
6327        max_new: usize,
6328        k: usize,
6329        graph_draft: bool,
6330        mut sess: Option<&mut SpecSession>,
6331        sampling: Option<SpecSampling>,
6332        mut constraint: Option<&mut dyn SpecConstraint>,
6333        mut on_commit: Option<&mut dyn FnMut(&[u32]) -> bool>,
6334        prime_split: Option<usize>,
6335        pipe: Option<&SpecPipeLane>,
6336    ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
6337        assert!(k >= 1, "k must be >= 1");
6338        if let Some(p) = pipe {
6339            p.setup_begin()?;
6340        }
6341        // sse-cadence flush cursor: everything in out[..flushed] has been handed to on_commit.
6342        let mut flushed = 0usize;
6343        // admission yield (2026-08-06): on_commit's continue-verdict; false = end the burst
6344        // at the next round boundary (same exit as max_new reached — the session tail runs).
6345        // Initialized by the unconditional post-prime flush below.
6346        let mut keep_going;
6347        let mtp = self
6348            .mtp
6349            .as_ref()
6350            .expect("generate_spec requires an MTP head (nextn_predict_layers>0)");
6351        let n_vocab = self.output.out_features();
6352        // FR-Spec: the draft head may be TRIMMED (fewer rows than n_vocab); the draft argmax runs
6353        // over the draft vocab and the winning index maps through d2t to a TARGET token id.
6354        // Everything downstream (verify/accept/commit) sees target ids only — exactness unchanged.
6355        let d_vocab = mtp
6356            .shared_head_head
6357            .as_ref()
6358            .unwrap_or(&self.output)
6359            .out_features();
6360        let n_embd = self.cfg.n_embd as usize;
6361        // SESSION MODE: reuse the live cache/scratch, prime only the suffix. `base` = tokens
6362        // already committed (their state is in the caches); 0 = fresh single-shot call.
6363        let session_mode = sess.is_some();
6364        let max_ctx = match sess.as_ref() {
6365            Some(s) => s.cache.max_ctx,
6366            None => prompt.len() + max_new + k + 8,
6367        };
6368        let mut own_cache;
6369        let mut own_scratch;
6370        // PREFIX-CACHE capture request threaded out of the session (lane/spec-prefix-cache):
6371        // (requested split, destination slot). Single-shot per burst; fresh calls have none.
6372        let mut sess_capture: Option<(Option<usize>, &mut Option<SpecBoundaryCapture>)> = None;
6373        let (
6374            cache,
6375            scratch,
6376            mut sess_tail,
6377            mut sess_draft_slot,
6378            mut sess_pending_slot,
6379            sess_ckpt_slot,
6380            sess_telem,
6381        ): (
6382            &mut Cache,
6383            &mut MtpScratch,
6384            Option<(
6385                &mut Vec<u32>,
6386                &mut Option<CudaSlice<f32>>,
6387                &mut Option<u32>,
6388                &mut u32,
6389                &mut u32,
6390            )>,
6391            Option<&mut Option<DraftGraphCtx>>,
6392            Option<&mut Option<u32>>,
6393            Option<&mut Option<SpecCheckpoint>>,
6394            Option<&SpecTelemetryCounters>,
6395        ) = match sess.take() {
6396            Some(sr) => {
6397                let SpecSession {
6398                    cache,
6399                    scratch,
6400                    committed,
6401                    last_h,
6402                    next_pred,
6403                    sctr: s_sctr,
6404                    uctr: s_uctr,
6405                    draft_ctx,
6406                    pending_tok,
6407                    turn_ckpt,
6408                    telem,
6409                    capture_at,
6410                    boundary_capture,
6411                } = sr;
6412                sess_capture = Some((capture_at.take(), boundary_capture));
6413                (
6414                    cache,
6415                    scratch,
6416                    Some((committed, last_h, next_pred, s_sctr, s_uctr)),
6417                    Some(draft_ctx),
6418                    Some(pending_tok),
6419                    Some(turn_ckpt),
6420                    Some(telem),
6421                )
6422            }
6423            None => {
6424                // STAGE-OWNED KV (lane/pp2-spec 2026-08-06) — see `new_session`. Door shut =
6425                // `Cache::new` verbatim.
6426                own_cache = crate::pp::new_cache(e, &self.cfg, max_ctx)?;
6427                // Persistent scratch = max_ctx rows (~2KB/token quantized).
6428                own_scratch = MtpScratch::new(
6429                    e,
6430                    &self.cfg,
6431                    max_ctx,
6432                    self.mtp.as_ref().and_then(|m| m.geom.as_ref()),
6433                )?;
6434                (
6435                    &mut own_cache,
6436                    &mut own_scratch,
6437                    None,
6438                    None,
6439                    None,
6440                    None,
6441                    None,
6442                )
6443            }
6444        };
6445        let base = cache.pos;
6446        // PENDING-CARRY consume (2026-08-01): a carried bonus reaches here only on the
6447        // empty-suffix GREEDY continuation path (generate_spec_session_sampled flushed every
6448        // other case). It enters the round loop as round-0's pending — verify col 0 — exactly
6449        // like a mid-burst full-accept boundary: no init feed, no tail commit pass.
6450        let carried_pending: Option<u32> = sess_pending_slot.as_mut().and_then(|s| s.take());
6451        // PERSISTENT DRAFT KV (the only mode since 2026-07-08 — the legacy round-local scratch,
6452        // MEMRA_SPEC_KVLOCAL, measured -35 acceptance pts on the 27B p3 sweep and was removed;
6453        // acceptance-only — exactness is verify's job either way).
6454        // HIDDEN-PAIRING CONVENTION (DEFAULT = predecessor-row, 2026-07-04 — the 27B acceptance
6455        // unlock, +16pts): the MTP head is TRAINED on rows pairing token x_p with the trunk
6456        // hidden of its PREDECESSOR h_{p-1} (the reference engine's mtp_update shifts the target
6457        // hiddens right by one; its draft step 0 feeds (id_last, TRUE hidden of the row id_last
6458        // was sampled from)). memra's historical convention paired SAME-ROW (x_p, h_p) in the fill
6459        // and seeded chain step 0 through an extra MTP pass on a duplicated token (the
6460        // pseudo-seed) — measured 27B p2 K=3 acceptance 0.569 vs 0.731, p3 0.445 vs 0.63+, and
6461        // the chain steps j>=1 were already predecessor-shaped, so ONLY the fill + step-0 seed
6462        // move. The fill shifts by one and the chain seeds from the predecessor's true hidden
6463        // DIRECTLY (vh_seed / vx[j-1]) — the pseudo pass disappears (one MTP-block pass saved
6464        // per round on top of the acceptance win). Draft-quality-only: exactness stays the
6465        // verify's job either way. (The legacy same-row pairing seam, MEMRA_SPEC_HSAME, and its
6466        // pseudo-seed passes were removed 2026-07-08 — predecessor pairing won by +16 acc pts;
6467        // the legacy round-local scratch, MEMRA_SPEC_KVLOCAL, went with it.)
6468        // REPLAY-FREE PARTIAL ACCEPT (default, 2026-07-03): partial rounds keep the verify's own
6469        // bit-identical committed-prefix state (KV truncate + recur rebuild from the VerifyCkpt)
6470        // and leave the bonus PENDING — no duplicate trunk pass (profiled ~0.54 extra full weight
6471        // reads/round at long ctx). MEMRA_SPEC_REPLAY=1 restores the legacy rollback+replay (A/B
6472        // + fallback seam).
6473        // Qwen35-MoE stays on the correctness reference path until its retained verify-state
6474        // commit is proven equivalent to sequential serving on the long-prompt gate. Replaying
6475        // every accepted round through the serving-class verifier is slower, but prevents a
6476        // numerically exact verify result from carrying a drifted recurrent cache into the next
6477        // round. DENSE qwen35 runs replay-free: its verify already executes the serving batched
6478        // class (qwen35_verify_batch_layers), and the serving-class replay loop below steps
6479        // per-row T=1 (replay.len() full weight reads/round — measured 69 -> 30 tok/s on
6480        // Qwen3.8-27B, 2026-08-15); the replay-free VerifyCkpt commit is gated bit-identical by
6481        // the spec-serve battery before release.
6482        let spec_replay = spec_replay_env_enabled()
6483            || matches!(self.cfg.arch, memra_gguf::config::Arch::Qwen35Moe);
6484        if constraint.is_some() && spec_replay {
6485            return Err(
6486                "constrained spec decode does not support MEMRA_SPEC_REPLAY=1 \
6487                        (legacy replay commits an unmasked bonus)"
6488                    .into(),
6489            );
6490        }
6491        // TRUE-HIDDEN REFRESH (default in persistent-draft-KV mode): every round overwrites the
6492        // committed positions' scratch entries from the verify's exact hiddens (mtp_kv_fill batch)
6493        // instead of keeping chain-approximate entries. MEMRA_SPEC_NOREFRESH=1 = legacy (A/B seam).
6494        let refresh = std::env::var("MEMRA_SPEC_NOREFRESH").is_err();
6495
6496        // prime: BATCHED cache prime (prime_cache — the measured #1 e2e gap: tokenwise primed at
6497        // ~102/38 tok/s vs the engine's ~2000-5900 tok/s batched prefill). prime_cache returns the
6498        // full pre-output_norm hidden stack [T, n_embd], which IS prompt_h (the persistent-draft-KV
6499        // mtp_kv_fill input) — no per-token collection needed. Prompts below PRIME_MIN_T, and
6500        // MEMRA_PRIME_TOKENWISE=1, and frozen Hy3 CPU/GPU expert splits take the tokenwise
6501        // decode_step_h loop. The latter avoids transient GPU staging of the spilled expert bank.
6502        // EMPTY-SUFFIX CONTINUATION (serve bursts): a session turn with NO new tokens resumes
6503        // generation exactly where the last turn stopped — no prime at all. The stashed
6504        // `next_pred` plays prime_logits' argmax role (it IS the argmax of the logits after
6505        // committed.last()); `last_h` seeds the predecessor pairing below. Fresh calls and
6506        // non-empty suffixes take the normal path.
6507        let continuation = prompt.is_empty();
6508        if continuation {
6509            assert!(session_mode, "empty prompt requires a session");
6510            assert!(
6511                sess_tail
6512                    .as_ref()
6513                    .map_or(false, |(c, lh, np, _, _)| !c.is_empty()
6514                        && lh.is_some()
6515                        && (np.is_some() || carried_pending.is_some())),
6516                "empty-suffix continuation needs a primed session (committed + last_h + next_pred|pending)"
6517            );
6518        }
6519        let mut prime_logits;
6520        let mut prompt_h: Option<CudaSlice<f32>> = None;
6521        let t_prime = std::time::Instant::now();
6522        let batched_prime = !continuation
6523            && prompt.len() >= crate::hybrid_forward::PRIME_MIN_T
6524            && std::env::var("MEMRA_PRIME_TOKENWISE").is_err()
6525            && !e.frozen_cpu_experts_prefer_tokenwise_prime();
6526        let prime_split = prime_split.filter(|&split| split > 0 && split < prompt.len());
6527        if prime_split.is_some() && (continuation || base != 0) {
6528            return Err("spec prime split is cold-session-only".into());
6529        }
6530        if continuation {
6531            prime_logits = Vec::new();
6532        } else if let Some(split) = prime_split {
6533            if split < crate::hybrid_forward::PRIME_MIN_T {
6534                return Err(format!(
6535                    "spec prime split {split} is below PRIME_MIN_T {}",
6536                    crate::hybrid_forward::PRIME_MIN_T,
6537                )
6538                .into());
6539            }
6540            // Mirror the plain worker's affinity boundary exactly. The prefix is a request-level
6541            // prime (`queued_after` keeps Step35 arm selection independent of this stop); a tail
6542            // below PRIME_MIN_T then takes the same eager tokenwise continuation as prefill_tick.
6543            // Retain every hidden row so the draft scratch fill remains one coherent prompt.
6544            let mut h_all = e.uninit(prompt.len() * n_embd)?;
6545            let (l, _, h_prefix) =
6546                self.prime_cache(e, &prompt[..split], &mut *cache, prompt.len() - split)?;
6547            e.copy_into(&mut h_all, 0, &h_prefix, split * n_embd)?;
6548            prime_logits = l;
6549            // PREFIX-CACHE BOUNDARY CAPTURE (lane/spec-prefix-cache): the GDN conv/ssm states
6550            // are about to be advanced in place by the tail prime, so this is the ONLY moment
6551            // the boundary's recurrent state exists. Capture iff the worker requested exactly
6552            // this split. cache.pos == split here (the prefix prime just finished). A failed
6553            // snapshot is silent (turn_ckpt convention) — publication is an optimization,
6554            // never a correctness dependency.
6555            if let Some((requested, slot)) = sess_capture.as_mut() {
6556                if *requested == Some(split) {
6557                    debug_assert_eq!(cache.pos, split, "boundary capture off the prime split");
6558                    if let Ok(snap) = cache.snapshot(e) {
6559                        **slot = Some(SpecBoundaryCapture {
6560                            snap,
6561                            pos: split,
6562                            logits: prime_logits.clone(),
6563                        });
6564                    }
6565                }
6566            }
6567            let tail = &prompt[split..];
6568            if tail.len() >= crate::hybrid_forward::PRIME_MIN_T
6569                && std::env::var("MEMRA_PRIME_TOKENWISE").is_err()
6570                && !e.frozen_cpu_experts_prefer_tokenwise_prime()
6571            {
6572                let (l, _, h_tail) = self.prime_cache(e, tail, &mut *cache, 0)?;
6573                e.copy_into(&mut h_all, split * n_embd, &h_tail, tail.len() * n_embd)?;
6574                prime_logits = l;
6575            } else {
6576                for (i, &tok) in tail.iter().enumerate() {
6577                    let (l, h) = self.decode_step_h(e, tok, &mut *cache)?;
6578                    e.copy_into(&mut h_all, (split + i) * n_embd, &h, n_embd)?;
6579                    prime_logits = l;
6580                }
6581            }
6582            if std::env::var("MEMRA_SPEC_STATS").as_deref() == Ok("1") {
6583                eprintln!("[spec-prime] affinity split={split} tail={}", tail.len());
6584            }
6585            prompt_h = Some(h_all);
6586        } else if batched_prime {
6587            let (l, _h_seed, hiddens) = self.prime_cache(e, prompt, &mut *cache, 0)?;
6588            prime_logits = l;
6589            prompt_h = Some(hiddens);
6590        } else {
6591            prime_logits = Vec::new();
6592            prompt_h = Some(e.uninit(prompt.len() * n_embd)?);
6593            for (i, &tok) in prompt.iter().enumerate() {
6594                let (l, h) = self.spec_target_step_h(e, tok, &mut *cache)?;
6595                if let Some(ph) = prompt_h.as_mut() {
6596                    e.copy_into(ph, i * n_embd, &h, n_embd)?;
6597                }
6598                prime_logits = l;
6599            }
6600        }
6601        e.stream().synchronize()?;
6602        // PREFIX-CACHE SEED CAPTURE (lane/spec-prefix-cache): boundary == prompt end (the seed
6603        // case — no shared-prefix split, publish the whole prompt). The prime just finished, so
6604        // cache.pos == base + prompt.len() and the recurrent state IS the boundary state;
6605        // prime_logits are the boundary logits. Cold sessions only (base == 0) — same law as
6606        // prime_split. The mid-prompt capture above already consumed the request if it matched.
6607        if !continuation && base == 0 {
6608            if let Some((requested, slot)) = sess_capture.as_mut() {
6609                if *requested == Some(prompt.len()) && slot.is_none() {
6610                    debug_assert_eq!(cache.pos, prompt.len(), "seed capture off prompt end");
6611                    if let Ok(snap) = cache.snapshot(e) {
6612                        **slot = Some(SpecBoundaryCapture {
6613                            snap,
6614                            pos: prompt.len(),
6615                            logits: prime_logits.clone(),
6616                        });
6617                    }
6618                }
6619            }
6620        }
6621        // Harness timing contract (see crate::PRIME_NANOS): gen-only throughput without the
6622        // prime-subtraction hack.
6623        crate::PRIME_NANOS.store(
6624            t_prime.elapsed().as_nanos() as u64,
6625            std::sync::atomic::Ordering::Relaxed,
6626        );
6627
6628        let (embd_qt, embd_rb) = self.embd.qt_and_row_bytes(n_embd);
6629        // Resident table is fastest when it fits. Large spill deployments can preserve that HBM
6630        // for expert-cache slots and gather only the exact rows needed by MTP/verify from host.
6631        let host_embd = spec_host_embd();
6632        let embd_gpu = if host_embd {
6633            None
6634        } else {
6635            Some(
6636                self.embd_gpu
6637                    .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload")),
6638            )
6639        };
6640        let embd_dev = embd_gpu.map(|g| (g, embd_qt, embd_rb));
6641        if host_embd {
6642            eprintln!(
6643                "[spec] host-row embedding: {} bytes kept off HBM",
6644                self.embd.raw.len()
6645            );
6646        }
6647        let mut out: Vec<u32> = Vec::with_capacity(max_new);
6648        let mut total_drafted = 0usize;
6649        let mut total_accepted = 0usize;
6650
6651        // First generated token = argmax of the prompt's last logits (== greedy's first token).
6652        // Emit it, then FEED it to establish the loop invariant below.
6653        // PENDING-CARRY: the carried bonus was already emitted by the LAST burst — it becomes
6654        // last_token WITHOUT re-emission, and round 0 consumes it as pending (no init feed).
6655        // CONSTRAINED entry rules: the first emitted token is the MASKED argmax of the
6656        // prompt's last logits (plain constrained-greedy identity); a continuation without
6657        // a carried pending would emit an UNMASKED stashed next_pred — refused loudly (the
6658        // worker never resumes constrained sessions from the pool, so this cannot fire).
6659        if let Some(c) = constraint.as_deref_mut() {
6660            if continuation && carried_pending.is_none() {
6661                return Err("constrained spec continuation requires a carried pending \
6662                            (pool resume is unconstrained-only)"
6663                    .into());
6664            }
6665            if !continuation {
6666                c.mask_logits(&mut prime_logits)
6667                    .map_err(|e2| format!("constraint: {e2}"))?;
6668            }
6669        }
6670        let mut last_token = if let Some(b) = carried_pending {
6671            b
6672        } else if continuation {
6673            sess_tail.as_ref().unwrap().2.unwrap()
6674        } else {
6675            argmax(&prime_logits) as u32
6676        };
6677        if carried_pending.is_none() {
6678            out.push(last_token);
6679            // grammar advances with every emitted token (carried pendings were consumed
6680            // by the burst that emitted them).
6681            if let Some(c) = constraint.as_deref_mut() {
6682                c.consume(last_token)
6683                    .map_err(|e2| format!("constraint: {e2}"))?;
6684            }
6685        }
6686        if continuation {
6687            // draft-KV invariant: entries [0..base) are the session's exact fills; truncate any
6688            // overhang so the chain's first append lands at slot base (== committed.len()).
6689            scratch.set_len(e, base)?;
6690        }
6691        // sse-cadence: hand the caller every not-yet-flushed token (disjoint in-order slices
6692        // concatenating to the full `out`). Called after the prime's first token and after each
6693        // round commit — emission timing only, token bytes untouched. The slice may be EMPTY
6694        // (poll-only boundary: zero-round folds commit nothing new); returns the caller's
6695        // continue-verdict (admission yield, 2026-08-06) — false ends the burst at this round.
6696        fn flush_commit(
6697            cb: &mut Option<&mut dyn FnMut(&[u32]) -> bool>,
6698            out: &[u32],
6699            flushed: &mut usize,
6700        ) -> bool {
6701            if let Some(f) = cb.as_mut() {
6702                let keep = f(&out[*flushed..]);
6703                *flushed = out.len();
6704                keep
6705            } else {
6706                true
6707            }
6708        }
6709        keep_going = flush_commit(&mut on_commit, &out, &mut flushed);
6710        // INVARIANT at loop top: `last_token` is the most-recently-committed/emitted token, its
6711        // KV+recur state IS in `cache` (cache.pos = position right AFTER last_token), `last_pred`
6712        // is the greedy ARGMAX of the logits that predict the token FOLLOWING last_token, and
6713        // `h_seed` = last_token's pre-output_norm hidden. Establish it by feeding last_token once
6714        // (mirrors plain greedy). DEVICE-ARGMAX lever: the accept walk only ever consumes the
6715        // argmax of those logits — never the full vector — so a host u32 replaces the Vec<f32>.
6716        // --- SAMPLED SPEC (MEMRA_SPEC_TEMP>0, research/sampled-spec-impl-map.md): rejection-
6717        // sampling verify (Leviathan/Chen) — accept draft x at u < p(x)/q(x), resample from
6718        // norm(max(0,p-q)) on reject, bonus sampled from p on full accept. Counter-based Philox
6719        // everywhere (seed, event) -> reproducible. temp==0/unset = the greedy path, untouched.
6720        let sp = sampling.unwrap_or_else(|| SpecSampling {
6721            temp: std::env::var("MEMRA_SPEC_TEMP")
6722                .ok()
6723                .and_then(|v| v.parse().ok())
6724                .unwrap_or(0.0),
6725            seed: std::env::var("MEMRA_SEED")
6726                .ok()
6727                .and_then(|v| v.parse().ok())
6728                .unwrap_or(42),
6729            top_k: std::env::var("MEMRA_TOP_K")
6730                .ok()
6731                .and_then(|v| v.parse().ok())
6732                .unwrap_or(0),
6733            top_p: std::env::var("MEMRA_TOP_P")
6734                .ok()
6735                .and_then(|v| v.parse().ok())
6736                .unwrap_or(1.0),
6737            min_p: std::env::var("MEMRA_MIN_P")
6738                .ok()
6739                .and_then(|v| v.parse().ok())
6740                .unwrap_or(0.0),
6741            penalty_last_n: std::env::var("MEMRA_PENALTY_LAST_N")
6742                .ok()
6743                .and_then(|v| v.parse().ok())
6744                .unwrap_or(0),
6745            penalty_repeat: std::env::var("MEMRA_PENALTY_REPEAT")
6746                .ok()
6747                .and_then(|v| v.parse().ok())
6748                .unwrap_or(1.0),
6749            penalty_freq: std::env::var("MEMRA_PENALTY_FREQ")
6750                .ok()
6751                .and_then(|v| v.parse().ok())
6752                .unwrap_or(0.0),
6753            penalty_present: std::env::var("MEMRA_PENALTY_PRESENT")
6754                .ok()
6755                .and_then(|v| v.parse().ok())
6756                .unwrap_or(0.0),
6757        });
6758        let (sp_temp, sp_seed) = (sp.temp, sp.seed);
6759        let sampled = sp_temp > 0.0;
6760        // Trimmed heads: q lives on the trimmed vocab; accept gathers use the TRIMMED index and
6761        // the residual scatters q into target-id space (q=-inf off-trim — the head cannot propose
6762        // those, so their residual mass is p(x), correct by construction).
6763        let d2t_dev: Option<CudaSlice<u32>> = if sampled || crate::spec::spec_stream() {
6764            match &mtp.d2t {
6765                Some(map) => Some(e.htod_u32_v(map)?),
6766                None => None,
6767            }
6768        } else {
6769            None
6770        };
6771        let mut q_full_buf: Option<CudaSlice<f32>> = None;
6772        // Counters resume from the session (burst continuity: randomness must never repeat
6773        // across generate_spec_session calls); one-shot callers start at (0,0). Read through
6774        // sess_tail — `sess` was take()n into it above, so sess.as_ref() here is always None.
6775        let mut sctr: u32 = sess_tail.as_ref().map(|(_, _, _, s, _)| **s).unwrap_or(0);
6776        let mut uctr: u32 = sess_tail.as_ref().map(|(_, _, _, _, u)| **u).unwrap_or(0);
6777        // host Philox4x32-10 (mirrors spec_sample.cu; independent stream via ctr_lo tag)
6778        let host_u01 = |seed: u64, ctr: u32| -> f32 {
6779            let (m0, m1) = (0xD2511F53u32, 0xCD9E8D57u32);
6780            let (mut c0, mut c1, mut c2, mut c3) = (0xFFFF_FFFEu32, ctr, 0u32, 0u32);
6781            let (mut k0, mut k1) = ((seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
6782            for _ in 0..10 {
6783                let (h0, l0) = (((m0 as u64 * c0 as u64) >> 32) as u32, m0.wrapping_mul(c0));
6784                let (h1, l1) = (((m1 as u64 * c2 as u64) >> 32) as u32, m1.wrapping_mul(c2));
6785                let (n0, n1, n2, n3) = (h1 ^ c1 ^ k0, l1, h0 ^ c3 ^ k1, l0);
6786                c0 = n0;
6787                c1 = n1;
6788                c2 = n2;
6789                c3 = n3;
6790                k0 = k0.wrapping_add(0x9E3779B9);
6791                k1 = k1.wrapping_add(0xBB67AE85);
6792            }
6793            (c0 as f32 + 1.0) * (1.0 / 4294967296.0)
6794        };
6795        let mut draft_logits: Vec<CudaSlice<f32>> = Vec::new(); // retained head logits (q), per slot
6796        let mut draft_stats: Vec<(f32, f32, f32)> = Vec::new(); // (row_max, th_e, z_e) per slot
6797        let mut perturb_buf: Option<CudaSlice<f32>> = None; // gumbel scratch (max(n_vocab,d_vocab))
6798        let mut sample_tok = e.alloc_u32_zeroed(1)?; // residual/bonus sample out
6799        let mut col_buf: Option<CudaSlice<f32>> = None; // materialized verify column
6800        // Penalties (v2.1): applied to COPIES of q rows and p columns symmetrically (exactness
6801        // for the penalized+filtered target). History = generated tokens, host-tracked window.
6802        let pen_on = sampled
6803            && sp.penalty_last_n > 0
6804            && (sp.penalty_repeat != 1.0 || sp.penalty_freq != 0.0 || sp.penalty_present != 0.0);
6805        let mut pen_hist: Vec<u32> = if pen_on {
6806            prompt.iter().rev().take(64).rev().cloned().collect() // llama-parity: history spans prompt tail too
6807        } else {
6808            Vec::new()
6809        };
6810        let mut pen_hist_d: Option<CudaSlice<u32>> = None;
6811        let mut pcol_buf: Option<CudaSlice<f32>> = None; // penalized p-column scratch
6812        // MEMRA_SPEC_SETUP_TRACE=1 (diagnostics): per-call wall decomposition of the burst
6813        // SETUP + TAIL segments (the round loop's internals are MEMRA_SPEC_PHASE's job) —
6814        // built to pin the serve per-burst fixed cost (research/spec-serving-20260801).
6815        let setup_trace = std::env::var("MEMRA_SPEC_SETUP_TRACE").as_deref() == Ok("1");
6816        let t_ent = std::time::Instant::now();
6817
6818        // SESSION-AFFINITY TURN CHECKPOINT (lane/session-affinity, 2026-08-05): capture the
6819        // PROMPT-END boundary state so a LATER turn can rewind here and re-prime only its own
6820        // delta instead of the whole conversation. See `SpecCheckpoint` for why this boundary is
6821        // the one that matters (a history-rewriting client mutates what the session GENERATED,
6822        // so the next turn's prompt agrees with this one up to exactly here).
6823        //
6824        // WHERE — AND WHY THIS EXACT LINE. Right after the trunk prime, BEFORE the init feed
6825        // (`decode_step_h(last_token)`) and before round 0: the last instant at which the caches
6826        // hold exactly `base + prompt.len()` rows and nothing generated.
6827        //
6828        // This was WRONG in the first cut of this lane: the capture sat after the draft-KV fill,
6829        // which is also after the init feed, so `cache.pos` was `base + prompt.len() + 1` — the
6830        // boundary included the FIRST GENERATED TOKEN. That token is the first thing inside the
6831        // `<think>` block the client strips, so every later turn's diff diverged exactly one
6832        // token below the checkpoint and affinity declined 100% of the time. Measured on the
6833        // owner regime: "history diverged at 12233 of checkpoint 12234". The off-by-one made the
6834        // whole mechanism inert while looking, from the outside, like a working
6835        // correctness-declines-safely path — hence the decline log carries the offsets.
6836        //
6837        // The full-attn planes are `len`-truncatable so the snapshot copies only the GDN conv/ssm
6838        // state (the reason a spec session could not rewind before). The draft scratch needs no
6839        // copy: rows below the boundary are rewritten by the next turn's own fill.
6840        //
6841        // WHEN: non-empty prime only. An empty-suffix continuation burst adds no prompt boundary
6842        // (its "prompt end" IS the previous checkpoint's, already held), so it keeps the existing
6843        // checkpoint rather than replacing it with a strictly worse one.
6844        //
6845        // FAILURE IS SILENT BY DESIGN: on a VRAM-tight rig the snapshot alloc can fail. That
6846        // costs the NEXT turn its rewind (it re-primes fully, today's behavior) and must never
6847        // fail the burst that is already running — so the error is swallowed, loud only under
6848        // MEMRA_DEBUG_SPEC.
6849        if let Some(slot) = sess_ckpt_slot {
6850            if !continuation {
6851                let pos = cache.pos;
6852                debug_assert_eq!(
6853                    pos,
6854                    base + prompt.len(),
6855                    "turn checkpoint must sit at the prompt end, before the init feed"
6856                );
6857                let anchor: Result<CudaSlice<f32>, Box<dyn std::error::Error>> =
6858                    if let Some(ph) = &prompt_h {
6859                        // hidden of the LAST primed row = the predecessor anchor at this
6860                        // boundary (exactly what a fresh prime of committed[..pos] leaves in
6861                        // last_h, and what the next prime's fill reads for its first row).
6862                        let np = prompt.len();
6863                        e.uninit(n_embd).and_then(|mut a| {
6864                            e.copy_view_into(
6865                                &mut a,
6866                                0,
6867                                &ph.slice((np - 1) * n_embd..np * n_embd),
6868                                n_embd,
6869                            )?;
6870                            Ok(a)
6871                        })
6872                    } else {
6873                        Err("no prompt hiddens".into())
6874                    };
6875                match (cache.snapshot(e), anchor) {
6876                    (Ok(snap), Ok(last_h)) => {
6877                        *slot = Some(SpecCheckpoint { snap, pos, last_h });
6878                    }
6879                    (s, a) => {
6880                        *slot = None; // a stale checkpoint would rewind to the WRONG boundary
6881                        if std::env::var("MEMRA_DEBUG_SPEC").is_ok() {
6882                            let err = s
6883                                .err()
6884                                .map(|e| e.to_string())
6885                                .or_else(|| a.err().map(|e| e.to_string()))
6886                                .unwrap_or_default();
6887                            eprintln!(
6888                                "[spec] turn checkpoint skipped ({err}); \
6889                                       next turn re-primes in full"
6890                            );
6891                        }
6892                    }
6893                }
6894            }
6895        }
6896        // INIT FEED — skipped on a pending carry: last_token (the carried bonus) is NOT in the
6897        // caches and must NOT be fed solo; round 0's batched verify commits it as col 0. Its
6898        // seed/anchor hidden is the carried last_h (copied below); last_pred is dead in the
6899        // pending path (t_pred reads verify col 0 — the accept walk overwrites it).
6900        let mut last_pred = 0u32;
6901        let mut last_col_logits: Option<CudaSlice<f32>> = None;
6902        // CONSTRAINED: the init feed's logits back the (n_acc==0, base==0) masked-argmax
6903        // recompute in the grammar-truncation walk — retained host-side, round 0 only.
6904        let mut init_logits_host: Option<Vec<f32>> = None;
6905        let h_seed0: CudaSlice<f32> = if carried_pending.is_none() {
6906            let (init_logits, h) = self.spec_target_step_h(e, last_token, &mut *cache)?;
6907            last_pred = argmax(&init_logits) as u32;
6908            if constraint.is_some() {
6909                init_logits_host = Some(init_logits.clone());
6910            }
6911            // sampled mode: p-distribution after last_token, for the j==0/base==0 accept test.
6912            if sampled {
6913                last_col_logits = Some(e.htod(&init_logits)?);
6914            }
6915            h
6916        } else {
6917            // predecessor-row anchor: hidden of the last COMMITTED row (the carry contract).
6918            let lh = sess_tail
6919                .as_ref()
6920                .unwrap()
6921                .1
6922                .as_ref()
6923                .expect("pending carry requires last_h");
6924            e.clone_dtod(lh)?
6925        };
6926        let t_init = t_ent.elapsed();
6927        let mut last_col_stats: Option<(f32, f32, f32)> = None;
6928        // PERSISTENT h_seed buffer (allocated BEFORE any graph capture so no captured scratch can
6929        // alias it): every path that updates the round seed copies INTO it — no per-round allocs,
6930        // stable pointer for the graph-draft round-start copy.
6931        let mut h_seed_buf = e.clone_dtod(&h_seed0)?;
6932        // Predecessor-pairing trackers: `fill_prev` = trunk hidden AT the last COMMITTED row (the
6933        // predecessor of the next verify's col 0 — the reference's carried pending-h analogue;
6934        // also the predecessor-row hidden for the round-0 legacy-replay seed). At round 0 that
6935        // row is last_token's own (h_seed0). The chain step-0 seed under the pairing default =
6936        // hidden of the row BEFORE last_token = the prompt's last row at round 0 (h_seed_buf
6937        // overwritten below).
6938        let mut fill_prev = e.clone_dtod(&h_seed0)?;
6939        {
6940            if let Some(ph) = &prompt_h {
6941                let np = prompt.len();
6942                e.copy_view_into(
6943                    &mut h_seed_buf,
6944                    0,
6945                    &ph.slice((np - 1) * n_embd..np * n_embd),
6946                    n_embd,
6947                )?;
6948            } else if continuation {
6949                if let Some((_, lh, _, _, _)) = sess_tail.as_ref() {
6950                    if let Some(lh) = lh.as_ref() {
6951                        e.copy_into(&mut h_seed_buf, 0, lh, n_embd)?;
6952                    }
6953                }
6954            }
6955        }
6956        // Persistent device prediction slots for the accept walk (max k+1 verify columns).
6957        let mut preds_d = e.alloc_u32_zeroed(k + 2)?;
6958
6959        let debug_spec = std::env::var("MEMRA_DEBUG_SPEC").is_ok();
6960        let fork_mode = OptiForkGateMode::configured();
6961        // MEMRA_SPEC_STATS=1: per-slot accept histogram + draft-length histogram, printed once at
6962        // the end. Metric normalization vs the reference engine: BOTH engines count
6963        // accepted/drafted where the chain stopped at p-min and the sub-threshold token is
6964        // discarded uncounted — per-slot decay + chain-length mix are the extra dimensions.
6965        let spec_stats = std::env::var("MEMRA_SPEC_STATS").is_ok();
6966        let mut st_drafted = vec![0usize; k];
6967        let mut st_accepted = vec![0usize; k];
6968        let mut st_len_hist = vec![0usize; k + 1];
6969        let mut st_full = 0usize;
6970        // P-MIN CONFIDENCE GATE (MEMRA_SPEC_PMIN, the serve script's --spec-draft-p-min mechanism):
6971        // stop the draft chain early when the head's softmax confidence in its own pick drops
6972        // below p_min. Hoisted above the loop: the graph capture bakes the prob kernels iff on.
6973        static PMIN: std::sync::OnceLock<f32> = std::sync::OnceLock::new();
6974        let p_min = *PMIN.get_or_init(|| {
6975            std::env::var("MEMRA_SPEC_PMIN")
6976                .ok()
6977                .and_then(|v| v.parse().ok())
6978                .unwrap_or(0.0)
6979        });
6980        // ZERO-DRAFT ROUNDS (MEMRA_SPEC_PMIN0=1, vendored from llama.cpp's draft gating): let the
6981        // p-min gate apply at j==0 too, so a low-confidence round drafts NOTHING and the verify
6982        // batch is just the pending bonus (m=1 = a plain decode step). llama's 35B win rides
6983        // exactly this — draft acceptance 76% at mean len 2.5 because unpredictable stretches
6984        // never pay draft+verify overhead. Only legal when a pending bonus exists (an empty
6985        // verify batch is not); the j==0 exemption stays for pending-less rounds.
6986        let pmin0 = std::env::var("MEMRA_SPEC_PMIN0")
6987            .map(|v| v == "1")
6988            .unwrap_or(false);
6989
6990        // --- GRAPH DRAFT setup: persistent I/O buffers + ONE capture (2 warmups inside). The
6991        // warmups mutate scratch len_d / pos / tok / seed — all reset at every round start, so the
6992        // only restore needed is the scratch counter. Capture failure (e.g. a non-capturable
6993        // cuBLAS path in an exotic head) falls back to the eager draft chain.
6994        // PER-SESSION PERSISTENCE (2026-08-01): session calls reuse the DraftGraphCtx parked on
6995        // the SpecSession — the capture (2 warmup head forwards + instantiate) ran ONCE at the
6996        // session's first burst, not per burst (measured ~16ms/burst fixed cost on H100 q27,
6997        // research/spec-serving-20260801). Reuse is pointer-exact: the graph bakes the session's
6998        // own scratch KV (never realloc'd), the model's resident embedding, the OnceLock p_min,
6999        // and the g_* buffers carried in the ctx — replay dispatch is identical to a fresh
7000        // capture, so draft tokens are bit-identical (drafts never decide exactness anyway; the
7001        // verify arbitrates). Single-shot calls (sess=None) build a fresh ctx and drop it.
7002        let mut dctx: DraftGraphCtx = match sess_draft_slot.as_mut().and_then(|s| s.take()) {
7003            Some(c) => c,
7004            None => DraftGraphCtx::new(e, n_embd, if sampled { d_vocab } else { 1 })?,
7005        };
7006        // A session that ran greedy bursts first sized g_q/g_perturb at 1; a sampled resume
7007        // needs d_vocab. Realloc is legal exactly while graph_s is None (nothing baked them).
7008        if sampled && dctx.g_q.len() < d_vocab {
7009            dctx.g_q = e.zeros(d_vocab)?;
7010            dctx.g_perturb = e.zeros(d_vocab)?;
7011        }
7012        // DRAFT-SIDE GRAMMAR MASK (lane/draft-mask, 2026-08-04): the drafter samples the
7013        // grammar's legal set, so proposals are legal BY CONSTRUCTION and the verify-side
7014        // truncation (the correctness backstop) stops cutting every tight-schema round.
7015        // The mask is one node inside the captured draft chain — presence is a CAPTURE-TIME
7016        // shape, so a parked graph of the other shape is dropped and recaptured.
7017        let dmask_on = constraint
7018            .as_deref()
7019            .is_some_and(|c| c.draft_mask_enabled());
7020        let dmask_words = if dmask_on { d_vocab.div_ceil(32) } else { 0 };
7021        if dmask_on && dctx.g_dmask.len() < dmask_words {
7022            dctx.g_dmask = e.alloc_u32_zeroed(dmask_words)?;
7023            dctx.graph = None; // the old capture baked the old (or no) mask pointer
7024            dctx.failed.clear_greedy();
7025            dctx.keeper.clear();
7026        }
7027        if dctx.graph.is_some() && dctx.graph_masked != dmask_on {
7028            dctx.graph = None;
7029            dctx.failed.clear_greedy();
7030            dctx.keeper.clear();
7031        }
7032        if graph_draft && !sampled && dctx.graph.is_none() && !dctx.failed.greedy_failed() {
7033            let DraftGraphCtx {
7034                g_tok,
7035                g_pos,
7036                g_seed,
7037                g_p,
7038                g_dmask,
7039                ..
7040            } = &mut dctx;
7041            // capture-time contents: ALL-ONES (ban nothing). A replay only ever runs after the
7042            // host uploads the position's real words, so the warmups stay grammar-free.
7043            if dmask_on {
7044                e.htod_u32_into(g_dmask, &vec![u32::MAX; dmask_words])?;
7045            }
7046            let g_dmask_ro: &CudaSlice<u32> = &*g_dmask;
7047            // CAPTURE-RETAIN (#68 fix): the warmup transients' pool addresses are baked into the
7048            // captured graph; the keeper pins them for the graph's lifetime. capture_graph (non-
7049            // retained) freed them at exit — safe for one-shot generate_spec (nothing else touches
7050            // the pool between replays) but WRONG for sessions: burst-boundary prime/fill/commit
7051            // passes (and, in serve, other sessions) recycle those addresses and the replay then
7052            // clobbers live buffers — the ST serve-spec corruption (research/serve-st-20260803).
7053            let cap_res = e.capture_graph_retained(|e| {
7054                self.mtp_head_forward_cap(
7055                    e,
7056                    mtp,
7057                    g_tok,
7058                    g_pos,
7059                    g_seed,
7060                    g_p,
7061                    &mut *scratch,
7062                    p_min > 0.0 || fork_mode == OptiForkGateMode::Controller,
7063                    true,
7064                    embd_gpu.expect("graph draft requires resident embedding"),
7065                    embd_qt,
7066                    embd_rb,
7067                    d_vocab,
7068                    None,
7069                    None,
7070                    if dmask_on {
7071                        Some((g_dmask_ro, dmask_words))
7072                    } else {
7073                        None
7074                    },
7075                )
7076            });
7077            match cap_res {
7078                Ok((g, keep)) => {
7079                    scratch.set_len(e, base)?;
7080                    dctx.graph = Some(g);
7081                    dctx.graph_masked = dmask_on;
7082                    dctx.keeper = keep;
7083                }
7084                Err(err) => {
7085                    scratch.set_len(e, base)?;
7086                    // LOUD flip (audit Q2): a dropped draft graph is a coverage loss, never
7087                    // silent. Once per flip — mark returns None on an already-failed ctx.
7088                    if let Some(line) = dctx.failed.mark_greedy(&err.to_string()) {
7089                        eprintln!("{line}");
7090                    }
7091                }
7092            }
7093        }
7094        // --- SAMPLED GRAPH DRAFT setup (step 3 of the sampled-spec arc): a SECOND capture, own
7095        // graph object, built only when sampled && graph-eligible — the greedy capture above is
7096        // untouched (and skipped when sampled: its graph would never be launched). Same head
7097        // forward, but the in-graph argmax reads GUMBEL-PERTURBED logits; the Philox event
7098        // counter lives in the persistent device g_ctr (bumped in-graph, host-seeded from sctr
7099        // once per round); the raw head logits land in the persistent g_q for the host's
7100        // per-replay async D2D into the round's q slot (q_slots, K x d_vocab, allocated once).
7101        // seed/temp are capture-time constants — baked into graph_s, so a pool-resumed request
7102        // with a different (seed, temp, k) drops the parked sampled graph and recaptures.
7103        // COST OF THE FRESH-SEED SERVE DEFAULT (dogfood F4, 2026-08-04): omitting `seed` on a
7104        // serve request now draws fresh per-request entropy (it used to default to a pinned 0),
7105        // so a seed-omitting request that RESUMES a parked spec session finds an s_key baked
7106        // with the PREVIOUS request's seed and pays one recapture. Bounded, and it does not
7107        // reopen the ~16ms/burst regression the persistent ctx exists to fix: a session's seed
7108        // is fixed for its whole lifetime (worker.rs reads s.sampler.seed() per burst), so
7109        // this compare misses at most ONCE per resumed request — the first burst recaptures
7110        // and every later burst in that request replays. A client that wants the parked graph
7111        // AND reproducibility supplies an explicit `seed`, honored exactly, which keeps s_key
7112        // stable across its whole conversation.
7113        // COMPOSITION RULE (fspec x gsd merge): the in-graph chain samples from the RAW
7114        // softmax — it can hold neither per-row filter stats nor the varying penalty history.
7115        // The sampled graph therefore engages only in the PURE-TEMP regime; filters/penalties
7116        // force the eager draft (which computes stats/penalties per row).
7117        let pure_temp = sp.top_k == 0 && sp.top_p >= 1.0 && sp.min_p <= 0.0 && !pen_on;
7118        let s_key = (sp_seed, sp_temp.to_bits(), k);
7119        if sampled && dctx.s_key.is_some_and(|old| old != s_key) {
7120            dctx.graph_s = None;
7121            dctx.failed.clear_sampled();
7122            dctx.s_key = None;
7123            dctx.q_slots.clear();
7124            dctx.keeper_s.clear();
7125        }
7126        if graph_draft
7127            && sampled
7128            && pure_temp
7129            && dctx.graph_s.is_none()
7130            && !dctx.failed.sampled_failed()
7131        {
7132            let DraftGraphCtx {
7133                g_tok,
7134                g_pos,
7135                g_seed,
7136                g_p,
7137                g_ctr,
7138                g_perturb,
7139                g_q,
7140                ..
7141            } = &mut dctx;
7142            // CAPTURE-RETAIN (#68 fix): same keeper contract as the greedy capture above.
7143            let cap_res = e.capture_graph_retained(|e| {
7144                self.mtp_head_forward_cap(
7145                    e,
7146                    mtp,
7147                    g_tok,
7148                    g_pos,
7149                    g_seed,
7150                    g_p,
7151                    &mut *scratch,
7152                    p_min > 0.0,
7153                    true,
7154                    embd_gpu.expect("graph draft requires resident embedding"),
7155                    embd_qt,
7156                    embd_rb,
7157                    d_vocab,
7158                    Some((g_ctr, g_perturb, g_q, sp_seed, sp_temp)),
7159                    None,
7160                    None, // constrained spec is greedy-only — sampled never carries a hook
7161                )
7162            });
7163            match cap_res {
7164                Ok((g, keep)) => {
7165                    scratch.set_len(e, base)?;
7166                    for _ in 0..k {
7167                        dctx.q_slots.push(e.zeros(d_vocab)?);
7168                    }
7169                    dctx.graph_s = Some(g);
7170                    dctx.s_key = Some(s_key);
7171                    dctx.keeper_s = keep;
7172                }
7173                Err(err) => {
7174                    scratch.set_len(e, base)?;
7175                    // LOUD flip (audit Q2): same contract as the greedy capture above.
7176                    if let Some(line) = dctx.failed.mark_sampled(&err.to_string()) {
7177                        eprintln!("{line}");
7178                    }
7179                }
7180            }
7181        }
7182        let t_cap = t_ent.elapsed();
7183        // PERSISTENT DRAFT KV: fill the MTP block's K/V for every prompt position from the exact
7184        // trunk hiddens collected during prime — ONE batched K/V-only pass (overwrites any
7185        // capture-warmup garbage; capture left len at 0). last_token (the init feed) needs no
7186        // fill: the first chain step processes it and appends its entry at slot prompt.len().
7187        if let Some(ph) = &prompt_h {
7188            // SESSION: rows [0..base) are the previous turns' exact fills (refresh overwrote them
7189            // with true verify hiddens) — truncate any draft overhang, fill ONLY the suffix at
7190            // global positions [base..base+tp). Fresh call: base==0, identical to before.
7191            scratch.set_len(e, base)?;
7192            // CHUNKED FILL (long-ctx OOM fix, 2026-07-05): mtp_kv_fill's transients scale with its
7193            // T (concat = T*2*n_embd*4B — 1.5GB at 40k) and its concat loop is 2*T launches. The
7194            // fill is a pure sequential append, so chunking is exact: each chunk appends its rows
7195            // at pos0=base+start with the identical per-row math. Same knob as the trunk prime.
7196            let tp = prompt.len();
7197            let fill_chunk: usize = if crate::cache::swa_ring_on() {
7198                crate::hybrid_forward::prime_chunk_tokens(tp, self.layers.len())
7199            } else {
7200                // Preserve the flag-OFF schedule byte-for-byte, including the legacy zero value
7201                // meaning one monolithic fill.
7202                std::env::var("MEMRA_PRIME_CHUNK")
7203                    .ok()
7204                    .and_then(|v| v.parse().ok())
7205                    .unwrap_or(4096)
7206            };
7207            let fill_chunk = if fill_chunk == 0 { tp } else { fill_chunk };
7208            let mut start = 0usize;
7209            while start < tp {
7210                let end = (start + fill_chunk).min(tp);
7211                let tc = end - start;
7212                {
7213                    // PREDECESSOR pairing: row i gets h[i-1]; global row 0 a zeros row (the
7214                    // reference engine's initial pending-h is zeroed too); a session turn's row 0
7215                    // gets the PREVIOUS turn's last committed hidden (sess.last_h). Per chunk:
7216                    // rows start..end read h[start-1..end-1] — one dtod into a chunk buffer.
7217                    let mut phs = e.zeros(tc * n_embd)?;
7218                    let (src_lo, dst_off) = if start == 0 {
7219                        (0, n_embd)
7220                    } else {
7221                        ((start - 1) * n_embd, 0)
7222                    };
7223                    let n_copy = if start == 0 {
7224                        (tc - 1) * n_embd
7225                    } else {
7226                        tc * n_embd
7227                    };
7228                    if start == 0 {
7229                        if let Some((_, lh, _, _, _)) = sess_tail.as_ref() {
7230                            if let Some(lh) = lh.as_ref() {
7231                                e.copy_into(&mut phs, 0, lh, n_embd)?;
7232                            }
7233                        }
7234                    }
7235                    if n_copy > 0 {
7236                        e.copy_view_into(
7237                            &mut phs,
7238                            dst_off,
7239                            &ph.slice(src_lo..src_lo + n_copy),
7240                            n_copy,
7241                        )?;
7242                    }
7243                    self.mtp_kv_fill(
7244                        e,
7245                        mtp,
7246                        &prompt[start..end],
7247                        &phs,
7248                        base + start,
7249                        &mut *scratch,
7250                        embd_dev,
7251                    )?;
7252                }
7253                start = end;
7254            }
7255        }
7256        // MEMRA_PROFILE_SPEC=2: profiler capture starts HERE — after the prime, so an
7257        // `nsys -c cudaProfilerApi` capture contains ONLY the round loop (draft/verify/commit).
7258        // (=1 brackets the whole call in run_spec.rs, prime included.)
7259        if std::env::var("MEMRA_PROFILE_SPEC").as_deref() == Ok("2") {
7260            unsafe extern "C" {
7261                fn cudaProfilerStart() -> i32;
7262            }
7263            unsafe {
7264                cudaProfilerStart();
7265            }
7266        }
7267        // ROUND-STREAM stage (c) 4 (MEMRA_SPEC_STREAM=1, experimental): pre-issued M-round
7268        // bursts with ZERO per-round host readbacks — the accept/seed/rollback/ring kernels
7269        // consume each other's device outputs; the host drains the ring every M rounds. v1
7270        // constraints: greedy, !spec_replay, single-shot, batched-linear layers, no refresh
7271        // fills (acceptance effect A/B-arbitrated), enters from round 1 (pending guaranteed).
7272        // NOTE: not gated on the caller's graph_draft (its trunk_dense conjunct turns the 35B
7273        // MoE off) — the stream capture encloses ONLY the dense MTP head; the head-dense /
7274        // full-prec / k gates are re-derived here and a failed capture degrades to stream-off.
7275        let stream_on = crate::spec::spec_stream()
7276            && !sampled
7277            && !spec_replay
7278            && constraint.is_none()
7279            && !session_mode
7280            && embd_gpu.is_some()
7281            && !crate::model::full_prec_enabled()
7282            && k + 2 < 96;
7283        let mut stream_graph: Option<cudarc::driver::CudaGraph> = None;
7284        let mut g_tokp2k = e.alloc_u32_zeroed(2 * k.max(1))?;
7285        if stream_on {
7286            let cap = e.capture_graph(|e| {
7287                for j in 0..k.max(1) {
7288                    self.mtp_head_forward_cap(
7289                        e,
7290                        mtp,
7291                        &mut dctx.g_tok,
7292                        &mut dctx.g_pos,
7293                        &mut dctx.g_seed,
7294                        &mut dctx.g_p,
7295                        &mut *scratch,
7296                        true,
7297                        true,
7298                        embd_gpu.expect("round stream requires resident embedding"),
7299                        embd_qt,
7300                        embd_rb,
7301                        d_vocab,
7302                        None,
7303                        Some((&mut g_tokp2k, j, d2t_dev.as_ref())),
7304                        None, // round-stream requires constraint.is_none() (see stream_on)
7305                    )?;
7306                }
7307                Ok(())
7308            });
7309            match cap {
7310                Ok(g) => {
7311                    scratch.set_len(e, 0)?;
7312                    stream_graph = Some(g);
7313                }
7314                Err(err) => {
7315                    scratch.set_len(e, 0)?;
7316                    if debug_spec {
7317                        eprintln!("[spec] stream-graph capture failed ({err}); stream off");
7318                    }
7319                }
7320            }
7321        }
7322        let stream_active = stream_on && stream_graph.is_some();
7323        if debug_spec {
7324            eprintln!(
7325                "[spec] stream_on={stream_on} env={} samp={sampled} dg={} captured={} active={stream_active} session={session_mode} replay={spec_replay}",
7326                crate::spec::spec_stream(),
7327                dctx.graph.is_some(),
7328                stream_graph.is_some()
7329            );
7330        }
7331        let t_v_s = k + 1;
7332        // ROUND-STREAM buffers + ptr tables now live in the model-generic round_stream
7333        // module (extracted 2026-07-12; the gemma burst reuses them).
7334        let sb = crate::round_stream::StreamBufs::new(e, k, crate::spec::spec_stream_m())?;
7335        let crate::round_stream::StreamBufs {
7336            mut vtok_d,
7337            mut brk_d,
7338            mut pend_d,
7339            last_pred_d,
7340            mut pos_ctr,
7341            mut pos_start_d,
7342            mut ring_d,
7343            acc_d: mut stream_acc,
7344            m_rounds,
7345            k: _,
7346        } = sb;
7347        let stream_ptrs: Option<CudaSlice<u64>> = if stream_active {
7348            Some(crate::round_stream::kv_len_ptr_table(
7349                e,
7350                cache,
7351                Some(&pos_ctr),
7352            )?)
7353        } else {
7354            None
7355        };
7356
7357        let t_fill = t_ent.elapsed();
7358        let mut round = 0usize;
7359        // ADAPTIVE DRAFT LENGTH (MEMRA_SPEC_ADAPT=1, opt-in — the gemma_spec accepted-run law,
7360        // ported 2026-08-01): next round's draft depth = last round's accepted run + 1, clamped
7361        // to [floor(pos), k_cap] — a miss shrinks the next draft to the miss point + 1,
7362        // full-accept streaks re-deepen one step per round. NOT the 2026-07-07 acceptance-EMA
7363        // (that arm measured an HONEST LOSS to static per-class optima — 115.0/85.8/73.4 vs
7364        // 121.6/92.7/75.6, EMA lag — and was removed 2026-07-08; rig5090.jsonl has the record).
7365        // The gemma law has no lag class: it reacts within one round, and was worth +7-20% on
7366        // the gemma cells at unchanged exactness (2026-07-10 flip; floor sweep 2026-07-25;
7367        // position key 2026-07-26). Signal = n_acc from the round's EXISTING accept readback —
7368        // zero new syncs; the draft graph is a SINGLE-STEP capture replayed per drafted token,
7369        // so a per-round depth needs no re-capture (unlike gemma's whole-chain graphs). qwen's
7370        // in-round p-min cut already shortens chains mid-round, so gemma's one-round-late p-min
7371        // fold into kc is unnecessary here — the accepted-run law sees the cut via n_acc.
7372        // Exactness is the verify's job at ANY depth (same contract as p-min variable rounds).
7373        // DEFAULT OFF on the qwen path until its cells gate a flip (gemma's is default-on).
7374        // MEASURED 2026-08-01 (H100 GPU-3, interleaved x3, NGEN=256, same-invocation plain
7375        // denominators; research/qwen-adaptive-k-20260801/): REFUTED on the tuned qwen configs.
7376        // q27 K=3+HPOST+PMIN=0.3: short +0.8% (noise; law ~idles, len_hist identical), board
7377        // -1.9%, agentic -0.5%; board PMIN=0 -2.1% (not p-min shadowing — the law itself);
7378        // floor=1 -2.8% (gemma's floor-collapse, reproduced). q35 K=2 board: -6.4% (52/136
7379        // rounds shrink to depth 1; no depth to reclaim at K=2). The gemma direction DOES
7380        // appear at untuned depth-K — q27 K=6 floor=4 +1.5% over fixed K=6 — but stays -3.7%
7381        // below fixed K=3: same verdict class as the retired EMA arm (honest loss to static
7382        // per-class optima). Acceptance-rate rises under the law while tokens/round falls —
7383        // it buys accept-% by adding rounds, and a round's fixed draft+verify cost wins.
7384        // K=1..8 self-consistency PASS both models with the law ON (exactness held).
7385        let adapt = std::env::var("MEMRA_SPEC_ADAPT").as_deref() == Ok("1");
7386        // floor: per-model default keyed on n_embd (gemma's tiering — models with an expensive
7387        // verify keep deep drafts after a miss); MEMRA_SPEC_ADAPT_FLOOR pins it everywhere.
7388        let adapt_floor_env: Option<usize> = std::env::var("MEMRA_SPEC_ADAPT_FLOOR")
7389            .ok()
7390            .and_then(|v| v.parse().ok());
7391        let adapt_floor_default: usize = if self.cfg.n_embd as usize >= 3500 {
7392            4
7393        } else if self.cfg.n_embd as usize >= 2500 {
7394            2
7395        } else {
7396            1
7397        };
7398        let adapt_floor: usize = adapt_floor_env.unwrap_or(adapt_floor_default);
7399        // position key: past floor_ctx a HIGH floor (>=4) relaxes to 1 — forced-deep drafts
7400        // turn net-negative at depth (gemma 31B d1736 evidence); MEMRA_SPEC_FLOOR_CTX moves
7401        // the boundary, an explicit MEMRA_SPEC_ADAPT_FLOOR pins the floor everywhere.
7402        let floor_ctx: usize = std::env::var("MEMRA_SPEC_FLOOR_CTX")
7403            .ok()
7404            .and_then(|v| v.parse().ok())
7405            .unwrap_or(1024);
7406        let floor_at = |pos: usize| -> usize {
7407            if adapt_floor_env.is_some() || pos < floor_ctx {
7408                adapt_floor
7409            } else if adapt_floor >= 4 {
7410                1
7411            } else {
7412                adapt_floor
7413            }
7414        };
7415        // cap: MEMRA_SPEC_CAPMAX (gemma semantics, default 7). Binds only under adapt — the
7416        // fixed-K default path is untouched by this whole block.
7417        let cap_max: usize = std::env::var("MEMRA_SPEC_CAPMAX")
7418            .ok()
7419            .and_then(|v| v.parse().ok())
7420            .unwrap_or(7);
7421        let k_cap = k.min(cap_max).max(1);
7422        let mut kc = k_cap;
7423        let mut opti_fork: Option<OptiForkState> = None;
7424        let mut fork_snapshot: Option<crate::cache::CacheSnapshot> = None;
7425        if fork_mode != OptiForkGateMode::Disabled {
7426            let fence = crate::pp::pp_cuts(self.layers.len());
7427            let refusal = if !session_mode {
7428                Some("not-session")
7429            } else if k != 1 || adapt {
7430                Some("requires-fixed-k1")
7431            } else if sampled || constraint.is_some() || spec_replay {
7432                Some("sampled-constrained-or-replay")
7433            } else if pipe.is_some() {
7434                Some("two-session-pipeline")
7435            } else if !spec_devacc() {
7436                Some("requires-device-accept")
7437            } else if stream_active || crate::spec::spec_stream() {
7438                Some("round-stream")
7439            } else if crate::cache::swa_ring_on() || cache.has_swa_ring() {
7440                Some("swa-ring")
7441            } else if crate::pp::pp_host_bounce_active() {
7442                Some("host-bounce")
7443            } else if fork_mode == OptiForkGateMode::Controller
7444                && cache.recur.iter().any(Option::is_some)
7445            {
7446                Some("controller-requires-zero-recurrent-state")
7447            } else if fence.as_ref().is_none_or(|f| f.len() != 3) {
7448                Some("requires-pp2")
7449            } else {
7450                None
7451            };
7452            if let Some(reason) = refusal {
7453                OPTI_FORK_REFUSALS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
7454                eprintln!("[opti-fork] refused reason={reason}");
7455            } else {
7456                let fence = fence.expect("validated PP-2 fence");
7457                let rt = crate::pp::PpNRt::get(e)?;
7458                let primary_stage0 = rt.engine(0, e).ctx().ordinal() == e.ctx().ordinal();
7459                let primary_stage1 = rt.engine(1, e).ctx().ordinal() == e.ctx().ordinal();
7460                let primary_supported =
7461                    primary_stage0 || (fork_mode == OptiForkGateMode::Controller && primary_stage1);
7462                if !rt.cross_device() || !primary_supported {
7463                    OPTI_FORK_REFUSALS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
7464                    eprintln!("[opti-fork] refused reason=requires-supported-primary-cross-device");
7465                } else {
7466                    // Both recurrent snapshots and both seed generations are allocated before
7467                    // the first fork, each through its owning PP stage. Allocation failure
7468                    // therefore happens before any optimistic state mutation can occur.
7469                    let current_snapshot = opti_snapshot_stage_owned(e, cache, rt, &fence)?;
7470                    let alternate_snapshot = opti_snapshot_stage_owned(e, cache, rt, &fence)?;
7471                    let fork = OptiForkState::new(
7472                        e,
7473                        cache,
7474                        fork_mode,
7475                        alternate_snapshot,
7476                        &h_seed_buf,
7477                        &fill_prev,
7478                        rt,
7479                        fence[1],
7480                        self.layers.len(),
7481                    )?;
7482                    eprintln!(
7483                        "[opti-fork] armed mode={fork_mode:?} snapshots=2 seeds=2 split={} \
7484                         payload_dev0={} payload_dev1={} q_threshold={:.3}",
7485                        fence[1],
7486                        fork.logical_payload_bytes[0],
7487                        fork.logical_payload_bytes[1],
7488                        fork.controller.map_or(0.0, |policy| policy.threshold),
7489                    );
7490                    fork_snapshot = Some(current_snapshot);
7491                    opti_fork = Some(fork);
7492                }
7493            }
7494        }
7495        // Persistent snapshot buffers are allocated once and refreshed in place. The fork arm
7496        // uses stage-owned snapshots; refused/disabled arms retain the existing generic helper.
7497        let mut snap = match fork_snapshot {
7498            Some(snapshot) => snapshot,
7499            None => cache.snapshot(e)?,
7500        };
7501        let mut carried_opti: Option<OptiControllerTicket> = None;
7502        // ROUND-STREAM stage (b) 3a: device table of per-layer kvl.len_d pointers (stable — the
7503        // cache never reallocates len_d; see cache.rs "stable pointer" note). 0 = no KV layer.
7504        let kv_len_ptrs: Option<CudaSlice<u64>> = if spec_devacc() && !spec_replay {
7505            Some(crate::round_stream::kv_len_ptr_table(e, cache, None)?)
7506        } else {
7507            None
7508        };
7509        // BONUS FOLD (2026-07-04): after a FULL accept the bonus token is NOT committed with a
7510        // separate T=1 trunk pass (a full weight read per round). It stays PENDING and rides as
7511        // column 0 of the NEXT round's verify batch. Under predecessor pairing the next chain
7512        // seeds from the bonus's predecessor's TRUE verify hidden (free — no extra
7513        // pass of any kind). Verify still
7514        // checks every emitted token against the target -> exactness holds by construction; only
7515        // DRAFT QUALITY can shift, which the acceptance numbers arbitrate.
7516        // bonus emitted but not yet committed to cache. A carried pending (see SpecSession::
7517        // pending_tok) enters round 0 directly — the burst boundary becomes a plain round edge.
7518        let mut pending: Option<u32> = carried_pending;
7519        // MEMRA_SPEC_PHASE=1: per-round wall decomposition (draft / verify / accept+commit) —
7520        // no tracing, no extra syncs (each phase is naturally sync-bounded: draft readbacks,
7521        // the verify accept readback). Printed once at loop end via spec-stats.
7522        let anatomy_on = std::env::var("MEMRA_SPEC_PP_ANATOMY").as_deref() == Ok("1");
7523        let phase_on = anatomy_on || std::env::var("MEMRA_SPEC_PHASE").as_deref() == Ok("1");
7524        // DRAFT-MASK receipt (lane/draft-mask): speculative-clone wall + rounds, printed with
7525        // spec-stats. The clone is the one cost the design adds per round — measured, not assumed.
7526        let (mut dm_clone_ns, mut dm_rounds) = (0u128, 0usize);
7527        // grammar-truncation counters: how many rounds the verify-side cut fired and how many
7528        // already-verified tokens it threw away. THIS is the quantity draft masking targets.
7529        let (mut dm_cuts, mut dm_cut_tokens) = (0usize, 0usize);
7530        let (mut ph_draft, mut ph_verify, mut ph_rest) = (0f64, 0f64, 0f64);
7531        let mut ph_wait = 0f64;
7532        let mut ph_commit = 0f64;
7533        let mut ph_t = std::time::Instant::now();
7534        let mut ph_mark = |acc: &mut f64, on: bool| {
7535            if on {
7536                let now = std::time::Instant::now();
7537                *acc += (now - ph_t).as_secs_f64();
7538                ph_t = now;
7539            }
7540        };
7541        if let Some(p) = pipe {
7542            p.setup_end();
7543        }
7544        while keep_going && out.len() < max_new {
7545            // ROUND-STREAM BURST: from round 1 (pending guaranteed by every non-replay arm),
7546            // issue M rounds with zero readbacks, then drain the ring + reconcile mirrors.
7547            if let (true, Some(sg), Some(ptrs)) = (
7548                stream_active && round >= 1 && pending.is_some(),
7549                &stream_graph,
7550                &stream_ptrs,
7551            ) {
7552                if debug_spec {
7553                    static ONCE: std::sync::Once = std::sync::Once::new();
7554                    ONCE.call_once(|| {
7555                        eprintln!("[memra] ROUND-STREAM burst engaged (M={m_rounds} k={k})")
7556                    });
7557                }
7558                e.set_i32_one(&mut pos_ctr, cache.pos as i32)?;
7559                e.set_u32_one(&mut pend_d, pending.unwrap())?;
7560                e.set_u32_one(&mut ring_d, 0)?; // ring count = 0 (writes element 0)
7561                for _mi in 0..m_rounds {
7562                    e.i32_copy_add(&pos_ctr, &mut pos_start_d, 0)?;
7563                    cache.snapshot_into(e, &mut snap)?; // device D2Ds, stream-ordered
7564                    e.i32_copy_add(&pos_ctr, &mut scratch.kv.len_d, 0)?; // draft-KV rollback
7565                    e.i32_copy_add(&pos_ctr, &mut dctx.g_pos, 1)?; // rope pos = pos + base
7566                    e.u32_copy(&pend_d, &mut dctx.g_tok)?;
7567                    e.copy_into(&mut dctx.g_seed, 0, &h_seed_buf, n_embd)?;
7568                    sg.launch()?;
7569                    e.spec_assemble_verify(
7570                        &g_tokp2k,
7571                        &pend_d,
7572                        d2t_dev.as_ref(),
7573                        &mut vtok_d,
7574                        &mut brk_d,
7575                        p_min,
7576                        k,
7577                        pmin0,
7578                    )?;
7579                    let mut ck = VerifyCkpt::new(self.layers.len());
7580                    let dummy = vec![0u32; t_v_s];
7581                    let (tl_d, vx) = self.decode_step_t_core_stream(
7582                        e,
7583                        &dummy,
7584                        0,
7585                        &mut *cache,
7586                        embd_dev,
7587                        Some(&mut ck),
7588                        Some((&vtok_d, &pos_ctr)),
7589                        None,
7590                    )?;
7591                    for j in 0..t_v_s {
7592                        e.argmax_token_device_col(&tl_d, j, n_vocab, &mut preds_d, j)?;
7593                    }
7594                    e.spec_accept_greedy_dc(
7595                        &preds_d,
7596                        &vtok_d,
7597                        &last_pred_d,
7598                        &brk_d,
7599                        &mut stream_acc,
7600                    )?;
7601                    e.spec_seed_gather(&vx, &fill_prev, &stream_acc, &mut h_seed_buf, 1, n_embd)?;
7602                    e.copy_into(&mut fill_prev, 0, &h_seed_buf, n_embd)?;
7603                    self.commit_verified_prefix_stream(
7604                        e,
7605                        &mut *cache,
7606                        &snap,
7607                        &ck,
7608                        &stream_acc,
7609                        1,
7610                        t_v_s,
7611                    )?;
7612                    e.spec_rollback_stream(
7613                        ptrs,
7614                        &pos_start_d,
7615                        &stream_acc,
7616                        1,
7617                        self.layers.len() + 1,
7618                    )?;
7619                    e.spec_ring_commit(&vtok_d, &stream_acc, &brk_d, &mut ring_d, &mut pend_d)?;
7620                }
7621                e.stream().synchronize()?;
7622                let ring_h = e.dtoh_u32(&ring_d)?;
7623                let cnt = ring_h[0] as usize;
7624                for i in 0..cnt {
7625                    if out.len() < max_new {
7626                        out.push(ring_h[1 + i]);
7627                    }
7628                }
7629                let pos_h = e.dtoh_i32(&pos_ctr)?[0] as usize;
7630                for il in 0..self.layers.len() {
7631                    if let Some(kvl) = cache.kv[il].as_mut() {
7632                        kvl.len = pos_h;
7633                    }
7634                }
7635                cache.pos = pos_h;
7636                scratch.kv.len = pos_h;
7637                pending = Some(ring_h[cnt]); // last drained token = the live bonus
7638                last_token = ring_h[cnt];
7639                total_drafted += k * m_rounds; // upper bound (p-min breaks uncounted)
7640                total_accepted += cnt.saturating_sub(m_rounds);
7641                if let Some(t) = sess_telem {
7642                    // totals only — the burst's per-round accept counts stayed on device
7643                    // (that is the point of the round-stream arm). pos_* untouched.
7644                    t.record_totals(m_rounds, k * m_rounds, cnt.saturating_sub(m_rounds));
7645                }
7646                round += m_rounds;
7647                // sse-cadence: the drained ring is committed — flush it at burst-drain cadence.
7648                keep_going = flush_commit(&mut on_commit, &out, &mut flushed);
7649                continue;
7650            }
7651            let pipe_draft = match pipe {
7652                Some(p) => Some(p.draft_begin(round)?),
7653                None => None,
7654            };
7655            let pos = cache.pos; // #tokens committed (EXCLUDES a pending bonus)
7656            let mut current_opti = carried_opti.take();
7657            let mut fork_generation = if current_opti.is_none() && pending.is_some() {
7658                match opti_fork.as_mut() {
7659                    Some(fork) if fork.mode.is_forced() => Some(fork.reserve(&mut snap)?),
7660                    None => None,
7661                    Some(_) => None,
7662                }
7663            } else {
7664                None
7665            };
7666            if current_opti.is_none() {
7667                if let Some(fork) = opti_fork.as_ref() {
7668                    opti_snapshot_stage_owned_into(e, cache, fork.rt, &fork.fence, &mut snap)?;
7669                } else {
7670                    cache.snapshot_into(e, &mut snap)?;
7671                }
7672            } else if snap.pos != pos {
7673                return Err(format!(
7674                    "optipipe carried snapshot pos {} != current pos {pos}",
7675                    snap.pos
7676                )
7677                .into());
7678            } // §C: snapshot BEFORE draft+verify (already retained for a carried successor)
7679            ph_mark(&mut ph_rest, phase_on);
7680
7681            // --- 1. DRAFT k tokens with the NextN head (autoregressive, T=1 each) ---
7682            // p-min semantics (both paths): stop the chain early when the head's confidence in
7683            // its own pick drops below p_min — the just-drafted token is DISCARDED, but its
7684            // scratch append stands (identical to the eager chain's ordering). j==0 always drafts.
7685            let base0 = if pending.is_some() { 1usize } else { 0usize };
7686            // fixed draft length by default; MEMRA_SPEC_ADAPT=1 drafts at last round's
7687            // accepted run + 1 (the gemma law — see the setup block above the loop).
7688            let k_this = if adapt { kc } else { k };
7689            let mut draft: Vec<u32> = Vec::with_capacity(k);
7690            let mut draft_idx: Vec<u32> = Vec::with_capacity(k); // trimmed-vocab ids (== draft when untrimmed)
7691            let mut controller_draft_prob: Option<f32> = None;
7692            let mut controller_eager_state: Option<(u32, CudaSlice<f32>)> = None;
7693            if let Some(ticket) = current_opti.as_mut() {
7694                let carried_pending = pending.ok_or("optipipe carried successor lost pending")?;
7695                if ticket.verify_tokens[0] != carried_pending {
7696                    return Err(format!(
7697                        "optipipe carried pending mismatch: ticket={} live={carried_pending}",
7698                        ticket.verify_tokens[0],
7699                    )
7700                    .into());
7701                }
7702                draft.push(ticket.verify_tokens[1]);
7703                controller_draft_prob = Some(ticket.draft_prob);
7704                controller_eager_state = ticket
7705                    .take_eager_seed()
7706                    .map(|seed| (ticket.verify_tokens[1], seed));
7707            } else {
7708                // Round-start draft-KV sync (BOTH paths). Persistent: truncate/align to the committed
7709                // history — slots 0..P hold entries for the tokens before last_token@P (P = pos +
7710                // base0 - 1); this single set_len IS the draft-side rollback (drops last round's
7711                // rejected drafts and p-min extras via the len mechanism).
7712                scratch.set_len(e, pos + base0 - 1)?;
7713                if pen_on {
7714                    let w0 = pen_hist.len().saturating_sub(sp.penalty_last_n);
7715                    pen_hist_d = Some(e.htod_u32_v(&pen_hist[w0..])?);
7716                }
7717                if sampled {
7718                    draft_logits.clear();
7719                    draft_stats.clear();
7720                }
7721                // DRAFT-SIDE GRAMMAR MASK: clone the committed grammar state ONCE per round; each
7722                // position's mask is computed on that clone and advanced by the PROPOSED token. The
7723                // real state moves only on emission (verify's job), so the emitted stream is
7724                // unchanged — the mask only removes tokens the verify would have truncated anyway.
7725                let mut dmask_live = dmask_on;
7726                if dmask_live {
7727                    let t_c = std::time::Instant::now();
7728                    constraint
7729                        .as_deref_mut()
7730                        .unwrap()
7731                        .draft_begin()
7732                        .map_err(|e2| format!("constraint: {e2}"))?;
7733                    dm_clone_ns += t_c.elapsed().as_nanos();
7734                    dm_rounds += 1;
7735                }
7736                if let (false, Some(gr)) = (sampled || pen_on, &dctx.graph) {
7737                    // GRAPH DRAFT: one dispatch per drafted token. The chain feeds itself on-device
7738                    // (in-graph argmax -> tok_d -> next replay's embed; h_nextn -> h_seed_d; pos_d
7739                    // inc'd in-graph); the host only reads 4B token (+4B p) and decides the break.
7740                    e.set_i32_one(&mut dctx.g_pos, (pos + base0) as i32)?;
7741                    e.set_u32_one(&mut dctx.g_tok, last_token)?;
7742                    e.copy_into(&mut dctx.g_seed, 0, &h_seed_buf, n_embd)?;
7743                    for j in 0..k_this {
7744                        // per-position mask upload (contents only — the graph's baked pointer is
7745                        // dctx.g_dmask). All-ones once masking goes dead mid-chain, so the captured
7746                        // mask node degrades to a no-op ban instead of needing a second graph.
7747                        if dmask_live
7748                            && !upload_draft_mask(
7749                                e,
7750                                constraint.as_deref_mut().unwrap(),
7751                                &mut dctx.g_dmask,
7752                                mtp.d2t.as_ref(),
7753                                d_vocab,
7754                                dmask_words,
7755                            )?
7756                        {
7757                            // no draft-vocab row is grammar-legal here (a trimmed FR-Spec head can
7758                            // genuinely miss the legal set): neutralize the captured mask node and
7759                            // finish the chain UNMASKED — exactly pre-lane behaviour, never worse.
7760                            e.htod_u32_into(&mut dctx.g_dmask, &vec![u32::MAX; dmask_words])?;
7761                            dmask_live = false;
7762                        }
7763                        gr.launch()?;
7764                        scratch.kv.len += 1; // host mirror (len_d advanced in-graph)
7765                        let idx = e.dtoh_u32_one(&dctx.g_tok)?;
7766                        // #87 SENTINEL TRAP: an all-NaN head-logits row leaves the device argmax's
7767                        // init sentinel (0x7FFFFFFF) in g_tok — feeding it onward dereferences
7768                        // embed_row(sentinel) = table + ~4.6TB (never mapped) inside the NEXT graph
7769                        // replay's embed node, and the MMU fault kills the CUDA context for the
7770                        // whole process (research/pp2spec-crash-20260807: 3 coredumps, byte-exact
7771                        // VA arithmetic). Refuse loudly instead; the diagnostics name the first-NaN
7772                        // buffer (g_seed = the verify-side handoff vs head-side compute).
7773                        if (idx as usize) >= d_vocab {
7774                            // g_seed is SELF-FED (the replay writes h_nextn back into it), so it
7775                            // reads as the head's OUTPUT at j; h_seed_buf is the round's INPUT
7776                            // seed, untouched since the round-start copy — the pair discriminates
7777                            // "seed arrived poisoned" from "head forward produced NaN".
7778                            let seed_h = e.dtoh(&dctx.g_seed)?;
7779                            let seed_nan = seed_h.iter().filter(|v| v.is_nan()).count();
7780                            let in_h = e.dtoh(&h_seed_buf)?;
7781                            let in_nan = in_h.iter().filter(|v| v.is_nan()).count();
7782                            return Err(format!(
7783                                "draft(graph) argmax sentinel 0x{idx:08x} >= d_vocab {d_vocab} at \
7784                             round {round} j={j} pos={pos}: head-out NaN {seed_nan}/{n_embd}, \
7785                             round-input-seed NaN {in_nan}/{n_embd} — refusing to dereference \
7786                             the embed row (#87 trap)"
7787                            )
7788                            .into());
7789                        }
7790                        // trimmed draft vocab -> target token id (identity when no d2t map)
7791                        let d = match &mtp.d2t {
7792                            Some(map) => map[idx as usize],
7793                            None => idx,
7794                        };
7795                        let draft_p = if p_min > 0.0
7796                            || opti_fork
7797                                .as_ref()
7798                                .is_some_and(|fork| fork.controller.is_some())
7799                        {
7800                            Some(e.dtoh(&dctx.g_p)?[0])
7801                        } else {
7802                            None
7803                        };
7804                        if j == 0 {
7805                            controller_draft_prob = draft_p;
7806                        }
7807                        if let Some(p) = draft_p.filter(|_| p_min > 0.0) {
7808                            if p < p_min && (j > 0 || (pmin0 && base0 == 1)) {
7809                                break;
7810                            }
7811                        }
7812                        draft.push(d);
7813                        // with a trimmed head the NEXT embed must read the TARGET id, not the draft
7814                        // index the argmax wrote — patch the persistent token buffer (4B htod).
7815                        if d != idx {
7816                            e.set_u32_one(&mut dctx.g_tok, d)?;
7817                        }
7818                        // advance the SPECULATIVE state with the proposal; a dead chain drops to
7819                        // unmasked drafting for the remaining positions (verify still arbitrates).
7820                        // speculative advance; a chain the grammar can no longer follow (EOS
7821                        // proposed) ends here. The captured mask node always runs, so a dead chain
7822                        // leaves the buffer NEUTRAL (all-ones = ban nothing) before it exits.
7823                        if dmask_live
7824                            && !constraint
7825                                .as_deref_mut()
7826                                .unwrap()
7827                                .draft_advance(d)
7828                                .map_err(|e2| format!("constraint: {e2}"))?
7829                        {
7830                            e.htod_u32_into(&mut dctx.g_dmask, &vec![u32::MAX; dmask_words])?;
7831                            break;
7832                        }
7833                    }
7834                } else if let (true, Some(gr)) = (sampled, &dctx.graph_s) {
7835                    // SAMPLED GRAPH DRAFT: one replay per drafted token — head forward + gumbel +
7836                    // argmax in ONE dispatch; the host reads 4B token (+4B p), D2Ds q into slot j,
7837                    // and decides the break. Event-counter continuity: g_ctr is host-seeded to
7838                    // sctr-1 ONCE per round (outside the graph); the in-graph bump runs BEFORE the
7839                    // perturb, so replay j consumes counter sctr+j — exactly the eager arm's Philox
7840                    // stream. Host sctr advances in lockstep (computed, no readback needed).
7841                    e.set_i32_one(&mut dctx.g_pos, (pos + base0) as i32)?;
7842                    e.set_u32_one(&mut dctx.g_tok, last_token)?;
7843                    e.copy_into(&mut dctx.g_seed, 0, &h_seed_buf, n_embd)?;
7844                    e.set_u32_one(&mut dctx.g_ctr, sctr.wrapping_sub(1))?;
7845                    for j in 0..k_this {
7846                        gr.launch()?;
7847                        scratch.kv.len += 1; // host mirror (len_d advanced in-graph)
7848                        sctr += 1; // mirrors the in-graph g_ctr bump (eager parity:
7849                        // counts the p-min-discarded token too)
7850                        // q retention: ONE async D2D of the persistent head-logits buffer into this
7851                        // round's slot j (stream-ordered after the replay, before the next one).
7852                        e.copy_into(&mut dctx.q_slots[j], 0, &dctx.g_q, d_vocab)?;
7853                        let idx = e.dtoh_u32_one(&dctx.g_tok)?;
7854                        // #87 SENTINEL TRAP (see the greedy graph arm above).
7855                        if (idx as usize) >= d_vocab {
7856                            let seed_h = e.dtoh(&dctx.g_seed)?;
7857                            let seed_nan = seed_h.iter().filter(|v| v.is_nan()).count();
7858                            return Err(format!(
7859                                "draft(graph-sampled) argmax sentinel 0x{idx:08x} >= d_vocab \
7860                             {d_vocab} at round {round} j={j} pos={pos}: round-seed NaN \
7861                             {seed_nan}/{n_embd} — refusing to dereference the embed row \
7862                             (#87 trap)"
7863                            )
7864                            .into());
7865                        }
7866                        let d = match &mtp.d2t {
7867                            Some(map) => map[idx as usize],
7868                            None => idx,
7869                        };
7870                        draft_idx.push(idx);
7871                        if p_min > 0.0 {
7872                            let p = e.dtoh(&dctx.g_p)?[0];
7873                            if p < p_min && (j > 0 || (pmin0 && base0 == 1)) {
7874                                break;
7875                            }
7876                        }
7877                        draft.push(d);
7878                        // trimmed head: the NEXT embed must read the TARGET id (see the greedy arm).
7879                        if d != idx {
7880                            e.set_u32_one(&mut dctx.g_tok, d)?;
7881                        }
7882                    }
7883                    // uniform accept path: fill draft_stats per used slot (pure-temp regime — the
7884                    // stats degenerate to th=0 / full-Z; one filter_stats launch per slot, tiny).
7885                    for j in 0..draft.len().max(draft_idx.len()) {
7886                        let rows0 = e.htod_i32(&[0])?;
7887                        let (mut th_d, mut z_d, mut mx_d) = (e.zeros(1)?, e.zeros(1)?, e.zeros(1)?);
7888                        e.filter_stats(
7889                            &dctx.q_slots[j],
7890                            d_vocab,
7891                            &rows0,
7892                            &mut th_d,
7893                            &mut z_d,
7894                            &mut mx_d,
7895                            d_vocab,
7896                            1,
7897                            sp_temp,
7898                            sp.top_k,
7899                            sp.top_p,
7900                            sp.min_p,
7901                        )?;
7902                        draft_stats.push((e.dtoh(&mx_d)?[0], e.dtoh(&th_d)?[0], e.dtoh(&z_d)?[0]));
7903                    }
7904                } else {
7905                    // EAGER DRAFT (fallback: MoE head/trunk, huge k, MEMRA_SPEC_NOGRAPH, capture fail).
7906                    let mut e_tok = last_token;
7907                    let mut d_seed = e.clone_dtod(&h_seed_buf)?;
7908                    for j in 0..k_this {
7909                        // GPU-ARGMAX DRAFT (2026-07-03): device logits + device argmax + 4-byte token
7910                        // read instead of the ~600KB full-vocab dtoh + host argmax per draft token.
7911                        let mtp_pos = pos + base0 + j;
7912                        // draft-side grammar mask (eager twin of the graph arm's in-graph node).
7913                        // A position with no legal draft-vocab row drops to unmasked drafting for
7914                        // the rest of the chain (pre-lane behaviour; verify still arbitrates).
7915                        if dmask_live {
7916                            dmask_live = upload_draft_mask(
7917                                e,
7918                                constraint.as_deref_mut().unwrap(),
7919                                &mut dctx.g_dmask,
7920                                mtp.d2t.as_ref(),
7921                                d_vocab,
7922                                dmask_words,
7923                            )?;
7924                        }
7925                        let (dl_d, h_nextn) = self.mtp_head_forward_dev(
7926                            e,
7927                            mtp,
7928                            e_tok,
7929                            &d_seed,
7930                            &mut *scratch,
7931                            mtp_pos,
7932                            embd_dev,
7933                            if dmask_live {
7934                                Some((&dctx.g_dmask, dmask_words))
7935                            } else {
7936                                None
7937                            },
7938                        )?;
7939                        let tok_d = if sampled {
7940                            // FILTERED Gumbel-max: stats -> masked perturb -> argmax = one draw from
7941                            // the filtered softmax (filters off => th=0, exact v1 semantics).
7942                            if perturb_buf.is_none() {
7943                                perturb_buf = Some(e.zeros(d_vocab.max(n_vocab))?);
7944                            }
7945                            let mut q_row = e.clone_dtod(&dl_d)?; // retained q (penalized when on)
7946                            if pen_on {
7947                                let h = pen_hist_d.as_ref().unwrap();
7948                                let nh = h.len();
7949                                e.penalize_logits(
7950                                    &mut q_row,
7951                                    h,
7952                                    nh,
7953                                    sp.penalty_repeat,
7954                                    sp.penalty_freq,
7955                                    sp.penalty_present,
7956                                    d_vocab,
7957                                )?;
7958                            }
7959                            let rows0 = e.htod_i32(&[0])?;
7960                            let (mut th_d, mut z_d, mut mx_d) =
7961                                (e.zeros(1)?, e.zeros(1)?, e.zeros(1)?);
7962                            e.filter_stats(
7963                                &q_row, d_vocab, &rows0, &mut th_d, &mut z_d, &mut mx_d, d_vocab,
7964                                1, sp_temp, sp.top_k, sp.top_p, sp.min_p,
7965                            )?;
7966                            let (th, z, mx) =
7967                                (e.dtoh(&th_d)?[0], e.dtoh(&z_d)?[0], e.dtoh(&mx_d)?[0]);
7968                            let pb = perturb_buf.as_mut().unwrap();
7969                            e.gumbel_perturb_filtered(
7970                                &q_row, pb, d_vocab, sp_seed, sctr, sp_temp, mx, th,
7971                            )?;
7972                            sctr += 1;
7973                            draft_logits.push(q_row);
7974                            draft_stats.push((mx, th, z));
7975                            e.argmax_token_device(pb, d_vocab)?
7976                        } else {
7977                            e.argmax_token_device(&dl_d, d_vocab)?
7978                        };
7979                        let idx = e.dtoh_u32_one(&tok_d)?;
7980                        // #87 SENTINEL TRAP (eager twin — see the graph arm). Extra diagnostics
7981                        // here because the eager chain's operands are all readable: dl_d (the head
7982                        // logits row) and d_seed (this step's h_seed) name the first-NaN buffer.
7983                        if (idx as usize) >= d_vocab {
7984                            let dl_h = e.dtoh(&dl_d)?;
7985                            let dl_nan = dl_h.iter().filter(|v| v.is_nan()).count();
7986                            let seed_h = e.dtoh(&d_seed)?;
7987                            let seed_nan = seed_h.iter().filter(|v| v.is_nan()).count();
7988                            return Err(format!(
7989                                "draft(eager) argmax sentinel 0x{idx:08x} >= d_vocab {d_vocab} at \
7990                             round {round} j={j} pos={pos}: head-logits NaN {dl_nan}/{d_vocab}, \
7991                             step-seed NaN {seed_nan}/{n_embd} — refusing to dereference the \
7992                             embed row (#87 trap)"
7993                            )
7994                            .into());
7995                        }
7996                        let d = match &mtp.d2t {
7997                            Some(map) => map[idx as usize],
7998                            None => idx,
7999                        };
8000                        if sampled {
8001                            draft_idx.push(idx);
8002                        }
8003                        let draft_p = if p_min > 0.0
8004                            || opti_fork
8005                                .as_ref()
8006                                .is_some_and(|fork| fork.controller.is_some())
8007                        {
8008                            let p_d = e.prob_of_token_device(&dl_d, &tok_d, d_vocab)?;
8009                            Some(e.dtoh(&p_d)?[0])
8010                        } else {
8011                            None
8012                        };
8013                        if j == 0 {
8014                            controller_draft_prob = draft_p;
8015                        }
8016                        if let Some(p) = draft_p.filter(|_| p_min > 0.0) {
8017                            if p < p_min && (j > 0 || (pmin0 && base0 == 1)) {
8018                                break;
8019                            }
8020                        }
8021                        draft.push(d);
8022                        e_tok = d;
8023                        d_seed = h_nextn;
8024                        // speculative advance; a chain the grammar can no longer follow (EOS
8025                        // proposed) ends here — the prefix already proposed still rides verify.
8026                        if dmask_live
8027                            && !constraint
8028                                .as_deref_mut()
8029                                .unwrap()
8030                                .draft_advance(d)
8031                                .map_err(|e2| format!("constraint: {e2}"))?
8032                        {
8033                            break;
8034                        }
8035                    }
8036                    if opti_fork
8037                        .as_ref()
8038                        .is_some_and(|fork| fork.controller.is_some())
8039                    {
8040                        controller_eager_state = Some((e_tok, d_seed));
8041                    }
8042                }
8043            }
8044            let k_round = draft.len();
8045            if let Some(p) = pipe {
8046                p.draft_end(round);
8047            }
8048            drop(pipe_draft);
8049
8050            ph_mark(&mut ph_draft, phase_on);
8051            // --- 2. VERIFY: one batched target forward. With a pending bonus, it rides as col 0
8052            //         (committing its KV/recur inside the SAME weight read); drafts follow. ---
8053            let verify_tokens: Vec<u32> = match pending {
8054                Some(b) => {
8055                    let mut v = Vec::with_capacity(k_round + 1);
8056                    v.push(b);
8057                    v.extend_from_slice(&draft);
8058                    v
8059                }
8060                None => draft.clone(),
8061            };
8062            let base = if pending.is_some() { 1 } else { 0 };
8063            // ckpt (REPLAY-FREE partial accept): retain per-layer state-rebuild inputs alongside
8064            // the verify. Pure buffer keep-alives + dtod clones — kernel work is unchanged.
8065            let mut ckpt = if let Some(ticket) = current_opti.as_mut() {
8066                Some(ticket.take_ckpt())
8067            } else if spec_replay {
8068                None
8069            } else {
8070                Some(VerifyCkpt::new(self.layers.len()))
8071            };
8072            let controller_can_probe = base == 1
8073                && k_round == 1
8074                && out.len().saturating_add(2) < max_new
8075                && controller_draft_prob.is_some()
8076                && opti_fork
8077                    .as_ref()
8078                    .and_then(|fork| fork.controller.as_ref())
8079                    .is_some_and(|policy| !policy.breaker_tripped);
8080            let mut successor_attempt: Option<OptiControllerTicket> = None;
8081            let mut rejected_probe: Option<(f32, u32)> = None;
8082            let mut controller_prepared: Option<OptiControllerPrepared> = None;
8083            if controller_can_probe {
8084                // Prepare d2/q and, on admission, d3 before either current verify half is
8085                // issued. N stage 0 can then be followed immediately by N+1 stage 0; once N's
8086                // boundary fires, those dev0 launches overlap N stage 1 on dev1. Preparing on
8087                // the primary stream after N stage 1 would serialize the supposed pipeline.
8088                let eager_pos = scratch.kv.len + 1;
8089                let (optimistic_pending, pending_probability) = self.opti_controller_draft_step(
8090                    e,
8091                    mtp,
8092                    &mut dctx,
8093                    &mut *scratch,
8094                    d_vocab,
8095                    &mut controller_eager_state,
8096                    eager_pos,
8097                    embd_dev,
8098                )?;
8099                let first_probability = controller_draft_prob
8100                    .ok_or("optipipe controller probe lost first-token probability")?;
8101                let q_proxy = first_probability * pending_probability;
8102                OPTI_GATE_CHECKS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
8103                OPTI_SHADOW_DRAFT_TOKENS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
8104                let admitted = opti_fork
8105                    .as_ref()
8106                    .and_then(|fork| fork.controller.as_ref())
8107                    .ok_or("optipipe controller policy disappeared")?
8108                    .admit(q_proxy);
8109                if admitted {
8110                    OPTI_GATE_ADMITS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
8111                    let eager_pos = scratch.kv.len + 1;
8112                    let (optimistic_draft, optimistic_draft_probability) = self
8113                        .opti_controller_draft_step(
8114                            e,
8115                            mtp,
8116                            &mut dctx,
8117                            &mut *scratch,
8118                            d_vocab,
8119                            &mut controller_eager_state,
8120                            eager_pos,
8121                            embd_dev,
8122                        )?;
8123                    OPTI_SHADOW_DRAFT_TOKENS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
8124                    let eager_seed = controller_eager_state.take().map(|(token, seed)| {
8125                        debug_assert_eq!(token, optimistic_draft);
8126                        seed
8127                    });
8128                    controller_prepared = Some(OptiControllerPrepared {
8129                        verify_tokens: [optimistic_pending, optimistic_draft],
8130                        draft_prob: optimistic_draft_probability,
8131                        eager_seed,
8132                        q_proxy,
8133                        scratch_len: scratch.kv.len,
8134                    });
8135                } else {
8136                    OPTI_GATE_REJECTS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
8137                    OPTI_WASTED_DRAFT_TOKENS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
8138                    rejected_probe = Some((q_proxy, optimistic_pending));
8139                    eprintln!(
8140                        "[opti-controller] reject q={q_proxy:.6} threshold={:.3}",
8141                        opti_fork
8142                            .as_ref()
8143                            .and_then(|fork| fork.controller.as_ref())
8144                            .expect("controller policy")
8145                            .threshold,
8146                    );
8147                }
8148            }
8149            let fork_attempt = match fork_generation.take() {
8150                Some(generation) if base == 1 && k_round == 1 => Some(generation),
8151                Some(generation) => {
8152                    opti_fork
8153                        .as_mut()
8154                        .expect("fork generation without fork state")
8155                        .retire(generation)?;
8156                    None
8157                }
8158                None => None,
8159            };
8160            let (tlogits_d, vx) = if let Some(p) = pipe {
8161                self.decode_step_t_core_pipelined(
8162                    e,
8163                    &verify_tokens,
8164                    pos,
8165                    &mut *cache,
8166                    embd_dev,
8167                    ckpt.as_mut(),
8168                    p,
8169                    round,
8170                )?
8171            } else if controller_can_probe {
8172                let fence = opti_fork
8173                    .as_ref()
8174                    .ok_or("optipipe controller probe lost fork state")?
8175                    .fence;
8176                let boundary = match current_opti.as_mut() {
8177                    Some(ticket) => ticket.take_boundary(),
8178                    None => self.verify_stage0_issue(
8179                        e,
8180                        &verify_tokens,
8181                        pos,
8182                        &mut *cache,
8183                        embd_dev,
8184                        ckpt.as_mut(),
8185                        None,
8186                        &fence,
8187                        Some(true),
8188                        None,
8189                    )?,
8190                };
8191                if let Some(prepared) = controller_prepared.take() {
8192                    let generation = {
8193                        let fork = opti_fork
8194                            .as_mut()
8195                            .ok_or("optipipe controller admission lost fork state")?;
8196                        let generation = fork.reserve_successor()?;
8197                        let rt = fork.rt;
8198                        let snapshot_fence = fork.fence;
8199                        opti_snapshot_one_stage_owned_into(
8200                            e,
8201                            cache,
8202                            rt,
8203                            &snapshot_fence,
8204                            0,
8205                            fork.successor_snapshot_mut(),
8206                        )?;
8207                        generation
8208                    };
8209                    let mut successor_ckpt = VerifyCkpt::new(self.layers.len());
8210                    let successor_boundary = self.verify_stage0_issue(
8211                        e,
8212                        &prepared.verify_tokens,
8213                        pos + verify_tokens.len(),
8214                        &mut *cache,
8215                        embd_dev,
8216                        Some(&mut successor_ckpt),
8217                        None,
8218                        &fence,
8219                        Some(false),
8220                        None,
8221                    )?;
8222                    OPTI_FORK_ATTEMPTS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
8223                    let fork = opti_fork
8224                        .as_ref()
8225                        .ok_or("optipipe controller ticket lost fork state")?;
8226                    successor_attempt = Some(fork.controller_ticket(
8227                        generation,
8228                        successor_boundary,
8229                        successor_ckpt,
8230                        prepared.verify_tokens,
8231                        prepared.draft_prob,
8232                        prepared.eager_seed,
8233                        prepared.q_proxy,
8234                        prepared.scratch_len,
8235                    ));
8236                    eprintln!(
8237                        "[opti-controller] issue generation={} q={:.6} threshold={:.3} \
8238                         verify={:?}",
8239                        generation.id,
8240                        prepared.q_proxy,
8241                        fork.controller.expect("controller policy").threshold,
8242                        prepared.verify_tokens,
8243                    );
8244                }
8245                let result = self.verify_stage1_finish(
8246                    e,
8247                    boundary,
8248                    &mut *cache,
8249                    ckpt.as_mut(),
8250                    None,
8251                    &fence,
8252                    successor_attempt.is_none(),
8253                )?;
8254                if let Some(ticket) = current_opti.as_mut() {
8255                    ticket.settle();
8256                }
8257                if successor_attempt.is_some() {
8258                    let fork = opti_fork
8259                        .as_mut()
8260                        .ok_or("optipipe successor snapshot lost fork state")?;
8261                    let rt = fork.rt;
8262                    let snapshot_fence = fork.fence;
8263                    opti_snapshot_one_stage_owned_into(
8264                        e,
8265                        cache,
8266                        rt,
8267                        &snapshot_fence,
8268                        1,
8269                        fork.successor_snapshot_mut(),
8270                    )?;
8271                    // Publish N only after both independent successor-state queues are complete.
8272                    fork.rt.publish_to(1, &e.stream())?;
8273                }
8274                result
8275            } else if let Some(ticket) = current_opti.as_mut() {
8276                let fork = opti_fork
8277                    .as_mut()
8278                    .ok_or("optipipe carried controller ticket lost fork state")?;
8279                let boundary = ticket.take_boundary();
8280                let result = self.verify_stage1_finish(
8281                    e,
8282                    boundary,
8283                    &mut *cache,
8284                    ckpt.as_mut(),
8285                    None,
8286                    &fork.fence,
8287                    true,
8288                )?;
8289                ticket.settle();
8290                result
8291            } else if let Some(generation) = fork_attempt {
8292                let fork = opti_fork
8293                    .as_mut()
8294                    .expect("fork generation without fork state");
8295                fork.capture_seed(e, generation, &h_seed_buf, &fill_prev, scratch.kv.len)?;
8296                let action = fork.mode.action(generation.id);
8297                let boundary = self.verify_stage0_issue(
8298                    e,
8299                    &verify_tokens,
8300                    pos,
8301                    &mut *cache,
8302                    embd_dev,
8303                    ckpt.as_mut(),
8304                    None,
8305                    &fork.fence,
8306                    Some(true),
8307                    None,
8308                )?;
8309                OPTI_FORK_ATTEMPTS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
8310                let mut ticket = fork.ticket(generation, boundary);
8311                if action == OptiForkAction::Abort {
8312                    return Err(format!(
8313                        "optipipe forced abort with generation {} stage0 in flight",
8314                        generation.id,
8315                    )
8316                    .into());
8317                }
8318                fork.reconcile(
8319                    e,
8320                    &mut *cache,
8321                    &mut *scratch,
8322                    &snap,
8323                    &mut h_seed_buf,
8324                    &mut fill_prev,
8325                    generation,
8326                    action,
8327                    verify_tokens[0],
8328                )?;
8329                let result = if action == OptiForkAction::Hit {
8330                    let boundary = ticket.take_boundary();
8331                    self.verify_stage1_finish(
8332                        e,
8333                        boundary,
8334                        &mut *cache,
8335                        ckpt.as_mut(),
8336                        None,
8337                        &fork.fence,
8338                        true,
8339                    )?
8340                } else {
8341                    // The optimistic boundary slot has no reader. Re-run the unchanged serial
8342                    // verify only after E_restart published the restored stage-0 state.
8343                    self.decode_step_t_core(
8344                        e,
8345                        &verify_tokens,
8346                        pos,
8347                        &mut *cache,
8348                        embd_dev,
8349                        ckpt.as_mut(),
8350                    )?
8351                };
8352                ticket.settle();
8353                debug_assert_eq!(ticket.generation, generation);
8354                fork.retire(generation)?;
8355                result
8356            } else {
8357                self.decode_step_t_core(
8358                    e,
8359                    &verify_tokens,
8360                    pos,
8361                    &mut *cache,
8362                    embd_dev,
8363                    ckpt.as_mut(),
8364                )?
8365            };
8366            let pipe_accept = match pipe {
8367                Some(p) => Some(p.accept_begin(round)?),
8368                None => None,
8369            };
8370
8371            ph_mark(&mut ph_verify, phase_on);
8372            // --- 3. GREEDY ACCEPT (walk prefix, stop at first mismatch) ---
8373            // DEVICE-ARGMAX ACCEPT: argmax every verify column ON DEVICE (same 2-pass kernels +
8374            // smallest-index tie-break as host argmax, argmax_gate-validated) and read back ONE
8375            // [T] u32 — replaces the T x n_vocab f32 dtoh + T host argmaxes per round.
8376            // t_pred[j] = target's greedy prediction for the slot after draft[j-1] (j>=1) or after
8377            // last_token (j==0). With a pending bonus, col 0 IS the prediction after last_token
8378            // (== the bonus), so every index shifts by `base` and last_pred is unused.
8379            let t_v = verify_tokens.len();
8380            let mut preds: Vec<u32> = Vec::new();
8381            if !sampled {
8382                for j in 0..t_v {
8383                    e.argmax_token_device_col(&tlogits_d, j, n_vocab, &mut preds_d, j)?;
8384                }
8385                preds = e.dtoh_u32(&preds_d)?; // <- the verify-GPU wait lands here
8386                // #87 SENTINEL TRAP, verify side: a sentinel pred becomes the round's bonus =
8387                // next round's last_token = the next chain's embed lookup. Catch it at the
8388                // source with the column named — an all-NaN VERIFY column implicates the
8389                // stage-split trunk (decode_step_t_core_ppn), not the draft head.
8390                if let Some(bad) = preds[..t_v].iter().position(|&p| (p as usize) >= n_vocab) {
8391                    let col = &tlogits_d.slice(bad * n_vocab..(bad + 1) * n_vocab);
8392                    let mut probe = e.zeros(n_vocab)?;
8393                    e.copy_view_into(&mut probe, 0, col, n_vocab)?;
8394                    let col_h = e.dtoh(&probe)?;
8395                    let col_nan = col_h.iter().filter(|v| v.is_nan()).count();
8396                    return Err(format!(
8397                        "verify argmax sentinel 0x{:08x} >= n_vocab {n_vocab} at round {round} \
8398                         col {bad}/{t_v} pos={pos}: verify-logits col NaN {col_nan}/{n_vocab} \
8399                         — the stage-split verify produced a poisoned column (#87 trap)",
8400                        preds[bad]
8401                    )
8402                    .into());
8403                }
8404            }
8405            ph_mark(&mut ph_wait, phase_on);
8406            let t_pred = |j: usize| -> u32 {
8407                if j == 0 && base == 0 {
8408                    last_pred
8409                } else {
8410                    preds[base + j - 1]
8411                }
8412            };
8413            let mut devacc_seeded = false;
8414            let mut devacc_acc: Option<CudaSlice<u32>> = None;
8415            let (n_acc, bonus) = if !sampled {
8416                // ROUND-STREAM stage (a) (MEMRA_SPEC_DEVACC=1 opt-in): the walk runs ON DEVICE
8417                // (spec_accept_greedy, verbatim rule) and the host reads back 8B (n_acc, bonus)
8418                // instead of the [T] preds. Same sync count — machinery for stages (b)/(c),
8419                // gated on token identity vs the host walk (the arms below are bit-equal rules).
8420                if crate::spec::spec_devacc() && k_round > 0 && !spec_replay && constraint.is_none()
8421                {
8422                    let draft_d = e.htod_u32_v(&draft)?;
8423                    let mut acc_out = e.alloc_u32_zeroed(2)?;
8424                    e.spec_accept_greedy(
8425                        &preds_d,
8426                        &draft_d,
8427                        last_pred,
8428                        base,
8429                        k_round,
8430                        &mut acc_out,
8431                    )?;
8432                    devacc_acc = Some(acc_out.clone());
8433                    // stage (b): next-round seed gathered ON DEVICE from acc_out before the host
8434                    // ever reads n_acc (j=base+n_acc -> vx col j-1; j==0 -> fill_prev). The three
8435                    // non-replay commit arms skip their host-offset seed copies (guarded below);
8436                    // the legacy spec_replay arm keeps its own rx-based seeding (excluded here).
8437                    // NOTE: fill_prev is NOT updated here — the commit arms' TRUE-HIDDEN
8438                    // REFRESH reads the OLD fill_prev (predecessor of this round's verify batch);
8439                    // the update lands after the arms (devacc_seeded guard below).
8440                    e.spec_seed_gather(&vx, &fill_prev, &acc_out, &mut h_seed_buf, base, n_embd)?;
8441                    // 3a: KV lens roll back on device (len = saved + base + n_acc, all arms'
8442                    // unified rule; full accept rewrites the verify-left value). Host mirrors
8443                    // update after the readback; commit_verified_prefix skips its len_d writes.
8444                    if let Some(successor) = successor_attempt.as_ref() {
8445                        opti_fork
8446                            .as_mut()
8447                            .ok_or("optipipe successor reconcile lost fork state")?
8448                            .queue_actual_reconcile(
8449                                e,
8450                                &snap,
8451                                &acc_out,
8452                                successor.verify_tokens[0],
8453                                base,
8454                            )?;
8455                    } else if let Some(ptrs) = &kv_len_ptrs {
8456                        let saved: Vec<i32> = (0..self.layers.len())
8457                            .map(|il| snap.kv_len[il].map(|v| v as i32).unwrap_or(0))
8458                            .collect();
8459                        let saved_d = e.htod_i32(&saved)?;
8460                        e.spec_rollback_kv(ptrs, &saved_d, &acc_out, base, self.layers.len())?;
8461                    }
8462                    devacc_seeded = true;
8463                    let ab = e.dtoh_u32(&acc_out)?;
8464                    (ab[0] as usize, ab[1])
8465                } else {
8466                    let mut n_acc = 0usize;
8467                    for j in 0..k_round {
8468                        if t_pred(j) == draft[j] {
8469                            n_acc += 1;
8470                        } else {
8471                            break;
8472                        }
8473                    }
8474                    // bonus = target's own token at the first non-accepted slot. n_acc in 0..=k; t_pred
8475                    // is defined for j in 0..=k (j==0 -> last_logits, j>=1 -> col j-1, last col = k-1).
8476                    (n_acc, t_pred(n_acc))
8477                }
8478            } else {
8479                // --- SAMPLED ACCEPT (rejection sampling): u_j < p_j(x_j)/q_j(x_j) walk ---
8480                if col_buf.is_none() {
8481                    col_buf = Some(e.zeros(n_vocab)?);
8482                }
8483                // FILTERED p_j: per-verify-col stats (one batched filter_stats call), then the
8484                // filtered gather. j==0&&base==0 reads last_col (its own stats row appended).
8485                let mut pj = vec![0f32; k_round.max(1)];
8486                let mut col_stats: Vec<(f32, f32, f32)> = Vec::new(); // (max, th, z) per verify col used
8487                if k_round > 0 {
8488                    let mut ids: Vec<u32> = Vec::new();
8489                    let mut rows: Vec<i32> = Vec::new();
8490                    for j in 0..k_round {
8491                        if j > 0 || base == 1 {
8492                            ids.push(draft[j]);
8493                            rows.push((base + j) as i32 - 1);
8494                        }
8495                    }
8496                    if !ids.is_empty() {
8497                        let nr = rows.len();
8498                        // penalties: materialize the used columns into one contiguous penalized
8499                        // buffer (rows remapped 0..nr) so stats+gathers see the penalized p.
8500                        // penalties: materialize used columns contiguously, penalize all rows in
8501                        // one launch, and point stats+gathers at the penalized buffer (rows 0..nr).
8502                        let p_rows: Vec<i32> = if pen_on {
8503                            (0..nr as i32).collect()
8504                        } else {
8505                            rows.clone()
8506                        };
8507                        if pen_on {
8508                            if pcol_buf.as_ref().map(|b| b.len()).unwrap_or(0) < nr * n_vocab {
8509                                pcol_buf = Some(e.zeros(nr * n_vocab)?);
8510                            }
8511                            let pc = pcol_buf.as_mut().unwrap();
8512                            for (i2, &r) in rows.iter().enumerate() {
8513                                let c = r as usize;
8514                                e.copy_view_into(
8515                                    pc,
8516                                    i2 * n_vocab,
8517                                    &tlogits_d.slice(c * n_vocab..(c + 1) * n_vocab),
8518                                    n_vocab,
8519                                )?;
8520                            }
8521                            let h = pen_hist_d.as_ref().unwrap();
8522                            let nh = h.len();
8523                            e.penalize_logits_rows(
8524                                pc,
8525                                h,
8526                                nh,
8527                                sp.penalty_repeat,
8528                                sp.penalty_freq,
8529                                sp.penalty_present,
8530                                n_vocab,
8531                                nr,
8532                            )?;
8533                        }
8534                        let p_src: &CudaSlice<f32> = if pen_on {
8535                            pcol_buf.as_ref().unwrap()
8536                        } else {
8537                            &tlogits_d
8538                        };
8539                        let rowsd = e.htod_i32(&p_rows)?;
8540                        let (mut th_d, mut z_d, mut mx_d) =
8541                            (e.zeros(nr)?, e.zeros(nr)?, e.zeros(nr)?);
8542                        e.filter_stats(
8543                            p_src, n_vocab, &rowsd, &mut th_d, &mut z_d, &mut mx_d, n_vocab, nr,
8544                            sp_temp, sp.top_k, sp.top_p, sp.min_p,
8545                        )?;
8546                        let idsd = e.htod_u32_v(&ids)?;
8547                        let mut outd = e.zeros(nr)?;
8548                        e.softmax_gather_filtered(
8549                            p_src, n_vocab, &idsd, &rowsd, &th_d, &z_d, &mut outd, n_vocab, nr,
8550                            sp_temp,
8551                        )?;
8552                        let outv = e.dtoh(&outd)?;
8553                        let (thv, zv, mxv) = (e.dtoh(&th_d)?, e.dtoh(&z_d)?, e.dtoh(&mx_d)?);
8554                        let mut oi = 0usize;
8555                        for j in 0..k_round {
8556                            if j > 0 || base == 1 {
8557                                pj[j] = outv[oi];
8558                                oi += 1;
8559                            }
8560                        }
8561                        col_stats = (0..nr).map(|i| (mxv[i], thv[i], zv[i])).collect();
8562                    }
8563                    if base == 0 {
8564                        let lc: &CudaSlice<f32> = if pen_on {
8565                            if col_buf.is_none() {
8566                                col_buf = Some(e.zeros(n_vocab)?);
8567                            }
8568                            let cb = col_buf.as_mut().unwrap();
8569                            e.copy_into(
8570                                cb,
8571                                0,
8572                                last_col_logits
8573                                    .as_ref()
8574                                    .expect("sampled: last_col_logits unset"),
8575                                n_vocab,
8576                            )?;
8577                            let h = pen_hist_d.as_ref().unwrap();
8578                            let nh = h.len();
8579                            e.penalize_logits(
8580                                cb,
8581                                h,
8582                                nh,
8583                                sp.penalty_repeat,
8584                                sp.penalty_freq,
8585                                sp.penalty_present,
8586                                n_vocab,
8587                            )?;
8588                            col_buf.as_ref().unwrap()
8589                        } else {
8590                            last_col_logits
8591                                .as_ref()
8592                                .expect("sampled: last_col_logits unset")
8593                        };
8594                        let rows0 = e.htod_i32(&[0])?;
8595                        let (mut th_d, mut z_d, mut mx_d) = (e.zeros(1)?, e.zeros(1)?, e.zeros(1)?);
8596                        e.filter_stats(
8597                            lc, n_vocab, &rows0, &mut th_d, &mut z_d, &mut mx_d, n_vocab, 1,
8598                            sp_temp, sp.top_k, sp.top_p, sp.min_p,
8599                        )?;
8600                        let idsd = e.htod_u32_v(&[draft[0]])?;
8601                        let mut outd = e.zeros(1)?;
8602                        e.softmax_gather_filtered(
8603                            lc, n_vocab, &idsd, &rows0, &th_d, &z_d, &mut outd, n_vocab, 1, sp_temp,
8604                        )?;
8605                        pj[0] = e.dtoh(&outd)?[0];
8606                        last_col_stats =
8607                            Some((e.dtoh(&mx_d)?[0], e.dtoh(&th_d)?[0], e.dtoh(&z_d)?[0]));
8608                    }
8609                }
8610                // q source: the graph arm retained the head logits in the persistent q_slots;
8611                // the eager arm in per-round draft_logits clones. Same raw-logit values either way.
8612                // FILTERED q_j: stats from draft_stats (eager pushes in-chain; the graph arm
8613                // computes them post-replay — graph engages only filter/penalty-free, so the
8614                // stats degenerate to th=0/full-Z there, keeping ONE accept path).
8615                let q_bufs: &[CudaSlice<f32>] = if dctx.graph_s.is_some() {
8616                    &dctx.q_slots
8617                } else {
8618                    &draft_logits
8619                };
8620                let mut n_acc = 0usize;
8621                for j in 0..k_round {
8622                    let (qmx, qth, qz) = draft_stats[j];
8623                    let idsd = e.htod_u32_v(&[draft_idx[j]])?;
8624                    let rowsd = e.htod_i32(&[0])?;
8625                    let thd = e.htod(&[qth])?;
8626                    let zd = e.htod(&[qz])?;
8627                    let _ = qmx;
8628                    let mut outd = e.zeros(1)?;
8629                    e.softmax_gather_filtered(
8630                        &q_bufs[j], d_vocab, &idsd, &rowsd, &thd, &zd, &mut outd, d_vocab, 1,
8631                        sp_temp,
8632                    )?;
8633                    let qj = e.dtoh(&outd)?[0];
8634                    let u = host_u01(sp_seed, uctr);
8635                    uctr += 1;
8636                    if (u as f64) * (qj as f64) < pj[j] as f64 {
8637                        n_acc += 1;
8638                    } else {
8639                        break;
8640                    }
8641                }
8642                let bonus = if n_acc == k_round {
8643                    // FULL ACCEPT: bonus ~ FILTERED softmax at the last verify column.
8644                    let col = base + k_round - 1;
8645                    let cb = col_buf.as_mut().unwrap();
8646                    e.copy_view_into(
8647                        cb,
8648                        0,
8649                        &tlogits_d.slice(col * n_vocab..(col + 1) * n_vocab),
8650                        n_vocab,
8651                    )?;
8652                    if pen_on {
8653                        let h = pen_hist_d.as_ref().unwrap();
8654                        let nh = h.len();
8655                        e.penalize_logits(
8656                            cb,
8657                            h,
8658                            nh,
8659                            sp.penalty_repeat,
8660                            sp.penalty_freq,
8661                            sp.penalty_present,
8662                            n_vocab,
8663                        )?;
8664                    }
8665                    if perturb_buf.is_none() {
8666                        perturb_buf = Some(e.zeros(d_vocab.max(n_vocab))?);
8667                    }
8668                    // STATS MUST COME FROM THIS COLUMN (bug fix 2026-08-05, lane/sampler-
8669                    // truncation-fix; receipts research/sampfix-20260805/). The old code reused
8670                    // `col_stats.last()` here, which is ALWAYS the wrong row: the gathered set
8671                    // covers verify columns 0..=(base+k_round-2) (rows pushed as base+j-1), while
8672                    // the full-accept bonus samples column base+k_round-1 — exactly ONE PAST the
8673                    // last gathered column, in both base arms. `th` is a threshold in e-units of
8674                    // its OWN row's max, so feeding a neighbour's (row_max, th) into
8675                    // gumbel_perturb_filtered mis-scales every e0 = exp((x-row_max)/T). When the
8676                    // donor column's peak is higher by more than T*ln(1/th), EVERY id fails
8677                    // `e0 >= th`, the whole perturbed row becomes -3.4e38, and the 2-pass argmax
8678                    // falls through to its smallest-index tie-break => token id 0 ("!") spliced
8679                    // mid-word. Fragility is ordered by how large th is: min_p pins th = min_p
8680                    // (0.05 => trigger at delta > 2.4 at T=0.8, fires constantly), top_p's
8681                    // mass-boundary th is smaller, top_k's k-th-largest th smaller still — which
8682                    // is why the head-to-head matrix saw min_p and top_p corrupt while top_k-only
8683                    // stayed clean. The pure-temp default regime is immune (th == 0 masks nothing,
8684                    // and row_max is unused once nothing is masked), so this fix is a byte-level
8685                    // no-op for the untruncated serve default. One extra one-block filter_stats
8686                    // per full-accept round is the whole cost.
8687                    let (mx, th) = {
8688                        let rows0 = e.htod_i32(&[0])?;
8689                        let (mut th_d, mut z_d, mut mx_d) = (e.zeros(1)?, e.zeros(1)?, e.zeros(1)?);
8690                        let cb0 = col_buf.as_ref().unwrap();
8691                        e.filter_stats(
8692                            cb0, n_vocab, &rows0, &mut th_d, &mut z_d, &mut mx_d, n_vocab, 1,
8693                            sp_temp, sp.top_k, sp.top_p, sp.min_p,
8694                        )?;
8695                        (e.dtoh(&mx_d)?[0], e.dtoh(&th_d)?[0])
8696                    };
8697                    let pb = perturb_buf.as_mut().unwrap();
8698                    let cb2 = col_buf.as_ref().unwrap();
8699                    e.gumbel_perturb_filtered(cb2, pb, n_vocab, sp_seed, sctr, sp_temp, mx, th)?;
8700                    sctr += 1;
8701                    let td = e.argmax_token_device(pb, n_vocab)?;
8702                    e.dtoh_u32_one(&td)?
8703                } else {
8704                    // REJECT at n_acc: bonus ~ norm(max(0, softmax_T(p) - softmax_T(q))).
8705                    let cb = col_buf.as_mut().unwrap();
8706                    if n_acc > 0 || base == 1 {
8707                        let col = base + n_acc - 1;
8708                        e.copy_view_into(
8709                            cb,
8710                            0,
8711                            &tlogits_d.slice(col * n_vocab..(col + 1) * n_vocab),
8712                            n_vocab,
8713                        )?;
8714                    } else {
8715                        let lc = last_col_logits.as_ref().unwrap();
8716                        e.copy_into(cb, 0, lc, n_vocab)?;
8717                    }
8718                    if pen_on {
8719                        let h = pen_hist_d.as_ref().unwrap();
8720                        let nh = h.len();
8721                        e.penalize_logits(
8722                            cb,
8723                            h,
8724                            nh,
8725                            sp.penalty_repeat,
8726                            sp.penalty_freq,
8727                            sp.penalty_present,
8728                            n_vocab,
8729                        )?;
8730                    }
8731                    let cb2 = col_buf.as_ref().unwrap();
8732                    let sc = sctr;
8733                    sctr += 1;
8734                    // p-stats for the reject column: from col_stats when the col was gathered,
8735                    // else (j==0&&base==0) from last_col_stats.
8736                    let p_stats = if n_acc > 0 || base == 1 {
8737                        // col index within the gathered set == number of gathered cols before n_acc
8738                        let gi = if base == 1 { n_acc } else { n_acc - 1 };
8739                        col_stats.get(gi).copied().unwrap_or_else(|| {
8740                            (0.0, 0.0, 1.0) // unreachable: gathered cols always cover the reject slot
8741                        })
8742                    } else {
8743                        last_col_stats.expect("sampled: last_col_stats unset at reject")
8744                    };
8745                    let q_stats = draft_stats[n_acc];
8746                    if let Some(map) = &d2t_dev {
8747                        if q_full_buf.is_none() {
8748                            q_full_buf = Some(e.zeros(n_vocab)?);
8749                        }
8750                        let qf = q_full_buf.as_mut().unwrap();
8751                        e.scatter_trim_logits(&q_bufs[n_acc], map, qf, d_vocab, n_vocab)?;
8752                        let qf2 = q_full_buf.as_ref().unwrap();
8753                        e.residual_sample_filtered(
8754                            cb2,
8755                            Some(qf2),
8756                            n_vocab,
8757                            sp_temp,
8758                            sp_seed,
8759                            sc,
8760                            p_stats,
8761                            q_stats,
8762                            &mut sample_tok,
8763                        )?;
8764                    } else {
8765                        e.residual_sample_filtered(
8766                            cb2,
8767                            Some(&q_bufs[n_acc]),
8768                            n_vocab,
8769                            sp_temp,
8770                            sp_seed,
8771                            sc,
8772                            p_stats,
8773                            q_stats,
8774                            &mut sample_tok,
8775                        )?;
8776                    }
8777                    e.dtoh_u32(&sample_tok)?[0]
8778                };
8779                (n_acc, bonus)
8780            };
8781            // --- 3b. GRAMMAR TRUNCATION (constrained spec, 2026-08-03): the grammar is
8782            // an extra rejection rule AFTER the exactness verify (the batched-verify-twins
8783            // ordering). Walk the accepted drafts through the grammar in commit order; the
8784            // first illegal token truncates acceptance at its slot, and that slot's emission
8785            // is recomputed as the MASKED argmax of the target's own verify column — token-
8786            // identical to constrained plain greedy decode (an unmasked argmax that is
8787            // grammar-legal IS the masked argmax: masking only removes competitors). The
8788            // column D2H (~1MB) is paid only when a cut fires — the tight-grammar cost,
8789            // measured in acceptance numbers, never hidden.
8790            let (n_acc, bonus) = match constraint.as_deref_mut() {
8791                None => (n_acc, bonus),
8792                Some(c) => {
8793                    fn ce(e2: String) -> Box<dyn std::error::Error> {
8794                        format!("constraint: {e2}").into()
8795                    }
8796                    let mut na = n_acc;
8797                    let mut cut = false;
8798                    for (j, &d) in draft.iter().enumerate().take(n_acc) {
8799                        if c.is_allowed(d).map_err(ce)? {
8800                            c.consume(d).map_err(ce)?;
8801                        } else {
8802                            na = j;
8803                            cut = true;
8804                            dm_cut_tokens += n_acc - j;
8805                            break;
8806                        }
8807                    }
8808                    if cut {
8809                        dm_cuts += 1;
8810                    }
8811                    let mut bo = bonus;
8812                    if cut || !c.is_allowed(bo).map_err(ce)? {
8813                        let mut row = if na == 0 && base == 0 {
8814                            init_logits_host
8815                                .clone()
8816                                .ok_or("constraint: init logits missing (round-0 cut)")?
8817                        } else {
8818                            e.dtoh_view(
8819                                &tlogits_d.slice((base + na - 1) * n_vocab..(base + na) * n_vocab),
8820                            )?
8821                        };
8822                        c.mask_logits(&mut row).map_err(ce)?;
8823                        bo = argmax(&row) as u32;
8824                    }
8825                    c.consume(bo).map_err(ce)?;
8826                    (na, bo)
8827                }
8828            };
8829            let mut successor_valid = false;
8830            if let Some((q_proxy, expected_d2)) = rejected_probe {
8831                let v_n = n_acc == 1 && bonus == expected_d2;
8832                eprintln!(
8833                    "[opti-controller] shadow q={q_proxy:.6} admitted=false v_n={v_n} \
8834                     expected_d2={expected_d2} n_acc={n_acc} bonus={bonus}",
8835                );
8836            }
8837            if let Some(successor) = successor_attempt.as_ref() {
8838                successor_valid = n_acc == 1 && bonus == successor.verify_tokens[0];
8839                let generation = successor.generation;
8840                let q_proxy = successor.q_proxy;
8841                let expected_pending = successor.verify_tokens[0];
8842                let resolution_ms = successor.issued_at.elapsed().as_secs_f64() * 1e3;
8843                let fork = opti_fork
8844                    .as_mut()
8845                    .ok_or("optipipe successor resolution lost fork state")?;
8846                fork.finish_actual_reconcile(e, &mut *cache, &snap, n_acc, base, successor_valid)?;
8847                if successor_valid {
8848                    OPTI_FORK_HITS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
8849                } else {
8850                    OPTI_FORK_MISSES.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
8851                    OPTI_RECONCILES.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
8852                    OPTI_WASTED_DRAFT_TOKENS.fetch_add(2, std::sync::atomic::Ordering::Relaxed);
8853                }
8854                let breaker_tripped = fork
8855                    .controller
8856                    .as_mut()
8857                    .expect("controller policy")
8858                    .resolve(successor_valid);
8859                if breaker_tripped {
8860                    OPTI_BREAKER_TRIPS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
8861                }
8862                eprintln!(
8863                    "[opti-controller] resolve generation={} hit={} q={q_proxy:.6} \
8864                     expected_pending={expected_pending} n_acc={n_acc} bonus={bonus} \
8865                     resolution_ms={resolution_ms:.3} reconcile={} breaker={}",
8866                    generation.id, successor_valid, !successor_valid, breaker_tripped,
8867                );
8868                if !successor_valid {
8869                    let mut successor = successor_attempt
8870                        .take()
8871                        .expect("controller successor disappeared on miss");
8872                    successor.settle();
8873                    fork.retire(generation)?;
8874                }
8875            }
8876            total_drafted += k_round;
8877            total_accepted += n_acc;
8878            if let Some(t) = sess_telem {
8879                // Greedy, rejection-sampling, and grammar truncation all converge here after
8880                // the accept decision is already on host. Fixed-size relaxed atomics only.
8881                t.record_round(k_round, n_acc);
8882            }
8883            if spec_stats {
8884                st_len_hist[k_round] += 1;
8885                for j in 0..k_round {
8886                    st_drafted[j] += 1;
8887                }
8888                for j in 0..n_acc {
8889                    st_accepted[j] += 1;
8890                }
8891                if n_acc == k_round {
8892                    st_full += 1;
8893                }
8894            }
8895
8896            if debug_spec {
8897                eprintln!(
8898                    "[R{round}] pos={pos} out_len={} last_tok={last_token} draft={draft:?} n_acc={n_acc} bonus={bonus} t_pred0={}",
8899                    out.len(),
8900                    t_pred(0)
8901                );
8902            }
8903
8904            // --- 4. COMMIT: draft[0..n_acc] then bonus (n_acc + 1 tokens) ---
8905            let commit_started = std::time::Instant::now();
8906            // SESSION MODE: every accepted column is already in the CACHE — `out` must carry all
8907            // of them (overshoot past max_new included) or `committed` under-counts the cache rows
8908            // and the next turn's continuation seeds one token off (gate-caught 2026-07-05). The
8909            // single-shot path keeps the cap (its caller truncates + drops the cache anyway).
8910            for j in 0..n_acc {
8911                if !session_mode && out.len() >= max_new {
8912                    break;
8913                }
8914                out.push(draft[j]);
8915            }
8916            if pen_on {
8917                pen_hist.extend_from_slice(&draft[0..n_acc]);
8918                pen_hist.push(bonus);
8919            }
8920            let bonus_emitted = session_mode || out.len() < max_new;
8921            if bonus_emitted {
8922                out.push(bonus);
8923            }
8924            last_token = bonus;
8925
8926            // --- 5. ROLLBACK + advance (§C) ---
8927            if n_acc == k_round && !spec_replay {
8928                // FULL ACCEPT, BONUS FOLD: all verify columns (pending? + drafts) are committed in
8929                // cache; the NEW bonus stays PENDING for the next round's verify batch — NO extra
8930                // T=1 trunk pass. The next draft chain seeds from the MTP block's h_nextn at the
8931                // bonus position: one MTP-block pass (~1/33 trunk cost) replaces the trunk read.
8932                // last_pred is dead in the pending path (t_pred reads verify col 0).
8933                //
8934                // PERSISTENT DRAFT KV, full-accept fill: the chain covered last_token +
8935                // draft[0..k_round-2] as INPUTS (slots P..P'-2); draft[k_round-1] (slot P'-1) was
8936                // only ever an output, so its entry is MISSING. Fill it from vh_seed — its EXACT
8937                // trunk hidden (the last verify column). set_len first: a p-min break may have
8938                // left one extra chain append at that slot. Partial accepts need NO fill (the
8939                // chain already covered every accepted position; round-start set_len truncates).
8940                let mut vh_seed = e.zeros(n_embd)?;
8941                e.copy_view_into(
8942                    &mut vh_seed,
8943                    0,
8944                    &vx.slice((t_v - 1) * n_embd..t_v * n_embd),
8945                    n_embd,
8946                )?;
8947                if refresh {
8948                    // TRUE-HIDDEN REFRESH (2026-07-03, the HANDOVER-listed acceptance lever):
8949                    // overwrite ALL committed positions' scratch entries with K/V from their EXACT
8950                    // verify hiddens — the reference engine's mtp_update fills from true hiddens;
8951                    // the full stack (vx) is already resident from the verify. Replaces both the
8952                    // chain-approximate entries AND the old last-token-only fill. Acceptance-only
8953                    // (draft attention quality); exactness stays the verify's job.
8954                    scratch.set_len(e, pos)?;
8955                    // PREDECESSOR pairing: row i gets vx[i-1]; row 0 the carried fill_prev
8956                    // (hidden of the last committed row before this verify batch).
8957                    let mut vxs = e.zeros(t_v * n_embd)?;
8958                    e.copy_into(&mut vxs, 0, &fill_prev, n_embd)?;
8959                    if t_v > 1 {
8960                        e.copy_view_into(
8961                            &mut vxs,
8962                            n_embd,
8963                            &vx.slice(0..(t_v - 1) * n_embd),
8964                            (t_v - 1) * n_embd,
8965                        )?;
8966                    }
8967                    self.mtp_kv_fill(e, mtp, &verify_tokens, &vxs, pos, &mut *scratch, embd_dev)?;
8968                } else {
8969                    scratch.set_len(e, pos + base + k_round - 1)?;
8970                    // predecessor of the last draft = verify col t_v-2 (or fill_prev at t_v==1)
8971                    let mut hp = e.zeros(n_embd)?;
8972                    if t_v >= 2 {
8973                        e.copy_view_into(
8974                            &mut hp,
8975                            0,
8976                            &vx.slice((t_v - 2) * n_embd..(t_v - 1) * n_embd),
8977                            n_embd,
8978                        )?;
8979                    } else {
8980                        e.copy_into(&mut hp, 0, &fill_prev, n_embd)?;
8981                    }
8982                    self.mtp_kv_fill(
8983                        e,
8984                        mtp,
8985                        &[draft[k_round - 1]],
8986                        &hp,
8987                        pos + base + k_round - 1,
8988                        &mut *scratch,
8989                        embd_dev,
8990                    )?;
8991                }
8992                // REFERENCE SEEDING: no pseudo pass — the next chain's step 0 IS the
8993                // reference's (id_last, h_prev) draft row; it appends the bonus's scratch
8994                // entry itself. Seed = TRUE hidden of the bonus's predecessor (last verify
8995                // col). Saves one MTP-block pass per round on top of the pairing fix.
8996                if !devacc_seeded {
8997                    e.copy_into(&mut h_seed_buf, 0, &vh_seed, n_embd)?;
8998                    e.copy_into(&mut fill_prev, 0, &vh_seed, n_embd)?;
8999                }
9000                pending = Some(bonus);
9001                if debug_spec {
9002                    eprintln!("  -> FULL ACCEPT (bonus pending, prev-h seed)");
9003                }
9004            } else if !spec_replay && base + n_acc >= 1 {
9005                // PARTIAL ACCEPT, REPLAY-FREE (2026-07-03 — the profiled #1 long-ctx spec cost):
9006                // the verify's first j = base+n_acc columns ARE the committed sequence, computed
9007                // bit-identically to eager (decode-exact contract) — so KEEP them: KV truncates to
9008                // pos+j, recurrent state rebuilds from the VerifyCkpt (same-kernel gdn prefix
9009                // re-run / pure state-clone restore), and the bonus stays PENDING exactly like the
9010                // full-accept path — the legacy duplicate trunk replay is gone. The next chain
9011                // seeds from the MTP pseudo-hidden of the bonus, whose seed = the TRUE verify
9012                // hidden of its predecessor (col j-1) — same one-hop pseudo structure as full
9013                // accept (never compounds: the next verify recomputes true hiddens for all
9014                // committed columns).
9015                let j = base + n_acc;
9016                self.commit_verified_prefix(
9017                    e,
9018                    &mut *cache,
9019                    &snap,
9020                    ckpt.as_ref().unwrap(),
9021                    j,
9022                    devacc_seeded,
9023                    if devacc_seeded {
9024                        devacc_acc.as_ref().map(|a| (a, base, t_v))
9025                    } else {
9026                        None
9027                    },
9028                )?;
9029                let mut seed = e.zeros(n_embd)?;
9030                e.copy_view_into(
9031                    &mut seed,
9032                    0,
9033                    &vx.slice((j - 1) * n_embd..j * n_embd),
9034                    n_embd,
9035                )?;
9036                // Draft scratch: TRUE-HIDDEN REFRESH of the committed prefix (see the full-accept
9037                // branch); without it the chain entries stand and only the tail truncates. Either
9038                // way len ends at pos+j so the pseudo append lands at the bonus's slot pos+j
9039                // (persistent mode), rope pos+j+1 (chain convention).
9040                if refresh {
9041                    scratch.set_len(e, pos)?;
9042                    let mut vxs = e.zeros(j * n_embd)?;
9043                    e.copy_into(&mut vxs, 0, &fill_prev, n_embd)?;
9044                    if j > 1 {
9045                        e.copy_view_into(
9046                            &mut vxs,
9047                            n_embd,
9048                            &vx.slice(0..(j - 1) * n_embd),
9049                            (j - 1) * n_embd,
9050                        )?;
9051                    }
9052                    self.mtp_kv_fill(
9053                        e,
9054                        mtp,
9055                        &verify_tokens[0..j],
9056                        &vxs,
9057                        pos,
9058                        &mut *scratch,
9059                        embd_dev,
9060                    )?;
9061                } else {
9062                    scratch.set_len(e, pos + j)?;
9063                }
9064                // REFERENCE SEEDING (see the full-accept branch): seed = TRUE hidden of the
9065                // bonus's predecessor (verify col j-1); no pseudo pass.
9066                if !devacc_seeded {
9067                    e.copy_into(&mut h_seed_buf, 0, &seed, n_embd)?;
9068                    e.copy_into(&mut fill_prev, 0, &seed, n_embd)?;
9069                }
9070                pending = Some(bonus);
9071                if debug_spec {
9072                    eprintln!("  -> PARTIAL(replay-free j={j}, bonus pending, prev-h seed)");
9073                }
9074            } else if !spec_replay {
9075                // ZERO ROUND FOLD (2026-07-10, verify-cost target #3): base+n_acc == 0 — a
9076                // pending-less round where nothing was accepted (PMIN0 zero-draft chains after a
9077                // replay/commit, or plain 0-accept rounds at round 0). The old path replayed
9078                // [bonus] through a FULL m=1 trunk+head forward (the 489us full-vocab head pass
9079                // measured at ~0.75/round on PMIN0 configs). Instead: restore the pre-round
9080                // snapshot and let the bonus ride the NEXT round's verify as col 0 — the existing
9081                // base=1 pending machinery, bit-identical by the decode-exact verify contract.
9082                // Seed: the bonus's predecessor is the last COMMITTED token, whose hidden
9083                // fill_prev already carries (same seeding as the 1-token-replay case it replaces).
9084                cache.rollback(e, &snap, 0)?;
9085                scratch.set_len(e, pos)?;
9086                e.copy_into(&mut h_seed_buf, 0, &fill_prev, n_embd)?;
9087                pending = Some(bonus);
9088                if debug_spec {
9089                    eprintln!("  -> ZERO-ROUND FOLD (bonus pending, fill_prev seed)");
9090                }
9091            } else {
9092                // PARTIAL ACCEPT, LEGACY REPLAY (seam MEMRA_SPEC_REPLAY=1 — or j==0: nothing of
9093                // this round survives, only possible before the first pending exists, ~round 0):
9094                // restore EVERYTHING to the pre-round snapshot (KV truncate to pos + recur
9095                // restore), then replay the committed prefix pending? ++ draft[0..n_acc] ++
9096                // [bonus] as ONE batched T forward — single weight read, bit-identical to greedy
9097                // (the verify-all-columns path is the same math). Commits the bonus with a TRUE
9098                // trunk hidden.
9099                cache.rollback(e, &snap, 0)?; // accept_len=0: KV len = pos, recur = snapshot
9100                let mut replay: Vec<u32> = Vec::with_capacity(base + n_acc + 1);
9101                if let Some(b) = pending.take() {
9102                    replay.push(b);
9103                }
9104                replay.extend_from_slice(&draft[0..n_acc]);
9105                replay.push(bonus);
9106                // Full-stack forward (decode_step_t_core = decode_step_t_h_emb_dev's body):
9107                // Predecessor pairing seeds from the PREDECESSOR row (col len-2) — the same-row path takes the
9108                // last col exactly as before (byte-identical to the old _h_emb_dev call).
9109                let (rl_d, rx) = if self.qwen35_serving_class() {
9110                    let mut logits = Vec::with_capacity(replay.len() * n_vocab);
9111                    let mut hidden = e.uninit(replay.len() * n_embd)?;
9112                    for (row, &token) in replay.iter().enumerate() {
9113                        let (row_logits, row_hidden) =
9114                            self.spec_target_step_h(e, token, &mut *cache)?;
9115                        logits.extend_from_slice(&row_logits);
9116                        e.dtod_copy_into(&row_hidden, &mut hidden, row * n_embd)?;
9117                    }
9118                    (e.htod(&logits)?, hidden)
9119                } else {
9120                    self.decode_step_t_core(e, &replay, pos, &mut *cache, embd_dev, None)?
9121                };
9122                // last_pred = argmax of the LAST column's logits (predicts the token after `bonus`)
9123                // — device argmax + one 4-byte read instead of the full-vocab column dtoh.
9124                e.argmax_token_device_col(&rl_d, replay.len() - 1, n_vocab, &mut preds_d, 0)?;
9125                last_pred = e.dtoh_u32(&preds_d)?[0];
9126                if sampled {
9127                    let lr0 = replay.len();
9128                    let lc = last_col_logits
9129                        .as_mut()
9130                        .expect("sampled: last_col_logits unset");
9131                    e.copy_view_into(
9132                        lc,
9133                        0,
9134                        &rl_d.slice((lr0 - 1) * n_vocab..lr0 * n_vocab),
9135                        n_vocab,
9136                    )?;
9137                }
9138                let lr = replay.len();
9139                if lr >= 2 {
9140                    e.copy_view_into(
9141                        &mut h_seed_buf,
9142                        0,
9143                        &rx.slice((lr - 2) * n_embd..(lr - 1) * n_embd),
9144                        n_embd,
9145                    )?;
9146                } else {
9147                    // 1-token replay (round-0 miss): the bonus's predecessor is the OLD
9148                    // last_token, whose own-row hidden fill_prev still holds.
9149                    e.copy_into(&mut h_seed_buf, 0, &fill_prev, n_embd)?;
9150                }
9151                // the bonus is COMMITTED here — it becomes the last committed row.
9152                let mut rh_last = e.zeros(n_embd)?;
9153                e.copy_view_into(
9154                    &mut rh_last,
9155                    0,
9156                    &rx.slice((lr - 1) * n_embd..lr * n_embd),
9157                    n_embd,
9158                )?;
9159                e.copy_into(&mut fill_prev, 0, &rh_last, n_embd)?;
9160                if debug_spec {
9161                    eprintln!("  -> PARTIAL(replay={replay:?}), next_pred={last_pred}");
9162                }
9163            }
9164            if devacc_seeded {
9165                // stage (b) epilogue: fill_prev takes the gathered seed AFTER the refresh fills
9166                // consumed the old value (both slots carry the same value in every non-replay arm).
9167                e.copy_into(&mut fill_prev, 0, &h_seed_buf, n_embd)?;
9168            }
9169            if successor_valid {
9170                let optimistic_scratch_len = successor_attempt
9171                    .as_ref()
9172                    .expect("valid controller successor disappeared")
9173                    .scratch_len;
9174                // The normal current-round commit refreshed/truncated the logical scratch tail.
9175                // Its optimistic successor row was already written physically, so restoring only
9176                // the retained logical length makes that row live for the carried round.
9177                scratch.set_len(e, optimistic_scratch_len)?;
9178            }
9179            if let Some(current) = current_opti.take() {
9180                opti_fork
9181                    .as_mut()
9182                    .ok_or("optipipe current retirement lost fork state")?
9183                    .retire(current.generation)?;
9184            }
9185            if successor_valid {
9186                let successor = successor_attempt
9187                    .take()
9188                    .expect("valid controller successor disappeared before promotion");
9189                let generation = successor.generation;
9190                opti_fork
9191                    .as_mut()
9192                    .ok_or("optipipe successor promotion lost fork state")?
9193                    .promote_successor_snapshot(&mut snap, generation);
9194                carried_opti = Some(successor);
9195            }
9196            if anatomy_on {
9197                // Commit/rollback is normally asynchronous on the primary/head stream. Bound it
9198                // only for this diagnostic so it does not disappear into the following draft's
9199                // first token readback.
9200                e.stream().synchronize()?;
9201                ph_commit += commit_started.elapsed().as_secs_f64();
9202            }
9203            // adaptive-K update (host math, zero syncs): next round drafts accepted-run + 1,
9204            // clamped to [floor(pos), k_cap]. cache.pos is post-rollback here (the round's
9205            // final position — the floor's position key reads the committed depth). Burst
9206            // rounds (`continue` above) draft the captured fixed depth and skip this, exactly
9207            // like gemma's burst arm.
9208            if adapt {
9209                let fl_now = floor_at(cache.pos);
9210                kc = (n_acc + 1).clamp(fl_now.min(k_cap), k_cap);
9211            }
9212            ph_mark(&mut ph_rest, phase_on);
9213            if let Some(p) = pipe {
9214                p.accept_end(round);
9215            }
9216            drop(pipe_accept);
9217            round += 1;
9218            // sse-cadence: this round's accepted drafts + bonus are committed (out is
9219            // append-only past step 4) — flush at round cadence.
9220            keep_going = flush_commit(&mut on_commit, &out, &mut flushed);
9221        }
9222        if let Some(mut ticket) = carried_opti.take() {
9223            opti_fork
9224                .as_mut()
9225                .ok_or("optipipe tail drain lost fork state")?
9226                .cancel_controller_ticket(e, &mut *cache, &mut *scratch, &snap, &mut ticket)?;
9227        }
9228        // sse-cadence: nothing below appends to `out`; flush any remainder (defensive).
9229        // (verdict ignored — the burst is over either way; the session tail runs unchanged.)
9230        let _ = flush_commit(&mut on_commit, &out, &mut flushed);
9231
9232        if spec_stats {
9233            let per_slot: Vec<String> = (0..k)
9234                .map(|j| {
9235                    if st_drafted[j] > 0 {
9236                        format!(
9237                            "{}/{}={:.3}",
9238                            st_accepted[j],
9239                            st_drafted[j],
9240                            st_accepted[j] as f64 / st_drafted[j] as f64
9241                        )
9242                    } else {
9243                        "0/0".into()
9244                    }
9245                })
9246                .collect();
9247            let acc = if total_drafted > 0 {
9248                total_accepted as f64 / total_drafted as f64
9249            } else {
9250                0.0
9251            };
9252            eprintln!(
9253                "[spec-stats] rounds={round} full_accept={st_full} len_hist={st_len_hist:?} \
9254                       per_slot=[{}] total={total_accepted}/{total_drafted}={acc:.3} \
9255                       tok_per_round={:.3}",
9256                per_slot.join(" "),
9257                (total_accepted + round) as f64 / round.max(1) as f64
9258            );
9259        }
9260        if constraint.is_some() {
9261            eprintln!(
9262                "[draft-mask] mask_rounds={dm_rounds} clone_total={:.3}ms \
9263                 clone_per_round={:.4}ms gram_cuts={dm_cuts}/{round} cut_tokens={dm_cut_tokens}",
9264                dm_clone_ns as f64 / 1e6,
9265                dm_clone_ns as f64 / 1e6 / dm_rounds.max(1) as f64
9266            );
9267        }
9268        if phase_on {
9269            let tot = ph_draft + ph_verify + ph_wait + ph_rest;
9270            eprintln!(
9271                "[spec-phase] draft={:.1}ms ({:.1}%) verify-issue={:.1}ms ({:.1}%) verify-wait={:.1}ms ({:.1}%) commit-host={:.1}ms ({:.1}%) rounds={round}",
9272                ph_draft * 1e3,
9273                ph_draft / tot * 100.0,
9274                ph_verify * 1e3,
9275                ph_verify / tot * 100.0,
9276                ph_wait * 1e3,
9277                ph_wait / tot * 100.0,
9278                ph_rest * 1e3,
9279                ph_rest / tot * 100.0
9280            );
9281        }
9282        if anatomy_on {
9283            let rounds_f = round.max(1) as f64;
9284            let other = (ph_rest - ph_commit).max(0.0);
9285            eprintln!(
9286                "[spec-anatomy] per-round draft={:.3}ms pp-verify={:.3}ms \
9287                 verify-accept={:.3}ms commit-rollback={:.3}ms other={:.3}ms rounds={round}",
9288                ph_draft * 1e3 / rounds_f,
9289                ph_verify * 1e3 / rounds_f,
9290                ph_wait * 1e3 / rounds_f,
9291                ph_commit * 1e3 / rounds_f,
9292                other * 1e3 / rounds_f,
9293            );
9294        }
9295        let _pipe_tail = pipe.map(|p| p.primary());
9296        // SESSION TAIL: leave the session in the exact invariant the next turn's suffix prime
9297        // expects — every row in `committed` has trunk KV/recur state AND an exact draft-KV row.
9298        // Park the draft-graph ctx back on the session (the serve-burst fixed-cost fix): the next
9299        // burst replays instead of recapturing. Error paths (`?` above) drop it — recaptured then.
9300        if let Some(slot) = sess_draft_slot.take() {
9301            *slot = Some(dctx);
9302        }
9303        let t_rounds = t_ent.elapsed();
9304        if let Some((committed, last_h, next_pred_slot, sctr_slot, uctr_slot)) = sess_tail.take() {
9305            *sctr_slot = sctr;
9306            *uctr_slot = uctr;
9307            *next_pred_slot = Some(last_pred);
9308            let mut stashed_pending = false;
9309            if let Some(b) = pending.take() {
9310                if !sampled {
9311                    // PENDING-CARRY (2026-08-01): stash the bonus on the session instead of
9312                    // committing it with a solo T=1 pass — the next empty-suffix greedy burst
9313                    // consumes it as round-0 verify col 0 (a plain round edge; the old tail
9314                    // commit + next burst's init feed were 11.6+11.5ms solo trunk passes per
9315                    // burst on H100 q27, [spec-setup] trace). b stays in `out` (emitted) but
9316                    // OUT of `committed` (cache rows == committed); the consuming call
9317                    // prepends it once its verify commits the row. next_pred is unknowable
9318                    // without the commit pass — None; callers gate on pending_tok too.
9319                    debug_assert_eq!(out.last(), Some(&b), "pending must be the last emitted");
9320                    if let Some(slot) = sess_pending_slot.take() {
9321                        *slot = Some(b);
9322                    }
9323                    *next_pred_slot = None;
9324                    // fill_prev = hidden of the last COMMITTED row (b's predecessor) — the
9325                    // exact chain-seed/fill anchor the consuming burst (or a flush) needs.
9326                    *last_h = Some(e.clone_dtod(&fill_prev)?);
9327                    stashed_pending = true;
9328                } else {
9329                    // SAMPLED tail (unchanged): commit the bonus (one T=1 pass) + draft fill —
9330                    // the sampled round-0 accept needs this pass's logits (last_col_logits).
9331                    let pos_b = cache.pos;
9332                    scratch.set_len(e, pos_b)?;
9333                    let (lg_b, hb) = self.spec_target_step_h(e, b, &mut *cache)?;
9334                    // after a FULL-accept exit `last_pred` is STALE (it predicted the bonus
9335                    // itself — the prediction AFTER the bonus never materialized; it would have
9336                    // been the next round's verify col 0). The commit's logits ARE that
9337                    // prediction.
9338                    *next_pred_slot = Some(argmax(&lg_b) as u32);
9339                    self.mtp_kv_fill(e, mtp, &[b], &fill_prev, pos_b, &mut *scratch, embd_dev)?;
9340                    *last_h = Some(hb);
9341                }
9342            } else {
9343                // fill_prev tracks the hidden of the last COMMITTED row throughout the loop.
9344                *last_h = Some(e.clone_dtod(&fill_prev)?);
9345            }
9346            committed.extend_from_slice(prompt);
9347            if let Some(cb) = carried_pending {
9348                // the consumed carry's cache row landed in round 0's verify (every pending
9349                // round commits col 0) — it joins `committed` here, in sequence order.
9350                committed.push(cb);
9351            }
9352            if stashed_pending {
9353                // ZERO-EMIT BURST (2026-08-06 c=8 serve panic, pre-existing since b4aea184):
9354                // `out.len() - 1` underflowed on an EMPTY `out` — "range end index
9355                // 18446744073709551615 out of range for slice of length 0", killing the
9356                // memra-gpu-worker and failing 31 of 32 concurrent requests with "worker closed
9357                // stream". Reachable because `pending` starts as `carried_pending` (a bonus
9358                // stashed by the PREVIOUS burst) while `out` starts empty, and the carry is
9359                // deliberately NOT pushed to `out` (line ~3239: the burst that emitted it already
9360                // did). So a burst that stashes a pending without emitting anything of its own —
9361                // the round loop exits before a push, e.g. the ring drain's `out.len() < max_new`
9362                // guard skipping every token under a tight budget — arrives here with
9363                // out.len() == 0 and stashed_pending == true.
9364                //
9365                // The invariant is unchanged: `committed` gets every emitted token EXCEPT the
9366                // stashed bonus. With nothing emitted, that is nothing — and the carry pushed
9367                // just above is already accounted. Saturating, not a min/assert: an empty `out`
9368                // here is a legitimate burst shape, not a corrupt state.
9369                let emitted = out.len().saturating_sub(1);
9370                committed.extend_from_slice(&out[..emitted]);
9371            } else {
9372                committed.extend_from_slice(&out); // FULL out incl. overshoot — all committed
9373            }
9374            debug_assert_eq!(
9375                cache.pos,
9376                committed.len(),
9377                "session invariant: cache rows == committed tokens"
9378            );
9379            if setup_trace {
9380                e.stream().synchronize()?; // bound the async tail fill in the trace
9381                let t_tail = t_ent.elapsed();
9382                eprintln!(
9383                    "[spec-setup] init={:.2}ms cap={:.2}ms fill={:.2}ms rounds={:.2}ms tail={:.2}ms total={:.2}ms out={} cont={}",
9384                    t_init.as_secs_f64() * 1e3,
9385                    (t_cap - t_init).as_secs_f64() * 1e3,
9386                    (t_fill - t_cap).as_secs_f64() * 1e3,
9387                    (t_rounds - t_fill).as_secs_f64() * 1e3,
9388                    (t_tail - t_rounds).as_secs_f64() * 1e3,
9389                    t_tail.as_secs_f64() * 1e3,
9390                    out.len(),
9391                    continuation
9392                );
9393            }
9394            return Ok((out, total_drafted, total_accepted));
9395        }
9396        out.truncate(max_new);
9397        Ok((out, total_drafted, total_accepted))
9398    }
9399
9400    /// Anchor-bounded DSpark target extraction. The trunk sees the exact generated token tape;
9401    /// only requested hidden rows and target-logit rows cross PCIe. An anchor token at p pairs
9402    /// with the pre-output-norm h[p-1] carrier, exactly as the existing replay/NextN path does.
9403    pub fn extract_dspark_anchors(
9404        &self,
9405        e: &Engine,
9406        tokens: &[u32],
9407        anchor_positions: &[usize],
9408        gamma: usize,
9409        top_k: usize,
9410        chunk: usize,
9411        temperature: f32,
9412    ) -> Result<Vec<DsparkAnchorRecord>, Box<dyn std::error::Error>> {
9413        if tokens.len() < gamma + 2 || gamma == 0 || chunk < 2 {
9414            return Err("DSpark extraction token tape/gamma/chunk is invalid".into());
9415        }
9416        if anchor_positions.windows(2).any(|pair| pair[0] >= pair[1]) {
9417            return Err("DSpark anchor positions must be sorted and unique".into());
9418        }
9419        for &position in anchor_positions {
9420            if position == 0 || position + gamma >= tokens.len() {
9421                return Err(format!(
9422                    "DSpark anchor {position} has no predecessor or cannot cover gamma={gamma} in {} tokens",
9423                    tokens.len()
9424                )
9425                .into());
9426            }
9427        }
9428
9429        let n_vocab = self.output.out_features();
9430        let n_embd = self.cfg.n_embd as usize;
9431        let mut cache = crate::pp::new_cache(e, &self.cfg, tokens.len() + gamma + 8)?;
9432        let (embd_qt, embd_rb) = self.embd.qt_and_row_bytes(n_embd);
9433        let embd_gpu = if spec_host_embd() {
9434            None
9435        } else {
9436            Some(
9437                self.embd_gpu
9438                    .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload")),
9439            )
9440        };
9441        let embd_dev = embd_gpu.map(|gpu| (gpu, embd_qt, embd_rb));
9442
9443        struct PendingRecord {
9444            position: usize,
9445            hidden: Option<Vec<f32>>,
9446            tokens: Vec<u32>,
9447            target_top_ids: Vec<Option<Vec<u32>>>,
9448            target_top_logits: Vec<Option<Vec<f32>>>,
9449            target_top_probs: Vec<Option<Vec<f32>>>,
9450            target_tail_probs: Vec<Option<f32>>,
9451        }
9452
9453        let mut pending: Vec<PendingRecord> = anchor_positions
9454            .iter()
9455            .map(|&position| PendingRecord {
9456                position,
9457                hidden: None,
9458                tokens: tokens[position..=position + gamma].to_vec(),
9459                target_top_ids: vec![None; gamma],
9460                target_top_logits: vec![None; gamma],
9461                target_top_probs: vec![None; gamma],
9462                target_tail_probs: vec![None; gamma],
9463            })
9464            .collect();
9465
9466        let mut start = 0usize;
9467        while start < tokens.len() {
9468            let end = (start + chunk).min(tokens.len());
9469            let chunk_tokens = &tokens[start..end];
9470            let (target_logits, hidden_rows) =
9471                self.decode_step_t_core(e, chunk_tokens, start, &mut cache, embd_dev, None)?;
9472            for record in &mut pending {
9473                let hidden_position = record.position - 1;
9474                if hidden_position >= start && hidden_position < end {
9475                    let local = hidden_position - start;
9476                    record.hidden = Some(
9477                        e.dtoh_view(&hidden_rows.slice(local * n_embd..(local + 1) * n_embd))?,
9478                    );
9479                }
9480                for slot in 0..gamma {
9481                    let target_row = record.position + slot;
9482                    if target_row < start || target_row >= end {
9483                        continue;
9484                    }
9485                    let local = target_row - start;
9486                    let logits =
9487                        e.dtoh_view(&target_logits.slice(local * n_vocab..(local + 1) * n_vocab))?;
9488                    let (ids, top_logits, probs, tail) =
9489                        dspark_sparse_softmax_topk(&logits, top_k, temperature)?;
9490                    record.target_top_ids[slot] = Some(ids);
9491                    record.target_top_logits[slot] = Some(top_logits);
9492                    record.target_top_probs[slot] = Some(probs);
9493                    record.target_tail_probs[slot] = Some(tail);
9494                }
9495            }
9496            start = end;
9497        }
9498
9499        pending
9500            .into_iter()
9501            .map(|record| {
9502                let hidden = record
9503                    .hidden
9504                    .ok_or_else(|| format!("missing DSpark hidden at {}", record.position))?;
9505                let target_top_ids =
9506                    flatten_dspark_rows(record.target_top_ids, record.position, "target ids")?;
9507                let target_top_logits = flatten_dspark_rows(
9508                    record.target_top_logits,
9509                    record.position,
9510                    "target logits",
9511                )?;
9512                let target_top_probs =
9513                    flatten_dspark_rows(record.target_top_probs, record.position, "target probs")?;
9514                let target_tail_probs = record
9515                    .target_tail_probs
9516                    .into_iter()
9517                    .enumerate()
9518                    .map(|(slot, value)| {
9519                        value.ok_or_else(|| {
9520                            format!("missing DSpark tail at {} slot {slot}", record.position)
9521                        })
9522                    })
9523                    .collect::<Result<Vec<_>, _>>()?;
9524                Ok(DsparkAnchorRecord {
9525                    position: record.position,
9526                    hidden,
9527                    tokens: record.tokens,
9528                    target_top_ids,
9529                    target_top_logits,
9530                    target_top_probs,
9531                    target_tail_probs,
9532                })
9533            })
9534            .collect()
9535    }
9536
9537    /// TEACHER-FORCED REPLAY ACCEPTANCE (hqmtp MTP-heal protocol): walk a FIXED token
9538    /// sequence and, at sampled positions, compare the MTP head's K-token draft chain against
9539    /// the trunk's own teacher-forced greedy predictions. Nothing is generated — the context is
9540    /// the corpus text itself, so (a) degenerate self-generated loops cannot inflate acceptance
9541    /// and (b) two arms (bf16 ceiling vs NVFP4) score on IDENTICAL contexts, isolating the
9542    /// quant-induced head/hidden-state mismatch from text drift.
9543    ///
9544    /// Per eval position p (context = tokens[0..=p], predecessor pairing as in spec decode):
9545    ///   draft_j  = chain token j from (tokens[p], h_{p-1}), then its own drafts — the exact
9546    ///              eager spec-decode chain (same mtp_head_forward_dev, same rope positions).
9547    ///   target_j = teacher-forced greedy pick for position p+1+j (argmax of the trunk logits
9548    ///              at forced context tokens[0..p+j]). For j==0 this equals live spec
9549    ///              acceptance; for j>=1 live verify would condition on the drafts, here it
9550    ///              conditions on the corpus — deterministic and arm-comparable by design.
9551    ///
9552    /// Returns (rows, bg): one (p, drafts[k], targets[k]) row per eval position (ascending p),
9553    /// plus the full teacher-forced greedy track bg (bg[i] = greedy pick for position i, i>=1)
9554    /// so harnesses can cross-check runs (e.g. different chunk sizes must give identical bg).
9555    ///
9556    /// `hdump`: when Some, every position's pre-output_norm trunk hidden (the exact rows the
9557    /// draft-KV fill pairs from) streams to the file as little-endian f32 [t_total, n_embd] —
9558    /// the head-distillation extraction (hqmtp): the ENGINE is the source of truth for trunk
9559    /// hiddens (HF torch reproductions of the hybrid trunk measured only ~0.5 greedy
9560    /// agreement vs this path — not usable as a training-data source).
9561    pub fn replay_acceptance(
9562        &self,
9563        e: &Engine,
9564        tokens: &[u32],
9565        k: usize,
9566        stride: usize,
9567        chunk: usize,
9568        mut hdump: Option<&mut std::fs::File>,
9569    ) -> Result<(Vec<(usize, Vec<u32>, Vec<u32>)>, Vec<u32>), Box<dyn std::error::Error>> {
9570        assert!(k >= 1 && stride >= 1 && chunk >= 2);
9571        let mtp = self
9572            .mtp
9573            .as_ref()
9574            .expect("replay_acceptance requires an MTP head");
9575        let n_vocab = self.output.out_features();
9576        let d_vocab = mtp
9577            .shared_head_head
9578            .as_ref()
9579            .unwrap_or(&self.output)
9580            .out_features();
9581        let n_embd = self.cfg.n_embd as usize;
9582        let t_total = tokens.len();
9583        assert!(t_total >= 8, "corpus too short ({t_total} tokens)");
9584        // STAGE-OWNED KV (lane/pp2-spec 2026-08-06) — see `new_session`. Door shut = `Cache::new`.
9585        let mut cache = crate::pp::new_cache(e, &self.cfg, t_total + k + 8)?;
9586        let mut scratch = MtpScratch::new(
9587            e,
9588            &self.cfg,
9589            t_total + k + 8,
9590            self.mtp.as_ref().and_then(|m| m.geom.as_ref()),
9591        )?;
9592        let (embd_qt, embd_rb) = self.embd.qt_and_row_bytes(n_embd);
9593        let embd_gpu = if spec_host_embd() {
9594            None
9595        } else {
9596            Some(
9597                self.embd_gpu
9598                    .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload")),
9599            )
9600        };
9601        let embd_dev = embd_gpu.map(|g| (g, embd_qt, embd_rb));
9602
9603        // bg[i] = the trunk's greedy pick for position i under the forced context (i >= 1).
9604        let mut bg: Vec<u32> = vec![0; t_total + 1];
9605        let mut rows: Vec<(usize, Vec<u32>, Vec<u32>)> = Vec::new();
9606        let mut prev_last_h = e.zeros(n_embd)?; // predecessor hidden entering the chunk
9607        let mut seed_buf = e.zeros(n_embd)?;
9608        let mut preds_d = e.alloc_u32_zeroed(chunk)?;
9609        let nll_on = std::env::var("MEMRA_REPLAY_NLL").as_deref() == Ok("1");
9610        let (mut nll_sum, mut nll_cnt) = (0f64, 0u64);
9611        let mut s = 0usize;
9612        while s < t_total {
9613            let cend = (s + chunk).min(t_total);
9614            let tc = cend - s;
9615            let ch = &tokens[s..cend];
9616            // 1. forced trunk pass — verify path (decode-exact contract): all-column logits +
9617            //    the chunk's true hiddens.
9618            let (tl_d, vx) = self.decode_step_t_core(e, ch, s, &mut cache, embd_dev, None)?;
9619            for j in 0..tc {
9620                e.argmax_token_device_col(&tl_d, j, n_vocab, &mut preds_d, j)?;
9621            }
9622            let preds = e.dtoh_u32(&preds_d)?;
9623            for j in 0..tc {
9624                bg[s + j + 1] = preds[j];
9625            }
9626            // MEMRA_REPLAY_NLL=1: teacher-forced NLL/perplexity over the same forced pass — the
9627            // checkpoint-quality metric (position j's logits score the GOLD next token).
9628            if nll_on {
9629                let jmax = if cend < t_total { tc } else { tc - 1 }; // last pos has no gold next
9630                if jmax > 0 {
9631                    let ids: Vec<u32> = (0..jmax).map(|j| tokens[s + j + 1]).collect();
9632                    let rows: Vec<i32> = (0..jmax as i32).collect();
9633                    let idsd = e.htod_u32_v(&ids)?;
9634                    let rowsd = e.htod_i32(&rows)?;
9635                    let mut outd = e.zeros(jmax)?;
9636                    e.softmax_gather(&tl_d, n_vocab, &idsd, &rowsd, &mut outd, n_vocab, jmax, 1.0)?;
9637                    for pr in e.dtoh(&outd)? {
9638                        nll_sum += -((pr.max(1e-30)) as f64).ln();
9639                        nll_cnt += 1;
9640                    }
9641                }
9642            }
9643            if let Some(f) = hdump.as_deref_mut() {
9644                use std::io::Write;
9645                let host: Vec<f32> = e.dtoh(&vx)?;
9646                // bf16 round-to-nearest-even — f32 doubled the disk bill at bulk
9647                // extraction scale (20M tokens x 4096 = 320GB f32 vs 160GB bf16).
9648                let mut bytes = Vec::with_capacity(tc * n_embd * 2);
9649                for v in &host[..tc * n_embd] {
9650                    let b = v.to_bits();
9651                    let r = b.wrapping_add(0x7FFF + ((b >> 16) & 1));
9652                    bytes.extend_from_slice(&((r >> 16) as u16).to_le_bytes());
9653                }
9654                f.write_all(&bytes)?;
9655            }
9656            // CHAINLESS extraction (stride > corpus, the bulk-hdump mode): no chunk ever
9657            // drafts, so the draft-KV fills are pure waste — skip them (2 MTP-block passes
9658            // per token saved; the forced trunk pass + hdump is all the mode needs).
9659            let chainless = stride > t_total;
9660            if chainless {
9661                e.copy_view_into(
9662                    &mut prev_last_h,
9663                    0,
9664                    &vx.slice((tc - 1) * n_embd..tc * n_embd),
9665                    n_embd,
9666                )?;
9667                s = cend;
9668                continue;
9669            }
9670            // 2. TRUE predecessor-paired draft-KV fill for the chunk (row i carries h_{i-1};
9671            //    row s reads the previous chunk's last true hidden, zeros at corpus start).
9672            let mut vxs = e.zeros(tc * n_embd)?;
9673            e.copy_into(&mut vxs, 0, &prev_last_h, n_embd)?;
9674            if tc > 1 {
9675                e.copy_view_into(
9676                    &mut vxs,
9677                    n_embd,
9678                    &vx.slice(0..(tc - 1) * n_embd),
9679                    (tc - 1) * n_embd,
9680                )?;
9681            }
9682            scratch.set_len(e, s)?;
9683            self.mtp_kv_fill(e, mtp, ch, &vxs, s, &mut scratch, embd_dev)?;
9684            // 3. draft chains at sampled positions, DESCENDING: a chain reads only slots
9685            //    [0..p) (true fills) and appends at >= p; the next (smaller-p) chain's set_len
9686            //    truncates those approximate appends before they can ever be read.
9687            let ps: Vec<usize> = (s..cend)
9688                .filter(|p| *p >= 1 && *p % stride == 0 && *p + k <= t_total)
9689                .collect();
9690            for &p in ps.iter().rev() {
9691                scratch.set_len(e, p)?;
9692                if p == s {
9693                    e.copy_into(&mut seed_buf, 0, &prev_last_h, n_embd)?;
9694                } else {
9695                    e.copy_view_into(
9696                        &mut seed_buf,
9697                        0,
9698                        &vx.slice((p - 1 - s) * n_embd..(p - s) * n_embd),
9699                        n_embd,
9700                    )?;
9701                }
9702                let mut e_tok = tokens[p];
9703                let mut d_seed = e.clone_dtod(&seed_buf)?;
9704                let mut drafts: Vec<u32> = Vec::with_capacity(k);
9705                for j in 0..k {
9706                    let (dl_d, h_nextn) = self.mtp_head_forward_dev(
9707                        e,
9708                        mtp,
9709                        e_tok,
9710                        &d_seed,
9711                        &mut scratch,
9712                        p + 1 + j,
9713                        embd_dev,
9714                        None, // acceptance-oracle walk: no grammar
9715                    )?;
9716                    let tok_d = e.argmax_token_device(&dl_d, d_vocab)?;
9717                    let idx = e.dtoh_u32_one(&tok_d)?;
9718                    let d = match &mtp.d2t {
9719                        Some(map) => map[idx as usize],
9720                        None => idx,
9721                    };
9722                    drafts.push(d);
9723                    e_tok = d;
9724                    d_seed = h_nextn;
9725                }
9726                // targets may live in a LATER chunk's bg — resolved after the walk.
9727                rows.push((p, drafts, Vec::new()));
9728            }
9729            // 4. restore TRUE entries for the whole chunk (the next chunk's chains and fills
9730            //    expect scratch.len == cend with exact rows).
9731            scratch.set_len(e, s)?;
9732            self.mtp_kv_fill(e, mtp, ch, &vxs, s, &mut scratch, embd_dev)?;
9733            e.copy_view_into(
9734                &mut prev_last_h,
9735                0,
9736                &vx.slice((tc - 1) * n_embd..tc * n_embd),
9737                n_embd,
9738            )?;
9739            s = cend;
9740        }
9741        for (p, drafts, targets) in rows.iter_mut() {
9742            for j in 0..drafts.len() {
9743                targets.push(bg[*p + 1 + j]);
9744            }
9745        }
9746        rows.sort_by_key(|r| r.0);
9747        if nll_cnt > 0 {
9748            let mean = nll_sum / nll_cnt as f64;
9749            println!(
9750                "[replay-nll] tokens={nll_cnt} nll/token={mean:.5} ppl={:.4}",
9751                mean.exp()
9752            );
9753        }
9754        Ok((rows, bg))
9755    }
9756}
9757
9758#[cfg(test)]
9759mod dspark_sparse_tests {
9760    use super::dspark_sparse_softmax_topk;
9761
9762    #[test]
9763    fn topk_keeps_full_softmax_mass_and_stable_ties() {
9764        let logits = [1.0f32, 3.0, 3.0, -2.0];
9765        let (ids, top_logits, probs, tail) = dspark_sparse_softmax_topk(&logits, 2, 1.0).unwrap();
9766        assert_eq!(ids, vec![1, 2]);
9767        assert_eq!(top_logits, vec![3.0, 3.0]);
9768        let denominator = logits.iter().map(|value| (value - 3.0).exp()).sum::<f32>();
9769        let expected = 1.0 / denominator;
9770        assert!((probs[0] - expected).abs() < 1.0e-6);
9771        assert!((probs[1] - expected).abs() < 1.0e-6);
9772        assert!((tail - (1.0 - 2.0 * expected)).abs() < 1.0e-6);
9773        assert!((probs.iter().sum::<f32>() + tail - 1.0).abs() < 1.0e-6);
9774    }
9775}
9776
9777#[cfg(test)]
9778mod spec_replay_env_tests {
9779    use super::spec_replay_env_on;
9780
9781    #[test]
9782    fn replay_requires_literal_one() {
9783        assert!(!spec_replay_env_on(None));
9784        assert!(!spec_replay_env_on(Some("")));
9785        assert!(!spec_replay_env_on(Some("0")));
9786        assert!(!spec_replay_env_on(Some("true")));
9787        assert!(!spec_replay_env_on(Some("2")));
9788        assert!(spec_replay_env_on(Some("1")));
9789    }
9790}
9791
9792#[cfg(test)]
9793mod telem_tests {
9794    use super::{SPEC_TELEM_POS, SpecTelemetry, SpecTelemetryCounters};
9795
9796    #[test]
9797    fn synthetic_accept_masks_produce_tau_and_position_histogram() {
9798        let counters = SpecTelemetryCounters::default();
9799        for mask in [
9800            [true, true, true],
9801            [true, true, false],
9802            [true, false, false],
9803            [false, false, false],
9804        ] {
9805            let accepted = mask.iter().take_while(|&&value| value).count();
9806            counters.record_round(mask.len(), accepted);
9807        }
9808
9809        let snapshot = counters.snapshot();
9810        assert_eq!(
9811            (snapshot.rounds, snapshot.drafted, snapshot.accepted),
9812            (4, 12, 6)
9813        );
9814        assert_eq!(&snapshot.pos_drafted[..3], &[4, 4, 4]);
9815        assert_eq!(&snapshot.pos_accepted[..3], &[3, 2, 1]);
9816        assert_eq!(snapshot.tau(), 1.5);
9817        assert_eq!(snapshot.pos_drafted[3..], [0; SPEC_TELEM_POS - 3]);
9818        assert_eq!(snapshot.pos_accepted[3..], [0; SPEC_TELEM_POS - 3]);
9819    }
9820
9821    /// The worker's per-burst pattern: stash, accumulate, diff — the delta must isolate
9822    /// exactly the burst's contribution (pool-resumed sessions carry prior requests' counts).
9823    #[test]
9824    fn delta_isolates_burst_contribution() {
9825        let mut t = SpecTelemetry::default();
9826        // "previous request": 2 rounds of k=3, accepts 3 then 1.
9827        for (kr, na) in [(3usize, 3usize), (3, 1)] {
9828            t.rounds += 1;
9829            t.drafted += kr as u64;
9830            t.accepted += na as u64;
9831            for j in 0..kr {
9832                t.pos_drafted[j] += 1;
9833            }
9834            for j in 0..na {
9835                t.pos_accepted[j] += 1;
9836            }
9837        }
9838        let before = t;
9839        // "this burst": 1 round k=3, accepts 2.
9840        t.rounds += 1;
9841        t.drafted += 3;
9842        t.accepted += 2;
9843        for j in 0..3 {
9844            t.pos_drafted[j] += 1;
9845        }
9846        for j in 0..2 {
9847            t.pos_accepted[j] += 1;
9848        }
9849        let d = t.delta_since(&before);
9850        assert_eq!((d.rounds, d.drafted, d.accepted), (1, 3, 2));
9851        assert_eq!(&d.pos_drafted[..3], &[1, 1, 1]);
9852        assert_eq!(&d.pos_accepted[..3], &[1, 1, 0]);
9853        assert_eq!(d.pos_drafted[3..], [0; SPEC_TELEM_POS - 3]);
9854    }
9855
9856    /// merge(delta) then merge(delta2) equals accumulating both — the per-model /metrics
9857    /// aggregation invariant.
9858    #[test]
9859    fn merge_accumulates_fieldwise() {
9860        let mut agg = SpecTelemetry::default();
9861        let mut d1 = SpecTelemetry {
9862            rounds: 2,
9863            drafted: 6,
9864            accepted: 4,
9865            ..Default::default()
9866        };
9867        d1.pos_drafted[0] = 2;
9868        d1.pos_accepted[0] = 2;
9869        let mut d2 = SpecTelemetry {
9870            rounds: 1,
9871            drafted: 3,
9872            accepted: 1,
9873            ..Default::default()
9874        };
9875        d2.pos_drafted[0] = 1;
9876        d2.pos_accepted[0] = 1;
9877        d2.pos_drafted[1] = 1;
9878        agg.merge(&d1);
9879        agg.merge(&d2);
9880        assert_eq!((agg.rounds, agg.drafted, agg.accepted), (3, 9, 5));
9881        assert_eq!(agg.pos_drafted[0], 3);
9882        assert_eq!(agg.pos_accepted[0], 3);
9883        assert_eq!(agg.pos_drafted[1], 1);
9884        assert_eq!(agg.pos_accepted[1], 0);
9885    }
9886
9887    /// Wrong-snapshot diff saturates to zero instead of wrapping — the counters feed a
9888    /// public metrics surface and must never publish a u64-wrapped garbage value.
9889    #[test]
9890    fn delta_saturates_never_wraps() {
9891        let small = SpecTelemetry {
9892            rounds: 1,
9893            drafted: 2,
9894            accepted: 1,
9895            ..Default::default()
9896        };
9897        let big = SpecTelemetry {
9898            rounds: 5,
9899            drafted: 15,
9900            accepted: 9,
9901            ..Default::default()
9902        };
9903        let d = small.delta_since(&big);
9904        assert_eq!((d.rounds, d.drafted, d.accepted), (0, 0, 0));
9905    }
9906}
9907
9908#[cfg(test)]
9909mod opti_fork_tests {
9910    use super::{
9911        OptiControllerPolicy, OptiForkAction, OptiForkGateMode, OptiForkGenerationTracker,
9912    };
9913
9914    #[test]
9915    fn controller_threshold_and_three_miss_breaker_are_exact() {
9916        let mut policy = OptiControllerPolicy {
9917            threshold: 0.7,
9918            consecutive_misses: 0,
9919            breaker_tripped: false,
9920        };
9921        assert!(!policy.admit(0.699_999));
9922        assert!(policy.admit(0.7));
9923        assert!(!policy.resolve(false));
9924        assert!(!policy.resolve(false));
9925        assert!(policy.resolve(false));
9926        assert!(policy.breaker_tripped);
9927        assert!(!policy.admit(1.0));
9928        assert!(
9929            !policy.resolve(true),
9930            "a resolved hit cannot re-arm a tripped request"
9931        );
9932        assert!(policy.breaker_tripped);
9933    }
9934
9935    #[test]
9936    fn zero_threshold_is_the_true_unconditional_measurement_arm() {
9937        let mut policy = OptiControllerPolicy {
9938            threshold: 0.0,
9939            consecutive_misses: 0,
9940            breaker_tripped: false,
9941        };
9942        for _ in 0..16 {
9943            assert!(policy.admit(0.0));
9944            assert!(!policy.resolve(false));
9945        }
9946        for invalid in [f32::NAN, f32::INFINITY, -0.01, 1.01] {
9947            assert!(
9948                !policy.admit(invalid),
9949                "invalid q proxy must fail closed: {invalid}"
9950            );
9951        }
9952        assert!(!policy.breaker_tripped);
9953        assert_eq!(policy.consecutive_misses, 0);
9954    }
9955
9956    #[test]
9957    fn alternating_mode_flips_by_generation_not_round_parity() {
9958        assert_eq!(OptiForkGateMode::Alternate.action(0), OptiForkAction::Hit);
9959        assert_eq!(OptiForkGateMode::Alternate.action(1), OptiForkAction::Miss);
9960        assert_eq!(OptiForkGateMode::Alternate.action(8), OptiForkAction::Hit);
9961        assert_eq!(OptiForkGateMode::Alternate.action(9), OptiForkAction::Miss);
9962    }
9963
9964    #[test]
9965    fn live_generation_cannot_be_overwritten() {
9966        let mut tracker = OptiForkGenerationTracker::default();
9967        let g0 = tracker.reserve().unwrap();
9968        let g1 = tracker.reserve().unwrap();
9969        let err = tracker.reserve().unwrap_err().to_string();
9970        assert!(
9971            err.contains("still owns generation 0"),
9972            "unexpected error: {err}"
9973        );
9974        tracker.retire(g0).unwrap();
9975        let g2 = tracker.reserve().unwrap();
9976        assert_eq!((g2.id, g2.slot), (2, 0));
9977        tracker.retire(g1).unwrap();
9978        tracker.retire(g2).unwrap();
9979    }
9980
9981    #[test]
9982    fn teardown_rejects_a_stale_generation_tag() {
9983        let mut tracker = OptiForkGenerationTracker::default();
9984        let g0 = tracker.reserve().unwrap();
9985        tracker.retire(g0).unwrap();
9986        let err = tracker.retire(g0).unwrap_err().to_string();
9987        assert!(err.contains("teardown mismatch"), "unexpected error: {err}");
9988    }
9989}
9990
9991#[cfg(test)]
9992mod draft_graph_fallback_tests {
9993    use super::DraftGraphFallback;
9994
9995    /// The Q2 contract, part (a): a fallback flip is LOUD — exactly once per flip.
9996    #[test]
9997    fn flip_is_loud_once_and_memoized_after() {
9998        let mut f = DraftGraphFallback::default();
9999        let line = f
10000            .mark_greedy("out of memory")
10001            .expect("first flip must return the warn line");
10002        assert!(
10003            line.contains("WARN"),
10004            "flip line must be warn-level: {line}"
10005        );
10006        assert!(
10007            line.contains("out of memory"),
10008            "flip line must carry the reason: {line}"
10009        );
10010        assert!(f.greedy_failed());
10011        // re-marking an already-failed graph is the memoization: quiet, still failed.
10012        assert!(f.mark_greedy("out of memory").is_none());
10013        assert!(f.greedy_failed());
10014        // the two graphs' flags are independent (greedy flip leaves sampled capturable).
10015        assert!(!f.sampled_failed());
10016        let line_s = f
10017            .mark_sampled("capture unsupported")
10018            .expect("sampled flip is its own flip");
10019        assert!(
10020            line_s.contains("sampled"),
10021            "sampled flip names itself: {line_s}"
10022        );
10023        assert!(f.mark_sampled("capture unsupported").is_none());
10024    }
10025
10026    /// The Q2 contract, part (b): resume-from-pool RESETS both flags (fresh capture chance),
10027    /// and says so exactly when there was something to reset.
10028    #[test]
10029    fn reset_on_resume_clears_flags_and_logs_once() {
10030        let mut f = DraftGraphFallback::default();
10031        // clean session: resume is silent, nothing to reset.
10032        assert!(f.reset_on_resume().is_none());
10033        f.mark_greedy("oom").unwrap();
10034        f.mark_sampled("oom").unwrap();
10035        let note = f
10036            .reset_on_resume()
10037            .expect("a set flag must produce the reset note");
10038        assert!(
10039            note.contains("greedy+sampled"),
10040            "note names what was reset: {note}"
10041        );
10042        assert!(
10043            !f.greedy_failed() && !f.sampled_failed(),
10044            "both flags cleared"
10045        );
10046        // and the NEXT failure after a reset is a fresh flip — loud again.
10047        assert!(f.mark_greedy("oom again").is_some());
10048        let note2 = f.reset_on_resume().expect("greedy-only reset");
10049        assert!(note2.contains("(greedy)"), "single-flag note: {note2}");
10050    }
10051
10052    /// Shape-change clears (dmask realloc / mask-shape mismatch / s_key change) stay silent —
10053    /// they precede a fresh capture attempt whose own failure re-flips loudly.
10054    #[test]
10055    fn shape_change_clears_are_silent() {
10056        let mut f = DraftGraphFallback::default();
10057        f.mark_greedy("oom").unwrap();
10058        f.clear_greedy();
10059        assert!(!f.greedy_failed());
10060        f.mark_sampled("oom").unwrap();
10061        f.clear_sampled();
10062        assert!(!f.sampled_failed());
10063        // after a silent clear there is nothing left for resume to report.
10064        assert!(f.reset_on_resume().is_none());
10065    }
10066}