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/// Engine-bundle slice 2 (DSF-ROUNDCOST-20260820 §1.1 host/device round trips + §2 rows 2-3),
179/// DEFAULT ON (`MEMRA_DSPARK_DEFER_READBACK=0` reverts): the dspark round's draft-chain DtoH
180/// is DEFERRED past verify dispatch and merged with the verify-argmax readback into ONE host
181/// sync (2 blocking DtoH/round -> 1). Verify embeds DEVICE tokens (`chain_d`) through the
182/// resident embed table — `embed_gather_u32_t`, bit-identical rows to the host gather by its
183/// own pinned contract. The host therefore dispatches snap + the whole verify while the DRAFT
184/// is still executing, instead of blocking ~1.7 ms on the chain and letting the device drain.
185/// Ladder arm only: the confidence policies size vt from a pre-verify head readback (their
186/// chain readback merges into that same sync instead). Exactness unchanged BY CONSTRUCTION —
187/// same tokens, same kernels, same order; E2E + accept-bank gates arbitrate.
188pub(crate) fn dspark_defer_readback_on() -> bool {
189 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
190 *ON.get_or_init(|| {
191 std::env::var("MEMRA_DSPARK_DEFER_READBACK")
192 .map(|v| v != "0")
193 .unwrap_or(true)
194 })
195}
196/// Engine-bundle slice 1 (DSF-ROUNDCOST-20260820 §1.1, lane/dspark-engine-bundle-20260820),
197/// DEFAULT ON (`MEMRA_STATE_COPY_BATCH=0` reverts): batch the dspark round's GDN state
198/// snapshot and partial-accept restore into single `copy_batch_uniform_f32` launches
199/// instead of ~2 memcpy dispatches (+2 alloc_zeros on the snap side) per linear layer per
200/// round — measured 0.67 ms/round snap + 0.25 ms/round commit of pure dispatch on the q38
201/// route. Launch-structure only: bytes, buffers and stream order are unchanged, so
202/// acceptance and streams stay bit-identical (E2E-gated on the B1 packs).
203pub(crate) fn state_copy_batch_on() -> bool {
204 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
205 *ON.get_or_init(|| {
206 std::env::var("MEMRA_STATE_COPY_BATCH")
207 .map(|v| v != "0")
208 .unwrap_or(true)
209 })
210}
211/// Engine-bundle slice 3 (DSF-ROUNDCOST-20260820 §5 rank 1, bucketed verify graphs),
212/// DEFAULT OFF — `MEMRA_DSPARK_VERIFY_GRAPH=1` opts in: the dspark verify's LINEAR-layer
213/// runs replay per-(segment, vt) CUDA graphs — see [`DsparkVerifyGraphs`]. Requires the
214/// slice-2 deferred path (device tokens); the eager walk is the byte-identical fallback.
215///
216/// MEASURED disposition (box6 card0, agentic pack, 2026-08-20): exactness holds (ALL
217/// EXACT, accept lines byte-match the bank) and steady-state prompts gain +0.4..+1.5
218/// tok/s, but the gate-scale MEAN is flat (109.2 vs 109.4): the AUTO_FREE launch scan
219/// costs 25.6 us x 16 graph launches/round (~0.41 ms — most of the eager-launch
220/// savings) and the first generation pays ~80 lazy captures. The residual verify gap
221/// lives in the FULL-ATTENTION per-row section (fa partial memsets + appends), i.e.
222/// the fa exec-update extension, not here. Off until one of: node-count reduction
223/// (ctx-scratch transients), WithParams UPLOAD instantiation, or the full-verify
224/// single-graph — each re-gated by the same battery.
225pub(crate) fn dspark_verify_graph_on() -> bool {
226 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
227 *ON.get_or_init(|| std::env::var("MEMRA_DSPARK_VERIFY_GRAPH").as_deref() == Ok("1"))
228}
229
230/// `t_pred0` for the `MEMRA_DEBUG_SPEC` per-round print, sampled-safe.
231///
232/// `generate_spec_inner2` fills its `preds` vector ONLY on the greedy path (`if !sampled`), and
233/// the per-round debug print was the sole consumer in the sampled arm: `t_pred(0)` survives round
234/// 0 (`base == 0` returns `last_pred`) and from round 1 (`base == 1`, a pending bonus) indexes an
235/// EMPTY vector — `index out of bounds: the len is 0 but the index is 0`, in the GPU worker
236/// thread, which then respawns and reloads weights while the request dies. So any sampled spec
237/// request longer than one round used to kill the worker whenever `MEMRA_DEBUG_SPEC` was set:
238/// the flag crashed precisely the regime it exists to investigate.
239///
240/// Fixed at the print site, not inside the closure, so the greedy accept walk keeps its strict
241/// indexing (an out-of-range pred there is a real bug and must still be loud).
242fn debug_t_pred0(sampled: bool, base: usize, last_pred: u32, preds: &[u32]) -> String {
243 if base == 0 {
244 return last_pred.to_string();
245 }
246 match preds.get(base - 1) {
247 Some(p) => p.to_string(),
248 // sampled: the greedy per-column argmax was never run for this round.
249 None => {
250 debug_assert!(
251 sampled,
252 "greedy spec: preds[{}] missing at base {base}",
253 base - 1
254 );
255 "n/a".to_string()
256 }
257 }
258}
259
260/// `MEMRA_SKEY_PROBE=1` — sampled-draft-graph key probe (lane/graph-s-key-exactness-20260819).
261///
262/// Reports, per burst and per round, which draft chain the sampled arm chose and under which
263/// filter regime, plus the ONE observable that separates a legal filtered draft from a stale
264/// pure-temp graph replayed under filters: an accept test whose gathered `q` is exactly 0.
265/// A draft token sampled from the FILTERED softmax can never gather q=0 (it was drawn from the
266/// kept set), so `q=0` in the verify means the draft came from a distribution the verify does
267/// not believe in — and `u * 0 < p` then accepts it unconditionally.
268///
269/// Its own env var, deliberately NOT `MEMRA_DEBUG_SPEC`: that flag panicked the GPU worker on
270/// any sampled spec request past round 0 until this lane fixed it (§2 of the bank note).
271pub(crate) fn skey_probe() -> bool {
272 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
273 *ON.get_or_init(|| std::env::var("MEMRA_SKEY_PROBE").as_deref() == Ok("1"))
274}
275
276/// GRAMMAR HOOK for constrained spec decode (lane/constrained-full, 2026-08-03). The engine
277/// stays llguidance-agnostic: the server adapts its per-session grammar state behind this
278/// trait. CONTRACT (the verify-side truncation rule — token-identical to constrained plain
279/// greedy decode): the exactness walk runs UNMASKED first; the hook then (a) truncates
280/// acceptance at the first grammar-illegal accepted token, and (b) when the truncation fired
281/// or the bonus is illegal, the engine recomputes that slot as the MASKED argmax of the
282/// target's own verify column (an unmasked argmax that is grammar-legal IS the masked argmax
283/// — masking only removes tokens — so the common case pays nothing). `consume` advances the
284/// state with each EMITTED token in order; EOS handling is the implementor's job (skip).
285pub trait SpecConstraint {
286 /// -inf the current state's banned ids on a HOST logits row (prompt-tail / init-feed
287 /// masked argmax).
288 fn mask_logits(&mut self, logits: &mut [f32]) -> Result<(), String>;
289 /// Packed 32-bit bitset words of the CURRENT state's allowed set (device-mask form).
290 fn mask_words(&mut self) -> Result<Vec<u32>, String>;
291 /// Is `tok` consumable in the CURRENT state?
292 fn is_allowed(&mut self, tok: u32) -> Result<bool, String>;
293 /// Advance the state with an emitted token.
294 fn consume(&mut self, tok: u32) -> Result<(), String>;
295
296 // --- DRAFT-SIDE MASKING (lane/draft-mask, 2026-08-04) ---
297 // The drafter proposed grammar-illegal tokens under tight schemas, so verify-side
298 // truncation cut nearly every round (measured acceptance 0.467-0.513 tight vs 0.62-0.82
299 // loose, research/constrained-full-20260803). These three methods let the engine mask the
300 // DRAFT model's own sampling with the grammar's legal set, so proposals are legal by
301 // construction. The state they walk is a SPECULATIVE CLONE of the session matcher — the
302 // real state is advanced only by `consume` (emitted tokens), so verify-side truncation
303 // stays the correctness backstop and the emitted stream is unchanged by construction
304 // (an accepted draft is the target's unmasked argmax AND grammar-legal, hence the masked
305 // argmax; a cut slot is recomputed as the masked argmax either way).
306 // Default impls = feature OFF (pre-lane behaviour: unmasked drafts).
307
308 /// Is draft-side masking available on this hook? Probed ONCE per burst, before the draft
309 /// graph is captured (the mask is an in-graph node — its presence is a capture-time shape).
310 fn draft_mask_enabled(&self) -> bool {
311 false
312 }
313 /// Start a draft chain: clone the CURRENT (committed) grammar state into the speculative
314 /// slot. Called once per spec round, before the first draft position.
315 fn draft_begin(&mut self) -> Result<(), String> {
316 Ok(())
317 }
318 /// Packed 32-bit bitset words of the SPECULATIVE state's allowed set (target-vocab ids),
319 /// for the draft position about to be sampled. `None` = draft masking off (no-op).
320 fn draft_mask_words(&mut self) -> Result<Option<Vec<u32>>, String> {
321 Ok(None)
322 }
323 /// Advance the SPECULATIVE state with a PROPOSED draft token. `false` = the chain cannot
324 /// continue (EOS proposed, or an unmasked position proposed something illegal) — the
325 /// engine stops drafting; the token already pushed still goes through verify.
326 fn draft_advance(&mut self, _tok: u32) -> Result<bool, String> {
327 Ok(false)
328 }
329}
330
331/// DRAFT-MASK UPLOAD (lane/draft-mask): pull the speculative state's allowed set (TARGET-id
332/// space) from the hook, project it into the DRAFT head's vocab space, and upload it into the
333/// stable device buffer the draft chain reads. Returns false when the chain must stop drafting:
334/// the hook handed out no mask, or NO draft-vocab row is grammar-legal at this position (a
335/// trimmed FR-Spec head genuinely cannot propose a legal token there — masking it would leave
336/// a fully-banned row whose argmax is meaningless, so the round drafts fewer tokens and the
337/// verify emits the masked argmax as usual).
338fn upload_draft_mask(
339 e: &Engine,
340 c: &mut dyn SpecConstraint,
341 dst: &mut CudaSlice<u32>,
342 d2t: Option<&Vec<u32>>,
343 d_vocab: usize,
344 words: usize,
345) -> Result<bool, Box<dyn std::error::Error>> {
346 let Some(tw) = c
347 .draft_mask_words()
348 .map_err(|e2| format!("constraint: {e2}"))?
349 else {
350 return Ok(false);
351 };
352 let bit = |t: usize| -> bool {
353 let w = t >> 5;
354 w < tw.len() && (tw[w] >> (t & 31)) & 1 == 1
355 };
356 let mut buf = vec![0u32; words];
357 match d2t {
358 // TRIMMED draft head: row i proposes target id d2t[i] — permute the mask accordingly.
359 Some(map) => {
360 for (i, &t) in map.iter().enumerate().take(d_vocab) {
361 if bit(t as usize) {
362 buf[i >> 5] |= 1u32 << (i & 31);
363 }
364 }
365 }
366 // UNTRIMMED: draft ids ARE target ids; the packed words transfer verbatim (a short
367 // mask leaves the padded tail zeroed == banned, same rule as constrained::apply_mask).
368 None => {
369 let n = tw.len().min(words);
370 buf[..n].copy_from_slice(&tw[..n]);
371 }
372 }
373 if buf.iter().all(|w| *w == 0) {
374 return Ok(false);
375 }
376 e.htod_u32_into(dst, &buf)?;
377 Ok(true)
378}
379
380/// Keep the full token-embedding table in host memory and upload only the rows needed by each
381/// MTP/verify step. This is an exact memory-capacity seam for very large BF16 vocab tables: host
382/// gather expands the same source bits to f32, and only O(T*n_embd) bytes cross PCIe per step.
383/// CUDA-graph/round-stream draft paths require device token ids and therefore stay disabled.
384pub(crate) fn spec_host_embd() -> bool {
385 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
386 *ON.get_or_init(|| std::env::var("MEMRA_SPEC_HOST_EMBD").as_deref() == Ok("1"))
387}
388
389/// VERIFY-TIER TRUNK LAUNCH-FUSION (default ON since 2026-07-09; MEMRA_SPEC_FUSED_T=0 reverts — lane/close35b): extend
390/// the t=1 fused2/fused3 Q8_0 trunk launches to the batched verify tier (t=2-4, the K=1..3
391/// verify shapes). At t>1 the trunk pairs/triples (35B wqkv+wqkv_gate, wq/wk/wv,
392/// gate_shexp+up_shexp) each run a separate `matmul_decode_exact` — one q8_1 re-quantize of the
393/// SAME activation plus one _b2/_b4 launch per tensor. The fused twins share ONE quantize and
394/// ONE launch per group; per (tensor,token,row) the kernel body is q8_0_mmvq_batched verbatim
395/// with the identical row mapping -> BIT-IDENTICAL by construction (kernel-check pins it,
396/// run-spec K=1..8 + acceptance identity arbitrate e2e).
397pub(crate) fn spec_fused_t() -> bool {
398 static F: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
399 // DEFAULT ON since 2026-07-09 (MEMRA_SPEC_FUSED_T=0 reverts): verify t=2-4 trunk launch-fusion
400 // (fused2/fused3 Q8_0 batched twins, bit-identical by construction — m=1 block-offset split on
401 // the batched body). m=2 marginal token 2117->1762us; 35B daily: p3 +3.7% (crosses llama), p2 +5%.
402 *F.get_or_init(|| {
403 std::env::var("MEMRA_SPEC_FUSED_T")
404 .map(|v| v != "0")
405 .unwrap_or(true)
406 })
407}
408
409/// zeros/uninit switch for verify-path buffers that are FULLY OVERWRITTEN before any read.
410/// Only call this on such buffers — the lean contract is "identical bytes by construction".
411fn vbuf(e: &Engine, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
412 if spec_lean() { e.uninit(n) } else { e.zeros(n) }
413}
414
415/// Scratch KV for the MTP block (one full-attn layer).
416///
417/// PERSISTENT MODE (default, 2026-07-03 — the acceptance lever): sized cap = max_ctx and kept in
418/// sync with the COMMITTED sequence — slot p holds the MTP block's K/V for committed token p
419/// (roped p+1, the chain's rope convention), so the draft chain's self-attention sees the FULL
420/// committed history instead of only the current round's 1..K+1 chain tokens (the reference
421/// engine's "mtp_update" design). Entries come from two sources:
422/// - chain appends: accepted positions KEEP their chain-computed entries (embedding exact,
423/// hidden chain-approximate — the reference engine accepts the same);
424/// - `mtp_kv_fill` batches: prompt positions + the last-draft position on full accept, computed
425/// from EXACT trunk hiddens (K/V-only MTP-block pass, no attention/FFN/lm_head).
426/// Rejected drafts / p-min extras / pseudo-seed appends are all discarded by the round-start
427/// `set_len` truncation (the KvLayer len mechanism — §C rollback for the draft side).
428/// Multi-turn spec-decode session (2026-07-05): trunk Cache + persistent MTP draft scratch +
429/// the committed token list, alive across generate_spec_session calls. Turn N+1 primes ONLY its
430/// suffix (chunked continuation prime over the quantized past) and mtp_kv_fill's its suffix rows,
431/// then the round loop runs unchanged. `last_h` carries the pre-output_norm hidden of the last
432/// committed row across turns (the predecessor-pairing seed + fill anchor).
433/// Per-request sampling config for the sampled-spec serve path.
434#[derive(Clone, Copy, Debug)]
435pub struct SpecSampling {
436 pub temp: f32,
437 pub seed: u64,
438 pub top_k: i32, // 0 = off
439 pub top_p: f32, // 1.0 = off
440 pub min_p: f32, // 0.0 = off
441 pub penalty_last_n: usize, // 0 = penalties off
442 pub penalty_repeat: f32,
443 pub penalty_freq: f32,
444 pub penalty_present: f32,
445}
446
447/// Tracked draft positions for [`SpecTelemetry`] (serve K defaults to 3; the run-spec gate
448/// sweeps K=1..8, and MEMRA_SPEC_CAPMAX defaults to 7 — 8 covers every tuned config).
449pub const SPEC_TELEM_POS: usize = 8;
450
451/// Always-on per-draft-position acceptance telemetry (lane/accept-telemetry, 2026-08-05 —
452/// the llama.cpp #26389 / vLLM spec-decode counter schema, upstream-sweeps 2026-08-05).
453/// Lives on the [`SpecSession`] and accumulates across bursts; the serve worker diffs a
454/// stashed copy per burst for its per-model /metrics aggregation and per-request usage.
455/// Same normalization as the `[spec-stats]` line: p-min-discarded chain tokens are counted
456/// in NEITHER drafted nor accepted.
457#[derive(Clone, Copy, Default, Debug)]
458pub struct SpecTelemetry {
459 /// verify rounds completed (a round-stream burst counts each of its M rounds).
460 pub rounds: u64,
461 /// tokens drafted / accepted across all rounds.
462 pub drafted: u64,
463 pub accepted: u64,
464 /// how often draft position j (0-based within a round's chain) was offered / accepted.
465 /// Positions >= SPEC_TELEM_POS are untracked (totals still count them). The opt-in
466 /// round-stream arm (MEMRA_SPEC_STREAM=1) reads back only totals, so under it these
467 /// arrays cover the standard-path rounds only and their sums may undercount the totals.
468 pub pos_drafted: [u64; SPEC_TELEM_POS],
469 pub pos_accepted: [u64; SPEC_TELEM_POS],
470}
471
472impl SpecTelemetry {
473 /// Fieldwise `self - prev` — the worker's per-burst delta off a copy stashed before the
474 /// burst call. Saturating: a caller diffing against the wrong snapshot gets zeros, not
475 /// a wrapped counter.
476 pub fn delta_since(&self, prev: &SpecTelemetry) -> SpecTelemetry {
477 let mut d = SpecTelemetry {
478 rounds: self.rounds.saturating_sub(prev.rounds),
479 drafted: self.drafted.saturating_sub(prev.drafted),
480 accepted: self.accepted.saturating_sub(prev.accepted),
481 ..Default::default()
482 };
483 for j in 0..SPEC_TELEM_POS {
484 d.pos_drafted[j] = self.pos_drafted[j].saturating_sub(prev.pos_drafted[j]);
485 d.pos_accepted[j] = self.pos_accepted[j].saturating_sub(prev.pos_accepted[j]);
486 }
487 d
488 }
489 /// Fieldwise `self += d` — the worker's per-model aggregation.
490 pub fn merge(&mut self, d: &SpecTelemetry) {
491 self.rounds += d.rounds;
492 self.drafted += d.drafted;
493 self.accepted += d.accepted;
494 for j in 0..SPEC_TELEM_POS {
495 self.pos_drafted[j] += d.pos_drafted[j];
496 self.pos_accepted[j] += d.pos_accepted[j];
497 }
498 }
499
500 /// Mean accepted draft-prefix length per verify round (tau).
501 pub fn tau(&self) -> f64 {
502 if self.rounds > 0 {
503 self.accepted as f64 / self.rounds as f64
504 } else {
505 0.0
506 }
507 }
508}
509
510/// Session-lifetime atomic acceptance counters. The verifier records only after the greedy or
511/// rejection-sampling walk has resolved on the host, so these relaxed increments add no GPU
512/// launch, synchronization, allocation, or ordering dependency to the numeric path.
513struct SpecTelemetryCounters {
514 rounds: AtomicU64,
515 drafted: AtomicU64,
516 accepted: AtomicU64,
517 pos_drafted: [AtomicU64; SPEC_TELEM_POS],
518 pos_accepted: [AtomicU64; SPEC_TELEM_POS],
519}
520
521impl Default for SpecTelemetryCounters {
522 fn default() -> Self {
523 Self {
524 rounds: AtomicU64::new(0),
525 drafted: AtomicU64::new(0),
526 accepted: AtomicU64::new(0),
527 pos_drafted: std::array::from_fn(|_| AtomicU64::new(0)),
528 pos_accepted: std::array::from_fn(|_| AtomicU64::new(0)),
529 }
530 }
531}
532
533impl SpecTelemetryCounters {
534 fn record_round(&self, drafted: usize, accepted: usize) {
535 debug_assert!(accepted <= drafted);
536 self.rounds.fetch_add(1, Ordering::Relaxed);
537 self.drafted.fetch_add(drafted as u64, Ordering::Relaxed);
538 self.accepted.fetch_add(accepted as u64, Ordering::Relaxed);
539 for counter in self.pos_drafted.iter().take(drafted) {
540 counter.fetch_add(1, Ordering::Relaxed);
541 }
542 for counter in self.pos_accepted.iter().take(accepted) {
543 counter.fetch_add(1, Ordering::Relaxed);
544 }
545 }
546
547 /// Round-stream keeps each round's accept length on device; retain exact scalar totals while
548 /// leaving the per-position arrays untouched, matching the pre-existing telemetry contract.
549 fn record_totals(&self, rounds: usize, drafted: usize, accepted: usize) {
550 self.rounds.fetch_add(rounds as u64, Ordering::Relaxed);
551 self.drafted.fetch_add(drafted as u64, Ordering::Relaxed);
552 self.accepted.fetch_add(accepted as u64, Ordering::Relaxed);
553 }
554
555 fn snapshot(&self) -> SpecTelemetry {
556 SpecTelemetry {
557 rounds: self.rounds.load(Ordering::Relaxed),
558 drafted: self.drafted.load(Ordering::Relaxed),
559 accepted: self.accepted.load(Ordering::Relaxed),
560 pos_drafted: std::array::from_fn(|j| self.pos_drafted[j].load(Ordering::Relaxed)),
561 pos_accepted: std::array::from_fn(|j| self.pos_accepted[j].load(Ordering::Relaxed)),
562 }
563 }
564}
565
566pub struct SpecSession {
567 pub(crate) cache: Cache,
568 pub(crate) scratch: MtpScratch,
569 /// Every token whose state the caches hold, in order (prompt turns + generated), INCLUDING
570 /// overshoot: spec commits accepted drafts past max_new; those rows are in the caches, so the
571 /// session must count them. Callers render output from this, not from their own echo.
572 pub committed: Vec<u32>,
573 /// Pre-output_norm hidden of the LAST committed row (device). None before the first turn.
574 pub(crate) last_h: Option<CudaSlice<f32>>,
575 /// Greedy argmax predicting the token AFTER committed.last() (from the last turn's final
576 /// logits). Fuels empty-suffix continuation bursts (serve): the next turn emits this token
577 /// first, feeds it, and the round loop resumes without any prime. None before the first turn.
578 pub next_pred: Option<u32>,
579 /// SAMPLED-SPEC stream continuity across bursts: Philox event counters persist here so a
580 /// session's randomness never repeats between generate_spec_session calls. (0,0) at admit.
581 pub sctr: u32,
582 pub uctr: u32,
583 /// PERSISTENT DRAFT-GRAPH CONTEXT (2026-08-01, the serve-burst fixed-cost fix): the captured
584 /// draft graph(s) + every device I/O buffer they bake, carried ACROSS generate_spec_session
585 /// calls. Before this, every serve burst re-captured the draft graph (2 warmup forwards +
586 /// instantiate) — measured ~16ms/burst on H100 q27 (MEMRA_SPEC_BURST sweep,
587 /// research/spec-serving-20260801). None before the first turn; error paths drop it
588 /// (next burst recaptures — serve retires errored sessions anyway).
589 pub(crate) draft_ctx: Option<DraftGraphCtx>,
590 /// PENDING-CARRY across bursts (2026-08-01, the serve burst-boundary fix): the bonus token
591 /// emitted by the last round but NOT committed to the caches. The old tail committed it with
592 /// a solo T=1 trunk pass (+ draft fill), and the next burst's setup fed the stashed next_pred
593 /// with ANOTHER solo pass — 2x ~11.5ms/burst measured on H100 q27 ([spec-setup] trace).
594 /// Carrying it lets the next empty-suffix greedy burst consume it as round-0 verify col 0,
595 /// exactly like a mid-burst full-accept boundary (no solo passes). INVARIANT: when set,
596 /// `committed` (== cache rows) EXCLUDES this token although it was already emitted in the
597 /// last burst's output, and `last_h` holds the hidden of the last COMMITTED row (its
598 /// predecessor — the chain-seed/fill anchor). `next_pred` is None (unknown without the
599 /// commit pass). Non-empty-suffix or sampled turns must flush first (spec_flush_pending);
600 /// generate_spec_session_sampled does this at entry, and serve parks only flushed sessions.
601 pub pending_tok: Option<u32>,
602 /// SESSION-AFFINITY TURN CHECKPOINT (lane/session-affinity, 2026-08-05): the state at this
603 /// turn's PROMPT-END boundary, retained so a later turn can REWIND here. See
604 /// [`SpecCheckpoint`]. Refreshed by every non-empty prime; None until the first one, and on
605 /// a rig too tight to hold it (a failed capture is silent — resume just isn't available).
606 pub(crate) turn_ckpt: Option<SpecCheckpoint>,
607 /// Session-lifetime acceptance telemetry. Relaxed atomics update at the host-side round
608 /// accounting the loop already does — no syncs, no allocation. NOTE a
609 /// pool-resumed session carries the PREVIOUS requests' counts; per-request consumers
610 /// diff with [`SpecTelemetry::delta_since`] around each burst.
611 telem: SpecTelemetryCounters,
612 /// PREFIX-CACHE publication request (lane/spec-prefix-cache): worker sets this to the
613 /// miss-LCP boundary before a cold burst; the prime captures at exactly that split (it must
614 /// coincide with the burst's `prime_split` or no capture happens). One-shot: consumed by the
615 /// prime, result lands in `boundary_capture`.
616 pub capture_at: Option<usize>,
617 /// The capture the last cold prime produced (see [`SpecBoundaryCapture`]). Worker takes it
618 /// post-burst to assemble the prefix entry. A failed capture is silent, like `turn_ckpt` —
619 /// publication just isn't available for that request.
620 pub boundary_capture: Option<SpecBoundaryCapture>,
621}
622impl SpecSession {
623 /// Context capacity of the session's caches (the server's ContextFull guard).
624 pub fn cache_max_ctx(&self) -> usize {
625 self.cache.max_ctx
626 }
627 /// Read access to the live trunk cache (lane/spec-prefix-cache): the worker slices
628 /// full-attn KV rows `[0..capture.pos)` out of it when publishing a boundary capture —
629 /// those rows are append-only for the session's lifetime (rollbacks never truncate below
630 /// the prime boundary), so no copy was taken at prime time.
631 pub fn cache_ref(&self) -> &Cache {
632 &self.cache
633 }
634 /// Read access to the persistent draft-scratch plane (lane/spec-on-cache-hit): the
635 /// worker slices rows `[0..capture.pos)` when publishing a boundary capture, exactly
636 /// like the trunk KV — draft rows below the prompt end are append-only for the
637 /// session's lifetime (the prime fill wrote them once; rollbacks reset `len_d` to the
638 /// committed length, never below the prime boundary, and the true-hidden refresh
639 /// rewrites generated positions only). Returns `(k, v, k_tok_bytes, v_tok_bytes)`.
640 /// None when the scratch is ring-backed (Step35 SWA — physical rows are not
641 /// prefix-addressable; the prefix cache already refuses that class end to end).
642 pub fn draft_plane_ref(&self) -> Option<(&CudaSlice<u8>, &CudaSlice<u8>, usize, usize)> {
643 if self.scratch.kv.ring.is_some() {
644 return None;
645 }
646 Some((
647 &self.scratch.kv.k,
648 &self.scratch.kv.v,
649 self.scratch.kv.k_tok_bytes,
650 self.scratch.kv.v_tok_bytes,
651 ))
652 }
653 /// Snapshot the session's process-local acceptance counters for per-burst diffing.
654 pub fn telemetry(&self) -> SpecTelemetry {
655 self.telem.snapshot()
656 }
657 /// Committed position this session can REWIND to (its retained prompt-end boundary), if any.
658 /// A request whose prompt matches `committed[..pos]` exactly can resume from here — see
659 /// `spec_rewind_to_checkpoint`.
660 pub fn rewind_pos(&self) -> Option<usize> {
661 self.turn_ckpt.as_ref().map(|c| c.pos)
662 }
663 /// Whether every ring-backed trunk/draft row needed by the retained checkpoint is resident.
664 pub fn rewind_is_resident(&self) -> bool {
665 self.turn_ckpt.as_ref().is_some_and(|ckpt| {
666 self.cache.can_rollback(&ckpt.snap, 0) && self.scratch.can_rewind_to(ckpt.pos)
667 })
668 }
669 /// Is this session in the DEMOTION-READY shape (see [`SpecSession::into_demoted`])?
670 /// `false` means a carried pending must be flushed first (`spec_flush_pending`), or the
671 /// session has never run a turn and has no prediction to hand over.
672 pub fn demote_ready(&self) -> bool {
673 self.pending_tok.is_none() && self.next_pred.is_some()
674 }
675 /// Does this session hold a carried pending bonus (flush required before a handoff/park)?
676 pub fn has_pending(&self) -> bool {
677 self.pending_tok.is_some()
678 }
679 /// Committed row count == cache rows (the session invariant), for the caller's own
680 /// `fed`-length cross-check at a handoff boundary.
681 pub fn committed_len(&self) -> usize {
682 self.committed.len()
683 }
684 /// DEMOTION HANDOFF (lane/spec-gate, 2026-08-07): consume this session and hand its trunk
685 /// cache + next-token prediction to the plain batched-decode path.
686 ///
687 /// WHY THIS IS EXACT (greedy). The invariant at a burst boundary is `cache.pos ==
688 /// committed.len()`: every committed row has trunk KV + recurrent state, exactly as a plain
689 /// tokenwise prime of the same `committed` sequence would have left it (that is the
690 /// session-tail contract, and the same property `spec_rewind_to_checkpoint` and the reuse
691 /// pool already rely on). `next_pred` is the argmax of the verify's logits for the LAST
692 /// committed row — and verify-column logits are bit-identical to plain decode's logits at
693 /// that position, because `matmul_decode_exact` bit-identity IS the basis of the greedy
694 /// accept walk. So handing (cache, next_pred) to the batched path continues the stream from
695 /// a state indistinguishable from one the batched path produced itself: the batched tick
696 /// emits `next_pred`, feeds it into this same cache, and decodes on.
697 ///
698 /// `None` when the session is not in the handoff shape — a carried pending (its bonus row is
699 /// NOT in the cache, so `spec_flush_pending` must commit it first) or no `next_pred` yet
700 /// (never bursted). Callers must not force it: a half-committed cache handed to the batched
701 /// path would silently skip a token.
702 ///
703 /// The MTP draft scratch, the persistent draft-graph context and the turn checkpoint are
704 /// DROPPED here (freeing their VRAM): the batched path never drafts, and this handoff is
705 /// one-way by design — there is no cheap symmetric re-promotion (rebuilding the draft KV
706 /// would mean an `mtp_kv_fill` over the whole committed history).
707 pub fn into_demoted(self) -> Option<(Cache, u32)> {
708 if self.pending_tok.is_some() {
709 return None;
710 }
711 let np = self.next_pred?;
712 debug_assert_eq!(
713 self.cache.pos,
714 self.committed.len(),
715 "demotion handoff: cache rows != committed tokens"
716 );
717 Some((self.cache, np))
718 }
719 /// Pool-resume hook (audit Q2): clear the parked draft-graph failure memoization so a
720 /// NEW request resuming this session gets one fresh capture chance — a transient-pressure
721 /// capture failure must not persist for the pool's whole lifetime (the TRT #16072 class).
722 /// Logs once iff a flag was actually set; a no-fallback resume is silent and free.
723 pub fn reset_graph_fallback_on_resume(&mut self) {
724 if let Some(line) = self
725 .draft_ctx
726 .as_mut()
727 .and_then(|c| c.failed.reset_on_resume())
728 {
729 eprintln!("{line}");
730 }
731 }
732}
733
734/// A session's PROMPT-END boundary state, the rewind target for session-affinity resume.
735///
736/// WHY THIS BOUNDARY, AND WHY IT IS THE ONLY ONE WORTH KEEPING. The rewrite class this lane
737/// exists for (a client that strips `<think>` blocks out of prior assistant turns) mutates the
738/// text the session GENERATED, never the prompt it was given. So turn N's prompt agrees with
739/// turn N-1's committed tokens up to almost exactly where turn N-1's generation began — the
740/// prompt-end boundary. Keeping a checkpoint there means the next turn re-primes only its own
741/// delta (the rewritten answer + the new user turn) instead of the whole conversation.
742///
743/// WHAT IT MUST HOLD. Full-attn KV is append-only and position-addressed, so rewinding it is a
744/// `len` truncation (no data). Linear-attn (GDN) conv/ssm state is mutated IN PLACE with no
745/// position index, so it must be a real device COPY — that copy is the entire reason a spec
746/// session could not previously rewind. The MTP draft scratch needs no copy either: its rows
747/// below the boundary were written by this turn's fill and are never revisited (the per-round
748/// true-hidden refresh only rewrites the CURRENT burst's committed positions), so rewinding it
749/// is also just a `len` reset. `last_h` is the hidden of the last row below the boundary — the
750/// predecessor-pairing anchor the next prime's fill reads for its first row.
751///
752/// COST: one `Cache::snapshot` per TURN, on a code path that already takes one per ROUND.
753pub(crate) struct SpecCheckpoint {
754 snap: crate::cache::CacheSnapshot,
755 /// Committed length at the boundary (== cache.pos there, the session invariant).
756 pos: usize,
757 /// Pre-output_norm hidden of row `pos - 1`.
758 last_h: CudaSlice<f32>,
759}
760
761/// PREFIX-CACHE BOUNDARY CAPTURE (lane/spec-prefix-cache, 2026-08-14): the state a spec session
762/// records at its cold-prime split so the WORKER can publish a cross-request prefix entry —
763/// the commit-gated-publication port (research/cache-spec-design-20260814/PORT-PLAN.md item 1).
764/// Only the pieces that are DESTROYED by continuing the prime need copies here: the in-place
765/// GDN conv/ssm states (via `Cache::snapshot`, same mechanism as [`SpecCheckpoint`]) and the
766/// boundary logits. Full-attn KV rows `[0..pos)` and draft-scratch rows `[0..pos)` are
767/// append-only for the session's lifetime (rollbacks never truncate below the prime boundary),
768/// so the worker slices those from the live caches post-burst instead of copying at prime time.
769pub struct SpecBoundaryCapture {
770 pub snap: crate::cache::CacheSnapshot,
771 /// Token boundary (== cache.pos at capture; == the worker's miss-LCP split).
772 pub pos: usize,
773 /// Full-vocab logits after the prefix prime — the entry's boundary logits.
774 pub logits: Vec<f32>,
775 /// Pre-output_norm trunk hidden of row `pos - 1` (lane/spec-on-cache-hit): the
776 /// predecessor-pairing anchor a RESTORED spec session's first suffix-fill row reads
777 /// (the `SpecSession::last_h` convention). Empty = unavailable (capture stays valid;
778 /// the fill's zeros row-0 fallback covers it at a bounded acceptance cost).
779 pub last_h: Vec<f32>,
780}
781
782/// D2H one hidden row out of a `[T, n_embd]` prime hidden stack — the boundary anchor a
783/// spec boundary capture carries for later restored-session fills. Failure is silent
784/// (`turn_ckpt` convention): the capture publishes without an anchor.
785fn capture_boundary_hidden(
786 e: &Engine,
787 h_rows: &CudaSlice<f32>,
788 pos: usize,
789 n_embd: usize,
790) -> Vec<f32> {
791 if pos == 0 || h_rows.len() < pos * n_embd {
792 return Vec::new();
793 }
794 let Ok(mut row) = e.uninit(n_embd) else {
795 return Vec::new();
796 };
797 if e.copy_view_into(
798 &mut row,
799 0,
800 &h_rows.slice((pos - 1) * n_embd..pos * n_embd),
801 n_embd,
802 )
803 .is_err()
804 {
805 return Vec::new();
806 }
807 e.dtoh(&row).unwrap_or_default()
808}
809
810/// ROLLBACK DOOR for sampled BOUNDARY tokens (lane/sampled-spec-quality, 2026-08-19).
811/// Default ON: the token a burst emits at its own boundary is drawn from the request's
812/// sampler. `MEMRA_SPEC_SAMPLED_BOUNDARY=0` restores the pre-lane posture (an ARGMAX at
813/// every boundary) without touching greedy, which is byte-unaffected either way.
814pub fn spec_sampled_boundary_on() -> bool {
815 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
816 *ON.get_or_init(|| std::env::var("MEMRA_SPEC_SAMPLED_BOUNDARY").as_deref() != Ok("0"))
817}
818
819/// ROLLBACK DOOR for SESSION-SPANNING penalty history (lane/sampled-spec-quality).
820/// Default ON: `pen_hist` is seeded from the session's committed tail, so repetition /
821/// frequency / presence penalties see the whole stream. `MEMRA_SPEC_PEN_SESSION=0`
822/// restores the pre-lane posture (each burst restarts the window from its own prompt
823/// slice, i.e. from NOTHING on a continuation burst) — and with the door shut the worker
824/// must keep refusing penalized sampled prefix-cache restores, because the restored
825/// session's continuation burst is handed no prompt slice at all.
826pub fn spec_pen_session_on() -> bool {
827 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
828 *ON.get_or_init(|| std::env::var("MEMRA_SPEC_PEN_SESSION").as_deref() != Ok("0"))
829}
830
831/// ROLLBACK DOOR for extended-entry publication from a RESTORED session
832/// (lane/sampled-spec-quality, Item 3). Default ON: a converted prefix-cache hit that fed a
833/// suffix captures its own prompt-end boundary so the NEXT turn can hit a longer prefix.
834/// `MEMRA_SPEC_RESTORE_REPUBLISH=0` restores the pre-lane posture (a namespace learns exactly
835/// one boundary and never advances it). Whole-entry semantics only — the boundary is the
836/// restored session's own prompt end, so `entry_pos != fed_len` still refuses on the way in.
837pub fn spec_restore_republish_on() -> bool {
838 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
839 *ON.get_or_init(|| std::env::var("MEMRA_SPEC_RESTORE_REPUBLISH").as_deref() != Ok("0"))
840}
841
842/// Diagnostics: name every boundary token on stderr (`MEMRA_SPEC_BOUNDARY_TRACE=1`), with
843/// the argmax the pre-lane code would have emitted from the same row. This is how the
844/// lane MEASURES the boundary rate and the deviation rate instead of estimating them.
845fn spec_boundary_trace() -> bool {
846 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
847 *ON.get_or_init(|| std::env::var("MEMRA_SPEC_BOUNDARY_TRACE").as_deref() == Ok("1"))
848}
849
850/// llama-parity floor for the penalty window when the request does not ask for a bigger
851/// one (`repeat_last_n` default). The serve API arms `penalty_last_n = usize::MAX` for any
852/// non-identity penalty, so this floor only matters to explicit small windows and to the
853/// CLI env path.
854const PEN_WINDOW_FLOOR: usize = 64;
855
856/// CEILING on the penalty window, and it is a COST bound, not a semantic preference.
857/// `penalize_logits_f32` (cu/spec_sample.cu) dedups on device by having thread `i` scan
858/// `hist[0..i]`, so a pass is O(n_hist²) and it runs ~3x per verify round (the q rows, the
859/// p column, the bonus column). The serve API arms `penalty_last_n = usize::MAX` for ANY
860/// non-identity penalty — "the whole context", llama's `repeat_last_n = -1` — so an
861/// uncapped session window would put a 128k-token history through that kernel: ~1.7e10
862/// comparisons per pass, tens of ms per round, i.e. penalties would silently destroy decode
863/// throughput on exactly the long-context requests that most want them. 8192 keeps a pass
864/// at ~7e7 comparisons (tens of microseconds) while still being **128x wider than the
865/// pre-lane effective window** (64 prompt-tail tokens + whatever the current burst had
866/// generated). A request that genuinely needs a window beyond this wants host-side dedup +
867/// counts through a new kernel signature — a follow-up lane, named here rather than hidden.
868const PEN_WINDOW_MAX: usize = 8192;
869
870/// Seed a penalty window over the SESSION, not the burst (lane/sampled-spec-quality,
871/// Item 2). The window is the last `max(penalty_last_n, 64)` tokens of
872/// `session_committed ++ burst_prompt` — for a cold turn-1 burst (`session_committed`
873/// empty, default `penalty_last_n`) that is byte-identically the pre-lane
874/// `prompt.iter().rev().take(64).rev()`; for a continuation burst it is the stream the
875/// client actually asked us to penalize, where the pre-lane code had NOTHING.
876fn pen_window_seed(
877 session_committed: &[u32],
878 burst_prompt: &[u32],
879 penalty_last_n: usize,
880) -> Vec<u32> {
881 let win = penalty_last_n.clamp(PEN_WINDOW_FLOOR, PEN_WINDOW_MAX);
882 let take_prompt = burst_prompt.len().min(win);
883 let take_sess = (win - take_prompt).min(session_committed.len());
884 let mut hist = Vec::with_capacity(take_sess + take_prompt);
885 hist.extend_from_slice(&session_committed[session_committed.len() - take_sess..]);
886 hist.extend_from_slice(&burst_prompt[burst_prompt.len() - take_prompt..]);
887 hist
888}
889
890/// Draw a BOUNDARY token from the target distribution the request asked for
891/// (lane/sampled-spec-quality, Item 1) — the fix for "sampled spec emits an ARGMAX token at
892/// every burst boundary".
893///
894/// WHY THIS EXISTS. A spec burst's first emitted token is not produced by the accept walk:
895/// it comes off a logits row that already exists (the prime's last row on a cold burst; the
896/// row after the last committed token on a continuation burst; the prefix-cache entry's
897/// boundary row on a restored one). Pre-lane that token was `argmax` in BOTH sampling
898/// regimes, so a sampled stream took a greedy token once per burst — measured, not
899/// estimated, in research/spec-cache-20260818/SAMPLED-QUALITY.md. At temperature > 0 the
900/// customer asked for a sampled token, so this draws one.
901///
902/// THE PROGRAM IS THE FULL-ACCEPT BONUS'S PROGRAM, deliberately: penalize the row (over the
903/// session's window), take this row's OWN filter stats (the sampfix-20260805 law — stats
904/// from a neighbour row mis-scale every `e0` and can wipe the row to token 0), gumbel-perturb
905/// with the session's Philox stream at `*sctr`, argmax the perturbed row. Reusing the bonus's
906/// composition means `sample_check`'s distributional oracle covers this draw too, and the
907/// boundary token is drawn from the same filtered/penalized `p` the accept walk targets.
908///
909/// THE STREAM IS THE SESSION'S, NOT A FRESH ONE. `sctr` is the caller's live counter and is
910/// advanced by exactly one, so a boundary draw consumes the next value in the same Philox
911/// stream the accept walk uses — never a second, independently seeded stream (which would be
912/// a new distributional bug: two streams from one seed correlate wherever their counters
913/// collide). That also makes a restored session's boundary draw at `sctr == 0` bit-identical
914/// to the cold session's own first draw from the same logits row, which is what preserves the
915/// sampled-hit lane's per-seed hit==cold byte identity.
916#[allow(clippy::too_many_arguments)]
917pub fn sample_boundary_token_dev(
918 e: &Engine,
919 logits: &CudaSlice<f32>,
920 n_vocab: usize,
921 sp: &SpecSampling,
922 pen_hist: &[u32],
923 sctr: &mut u32,
924 site: &str,
925) -> Result<u32, Box<dyn std::error::Error>> {
926 debug_assert!(
927 sp.temp > 0.0,
928 "boundary sampling is the sampled regime only"
929 );
930 // Own copy: penalize_logits mutates in place and the caller's row is live state
931 // (prime_logits back the constrained recompute; last_col_logits backs round 0's accept).
932 let mut col = e.zeros(n_vocab)?;
933 e.copy_into(&mut col, 0, logits, n_vocab)?;
934 let pen_on = sp.penalty_last_n > 0
935 && (sp.penalty_repeat != 1.0 || sp.penalty_freq != 0.0 || sp.penalty_present != 0.0);
936 if pen_on && !pen_hist.is_empty() {
937 // window trim mirrors the round loop's own upload (`pen_hist[w0..]`), cap included.
938 let w0 = pen_hist
939 .len()
940 .saturating_sub(sp.penalty_last_n.min(PEN_WINDOW_MAX));
941 let hist = &pen_hist[w0..];
942 let hd = e.htod_u32_v(hist)?;
943 e.penalize_logits(
944 &mut col,
945 &hd,
946 hist.len(),
947 sp.penalty_repeat,
948 sp.penalty_freq,
949 sp.penalty_present,
950 n_vocab,
951 )?;
952 }
953 let rows0 = e.htod_i32(&[0])?;
954 let (mut th_d, mut z_d, mut mx_d) = (e.zeros(1)?, e.zeros(1)?, e.zeros(1)?);
955 e.filter_stats(
956 &col, n_vocab, &rows0, &mut th_d, &mut z_d, &mut mx_d, n_vocab, 1, sp.temp, sp.top_k,
957 sp.top_p, sp.min_p,
958 )?;
959 let (th, mx) = (e.dtoh(&th_d)?[0], e.dtoh(&mx_d)?[0]);
960 let mut perturb = e.zeros(n_vocab)?;
961 e.gumbel_perturb_filtered(&col, &mut perturb, n_vocab, sp.seed, *sctr, sp.temp, mx, th)?;
962 *sctr = sctr.wrapping_add(1);
963 let td = e.argmax_token_device(&perturb, n_vocab)?;
964 let tok = e.dtoh_u32_one(&td)?;
965 if spec_boundary_trace() {
966 // the pre-lane token, from the SAME row, so the deviation rate is measurable.
967 let raw = e.argmax_token_device(logits, n_vocab)?;
968 let greedy = e.dtoh_u32_one(&raw)?;
969 eprintln!(
970 "[spec-boundary] site={site} sampled={tok} argmax={greedy} \
971 deviates={} temp={} sctr={}",
972 (tok != greedy) as u8,
973 sp.temp,
974 sctr.wrapping_sub(1),
975 );
976 }
977 Ok(tok)
978}
979
980/// Host-row twin of [`sample_boundary_token_dev`] (the prime / feed / entry rows arrive as
981/// host `Vec<f32>`).
982#[allow(clippy::too_many_arguments)]
983pub fn sample_boundary_token(
984 e: &Engine,
985 logits: &[f32],
986 sp: &SpecSampling,
987 pen_hist: &[u32],
988 sctr: &mut u32,
989 site: &str,
990) -> Result<u32, Box<dyn std::error::Error>> {
991 let n_vocab = logits.len();
992 let d = e.htod(logits)?;
993 sample_boundary_token_dev(e, &d, n_vocab, sp, pen_hist, sctr, site)
994}
995
996struct SpecPipeTraceClock {
997 pair: usize,
998 started: std::time::Instant,
999}
1000
1001#[derive(Clone)]
1002struct SpecPipeTraceCtx {
1003 clock: std::sync::Arc<SpecPipeTraceClock>,
1004 round: usize,
1005 lane: usize,
1006}
1007
1008struct SpecPipeTraceMarker {
1009 trace: SpecPipeTraceCtx,
1010 phase: &'static str,
1011 edge: &'static str,
1012 slot: Option<usize>,
1013}
1014
1015unsafe extern "C" fn spec_pipe_trace_marker(raw: *mut std::ffi::c_void) {
1016 let marker = unsafe { Box::from_raw(raw.cast::<SpecPipeTraceMarker>()) };
1017 let lane = if marker.trace.lane == 0 { "A" } else { "B" };
1018 let slot = marker
1019 .slot
1020 .map(|v| v.to_string())
1021 .unwrap_or_else(|| "-".into());
1022 let t_ms = marker.trace.clock.started.elapsed().as_secs_f64() * 1e3;
1023 use std::io::Write as _;
1024 let stderr = std::io::stderr();
1025 let mut stderr = stderr.lock();
1026 let _ = writeln!(
1027 stderr,
1028 "[spec-pipe-timeline] pair={} round={} lane={lane} phase={} edge={} \
1029 slot={slot} t_ms={t_ms:.3}",
1030 marker.trace.clock.pair, marker.trace.round, marker.phase, marker.edge,
1031 );
1032}
1033
1034fn enqueue_spec_pipe_trace_marker(
1035 stream: &cudarc::driver::CudaStream,
1036 trace: Option<&SpecPipeTraceCtx>,
1037 phase: &'static str,
1038 edge: &'static str,
1039 slot: Option<usize>,
1040) -> Result<(), Box<dyn std::error::Error>> {
1041 let Some(trace) = trace else {
1042 return Ok(());
1043 };
1044 let marker = Box::new(SpecPipeTraceMarker {
1045 trace: trace.clone(),
1046 phase,
1047 edge,
1048 slot,
1049 });
1050 let raw = Box::into_raw(marker);
1051 let result = unsafe {
1052 cudarc::driver::result::stream::launch_host_function(
1053 stream.cu_stream(),
1054 spec_pipe_trace_marker,
1055 raw.cast(),
1056 )
1057 };
1058 if let Err(err) = result {
1059 unsafe {
1060 drop(Box::from_raw(raw));
1061 }
1062 return Err(err.into());
1063 }
1064 Ok(())
1065}
1066
1067#[derive(Default)]
1068struct SpecPipeProgress {
1069 setup_done: [bool; 2],
1070 draft_done: [usize; 2],
1071 stage0_done: [usize; 2],
1072 verify_done: [usize; 2],
1073 accept_done: [usize; 2],
1074 finished: [bool; 2],
1075 aborted: bool,
1076}
1077
1078/// Host-side issue coordinator for the reduced two-session speculative pipeline. Each session
1079/// keeps its existing call stack and round locals; this object only orders phase entry. The
1080/// primary mutex spans whole draft/accept/tail issue regions so Engine's single-stream scratch
1081/// cannot be interleaved by the two host threads.
1082struct SpecPipeSync {
1083 progress: std::sync::Mutex<SpecPipeProgress>,
1084 changed: std::sync::Condvar,
1085 primary: std::sync::Mutex<()>,
1086 trace: Option<std::sync::Arc<SpecPipeTraceClock>>,
1087}
1088
1089impl SpecPipeSync {
1090 fn new() -> Self {
1091 static TRACE_PAIR: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
1092 let trace = (std::env::var("MEMRA_SPEC_PIPE_TRACE").as_deref() == Ok("1")).then(|| {
1093 std::sync::Arc::new(SpecPipeTraceClock {
1094 pair: TRACE_PAIR.fetch_add(1, std::sync::atomic::Ordering::Relaxed) + 1,
1095 started: std::time::Instant::now(),
1096 })
1097 });
1098 Self {
1099 progress: std::sync::Mutex::new(SpecPipeProgress::default()),
1100 changed: std::sync::Condvar::new(),
1101 primary: std::sync::Mutex::new(()),
1102 trace,
1103 }
1104 }
1105}
1106
1107#[derive(Clone)]
1108struct SpecPipeLane {
1109 sync: std::sync::Arc<SpecPipeSync>,
1110 lane: usize,
1111}
1112
1113impl SpecPipeLane {
1114 fn peer(&self) -> usize {
1115 1 - self.lane
1116 }
1117
1118 fn aborted() -> Box<dyn std::error::Error> {
1119 "paired speculative peer aborted".into()
1120 }
1121
1122 fn trace(&self, round: usize) -> Option<SpecPipeTraceCtx> {
1123 self.sync.trace.as_ref().map(|clock| SpecPipeTraceCtx {
1124 clock: clock.clone(),
1125 round,
1126 lane: self.lane,
1127 })
1128 }
1129
1130 fn setup_begin(&self) -> Result<(), Box<dyn std::error::Error>> {
1131 let mut p = self.sync.progress.lock().unwrap();
1132 while !p.aborted && self.lane == 1 && !p.setup_done[0] && !p.finished[0] {
1133 p = self.sync.changed.wait(p).unwrap();
1134 }
1135 if p.aborted {
1136 Err(Self::aborted())
1137 } else {
1138 Ok(())
1139 }
1140 }
1141
1142 fn setup_end(&self) {
1143 let mut p = self.sync.progress.lock().unwrap();
1144 p.setup_done[self.lane] = true;
1145 self.sync.changed.notify_all();
1146 }
1147
1148 fn draft_begin(
1149 &self,
1150 round: usize,
1151 ) -> Result<std::sync::MutexGuard<'_, ()>, Box<dyn std::error::Error>> {
1152 let peer = self.peer();
1153 let mut p = self.sync.progress.lock().unwrap();
1154 loop {
1155 if p.aborted {
1156 return Err(Self::aborted());
1157 }
1158 let setup_ready =
1159 (p.setup_done[0] || p.finished[0]) && (p.setup_done[1] || p.finished[1]);
1160 let prior_ready = p.accept_done[self.lane] >= round
1161 && (p.accept_done[peer] >= round || p.finished[peer]);
1162 let turn_ready = if self.lane == 0 {
1163 true
1164 } else {
1165 p.draft_done[0] > round || p.finished[0]
1166 };
1167 if setup_ready && prior_ready && turn_ready {
1168 break;
1169 }
1170 p = self.sync.changed.wait(p).unwrap();
1171 }
1172 drop(p);
1173 Ok(self.sync.primary.lock().unwrap())
1174 }
1175
1176 fn draft_end(&self, round: usize) {
1177 let mut p = self.sync.progress.lock().unwrap();
1178 p.draft_done[self.lane] = round + 1;
1179 self.sync.changed.notify_all();
1180 }
1181
1182 /// Admit stage 0 and return whether this lane owns the interval's one reverse fence.
1183 /// Lane B releases as soon as lane A has issued its boundary TX, not after A's full body.
1184 fn stage0_begin(&self, round: usize) -> Result<bool, Box<dyn std::error::Error>> {
1185 let peer = self.peer();
1186 let mut p = self.sync.progress.lock().unwrap();
1187 loop {
1188 if p.aborted {
1189 return Err(Self::aborted());
1190 }
1191 let ready = if self.lane == 0 {
1192 p.draft_done[0] > round && (p.draft_done[1] > round || p.finished[1])
1193 } else {
1194 p.draft_done[1] > round && (p.stage0_done[0] > round || p.finished[0])
1195 };
1196 if ready {
1197 return Ok(self.lane == 0 || p.finished[peer]);
1198 }
1199 p = self.sync.changed.wait(p).unwrap();
1200 }
1201 }
1202
1203 fn stage0_end(&self, round: usize) {
1204 let mut p = self.sync.progress.lock().unwrap();
1205 p.stage0_done[self.lane] = round + 1;
1206 self.sync.changed.notify_all();
1207 }
1208
1209 /// Stage 1 is single-owner per engine. A proceeds immediately after its own ticket; B waits
1210 /// for A's full stage1/head issue so only A.S1 and B.S0 can overlap.
1211 fn stage1_begin(&self, round: usize) -> Result<(), Box<dyn std::error::Error>> {
1212 let mut p = self.sync.progress.lock().unwrap();
1213 while !p.aborted
1214 && !(p.stage0_done[self.lane] > round
1215 && (self.lane == 0 || p.verify_done[0] > round || p.finished[0]))
1216 {
1217 p = self.sync.changed.wait(p).unwrap();
1218 }
1219 if p.aborted {
1220 Err(Self::aborted())
1221 } else {
1222 Ok(())
1223 }
1224 }
1225
1226 fn verify_end(&self, round: usize) {
1227 let mut p = self.sync.progress.lock().unwrap();
1228 p.verify_done[self.lane] = round + 1;
1229 self.sync.changed.notify_all();
1230 }
1231
1232 fn accept_begin(
1233 &self,
1234 round: usize,
1235 ) -> Result<std::sync::MutexGuard<'_, ()>, Box<dyn std::error::Error>> {
1236 let mut p = self.sync.progress.lock().unwrap();
1237 loop {
1238 if p.aborted {
1239 return Err(Self::aborted());
1240 }
1241 let ready = if self.lane == 0 {
1242 p.verify_done[0] > round && (p.verify_done[1] > round || p.finished[1])
1243 } else {
1244 p.verify_done[1] > round && (p.accept_done[0] > round || p.finished[0])
1245 };
1246 if ready {
1247 break;
1248 }
1249 p = self.sync.changed.wait(p).unwrap();
1250 }
1251 drop(p);
1252 Ok(self.sync.primary.lock().unwrap())
1253 }
1254
1255 fn accept_end(&self, round: usize) {
1256 let mut p = self.sync.progress.lock().unwrap();
1257 p.accept_done[self.lane] = round + 1;
1258 self.sync.changed.notify_all();
1259 }
1260
1261 fn primary(&self) -> std::sync::MutexGuard<'_, ()> {
1262 self.sync.primary.lock().unwrap()
1263 }
1264
1265 fn finish(&self, failed: bool) {
1266 let mut p = self.sync.progress.lock().unwrap();
1267 p.finished[self.lane] = true;
1268 p.aborted |= failed;
1269 self.sync.changed.notify_all();
1270 }
1271}
1272
1273struct SpecPipeFinish<'a> {
1274 lane: &'a SpecPipeLane,
1275 closed: bool,
1276}
1277
1278impl<'a> SpecPipeFinish<'a> {
1279 fn new(lane: &'a SpecPipeLane) -> Self {
1280 Self {
1281 lane,
1282 closed: false,
1283 }
1284 }
1285
1286 fn close(&mut self, failed: bool) {
1287 self.lane.finish(failed);
1288 self.closed = true;
1289 }
1290}
1291
1292impl Drop for SpecPipeFinish<'_> {
1293 fn drop(&mut self) {
1294 if !self.closed {
1295 self.lane.finish(true);
1296 }
1297 }
1298}
1299
1300/// Scoped transfer of one exclusively-borrowed session to the second host issue thread.
1301/// `CudaGraph` is not marked Send by cudarc because its raw driver handles carry no automatic
1302/// trait. CUDA driver graph handles are context-scoped rather than OS-thread-affine; the caller
1303/// binds that context before touching the session, joins before returning, and never aliases the
1304/// pointer. Keep this exception local to the experimental pair call instead of marking the public
1305/// session type Send.
1306struct SpecPipeSessionPtr(*mut SpecSession);
1307
1308unsafe impl Send for SpecPipeSessionPtr {}
1309
1310impl SpecPipeSessionPtr {
1311 unsafe fn get_mut(&mut self) -> &mut SpecSession {
1312 unsafe { &mut *self.0 }
1313 }
1314}
1315
1316/// Per-session persistent draft-graph context: the captured CUDA graph(s) plus the device
1317/// buffers whose POINTERS the capture bakes. Reuse legality: the greedy capture bakes only
1318/// session-stable pointers (the session's own MtpScratch KV — allocated once, never realloc'd;
1319/// the model's resident embedding; the process-wide OnceLock p_min) and the g_* buffers held
1320/// HERE — so one capture serves the session's whole lifetime. The sampled capture additionally
1321/// bakes (seed, temp) as capture-time constants and needs k q-slots — keyed by `s_key`, dropped
1322/// and recaptured when a pool-resumed request changes them. `*_failed` memoizes a failed capture
1323/// so the eager fallback doesn't pay a doomed capture attempt every burst.
1324/// Capture identity of the parked SAMPLED draft graph (`DraftGraphCtx::graph_s`).
1325///
1326/// EXACTNESS, not perf (lane/graph-s-key-exactness-20260819; receipts
1327/// `research/spec-cache-20260818/GRAPH-S-KEY.md`). Two classes of field live here, both
1328/// load-bearing:
1329///
1330/// - **Baked constants.** `seed` and `temp` are capture-time constants INSIDE the graph and `k`
1331/// sizes the q slots its replays write. A resumed request changing any of them must recapture.
1332/// This is all the key used to carry.
1333/// - **Regime fields.** `top_k`/`top_p`/`min_p`/`pen_on` are not baked, but they decide whether
1334/// the captured graph is a legal draft chain AT ALL. The in-graph draw is one gumbel-max over
1335/// the RAW softmax (`gumbel_perturb_ctr`, unfiltered by construction), while the verify builds
1336/// the accept test's `q` from `filter_stats(q_slots, top_k, top_p, min_p)`. If those disagree
1337/// the accept test evaluates a distribution the draft was never sampled from: a draft token
1338/// below the filter threshold gathers `q = 0` (`softmax_gather_filtered_f32`,
1339/// `cu/spec_sample.cu`) and `u * 0 < p` accepts it UNCONDITIONALLY.
1340///
1341/// Omitting the regime fields was reachable — not through the prefix-cache spec restore (that
1342/// path is greedy-only, `memra-server` `spec_restore_convertible`), but through WHOLE-SESSION
1343/// spec reuse: a parked `SpecSession` carries this `DraftGraphCtx`, and the pool-resume probe
1344/// applies no sampler predicate at all. Turn 1 pure-temp parks a graph; turn 2 of the same
1345/// conversation, same explicit seed and temperature, adds `top_p`/`top_k` and inherits it.
1346#[derive(Clone, Copy, PartialEq, Eq, Debug)]
1347pub(crate) struct SampledGraphKey {
1348 seed: u64,
1349 temp_bits: u32,
1350 k: usize,
1351 top_k: i32,
1352 top_p_bits: u32,
1353 min_p_bits: u32,
1354 pen_on: bool,
1355}
1356
1357impl SampledGraphKey {
1358 pub(crate) fn new(
1359 seed: u64,
1360 temp: f32,
1361 k: usize,
1362 top_k: i32,
1363 top_p: f32,
1364 min_p: f32,
1365 pen_on: bool,
1366 ) -> Self {
1367 SampledGraphKey {
1368 seed,
1369 temp_bits: temp.to_bits(),
1370 k,
1371 top_k,
1372 top_p_bits: top_p.to_bits(),
1373 min_p_bits: min_p.to_bits(),
1374 pen_on,
1375 }
1376 }
1377
1378 /// The one regime the in-graph sampled chain may stand in for the eager one: nothing but
1379 /// temperature shapes `q`. Computed FROM THE KEY so the capture guard, the launch guard and
1380 /// the key can never drift apart (they were three separate expressions before this lane, and
1381 /// the launch site simply forgot to ask).
1382 pub(crate) fn pure_temp(&self) -> bool {
1383 self.top_k == 0
1384 && f32::from_bits(self.top_p_bits) >= 1.0
1385 && f32::from_bits(self.min_p_bits) <= 0.0
1386 && !self.pen_on
1387 }
1388}
1389
1390pub(crate) struct DraftGraphCtx {
1391 g_tok: CudaSlice<u32>,
1392 g_pos: CudaSlice<i32>,
1393 g_seed: CudaSlice<f32>,
1394 g_p: CudaSlice<f32>,
1395 g_ctr: CudaSlice<u32>,
1396 g_q: CudaSlice<f32>,
1397 g_perturb: CudaSlice<f32>,
1398 q_slots: Vec<CudaSlice<f32>>,
1399 /// DRAFT-SIDE GRAMMAR MASK (lane/draft-mask): packed allowed-set words over the DRAFT
1400 /// head's vocab, at a STABLE address so the captured draft graph's mask node reads the
1401 /// per-position contents the host re-uploads before each replay (the graph-promote
1402 /// pattern from decode.rs). Empty unless the session drafts under a grammar.
1403 g_dmask: CudaSlice<u32>,
1404 /// was `graph` captured WITH the mask node? A parked graph of the wrong shape is dropped.
1405 graph_masked: bool,
1406 graph: Option<cudarc::driver::CudaGraph>,
1407 graph_s: Option<cudarc::driver::CudaGraph>,
1408 /// Failed-capture memoization for both graphs — LOUD on flip, cleared on pool resume
1409 /// (audit Q2, the TRT #16072 silent-permanent-coverage-loss class).
1410 failed: DraftGraphFallback,
1411 /// Capture identity of `graph_s` — see [`SampledGraphKey`]. `None` iff no sampled graph is
1412 /// parked; a request whose key differs drops the parked graph (and its q slots/keeper).
1413 s_key: Option<SampledGraphKey>,
1414 /// CAPTURE-RETAIN keepers (#68 root cause, 2026-08-04): the warmup-run transients whose
1415 /// pool addresses the captured graph(s) bake. Without these, the transients return to the
1416 /// pool at capture-body exit and later work (burst-boundary prime/fill/commit passes, or a
1417 /// co-served session in the worker) reuses those addresses — the persisted graph's replay
1418 /// then reads/writes live unrelated buffers (exactness corruption, first seen as the ST
1419 /// serve-spec 4B graph-arm corruption; one-shot CLI calls never re-shuffled the pool, which
1420 /// is why run-spec K=1..8 passed on the same checkpoint). Same fix class as
1421 /// capture_graph_retained's gemma/decode.rs sites — hold as long as the graph replays.
1422 keeper: Vec<Box<dyn std::any::Any + Send>>,
1423 keeper_s: Vec<Box<dyn std::any::Any + Send>>,
1424}
1425
1426/// Failed-capture memoization for the two draft graphs (audit Q2, 2026-08-05 — the
1427/// TRT #16072 trap class: pressure-triggered, silent, long-lived coverage loss).
1428///
1429/// Three contracts:
1430/// - LOUD FLIP: `mark_*` returns the warn line exactly on the false→true transition
1431/// (returned, not printed, so the once-per-flip contract is unit-testable); the caller
1432/// `eprintln!`s it UNCONDITIONALLY — a dropped draft graph is never silent. Re-marking
1433/// an already-failed graph returns None (the per-burst memoization that keeps the eager
1434/// fallback from paying a doomed capture attempt every burst).
1435/// - RESET ON RESUME: `reset_on_resume` clears both flags — a parked session resumed by a
1436/// NEW request gets one fresh capture chance instead of carrying a transient-pressure
1437/// failure for the pool's whole lifetime. Returns the note line only when a flag was
1438/// actually set (quiet on the common clean-resume path).
1439/// - Shape-change clears (`clear_*`) stay silent, exactly as before: they precede a fresh
1440/// capture attempt whose own failure would re-flip loudly.
1441#[derive(Default)]
1442pub(crate) struct DraftGraphFallback {
1443 greedy: bool,
1444 sampled: bool,
1445}
1446impl DraftGraphFallback {
1447 fn mark_greedy(&mut self, reason: &str) -> Option<String> {
1448 if self.greedy {
1449 return None;
1450 }
1451 self.greedy = true;
1452 Some(format!(
1453 "[spec] WARN: draft-graph capture failed ({reason}); eager fallback until session resume"
1454 ))
1455 }
1456 fn mark_sampled(&mut self, reason: &str) -> Option<String> {
1457 if self.sampled {
1458 return None;
1459 }
1460 self.sampled = true;
1461 Some(format!(
1462 "[spec] WARN: sampled draft-graph capture failed ({reason}); eager fallback until session resume"
1463 ))
1464 }
1465 fn greedy_failed(&self) -> bool {
1466 self.greedy
1467 }
1468 fn sampled_failed(&self) -> bool {
1469 self.sampled
1470 }
1471 fn clear_greedy(&mut self) {
1472 self.greedy = false;
1473 }
1474 fn clear_sampled(&mut self) {
1475 self.sampled = false;
1476 }
1477 /// Pool-resume reset: both graphs get a fresh capture chance. Some(note) iff any flag
1478 /// was set (so clean resumes stay quiet).
1479 pub(crate) fn reset_on_resume(&mut self) -> Option<String> {
1480 if !self.greedy && !self.sampled {
1481 return None;
1482 }
1483 let which = match (self.greedy, self.sampled) {
1484 (true, true) => "greedy+sampled",
1485 (true, false) => "greedy",
1486 _ => "sampled",
1487 };
1488 self.greedy = false;
1489 self.sampled = false;
1490 Some(format!(
1491 "[spec] draft-graph fallback reset on session resume ({which}); recapture eligible"
1492 ))
1493 }
1494}
1495
1496impl DraftGraphCtx {
1497 fn new(e: &Engine, n_embd: usize, qlen: usize) -> Result<Self, Box<dyn std::error::Error>> {
1498 Ok(DraftGraphCtx {
1499 g_tok: e.alloc_u32_zeroed(1)?,
1500 g_pos: e.htod_i32(&[0])?,
1501 g_seed: e.zeros(n_embd)?,
1502 g_p: e.zeros(1)?,
1503 g_ctr: e.alloc_u32_zeroed(1)?,
1504 g_q: e.zeros(qlen)?,
1505 g_perturb: e.zeros(qlen)?,
1506 q_slots: Vec::new(),
1507 g_dmask: e.alloc_u32_zeroed(1)?,
1508 graph_masked: false,
1509 graph: None,
1510 graph_s: None,
1511 failed: DraftGraphFallback::default(),
1512 s_key: None,
1513 keeper: Vec::new(),
1514 keeper_s: Vec::new(),
1515 })
1516 }
1517}
1518
1519pub(crate) struct MtpScratch {
1520 kv: KvLayer,
1521 /// Logical row capacity. On the graph/DC draft path it also doubles as fa_decode_dc's
1522 /// bucket_max: n_splits is sized from it ONCE, so the graph captured at round 0 stays valid
1523 /// for every later t_kv. Step35 refuses that path and may back this logical extent with the
1524 /// smaller host-indexed SWA ring instead.
1525 cap: usize,
1526}
1527
1528fn mtp_scratch_layout(
1529 cfg: &memra_gguf::config::ModelConfig,
1530 geom: Option<&crate::hybrid::DraftGeom>,
1531) -> (usize, usize, usize, usize) {
1532 // Student draft heads carry fewer KV heads (head_dim unchanged) -> smaller scratch rows.
1533 let n_head_kv = geom.map(|g| g.n_head_kv).unwrap_or(cfg.n_head_kv as usize);
1534 let head_dim_k = cfg.head_dim_k as usize;
1535 let head_dim_v = cfg.head_dim_v as usize;
1536 assert!(
1537 head_dim_k % 32 == 0 && head_dim_v % 32 == 0,
1538 "KVQUANT requires head_dim%32==0 (MTP scratch)"
1539 );
1540 let kv_dim_k = head_dim_k * n_head_kv;
1541 let kv_dim_v = head_dim_v * n_head_kv;
1542 // The fp8-KV arm deliberately does not reach the draft scratch; keep the exact format
1543 // policy shared with `MtpScratch::new` so admission scales the same allocation.
1544 let (kbb, vbb) = crate::kv_blk_bytes();
1545 let k_tok_bytes = (kv_dim_k / 32) * kbb;
1546 let v_tok_bytes = (kv_dim_v / 32) * vbb;
1547 (kv_dim_k, kv_dim_v, k_tok_bytes, v_tok_bytes)
1548}
1549
1550impl MtpScratch {
1551 fn new(
1552 e: &Engine,
1553 cfg: &memra_gguf::config::ModelConfig,
1554 cap: usize,
1555 geom: Option<&crate::hybrid::DraftGeom>,
1556 ) -> Result<Self, Box<dyn std::error::Error>> {
1557 // env-selected KV formats (default 34/24). The fp8-KV arm (MEMRA_KV_FP8) deliberately
1558 // does NOT reach the draft scratch: fp8 drafts drifted acceptance 69-88% -> 46%
1559 // (2026-07-12 A/B); the scratch is tiny, so it keeps baseline q8_0/q5_1 numerics
1560 // while the TRUNK cache carries the fp8 depth win. Scratch append/fa pass g=false.
1561 let (kv_dim_k, kv_dim_v, k_tok_bytes, v_tok_bytes) = mtp_scratch_layout(cfg, geom);
1562 let ring = if crate::cache::swa_ring_on() && cfg.arch.is_step35() {
1563 let window = cfg.step35.as_ref().unwrap().sliding_window as usize;
1564 Some(crate::cache::KvRing::new(
1565 crate::cache::swa_ring_rows(window, cap),
1566 window,
1567 ))
1568 } else {
1569 None
1570 };
1571 let alloc_rows = ring.as_ref().map(crate::cache::KvRing::rows).unwrap_or(cap);
1572 Ok(MtpScratch {
1573 kv: KvLayer {
1574 k: e.alloc_u8(alloc_rows * k_tok_bytes)?,
1575 v: e.alloc_u8(alloc_rows * v_tok_bytes)?,
1576 kv_dim_k,
1577 kv_dim_v,
1578 k_tok_bytes,
1579 v_tok_bytes,
1580 len: 0,
1581 ring,
1582 len_d: e.htod_i32(&[0])?,
1583 },
1584 cap,
1585 })
1586 }
1587 /// Set BOTH length counters: the host mirror AND the device len_d the captured append/fa read
1588 /// (a 4-byte in-place htod — the counter pointer is baked into the graph, never realloc'd).
1589 /// This is the ONLY truncation/rollback mechanism the persistent draft KV needs.
1590 fn set_len(&mut self, e: &Engine, n: usize) -> Result<(), Box<dyn std::error::Error>> {
1591 if self
1592 .kv
1593 .ring
1594 .as_ref()
1595 .is_some_and(|ring| !ring.can_rewind_to(n))
1596 {
1597 return Err("SWA ring MTP checkpoint has been lapped; full re-prime required".into());
1598 }
1599 self.kv.len = n;
1600 e.set_i32_one(&mut self.kv.len_d, n as i32)
1601 }
1602
1603 fn can_rewind_to(&self, n: usize) -> bool {
1604 self.kv
1605 .ring
1606 .as_ref()
1607 .is_none_or(|ring| ring.can_rewind_to(n))
1608 }
1609}
1610
1611/// Retained verify intermediates for the REPLAY-FREE partial accept (2026-07-03, the profiled
1612/// #1 spec cost at long ctx: the partial-accept replay was a DUPLICATE trunk pass — ~0.54 extra
1613/// full weight reads per round — recomputing columns the verify had already produced
1614/// bit-identically). Holds, per linear layer, everything needed to rebuild its recurrent state
1615/// to "after the first j verify columns" WITHOUT re-running the trunk:
1616/// - BATCHED-path layers (`gdn`): the exact token-major inputs the round's ONE gdn_scan
1617/// consumed. A prefix re-run of the SAME kernel (t=j) from the snapshot state is bit-identical
1618/// to the first j iterations of the verify's scan — the kernel's t-loop carries state in
1619/// registers and iteration t never depends on T. `qkv_mixed` (the conv input) feeds the
1620/// pure-copy ring rebuild.
1621/// - PER-COLUMN-path layers (`cols`): dtod clones of (conv_state, ssm_state) taken after each
1622/// column 0..t-2 — pure copies of the actual chain states (the last column is never a rebuild
1623/// target: j <= t-1).
1624/// Full-attn layers need nothing: their verify KV rows are bit-identical to eager's (the
1625/// decode-exact contract; verify-probe pins it), so rollback = len truncation.
1626struct GdnStash {
1627 qkv_mixed: CudaSlice<f32>, // [t, conv_dim] token-major (conv input)
1628 q_l2: CudaSlice<f32>,
1629 k_l2: CudaSlice<f32>,
1630 v_g: CudaSlice<f32>, // [t, num_v, d_state]
1631 g_log: CudaSlice<f32>,
1632 beta: CudaSlice<f32>, // [t, num_v]
1633}
1634struct VerifyCkpt {
1635 gdn: Vec<Option<GdnStash>>, // [n_layer], Some iff batched linear path ran
1636 cols: Vec<Option<Vec<(CudaSlice<f32>, CudaSlice<f32>)>>>, // [n_layer][col] = (conv, ssm) after col
1637}
1638/// Opaque handle for the dspark round (dflash.rs) — VerifyCkpt stays spec-private.
1639pub(crate) struct DsparkVerifyCkpt(VerifyCkpt);
1640
1641/// Engine-bundle slice 3 (DSF-ROUNDCOST-20260820 §2 row 4 / §5 rank 1): bucketed CUDA
1642/// graphs for the dspark verify's LINEAR-layer segments. The measured verify is ~2,800
1643/// eager launches whose residual cost is DEVICE-side per-launch overhead (slice 2 proved
1644/// host dispatch is not the binder: fully-deferred dispatch bought ~0 wall). The 48 GDN
1645/// layers between full-attention layers are shape-static given vt — no positions, no
1646/// t_kv, state addressed through pointer tables — so runs of them capture per
1647/// (segment, vt) and replay as ONE graph launch each. Full-attention layers stay eager
1648/// (their per-row append/fa arm picks are t_kv-driven — the exec-update extension).
1649///
1650/// Per round out-of-graph: one pointer-table refresh (gdn ping-pong moves the canonical
1651/// handles), one input-staging copy per segment, host parity bookkeeping. Captured via
1652/// `capture_graph_retained` (2 warmups + capture, keeper retains warmup transients so
1653/// pool addresses stay stable); the warmups EXECUTE, so segment conv/ssm state is saved
1654/// before and restored after — the graph's first real launch starts from the exact
1655/// pre-round state. The ckpt column stash rides persistent slabs (written inside the
1656/// graph as memcpy nodes); commit reads them via `dspark_commit_prefix_slab`.
1657/// `MEMRA_DSPARK_VERIFY_GRAPH=0` reverts to the eager walk (byte-identical body).
1658pub(crate) struct DsparkVerifyGraphs {
1659 /// Linear-attention layer indices ascending; `lin_pos[il]` = index into the vecs.
1660 lin: Vec<usize>,
1661 lin_pos: std::collections::HashMap<usize, usize>,
1662 /// [n_lin x 6] pointer table (conv, s0, s1, conv, s1, s0 per layer), refreshed per
1663 /// verify from the live handles; layer il's slice starts at lin_pos[il]*6.
1664 table_all: CudaSlice<u64>,
1665 host_table: Vec<u64>,
1666 /// Persistent per-layer ckpt stash slabs: row r of the verify at slab offset
1667 /// r*words. Shared by every (segment, vt) bucket — one verify runs at a time.
1668 stash_conv: Vec<CudaSlice<f32>>,
1669 stash_ssm: Vec<CudaSlice<f32>>,
1670 conv_words: usize,
1671 ssm_words: usize,
1672 /// Per-vt input/output staging (stable addresses the graphs bake).
1673 stage: std::collections::HashMap<usize, (CudaSlice<f32>, CudaSlice<f32>)>,
1674 /// Per-vt dflash tap-sink buffers — the captured segments bake the tap dst address,
1675 /// so the sink buffer must live (and persist) with the graphs, not with the round.
1676 pub(crate) tap_bufs: std::collections::HashMap<usize, CudaSlice<f32>>,
1677 graphs: std::collections::HashMap<(usize, usize), DsparkSegGraph>,
1678 /// Warmup-corruption guard scratch: pre-capture conv/ssm of one segment.
1679 save_conv: CudaSlice<f32>,
1680 save_ssm: CudaSlice<f32>,
1681 max_run: usize,
1682 n_embd: usize,
1683 /// Set by the verify walk: this round's linear ckpt lives in the slabs (the caller
1684 /// commits through `dspark_commit_prefix_slab` instead of the cols arm).
1685 pub(crate) round_slab: bool,
1686}
1687
1688struct DsparkSegGraph {
1689 graph: cudarc::driver::CudaGraph,
1690 _keeper: Vec<Box<dyn std::any::Any + Send>>,
1691}
1692
1693// SAFETY: `CudaGraph` is not marked Send by cudarc because its raw driver handles carry
1694// no automatic trait; CUDA driver graph handles are context-scoped rather than
1695// OS-thread-affine (the SpecPipeSessionPtr precedent above). The ctx lives in
1696// `HybridModel::dspark_vgraphs` behind a Mutex and every touch happens on the engine's
1697// single decode-stream thread.
1698unsafe impl Send for DsparkVerifyGraphs {}
1699
1700impl DsparkVerifyGraphs {
1701 /// Build for this cache's shape. None when there are no linear layers, sizes are
1702 /// non-uniform, or the trunk keeps a gemma4 config (never on the qwen35 family).
1703 pub(crate) fn new(
1704 e: &Engine,
1705 cache: &Cache,
1706 t_max: usize,
1707 n_embd: usize,
1708 ) -> Result<Option<Self>, Box<dyn std::error::Error>> {
1709 let lin: Vec<usize> = (0..cache.recur.len())
1710 .filter(|&il| cache.recur[il].is_some())
1711 .collect();
1712 if lin.is_empty() || t_max < 2 {
1713 return Ok(None);
1714 }
1715 let first = cache.recur[lin[0]].as_ref().unwrap();
1716 let (conv_words, ssm_words) = (first.conv_state.len(), first.ssm_state.len());
1717 for &il in &lin {
1718 let rl = cache.recur[il].as_ref().unwrap();
1719 if rl.conv_state.len() != conv_words || rl.ssm_state.len() != ssm_words {
1720 return Ok(None);
1721 }
1722 }
1723 let n = lin.len();
1724 let mut lin_pos = std::collections::HashMap::with_capacity(n);
1725 for (k, &il) in lin.iter().enumerate() {
1726 lin_pos.insert(il, k);
1727 }
1728 // longest run of consecutive linear layers (save-scratch sizing)
1729 let mut max_run = 1usize;
1730 let mut run = 1usize;
1731 for w in lin.windows(2) {
1732 if w[1] == w[0] + 1 {
1733 run += 1;
1734 max_run = max_run.max(run);
1735 } else {
1736 run = 1;
1737 }
1738 }
1739 let rows = t_max - 1;
1740 let mut stash_conv = Vec::with_capacity(n);
1741 let mut stash_ssm = Vec::with_capacity(n);
1742 for _ in 0..n {
1743 stash_conv.push(e.uninit(rows * conv_words)?);
1744 stash_ssm.push(e.uninit(rows * ssm_words)?);
1745 }
1746 let host_table = vec![0u64; n * 6];
1747 let table_all = e.htod_u64(&host_table)?;
1748 Ok(Some(Self {
1749 lin,
1750 lin_pos,
1751 table_all,
1752 host_table,
1753 stash_conv,
1754 stash_ssm,
1755 conv_words,
1756 ssm_words,
1757 stage: std::collections::HashMap::new(),
1758 tap_bufs: std::collections::HashMap::new(),
1759 graphs: std::collections::HashMap::new(),
1760 save_conv: e.uninit(max_run * conv_words)?,
1761 save_ssm: e.uninit(max_run * ssm_words)?,
1762 max_run,
1763 n_embd,
1764 round_slab: false,
1765 }))
1766 }
1767
1768 /// Rebuild the pointer table from the live handles (once per verify — the gdn
1769 /// ping-pong swaps the canonical/alt handles between rounds; a stale table would
1770 /// read the wrong parity's state).
1771 pub(crate) fn refresh_tables(
1772 &mut self,
1773 e: &Engine,
1774 cache: &Cache,
1775 ) -> Result<(), Box<dyn std::error::Error>> {
1776 use cudarc::driver::DevicePtr;
1777 {
1778 let s = &e.gpu.stream();
1779 for (k, &il) in self.lin.iter().enumerate() {
1780 let rl = cache.recur[il].as_ref().unwrap();
1781 let (pc, _g0) = rl.conv_state.device_ptr(s);
1782 let (p0, _g1) = rl.ssm_state.device_ptr(s);
1783 let (p1, _g2) = rl.ssm_state_alt.device_ptr(s);
1784 let o = k * 6;
1785 self.host_table[o] = pc as u64;
1786 self.host_table[o + 1] = p0 as u64;
1787 self.host_table[o + 2] = p1 as u64;
1788 self.host_table[o + 3] = pc as u64;
1789 self.host_table[o + 4] = p1 as u64;
1790 self.host_table[o + 5] = p0 as u64;
1791 }
1792 }
1793 e.htod_u64_into(&self.host_table, &mut self.table_all)?;
1794 Ok(())
1795 }
1796
1797 /// Run layers [start, end) (all linear) as one captured graph at this vt: stage the
1798 /// residual into the bucket's x_in, capture on first encounter (2 executing warmups
1799 /// bracketed by a segment state save/restore), launch, then apply the host parity
1800 /// bookkeeping the captured body would have done. Returns the fresh residual.
1801 #[allow(clippy::too_many_arguments)]
1802 fn run_segment(
1803 &mut self,
1804 model: &crate::hybrid::HybridModel,
1805 e: &Engine,
1806 start: usize,
1807 end: usize,
1808 x: &CudaSlice<f32>,
1809 t: usize,
1810 cache: &mut Cache,
1811 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
1812 let n_embd = self.n_embd;
1813 debug_assert!(end - start <= self.max_run);
1814 if !self.stage.contains_key(&t) {
1815 let xin = e.uninit(t * n_embd)?;
1816 let xout = e.uninit(t * n_embd)?;
1817 self.stage.insert(t, (xin, xout));
1818 }
1819 // Stage the residual at the bucket's baked input address.
1820 {
1821 let (xin, _) = self.stage.get_mut(&t).unwrap();
1822 e.copy_into(xin, 0, x, t * n_embd)?;
1823 }
1824 let key = (start, t);
1825 if !self.graphs.contains_key(&key) {
1826 // The 2 warmups EXECUTE the segment on live state — save conv + the canonical
1827 // ssm of every segment layer first, restore after, so the graph's first real
1828 // launch starts from the exact pre-round state (bytes gated e2e).
1829 for (k, il) in (start..end).enumerate() {
1830 let rl = cache.recur[il].as_ref().unwrap();
1831 e.copy_into(
1832 &mut self.save_conv,
1833 k * self.conv_words,
1834 &rl.conv_state,
1835 self.conv_words,
1836 )?;
1837 e.copy_into(
1838 &mut self.save_ssm,
1839 k * self.ssm_words,
1840 &rl.ssm_state,
1841 self.ssm_words,
1842 )?;
1843 }
1844 let (graph, keeper) = {
1845 let table_all = &self.table_all;
1846 let lin_pos = &self.lin_pos;
1847 let stash_conv = &mut self.stash_conv;
1848 let stash_ssm = &mut self.stash_ssm;
1849 let (xin, xout) = self
1850 .stage
1851 .get_mut(&t)
1852 .map(|(a, b)| (&*a, b))
1853 .expect("stage bucket created above");
1854 let cache_ref: &mut Cache = cache;
1855 // AUTO_FREE_ON_LAUNCH (the retained default): UPLOAD via
1856 // cuGraphInstantiateWithFlags is CUDA_ERROR_INVALID_VALUE (the flag is
1857 // WithParams-only), and the alloc nodes need the auto-free semantics.
1858 // Its launch-time mem-pool scan is the measured limiter — 25.6 us per
1859 // cuGraphLaunch x 16 segments = ~0.41 ms/round, most of the
1860 // eager-launch savings — which is why this door is OPT-IN until the
1861 // node count drops (ctx-scratch transients / fused state chain) or the
1862 // full-verify single-graph (fa exec-update) lands.
1863 e.capture_graph_retained_flags(
1864 cudarc::driver::sys::CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH,
1865 move |e| {
1866 let mut xc: Option<CudaSlice<f32>> = None;
1867 for il in start..end {
1868 let k = lin_pos[&il];
1869 let xr: &CudaSlice<f32> = xc.as_ref().unwrap_or(xin);
1870 let nx = model.qwen35_tparallel_linear_layer(
1871 e,
1872 il,
1873 xr,
1874 t,
1875 cache_ref,
1876 None,
1877 Some((&mut stash_conv[k], &mut stash_ssm[k])),
1878 Some((table_all, k * 6)),
1879 )?;
1880 xc = Some(nx);
1881 }
1882 e.copy_into(xout, 0, xc.as_ref().unwrap(), t * n_embd)?;
1883 Ok(())
1884 },
1885 )?
1886 };
1887 // Undo the net host parity motion of the 3 body runs (each run swaps iff t
1888 // is odd -> 3 runs = net one swap), then restore the device state the
1889 // warmups consumed. The launch below then behaves exactly like one run.
1890 if t % 2 == 1 {
1891 for il in start..end {
1892 let rl = cache.recur[il].as_mut().unwrap();
1893 std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
1894 }
1895 }
1896 for (k, il) in (start..end).enumerate() {
1897 let rl = cache.recur[il].as_mut().unwrap();
1898 let (cw, sw) = (self.conv_words, self.ssm_words);
1899 {
1900 let sv = e.view(&self.save_conv, self.max_run * cw);
1901 let win = sv.slice(k * cw..(k + 1) * cw);
1902 e.copy_view_into(&mut rl.conv_state, 0, &win, cw)?;
1903 }
1904 {
1905 let sv = e.view(&self.save_ssm, self.max_run * sw);
1906 let win = sv.slice(k * sw..(k + 1) * sw);
1907 e.copy_view_into(&mut rl.ssm_state, 0, &win, sw)?;
1908 }
1909 }
1910 self.graphs.insert(
1911 key,
1912 DsparkSegGraph {
1913 graph,
1914 _keeper: keeper,
1915 },
1916 );
1917 }
1918 self.graphs[&key].graph.launch()?;
1919 // Host parity bookkeeping for the replayed body (the captured host swaps do not
1920 // re-run at replay).
1921 if t % 2 == 1 {
1922 for il in start..end {
1923 let rl = cache.recur[il].as_mut().unwrap();
1924 std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
1925 }
1926 }
1927 let (_, xout) = self.stage.get(&t).unwrap();
1928 let mut out = e.uninit(t * n_embd)?;
1929 e.copy_into(&mut out, 0, xout, t * n_embd)?;
1930 Ok(out)
1931 }
1932
1933 /// Slab row (conv, ssm) device pointers + lengths for the commit restore of column
1934 /// `row` (0-based) of layer `il`. None for non-linear layers.
1935 pub(crate) fn slab_row(
1936 &self,
1937 e: &Engine,
1938 il: usize,
1939 row: usize,
1940 ) -> Option<(u64, u64, usize, usize)> {
1941 use cudarc::driver::DevicePtr;
1942 let k = *self.lin_pos.get(&il)?;
1943 let s = &e.gpu.stream();
1944 let (pc, _g0) = self.stash_conv[k].device_ptr(s);
1945 let (ps, _g1) = self.stash_ssm[k].device_ptr(s);
1946 Some((
1947 pc as u64 + (row * self.conv_words * 4) as u64,
1948 ps as u64 + (row * self.ssm_words * 4) as u64,
1949 self.conv_words,
1950 self.ssm_words,
1951 ))
1952 }
1953}
1954
1955impl VerifyCkpt {
1956 fn new(n_layer: usize) -> Self {
1957 VerifyCkpt {
1958 gdn: (0..n_layer).map(|_| None).collect(),
1959 cols: (0..n_layer).map(|_| None).collect(),
1960 }
1961 }
1962}
1963
1964/// The stage-0/TX half of one PP verify. The boundary slot is the ownership token: stage 1
1965/// consumes exactly the slot selected by `tx()` / `tx_pipelined()`, never a slot inferred from
1966/// a logical round number.
1967struct VerifyBoundaryTicket {
1968 rt: &'static crate::pp::PpNRt,
1969 caller_stream: std::sync::Arc<cudarc::driver::CudaStream>,
1970 slot: usize,
1971 pos0: usize,
1972 t: usize,
1973 payload: usize,
1974 n_st: usize,
1975 pipelined: bool,
1976 pp_anatomy: bool,
1977 pp_started: std::time::Instant,
1978 reverse_ms: f64,
1979 stage0_ms: f64,
1980 tx_ms: f64,
1981 trace: Option<SpecPipeTraceCtx>,
1982}
1983
1984/// Explicit OPTIPIPE diagnostic control. Forced modes are set only by `optipipe-gate`; the
1985/// increment-2 controller can also be armed by the server's fresh-process research door.
1986#[derive(Clone, Copy, Debug, PartialEq, Eq)]
1987pub enum OptiForkGateMode {
1988 Disabled,
1989 Hit,
1990 Miss,
1991 Alternate,
1992 Abort,
1993 Controller,
1994}
1995
1996static OPTI_FORK_GATE_MODE: std::sync::atomic::AtomicU8 = std::sync::atomic::AtomicU8::new(0);
1997static OPTI_CONTROLLER_THRESHOLD: std::sync::atomic::AtomicU32 =
1998 std::sync::atomic::AtomicU32::new(0);
1999static OPTI_FORK_ATTEMPTS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
2000static OPTI_FORK_HITS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
2001static OPTI_FORK_MISSES: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
2002static OPTI_FORK_ABORT_DRAINS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
2003static OPTI_FORK_REFUSALS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
2004static OPTI_GATE_CHECKS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
2005static OPTI_GATE_ADMITS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
2006static OPTI_GATE_REJECTS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
2007static OPTI_RECONCILES: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
2008static OPTI_WASTED_DRAFT_TOKENS: std::sync::atomic::AtomicU64 =
2009 std::sync::atomic::AtomicU64::new(0);
2010static OPTI_SHADOW_DRAFT_TOKENS: std::sync::atomic::AtomicU64 =
2011 std::sync::atomic::AtomicU64::new(0);
2012static OPTI_BREAKER_TRIPS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
2013
2014impl OptiForkGateMode {
2015 fn code(self) -> u8 {
2016 match self {
2017 Self::Disabled => 0,
2018 Self::Hit => 1,
2019 Self::Miss => 2,
2020 Self::Alternate => 3,
2021 Self::Abort => 4,
2022 Self::Controller => 5,
2023 }
2024 }
2025
2026 fn configured() -> Self {
2027 match OPTI_FORK_GATE_MODE.load(std::sync::atomic::Ordering::Relaxed) {
2028 1 => Self::Hit,
2029 2 => Self::Miss,
2030 3 => Self::Alternate,
2031 4 => Self::Abort,
2032 5 => Self::Controller,
2033 _ => Self::Disabled,
2034 }
2035 }
2036
2037 fn action(self, generation: u64) -> OptiForkAction {
2038 match self {
2039 Self::Hit => OptiForkAction::Hit,
2040 Self::Miss => OptiForkAction::Miss,
2041 Self::Alternate if generation & 1 == 0 => OptiForkAction::Hit,
2042 Self::Alternate => OptiForkAction::Miss,
2043 Self::Abort => OptiForkAction::Abort,
2044 Self::Disabled | Self::Controller => {
2045 unreachable!("non-forced mode cannot choose a forced fork action")
2046 }
2047 }
2048 }
2049
2050 fn is_forced(self) -> bool {
2051 matches!(self, Self::Hit | Self::Miss | Self::Alternate | Self::Abort)
2052 }
2053}
2054
2055/// Arm or disarm the forced harness. Serving uses only `set_optipipe_controller_threshold`.
2056pub fn set_optipipe_gate_mode(mode: OptiForkGateMode) {
2057 OPTI_FORK_GATE_MODE.store(mode.code(), std::sync::atomic::Ordering::Relaxed);
2058}
2059
2060/// Arm the increment-2 diagnostic controller. The threshold applies to the uncalibrated
2061/// two-token draft-probability product. Serving can call this only through its explicit
2062/// fresh-process research door; the absent-door default remains byte-for-byte disabled.
2063pub fn set_optipipe_controller_threshold(threshold: f32) {
2064 assert!(threshold.is_finite() && (0.0..=1.0).contains(&threshold));
2065 OPTI_CONTROLLER_THRESHOLD.store(threshold.to_bits(), std::sync::atomic::Ordering::Relaxed);
2066 set_optipipe_gate_mode(OptiForkGateMode::Controller);
2067}
2068
2069#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
2070pub struct OptiForkGateStats {
2071 pub attempts: u64,
2072 pub hits: u64,
2073 pub misses: u64,
2074 pub abort_drains: u64,
2075 pub refusals: u64,
2076 pub gate_checks: u64,
2077 pub gate_admits: u64,
2078 pub gate_rejects: u64,
2079 pub reconciles: u64,
2080 pub wasted_draft_tokens: u64,
2081 pub shadow_draft_tokens: u64,
2082 pub breaker_trips: u64,
2083}
2084
2085#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
2086pub struct OptiForkStateIdentity {
2087 pub trunk_kv_bytes: usize,
2088 pub recurrent_bytes: usize,
2089 pub scratch_kv_bytes: usize,
2090 pub hidden_bytes: usize,
2091}
2092
2093pub fn reset_optipipe_gate_stats() {
2094 for counter in [
2095 &OPTI_FORK_ATTEMPTS,
2096 &OPTI_FORK_HITS,
2097 &OPTI_FORK_MISSES,
2098 &OPTI_FORK_ABORT_DRAINS,
2099 &OPTI_FORK_REFUSALS,
2100 &OPTI_GATE_CHECKS,
2101 &OPTI_GATE_ADMITS,
2102 &OPTI_GATE_REJECTS,
2103 &OPTI_RECONCILES,
2104 &OPTI_WASTED_DRAFT_TOKENS,
2105 &OPTI_SHADOW_DRAFT_TOKENS,
2106 &OPTI_BREAKER_TRIPS,
2107 ] {
2108 counter.store(0, std::sync::atomic::Ordering::Relaxed);
2109 }
2110}
2111
2112pub fn optipipe_gate_stats() -> OptiForkGateStats {
2113 let load = |v: &std::sync::atomic::AtomicU64| v.load(std::sync::atomic::Ordering::Relaxed);
2114 OptiForkGateStats {
2115 attempts: load(&OPTI_FORK_ATTEMPTS),
2116 hits: load(&OPTI_FORK_HITS),
2117 misses: load(&OPTI_FORK_MISSES),
2118 abort_drains: load(&OPTI_FORK_ABORT_DRAINS),
2119 refusals: load(&OPTI_FORK_REFUSALS),
2120 gate_checks: load(&OPTI_GATE_CHECKS),
2121 gate_admits: load(&OPTI_GATE_ADMITS),
2122 gate_rejects: load(&OPTI_GATE_REJECTS),
2123 reconciles: load(&OPTI_RECONCILES),
2124 wasted_draft_tokens: load(&OPTI_WASTED_DRAFT_TOKENS),
2125 shadow_draft_tokens: load(&OPTI_SHADOW_DRAFT_TOKENS),
2126 breaker_trips: load(&OPTI_BREAKER_TRIPS),
2127 }
2128}
2129
2130#[derive(Clone, Copy, Debug)]
2131struct OptiControllerPolicy {
2132 threshold: f32,
2133 consecutive_misses: u8,
2134 breaker_tripped: bool,
2135}
2136
2137impl OptiControllerPolicy {
2138 fn configured() -> Self {
2139 Self {
2140 threshold: f32::from_bits(
2141 OPTI_CONTROLLER_THRESHOLD.load(std::sync::atomic::Ordering::Relaxed),
2142 ),
2143 consecutive_misses: 0,
2144 breaker_tripped: false,
2145 }
2146 }
2147
2148 fn admit(&self, q_proxy: f32) -> bool {
2149 q_proxy.is_finite()
2150 && (0.0..=1.0).contains(&q_proxy)
2151 && (self.threshold == 0.0 || (!self.breaker_tripped && q_proxy >= self.threshold))
2152 }
2153
2154 /// Returns true exactly when this resolution newly trips the three-miss breaker.
2155 fn resolve(&mut self, hit: bool) -> bool {
2156 // q*=0 is the lane's explicit unconditional measurement arm. Its purpose is to price
2157 // every optimistic opportunity, so the safety breaker is measured separately and must
2158 // not silently turn this arm into "three attempts then serial".
2159 if self.threshold == 0.0 {
2160 self.consecutive_misses = 0;
2161 return false;
2162 }
2163 if hit {
2164 self.consecutive_misses = 0;
2165 return false;
2166 }
2167 self.consecutive_misses = self.consecutive_misses.saturating_add(1);
2168 if !self.breaker_tripped && self.consecutive_misses >= 3 {
2169 self.breaker_tripped = true;
2170 return true;
2171 }
2172 false
2173 }
2174}
2175
2176#[derive(Clone, Copy, Debug, PartialEq, Eq)]
2177enum OptiForkAction {
2178 Hit,
2179 Miss,
2180 Abort,
2181}
2182
2183#[derive(Clone, Copy, Debug, PartialEq, Eq)]
2184struct OptiForkGeneration {
2185 id: u64,
2186 slot: usize,
2187}
2188
2189#[derive(Default)]
2190struct OptiForkGenerationTracker {
2191 next: u64,
2192 live: [Option<u64>; 2],
2193}
2194
2195impl OptiForkGenerationTracker {
2196 fn reserve(&mut self) -> Result<OptiForkGeneration, Box<dyn std::error::Error>> {
2197 let generation = OptiForkGeneration {
2198 id: self.next,
2199 slot: (self.next & 1) as usize,
2200 };
2201 if let Some(live) = self.live[generation.slot] {
2202 return Err(format!(
2203 "optipipe snapshot slot {} still owns generation {live}; refusing to overwrite it",
2204 generation.slot,
2205 )
2206 .into());
2207 }
2208 self.next += 1;
2209 self.live[generation.slot] = Some(generation.id);
2210 Ok(generation)
2211 }
2212
2213 fn retire(&mut self, generation: OptiForkGeneration) -> Result<(), Box<dyn std::error::Error>> {
2214 match self.live[generation.slot] {
2215 Some(id) if id == generation.id => {
2216 self.live[generation.slot] = None;
2217 Ok(())
2218 }
2219 other => Err(format!(
2220 "optipipe generation teardown mismatch: ticket={} slot={} live={other:?}",
2221 generation.id, generation.slot,
2222 )
2223 .into()),
2224 }
2225 }
2226}
2227
2228struct OptiForkSeedGeneration {
2229 h_seed: CudaSlice<f32>,
2230 fill_prev: CudaSlice<f32>,
2231 scratch_len: usize,
2232}
2233
2234/// Allocate or refresh one full checkpoint through the engine that owns each PP stage. The
2235/// generic cache helper accepts one device and therefore cannot copy GDN state split across
2236/// devices. KV lengths and position stay host metadata; only recurrent buffers need stage-local
2237/// device ownership.
2238fn opti_snapshot_stage_owned(
2239 e: &Engine,
2240 cache: &Cache,
2241 rt: &'static crate::pp::PpNRt,
2242 fence: &[usize],
2243) -> Result<crate::cache::CacheSnapshot, Box<dyn std::error::Error>> {
2244 let n = cache.kv.len();
2245 let mut snapshot = crate::cache::CacheSnapshot {
2246 kv_len: vec![None; n],
2247 conv: (0..n).map(|_| None).collect(),
2248 ssm: (0..n).map(|_| None).collect(),
2249 pos: cache.pos,
2250 };
2251 opti_snapshot_stage_owned_into(e, cache, rt, fence, &mut snapshot)?;
2252 Ok(snapshot)
2253}
2254
2255fn opti_snapshot_stage_owned_into(
2256 e: &Engine,
2257 cache: &Cache,
2258 rt: &'static crate::pp::PpNRt,
2259 fence: &[usize],
2260 snapshot: &mut crate::cache::CacheSnapshot,
2261) -> Result<(), Box<dyn std::error::Error>> {
2262 if fence.len() != rt.n_stages() + 1 || snapshot.kv_len.len() != cache.kv.len() {
2263 return Err("optipipe stage-owned snapshot shape mismatch".into());
2264 }
2265 for stage in 0..rt.n_stages() {
2266 opti_snapshot_one_stage_owned_into(e, cache, rt, fence, stage, snapshot)?;
2267 }
2268 snapshot.pos = cache.pos;
2269 Ok(())
2270}
2271
2272/// Refresh one PP stage of a checkpoint. Increment 2 uses this split form so stage 0's
2273/// optimistic post-N state is captured before N+1 stage 0 is queued, while stage 1's matching
2274/// post-N state is captured only after N stage 1 is enqueued. Calling the all-stage helper at
2275/// either point would capture one side of the fork at the wrong generation.
2276fn opti_snapshot_one_stage_owned_into(
2277 e: &Engine,
2278 cache: &Cache,
2279 rt: &'static crate::pp::PpNRt,
2280 fence: &[usize],
2281 stage: usize,
2282 snapshot: &mut crate::cache::CacheSnapshot,
2283) -> Result<(), Box<dyn std::error::Error>> {
2284 if fence.len() != rt.n_stages() + 1
2285 || snapshot.kv_len.len() != cache.kv.len()
2286 || stage >= rt.n_stages()
2287 {
2288 return Err("optipipe single-stage snapshot shape mismatch".into());
2289 }
2290 let _scope = rt.enter(stage);
2291 let owner = rt.engine(stage, e);
2292 for il in fence[stage]..fence[stage + 1] {
2293 snapshot.kv_len[il] = cache.kv[il].as_ref().map(|kv| kv.len);
2294 match &cache.recur[il] {
2295 Some(recur) => {
2296 match snapshot.conv[il].as_mut() {
2297 Some(dst) => {
2298 owner.copy_into(dst, 0, &recur.conv_state, recur.conv_state.len())?
2299 }
2300 None => snapshot.conv[il] = Some(owner.clone_dtod(&recur.conv_state)?),
2301 }
2302 match snapshot.ssm[il].as_mut() {
2303 Some(dst) => {
2304 owner.copy_into(dst, 0, &recur.ssm_state, recur.ssm_state.len())?
2305 }
2306 None => snapshot.ssm[il] = Some(owner.clone_dtod(&recur.ssm_state)?),
2307 }
2308 }
2309 None if snapshot.conv[il].is_some() || snapshot.ssm[il].is_some() => {
2310 return Err(
2311 format!("optipipe stage-owned snapshot layer {il} changed shape").into(),
2312 );
2313 }
2314 None => {}
2315 }
2316 }
2317 snapshot.pos = cache.pos;
2318 Ok(())
2319}
2320
2321/// Increment-1 persistent fork state. Exactly two snapshot/seed slots alternate; a live ticket
2322/// names its generation and keeps teardown fail-closed. Only stage 0 is allowed to mutate before
2323/// resolve, so the reconcile tables and conditional restores are stage-local.
2324struct OptiForkState {
2325 mode: OptiForkGateMode,
2326 controller: Option<OptiControllerPolicy>,
2327 generations: OptiForkGenerationTracker,
2328 active_snapshot_slot: usize,
2329 alternate_snapshot: crate::cache::CacheSnapshot,
2330 seeds: [OptiForkSeedGeneration; 2],
2331 rt: &'static crate::pp::PpNRt,
2332 fence: [usize; 3],
2333 split: usize,
2334 len_ptrs: CudaSlice<u64>,
2335 saved_lens: CudaSlice<i32>,
2336 forced_acc: CudaSlice<u32>,
2337 valid: CudaSlice<u32>,
2338 stage0_stream: std::sync::Arc<cudarc::driver::CudaStream>,
2339 logical_payload_bytes: [usize; 2],
2340}
2341
2342struct OptiForkTicket {
2343 generation: OptiForkGeneration,
2344 boundary: Option<VerifyBoundaryTicket>,
2345 drain: std::sync::Arc<cudarc::driver::CudaStream>,
2346 settled: bool,
2347}
2348
2349struct OptiControllerTicket {
2350 generation: OptiForkGeneration,
2351 boundary: Option<VerifyBoundaryTicket>,
2352 ckpt: Option<VerifyCkpt>,
2353 verify_tokens: [u32; 2],
2354 draft_prob: f32,
2355 eager_seed: Option<CudaSlice<f32>>,
2356 q_proxy: f32,
2357 scratch_len: usize,
2358 issued_at: std::time::Instant,
2359 drain: std::sync::Arc<cudarc::driver::CudaStream>,
2360 settled: bool,
2361}
2362
2363struct OptiControllerPrepared {
2364 verify_tokens: [u32; 2],
2365 draft_prob: f32,
2366 eager_seed: Option<CudaSlice<f32>>,
2367 q_proxy: f32,
2368 scratch_len: usize,
2369}
2370
2371impl OptiControllerTicket {
2372 fn take_boundary(&mut self) -> VerifyBoundaryTicket {
2373 self.boundary
2374 .take()
2375 .expect("controller boundary ticket already consumed")
2376 }
2377
2378 fn take_ckpt(&mut self) -> VerifyCkpt {
2379 self.ckpt
2380 .take()
2381 .expect("controller verify checkpoint already consumed")
2382 }
2383
2384 fn take_eager_seed(&mut self) -> Option<CudaSlice<f32>> {
2385 self.eager_seed.take()
2386 }
2387
2388 fn settle(&mut self) {
2389 self.settled = true;
2390 }
2391}
2392
2393impl Drop for OptiControllerTicket {
2394 fn drop(&mut self) {
2395 if !self.settled {
2396 let _ = self.drain.synchronize();
2397 OPTI_FORK_ABORT_DRAINS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2398 }
2399 }
2400}
2401
2402impl OptiForkTicket {
2403 fn take_boundary(&mut self) -> VerifyBoundaryTicket {
2404 self.boundary
2405 .take()
2406 .expect("fork ticket boundary already consumed")
2407 }
2408
2409 fn settle(&mut self) {
2410 self.settled = true;
2411 }
2412}
2413
2414impl Drop for OptiForkTicket {
2415 fn drop(&mut self) {
2416 if !self.settled {
2417 let _ = self.drain.synchronize();
2418 OPTI_FORK_ABORT_DRAINS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2419 }
2420 }
2421}
2422
2423impl OptiForkState {
2424 #[allow(clippy::too_many_arguments)]
2425 fn new(
2426 e: &Engine,
2427 cache: &Cache,
2428 mode: OptiForkGateMode,
2429 alternate_snapshot: crate::cache::CacheSnapshot,
2430 h_seed: &CudaSlice<f32>,
2431 fill_prev: &CudaSlice<f32>,
2432 rt: &'static crate::pp::PpNRt,
2433 split: usize,
2434 n_layer: usize,
2435 ) -> Result<Self, Box<dyn std::error::Error>> {
2436 let fence = [0, split, n_layer];
2437 let mut logical_payload_bytes = [0usize; 2];
2438 for stage in 0..2 {
2439 for il in fence[stage]..fence[stage + 1] {
2440 logical_payload_bytes[stage] += alternate_snapshot.conv[il]
2441 .as_ref()
2442 .map_or(0, |v| v.len() * std::mem::size_of::<f32>());
2443 logical_payload_bytes[stage] += alternate_snapshot.ssm[il]
2444 .as_ref()
2445 .map_or(0, |v| v.len() * std::mem::size_of::<f32>());
2446 }
2447 }
2448 let seeds = [
2449 OptiForkSeedGeneration {
2450 h_seed: e.clone_dtod(h_seed)?,
2451 fill_prev: e.clone_dtod(fill_prev)?,
2452 scratch_len: 0,
2453 },
2454 OptiForkSeedGeneration {
2455 h_seed: e.clone_dtod(h_seed)?,
2456 fill_prev: e.clone_dtod(fill_prev)?,
2457 scratch_len: 0,
2458 },
2459 ];
2460 let (len_ptrs, saved_lens, forced_acc, valid, stage0_stream) = {
2461 let _stage = rt.enter(0);
2462 let e0 = rt.engine(0, e);
2463 (
2464 crate::round_stream::kv_len_ptr_table_range(e0, cache, 0..split, None)?,
2465 e0.htod_i32(&vec![0; split])?,
2466 e0.alloc_u32_zeroed(2)?,
2467 e0.alloc_u32_zeroed(1)?,
2468 e0.stream(),
2469 )
2470 };
2471 logical_payload_bytes[0] += seeds
2472 .iter()
2473 .map(|seed| (seed.h_seed.len() + seed.fill_prev.len()) * std::mem::size_of::<f32>())
2474 .sum::<usize>();
2475 logical_payload_bytes[0] += len_ptrs.len() * std::mem::size_of::<u64>()
2476 + saved_lens.len() * std::mem::size_of::<i32>()
2477 + forced_acc.len() * std::mem::size_of::<u32>()
2478 + valid.len() * std::mem::size_of::<u32>();
2479 Ok(Self {
2480 mode,
2481 controller: (mode == OptiForkGateMode::Controller)
2482 .then(OptiControllerPolicy::configured),
2483 generations: OptiForkGenerationTracker::default(),
2484 active_snapshot_slot: 0,
2485 alternate_snapshot,
2486 seeds,
2487 rt,
2488 fence,
2489 split,
2490 len_ptrs,
2491 saved_lens,
2492 forced_acc,
2493 valid,
2494 stage0_stream,
2495 logical_payload_bytes,
2496 })
2497 }
2498
2499 fn reserve(
2500 &mut self,
2501 current_snapshot: &mut crate::cache::CacheSnapshot,
2502 ) -> Result<OptiForkGeneration, Box<dyn std::error::Error>> {
2503 let generation = self.generations.reserve()?;
2504 if generation.slot != self.active_snapshot_slot {
2505 std::mem::swap(current_snapshot, &mut self.alternate_snapshot);
2506 self.active_snapshot_slot = generation.slot;
2507 }
2508 Ok(generation)
2509 }
2510
2511 fn capture_seed(
2512 &mut self,
2513 e: &Engine,
2514 generation: OptiForkGeneration,
2515 h_seed: &CudaSlice<f32>,
2516 fill_prev: &CudaSlice<f32>,
2517 scratch_len: usize,
2518 ) -> Result<(), Box<dyn std::error::Error>> {
2519 let seed = &mut self.seeds[generation.slot];
2520 e.copy_into(&mut seed.h_seed, 0, h_seed, h_seed.len())?;
2521 e.copy_into(&mut seed.fill_prev, 0, fill_prev, fill_prev.len())?;
2522 seed.scratch_len = scratch_len;
2523 Ok(())
2524 }
2525
2526 fn ticket(
2527 &self,
2528 generation: OptiForkGeneration,
2529 boundary: VerifyBoundaryTicket,
2530 ) -> OptiForkTicket {
2531 OptiForkTicket {
2532 generation,
2533 boundary: Some(boundary),
2534 drain: self.stage0_stream.clone(),
2535 settled: false,
2536 }
2537 }
2538
2539 #[allow(clippy::too_many_arguments)]
2540 fn controller_ticket(
2541 &self,
2542 generation: OptiForkGeneration,
2543 boundary: VerifyBoundaryTicket,
2544 ckpt: VerifyCkpt,
2545 verify_tokens: [u32; 2],
2546 draft_prob: f32,
2547 eager_seed: Option<CudaSlice<f32>>,
2548 q_proxy: f32,
2549 scratch_len: usize,
2550 ) -> OptiControllerTicket {
2551 OptiControllerTicket {
2552 generation,
2553 boundary: Some(boundary),
2554 ckpt: Some(ckpt),
2555 verify_tokens,
2556 draft_prob,
2557 eager_seed,
2558 q_proxy,
2559 scratch_len,
2560 issued_at: std::time::Instant::now(),
2561 drain: self.stage0_stream.clone(),
2562 settled: false,
2563 }
2564 }
2565
2566 fn reserve_successor(&mut self) -> Result<OptiForkGeneration, Box<dyn std::error::Error>> {
2567 self.generations.reserve()
2568 }
2569
2570 fn successor_snapshot_mut(&mut self) -> &mut crate::cache::CacheSnapshot {
2571 &mut self.alternate_snapshot
2572 }
2573
2574 fn promote_successor_snapshot(
2575 &mut self,
2576 current_snapshot: &mut crate::cache::CacheSnapshot,
2577 generation: OptiForkGeneration,
2578 ) {
2579 std::mem::swap(current_snapshot, &mut self.alternate_snapshot);
2580 self.active_snapshot_slot = generation.slot;
2581 }
2582
2583 fn queue_actual_reconcile(
2584 &mut self,
2585 e: &Engine,
2586 snapshot: &crate::cache::CacheSnapshot,
2587 acc: &CudaSlice<u32>,
2588 optimistic_pending: u32,
2589 base: usize,
2590 ) -> Result<(), Box<dyn std::error::Error>> {
2591 let saved: Vec<i32> = (0..self.split)
2592 .map(|il| snapshot.kv_len[il].map(|v| v as i32).unwrap_or(0))
2593 .collect();
2594 // Serving keeps the caller/accept walk on the head (stage-1) device. Record the accept
2595 // decision point there and append a wait to stage 0 after its optimistic successor/TX;
2596 // the validity/reconcile kernels must never peer-read acc before it is written. The
2597 // increment-1 harness uses primary stage 0, where stream order already provides this.
2598 if self.rt.engine(0, e).ctx().ordinal() != e.ctx().ordinal() {
2599 self.rt.fence_stages_behind(&e.stream())?;
2600 }
2601 let _stage = self.rt.enter(0);
2602 let e0 = self.rt.engine(0, e);
2603 e0.htod_i32_into(&mut self.saved_lens, &saved)?;
2604 e0.spec_fork_valid(acc, optimistic_pending, &mut self.valid)?;
2605 e0.spec_fork_reconcile_kv(
2606 &self.len_ptrs,
2607 &self.saved_lens,
2608 acc,
2609 &self.valid,
2610 base,
2611 self.split,
2612 )
2613 }
2614
2615 fn finish_actual_reconcile(
2616 &mut self,
2617 e: &Engine,
2618 cache: &mut Cache,
2619 snapshot: &crate::cache::CacheSnapshot,
2620 n_acc: usize,
2621 base: usize,
2622 hit: bool,
2623 ) -> Result<(), Box<dyn std::error::Error>> {
2624 if hit {
2625 return Ok(());
2626 }
2627 let len_delta = base + n_acc;
2628 for il in 0..self.split {
2629 if let (Some(kv), Some(saved)) = (cache.kv[il].as_mut(), snapshot.kv_len[il]) {
2630 kv.len = saved + len_delta;
2631 }
2632 }
2633 {
2634 let _stage = self.rt.enter(1);
2635 let e1 = self.rt.engine(1, e);
2636 for il in self.split..self.fence[2] {
2637 if let (Some(kv), Some(saved)) = (cache.kv[il].as_mut(), snapshot.kv_len[il]) {
2638 kv.len = saved + len_delta;
2639 e1.set_i32_one(&mut kv.len_d, kv.len as i32)?;
2640 }
2641 }
2642 }
2643 self.rt.publish_to(0, &e.stream())?;
2644 Ok(())
2645 }
2646
2647 fn cancel_controller_ticket(
2648 &mut self,
2649 e: &Engine,
2650 cache: &mut Cache,
2651 scratch: &mut MtpScratch,
2652 snapshot: &crate::cache::CacheSnapshot,
2653 ticket: &mut OptiControllerTicket,
2654 ) -> Result<(), Box<dyn std::error::Error>> {
2655 {
2656 let _stage = self.rt.enter(0);
2657 let e0 = self.rt.engine(0, e);
2658 for il in 0..self.split {
2659 if let (Some(kv), Some(saved)) = (cache.kv[il].as_mut(), snapshot.kv_len[il]) {
2660 kv.len = saved;
2661 e0.set_i32_one(&mut kv.len_d, saved as i32)?;
2662 }
2663 }
2664 }
2665 scratch.set_len(e, snapshot.pos)?;
2666 ticket.settle();
2667 self.generations.retire(ticket.generation)?;
2668 OPTI_FORK_ABORT_DRAINS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2669 OPTI_WASTED_DRAFT_TOKENS.fetch_add(2, std::sync::atomic::Ordering::Relaxed);
2670 eprintln!(
2671 "[opti-controller] tail-drain generation={} slot={}",
2672 ticket.generation.id, ticket.generation.slot,
2673 );
2674 Ok(())
2675 }
2676
2677 #[allow(clippy::too_many_arguments)]
2678 fn reconcile(
2679 &mut self,
2680 e: &Engine,
2681 cache: &mut Cache,
2682 scratch: &mut MtpScratch,
2683 snapshot: &crate::cache::CacheSnapshot,
2684 h_seed: &mut CudaSlice<f32>,
2685 fill_prev: &mut CudaSlice<f32>,
2686 generation: OptiForkGeneration,
2687 action: OptiForkAction,
2688 optimistic_pending: u32,
2689 ) -> Result<(), Box<dyn std::error::Error>> {
2690 debug_assert!(action != OptiForkAction::Abort);
2691 let miss_started = std::time::Instant::now();
2692 let keep = action == OptiForkAction::Hit;
2693 let saved: Vec<i32> = (0..self.split)
2694 .map(|il| snapshot.kv_len[il].map(|v| v as i32).unwrap_or(0))
2695 .collect();
2696 let seed = &self.seeds[generation.slot];
2697 {
2698 let _stage = self.rt.enter(0);
2699 let e0 = self.rt.engine(0, e);
2700 e0.htod_i32_into(&mut self.saved_lens, &saved)?;
2701 let forced = if keep {
2702 [1u32, optimistic_pending]
2703 } else {
2704 [0u32, optimistic_pending]
2705 };
2706 e0.htod_u32_into(&mut self.forced_acc, &forced)?;
2707 e0.spec_fork_valid(&self.forced_acc, optimistic_pending, &mut self.valid)?;
2708 e0.spec_fork_reconcile_kv(
2709 &self.len_ptrs,
2710 &self.saved_lens,
2711 &self.forced_acc,
2712 &self.valid,
2713 0,
2714 self.split,
2715 )?;
2716 for il in 0..self.split {
2717 if let Some(recur) = cache.recur[il].as_mut() {
2718 let conv = snapshot.conv[il]
2719 .as_ref()
2720 .ok_or("optipipe stage0 snapshot missing conv state")?;
2721 let ssm = snapshot.ssm[il]
2722 .as_ref()
2723 .ok_or("optipipe stage0 snapshot missing ssm state")?;
2724 e0.spec_fork_restore_f32(conv, &mut recur.conv_state, &self.valid)?;
2725 e0.spec_fork_restore_f32(ssm, &mut recur.ssm_state, &self.valid)?;
2726 }
2727 }
2728 e0.spec_fork_restore_f32(&seed.h_seed, h_seed, &self.valid)?;
2729 e0.spec_fork_restore_f32(&seed.fill_prev, fill_prev, &self.valid)?;
2730 }
2731
2732 if keep {
2733 OPTI_FORK_HITS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2734 return Ok(());
2735 }
2736
2737 for il in 0..self.split {
2738 if let (Some(kv), Some(saved)) = (cache.kv[il].as_mut(), snapshot.kv_len[il]) {
2739 kv.len = saved;
2740 }
2741 }
2742 scratch.set_len(e, seed.scratch_len)?;
2743 // Targeted E_restart: publish only stage 0's reconcile to the caller, then bound the
2744 // forced diagnostic so the retained number is the actual miss cost, not enqueue time.
2745 let caller = e.stream();
2746 self.rt.publish_to(0, &caller)?;
2747 caller.synchronize()?;
2748 let miss_ms = miss_started.elapsed().as_secs_f64() * 1e3;
2749 eprintln!(
2750 "[opti-fork-reconcile] generation={} slot={} miss_ms={miss_ms:.3}",
2751 generation.id, generation.slot,
2752 );
2753 OPTI_FORK_MISSES.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2754 Ok(())
2755 }
2756
2757 fn retire(&mut self, generation: OptiForkGeneration) -> Result<(), Box<dyn std::error::Error>> {
2758 self.generations.retire(generation)
2759 }
2760}
2761
2762impl HybridModel {
2763 fn opti_graph_draft_step(
2764 &self,
2765 e: &Engine,
2766 mtp: &MtpHead,
2767 dctx: &mut DraftGraphCtx,
2768 scratch: &mut MtpScratch,
2769 d_vocab: usize,
2770 ) -> Result<(u32, f32), Box<dyn std::error::Error>> {
2771 dctx.graph
2772 .as_ref()
2773 .ok_or("optipipe controller requires the greedy draft graph")?
2774 .launch()?;
2775 scratch.kv.len += 1;
2776 let idx = e.dtoh_u32_one(&dctx.g_tok)?;
2777 if (idx as usize) >= d_vocab {
2778 return Err(
2779 format!("optipipe draft argmax sentinel 0x{idx:08x} >= d_vocab {d_vocab}").into(),
2780 );
2781 }
2782 let probability = e.dtoh(&dctx.g_p)?[0];
2783 if !probability.is_finite() || !(0.0..=1.0).contains(&probability) {
2784 return Err(format!("optipipe draft probability is invalid: {probability}").into());
2785 }
2786 let token = match &mtp.d2t {
2787 Some(map) => map[idx as usize],
2788 None => idx,
2789 };
2790 if token != idx {
2791 e.set_u32_one(&mut dctx.g_tok, token)?;
2792 }
2793 Ok((token, probability))
2794 }
2795
2796 #[allow(clippy::too_many_arguments)]
2797 fn opti_controller_draft_step(
2798 &self,
2799 e: &Engine,
2800 mtp: &MtpHead,
2801 dctx: &mut DraftGraphCtx,
2802 scratch: &mut MtpScratch,
2803 d_vocab: usize,
2804 eager_state: &mut Option<(u32, CudaSlice<f32>)>,
2805 eager_pos: usize,
2806 embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
2807 ) -> Result<(u32, f32), Box<dyn std::error::Error>> {
2808 if dctx.graph.is_some() {
2809 return self.opti_graph_draft_step(e, mtp, dctx, scratch, d_vocab);
2810 }
2811 let (input_token, input_seed) = eager_state
2812 .take()
2813 .ok_or("optipipe eager continuation seed is unavailable")?;
2814 let (logits, next_seed) = self.mtp_head_forward_dev(
2815 e,
2816 mtp,
2817 input_token,
2818 &input_seed,
2819 scratch,
2820 eager_pos,
2821 embd_dev,
2822 None,
2823 )?;
2824 let token_d = e.argmax_token_device(&logits, d_vocab)?;
2825 let idx = e.dtoh_u32_one(&token_d)?;
2826 if (idx as usize) >= d_vocab {
2827 return Err(format!(
2828 "optipipe eager draft argmax sentinel 0x{idx:08x} >= d_vocab {d_vocab}"
2829 )
2830 .into());
2831 }
2832 let probability_d = e.prob_of_token_device(&logits, &token_d, d_vocab)?;
2833 let probability = e.dtoh(&probability_d)?[0];
2834 if !probability.is_finite() || !(0.0..=1.0).contains(&probability) {
2835 return Err(
2836 format!("optipipe eager draft probability is invalid: {probability}").into(),
2837 );
2838 }
2839 let token = match &mtp.d2t {
2840 Some(map) => map[idx as usize],
2841 None => idx,
2842 };
2843 *eager_state = Some((token, next_seed));
2844 Ok((token, probability))
2845 }
2846
2847 /// NextN head forward for ONE draft token (§A ops 1-13, T=1).
2848 /// Inputs: `e_tok` = the token to predict FROM (last committed / previous draft); `h_seed` =
2849 /// the trunk's pre-output_norm hidden of that token (§A op 2 input). `mtp_pos` = absolute
2850 /// position of the token being predicted from. Returns (draft_logits[n_vocab] host, h_nextn dev).
2851 /// `h_nextn` (§A op 10) becomes `h_seed` for the next autoregressive draft step.
2852 /// Device-resident: returns draft logits ON DEVICE (no [n_vocab] dtoh). The greedy draft
2853 /// loop only needs argmax — paired with `argmax_token_device` this cuts the ~600KB logits
2854 /// transfer + host argmax per draft token from the K-token draft chain.
2855 #[allow(clippy::too_many_arguments)]
2856 fn mtp_head_forward_dev(
2857 &self,
2858 e: &Engine,
2859 mtp: &MtpHead,
2860 e_tok: u32,
2861 h_seed: &CudaSlice<f32>,
2862 scratch: &mut MtpScratch,
2863 mtp_pos: usize,
2864 embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
2865 // DRAFT-SIDE GRAMMAR MASK (lane/draft-mask): (packed draft-vocab allowed set, words).
2866 // Applied to the head logits BEFORE they are returned, so every consumer (argmax,
2867 // gumbel draw, p-min prob) sees the grammar-legal row. None = unmasked (pre-lane).
2868 mask: Option<(&CudaSlice<u32>, usize)>,
2869 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
2870 // MEMRA_SPEC_ANATOMY=1 — eager-step phase timers (diagnostic only). Phase boundaries
2871 // sync the stream, so absolute time inflates; the BREAKDOWN is the signal. Cumulative
2872 // summary on stderr every 128 steps: glue (embed..attn_norm), attn, ffn, head.
2873 use std::sync::atomic::{AtomicU64, Ordering::Relaxed};
2874 static ANAT_NS: [AtomicU64; 5] = [
2875 AtomicU64::new(0),
2876 AtomicU64::new(0),
2877 AtomicU64::new(0),
2878 AtomicU64::new(0),
2879 AtomicU64::new(0),
2880 ];
2881 static ANAT_STEPS: AtomicU64 = AtomicU64::new(0);
2882 let anat = {
2883 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2884 *ON.get_or_init(|| std::env::var("MEMRA_SPEC_ANATOMY").as_deref() == Ok("1"))
2885 };
2886 if anat {
2887 e.stream().synchronize()?; // drain prior queue so phase 0 starts clean
2888 }
2889 let t_all = std::time::Instant::now();
2890 let mut t_ph = std::time::Instant::now();
2891 let mut anat_mark = |i: usize,
2892 e: &Engine,
2893 t: &mut std::time::Instant|
2894 -> Result<(), Box<dyn std::error::Error>> {
2895 if anat {
2896 e.stream().synchronize()?;
2897 ANAT_NS[i].fetch_add(t.elapsed().as_nanos() as u64, Relaxed);
2898 *t = std::time::Instant::now();
2899 }
2900 Ok(())
2901 };
2902 let cfg = &self.cfg;
2903 let n_embd = cfg.n_embd as usize;
2904 // Distilled-student geometry: the block runs at the INNER width `di` (eh_proj out /
2905 // attn / ffn); the n_embd interface (embed, norms in, carrier out, head in) is unchanged.
2906 let di = mtp.geom.as_ref().map(|g| g.d_inner).unwrap_or(n_embd);
2907 let eps = cfg.rms_eps;
2908 let pos_d = e.htod_i32(&[mtp_pos as i32])?;
2909
2910 // op A: a resident table transfers one 4B token id. The exact host-row capacity path
2911 // expands this one row on CPU and transfers n_embd f32 values instead.
2912 let e_emb = match embd_dev {
2913 Some((g, qt, rb)) => e.embed_gather_device_t(g, &[e_tok], n_embd, qt, rb)?,
2914 None => e.htod(&self.embd.gather(n_embd, &[e_tok]))?,
2915 };
2916
2917 // op 1/2: e_norm = RMSNorm(e, enorm); h_norm = RMSNorm(h_seed, hnorm)
2918 let mut e_norm = e.zeros(n_embd)?;
2919 e.rms_norm(&e_emb, mtp.enorm.float_data(), &mut e_norm, n_embd, 1, eps)?;
2920 let mut h_norm = e.zeros(n_embd)?;
2921 e.rms_norm(h_seed, mtp.hnorm.float_data(), &mut h_norm, n_embd, 1, eps)?;
2922
2923 // op 3: concat = [e_norm ; h_norm] -> [2*n_embd], e_norm in [0,n_embd), h_norm in [n_embd,2n_embd)
2924 let mut concat = e.zeros(2 * n_embd)?;
2925 e.copy_into(&mut concat, 0, &e_norm, n_embd)?;
2926 e.copy_into(&mut concat, n_embd, &h_norm, n_embd)?;
2927
2928 // op 4: inpSA = eh_proj @ concat (eh_proj [2*n_embd, n_embd]) -> [n_embd]
2929 let inp_sa = e.matmul(&mtp.eh_proj, &concat, 1)?;
2930
2931 // op 5: a_norm = RMSNorm(inpSA, attn_norm)
2932 let mut a_norm = e.zeros(di)?;
2933 e.rms_norm(&inp_sa, mtp.attn_norm.float_data(), &mut a_norm, di, 1, eps)?;
2934 anat_mark(0, e, &mut t_ph)?;
2935
2936 // op 6: attention on the scratch KV. SAME dc launcher as the graph path (bucket_max =
2937 // scratch.cap, length from the device len_d) so eager drafts match graph drafts
2938 // bit-for-bit at any t_kv (the parity gate). Host len mirrored here (the dc append
2939 // advances only the device counter).
2940 let attn_out = match (&mtp.mixer, mtp.step35.as_ref()) {
2941 // step35 MTP block: PER-LAYER geometry + a separate head-wise gate + an SWA window,
2942 // none of which the dc launcher can express (see `mtp_step35_attn`). Host-len arm.
2943 // Advances BOTH the host len and the device counter itself (unlike the dc arm,
2944 // whose host-side mirror the caller does).
2945 (Mixer::Full(fa), Some(g)) => {
2946 self.mtp_step35_attn(e, fa, g, &a_norm, &pos_d, scratch)?
2947 }
2948 (Mixer::Full(fa), None) => {
2949 let out =
2950 self.mtp_full_attn_dc(e, fa, &a_norm, &pos_d, scratch, mtp.geom.as_ref())?;
2951 scratch.kv.len += 1;
2952 out
2953 }
2954 (Mixer::Linear(_), _) => {
2955 panic!("MTP block is full-attn in qwen35; linear MTP not supported")
2956 }
2957 (Mixer::Mla(_), _) => crate::hybrid::mla_forward_unimplemented(),
2958 };
2959 anat_mark(1, e, &mut t_ph)?;
2960
2961 // op 7: x1 = inpSA + attn_out
2962 let mut x1 = e.zeros(di)?;
2963 e.add(&inp_sa, &attn_out, &mut x1, di)?;
2964
2965 // op 8: z = RMSNorm(x1, post_attn_norm) (pre-FFN norm)
2966 let mut z = e.zeros(di)?;
2967 e.rms_norm(&x1, mtp.post_attn_norm.float_data(), &mut z, di, 1, eps)?;
2968
2969 // op 9: FFN (Dense or MoE) — same as the trunk decode FFN
2970 let ffn_out = match &mtp.ffn {
2971 crate::hybrid::Ffn::Dense {
2972 ffn_gate,
2973 ffn_up,
2974 ffn_down,
2975 } => {
2976 let n_ff = ffn_gate.out_features();
2977 let (gate, up) = if e.uses_q8_1_fast(ffn_gate) && e.uses_q8_1_fast(ffn_up) {
2978 let (zq, zd) = e.quantize_q8_1(&z, 1, di)?;
2979 (
2980 e.matmul_pre(ffn_gate, &zq, &zd, &z, 1)?,
2981 e.matmul_pre(ffn_up, &zq, &zd, &z, 1)?,
2982 )
2983 } else {
2984 (e.matmul(ffn_gate, &z, 1)?, e.matmul(ffn_up, &z, 1)?)
2985 };
2986 let mut act = e.zeros(n_ff)?;
2987 // step35: a DENSE FFN reads the per-layer SHEXP clamp (upstream's one `build_ffn`
2988 // serves the dense MLP and the shared expert off `swiglu_clamp_shexp` —
2989 // llama-graph.cpp:1751), resolved for the MTP block's OWN index. Every other arch
2990 // passes None, which is `ffn_act`'s dispatch verbatim.
2991 Self::ffn_act_lim(
2992 e,
2993 &self.cfg,
2994 &gate,
2995 &up,
2996 1.0,
2997 1.0,
2998 mtp.step35.as_ref().and_then(|s| s.clamp_shexp),
2999 &mut act,
3000 n_ff,
3001 )?;
3002 e.matmul(ffn_down, &act, 1)?
3003 }
3004 // MTP head is a distinct block — key its experts under a separate layer index (u16::MAX)
3005 // so they never alias trunk layer 0's cache keys.
3006 crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il(e, m, &z, 1, u16::MAX)?,
3007 };
3008 anat_mark(2, e, &mut t_ph)?;
3009
3010 // op 10: h_nextn = x1 + ffn_out (at di)
3011 let mut h_inner = e.zeros(di)?;
3012 e.add(&x1, &ffn_out, &mut h_inner, di)?;
3013
3014 // op 10.5 (student): up-project the inner hidden back to n_embd — training semantics:
3015 // the chain carrier AND the head input are out_up(h_inner) (pre-final-norm).
3016 let h_nextn = match mtp.geom.as_ref() {
3017 Some(g) => e.matmul(&g.out_up, &h_inner, 1)?,
3018 None => h_inner,
3019 };
3020
3021 // op 11: final = RMSNorm(h_nextn, shared_head_norm OR output_norm)
3022 let final_norm = mtp.shared_head_norm.as_ref().unwrap_or(&self.output_norm);
3023 let mut final_h = e.zeros(n_embd)?;
3024 e.rms_norm(
3025 &h_nextn,
3026 final_norm.float_data(),
3027 &mut final_h,
3028 n_embd,
3029 1,
3030 eps,
3031 )?;
3032
3033 // op 12: draft_logits = (shared_head_head OR output) @ final — stays ON DEVICE.
3034 let head = mtp.shared_head_head.as_ref().unwrap_or(&self.output);
3035 let mut logits = e.matmul(head, &final_h, 1)?;
3036 // op 12b (lane/draft-mask): grammar mask over the DRAFT vocab, applied here so the
3037 // caller's argmax / gumbel draw / p-min prob all read the grammar-legal row.
3038 if let Some((mask_d, mw)) = mask {
3039 let d_vocab = head.out_features();
3040 e.mask_logits_col(&mut logits, mask_d, 0, d_vocab, mw)?;
3041 }
3042 anat_mark(3, e, &mut t_ph)?;
3043 if anat {
3044 ANAT_NS[4].fetch_add(t_all.elapsed().as_nanos() as u64, Relaxed);
3045 let n = ANAT_STEPS.fetch_add(1, Relaxed) + 1;
3046 if n % 128 == 0 {
3047 let us = |i: usize| ANAT_NS[i].load(Relaxed) / n / 1000;
3048 eprintln!(
3049 "[spec-anatomy] steps={n} avg us/step: glue={} attn={} ffn={} head={} total={}",
3050 us(0),
3051 us(1),
3052 us(2),
3053 us(3),
3054 us(4)
3055 );
3056 }
3057 }
3058 // Chain recurrence hand-over: pre-norm h_nextn (default) or post-norm final_h
3059 // (MEMRA_SPEC_HPOST — llama.cpp #24025's t_h_nextn is taken AFTER the head norm).
3060 Ok((logits, if spec_hpost() { final_h } else { h_nextn }))
3061 }
3062
3063 /// step35 MTP-block attention, T=1, on the scratch KV — the EAGER-ONLY twin of
3064 /// `mtp_full_attn_dc`. Three things force a separate arm rather than a geometry parameter on
3065 /// the dc path, and all three are properties of this arch's MTP block:
3066 ///
3067 /// 1. **The SWA window.** Block 45 is an SWA-type block (`sliding_window_pattern[45]=true`,
3068 /// window 512). Windowed decode in memra is a token-aligned VIEW OFFSET into the quantized
3069 /// cache (the gemma4 R6 / `step35_decode_attn` pattern: keys carry absolute rope and the
3070 /// mask is purely positional, so one query at `len-1` attending the last `win` rows IS the
3071 /// windowed result). `fa_decode_dc` takes the key count from a DEVICE counter and always
3072 /// starts at row 0 — it cannot express a nonzero offset, so a windowed dc arm would need a
3073 /// new kernel. That is deliberately not built here: see the CUDA-graph note below.
3074 /// 2. **Per-layer head count.** 96 q heads over 8 KV (GQA 12) at this block, vs the trunk's 64
3075 /// on its full-attn layers. The trunk cfg's `n_head` scalar is the MAX over layers, and the
3076 /// trunk ARTIFACT's per-layer arrays stop at index 44 — so the count must come from the
3077 /// resolved `Step35MtpGeom`, never from `cfg`.
3078 /// 3. **The separate head-wise gate.** `blk.45.attn_gate.weight [n_embd, 96]` produces one
3079 /// sigmoid scalar per head (broadcast over head_dim) — `attn_head_gate`, not the qwen35
3080 /// fused-into-wq `q_gate_split` form the dc arm handles.
3081 ///
3082 /// WHY EAGER-ONLY IS NOT A GAP TODAY: `mtp_head_forward_cap` refuses step35 heads
3083 /// explicitly (the SWA-window refusal below), so the graph draft never engages for step35
3084 /// models regardless of eligibility — the eager chain IS the served path. `mtp_head_forward_cap` refuses step35 explicitly rather
3085 /// than silently capturing a window-less (wrong past `win` draft rows) graph.
3086 ///
3087 /// Unlike the dc arm this advances BOTH the host `kv.len` and the device counter, so the
3088 /// caller must not mirror.
3089 fn mtp_step35_attn(
3090 &self,
3091 e: &Engine,
3092 fa: &FullAttnLayer,
3093 g: &crate::hybrid::Step35MtpGeom,
3094 h: &CudaSlice<f32>,
3095 pos_d: &CudaSlice<i32>,
3096 scratch: &mut MtpScratch,
3097 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3098 let (nh, nkv, hd) = (g.n_head, g.n_head_kv, self.cfg.head_dim_k as usize);
3099 let eps = self.cfg.rms_eps;
3100 let scale = 1.0 / (hd as f32).sqrt(); // step35.cpp:255 kq_scale
3101 let n_embd = self.cfg.n_embd as usize;
3102 let gw = fa
3103 .attn_gate
3104 .as_ref()
3105 .ok_or("step35 MTP block is missing attn_gate.weight (head-wise attention gate)")?;
3106
3107 let (q0, k0, v0, gt) = if e.uses_q8_1_fast(&fa.wq)
3108 && e.uses_q8_1_fast(&fa.wk)
3109 && e.uses_q8_1_fast(&fa.wv)
3110 && e.uses_q8_1_fast(gw)
3111 {
3112 let (hq, hdq) = e.quantize_q8_1(h, 1, n_embd)?;
3113 let (a, b, c) = match e.matmul_q8_fused3(&fa.wq, &fa.wk, &fa.wv, &hq, &hdq)? {
3114 Some(t3) => t3,
3115 None => (
3116 e.matmul_pre(&fa.wq, &hq, &hdq, h, 1)?,
3117 e.matmul_pre(&fa.wk, &hq, &hdq, h, 1)?,
3118 e.matmul_pre(&fa.wv, &hq, &hdq, h, 1)?,
3119 ),
3120 };
3121 (a, b, c, e.matmul_pre(gw, &hq, &hdq, h, 1)?)
3122 } else {
3123 (
3124 e.matmul(&fa.wq, h, 1)?,
3125 e.matmul(&fa.wk, h, 1)?,
3126 e.matmul(&fa.wv, h, 1)?,
3127 e.matmul(gw, h, 1)?,
3128 )
3129 };
3130
3131 let mut q = e.uninit(nh * hd)?;
3132 e.rms_norm(&q0, fa.q_norm.float_data(), &mut q, hd, nh, eps)?;
3133 let mut k = e.uninit(nkv * hd)?;
3134 e.rms_norm(&k0, fa.k_norm.float_data(), &mut k, hd, nkv, eps)?;
3135 // `rope_freqs.weight` (llama3 factors) applies to the FULL-attn layers ONLY; SWA passes
3136 // null (llama-hparams / step35.cpp). Block 45 is SWA, so `ff` is None there — but read
3137 // the resolved flag, not the constant, so an all-full sibling stays correct.
3138 let ff = if g.swa {
3139 None
3140 } else {
3141 self.step35_aux.as_ref().and_then(|a| a.rope_freqs(e))
3142 };
3143 #[cfg(debug_assertions)]
3144 if let Some(ff) = ff {
3145 crate::debug_assert_tensor_stream_device(ff, &e.stream(), "mtp_step35_attn.rope_freqs");
3146 }
3147 e.rope_neox2(
3148 &mut q,
3149 &mut k,
3150 pos_d,
3151 hd,
3152 g.n_rot,
3153 nh,
3154 nkv,
3155 1,
3156 g.rope_base,
3157 1.0,
3158 ff,
3159 )?;
3160
3161 // Append at the HOST slot, then re-stamp the device counter: the eager chain has the
3162 // length on the host anyway, and the windowed view below needs it there to compute the
3163 // offset. The device counter is kept in lockstep so `mtp_kv_fill`'s `set_i32_one` and any
3164 // dc-family consumer of this scratch still agree.
3165 let kv = &mut scratch.kv;
3166 assert!(
3167 kv.len < scratch.cap,
3168 "step35 MTP scratch overflow ({} >= {})",
3169 kv.len,
3170 scratch.cap
3171 );
3172 let next_len = kv.len + 1;
3173 let (off, t_kv) = if g.swa && next_len > g.window {
3174 (next_len - g.window, g.window)
3175 } else {
3176 (0, next_len)
3177 };
3178 let write_row = e.prepare_kv_append(kv, off & !31usize, 1)?;
3179 e.append_kv_quantized(
3180 &k,
3181 &v0,
3182 &mut kv.k,
3183 &mut kv.v,
3184 write_row,
3185 kv.kv_dim_k,
3186 kv.kv_dim_v,
3187 kv.k_tok_bytes,
3188 kv.v_tok_bytes,
3189 false,
3190 )?;
3191 kv.len = next_len;
3192 e.set_i32_one(&mut kv.len_d, kv.len as i32)?;
3193 // SWA view offset (see note 1). The draft chain is short (k+2 rows), but the scratch is
3194 // PERSISTENT across rounds — `mtp_kv_fill` leaves one row per committed token behind, so
3195 // `kv.len` tracks absolute position and crosses 512 in any real generation. The window is
3196 // therefore live, not theoretical.
3197 let physical = kv.physical_rows(off, off + t_kv)?;
3198 let k_view = e.view_u8_range(
3199 &kv.k,
3200 physical.start * kv.k_tok_bytes,
3201 physical.end * kv.k_tok_bytes,
3202 );
3203 let v_view = e.view_u8_range(
3204 &kv.v,
3205 physical.start * kv.v_tok_bytes,
3206 physical.end * kv.v_tok_bytes,
3207 );
3208 let mut attn = e.uninit(nh * hd)?;
3209 e.fa_decode_kvmod(
3210 &q,
3211 &k_view,
3212 &v_view,
3213 &mut attn,
3214 hd,
3215 nh,
3216 nkv,
3217 t_kv,
3218 scale,
3219 kv.k_tok_bytes,
3220 kv.v_tok_bytes,
3221 false,
3222 )?;
3223
3224 let mut ag = e.uninit(nh * hd)?;
3225 e.attn_head_gate(&attn, >, &mut ag, None, hd, nh, 1)?;
3226 Ok(e.matmul(&fa.wo, &ag, 1)?)
3227 }
3228
3229 /// MTP-block full attention, T=1, on the scratch KV (BOTH draft paths — eager and graph):
3230 /// the scratch write slot and the attention bound come from `scratch.kv.len_d` (device i32[1])
3231 /// so the launch args are FIXED across draft steps — ONE captured graph serves the whole
3232 /// chain, and replays keep seeing KV growth through the device counter (no recapture).
3233 /// Geometry contract: n_splits is sized from `scratch.cap` (the persistent capacity); splits
3234 /// whose key range lies beyond the device t_kv exit empty and the shared combine skips them
3235 /// (fa_decode_dc bit-correct-for-any-t_kv<=bucket_max contract). The eager path uses the SAME
3236 /// launcher with the SAME bucket_max -> identical dispatch -> bit-identical draft tokens (the
3237 /// graph-vs-eager parity gate). Host len is NOT advanced here (graph contract); callers mirror.
3238 fn mtp_full_attn_dc(
3239 &self,
3240 e: &Engine,
3241 fa: &FullAttnLayer,
3242 h: &CudaSlice<f32>,
3243 pos_d: &CudaSlice<i32>,
3244 scratch: &mut MtpScratch,
3245 geom: Option<&crate::hybrid::DraftGeom>,
3246 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3247 let cfg = &self.cfg;
3248 let mtp_il = cfg.n_layer.saturating_sub(cfg.nextn_predict_layers);
3249 let geometry = cfg.full_attention_geometry_at(mtp_il);
3250 let n_head = geom.map(|g| g.n_head).unwrap_or(geometry.n_head as usize);
3251 let n_head_kv = geom
3252 .map(|g| g.n_head_kv)
3253 .unwrap_or(geometry.n_head_kv as usize);
3254 let head_dim = geometry.head_dim_k as usize;
3255 let eps = cfg.rms_eps;
3256 let scale = geometry.attention_scale();
3257 let n_embd = geom.map(|g| g.d_inner).unwrap_or(cfg.n_embd as usize);
3258 let bucket_max = scratch.cap; // < 96 guaranteed by the graph_draft eligibility gate
3259
3260 let (qf, mut k, v) =
3261 if e.uses_q8_1_fast(&fa.wq) && e.uses_q8_1_fast(&fa.wk) && e.uses_q8_1_fast(&fa.wv) {
3262 let (hq, hd) = e.quantize_q8_1(h, 1, n_embd)?;
3263 (
3264 e.matmul_pre(&fa.wq, &hq, &hd, h, 1)?,
3265 e.matmul_pre(&fa.wk, &hq, &hd, h, 1)?,
3266 e.matmul_pre(&fa.wv, &hq, &hd, h, 1)?,
3267 )
3268 } else {
3269 (
3270 e.matmul(&fa.wq, h, 1)?,
3271 e.matmul(&fa.wk, h, 1)?,
3272 e.matmul(&fa.wv, h, 1)?,
3273 )
3274 };
3275 // M3/Hy3 have no attention output gate — wq out is exactly q; skip the split.
3276 let gated = geometry.attention_gate == memra_gguf::config::AttentionGateKind::FusedQ;
3277 let (mut q, gate) = if gated {
3278 let mut q = e.zeros(n_head * head_dim)?;
3279 let mut gate = e.zeros(n_head * head_dim)?;
3280 e.q_gate_split(&qf, &mut q, &mut gate, head_dim, n_head, 1)?;
3281 (q, Some(gate))
3282 } else {
3283 (qf, None)
3284 };
3285
3286 let mut qn = e.zeros(n_head * head_dim)?;
3287 e.rms_norm(&q, fa.q_norm.float_data(), &mut qn, head_dim, n_head, eps)?;
3288 q = qn;
3289 let mut kn = e.zeros(n_head_kv * head_dim)?;
3290 e.rms_norm(
3291 &k,
3292 fa.k_norm.float_data(),
3293 &mut kn,
3294 head_dim,
3295 n_head_kv,
3296 eps,
3297 )?;
3298 k = kn;
3299 let rope_dims = geometry.n_rot as usize;
3300 e.rope_neox(
3301 &mut q,
3302 pos_d,
3303 head_dim,
3304 rope_dims,
3305 n_head,
3306 1,
3307 geometry.rope_base,
3308 1.0,
3309 )?;
3310 e.rope_neox(
3311 &mut k,
3312 pos_d,
3313 head_dim,
3314 rope_dims,
3315 n_head_kv,
3316 1,
3317 geometry.rope_base,
3318 1.0,
3319 )?;
3320
3321 let kv = &mut scratch.kv;
3322 // append at the DEVICE slot (kv.len_d == old len), then advance the counter in-graph.
3323 e.append_kv_quantized_dc(
3324 &k,
3325 &v,
3326 &mut kv.k,
3327 &mut kv.v,
3328 &kv.len_d,
3329 kv.kv_dim_k,
3330 kv.kv_dim_v,
3331 kv.k_tok_bytes,
3332 kv.v_tok_bytes,
3333 false,
3334 )?;
3335 e.inc_seqlen(&mut kv.len_d)?;
3336 // full-buffer views (any in-round t_kv stays in range on replay); the kernel bounds the
3337 // key range from the device counter.
3338 let k_view = e.view_u8(&kv.k, kv.k.len());
3339 let v_view = e.view_u8(&kv.v, kv.v.len());
3340 let (ktb, vtb) = (kv.k_tok_bytes, kv.v_tok_bytes);
3341 let mut attn = e.zeros(n_head * head_dim)?;
3342 e.fa_decode_dc(
3343 &q, &k_view, &v_view, &mut attn, head_dim, n_head, n_head_kv, &kv.len_d, bucket_max,
3344 scale, ktb, vtb, false,
3345 )?;
3346
3347 let attn_g = match &gate {
3348 Some(gate) => {
3349 let mut gsig = e.zeros(n_head * head_dim)?;
3350 e.sigmoid(gate, &mut gsig, n_head * head_dim)?;
3351 let mut ag = e.zeros(n_head * head_dim)?;
3352 e.mul(&attn, &gsig, &mut ag, n_head * head_dim)?;
3353 ag
3354 }
3355 None => attn,
3356 };
3357 Ok(e.matmul(&fa.wo, &attn_g, 1)?)
3358 }
3359
3360 /// PERSISTENT-DRAFT-KV fill (the reference engine's "mtp_update" analogue): compute the MTP
3361 /// block's K/V for `tokens` (committed tokens at positions pos0..pos0+T) from their EXACT
3362 /// trunk hiddens `h` ([T, n_embd] token-major, pre-output_norm) and append at slots pos0.. of
3363 /// the scratch KV. K/V-ONLY — ops A/1-5 plus the K-side of op 6 (wk/wv + k_norm + rope +
3364 /// quantized append); no wq/attention/FFN/lm_head, so per-token cost ~= eh_proj + wk/wv (a
3365 /// small fraction of one trunk layer), T-batched. Rope follows the chain convention
3366 /// rope(token@p) = p+1. Runs at round boundaries OUTSIDE the captured graph in BOTH draft
3367 /// modes -> draft parity by construction. Caller must have scratch.kv.len == pos0.
3368 #[allow(clippy::too_many_arguments)]
3369 fn mtp_kv_fill(
3370 &self,
3371 e: &Engine,
3372 mtp: &MtpHead,
3373 tokens: &[u32],
3374 h: &CudaSlice<f32>,
3375 pos0: usize,
3376 scratch: &mut MtpScratch,
3377 embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
3378 ) -> Result<(), Box<dyn std::error::Error>> {
3379 let cfg = &self.cfg;
3380 let n_embd = cfg.n_embd as usize;
3381 let eps = cfg.rms_eps;
3382 let t = tokens.len();
3383 assert_eq!(scratch.kv.len, pos0, "mtp_kv_fill: append slot mismatch");
3384 assert!(pos0 + t <= scratch.cap, "mtp_kv_fill: scratch overflow");
3385 let Mixer::Full(fa) = &mtp.mixer else {
3386 panic!("MTP block is full-attn in qwen35; linear MTP not supported")
3387 };
3388 let pos_vec: Vec<i32> = (0..t).map(|i| (pos0 + i + 1) as i32).collect();
3389 let pos_d = e.htod_i32(&pos_vec)?;
3390
3391 // ops A/1/2: embed + the two input norms, T-wide.
3392 let e_emb = match embd_dev {
3393 Some((g, qt, rb)) => e.embed_gather_device_t(g, tokens, n_embd, qt, rb)?,
3394 None => e.htod(&self.embd.gather(n_embd, tokens))?,
3395 };
3396 let mut e_norm = e.zeros(t * n_embd)?;
3397 e.rms_norm(&e_emb, mtp.enorm.float_data(), &mut e_norm, n_embd, t, eps)?;
3398 let mut h_norm = e.zeros(t * n_embd)?;
3399 e.rms_norm(h, mtp.hnorm.float_data(), &mut h_norm, n_embd, t, eps)?;
3400
3401 // op 3: per-row [e_norm ; h_norm] concat, token-major [T, 2*n_embd].
3402 let mut concat = e.zeros(t * 2 * n_embd)?;
3403 for i in 0..t {
3404 e.copy_view_into(
3405 &mut concat,
3406 i * 2 * n_embd,
3407 &e_norm.slice(i * n_embd..(i + 1) * n_embd),
3408 n_embd,
3409 )?;
3410 e.copy_view_into(
3411 &mut concat,
3412 i * 2 * n_embd + n_embd,
3413 &h_norm.slice(i * n_embd..(i + 1) * n_embd),
3414 n_embd,
3415 )?;
3416 }
3417
3418 // ops 4/5: eh_proj + attn_norm, T-wide (at the student inner width when geom is set).
3419 let di = mtp.geom.as_ref().map(|g| g.d_inner).unwrap_or(n_embd);
3420 let inp_sa = e.matmul(&mtp.eh_proj, &concat, t)?;
3421 let mut a_norm = e.zeros(t * di)?;
3422 e.rms_norm(&inp_sa, mtp.attn_norm.float_data(), &mut a_norm, di, t, eps)?;
3423
3424 // op 6 (K/V half): wk/wv + k_norm + rope + per-row quantized append. No wq/attention —
3425 // the fill only has to leave correct K/V rows behind for later chains to attend over.
3426 let n_head_kv = mtp
3427 .geom
3428 .as_ref()
3429 .map(|g| g.n_head_kv)
3430 .or(mtp.step35.as_ref().map(|s| s.n_head_kv))
3431 .unwrap_or_else(|| {
3432 let mtp_il = cfg.n_layer.saturating_sub(cfg.nextn_predict_layers);
3433 cfg.full_attention_geometry_at(mtp_il).n_head_kv as usize
3434 });
3435 let mtp_il = cfg.n_layer.saturating_sub(cfg.nextn_predict_layers);
3436 let geometry = cfg.full_attention_geometry_at(mtp_il);
3437 let head_dim = geometry.head_dim_k as usize;
3438 let mut k = e.matmul(&fa.wk, &a_norm, t)?;
3439 let v = e.matmul(&fa.wv, &a_norm, t)?;
3440 let mut kn = e.zeros(t * n_head_kv * head_dim)?;
3441 e.rms_norm(
3442 &k,
3443 fa.k_norm.float_data(),
3444 &mut kn,
3445 head_dim,
3446 n_head_kv * t,
3447 eps,
3448 )?;
3449 k = kn;
3450 // step35: rotary width AND base are per-layer, and the MTP block's values come from the
3451 // resolved `Step35MtpGeom` — NOT from `cfg.rope_dim_count`/`cfg.rope_freq_base`, which
3452 // carry the arch defaults (128 / 5e6, i.e. the FULL-attn layers' base). Getting this wrong
3453 // writes K rows the attention arm then re-derives at a different theta: correct-looking
3454 // output with dead acceptance, invisible to the exactness gates.
3455 let (rope_dims, rope_base, ff) = match mtp.step35.as_ref() {
3456 Some(s) => (
3457 s.n_rot,
3458 s.rope_base,
3459 if s.swa {
3460 None
3461 } else {
3462 self.step35_aux.as_ref().and_then(|a| a.rope_freqs(e))
3463 },
3464 ),
3465 None => (geometry.n_rot as usize, geometry.rope_base, None),
3466 };
3467 #[cfg(debug_assertions)]
3468 if let Some(ff) = ff {
3469 crate::debug_assert_tensor_stream_device(ff, &e.stream(), "mtp_kv_fill.rope_freqs");
3470 }
3471 match ff {
3472 Some(f) => e.rope_neox_ff(
3473 &mut k, &pos_d, head_dim, rope_dims, n_head_kv, t, rope_base, 1.0, f,
3474 )?,
3475 None => e.rope_neox(
3476 &mut k, &pos_d, head_dim, rope_dims, n_head_kv, t, rope_base, 1.0,
3477 )?,
3478 }
3479
3480 let kv = &mut scratch.kv;
3481 // Match the trunk prime contract: a chunk may need the aligned window immediately before
3482 // its first row, so preserve that prefix when the physical tail rebases at wrap.
3483 let retain_from = kv
3484 .ring
3485 .as_ref()
3486 .map(|ring| pos0.saturating_sub(ring.window() - 1) & !31usize)
3487 .unwrap_or(0);
3488 let write_row = e.prepare_kv_append(kv, retain_from, t)?;
3489 for i in 0..t {
3490 let k_row = k.slice(i * kv.kv_dim_k..(i + 1) * kv.kv_dim_k);
3491 let v_row = v.slice(i * kv.kv_dim_v..(i + 1) * kv.kv_dim_v);
3492 e.append_kv_quantized_view(
3493 &k_row,
3494 &v_row,
3495 &mut kv.k,
3496 &mut kv.v,
3497 write_row + i,
3498 kv.kv_dim_k,
3499 kv.kv_dim_v,
3500 kv.k_tok_bytes,
3501 kv.v_tok_bytes,
3502 false,
3503 )?;
3504 }
3505 kv.len = pos0 + t;
3506 e.set_i32_one(&mut kv.len_d, kv.len as i32)?;
3507 Ok(())
3508 }
3509
3510 /// CAPTURE body for the GRAPH DRAFT (stage 2 of graph-grade spec): ONE MTP head forward with
3511 /// every varying input device-resident —
3512 /// - token id from the persistent `tok_d` (the previous replay's in-graph argmax wrote it,
3513 /// so the chain feeds itself; the host reads the same 4 bytes for the draft list),
3514 /// - h_seed from the persistent `h_seed_d` (h_nextn is copied BACK into it at the end),
3515 /// - rope pos from the persistent `pos_d` counter (inc'd in-graph),
3516 /// - scratch KV slot/bound from `scratch.kv.len_d` (see mtp_full_attn_dc).
3517 /// The p-min confidence lands in the persistent `p_d` iff `with_prob` (env is fixed per run).
3518 /// Same kernels, same dispatch as the eager mtp_head_forward_dev chain -> same draft tokens
3519 /// (exactness never depends on drafts — the verify arbitrates — but acceptance parity does).
3520 /// `with_head=false` captures the HEAD-LESS twin for the pseudo-seed replay (2026-07-03):
3521 /// the pseudo pass only needs h_nextn (op 10) + the scratch append — the lm_head read
3522 /// (~1.06ms q6_K on the 9B), argmax and prob are dead weight there. h_nextn's inputs are
3523 /// untouched, so the seed value is identical; round-start resets overwrite tok_d/p_d anyway.
3524 /// `sampled_cap` = Some((ctr_d, perturb_d, q_out_d, seed, temp)) captures the SAMPLED twin
3525 /// (step 3 of the sampled-spec arc): head logits are retained in the persistent `q_out_d`
3526 /// (host D2Ds them to the round's q slot after each replay), the DEVICE event counter is
3527 /// bumped in-graph, and the argmax reads GUMBEL-PERTURBED logits — one categorical draw per
3528 /// replay, bit-identical to the eager arm's gumbel_perturb at the same (seed, sctr, temp).
3529 /// seed/temp are capture-time constants (fixed per generate call, like p_min).
3530 #[allow(clippy::too_many_arguments)]
3531 fn mtp_head_forward_cap(
3532 &self,
3533 e: &Engine,
3534 mtp: &MtpHead,
3535 tok_d: &mut CudaSlice<u32>,
3536 pos_d: &mut CudaSlice<i32>,
3537 h_seed_d: &mut CudaSlice<f32>,
3538 p_d: &mut CudaSlice<f32>,
3539 scratch: &mut MtpScratch,
3540 with_prob: bool,
3541 with_head: bool,
3542 embd_gpu: &CudaSlice<u8>,
3543 embd_qt: i32,
3544 embd_rb: usize,
3545 d_vocab: usize,
3546 sampled_cap: Option<(
3547 &mut CudaSlice<u32>,
3548 &mut CudaSlice<f32>,
3549 &mut CudaSlice<f32>,
3550 u64,
3551 f32,
3552 )>,
3553 stream_pack: Option<(&mut CudaSlice<u32>, usize, Option<&CudaSlice<u32>>)>,
3554 // DRAFT-SIDE GRAMMAR MASK (lane/draft-mask): (packed draft-vocab allowed-set buffer,
3555 // word count). Captured as ONE mask_logits_f32 node between the head matmul and the
3556 // in-graph argmax; the buffer address is baked, its CONTENTS are re-uploaded by the
3557 // host before every replay (the decode.rs graph-mask pattern). All-ones contents = a
3558 // no-op ban, so a position the grammar cannot constrain costs one pass over the row.
3559 mask_cap: Option<(&CudaSlice<u32>, usize)>,
3560 ) -> Result<(), Box<dyn std::error::Error>> {
3561 let cfg = &self.cfg;
3562 let n_embd = cfg.n_embd as usize;
3563 // step35 REFUSAL (deliberate, named): the graph body's attention is `mtp_full_attn_dc`,
3564 // whose device-counter key bound always starts at row 0 — it cannot express this block's
3565 // SWA view offset, so a captured chain would silently attend OUTSIDE the window once the
3566 // persistent scratch passes 512 rows. Nothing is lost today: `mtp_head_forward_cap`
3567 // refuses step35 heads explicitly (SWA refusal), so the eager chain
3568 // (`mtp_head_forward_dev` -> `mtp_step35_attn`) is the served path. Returning Err (not a
3569 // panic) is what the two capture sites and the round-stream capture already handle by
3570 // degrading to eager / stream-off.
3571 if mtp.step35.is_some() {
3572 return Err(
3573 "step35 has no captured draft chain (fa_decode_dc cannot express the MTP \
3574 block's SWA view offset; same root cause as the dc decode refusal) — the \
3575 eager draft chain serves this arch"
3576 .into(),
3577 );
3578 }
3579 // student inner width (see mtp_head_forward_dev) — interface dims stay n_embd.
3580 let di = mtp.geom.as_ref().map(|g| g.d_inner).unwrap_or(n_embd);
3581 let eps = cfg.rms_eps;
3582 let e_emb = e.embed_gather_device(embd_gpu, tok_d, n_embd, embd_qt, embd_rb)?;
3583 let mut e_norm = e.zeros(n_embd)?;
3584 e.rms_norm(&e_emb, mtp.enorm.float_data(), &mut e_norm, n_embd, 1, eps)?;
3585 let mut h_norm = e.zeros(n_embd)?;
3586 e.rms_norm(
3587 &*h_seed_d,
3588 mtp.hnorm.float_data(),
3589 &mut h_norm,
3590 n_embd,
3591 1,
3592 eps,
3593 )?;
3594 let mut concat = e.zeros(2 * n_embd)?;
3595 e.copy_into(&mut concat, 0, &e_norm, n_embd)?;
3596 e.copy_into(&mut concat, n_embd, &h_norm, n_embd)?;
3597 let inp_sa = e.matmul(&mtp.eh_proj, &concat, 1)?;
3598 let mut a_norm = e.zeros(di)?;
3599 e.rms_norm(&inp_sa, mtp.attn_norm.float_data(), &mut a_norm, di, 1, eps)?;
3600 let attn_out = match &mtp.mixer {
3601 Mixer::Full(fa) => {
3602 self.mtp_full_attn_dc(e, fa, &a_norm, pos_d, scratch, mtp.geom.as_ref())?
3603 }
3604 Mixer::Linear(_) => {
3605 panic!("MTP block is full-attn in qwen35; linear MTP not supported")
3606 }
3607 Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
3608 };
3609 let mut x1 = e.zeros(di)?;
3610 e.add(&inp_sa, &attn_out, &mut x1, di)?;
3611 let mut z = e.zeros(di)?;
3612 e.rms_norm(&x1, mtp.post_attn_norm.float_data(), &mut z, di, 1, eps)?;
3613 let ffn_out = match &mtp.ffn {
3614 crate::hybrid::Ffn::Dense {
3615 ffn_gate,
3616 ffn_up,
3617 ffn_down,
3618 } => {
3619 let n_ff = ffn_gate.out_features();
3620 let (gate, up) = if e.uses_q8_1_fast(ffn_gate) && e.uses_q8_1_fast(ffn_up) {
3621 let (zq, zd) = e.quantize_q8_1(&z, 1, di)?;
3622 (
3623 e.matmul_pre(ffn_gate, &zq, &zd, &z, 1)?,
3624 e.matmul_pre(ffn_up, &zq, &zd, &z, 1)?,
3625 )
3626 } else {
3627 (e.matmul(ffn_gate, &z, 1)?, e.matmul(ffn_up, &z, 1)?)
3628 };
3629 let mut act = e.zeros(n_ff)?;
3630 Self::ffn_act(e, &self.cfg, &gate, &up, &mut act, n_ff)?;
3631 e.matmul(ffn_down, &act, 1)?
3632 }
3633 // ROUND-STREAM: the 35B NextN block carries a MoE FFN. With RESIDENT experts the
3634 // dev path is pure device launches (device top-k + rows kernels, ZERO-DtoH by
3635 // design) — capture-legal. Non-resident (SLRU-lock) stays rejected: the capture
3636 // error arm degrades the caller to eager/stream-off.
3637 crate::hybrid::Ffn::Moe(m) if m.dev_exps.is_some() => {
3638 self.moe_ffn_il(e, m, &z, 1, u16::MAX)?
3639 }
3640 crate::hybrid::Ffn::Moe(_) => {
3641 return Err("graph draft requires a Dense (or resident-MoE) MTP FFN".into());
3642 }
3643 };
3644 let mut h_inner = e.zeros(di)?;
3645 e.add(&x1, &ffn_out, &mut h_inner, di)?;
3646 // student: up-project back to n_embd (carrier + head input; see mtp_head_forward_dev).
3647 let h_nextn = match mtp.geom.as_ref() {
3648 Some(g) => e.matmul(&g.out_up, &h_inner, 1)?,
3649 None => h_inner,
3650 };
3651 // MEMRA_SPEC_HPOST needs final_h even head-less (it IS the next seed under that convention).
3652 let final_h = if with_head || spec_hpost() {
3653 let final_norm = mtp.shared_head_norm.as_ref().unwrap_or(&self.output_norm);
3654 let mut fh = e.zeros(n_embd)?;
3655 e.rms_norm(&h_nextn, final_norm.float_data(), &mut fh, n_embd, 1, eps)?;
3656 Some(fh)
3657 } else {
3658 None
3659 };
3660 if with_head {
3661 let head = mtp.shared_head_head.as_ref().unwrap_or(&self.output);
3662 let mut logits = e.matmul(head, final_h.as_ref().unwrap(), 1)?;
3663 // DRAFT-SIDE GRAMMAR MASK: ban the grammar-illegal draft ids IN the captured chain,
3664 // before the argmax — proposals become legal by construction. Contents-only
3665 // per-replay upload keeps the capture valid.
3666 if let Some((mask_d, mw)) = mask_cap {
3667 e.mask_logits_col(&mut logits, mask_d, 0, d_vocab, mw)?;
3668 }
3669 if let Some((ctr_d, perturb_d, q_out_d, seed, temp)) = sampled_cap {
3670 // SAMPLED chain: retain q (raw head logits -> persistent q_out_d; the matmul's
3671 // own buffer is pool-recycled after the capture body returns, so it can't be the
3672 // retention target), bump the device event counter, gumbel-perturb reading it,
3673 // and argmax the PERTURBED logits into tok_d — the in-graph categorical draw.
3674 e.copy_into(q_out_d, 0, &logits, d_vocab)?;
3675 e.sctr_inc(ctr_d)?;
3676 e.gumbel_perturb_ctr(&logits, perturb_d, d_vocab, seed, ctr_d, temp)?;
3677 e.argmax_token_device_into(perturb_d, tok_d, d_vocab)?;
3678 // p-min prob = the head's RAW softmax confidence in the SAMPLED pick — same
3679 // semantics as the eager sampled arm's prob_of_token_device(dl_d, tok_d).
3680 if with_prob {
3681 e.prob_of_token_device_into(&logits, tok_d, p_d, d_vocab)?;
3682 }
3683 } else {
3684 // draft token -> persistent tok_d (next replay's embed reads it; host reads the 4 bytes).
3685 e.argmax_token_device_into(&logits, tok_d, d_vocab)?;
3686 // p-min under a draft mask reads the MASKED row: confidence relative to the
3687 // grammar-LEGAL alternatives (illegal ids leave the softmax denominator), which
3688 // is the right semantics for "does the drafter know what comes next here" and
3689 // the same row the pick came from. Draft-quality only — verify arbitrates.
3690 if with_prob {
3691 e.prob_of_token_device_into(&logits, tok_d, p_d, d_vocab)?;
3692 }
3693 }
3694 }
3695 // ROUND-STREAM K-chain: pack (tok, p) into slot j, then remap tok through d2t so the
3696 // NEXT chained body's embed reads the TARGET id — zero host involvement per step.
3697 if let Some((out, slot, d2t)) = stream_pack {
3698 e.pack_tok_p(tok_d, p_d, out, slot)?;
3699 if let Some(map) = d2t {
3700 e.tok_map_u32(tok_d, map)?;
3701 }
3702 }
3703 // Next draft step's h_seed: pre-norm h_nextn (default) or post-norm final_h (HPOST).
3704 if spec_hpost() {
3705 e.copy_into(h_seed_d, 0, final_h.as_ref().unwrap(), n_embd)?;
3706 } else {
3707 e.copy_into(h_seed_d, 0, &h_nextn, n_embd)?;
3708 }
3709 // advance the draft rope position in-graph.
3710 e.inc_seqlen(pos_d)?;
3711 Ok(())
3712 }
3713
3714 /// Batched target verify forward over `tokens` at positions `pos0..pos0+T` (§D.3, T=K+1).
3715 /// Returns ALL T logit columns (host f32, [T*n_vocab]); appends T cols to every full-attn KV
3716 /// and advances every linear-attn recur state by T steps (the recur steps are SEQUENTIAL T=1).
3717 /// Advances `cache.pos` by T.
3718 pub fn decode_step_t(
3719 &self,
3720 e: &Engine,
3721 tokens: &[u32],
3722 pos0: usize,
3723 cache: &mut Cache,
3724 ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
3725 if self.is_gemma4_e4b() {
3726 return Ok(self.gemma4_e4b_decode_step_t_h(e, tokens, pos0, cache)?.0);
3727 }
3728 if self.cfg.gemma4.is_some() {
3729 return self.gemma4_decode_step_t(e, tokens, pos0, cache);
3730 }
3731 Ok(self.decode_step_t_h(e, tokens, pos0, cache)?.0)
3732 }
3733
3734 /// Like `decode_step_t` but ALSO returns the LAST column's pre-output_norm hidden (h_seed for
3735 /// the next draft round). This lets partial-accept replay run as ONE batched T=(n_acc+1) forward
3736 /// (single weight read) instead of n_acc+1 separate T=1 decode_steps (n_acc+1 weight reads).
3737 /// At batch=1 decode is bandwidth-bound, so batching the replay is THE MTP profitability lever.
3738 pub fn decode_step_t_h(
3739 &self,
3740 e: &Engine,
3741 tokens: &[u32],
3742 pos0: usize,
3743 cache: &mut Cache,
3744 ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
3745 self.decode_step_t_h_emb(e, tokens, pos0, cache, None)
3746 }
3747
3748 /// Like `decode_step_t_h` with an optional RESIDENT embed table (spec hot loop): device
3749 /// gather instead of host dequant + [T, n_embd] f32 htod. Bit-identical rows.
3750 pub fn decode_step_t_h_emb(
3751 &self,
3752 e: &Engine,
3753 tokens: &[u32],
3754 pos0: usize,
3755 cache: &mut Cache,
3756 embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
3757 ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
3758 let (logits_d, h_seed) = self.decode_step_t_h_emb_dev(e, tokens, pos0, cache, embd_dev)?;
3759 Ok((e.dtoh(&logits_d)?, h_seed))
3760 }
3761
3762 /// DEVICE-LOGITS verify forward (spec device-argmax lever): identical kernel chain to
3763 /// `decode_step_t_h_emb` but returns the [T, n_vocab] logits ON DEVICE — the accept walk
3764 /// argmaxes each column on-device and reads back ONE [T] u32 instead of dtoh'ing the full
3765 /// T x n_vocab f32 block (~1-4 MB + T host argmaxes, every round). Kernel dispatch is
3766 /// UNCHANGED (same decode-exact kernels); only the post-logits transfer moves.
3767 pub fn decode_step_t_h_emb_dev(
3768 &self,
3769 e: &Engine,
3770 tokens: &[u32],
3771 pos0: usize,
3772 cache: &mut Cache,
3773 embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
3774 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
3775 let n_embd = self.cfg.n_embd as usize;
3776 let t = tokens.len();
3777 let (logits, x) = self.decode_step_t_core(e, tokens, pos0, cache, embd_dev, None)?;
3778 // h_seed for the next round = LAST column's pre-output_norm hidden ([n_embd]).
3779 let mut hs = vbuf(e, n_embd)?; // fully written by copy_view_into below
3780 e.copy_view_into(&mut hs, 0, &x.slice((t - 1) * n_embd..t * n_embd), n_embd)?;
3781 Ok((logits, hs))
3782 }
3783
3784 /// CORE verify forward: the `decode_step_t_h_emb_dev` kernel chain, returning the FULL
3785 /// pre-output_norm hidden stack x ([T, n_embd], any column extractable) and optionally
3786 /// filling a `VerifyCkpt` (retained per-layer state-rebuild inputs) for the REPLAY-FREE
3787 /// partial accept. `ckpt: None` => byte-for-byte the old behavior (the ckpt writes are pure
3788 /// retains/copies — they never change what any kernel computes).
3789 fn decode_step_t_core(
3790 &self,
3791 e: &Engine,
3792 tokens: &[u32],
3793 pos0: usize,
3794 cache: &mut Cache,
3795 embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
3796 mut ckpt: Option<&mut VerifyCkpt>,
3797 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
3798 self.decode_step_t_core_stream(
3799 e,
3800 tokens,
3801 pos0,
3802 cache,
3803 embd_dev,
3804 ckpt.take(),
3805 None,
3806 None,
3807 None,
3808 None,
3809 )
3810 }
3811
3812 /// Increment-0 two-session PP seam: release the peer after this lane's stage-0 boundary TX.
3813 /// The two independent sessions keep their own cache/checkpoint state; only issue order moves.
3814 fn decode_step_t_core_pipelined(
3815 &self,
3816 e: &Engine,
3817 tokens: &[u32],
3818 pos0: usize,
3819 cache: &mut Cache,
3820 embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
3821 mut ckpt: Option<&mut VerifyCkpt>,
3822 pipe: &SpecPipeLane,
3823 round: usize,
3824 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
3825 let fence = crate::pp::pp_cuts(self.layers.len())
3826 .ok_or("two-session speculative pipeline requires a PP stage cut")?;
3827 if crate::pp::pp2_streams_off() || !crate::pp::spec_pp_on() {
3828 return Err("two-session speculative pipeline requires the PP verify split".into());
3829 }
3830 let interval_fence = pipe.stage0_begin(round)?;
3831 let ticket = self.verify_stage0_issue(
3832 e,
3833 tokens,
3834 pos0,
3835 cache,
3836 embd_dev,
3837 ckpt.as_deref_mut(),
3838 None,
3839 &fence,
3840 Some(interval_fence),
3841 pipe.trace(round),
3842 )?;
3843 pipe.stage0_end(round);
3844 pipe.stage1_begin(round)?;
3845 let result = self.verify_stage1_finish(e, ticket, cache, ckpt, None, &fence, true)?;
3846 pipe.verify_end(round);
3847 Ok(result)
3848 }
3849
3850 /// ROUND-STREAM stage (c) 4: `stream` = (device verify tokens [t], device pos counter) —
3851 /// when Some, rope positions come from pos_iota over the counter, the embed gathers the
3852 /// device tokens, and full_attn_verify routes appends/FA through the _dc twins reading the
3853 /// SAME counter (every layer's kvl.len == cache.pos, one counter drives all three). The
3854 /// host `tokens`/`pos0` args still size buffers (t is FIXED K+1 in stream mode).
3855 /// `vtok_dev` (engine-bundle slice 2): device verify tokens for the EMBED only —
3856 /// unlike `stream` mode it changes nothing else (host pos iota, host-len KV appends).
3857 /// `tokens` then only sizes buffers (the dummy-slice pattern the round-stream arm uses).
3858 #[allow(clippy::too_many_arguments)]
3859 fn decode_step_t_core_stream(
3860 &self,
3861 e: &Engine,
3862 tokens: &[u32],
3863 pos0: usize,
3864 cache: &mut Cache,
3865 embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
3866 mut ckpt: Option<&mut VerifyCkpt>,
3867 stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
3868 pp_pipe: Option<bool>,
3869 vtok_dev: Option<&CudaSlice<u32>>,
3870 graphs: Option<&mut DsparkVerifyGraphs>,
3871 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
3872 // PP DOOR (lane/pp2-spec 2026-08-06): the verify trunk now takes its OWN stage split,
3873 // exactly as the eager and batched steps do. This is the single funnel every verify
3874 // forward reaches (decode_step_t / _h / _h_emb / _h_emb_dev / _core all land here), so
3875 // wiring it here wires the whole spec surface — the draft/accept/commit machinery above
3876 // is untouched.
3877 //
3878 // History: pp2-hardening (2026-08-06) made this funnel FAIL CLOSED, because its trunk
3879 // walk was unsplit on one stream and a sharded cross-device placement peer-read every
3880 // remote layer's weights on every spec round (measured 13.9-28x on the batched twin).
3881 // The refusal below survives to cover the residue — MEMRA_SPEC_PP=0, MEMRA_PP_STREAMS=0,
3882 // or a placement whose PpNRt fails to build — so a config that would still walk the
3883 // whole trunk on one stream refuses instead of regressing 28x.
3884 if let Some(fence) = crate::pp::pp_cuts(self.layers.len()) {
3885 if !crate::pp::pp2_streams_off() && crate::pp::spec_pp_on() {
3886 if vtok_dev.is_some() {
3887 return Err(
3888 "device-token dspark verify (slice-2 deferred readback) has no PP \
3889 stage-split arm; set MEMRA_DSPARK_DEFER_READBACK=0 or run the dspark \
3890 route on one device"
3891 .into(),
3892 );
3893 }
3894 return self.decode_step_t_core_ppn(
3895 e,
3896 tokens,
3897 pos0,
3898 cache,
3899 embd_dev,
3900 ckpt.take(),
3901 stream,
3902 &fence,
3903 pp_pipe,
3904 );
3905 }
3906 }
3907 crate::pp::refuse_unsplit_if_remote(
3908 "decode_step_t (spec verify)",
3909 "drop MEMRA_SPEC_PP=0 / MEMRA_PP_STREAMS=0 so the verify trunk takes its OWN stage \
3910 split (decode_step_t_core_ppn); or run spec on one device",
3911 )?;
3912 let cfg = &self.cfg;
3913 let n_embd = cfg.n_embd as usize;
3914 let eps = cfg.rms_eps;
3915 let t = tokens.len();
3916 let pos_d = match stream {
3917 Some((_, ctr)) => {
3918 let mut p = e.alloc_uninit::<i32>(t)?;
3919 e.pos_iota(ctr, &mut p, t)?;
3920 p
3921 }
3922 None => {
3923 let pos_vec: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
3924 e.htod_i32(&pos_vec)?
3925 }
3926 };
3927
3928 // embed T tokens -> [T, n_embd] token-major (device gather on the spec hot loop)
3929 let x = match (stream, embd_dev) {
3930 (Some((vtok, _)), Some((g, qt, rb))) => {
3931 e.embed_gather_device_td(g, vtok, t, n_embd, qt, rb)?
3932 }
3933 (None, Some((g, qt, rb))) => match vtok_dev {
3934 // slice 2: device verify tokens, same embed_gather_u32_t kernel —
3935 // bit-identical rows to the host-token arm (same per-dtype deq).
3936 Some(vt_d) => e.embed_gather_device_td(g, vt_d, t, n_embd, qt, rb)?,
3937 None => e.embed_gather_device_t(g, tokens, n_embd, qt, rb)?,
3938 },
3939 _ => {
3940 assert!(
3941 vtok_dev.is_none(),
3942 "device-token verify requires the resident embed table (embd_dev)"
3943 );
3944 e.htod(&self.embd.gather(n_embd, tokens))?
3945 }
3946 };
3947
3948 // TRUNK WALK: layers [0, n_layers) through the SAME range-scoped subgraph the PP-N
3949 // stage split calls per stage (`verify_layers`) — one code path, so the split cannot
3950 // drift from the unsplit dispatch mirroring. lane/pp2-spec 2026-08-06.
3951 let x = self.verify_layers(
3952 e,
3953 x,
3954 0,
3955 self.layers.len(),
3956 &pos_d,
3957 pos0,
3958 t,
3959 cache,
3960 ckpt.take(),
3961 stream,
3962 graphs,
3963 )?;
3964
3965 let mut hn = vbuf(e, t * n_embd)?;
3966 let serving_head = self.cfg.step35.is_some() || self.qwen35_serving_class();
3967 let logits = if serving_head {
3968 // Step35 and the qwen35 family (MoE 2026-08-14 AM, dense-hybrid same day PM — the
3969 // Q3.8 bring-up reproduced the identical near-tie class on dense: eager-class verify
3970 // vs batched-class live serving, ULP drift amplified through the GDN recurrence)
3971 // serve one batched numeric class at every live width, including B=1. Keep the
3972 // verify head in that same class; other generic families retain the decode-exact
3973 // head that their run-spec contract pins.
3974 e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
3975 e.matmul(&self.output, &hn, t)?
3976 } else {
3977 e.rms_norm_decode(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
3978 e.matmul_decode_exact(&self.output, &hn, t)?
3979 };
3980 // stream: the device pos counter owns position; host mirror reconciles at drain.
3981 if stream.is_none() {
3982 cache.pos += t;
3983 }
3984 // Hidden stack for seeds/refresh-fills: pre-norm x (default) or post-norm hn (HPOST).
3985 Ok((logits, if spec_hpost() { hn } else { x }))
3986 }
3987
3988 /// THE VERIFY TRUNK OVER PP-N (lane/pp2-spec 2026-08-06): `decode_step_t_core_stream`'s walk
3989 /// as N stage subgraphs, each on its own engine/stream (and, under `MEMRA_PP_DEVICES`, its own
3990 /// device), with a `[T, n_embd]` boundary transfer between them. T = K+1 (the verify batch),
3991 /// so this is the batched-boundary shape the pp2-batch lane's grow-only slots already handle
3992 /// (`tx(b, x, t*n_embd)`; the slot grows to the high-water T and the transport moves exactly
3993 /// the payload).
3994 ///
3995 /// Structure is `decode_step_batch_ppn`'s, which is `decode_step_h_ppn`'s. FOUR THINGS ARE
3996 /// PER-STAGE and each for a measured reason (see `decode_step_batch_ppn`'s header for the
3997 /// receipts):
3998 ///
3999 /// 1. THE ENGINE (`rt.engine(s, e)`) — `Engine` owns lazily-grown stable-pointer scratch
4000 /// (`fa_part_pool`, `fa_vf16_scratch`, `argmax_partials`) that is single-stream-safe BY
4001 /// DESIGN. Two stage streams through one Engine is the 2026-08-02 shared-scratch race
4002 /// (35% flake, nondeterministic all-logits divergence). `PpNRt::build` gives every stage
4003 /// s>0 its own Engine even on the primary device; honouring it here is what scopes the
4004 /// pools. The verify path allocates MORE of that scratch than eager decode does (FA at
4005 /// m=T, and the per-layer `GdnStash` retains), so this is load-bearing, not inherited.
4006 ///
4007 /// 2. `pos_d` — each stage uploads its OWN copy of the T rope positions on ITS stream, so the
4008 /// buffer is allocated, consumed and freed on one stream. In `stream` mode that means each
4009 /// stage runs its own `pos_iota` over the SHARED device counter (`pos_ctr`): the counter is
4010 /// read-only during the forward (the round's `inc`/`copy_add` happen outside it), so every
4011 /// stage derives the identical iota, and each stage's own output buffer is stream-local.
4012 ///
4013 /// 3. THE EMBED lives with stage 0 (`self.embd` / `embd_gpu` are host/primary-side; the
4014 /// sharded loader leaves the table with stage 0 by construction).
4015 ///
4016 /// 4. THE HEAD (`output_norm` + `output`) runs on the LAST stage — the sharded loader uploaded
4017 /// both through that stage's engine (`hybrid.rs`: `e_head = layer_engine(e, n_trunk,
4018 /// n_trunk-1)`), so reading them anywhere else is a peer read of the biggest tensor in the
4019 /// model, every round.
4020 ///
4021 /// WHAT STAYS ON THE PRIMARY, deliberately: the returned logits and hidden stack `x`. Both are
4022 /// last-stage-allocated device buffers, and every consumer (the device argmax walk, the accept
4023 /// kernels, `spec_seed_gather`, the ckpt rebuild in `commit_verified_prefix`) reads them
4024 /// through the primary context by UVA — the same read the batched serving epilogue's
4025 /// `last_logits_dev` park does. Those consumers are per-round O(T x n_vocab) and O(n_embd),
4026 /// not per-layer, so they are not the 28x class; splitting them is a separate lane.
4027 ///
4028 /// The MTP HEAD (draft side) is NOT split: it is one block, it lives wherever the loader put
4029 /// it (`load_mtp` uses the primary engine), and it is ~1-2 GB against the trunk's tens. Draft
4030 /// placement is measured, not assumed — see `research/pp2-spec-20260806`.
4031 ///
4032 /// EXACTNESS: PP-N adds ZERO deviation. Each stage runs the SAME kernels on the SAME bytes in
4033 /// the same order via the SAME `verify_layers` the unsplit body calls; the only change is
4034 /// where the residual is materialized, and the boundary is a straight f32 copy (dtod
4035 /// same-device / `cudaMemcpyPeerAsync` cross-device, no conversion). So the split MUST be
4036 /// BIT-IDENTICAL to the unsplit verify at the same T, in both placement orders. Gate:
4037 /// `decode-batch-gate --mode ppspec`. Acceptance counts are a DERIVED consequence — greedy
4038 /// accept argmaxes these logits, so bit-identical logits force identical accept walks; the
4039 /// `run-spec` K=1..8 arm checks that end-to-end rather than trusting the implication.
4040 #[allow(clippy::too_many_arguments)]
4041 fn decode_step_t_core_ppn(
4042 &self,
4043 e: &Engine,
4044 tokens: &[u32],
4045 pos0: usize,
4046 cache: &mut Cache,
4047 embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
4048 mut ckpt: Option<&mut VerifyCkpt>,
4049 stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
4050 fence: &[usize],
4051 pp_pipe: Option<bool>,
4052 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4053 let ticket = self.verify_stage0_issue(
4054 e,
4055 tokens,
4056 pos0,
4057 cache,
4058 embd_dev,
4059 ckpt.as_deref_mut(),
4060 stream,
4061 fence,
4062 pp_pipe,
4063 None,
4064 )?;
4065 self.verify_stage1_finish(e, ticket, cache, ckpt, stream, fence, true)
4066 }
4067
4068 /// Enqueue embed, stage 0, and the first boundary TX, then return the actual boundary slot.
4069 /// The ordinary PP verify wrapper calls `verify_stage1_finish` immediately after this return.
4070 #[allow(clippy::too_many_arguments)]
4071 fn verify_stage0_issue(
4072 &self,
4073 e: &Engine,
4074 tokens: &[u32],
4075 pos0: usize,
4076 cache: &mut Cache,
4077 embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
4078 mut ckpt: Option<&mut VerifyCkpt>,
4079 stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
4080 fence: &[usize],
4081 pp_pipe: Option<bool>,
4082 trace: Option<SpecPipeTraceCtx>,
4083 ) -> Result<VerifyBoundaryTicket, Box<dyn std::error::Error>> {
4084 assert!(
4085 !self.is_gemma4_e4b() && self.cfg.gemma4.is_none(),
4086 "decode_step_t_core_ppn covers the hybrid non-gemma4 verify trunk only \
4087 (the gemma4 arms have their own decode_step_t twins)"
4088 );
4089 if crate::pp::pp_host_bounce_active() && (stream.is_some() || embd_dev.is_some()) {
4090 return Err(
4091 "decode_step_t_core_ppn: refused with MEMRA_PP_HOST_BOUNCE=1 — the trunk \
4092 boundary itself is host-staged, but device-resident verify still peer-reads \
4093 primary-device token/position/embedding buffers from stage 0. Run plain PP \
4094 serving on this host class; spec requires local per-stage inputs first."
4095 .into(),
4096 );
4097 }
4098 let rt = crate::pp::PpNRt::get(e)?;
4099 let n_st = fence.len() - 1;
4100 assert_eq!(
4101 rt.n_stages(),
4102 n_st,
4103 "PpNRt stage count {} != fence stages {n_st}",
4104 rt.n_stages()
4105 );
4106 let n_embd = self.cfg.n_embd as usize;
4107 let t = tokens.len();
4108 let payload = t * n_embd;
4109 if pp_pipe.is_some() {
4110 assert_eq!(n_st, 2, "spec pipeline requires exactly two PP stages");
4111 }
4112 // One-shot lane diagnostic: force natural PP-2 boundaries to completion so the server
4113 // log can price stage 0, the peer hop, the RX copy, and stage 1 + head separately without
4114 // nsys. The ordinary path keeps every enqueue asynchronous. N>2 is deliberately excluded:
4115 // the report below names exactly two stages and must never imply it measured middle ones.
4116 let pp_anatomy = n_st == 2 && std::env::var("MEMRA_SPEC_PP_ANATOMY").as_deref() == Ok("1");
4117 let pp_started = std::time::Instant::now();
4118 let (mut reverse_ms, mut stage0_ms, mut tx_ms) = (0.0f64, 0.0f64, 0.0f64);
4119 // The CALLER's ambient stream, captured BEFORE any `rt.enter()` pushes a stage stream:
4120 // this body returns DEVICE-RESIDENT buffers (the device-argmax accept walk's contract),
4121 // so the exit needs the same publication the boundaries get — see `PpNRt::publish_to`.
4122 // Taken here, not at the end, because inside the last-stage scope `e.stream()` IS the
4123 // stage stream and the wait would self-order into a no-op.
4124 let caller_stream = e.stream();
4125 // #87 ROOT-CAUSE FENCE (lane/pp2spec-crash): the PREVIOUS round's stage-allocated
4126 // outputs (logits/hidden/ckpt stashes) freed stream-ordered on the STAGE streams while
4127 // the primary stream still holds queued reads of them — with event tracking elided,
4128 // nothing stops the pool from reusing those blocks for THIS round's stage allocations,
4129 // whose writes then race the queued reads (measured: 13/4096-NaN random-bits garbage in
4130 // the spec round seed; the full anatomy is on `PpNRt::fence_stages_behind`). Order every
4131 // stage stream behind the caller before enqueueing new stage work.
4132 let reverse_started = std::time::Instant::now();
4133 if pp_pipe != Some(false) {
4134 rt.fence_stages_behind(&caller_stream)?;
4135 }
4136 if pp_pipe == Some(true) {
4137 // Both session verifies must alternate boundary slots even when the ordinary
4138 // decode overlap experiment is off. Prewarm before A's stage 0 so B cannot grow
4139 // slot 1 by synchronizing the RX stream while A's stage 1 is in flight.
4140 rt.prepare_overlap_slots(0, payload)?;
4141 }
4142 if pp_anatomy {
4143 // Drain the reverse-publication dependency before timing stage 0 itself. At c=1 this
4144 // prices any primary-stream rollback/refresh tail inherited from the prior round.
4145 for s in 0..n_st {
4146 let _st = rt.enter(s);
4147 rt.engine(s, e).stream().synchronize()?;
4148 }
4149 reverse_ms = reverse_started.elapsed().as_secs_f64() * 1e3;
4150 }
4151
4152 // Per-stage rope positions: in host mode the same [T] iota each stage uploads itself; in
4153 // stream mode each stage's own `pos_iota` over the shared read-only device counter.
4154 let stage_pos = |es: &Engine| -> Result<CudaSlice<i32>, Box<dyn std::error::Error>> {
4155 match stream {
4156 Some((_, ctr)) => {
4157 let mut p = es.alloc_uninit::<i32>(t)?;
4158 es.pos_iota(ctr, &mut p, t)?;
4159 Ok(p)
4160 }
4161 None => {
4162 let pos_vec: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
4163 es.htod_i32(&pos_vec)
4164 }
4165 }
4166 };
4167
4168 // ---- STAGE 0: embed (the table lives with stage 0) + layers [0, fence[1]) + TX ----
4169 let slot = {
4170 let _st0 = rt.enter(0);
4171 let e0 = rt.engine(0, e);
4172 enqueue_spec_pipe_trace_marker(&e0.stream(), trace.as_ref(), "S0", "start", None)?;
4173 let stage0_started = std::time::Instant::now();
4174 let pos_d = stage_pos(e0)?;
4175 let x = match (stream, embd_dev) {
4176 (Some((vtok, _)), Some((g, qt, rb))) => {
4177 e0.embed_gather_device_td(g, vtok, t, n_embd, qt, rb)?
4178 }
4179 (None, Some((g, qt, rb))) => e0.embed_gather_device_t(g, tokens, n_embd, qt, rb)?,
4180 _ => e0.htod(&self.embd.gather(n_embd, tokens))?,
4181 };
4182 let x = self.verify_layers(
4183 e0,
4184 x,
4185 fence[0],
4186 fence[1],
4187 &pos_d,
4188 pos0,
4189 t,
4190 cache,
4191 ckpt.as_deref_mut(),
4192 stream,
4193 None,
4194 )?;
4195 if pp_anatomy {
4196 e0.stream().synchronize()?;
4197 stage0_ms = stage0_started.elapsed().as_secs_f64() * 1e3;
4198 }
4199 let tx_started = std::time::Instant::now();
4200 let slot = if pp_pipe.is_some() {
4201 rt.tx_pipelined(0, &x, payload)?
4202 } else {
4203 rt.tx(0, &x, payload)?
4204 };
4205 enqueue_spec_pipe_trace_marker(&e0.stream(), trace.as_ref(), "S0", "end", Some(slot))?;
4206 if pp_anatomy {
4207 e0.stream().synchronize()?;
4208 tx_ms = tx_started.elapsed().as_secs_f64() * 1e3;
4209 }
4210 slot
4211 // x + pos_d drop here: freed stream-ordered on stage-0's stream after use.
4212 };
4213
4214 Ok(VerifyBoundaryTicket {
4215 rt,
4216 caller_stream,
4217 slot,
4218 pos0,
4219 t,
4220 payload,
4221 n_st,
4222 pipelined: pp_pipe.is_some(),
4223 pp_anatomy,
4224 pp_started,
4225 reverse_ms,
4226 stage0_ms,
4227 tx_ms,
4228 trace,
4229 })
4230 }
4231
4232 /// Consume a stage-0 boundary ticket and enqueue the remaining PP stages plus the head.
4233 /// On PP-2 this is exactly stage 1; PP-N keeps its pre-existing middle-stage walk here.
4234 #[allow(clippy::too_many_arguments)]
4235 fn verify_stage1_finish(
4236 &self,
4237 e: &Engine,
4238 ticket: VerifyBoundaryTicket,
4239 cache: &mut Cache,
4240 mut ckpt: Option<&mut VerifyCkpt>,
4241 stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
4242 fence: &[usize],
4243 publish_to_caller: bool,
4244 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4245 let VerifyBoundaryTicket {
4246 rt,
4247 caller_stream,
4248 slot,
4249 pos0,
4250 t,
4251 payload,
4252 n_st,
4253 pipelined,
4254 pp_anatomy,
4255 pp_started,
4256 reverse_ms,
4257 stage0_ms,
4258 tx_ms,
4259 trace,
4260 } = ticket;
4261 let n_embd = self.cfg.n_embd as usize;
4262 let eps = self.cfg.rms_eps;
4263 let mut slot = slot;
4264 let (mut rx_ms, mut stage1_ms) = (0.0f64, 0.0f64);
4265 let stage_pos = |es: &Engine| -> Result<CudaSlice<i32>, Box<dyn std::error::Error>> {
4266 match stream {
4267 Some((_, ctr)) => {
4268 let mut p = es.alloc_uninit::<i32>(t)?;
4269 es.pos_iota(ctr, &mut p, t)?;
4270 Ok(p)
4271 }
4272 None => {
4273 let pos_vec: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
4274 es.htod_i32(&pos_vec)
4275 }
4276 }
4277 };
4278
4279 // ---- MIDDLE STAGES: RX boundary s-1 -> range -> TX boundary s ----
4280 for s in 1..n_st - 1 {
4281 let _st = rt.enter(s);
4282 let es = rt.engine(s, e);
4283 let pos_d = stage_pos(es)?;
4284 let x = rt.rx(s - 1, slot, payload)?;
4285 let x = self.verify_layers(
4286 es,
4287 x,
4288 fence[s],
4289 fence[s + 1],
4290 &pos_d,
4291 pos0,
4292 t,
4293 cache,
4294 ckpt.as_deref_mut(),
4295 stream,
4296 None,
4297 )?;
4298 slot = if pipelined {
4299 rt.tx_pipelined(s, &x, payload)?
4300 } else {
4301 rt.tx(s, &x, payload)?
4302 };
4303 }
4304
4305 // ---- LAST STAGE: RX + final range + output_norm + lm head ----
4306 let _stl = rt.enter(n_st - 1);
4307 let el = rt.engine(n_st - 1, e);
4308 let pos_d = stage_pos(el)?;
4309 let rx_started = std::time::Instant::now();
4310 let x = rt.rx(n_st - 2, slot, payload)?;
4311 if pp_anatomy {
4312 el.stream().synchronize()?;
4313 rx_ms = rx_started.elapsed().as_secs_f64() * 1e3;
4314 }
4315 enqueue_spec_pipe_trace_marker(&el.stream(), trace.as_ref(), "S1", "start", Some(slot))?;
4316 let stage1_started = std::time::Instant::now();
4317 let x = self.verify_layers(
4318 el,
4319 x,
4320 fence[n_st - 1],
4321 fence[n_st],
4322 &pos_d,
4323 pos0,
4324 t,
4325 cache,
4326 ckpt.as_deref_mut(),
4327 stream,
4328 None,
4329 )?;
4330
4331 let mut hn = vbuf(el, payload)?;
4332 let logits = if self.cfg.step35.is_some() {
4333 // The PP Step35 serving path uses rms_norm + matmul for B=1 as well as B>1.
4334 // Verify must not switch numeric class merely because the same session speculates.
4335 el.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
4336 el.matmul(&self.output, &hn, t)?
4337 } else {
4338 el.rms_norm_decode(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
4339 el.matmul_decode_exact(&self.output, &hn, t)?
4340 };
4341 enqueue_spec_pipe_trace_marker(&el.stream(), trace.as_ref(), "S1", "end", Some(slot))?;
4342 if pp_anatomy {
4343 el.stream().synchronize()?;
4344 stage1_ms = stage1_started.elapsed().as_secs_f64() * 1e3;
4345 }
4346 // EXIT PUBLICATION: both returned buffers are still being produced on the last stage's
4347 // stream. Order the caller's stream behind that work before the buffers escape this
4348 // scope (the 2026-08-06 same-device ppspec find: without it the caller's primary-stream
4349 // consumer read unwritten logits — nondeterministic, one-device-only, and it poisoned
4350 // the following arm's KV in the same process).
4351 if publish_to_caller {
4352 rt.publish_to(n_st - 1, &caller_stream)?;
4353 }
4354 if pp_anatomy {
4355 if publish_to_caller {
4356 caller_stream.synchronize()?;
4357 }
4358 eprintln!(
4359 "[spec-pp-anatomy] t={t} reverse={reverse_ms:.3}ms stage0={stage0_ms:.3}ms \
4360 tx={tx_ms:.3}ms rx={rx_ms:.3}ms stage1-head={stage1_ms:.3}ms total={:.3}ms",
4361 pp_started.elapsed().as_secs_f64() * 1e3,
4362 );
4363 }
4364 // stream: the device pos counter owns position; host mirror reconciles at drain.
4365 if stream.is_none() {
4366 cache.pos += t;
4367 }
4368 Ok((logits, if spec_hpost() { hn } else { x }))
4369 }
4370
4371 /// Step3.5/Step3.7 verify trunk in the serving batched numeric class.
4372 ///
4373 /// `step35_decode_batch_layers` is now authoritative at every live serving width, including
4374 /// B=1 (lane/cx-b1fix). The older verify walk deliberately mirrored the eager T=1 class:
4375 /// it replayed `step35_decode_attn` per row and used the eager/decode-exact FFN dispatch.
4376 /// Those classes are individually stable, but a near-tie prompt can choose different greedy
4377 /// bytes when a request moves from batched plain serving into speculative verify. Run the
4378 /// same authoritative B=1 stage subgraph for each verify row here. Rows still advance
4379 /// layer-by-layer, so every layer sees the preceding verify rows in its attention cache while
4380 /// every norm/projection/FFN uses exactly the live serving dispatch.
4381 #[allow(clippy::too_many_arguments)]
4382 fn step35_verify_batch_layers(
4383 &self,
4384 e: &Engine,
4385 mut x: CudaSlice<f32>,
4386 lo: usize,
4387 hi: usize,
4388 pos0: usize,
4389 t: usize,
4390 cache: &mut Cache,
4391 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4392 let n_embd = self.cfg.n_embd as usize;
4393 self.cfg
4394 .step35
4395 .as_ref()
4396 .ok_or("step35 verify batch requires step35 cfg")?;
4397 let mut ph_last = std::time::Instant::now();
4398 for il in lo..hi {
4399 let mut next = e.uninit(t * n_embd)?;
4400 for r in 0..t {
4401 let mut row = e.uninit(n_embd)?;
4402 e.dtod_copy_view(&x.slice(r * n_embd..(r + 1) * n_embd), &mut row)?;
4403 // The caller owns this verify's position. During controller overlap, cache.pos
4404 // still describes generation N while this stage-0 walk belongs to N+1.
4405 let row_pos = e.htod_i32(&[(pos0 + r) as i32])?;
4406 let mut one = [&mut *cache];
4407 let out = self.step35_decode_batch_layers(
4408 e,
4409 row,
4410 &mut one,
4411 &row_pos,
4412 il,
4413 il + 1,
4414 &mut ph_last,
4415 )?;
4416 e.dtod_copy_into(&out, &mut next, r * n_embd)?;
4417 }
4418 self.dflash_tap(e, cache, il, &next, t)?;
4419 x = next;
4420 }
4421 Ok(x)
4422 }
4423
4424 /// DSpark drafter verify (lane/dspark-q38-recover): one t-row forward through the
4425 /// SERVING-CLASS verify funnel (`decode_step_t_core_stream` — the same numeric class
4426 /// MTP verify rides, GDN state advanced in place), returning per-row argmax tokens.
4427 /// Advances `cache.pos += t`; the caller owns snapshot/rollback (block acceptance is
4428 /// prefix-keep, not all-or-nothing).
4429 pub(crate) fn dspark_verify_t_am(
4430 &self,
4431 e: &Engine,
4432 tokens: &[u32],
4433 pos0: usize,
4434 cache: &mut Cache,
4435 ) -> Result<Vec<u32>, Box<dyn std::error::Error>> {
4436 let (logits, _hn) = self.decode_step_t_core_stream(
4437 e, tokens, pos0, cache, None, None, None, None, None, None,
4438 )?;
4439 let t = tokens.len();
4440 let v = self.output.out_features();
4441 let mut am_d = e.stream().alloc_zeros::<u32>(t)?;
4442 for r in 0..t {
4443 e.argmax_token_device_col(&logits, r, v, &mut am_d, r)?;
4444 }
4445 Ok(e.dtoh_u32(&am_d)?)
4446 }
4447
4448 /// DSpark verify with the MTP column-stash armed: identical forward to
4449 /// `dspark_verify_t_am`, but fills a `VerifyCkpt` so a partial accept can restore
4450 /// column state directly (`dspark_commit_prefix`) instead of snapshot-replay.
4451 /// The ckpt type is opaque outside spec.rs (newtype) — dflash.rs threads it through.
4452 pub(crate) fn dspark_verify_t_am_ckpt(
4453 &self,
4454 e: &Engine,
4455 tokens: &[u32],
4456 pos0: usize,
4457 cache: &mut Cache,
4458 ) -> Result<(Vec<u32>, DsparkVerifyCkpt), Box<dyn std::error::Error>> {
4459 let mut ck = VerifyCkpt::new(self.layers.len());
4460 let (logits, _hn) = self.decode_step_t_core_stream(
4461 e,
4462 tokens,
4463 pos0,
4464 cache,
4465 None,
4466 Some(&mut ck),
4467 None,
4468 None,
4469 None,
4470 None,
4471 )?;
4472 let t = tokens.len();
4473 let v = self.output.out_features();
4474 let mut am_d = e.stream().alloc_zeros::<u32>(t)?;
4475 for r in 0..t {
4476 e.argmax_token_device_col(&logits, r, v, &mut am_d, r)?;
4477 }
4478 Ok((e.dtoh_u32(&am_d)?, DsparkVerifyCkpt(ck)))
4479 }
4480
4481 /// Engine-bundle slice 2: `dspark_verify_t_am_ckpt` with DEVICE tokens and NO readback.
4482 /// The verify tokens are the round's `chain_d` (cand layout: [anchor, drafts...]); the
4483 /// embed gathers its first `t` entries on-device (`embed_gather_u32_t` — bit-identical
4484 /// rows to the host arm), so the host never blocks on the draft chain before dispatching
4485 /// verify. Returns the device per-row argmax buffer; the caller merges its readback with
4486 /// the chain's into ONE sync. Forward, ckpt fill and argmax walk are `_ckpt` verbatim.
4487 pub(crate) fn dspark_verify_t_am_ckpt_dev(
4488 &self,
4489 e: &Engine,
4490 vtok: &CudaSlice<u32>,
4491 t: usize,
4492 pos0: usize,
4493 cache: &mut Cache,
4494 embd_dev: (&CudaSlice<u8>, i32, usize),
4495 graphs: Option<&mut DsparkVerifyGraphs>,
4496 ) -> Result<(CudaSlice<u32>, DsparkVerifyCkpt), Box<dyn std::error::Error>> {
4497 debug_assert!(
4498 vtok.len() >= t,
4499 "verify window exceeds the device token buffer"
4500 );
4501 // The slab flag is a per-round statement: clear it here so a verify that never
4502 // reaches the graphs door (rowwise env, a non-tparallel arm) cannot leave a
4503 // stale `true` steering the commit at slabs the round never wrote.
4504 let mut graphs = graphs;
4505 if let Some(g) = graphs.as_deref_mut() {
4506 g.round_slab = false;
4507 }
4508 let mut ck = VerifyCkpt::new(self.layers.len());
4509 // Dummy host tokens size the funnel; the embed reads `vtok` (the round-stream
4510 // arm's established pattern — spec.rs stream-mode verify does the same).
4511 let dummy = vec![0u32; t];
4512 let (logits, _hn) = self.decode_step_t_core_stream(
4513 e,
4514 &dummy,
4515 pos0,
4516 cache,
4517 Some(embd_dev),
4518 Some(&mut ck),
4519 None,
4520 None,
4521 Some(vtok),
4522 graphs,
4523 )?;
4524 let v = self.output.out_features();
4525 let mut am_d = e.stream().alloc_zeros::<u32>(t)?;
4526 for r in 0..t {
4527 e.argmax_token_device_col(&logits, r, v, &mut am_d, r)?;
4528 }
4529 Ok((am_d, DsparkVerifyCkpt(ck)))
4530 }
4531
4532 /// Restore the round to `keep` accepted columns from the verify stash: KV lens and
4533 /// pos from the pre-verify snapshot + keep, GDN conv/ssm from the stashed column
4534 /// state — no replay forward. The exact `commit_verified_prefix` the MTP path ships.
4535 pub(crate) fn dspark_commit_prefix(
4536 &self,
4537 e: &Engine,
4538 cache: &mut Cache,
4539 snap: &crate::cache::CacheSnapshot,
4540 ckpt: &DsparkVerifyCkpt,
4541 keep: usize,
4542 ) -> Result<(), Box<dyn std::error::Error>> {
4543 self.commit_verified_prefix(e, cache, snap, &ckpt.0, keep, false, None)
4544 }
4545
4546 /// Slice-3 commit twin: restore to `keep` accepted columns when the round's linear
4547 /// column stash lives in the graphs ctx's persistent slabs (`DsparkVerifyGraphs`) —
4548 /// the cols arm's exact semantics (KV lens + pos from the snapshot, GDN conv/ssm
4549 /// from the stash of column keep-1), slab-addressed and batched into two copy
4550 /// launches. `MEMRA_STATE_COPY_BATCH=0` falls back to per-layer view copies.
4551 pub(crate) fn dspark_commit_prefix_slab(
4552 &self,
4553 e: &Engine,
4554 cache: &mut Cache,
4555 snap: &crate::cache::CacheSnapshot,
4556 ctx: &DsparkVerifyGraphs,
4557 keep: usize,
4558 ) -> Result<(), Box<dyn std::error::Error>> {
4559 use cudarc::driver::DevicePtr;
4560 debug_assert!(keep >= 1, "keep==0 rounds take the legacy rollback");
4561 let mut conv_src: Vec<u64> = Vec::new();
4562 let mut ssm_src: Vec<u64> = Vec::new();
4563 let mut conv_dst: Vec<u64> = Vec::new();
4564 let mut ssm_dst: Vec<u64> = Vec::new();
4565 for il in 0..self.layers.len() {
4566 if let (Some(kvl), Some(saved)) = (cache.kv[il].as_mut(), snap.kv_len[il]) {
4567 kvl.len = saved + keep;
4568 e.set_i32_one(&mut kvl.len_d, kvl.len as i32)?;
4569 }
4570 if let Some(rl) = cache.recur[il].as_ref() {
4571 let (pc, ps, _cw, _sw) = ctx
4572 .slab_row(e, il, keep - 1)
4573 .ok_or("slab commit: linear layer missing from the graphs ctx")?;
4574 conv_src.push(pc);
4575 ssm_src.push(ps);
4576 let st = &e.gpu.stream();
4577 let (dc, _g0) = rl.conv_state.device_ptr(st);
4578 let (ds, _g1) = rl.ssm_state.device_ptr(st);
4579 conv_dst.push(dc as u64);
4580 ssm_dst.push(ds as u64);
4581 }
4582 }
4583 let n = conv_src.len();
4584 if n > 0 {
4585 if state_copy_batch_on() {
4586 let mut tt = vec![0u64; 2 * n];
4587 tt[..n].copy_from_slice(&conv_src);
4588 tt[n..].copy_from_slice(&conv_dst);
4589 let ct = e.htod_u64(&tt)?;
4590 tt[..n].copy_from_slice(&ssm_src);
4591 tt[n..].copy_from_slice(&ssm_dst);
4592 let st = e.htod_u64(&tt)?;
4593 e.copy_batch_uniform_f32(&ct, n, ctx.conv_words)?;
4594 e.copy_batch_uniform_f32(&st, n, ctx.ssm_words)?;
4595 } else {
4596 let (cw, sw) = (ctx.conv_words, ctx.ssm_words);
4597 let row = keep - 1;
4598 for il in 0..self.layers.len() {
4599 let Some(rl) = cache.recur[il].as_mut() else {
4600 continue;
4601 };
4602 let k = ctx.lin_pos[&il];
4603 {
4604 let sv = e.view(&ctx.stash_conv[k], (row + 1) * cw);
4605 let win = sv.slice(row * cw..(row + 1) * cw);
4606 e.copy_view_into(&mut rl.conv_state, 0, &win, cw)?;
4607 }
4608 {
4609 let sv = e.view(&ctx.stash_ssm[k], (row + 1) * sw);
4610 let win = sv.slice(row * sw..(row + 1) * sw);
4611 e.copy_view_into(&mut rl.ssm_state, 0, &win, sw)?;
4612 }
4613 }
4614 }
4615 }
4616 cache.pos = snap.pos + keep;
4617 Ok(())
4618 }
4619
4620 /// Qwen35-family verify trunk in the live serving numeric class.
4621 ///
4622 /// Serving intentionally keeps this architecture in the generic batched program even at
4623 /// B=1. The older verify walk used its own mirrored dispatch and can flip near-tie argmaxes.
4624 ///
4625 /// Two arms, one numeric class:
4626 /// - DENSE (`Arch::Qwen35`, t<=16): `qwen35_verify_tparallel` — the weight ops (norms,
4627 /// projections, FFN) hoist to m=T through the exact-tier batched kernels whose per-row
4628 /// program IS the m=1 program (`matmul_pre == fused2 per (tensor,row); _bN mmvq per-row
4629 /// == m=1` — decode_batch.rs v2 note), while the state ops (conv ring, gdn scan, KV
4630 /// append, fa decode) stay a per-row loop running the b_n=1 serving kernels with each
4631 /// row's own t_kv-driven arm pick (the straddle law: every row executes the exact
4632 /// program its isolated serving step would). One weight read per layer per round
4633 /// instead of T — this is what makes MTP profitable in the exact class (the per-row
4634 /// walk measured verify(K+1) ~= (K+1) plain steps: 69 -> 44 tok/s served, 2026-08-15).
4635 /// - MoE / t>16 / `MEMRA_SPEC_VERIFY_ROWWISE=1`: the per-row replay of the authoritative
4636 /// serving layer body, preserving single-session autoregressive cache order (the
4637 /// correctness reference; also the rollback seam for the t-parallel arm).
4638 ///
4639 /// Bit-identity of the t-parallel arm vs the rowwise arm is gated by spec-serve-gate
4640 /// (zero differing logits at T=1..4, K arms) + the 8-prompt ON/OFF canary before ship.
4641 #[allow(clippy::too_many_arguments)]
4642 fn qwen35_verify_batch_layers(
4643 &self,
4644 e: &Engine,
4645 x: CudaSlice<f32>,
4646 lo: usize,
4647 hi: usize,
4648 pos0: usize,
4649 t: usize,
4650 cache: &mut Cache,
4651 ckpt: Option<&mut VerifyCkpt>,
4652 stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
4653 graphs: Option<&mut DsparkVerifyGraphs>,
4654 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4655 // Qwen35Moe admitted 2026-08-20 (lane/draftcost-moe): the t-parallel arm already
4656 // carries the MoE FFN (`moe_ffn_il_zq8` at m=T) and the GDN per-row state loop; the
4657 // arch fence was a qualification gate, not a mechanism gap. Measured disease on the
4658 // 35B-A3B class: rowwise verify ~= 5.6 ms per drafted token (one full trunk step
4659 // each) — the same (K+1)-plain-steps wall the dense admission fixed on 2026-08-15.
4660 // Rollback seam unchanged: MEMRA_SPEC_VERIFY_ROWWISE=1.
4661 let rowwise = std::env::var("MEMRA_SPEC_VERIFY_ROWWISE").as_deref() == Ok("1")
4662 || !matches!(
4663 self.cfg.arch,
4664 memra_gguf::config::Arch::Qwen35 | memra_gguf::config::Arch::Qwen35Moe
4665 )
4666 || t > 16;
4667 if rowwise {
4668 if stream.is_some() {
4669 // rowwise replays per row with host cache.pos — irreconcilable with a
4670 // device position counter. Burst callers must keep t <= 16 and the
4671 // ROWWISE env unset; refusing beats silently mispositioned rows.
4672 return Err("qwen35 rowwise verify has no ROUND-STREAM arm \
4673 (t > 16 or MEMRA_SPEC_VERIFY_ROWWISE=1)"
4674 .into());
4675 }
4676 self.qwen35_verify_rowwise(e, x, lo, hi, pos0, t, cache, ckpt)
4677 } else {
4678 self.qwen35_verify_tparallel(e, x, lo, hi, pos0, t, cache, ckpt, stream, graphs)
4679 }
4680 }
4681
4682 /// The per-row correctness reference: replay each verify row through the authoritative
4683 /// serving layer body (`decode_batch_layers` at b_n=1). T full weight reads per layer.
4684 #[allow(clippy::too_many_arguments)]
4685 fn qwen35_verify_rowwise(
4686 &self,
4687 e: &Engine,
4688 mut x: CudaSlice<f32>,
4689 lo: usize,
4690 hi: usize,
4691 pos0: usize,
4692 t: usize,
4693 cache: &mut Cache,
4694 mut ckpt: Option<&mut VerifyCkpt>,
4695 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4696 let n_embd = self.cfg.n_embd as usize;
4697 let saved_pos = cache.pos;
4698 let mut ph_last = std::time::Instant::now();
4699 for il in lo..hi {
4700 let mut next = e.uninit(t * n_embd)?;
4701 let mut col_states: Option<Vec<(CudaSlice<f32>, CudaSlice<f32>)>> =
4702 if ckpt.is_some() && t >= 2 && matches!(self.layers[il].mixer, Mixer::Linear(_)) {
4703 Some(Vec::with_capacity(t - 1))
4704 } else {
4705 None
4706 };
4707 for r in 0..t {
4708 cache.pos = pos0 + r;
4709 let mut row = e.uninit(n_embd)?;
4710 e.dtod_copy_view(&x.slice(r * n_embd..(r + 1) * n_embd), &mut row)?;
4711 let row_pos = e.htod_i32(&[(pos0 + r) as i32])?;
4712 let mut one = [&mut *cache];
4713 let ctx = self.batch_layer_ctx(e, &one, il, il + 1)?;
4714 let out = match self.decode_batch_layers(
4715 e,
4716 row,
4717 &mut one,
4718 &ctx,
4719 &row_pos,
4720 &mut ph_last,
4721 ) {
4722 Ok(out) => out,
4723 Err(error) => {
4724 cache.pos = saved_pos;
4725 return Err(error);
4726 }
4727 };
4728 e.dtod_copy_into(&out, &mut next, r * n_embd)?;
4729 if r + 1 < t {
4730 if let Some(states) = col_states.as_mut() {
4731 let recur = cache.recur[il]
4732 .as_ref()
4733 .ok_or("Qwen35-MoE linear verify layer has no recurrent state")?;
4734 states.push((
4735 e.clone_dtod(&recur.conv_state)?,
4736 e.clone_dtod(&recur.ssm_state)?,
4737 ));
4738 }
4739 }
4740 }
4741 if let (Some(checkpoint), Some(states)) = (ckpt.as_deref_mut(), col_states) {
4742 checkpoint.cols[il] = Some(states);
4743 }
4744 x = next;
4745 }
4746 cache.pos = saved_pos;
4747 Ok(x)
4748 }
4749
4750 /// T-PARALLEL VERIFY IN THE SERVING NUMERIC CLASS (lane/tparallel-verify, 2026-08-15).
4751 ///
4752 /// The weight ops run ONCE per layer at m=T; the state ops run per row through the same
4753 /// b_n=1 serving kernels the rowwise replay uses. Per-row bit-identity rests on the two
4754 /// pins the serving batch tier already carries:
4755 /// * `matmul_pre` / `_bN` mmvq: per-row program == m=1 program (decode_batch.rs v2 note,
4756 /// kernel-check pinned) — so a [T, n_embd] projection row equals the row projected
4757 /// alone;
4758 /// * row-indexed norms/elementwise (`rms_norm`, `quantize_q8_1`, `add_rms_norm`,
4759 /// `gated_rmsnorm[_q8_1]`, `silu_mul`, `rope_neox` with per-row positions): the T-row
4760 /// launch is the per-row program (same pin the generic verify's fused norms rely on).
4761 /// The sequential dependencies keep their exact serving order: the conv ring / gdn scan
4762 /// chain state row -> row through the `_b` kernels at b_n=1 (ping-pong via a 6-entry
4763 /// alternating pointer table, host handles swapped per row so VerifyCkpt clones the
4764 /// canonical state exactly as the rowwise arm does), and each row's KV append + fa decode
4765 /// picks its arm from ITS OWN t_kv (append: format-only; fa: `fa_seqs_eligible` + its own
4766 /// `fa_split_keys` rung at b_n=1) — the straddle law per row, so every row executes the
4767 /// program its isolated B=1 serving step would.
4768 ///
4769 /// Cost: 1 weight read per layer per round + T state micro-launches, vs the rowwise arm's
4770 /// T weight reads. Gated bit-identical vs the rowwise arm by spec-serve-gate + canary.
4771 #[allow(clippy::too_many_arguments)]
4772 fn qwen35_verify_tparallel(
4773 &self,
4774 e: &Engine,
4775 mut x: CudaSlice<f32>,
4776 lo: usize,
4777 hi: usize,
4778 pos0: usize,
4779 t: usize,
4780 cache: &mut Cache,
4781 mut ckpt: Option<&mut VerifyCkpt>,
4782 stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
4783 mut graphs: Option<&mut DsparkVerifyGraphs>,
4784 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4785 use cudarc::driver::DevicePtr;
4786 let cfg = &self.cfg;
4787 let n_embd = cfg.n_embd as usize;
4788 let eps = cfg.rms_eps;
4789 let head_dim_global = cfg.head_dim_k as usize;
4790 // STREAM (2b, lane/draftcost-moe): positions come from the device round counter
4791 // (pos_iota / i32_copy_add) so a burst round needs no host position knowledge.
4792 let pos_d = match stream {
4793 Some((_, ctr)) => {
4794 let mut p = e.alloc_uninit::<i32>(t)?;
4795 e.pos_iota(ctr, &mut p, t)?;
4796 p
4797 }
4798 None => {
4799 let pos_host: Vec<i32> = (0..t).map(|r| (pos0 + r) as i32).collect();
4800 e.htod_i32(&pos_host)?
4801 }
4802 };
4803 // Per-row 1-element position buffers, built ONCE per verify (the append/fa wrappers
4804 // take owned pos slices; building these inside the layer x row loops cost 16xT H2Ds).
4805 let pos_rows: Vec<CudaSlice<i32>> = match stream {
4806 Some((_, ctr)) => (0..t)
4807 .map(|r| {
4808 let mut b = e.alloc_uninit::<i32>(1)?;
4809 e.i32_copy_add(ctr, &mut b, r as i32)?;
4810 Ok(b)
4811 })
4812 .collect::<Result<_, Box<dyn std::error::Error>>>()?,
4813 None => (0..t)
4814 .map(|r| e.htod_i32(&[(pos0 + r) as i32]))
4815 .collect::<Result<_, _>>()?,
4816 };
4817 let seqs_append =
4818 std::env::var("MEMRA_BATCH_APPEND").as_deref() != Ok("0") && !Engine::kv_fp8_on();
4819 let batch_fa_on = std::env::var("MEMRA_BATCH_FA").as_deref() != Ok("0");
4820
4821 // Engine-bundle slice 3: with a graphs ctx armed, runs of consecutive LINEAR layers
4822 // replay per-(segment, vt) CUDA graphs (captured from the SAME
4823 // `qwen35_tparallel_linear_layer` body the eager arm runs — no second copy of the
4824 // math). The full-attention layers stay eager: their per-row append/fa arm picks are
4825 // t_kv-driven (the straddle law) and belong to the exec-update extension, not this
4826 // slice. Pointer tables are refreshed once per verify (gdn ping-pong moves handles).
4827 // Merge guard (v0.98 train): the ROUND-STREAM arm (lane/draftcost-moe, device
4828 // position counter) and the dspark verify graphs (engine-bundle slice 3) have no
4829 // common caller — stream rides the qwen35moe burst, graphs ride the dspark route.
4830 // If a future caller arms both, refuse loudly instead of silently dropping the
4831 // graphs ctx (the stream linear arm takes linear_attn_verify_t, not the graphed
4832 // segment body).
4833 if stream.is_some() && graphs.is_some() {
4834 return Err(
4835 "qwen35 tparallel verify: ROUND-STREAM and dspark verify graphs \
4836 cannot arm together"
4837 .into(),
4838 );
4839 }
4840 if let Some(g) = graphs.as_deref_mut() {
4841 g.refresh_tables(e, cache)?;
4842 g.round_slab = false;
4843 }
4844 let mut il = lo;
4845 while il < hi {
4846 if graphs.is_some() && matches!(self.layers[il].mixer, Mixer::Linear(_)) {
4847 let mut end = il;
4848 while end < hi && matches!(self.layers[end].mixer, Mixer::Linear(_)) {
4849 end += 1;
4850 }
4851 let g = graphs.as_deref_mut().expect("checked above");
4852 x = g.run_segment(self, e, il, end, &x, t, cache)?;
4853 g.round_slab = true;
4854 il = end;
4855 continue;
4856 }
4857 let layer = &self.layers[il];
4858 if stream.is_none() && matches!(layer.mixer, Mixer::Linear(_)) {
4859 // Eager linear layer (no graphs ctx): the shared body, legacy cols-ckpt arm.
4860 // Under ROUND-STREAM the linear layers ride the match's stream arm below
4861 // (linear_attn_verify_t — the stream COMMIT needs its GdnStash).
4862 x = self.qwen35_tparallel_linear_layer(
4863 e,
4864 il,
4865 &x,
4866 t,
4867 cache,
4868 ckpt.as_deref_mut(),
4869 None,
4870 None,
4871 )?;
4872 il += 1;
4873 continue;
4874 }
4875 // ---- attn_norm + q8_1 quantize at m=T (row-indexed == per-row) ----
4876 let anorm = layer.attn_norm.float_data();
4877 let mut xn = e.uninit(t * n_embd)?;
4878 e.rms_norm(&x, anorm, &mut xn, n_embd, t, eps)?;
4879 let (hq, hd) = e.quantize_q8_1(&xn, t, n_embd)?;
4880
4881 let mixed: CudaSlice<f32> = match &layer.mixer {
4882 Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
4883 Mixer::Full(fa) => {
4884 let geometry = cfg.full_attention_geometry_at(il as u32);
4885 let n_head = geometry.n_head as usize;
4886 let n_head_kv = geometry.n_head_kv as usize;
4887 let head_dim = geometry.head_dim_k as usize;
4888 let rope_dims = geometry.n_rot as usize;
4889 let rope_base = geometry.rope_base;
4890 let scale = geometry.attention_scale();
4891 // Batched projections: one weight read serves all T rows.
4892 let qf = e.matmul_pre(&fa.wq, &hq, &hd, &xn, t)?;
4893 let mut k = e.matmul_pre(&fa.wk, &hq, &hd, &xn, t)?;
4894 let v = e.matmul_pre(&fa.wv, &hq, &hd, &xn, t)?;
4895 let gated =
4896 geometry.attention_gate == memra_gguf::config::AttentionGateKind::FusedQ;
4897 let (mut q, gate) = if gated {
4898 let mut qs = e.uninit(t * n_head * head_dim)?;
4899 let mut gs = e.uninit(t * n_head * head_dim)?;
4900 e.q_gate_split(&qf, &mut qs, &mut gs, head_dim, n_head, t)?;
4901 (qs, Some(gs))
4902 } else {
4903 (qf, None)
4904 };
4905 let mut qn = e.uninit(t * n_head * head_dim)?;
4906 e.rms_norm(
4907 &q,
4908 fa.q_norm.float_data(),
4909 &mut qn,
4910 head_dim,
4911 t * n_head,
4912 eps,
4913 )?;
4914 q = qn;
4915 let mut kn = e.uninit(t * n_head_kv * head_dim)?;
4916 e.rms_norm(
4917 &k,
4918 fa.k_norm.float_data(),
4919 &mut kn,
4920 head_dim,
4921 t * n_head_kv,
4922 eps,
4923 )?;
4924 k = kn;
4925 e.rope_neox(
4926 &mut q, &pos_d, head_dim, rope_dims, n_head, t, rope_base, 1.0,
4927 )?;
4928 e.rope_neox(
4929 &mut k, &pos_d, head_dim, rope_dims, n_head_kv, t, rope_base, 1.0,
4930 )?;
4931
4932 // Per-row append + attend: row r sees rows 0..r in KV (causal within the
4933 // draft), each through the b_n=1 serving kernels at its own t_kv.
4934 let q_dim = n_head * head_dim;
4935 let kv_dim = n_head_kv * head_dim;
4936 let mut attn = e.uninit(t * q_dim)?;
4937 let (kdk, kdv, ktb, vtb, kv_view) = {
4938 let kvl = cache.kv[il].as_ref().unwrap();
4939 let s = &e.gpu.stream();
4940 let (pk, _g) = kvl.k.device_ptr(s);
4941 let (pv, _g2) = kvl.v.device_ptr(s);
4942 (
4943 kvl.kv_dim_k,
4944 kvl.kv_dim_v,
4945 kvl.k_tok_bytes,
4946 kvl.v_tok_bytes,
4947 e.htod_u64(&[pk as u64, pv as u64])?,
4948 )
4949 };
4950 if let Some((_, ctr)) = stream {
4951 // STREAM ARM (2b): one batched dc append + the multi-row dc attention
4952 // — the generic stream arm's exact shape (rows kernels are pinned
4953 // byte-identical to the per-row programs by kernel-check). Host len
4954 // stays a stale lower bound; the burst drain reconciles it.
4955 let kvl = cache.kv[il].as_mut().unwrap();
4956 e.append_kv_quantized_rows_dc(
4957 &k,
4958 &v,
4959 &mut kvl.k,
4960 &mut kvl.v,
4961 ctr,
4962 t,
4963 kdk,
4964 kdv,
4965 ktb,
4966 vtb,
4967 Engine::kv_fp8_on(),
4968 )?;
4969 let upper = (kvl.len + t + 64).min(cache.max_ctx);
4970 let k_view = e.view_u8(&kvl.k, upper * ktb);
4971 let v_view = e.view_u8(&kvl.v, upper * vtb);
4972 e.fa_decode_rows_dc(
4973 &q, &k_view, &v_view, &mut attn, head_dim, n_head, n_head_kv, ctr,
4974 upper, t, scale, ktb, vtb, 0, false,
4975 )?;
4976 } else {
4977 for r in 0..t {
4978 // Owned per-row scratch: the b_n=1 kernels take packed batch buffers
4979 // whose row 0 is this row (arithmetic-free materialization copies,
4980 // same as decode's per-seq fallback arm).
4981 let mut k_row = e.uninit(kv_dim)?;
4982 e.dtod_copy_view(&k.slice(r * kv_dim..(r + 1) * kv_dim), &mut k_row)?;
4983 let mut v_row = e.uninit(kv_dim)?;
4984 e.dtod_copy_view(&v.slice(r * kv_dim..(r + 1) * kv_dim), &mut v_row)?;
4985 let pos_row = &pos_rows[r];
4986 let kvl = cache.kv[il].as_mut().unwrap();
4987 if seqs_append {
4988 e.append_kv_quantized_seqs(
4989 &k_row,
4990 &v_row,
4991 &kv_view.slice(0..2),
4992 pos_row,
4993 1,
4994 kdk,
4995 kdv,
4996 ktb,
4997 vtb,
4998 )?;
4999 kvl.len += 1;
5000 } else {
5001 e.append_kv_quantized_view(
5002 &k_row.slice(0..kv_dim),
5003 &v_row.slice(0..kv_dim),
5004 &mut kvl.k,
5005 &mut kvl.v,
5006 kvl.len,
5007 kvl.kv_dim_k,
5008 kvl.kv_dim_v,
5009 kvl.k_tok_bytes,
5010 kvl.v_tok_bytes,
5011 Engine::kv_fp8_on(),
5012 )?;
5013 kvl.len += 1;
5014 }
5015 let t_kv = kvl.len;
5016 let mut q_row = e.uninit(q_dim)?;
5017 e.dtod_copy_view(&q.slice(r * q_dim..(r + 1) * q_dim), &mut q_row)?;
5018 let mut a_row = e.uninit(q_dim)?;
5019 if batch_fa_on && crate::fa_seqs_eligible(t_kv, head_dim_global) {
5020 let sp0_r = crate::fa_split_keys(t_kv, cfg.n_head_kv as usize);
5021 e.fa_decode_batch_seqs_v4(
5022 &q_row,
5023 &kv_view.slice(0..2),
5024 pos_row,
5025 &mut a_row,
5026 head_dim,
5027 n_head,
5028 n_head_kv,
5029 1,
5030 t_kv,
5031 scale,
5032 sp0_r,
5033 ktb,
5034 vtb,
5035 )?;
5036 } else {
5037 let k_view = e.view_u8(&kvl.k, t_kv * kvl.k_tok_bytes);
5038 let v_view = e.view_u8(&kvl.v, t_kv * kvl.v_tok_bytes);
5039 let mut a_view = a_row.slice_mut(0..q_dim);
5040 e.fa_decode_kvmod_view(
5041 &q_row.slice(0..q_dim),
5042 &k_view,
5043 &v_view,
5044 &mut a_view,
5045 head_dim,
5046 n_head,
5047 n_head_kv,
5048 t_kv,
5049 scale,
5050 kvl.k_tok_bytes,
5051 kvl.v_tok_bytes,
5052 Engine::kv_fp8_on(),
5053 )?;
5054 }
5055 e.dtod_copy_into(&a_row, &mut attn, r * q_dim)?;
5056 }
5057 }
5058
5059 // Output gate (element-wise) + o-proj at m=T.
5060 let attn_g = match &gate {
5061 Some(g) => {
5062 let n = t * q_dim;
5063 let mut gsig = e.uninit(n)?;
5064 e.sigmoid(g, &mut gsig, n)?;
5065 let mut ag = e.uninit(n)?;
5066 e.mul(&attn, &gsig, &mut ag, n)?;
5067 ag
5068 }
5069 None => attn,
5070 };
5071 e.matmul(&fa.wo, &attn_g, t)?
5072 }
5073 // STREAM ARM (2b, lane/draftcost-moe): under a device position counter the
5074 // per-row serving-kernel chain cannot run (host state swaps keyed on host
5075 // row index are fine, but the stream COMMIT needs the GdnStash for its _dc
5076 // rebuild — the per-row chain only produces per-column clones). GDN rides
5077 // `linear_attn_verify_t`: batched q8_1-class projections, stash-producing,
5078 // and its one-scan recurrence is pinned bit-identical to T chained T=1
5079 // steps (its header + kernel-check). Position-independent, so no counter
5080 // plumbing is needed. Guards mirror the generic call site exactly.
5081 Mixer::Linear(la) if stream.is_some() => {
5082 if !(t >= 3 || (t == 2 && spec_m2()))
5083 || !self.mixer_in_q8_1_fast(e, &layer.mixer)
5084 || !e.uses_q8_1_fast(&la.ssm_out)
5085 {
5086 return Err("qwen35 stream verify: GDN batched arm requires t>=3 \
5087 (or MEMRA_SPEC_M2 at t=2) and q8_1-fast projections"
5088 .into());
5089 }
5090 let want = ckpt.is_some();
5091 let (out, stash) = self.linear_attn_verify_t(
5092 e,
5093 la,
5094 &xn,
5095 Some((&hq, &hd)),
5096 t,
5097 cache,
5098 il,
5099 want,
5100 )?;
5101 if let (Some(ck), Some(st)) = (ckpt.as_deref_mut(), stash) {
5102 ck.gdn[il] = Some(st);
5103 }
5104 out
5105 }
5106 Mixer::Linear(_) => {
5107 unreachable!("linear layers ride qwen35_tparallel_linear_layer")
5108 }
5109 };
5110
5111 // ---- residual add + post_attn_norm + FFN at m=T (serving dispatch verbatim) ----
5112 let pnorm = layer.post_attn_norm.float_data();
5113 let mut x1 = e.uninit(t * n_embd)?;
5114 let mut zn = e.uninit(t * n_embd)?;
5115 e.add_rms_norm(&x, &mixed, pnorm, &mut x1, &mut zn, n_embd, t, eps)?;
5116 let ffn_out = match &layer.ffn {
5117 crate::hybrid::Ffn::Dense {
5118 ffn_gate,
5119 ffn_up,
5120 ffn_down,
5121 } => {
5122 assert!(
5123 self.cfg.m3.is_none(),
5124 "qwen35 t-parallel verify: M3 swigluoai FFN not yet batched"
5125 );
5126 let n_ff = ffn_gate.out_features();
5127 let (zq, zd) = e.quantize_q8_1(&zn, t, n_embd)?;
5128 let g = e.matmul_pre(ffn_gate, &zq, &zd, &zn, t)?;
5129 let u = e.matmul_pre(ffn_up, &zq, &zd, &zn, t)?;
5130 let mut act = e.uninit(t * n_ff)?;
5131 e.silu_mul(&g, &u, &mut act, t * n_ff)?;
5132 let (aq, ad) = e.quantize_q8_1(&act, t, n_ff)?;
5133 e.matmul_pre(ffn_down, &aq, &ad, &act, t)?
5134 }
5135 crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il_zq8(e, m, &zn, None, t, il as u16)?,
5136 };
5137 let mut x2 = e.uninit(t * n_embd)?;
5138 e.add(&x1, &ffn_out, &mut x2, t * n_embd)?;
5139 // dspark drafter tap (no-op when no sink armed): post-layer residual verify rows
5140 self.dflash_tap(e, cache, il, &x2, t)?;
5141 x = x2;
5142 il += 1;
5143 }
5144 Ok(x)
5145 }
5146
5147 /// ONE t-parallel LINEAR layer (attn_norm + gdn mixer + post_attn_norm + FFN + tap) —
5148 /// the exact body the old in-loop Linear arm ran, extracted so the eager walk and the
5149 /// slice-3 captured segments execute the SAME code (a second copy is how dispatch
5150 /// mirrors drift — the verify_layers extraction lesson). Two deliberate changes, both
5151 /// bit-identical by construction:
5152 /// - the gdn ping-pong host swap moves from per-row to ONE end-of-body swap (t odd):
5153 /// the device sequence is driven entirely by the 6-entry pointer table, which
5154 /// already encodes both parities; the ckpt stash reads name row r's out buffer
5155 /// directly (r even -> alt handle, odd -> canonical) — the same physical bytes the
5156 /// legacy post-swap clone read.
5157 /// - `stash` (slice-3 ctx): persistent per-layer slabs written by copy_into instead of
5158 /// per-row clone_dtod allocs — same bytes, capture-legal (no per-round host objects).
5159 /// `table_src` = (persistent pointer table, offset) when the ctx owns the tables;
5160 /// None builds the per-verify table exactly as before.
5161 #[allow(clippy::too_many_arguments)]
5162 fn qwen35_tparallel_linear_layer(
5163 &self,
5164 e: &Engine,
5165 il: usize,
5166 x: &CudaSlice<f32>,
5167 t: usize,
5168 cache: &mut Cache,
5169 mut ckpt: Option<&mut VerifyCkpt>,
5170 stash: Option<(&mut CudaSlice<f32>, &mut CudaSlice<f32>)>,
5171 table_src: Option<(&CudaSlice<u64>, usize)>,
5172 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5173 use cudarc::driver::DevicePtr;
5174 let cfg = &self.cfg;
5175 let n_embd = cfg.n_embd as usize;
5176 let eps = cfg.rms_eps;
5177 let layer = &self.layers[il];
5178 let Mixer::Linear(la) = &layer.mixer else {
5179 return Err("qwen35_tparallel_linear_layer on a non-linear layer".into());
5180 };
5181 // ---- attn_norm + q8_1 quantize at m=T (row-indexed == per-row) ----
5182 let anorm = layer.attn_norm.float_data();
5183 let mut xn = e.uninit(t * n_embd)?;
5184 e.rms_norm(x, anorm, &mut xn, n_embd, t, eps)?;
5185 let (hq, hd) = e.quantize_q8_1(&xn, t, n_embd)?;
5186
5187 let ssm = cfg.ssm.as_ref().expect("linear mixer requires ssm cfg");
5188 let d_state = ssm.state_size as usize;
5189 let num_k = ssm.group_count as usize;
5190 let num_v = ssm.time_step_rank as usize;
5191 let d_conv = ssm.conv_kernel as usize;
5192 let key_dim = d_state * num_k;
5193 let value_dim = d_state * num_v;
5194 let conv_dim = key_dim * 2 + value_dim;
5195 let gdn_scale = 1.0 / (d_state as f32).sqrt();
5196
5197 // ---- batched projections: one weight read for all T rows ----
5198 let qkv_mixed = e.matmul_pre(&la.wqkv, &hq, &hd, &xn, t)?;
5199 let z = e.matmul_pre(&la.wqkv_gate, &hq, &hd, &xn, t)?;
5200 let beta_raw = e.matmul_pre(&la.ssm_beta, &hq, &hd, &xn, t)?;
5201 let alpha = e.matmul_pre(&la.ssm_alpha, &hq, &hd, &xn, t)?;
5202 let beta_w = la.ssm_beta.out_features();
5203 let alpha_w = la.ssm_alpha.out_features();
5204 let qkv_w = la.wqkv.out_features();
5205
5206 // ---- per-row state chain through the b_n=1 serving kernels ----
5207 // 6-entry alternating pointer table expresses the ping-pong without a rebuild per
5208 // row: even rows scan s0 -> s1, odd rows s1 -> s0.
5209 let table_local: Option<CudaSlice<u64>> = match table_src {
5210 Some(_) => None,
5211 None => {
5212 let rl = cache.recur[il].as_ref().unwrap();
5213 let s = &e.gpu.stream();
5214 let (pc, _g0) = rl.conv_state.device_ptr(s);
5215 let (p0, _g1) = rl.ssm_state.device_ptr(s);
5216 let (p1, _g2) = rl.ssm_state_alt.device_ptr(s);
5217 Some(e.htod_u64(&[
5218 pc as u64, p0 as u64, p1 as u64, pc as u64, p1 as u64, p0 as u64,
5219 ])?)
5220 }
5221 };
5222 let (table, toff): (&CudaSlice<u64>, usize) = match table_src {
5223 Some((tb, off)) => (tb, off),
5224 None => (table_local.as_ref().unwrap(), 0),
5225 };
5226 let mut o_all = e.uninit(t * value_dim)?;
5227 let mut col_states: Option<Vec<(CudaSlice<f32>, CudaSlice<f32>)>> =
5228 if ckpt.is_some() && stash.is_none() && t >= 2 {
5229 Some(Vec::with_capacity(t - 1))
5230 } else {
5231 None
5232 };
5233 let mut stash = stash;
5234 // Per-row scratch reused across rows (uninit is cheap but not free at
5235 // 48 layers x T rows); row inputs/outputs pass as VIEWS into the packed
5236 // [T, ...] buffers — zero arithmetic-free copies in this loop.
5237 let mut conv_out = e.uninit(conv_dim)?;
5238 let mut q_l2 = e.uninit(value_dim)?;
5239 let mut k_l2 = e.uninit(value_dim)?;
5240 let mut v_gd = e.uninit(value_dim)?;
5241 let mut beta_b = e.uninit(num_v)?;
5242 let mut g_log = e.uninit(num_v)?;
5243 for r in 0..t {
5244 let base = toff + if r % 2 == 0 { 0 } else { 3 };
5245 let conv_view = table.slice(base..base + 1);
5246 let in_view = table.slice(base + 1..base + 2);
5247 let out_view = table.slice(base + 2..base + 3);
5248 e.ssm_conv1d_fused_decode_b_view(
5249 &qkv_mixed.slice(r * qkv_w..(r + 1) * qkv_w),
5250 &conv_view,
5251 la.ssm_conv1d.float_data(),
5252 &mut conv_out,
5253 conv_dim,
5254 d_conv,
5255 1,
5256 )?;
5257 e.gdn_prep_decode_b_view(
5258 &conv_out,
5259 &beta_raw.slice(r * beta_w..(r + 1) * beta_w),
5260 &alpha.slice(r * alpha_w..(r + 1) * alpha_w),
5261 la.ssm_dt.float_data(),
5262 la.ssm_a.float_data(),
5263 &mut q_l2,
5264 &mut k_l2,
5265 &mut v_gd,
5266 &mut beta_b,
5267 &mut g_log,
5268 d_state,
5269 num_v,
5270 num_k,
5271 key_dim,
5272 eps,
5273 conv_dim,
5274 1,
5275 )?;
5276 let mut o_row = o_all.slice_mut(r * value_dim..(r + 1) * value_dim);
5277 e.gdn_scan_s128_batched_view(
5278 &q_l2, &k_l2, &v_gd, &g_log, &beta_b, &in_view, &out_view, &mut o_row, num_v, 1,
5279 gdn_scale,
5280 )?;
5281 if r + 1 < t {
5282 // Row r's out buffer: even rows write s1 (the alt handle — no swaps ran),
5283 // odd rows write s0 — the same physical state the legacy post-swap
5284 // canonical clone read.
5285 let rl = cache.recur[il]
5286 .as_ref()
5287 .ok_or("qwen35 linear verify layer has no recurrent state")?;
5288 let ssm_src = if r % 2 == 0 {
5289 &rl.ssm_state_alt
5290 } else {
5291 &rl.ssm_state
5292 };
5293 match stash.as_mut() {
5294 Some((conv_slab, ssm_slab)) => {
5295 // BOTH stash reads go through the pointer table at run time: the
5296 // ssm handles ping-pong between rounds, and the ctx (with its
5297 // captured graphs) outlives the Cache — a fresh generation's
5298 // conv/ssm buffers land at new addresses that only the per-round
5299 // table refresh knows. A baked direct copy would read freed
5300 // memory (parity was the slice-3 smoke divergence; cache
5301 // lifetime is the cross-generation twin).
5302 e.copy_indirect_src_f32(
5303 &conv_view,
5304 conv_slab,
5305 r * conv_dim * (d_conv - 1),
5306 conv_dim * (d_conv - 1),
5307 )?;
5308 // The ssm handles PING-PONG between rounds: a captured direct
5309 // copy would bake the capture-time physical buffer and read the
5310 // wrong parity after any odd-vt round (the slice-3 smoke
5311 // divergence). Read the src address from row r's OUT table
5312 // entry at run time — the same entry the scan just wrote.
5313 e.copy_indirect_src_f32(
5314 &out_view,
5315 ssm_slab,
5316 r * d_state * d_state * num_v,
5317 d_state * d_state * num_v,
5318 )?;
5319 }
5320 None => {
5321 if let Some(states) = col_states.as_mut() {
5322 states.push((e.clone_dtod(&rl.conv_state)?, e.clone_dtod(ssm_src)?));
5323 }
5324 }
5325 }
5326 }
5327 }
5328 // ONE end-of-body parity swap (t odd) — the legacy loop swapped per row; the net
5329 // handle motion is identical and the device sequence never read the handles.
5330 if t % 2 == 1 {
5331 let rl = cache.recur[il].as_mut().unwrap();
5332 std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
5333 }
5334 if let (Some(checkpoint), Some(states)) = (ckpt.as_deref_mut(), col_states) {
5335 checkpoint.cols[il] = Some(states);
5336 }
5337
5338 // ---- batched gated norm + out-projection at m=T ----
5339 let mixed = if e.uses_q8_1_fast(&la.ssm_out) {
5340 let (gq, gd) = e.gated_rmsnorm_q8_1(
5341 &o_all,
5342 la.ssm_norm.float_data(),
5343 &z,
5344 d_state,
5345 t * num_v,
5346 eps,
5347 )?;
5348 let g0 = e.zeros(0)?;
5349 e.matmul_pre(&la.ssm_out, &gq, &gd, &g0, t)?
5350 } else {
5351 let mut gn = e.uninit(t * value_dim)?;
5352 e.gated_rmsnorm(
5353 &o_all,
5354 la.ssm_norm.float_data(),
5355 &z,
5356 &mut gn,
5357 d_state,
5358 t * num_v,
5359 eps,
5360 )?;
5361 e.matmul(&la.ssm_out, &gn, t)?
5362 };
5363
5364 // ---- residual add + post_attn_norm + FFN at m=T (serving dispatch verbatim) ----
5365 let pnorm = layer.post_attn_norm.float_data();
5366 let mut x1 = e.uninit(t * n_embd)?;
5367 let mut zn = e.uninit(t * n_embd)?;
5368 e.add_rms_norm(x, &mixed, pnorm, &mut x1, &mut zn, n_embd, t, eps)?;
5369 let ffn_out = match &layer.ffn {
5370 crate::hybrid::Ffn::Dense {
5371 ffn_gate,
5372 ffn_up,
5373 ffn_down,
5374 } => {
5375 assert!(
5376 self.cfg.m3.is_none(),
5377 "qwen35 t-parallel verify: M3 swigluoai FFN not yet batched"
5378 );
5379 let n_ff = ffn_gate.out_features();
5380 let (zq, zd) = e.quantize_q8_1(&zn, t, n_embd)?;
5381 let g = e.matmul_pre(ffn_gate, &zq, &zd, &zn, t)?;
5382 let u = e.matmul_pre(ffn_up, &zq, &zd, &zn, t)?;
5383 let mut act = e.uninit(t * n_ff)?;
5384 e.silu_mul(&g, &u, &mut act, t * n_ff)?;
5385 let (aq, ad) = e.quantize_q8_1(&act, t, n_ff)?;
5386 e.matmul_pre(ffn_down, &aq, &ad, &act, t)?
5387 }
5388 crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il_zq8(e, m, &zn, None, t, il as u16)?,
5389 };
5390 let mut x2 = e.uninit(t * n_embd)?;
5391 e.add(&x1, &ffn_out, &mut x2, t * n_embd)?;
5392 // dspark drafter tap (no-op when no sink armed): post-layer residual verify rows
5393 self.dflash_tap(e, cache, il, &x2, t)?;
5394 Ok(x2)
5395 }
5396
5397 /// PP-N STAGE SUBGRAPH of the verify trunk: layers `[lo, hi)` of `decode_step_t_core_stream`'s
5398 /// walk, verbatim. Enters with a MATERIALIZED `[T, n_embd]` residual (no pending fusion pair
5399 /// carried in from outside the range) and exits with the range's final residual materialized
5400 /// (the trailing add executed) — exactly the `decode_layers_eager(lo, hi)` contract, T rows
5401 /// instead of one.
5402 ///
5403 /// EXTRACTED (lane/pp2-spec 2026-08-06) rather than duplicated: `decode_step_t_core_stream` IS
5404 /// the single funnel every verify forward reaches, and its per-layer dispatch MIRRORING (norm
5405 /// fusion per layer, the t>=3/spec_m2 batched-linear window, the fused-q8 FFN chain, the
5406 /// decode-exact projections) is what makes verify bit-identical to eager decode. A second copy
5407 /// for the split arm is how those mirrors drift apart on the next lever. The unsplit body now
5408 /// calls this with `(0, n_layers)`, so the whole-trunk path and every stage range run the SAME
5409 /// code — there is no "split version" of the verify math.
5410 ///
5411 /// Bit-identity of a cut rests on the same kernel-check-pinned identity the eager arm's cut
5412 /// does — `add_rms_norm_q8_1 == add then rms_norm_q8_1` at nrows=T — because the ONLY thing a
5413 /// fence changes is that the cross-layer fusion carry breaks at `hi-1` and is re-materialized
5414 /// as an explicit `add`. `decode-batch-gate --mode ppspec` verifies end-to-end on real weights.
5415 #[allow(clippy::too_many_arguments)]
5416 fn verify_layers(
5417 &self,
5418 e: &Engine,
5419 mut x: CudaSlice<f32>,
5420 lo: usize,
5421 hi: usize,
5422 pos_d: &CudaSlice<i32>,
5423 pos0: usize,
5424 t: usize,
5425 cache: &mut Cache,
5426 mut ckpt: Option<&mut VerifyCkpt>,
5427 stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
5428 graphs: Option<&mut DsparkVerifyGraphs>,
5429 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5430 if self.cfg.step35.is_some() {
5431 if stream.is_some() {
5432 return Err(
5433 "step35 has no ROUND-STREAM verify arm (the device-counter _dc twins \
5434 cannot express the SWA offset KV view)"
5435 .into(),
5436 );
5437 }
5438 return self.step35_verify_batch_layers(e, x, lo, hi, pos0, t, cache);
5439 }
5440 if self.qwen35_serving_class() {
5441 return self.qwen35_verify_batch_layers(
5442 e,
5443 x,
5444 lo,
5445 hi,
5446 pos0,
5447 t,
5448 cache,
5449 ckpt.take(),
5450 stream,
5451 graphs,
5452 );
5453 }
5454 let n_embd = self.cfg.n_embd as usize;
5455 let eps = self.cfg.rms_eps;
5456 // CROSS-LAYER ADD+NORM FUSION (lane/vt-fixes fix 2, mirroring decode_step_h's
5457 // launch-arc form): layer il's post-FFN residual add (x2 = x1 + ffn_out) and layer
5458 // il+1's attn_norm(+quantize) are consecutive row-wise ops — ONE add_rms_norm_q8_1
5459 // launch at nrows=t does all three (bit-identity pinned by the T-row kernel-check
5460 // arms). Carry the un-added (x1, ffn_out) pair; the fused launch materializes x2 (the
5461 // residual the next layer needs) as its `res` output. Falls back to the separate add
5462 // when the next layer is off the fused-q8 path.
5463 let mut pending: Option<(CudaSlice<f32>, CudaSlice<f32>)> = None;
5464 for il in lo..hi {
5465 let layer = &self.layers[il];
5466 // DISPATCH-MIRRORED attn-input RMSNorm (FP-order lesson #8): eager decode fuses the
5467 // 1024-thread rms_norm_q8_1 ONLY when every mixer projection is q8_1-fast; layers with
5468 // Float projections (ssm_beta/ssm_alpha on layers 1/2/4 of the 9B NVFP4 GGUF) take the
5469 // UNFUSED 256-thread rms_norm. The verify norm must mirror that PER-LAYER choice —
5470 // blockDim changes the sum-of-squares reduce order, and the ULP shift amplifies through
5471 // the GDN recurrence into argmax flips (measured: 9B text prompt, 1 ULP at layer 2 ->
5472 // 2.3e-1 logit maxdiff at the head -> K=1..8 divergence at a 0.03-margin token).
5473 let mixer_fast = self.mixer_in_q8_1_fast(e, &layer.mixer);
5474 let norm_fused = std::env::var("MEMRA_NO_FUSE_NORMQ").is_err() && mixer_fast;
5475 // BATCHED EPILOGUE RE-FUSE (lane/vt-fixes fix 2, 2026-08-03): when the norm is
5476 // dispatch-fused AND every consumer of `h` reads only its q8_1 form (Full mixer:
5477 // projections only; Linear mixer: the batched arm — the per-column fallback needs
5478 // f32 h), emit the attn-input norm DIRECTLY as q8_1 via `rms_norm_q8_1` at nrows=t
5479 // (row-indexed kernel — the T-row launch is the per-row m=1 program, kernel-check
5480 // pins bit-identity vs rms_norm_decode -> quantize_q8_1). Kills the standalone
5481 // quantize launch(es) + the f32 h HBM round-trip that decode never pays.
5482 // step35 (Full mixer) is the third case that needs f32 `h`: its verify arm is a
5483 // per-ROW replay of the eager decode mixer, whose `pre_q` contract is a single row —
5484 // a T-row q8_1 pair cannot be handed to it, and re-deriving per-row q8_1 from the f32
5485 // rows is exactly the dispatch being mirrored. Keep step35 on the unfused arm.
5486 let lin_q8_only = match &layer.mixer {
5487 Mixer::Linear(la) => {
5488 (t >= 3 || (t == 2 && spec_m2())) && e.uses_q8_1_fast(&la.ssm_out)
5489 }
5490 Mixer::Full(_) if self.cfg.step35.is_some() => false,
5491 _ => true,
5492 };
5493 // NOTE decode.rs's take()-first lesson: take the pending pair BEFORE branching so
5494 // a non-fused layer still performs the residual add.
5495 let taken = pending.take();
5496 let (h, h_q8) = if norm_fused && lin_q8_only {
5497 let pair = match taken {
5498 // fused add + attn_norm + q8_1: ONE launch resolves the carried residual
5499 // AND emits this layer's mixer input pre-quantized. res -> x2 (= new x).
5500 Some((x1p, f1p)) => {
5501 let mut x2 = vbuf(e, t * n_embd)?; // fully written (res output)
5502 let p = e.add_rms_norm_q8_1(
5503 &x1p,
5504 &f1p,
5505 layer.attn_norm.float_data(),
5506 &mut x2,
5507 n_embd,
5508 t,
5509 eps,
5510 )?;
5511 x = x2;
5512 p
5513 }
5514 None => e.rms_norm_q8_1(&x, layer.attn_norm.float_data(), n_embd, t, eps)?,
5515 };
5516 (e.zeros(0)?, Some(pair)) // h unused on this path (q8-only consumers)
5517 } else {
5518 if let Some((x1p, f1p)) = taken {
5519 let mut x2 = vbuf(e, t * n_embd)?; // fully written by add
5520 e.add(&x1p, &f1p, &mut x2, t * n_embd)?;
5521 x = x2;
5522 }
5523 let mut h = vbuf(e, t * n_embd)?; // fully written by either rms_norm arm
5524 if norm_fused {
5525 e.rms_norm_decode(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
5526 } else {
5527 e.rms_norm(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
5528 }
5529 (h, None)
5530 };
5531 let h_q8_ref = h_q8.as_ref().map(|(q, d)| (q, d));
5532
5533 let mixed = match &layer.mixer {
5534 Mixer::Full(fa) => self.full_attn_verify(
5535 e,
5536 fa,
5537 &h,
5538 h_q8_ref,
5539 pos_d,
5540 t,
5541 cache,
5542 il,
5543 stream.map(|(_, c)| c),
5544 )?,
5545 Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
5546 Mixer::Linear(la) => {
5547 // BATCHED linear verify (2026-07-03, the MTP-profit lever): one T-token pass —
5548 // batched projections (weight read ONCE, hits the m=2-4 weight-resident matvec),
5549 // carried-state conv (ssm_conv1d_tm_state), GDN prep on the prefill kernels, and
5550 // ONE gdn_scan whose internal sequential t-loop is the SAME recurrence as T
5551 // chained T=1 steps (bit-identical). Falls back to the sequential per-column
5552 // chain when T < d_conv-1 (conv ring update needs T >= pad) — or when ANY
5553 // projection is off the q8_1 fast path: matmul_decode_exact would route a Float
5554 // tensor to cuBLAS at m=t (different FP accumulation than eager's per-token
5555 // GEMV), so mixed-dtype layers stay on the eager-identical per-column chain.
5556 // MEMRA_SPEC_M2 (lane/spec-m2): the t==2 batch rides the same arm — the conv
5557 // wrapper handles t<pad with a pure-copy ring rebuild; see spec_m2() header.
5558 if (t >= 3 || (t == 2 && spec_m2()))
5559 && mixer_fast
5560 && e.uses_q8_1_fast(&la.ssm_out)
5561 {
5562 let want = ckpt.is_some();
5563 let (out, stash) =
5564 self.linear_attn_verify_t(e, la, &h, h_q8_ref, t, cache, il, want)?;
5565 if let (Some(ck), Some(st)) = (ckpt.as_deref_mut(), stash) {
5566 ck.gdn[il] = Some(st);
5567 }
5568 out
5569 } else {
5570 let mut out = vbuf(e, t * n_embd)?; // every col written by copy_into
5571 let mut col_states: Option<Vec<(CudaSlice<f32>, CudaSlice<f32>)>> =
5572 if ckpt.is_some() && t >= 2 {
5573 Some(Vec::with_capacity(t - 1))
5574 } else {
5575 None
5576 };
5577 for col in 0..t {
5578 let mut h_col = vbuf(e, n_embd)?; // fully written by copy_view_into
5579 let src = h.slice(col * n_embd..(col + 1) * n_embd);
5580 e.copy_view_into(&mut h_col, 0, &src, n_embd)?;
5581 let m_col = self.linear_attn_decode(e, la, &h_col, cache, il)?;
5582 e.copy_into(&mut out, col * n_embd, &m_col, n_embd)?;
5583 // REPLAY-FREE ckpt: clone the chain's ACTUAL state after this column
5584 // (pure dtod — cannot change any computed value). Last column skipped:
5585 // rebuild targets are j <= t-1 columns.
5586 if let Some(cs) = col_states.as_mut() {
5587 if col + 1 < t {
5588 let rl = cache.recur[il].as_ref().unwrap();
5589 cs.push((
5590 e.clone_dtod(&rl.conv_state)?,
5591 e.clone_dtod(&rl.ssm_state)?,
5592 ));
5593 }
5594 }
5595 }
5596 if let (Some(ck), Some(cs)) = (ckpt.as_deref_mut(), col_states) {
5597 // ReplaySSM-assessment instrumentation (2026-07-30): the
5598 // per-column clones are the only true state snapshots left in
5599 // the verify (the batched path stashes INPUTS and replays).
5600 if std::env::var("MEMRA_SPEC_STATS").as_deref() == Ok("1") {
5601 static ONCE: std::sync::Once = std::sync::Once::new();
5602 let bytes: usize =
5603 cs.iter().map(|(c, s)| (c.len() + s.len()) * 4).sum();
5604 ONCE.call_once(|| eprintln!(
5605 "[verify-ckpt] per-column layer il={il}: {} clones, {:.2} MB/layer/round",
5606 cs.len(), bytes as f64 / 1e6));
5607 }
5608 ck.cols[il] = Some(cs);
5609 }
5610 out
5611 }
5612 }
5613 };
5614
5615 // DISPATCH-MIRRORED post-attn norm: eager residual_norm_ffn fuses add+norm+quant
5616 // (1024-thread add_rms_norm_q8_1) only for Dense FFNs whose gate+up are q8_1-fast;
5617 // otherwise (and for MoE) it runs the 256-thread fused add_rms_norm. Mirror per layer.
5618 let ffn_fuse = match &layer.ffn {
5619 crate::hybrid::Ffn::Dense {
5620 ffn_gate, ffn_up, ..
5621 } => {
5622 std::env::var("MEMRA_NO_FUSE_NORMQ").is_err()
5623 && e.uses_q8_1_fast(ffn_gate)
5624 && e.uses_q8_1_fast(ffn_up)
5625 }
5626 crate::hybrid::Ffn::Moe(_) => false,
5627 };
5628 // BATCHED EPILOGUE RE-FUSE (lane/vt-fixes fix 2): on the ffn_fuse path (Dense,
5629 // gate+up q8_1-fast, non-M3) the FFN input is emitted DIRECTLY as q8_1 by ONE
5630 // add_rms_norm_q8_1 launch at nrows=t (row-indexed kernel: the T-row launch is the
5631 // per-row m=1 program; kernel-check pins bit-identity vs the unfused
5632 // add_f32 -> rms_norm_decode -> quantize_q8_1 chain at T=2/4/5/8) — replacing the
5633 // add + rms_norm_decode launches AND the dual/singles' internal re-quantize.
5634 // M3's swigluoai must keep the f32 chain (the fused SwiGLU epilogue encodes plain
5635 // SiLU), mirroring residual_norm_ffn's m3 guard on the decode path.
5636 // step35: same guard per LAYER. A dense FFN's clamp is the SHEXP array (upstream's
5637 // one build_ffn serves dense + shared expert, llama-graph.cpp:1751), and verify MUST
5638 // mirror decode's dispatch or spec self-consistency fails.
5639 let dense_lim = self.cfg.clamp_shexp_at(il as u32);
5640 let fuse_q8 = ffn_fuse && self.cfg.m3.is_none() && dense_lim.is_none();
5641 let mut x1 = vbuf(e, t * n_embd)?; // fully written by add / add_rms_norm*
5642 let mut z = e.zeros(0)?; // replaced below on the unfused arms
5643 let z_q8 = if fuse_q8 {
5644 Some(e.add_rms_norm_q8_1(
5645 &x,
5646 &mixed,
5647 layer.post_attn_norm.float_data(),
5648 &mut x1,
5649 n_embd,
5650 t,
5651 eps,
5652 )?)
5653 } else {
5654 let mut zf = vbuf(e, t * n_embd)?; // fully written by rms_norm_decode / add_rms_norm
5655 if ffn_fuse {
5656 e.add(&x, &mixed, &mut x1, t * n_embd)?;
5657 e.rms_norm_decode(
5658 &x1,
5659 layer.post_attn_norm.float_data(),
5660 &mut zf,
5661 n_embd,
5662 t,
5663 eps,
5664 )?;
5665 } else {
5666 e.add_rms_norm(
5667 &x,
5668 &mixed,
5669 layer.post_attn_norm.float_data(),
5670 &mut x1,
5671 &mut zf,
5672 n_embd,
5673 t,
5674 eps,
5675 )?;
5676 }
5677 z = zf;
5678 None
5679 };
5680 // DECODE-EXACT FFN projections: force MMVQ for gate/up/down at any T to match the
5681 // T=1 decode FP accumulation order. At T>=5 the generic matmul/matmul_pre falls to dp4a
5682 // (128-thread, different FP sum order). At T=2-4 the batched MMVQ is already bit-identical.
5683 let ffn_out = match &layer.ffn {
5684 crate::hybrid::Ffn::Dense {
5685 ffn_gate,
5686 ffn_up,
5687 ffn_down,
5688 } => {
5689 let n_ff = ffn_gate.out_features();
5690 if let Some((zq, zd)) = z_q8.as_ref() {
5691 // FUSED CHAIN (fix 2): pre-quantized z feeds the projections; the SwiGLU
5692 // epilogue emits act pre-quantized for ffn_down (silu_mul_scaled_q8_1,
5693 // bit-identical to silu_mul + quantize — kernel-check-pinned) with the
5694 // NVFP4 macro-scales folded (deferred-scale dual: y*s inline == the
5695 // scale_inplace store, value-exact) — the exact m=1 decode epilogue
5696 // structure at nrows=t.
5697 let pair =
5698 match e.matmul_decode_exact_dual_pre(ffn_gate, ffn_up, zq, zd, t)? {
5699 Some(((g, gs), (u, us))) => Some((g, gs, u, us)),
5700 None => None,
5701 };
5702 let (gate, gs, up, us) = match pair {
5703 Some(x4) => x4,
5704 None => (
5705 e.matmul_decode_exact_pre(ffn_gate, zq, zd, t)?,
5706 1.0, // scale already applied inside _pre
5707 e.matmul_decode_exact_pre(ffn_up, zq, zd, t)?,
5708 1.0,
5709 ),
5710 };
5711 if e.uses_q8_1_fast(ffn_down) {
5712 let (aq, ad) = e.silu_mul_scaled_q8_1(&gate, &up, gs, us, t * n_ff)?;
5713 e.matmul_decode_exact_pre(ffn_down, &aq, &ad, t)?
5714 } else {
5715 let mut act = vbuf(e, t * n_ff)?;
5716 e.silu_mul_scaled(&gate, &up, gs, us, &mut act, t * n_ff)?;
5717 e.matmul_decode_exact(ffn_down, &act, t)?
5718 }
5719 } else {
5720 // UNFUSED (pre-fix) chain — MoE-adjacent/M3/off-fast layers, unchanged.
5721 // DUAL gate+up batched twin (lane/verify-economics, 2026-08-02): one launch
5722 // for the pair at t=2..8 — bit-identical per (tensor,token,row) to the two
5723 // singles (kernel-check pins bitwise; MEMRA_SPEC_DUAL_T=0 reverts). None
5724 // (non-NVFP4 / t outside the tier / seam off) -> the two singles, unchanged.
5725 let (gate, up) =
5726 match e.matmul_decode_exact_dual(ffn_gate, ffn_up, &z, t)? {
5727 Some(pair) => pair,
5728 None => (
5729 e.matmul_decode_exact(ffn_gate, &z, t)?,
5730 e.matmul_decode_exact(ffn_up, &z, t)?,
5731 ),
5732 };
5733 let mut act = vbuf(e, t * n_ff)?; // fully written by ffn_act_lim
5734 Self::ffn_act_lim(
5735 e,
5736 &self.cfg,
5737 &gate,
5738 &up,
5739 1.0,
5740 1.0,
5741 dense_lim,
5742 &mut act,
5743 t * n_ff,
5744 )?;
5745 e.matmul_decode_exact(ffn_down, &act, t)?
5746 }
5747 }
5748 crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il(e, m, &z, t, il as u16)?,
5749 };
5750 // CROSS-LAYER fusion: defer this layer's post-FFN residual add — the next layer's
5751 // fused-q8 attn norm folds it in (add_rms_norm_q8_1 == add; rms_norm; quantize,
5752 // kernel-check-pinned at nrows=T). Non-fused next layers add explicitly above.
5753 pending = Some((x1, ffn_out));
5754 }
5755 // RANGE's final add (no next norm INSIDE the range to fuse with; for the
5756 // whole-trunk call that is the last layer, whose next norm is output_norm — f32-out).
5757 if let Some((x1p, f1p)) = pending.take() {
5758 let mut x2 = vbuf(e, t * n_embd)?; // fully written by add
5759 e.add(&x1p, &f1p, &mut x2, t * n_embd)?;
5760 x = x2;
5761 }
5762 Ok(x)
5763 }
5764 /// BATCHED linear-attn verify (T=K+1): the whole layer in ~10 launches instead of T x the
5765 /// T=1 decode chain (T x ~12 launches + T weight reads of the four projections). The GDN
5766 /// recurrence itself is inherently sequential — gdn_scan_s128 runs its internal t-loop with
5767 /// the SAME per-token math as chained T=1 calls (bit-identical state evolution); everything
5768 /// around it (projections, conv, prep, gated norm, out-proj) batches. Advances conv ring +
5769 /// ssm state exactly like T sequential decode steps.
5770 /// `want_stash`: additionally RETAIN the gdn-scan inputs (pure buffer keep-alives, zero extra
5771 /// kernels) so a partial accept can rebuild the state after any column prefix (REPLAY-FREE).
5772 #[allow(clippy::too_many_arguments)]
5773 fn linear_attn_verify_t(
5774 &self,
5775 e: &Engine,
5776 la: &LinearAttnLayer,
5777 h: &CudaSlice<f32>,
5778 h_q8: Option<(&CudaSlice<i8>, &CudaSlice<f32>)>,
5779 t: usize,
5780 cache: &mut Cache,
5781 il: usize,
5782 want_stash: bool,
5783 ) -> Result<(CudaSlice<f32>, Option<GdnStash>), Box<dyn std::error::Error>> {
5784 let cfg = &self.cfg;
5785 let ssm = cfg.ssm.as_ref().unwrap();
5786 let d_state = ssm.state_size as usize;
5787 let num_k = ssm.group_count as usize;
5788 let num_v = ssm.time_step_rank as usize;
5789 let d_conv = ssm.conv_kernel as usize;
5790 let key_dim = d_state * num_k;
5791 let conv_dim = key_dim * 2 + d_state * num_v;
5792 let eps = cfg.rms_eps;
5793 let scale = 1.0 / (d_state as f32).sqrt();
5794
5795 // DECODE-EXACT projections: matmul_decode_exact forces the MMVQ (warp-per-row, 32-thread)
5796 // accumulation order for EVERY m, matching the T=1 decode path bit-for-bit. The generic
5797 // `matmul` at m>=5 falls to dp4a (128-thread, two-level reduce) which has a different FP
5798 // sum order — ULP differences propagate through gdn_scan and flip argmax on the 27B.
5799 // Q8 TRUNK-FUSION at T=1 (35B: wqkv+wqkv_gate both Q8_0): one fused2 launch, bit-identical
5800 // per (tensor,row) to the two m=1 MMVQ dispatches below — decode-exact contract holds.
5801 // VERIFY-TIER TRUNK FUSION (MEMRA_SPEC_FUSED_T, t=2-4): quantize h ONCE for every
5802 // fused-eligible same-input Q8_0 pair of this layer (35B wqkv+wqkv_gate; 9B
5803 // ssm_beta+ssm_alpha) — each fused2 batched launch then replaces two decode-exact
5804 // calls (each of which re-quantizes the same h + runs its own _b2/_b4 launch).
5805 // Bit-identical per (tensor,token,row) — see spec_fused_t().
5806 // BATCHED EPILOGUE RE-FUSE (lane/vt-fixes fix 2): `h_q8` = the attn-input norm emitted
5807 // directly as q8_1 by the caller's fused rms_norm_q8_1 (bit-identical to the unfused
5808 // chain, kernel-check-pinned). When present it REPLACES the standalone quantize below
5809 // and feeds every projection; the caller guaranteed all four input projections are
5810 // q8_1-fast. When absent, the old shared-quantize (fused-t window) stands.
5811 let h_q8_t = if h_q8.is_none()
5812 && spec_fused_t()
5813 && (2..=4).contains(&t)
5814 && ((e.uses_q8_1_fast(&la.wqkv) && e.uses_q8_1_fast(&la.wqkv_gate))
5815 || (e.uses_q8_1_fast(&la.ssm_beta) && e.uses_q8_1_fast(&la.ssm_alpha)))
5816 {
5817 Some(e.quantize_q8_1(h, t, cfg.n_embd as usize)?)
5818 } else {
5819 None
5820 };
5821 // one view: the caller's fused-norm q8 or this fn's own shared quantize.
5822 let hq8_any: Option<(&CudaSlice<i8>, &CudaSlice<f32>)> =
5823 h_q8.or(h_q8_t.as_ref().map(|(q, d)| (q, d)));
5824 let (qkv_mixed, z) = {
5825 let mut fused = None;
5826 if t == 1 && e.uses_q8_1_fast(&la.wqkv) && e.uses_q8_1_fast(&la.wqkv_gate) {
5827 let (hq, hd) = e.quantize_q8_1(h, 1, cfg.n_embd as usize)?;
5828 fused = e.matmul_q8_fused2(&la.wqkv, &la.wqkv_gate, &hq, &hd)?;
5829 } else if let Some((hq, hd)) = hq8_any {
5830 if spec_fused_t() && (2..=4).contains(&t) {
5831 fused = e.matmul_q8_fused2_t(&la.wqkv, &la.wqkv_gate, hq, hd, t)?;
5832 }
5833 }
5834 match (fused, hq8_any) {
5835 (Some(pair), _) => pair,
5836 (None, Some((hq, hd))) if h_q8.is_some() => (
5837 e.matmul_decode_exact_pre(&la.wqkv, hq, hd, t)?,
5838 e.matmul_decode_exact_pre(&la.wqkv_gate, hq, hd, t)?,
5839 ),
5840 (None, _) => (
5841 e.matmul_decode_exact(&la.wqkv, h, t)?,
5842 e.matmul_decode_exact(&la.wqkv_gate, h, t)?,
5843 ),
5844 }
5845 };
5846 // beta+alpha DUAL at T=1 (75% of p3 rounds run T=1 verify — p-min chain cuts): the dual
5847 // mr2 kernel is bit-identical per element to the m=1 MMVQ matmul_decode_exact dispatches
5848 // (same warp-per-row body, blockIdx.y picks the weight), so the decode-exact contract
5849 // holds; the run-spec battery is the arbiter. T>1 keeps the per-tensor decode-exact path.
5850 let (beta_raw, alpha) = if t == 1 {
5851 let (hq, hd) = e.quantize_q8_1(h, 1, cfg.n_embd as usize)?;
5852 match e.matmul_pre_dual_noscale(&la.ssm_beta, &la.ssm_alpha, &hq, &hd, 1)? {
5853 Some(((mut b, bs), (mut a, as_))) => {
5854 if bs != 1.0 {
5855 e.scale_inplace(&mut b, bs, la.ssm_beta.out_features())?;
5856 }
5857 if as_ != 1.0 {
5858 e.scale_inplace(&mut a, as_, la.ssm_alpha.out_features())?;
5859 }
5860 (b, a)
5861 }
5862 // Q8_0 fused2 twin (9B stores beta/alpha as Q8_0): DISPATCH-MIRRORS the eager
5863 // decode's beta_alpha closure — the fused body is qmatvec_q8_0_mmvq verbatim,
5864 // bit-identical per row (kernel-check rel=0.00e0 gate), so decode==verify holds.
5865 None => match e.matmul_q8_fused2(&la.ssm_beta, &la.ssm_alpha, &hq, &hd)? {
5866 Some((b, a)) => (b, a),
5867 None => (
5868 e.matmul_decode_exact(&la.ssm_beta, h, 1)?,
5869 e.matmul_decode_exact(&la.ssm_alpha, h, 1)?,
5870 ),
5871 },
5872 }
5873 } else {
5874 // fused-t twin (9B stores beta/alpha as Q8_0): same shared-quantize + one launch
5875 // contract as the wqkv pair above; 35B beta/alpha are Float -> None -> fallback.
5876 let mut nvfp4_fused = None;
5877 let mut q8_fused = None;
5878 if let Some((hq, hd)) = hq8_any {
5879 if t == 3 && std::env::var("MEMRA_NVFP4_AUX_DUAL").as_deref() != Ok("0") {
5880 nvfp4_fused =
5881 e.matmul_decode_exact_dual_pre(&la.ssm_beta, &la.ssm_alpha, hq, hd, t)?;
5882 if nvfp4_fused.is_some() && std::env::var("MEMRA_DEBUG").is_ok() {
5883 static ONCE: std::sync::Once = std::sync::Once::new();
5884 ONCE.call_once(|| {
5885 eprintln!("[memra] NVFP4 beta+alpha batched aux dual ENGAGED (t={t})")
5886 });
5887 }
5888 }
5889 if nvfp4_fused.is_none() && spec_fused_t() && (2..=4).contains(&t) {
5890 q8_fused = e.matmul_q8_fused2_t(&la.ssm_beta, &la.ssm_alpha, hq, hd, t)?;
5891 }
5892 }
5893 if let Some(((mut b, bs), (mut a, as_))) = nvfp4_fused {
5894 if bs != 1.0 {
5895 e.scale_inplace(&mut b, bs, t * la.ssm_beta.out_features())?;
5896 }
5897 if as_ != 1.0 {
5898 e.scale_inplace(&mut a, as_, t * la.ssm_alpha.out_features())?;
5899 }
5900 (b, a)
5901 } else if let Some(pair) = q8_fused {
5902 pair
5903 } else {
5904 match hq8_any {
5905 Some((hq, hd)) if h_q8.is_some() => (
5906 e.matmul_decode_exact_pre(&la.ssm_beta, hq, hd, t)?,
5907 e.matmul_decode_exact_pre(&la.ssm_alpha, hq, hd, t)?,
5908 ),
5909 _ => (
5910 e.matmul_decode_exact(&la.ssm_beta, h, t)?,
5911 e.matmul_decode_exact(&la.ssm_alpha, h, t)?,
5912 ),
5913 }
5914 }
5915 };
5916
5917 // conv with CARRIED state + ring roll (T >= pad rides the input-column update kernel;
5918 // T < pad — the MEMRA_SPEC_M2 t=2 arm — rolls via the pure-copy ring rebuild).
5919 let rl = cache.recur[il].as_mut().unwrap();
5920 let mut conv_out = e.uninit(conv_dim * t)?;
5921 e.ssm_conv1d_tm_state(
5922 &qkv_mixed,
5923 &mut rl.conv_state,
5924 la.ssm_conv1d.float_data(),
5925 &mut conv_out,
5926 conv_dim,
5927 t,
5928 d_conv,
5929 )?;
5930
5931 // GDN prep via the prefill kernels (repack + L2 + sigmoid + glog), T-wide.
5932 let mut q_g = e.uninit(d_state * num_v * t)?;
5933 let mut k_g = e.uninit(d_state * num_v * t)?;
5934 let mut v_g = e.uninit(d_state * num_v * t)?;
5935 e.qkv_to_gdn_repack(
5936 &conv_out, &mut q_g, &mut k_g, &mut v_g, d_state, num_v, num_k, key_dim, t,
5937 )?;
5938 let mut q_l2 = e.uninit(d_state * num_v * t)?;
5939 e.l2_norm_decode(&q_g, &mut q_l2, d_state, num_v * t, eps)?;
5940 let mut k_l2 = e.uninit(d_state * num_v * t)?;
5941 e.l2_norm_decode(&k_g, &mut k_l2, d_state, num_v * t, eps)?;
5942 let mut beta = e.uninit(t * num_v)?;
5943 e.sigmoid(&beta_raw, &mut beta, t * num_v)?;
5944 let mut g_log = e.uninit(t * num_v)?;
5945 e.gdn_glog(
5946 &alpha,
5947 la.ssm_dt.float_data(),
5948 la.ssm_a.float_data(),
5949 &mut g_log,
5950 num_v,
5951 t,
5952 )?;
5953
5954 // ONE gdn_scan over T tokens from the carried state (internal sequential loop ==
5955 // T chained T=1 steps). Ping-pong the resident buffers like eager decode.
5956 let mut o = e.uninit(d_state * num_v * t)?;
5957 {
5958 let crate::cache::RecurLayer {
5959 ssm_state,
5960 ssm_state_alt,
5961 ..
5962 } = rl;
5963 e.gdn_scan_s128(
5964 &q_l2,
5965 &k_l2,
5966 &v_g,
5967 &g_log,
5968 &beta,
5969 ssm_state,
5970 ssm_state_alt,
5971 &mut o,
5972 num_v,
5973 t,
5974 scale,
5975 )?;
5976 }
5977 std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
5978
5979 // gated RMSNorm + out projection, T-wide. FUSED-QUANTIZE ARM (lane/vt-fixes fix 2,
5980 // mirroring the T=1 decode's launch-arc form): when ssm_out rides the q8_1 fast path,
5981 // emit q8_1 straight from the gated norm at nrows=num_v*t (row-indexed kernel, the
5982 // T-wide launch is the per-row program; kernel-check pins bit-identity vs
5983 // gated_rmsnorm -> quantize_q8_1 at T=1 and T=5) and feed the decode-exact dispatch
5984 // pre-quantized — one launch replaces norm + quantize. Fallback = the f32 chain.
5985 let out = if e.uses_q8_1_fast(&la.ssm_out) {
5986 let (gq, gd) =
5987 e.gated_rmsnorm_q8_1(&o, la.ssm_norm.float_data(), &z, d_state, num_v * t, eps)?;
5988 e.matmul_decode_exact_pre(&la.ssm_out, &gq, &gd, t)?
5989 } else {
5990 let mut gn = e.uninit(d_state * num_v * t)?;
5991 e.gated_rmsnorm(
5992 &o,
5993 la.ssm_norm.float_data(),
5994 &z,
5995 &mut gn,
5996 d_state,
5997 num_v * t,
5998 eps,
5999 )?;
6000 // DECODE-EXACT out-projection: same MMVQ path as the T=1 decode (ssm_out at m>=5
6001 // would fall to dp4a with a different FP reduction order — same class of bug as
6002 // the input projs).
6003 e.matmul_decode_exact(&la.ssm_out, &gn, t)?
6004 };
6005 let stash = if want_stash {
6006 Some(GdnStash {
6007 qkv_mixed,
6008 q_l2,
6009 k_l2,
6010 v_g,
6011 g_log,
6012 beta,
6013 })
6014 } else {
6015 None
6016 };
6017 Ok((out, stash))
6018 }
6019
6020 /// REPLAY-FREE partial-accept commit (2026-07-03): make the cache state == "committed through
6021 /// the first `j` verify columns" WITHOUT the legacy rollback + duplicate trunk replay.
6022 /// - Full-attn KV: truncate len to snapshot + j. The verify's appended rows for those columns
6023 /// are bit-identical to what an eager T=1 chain writes (the decode-exact contract the
6024 /// verify-probe gates), so keeping them == replaying them.
6025 /// - Linear layers, batched path: rebuild the conv ring by PURE COPIES (ring holds raw input
6026 /// columns) and the ssm state by a prefix re-run of the SAME gdn_scan kernel (t=j) from the
6027 /// snapshot state over the stash's identical inputs — the kernel's t-loop carries state in
6028 /// registers and writes it once at the end, so iterations 0..j-1 are independent of T:
6029 /// bit-identical to the verify's own state after j tokens == the eager chain state.
6030 /// - Linear layers, per-column path: restore the cloned actual state after column j-1.
6031 /// Caller guarantees 1 <= j <= t-1 (j==0 rounds take the legacy rollback; j==t is full accept).
6032 fn commit_verified_prefix(
6033 &self,
6034 e: &Engine,
6035 cache: &mut Cache,
6036 snap: &crate::cache::CacheSnapshot,
6037 ckpt: &VerifyCkpt,
6038 j: usize,
6039 kv_lens_done: bool,
6040 dev_j: Option<(&CudaSlice<u32>, usize, usize)>,
6041 ) -> Result<(), Box<dyn std::error::Error>> {
6042 let cfg = &self.cfg;
6043 let ssm = cfg.ssm.as_ref().unwrap();
6044 let d_state = ssm.state_size as usize;
6045 let num_k = ssm.group_count as usize;
6046 let num_v = ssm.time_step_rank as usize;
6047 let d_conv = ssm.conv_kernel as usize;
6048 let conv_dim = d_state * num_k * 2 + d_state * num_v;
6049 let scale = 1.0 / (d_state as f32).sqrt();
6050 // Engine-bundle slice 1 (DSF-ROUNDCOST-20260820 §1.1): the per-column-arm restores
6051 // are 2 tiny D2D copies per linear layer (~96 dispatches/partial round on the q38
6052 // route). When every cols-arm layer shares uniform state sizes (single ssm cfg —
6053 // always true today), batch them into two `copy_batch_uniform_f32` launches. Bytes,
6054 // buffers and stream order are identical to the per-layer memcpy sequence; the
6055 // kernel-rebuild (gdn-stash) arm below is untouched. MEMRA_STATE_COPY_BATCH=0 reverts.
6056 let mut batched_cols = false;
6057 if state_copy_batch_on() && dev_j.is_none() {
6058 use cudarc::driver::DevicePtr;
6059 let s = &e.gpu.stream();
6060 let mut conv_pairs: Vec<(u64, u64)> = Vec::new();
6061 let mut ssm_pairs: Vec<(u64, u64)> = Vec::new();
6062 let (mut conv_words, mut ssm_words) = (0usize, 0usize);
6063 let mut uniform = true;
6064 for il in 0..self.layers.len() {
6065 let Some(rl) = cache.recur[il].as_ref() else {
6066 continue;
6067 };
6068 if ckpt.gdn[il].is_some() {
6069 continue; // kernel-rebuild arm restores below, per layer
6070 }
6071 let Some(cols) = &ckpt.cols[il] else {
6072 continue; // missing-ckpt error surfaces in the main loop
6073 };
6074 let (c, st) = &cols[j - 1];
6075 if conv_pairs.is_empty() {
6076 conv_words = c.len();
6077 ssm_words = st.len();
6078 } else if c.len() != conv_words || st.len() != ssm_words {
6079 uniform = false;
6080 break;
6081 }
6082 let (pc, _g0) = c.device_ptr(s);
6083 let (dc, _g1) = rl.conv_state.device_ptr(s);
6084 let (ps, _g2) = st.device_ptr(s);
6085 let (ds, _g3) = rl.ssm_state.device_ptr(s);
6086 conv_pairs.push((pc as u64, dc as u64));
6087 ssm_pairs.push((ps as u64, ds as u64));
6088 }
6089 if uniform && !conv_pairs.is_empty() {
6090 let n = conv_pairs.len();
6091 let mut t = vec![0u64; 2 * n];
6092 for (k, &(src, dst)) in conv_pairs.iter().enumerate() {
6093 t[k] = src;
6094 t[n + k] = dst;
6095 }
6096 let conv_t = e.htod_u64(&t)?;
6097 for (k, &(src, dst)) in ssm_pairs.iter().enumerate() {
6098 t[k] = src;
6099 t[n + k] = dst;
6100 }
6101 let ssm_t = e.htod_u64(&t)?;
6102 e.copy_batch_uniform_f32(&conv_t, n, conv_words)?;
6103 e.copy_batch_uniform_f32(&ssm_t, n, ssm_words)?;
6104 batched_cols = true;
6105 }
6106 }
6107 for il in 0..self.layers.len() {
6108 if let (Some(kvl), Some(saved)) = (cache.kv[il].as_mut(), snap.kv_len[il]) {
6109 kvl.len = saved + j;
6110 // devacc 3a: spec_rollback_kv already wrote len_d on-device (same value).
6111 if !kv_lens_done {
6112 e.set_i32_one(&mut kvl.len_d, kvl.len as i32)?;
6113 }
6114 }
6115 if let Some(rl) = cache.recur[il].as_mut() {
6116 if let Some(st) = &ckpt.gdn[il] {
6117 let ring_old = snap.conv[il].as_ref().expect("snapshot missing conv");
6118 let state_in = snap.ssm[il].as_ref().expect("snapshot missing ssm");
6119 if let Some((acc, base, t_v)) = dev_j {
6120 // 3b: j read on-device (_dc twins, same bodies; full accept early-exits).
6121 e.ssm_conv_ring_rebuild_dc(
6122 &st.qkv_mixed,
6123 ring_old,
6124 &mut rl.conv_state,
6125 conv_dim,
6126 acc,
6127 base,
6128 t_v,
6129 d_conv,
6130 )?;
6131 let mut o = e.uninit(d_state * num_v * j.max(1))?;
6132 e.gdn_scan_s128_dc(
6133 &st.q_l2,
6134 &st.k_l2,
6135 &st.v_g,
6136 &st.g_log,
6137 &st.beta,
6138 state_in,
6139 &mut rl.ssm_state,
6140 &mut o,
6141 num_v,
6142 acc,
6143 base,
6144 t_v,
6145 scale,
6146 )?;
6147 } else {
6148 e.ssm_conv_ring_rebuild(
6149 &st.qkv_mixed,
6150 ring_old,
6151 &mut rl.conv_state,
6152 conv_dim,
6153 j,
6154 d_conv,
6155 )?;
6156 let mut o = e.uninit(d_state * num_v * j)?; // scan output, discarded
6157 e.gdn_scan_s128(
6158 &st.q_l2,
6159 &st.k_l2,
6160 &st.v_g,
6161 &st.g_log,
6162 &st.beta,
6163 state_in,
6164 &mut rl.ssm_state,
6165 &mut o,
6166 num_v,
6167 j,
6168 scale,
6169 )?;
6170 }
6171 } else if let Some(cols) = &ckpt.cols[il] {
6172 if !batched_cols {
6173 let (c, s) = &cols[j - 1];
6174 e.copy_into(&mut rl.conv_state, 0, c, c.len())?;
6175 e.copy_into(&mut rl.ssm_state, 0, s, s.len())?;
6176 }
6177 } else {
6178 return Err(
6179 "commit_verified_prefix: verify ckpt missing for linear layer".into(),
6180 );
6181 }
6182 }
6183 }
6184 cache.pos = snap.pos + j;
6185 Ok(())
6186 }
6187
6188 /// ROUND-STREAM: recur restore with device-j (the _dc twins; full accept early-exits
6189 /// in-kernel). Requires the batched-linear stash on every linear layer (stream gate).
6190 fn commit_verified_prefix_stream(
6191 &self,
6192 e: &Engine,
6193 cache: &mut Cache,
6194 snap: &crate::cache::CacheSnapshot,
6195 ckpt: &VerifyCkpt,
6196 acc: &CudaSlice<u32>,
6197 base: usize,
6198 t_v: usize,
6199 ) -> Result<(), Box<dyn std::error::Error>> {
6200 let cfg = &self.cfg;
6201 let ssm = cfg.ssm.as_ref().unwrap();
6202 let d_state = ssm.state_size as usize;
6203 let num_k = ssm.group_count as usize;
6204 let num_v = ssm.time_step_rank as usize;
6205 let d_conv = ssm.conv_kernel as usize;
6206 let conv_dim = d_state * num_k * 2 + d_state * num_v;
6207 let scale = 1.0 / (d_state as f32).sqrt();
6208 for il in 0..self.layers.len() {
6209 if let Some(rl) = cache.recur[il].as_mut() {
6210 let st = ckpt.gdn[il]
6211 .as_ref()
6212 .ok_or("stream restore: batched-linear stash missing")?;
6213 let ring_old = snap.conv[il].as_ref().expect("snapshot missing conv");
6214 let state_in = snap.ssm[il].as_ref().expect("snapshot missing ssm");
6215 e.ssm_conv_ring_rebuild_dc(
6216 &st.qkv_mixed,
6217 ring_old,
6218 &mut rl.conv_state,
6219 conv_dim,
6220 acc,
6221 base,
6222 t_v,
6223 d_conv,
6224 )?;
6225 let mut o = e.uninit(d_state * num_v * t_v)?;
6226 e.gdn_scan_s128_dc(
6227 &st.q_l2,
6228 &st.k_l2,
6229 &st.v_g,
6230 &st.g_log,
6231 &st.beta,
6232 state_in,
6233 &mut rl.ssm_state,
6234 &mut o,
6235 num_v,
6236 acc,
6237 base,
6238 t_v,
6239 scale,
6240 )?;
6241 }
6242 }
6243 Ok(())
6244 }
6245
6246 /// EAGLE3 aux-capturing verify forward over `tokens` (T) — mirrors `decode_step_t_h` exactly
6247 /// (same KV append, same causal verify, same recur advance) but ALSO clones the aux residual-
6248 /// stream hiddens (blocks in `aux_layers`) for TWO columns: the LAST column (always) and the
6249 /// optional `pred_col` (the EAGLE seed = bonus's predecessor). Returns
6250 /// (all_T_logits host, last_col_aux, pred_col_aux?). Used by the EAGLE3 orchestrator's commit.
6251 pub fn decode_step_t_aux2(
6252 &self,
6253 e: &Engine,
6254 tokens: &[u32],
6255 pos0: usize,
6256 cache: &mut Cache,
6257 aux_layers: &[usize],
6258 pred_col: Option<usize>,
6259 ) -> Result<
6260 (Vec<f32>, Vec<CudaSlice<f32>>, Option<Vec<CudaSlice<f32>>>),
6261 Box<dyn std::error::Error>,
6262 > {
6263 let cfg = &self.cfg;
6264 let n_embd = cfg.n_embd as usize;
6265 let eps = cfg.rms_eps;
6266 let t = tokens.len();
6267 let pos_vec: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
6268 let pos_d = e.htod_i32(&pos_vec)?;
6269 let mut x = e.htod(&self.embd.gather(n_embd, tokens))?;
6270 let mut aux_last: Vec<CudaSlice<f32>> = Vec::with_capacity(aux_layers.len());
6271 let mut aux_pred: Vec<CudaSlice<f32>> = Vec::new();
6272 let want_pred = pred_col.is_some();
6273
6274 for (il, layer) in self.layers.iter().enumerate() {
6275 // DISPATCH-MIRRORED norms (FP-order lesson #8) — see decode_step_t_h_emb.
6276 let mixer_fast = self.mixer_in_q8_1_fast(e, &layer.mixer);
6277 let norm_fused = std::env::var("MEMRA_NO_FUSE_NORMQ").is_err() && mixer_fast;
6278 let mut h = vbuf(e, t * n_embd)?; // fully written by either rms_norm arm
6279 if norm_fused {
6280 e.rms_norm_decode(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
6281 } else {
6282 e.rms_norm(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
6283 }
6284 let mixed = match &layer.mixer {
6285 Mixer::Full(fa) => {
6286 self.full_attn_verify(e, fa, &h, None, &pos_d, t, cache, il, None)?
6287 }
6288 Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
6289 Mixer::Linear(la) => {
6290 let mut out = e.zeros(t * n_embd)?;
6291 for col in 0..t {
6292 let mut h_col = e.zeros(n_embd)?;
6293 let src = h.slice(col * n_embd..(col + 1) * n_embd);
6294 e.copy_view_into(&mut h_col, 0, &src, n_embd)?;
6295 let m_col = self.linear_attn_decode(e, la, &h_col, cache, il)?;
6296 e.copy_into(&mut out, col * n_embd, &m_col, n_embd)?;
6297 }
6298 out
6299 }
6300 };
6301 let ffn_fuse = match &layer.ffn {
6302 crate::hybrid::Ffn::Dense {
6303 ffn_gate, ffn_up, ..
6304 } => {
6305 std::env::var("MEMRA_NO_FUSE_NORMQ").is_err()
6306 && e.uses_q8_1_fast(ffn_gate)
6307 && e.uses_q8_1_fast(ffn_up)
6308 }
6309 crate::hybrid::Ffn::Moe(_) => false,
6310 };
6311 let mut x1 = vbuf(e, t * n_embd)?; // fully written by add / add_rms_norm
6312 let mut z = vbuf(e, t * n_embd)?; // fully written by rms_norm_decode / add_rms_norm
6313 if ffn_fuse {
6314 e.add(&x, &mixed, &mut x1, t * n_embd)?;
6315 e.rms_norm_decode(
6316 &x1,
6317 layer.post_attn_norm.float_data(),
6318 &mut z,
6319 n_embd,
6320 t,
6321 eps,
6322 )?;
6323 } else {
6324 e.add_rms_norm(
6325 &x,
6326 &mixed,
6327 layer.post_attn_norm.float_data(),
6328 &mut x1,
6329 &mut z,
6330 n_embd,
6331 t,
6332 eps,
6333 )?;
6334 }
6335 let ffn_out = match &layer.ffn {
6336 crate::hybrid::Ffn::Dense {
6337 ffn_gate,
6338 ffn_up,
6339 ffn_down,
6340 } => {
6341 let n_ff = ffn_gate.out_features();
6342 let gate = e.matmul_decode_exact(ffn_gate, &z, t)?;
6343 let up = e.matmul_decode_exact(ffn_up, &z, t)?;
6344 let mut act = vbuf(e, t * n_ff)?; // fully written by ffn_act_lim
6345 // dense FFN clamp = the SHEXP array (upstream build_ffn serves both).
6346 Self::ffn_act_lim(
6347 e,
6348 &self.cfg,
6349 &gate,
6350 &up,
6351 1.0,
6352 1.0,
6353 self.cfg.clamp_shexp_at(il as u32),
6354 &mut act,
6355 t * n_ff,
6356 )?;
6357 e.matmul_decode_exact(ffn_down, &act, t)?
6358 }
6359 crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il(e, m, &z, t, il as u16)?,
6360 };
6361 let mut x2 = vbuf(e, t * n_embd)?; // fully written by add
6362 e.add(&x1, &ffn_out, &mut x2, t * n_embd)?;
6363 if aux_layers.contains(&il) {
6364 let mut a = e.zeros(n_embd)?;
6365 e.copy_view_into(&mut a, 0, &x2.slice((t - 1) * n_embd..t * n_embd), n_embd)?;
6366 aux_last.push(a);
6367 if let Some(pc) = pred_col {
6368 let mut ap = e.zeros(n_embd)?;
6369 e.copy_view_into(
6370 &mut ap,
6371 0,
6372 &x2.slice(pc * n_embd..(pc + 1) * n_embd),
6373 n_embd,
6374 )?;
6375 aux_pred.push(ap);
6376 }
6377 }
6378 x = x2;
6379 }
6380 let mut hn = vbuf(e, t * n_embd)?; // fully written by rms_norm_decode
6381 e.rms_norm_decode(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
6382 let logits = e.matmul_decode_exact(&self.output, &hn, t)?;
6383 let host = e.dtoh(&logits)?;
6384 cache.pos += t;
6385 Ok((
6386 host,
6387 aux_last,
6388 if want_pred { Some(aux_pred) } else { None },
6389 ))
6390 }
6391
6392 /// step35 SPEC-VERIFY attention over T query tokens — a per-row REPLAY of the eager
6393 /// `step35_decode_attn`.
6394 ///
6395 /// WHY A REPLAY AND NOT A BATCHED TWIN. The verify's whole job is to be bit-identical to what
6396 /// the eager decode would have computed for the same tokens; that is what makes greedy spec
6397 /// decode exact (run-spec asserts token identity for K=1..8). Every other verify arm in this
6398 /// file earns that identity by carefully mirroring dispatch (`matmul_decode_exact` to force
6399 /// MMVQ at any m, per-layer `ffn_fuse` mirroring, per-row `fa_decode` key bounds). step35
6400 /// stacks FOUR more per-layer degrees of freedom on top of that — per-layer `n_head`
6401 /// (64 full / 96 SWA), per-layer rotary width (64 full / 128 SWA), per-layer rope base, and a
6402 /// SEPARATE `attn_gate` tensor whose projection shares the attn-normed input — and its SWA
6403 /// layers attend through a token-OFFSET view whose offset is a function of the ABSOLUTE
6404 /// position of each query row. A batched twin would have to reproduce all of that AND the
6405 /// per-row offset in one launch; the offset alone rules out the existing rows kernels (they
6406 /// take one `base_len`, not a per-row offset).
6407 ///
6408 /// So this arm calls the eager path itself, once per row, on the same cache. Identity is then
6409 /// true BY CONSTRUCTION rather than by mirroring: row r runs exactly the kernel sequence that
6410 /// eager decode step r runs (same projections, same q8_1 fusion decision, same append, same
6411 /// view arithmetic, same `fa_decode_kvmod`, same gate), because it IS that code. Cost: T x the
6412 /// eager decode mixer instead of one batched pass — the same trade the generic arm's `else`
6413 /// per-row loop already accepts when `fa_rows_eligible` says no. Correctness first; a batched
6414 /// step35 twin is a perf lane's job and must be gated against this arm.
6415 ///
6416 /// The `h_q8` pre-quantized pair from the caller's fused norm is NOT forwarded: it is a
6417 /// T-row buffer and `step35_decode_attn`'s `pre_q` contract is one row. Instead each row's
6418 /// f32 `h` slice is handed over and the callee re-derives its own q8_1 exactly as eager decode
6419 /// does (`quantize_q8_1(h, 1, n_embd)`) — which is the dispatch being mirrored. Callers that
6420 /// took the fused arm therefore MUST still pass a live `h`; `step35_verify` asserts that.
6421 #[allow(clippy::too_many_arguments)]
6422 fn step35_verify(
6423 &self,
6424 e: &Engine,
6425 fa: &FullAttnLayer,
6426 h: &CudaSlice<f32>,
6427 h_q8: Option<(&CudaSlice<i8>, &CudaSlice<f32>)>,
6428 t: usize,
6429 cache: &mut Cache,
6430 il: usize,
6431 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6432 let n_embd = self.cfg.n_embd as usize;
6433 // The fused (h-less) attn-norm arm hands `h` as a zero-length placeholder. This arm needs
6434 // the f32 rows, so the caller must not take that lever for step35 — enforced at the call
6435 // site by the `Mixer::Full(_) if self.cfg.step35.is_some() => false` arm of
6436 // `lin_q8_only` in `decode_step_t_core_stream`, and asserted here so a future caller
6437 // cannot regress it into silently reading an empty buffer.
6438 assert_eq!(
6439 h.len(),
6440 t * n_embd,
6441 "step35_verify needs the f32 attn-normed rows ([t*n_embd]); the caller took the \
6442 fused q8-only norm arm (h_q8={}) — step35 must stay on the unfused arm",
6443 h_q8.is_some()
6444 );
6445 // ROW WIDTH IS n_embd, NOT n_head*head_dim: `step35_decode_attn` returns the mixer output
6446 // AFTER `wo`, so a row is [n_embd] — the same contract the generic arm's
6447 // `matmul_decode_exact(&fa.wo, &attn_g, t)` return has. Sizing this buffer from the
6448 // per-layer head geometry (8192 on full-attn, 12288 on SWA) instead overran the row on the
6449 // FIRST copy and panicked inside `copy_into`'s `CudaView::slice` unwrap
6450 // (raw/mtp-bt-20260806T212127Z.log frames 12-13).
6451 let mut out = vbuf(e, t * n_embd)?; // each row fully written by the copy below
6452 for r in 0..t {
6453 // Absolute position of this query row. `cache.pos` is the committed length at round
6454 // start and every row before r has already been appended by this loop, so the r-th
6455 // verify token sits at cache.pos + r — the same position eager decode would give it.
6456 let pos_d = e.htod_i32(&[(cache.pos + r) as i32])?;
6457 let mut h_row = vbuf(e, n_embd)?; // fully written by copy_view_into
6458 e.copy_view_into(
6459 &mut h_row,
6460 0,
6461 &h.slice(r * n_embd..(r + 1) * n_embd),
6462 n_embd,
6463 )?;
6464 // THE eager decode mixer: appends this row's K/V at kvl.len, advances it, then
6465 // attends over the (SWA-offset) view. Post-`wo`, same contract as this fn returns.
6466 let o = self.step35_decode_attn(e, fa, il, &h_row, None, &pos_d, cache)?;
6467 debug_assert_eq!(
6468 o.len(),
6469 n_embd,
6470 "step35_decode_attn returns post-wo [n_embd]"
6471 );
6472 e.copy_into(&mut out, r * n_embd, &o, n_embd)?;
6473 }
6474 Ok(out)
6475 }
6476
6477 /// Full-attention mixer over T query tokens with a GROWING resident KV (verify path, §D.3).
6478 /// Appends the T new K/V columns to cache.kv[il] then attends causally over [0..len) via
6479 /// fa_prefill. Token-major [T, kv_dim] projection layout == cache row layout (single copy).
6480 #[allow(clippy::too_many_arguments)]
6481 fn full_attn_verify(
6482 &self,
6483 e: &Engine,
6484 fa: &FullAttnLayer,
6485 h: &CudaSlice<f32>,
6486 h_q8: Option<(&CudaSlice<i8>, &CudaSlice<f32>)>,
6487 pos_d: &CudaSlice<i32>,
6488 t: usize,
6489 cache: &mut Cache,
6490 il: usize,
6491 stream_ctr: Option<&CudaSlice<i32>>,
6492 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6493 // step35: the generic geometry below is wrong for this arch (per-layer n_head, partial
6494 // per-layer rope, the SWA offset view, and a SEPARATE head-wise gate tensor), so it takes
6495 // its own arm. A verify that silently computes different attention than decode defeats the
6496 // whole self-consistency gate, so the arm is a per-row REPLAY of `step35_decode_attn`
6497 // rather than a batched twin — see `step35_verify` for why that is the exactness-correct
6498 // shape and not laziness.
6499 if self.cfg.step35.is_some() {
6500 if stream_ctr.is_some() {
6501 return Err(
6502 "step35 has no ROUND-STREAM verify arm (the device-counter _dc twins \
6503 cannot express the SWA offset KV view; same root cause as the dc \
6504 decode refusal) — run spec without the stream arm"
6505 .into(),
6506 );
6507 }
6508 return self.step35_verify(e, fa, h, h_q8, t, cache, il);
6509 }
6510 let cfg = &self.cfg;
6511 let geometry = cfg.full_attention_geometry_at(il as u32);
6512 let n_head = geometry.n_head as usize;
6513 let n_head_kv = geometry.n_head_kv as usize;
6514 let head_dim = geometry.head_dim_k as usize;
6515 let eps = cfg.rms_eps;
6516 let scale = geometry.attention_scale();
6517 let n_embd = cfg.n_embd as usize;
6518
6519 // DECODE-EXACT Q/K/V projections: matmul_decode_exact forces the MMVQ (warp-per-row) path
6520 // for every m, matching the T=1 decode's FP accumulation order. matmul_pre at m>=5 would
6521 // fall to dp4a (128-thread, two-level reduce) with a different FP sum order.
6522 // Q8 TRUNK-FUSION at T=1: DISPATCH-MIRRORS the eager decode's fused3 (bit-identical body).
6523 // BATCHED EPILOGUE RE-FUSE (lane/vt-fixes fix 2): `h_q8` = the attn-input norm's q8_1
6524 // form emitted by the fused rms_norm_q8_1 (bit-identical to rms_norm_decode ->
6525 // quantize_q8_1, kernel-check-pinned). When present (caller checked mixer q8_1-fast),
6526 // every projection consumes it — `h` may be a zero-len placeholder and must not be read.
6527 let (qf, mut k, v) = {
6528 let mut fused = None;
6529 let qkv_fast =
6530 e.uses_q8_1_fast(&fa.wq) && e.uses_q8_1_fast(&fa.wk) && e.uses_q8_1_fast(&fa.wv);
6531 if t == 1 && qkv_fast {
6532 let (hq_o, hd_o);
6533 let (hq, hd): (&CudaSlice<i8>, &CudaSlice<f32>) = match h_q8 {
6534 Some(p) => p,
6535 None => {
6536 (hq_o, hd_o) = e.quantize_q8_1(h, 1, n_embd)?;
6537 (&hq_o, &hd_o)
6538 }
6539 };
6540 fused = e.matmul_q8_fused3(&fa.wq, &fa.wk, &fa.wv, hq, hd)?;
6541 } else if spec_fused_t() && (2..=4).contains(&t) && qkv_fast {
6542 // VERIFY-TIER TRUNK FUSION (MEMRA_SPEC_FUSED_T): one shared quantize + one
6543 // fused3 batched launch replaces three decode-exact calls (3 re-quantizes of
6544 // the same h + 3 _b2/_b4 launches). Bit-identical per (tensor,token,row).
6545 let (hq_o, hd_o);
6546 let (hq, hd): (&CudaSlice<i8>, &CudaSlice<f32>) = match h_q8 {
6547 Some(p) => p,
6548 None => {
6549 (hq_o, hd_o) = e.quantize_q8_1(h, t, n_embd)?;
6550 (&hq_o, &hd_o)
6551 }
6552 };
6553 fused = e.matmul_q8_fused3_t(&fa.wq, &fa.wk, &fa.wv, hq, hd, t)?;
6554 }
6555 match (fused, h_q8) {
6556 (Some(triple), _) => triple,
6557 // shared pre-quantized activation (q8_1-fast guaranteed by the caller): the
6558 // decode-exact dispatch consumes (hq, hd) instead of re-quantizing 3x.
6559 (None, Some((hq, hd))) if qkv_fast => (
6560 e.matmul_decode_exact_pre(&fa.wq, hq, hd, t)?,
6561 e.matmul_decode_exact_pre(&fa.wk, hq, hd, t)?,
6562 e.matmul_decode_exact_pre(&fa.wv, hq, hd, t)?,
6563 ),
6564 (None, _) => (
6565 e.matmul_decode_exact(&fa.wq, h, t)?,
6566 e.matmul_decode_exact(&fa.wk, h, t)?,
6567 e.matmul_decode_exact(&fa.wv, h, t)?,
6568 ),
6569 }
6570 };
6571 // M3/Hy3 have no attention output gate — wq out is exactly q; skip the split.
6572 let gated = geometry.attention_gate == memra_gguf::config::AttentionGateKind::FusedQ;
6573 let (mut q, gate) = if gated {
6574 let mut q = vbuf(e, t * n_head * head_dim)?; // fully written by q_gate_split
6575 let mut gate = vbuf(e, t * n_head * head_dim)?; // fully written by q_gate_split
6576 e.q_gate_split(&qf, &mut q, &mut gate, head_dim, n_head, t)?;
6577 (q, Some(gate))
6578 } else {
6579 (qf, None)
6580 };
6581
6582 let mut qn = vbuf(e, t * n_head * head_dim)?; // fully written by rms_norm
6583 e.rms_norm(
6584 &q,
6585 fa.q_norm.float_data(),
6586 &mut qn,
6587 head_dim,
6588 n_head * t,
6589 eps,
6590 )?;
6591 q = qn;
6592 let mut kn = vbuf(e, t * n_head_kv * head_dim)?; // fully written by rms_norm
6593 e.rms_norm(
6594 &k,
6595 fa.k_norm.float_data(),
6596 &mut kn,
6597 head_dim,
6598 n_head_kv * t,
6599 eps,
6600 )?;
6601 k = kn;
6602 let rope_dims = geometry.n_rot as usize;
6603 e.rope_neox(
6604 &mut q,
6605 pos_d,
6606 head_dim,
6607 rope_dims,
6608 n_head,
6609 t,
6610 geometry.rope_base,
6611 1.0,
6612 )?;
6613 e.rope_neox(
6614 &mut k,
6615 pos_d,
6616 head_dim,
6617 rope_dims,
6618 n_head_kv,
6619 t,
6620 geometry.rope_base,
6621 1.0,
6622 )?;
6623
6624 // append T new K/V columns to the resident QUANTIZED cache. k/v are token-major [T, kv_dim]
6625 // f32; append-quantize each of the T token rows into the byte cache (q8_0 K / q5_1 V).
6626 let kvl = cache.kv[il].as_mut().unwrap();
6627 let (kv_dim_k, kv_dim_v, ktb, vtb) =
6628 (kvl.kv_dim_k, kvl.kv_dim_v, kvl.k_tok_bytes, kvl.v_tok_bytes);
6629 if let Some(ctr) = stream_ctr {
6630 // stream: ONE batched append at the device counter (rows kernel = the per-view warp
6631 // math on a (block, token) grid, documented byte-identical); host len is a stale
6632 // LOWER BOUND under pre-issue (drain reconciles it).
6633 e.append_kv_quantized_rows_dc(
6634 &k,
6635 &v,
6636 &mut kvl.k,
6637 &mut kvl.v,
6638 ctr,
6639 t,
6640 kv_dim_k,
6641 kv_dim_v,
6642 ktb,
6643 vtb,
6644 crate::Engine::kv_fp8_on(),
6645 )?;
6646 } else {
6647 for i in 0..t {
6648 let k_row = k.slice(i * kv_dim_k..(i + 1) * kv_dim_k);
6649 let v_row = v.slice(i * kv_dim_v..(i + 1) * kv_dim_v);
6650 e.append_kv_quantized_view(
6651 &k_row,
6652 &v_row,
6653 &mut kvl.k,
6654 &mut kvl.v,
6655 kvl.len + i,
6656 kv_dim_k,
6657 kv_dim_v,
6658 ktb,
6659 vtb,
6660 crate::Engine::kv_fp8_on(),
6661 )?;
6662 }
6663 kvl.len += t;
6664 }
6665
6666 // BIT-IDENTICAL VERIFY ATTENTION (spec-exactness fix): the FP accumulation order must be
6667 // byte-for-byte identical to the eager decode path. fa_prefill uses a different tile size
6668 // (BLOCK_Q=64, BK=32) and online-softmax structure than fa_decode's split-K + combine,
6669 // which changes FP summation order and can flip argmax at tight logit margins. Query row r
6670 // attends to keys [0..base_len+r+1) — each successive row sees one more key (the causal
6671 // property). This matches eager: decode appends k at len, then fa_decode sees t_kv = len+1
6672 // keys. The verify appends all T tokens first but bounds the key range per row.
6673 //
6674 // MULTI-ROW FUSED PATH (the long-ctx spec fix, 2026-07-03): when every row takes the vec
6675 // kernel (base_len+1 >= FA_VEC_MIN_TKV), ONE fa_decode_rows launch executes the exact
6676 // per-row program for all T rows (grid.z = row, per-row n_splits from the same
6677 // fa_split_keys formula) — replacing T x (2 launches + 2 dtod copies + 5 partial allocs)
6678 // and multiplying resident CTAs by T on a latency-bound kernel. Bit-identical per row by
6679 // construction; kernel-check pins rows-vs-loop byte identity, run-spec is the end gate.
6680 // Short ctx (any row below the vec crossover) and MEMRA_NO_FA_VEC/MEMRA_FA_ROWS_OFF keep the
6681 // per-row loop (whose fa_decode picks scalar/vec per row exactly like eager decode).
6682 let mut attn = vbuf(e, t * n_head * head_dim)?; // fully written by every FA arm below
6683 let base_len = kvl.len - t; // KV len BEFORE this round's T tokens were appended
6684 // T=1 INCLUDED (2026-07-05): p-min cuts the draft to 1 in ~75% of rounds on hard
6685 // (agentic) content — the old t>1 gate sent those rounds to the per-row loop (262us/row
6686 // + q-row copy + per-row allocs vs 93us/row through the fused kernel at grid.z=1, same
6687 // program). nsys accounting: 1088 of 1456 verify FA launches were T=1 escapees.
6688 // LEAN T=1 ARM (MEMRA_SPEC_LEAN, close35): at t==1, q IS one row and fa_decode on it is
6689 // the EXACT eager decode dispatch (vec_q_v2 + combine_f32; the rows pair measured +50us
6690 // at m=1). Byte-identical: kernel-check pins rows-vs-loop identity, and the per-row loop
6691 // at t=1 is fa_decode on the same q with zero-offset copies. Gates arbitrate.
6692 if let Some(ctr) = stream_ctr {
6693 // STREAM ARM: causal base from the device counter; host kvl.len is a stale lower
6694 // bound used only for the split-sizing upper bound (+64 slack covers M pre-issued
6695 // rounds at K<=8). Views span the bound; per-row limits derive in-kernel.
6696 let upper = kvl.len + t + 64;
6697 let k_view = e.view_u8(&kvl.k, (upper.min(cache.max_ctx)) * ktb);
6698 let v_view = e.view_u8(&kvl.v, (upper.min(cache.max_ctx)) * vtb);
6699 e.fa_decode_rows_dc(
6700 &q,
6701 &k_view,
6702 &v_view,
6703 &mut attn,
6704 head_dim,
6705 n_head,
6706 n_head_kv,
6707 ctr,
6708 upper.min(cache.max_ctx),
6709 t,
6710 scale,
6711 ktb,
6712 vtb,
6713 0,
6714 false,
6715 )?;
6716 } else if spec_lean() && t == 1 {
6717 let t_kv = base_len + 1;
6718 let k_view = e.view_u8(&kvl.k, t_kv * ktb);
6719 let v_view = e.view_u8(&kvl.v, t_kv * vtb);
6720 e.fa_decode_kvmod(
6721 &q,
6722 &k_view,
6723 &v_view,
6724 &mut attn,
6725 head_dim,
6726 n_head,
6727 n_head_kv,
6728 t_kv,
6729 scale,
6730 ktb,
6731 vtb,
6732 crate::Engine::kv_fp8_on(),
6733 )?;
6734 } else if e.fa_rows_eligible(base_len, head_dim) {
6735 let k_view = e.view_u8(&kvl.k, (base_len + t) * ktb);
6736 let v_view = e.view_u8(&kvl.v, (base_len + t) * vtb);
6737 e.fa_decode_rows(
6738 &q,
6739 &k_view,
6740 &v_view,
6741 &mut attn,
6742 head_dim,
6743 n_head,
6744 n_head_kv,
6745 base_len,
6746 t,
6747 scale,
6748 ktb,
6749 vtb,
6750 None,
6751 false,
6752 crate::Engine::kv_fp8_on(),
6753 None,
6754 )?;
6755 } else {
6756 for r in 0..t {
6757 let t_kv_r = base_len + r + 1; // this row sees keys [0..t_kv_r)
6758 let k_view_r = e.view_u8(&kvl.k, t_kv_r * ktb);
6759 let v_view_r = e.view_u8(&kvl.v, t_kv_r * vtb);
6760 // copy q row into an owned buffer (fa_decode takes &CudaSlice, not CudaView)
6761 let mut q_row = vbuf(e, n_head * head_dim)?; // fully written by copy_view_into
6762 let q_src = q.slice(r * n_head * head_dim..(r + 1) * n_head * head_dim);
6763 e.copy_view_into(&mut q_row, 0, &q_src, n_head * head_dim)?;
6764 let mut attn_row = vbuf(e, n_head * head_dim)?; // fully written by fa_decode
6765 e.fa_decode_kvmod(
6766 &q_row,
6767 &k_view_r,
6768 &v_view_r,
6769 &mut attn_row,
6770 head_dim,
6771 n_head,
6772 n_head_kv,
6773 t_kv_r,
6774 scale,
6775 ktb,
6776 vtb,
6777 crate::Engine::kv_fp8_on(),
6778 )?;
6779 e.copy_into(
6780 &mut attn,
6781 r * n_head * head_dim,
6782 &attn_row,
6783 n_head * head_dim,
6784 )?;
6785 }
6786 }
6787
6788 let attn_g = match &gate {
6789 Some(gate) => {
6790 let mut gsig = vbuf(e, t * n_head * head_dim)?; // fully written by sigmoid
6791 e.sigmoid(gate, &mut gsig, t * n_head * head_dim)?;
6792 let mut ag = vbuf(e, t * n_head * head_dim)?; // fully written by mul
6793 e.mul(&attn, &gsig, &mut ag, t * n_head * head_dim)?;
6794 ag
6795 }
6796 None => attn,
6797 };
6798 // DECODE-EXACT wo projection: at m>=5 (K=4+ with pending) the generic matmul would use dp4a
6799 // (128-thread, different FP sum order than MMVQ). Force MMVQ for bit-identity with decode.
6800 Ok(e.matmul_decode_exact(&fa.wo, &attn_g, t)?)
6801 }
6802
6803 /// Context-linear bytes for a plain serving session's trunk cache.
6804 pub fn plain_session_kv_bytes_per_token(&self) -> usize {
6805 crate::cache::cache_bytes_per_token(&self.cfg)
6806 }
6807
6808 /// `(logical bytes/token, ring-capped bytes/token, ring row cap)` for exact admission.
6809 pub fn plain_session_kv_shape(&self) -> (usize, usize, usize) {
6810 (
6811 self.plain_session_kv_bytes_per_token(),
6812 crate::cache::cache_ring_bytes_per_token(&self.cfg),
6813 crate::cache::cache_ring_row_cap(&self.cfg),
6814 )
6815 }
6816
6817 /// Context-linear bytes for a speculative serving session: trunk cache plus persistent MTP
6818 /// scratch. With no MTP head this equals the plain coefficient.
6819 pub fn spec_session_kv_bytes_per_token(&self) -> usize {
6820 let scratch = self
6821 .mtp
6822 .as_ref()
6823 .map(|mtp| {
6824 let (_, _, k, v) = mtp_scratch_layout(&self.cfg, mtp.geom.as_ref());
6825 k + v
6826 })
6827 .unwrap_or(0);
6828 self.plain_session_kv_bytes_per_token()
6829 .saturating_add(scratch)
6830 }
6831
6832 /// Spec twin of [`HybridModel::plain_session_kv_shape`]; Step35's persistent MTP scratch is
6833 /// capped by the same SWA ring rows as the trunk.
6834 pub fn spec_session_kv_shape(&self) -> (usize, usize, usize) {
6835 let total = self.spec_session_kv_bytes_per_token();
6836 let (_, mut ring, rows) = self.plain_session_kv_shape();
6837 if rows > 0 {
6838 ring = ring.saturating_add(
6839 self.mtp
6840 .as_ref()
6841 .map(|mtp| {
6842 let (_, _, k, v) = mtp_scratch_layout(&self.cfg, mtp.geom.as_ref());
6843 k + v
6844 })
6845 .unwrap_or(0),
6846 );
6847 }
6848 (total, ring, rows)
6849 }
6850
6851 /// Greedy MTP speculative decode (§B). Token-identical to `generate(prompt, max_new)` but uses
6852 /// the NextN head to draft K tokens then verifies them in one batched target forward.
6853 /// Returns (generated tokens, total_drafted, total_accepted) so the caller can report
6854 /// acceptance rate. `k` = draft length per round.
6855 ///
6856 /// GRAPH DRAFT (stage 2 of graph-grade spec): when the model is all-Dense and the MTP head is
6857 /// Dense (no MoE host readbacks), the fixed-shape T=1 MTP forward is CUDA-graph-captured ONCE
6858 /// and replayed per draft step — the ~40 eager launches per drafted token collapse into one
6859 /// graph dispatch; only the 4-byte token id (and 4-byte p-min confidence) round-trip per step.
6860 /// Event tracking is disabled for the whole call (generate_graph pattern) so every buffer the
6861 /// captured graph references is event-free; the spec loop is strictly single-stream.
6862 /// MEMRA_SPEC_NOGRAPH=1 forces the eager draft chain.
6863 /// SAMPLED mode (MEMRA_SPEC_TEMP>0) has its OWN capture (gumbel-perturbed in-graph argmax,
6864 /// device Philox event counter, persistent q retention) — graph-vs-eager sampled streams are
6865 /// bit-identical for the same (seed, prompt, K, temp); see the sampled-graph setup in
6866 /// generate_spec_inner2.
6867 /// Multi-turn session: trunk cache + MTP draft scratch persist across generate calls, so
6868 /// turn N+1 primes ONLY its new suffix (the 124k-conversation daily pattern — re-priming a
6869 /// 32k history costs ~54s; a suffix prime costs seconds). APPEND-ONLY by construction: the
6870 /// hybrid linear-attn states are in-place (no position index), so a session can extend but
6871 /// never rewind — `committed` is the exact token list whose state the caches hold (includes
6872 /// any overshoot tokens past max_new; the caller renders from `committed`, not its own echo).
6873 pub fn new_session(
6874 &self,
6875 e: &Engine,
6876 max_ctx: usize,
6877 ) -> Result<SpecSession, Box<dyn std::error::Error>> {
6878 Ok(SpecSession {
6879 // STAGE-OWNED KV (lane/pp2-spec 2026-08-06): `pp::new_cache`, not `Cache::new`. This
6880 // is the SERVING spec-session path, and with the ppN door open across two cards a
6881 // primary-homed cache makes every remote stage peer-read its OWN KV on every verify
6882 // round — the wrong-card class already fixed on the two batched serving paths
6883 // (worker.rs 2483 / 2837). With the door shut `new_cache` IS `Cache::new` (same
6884 // branch, same allocations), so single-device behavior is byte-unchanged.
6885 cache: crate::pp::new_cache(e, &self.cfg, max_ctx)?,
6886 scratch: MtpScratch::new(
6887 e,
6888 &self.cfg,
6889 max_ctx,
6890 self.mtp.as_ref().and_then(|m| m.geom.as_ref()),
6891 )?,
6892 committed: Vec::new(),
6893 last_h: None,
6894 next_pred: None,
6895 sctr: 0,
6896 uctr: 0,
6897 draft_ctx: None,
6898 pending_tok: None,
6899 turn_ckpt: None,
6900 telem: SpecTelemetryCounters::default(),
6901 capture_at: None,
6902 boundary_capture: None,
6903 })
6904 }
6905
6906 /// SPEC-ON-CACHE-HIT restore (lane/spec-on-cache-hit, 2026-08-18 — PORT-PLAN item 3,
6907 /// research/cache-spec-design-20260814, scoped to WHOLE-ENTRY restores only): build a
6908 /// SpecSession around a trunk cache the worker already restored from a prefix-cache
6909 /// entry, re-installing the entry's published draft plane as the MTP scratch rows
6910 /// `[0..prefix.len())` and the entry's boundary hidden as `last_h`, then feeding the
6911 /// prompt SUFFIX here — through EXACTLY the plain path's program selection — so the
6912 /// worker always receives a fully-warm continuation session (committed = whole
6913 /// prompt, `next_pred` + `last_h` set; caller sets `next_pred` from the entry's
6914 /// boundary logits on the empty-suffix shape).
6915 ///
6916 /// PROGRAM LAW (the splitiso two-programs class, learned AGAIN in this lane's own
6917 /// gate): the identity target for a converted hit is the PLAIN hit serving the same
6918 /// request, and plain feeds a carried suffix via eager `decode_step` below
6919 /// PRIME_MIN_T and via `prime_cache` at/above it (prefill_tick's arms). The generate
6920 /// path's tokenwise arm routes qwen35-class through the BATCHED T=1 program
6921 /// (`spec_target_step_h`) instead — ULP-different suffix rows, and the gate measured
6922 /// the near-tie flip at generated token ~8 (research/spec-cache-20260818, qwen r3).
6923 /// So the suffix is fed HERE, mirroring prefill_tick arm-for-arm, not handed to the
6924 /// burst prime.
6925 ///
6926 /// SEED RULE (both sampling regimes; lane/sampled-hit-spec 2026-08-19, sampled draw
6927 /// added by lane/sampled-spec-quality 2026-08-19). The boundary token is produced by
6928 /// EXACTLY the rule the cold burst entry applies to its own first token from the same
6929 /// logits row: `argmax` when greedy, and a `sample_boundary_token` draw at Philox
6930 /// counter 0 when sampled. Both shapes are covered — the entry's boundary logits on a
6931 /// full-cover (empty-suffix) hit, this feed's own boundary logits on a suffix hit.
6932 /// That is what keeps a restored session seed-identical to a cold one PER SEED: the
6933 /// cold session draws from the identical row at counter 0 and then runs its rounds from
6934 /// counter 1, so the restored session admits with `sctr = 1` after its own draw.
6935 /// The WORKER owns the one refusal this constructor cannot see — a constrained request.
6936 /// (The penalized-sampled refusal was LIFTED once the burst's penalty window learned to
6937 /// span the session: `committed` here is the WHOLE prompt, so the restored session's
6938 /// window is the cold session's window. It comes back if `MEMRA_SPEC_PEN_SESSION=0`.)
6939 ///
6940 /// NOT the rolled-back partial-restore hazard: the caller restores at exactly the
6941 /// entry's captured endpoint (`e.pos`) through the shipping whole-entry path;
6942 /// mid-entry (`at < e.pos`) trunk restores stay behind MEMRA_PREFIX_PARTIAL_RESTORE
6943 /// and are never routed here.
6944 ///
6945 /// Failure contract: `Err((Some(cache), why))` before any trunk mutation — the
6946 /// worker rebuilds the plain carrier and the hit serves plain, byte-unchanged.
6947 /// `Err((None, why))` after the suffix feed began — the carrier is part-fed and
6948 /// UNUSABLE; the worker serves the request cold-plain (correct, slower) and the
6949 /// entry stays published for the next request.
6950 #[allow(clippy::too_many_arguments)]
6951 pub fn spec_session_from_restored(
6952 &self,
6953 e: &Engine,
6954 mut cache: Cache,
6955 prefix: Vec<u32>,
6956 suffix: &[u32],
6957 draft_k: &CudaSlice<u8>,
6958 draft_v: &CudaSlice<u8>,
6959 draft_k_tok_bytes: usize,
6960 draft_v_tok_bytes: usize,
6961 draft_len: usize,
6962 last_h: &[f32],
6963 // The ENTRY's boundary logits row (the full-cover shape's seed source). May be empty
6964 // when a suffix follows — the feed's own logits are the boundary then.
6965 boundary_logits: &[f32],
6966 // The request's sampler, or None for greedy. Owned here so the seed rule lives in
6967 // ONE place instead of being half-applied by the worker.
6968 sampling: Option<SpecSampling>,
6969 require_anchor: bool,
6970 max_ctx: usize,
6971 ) -> Result<SpecSession, (Option<Cache>, String)> {
6972 let pos = prefix.len();
6973 let fail = |cache: Cache, msg: String| -> Result<SpecSession, (Option<Cache>, String)> {
6974 Err((Some(cache), msg))
6975 };
6976 if self.mtp.is_none() {
6977 return fail(cache, "no MTP head attached (nothing to draft with)".into());
6978 }
6979 if pos == 0 {
6980 return fail(cache, "empty committed prefix".into());
6981 }
6982 if cache.pos != pos {
6983 let msg = format!(
6984 "restored cache pos {} != restored prefix len {pos}",
6985 cache.pos
6986 );
6987 return fail(cache, msg);
6988 }
6989 if draft_len != pos {
6990 return fail(
6991 cache,
6992 format!("draft plane len {draft_len} != restored prefix len {pos}"),
6993 );
6994 }
6995 if pos + suffix.len() >= max_ctx {
6996 return fail(
6997 cache,
6998 format!(
6999 "prompt {} + suffix would not leave generation room in ctx {max_ctx}",
7000 pos + suffix.len(),
7001 ),
7002 );
7003 }
7004 let mut scratch = match MtpScratch::new(
7005 e,
7006 &self.cfg,
7007 max_ctx,
7008 self.mtp.as_ref().and_then(|m| m.geom.as_ref()),
7009 ) {
7010 Ok(s) => s,
7011 Err(err) => return fail(cache, format!("draft scratch alloc failed: {err}")),
7012 };
7013 if scratch.kv.ring.is_some() {
7014 return fail(
7015 cache,
7016 "ring-backed draft scratch (Step35 SWA) cannot take a flat prefix restore".into(),
7017 );
7018 }
7019 if scratch.kv.k_tok_bytes != draft_k_tok_bytes
7020 || scratch.kv.v_tok_bytes != draft_v_tok_bytes
7021 {
7022 return fail(
7023 cache,
7024 format!(
7025 "draft plane layout {draft_k_tok_bytes}/{draft_v_tok_bytes} != scratch \
7026 {}/{} bytes/token (stale entry across a format change)",
7027 scratch.kv.k_tok_bytes, scratch.kv.v_tok_bytes,
7028 ),
7029 );
7030 }
7031 if pos > scratch.cap {
7032 return fail(
7033 cache,
7034 format!(
7035 "draft plane rows {pos} exceed scratch capacity {}",
7036 scratch.cap
7037 ),
7038 );
7039 }
7040 let kb = pos * draft_k_tok_bytes;
7041 let vb = pos * draft_v_tok_bytes;
7042 if draft_k.len() < kb || draft_v.len() < vb {
7043 return fail(
7044 cache,
7045 format!(
7046 "truncated draft plane: K {} < {kb} or V {} < {vb} bytes",
7047 draft_k.len(),
7048 draft_v.len(),
7049 ),
7050 );
7051 }
7052 if kb > 0 {
7053 if let Err(err) = e.copy_u8_into(&mut scratch.kv.k, 0, draft_k, kb) {
7054 return fail(cache, format!("draft K restore copy failed: {err}"));
7055 }
7056 }
7057 if vb > 0 {
7058 if let Err(err) = e.copy_u8_into(&mut scratch.kv.v, 0, draft_v, vb) {
7059 return fail(cache, format!("draft V restore copy failed: {err}"));
7060 }
7061 }
7062 if let Err(err) = scratch.set_len(e, pos) {
7063 return fail(cache, format!("draft scratch len set failed: {err}"));
7064 }
7065 let mut last_h_dev = if last_h.len() == self.cfg.n_embd as usize {
7066 // anchor upload failure is acceptance-only when a suffix feed follows (fill
7067 // row-0 falls back to zeros) but FATAL for an empty-suffix continuation (the
7068 // burst entry asserts committed + last_h + next_pred) — the caller says which.
7069 e.htod(last_h).ok()
7070 } else {
7071 None
7072 };
7073 if require_anchor && last_h_dev.is_none() {
7074 return fail(
7075 cache,
7076 "empty-suffix continuation requires the entry's boundary hidden anchor".into(),
7077 );
7078 }
7079 let mut committed = prefix;
7080 // Set on BOTH shapes below (suffix-fed and full-cover) — never left None, which is
7081 // what the empty-suffix continuation assert in the burst entry requires.
7082 let next_pred;
7083 // Philox: (0,0) at admit exactly like a fresh session; a sampled boundary draw below
7084 // consumes counter 0 and leaves 1, which is the state a cold session reaches after
7085 // drawing its own first token from the same row.
7086 let mut sctr = 0u32;
7087 let sampled = sampling.is_some_and(|s| s.temp > 0.0) && spec_sampled_boundary_on();
7088 // Penalty window for the boundary draw: the last `penalty_last_n` tokens of the WHOLE
7089 // prompt, which is what the cold session's own burst sees (Item 2's window). Built
7090 // after the suffix joins `committed` below.
7091 let mut boundary_capture: Option<SpecBoundaryCapture> = None;
7092 if !suffix.is_empty() {
7093 // ---- SUFFIX FEED, mirroring prefill_tick's program selection exactly ----
7094 // From here on the trunk cache mutates: failures return Err((None, _)) and
7095 // the worker serves the request cold-plain instead of reusing the carrier.
7096 let dirty =
7097 |msg: String| -> Result<SpecSession, (Option<Cache>, String)> { Err((None, msg)) };
7098 let n_embd = self.cfg.n_embd as usize;
7099 let t = suffix.len();
7100 let mut h_rows = match e.uninit(t * n_embd) {
7101 Ok(b) => b,
7102 Err(err) => return fail(cache, format!("suffix hidden buffer alloc: {err}")),
7103 };
7104 let mut feed_logits = Vec::new();
7105 let batched = t >= crate::hybrid_forward::PRIME_MIN_T
7106 && std::env::var("MEMRA_PRIME_TOKENWISE").is_err()
7107 && !e.frozen_cpu_experts_prefer_tokenwise_prime();
7108 if batched {
7109 // prefill_tick's prime arm: one request-level prime_cache call.
7110 match self.prime_cache(e, suffix, &mut cache, 0) {
7111 Ok((l, _h_seed, hiddens)) => {
7112 if let Err(err) = e.copy_into(&mut h_rows, 0, &hiddens, t * n_embd) {
7113 return dirty(format!("suffix hidden copy: {err}"));
7114 }
7115 feed_logits = l;
7116 }
7117 Err(err) => return dirty(format!("suffix prime failed: {err}")),
7118 }
7119 } else {
7120 // prefill_tick's tokenwise arm: eager decode_step, one token at a time.
7121 for (i, &tok) in suffix.iter().enumerate() {
7122 match self.decode_step_h(e, tok, &mut cache) {
7123 Ok((l, h)) => {
7124 if let Err(err) = e.copy_into(&mut h_rows, i * n_embd, &h, n_embd) {
7125 return dirty(format!("suffix hidden copy: {err}"));
7126 }
7127 feed_logits = l;
7128 }
7129 Err(err) => return dirty(format!("suffix decode_step failed: {err}")),
7130 }
7131 }
7132 }
7133 // Draft-scratch fill for the suffix rows, predecessor-paired: row `pos` reads
7134 // the entry's boundary anchor (zeros fallback — acceptance-only), row `pos+i`
7135 // reads h_rows[i-1]. Chunked like the generate path's fill (transients scale
7136 // with T). Fill failures are acceptance-only — truncate to the restored rows
7137 // and continue; the burst's own set_len keeps the invariant.
7138 let mtp = self.mtp.as_ref().expect("mtp checked above");
7139 let (embd_qt, embd_rb) = self.embd.qt_and_row_bytes(n_embd);
7140 let embd_gpu = if spec_host_embd() {
7141 None
7142 } else {
7143 Some(
7144 self.embd_gpu
7145 .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload")),
7146 )
7147 };
7148 let embd_dev = embd_gpu.map(|g| (g, embd_qt, embd_rb));
7149 let fill_chunk = 4096usize;
7150 let mut filled = true;
7151 let mut start = 0usize;
7152 'fill: while start < t {
7153 let end = (start + fill_chunk).min(t);
7154 let tc = end - start;
7155 let Ok(mut phs) = e.zeros(tc * n_embd) else {
7156 filled = false;
7157 break 'fill;
7158 };
7159 let (src_lo, dst_off, n_copy) = if start == 0 {
7160 (0, n_embd, (tc - 1) * n_embd)
7161 } else {
7162 ((start - 1) * n_embd, 0, tc * n_embd)
7163 };
7164 if start == 0 {
7165 if let Some(lh) = last_h_dev.as_ref() {
7166 if e.copy_into(&mut phs, 0, lh, n_embd).is_err() {
7167 filled = false;
7168 break 'fill;
7169 }
7170 }
7171 }
7172 if n_copy > 0
7173 && e.copy_view_into(
7174 &mut phs,
7175 dst_off,
7176 &h_rows.slice(src_lo..src_lo + n_copy),
7177 n_copy,
7178 )
7179 .is_err()
7180 {
7181 filled = false;
7182 break 'fill;
7183 }
7184 if self
7185 .mtp_kv_fill(
7186 e,
7187 mtp,
7188 &suffix[start..end],
7189 &phs,
7190 pos + start,
7191 &mut scratch,
7192 embd_dev,
7193 )
7194 .is_err()
7195 {
7196 filled = false;
7197 break 'fill;
7198 }
7199 start = end;
7200 }
7201 if !filled {
7202 // acceptance-only: drafts over missing suffix rows are cheap and wrong,
7203 // so keep only the restored rows resident and let verify arbitrate.
7204 if let Err(err) = scratch.set_len(e, pos) {
7205 return dirty(format!("scratch truncation after failed fill: {err}"));
7206 }
7207 }
7208 // EXTENDED-ENTRY PUBLICATION (lane/sampled-spec-quality, Item 3 — the fix for
7209 // "a restored spec session never publishes an extended entry", SAMPLED-HIT.md
7210 // finding (d)). Pre-lane, publication was armed only for COLD sessions
7211 // (`spec_resumed == 0` in the worker) and both engine capture sites require a
7212 // non-continuation burst — but a converted hit's first burst IS a continuation,
7213 // so a growing conversation learned exactly ONE boundary and turn 3 could never
7214 // hit a longer prefix than turn 2 did.
7215 //
7216 // WHERE, and why it is safe here: `cache.pos == prefix + suffix` at this exact
7217 // line — the trunk is primed over the whole prompt, nothing is generated, and the
7218 // draft plane rows [0..prompt) are filled just above. That is a complete
7219 // whole-entry boundary (`pos == fed_len`), the same shape the cold seed capture
7220 // publishes; the worker's existing publication sweep picks it up because it is
7221 // keyed on `boundary_capture.is_some()` and is sampler- and resume-independent.
7222 // NOT the partial-restore hazard: the boundary is this session's own prompt END,
7223 // never mid-entry, so `entry_pos != fed_len` still refuses on the way back in.
7224 // Failure is SILENT by design (the turn_ckpt / boundary-capture convention):
7225 // publication is an optimization, never a correctness dependency.
7226 if spec_restore_republish_on() {
7227 debug_assert_eq!(
7228 cache.pos,
7229 pos + t,
7230 "extended-entry capture must sit at the restored session's prompt end",
7231 );
7232 if let Ok(snap) = cache.snapshot(e) {
7233 boundary_capture = Some(SpecBoundaryCapture {
7234 snap,
7235 pos: pos + t,
7236 logits: feed_logits.clone(),
7237 last_h: capture_boundary_hidden(e, &h_rows, t, n_embd),
7238 });
7239 }
7240 }
7241 // continuation seed: the feed's boundary logits ARE the plain path's boundary
7242 // logits (same program), so greedy's argmax here is plain's first emitted token,
7243 // and the sampled draw is the cold sampled session's own first token.
7244 next_pred = Some(if sampled {
7245 let sp = sampling.expect("sampled implies a sampler");
7246 // `committed` is still the restored prefix here; the suffix joins it below —
7247 // so this is the last-N window over the WHOLE prompt, exactly the cold
7248 // session's own window at its first token.
7249 let hist = pen_window_seed(&committed, suffix, sp.penalty_last_n);
7250 match sample_boundary_token(
7251 e,
7252 &feed_logits,
7253 &sp,
7254 &hist,
7255 &mut sctr,
7256 "restore-suffix-feed",
7257 ) {
7258 Ok(t) => t,
7259 // the trunk is already fed: hand nothing back, the worker serves the
7260 // request cold-plain. Never fall back to an argmax — that would put a
7261 // greedy token in a sampled stream to save a slow path.
7262 Err(err) => {
7263 return dirty(format!("boundary token draw failed: {err}"));
7264 }
7265 }
7266 } else {
7267 argmax(&feed_logits) as u32
7268 });
7269 let mut lh = match e.uninit(n_embd) {
7270 Ok(b) => b,
7271 Err(err) => return dirty(format!("boundary hidden alloc: {err}")),
7272 };
7273 if let Err(err) = e.copy_view_into(
7274 &mut lh,
7275 0,
7276 &h_rows.slice((t - 1) * n_embd..t * n_embd),
7277 n_embd,
7278 ) {
7279 return dirty(format!("boundary hidden copy: {err}"));
7280 }
7281 last_h_dev = Some(lh);
7282 committed.extend_from_slice(suffix);
7283 } else {
7284 // FULL-COVER shape (empty suffix — the identical-repeat / agent-loop shape): the
7285 // ENTRY's boundary logits are the boundary row, and this is the token the cold
7286 // session emits from that same row. Owned here rather than in the worker so the
7287 // sampled draw cannot be half-applied on one shape (the worker used to argmax it).
7288 if boundary_logits.is_empty() {
7289 return fail(
7290 cache,
7291 "full-cover restore without the entry's boundary logits".into(),
7292 );
7293 }
7294 next_pred = Some(if sampled {
7295 let sp = sampling.expect("sampled implies a sampler");
7296 let hist = pen_window_seed(&committed, &[], sp.penalty_last_n);
7297 match sample_boundary_token(
7298 e,
7299 boundary_logits,
7300 &sp,
7301 &hist,
7302 &mut sctr,
7303 "restore-full-cover",
7304 ) {
7305 Ok(t) => t,
7306 // nothing has been mutated on this shape — hand the carrier back and let
7307 // the hit serve PLAIN (the banked pre-lane path).
7308 Err(err) => {
7309 return fail(cache, format!("boundary token draw failed: {err}"));
7310 }
7311 }
7312 } else {
7313 argmax(boundary_logits) as u32
7314 });
7315 }
7316 Ok(SpecSession {
7317 cache,
7318 scratch,
7319 committed,
7320 last_h: last_h_dev,
7321 next_pred,
7322 sctr,
7323 uctr: 0,
7324 draft_ctx: None,
7325 pending_tok: None,
7326 turn_ckpt: None,
7327 telem: SpecTelemetryCounters::default(),
7328 capture_at: None,
7329 boundary_capture,
7330 })
7331 }
7332
7333 /// Forced-gate exact state comparison. This intentionally reads the real live prefixes from
7334 /// their owning PP devices: matching emitted ids alone would miss a stale `len_d`, recurrent
7335 /// snapshot, or draft-KV row that only corrupts the following round.
7336 pub fn optipipe_compare_session_state(
7337 &self,
7338 e: &Engine,
7339 reference: &SpecSession,
7340 candidate: &SpecSession,
7341 ) -> Result<OptiForkStateIdentity, Box<dyn std::error::Error>> {
7342 fn fail(what: &str) -> Box<dyn std::error::Error> {
7343 format!("optipipe state mismatch: {what}").into()
7344 }
7345 fn same_f32(a: &[f32], b: &[f32]) -> bool {
7346 a.len() == b.len() && a.iter().zip(b).all(|(x, y)| x.to_bits() == y.to_bits())
7347 }
7348 fn compare_layers(
7349 es: &Engine,
7350 range: std::ops::Range<usize>,
7351 reference: &SpecSession,
7352 candidate: &SpecSession,
7353 report: &mut OptiForkStateIdentity,
7354 ) -> Result<(), Box<dyn std::error::Error>> {
7355 for il in range {
7356 match (&reference.cache.kv[il], &candidate.cache.kv[il]) {
7357 (Some(a), Some(b)) => {
7358 if a.len != b.len {
7359 return Err(fail(&format!(
7360 "layer {il} host KV len {} != {}",
7361 a.len, b.len
7362 )));
7363 }
7364 let ad = es.dtoh_i32(&a.len_d)?;
7365 let bd = es.dtoh_i32(&b.len_d)?;
7366 if ad != bd || ad.first().copied() != Some(a.len as i32) {
7367 return Err(fail(&format!(
7368 "layer {il} device KV len {ad:?} != {bd:?} (host={})",
7369 a.len,
7370 )));
7371 }
7372 let kb = a.len * a.k_tok_bytes;
7373 let vb = a.len * a.v_tok_bytes;
7374 if kb > 0 {
7375 let ak = es.dtoh_u8_view(&a.k.slice(0..kb))?;
7376 let bk = es.dtoh_u8_view(&b.k.slice(0..kb))?;
7377 if ak != bk {
7378 let at = ak.iter().zip(&bk).position(|(x, y)| x != y).unwrap();
7379 return Err(fail(&format!(
7380 "layer {il} K bytes at byte {at} row {} offset {}: {} != {}",
7381 at / a.k_tok_bytes,
7382 at % a.k_tok_bytes,
7383 ak[at],
7384 bk[at],
7385 )));
7386 }
7387 }
7388 if vb > 0 {
7389 let av = es.dtoh_u8_view(&a.v.slice(0..vb))?;
7390 let bv = es.dtoh_u8_view(&b.v.slice(0..vb))?;
7391 if av != bv {
7392 let at = av.iter().zip(&bv).position(|(x, y)| x != y).unwrap();
7393 return Err(fail(&format!(
7394 "layer {il} V bytes at byte {at} row {} offset {}: {} != {}",
7395 at / a.v_tok_bytes,
7396 at % a.v_tok_bytes,
7397 av[at],
7398 bv[at],
7399 )));
7400 }
7401 }
7402 report.trunk_kv_bytes += kb + vb;
7403 }
7404 (None, None) => {}
7405 _ => return Err(fail(&format!("layer {il} KV presence"))),
7406 }
7407 match (&reference.cache.recur[il], &candidate.cache.recur[il]) {
7408 (Some(a), Some(b)) => {
7409 let ac = es.dtoh(&a.conv_state)?;
7410 let bc = es.dtoh(&b.conv_state)?;
7411 if !same_f32(&ac, &bc) {
7412 return Err(fail(&format!("layer {il} conv state")));
7413 }
7414 let as_ = es.dtoh(&a.ssm_state)?;
7415 let bs = es.dtoh(&b.ssm_state)?;
7416 if !same_f32(&as_, &bs) {
7417 return Err(fail(&format!("layer {il} SSM state")));
7418 }
7419 report.recurrent_bytes += (ac.len() + as_.len()) * 4;
7420 }
7421 (None, None) => {}
7422 _ => return Err(fail(&format!("layer {il} recurrent presence"))),
7423 }
7424 }
7425 Ok(())
7426 }
7427
7428 if reference.committed != candidate.committed {
7429 return Err(fail("committed token ids"));
7430 }
7431 if reference.cache.pos != candidate.cache.pos
7432 || reference.cache.max_ctx != candidate.cache.max_ctx
7433 {
7434 return Err(fail("cache pos/capacity"));
7435 }
7436 if reference.pending_tok != candidate.pending_tok
7437 || reference.next_pred != candidate.next_pred
7438 || reference.sctr != candidate.sctr
7439 || reference.uctr != candidate.uctr
7440 {
7441 return Err(fail("pending/prediction/counter tail"));
7442 }
7443
7444 let mut report = OptiForkStateIdentity::default();
7445 if let Some(fence) = crate::pp::pp_cuts(self.layers.len()) {
7446 let rt = crate::pp::PpNRt::get(e)?;
7447 for stage in 0..rt.n_stages() {
7448 let _scope = rt.enter(stage);
7449 compare_layers(
7450 rt.engine(stage, e),
7451 fence[stage]..fence[stage + 1],
7452 reference,
7453 candidate,
7454 &mut report,
7455 )?;
7456 }
7457 } else {
7458 compare_layers(e, 0..self.layers.len(), reference, candidate, &mut report)?;
7459 }
7460
7461 let (a, b) = (&reference.scratch.kv, &candidate.scratch.kv);
7462 if a.len != b.len || e.dtoh_i32(&a.len_d)? != e.dtoh_i32(&b.len_d)? {
7463 return Err(fail("draft scratch length"));
7464 }
7465 let kb = a.len * a.k_tok_bytes;
7466 let vb = a.len * a.v_tok_bytes;
7467 if kb > 0 && e.dtoh_u8_view(&a.k.slice(0..kb))? != e.dtoh_u8_view(&b.k.slice(0..kb))? {
7468 return Err(fail("draft scratch K bytes"));
7469 }
7470 if vb > 0 && e.dtoh_u8_view(&a.v.slice(0..vb))? != e.dtoh_u8_view(&b.v.slice(0..vb))? {
7471 return Err(fail("draft scratch V bytes"));
7472 }
7473 report.scratch_kv_bytes = kb + vb;
7474
7475 match (&reference.last_h, &candidate.last_h) {
7476 (Some(a), Some(b)) => {
7477 let ah = e.dtoh(a)?;
7478 let bh = e.dtoh(b)?;
7479 if !same_f32(&ah, &bh) {
7480 return Err(fail("last hidden/seed bytes"));
7481 }
7482 report.hidden_bytes = ah.len() * 4;
7483 }
7484 (None, None) => {}
7485 _ => return Err(fail("last hidden/seed presence")),
7486 }
7487 Ok(report)
7488 }
7489
7490 /// SESSION-AFFINITY REWIND (lane/session-affinity, 2026-08-05): roll `sess` back to its
7491 /// retained prompt-end checkpoint, so a request whose prompt matches
7492 /// `committed[..rewind_pos()]` exactly can resume there and prime only its own delta.
7493 ///
7494 /// EXACTNESS. After this returns, the session is byte-for-byte the state it was in AT that
7495 /// boundary: full-attn KV truncated to it (append-only, position-addressed), GDN conv/ssm
7496 /// restored from the device copy taken there, draft scratch length reset, `committed`
7497 /// truncated, `last_h` = the boundary's predecessor anchor. That is precisely the state a
7498 /// fresh prime of `committed[..pos]` would have produced, so the following suffix prime and
7499 /// every burst after it are identical to a cold run of the same token stream — the
7500 /// committed-tokens-authoritative contract.
7501 ///
7502 /// `next_pred` and `pending_tok` are CLEARED: both describe generation past the boundary,
7503 /// which the rewind discards. The caller therefore must supply a non-empty suffix (a
7504 /// rewound session cannot serve an empty-suffix continuation burst — there is nothing to
7505 /// continue). The persistent draft graph survives: it bakes only session-stable pointers
7506 /// (the scratch KV, the resident embedding), none of which the rewind moves.
7507 ///
7508 /// The checkpoint is CONSUMED (`turn_ckpt` taken): its snapshot buffers are freed here, and
7509 /// this turn's own prime installs a fresh one at the new prompt end. Returns the position
7510 /// rewound to, or `None` when the session holds no checkpoint (caller: full re-prime).
7511 pub fn spec_rewind_to_checkpoint(
7512 &self,
7513 e: &Engine,
7514 sess: &mut SpecSession,
7515 ) -> Result<Option<usize>, Box<dyn std::error::Error>> {
7516 if sess.turn_ckpt.as_ref().is_some_and(|ckpt| {
7517 !sess.cache.can_rollback(&ckpt.snap, 0) || !sess.scratch.can_rewind_to(ckpt.pos)
7518 }) {
7519 return Err(
7520 "SWA ring rewind checkpoint has been lapped; full re-prime required".into(),
7521 );
7522 }
7523 let Some(ckpt) = sess.turn_ckpt.take() else {
7524 return Ok(None);
7525 };
7526 assert!(
7527 ckpt.pos <= sess.committed.len(),
7528 "checkpoint past committed ({} > {})",
7529 ckpt.pos,
7530 sess.committed.len()
7531 );
7532 // Restore through each layer's owning engine. A single primary-engine rollback is not
7533 // sufficient when the serving cache is stage-owned under cross-device PP.
7534 crate::pp::restore_cache_checkpoint(e, &self.cfg, None, &mut sess.cache, &ckpt.snap)?;
7535 debug_assert_eq!(
7536 sess.cache.pos, ckpt.pos,
7537 "rollback landed off the checkpoint"
7538 );
7539 sess.scratch.set_len(e, ckpt.pos)?;
7540 sess.committed.truncate(ckpt.pos);
7541 sess.last_h = Some(ckpt.last_h);
7542 sess.next_pred = None;
7543 sess.pending_tok = None;
7544 Ok(Some(ckpt.pos))
7545 }
7546
7547 /// Grow a parked speculative session to `target_cap` and rewind it to its retained turn
7548 /// checkpoint without re-priming the checkpoint prefix.
7549 ///
7550 /// The trunk cache is restored exactly like a plain grown cache: append-only full-attention
7551 /// KV rows come from the parked cache, while recurrent state comes from the checkpoint's
7552 /// owned snapshot. The MTP scratch is also context-linear and its rows below the checkpoint
7553 /// remain authoritative, so they are copied into a fresh larger scratch before its length is
7554 /// truncated. Pointer-baking draft graphs are dropped and recaptured on the next burst.
7555 ///
7556 /// All fallible work completes before `sess` is mutated. A failed allocation or copy leaves
7557 /// the parked session intact, allowing the caller one reclaim-and-retry attempt.
7558 pub fn spec_grow_and_rewind_to_checkpoint(
7559 &self,
7560 e: &Engine,
7561 sess: &mut SpecSession,
7562 target_cap: usize,
7563 ) -> Result<Option<usize>, Box<dyn std::error::Error>> {
7564 if target_cap <= sess.cache.max_ctx {
7565 return self.spec_rewind_to_checkpoint(e, sess);
7566 }
7567 let Some(ckpt) = sess.turn_ckpt.as_ref() else {
7568 return Ok(None);
7569 };
7570 if ckpt.pos == 0 || ckpt.pos > sess.committed.len() {
7571 return Err(format!(
7572 "checkpoint pos {} outside committed length {}",
7573 ckpt.pos,
7574 sess.committed.len(),
7575 )
7576 .into());
7577 }
7578 if ckpt.pos > target_cap {
7579 return Err(format!(
7580 "checkpoint pos {} exceeds grown capacity {target_cap}",
7581 ckpt.pos,
7582 )
7583 .into());
7584 }
7585
7586 let mut grown_cache = crate::pp::new_cache(e, &self.cfg, target_cap)?;
7587 let mut grown_scratch = MtpScratch::new(
7588 e,
7589 &self.cfg,
7590 target_cap,
7591 self.mtp.as_ref().and_then(|m| m.geom.as_ref()),
7592 )?;
7593 crate::pp::restore_cache_checkpoint(
7594 e,
7595 &self.cfg,
7596 Some(&sess.cache),
7597 &mut grown_cache,
7598 &ckpt.snap,
7599 )?;
7600
7601 let src = &sess.scratch.kv;
7602 let dst = &mut grown_scratch.kv;
7603 if ckpt.pos > src.len
7604 || src.kv_dim_k != dst.kv_dim_k
7605 || src.kv_dim_v != dst.kv_dim_v
7606 || src.k_tok_bytes != dst.k_tok_bytes
7607 || src.v_tok_bytes != dst.v_tok_bytes
7608 {
7609 return Err(format!(
7610 "checkpoint draft layout mismatch (pos {}, source len {})",
7611 ckpt.pos, src.len,
7612 )
7613 .into());
7614 }
7615 let kb = ckpt.pos * src.k_tok_bytes;
7616 let vb = ckpt.pos * src.v_tok_bytes;
7617 if kb > 0 {
7618 e.copy_u8_into(&mut dst.k, 0, &src.k, kb)?;
7619 }
7620 if vb > 0 {
7621 e.copy_u8_into(&mut dst.v, 0, &src.v, vb)?;
7622 }
7623 grown_scratch.set_len(e, ckpt.pos)?;
7624 // The old scratch is dropped immediately after publication below. Bound its D2D reads
7625 // first; growth happens once per rewritten turn, outside the decode hot loop.
7626 e.stream().synchronize()?;
7627
7628 let ckpt = sess
7629 .turn_ckpt
7630 .take()
7631 .expect("checkpoint remained present through transactional grow");
7632 let pos = ckpt.pos;
7633 sess.cache = grown_cache;
7634 sess.scratch = grown_scratch;
7635 sess.committed.truncate(pos);
7636 sess.last_h = Some(ckpt.last_h);
7637 sess.next_pred = None;
7638 sess.pending_tok = None;
7639 sess.draft_ctx = None;
7640 debug_assert_eq!(sess.cache.pos, pos, "grown rewind landed off checkpoint");
7641 debug_assert_eq!(
7642 sess.scratch.kv.len, pos,
7643 "grown draft rewind landed off checkpoint"
7644 );
7645 Ok(Some(pos))
7646 }
7647
7648 /// Commit a carried pending bonus (see SpecSession::pending_tok): one T=1 trunk pass
7649 /// (its logits' argmax becomes next_pred) + the draft-KV fill at the carried anchor —
7650 /// byte-identical to the pre-carry session tail. Required before a non-empty-suffix
7651 /// prime, a sampled turn, or parking a session for pool reuse. No-op without a pending.
7652 /// `sampling` is the sampler of the request that will CONSUME the resulting `next_pred`
7653 /// (lane/sampled-spec-quality): this is a boundary site like any other, so a sampled
7654 /// consumer must get a DRAWN token, not an argmax. Pass `None` from the park/demote
7655 /// callers — a pending only ever exists on the GREEDY tail, and the consumer of a
7656 /// park-time flush is a future request whose sampler is not knowable here (residual
7657 /// named at the pool-resume probe in worker.rs and in SAMPLED-QUALITY.md).
7658 pub fn spec_flush_pending(
7659 &self,
7660 e: &Engine,
7661 sess: &mut SpecSession,
7662 sampling: Option<SpecSampling>,
7663 ) -> Result<(), Box<dyn std::error::Error>> {
7664 let Some(b) = sess.pending_tok.take() else {
7665 return Ok(());
7666 };
7667 let mtp = self
7668 .mtp
7669 .as_ref()
7670 .expect("pending carry requires an MTP head");
7671 let n_embd = self.cfg.n_embd as usize;
7672 let (embd_qt, embd_rb) = self.embd.qt_and_row_bytes(n_embd);
7673 let embd_gpu = if spec_host_embd() {
7674 None
7675 } else {
7676 Some(
7677 self.embd_gpu
7678 .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload")),
7679 )
7680 };
7681 let embd_dev = embd_gpu.map(|g| (g, embd_qt, embd_rb));
7682 let pos_b = sess.cache.pos;
7683 sess.scratch.set_len(e, pos_b)?;
7684 let (lg_b, hb) = self.spec_target_step_h(e, b, &mut sess.cache)?;
7685 sess.next_pred = Some(match sampling {
7686 Some(sp) if sp.temp > 0.0 && spec_sampled_boundary_on() => {
7687 // window includes `b` itself: it is committed by this pass, and the pre-lane
7688 // code never counted a boundary token in the penalty history at all.
7689 let hist = pen_window_seed(&sess.committed, &[b], sp.penalty_last_n);
7690 sample_boundary_token(e, &lg_b, &sp, &hist, &mut sess.sctr, "flush-pending")?
7691 }
7692 _ => argmax(&lg_b) as u32,
7693 });
7694 let anchor = sess
7695 .last_h
7696 .as_ref()
7697 .expect("pending carry requires last_h (the predecessor-row anchor)");
7698 self.mtp_kv_fill(e, mtp, &[b], anchor, pos_b, &mut sess.scratch, embd_dev)?;
7699 sess.last_h = Some(hb);
7700 sess.committed.push(b);
7701 Ok(())
7702 }
7703
7704 /// Solo target feed used only at speculative round boundaries. Step35 serving made its
7705 /// staged batched B=1 graph authoritative, so a speculative session must enter and leave
7706 /// rounds through that same graph. Other model families keep their eager T=1 contract.
7707 fn spec_target_step_h(
7708 &self,
7709 e: &Engine,
7710 token: u32,
7711 cache: &mut Cache,
7712 ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
7713 if self.cfg.step35.is_none() && !self.qwen35_serving_class() {
7714 return self.decode_step_h(e, token, cache);
7715 }
7716 let pos0 = cache.pos;
7717 let (logits, hidden) = self.decode_step_t_core(e, &[token], pos0, cache, None, None)?;
7718 Ok((e.dtoh(&logits)?, hidden))
7719 }
7720
7721 /// The archs whose LIVE B=1 serving runs the generic BATCHED numeric class (decode_step_batch
7722 /// walk + batched head), so their spec verify must run the SAME class. MoE learned this
7723 /// 2026-08-14 AM (4b777ccc5); the dense hybrid reproduced the identical near-tie flip class
7724 /// the same day on Qwen3.8-27B — eager-class verify logits drift from batched-class serving
7725 /// logits ("1 ULP at layer 2 → 2.3e-1 logit maxdiff at the head"), and the GDN recurrence
7726 /// carries the drift until a near-tie flips deep in generation. One predicate so the five
7727 /// dispatch sites cannot drift apart again.
7728 /// Draft-graph head admissibility (lane/draftcost-moe, 2026-08-20): the capture body
7729 /// (`mtp_head_forward_cap`) supports Dense heads AND resident-MoE heads
7730 /// (`Ffn::Moe(m) if m.dev_exps.is_some()`); non-resident MoE still refuses inside the
7731 /// capture and the caller falls back to the eager chain by design. Trunk FFN class is
7732 /// irrelevant — the graph body is the HEAD forward only. One predicate for all three
7733 /// eligibility sites so they cannot drift (the qwen35_serving_class lesson).
7734 fn mtp_graph_capturable(&self) -> bool {
7735 self.mtp
7736 .as_ref()
7737 .map(|m| match &m.ffn {
7738 crate::hybrid::Ffn::Dense { .. } => true,
7739 crate::hybrid::Ffn::Moe(mo) => mo.dev_exps.is_some(),
7740 })
7741 .unwrap_or(false)
7742 }
7743
7744 fn qwen35_serving_class(&self) -> bool {
7745 matches!(
7746 self.cfg.arch,
7747 memra_gguf::config::Arch::Qwen35 | memra_gguf::config::Arch::Qwen35Moe
7748 )
7749 }
7750
7751 /// Reduced-matrix admission for increment 1. This deliberately does not change the PP-2
7752 /// serving policy: the worker calls it only after `MEMRA_SPEC_PIPE=1` and an explicit spec
7753 /// session already exist.
7754 pub fn spec_pipe_available(&self, e: &Engine) -> bool {
7755 if std::env::var("MEMRA_SPEC_PIPE").as_deref() != Ok("1")
7756 || !spec_devacc()
7757 || spec_replay_env_enabled()
7758 || spec_stream()
7759 || std::env::var("MEMRA_SPEC_ADAPT").as_deref() == Ok("1")
7760 || std::env::var("MEMRA_SPEC_PMIN0").as_deref() == Ok("1")
7761 || std::env::var("MEMRA_SPEC_PP_ANATOMY").as_deref() == Ok("1")
7762 || std::env::var("MEMRA_SPEC_PMIN")
7763 .ok()
7764 .and_then(|v| v.parse::<f32>().ok())
7765 .unwrap_or(0.0)
7766 > 0.0
7767 || self.is_gemma4_e4b()
7768 || self.cfg.gemma4.is_some()
7769 || self.mtp.is_none()
7770 {
7771 return false;
7772 }
7773 let Some(cuts) = crate::pp::pp_cuts(self.layers.len()) else {
7774 return false;
7775 };
7776 if cuts.len() != 3 || crate::pp::pp2_streams_off() || !crate::pp::spec_pp_on() {
7777 return false;
7778 }
7779 crate::pp::PpNRt::get(e)
7780 .map(|rt| rt.n_stages() == 2 && rt.cross_device())
7781 .unwrap_or(false)
7782 }
7783
7784 /// Two warm greedy continuation bursts over one PP-2 interval coordinator. The two existing
7785 /// `generate_spec_inner2` call stacks own all per-session round locals; only phase issue order
7786 /// changes. No callback is accepted in increment 1 — the worker publishes each completed burst.
7787 #[allow(clippy::too_many_arguments)]
7788 pub fn generate_spec_session_pair(
7789 &self,
7790 e: &Engine,
7791 sess_a: &mut SpecSession,
7792 max_new_a: usize,
7793 k_a: usize,
7794 sess_b: &mut SpecSession,
7795 max_new_b: usize,
7796 k_b: usize,
7797 ) -> Result<((Vec<u32>, usize, usize), (Vec<u32>, usize, usize)), Box<dyn std::error::Error>>
7798 {
7799 if !self.spec_pipe_available(e) {
7800 return Err("two-session speculative pipeline is outside its reduced matrix".into());
7801 }
7802 if max_new_a == 0 || max_new_b == 0 || k_a == 0 || k_b == 0 {
7803 return Err(
7804 "two-session speculative pipeline requires non-empty positive-K bursts".into(),
7805 );
7806 }
7807 for sess in [&*sess_a, &*sess_b] {
7808 if sess.committed.is_empty()
7809 || sess.last_h.is_none()
7810 || (sess.next_pred.is_none() && sess.pending_tok.is_none())
7811 {
7812 return Err("two-session speculative pipeline requires warm continuations".into());
7813 }
7814 }
7815
7816 let graph_ok = std::env::var("MEMRA_SPEC_NOGRAPH").is_err()
7817 && !spec_host_embd()
7818 && self.mtp_graph_capturable()
7819 && !crate::model::full_prec_enabled();
7820 let graph_a = graph_ok && k_a + 2 < 96;
7821 let graph_b = graph_ok && k_b + 2 < 96;
7822 let was_tracking = e.ctx().is_event_tracking();
7823 if (graph_a || graph_b) && was_tracking {
7824 unsafe {
7825 e.ctx().disable_event_tracking();
7826 }
7827 }
7828
7829 static LOGGED: std::sync::Once = std::sync::Once::new();
7830 LOGGED.call_once(|| {
7831 eprintln!("[spec-pipe] two-session PP-2 continuation pipeline engaged");
7832 });
7833 let sync = std::sync::Arc::new(SpecPipeSync::new());
7834 let lane_a = SpecPipeLane {
7835 sync: sync.clone(),
7836 lane: 0,
7837 };
7838 let lane_b = SpecPipeLane { sync, lane: 1 };
7839 let mut sess_b_ptr = SpecPipeSessionPtr(sess_b as *mut SpecSession);
7840 let (result_a, result_b) = std::thread::scope(|scope| {
7841 let b = scope.spawn(move || {
7842 let mut finish = SpecPipeFinish::new(&lane_b);
7843 let sess_b = unsafe { sess_b_ptr.get_mut() };
7844 let result = e
7845 .ctx()
7846 .bind_to_thread()
7847 .map_err(|err| err.to_string())
7848 .and_then(|_| {
7849 self.generate_spec_inner2(
7850 e,
7851 &[],
7852 max_new_b,
7853 k_b,
7854 graph_b,
7855 Some(sess_b),
7856 None,
7857 None,
7858 None,
7859 None,
7860 Some(&lane_b),
7861 )
7862 .map_err(|err| err.to_string())
7863 });
7864 finish.close(result.is_err());
7865 result
7866 });
7867 let mut finish = SpecPipeFinish::new(&lane_a);
7868 let result_a = self.generate_spec_inner2(
7869 e,
7870 &[],
7871 max_new_a,
7872 k_a,
7873 graph_a,
7874 Some(sess_a),
7875 None,
7876 None,
7877 None,
7878 None,
7879 Some(&lane_a),
7880 );
7881 finish.close(result_a.is_err());
7882 let result_b = b
7883 .join()
7884 .map_err(|_| "paired speculative session B panicked".to_string())
7885 .and_then(|r| r);
7886 (result_a, result_b)
7887 });
7888
7889 if (graph_a || graph_b) && was_tracking {
7890 unsafe {
7891 e.ctx().enable_event_tracking();
7892 }
7893 }
7894 let result_a = result_a?;
7895 let result_b = result_b.map_err(|err| -> Box<dyn std::error::Error> { err.into() })?;
7896 Ok((result_a, result_b))
7897 }
7898
7899 /// One spec-decode turn on a live session. `suffix` = the NEW tokens only (turn N+1's user
7900 /// message rendered through the chat template continuation). Returns (new tokens emitted,
7901 /// drafted, accepted); session.committed grows by suffix + emitted.
7902 pub fn generate_spec_session(
7903 &self,
7904 e: &Engine,
7905 sess: &mut SpecSession,
7906 suffix: &[u32],
7907 max_new: usize,
7908 k: usize,
7909 ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
7910 self.generate_spec_session_sampled(e, sess, suffix, max_new, k, None, None)
7911 }
7912
7913 /// Serve-path sampled spec: routes the burst through the rejection-sampling verify with
7914 /// per-SESSION Philox continuity (sess.sctr/uctr). None = env-driven (CLI) or greedy.
7915 /// Filters (top-k/p/min-p) apply SYMMETRICALLY to draft q and verify p — distribution-exact
7916 /// for the filtered target (feat/filtered-spec).
7917 ///
7918 /// `on_commit` (sse-cadence, 2026-08-05): called with each newly-emitted slice of the
7919 /// output — once right after the prime's first token, then once per round commit — so a
7920 /// streaming caller can flush text at round cadence instead of once per burst. The slices
7921 /// are disjoint, in order, and concatenate to exactly the returned token vec. Emission-
7922 /// timing only: token bytes, session state, and exactness are untouched.
7923 ///
7924 /// The returned bool is a CONTINUE-VERDICT (admission yield, 2026-08-06): `false` ends
7925 /// the burst at the current round boundary, exactly as if `max_new` had been reached —
7926 /// the caller's scheduler regains control without waiting the burst out. Burst size is
7927 /// content-neutral (spec-levers battery), so an early exit moves WHEN the burst returns,
7928 /// never what tokens say. The slice may be EMPTY (a poll-only boundary — round-stream
7929 /// drains and the defensive tail flush can land with nothing new committed).
7930 #[allow(clippy::too_many_arguments)]
7931 pub fn generate_spec_session_sampled(
7932 &self,
7933 e: &Engine,
7934 sess: &mut SpecSession,
7935 suffix: &[u32],
7936 max_new: usize,
7937 k: usize,
7938 sampling: Option<SpecSampling>,
7939 on_commit: Option<&mut dyn FnMut(&[u32]) -> bool>,
7940 ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
7941 self.generate_spec_session_sampled_prime_split(
7942 e, sess, suffix, max_new, k, sampling, None, on_commit,
7943 )
7944 }
7945
7946 /// Serve-only cold-prime segmentation twin. `prime_split` is the same stable boundary the
7947 /// plain worker would honor before entering its sub-floor tokenwise tail; warm continuations
7948 /// pass `None` and stay on the existing zero-prime path.
7949 #[allow(clippy::too_many_arguments)]
7950 pub fn generate_spec_session_sampled_prime_split(
7951 &self,
7952 e: &Engine,
7953 sess: &mut SpecSession,
7954 suffix: &[u32],
7955 max_new: usize,
7956 k: usize,
7957 sampling: Option<SpecSampling>,
7958 prime_split: Option<usize>,
7959 on_commit: Option<&mut dyn FnMut(&[u32]) -> bool>,
7960 ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
7961 self.generate_spec_session_constrained_prime_split(
7962 e,
7963 sess,
7964 suffix,
7965 max_new,
7966 k,
7967 sampling,
7968 None,
7969 prime_split,
7970 on_commit,
7971 )
7972 }
7973
7974 /// `generate_spec_session_sampled` + GRAMMAR (constrained decoding, 2026-08-03): the
7975 /// hook truncates acceptance at the first grammar-illegal token AFTER the exactness
7976 /// verify (grammar is an extra rejection rule, ordering like the batched-verify twins)
7977 /// and replaces an illegal bonus with the MASKED argmax of the target's own verify
7978 /// column — token-identical to constrained plain greedy decode. GREEDY only (the
7979 /// worker routes sampled constrained to plain decode). Acceptance under tight grammars
7980 /// may drop (drafter is unconstrained); that is measured, not hidden.
7981 #[allow(clippy::too_many_arguments)]
7982 pub fn generate_spec_session_constrained(
7983 &self,
7984 e: &Engine,
7985 sess: &mut SpecSession,
7986 suffix: &[u32],
7987 max_new: usize,
7988 k: usize,
7989 sampling: Option<SpecSampling>,
7990 constraint: Option<&mut dyn SpecConstraint>,
7991 on_commit: Option<&mut dyn FnMut(&[u32]) -> bool>,
7992 ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
7993 self.generate_spec_session_constrained_prime_split(
7994 e, sess, suffix, max_new, k, sampling, constraint, None, on_commit,
7995 )
7996 }
7997
7998 #[allow(clippy::too_many_arguments)]
7999 pub fn generate_spec_session_constrained_prime_split(
8000 &self,
8001 e: &Engine,
8002 sess: &mut SpecSession,
8003 suffix: &[u32],
8004 max_new: usize,
8005 k: usize,
8006 sampling: Option<SpecSampling>,
8007 constraint: Option<&mut dyn SpecConstraint>,
8008 prime_split: Option<usize>,
8009 on_commit: Option<&mut dyn FnMut(&[u32]) -> bool>,
8010 ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
8011 if constraint.is_some() && sampling.is_some_and(|s| s.temp > 0.0) {
8012 return Err(
8013 "constrained spec decode is greedy-only (worker routes sampled \
8014 constrained to plain decode)"
8015 .into(),
8016 );
8017 }
8018 // PENDING-CARRY entry flush: a carried bonus precedes any new suffix in the sequence,
8019 // so it must commit BEFORE the suffix primes; the sampled path doesn't carry (its
8020 // round-0 accept needs the commit pass's logits). Empty-suffix greedy bursts — the
8021 // serve continuation case — consume the carry in-loop with zero solo passes.
8022 if sess.pending_tok.is_some()
8023 && (!suffix.is_empty() || sampling.map_or(false, |s| s.temp > 0.0))
8024 {
8025 self.spec_flush_pending(e, sess, sampling)?;
8026 }
8027
8028 // FULL_PREC forces the EAGER draft: the graph capture would enclose cuBLASLt f32 GEMV
8029 // (the FloatBf16 else-branches) and a bf16_to_f32 dequant alloc — neither is stream-capture
8030 // safe. Eager rides matmul/matmul_decode_exact, which dequant FloatBf16 on use. (§item 2.)
8031 let graph_draft = std::env::var("MEMRA_SPEC_NOGRAPH").is_err()
8032 && !spec_host_embd()
8033 && self.mtp_graph_capturable()
8034 && k + 2 < 96
8035 && !crate::model::full_prec_enabled();
8036 let was_tracking = e.ctx().is_event_tracking();
8037 if graph_draft && was_tracking {
8038 unsafe {
8039 e.ctx().disable_event_tracking();
8040 }
8041 }
8042 let r = self.generate_spec_inner2(
8043 e,
8044 suffix,
8045 max_new,
8046 k,
8047 graph_draft,
8048 Some(sess),
8049 sampling,
8050 constraint,
8051 on_commit,
8052 prime_split,
8053 None,
8054 );
8055 if graph_draft && was_tracking {
8056 unsafe {
8057 e.ctx().enable_event_tracking();
8058 }
8059 }
8060 let (out, d, a) = r?;
8061 Ok((out, d, a))
8062 }
8063
8064 pub fn generate_spec(
8065 &self,
8066 e: &Engine,
8067 prompt: &[u32],
8068 max_new: usize,
8069 k: usize,
8070 ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
8071 // FULL_PREC forces eager (see generate_spec_session note): CUDA graph capture cannot
8072 // enclose cuBLASLt f32 GEMV or the bf16_to_f32 dequant alloc the FloatBf16 path needs.
8073 let graph_draft = std::env::var("MEMRA_SPEC_NOGRAPH").is_err()
8074 && !spec_host_embd()
8075 && self.mtp_graph_capturable()
8076 && k + 2 < 96
8077 && !crate::model::full_prec_enabled();
8078 if !graph_draft {
8079 return self.generate_spec_inner2(
8080 e, prompt, max_new, k, false, None, None, None, None, None, None,
8081 );
8082 }
8083 let was_tracking = e.ctx().is_event_tracking();
8084 if was_tracking {
8085 unsafe {
8086 e.ctx().disable_event_tracking();
8087 }
8088 }
8089 let r = self.generate_spec_inner2(
8090 e, prompt, max_new, k, true, None, None, None, None, None, None,
8091 );
8092 if was_tracking {
8093 unsafe {
8094 e.ctx().enable_event_tracking();
8095 }
8096 }
8097 r
8098 }
8099
8100 fn generate_spec_inner2(
8101 &self,
8102 e: &Engine,
8103 prompt: &[u32],
8104 max_new: usize,
8105 k: usize,
8106 graph_draft: bool,
8107 mut sess: Option<&mut SpecSession>,
8108 sampling: Option<SpecSampling>,
8109 mut constraint: Option<&mut dyn SpecConstraint>,
8110 mut on_commit: Option<&mut dyn FnMut(&[u32]) -> bool>,
8111 prime_split: Option<usize>,
8112 pipe: Option<&SpecPipeLane>,
8113 ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
8114 assert!(k >= 1, "k must be >= 1");
8115 if let Some(p) = pipe {
8116 p.setup_begin()?;
8117 }
8118 // sse-cadence flush cursor: everything in out[..flushed] has been handed to on_commit.
8119 let mut flushed = 0usize;
8120 // admission yield (2026-08-06): on_commit's continue-verdict; false = end the burst
8121 // at the next round boundary (same exit as max_new reached — the session tail runs).
8122 // Initialized by the unconditional post-prime flush below.
8123 let mut keep_going;
8124 let mtp = self
8125 .mtp
8126 .as_ref()
8127 .expect("generate_spec requires an MTP head (nextn_predict_layers>0)");
8128 let n_vocab = self.output.out_features();
8129 // FR-Spec: the draft head may be TRIMMED (fewer rows than n_vocab); the draft argmax runs
8130 // over the draft vocab and the winning index maps through d2t to a TARGET token id.
8131 // Everything downstream (verify/accept/commit) sees target ids only — exactness unchanged.
8132 let d_vocab = mtp
8133 .shared_head_head
8134 .as_ref()
8135 .unwrap_or(&self.output)
8136 .out_features();
8137 let n_embd = self.cfg.n_embd as usize;
8138 // SESSION MODE: reuse the live cache/scratch, prime only the suffix. `base` = tokens
8139 // already committed (their state is in the caches); 0 = fresh single-shot call.
8140 let session_mode = sess.is_some();
8141 let max_ctx = match sess.as_ref() {
8142 Some(s) => s.cache.max_ctx,
8143 None => prompt.len() + max_new + k + 8,
8144 };
8145 let mut own_cache;
8146 let mut own_scratch;
8147 // PREFIX-CACHE capture request threaded out of the session (lane/spec-prefix-cache):
8148 // (requested split, destination slot). Single-shot per burst; fresh calls have none.
8149 let mut sess_capture: Option<(Option<usize>, &mut Option<SpecBoundaryCapture>)> = None;
8150 let (
8151 cache,
8152 scratch,
8153 mut sess_tail,
8154 mut sess_draft_slot,
8155 mut sess_pending_slot,
8156 sess_ckpt_slot,
8157 sess_telem,
8158 ): (
8159 &mut Cache,
8160 &mut MtpScratch,
8161 Option<(
8162 &mut Vec<u32>,
8163 &mut Option<CudaSlice<f32>>,
8164 &mut Option<u32>,
8165 &mut u32,
8166 &mut u32,
8167 )>,
8168 Option<&mut Option<DraftGraphCtx>>,
8169 Option<&mut Option<u32>>,
8170 Option<&mut Option<SpecCheckpoint>>,
8171 Option<&SpecTelemetryCounters>,
8172 ) = match sess.take() {
8173 Some(sr) => {
8174 let SpecSession {
8175 cache,
8176 scratch,
8177 committed,
8178 last_h,
8179 next_pred,
8180 sctr: s_sctr,
8181 uctr: s_uctr,
8182 draft_ctx,
8183 pending_tok,
8184 turn_ckpt,
8185 telem,
8186 capture_at,
8187 boundary_capture,
8188 } = sr;
8189 sess_capture = Some((capture_at.take(), boundary_capture));
8190 (
8191 cache,
8192 scratch,
8193 Some((committed, last_h, next_pred, s_sctr, s_uctr)),
8194 Some(draft_ctx),
8195 Some(pending_tok),
8196 Some(turn_ckpt),
8197 Some(telem),
8198 )
8199 }
8200 None => {
8201 // STAGE-OWNED KV (lane/pp2-spec 2026-08-06) — see `new_session`. Door shut =
8202 // `Cache::new` verbatim.
8203 own_cache = crate::pp::new_cache(e, &self.cfg, max_ctx)?;
8204 // Persistent scratch = max_ctx rows (~2KB/token quantized).
8205 own_scratch = MtpScratch::new(
8206 e,
8207 &self.cfg,
8208 max_ctx,
8209 self.mtp.as_ref().and_then(|m| m.geom.as_ref()),
8210 )?;
8211 (
8212 &mut own_cache,
8213 &mut own_scratch,
8214 None,
8215 None,
8216 None,
8217 None,
8218 None,
8219 )
8220 }
8221 };
8222 let base = cache.pos;
8223 // PENDING-CARRY consume (2026-08-01): a carried bonus reaches here only on the
8224 // empty-suffix GREEDY continuation path (generate_spec_session_sampled flushed every
8225 // other case). It enters the round loop as round-0's pending — verify col 0 — exactly
8226 // like a mid-burst full-accept boundary: no init feed, no tail commit pass.
8227 let carried_pending: Option<u32> = sess_pending_slot.as_mut().and_then(|s| s.take());
8228 // PERSISTENT DRAFT KV (the only mode since 2026-07-08 — the legacy round-local scratch,
8229 // MEMRA_SPEC_KVLOCAL, measured -35 acceptance pts on the 27B p3 sweep and was removed;
8230 // acceptance-only — exactness is verify's job either way).
8231 // HIDDEN-PAIRING CONVENTION (DEFAULT = predecessor-row, 2026-07-04 — the 27B acceptance
8232 // unlock, +16pts): the MTP head is TRAINED on rows pairing token x_p with the trunk
8233 // hidden of its PREDECESSOR h_{p-1} (the reference engine's mtp_update shifts the target
8234 // hiddens right by one; its draft step 0 feeds (id_last, TRUE hidden of the row id_last
8235 // was sampled from)). memra's historical convention paired SAME-ROW (x_p, h_p) in the fill
8236 // and seeded chain step 0 through an extra MTP pass on a duplicated token (the
8237 // pseudo-seed) — measured 27B p2 K=3 acceptance 0.569 vs 0.731, p3 0.445 vs 0.63+, and
8238 // the chain steps j>=1 were already predecessor-shaped, so ONLY the fill + step-0 seed
8239 // move. The fill shifts by one and the chain seeds from the predecessor's true hidden
8240 // DIRECTLY (vh_seed / vx[j-1]) — the pseudo pass disappears (one MTP-block pass saved
8241 // per round on top of the acceptance win). Draft-quality-only: exactness stays the
8242 // verify's job either way. (The legacy same-row pairing seam, MEMRA_SPEC_HSAME, and its
8243 // pseudo-seed passes were removed 2026-07-08 — predecessor pairing won by +16 acc pts;
8244 // the legacy round-local scratch, MEMRA_SPEC_KVLOCAL, went with it.)
8245 // REPLAY-FREE PARTIAL ACCEPT (default, 2026-07-03): partial rounds keep the verify's own
8246 // bit-identical committed-prefix state (KV truncate + recur rebuild from the VerifyCkpt)
8247 // and leave the bonus PENDING — no duplicate trunk pass (profiled ~0.54 extra full weight
8248 // reads/round at long ctx). MEMRA_SPEC_REPLAY=1 restores the legacy rollback+replay (A/B
8249 // + fallback seam).
8250 // Qwen35-MoE replay pin LIFTED (lane/draftcost-moe, 2026-08-20). The pin's stated
8251 // bar — the retained verify-state commit proven equivalent to sequential serving —
8252 // was waiting on this arch running the serving batched verify class, which the
8253 // t-parallel admission (this lane, increment 1) provided: the VerifyCkpt the
8254 // replay-free commit consumes is now produced by the SAME serving-class verify that
8255 // qualified dense qwen35 on 2026-08-15 (where the per-round duplicate replay
8256 // measured 69 -> 30 tok/s). Qualification receipts (run-spec K=1..8 both arms,
8257 // 8-prompt replay-vs-replay-free canary, long-prompt cell):
8258 // research/draftcost-moe-20260820/RECEIPTS.md. MEMRA_SPEC_REPLAY=1 stays the
8259 // rollback + A/B seam.
8260 let spec_replay = spec_replay_env_enabled();
8261 if constraint.is_some() && spec_replay {
8262 return Err(
8263 "constrained spec decode does not support MEMRA_SPEC_REPLAY=1 \
8264 (legacy replay commits an unmasked bonus)"
8265 .into(),
8266 );
8267 }
8268 // TRUE-HIDDEN REFRESH (default in persistent-draft-KV mode): every round overwrites the
8269 // committed positions' scratch entries from the verify's exact hiddens (mtp_kv_fill batch)
8270 // instead of keeping chain-approximate entries. MEMRA_SPEC_NOREFRESH=1 = legacy (A/B seam).
8271 let refresh = std::env::var("MEMRA_SPEC_NOREFRESH").is_err();
8272
8273 // prime: BATCHED cache prime (prime_cache — the measured #1 e2e gap: tokenwise primed at
8274 // ~102/38 tok/s vs the engine's ~2000-5900 tok/s batched prefill). prime_cache returns the
8275 // full pre-output_norm hidden stack [T, n_embd], which IS prompt_h (the persistent-draft-KV
8276 // mtp_kv_fill input) — no per-token collection needed. Prompts below PRIME_MIN_T, and
8277 // MEMRA_PRIME_TOKENWISE=1, and frozen Hy3 CPU/GPU expert splits take the tokenwise
8278 // decode_step_h loop. The latter avoids transient GPU staging of the spilled expert bank.
8279 // EMPTY-SUFFIX CONTINUATION (serve bursts): a session turn with NO new tokens resumes
8280 // generation exactly where the last turn stopped — no prime at all. The stashed
8281 // `next_pred` plays prime_logits' role: it is the token produced from the logits after
8282 // committed.last() by the same rule this entry applies to a cold prime's last row —
8283 // an argmax when greedy, a `sample_boundary_token` draw when sampled (the burst tail,
8284 // or `spec_session_from_restored` for a converted prefix-cache hit, did the drawing
8285 // where the sampler and the session's Philox counters were live). `last_h` seeds the
8286 // predecessor pairing below. Fresh calls and non-empty suffixes take the normal path.
8287 let continuation = prompt.is_empty();
8288 if continuation {
8289 assert!(session_mode, "empty prompt requires a session");
8290 assert!(
8291 sess_tail
8292 .as_ref()
8293 .map_or(false, |(c, lh, np, _, _)| !c.is_empty()
8294 && lh.is_some()
8295 && (np.is_some() || carried_pending.is_some())),
8296 "empty-suffix continuation needs a primed session (committed + last_h + next_pred|pending)"
8297 );
8298 }
8299 let mut prime_logits;
8300 let mut prompt_h: Option<CudaSlice<f32>> = None;
8301 let t_prime = std::time::Instant::now();
8302 let batched_prime = !continuation
8303 && prompt.len() >= crate::hybrid_forward::PRIME_MIN_T
8304 && std::env::var("MEMRA_PRIME_TOKENWISE").is_err()
8305 && !e.frozen_cpu_experts_prefer_tokenwise_prime();
8306 let prime_split = prime_split.filter(|&split| split > 0 && split < prompt.len());
8307 if prime_split.is_some() && (continuation || base != 0) {
8308 return Err("spec prime split is cold-session-only".into());
8309 }
8310 if continuation {
8311 prime_logits = Vec::new();
8312 } else if let Some(split) = prime_split {
8313 if split < crate::hybrid_forward::PRIME_MIN_T {
8314 return Err(format!(
8315 "spec prime split {split} is below PRIME_MIN_T {}",
8316 crate::hybrid_forward::PRIME_MIN_T,
8317 )
8318 .into());
8319 }
8320 // Mirror the plain worker's affinity boundary exactly. The prefix is a request-level
8321 // prime (`queued_after` keeps Step35 arm selection independent of this stop); a tail
8322 // below PRIME_MIN_T then takes the same eager tokenwise continuation as prefill_tick.
8323 // Retain every hidden row so the draft scratch fill remains one coherent prompt.
8324 let mut h_all = e.uninit(prompt.len() * n_embd)?;
8325 let (l, _, h_prefix) =
8326 self.prime_cache(e, &prompt[..split], &mut *cache, prompt.len() - split)?;
8327 e.copy_into(&mut h_all, 0, &h_prefix, split * n_embd)?;
8328 prime_logits = l;
8329 // PREFIX-CACHE BOUNDARY CAPTURE (lane/spec-prefix-cache): the GDN conv/ssm states
8330 // are about to be advanced in place by the tail prime, so this is the ONLY moment
8331 // the boundary's recurrent state exists. Capture iff the worker requested exactly
8332 // this split. cache.pos == split here (the prefix prime just finished). A failed
8333 // snapshot is silent (turn_ckpt convention) — publication is an optimization,
8334 // never a correctness dependency.
8335 if let Some((requested, slot)) = sess_capture.as_mut() {
8336 if *requested == Some(split) {
8337 debug_assert_eq!(cache.pos, split, "boundary capture off the prime split");
8338 if let Ok(snap) = cache.snapshot(e) {
8339 **slot = Some(SpecBoundaryCapture {
8340 snap,
8341 pos: split,
8342 logits: prime_logits.clone(),
8343 // rows [0..split) of h_all are the prefix prime's hiddens — copied
8344 // just above, before the tail prime overwrites nothing (append-only).
8345 last_h: capture_boundary_hidden(e, &h_all, split, n_embd),
8346 });
8347 }
8348 }
8349 }
8350 let tail = &prompt[split..];
8351 if tail.len() >= crate::hybrid_forward::PRIME_MIN_T
8352 && std::env::var("MEMRA_PRIME_TOKENWISE").is_err()
8353 && !e.frozen_cpu_experts_prefer_tokenwise_prime()
8354 {
8355 let (l, _, h_tail) = self.prime_cache(e, tail, &mut *cache, 0)?;
8356 e.copy_into(&mut h_all, split * n_embd, &h_tail, tail.len() * n_embd)?;
8357 prime_logits = l;
8358 } else {
8359 for (i, &tok) in tail.iter().enumerate() {
8360 let (l, h) = self.decode_step_h(e, tok, &mut *cache)?;
8361 e.copy_into(&mut h_all, (split + i) * n_embd, &h, n_embd)?;
8362 prime_logits = l;
8363 }
8364 }
8365 if std::env::var("MEMRA_SPEC_STATS").as_deref() == Ok("1") {
8366 eprintln!("[spec-prime] affinity split={split} tail={}", tail.len());
8367 }
8368 prompt_h = Some(h_all);
8369 } else if batched_prime {
8370 let (l, _h_seed, hiddens) = self.prime_cache(e, prompt, &mut *cache, 0)?;
8371 prime_logits = l;
8372 prompt_h = Some(hiddens);
8373 } else {
8374 prime_logits = Vec::new();
8375 prompt_h = Some(e.uninit(prompt.len() * n_embd)?);
8376 for (i, &tok) in prompt.iter().enumerate() {
8377 let (l, h) = self.spec_target_step_h(e, tok, &mut *cache)?;
8378 if let Some(ph) = prompt_h.as_mut() {
8379 e.copy_into(ph, i * n_embd, &h, n_embd)?;
8380 }
8381 prime_logits = l;
8382 }
8383 }
8384 e.stream().synchronize()?;
8385 // PREFIX-CACHE SEED CAPTURE (lane/spec-prefix-cache): boundary == prompt end (the seed
8386 // case — no shared-prefix split, publish the whole prompt). The prime just finished, so
8387 // cache.pos == base + prompt.len() and the recurrent state IS the boundary state;
8388 // prime_logits are the boundary logits. Cold sessions only (base == 0) — same law as
8389 // prime_split. The mid-prompt capture above already consumed the request if it matched.
8390 if !continuation && base == 0 {
8391 if let Some((requested, slot)) = sess_capture.as_mut() {
8392 if *requested == Some(prompt.len()) && slot.is_none() {
8393 debug_assert_eq!(cache.pos, prompt.len(), "seed capture off prompt end");
8394 if let Ok(snap) = cache.snapshot(e) {
8395 **slot = Some(SpecBoundaryCapture {
8396 snap,
8397 pos: prompt.len(),
8398 logits: prime_logits.clone(),
8399 last_h: prompt_h
8400 .as_ref()
8401 .map(|ph| capture_boundary_hidden(e, ph, prompt.len(), n_embd))
8402 .unwrap_or_default(),
8403 });
8404 }
8405 }
8406 }
8407 }
8408 // Harness timing contract (see crate::PRIME_NANOS): gen-only throughput without the
8409 // prime-subtraction hack.
8410 crate::PRIME_NANOS.store(
8411 t_prime.elapsed().as_nanos() as u64,
8412 std::sync::atomic::Ordering::Relaxed,
8413 );
8414
8415 let (embd_qt, embd_rb) = self.embd.qt_and_row_bytes(n_embd);
8416 // Resident table is fastest when it fits. Large spill deployments can preserve that HBM
8417 // for expert-cache slots and gather only the exact rows needed by MTP/verify from host.
8418 let host_embd = spec_host_embd();
8419 let embd_gpu = if host_embd {
8420 None
8421 } else {
8422 Some(
8423 self.embd_gpu
8424 .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload")),
8425 )
8426 };
8427 let embd_dev = embd_gpu.map(|g| (g, embd_qt, embd_rb));
8428 if host_embd {
8429 eprintln!(
8430 "[spec] host-row embedding: {} bytes kept off HBM",
8431 self.embd.raw.len()
8432 );
8433 }
8434 let mut out: Vec<u32> = Vec::with_capacity(max_new);
8435 let mut total_drafted = 0usize;
8436 let mut total_accepted = 0usize;
8437
8438 // --- SAMPLER FIRST (lane/sampled-spec-quality, 2026-08-19) ---
8439 // The sampler config, the session's Philox counters and the penalty window are parsed
8440 // HERE, above the boundary-token selection, because the boundary token must be drawn
8441 // from the sampler the request asked for. Pre-lane this block sat ~50 lines BELOW the
8442 // selection, which is the whole mechanical reason the boundary token was an argmax:
8443 // the sampler state was not in scope yet. Nothing here depends on the round loop, so
8444 // moving it up is a pure reordering for greedy (`sampled == false` ⇒ every branch
8445 // below takes the argmax path it always took).
8446 // --- SAMPLED SPEC (MEMRA_SPEC_TEMP>0, research/sampled-spec-impl-map.md): rejection-
8447 // sampling verify (Leviathan/Chen) — accept draft x at u < p(x)/q(x), resample from
8448 // norm(max(0,p-q)) on reject, bonus sampled from p on full accept. Counter-based Philox
8449 // everywhere (seed, event) -> reproducible. temp==0/unset = the greedy path, untouched.
8450 let sp = sampling.unwrap_or_else(|| SpecSampling {
8451 temp: std::env::var("MEMRA_SPEC_TEMP")
8452 .ok()
8453 .and_then(|v| v.parse().ok())
8454 .unwrap_or(0.0),
8455 seed: std::env::var("MEMRA_SEED")
8456 .ok()
8457 .and_then(|v| v.parse().ok())
8458 .unwrap_or(42),
8459 top_k: std::env::var("MEMRA_TOP_K")
8460 .ok()
8461 .and_then(|v| v.parse().ok())
8462 .unwrap_or(0),
8463 top_p: std::env::var("MEMRA_TOP_P")
8464 .ok()
8465 .and_then(|v| v.parse().ok())
8466 .unwrap_or(1.0),
8467 min_p: std::env::var("MEMRA_MIN_P")
8468 .ok()
8469 .and_then(|v| v.parse().ok())
8470 .unwrap_or(0.0),
8471 penalty_last_n: std::env::var("MEMRA_PENALTY_LAST_N")
8472 .ok()
8473 .and_then(|v| v.parse().ok())
8474 .unwrap_or(0),
8475 penalty_repeat: std::env::var("MEMRA_PENALTY_REPEAT")
8476 .ok()
8477 .and_then(|v| v.parse().ok())
8478 .unwrap_or(1.0),
8479 penalty_freq: std::env::var("MEMRA_PENALTY_FREQ")
8480 .ok()
8481 .and_then(|v| v.parse().ok())
8482 .unwrap_or(0.0),
8483 penalty_present: std::env::var("MEMRA_PENALTY_PRESENT")
8484 .ok()
8485 .and_then(|v| v.parse().ok())
8486 .unwrap_or(0.0),
8487 });
8488 let (sp_temp, sp_seed) = (sp.temp, sp.seed);
8489 let sampled = sp_temp > 0.0;
8490 // Counters resume from the session (burst continuity: randomness must never repeat
8491 // across generate_spec_session calls); one-shot callers start at (0,0). Read through
8492 // sess_tail — `sess` was take()n into it above, so sess.as_ref() here is always None.
8493 let mut sctr: u32 = sess_tail.as_ref().map(|(_, _, _, s, _)| **s).unwrap_or(0);
8494 let mut uctr: u32 = sess_tail.as_ref().map(|(_, _, _, _, u)| **u).unwrap_or(0);
8495 // Penalties (v2.1): applied to COPIES of q rows and p columns symmetrically (exactness
8496 // for the penalized+filtered target). History = generated tokens, host-tracked window.
8497 let pen_on = sampled
8498 && sp.penalty_last_n > 0
8499 && (sp.penalty_repeat != 1.0 || sp.penalty_freq != 0.0 || sp.penalty_present != 0.0);
8500 // SESSION-SPANNING PENALTY WINDOW (Item 2). Pre-lane this was
8501 // `prompt.iter().rev().take(64).rev()` — the BURST's suffix slice — so a continuation
8502 // burst (the majority of a stream's tokens, and ALL of a converted cache hit's) started
8503 // with an EMPTY penalty history and the client's repetition/frequency/presence penalties
8504 // silently reset at every burst boundary. The window now spans `committed ++ prompt`,
8505 // which is what the API contract says and what the plain sampler's own `history` does.
8506 // Byte-identical to the pre-lane seed for a cold turn-1 burst at the default window.
8507 let mut pen_hist: Vec<u32> = if pen_on {
8508 let sess_hist: &[u32] = if spec_pen_session_on() {
8509 sess_tail
8510 .as_ref()
8511 .map(|(c, ..)| c.as_slice())
8512 .unwrap_or(&[])
8513 } else {
8514 &[] // MEMRA_SPEC_PEN_SESSION=0: pre-lane burst-local window
8515 };
8516 pen_window_seed(sess_hist, prompt, sp.penalty_last_n)
8517 } else {
8518 Vec::new()
8519 };
8520 // First generated token = the BOUNDARY token: greedy takes the argmax of the prompt's
8521 // last logits (== greedy's first token, byte-contract); SAMPLED draws it from the
8522 // request's own filtered/penalized target through the session's Philox stream
8523 // (`sample_boundary_token`, lane/sampled-spec-quality Item 1 — pre-lane this was an
8524 // argmax in both regimes, so ~1 token per burst of a sampled stream was greedy).
8525 // Emit it, then FEED it to establish the loop invariant below.
8526 // PENDING-CARRY: the carried bonus was already emitted by the LAST burst — it becomes
8527 // last_token WITHOUT re-emission, and round 0 consumes it as pending (no init feed).
8528 // CONSTRAINED entry rules: the first emitted token is the MASKED argmax of the
8529 // prompt's last logits (plain constrained-greedy identity); a continuation without
8530 // a carried pending would emit an UNMASKED stashed next_pred — refused loudly (the
8531 // worker never resumes constrained sessions from the pool, so this cannot fire).
8532 if let Some(c) = constraint.as_deref_mut() {
8533 if continuation && carried_pending.is_none() {
8534 return Err("constrained spec continuation requires a carried pending \
8535 (pool resume is unconstrained-only)"
8536 .into());
8537 }
8538 if !continuation {
8539 c.mask_logits(&mut prime_logits)
8540 .map_err(|e2| format!("constraint: {e2}"))?;
8541 }
8542 }
8543 let mut last_token = if let Some(b) = carried_pending {
8544 b
8545 } else if continuation {
8546 // A continuation's boundary token was DRAWN by the burst that stashed it (the
8547 // session tail below), or by `spec_session_from_restored` for a converted
8548 // prefix-cache hit — in both cases from the correct logits row with this same
8549 // session's Philox stream, which is why it can be consumed here as-is.
8550 sess_tail.as_ref().unwrap().2.unwrap()
8551 } else if sampled && constraint.is_none() && spec_sampled_boundary_on() {
8552 sample_boundary_token(e, &prime_logits, &sp, &pen_hist, &mut sctr, "cold-prime")?
8553 } else {
8554 // greedy (byte contract), the rollback door, or constrained (masked-argmax
8555 // identity — the worker routes sampled+constrained to the plain path, and this
8556 // function refuses the combination outright above).
8557 argmax(&prime_logits) as u32
8558 };
8559 if pen_on {
8560 // The boundary token is a GENERATED token: the plain sampler `accept()`s every
8561 // emitted token into its penalty history, and pre-lane the burst's first token
8562 // was invisible to penalties forever (never pushed, and never in `committed`
8563 // until this burst's tail). Covers the carry/continuation seeds too — neither is
8564 // in `committed` yet.
8565 pen_hist.push(last_token);
8566 }
8567 if carried_pending.is_none() {
8568 out.push(last_token);
8569 // grammar advances with every emitted token (carried pendings were consumed
8570 // by the burst that emitted them).
8571 if let Some(c) = constraint.as_deref_mut() {
8572 c.consume(last_token)
8573 .map_err(|e2| format!("constraint: {e2}"))?;
8574 }
8575 }
8576 if continuation {
8577 // draft-KV invariant: entries [0..base) are the session's exact fills; truncate any
8578 // overhang so the chain's first append lands at slot base (== committed.len()).
8579 scratch.set_len(e, base)?;
8580 }
8581 // sse-cadence: hand the caller every not-yet-flushed token (disjoint in-order slices
8582 // concatenating to the full `out`). Called after the prime's first token and after each
8583 // round commit — emission timing only, token bytes untouched. The slice may be EMPTY
8584 // (poll-only boundary: zero-round folds commit nothing new); returns the caller's
8585 // continue-verdict (admission yield, 2026-08-06) — false ends the burst at this round.
8586 fn flush_commit(
8587 cb: &mut Option<&mut dyn FnMut(&[u32]) -> bool>,
8588 out: &[u32],
8589 flushed: &mut usize,
8590 ) -> bool {
8591 if let Some(f) = cb.as_mut() {
8592 let keep = f(&out[*flushed..]);
8593 *flushed = out.len();
8594 keep
8595 } else {
8596 true
8597 }
8598 }
8599 keep_going = flush_commit(&mut on_commit, &out, &mut flushed);
8600 // INVARIANT at loop top: `last_token` is the most-recently-committed/emitted token, its
8601 // KV+recur state IS in `cache` (cache.pos = position right AFTER last_token), `last_pred`
8602 // is the greedy ARGMAX of the logits that predict the token FOLLOWING last_token, and
8603 // `h_seed` = last_token's pre-output_norm hidden. Establish it by feeding last_token once
8604 // (mirrors plain greedy). DEVICE-ARGMAX lever: the accept walk only ever consumes the
8605 // argmax of those logits — never the full vector — so a host u32 replaces the Vec<f32>.
8606 // Trimmed heads: q lives on the trimmed vocab; accept gathers use the TRIMMED index and
8607 // the residual scatters q into target-id space (q=-inf off-trim — the head cannot propose
8608 // those, so their residual mass is p(x), correct by construction).
8609 let d2t_dev: Option<CudaSlice<u32>> = if sampled || crate::spec::spec_stream() {
8610 match &mtp.d2t {
8611 Some(map) => Some(e.htod_u32_v(map)?),
8612 None => None,
8613 }
8614 } else {
8615 None
8616 };
8617 let mut q_full_buf: Option<CudaSlice<f32>> = None;
8618 // host Philox4x32-10 (mirrors spec_sample.cu; independent stream via ctr_lo tag)
8619 let host_u01 = |seed: u64, ctr: u32| -> f32 {
8620 let (m0, m1) = (0xD2511F53u32, 0xCD9E8D57u32);
8621 let (mut c0, mut c1, mut c2, mut c3) = (0xFFFF_FFFEu32, ctr, 0u32, 0u32);
8622 let (mut k0, mut k1) = ((seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
8623 for _ in 0..10 {
8624 let (h0, l0) = (((m0 as u64 * c0 as u64) >> 32) as u32, m0.wrapping_mul(c0));
8625 let (h1, l1) = (((m1 as u64 * c2 as u64) >> 32) as u32, m1.wrapping_mul(c2));
8626 let (n0, n1, n2, n3) = (h1 ^ c1 ^ k0, l1, h0 ^ c3 ^ k1, l0);
8627 c0 = n0;
8628 c1 = n1;
8629 c2 = n2;
8630 c3 = n3;
8631 k0 = k0.wrapping_add(0x9E3779B9);
8632 k1 = k1.wrapping_add(0xBB67AE85);
8633 }
8634 (c0 as f32 + 1.0) * (1.0 / 4294967296.0)
8635 };
8636 let mut draft_logits: Vec<CudaSlice<f32>> = Vec::new(); // retained head logits (q), per slot
8637 let mut draft_stats: Vec<(f32, f32, f32)> = Vec::new(); // (row_max, th_e, z_e) per slot
8638 let mut perturb_buf: Option<CudaSlice<f32>> = None; // gumbel scratch (max(n_vocab,d_vocab))
8639 let mut sample_tok = e.alloc_u32_zeroed(1)?; // residual/bonus sample out
8640 let mut col_buf: Option<CudaSlice<f32>> = None; // materialized verify column
8641 let mut pen_hist_d: Option<CudaSlice<u32>> = None;
8642 let mut pcol_buf: Option<CudaSlice<f32>> = None; // penalized p-column scratch
8643 // MEMRA_SPEC_SETUP_TRACE=1 (diagnostics): per-call wall decomposition of the burst
8644 // SETUP + TAIL segments (the round loop's internals are MEMRA_SPEC_PHASE's job) —
8645 // built to pin the serve per-burst fixed cost (research/spec-serving-20260801).
8646 let setup_trace = std::env::var("MEMRA_SPEC_SETUP_TRACE").as_deref() == Ok("1");
8647 let t_ent = std::time::Instant::now();
8648
8649 // SESSION-AFFINITY TURN CHECKPOINT (lane/session-affinity, 2026-08-05): capture the
8650 // PROMPT-END boundary state so a LATER turn can rewind here and re-prime only its own
8651 // delta instead of the whole conversation. See `SpecCheckpoint` for why this boundary is
8652 // the one that matters (a history-rewriting client mutates what the session GENERATED,
8653 // so the next turn's prompt agrees with this one up to exactly here).
8654 //
8655 // WHERE — AND WHY THIS EXACT LINE. Right after the trunk prime, BEFORE the init feed
8656 // (`decode_step_h(last_token)`) and before round 0: the last instant at which the caches
8657 // hold exactly `base + prompt.len()` rows and nothing generated.
8658 //
8659 // This was WRONG in the first cut of this lane: the capture sat after the draft-KV fill,
8660 // which is also after the init feed, so `cache.pos` was `base + prompt.len() + 1` — the
8661 // boundary included the FIRST GENERATED TOKEN. That token is the first thing inside the
8662 // `<think>` block the client strips, so every later turn's diff diverged exactly one
8663 // token below the checkpoint and affinity declined 100% of the time. Measured on the
8664 // owner regime: "history diverged at 12233 of checkpoint 12234". The off-by-one made the
8665 // whole mechanism inert while looking, from the outside, like a working
8666 // correctness-declines-safely path — hence the decline log carries the offsets.
8667 //
8668 // The full-attn planes are `len`-truncatable so the snapshot copies only the GDN conv/ssm
8669 // state (the reason a spec session could not rewind before). The draft scratch needs no
8670 // copy: rows below the boundary are rewritten by the next turn's own fill.
8671 //
8672 // WHEN: non-empty prime only. An empty-suffix continuation burst adds no prompt boundary
8673 // (its "prompt end" IS the previous checkpoint's, already held), so it keeps the existing
8674 // checkpoint rather than replacing it with a strictly worse one.
8675 //
8676 // FAILURE IS SILENT BY DESIGN: on a VRAM-tight rig the snapshot alloc can fail. That
8677 // costs the NEXT turn its rewind (it re-primes fully, today's behavior) and must never
8678 // fail the burst that is already running — so the error is swallowed, loud only under
8679 // MEMRA_DEBUG_SPEC.
8680 if let Some(slot) = sess_ckpt_slot {
8681 if !continuation {
8682 let pos = cache.pos;
8683 debug_assert_eq!(
8684 pos,
8685 base + prompt.len(),
8686 "turn checkpoint must sit at the prompt end, before the init feed"
8687 );
8688 let anchor: Result<CudaSlice<f32>, Box<dyn std::error::Error>> =
8689 if let Some(ph) = &prompt_h {
8690 // hidden of the LAST primed row = the predecessor anchor at this
8691 // boundary (exactly what a fresh prime of committed[..pos] leaves in
8692 // last_h, and what the next prime's fill reads for its first row).
8693 let np = prompt.len();
8694 e.uninit(n_embd).and_then(|mut a| {
8695 e.copy_view_into(
8696 &mut a,
8697 0,
8698 &ph.slice((np - 1) * n_embd..np * n_embd),
8699 n_embd,
8700 )?;
8701 Ok(a)
8702 })
8703 } else {
8704 Err("no prompt hiddens".into())
8705 };
8706 match (cache.snapshot(e), anchor) {
8707 (Ok(snap), Ok(last_h)) => {
8708 *slot = Some(SpecCheckpoint { snap, pos, last_h });
8709 }
8710 (s, a) => {
8711 *slot = None; // a stale checkpoint would rewind to the WRONG boundary
8712 if std::env::var("MEMRA_DEBUG_SPEC").is_ok() {
8713 let err = s
8714 .err()
8715 .map(|e| e.to_string())
8716 .or_else(|| a.err().map(|e| e.to_string()))
8717 .unwrap_or_default();
8718 eprintln!(
8719 "[spec] turn checkpoint skipped ({err}); \
8720 next turn re-primes in full"
8721 );
8722 }
8723 }
8724 }
8725 }
8726 }
8727 // INIT FEED — skipped on a pending carry: last_token (the carried bonus) is NOT in the
8728 // caches and must NOT be fed solo; round 0's batched verify commits it as col 0. Its
8729 // seed/anchor hidden is the carried last_h (copied below); last_pred is dead in the
8730 // pending path (t_pred reads verify col 0 — the accept walk overwrites it).
8731 let mut last_pred = 0u32;
8732 let mut last_col_logits: Option<CudaSlice<f32>> = None;
8733 // CONSTRAINED: the init feed's logits back the (n_acc==0, base==0) masked-argmax
8734 // recompute in the grammar-truncation walk — retained host-side, round 0 only.
8735 let mut init_logits_host: Option<Vec<f32>> = None;
8736 let h_seed0: CudaSlice<f32> = if carried_pending.is_none() {
8737 let (init_logits, h) = self.spec_target_step_h(e, last_token, &mut *cache)?;
8738 last_pred = argmax(&init_logits) as u32;
8739 if constraint.is_some() {
8740 init_logits_host = Some(init_logits.clone());
8741 }
8742 // sampled mode: p-distribution after last_token, for the j==0/base==0 accept test.
8743 if sampled {
8744 last_col_logits = Some(e.htod(&init_logits)?);
8745 }
8746 h
8747 } else {
8748 // predecessor-row anchor: hidden of the last COMMITTED row (the carry contract).
8749 let lh = sess_tail
8750 .as_ref()
8751 .unwrap()
8752 .1
8753 .as_ref()
8754 .expect("pending carry requires last_h");
8755 e.clone_dtod(lh)?
8756 };
8757 let t_init = t_ent.elapsed();
8758 let mut last_col_stats: Option<(f32, f32, f32)> = None;
8759 // PERSISTENT h_seed buffer (allocated BEFORE any graph capture so no captured scratch can
8760 // alias it): every path that updates the round seed copies INTO it — no per-round allocs,
8761 // stable pointer for the graph-draft round-start copy.
8762 let mut h_seed_buf = e.clone_dtod(&h_seed0)?;
8763 // Predecessor-pairing trackers: `fill_prev` = trunk hidden AT the last COMMITTED row (the
8764 // predecessor of the next verify's col 0 — the reference's carried pending-h analogue;
8765 // also the predecessor-row hidden for the round-0 legacy-replay seed). At round 0 that
8766 // row is last_token's own (h_seed0). The chain step-0 seed under the pairing default =
8767 // hidden of the row BEFORE last_token = the prompt's last row at round 0 (h_seed_buf
8768 // overwritten below).
8769 let mut fill_prev = e.clone_dtod(&h_seed0)?;
8770 {
8771 if let Some(ph) = &prompt_h {
8772 let np = prompt.len();
8773 e.copy_view_into(
8774 &mut h_seed_buf,
8775 0,
8776 &ph.slice((np - 1) * n_embd..np * n_embd),
8777 n_embd,
8778 )?;
8779 } else if continuation {
8780 if let Some((_, lh, _, _, _)) = sess_tail.as_ref() {
8781 if let Some(lh) = lh.as_ref() {
8782 e.copy_into(&mut h_seed_buf, 0, lh, n_embd)?;
8783 }
8784 }
8785 }
8786 }
8787 // Persistent device prediction slots for the accept walk (max k+1 verify columns).
8788 let mut preds_d = e.alloc_u32_zeroed(k + 2)?;
8789
8790 let debug_spec = std::env::var("MEMRA_DEBUG_SPEC").is_ok();
8791 let fork_mode = OptiForkGateMode::configured();
8792 // MEMRA_SPEC_STATS=1: per-slot accept histogram + draft-length histogram, printed once at
8793 // the end. Metric normalization vs the reference engine: BOTH engines count
8794 // accepted/drafted where the chain stopped at p-min and the sub-threshold token is
8795 // discarded uncounted — per-slot decay + chain-length mix are the extra dimensions.
8796 let spec_stats = std::env::var("MEMRA_SPEC_STATS").is_ok();
8797 let mut st_drafted = vec![0usize; k];
8798 let mut st_accepted = vec![0usize; k];
8799 let mut st_len_hist = vec![0usize; k + 1];
8800 let mut st_full = 0usize;
8801 // P-MIN CONFIDENCE GATE (MEMRA_SPEC_PMIN, the serve script's --spec-draft-p-min mechanism):
8802 // stop the draft chain early when the head's softmax confidence in its own pick drops
8803 // below p_min. Hoisted above the loop: the graph capture bakes the prob kernels iff on.
8804 static PMIN: std::sync::OnceLock<f32> = std::sync::OnceLock::new();
8805 let p_min = *PMIN.get_or_init(|| {
8806 std::env::var("MEMRA_SPEC_PMIN")
8807 .ok()
8808 .and_then(|v| v.parse().ok())
8809 .unwrap_or(0.0)
8810 });
8811 // ZERO-DRAFT ROUNDS (MEMRA_SPEC_PMIN0=1, vendored from llama.cpp's draft gating): let the
8812 // p-min gate apply at j==0 too, so a low-confidence round drafts NOTHING and the verify
8813 // batch is just the pending bonus (m=1 = a plain decode step). llama's 35B win rides
8814 // exactly this — draft acceptance 76% at mean len 2.5 because unpredictable stretches
8815 // never pay draft+verify overhead. Only legal when a pending bonus exists (an empty
8816 // verify batch is not); the j==0 exemption stays for pending-less rounds.
8817 let pmin0 = std::env::var("MEMRA_SPEC_PMIN0")
8818 .map(|v| v == "1")
8819 .unwrap_or(false);
8820
8821 // --- GRAPH DRAFT setup: persistent I/O buffers + ONE capture (2 warmups inside). The
8822 // warmups mutate scratch len_d / pos / tok / seed — all reset at every round start, so the
8823 // only restore needed is the scratch counter. Capture failure (e.g. a non-capturable
8824 // cuBLAS path in an exotic head) falls back to the eager draft chain.
8825 // PER-SESSION PERSISTENCE (2026-08-01): session calls reuse the DraftGraphCtx parked on
8826 // the SpecSession — the capture (2 warmup head forwards + instantiate) ran ONCE at the
8827 // session's first burst, not per burst (measured ~16ms/burst fixed cost on H100 q27,
8828 // research/spec-serving-20260801). Reuse is pointer-exact: the graph bakes the session's
8829 // own scratch KV (never realloc'd), the model's resident embedding, the OnceLock p_min,
8830 // and the g_* buffers carried in the ctx — replay dispatch is identical to a fresh
8831 // capture, so draft tokens are bit-identical (drafts never decide exactness anyway; the
8832 // verify arbitrates). Single-shot calls (sess=None) build a fresh ctx and drop it.
8833 let mut dctx: DraftGraphCtx = match sess_draft_slot.as_mut().and_then(|s| s.take()) {
8834 Some(c) => c,
8835 None => DraftGraphCtx::new(e, n_embd, if sampled { d_vocab } else { 1 })?,
8836 };
8837 // A session that ran greedy bursts first sized g_q/g_perturb at 1; a sampled resume
8838 // needs d_vocab. Realloc is legal exactly while graph_s is None (nothing baked them).
8839 if sampled && dctx.g_q.len() < d_vocab {
8840 dctx.g_q = e.zeros(d_vocab)?;
8841 dctx.g_perturb = e.zeros(d_vocab)?;
8842 }
8843 // DRAFT-SIDE GRAMMAR MASK (lane/draft-mask, 2026-08-04): the drafter samples the
8844 // grammar's legal set, so proposals are legal BY CONSTRUCTION and the verify-side
8845 // truncation (the correctness backstop) stops cutting every tight-schema round.
8846 // The mask is one node inside the captured draft chain — presence is a CAPTURE-TIME
8847 // shape, so a parked graph of the other shape is dropped and recaptured.
8848 let dmask_on = constraint
8849 .as_deref()
8850 .is_some_and(|c| c.draft_mask_enabled());
8851 let dmask_words = if dmask_on { d_vocab.div_ceil(32) } else { 0 };
8852 if dmask_on && dctx.g_dmask.len() < dmask_words {
8853 dctx.g_dmask = e.alloc_u32_zeroed(dmask_words)?;
8854 dctx.graph = None; // the old capture baked the old (or no) mask pointer
8855 dctx.failed.clear_greedy();
8856 dctx.keeper.clear();
8857 }
8858 if dctx.graph.is_some() && dctx.graph_masked != dmask_on {
8859 dctx.graph = None;
8860 dctx.failed.clear_greedy();
8861 dctx.keeper.clear();
8862 }
8863 if graph_draft && !sampled && dctx.graph.is_none() && !dctx.failed.greedy_failed() {
8864 let DraftGraphCtx {
8865 g_tok,
8866 g_pos,
8867 g_seed,
8868 g_p,
8869 g_dmask,
8870 ..
8871 } = &mut dctx;
8872 // capture-time contents: ALL-ONES (ban nothing). A replay only ever runs after the
8873 // host uploads the position's real words, so the warmups stay grammar-free.
8874 if dmask_on {
8875 e.htod_u32_into(g_dmask, &vec![u32::MAX; dmask_words])?;
8876 }
8877 let g_dmask_ro: &CudaSlice<u32> = &*g_dmask;
8878 // CAPTURE-RETAIN (#68 fix): the warmup transients' pool addresses are baked into the
8879 // captured graph; the keeper pins them for the graph's lifetime. capture_graph (non-
8880 // retained) freed them at exit — safe for one-shot generate_spec (nothing else touches
8881 // the pool between replays) but WRONG for sessions: burst-boundary prime/fill/commit
8882 // passes (and, in serve, other sessions) recycle those addresses and the replay then
8883 // clobbers live buffers — the ST serve-spec corruption (research/serve-st-20260803).
8884 let cap_res = e.capture_graph_retained(|e| {
8885 self.mtp_head_forward_cap(
8886 e,
8887 mtp,
8888 g_tok,
8889 g_pos,
8890 g_seed,
8891 g_p,
8892 &mut *scratch,
8893 p_min > 0.0 || fork_mode == OptiForkGateMode::Controller,
8894 true,
8895 embd_gpu.expect("graph draft requires resident embedding"),
8896 embd_qt,
8897 embd_rb,
8898 d_vocab,
8899 None,
8900 None,
8901 if dmask_on {
8902 Some((g_dmask_ro, dmask_words))
8903 } else {
8904 None
8905 },
8906 )
8907 });
8908 match cap_res {
8909 Ok((g, keep)) => {
8910 scratch.set_len(e, base)?;
8911 dctx.graph = Some(g);
8912 dctx.graph_masked = dmask_on;
8913 dctx.keeper = keep;
8914 }
8915 Err(err) => {
8916 scratch.set_len(e, base)?;
8917 // LOUD flip (audit Q2): a dropped draft graph is a coverage loss, never
8918 // silent. Once per flip — mark returns None on an already-failed ctx.
8919 if let Some(line) = dctx.failed.mark_greedy(&err.to_string()) {
8920 eprintln!("{line}");
8921 }
8922 }
8923 }
8924 }
8925 // --- SAMPLED GRAPH DRAFT setup (step 3 of the sampled-spec arc): a SECOND capture, own
8926 // graph object, built only when sampled && graph-eligible — the greedy capture above is
8927 // untouched (and skipped when sampled: its graph would never be launched). Same head
8928 // forward, but the in-graph argmax reads GUMBEL-PERTURBED logits; the Philox event
8929 // counter lives in the persistent device g_ctr (bumped in-graph, host-seeded from sctr
8930 // once per round); the raw head logits land in the persistent g_q for the host's
8931 // per-replay async D2D into the round's q slot (q_slots, K x d_vocab, allocated once).
8932 // seed/temp are capture-time constants — baked into graph_s, so a pool-resumed request
8933 // with a different (seed, temp, k) drops the parked sampled graph and recaptures.
8934 // COST OF THE FRESH-SEED SERVE DEFAULT (dogfood F4, 2026-08-04): omitting `seed` on a
8935 // serve request now draws fresh per-request entropy (it used to default to a pinned 0),
8936 // so a seed-omitting request that RESUMES a parked spec session finds an s_key baked
8937 // with the PREVIOUS request's seed and pays one recapture. Bounded, and it does not
8938 // reopen the ~16ms/burst regression the persistent ctx exists to fix: a session's seed
8939 // is fixed for its whole lifetime (worker.rs reads s.sampler.seed() per burst), so
8940 // this compare misses at most ONCE per resumed request — the first burst recaptures
8941 // and every later burst in that request replays. A client that wants the parked graph
8942 // AND reproducibility supplies an explicit `seed`, honored exactly, which keeps s_key
8943 // stable across its whole conversation.
8944 // COMPOSITION RULE (fspec x gsd merge): the in-graph chain samples from the RAW
8945 // softmax — it can hold neither per-row filter stats nor the varying penalty history.
8946 // The sampled graph therefore engages only in the PURE-TEMP regime; filters/penalties
8947 // force the eager draft (which computes stats/penalties per row).
8948 // KEY THE WHOLE REGIME, not just the baked constants (lane/graph-s-key-exactness-
8949 // 20260819). `s_key` used to be `(seed, temp, k)`; the filters and penalties were left
8950 // out, so a filtered request resuming a session that parked a PURE-TEMP graph kept it —
8951 // and the launch site never re-asked `pure_temp`. See [`SampledGraphKey`] for what that
8952 // costs (an unconditional accept of out-of-head draft tokens, i.e. an exactness bug on
8953 // the request shape the vendor-default flip makes the majority).
8954 let s_key = SampledGraphKey::new(sp_seed, sp_temp, k, sp.top_k, sp.top_p, sp.min_p, pen_on);
8955 let pure_temp = s_key.pure_temp();
8956 if sampled && dctx.s_key.is_some_and(|old| old != s_key) {
8957 dctx.graph_s = None;
8958 dctx.failed.clear_sampled();
8959 dctx.s_key = None;
8960 dctx.q_slots.clear();
8961 dctx.keeper_s.clear();
8962 }
8963 if graph_draft
8964 && sampled
8965 && pure_temp
8966 && dctx.graph_s.is_none()
8967 && !dctx.failed.sampled_failed()
8968 {
8969 let DraftGraphCtx {
8970 g_tok,
8971 g_pos,
8972 g_seed,
8973 g_p,
8974 g_ctr,
8975 g_perturb,
8976 g_q,
8977 ..
8978 } = &mut dctx;
8979 // CAPTURE-RETAIN (#68 fix): same keeper contract as the greedy capture above.
8980 let cap_res = e.capture_graph_retained(|e| {
8981 self.mtp_head_forward_cap(
8982 e,
8983 mtp,
8984 g_tok,
8985 g_pos,
8986 g_seed,
8987 g_p,
8988 &mut *scratch,
8989 p_min > 0.0,
8990 true,
8991 embd_gpu.expect("graph draft requires resident embedding"),
8992 embd_qt,
8993 embd_rb,
8994 d_vocab,
8995 Some((g_ctr, g_perturb, g_q, sp_seed, sp_temp)),
8996 None,
8997 None, // constrained spec is greedy-only — sampled never carries a hook
8998 )
8999 });
9000 match cap_res {
9001 Ok((g, keep)) => {
9002 scratch.set_len(e, base)?;
9003 for _ in 0..k {
9004 dctx.q_slots.push(e.zeros(d_vocab)?);
9005 }
9006 dctx.graph_s = Some(g);
9007 dctx.s_key = Some(s_key);
9008 dctx.keeper_s = keep;
9009 }
9010 Err(err) => {
9011 scratch.set_len(e, base)?;
9012 // LOUD flip (audit Q2): same contract as the greedy capture above.
9013 if let Some(line) = dctx.failed.mark_sampled(&err.to_string()) {
9014 eprintln!("{line}");
9015 }
9016 }
9017 }
9018 }
9019 // ---- EXACTNESS GUARD, the enforceable half (lane/graph-s-key-exactness-20260819) ----
9020 // With the filters and penalties in `s_key`, a graph that SURVIVED the drop above was
9021 // captured under this request's exact regime, and capture requires `pure_temp` — so a
9022 // parked graph implies `pure_temp`. That implication is the whole exactness argument for
9023 // the graph arm, so it is asserted here rather than assumed: a future change that widens
9024 // the capture condition, narrows the key, or copies a `DraftGraphCtx` across regimes
9025 // fails LOUDLY at this line instead of silently drafting from the raw softmax while the
9026 // verify applies filter stats. Release builds refuse the graph (drop it, draft eager)
9027 // rather than launching it; the launch site re-tests `pure_temp` independently.
9028 if sampled && !pure_temp && dctx.graph_s.is_some() {
9029 debug_assert!(
9030 false,
9031 "sampled draft graph parked under {:?} survived into a FILTERED request \
9032 (top_k={} top_p={} min_p={} pen_on={}): the in-graph chain draws from the RAW \
9033 softmax, so the verify's filtered q would test a distribution the draft was \
9034 never sampled from",
9035 dctx.s_key, sp.top_k, sp.top_p, sp.min_p, pen_on,
9036 );
9037 eprintln!(
9038 "[spec] BUG: dropping a parked sampled draft graph that outlived its capture \
9039 regime (s_key={:?}, request top_k={} top_p={} min_p={} pen_on={}); drafting \
9040 EAGER — the key must carry every field that shapes q",
9041 dctx.s_key, sp.top_k, sp.top_p, sp.min_p, pen_on,
9042 );
9043 dctx.graph_s = None;
9044 dctx.s_key = None;
9045 dctx.q_slots.clear();
9046 dctx.keeper_s.clear();
9047 }
9048 // SKEY PROBE (MEMRA_SKEY_PROBE=1): the burst-entry facts the reachability question turns
9049 // on — is this request sampled, is it in the pure-temp regime the sampled graph is only
9050 // legal in, and is a graph PARKED from an earlier request of the same session? The launch
9051 // arms below print which chain actually ran, so the probe never restates the condition.
9052 if skey_probe() {
9053 eprintln!(
9054 "[skey] burst sampled={} pure_temp={} temp={} top_k={} top_p={} min_p={} \
9055 pen_on={} k={} graph_draft={} graph_s_parked={} s_key_parked={:?}",
9056 sampled as u8,
9057 pure_temp as u8,
9058 sp_temp,
9059 sp.top_k,
9060 sp.top_p,
9061 sp.min_p,
9062 pen_on as u8,
9063 k,
9064 graph_draft as u8,
9065 dctx.graph_s.is_some() as u8,
9066 dctx.s_key,
9067 );
9068 }
9069 let t_cap = t_ent.elapsed();
9070 // PERSISTENT DRAFT KV: fill the MTP block's K/V for every prompt position from the exact
9071 // trunk hiddens collected during prime — ONE batched K/V-only pass (overwrites any
9072 // capture-warmup garbage; capture left len at 0). last_token (the init feed) needs no
9073 // fill: the first chain step processes it and appends its entry at slot prompt.len().
9074 if let Some(ph) = &prompt_h {
9075 // SESSION: rows [0..base) are the previous turns' exact fills (refresh overwrote them
9076 // with true verify hiddens) — truncate any draft overhang, fill ONLY the suffix at
9077 // global positions [base..base+tp). Fresh call: base==0, identical to before.
9078 scratch.set_len(e, base)?;
9079 // CHUNKED FILL (long-ctx OOM fix, 2026-07-05): mtp_kv_fill's transients scale with its
9080 // T (concat = T*2*n_embd*4B — 1.5GB at 40k) and its concat loop is 2*T launches. The
9081 // fill is a pure sequential append, so chunking is exact: each chunk appends its rows
9082 // at pos0=base+start with the identical per-row math. Same knob as the trunk prime.
9083 let tp = prompt.len();
9084 let fill_chunk: usize = if crate::cache::swa_ring_on() {
9085 crate::hybrid_forward::prime_chunk_tokens(tp, self.layers.len())
9086 } else {
9087 // Preserve the flag-OFF schedule byte-for-byte, including the legacy zero value
9088 // meaning one monolithic fill.
9089 std::env::var("MEMRA_PRIME_CHUNK")
9090 .ok()
9091 .and_then(|v| v.parse().ok())
9092 .unwrap_or(4096)
9093 };
9094 let fill_chunk = if fill_chunk == 0 { tp } else { fill_chunk };
9095 let mut start = 0usize;
9096 while start < tp {
9097 let end = (start + fill_chunk).min(tp);
9098 let tc = end - start;
9099 {
9100 // PREDECESSOR pairing: row i gets h[i-1]; global row 0 a zeros row (the
9101 // reference engine's initial pending-h is zeroed too); a session turn's row 0
9102 // gets the PREVIOUS turn's last committed hidden (sess.last_h). Per chunk:
9103 // rows start..end read h[start-1..end-1] — one dtod into a chunk buffer.
9104 let mut phs = e.zeros(tc * n_embd)?;
9105 let (src_lo, dst_off) = if start == 0 {
9106 (0, n_embd)
9107 } else {
9108 ((start - 1) * n_embd, 0)
9109 };
9110 let n_copy = if start == 0 {
9111 (tc - 1) * n_embd
9112 } else {
9113 tc * n_embd
9114 };
9115 if start == 0 {
9116 if let Some((_, lh, _, _, _)) = sess_tail.as_ref() {
9117 if let Some(lh) = lh.as_ref() {
9118 e.copy_into(&mut phs, 0, lh, n_embd)?;
9119 }
9120 }
9121 }
9122 if n_copy > 0 {
9123 e.copy_view_into(
9124 &mut phs,
9125 dst_off,
9126 &ph.slice(src_lo..src_lo + n_copy),
9127 n_copy,
9128 )?;
9129 }
9130 self.mtp_kv_fill(
9131 e,
9132 mtp,
9133 &prompt[start..end],
9134 &phs,
9135 base + start,
9136 &mut *scratch,
9137 embd_dev,
9138 )?;
9139 }
9140 start = end;
9141 }
9142 }
9143 // MEMRA_PROFILE_SPEC=2: profiler capture starts HERE — after the prime, so an
9144 // `nsys -c cudaProfilerApi` capture contains ONLY the round loop (draft/verify/commit).
9145 // (=1 brackets the whole call in run_spec.rs, prime included.)
9146 if std::env::var("MEMRA_PROFILE_SPEC").as_deref() == Ok("2") {
9147 unsafe extern "C" {
9148 fn cudaProfilerStart() -> i32;
9149 }
9150 unsafe {
9151 cudaProfilerStart();
9152 }
9153 }
9154 // ROUND-STREAM stage (c) 4 (MEMRA_SPEC_STREAM=1, experimental): pre-issued M-round
9155 // bursts with ZERO per-round host readbacks — the accept/seed/rollback/ring kernels
9156 // consume each other's device outputs; the host drains the ring every M rounds. v1
9157 // constraints: greedy, !spec_replay, single-shot, batched-linear layers, no refresh
9158 // fills (acceptance effect A/B-arbitrated), enters from round 1 (pending guaranteed).
9159 // NOTE: not gated on the caller's graph_draft (its trunk_dense conjunct turns the 35B
9160 // MoE off) — the stream capture encloses ONLY the dense MTP head; the head-dense /
9161 // full-prec / k gates are re-derived here and a failed capture degrades to stream-off.
9162 let stream_on = crate::spec::spec_stream()
9163 && !sampled
9164 && !spec_replay
9165 && constraint.is_none()
9166 && !session_mode
9167 && embd_gpu.is_some()
9168 && !crate::model::full_prec_enabled()
9169 && k + 2 < 96;
9170 let mut stream_graph: Option<cudarc::driver::CudaGraph> = None;
9171 let mut g_tokp2k = e.alloc_u32_zeroed(2 * k.max(1))?;
9172 if stream_on {
9173 let cap = e.capture_graph(|e| {
9174 for j in 0..k.max(1) {
9175 self.mtp_head_forward_cap(
9176 e,
9177 mtp,
9178 &mut dctx.g_tok,
9179 &mut dctx.g_pos,
9180 &mut dctx.g_seed,
9181 &mut dctx.g_p,
9182 &mut *scratch,
9183 true,
9184 true,
9185 embd_gpu.expect("round stream requires resident embedding"),
9186 embd_qt,
9187 embd_rb,
9188 d_vocab,
9189 None,
9190 Some((&mut g_tokp2k, j, d2t_dev.as_ref())),
9191 None, // round-stream requires constraint.is_none() (see stream_on)
9192 )?;
9193 }
9194 Ok(())
9195 });
9196 match cap {
9197 Ok(g) => {
9198 scratch.set_len(e, 0)?;
9199 stream_graph = Some(g);
9200 }
9201 Err(err) => {
9202 scratch.set_len(e, 0)?;
9203 if debug_spec {
9204 eprintln!("[spec] stream-graph capture failed ({err}); stream off");
9205 }
9206 }
9207 }
9208 }
9209 let stream_active = stream_on && stream_graph.is_some();
9210 if debug_spec {
9211 eprintln!(
9212 "[spec] stream_on={stream_on} env={} samp={sampled} dg={} captured={} active={stream_active} session={session_mode} replay={spec_replay}",
9213 crate::spec::spec_stream(),
9214 dctx.graph.is_some(),
9215 stream_graph.is_some()
9216 );
9217 }
9218 let t_v_s = k + 1;
9219 // ROUND-STREAM buffers + ptr tables now live in the model-generic round_stream
9220 // module (extracted 2026-07-12; the gemma burst reuses them).
9221 let sb = crate::round_stream::StreamBufs::new(e, k, crate::spec::spec_stream_m())?;
9222 let crate::round_stream::StreamBufs {
9223 mut vtok_d,
9224 mut brk_d,
9225 mut pend_d,
9226 last_pred_d,
9227 mut pos_ctr,
9228 mut pos_start_d,
9229 mut ring_d,
9230 acc_d: mut stream_acc,
9231 m_rounds,
9232 k: _,
9233 } = sb;
9234 let stream_ptrs: Option<CudaSlice<u64>> = if stream_active {
9235 Some(crate::round_stream::kv_len_ptr_table(
9236 e,
9237 cache,
9238 Some(&pos_ctr),
9239 )?)
9240 } else {
9241 None
9242 };
9243
9244 let t_fill = t_ent.elapsed();
9245 let mut round = 0usize;
9246 // ADAPTIVE DRAFT LENGTH (MEMRA_SPEC_ADAPT=1, opt-in — the gemma_spec accepted-run law,
9247 // ported 2026-08-01): next round's draft depth = last round's accepted run + 1, clamped
9248 // to [floor(pos), k_cap] — a miss shrinks the next draft to the miss point + 1,
9249 // full-accept streaks re-deepen one step per round. NOT the 2026-07-07 acceptance-EMA
9250 // (that arm measured an HONEST LOSS to static per-class optima — 115.0/85.8/73.4 vs
9251 // 121.6/92.7/75.6, EMA lag — and was removed 2026-07-08; rig5090.jsonl has the record).
9252 // The gemma law has no lag class: it reacts within one round, and was worth +7-20% on
9253 // the gemma cells at unchanged exactness (2026-07-10 flip; floor sweep 2026-07-25;
9254 // position key 2026-07-26). Signal = n_acc from the round's EXISTING accept readback —
9255 // zero new syncs; the draft graph is a SINGLE-STEP capture replayed per drafted token,
9256 // so a per-round depth needs no re-capture (unlike gemma's whole-chain graphs). qwen's
9257 // in-round p-min cut already shortens chains mid-round, so gemma's one-round-late p-min
9258 // fold into kc is unnecessary here — the accepted-run law sees the cut via n_acc.
9259 // Exactness is the verify's job at ANY depth (same contract as p-min variable rounds).
9260 // DEFAULT OFF on the qwen path until its cells gate a flip (gemma's is default-on).
9261 // MEASURED 2026-08-01 (H100 GPU-3, interleaved x3, NGEN=256, same-invocation plain
9262 // denominators; research/qwen-adaptive-k-20260801/): REFUTED on the tuned qwen configs.
9263 // q27 K=3+HPOST+PMIN=0.3: short +0.8% (noise; law ~idles, len_hist identical), board
9264 // -1.9%, agentic -0.5%; board PMIN=0 -2.1% (not p-min shadowing — the law itself);
9265 // floor=1 -2.8% (gemma's floor-collapse, reproduced). q35 K=2 board: -6.4% (52/136
9266 // rounds shrink to depth 1; no depth to reclaim at K=2). The gemma direction DOES
9267 // appear at untuned depth-K — q27 K=6 floor=4 +1.5% over fixed K=6 — but stays -3.7%
9268 // below fixed K=3: same verdict class as the retired EMA arm (honest loss to static
9269 // per-class optima). Acceptance-rate rises under the law while tokens/round falls —
9270 // it buys accept-% by adding rounds, and a round's fixed draft+verify cost wins.
9271 // K=1..8 self-consistency PASS both models with the law ON (exactness held).
9272 let adapt = std::env::var("MEMRA_SPEC_ADAPT").as_deref() == Ok("1");
9273 // floor: per-model default keyed on n_embd (gemma's tiering — models with an expensive
9274 // verify keep deep drafts after a miss); MEMRA_SPEC_ADAPT_FLOOR pins it everywhere.
9275 let adapt_floor_env: Option<usize> = std::env::var("MEMRA_SPEC_ADAPT_FLOOR")
9276 .ok()
9277 .and_then(|v| v.parse().ok());
9278 let adapt_floor_default: usize = if self.cfg.n_embd as usize >= 3500 {
9279 4
9280 } else if self.cfg.n_embd as usize >= 2500 {
9281 2
9282 } else {
9283 1
9284 };
9285 let adapt_floor: usize = adapt_floor_env.unwrap_or(adapt_floor_default);
9286 // position key: past floor_ctx a HIGH floor (>=4) relaxes to 1 — forced-deep drafts
9287 // turn net-negative at depth (gemma 31B d1736 evidence); MEMRA_SPEC_FLOOR_CTX moves
9288 // the boundary, an explicit MEMRA_SPEC_ADAPT_FLOOR pins the floor everywhere.
9289 let floor_ctx: usize = std::env::var("MEMRA_SPEC_FLOOR_CTX")
9290 .ok()
9291 .and_then(|v| v.parse().ok())
9292 .unwrap_or(1024);
9293 let floor_at = |pos: usize| -> usize {
9294 if adapt_floor_env.is_some() || pos < floor_ctx {
9295 adapt_floor
9296 } else if adapt_floor >= 4 {
9297 1
9298 } else {
9299 adapt_floor
9300 }
9301 };
9302 // cap: MEMRA_SPEC_CAPMAX (gemma semantics, default 7). Binds only under adapt — the
9303 // fixed-K default path is untouched by this whole block.
9304 let cap_max: usize = std::env::var("MEMRA_SPEC_CAPMAX")
9305 .ok()
9306 .and_then(|v| v.parse().ok())
9307 .unwrap_or(7);
9308 let k_cap = k.min(cap_max).max(1);
9309 let mut kc = k_cap;
9310 let mut opti_fork: Option<OptiForkState> = None;
9311 let mut fork_snapshot: Option<crate::cache::CacheSnapshot> = None;
9312 if fork_mode != OptiForkGateMode::Disabled {
9313 let fence = crate::pp::pp_cuts(self.layers.len());
9314 let refusal = if !session_mode {
9315 Some("not-session")
9316 } else if k != 1 || adapt {
9317 Some("requires-fixed-k1")
9318 } else if sampled || constraint.is_some() || spec_replay {
9319 Some("sampled-constrained-or-replay")
9320 } else if pipe.is_some() {
9321 Some("two-session-pipeline")
9322 } else if !spec_devacc() {
9323 Some("requires-device-accept")
9324 } else if stream_active || crate::spec::spec_stream() {
9325 Some("round-stream")
9326 } else if crate::cache::swa_ring_on() || cache.has_swa_ring() {
9327 Some("swa-ring")
9328 } else if crate::pp::pp_host_bounce_active() {
9329 Some("host-bounce")
9330 } else if fork_mode == OptiForkGateMode::Controller
9331 && cache.recur.iter().any(Option::is_some)
9332 {
9333 Some("controller-requires-zero-recurrent-state")
9334 } else if fence.as_ref().is_none_or(|f| f.len() != 3) {
9335 Some("requires-pp2")
9336 } else {
9337 None
9338 };
9339 if let Some(reason) = refusal {
9340 OPTI_FORK_REFUSALS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
9341 eprintln!("[opti-fork] refused reason={reason}");
9342 } else {
9343 let fence = fence.expect("validated PP-2 fence");
9344 let rt = crate::pp::PpNRt::get(e)?;
9345 let primary_stage0 = rt.engine(0, e).ctx().ordinal() == e.ctx().ordinal();
9346 let primary_stage1 = rt.engine(1, e).ctx().ordinal() == e.ctx().ordinal();
9347 let primary_supported =
9348 primary_stage0 || (fork_mode == OptiForkGateMode::Controller && primary_stage1);
9349 if !rt.cross_device() || !primary_supported {
9350 OPTI_FORK_REFUSALS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
9351 eprintln!("[opti-fork] refused reason=requires-supported-primary-cross-device");
9352 } else {
9353 // Both recurrent snapshots and both seed generations are allocated before
9354 // the first fork, each through its owning PP stage. Allocation failure
9355 // therefore happens before any optimistic state mutation can occur.
9356 let current_snapshot = opti_snapshot_stage_owned(e, cache, rt, &fence)?;
9357 let alternate_snapshot = opti_snapshot_stage_owned(e, cache, rt, &fence)?;
9358 let fork = OptiForkState::new(
9359 e,
9360 cache,
9361 fork_mode,
9362 alternate_snapshot,
9363 &h_seed_buf,
9364 &fill_prev,
9365 rt,
9366 fence[1],
9367 self.layers.len(),
9368 )?;
9369 eprintln!(
9370 "[opti-fork] armed mode={fork_mode:?} snapshots=2 seeds=2 split={} \
9371 payload_dev0={} payload_dev1={} q_threshold={:.3}",
9372 fence[1],
9373 fork.logical_payload_bytes[0],
9374 fork.logical_payload_bytes[1],
9375 fork.controller.map_or(0.0, |policy| policy.threshold),
9376 );
9377 fork_snapshot = Some(current_snapshot);
9378 opti_fork = Some(fork);
9379 }
9380 }
9381 }
9382 // Persistent snapshot buffers are allocated once and refreshed in place. The fork arm
9383 // uses stage-owned snapshots; refused/disabled arms retain the existing generic helper.
9384 let mut snap = match fork_snapshot {
9385 Some(snapshot) => snapshot,
9386 None => cache.snapshot(e)?,
9387 };
9388 let mut carried_opti: Option<OptiControllerTicket> = None;
9389 // ROUND-STREAM stage (b) 3a: device table of per-layer kvl.len_d pointers (stable — the
9390 // cache never reallocates len_d; see cache.rs "stable pointer" note). 0 = no KV layer.
9391 let kv_len_ptrs: Option<CudaSlice<u64>> = if spec_devacc() && !spec_replay {
9392 Some(crate::round_stream::kv_len_ptr_table(e, cache, None)?)
9393 } else {
9394 None
9395 };
9396 // BONUS FOLD (2026-07-04): after a FULL accept the bonus token is NOT committed with a
9397 // separate T=1 trunk pass (a full weight read per round). It stays PENDING and rides as
9398 // column 0 of the NEXT round's verify batch. Under predecessor pairing the next chain
9399 // seeds from the bonus's predecessor's TRUE verify hidden (free — no extra
9400 // pass of any kind). Verify still
9401 // checks every emitted token against the target -> exactness holds by construction; only
9402 // DRAFT QUALITY can shift, which the acceptance numbers arbitrate.
9403 // bonus emitted but not yet committed to cache. A carried pending (see SpecSession::
9404 // pending_tok) enters round 0 directly — the burst boundary becomes a plain round edge.
9405 let mut pending: Option<u32> = carried_pending;
9406 // MEMRA_SPEC_PHASE=1: per-round wall decomposition (draft / verify / accept+commit) —
9407 // no tracing, no extra syncs (each phase is naturally sync-bounded: draft readbacks,
9408 // the verify accept readback). Printed once at loop end via spec-stats.
9409 let anatomy_on = std::env::var("MEMRA_SPEC_PP_ANATOMY").as_deref() == Ok("1");
9410 let phase_on = anatomy_on || std::env::var("MEMRA_SPEC_PHASE").as_deref() == Ok("1");
9411 // DRAFT-MASK receipt (lane/draft-mask): speculative-clone wall + rounds, printed with
9412 // spec-stats. The clone is the one cost the design adds per round — measured, not assumed.
9413 let (mut dm_clone_ns, mut dm_rounds) = (0u128, 0usize);
9414 // grammar-truncation counters: how many rounds the verify-side cut fired and how many
9415 // already-verified tokens it threw away. THIS is the quantity draft masking targets.
9416 let (mut dm_cuts, mut dm_cut_tokens) = (0usize, 0usize);
9417 let (mut ph_draft, mut ph_verify, mut ph_rest) = (0f64, 0f64, 0f64);
9418 let mut ph_wait = 0f64;
9419 let mut ph_commit = 0f64;
9420 let mut ph_t = std::time::Instant::now();
9421 let mut ph_mark = |acc: &mut f64, on: bool| {
9422 if on {
9423 let now = std::time::Instant::now();
9424 *acc += (now - ph_t).as_secs_f64();
9425 ph_t = now;
9426 }
9427 };
9428 if let Some(p) = pipe {
9429 p.setup_end();
9430 }
9431 while keep_going && out.len() < max_new {
9432 // ROUND-STREAM BURST: from round 1 (pending guaranteed by every non-replay arm),
9433 // issue M rounds with zero readbacks, then drain the ring + reconcile mirrors.
9434 if let (true, Some(sg), Some(ptrs)) = (
9435 stream_active && round >= 1 && pending.is_some(),
9436 &stream_graph,
9437 &stream_ptrs,
9438 ) {
9439 if debug_spec {
9440 static ONCE: std::sync::Once = std::sync::Once::new();
9441 ONCE.call_once(|| {
9442 eprintln!("[memra] ROUND-STREAM burst engaged (M={m_rounds} k={k})")
9443 });
9444 }
9445 e.set_i32_one(&mut pos_ctr, cache.pos as i32)?;
9446 e.set_u32_one(&mut pend_d, pending.unwrap())?;
9447 e.set_u32_one(&mut ring_d, 0)?; // ring count = 0 (writes element 0)
9448 for _mi in 0..m_rounds {
9449 e.i32_copy_add(&pos_ctr, &mut pos_start_d, 0)?;
9450 cache.snapshot_into(e, &mut snap)?; // device D2Ds, stream-ordered
9451 e.i32_copy_add(&pos_ctr, &mut scratch.kv.len_d, 0)?; // draft-KV rollback
9452 e.i32_copy_add(&pos_ctr, &mut dctx.g_pos, 1)?; // rope pos = pos + base
9453 e.u32_copy(&pend_d, &mut dctx.g_tok)?;
9454 e.copy_into(&mut dctx.g_seed, 0, &h_seed_buf, n_embd)?;
9455 sg.launch()?;
9456 e.spec_assemble_verify(
9457 &g_tokp2k,
9458 &pend_d,
9459 d2t_dev.as_ref(),
9460 &mut vtok_d,
9461 &mut brk_d,
9462 p_min,
9463 k,
9464 pmin0,
9465 )?;
9466 let mut ck = VerifyCkpt::new(self.layers.len());
9467 let dummy = vec![0u32; t_v_s];
9468 let (tl_d, vx) = self.decode_step_t_core_stream(
9469 e,
9470 &dummy,
9471 0,
9472 &mut *cache,
9473 embd_dev,
9474 Some(&mut ck),
9475 Some((&vtok_d, &pos_ctr)),
9476 None,
9477 None,
9478 None,
9479 )?;
9480 for j in 0..t_v_s {
9481 e.argmax_token_device_col(&tl_d, j, n_vocab, &mut preds_d, j)?;
9482 }
9483 e.spec_accept_greedy_dc(
9484 &preds_d,
9485 &vtok_d,
9486 &last_pred_d,
9487 &brk_d,
9488 &mut stream_acc,
9489 )?;
9490 e.spec_seed_gather(&vx, &fill_prev, &stream_acc, &mut h_seed_buf, 1, n_embd)?;
9491 e.copy_into(&mut fill_prev, 0, &h_seed_buf, n_embd)?;
9492 self.commit_verified_prefix_stream(
9493 e,
9494 &mut *cache,
9495 &snap,
9496 &ck,
9497 &stream_acc,
9498 1,
9499 t_v_s,
9500 )?;
9501 e.spec_rollback_stream(
9502 ptrs,
9503 &pos_start_d,
9504 &stream_acc,
9505 1,
9506 self.layers.len() + 1,
9507 )?;
9508 e.spec_ring_commit(&vtok_d, &stream_acc, &brk_d, &mut ring_d, &mut pend_d)?;
9509 }
9510 e.stream().synchronize()?;
9511 let ring_h = e.dtoh_u32(&ring_d)?;
9512 let cnt = ring_h[0] as usize;
9513 for i in 0..cnt {
9514 if out.len() < max_new {
9515 out.push(ring_h[1 + i]);
9516 }
9517 }
9518 let pos_h = e.dtoh_i32(&pos_ctr)?[0] as usize;
9519 for il in 0..self.layers.len() {
9520 if let Some(kvl) = cache.kv[il].as_mut() {
9521 kvl.len = pos_h;
9522 }
9523 }
9524 cache.pos = pos_h;
9525 scratch.kv.len = pos_h;
9526 pending = Some(ring_h[cnt]); // last drained token = the live bonus
9527 last_token = ring_h[cnt];
9528 total_drafted += k * m_rounds; // upper bound (p-min breaks uncounted)
9529 total_accepted += cnt.saturating_sub(m_rounds);
9530 if let Some(t) = sess_telem {
9531 // totals only — the burst's per-round accept counts stayed on device
9532 // (that is the point of the round-stream arm). pos_* untouched.
9533 t.record_totals(m_rounds, k * m_rounds, cnt.saturating_sub(m_rounds));
9534 }
9535 round += m_rounds;
9536 // sse-cadence: the drained ring is committed — flush it at burst-drain cadence.
9537 keep_going = flush_commit(&mut on_commit, &out, &mut flushed);
9538 continue;
9539 }
9540 let pipe_draft = match pipe {
9541 Some(p) => Some(p.draft_begin(round)?),
9542 None => None,
9543 };
9544 let pos = cache.pos; // #tokens committed (EXCLUDES a pending bonus)
9545 let mut current_opti = carried_opti.take();
9546 let mut fork_generation = if current_opti.is_none() && pending.is_some() {
9547 match opti_fork.as_mut() {
9548 Some(fork) if fork.mode.is_forced() => Some(fork.reserve(&mut snap)?),
9549 None => None,
9550 Some(_) => None,
9551 }
9552 } else {
9553 None
9554 };
9555 if current_opti.is_none() {
9556 if let Some(fork) = opti_fork.as_ref() {
9557 opti_snapshot_stage_owned_into(e, cache, fork.rt, &fork.fence, &mut snap)?;
9558 } else {
9559 cache.snapshot_into(e, &mut snap)?;
9560 }
9561 } else if snap.pos != pos {
9562 return Err(format!(
9563 "optipipe carried snapshot pos {} != current pos {pos}",
9564 snap.pos
9565 )
9566 .into());
9567 } // §C: snapshot BEFORE draft+verify (already retained for a carried successor)
9568 ph_mark(&mut ph_rest, phase_on);
9569
9570 // --- 1. DRAFT k tokens with the NextN head (autoregressive, T=1 each) ---
9571 // p-min semantics (both paths): stop the chain early when the head's confidence in
9572 // its own pick drops below p_min — the just-drafted token is DISCARDED, but its
9573 // scratch append stands (identical to the eager chain's ordering). j==0 always drafts.
9574 let base0 = if pending.is_some() { 1usize } else { 0usize };
9575 // fixed draft length by default; MEMRA_SPEC_ADAPT=1 drafts at last round's
9576 // accepted run + 1 (the gemma law — see the setup block above the loop).
9577 let k_this = if adapt { kc } else { k };
9578 let mut draft: Vec<u32> = Vec::with_capacity(k);
9579 let mut draft_idx: Vec<u32> = Vec::with_capacity(k); // trimmed-vocab ids (== draft when untrimmed)
9580 let mut controller_draft_prob: Option<f32> = None;
9581 let mut controller_eager_state: Option<(u32, CudaSlice<f32>)> = None;
9582 if let Some(ticket) = current_opti.as_mut() {
9583 let carried_pending = pending.ok_or("optipipe carried successor lost pending")?;
9584 if ticket.verify_tokens[0] != carried_pending {
9585 return Err(format!(
9586 "optipipe carried pending mismatch: ticket={} live={carried_pending}",
9587 ticket.verify_tokens[0],
9588 )
9589 .into());
9590 }
9591 draft.push(ticket.verify_tokens[1]);
9592 controller_draft_prob = Some(ticket.draft_prob);
9593 controller_eager_state = ticket
9594 .take_eager_seed()
9595 .map(|seed| (ticket.verify_tokens[1], seed));
9596 } else {
9597 // Round-start draft-KV sync (BOTH paths). Persistent: truncate/align to the committed
9598 // history — slots 0..P hold entries for the tokens before last_token@P (P = pos +
9599 // base0 - 1); this single set_len IS the draft-side rollback (drops last round's
9600 // rejected drafts and p-min extras via the len mechanism).
9601 scratch.set_len(e, pos + base0 - 1)?;
9602 if pen_on {
9603 // PEN_WINDOW_MAX also bounds the per-round upload and the O(n_hist^2)
9604 // device dedup: `penalty_last_n` is usize::MAX for any serve request with
9605 // a penalty, so without the cap this grew with the whole session.
9606 let win = sp.penalty_last_n.min(PEN_WINDOW_MAX);
9607 let w0 = pen_hist.len().saturating_sub(win);
9608 pen_hist_d = Some(e.htod_u32_v(&pen_hist[w0..])?);
9609 }
9610 if sampled {
9611 draft_logits.clear();
9612 draft_stats.clear();
9613 }
9614 // DRAFT-SIDE GRAMMAR MASK: clone the committed grammar state ONCE per round; each
9615 // position's mask is computed on that clone and advanced by the PROPOSED token. The
9616 // real state moves only on emission (verify's job), so the emitted stream is
9617 // unchanged — the mask only removes tokens the verify would have truncated anyway.
9618 let mut dmask_live = dmask_on;
9619 if dmask_live {
9620 let t_c = std::time::Instant::now();
9621 constraint
9622 .as_deref_mut()
9623 .unwrap()
9624 .draft_begin()
9625 .map_err(|e2| format!("constraint: {e2}"))?;
9626 dm_clone_ns += t_c.elapsed().as_nanos();
9627 dm_rounds += 1;
9628 }
9629 if let (false, Some(gr)) = (sampled || pen_on, &dctx.graph) {
9630 // GRAPH DRAFT: one dispatch per drafted token. The chain feeds itself on-device
9631 // (in-graph argmax -> tok_d -> next replay's embed; h_nextn -> h_seed_d; pos_d
9632 // inc'd in-graph); the host only reads 4B token (+4B p) and decides the break.
9633 e.set_i32_one(&mut dctx.g_pos, (pos + base0) as i32)?;
9634 e.set_u32_one(&mut dctx.g_tok, last_token)?;
9635 e.copy_into(&mut dctx.g_seed, 0, &h_seed_buf, n_embd)?;
9636 for j in 0..k_this {
9637 // per-position mask upload (contents only — the graph's baked pointer is
9638 // dctx.g_dmask). All-ones once masking goes dead mid-chain, so the captured
9639 // mask node degrades to a no-op ban instead of needing a second graph.
9640 if dmask_live
9641 && !upload_draft_mask(
9642 e,
9643 constraint.as_deref_mut().unwrap(),
9644 &mut dctx.g_dmask,
9645 mtp.d2t.as_ref(),
9646 d_vocab,
9647 dmask_words,
9648 )?
9649 {
9650 // no draft-vocab row is grammar-legal here (a trimmed FR-Spec head can
9651 // genuinely miss the legal set): neutralize the captured mask node and
9652 // finish the chain UNMASKED — exactly pre-lane behaviour, never worse.
9653 e.htod_u32_into(&mut dctx.g_dmask, &vec![u32::MAX; dmask_words])?;
9654 dmask_live = false;
9655 }
9656 gr.launch()?;
9657 scratch.kv.len += 1; // host mirror (len_d advanced in-graph)
9658 let idx = e.dtoh_u32_one(&dctx.g_tok)?;
9659 // #87 SENTINEL TRAP: an all-NaN head-logits row leaves the device argmax's
9660 // init sentinel (0x7FFFFFFF) in g_tok — feeding it onward dereferences
9661 // embed_row(sentinel) = table + ~4.6TB (never mapped) inside the NEXT graph
9662 // replay's embed node, and the MMU fault kills the CUDA context for the
9663 // whole process (research/pp2spec-crash-20260807: 3 coredumps, byte-exact
9664 // VA arithmetic). Refuse loudly instead; the diagnostics name the first-NaN
9665 // buffer (g_seed = the verify-side handoff vs head-side compute).
9666 if (idx as usize) >= d_vocab {
9667 // g_seed is SELF-FED (the replay writes h_nextn back into it), so it
9668 // reads as the head's OUTPUT at j; h_seed_buf is the round's INPUT
9669 // seed, untouched since the round-start copy — the pair discriminates
9670 // "seed arrived poisoned" from "head forward produced NaN".
9671 let seed_h = e.dtoh(&dctx.g_seed)?;
9672 let seed_nan = seed_h.iter().filter(|v| v.is_nan()).count();
9673 let in_h = e.dtoh(&h_seed_buf)?;
9674 let in_nan = in_h.iter().filter(|v| v.is_nan()).count();
9675 return Err(format!(
9676 "draft(graph) argmax sentinel 0x{idx:08x} >= d_vocab {d_vocab} at \
9677 round {round} j={j} pos={pos}: head-out NaN {seed_nan}/{n_embd}, \
9678 round-input-seed NaN {in_nan}/{n_embd} — refusing to dereference \
9679 the embed row (#87 trap)"
9680 )
9681 .into());
9682 }
9683 // trimmed draft vocab -> target token id (identity when no d2t map)
9684 let d = match &mtp.d2t {
9685 Some(map) => map[idx as usize],
9686 None => idx,
9687 };
9688 let draft_p = if p_min > 0.0
9689 || opti_fork
9690 .as_ref()
9691 .is_some_and(|fork| fork.controller.is_some())
9692 {
9693 Some(e.dtoh(&dctx.g_p)?[0])
9694 } else {
9695 None
9696 };
9697 if j == 0 {
9698 controller_draft_prob = draft_p;
9699 }
9700 if let Some(p) = draft_p.filter(|_| p_min > 0.0) {
9701 if p < p_min && (j > 0 || (pmin0 && base0 == 1)) {
9702 break;
9703 }
9704 }
9705 draft.push(d);
9706 // with a trimmed head the NEXT embed must read the TARGET id, not the draft
9707 // index the argmax wrote — patch the persistent token buffer (4B htod).
9708 if d != idx {
9709 e.set_u32_one(&mut dctx.g_tok, d)?;
9710 }
9711 // advance the SPECULATIVE state with the proposal; a dead chain drops to
9712 // unmasked drafting for the remaining positions (verify still arbitrates).
9713 // speculative advance; a chain the grammar can no longer follow (EOS
9714 // proposed) ends here. The captured mask node always runs, so a dead chain
9715 // leaves the buffer NEUTRAL (all-ones = ban nothing) before it exits.
9716 if dmask_live
9717 && !constraint
9718 .as_deref_mut()
9719 .unwrap()
9720 .draft_advance(d)
9721 .map_err(|e2| format!("constraint: {e2}"))?
9722 {
9723 e.htod_u32_into(&mut dctx.g_dmask, &vec![u32::MAX; dmask_words])?;
9724 break;
9725 }
9726 }
9727 // PURE-TEMP RE-TEST (lane/graph-s-key-exactness-20260819): the sampled graph is
9728 // legal ONLY in the regime it was captured in. The condition used to read
9729 // `(sampled, &dctx.graph_s)` and trusted `s_key` to have dropped anything else —
9730 // which it could not, because the key omitted the filters. Both halves are now
9731 // enforced: the key drops a stale graph, and this site refuses to launch one.
9732 } else if let (true, Some(gr)) = (sampled && pure_temp, &dctx.graph_s) {
9733 if skey_probe() {
9734 eprintln!(
9735 "[skey] chain=graph_s round={round} pure_temp={} top_k={} \
9736 top_p={} min_p={} s_key_parked={:?}",
9737 pure_temp as u8, sp.top_k, sp.top_p, sp.min_p, dctx.s_key,
9738 );
9739 }
9740 // SAMPLED GRAPH DRAFT: one replay per drafted token — head forward + gumbel +
9741 // argmax in ONE dispatch; the host reads 4B token (+4B p), D2Ds q into slot j,
9742 // and decides the break. Event-counter continuity: g_ctr is host-seeded to
9743 // sctr-1 ONCE per round (outside the graph); the in-graph bump runs BEFORE the
9744 // perturb, so replay j consumes counter sctr+j — exactly the eager arm's Philox
9745 // stream. Host sctr advances in lockstep (computed, no readback needed).
9746 e.set_i32_one(&mut dctx.g_pos, (pos + base0) as i32)?;
9747 e.set_u32_one(&mut dctx.g_tok, last_token)?;
9748 e.copy_into(&mut dctx.g_seed, 0, &h_seed_buf, n_embd)?;
9749 e.set_u32_one(&mut dctx.g_ctr, sctr.wrapping_sub(1))?;
9750 for j in 0..k_this {
9751 gr.launch()?;
9752 scratch.kv.len += 1; // host mirror (len_d advanced in-graph)
9753 sctr += 1; // mirrors the in-graph g_ctr bump (eager parity:
9754 // counts the p-min-discarded token too)
9755 // q retention: ONE async D2D of the persistent head-logits buffer into this
9756 // round's slot j (stream-ordered after the replay, before the next one).
9757 e.copy_into(&mut dctx.q_slots[j], 0, &dctx.g_q, d_vocab)?;
9758 let idx = e.dtoh_u32_one(&dctx.g_tok)?;
9759 // #87 SENTINEL TRAP (see the greedy graph arm above).
9760 if (idx as usize) >= d_vocab {
9761 let seed_h = e.dtoh(&dctx.g_seed)?;
9762 let seed_nan = seed_h.iter().filter(|v| v.is_nan()).count();
9763 return Err(format!(
9764 "draft(graph-sampled) argmax sentinel 0x{idx:08x} >= d_vocab \
9765 {d_vocab} at round {round} j={j} pos={pos}: round-seed NaN \
9766 {seed_nan}/{n_embd} — refusing to dereference the embed row \
9767 (#87 trap)"
9768 )
9769 .into());
9770 }
9771 let d = match &mtp.d2t {
9772 Some(map) => map[idx as usize],
9773 None => idx,
9774 };
9775 draft_idx.push(idx);
9776 if p_min > 0.0 {
9777 let p = e.dtoh(&dctx.g_p)?[0];
9778 if p < p_min && (j > 0 || (pmin0 && base0 == 1)) {
9779 break;
9780 }
9781 }
9782 draft.push(d);
9783 // trimmed head: the NEXT embed must read the TARGET id (see the greedy arm).
9784 if d != idx {
9785 e.set_u32_one(&mut dctx.g_tok, d)?;
9786 }
9787 }
9788 // uniform accept path: fill draft_stats per used slot (pure-temp regime — the
9789 // stats degenerate to th=0 / full-Z; one filter_stats launch per slot, tiny).
9790 for j in 0..draft.len().max(draft_idx.len()) {
9791 let rows0 = e.htod_i32(&[0])?;
9792 let (mut th_d, mut z_d, mut mx_d) = (e.zeros(1)?, e.zeros(1)?, e.zeros(1)?);
9793 e.filter_stats(
9794 &dctx.q_slots[j],
9795 d_vocab,
9796 &rows0,
9797 &mut th_d,
9798 &mut z_d,
9799 &mut mx_d,
9800 d_vocab,
9801 1,
9802 sp_temp,
9803 sp.top_k,
9804 sp.top_p,
9805 sp.min_p,
9806 )?;
9807 draft_stats.push((e.dtoh(&mx_d)?[0], e.dtoh(&th_d)?[0], e.dtoh(&z_d)?[0]));
9808 }
9809 } else {
9810 if skey_probe() && sampled {
9811 eprintln!(
9812 "[skey] chain=eager round={round} pure_temp={} top_k={} \
9813 top_p={} min_p={} s_key_parked={:?}",
9814 pure_temp as u8, sp.top_k, sp.top_p, sp.min_p, dctx.s_key,
9815 );
9816 }
9817 // EAGER DRAFT (fallback: MoE head/trunk, huge k, MEMRA_SPEC_NOGRAPH, capture fail).
9818 let mut e_tok = last_token;
9819 let mut d_seed = e.clone_dtod(&h_seed_buf)?;
9820 for j in 0..k_this {
9821 // GPU-ARGMAX DRAFT (2026-07-03): device logits + device argmax + 4-byte token
9822 // read instead of the ~600KB full-vocab dtoh + host argmax per draft token.
9823 let mtp_pos = pos + base0 + j;
9824 // draft-side grammar mask (eager twin of the graph arm's in-graph node).
9825 // A position with no legal draft-vocab row drops to unmasked drafting for
9826 // the rest of the chain (pre-lane behaviour; verify still arbitrates).
9827 if dmask_live {
9828 dmask_live = upload_draft_mask(
9829 e,
9830 constraint.as_deref_mut().unwrap(),
9831 &mut dctx.g_dmask,
9832 mtp.d2t.as_ref(),
9833 d_vocab,
9834 dmask_words,
9835 )?;
9836 }
9837 let (dl_d, h_nextn) = self.mtp_head_forward_dev(
9838 e,
9839 mtp,
9840 e_tok,
9841 &d_seed,
9842 &mut *scratch,
9843 mtp_pos,
9844 embd_dev,
9845 if dmask_live {
9846 Some((&dctx.g_dmask, dmask_words))
9847 } else {
9848 None
9849 },
9850 )?;
9851 let tok_d = if sampled {
9852 // FILTERED Gumbel-max: stats -> masked perturb -> argmax = one draw from
9853 // the filtered softmax (filters off => th=0, exact v1 semantics).
9854 if perturb_buf.is_none() {
9855 perturb_buf = Some(e.zeros(d_vocab.max(n_vocab))?);
9856 }
9857 let mut q_row = e.clone_dtod(&dl_d)?; // retained q (penalized when on)
9858 if pen_on {
9859 let h = pen_hist_d.as_ref().unwrap();
9860 let nh = h.len();
9861 e.penalize_logits(
9862 &mut q_row,
9863 h,
9864 nh,
9865 sp.penalty_repeat,
9866 sp.penalty_freq,
9867 sp.penalty_present,
9868 d_vocab,
9869 )?;
9870 }
9871 let rows0 = e.htod_i32(&[0])?;
9872 let (mut th_d, mut z_d, mut mx_d) =
9873 (e.zeros(1)?, e.zeros(1)?, e.zeros(1)?);
9874 e.filter_stats(
9875 &q_row, d_vocab, &rows0, &mut th_d, &mut z_d, &mut mx_d, d_vocab,
9876 1, sp_temp, sp.top_k, sp.top_p, sp.min_p,
9877 )?;
9878 let (th, z, mx) =
9879 (e.dtoh(&th_d)?[0], e.dtoh(&z_d)?[0], e.dtoh(&mx_d)?[0]);
9880 let pb = perturb_buf.as_mut().unwrap();
9881 e.gumbel_perturb_filtered(
9882 &q_row, pb, d_vocab, sp_seed, sctr, sp_temp, mx, th,
9883 )?;
9884 sctr += 1;
9885 draft_logits.push(q_row);
9886 draft_stats.push((mx, th, z));
9887 e.argmax_token_device(pb, d_vocab)?
9888 } else {
9889 e.argmax_token_device(&dl_d, d_vocab)?
9890 };
9891 let idx = e.dtoh_u32_one(&tok_d)?;
9892 // #87 SENTINEL TRAP (eager twin — see the graph arm). Extra diagnostics
9893 // here because the eager chain's operands are all readable: dl_d (the head
9894 // logits row) and d_seed (this step's h_seed) name the first-NaN buffer.
9895 if (idx as usize) >= d_vocab {
9896 let dl_h = e.dtoh(&dl_d)?;
9897 let dl_nan = dl_h.iter().filter(|v| v.is_nan()).count();
9898 let seed_h = e.dtoh(&d_seed)?;
9899 let seed_nan = seed_h.iter().filter(|v| v.is_nan()).count();
9900 return Err(format!(
9901 "draft(eager) argmax sentinel 0x{idx:08x} >= d_vocab {d_vocab} at \
9902 round {round} j={j} pos={pos}: head-logits NaN {dl_nan}/{d_vocab}, \
9903 step-seed NaN {seed_nan}/{n_embd} — refusing to dereference the \
9904 embed row (#87 trap)"
9905 )
9906 .into());
9907 }
9908 let d = match &mtp.d2t {
9909 Some(map) => map[idx as usize],
9910 None => idx,
9911 };
9912 if sampled {
9913 draft_idx.push(idx);
9914 }
9915 let draft_p = if p_min > 0.0
9916 || opti_fork
9917 .as_ref()
9918 .is_some_and(|fork| fork.controller.is_some())
9919 {
9920 let p_d = e.prob_of_token_device(&dl_d, &tok_d, d_vocab)?;
9921 Some(e.dtoh(&p_d)?[0])
9922 } else {
9923 None
9924 };
9925 if j == 0 {
9926 controller_draft_prob = draft_p;
9927 }
9928 if let Some(p) = draft_p.filter(|_| p_min > 0.0) {
9929 if p < p_min && (j > 0 || (pmin0 && base0 == 1)) {
9930 break;
9931 }
9932 }
9933 draft.push(d);
9934 e_tok = d;
9935 d_seed = h_nextn;
9936 // speculative advance; a chain the grammar can no longer follow (EOS
9937 // proposed) ends here — the prefix already proposed still rides verify.
9938 if dmask_live
9939 && !constraint
9940 .as_deref_mut()
9941 .unwrap()
9942 .draft_advance(d)
9943 .map_err(|e2| format!("constraint: {e2}"))?
9944 {
9945 break;
9946 }
9947 }
9948 if opti_fork
9949 .as_ref()
9950 .is_some_and(|fork| fork.controller.is_some())
9951 {
9952 controller_eager_state = Some((e_tok, d_seed));
9953 }
9954 }
9955 }
9956 let k_round = draft.len();
9957 if let Some(p) = pipe {
9958 p.draft_end(round);
9959 }
9960 drop(pipe_draft);
9961
9962 ph_mark(&mut ph_draft, phase_on);
9963 // --- 2. VERIFY: one batched target forward. With a pending bonus, it rides as col 0
9964 // (committing its KV/recur inside the SAME weight read); drafts follow. ---
9965 let verify_tokens: Vec<u32> = match pending {
9966 Some(b) => {
9967 let mut v = Vec::with_capacity(k_round + 1);
9968 v.push(b);
9969 v.extend_from_slice(&draft);
9970 v
9971 }
9972 None => draft.clone(),
9973 };
9974 let base = if pending.is_some() { 1 } else { 0 };
9975 // ckpt (REPLAY-FREE partial accept): retain per-layer state-rebuild inputs alongside
9976 // the verify. Pure buffer keep-alives + dtod clones — kernel work is unchanged.
9977 let mut ckpt = if let Some(ticket) = current_opti.as_mut() {
9978 Some(ticket.take_ckpt())
9979 } else if spec_replay {
9980 None
9981 } else {
9982 Some(VerifyCkpt::new(self.layers.len()))
9983 };
9984 let controller_can_probe = base == 1
9985 && k_round == 1
9986 && out.len().saturating_add(2) < max_new
9987 && controller_draft_prob.is_some()
9988 && opti_fork
9989 .as_ref()
9990 .and_then(|fork| fork.controller.as_ref())
9991 .is_some_and(|policy| !policy.breaker_tripped);
9992 let mut successor_attempt: Option<OptiControllerTicket> = None;
9993 let mut rejected_probe: Option<(f32, u32)> = None;
9994 let mut controller_prepared: Option<OptiControllerPrepared> = None;
9995 if controller_can_probe {
9996 // Prepare d2/q and, on admission, d3 before either current verify half is
9997 // issued. N stage 0 can then be followed immediately by N+1 stage 0; once N's
9998 // boundary fires, those dev0 launches overlap N stage 1 on dev1. Preparing on
9999 // the primary stream after N stage 1 would serialize the supposed pipeline.
10000 let eager_pos = scratch.kv.len + 1;
10001 let (optimistic_pending, pending_probability) = self.opti_controller_draft_step(
10002 e,
10003 mtp,
10004 &mut dctx,
10005 &mut *scratch,
10006 d_vocab,
10007 &mut controller_eager_state,
10008 eager_pos,
10009 embd_dev,
10010 )?;
10011 let first_probability = controller_draft_prob
10012 .ok_or("optipipe controller probe lost first-token probability")?;
10013 let q_proxy = first_probability * pending_probability;
10014 OPTI_GATE_CHECKS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
10015 OPTI_SHADOW_DRAFT_TOKENS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
10016 let admitted = opti_fork
10017 .as_ref()
10018 .and_then(|fork| fork.controller.as_ref())
10019 .ok_or("optipipe controller policy disappeared")?
10020 .admit(q_proxy);
10021 if admitted {
10022 OPTI_GATE_ADMITS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
10023 let eager_pos = scratch.kv.len + 1;
10024 let (optimistic_draft, optimistic_draft_probability) = self
10025 .opti_controller_draft_step(
10026 e,
10027 mtp,
10028 &mut dctx,
10029 &mut *scratch,
10030 d_vocab,
10031 &mut controller_eager_state,
10032 eager_pos,
10033 embd_dev,
10034 )?;
10035 OPTI_SHADOW_DRAFT_TOKENS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
10036 let eager_seed = controller_eager_state.take().map(|(token, seed)| {
10037 debug_assert_eq!(token, optimistic_draft);
10038 seed
10039 });
10040 controller_prepared = Some(OptiControllerPrepared {
10041 verify_tokens: [optimistic_pending, optimistic_draft],
10042 draft_prob: optimistic_draft_probability,
10043 eager_seed,
10044 q_proxy,
10045 scratch_len: scratch.kv.len,
10046 });
10047 } else {
10048 OPTI_GATE_REJECTS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
10049 OPTI_WASTED_DRAFT_TOKENS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
10050 rejected_probe = Some((q_proxy, optimistic_pending));
10051 eprintln!(
10052 "[opti-controller] reject q={q_proxy:.6} threshold={:.3}",
10053 opti_fork
10054 .as_ref()
10055 .and_then(|fork| fork.controller.as_ref())
10056 .expect("controller policy")
10057 .threshold,
10058 );
10059 }
10060 }
10061 let fork_attempt = match fork_generation.take() {
10062 Some(generation) if base == 1 && k_round == 1 => Some(generation),
10063 Some(generation) => {
10064 opti_fork
10065 .as_mut()
10066 .expect("fork generation without fork state")
10067 .retire(generation)?;
10068 None
10069 }
10070 None => None,
10071 };
10072 let (tlogits_d, vx) = if let Some(p) = pipe {
10073 self.decode_step_t_core_pipelined(
10074 e,
10075 &verify_tokens,
10076 pos,
10077 &mut *cache,
10078 embd_dev,
10079 ckpt.as_mut(),
10080 p,
10081 round,
10082 )?
10083 } else if controller_can_probe {
10084 let fence = opti_fork
10085 .as_ref()
10086 .ok_or("optipipe controller probe lost fork state")?
10087 .fence;
10088 let boundary = match current_opti.as_mut() {
10089 Some(ticket) => ticket.take_boundary(),
10090 None => self.verify_stage0_issue(
10091 e,
10092 &verify_tokens,
10093 pos,
10094 &mut *cache,
10095 embd_dev,
10096 ckpt.as_mut(),
10097 None,
10098 &fence,
10099 Some(true),
10100 None,
10101 )?,
10102 };
10103 if let Some(prepared) = controller_prepared.take() {
10104 let generation = {
10105 let fork = opti_fork
10106 .as_mut()
10107 .ok_or("optipipe controller admission lost fork state")?;
10108 let generation = fork.reserve_successor()?;
10109 let rt = fork.rt;
10110 let snapshot_fence = fork.fence;
10111 opti_snapshot_one_stage_owned_into(
10112 e,
10113 cache,
10114 rt,
10115 &snapshot_fence,
10116 0,
10117 fork.successor_snapshot_mut(),
10118 )?;
10119 generation
10120 };
10121 let mut successor_ckpt = VerifyCkpt::new(self.layers.len());
10122 let successor_boundary = self.verify_stage0_issue(
10123 e,
10124 &prepared.verify_tokens,
10125 pos + verify_tokens.len(),
10126 &mut *cache,
10127 embd_dev,
10128 Some(&mut successor_ckpt),
10129 None,
10130 &fence,
10131 Some(false),
10132 None,
10133 )?;
10134 OPTI_FORK_ATTEMPTS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
10135 let fork = opti_fork
10136 .as_ref()
10137 .ok_or("optipipe controller ticket lost fork state")?;
10138 successor_attempt = Some(fork.controller_ticket(
10139 generation,
10140 successor_boundary,
10141 successor_ckpt,
10142 prepared.verify_tokens,
10143 prepared.draft_prob,
10144 prepared.eager_seed,
10145 prepared.q_proxy,
10146 prepared.scratch_len,
10147 ));
10148 eprintln!(
10149 "[opti-controller] issue generation={} q={:.6} threshold={:.3} \
10150 verify={:?}",
10151 generation.id,
10152 prepared.q_proxy,
10153 fork.controller.expect("controller policy").threshold,
10154 prepared.verify_tokens,
10155 );
10156 }
10157 let result = self.verify_stage1_finish(
10158 e,
10159 boundary,
10160 &mut *cache,
10161 ckpt.as_mut(),
10162 None,
10163 &fence,
10164 successor_attempt.is_none(),
10165 )?;
10166 if let Some(ticket) = current_opti.as_mut() {
10167 ticket.settle();
10168 }
10169 if successor_attempt.is_some() {
10170 let fork = opti_fork
10171 .as_mut()
10172 .ok_or("optipipe successor snapshot lost fork state")?;
10173 let rt = fork.rt;
10174 let snapshot_fence = fork.fence;
10175 opti_snapshot_one_stage_owned_into(
10176 e,
10177 cache,
10178 rt,
10179 &snapshot_fence,
10180 1,
10181 fork.successor_snapshot_mut(),
10182 )?;
10183 // Publish N only after both independent successor-state queues are complete.
10184 fork.rt.publish_to(1, &e.stream())?;
10185 }
10186 result
10187 } else if let Some(ticket) = current_opti.as_mut() {
10188 let fork = opti_fork
10189 .as_mut()
10190 .ok_or("optipipe carried controller ticket lost fork state")?;
10191 let boundary = ticket.take_boundary();
10192 let result = self.verify_stage1_finish(
10193 e,
10194 boundary,
10195 &mut *cache,
10196 ckpt.as_mut(),
10197 None,
10198 &fork.fence,
10199 true,
10200 )?;
10201 ticket.settle();
10202 result
10203 } else if let Some(generation) = fork_attempt {
10204 let fork = opti_fork
10205 .as_mut()
10206 .expect("fork generation without fork state");
10207 fork.capture_seed(e, generation, &h_seed_buf, &fill_prev, scratch.kv.len)?;
10208 let action = fork.mode.action(generation.id);
10209 let boundary = self.verify_stage0_issue(
10210 e,
10211 &verify_tokens,
10212 pos,
10213 &mut *cache,
10214 embd_dev,
10215 ckpt.as_mut(),
10216 None,
10217 &fork.fence,
10218 Some(true),
10219 None,
10220 )?;
10221 OPTI_FORK_ATTEMPTS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
10222 let mut ticket = fork.ticket(generation, boundary);
10223 if action == OptiForkAction::Abort {
10224 return Err(format!(
10225 "optipipe forced abort with generation {} stage0 in flight",
10226 generation.id,
10227 )
10228 .into());
10229 }
10230 fork.reconcile(
10231 e,
10232 &mut *cache,
10233 &mut *scratch,
10234 &snap,
10235 &mut h_seed_buf,
10236 &mut fill_prev,
10237 generation,
10238 action,
10239 verify_tokens[0],
10240 )?;
10241 let result = if action == OptiForkAction::Hit {
10242 let boundary = ticket.take_boundary();
10243 self.verify_stage1_finish(
10244 e,
10245 boundary,
10246 &mut *cache,
10247 ckpt.as_mut(),
10248 None,
10249 &fork.fence,
10250 true,
10251 )?
10252 } else {
10253 // The optimistic boundary slot has no reader. Re-run the unchanged serial
10254 // verify only after E_restart published the restored stage-0 state.
10255 self.decode_step_t_core(
10256 e,
10257 &verify_tokens,
10258 pos,
10259 &mut *cache,
10260 embd_dev,
10261 ckpt.as_mut(),
10262 )?
10263 };
10264 ticket.settle();
10265 debug_assert_eq!(ticket.generation, generation);
10266 fork.retire(generation)?;
10267 result
10268 } else {
10269 self.decode_step_t_core(
10270 e,
10271 &verify_tokens,
10272 pos,
10273 &mut *cache,
10274 embd_dev,
10275 ckpt.as_mut(),
10276 )?
10277 };
10278 let pipe_accept = match pipe {
10279 Some(p) => Some(p.accept_begin(round)?),
10280 None => None,
10281 };
10282
10283 ph_mark(&mut ph_verify, phase_on);
10284 // --- 3. GREEDY ACCEPT (walk prefix, stop at first mismatch) ---
10285 // DEVICE-ARGMAX ACCEPT: argmax every verify column ON DEVICE (same 2-pass kernels +
10286 // smallest-index tie-break as host argmax, argmax_gate-validated) and read back ONE
10287 // [T] u32 — replaces the T x n_vocab f32 dtoh + T host argmaxes per round.
10288 // t_pred[j] = target's greedy prediction for the slot after draft[j-1] (j>=1) or after
10289 // last_token (j==0). With a pending bonus, col 0 IS the prediction after last_token
10290 // (== the bonus), so every index shifts by `base` and last_pred is unused.
10291 let t_v = verify_tokens.len();
10292 let mut preds: Vec<u32> = Vec::new();
10293 if !sampled {
10294 for j in 0..t_v {
10295 e.argmax_token_device_col(&tlogits_d, j, n_vocab, &mut preds_d, j)?;
10296 }
10297 preds = e.dtoh_u32(&preds_d)?; // <- the verify-GPU wait lands here
10298 // #87 SENTINEL TRAP, verify side: a sentinel pred becomes the round's bonus =
10299 // next round's last_token = the next chain's embed lookup. Catch it at the
10300 // source with the column named — an all-NaN VERIFY column implicates the
10301 // stage-split trunk (decode_step_t_core_ppn), not the draft head.
10302 if let Some(bad) = preds[..t_v].iter().position(|&p| (p as usize) >= n_vocab) {
10303 let col = &tlogits_d.slice(bad * n_vocab..(bad + 1) * n_vocab);
10304 let mut probe = e.zeros(n_vocab)?;
10305 e.copy_view_into(&mut probe, 0, col, n_vocab)?;
10306 let col_h = e.dtoh(&probe)?;
10307 let col_nan = col_h.iter().filter(|v| v.is_nan()).count();
10308 return Err(format!(
10309 "verify argmax sentinel 0x{:08x} >= n_vocab {n_vocab} at round {round} \
10310 col {bad}/{t_v} pos={pos}: verify-logits col NaN {col_nan}/{n_vocab} \
10311 — the stage-split verify produced a poisoned column (#87 trap)",
10312 preds[bad]
10313 )
10314 .into());
10315 }
10316 }
10317 ph_mark(&mut ph_wait, phase_on);
10318 let t_pred = |j: usize| -> u32 {
10319 if j == 0 && base == 0 {
10320 last_pred
10321 } else {
10322 // GREEDY-ONLY: `preds` is filled under `if !sampled` above. The debug print
10323 // used to call this from the sampled arm and panicked the worker; it now goes
10324 // through `debug_t_pred0`. Keep the strict index here — in the greedy walk an
10325 // out-of-range pred is a real bug, not something to paper over.
10326 debug_assert!(
10327 !sampled,
10328 "t_pred is greedy-only: `preds` is empty in the sampled arm"
10329 );
10330 preds[base + j - 1]
10331 }
10332 };
10333 let mut devacc_seeded = false;
10334 let mut devacc_acc: Option<CudaSlice<u32>> = None;
10335 let (n_acc, bonus) = if !sampled {
10336 // ROUND-STREAM stage (a) (MEMRA_SPEC_DEVACC=1 opt-in): the walk runs ON DEVICE
10337 // (spec_accept_greedy, verbatim rule) and the host reads back 8B (n_acc, bonus)
10338 // instead of the [T] preds. Same sync count — machinery for stages (b)/(c),
10339 // gated on token identity vs the host walk (the arms below are bit-equal rules).
10340 if crate::spec::spec_devacc() && k_round > 0 && !spec_replay && constraint.is_none()
10341 {
10342 let draft_d = e.htod_u32_v(&draft)?;
10343 let mut acc_out = e.alloc_u32_zeroed(2)?;
10344 e.spec_accept_greedy(
10345 &preds_d,
10346 &draft_d,
10347 last_pred,
10348 base,
10349 k_round,
10350 &mut acc_out,
10351 )?;
10352 devacc_acc = Some(acc_out.clone());
10353 // stage (b): next-round seed gathered ON DEVICE from acc_out before the host
10354 // ever reads n_acc (j=base+n_acc -> vx col j-1; j==0 -> fill_prev). The three
10355 // non-replay commit arms skip their host-offset seed copies (guarded below);
10356 // the legacy spec_replay arm keeps its own rx-based seeding (excluded here).
10357 // NOTE: fill_prev is NOT updated here — the commit arms' TRUE-HIDDEN
10358 // REFRESH reads the OLD fill_prev (predecessor of this round's verify batch);
10359 // the update lands after the arms (devacc_seeded guard below).
10360 e.spec_seed_gather(&vx, &fill_prev, &acc_out, &mut h_seed_buf, base, n_embd)?;
10361 // 3a: KV lens roll back on device (len = saved + base + n_acc, all arms'
10362 // unified rule; full accept rewrites the verify-left value). Host mirrors
10363 // update after the readback; commit_verified_prefix skips its len_d writes.
10364 if let Some(successor) = successor_attempt.as_ref() {
10365 opti_fork
10366 .as_mut()
10367 .ok_or("optipipe successor reconcile lost fork state")?
10368 .queue_actual_reconcile(
10369 e,
10370 &snap,
10371 &acc_out,
10372 successor.verify_tokens[0],
10373 base,
10374 )?;
10375 } else if let Some(ptrs) = &kv_len_ptrs {
10376 let saved: Vec<i32> = (0..self.layers.len())
10377 .map(|il| snap.kv_len[il].map(|v| v as i32).unwrap_or(0))
10378 .collect();
10379 let saved_d = e.htod_i32(&saved)?;
10380 e.spec_rollback_kv(ptrs, &saved_d, &acc_out, base, self.layers.len())?;
10381 }
10382 devacc_seeded = true;
10383 let ab = e.dtoh_u32(&acc_out)?;
10384 (ab[0] as usize, ab[1])
10385 } else {
10386 let mut n_acc = 0usize;
10387 for j in 0..k_round {
10388 if t_pred(j) == draft[j] {
10389 n_acc += 1;
10390 } else {
10391 break;
10392 }
10393 }
10394 // bonus = target's own token at the first non-accepted slot. n_acc in 0..=k; t_pred
10395 // is defined for j in 0..=k (j==0 -> last_logits, j>=1 -> col j-1, last col = k-1).
10396 (n_acc, t_pred(n_acc))
10397 }
10398 } else {
10399 // --- SAMPLED ACCEPT (rejection sampling): u_j < p_j(x_j)/q_j(x_j) walk ---
10400 if col_buf.is_none() {
10401 col_buf = Some(e.zeros(n_vocab)?);
10402 }
10403 // FILTERED p_j: per-verify-col stats (one batched filter_stats call), then the
10404 // filtered gather. j==0&&base==0 reads last_col (its own stats row appended).
10405 let mut pj = vec![0f32; k_round.max(1)];
10406 let mut col_stats: Vec<(f32, f32, f32)> = Vec::new(); // (max, th, z) per verify col used
10407 if k_round > 0 {
10408 let mut ids: Vec<u32> = Vec::new();
10409 let mut rows: Vec<i32> = Vec::new();
10410 for j in 0..k_round {
10411 if j > 0 || base == 1 {
10412 ids.push(draft[j]);
10413 rows.push((base + j) as i32 - 1);
10414 }
10415 }
10416 if !ids.is_empty() {
10417 let nr = rows.len();
10418 // penalties: materialize the used columns into one contiguous penalized
10419 // buffer (rows remapped 0..nr) so stats+gathers see the penalized p.
10420 // penalties: materialize used columns contiguously, penalize all rows in
10421 // one launch, and point stats+gathers at the penalized buffer (rows 0..nr).
10422 let p_rows: Vec<i32> = if pen_on {
10423 (0..nr as i32).collect()
10424 } else {
10425 rows.clone()
10426 };
10427 if pen_on {
10428 if pcol_buf.as_ref().map(|b| b.len()).unwrap_or(0) < nr * n_vocab {
10429 pcol_buf = Some(e.zeros(nr * n_vocab)?);
10430 }
10431 let pc = pcol_buf.as_mut().unwrap();
10432 for (i2, &r) in rows.iter().enumerate() {
10433 let c = r as usize;
10434 e.copy_view_into(
10435 pc,
10436 i2 * n_vocab,
10437 &tlogits_d.slice(c * n_vocab..(c + 1) * n_vocab),
10438 n_vocab,
10439 )?;
10440 }
10441 let h = pen_hist_d.as_ref().unwrap();
10442 let nh = h.len();
10443 e.penalize_logits_rows(
10444 pc,
10445 h,
10446 nh,
10447 sp.penalty_repeat,
10448 sp.penalty_freq,
10449 sp.penalty_present,
10450 n_vocab,
10451 nr,
10452 )?;
10453 }
10454 let p_src: &CudaSlice<f32> = if pen_on {
10455 pcol_buf.as_ref().unwrap()
10456 } else {
10457 &tlogits_d
10458 };
10459 let rowsd = e.htod_i32(&p_rows)?;
10460 let (mut th_d, mut z_d, mut mx_d) =
10461 (e.zeros(nr)?, e.zeros(nr)?, e.zeros(nr)?);
10462 e.filter_stats(
10463 p_src, n_vocab, &rowsd, &mut th_d, &mut z_d, &mut mx_d, n_vocab, nr,
10464 sp_temp, sp.top_k, sp.top_p, sp.min_p,
10465 )?;
10466 let idsd = e.htod_u32_v(&ids)?;
10467 let mut outd = e.zeros(nr)?;
10468 e.softmax_gather_filtered(
10469 p_src, n_vocab, &idsd, &rowsd, &th_d, &z_d, &mut outd, n_vocab, nr,
10470 sp_temp,
10471 )?;
10472 let outv = e.dtoh(&outd)?;
10473 let (thv, zv, mxv) = (e.dtoh(&th_d)?, e.dtoh(&z_d)?, e.dtoh(&mx_d)?);
10474 let mut oi = 0usize;
10475 for j in 0..k_round {
10476 if j > 0 || base == 1 {
10477 pj[j] = outv[oi];
10478 oi += 1;
10479 }
10480 }
10481 col_stats = (0..nr).map(|i| (mxv[i], thv[i], zv[i])).collect();
10482 }
10483 if base == 0 {
10484 let lc: &CudaSlice<f32> = if pen_on {
10485 if col_buf.is_none() {
10486 col_buf = Some(e.zeros(n_vocab)?);
10487 }
10488 let cb = col_buf.as_mut().unwrap();
10489 e.copy_into(
10490 cb,
10491 0,
10492 last_col_logits
10493 .as_ref()
10494 .expect("sampled: last_col_logits unset"),
10495 n_vocab,
10496 )?;
10497 let h = pen_hist_d.as_ref().unwrap();
10498 let nh = h.len();
10499 e.penalize_logits(
10500 cb,
10501 h,
10502 nh,
10503 sp.penalty_repeat,
10504 sp.penalty_freq,
10505 sp.penalty_present,
10506 n_vocab,
10507 )?;
10508 col_buf.as_ref().unwrap()
10509 } else {
10510 last_col_logits
10511 .as_ref()
10512 .expect("sampled: last_col_logits unset")
10513 };
10514 let rows0 = e.htod_i32(&[0])?;
10515 let (mut th_d, mut z_d, mut mx_d) = (e.zeros(1)?, e.zeros(1)?, e.zeros(1)?);
10516 e.filter_stats(
10517 lc, n_vocab, &rows0, &mut th_d, &mut z_d, &mut mx_d, n_vocab, 1,
10518 sp_temp, sp.top_k, sp.top_p, sp.min_p,
10519 )?;
10520 let idsd = e.htod_u32_v(&[draft[0]])?;
10521 let mut outd = e.zeros(1)?;
10522 e.softmax_gather_filtered(
10523 lc, n_vocab, &idsd, &rows0, &th_d, &z_d, &mut outd, n_vocab, 1, sp_temp,
10524 )?;
10525 pj[0] = e.dtoh(&outd)?[0];
10526 last_col_stats =
10527 Some((e.dtoh(&mx_d)?[0], e.dtoh(&th_d)?[0], e.dtoh(&z_d)?[0]));
10528 }
10529 }
10530 // q source: the graph arm retained the head logits in the persistent q_slots;
10531 // the eager arm in per-round draft_logits clones. Same raw-logit values either way.
10532 // FILTERED q_j: stats from draft_stats (eager pushes in-chain; the graph arm
10533 // computes them post-replay — graph engages only filter/penalty-free, so the
10534 // stats degenerate to th=0/full-Z there, keeping ONE accept path).
10535 let q_bufs: &[CudaSlice<f32>] = if dctx.graph_s.is_some() {
10536 &dctx.q_slots
10537 } else {
10538 &draft_logits
10539 };
10540 let mut n_acc = 0usize;
10541 for j in 0..k_round {
10542 let (qmx, qth, qz) = draft_stats[j];
10543 let idsd = e.htod_u32_v(&[draft_idx[j]])?;
10544 let rowsd = e.htod_i32(&[0])?;
10545 let thd = e.htod(&[qth])?;
10546 let zd = e.htod(&[qz])?;
10547 let _ = qmx;
10548 let mut outd = e.zeros(1)?;
10549 e.softmax_gather_filtered(
10550 &q_bufs[j], d_vocab, &idsd, &rowsd, &thd, &zd, &mut outd, d_vocab, 1,
10551 sp_temp,
10552 )?;
10553 let qj = e.dtoh(&outd)?[0];
10554 let u = host_u01(sp_seed, uctr);
10555 uctr += 1;
10556 let accept = (u as f64) * (qj as f64) < pj[j] as f64;
10557 // SKEY PROBE: q == 0 for the token the draft actually proposed is the
10558 // exactness signature (see `skey_probe`). Impossible when the draft was
10559 // drawn from the same filtered distribution the verify reconstructs here;
10560 // `u * 0 < p` makes it an UNCONDITIONAL accept whenever p > 0.
10561 if skey_probe() && qj == 0.0 {
10562 eprintln!(
10563 "[skey] EXACTNESS q=0 round={round} j={j} draft_tok={} \
10564 draft_idx={} p={:e} u={u} accepted={} th_z={:?}",
10565 draft[j], draft_idx[j], pj[j], accept as u8, draft_stats[j],
10566 );
10567 }
10568 if accept {
10569 n_acc += 1;
10570 } else {
10571 break;
10572 }
10573 }
10574 let bonus = if n_acc == k_round {
10575 // FULL ACCEPT: bonus ~ FILTERED softmax at the last verify column.
10576 let col = base + k_round - 1;
10577 let cb = col_buf.as_mut().unwrap();
10578 e.copy_view_into(
10579 cb,
10580 0,
10581 &tlogits_d.slice(col * n_vocab..(col + 1) * n_vocab),
10582 n_vocab,
10583 )?;
10584 if pen_on {
10585 let h = pen_hist_d.as_ref().unwrap();
10586 let nh = h.len();
10587 e.penalize_logits(
10588 cb,
10589 h,
10590 nh,
10591 sp.penalty_repeat,
10592 sp.penalty_freq,
10593 sp.penalty_present,
10594 n_vocab,
10595 )?;
10596 }
10597 if perturb_buf.is_none() {
10598 perturb_buf = Some(e.zeros(d_vocab.max(n_vocab))?);
10599 }
10600 // STATS MUST COME FROM THIS COLUMN (bug fix 2026-08-05, lane/sampler-
10601 // truncation-fix; receipts research/sampfix-20260805/). The old code reused
10602 // `col_stats.last()` here, which is ALWAYS the wrong row: the gathered set
10603 // covers verify columns 0..=(base+k_round-2) (rows pushed as base+j-1), while
10604 // the full-accept bonus samples column base+k_round-1 — exactly ONE PAST the
10605 // last gathered column, in both base arms. `th` is a threshold in e-units of
10606 // its OWN row's max, so feeding a neighbour's (row_max, th) into
10607 // gumbel_perturb_filtered mis-scales every e0 = exp((x-row_max)/T). When the
10608 // donor column's peak is higher by more than T*ln(1/th), EVERY id fails
10609 // `e0 >= th`, the whole perturbed row becomes -3.4e38, and the 2-pass argmax
10610 // falls through to its smallest-index tie-break => token id 0 ("!") spliced
10611 // mid-word. Fragility is ordered by how large th is: min_p pins th = min_p
10612 // (0.05 => trigger at delta > 2.4 at T=0.8, fires constantly), top_p's
10613 // mass-boundary th is smaller, top_k's k-th-largest th smaller still — which
10614 // is why the head-to-head matrix saw min_p and top_p corrupt while top_k-only
10615 // stayed clean. The pure-temp default regime is immune (th == 0 masks nothing,
10616 // and row_max is unused once nothing is masked), so this fix is a byte-level
10617 // no-op for the untruncated serve default. One extra one-block filter_stats
10618 // per full-accept round is the whole cost.
10619 let (mx, th) = {
10620 let rows0 = e.htod_i32(&[0])?;
10621 let (mut th_d, mut z_d, mut mx_d) = (e.zeros(1)?, e.zeros(1)?, e.zeros(1)?);
10622 let cb0 = col_buf.as_ref().unwrap();
10623 e.filter_stats(
10624 cb0, n_vocab, &rows0, &mut th_d, &mut z_d, &mut mx_d, n_vocab, 1,
10625 sp_temp, sp.top_k, sp.top_p, sp.min_p,
10626 )?;
10627 (e.dtoh(&mx_d)?[0], e.dtoh(&th_d)?[0])
10628 };
10629 let pb = perturb_buf.as_mut().unwrap();
10630 let cb2 = col_buf.as_ref().unwrap();
10631 e.gumbel_perturb_filtered(cb2, pb, n_vocab, sp_seed, sctr, sp_temp, mx, th)?;
10632 sctr += 1;
10633 let td = e.argmax_token_device(pb, n_vocab)?;
10634 e.dtoh_u32_one(&td)?
10635 } else {
10636 // REJECT at n_acc: bonus ~ norm(max(0, softmax_T(p) - softmax_T(q))).
10637 let cb = col_buf.as_mut().unwrap();
10638 if n_acc > 0 || base == 1 {
10639 let col = base + n_acc - 1;
10640 e.copy_view_into(
10641 cb,
10642 0,
10643 &tlogits_d.slice(col * n_vocab..(col + 1) * n_vocab),
10644 n_vocab,
10645 )?;
10646 } else {
10647 let lc = last_col_logits.as_ref().unwrap();
10648 e.copy_into(cb, 0, lc, n_vocab)?;
10649 }
10650 if pen_on {
10651 let h = pen_hist_d.as_ref().unwrap();
10652 let nh = h.len();
10653 e.penalize_logits(
10654 cb,
10655 h,
10656 nh,
10657 sp.penalty_repeat,
10658 sp.penalty_freq,
10659 sp.penalty_present,
10660 n_vocab,
10661 )?;
10662 }
10663 let cb2 = col_buf.as_ref().unwrap();
10664 let sc = sctr;
10665 sctr += 1;
10666 // p-stats for the reject column: from col_stats when the col was gathered,
10667 // else (j==0&&base==0) from last_col_stats.
10668 let p_stats = if n_acc > 0 || base == 1 {
10669 // col index within the gathered set == number of gathered cols before n_acc
10670 let gi = if base == 1 { n_acc } else { n_acc - 1 };
10671 col_stats.get(gi).copied().unwrap_or_else(|| {
10672 (0.0, 0.0, 1.0) // unreachable: gathered cols always cover the reject slot
10673 })
10674 } else {
10675 last_col_stats.expect("sampled: last_col_stats unset at reject")
10676 };
10677 let q_stats = draft_stats[n_acc];
10678 if let Some(map) = &d2t_dev {
10679 if q_full_buf.is_none() {
10680 q_full_buf = Some(e.zeros(n_vocab)?);
10681 }
10682 let qf = q_full_buf.as_mut().unwrap();
10683 e.scatter_trim_logits(&q_bufs[n_acc], map, qf, d_vocab, n_vocab)?;
10684 let qf2 = q_full_buf.as_ref().unwrap();
10685 e.residual_sample_filtered(
10686 cb2,
10687 Some(qf2),
10688 n_vocab,
10689 sp_temp,
10690 sp_seed,
10691 sc,
10692 p_stats,
10693 q_stats,
10694 &mut sample_tok,
10695 )?;
10696 } else {
10697 e.residual_sample_filtered(
10698 cb2,
10699 Some(&q_bufs[n_acc]),
10700 n_vocab,
10701 sp_temp,
10702 sp_seed,
10703 sc,
10704 p_stats,
10705 q_stats,
10706 &mut sample_tok,
10707 )?;
10708 }
10709 e.dtoh_u32(&sample_tok)?[0]
10710 };
10711 (n_acc, bonus)
10712 };
10713 // --- 3b. GRAMMAR TRUNCATION (constrained spec, 2026-08-03): the grammar is
10714 // an extra rejection rule AFTER the exactness verify (the batched-verify-twins
10715 // ordering). Walk the accepted drafts through the grammar in commit order; the
10716 // first illegal token truncates acceptance at its slot, and that slot's emission
10717 // is recomputed as the MASKED argmax of the target's own verify column — token-
10718 // identical to constrained plain greedy decode (an unmasked argmax that is
10719 // grammar-legal IS the masked argmax: masking only removes competitors). The
10720 // column D2H (~1MB) is paid only when a cut fires — the tight-grammar cost,
10721 // measured in acceptance numbers, never hidden.
10722 let (n_acc, bonus) = match constraint.as_deref_mut() {
10723 None => (n_acc, bonus),
10724 Some(c) => {
10725 fn ce(e2: String) -> Box<dyn std::error::Error> {
10726 format!("constraint: {e2}").into()
10727 }
10728 let mut na = n_acc;
10729 let mut cut = false;
10730 for (j, &d) in draft.iter().enumerate().take(n_acc) {
10731 if c.is_allowed(d).map_err(ce)? {
10732 c.consume(d).map_err(ce)?;
10733 } else {
10734 na = j;
10735 cut = true;
10736 dm_cut_tokens += n_acc - j;
10737 break;
10738 }
10739 }
10740 if cut {
10741 dm_cuts += 1;
10742 }
10743 let mut bo = bonus;
10744 if cut || !c.is_allowed(bo).map_err(ce)? {
10745 let mut row = if na == 0 && base == 0 {
10746 init_logits_host
10747 .clone()
10748 .ok_or("constraint: init logits missing (round-0 cut)")?
10749 } else {
10750 e.dtoh_view(
10751 &tlogits_d.slice((base + na - 1) * n_vocab..(base + na) * n_vocab),
10752 )?
10753 };
10754 c.mask_logits(&mut row).map_err(ce)?;
10755 bo = argmax(&row) as u32;
10756 }
10757 c.consume(bo).map_err(ce)?;
10758 (na, bo)
10759 }
10760 };
10761 let mut successor_valid = false;
10762 if let Some((q_proxy, expected_d2)) = rejected_probe {
10763 let v_n = n_acc == 1 && bonus == expected_d2;
10764 eprintln!(
10765 "[opti-controller] shadow q={q_proxy:.6} admitted=false v_n={v_n} \
10766 expected_d2={expected_d2} n_acc={n_acc} bonus={bonus}",
10767 );
10768 }
10769 if let Some(successor) = successor_attempt.as_ref() {
10770 successor_valid = n_acc == 1 && bonus == successor.verify_tokens[0];
10771 let generation = successor.generation;
10772 let q_proxy = successor.q_proxy;
10773 let expected_pending = successor.verify_tokens[0];
10774 let resolution_ms = successor.issued_at.elapsed().as_secs_f64() * 1e3;
10775 let fork = opti_fork
10776 .as_mut()
10777 .ok_or("optipipe successor resolution lost fork state")?;
10778 fork.finish_actual_reconcile(e, &mut *cache, &snap, n_acc, base, successor_valid)?;
10779 if successor_valid {
10780 OPTI_FORK_HITS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
10781 } else {
10782 OPTI_FORK_MISSES.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
10783 OPTI_RECONCILES.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
10784 OPTI_WASTED_DRAFT_TOKENS.fetch_add(2, std::sync::atomic::Ordering::Relaxed);
10785 }
10786 let breaker_tripped = fork
10787 .controller
10788 .as_mut()
10789 .expect("controller policy")
10790 .resolve(successor_valid);
10791 if breaker_tripped {
10792 OPTI_BREAKER_TRIPS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
10793 }
10794 eprintln!(
10795 "[opti-controller] resolve generation={} hit={} q={q_proxy:.6} \
10796 expected_pending={expected_pending} n_acc={n_acc} bonus={bonus} \
10797 resolution_ms={resolution_ms:.3} reconcile={} breaker={}",
10798 generation.id, successor_valid, !successor_valid, breaker_tripped,
10799 );
10800 if !successor_valid {
10801 let mut successor = successor_attempt
10802 .take()
10803 .expect("controller successor disappeared on miss");
10804 successor.settle();
10805 fork.retire(generation)?;
10806 }
10807 }
10808 total_drafted += k_round;
10809 total_accepted += n_acc;
10810 if let Some(t) = sess_telem {
10811 // Greedy, rejection-sampling, and grammar truncation all converge here after
10812 // the accept decision is already on host. Fixed-size relaxed atomics only.
10813 t.record_round(k_round, n_acc);
10814 }
10815 if spec_stats {
10816 st_len_hist[k_round] += 1;
10817 for j in 0..k_round {
10818 st_drafted[j] += 1;
10819 }
10820 for j in 0..n_acc {
10821 st_accepted[j] += 1;
10822 }
10823 if n_acc == k_round {
10824 st_full += 1;
10825 }
10826 }
10827
10828 if debug_spec {
10829 eprintln!(
10830 "[R{round}] pos={pos} out_len={} last_tok={last_token} draft={draft:?} n_acc={n_acc} bonus={bonus} t_pred0={}",
10831 out.len(),
10832 // NOT `t_pred(0)`: `preds` is filled only under `if !sampled` above, so on a
10833 // sampled request round >= 1 (base == 1) indexed an EMPTY vector and PANICKED
10834 // the GPU worker thread — a debug flag that killed the exact regime you would
10835 // set it to investigate. See `debug_t_pred0`.
10836 debug_t_pred0(sampled, base, last_pred, &preds)
10837 );
10838 }
10839
10840 // --- 4. COMMIT: draft[0..n_acc] then bonus (n_acc + 1 tokens) ---
10841 let commit_started = std::time::Instant::now();
10842 // SESSION MODE: every accepted column is already in the CACHE — `out` must carry all
10843 // of them (overshoot past max_new included) or `committed` under-counts the cache rows
10844 // and the next turn's continuation seeds one token off (gate-caught 2026-07-05). The
10845 // single-shot path keeps the cap (its caller truncates + drops the cache anyway).
10846 for j in 0..n_acc {
10847 if !session_mode && out.len() >= max_new {
10848 break;
10849 }
10850 out.push(draft[j]);
10851 }
10852 if pen_on {
10853 pen_hist.extend_from_slice(&draft[0..n_acc]);
10854 pen_hist.push(bonus);
10855 }
10856 let bonus_emitted = session_mode || out.len() < max_new;
10857 if bonus_emitted {
10858 out.push(bonus);
10859 }
10860 last_token = bonus;
10861
10862 // --- 5. ROLLBACK + advance (§C) ---
10863 if n_acc == k_round && !spec_replay {
10864 // FULL ACCEPT, BONUS FOLD: all verify columns (pending? + drafts) are committed in
10865 // cache; the NEW bonus stays PENDING for the next round's verify batch — NO extra
10866 // T=1 trunk pass. The next draft chain seeds from the MTP block's h_nextn at the
10867 // bonus position: one MTP-block pass (~1/33 trunk cost) replaces the trunk read.
10868 // last_pred is dead in the pending path (t_pred reads verify col 0).
10869 //
10870 // PERSISTENT DRAFT KV, full-accept fill: the chain covered last_token +
10871 // draft[0..k_round-2] as INPUTS (slots P..P'-2); draft[k_round-1] (slot P'-1) was
10872 // only ever an output, so its entry is MISSING. Fill it from vh_seed — its EXACT
10873 // trunk hidden (the last verify column). set_len first: a p-min break may have
10874 // left one extra chain append at that slot. Partial accepts need NO fill (the
10875 // chain already covered every accepted position; round-start set_len truncates).
10876 let mut vh_seed = e.zeros(n_embd)?;
10877 e.copy_view_into(
10878 &mut vh_seed,
10879 0,
10880 &vx.slice((t_v - 1) * n_embd..t_v * n_embd),
10881 n_embd,
10882 )?;
10883 if refresh {
10884 // TRUE-HIDDEN REFRESH (2026-07-03, the HANDOVER-listed acceptance lever):
10885 // overwrite ALL committed positions' scratch entries with K/V from their EXACT
10886 // verify hiddens — the reference engine's mtp_update fills from true hiddens;
10887 // the full stack (vx) is already resident from the verify. Replaces both the
10888 // chain-approximate entries AND the old last-token-only fill. Acceptance-only
10889 // (draft attention quality); exactness stays the verify's job.
10890 scratch.set_len(e, pos)?;
10891 // PREDECESSOR pairing: row i gets vx[i-1]; row 0 the carried fill_prev
10892 // (hidden of the last committed row before this verify batch).
10893 let mut vxs = e.zeros(t_v * n_embd)?;
10894 e.copy_into(&mut vxs, 0, &fill_prev, n_embd)?;
10895 if t_v > 1 {
10896 e.copy_view_into(
10897 &mut vxs,
10898 n_embd,
10899 &vx.slice(0..(t_v - 1) * n_embd),
10900 (t_v - 1) * n_embd,
10901 )?;
10902 }
10903 self.mtp_kv_fill(e, mtp, &verify_tokens, &vxs, pos, &mut *scratch, embd_dev)?;
10904 } else {
10905 scratch.set_len(e, pos + base + k_round - 1)?;
10906 // predecessor of the last draft = verify col t_v-2 (or fill_prev at t_v==1)
10907 let mut hp = e.zeros(n_embd)?;
10908 if t_v >= 2 {
10909 e.copy_view_into(
10910 &mut hp,
10911 0,
10912 &vx.slice((t_v - 2) * n_embd..(t_v - 1) * n_embd),
10913 n_embd,
10914 )?;
10915 } else {
10916 e.copy_into(&mut hp, 0, &fill_prev, n_embd)?;
10917 }
10918 self.mtp_kv_fill(
10919 e,
10920 mtp,
10921 &[draft[k_round - 1]],
10922 &hp,
10923 pos + base + k_round - 1,
10924 &mut *scratch,
10925 embd_dev,
10926 )?;
10927 }
10928 // REFERENCE SEEDING: no pseudo pass — the next chain's step 0 IS the
10929 // reference's (id_last, h_prev) draft row; it appends the bonus's scratch
10930 // entry itself. Seed = TRUE hidden of the bonus's predecessor (last verify
10931 // col). Saves one MTP-block pass per round on top of the pairing fix.
10932 if !devacc_seeded {
10933 e.copy_into(&mut h_seed_buf, 0, &vh_seed, n_embd)?;
10934 e.copy_into(&mut fill_prev, 0, &vh_seed, n_embd)?;
10935 }
10936 pending = Some(bonus);
10937 if debug_spec {
10938 eprintln!(" -> FULL ACCEPT (bonus pending, prev-h seed)");
10939 }
10940 } else if !spec_replay && base + n_acc >= 1 {
10941 // PARTIAL ACCEPT, REPLAY-FREE (2026-07-03 — the profiled #1 long-ctx spec cost):
10942 // the verify's first j = base+n_acc columns ARE the committed sequence, computed
10943 // bit-identically to eager (decode-exact contract) — so KEEP them: KV truncates to
10944 // pos+j, recurrent state rebuilds from the VerifyCkpt (same-kernel gdn prefix
10945 // re-run / pure state-clone restore), and the bonus stays PENDING exactly like the
10946 // full-accept path — the legacy duplicate trunk replay is gone. The next chain
10947 // seeds from the MTP pseudo-hidden of the bonus, whose seed = the TRUE verify
10948 // hidden of its predecessor (col j-1) — same one-hop pseudo structure as full
10949 // accept (never compounds: the next verify recomputes true hiddens for all
10950 // committed columns).
10951 let j = base + n_acc;
10952 self.commit_verified_prefix(
10953 e,
10954 &mut *cache,
10955 &snap,
10956 ckpt.as_ref().unwrap(),
10957 j,
10958 devacc_seeded,
10959 if devacc_seeded {
10960 devacc_acc.as_ref().map(|a| (a, base, t_v))
10961 } else {
10962 None
10963 },
10964 )?;
10965 let mut seed = e.zeros(n_embd)?;
10966 e.copy_view_into(
10967 &mut seed,
10968 0,
10969 &vx.slice((j - 1) * n_embd..j * n_embd),
10970 n_embd,
10971 )?;
10972 // Draft scratch: TRUE-HIDDEN REFRESH of the committed prefix (see the full-accept
10973 // branch); without it the chain entries stand and only the tail truncates. Either
10974 // way len ends at pos+j so the pseudo append lands at the bonus's slot pos+j
10975 // (persistent mode), rope pos+j+1 (chain convention).
10976 if refresh {
10977 scratch.set_len(e, pos)?;
10978 let mut vxs = e.zeros(j * n_embd)?;
10979 e.copy_into(&mut vxs, 0, &fill_prev, n_embd)?;
10980 if j > 1 {
10981 e.copy_view_into(
10982 &mut vxs,
10983 n_embd,
10984 &vx.slice(0..(j - 1) * n_embd),
10985 (j - 1) * n_embd,
10986 )?;
10987 }
10988 self.mtp_kv_fill(
10989 e,
10990 mtp,
10991 &verify_tokens[0..j],
10992 &vxs,
10993 pos,
10994 &mut *scratch,
10995 embd_dev,
10996 )?;
10997 } else {
10998 scratch.set_len(e, pos + j)?;
10999 }
11000 // REFERENCE SEEDING (see the full-accept branch): seed = TRUE hidden of the
11001 // bonus's predecessor (verify col j-1); no pseudo pass.
11002 if !devacc_seeded {
11003 e.copy_into(&mut h_seed_buf, 0, &seed, n_embd)?;
11004 e.copy_into(&mut fill_prev, 0, &seed, n_embd)?;
11005 }
11006 pending = Some(bonus);
11007 if debug_spec {
11008 eprintln!(" -> PARTIAL(replay-free j={j}, bonus pending, prev-h seed)");
11009 }
11010 } else if !spec_replay {
11011 // ZERO ROUND FOLD (2026-07-10, verify-cost target #3): base+n_acc == 0 — a
11012 // pending-less round where nothing was accepted (PMIN0 zero-draft chains after a
11013 // replay/commit, or plain 0-accept rounds at round 0). The old path replayed
11014 // [bonus] through a FULL m=1 trunk+head forward (the 489us full-vocab head pass
11015 // measured at ~0.75/round on PMIN0 configs). Instead: restore the pre-round
11016 // snapshot and let the bonus ride the NEXT round's verify as col 0 — the existing
11017 // base=1 pending machinery, bit-identical by the decode-exact verify contract.
11018 // Seed: the bonus's predecessor is the last COMMITTED token, whose hidden
11019 // fill_prev already carries (same seeding as the 1-token-replay case it replaces).
11020 cache.rollback(e, &snap, 0)?;
11021 scratch.set_len(e, pos)?;
11022 e.copy_into(&mut h_seed_buf, 0, &fill_prev, n_embd)?;
11023 pending = Some(bonus);
11024 if debug_spec {
11025 eprintln!(" -> ZERO-ROUND FOLD (bonus pending, fill_prev seed)");
11026 }
11027 } else {
11028 // PARTIAL ACCEPT, LEGACY REPLAY (seam MEMRA_SPEC_REPLAY=1 — or j==0: nothing of
11029 // this round survives, only possible before the first pending exists, ~round 0):
11030 // restore EVERYTHING to the pre-round snapshot (KV truncate to pos + recur
11031 // restore), then replay the committed prefix pending? ++ draft[0..n_acc] ++
11032 // [bonus] as ONE batched T forward — single weight read, bit-identical to greedy
11033 // (the verify-all-columns path is the same math). Commits the bonus with a TRUE
11034 // trunk hidden.
11035 cache.rollback(e, &snap, 0)?; // accept_len=0: KV len = pos, recur = snapshot
11036 let mut replay: Vec<u32> = Vec::with_capacity(base + n_acc + 1);
11037 if let Some(b) = pending.take() {
11038 replay.push(b);
11039 }
11040 replay.extend_from_slice(&draft[0..n_acc]);
11041 replay.push(bonus);
11042 // Full-stack forward (decode_step_t_core = decode_step_t_h_emb_dev's body):
11043 // Predecessor pairing seeds from the PREDECESSOR row (col len-2) — the same-row path takes the
11044 // last col exactly as before (byte-identical to the old _h_emb_dev call).
11045 let (rl_d, rx) = if self.qwen35_serving_class() {
11046 let mut logits = Vec::with_capacity(replay.len() * n_vocab);
11047 let mut hidden = e.uninit(replay.len() * n_embd)?;
11048 for (row, &token) in replay.iter().enumerate() {
11049 let (row_logits, row_hidden) =
11050 self.spec_target_step_h(e, token, &mut *cache)?;
11051 logits.extend_from_slice(&row_logits);
11052 e.dtod_copy_into(&row_hidden, &mut hidden, row * n_embd)?;
11053 }
11054 (e.htod(&logits)?, hidden)
11055 } else {
11056 self.decode_step_t_core(e, &replay, pos, &mut *cache, embd_dev, None)?
11057 };
11058 // last_pred = argmax of the LAST column's logits (predicts the token after `bonus`)
11059 // — device argmax + one 4-byte read instead of the full-vocab column dtoh.
11060 e.argmax_token_device_col(&rl_d, replay.len() - 1, n_vocab, &mut preds_d, 0)?;
11061 last_pred = e.dtoh_u32(&preds_d)?[0];
11062 if sampled {
11063 let lr0 = replay.len();
11064 let lc = last_col_logits
11065 .as_mut()
11066 .expect("sampled: last_col_logits unset");
11067 e.copy_view_into(
11068 lc,
11069 0,
11070 &rl_d.slice((lr0 - 1) * n_vocab..lr0 * n_vocab),
11071 n_vocab,
11072 )?;
11073 }
11074 let lr = replay.len();
11075 if lr >= 2 {
11076 e.copy_view_into(
11077 &mut h_seed_buf,
11078 0,
11079 &rx.slice((lr - 2) * n_embd..(lr - 1) * n_embd),
11080 n_embd,
11081 )?;
11082 } else {
11083 // 1-token replay (round-0 miss): the bonus's predecessor is the OLD
11084 // last_token, whose own-row hidden fill_prev still holds.
11085 e.copy_into(&mut h_seed_buf, 0, &fill_prev, n_embd)?;
11086 }
11087 // the bonus is COMMITTED here — it becomes the last committed row.
11088 let mut rh_last = e.zeros(n_embd)?;
11089 e.copy_view_into(
11090 &mut rh_last,
11091 0,
11092 &rx.slice((lr - 1) * n_embd..lr * n_embd),
11093 n_embd,
11094 )?;
11095 e.copy_into(&mut fill_prev, 0, &rh_last, n_embd)?;
11096 if debug_spec {
11097 eprintln!(" -> PARTIAL(replay={replay:?}), next_pred={last_pred}");
11098 }
11099 }
11100 if devacc_seeded {
11101 // stage (b) epilogue: fill_prev takes the gathered seed AFTER the refresh fills
11102 // consumed the old value (both slots carry the same value in every non-replay arm).
11103 e.copy_into(&mut fill_prev, 0, &h_seed_buf, n_embd)?;
11104 }
11105 if successor_valid {
11106 let optimistic_scratch_len = successor_attempt
11107 .as_ref()
11108 .expect("valid controller successor disappeared")
11109 .scratch_len;
11110 // The normal current-round commit refreshed/truncated the logical scratch tail.
11111 // Its optimistic successor row was already written physically, so restoring only
11112 // the retained logical length makes that row live for the carried round.
11113 scratch.set_len(e, optimistic_scratch_len)?;
11114 }
11115 if let Some(current) = current_opti.take() {
11116 opti_fork
11117 .as_mut()
11118 .ok_or("optipipe current retirement lost fork state")?
11119 .retire(current.generation)?;
11120 }
11121 if successor_valid {
11122 let successor = successor_attempt
11123 .take()
11124 .expect("valid controller successor disappeared before promotion");
11125 let generation = successor.generation;
11126 opti_fork
11127 .as_mut()
11128 .ok_or("optipipe successor promotion lost fork state")?
11129 .promote_successor_snapshot(&mut snap, generation);
11130 carried_opti = Some(successor);
11131 }
11132 if anatomy_on {
11133 // Commit/rollback is normally asynchronous on the primary/head stream. Bound it
11134 // only for this diagnostic so it does not disappear into the following draft's
11135 // first token readback.
11136 e.stream().synchronize()?;
11137 ph_commit += commit_started.elapsed().as_secs_f64();
11138 }
11139 // adaptive-K update (host math, zero syncs): next round drafts accepted-run + 1,
11140 // clamped to [floor(pos), k_cap]. cache.pos is post-rollback here (the round's
11141 // final position — the floor's position key reads the committed depth). Burst
11142 // rounds (`continue` above) draft the captured fixed depth and skip this, exactly
11143 // like gemma's burst arm.
11144 if adapt {
11145 let fl_now = floor_at(cache.pos);
11146 kc = (n_acc + 1).clamp(fl_now.min(k_cap), k_cap);
11147 }
11148 ph_mark(&mut ph_rest, phase_on);
11149 if let Some(p) = pipe {
11150 p.accept_end(round);
11151 }
11152 drop(pipe_accept);
11153 round += 1;
11154 // sse-cadence: this round's accepted drafts + bonus are committed (out is
11155 // append-only past step 4) — flush at round cadence.
11156 keep_going = flush_commit(&mut on_commit, &out, &mut flushed);
11157 }
11158 if let Some(mut ticket) = carried_opti.take() {
11159 opti_fork
11160 .as_mut()
11161 .ok_or("optipipe tail drain lost fork state")?
11162 .cancel_controller_ticket(e, &mut *cache, &mut *scratch, &snap, &mut ticket)?;
11163 }
11164 // sse-cadence: nothing below appends to `out`; flush any remainder (defensive).
11165 // (verdict ignored — the burst is over either way; the session tail runs unchanged.)
11166 let _ = flush_commit(&mut on_commit, &out, &mut flushed);
11167
11168 if spec_stats {
11169 let per_slot: Vec<String> = (0..k)
11170 .map(|j| {
11171 if st_drafted[j] > 0 {
11172 format!(
11173 "{}/{}={:.3}",
11174 st_accepted[j],
11175 st_drafted[j],
11176 st_accepted[j] as f64 / st_drafted[j] as f64
11177 )
11178 } else {
11179 "0/0".into()
11180 }
11181 })
11182 .collect();
11183 let acc = if total_drafted > 0 {
11184 total_accepted as f64 / total_drafted as f64
11185 } else {
11186 0.0
11187 };
11188 eprintln!(
11189 "[spec-stats] rounds={round} full_accept={st_full} len_hist={st_len_hist:?} \
11190 per_slot=[{}] total={total_accepted}/{total_drafted}={acc:.3} \
11191 tok_per_round={:.3}",
11192 per_slot.join(" "),
11193 (total_accepted + round) as f64 / round.max(1) as f64
11194 );
11195 }
11196 if constraint.is_some() {
11197 eprintln!(
11198 "[draft-mask] mask_rounds={dm_rounds} clone_total={:.3}ms \
11199 clone_per_round={:.4}ms gram_cuts={dm_cuts}/{round} cut_tokens={dm_cut_tokens}",
11200 dm_clone_ns as f64 / 1e6,
11201 dm_clone_ns as f64 / 1e6 / dm_rounds.max(1) as f64
11202 );
11203 }
11204 if phase_on {
11205 let tot = ph_draft + ph_verify + ph_wait + ph_rest;
11206 eprintln!(
11207 "[spec-phase] draft={:.1}ms ({:.1}%) verify-issue={:.1}ms ({:.1}%) verify-wait={:.1}ms ({:.1}%) commit-host={:.1}ms ({:.1}%) rounds={round}",
11208 ph_draft * 1e3,
11209 ph_draft / tot * 100.0,
11210 ph_verify * 1e3,
11211 ph_verify / tot * 100.0,
11212 ph_wait * 1e3,
11213 ph_wait / tot * 100.0,
11214 ph_rest * 1e3,
11215 ph_rest / tot * 100.0
11216 );
11217 }
11218 if anatomy_on {
11219 let rounds_f = round.max(1) as f64;
11220 let other = (ph_rest - ph_commit).max(0.0);
11221 eprintln!(
11222 "[spec-anatomy] per-round draft={:.3}ms pp-verify={:.3}ms \
11223 verify-accept={:.3}ms commit-rollback={:.3}ms other={:.3}ms rounds={round}",
11224 ph_draft * 1e3 / rounds_f,
11225 ph_verify * 1e3 / rounds_f,
11226 ph_wait * 1e3 / rounds_f,
11227 ph_commit * 1e3 / rounds_f,
11228 other * 1e3 / rounds_f,
11229 );
11230 }
11231 let _pipe_tail = pipe.map(|p| p.primary());
11232 // SESSION TAIL: leave the session in the exact invariant the next turn's suffix prime
11233 // expects — every row in `committed` has trunk KV/recur state AND an exact draft-KV row.
11234 // Park the draft-graph ctx back on the session (the serve-burst fixed-cost fix): the next
11235 // burst replays instead of recapturing. Error paths (`?` above) drop it — recaptured then.
11236 if let Some(slot) = sess_draft_slot.take() {
11237 *slot = Some(dctx);
11238 }
11239 let t_rounds = t_ent.elapsed();
11240 if let Some((committed, last_h, next_pred_slot, sctr_slot, uctr_slot)) = sess_tail.take() {
11241 // NEXT BURST'S BOUNDARY TOKEN (lane/sampled-spec-quality, Item 1). Greedy stashes
11242 // the argmax `last_pred` exactly as before (byte contract). SAMPLED draws the token
11243 // HERE, where the sampler, the session Philox counters and the penalty window are
11244 // all live and the boundary logits row still exists — that is the "make the state
11245 // available" half of the fix; the consuming burst then just emits it. `sctr` is
11246 // written to the session BELOW the draws so the advance is never lost.
11247 *next_pred_slot = Some(last_pred);
11248 let sample_boundary = sampled && constraint.is_none() && spec_sampled_boundary_on();
11249 let mut stashed_pending = false;
11250 if let Some(b) = pending.take() {
11251 if !sampled {
11252 // PENDING-CARRY (2026-08-01): stash the bonus on the session instead of
11253 // committing it with a solo T=1 pass — the next empty-suffix greedy burst
11254 // consumes it as round-0 verify col 0 (a plain round edge; the old tail
11255 // commit + next burst's init feed were 11.6+11.5ms solo trunk passes per
11256 // burst on H100 q27, [spec-setup] trace). b stays in `out` (emitted) but
11257 // OUT of `committed` (cache rows == committed); the consuming call
11258 // prepends it once its verify commits the row. next_pred is unknowable
11259 // without the commit pass — None; callers gate on pending_tok too.
11260 debug_assert_eq!(out.last(), Some(&b), "pending must be the last emitted");
11261 if let Some(slot) = sess_pending_slot.take() {
11262 *slot = Some(b);
11263 }
11264 *next_pred_slot = None;
11265 // fill_prev = hidden of the last COMMITTED row (b's predecessor) — the
11266 // exact chain-seed/fill anchor the consuming burst (or a flush) needs.
11267 *last_h = Some(e.clone_dtod(&fill_prev)?);
11268 stashed_pending = true;
11269 } else {
11270 // SAMPLED tail (unchanged): commit the bonus (one T=1 pass) + draft fill —
11271 // the sampled round-0 accept needs this pass's logits (last_col_logits).
11272 let pos_b = cache.pos;
11273 scratch.set_len(e, pos_b)?;
11274 let (lg_b, hb) = self.spec_target_step_h(e, b, &mut *cache)?;
11275 // after a FULL-accept exit `last_pred` is STALE (it predicted the bonus
11276 // itself — the prediction AFTER the bonus never materialized; it would have
11277 // been the next round's verify col 0). The commit's logits ARE that
11278 // prediction — so they are also the row the next burst's boundary token
11279 // comes off, and (lane/sampled-spec-quality) it is DRAWN from them here.
11280 *next_pred_slot = Some(if sample_boundary {
11281 sample_boundary_token(
11282 e,
11283 &lg_b,
11284 &sp,
11285 &pen_hist,
11286 &mut sctr,
11287 "burst-tail-commit",
11288 )?
11289 } else {
11290 argmax(&lg_b) as u32
11291 });
11292 self.mtp_kv_fill(e, mtp, &[b], &fill_prev, pos_b, &mut *scratch, embd_dev)?;
11293 *last_h = Some(hb);
11294 }
11295 } else {
11296 // fill_prev tracks the hidden of the last COMMITTED row throughout the loop.
11297 *last_h = Some(e.clone_dtod(&fill_prev)?);
11298 if sample_boundary {
11299 // No pending to commit, so the boundary row is the one `last_pred` was
11300 // argmaxed from and the sampled path keeps it on device: the init feed's
11301 // logits when the burst ran zero rounds, else the legacy-replay path's
11302 // last verify column (both predict the token AFTER the last committed
11303 // row). It is retained precisely because round 0's accept test needs it,
11304 // so the draw costs no extra D2H of the [n_vocab] row.
11305 match last_col_logits.as_ref() {
11306 Some(lc) => {
11307 *next_pred_slot = Some(sample_boundary_token_dev(
11308 e,
11309 lc,
11310 n_vocab,
11311 &sp,
11312 &pen_hist,
11313 &mut sctr,
11314 "burst-tail-nopending",
11315 )?);
11316 }
11317 // NAME THE FALLBACK (house standard): unreachable today — a sampled
11318 // burst always feeds or replays, so the row exists — but if it ever
11319 // is, the stream takes a greedy token and SAYS so rather than
11320 // silently regressing to the pre-lane behaviour.
11321 None => eprintln!(
11322 "[spec-boundary] sampled tail kept the ARGMAX boundary token \
11323 (reason: no retained boundary logits row)"
11324 ),
11325 }
11326 }
11327 }
11328 *sctr_slot = sctr;
11329 *uctr_slot = uctr;
11330 committed.extend_from_slice(prompt);
11331 if let Some(cb) = carried_pending {
11332 // the consumed carry's cache row landed in round 0's verify (every pending
11333 // round commits col 0) — it joins `committed` here, in sequence order.
11334 committed.push(cb);
11335 }
11336 if stashed_pending {
11337 // ZERO-EMIT BURST (2026-08-06 c=8 serve panic, pre-existing since b4aea184):
11338 // `out.len() - 1` underflowed on an EMPTY `out` — "range end index
11339 // 18446744073709551615 out of range for slice of length 0", killing the
11340 // memra-gpu-worker and failing 31 of 32 concurrent requests with "worker closed
11341 // stream". Reachable because `pending` starts as `carried_pending` (a bonus
11342 // stashed by the PREVIOUS burst) while `out` starts empty, and the carry is
11343 // deliberately NOT pushed to `out` (line ~3239: the burst that emitted it already
11344 // did). So a burst that stashes a pending without emitting anything of its own —
11345 // the round loop exits before a push, e.g. the ring drain's `out.len() < max_new`
11346 // guard skipping every token under a tight budget — arrives here with
11347 // out.len() == 0 and stashed_pending == true.
11348 //
11349 // The invariant is unchanged: `committed` gets every emitted token EXCEPT the
11350 // stashed bonus. With nothing emitted, that is nothing — and the carry pushed
11351 // just above is already accounted. Saturating, not a min/assert: an empty `out`
11352 // here is a legitimate burst shape, not a corrupt state.
11353 let emitted = out.len().saturating_sub(1);
11354 committed.extend_from_slice(&out[..emitted]);
11355 } else {
11356 committed.extend_from_slice(&out); // FULL out incl. overshoot — all committed
11357 }
11358 debug_assert_eq!(
11359 cache.pos,
11360 committed.len(),
11361 "session invariant: cache rows == committed tokens"
11362 );
11363 if setup_trace {
11364 e.stream().synchronize()?; // bound the async tail fill in the trace
11365 let t_tail = t_ent.elapsed();
11366 eprintln!(
11367 "[spec-setup] init={:.2}ms cap={:.2}ms fill={:.2}ms rounds={:.2}ms tail={:.2}ms total={:.2}ms out={} cont={}",
11368 t_init.as_secs_f64() * 1e3,
11369 (t_cap - t_init).as_secs_f64() * 1e3,
11370 (t_fill - t_cap).as_secs_f64() * 1e3,
11371 (t_rounds - t_fill).as_secs_f64() * 1e3,
11372 (t_tail - t_rounds).as_secs_f64() * 1e3,
11373 t_tail.as_secs_f64() * 1e3,
11374 out.len(),
11375 continuation
11376 );
11377 }
11378 return Ok((out, total_drafted, total_accepted));
11379 }
11380 out.truncate(max_new);
11381 Ok((out, total_drafted, total_accepted))
11382 }
11383
11384 /// Anchor-bounded DSpark target extraction. The trunk sees the exact generated token tape;
11385 /// only requested hidden rows and target-logit rows cross PCIe. An anchor token at p pairs
11386 /// with the pre-output-norm h[p-1] carrier, exactly as the existing replay/NextN path does.
11387 pub fn extract_dspark_anchors(
11388 &self,
11389 e: &Engine,
11390 tokens: &[u32],
11391 anchor_positions: &[usize],
11392 gamma: usize,
11393 top_k: usize,
11394 chunk: usize,
11395 temperature: f32,
11396 ) -> Result<Vec<DsparkAnchorRecord>, Box<dyn std::error::Error>> {
11397 if tokens.len() < gamma + 2 || gamma == 0 || chunk < 2 {
11398 return Err("DSpark extraction token tape/gamma/chunk is invalid".into());
11399 }
11400 if anchor_positions.windows(2).any(|pair| pair[0] >= pair[1]) {
11401 return Err("DSpark anchor positions must be sorted and unique".into());
11402 }
11403 for &position in anchor_positions {
11404 if position == 0 || position + gamma >= tokens.len() {
11405 return Err(format!(
11406 "DSpark anchor {position} has no predecessor or cannot cover gamma={gamma} in {} tokens",
11407 tokens.len()
11408 )
11409 .into());
11410 }
11411 }
11412
11413 let n_vocab = self.output.out_features();
11414 let n_embd = self.cfg.n_embd as usize;
11415 let mut cache = crate::pp::new_cache(e, &self.cfg, tokens.len() + gamma + 8)?;
11416 let (embd_qt, embd_rb) = self.embd.qt_and_row_bytes(n_embd);
11417 let embd_gpu = if spec_host_embd() {
11418 None
11419 } else {
11420 Some(
11421 self.embd_gpu
11422 .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload")),
11423 )
11424 };
11425 let embd_dev = embd_gpu.map(|gpu| (gpu, embd_qt, embd_rb));
11426
11427 struct PendingRecord {
11428 position: usize,
11429 hidden: Option<Vec<f32>>,
11430 tokens: Vec<u32>,
11431 target_top_ids: Vec<Option<Vec<u32>>>,
11432 target_top_logits: Vec<Option<Vec<f32>>>,
11433 target_top_probs: Vec<Option<Vec<f32>>>,
11434 target_tail_probs: Vec<Option<f32>>,
11435 }
11436
11437 let mut pending: Vec<PendingRecord> = anchor_positions
11438 .iter()
11439 .map(|&position| PendingRecord {
11440 position,
11441 hidden: None,
11442 tokens: tokens[position..=position + gamma].to_vec(),
11443 target_top_ids: vec![None; gamma],
11444 target_top_logits: vec![None; gamma],
11445 target_top_probs: vec![None; gamma],
11446 target_tail_probs: vec![None; gamma],
11447 })
11448 .collect();
11449
11450 let mut start = 0usize;
11451 while start < tokens.len() {
11452 let end = (start + chunk).min(tokens.len());
11453 let chunk_tokens = &tokens[start..end];
11454 let (target_logits, hidden_rows) =
11455 self.decode_step_t_core(e, chunk_tokens, start, &mut cache, embd_dev, None)?;
11456 for record in &mut pending {
11457 let hidden_position = record.position - 1;
11458 if hidden_position >= start && hidden_position < end {
11459 let local = hidden_position - start;
11460 record.hidden = Some(
11461 e.dtoh_view(&hidden_rows.slice(local * n_embd..(local + 1) * n_embd))?,
11462 );
11463 }
11464 for slot in 0..gamma {
11465 let target_row = record.position + slot;
11466 if target_row < start || target_row >= end {
11467 continue;
11468 }
11469 let local = target_row - start;
11470 let logits =
11471 e.dtoh_view(&target_logits.slice(local * n_vocab..(local + 1) * n_vocab))?;
11472 let (ids, top_logits, probs, tail) =
11473 dspark_sparse_softmax_topk(&logits, top_k, temperature)?;
11474 record.target_top_ids[slot] = Some(ids);
11475 record.target_top_logits[slot] = Some(top_logits);
11476 record.target_top_probs[slot] = Some(probs);
11477 record.target_tail_probs[slot] = Some(tail);
11478 }
11479 }
11480 start = end;
11481 }
11482
11483 pending
11484 .into_iter()
11485 .map(|record| {
11486 let hidden = record
11487 .hidden
11488 .ok_or_else(|| format!("missing DSpark hidden at {}", record.position))?;
11489 let target_top_ids =
11490 flatten_dspark_rows(record.target_top_ids, record.position, "target ids")?;
11491 let target_top_logits = flatten_dspark_rows(
11492 record.target_top_logits,
11493 record.position,
11494 "target logits",
11495 )?;
11496 let target_top_probs =
11497 flatten_dspark_rows(record.target_top_probs, record.position, "target probs")?;
11498 let target_tail_probs = record
11499 .target_tail_probs
11500 .into_iter()
11501 .enumerate()
11502 .map(|(slot, value)| {
11503 value.ok_or_else(|| {
11504 format!("missing DSpark tail at {} slot {slot}", record.position)
11505 })
11506 })
11507 .collect::<Result<Vec<_>, _>>()?;
11508 Ok(DsparkAnchorRecord {
11509 position: record.position,
11510 hidden,
11511 tokens: record.tokens,
11512 target_top_ids,
11513 target_top_logits,
11514 target_top_probs,
11515 target_tail_probs,
11516 })
11517 })
11518 .collect()
11519 }
11520
11521 /// TEACHER-FORCED REPLAY ACCEPTANCE (hqmtp MTP-heal protocol): walk a FIXED token
11522 /// sequence and, at sampled positions, compare the MTP head's K-token draft chain against
11523 /// the trunk's own teacher-forced greedy predictions. Nothing is generated — the context is
11524 /// the corpus text itself, so (a) degenerate self-generated loops cannot inflate acceptance
11525 /// and (b) two arms (bf16 ceiling vs NVFP4) score on IDENTICAL contexts, isolating the
11526 /// quant-induced head/hidden-state mismatch from text drift.
11527 ///
11528 /// Per eval position p (context = tokens[0..=p], predecessor pairing as in spec decode):
11529 /// draft_j = chain token j from (tokens[p], h_{p-1}), then its own drafts — the exact
11530 /// eager spec-decode chain (same mtp_head_forward_dev, same rope positions).
11531 /// target_j = teacher-forced greedy pick for position p+1+j (argmax of the trunk logits
11532 /// at forced context tokens[0..p+j]). For j==0 this equals live spec
11533 /// acceptance; for j>=1 live verify would condition on the drafts, here it
11534 /// conditions on the corpus — deterministic and arm-comparable by design.
11535 ///
11536 /// Returns (rows, bg): one (p, drafts[k], targets[k]) row per eval position (ascending p),
11537 /// plus the full teacher-forced greedy track bg (bg[i] = greedy pick for position i, i>=1)
11538 /// so harnesses can cross-check runs (e.g. different chunk sizes must give identical bg).
11539 ///
11540 /// `hdump`: when Some, every position's pre-output_norm trunk hidden (the exact rows the
11541 /// draft-KV fill pairs from) streams to the file as little-endian f32 [t_total, n_embd] —
11542 /// the head-distillation extraction (hqmtp): the ENGINE is the source of truth for trunk
11543 /// hiddens (HF torch reproductions of the hybrid trunk measured only ~0.5 greedy
11544 /// agreement vs this path — not usable as a training-data source).
11545 pub fn replay_acceptance(
11546 &self,
11547 e: &Engine,
11548 tokens: &[u32],
11549 k: usize,
11550 stride: usize,
11551 chunk: usize,
11552 mut hdump: Option<&mut std::fs::File>,
11553 ) -> Result<(Vec<(usize, Vec<u32>, Vec<u32>)>, Vec<u32>), Box<dyn std::error::Error>> {
11554 assert!(k >= 1 && stride >= 1 && chunk >= 2);
11555 let mtp = self
11556 .mtp
11557 .as_ref()
11558 .expect("replay_acceptance requires an MTP head");
11559 let n_vocab = self.output.out_features();
11560 let d_vocab = mtp
11561 .shared_head_head
11562 .as_ref()
11563 .unwrap_or(&self.output)
11564 .out_features();
11565 let n_embd = self.cfg.n_embd as usize;
11566 let t_total = tokens.len();
11567 assert!(t_total >= 8, "corpus too short ({t_total} tokens)");
11568 // STAGE-OWNED KV (lane/pp2-spec 2026-08-06) — see `new_session`. Door shut = `Cache::new`.
11569 let mut cache = crate::pp::new_cache(e, &self.cfg, t_total + k + 8)?;
11570 let mut scratch = MtpScratch::new(
11571 e,
11572 &self.cfg,
11573 t_total + k + 8,
11574 self.mtp.as_ref().and_then(|m| m.geom.as_ref()),
11575 )?;
11576 let (embd_qt, embd_rb) = self.embd.qt_and_row_bytes(n_embd);
11577 let embd_gpu = if spec_host_embd() {
11578 None
11579 } else {
11580 Some(
11581 self.embd_gpu
11582 .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload")),
11583 )
11584 };
11585 let embd_dev = embd_gpu.map(|g| (g, embd_qt, embd_rb));
11586
11587 // bg[i] = the trunk's greedy pick for position i under the forced context (i >= 1).
11588 let mut bg: Vec<u32> = vec![0; t_total + 1];
11589 let mut rows: Vec<(usize, Vec<u32>, Vec<u32>)> = Vec::new();
11590 let mut prev_last_h = e.zeros(n_embd)?; // predecessor hidden entering the chunk
11591 let mut seed_buf = e.zeros(n_embd)?;
11592 let mut preds_d = e.alloc_u32_zeroed(chunk)?;
11593 let nll_on = std::env::var("MEMRA_REPLAY_NLL").as_deref() == Ok("1");
11594 let (mut nll_sum, mut nll_cnt) = (0f64, 0u64);
11595 let mut s = 0usize;
11596 while s < t_total {
11597 let cend = (s + chunk).min(t_total);
11598 let tc = cend - s;
11599 let ch = &tokens[s..cend];
11600 // 1. forced trunk pass — verify path (decode-exact contract): all-column logits +
11601 // the chunk's true hiddens.
11602 let (tl_d, vx) = self.decode_step_t_core(e, ch, s, &mut cache, embd_dev, None)?;
11603 for j in 0..tc {
11604 e.argmax_token_device_col(&tl_d, j, n_vocab, &mut preds_d, j)?;
11605 }
11606 let preds = e.dtoh_u32(&preds_d)?;
11607 for j in 0..tc {
11608 bg[s + j + 1] = preds[j];
11609 }
11610 // MEMRA_REPLAY_NLL=1: teacher-forced NLL/perplexity over the same forced pass — the
11611 // checkpoint-quality metric (position j's logits score the GOLD next token).
11612 if nll_on {
11613 let jmax = if cend < t_total { tc } else { tc - 1 }; // last pos has no gold next
11614 if jmax > 0 {
11615 let ids: Vec<u32> = (0..jmax).map(|j| tokens[s + j + 1]).collect();
11616 let rows: Vec<i32> = (0..jmax as i32).collect();
11617 let idsd = e.htod_u32_v(&ids)?;
11618 let rowsd = e.htod_i32(&rows)?;
11619 let mut outd = e.zeros(jmax)?;
11620 e.softmax_gather(&tl_d, n_vocab, &idsd, &rowsd, &mut outd, n_vocab, jmax, 1.0)?;
11621 for pr in e.dtoh(&outd)? {
11622 nll_sum += -((pr.max(1e-30)) as f64).ln();
11623 nll_cnt += 1;
11624 }
11625 }
11626 }
11627 if let Some(f) = hdump.as_deref_mut() {
11628 use std::io::Write;
11629 let host: Vec<f32> = e.dtoh(&vx)?;
11630 // bf16 round-to-nearest-even — f32 doubled the disk bill at bulk
11631 // extraction scale (20M tokens x 4096 = 320GB f32 vs 160GB bf16).
11632 let mut bytes = Vec::with_capacity(tc * n_embd * 2);
11633 for v in &host[..tc * n_embd] {
11634 let b = v.to_bits();
11635 let r = b.wrapping_add(0x7FFF + ((b >> 16) & 1));
11636 bytes.extend_from_slice(&((r >> 16) as u16).to_le_bytes());
11637 }
11638 f.write_all(&bytes)?;
11639 }
11640 // CHAINLESS extraction (stride > corpus, the bulk-hdump mode): no chunk ever
11641 // drafts, so the draft-KV fills are pure waste — skip them (2 MTP-block passes
11642 // per token saved; the forced trunk pass + hdump is all the mode needs).
11643 let chainless = stride > t_total;
11644 if chainless {
11645 e.copy_view_into(
11646 &mut prev_last_h,
11647 0,
11648 &vx.slice((tc - 1) * n_embd..tc * n_embd),
11649 n_embd,
11650 )?;
11651 s = cend;
11652 continue;
11653 }
11654 // 2. TRUE predecessor-paired draft-KV fill for the chunk (row i carries h_{i-1};
11655 // row s reads the previous chunk's last true hidden, zeros at corpus start).
11656 let mut vxs = e.zeros(tc * n_embd)?;
11657 e.copy_into(&mut vxs, 0, &prev_last_h, n_embd)?;
11658 if tc > 1 {
11659 e.copy_view_into(
11660 &mut vxs,
11661 n_embd,
11662 &vx.slice(0..(tc - 1) * n_embd),
11663 (tc - 1) * n_embd,
11664 )?;
11665 }
11666 scratch.set_len(e, s)?;
11667 self.mtp_kv_fill(e, mtp, ch, &vxs, s, &mut scratch, embd_dev)?;
11668 // 3. draft chains at sampled positions, DESCENDING: a chain reads only slots
11669 // [0..p) (true fills) and appends at >= p; the next (smaller-p) chain's set_len
11670 // truncates those approximate appends before they can ever be read.
11671 let ps: Vec<usize> = (s..cend)
11672 .filter(|p| *p >= 1 && *p % stride == 0 && *p + k <= t_total)
11673 .collect();
11674 for &p in ps.iter().rev() {
11675 scratch.set_len(e, p)?;
11676 if p == s {
11677 e.copy_into(&mut seed_buf, 0, &prev_last_h, n_embd)?;
11678 } else {
11679 e.copy_view_into(
11680 &mut seed_buf,
11681 0,
11682 &vx.slice((p - 1 - s) * n_embd..(p - s) * n_embd),
11683 n_embd,
11684 )?;
11685 }
11686 let mut e_tok = tokens[p];
11687 let mut d_seed = e.clone_dtod(&seed_buf)?;
11688 let mut drafts: Vec<u32> = Vec::with_capacity(k);
11689 for j in 0..k {
11690 let (dl_d, h_nextn) = self.mtp_head_forward_dev(
11691 e,
11692 mtp,
11693 e_tok,
11694 &d_seed,
11695 &mut scratch,
11696 p + 1 + j,
11697 embd_dev,
11698 None, // acceptance-oracle walk: no grammar
11699 )?;
11700 let tok_d = e.argmax_token_device(&dl_d, d_vocab)?;
11701 let idx = e.dtoh_u32_one(&tok_d)?;
11702 let d = match &mtp.d2t {
11703 Some(map) => map[idx as usize],
11704 None => idx,
11705 };
11706 drafts.push(d);
11707 e_tok = d;
11708 d_seed = h_nextn;
11709 }
11710 // targets may live in a LATER chunk's bg — resolved after the walk.
11711 rows.push((p, drafts, Vec::new()));
11712 }
11713 // 4. restore TRUE entries for the whole chunk (the next chunk's chains and fills
11714 // expect scratch.len == cend with exact rows).
11715 scratch.set_len(e, s)?;
11716 self.mtp_kv_fill(e, mtp, ch, &vxs, s, &mut scratch, embd_dev)?;
11717 e.copy_view_into(
11718 &mut prev_last_h,
11719 0,
11720 &vx.slice((tc - 1) * n_embd..tc * n_embd),
11721 n_embd,
11722 )?;
11723 s = cend;
11724 }
11725 for (p, drafts, targets) in rows.iter_mut() {
11726 for j in 0..drafts.len() {
11727 targets.push(bg[*p + 1 + j]);
11728 }
11729 }
11730 rows.sort_by_key(|r| r.0);
11731 if nll_cnt > 0 {
11732 let mean = nll_sum / nll_cnt as f64;
11733 println!(
11734 "[replay-nll] tokens={nll_cnt} nll/token={mean:.5} ppl={:.4}",
11735 mean.exp()
11736 );
11737 }
11738 Ok((rows, bg))
11739 }
11740}
11741
11742#[cfg(test)]
11743mod dspark_sparse_tests {
11744 use super::dspark_sparse_softmax_topk;
11745
11746 #[test]
11747 fn topk_keeps_full_softmax_mass_and_stable_ties() {
11748 let logits = [1.0f32, 3.0, 3.0, -2.0];
11749 let (ids, top_logits, probs, tail) = dspark_sparse_softmax_topk(&logits, 2, 1.0).unwrap();
11750 assert_eq!(ids, vec![1, 2]);
11751 assert_eq!(top_logits, vec![3.0, 3.0]);
11752 let denominator = logits.iter().map(|value| (value - 3.0).exp()).sum::<f32>();
11753 let expected = 1.0 / denominator;
11754 assert!((probs[0] - expected).abs() < 1.0e-6);
11755 assert!((probs[1] - expected).abs() < 1.0e-6);
11756 assert!((tail - (1.0 - 2.0 * expected)).abs() < 1.0e-6);
11757 assert!((probs.iter().sum::<f32>() + tail - 1.0).abs() < 1.0e-6);
11758 }
11759}
11760
11761#[cfg(test)]
11762mod spec_replay_env_tests {
11763 use super::spec_replay_env_on;
11764
11765 #[test]
11766 fn replay_requires_literal_one() {
11767 assert!(!spec_replay_env_on(None));
11768 assert!(!spec_replay_env_on(Some("")));
11769 assert!(!spec_replay_env_on(Some("0")));
11770 assert!(!spec_replay_env_on(Some("true")));
11771 assert!(!spec_replay_env_on(Some("2")));
11772 assert!(spec_replay_env_on(Some("1")));
11773 }
11774}
11775
11776#[cfg(test)]
11777mod telem_tests {
11778 use super::{SPEC_TELEM_POS, SpecTelemetry, SpecTelemetryCounters};
11779
11780 #[test]
11781 fn synthetic_accept_masks_produce_tau_and_position_histogram() {
11782 let counters = SpecTelemetryCounters::default();
11783 for mask in [
11784 [true, true, true],
11785 [true, true, false],
11786 [true, false, false],
11787 [false, false, false],
11788 ] {
11789 let accepted = mask.iter().take_while(|&&value| value).count();
11790 counters.record_round(mask.len(), accepted);
11791 }
11792
11793 let snapshot = counters.snapshot();
11794 assert_eq!(
11795 (snapshot.rounds, snapshot.drafted, snapshot.accepted),
11796 (4, 12, 6)
11797 );
11798 assert_eq!(&snapshot.pos_drafted[..3], &[4, 4, 4]);
11799 assert_eq!(&snapshot.pos_accepted[..3], &[3, 2, 1]);
11800 assert_eq!(snapshot.tau(), 1.5);
11801 assert_eq!(snapshot.pos_drafted[3..], [0; SPEC_TELEM_POS - 3]);
11802 assert_eq!(snapshot.pos_accepted[3..], [0; SPEC_TELEM_POS - 3]);
11803 }
11804
11805 /// The worker's per-burst pattern: stash, accumulate, diff — the delta must isolate
11806 /// exactly the burst's contribution (pool-resumed sessions carry prior requests' counts).
11807 #[test]
11808 fn delta_isolates_burst_contribution() {
11809 let mut t = SpecTelemetry::default();
11810 // "previous request": 2 rounds of k=3, accepts 3 then 1.
11811 for (kr, na) in [(3usize, 3usize), (3, 1)] {
11812 t.rounds += 1;
11813 t.drafted += kr as u64;
11814 t.accepted += na as u64;
11815 for j in 0..kr {
11816 t.pos_drafted[j] += 1;
11817 }
11818 for j in 0..na {
11819 t.pos_accepted[j] += 1;
11820 }
11821 }
11822 let before = t;
11823 // "this burst": 1 round k=3, accepts 2.
11824 t.rounds += 1;
11825 t.drafted += 3;
11826 t.accepted += 2;
11827 for j in 0..3 {
11828 t.pos_drafted[j] += 1;
11829 }
11830 for j in 0..2 {
11831 t.pos_accepted[j] += 1;
11832 }
11833 let d = t.delta_since(&before);
11834 assert_eq!((d.rounds, d.drafted, d.accepted), (1, 3, 2));
11835 assert_eq!(&d.pos_drafted[..3], &[1, 1, 1]);
11836 assert_eq!(&d.pos_accepted[..3], &[1, 1, 0]);
11837 assert_eq!(d.pos_drafted[3..], [0; SPEC_TELEM_POS - 3]);
11838 }
11839
11840 /// merge(delta) then merge(delta2) equals accumulating both — the per-model /metrics
11841 /// aggregation invariant.
11842 #[test]
11843 fn merge_accumulates_fieldwise() {
11844 let mut agg = SpecTelemetry::default();
11845 let mut d1 = SpecTelemetry {
11846 rounds: 2,
11847 drafted: 6,
11848 accepted: 4,
11849 ..Default::default()
11850 };
11851 d1.pos_drafted[0] = 2;
11852 d1.pos_accepted[0] = 2;
11853 let mut d2 = SpecTelemetry {
11854 rounds: 1,
11855 drafted: 3,
11856 accepted: 1,
11857 ..Default::default()
11858 };
11859 d2.pos_drafted[0] = 1;
11860 d2.pos_accepted[0] = 1;
11861 d2.pos_drafted[1] = 1;
11862 agg.merge(&d1);
11863 agg.merge(&d2);
11864 assert_eq!((agg.rounds, agg.drafted, agg.accepted), (3, 9, 5));
11865 assert_eq!(agg.pos_drafted[0], 3);
11866 assert_eq!(agg.pos_accepted[0], 3);
11867 assert_eq!(agg.pos_drafted[1], 1);
11868 assert_eq!(agg.pos_accepted[1], 0);
11869 }
11870
11871 /// Wrong-snapshot diff saturates to zero instead of wrapping — the counters feed a
11872 /// public metrics surface and must never publish a u64-wrapped garbage value.
11873 #[test]
11874 fn delta_saturates_never_wraps() {
11875 let small = SpecTelemetry {
11876 rounds: 1,
11877 drafted: 2,
11878 accepted: 1,
11879 ..Default::default()
11880 };
11881 let big = SpecTelemetry {
11882 rounds: 5,
11883 drafted: 15,
11884 accepted: 9,
11885 ..Default::default()
11886 };
11887 let d = small.delta_since(&big);
11888 assert_eq!((d.rounds, d.drafted, d.accepted), (0, 0, 0));
11889 }
11890}
11891
11892#[cfg(test)]
11893mod opti_fork_tests {
11894 use super::{
11895 OptiControllerPolicy, OptiForkAction, OptiForkGateMode, OptiForkGenerationTracker,
11896 };
11897
11898 #[test]
11899 fn controller_threshold_and_three_miss_breaker_are_exact() {
11900 let mut policy = OptiControllerPolicy {
11901 threshold: 0.7,
11902 consecutive_misses: 0,
11903 breaker_tripped: false,
11904 };
11905 assert!(!policy.admit(0.699_999));
11906 assert!(policy.admit(0.7));
11907 assert!(!policy.resolve(false));
11908 assert!(!policy.resolve(false));
11909 assert!(policy.resolve(false));
11910 assert!(policy.breaker_tripped);
11911 assert!(!policy.admit(1.0));
11912 assert!(
11913 !policy.resolve(true),
11914 "a resolved hit cannot re-arm a tripped request"
11915 );
11916 assert!(policy.breaker_tripped);
11917 }
11918
11919 #[test]
11920 fn zero_threshold_is_the_true_unconditional_measurement_arm() {
11921 let mut policy = OptiControllerPolicy {
11922 threshold: 0.0,
11923 consecutive_misses: 0,
11924 breaker_tripped: false,
11925 };
11926 for _ in 0..16 {
11927 assert!(policy.admit(0.0));
11928 assert!(!policy.resolve(false));
11929 }
11930 for invalid in [f32::NAN, f32::INFINITY, -0.01, 1.01] {
11931 assert!(
11932 !policy.admit(invalid),
11933 "invalid q proxy must fail closed: {invalid}"
11934 );
11935 }
11936 assert!(!policy.breaker_tripped);
11937 assert_eq!(policy.consecutive_misses, 0);
11938 }
11939
11940 #[test]
11941 fn alternating_mode_flips_by_generation_not_round_parity() {
11942 assert_eq!(OptiForkGateMode::Alternate.action(0), OptiForkAction::Hit);
11943 assert_eq!(OptiForkGateMode::Alternate.action(1), OptiForkAction::Miss);
11944 assert_eq!(OptiForkGateMode::Alternate.action(8), OptiForkAction::Hit);
11945 assert_eq!(OptiForkGateMode::Alternate.action(9), OptiForkAction::Miss);
11946 }
11947
11948 #[test]
11949 fn live_generation_cannot_be_overwritten() {
11950 let mut tracker = OptiForkGenerationTracker::default();
11951 let g0 = tracker.reserve().unwrap();
11952 let g1 = tracker.reserve().unwrap();
11953 let err = tracker.reserve().unwrap_err().to_string();
11954 assert!(
11955 err.contains("still owns generation 0"),
11956 "unexpected error: {err}"
11957 );
11958 tracker.retire(g0).unwrap();
11959 let g2 = tracker.reserve().unwrap();
11960 assert_eq!((g2.id, g2.slot), (2, 0));
11961 tracker.retire(g1).unwrap();
11962 tracker.retire(g2).unwrap();
11963 }
11964
11965 #[test]
11966 fn teardown_rejects_a_stale_generation_tag() {
11967 let mut tracker = OptiForkGenerationTracker::default();
11968 let g0 = tracker.reserve().unwrap();
11969 tracker.retire(g0).unwrap();
11970 let err = tracker.retire(g0).unwrap_err().to_string();
11971 assert!(err.contains("teardown mismatch"), "unexpected error: {err}");
11972 }
11973}
11974
11975#[cfg(test)]
11976mod draft_graph_fallback_tests {
11977 use super::DraftGraphFallback;
11978
11979 /// The Q2 contract, part (a): a fallback flip is LOUD — exactly once per flip.
11980 #[test]
11981 fn flip_is_loud_once_and_memoized_after() {
11982 let mut f = DraftGraphFallback::default();
11983 let line = f
11984 .mark_greedy("out of memory")
11985 .expect("first flip must return the warn line");
11986 assert!(
11987 line.contains("WARN"),
11988 "flip line must be warn-level: {line}"
11989 );
11990 assert!(
11991 line.contains("out of memory"),
11992 "flip line must carry the reason: {line}"
11993 );
11994 assert!(f.greedy_failed());
11995 // re-marking an already-failed graph is the memoization: quiet, still failed.
11996 assert!(f.mark_greedy("out of memory").is_none());
11997 assert!(f.greedy_failed());
11998 // the two graphs' flags are independent (greedy flip leaves sampled capturable).
11999 assert!(!f.sampled_failed());
12000 let line_s = f
12001 .mark_sampled("capture unsupported")
12002 .expect("sampled flip is its own flip");
12003 assert!(
12004 line_s.contains("sampled"),
12005 "sampled flip names itself: {line_s}"
12006 );
12007 assert!(f.mark_sampled("capture unsupported").is_none());
12008 }
12009
12010 /// The Q2 contract, part (b): resume-from-pool RESETS both flags (fresh capture chance),
12011 /// and says so exactly when there was something to reset.
12012 #[test]
12013 fn reset_on_resume_clears_flags_and_logs_once() {
12014 let mut f = DraftGraphFallback::default();
12015 // clean session: resume is silent, nothing to reset.
12016 assert!(f.reset_on_resume().is_none());
12017 f.mark_greedy("oom").unwrap();
12018 f.mark_sampled("oom").unwrap();
12019 let note = f
12020 .reset_on_resume()
12021 .expect("a set flag must produce the reset note");
12022 assert!(
12023 note.contains("greedy+sampled"),
12024 "note names what was reset: {note}"
12025 );
12026 assert!(
12027 !f.greedy_failed() && !f.sampled_failed(),
12028 "both flags cleared"
12029 );
12030 // and the NEXT failure after a reset is a fresh flip — loud again.
12031 assert!(f.mark_greedy("oom again").is_some());
12032 let note2 = f.reset_on_resume().expect("greedy-only reset");
12033 assert!(note2.contains("(greedy)"), "single-flag note: {note2}");
12034 }
12035
12036 /// Shape-change clears (dmask realloc / mask-shape mismatch / s_key change) stay silent —
12037 /// they precede a fresh capture attempt whose own failure re-flips loudly.
12038 #[test]
12039 fn shape_change_clears_are_silent() {
12040 let mut f = DraftGraphFallback::default();
12041 f.mark_greedy("oom").unwrap();
12042 f.clear_greedy();
12043 assert!(!f.greedy_failed());
12044 f.mark_sampled("oom").unwrap();
12045 f.clear_sampled();
12046 assert!(!f.sampled_failed());
12047 // after a silent clear there is nothing left for resume to report.
12048 assert!(f.reset_on_resume().is_none());
12049 }
12050}
12051
12052/// SAMPLED DRAFT-GRAPH KEY (lane/graph-s-key-exactness-20260819).
12053///
12054/// These are the CPU teeth for an exactness bug whose live reproduction needs a GPU, a trunk, a
12055/// drafter and a two-turn session: the key itself. Every test below fails against the pre-fix key
12056/// `(seed, temp.to_bits(), k)` — `legacy_key` restates it so the collision is explicit rather
12057/// than remembered.
12058#[cfg(test)]
12059mod sampled_graph_key_tests {
12060 use super::{SampledGraphKey, debug_t_pred0};
12061
12062 /// The pre-fix key, verbatim: `let s_key = (sp_seed, sp_temp.to_bits(), k);`
12063 fn legacy_key(k: &SampledGraphKey) -> (u64, u32, usize) {
12064 (k.seed, k.temp_bits, k.k)
12065 }
12066
12067 fn pure_temp_key() -> SampledGraphKey {
12068 // temperature 1.0, filters off — today's serve default, the shape that parks a graph.
12069 SampledGraphKey::new(12345, 1.0, 3, 0, 1.0, 0.0, false)
12070 }
12071
12072 /// THE COLLISION. Two requests that differ ONLY in the truncation filters shared one key, so
12073 /// a parked pure-temp graph survived into a filtered request and the launch site launched it.
12074 #[test]
12075 fn vendor_filters_change_the_key() {
12076 let parked = pure_temp_key();
12077 // qwen3.8 generation_config.json — what the vendor-default flip makes the default shape.
12078 let vendor = SampledGraphKey::new(12345, 1.0, 3, 20, 0.95, 0.0, false);
12079 assert_eq!(
12080 legacy_key(&parked),
12081 legacy_key(&vendor),
12082 "pre-fix key collided: this is the bug, and the reason a test asserts on it",
12083 );
12084 assert_ne!(parked, vendor, "post-fix key must separate the two regimes");
12085 assert!(parked.pure_temp());
12086 assert!(!vendor.pure_temp());
12087 }
12088
12089 /// Each distribution-shaping field alone is enough to drop the parked graph.
12090 #[test]
12091 fn every_filter_field_is_keyed() {
12092 let base = pure_temp_key();
12093 for (what, other) in [
12094 (
12095 "top_k",
12096 SampledGraphKey::new(12345, 1.0, 3, 20, 1.0, 0.0, false),
12097 ),
12098 (
12099 "top_p",
12100 SampledGraphKey::new(12345, 1.0, 3, 0, 0.95, 0.0, false),
12101 ),
12102 (
12103 "min_p",
12104 SampledGraphKey::new(12345, 1.0, 3, 0, 1.0, 0.05, false),
12105 ),
12106 (
12107 "penalties",
12108 SampledGraphKey::new(12345, 1.0, 3, 0, 1.0, 0.0, true),
12109 ),
12110 ] {
12111 assert_ne!(base, other, "{what} must be part of the key");
12112 assert!(!other.pure_temp(), "{what} leaves the pure-temp regime");
12113 assert_eq!(
12114 legacy_key(&base),
12115 legacy_key(&other),
12116 "{what} was invisible to the pre-fix key",
12117 );
12118 }
12119 }
12120
12121 /// The baked constants stay keyed (this half was always right — regression cover for it).
12122 #[test]
12123 fn baked_constants_stay_keyed() {
12124 let base = pure_temp_key();
12125 assert_ne!(
12126 base,
12127 SampledGraphKey::new(999, 1.0, 3, 0, 1.0, 0.0, false),
12128 "seed"
12129 );
12130 assert_ne!(
12131 base,
12132 SampledGraphKey::new(12345, 0.7, 3, 0, 1.0, 0.0, false),
12133 "temp"
12134 );
12135 assert_ne!(
12136 base,
12137 SampledGraphKey::new(12345, 1.0, 4, 0, 1.0, 0.0, false),
12138 "k"
12139 );
12140 // bitwise on temperature: 0.7f32 vs the same value re-derived must NOT differ.
12141 assert_eq!(
12142 SampledGraphKey::new(1, 0.7, 3, 0, 1.0, 0.0, false),
12143 SampledGraphKey::new(1, 7.0 / 10.0, 3, 0, 1.0, 0.0, false),
12144 );
12145 }
12146
12147 /// THE LOAD-BEARING HALF OF THE SEED DECISION (lane/session-resume-sampler-predicate-
12148 /// 20260820). The whole-session resume predicate deliberately does NOT compare `seed`: an
12149 /// omitted serve `seed` draws fresh per-request entropy, so comparing it would refuse every
12150 /// seed-omitting sampled conversation. That is only sound because the one piece of parked state
12151 /// that BAKES the seed — this graph — is re-keyed on it, so a seed change drops and recaptures.
12152 ///
12153 /// This test is the other end of that argument, asserted here rather than remembered in a
12154 /// comment: if a future change dropped `seed` from the key, the resume predicate's exclusion
12155 /// would silently become the unsound thing it is documented not to be.
12156 /// (Paired with `seed_alone_does_not_refuse` in `memra-sampling`.)
12157 #[test]
12158 fn seed_alone_still_rekeys_the_draft_graph() {
12159 let parked = pure_temp_key();
12160 let reseeded = SampledGraphKey::new(999, 1.0, 3, 0, 1.0, 0.0, false);
12161 assert_ne!(
12162 parked, reseeded,
12163 "a seed-only change MUST drop the parked sampled graph — the resume predicate's \
12164 decision not to compare seed rests on exactly this",
12165 );
12166 // Same regime on both sides: the drop is a recapture, not a fall to the eager chain
12167 // because of a filter difference.
12168 assert!(parked.pure_temp() && reseeded.pure_temp());
12169 }
12170
12171 /// `pure_temp()` is the capture guard's predicate, computed from the key so the two cannot
12172 /// drift. The equality below is the invariant the launch-site guard asserts: identical keys
12173 /// agree on the regime, so a graph that survives the drop is legal to launch.
12174 #[test]
12175 fn equal_keys_agree_on_the_regime() {
12176 let a = SampledGraphKey::new(7, 0.8, 3, 20, 0.95, 0.0, false);
12177 let b = SampledGraphKey::new(7, 0.8, 3, 20, 0.95, 0.0, false);
12178 assert_eq!(a, b);
12179 assert_eq!(a.pure_temp(), b.pure_temp());
12180 // top_p slightly above 1.0 (a client sending 1.0 exactly, or an operator default) is
12181 // still the unfiltered regime, matching the original `sp.top_p >= 1.0` test.
12182 assert!(SampledGraphKey::new(7, 0.8, 3, 0, 1.0, 0.0, false).pure_temp());
12183 assert!(SampledGraphKey::new(7, 0.8, 3, 0, 1.5, -1.0, false).pure_temp());
12184 }
12185
12186 /// MEMRA_DEBUG_SPEC on a SAMPLED spec request past round 0: the print must render without
12187 /// indexing the empty greedy `preds` vector (it panicked the GPU worker before this lane).
12188 #[test]
12189 fn debug_print_survives_the_sampled_arm() {
12190 // round >= 1 with a pending bonus == base 1, sampled == `preds` empty.
12191 assert_eq!(debug_t_pred0(true, 1, 4242, &[]), "n/a");
12192 assert_eq!(debug_t_pred0(true, 2, 4242, &[]), "n/a");
12193 // round 0 without a pending bonus still reports last_pred, in both arms.
12194 assert_eq!(debug_t_pred0(true, 0, 4242, &[]), "4242");
12195 assert_eq!(debug_t_pred0(false, 0, 4242, &[7, 8]), "4242");
12196 // greedy keeps the real prediction it always printed.
12197 assert_eq!(debug_t_pred0(false, 1, 4242, &[7, 8]), "7");
12198 assert_eq!(debug_t_pred0(false, 2, 4242, &[7, 8]), "8");
12199 }
12200}