memra_engine/spec.rs
1//! Qwen3.5 MTP (NextN) greedy speculative decode (research/mtp/MTP-PLAN.md §A/§B/§C/§D).
2//!
3//! Greedy spec decode is MATHEMATICALLY EXACT: the accepted+bonus token stream is token-for-token
4//! identical to plain greedy `generate`. This module provides:
5//! - `mtp_head_forward` (§A, T=1): one NextN draft-token forward.
6//! - `decode_step_t` (§D.3, T=K+1): batched target verify forward, all-column logits.
7//! - `generate_spec` (§B): the draft/verify/accept/rollback orchestrator.
8//! Cache snapshot/rollback lives in cache.rs (§D.4). The MTP head uses its OWN scratch KV (§D.6),
9//! PERSISTENT over the committed sequence (see `MtpScratch`).
10
11use crate::Engine;
12use crate::cache::{Cache, KvLayer};
13use crate::forward::argmax;
14use crate::hybrid::{FullAttnLayer, HybridModel, LinearAttnLayer, Mixer, MtpHead};
15use cudarc::driver::CudaSlice;
16use std::sync::atomic::{AtomicU64, Ordering};
17
18/// Parse the documented `MEMRA_SPEC_REPLAY=1` rollback seam.
19///
20/// Keep this shared with serving admission so `=0` cannot select replay in one
21/// layer while another layer treats it as disabled.
22pub fn spec_replay_env_on(value: Option<&str>) -> bool {
23 value == Some("1")
24}
25
26pub fn spec_replay_env_enabled() -> bool {
27 let value = std::env::var("MEMRA_SPEC_REPLAY").ok();
28 spec_replay_env_on(value.as_deref())
29}
30
31/// One compact, anchor-bounded DSpark supervision record. `tokens[0]` is the anchor at p and
32/// `hidden` is its predecessor carrier h[p-1], matching the live NextN/DSpark pairing. Target
33/// rows p..p+gamma-1 score tokens p+1..p+gamma. They are the full-target softmax's top-k
34/// entries; `target_tail_probs[j]` is the probability mass outside those rows. All flattened
35/// target arrays are `[gamma, top_k]` in row-major order.
36pub struct DsparkAnchorRecord {
37 pub position: usize,
38 pub hidden: Vec<f32>,
39 pub tokens: Vec<u32>,
40 pub target_top_ids: Vec<u32>,
41 pub target_top_logits: Vec<f32>,
42 pub target_top_probs: Vec<f32>,
43 pub target_tail_probs: Vec<f32>,
44}
45
46fn dspark_sparse_softmax_topk(
47 logits: &[f32],
48 top_k: usize,
49 temperature: f32,
50) -> Result<(Vec<u32>, Vec<f32>, Vec<f32>, f32), Box<dyn std::error::Error>> {
51 if logits.is_empty() || top_k == 0 || top_k > logits.len() || temperature <= 0.0 {
52 return Err("invalid DSpark sparse-softmax shape or temperature".into());
53 }
54 if logits.iter().any(|value| !value.is_finite()) {
55 return Err("DSpark target logits contain a non-finite value".into());
56 }
57 let mut ranked: Vec<(u32, f32)> = logits
58 .iter()
59 .copied()
60 .enumerate()
61 .map(|(index, value)| (index as u32, value))
62 .collect();
63 let compare = |left: &(u32, f32), right: &(u32, f32)| {
64 right.1.total_cmp(&left.1).then(left.0.cmp(&right.0))
65 };
66 ranked.select_nth_unstable_by(top_k - 1, compare);
67 ranked[..top_k].sort_unstable_by(compare);
68
69 let max_logit = logits.iter().copied().fold(f32::NEG_INFINITY, f32::max);
70 let inv_temperature = 1.0f64 / temperature as f64;
71 let denominator: f64 = logits
72 .iter()
73 .map(|value| (((*value - max_logit) as f64) * inv_temperature).exp())
74 .sum();
75 let ids: Vec<u32> = ranked[..top_k].iter().map(|(index, _)| *index).collect();
76 let top_logits: Vec<f32> = ranked[..top_k].iter().map(|(_, value)| *value).collect();
77 let top_probs: Vec<f32> = top_logits
78 .iter()
79 .map(|value| ((((value - max_logit) as f64) * inv_temperature).exp() / denominator) as f32)
80 .collect();
81 let top_mass: f64 = top_probs.iter().map(|value| *value as f64).sum();
82 let tail = (1.0f64 - top_mass).clamp(0.0, 1.0) as f32;
83 Ok((ids, top_logits, top_probs, tail))
84}
85
86fn flatten_dspark_rows<T>(
87 rows: Vec<Option<Vec<T>>>,
88 position: usize,
89 label: &str,
90) -> Result<Vec<T>, Box<dyn std::error::Error>> {
91 let mut flattened = Vec::new();
92 for (slot, row) in rows.into_iter().enumerate() {
93 flattened.extend(
94 row.ok_or_else(|| format!("missing DSpark {label} at {position} slot {slot}"))?,
95 );
96 }
97 Ok(flattened)
98}
99
100/// H-SEED CONVENTION (MEMRA_SPEC_HPOST=1): feed the MTP head the POST-norm hidden — trunk rows
101/// hand over `output_norm(x)` and the draft chain recurrence hands over `shared_head_norm(h_nextn)`
102/// (= final_h) — matching the reference engines: llama.cpp #24025 ("qwen35: use post-norm hidden
103/// state for MTP", t_h_nextn is taken AFTER the final norm in both trunk and MTP graphs) and
104/// SGLang's qwen3_5_mtp (spec_info.hidden_states = the target model's post-norm output). memra's
105/// historical convention (default, MTP-PLAN §A) is PRE-norm x. Draft-quality-only: exactness is
106/// the verify's job either way; acceptance arbitrates. OnceLock: read once, hot-loop safe.
107pub(crate) fn spec_hpost() -> bool {
108 static H: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
109 *H.get_or_init(|| {
110 std::env::var("MEMRA_SPEC_HPOST")
111 .map(|v| v != "0")
112 .unwrap_or(false)
113 })
114}
115
116/// LEAN VERIFY (default ON since 2026-07-08; MEMRA_SPEC_LEAN=0 reverts — close35 lane): the verify m-scaling
117/// probe + nsys diff showed the verify t-path pays ~1.0ms/call at m=1 over eager decode on the
118/// 35B, and the kernels are NOT the cause (dev-MoE identical, kernel-time delta only +179us).
119/// The overhead is (a) ~250 extra cuMemsetD8Async/call from `e.zeros()` on buffers every kernel
120/// fully overwrites (~0.9ms host issue + ~0.35ms GPU) and (b) the t=1 FA rows dispatch (rows_v2 +
121/// combine_rows, +50us vs the eager fa_decode pair). This flag switches (a) fully-overwritten
122/// verify buffers to `e.uninit` (identical bytes: every element is written before read) and
123/// (b) t==1 verify FA to the eager `fa_decode` entry (byte-identical: kernel-check pins the
124/// rows-vs-loop identity and the per-row loop at t=1 IS fa_decode on the same q). Gates arbitrate.
125pub(crate) fn spec_lean() -> bool {
126 static L: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
127 // DEFAULT ON since 2026-07-08 (MEMRA_SPEC_LEAN=0 reverts): bit-identical (buffers fully
128 // overwritten; gates green incl maxdiff-identical run-gen) and measured +2.4% e2e p3 /
129 // +1.5% p2 at the daily 35B config. m=1 verify now costs eager-decode parity.
130 *L.get_or_init(|| {
131 std::env::var("MEMRA_SPEC_LEAN")
132 .map(|v| v != "0")
133 .unwrap_or(true)
134 })
135}
136
137/// SMALL-M BATCHED VERIFY (default ON since 2026-07-09; MEMRA_SPEC_M2=0 reverts — lane/spec-m2): extend the
138/// batched linear-attn verify arm down to t=2 and batch the MoE dev token loop over a
139/// grid.z=token axis at every verify t. The close35 m-scaling probe put the m=2 verify tier at
140/// x1.54 of m=1 (llama x1.14); the per-column linear chain (t<3) and the serial MoE dev token
141/// loop are the two launch-structure causes. Both changes are LAUNCH-STRUCTURE ONLY:
142/// (a) the batched conv's t<pad ring update is pure copies (ssm_conv_ring_rebuild from a cloned
143/// ring — the ring stores raw input columns); every arithmetic kernel is the same one the
144/// t>=3 arm already runs (matmul_decode_exact bit-identical at m=2-4, gdn_scan's internal
145/// t-loop == chained T=1 steps);
146/// (b) the MoE dev-rows twins run the serial loop's per-token warp program with tok-offset
147/// pointers (same sel/w/aq/ad bytes, same dot order, same slot-ordered FMA chain).
148/// Gates arbitrate: run-spec K=1..8 self-consistency (35B+9B), kernel-check, run-gen argmax.
149pub(crate) fn spec_m2() -> bool {
150 static M: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
151 // DEFAULT ON since 2026-07-09 (MEMRA_SPEC_M2=0 reverts): launch-structure only — t=2
152 // batched linear arm (ring-roll copies, zero new FP order) + MoE dev-rows kernels
153 // (grid.z=token, 4 launches/layer at any verify t). Acceptance bit-identical at every K;
154 // 35B p2 +3.4% / p3 +3.6%; the profitable-K plateau widens (new optimum K=3 at 223).
155 *M.get_or_init(|| {
156 std::env::var("MEMRA_SPEC_M2")
157 .map(|v| v != "0")
158 .unwrap_or(true)
159 })
160}
161pub(crate) fn spec_stream() -> bool {
162 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
163 *ON.get_or_init(|| std::env::var("MEMRA_SPEC_STREAM").as_deref() == Ok("1"))
164}
165pub(crate) fn spec_stream_m() -> usize {
166 static M: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
167 *M.get_or_init(|| {
168 std::env::var("MEMRA_SPEC_STREAM_M")
169 .ok()
170 .and_then(|v| v.parse().ok())
171 .unwrap_or(4)
172 })
173}
174pub(crate) fn spec_devacc() -> bool {
175 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
176 *ON.get_or_init(|| std::env::var("MEMRA_SPEC_DEVACC").as_deref() == Ok("1"))
177}
178
179/// `t_pred0` for the `MEMRA_DEBUG_SPEC` per-round print, sampled-safe.
180///
181/// `generate_spec_inner2` fills its `preds` vector ONLY on the greedy path (`if !sampled`), and
182/// the per-round debug print was the sole consumer in the sampled arm: `t_pred(0)` survives round
183/// 0 (`base == 0` returns `last_pred`) and from round 1 (`base == 1`, a pending bonus) indexes an
184/// EMPTY vector — `index out of bounds: the len is 0 but the index is 0`, in the GPU worker
185/// thread, which then respawns and reloads weights while the request dies. So any sampled spec
186/// request longer than one round used to kill the worker whenever `MEMRA_DEBUG_SPEC` was set:
187/// the flag crashed precisely the regime it exists to investigate.
188///
189/// Fixed at the print site, not inside the closure, so the greedy accept walk keeps its strict
190/// indexing (an out-of-range pred there is a real bug and must still be loud).
191fn debug_t_pred0(sampled: bool, base: usize, last_pred: u32, preds: &[u32]) -> String {
192 if base == 0 {
193 return last_pred.to_string();
194 }
195 match preds.get(base - 1) {
196 Some(p) => p.to_string(),
197 // sampled: the greedy per-column argmax was never run for this round.
198 None => {
199 debug_assert!(
200 sampled,
201 "greedy spec: preds[{}] missing at base {base}",
202 base - 1
203 );
204 "n/a".to_string()
205 }
206 }
207}
208
209/// `MEMRA_SKEY_PROBE=1` — sampled-draft-graph key probe (lane/graph-s-key-exactness-20260819).
210///
211/// Reports, per burst and per round, which draft chain the sampled arm chose and under which
212/// filter regime, plus the ONE observable that separates a legal filtered draft from a stale
213/// pure-temp graph replayed under filters: an accept test whose gathered `q` is exactly 0.
214/// A draft token sampled from the FILTERED softmax can never gather q=0 (it was drawn from the
215/// kept set), so `q=0` in the verify means the draft came from a distribution the verify does
216/// not believe in — and `u * 0 < p` then accepts it unconditionally.
217///
218/// Its own env var, deliberately NOT `MEMRA_DEBUG_SPEC`: that flag panicked the GPU worker on
219/// any sampled spec request past round 0 until this lane fixed it (§2 of the bank note).
220pub(crate) fn skey_probe() -> bool {
221 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
222 *ON.get_or_init(|| std::env::var("MEMRA_SKEY_PROBE").as_deref() == Ok("1"))
223}
224
225/// GRAMMAR HOOK for constrained spec decode (lane/constrained-full, 2026-08-03). The engine
226/// stays llguidance-agnostic: the server adapts its per-session grammar state behind this
227/// trait. CONTRACT (the verify-side truncation rule — token-identical to constrained plain
228/// greedy decode): the exactness walk runs UNMASKED first; the hook then (a) truncates
229/// acceptance at the first grammar-illegal accepted token, and (b) when the truncation fired
230/// or the bonus is illegal, the engine recomputes that slot as the MASKED argmax of the
231/// target's own verify column (an unmasked argmax that is grammar-legal IS the masked argmax
232/// — masking only removes tokens — so the common case pays nothing). `consume` advances the
233/// state with each EMITTED token in order; EOS handling is the implementor's job (skip).
234pub trait SpecConstraint {
235 /// -inf the current state's banned ids on a HOST logits row (prompt-tail / init-feed
236 /// masked argmax).
237 fn mask_logits(&mut self, logits: &mut [f32]) -> Result<(), String>;
238 /// Packed 32-bit bitset words of the CURRENT state's allowed set (device-mask form).
239 fn mask_words(&mut self) -> Result<Vec<u32>, String>;
240 /// Is `tok` consumable in the CURRENT state?
241 fn is_allowed(&mut self, tok: u32) -> Result<bool, String>;
242 /// Advance the state with an emitted token.
243 fn consume(&mut self, tok: u32) -> Result<(), String>;
244
245 // --- DRAFT-SIDE MASKING (lane/draft-mask, 2026-08-04) ---
246 // The drafter proposed grammar-illegal tokens under tight schemas, so verify-side
247 // truncation cut nearly every round (measured acceptance 0.467-0.513 tight vs 0.62-0.82
248 // loose, research/constrained-full-20260803). These three methods let the engine mask the
249 // DRAFT model's own sampling with the grammar's legal set, so proposals are legal by
250 // construction. The state they walk is a SPECULATIVE CLONE of the session matcher — the
251 // real state is advanced only by `consume` (emitted tokens), so verify-side truncation
252 // stays the correctness backstop and the emitted stream is unchanged by construction
253 // (an accepted draft is the target's unmasked argmax AND grammar-legal, hence the masked
254 // argmax; a cut slot is recomputed as the masked argmax either way).
255 // Default impls = feature OFF (pre-lane behaviour: unmasked drafts).
256
257 /// Is draft-side masking available on this hook? Probed ONCE per burst, before the draft
258 /// graph is captured (the mask is an in-graph node — its presence is a capture-time shape).
259 fn draft_mask_enabled(&self) -> bool {
260 false
261 }
262 /// Start a draft chain: clone the CURRENT (committed) grammar state into the speculative
263 /// slot. Called once per spec round, before the first draft position.
264 fn draft_begin(&mut self) -> Result<(), String> {
265 Ok(())
266 }
267 /// Packed 32-bit bitset words of the SPECULATIVE state's allowed set (target-vocab ids),
268 /// for the draft position about to be sampled. `None` = draft masking off (no-op).
269 fn draft_mask_words(&mut self) -> Result<Option<Vec<u32>>, String> {
270 Ok(None)
271 }
272 /// Advance the SPECULATIVE state with a PROPOSED draft token. `false` = the chain cannot
273 /// continue (EOS proposed, or an unmasked position proposed something illegal) — the
274 /// engine stops drafting; the token already pushed still goes through verify.
275 fn draft_advance(&mut self, _tok: u32) -> Result<bool, String> {
276 Ok(false)
277 }
278}
279
280/// DRAFT-MASK UPLOAD (lane/draft-mask): pull the speculative state's allowed set (TARGET-id
281/// space) from the hook, project it into the DRAFT head's vocab space, and upload it into the
282/// stable device buffer the draft chain reads. Returns false when the chain must stop drafting:
283/// the hook handed out no mask, or NO draft-vocab row is grammar-legal at this position (a
284/// trimmed FR-Spec head genuinely cannot propose a legal token there — masking it would leave
285/// a fully-banned row whose argmax is meaningless, so the round drafts fewer tokens and the
286/// verify emits the masked argmax as usual).
287fn upload_draft_mask(
288 e: &Engine,
289 c: &mut dyn SpecConstraint,
290 dst: &mut CudaSlice<u32>,
291 d2t: Option<&Vec<u32>>,
292 d_vocab: usize,
293 words: usize,
294) -> Result<bool, Box<dyn std::error::Error>> {
295 let Some(tw) = c
296 .draft_mask_words()
297 .map_err(|e2| format!("constraint: {e2}"))?
298 else {
299 return Ok(false);
300 };
301 let bit = |t: usize| -> bool {
302 let w = t >> 5;
303 w < tw.len() && (tw[w] >> (t & 31)) & 1 == 1
304 };
305 let mut buf = vec![0u32; words];
306 match d2t {
307 // TRIMMED draft head: row i proposes target id d2t[i] — permute the mask accordingly.
308 Some(map) => {
309 for (i, &t) in map.iter().enumerate().take(d_vocab) {
310 if bit(t as usize) {
311 buf[i >> 5] |= 1u32 << (i & 31);
312 }
313 }
314 }
315 // UNTRIMMED: draft ids ARE target ids; the packed words transfer verbatim (a short
316 // mask leaves the padded tail zeroed == banned, same rule as constrained::apply_mask).
317 None => {
318 let n = tw.len().min(words);
319 buf[..n].copy_from_slice(&tw[..n]);
320 }
321 }
322 if buf.iter().all(|w| *w == 0) {
323 return Ok(false);
324 }
325 e.htod_u32_into(dst, &buf)?;
326 Ok(true)
327}
328
329/// Keep the full token-embedding table in host memory and upload only the rows needed by each
330/// MTP/verify step. This is an exact memory-capacity seam for very large BF16 vocab tables: host
331/// gather expands the same source bits to f32, and only O(T*n_embd) bytes cross PCIe per step.
332/// CUDA-graph/round-stream draft paths require device token ids and therefore stay disabled.
333pub(crate) fn spec_host_embd() -> bool {
334 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
335 *ON.get_or_init(|| std::env::var("MEMRA_SPEC_HOST_EMBD").as_deref() == Ok("1"))
336}
337
338/// VERIFY-TIER TRUNK LAUNCH-FUSION (default ON since 2026-07-09; MEMRA_SPEC_FUSED_T=0 reverts — lane/close35b): extend
339/// the t=1 fused2/fused3 Q8_0 trunk launches to the batched verify tier (t=2-4, the K=1..3
340/// verify shapes). At t>1 the trunk pairs/triples (35B wqkv+wqkv_gate, wq/wk/wv,
341/// gate_shexp+up_shexp) each run a separate `matmul_decode_exact` — one q8_1 re-quantize of the
342/// SAME activation plus one _b2/_b4 launch per tensor. The fused twins share ONE quantize and
343/// ONE launch per group; per (tensor,token,row) the kernel body is q8_0_mmvq_batched verbatim
344/// with the identical row mapping -> BIT-IDENTICAL by construction (kernel-check pins it,
345/// run-spec K=1..8 + acceptance identity arbitrate e2e).
346pub(crate) fn spec_fused_t() -> bool {
347 static F: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
348 // DEFAULT ON since 2026-07-09 (MEMRA_SPEC_FUSED_T=0 reverts): verify t=2-4 trunk launch-fusion
349 // (fused2/fused3 Q8_0 batched twins, bit-identical by construction — m=1 block-offset split on
350 // the batched body). m=2 marginal token 2117->1762us; 35B daily: p3 +3.7% (crosses llama), p2 +5%.
351 *F.get_or_init(|| {
352 std::env::var("MEMRA_SPEC_FUSED_T")
353 .map(|v| v != "0")
354 .unwrap_or(true)
355 })
356}
357
358/// zeros/uninit switch for verify-path buffers that are FULLY OVERWRITTEN before any read.
359/// Only call this on such buffers — the lean contract is "identical bytes by construction".
360fn vbuf(e: &Engine, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
361 if spec_lean() { e.uninit(n) } else { e.zeros(n) }
362}
363
364/// Scratch KV for the MTP block (one full-attn layer).
365///
366/// PERSISTENT MODE (default, 2026-07-03 — the acceptance lever): sized cap = max_ctx and kept in
367/// sync with the COMMITTED sequence — slot p holds the MTP block's K/V for committed token p
368/// (roped p+1, the chain's rope convention), so the draft chain's self-attention sees the FULL
369/// committed history instead of only the current round's 1..K+1 chain tokens (the reference
370/// engine's "mtp_update" design). Entries come from two sources:
371/// - chain appends: accepted positions KEEP their chain-computed entries (embedding exact,
372/// hidden chain-approximate — the reference engine accepts the same);
373/// - `mtp_kv_fill` batches: prompt positions + the last-draft position on full accept, computed
374/// from EXACT trunk hiddens (K/V-only MTP-block pass, no attention/FFN/lm_head).
375/// Rejected drafts / p-min extras / pseudo-seed appends are all discarded by the round-start
376/// `set_len` truncation (the KvLayer len mechanism — §C rollback for the draft side).
377/// Multi-turn spec-decode session (2026-07-05): trunk Cache + persistent MTP draft scratch +
378/// the committed token list, alive across generate_spec_session calls. Turn N+1 primes ONLY its
379/// suffix (chunked continuation prime over the quantized past) and mtp_kv_fill's its suffix rows,
380/// then the round loop runs unchanged. `last_h` carries the pre-output_norm hidden of the last
381/// committed row across turns (the predecessor-pairing seed + fill anchor).
382/// Per-request sampling config for the sampled-spec serve path.
383#[derive(Clone, Copy, Debug)]
384pub struct SpecSampling {
385 pub temp: f32,
386 pub seed: u64,
387 pub top_k: i32, // 0 = off
388 pub top_p: f32, // 1.0 = off
389 pub min_p: f32, // 0.0 = off
390 pub penalty_last_n: usize, // 0 = penalties off
391 pub penalty_repeat: f32,
392 pub penalty_freq: f32,
393 pub penalty_present: f32,
394}
395
396/// Tracked draft positions for [`SpecTelemetry`] (serve K defaults to 3; the run-spec gate
397/// sweeps K=1..8, and MEMRA_SPEC_CAPMAX defaults to 7 — 8 covers every tuned config).
398pub const SPEC_TELEM_POS: usize = 8;
399
400/// Always-on per-draft-position acceptance telemetry (lane/accept-telemetry, 2026-08-05 —
401/// the llama.cpp #26389 / vLLM spec-decode counter schema, upstream-sweeps 2026-08-05).
402/// Lives on the [`SpecSession`] and accumulates across bursts; the serve worker diffs a
403/// stashed copy per burst for its per-model /metrics aggregation and per-request usage.
404/// Same normalization as the `[spec-stats]` line: p-min-discarded chain tokens are counted
405/// in NEITHER drafted nor accepted.
406#[derive(Clone, Copy, Default, Debug)]
407pub struct SpecTelemetry {
408 /// verify rounds completed (a round-stream burst counts each of its M rounds).
409 pub rounds: u64,
410 /// tokens drafted / accepted across all rounds.
411 pub drafted: u64,
412 pub accepted: u64,
413 /// how often draft position j (0-based within a round's chain) was offered / accepted.
414 /// Positions >= SPEC_TELEM_POS are untracked (totals still count them). The opt-in
415 /// round-stream arm (MEMRA_SPEC_STREAM=1) reads back only totals, so under it these
416 /// arrays cover the standard-path rounds only and their sums may undercount the totals.
417 pub pos_drafted: [u64; SPEC_TELEM_POS],
418 pub pos_accepted: [u64; SPEC_TELEM_POS],
419}
420
421impl SpecTelemetry {
422 /// Fieldwise `self - prev` — the worker's per-burst delta off a copy stashed before the
423 /// burst call. Saturating: a caller diffing against the wrong snapshot gets zeros, not
424 /// a wrapped counter.
425 pub fn delta_since(&self, prev: &SpecTelemetry) -> SpecTelemetry {
426 let mut d = SpecTelemetry {
427 rounds: self.rounds.saturating_sub(prev.rounds),
428 drafted: self.drafted.saturating_sub(prev.drafted),
429 accepted: self.accepted.saturating_sub(prev.accepted),
430 ..Default::default()
431 };
432 for j in 0..SPEC_TELEM_POS {
433 d.pos_drafted[j] = self.pos_drafted[j].saturating_sub(prev.pos_drafted[j]);
434 d.pos_accepted[j] = self.pos_accepted[j].saturating_sub(prev.pos_accepted[j]);
435 }
436 d
437 }
438 /// Fieldwise `self += d` — the worker's per-model aggregation.
439 pub fn merge(&mut self, d: &SpecTelemetry) {
440 self.rounds += d.rounds;
441 self.drafted += d.drafted;
442 self.accepted += d.accepted;
443 for j in 0..SPEC_TELEM_POS {
444 self.pos_drafted[j] += d.pos_drafted[j];
445 self.pos_accepted[j] += d.pos_accepted[j];
446 }
447 }
448
449 /// Mean accepted draft-prefix length per verify round (tau).
450 pub fn tau(&self) -> f64 {
451 if self.rounds > 0 {
452 self.accepted as f64 / self.rounds as f64
453 } else {
454 0.0
455 }
456 }
457}
458
459/// Session-lifetime atomic acceptance counters. The verifier records only after the greedy or
460/// rejection-sampling walk has resolved on the host, so these relaxed increments add no GPU
461/// launch, synchronization, allocation, or ordering dependency to the numeric path.
462struct SpecTelemetryCounters {
463 rounds: AtomicU64,
464 drafted: AtomicU64,
465 accepted: AtomicU64,
466 pos_drafted: [AtomicU64; SPEC_TELEM_POS],
467 pos_accepted: [AtomicU64; SPEC_TELEM_POS],
468}
469
470impl Default for SpecTelemetryCounters {
471 fn default() -> Self {
472 Self {
473 rounds: AtomicU64::new(0),
474 drafted: AtomicU64::new(0),
475 accepted: AtomicU64::new(0),
476 pos_drafted: std::array::from_fn(|_| AtomicU64::new(0)),
477 pos_accepted: std::array::from_fn(|_| AtomicU64::new(0)),
478 }
479 }
480}
481
482impl SpecTelemetryCounters {
483 fn record_round(&self, drafted: usize, accepted: usize) {
484 debug_assert!(accepted <= drafted);
485 self.rounds.fetch_add(1, Ordering::Relaxed);
486 self.drafted.fetch_add(drafted as u64, Ordering::Relaxed);
487 self.accepted.fetch_add(accepted as u64, Ordering::Relaxed);
488 for counter in self.pos_drafted.iter().take(drafted) {
489 counter.fetch_add(1, Ordering::Relaxed);
490 }
491 for counter in self.pos_accepted.iter().take(accepted) {
492 counter.fetch_add(1, Ordering::Relaxed);
493 }
494 }
495
496 /// Round-stream keeps each round's accept length on device; retain exact scalar totals while
497 /// leaving the per-position arrays untouched, matching the pre-existing telemetry contract.
498 fn record_totals(&self, rounds: usize, drafted: usize, accepted: usize) {
499 self.rounds.fetch_add(rounds as u64, Ordering::Relaxed);
500 self.drafted.fetch_add(drafted as u64, Ordering::Relaxed);
501 self.accepted.fetch_add(accepted as u64, Ordering::Relaxed);
502 }
503
504 fn snapshot(&self) -> SpecTelemetry {
505 SpecTelemetry {
506 rounds: self.rounds.load(Ordering::Relaxed),
507 drafted: self.drafted.load(Ordering::Relaxed),
508 accepted: self.accepted.load(Ordering::Relaxed),
509 pos_drafted: std::array::from_fn(|j| self.pos_drafted[j].load(Ordering::Relaxed)),
510 pos_accepted: std::array::from_fn(|j| self.pos_accepted[j].load(Ordering::Relaxed)),
511 }
512 }
513}
514
515pub struct SpecSession {
516 pub(crate) cache: Cache,
517 pub(crate) scratch: MtpScratch,
518 /// Every token whose state the caches hold, in order (prompt turns + generated), INCLUDING
519 /// overshoot: spec commits accepted drafts past max_new; those rows are in the caches, so the
520 /// session must count them. Callers render output from this, not from their own echo.
521 pub committed: Vec<u32>,
522 /// Pre-output_norm hidden of the LAST committed row (device). None before the first turn.
523 pub(crate) last_h: Option<CudaSlice<f32>>,
524 /// Greedy argmax predicting the token AFTER committed.last() (from the last turn's final
525 /// logits). Fuels empty-suffix continuation bursts (serve): the next turn emits this token
526 /// first, feeds it, and the round loop resumes without any prime. None before the first turn.
527 pub next_pred: Option<u32>,
528 /// SAMPLED-SPEC stream continuity across bursts: Philox event counters persist here so a
529 /// session's randomness never repeats between generate_spec_session calls. (0,0) at admit.
530 pub sctr: u32,
531 pub uctr: u32,
532 /// PERSISTENT DRAFT-GRAPH CONTEXT (2026-08-01, the serve-burst fixed-cost fix): the captured
533 /// draft graph(s) + every device I/O buffer they bake, carried ACROSS generate_spec_session
534 /// calls. Before this, every serve burst re-captured the draft graph (2 warmup forwards +
535 /// instantiate) — measured ~16ms/burst on H100 q27 (MEMRA_SPEC_BURST sweep,
536 /// research/spec-serving-20260801). None before the first turn; error paths drop it
537 /// (next burst recaptures — serve retires errored sessions anyway).
538 pub(crate) draft_ctx: Option<DraftGraphCtx>,
539 /// PENDING-CARRY across bursts (2026-08-01, the serve burst-boundary fix): the bonus token
540 /// emitted by the last round but NOT committed to the caches. The old tail committed it with
541 /// a solo T=1 trunk pass (+ draft fill), and the next burst's setup fed the stashed next_pred
542 /// with ANOTHER solo pass — 2x ~11.5ms/burst measured on H100 q27 ([spec-setup] trace).
543 /// Carrying it lets the next empty-suffix greedy burst consume it as round-0 verify col 0,
544 /// exactly like a mid-burst full-accept boundary (no solo passes). INVARIANT: when set,
545 /// `committed` (== cache rows) EXCLUDES this token although it was already emitted in the
546 /// last burst's output, and `last_h` holds the hidden of the last COMMITTED row (its
547 /// predecessor — the chain-seed/fill anchor). `next_pred` is None (unknown without the
548 /// commit pass). Non-empty-suffix or sampled turns must flush first (spec_flush_pending);
549 /// generate_spec_session_sampled does this at entry, and serve parks only flushed sessions.
550 pub pending_tok: Option<u32>,
551 /// SESSION-AFFINITY TURN CHECKPOINT (lane/session-affinity, 2026-08-05): the state at this
552 /// turn's PROMPT-END boundary, retained so a later turn can REWIND here. See
553 /// [`SpecCheckpoint`]. Refreshed by every non-empty prime; None until the first one, and on
554 /// a rig too tight to hold it (a failed capture is silent — resume just isn't available).
555 pub(crate) turn_ckpt: Option<SpecCheckpoint>,
556 /// Session-lifetime acceptance telemetry. Relaxed atomics update at the host-side round
557 /// accounting the loop already does — no syncs, no allocation. NOTE a
558 /// pool-resumed session carries the PREVIOUS requests' counts; per-request consumers
559 /// diff with [`SpecTelemetry::delta_since`] around each burst.
560 telem: SpecTelemetryCounters,
561 /// PREFIX-CACHE publication request (lane/spec-prefix-cache): worker sets this to the
562 /// miss-LCP boundary before a cold burst; the prime captures at exactly that split (it must
563 /// coincide with the burst's `prime_split` or no capture happens). One-shot: consumed by the
564 /// prime, result lands in `boundary_capture`.
565 pub capture_at: Option<usize>,
566 /// The capture the last cold prime produced (see [`SpecBoundaryCapture`]). Worker takes it
567 /// post-burst to assemble the prefix entry. A failed capture is silent, like `turn_ckpt` —
568 /// publication just isn't available for that request.
569 pub boundary_capture: Option<SpecBoundaryCapture>,
570}
571impl SpecSession {
572 /// Context capacity of the session's caches (the server's ContextFull guard).
573 pub fn cache_max_ctx(&self) -> usize {
574 self.cache.max_ctx
575 }
576 /// Read access to the live trunk cache (lane/spec-prefix-cache): the worker slices
577 /// full-attn KV rows `[0..capture.pos)` out of it when publishing a boundary capture —
578 /// those rows are append-only for the session's lifetime (rollbacks never truncate below
579 /// the prime boundary), so no copy was taken at prime time.
580 pub fn cache_ref(&self) -> &Cache {
581 &self.cache
582 }
583 /// Read access to the persistent draft-scratch plane (lane/spec-on-cache-hit): the
584 /// worker slices rows `[0..capture.pos)` when publishing a boundary capture, exactly
585 /// like the trunk KV — draft rows below the prompt end are append-only for the
586 /// session's lifetime (the prime fill wrote them once; rollbacks reset `len_d` to the
587 /// committed length, never below the prime boundary, and the true-hidden refresh
588 /// rewrites generated positions only). Returns `(k, v, k_tok_bytes, v_tok_bytes)`.
589 /// None when the scratch is ring-backed (Step35 SWA — physical rows are not
590 /// prefix-addressable; the prefix cache already refuses that class end to end).
591 pub fn draft_plane_ref(&self) -> Option<(&CudaSlice<u8>, &CudaSlice<u8>, usize, usize)> {
592 if self.scratch.kv.ring.is_some() {
593 return None;
594 }
595 Some((
596 &self.scratch.kv.k,
597 &self.scratch.kv.v,
598 self.scratch.kv.k_tok_bytes,
599 self.scratch.kv.v_tok_bytes,
600 ))
601 }
602 /// Snapshot the session's process-local acceptance counters for per-burst diffing.
603 pub fn telemetry(&self) -> SpecTelemetry {
604 self.telem.snapshot()
605 }
606 /// Committed position this session can REWIND to (its retained prompt-end boundary), if any.
607 /// A request whose prompt matches `committed[..pos]` exactly can resume from here — see
608 /// `spec_rewind_to_checkpoint`.
609 pub fn rewind_pos(&self) -> Option<usize> {
610 self.turn_ckpt.as_ref().map(|c| c.pos)
611 }
612 /// Whether every ring-backed trunk/draft row needed by the retained checkpoint is resident.
613 pub fn rewind_is_resident(&self) -> bool {
614 self.turn_ckpt.as_ref().is_some_and(|ckpt| {
615 self.cache.can_rollback(&ckpt.snap, 0) && self.scratch.can_rewind_to(ckpt.pos)
616 })
617 }
618 /// Is this session in the DEMOTION-READY shape (see [`SpecSession::into_demoted`])?
619 /// `false` means a carried pending must be flushed first (`spec_flush_pending`), or the
620 /// session has never run a turn and has no prediction to hand over.
621 pub fn demote_ready(&self) -> bool {
622 self.pending_tok.is_none() && self.next_pred.is_some()
623 }
624 /// Does this session hold a carried pending bonus (flush required before a handoff/park)?
625 pub fn has_pending(&self) -> bool {
626 self.pending_tok.is_some()
627 }
628 /// Committed row count == cache rows (the session invariant), for the caller's own
629 /// `fed`-length cross-check at a handoff boundary.
630 pub fn committed_len(&self) -> usize {
631 self.committed.len()
632 }
633 /// DEMOTION HANDOFF (lane/spec-gate, 2026-08-07): consume this session and hand its trunk
634 /// cache + next-token prediction to the plain batched-decode path.
635 ///
636 /// WHY THIS IS EXACT (greedy). The invariant at a burst boundary is `cache.pos ==
637 /// committed.len()`: every committed row has trunk KV + recurrent state, exactly as a plain
638 /// tokenwise prime of the same `committed` sequence would have left it (that is the
639 /// session-tail contract, and the same property `spec_rewind_to_checkpoint` and the reuse
640 /// pool already rely on). `next_pred` is the argmax of the verify's logits for the LAST
641 /// committed row — and verify-column logits are bit-identical to plain decode's logits at
642 /// that position, because `matmul_decode_exact` bit-identity IS the basis of the greedy
643 /// accept walk. So handing (cache, next_pred) to the batched path continues the stream from
644 /// a state indistinguishable from one the batched path produced itself: the batched tick
645 /// emits `next_pred`, feeds it into this same cache, and decodes on.
646 ///
647 /// `None` when the session is not in the handoff shape — a carried pending (its bonus row is
648 /// NOT in the cache, so `spec_flush_pending` must commit it first) or no `next_pred` yet
649 /// (never bursted). Callers must not force it: a half-committed cache handed to the batched
650 /// path would silently skip a token.
651 ///
652 /// The MTP draft scratch, the persistent draft-graph context and the turn checkpoint are
653 /// DROPPED here (freeing their VRAM): the batched path never drafts, and this handoff is
654 /// one-way by design — there is no cheap symmetric re-promotion (rebuilding the draft KV
655 /// would mean an `mtp_kv_fill` over the whole committed history).
656 pub fn into_demoted(self) -> Option<(Cache, u32)> {
657 if self.pending_tok.is_some() {
658 return None;
659 }
660 let np = self.next_pred?;
661 debug_assert_eq!(
662 self.cache.pos,
663 self.committed.len(),
664 "demotion handoff: cache rows != committed tokens"
665 );
666 Some((self.cache, np))
667 }
668 /// Pool-resume hook (audit Q2): clear the parked draft-graph failure memoization so a
669 /// NEW request resuming this session gets one fresh capture chance — a transient-pressure
670 /// capture failure must not persist for the pool's whole lifetime (the TRT #16072 class).
671 /// Logs once iff a flag was actually set; a no-fallback resume is silent and free.
672 pub fn reset_graph_fallback_on_resume(&mut self) {
673 if let Some(line) = self
674 .draft_ctx
675 .as_mut()
676 .and_then(|c| c.failed.reset_on_resume())
677 {
678 eprintln!("{line}");
679 }
680 }
681}
682
683/// A session's PROMPT-END boundary state, the rewind target for session-affinity resume.
684///
685/// WHY THIS BOUNDARY, AND WHY IT IS THE ONLY ONE WORTH KEEPING. The rewrite class this lane
686/// exists for (a client that strips `<think>` blocks out of prior assistant turns) mutates the
687/// text the session GENERATED, never the prompt it was given. So turn N's prompt agrees with
688/// turn N-1's committed tokens up to almost exactly where turn N-1's generation began — the
689/// prompt-end boundary. Keeping a checkpoint there means the next turn re-primes only its own
690/// delta (the rewritten answer + the new user turn) instead of the whole conversation.
691///
692/// WHAT IT MUST HOLD. Full-attn KV is append-only and position-addressed, so rewinding it is a
693/// `len` truncation (no data). Linear-attn (GDN) conv/ssm state is mutated IN PLACE with no
694/// position index, so it must be a real device COPY — that copy is the entire reason a spec
695/// session could not previously rewind. The MTP draft scratch needs no copy either: its rows
696/// below the boundary were written by this turn's fill and are never revisited (the per-round
697/// true-hidden refresh only rewrites the CURRENT burst's committed positions), so rewinding it
698/// is also just a `len` reset. `last_h` is the hidden of the last row below the boundary — the
699/// predecessor-pairing anchor the next prime's fill reads for its first row.
700///
701/// COST: one `Cache::snapshot` per TURN, on a code path that already takes one per ROUND.
702pub(crate) struct SpecCheckpoint {
703 snap: crate::cache::CacheSnapshot,
704 /// Committed length at the boundary (== cache.pos there, the session invariant).
705 pos: usize,
706 /// Pre-output_norm hidden of row `pos - 1`.
707 last_h: CudaSlice<f32>,
708}
709
710/// PREFIX-CACHE BOUNDARY CAPTURE (lane/spec-prefix-cache, 2026-08-14): the state a spec session
711/// records at its cold-prime split so the WORKER can publish a cross-request prefix entry —
712/// the commit-gated-publication port (research/cache-spec-design-20260814/PORT-PLAN.md item 1).
713/// Only the pieces that are DESTROYED by continuing the prime need copies here: the in-place
714/// GDN conv/ssm states (via `Cache::snapshot`, same mechanism as [`SpecCheckpoint`]) and the
715/// boundary logits. Full-attn KV rows `[0..pos)` and draft-scratch rows `[0..pos)` are
716/// append-only for the session's lifetime (rollbacks never truncate below the prime boundary),
717/// so the worker slices those from the live caches post-burst instead of copying at prime time.
718pub struct SpecBoundaryCapture {
719 pub snap: crate::cache::CacheSnapshot,
720 /// Token boundary (== cache.pos at capture; == the worker's miss-LCP split).
721 pub pos: usize,
722 /// Full-vocab logits after the prefix prime — the entry's boundary logits.
723 pub logits: Vec<f32>,
724 /// Pre-output_norm trunk hidden of row `pos - 1` (lane/spec-on-cache-hit): the
725 /// predecessor-pairing anchor a RESTORED spec session's first suffix-fill row reads
726 /// (the `SpecSession::last_h` convention). Empty = unavailable (capture stays valid;
727 /// the fill's zeros row-0 fallback covers it at a bounded acceptance cost).
728 pub last_h: Vec<f32>,
729}
730
731/// D2H one hidden row out of a `[T, n_embd]` prime hidden stack — the boundary anchor a
732/// spec boundary capture carries for later restored-session fills. Failure is silent
733/// (`turn_ckpt` convention): the capture publishes without an anchor.
734fn capture_boundary_hidden(
735 e: &Engine,
736 h_rows: &CudaSlice<f32>,
737 pos: usize,
738 n_embd: usize,
739) -> Vec<f32> {
740 if pos == 0 || h_rows.len() < pos * n_embd {
741 return Vec::new();
742 }
743 let Ok(mut row) = e.uninit(n_embd) else {
744 return Vec::new();
745 };
746 if e.copy_view_into(
747 &mut row,
748 0,
749 &h_rows.slice((pos - 1) * n_embd..pos * n_embd),
750 n_embd,
751 )
752 .is_err()
753 {
754 return Vec::new();
755 }
756 e.dtoh(&row).unwrap_or_default()
757}
758
759/// ROLLBACK DOOR for sampled BOUNDARY tokens (lane/sampled-spec-quality, 2026-08-19).
760/// Default ON: the token a burst emits at its own boundary is drawn from the request's
761/// sampler. `MEMRA_SPEC_SAMPLED_BOUNDARY=0` restores the pre-lane posture (an ARGMAX at
762/// every boundary) without touching greedy, which is byte-unaffected either way.
763pub fn spec_sampled_boundary_on() -> bool {
764 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
765 *ON.get_or_init(|| std::env::var("MEMRA_SPEC_SAMPLED_BOUNDARY").as_deref() != Ok("0"))
766}
767
768/// ROLLBACK DOOR for SESSION-SPANNING penalty history (lane/sampled-spec-quality).
769/// Default ON: `pen_hist` is seeded from the session's committed tail, so repetition /
770/// frequency / presence penalties see the whole stream. `MEMRA_SPEC_PEN_SESSION=0`
771/// restores the pre-lane posture (each burst restarts the window from its own prompt
772/// slice, i.e. from NOTHING on a continuation burst) — and with the door shut the worker
773/// must keep refusing penalized sampled prefix-cache restores, because the restored
774/// session's continuation burst is handed no prompt slice at all.
775pub fn spec_pen_session_on() -> bool {
776 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
777 *ON.get_or_init(|| std::env::var("MEMRA_SPEC_PEN_SESSION").as_deref() != Ok("0"))
778}
779
780/// ROLLBACK DOOR for extended-entry publication from a RESTORED session
781/// (lane/sampled-spec-quality, Item 3). Default ON: a converted prefix-cache hit that fed a
782/// suffix captures its own prompt-end boundary so the NEXT turn can hit a longer prefix.
783/// `MEMRA_SPEC_RESTORE_REPUBLISH=0` restores the pre-lane posture (a namespace learns exactly
784/// one boundary and never advances it). Whole-entry semantics only — the boundary is the
785/// restored session's own prompt end, so `entry_pos != fed_len` still refuses on the way in.
786pub fn spec_restore_republish_on() -> bool {
787 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
788 *ON.get_or_init(|| std::env::var("MEMRA_SPEC_RESTORE_REPUBLISH").as_deref() != Ok("0"))
789}
790
791/// Diagnostics: name every boundary token on stderr (`MEMRA_SPEC_BOUNDARY_TRACE=1`), with
792/// the argmax the pre-lane code would have emitted from the same row. This is how the
793/// lane MEASURES the boundary rate and the deviation rate instead of estimating them.
794fn spec_boundary_trace() -> bool {
795 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
796 *ON.get_or_init(|| std::env::var("MEMRA_SPEC_BOUNDARY_TRACE").as_deref() == Ok("1"))
797}
798
799/// llama-parity floor for the penalty window when the request does not ask for a bigger
800/// one (`repeat_last_n` default). The serve API arms `penalty_last_n = usize::MAX` for any
801/// non-identity penalty, so this floor only matters to explicit small windows and to the
802/// CLI env path.
803const PEN_WINDOW_FLOOR: usize = 64;
804
805/// CEILING on the penalty window, and it is a COST bound, not a semantic preference.
806/// `penalize_logits_f32` (cu/spec_sample.cu) dedups on device by having thread `i` scan
807/// `hist[0..i]`, so a pass is O(n_hist²) and it runs ~3x per verify round (the q rows, the
808/// p column, the bonus column). The serve API arms `penalty_last_n = usize::MAX` for ANY
809/// non-identity penalty — "the whole context", llama's `repeat_last_n = -1` — so an
810/// uncapped session window would put a 128k-token history through that kernel: ~1.7e10
811/// comparisons per pass, tens of ms per round, i.e. penalties would silently destroy decode
812/// throughput on exactly the long-context requests that most want them. 8192 keeps a pass
813/// at ~7e7 comparisons (tens of microseconds) while still being **128x wider than the
814/// pre-lane effective window** (64 prompt-tail tokens + whatever the current burst had
815/// generated). A request that genuinely needs a window beyond this wants host-side dedup +
816/// counts through a new kernel signature — a follow-up lane, named here rather than hidden.
817const PEN_WINDOW_MAX: usize = 8192;
818
819/// Seed a penalty window over the SESSION, not the burst (lane/sampled-spec-quality,
820/// Item 2). The window is the last `max(penalty_last_n, 64)` tokens of
821/// `session_committed ++ burst_prompt` — for a cold turn-1 burst (`session_committed`
822/// empty, default `penalty_last_n`) that is byte-identically the pre-lane
823/// `prompt.iter().rev().take(64).rev()`; for a continuation burst it is the stream the
824/// client actually asked us to penalize, where the pre-lane code had NOTHING.
825fn pen_window_seed(
826 session_committed: &[u32],
827 burst_prompt: &[u32],
828 penalty_last_n: usize,
829) -> Vec<u32> {
830 let win = penalty_last_n.clamp(PEN_WINDOW_FLOOR, PEN_WINDOW_MAX);
831 let take_prompt = burst_prompt.len().min(win);
832 let take_sess = (win - take_prompt).min(session_committed.len());
833 let mut hist = Vec::with_capacity(take_sess + take_prompt);
834 hist.extend_from_slice(&session_committed[session_committed.len() - take_sess..]);
835 hist.extend_from_slice(&burst_prompt[burst_prompt.len() - take_prompt..]);
836 hist
837}
838
839/// Draw a BOUNDARY token from the target distribution the request asked for
840/// (lane/sampled-spec-quality, Item 1) — the fix for "sampled spec emits an ARGMAX token at
841/// every burst boundary".
842///
843/// WHY THIS EXISTS. A spec burst's first emitted token is not produced by the accept walk:
844/// it comes off a logits row that already exists (the prime's last row on a cold burst; the
845/// row after the last committed token on a continuation burst; the prefix-cache entry's
846/// boundary row on a restored one). Pre-lane that token was `argmax` in BOTH sampling
847/// regimes, so a sampled stream took a greedy token once per burst — measured, not
848/// estimated, in research/spec-cache-20260818/SAMPLED-QUALITY.md. At temperature > 0 the
849/// customer asked for a sampled token, so this draws one.
850///
851/// THE PROGRAM IS THE FULL-ACCEPT BONUS'S PROGRAM, deliberately: penalize the row (over the
852/// session's window), take this row's OWN filter stats (the sampfix-20260805 law — stats
853/// from a neighbour row mis-scale every `e0` and can wipe the row to token 0), gumbel-perturb
854/// with the session's Philox stream at `*sctr`, argmax the perturbed row. Reusing the bonus's
855/// composition means `sample_check`'s distributional oracle covers this draw too, and the
856/// boundary token is drawn from the same filtered/penalized `p` the accept walk targets.
857///
858/// THE STREAM IS THE SESSION'S, NOT A FRESH ONE. `sctr` is the caller's live counter and is
859/// advanced by exactly one, so a boundary draw consumes the next value in the same Philox
860/// stream the accept walk uses — never a second, independently seeded stream (which would be
861/// a new distributional bug: two streams from one seed correlate wherever their counters
862/// collide). That also makes a restored session's boundary draw at `sctr == 0` bit-identical
863/// to the cold session's own first draw from the same logits row, which is what preserves the
864/// sampled-hit lane's per-seed hit==cold byte identity.
865#[allow(clippy::too_many_arguments)]
866pub fn sample_boundary_token_dev(
867 e: &Engine,
868 logits: &CudaSlice<f32>,
869 n_vocab: usize,
870 sp: &SpecSampling,
871 pen_hist: &[u32],
872 sctr: &mut u32,
873 site: &str,
874) -> Result<u32, Box<dyn std::error::Error>> {
875 debug_assert!(
876 sp.temp > 0.0,
877 "boundary sampling is the sampled regime only"
878 );
879 // Own copy: penalize_logits mutates in place and the caller's row is live state
880 // (prime_logits back the constrained recompute; last_col_logits backs round 0's accept).
881 let mut col = e.zeros(n_vocab)?;
882 e.copy_into(&mut col, 0, logits, n_vocab)?;
883 let pen_on = sp.penalty_last_n > 0
884 && (sp.penalty_repeat != 1.0 || sp.penalty_freq != 0.0 || sp.penalty_present != 0.0);
885 if pen_on && !pen_hist.is_empty() {
886 // window trim mirrors the round loop's own upload (`pen_hist[w0..]`), cap included.
887 let w0 = pen_hist
888 .len()
889 .saturating_sub(sp.penalty_last_n.min(PEN_WINDOW_MAX));
890 let hist = &pen_hist[w0..];
891 let hd = e.htod_u32_v(hist)?;
892 e.penalize_logits(
893 &mut col,
894 &hd,
895 hist.len(),
896 sp.penalty_repeat,
897 sp.penalty_freq,
898 sp.penalty_present,
899 n_vocab,
900 )?;
901 }
902 let rows0 = e.htod_i32(&[0])?;
903 let (mut th_d, mut z_d, mut mx_d) = (e.zeros(1)?, e.zeros(1)?, e.zeros(1)?);
904 e.filter_stats(
905 &col, n_vocab, &rows0, &mut th_d, &mut z_d, &mut mx_d, n_vocab, 1, sp.temp, sp.top_k,
906 sp.top_p, sp.min_p,
907 )?;
908 let (th, mx) = (e.dtoh(&th_d)?[0], e.dtoh(&mx_d)?[0]);
909 let mut perturb = e.zeros(n_vocab)?;
910 e.gumbel_perturb_filtered(&col, &mut perturb, n_vocab, sp.seed, *sctr, sp.temp, mx, th)?;
911 *sctr = sctr.wrapping_add(1);
912 let td = e.argmax_token_device(&perturb, n_vocab)?;
913 let tok = e.dtoh_u32_one(&td)?;
914 if spec_boundary_trace() {
915 // the pre-lane token, from the SAME row, so the deviation rate is measurable.
916 let raw = e.argmax_token_device(logits, n_vocab)?;
917 let greedy = e.dtoh_u32_one(&raw)?;
918 eprintln!(
919 "[spec-boundary] site={site} sampled={tok} argmax={greedy} \
920 deviates={} temp={} sctr={}",
921 (tok != greedy) as u8,
922 sp.temp,
923 sctr.wrapping_sub(1),
924 );
925 }
926 Ok(tok)
927}
928
929/// Host-row twin of [`sample_boundary_token_dev`] (the prime / feed / entry rows arrive as
930/// host `Vec<f32>`).
931#[allow(clippy::too_many_arguments)]
932pub fn sample_boundary_token(
933 e: &Engine,
934 logits: &[f32],
935 sp: &SpecSampling,
936 pen_hist: &[u32],
937 sctr: &mut u32,
938 site: &str,
939) -> Result<u32, Box<dyn std::error::Error>> {
940 let n_vocab = logits.len();
941 let d = e.htod(logits)?;
942 sample_boundary_token_dev(e, &d, n_vocab, sp, pen_hist, sctr, site)
943}
944
945struct SpecPipeTraceClock {
946 pair: usize,
947 started: std::time::Instant,
948}
949
950#[derive(Clone)]
951struct SpecPipeTraceCtx {
952 clock: std::sync::Arc<SpecPipeTraceClock>,
953 round: usize,
954 lane: usize,
955}
956
957struct SpecPipeTraceMarker {
958 trace: SpecPipeTraceCtx,
959 phase: &'static str,
960 edge: &'static str,
961 slot: Option<usize>,
962}
963
964unsafe extern "C" fn spec_pipe_trace_marker(raw: *mut std::ffi::c_void) {
965 let marker = unsafe { Box::from_raw(raw.cast::<SpecPipeTraceMarker>()) };
966 let lane = if marker.trace.lane == 0 { "A" } else { "B" };
967 let slot = marker
968 .slot
969 .map(|v| v.to_string())
970 .unwrap_or_else(|| "-".into());
971 let t_ms = marker.trace.clock.started.elapsed().as_secs_f64() * 1e3;
972 use std::io::Write as _;
973 let stderr = std::io::stderr();
974 let mut stderr = stderr.lock();
975 let _ = writeln!(
976 stderr,
977 "[spec-pipe-timeline] pair={} round={} lane={lane} phase={} edge={} \
978 slot={slot} t_ms={t_ms:.3}",
979 marker.trace.clock.pair, marker.trace.round, marker.phase, marker.edge,
980 );
981}
982
983fn enqueue_spec_pipe_trace_marker(
984 stream: &cudarc::driver::CudaStream,
985 trace: Option<&SpecPipeTraceCtx>,
986 phase: &'static str,
987 edge: &'static str,
988 slot: Option<usize>,
989) -> Result<(), Box<dyn std::error::Error>> {
990 let Some(trace) = trace else {
991 return Ok(());
992 };
993 let marker = Box::new(SpecPipeTraceMarker {
994 trace: trace.clone(),
995 phase,
996 edge,
997 slot,
998 });
999 let raw = Box::into_raw(marker);
1000 let result = unsafe {
1001 cudarc::driver::result::stream::launch_host_function(
1002 stream.cu_stream(),
1003 spec_pipe_trace_marker,
1004 raw.cast(),
1005 )
1006 };
1007 if let Err(err) = result {
1008 unsafe {
1009 drop(Box::from_raw(raw));
1010 }
1011 return Err(err.into());
1012 }
1013 Ok(())
1014}
1015
1016#[derive(Default)]
1017struct SpecPipeProgress {
1018 setup_done: [bool; 2],
1019 draft_done: [usize; 2],
1020 stage0_done: [usize; 2],
1021 verify_done: [usize; 2],
1022 accept_done: [usize; 2],
1023 finished: [bool; 2],
1024 aborted: bool,
1025}
1026
1027/// Host-side issue coordinator for the reduced two-session speculative pipeline. Each session
1028/// keeps its existing call stack and round locals; this object only orders phase entry. The
1029/// primary mutex spans whole draft/accept/tail issue regions so Engine's single-stream scratch
1030/// cannot be interleaved by the two host threads.
1031struct SpecPipeSync {
1032 progress: std::sync::Mutex<SpecPipeProgress>,
1033 changed: std::sync::Condvar,
1034 primary: std::sync::Mutex<()>,
1035 trace: Option<std::sync::Arc<SpecPipeTraceClock>>,
1036}
1037
1038impl SpecPipeSync {
1039 fn new() -> Self {
1040 static TRACE_PAIR: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
1041 let trace = (std::env::var("MEMRA_SPEC_PIPE_TRACE").as_deref() == Ok("1")).then(|| {
1042 std::sync::Arc::new(SpecPipeTraceClock {
1043 pair: TRACE_PAIR.fetch_add(1, std::sync::atomic::Ordering::Relaxed) + 1,
1044 started: std::time::Instant::now(),
1045 })
1046 });
1047 Self {
1048 progress: std::sync::Mutex::new(SpecPipeProgress::default()),
1049 changed: std::sync::Condvar::new(),
1050 primary: std::sync::Mutex::new(()),
1051 trace,
1052 }
1053 }
1054}
1055
1056#[derive(Clone)]
1057struct SpecPipeLane {
1058 sync: std::sync::Arc<SpecPipeSync>,
1059 lane: usize,
1060}
1061
1062impl SpecPipeLane {
1063 fn peer(&self) -> usize {
1064 1 - self.lane
1065 }
1066
1067 fn aborted() -> Box<dyn std::error::Error> {
1068 "paired speculative peer aborted".into()
1069 }
1070
1071 fn trace(&self, round: usize) -> Option<SpecPipeTraceCtx> {
1072 self.sync.trace.as_ref().map(|clock| SpecPipeTraceCtx {
1073 clock: clock.clone(),
1074 round,
1075 lane: self.lane,
1076 })
1077 }
1078
1079 fn setup_begin(&self) -> Result<(), Box<dyn std::error::Error>> {
1080 let mut p = self.sync.progress.lock().unwrap();
1081 while !p.aborted && self.lane == 1 && !p.setup_done[0] && !p.finished[0] {
1082 p = self.sync.changed.wait(p).unwrap();
1083 }
1084 if p.aborted {
1085 Err(Self::aborted())
1086 } else {
1087 Ok(())
1088 }
1089 }
1090
1091 fn setup_end(&self) {
1092 let mut p = self.sync.progress.lock().unwrap();
1093 p.setup_done[self.lane] = true;
1094 self.sync.changed.notify_all();
1095 }
1096
1097 fn draft_begin(
1098 &self,
1099 round: usize,
1100 ) -> Result<std::sync::MutexGuard<'_, ()>, Box<dyn std::error::Error>> {
1101 let peer = self.peer();
1102 let mut p = self.sync.progress.lock().unwrap();
1103 loop {
1104 if p.aborted {
1105 return Err(Self::aborted());
1106 }
1107 let setup_ready =
1108 (p.setup_done[0] || p.finished[0]) && (p.setup_done[1] || p.finished[1]);
1109 let prior_ready = p.accept_done[self.lane] >= round
1110 && (p.accept_done[peer] >= round || p.finished[peer]);
1111 let turn_ready = if self.lane == 0 {
1112 true
1113 } else {
1114 p.draft_done[0] > round || p.finished[0]
1115 };
1116 if setup_ready && prior_ready && turn_ready {
1117 break;
1118 }
1119 p = self.sync.changed.wait(p).unwrap();
1120 }
1121 drop(p);
1122 Ok(self.sync.primary.lock().unwrap())
1123 }
1124
1125 fn draft_end(&self, round: usize) {
1126 let mut p = self.sync.progress.lock().unwrap();
1127 p.draft_done[self.lane] = round + 1;
1128 self.sync.changed.notify_all();
1129 }
1130
1131 /// Admit stage 0 and return whether this lane owns the interval's one reverse fence.
1132 /// Lane B releases as soon as lane A has issued its boundary TX, not after A's full body.
1133 fn stage0_begin(&self, round: usize) -> Result<bool, Box<dyn std::error::Error>> {
1134 let peer = self.peer();
1135 let mut p = self.sync.progress.lock().unwrap();
1136 loop {
1137 if p.aborted {
1138 return Err(Self::aborted());
1139 }
1140 let ready = if self.lane == 0 {
1141 p.draft_done[0] > round && (p.draft_done[1] > round || p.finished[1])
1142 } else {
1143 p.draft_done[1] > round && (p.stage0_done[0] > round || p.finished[0])
1144 };
1145 if ready {
1146 return Ok(self.lane == 0 || p.finished[peer]);
1147 }
1148 p = self.sync.changed.wait(p).unwrap();
1149 }
1150 }
1151
1152 fn stage0_end(&self, round: usize) {
1153 let mut p = self.sync.progress.lock().unwrap();
1154 p.stage0_done[self.lane] = round + 1;
1155 self.sync.changed.notify_all();
1156 }
1157
1158 /// Stage 1 is single-owner per engine. A proceeds immediately after its own ticket; B waits
1159 /// for A's full stage1/head issue so only A.S1 and B.S0 can overlap.
1160 fn stage1_begin(&self, round: usize) -> Result<(), Box<dyn std::error::Error>> {
1161 let mut p = self.sync.progress.lock().unwrap();
1162 while !p.aborted
1163 && !(p.stage0_done[self.lane] > round
1164 && (self.lane == 0 || p.verify_done[0] > round || p.finished[0]))
1165 {
1166 p = self.sync.changed.wait(p).unwrap();
1167 }
1168 if p.aborted {
1169 Err(Self::aborted())
1170 } else {
1171 Ok(())
1172 }
1173 }
1174
1175 fn verify_end(&self, round: usize) {
1176 let mut p = self.sync.progress.lock().unwrap();
1177 p.verify_done[self.lane] = round + 1;
1178 self.sync.changed.notify_all();
1179 }
1180
1181 fn accept_begin(
1182 &self,
1183 round: usize,
1184 ) -> Result<std::sync::MutexGuard<'_, ()>, Box<dyn std::error::Error>> {
1185 let mut p = self.sync.progress.lock().unwrap();
1186 loop {
1187 if p.aborted {
1188 return Err(Self::aborted());
1189 }
1190 let ready = if self.lane == 0 {
1191 p.verify_done[0] > round && (p.verify_done[1] > round || p.finished[1])
1192 } else {
1193 p.verify_done[1] > round && (p.accept_done[0] > round || p.finished[0])
1194 };
1195 if ready {
1196 break;
1197 }
1198 p = self.sync.changed.wait(p).unwrap();
1199 }
1200 drop(p);
1201 Ok(self.sync.primary.lock().unwrap())
1202 }
1203
1204 fn accept_end(&self, round: usize) {
1205 let mut p = self.sync.progress.lock().unwrap();
1206 p.accept_done[self.lane] = round + 1;
1207 self.sync.changed.notify_all();
1208 }
1209
1210 fn primary(&self) -> std::sync::MutexGuard<'_, ()> {
1211 self.sync.primary.lock().unwrap()
1212 }
1213
1214 fn finish(&self, failed: bool) {
1215 let mut p = self.sync.progress.lock().unwrap();
1216 p.finished[self.lane] = true;
1217 p.aborted |= failed;
1218 self.sync.changed.notify_all();
1219 }
1220}
1221
1222struct SpecPipeFinish<'a> {
1223 lane: &'a SpecPipeLane,
1224 closed: bool,
1225}
1226
1227impl<'a> SpecPipeFinish<'a> {
1228 fn new(lane: &'a SpecPipeLane) -> Self {
1229 Self {
1230 lane,
1231 closed: false,
1232 }
1233 }
1234
1235 fn close(&mut self, failed: bool) {
1236 self.lane.finish(failed);
1237 self.closed = true;
1238 }
1239}
1240
1241impl Drop for SpecPipeFinish<'_> {
1242 fn drop(&mut self) {
1243 if !self.closed {
1244 self.lane.finish(true);
1245 }
1246 }
1247}
1248
1249/// Scoped transfer of one exclusively-borrowed session to the second host issue thread.
1250/// `CudaGraph` is not marked Send by cudarc because its raw driver handles carry no automatic
1251/// trait. CUDA driver graph handles are context-scoped rather than OS-thread-affine; the caller
1252/// binds that context before touching the session, joins before returning, and never aliases the
1253/// pointer. Keep this exception local to the experimental pair call instead of marking the public
1254/// session type Send.
1255struct SpecPipeSessionPtr(*mut SpecSession);
1256
1257unsafe impl Send for SpecPipeSessionPtr {}
1258
1259impl SpecPipeSessionPtr {
1260 unsafe fn get_mut(&mut self) -> &mut SpecSession {
1261 unsafe { &mut *self.0 }
1262 }
1263}
1264
1265/// Per-session persistent draft-graph context: the captured CUDA graph(s) plus the device
1266/// buffers whose POINTERS the capture bakes. Reuse legality: the greedy capture bakes only
1267/// session-stable pointers (the session's own MtpScratch KV — allocated once, never realloc'd;
1268/// the model's resident embedding; the process-wide OnceLock p_min) and the g_* buffers held
1269/// HERE — so one capture serves the session's whole lifetime. The sampled capture additionally
1270/// bakes (seed, temp) as capture-time constants and needs k q-slots — keyed by `s_key`, dropped
1271/// and recaptured when a pool-resumed request changes them. `*_failed` memoizes a failed capture
1272/// so the eager fallback doesn't pay a doomed capture attempt every burst.
1273/// Capture identity of the parked SAMPLED draft graph (`DraftGraphCtx::graph_s`).
1274///
1275/// EXACTNESS, not perf (lane/graph-s-key-exactness-20260819; receipts
1276/// `research/spec-cache-20260818/GRAPH-S-KEY.md`). Two classes of field live here, both
1277/// load-bearing:
1278///
1279/// - **Baked constants.** `seed` and `temp` are capture-time constants INSIDE the graph and `k`
1280/// sizes the q slots its replays write. A resumed request changing any of them must recapture.
1281/// This is all the key used to carry.
1282/// - **Regime fields.** `top_k`/`top_p`/`min_p`/`pen_on` are not baked, but they decide whether
1283/// the captured graph is a legal draft chain AT ALL. The in-graph draw is one gumbel-max over
1284/// the RAW softmax (`gumbel_perturb_ctr`, unfiltered by construction), while the verify builds
1285/// the accept test's `q` from `filter_stats(q_slots, top_k, top_p, min_p)`. If those disagree
1286/// the accept test evaluates a distribution the draft was never sampled from: a draft token
1287/// below the filter threshold gathers `q = 0` (`softmax_gather_filtered_f32`,
1288/// `cu/spec_sample.cu`) and `u * 0 < p` accepts it UNCONDITIONALLY.
1289///
1290/// Omitting the regime fields was reachable — not through the prefix-cache spec restore (that
1291/// path is greedy-only, `memra-server` `spec_restore_convertible`), but through WHOLE-SESSION
1292/// spec reuse: a parked `SpecSession` carries this `DraftGraphCtx`, and the pool-resume probe
1293/// applies no sampler predicate at all. Turn 1 pure-temp parks a graph; turn 2 of the same
1294/// conversation, same explicit seed and temperature, adds `top_p`/`top_k` and inherits it.
1295#[derive(Clone, Copy, PartialEq, Eq, Debug)]
1296pub(crate) struct SampledGraphKey {
1297 seed: u64,
1298 temp_bits: u32,
1299 k: usize,
1300 top_k: i32,
1301 top_p_bits: u32,
1302 min_p_bits: u32,
1303 pen_on: bool,
1304}
1305
1306impl SampledGraphKey {
1307 pub(crate) fn new(
1308 seed: u64,
1309 temp: f32,
1310 k: usize,
1311 top_k: i32,
1312 top_p: f32,
1313 min_p: f32,
1314 pen_on: bool,
1315 ) -> Self {
1316 SampledGraphKey {
1317 seed,
1318 temp_bits: temp.to_bits(),
1319 k,
1320 top_k,
1321 top_p_bits: top_p.to_bits(),
1322 min_p_bits: min_p.to_bits(),
1323 pen_on,
1324 }
1325 }
1326
1327 /// The one regime the in-graph sampled chain may stand in for the eager one: nothing but
1328 /// temperature shapes `q`. Computed FROM THE KEY so the capture guard, the launch guard and
1329 /// the key can never drift apart (they were three separate expressions before this lane, and
1330 /// the launch site simply forgot to ask).
1331 pub(crate) fn pure_temp(&self) -> bool {
1332 self.top_k == 0
1333 && f32::from_bits(self.top_p_bits) >= 1.0
1334 && f32::from_bits(self.min_p_bits) <= 0.0
1335 && !self.pen_on
1336 }
1337}
1338
1339pub(crate) struct DraftGraphCtx {
1340 g_tok: CudaSlice<u32>,
1341 g_pos: CudaSlice<i32>,
1342 g_seed: CudaSlice<f32>,
1343 g_p: CudaSlice<f32>,
1344 g_ctr: CudaSlice<u32>,
1345 g_q: CudaSlice<f32>,
1346 g_perturb: CudaSlice<f32>,
1347 q_slots: Vec<CudaSlice<f32>>,
1348 /// DRAFT-SIDE GRAMMAR MASK (lane/draft-mask): packed allowed-set words over the DRAFT
1349 /// head's vocab, at a STABLE address so the captured draft graph's mask node reads the
1350 /// per-position contents the host re-uploads before each replay (the graph-promote
1351 /// pattern from decode.rs). Empty unless the session drafts under a grammar.
1352 g_dmask: CudaSlice<u32>,
1353 /// was `graph` captured WITH the mask node? A parked graph of the wrong shape is dropped.
1354 graph_masked: bool,
1355 graph: Option<cudarc::driver::CudaGraph>,
1356 graph_s: Option<cudarc::driver::CudaGraph>,
1357 /// Failed-capture memoization for both graphs — LOUD on flip, cleared on pool resume
1358 /// (audit Q2, the TRT #16072 silent-permanent-coverage-loss class).
1359 failed: DraftGraphFallback,
1360 /// Capture identity of `graph_s` — see [`SampledGraphKey`]. `None` iff no sampled graph is
1361 /// parked; a request whose key differs drops the parked graph (and its q slots/keeper).
1362 s_key: Option<SampledGraphKey>,
1363 /// CAPTURE-RETAIN keepers (#68 root cause, 2026-08-04): the warmup-run transients whose
1364 /// pool addresses the captured graph(s) bake. Without these, the transients return to the
1365 /// pool at capture-body exit and later work (burst-boundary prime/fill/commit passes, or a
1366 /// co-served session in the worker) reuses those addresses — the persisted graph's replay
1367 /// then reads/writes live unrelated buffers (exactness corruption, first seen as the ST
1368 /// serve-spec 4B graph-arm corruption; one-shot CLI calls never re-shuffled the pool, which
1369 /// is why run-spec K=1..8 passed on the same checkpoint). Same fix class as
1370 /// capture_graph_retained's gemma/decode.rs sites — hold as long as the graph replays.
1371 keeper: Vec<Box<dyn std::any::Any + Send>>,
1372 keeper_s: Vec<Box<dyn std::any::Any + Send>>,
1373}
1374
1375/// Failed-capture memoization for the two draft graphs (audit Q2, 2026-08-05 — the
1376/// TRT #16072 trap class: pressure-triggered, silent, long-lived coverage loss).
1377///
1378/// Three contracts:
1379/// - LOUD FLIP: `mark_*` returns the warn line exactly on the false→true transition
1380/// (returned, not printed, so the once-per-flip contract is unit-testable); the caller
1381/// `eprintln!`s it UNCONDITIONALLY — a dropped draft graph is never silent. Re-marking
1382/// an already-failed graph returns None (the per-burst memoization that keeps the eager
1383/// fallback from paying a doomed capture attempt every burst).
1384/// - RESET ON RESUME: `reset_on_resume` clears both flags — a parked session resumed by a
1385/// NEW request gets one fresh capture chance instead of carrying a transient-pressure
1386/// failure for the pool's whole lifetime. Returns the note line only when a flag was
1387/// actually set (quiet on the common clean-resume path).
1388/// - Shape-change clears (`clear_*`) stay silent, exactly as before: they precede a fresh
1389/// capture attempt whose own failure would re-flip loudly.
1390#[derive(Default)]
1391pub(crate) struct DraftGraphFallback {
1392 greedy: bool,
1393 sampled: bool,
1394}
1395impl DraftGraphFallback {
1396 fn mark_greedy(&mut self, reason: &str) -> Option<String> {
1397 if self.greedy {
1398 return None;
1399 }
1400 self.greedy = true;
1401 Some(format!(
1402 "[spec] WARN: draft-graph capture failed ({reason}); eager fallback until session resume"
1403 ))
1404 }
1405 fn mark_sampled(&mut self, reason: &str) -> Option<String> {
1406 if self.sampled {
1407 return None;
1408 }
1409 self.sampled = true;
1410 Some(format!(
1411 "[spec] WARN: sampled draft-graph capture failed ({reason}); eager fallback until session resume"
1412 ))
1413 }
1414 fn greedy_failed(&self) -> bool {
1415 self.greedy
1416 }
1417 fn sampled_failed(&self) -> bool {
1418 self.sampled
1419 }
1420 fn clear_greedy(&mut self) {
1421 self.greedy = false;
1422 }
1423 fn clear_sampled(&mut self) {
1424 self.sampled = false;
1425 }
1426 /// Pool-resume reset: both graphs get a fresh capture chance. Some(note) iff any flag
1427 /// was set (so clean resumes stay quiet).
1428 pub(crate) fn reset_on_resume(&mut self) -> Option<String> {
1429 if !self.greedy && !self.sampled {
1430 return None;
1431 }
1432 let which = match (self.greedy, self.sampled) {
1433 (true, true) => "greedy+sampled",
1434 (true, false) => "greedy",
1435 _ => "sampled",
1436 };
1437 self.greedy = false;
1438 self.sampled = false;
1439 Some(format!(
1440 "[spec] draft-graph fallback reset on session resume ({which}); recapture eligible"
1441 ))
1442 }
1443}
1444
1445impl DraftGraphCtx {
1446 fn new(e: &Engine, n_embd: usize, qlen: usize) -> Result<Self, Box<dyn std::error::Error>> {
1447 Ok(DraftGraphCtx {
1448 g_tok: e.alloc_u32_zeroed(1)?,
1449 g_pos: e.htod_i32(&[0])?,
1450 g_seed: e.zeros(n_embd)?,
1451 g_p: e.zeros(1)?,
1452 g_ctr: e.alloc_u32_zeroed(1)?,
1453 g_q: e.zeros(qlen)?,
1454 g_perturb: e.zeros(qlen)?,
1455 q_slots: Vec::new(),
1456 g_dmask: e.alloc_u32_zeroed(1)?,
1457 graph_masked: false,
1458 graph: None,
1459 graph_s: None,
1460 failed: DraftGraphFallback::default(),
1461 s_key: None,
1462 keeper: Vec::new(),
1463 keeper_s: Vec::new(),
1464 })
1465 }
1466}
1467
1468pub(crate) struct MtpScratch {
1469 kv: KvLayer,
1470 /// Logical row capacity. On the graph/DC draft path it also doubles as fa_decode_dc's
1471 /// bucket_max: n_splits is sized from it ONCE, so the graph captured at round 0 stays valid
1472 /// for every later t_kv. Step35 refuses that path and may back this logical extent with the
1473 /// smaller host-indexed SWA ring instead.
1474 cap: usize,
1475}
1476
1477fn mtp_scratch_layout(
1478 cfg: &memra_gguf::config::ModelConfig,
1479 geom: Option<&crate::hybrid::DraftGeom>,
1480) -> (usize, usize, usize, usize) {
1481 // Student draft heads carry fewer KV heads (head_dim unchanged) -> smaller scratch rows.
1482 let n_head_kv = geom.map(|g| g.n_head_kv).unwrap_or(cfg.n_head_kv as usize);
1483 let head_dim_k = cfg.head_dim_k as usize;
1484 let head_dim_v = cfg.head_dim_v as usize;
1485 assert!(
1486 head_dim_k % 32 == 0 && head_dim_v % 32 == 0,
1487 "KVQUANT requires head_dim%32==0 (MTP scratch)"
1488 );
1489 let kv_dim_k = head_dim_k * n_head_kv;
1490 let kv_dim_v = head_dim_v * n_head_kv;
1491 // The fp8-KV arm deliberately does not reach the draft scratch; keep the exact format
1492 // policy shared with `MtpScratch::new` so admission scales the same allocation.
1493 let (kbb, vbb) = crate::kv_blk_bytes();
1494 let k_tok_bytes = (kv_dim_k / 32) * kbb;
1495 let v_tok_bytes = (kv_dim_v / 32) * vbb;
1496 (kv_dim_k, kv_dim_v, k_tok_bytes, v_tok_bytes)
1497}
1498
1499impl MtpScratch {
1500 fn new(
1501 e: &Engine,
1502 cfg: &memra_gguf::config::ModelConfig,
1503 cap: usize,
1504 geom: Option<&crate::hybrid::DraftGeom>,
1505 ) -> Result<Self, Box<dyn std::error::Error>> {
1506 // env-selected KV formats (default 34/24). The fp8-KV arm (MEMRA_KV_FP8) deliberately
1507 // does NOT reach the draft scratch: fp8 drafts drifted acceptance 69-88% -> 46%
1508 // (2026-07-12 A/B); the scratch is tiny, so it keeps baseline q8_0/q5_1 numerics
1509 // while the TRUNK cache carries the fp8 depth win. Scratch append/fa pass g=false.
1510 let (kv_dim_k, kv_dim_v, k_tok_bytes, v_tok_bytes) = mtp_scratch_layout(cfg, geom);
1511 let ring = if crate::cache::swa_ring_on() && cfg.arch.is_step35() {
1512 let window = cfg.step35.as_ref().unwrap().sliding_window as usize;
1513 Some(crate::cache::KvRing::new(
1514 crate::cache::swa_ring_rows(window, cap),
1515 window,
1516 ))
1517 } else {
1518 None
1519 };
1520 let alloc_rows = ring.as_ref().map(crate::cache::KvRing::rows).unwrap_or(cap);
1521 Ok(MtpScratch {
1522 kv: KvLayer {
1523 k: e.alloc_u8(alloc_rows * k_tok_bytes)?,
1524 v: e.alloc_u8(alloc_rows * v_tok_bytes)?,
1525 kv_dim_k,
1526 kv_dim_v,
1527 k_tok_bytes,
1528 v_tok_bytes,
1529 len: 0,
1530 ring,
1531 len_d: e.htod_i32(&[0])?,
1532 },
1533 cap,
1534 })
1535 }
1536 /// Set BOTH length counters: the host mirror AND the device len_d the captured append/fa read
1537 /// (a 4-byte in-place htod — the counter pointer is baked into the graph, never realloc'd).
1538 /// This is the ONLY truncation/rollback mechanism the persistent draft KV needs.
1539 fn set_len(&mut self, e: &Engine, n: usize) -> Result<(), Box<dyn std::error::Error>> {
1540 if self
1541 .kv
1542 .ring
1543 .as_ref()
1544 .is_some_and(|ring| !ring.can_rewind_to(n))
1545 {
1546 return Err("SWA ring MTP checkpoint has been lapped; full re-prime required".into());
1547 }
1548 self.kv.len = n;
1549 e.set_i32_one(&mut self.kv.len_d, n as i32)
1550 }
1551
1552 fn can_rewind_to(&self, n: usize) -> bool {
1553 self.kv
1554 .ring
1555 .as_ref()
1556 .is_none_or(|ring| ring.can_rewind_to(n))
1557 }
1558}
1559
1560/// Retained verify intermediates for the REPLAY-FREE partial accept (2026-07-03, the profiled
1561/// #1 spec cost at long ctx: the partial-accept replay was a DUPLICATE trunk pass — ~0.54 extra
1562/// full weight reads per round — recomputing columns the verify had already produced
1563/// bit-identically). Holds, per linear layer, everything needed to rebuild its recurrent state
1564/// to "after the first j verify columns" WITHOUT re-running the trunk:
1565/// - BATCHED-path layers (`gdn`): the exact token-major inputs the round's ONE gdn_scan
1566/// consumed. A prefix re-run of the SAME kernel (t=j) from the snapshot state is bit-identical
1567/// to the first j iterations of the verify's scan — the kernel's t-loop carries state in
1568/// registers and iteration t never depends on T. `qkv_mixed` (the conv input) feeds the
1569/// pure-copy ring rebuild.
1570/// - PER-COLUMN-path layers (`cols`): dtod clones of (conv_state, ssm_state) taken after each
1571/// column 0..t-2 — pure copies of the actual chain states (the last column is never a rebuild
1572/// target: j <= t-1).
1573/// Full-attn layers need nothing: their verify KV rows are bit-identical to eager's (the
1574/// decode-exact contract; verify-probe pins it), so rollback = len truncation.
1575struct GdnStash {
1576 qkv_mixed: CudaSlice<f32>, // [t, conv_dim] token-major (conv input)
1577 q_l2: CudaSlice<f32>,
1578 k_l2: CudaSlice<f32>,
1579 v_g: CudaSlice<f32>, // [t, num_v, d_state]
1580 g_log: CudaSlice<f32>,
1581 beta: CudaSlice<f32>, // [t, num_v]
1582}
1583struct VerifyCkpt {
1584 gdn: Vec<Option<GdnStash>>, // [n_layer], Some iff batched linear path ran
1585 cols: Vec<Option<Vec<(CudaSlice<f32>, CudaSlice<f32>)>>>, // [n_layer][col] = (conv, ssm) after col
1586}
1587/// Opaque handle for the dspark round (dflash.rs) — VerifyCkpt stays spec-private.
1588pub(crate) struct DsparkVerifyCkpt(VerifyCkpt);
1589
1590impl VerifyCkpt {
1591 fn new(n_layer: usize) -> Self {
1592 VerifyCkpt {
1593 gdn: (0..n_layer).map(|_| None).collect(),
1594 cols: (0..n_layer).map(|_| None).collect(),
1595 }
1596 }
1597}
1598
1599/// The stage-0/TX half of one PP verify. The boundary slot is the ownership token: stage 1
1600/// consumes exactly the slot selected by `tx()` / `tx_pipelined()`, never a slot inferred from
1601/// a logical round number.
1602struct VerifyBoundaryTicket {
1603 rt: &'static crate::pp::PpNRt,
1604 caller_stream: std::sync::Arc<cudarc::driver::CudaStream>,
1605 slot: usize,
1606 pos0: usize,
1607 t: usize,
1608 payload: usize,
1609 n_st: usize,
1610 pipelined: bool,
1611 pp_anatomy: bool,
1612 pp_started: std::time::Instant,
1613 reverse_ms: f64,
1614 stage0_ms: f64,
1615 tx_ms: f64,
1616 trace: Option<SpecPipeTraceCtx>,
1617}
1618
1619/// Explicit OPTIPIPE diagnostic control. Forced modes are set only by `optipipe-gate`; the
1620/// increment-2 controller can also be armed by the server's fresh-process research door.
1621#[derive(Clone, Copy, Debug, PartialEq, Eq)]
1622pub enum OptiForkGateMode {
1623 Disabled,
1624 Hit,
1625 Miss,
1626 Alternate,
1627 Abort,
1628 Controller,
1629}
1630
1631static OPTI_FORK_GATE_MODE: std::sync::atomic::AtomicU8 = std::sync::atomic::AtomicU8::new(0);
1632static OPTI_CONTROLLER_THRESHOLD: std::sync::atomic::AtomicU32 =
1633 std::sync::atomic::AtomicU32::new(0);
1634static OPTI_FORK_ATTEMPTS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
1635static OPTI_FORK_HITS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
1636static OPTI_FORK_MISSES: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
1637static OPTI_FORK_ABORT_DRAINS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
1638static OPTI_FORK_REFUSALS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
1639static OPTI_GATE_CHECKS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
1640static OPTI_GATE_ADMITS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
1641static OPTI_GATE_REJECTS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
1642static OPTI_RECONCILES: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
1643static OPTI_WASTED_DRAFT_TOKENS: std::sync::atomic::AtomicU64 =
1644 std::sync::atomic::AtomicU64::new(0);
1645static OPTI_SHADOW_DRAFT_TOKENS: std::sync::atomic::AtomicU64 =
1646 std::sync::atomic::AtomicU64::new(0);
1647static OPTI_BREAKER_TRIPS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
1648
1649impl OptiForkGateMode {
1650 fn code(self) -> u8 {
1651 match self {
1652 Self::Disabled => 0,
1653 Self::Hit => 1,
1654 Self::Miss => 2,
1655 Self::Alternate => 3,
1656 Self::Abort => 4,
1657 Self::Controller => 5,
1658 }
1659 }
1660
1661 fn configured() -> Self {
1662 match OPTI_FORK_GATE_MODE.load(std::sync::atomic::Ordering::Relaxed) {
1663 1 => Self::Hit,
1664 2 => Self::Miss,
1665 3 => Self::Alternate,
1666 4 => Self::Abort,
1667 5 => Self::Controller,
1668 _ => Self::Disabled,
1669 }
1670 }
1671
1672 fn action(self, generation: u64) -> OptiForkAction {
1673 match self {
1674 Self::Hit => OptiForkAction::Hit,
1675 Self::Miss => OptiForkAction::Miss,
1676 Self::Alternate if generation & 1 == 0 => OptiForkAction::Hit,
1677 Self::Alternate => OptiForkAction::Miss,
1678 Self::Abort => OptiForkAction::Abort,
1679 Self::Disabled | Self::Controller => {
1680 unreachable!("non-forced mode cannot choose a forced fork action")
1681 }
1682 }
1683 }
1684
1685 fn is_forced(self) -> bool {
1686 matches!(self, Self::Hit | Self::Miss | Self::Alternate | Self::Abort)
1687 }
1688}
1689
1690/// Arm or disarm the forced harness. Serving uses only `set_optipipe_controller_threshold`.
1691pub fn set_optipipe_gate_mode(mode: OptiForkGateMode) {
1692 OPTI_FORK_GATE_MODE.store(mode.code(), std::sync::atomic::Ordering::Relaxed);
1693}
1694
1695/// Arm the increment-2 diagnostic controller. The threshold applies to the uncalibrated
1696/// two-token draft-probability product. Serving can call this only through its explicit
1697/// fresh-process research door; the absent-door default remains byte-for-byte disabled.
1698pub fn set_optipipe_controller_threshold(threshold: f32) {
1699 assert!(threshold.is_finite() && (0.0..=1.0).contains(&threshold));
1700 OPTI_CONTROLLER_THRESHOLD.store(threshold.to_bits(), std::sync::atomic::Ordering::Relaxed);
1701 set_optipipe_gate_mode(OptiForkGateMode::Controller);
1702}
1703
1704#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
1705pub struct OptiForkGateStats {
1706 pub attempts: u64,
1707 pub hits: u64,
1708 pub misses: u64,
1709 pub abort_drains: u64,
1710 pub refusals: u64,
1711 pub gate_checks: u64,
1712 pub gate_admits: u64,
1713 pub gate_rejects: u64,
1714 pub reconciles: u64,
1715 pub wasted_draft_tokens: u64,
1716 pub shadow_draft_tokens: u64,
1717 pub breaker_trips: u64,
1718}
1719
1720#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
1721pub struct OptiForkStateIdentity {
1722 pub trunk_kv_bytes: usize,
1723 pub recurrent_bytes: usize,
1724 pub scratch_kv_bytes: usize,
1725 pub hidden_bytes: usize,
1726}
1727
1728pub fn reset_optipipe_gate_stats() {
1729 for counter in [
1730 &OPTI_FORK_ATTEMPTS,
1731 &OPTI_FORK_HITS,
1732 &OPTI_FORK_MISSES,
1733 &OPTI_FORK_ABORT_DRAINS,
1734 &OPTI_FORK_REFUSALS,
1735 &OPTI_GATE_CHECKS,
1736 &OPTI_GATE_ADMITS,
1737 &OPTI_GATE_REJECTS,
1738 &OPTI_RECONCILES,
1739 &OPTI_WASTED_DRAFT_TOKENS,
1740 &OPTI_SHADOW_DRAFT_TOKENS,
1741 &OPTI_BREAKER_TRIPS,
1742 ] {
1743 counter.store(0, std::sync::atomic::Ordering::Relaxed);
1744 }
1745}
1746
1747pub fn optipipe_gate_stats() -> OptiForkGateStats {
1748 let load = |v: &std::sync::atomic::AtomicU64| v.load(std::sync::atomic::Ordering::Relaxed);
1749 OptiForkGateStats {
1750 attempts: load(&OPTI_FORK_ATTEMPTS),
1751 hits: load(&OPTI_FORK_HITS),
1752 misses: load(&OPTI_FORK_MISSES),
1753 abort_drains: load(&OPTI_FORK_ABORT_DRAINS),
1754 refusals: load(&OPTI_FORK_REFUSALS),
1755 gate_checks: load(&OPTI_GATE_CHECKS),
1756 gate_admits: load(&OPTI_GATE_ADMITS),
1757 gate_rejects: load(&OPTI_GATE_REJECTS),
1758 reconciles: load(&OPTI_RECONCILES),
1759 wasted_draft_tokens: load(&OPTI_WASTED_DRAFT_TOKENS),
1760 shadow_draft_tokens: load(&OPTI_SHADOW_DRAFT_TOKENS),
1761 breaker_trips: load(&OPTI_BREAKER_TRIPS),
1762 }
1763}
1764
1765#[derive(Clone, Copy, Debug)]
1766struct OptiControllerPolicy {
1767 threshold: f32,
1768 consecutive_misses: u8,
1769 breaker_tripped: bool,
1770}
1771
1772impl OptiControllerPolicy {
1773 fn configured() -> Self {
1774 Self {
1775 threshold: f32::from_bits(
1776 OPTI_CONTROLLER_THRESHOLD.load(std::sync::atomic::Ordering::Relaxed),
1777 ),
1778 consecutive_misses: 0,
1779 breaker_tripped: false,
1780 }
1781 }
1782
1783 fn admit(&self, q_proxy: f32) -> bool {
1784 q_proxy.is_finite()
1785 && (0.0..=1.0).contains(&q_proxy)
1786 && (self.threshold == 0.0 || (!self.breaker_tripped && q_proxy >= self.threshold))
1787 }
1788
1789 /// Returns true exactly when this resolution newly trips the three-miss breaker.
1790 fn resolve(&mut self, hit: bool) -> bool {
1791 // q*=0 is the lane's explicit unconditional measurement arm. Its purpose is to price
1792 // every optimistic opportunity, so the safety breaker is measured separately and must
1793 // not silently turn this arm into "three attempts then serial".
1794 if self.threshold == 0.0 {
1795 self.consecutive_misses = 0;
1796 return false;
1797 }
1798 if hit {
1799 self.consecutive_misses = 0;
1800 return false;
1801 }
1802 self.consecutive_misses = self.consecutive_misses.saturating_add(1);
1803 if !self.breaker_tripped && self.consecutive_misses >= 3 {
1804 self.breaker_tripped = true;
1805 return true;
1806 }
1807 false
1808 }
1809}
1810
1811#[derive(Clone, Copy, Debug, PartialEq, Eq)]
1812enum OptiForkAction {
1813 Hit,
1814 Miss,
1815 Abort,
1816}
1817
1818#[derive(Clone, Copy, Debug, PartialEq, Eq)]
1819struct OptiForkGeneration {
1820 id: u64,
1821 slot: usize,
1822}
1823
1824#[derive(Default)]
1825struct OptiForkGenerationTracker {
1826 next: u64,
1827 live: [Option<u64>; 2],
1828}
1829
1830impl OptiForkGenerationTracker {
1831 fn reserve(&mut self) -> Result<OptiForkGeneration, Box<dyn std::error::Error>> {
1832 let generation = OptiForkGeneration {
1833 id: self.next,
1834 slot: (self.next & 1) as usize,
1835 };
1836 if let Some(live) = self.live[generation.slot] {
1837 return Err(format!(
1838 "optipipe snapshot slot {} still owns generation {live}; refusing to overwrite it",
1839 generation.slot,
1840 )
1841 .into());
1842 }
1843 self.next += 1;
1844 self.live[generation.slot] = Some(generation.id);
1845 Ok(generation)
1846 }
1847
1848 fn retire(&mut self, generation: OptiForkGeneration) -> Result<(), Box<dyn std::error::Error>> {
1849 match self.live[generation.slot] {
1850 Some(id) if id == generation.id => {
1851 self.live[generation.slot] = None;
1852 Ok(())
1853 }
1854 other => Err(format!(
1855 "optipipe generation teardown mismatch: ticket={} slot={} live={other:?}",
1856 generation.id, generation.slot,
1857 )
1858 .into()),
1859 }
1860 }
1861}
1862
1863struct OptiForkSeedGeneration {
1864 h_seed: CudaSlice<f32>,
1865 fill_prev: CudaSlice<f32>,
1866 scratch_len: usize,
1867}
1868
1869/// Allocate or refresh one full checkpoint through the engine that owns each PP stage. The
1870/// generic cache helper accepts one device and therefore cannot copy GDN state split across
1871/// devices. KV lengths and position stay host metadata; only recurrent buffers need stage-local
1872/// device ownership.
1873fn opti_snapshot_stage_owned(
1874 e: &Engine,
1875 cache: &Cache,
1876 rt: &'static crate::pp::PpNRt,
1877 fence: &[usize],
1878) -> Result<crate::cache::CacheSnapshot, Box<dyn std::error::Error>> {
1879 let n = cache.kv.len();
1880 let mut snapshot = crate::cache::CacheSnapshot {
1881 kv_len: vec![None; n],
1882 conv: (0..n).map(|_| None).collect(),
1883 ssm: (0..n).map(|_| None).collect(),
1884 pos: cache.pos,
1885 };
1886 opti_snapshot_stage_owned_into(e, cache, rt, fence, &mut snapshot)?;
1887 Ok(snapshot)
1888}
1889
1890fn opti_snapshot_stage_owned_into(
1891 e: &Engine,
1892 cache: &Cache,
1893 rt: &'static crate::pp::PpNRt,
1894 fence: &[usize],
1895 snapshot: &mut crate::cache::CacheSnapshot,
1896) -> Result<(), Box<dyn std::error::Error>> {
1897 if fence.len() != rt.n_stages() + 1 || snapshot.kv_len.len() != cache.kv.len() {
1898 return Err("optipipe stage-owned snapshot shape mismatch".into());
1899 }
1900 for stage in 0..rt.n_stages() {
1901 opti_snapshot_one_stage_owned_into(e, cache, rt, fence, stage, snapshot)?;
1902 }
1903 snapshot.pos = cache.pos;
1904 Ok(())
1905}
1906
1907/// Refresh one PP stage of a checkpoint. Increment 2 uses this split form so stage 0's
1908/// optimistic post-N state is captured before N+1 stage 0 is queued, while stage 1's matching
1909/// post-N state is captured only after N stage 1 is enqueued. Calling the all-stage helper at
1910/// either point would capture one side of the fork at the wrong generation.
1911fn opti_snapshot_one_stage_owned_into(
1912 e: &Engine,
1913 cache: &Cache,
1914 rt: &'static crate::pp::PpNRt,
1915 fence: &[usize],
1916 stage: usize,
1917 snapshot: &mut crate::cache::CacheSnapshot,
1918) -> Result<(), Box<dyn std::error::Error>> {
1919 if fence.len() != rt.n_stages() + 1
1920 || snapshot.kv_len.len() != cache.kv.len()
1921 || stage >= rt.n_stages()
1922 {
1923 return Err("optipipe single-stage snapshot shape mismatch".into());
1924 }
1925 let _scope = rt.enter(stage);
1926 let owner = rt.engine(stage, e);
1927 for il in fence[stage]..fence[stage + 1] {
1928 snapshot.kv_len[il] = cache.kv[il].as_ref().map(|kv| kv.len);
1929 match &cache.recur[il] {
1930 Some(recur) => {
1931 match snapshot.conv[il].as_mut() {
1932 Some(dst) => {
1933 owner.copy_into(dst, 0, &recur.conv_state, recur.conv_state.len())?
1934 }
1935 None => snapshot.conv[il] = Some(owner.clone_dtod(&recur.conv_state)?),
1936 }
1937 match snapshot.ssm[il].as_mut() {
1938 Some(dst) => {
1939 owner.copy_into(dst, 0, &recur.ssm_state, recur.ssm_state.len())?
1940 }
1941 None => snapshot.ssm[il] = Some(owner.clone_dtod(&recur.ssm_state)?),
1942 }
1943 }
1944 None if snapshot.conv[il].is_some() || snapshot.ssm[il].is_some() => {
1945 return Err(
1946 format!("optipipe stage-owned snapshot layer {il} changed shape").into(),
1947 );
1948 }
1949 None => {}
1950 }
1951 }
1952 snapshot.pos = cache.pos;
1953 Ok(())
1954}
1955
1956/// Increment-1 persistent fork state. Exactly two snapshot/seed slots alternate; a live ticket
1957/// names its generation and keeps teardown fail-closed. Only stage 0 is allowed to mutate before
1958/// resolve, so the reconcile tables and conditional restores are stage-local.
1959struct OptiForkState {
1960 mode: OptiForkGateMode,
1961 controller: Option<OptiControllerPolicy>,
1962 generations: OptiForkGenerationTracker,
1963 active_snapshot_slot: usize,
1964 alternate_snapshot: crate::cache::CacheSnapshot,
1965 seeds: [OptiForkSeedGeneration; 2],
1966 rt: &'static crate::pp::PpNRt,
1967 fence: [usize; 3],
1968 split: usize,
1969 len_ptrs: CudaSlice<u64>,
1970 saved_lens: CudaSlice<i32>,
1971 forced_acc: CudaSlice<u32>,
1972 valid: CudaSlice<u32>,
1973 stage0_stream: std::sync::Arc<cudarc::driver::CudaStream>,
1974 logical_payload_bytes: [usize; 2],
1975}
1976
1977struct OptiForkTicket {
1978 generation: OptiForkGeneration,
1979 boundary: Option<VerifyBoundaryTicket>,
1980 drain: std::sync::Arc<cudarc::driver::CudaStream>,
1981 settled: bool,
1982}
1983
1984struct OptiControllerTicket {
1985 generation: OptiForkGeneration,
1986 boundary: Option<VerifyBoundaryTicket>,
1987 ckpt: Option<VerifyCkpt>,
1988 verify_tokens: [u32; 2],
1989 draft_prob: f32,
1990 eager_seed: Option<CudaSlice<f32>>,
1991 q_proxy: f32,
1992 scratch_len: usize,
1993 issued_at: std::time::Instant,
1994 drain: std::sync::Arc<cudarc::driver::CudaStream>,
1995 settled: bool,
1996}
1997
1998struct OptiControllerPrepared {
1999 verify_tokens: [u32; 2],
2000 draft_prob: f32,
2001 eager_seed: Option<CudaSlice<f32>>,
2002 q_proxy: f32,
2003 scratch_len: usize,
2004}
2005
2006impl OptiControllerTicket {
2007 fn take_boundary(&mut self) -> VerifyBoundaryTicket {
2008 self.boundary
2009 .take()
2010 .expect("controller boundary ticket already consumed")
2011 }
2012
2013 fn take_ckpt(&mut self) -> VerifyCkpt {
2014 self.ckpt
2015 .take()
2016 .expect("controller verify checkpoint already consumed")
2017 }
2018
2019 fn take_eager_seed(&mut self) -> Option<CudaSlice<f32>> {
2020 self.eager_seed.take()
2021 }
2022
2023 fn settle(&mut self) {
2024 self.settled = true;
2025 }
2026}
2027
2028impl Drop for OptiControllerTicket {
2029 fn drop(&mut self) {
2030 if !self.settled {
2031 let _ = self.drain.synchronize();
2032 OPTI_FORK_ABORT_DRAINS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2033 }
2034 }
2035}
2036
2037impl OptiForkTicket {
2038 fn take_boundary(&mut self) -> VerifyBoundaryTicket {
2039 self.boundary
2040 .take()
2041 .expect("fork ticket boundary already consumed")
2042 }
2043
2044 fn settle(&mut self) {
2045 self.settled = true;
2046 }
2047}
2048
2049impl Drop for OptiForkTicket {
2050 fn drop(&mut self) {
2051 if !self.settled {
2052 let _ = self.drain.synchronize();
2053 OPTI_FORK_ABORT_DRAINS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2054 }
2055 }
2056}
2057
2058impl OptiForkState {
2059 #[allow(clippy::too_many_arguments)]
2060 fn new(
2061 e: &Engine,
2062 cache: &Cache,
2063 mode: OptiForkGateMode,
2064 alternate_snapshot: crate::cache::CacheSnapshot,
2065 h_seed: &CudaSlice<f32>,
2066 fill_prev: &CudaSlice<f32>,
2067 rt: &'static crate::pp::PpNRt,
2068 split: usize,
2069 n_layer: usize,
2070 ) -> Result<Self, Box<dyn std::error::Error>> {
2071 let fence = [0, split, n_layer];
2072 let mut logical_payload_bytes = [0usize; 2];
2073 for stage in 0..2 {
2074 for il in fence[stage]..fence[stage + 1] {
2075 logical_payload_bytes[stage] += alternate_snapshot.conv[il]
2076 .as_ref()
2077 .map_or(0, |v| v.len() * std::mem::size_of::<f32>());
2078 logical_payload_bytes[stage] += alternate_snapshot.ssm[il]
2079 .as_ref()
2080 .map_or(0, |v| v.len() * std::mem::size_of::<f32>());
2081 }
2082 }
2083 let seeds = [
2084 OptiForkSeedGeneration {
2085 h_seed: e.clone_dtod(h_seed)?,
2086 fill_prev: e.clone_dtod(fill_prev)?,
2087 scratch_len: 0,
2088 },
2089 OptiForkSeedGeneration {
2090 h_seed: e.clone_dtod(h_seed)?,
2091 fill_prev: e.clone_dtod(fill_prev)?,
2092 scratch_len: 0,
2093 },
2094 ];
2095 let (len_ptrs, saved_lens, forced_acc, valid, stage0_stream) = {
2096 let _stage = rt.enter(0);
2097 let e0 = rt.engine(0, e);
2098 (
2099 crate::round_stream::kv_len_ptr_table_range(e0, cache, 0..split, None)?,
2100 e0.htod_i32(&vec![0; split])?,
2101 e0.alloc_u32_zeroed(2)?,
2102 e0.alloc_u32_zeroed(1)?,
2103 e0.stream(),
2104 )
2105 };
2106 logical_payload_bytes[0] += seeds
2107 .iter()
2108 .map(|seed| (seed.h_seed.len() + seed.fill_prev.len()) * std::mem::size_of::<f32>())
2109 .sum::<usize>();
2110 logical_payload_bytes[0] += len_ptrs.len() * std::mem::size_of::<u64>()
2111 + saved_lens.len() * std::mem::size_of::<i32>()
2112 + forced_acc.len() * std::mem::size_of::<u32>()
2113 + valid.len() * std::mem::size_of::<u32>();
2114 Ok(Self {
2115 mode,
2116 controller: (mode == OptiForkGateMode::Controller)
2117 .then(OptiControllerPolicy::configured),
2118 generations: OptiForkGenerationTracker::default(),
2119 active_snapshot_slot: 0,
2120 alternate_snapshot,
2121 seeds,
2122 rt,
2123 fence,
2124 split,
2125 len_ptrs,
2126 saved_lens,
2127 forced_acc,
2128 valid,
2129 stage0_stream,
2130 logical_payload_bytes,
2131 })
2132 }
2133
2134 fn reserve(
2135 &mut self,
2136 current_snapshot: &mut crate::cache::CacheSnapshot,
2137 ) -> Result<OptiForkGeneration, Box<dyn std::error::Error>> {
2138 let generation = self.generations.reserve()?;
2139 if generation.slot != self.active_snapshot_slot {
2140 std::mem::swap(current_snapshot, &mut self.alternate_snapshot);
2141 self.active_snapshot_slot = generation.slot;
2142 }
2143 Ok(generation)
2144 }
2145
2146 fn capture_seed(
2147 &mut self,
2148 e: &Engine,
2149 generation: OptiForkGeneration,
2150 h_seed: &CudaSlice<f32>,
2151 fill_prev: &CudaSlice<f32>,
2152 scratch_len: usize,
2153 ) -> Result<(), Box<dyn std::error::Error>> {
2154 let seed = &mut self.seeds[generation.slot];
2155 e.copy_into(&mut seed.h_seed, 0, h_seed, h_seed.len())?;
2156 e.copy_into(&mut seed.fill_prev, 0, fill_prev, fill_prev.len())?;
2157 seed.scratch_len = scratch_len;
2158 Ok(())
2159 }
2160
2161 fn ticket(
2162 &self,
2163 generation: OptiForkGeneration,
2164 boundary: VerifyBoundaryTicket,
2165 ) -> OptiForkTicket {
2166 OptiForkTicket {
2167 generation,
2168 boundary: Some(boundary),
2169 drain: self.stage0_stream.clone(),
2170 settled: false,
2171 }
2172 }
2173
2174 #[allow(clippy::too_many_arguments)]
2175 fn controller_ticket(
2176 &self,
2177 generation: OptiForkGeneration,
2178 boundary: VerifyBoundaryTicket,
2179 ckpt: VerifyCkpt,
2180 verify_tokens: [u32; 2],
2181 draft_prob: f32,
2182 eager_seed: Option<CudaSlice<f32>>,
2183 q_proxy: f32,
2184 scratch_len: usize,
2185 ) -> OptiControllerTicket {
2186 OptiControllerTicket {
2187 generation,
2188 boundary: Some(boundary),
2189 ckpt: Some(ckpt),
2190 verify_tokens,
2191 draft_prob,
2192 eager_seed,
2193 q_proxy,
2194 scratch_len,
2195 issued_at: std::time::Instant::now(),
2196 drain: self.stage0_stream.clone(),
2197 settled: false,
2198 }
2199 }
2200
2201 fn reserve_successor(&mut self) -> Result<OptiForkGeneration, Box<dyn std::error::Error>> {
2202 self.generations.reserve()
2203 }
2204
2205 fn successor_snapshot_mut(&mut self) -> &mut crate::cache::CacheSnapshot {
2206 &mut self.alternate_snapshot
2207 }
2208
2209 fn promote_successor_snapshot(
2210 &mut self,
2211 current_snapshot: &mut crate::cache::CacheSnapshot,
2212 generation: OptiForkGeneration,
2213 ) {
2214 std::mem::swap(current_snapshot, &mut self.alternate_snapshot);
2215 self.active_snapshot_slot = generation.slot;
2216 }
2217
2218 fn queue_actual_reconcile(
2219 &mut self,
2220 e: &Engine,
2221 snapshot: &crate::cache::CacheSnapshot,
2222 acc: &CudaSlice<u32>,
2223 optimistic_pending: u32,
2224 base: usize,
2225 ) -> Result<(), Box<dyn std::error::Error>> {
2226 let saved: Vec<i32> = (0..self.split)
2227 .map(|il| snapshot.kv_len[il].map(|v| v as i32).unwrap_or(0))
2228 .collect();
2229 // Serving keeps the caller/accept walk on the head (stage-1) device. Record the accept
2230 // decision point there and append a wait to stage 0 after its optimistic successor/TX;
2231 // the validity/reconcile kernels must never peer-read acc before it is written. The
2232 // increment-1 harness uses primary stage 0, where stream order already provides this.
2233 if self.rt.engine(0, e).ctx().ordinal() != e.ctx().ordinal() {
2234 self.rt.fence_stages_behind(&e.stream())?;
2235 }
2236 let _stage = self.rt.enter(0);
2237 let e0 = self.rt.engine(0, e);
2238 e0.htod_i32_into(&mut self.saved_lens, &saved)?;
2239 e0.spec_fork_valid(acc, optimistic_pending, &mut self.valid)?;
2240 e0.spec_fork_reconcile_kv(
2241 &self.len_ptrs,
2242 &self.saved_lens,
2243 acc,
2244 &self.valid,
2245 base,
2246 self.split,
2247 )
2248 }
2249
2250 fn finish_actual_reconcile(
2251 &mut self,
2252 e: &Engine,
2253 cache: &mut Cache,
2254 snapshot: &crate::cache::CacheSnapshot,
2255 n_acc: usize,
2256 base: usize,
2257 hit: bool,
2258 ) -> Result<(), Box<dyn std::error::Error>> {
2259 if hit {
2260 return Ok(());
2261 }
2262 let len_delta = base + n_acc;
2263 for il in 0..self.split {
2264 if let (Some(kv), Some(saved)) = (cache.kv[il].as_mut(), snapshot.kv_len[il]) {
2265 kv.len = saved + len_delta;
2266 }
2267 }
2268 {
2269 let _stage = self.rt.enter(1);
2270 let e1 = self.rt.engine(1, e);
2271 for il in self.split..self.fence[2] {
2272 if let (Some(kv), Some(saved)) = (cache.kv[il].as_mut(), snapshot.kv_len[il]) {
2273 kv.len = saved + len_delta;
2274 e1.set_i32_one(&mut kv.len_d, kv.len as i32)?;
2275 }
2276 }
2277 }
2278 self.rt.publish_to(0, &e.stream())?;
2279 Ok(())
2280 }
2281
2282 fn cancel_controller_ticket(
2283 &mut self,
2284 e: &Engine,
2285 cache: &mut Cache,
2286 scratch: &mut MtpScratch,
2287 snapshot: &crate::cache::CacheSnapshot,
2288 ticket: &mut OptiControllerTicket,
2289 ) -> Result<(), Box<dyn std::error::Error>> {
2290 {
2291 let _stage = self.rt.enter(0);
2292 let e0 = self.rt.engine(0, e);
2293 for il in 0..self.split {
2294 if let (Some(kv), Some(saved)) = (cache.kv[il].as_mut(), snapshot.kv_len[il]) {
2295 kv.len = saved;
2296 e0.set_i32_one(&mut kv.len_d, saved as i32)?;
2297 }
2298 }
2299 }
2300 scratch.set_len(e, snapshot.pos)?;
2301 ticket.settle();
2302 self.generations.retire(ticket.generation)?;
2303 OPTI_FORK_ABORT_DRAINS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2304 OPTI_WASTED_DRAFT_TOKENS.fetch_add(2, std::sync::atomic::Ordering::Relaxed);
2305 eprintln!(
2306 "[opti-controller] tail-drain generation={} slot={}",
2307 ticket.generation.id, ticket.generation.slot,
2308 );
2309 Ok(())
2310 }
2311
2312 #[allow(clippy::too_many_arguments)]
2313 fn reconcile(
2314 &mut self,
2315 e: &Engine,
2316 cache: &mut Cache,
2317 scratch: &mut MtpScratch,
2318 snapshot: &crate::cache::CacheSnapshot,
2319 h_seed: &mut CudaSlice<f32>,
2320 fill_prev: &mut CudaSlice<f32>,
2321 generation: OptiForkGeneration,
2322 action: OptiForkAction,
2323 optimistic_pending: u32,
2324 ) -> Result<(), Box<dyn std::error::Error>> {
2325 debug_assert!(action != OptiForkAction::Abort);
2326 let miss_started = std::time::Instant::now();
2327 let keep = action == OptiForkAction::Hit;
2328 let saved: Vec<i32> = (0..self.split)
2329 .map(|il| snapshot.kv_len[il].map(|v| v as i32).unwrap_or(0))
2330 .collect();
2331 let seed = &self.seeds[generation.slot];
2332 {
2333 let _stage = self.rt.enter(0);
2334 let e0 = self.rt.engine(0, e);
2335 e0.htod_i32_into(&mut self.saved_lens, &saved)?;
2336 let forced = if keep {
2337 [1u32, optimistic_pending]
2338 } else {
2339 [0u32, optimistic_pending]
2340 };
2341 e0.htod_u32_into(&mut self.forced_acc, &forced)?;
2342 e0.spec_fork_valid(&self.forced_acc, optimistic_pending, &mut self.valid)?;
2343 e0.spec_fork_reconcile_kv(
2344 &self.len_ptrs,
2345 &self.saved_lens,
2346 &self.forced_acc,
2347 &self.valid,
2348 0,
2349 self.split,
2350 )?;
2351 for il in 0..self.split {
2352 if let Some(recur) = cache.recur[il].as_mut() {
2353 let conv = snapshot.conv[il]
2354 .as_ref()
2355 .ok_or("optipipe stage0 snapshot missing conv state")?;
2356 let ssm = snapshot.ssm[il]
2357 .as_ref()
2358 .ok_or("optipipe stage0 snapshot missing ssm state")?;
2359 e0.spec_fork_restore_f32(conv, &mut recur.conv_state, &self.valid)?;
2360 e0.spec_fork_restore_f32(ssm, &mut recur.ssm_state, &self.valid)?;
2361 }
2362 }
2363 e0.spec_fork_restore_f32(&seed.h_seed, h_seed, &self.valid)?;
2364 e0.spec_fork_restore_f32(&seed.fill_prev, fill_prev, &self.valid)?;
2365 }
2366
2367 if keep {
2368 OPTI_FORK_HITS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2369 return Ok(());
2370 }
2371
2372 for il in 0..self.split {
2373 if let (Some(kv), Some(saved)) = (cache.kv[il].as_mut(), snapshot.kv_len[il]) {
2374 kv.len = saved;
2375 }
2376 }
2377 scratch.set_len(e, seed.scratch_len)?;
2378 // Targeted E_restart: publish only stage 0's reconcile to the caller, then bound the
2379 // forced diagnostic so the retained number is the actual miss cost, not enqueue time.
2380 let caller = e.stream();
2381 self.rt.publish_to(0, &caller)?;
2382 caller.synchronize()?;
2383 let miss_ms = miss_started.elapsed().as_secs_f64() * 1e3;
2384 eprintln!(
2385 "[opti-fork-reconcile] generation={} slot={} miss_ms={miss_ms:.3}",
2386 generation.id, generation.slot,
2387 );
2388 OPTI_FORK_MISSES.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2389 Ok(())
2390 }
2391
2392 fn retire(&mut self, generation: OptiForkGeneration) -> Result<(), Box<dyn std::error::Error>> {
2393 self.generations.retire(generation)
2394 }
2395}
2396
2397impl HybridModel {
2398 fn opti_graph_draft_step(
2399 &self,
2400 e: &Engine,
2401 mtp: &MtpHead,
2402 dctx: &mut DraftGraphCtx,
2403 scratch: &mut MtpScratch,
2404 d_vocab: usize,
2405 ) -> Result<(u32, f32), Box<dyn std::error::Error>> {
2406 dctx.graph
2407 .as_ref()
2408 .ok_or("optipipe controller requires the greedy draft graph")?
2409 .launch()?;
2410 scratch.kv.len += 1;
2411 let idx = e.dtoh_u32_one(&dctx.g_tok)?;
2412 if (idx as usize) >= d_vocab {
2413 return Err(
2414 format!("optipipe draft argmax sentinel 0x{idx:08x} >= d_vocab {d_vocab}").into(),
2415 );
2416 }
2417 let probability = e.dtoh(&dctx.g_p)?[0];
2418 if !probability.is_finite() || !(0.0..=1.0).contains(&probability) {
2419 return Err(format!("optipipe draft probability is invalid: {probability}").into());
2420 }
2421 let token = match &mtp.d2t {
2422 Some(map) => map[idx as usize],
2423 None => idx,
2424 };
2425 if token != idx {
2426 e.set_u32_one(&mut dctx.g_tok, token)?;
2427 }
2428 Ok((token, probability))
2429 }
2430
2431 #[allow(clippy::too_many_arguments)]
2432 fn opti_controller_draft_step(
2433 &self,
2434 e: &Engine,
2435 mtp: &MtpHead,
2436 dctx: &mut DraftGraphCtx,
2437 scratch: &mut MtpScratch,
2438 d_vocab: usize,
2439 eager_state: &mut Option<(u32, CudaSlice<f32>)>,
2440 eager_pos: usize,
2441 embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
2442 ) -> Result<(u32, f32), Box<dyn std::error::Error>> {
2443 if dctx.graph.is_some() {
2444 return self.opti_graph_draft_step(e, mtp, dctx, scratch, d_vocab);
2445 }
2446 let (input_token, input_seed) = eager_state
2447 .take()
2448 .ok_or("optipipe eager continuation seed is unavailable")?;
2449 let (logits, next_seed) = self.mtp_head_forward_dev(
2450 e,
2451 mtp,
2452 input_token,
2453 &input_seed,
2454 scratch,
2455 eager_pos,
2456 embd_dev,
2457 None,
2458 )?;
2459 let token_d = e.argmax_token_device(&logits, d_vocab)?;
2460 let idx = e.dtoh_u32_one(&token_d)?;
2461 if (idx as usize) >= d_vocab {
2462 return Err(format!(
2463 "optipipe eager draft argmax sentinel 0x{idx:08x} >= d_vocab {d_vocab}"
2464 )
2465 .into());
2466 }
2467 let probability_d = e.prob_of_token_device(&logits, &token_d, d_vocab)?;
2468 let probability = e.dtoh(&probability_d)?[0];
2469 if !probability.is_finite() || !(0.0..=1.0).contains(&probability) {
2470 return Err(
2471 format!("optipipe eager draft probability is invalid: {probability}").into(),
2472 );
2473 }
2474 let token = match &mtp.d2t {
2475 Some(map) => map[idx as usize],
2476 None => idx,
2477 };
2478 *eager_state = Some((token, next_seed));
2479 Ok((token, probability))
2480 }
2481
2482 /// NextN head forward for ONE draft token (§A ops 1-13, T=1).
2483 /// Inputs: `e_tok` = the token to predict FROM (last committed / previous draft); `h_seed` =
2484 /// the trunk's pre-output_norm hidden of that token (§A op 2 input). `mtp_pos` = absolute
2485 /// position of the token being predicted from. Returns (draft_logits[n_vocab] host, h_nextn dev).
2486 /// `h_nextn` (§A op 10) becomes `h_seed` for the next autoregressive draft step.
2487 /// Device-resident: returns draft logits ON DEVICE (no [n_vocab] dtoh). The greedy draft
2488 /// loop only needs argmax — paired with `argmax_token_device` this cuts the ~600KB logits
2489 /// transfer + host argmax per draft token from the K-token draft chain.
2490 #[allow(clippy::too_many_arguments)]
2491 fn mtp_head_forward_dev(
2492 &self,
2493 e: &Engine,
2494 mtp: &MtpHead,
2495 e_tok: u32,
2496 h_seed: &CudaSlice<f32>,
2497 scratch: &mut MtpScratch,
2498 mtp_pos: usize,
2499 embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
2500 // DRAFT-SIDE GRAMMAR MASK (lane/draft-mask): (packed draft-vocab allowed set, words).
2501 // Applied to the head logits BEFORE they are returned, so every consumer (argmax,
2502 // gumbel draw, p-min prob) sees the grammar-legal row. None = unmasked (pre-lane).
2503 mask: Option<(&CudaSlice<u32>, usize)>,
2504 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
2505 // MEMRA_SPEC_ANATOMY=1 — eager-step phase timers (diagnostic only). Phase boundaries
2506 // sync the stream, so absolute time inflates; the BREAKDOWN is the signal. Cumulative
2507 // summary on stderr every 128 steps: glue (embed..attn_norm), attn, ffn, head.
2508 use std::sync::atomic::{AtomicU64, Ordering::Relaxed};
2509 static ANAT_NS: [AtomicU64; 5] = [
2510 AtomicU64::new(0),
2511 AtomicU64::new(0),
2512 AtomicU64::new(0),
2513 AtomicU64::new(0),
2514 AtomicU64::new(0),
2515 ];
2516 static ANAT_STEPS: AtomicU64 = AtomicU64::new(0);
2517 let anat = {
2518 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2519 *ON.get_or_init(|| std::env::var("MEMRA_SPEC_ANATOMY").as_deref() == Ok("1"))
2520 };
2521 if anat {
2522 e.stream().synchronize()?; // drain prior queue so phase 0 starts clean
2523 }
2524 let t_all = std::time::Instant::now();
2525 let mut t_ph = std::time::Instant::now();
2526 let mut anat_mark = |i: usize,
2527 e: &Engine,
2528 t: &mut std::time::Instant|
2529 -> Result<(), Box<dyn std::error::Error>> {
2530 if anat {
2531 e.stream().synchronize()?;
2532 ANAT_NS[i].fetch_add(t.elapsed().as_nanos() as u64, Relaxed);
2533 *t = std::time::Instant::now();
2534 }
2535 Ok(())
2536 };
2537 let cfg = &self.cfg;
2538 let n_embd = cfg.n_embd as usize;
2539 // Distilled-student geometry: the block runs at the INNER width `di` (eh_proj out /
2540 // attn / ffn); the n_embd interface (embed, norms in, carrier out, head in) is unchanged.
2541 let di = mtp.geom.as_ref().map(|g| g.d_inner).unwrap_or(n_embd);
2542 let eps = cfg.rms_eps;
2543 let pos_d = e.htod_i32(&[mtp_pos as i32])?;
2544
2545 // op A: a resident table transfers one 4B token id. The exact host-row capacity path
2546 // expands this one row on CPU and transfers n_embd f32 values instead.
2547 let e_emb = match embd_dev {
2548 Some((g, qt, rb)) => e.embed_gather_device_t(g, &[e_tok], n_embd, qt, rb)?,
2549 None => e.htod(&self.embd.gather(n_embd, &[e_tok]))?,
2550 };
2551
2552 // op 1/2: e_norm = RMSNorm(e, enorm); h_norm = RMSNorm(h_seed, hnorm)
2553 let mut e_norm = e.zeros(n_embd)?;
2554 e.rms_norm(&e_emb, mtp.enorm.float_data(), &mut e_norm, n_embd, 1, eps)?;
2555 let mut h_norm = e.zeros(n_embd)?;
2556 e.rms_norm(h_seed, mtp.hnorm.float_data(), &mut h_norm, n_embd, 1, eps)?;
2557
2558 // op 3: concat = [e_norm ; h_norm] -> [2*n_embd], e_norm in [0,n_embd), h_norm in [n_embd,2n_embd)
2559 let mut concat = e.zeros(2 * n_embd)?;
2560 e.copy_into(&mut concat, 0, &e_norm, n_embd)?;
2561 e.copy_into(&mut concat, n_embd, &h_norm, n_embd)?;
2562
2563 // op 4: inpSA = eh_proj @ concat (eh_proj [2*n_embd, n_embd]) -> [n_embd]
2564 let inp_sa = e.matmul(&mtp.eh_proj, &concat, 1)?;
2565
2566 // op 5: a_norm = RMSNorm(inpSA, attn_norm)
2567 let mut a_norm = e.zeros(di)?;
2568 e.rms_norm(&inp_sa, mtp.attn_norm.float_data(), &mut a_norm, di, 1, eps)?;
2569 anat_mark(0, e, &mut t_ph)?;
2570
2571 // op 6: attention on the scratch KV. SAME dc launcher as the graph path (bucket_max =
2572 // scratch.cap, length from the device len_d) so eager drafts match graph drafts
2573 // bit-for-bit at any t_kv (the parity gate). Host len mirrored here (the dc append
2574 // advances only the device counter).
2575 let attn_out = match (&mtp.mixer, mtp.step35.as_ref()) {
2576 // step35 MTP block: PER-LAYER geometry + a separate head-wise gate + an SWA window,
2577 // none of which the dc launcher can express (see `mtp_step35_attn`). Host-len arm.
2578 // Advances BOTH the host len and the device counter itself (unlike the dc arm,
2579 // whose host-side mirror the caller does).
2580 (Mixer::Full(fa), Some(g)) => {
2581 self.mtp_step35_attn(e, fa, g, &a_norm, &pos_d, scratch)?
2582 }
2583 (Mixer::Full(fa), None) => {
2584 let out =
2585 self.mtp_full_attn_dc(e, fa, &a_norm, &pos_d, scratch, mtp.geom.as_ref())?;
2586 scratch.kv.len += 1;
2587 out
2588 }
2589 (Mixer::Linear(_), _) => {
2590 panic!("MTP block is full-attn in qwen35; linear MTP not supported")
2591 }
2592 (Mixer::Mla(_), _) => crate::hybrid::mla_forward_unimplemented(),
2593 };
2594 anat_mark(1, e, &mut t_ph)?;
2595
2596 // op 7: x1 = inpSA + attn_out
2597 let mut x1 = e.zeros(di)?;
2598 e.add(&inp_sa, &attn_out, &mut x1, di)?;
2599
2600 // op 8: z = RMSNorm(x1, post_attn_norm) (pre-FFN norm)
2601 let mut z = e.zeros(di)?;
2602 e.rms_norm(&x1, mtp.post_attn_norm.float_data(), &mut z, di, 1, eps)?;
2603
2604 // op 9: FFN (Dense or MoE) — same as the trunk decode FFN
2605 let ffn_out = match &mtp.ffn {
2606 crate::hybrid::Ffn::Dense {
2607 ffn_gate,
2608 ffn_up,
2609 ffn_down,
2610 } => {
2611 let n_ff = ffn_gate.out_features();
2612 let (gate, up) = if e.uses_q8_1_fast(ffn_gate) && e.uses_q8_1_fast(ffn_up) {
2613 let (zq, zd) = e.quantize_q8_1(&z, 1, di)?;
2614 (
2615 e.matmul_pre(ffn_gate, &zq, &zd, &z, 1)?,
2616 e.matmul_pre(ffn_up, &zq, &zd, &z, 1)?,
2617 )
2618 } else {
2619 (e.matmul(ffn_gate, &z, 1)?, e.matmul(ffn_up, &z, 1)?)
2620 };
2621 let mut act = e.zeros(n_ff)?;
2622 // step35: a DENSE FFN reads the per-layer SHEXP clamp (upstream's one `build_ffn`
2623 // serves the dense MLP and the shared expert off `swiglu_clamp_shexp` —
2624 // llama-graph.cpp:1751), resolved for the MTP block's OWN index. Every other arch
2625 // passes None, which is `ffn_act`'s dispatch verbatim.
2626 Self::ffn_act_lim(
2627 e,
2628 &self.cfg,
2629 &gate,
2630 &up,
2631 1.0,
2632 1.0,
2633 mtp.step35.as_ref().and_then(|s| s.clamp_shexp),
2634 &mut act,
2635 n_ff,
2636 )?;
2637 e.matmul(ffn_down, &act, 1)?
2638 }
2639 // MTP head is a distinct block — key its experts under a separate layer index (u16::MAX)
2640 // so they never alias trunk layer 0's cache keys.
2641 crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il(e, m, &z, 1, u16::MAX)?,
2642 };
2643 anat_mark(2, e, &mut t_ph)?;
2644
2645 // op 10: h_nextn = x1 + ffn_out (at di)
2646 let mut h_inner = e.zeros(di)?;
2647 e.add(&x1, &ffn_out, &mut h_inner, di)?;
2648
2649 // op 10.5 (student): up-project the inner hidden back to n_embd — training semantics:
2650 // the chain carrier AND the head input are out_up(h_inner) (pre-final-norm).
2651 let h_nextn = match mtp.geom.as_ref() {
2652 Some(g) => e.matmul(&g.out_up, &h_inner, 1)?,
2653 None => h_inner,
2654 };
2655
2656 // op 11: final = RMSNorm(h_nextn, shared_head_norm OR output_norm)
2657 let final_norm = mtp.shared_head_norm.as_ref().unwrap_or(&self.output_norm);
2658 let mut final_h = e.zeros(n_embd)?;
2659 e.rms_norm(
2660 &h_nextn,
2661 final_norm.float_data(),
2662 &mut final_h,
2663 n_embd,
2664 1,
2665 eps,
2666 )?;
2667
2668 // op 12: draft_logits = (shared_head_head OR output) @ final — stays ON DEVICE.
2669 let head = mtp.shared_head_head.as_ref().unwrap_or(&self.output);
2670 let mut logits = e.matmul(head, &final_h, 1)?;
2671 // op 12b (lane/draft-mask): grammar mask over the DRAFT vocab, applied here so the
2672 // caller's argmax / gumbel draw / p-min prob all read the grammar-legal row.
2673 if let Some((mask_d, mw)) = mask {
2674 let d_vocab = head.out_features();
2675 e.mask_logits_col(&mut logits, mask_d, 0, d_vocab, mw)?;
2676 }
2677 anat_mark(3, e, &mut t_ph)?;
2678 if anat {
2679 ANAT_NS[4].fetch_add(t_all.elapsed().as_nanos() as u64, Relaxed);
2680 let n = ANAT_STEPS.fetch_add(1, Relaxed) + 1;
2681 if n % 128 == 0 {
2682 let us = |i: usize| ANAT_NS[i].load(Relaxed) / n / 1000;
2683 eprintln!(
2684 "[spec-anatomy] steps={n} avg us/step: glue={} attn={} ffn={} head={} total={}",
2685 us(0),
2686 us(1),
2687 us(2),
2688 us(3),
2689 us(4)
2690 );
2691 }
2692 }
2693 // Chain recurrence hand-over: pre-norm h_nextn (default) or post-norm final_h
2694 // (MEMRA_SPEC_HPOST — llama.cpp #24025's t_h_nextn is taken AFTER the head norm).
2695 Ok((logits, if spec_hpost() { final_h } else { h_nextn }))
2696 }
2697
2698 /// step35 MTP-block attention, T=1, on the scratch KV — the EAGER-ONLY twin of
2699 /// `mtp_full_attn_dc`. Three things force a separate arm rather than a geometry parameter on
2700 /// the dc path, and all three are properties of this arch's MTP block:
2701 ///
2702 /// 1. **The SWA window.** Block 45 is an SWA-type block (`sliding_window_pattern[45]=true`,
2703 /// window 512). Windowed decode in memra is a token-aligned VIEW OFFSET into the quantized
2704 /// cache (the gemma4 R6 / `step35_decode_attn` pattern: keys carry absolute rope and the
2705 /// mask is purely positional, so one query at `len-1` attending the last `win` rows IS the
2706 /// windowed result). `fa_decode_dc` takes the key count from a DEVICE counter and always
2707 /// starts at row 0 — it cannot express a nonzero offset, so a windowed dc arm would need a
2708 /// new kernel. That is deliberately not built here: see the CUDA-graph note below.
2709 /// 2. **Per-layer head count.** 96 q heads over 8 KV (GQA 12) at this block, vs the trunk's 64
2710 /// on its full-attn layers. The trunk cfg's `n_head` scalar is the MAX over layers, and the
2711 /// trunk ARTIFACT's per-layer arrays stop at index 44 — so the count must come from the
2712 /// resolved `Step35MtpGeom`, never from `cfg`.
2713 /// 3. **The separate head-wise gate.** `blk.45.attn_gate.weight [n_embd, 96]` produces one
2714 /// sigmoid scalar per head (broadcast over head_dim) — `attn_head_gate`, not the qwen35
2715 /// fused-into-wq `q_gate_split` form the dc arm handles.
2716 ///
2717 /// WHY EAGER-ONLY IS NOT A GAP TODAY: `mtp_head_forward_cap` refuses step35 heads
2718 /// explicitly (the SWA-window refusal below), so the graph draft never engages for step35
2719 /// models regardless of eligibility — the eager chain IS the served path. `mtp_head_forward_cap` refuses step35 explicitly rather
2720 /// than silently capturing a window-less (wrong past `win` draft rows) graph.
2721 ///
2722 /// Unlike the dc arm this advances BOTH the host `kv.len` and the device counter, so the
2723 /// caller must not mirror.
2724 fn mtp_step35_attn(
2725 &self,
2726 e: &Engine,
2727 fa: &FullAttnLayer,
2728 g: &crate::hybrid::Step35MtpGeom,
2729 h: &CudaSlice<f32>,
2730 pos_d: &CudaSlice<i32>,
2731 scratch: &mut MtpScratch,
2732 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2733 let (nh, nkv, hd) = (g.n_head, g.n_head_kv, self.cfg.head_dim_k as usize);
2734 let eps = self.cfg.rms_eps;
2735 let scale = 1.0 / (hd as f32).sqrt(); // step35.cpp:255 kq_scale
2736 let n_embd = self.cfg.n_embd as usize;
2737 let gw = fa
2738 .attn_gate
2739 .as_ref()
2740 .ok_or("step35 MTP block is missing attn_gate.weight (head-wise attention gate)")?;
2741
2742 let (q0, k0, v0, gt) = if e.uses_q8_1_fast(&fa.wq)
2743 && e.uses_q8_1_fast(&fa.wk)
2744 && e.uses_q8_1_fast(&fa.wv)
2745 && e.uses_q8_1_fast(gw)
2746 {
2747 let (hq, hdq) = e.quantize_q8_1(h, 1, n_embd)?;
2748 let (a, b, c) = match e.matmul_q8_fused3(&fa.wq, &fa.wk, &fa.wv, &hq, &hdq)? {
2749 Some(t3) => t3,
2750 None => (
2751 e.matmul_pre(&fa.wq, &hq, &hdq, h, 1)?,
2752 e.matmul_pre(&fa.wk, &hq, &hdq, h, 1)?,
2753 e.matmul_pre(&fa.wv, &hq, &hdq, h, 1)?,
2754 ),
2755 };
2756 (a, b, c, e.matmul_pre(gw, &hq, &hdq, h, 1)?)
2757 } else {
2758 (
2759 e.matmul(&fa.wq, h, 1)?,
2760 e.matmul(&fa.wk, h, 1)?,
2761 e.matmul(&fa.wv, h, 1)?,
2762 e.matmul(gw, h, 1)?,
2763 )
2764 };
2765
2766 let mut q = e.uninit(nh * hd)?;
2767 e.rms_norm(&q0, fa.q_norm.float_data(), &mut q, hd, nh, eps)?;
2768 let mut k = e.uninit(nkv * hd)?;
2769 e.rms_norm(&k0, fa.k_norm.float_data(), &mut k, hd, nkv, eps)?;
2770 // `rope_freqs.weight` (llama3 factors) applies to the FULL-attn layers ONLY; SWA passes
2771 // null (llama-hparams / step35.cpp). Block 45 is SWA, so `ff` is None there — but read
2772 // the resolved flag, not the constant, so an all-full sibling stays correct.
2773 let ff = if g.swa {
2774 None
2775 } else {
2776 self.step35_aux.as_ref().and_then(|a| a.rope_freqs(e))
2777 };
2778 #[cfg(debug_assertions)]
2779 if let Some(ff) = ff {
2780 crate::debug_assert_tensor_stream_device(ff, &e.stream(), "mtp_step35_attn.rope_freqs");
2781 }
2782 e.rope_neox2(
2783 &mut q,
2784 &mut k,
2785 pos_d,
2786 hd,
2787 g.n_rot,
2788 nh,
2789 nkv,
2790 1,
2791 g.rope_base,
2792 1.0,
2793 ff,
2794 )?;
2795
2796 // Append at the HOST slot, then re-stamp the device counter: the eager chain has the
2797 // length on the host anyway, and the windowed view below needs it there to compute the
2798 // offset. The device counter is kept in lockstep so `mtp_kv_fill`'s `set_i32_one` and any
2799 // dc-family consumer of this scratch still agree.
2800 let kv = &mut scratch.kv;
2801 assert!(
2802 kv.len < scratch.cap,
2803 "step35 MTP scratch overflow ({} >= {})",
2804 kv.len,
2805 scratch.cap
2806 );
2807 let next_len = kv.len + 1;
2808 let (off, t_kv) = if g.swa && next_len > g.window {
2809 (next_len - g.window, g.window)
2810 } else {
2811 (0, next_len)
2812 };
2813 let write_row = e.prepare_kv_append(kv, off & !31usize, 1)?;
2814 e.append_kv_quantized(
2815 &k,
2816 &v0,
2817 &mut kv.k,
2818 &mut kv.v,
2819 write_row,
2820 kv.kv_dim_k,
2821 kv.kv_dim_v,
2822 kv.k_tok_bytes,
2823 kv.v_tok_bytes,
2824 false,
2825 )?;
2826 kv.len = next_len;
2827 e.set_i32_one(&mut kv.len_d, kv.len as i32)?;
2828 // SWA view offset (see note 1). The draft chain is short (k+2 rows), but the scratch is
2829 // PERSISTENT across rounds — `mtp_kv_fill` leaves one row per committed token behind, so
2830 // `kv.len` tracks absolute position and crosses 512 in any real generation. The window is
2831 // therefore live, not theoretical.
2832 let physical = kv.physical_rows(off, off + t_kv)?;
2833 let k_view = e.view_u8_range(
2834 &kv.k,
2835 physical.start * kv.k_tok_bytes,
2836 physical.end * kv.k_tok_bytes,
2837 );
2838 let v_view = e.view_u8_range(
2839 &kv.v,
2840 physical.start * kv.v_tok_bytes,
2841 physical.end * kv.v_tok_bytes,
2842 );
2843 let mut attn = e.uninit(nh * hd)?;
2844 e.fa_decode_kvmod(
2845 &q,
2846 &k_view,
2847 &v_view,
2848 &mut attn,
2849 hd,
2850 nh,
2851 nkv,
2852 t_kv,
2853 scale,
2854 kv.k_tok_bytes,
2855 kv.v_tok_bytes,
2856 false,
2857 )?;
2858
2859 let mut ag = e.uninit(nh * hd)?;
2860 e.attn_head_gate(&attn, >, &mut ag, None, hd, nh, 1)?;
2861 Ok(e.matmul(&fa.wo, &ag, 1)?)
2862 }
2863
2864 /// MTP-block full attention, T=1, on the scratch KV (BOTH draft paths — eager and graph):
2865 /// the scratch write slot and the attention bound come from `scratch.kv.len_d` (device i32[1])
2866 /// so the launch args are FIXED across draft steps — ONE captured graph serves the whole
2867 /// chain, and replays keep seeing KV growth through the device counter (no recapture).
2868 /// Geometry contract: n_splits is sized from `scratch.cap` (the persistent capacity); splits
2869 /// whose key range lies beyond the device t_kv exit empty and the shared combine skips them
2870 /// (fa_decode_dc bit-correct-for-any-t_kv<=bucket_max contract). The eager path uses the SAME
2871 /// launcher with the SAME bucket_max -> identical dispatch -> bit-identical draft tokens (the
2872 /// graph-vs-eager parity gate). Host len is NOT advanced here (graph contract); callers mirror.
2873 fn mtp_full_attn_dc(
2874 &self,
2875 e: &Engine,
2876 fa: &FullAttnLayer,
2877 h: &CudaSlice<f32>,
2878 pos_d: &CudaSlice<i32>,
2879 scratch: &mut MtpScratch,
2880 geom: Option<&crate::hybrid::DraftGeom>,
2881 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2882 let cfg = &self.cfg;
2883 let mtp_il = cfg.n_layer.saturating_sub(cfg.nextn_predict_layers);
2884 let geometry = cfg.full_attention_geometry_at(mtp_il);
2885 let n_head = geom.map(|g| g.n_head).unwrap_or(geometry.n_head as usize);
2886 let n_head_kv = geom
2887 .map(|g| g.n_head_kv)
2888 .unwrap_or(geometry.n_head_kv as usize);
2889 let head_dim = geometry.head_dim_k as usize;
2890 let eps = cfg.rms_eps;
2891 let scale = geometry.attention_scale();
2892 let n_embd = geom.map(|g| g.d_inner).unwrap_or(cfg.n_embd as usize);
2893 let bucket_max = scratch.cap; // < 96 guaranteed by the graph_draft eligibility gate
2894
2895 let (qf, mut k, v) =
2896 if e.uses_q8_1_fast(&fa.wq) && e.uses_q8_1_fast(&fa.wk) && e.uses_q8_1_fast(&fa.wv) {
2897 let (hq, hd) = e.quantize_q8_1(h, 1, n_embd)?;
2898 (
2899 e.matmul_pre(&fa.wq, &hq, &hd, h, 1)?,
2900 e.matmul_pre(&fa.wk, &hq, &hd, h, 1)?,
2901 e.matmul_pre(&fa.wv, &hq, &hd, h, 1)?,
2902 )
2903 } else {
2904 (
2905 e.matmul(&fa.wq, h, 1)?,
2906 e.matmul(&fa.wk, h, 1)?,
2907 e.matmul(&fa.wv, h, 1)?,
2908 )
2909 };
2910 // M3/Hy3 have no attention output gate — wq out is exactly q; skip the split.
2911 let gated = geometry.attention_gate == memra_gguf::config::AttentionGateKind::FusedQ;
2912 let (mut q, gate) = if gated {
2913 let mut q = e.zeros(n_head * head_dim)?;
2914 let mut gate = e.zeros(n_head * head_dim)?;
2915 e.q_gate_split(&qf, &mut q, &mut gate, head_dim, n_head, 1)?;
2916 (q, Some(gate))
2917 } else {
2918 (qf, None)
2919 };
2920
2921 let mut qn = e.zeros(n_head * head_dim)?;
2922 e.rms_norm(&q, fa.q_norm.float_data(), &mut qn, head_dim, n_head, eps)?;
2923 q = qn;
2924 let mut kn = e.zeros(n_head_kv * head_dim)?;
2925 e.rms_norm(
2926 &k,
2927 fa.k_norm.float_data(),
2928 &mut kn,
2929 head_dim,
2930 n_head_kv,
2931 eps,
2932 )?;
2933 k = kn;
2934 let rope_dims = geometry.n_rot as usize;
2935 e.rope_neox(
2936 &mut q,
2937 pos_d,
2938 head_dim,
2939 rope_dims,
2940 n_head,
2941 1,
2942 geometry.rope_base,
2943 1.0,
2944 )?;
2945 e.rope_neox(
2946 &mut k,
2947 pos_d,
2948 head_dim,
2949 rope_dims,
2950 n_head_kv,
2951 1,
2952 geometry.rope_base,
2953 1.0,
2954 )?;
2955
2956 let kv = &mut scratch.kv;
2957 // append at the DEVICE slot (kv.len_d == old len), then advance the counter in-graph.
2958 e.append_kv_quantized_dc(
2959 &k,
2960 &v,
2961 &mut kv.k,
2962 &mut kv.v,
2963 &kv.len_d,
2964 kv.kv_dim_k,
2965 kv.kv_dim_v,
2966 kv.k_tok_bytes,
2967 kv.v_tok_bytes,
2968 false,
2969 )?;
2970 e.inc_seqlen(&mut kv.len_d)?;
2971 // full-buffer views (any in-round t_kv stays in range on replay); the kernel bounds the
2972 // key range from the device counter.
2973 let k_view = e.view_u8(&kv.k, kv.k.len());
2974 let v_view = e.view_u8(&kv.v, kv.v.len());
2975 let (ktb, vtb) = (kv.k_tok_bytes, kv.v_tok_bytes);
2976 let mut attn = e.zeros(n_head * head_dim)?;
2977 e.fa_decode_dc(
2978 &q, &k_view, &v_view, &mut attn, head_dim, n_head, n_head_kv, &kv.len_d, bucket_max,
2979 scale, ktb, vtb, false,
2980 )?;
2981
2982 let attn_g = match &gate {
2983 Some(gate) => {
2984 let mut gsig = e.zeros(n_head * head_dim)?;
2985 e.sigmoid(gate, &mut gsig, n_head * head_dim)?;
2986 let mut ag = e.zeros(n_head * head_dim)?;
2987 e.mul(&attn, &gsig, &mut ag, n_head * head_dim)?;
2988 ag
2989 }
2990 None => attn,
2991 };
2992 Ok(e.matmul(&fa.wo, &attn_g, 1)?)
2993 }
2994
2995 /// PERSISTENT-DRAFT-KV fill (the reference engine's "mtp_update" analogue): compute the MTP
2996 /// block's K/V for `tokens` (committed tokens at positions pos0..pos0+T) from their EXACT
2997 /// trunk hiddens `h` ([T, n_embd] token-major, pre-output_norm) and append at slots pos0.. of
2998 /// the scratch KV. K/V-ONLY — ops A/1-5 plus the K-side of op 6 (wk/wv + k_norm + rope +
2999 /// quantized append); no wq/attention/FFN/lm_head, so per-token cost ~= eh_proj + wk/wv (a
3000 /// small fraction of one trunk layer), T-batched. Rope follows the chain convention
3001 /// rope(token@p) = p+1. Runs at round boundaries OUTSIDE the captured graph in BOTH draft
3002 /// modes -> draft parity by construction. Caller must have scratch.kv.len == pos0.
3003 #[allow(clippy::too_many_arguments)]
3004 fn mtp_kv_fill(
3005 &self,
3006 e: &Engine,
3007 mtp: &MtpHead,
3008 tokens: &[u32],
3009 h: &CudaSlice<f32>,
3010 pos0: usize,
3011 scratch: &mut MtpScratch,
3012 embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
3013 ) -> Result<(), Box<dyn std::error::Error>> {
3014 let cfg = &self.cfg;
3015 let n_embd = cfg.n_embd as usize;
3016 let eps = cfg.rms_eps;
3017 let t = tokens.len();
3018 assert_eq!(scratch.kv.len, pos0, "mtp_kv_fill: append slot mismatch");
3019 assert!(pos0 + t <= scratch.cap, "mtp_kv_fill: scratch overflow");
3020 let Mixer::Full(fa) = &mtp.mixer else {
3021 panic!("MTP block is full-attn in qwen35; linear MTP not supported")
3022 };
3023 let pos_vec: Vec<i32> = (0..t).map(|i| (pos0 + i + 1) as i32).collect();
3024 let pos_d = e.htod_i32(&pos_vec)?;
3025
3026 // ops A/1/2: embed + the two input norms, T-wide.
3027 let e_emb = match embd_dev {
3028 Some((g, qt, rb)) => e.embed_gather_device_t(g, tokens, n_embd, qt, rb)?,
3029 None => e.htod(&self.embd.gather(n_embd, tokens))?,
3030 };
3031 let mut e_norm = e.zeros(t * n_embd)?;
3032 e.rms_norm(&e_emb, mtp.enorm.float_data(), &mut e_norm, n_embd, t, eps)?;
3033 let mut h_norm = e.zeros(t * n_embd)?;
3034 e.rms_norm(h, mtp.hnorm.float_data(), &mut h_norm, n_embd, t, eps)?;
3035
3036 // op 3: per-row [e_norm ; h_norm] concat, token-major [T, 2*n_embd].
3037 let mut concat = e.zeros(t * 2 * n_embd)?;
3038 for i in 0..t {
3039 e.copy_view_into(
3040 &mut concat,
3041 i * 2 * n_embd,
3042 &e_norm.slice(i * n_embd..(i + 1) * n_embd),
3043 n_embd,
3044 )?;
3045 e.copy_view_into(
3046 &mut concat,
3047 i * 2 * n_embd + n_embd,
3048 &h_norm.slice(i * n_embd..(i + 1) * n_embd),
3049 n_embd,
3050 )?;
3051 }
3052
3053 // ops 4/5: eh_proj + attn_norm, T-wide (at the student inner width when geom is set).
3054 let di = mtp.geom.as_ref().map(|g| g.d_inner).unwrap_or(n_embd);
3055 let inp_sa = e.matmul(&mtp.eh_proj, &concat, t)?;
3056 let mut a_norm = e.zeros(t * di)?;
3057 e.rms_norm(&inp_sa, mtp.attn_norm.float_data(), &mut a_norm, di, t, eps)?;
3058
3059 // op 6 (K/V half): wk/wv + k_norm + rope + per-row quantized append. No wq/attention —
3060 // the fill only has to leave correct K/V rows behind for later chains to attend over.
3061 let n_head_kv = mtp
3062 .geom
3063 .as_ref()
3064 .map(|g| g.n_head_kv)
3065 .or(mtp.step35.as_ref().map(|s| s.n_head_kv))
3066 .unwrap_or_else(|| {
3067 let mtp_il = cfg.n_layer.saturating_sub(cfg.nextn_predict_layers);
3068 cfg.full_attention_geometry_at(mtp_il).n_head_kv as usize
3069 });
3070 let mtp_il = cfg.n_layer.saturating_sub(cfg.nextn_predict_layers);
3071 let geometry = cfg.full_attention_geometry_at(mtp_il);
3072 let head_dim = geometry.head_dim_k as usize;
3073 let mut k = e.matmul(&fa.wk, &a_norm, t)?;
3074 let v = e.matmul(&fa.wv, &a_norm, t)?;
3075 let mut kn = e.zeros(t * n_head_kv * head_dim)?;
3076 e.rms_norm(
3077 &k,
3078 fa.k_norm.float_data(),
3079 &mut kn,
3080 head_dim,
3081 n_head_kv * t,
3082 eps,
3083 )?;
3084 k = kn;
3085 // step35: rotary width AND base are per-layer, and the MTP block's values come from the
3086 // resolved `Step35MtpGeom` — NOT from `cfg.rope_dim_count`/`cfg.rope_freq_base`, which
3087 // carry the arch defaults (128 / 5e6, i.e. the FULL-attn layers' base). Getting this wrong
3088 // writes K rows the attention arm then re-derives at a different theta: correct-looking
3089 // output with dead acceptance, invisible to the exactness gates.
3090 let (rope_dims, rope_base, ff) = match mtp.step35.as_ref() {
3091 Some(s) => (
3092 s.n_rot,
3093 s.rope_base,
3094 if s.swa {
3095 None
3096 } else {
3097 self.step35_aux.as_ref().and_then(|a| a.rope_freqs(e))
3098 },
3099 ),
3100 None => (geometry.n_rot as usize, geometry.rope_base, None),
3101 };
3102 #[cfg(debug_assertions)]
3103 if let Some(ff) = ff {
3104 crate::debug_assert_tensor_stream_device(ff, &e.stream(), "mtp_kv_fill.rope_freqs");
3105 }
3106 match ff {
3107 Some(f) => e.rope_neox_ff(
3108 &mut k, &pos_d, head_dim, rope_dims, n_head_kv, t, rope_base, 1.0, f,
3109 )?,
3110 None => e.rope_neox(
3111 &mut k, &pos_d, head_dim, rope_dims, n_head_kv, t, rope_base, 1.0,
3112 )?,
3113 }
3114
3115 let kv = &mut scratch.kv;
3116 // Match the trunk prime contract: a chunk may need the aligned window immediately before
3117 // its first row, so preserve that prefix when the physical tail rebases at wrap.
3118 let retain_from = kv
3119 .ring
3120 .as_ref()
3121 .map(|ring| pos0.saturating_sub(ring.window() - 1) & !31usize)
3122 .unwrap_or(0);
3123 let write_row = e.prepare_kv_append(kv, retain_from, t)?;
3124 for i in 0..t {
3125 let k_row = k.slice(i * kv.kv_dim_k..(i + 1) * kv.kv_dim_k);
3126 let v_row = v.slice(i * kv.kv_dim_v..(i + 1) * kv.kv_dim_v);
3127 e.append_kv_quantized_view(
3128 &k_row,
3129 &v_row,
3130 &mut kv.k,
3131 &mut kv.v,
3132 write_row + i,
3133 kv.kv_dim_k,
3134 kv.kv_dim_v,
3135 kv.k_tok_bytes,
3136 kv.v_tok_bytes,
3137 false,
3138 )?;
3139 }
3140 kv.len = pos0 + t;
3141 e.set_i32_one(&mut kv.len_d, kv.len as i32)?;
3142 Ok(())
3143 }
3144
3145 /// CAPTURE body for the GRAPH DRAFT (stage 2 of graph-grade spec): ONE MTP head forward with
3146 /// every varying input device-resident —
3147 /// - token id from the persistent `tok_d` (the previous replay's in-graph argmax wrote it,
3148 /// so the chain feeds itself; the host reads the same 4 bytes for the draft list),
3149 /// - h_seed from the persistent `h_seed_d` (h_nextn is copied BACK into it at the end),
3150 /// - rope pos from the persistent `pos_d` counter (inc'd in-graph),
3151 /// - scratch KV slot/bound from `scratch.kv.len_d` (see mtp_full_attn_dc).
3152 /// The p-min confidence lands in the persistent `p_d` iff `with_prob` (env is fixed per run).
3153 /// Same kernels, same dispatch as the eager mtp_head_forward_dev chain -> same draft tokens
3154 /// (exactness never depends on drafts — the verify arbitrates — but acceptance parity does).
3155 /// `with_head=false` captures the HEAD-LESS twin for the pseudo-seed replay (2026-07-03):
3156 /// the pseudo pass only needs h_nextn (op 10) + the scratch append — the lm_head read
3157 /// (~1.06ms q6_K on the 9B), argmax and prob are dead weight there. h_nextn's inputs are
3158 /// untouched, so the seed value is identical; round-start resets overwrite tok_d/p_d anyway.
3159 /// `sampled_cap` = Some((ctr_d, perturb_d, q_out_d, seed, temp)) captures the SAMPLED twin
3160 /// (step 3 of the sampled-spec arc): head logits are retained in the persistent `q_out_d`
3161 /// (host D2Ds them to the round's q slot after each replay), the DEVICE event counter is
3162 /// bumped in-graph, and the argmax reads GUMBEL-PERTURBED logits — one categorical draw per
3163 /// replay, bit-identical to the eager arm's gumbel_perturb at the same (seed, sctr, temp).
3164 /// seed/temp are capture-time constants (fixed per generate call, like p_min).
3165 #[allow(clippy::too_many_arguments)]
3166 fn mtp_head_forward_cap(
3167 &self,
3168 e: &Engine,
3169 mtp: &MtpHead,
3170 tok_d: &mut CudaSlice<u32>,
3171 pos_d: &mut CudaSlice<i32>,
3172 h_seed_d: &mut CudaSlice<f32>,
3173 p_d: &mut CudaSlice<f32>,
3174 scratch: &mut MtpScratch,
3175 with_prob: bool,
3176 with_head: bool,
3177 embd_gpu: &CudaSlice<u8>,
3178 embd_qt: i32,
3179 embd_rb: usize,
3180 d_vocab: usize,
3181 sampled_cap: Option<(
3182 &mut CudaSlice<u32>,
3183 &mut CudaSlice<f32>,
3184 &mut CudaSlice<f32>,
3185 u64,
3186 f32,
3187 )>,
3188 stream_pack: Option<(&mut CudaSlice<u32>, usize, Option<&CudaSlice<u32>>)>,
3189 // DRAFT-SIDE GRAMMAR MASK (lane/draft-mask): (packed draft-vocab allowed-set buffer,
3190 // word count). Captured as ONE mask_logits_f32 node between the head matmul and the
3191 // in-graph argmax; the buffer address is baked, its CONTENTS are re-uploaded by the
3192 // host before every replay (the decode.rs graph-mask pattern). All-ones contents = a
3193 // no-op ban, so a position the grammar cannot constrain costs one pass over the row.
3194 mask_cap: Option<(&CudaSlice<u32>, usize)>,
3195 ) -> Result<(), Box<dyn std::error::Error>> {
3196 let cfg = &self.cfg;
3197 let n_embd = cfg.n_embd as usize;
3198 // step35 REFUSAL (deliberate, named): the graph body's attention is `mtp_full_attn_dc`,
3199 // whose device-counter key bound always starts at row 0 — it cannot express this block's
3200 // SWA view offset, so a captured chain would silently attend OUTSIDE the window once the
3201 // persistent scratch passes 512 rows. Nothing is lost today: `mtp_head_forward_cap`
3202 // refuses step35 heads explicitly (SWA refusal), so the eager chain
3203 // (`mtp_head_forward_dev` -> `mtp_step35_attn`) is the served path. Returning Err (not a
3204 // panic) is what the two capture sites and the round-stream capture already handle by
3205 // degrading to eager / stream-off.
3206 if mtp.step35.is_some() {
3207 return Err(
3208 "step35 has no captured draft chain (fa_decode_dc cannot express the MTP \
3209 block's SWA view offset; same root cause as the dc decode refusal) — the \
3210 eager draft chain serves this arch"
3211 .into(),
3212 );
3213 }
3214 // student inner width (see mtp_head_forward_dev) — interface dims stay n_embd.
3215 let di = mtp.geom.as_ref().map(|g| g.d_inner).unwrap_or(n_embd);
3216 let eps = cfg.rms_eps;
3217 let e_emb = e.embed_gather_device(embd_gpu, tok_d, n_embd, embd_qt, embd_rb)?;
3218 let mut e_norm = e.zeros(n_embd)?;
3219 e.rms_norm(&e_emb, mtp.enorm.float_data(), &mut e_norm, n_embd, 1, eps)?;
3220 let mut h_norm = e.zeros(n_embd)?;
3221 e.rms_norm(
3222 &*h_seed_d,
3223 mtp.hnorm.float_data(),
3224 &mut h_norm,
3225 n_embd,
3226 1,
3227 eps,
3228 )?;
3229 let mut concat = e.zeros(2 * n_embd)?;
3230 e.copy_into(&mut concat, 0, &e_norm, n_embd)?;
3231 e.copy_into(&mut concat, n_embd, &h_norm, n_embd)?;
3232 let inp_sa = e.matmul(&mtp.eh_proj, &concat, 1)?;
3233 let mut a_norm = e.zeros(di)?;
3234 e.rms_norm(&inp_sa, mtp.attn_norm.float_data(), &mut a_norm, di, 1, eps)?;
3235 let attn_out = match &mtp.mixer {
3236 Mixer::Full(fa) => {
3237 self.mtp_full_attn_dc(e, fa, &a_norm, pos_d, scratch, mtp.geom.as_ref())?
3238 }
3239 Mixer::Linear(_) => {
3240 panic!("MTP block is full-attn in qwen35; linear MTP not supported")
3241 }
3242 Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
3243 };
3244 let mut x1 = e.zeros(di)?;
3245 e.add(&inp_sa, &attn_out, &mut x1, di)?;
3246 let mut z = e.zeros(di)?;
3247 e.rms_norm(&x1, mtp.post_attn_norm.float_data(), &mut z, di, 1, eps)?;
3248 let ffn_out = match &mtp.ffn {
3249 crate::hybrid::Ffn::Dense {
3250 ffn_gate,
3251 ffn_up,
3252 ffn_down,
3253 } => {
3254 let n_ff = ffn_gate.out_features();
3255 let (gate, up) = if e.uses_q8_1_fast(ffn_gate) && e.uses_q8_1_fast(ffn_up) {
3256 let (zq, zd) = e.quantize_q8_1(&z, 1, di)?;
3257 (
3258 e.matmul_pre(ffn_gate, &zq, &zd, &z, 1)?,
3259 e.matmul_pre(ffn_up, &zq, &zd, &z, 1)?,
3260 )
3261 } else {
3262 (e.matmul(ffn_gate, &z, 1)?, e.matmul(ffn_up, &z, 1)?)
3263 };
3264 let mut act = e.zeros(n_ff)?;
3265 Self::ffn_act(e, &self.cfg, &gate, &up, &mut act, n_ff)?;
3266 e.matmul(ffn_down, &act, 1)?
3267 }
3268 // ROUND-STREAM: the 35B NextN block carries a MoE FFN. With RESIDENT experts the
3269 // dev path is pure device launches (device top-k + rows kernels, ZERO-DtoH by
3270 // design) — capture-legal. Non-resident (SLRU-lock) stays rejected: the capture
3271 // error arm degrades the caller to eager/stream-off.
3272 crate::hybrid::Ffn::Moe(m) if m.dev_exps.is_some() => {
3273 self.moe_ffn_il(e, m, &z, 1, u16::MAX)?
3274 }
3275 crate::hybrid::Ffn::Moe(_) => {
3276 return Err("graph draft requires a Dense (or resident-MoE) MTP FFN".into());
3277 }
3278 };
3279 let mut h_inner = e.zeros(di)?;
3280 e.add(&x1, &ffn_out, &mut h_inner, di)?;
3281 // student: up-project back to n_embd (carrier + head input; see mtp_head_forward_dev).
3282 let h_nextn = match mtp.geom.as_ref() {
3283 Some(g) => e.matmul(&g.out_up, &h_inner, 1)?,
3284 None => h_inner,
3285 };
3286 // MEMRA_SPEC_HPOST needs final_h even head-less (it IS the next seed under that convention).
3287 let final_h = if with_head || spec_hpost() {
3288 let final_norm = mtp.shared_head_norm.as_ref().unwrap_or(&self.output_norm);
3289 let mut fh = e.zeros(n_embd)?;
3290 e.rms_norm(&h_nextn, final_norm.float_data(), &mut fh, n_embd, 1, eps)?;
3291 Some(fh)
3292 } else {
3293 None
3294 };
3295 if with_head {
3296 let head = mtp.shared_head_head.as_ref().unwrap_or(&self.output);
3297 let mut logits = e.matmul(head, final_h.as_ref().unwrap(), 1)?;
3298 // DRAFT-SIDE GRAMMAR MASK: ban the grammar-illegal draft ids IN the captured chain,
3299 // before the argmax — proposals become legal by construction. Contents-only
3300 // per-replay upload keeps the capture valid.
3301 if let Some((mask_d, mw)) = mask_cap {
3302 e.mask_logits_col(&mut logits, mask_d, 0, d_vocab, mw)?;
3303 }
3304 if let Some((ctr_d, perturb_d, q_out_d, seed, temp)) = sampled_cap {
3305 // SAMPLED chain: retain q (raw head logits -> persistent q_out_d; the matmul's
3306 // own buffer is pool-recycled after the capture body returns, so it can't be the
3307 // retention target), bump the device event counter, gumbel-perturb reading it,
3308 // and argmax the PERTURBED logits into tok_d — the in-graph categorical draw.
3309 e.copy_into(q_out_d, 0, &logits, d_vocab)?;
3310 e.sctr_inc(ctr_d)?;
3311 e.gumbel_perturb_ctr(&logits, perturb_d, d_vocab, seed, ctr_d, temp)?;
3312 e.argmax_token_device_into(perturb_d, tok_d, d_vocab)?;
3313 // p-min prob = the head's RAW softmax confidence in the SAMPLED pick — same
3314 // semantics as the eager sampled arm's prob_of_token_device(dl_d, tok_d).
3315 if with_prob {
3316 e.prob_of_token_device_into(&logits, tok_d, p_d, d_vocab)?;
3317 }
3318 } else {
3319 // draft token -> persistent tok_d (next replay's embed reads it; host reads the 4 bytes).
3320 e.argmax_token_device_into(&logits, tok_d, d_vocab)?;
3321 // p-min under a draft mask reads the MASKED row: confidence relative to the
3322 // grammar-LEGAL alternatives (illegal ids leave the softmax denominator), which
3323 // is the right semantics for "does the drafter know what comes next here" and
3324 // the same row the pick came from. Draft-quality only — verify arbitrates.
3325 if with_prob {
3326 e.prob_of_token_device_into(&logits, tok_d, p_d, d_vocab)?;
3327 }
3328 }
3329 }
3330 // ROUND-STREAM K-chain: pack (tok, p) into slot j, then remap tok through d2t so the
3331 // NEXT chained body's embed reads the TARGET id — zero host involvement per step.
3332 if let Some((out, slot, d2t)) = stream_pack {
3333 e.pack_tok_p(tok_d, p_d, out, slot)?;
3334 if let Some(map) = d2t {
3335 e.tok_map_u32(tok_d, map)?;
3336 }
3337 }
3338 // Next draft step's h_seed: pre-norm h_nextn (default) or post-norm final_h (HPOST).
3339 if spec_hpost() {
3340 e.copy_into(h_seed_d, 0, final_h.as_ref().unwrap(), n_embd)?;
3341 } else {
3342 e.copy_into(h_seed_d, 0, &h_nextn, n_embd)?;
3343 }
3344 // advance the draft rope position in-graph.
3345 e.inc_seqlen(pos_d)?;
3346 Ok(())
3347 }
3348
3349 /// Batched target verify forward over `tokens` at positions `pos0..pos0+T` (§D.3, T=K+1).
3350 /// Returns ALL T logit columns (host f32, [T*n_vocab]); appends T cols to every full-attn KV
3351 /// and advances every linear-attn recur state by T steps (the recur steps are SEQUENTIAL T=1).
3352 /// Advances `cache.pos` by T.
3353 pub fn decode_step_t(
3354 &self,
3355 e: &Engine,
3356 tokens: &[u32],
3357 pos0: usize,
3358 cache: &mut Cache,
3359 ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
3360 if self.is_gemma4_e4b() {
3361 return Ok(self.gemma4_e4b_decode_step_t_h(e, tokens, pos0, cache)?.0);
3362 }
3363 if self.cfg.gemma4.is_some() {
3364 return self.gemma4_decode_step_t(e, tokens, pos0, cache);
3365 }
3366 Ok(self.decode_step_t_h(e, tokens, pos0, cache)?.0)
3367 }
3368
3369 /// Like `decode_step_t` but ALSO returns the LAST column's pre-output_norm hidden (h_seed for
3370 /// the next draft round). This lets partial-accept replay run as ONE batched T=(n_acc+1) forward
3371 /// (single weight read) instead of n_acc+1 separate T=1 decode_steps (n_acc+1 weight reads).
3372 /// At batch=1 decode is bandwidth-bound, so batching the replay is THE MTP profitability lever.
3373 pub fn decode_step_t_h(
3374 &self,
3375 e: &Engine,
3376 tokens: &[u32],
3377 pos0: usize,
3378 cache: &mut Cache,
3379 ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
3380 self.decode_step_t_h_emb(e, tokens, pos0, cache, None)
3381 }
3382
3383 /// Like `decode_step_t_h` with an optional RESIDENT embed table (spec hot loop): device
3384 /// gather instead of host dequant + [T, n_embd] f32 htod. Bit-identical rows.
3385 pub fn decode_step_t_h_emb(
3386 &self,
3387 e: &Engine,
3388 tokens: &[u32],
3389 pos0: usize,
3390 cache: &mut Cache,
3391 embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
3392 ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
3393 let (logits_d, h_seed) = self.decode_step_t_h_emb_dev(e, tokens, pos0, cache, embd_dev)?;
3394 Ok((e.dtoh(&logits_d)?, h_seed))
3395 }
3396
3397 /// DEVICE-LOGITS verify forward (spec device-argmax lever): identical kernel chain to
3398 /// `decode_step_t_h_emb` but returns the [T, n_vocab] logits ON DEVICE — the accept walk
3399 /// argmaxes each column on-device and reads back ONE [T] u32 instead of dtoh'ing the full
3400 /// T x n_vocab f32 block (~1-4 MB + T host argmaxes, every round). Kernel dispatch is
3401 /// UNCHANGED (same decode-exact kernels); only the post-logits transfer moves.
3402 pub fn decode_step_t_h_emb_dev(
3403 &self,
3404 e: &Engine,
3405 tokens: &[u32],
3406 pos0: usize,
3407 cache: &mut Cache,
3408 embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
3409 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
3410 let n_embd = self.cfg.n_embd as usize;
3411 let t = tokens.len();
3412 let (logits, x) = self.decode_step_t_core(e, tokens, pos0, cache, embd_dev, None)?;
3413 // h_seed for the next round = LAST column's pre-output_norm hidden ([n_embd]).
3414 let mut hs = vbuf(e, n_embd)?; // fully written by copy_view_into below
3415 e.copy_view_into(&mut hs, 0, &x.slice((t - 1) * n_embd..t * n_embd), n_embd)?;
3416 Ok((logits, hs))
3417 }
3418
3419 /// CORE verify forward: the `decode_step_t_h_emb_dev` kernel chain, returning the FULL
3420 /// pre-output_norm hidden stack x ([T, n_embd], any column extractable) and optionally
3421 /// filling a `VerifyCkpt` (retained per-layer state-rebuild inputs) for the REPLAY-FREE
3422 /// partial accept. `ckpt: None` => byte-for-byte the old behavior (the ckpt writes are pure
3423 /// retains/copies — they never change what any kernel computes).
3424 fn decode_step_t_core(
3425 &self,
3426 e: &Engine,
3427 tokens: &[u32],
3428 pos0: usize,
3429 cache: &mut Cache,
3430 embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
3431 mut ckpt: Option<&mut VerifyCkpt>,
3432 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
3433 self.decode_step_t_core_stream(e, tokens, pos0, cache, embd_dev, ckpt.take(), None, None)
3434 }
3435
3436 /// Increment-0 two-session PP seam: release the peer after this lane's stage-0 boundary TX.
3437 /// The two independent sessions keep their own cache/checkpoint state; only issue order moves.
3438 fn decode_step_t_core_pipelined(
3439 &self,
3440 e: &Engine,
3441 tokens: &[u32],
3442 pos0: usize,
3443 cache: &mut Cache,
3444 embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
3445 mut ckpt: Option<&mut VerifyCkpt>,
3446 pipe: &SpecPipeLane,
3447 round: usize,
3448 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
3449 let fence = crate::pp::pp_cuts(self.layers.len())
3450 .ok_or("two-session speculative pipeline requires a PP stage cut")?;
3451 if crate::pp::pp2_streams_off() || !crate::pp::spec_pp_on() {
3452 return Err("two-session speculative pipeline requires the PP verify split".into());
3453 }
3454 let interval_fence = pipe.stage0_begin(round)?;
3455 let ticket = self.verify_stage0_issue(
3456 e,
3457 tokens,
3458 pos0,
3459 cache,
3460 embd_dev,
3461 ckpt.as_deref_mut(),
3462 None,
3463 &fence,
3464 Some(interval_fence),
3465 pipe.trace(round),
3466 )?;
3467 pipe.stage0_end(round);
3468 pipe.stage1_begin(round)?;
3469 let result = self.verify_stage1_finish(e, ticket, cache, ckpt, None, &fence, true)?;
3470 pipe.verify_end(round);
3471 Ok(result)
3472 }
3473
3474 /// ROUND-STREAM stage (c) 4: `stream` = (device verify tokens [t], device pos counter) —
3475 /// when Some, rope positions come from pos_iota over the counter, the embed gathers the
3476 /// device tokens, and full_attn_verify routes appends/FA through the _dc twins reading the
3477 /// SAME counter (every layer's kvl.len == cache.pos, one counter drives all three). The
3478 /// host `tokens`/`pos0` args still size buffers (t is FIXED K+1 in stream mode).
3479 #[allow(clippy::too_many_arguments)]
3480 fn decode_step_t_core_stream(
3481 &self,
3482 e: &Engine,
3483 tokens: &[u32],
3484 pos0: usize,
3485 cache: &mut Cache,
3486 embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
3487 mut ckpt: Option<&mut VerifyCkpt>,
3488 stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
3489 pp_pipe: Option<bool>,
3490 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
3491 // PP DOOR (lane/pp2-spec 2026-08-06): the verify trunk now takes its OWN stage split,
3492 // exactly as the eager and batched steps do. This is the single funnel every verify
3493 // forward reaches (decode_step_t / _h / _h_emb / _h_emb_dev / _core all land here), so
3494 // wiring it here wires the whole spec surface — the draft/accept/commit machinery above
3495 // is untouched.
3496 //
3497 // History: pp2-hardening (2026-08-06) made this funnel FAIL CLOSED, because its trunk
3498 // walk was unsplit on one stream and a sharded cross-device placement peer-read every
3499 // remote layer's weights on every spec round (measured 13.9-28x on the batched twin).
3500 // The refusal below survives to cover the residue — MEMRA_SPEC_PP=0, MEMRA_PP_STREAMS=0,
3501 // or a placement whose PpNRt fails to build — so a config that would still walk the
3502 // whole trunk on one stream refuses instead of regressing 28x.
3503 if let Some(fence) = crate::pp::pp_cuts(self.layers.len()) {
3504 if !crate::pp::pp2_streams_off() && crate::pp::spec_pp_on() {
3505 return self.decode_step_t_core_ppn(
3506 e,
3507 tokens,
3508 pos0,
3509 cache,
3510 embd_dev,
3511 ckpt.take(),
3512 stream,
3513 &fence,
3514 pp_pipe,
3515 );
3516 }
3517 }
3518 crate::pp::refuse_unsplit_if_remote(
3519 "decode_step_t (spec verify)",
3520 "drop MEMRA_SPEC_PP=0 / MEMRA_PP_STREAMS=0 so the verify trunk takes its OWN stage \
3521 split (decode_step_t_core_ppn); or run spec on one device",
3522 )?;
3523 let cfg = &self.cfg;
3524 let n_embd = cfg.n_embd as usize;
3525 let eps = cfg.rms_eps;
3526 let t = tokens.len();
3527 let pos_d = match stream {
3528 Some((_, ctr)) => {
3529 let mut p = e.alloc_uninit::<i32>(t)?;
3530 e.pos_iota(ctr, &mut p, t)?;
3531 p
3532 }
3533 None => {
3534 let pos_vec: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
3535 e.htod_i32(&pos_vec)?
3536 }
3537 };
3538
3539 // embed T tokens -> [T, n_embd] token-major (device gather on the spec hot loop)
3540 let x = match (stream, embd_dev) {
3541 (Some((vtok, _)), Some((g, qt, rb))) => {
3542 e.embed_gather_device_td(g, vtok, t, n_embd, qt, rb)?
3543 }
3544 (None, Some((g, qt, rb))) => e.embed_gather_device_t(g, tokens, n_embd, qt, rb)?,
3545 _ => e.htod(&self.embd.gather(n_embd, tokens))?,
3546 };
3547
3548 // TRUNK WALK: layers [0, n_layers) through the SAME range-scoped subgraph the PP-N
3549 // stage split calls per stage (`verify_layers`) — one code path, so the split cannot
3550 // drift from the unsplit dispatch mirroring. lane/pp2-spec 2026-08-06.
3551 let x = self.verify_layers(
3552 e,
3553 x,
3554 0,
3555 self.layers.len(),
3556 &pos_d,
3557 pos0,
3558 t,
3559 cache,
3560 ckpt.take(),
3561 stream,
3562 )?;
3563
3564 let mut hn = vbuf(e, t * n_embd)?;
3565 let serving_head = self.cfg.step35.is_some() || self.qwen35_serving_class();
3566 let logits = if serving_head {
3567 // Step35 and the qwen35 family (MoE 2026-08-14 AM, dense-hybrid same day PM — the
3568 // Q3.8 bring-up reproduced the identical near-tie class on dense: eager-class verify
3569 // vs batched-class live serving, ULP drift amplified through the GDN recurrence)
3570 // serve one batched numeric class at every live width, including B=1. Keep the
3571 // verify head in that same class; other generic families retain the decode-exact
3572 // head that their run-spec contract pins.
3573 e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
3574 e.matmul(&self.output, &hn, t)?
3575 } else {
3576 e.rms_norm_decode(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
3577 e.matmul_decode_exact(&self.output, &hn, t)?
3578 };
3579 // stream: the device pos counter owns position; host mirror reconciles at drain.
3580 if stream.is_none() {
3581 cache.pos += t;
3582 }
3583 // Hidden stack for seeds/refresh-fills: pre-norm x (default) or post-norm hn (HPOST).
3584 Ok((logits, if spec_hpost() { hn } else { x }))
3585 }
3586
3587 /// THE VERIFY TRUNK OVER PP-N (lane/pp2-spec 2026-08-06): `decode_step_t_core_stream`'s walk
3588 /// as N stage subgraphs, each on its own engine/stream (and, under `MEMRA_PP_DEVICES`, its own
3589 /// device), with a `[T, n_embd]` boundary transfer between them. T = K+1 (the verify batch),
3590 /// so this is the batched-boundary shape the pp2-batch lane's grow-only slots already handle
3591 /// (`tx(b, x, t*n_embd)`; the slot grows to the high-water T and the transport moves exactly
3592 /// the payload).
3593 ///
3594 /// Structure is `decode_step_batch_ppn`'s, which is `decode_step_h_ppn`'s. FOUR THINGS ARE
3595 /// PER-STAGE and each for a measured reason (see `decode_step_batch_ppn`'s header for the
3596 /// receipts):
3597 ///
3598 /// 1. THE ENGINE (`rt.engine(s, e)`) — `Engine` owns lazily-grown stable-pointer scratch
3599 /// (`fa_part_pool`, `fa_vf16_scratch`, `argmax_partials`) that is single-stream-safe BY
3600 /// DESIGN. Two stage streams through one Engine is the 2026-08-02 shared-scratch race
3601 /// (35% flake, nondeterministic all-logits divergence). `PpNRt::build` gives every stage
3602 /// s>0 its own Engine even on the primary device; honouring it here is what scopes the
3603 /// pools. The verify path allocates MORE of that scratch than eager decode does (FA at
3604 /// m=T, and the per-layer `GdnStash` retains), so this is load-bearing, not inherited.
3605 ///
3606 /// 2. `pos_d` — each stage uploads its OWN copy of the T rope positions on ITS stream, so the
3607 /// buffer is allocated, consumed and freed on one stream. In `stream` mode that means each
3608 /// stage runs its own `pos_iota` over the SHARED device counter (`pos_ctr`): the counter is
3609 /// read-only during the forward (the round's `inc`/`copy_add` happen outside it), so every
3610 /// stage derives the identical iota, and each stage's own output buffer is stream-local.
3611 ///
3612 /// 3. THE EMBED lives with stage 0 (`self.embd` / `embd_gpu` are host/primary-side; the
3613 /// sharded loader leaves the table with stage 0 by construction).
3614 ///
3615 /// 4. THE HEAD (`output_norm` + `output`) runs on the LAST stage — the sharded loader uploaded
3616 /// both through that stage's engine (`hybrid.rs`: `e_head = layer_engine(e, n_trunk,
3617 /// n_trunk-1)`), so reading them anywhere else is a peer read of the biggest tensor in the
3618 /// model, every round.
3619 ///
3620 /// WHAT STAYS ON THE PRIMARY, deliberately: the returned logits and hidden stack `x`. Both are
3621 /// last-stage-allocated device buffers, and every consumer (the device argmax walk, the accept
3622 /// kernels, `spec_seed_gather`, the ckpt rebuild in `commit_verified_prefix`) reads them
3623 /// through the primary context by UVA — the same read the batched serving epilogue's
3624 /// `last_logits_dev` park does. Those consumers are per-round O(T x n_vocab) and O(n_embd),
3625 /// not per-layer, so they are not the 28x class; splitting them is a separate lane.
3626 ///
3627 /// The MTP HEAD (draft side) is NOT split: it is one block, it lives wherever the loader put
3628 /// it (`load_mtp` uses the primary engine), and it is ~1-2 GB against the trunk's tens. Draft
3629 /// placement is measured, not assumed — see `research/pp2-spec-20260806`.
3630 ///
3631 /// EXACTNESS: PP-N adds ZERO deviation. Each stage runs the SAME kernels on the SAME bytes in
3632 /// the same order via the SAME `verify_layers` the unsplit body calls; the only change is
3633 /// where the residual is materialized, and the boundary is a straight f32 copy (dtod
3634 /// same-device / `cudaMemcpyPeerAsync` cross-device, no conversion). So the split MUST be
3635 /// BIT-IDENTICAL to the unsplit verify at the same T, in both placement orders. Gate:
3636 /// `decode-batch-gate --mode ppspec`. Acceptance counts are a DERIVED consequence — greedy
3637 /// accept argmaxes these logits, so bit-identical logits force identical accept walks; the
3638 /// `run-spec` K=1..8 arm checks that end-to-end rather than trusting the implication.
3639 #[allow(clippy::too_many_arguments)]
3640 fn decode_step_t_core_ppn(
3641 &self,
3642 e: &Engine,
3643 tokens: &[u32],
3644 pos0: usize,
3645 cache: &mut Cache,
3646 embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
3647 mut ckpt: Option<&mut VerifyCkpt>,
3648 stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
3649 fence: &[usize],
3650 pp_pipe: Option<bool>,
3651 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
3652 let ticket = self.verify_stage0_issue(
3653 e,
3654 tokens,
3655 pos0,
3656 cache,
3657 embd_dev,
3658 ckpt.as_deref_mut(),
3659 stream,
3660 fence,
3661 pp_pipe,
3662 None,
3663 )?;
3664 self.verify_stage1_finish(e, ticket, cache, ckpt, stream, fence, true)
3665 }
3666
3667 /// Enqueue embed, stage 0, and the first boundary TX, then return the actual boundary slot.
3668 /// The ordinary PP verify wrapper calls `verify_stage1_finish` immediately after this return.
3669 #[allow(clippy::too_many_arguments)]
3670 fn verify_stage0_issue(
3671 &self,
3672 e: &Engine,
3673 tokens: &[u32],
3674 pos0: usize,
3675 cache: &mut Cache,
3676 embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
3677 mut ckpt: Option<&mut VerifyCkpt>,
3678 stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
3679 fence: &[usize],
3680 pp_pipe: Option<bool>,
3681 trace: Option<SpecPipeTraceCtx>,
3682 ) -> Result<VerifyBoundaryTicket, Box<dyn std::error::Error>> {
3683 assert!(
3684 !self.is_gemma4_e4b() && self.cfg.gemma4.is_none(),
3685 "decode_step_t_core_ppn covers the hybrid non-gemma4 verify trunk only \
3686 (the gemma4 arms have their own decode_step_t twins)"
3687 );
3688 if crate::pp::pp_host_bounce_active() && (stream.is_some() || embd_dev.is_some()) {
3689 return Err(
3690 "decode_step_t_core_ppn: refused with MEMRA_PP_HOST_BOUNCE=1 — the trunk \
3691 boundary itself is host-staged, but device-resident verify still peer-reads \
3692 primary-device token/position/embedding buffers from stage 0. Run plain PP \
3693 serving on this host class; spec requires local per-stage inputs first."
3694 .into(),
3695 );
3696 }
3697 let rt = crate::pp::PpNRt::get(e)?;
3698 let n_st = fence.len() - 1;
3699 assert_eq!(
3700 rt.n_stages(),
3701 n_st,
3702 "PpNRt stage count {} != fence stages {n_st}",
3703 rt.n_stages()
3704 );
3705 let n_embd = self.cfg.n_embd as usize;
3706 let t = tokens.len();
3707 let payload = t * n_embd;
3708 if pp_pipe.is_some() {
3709 assert_eq!(n_st, 2, "spec pipeline requires exactly two PP stages");
3710 }
3711 // One-shot lane diagnostic: force natural PP-2 boundaries to completion so the server
3712 // log can price stage 0, the peer hop, the RX copy, and stage 1 + head separately without
3713 // nsys. The ordinary path keeps every enqueue asynchronous. N>2 is deliberately excluded:
3714 // the report below names exactly two stages and must never imply it measured middle ones.
3715 let pp_anatomy = n_st == 2 && std::env::var("MEMRA_SPEC_PP_ANATOMY").as_deref() == Ok("1");
3716 let pp_started = std::time::Instant::now();
3717 let (mut reverse_ms, mut stage0_ms, mut tx_ms) = (0.0f64, 0.0f64, 0.0f64);
3718 // The CALLER's ambient stream, captured BEFORE any `rt.enter()` pushes a stage stream:
3719 // this body returns DEVICE-RESIDENT buffers (the device-argmax accept walk's contract),
3720 // so the exit needs the same publication the boundaries get — see `PpNRt::publish_to`.
3721 // Taken here, not at the end, because inside the last-stage scope `e.stream()` IS the
3722 // stage stream and the wait would self-order into a no-op.
3723 let caller_stream = e.stream();
3724 // #87 ROOT-CAUSE FENCE (lane/pp2spec-crash): the PREVIOUS round's stage-allocated
3725 // outputs (logits/hidden/ckpt stashes) freed stream-ordered on the STAGE streams while
3726 // the primary stream still holds queued reads of them — with event tracking elided,
3727 // nothing stops the pool from reusing those blocks for THIS round's stage allocations,
3728 // whose writes then race the queued reads (measured: 13/4096-NaN random-bits garbage in
3729 // the spec round seed; the full anatomy is on `PpNRt::fence_stages_behind`). Order every
3730 // stage stream behind the caller before enqueueing new stage work.
3731 let reverse_started = std::time::Instant::now();
3732 if pp_pipe != Some(false) {
3733 rt.fence_stages_behind(&caller_stream)?;
3734 }
3735 if pp_pipe == Some(true) {
3736 // Both session verifies must alternate boundary slots even when the ordinary
3737 // decode overlap experiment is off. Prewarm before A's stage 0 so B cannot grow
3738 // slot 1 by synchronizing the RX stream while A's stage 1 is in flight.
3739 rt.prepare_overlap_slots(0, payload)?;
3740 }
3741 if pp_anatomy {
3742 // Drain the reverse-publication dependency before timing stage 0 itself. At c=1 this
3743 // prices any primary-stream rollback/refresh tail inherited from the prior round.
3744 for s in 0..n_st {
3745 let _st = rt.enter(s);
3746 rt.engine(s, e).stream().synchronize()?;
3747 }
3748 reverse_ms = reverse_started.elapsed().as_secs_f64() * 1e3;
3749 }
3750
3751 // Per-stage rope positions: in host mode the same [T] iota each stage uploads itself; in
3752 // stream mode each stage's own `pos_iota` over the shared read-only device counter.
3753 let stage_pos = |es: &Engine| -> Result<CudaSlice<i32>, Box<dyn std::error::Error>> {
3754 match stream {
3755 Some((_, ctr)) => {
3756 let mut p = es.alloc_uninit::<i32>(t)?;
3757 es.pos_iota(ctr, &mut p, t)?;
3758 Ok(p)
3759 }
3760 None => {
3761 let pos_vec: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
3762 es.htod_i32(&pos_vec)
3763 }
3764 }
3765 };
3766
3767 // ---- STAGE 0: embed (the table lives with stage 0) + layers [0, fence[1]) + TX ----
3768 let slot = {
3769 let _st0 = rt.enter(0);
3770 let e0 = rt.engine(0, e);
3771 enqueue_spec_pipe_trace_marker(&e0.stream(), trace.as_ref(), "S0", "start", None)?;
3772 let stage0_started = std::time::Instant::now();
3773 let pos_d = stage_pos(e0)?;
3774 let x = match (stream, embd_dev) {
3775 (Some((vtok, _)), Some((g, qt, rb))) => {
3776 e0.embed_gather_device_td(g, vtok, t, n_embd, qt, rb)?
3777 }
3778 (None, Some((g, qt, rb))) => e0.embed_gather_device_t(g, tokens, n_embd, qt, rb)?,
3779 _ => e0.htod(&self.embd.gather(n_embd, tokens))?,
3780 };
3781 let x = self.verify_layers(
3782 e0,
3783 x,
3784 fence[0],
3785 fence[1],
3786 &pos_d,
3787 pos0,
3788 t,
3789 cache,
3790 ckpt.as_deref_mut(),
3791 stream,
3792 )?;
3793 if pp_anatomy {
3794 e0.stream().synchronize()?;
3795 stage0_ms = stage0_started.elapsed().as_secs_f64() * 1e3;
3796 }
3797 let tx_started = std::time::Instant::now();
3798 let slot = if pp_pipe.is_some() {
3799 rt.tx_pipelined(0, &x, payload)?
3800 } else {
3801 rt.tx(0, &x, payload)?
3802 };
3803 enqueue_spec_pipe_trace_marker(&e0.stream(), trace.as_ref(), "S0", "end", Some(slot))?;
3804 if pp_anatomy {
3805 e0.stream().synchronize()?;
3806 tx_ms = tx_started.elapsed().as_secs_f64() * 1e3;
3807 }
3808 slot
3809 // x + pos_d drop here: freed stream-ordered on stage-0's stream after use.
3810 };
3811
3812 Ok(VerifyBoundaryTicket {
3813 rt,
3814 caller_stream,
3815 slot,
3816 pos0,
3817 t,
3818 payload,
3819 n_st,
3820 pipelined: pp_pipe.is_some(),
3821 pp_anatomy,
3822 pp_started,
3823 reverse_ms,
3824 stage0_ms,
3825 tx_ms,
3826 trace,
3827 })
3828 }
3829
3830 /// Consume a stage-0 boundary ticket and enqueue the remaining PP stages plus the head.
3831 /// On PP-2 this is exactly stage 1; PP-N keeps its pre-existing middle-stage walk here.
3832 #[allow(clippy::too_many_arguments)]
3833 fn verify_stage1_finish(
3834 &self,
3835 e: &Engine,
3836 ticket: VerifyBoundaryTicket,
3837 cache: &mut Cache,
3838 mut ckpt: Option<&mut VerifyCkpt>,
3839 stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
3840 fence: &[usize],
3841 publish_to_caller: bool,
3842 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
3843 let VerifyBoundaryTicket {
3844 rt,
3845 caller_stream,
3846 slot,
3847 pos0,
3848 t,
3849 payload,
3850 n_st,
3851 pipelined,
3852 pp_anatomy,
3853 pp_started,
3854 reverse_ms,
3855 stage0_ms,
3856 tx_ms,
3857 trace,
3858 } = ticket;
3859 let n_embd = self.cfg.n_embd as usize;
3860 let eps = self.cfg.rms_eps;
3861 let mut slot = slot;
3862 let (mut rx_ms, mut stage1_ms) = (0.0f64, 0.0f64);
3863 let stage_pos = |es: &Engine| -> Result<CudaSlice<i32>, Box<dyn std::error::Error>> {
3864 match stream {
3865 Some((_, ctr)) => {
3866 let mut p = es.alloc_uninit::<i32>(t)?;
3867 es.pos_iota(ctr, &mut p, t)?;
3868 Ok(p)
3869 }
3870 None => {
3871 let pos_vec: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
3872 es.htod_i32(&pos_vec)
3873 }
3874 }
3875 };
3876
3877 // ---- MIDDLE STAGES: RX boundary s-1 -> range -> TX boundary s ----
3878 for s in 1..n_st - 1 {
3879 let _st = rt.enter(s);
3880 let es = rt.engine(s, e);
3881 let pos_d = stage_pos(es)?;
3882 let x = rt.rx(s - 1, slot, payload)?;
3883 let x = self.verify_layers(
3884 es,
3885 x,
3886 fence[s],
3887 fence[s + 1],
3888 &pos_d,
3889 pos0,
3890 t,
3891 cache,
3892 ckpt.as_deref_mut(),
3893 stream,
3894 )?;
3895 slot = if pipelined {
3896 rt.tx_pipelined(s, &x, payload)?
3897 } else {
3898 rt.tx(s, &x, payload)?
3899 };
3900 }
3901
3902 // ---- LAST STAGE: RX + final range + output_norm + lm head ----
3903 let _stl = rt.enter(n_st - 1);
3904 let el = rt.engine(n_st - 1, e);
3905 let pos_d = stage_pos(el)?;
3906 let rx_started = std::time::Instant::now();
3907 let x = rt.rx(n_st - 2, slot, payload)?;
3908 if pp_anatomy {
3909 el.stream().synchronize()?;
3910 rx_ms = rx_started.elapsed().as_secs_f64() * 1e3;
3911 }
3912 enqueue_spec_pipe_trace_marker(&el.stream(), trace.as_ref(), "S1", "start", Some(slot))?;
3913 let stage1_started = std::time::Instant::now();
3914 let x = self.verify_layers(
3915 el,
3916 x,
3917 fence[n_st - 1],
3918 fence[n_st],
3919 &pos_d,
3920 pos0,
3921 t,
3922 cache,
3923 ckpt.as_deref_mut(),
3924 stream,
3925 )?;
3926
3927 let mut hn = vbuf(el, payload)?;
3928 let logits = if self.cfg.step35.is_some() {
3929 // The PP Step35 serving path uses rms_norm + matmul for B=1 as well as B>1.
3930 // Verify must not switch numeric class merely because the same session speculates.
3931 el.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
3932 el.matmul(&self.output, &hn, t)?
3933 } else {
3934 el.rms_norm_decode(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
3935 el.matmul_decode_exact(&self.output, &hn, t)?
3936 };
3937 enqueue_spec_pipe_trace_marker(&el.stream(), trace.as_ref(), "S1", "end", Some(slot))?;
3938 if pp_anatomy {
3939 el.stream().synchronize()?;
3940 stage1_ms = stage1_started.elapsed().as_secs_f64() * 1e3;
3941 }
3942 // EXIT PUBLICATION: both returned buffers are still being produced on the last stage's
3943 // stream. Order the caller's stream behind that work before the buffers escape this
3944 // scope (the 2026-08-06 same-device ppspec find: without it the caller's primary-stream
3945 // consumer read unwritten logits — nondeterministic, one-device-only, and it poisoned
3946 // the following arm's KV in the same process).
3947 if publish_to_caller {
3948 rt.publish_to(n_st - 1, &caller_stream)?;
3949 }
3950 if pp_anatomy {
3951 if publish_to_caller {
3952 caller_stream.synchronize()?;
3953 }
3954 eprintln!(
3955 "[spec-pp-anatomy] t={t} reverse={reverse_ms:.3}ms stage0={stage0_ms:.3}ms \
3956 tx={tx_ms:.3}ms rx={rx_ms:.3}ms stage1-head={stage1_ms:.3}ms total={:.3}ms",
3957 pp_started.elapsed().as_secs_f64() * 1e3,
3958 );
3959 }
3960 // stream: the device pos counter owns position; host mirror reconciles at drain.
3961 if stream.is_none() {
3962 cache.pos += t;
3963 }
3964 Ok((logits, if spec_hpost() { hn } else { x }))
3965 }
3966
3967 /// Step3.5/Step3.7 verify trunk in the serving batched numeric class.
3968 ///
3969 /// `step35_decode_batch_layers` is now authoritative at every live serving width, including
3970 /// B=1 (lane/cx-b1fix). The older verify walk deliberately mirrored the eager T=1 class:
3971 /// it replayed `step35_decode_attn` per row and used the eager/decode-exact FFN dispatch.
3972 /// Those classes are individually stable, but a near-tie prompt can choose different greedy
3973 /// bytes when a request moves from batched plain serving into speculative verify. Run the
3974 /// same authoritative B=1 stage subgraph for each verify row here. Rows still advance
3975 /// layer-by-layer, so every layer sees the preceding verify rows in its attention cache while
3976 /// every norm/projection/FFN uses exactly the live serving dispatch.
3977 #[allow(clippy::too_many_arguments)]
3978 fn step35_verify_batch_layers(
3979 &self,
3980 e: &Engine,
3981 mut x: CudaSlice<f32>,
3982 lo: usize,
3983 hi: usize,
3984 pos0: usize,
3985 t: usize,
3986 cache: &mut Cache,
3987 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3988 let n_embd = self.cfg.n_embd as usize;
3989 self.cfg
3990 .step35
3991 .as_ref()
3992 .ok_or("step35 verify batch requires step35 cfg")?;
3993 let mut ph_last = std::time::Instant::now();
3994 for il in lo..hi {
3995 let mut next = e.uninit(t * n_embd)?;
3996 for r in 0..t {
3997 let mut row = e.uninit(n_embd)?;
3998 e.dtod_copy_view(&x.slice(r * n_embd..(r + 1) * n_embd), &mut row)?;
3999 // The caller owns this verify's position. During controller overlap, cache.pos
4000 // still describes generation N while this stage-0 walk belongs to N+1.
4001 let row_pos = e.htod_i32(&[(pos0 + r) as i32])?;
4002 let mut one = [&mut *cache];
4003 let out = self.step35_decode_batch_layers(
4004 e,
4005 row,
4006 &mut one,
4007 &row_pos,
4008 il,
4009 il + 1,
4010 &mut ph_last,
4011 )?;
4012 e.dtod_copy_into(&out, &mut next, r * n_embd)?;
4013 }
4014 self.dflash_tap(e, cache, il, &next, t)?;
4015 x = next;
4016 }
4017 Ok(x)
4018 }
4019
4020 /// DSpark drafter verify (lane/dspark-q38-recover): one t-row forward through the
4021 /// SERVING-CLASS verify funnel (`decode_step_t_core_stream` — the same numeric class
4022 /// MTP verify rides, GDN state advanced in place), returning per-row argmax tokens.
4023 /// Advances `cache.pos += t`; the caller owns snapshot/rollback (block acceptance is
4024 /// prefix-keep, not all-or-nothing).
4025 pub(crate) fn dspark_verify_t_am(
4026 &self,
4027 e: &Engine,
4028 tokens: &[u32],
4029 pos0: usize,
4030 cache: &mut Cache,
4031 ) -> Result<Vec<u32>, Box<dyn std::error::Error>> {
4032 let (logits, _hn) =
4033 self.decode_step_t_core_stream(e, tokens, pos0, cache, None, None, None, None)?;
4034 let t = tokens.len();
4035 let v = self.output.out_features();
4036 let mut am_d = e.stream().alloc_zeros::<u32>(t)?;
4037 for r in 0..t {
4038 e.argmax_token_device_col(&logits, r, v, &mut am_d, r)?;
4039 }
4040 Ok(e.dtoh_u32(&am_d)?)
4041 }
4042
4043 /// DSpark verify with the MTP column-stash armed: identical forward to
4044 /// `dspark_verify_t_am`, but fills a `VerifyCkpt` so a partial accept can restore
4045 /// column state directly (`dspark_commit_prefix`) instead of snapshot-replay.
4046 /// The ckpt type is opaque outside spec.rs (newtype) — dflash.rs threads it through.
4047 pub(crate) fn dspark_verify_t_am_ckpt(
4048 &self,
4049 e: &Engine,
4050 tokens: &[u32],
4051 pos0: usize,
4052 cache: &mut Cache,
4053 ) -> Result<(Vec<u32>, DsparkVerifyCkpt), Box<dyn std::error::Error>> {
4054 let mut ck = VerifyCkpt::new(self.layers.len());
4055 let (logits, _hn) = self.decode_step_t_core_stream(
4056 e,
4057 tokens,
4058 pos0,
4059 cache,
4060 None,
4061 Some(&mut ck),
4062 None,
4063 None,
4064 )?;
4065 let t = tokens.len();
4066 let v = self.output.out_features();
4067 let mut am_d = e.stream().alloc_zeros::<u32>(t)?;
4068 for r in 0..t {
4069 e.argmax_token_device_col(&logits, r, v, &mut am_d, r)?;
4070 }
4071 Ok((e.dtoh_u32(&am_d)?, DsparkVerifyCkpt(ck)))
4072 }
4073
4074 /// Restore the round to `keep` accepted columns from the verify stash: KV lens and
4075 /// pos from the pre-verify snapshot + keep, GDN conv/ssm from the stashed column
4076 /// state — no replay forward. The exact `commit_verified_prefix` the MTP path ships.
4077 pub(crate) fn dspark_commit_prefix(
4078 &self,
4079 e: &Engine,
4080 cache: &mut Cache,
4081 snap: &crate::cache::CacheSnapshot,
4082 ckpt: &DsparkVerifyCkpt,
4083 keep: usize,
4084 ) -> Result<(), Box<dyn std::error::Error>> {
4085 self.commit_verified_prefix(e, cache, snap, &ckpt.0, keep, false, None)
4086 }
4087
4088 /// Qwen35-family verify trunk in the live serving numeric class.
4089 ///
4090 /// Serving intentionally keeps this architecture in the generic batched program even at
4091 /// B=1. The older verify walk used its own mirrored dispatch and can flip near-tie argmaxes.
4092 ///
4093 /// Two arms, one numeric class:
4094 /// - DENSE (`Arch::Qwen35`, t<=16): `qwen35_verify_tparallel` — the weight ops (norms,
4095 /// projections, FFN) hoist to m=T through the exact-tier batched kernels whose per-row
4096 /// program IS the m=1 program (`matmul_pre == fused2 per (tensor,row); _bN mmvq per-row
4097 /// == m=1` — decode_batch.rs v2 note), while the state ops (conv ring, gdn scan, KV
4098 /// append, fa decode) stay a per-row loop running the b_n=1 serving kernels with each
4099 /// row's own t_kv-driven arm pick (the straddle law: every row executes the exact
4100 /// program its isolated serving step would). One weight read per layer per round
4101 /// instead of T — this is what makes MTP profitable in the exact class (the per-row
4102 /// walk measured verify(K+1) ~= (K+1) plain steps: 69 -> 44 tok/s served, 2026-08-15).
4103 /// - MoE / t>16 / `MEMRA_SPEC_VERIFY_ROWWISE=1`: the per-row replay of the authoritative
4104 /// serving layer body, preserving single-session autoregressive cache order (the
4105 /// correctness reference; also the rollback seam for the t-parallel arm).
4106 ///
4107 /// Bit-identity of the t-parallel arm vs the rowwise arm is gated by spec-serve-gate
4108 /// (zero differing logits at T=1..4, K arms) + the 8-prompt ON/OFF canary before ship.
4109 #[allow(clippy::too_many_arguments)]
4110 fn qwen35_verify_batch_layers(
4111 &self,
4112 e: &Engine,
4113 x: CudaSlice<f32>,
4114 lo: usize,
4115 hi: usize,
4116 pos0: usize,
4117 t: usize,
4118 cache: &mut Cache,
4119 ckpt: Option<&mut VerifyCkpt>,
4120 stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
4121 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4122 // Qwen35Moe admitted 2026-08-20 (lane/draftcost-moe): the t-parallel arm already
4123 // carries the MoE FFN (`moe_ffn_il_zq8` at m=T) and the GDN per-row state loop; the
4124 // arch fence was a qualification gate, not a mechanism gap. Measured disease on the
4125 // 35B-A3B class: rowwise verify ~= 5.6 ms per drafted token (one full trunk step
4126 // each) — the same (K+1)-plain-steps wall the dense admission fixed on 2026-08-15.
4127 // Rollback seam unchanged: MEMRA_SPEC_VERIFY_ROWWISE=1.
4128 let rowwise = std::env::var("MEMRA_SPEC_VERIFY_ROWWISE").as_deref() == Ok("1")
4129 || !matches!(
4130 self.cfg.arch,
4131 memra_gguf::config::Arch::Qwen35 | memra_gguf::config::Arch::Qwen35Moe
4132 )
4133 || t > 16;
4134 if rowwise {
4135 if stream.is_some() {
4136 // rowwise replays per row with host cache.pos — irreconcilable with a
4137 // device position counter. Burst callers must keep t <= 16 and the
4138 // ROWWISE env unset; refusing beats silently mispositioned rows.
4139 return Err("qwen35 rowwise verify has no ROUND-STREAM arm \
4140 (t > 16 or MEMRA_SPEC_VERIFY_ROWWISE=1)"
4141 .into());
4142 }
4143 self.qwen35_verify_rowwise(e, x, lo, hi, pos0, t, cache, ckpt)
4144 } else {
4145 self.qwen35_verify_tparallel(e, x, lo, hi, pos0, t, cache, ckpt, stream)
4146 }
4147 }
4148
4149 /// The per-row correctness reference: replay each verify row through the authoritative
4150 /// serving layer body (`decode_batch_layers` at b_n=1). T full weight reads per layer.
4151 #[allow(clippy::too_many_arguments)]
4152 fn qwen35_verify_rowwise(
4153 &self,
4154 e: &Engine,
4155 mut x: CudaSlice<f32>,
4156 lo: usize,
4157 hi: usize,
4158 pos0: usize,
4159 t: usize,
4160 cache: &mut Cache,
4161 mut ckpt: Option<&mut VerifyCkpt>,
4162 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4163 let n_embd = self.cfg.n_embd as usize;
4164 let saved_pos = cache.pos;
4165 let mut ph_last = std::time::Instant::now();
4166 for il in lo..hi {
4167 let mut next = e.uninit(t * n_embd)?;
4168 let mut col_states: Option<Vec<(CudaSlice<f32>, CudaSlice<f32>)>> =
4169 if ckpt.is_some() && t >= 2 && matches!(self.layers[il].mixer, Mixer::Linear(_)) {
4170 Some(Vec::with_capacity(t - 1))
4171 } else {
4172 None
4173 };
4174 for r in 0..t {
4175 cache.pos = pos0 + r;
4176 let mut row = e.uninit(n_embd)?;
4177 e.dtod_copy_view(&x.slice(r * n_embd..(r + 1) * n_embd), &mut row)?;
4178 let row_pos = e.htod_i32(&[(pos0 + r) as i32])?;
4179 let mut one = [&mut *cache];
4180 let ctx = self.batch_layer_ctx(e, &one, il, il + 1)?;
4181 let out = match self.decode_batch_layers(
4182 e,
4183 row,
4184 &mut one,
4185 &ctx,
4186 &row_pos,
4187 &mut ph_last,
4188 ) {
4189 Ok(out) => out,
4190 Err(error) => {
4191 cache.pos = saved_pos;
4192 return Err(error);
4193 }
4194 };
4195 e.dtod_copy_into(&out, &mut next, r * n_embd)?;
4196 if r + 1 < t {
4197 if let Some(states) = col_states.as_mut() {
4198 let recur = cache.recur[il]
4199 .as_ref()
4200 .ok_or("Qwen35-MoE linear verify layer has no recurrent state")?;
4201 states.push((
4202 e.clone_dtod(&recur.conv_state)?,
4203 e.clone_dtod(&recur.ssm_state)?,
4204 ));
4205 }
4206 }
4207 }
4208 if let (Some(checkpoint), Some(states)) = (ckpt.as_deref_mut(), col_states) {
4209 checkpoint.cols[il] = Some(states);
4210 }
4211 x = next;
4212 }
4213 cache.pos = saved_pos;
4214 Ok(x)
4215 }
4216
4217 /// T-PARALLEL VERIFY IN THE SERVING NUMERIC CLASS (lane/tparallel-verify, 2026-08-15).
4218 ///
4219 /// The weight ops run ONCE per layer at m=T; the state ops run per row through the same
4220 /// b_n=1 serving kernels the rowwise replay uses. Per-row bit-identity rests on the two
4221 /// pins the serving batch tier already carries:
4222 /// * `matmul_pre` / `_bN` mmvq: per-row program == m=1 program (decode_batch.rs v2 note,
4223 /// kernel-check pinned) — so a [T, n_embd] projection row equals the row projected
4224 /// alone;
4225 /// * row-indexed norms/elementwise (`rms_norm`, `quantize_q8_1`, `add_rms_norm`,
4226 /// `gated_rmsnorm[_q8_1]`, `silu_mul`, `rope_neox` with per-row positions): the T-row
4227 /// launch is the per-row program (same pin the generic verify's fused norms rely on).
4228 /// The sequential dependencies keep their exact serving order: the conv ring / gdn scan
4229 /// chain state row -> row through the `_b` kernels at b_n=1 (ping-pong via a 6-entry
4230 /// alternating pointer table, host handles swapped per row so VerifyCkpt clones the
4231 /// canonical state exactly as the rowwise arm does), and each row's KV append + fa decode
4232 /// picks its arm from ITS OWN t_kv (append: format-only; fa: `fa_seqs_eligible` + its own
4233 /// `fa_split_keys` rung at b_n=1) — the straddle law per row, so every row executes the
4234 /// program its isolated B=1 serving step would.
4235 ///
4236 /// Cost: 1 weight read per layer per round + T state micro-launches, vs the rowwise arm's
4237 /// T weight reads. Gated bit-identical vs the rowwise arm by spec-serve-gate + canary.
4238 #[allow(clippy::too_many_arguments)]
4239 fn qwen35_verify_tparallel(
4240 &self,
4241 e: &Engine,
4242 mut x: CudaSlice<f32>,
4243 lo: usize,
4244 hi: usize,
4245 pos0: usize,
4246 t: usize,
4247 cache: &mut Cache,
4248 mut ckpt: Option<&mut VerifyCkpt>,
4249 stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
4250 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4251 use cudarc::driver::DevicePtr;
4252 let cfg = &self.cfg;
4253 let n_embd = cfg.n_embd as usize;
4254 let eps = cfg.rms_eps;
4255 let head_dim_global = cfg.head_dim_k as usize;
4256 // STREAM (2b, lane/draftcost-moe): positions come from the device round counter
4257 // (pos_iota / i32_copy_add) so a burst round needs no host position knowledge.
4258 let pos_d = match stream {
4259 Some((_, ctr)) => {
4260 let mut p = e.alloc_uninit::<i32>(t)?;
4261 e.pos_iota(ctr, &mut p, t)?;
4262 p
4263 }
4264 None => {
4265 let pos_host: Vec<i32> = (0..t).map(|r| (pos0 + r) as i32).collect();
4266 e.htod_i32(&pos_host)?
4267 }
4268 };
4269 // Per-row 1-element position buffers, built ONCE per verify (the append/fa wrappers
4270 // take owned pos slices; building these inside the layer x row loops cost 16xT H2Ds).
4271 let pos_rows: Vec<CudaSlice<i32>> = match stream {
4272 Some((_, ctr)) => (0..t)
4273 .map(|r| {
4274 let mut b = e.alloc_uninit::<i32>(1)?;
4275 e.i32_copy_add(ctr, &mut b, r as i32)?;
4276 Ok(b)
4277 })
4278 .collect::<Result<_, Box<dyn std::error::Error>>>()?,
4279 None => (0..t)
4280 .map(|r| e.htod_i32(&[(pos0 + r) as i32]))
4281 .collect::<Result<_, _>>()?,
4282 };
4283 let seqs_append =
4284 std::env::var("MEMRA_BATCH_APPEND").as_deref() != Ok("0") && !Engine::kv_fp8_on();
4285 let batch_fa_on = std::env::var("MEMRA_BATCH_FA").as_deref() != Ok("0");
4286
4287 for il in lo..hi {
4288 let layer = &self.layers[il];
4289 // ---- attn_norm + q8_1 quantize at m=T (row-indexed == per-row) ----
4290 let anorm = layer.attn_norm.float_data();
4291 let mut xn = e.uninit(t * n_embd)?;
4292 e.rms_norm(&x, anorm, &mut xn, n_embd, t, eps)?;
4293 let (hq, hd) = e.quantize_q8_1(&xn, t, n_embd)?;
4294
4295 let mixed: CudaSlice<f32> = match &layer.mixer {
4296 Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
4297 Mixer::Full(fa) => {
4298 let geometry = cfg.full_attention_geometry_at(il as u32);
4299 let n_head = geometry.n_head as usize;
4300 let n_head_kv = geometry.n_head_kv as usize;
4301 let head_dim = geometry.head_dim_k as usize;
4302 let rope_dims = geometry.n_rot as usize;
4303 let rope_base = geometry.rope_base;
4304 let scale = geometry.attention_scale();
4305 // Batched projections: one weight read serves all T rows.
4306 let qf = e.matmul_pre(&fa.wq, &hq, &hd, &xn, t)?;
4307 let mut k = e.matmul_pre(&fa.wk, &hq, &hd, &xn, t)?;
4308 let v = e.matmul_pre(&fa.wv, &hq, &hd, &xn, t)?;
4309 let gated =
4310 geometry.attention_gate == memra_gguf::config::AttentionGateKind::FusedQ;
4311 let (mut q, gate) = if gated {
4312 let mut qs = e.uninit(t * n_head * head_dim)?;
4313 let mut gs = e.uninit(t * n_head * head_dim)?;
4314 e.q_gate_split(&qf, &mut qs, &mut gs, head_dim, n_head, t)?;
4315 (qs, Some(gs))
4316 } else {
4317 (qf, None)
4318 };
4319 let mut qn = e.uninit(t * n_head * head_dim)?;
4320 e.rms_norm(
4321 &q,
4322 fa.q_norm.float_data(),
4323 &mut qn,
4324 head_dim,
4325 t * n_head,
4326 eps,
4327 )?;
4328 q = qn;
4329 let mut kn = e.uninit(t * n_head_kv * head_dim)?;
4330 e.rms_norm(
4331 &k,
4332 fa.k_norm.float_data(),
4333 &mut kn,
4334 head_dim,
4335 t * n_head_kv,
4336 eps,
4337 )?;
4338 k = kn;
4339 e.rope_neox(
4340 &mut q, &pos_d, head_dim, rope_dims, n_head, t, rope_base, 1.0,
4341 )?;
4342 e.rope_neox(
4343 &mut k, &pos_d, head_dim, rope_dims, n_head_kv, t, rope_base, 1.0,
4344 )?;
4345
4346 // Per-row append + attend: row r sees rows 0..r in KV (causal within the
4347 // draft), each through the b_n=1 serving kernels at its own t_kv.
4348 let q_dim = n_head * head_dim;
4349 let kv_dim = n_head_kv * head_dim;
4350 let mut attn = e.uninit(t * q_dim)?;
4351 let (kdk, kdv, ktb, vtb, kv_view) = {
4352 let kvl = cache.kv[il].as_ref().unwrap();
4353 let s = &e.gpu.stream();
4354 let (pk, _g) = kvl.k.device_ptr(s);
4355 let (pv, _g2) = kvl.v.device_ptr(s);
4356 (
4357 kvl.kv_dim_k,
4358 kvl.kv_dim_v,
4359 kvl.k_tok_bytes,
4360 kvl.v_tok_bytes,
4361 e.htod_u64(&[pk as u64, pv as u64])?,
4362 )
4363 };
4364 if let Some((_, ctr)) = stream {
4365 // STREAM ARM (2b): one batched dc append + the multi-row dc attention
4366 // — the generic stream arm's exact shape (rows kernels are pinned
4367 // byte-identical to the per-row programs by kernel-check). Host len
4368 // stays a stale lower bound; the burst drain reconciles it.
4369 let kvl = cache.kv[il].as_mut().unwrap();
4370 e.append_kv_quantized_rows_dc(
4371 &k,
4372 &v,
4373 &mut kvl.k,
4374 &mut kvl.v,
4375 ctr,
4376 t,
4377 kdk,
4378 kdv,
4379 ktb,
4380 vtb,
4381 Engine::kv_fp8_on(),
4382 )?;
4383 let upper = (kvl.len + t + 64).min(cache.max_ctx);
4384 let k_view = e.view_u8(&kvl.k, upper * ktb);
4385 let v_view = e.view_u8(&kvl.v, upper * vtb);
4386 e.fa_decode_rows_dc(
4387 &q, &k_view, &v_view, &mut attn, head_dim, n_head, n_head_kv, ctr,
4388 upper, t, scale, ktb, vtb, 0, false,
4389 )?;
4390 } else {
4391 for r in 0..t {
4392 // Owned per-row scratch: the b_n=1 kernels take packed batch buffers
4393 // whose row 0 is this row (arithmetic-free materialization copies,
4394 // same as decode's per-seq fallback arm).
4395 let mut k_row = e.uninit(kv_dim)?;
4396 e.dtod_copy_view(&k.slice(r * kv_dim..(r + 1) * kv_dim), &mut k_row)?;
4397 let mut v_row = e.uninit(kv_dim)?;
4398 e.dtod_copy_view(&v.slice(r * kv_dim..(r + 1) * kv_dim), &mut v_row)?;
4399 let pos_row = &pos_rows[r];
4400 let kvl = cache.kv[il].as_mut().unwrap();
4401 if seqs_append {
4402 e.append_kv_quantized_seqs(
4403 &k_row,
4404 &v_row,
4405 &kv_view.slice(0..2),
4406 pos_row,
4407 1,
4408 kdk,
4409 kdv,
4410 ktb,
4411 vtb,
4412 )?;
4413 kvl.len += 1;
4414 } else {
4415 e.append_kv_quantized_view(
4416 &k_row.slice(0..kv_dim),
4417 &v_row.slice(0..kv_dim),
4418 &mut kvl.k,
4419 &mut kvl.v,
4420 kvl.len,
4421 kvl.kv_dim_k,
4422 kvl.kv_dim_v,
4423 kvl.k_tok_bytes,
4424 kvl.v_tok_bytes,
4425 Engine::kv_fp8_on(),
4426 )?;
4427 kvl.len += 1;
4428 }
4429 let t_kv = kvl.len;
4430 let mut q_row = e.uninit(q_dim)?;
4431 e.dtod_copy_view(&q.slice(r * q_dim..(r + 1) * q_dim), &mut q_row)?;
4432 let mut a_row = e.uninit(q_dim)?;
4433 if batch_fa_on && crate::fa_seqs_eligible(t_kv, head_dim_global) {
4434 let sp0_r = crate::fa_split_keys(t_kv, cfg.n_head_kv as usize);
4435 e.fa_decode_batch_seqs_v4(
4436 &q_row,
4437 &kv_view.slice(0..2),
4438 pos_row,
4439 &mut a_row,
4440 head_dim,
4441 n_head,
4442 n_head_kv,
4443 1,
4444 t_kv,
4445 scale,
4446 sp0_r,
4447 ktb,
4448 vtb,
4449 )?;
4450 } else {
4451 let k_view = e.view_u8(&kvl.k, t_kv * kvl.k_tok_bytes);
4452 let v_view = e.view_u8(&kvl.v, t_kv * kvl.v_tok_bytes);
4453 let mut a_view = a_row.slice_mut(0..q_dim);
4454 e.fa_decode_kvmod_view(
4455 &q_row.slice(0..q_dim),
4456 &k_view,
4457 &v_view,
4458 &mut a_view,
4459 head_dim,
4460 n_head,
4461 n_head_kv,
4462 t_kv,
4463 scale,
4464 kvl.k_tok_bytes,
4465 kvl.v_tok_bytes,
4466 Engine::kv_fp8_on(),
4467 )?;
4468 }
4469 e.dtod_copy_into(&a_row, &mut attn, r * q_dim)?;
4470 }
4471 }
4472
4473 // Output gate (element-wise) + o-proj at m=T.
4474 let attn_g = match &gate {
4475 Some(g) => {
4476 let n = t * q_dim;
4477 let mut gsig = e.uninit(n)?;
4478 e.sigmoid(g, &mut gsig, n)?;
4479 let mut ag = e.uninit(n)?;
4480 e.mul(&attn, &gsig, &mut ag, n)?;
4481 ag
4482 }
4483 None => attn,
4484 };
4485 e.matmul(&fa.wo, &attn_g, t)?
4486 }
4487 // STREAM ARM (2b, lane/draftcost-moe): under a device position counter the
4488 // per-row serving-kernel chain cannot run (host state swaps keyed on host
4489 // row index are fine, but the stream COMMIT needs the GdnStash for its _dc
4490 // rebuild — the per-row chain only produces per-column clones). GDN rides
4491 // `linear_attn_verify_t`: batched q8_1-class projections, stash-producing,
4492 // and its one-scan recurrence is pinned bit-identical to T chained T=1
4493 // steps (its header + kernel-check). Position-independent, so no counter
4494 // plumbing is needed. Guards mirror the generic call site exactly.
4495 Mixer::Linear(la) if stream.is_some() => {
4496 if !(t >= 3 || (t == 2 && spec_m2()))
4497 || !self.mixer_in_q8_1_fast(e, &layer.mixer)
4498 || !e.uses_q8_1_fast(&la.ssm_out)
4499 {
4500 return Err("qwen35 stream verify: GDN batched arm requires t>=3 \
4501 (or MEMRA_SPEC_M2 at t=2) and q8_1-fast projections"
4502 .into());
4503 }
4504 let want = ckpt.is_some();
4505 let (out, stash) = self.linear_attn_verify_t(
4506 e,
4507 la,
4508 &xn,
4509 Some((&hq, &hd)),
4510 t,
4511 cache,
4512 il,
4513 want,
4514 )?;
4515 if let (Some(ck), Some(st)) = (ckpt.as_deref_mut(), stash) {
4516 ck.gdn[il] = Some(st);
4517 }
4518 out
4519 }
4520 Mixer::Linear(la) => {
4521 let ssm = cfg.ssm.as_ref().expect("linear mixer requires ssm cfg");
4522 let d_state = ssm.state_size as usize;
4523 let num_k = ssm.group_count as usize;
4524 let num_v = ssm.time_step_rank as usize;
4525 let d_conv = ssm.conv_kernel as usize;
4526 let key_dim = d_state * num_k;
4527 let value_dim = d_state * num_v;
4528 let conv_dim = key_dim * 2 + value_dim;
4529 let gdn_scale = 1.0 / (d_state as f32).sqrt();
4530
4531 // ---- batched projections: one weight read for all T rows ----
4532 let qkv_mixed = e.matmul_pre(&la.wqkv, &hq, &hd, &xn, t)?;
4533 let z = e.matmul_pre(&la.wqkv_gate, &hq, &hd, &xn, t)?;
4534 let beta_raw = e.matmul_pre(&la.ssm_beta, &hq, &hd, &xn, t)?;
4535 let alpha = e.matmul_pre(&la.ssm_alpha, &hq, &hd, &xn, t)?;
4536 let beta_w = la.ssm_beta.out_features();
4537 let alpha_w = la.ssm_alpha.out_features();
4538 let qkv_w = la.wqkv.out_features();
4539
4540 // ---- per-row state chain through the b_n=1 serving kernels ----
4541 // 6-entry alternating pointer table expresses the ping-pong without a
4542 // rebuild per row: even rows scan s0 -> s1, odd rows s1 -> s0. Host
4543 // handles swap per row so ckpt clones the canonical state (and the
4544 // post-verify canonical handle matches the last write), exactly as the
4545 // rowwise arm leaves them.
4546 let table = {
4547 let rl = cache.recur[il].as_ref().unwrap();
4548 let s = &e.gpu.stream();
4549 let (pc, _g0) = rl.conv_state.device_ptr(s);
4550 let (p0, _g1) = rl.ssm_state.device_ptr(s);
4551 let (p1, _g2) = rl.ssm_state_alt.device_ptr(s);
4552 e.htod_u64(&[
4553 pc as u64, p0 as u64, p1 as u64, pc as u64, p1 as u64, p0 as u64,
4554 ])?
4555 };
4556 let mut o_all = e.uninit(t * value_dim)?;
4557 let mut col_states: Option<Vec<(CudaSlice<f32>, CudaSlice<f32>)>> =
4558 if ckpt.is_some() && t >= 2 {
4559 Some(Vec::with_capacity(t - 1))
4560 } else {
4561 None
4562 };
4563 // Per-row scratch reused across rows (uninit is cheap but not free at
4564 // 48 layers x T rows); row inputs/outputs pass as VIEWS into the packed
4565 // [T, ...] buffers — zero arithmetic-free copies in this loop.
4566 let mut conv_out = e.uninit(conv_dim)?;
4567 let mut q_l2 = e.uninit(value_dim)?;
4568 let mut k_l2 = e.uninit(value_dim)?;
4569 let mut v_gd = e.uninit(value_dim)?;
4570 let mut beta_b = e.uninit(num_v)?;
4571 let mut g_log = e.uninit(num_v)?;
4572 for r in 0..t {
4573 let base = if r % 2 == 0 { 0 } else { 3 };
4574 let conv_view = table.slice(base..base + 1);
4575 let in_view = table.slice(base + 1..base + 2);
4576 let out_view = table.slice(base + 2..base + 3);
4577 e.ssm_conv1d_fused_decode_b_view(
4578 &qkv_mixed.slice(r * qkv_w..(r + 1) * qkv_w),
4579 &conv_view,
4580 la.ssm_conv1d.float_data(),
4581 &mut conv_out,
4582 conv_dim,
4583 d_conv,
4584 1,
4585 )?;
4586 e.gdn_prep_decode_b_view(
4587 &conv_out,
4588 &beta_raw.slice(r * beta_w..(r + 1) * beta_w),
4589 &alpha.slice(r * alpha_w..(r + 1) * alpha_w),
4590 la.ssm_dt.float_data(),
4591 la.ssm_a.float_data(),
4592 &mut q_l2,
4593 &mut k_l2,
4594 &mut v_gd,
4595 &mut beta_b,
4596 &mut g_log,
4597 d_state,
4598 num_v,
4599 num_k,
4600 key_dim,
4601 eps,
4602 conv_dim,
4603 1,
4604 )?;
4605 let mut o_row = o_all.slice_mut(r * value_dim..(r + 1) * value_dim);
4606 e.gdn_scan_s128_batched_view(
4607 &q_l2, &k_l2, &v_gd, &g_log, &beta_b, &in_view, &out_view, &mut o_row,
4608 num_v, 1, gdn_scale,
4609 )?;
4610 {
4611 let rl = cache.recur[il].as_mut().unwrap();
4612 std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
4613 }
4614 if r + 1 < t {
4615 if let Some(states) = col_states.as_mut() {
4616 let recur = cache.recur[il]
4617 .as_ref()
4618 .ok_or("qwen35 linear verify layer has no recurrent state")?;
4619 states.push((
4620 e.clone_dtod(&recur.conv_state)?,
4621 e.clone_dtod(&recur.ssm_state)?,
4622 ));
4623 }
4624 }
4625 }
4626 if let (Some(checkpoint), Some(states)) = (ckpt.as_deref_mut(), col_states) {
4627 checkpoint.cols[il] = Some(states);
4628 }
4629
4630 // ---- batched gated norm + out-projection at m=T ----
4631 if e.uses_q8_1_fast(&la.ssm_out) {
4632 let (gq, gd) = e.gated_rmsnorm_q8_1(
4633 &o_all,
4634 la.ssm_norm.float_data(),
4635 &z,
4636 d_state,
4637 t * num_v,
4638 eps,
4639 )?;
4640 let g0 = e.zeros(0)?;
4641 e.matmul_pre(&la.ssm_out, &gq, &gd, &g0, t)?
4642 } else {
4643 let mut gn = e.uninit(t * value_dim)?;
4644 e.gated_rmsnorm(
4645 &o_all,
4646 la.ssm_norm.float_data(),
4647 &z,
4648 &mut gn,
4649 d_state,
4650 t * num_v,
4651 eps,
4652 )?;
4653 e.matmul(&la.ssm_out, &gn, t)?
4654 }
4655 }
4656 };
4657
4658 // ---- residual add + post_attn_norm + FFN at m=T (serving dispatch verbatim) ----
4659 let pnorm = layer.post_attn_norm.float_data();
4660 let mut x1 = e.uninit(t * n_embd)?;
4661 let mut zn = e.uninit(t * n_embd)?;
4662 e.add_rms_norm(&x, &mixed, pnorm, &mut x1, &mut zn, n_embd, t, eps)?;
4663 let ffn_out = match &layer.ffn {
4664 crate::hybrid::Ffn::Dense {
4665 ffn_gate,
4666 ffn_up,
4667 ffn_down,
4668 } => {
4669 assert!(
4670 self.cfg.m3.is_none(),
4671 "qwen35 t-parallel verify: M3 swigluoai FFN not yet batched"
4672 );
4673 let n_ff = ffn_gate.out_features();
4674 let (zq, zd) = e.quantize_q8_1(&zn, t, n_embd)?;
4675 let g = e.matmul_pre(ffn_gate, &zq, &zd, &zn, t)?;
4676 let u = e.matmul_pre(ffn_up, &zq, &zd, &zn, t)?;
4677 let mut act = e.uninit(t * n_ff)?;
4678 e.silu_mul(&g, &u, &mut act, t * n_ff)?;
4679 let (aq, ad) = e.quantize_q8_1(&act, t, n_ff)?;
4680 e.matmul_pre(ffn_down, &aq, &ad, &act, t)?
4681 }
4682 crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il_zq8(e, m, &zn, None, t, il as u16)?,
4683 };
4684 let mut x2 = e.uninit(t * n_embd)?;
4685 e.add(&x1, &ffn_out, &mut x2, t * n_embd)?;
4686 // dspark drafter tap (no-op when no sink armed): post-layer residual verify rows
4687 self.dflash_tap(e, cache, il, &x2, t)?;
4688 x = x2;
4689 }
4690 Ok(x)
4691 }
4692
4693 /// PP-N STAGE SUBGRAPH of the verify trunk: layers `[lo, hi)` of `decode_step_t_core_stream`'s
4694 /// walk, verbatim. Enters with a MATERIALIZED `[T, n_embd]` residual (no pending fusion pair
4695 /// carried in from outside the range) and exits with the range's final residual materialized
4696 /// (the trailing add executed) — exactly the `decode_layers_eager(lo, hi)` contract, T rows
4697 /// instead of one.
4698 ///
4699 /// EXTRACTED (lane/pp2-spec 2026-08-06) rather than duplicated: `decode_step_t_core_stream` IS
4700 /// the single funnel every verify forward reaches, and its per-layer dispatch MIRRORING (norm
4701 /// fusion per layer, the t>=3/spec_m2 batched-linear window, the fused-q8 FFN chain, the
4702 /// decode-exact projections) is what makes verify bit-identical to eager decode. A second copy
4703 /// for the split arm is how those mirrors drift apart on the next lever. The unsplit body now
4704 /// calls this with `(0, n_layers)`, so the whole-trunk path and every stage range run the SAME
4705 /// code — there is no "split version" of the verify math.
4706 ///
4707 /// Bit-identity of a cut rests on the same kernel-check-pinned identity the eager arm's cut
4708 /// does — `add_rms_norm_q8_1 == add then rms_norm_q8_1` at nrows=T — because the ONLY thing a
4709 /// fence changes is that the cross-layer fusion carry breaks at `hi-1` and is re-materialized
4710 /// as an explicit `add`. `decode-batch-gate --mode ppspec` verifies end-to-end on real weights.
4711 #[allow(clippy::too_many_arguments)]
4712 fn verify_layers(
4713 &self,
4714 e: &Engine,
4715 mut x: CudaSlice<f32>,
4716 lo: usize,
4717 hi: usize,
4718 pos_d: &CudaSlice<i32>,
4719 pos0: usize,
4720 t: usize,
4721 cache: &mut Cache,
4722 mut ckpt: Option<&mut VerifyCkpt>,
4723 stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
4724 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4725 if self.cfg.step35.is_some() {
4726 if stream.is_some() {
4727 return Err(
4728 "step35 has no ROUND-STREAM verify arm (the device-counter _dc twins \
4729 cannot express the SWA offset KV view)"
4730 .into(),
4731 );
4732 }
4733 return self.step35_verify_batch_layers(e, x, lo, hi, pos0, t, cache);
4734 }
4735 if self.qwen35_serving_class() {
4736 return self.qwen35_verify_batch_layers(
4737 e,
4738 x,
4739 lo,
4740 hi,
4741 pos0,
4742 t,
4743 cache,
4744 ckpt.take(),
4745 stream,
4746 );
4747 }
4748 let n_embd = self.cfg.n_embd as usize;
4749 let eps = self.cfg.rms_eps;
4750 // CROSS-LAYER ADD+NORM FUSION (lane/vt-fixes fix 2, mirroring decode_step_h's
4751 // launch-arc form): layer il's post-FFN residual add (x2 = x1 + ffn_out) and layer
4752 // il+1's attn_norm(+quantize) are consecutive row-wise ops — ONE add_rms_norm_q8_1
4753 // launch at nrows=t does all three (bit-identity pinned by the T-row kernel-check
4754 // arms). Carry the un-added (x1, ffn_out) pair; the fused launch materializes x2 (the
4755 // residual the next layer needs) as its `res` output. Falls back to the separate add
4756 // when the next layer is off the fused-q8 path.
4757 let mut pending: Option<(CudaSlice<f32>, CudaSlice<f32>)> = None;
4758 for il in lo..hi {
4759 let layer = &self.layers[il];
4760 // DISPATCH-MIRRORED attn-input RMSNorm (FP-order lesson #8): eager decode fuses the
4761 // 1024-thread rms_norm_q8_1 ONLY when every mixer projection is q8_1-fast; layers with
4762 // Float projections (ssm_beta/ssm_alpha on layers 1/2/4 of the 9B NVFP4 GGUF) take the
4763 // UNFUSED 256-thread rms_norm. The verify norm must mirror that PER-LAYER choice —
4764 // blockDim changes the sum-of-squares reduce order, and the ULP shift amplifies through
4765 // the GDN recurrence into argmax flips (measured: 9B text prompt, 1 ULP at layer 2 ->
4766 // 2.3e-1 logit maxdiff at the head -> K=1..8 divergence at a 0.03-margin token).
4767 let mixer_fast = self.mixer_in_q8_1_fast(e, &layer.mixer);
4768 let norm_fused = std::env::var("MEMRA_NO_FUSE_NORMQ").is_err() && mixer_fast;
4769 // BATCHED EPILOGUE RE-FUSE (lane/vt-fixes fix 2, 2026-08-03): when the norm is
4770 // dispatch-fused AND every consumer of `h` reads only its q8_1 form (Full mixer:
4771 // projections only; Linear mixer: the batched arm — the per-column fallback needs
4772 // f32 h), emit the attn-input norm DIRECTLY as q8_1 via `rms_norm_q8_1` at nrows=t
4773 // (row-indexed kernel — the T-row launch is the per-row m=1 program, kernel-check
4774 // pins bit-identity vs rms_norm_decode -> quantize_q8_1). Kills the standalone
4775 // quantize launch(es) + the f32 h HBM round-trip that decode never pays.
4776 // step35 (Full mixer) is the third case that needs f32 `h`: its verify arm is a
4777 // per-ROW replay of the eager decode mixer, whose `pre_q` contract is a single row —
4778 // a T-row q8_1 pair cannot be handed to it, and re-deriving per-row q8_1 from the f32
4779 // rows is exactly the dispatch being mirrored. Keep step35 on the unfused arm.
4780 let lin_q8_only = match &layer.mixer {
4781 Mixer::Linear(la) => {
4782 (t >= 3 || (t == 2 && spec_m2())) && e.uses_q8_1_fast(&la.ssm_out)
4783 }
4784 Mixer::Full(_) if self.cfg.step35.is_some() => false,
4785 _ => true,
4786 };
4787 // NOTE decode.rs's take()-first lesson: take the pending pair BEFORE branching so
4788 // a non-fused layer still performs the residual add.
4789 let taken = pending.take();
4790 let (h, h_q8) = if norm_fused && lin_q8_only {
4791 let pair = match taken {
4792 // fused add + attn_norm + q8_1: ONE launch resolves the carried residual
4793 // AND emits this layer's mixer input pre-quantized. res -> x2 (= new x).
4794 Some((x1p, f1p)) => {
4795 let mut x2 = vbuf(e, t * n_embd)?; // fully written (res output)
4796 let p = e.add_rms_norm_q8_1(
4797 &x1p,
4798 &f1p,
4799 layer.attn_norm.float_data(),
4800 &mut x2,
4801 n_embd,
4802 t,
4803 eps,
4804 )?;
4805 x = x2;
4806 p
4807 }
4808 None => e.rms_norm_q8_1(&x, layer.attn_norm.float_data(), n_embd, t, eps)?,
4809 };
4810 (e.zeros(0)?, Some(pair)) // h unused on this path (q8-only consumers)
4811 } else {
4812 if let Some((x1p, f1p)) = taken {
4813 let mut x2 = vbuf(e, t * n_embd)?; // fully written by add
4814 e.add(&x1p, &f1p, &mut x2, t * n_embd)?;
4815 x = x2;
4816 }
4817 let mut h = vbuf(e, t * n_embd)?; // fully written by either rms_norm arm
4818 if norm_fused {
4819 e.rms_norm_decode(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
4820 } else {
4821 e.rms_norm(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
4822 }
4823 (h, None)
4824 };
4825 let h_q8_ref = h_q8.as_ref().map(|(q, d)| (q, d));
4826
4827 let mixed = match &layer.mixer {
4828 Mixer::Full(fa) => self.full_attn_verify(
4829 e,
4830 fa,
4831 &h,
4832 h_q8_ref,
4833 pos_d,
4834 t,
4835 cache,
4836 il,
4837 stream.map(|(_, c)| c),
4838 )?,
4839 Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
4840 Mixer::Linear(la) => {
4841 // BATCHED linear verify (2026-07-03, the MTP-profit lever): one T-token pass —
4842 // batched projections (weight read ONCE, hits the m=2-4 weight-resident matvec),
4843 // carried-state conv (ssm_conv1d_tm_state), GDN prep on the prefill kernels, and
4844 // ONE gdn_scan whose internal sequential t-loop is the SAME recurrence as T
4845 // chained T=1 steps (bit-identical). Falls back to the sequential per-column
4846 // chain when T < d_conv-1 (conv ring update needs T >= pad) — or when ANY
4847 // projection is off the q8_1 fast path: matmul_decode_exact would route a Float
4848 // tensor to cuBLAS at m=t (different FP accumulation than eager's per-token
4849 // GEMV), so mixed-dtype layers stay on the eager-identical per-column chain.
4850 // MEMRA_SPEC_M2 (lane/spec-m2): the t==2 batch rides the same arm — the conv
4851 // wrapper handles t<pad with a pure-copy ring rebuild; see spec_m2() header.
4852 if (t >= 3 || (t == 2 && spec_m2()))
4853 && mixer_fast
4854 && e.uses_q8_1_fast(&la.ssm_out)
4855 {
4856 let want = ckpt.is_some();
4857 let (out, stash) =
4858 self.linear_attn_verify_t(e, la, &h, h_q8_ref, t, cache, il, want)?;
4859 if let (Some(ck), Some(st)) = (ckpt.as_deref_mut(), stash) {
4860 ck.gdn[il] = Some(st);
4861 }
4862 out
4863 } else {
4864 let mut out = vbuf(e, t * n_embd)?; // every col written by copy_into
4865 let mut col_states: Option<Vec<(CudaSlice<f32>, CudaSlice<f32>)>> =
4866 if ckpt.is_some() && t >= 2 {
4867 Some(Vec::with_capacity(t - 1))
4868 } else {
4869 None
4870 };
4871 for col in 0..t {
4872 let mut h_col = vbuf(e, n_embd)?; // fully written by copy_view_into
4873 let src = h.slice(col * n_embd..(col + 1) * n_embd);
4874 e.copy_view_into(&mut h_col, 0, &src, n_embd)?;
4875 let m_col = self.linear_attn_decode(e, la, &h_col, cache, il)?;
4876 e.copy_into(&mut out, col * n_embd, &m_col, n_embd)?;
4877 // REPLAY-FREE ckpt: clone the chain's ACTUAL state after this column
4878 // (pure dtod — cannot change any computed value). Last column skipped:
4879 // rebuild targets are j <= t-1 columns.
4880 if let Some(cs) = col_states.as_mut() {
4881 if col + 1 < t {
4882 let rl = cache.recur[il].as_ref().unwrap();
4883 cs.push((
4884 e.clone_dtod(&rl.conv_state)?,
4885 e.clone_dtod(&rl.ssm_state)?,
4886 ));
4887 }
4888 }
4889 }
4890 if let (Some(ck), Some(cs)) = (ckpt.as_deref_mut(), col_states) {
4891 // ReplaySSM-assessment instrumentation (2026-07-30): the
4892 // per-column clones are the only true state snapshots left in
4893 // the verify (the batched path stashes INPUTS and replays).
4894 if std::env::var("MEMRA_SPEC_STATS").as_deref() == Ok("1") {
4895 static ONCE: std::sync::Once = std::sync::Once::new();
4896 let bytes: usize =
4897 cs.iter().map(|(c, s)| (c.len() + s.len()) * 4).sum();
4898 ONCE.call_once(|| eprintln!(
4899 "[verify-ckpt] per-column layer il={il}: {} clones, {:.2} MB/layer/round",
4900 cs.len(), bytes as f64 / 1e6));
4901 }
4902 ck.cols[il] = Some(cs);
4903 }
4904 out
4905 }
4906 }
4907 };
4908
4909 // DISPATCH-MIRRORED post-attn norm: eager residual_norm_ffn fuses add+norm+quant
4910 // (1024-thread add_rms_norm_q8_1) only for Dense FFNs whose gate+up are q8_1-fast;
4911 // otherwise (and for MoE) it runs the 256-thread fused add_rms_norm. Mirror per layer.
4912 let ffn_fuse = match &layer.ffn {
4913 crate::hybrid::Ffn::Dense {
4914 ffn_gate, ffn_up, ..
4915 } => {
4916 std::env::var("MEMRA_NO_FUSE_NORMQ").is_err()
4917 && e.uses_q8_1_fast(ffn_gate)
4918 && e.uses_q8_1_fast(ffn_up)
4919 }
4920 crate::hybrid::Ffn::Moe(_) => false,
4921 };
4922 // BATCHED EPILOGUE RE-FUSE (lane/vt-fixes fix 2): on the ffn_fuse path (Dense,
4923 // gate+up q8_1-fast, non-M3) the FFN input is emitted DIRECTLY as q8_1 by ONE
4924 // add_rms_norm_q8_1 launch at nrows=t (row-indexed kernel: the T-row launch is the
4925 // per-row m=1 program; kernel-check pins bit-identity vs the unfused
4926 // add_f32 -> rms_norm_decode -> quantize_q8_1 chain at T=2/4/5/8) — replacing the
4927 // add + rms_norm_decode launches AND the dual/singles' internal re-quantize.
4928 // M3's swigluoai must keep the f32 chain (the fused SwiGLU epilogue encodes plain
4929 // SiLU), mirroring residual_norm_ffn's m3 guard on the decode path.
4930 // step35: same guard per LAYER. A dense FFN's clamp is the SHEXP array (upstream's
4931 // one build_ffn serves dense + shared expert, llama-graph.cpp:1751), and verify MUST
4932 // mirror decode's dispatch or spec self-consistency fails.
4933 let dense_lim = self.cfg.clamp_shexp_at(il as u32);
4934 let fuse_q8 = ffn_fuse && self.cfg.m3.is_none() && dense_lim.is_none();
4935 let mut x1 = vbuf(e, t * n_embd)?; // fully written by add / add_rms_norm*
4936 let mut z = e.zeros(0)?; // replaced below on the unfused arms
4937 let z_q8 = if fuse_q8 {
4938 Some(e.add_rms_norm_q8_1(
4939 &x,
4940 &mixed,
4941 layer.post_attn_norm.float_data(),
4942 &mut x1,
4943 n_embd,
4944 t,
4945 eps,
4946 )?)
4947 } else {
4948 let mut zf = vbuf(e, t * n_embd)?; // fully written by rms_norm_decode / add_rms_norm
4949 if ffn_fuse {
4950 e.add(&x, &mixed, &mut x1, t * n_embd)?;
4951 e.rms_norm_decode(
4952 &x1,
4953 layer.post_attn_norm.float_data(),
4954 &mut zf,
4955 n_embd,
4956 t,
4957 eps,
4958 )?;
4959 } else {
4960 e.add_rms_norm(
4961 &x,
4962 &mixed,
4963 layer.post_attn_norm.float_data(),
4964 &mut x1,
4965 &mut zf,
4966 n_embd,
4967 t,
4968 eps,
4969 )?;
4970 }
4971 z = zf;
4972 None
4973 };
4974 // DECODE-EXACT FFN projections: force MMVQ for gate/up/down at any T to match the
4975 // T=1 decode FP accumulation order. At T>=5 the generic matmul/matmul_pre falls to dp4a
4976 // (128-thread, different FP sum order). At T=2-4 the batched MMVQ is already bit-identical.
4977 let ffn_out = match &layer.ffn {
4978 crate::hybrid::Ffn::Dense {
4979 ffn_gate,
4980 ffn_up,
4981 ffn_down,
4982 } => {
4983 let n_ff = ffn_gate.out_features();
4984 if let Some((zq, zd)) = z_q8.as_ref() {
4985 // FUSED CHAIN (fix 2): pre-quantized z feeds the projections; the SwiGLU
4986 // epilogue emits act pre-quantized for ffn_down (silu_mul_scaled_q8_1,
4987 // bit-identical to silu_mul + quantize — kernel-check-pinned) with the
4988 // NVFP4 macro-scales folded (deferred-scale dual: y*s inline == the
4989 // scale_inplace store, value-exact) — the exact m=1 decode epilogue
4990 // structure at nrows=t.
4991 let pair =
4992 match e.matmul_decode_exact_dual_pre(ffn_gate, ffn_up, zq, zd, t)? {
4993 Some(((g, gs), (u, us))) => Some((g, gs, u, us)),
4994 None => None,
4995 };
4996 let (gate, gs, up, us) = match pair {
4997 Some(x4) => x4,
4998 None => (
4999 e.matmul_decode_exact_pre(ffn_gate, zq, zd, t)?,
5000 1.0, // scale already applied inside _pre
5001 e.matmul_decode_exact_pre(ffn_up, zq, zd, t)?,
5002 1.0,
5003 ),
5004 };
5005 if e.uses_q8_1_fast(ffn_down) {
5006 let (aq, ad) = e.silu_mul_scaled_q8_1(&gate, &up, gs, us, t * n_ff)?;
5007 e.matmul_decode_exact_pre(ffn_down, &aq, &ad, t)?
5008 } else {
5009 let mut act = vbuf(e, t * n_ff)?;
5010 e.silu_mul_scaled(&gate, &up, gs, us, &mut act, t * n_ff)?;
5011 e.matmul_decode_exact(ffn_down, &act, t)?
5012 }
5013 } else {
5014 // UNFUSED (pre-fix) chain — MoE-adjacent/M3/off-fast layers, unchanged.
5015 // DUAL gate+up batched twin (lane/verify-economics, 2026-08-02): one launch
5016 // for the pair at t=2..8 — bit-identical per (tensor,token,row) to the two
5017 // singles (kernel-check pins bitwise; MEMRA_SPEC_DUAL_T=0 reverts). None
5018 // (non-NVFP4 / t outside the tier / seam off) -> the two singles, unchanged.
5019 let (gate, up) =
5020 match e.matmul_decode_exact_dual(ffn_gate, ffn_up, &z, t)? {
5021 Some(pair) => pair,
5022 None => (
5023 e.matmul_decode_exact(ffn_gate, &z, t)?,
5024 e.matmul_decode_exact(ffn_up, &z, t)?,
5025 ),
5026 };
5027 let mut act = vbuf(e, t * n_ff)?; // fully written by ffn_act_lim
5028 Self::ffn_act_lim(
5029 e,
5030 &self.cfg,
5031 &gate,
5032 &up,
5033 1.0,
5034 1.0,
5035 dense_lim,
5036 &mut act,
5037 t * n_ff,
5038 )?;
5039 e.matmul_decode_exact(ffn_down, &act, t)?
5040 }
5041 }
5042 crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il(e, m, &z, t, il as u16)?,
5043 };
5044 // CROSS-LAYER fusion: defer this layer's post-FFN residual add — the next layer's
5045 // fused-q8 attn norm folds it in (add_rms_norm_q8_1 == add; rms_norm; quantize,
5046 // kernel-check-pinned at nrows=T). Non-fused next layers add explicitly above.
5047 pending = Some((x1, ffn_out));
5048 }
5049 // RANGE's final add (no next norm INSIDE the range to fuse with; for the
5050 // whole-trunk call that is the last layer, whose next norm is output_norm — f32-out).
5051 if let Some((x1p, f1p)) = pending.take() {
5052 let mut x2 = vbuf(e, t * n_embd)?; // fully written by add
5053 e.add(&x1p, &f1p, &mut x2, t * n_embd)?;
5054 x = x2;
5055 }
5056 Ok(x)
5057 }
5058 /// BATCHED linear-attn verify (T=K+1): the whole layer in ~10 launches instead of T x the
5059 /// T=1 decode chain (T x ~12 launches + T weight reads of the four projections). The GDN
5060 /// recurrence itself is inherently sequential — gdn_scan_s128 runs its internal t-loop with
5061 /// the SAME per-token math as chained T=1 calls (bit-identical state evolution); everything
5062 /// around it (projections, conv, prep, gated norm, out-proj) batches. Advances conv ring +
5063 /// ssm state exactly like T sequential decode steps.
5064 /// `want_stash`: additionally RETAIN the gdn-scan inputs (pure buffer keep-alives, zero extra
5065 /// kernels) so a partial accept can rebuild the state after any column prefix (REPLAY-FREE).
5066 #[allow(clippy::too_many_arguments)]
5067 fn linear_attn_verify_t(
5068 &self,
5069 e: &Engine,
5070 la: &LinearAttnLayer,
5071 h: &CudaSlice<f32>,
5072 h_q8: Option<(&CudaSlice<i8>, &CudaSlice<f32>)>,
5073 t: usize,
5074 cache: &mut Cache,
5075 il: usize,
5076 want_stash: bool,
5077 ) -> Result<(CudaSlice<f32>, Option<GdnStash>), Box<dyn std::error::Error>> {
5078 let cfg = &self.cfg;
5079 let ssm = cfg.ssm.as_ref().unwrap();
5080 let d_state = ssm.state_size as usize;
5081 let num_k = ssm.group_count as usize;
5082 let num_v = ssm.time_step_rank as usize;
5083 let d_conv = ssm.conv_kernel as usize;
5084 let key_dim = d_state * num_k;
5085 let conv_dim = key_dim * 2 + d_state * num_v;
5086 let eps = cfg.rms_eps;
5087 let scale = 1.0 / (d_state as f32).sqrt();
5088
5089 // DECODE-EXACT projections: matmul_decode_exact forces the MMVQ (warp-per-row, 32-thread)
5090 // accumulation order for EVERY m, matching the T=1 decode path bit-for-bit. The generic
5091 // `matmul` at m>=5 falls to dp4a (128-thread, two-level reduce) which has a different FP
5092 // sum order — ULP differences propagate through gdn_scan and flip argmax on the 27B.
5093 // Q8 TRUNK-FUSION at T=1 (35B: wqkv+wqkv_gate both Q8_0): one fused2 launch, bit-identical
5094 // per (tensor,row) to the two m=1 MMVQ dispatches below — decode-exact contract holds.
5095 // VERIFY-TIER TRUNK FUSION (MEMRA_SPEC_FUSED_T, t=2-4): quantize h ONCE for every
5096 // fused-eligible same-input Q8_0 pair of this layer (35B wqkv+wqkv_gate; 9B
5097 // ssm_beta+ssm_alpha) — each fused2 batched launch then replaces two decode-exact
5098 // calls (each of which re-quantizes the same h + runs its own _b2/_b4 launch).
5099 // Bit-identical per (tensor,token,row) — see spec_fused_t().
5100 // BATCHED EPILOGUE RE-FUSE (lane/vt-fixes fix 2): `h_q8` = the attn-input norm emitted
5101 // directly as q8_1 by the caller's fused rms_norm_q8_1 (bit-identical to the unfused
5102 // chain, kernel-check-pinned). When present it REPLACES the standalone quantize below
5103 // and feeds every projection; the caller guaranteed all four input projections are
5104 // q8_1-fast. When absent, the old shared-quantize (fused-t window) stands.
5105 let h_q8_t = if h_q8.is_none()
5106 && spec_fused_t()
5107 && (2..=4).contains(&t)
5108 && ((e.uses_q8_1_fast(&la.wqkv) && e.uses_q8_1_fast(&la.wqkv_gate))
5109 || (e.uses_q8_1_fast(&la.ssm_beta) && e.uses_q8_1_fast(&la.ssm_alpha)))
5110 {
5111 Some(e.quantize_q8_1(h, t, cfg.n_embd as usize)?)
5112 } else {
5113 None
5114 };
5115 // one view: the caller's fused-norm q8 or this fn's own shared quantize.
5116 let hq8_any: Option<(&CudaSlice<i8>, &CudaSlice<f32>)> =
5117 h_q8.or(h_q8_t.as_ref().map(|(q, d)| (q, d)));
5118 let (qkv_mixed, z) = {
5119 let mut fused = None;
5120 if t == 1 && e.uses_q8_1_fast(&la.wqkv) && e.uses_q8_1_fast(&la.wqkv_gate) {
5121 let (hq, hd) = e.quantize_q8_1(h, 1, cfg.n_embd as usize)?;
5122 fused = e.matmul_q8_fused2(&la.wqkv, &la.wqkv_gate, &hq, &hd)?;
5123 } else if let Some((hq, hd)) = hq8_any {
5124 if spec_fused_t() && (2..=4).contains(&t) {
5125 fused = e.matmul_q8_fused2_t(&la.wqkv, &la.wqkv_gate, hq, hd, t)?;
5126 }
5127 }
5128 match (fused, hq8_any) {
5129 (Some(pair), _) => pair,
5130 (None, Some((hq, hd))) if h_q8.is_some() => (
5131 e.matmul_decode_exact_pre(&la.wqkv, hq, hd, t)?,
5132 e.matmul_decode_exact_pre(&la.wqkv_gate, hq, hd, t)?,
5133 ),
5134 (None, _) => (
5135 e.matmul_decode_exact(&la.wqkv, h, t)?,
5136 e.matmul_decode_exact(&la.wqkv_gate, h, t)?,
5137 ),
5138 }
5139 };
5140 // beta+alpha DUAL at T=1 (75% of p3 rounds run T=1 verify — p-min chain cuts): the dual
5141 // mr2 kernel is bit-identical per element to the m=1 MMVQ matmul_decode_exact dispatches
5142 // (same warp-per-row body, blockIdx.y picks the weight), so the decode-exact contract
5143 // holds; the run-spec battery is the arbiter. T>1 keeps the per-tensor decode-exact path.
5144 let (beta_raw, alpha) = if t == 1 {
5145 let (hq, hd) = e.quantize_q8_1(h, 1, cfg.n_embd as usize)?;
5146 match e.matmul_pre_dual_noscale(&la.ssm_beta, &la.ssm_alpha, &hq, &hd, 1)? {
5147 Some(((mut b, bs), (mut a, as_))) => {
5148 if bs != 1.0 {
5149 e.scale_inplace(&mut b, bs, la.ssm_beta.out_features())?;
5150 }
5151 if as_ != 1.0 {
5152 e.scale_inplace(&mut a, as_, la.ssm_alpha.out_features())?;
5153 }
5154 (b, a)
5155 }
5156 // Q8_0 fused2 twin (9B stores beta/alpha as Q8_0): DISPATCH-MIRRORS the eager
5157 // decode's beta_alpha closure — the fused body is qmatvec_q8_0_mmvq verbatim,
5158 // bit-identical per row (kernel-check rel=0.00e0 gate), so decode==verify holds.
5159 None => match e.matmul_q8_fused2(&la.ssm_beta, &la.ssm_alpha, &hq, &hd)? {
5160 Some((b, a)) => (b, a),
5161 None => (
5162 e.matmul_decode_exact(&la.ssm_beta, h, 1)?,
5163 e.matmul_decode_exact(&la.ssm_alpha, h, 1)?,
5164 ),
5165 },
5166 }
5167 } else {
5168 // fused-t twin (9B stores beta/alpha as Q8_0): same shared-quantize + one launch
5169 // contract as the wqkv pair above; 35B beta/alpha are Float -> None -> fallback.
5170 let mut nvfp4_fused = None;
5171 let mut q8_fused = None;
5172 if let Some((hq, hd)) = hq8_any {
5173 if t == 3 && std::env::var("MEMRA_NVFP4_AUX_DUAL").as_deref() != Ok("0") {
5174 nvfp4_fused =
5175 e.matmul_decode_exact_dual_pre(&la.ssm_beta, &la.ssm_alpha, hq, hd, t)?;
5176 if nvfp4_fused.is_some() && std::env::var("MEMRA_DEBUG").is_ok() {
5177 static ONCE: std::sync::Once = std::sync::Once::new();
5178 ONCE.call_once(|| {
5179 eprintln!("[memra] NVFP4 beta+alpha batched aux dual ENGAGED (t={t})")
5180 });
5181 }
5182 }
5183 if nvfp4_fused.is_none() && spec_fused_t() && (2..=4).contains(&t) {
5184 q8_fused = e.matmul_q8_fused2_t(&la.ssm_beta, &la.ssm_alpha, hq, hd, t)?;
5185 }
5186 }
5187 if let Some(((mut b, bs), (mut a, as_))) = nvfp4_fused {
5188 if bs != 1.0 {
5189 e.scale_inplace(&mut b, bs, t * la.ssm_beta.out_features())?;
5190 }
5191 if as_ != 1.0 {
5192 e.scale_inplace(&mut a, as_, t * la.ssm_alpha.out_features())?;
5193 }
5194 (b, a)
5195 } else if let Some(pair) = q8_fused {
5196 pair
5197 } else {
5198 match hq8_any {
5199 Some((hq, hd)) if h_q8.is_some() => (
5200 e.matmul_decode_exact_pre(&la.ssm_beta, hq, hd, t)?,
5201 e.matmul_decode_exact_pre(&la.ssm_alpha, hq, hd, t)?,
5202 ),
5203 _ => (
5204 e.matmul_decode_exact(&la.ssm_beta, h, t)?,
5205 e.matmul_decode_exact(&la.ssm_alpha, h, t)?,
5206 ),
5207 }
5208 }
5209 };
5210
5211 // conv with CARRIED state + ring roll (T >= pad rides the input-column update kernel;
5212 // T < pad — the MEMRA_SPEC_M2 t=2 arm — rolls via the pure-copy ring rebuild).
5213 let rl = cache.recur[il].as_mut().unwrap();
5214 let mut conv_out = e.uninit(conv_dim * t)?;
5215 e.ssm_conv1d_tm_state(
5216 &qkv_mixed,
5217 &mut rl.conv_state,
5218 la.ssm_conv1d.float_data(),
5219 &mut conv_out,
5220 conv_dim,
5221 t,
5222 d_conv,
5223 )?;
5224
5225 // GDN prep via the prefill kernels (repack + L2 + sigmoid + glog), T-wide.
5226 let mut q_g = e.uninit(d_state * num_v * t)?;
5227 let mut k_g = e.uninit(d_state * num_v * t)?;
5228 let mut v_g = e.uninit(d_state * num_v * t)?;
5229 e.qkv_to_gdn_repack(
5230 &conv_out, &mut q_g, &mut k_g, &mut v_g, d_state, num_v, num_k, key_dim, t,
5231 )?;
5232 let mut q_l2 = e.uninit(d_state * num_v * t)?;
5233 e.l2_norm_decode(&q_g, &mut q_l2, d_state, num_v * t, eps)?;
5234 let mut k_l2 = e.uninit(d_state * num_v * t)?;
5235 e.l2_norm_decode(&k_g, &mut k_l2, d_state, num_v * t, eps)?;
5236 let mut beta = e.uninit(t * num_v)?;
5237 e.sigmoid(&beta_raw, &mut beta, t * num_v)?;
5238 let mut g_log = e.uninit(t * num_v)?;
5239 e.gdn_glog(
5240 &alpha,
5241 la.ssm_dt.float_data(),
5242 la.ssm_a.float_data(),
5243 &mut g_log,
5244 num_v,
5245 t,
5246 )?;
5247
5248 // ONE gdn_scan over T tokens from the carried state (internal sequential loop ==
5249 // T chained T=1 steps). Ping-pong the resident buffers like eager decode.
5250 let mut o = e.uninit(d_state * num_v * t)?;
5251 {
5252 let crate::cache::RecurLayer {
5253 ssm_state,
5254 ssm_state_alt,
5255 ..
5256 } = rl;
5257 e.gdn_scan_s128(
5258 &q_l2,
5259 &k_l2,
5260 &v_g,
5261 &g_log,
5262 &beta,
5263 ssm_state,
5264 ssm_state_alt,
5265 &mut o,
5266 num_v,
5267 t,
5268 scale,
5269 )?;
5270 }
5271 std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
5272
5273 // gated RMSNorm + out projection, T-wide. FUSED-QUANTIZE ARM (lane/vt-fixes fix 2,
5274 // mirroring the T=1 decode's launch-arc form): when ssm_out rides the q8_1 fast path,
5275 // emit q8_1 straight from the gated norm at nrows=num_v*t (row-indexed kernel, the
5276 // T-wide launch is the per-row program; kernel-check pins bit-identity vs
5277 // gated_rmsnorm -> quantize_q8_1 at T=1 and T=5) and feed the decode-exact dispatch
5278 // pre-quantized — one launch replaces norm + quantize. Fallback = the f32 chain.
5279 let out = if e.uses_q8_1_fast(&la.ssm_out) {
5280 let (gq, gd) =
5281 e.gated_rmsnorm_q8_1(&o, la.ssm_norm.float_data(), &z, d_state, num_v * t, eps)?;
5282 e.matmul_decode_exact_pre(&la.ssm_out, &gq, &gd, t)?
5283 } else {
5284 let mut gn = e.uninit(d_state * num_v * t)?;
5285 e.gated_rmsnorm(
5286 &o,
5287 la.ssm_norm.float_data(),
5288 &z,
5289 &mut gn,
5290 d_state,
5291 num_v * t,
5292 eps,
5293 )?;
5294 // DECODE-EXACT out-projection: same MMVQ path as the T=1 decode (ssm_out at m>=5
5295 // would fall to dp4a with a different FP reduction order — same class of bug as
5296 // the input projs).
5297 e.matmul_decode_exact(&la.ssm_out, &gn, t)?
5298 };
5299 let stash = if want_stash {
5300 Some(GdnStash {
5301 qkv_mixed,
5302 q_l2,
5303 k_l2,
5304 v_g,
5305 g_log,
5306 beta,
5307 })
5308 } else {
5309 None
5310 };
5311 Ok((out, stash))
5312 }
5313
5314 /// REPLAY-FREE partial-accept commit (2026-07-03): make the cache state == "committed through
5315 /// the first `j` verify columns" WITHOUT the legacy rollback + duplicate trunk replay.
5316 /// - Full-attn KV: truncate len to snapshot + j. The verify's appended rows for those columns
5317 /// are bit-identical to what an eager T=1 chain writes (the decode-exact contract the
5318 /// verify-probe gates), so keeping them == replaying them.
5319 /// - Linear layers, batched path: rebuild the conv ring by PURE COPIES (ring holds raw input
5320 /// columns) and the ssm state by a prefix re-run of the SAME gdn_scan kernel (t=j) from the
5321 /// snapshot state over the stash's identical inputs — the kernel's t-loop carries state in
5322 /// registers and writes it once at the end, so iterations 0..j-1 are independent of T:
5323 /// bit-identical to the verify's own state after j tokens == the eager chain state.
5324 /// - Linear layers, per-column path: restore the cloned actual state after column j-1.
5325 /// Caller guarantees 1 <= j <= t-1 (j==0 rounds take the legacy rollback; j==t is full accept).
5326 fn commit_verified_prefix(
5327 &self,
5328 e: &Engine,
5329 cache: &mut Cache,
5330 snap: &crate::cache::CacheSnapshot,
5331 ckpt: &VerifyCkpt,
5332 j: usize,
5333 kv_lens_done: bool,
5334 dev_j: Option<(&CudaSlice<u32>, usize, usize)>,
5335 ) -> Result<(), Box<dyn std::error::Error>> {
5336 let cfg = &self.cfg;
5337 let ssm = cfg.ssm.as_ref().unwrap();
5338 let d_state = ssm.state_size as usize;
5339 let num_k = ssm.group_count as usize;
5340 let num_v = ssm.time_step_rank as usize;
5341 let d_conv = ssm.conv_kernel as usize;
5342 let conv_dim = d_state * num_k * 2 + d_state * num_v;
5343 let scale = 1.0 / (d_state as f32).sqrt();
5344 for il in 0..self.layers.len() {
5345 if let (Some(kvl), Some(saved)) = (cache.kv[il].as_mut(), snap.kv_len[il]) {
5346 kvl.len = saved + j;
5347 // devacc 3a: spec_rollback_kv already wrote len_d on-device (same value).
5348 if !kv_lens_done {
5349 e.set_i32_one(&mut kvl.len_d, kvl.len as i32)?;
5350 }
5351 }
5352 if let Some(rl) = cache.recur[il].as_mut() {
5353 if let Some(st) = &ckpt.gdn[il] {
5354 let ring_old = snap.conv[il].as_ref().expect("snapshot missing conv");
5355 let state_in = snap.ssm[il].as_ref().expect("snapshot missing ssm");
5356 if let Some((acc, base, t_v)) = dev_j {
5357 // 3b: j read on-device (_dc twins, same bodies; full accept early-exits).
5358 e.ssm_conv_ring_rebuild_dc(
5359 &st.qkv_mixed,
5360 ring_old,
5361 &mut rl.conv_state,
5362 conv_dim,
5363 acc,
5364 base,
5365 t_v,
5366 d_conv,
5367 )?;
5368 let mut o = e.uninit(d_state * num_v * j.max(1))?;
5369 e.gdn_scan_s128_dc(
5370 &st.q_l2,
5371 &st.k_l2,
5372 &st.v_g,
5373 &st.g_log,
5374 &st.beta,
5375 state_in,
5376 &mut rl.ssm_state,
5377 &mut o,
5378 num_v,
5379 acc,
5380 base,
5381 t_v,
5382 scale,
5383 )?;
5384 } else {
5385 e.ssm_conv_ring_rebuild(
5386 &st.qkv_mixed,
5387 ring_old,
5388 &mut rl.conv_state,
5389 conv_dim,
5390 j,
5391 d_conv,
5392 )?;
5393 let mut o = e.uninit(d_state * num_v * j)?; // scan output, discarded
5394 e.gdn_scan_s128(
5395 &st.q_l2,
5396 &st.k_l2,
5397 &st.v_g,
5398 &st.g_log,
5399 &st.beta,
5400 state_in,
5401 &mut rl.ssm_state,
5402 &mut o,
5403 num_v,
5404 j,
5405 scale,
5406 )?;
5407 }
5408 } else if let Some(cols) = &ckpt.cols[il] {
5409 let (c, s) = &cols[j - 1];
5410 e.copy_into(&mut rl.conv_state, 0, c, c.len())?;
5411 e.copy_into(&mut rl.ssm_state, 0, s, s.len())?;
5412 } else {
5413 return Err(
5414 "commit_verified_prefix: verify ckpt missing for linear layer".into(),
5415 );
5416 }
5417 }
5418 }
5419 cache.pos = snap.pos + j;
5420 Ok(())
5421 }
5422
5423 /// ROUND-STREAM: recur restore with device-j (the _dc twins; full accept early-exits
5424 /// in-kernel). Requires the batched-linear stash on every linear layer (stream gate).
5425 fn commit_verified_prefix_stream(
5426 &self,
5427 e: &Engine,
5428 cache: &mut Cache,
5429 snap: &crate::cache::CacheSnapshot,
5430 ckpt: &VerifyCkpt,
5431 acc: &CudaSlice<u32>,
5432 base: usize,
5433 t_v: usize,
5434 ) -> Result<(), Box<dyn std::error::Error>> {
5435 let cfg = &self.cfg;
5436 let ssm = cfg.ssm.as_ref().unwrap();
5437 let d_state = ssm.state_size as usize;
5438 let num_k = ssm.group_count as usize;
5439 let num_v = ssm.time_step_rank as usize;
5440 let d_conv = ssm.conv_kernel as usize;
5441 let conv_dim = d_state * num_k * 2 + d_state * num_v;
5442 let scale = 1.0 / (d_state as f32).sqrt();
5443 for il in 0..self.layers.len() {
5444 if let Some(rl) = cache.recur[il].as_mut() {
5445 let st = ckpt.gdn[il]
5446 .as_ref()
5447 .ok_or("stream restore: batched-linear stash missing")?;
5448 let ring_old = snap.conv[il].as_ref().expect("snapshot missing conv");
5449 let state_in = snap.ssm[il].as_ref().expect("snapshot missing ssm");
5450 e.ssm_conv_ring_rebuild_dc(
5451 &st.qkv_mixed,
5452 ring_old,
5453 &mut rl.conv_state,
5454 conv_dim,
5455 acc,
5456 base,
5457 t_v,
5458 d_conv,
5459 )?;
5460 let mut o = e.uninit(d_state * num_v * t_v)?;
5461 e.gdn_scan_s128_dc(
5462 &st.q_l2,
5463 &st.k_l2,
5464 &st.v_g,
5465 &st.g_log,
5466 &st.beta,
5467 state_in,
5468 &mut rl.ssm_state,
5469 &mut o,
5470 num_v,
5471 acc,
5472 base,
5473 t_v,
5474 scale,
5475 )?;
5476 }
5477 }
5478 Ok(())
5479 }
5480
5481 /// EAGLE3 aux-capturing verify forward over `tokens` (T) — mirrors `decode_step_t_h` exactly
5482 /// (same KV append, same causal verify, same recur advance) but ALSO clones the aux residual-
5483 /// stream hiddens (blocks in `aux_layers`) for TWO columns: the LAST column (always) and the
5484 /// optional `pred_col` (the EAGLE seed = bonus's predecessor). Returns
5485 /// (all_T_logits host, last_col_aux, pred_col_aux?). Used by the EAGLE3 orchestrator's commit.
5486 pub fn decode_step_t_aux2(
5487 &self,
5488 e: &Engine,
5489 tokens: &[u32],
5490 pos0: usize,
5491 cache: &mut Cache,
5492 aux_layers: &[usize],
5493 pred_col: Option<usize>,
5494 ) -> Result<
5495 (Vec<f32>, Vec<CudaSlice<f32>>, Option<Vec<CudaSlice<f32>>>),
5496 Box<dyn std::error::Error>,
5497 > {
5498 let cfg = &self.cfg;
5499 let n_embd = cfg.n_embd as usize;
5500 let eps = cfg.rms_eps;
5501 let t = tokens.len();
5502 let pos_vec: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
5503 let pos_d = e.htod_i32(&pos_vec)?;
5504 let mut x = e.htod(&self.embd.gather(n_embd, tokens))?;
5505 let mut aux_last: Vec<CudaSlice<f32>> = Vec::with_capacity(aux_layers.len());
5506 let mut aux_pred: Vec<CudaSlice<f32>> = Vec::new();
5507 let want_pred = pred_col.is_some();
5508
5509 for (il, layer) in self.layers.iter().enumerate() {
5510 // DISPATCH-MIRRORED norms (FP-order lesson #8) — see decode_step_t_h_emb.
5511 let mixer_fast = self.mixer_in_q8_1_fast(e, &layer.mixer);
5512 let norm_fused = std::env::var("MEMRA_NO_FUSE_NORMQ").is_err() && mixer_fast;
5513 let mut h = vbuf(e, t * n_embd)?; // fully written by either rms_norm arm
5514 if norm_fused {
5515 e.rms_norm_decode(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
5516 } else {
5517 e.rms_norm(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
5518 }
5519 let mixed = match &layer.mixer {
5520 Mixer::Full(fa) => {
5521 self.full_attn_verify(e, fa, &h, None, &pos_d, t, cache, il, None)?
5522 }
5523 Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
5524 Mixer::Linear(la) => {
5525 let mut out = e.zeros(t * n_embd)?;
5526 for col in 0..t {
5527 let mut h_col = e.zeros(n_embd)?;
5528 let src = h.slice(col * n_embd..(col + 1) * n_embd);
5529 e.copy_view_into(&mut h_col, 0, &src, n_embd)?;
5530 let m_col = self.linear_attn_decode(e, la, &h_col, cache, il)?;
5531 e.copy_into(&mut out, col * n_embd, &m_col, n_embd)?;
5532 }
5533 out
5534 }
5535 };
5536 let ffn_fuse = match &layer.ffn {
5537 crate::hybrid::Ffn::Dense {
5538 ffn_gate, ffn_up, ..
5539 } => {
5540 std::env::var("MEMRA_NO_FUSE_NORMQ").is_err()
5541 && e.uses_q8_1_fast(ffn_gate)
5542 && e.uses_q8_1_fast(ffn_up)
5543 }
5544 crate::hybrid::Ffn::Moe(_) => false,
5545 };
5546 let mut x1 = vbuf(e, t * n_embd)?; // fully written by add / add_rms_norm
5547 let mut z = vbuf(e, t * n_embd)?; // fully written by rms_norm_decode / add_rms_norm
5548 if ffn_fuse {
5549 e.add(&x, &mixed, &mut x1, t * n_embd)?;
5550 e.rms_norm_decode(
5551 &x1,
5552 layer.post_attn_norm.float_data(),
5553 &mut z,
5554 n_embd,
5555 t,
5556 eps,
5557 )?;
5558 } else {
5559 e.add_rms_norm(
5560 &x,
5561 &mixed,
5562 layer.post_attn_norm.float_data(),
5563 &mut x1,
5564 &mut z,
5565 n_embd,
5566 t,
5567 eps,
5568 )?;
5569 }
5570 let ffn_out = match &layer.ffn {
5571 crate::hybrid::Ffn::Dense {
5572 ffn_gate,
5573 ffn_up,
5574 ffn_down,
5575 } => {
5576 let n_ff = ffn_gate.out_features();
5577 let gate = e.matmul_decode_exact(ffn_gate, &z, t)?;
5578 let up = e.matmul_decode_exact(ffn_up, &z, t)?;
5579 let mut act = vbuf(e, t * n_ff)?; // fully written by ffn_act_lim
5580 // dense FFN clamp = the SHEXP array (upstream build_ffn serves both).
5581 Self::ffn_act_lim(
5582 e,
5583 &self.cfg,
5584 &gate,
5585 &up,
5586 1.0,
5587 1.0,
5588 self.cfg.clamp_shexp_at(il as u32),
5589 &mut act,
5590 t * n_ff,
5591 )?;
5592 e.matmul_decode_exact(ffn_down, &act, t)?
5593 }
5594 crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il(e, m, &z, t, il as u16)?,
5595 };
5596 let mut x2 = vbuf(e, t * n_embd)?; // fully written by add
5597 e.add(&x1, &ffn_out, &mut x2, t * n_embd)?;
5598 if aux_layers.contains(&il) {
5599 let mut a = e.zeros(n_embd)?;
5600 e.copy_view_into(&mut a, 0, &x2.slice((t - 1) * n_embd..t * n_embd), n_embd)?;
5601 aux_last.push(a);
5602 if let Some(pc) = pred_col {
5603 let mut ap = e.zeros(n_embd)?;
5604 e.copy_view_into(
5605 &mut ap,
5606 0,
5607 &x2.slice(pc * n_embd..(pc + 1) * n_embd),
5608 n_embd,
5609 )?;
5610 aux_pred.push(ap);
5611 }
5612 }
5613 x = x2;
5614 }
5615 let mut hn = vbuf(e, t * n_embd)?; // fully written by rms_norm_decode
5616 e.rms_norm_decode(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
5617 let logits = e.matmul_decode_exact(&self.output, &hn, t)?;
5618 let host = e.dtoh(&logits)?;
5619 cache.pos += t;
5620 Ok((
5621 host,
5622 aux_last,
5623 if want_pred { Some(aux_pred) } else { None },
5624 ))
5625 }
5626
5627 /// step35 SPEC-VERIFY attention over T query tokens — a per-row REPLAY of the eager
5628 /// `step35_decode_attn`.
5629 ///
5630 /// WHY A REPLAY AND NOT A BATCHED TWIN. The verify's whole job is to be bit-identical to what
5631 /// the eager decode would have computed for the same tokens; that is what makes greedy spec
5632 /// decode exact (run-spec asserts token identity for K=1..8). Every other verify arm in this
5633 /// file earns that identity by carefully mirroring dispatch (`matmul_decode_exact` to force
5634 /// MMVQ at any m, per-layer `ffn_fuse` mirroring, per-row `fa_decode` key bounds). step35
5635 /// stacks FOUR more per-layer degrees of freedom on top of that — per-layer `n_head`
5636 /// (64 full / 96 SWA), per-layer rotary width (64 full / 128 SWA), per-layer rope base, and a
5637 /// SEPARATE `attn_gate` tensor whose projection shares the attn-normed input — and its SWA
5638 /// layers attend through a token-OFFSET view whose offset is a function of the ABSOLUTE
5639 /// position of each query row. A batched twin would have to reproduce all of that AND the
5640 /// per-row offset in one launch; the offset alone rules out the existing rows kernels (they
5641 /// take one `base_len`, not a per-row offset).
5642 ///
5643 /// So this arm calls the eager path itself, once per row, on the same cache. Identity is then
5644 /// true BY CONSTRUCTION rather than by mirroring: row r runs exactly the kernel sequence that
5645 /// eager decode step r runs (same projections, same q8_1 fusion decision, same append, same
5646 /// view arithmetic, same `fa_decode_kvmod`, same gate), because it IS that code. Cost: T x the
5647 /// eager decode mixer instead of one batched pass — the same trade the generic arm's `else`
5648 /// per-row loop already accepts when `fa_rows_eligible` says no. Correctness first; a batched
5649 /// step35 twin is a perf lane's job and must be gated against this arm.
5650 ///
5651 /// The `h_q8` pre-quantized pair from the caller's fused norm is NOT forwarded: it is a
5652 /// T-row buffer and `step35_decode_attn`'s `pre_q` contract is one row. Instead each row's
5653 /// f32 `h` slice is handed over and the callee re-derives its own q8_1 exactly as eager decode
5654 /// does (`quantize_q8_1(h, 1, n_embd)`) — which is the dispatch being mirrored. Callers that
5655 /// took the fused arm therefore MUST still pass a live `h`; `step35_verify` asserts that.
5656 #[allow(clippy::too_many_arguments)]
5657 fn step35_verify(
5658 &self,
5659 e: &Engine,
5660 fa: &FullAttnLayer,
5661 h: &CudaSlice<f32>,
5662 h_q8: Option<(&CudaSlice<i8>, &CudaSlice<f32>)>,
5663 t: usize,
5664 cache: &mut Cache,
5665 il: usize,
5666 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5667 let n_embd = self.cfg.n_embd as usize;
5668 // The fused (h-less) attn-norm arm hands `h` as a zero-length placeholder. This arm needs
5669 // the f32 rows, so the caller must not take that lever for step35 — enforced at the call
5670 // site by the `Mixer::Full(_) if self.cfg.step35.is_some() => false` arm of
5671 // `lin_q8_only` in `decode_step_t_core_stream`, and asserted here so a future caller
5672 // cannot regress it into silently reading an empty buffer.
5673 assert_eq!(
5674 h.len(),
5675 t * n_embd,
5676 "step35_verify needs the f32 attn-normed rows ([t*n_embd]); the caller took the \
5677 fused q8-only norm arm (h_q8={}) — step35 must stay on the unfused arm",
5678 h_q8.is_some()
5679 );
5680 // ROW WIDTH IS n_embd, NOT n_head*head_dim: `step35_decode_attn` returns the mixer output
5681 // AFTER `wo`, so a row is [n_embd] — the same contract the generic arm's
5682 // `matmul_decode_exact(&fa.wo, &attn_g, t)` return has. Sizing this buffer from the
5683 // per-layer head geometry (8192 on full-attn, 12288 on SWA) instead overran the row on the
5684 // FIRST copy and panicked inside `copy_into`'s `CudaView::slice` unwrap
5685 // (raw/mtp-bt-20260806T212127Z.log frames 12-13).
5686 let mut out = vbuf(e, t * n_embd)?; // each row fully written by the copy below
5687 for r in 0..t {
5688 // Absolute position of this query row. `cache.pos` is the committed length at round
5689 // start and every row before r has already been appended by this loop, so the r-th
5690 // verify token sits at cache.pos + r — the same position eager decode would give it.
5691 let pos_d = e.htod_i32(&[(cache.pos + r) as i32])?;
5692 let mut h_row = vbuf(e, n_embd)?; // fully written by copy_view_into
5693 e.copy_view_into(
5694 &mut h_row,
5695 0,
5696 &h.slice(r * n_embd..(r + 1) * n_embd),
5697 n_embd,
5698 )?;
5699 // THE eager decode mixer: appends this row's K/V at kvl.len, advances it, then
5700 // attends over the (SWA-offset) view. Post-`wo`, same contract as this fn returns.
5701 let o = self.step35_decode_attn(e, fa, il, &h_row, None, &pos_d, cache)?;
5702 debug_assert_eq!(
5703 o.len(),
5704 n_embd,
5705 "step35_decode_attn returns post-wo [n_embd]"
5706 );
5707 e.copy_into(&mut out, r * n_embd, &o, n_embd)?;
5708 }
5709 Ok(out)
5710 }
5711
5712 /// Full-attention mixer over T query tokens with a GROWING resident KV (verify path, §D.3).
5713 /// Appends the T new K/V columns to cache.kv[il] then attends causally over [0..len) via
5714 /// fa_prefill. Token-major [T, kv_dim] projection layout == cache row layout (single copy).
5715 #[allow(clippy::too_many_arguments)]
5716 fn full_attn_verify(
5717 &self,
5718 e: &Engine,
5719 fa: &FullAttnLayer,
5720 h: &CudaSlice<f32>,
5721 h_q8: Option<(&CudaSlice<i8>, &CudaSlice<f32>)>,
5722 pos_d: &CudaSlice<i32>,
5723 t: usize,
5724 cache: &mut Cache,
5725 il: usize,
5726 stream_ctr: Option<&CudaSlice<i32>>,
5727 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5728 // step35: the generic geometry below is wrong for this arch (per-layer n_head, partial
5729 // per-layer rope, the SWA offset view, and a SEPARATE head-wise gate tensor), so it takes
5730 // its own arm. A verify that silently computes different attention than decode defeats the
5731 // whole self-consistency gate, so the arm is a per-row REPLAY of `step35_decode_attn`
5732 // rather than a batched twin — see `step35_verify` for why that is the exactness-correct
5733 // shape and not laziness.
5734 if self.cfg.step35.is_some() {
5735 if stream_ctr.is_some() {
5736 return Err(
5737 "step35 has no ROUND-STREAM verify arm (the device-counter _dc twins \
5738 cannot express the SWA offset KV view; same root cause as the dc \
5739 decode refusal) — run spec without the stream arm"
5740 .into(),
5741 );
5742 }
5743 return self.step35_verify(e, fa, h, h_q8, t, cache, il);
5744 }
5745 let cfg = &self.cfg;
5746 let geometry = cfg.full_attention_geometry_at(il as u32);
5747 let n_head = geometry.n_head as usize;
5748 let n_head_kv = geometry.n_head_kv as usize;
5749 let head_dim = geometry.head_dim_k as usize;
5750 let eps = cfg.rms_eps;
5751 let scale = geometry.attention_scale();
5752 let n_embd = cfg.n_embd as usize;
5753
5754 // DECODE-EXACT Q/K/V projections: matmul_decode_exact forces the MMVQ (warp-per-row) path
5755 // for every m, matching the T=1 decode's FP accumulation order. matmul_pre at m>=5 would
5756 // fall to dp4a (128-thread, two-level reduce) with a different FP sum order.
5757 // Q8 TRUNK-FUSION at T=1: DISPATCH-MIRRORS the eager decode's fused3 (bit-identical body).
5758 // BATCHED EPILOGUE RE-FUSE (lane/vt-fixes fix 2): `h_q8` = the attn-input norm's q8_1
5759 // form emitted by the fused rms_norm_q8_1 (bit-identical to rms_norm_decode ->
5760 // quantize_q8_1, kernel-check-pinned). When present (caller checked mixer q8_1-fast),
5761 // every projection consumes it — `h` may be a zero-len placeholder and must not be read.
5762 let (qf, mut k, v) = {
5763 let mut fused = None;
5764 let qkv_fast =
5765 e.uses_q8_1_fast(&fa.wq) && e.uses_q8_1_fast(&fa.wk) && e.uses_q8_1_fast(&fa.wv);
5766 if t == 1 && qkv_fast {
5767 let (hq_o, hd_o);
5768 let (hq, hd): (&CudaSlice<i8>, &CudaSlice<f32>) = match h_q8 {
5769 Some(p) => p,
5770 None => {
5771 (hq_o, hd_o) = e.quantize_q8_1(h, 1, n_embd)?;
5772 (&hq_o, &hd_o)
5773 }
5774 };
5775 fused = e.matmul_q8_fused3(&fa.wq, &fa.wk, &fa.wv, hq, hd)?;
5776 } else if spec_fused_t() && (2..=4).contains(&t) && qkv_fast {
5777 // VERIFY-TIER TRUNK FUSION (MEMRA_SPEC_FUSED_T): one shared quantize + one
5778 // fused3 batched launch replaces three decode-exact calls (3 re-quantizes of
5779 // the same h + 3 _b2/_b4 launches). Bit-identical per (tensor,token,row).
5780 let (hq_o, hd_o);
5781 let (hq, hd): (&CudaSlice<i8>, &CudaSlice<f32>) = match h_q8 {
5782 Some(p) => p,
5783 None => {
5784 (hq_o, hd_o) = e.quantize_q8_1(h, t, n_embd)?;
5785 (&hq_o, &hd_o)
5786 }
5787 };
5788 fused = e.matmul_q8_fused3_t(&fa.wq, &fa.wk, &fa.wv, hq, hd, t)?;
5789 }
5790 match (fused, h_q8) {
5791 (Some(triple), _) => triple,
5792 // shared pre-quantized activation (q8_1-fast guaranteed by the caller): the
5793 // decode-exact dispatch consumes (hq, hd) instead of re-quantizing 3x.
5794 (None, Some((hq, hd))) if qkv_fast => (
5795 e.matmul_decode_exact_pre(&fa.wq, hq, hd, t)?,
5796 e.matmul_decode_exact_pre(&fa.wk, hq, hd, t)?,
5797 e.matmul_decode_exact_pre(&fa.wv, hq, hd, t)?,
5798 ),
5799 (None, _) => (
5800 e.matmul_decode_exact(&fa.wq, h, t)?,
5801 e.matmul_decode_exact(&fa.wk, h, t)?,
5802 e.matmul_decode_exact(&fa.wv, h, t)?,
5803 ),
5804 }
5805 };
5806 // M3/Hy3 have no attention output gate — wq out is exactly q; skip the split.
5807 let gated = geometry.attention_gate == memra_gguf::config::AttentionGateKind::FusedQ;
5808 let (mut q, gate) = if gated {
5809 let mut q = vbuf(e, t * n_head * head_dim)?; // fully written by q_gate_split
5810 let mut gate = vbuf(e, t * n_head * head_dim)?; // fully written by q_gate_split
5811 e.q_gate_split(&qf, &mut q, &mut gate, head_dim, n_head, t)?;
5812 (q, Some(gate))
5813 } else {
5814 (qf, None)
5815 };
5816
5817 let mut qn = vbuf(e, t * n_head * head_dim)?; // fully written by rms_norm
5818 e.rms_norm(
5819 &q,
5820 fa.q_norm.float_data(),
5821 &mut qn,
5822 head_dim,
5823 n_head * t,
5824 eps,
5825 )?;
5826 q = qn;
5827 let mut kn = vbuf(e, t * n_head_kv * head_dim)?; // fully written by rms_norm
5828 e.rms_norm(
5829 &k,
5830 fa.k_norm.float_data(),
5831 &mut kn,
5832 head_dim,
5833 n_head_kv * t,
5834 eps,
5835 )?;
5836 k = kn;
5837 let rope_dims = geometry.n_rot as usize;
5838 e.rope_neox(
5839 &mut q,
5840 pos_d,
5841 head_dim,
5842 rope_dims,
5843 n_head,
5844 t,
5845 geometry.rope_base,
5846 1.0,
5847 )?;
5848 e.rope_neox(
5849 &mut k,
5850 pos_d,
5851 head_dim,
5852 rope_dims,
5853 n_head_kv,
5854 t,
5855 geometry.rope_base,
5856 1.0,
5857 )?;
5858
5859 // append T new K/V columns to the resident QUANTIZED cache. k/v are token-major [T, kv_dim]
5860 // f32; append-quantize each of the T token rows into the byte cache (q8_0 K / q5_1 V).
5861 let kvl = cache.kv[il].as_mut().unwrap();
5862 let (kv_dim_k, kv_dim_v, ktb, vtb) =
5863 (kvl.kv_dim_k, kvl.kv_dim_v, kvl.k_tok_bytes, kvl.v_tok_bytes);
5864 if let Some(ctr) = stream_ctr {
5865 // stream: ONE batched append at the device counter (rows kernel = the per-view warp
5866 // math on a (block, token) grid, documented byte-identical); host len is a stale
5867 // LOWER BOUND under pre-issue (drain reconciles it).
5868 e.append_kv_quantized_rows_dc(
5869 &k,
5870 &v,
5871 &mut kvl.k,
5872 &mut kvl.v,
5873 ctr,
5874 t,
5875 kv_dim_k,
5876 kv_dim_v,
5877 ktb,
5878 vtb,
5879 crate::Engine::kv_fp8_on(),
5880 )?;
5881 } else {
5882 for i in 0..t {
5883 let k_row = k.slice(i * kv_dim_k..(i + 1) * kv_dim_k);
5884 let v_row = v.slice(i * kv_dim_v..(i + 1) * kv_dim_v);
5885 e.append_kv_quantized_view(
5886 &k_row,
5887 &v_row,
5888 &mut kvl.k,
5889 &mut kvl.v,
5890 kvl.len + i,
5891 kv_dim_k,
5892 kv_dim_v,
5893 ktb,
5894 vtb,
5895 crate::Engine::kv_fp8_on(),
5896 )?;
5897 }
5898 kvl.len += t;
5899 }
5900
5901 // BIT-IDENTICAL VERIFY ATTENTION (spec-exactness fix): the FP accumulation order must be
5902 // byte-for-byte identical to the eager decode path. fa_prefill uses a different tile size
5903 // (BLOCK_Q=64, BK=32) and online-softmax structure than fa_decode's split-K + combine,
5904 // which changes FP summation order and can flip argmax at tight logit margins. Query row r
5905 // attends to keys [0..base_len+r+1) — each successive row sees one more key (the causal
5906 // property). This matches eager: decode appends k at len, then fa_decode sees t_kv = len+1
5907 // keys. The verify appends all T tokens first but bounds the key range per row.
5908 //
5909 // MULTI-ROW FUSED PATH (the long-ctx spec fix, 2026-07-03): when every row takes the vec
5910 // kernel (base_len+1 >= FA_VEC_MIN_TKV), ONE fa_decode_rows launch executes the exact
5911 // per-row program for all T rows (grid.z = row, per-row n_splits from the same
5912 // fa_split_keys formula) — replacing T x (2 launches + 2 dtod copies + 5 partial allocs)
5913 // and multiplying resident CTAs by T on a latency-bound kernel. Bit-identical per row by
5914 // construction; kernel-check pins rows-vs-loop byte identity, run-spec is the end gate.
5915 // Short ctx (any row below the vec crossover) and MEMRA_NO_FA_VEC/MEMRA_FA_ROWS_OFF keep the
5916 // per-row loop (whose fa_decode picks scalar/vec per row exactly like eager decode).
5917 let mut attn = vbuf(e, t * n_head * head_dim)?; // fully written by every FA arm below
5918 let base_len = kvl.len - t; // KV len BEFORE this round's T tokens were appended
5919 // T=1 INCLUDED (2026-07-05): p-min cuts the draft to 1 in ~75% of rounds on hard
5920 // (agentic) content — the old t>1 gate sent those rounds to the per-row loop (262us/row
5921 // + q-row copy + per-row allocs vs 93us/row through the fused kernel at grid.z=1, same
5922 // program). nsys accounting: 1088 of 1456 verify FA launches were T=1 escapees.
5923 // LEAN T=1 ARM (MEMRA_SPEC_LEAN, close35): at t==1, q IS one row and fa_decode on it is
5924 // the EXACT eager decode dispatch (vec_q_v2 + combine_f32; the rows pair measured +50us
5925 // at m=1). Byte-identical: kernel-check pins rows-vs-loop identity, and the per-row loop
5926 // at t=1 is fa_decode on the same q with zero-offset copies. Gates arbitrate.
5927 if let Some(ctr) = stream_ctr {
5928 // STREAM ARM: causal base from the device counter; host kvl.len is a stale lower
5929 // bound used only for the split-sizing upper bound (+64 slack covers M pre-issued
5930 // rounds at K<=8). Views span the bound; per-row limits derive in-kernel.
5931 let upper = kvl.len + t + 64;
5932 let k_view = e.view_u8(&kvl.k, (upper.min(cache.max_ctx)) * ktb);
5933 let v_view = e.view_u8(&kvl.v, (upper.min(cache.max_ctx)) * vtb);
5934 e.fa_decode_rows_dc(
5935 &q,
5936 &k_view,
5937 &v_view,
5938 &mut attn,
5939 head_dim,
5940 n_head,
5941 n_head_kv,
5942 ctr,
5943 upper.min(cache.max_ctx),
5944 t,
5945 scale,
5946 ktb,
5947 vtb,
5948 0,
5949 false,
5950 )?;
5951 } else if spec_lean() && t == 1 {
5952 let t_kv = base_len + 1;
5953 let k_view = e.view_u8(&kvl.k, t_kv * ktb);
5954 let v_view = e.view_u8(&kvl.v, t_kv * vtb);
5955 e.fa_decode_kvmod(
5956 &q,
5957 &k_view,
5958 &v_view,
5959 &mut attn,
5960 head_dim,
5961 n_head,
5962 n_head_kv,
5963 t_kv,
5964 scale,
5965 ktb,
5966 vtb,
5967 crate::Engine::kv_fp8_on(),
5968 )?;
5969 } else if e.fa_rows_eligible(base_len, head_dim) {
5970 let k_view = e.view_u8(&kvl.k, (base_len + t) * ktb);
5971 let v_view = e.view_u8(&kvl.v, (base_len + t) * vtb);
5972 e.fa_decode_rows(
5973 &q,
5974 &k_view,
5975 &v_view,
5976 &mut attn,
5977 head_dim,
5978 n_head,
5979 n_head_kv,
5980 base_len,
5981 t,
5982 scale,
5983 ktb,
5984 vtb,
5985 None,
5986 false,
5987 crate::Engine::kv_fp8_on(),
5988 None,
5989 )?;
5990 } else {
5991 for r in 0..t {
5992 let t_kv_r = base_len + r + 1; // this row sees keys [0..t_kv_r)
5993 let k_view_r = e.view_u8(&kvl.k, t_kv_r * ktb);
5994 let v_view_r = e.view_u8(&kvl.v, t_kv_r * vtb);
5995 // copy q row into an owned buffer (fa_decode takes &CudaSlice, not CudaView)
5996 let mut q_row = vbuf(e, n_head * head_dim)?; // fully written by copy_view_into
5997 let q_src = q.slice(r * n_head * head_dim..(r + 1) * n_head * head_dim);
5998 e.copy_view_into(&mut q_row, 0, &q_src, n_head * head_dim)?;
5999 let mut attn_row = vbuf(e, n_head * head_dim)?; // fully written by fa_decode
6000 e.fa_decode_kvmod(
6001 &q_row,
6002 &k_view_r,
6003 &v_view_r,
6004 &mut attn_row,
6005 head_dim,
6006 n_head,
6007 n_head_kv,
6008 t_kv_r,
6009 scale,
6010 ktb,
6011 vtb,
6012 crate::Engine::kv_fp8_on(),
6013 )?;
6014 e.copy_into(
6015 &mut attn,
6016 r * n_head * head_dim,
6017 &attn_row,
6018 n_head * head_dim,
6019 )?;
6020 }
6021 }
6022
6023 let attn_g = match &gate {
6024 Some(gate) => {
6025 let mut gsig = vbuf(e, t * n_head * head_dim)?; // fully written by sigmoid
6026 e.sigmoid(gate, &mut gsig, t * n_head * head_dim)?;
6027 let mut ag = vbuf(e, t * n_head * head_dim)?; // fully written by mul
6028 e.mul(&attn, &gsig, &mut ag, t * n_head * head_dim)?;
6029 ag
6030 }
6031 None => attn,
6032 };
6033 // DECODE-EXACT wo projection: at m>=5 (K=4+ with pending) the generic matmul would use dp4a
6034 // (128-thread, different FP sum order than MMVQ). Force MMVQ for bit-identity with decode.
6035 Ok(e.matmul_decode_exact(&fa.wo, &attn_g, t)?)
6036 }
6037
6038 /// Context-linear bytes for a plain serving session's trunk cache.
6039 pub fn plain_session_kv_bytes_per_token(&self) -> usize {
6040 crate::cache::cache_bytes_per_token(&self.cfg)
6041 }
6042
6043 /// `(logical bytes/token, ring-capped bytes/token, ring row cap)` for exact admission.
6044 pub fn plain_session_kv_shape(&self) -> (usize, usize, usize) {
6045 (
6046 self.plain_session_kv_bytes_per_token(),
6047 crate::cache::cache_ring_bytes_per_token(&self.cfg),
6048 crate::cache::cache_ring_row_cap(&self.cfg),
6049 )
6050 }
6051
6052 /// Context-linear bytes for a speculative serving session: trunk cache plus persistent MTP
6053 /// scratch. With no MTP head this equals the plain coefficient.
6054 pub fn spec_session_kv_bytes_per_token(&self) -> usize {
6055 let scratch = self
6056 .mtp
6057 .as_ref()
6058 .map(|mtp| {
6059 let (_, _, k, v) = mtp_scratch_layout(&self.cfg, mtp.geom.as_ref());
6060 k + v
6061 })
6062 .unwrap_or(0);
6063 self.plain_session_kv_bytes_per_token()
6064 .saturating_add(scratch)
6065 }
6066
6067 /// Spec twin of [`HybridModel::plain_session_kv_shape`]; Step35's persistent MTP scratch is
6068 /// capped by the same SWA ring rows as the trunk.
6069 pub fn spec_session_kv_shape(&self) -> (usize, usize, usize) {
6070 let total = self.spec_session_kv_bytes_per_token();
6071 let (_, mut ring, rows) = self.plain_session_kv_shape();
6072 if rows > 0 {
6073 ring = ring.saturating_add(
6074 self.mtp
6075 .as_ref()
6076 .map(|mtp| {
6077 let (_, _, k, v) = mtp_scratch_layout(&self.cfg, mtp.geom.as_ref());
6078 k + v
6079 })
6080 .unwrap_or(0),
6081 );
6082 }
6083 (total, ring, rows)
6084 }
6085
6086 /// Greedy MTP speculative decode (§B). Token-identical to `generate(prompt, max_new)` but uses
6087 /// the NextN head to draft K tokens then verifies them in one batched target forward.
6088 /// Returns (generated tokens, total_drafted, total_accepted) so the caller can report
6089 /// acceptance rate. `k` = draft length per round.
6090 ///
6091 /// GRAPH DRAFT (stage 2 of graph-grade spec): when the model is all-Dense and the MTP head is
6092 /// Dense (no MoE host readbacks), the fixed-shape T=1 MTP forward is CUDA-graph-captured ONCE
6093 /// and replayed per draft step — the ~40 eager launches per drafted token collapse into one
6094 /// graph dispatch; only the 4-byte token id (and 4-byte p-min confidence) round-trip per step.
6095 /// Event tracking is disabled for the whole call (generate_graph pattern) so every buffer the
6096 /// captured graph references is event-free; the spec loop is strictly single-stream.
6097 /// MEMRA_SPEC_NOGRAPH=1 forces the eager draft chain.
6098 /// SAMPLED mode (MEMRA_SPEC_TEMP>0) has its OWN capture (gumbel-perturbed in-graph argmax,
6099 /// device Philox event counter, persistent q retention) — graph-vs-eager sampled streams are
6100 /// bit-identical for the same (seed, prompt, K, temp); see the sampled-graph setup in
6101 /// generate_spec_inner2.
6102 /// Multi-turn session: trunk cache + MTP draft scratch persist across generate calls, so
6103 /// turn N+1 primes ONLY its new suffix (the 124k-conversation daily pattern — re-priming a
6104 /// 32k history costs ~54s; a suffix prime costs seconds). APPEND-ONLY by construction: the
6105 /// hybrid linear-attn states are in-place (no position index), so a session can extend but
6106 /// never rewind — `committed` is the exact token list whose state the caches hold (includes
6107 /// any overshoot tokens past max_new; the caller renders from `committed`, not its own echo).
6108 pub fn new_session(
6109 &self,
6110 e: &Engine,
6111 max_ctx: usize,
6112 ) -> Result<SpecSession, Box<dyn std::error::Error>> {
6113 Ok(SpecSession {
6114 // STAGE-OWNED KV (lane/pp2-spec 2026-08-06): `pp::new_cache`, not `Cache::new`. This
6115 // is the SERVING spec-session path, and with the ppN door open across two cards a
6116 // primary-homed cache makes every remote stage peer-read its OWN KV on every verify
6117 // round — the wrong-card class already fixed on the two batched serving paths
6118 // (worker.rs 2483 / 2837). With the door shut `new_cache` IS `Cache::new` (same
6119 // branch, same allocations), so single-device behavior is byte-unchanged.
6120 cache: crate::pp::new_cache(e, &self.cfg, max_ctx)?,
6121 scratch: MtpScratch::new(
6122 e,
6123 &self.cfg,
6124 max_ctx,
6125 self.mtp.as_ref().and_then(|m| m.geom.as_ref()),
6126 )?,
6127 committed: Vec::new(),
6128 last_h: None,
6129 next_pred: None,
6130 sctr: 0,
6131 uctr: 0,
6132 draft_ctx: None,
6133 pending_tok: None,
6134 turn_ckpt: None,
6135 telem: SpecTelemetryCounters::default(),
6136 capture_at: None,
6137 boundary_capture: None,
6138 })
6139 }
6140
6141 /// SPEC-ON-CACHE-HIT restore (lane/spec-on-cache-hit, 2026-08-18 — PORT-PLAN item 3,
6142 /// research/cache-spec-design-20260814, scoped to WHOLE-ENTRY restores only): build a
6143 /// SpecSession around a trunk cache the worker already restored from a prefix-cache
6144 /// entry, re-installing the entry's published draft plane as the MTP scratch rows
6145 /// `[0..prefix.len())` and the entry's boundary hidden as `last_h`, then feeding the
6146 /// prompt SUFFIX here — through EXACTLY the plain path's program selection — so the
6147 /// worker always receives a fully-warm continuation session (committed = whole
6148 /// prompt, `next_pred` + `last_h` set; caller sets `next_pred` from the entry's
6149 /// boundary logits on the empty-suffix shape).
6150 ///
6151 /// PROGRAM LAW (the splitiso two-programs class, learned AGAIN in this lane's own
6152 /// gate): the identity target for a converted hit is the PLAIN hit serving the same
6153 /// request, and plain feeds a carried suffix via eager `decode_step` below
6154 /// PRIME_MIN_T and via `prime_cache` at/above it (prefill_tick's arms). The generate
6155 /// path's tokenwise arm routes qwen35-class through the BATCHED T=1 program
6156 /// (`spec_target_step_h`) instead — ULP-different suffix rows, and the gate measured
6157 /// the near-tie flip at generated token ~8 (research/spec-cache-20260818, qwen r3).
6158 /// So the suffix is fed HERE, mirroring prefill_tick arm-for-arm, not handed to the
6159 /// burst prime.
6160 ///
6161 /// SEED RULE (both sampling regimes; lane/sampled-hit-spec 2026-08-19, sampled draw
6162 /// added by lane/sampled-spec-quality 2026-08-19). The boundary token is produced by
6163 /// EXACTLY the rule the cold burst entry applies to its own first token from the same
6164 /// logits row: `argmax` when greedy, and a `sample_boundary_token` draw at Philox
6165 /// counter 0 when sampled. Both shapes are covered — the entry's boundary logits on a
6166 /// full-cover (empty-suffix) hit, this feed's own boundary logits on a suffix hit.
6167 /// That is what keeps a restored session seed-identical to a cold one PER SEED: the
6168 /// cold session draws from the identical row at counter 0 and then runs its rounds from
6169 /// counter 1, so the restored session admits with `sctr = 1` after its own draw.
6170 /// The WORKER owns the one refusal this constructor cannot see — a constrained request.
6171 /// (The penalized-sampled refusal was LIFTED once the burst's penalty window learned to
6172 /// span the session: `committed` here is the WHOLE prompt, so the restored session's
6173 /// window is the cold session's window. It comes back if `MEMRA_SPEC_PEN_SESSION=0`.)
6174 ///
6175 /// NOT the rolled-back partial-restore hazard: the caller restores at exactly the
6176 /// entry's captured endpoint (`e.pos`) through the shipping whole-entry path;
6177 /// mid-entry (`at < e.pos`) trunk restores stay behind MEMRA_PREFIX_PARTIAL_RESTORE
6178 /// and are never routed here.
6179 ///
6180 /// Failure contract: `Err((Some(cache), why))` before any trunk mutation — the
6181 /// worker rebuilds the plain carrier and the hit serves plain, byte-unchanged.
6182 /// `Err((None, why))` after the suffix feed began — the carrier is part-fed and
6183 /// UNUSABLE; the worker serves the request cold-plain (correct, slower) and the
6184 /// entry stays published for the next request.
6185 #[allow(clippy::too_many_arguments)]
6186 pub fn spec_session_from_restored(
6187 &self,
6188 e: &Engine,
6189 mut cache: Cache,
6190 prefix: Vec<u32>,
6191 suffix: &[u32],
6192 draft_k: &CudaSlice<u8>,
6193 draft_v: &CudaSlice<u8>,
6194 draft_k_tok_bytes: usize,
6195 draft_v_tok_bytes: usize,
6196 draft_len: usize,
6197 last_h: &[f32],
6198 // The ENTRY's boundary logits row (the full-cover shape's seed source). May be empty
6199 // when a suffix follows — the feed's own logits are the boundary then.
6200 boundary_logits: &[f32],
6201 // The request's sampler, or None for greedy. Owned here so the seed rule lives in
6202 // ONE place instead of being half-applied by the worker.
6203 sampling: Option<SpecSampling>,
6204 require_anchor: bool,
6205 max_ctx: usize,
6206 ) -> Result<SpecSession, (Option<Cache>, String)> {
6207 let pos = prefix.len();
6208 let fail = |cache: Cache, msg: String| -> Result<SpecSession, (Option<Cache>, String)> {
6209 Err((Some(cache), msg))
6210 };
6211 if self.mtp.is_none() {
6212 return fail(cache, "no MTP head attached (nothing to draft with)".into());
6213 }
6214 if pos == 0 {
6215 return fail(cache, "empty committed prefix".into());
6216 }
6217 if cache.pos != pos {
6218 let msg = format!(
6219 "restored cache pos {} != restored prefix len {pos}",
6220 cache.pos
6221 );
6222 return fail(cache, msg);
6223 }
6224 if draft_len != pos {
6225 return fail(
6226 cache,
6227 format!("draft plane len {draft_len} != restored prefix len {pos}"),
6228 );
6229 }
6230 if pos + suffix.len() >= max_ctx {
6231 return fail(
6232 cache,
6233 format!(
6234 "prompt {} + suffix would not leave generation room in ctx {max_ctx}",
6235 pos + suffix.len(),
6236 ),
6237 );
6238 }
6239 let mut scratch = match MtpScratch::new(
6240 e,
6241 &self.cfg,
6242 max_ctx,
6243 self.mtp.as_ref().and_then(|m| m.geom.as_ref()),
6244 ) {
6245 Ok(s) => s,
6246 Err(err) => return fail(cache, format!("draft scratch alloc failed: {err}")),
6247 };
6248 if scratch.kv.ring.is_some() {
6249 return fail(
6250 cache,
6251 "ring-backed draft scratch (Step35 SWA) cannot take a flat prefix restore".into(),
6252 );
6253 }
6254 if scratch.kv.k_tok_bytes != draft_k_tok_bytes
6255 || scratch.kv.v_tok_bytes != draft_v_tok_bytes
6256 {
6257 return fail(
6258 cache,
6259 format!(
6260 "draft plane layout {draft_k_tok_bytes}/{draft_v_tok_bytes} != scratch \
6261 {}/{} bytes/token (stale entry across a format change)",
6262 scratch.kv.k_tok_bytes, scratch.kv.v_tok_bytes,
6263 ),
6264 );
6265 }
6266 if pos > scratch.cap {
6267 return fail(
6268 cache,
6269 format!(
6270 "draft plane rows {pos} exceed scratch capacity {}",
6271 scratch.cap
6272 ),
6273 );
6274 }
6275 let kb = pos * draft_k_tok_bytes;
6276 let vb = pos * draft_v_tok_bytes;
6277 if draft_k.len() < kb || draft_v.len() < vb {
6278 return fail(
6279 cache,
6280 format!(
6281 "truncated draft plane: K {} < {kb} or V {} < {vb} bytes",
6282 draft_k.len(),
6283 draft_v.len(),
6284 ),
6285 );
6286 }
6287 if kb > 0 {
6288 if let Err(err) = e.copy_u8_into(&mut scratch.kv.k, 0, draft_k, kb) {
6289 return fail(cache, format!("draft K restore copy failed: {err}"));
6290 }
6291 }
6292 if vb > 0 {
6293 if let Err(err) = e.copy_u8_into(&mut scratch.kv.v, 0, draft_v, vb) {
6294 return fail(cache, format!("draft V restore copy failed: {err}"));
6295 }
6296 }
6297 if let Err(err) = scratch.set_len(e, pos) {
6298 return fail(cache, format!("draft scratch len set failed: {err}"));
6299 }
6300 let mut last_h_dev = if last_h.len() == self.cfg.n_embd as usize {
6301 // anchor upload failure is acceptance-only when a suffix feed follows (fill
6302 // row-0 falls back to zeros) but FATAL for an empty-suffix continuation (the
6303 // burst entry asserts committed + last_h + next_pred) — the caller says which.
6304 e.htod(last_h).ok()
6305 } else {
6306 None
6307 };
6308 if require_anchor && last_h_dev.is_none() {
6309 return fail(
6310 cache,
6311 "empty-suffix continuation requires the entry's boundary hidden anchor".into(),
6312 );
6313 }
6314 let mut committed = prefix;
6315 // Set on BOTH shapes below (suffix-fed and full-cover) — never left None, which is
6316 // what the empty-suffix continuation assert in the burst entry requires.
6317 let next_pred;
6318 // Philox: (0,0) at admit exactly like a fresh session; a sampled boundary draw below
6319 // consumes counter 0 and leaves 1, which is the state a cold session reaches after
6320 // drawing its own first token from the same row.
6321 let mut sctr = 0u32;
6322 let sampled = sampling.is_some_and(|s| s.temp > 0.0) && spec_sampled_boundary_on();
6323 // Penalty window for the boundary draw: the last `penalty_last_n` tokens of the WHOLE
6324 // prompt, which is what the cold session's own burst sees (Item 2's window). Built
6325 // after the suffix joins `committed` below.
6326 let mut boundary_capture: Option<SpecBoundaryCapture> = None;
6327 if !suffix.is_empty() {
6328 // ---- SUFFIX FEED, mirroring prefill_tick's program selection exactly ----
6329 // From here on the trunk cache mutates: failures return Err((None, _)) and
6330 // the worker serves the request cold-plain instead of reusing the carrier.
6331 let dirty =
6332 |msg: String| -> Result<SpecSession, (Option<Cache>, String)> { Err((None, msg)) };
6333 let n_embd = self.cfg.n_embd as usize;
6334 let t = suffix.len();
6335 let mut h_rows = match e.uninit(t * n_embd) {
6336 Ok(b) => b,
6337 Err(err) => return fail(cache, format!("suffix hidden buffer alloc: {err}")),
6338 };
6339 let mut feed_logits = Vec::new();
6340 let batched = t >= crate::hybrid_forward::PRIME_MIN_T
6341 && std::env::var("MEMRA_PRIME_TOKENWISE").is_err()
6342 && !e.frozen_cpu_experts_prefer_tokenwise_prime();
6343 if batched {
6344 // prefill_tick's prime arm: one request-level prime_cache call.
6345 match self.prime_cache(e, suffix, &mut cache, 0) {
6346 Ok((l, _h_seed, hiddens)) => {
6347 if let Err(err) = e.copy_into(&mut h_rows, 0, &hiddens, t * n_embd) {
6348 return dirty(format!("suffix hidden copy: {err}"));
6349 }
6350 feed_logits = l;
6351 }
6352 Err(err) => return dirty(format!("suffix prime failed: {err}")),
6353 }
6354 } else {
6355 // prefill_tick's tokenwise arm: eager decode_step, one token at a time.
6356 for (i, &tok) in suffix.iter().enumerate() {
6357 match self.decode_step_h(e, tok, &mut cache) {
6358 Ok((l, h)) => {
6359 if let Err(err) = e.copy_into(&mut h_rows, i * n_embd, &h, n_embd) {
6360 return dirty(format!("suffix hidden copy: {err}"));
6361 }
6362 feed_logits = l;
6363 }
6364 Err(err) => return dirty(format!("suffix decode_step failed: {err}")),
6365 }
6366 }
6367 }
6368 // Draft-scratch fill for the suffix rows, predecessor-paired: row `pos` reads
6369 // the entry's boundary anchor (zeros fallback — acceptance-only), row `pos+i`
6370 // reads h_rows[i-1]. Chunked like the generate path's fill (transients scale
6371 // with T). Fill failures are acceptance-only — truncate to the restored rows
6372 // and continue; the burst's own set_len keeps the invariant.
6373 let mtp = self.mtp.as_ref().expect("mtp checked above");
6374 let (embd_qt, embd_rb) = self.embd.qt_and_row_bytes(n_embd);
6375 let embd_gpu = if spec_host_embd() {
6376 None
6377 } else {
6378 Some(
6379 self.embd_gpu
6380 .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload")),
6381 )
6382 };
6383 let embd_dev = embd_gpu.map(|g| (g, embd_qt, embd_rb));
6384 let fill_chunk = 4096usize;
6385 let mut filled = true;
6386 let mut start = 0usize;
6387 'fill: while start < t {
6388 let end = (start + fill_chunk).min(t);
6389 let tc = end - start;
6390 let Ok(mut phs) = e.zeros(tc * n_embd) else {
6391 filled = false;
6392 break 'fill;
6393 };
6394 let (src_lo, dst_off, n_copy) = if start == 0 {
6395 (0, n_embd, (tc - 1) * n_embd)
6396 } else {
6397 ((start - 1) * n_embd, 0, tc * n_embd)
6398 };
6399 if start == 0 {
6400 if let Some(lh) = last_h_dev.as_ref() {
6401 if e.copy_into(&mut phs, 0, lh, n_embd).is_err() {
6402 filled = false;
6403 break 'fill;
6404 }
6405 }
6406 }
6407 if n_copy > 0
6408 && e.copy_view_into(
6409 &mut phs,
6410 dst_off,
6411 &h_rows.slice(src_lo..src_lo + n_copy),
6412 n_copy,
6413 )
6414 .is_err()
6415 {
6416 filled = false;
6417 break 'fill;
6418 }
6419 if self
6420 .mtp_kv_fill(
6421 e,
6422 mtp,
6423 &suffix[start..end],
6424 &phs,
6425 pos + start,
6426 &mut scratch,
6427 embd_dev,
6428 )
6429 .is_err()
6430 {
6431 filled = false;
6432 break 'fill;
6433 }
6434 start = end;
6435 }
6436 if !filled {
6437 // acceptance-only: drafts over missing suffix rows are cheap and wrong,
6438 // so keep only the restored rows resident and let verify arbitrate.
6439 if let Err(err) = scratch.set_len(e, pos) {
6440 return dirty(format!("scratch truncation after failed fill: {err}"));
6441 }
6442 }
6443 // EXTENDED-ENTRY PUBLICATION (lane/sampled-spec-quality, Item 3 — the fix for
6444 // "a restored spec session never publishes an extended entry", SAMPLED-HIT.md
6445 // finding (d)). Pre-lane, publication was armed only for COLD sessions
6446 // (`spec_resumed == 0` in the worker) and both engine capture sites require a
6447 // non-continuation burst — but a converted hit's first burst IS a continuation,
6448 // so a growing conversation learned exactly ONE boundary and turn 3 could never
6449 // hit a longer prefix than turn 2 did.
6450 //
6451 // WHERE, and why it is safe here: `cache.pos == prefix + suffix` at this exact
6452 // line — the trunk is primed over the whole prompt, nothing is generated, and the
6453 // draft plane rows [0..prompt) are filled just above. That is a complete
6454 // whole-entry boundary (`pos == fed_len`), the same shape the cold seed capture
6455 // publishes; the worker's existing publication sweep picks it up because it is
6456 // keyed on `boundary_capture.is_some()` and is sampler- and resume-independent.
6457 // NOT the partial-restore hazard: the boundary is this session's own prompt END,
6458 // never mid-entry, so `entry_pos != fed_len` still refuses on the way back in.
6459 // Failure is SILENT by design (the turn_ckpt / boundary-capture convention):
6460 // publication is an optimization, never a correctness dependency.
6461 if spec_restore_republish_on() {
6462 debug_assert_eq!(
6463 cache.pos,
6464 pos + t,
6465 "extended-entry capture must sit at the restored session's prompt end",
6466 );
6467 if let Ok(snap) = cache.snapshot(e) {
6468 boundary_capture = Some(SpecBoundaryCapture {
6469 snap,
6470 pos: pos + t,
6471 logits: feed_logits.clone(),
6472 last_h: capture_boundary_hidden(e, &h_rows, t, n_embd),
6473 });
6474 }
6475 }
6476 // continuation seed: the feed's boundary logits ARE the plain path's boundary
6477 // logits (same program), so greedy's argmax here is plain's first emitted token,
6478 // and the sampled draw is the cold sampled session's own first token.
6479 next_pred = Some(if sampled {
6480 let sp = sampling.expect("sampled implies a sampler");
6481 // `committed` is still the restored prefix here; the suffix joins it below —
6482 // so this is the last-N window over the WHOLE prompt, exactly the cold
6483 // session's own window at its first token.
6484 let hist = pen_window_seed(&committed, suffix, sp.penalty_last_n);
6485 match sample_boundary_token(
6486 e,
6487 &feed_logits,
6488 &sp,
6489 &hist,
6490 &mut sctr,
6491 "restore-suffix-feed",
6492 ) {
6493 Ok(t) => t,
6494 // the trunk is already fed: hand nothing back, the worker serves the
6495 // request cold-plain. Never fall back to an argmax — that would put a
6496 // greedy token in a sampled stream to save a slow path.
6497 Err(err) => {
6498 return dirty(format!("boundary token draw failed: {err}"));
6499 }
6500 }
6501 } else {
6502 argmax(&feed_logits) as u32
6503 });
6504 let mut lh = match e.uninit(n_embd) {
6505 Ok(b) => b,
6506 Err(err) => return dirty(format!("boundary hidden alloc: {err}")),
6507 };
6508 if let Err(err) = e.copy_view_into(
6509 &mut lh,
6510 0,
6511 &h_rows.slice((t - 1) * n_embd..t * n_embd),
6512 n_embd,
6513 ) {
6514 return dirty(format!("boundary hidden copy: {err}"));
6515 }
6516 last_h_dev = Some(lh);
6517 committed.extend_from_slice(suffix);
6518 } else {
6519 // FULL-COVER shape (empty suffix — the identical-repeat / agent-loop shape): the
6520 // ENTRY's boundary logits are the boundary row, and this is the token the cold
6521 // session emits from that same row. Owned here rather than in the worker so the
6522 // sampled draw cannot be half-applied on one shape (the worker used to argmax it).
6523 if boundary_logits.is_empty() {
6524 return fail(
6525 cache,
6526 "full-cover restore without the entry's boundary logits".into(),
6527 );
6528 }
6529 next_pred = Some(if sampled {
6530 let sp = sampling.expect("sampled implies a sampler");
6531 let hist = pen_window_seed(&committed, &[], sp.penalty_last_n);
6532 match sample_boundary_token(
6533 e,
6534 boundary_logits,
6535 &sp,
6536 &hist,
6537 &mut sctr,
6538 "restore-full-cover",
6539 ) {
6540 Ok(t) => t,
6541 // nothing has been mutated on this shape — hand the carrier back and let
6542 // the hit serve PLAIN (the banked pre-lane path).
6543 Err(err) => {
6544 return fail(cache, format!("boundary token draw failed: {err}"));
6545 }
6546 }
6547 } else {
6548 argmax(boundary_logits) as u32
6549 });
6550 }
6551 Ok(SpecSession {
6552 cache,
6553 scratch,
6554 committed,
6555 last_h: last_h_dev,
6556 next_pred,
6557 sctr,
6558 uctr: 0,
6559 draft_ctx: None,
6560 pending_tok: None,
6561 turn_ckpt: None,
6562 telem: SpecTelemetryCounters::default(),
6563 capture_at: None,
6564 boundary_capture,
6565 })
6566 }
6567
6568 /// Forced-gate exact state comparison. This intentionally reads the real live prefixes from
6569 /// their owning PP devices: matching emitted ids alone would miss a stale `len_d`, recurrent
6570 /// snapshot, or draft-KV row that only corrupts the following round.
6571 pub fn optipipe_compare_session_state(
6572 &self,
6573 e: &Engine,
6574 reference: &SpecSession,
6575 candidate: &SpecSession,
6576 ) -> Result<OptiForkStateIdentity, Box<dyn std::error::Error>> {
6577 fn fail(what: &str) -> Box<dyn std::error::Error> {
6578 format!("optipipe state mismatch: {what}").into()
6579 }
6580 fn same_f32(a: &[f32], b: &[f32]) -> bool {
6581 a.len() == b.len() && a.iter().zip(b).all(|(x, y)| x.to_bits() == y.to_bits())
6582 }
6583 fn compare_layers(
6584 es: &Engine,
6585 range: std::ops::Range<usize>,
6586 reference: &SpecSession,
6587 candidate: &SpecSession,
6588 report: &mut OptiForkStateIdentity,
6589 ) -> Result<(), Box<dyn std::error::Error>> {
6590 for il in range {
6591 match (&reference.cache.kv[il], &candidate.cache.kv[il]) {
6592 (Some(a), Some(b)) => {
6593 if a.len != b.len {
6594 return Err(fail(&format!(
6595 "layer {il} host KV len {} != {}",
6596 a.len, b.len
6597 )));
6598 }
6599 let ad = es.dtoh_i32(&a.len_d)?;
6600 let bd = es.dtoh_i32(&b.len_d)?;
6601 if ad != bd || ad.first().copied() != Some(a.len as i32) {
6602 return Err(fail(&format!(
6603 "layer {il} device KV len {ad:?} != {bd:?} (host={})",
6604 a.len,
6605 )));
6606 }
6607 let kb = a.len * a.k_tok_bytes;
6608 let vb = a.len * a.v_tok_bytes;
6609 if kb > 0 {
6610 let ak = es.dtoh_u8_view(&a.k.slice(0..kb))?;
6611 let bk = es.dtoh_u8_view(&b.k.slice(0..kb))?;
6612 if ak != bk {
6613 let at = ak.iter().zip(&bk).position(|(x, y)| x != y).unwrap();
6614 return Err(fail(&format!(
6615 "layer {il} K bytes at byte {at} row {} offset {}: {} != {}",
6616 at / a.k_tok_bytes,
6617 at % a.k_tok_bytes,
6618 ak[at],
6619 bk[at],
6620 )));
6621 }
6622 }
6623 if vb > 0 {
6624 let av = es.dtoh_u8_view(&a.v.slice(0..vb))?;
6625 let bv = es.dtoh_u8_view(&b.v.slice(0..vb))?;
6626 if av != bv {
6627 let at = av.iter().zip(&bv).position(|(x, y)| x != y).unwrap();
6628 return Err(fail(&format!(
6629 "layer {il} V bytes at byte {at} row {} offset {}: {} != {}",
6630 at / a.v_tok_bytes,
6631 at % a.v_tok_bytes,
6632 av[at],
6633 bv[at],
6634 )));
6635 }
6636 }
6637 report.trunk_kv_bytes += kb + vb;
6638 }
6639 (None, None) => {}
6640 _ => return Err(fail(&format!("layer {il} KV presence"))),
6641 }
6642 match (&reference.cache.recur[il], &candidate.cache.recur[il]) {
6643 (Some(a), Some(b)) => {
6644 let ac = es.dtoh(&a.conv_state)?;
6645 let bc = es.dtoh(&b.conv_state)?;
6646 if !same_f32(&ac, &bc) {
6647 return Err(fail(&format!("layer {il} conv state")));
6648 }
6649 let as_ = es.dtoh(&a.ssm_state)?;
6650 let bs = es.dtoh(&b.ssm_state)?;
6651 if !same_f32(&as_, &bs) {
6652 return Err(fail(&format!("layer {il} SSM state")));
6653 }
6654 report.recurrent_bytes += (ac.len() + as_.len()) * 4;
6655 }
6656 (None, None) => {}
6657 _ => return Err(fail(&format!("layer {il} recurrent presence"))),
6658 }
6659 }
6660 Ok(())
6661 }
6662
6663 if reference.committed != candidate.committed {
6664 return Err(fail("committed token ids"));
6665 }
6666 if reference.cache.pos != candidate.cache.pos
6667 || reference.cache.max_ctx != candidate.cache.max_ctx
6668 {
6669 return Err(fail("cache pos/capacity"));
6670 }
6671 if reference.pending_tok != candidate.pending_tok
6672 || reference.next_pred != candidate.next_pred
6673 || reference.sctr != candidate.sctr
6674 || reference.uctr != candidate.uctr
6675 {
6676 return Err(fail("pending/prediction/counter tail"));
6677 }
6678
6679 let mut report = OptiForkStateIdentity::default();
6680 if let Some(fence) = crate::pp::pp_cuts(self.layers.len()) {
6681 let rt = crate::pp::PpNRt::get(e)?;
6682 for stage in 0..rt.n_stages() {
6683 let _scope = rt.enter(stage);
6684 compare_layers(
6685 rt.engine(stage, e),
6686 fence[stage]..fence[stage + 1],
6687 reference,
6688 candidate,
6689 &mut report,
6690 )?;
6691 }
6692 } else {
6693 compare_layers(e, 0..self.layers.len(), reference, candidate, &mut report)?;
6694 }
6695
6696 let (a, b) = (&reference.scratch.kv, &candidate.scratch.kv);
6697 if a.len != b.len || e.dtoh_i32(&a.len_d)? != e.dtoh_i32(&b.len_d)? {
6698 return Err(fail("draft scratch length"));
6699 }
6700 let kb = a.len * a.k_tok_bytes;
6701 let vb = a.len * a.v_tok_bytes;
6702 if kb > 0 && e.dtoh_u8_view(&a.k.slice(0..kb))? != e.dtoh_u8_view(&b.k.slice(0..kb))? {
6703 return Err(fail("draft scratch K bytes"));
6704 }
6705 if vb > 0 && e.dtoh_u8_view(&a.v.slice(0..vb))? != e.dtoh_u8_view(&b.v.slice(0..vb))? {
6706 return Err(fail("draft scratch V bytes"));
6707 }
6708 report.scratch_kv_bytes = kb + vb;
6709
6710 match (&reference.last_h, &candidate.last_h) {
6711 (Some(a), Some(b)) => {
6712 let ah = e.dtoh(a)?;
6713 let bh = e.dtoh(b)?;
6714 if !same_f32(&ah, &bh) {
6715 return Err(fail("last hidden/seed bytes"));
6716 }
6717 report.hidden_bytes = ah.len() * 4;
6718 }
6719 (None, None) => {}
6720 _ => return Err(fail("last hidden/seed presence")),
6721 }
6722 Ok(report)
6723 }
6724
6725 /// SESSION-AFFINITY REWIND (lane/session-affinity, 2026-08-05): roll `sess` back to its
6726 /// retained prompt-end checkpoint, so a request whose prompt matches
6727 /// `committed[..rewind_pos()]` exactly can resume there and prime only its own delta.
6728 ///
6729 /// EXACTNESS. After this returns, the session is byte-for-byte the state it was in AT that
6730 /// boundary: full-attn KV truncated to it (append-only, position-addressed), GDN conv/ssm
6731 /// restored from the device copy taken there, draft scratch length reset, `committed`
6732 /// truncated, `last_h` = the boundary's predecessor anchor. That is precisely the state a
6733 /// fresh prime of `committed[..pos]` would have produced, so the following suffix prime and
6734 /// every burst after it are identical to a cold run of the same token stream — the
6735 /// committed-tokens-authoritative contract.
6736 ///
6737 /// `next_pred` and `pending_tok` are CLEARED: both describe generation past the boundary,
6738 /// which the rewind discards. The caller therefore must supply a non-empty suffix (a
6739 /// rewound session cannot serve an empty-suffix continuation burst — there is nothing to
6740 /// continue). The persistent draft graph survives: it bakes only session-stable pointers
6741 /// (the scratch KV, the resident embedding), none of which the rewind moves.
6742 ///
6743 /// The checkpoint is CONSUMED (`turn_ckpt` taken): its snapshot buffers are freed here, and
6744 /// this turn's own prime installs a fresh one at the new prompt end. Returns the position
6745 /// rewound to, or `None` when the session holds no checkpoint (caller: full re-prime).
6746 pub fn spec_rewind_to_checkpoint(
6747 &self,
6748 e: &Engine,
6749 sess: &mut SpecSession,
6750 ) -> Result<Option<usize>, Box<dyn std::error::Error>> {
6751 if sess.turn_ckpt.as_ref().is_some_and(|ckpt| {
6752 !sess.cache.can_rollback(&ckpt.snap, 0) || !sess.scratch.can_rewind_to(ckpt.pos)
6753 }) {
6754 return Err(
6755 "SWA ring rewind checkpoint has been lapped; full re-prime required".into(),
6756 );
6757 }
6758 let Some(ckpt) = sess.turn_ckpt.take() else {
6759 return Ok(None);
6760 };
6761 assert!(
6762 ckpt.pos <= sess.committed.len(),
6763 "checkpoint past committed ({} > {})",
6764 ckpt.pos,
6765 sess.committed.len()
6766 );
6767 // Restore through each layer's owning engine. A single primary-engine rollback is not
6768 // sufficient when the serving cache is stage-owned under cross-device PP.
6769 crate::pp::restore_cache_checkpoint(e, &self.cfg, None, &mut sess.cache, &ckpt.snap)?;
6770 debug_assert_eq!(
6771 sess.cache.pos, ckpt.pos,
6772 "rollback landed off the checkpoint"
6773 );
6774 sess.scratch.set_len(e, ckpt.pos)?;
6775 sess.committed.truncate(ckpt.pos);
6776 sess.last_h = Some(ckpt.last_h);
6777 sess.next_pred = None;
6778 sess.pending_tok = None;
6779 Ok(Some(ckpt.pos))
6780 }
6781
6782 /// Grow a parked speculative session to `target_cap` and rewind it to its retained turn
6783 /// checkpoint without re-priming the checkpoint prefix.
6784 ///
6785 /// The trunk cache is restored exactly like a plain grown cache: append-only full-attention
6786 /// KV rows come from the parked cache, while recurrent state comes from the checkpoint's
6787 /// owned snapshot. The MTP scratch is also context-linear and its rows below the checkpoint
6788 /// remain authoritative, so they are copied into a fresh larger scratch before its length is
6789 /// truncated. Pointer-baking draft graphs are dropped and recaptured on the next burst.
6790 ///
6791 /// All fallible work completes before `sess` is mutated. A failed allocation or copy leaves
6792 /// the parked session intact, allowing the caller one reclaim-and-retry attempt.
6793 pub fn spec_grow_and_rewind_to_checkpoint(
6794 &self,
6795 e: &Engine,
6796 sess: &mut SpecSession,
6797 target_cap: usize,
6798 ) -> Result<Option<usize>, Box<dyn std::error::Error>> {
6799 if target_cap <= sess.cache.max_ctx {
6800 return self.spec_rewind_to_checkpoint(e, sess);
6801 }
6802 let Some(ckpt) = sess.turn_ckpt.as_ref() else {
6803 return Ok(None);
6804 };
6805 if ckpt.pos == 0 || ckpt.pos > sess.committed.len() {
6806 return Err(format!(
6807 "checkpoint pos {} outside committed length {}",
6808 ckpt.pos,
6809 sess.committed.len(),
6810 )
6811 .into());
6812 }
6813 if ckpt.pos > target_cap {
6814 return Err(format!(
6815 "checkpoint pos {} exceeds grown capacity {target_cap}",
6816 ckpt.pos,
6817 )
6818 .into());
6819 }
6820
6821 let mut grown_cache = crate::pp::new_cache(e, &self.cfg, target_cap)?;
6822 let mut grown_scratch = MtpScratch::new(
6823 e,
6824 &self.cfg,
6825 target_cap,
6826 self.mtp.as_ref().and_then(|m| m.geom.as_ref()),
6827 )?;
6828 crate::pp::restore_cache_checkpoint(
6829 e,
6830 &self.cfg,
6831 Some(&sess.cache),
6832 &mut grown_cache,
6833 &ckpt.snap,
6834 )?;
6835
6836 let src = &sess.scratch.kv;
6837 let dst = &mut grown_scratch.kv;
6838 if ckpt.pos > src.len
6839 || src.kv_dim_k != dst.kv_dim_k
6840 || src.kv_dim_v != dst.kv_dim_v
6841 || src.k_tok_bytes != dst.k_tok_bytes
6842 || src.v_tok_bytes != dst.v_tok_bytes
6843 {
6844 return Err(format!(
6845 "checkpoint draft layout mismatch (pos {}, source len {})",
6846 ckpt.pos, src.len,
6847 )
6848 .into());
6849 }
6850 let kb = ckpt.pos * src.k_tok_bytes;
6851 let vb = ckpt.pos * src.v_tok_bytes;
6852 if kb > 0 {
6853 e.copy_u8_into(&mut dst.k, 0, &src.k, kb)?;
6854 }
6855 if vb > 0 {
6856 e.copy_u8_into(&mut dst.v, 0, &src.v, vb)?;
6857 }
6858 grown_scratch.set_len(e, ckpt.pos)?;
6859 // The old scratch is dropped immediately after publication below. Bound its D2D reads
6860 // first; growth happens once per rewritten turn, outside the decode hot loop.
6861 e.stream().synchronize()?;
6862
6863 let ckpt = sess
6864 .turn_ckpt
6865 .take()
6866 .expect("checkpoint remained present through transactional grow");
6867 let pos = ckpt.pos;
6868 sess.cache = grown_cache;
6869 sess.scratch = grown_scratch;
6870 sess.committed.truncate(pos);
6871 sess.last_h = Some(ckpt.last_h);
6872 sess.next_pred = None;
6873 sess.pending_tok = None;
6874 sess.draft_ctx = None;
6875 debug_assert_eq!(sess.cache.pos, pos, "grown rewind landed off checkpoint");
6876 debug_assert_eq!(
6877 sess.scratch.kv.len, pos,
6878 "grown draft rewind landed off checkpoint"
6879 );
6880 Ok(Some(pos))
6881 }
6882
6883 /// Commit a carried pending bonus (see SpecSession::pending_tok): one T=1 trunk pass
6884 /// (its logits' argmax becomes next_pred) + the draft-KV fill at the carried anchor —
6885 /// byte-identical to the pre-carry session tail. Required before a non-empty-suffix
6886 /// prime, a sampled turn, or parking a session for pool reuse. No-op without a pending.
6887 /// `sampling` is the sampler of the request that will CONSUME the resulting `next_pred`
6888 /// (lane/sampled-spec-quality): this is a boundary site like any other, so a sampled
6889 /// consumer must get a DRAWN token, not an argmax. Pass `None` from the park/demote
6890 /// callers — a pending only ever exists on the GREEDY tail, and the consumer of a
6891 /// park-time flush is a future request whose sampler is not knowable here (residual
6892 /// named at the pool-resume probe in worker.rs and in SAMPLED-QUALITY.md).
6893 pub fn spec_flush_pending(
6894 &self,
6895 e: &Engine,
6896 sess: &mut SpecSession,
6897 sampling: Option<SpecSampling>,
6898 ) -> Result<(), Box<dyn std::error::Error>> {
6899 let Some(b) = sess.pending_tok.take() else {
6900 return Ok(());
6901 };
6902 let mtp = self
6903 .mtp
6904 .as_ref()
6905 .expect("pending carry requires an MTP head");
6906 let n_embd = self.cfg.n_embd as usize;
6907 let (embd_qt, embd_rb) = self.embd.qt_and_row_bytes(n_embd);
6908 let embd_gpu = if spec_host_embd() {
6909 None
6910 } else {
6911 Some(
6912 self.embd_gpu
6913 .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload")),
6914 )
6915 };
6916 let embd_dev = embd_gpu.map(|g| (g, embd_qt, embd_rb));
6917 let pos_b = sess.cache.pos;
6918 sess.scratch.set_len(e, pos_b)?;
6919 let (lg_b, hb) = self.spec_target_step_h(e, b, &mut sess.cache)?;
6920 sess.next_pred = Some(match sampling {
6921 Some(sp) if sp.temp > 0.0 && spec_sampled_boundary_on() => {
6922 // window includes `b` itself: it is committed by this pass, and the pre-lane
6923 // code never counted a boundary token in the penalty history at all.
6924 let hist = pen_window_seed(&sess.committed, &[b], sp.penalty_last_n);
6925 sample_boundary_token(e, &lg_b, &sp, &hist, &mut sess.sctr, "flush-pending")?
6926 }
6927 _ => argmax(&lg_b) as u32,
6928 });
6929 let anchor = sess
6930 .last_h
6931 .as_ref()
6932 .expect("pending carry requires last_h (the predecessor-row anchor)");
6933 self.mtp_kv_fill(e, mtp, &[b], anchor, pos_b, &mut sess.scratch, embd_dev)?;
6934 sess.last_h = Some(hb);
6935 sess.committed.push(b);
6936 Ok(())
6937 }
6938
6939 /// Solo target feed used only at speculative round boundaries. Step35 serving made its
6940 /// staged batched B=1 graph authoritative, so a speculative session must enter and leave
6941 /// rounds through that same graph. Other model families keep their eager T=1 contract.
6942 fn spec_target_step_h(
6943 &self,
6944 e: &Engine,
6945 token: u32,
6946 cache: &mut Cache,
6947 ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
6948 if self.cfg.step35.is_none() && !self.qwen35_serving_class() {
6949 return self.decode_step_h(e, token, cache);
6950 }
6951 let pos0 = cache.pos;
6952 let (logits, hidden) = self.decode_step_t_core(e, &[token], pos0, cache, None, None)?;
6953 Ok((e.dtoh(&logits)?, hidden))
6954 }
6955
6956 /// The archs whose LIVE B=1 serving runs the generic BATCHED numeric class (decode_step_batch
6957 /// walk + batched head), so their spec verify must run the SAME class. MoE learned this
6958 /// 2026-08-14 AM (4b777ccc5); the dense hybrid reproduced the identical near-tie flip class
6959 /// the same day on Qwen3.8-27B — eager-class verify logits drift from batched-class serving
6960 /// logits ("1 ULP at layer 2 → 2.3e-1 logit maxdiff at the head"), and the GDN recurrence
6961 /// carries the drift until a near-tie flips deep in generation. One predicate so the five
6962 /// dispatch sites cannot drift apart again.
6963 /// Draft-graph head admissibility (lane/draftcost-moe, 2026-08-20): the capture body
6964 /// (`mtp_head_forward_cap`) supports Dense heads AND resident-MoE heads
6965 /// (`Ffn::Moe(m) if m.dev_exps.is_some()`); non-resident MoE still refuses inside the
6966 /// capture and the caller falls back to the eager chain by design. Trunk FFN class is
6967 /// irrelevant — the graph body is the HEAD forward only. One predicate for all three
6968 /// eligibility sites so they cannot drift (the qwen35_serving_class lesson).
6969 fn mtp_graph_capturable(&self) -> bool {
6970 self.mtp
6971 .as_ref()
6972 .map(|m| match &m.ffn {
6973 crate::hybrid::Ffn::Dense { .. } => true,
6974 crate::hybrid::Ffn::Moe(mo) => mo.dev_exps.is_some(),
6975 })
6976 .unwrap_or(false)
6977 }
6978
6979 fn qwen35_serving_class(&self) -> bool {
6980 matches!(
6981 self.cfg.arch,
6982 memra_gguf::config::Arch::Qwen35 | memra_gguf::config::Arch::Qwen35Moe
6983 )
6984 }
6985
6986 /// Reduced-matrix admission for increment 1. This deliberately does not change the PP-2
6987 /// serving policy: the worker calls it only after `MEMRA_SPEC_PIPE=1` and an explicit spec
6988 /// session already exist.
6989 pub fn spec_pipe_available(&self, e: &Engine) -> bool {
6990 if std::env::var("MEMRA_SPEC_PIPE").as_deref() != Ok("1")
6991 || !spec_devacc()
6992 || spec_replay_env_enabled()
6993 || spec_stream()
6994 || std::env::var("MEMRA_SPEC_ADAPT").as_deref() == Ok("1")
6995 || std::env::var("MEMRA_SPEC_PMIN0").as_deref() == Ok("1")
6996 || std::env::var("MEMRA_SPEC_PP_ANATOMY").as_deref() == Ok("1")
6997 || std::env::var("MEMRA_SPEC_PMIN")
6998 .ok()
6999 .and_then(|v| v.parse::<f32>().ok())
7000 .unwrap_or(0.0)
7001 > 0.0
7002 || self.is_gemma4_e4b()
7003 || self.cfg.gemma4.is_some()
7004 || self.mtp.is_none()
7005 {
7006 return false;
7007 }
7008 let Some(cuts) = crate::pp::pp_cuts(self.layers.len()) else {
7009 return false;
7010 };
7011 if cuts.len() != 3 || crate::pp::pp2_streams_off() || !crate::pp::spec_pp_on() {
7012 return false;
7013 }
7014 crate::pp::PpNRt::get(e)
7015 .map(|rt| rt.n_stages() == 2 && rt.cross_device())
7016 .unwrap_or(false)
7017 }
7018
7019 /// Two warm greedy continuation bursts over one PP-2 interval coordinator. The two existing
7020 /// `generate_spec_inner2` call stacks own all per-session round locals; only phase issue order
7021 /// changes. No callback is accepted in increment 1 — the worker publishes each completed burst.
7022 #[allow(clippy::too_many_arguments)]
7023 pub fn generate_spec_session_pair(
7024 &self,
7025 e: &Engine,
7026 sess_a: &mut SpecSession,
7027 max_new_a: usize,
7028 k_a: usize,
7029 sess_b: &mut SpecSession,
7030 max_new_b: usize,
7031 k_b: usize,
7032 ) -> Result<((Vec<u32>, usize, usize), (Vec<u32>, usize, usize)), Box<dyn std::error::Error>>
7033 {
7034 if !self.spec_pipe_available(e) {
7035 return Err("two-session speculative pipeline is outside its reduced matrix".into());
7036 }
7037 if max_new_a == 0 || max_new_b == 0 || k_a == 0 || k_b == 0 {
7038 return Err(
7039 "two-session speculative pipeline requires non-empty positive-K bursts".into(),
7040 );
7041 }
7042 for sess in [&*sess_a, &*sess_b] {
7043 if sess.committed.is_empty()
7044 || sess.last_h.is_none()
7045 || (sess.next_pred.is_none() && sess.pending_tok.is_none())
7046 {
7047 return Err("two-session speculative pipeline requires warm continuations".into());
7048 }
7049 }
7050
7051 let graph_ok = std::env::var("MEMRA_SPEC_NOGRAPH").is_err()
7052 && !spec_host_embd()
7053 && self.mtp_graph_capturable()
7054 && !crate::model::full_prec_enabled();
7055 let graph_a = graph_ok && k_a + 2 < 96;
7056 let graph_b = graph_ok && k_b + 2 < 96;
7057 let was_tracking = e.ctx().is_event_tracking();
7058 if (graph_a || graph_b) && was_tracking {
7059 unsafe {
7060 e.ctx().disable_event_tracking();
7061 }
7062 }
7063
7064 static LOGGED: std::sync::Once = std::sync::Once::new();
7065 LOGGED.call_once(|| {
7066 eprintln!("[spec-pipe] two-session PP-2 continuation pipeline engaged");
7067 });
7068 let sync = std::sync::Arc::new(SpecPipeSync::new());
7069 let lane_a = SpecPipeLane {
7070 sync: sync.clone(),
7071 lane: 0,
7072 };
7073 let lane_b = SpecPipeLane { sync, lane: 1 };
7074 let mut sess_b_ptr = SpecPipeSessionPtr(sess_b as *mut SpecSession);
7075 let (result_a, result_b) = std::thread::scope(|scope| {
7076 let b = scope.spawn(move || {
7077 let mut finish = SpecPipeFinish::new(&lane_b);
7078 let sess_b = unsafe { sess_b_ptr.get_mut() };
7079 let result = e
7080 .ctx()
7081 .bind_to_thread()
7082 .map_err(|err| err.to_string())
7083 .and_then(|_| {
7084 self.generate_spec_inner2(
7085 e,
7086 &[],
7087 max_new_b,
7088 k_b,
7089 graph_b,
7090 Some(sess_b),
7091 None,
7092 None,
7093 None,
7094 None,
7095 Some(&lane_b),
7096 )
7097 .map_err(|err| err.to_string())
7098 });
7099 finish.close(result.is_err());
7100 result
7101 });
7102 let mut finish = SpecPipeFinish::new(&lane_a);
7103 let result_a = self.generate_spec_inner2(
7104 e,
7105 &[],
7106 max_new_a,
7107 k_a,
7108 graph_a,
7109 Some(sess_a),
7110 None,
7111 None,
7112 None,
7113 None,
7114 Some(&lane_a),
7115 );
7116 finish.close(result_a.is_err());
7117 let result_b = b
7118 .join()
7119 .map_err(|_| "paired speculative session B panicked".to_string())
7120 .and_then(|r| r);
7121 (result_a, result_b)
7122 });
7123
7124 if (graph_a || graph_b) && was_tracking {
7125 unsafe {
7126 e.ctx().enable_event_tracking();
7127 }
7128 }
7129 let result_a = result_a?;
7130 let result_b = result_b.map_err(|err| -> Box<dyn std::error::Error> { err.into() })?;
7131 Ok((result_a, result_b))
7132 }
7133
7134 /// One spec-decode turn on a live session. `suffix` = the NEW tokens only (turn N+1's user
7135 /// message rendered through the chat template continuation). Returns (new tokens emitted,
7136 /// drafted, accepted); session.committed grows by suffix + emitted.
7137 pub fn generate_spec_session(
7138 &self,
7139 e: &Engine,
7140 sess: &mut SpecSession,
7141 suffix: &[u32],
7142 max_new: usize,
7143 k: usize,
7144 ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
7145 self.generate_spec_session_sampled(e, sess, suffix, max_new, k, None, None)
7146 }
7147
7148 /// Serve-path sampled spec: routes the burst through the rejection-sampling verify with
7149 /// per-SESSION Philox continuity (sess.sctr/uctr). None = env-driven (CLI) or greedy.
7150 /// Filters (top-k/p/min-p) apply SYMMETRICALLY to draft q and verify p — distribution-exact
7151 /// for the filtered target (feat/filtered-spec).
7152 ///
7153 /// `on_commit` (sse-cadence, 2026-08-05): called with each newly-emitted slice of the
7154 /// output — once right after the prime's first token, then once per round commit — so a
7155 /// streaming caller can flush text at round cadence instead of once per burst. The slices
7156 /// are disjoint, in order, and concatenate to exactly the returned token vec. Emission-
7157 /// timing only: token bytes, session state, and exactness are untouched.
7158 ///
7159 /// The returned bool is a CONTINUE-VERDICT (admission yield, 2026-08-06): `false` ends
7160 /// the burst at the current round boundary, exactly as if `max_new` had been reached —
7161 /// the caller's scheduler regains control without waiting the burst out. Burst size is
7162 /// content-neutral (spec-levers battery), so an early exit moves WHEN the burst returns,
7163 /// never what tokens say. The slice may be EMPTY (a poll-only boundary — round-stream
7164 /// drains and the defensive tail flush can land with nothing new committed).
7165 #[allow(clippy::too_many_arguments)]
7166 pub fn generate_spec_session_sampled(
7167 &self,
7168 e: &Engine,
7169 sess: &mut SpecSession,
7170 suffix: &[u32],
7171 max_new: usize,
7172 k: usize,
7173 sampling: Option<SpecSampling>,
7174 on_commit: Option<&mut dyn FnMut(&[u32]) -> bool>,
7175 ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
7176 self.generate_spec_session_sampled_prime_split(
7177 e, sess, suffix, max_new, k, sampling, None, on_commit,
7178 )
7179 }
7180
7181 /// Serve-only cold-prime segmentation twin. `prime_split` is the same stable boundary the
7182 /// plain worker would honor before entering its sub-floor tokenwise tail; warm continuations
7183 /// pass `None` and stay on the existing zero-prime path.
7184 #[allow(clippy::too_many_arguments)]
7185 pub fn generate_spec_session_sampled_prime_split(
7186 &self,
7187 e: &Engine,
7188 sess: &mut SpecSession,
7189 suffix: &[u32],
7190 max_new: usize,
7191 k: usize,
7192 sampling: Option<SpecSampling>,
7193 prime_split: Option<usize>,
7194 on_commit: Option<&mut dyn FnMut(&[u32]) -> bool>,
7195 ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
7196 self.generate_spec_session_constrained_prime_split(
7197 e,
7198 sess,
7199 suffix,
7200 max_new,
7201 k,
7202 sampling,
7203 None,
7204 prime_split,
7205 on_commit,
7206 )
7207 }
7208
7209 /// `generate_spec_session_sampled` + GRAMMAR (constrained decoding, 2026-08-03): the
7210 /// hook truncates acceptance at the first grammar-illegal token AFTER the exactness
7211 /// verify (grammar is an extra rejection rule, ordering like the batched-verify twins)
7212 /// and replaces an illegal bonus with the MASKED argmax of the target's own verify
7213 /// column — token-identical to constrained plain greedy decode. GREEDY only (the
7214 /// worker routes sampled constrained to plain decode). Acceptance under tight grammars
7215 /// may drop (drafter is unconstrained); that is measured, not hidden.
7216 #[allow(clippy::too_many_arguments)]
7217 pub fn generate_spec_session_constrained(
7218 &self,
7219 e: &Engine,
7220 sess: &mut SpecSession,
7221 suffix: &[u32],
7222 max_new: usize,
7223 k: usize,
7224 sampling: Option<SpecSampling>,
7225 constraint: Option<&mut dyn SpecConstraint>,
7226 on_commit: Option<&mut dyn FnMut(&[u32]) -> bool>,
7227 ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
7228 self.generate_spec_session_constrained_prime_split(
7229 e, sess, suffix, max_new, k, sampling, constraint, None, on_commit,
7230 )
7231 }
7232
7233 #[allow(clippy::too_many_arguments)]
7234 pub fn generate_spec_session_constrained_prime_split(
7235 &self,
7236 e: &Engine,
7237 sess: &mut SpecSession,
7238 suffix: &[u32],
7239 max_new: usize,
7240 k: usize,
7241 sampling: Option<SpecSampling>,
7242 constraint: Option<&mut dyn SpecConstraint>,
7243 prime_split: Option<usize>,
7244 on_commit: Option<&mut dyn FnMut(&[u32]) -> bool>,
7245 ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
7246 if constraint.is_some() && sampling.is_some_and(|s| s.temp > 0.0) {
7247 return Err(
7248 "constrained spec decode is greedy-only (worker routes sampled \
7249 constrained to plain decode)"
7250 .into(),
7251 );
7252 }
7253 // PENDING-CARRY entry flush: a carried bonus precedes any new suffix in the sequence,
7254 // so it must commit BEFORE the suffix primes; the sampled path doesn't carry (its
7255 // round-0 accept needs the commit pass's logits). Empty-suffix greedy bursts — the
7256 // serve continuation case — consume the carry in-loop with zero solo passes.
7257 if sess.pending_tok.is_some()
7258 && (!suffix.is_empty() || sampling.map_or(false, |s| s.temp > 0.0))
7259 {
7260 self.spec_flush_pending(e, sess, sampling)?;
7261 }
7262
7263 // FULL_PREC forces the EAGER draft: the graph capture would enclose cuBLASLt f32 GEMV
7264 // (the FloatBf16 else-branches) and a bf16_to_f32 dequant alloc — neither is stream-capture
7265 // safe. Eager rides matmul/matmul_decode_exact, which dequant FloatBf16 on use. (§item 2.)
7266 let graph_draft = std::env::var("MEMRA_SPEC_NOGRAPH").is_err()
7267 && !spec_host_embd()
7268 && self.mtp_graph_capturable()
7269 && k + 2 < 96
7270 && !crate::model::full_prec_enabled();
7271 let was_tracking = e.ctx().is_event_tracking();
7272 if graph_draft && was_tracking {
7273 unsafe {
7274 e.ctx().disable_event_tracking();
7275 }
7276 }
7277 let r = self.generate_spec_inner2(
7278 e,
7279 suffix,
7280 max_new,
7281 k,
7282 graph_draft,
7283 Some(sess),
7284 sampling,
7285 constraint,
7286 on_commit,
7287 prime_split,
7288 None,
7289 );
7290 if graph_draft && was_tracking {
7291 unsafe {
7292 e.ctx().enable_event_tracking();
7293 }
7294 }
7295 let (out, d, a) = r?;
7296 Ok((out, d, a))
7297 }
7298
7299 pub fn generate_spec(
7300 &self,
7301 e: &Engine,
7302 prompt: &[u32],
7303 max_new: usize,
7304 k: usize,
7305 ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
7306 // FULL_PREC forces eager (see generate_spec_session note): CUDA graph capture cannot
7307 // enclose cuBLASLt f32 GEMV or the bf16_to_f32 dequant alloc the FloatBf16 path needs.
7308 let graph_draft = std::env::var("MEMRA_SPEC_NOGRAPH").is_err()
7309 && !spec_host_embd()
7310 && self.mtp_graph_capturable()
7311 && k + 2 < 96
7312 && !crate::model::full_prec_enabled();
7313 if !graph_draft {
7314 return self.generate_spec_inner2(
7315 e, prompt, max_new, k, false, None, None, None, None, None, None,
7316 );
7317 }
7318 let was_tracking = e.ctx().is_event_tracking();
7319 if was_tracking {
7320 unsafe {
7321 e.ctx().disable_event_tracking();
7322 }
7323 }
7324 let r = self.generate_spec_inner2(
7325 e, prompt, max_new, k, true, None, None, None, None, None, None,
7326 );
7327 if was_tracking {
7328 unsafe {
7329 e.ctx().enable_event_tracking();
7330 }
7331 }
7332 r
7333 }
7334
7335 fn generate_spec_inner2(
7336 &self,
7337 e: &Engine,
7338 prompt: &[u32],
7339 max_new: usize,
7340 k: usize,
7341 graph_draft: bool,
7342 mut sess: Option<&mut SpecSession>,
7343 sampling: Option<SpecSampling>,
7344 mut constraint: Option<&mut dyn SpecConstraint>,
7345 mut on_commit: Option<&mut dyn FnMut(&[u32]) -> bool>,
7346 prime_split: Option<usize>,
7347 pipe: Option<&SpecPipeLane>,
7348 ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
7349 assert!(k >= 1, "k must be >= 1");
7350 if let Some(p) = pipe {
7351 p.setup_begin()?;
7352 }
7353 // sse-cadence flush cursor: everything in out[..flushed] has been handed to on_commit.
7354 let mut flushed = 0usize;
7355 // admission yield (2026-08-06): on_commit's continue-verdict; false = end the burst
7356 // at the next round boundary (same exit as max_new reached — the session tail runs).
7357 // Initialized by the unconditional post-prime flush below.
7358 let mut keep_going;
7359 let mtp = self
7360 .mtp
7361 .as_ref()
7362 .expect("generate_spec requires an MTP head (nextn_predict_layers>0)");
7363 let n_vocab = self.output.out_features();
7364 // FR-Spec: the draft head may be TRIMMED (fewer rows than n_vocab); the draft argmax runs
7365 // over the draft vocab and the winning index maps through d2t to a TARGET token id.
7366 // Everything downstream (verify/accept/commit) sees target ids only — exactness unchanged.
7367 let d_vocab = mtp
7368 .shared_head_head
7369 .as_ref()
7370 .unwrap_or(&self.output)
7371 .out_features();
7372 let n_embd = self.cfg.n_embd as usize;
7373 // SESSION MODE: reuse the live cache/scratch, prime only the suffix. `base` = tokens
7374 // already committed (their state is in the caches); 0 = fresh single-shot call.
7375 let session_mode = sess.is_some();
7376 let max_ctx = match sess.as_ref() {
7377 Some(s) => s.cache.max_ctx,
7378 None => prompt.len() + max_new + k + 8,
7379 };
7380 let mut own_cache;
7381 let mut own_scratch;
7382 // PREFIX-CACHE capture request threaded out of the session (lane/spec-prefix-cache):
7383 // (requested split, destination slot). Single-shot per burst; fresh calls have none.
7384 let mut sess_capture: Option<(Option<usize>, &mut Option<SpecBoundaryCapture>)> = None;
7385 let (
7386 cache,
7387 scratch,
7388 mut sess_tail,
7389 mut sess_draft_slot,
7390 mut sess_pending_slot,
7391 sess_ckpt_slot,
7392 sess_telem,
7393 ): (
7394 &mut Cache,
7395 &mut MtpScratch,
7396 Option<(
7397 &mut Vec<u32>,
7398 &mut Option<CudaSlice<f32>>,
7399 &mut Option<u32>,
7400 &mut u32,
7401 &mut u32,
7402 )>,
7403 Option<&mut Option<DraftGraphCtx>>,
7404 Option<&mut Option<u32>>,
7405 Option<&mut Option<SpecCheckpoint>>,
7406 Option<&SpecTelemetryCounters>,
7407 ) = match sess.take() {
7408 Some(sr) => {
7409 let SpecSession {
7410 cache,
7411 scratch,
7412 committed,
7413 last_h,
7414 next_pred,
7415 sctr: s_sctr,
7416 uctr: s_uctr,
7417 draft_ctx,
7418 pending_tok,
7419 turn_ckpt,
7420 telem,
7421 capture_at,
7422 boundary_capture,
7423 } = sr;
7424 sess_capture = Some((capture_at.take(), boundary_capture));
7425 (
7426 cache,
7427 scratch,
7428 Some((committed, last_h, next_pred, s_sctr, s_uctr)),
7429 Some(draft_ctx),
7430 Some(pending_tok),
7431 Some(turn_ckpt),
7432 Some(telem),
7433 )
7434 }
7435 None => {
7436 // STAGE-OWNED KV (lane/pp2-spec 2026-08-06) — see `new_session`. Door shut =
7437 // `Cache::new` verbatim.
7438 own_cache = crate::pp::new_cache(e, &self.cfg, max_ctx)?;
7439 // Persistent scratch = max_ctx rows (~2KB/token quantized).
7440 own_scratch = MtpScratch::new(
7441 e,
7442 &self.cfg,
7443 max_ctx,
7444 self.mtp.as_ref().and_then(|m| m.geom.as_ref()),
7445 )?;
7446 (
7447 &mut own_cache,
7448 &mut own_scratch,
7449 None,
7450 None,
7451 None,
7452 None,
7453 None,
7454 )
7455 }
7456 };
7457 let base = cache.pos;
7458 // PENDING-CARRY consume (2026-08-01): a carried bonus reaches here only on the
7459 // empty-suffix GREEDY continuation path (generate_spec_session_sampled flushed every
7460 // other case). It enters the round loop as round-0's pending — verify col 0 — exactly
7461 // like a mid-burst full-accept boundary: no init feed, no tail commit pass.
7462 let carried_pending: Option<u32> = sess_pending_slot.as_mut().and_then(|s| s.take());
7463 // PERSISTENT DRAFT KV (the only mode since 2026-07-08 — the legacy round-local scratch,
7464 // MEMRA_SPEC_KVLOCAL, measured -35 acceptance pts on the 27B p3 sweep and was removed;
7465 // acceptance-only — exactness is verify's job either way).
7466 // HIDDEN-PAIRING CONVENTION (DEFAULT = predecessor-row, 2026-07-04 — the 27B acceptance
7467 // unlock, +16pts): the MTP head is TRAINED on rows pairing token x_p with the trunk
7468 // hidden of its PREDECESSOR h_{p-1} (the reference engine's mtp_update shifts the target
7469 // hiddens right by one; its draft step 0 feeds (id_last, TRUE hidden of the row id_last
7470 // was sampled from)). memra's historical convention paired SAME-ROW (x_p, h_p) in the fill
7471 // and seeded chain step 0 through an extra MTP pass on a duplicated token (the
7472 // pseudo-seed) — measured 27B p2 K=3 acceptance 0.569 vs 0.731, p3 0.445 vs 0.63+, and
7473 // the chain steps j>=1 were already predecessor-shaped, so ONLY the fill + step-0 seed
7474 // move. The fill shifts by one and the chain seeds from the predecessor's true hidden
7475 // DIRECTLY (vh_seed / vx[j-1]) — the pseudo pass disappears (one MTP-block pass saved
7476 // per round on top of the acceptance win). Draft-quality-only: exactness stays the
7477 // verify's job either way. (The legacy same-row pairing seam, MEMRA_SPEC_HSAME, and its
7478 // pseudo-seed passes were removed 2026-07-08 — predecessor pairing won by +16 acc pts;
7479 // the legacy round-local scratch, MEMRA_SPEC_KVLOCAL, went with it.)
7480 // REPLAY-FREE PARTIAL ACCEPT (default, 2026-07-03): partial rounds keep the verify's own
7481 // bit-identical committed-prefix state (KV truncate + recur rebuild from the VerifyCkpt)
7482 // and leave the bonus PENDING — no duplicate trunk pass (profiled ~0.54 extra full weight
7483 // reads/round at long ctx). MEMRA_SPEC_REPLAY=1 restores the legacy rollback+replay (A/B
7484 // + fallback seam).
7485 // Qwen35-MoE replay pin LIFTED (lane/draftcost-moe, 2026-08-20). The pin's stated
7486 // bar — the retained verify-state commit proven equivalent to sequential serving —
7487 // was waiting on this arch running the serving batched verify class, which the
7488 // t-parallel admission (this lane, increment 1) provided: the VerifyCkpt the
7489 // replay-free commit consumes is now produced by the SAME serving-class verify that
7490 // qualified dense qwen35 on 2026-08-15 (where the per-round duplicate replay
7491 // measured 69 -> 30 tok/s). Qualification receipts (run-spec K=1..8 both arms,
7492 // 8-prompt replay-vs-replay-free canary, long-prompt cell):
7493 // research/draftcost-moe-20260820/RECEIPTS.md. MEMRA_SPEC_REPLAY=1 stays the
7494 // rollback + A/B seam.
7495 let spec_replay = spec_replay_env_enabled();
7496 if constraint.is_some() && spec_replay {
7497 return Err(
7498 "constrained spec decode does not support MEMRA_SPEC_REPLAY=1 \
7499 (legacy replay commits an unmasked bonus)"
7500 .into(),
7501 );
7502 }
7503 // TRUE-HIDDEN REFRESH (default in persistent-draft-KV mode): every round overwrites the
7504 // committed positions' scratch entries from the verify's exact hiddens (mtp_kv_fill batch)
7505 // instead of keeping chain-approximate entries. MEMRA_SPEC_NOREFRESH=1 = legacy (A/B seam).
7506 let refresh = std::env::var("MEMRA_SPEC_NOREFRESH").is_err();
7507
7508 // prime: BATCHED cache prime (prime_cache — the measured #1 e2e gap: tokenwise primed at
7509 // ~102/38 tok/s vs the engine's ~2000-5900 tok/s batched prefill). prime_cache returns the
7510 // full pre-output_norm hidden stack [T, n_embd], which IS prompt_h (the persistent-draft-KV
7511 // mtp_kv_fill input) — no per-token collection needed. Prompts below PRIME_MIN_T, and
7512 // MEMRA_PRIME_TOKENWISE=1, and frozen Hy3 CPU/GPU expert splits take the tokenwise
7513 // decode_step_h loop. The latter avoids transient GPU staging of the spilled expert bank.
7514 // EMPTY-SUFFIX CONTINUATION (serve bursts): a session turn with NO new tokens resumes
7515 // generation exactly where the last turn stopped — no prime at all. The stashed
7516 // `next_pred` plays prime_logits' role: it is the token produced from the logits after
7517 // committed.last() by the same rule this entry applies to a cold prime's last row —
7518 // an argmax when greedy, a `sample_boundary_token` draw when sampled (the burst tail,
7519 // or `spec_session_from_restored` for a converted prefix-cache hit, did the drawing
7520 // where the sampler and the session's Philox counters were live). `last_h` seeds the
7521 // predecessor pairing below. Fresh calls and non-empty suffixes take the normal path.
7522 let continuation = prompt.is_empty();
7523 if continuation {
7524 assert!(session_mode, "empty prompt requires a session");
7525 assert!(
7526 sess_tail
7527 .as_ref()
7528 .map_or(false, |(c, lh, np, _, _)| !c.is_empty()
7529 && lh.is_some()
7530 && (np.is_some() || carried_pending.is_some())),
7531 "empty-suffix continuation needs a primed session (committed + last_h + next_pred|pending)"
7532 );
7533 }
7534 let mut prime_logits;
7535 let mut prompt_h: Option<CudaSlice<f32>> = None;
7536 let t_prime = std::time::Instant::now();
7537 let batched_prime = !continuation
7538 && prompt.len() >= crate::hybrid_forward::PRIME_MIN_T
7539 && std::env::var("MEMRA_PRIME_TOKENWISE").is_err()
7540 && !e.frozen_cpu_experts_prefer_tokenwise_prime();
7541 let prime_split = prime_split.filter(|&split| split > 0 && split < prompt.len());
7542 if prime_split.is_some() && (continuation || base != 0) {
7543 return Err("spec prime split is cold-session-only".into());
7544 }
7545 if continuation {
7546 prime_logits = Vec::new();
7547 } else if let Some(split) = prime_split {
7548 if split < crate::hybrid_forward::PRIME_MIN_T {
7549 return Err(format!(
7550 "spec prime split {split} is below PRIME_MIN_T {}",
7551 crate::hybrid_forward::PRIME_MIN_T,
7552 )
7553 .into());
7554 }
7555 // Mirror the plain worker's affinity boundary exactly. The prefix is a request-level
7556 // prime (`queued_after` keeps Step35 arm selection independent of this stop); a tail
7557 // below PRIME_MIN_T then takes the same eager tokenwise continuation as prefill_tick.
7558 // Retain every hidden row so the draft scratch fill remains one coherent prompt.
7559 let mut h_all = e.uninit(prompt.len() * n_embd)?;
7560 let (l, _, h_prefix) =
7561 self.prime_cache(e, &prompt[..split], &mut *cache, prompt.len() - split)?;
7562 e.copy_into(&mut h_all, 0, &h_prefix, split * n_embd)?;
7563 prime_logits = l;
7564 // PREFIX-CACHE BOUNDARY CAPTURE (lane/spec-prefix-cache): the GDN conv/ssm states
7565 // are about to be advanced in place by the tail prime, so this is the ONLY moment
7566 // the boundary's recurrent state exists. Capture iff the worker requested exactly
7567 // this split. cache.pos == split here (the prefix prime just finished). A failed
7568 // snapshot is silent (turn_ckpt convention) — publication is an optimization,
7569 // never a correctness dependency.
7570 if let Some((requested, slot)) = sess_capture.as_mut() {
7571 if *requested == Some(split) {
7572 debug_assert_eq!(cache.pos, split, "boundary capture off the prime split");
7573 if let Ok(snap) = cache.snapshot(e) {
7574 **slot = Some(SpecBoundaryCapture {
7575 snap,
7576 pos: split,
7577 logits: prime_logits.clone(),
7578 // rows [0..split) of h_all are the prefix prime's hiddens — copied
7579 // just above, before the tail prime overwrites nothing (append-only).
7580 last_h: capture_boundary_hidden(e, &h_all, split, n_embd),
7581 });
7582 }
7583 }
7584 }
7585 let tail = &prompt[split..];
7586 if tail.len() >= crate::hybrid_forward::PRIME_MIN_T
7587 && std::env::var("MEMRA_PRIME_TOKENWISE").is_err()
7588 && !e.frozen_cpu_experts_prefer_tokenwise_prime()
7589 {
7590 let (l, _, h_tail) = self.prime_cache(e, tail, &mut *cache, 0)?;
7591 e.copy_into(&mut h_all, split * n_embd, &h_tail, tail.len() * n_embd)?;
7592 prime_logits = l;
7593 } else {
7594 for (i, &tok) in tail.iter().enumerate() {
7595 let (l, h) = self.decode_step_h(e, tok, &mut *cache)?;
7596 e.copy_into(&mut h_all, (split + i) * n_embd, &h, n_embd)?;
7597 prime_logits = l;
7598 }
7599 }
7600 if std::env::var("MEMRA_SPEC_STATS").as_deref() == Ok("1") {
7601 eprintln!("[spec-prime] affinity split={split} tail={}", tail.len());
7602 }
7603 prompt_h = Some(h_all);
7604 } else if batched_prime {
7605 let (l, _h_seed, hiddens) = self.prime_cache(e, prompt, &mut *cache, 0)?;
7606 prime_logits = l;
7607 prompt_h = Some(hiddens);
7608 } else {
7609 prime_logits = Vec::new();
7610 prompt_h = Some(e.uninit(prompt.len() * n_embd)?);
7611 for (i, &tok) in prompt.iter().enumerate() {
7612 let (l, h) = self.spec_target_step_h(e, tok, &mut *cache)?;
7613 if let Some(ph) = prompt_h.as_mut() {
7614 e.copy_into(ph, i * n_embd, &h, n_embd)?;
7615 }
7616 prime_logits = l;
7617 }
7618 }
7619 e.stream().synchronize()?;
7620 // PREFIX-CACHE SEED CAPTURE (lane/spec-prefix-cache): boundary == prompt end (the seed
7621 // case — no shared-prefix split, publish the whole prompt). The prime just finished, so
7622 // cache.pos == base + prompt.len() and the recurrent state IS the boundary state;
7623 // prime_logits are the boundary logits. Cold sessions only (base == 0) — same law as
7624 // prime_split. The mid-prompt capture above already consumed the request if it matched.
7625 if !continuation && base == 0 {
7626 if let Some((requested, slot)) = sess_capture.as_mut() {
7627 if *requested == Some(prompt.len()) && slot.is_none() {
7628 debug_assert_eq!(cache.pos, prompt.len(), "seed capture off prompt end");
7629 if let Ok(snap) = cache.snapshot(e) {
7630 **slot = Some(SpecBoundaryCapture {
7631 snap,
7632 pos: prompt.len(),
7633 logits: prime_logits.clone(),
7634 last_h: prompt_h
7635 .as_ref()
7636 .map(|ph| capture_boundary_hidden(e, ph, prompt.len(), n_embd))
7637 .unwrap_or_default(),
7638 });
7639 }
7640 }
7641 }
7642 }
7643 // Harness timing contract (see crate::PRIME_NANOS): gen-only throughput without the
7644 // prime-subtraction hack.
7645 crate::PRIME_NANOS.store(
7646 t_prime.elapsed().as_nanos() as u64,
7647 std::sync::atomic::Ordering::Relaxed,
7648 );
7649
7650 let (embd_qt, embd_rb) = self.embd.qt_and_row_bytes(n_embd);
7651 // Resident table is fastest when it fits. Large spill deployments can preserve that HBM
7652 // for expert-cache slots and gather only the exact rows needed by MTP/verify from host.
7653 let host_embd = spec_host_embd();
7654 let embd_gpu = if host_embd {
7655 None
7656 } else {
7657 Some(
7658 self.embd_gpu
7659 .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload")),
7660 )
7661 };
7662 let embd_dev = embd_gpu.map(|g| (g, embd_qt, embd_rb));
7663 if host_embd {
7664 eprintln!(
7665 "[spec] host-row embedding: {} bytes kept off HBM",
7666 self.embd.raw.len()
7667 );
7668 }
7669 let mut out: Vec<u32> = Vec::with_capacity(max_new);
7670 let mut total_drafted = 0usize;
7671 let mut total_accepted = 0usize;
7672
7673 // --- SAMPLER FIRST (lane/sampled-spec-quality, 2026-08-19) ---
7674 // The sampler config, the session's Philox counters and the penalty window are parsed
7675 // HERE, above the boundary-token selection, because the boundary token must be drawn
7676 // from the sampler the request asked for. Pre-lane this block sat ~50 lines BELOW the
7677 // selection, which is the whole mechanical reason the boundary token was an argmax:
7678 // the sampler state was not in scope yet. Nothing here depends on the round loop, so
7679 // moving it up is a pure reordering for greedy (`sampled == false` ⇒ every branch
7680 // below takes the argmax path it always took).
7681 // --- SAMPLED SPEC (MEMRA_SPEC_TEMP>0, research/sampled-spec-impl-map.md): rejection-
7682 // sampling verify (Leviathan/Chen) — accept draft x at u < p(x)/q(x), resample from
7683 // norm(max(0,p-q)) on reject, bonus sampled from p on full accept. Counter-based Philox
7684 // everywhere (seed, event) -> reproducible. temp==0/unset = the greedy path, untouched.
7685 let sp = sampling.unwrap_or_else(|| SpecSampling {
7686 temp: std::env::var("MEMRA_SPEC_TEMP")
7687 .ok()
7688 .and_then(|v| v.parse().ok())
7689 .unwrap_or(0.0),
7690 seed: std::env::var("MEMRA_SEED")
7691 .ok()
7692 .and_then(|v| v.parse().ok())
7693 .unwrap_or(42),
7694 top_k: std::env::var("MEMRA_TOP_K")
7695 .ok()
7696 .and_then(|v| v.parse().ok())
7697 .unwrap_or(0),
7698 top_p: std::env::var("MEMRA_TOP_P")
7699 .ok()
7700 .and_then(|v| v.parse().ok())
7701 .unwrap_or(1.0),
7702 min_p: std::env::var("MEMRA_MIN_P")
7703 .ok()
7704 .and_then(|v| v.parse().ok())
7705 .unwrap_or(0.0),
7706 penalty_last_n: std::env::var("MEMRA_PENALTY_LAST_N")
7707 .ok()
7708 .and_then(|v| v.parse().ok())
7709 .unwrap_or(0),
7710 penalty_repeat: std::env::var("MEMRA_PENALTY_REPEAT")
7711 .ok()
7712 .and_then(|v| v.parse().ok())
7713 .unwrap_or(1.0),
7714 penalty_freq: std::env::var("MEMRA_PENALTY_FREQ")
7715 .ok()
7716 .and_then(|v| v.parse().ok())
7717 .unwrap_or(0.0),
7718 penalty_present: std::env::var("MEMRA_PENALTY_PRESENT")
7719 .ok()
7720 .and_then(|v| v.parse().ok())
7721 .unwrap_or(0.0),
7722 });
7723 let (sp_temp, sp_seed) = (sp.temp, sp.seed);
7724 let sampled = sp_temp > 0.0;
7725 // Counters resume from the session (burst continuity: randomness must never repeat
7726 // across generate_spec_session calls); one-shot callers start at (0,0). Read through
7727 // sess_tail — `sess` was take()n into it above, so sess.as_ref() here is always None.
7728 let mut sctr: u32 = sess_tail.as_ref().map(|(_, _, _, s, _)| **s).unwrap_or(0);
7729 let mut uctr: u32 = sess_tail.as_ref().map(|(_, _, _, _, u)| **u).unwrap_or(0);
7730 // Penalties (v2.1): applied to COPIES of q rows and p columns symmetrically (exactness
7731 // for the penalized+filtered target). History = generated tokens, host-tracked window.
7732 let pen_on = sampled
7733 && sp.penalty_last_n > 0
7734 && (sp.penalty_repeat != 1.0 || sp.penalty_freq != 0.0 || sp.penalty_present != 0.0);
7735 // SESSION-SPANNING PENALTY WINDOW (Item 2). Pre-lane this was
7736 // `prompt.iter().rev().take(64).rev()` — the BURST's suffix slice — so a continuation
7737 // burst (the majority of a stream's tokens, and ALL of a converted cache hit's) started
7738 // with an EMPTY penalty history and the client's repetition/frequency/presence penalties
7739 // silently reset at every burst boundary. The window now spans `committed ++ prompt`,
7740 // which is what the API contract says and what the plain sampler's own `history` does.
7741 // Byte-identical to the pre-lane seed for a cold turn-1 burst at the default window.
7742 let mut pen_hist: Vec<u32> = if pen_on {
7743 let sess_hist: &[u32] = if spec_pen_session_on() {
7744 sess_tail
7745 .as_ref()
7746 .map(|(c, ..)| c.as_slice())
7747 .unwrap_or(&[])
7748 } else {
7749 &[] // MEMRA_SPEC_PEN_SESSION=0: pre-lane burst-local window
7750 };
7751 pen_window_seed(sess_hist, prompt, sp.penalty_last_n)
7752 } else {
7753 Vec::new()
7754 };
7755 // First generated token = the BOUNDARY token: greedy takes the argmax of the prompt's
7756 // last logits (== greedy's first token, byte-contract); SAMPLED draws it from the
7757 // request's own filtered/penalized target through the session's Philox stream
7758 // (`sample_boundary_token`, lane/sampled-spec-quality Item 1 — pre-lane this was an
7759 // argmax in both regimes, so ~1 token per burst of a sampled stream was greedy).
7760 // Emit it, then FEED it to establish the loop invariant below.
7761 // PENDING-CARRY: the carried bonus was already emitted by the LAST burst — it becomes
7762 // last_token WITHOUT re-emission, and round 0 consumes it as pending (no init feed).
7763 // CONSTRAINED entry rules: the first emitted token is the MASKED argmax of the
7764 // prompt's last logits (plain constrained-greedy identity); a continuation without
7765 // a carried pending would emit an UNMASKED stashed next_pred — refused loudly (the
7766 // worker never resumes constrained sessions from the pool, so this cannot fire).
7767 if let Some(c) = constraint.as_deref_mut() {
7768 if continuation && carried_pending.is_none() {
7769 return Err("constrained spec continuation requires a carried pending \
7770 (pool resume is unconstrained-only)"
7771 .into());
7772 }
7773 if !continuation {
7774 c.mask_logits(&mut prime_logits)
7775 .map_err(|e2| format!("constraint: {e2}"))?;
7776 }
7777 }
7778 let mut last_token = if let Some(b) = carried_pending {
7779 b
7780 } else if continuation {
7781 // A continuation's boundary token was DRAWN by the burst that stashed it (the
7782 // session tail below), or by `spec_session_from_restored` for a converted
7783 // prefix-cache hit — in both cases from the correct logits row with this same
7784 // session's Philox stream, which is why it can be consumed here as-is.
7785 sess_tail.as_ref().unwrap().2.unwrap()
7786 } else if sampled && constraint.is_none() && spec_sampled_boundary_on() {
7787 sample_boundary_token(e, &prime_logits, &sp, &pen_hist, &mut sctr, "cold-prime")?
7788 } else {
7789 // greedy (byte contract), the rollback door, or constrained (masked-argmax
7790 // identity — the worker routes sampled+constrained to the plain path, and this
7791 // function refuses the combination outright above).
7792 argmax(&prime_logits) as u32
7793 };
7794 if pen_on {
7795 // The boundary token is a GENERATED token: the plain sampler `accept()`s every
7796 // emitted token into its penalty history, and pre-lane the burst's first token
7797 // was invisible to penalties forever (never pushed, and never in `committed`
7798 // until this burst's tail). Covers the carry/continuation seeds too — neither is
7799 // in `committed` yet.
7800 pen_hist.push(last_token);
7801 }
7802 if carried_pending.is_none() {
7803 out.push(last_token);
7804 // grammar advances with every emitted token (carried pendings were consumed
7805 // by the burst that emitted them).
7806 if let Some(c) = constraint.as_deref_mut() {
7807 c.consume(last_token)
7808 .map_err(|e2| format!("constraint: {e2}"))?;
7809 }
7810 }
7811 if continuation {
7812 // draft-KV invariant: entries [0..base) are the session's exact fills; truncate any
7813 // overhang so the chain's first append lands at slot base (== committed.len()).
7814 scratch.set_len(e, base)?;
7815 }
7816 // sse-cadence: hand the caller every not-yet-flushed token (disjoint in-order slices
7817 // concatenating to the full `out`). Called after the prime's first token and after each
7818 // round commit — emission timing only, token bytes untouched. The slice may be EMPTY
7819 // (poll-only boundary: zero-round folds commit nothing new); returns the caller's
7820 // continue-verdict (admission yield, 2026-08-06) — false ends the burst at this round.
7821 fn flush_commit(
7822 cb: &mut Option<&mut dyn FnMut(&[u32]) -> bool>,
7823 out: &[u32],
7824 flushed: &mut usize,
7825 ) -> bool {
7826 if let Some(f) = cb.as_mut() {
7827 let keep = f(&out[*flushed..]);
7828 *flushed = out.len();
7829 keep
7830 } else {
7831 true
7832 }
7833 }
7834 keep_going = flush_commit(&mut on_commit, &out, &mut flushed);
7835 // INVARIANT at loop top: `last_token` is the most-recently-committed/emitted token, its
7836 // KV+recur state IS in `cache` (cache.pos = position right AFTER last_token), `last_pred`
7837 // is the greedy ARGMAX of the logits that predict the token FOLLOWING last_token, and
7838 // `h_seed` = last_token's pre-output_norm hidden. Establish it by feeding last_token once
7839 // (mirrors plain greedy). DEVICE-ARGMAX lever: the accept walk only ever consumes the
7840 // argmax of those logits — never the full vector — so a host u32 replaces the Vec<f32>.
7841 // Trimmed heads: q lives on the trimmed vocab; accept gathers use the TRIMMED index and
7842 // the residual scatters q into target-id space (q=-inf off-trim — the head cannot propose
7843 // those, so their residual mass is p(x), correct by construction).
7844 let d2t_dev: Option<CudaSlice<u32>> = if sampled || crate::spec::spec_stream() {
7845 match &mtp.d2t {
7846 Some(map) => Some(e.htod_u32_v(map)?),
7847 None => None,
7848 }
7849 } else {
7850 None
7851 };
7852 let mut q_full_buf: Option<CudaSlice<f32>> = None;
7853 // host Philox4x32-10 (mirrors spec_sample.cu; independent stream via ctr_lo tag)
7854 let host_u01 = |seed: u64, ctr: u32| -> f32 {
7855 let (m0, m1) = (0xD2511F53u32, 0xCD9E8D57u32);
7856 let (mut c0, mut c1, mut c2, mut c3) = (0xFFFF_FFFEu32, ctr, 0u32, 0u32);
7857 let (mut k0, mut k1) = ((seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
7858 for _ in 0..10 {
7859 let (h0, l0) = (((m0 as u64 * c0 as u64) >> 32) as u32, m0.wrapping_mul(c0));
7860 let (h1, l1) = (((m1 as u64 * c2 as u64) >> 32) as u32, m1.wrapping_mul(c2));
7861 let (n0, n1, n2, n3) = (h1 ^ c1 ^ k0, l1, h0 ^ c3 ^ k1, l0);
7862 c0 = n0;
7863 c1 = n1;
7864 c2 = n2;
7865 c3 = n3;
7866 k0 = k0.wrapping_add(0x9E3779B9);
7867 k1 = k1.wrapping_add(0xBB67AE85);
7868 }
7869 (c0 as f32 + 1.0) * (1.0 / 4294967296.0)
7870 };
7871 let mut draft_logits: Vec<CudaSlice<f32>> = Vec::new(); // retained head logits (q), per slot
7872 let mut draft_stats: Vec<(f32, f32, f32)> = Vec::new(); // (row_max, th_e, z_e) per slot
7873 let mut perturb_buf: Option<CudaSlice<f32>> = None; // gumbel scratch (max(n_vocab,d_vocab))
7874 let mut sample_tok = e.alloc_u32_zeroed(1)?; // residual/bonus sample out
7875 let mut col_buf: Option<CudaSlice<f32>> = None; // materialized verify column
7876 let mut pen_hist_d: Option<CudaSlice<u32>> = None;
7877 let mut pcol_buf: Option<CudaSlice<f32>> = None; // penalized p-column scratch
7878 // MEMRA_SPEC_SETUP_TRACE=1 (diagnostics): per-call wall decomposition of the burst
7879 // SETUP + TAIL segments (the round loop's internals are MEMRA_SPEC_PHASE's job) —
7880 // built to pin the serve per-burst fixed cost (research/spec-serving-20260801).
7881 let setup_trace = std::env::var("MEMRA_SPEC_SETUP_TRACE").as_deref() == Ok("1");
7882 let t_ent = std::time::Instant::now();
7883
7884 // SESSION-AFFINITY TURN CHECKPOINT (lane/session-affinity, 2026-08-05): capture the
7885 // PROMPT-END boundary state so a LATER turn can rewind here and re-prime only its own
7886 // delta instead of the whole conversation. See `SpecCheckpoint` for why this boundary is
7887 // the one that matters (a history-rewriting client mutates what the session GENERATED,
7888 // so the next turn's prompt agrees with this one up to exactly here).
7889 //
7890 // WHERE — AND WHY THIS EXACT LINE. Right after the trunk prime, BEFORE the init feed
7891 // (`decode_step_h(last_token)`) and before round 0: the last instant at which the caches
7892 // hold exactly `base + prompt.len()` rows and nothing generated.
7893 //
7894 // This was WRONG in the first cut of this lane: the capture sat after the draft-KV fill,
7895 // which is also after the init feed, so `cache.pos` was `base + prompt.len() + 1` — the
7896 // boundary included the FIRST GENERATED TOKEN. That token is the first thing inside the
7897 // `<think>` block the client strips, so every later turn's diff diverged exactly one
7898 // token below the checkpoint and affinity declined 100% of the time. Measured on the
7899 // owner regime: "history diverged at 12233 of checkpoint 12234". The off-by-one made the
7900 // whole mechanism inert while looking, from the outside, like a working
7901 // correctness-declines-safely path — hence the decline log carries the offsets.
7902 //
7903 // The full-attn planes are `len`-truncatable so the snapshot copies only the GDN conv/ssm
7904 // state (the reason a spec session could not rewind before). The draft scratch needs no
7905 // copy: rows below the boundary are rewritten by the next turn's own fill.
7906 //
7907 // WHEN: non-empty prime only. An empty-suffix continuation burst adds no prompt boundary
7908 // (its "prompt end" IS the previous checkpoint's, already held), so it keeps the existing
7909 // checkpoint rather than replacing it with a strictly worse one.
7910 //
7911 // FAILURE IS SILENT BY DESIGN: on a VRAM-tight rig the snapshot alloc can fail. That
7912 // costs the NEXT turn its rewind (it re-primes fully, today's behavior) and must never
7913 // fail the burst that is already running — so the error is swallowed, loud only under
7914 // MEMRA_DEBUG_SPEC.
7915 if let Some(slot) = sess_ckpt_slot {
7916 if !continuation {
7917 let pos = cache.pos;
7918 debug_assert_eq!(
7919 pos,
7920 base + prompt.len(),
7921 "turn checkpoint must sit at the prompt end, before the init feed"
7922 );
7923 let anchor: Result<CudaSlice<f32>, Box<dyn std::error::Error>> =
7924 if let Some(ph) = &prompt_h {
7925 // hidden of the LAST primed row = the predecessor anchor at this
7926 // boundary (exactly what a fresh prime of committed[..pos] leaves in
7927 // last_h, and what the next prime's fill reads for its first row).
7928 let np = prompt.len();
7929 e.uninit(n_embd).and_then(|mut a| {
7930 e.copy_view_into(
7931 &mut a,
7932 0,
7933 &ph.slice((np - 1) * n_embd..np * n_embd),
7934 n_embd,
7935 )?;
7936 Ok(a)
7937 })
7938 } else {
7939 Err("no prompt hiddens".into())
7940 };
7941 match (cache.snapshot(e), anchor) {
7942 (Ok(snap), Ok(last_h)) => {
7943 *slot = Some(SpecCheckpoint { snap, pos, last_h });
7944 }
7945 (s, a) => {
7946 *slot = None; // a stale checkpoint would rewind to the WRONG boundary
7947 if std::env::var("MEMRA_DEBUG_SPEC").is_ok() {
7948 let err = s
7949 .err()
7950 .map(|e| e.to_string())
7951 .or_else(|| a.err().map(|e| e.to_string()))
7952 .unwrap_or_default();
7953 eprintln!(
7954 "[spec] turn checkpoint skipped ({err}); \
7955 next turn re-primes in full"
7956 );
7957 }
7958 }
7959 }
7960 }
7961 }
7962 // INIT FEED — skipped on a pending carry: last_token (the carried bonus) is NOT in the
7963 // caches and must NOT be fed solo; round 0's batched verify commits it as col 0. Its
7964 // seed/anchor hidden is the carried last_h (copied below); last_pred is dead in the
7965 // pending path (t_pred reads verify col 0 — the accept walk overwrites it).
7966 let mut last_pred = 0u32;
7967 let mut last_col_logits: Option<CudaSlice<f32>> = None;
7968 // CONSTRAINED: the init feed's logits back the (n_acc==0, base==0) masked-argmax
7969 // recompute in the grammar-truncation walk — retained host-side, round 0 only.
7970 let mut init_logits_host: Option<Vec<f32>> = None;
7971 let h_seed0: CudaSlice<f32> = if carried_pending.is_none() {
7972 let (init_logits, h) = self.spec_target_step_h(e, last_token, &mut *cache)?;
7973 last_pred = argmax(&init_logits) as u32;
7974 if constraint.is_some() {
7975 init_logits_host = Some(init_logits.clone());
7976 }
7977 // sampled mode: p-distribution after last_token, for the j==0/base==0 accept test.
7978 if sampled {
7979 last_col_logits = Some(e.htod(&init_logits)?);
7980 }
7981 h
7982 } else {
7983 // predecessor-row anchor: hidden of the last COMMITTED row (the carry contract).
7984 let lh = sess_tail
7985 .as_ref()
7986 .unwrap()
7987 .1
7988 .as_ref()
7989 .expect("pending carry requires last_h");
7990 e.clone_dtod(lh)?
7991 };
7992 let t_init = t_ent.elapsed();
7993 let mut last_col_stats: Option<(f32, f32, f32)> = None;
7994 // PERSISTENT h_seed buffer (allocated BEFORE any graph capture so no captured scratch can
7995 // alias it): every path that updates the round seed copies INTO it — no per-round allocs,
7996 // stable pointer for the graph-draft round-start copy.
7997 let mut h_seed_buf = e.clone_dtod(&h_seed0)?;
7998 // Predecessor-pairing trackers: `fill_prev` = trunk hidden AT the last COMMITTED row (the
7999 // predecessor of the next verify's col 0 — the reference's carried pending-h analogue;
8000 // also the predecessor-row hidden for the round-0 legacy-replay seed). At round 0 that
8001 // row is last_token's own (h_seed0). The chain step-0 seed under the pairing default =
8002 // hidden of the row BEFORE last_token = the prompt's last row at round 0 (h_seed_buf
8003 // overwritten below).
8004 let mut fill_prev = e.clone_dtod(&h_seed0)?;
8005 {
8006 if let Some(ph) = &prompt_h {
8007 let np = prompt.len();
8008 e.copy_view_into(
8009 &mut h_seed_buf,
8010 0,
8011 &ph.slice((np - 1) * n_embd..np * n_embd),
8012 n_embd,
8013 )?;
8014 } else if continuation {
8015 if let Some((_, lh, _, _, _)) = sess_tail.as_ref() {
8016 if let Some(lh) = lh.as_ref() {
8017 e.copy_into(&mut h_seed_buf, 0, lh, n_embd)?;
8018 }
8019 }
8020 }
8021 }
8022 // Persistent device prediction slots for the accept walk (max k+1 verify columns).
8023 let mut preds_d = e.alloc_u32_zeroed(k + 2)?;
8024
8025 let debug_spec = std::env::var("MEMRA_DEBUG_SPEC").is_ok();
8026 let fork_mode = OptiForkGateMode::configured();
8027 // MEMRA_SPEC_STATS=1: per-slot accept histogram + draft-length histogram, printed once at
8028 // the end. Metric normalization vs the reference engine: BOTH engines count
8029 // accepted/drafted where the chain stopped at p-min and the sub-threshold token is
8030 // discarded uncounted — per-slot decay + chain-length mix are the extra dimensions.
8031 let spec_stats = std::env::var("MEMRA_SPEC_STATS").is_ok();
8032 let mut st_drafted = vec![0usize; k];
8033 let mut st_accepted = vec![0usize; k];
8034 let mut st_len_hist = vec![0usize; k + 1];
8035 let mut st_full = 0usize;
8036 // P-MIN CONFIDENCE GATE (MEMRA_SPEC_PMIN, the serve script's --spec-draft-p-min mechanism):
8037 // stop the draft chain early when the head's softmax confidence in its own pick drops
8038 // below p_min. Hoisted above the loop: the graph capture bakes the prob kernels iff on.
8039 static PMIN: std::sync::OnceLock<f32> = std::sync::OnceLock::new();
8040 let p_min = *PMIN.get_or_init(|| {
8041 std::env::var("MEMRA_SPEC_PMIN")
8042 .ok()
8043 .and_then(|v| v.parse().ok())
8044 .unwrap_or(0.0)
8045 });
8046 // ZERO-DRAFT ROUNDS (MEMRA_SPEC_PMIN0=1, vendored from llama.cpp's draft gating): let the
8047 // p-min gate apply at j==0 too, so a low-confidence round drafts NOTHING and the verify
8048 // batch is just the pending bonus (m=1 = a plain decode step). llama's 35B win rides
8049 // exactly this — draft acceptance 76% at mean len 2.5 because unpredictable stretches
8050 // never pay draft+verify overhead. Only legal when a pending bonus exists (an empty
8051 // verify batch is not); the j==0 exemption stays for pending-less rounds.
8052 let pmin0 = std::env::var("MEMRA_SPEC_PMIN0")
8053 .map(|v| v == "1")
8054 .unwrap_or(false);
8055
8056 // --- GRAPH DRAFT setup: persistent I/O buffers + ONE capture (2 warmups inside). The
8057 // warmups mutate scratch len_d / pos / tok / seed — all reset at every round start, so the
8058 // only restore needed is the scratch counter. Capture failure (e.g. a non-capturable
8059 // cuBLAS path in an exotic head) falls back to the eager draft chain.
8060 // PER-SESSION PERSISTENCE (2026-08-01): session calls reuse the DraftGraphCtx parked on
8061 // the SpecSession — the capture (2 warmup head forwards + instantiate) ran ONCE at the
8062 // session's first burst, not per burst (measured ~16ms/burst fixed cost on H100 q27,
8063 // research/spec-serving-20260801). Reuse is pointer-exact: the graph bakes the session's
8064 // own scratch KV (never realloc'd), the model's resident embedding, the OnceLock p_min,
8065 // and the g_* buffers carried in the ctx — replay dispatch is identical to a fresh
8066 // capture, so draft tokens are bit-identical (drafts never decide exactness anyway; the
8067 // verify arbitrates). Single-shot calls (sess=None) build a fresh ctx and drop it.
8068 let mut dctx: DraftGraphCtx = match sess_draft_slot.as_mut().and_then(|s| s.take()) {
8069 Some(c) => c,
8070 None => DraftGraphCtx::new(e, n_embd, if sampled { d_vocab } else { 1 })?,
8071 };
8072 // A session that ran greedy bursts first sized g_q/g_perturb at 1; a sampled resume
8073 // needs d_vocab. Realloc is legal exactly while graph_s is None (nothing baked them).
8074 if sampled && dctx.g_q.len() < d_vocab {
8075 dctx.g_q = e.zeros(d_vocab)?;
8076 dctx.g_perturb = e.zeros(d_vocab)?;
8077 }
8078 // DRAFT-SIDE GRAMMAR MASK (lane/draft-mask, 2026-08-04): the drafter samples the
8079 // grammar's legal set, so proposals are legal BY CONSTRUCTION and the verify-side
8080 // truncation (the correctness backstop) stops cutting every tight-schema round.
8081 // The mask is one node inside the captured draft chain — presence is a CAPTURE-TIME
8082 // shape, so a parked graph of the other shape is dropped and recaptured.
8083 let dmask_on = constraint
8084 .as_deref()
8085 .is_some_and(|c| c.draft_mask_enabled());
8086 let dmask_words = if dmask_on { d_vocab.div_ceil(32) } else { 0 };
8087 if dmask_on && dctx.g_dmask.len() < dmask_words {
8088 dctx.g_dmask = e.alloc_u32_zeroed(dmask_words)?;
8089 dctx.graph = None; // the old capture baked the old (or no) mask pointer
8090 dctx.failed.clear_greedy();
8091 dctx.keeper.clear();
8092 }
8093 if dctx.graph.is_some() && dctx.graph_masked != dmask_on {
8094 dctx.graph = None;
8095 dctx.failed.clear_greedy();
8096 dctx.keeper.clear();
8097 }
8098 if graph_draft && !sampled && dctx.graph.is_none() && !dctx.failed.greedy_failed() {
8099 let DraftGraphCtx {
8100 g_tok,
8101 g_pos,
8102 g_seed,
8103 g_p,
8104 g_dmask,
8105 ..
8106 } = &mut dctx;
8107 // capture-time contents: ALL-ONES (ban nothing). A replay only ever runs after the
8108 // host uploads the position's real words, so the warmups stay grammar-free.
8109 if dmask_on {
8110 e.htod_u32_into(g_dmask, &vec![u32::MAX; dmask_words])?;
8111 }
8112 let g_dmask_ro: &CudaSlice<u32> = &*g_dmask;
8113 // CAPTURE-RETAIN (#68 fix): the warmup transients' pool addresses are baked into the
8114 // captured graph; the keeper pins them for the graph's lifetime. capture_graph (non-
8115 // retained) freed them at exit — safe for one-shot generate_spec (nothing else touches
8116 // the pool between replays) but WRONG for sessions: burst-boundary prime/fill/commit
8117 // passes (and, in serve, other sessions) recycle those addresses and the replay then
8118 // clobbers live buffers — the ST serve-spec corruption (research/serve-st-20260803).
8119 let cap_res = e.capture_graph_retained(|e| {
8120 self.mtp_head_forward_cap(
8121 e,
8122 mtp,
8123 g_tok,
8124 g_pos,
8125 g_seed,
8126 g_p,
8127 &mut *scratch,
8128 p_min > 0.0 || fork_mode == OptiForkGateMode::Controller,
8129 true,
8130 embd_gpu.expect("graph draft requires resident embedding"),
8131 embd_qt,
8132 embd_rb,
8133 d_vocab,
8134 None,
8135 None,
8136 if dmask_on {
8137 Some((g_dmask_ro, dmask_words))
8138 } else {
8139 None
8140 },
8141 )
8142 });
8143 match cap_res {
8144 Ok((g, keep)) => {
8145 scratch.set_len(e, base)?;
8146 dctx.graph = Some(g);
8147 dctx.graph_masked = dmask_on;
8148 dctx.keeper = keep;
8149 }
8150 Err(err) => {
8151 scratch.set_len(e, base)?;
8152 // LOUD flip (audit Q2): a dropped draft graph is a coverage loss, never
8153 // silent. Once per flip — mark returns None on an already-failed ctx.
8154 if let Some(line) = dctx.failed.mark_greedy(&err.to_string()) {
8155 eprintln!("{line}");
8156 }
8157 }
8158 }
8159 }
8160 // --- SAMPLED GRAPH DRAFT setup (step 3 of the sampled-spec arc): a SECOND capture, own
8161 // graph object, built only when sampled && graph-eligible — the greedy capture above is
8162 // untouched (and skipped when sampled: its graph would never be launched). Same head
8163 // forward, but the in-graph argmax reads GUMBEL-PERTURBED logits; the Philox event
8164 // counter lives in the persistent device g_ctr (bumped in-graph, host-seeded from sctr
8165 // once per round); the raw head logits land in the persistent g_q for the host's
8166 // per-replay async D2D into the round's q slot (q_slots, K x d_vocab, allocated once).
8167 // seed/temp are capture-time constants — baked into graph_s, so a pool-resumed request
8168 // with a different (seed, temp, k) drops the parked sampled graph and recaptures.
8169 // COST OF THE FRESH-SEED SERVE DEFAULT (dogfood F4, 2026-08-04): omitting `seed` on a
8170 // serve request now draws fresh per-request entropy (it used to default to a pinned 0),
8171 // so a seed-omitting request that RESUMES a parked spec session finds an s_key baked
8172 // with the PREVIOUS request's seed and pays one recapture. Bounded, and it does not
8173 // reopen the ~16ms/burst regression the persistent ctx exists to fix: a session's seed
8174 // is fixed for its whole lifetime (worker.rs reads s.sampler.seed() per burst), so
8175 // this compare misses at most ONCE per resumed request — the first burst recaptures
8176 // and every later burst in that request replays. A client that wants the parked graph
8177 // AND reproducibility supplies an explicit `seed`, honored exactly, which keeps s_key
8178 // stable across its whole conversation.
8179 // COMPOSITION RULE (fspec x gsd merge): the in-graph chain samples from the RAW
8180 // softmax — it can hold neither per-row filter stats nor the varying penalty history.
8181 // The sampled graph therefore engages only in the PURE-TEMP regime; filters/penalties
8182 // force the eager draft (which computes stats/penalties per row).
8183 // KEY THE WHOLE REGIME, not just the baked constants (lane/graph-s-key-exactness-
8184 // 20260819). `s_key` used to be `(seed, temp, k)`; the filters and penalties were left
8185 // out, so a filtered request resuming a session that parked a PURE-TEMP graph kept it —
8186 // and the launch site never re-asked `pure_temp`. See [`SampledGraphKey`] for what that
8187 // costs (an unconditional accept of out-of-head draft tokens, i.e. an exactness bug on
8188 // the request shape the vendor-default flip makes the majority).
8189 let s_key = SampledGraphKey::new(sp_seed, sp_temp, k, sp.top_k, sp.top_p, sp.min_p, pen_on);
8190 let pure_temp = s_key.pure_temp();
8191 if sampled && dctx.s_key.is_some_and(|old| old != s_key) {
8192 dctx.graph_s = None;
8193 dctx.failed.clear_sampled();
8194 dctx.s_key = None;
8195 dctx.q_slots.clear();
8196 dctx.keeper_s.clear();
8197 }
8198 if graph_draft
8199 && sampled
8200 && pure_temp
8201 && dctx.graph_s.is_none()
8202 && !dctx.failed.sampled_failed()
8203 {
8204 let DraftGraphCtx {
8205 g_tok,
8206 g_pos,
8207 g_seed,
8208 g_p,
8209 g_ctr,
8210 g_perturb,
8211 g_q,
8212 ..
8213 } = &mut dctx;
8214 // CAPTURE-RETAIN (#68 fix): same keeper contract as the greedy capture above.
8215 let cap_res = e.capture_graph_retained(|e| {
8216 self.mtp_head_forward_cap(
8217 e,
8218 mtp,
8219 g_tok,
8220 g_pos,
8221 g_seed,
8222 g_p,
8223 &mut *scratch,
8224 p_min > 0.0,
8225 true,
8226 embd_gpu.expect("graph draft requires resident embedding"),
8227 embd_qt,
8228 embd_rb,
8229 d_vocab,
8230 Some((g_ctr, g_perturb, g_q, sp_seed, sp_temp)),
8231 None,
8232 None, // constrained spec is greedy-only — sampled never carries a hook
8233 )
8234 });
8235 match cap_res {
8236 Ok((g, keep)) => {
8237 scratch.set_len(e, base)?;
8238 for _ in 0..k {
8239 dctx.q_slots.push(e.zeros(d_vocab)?);
8240 }
8241 dctx.graph_s = Some(g);
8242 dctx.s_key = Some(s_key);
8243 dctx.keeper_s = keep;
8244 }
8245 Err(err) => {
8246 scratch.set_len(e, base)?;
8247 // LOUD flip (audit Q2): same contract as the greedy capture above.
8248 if let Some(line) = dctx.failed.mark_sampled(&err.to_string()) {
8249 eprintln!("{line}");
8250 }
8251 }
8252 }
8253 }
8254 // ---- EXACTNESS GUARD, the enforceable half (lane/graph-s-key-exactness-20260819) ----
8255 // With the filters and penalties in `s_key`, a graph that SURVIVED the drop above was
8256 // captured under this request's exact regime, and capture requires `pure_temp` — so a
8257 // parked graph implies `pure_temp`. That implication is the whole exactness argument for
8258 // the graph arm, so it is asserted here rather than assumed: a future change that widens
8259 // the capture condition, narrows the key, or copies a `DraftGraphCtx` across regimes
8260 // fails LOUDLY at this line instead of silently drafting from the raw softmax while the
8261 // verify applies filter stats. Release builds refuse the graph (drop it, draft eager)
8262 // rather than launching it; the launch site re-tests `pure_temp` independently.
8263 if sampled && !pure_temp && dctx.graph_s.is_some() {
8264 debug_assert!(
8265 false,
8266 "sampled draft graph parked under {:?} survived into a FILTERED request \
8267 (top_k={} top_p={} min_p={} pen_on={}): the in-graph chain draws from the RAW \
8268 softmax, so the verify's filtered q would test a distribution the draft was \
8269 never sampled from",
8270 dctx.s_key, sp.top_k, sp.top_p, sp.min_p, pen_on,
8271 );
8272 eprintln!(
8273 "[spec] BUG: dropping a parked sampled draft graph that outlived its capture \
8274 regime (s_key={:?}, request top_k={} top_p={} min_p={} pen_on={}); drafting \
8275 EAGER — the key must carry every field that shapes q",
8276 dctx.s_key, sp.top_k, sp.top_p, sp.min_p, pen_on,
8277 );
8278 dctx.graph_s = None;
8279 dctx.s_key = None;
8280 dctx.q_slots.clear();
8281 dctx.keeper_s.clear();
8282 }
8283 // SKEY PROBE (MEMRA_SKEY_PROBE=1): the burst-entry facts the reachability question turns
8284 // on — is this request sampled, is it in the pure-temp regime the sampled graph is only
8285 // legal in, and is a graph PARKED from an earlier request of the same session? The launch
8286 // arms below print which chain actually ran, so the probe never restates the condition.
8287 if skey_probe() {
8288 eprintln!(
8289 "[skey] burst sampled={} pure_temp={} temp={} top_k={} top_p={} min_p={} \
8290 pen_on={} k={} graph_draft={} graph_s_parked={} s_key_parked={:?}",
8291 sampled as u8,
8292 pure_temp as u8,
8293 sp_temp,
8294 sp.top_k,
8295 sp.top_p,
8296 sp.min_p,
8297 pen_on as u8,
8298 k,
8299 graph_draft as u8,
8300 dctx.graph_s.is_some() as u8,
8301 dctx.s_key,
8302 );
8303 }
8304 let t_cap = t_ent.elapsed();
8305 // PERSISTENT DRAFT KV: fill the MTP block's K/V for every prompt position from the exact
8306 // trunk hiddens collected during prime — ONE batched K/V-only pass (overwrites any
8307 // capture-warmup garbage; capture left len at 0). last_token (the init feed) needs no
8308 // fill: the first chain step processes it and appends its entry at slot prompt.len().
8309 if let Some(ph) = &prompt_h {
8310 // SESSION: rows [0..base) are the previous turns' exact fills (refresh overwrote them
8311 // with true verify hiddens) — truncate any draft overhang, fill ONLY the suffix at
8312 // global positions [base..base+tp). Fresh call: base==0, identical to before.
8313 scratch.set_len(e, base)?;
8314 // CHUNKED FILL (long-ctx OOM fix, 2026-07-05): mtp_kv_fill's transients scale with its
8315 // T (concat = T*2*n_embd*4B — 1.5GB at 40k) and its concat loop is 2*T launches. The
8316 // fill is a pure sequential append, so chunking is exact: each chunk appends its rows
8317 // at pos0=base+start with the identical per-row math. Same knob as the trunk prime.
8318 let tp = prompt.len();
8319 let fill_chunk: usize = if crate::cache::swa_ring_on() {
8320 crate::hybrid_forward::prime_chunk_tokens(tp, self.layers.len())
8321 } else {
8322 // Preserve the flag-OFF schedule byte-for-byte, including the legacy zero value
8323 // meaning one monolithic fill.
8324 std::env::var("MEMRA_PRIME_CHUNK")
8325 .ok()
8326 .and_then(|v| v.parse().ok())
8327 .unwrap_or(4096)
8328 };
8329 let fill_chunk = if fill_chunk == 0 { tp } else { fill_chunk };
8330 let mut start = 0usize;
8331 while start < tp {
8332 let end = (start + fill_chunk).min(tp);
8333 let tc = end - start;
8334 {
8335 // PREDECESSOR pairing: row i gets h[i-1]; global row 0 a zeros row (the
8336 // reference engine's initial pending-h is zeroed too); a session turn's row 0
8337 // gets the PREVIOUS turn's last committed hidden (sess.last_h). Per chunk:
8338 // rows start..end read h[start-1..end-1] — one dtod into a chunk buffer.
8339 let mut phs = e.zeros(tc * n_embd)?;
8340 let (src_lo, dst_off) = if start == 0 {
8341 (0, n_embd)
8342 } else {
8343 ((start - 1) * n_embd, 0)
8344 };
8345 let n_copy = if start == 0 {
8346 (tc - 1) * n_embd
8347 } else {
8348 tc * n_embd
8349 };
8350 if start == 0 {
8351 if let Some((_, lh, _, _, _)) = sess_tail.as_ref() {
8352 if let Some(lh) = lh.as_ref() {
8353 e.copy_into(&mut phs, 0, lh, n_embd)?;
8354 }
8355 }
8356 }
8357 if n_copy > 0 {
8358 e.copy_view_into(
8359 &mut phs,
8360 dst_off,
8361 &ph.slice(src_lo..src_lo + n_copy),
8362 n_copy,
8363 )?;
8364 }
8365 self.mtp_kv_fill(
8366 e,
8367 mtp,
8368 &prompt[start..end],
8369 &phs,
8370 base + start,
8371 &mut *scratch,
8372 embd_dev,
8373 )?;
8374 }
8375 start = end;
8376 }
8377 }
8378 // MEMRA_PROFILE_SPEC=2: profiler capture starts HERE — after the prime, so an
8379 // `nsys -c cudaProfilerApi` capture contains ONLY the round loop (draft/verify/commit).
8380 // (=1 brackets the whole call in run_spec.rs, prime included.)
8381 if std::env::var("MEMRA_PROFILE_SPEC").as_deref() == Ok("2") {
8382 unsafe extern "C" {
8383 fn cudaProfilerStart() -> i32;
8384 }
8385 unsafe {
8386 cudaProfilerStart();
8387 }
8388 }
8389 // ROUND-STREAM stage (c) 4 (MEMRA_SPEC_STREAM=1, experimental): pre-issued M-round
8390 // bursts with ZERO per-round host readbacks — the accept/seed/rollback/ring kernels
8391 // consume each other's device outputs; the host drains the ring every M rounds. v1
8392 // constraints: greedy, !spec_replay, single-shot, batched-linear layers, no refresh
8393 // fills (acceptance effect A/B-arbitrated), enters from round 1 (pending guaranteed).
8394 // NOTE: not gated on the caller's graph_draft (its trunk_dense conjunct turns the 35B
8395 // MoE off) — the stream capture encloses ONLY the dense MTP head; the head-dense /
8396 // full-prec / k gates are re-derived here and a failed capture degrades to stream-off.
8397 let stream_on = crate::spec::spec_stream()
8398 && !sampled
8399 && !spec_replay
8400 && constraint.is_none()
8401 && !session_mode
8402 && embd_gpu.is_some()
8403 && !crate::model::full_prec_enabled()
8404 && k + 2 < 96;
8405 let mut stream_graph: Option<cudarc::driver::CudaGraph> = None;
8406 let mut g_tokp2k = e.alloc_u32_zeroed(2 * k.max(1))?;
8407 if stream_on {
8408 let cap = e.capture_graph(|e| {
8409 for j in 0..k.max(1) {
8410 self.mtp_head_forward_cap(
8411 e,
8412 mtp,
8413 &mut dctx.g_tok,
8414 &mut dctx.g_pos,
8415 &mut dctx.g_seed,
8416 &mut dctx.g_p,
8417 &mut *scratch,
8418 true,
8419 true,
8420 embd_gpu.expect("round stream requires resident embedding"),
8421 embd_qt,
8422 embd_rb,
8423 d_vocab,
8424 None,
8425 Some((&mut g_tokp2k, j, d2t_dev.as_ref())),
8426 None, // round-stream requires constraint.is_none() (see stream_on)
8427 )?;
8428 }
8429 Ok(())
8430 });
8431 match cap {
8432 Ok(g) => {
8433 scratch.set_len(e, 0)?;
8434 stream_graph = Some(g);
8435 }
8436 Err(err) => {
8437 scratch.set_len(e, 0)?;
8438 if debug_spec {
8439 eprintln!("[spec] stream-graph capture failed ({err}); stream off");
8440 }
8441 }
8442 }
8443 }
8444 let stream_active = stream_on && stream_graph.is_some();
8445 if debug_spec {
8446 eprintln!(
8447 "[spec] stream_on={stream_on} env={} samp={sampled} dg={} captured={} active={stream_active} session={session_mode} replay={spec_replay}",
8448 crate::spec::spec_stream(),
8449 dctx.graph.is_some(),
8450 stream_graph.is_some()
8451 );
8452 }
8453 let t_v_s = k + 1;
8454 // ROUND-STREAM buffers + ptr tables now live in the model-generic round_stream
8455 // module (extracted 2026-07-12; the gemma burst reuses them).
8456 let sb = crate::round_stream::StreamBufs::new(e, k, crate::spec::spec_stream_m())?;
8457 let crate::round_stream::StreamBufs {
8458 mut vtok_d,
8459 mut brk_d,
8460 mut pend_d,
8461 last_pred_d,
8462 mut pos_ctr,
8463 mut pos_start_d,
8464 mut ring_d,
8465 acc_d: mut stream_acc,
8466 m_rounds,
8467 k: _,
8468 } = sb;
8469 let stream_ptrs: Option<CudaSlice<u64>> = if stream_active {
8470 Some(crate::round_stream::kv_len_ptr_table(
8471 e,
8472 cache,
8473 Some(&pos_ctr),
8474 )?)
8475 } else {
8476 None
8477 };
8478
8479 let t_fill = t_ent.elapsed();
8480 let mut round = 0usize;
8481 // ADAPTIVE DRAFT LENGTH (MEMRA_SPEC_ADAPT=1, opt-in — the gemma_spec accepted-run law,
8482 // ported 2026-08-01): next round's draft depth = last round's accepted run + 1, clamped
8483 // to [floor(pos), k_cap] — a miss shrinks the next draft to the miss point + 1,
8484 // full-accept streaks re-deepen one step per round. NOT the 2026-07-07 acceptance-EMA
8485 // (that arm measured an HONEST LOSS to static per-class optima — 115.0/85.8/73.4 vs
8486 // 121.6/92.7/75.6, EMA lag — and was removed 2026-07-08; rig5090.jsonl has the record).
8487 // The gemma law has no lag class: it reacts within one round, and was worth +7-20% on
8488 // the gemma cells at unchanged exactness (2026-07-10 flip; floor sweep 2026-07-25;
8489 // position key 2026-07-26). Signal = n_acc from the round's EXISTING accept readback —
8490 // zero new syncs; the draft graph is a SINGLE-STEP capture replayed per drafted token,
8491 // so a per-round depth needs no re-capture (unlike gemma's whole-chain graphs). qwen's
8492 // in-round p-min cut already shortens chains mid-round, so gemma's one-round-late p-min
8493 // fold into kc is unnecessary here — the accepted-run law sees the cut via n_acc.
8494 // Exactness is the verify's job at ANY depth (same contract as p-min variable rounds).
8495 // DEFAULT OFF on the qwen path until its cells gate a flip (gemma's is default-on).
8496 // MEASURED 2026-08-01 (H100 GPU-3, interleaved x3, NGEN=256, same-invocation plain
8497 // denominators; research/qwen-adaptive-k-20260801/): REFUTED on the tuned qwen configs.
8498 // q27 K=3+HPOST+PMIN=0.3: short +0.8% (noise; law ~idles, len_hist identical), board
8499 // -1.9%, agentic -0.5%; board PMIN=0 -2.1% (not p-min shadowing — the law itself);
8500 // floor=1 -2.8% (gemma's floor-collapse, reproduced). q35 K=2 board: -6.4% (52/136
8501 // rounds shrink to depth 1; no depth to reclaim at K=2). The gemma direction DOES
8502 // appear at untuned depth-K — q27 K=6 floor=4 +1.5% over fixed K=6 — but stays -3.7%
8503 // below fixed K=3: same verdict class as the retired EMA arm (honest loss to static
8504 // per-class optima). Acceptance-rate rises under the law while tokens/round falls —
8505 // it buys accept-% by adding rounds, and a round's fixed draft+verify cost wins.
8506 // K=1..8 self-consistency PASS both models with the law ON (exactness held).
8507 let adapt = std::env::var("MEMRA_SPEC_ADAPT").as_deref() == Ok("1");
8508 // floor: per-model default keyed on n_embd (gemma's tiering — models with an expensive
8509 // verify keep deep drafts after a miss); MEMRA_SPEC_ADAPT_FLOOR pins it everywhere.
8510 let adapt_floor_env: Option<usize> = std::env::var("MEMRA_SPEC_ADAPT_FLOOR")
8511 .ok()
8512 .and_then(|v| v.parse().ok());
8513 let adapt_floor_default: usize = if self.cfg.n_embd as usize >= 3500 {
8514 4
8515 } else if self.cfg.n_embd as usize >= 2500 {
8516 2
8517 } else {
8518 1
8519 };
8520 let adapt_floor: usize = adapt_floor_env.unwrap_or(adapt_floor_default);
8521 // position key: past floor_ctx a HIGH floor (>=4) relaxes to 1 — forced-deep drafts
8522 // turn net-negative at depth (gemma 31B d1736 evidence); MEMRA_SPEC_FLOOR_CTX moves
8523 // the boundary, an explicit MEMRA_SPEC_ADAPT_FLOOR pins the floor everywhere.
8524 let floor_ctx: usize = std::env::var("MEMRA_SPEC_FLOOR_CTX")
8525 .ok()
8526 .and_then(|v| v.parse().ok())
8527 .unwrap_or(1024);
8528 let floor_at = |pos: usize| -> usize {
8529 if adapt_floor_env.is_some() || pos < floor_ctx {
8530 adapt_floor
8531 } else if adapt_floor >= 4 {
8532 1
8533 } else {
8534 adapt_floor
8535 }
8536 };
8537 // cap: MEMRA_SPEC_CAPMAX (gemma semantics, default 7). Binds only under adapt — the
8538 // fixed-K default path is untouched by this whole block.
8539 let cap_max: usize = std::env::var("MEMRA_SPEC_CAPMAX")
8540 .ok()
8541 .and_then(|v| v.parse().ok())
8542 .unwrap_or(7);
8543 let k_cap = k.min(cap_max).max(1);
8544 let mut kc = k_cap;
8545 let mut opti_fork: Option<OptiForkState> = None;
8546 let mut fork_snapshot: Option<crate::cache::CacheSnapshot> = None;
8547 if fork_mode != OptiForkGateMode::Disabled {
8548 let fence = crate::pp::pp_cuts(self.layers.len());
8549 let refusal = if !session_mode {
8550 Some("not-session")
8551 } else if k != 1 || adapt {
8552 Some("requires-fixed-k1")
8553 } else if sampled || constraint.is_some() || spec_replay {
8554 Some("sampled-constrained-or-replay")
8555 } else if pipe.is_some() {
8556 Some("two-session-pipeline")
8557 } else if !spec_devacc() {
8558 Some("requires-device-accept")
8559 } else if stream_active || crate::spec::spec_stream() {
8560 Some("round-stream")
8561 } else if crate::cache::swa_ring_on() || cache.has_swa_ring() {
8562 Some("swa-ring")
8563 } else if crate::pp::pp_host_bounce_active() {
8564 Some("host-bounce")
8565 } else if fork_mode == OptiForkGateMode::Controller
8566 && cache.recur.iter().any(Option::is_some)
8567 {
8568 Some("controller-requires-zero-recurrent-state")
8569 } else if fence.as_ref().is_none_or(|f| f.len() != 3) {
8570 Some("requires-pp2")
8571 } else {
8572 None
8573 };
8574 if let Some(reason) = refusal {
8575 OPTI_FORK_REFUSALS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
8576 eprintln!("[opti-fork] refused reason={reason}");
8577 } else {
8578 let fence = fence.expect("validated PP-2 fence");
8579 let rt = crate::pp::PpNRt::get(e)?;
8580 let primary_stage0 = rt.engine(0, e).ctx().ordinal() == e.ctx().ordinal();
8581 let primary_stage1 = rt.engine(1, e).ctx().ordinal() == e.ctx().ordinal();
8582 let primary_supported =
8583 primary_stage0 || (fork_mode == OptiForkGateMode::Controller && primary_stage1);
8584 if !rt.cross_device() || !primary_supported {
8585 OPTI_FORK_REFUSALS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
8586 eprintln!("[opti-fork] refused reason=requires-supported-primary-cross-device");
8587 } else {
8588 // Both recurrent snapshots and both seed generations are allocated before
8589 // the first fork, each through its owning PP stage. Allocation failure
8590 // therefore happens before any optimistic state mutation can occur.
8591 let current_snapshot = opti_snapshot_stage_owned(e, cache, rt, &fence)?;
8592 let alternate_snapshot = opti_snapshot_stage_owned(e, cache, rt, &fence)?;
8593 let fork = OptiForkState::new(
8594 e,
8595 cache,
8596 fork_mode,
8597 alternate_snapshot,
8598 &h_seed_buf,
8599 &fill_prev,
8600 rt,
8601 fence[1],
8602 self.layers.len(),
8603 )?;
8604 eprintln!(
8605 "[opti-fork] armed mode={fork_mode:?} snapshots=2 seeds=2 split={} \
8606 payload_dev0={} payload_dev1={} q_threshold={:.3}",
8607 fence[1],
8608 fork.logical_payload_bytes[0],
8609 fork.logical_payload_bytes[1],
8610 fork.controller.map_or(0.0, |policy| policy.threshold),
8611 );
8612 fork_snapshot = Some(current_snapshot);
8613 opti_fork = Some(fork);
8614 }
8615 }
8616 }
8617 // Persistent snapshot buffers are allocated once and refreshed in place. The fork arm
8618 // uses stage-owned snapshots; refused/disabled arms retain the existing generic helper.
8619 let mut snap = match fork_snapshot {
8620 Some(snapshot) => snapshot,
8621 None => cache.snapshot(e)?,
8622 };
8623 let mut carried_opti: Option<OptiControllerTicket> = None;
8624 // ROUND-STREAM stage (b) 3a: device table of per-layer kvl.len_d pointers (stable — the
8625 // cache never reallocates len_d; see cache.rs "stable pointer" note). 0 = no KV layer.
8626 let kv_len_ptrs: Option<CudaSlice<u64>> = if spec_devacc() && !spec_replay {
8627 Some(crate::round_stream::kv_len_ptr_table(e, cache, None)?)
8628 } else {
8629 None
8630 };
8631 // BONUS FOLD (2026-07-04): after a FULL accept the bonus token is NOT committed with a
8632 // separate T=1 trunk pass (a full weight read per round). It stays PENDING and rides as
8633 // column 0 of the NEXT round's verify batch. Under predecessor pairing the next chain
8634 // seeds from the bonus's predecessor's TRUE verify hidden (free — no extra
8635 // pass of any kind). Verify still
8636 // checks every emitted token against the target -> exactness holds by construction; only
8637 // DRAFT QUALITY can shift, which the acceptance numbers arbitrate.
8638 // bonus emitted but not yet committed to cache. A carried pending (see SpecSession::
8639 // pending_tok) enters round 0 directly — the burst boundary becomes a plain round edge.
8640 let mut pending: Option<u32> = carried_pending;
8641 // MEMRA_SPEC_PHASE=1: per-round wall decomposition (draft / verify / accept+commit) —
8642 // no tracing, no extra syncs (each phase is naturally sync-bounded: draft readbacks,
8643 // the verify accept readback). Printed once at loop end via spec-stats.
8644 let anatomy_on = std::env::var("MEMRA_SPEC_PP_ANATOMY").as_deref() == Ok("1");
8645 let phase_on = anatomy_on || std::env::var("MEMRA_SPEC_PHASE").as_deref() == Ok("1");
8646 // DRAFT-MASK receipt (lane/draft-mask): speculative-clone wall + rounds, printed with
8647 // spec-stats. The clone is the one cost the design adds per round — measured, not assumed.
8648 let (mut dm_clone_ns, mut dm_rounds) = (0u128, 0usize);
8649 // grammar-truncation counters: how many rounds the verify-side cut fired and how many
8650 // already-verified tokens it threw away. THIS is the quantity draft masking targets.
8651 let (mut dm_cuts, mut dm_cut_tokens) = (0usize, 0usize);
8652 let (mut ph_draft, mut ph_verify, mut ph_rest) = (0f64, 0f64, 0f64);
8653 let mut ph_wait = 0f64;
8654 let mut ph_commit = 0f64;
8655 let mut ph_t = std::time::Instant::now();
8656 let mut ph_mark = |acc: &mut f64, on: bool| {
8657 if on {
8658 let now = std::time::Instant::now();
8659 *acc += (now - ph_t).as_secs_f64();
8660 ph_t = now;
8661 }
8662 };
8663 if let Some(p) = pipe {
8664 p.setup_end();
8665 }
8666 while keep_going && out.len() < max_new {
8667 // ROUND-STREAM BURST: from round 1 (pending guaranteed by every non-replay arm),
8668 // issue M rounds with zero readbacks, then drain the ring + reconcile mirrors.
8669 if let (true, Some(sg), Some(ptrs)) = (
8670 stream_active && round >= 1 && pending.is_some(),
8671 &stream_graph,
8672 &stream_ptrs,
8673 ) {
8674 if debug_spec {
8675 static ONCE: std::sync::Once = std::sync::Once::new();
8676 ONCE.call_once(|| {
8677 eprintln!("[memra] ROUND-STREAM burst engaged (M={m_rounds} k={k})")
8678 });
8679 }
8680 e.set_i32_one(&mut pos_ctr, cache.pos as i32)?;
8681 e.set_u32_one(&mut pend_d, pending.unwrap())?;
8682 e.set_u32_one(&mut ring_d, 0)?; // ring count = 0 (writes element 0)
8683 for _mi in 0..m_rounds {
8684 e.i32_copy_add(&pos_ctr, &mut pos_start_d, 0)?;
8685 cache.snapshot_into(e, &mut snap)?; // device D2Ds, stream-ordered
8686 e.i32_copy_add(&pos_ctr, &mut scratch.kv.len_d, 0)?; // draft-KV rollback
8687 e.i32_copy_add(&pos_ctr, &mut dctx.g_pos, 1)?; // rope pos = pos + base
8688 e.u32_copy(&pend_d, &mut dctx.g_tok)?;
8689 e.copy_into(&mut dctx.g_seed, 0, &h_seed_buf, n_embd)?;
8690 sg.launch()?;
8691 e.spec_assemble_verify(
8692 &g_tokp2k,
8693 &pend_d,
8694 d2t_dev.as_ref(),
8695 &mut vtok_d,
8696 &mut brk_d,
8697 p_min,
8698 k,
8699 pmin0,
8700 )?;
8701 let mut ck = VerifyCkpt::new(self.layers.len());
8702 let dummy = vec![0u32; t_v_s];
8703 let (tl_d, vx) = self.decode_step_t_core_stream(
8704 e,
8705 &dummy,
8706 0,
8707 &mut *cache,
8708 embd_dev,
8709 Some(&mut ck),
8710 Some((&vtok_d, &pos_ctr)),
8711 None,
8712 )?;
8713 for j in 0..t_v_s {
8714 e.argmax_token_device_col(&tl_d, j, n_vocab, &mut preds_d, j)?;
8715 }
8716 e.spec_accept_greedy_dc(
8717 &preds_d,
8718 &vtok_d,
8719 &last_pred_d,
8720 &brk_d,
8721 &mut stream_acc,
8722 )?;
8723 e.spec_seed_gather(&vx, &fill_prev, &stream_acc, &mut h_seed_buf, 1, n_embd)?;
8724 e.copy_into(&mut fill_prev, 0, &h_seed_buf, n_embd)?;
8725 self.commit_verified_prefix_stream(
8726 e,
8727 &mut *cache,
8728 &snap,
8729 &ck,
8730 &stream_acc,
8731 1,
8732 t_v_s,
8733 )?;
8734 e.spec_rollback_stream(
8735 ptrs,
8736 &pos_start_d,
8737 &stream_acc,
8738 1,
8739 self.layers.len() + 1,
8740 )?;
8741 e.spec_ring_commit(&vtok_d, &stream_acc, &brk_d, &mut ring_d, &mut pend_d)?;
8742 }
8743 e.stream().synchronize()?;
8744 let ring_h = e.dtoh_u32(&ring_d)?;
8745 let cnt = ring_h[0] as usize;
8746 for i in 0..cnt {
8747 if out.len() < max_new {
8748 out.push(ring_h[1 + i]);
8749 }
8750 }
8751 let pos_h = e.dtoh_i32(&pos_ctr)?[0] as usize;
8752 for il in 0..self.layers.len() {
8753 if let Some(kvl) = cache.kv[il].as_mut() {
8754 kvl.len = pos_h;
8755 }
8756 }
8757 cache.pos = pos_h;
8758 scratch.kv.len = pos_h;
8759 pending = Some(ring_h[cnt]); // last drained token = the live bonus
8760 last_token = ring_h[cnt];
8761 total_drafted += k * m_rounds; // upper bound (p-min breaks uncounted)
8762 total_accepted += cnt.saturating_sub(m_rounds);
8763 if let Some(t) = sess_telem {
8764 // totals only — the burst's per-round accept counts stayed on device
8765 // (that is the point of the round-stream arm). pos_* untouched.
8766 t.record_totals(m_rounds, k * m_rounds, cnt.saturating_sub(m_rounds));
8767 }
8768 round += m_rounds;
8769 // sse-cadence: the drained ring is committed — flush it at burst-drain cadence.
8770 keep_going = flush_commit(&mut on_commit, &out, &mut flushed);
8771 continue;
8772 }
8773 let pipe_draft = match pipe {
8774 Some(p) => Some(p.draft_begin(round)?),
8775 None => None,
8776 };
8777 let pos = cache.pos; // #tokens committed (EXCLUDES a pending bonus)
8778 let mut current_opti = carried_opti.take();
8779 let mut fork_generation = if current_opti.is_none() && pending.is_some() {
8780 match opti_fork.as_mut() {
8781 Some(fork) if fork.mode.is_forced() => Some(fork.reserve(&mut snap)?),
8782 None => None,
8783 Some(_) => None,
8784 }
8785 } else {
8786 None
8787 };
8788 if current_opti.is_none() {
8789 if let Some(fork) = opti_fork.as_ref() {
8790 opti_snapshot_stage_owned_into(e, cache, fork.rt, &fork.fence, &mut snap)?;
8791 } else {
8792 cache.snapshot_into(e, &mut snap)?;
8793 }
8794 } else if snap.pos != pos {
8795 return Err(format!(
8796 "optipipe carried snapshot pos {} != current pos {pos}",
8797 snap.pos
8798 )
8799 .into());
8800 } // §C: snapshot BEFORE draft+verify (already retained for a carried successor)
8801 ph_mark(&mut ph_rest, phase_on);
8802
8803 // --- 1. DRAFT k tokens with the NextN head (autoregressive, T=1 each) ---
8804 // p-min semantics (both paths): stop the chain early when the head's confidence in
8805 // its own pick drops below p_min — the just-drafted token is DISCARDED, but its
8806 // scratch append stands (identical to the eager chain's ordering). j==0 always drafts.
8807 let base0 = if pending.is_some() { 1usize } else { 0usize };
8808 // fixed draft length by default; MEMRA_SPEC_ADAPT=1 drafts at last round's
8809 // accepted run + 1 (the gemma law — see the setup block above the loop).
8810 let k_this = if adapt { kc } else { k };
8811 let mut draft: Vec<u32> = Vec::with_capacity(k);
8812 let mut draft_idx: Vec<u32> = Vec::with_capacity(k); // trimmed-vocab ids (== draft when untrimmed)
8813 let mut controller_draft_prob: Option<f32> = None;
8814 let mut controller_eager_state: Option<(u32, CudaSlice<f32>)> = None;
8815 if let Some(ticket) = current_opti.as_mut() {
8816 let carried_pending = pending.ok_or("optipipe carried successor lost pending")?;
8817 if ticket.verify_tokens[0] != carried_pending {
8818 return Err(format!(
8819 "optipipe carried pending mismatch: ticket={} live={carried_pending}",
8820 ticket.verify_tokens[0],
8821 )
8822 .into());
8823 }
8824 draft.push(ticket.verify_tokens[1]);
8825 controller_draft_prob = Some(ticket.draft_prob);
8826 controller_eager_state = ticket
8827 .take_eager_seed()
8828 .map(|seed| (ticket.verify_tokens[1], seed));
8829 } else {
8830 // Round-start draft-KV sync (BOTH paths). Persistent: truncate/align to the committed
8831 // history — slots 0..P hold entries for the tokens before last_token@P (P = pos +
8832 // base0 - 1); this single set_len IS the draft-side rollback (drops last round's
8833 // rejected drafts and p-min extras via the len mechanism).
8834 scratch.set_len(e, pos + base0 - 1)?;
8835 if pen_on {
8836 // PEN_WINDOW_MAX also bounds the per-round upload and the O(n_hist^2)
8837 // device dedup: `penalty_last_n` is usize::MAX for any serve request with
8838 // a penalty, so without the cap this grew with the whole session.
8839 let win = sp.penalty_last_n.min(PEN_WINDOW_MAX);
8840 let w0 = pen_hist.len().saturating_sub(win);
8841 pen_hist_d = Some(e.htod_u32_v(&pen_hist[w0..])?);
8842 }
8843 if sampled {
8844 draft_logits.clear();
8845 draft_stats.clear();
8846 }
8847 // DRAFT-SIDE GRAMMAR MASK: clone the committed grammar state ONCE per round; each
8848 // position's mask is computed on that clone and advanced by the PROPOSED token. The
8849 // real state moves only on emission (verify's job), so the emitted stream is
8850 // unchanged — the mask only removes tokens the verify would have truncated anyway.
8851 let mut dmask_live = dmask_on;
8852 if dmask_live {
8853 let t_c = std::time::Instant::now();
8854 constraint
8855 .as_deref_mut()
8856 .unwrap()
8857 .draft_begin()
8858 .map_err(|e2| format!("constraint: {e2}"))?;
8859 dm_clone_ns += t_c.elapsed().as_nanos();
8860 dm_rounds += 1;
8861 }
8862 if let (false, Some(gr)) = (sampled || pen_on, &dctx.graph) {
8863 // GRAPH DRAFT: one dispatch per drafted token. The chain feeds itself on-device
8864 // (in-graph argmax -> tok_d -> next replay's embed; h_nextn -> h_seed_d; pos_d
8865 // inc'd in-graph); the host only reads 4B token (+4B p) and decides the break.
8866 e.set_i32_one(&mut dctx.g_pos, (pos + base0) as i32)?;
8867 e.set_u32_one(&mut dctx.g_tok, last_token)?;
8868 e.copy_into(&mut dctx.g_seed, 0, &h_seed_buf, n_embd)?;
8869 for j in 0..k_this {
8870 // per-position mask upload (contents only — the graph's baked pointer is
8871 // dctx.g_dmask). All-ones once masking goes dead mid-chain, so the captured
8872 // mask node degrades to a no-op ban instead of needing a second graph.
8873 if dmask_live
8874 && !upload_draft_mask(
8875 e,
8876 constraint.as_deref_mut().unwrap(),
8877 &mut dctx.g_dmask,
8878 mtp.d2t.as_ref(),
8879 d_vocab,
8880 dmask_words,
8881 )?
8882 {
8883 // no draft-vocab row is grammar-legal here (a trimmed FR-Spec head can
8884 // genuinely miss the legal set): neutralize the captured mask node and
8885 // finish the chain UNMASKED — exactly pre-lane behaviour, never worse.
8886 e.htod_u32_into(&mut dctx.g_dmask, &vec![u32::MAX; dmask_words])?;
8887 dmask_live = false;
8888 }
8889 gr.launch()?;
8890 scratch.kv.len += 1; // host mirror (len_d advanced in-graph)
8891 let idx = e.dtoh_u32_one(&dctx.g_tok)?;
8892 // #87 SENTINEL TRAP: an all-NaN head-logits row leaves the device argmax's
8893 // init sentinel (0x7FFFFFFF) in g_tok — feeding it onward dereferences
8894 // embed_row(sentinel) = table + ~4.6TB (never mapped) inside the NEXT graph
8895 // replay's embed node, and the MMU fault kills the CUDA context for the
8896 // whole process (research/pp2spec-crash-20260807: 3 coredumps, byte-exact
8897 // VA arithmetic). Refuse loudly instead; the diagnostics name the first-NaN
8898 // buffer (g_seed = the verify-side handoff vs head-side compute).
8899 if (idx as usize) >= d_vocab {
8900 // g_seed is SELF-FED (the replay writes h_nextn back into it), so it
8901 // reads as the head's OUTPUT at j; h_seed_buf is the round's INPUT
8902 // seed, untouched since the round-start copy — the pair discriminates
8903 // "seed arrived poisoned" from "head forward produced NaN".
8904 let seed_h = e.dtoh(&dctx.g_seed)?;
8905 let seed_nan = seed_h.iter().filter(|v| v.is_nan()).count();
8906 let in_h = e.dtoh(&h_seed_buf)?;
8907 let in_nan = in_h.iter().filter(|v| v.is_nan()).count();
8908 return Err(format!(
8909 "draft(graph) argmax sentinel 0x{idx:08x} >= d_vocab {d_vocab} at \
8910 round {round} j={j} pos={pos}: head-out NaN {seed_nan}/{n_embd}, \
8911 round-input-seed NaN {in_nan}/{n_embd} — refusing to dereference \
8912 the embed row (#87 trap)"
8913 )
8914 .into());
8915 }
8916 // trimmed draft vocab -> target token id (identity when no d2t map)
8917 let d = match &mtp.d2t {
8918 Some(map) => map[idx as usize],
8919 None => idx,
8920 };
8921 let draft_p = if p_min > 0.0
8922 || opti_fork
8923 .as_ref()
8924 .is_some_and(|fork| fork.controller.is_some())
8925 {
8926 Some(e.dtoh(&dctx.g_p)?[0])
8927 } else {
8928 None
8929 };
8930 if j == 0 {
8931 controller_draft_prob = draft_p;
8932 }
8933 if let Some(p) = draft_p.filter(|_| p_min > 0.0) {
8934 if p < p_min && (j > 0 || (pmin0 && base0 == 1)) {
8935 break;
8936 }
8937 }
8938 draft.push(d);
8939 // with a trimmed head the NEXT embed must read the TARGET id, not the draft
8940 // index the argmax wrote — patch the persistent token buffer (4B htod).
8941 if d != idx {
8942 e.set_u32_one(&mut dctx.g_tok, d)?;
8943 }
8944 // advance the SPECULATIVE state with the proposal; a dead chain drops to
8945 // unmasked drafting for the remaining positions (verify still arbitrates).
8946 // speculative advance; a chain the grammar can no longer follow (EOS
8947 // proposed) ends here. The captured mask node always runs, so a dead chain
8948 // leaves the buffer NEUTRAL (all-ones = ban nothing) before it exits.
8949 if dmask_live
8950 && !constraint
8951 .as_deref_mut()
8952 .unwrap()
8953 .draft_advance(d)
8954 .map_err(|e2| format!("constraint: {e2}"))?
8955 {
8956 e.htod_u32_into(&mut dctx.g_dmask, &vec![u32::MAX; dmask_words])?;
8957 break;
8958 }
8959 }
8960 // PURE-TEMP RE-TEST (lane/graph-s-key-exactness-20260819): the sampled graph is
8961 // legal ONLY in the regime it was captured in. The condition used to read
8962 // `(sampled, &dctx.graph_s)` and trusted `s_key` to have dropped anything else —
8963 // which it could not, because the key omitted the filters. Both halves are now
8964 // enforced: the key drops a stale graph, and this site refuses to launch one.
8965 } else if let (true, Some(gr)) = (sampled && pure_temp, &dctx.graph_s) {
8966 if skey_probe() {
8967 eprintln!(
8968 "[skey] chain=graph_s round={round} pure_temp={} top_k={} \
8969 top_p={} min_p={} s_key_parked={:?}",
8970 pure_temp as u8, sp.top_k, sp.top_p, sp.min_p, dctx.s_key,
8971 );
8972 }
8973 // SAMPLED GRAPH DRAFT: one replay per drafted token — head forward + gumbel +
8974 // argmax in ONE dispatch; the host reads 4B token (+4B p), D2Ds q into slot j,
8975 // and decides the break. Event-counter continuity: g_ctr is host-seeded to
8976 // sctr-1 ONCE per round (outside the graph); the in-graph bump runs BEFORE the
8977 // perturb, so replay j consumes counter sctr+j — exactly the eager arm's Philox
8978 // stream. Host sctr advances in lockstep (computed, no readback needed).
8979 e.set_i32_one(&mut dctx.g_pos, (pos + base0) as i32)?;
8980 e.set_u32_one(&mut dctx.g_tok, last_token)?;
8981 e.copy_into(&mut dctx.g_seed, 0, &h_seed_buf, n_embd)?;
8982 e.set_u32_one(&mut dctx.g_ctr, sctr.wrapping_sub(1))?;
8983 for j in 0..k_this {
8984 gr.launch()?;
8985 scratch.kv.len += 1; // host mirror (len_d advanced in-graph)
8986 sctr += 1; // mirrors the in-graph g_ctr bump (eager parity:
8987 // counts the p-min-discarded token too)
8988 // q retention: ONE async D2D of the persistent head-logits buffer into this
8989 // round's slot j (stream-ordered after the replay, before the next one).
8990 e.copy_into(&mut dctx.q_slots[j], 0, &dctx.g_q, d_vocab)?;
8991 let idx = e.dtoh_u32_one(&dctx.g_tok)?;
8992 // #87 SENTINEL TRAP (see the greedy graph arm above).
8993 if (idx as usize) >= d_vocab {
8994 let seed_h = e.dtoh(&dctx.g_seed)?;
8995 let seed_nan = seed_h.iter().filter(|v| v.is_nan()).count();
8996 return Err(format!(
8997 "draft(graph-sampled) argmax sentinel 0x{idx:08x} >= d_vocab \
8998 {d_vocab} at round {round} j={j} pos={pos}: round-seed NaN \
8999 {seed_nan}/{n_embd} — refusing to dereference the embed row \
9000 (#87 trap)"
9001 )
9002 .into());
9003 }
9004 let d = match &mtp.d2t {
9005 Some(map) => map[idx as usize],
9006 None => idx,
9007 };
9008 draft_idx.push(idx);
9009 if p_min > 0.0 {
9010 let p = e.dtoh(&dctx.g_p)?[0];
9011 if p < p_min && (j > 0 || (pmin0 && base0 == 1)) {
9012 break;
9013 }
9014 }
9015 draft.push(d);
9016 // trimmed head: the NEXT embed must read the TARGET id (see the greedy arm).
9017 if d != idx {
9018 e.set_u32_one(&mut dctx.g_tok, d)?;
9019 }
9020 }
9021 // uniform accept path: fill draft_stats per used slot (pure-temp regime — the
9022 // stats degenerate to th=0 / full-Z; one filter_stats launch per slot, tiny).
9023 for j in 0..draft.len().max(draft_idx.len()) {
9024 let rows0 = e.htod_i32(&[0])?;
9025 let (mut th_d, mut z_d, mut mx_d) = (e.zeros(1)?, e.zeros(1)?, e.zeros(1)?);
9026 e.filter_stats(
9027 &dctx.q_slots[j],
9028 d_vocab,
9029 &rows0,
9030 &mut th_d,
9031 &mut z_d,
9032 &mut mx_d,
9033 d_vocab,
9034 1,
9035 sp_temp,
9036 sp.top_k,
9037 sp.top_p,
9038 sp.min_p,
9039 )?;
9040 draft_stats.push((e.dtoh(&mx_d)?[0], e.dtoh(&th_d)?[0], e.dtoh(&z_d)?[0]));
9041 }
9042 } else {
9043 if skey_probe() && sampled {
9044 eprintln!(
9045 "[skey] chain=eager round={round} pure_temp={} top_k={} \
9046 top_p={} min_p={} s_key_parked={:?}",
9047 pure_temp as u8, sp.top_k, sp.top_p, sp.min_p, dctx.s_key,
9048 );
9049 }
9050 // EAGER DRAFT (fallback: MoE head/trunk, huge k, MEMRA_SPEC_NOGRAPH, capture fail).
9051 let mut e_tok = last_token;
9052 let mut d_seed = e.clone_dtod(&h_seed_buf)?;
9053 for j in 0..k_this {
9054 // GPU-ARGMAX DRAFT (2026-07-03): device logits + device argmax + 4-byte token
9055 // read instead of the ~600KB full-vocab dtoh + host argmax per draft token.
9056 let mtp_pos = pos + base0 + j;
9057 // draft-side grammar mask (eager twin of the graph arm's in-graph node).
9058 // A position with no legal draft-vocab row drops to unmasked drafting for
9059 // the rest of the chain (pre-lane behaviour; verify still arbitrates).
9060 if dmask_live {
9061 dmask_live = upload_draft_mask(
9062 e,
9063 constraint.as_deref_mut().unwrap(),
9064 &mut dctx.g_dmask,
9065 mtp.d2t.as_ref(),
9066 d_vocab,
9067 dmask_words,
9068 )?;
9069 }
9070 let (dl_d, h_nextn) = self.mtp_head_forward_dev(
9071 e,
9072 mtp,
9073 e_tok,
9074 &d_seed,
9075 &mut *scratch,
9076 mtp_pos,
9077 embd_dev,
9078 if dmask_live {
9079 Some((&dctx.g_dmask, dmask_words))
9080 } else {
9081 None
9082 },
9083 )?;
9084 let tok_d = if sampled {
9085 // FILTERED Gumbel-max: stats -> masked perturb -> argmax = one draw from
9086 // the filtered softmax (filters off => th=0, exact v1 semantics).
9087 if perturb_buf.is_none() {
9088 perturb_buf = Some(e.zeros(d_vocab.max(n_vocab))?);
9089 }
9090 let mut q_row = e.clone_dtod(&dl_d)?; // retained q (penalized when on)
9091 if pen_on {
9092 let h = pen_hist_d.as_ref().unwrap();
9093 let nh = h.len();
9094 e.penalize_logits(
9095 &mut q_row,
9096 h,
9097 nh,
9098 sp.penalty_repeat,
9099 sp.penalty_freq,
9100 sp.penalty_present,
9101 d_vocab,
9102 )?;
9103 }
9104 let rows0 = e.htod_i32(&[0])?;
9105 let (mut th_d, mut z_d, mut mx_d) =
9106 (e.zeros(1)?, e.zeros(1)?, e.zeros(1)?);
9107 e.filter_stats(
9108 &q_row, d_vocab, &rows0, &mut th_d, &mut z_d, &mut mx_d, d_vocab,
9109 1, sp_temp, sp.top_k, sp.top_p, sp.min_p,
9110 )?;
9111 let (th, z, mx) =
9112 (e.dtoh(&th_d)?[0], e.dtoh(&z_d)?[0], e.dtoh(&mx_d)?[0]);
9113 let pb = perturb_buf.as_mut().unwrap();
9114 e.gumbel_perturb_filtered(
9115 &q_row, pb, d_vocab, sp_seed, sctr, sp_temp, mx, th,
9116 )?;
9117 sctr += 1;
9118 draft_logits.push(q_row);
9119 draft_stats.push((mx, th, z));
9120 e.argmax_token_device(pb, d_vocab)?
9121 } else {
9122 e.argmax_token_device(&dl_d, d_vocab)?
9123 };
9124 let idx = e.dtoh_u32_one(&tok_d)?;
9125 // #87 SENTINEL TRAP (eager twin — see the graph arm). Extra diagnostics
9126 // here because the eager chain's operands are all readable: dl_d (the head
9127 // logits row) and d_seed (this step's h_seed) name the first-NaN buffer.
9128 if (idx as usize) >= d_vocab {
9129 let dl_h = e.dtoh(&dl_d)?;
9130 let dl_nan = dl_h.iter().filter(|v| v.is_nan()).count();
9131 let seed_h = e.dtoh(&d_seed)?;
9132 let seed_nan = seed_h.iter().filter(|v| v.is_nan()).count();
9133 return Err(format!(
9134 "draft(eager) argmax sentinel 0x{idx:08x} >= d_vocab {d_vocab} at \
9135 round {round} j={j} pos={pos}: head-logits NaN {dl_nan}/{d_vocab}, \
9136 step-seed NaN {seed_nan}/{n_embd} — refusing to dereference the \
9137 embed row (#87 trap)"
9138 )
9139 .into());
9140 }
9141 let d = match &mtp.d2t {
9142 Some(map) => map[idx as usize],
9143 None => idx,
9144 };
9145 if sampled {
9146 draft_idx.push(idx);
9147 }
9148 let draft_p = if p_min > 0.0
9149 || opti_fork
9150 .as_ref()
9151 .is_some_and(|fork| fork.controller.is_some())
9152 {
9153 let p_d = e.prob_of_token_device(&dl_d, &tok_d, d_vocab)?;
9154 Some(e.dtoh(&p_d)?[0])
9155 } else {
9156 None
9157 };
9158 if j == 0 {
9159 controller_draft_prob = draft_p;
9160 }
9161 if let Some(p) = draft_p.filter(|_| p_min > 0.0) {
9162 if p < p_min && (j > 0 || (pmin0 && base0 == 1)) {
9163 break;
9164 }
9165 }
9166 draft.push(d);
9167 e_tok = d;
9168 d_seed = h_nextn;
9169 // speculative advance; a chain the grammar can no longer follow (EOS
9170 // proposed) ends here — the prefix already proposed still rides verify.
9171 if dmask_live
9172 && !constraint
9173 .as_deref_mut()
9174 .unwrap()
9175 .draft_advance(d)
9176 .map_err(|e2| format!("constraint: {e2}"))?
9177 {
9178 break;
9179 }
9180 }
9181 if opti_fork
9182 .as_ref()
9183 .is_some_and(|fork| fork.controller.is_some())
9184 {
9185 controller_eager_state = Some((e_tok, d_seed));
9186 }
9187 }
9188 }
9189 let k_round = draft.len();
9190 if let Some(p) = pipe {
9191 p.draft_end(round);
9192 }
9193 drop(pipe_draft);
9194
9195 ph_mark(&mut ph_draft, phase_on);
9196 // --- 2. VERIFY: one batched target forward. With a pending bonus, it rides as col 0
9197 // (committing its KV/recur inside the SAME weight read); drafts follow. ---
9198 let verify_tokens: Vec<u32> = match pending {
9199 Some(b) => {
9200 let mut v = Vec::with_capacity(k_round + 1);
9201 v.push(b);
9202 v.extend_from_slice(&draft);
9203 v
9204 }
9205 None => draft.clone(),
9206 };
9207 let base = if pending.is_some() { 1 } else { 0 };
9208 // ckpt (REPLAY-FREE partial accept): retain per-layer state-rebuild inputs alongside
9209 // the verify. Pure buffer keep-alives + dtod clones — kernel work is unchanged.
9210 let mut ckpt = if let Some(ticket) = current_opti.as_mut() {
9211 Some(ticket.take_ckpt())
9212 } else if spec_replay {
9213 None
9214 } else {
9215 Some(VerifyCkpt::new(self.layers.len()))
9216 };
9217 let controller_can_probe = base == 1
9218 && k_round == 1
9219 && out.len().saturating_add(2) < max_new
9220 && controller_draft_prob.is_some()
9221 && opti_fork
9222 .as_ref()
9223 .and_then(|fork| fork.controller.as_ref())
9224 .is_some_and(|policy| !policy.breaker_tripped);
9225 let mut successor_attempt: Option<OptiControllerTicket> = None;
9226 let mut rejected_probe: Option<(f32, u32)> = None;
9227 let mut controller_prepared: Option<OptiControllerPrepared> = None;
9228 if controller_can_probe {
9229 // Prepare d2/q and, on admission, d3 before either current verify half is
9230 // issued. N stage 0 can then be followed immediately by N+1 stage 0; once N's
9231 // boundary fires, those dev0 launches overlap N stage 1 on dev1. Preparing on
9232 // the primary stream after N stage 1 would serialize the supposed pipeline.
9233 let eager_pos = scratch.kv.len + 1;
9234 let (optimistic_pending, pending_probability) = self.opti_controller_draft_step(
9235 e,
9236 mtp,
9237 &mut dctx,
9238 &mut *scratch,
9239 d_vocab,
9240 &mut controller_eager_state,
9241 eager_pos,
9242 embd_dev,
9243 )?;
9244 let first_probability = controller_draft_prob
9245 .ok_or("optipipe controller probe lost first-token probability")?;
9246 let q_proxy = first_probability * pending_probability;
9247 OPTI_GATE_CHECKS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
9248 OPTI_SHADOW_DRAFT_TOKENS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
9249 let admitted = opti_fork
9250 .as_ref()
9251 .and_then(|fork| fork.controller.as_ref())
9252 .ok_or("optipipe controller policy disappeared")?
9253 .admit(q_proxy);
9254 if admitted {
9255 OPTI_GATE_ADMITS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
9256 let eager_pos = scratch.kv.len + 1;
9257 let (optimistic_draft, optimistic_draft_probability) = self
9258 .opti_controller_draft_step(
9259 e,
9260 mtp,
9261 &mut dctx,
9262 &mut *scratch,
9263 d_vocab,
9264 &mut controller_eager_state,
9265 eager_pos,
9266 embd_dev,
9267 )?;
9268 OPTI_SHADOW_DRAFT_TOKENS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
9269 let eager_seed = controller_eager_state.take().map(|(token, seed)| {
9270 debug_assert_eq!(token, optimistic_draft);
9271 seed
9272 });
9273 controller_prepared = Some(OptiControllerPrepared {
9274 verify_tokens: [optimistic_pending, optimistic_draft],
9275 draft_prob: optimistic_draft_probability,
9276 eager_seed,
9277 q_proxy,
9278 scratch_len: scratch.kv.len,
9279 });
9280 } else {
9281 OPTI_GATE_REJECTS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
9282 OPTI_WASTED_DRAFT_TOKENS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
9283 rejected_probe = Some((q_proxy, optimistic_pending));
9284 eprintln!(
9285 "[opti-controller] reject q={q_proxy:.6} threshold={:.3}",
9286 opti_fork
9287 .as_ref()
9288 .and_then(|fork| fork.controller.as_ref())
9289 .expect("controller policy")
9290 .threshold,
9291 );
9292 }
9293 }
9294 let fork_attempt = match fork_generation.take() {
9295 Some(generation) if base == 1 && k_round == 1 => Some(generation),
9296 Some(generation) => {
9297 opti_fork
9298 .as_mut()
9299 .expect("fork generation without fork state")
9300 .retire(generation)?;
9301 None
9302 }
9303 None => None,
9304 };
9305 let (tlogits_d, vx) = if let Some(p) = pipe {
9306 self.decode_step_t_core_pipelined(
9307 e,
9308 &verify_tokens,
9309 pos,
9310 &mut *cache,
9311 embd_dev,
9312 ckpt.as_mut(),
9313 p,
9314 round,
9315 )?
9316 } else if controller_can_probe {
9317 let fence = opti_fork
9318 .as_ref()
9319 .ok_or("optipipe controller probe lost fork state")?
9320 .fence;
9321 let boundary = match current_opti.as_mut() {
9322 Some(ticket) => ticket.take_boundary(),
9323 None => self.verify_stage0_issue(
9324 e,
9325 &verify_tokens,
9326 pos,
9327 &mut *cache,
9328 embd_dev,
9329 ckpt.as_mut(),
9330 None,
9331 &fence,
9332 Some(true),
9333 None,
9334 )?,
9335 };
9336 if let Some(prepared) = controller_prepared.take() {
9337 let generation = {
9338 let fork = opti_fork
9339 .as_mut()
9340 .ok_or("optipipe controller admission lost fork state")?;
9341 let generation = fork.reserve_successor()?;
9342 let rt = fork.rt;
9343 let snapshot_fence = fork.fence;
9344 opti_snapshot_one_stage_owned_into(
9345 e,
9346 cache,
9347 rt,
9348 &snapshot_fence,
9349 0,
9350 fork.successor_snapshot_mut(),
9351 )?;
9352 generation
9353 };
9354 let mut successor_ckpt = VerifyCkpt::new(self.layers.len());
9355 let successor_boundary = self.verify_stage0_issue(
9356 e,
9357 &prepared.verify_tokens,
9358 pos + verify_tokens.len(),
9359 &mut *cache,
9360 embd_dev,
9361 Some(&mut successor_ckpt),
9362 None,
9363 &fence,
9364 Some(false),
9365 None,
9366 )?;
9367 OPTI_FORK_ATTEMPTS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
9368 let fork = opti_fork
9369 .as_ref()
9370 .ok_or("optipipe controller ticket lost fork state")?;
9371 successor_attempt = Some(fork.controller_ticket(
9372 generation,
9373 successor_boundary,
9374 successor_ckpt,
9375 prepared.verify_tokens,
9376 prepared.draft_prob,
9377 prepared.eager_seed,
9378 prepared.q_proxy,
9379 prepared.scratch_len,
9380 ));
9381 eprintln!(
9382 "[opti-controller] issue generation={} q={:.6} threshold={:.3} \
9383 verify={:?}",
9384 generation.id,
9385 prepared.q_proxy,
9386 fork.controller.expect("controller policy").threshold,
9387 prepared.verify_tokens,
9388 );
9389 }
9390 let result = self.verify_stage1_finish(
9391 e,
9392 boundary,
9393 &mut *cache,
9394 ckpt.as_mut(),
9395 None,
9396 &fence,
9397 successor_attempt.is_none(),
9398 )?;
9399 if let Some(ticket) = current_opti.as_mut() {
9400 ticket.settle();
9401 }
9402 if successor_attempt.is_some() {
9403 let fork = opti_fork
9404 .as_mut()
9405 .ok_or("optipipe successor snapshot lost fork state")?;
9406 let rt = fork.rt;
9407 let snapshot_fence = fork.fence;
9408 opti_snapshot_one_stage_owned_into(
9409 e,
9410 cache,
9411 rt,
9412 &snapshot_fence,
9413 1,
9414 fork.successor_snapshot_mut(),
9415 )?;
9416 // Publish N only after both independent successor-state queues are complete.
9417 fork.rt.publish_to(1, &e.stream())?;
9418 }
9419 result
9420 } else if let Some(ticket) = current_opti.as_mut() {
9421 let fork = opti_fork
9422 .as_mut()
9423 .ok_or("optipipe carried controller ticket lost fork state")?;
9424 let boundary = ticket.take_boundary();
9425 let result = self.verify_stage1_finish(
9426 e,
9427 boundary,
9428 &mut *cache,
9429 ckpt.as_mut(),
9430 None,
9431 &fork.fence,
9432 true,
9433 )?;
9434 ticket.settle();
9435 result
9436 } else if let Some(generation) = fork_attempt {
9437 let fork = opti_fork
9438 .as_mut()
9439 .expect("fork generation without fork state");
9440 fork.capture_seed(e, generation, &h_seed_buf, &fill_prev, scratch.kv.len)?;
9441 let action = fork.mode.action(generation.id);
9442 let boundary = self.verify_stage0_issue(
9443 e,
9444 &verify_tokens,
9445 pos,
9446 &mut *cache,
9447 embd_dev,
9448 ckpt.as_mut(),
9449 None,
9450 &fork.fence,
9451 Some(true),
9452 None,
9453 )?;
9454 OPTI_FORK_ATTEMPTS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
9455 let mut ticket = fork.ticket(generation, boundary);
9456 if action == OptiForkAction::Abort {
9457 return Err(format!(
9458 "optipipe forced abort with generation {} stage0 in flight",
9459 generation.id,
9460 )
9461 .into());
9462 }
9463 fork.reconcile(
9464 e,
9465 &mut *cache,
9466 &mut *scratch,
9467 &snap,
9468 &mut h_seed_buf,
9469 &mut fill_prev,
9470 generation,
9471 action,
9472 verify_tokens[0],
9473 )?;
9474 let result = if action == OptiForkAction::Hit {
9475 let boundary = ticket.take_boundary();
9476 self.verify_stage1_finish(
9477 e,
9478 boundary,
9479 &mut *cache,
9480 ckpt.as_mut(),
9481 None,
9482 &fork.fence,
9483 true,
9484 )?
9485 } else {
9486 // The optimistic boundary slot has no reader. Re-run the unchanged serial
9487 // verify only after E_restart published the restored stage-0 state.
9488 self.decode_step_t_core(
9489 e,
9490 &verify_tokens,
9491 pos,
9492 &mut *cache,
9493 embd_dev,
9494 ckpt.as_mut(),
9495 )?
9496 };
9497 ticket.settle();
9498 debug_assert_eq!(ticket.generation, generation);
9499 fork.retire(generation)?;
9500 result
9501 } else {
9502 self.decode_step_t_core(
9503 e,
9504 &verify_tokens,
9505 pos,
9506 &mut *cache,
9507 embd_dev,
9508 ckpt.as_mut(),
9509 )?
9510 };
9511 let pipe_accept = match pipe {
9512 Some(p) => Some(p.accept_begin(round)?),
9513 None => None,
9514 };
9515
9516 ph_mark(&mut ph_verify, phase_on);
9517 // --- 3. GREEDY ACCEPT (walk prefix, stop at first mismatch) ---
9518 // DEVICE-ARGMAX ACCEPT: argmax every verify column ON DEVICE (same 2-pass kernels +
9519 // smallest-index tie-break as host argmax, argmax_gate-validated) and read back ONE
9520 // [T] u32 — replaces the T x n_vocab f32 dtoh + T host argmaxes per round.
9521 // t_pred[j] = target's greedy prediction for the slot after draft[j-1] (j>=1) or after
9522 // last_token (j==0). With a pending bonus, col 0 IS the prediction after last_token
9523 // (== the bonus), so every index shifts by `base` and last_pred is unused.
9524 let t_v = verify_tokens.len();
9525 let mut preds: Vec<u32> = Vec::new();
9526 if !sampled {
9527 for j in 0..t_v {
9528 e.argmax_token_device_col(&tlogits_d, j, n_vocab, &mut preds_d, j)?;
9529 }
9530 preds = e.dtoh_u32(&preds_d)?; // <- the verify-GPU wait lands here
9531 // #87 SENTINEL TRAP, verify side: a sentinel pred becomes the round's bonus =
9532 // next round's last_token = the next chain's embed lookup. Catch it at the
9533 // source with the column named — an all-NaN VERIFY column implicates the
9534 // stage-split trunk (decode_step_t_core_ppn), not the draft head.
9535 if let Some(bad) = preds[..t_v].iter().position(|&p| (p as usize) >= n_vocab) {
9536 let col = &tlogits_d.slice(bad * n_vocab..(bad + 1) * n_vocab);
9537 let mut probe = e.zeros(n_vocab)?;
9538 e.copy_view_into(&mut probe, 0, col, n_vocab)?;
9539 let col_h = e.dtoh(&probe)?;
9540 let col_nan = col_h.iter().filter(|v| v.is_nan()).count();
9541 return Err(format!(
9542 "verify argmax sentinel 0x{:08x} >= n_vocab {n_vocab} at round {round} \
9543 col {bad}/{t_v} pos={pos}: verify-logits col NaN {col_nan}/{n_vocab} \
9544 — the stage-split verify produced a poisoned column (#87 trap)",
9545 preds[bad]
9546 )
9547 .into());
9548 }
9549 }
9550 ph_mark(&mut ph_wait, phase_on);
9551 let t_pred = |j: usize| -> u32 {
9552 if j == 0 && base == 0 {
9553 last_pred
9554 } else {
9555 // GREEDY-ONLY: `preds` is filled under `if !sampled` above. The debug print
9556 // used to call this from the sampled arm and panicked the worker; it now goes
9557 // through `debug_t_pred0`. Keep the strict index here — in the greedy walk an
9558 // out-of-range pred is a real bug, not something to paper over.
9559 debug_assert!(
9560 !sampled,
9561 "t_pred is greedy-only: `preds` is empty in the sampled arm"
9562 );
9563 preds[base + j - 1]
9564 }
9565 };
9566 let mut devacc_seeded = false;
9567 let mut devacc_acc: Option<CudaSlice<u32>> = None;
9568 let (n_acc, bonus) = if !sampled {
9569 // ROUND-STREAM stage (a) (MEMRA_SPEC_DEVACC=1 opt-in): the walk runs ON DEVICE
9570 // (spec_accept_greedy, verbatim rule) and the host reads back 8B (n_acc, bonus)
9571 // instead of the [T] preds. Same sync count — machinery for stages (b)/(c),
9572 // gated on token identity vs the host walk (the arms below are bit-equal rules).
9573 if crate::spec::spec_devacc() && k_round > 0 && !spec_replay && constraint.is_none()
9574 {
9575 let draft_d = e.htod_u32_v(&draft)?;
9576 let mut acc_out = e.alloc_u32_zeroed(2)?;
9577 e.spec_accept_greedy(
9578 &preds_d,
9579 &draft_d,
9580 last_pred,
9581 base,
9582 k_round,
9583 &mut acc_out,
9584 )?;
9585 devacc_acc = Some(acc_out.clone());
9586 // stage (b): next-round seed gathered ON DEVICE from acc_out before the host
9587 // ever reads n_acc (j=base+n_acc -> vx col j-1; j==0 -> fill_prev). The three
9588 // non-replay commit arms skip their host-offset seed copies (guarded below);
9589 // the legacy spec_replay arm keeps its own rx-based seeding (excluded here).
9590 // NOTE: fill_prev is NOT updated here — the commit arms' TRUE-HIDDEN
9591 // REFRESH reads the OLD fill_prev (predecessor of this round's verify batch);
9592 // the update lands after the arms (devacc_seeded guard below).
9593 e.spec_seed_gather(&vx, &fill_prev, &acc_out, &mut h_seed_buf, base, n_embd)?;
9594 // 3a: KV lens roll back on device (len = saved + base + n_acc, all arms'
9595 // unified rule; full accept rewrites the verify-left value). Host mirrors
9596 // update after the readback; commit_verified_prefix skips its len_d writes.
9597 if let Some(successor) = successor_attempt.as_ref() {
9598 opti_fork
9599 .as_mut()
9600 .ok_or("optipipe successor reconcile lost fork state")?
9601 .queue_actual_reconcile(
9602 e,
9603 &snap,
9604 &acc_out,
9605 successor.verify_tokens[0],
9606 base,
9607 )?;
9608 } else if let Some(ptrs) = &kv_len_ptrs {
9609 let saved: Vec<i32> = (0..self.layers.len())
9610 .map(|il| snap.kv_len[il].map(|v| v as i32).unwrap_or(0))
9611 .collect();
9612 let saved_d = e.htod_i32(&saved)?;
9613 e.spec_rollback_kv(ptrs, &saved_d, &acc_out, base, self.layers.len())?;
9614 }
9615 devacc_seeded = true;
9616 let ab = e.dtoh_u32(&acc_out)?;
9617 (ab[0] as usize, ab[1])
9618 } else {
9619 let mut n_acc = 0usize;
9620 for j in 0..k_round {
9621 if t_pred(j) == draft[j] {
9622 n_acc += 1;
9623 } else {
9624 break;
9625 }
9626 }
9627 // bonus = target's own token at the first non-accepted slot. n_acc in 0..=k; t_pred
9628 // is defined for j in 0..=k (j==0 -> last_logits, j>=1 -> col j-1, last col = k-1).
9629 (n_acc, t_pred(n_acc))
9630 }
9631 } else {
9632 // --- SAMPLED ACCEPT (rejection sampling): u_j < p_j(x_j)/q_j(x_j) walk ---
9633 if col_buf.is_none() {
9634 col_buf = Some(e.zeros(n_vocab)?);
9635 }
9636 // FILTERED p_j: per-verify-col stats (one batched filter_stats call), then the
9637 // filtered gather. j==0&&base==0 reads last_col (its own stats row appended).
9638 let mut pj = vec![0f32; k_round.max(1)];
9639 let mut col_stats: Vec<(f32, f32, f32)> = Vec::new(); // (max, th, z) per verify col used
9640 if k_round > 0 {
9641 let mut ids: Vec<u32> = Vec::new();
9642 let mut rows: Vec<i32> = Vec::new();
9643 for j in 0..k_round {
9644 if j > 0 || base == 1 {
9645 ids.push(draft[j]);
9646 rows.push((base + j) as i32 - 1);
9647 }
9648 }
9649 if !ids.is_empty() {
9650 let nr = rows.len();
9651 // penalties: materialize the used columns into one contiguous penalized
9652 // buffer (rows remapped 0..nr) so stats+gathers see the penalized p.
9653 // penalties: materialize used columns contiguously, penalize all rows in
9654 // one launch, and point stats+gathers at the penalized buffer (rows 0..nr).
9655 let p_rows: Vec<i32> = if pen_on {
9656 (0..nr as i32).collect()
9657 } else {
9658 rows.clone()
9659 };
9660 if pen_on {
9661 if pcol_buf.as_ref().map(|b| b.len()).unwrap_or(0) < nr * n_vocab {
9662 pcol_buf = Some(e.zeros(nr * n_vocab)?);
9663 }
9664 let pc = pcol_buf.as_mut().unwrap();
9665 for (i2, &r) in rows.iter().enumerate() {
9666 let c = r as usize;
9667 e.copy_view_into(
9668 pc,
9669 i2 * n_vocab,
9670 &tlogits_d.slice(c * n_vocab..(c + 1) * n_vocab),
9671 n_vocab,
9672 )?;
9673 }
9674 let h = pen_hist_d.as_ref().unwrap();
9675 let nh = h.len();
9676 e.penalize_logits_rows(
9677 pc,
9678 h,
9679 nh,
9680 sp.penalty_repeat,
9681 sp.penalty_freq,
9682 sp.penalty_present,
9683 n_vocab,
9684 nr,
9685 )?;
9686 }
9687 let p_src: &CudaSlice<f32> = if pen_on {
9688 pcol_buf.as_ref().unwrap()
9689 } else {
9690 &tlogits_d
9691 };
9692 let rowsd = e.htod_i32(&p_rows)?;
9693 let (mut th_d, mut z_d, mut mx_d) =
9694 (e.zeros(nr)?, e.zeros(nr)?, e.zeros(nr)?);
9695 e.filter_stats(
9696 p_src, n_vocab, &rowsd, &mut th_d, &mut z_d, &mut mx_d, n_vocab, nr,
9697 sp_temp, sp.top_k, sp.top_p, sp.min_p,
9698 )?;
9699 let idsd = e.htod_u32_v(&ids)?;
9700 let mut outd = e.zeros(nr)?;
9701 e.softmax_gather_filtered(
9702 p_src, n_vocab, &idsd, &rowsd, &th_d, &z_d, &mut outd, n_vocab, nr,
9703 sp_temp,
9704 )?;
9705 let outv = e.dtoh(&outd)?;
9706 let (thv, zv, mxv) = (e.dtoh(&th_d)?, e.dtoh(&z_d)?, e.dtoh(&mx_d)?);
9707 let mut oi = 0usize;
9708 for j in 0..k_round {
9709 if j > 0 || base == 1 {
9710 pj[j] = outv[oi];
9711 oi += 1;
9712 }
9713 }
9714 col_stats = (0..nr).map(|i| (mxv[i], thv[i], zv[i])).collect();
9715 }
9716 if base == 0 {
9717 let lc: &CudaSlice<f32> = if pen_on {
9718 if col_buf.is_none() {
9719 col_buf = Some(e.zeros(n_vocab)?);
9720 }
9721 let cb = col_buf.as_mut().unwrap();
9722 e.copy_into(
9723 cb,
9724 0,
9725 last_col_logits
9726 .as_ref()
9727 .expect("sampled: last_col_logits unset"),
9728 n_vocab,
9729 )?;
9730 let h = pen_hist_d.as_ref().unwrap();
9731 let nh = h.len();
9732 e.penalize_logits(
9733 cb,
9734 h,
9735 nh,
9736 sp.penalty_repeat,
9737 sp.penalty_freq,
9738 sp.penalty_present,
9739 n_vocab,
9740 )?;
9741 col_buf.as_ref().unwrap()
9742 } else {
9743 last_col_logits
9744 .as_ref()
9745 .expect("sampled: last_col_logits unset")
9746 };
9747 let rows0 = e.htod_i32(&[0])?;
9748 let (mut th_d, mut z_d, mut mx_d) = (e.zeros(1)?, e.zeros(1)?, e.zeros(1)?);
9749 e.filter_stats(
9750 lc, n_vocab, &rows0, &mut th_d, &mut z_d, &mut mx_d, n_vocab, 1,
9751 sp_temp, sp.top_k, sp.top_p, sp.min_p,
9752 )?;
9753 let idsd = e.htod_u32_v(&[draft[0]])?;
9754 let mut outd = e.zeros(1)?;
9755 e.softmax_gather_filtered(
9756 lc, n_vocab, &idsd, &rows0, &th_d, &z_d, &mut outd, n_vocab, 1, sp_temp,
9757 )?;
9758 pj[0] = e.dtoh(&outd)?[0];
9759 last_col_stats =
9760 Some((e.dtoh(&mx_d)?[0], e.dtoh(&th_d)?[0], e.dtoh(&z_d)?[0]));
9761 }
9762 }
9763 // q source: the graph arm retained the head logits in the persistent q_slots;
9764 // the eager arm in per-round draft_logits clones. Same raw-logit values either way.
9765 // FILTERED q_j: stats from draft_stats (eager pushes in-chain; the graph arm
9766 // computes them post-replay — graph engages only filter/penalty-free, so the
9767 // stats degenerate to th=0/full-Z there, keeping ONE accept path).
9768 let q_bufs: &[CudaSlice<f32>] = if dctx.graph_s.is_some() {
9769 &dctx.q_slots
9770 } else {
9771 &draft_logits
9772 };
9773 let mut n_acc = 0usize;
9774 for j in 0..k_round {
9775 let (qmx, qth, qz) = draft_stats[j];
9776 let idsd = e.htod_u32_v(&[draft_idx[j]])?;
9777 let rowsd = e.htod_i32(&[0])?;
9778 let thd = e.htod(&[qth])?;
9779 let zd = e.htod(&[qz])?;
9780 let _ = qmx;
9781 let mut outd = e.zeros(1)?;
9782 e.softmax_gather_filtered(
9783 &q_bufs[j], d_vocab, &idsd, &rowsd, &thd, &zd, &mut outd, d_vocab, 1,
9784 sp_temp,
9785 )?;
9786 let qj = e.dtoh(&outd)?[0];
9787 let u = host_u01(sp_seed, uctr);
9788 uctr += 1;
9789 let accept = (u as f64) * (qj as f64) < pj[j] as f64;
9790 // SKEY PROBE: q == 0 for the token the draft actually proposed is the
9791 // exactness signature (see `skey_probe`). Impossible when the draft was
9792 // drawn from the same filtered distribution the verify reconstructs here;
9793 // `u * 0 < p` makes it an UNCONDITIONAL accept whenever p > 0.
9794 if skey_probe() && qj == 0.0 {
9795 eprintln!(
9796 "[skey] EXACTNESS q=0 round={round} j={j} draft_tok={} \
9797 draft_idx={} p={:e} u={u} accepted={} th_z={:?}",
9798 draft[j], draft_idx[j], pj[j], accept as u8, draft_stats[j],
9799 );
9800 }
9801 if accept {
9802 n_acc += 1;
9803 } else {
9804 break;
9805 }
9806 }
9807 let bonus = if n_acc == k_round {
9808 // FULL ACCEPT: bonus ~ FILTERED softmax at the last verify column.
9809 let col = base + k_round - 1;
9810 let cb = col_buf.as_mut().unwrap();
9811 e.copy_view_into(
9812 cb,
9813 0,
9814 &tlogits_d.slice(col * n_vocab..(col + 1) * n_vocab),
9815 n_vocab,
9816 )?;
9817 if pen_on {
9818 let h = pen_hist_d.as_ref().unwrap();
9819 let nh = h.len();
9820 e.penalize_logits(
9821 cb,
9822 h,
9823 nh,
9824 sp.penalty_repeat,
9825 sp.penalty_freq,
9826 sp.penalty_present,
9827 n_vocab,
9828 )?;
9829 }
9830 if perturb_buf.is_none() {
9831 perturb_buf = Some(e.zeros(d_vocab.max(n_vocab))?);
9832 }
9833 // STATS MUST COME FROM THIS COLUMN (bug fix 2026-08-05, lane/sampler-
9834 // truncation-fix; receipts research/sampfix-20260805/). The old code reused
9835 // `col_stats.last()` here, which is ALWAYS the wrong row: the gathered set
9836 // covers verify columns 0..=(base+k_round-2) (rows pushed as base+j-1), while
9837 // the full-accept bonus samples column base+k_round-1 — exactly ONE PAST the
9838 // last gathered column, in both base arms. `th` is a threshold in e-units of
9839 // its OWN row's max, so feeding a neighbour's (row_max, th) into
9840 // gumbel_perturb_filtered mis-scales every e0 = exp((x-row_max)/T). When the
9841 // donor column's peak is higher by more than T*ln(1/th), EVERY id fails
9842 // `e0 >= th`, the whole perturbed row becomes -3.4e38, and the 2-pass argmax
9843 // falls through to its smallest-index tie-break => token id 0 ("!") spliced
9844 // mid-word. Fragility is ordered by how large th is: min_p pins th = min_p
9845 // (0.05 => trigger at delta > 2.4 at T=0.8, fires constantly), top_p's
9846 // mass-boundary th is smaller, top_k's k-th-largest th smaller still — which
9847 // is why the head-to-head matrix saw min_p and top_p corrupt while top_k-only
9848 // stayed clean. The pure-temp default regime is immune (th == 0 masks nothing,
9849 // and row_max is unused once nothing is masked), so this fix is a byte-level
9850 // no-op for the untruncated serve default. One extra one-block filter_stats
9851 // per full-accept round is the whole cost.
9852 let (mx, th) = {
9853 let rows0 = e.htod_i32(&[0])?;
9854 let (mut th_d, mut z_d, mut mx_d) = (e.zeros(1)?, e.zeros(1)?, e.zeros(1)?);
9855 let cb0 = col_buf.as_ref().unwrap();
9856 e.filter_stats(
9857 cb0, n_vocab, &rows0, &mut th_d, &mut z_d, &mut mx_d, n_vocab, 1,
9858 sp_temp, sp.top_k, sp.top_p, sp.min_p,
9859 )?;
9860 (e.dtoh(&mx_d)?[0], e.dtoh(&th_d)?[0])
9861 };
9862 let pb = perturb_buf.as_mut().unwrap();
9863 let cb2 = col_buf.as_ref().unwrap();
9864 e.gumbel_perturb_filtered(cb2, pb, n_vocab, sp_seed, sctr, sp_temp, mx, th)?;
9865 sctr += 1;
9866 let td = e.argmax_token_device(pb, n_vocab)?;
9867 e.dtoh_u32_one(&td)?
9868 } else {
9869 // REJECT at n_acc: bonus ~ norm(max(0, softmax_T(p) - softmax_T(q))).
9870 let cb = col_buf.as_mut().unwrap();
9871 if n_acc > 0 || base == 1 {
9872 let col = base + n_acc - 1;
9873 e.copy_view_into(
9874 cb,
9875 0,
9876 &tlogits_d.slice(col * n_vocab..(col + 1) * n_vocab),
9877 n_vocab,
9878 )?;
9879 } else {
9880 let lc = last_col_logits.as_ref().unwrap();
9881 e.copy_into(cb, 0, lc, n_vocab)?;
9882 }
9883 if pen_on {
9884 let h = pen_hist_d.as_ref().unwrap();
9885 let nh = h.len();
9886 e.penalize_logits(
9887 cb,
9888 h,
9889 nh,
9890 sp.penalty_repeat,
9891 sp.penalty_freq,
9892 sp.penalty_present,
9893 n_vocab,
9894 )?;
9895 }
9896 let cb2 = col_buf.as_ref().unwrap();
9897 let sc = sctr;
9898 sctr += 1;
9899 // p-stats for the reject column: from col_stats when the col was gathered,
9900 // else (j==0&&base==0) from last_col_stats.
9901 let p_stats = if n_acc > 0 || base == 1 {
9902 // col index within the gathered set == number of gathered cols before n_acc
9903 let gi = if base == 1 { n_acc } else { n_acc - 1 };
9904 col_stats.get(gi).copied().unwrap_or_else(|| {
9905 (0.0, 0.0, 1.0) // unreachable: gathered cols always cover the reject slot
9906 })
9907 } else {
9908 last_col_stats.expect("sampled: last_col_stats unset at reject")
9909 };
9910 let q_stats = draft_stats[n_acc];
9911 if let Some(map) = &d2t_dev {
9912 if q_full_buf.is_none() {
9913 q_full_buf = Some(e.zeros(n_vocab)?);
9914 }
9915 let qf = q_full_buf.as_mut().unwrap();
9916 e.scatter_trim_logits(&q_bufs[n_acc], map, qf, d_vocab, n_vocab)?;
9917 let qf2 = q_full_buf.as_ref().unwrap();
9918 e.residual_sample_filtered(
9919 cb2,
9920 Some(qf2),
9921 n_vocab,
9922 sp_temp,
9923 sp_seed,
9924 sc,
9925 p_stats,
9926 q_stats,
9927 &mut sample_tok,
9928 )?;
9929 } else {
9930 e.residual_sample_filtered(
9931 cb2,
9932 Some(&q_bufs[n_acc]),
9933 n_vocab,
9934 sp_temp,
9935 sp_seed,
9936 sc,
9937 p_stats,
9938 q_stats,
9939 &mut sample_tok,
9940 )?;
9941 }
9942 e.dtoh_u32(&sample_tok)?[0]
9943 };
9944 (n_acc, bonus)
9945 };
9946 // --- 3b. GRAMMAR TRUNCATION (constrained spec, 2026-08-03): the grammar is
9947 // an extra rejection rule AFTER the exactness verify (the batched-verify-twins
9948 // ordering). Walk the accepted drafts through the grammar in commit order; the
9949 // first illegal token truncates acceptance at its slot, and that slot's emission
9950 // is recomputed as the MASKED argmax of the target's own verify column — token-
9951 // identical to constrained plain greedy decode (an unmasked argmax that is
9952 // grammar-legal IS the masked argmax: masking only removes competitors). The
9953 // column D2H (~1MB) is paid only when a cut fires — the tight-grammar cost,
9954 // measured in acceptance numbers, never hidden.
9955 let (n_acc, bonus) = match constraint.as_deref_mut() {
9956 None => (n_acc, bonus),
9957 Some(c) => {
9958 fn ce(e2: String) -> Box<dyn std::error::Error> {
9959 format!("constraint: {e2}").into()
9960 }
9961 let mut na = n_acc;
9962 let mut cut = false;
9963 for (j, &d) in draft.iter().enumerate().take(n_acc) {
9964 if c.is_allowed(d).map_err(ce)? {
9965 c.consume(d).map_err(ce)?;
9966 } else {
9967 na = j;
9968 cut = true;
9969 dm_cut_tokens += n_acc - j;
9970 break;
9971 }
9972 }
9973 if cut {
9974 dm_cuts += 1;
9975 }
9976 let mut bo = bonus;
9977 if cut || !c.is_allowed(bo).map_err(ce)? {
9978 let mut row = if na == 0 && base == 0 {
9979 init_logits_host
9980 .clone()
9981 .ok_or("constraint: init logits missing (round-0 cut)")?
9982 } else {
9983 e.dtoh_view(
9984 &tlogits_d.slice((base + na - 1) * n_vocab..(base + na) * n_vocab),
9985 )?
9986 };
9987 c.mask_logits(&mut row).map_err(ce)?;
9988 bo = argmax(&row) as u32;
9989 }
9990 c.consume(bo).map_err(ce)?;
9991 (na, bo)
9992 }
9993 };
9994 let mut successor_valid = false;
9995 if let Some((q_proxy, expected_d2)) = rejected_probe {
9996 let v_n = n_acc == 1 && bonus == expected_d2;
9997 eprintln!(
9998 "[opti-controller] shadow q={q_proxy:.6} admitted=false v_n={v_n} \
9999 expected_d2={expected_d2} n_acc={n_acc} bonus={bonus}",
10000 );
10001 }
10002 if let Some(successor) = successor_attempt.as_ref() {
10003 successor_valid = n_acc == 1 && bonus == successor.verify_tokens[0];
10004 let generation = successor.generation;
10005 let q_proxy = successor.q_proxy;
10006 let expected_pending = successor.verify_tokens[0];
10007 let resolution_ms = successor.issued_at.elapsed().as_secs_f64() * 1e3;
10008 let fork = opti_fork
10009 .as_mut()
10010 .ok_or("optipipe successor resolution lost fork state")?;
10011 fork.finish_actual_reconcile(e, &mut *cache, &snap, n_acc, base, successor_valid)?;
10012 if successor_valid {
10013 OPTI_FORK_HITS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
10014 } else {
10015 OPTI_FORK_MISSES.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
10016 OPTI_RECONCILES.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
10017 OPTI_WASTED_DRAFT_TOKENS.fetch_add(2, std::sync::atomic::Ordering::Relaxed);
10018 }
10019 let breaker_tripped = fork
10020 .controller
10021 .as_mut()
10022 .expect("controller policy")
10023 .resolve(successor_valid);
10024 if breaker_tripped {
10025 OPTI_BREAKER_TRIPS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
10026 }
10027 eprintln!(
10028 "[opti-controller] resolve generation={} hit={} q={q_proxy:.6} \
10029 expected_pending={expected_pending} n_acc={n_acc} bonus={bonus} \
10030 resolution_ms={resolution_ms:.3} reconcile={} breaker={}",
10031 generation.id, successor_valid, !successor_valid, breaker_tripped,
10032 );
10033 if !successor_valid {
10034 let mut successor = successor_attempt
10035 .take()
10036 .expect("controller successor disappeared on miss");
10037 successor.settle();
10038 fork.retire(generation)?;
10039 }
10040 }
10041 total_drafted += k_round;
10042 total_accepted += n_acc;
10043 if let Some(t) = sess_telem {
10044 // Greedy, rejection-sampling, and grammar truncation all converge here after
10045 // the accept decision is already on host. Fixed-size relaxed atomics only.
10046 t.record_round(k_round, n_acc);
10047 }
10048 if spec_stats {
10049 st_len_hist[k_round] += 1;
10050 for j in 0..k_round {
10051 st_drafted[j] += 1;
10052 }
10053 for j in 0..n_acc {
10054 st_accepted[j] += 1;
10055 }
10056 if n_acc == k_round {
10057 st_full += 1;
10058 }
10059 }
10060
10061 if debug_spec {
10062 eprintln!(
10063 "[R{round}] pos={pos} out_len={} last_tok={last_token} draft={draft:?} n_acc={n_acc} bonus={bonus} t_pred0={}",
10064 out.len(),
10065 // NOT `t_pred(0)`: `preds` is filled only under `if !sampled` above, so on a
10066 // sampled request round >= 1 (base == 1) indexed an EMPTY vector and PANICKED
10067 // the GPU worker thread — a debug flag that killed the exact regime you would
10068 // set it to investigate. See `debug_t_pred0`.
10069 debug_t_pred0(sampled, base, last_pred, &preds)
10070 );
10071 }
10072
10073 // --- 4. COMMIT: draft[0..n_acc] then bonus (n_acc + 1 tokens) ---
10074 let commit_started = std::time::Instant::now();
10075 // SESSION MODE: every accepted column is already in the CACHE — `out` must carry all
10076 // of them (overshoot past max_new included) or `committed` under-counts the cache rows
10077 // and the next turn's continuation seeds one token off (gate-caught 2026-07-05). The
10078 // single-shot path keeps the cap (its caller truncates + drops the cache anyway).
10079 for j in 0..n_acc {
10080 if !session_mode && out.len() >= max_new {
10081 break;
10082 }
10083 out.push(draft[j]);
10084 }
10085 if pen_on {
10086 pen_hist.extend_from_slice(&draft[0..n_acc]);
10087 pen_hist.push(bonus);
10088 }
10089 let bonus_emitted = session_mode || out.len() < max_new;
10090 if bonus_emitted {
10091 out.push(bonus);
10092 }
10093 last_token = bonus;
10094
10095 // --- 5. ROLLBACK + advance (§C) ---
10096 if n_acc == k_round && !spec_replay {
10097 // FULL ACCEPT, BONUS FOLD: all verify columns (pending? + drafts) are committed in
10098 // cache; the NEW bonus stays PENDING for the next round's verify batch — NO extra
10099 // T=1 trunk pass. The next draft chain seeds from the MTP block's h_nextn at the
10100 // bonus position: one MTP-block pass (~1/33 trunk cost) replaces the trunk read.
10101 // last_pred is dead in the pending path (t_pred reads verify col 0).
10102 //
10103 // PERSISTENT DRAFT KV, full-accept fill: the chain covered last_token +
10104 // draft[0..k_round-2] as INPUTS (slots P..P'-2); draft[k_round-1] (slot P'-1) was
10105 // only ever an output, so its entry is MISSING. Fill it from vh_seed — its EXACT
10106 // trunk hidden (the last verify column). set_len first: a p-min break may have
10107 // left one extra chain append at that slot. Partial accepts need NO fill (the
10108 // chain already covered every accepted position; round-start set_len truncates).
10109 let mut vh_seed = e.zeros(n_embd)?;
10110 e.copy_view_into(
10111 &mut vh_seed,
10112 0,
10113 &vx.slice((t_v - 1) * n_embd..t_v * n_embd),
10114 n_embd,
10115 )?;
10116 if refresh {
10117 // TRUE-HIDDEN REFRESH (2026-07-03, the HANDOVER-listed acceptance lever):
10118 // overwrite ALL committed positions' scratch entries with K/V from their EXACT
10119 // verify hiddens — the reference engine's mtp_update fills from true hiddens;
10120 // the full stack (vx) is already resident from the verify. Replaces both the
10121 // chain-approximate entries AND the old last-token-only fill. Acceptance-only
10122 // (draft attention quality); exactness stays the verify's job.
10123 scratch.set_len(e, pos)?;
10124 // PREDECESSOR pairing: row i gets vx[i-1]; row 0 the carried fill_prev
10125 // (hidden of the last committed row before this verify batch).
10126 let mut vxs = e.zeros(t_v * n_embd)?;
10127 e.copy_into(&mut vxs, 0, &fill_prev, n_embd)?;
10128 if t_v > 1 {
10129 e.copy_view_into(
10130 &mut vxs,
10131 n_embd,
10132 &vx.slice(0..(t_v - 1) * n_embd),
10133 (t_v - 1) * n_embd,
10134 )?;
10135 }
10136 self.mtp_kv_fill(e, mtp, &verify_tokens, &vxs, pos, &mut *scratch, embd_dev)?;
10137 } else {
10138 scratch.set_len(e, pos + base + k_round - 1)?;
10139 // predecessor of the last draft = verify col t_v-2 (or fill_prev at t_v==1)
10140 let mut hp = e.zeros(n_embd)?;
10141 if t_v >= 2 {
10142 e.copy_view_into(
10143 &mut hp,
10144 0,
10145 &vx.slice((t_v - 2) * n_embd..(t_v - 1) * n_embd),
10146 n_embd,
10147 )?;
10148 } else {
10149 e.copy_into(&mut hp, 0, &fill_prev, n_embd)?;
10150 }
10151 self.mtp_kv_fill(
10152 e,
10153 mtp,
10154 &[draft[k_round - 1]],
10155 &hp,
10156 pos + base + k_round - 1,
10157 &mut *scratch,
10158 embd_dev,
10159 )?;
10160 }
10161 // REFERENCE SEEDING: no pseudo pass — the next chain's step 0 IS the
10162 // reference's (id_last, h_prev) draft row; it appends the bonus's scratch
10163 // entry itself. Seed = TRUE hidden of the bonus's predecessor (last verify
10164 // col). Saves one MTP-block pass per round on top of the pairing fix.
10165 if !devacc_seeded {
10166 e.copy_into(&mut h_seed_buf, 0, &vh_seed, n_embd)?;
10167 e.copy_into(&mut fill_prev, 0, &vh_seed, n_embd)?;
10168 }
10169 pending = Some(bonus);
10170 if debug_spec {
10171 eprintln!(" -> FULL ACCEPT (bonus pending, prev-h seed)");
10172 }
10173 } else if !spec_replay && base + n_acc >= 1 {
10174 // PARTIAL ACCEPT, REPLAY-FREE (2026-07-03 — the profiled #1 long-ctx spec cost):
10175 // the verify's first j = base+n_acc columns ARE the committed sequence, computed
10176 // bit-identically to eager (decode-exact contract) — so KEEP them: KV truncates to
10177 // pos+j, recurrent state rebuilds from the VerifyCkpt (same-kernel gdn prefix
10178 // re-run / pure state-clone restore), and the bonus stays PENDING exactly like the
10179 // full-accept path — the legacy duplicate trunk replay is gone. The next chain
10180 // seeds from the MTP pseudo-hidden of the bonus, whose seed = the TRUE verify
10181 // hidden of its predecessor (col j-1) — same one-hop pseudo structure as full
10182 // accept (never compounds: the next verify recomputes true hiddens for all
10183 // committed columns).
10184 let j = base + n_acc;
10185 self.commit_verified_prefix(
10186 e,
10187 &mut *cache,
10188 &snap,
10189 ckpt.as_ref().unwrap(),
10190 j,
10191 devacc_seeded,
10192 if devacc_seeded {
10193 devacc_acc.as_ref().map(|a| (a, base, t_v))
10194 } else {
10195 None
10196 },
10197 )?;
10198 let mut seed = e.zeros(n_embd)?;
10199 e.copy_view_into(
10200 &mut seed,
10201 0,
10202 &vx.slice((j - 1) * n_embd..j * n_embd),
10203 n_embd,
10204 )?;
10205 // Draft scratch: TRUE-HIDDEN REFRESH of the committed prefix (see the full-accept
10206 // branch); without it the chain entries stand and only the tail truncates. Either
10207 // way len ends at pos+j so the pseudo append lands at the bonus's slot pos+j
10208 // (persistent mode), rope pos+j+1 (chain convention).
10209 if refresh {
10210 scratch.set_len(e, pos)?;
10211 let mut vxs = e.zeros(j * n_embd)?;
10212 e.copy_into(&mut vxs, 0, &fill_prev, n_embd)?;
10213 if j > 1 {
10214 e.copy_view_into(
10215 &mut vxs,
10216 n_embd,
10217 &vx.slice(0..(j - 1) * n_embd),
10218 (j - 1) * n_embd,
10219 )?;
10220 }
10221 self.mtp_kv_fill(
10222 e,
10223 mtp,
10224 &verify_tokens[0..j],
10225 &vxs,
10226 pos,
10227 &mut *scratch,
10228 embd_dev,
10229 )?;
10230 } else {
10231 scratch.set_len(e, pos + j)?;
10232 }
10233 // REFERENCE SEEDING (see the full-accept branch): seed = TRUE hidden of the
10234 // bonus's predecessor (verify col j-1); no pseudo pass.
10235 if !devacc_seeded {
10236 e.copy_into(&mut h_seed_buf, 0, &seed, n_embd)?;
10237 e.copy_into(&mut fill_prev, 0, &seed, n_embd)?;
10238 }
10239 pending = Some(bonus);
10240 if debug_spec {
10241 eprintln!(" -> PARTIAL(replay-free j={j}, bonus pending, prev-h seed)");
10242 }
10243 } else if !spec_replay {
10244 // ZERO ROUND FOLD (2026-07-10, verify-cost target #3): base+n_acc == 0 — a
10245 // pending-less round where nothing was accepted (PMIN0 zero-draft chains after a
10246 // replay/commit, or plain 0-accept rounds at round 0). The old path replayed
10247 // [bonus] through a FULL m=1 trunk+head forward (the 489us full-vocab head pass
10248 // measured at ~0.75/round on PMIN0 configs). Instead: restore the pre-round
10249 // snapshot and let the bonus ride the NEXT round's verify as col 0 — the existing
10250 // base=1 pending machinery, bit-identical by the decode-exact verify contract.
10251 // Seed: the bonus's predecessor is the last COMMITTED token, whose hidden
10252 // fill_prev already carries (same seeding as the 1-token-replay case it replaces).
10253 cache.rollback(e, &snap, 0)?;
10254 scratch.set_len(e, pos)?;
10255 e.copy_into(&mut h_seed_buf, 0, &fill_prev, n_embd)?;
10256 pending = Some(bonus);
10257 if debug_spec {
10258 eprintln!(" -> ZERO-ROUND FOLD (bonus pending, fill_prev seed)");
10259 }
10260 } else {
10261 // PARTIAL ACCEPT, LEGACY REPLAY (seam MEMRA_SPEC_REPLAY=1 — or j==0: nothing of
10262 // this round survives, only possible before the first pending exists, ~round 0):
10263 // restore EVERYTHING to the pre-round snapshot (KV truncate to pos + recur
10264 // restore), then replay the committed prefix pending? ++ draft[0..n_acc] ++
10265 // [bonus] as ONE batched T forward — single weight read, bit-identical to greedy
10266 // (the verify-all-columns path is the same math). Commits the bonus with a TRUE
10267 // trunk hidden.
10268 cache.rollback(e, &snap, 0)?; // accept_len=0: KV len = pos, recur = snapshot
10269 let mut replay: Vec<u32> = Vec::with_capacity(base + n_acc + 1);
10270 if let Some(b) = pending.take() {
10271 replay.push(b);
10272 }
10273 replay.extend_from_slice(&draft[0..n_acc]);
10274 replay.push(bonus);
10275 // Full-stack forward (decode_step_t_core = decode_step_t_h_emb_dev's body):
10276 // Predecessor pairing seeds from the PREDECESSOR row (col len-2) — the same-row path takes the
10277 // last col exactly as before (byte-identical to the old _h_emb_dev call).
10278 let (rl_d, rx) = if self.qwen35_serving_class() {
10279 let mut logits = Vec::with_capacity(replay.len() * n_vocab);
10280 let mut hidden = e.uninit(replay.len() * n_embd)?;
10281 for (row, &token) in replay.iter().enumerate() {
10282 let (row_logits, row_hidden) =
10283 self.spec_target_step_h(e, token, &mut *cache)?;
10284 logits.extend_from_slice(&row_logits);
10285 e.dtod_copy_into(&row_hidden, &mut hidden, row * n_embd)?;
10286 }
10287 (e.htod(&logits)?, hidden)
10288 } else {
10289 self.decode_step_t_core(e, &replay, pos, &mut *cache, embd_dev, None)?
10290 };
10291 // last_pred = argmax of the LAST column's logits (predicts the token after `bonus`)
10292 // — device argmax + one 4-byte read instead of the full-vocab column dtoh.
10293 e.argmax_token_device_col(&rl_d, replay.len() - 1, n_vocab, &mut preds_d, 0)?;
10294 last_pred = e.dtoh_u32(&preds_d)?[0];
10295 if sampled {
10296 let lr0 = replay.len();
10297 let lc = last_col_logits
10298 .as_mut()
10299 .expect("sampled: last_col_logits unset");
10300 e.copy_view_into(
10301 lc,
10302 0,
10303 &rl_d.slice((lr0 - 1) * n_vocab..lr0 * n_vocab),
10304 n_vocab,
10305 )?;
10306 }
10307 let lr = replay.len();
10308 if lr >= 2 {
10309 e.copy_view_into(
10310 &mut h_seed_buf,
10311 0,
10312 &rx.slice((lr - 2) * n_embd..(lr - 1) * n_embd),
10313 n_embd,
10314 )?;
10315 } else {
10316 // 1-token replay (round-0 miss): the bonus's predecessor is the OLD
10317 // last_token, whose own-row hidden fill_prev still holds.
10318 e.copy_into(&mut h_seed_buf, 0, &fill_prev, n_embd)?;
10319 }
10320 // the bonus is COMMITTED here — it becomes the last committed row.
10321 let mut rh_last = e.zeros(n_embd)?;
10322 e.copy_view_into(
10323 &mut rh_last,
10324 0,
10325 &rx.slice((lr - 1) * n_embd..lr * n_embd),
10326 n_embd,
10327 )?;
10328 e.copy_into(&mut fill_prev, 0, &rh_last, n_embd)?;
10329 if debug_spec {
10330 eprintln!(" -> PARTIAL(replay={replay:?}), next_pred={last_pred}");
10331 }
10332 }
10333 if devacc_seeded {
10334 // stage (b) epilogue: fill_prev takes the gathered seed AFTER the refresh fills
10335 // consumed the old value (both slots carry the same value in every non-replay arm).
10336 e.copy_into(&mut fill_prev, 0, &h_seed_buf, n_embd)?;
10337 }
10338 if successor_valid {
10339 let optimistic_scratch_len = successor_attempt
10340 .as_ref()
10341 .expect("valid controller successor disappeared")
10342 .scratch_len;
10343 // The normal current-round commit refreshed/truncated the logical scratch tail.
10344 // Its optimistic successor row was already written physically, so restoring only
10345 // the retained logical length makes that row live for the carried round.
10346 scratch.set_len(e, optimistic_scratch_len)?;
10347 }
10348 if let Some(current) = current_opti.take() {
10349 opti_fork
10350 .as_mut()
10351 .ok_or("optipipe current retirement lost fork state")?
10352 .retire(current.generation)?;
10353 }
10354 if successor_valid {
10355 let successor = successor_attempt
10356 .take()
10357 .expect("valid controller successor disappeared before promotion");
10358 let generation = successor.generation;
10359 opti_fork
10360 .as_mut()
10361 .ok_or("optipipe successor promotion lost fork state")?
10362 .promote_successor_snapshot(&mut snap, generation);
10363 carried_opti = Some(successor);
10364 }
10365 if anatomy_on {
10366 // Commit/rollback is normally asynchronous on the primary/head stream. Bound it
10367 // only for this diagnostic so it does not disappear into the following draft's
10368 // first token readback.
10369 e.stream().synchronize()?;
10370 ph_commit += commit_started.elapsed().as_secs_f64();
10371 }
10372 // adaptive-K update (host math, zero syncs): next round drafts accepted-run + 1,
10373 // clamped to [floor(pos), k_cap]. cache.pos is post-rollback here (the round's
10374 // final position — the floor's position key reads the committed depth). Burst
10375 // rounds (`continue` above) draft the captured fixed depth and skip this, exactly
10376 // like gemma's burst arm.
10377 if adapt {
10378 let fl_now = floor_at(cache.pos);
10379 kc = (n_acc + 1).clamp(fl_now.min(k_cap), k_cap);
10380 }
10381 ph_mark(&mut ph_rest, phase_on);
10382 if let Some(p) = pipe {
10383 p.accept_end(round);
10384 }
10385 drop(pipe_accept);
10386 round += 1;
10387 // sse-cadence: this round's accepted drafts + bonus are committed (out is
10388 // append-only past step 4) — flush at round cadence.
10389 keep_going = flush_commit(&mut on_commit, &out, &mut flushed);
10390 }
10391 if let Some(mut ticket) = carried_opti.take() {
10392 opti_fork
10393 .as_mut()
10394 .ok_or("optipipe tail drain lost fork state")?
10395 .cancel_controller_ticket(e, &mut *cache, &mut *scratch, &snap, &mut ticket)?;
10396 }
10397 // sse-cadence: nothing below appends to `out`; flush any remainder (defensive).
10398 // (verdict ignored — the burst is over either way; the session tail runs unchanged.)
10399 let _ = flush_commit(&mut on_commit, &out, &mut flushed);
10400
10401 if spec_stats {
10402 let per_slot: Vec<String> = (0..k)
10403 .map(|j| {
10404 if st_drafted[j] > 0 {
10405 format!(
10406 "{}/{}={:.3}",
10407 st_accepted[j],
10408 st_drafted[j],
10409 st_accepted[j] as f64 / st_drafted[j] as f64
10410 )
10411 } else {
10412 "0/0".into()
10413 }
10414 })
10415 .collect();
10416 let acc = if total_drafted > 0 {
10417 total_accepted as f64 / total_drafted as f64
10418 } else {
10419 0.0
10420 };
10421 eprintln!(
10422 "[spec-stats] rounds={round} full_accept={st_full} len_hist={st_len_hist:?} \
10423 per_slot=[{}] total={total_accepted}/{total_drafted}={acc:.3} \
10424 tok_per_round={:.3}",
10425 per_slot.join(" "),
10426 (total_accepted + round) as f64 / round.max(1) as f64
10427 );
10428 }
10429 if constraint.is_some() {
10430 eprintln!(
10431 "[draft-mask] mask_rounds={dm_rounds} clone_total={:.3}ms \
10432 clone_per_round={:.4}ms gram_cuts={dm_cuts}/{round} cut_tokens={dm_cut_tokens}",
10433 dm_clone_ns as f64 / 1e6,
10434 dm_clone_ns as f64 / 1e6 / dm_rounds.max(1) as f64
10435 );
10436 }
10437 if phase_on {
10438 let tot = ph_draft + ph_verify + ph_wait + ph_rest;
10439 eprintln!(
10440 "[spec-phase] draft={:.1}ms ({:.1}%) verify-issue={:.1}ms ({:.1}%) verify-wait={:.1}ms ({:.1}%) commit-host={:.1}ms ({:.1}%) rounds={round}",
10441 ph_draft * 1e3,
10442 ph_draft / tot * 100.0,
10443 ph_verify * 1e3,
10444 ph_verify / tot * 100.0,
10445 ph_wait * 1e3,
10446 ph_wait / tot * 100.0,
10447 ph_rest * 1e3,
10448 ph_rest / tot * 100.0
10449 );
10450 }
10451 if anatomy_on {
10452 let rounds_f = round.max(1) as f64;
10453 let other = (ph_rest - ph_commit).max(0.0);
10454 eprintln!(
10455 "[spec-anatomy] per-round draft={:.3}ms pp-verify={:.3}ms \
10456 verify-accept={:.3}ms commit-rollback={:.3}ms other={:.3}ms rounds={round}",
10457 ph_draft * 1e3 / rounds_f,
10458 ph_verify * 1e3 / rounds_f,
10459 ph_wait * 1e3 / rounds_f,
10460 ph_commit * 1e3 / rounds_f,
10461 other * 1e3 / rounds_f,
10462 );
10463 }
10464 let _pipe_tail = pipe.map(|p| p.primary());
10465 // SESSION TAIL: leave the session in the exact invariant the next turn's suffix prime
10466 // expects — every row in `committed` has trunk KV/recur state AND an exact draft-KV row.
10467 // Park the draft-graph ctx back on the session (the serve-burst fixed-cost fix): the next
10468 // burst replays instead of recapturing. Error paths (`?` above) drop it — recaptured then.
10469 if let Some(slot) = sess_draft_slot.take() {
10470 *slot = Some(dctx);
10471 }
10472 let t_rounds = t_ent.elapsed();
10473 if let Some((committed, last_h, next_pred_slot, sctr_slot, uctr_slot)) = sess_tail.take() {
10474 // NEXT BURST'S BOUNDARY TOKEN (lane/sampled-spec-quality, Item 1). Greedy stashes
10475 // the argmax `last_pred` exactly as before (byte contract). SAMPLED draws the token
10476 // HERE, where the sampler, the session Philox counters and the penalty window are
10477 // all live and the boundary logits row still exists — that is the "make the state
10478 // available" half of the fix; the consuming burst then just emits it. `sctr` is
10479 // written to the session BELOW the draws so the advance is never lost.
10480 *next_pred_slot = Some(last_pred);
10481 let sample_boundary = sampled && constraint.is_none() && spec_sampled_boundary_on();
10482 let mut stashed_pending = false;
10483 if let Some(b) = pending.take() {
10484 if !sampled {
10485 // PENDING-CARRY (2026-08-01): stash the bonus on the session instead of
10486 // committing it with a solo T=1 pass — the next empty-suffix greedy burst
10487 // consumes it as round-0 verify col 0 (a plain round edge; the old tail
10488 // commit + next burst's init feed were 11.6+11.5ms solo trunk passes per
10489 // burst on H100 q27, [spec-setup] trace). b stays in `out` (emitted) but
10490 // OUT of `committed` (cache rows == committed); the consuming call
10491 // prepends it once its verify commits the row. next_pred is unknowable
10492 // without the commit pass — None; callers gate on pending_tok too.
10493 debug_assert_eq!(out.last(), Some(&b), "pending must be the last emitted");
10494 if let Some(slot) = sess_pending_slot.take() {
10495 *slot = Some(b);
10496 }
10497 *next_pred_slot = None;
10498 // fill_prev = hidden of the last COMMITTED row (b's predecessor) — the
10499 // exact chain-seed/fill anchor the consuming burst (or a flush) needs.
10500 *last_h = Some(e.clone_dtod(&fill_prev)?);
10501 stashed_pending = true;
10502 } else {
10503 // SAMPLED tail (unchanged): commit the bonus (one T=1 pass) + draft fill —
10504 // the sampled round-0 accept needs this pass's logits (last_col_logits).
10505 let pos_b = cache.pos;
10506 scratch.set_len(e, pos_b)?;
10507 let (lg_b, hb) = self.spec_target_step_h(e, b, &mut *cache)?;
10508 // after a FULL-accept exit `last_pred` is STALE (it predicted the bonus
10509 // itself — the prediction AFTER the bonus never materialized; it would have
10510 // been the next round's verify col 0). The commit's logits ARE that
10511 // prediction — so they are also the row the next burst's boundary token
10512 // comes off, and (lane/sampled-spec-quality) it is DRAWN from them here.
10513 *next_pred_slot = Some(if sample_boundary {
10514 sample_boundary_token(
10515 e,
10516 &lg_b,
10517 &sp,
10518 &pen_hist,
10519 &mut sctr,
10520 "burst-tail-commit",
10521 )?
10522 } else {
10523 argmax(&lg_b) as u32
10524 });
10525 self.mtp_kv_fill(e, mtp, &[b], &fill_prev, pos_b, &mut *scratch, embd_dev)?;
10526 *last_h = Some(hb);
10527 }
10528 } else {
10529 // fill_prev tracks the hidden of the last COMMITTED row throughout the loop.
10530 *last_h = Some(e.clone_dtod(&fill_prev)?);
10531 if sample_boundary {
10532 // No pending to commit, so the boundary row is the one `last_pred` was
10533 // argmaxed from and the sampled path keeps it on device: the init feed's
10534 // logits when the burst ran zero rounds, else the legacy-replay path's
10535 // last verify column (both predict the token AFTER the last committed
10536 // row). It is retained precisely because round 0's accept test needs it,
10537 // so the draw costs no extra D2H of the [n_vocab] row.
10538 match last_col_logits.as_ref() {
10539 Some(lc) => {
10540 *next_pred_slot = Some(sample_boundary_token_dev(
10541 e,
10542 lc,
10543 n_vocab,
10544 &sp,
10545 &pen_hist,
10546 &mut sctr,
10547 "burst-tail-nopending",
10548 )?);
10549 }
10550 // NAME THE FALLBACK (house standard): unreachable today — a sampled
10551 // burst always feeds or replays, so the row exists — but if it ever
10552 // is, the stream takes a greedy token and SAYS so rather than
10553 // silently regressing to the pre-lane behaviour.
10554 None => eprintln!(
10555 "[spec-boundary] sampled tail kept the ARGMAX boundary token \
10556 (reason: no retained boundary logits row)"
10557 ),
10558 }
10559 }
10560 }
10561 *sctr_slot = sctr;
10562 *uctr_slot = uctr;
10563 committed.extend_from_slice(prompt);
10564 if let Some(cb) = carried_pending {
10565 // the consumed carry's cache row landed in round 0's verify (every pending
10566 // round commits col 0) — it joins `committed` here, in sequence order.
10567 committed.push(cb);
10568 }
10569 if stashed_pending {
10570 // ZERO-EMIT BURST (2026-08-06 c=8 serve panic, pre-existing since b4aea184):
10571 // `out.len() - 1` underflowed on an EMPTY `out` — "range end index
10572 // 18446744073709551615 out of range for slice of length 0", killing the
10573 // memra-gpu-worker and failing 31 of 32 concurrent requests with "worker closed
10574 // stream". Reachable because `pending` starts as `carried_pending` (a bonus
10575 // stashed by the PREVIOUS burst) while `out` starts empty, and the carry is
10576 // deliberately NOT pushed to `out` (line ~3239: the burst that emitted it already
10577 // did). So a burst that stashes a pending without emitting anything of its own —
10578 // the round loop exits before a push, e.g. the ring drain's `out.len() < max_new`
10579 // guard skipping every token under a tight budget — arrives here with
10580 // out.len() == 0 and stashed_pending == true.
10581 //
10582 // The invariant is unchanged: `committed` gets every emitted token EXCEPT the
10583 // stashed bonus. With nothing emitted, that is nothing — and the carry pushed
10584 // just above is already accounted. Saturating, not a min/assert: an empty `out`
10585 // here is a legitimate burst shape, not a corrupt state.
10586 let emitted = out.len().saturating_sub(1);
10587 committed.extend_from_slice(&out[..emitted]);
10588 } else {
10589 committed.extend_from_slice(&out); // FULL out incl. overshoot — all committed
10590 }
10591 debug_assert_eq!(
10592 cache.pos,
10593 committed.len(),
10594 "session invariant: cache rows == committed tokens"
10595 );
10596 if setup_trace {
10597 e.stream().synchronize()?; // bound the async tail fill in the trace
10598 let t_tail = t_ent.elapsed();
10599 eprintln!(
10600 "[spec-setup] init={:.2}ms cap={:.2}ms fill={:.2}ms rounds={:.2}ms tail={:.2}ms total={:.2}ms out={} cont={}",
10601 t_init.as_secs_f64() * 1e3,
10602 (t_cap - t_init).as_secs_f64() * 1e3,
10603 (t_fill - t_cap).as_secs_f64() * 1e3,
10604 (t_rounds - t_fill).as_secs_f64() * 1e3,
10605 (t_tail - t_rounds).as_secs_f64() * 1e3,
10606 t_tail.as_secs_f64() * 1e3,
10607 out.len(),
10608 continuation
10609 );
10610 }
10611 return Ok((out, total_drafted, total_accepted));
10612 }
10613 out.truncate(max_new);
10614 Ok((out, total_drafted, total_accepted))
10615 }
10616
10617 /// Anchor-bounded DSpark target extraction. The trunk sees the exact generated token tape;
10618 /// only requested hidden rows and target-logit rows cross PCIe. An anchor token at p pairs
10619 /// with the pre-output-norm h[p-1] carrier, exactly as the existing replay/NextN path does.
10620 pub fn extract_dspark_anchors(
10621 &self,
10622 e: &Engine,
10623 tokens: &[u32],
10624 anchor_positions: &[usize],
10625 gamma: usize,
10626 top_k: usize,
10627 chunk: usize,
10628 temperature: f32,
10629 ) -> Result<Vec<DsparkAnchorRecord>, Box<dyn std::error::Error>> {
10630 if tokens.len() < gamma + 2 || gamma == 0 || chunk < 2 {
10631 return Err("DSpark extraction token tape/gamma/chunk is invalid".into());
10632 }
10633 if anchor_positions.windows(2).any(|pair| pair[0] >= pair[1]) {
10634 return Err("DSpark anchor positions must be sorted and unique".into());
10635 }
10636 for &position in anchor_positions {
10637 if position == 0 || position + gamma >= tokens.len() {
10638 return Err(format!(
10639 "DSpark anchor {position} has no predecessor or cannot cover gamma={gamma} in {} tokens",
10640 tokens.len()
10641 )
10642 .into());
10643 }
10644 }
10645
10646 let n_vocab = self.output.out_features();
10647 let n_embd = self.cfg.n_embd as usize;
10648 let mut cache = crate::pp::new_cache(e, &self.cfg, tokens.len() + gamma + 8)?;
10649 let (embd_qt, embd_rb) = self.embd.qt_and_row_bytes(n_embd);
10650 let embd_gpu = if spec_host_embd() {
10651 None
10652 } else {
10653 Some(
10654 self.embd_gpu
10655 .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload")),
10656 )
10657 };
10658 let embd_dev = embd_gpu.map(|gpu| (gpu, embd_qt, embd_rb));
10659
10660 struct PendingRecord {
10661 position: usize,
10662 hidden: Option<Vec<f32>>,
10663 tokens: Vec<u32>,
10664 target_top_ids: Vec<Option<Vec<u32>>>,
10665 target_top_logits: Vec<Option<Vec<f32>>>,
10666 target_top_probs: Vec<Option<Vec<f32>>>,
10667 target_tail_probs: Vec<Option<f32>>,
10668 }
10669
10670 let mut pending: Vec<PendingRecord> = anchor_positions
10671 .iter()
10672 .map(|&position| PendingRecord {
10673 position,
10674 hidden: None,
10675 tokens: tokens[position..=position + gamma].to_vec(),
10676 target_top_ids: vec![None; gamma],
10677 target_top_logits: vec![None; gamma],
10678 target_top_probs: vec![None; gamma],
10679 target_tail_probs: vec![None; gamma],
10680 })
10681 .collect();
10682
10683 let mut start = 0usize;
10684 while start < tokens.len() {
10685 let end = (start + chunk).min(tokens.len());
10686 let chunk_tokens = &tokens[start..end];
10687 let (target_logits, hidden_rows) =
10688 self.decode_step_t_core(e, chunk_tokens, start, &mut cache, embd_dev, None)?;
10689 for record in &mut pending {
10690 let hidden_position = record.position - 1;
10691 if hidden_position >= start && hidden_position < end {
10692 let local = hidden_position - start;
10693 record.hidden = Some(
10694 e.dtoh_view(&hidden_rows.slice(local * n_embd..(local + 1) * n_embd))?,
10695 );
10696 }
10697 for slot in 0..gamma {
10698 let target_row = record.position + slot;
10699 if target_row < start || target_row >= end {
10700 continue;
10701 }
10702 let local = target_row - start;
10703 let logits =
10704 e.dtoh_view(&target_logits.slice(local * n_vocab..(local + 1) * n_vocab))?;
10705 let (ids, top_logits, probs, tail) =
10706 dspark_sparse_softmax_topk(&logits, top_k, temperature)?;
10707 record.target_top_ids[slot] = Some(ids);
10708 record.target_top_logits[slot] = Some(top_logits);
10709 record.target_top_probs[slot] = Some(probs);
10710 record.target_tail_probs[slot] = Some(tail);
10711 }
10712 }
10713 start = end;
10714 }
10715
10716 pending
10717 .into_iter()
10718 .map(|record| {
10719 let hidden = record
10720 .hidden
10721 .ok_or_else(|| format!("missing DSpark hidden at {}", record.position))?;
10722 let target_top_ids =
10723 flatten_dspark_rows(record.target_top_ids, record.position, "target ids")?;
10724 let target_top_logits = flatten_dspark_rows(
10725 record.target_top_logits,
10726 record.position,
10727 "target logits",
10728 )?;
10729 let target_top_probs =
10730 flatten_dspark_rows(record.target_top_probs, record.position, "target probs")?;
10731 let target_tail_probs = record
10732 .target_tail_probs
10733 .into_iter()
10734 .enumerate()
10735 .map(|(slot, value)| {
10736 value.ok_or_else(|| {
10737 format!("missing DSpark tail at {} slot {slot}", record.position)
10738 })
10739 })
10740 .collect::<Result<Vec<_>, _>>()?;
10741 Ok(DsparkAnchorRecord {
10742 position: record.position,
10743 hidden,
10744 tokens: record.tokens,
10745 target_top_ids,
10746 target_top_logits,
10747 target_top_probs,
10748 target_tail_probs,
10749 })
10750 })
10751 .collect()
10752 }
10753
10754 /// TEACHER-FORCED REPLAY ACCEPTANCE (hqmtp MTP-heal protocol): walk a FIXED token
10755 /// sequence and, at sampled positions, compare the MTP head's K-token draft chain against
10756 /// the trunk's own teacher-forced greedy predictions. Nothing is generated — the context is
10757 /// the corpus text itself, so (a) degenerate self-generated loops cannot inflate acceptance
10758 /// and (b) two arms (bf16 ceiling vs NVFP4) score on IDENTICAL contexts, isolating the
10759 /// quant-induced head/hidden-state mismatch from text drift.
10760 ///
10761 /// Per eval position p (context = tokens[0..=p], predecessor pairing as in spec decode):
10762 /// draft_j = chain token j from (tokens[p], h_{p-1}), then its own drafts — the exact
10763 /// eager spec-decode chain (same mtp_head_forward_dev, same rope positions).
10764 /// target_j = teacher-forced greedy pick for position p+1+j (argmax of the trunk logits
10765 /// at forced context tokens[0..p+j]). For j==0 this equals live spec
10766 /// acceptance; for j>=1 live verify would condition on the drafts, here it
10767 /// conditions on the corpus — deterministic and arm-comparable by design.
10768 ///
10769 /// Returns (rows, bg): one (p, drafts[k], targets[k]) row per eval position (ascending p),
10770 /// plus the full teacher-forced greedy track bg (bg[i] = greedy pick for position i, i>=1)
10771 /// so harnesses can cross-check runs (e.g. different chunk sizes must give identical bg).
10772 ///
10773 /// `hdump`: when Some, every position's pre-output_norm trunk hidden (the exact rows the
10774 /// draft-KV fill pairs from) streams to the file as little-endian f32 [t_total, n_embd] —
10775 /// the head-distillation extraction (hqmtp): the ENGINE is the source of truth for trunk
10776 /// hiddens (HF torch reproductions of the hybrid trunk measured only ~0.5 greedy
10777 /// agreement vs this path — not usable as a training-data source).
10778 pub fn replay_acceptance(
10779 &self,
10780 e: &Engine,
10781 tokens: &[u32],
10782 k: usize,
10783 stride: usize,
10784 chunk: usize,
10785 mut hdump: Option<&mut std::fs::File>,
10786 ) -> Result<(Vec<(usize, Vec<u32>, Vec<u32>)>, Vec<u32>), Box<dyn std::error::Error>> {
10787 assert!(k >= 1 && stride >= 1 && chunk >= 2);
10788 let mtp = self
10789 .mtp
10790 .as_ref()
10791 .expect("replay_acceptance requires an MTP head");
10792 let n_vocab = self.output.out_features();
10793 let d_vocab = mtp
10794 .shared_head_head
10795 .as_ref()
10796 .unwrap_or(&self.output)
10797 .out_features();
10798 let n_embd = self.cfg.n_embd as usize;
10799 let t_total = tokens.len();
10800 assert!(t_total >= 8, "corpus too short ({t_total} tokens)");
10801 // STAGE-OWNED KV (lane/pp2-spec 2026-08-06) — see `new_session`. Door shut = `Cache::new`.
10802 let mut cache = crate::pp::new_cache(e, &self.cfg, t_total + k + 8)?;
10803 let mut scratch = MtpScratch::new(
10804 e,
10805 &self.cfg,
10806 t_total + k + 8,
10807 self.mtp.as_ref().and_then(|m| m.geom.as_ref()),
10808 )?;
10809 let (embd_qt, embd_rb) = self.embd.qt_and_row_bytes(n_embd);
10810 let embd_gpu = if spec_host_embd() {
10811 None
10812 } else {
10813 Some(
10814 self.embd_gpu
10815 .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload")),
10816 )
10817 };
10818 let embd_dev = embd_gpu.map(|g| (g, embd_qt, embd_rb));
10819
10820 // bg[i] = the trunk's greedy pick for position i under the forced context (i >= 1).
10821 let mut bg: Vec<u32> = vec![0; t_total + 1];
10822 let mut rows: Vec<(usize, Vec<u32>, Vec<u32>)> = Vec::new();
10823 let mut prev_last_h = e.zeros(n_embd)?; // predecessor hidden entering the chunk
10824 let mut seed_buf = e.zeros(n_embd)?;
10825 let mut preds_d = e.alloc_u32_zeroed(chunk)?;
10826 let nll_on = std::env::var("MEMRA_REPLAY_NLL").as_deref() == Ok("1");
10827 let (mut nll_sum, mut nll_cnt) = (0f64, 0u64);
10828 let mut s = 0usize;
10829 while s < t_total {
10830 let cend = (s + chunk).min(t_total);
10831 let tc = cend - s;
10832 let ch = &tokens[s..cend];
10833 // 1. forced trunk pass — verify path (decode-exact contract): all-column logits +
10834 // the chunk's true hiddens.
10835 let (tl_d, vx) = self.decode_step_t_core(e, ch, s, &mut cache, embd_dev, None)?;
10836 for j in 0..tc {
10837 e.argmax_token_device_col(&tl_d, j, n_vocab, &mut preds_d, j)?;
10838 }
10839 let preds = e.dtoh_u32(&preds_d)?;
10840 for j in 0..tc {
10841 bg[s + j + 1] = preds[j];
10842 }
10843 // MEMRA_REPLAY_NLL=1: teacher-forced NLL/perplexity over the same forced pass — the
10844 // checkpoint-quality metric (position j's logits score the GOLD next token).
10845 if nll_on {
10846 let jmax = if cend < t_total { tc } else { tc - 1 }; // last pos has no gold next
10847 if jmax > 0 {
10848 let ids: Vec<u32> = (0..jmax).map(|j| tokens[s + j + 1]).collect();
10849 let rows: Vec<i32> = (0..jmax as i32).collect();
10850 let idsd = e.htod_u32_v(&ids)?;
10851 let rowsd = e.htod_i32(&rows)?;
10852 let mut outd = e.zeros(jmax)?;
10853 e.softmax_gather(&tl_d, n_vocab, &idsd, &rowsd, &mut outd, n_vocab, jmax, 1.0)?;
10854 for pr in e.dtoh(&outd)? {
10855 nll_sum += -((pr.max(1e-30)) as f64).ln();
10856 nll_cnt += 1;
10857 }
10858 }
10859 }
10860 if let Some(f) = hdump.as_deref_mut() {
10861 use std::io::Write;
10862 let host: Vec<f32> = e.dtoh(&vx)?;
10863 // bf16 round-to-nearest-even — f32 doubled the disk bill at bulk
10864 // extraction scale (20M tokens x 4096 = 320GB f32 vs 160GB bf16).
10865 let mut bytes = Vec::with_capacity(tc * n_embd * 2);
10866 for v in &host[..tc * n_embd] {
10867 let b = v.to_bits();
10868 let r = b.wrapping_add(0x7FFF + ((b >> 16) & 1));
10869 bytes.extend_from_slice(&((r >> 16) as u16).to_le_bytes());
10870 }
10871 f.write_all(&bytes)?;
10872 }
10873 // CHAINLESS extraction (stride > corpus, the bulk-hdump mode): no chunk ever
10874 // drafts, so the draft-KV fills are pure waste — skip them (2 MTP-block passes
10875 // per token saved; the forced trunk pass + hdump is all the mode needs).
10876 let chainless = stride > t_total;
10877 if chainless {
10878 e.copy_view_into(
10879 &mut prev_last_h,
10880 0,
10881 &vx.slice((tc - 1) * n_embd..tc * n_embd),
10882 n_embd,
10883 )?;
10884 s = cend;
10885 continue;
10886 }
10887 // 2. TRUE predecessor-paired draft-KV fill for the chunk (row i carries h_{i-1};
10888 // row s reads the previous chunk's last true hidden, zeros at corpus start).
10889 let mut vxs = e.zeros(tc * n_embd)?;
10890 e.copy_into(&mut vxs, 0, &prev_last_h, n_embd)?;
10891 if tc > 1 {
10892 e.copy_view_into(
10893 &mut vxs,
10894 n_embd,
10895 &vx.slice(0..(tc - 1) * n_embd),
10896 (tc - 1) * n_embd,
10897 )?;
10898 }
10899 scratch.set_len(e, s)?;
10900 self.mtp_kv_fill(e, mtp, ch, &vxs, s, &mut scratch, embd_dev)?;
10901 // 3. draft chains at sampled positions, DESCENDING: a chain reads only slots
10902 // [0..p) (true fills) and appends at >= p; the next (smaller-p) chain's set_len
10903 // truncates those approximate appends before they can ever be read.
10904 let ps: Vec<usize> = (s..cend)
10905 .filter(|p| *p >= 1 && *p % stride == 0 && *p + k <= t_total)
10906 .collect();
10907 for &p in ps.iter().rev() {
10908 scratch.set_len(e, p)?;
10909 if p == s {
10910 e.copy_into(&mut seed_buf, 0, &prev_last_h, n_embd)?;
10911 } else {
10912 e.copy_view_into(
10913 &mut seed_buf,
10914 0,
10915 &vx.slice((p - 1 - s) * n_embd..(p - s) * n_embd),
10916 n_embd,
10917 )?;
10918 }
10919 let mut e_tok = tokens[p];
10920 let mut d_seed = e.clone_dtod(&seed_buf)?;
10921 let mut drafts: Vec<u32> = Vec::with_capacity(k);
10922 for j in 0..k {
10923 let (dl_d, h_nextn) = self.mtp_head_forward_dev(
10924 e,
10925 mtp,
10926 e_tok,
10927 &d_seed,
10928 &mut scratch,
10929 p + 1 + j,
10930 embd_dev,
10931 None, // acceptance-oracle walk: no grammar
10932 )?;
10933 let tok_d = e.argmax_token_device(&dl_d, d_vocab)?;
10934 let idx = e.dtoh_u32_one(&tok_d)?;
10935 let d = match &mtp.d2t {
10936 Some(map) => map[idx as usize],
10937 None => idx,
10938 };
10939 drafts.push(d);
10940 e_tok = d;
10941 d_seed = h_nextn;
10942 }
10943 // targets may live in a LATER chunk's bg — resolved after the walk.
10944 rows.push((p, drafts, Vec::new()));
10945 }
10946 // 4. restore TRUE entries for the whole chunk (the next chunk's chains and fills
10947 // expect scratch.len == cend with exact rows).
10948 scratch.set_len(e, s)?;
10949 self.mtp_kv_fill(e, mtp, ch, &vxs, s, &mut scratch, embd_dev)?;
10950 e.copy_view_into(
10951 &mut prev_last_h,
10952 0,
10953 &vx.slice((tc - 1) * n_embd..tc * n_embd),
10954 n_embd,
10955 )?;
10956 s = cend;
10957 }
10958 for (p, drafts, targets) in rows.iter_mut() {
10959 for j in 0..drafts.len() {
10960 targets.push(bg[*p + 1 + j]);
10961 }
10962 }
10963 rows.sort_by_key(|r| r.0);
10964 if nll_cnt > 0 {
10965 let mean = nll_sum / nll_cnt as f64;
10966 println!(
10967 "[replay-nll] tokens={nll_cnt} nll/token={mean:.5} ppl={:.4}",
10968 mean.exp()
10969 );
10970 }
10971 Ok((rows, bg))
10972 }
10973}
10974
10975#[cfg(test)]
10976mod dspark_sparse_tests {
10977 use super::dspark_sparse_softmax_topk;
10978
10979 #[test]
10980 fn topk_keeps_full_softmax_mass_and_stable_ties() {
10981 let logits = [1.0f32, 3.0, 3.0, -2.0];
10982 let (ids, top_logits, probs, tail) = dspark_sparse_softmax_topk(&logits, 2, 1.0).unwrap();
10983 assert_eq!(ids, vec![1, 2]);
10984 assert_eq!(top_logits, vec![3.0, 3.0]);
10985 let denominator = logits.iter().map(|value| (value - 3.0).exp()).sum::<f32>();
10986 let expected = 1.0 / denominator;
10987 assert!((probs[0] - expected).abs() < 1.0e-6);
10988 assert!((probs[1] - expected).abs() < 1.0e-6);
10989 assert!((tail - (1.0 - 2.0 * expected)).abs() < 1.0e-6);
10990 assert!((probs.iter().sum::<f32>() + tail - 1.0).abs() < 1.0e-6);
10991 }
10992}
10993
10994#[cfg(test)]
10995mod spec_replay_env_tests {
10996 use super::spec_replay_env_on;
10997
10998 #[test]
10999 fn replay_requires_literal_one() {
11000 assert!(!spec_replay_env_on(None));
11001 assert!(!spec_replay_env_on(Some("")));
11002 assert!(!spec_replay_env_on(Some("0")));
11003 assert!(!spec_replay_env_on(Some("true")));
11004 assert!(!spec_replay_env_on(Some("2")));
11005 assert!(spec_replay_env_on(Some("1")));
11006 }
11007}
11008
11009#[cfg(test)]
11010mod telem_tests {
11011 use super::{SPEC_TELEM_POS, SpecTelemetry, SpecTelemetryCounters};
11012
11013 #[test]
11014 fn synthetic_accept_masks_produce_tau_and_position_histogram() {
11015 let counters = SpecTelemetryCounters::default();
11016 for mask in [
11017 [true, true, true],
11018 [true, true, false],
11019 [true, false, false],
11020 [false, false, false],
11021 ] {
11022 let accepted = mask.iter().take_while(|&&value| value).count();
11023 counters.record_round(mask.len(), accepted);
11024 }
11025
11026 let snapshot = counters.snapshot();
11027 assert_eq!(
11028 (snapshot.rounds, snapshot.drafted, snapshot.accepted),
11029 (4, 12, 6)
11030 );
11031 assert_eq!(&snapshot.pos_drafted[..3], &[4, 4, 4]);
11032 assert_eq!(&snapshot.pos_accepted[..3], &[3, 2, 1]);
11033 assert_eq!(snapshot.tau(), 1.5);
11034 assert_eq!(snapshot.pos_drafted[3..], [0; SPEC_TELEM_POS - 3]);
11035 assert_eq!(snapshot.pos_accepted[3..], [0; SPEC_TELEM_POS - 3]);
11036 }
11037
11038 /// The worker's per-burst pattern: stash, accumulate, diff — the delta must isolate
11039 /// exactly the burst's contribution (pool-resumed sessions carry prior requests' counts).
11040 #[test]
11041 fn delta_isolates_burst_contribution() {
11042 let mut t = SpecTelemetry::default();
11043 // "previous request": 2 rounds of k=3, accepts 3 then 1.
11044 for (kr, na) in [(3usize, 3usize), (3, 1)] {
11045 t.rounds += 1;
11046 t.drafted += kr as u64;
11047 t.accepted += na as u64;
11048 for j in 0..kr {
11049 t.pos_drafted[j] += 1;
11050 }
11051 for j in 0..na {
11052 t.pos_accepted[j] += 1;
11053 }
11054 }
11055 let before = t;
11056 // "this burst": 1 round k=3, accepts 2.
11057 t.rounds += 1;
11058 t.drafted += 3;
11059 t.accepted += 2;
11060 for j in 0..3 {
11061 t.pos_drafted[j] += 1;
11062 }
11063 for j in 0..2 {
11064 t.pos_accepted[j] += 1;
11065 }
11066 let d = t.delta_since(&before);
11067 assert_eq!((d.rounds, d.drafted, d.accepted), (1, 3, 2));
11068 assert_eq!(&d.pos_drafted[..3], &[1, 1, 1]);
11069 assert_eq!(&d.pos_accepted[..3], &[1, 1, 0]);
11070 assert_eq!(d.pos_drafted[3..], [0; SPEC_TELEM_POS - 3]);
11071 }
11072
11073 /// merge(delta) then merge(delta2) equals accumulating both — the per-model /metrics
11074 /// aggregation invariant.
11075 #[test]
11076 fn merge_accumulates_fieldwise() {
11077 let mut agg = SpecTelemetry::default();
11078 let mut d1 = SpecTelemetry {
11079 rounds: 2,
11080 drafted: 6,
11081 accepted: 4,
11082 ..Default::default()
11083 };
11084 d1.pos_drafted[0] = 2;
11085 d1.pos_accepted[0] = 2;
11086 let mut d2 = SpecTelemetry {
11087 rounds: 1,
11088 drafted: 3,
11089 accepted: 1,
11090 ..Default::default()
11091 };
11092 d2.pos_drafted[0] = 1;
11093 d2.pos_accepted[0] = 1;
11094 d2.pos_drafted[1] = 1;
11095 agg.merge(&d1);
11096 agg.merge(&d2);
11097 assert_eq!((agg.rounds, agg.drafted, agg.accepted), (3, 9, 5));
11098 assert_eq!(agg.pos_drafted[0], 3);
11099 assert_eq!(agg.pos_accepted[0], 3);
11100 assert_eq!(agg.pos_drafted[1], 1);
11101 assert_eq!(agg.pos_accepted[1], 0);
11102 }
11103
11104 /// Wrong-snapshot diff saturates to zero instead of wrapping — the counters feed a
11105 /// public metrics surface and must never publish a u64-wrapped garbage value.
11106 #[test]
11107 fn delta_saturates_never_wraps() {
11108 let small = SpecTelemetry {
11109 rounds: 1,
11110 drafted: 2,
11111 accepted: 1,
11112 ..Default::default()
11113 };
11114 let big = SpecTelemetry {
11115 rounds: 5,
11116 drafted: 15,
11117 accepted: 9,
11118 ..Default::default()
11119 };
11120 let d = small.delta_since(&big);
11121 assert_eq!((d.rounds, d.drafted, d.accepted), (0, 0, 0));
11122 }
11123}
11124
11125#[cfg(test)]
11126mod opti_fork_tests {
11127 use super::{
11128 OptiControllerPolicy, OptiForkAction, OptiForkGateMode, OptiForkGenerationTracker,
11129 };
11130
11131 #[test]
11132 fn controller_threshold_and_three_miss_breaker_are_exact() {
11133 let mut policy = OptiControllerPolicy {
11134 threshold: 0.7,
11135 consecutive_misses: 0,
11136 breaker_tripped: false,
11137 };
11138 assert!(!policy.admit(0.699_999));
11139 assert!(policy.admit(0.7));
11140 assert!(!policy.resolve(false));
11141 assert!(!policy.resolve(false));
11142 assert!(policy.resolve(false));
11143 assert!(policy.breaker_tripped);
11144 assert!(!policy.admit(1.0));
11145 assert!(
11146 !policy.resolve(true),
11147 "a resolved hit cannot re-arm a tripped request"
11148 );
11149 assert!(policy.breaker_tripped);
11150 }
11151
11152 #[test]
11153 fn zero_threshold_is_the_true_unconditional_measurement_arm() {
11154 let mut policy = OptiControllerPolicy {
11155 threshold: 0.0,
11156 consecutive_misses: 0,
11157 breaker_tripped: false,
11158 };
11159 for _ in 0..16 {
11160 assert!(policy.admit(0.0));
11161 assert!(!policy.resolve(false));
11162 }
11163 for invalid in [f32::NAN, f32::INFINITY, -0.01, 1.01] {
11164 assert!(
11165 !policy.admit(invalid),
11166 "invalid q proxy must fail closed: {invalid}"
11167 );
11168 }
11169 assert!(!policy.breaker_tripped);
11170 assert_eq!(policy.consecutive_misses, 0);
11171 }
11172
11173 #[test]
11174 fn alternating_mode_flips_by_generation_not_round_parity() {
11175 assert_eq!(OptiForkGateMode::Alternate.action(0), OptiForkAction::Hit);
11176 assert_eq!(OptiForkGateMode::Alternate.action(1), OptiForkAction::Miss);
11177 assert_eq!(OptiForkGateMode::Alternate.action(8), OptiForkAction::Hit);
11178 assert_eq!(OptiForkGateMode::Alternate.action(9), OptiForkAction::Miss);
11179 }
11180
11181 #[test]
11182 fn live_generation_cannot_be_overwritten() {
11183 let mut tracker = OptiForkGenerationTracker::default();
11184 let g0 = tracker.reserve().unwrap();
11185 let g1 = tracker.reserve().unwrap();
11186 let err = tracker.reserve().unwrap_err().to_string();
11187 assert!(
11188 err.contains("still owns generation 0"),
11189 "unexpected error: {err}"
11190 );
11191 tracker.retire(g0).unwrap();
11192 let g2 = tracker.reserve().unwrap();
11193 assert_eq!((g2.id, g2.slot), (2, 0));
11194 tracker.retire(g1).unwrap();
11195 tracker.retire(g2).unwrap();
11196 }
11197
11198 #[test]
11199 fn teardown_rejects_a_stale_generation_tag() {
11200 let mut tracker = OptiForkGenerationTracker::default();
11201 let g0 = tracker.reserve().unwrap();
11202 tracker.retire(g0).unwrap();
11203 let err = tracker.retire(g0).unwrap_err().to_string();
11204 assert!(err.contains("teardown mismatch"), "unexpected error: {err}");
11205 }
11206}
11207
11208#[cfg(test)]
11209mod draft_graph_fallback_tests {
11210 use super::DraftGraphFallback;
11211
11212 /// The Q2 contract, part (a): a fallback flip is LOUD — exactly once per flip.
11213 #[test]
11214 fn flip_is_loud_once_and_memoized_after() {
11215 let mut f = DraftGraphFallback::default();
11216 let line = f
11217 .mark_greedy("out of memory")
11218 .expect("first flip must return the warn line");
11219 assert!(
11220 line.contains("WARN"),
11221 "flip line must be warn-level: {line}"
11222 );
11223 assert!(
11224 line.contains("out of memory"),
11225 "flip line must carry the reason: {line}"
11226 );
11227 assert!(f.greedy_failed());
11228 // re-marking an already-failed graph is the memoization: quiet, still failed.
11229 assert!(f.mark_greedy("out of memory").is_none());
11230 assert!(f.greedy_failed());
11231 // the two graphs' flags are independent (greedy flip leaves sampled capturable).
11232 assert!(!f.sampled_failed());
11233 let line_s = f
11234 .mark_sampled("capture unsupported")
11235 .expect("sampled flip is its own flip");
11236 assert!(
11237 line_s.contains("sampled"),
11238 "sampled flip names itself: {line_s}"
11239 );
11240 assert!(f.mark_sampled("capture unsupported").is_none());
11241 }
11242
11243 /// The Q2 contract, part (b): resume-from-pool RESETS both flags (fresh capture chance),
11244 /// and says so exactly when there was something to reset.
11245 #[test]
11246 fn reset_on_resume_clears_flags_and_logs_once() {
11247 let mut f = DraftGraphFallback::default();
11248 // clean session: resume is silent, nothing to reset.
11249 assert!(f.reset_on_resume().is_none());
11250 f.mark_greedy("oom").unwrap();
11251 f.mark_sampled("oom").unwrap();
11252 let note = f
11253 .reset_on_resume()
11254 .expect("a set flag must produce the reset note");
11255 assert!(
11256 note.contains("greedy+sampled"),
11257 "note names what was reset: {note}"
11258 );
11259 assert!(
11260 !f.greedy_failed() && !f.sampled_failed(),
11261 "both flags cleared"
11262 );
11263 // and the NEXT failure after a reset is a fresh flip — loud again.
11264 assert!(f.mark_greedy("oom again").is_some());
11265 let note2 = f.reset_on_resume().expect("greedy-only reset");
11266 assert!(note2.contains("(greedy)"), "single-flag note: {note2}");
11267 }
11268
11269 /// Shape-change clears (dmask realloc / mask-shape mismatch / s_key change) stay silent —
11270 /// they precede a fresh capture attempt whose own failure re-flips loudly.
11271 #[test]
11272 fn shape_change_clears_are_silent() {
11273 let mut f = DraftGraphFallback::default();
11274 f.mark_greedy("oom").unwrap();
11275 f.clear_greedy();
11276 assert!(!f.greedy_failed());
11277 f.mark_sampled("oom").unwrap();
11278 f.clear_sampled();
11279 assert!(!f.sampled_failed());
11280 // after a silent clear there is nothing left for resume to report.
11281 assert!(f.reset_on_resume().is_none());
11282 }
11283}
11284
11285/// SAMPLED DRAFT-GRAPH KEY (lane/graph-s-key-exactness-20260819).
11286///
11287/// These are the CPU teeth for an exactness bug whose live reproduction needs a GPU, a trunk, a
11288/// drafter and a two-turn session: the key itself. Every test below fails against the pre-fix key
11289/// `(seed, temp.to_bits(), k)` — `legacy_key` restates it so the collision is explicit rather
11290/// than remembered.
11291#[cfg(test)]
11292mod sampled_graph_key_tests {
11293 use super::{SampledGraphKey, debug_t_pred0};
11294
11295 /// The pre-fix key, verbatim: `let s_key = (sp_seed, sp_temp.to_bits(), k);`
11296 fn legacy_key(k: &SampledGraphKey) -> (u64, u32, usize) {
11297 (k.seed, k.temp_bits, k.k)
11298 }
11299
11300 fn pure_temp_key() -> SampledGraphKey {
11301 // temperature 1.0, filters off — today's serve default, the shape that parks a graph.
11302 SampledGraphKey::new(12345, 1.0, 3, 0, 1.0, 0.0, false)
11303 }
11304
11305 /// THE COLLISION. Two requests that differ ONLY in the truncation filters shared one key, so
11306 /// a parked pure-temp graph survived into a filtered request and the launch site launched it.
11307 #[test]
11308 fn vendor_filters_change_the_key() {
11309 let parked = pure_temp_key();
11310 // qwen3.8 generation_config.json — what the vendor-default flip makes the default shape.
11311 let vendor = SampledGraphKey::new(12345, 1.0, 3, 20, 0.95, 0.0, false);
11312 assert_eq!(
11313 legacy_key(&parked),
11314 legacy_key(&vendor),
11315 "pre-fix key collided: this is the bug, and the reason a test asserts on it",
11316 );
11317 assert_ne!(parked, vendor, "post-fix key must separate the two regimes");
11318 assert!(parked.pure_temp());
11319 assert!(!vendor.pure_temp());
11320 }
11321
11322 /// Each distribution-shaping field alone is enough to drop the parked graph.
11323 #[test]
11324 fn every_filter_field_is_keyed() {
11325 let base = pure_temp_key();
11326 for (what, other) in [
11327 (
11328 "top_k",
11329 SampledGraphKey::new(12345, 1.0, 3, 20, 1.0, 0.0, false),
11330 ),
11331 (
11332 "top_p",
11333 SampledGraphKey::new(12345, 1.0, 3, 0, 0.95, 0.0, false),
11334 ),
11335 (
11336 "min_p",
11337 SampledGraphKey::new(12345, 1.0, 3, 0, 1.0, 0.05, false),
11338 ),
11339 (
11340 "penalties",
11341 SampledGraphKey::new(12345, 1.0, 3, 0, 1.0, 0.0, true),
11342 ),
11343 ] {
11344 assert_ne!(base, other, "{what} must be part of the key");
11345 assert!(!other.pure_temp(), "{what} leaves the pure-temp regime");
11346 assert_eq!(
11347 legacy_key(&base),
11348 legacy_key(&other),
11349 "{what} was invisible to the pre-fix key",
11350 );
11351 }
11352 }
11353
11354 /// The baked constants stay keyed (this half was always right — regression cover for it).
11355 #[test]
11356 fn baked_constants_stay_keyed() {
11357 let base = pure_temp_key();
11358 assert_ne!(
11359 base,
11360 SampledGraphKey::new(999, 1.0, 3, 0, 1.0, 0.0, false),
11361 "seed"
11362 );
11363 assert_ne!(
11364 base,
11365 SampledGraphKey::new(12345, 0.7, 3, 0, 1.0, 0.0, false),
11366 "temp"
11367 );
11368 assert_ne!(
11369 base,
11370 SampledGraphKey::new(12345, 1.0, 4, 0, 1.0, 0.0, false),
11371 "k"
11372 );
11373 // bitwise on temperature: 0.7f32 vs the same value re-derived must NOT differ.
11374 assert_eq!(
11375 SampledGraphKey::new(1, 0.7, 3, 0, 1.0, 0.0, false),
11376 SampledGraphKey::new(1, 7.0 / 10.0, 3, 0, 1.0, 0.0, false),
11377 );
11378 }
11379
11380 /// THE LOAD-BEARING HALF OF THE SEED DECISION (lane/session-resume-sampler-predicate-
11381 /// 20260820). The whole-session resume predicate deliberately does NOT compare `seed`: an
11382 /// omitted serve `seed` draws fresh per-request entropy, so comparing it would refuse every
11383 /// seed-omitting sampled conversation. That is only sound because the one piece of parked state
11384 /// that BAKES the seed — this graph — is re-keyed on it, so a seed change drops and recaptures.
11385 ///
11386 /// This test is the other end of that argument, asserted here rather than remembered in a
11387 /// comment: if a future change dropped `seed` from the key, the resume predicate's exclusion
11388 /// would silently become the unsound thing it is documented not to be.
11389 /// (Paired with `seed_alone_does_not_refuse` in `memra-sampling`.)
11390 #[test]
11391 fn seed_alone_still_rekeys_the_draft_graph() {
11392 let parked = pure_temp_key();
11393 let reseeded = SampledGraphKey::new(999, 1.0, 3, 0, 1.0, 0.0, false);
11394 assert_ne!(
11395 parked, reseeded,
11396 "a seed-only change MUST drop the parked sampled graph — the resume predicate's \
11397 decision not to compare seed rests on exactly this",
11398 );
11399 // Same regime on both sides: the drop is a recapture, not a fall to the eager chain
11400 // because of a filter difference.
11401 assert!(parked.pure_temp() && reseeded.pure_temp());
11402 }
11403
11404 /// `pure_temp()` is the capture guard's predicate, computed from the key so the two cannot
11405 /// drift. The equality below is the invariant the launch-site guard asserts: identical keys
11406 /// agree on the regime, so a graph that survives the drop is legal to launch.
11407 #[test]
11408 fn equal_keys_agree_on_the_regime() {
11409 let a = SampledGraphKey::new(7, 0.8, 3, 20, 0.95, 0.0, false);
11410 let b = SampledGraphKey::new(7, 0.8, 3, 20, 0.95, 0.0, false);
11411 assert_eq!(a, b);
11412 assert_eq!(a.pure_temp(), b.pure_temp());
11413 // top_p slightly above 1.0 (a client sending 1.0 exactly, or an operator default) is
11414 // still the unfiltered regime, matching the original `sp.top_p >= 1.0` test.
11415 assert!(SampledGraphKey::new(7, 0.8, 3, 0, 1.0, 0.0, false).pure_temp());
11416 assert!(SampledGraphKey::new(7, 0.8, 3, 0, 1.5, -1.0, false).pure_temp());
11417 }
11418
11419 /// MEMRA_DEBUG_SPEC on a SAMPLED spec request past round 0: the print must render without
11420 /// indexing the empty greedy `preds` vector (it panicked the GPU worker before this lane).
11421 #[test]
11422 fn debug_print_survives_the_sampled_arm() {
11423 // round >= 1 with a pending bonus == base 1, sampled == `preds` empty.
11424 assert_eq!(debug_t_pred0(true, 1, 4242, &[]), "n/a");
11425 assert_eq!(debug_t_pred0(true, 2, 4242, &[]), "n/a");
11426 // round 0 without a pending bonus still reports last_pred, in both arms.
11427 assert_eq!(debug_t_pred0(true, 0, 4242, &[]), "4242");
11428 assert_eq!(debug_t_pred0(false, 0, 4242, &[7, 8]), "4242");
11429 // greedy keeps the real prediction it always printed.
11430 assert_eq!(debug_t_pred0(false, 1, 4242, &[7, 8]), "7");
11431 assert_eq!(debug_t_pred0(false, 2, 4242, &[7, 8]), "8");
11432 }
11433}