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::cache::{Cache, KvLayer};
12use crate::forward::argmax;
13use crate::hybrid::{FullAttnLayer, HybridModel, LinearAttnLayer, Mixer, MtpHead};
14use crate::Engine;
15use cudarc::driver::CudaSlice;
16use std::sync::atomic::{AtomicU64, Ordering};
17
18/// One compact, anchor-bounded DSpark supervision record. `tokens[0]` is the anchor at p and
19/// `hidden` is its predecessor carrier h[p-1], matching the live NextN/DSpark pairing. Target
20/// rows p..p+gamma-1 score tokens p+1..p+gamma. They are the full-target softmax's top-k
21/// entries; `target_tail_probs[j]` is the probability mass outside those rows. All flattened
22/// target arrays are `[gamma, top_k]` in row-major order.
23pub struct DsparkAnchorRecord {
24 pub position: usize,
25 pub hidden: Vec<f32>,
26 pub tokens: Vec<u32>,
27 pub target_top_ids: Vec<u32>,
28 pub target_top_logits: Vec<f32>,
29 pub target_top_probs: Vec<f32>,
30 pub target_tail_probs: Vec<f32>,
31}
32
33fn dspark_sparse_softmax_topk(
34 logits: &[f32],
35 top_k: usize,
36 temperature: f32,
37) -> Result<(Vec<u32>, Vec<f32>, Vec<f32>, f32), Box<dyn std::error::Error>> {
38 if logits.is_empty() || top_k == 0 || top_k > logits.len() || temperature <= 0.0 {
39 return Err("invalid DSpark sparse-softmax shape or temperature".into());
40 }
41 if logits.iter().any(|value| !value.is_finite()) {
42 return Err("DSpark target logits contain a non-finite value".into());
43 }
44 let mut ranked: Vec<(u32, f32)> = logits
45 .iter()
46 .copied()
47 .enumerate()
48 .map(|(index, value)| (index as u32, value))
49 .collect();
50 let compare = |left: &(u32, f32), right: &(u32, f32)| {
51 right.1.total_cmp(&left.1).then(left.0.cmp(&right.0))
52 };
53 ranked.select_nth_unstable_by(top_k - 1, compare);
54 ranked[..top_k].sort_unstable_by(compare);
55
56 let max_logit = logits.iter().copied().fold(f32::NEG_INFINITY, f32::max);
57 let inv_temperature = 1.0f64 / temperature as f64;
58 let denominator: f64 = logits
59 .iter()
60 .map(|value| (((*value - max_logit) as f64) * inv_temperature).exp())
61 .sum();
62 let ids: Vec<u32> = ranked[..top_k].iter().map(|(index, _)| *index).collect();
63 let top_logits: Vec<f32> = ranked[..top_k].iter().map(|(_, value)| *value).collect();
64 let top_probs: Vec<f32> = top_logits
65 .iter()
66 .map(|value| {
67 ((((value - max_logit) as f64) * inv_temperature).exp() / denominator) as f32
68 })
69 .collect();
70 let top_mass: f64 = top_probs.iter().map(|value| *value as f64).sum();
71 let tail = (1.0f64 - top_mass).clamp(0.0, 1.0) as f32;
72 Ok((ids, top_logits, top_probs, tail))
73}
74
75fn flatten_dspark_rows<T>(
76 rows: Vec<Option<Vec<T>>>,
77 position: usize,
78 label: &str,
79) -> Result<Vec<T>, Box<dyn std::error::Error>> {
80 let mut flattened = Vec::new();
81 for (slot, row) in rows.into_iter().enumerate() {
82 flattened.extend(
83 row.ok_or_else(|| format!("missing DSpark {label} at {position} slot {slot}"))?,
84 );
85 }
86 Ok(flattened)
87}
88
89/// H-SEED CONVENTION (MEMRA_SPEC_HPOST=1): feed the MTP head the POST-norm hidden — trunk rows
90/// hand over `output_norm(x)` and the draft chain recurrence hands over `shared_head_norm(h_nextn)`
91/// (= final_h) — matching the reference engines: llama.cpp #24025 ("qwen35: use post-norm hidden
92/// state for MTP", t_h_nextn is taken AFTER the final norm in both trunk and MTP graphs) and
93/// SGLang's qwen3_5_mtp (spec_info.hidden_states = the target model's post-norm output). memra's
94/// historical convention (default, MTP-PLAN §A) is PRE-norm x. Draft-quality-only: exactness is
95/// the verify's job either way; acceptance arbitrates. OnceLock: read once, hot-loop safe.
96pub(crate) fn spec_hpost() -> bool {
97 static H: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
98 *H.get_or_init(|| {
99 std::env::var("MEMRA_SPEC_HPOST")
100 .map(|v| v != "0")
101 .unwrap_or(false)
102 })
103}
104
105/// LEAN VERIFY (default ON since 2026-07-08; MEMRA_SPEC_LEAN=0 reverts — close35 lane): the verify m-scaling
106/// probe + nsys diff showed the verify t-path pays ~1.0ms/call at m=1 over eager decode on the
107/// 35B, and the kernels are NOT the cause (dev-MoE identical, kernel-time delta only +179us).
108/// The overhead is (a) ~250 extra cuMemsetD8Async/call from `e.zeros()` on buffers every kernel
109/// fully overwrites (~0.9ms host issue + ~0.35ms GPU) and (b) the t=1 FA rows dispatch (rows_v2 +
110/// combine_rows, +50us vs the eager fa_decode pair). This flag switches (a) fully-overwritten
111/// verify buffers to `e.uninit` (identical bytes: every element is written before read) and
112/// (b) t==1 verify FA to the eager `fa_decode` entry (byte-identical: kernel-check pins the
113/// rows-vs-loop identity and the per-row loop at t=1 IS fa_decode on the same q). Gates arbitrate.
114pub(crate) fn spec_lean() -> bool {
115 static L: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
116 // DEFAULT ON since 2026-07-08 (MEMRA_SPEC_LEAN=0 reverts): bit-identical (buffers fully
117 // overwritten; gates green incl maxdiff-identical run-gen) and measured +2.4% e2e p3 /
118 // +1.5% p2 at the daily 35B config. m=1 verify now costs eager-decode parity.
119 *L.get_or_init(|| {
120 std::env::var("MEMRA_SPEC_LEAN")
121 .map(|v| v != "0")
122 .unwrap_or(true)
123 })
124}
125
126/// SMALL-M BATCHED VERIFY (default ON since 2026-07-09; MEMRA_SPEC_M2=0 reverts — lane/spec-m2): extend the
127/// batched linear-attn verify arm down to t=2 and batch the MoE dev token loop over a
128/// grid.z=token axis at every verify t. The close35 m-scaling probe put the m=2 verify tier at
129/// x1.54 of m=1 (llama x1.14); the per-column linear chain (t<3) and the serial MoE dev token
130/// loop are the two launch-structure causes. Both changes are LAUNCH-STRUCTURE ONLY:
131/// (a) the batched conv's t<pad ring update is pure copies (ssm_conv_ring_rebuild from a cloned
132/// ring — the ring stores raw input columns); every arithmetic kernel is the same one the
133/// t>=3 arm already runs (matmul_decode_exact bit-identical at m=2-4, gdn_scan's internal
134/// t-loop == chained T=1 steps);
135/// (b) the MoE dev-rows twins run the serial loop's per-token warp program with tok-offset
136/// pointers (same sel/w/aq/ad bytes, same dot order, same slot-ordered FMA chain).
137/// Gates arbitrate: run-spec K=1..8 self-consistency (35B+9B), kernel-check, run-gen argmax.
138pub(crate) fn spec_m2() -> bool {
139 static M: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
140 // DEFAULT ON since 2026-07-09 (MEMRA_SPEC_M2=0 reverts): launch-structure only — t=2
141 // batched linear arm (ring-roll copies, zero new FP order) + MoE dev-rows kernels
142 // (grid.z=token, 4 launches/layer at any verify t). Acceptance bit-identical at every K;
143 // 35B p2 +3.4% / p3 +3.6%; the profitable-K plateau widens (new optimum K=3 at 223).
144 *M.get_or_init(|| {
145 std::env::var("MEMRA_SPEC_M2")
146 .map(|v| v != "0")
147 .unwrap_or(true)
148 })
149}
150pub(crate) fn spec_stream() -> bool {
151 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
152 *ON.get_or_init(|| std::env::var("MEMRA_SPEC_STREAM").as_deref() == Ok("1"))
153}
154pub(crate) fn spec_stream_m() -> usize {
155 static M: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
156 *M.get_or_init(|| {
157 std::env::var("MEMRA_SPEC_STREAM_M")
158 .ok()
159 .and_then(|v| v.parse().ok())
160 .unwrap_or(4)
161 })
162}
163pub(crate) fn spec_devacc() -> bool {
164 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
165 *ON.get_or_init(|| std::env::var("MEMRA_SPEC_DEVACC").as_deref() == Ok("1"))
166}
167
168/// GRAMMAR HOOK for constrained spec decode (lane/constrained-full, 2026-08-03). The engine
169/// stays llguidance-agnostic: the server adapts its per-session grammar state behind this
170/// trait. CONTRACT (the verify-side truncation rule — token-identical to constrained plain
171/// greedy decode): the exactness walk runs UNMASKED first; the hook then (a) truncates
172/// acceptance at the first grammar-illegal accepted token, and (b) when the truncation fired
173/// or the bonus is illegal, the engine recomputes that slot as the MASKED argmax of the
174/// target's own verify column (an unmasked argmax that is grammar-legal IS the masked argmax
175/// — masking only removes tokens — so the common case pays nothing). `consume` advances the
176/// state with each EMITTED token in order; EOS handling is the implementor's job (skip).
177pub trait SpecConstraint {
178 /// -inf the current state's banned ids on a HOST logits row (prompt-tail / init-feed
179 /// masked argmax).
180 fn mask_logits(&mut self, logits: &mut [f32]) -> Result<(), String>;
181 /// Packed 32-bit bitset words of the CURRENT state's allowed set (device-mask form).
182 fn mask_words(&mut self) -> Result<Vec<u32>, String>;
183 /// Is `tok` consumable in the CURRENT state?
184 fn is_allowed(&mut self, tok: u32) -> Result<bool, String>;
185 /// Advance the state with an emitted token.
186 fn consume(&mut self, tok: u32) -> Result<(), String>;
187
188 // --- DRAFT-SIDE MASKING (lane/draft-mask, 2026-08-04) ---
189 // The drafter proposed grammar-illegal tokens under tight schemas, so verify-side
190 // truncation cut nearly every round (measured acceptance 0.467-0.513 tight vs 0.62-0.82
191 // loose, research/constrained-full-20260803). These three methods let the engine mask the
192 // DRAFT model's own sampling with the grammar's legal set, so proposals are legal by
193 // construction. The state they walk is a SPECULATIVE CLONE of the session matcher — the
194 // real state is advanced only by `consume` (emitted tokens), so verify-side truncation
195 // stays the correctness backstop and the emitted stream is unchanged by construction
196 // (an accepted draft is the target's unmasked argmax AND grammar-legal, hence the masked
197 // argmax; a cut slot is recomputed as the masked argmax either way).
198 // Default impls = feature OFF (pre-lane behaviour: unmasked drafts).
199
200 /// Is draft-side masking available on this hook? Probed ONCE per burst, before the draft
201 /// graph is captured (the mask is an in-graph node — its presence is a capture-time shape).
202 fn draft_mask_enabled(&self) -> bool {
203 false
204 }
205 /// Start a draft chain: clone the CURRENT (committed) grammar state into the speculative
206 /// slot. Called once per spec round, before the first draft position.
207 fn draft_begin(&mut self) -> Result<(), String> {
208 Ok(())
209 }
210 /// Packed 32-bit bitset words of the SPECULATIVE state's allowed set (target-vocab ids),
211 /// for the draft position about to be sampled. `None` = draft masking off (no-op).
212 fn draft_mask_words(&mut self) -> Result<Option<Vec<u32>>, String> {
213 Ok(None)
214 }
215 /// Advance the SPECULATIVE state with a PROPOSED draft token. `false` = the chain cannot
216 /// continue (EOS proposed, or an unmasked position proposed something illegal) — the
217 /// engine stops drafting; the token already pushed still goes through verify.
218 fn draft_advance(&mut self, _tok: u32) -> Result<bool, String> {
219 Ok(false)
220 }
221}
222
223/// DRAFT-MASK UPLOAD (lane/draft-mask): pull the speculative state's allowed set (TARGET-id
224/// space) from the hook, project it into the DRAFT head's vocab space, and upload it into the
225/// stable device buffer the draft chain reads. Returns false when the chain must stop drafting:
226/// the hook handed out no mask, or NO draft-vocab row is grammar-legal at this position (a
227/// trimmed FR-Spec head genuinely cannot propose a legal token there — masking it would leave
228/// a fully-banned row whose argmax is meaningless, so the round drafts fewer tokens and the
229/// verify emits the masked argmax as usual).
230fn upload_draft_mask(
231 e: &Engine,
232 c: &mut dyn SpecConstraint,
233 dst: &mut CudaSlice<u32>,
234 d2t: Option<&Vec<u32>>,
235 d_vocab: usize,
236 words: usize,
237) -> Result<bool, Box<dyn std::error::Error>> {
238 let Some(tw) = c.draft_mask_words().map_err(|e2| format!("constraint: {e2}"))? else {
239 return Ok(false);
240 };
241 let bit = |t: usize| -> bool {
242 let w = t >> 5;
243 w < tw.len() && (tw[w] >> (t & 31)) & 1 == 1
244 };
245 let mut buf = vec![0u32; words];
246 match d2t {
247 // TRIMMED draft head: row i proposes target id d2t[i] — permute the mask accordingly.
248 Some(map) => {
249 for (i, &t) in map.iter().enumerate().take(d_vocab) {
250 if bit(t as usize) {
251 buf[i >> 5] |= 1u32 << (i & 31);
252 }
253 }
254 }
255 // UNTRIMMED: draft ids ARE target ids; the packed words transfer verbatim (a short
256 // mask leaves the padded tail zeroed == banned, same rule as constrained::apply_mask).
257 None => {
258 let n = tw.len().min(words);
259 buf[..n].copy_from_slice(&tw[..n]);
260 }
261 }
262 if buf.iter().all(|w| *w == 0) {
263 return Ok(false);
264 }
265 e.htod_u32_into(dst, &buf)?;
266 Ok(true)
267}
268
269/// Keep the full token-embedding table in host memory and upload only the rows needed by each
270/// MTP/verify step. This is an exact memory-capacity seam for very large BF16 vocab tables: host
271/// gather expands the same source bits to f32, and only O(T*n_embd) bytes cross PCIe per step.
272/// CUDA-graph/round-stream draft paths require device token ids and therefore stay disabled.
273pub(crate) fn spec_host_embd() -> bool {
274 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
275 *ON.get_or_init(|| std::env::var("MEMRA_SPEC_HOST_EMBD").as_deref() == Ok("1"))
276}
277
278/// VERIFY-TIER TRUNK LAUNCH-FUSION (default ON since 2026-07-09; MEMRA_SPEC_FUSED_T=0 reverts — lane/close35b): extend
279/// the t=1 fused2/fused3 Q8_0 trunk launches to the batched verify tier (t=2-4, the K=1..3
280/// verify shapes). At t>1 the trunk pairs/triples (35B wqkv+wqkv_gate, wq/wk/wv,
281/// gate_shexp+up_shexp) each run a separate `matmul_decode_exact` — one q8_1 re-quantize of the
282/// SAME activation plus one _b2/_b4 launch per tensor. The fused twins share ONE quantize and
283/// ONE launch per group; per (tensor,token,row) the kernel body is q8_0_mmvq_batched verbatim
284/// with the identical row mapping -> BIT-IDENTICAL by construction (kernel-check pins it,
285/// run-spec K=1..8 + acceptance identity arbitrate e2e).
286pub(crate) fn spec_fused_t() -> bool {
287 static F: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
288 // DEFAULT ON since 2026-07-09 (MEMRA_SPEC_FUSED_T=0 reverts): verify t=2-4 trunk launch-fusion
289 // (fused2/fused3 Q8_0 batched twins, bit-identical by construction — m=1 block-offset split on
290 // the batched body). m=2 marginal token 2117->1762us; 35B daily: p3 +3.7% (crosses llama), p2 +5%.
291 *F.get_or_init(|| {
292 std::env::var("MEMRA_SPEC_FUSED_T")
293 .map(|v| v != "0")
294 .unwrap_or(true)
295 })
296}
297
298/// zeros/uninit switch for verify-path buffers that are FULLY OVERWRITTEN before any read.
299/// Only call this on such buffers — the lean contract is "identical bytes by construction".
300fn vbuf(e: &Engine, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
301 if spec_lean() {
302 e.uninit(n)
303 } else {
304 e.zeros(n)
305 }
306}
307
308/// Scratch KV for the MTP block (one full-attn layer).
309///
310/// PERSISTENT MODE (default, 2026-07-03 — the acceptance lever): sized cap = max_ctx and kept in
311/// sync with the COMMITTED sequence — slot p holds the MTP block's K/V for committed token p
312/// (roped p+1, the chain's rope convention), so the draft chain's self-attention sees the FULL
313/// committed history instead of only the current round's 1..K+1 chain tokens (the reference
314/// engine's "mtp_update" design). Entries come from two sources:
315/// - chain appends: accepted positions KEEP their chain-computed entries (embedding exact,
316/// hidden chain-approximate — the reference engine accepts the same);
317/// - `mtp_kv_fill` batches: prompt positions + the last-draft position on full accept, computed
318/// from EXACT trunk hiddens (K/V-only MTP-block pass, no attention/FFN/lm_head).
319/// Rejected drafts / p-min extras / pseudo-seed appends are all discarded by the round-start
320/// `set_len` truncation (the KvLayer len mechanism — §C rollback for the draft side).
321/// Multi-turn spec-decode session (2026-07-05): trunk Cache + persistent MTP draft scratch +
322/// the committed token list, alive across generate_spec_session calls. Turn N+1 primes ONLY its
323/// suffix (chunked continuation prime over the quantized past) and mtp_kv_fill's its suffix rows,
324/// then the round loop runs unchanged. `last_h` carries the pre-output_norm hidden of the last
325/// committed row across turns (the predecessor-pairing seed + fill anchor).
326/// Per-request sampling config for the sampled-spec serve path.
327#[derive(Clone, Copy, Debug)]
328pub struct SpecSampling {
329 pub temp: f32,
330 pub seed: u64,
331 pub top_k: i32, // 0 = off
332 pub top_p: f32, // 1.0 = off
333 pub min_p: f32, // 0.0 = off
334 pub penalty_last_n: usize, // 0 = penalties off
335 pub penalty_repeat: f32,
336 pub penalty_freq: f32,
337 pub penalty_present: f32,
338}
339
340/// Tracked draft positions for [`SpecTelemetry`] (serve K defaults to 3; the run-spec gate
341/// sweeps K=1..8, and MEMRA_SPEC_CAPMAX defaults to 7 — 8 covers every tuned config).
342pub const SPEC_TELEM_POS: usize = 8;
343
344/// Always-on per-draft-position acceptance telemetry (lane/accept-telemetry, 2026-08-05 —
345/// the llama.cpp #26389 / vLLM spec-decode counter schema, upstream-sweeps 2026-08-05).
346/// Lives on the [`SpecSession`] and accumulates across bursts; the serve worker diffs a
347/// stashed copy per burst for its per-model /metrics aggregation and per-request usage.
348/// Same normalization as the `[spec-stats]` line: p-min-discarded chain tokens are counted
349/// in NEITHER drafted nor accepted.
350#[derive(Clone, Copy, Default, Debug)]
351pub struct SpecTelemetry {
352 /// verify rounds completed (a round-stream burst counts each of its M rounds).
353 pub rounds: u64,
354 /// tokens drafted / accepted across all rounds.
355 pub drafted: u64,
356 pub accepted: u64,
357 /// how often draft position j (0-based within a round's chain) was offered / accepted.
358 /// Positions >= SPEC_TELEM_POS are untracked (totals still count them). The opt-in
359 /// round-stream arm (MEMRA_SPEC_STREAM=1) reads back only totals, so under it these
360 /// arrays cover the standard-path rounds only and their sums may undercount the totals.
361 pub pos_drafted: [u64; SPEC_TELEM_POS],
362 pub pos_accepted: [u64; SPEC_TELEM_POS],
363}
364
365impl SpecTelemetry {
366 /// Fieldwise `self - prev` — the worker's per-burst delta off a copy stashed before the
367 /// burst call. Saturating: a caller diffing against the wrong snapshot gets zeros, not
368 /// a wrapped counter.
369 pub fn delta_since(&self, prev: &SpecTelemetry) -> SpecTelemetry {
370 let mut d = SpecTelemetry {
371 rounds: self.rounds.saturating_sub(prev.rounds),
372 drafted: self.drafted.saturating_sub(prev.drafted),
373 accepted: self.accepted.saturating_sub(prev.accepted),
374 ..Default::default()
375 };
376 for j in 0..SPEC_TELEM_POS {
377 d.pos_drafted[j] = self.pos_drafted[j].saturating_sub(prev.pos_drafted[j]);
378 d.pos_accepted[j] = self.pos_accepted[j].saturating_sub(prev.pos_accepted[j]);
379 }
380 d
381 }
382 /// Fieldwise `self += d` — the worker's per-model aggregation.
383 pub fn merge(&mut self, d: &SpecTelemetry) {
384 self.rounds += d.rounds;
385 self.drafted += d.drafted;
386 self.accepted += d.accepted;
387 for j in 0..SPEC_TELEM_POS {
388 self.pos_drafted[j] += d.pos_drafted[j];
389 self.pos_accepted[j] += d.pos_accepted[j];
390 }
391 }
392
393 /// Mean accepted draft-prefix length per verify round (tau).
394 pub fn tau(&self) -> f64 {
395 if self.rounds > 0 {
396 self.accepted as f64 / self.rounds as f64
397 } else {
398 0.0
399 }
400 }
401}
402
403/// Session-lifetime atomic acceptance counters. The verifier records only after the greedy or
404/// rejection-sampling walk has resolved on the host, so these relaxed increments add no GPU
405/// launch, synchronization, allocation, or ordering dependency to the numeric path.
406struct SpecTelemetryCounters {
407 rounds: AtomicU64,
408 drafted: AtomicU64,
409 accepted: AtomicU64,
410 pos_drafted: [AtomicU64; SPEC_TELEM_POS],
411 pos_accepted: [AtomicU64; SPEC_TELEM_POS],
412}
413
414impl Default for SpecTelemetryCounters {
415 fn default() -> Self {
416 Self {
417 rounds: AtomicU64::new(0),
418 drafted: AtomicU64::new(0),
419 accepted: AtomicU64::new(0),
420 pos_drafted: std::array::from_fn(|_| AtomicU64::new(0)),
421 pos_accepted: std::array::from_fn(|_| AtomicU64::new(0)),
422 }
423 }
424}
425
426impl SpecTelemetryCounters {
427 fn record_round(&self, drafted: usize, accepted: usize) {
428 debug_assert!(accepted <= drafted);
429 self.rounds.fetch_add(1, Ordering::Relaxed);
430 self.drafted.fetch_add(drafted as u64, Ordering::Relaxed);
431 self.accepted.fetch_add(accepted as u64, Ordering::Relaxed);
432 for counter in self.pos_drafted.iter().take(drafted) {
433 counter.fetch_add(1, Ordering::Relaxed);
434 }
435 for counter in self.pos_accepted.iter().take(accepted) {
436 counter.fetch_add(1, Ordering::Relaxed);
437 }
438 }
439
440 /// Round-stream keeps each round's accept length on device; retain exact scalar totals while
441 /// leaving the per-position arrays untouched, matching the pre-existing telemetry contract.
442 fn record_totals(&self, rounds: usize, drafted: usize, accepted: usize) {
443 self.rounds.fetch_add(rounds as u64, Ordering::Relaxed);
444 self.drafted.fetch_add(drafted as u64, Ordering::Relaxed);
445 self.accepted.fetch_add(accepted as u64, Ordering::Relaxed);
446 }
447
448 fn snapshot(&self) -> SpecTelemetry {
449 SpecTelemetry {
450 rounds: self.rounds.load(Ordering::Relaxed),
451 drafted: self.drafted.load(Ordering::Relaxed),
452 accepted: self.accepted.load(Ordering::Relaxed),
453 pos_drafted: std::array::from_fn(|j| {
454 self.pos_drafted[j].load(Ordering::Relaxed)
455 }),
456 pos_accepted: std::array::from_fn(|j| {
457 self.pos_accepted[j].load(Ordering::Relaxed)
458 }),
459 }
460 }
461}
462
463pub struct SpecSession {
464 pub(crate) cache: Cache,
465 pub(crate) scratch: MtpScratch,
466 /// Every token whose state the caches hold, in order (prompt turns + generated), INCLUDING
467 /// overshoot: spec commits accepted drafts past max_new; those rows are in the caches, so the
468 /// session must count them. Callers render output from this, not from their own echo.
469 pub committed: Vec<u32>,
470 /// Pre-output_norm hidden of the LAST committed row (device). None before the first turn.
471 pub(crate) last_h: Option<CudaSlice<f32>>,
472 /// Greedy argmax predicting the token AFTER committed.last() (from the last turn's final
473 /// logits). Fuels empty-suffix continuation bursts (serve): the next turn emits this token
474 /// first, feeds it, and the round loop resumes without any prime. None before the first turn.
475 pub next_pred: Option<u32>,
476 /// SAMPLED-SPEC stream continuity across bursts: Philox event counters persist here so a
477 /// session's randomness never repeats between generate_spec_session calls. (0,0) at admit.
478 pub sctr: u32,
479 pub uctr: u32,
480 /// PERSISTENT DRAFT-GRAPH CONTEXT (2026-08-01, the serve-burst fixed-cost fix): the captured
481 /// draft graph(s) + every device I/O buffer they bake, carried ACROSS generate_spec_session
482 /// calls. Before this, every serve burst re-captured the draft graph (2 warmup forwards +
483 /// instantiate) — measured ~16ms/burst on H100 q27 (MEMRA_SPEC_BURST sweep,
484 /// research/spec-serving-20260801). None before the first turn; error paths drop it
485 /// (next burst recaptures — serve retires errored sessions anyway).
486 pub(crate) draft_ctx: Option<DraftGraphCtx>,
487 /// PENDING-CARRY across bursts (2026-08-01, the serve burst-boundary fix): the bonus token
488 /// emitted by the last round but NOT committed to the caches. The old tail committed it with
489 /// a solo T=1 trunk pass (+ draft fill), and the next burst's setup fed the stashed next_pred
490 /// with ANOTHER solo pass — 2x ~11.5ms/burst measured on H100 q27 ([spec-setup] trace).
491 /// Carrying it lets the next empty-suffix greedy burst consume it as round-0 verify col 0,
492 /// exactly like a mid-burst full-accept boundary (no solo passes). INVARIANT: when set,
493 /// `committed` (== cache rows) EXCLUDES this token although it was already emitted in the
494 /// last burst's output, and `last_h` holds the hidden of the last COMMITTED row (its
495 /// predecessor — the chain-seed/fill anchor). `next_pred` is None (unknown without the
496 /// commit pass). Non-empty-suffix or sampled turns must flush first (spec_flush_pending);
497 /// generate_spec_session_sampled does this at entry, and serve parks only flushed sessions.
498 pub pending_tok: Option<u32>,
499 /// SESSION-AFFINITY TURN CHECKPOINT (lane/session-affinity, 2026-08-05): the state at this
500 /// turn's PROMPT-END boundary, retained so a later turn can REWIND here. See
501 /// [`SpecCheckpoint`]. Refreshed by every non-empty prime; None until the first one, and on
502 /// a rig too tight to hold it (a failed capture is silent — resume just isn't available).
503 pub(crate) turn_ckpt: Option<SpecCheckpoint>,
504 /// Session-lifetime acceptance telemetry. Relaxed atomics update at the host-side round
505 /// accounting the loop already does — no syncs, no allocation. NOTE a
506 /// pool-resumed session carries the PREVIOUS requests' counts; per-request consumers
507 /// diff with [`SpecTelemetry::delta_since`] around each burst.
508 telem: SpecTelemetryCounters,
509}
510impl SpecSession {
511 /// Context capacity of the session's caches (the server's ContextFull guard).
512 pub fn cache_max_ctx(&self) -> usize {
513 self.cache.max_ctx
514 }
515 /// Snapshot the session's process-local acceptance counters for per-burst diffing.
516 pub fn telemetry(&self) -> SpecTelemetry {
517 self.telem.snapshot()
518 }
519 /// Committed position this session can REWIND to (its retained prompt-end boundary), if any.
520 /// A request whose prompt matches `committed[..pos]` exactly can resume from here — see
521 /// `spec_rewind_to_checkpoint`.
522 pub fn rewind_pos(&self) -> Option<usize> {
523 self.turn_ckpt.as_ref().map(|c| c.pos)
524 }
525 /// Whether every ring-backed trunk/draft row needed by the retained checkpoint is resident.
526 pub fn rewind_is_resident(&self) -> bool {
527 self.turn_ckpt.as_ref().is_some_and(|ckpt| {
528 self.cache.can_rollback(&ckpt.snap, 0) && self.scratch.can_rewind_to(ckpt.pos)
529 })
530 }
531 /// Is this session in the DEMOTION-READY shape (see [`SpecSession::into_demoted`])?
532 /// `false` means a carried pending must be flushed first (`spec_flush_pending`), or the
533 /// session has never run a turn and has no prediction to hand over.
534 pub fn demote_ready(&self) -> bool {
535 self.pending_tok.is_none() && self.next_pred.is_some()
536 }
537 /// Does this session hold a carried pending bonus (flush required before a handoff/park)?
538 pub fn has_pending(&self) -> bool {
539 self.pending_tok.is_some()
540 }
541 /// Committed row count == cache rows (the session invariant), for the caller's own
542 /// `fed`-length cross-check at a handoff boundary.
543 pub fn committed_len(&self) -> usize {
544 self.committed.len()
545 }
546 /// DEMOTION HANDOFF (lane/spec-gate, 2026-08-07): consume this session and hand its trunk
547 /// cache + next-token prediction to the plain batched-decode path.
548 ///
549 /// WHY THIS IS EXACT (greedy). The invariant at a burst boundary is `cache.pos ==
550 /// committed.len()`: every committed row has trunk KV + recurrent state, exactly as a plain
551 /// tokenwise prime of the same `committed` sequence would have left it (that is the
552 /// session-tail contract, and the same property `spec_rewind_to_checkpoint` and the reuse
553 /// pool already rely on). `next_pred` is the argmax of the verify's logits for the LAST
554 /// committed row — and verify-column logits are bit-identical to plain decode's logits at
555 /// that position, because `matmul_decode_exact` bit-identity IS the basis of the greedy
556 /// accept walk. So handing (cache, next_pred) to the batched path continues the stream from
557 /// a state indistinguishable from one the batched path produced itself: the batched tick
558 /// emits `next_pred`, feeds it into this same cache, and decodes on.
559 ///
560 /// `None` when the session is not in the handoff shape — a carried pending (its bonus row is
561 /// NOT in the cache, so `spec_flush_pending` must commit it first) or no `next_pred` yet
562 /// (never bursted). Callers must not force it: a half-committed cache handed to the batched
563 /// path would silently skip a token.
564 ///
565 /// The MTP draft scratch, the persistent draft-graph context and the turn checkpoint are
566 /// DROPPED here (freeing their VRAM): the batched path never drafts, and this handoff is
567 /// one-way by design — there is no cheap symmetric re-promotion (rebuilding the draft KV
568 /// would mean an `mtp_kv_fill` over the whole committed history).
569 pub fn into_demoted(self) -> Option<(Cache, u32)> {
570 if self.pending_tok.is_some() {
571 return None;
572 }
573 let np = self.next_pred?;
574 debug_assert_eq!(
575 self.cache.pos,
576 self.committed.len(),
577 "demotion handoff: cache rows != committed tokens"
578 );
579 Some((self.cache, np))
580 }
581 /// Pool-resume hook (audit Q2): clear the parked draft-graph failure memoization so a
582 /// NEW request resuming this session gets one fresh capture chance — a transient-pressure
583 /// capture failure must not persist for the pool's whole lifetime (the TRT #16072 class).
584 /// Logs once iff a flag was actually set; a no-fallback resume is silent and free.
585 pub fn reset_graph_fallback_on_resume(&mut self) {
586 if let Some(line) = self
587 .draft_ctx
588 .as_mut()
589 .and_then(|c| c.failed.reset_on_resume())
590 {
591 eprintln!("{line}");
592 }
593 }
594}
595
596/// A session's PROMPT-END boundary state, the rewind target for session-affinity resume.
597///
598/// WHY THIS BOUNDARY, AND WHY IT IS THE ONLY ONE WORTH KEEPING. The rewrite class this lane
599/// exists for (a client that strips `<think>` blocks out of prior assistant turns) mutates the
600/// text the session GENERATED, never the prompt it was given. So turn N's prompt agrees with
601/// turn N-1's committed tokens up to almost exactly where turn N-1's generation began — the
602/// prompt-end boundary. Keeping a checkpoint there means the next turn re-primes only its own
603/// delta (the rewritten answer + the new user turn) instead of the whole conversation.
604///
605/// WHAT IT MUST HOLD. Full-attn KV is append-only and position-addressed, so rewinding it is a
606/// `len` truncation (no data). Linear-attn (GDN) conv/ssm state is mutated IN PLACE with no
607/// position index, so it must be a real device COPY — that copy is the entire reason a spec
608/// session could not previously rewind. The MTP draft scratch needs no copy either: its rows
609/// below the boundary were written by this turn's fill and are never revisited (the per-round
610/// true-hidden refresh only rewrites the CURRENT burst's committed positions), so rewinding it
611/// is also just a `len` reset. `last_h` is the hidden of the last row below the boundary — the
612/// predecessor-pairing anchor the next prime's fill reads for its first row.
613///
614/// COST: one `Cache::snapshot` per TURN, on a code path that already takes one per ROUND.
615pub(crate) struct SpecCheckpoint {
616 snap: crate::cache::CacheSnapshot,
617 /// Committed length at the boundary (== cache.pos there, the session invariant).
618 pos: usize,
619 /// Pre-output_norm hidden of row `pos - 1`.
620 last_h: CudaSlice<f32>,
621}
622
623struct SpecPipeTraceClock {
624 pair: usize,
625 started: std::time::Instant,
626}
627
628#[derive(Clone)]
629struct SpecPipeTraceCtx {
630 clock: std::sync::Arc<SpecPipeTraceClock>,
631 round: usize,
632 lane: usize,
633}
634
635struct SpecPipeTraceMarker {
636 trace: SpecPipeTraceCtx,
637 phase: &'static str,
638 edge: &'static str,
639 slot: Option<usize>,
640}
641
642unsafe extern "C" fn spec_pipe_trace_marker(raw: *mut std::ffi::c_void) {
643 let marker = unsafe { Box::from_raw(raw.cast::<SpecPipeTraceMarker>()) };
644 let lane = if marker.trace.lane == 0 { "A" } else { "B" };
645 let slot = marker
646 .slot
647 .map(|v| v.to_string())
648 .unwrap_or_else(|| "-".into());
649 let t_ms = marker.trace.clock.started.elapsed().as_secs_f64() * 1e3;
650 use std::io::Write as _;
651 let stderr = std::io::stderr();
652 let mut stderr = stderr.lock();
653 let _ = writeln!(
654 stderr,
655 "[spec-pipe-timeline] pair={} round={} lane={lane} phase={} edge={} \
656 slot={slot} t_ms={t_ms:.3}",
657 marker.trace.clock.pair,
658 marker.trace.round,
659 marker.phase,
660 marker.edge,
661 );
662}
663
664fn enqueue_spec_pipe_trace_marker(
665 stream: &cudarc::driver::CudaStream,
666 trace: Option<&SpecPipeTraceCtx>,
667 phase: &'static str,
668 edge: &'static str,
669 slot: Option<usize>,
670) -> Result<(), Box<dyn std::error::Error>> {
671 let Some(trace) = trace else {
672 return Ok(());
673 };
674 let marker = Box::new(SpecPipeTraceMarker {
675 trace: trace.clone(),
676 phase,
677 edge,
678 slot,
679 });
680 let raw = Box::into_raw(marker);
681 let result = unsafe {
682 cudarc::driver::result::stream::launch_host_function(
683 stream.cu_stream(),
684 spec_pipe_trace_marker,
685 raw.cast(),
686 )
687 };
688 if let Err(err) = result {
689 unsafe {
690 drop(Box::from_raw(raw));
691 }
692 return Err(err.into());
693 }
694 Ok(())
695}
696
697#[derive(Default)]
698struct SpecPipeProgress {
699 setup_done: [bool; 2],
700 draft_done: [usize; 2],
701 stage0_done: [usize; 2],
702 verify_done: [usize; 2],
703 accept_done: [usize; 2],
704 finished: [bool; 2],
705 aborted: bool,
706}
707
708/// Host-side issue coordinator for the reduced two-session speculative pipeline. Each session
709/// keeps its existing call stack and round locals; this object only orders phase entry. The
710/// primary mutex spans whole draft/accept/tail issue regions so Engine's single-stream scratch
711/// cannot be interleaved by the two host threads.
712struct SpecPipeSync {
713 progress: std::sync::Mutex<SpecPipeProgress>,
714 changed: std::sync::Condvar,
715 primary: std::sync::Mutex<()>,
716 trace: Option<std::sync::Arc<SpecPipeTraceClock>>,
717}
718
719impl SpecPipeSync {
720 fn new() -> Self {
721 static TRACE_PAIR: std::sync::atomic::AtomicUsize =
722 std::sync::atomic::AtomicUsize::new(0);
723 let trace = (std::env::var("MEMRA_SPEC_PIPE_TRACE").as_deref() == Ok("1")).then(|| {
724 std::sync::Arc::new(SpecPipeTraceClock {
725 pair: TRACE_PAIR.fetch_add(1, std::sync::atomic::Ordering::Relaxed) + 1,
726 started: std::time::Instant::now(),
727 })
728 });
729 Self {
730 progress: std::sync::Mutex::new(SpecPipeProgress::default()),
731 changed: std::sync::Condvar::new(),
732 primary: std::sync::Mutex::new(()),
733 trace,
734 }
735 }
736}
737
738#[derive(Clone)]
739struct SpecPipeLane {
740 sync: std::sync::Arc<SpecPipeSync>,
741 lane: usize,
742}
743
744impl SpecPipeLane {
745 fn peer(&self) -> usize {
746 1 - self.lane
747 }
748
749 fn aborted() -> Box<dyn std::error::Error> {
750 "paired speculative peer aborted".into()
751 }
752
753 fn trace(&self, round: usize) -> Option<SpecPipeTraceCtx> {
754 self.sync.trace.as_ref().map(|clock| SpecPipeTraceCtx {
755 clock: clock.clone(),
756 round,
757 lane: self.lane,
758 })
759 }
760
761 fn setup_begin(&self) -> Result<(), Box<dyn std::error::Error>> {
762 let mut p = self.sync.progress.lock().unwrap();
763 while !p.aborted
764 && self.lane == 1
765 && !p.setup_done[0]
766 && !p.finished[0]
767 {
768 p = self.sync.changed.wait(p).unwrap();
769 }
770 if p.aborted { Err(Self::aborted()) } else { Ok(()) }
771 }
772
773 fn setup_end(&self) {
774 let mut p = self.sync.progress.lock().unwrap();
775 p.setup_done[self.lane] = true;
776 self.sync.changed.notify_all();
777 }
778
779 fn draft_begin(
780 &self,
781 round: usize,
782 ) -> Result<std::sync::MutexGuard<'_, ()>, Box<dyn std::error::Error>> {
783 let peer = self.peer();
784 let mut p = self.sync.progress.lock().unwrap();
785 loop {
786 if p.aborted {
787 return Err(Self::aborted());
788 }
789 let setup_ready = (p.setup_done[0] || p.finished[0])
790 && (p.setup_done[1] || p.finished[1]);
791 let prior_ready = p.accept_done[self.lane] >= round
792 && (p.accept_done[peer] >= round || p.finished[peer]);
793 let turn_ready = if self.lane == 0 {
794 true
795 } else {
796 p.draft_done[0] > round || p.finished[0]
797 };
798 if setup_ready && prior_ready && turn_ready {
799 break;
800 }
801 p = self.sync.changed.wait(p).unwrap();
802 }
803 drop(p);
804 Ok(self.sync.primary.lock().unwrap())
805 }
806
807 fn draft_end(&self, round: usize) {
808 let mut p = self.sync.progress.lock().unwrap();
809 p.draft_done[self.lane] = round + 1;
810 self.sync.changed.notify_all();
811 }
812
813 /// Admit stage 0 and return whether this lane owns the interval's one reverse fence.
814 /// Lane B releases as soon as lane A has issued its boundary TX, not after A's full body.
815 fn stage0_begin(&self, round: usize) -> Result<bool, Box<dyn std::error::Error>> {
816 let peer = self.peer();
817 let mut p = self.sync.progress.lock().unwrap();
818 loop {
819 if p.aborted {
820 return Err(Self::aborted());
821 }
822 let ready = if self.lane == 0 {
823 p.draft_done[0] > round
824 && (p.draft_done[1] > round || p.finished[1])
825 } else {
826 p.draft_done[1] > round
827 && (p.stage0_done[0] > round || p.finished[0])
828 };
829 if ready {
830 return Ok(self.lane == 0 || p.finished[peer]);
831 }
832 p = self.sync.changed.wait(p).unwrap();
833 }
834 }
835
836 fn stage0_end(&self, round: usize) {
837 let mut p = self.sync.progress.lock().unwrap();
838 p.stage0_done[self.lane] = round + 1;
839 self.sync.changed.notify_all();
840 }
841
842 /// Stage 1 is single-owner per engine. A proceeds immediately after its own ticket; B waits
843 /// for A's full stage1/head issue so only A.S1 and B.S0 can overlap.
844 fn stage1_begin(&self, round: usize) -> Result<(), Box<dyn std::error::Error>> {
845 let mut p = self.sync.progress.lock().unwrap();
846 while !p.aborted
847 && !(p.stage0_done[self.lane] > round
848 && (self.lane == 0 || p.verify_done[0] > round || p.finished[0]))
849 {
850 p = self.sync.changed.wait(p).unwrap();
851 }
852 if p.aborted { Err(Self::aborted()) } else { Ok(()) }
853 }
854
855 fn verify_end(&self, round: usize) {
856 let mut p = self.sync.progress.lock().unwrap();
857 p.verify_done[self.lane] = round + 1;
858 self.sync.changed.notify_all();
859 }
860
861 fn accept_begin(
862 &self,
863 round: usize,
864 ) -> Result<std::sync::MutexGuard<'_, ()>, Box<dyn std::error::Error>> {
865 let mut p = self.sync.progress.lock().unwrap();
866 loop {
867 if p.aborted {
868 return Err(Self::aborted());
869 }
870 let ready = if self.lane == 0 {
871 p.verify_done[0] > round
872 && (p.verify_done[1] > round || p.finished[1])
873 } else {
874 p.verify_done[1] > round
875 && (p.accept_done[0] > round || p.finished[0])
876 };
877 if ready {
878 break;
879 }
880 p = self.sync.changed.wait(p).unwrap();
881 }
882 drop(p);
883 Ok(self.sync.primary.lock().unwrap())
884 }
885
886 fn accept_end(&self, round: usize) {
887 let mut p = self.sync.progress.lock().unwrap();
888 p.accept_done[self.lane] = round + 1;
889 self.sync.changed.notify_all();
890 }
891
892 fn primary(&self) -> std::sync::MutexGuard<'_, ()> {
893 self.sync.primary.lock().unwrap()
894 }
895
896 fn finish(&self, failed: bool) {
897 let mut p = self.sync.progress.lock().unwrap();
898 p.finished[self.lane] = true;
899 p.aborted |= failed;
900 self.sync.changed.notify_all();
901 }
902}
903
904struct SpecPipeFinish<'a> {
905 lane: &'a SpecPipeLane,
906 closed: bool,
907}
908
909impl<'a> SpecPipeFinish<'a> {
910 fn new(lane: &'a SpecPipeLane) -> Self {
911 Self { lane, closed: false }
912 }
913
914 fn close(&mut self, failed: bool) {
915 self.lane.finish(failed);
916 self.closed = true;
917 }
918}
919
920impl Drop for SpecPipeFinish<'_> {
921 fn drop(&mut self) {
922 if !self.closed {
923 self.lane.finish(true);
924 }
925 }
926}
927
928/// Scoped transfer of one exclusively-borrowed session to the second host issue thread.
929/// `CudaGraph` is not marked Send by cudarc because its raw driver handles carry no automatic
930/// trait. CUDA driver graph handles are context-scoped rather than OS-thread-affine; the caller
931/// binds that context before touching the session, joins before returning, and never aliases the
932/// pointer. Keep this exception local to the experimental pair call instead of marking the public
933/// session type Send.
934struct SpecPipeSessionPtr(*mut SpecSession);
935
936unsafe impl Send for SpecPipeSessionPtr {}
937
938impl SpecPipeSessionPtr {
939 unsafe fn get_mut(&mut self) -> &mut SpecSession {
940 unsafe { &mut *self.0 }
941 }
942}
943
944/// Per-session persistent draft-graph context: the captured CUDA graph(s) plus the device
945/// buffers whose POINTERS the capture bakes. Reuse legality: the greedy capture bakes only
946/// session-stable pointers (the session's own MtpScratch KV — allocated once, never realloc'd;
947/// the model's resident embedding; the process-wide OnceLock p_min) and the g_* buffers held
948/// HERE — so one capture serves the session's whole lifetime. The sampled capture additionally
949/// bakes (seed, temp) as capture-time constants and needs k q-slots — keyed by `s_key`, dropped
950/// and recaptured when a pool-resumed request changes them. `*_failed` memoizes a failed capture
951/// so the eager fallback doesn't pay a doomed capture attempt every burst.
952pub(crate) struct DraftGraphCtx {
953 g_tok: CudaSlice<u32>,
954 g_pos: CudaSlice<i32>,
955 g_seed: CudaSlice<f32>,
956 g_p: CudaSlice<f32>,
957 g_ctr: CudaSlice<u32>,
958 g_q: CudaSlice<f32>,
959 g_perturb: CudaSlice<f32>,
960 q_slots: Vec<CudaSlice<f32>>,
961 /// DRAFT-SIDE GRAMMAR MASK (lane/draft-mask): packed allowed-set words over the DRAFT
962 /// head's vocab, at a STABLE address so the captured draft graph's mask node reads the
963 /// per-position contents the host re-uploads before each replay (the graph-promote
964 /// pattern from decode.rs). Empty unless the session drafts under a grammar.
965 g_dmask: CudaSlice<u32>,
966 /// was `graph` captured WITH the mask node? A parked graph of the wrong shape is dropped.
967 graph_masked: bool,
968 graph: Option<cudarc::driver::CudaGraph>,
969 graph_s: Option<cudarc::driver::CudaGraph>,
970 /// Failed-capture memoization for both graphs — LOUD on flip, cleared on pool resume
971 /// (audit Q2, the TRT #16072 silent-permanent-coverage-loss class).
972 failed: DraftGraphFallback,
973 /// (seed, temp.to_bits(), k) baked into graph_s at its capture.
974 s_key: Option<(u64, u32, usize)>,
975 /// CAPTURE-RETAIN keepers (#68 root cause, 2026-08-04): the warmup-run transients whose
976 /// pool addresses the captured graph(s) bake. Without these, the transients return to the
977 /// pool at capture-body exit and later work (burst-boundary prime/fill/commit passes, or a
978 /// co-served session in the worker) reuses those addresses — the persisted graph's replay
979 /// then reads/writes live unrelated buffers (exactness corruption, first seen as the ST
980 /// serve-spec 4B graph-arm corruption; one-shot CLI calls never re-shuffled the pool, which
981 /// is why run-spec K=1..8 passed on the same checkpoint). Same fix class as
982 /// capture_graph_retained's gemma/decode.rs sites — hold as long as the graph replays.
983 keeper: Vec<Box<dyn std::any::Any + Send>>,
984 keeper_s: Vec<Box<dyn std::any::Any + Send>>,
985}
986
987/// Failed-capture memoization for the two draft graphs (audit Q2, 2026-08-05 — the
988/// TRT #16072 trap class: pressure-triggered, silent, long-lived coverage loss).
989///
990/// Three contracts:
991/// - LOUD FLIP: `mark_*` returns the warn line exactly on the false→true transition
992/// (returned, not printed, so the once-per-flip contract is unit-testable); the caller
993/// `eprintln!`s it UNCONDITIONALLY — a dropped draft graph is never silent. Re-marking
994/// an already-failed graph returns None (the per-burst memoization that keeps the eager
995/// fallback from paying a doomed capture attempt every burst).
996/// - RESET ON RESUME: `reset_on_resume` clears both flags — a parked session resumed by a
997/// NEW request gets one fresh capture chance instead of carrying a transient-pressure
998/// failure for the pool's whole lifetime. Returns the note line only when a flag was
999/// actually set (quiet on the common clean-resume path).
1000/// - Shape-change clears (`clear_*`) stay silent, exactly as before: they precede a fresh
1001/// capture attempt whose own failure would re-flip loudly.
1002#[derive(Default)]
1003pub(crate) struct DraftGraphFallback {
1004 greedy: bool,
1005 sampled: bool,
1006}
1007impl DraftGraphFallback {
1008 fn mark_greedy(&mut self, reason: &str) -> Option<String> {
1009 if self.greedy {
1010 return None;
1011 }
1012 self.greedy = true;
1013 Some(format!(
1014 "[spec] WARN: draft-graph capture failed ({reason}); eager fallback until session resume"
1015 ))
1016 }
1017 fn mark_sampled(&mut self, reason: &str) -> Option<String> {
1018 if self.sampled {
1019 return None;
1020 }
1021 self.sampled = true;
1022 Some(format!(
1023 "[spec] WARN: sampled draft-graph capture failed ({reason}); eager fallback until session resume"
1024 ))
1025 }
1026 fn greedy_failed(&self) -> bool {
1027 self.greedy
1028 }
1029 fn sampled_failed(&self) -> bool {
1030 self.sampled
1031 }
1032 fn clear_greedy(&mut self) {
1033 self.greedy = false;
1034 }
1035 fn clear_sampled(&mut self) {
1036 self.sampled = false;
1037 }
1038 /// Pool-resume reset: both graphs get a fresh capture chance. Some(note) iff any flag
1039 /// was set (so clean resumes stay quiet).
1040 pub(crate) fn reset_on_resume(&mut self) -> Option<String> {
1041 if !self.greedy && !self.sampled {
1042 return None;
1043 }
1044 let which = match (self.greedy, self.sampled) {
1045 (true, true) => "greedy+sampled",
1046 (true, false) => "greedy",
1047 _ => "sampled",
1048 };
1049 self.greedy = false;
1050 self.sampled = false;
1051 Some(format!(
1052 "[spec] draft-graph fallback reset on session resume ({which}); recapture eligible"
1053 ))
1054 }
1055}
1056
1057impl DraftGraphCtx {
1058 fn new(e: &Engine, n_embd: usize, qlen: usize) -> Result<Self, Box<dyn std::error::Error>> {
1059 Ok(DraftGraphCtx {
1060 g_tok: e.alloc_u32_zeroed(1)?,
1061 g_pos: e.htod_i32(&[0])?,
1062 g_seed: e.zeros(n_embd)?,
1063 g_p: e.zeros(1)?,
1064 g_ctr: e.alloc_u32_zeroed(1)?,
1065 g_q: e.zeros(qlen)?,
1066 g_perturb: e.zeros(qlen)?,
1067 q_slots: Vec::new(),
1068 g_dmask: e.alloc_u32_zeroed(1)?,
1069 graph_masked: false,
1070 graph: None,
1071 graph_s: None,
1072 failed: DraftGraphFallback::default(),
1073 s_key: None,
1074 keeper: Vec::new(),
1075 keeper_s: Vec::new(),
1076 })
1077 }
1078}
1079
1080pub(crate) struct MtpScratch {
1081 kv: KvLayer,
1082 /// Logical row capacity. On the graph/DC draft path it also doubles as fa_decode_dc's
1083 /// bucket_max: n_splits is sized from it ONCE, so the graph captured at round 0 stays valid
1084 /// for every later t_kv. Step35 refuses that path and may back this logical extent with the
1085 /// smaller host-indexed SWA ring instead.
1086 cap: usize,
1087}
1088
1089fn mtp_scratch_layout(
1090 cfg: &memra_gguf::config::ModelConfig,
1091 geom: Option<&crate::hybrid::DraftGeom>,
1092) -> (usize, usize, usize, usize) {
1093 // Student draft heads carry fewer KV heads (head_dim unchanged) -> smaller scratch rows.
1094 let n_head_kv = geom.map(|g| g.n_head_kv).unwrap_or(cfg.n_head_kv as usize);
1095 let head_dim_k = cfg.head_dim_k as usize;
1096 let head_dim_v = cfg.head_dim_v as usize;
1097 assert!(
1098 head_dim_k % 32 == 0 && head_dim_v % 32 == 0,
1099 "KVQUANT requires head_dim%32==0 (MTP scratch)"
1100 );
1101 let kv_dim_k = head_dim_k * n_head_kv;
1102 let kv_dim_v = head_dim_v * n_head_kv;
1103 // The fp8-KV arm deliberately does not reach the draft scratch; keep the exact format
1104 // policy shared with `MtpScratch::new` so admission scales the same allocation.
1105 let (kbb, vbb) = crate::kv_blk_bytes();
1106 let k_tok_bytes = (kv_dim_k / 32) * kbb;
1107 let v_tok_bytes = (kv_dim_v / 32) * vbb;
1108 (kv_dim_k, kv_dim_v, k_tok_bytes, v_tok_bytes)
1109}
1110
1111impl MtpScratch {
1112 fn new(
1113 e: &Engine,
1114 cfg: &memra_gguf::config::ModelConfig,
1115 cap: usize,
1116 geom: Option<&crate::hybrid::DraftGeom>,
1117 ) -> Result<Self, Box<dyn std::error::Error>> {
1118 // env-selected KV formats (default 34/24). The fp8-KV arm (MEMRA_KV_FP8) deliberately
1119 // does NOT reach the draft scratch: fp8 drafts drifted acceptance 69-88% -> 46%
1120 // (2026-07-12 A/B); the scratch is tiny, so it keeps baseline q8_0/q5_1 numerics
1121 // while the TRUNK cache carries the fp8 depth win. Scratch append/fa pass g=false.
1122 let (kv_dim_k, kv_dim_v, k_tok_bytes, v_tok_bytes) =
1123 mtp_scratch_layout(cfg, geom);
1124 let ring = if crate::cache::swa_ring_on() && cfg.arch.is_step35() {
1125 let window = cfg.step35.as_ref().unwrap().sliding_window as usize;
1126 Some(crate::cache::KvRing::new(
1127 crate::cache::swa_ring_rows(window, cap),
1128 window,
1129 ))
1130 } else {
1131 None
1132 };
1133 let alloc_rows = ring.as_ref().map(crate::cache::KvRing::rows).unwrap_or(cap);
1134 Ok(MtpScratch {
1135 kv: KvLayer {
1136 k: e.alloc_u8(alloc_rows * k_tok_bytes)?,
1137 v: e.alloc_u8(alloc_rows * v_tok_bytes)?,
1138 kv_dim_k,
1139 kv_dim_v,
1140 k_tok_bytes,
1141 v_tok_bytes,
1142 len: 0,
1143 ring,
1144 len_d: e.htod_i32(&[0])?,
1145 },
1146 cap,
1147 })
1148 }
1149 /// Set BOTH length counters: the host mirror AND the device len_d the captured append/fa read
1150 /// (a 4-byte in-place htod — the counter pointer is baked into the graph, never realloc'd).
1151 /// This is the ONLY truncation/rollback mechanism the persistent draft KV needs.
1152 fn set_len(&mut self, e: &Engine, n: usize) -> Result<(), Box<dyn std::error::Error>> {
1153 if self.kv.ring.as_ref().is_some_and(|ring| !ring.can_rewind_to(n)) {
1154 return Err("SWA ring MTP checkpoint has been lapped; full re-prime required".into());
1155 }
1156 self.kv.len = n;
1157 e.set_i32_one(&mut self.kv.len_d, n as i32)
1158 }
1159
1160 fn can_rewind_to(&self, n: usize) -> bool {
1161 self.kv.ring.as_ref().is_none_or(|ring| ring.can_rewind_to(n))
1162 }
1163}
1164
1165/// Retained verify intermediates for the REPLAY-FREE partial accept (2026-07-03, the profiled
1166/// #1 spec cost at long ctx: the partial-accept replay was a DUPLICATE trunk pass — ~0.54 extra
1167/// full weight reads per round — recomputing columns the verify had already produced
1168/// bit-identically). Holds, per linear layer, everything needed to rebuild its recurrent state
1169/// to "after the first j verify columns" WITHOUT re-running the trunk:
1170/// - BATCHED-path layers (`gdn`): the exact token-major inputs the round's ONE gdn_scan
1171/// consumed. A prefix re-run of the SAME kernel (t=j) from the snapshot state is bit-identical
1172/// to the first j iterations of the verify's scan — the kernel's t-loop carries state in
1173/// registers and iteration t never depends on T. `qkv_mixed` (the conv input) feeds the
1174/// pure-copy ring rebuild.
1175/// - PER-COLUMN-path layers (`cols`): dtod clones of (conv_state, ssm_state) taken after each
1176/// column 0..t-2 — pure copies of the actual chain states (the last column is never a rebuild
1177/// target: j <= t-1).
1178/// Full-attn layers need nothing: their verify KV rows are bit-identical to eager's (the
1179/// decode-exact contract; verify-probe pins it), so rollback = len truncation.
1180struct GdnStash {
1181 qkv_mixed: CudaSlice<f32>, // [t, conv_dim] token-major (conv input)
1182 q_l2: CudaSlice<f32>,
1183 k_l2: CudaSlice<f32>,
1184 v_g: CudaSlice<f32>, // [t, num_v, d_state]
1185 g_log: CudaSlice<f32>,
1186 beta: CudaSlice<f32>, // [t, num_v]
1187}
1188struct VerifyCkpt {
1189 gdn: Vec<Option<GdnStash>>, // [n_layer], Some iff batched linear path ran
1190 cols: Vec<Option<Vec<(CudaSlice<f32>, CudaSlice<f32>)>>>, // [n_layer][col] = (conv, ssm) after col
1191}
1192impl VerifyCkpt {
1193 fn new(n_layer: usize) -> Self {
1194 VerifyCkpt {
1195 gdn: (0..n_layer).map(|_| None).collect(),
1196 cols: (0..n_layer).map(|_| None).collect(),
1197 }
1198 }
1199}
1200
1201/// The stage-0/TX half of one PP verify. The boundary slot is the ownership token: stage 1
1202/// consumes exactly the slot selected by `tx()` / `tx_pipelined()`, never a slot inferred from
1203/// a logical round number.
1204struct VerifyBoundaryTicket {
1205 rt: &'static crate::pp::PpNRt,
1206 caller_stream: std::sync::Arc<cudarc::driver::CudaStream>,
1207 slot: usize,
1208 pos0: usize,
1209 t: usize,
1210 payload: usize,
1211 n_st: usize,
1212 pipelined: bool,
1213 pp_anatomy: bool,
1214 pp_started: std::time::Instant,
1215 reverse_ms: f64,
1216 stage0_ms: f64,
1217 tx_ms: f64,
1218 trace: Option<SpecPipeTraceCtx>,
1219}
1220
1221/// Explicit OPTIPIPE diagnostic control. Forced modes are set only by `optipipe-gate`; the
1222/// increment-2 controller can also be armed by the server's fresh-process research door.
1223#[derive(Clone, Copy, Debug, PartialEq, Eq)]
1224pub enum OptiForkGateMode {
1225 Disabled,
1226 Hit,
1227 Miss,
1228 Alternate,
1229 Abort,
1230 Controller,
1231}
1232
1233static OPTI_FORK_GATE_MODE: std::sync::atomic::AtomicU8 =
1234 std::sync::atomic::AtomicU8::new(0);
1235static OPTI_CONTROLLER_THRESHOLD: std::sync::atomic::AtomicU32 =
1236 std::sync::atomic::AtomicU32::new(0);
1237static OPTI_FORK_ATTEMPTS: std::sync::atomic::AtomicU64 =
1238 std::sync::atomic::AtomicU64::new(0);
1239static OPTI_FORK_HITS: std::sync::atomic::AtomicU64 =
1240 std::sync::atomic::AtomicU64::new(0);
1241static OPTI_FORK_MISSES: std::sync::atomic::AtomicU64 =
1242 std::sync::atomic::AtomicU64::new(0);
1243static OPTI_FORK_ABORT_DRAINS: std::sync::atomic::AtomicU64 =
1244 std::sync::atomic::AtomicU64::new(0);
1245static OPTI_FORK_REFUSALS: std::sync::atomic::AtomicU64 =
1246 std::sync::atomic::AtomicU64::new(0);
1247static OPTI_GATE_CHECKS: std::sync::atomic::AtomicU64 =
1248 std::sync::atomic::AtomicU64::new(0);
1249static OPTI_GATE_ADMITS: std::sync::atomic::AtomicU64 =
1250 std::sync::atomic::AtomicU64::new(0);
1251static OPTI_GATE_REJECTS: std::sync::atomic::AtomicU64 =
1252 std::sync::atomic::AtomicU64::new(0);
1253static OPTI_RECONCILES: std::sync::atomic::AtomicU64 =
1254 std::sync::atomic::AtomicU64::new(0);
1255static OPTI_WASTED_DRAFT_TOKENS: std::sync::atomic::AtomicU64 =
1256 std::sync::atomic::AtomicU64::new(0);
1257static OPTI_SHADOW_DRAFT_TOKENS: std::sync::atomic::AtomicU64 =
1258 std::sync::atomic::AtomicU64::new(0);
1259static OPTI_BREAKER_TRIPS: std::sync::atomic::AtomicU64 =
1260 std::sync::atomic::AtomicU64::new(0);
1261
1262impl OptiForkGateMode {
1263 fn code(self) -> u8 {
1264 match self {
1265 Self::Disabled => 0,
1266 Self::Hit => 1,
1267 Self::Miss => 2,
1268 Self::Alternate => 3,
1269 Self::Abort => 4,
1270 Self::Controller => 5,
1271 }
1272 }
1273
1274 fn configured() -> Self {
1275 match OPTI_FORK_GATE_MODE.load(std::sync::atomic::Ordering::Relaxed) {
1276 1 => Self::Hit,
1277 2 => Self::Miss,
1278 3 => Self::Alternate,
1279 4 => Self::Abort,
1280 5 => Self::Controller,
1281 _ => Self::Disabled,
1282 }
1283 }
1284
1285 fn action(self, generation: u64) -> OptiForkAction {
1286 match self {
1287 Self::Hit => OptiForkAction::Hit,
1288 Self::Miss => OptiForkAction::Miss,
1289 Self::Alternate if generation & 1 == 0 => OptiForkAction::Hit,
1290 Self::Alternate => OptiForkAction::Miss,
1291 Self::Abort => OptiForkAction::Abort,
1292 Self::Disabled | Self::Controller => {
1293 unreachable!("non-forced mode cannot choose a forced fork action")
1294 }
1295 }
1296 }
1297
1298 fn is_forced(self) -> bool {
1299 matches!(self, Self::Hit | Self::Miss | Self::Alternate | Self::Abort)
1300 }
1301}
1302
1303/// Arm or disarm the forced harness. Serving uses only `set_optipipe_controller_threshold`.
1304pub fn set_optipipe_gate_mode(mode: OptiForkGateMode) {
1305 OPTI_FORK_GATE_MODE.store(mode.code(), std::sync::atomic::Ordering::Relaxed);
1306}
1307
1308/// Arm the increment-2 diagnostic controller. The threshold applies to the uncalibrated
1309/// two-token draft-probability product. Serving can call this only through its explicit
1310/// fresh-process research door; the absent-door default remains byte-for-byte disabled.
1311pub fn set_optipipe_controller_threshold(threshold: f32) {
1312 assert!(threshold.is_finite() && (0.0..=1.0).contains(&threshold));
1313 OPTI_CONTROLLER_THRESHOLD.store(threshold.to_bits(), std::sync::atomic::Ordering::Relaxed);
1314 set_optipipe_gate_mode(OptiForkGateMode::Controller);
1315}
1316
1317#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
1318pub struct OptiForkGateStats {
1319 pub attempts: u64,
1320 pub hits: u64,
1321 pub misses: u64,
1322 pub abort_drains: u64,
1323 pub refusals: u64,
1324 pub gate_checks: u64,
1325 pub gate_admits: u64,
1326 pub gate_rejects: u64,
1327 pub reconciles: u64,
1328 pub wasted_draft_tokens: u64,
1329 pub shadow_draft_tokens: u64,
1330 pub breaker_trips: u64,
1331}
1332
1333#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
1334pub struct OptiForkStateIdentity {
1335 pub trunk_kv_bytes: usize,
1336 pub recurrent_bytes: usize,
1337 pub scratch_kv_bytes: usize,
1338 pub hidden_bytes: usize,
1339}
1340
1341pub fn reset_optipipe_gate_stats() {
1342 for counter in [
1343 &OPTI_FORK_ATTEMPTS,
1344 &OPTI_FORK_HITS,
1345 &OPTI_FORK_MISSES,
1346 &OPTI_FORK_ABORT_DRAINS,
1347 &OPTI_FORK_REFUSALS,
1348 &OPTI_GATE_CHECKS,
1349 &OPTI_GATE_ADMITS,
1350 &OPTI_GATE_REJECTS,
1351 &OPTI_RECONCILES,
1352 &OPTI_WASTED_DRAFT_TOKENS,
1353 &OPTI_SHADOW_DRAFT_TOKENS,
1354 &OPTI_BREAKER_TRIPS,
1355 ] {
1356 counter.store(0, std::sync::atomic::Ordering::Relaxed);
1357 }
1358}
1359
1360pub fn optipipe_gate_stats() -> OptiForkGateStats {
1361 let load = |v: &std::sync::atomic::AtomicU64| {
1362 v.load(std::sync::atomic::Ordering::Relaxed)
1363 };
1364 OptiForkGateStats {
1365 attempts: load(&OPTI_FORK_ATTEMPTS),
1366 hits: load(&OPTI_FORK_HITS),
1367 misses: load(&OPTI_FORK_MISSES),
1368 abort_drains: load(&OPTI_FORK_ABORT_DRAINS),
1369 refusals: load(&OPTI_FORK_REFUSALS),
1370 gate_checks: load(&OPTI_GATE_CHECKS),
1371 gate_admits: load(&OPTI_GATE_ADMITS),
1372 gate_rejects: load(&OPTI_GATE_REJECTS),
1373 reconciles: load(&OPTI_RECONCILES),
1374 wasted_draft_tokens: load(&OPTI_WASTED_DRAFT_TOKENS),
1375 shadow_draft_tokens: load(&OPTI_SHADOW_DRAFT_TOKENS),
1376 breaker_trips: load(&OPTI_BREAKER_TRIPS),
1377 }
1378}
1379
1380#[derive(Clone, Copy, Debug)]
1381struct OptiControllerPolicy {
1382 threshold: f32,
1383 consecutive_misses: u8,
1384 breaker_tripped: bool,
1385}
1386
1387impl OptiControllerPolicy {
1388 fn configured() -> Self {
1389 Self {
1390 threshold: f32::from_bits(
1391 OPTI_CONTROLLER_THRESHOLD.load(std::sync::atomic::Ordering::Relaxed),
1392 ),
1393 consecutive_misses: 0,
1394 breaker_tripped: false,
1395 }
1396 }
1397
1398 fn admit(&self, q_proxy: f32) -> bool {
1399 q_proxy.is_finite()
1400 && (0.0..=1.0).contains(&q_proxy)
1401 && (self.threshold == 0.0 || (!self.breaker_tripped && q_proxy >= self.threshold))
1402 }
1403
1404 /// Returns true exactly when this resolution newly trips the three-miss breaker.
1405 fn resolve(&mut self, hit: bool) -> bool {
1406 // q*=0 is the lane's explicit unconditional measurement arm. Its purpose is to price
1407 // every optimistic opportunity, so the safety breaker is measured separately and must
1408 // not silently turn this arm into "three attempts then serial".
1409 if self.threshold == 0.0 {
1410 self.consecutive_misses = 0;
1411 return false;
1412 }
1413 if hit {
1414 self.consecutive_misses = 0;
1415 return false;
1416 }
1417 self.consecutive_misses = self.consecutive_misses.saturating_add(1);
1418 if !self.breaker_tripped && self.consecutive_misses >= 3 {
1419 self.breaker_tripped = true;
1420 return true;
1421 }
1422 false
1423 }
1424}
1425
1426#[derive(Clone, Copy, Debug, PartialEq, Eq)]
1427enum OptiForkAction {
1428 Hit,
1429 Miss,
1430 Abort,
1431}
1432
1433#[derive(Clone, Copy, Debug, PartialEq, Eq)]
1434struct OptiForkGeneration {
1435 id: u64,
1436 slot: usize,
1437}
1438
1439#[derive(Default)]
1440struct OptiForkGenerationTracker {
1441 next: u64,
1442 live: [Option<u64>; 2],
1443}
1444
1445impl OptiForkGenerationTracker {
1446 fn reserve(&mut self) -> Result<OptiForkGeneration, Box<dyn std::error::Error>> {
1447 let generation = OptiForkGeneration {
1448 id: self.next,
1449 slot: (self.next & 1) as usize,
1450 };
1451 if let Some(live) = self.live[generation.slot] {
1452 return Err(format!(
1453 "optipipe snapshot slot {} still owns generation {live}; refusing to overwrite it",
1454 generation.slot,
1455 )
1456 .into());
1457 }
1458 self.next += 1;
1459 self.live[generation.slot] = Some(generation.id);
1460 Ok(generation)
1461 }
1462
1463 fn retire(&mut self, generation: OptiForkGeneration)
1464 -> Result<(), Box<dyn std::error::Error>> {
1465 match self.live[generation.slot] {
1466 Some(id) if id == generation.id => {
1467 self.live[generation.slot] = None;
1468 Ok(())
1469 }
1470 other => Err(format!(
1471 "optipipe generation teardown mismatch: ticket={} slot={} live={other:?}",
1472 generation.id, generation.slot,
1473 )
1474 .into()),
1475 }
1476 }
1477}
1478
1479struct OptiForkSeedGeneration {
1480 h_seed: CudaSlice<f32>,
1481 fill_prev: CudaSlice<f32>,
1482 scratch_len: usize,
1483}
1484
1485/// Allocate or refresh one full checkpoint through the engine that owns each PP stage. The
1486/// generic cache helper accepts one device and therefore cannot copy GDN state split across
1487/// devices. KV lengths and position stay host metadata; only recurrent buffers need stage-local
1488/// device ownership.
1489fn opti_snapshot_stage_owned(
1490 e: &Engine,
1491 cache: &Cache,
1492 rt: &'static crate::pp::PpNRt,
1493 fence: &[usize],
1494) -> Result<crate::cache::CacheSnapshot, Box<dyn std::error::Error>> {
1495 let n = cache.kv.len();
1496 let mut snapshot = crate::cache::CacheSnapshot {
1497 kv_len: vec![None; n],
1498 conv: (0..n).map(|_| None).collect(),
1499 ssm: (0..n).map(|_| None).collect(),
1500 pos: cache.pos,
1501 };
1502 opti_snapshot_stage_owned_into(e, cache, rt, fence, &mut snapshot)?;
1503 Ok(snapshot)
1504}
1505
1506fn opti_snapshot_stage_owned_into(
1507 e: &Engine,
1508 cache: &Cache,
1509 rt: &'static crate::pp::PpNRt,
1510 fence: &[usize],
1511 snapshot: &mut crate::cache::CacheSnapshot,
1512) -> Result<(), Box<dyn std::error::Error>> {
1513 if fence.len() != rt.n_stages() + 1 || snapshot.kv_len.len() != cache.kv.len() {
1514 return Err("optipipe stage-owned snapshot shape mismatch".into());
1515 }
1516 for stage in 0..rt.n_stages() {
1517 opti_snapshot_one_stage_owned_into(e, cache, rt, fence, stage, snapshot)?;
1518 }
1519 snapshot.pos = cache.pos;
1520 Ok(())
1521}
1522
1523/// Refresh one PP stage of a checkpoint. Increment 2 uses this split form so stage 0's
1524/// optimistic post-N state is captured before N+1 stage 0 is queued, while stage 1's matching
1525/// post-N state is captured only after N stage 1 is enqueued. Calling the all-stage helper at
1526/// either point would capture one side of the fork at the wrong generation.
1527fn opti_snapshot_one_stage_owned_into(
1528 e: &Engine,
1529 cache: &Cache,
1530 rt: &'static crate::pp::PpNRt,
1531 fence: &[usize],
1532 stage: usize,
1533 snapshot: &mut crate::cache::CacheSnapshot,
1534) -> Result<(), Box<dyn std::error::Error>> {
1535 if fence.len() != rt.n_stages() + 1
1536 || snapshot.kv_len.len() != cache.kv.len()
1537 || stage >= rt.n_stages()
1538 {
1539 return Err("optipipe single-stage snapshot shape mismatch".into());
1540 }
1541 let _scope = rt.enter(stage);
1542 let owner = rt.engine(stage, e);
1543 for il in fence[stage]..fence[stage + 1] {
1544 snapshot.kv_len[il] = cache.kv[il].as_ref().map(|kv| kv.len);
1545 match &cache.recur[il] {
1546 Some(recur) => {
1547 match snapshot.conv[il].as_mut() {
1548 Some(dst) => owner.copy_into(
1549 dst,
1550 0,
1551 &recur.conv_state,
1552 recur.conv_state.len(),
1553 )?,
1554 None => snapshot.conv[il] = Some(owner.clone_dtod(&recur.conv_state)?),
1555 }
1556 match snapshot.ssm[il].as_mut() {
1557 Some(dst) => owner.copy_into(
1558 dst,
1559 0,
1560 &recur.ssm_state,
1561 recur.ssm_state.len(),
1562 )?,
1563 None => snapshot.ssm[il] = Some(owner.clone_dtod(&recur.ssm_state)?),
1564 }
1565 }
1566 None if snapshot.conv[il].is_some() || snapshot.ssm[il].is_some() => {
1567 return Err(
1568 format!("optipipe stage-owned snapshot layer {il} changed shape").into()
1569 );
1570 }
1571 None => {}
1572 }
1573 }
1574 snapshot.pos = cache.pos;
1575 Ok(())
1576}
1577
1578/// Increment-1 persistent fork state. Exactly two snapshot/seed slots alternate; a live ticket
1579/// names its generation and keeps teardown fail-closed. Only stage 0 is allowed to mutate before
1580/// resolve, so the reconcile tables and conditional restores are stage-local.
1581struct OptiForkState {
1582 mode: OptiForkGateMode,
1583 controller: Option<OptiControllerPolicy>,
1584 generations: OptiForkGenerationTracker,
1585 active_snapshot_slot: usize,
1586 alternate_snapshot: crate::cache::CacheSnapshot,
1587 seeds: [OptiForkSeedGeneration; 2],
1588 rt: &'static crate::pp::PpNRt,
1589 fence: [usize; 3],
1590 split: usize,
1591 len_ptrs: CudaSlice<u64>,
1592 saved_lens: CudaSlice<i32>,
1593 forced_acc: CudaSlice<u32>,
1594 valid: CudaSlice<u32>,
1595 stage0_stream: std::sync::Arc<cudarc::driver::CudaStream>,
1596 logical_payload_bytes: [usize; 2],
1597}
1598
1599struct OptiForkTicket {
1600 generation: OptiForkGeneration,
1601 boundary: Option<VerifyBoundaryTicket>,
1602 drain: std::sync::Arc<cudarc::driver::CudaStream>,
1603 settled: bool,
1604}
1605
1606struct OptiControllerTicket {
1607 generation: OptiForkGeneration,
1608 boundary: Option<VerifyBoundaryTicket>,
1609 ckpt: Option<VerifyCkpt>,
1610 verify_tokens: [u32; 2],
1611 draft_prob: f32,
1612 eager_seed: Option<CudaSlice<f32>>,
1613 q_proxy: f32,
1614 scratch_len: usize,
1615 issued_at: std::time::Instant,
1616 drain: std::sync::Arc<cudarc::driver::CudaStream>,
1617 settled: bool,
1618}
1619
1620struct OptiControllerPrepared {
1621 verify_tokens: [u32; 2],
1622 draft_prob: f32,
1623 eager_seed: Option<CudaSlice<f32>>,
1624 q_proxy: f32,
1625 scratch_len: usize,
1626}
1627
1628impl OptiControllerTicket {
1629 fn take_boundary(&mut self) -> VerifyBoundaryTicket {
1630 self.boundary
1631 .take()
1632 .expect("controller boundary ticket already consumed")
1633 }
1634
1635 fn take_ckpt(&mut self) -> VerifyCkpt {
1636 self.ckpt
1637 .take()
1638 .expect("controller verify checkpoint already consumed")
1639 }
1640
1641 fn take_eager_seed(&mut self) -> Option<CudaSlice<f32>> {
1642 self.eager_seed.take()
1643 }
1644
1645 fn settle(&mut self) {
1646 self.settled = true;
1647 }
1648}
1649
1650impl Drop for OptiControllerTicket {
1651 fn drop(&mut self) {
1652 if !self.settled {
1653 let _ = self.drain.synchronize();
1654 OPTI_FORK_ABORT_DRAINS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
1655 }
1656 }
1657}
1658
1659impl OptiForkTicket {
1660 fn take_boundary(&mut self) -> VerifyBoundaryTicket {
1661 self.boundary.take().expect("fork ticket boundary already consumed")
1662 }
1663
1664 fn settle(&mut self) {
1665 self.settled = true;
1666 }
1667}
1668
1669impl Drop for OptiForkTicket {
1670 fn drop(&mut self) {
1671 if !self.settled {
1672 let _ = self.drain.synchronize();
1673 OPTI_FORK_ABORT_DRAINS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
1674 }
1675 }
1676}
1677
1678impl OptiForkState {
1679 #[allow(clippy::too_many_arguments)]
1680 fn new(
1681 e: &Engine,
1682 cache: &Cache,
1683 mode: OptiForkGateMode,
1684 alternate_snapshot: crate::cache::CacheSnapshot,
1685 h_seed: &CudaSlice<f32>,
1686 fill_prev: &CudaSlice<f32>,
1687 rt: &'static crate::pp::PpNRt,
1688 split: usize,
1689 n_layer: usize,
1690 ) -> Result<Self, Box<dyn std::error::Error>> {
1691 let fence = [0, split, n_layer];
1692 let mut logical_payload_bytes = [0usize; 2];
1693 for stage in 0..2 {
1694 for il in fence[stage]..fence[stage + 1] {
1695 logical_payload_bytes[stage] += alternate_snapshot.conv[il]
1696 .as_ref()
1697 .map_or(0, |v| v.len() * std::mem::size_of::<f32>());
1698 logical_payload_bytes[stage] += alternate_snapshot.ssm[il]
1699 .as_ref()
1700 .map_or(0, |v| v.len() * std::mem::size_of::<f32>());
1701 }
1702 }
1703 let seeds = [
1704 OptiForkSeedGeneration {
1705 h_seed: e.clone_dtod(h_seed)?,
1706 fill_prev: e.clone_dtod(fill_prev)?,
1707 scratch_len: 0,
1708 },
1709 OptiForkSeedGeneration {
1710 h_seed: e.clone_dtod(h_seed)?,
1711 fill_prev: e.clone_dtod(fill_prev)?,
1712 scratch_len: 0,
1713 },
1714 ];
1715 let (len_ptrs, saved_lens, forced_acc, valid, stage0_stream) = {
1716 let _stage = rt.enter(0);
1717 let e0 = rt.engine(0, e);
1718 (
1719 crate::round_stream::kv_len_ptr_table_range(e0, cache, 0..split, None)?,
1720 e0.htod_i32(&vec![0; split])?,
1721 e0.alloc_u32_zeroed(2)?,
1722 e0.alloc_u32_zeroed(1)?,
1723 e0.stream(),
1724 )
1725 };
1726 logical_payload_bytes[0] += seeds
1727 .iter()
1728 .map(|seed| (seed.h_seed.len() + seed.fill_prev.len()) * std::mem::size_of::<f32>())
1729 .sum::<usize>();
1730 logical_payload_bytes[0] += len_ptrs.len() * std::mem::size_of::<u64>()
1731 + saved_lens.len() * std::mem::size_of::<i32>()
1732 + forced_acc.len() * std::mem::size_of::<u32>()
1733 + valid.len() * std::mem::size_of::<u32>();
1734 Ok(Self {
1735 mode,
1736 controller: (mode == OptiForkGateMode::Controller)
1737 .then(OptiControllerPolicy::configured),
1738 generations: OptiForkGenerationTracker::default(),
1739 active_snapshot_slot: 0,
1740 alternate_snapshot,
1741 seeds,
1742 rt,
1743 fence,
1744 split,
1745 len_ptrs,
1746 saved_lens,
1747 forced_acc,
1748 valid,
1749 stage0_stream,
1750 logical_payload_bytes,
1751 })
1752 }
1753
1754 fn reserve(&mut self, current_snapshot: &mut crate::cache::CacheSnapshot)
1755 -> Result<OptiForkGeneration, Box<dyn std::error::Error>> {
1756 let generation = self.generations.reserve()?;
1757 if generation.slot != self.active_snapshot_slot {
1758 std::mem::swap(current_snapshot, &mut self.alternate_snapshot);
1759 self.active_snapshot_slot = generation.slot;
1760 }
1761 Ok(generation)
1762 }
1763
1764 fn capture_seed(
1765 &mut self,
1766 e: &Engine,
1767 generation: OptiForkGeneration,
1768 h_seed: &CudaSlice<f32>,
1769 fill_prev: &CudaSlice<f32>,
1770 scratch_len: usize,
1771 ) -> Result<(), Box<dyn std::error::Error>> {
1772 let seed = &mut self.seeds[generation.slot];
1773 e.copy_into(&mut seed.h_seed, 0, h_seed, h_seed.len())?;
1774 e.copy_into(&mut seed.fill_prev, 0, fill_prev, fill_prev.len())?;
1775 seed.scratch_len = scratch_len;
1776 Ok(())
1777 }
1778
1779 fn ticket(&self, generation: OptiForkGeneration, boundary: VerifyBoundaryTicket)
1780 -> OptiForkTicket {
1781 OptiForkTicket {
1782 generation,
1783 boundary: Some(boundary),
1784 drain: self.stage0_stream.clone(),
1785 settled: false,
1786 }
1787 }
1788
1789 #[allow(clippy::too_many_arguments)]
1790 fn controller_ticket(
1791 &self,
1792 generation: OptiForkGeneration,
1793 boundary: VerifyBoundaryTicket,
1794 ckpt: VerifyCkpt,
1795 verify_tokens: [u32; 2],
1796 draft_prob: f32,
1797 eager_seed: Option<CudaSlice<f32>>,
1798 q_proxy: f32,
1799 scratch_len: usize,
1800 ) -> OptiControllerTicket {
1801 OptiControllerTicket {
1802 generation,
1803 boundary: Some(boundary),
1804 ckpt: Some(ckpt),
1805 verify_tokens,
1806 draft_prob,
1807 eager_seed,
1808 q_proxy,
1809 scratch_len,
1810 issued_at: std::time::Instant::now(),
1811 drain: self.stage0_stream.clone(),
1812 settled: false,
1813 }
1814 }
1815
1816 fn reserve_successor(&mut self)
1817 -> Result<OptiForkGeneration, Box<dyn std::error::Error>> {
1818 self.generations.reserve()
1819 }
1820
1821 fn successor_snapshot_mut(&mut self) -> &mut crate::cache::CacheSnapshot {
1822 &mut self.alternate_snapshot
1823 }
1824
1825 fn promote_successor_snapshot(
1826 &mut self,
1827 current_snapshot: &mut crate::cache::CacheSnapshot,
1828 generation: OptiForkGeneration,
1829 ) {
1830 std::mem::swap(current_snapshot, &mut self.alternate_snapshot);
1831 self.active_snapshot_slot = generation.slot;
1832 }
1833
1834 fn queue_actual_reconcile(
1835 &mut self,
1836 e: &Engine,
1837 snapshot: &crate::cache::CacheSnapshot,
1838 acc: &CudaSlice<u32>,
1839 optimistic_pending: u32,
1840 base: usize,
1841 ) -> Result<(), Box<dyn std::error::Error>> {
1842 let saved: Vec<i32> = (0..self.split)
1843 .map(|il| snapshot.kv_len[il].map(|v| v as i32).unwrap_or(0))
1844 .collect();
1845 // Serving keeps the caller/accept walk on the head (stage-1) device. Record the accept
1846 // decision point there and append a wait to stage 0 after its optimistic successor/TX;
1847 // the validity/reconcile kernels must never peer-read acc before it is written. The
1848 // increment-1 harness uses primary stage 0, where stream order already provides this.
1849 if self.rt.engine(0, e).ctx().ordinal() != e.ctx().ordinal() {
1850 self.rt.fence_stages_behind(&e.stream())?;
1851 }
1852 let _stage = self.rt.enter(0);
1853 let e0 = self.rt.engine(0, e);
1854 e0.htod_i32_into(&mut self.saved_lens, &saved)?;
1855 e0.spec_fork_valid(acc, optimistic_pending, &mut self.valid)?;
1856 e0.spec_fork_reconcile_kv(
1857 &self.len_ptrs,
1858 &self.saved_lens,
1859 acc,
1860 &self.valid,
1861 base,
1862 self.split,
1863 )
1864 }
1865
1866 fn finish_actual_reconcile(
1867 &mut self,
1868 e: &Engine,
1869 cache: &mut Cache,
1870 snapshot: &crate::cache::CacheSnapshot,
1871 n_acc: usize,
1872 base: usize,
1873 hit: bool,
1874 ) -> Result<(), Box<dyn std::error::Error>> {
1875 if hit {
1876 return Ok(());
1877 }
1878 let len_delta = base + n_acc;
1879 for il in 0..self.split {
1880 if let (Some(kv), Some(saved)) = (cache.kv[il].as_mut(), snapshot.kv_len[il]) {
1881 kv.len = saved + len_delta;
1882 }
1883 }
1884 {
1885 let _stage = self.rt.enter(1);
1886 let e1 = self.rt.engine(1, e);
1887 for il in self.split..self.fence[2] {
1888 if let (Some(kv), Some(saved)) = (cache.kv[il].as_mut(), snapshot.kv_len[il]) {
1889 kv.len = saved + len_delta;
1890 e1.set_i32_one(&mut kv.len_d, kv.len as i32)?;
1891 }
1892 }
1893 }
1894 self.rt.publish_to(0, &e.stream())?;
1895 Ok(())
1896 }
1897
1898 fn cancel_controller_ticket(
1899 &mut self,
1900 e: &Engine,
1901 cache: &mut Cache,
1902 scratch: &mut MtpScratch,
1903 snapshot: &crate::cache::CacheSnapshot,
1904 ticket: &mut OptiControllerTicket,
1905 ) -> Result<(), Box<dyn std::error::Error>> {
1906 {
1907 let _stage = self.rt.enter(0);
1908 let e0 = self.rt.engine(0, e);
1909 for il in 0..self.split {
1910 if let (Some(kv), Some(saved)) = (cache.kv[il].as_mut(), snapshot.kv_len[il]) {
1911 kv.len = saved;
1912 e0.set_i32_one(&mut kv.len_d, saved as i32)?;
1913 }
1914 }
1915 }
1916 scratch.set_len(e, snapshot.pos)?;
1917 ticket.settle();
1918 self.generations.retire(ticket.generation)?;
1919 OPTI_FORK_ABORT_DRAINS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
1920 OPTI_WASTED_DRAFT_TOKENS.fetch_add(2, std::sync::atomic::Ordering::Relaxed);
1921 eprintln!(
1922 "[opti-controller] tail-drain generation={} slot={}",
1923 ticket.generation.id, ticket.generation.slot,
1924 );
1925 Ok(())
1926 }
1927
1928 #[allow(clippy::too_many_arguments)]
1929 fn reconcile(
1930 &mut self,
1931 e: &Engine,
1932 cache: &mut Cache,
1933 scratch: &mut MtpScratch,
1934 snapshot: &crate::cache::CacheSnapshot,
1935 h_seed: &mut CudaSlice<f32>,
1936 fill_prev: &mut CudaSlice<f32>,
1937 generation: OptiForkGeneration,
1938 action: OptiForkAction,
1939 optimistic_pending: u32,
1940 ) -> Result<(), Box<dyn std::error::Error>> {
1941 debug_assert!(action != OptiForkAction::Abort);
1942 let miss_started = std::time::Instant::now();
1943 let keep = action == OptiForkAction::Hit;
1944 let saved: Vec<i32> = (0..self.split)
1945 .map(|il| snapshot.kv_len[il].map(|v| v as i32).unwrap_or(0))
1946 .collect();
1947 let seed = &self.seeds[generation.slot];
1948 {
1949 let _stage = self.rt.enter(0);
1950 let e0 = self.rt.engine(0, e);
1951 e0.htod_i32_into(&mut self.saved_lens, &saved)?;
1952 let forced = if keep {
1953 [1u32, optimistic_pending]
1954 } else {
1955 [0u32, optimistic_pending]
1956 };
1957 e0.htod_u32_into(&mut self.forced_acc, &forced)?;
1958 e0.spec_fork_valid(&self.forced_acc, optimistic_pending, &mut self.valid)?;
1959 e0.spec_fork_reconcile_kv(
1960 &self.len_ptrs,
1961 &self.saved_lens,
1962 &self.forced_acc,
1963 &self.valid,
1964 0,
1965 self.split,
1966 )?;
1967 for il in 0..self.split {
1968 if let Some(recur) = cache.recur[il].as_mut() {
1969 let conv = snapshot.conv[il]
1970 .as_ref()
1971 .ok_or("optipipe stage0 snapshot missing conv state")?;
1972 let ssm = snapshot.ssm[il]
1973 .as_ref()
1974 .ok_or("optipipe stage0 snapshot missing ssm state")?;
1975 e0.spec_fork_restore_f32(conv, &mut recur.conv_state, &self.valid)?;
1976 e0.spec_fork_restore_f32(ssm, &mut recur.ssm_state, &self.valid)?;
1977 }
1978 }
1979 e0.spec_fork_restore_f32(&seed.h_seed, h_seed, &self.valid)?;
1980 e0.spec_fork_restore_f32(&seed.fill_prev, fill_prev, &self.valid)?;
1981 }
1982
1983 if keep {
1984 OPTI_FORK_HITS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
1985 return Ok(());
1986 }
1987
1988 for il in 0..self.split {
1989 if let (Some(kv), Some(saved)) = (cache.kv[il].as_mut(), snapshot.kv_len[il]) {
1990 kv.len = saved;
1991 }
1992 }
1993 scratch.set_len(e, seed.scratch_len)?;
1994 // Targeted E_restart: publish only stage 0's reconcile to the caller, then bound the
1995 // forced diagnostic so the retained number is the actual miss cost, not enqueue time.
1996 let caller = e.stream();
1997 self.rt.publish_to(0, &caller)?;
1998 caller.synchronize()?;
1999 let miss_ms = miss_started.elapsed().as_secs_f64() * 1e3;
2000 eprintln!(
2001 "[opti-fork-reconcile] generation={} slot={} miss_ms={miss_ms:.3}",
2002 generation.id, generation.slot,
2003 );
2004 OPTI_FORK_MISSES.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2005 Ok(())
2006 }
2007
2008 fn retire(&mut self, generation: OptiForkGeneration)
2009 -> Result<(), Box<dyn std::error::Error>> {
2010 self.generations.retire(generation)
2011 }
2012}
2013
2014impl HybridModel {
2015 fn opti_graph_draft_step(
2016 &self,
2017 e: &Engine,
2018 mtp: &MtpHead,
2019 dctx: &mut DraftGraphCtx,
2020 scratch: &mut MtpScratch,
2021 d_vocab: usize,
2022 ) -> Result<(u32, f32), Box<dyn std::error::Error>> {
2023 dctx.graph
2024 .as_ref()
2025 .ok_or("optipipe controller requires the greedy draft graph")?
2026 .launch()?;
2027 scratch.kv.len += 1;
2028 let idx = e.dtoh_u32_one(&dctx.g_tok)?;
2029 if (idx as usize) >= d_vocab {
2030 return Err(format!(
2031 "optipipe draft argmax sentinel 0x{idx:08x} >= d_vocab {d_vocab}"
2032 )
2033 .into());
2034 }
2035 let probability = e.dtoh(&dctx.g_p)?[0];
2036 if !probability.is_finite() || !(0.0..=1.0).contains(&probability) {
2037 return Err(format!("optipipe draft probability is invalid: {probability}").into());
2038 }
2039 let token = match &mtp.d2t {
2040 Some(map) => map[idx as usize],
2041 None => idx,
2042 };
2043 if token != idx {
2044 e.set_u32_one(&mut dctx.g_tok, token)?;
2045 }
2046 Ok((token, probability))
2047 }
2048
2049 #[allow(clippy::too_many_arguments)]
2050 fn opti_controller_draft_step(
2051 &self,
2052 e: &Engine,
2053 mtp: &MtpHead,
2054 dctx: &mut DraftGraphCtx,
2055 scratch: &mut MtpScratch,
2056 d_vocab: usize,
2057 eager_state: &mut Option<(u32, CudaSlice<f32>)>,
2058 eager_pos: usize,
2059 embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
2060 ) -> Result<(u32, f32), Box<dyn std::error::Error>> {
2061 if dctx.graph.is_some() {
2062 return self.opti_graph_draft_step(e, mtp, dctx, scratch, d_vocab);
2063 }
2064 let (input_token, input_seed) = eager_state
2065 .take()
2066 .ok_or("optipipe eager continuation seed is unavailable")?;
2067 let (logits, next_seed) = self.mtp_head_forward_dev(
2068 e,
2069 mtp,
2070 input_token,
2071 &input_seed,
2072 scratch,
2073 eager_pos,
2074 embd_dev,
2075 None,
2076 )?;
2077 let token_d = e.argmax_token_device(&logits, d_vocab)?;
2078 let idx = e.dtoh_u32_one(&token_d)?;
2079 if (idx as usize) >= d_vocab {
2080 return Err(format!(
2081 "optipipe eager draft argmax sentinel 0x{idx:08x} >= d_vocab {d_vocab}"
2082 )
2083 .into());
2084 }
2085 let probability_d = e.prob_of_token_device(&logits, &token_d, d_vocab)?;
2086 let probability = e.dtoh(&probability_d)?[0];
2087 if !probability.is_finite() || !(0.0..=1.0).contains(&probability) {
2088 return Err(
2089 format!("optipipe eager draft probability is invalid: {probability}").into()
2090 );
2091 }
2092 let token = match &mtp.d2t {
2093 Some(map) => map[idx as usize],
2094 None => idx,
2095 };
2096 *eager_state = Some((token, next_seed));
2097 Ok((token, probability))
2098 }
2099
2100 /// NextN head forward for ONE draft token (§A ops 1-13, T=1).
2101 /// Inputs: `e_tok` = the token to predict FROM (last committed / previous draft); `h_seed` =
2102 /// the trunk's pre-output_norm hidden of that token (§A op 2 input). `mtp_pos` = absolute
2103 /// position of the token being predicted from. Returns (draft_logits[n_vocab] host, h_nextn dev).
2104 /// `h_nextn` (§A op 10) becomes `h_seed` for the next autoregressive draft step.
2105 /// Device-resident: returns draft logits ON DEVICE (no [n_vocab] dtoh). The greedy draft
2106 /// loop only needs argmax — paired with `argmax_token_device` this cuts the ~600KB logits
2107 /// transfer + host argmax per draft token from the K-token draft chain.
2108 #[allow(clippy::too_many_arguments)]
2109 fn mtp_head_forward_dev(
2110 &self,
2111 e: &Engine,
2112 mtp: &MtpHead,
2113 e_tok: u32,
2114 h_seed: &CudaSlice<f32>,
2115 scratch: &mut MtpScratch,
2116 mtp_pos: usize,
2117 embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
2118 // DRAFT-SIDE GRAMMAR MASK (lane/draft-mask): (packed draft-vocab allowed set, words).
2119 // Applied to the head logits BEFORE they are returned, so every consumer (argmax,
2120 // gumbel draw, p-min prob) sees the grammar-legal row. None = unmasked (pre-lane).
2121 mask: Option<(&CudaSlice<u32>, usize)>,
2122 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
2123 let cfg = &self.cfg;
2124 let n_embd = cfg.n_embd as usize;
2125 // Distilled-student geometry: the block runs at the INNER width `di` (eh_proj out /
2126 // attn / ffn); the n_embd interface (embed, norms in, carrier out, head in) is unchanged.
2127 let di = mtp.geom.as_ref().map(|g| g.d_inner).unwrap_or(n_embd);
2128 let eps = cfg.rms_eps;
2129 let pos_d = e.htod_i32(&[mtp_pos as i32])?;
2130
2131 // op A: a resident table transfers one 4B token id. The exact host-row capacity path
2132 // expands this one row on CPU and transfers n_embd f32 values instead.
2133 let e_emb = match embd_dev {
2134 Some((g, qt, rb)) => e.embed_gather_device_t(g, &[e_tok], n_embd, qt, rb)?,
2135 None => e.htod(&self.embd.gather(n_embd, &[e_tok]))?,
2136 };
2137
2138 // op 1/2: e_norm = RMSNorm(e, enorm); h_norm = RMSNorm(h_seed, hnorm)
2139 let mut e_norm = e.zeros(n_embd)?;
2140 e.rms_norm(&e_emb, mtp.enorm.float_data(), &mut e_norm, n_embd, 1, eps)?;
2141 let mut h_norm = e.zeros(n_embd)?;
2142 e.rms_norm(h_seed, mtp.hnorm.float_data(), &mut h_norm, n_embd, 1, eps)?;
2143
2144 // op 3: concat = [e_norm ; h_norm] -> [2*n_embd], e_norm in [0,n_embd), h_norm in [n_embd,2n_embd)
2145 let mut concat = e.zeros(2 * n_embd)?;
2146 e.copy_into(&mut concat, 0, &e_norm, n_embd)?;
2147 e.copy_into(&mut concat, n_embd, &h_norm, n_embd)?;
2148
2149 // op 4: inpSA = eh_proj @ concat (eh_proj [2*n_embd, n_embd]) -> [n_embd]
2150 let inp_sa = e.matmul(&mtp.eh_proj, &concat, 1)?;
2151
2152 // op 5: a_norm = RMSNorm(inpSA, attn_norm)
2153 let mut a_norm = e.zeros(di)?;
2154 e.rms_norm(&inp_sa, mtp.attn_norm.float_data(), &mut a_norm, di, 1, eps)?;
2155
2156 // op 6: attention on the scratch KV. SAME dc launcher as the graph path (bucket_max =
2157 // scratch.cap, length from the device len_d) so eager drafts match graph drafts
2158 // bit-for-bit at any t_kv (the parity gate). Host len mirrored here (the dc append
2159 // advances only the device counter).
2160 let attn_out = match (&mtp.mixer, mtp.step35.as_ref()) {
2161 // step35 MTP block: PER-LAYER geometry + a separate head-wise gate + an SWA window,
2162 // none of which the dc launcher can express (see `mtp_step35_attn`). Host-len arm.
2163 // Advances BOTH the host len and the device counter itself (unlike the dc arm,
2164 // whose host-side mirror the caller does).
2165 (Mixer::Full(fa), Some(g)) => self.mtp_step35_attn(e, fa, g, &a_norm, &pos_d, scratch)?,
2166 (Mixer::Full(fa), None) => {
2167 let out =
2168 self.mtp_full_attn_dc(e, fa, &a_norm, &pos_d, scratch, mtp.geom.as_ref())?;
2169 scratch.kv.len += 1;
2170 out
2171 }
2172 (Mixer::Linear(_), _) => {
2173 panic!("MTP block is full-attn in qwen35; linear MTP not supported")
2174 }
2175 (Mixer::Mla(_), _) => crate::hybrid::mla_forward_unimplemented(),
2176 };
2177
2178 // op 7: x1 = inpSA + attn_out
2179 let mut x1 = e.zeros(di)?;
2180 e.add(&inp_sa, &attn_out, &mut x1, di)?;
2181
2182 // op 8: z = RMSNorm(x1, post_attn_norm) (pre-FFN norm)
2183 let mut z = e.zeros(di)?;
2184 e.rms_norm(&x1, mtp.post_attn_norm.float_data(), &mut z, di, 1, eps)?;
2185
2186 // op 9: FFN (Dense or MoE) — same as the trunk decode FFN
2187 let ffn_out = match &mtp.ffn {
2188 crate::hybrid::Ffn::Dense {
2189 ffn_gate,
2190 ffn_up,
2191 ffn_down,
2192 } => {
2193 let n_ff = ffn_gate.out_features();
2194 let (gate, up) = if e.uses_q8_1_fast(ffn_gate) && e.uses_q8_1_fast(ffn_up) {
2195 let (zq, zd) = e.quantize_q8_1(&z, 1, di)?;
2196 (
2197 e.matmul_pre(ffn_gate, &zq, &zd, &z, 1)?,
2198 e.matmul_pre(ffn_up, &zq, &zd, &z, 1)?,
2199 )
2200 } else {
2201 (e.matmul(ffn_gate, &z, 1)?, e.matmul(ffn_up, &z, 1)?)
2202 };
2203 let mut act = e.zeros(n_ff)?;
2204 // step35: a DENSE FFN reads the per-layer SHEXP clamp (upstream's one `build_ffn`
2205 // serves the dense MLP and the shared expert off `swiglu_clamp_shexp` —
2206 // llama-graph.cpp:1751), resolved for the MTP block's OWN index. Every other arch
2207 // passes None, which is `ffn_act`'s dispatch verbatim.
2208 Self::ffn_act_lim(e, &self.cfg, &gate, &up, 1.0, 1.0,
2209 mtp.step35.as_ref().and_then(|s| s.clamp_shexp),
2210 &mut act, n_ff)?;
2211 e.matmul(ffn_down, &act, 1)?
2212 }
2213 // MTP head is a distinct block — key its experts under a separate layer index (u16::MAX)
2214 // so they never alias trunk layer 0's cache keys.
2215 crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il(e, m, &z, 1, u16::MAX)?,
2216 };
2217
2218 // op 10: h_nextn = x1 + ffn_out (at di)
2219 let mut h_inner = e.zeros(di)?;
2220 e.add(&x1, &ffn_out, &mut h_inner, di)?;
2221
2222 // op 10.5 (student): up-project the inner hidden back to n_embd — training semantics:
2223 // the chain carrier AND the head input are out_up(h_inner) (pre-final-norm).
2224 let h_nextn = match mtp.geom.as_ref() {
2225 Some(g) => e.matmul(&g.out_up, &h_inner, 1)?,
2226 None => h_inner,
2227 };
2228
2229 // op 11: final = RMSNorm(h_nextn, shared_head_norm OR output_norm)
2230 let final_norm = mtp.shared_head_norm.as_ref().unwrap_or(&self.output_norm);
2231 let mut final_h = e.zeros(n_embd)?;
2232 e.rms_norm(
2233 &h_nextn,
2234 final_norm.float_data(),
2235 &mut final_h,
2236 n_embd,
2237 1,
2238 eps,
2239 )?;
2240
2241 // op 12: draft_logits = (shared_head_head OR output) @ final — stays ON DEVICE.
2242 let head = mtp.shared_head_head.as_ref().unwrap_or(&self.output);
2243 let mut logits = e.matmul(head, &final_h, 1)?;
2244 // op 12b (lane/draft-mask): grammar mask over the DRAFT vocab, applied here so the
2245 // caller's argmax / gumbel draw / p-min prob all read the grammar-legal row.
2246 if let Some((mask_d, mw)) = mask {
2247 let d_vocab = head.out_features();
2248 e.mask_logits_col(&mut logits, mask_d, 0, d_vocab, mw)?;
2249 }
2250 // Chain recurrence hand-over: pre-norm h_nextn (default) or post-norm final_h
2251 // (MEMRA_SPEC_HPOST — llama.cpp #24025's t_h_nextn is taken AFTER the head norm).
2252 Ok((logits, if spec_hpost() { final_h } else { h_nextn }))
2253 }
2254
2255 /// step35 MTP-block attention, T=1, on the scratch KV — the EAGER-ONLY twin of
2256 /// `mtp_full_attn_dc`. Three things force a separate arm rather than a geometry parameter on
2257 /// the dc path, and all three are properties of this arch's MTP block:
2258 ///
2259 /// 1. **The SWA window.** Block 45 is an SWA-type block (`sliding_window_pattern[45]=true`,
2260 /// window 512). Windowed decode in memra is a token-aligned VIEW OFFSET into the quantized
2261 /// cache (the gemma4 R6 / `step35_decode_attn` pattern: keys carry absolute rope and the
2262 /// mask is purely positional, so one query at `len-1` attending the last `win` rows IS the
2263 /// windowed result). `fa_decode_dc` takes the key count from a DEVICE counter and always
2264 /// starts at row 0 — it cannot express a nonzero offset, so a windowed dc arm would need a
2265 /// new kernel. That is deliberately not built here: see the CUDA-graph note below.
2266 /// 2. **Per-layer head count.** 96 q heads over 8 KV (GQA 12) at this block, vs the trunk's 64
2267 /// on its full-attn layers. The trunk cfg's `n_head` scalar is the MAX over layers, and the
2268 /// trunk ARTIFACT's per-layer arrays stop at index 44 — so the count must come from the
2269 /// resolved `Step35MtpGeom`, never from `cfg`.
2270 /// 3. **The separate head-wise gate.** `blk.45.attn_gate.weight [n_embd, 96]` produces one
2271 /// sigmoid scalar per head (broadcast over head_dim) — `attn_head_gate`, not the qwen35
2272 /// fused-into-wq `q_gate_split` form the dc arm handles.
2273 ///
2274 /// WHY EAGER-ONLY IS NOT A GAP TODAY: `graph_draft` requires `trunk_dense`, and this SKU's
2275 /// trunk is a 288-expert MoE, so the graph draft is already off for every step35 model — the
2276 /// eager chain IS the served path. `mtp_head_forward_cap` refuses step35 explicitly rather
2277 /// than silently capturing a window-less (wrong past `win` draft rows) graph.
2278 ///
2279 /// Unlike the dc arm this advances BOTH the host `kv.len` and the device counter, so the
2280 /// caller must not mirror.
2281 fn mtp_step35_attn(
2282 &self,
2283 e: &Engine,
2284 fa: &FullAttnLayer,
2285 g: &crate::hybrid::Step35MtpGeom,
2286 h: &CudaSlice<f32>,
2287 pos_d: &CudaSlice<i32>,
2288 scratch: &mut MtpScratch,
2289 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2290 let (nh, nkv, hd) = (g.n_head, g.n_head_kv, self.cfg.head_dim_k as usize);
2291 let eps = self.cfg.rms_eps;
2292 let scale = 1.0 / (hd as f32).sqrt(); // step35.cpp:255 kq_scale
2293 let n_embd = self.cfg.n_embd as usize;
2294 let gw = fa.attn_gate.as_ref()
2295 .ok_or("step35 MTP block is missing attn_gate.weight (head-wise attention gate)")?;
2296
2297 let (q0, k0, v0, gt) = if e.uses_q8_1_fast(&fa.wq) && e.uses_q8_1_fast(&fa.wk)
2298 && e.uses_q8_1_fast(&fa.wv) && e.uses_q8_1_fast(gw)
2299 {
2300 let (hq, hdq) = e.quantize_q8_1(h, 1, n_embd)?;
2301 let (a, b, c) = match e.matmul_q8_fused3(&fa.wq, &fa.wk, &fa.wv, &hq, &hdq)? {
2302 Some(t3) => t3,
2303 None => (e.matmul_pre(&fa.wq, &hq, &hdq, h, 1)?,
2304 e.matmul_pre(&fa.wk, &hq, &hdq, h, 1)?,
2305 e.matmul_pre(&fa.wv, &hq, &hdq, h, 1)?),
2306 };
2307 (a, b, c, e.matmul_pre(gw, &hq, &hdq, h, 1)?)
2308 } else {
2309 (e.matmul(&fa.wq, h, 1)?, e.matmul(&fa.wk, h, 1)?,
2310 e.matmul(&fa.wv, h, 1)?, e.matmul(gw, h, 1)?)
2311 };
2312
2313 let mut q = e.uninit(nh * hd)?;
2314 e.rms_norm(&q0, fa.q_norm.float_data(), &mut q, hd, nh, eps)?;
2315 let mut k = e.uninit(nkv * hd)?;
2316 e.rms_norm(&k0, fa.k_norm.float_data(), &mut k, hd, nkv, eps)?;
2317 // `rope_freqs.weight` (llama3 factors) applies to the FULL-attn layers ONLY; SWA passes
2318 // null (llama-hparams / step35.cpp). Block 45 is SWA, so `ff` is None there — but read
2319 // the resolved flag, not the constant, so an all-full sibling stays correct.
2320 let ff = if g.swa { None } else {
2321 self.step35_aux.as_ref().and_then(|a| a.rope_freqs(e))
2322 };
2323 #[cfg(debug_assertions)]
2324 if let Some(ff) = ff {
2325 crate::debug_assert_tensor_stream_device(ff, &e.stream(),
2326 "mtp_step35_attn.rope_freqs");
2327 }
2328 e.rope_neox2(&mut q, &mut k, pos_d, hd, g.n_rot, nh, nkv, 1, g.rope_base, 1.0, ff)?;
2329
2330 // Append at the HOST slot, then re-stamp the device counter: the eager chain has the
2331 // length on the host anyway, and the windowed view below needs it there to compute the
2332 // offset. The device counter is kept in lockstep so `mtp_kv_fill`'s `set_i32_one` and any
2333 // dc-family consumer of this scratch still agree.
2334 let kv = &mut scratch.kv;
2335 assert!(kv.len < scratch.cap, "step35 MTP scratch overflow ({} >= {})", kv.len, scratch.cap);
2336 let next_len = kv.len + 1;
2337 let (off, t_kv) = if g.swa && next_len > g.window {
2338 (next_len - g.window, g.window)
2339 } else {
2340 (0, next_len)
2341 };
2342 let write_row = e.prepare_kv_append(kv, off & !31usize, 1)?;
2343 e.append_kv_quantized(&k, &v0, &mut kv.k, &mut kv.v, write_row,
2344 kv.kv_dim_k, kv.kv_dim_v, kv.k_tok_bytes, kv.v_tok_bytes, false)?;
2345 kv.len = next_len;
2346 e.set_i32_one(&mut kv.len_d, kv.len as i32)?;
2347 // SWA view offset (see note 1). The draft chain is short (k+2 rows), but the scratch is
2348 // PERSISTENT across rounds — `mtp_kv_fill` leaves one row per committed token behind, so
2349 // `kv.len` tracks absolute position and crosses 512 in any real generation. The window is
2350 // therefore live, not theoretical.
2351 let physical = kv.physical_rows(off, off + t_kv)?;
2352 let k_view = e.view_u8_range(&kv.k, physical.start * kv.k_tok_bytes,
2353 physical.end * kv.k_tok_bytes);
2354 let v_view = e.view_u8_range(&kv.v, physical.start * kv.v_tok_bytes,
2355 physical.end * kv.v_tok_bytes);
2356 let mut attn = e.uninit(nh * hd)?;
2357 e.fa_decode_kvmod(&q, &k_view, &v_view, &mut attn, hd, nh, nkv, t_kv, scale,
2358 kv.k_tok_bytes, kv.v_tok_bytes, false)?;
2359
2360 let mut ag = e.uninit(nh * hd)?;
2361 e.attn_head_gate(&attn, >, &mut ag, None, hd, nh, 1)?;
2362 Ok(e.matmul(&fa.wo, &ag, 1)?)
2363 }
2364
2365 /// MTP-block full attention, T=1, on the scratch KV (BOTH draft paths — eager and graph):
2366 /// the scratch write slot and the attention bound come from `scratch.kv.len_d` (device i32[1])
2367 /// so the launch args are FIXED across draft steps — ONE captured graph serves the whole
2368 /// chain, and replays keep seeing KV growth through the device counter (no recapture).
2369 /// Geometry contract: n_splits is sized from `scratch.cap` (the persistent capacity); splits
2370 /// whose key range lies beyond the device t_kv exit empty and the shared combine skips them
2371 /// (fa_decode_dc bit-correct-for-any-t_kv<=bucket_max contract). The eager path uses the SAME
2372 /// launcher with the SAME bucket_max -> identical dispatch -> bit-identical draft tokens (the
2373 /// graph-vs-eager parity gate). Host len is NOT advanced here (graph contract); callers mirror.
2374 fn mtp_full_attn_dc(
2375 &self,
2376 e: &Engine,
2377 fa: &FullAttnLayer,
2378 h: &CudaSlice<f32>,
2379 pos_d: &CudaSlice<i32>,
2380 scratch: &mut MtpScratch,
2381 geom: Option<&crate::hybrid::DraftGeom>,
2382 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2383 let cfg = &self.cfg;
2384 let mtp_il = cfg.n_layer.saturating_sub(cfg.nextn_predict_layers);
2385 let geometry = cfg.full_attention_geometry_at(mtp_il);
2386 let n_head = geom.map(|g| g.n_head).unwrap_or(geometry.n_head as usize);
2387 let n_head_kv = geom.map(|g| g.n_head_kv).unwrap_or(geometry.n_head_kv as usize);
2388 let head_dim = geometry.head_dim_k as usize;
2389 let eps = cfg.rms_eps;
2390 let scale = geometry.attention_scale();
2391 let n_embd = geom.map(|g| g.d_inner).unwrap_or(cfg.n_embd as usize);
2392 let bucket_max = scratch.cap; // < 96 guaranteed by the graph_draft eligibility gate
2393
2394 let (qf, mut k, v) =
2395 if e.uses_q8_1_fast(&fa.wq) && e.uses_q8_1_fast(&fa.wk) && e.uses_q8_1_fast(&fa.wv) {
2396 let (hq, hd) = e.quantize_q8_1(h, 1, n_embd)?;
2397 (
2398 e.matmul_pre(&fa.wq, &hq, &hd, h, 1)?,
2399 e.matmul_pre(&fa.wk, &hq, &hd, h, 1)?,
2400 e.matmul_pre(&fa.wv, &hq, &hd, h, 1)?,
2401 )
2402 } else {
2403 (
2404 e.matmul(&fa.wq, h, 1)?,
2405 e.matmul(&fa.wk, h, 1)?,
2406 e.matmul(&fa.wv, h, 1)?,
2407 )
2408 };
2409 // M3/Hy3 have no attention output gate — wq out is exactly q; skip the split.
2410 let gated = geometry.attention_gate
2411 == memra_gguf::config::AttentionGateKind::FusedQ;
2412 let (mut q, gate) = if gated {
2413 let mut q = e.zeros(n_head * head_dim)?;
2414 let mut gate = e.zeros(n_head * head_dim)?;
2415 e.q_gate_split(&qf, &mut q, &mut gate, head_dim, n_head, 1)?;
2416 (q, Some(gate))
2417 } else {
2418 (qf, None)
2419 };
2420
2421 let mut qn = e.zeros(n_head * head_dim)?;
2422 e.rms_norm(&q, fa.q_norm.float_data(), &mut qn, head_dim, n_head, eps)?;
2423 q = qn;
2424 let mut kn = e.zeros(n_head_kv * head_dim)?;
2425 e.rms_norm(
2426 &k,
2427 fa.k_norm.float_data(),
2428 &mut kn,
2429 head_dim,
2430 n_head_kv,
2431 eps,
2432 )?;
2433 k = kn;
2434 let rope_dims = geometry.n_rot as usize;
2435 e.rope_neox(
2436 &mut q,
2437 pos_d,
2438 head_dim,
2439 rope_dims,
2440 n_head,
2441 1,
2442 geometry.rope_base,
2443 1.0,
2444 )?;
2445 e.rope_neox(
2446 &mut k,
2447 pos_d,
2448 head_dim,
2449 rope_dims,
2450 n_head_kv,
2451 1,
2452 geometry.rope_base,
2453 1.0,
2454 )?;
2455
2456 let kv = &mut scratch.kv;
2457 // append at the DEVICE slot (kv.len_d == old len), then advance the counter in-graph.
2458 e.append_kv_quantized_dc(
2459 &k,
2460 &v,
2461 &mut kv.k,
2462 &mut kv.v,
2463 &kv.len_d,
2464 kv.kv_dim_k,
2465 kv.kv_dim_v,
2466 kv.k_tok_bytes,
2467 kv.v_tok_bytes,
2468 false,
2469 )?;
2470 e.inc_seqlen(&mut kv.len_d)?;
2471 // full-buffer views (any in-round t_kv stays in range on replay); the kernel bounds the
2472 // key range from the device counter.
2473 let k_view = e.view_u8(&kv.k, kv.k.len());
2474 let v_view = e.view_u8(&kv.v, kv.v.len());
2475 let (ktb, vtb) = (kv.k_tok_bytes, kv.v_tok_bytes);
2476 let mut attn = e.zeros(n_head * head_dim)?;
2477 e.fa_decode_dc(
2478 &q, &k_view, &v_view, &mut attn, head_dim, n_head, n_head_kv, &kv.len_d, bucket_max,
2479 scale, ktb, vtb, false,
2480 )?;
2481
2482 let attn_g = match &gate {
2483 Some(gate) => {
2484 let mut gsig = e.zeros(n_head * head_dim)?;
2485 e.sigmoid(gate, &mut gsig, n_head * head_dim)?;
2486 let mut ag = e.zeros(n_head * head_dim)?;
2487 e.mul(&attn, &gsig, &mut ag, n_head * head_dim)?;
2488 ag
2489 }
2490 None => attn,
2491 };
2492 Ok(e.matmul(&fa.wo, &attn_g, 1)?)
2493 }
2494
2495 /// PERSISTENT-DRAFT-KV fill (the reference engine's "mtp_update" analogue): compute the MTP
2496 /// block's K/V for `tokens` (committed tokens at positions pos0..pos0+T) from their EXACT
2497 /// trunk hiddens `h` ([T, n_embd] token-major, pre-output_norm) and append at slots pos0.. of
2498 /// the scratch KV. K/V-ONLY — ops A/1-5 plus the K-side of op 6 (wk/wv + k_norm + rope +
2499 /// quantized append); no wq/attention/FFN/lm_head, so per-token cost ~= eh_proj + wk/wv (a
2500 /// small fraction of one trunk layer), T-batched. Rope follows the chain convention
2501 /// rope(token@p) = p+1. Runs at round boundaries OUTSIDE the captured graph in BOTH draft
2502 /// modes -> draft parity by construction. Caller must have scratch.kv.len == pos0.
2503 #[allow(clippy::too_many_arguments)]
2504 fn mtp_kv_fill(
2505 &self,
2506 e: &Engine,
2507 mtp: &MtpHead,
2508 tokens: &[u32],
2509 h: &CudaSlice<f32>,
2510 pos0: usize,
2511 scratch: &mut MtpScratch,
2512 embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
2513 ) -> Result<(), Box<dyn std::error::Error>> {
2514 let cfg = &self.cfg;
2515 let n_embd = cfg.n_embd as usize;
2516 let eps = cfg.rms_eps;
2517 let t = tokens.len();
2518 assert_eq!(scratch.kv.len, pos0, "mtp_kv_fill: append slot mismatch");
2519 assert!(pos0 + t <= scratch.cap, "mtp_kv_fill: scratch overflow");
2520 let Mixer::Full(fa) = &mtp.mixer else {
2521 panic!("MTP block is full-attn in qwen35; linear MTP not supported")
2522 };
2523 let pos_vec: Vec<i32> = (0..t).map(|i| (pos0 + i + 1) as i32).collect();
2524 let pos_d = e.htod_i32(&pos_vec)?;
2525
2526 // ops A/1/2: embed + the two input norms, T-wide.
2527 let e_emb = match embd_dev {
2528 Some((g, qt, rb)) => e.embed_gather_device_t(g, tokens, n_embd, qt, rb)?,
2529 None => e.htod(&self.embd.gather(n_embd, tokens))?,
2530 };
2531 let mut e_norm = e.zeros(t * n_embd)?;
2532 e.rms_norm(&e_emb, mtp.enorm.float_data(), &mut e_norm, n_embd, t, eps)?;
2533 let mut h_norm = e.zeros(t * n_embd)?;
2534 e.rms_norm(h, mtp.hnorm.float_data(), &mut h_norm, n_embd, t, eps)?;
2535
2536 // op 3: per-row [e_norm ; h_norm] concat, token-major [T, 2*n_embd].
2537 let mut concat = e.zeros(t * 2 * n_embd)?;
2538 for i in 0..t {
2539 e.copy_view_into(
2540 &mut concat,
2541 i * 2 * n_embd,
2542 &e_norm.slice(i * n_embd..(i + 1) * n_embd),
2543 n_embd,
2544 )?;
2545 e.copy_view_into(
2546 &mut concat,
2547 i * 2 * n_embd + n_embd,
2548 &h_norm.slice(i * n_embd..(i + 1) * n_embd),
2549 n_embd,
2550 )?;
2551 }
2552
2553 // ops 4/5: eh_proj + attn_norm, T-wide (at the student inner width when geom is set).
2554 let di = mtp.geom.as_ref().map(|g| g.d_inner).unwrap_or(n_embd);
2555 let inp_sa = e.matmul(&mtp.eh_proj, &concat, t)?;
2556 let mut a_norm = e.zeros(t * di)?;
2557 e.rms_norm(&inp_sa, mtp.attn_norm.float_data(), &mut a_norm, di, t, eps)?;
2558
2559 // op 6 (K/V half): wk/wv + k_norm + rope + per-row quantized append. No wq/attention —
2560 // the fill only has to leave correct K/V rows behind for later chains to attend over.
2561 let n_head_kv = mtp
2562 .geom
2563 .as_ref()
2564 .map(|g| g.n_head_kv)
2565 .or(mtp.step35.as_ref().map(|s| s.n_head_kv))
2566 .unwrap_or_else(|| {
2567 let mtp_il = cfg.n_layer.saturating_sub(cfg.nextn_predict_layers);
2568 cfg.full_attention_geometry_at(mtp_il).n_head_kv as usize
2569 });
2570 let mtp_il = cfg.n_layer.saturating_sub(cfg.nextn_predict_layers);
2571 let geometry = cfg.full_attention_geometry_at(mtp_il);
2572 let head_dim = geometry.head_dim_k as usize;
2573 let mut k = e.matmul(&fa.wk, &a_norm, t)?;
2574 let v = e.matmul(&fa.wv, &a_norm, t)?;
2575 let mut kn = e.zeros(t * n_head_kv * head_dim)?;
2576 e.rms_norm(
2577 &k,
2578 fa.k_norm.float_data(),
2579 &mut kn,
2580 head_dim,
2581 n_head_kv * t,
2582 eps,
2583 )?;
2584 k = kn;
2585 // step35: rotary width AND base are per-layer, and the MTP block's values come from the
2586 // resolved `Step35MtpGeom` — NOT from `cfg.rope_dim_count`/`cfg.rope_freq_base`, which
2587 // carry the arch defaults (128 / 5e6, i.e. the FULL-attn layers' base). Getting this wrong
2588 // writes K rows the attention arm then re-derives at a different theta: correct-looking
2589 // output with dead acceptance, invisible to the exactness gates.
2590 let (rope_dims, rope_base, ff) = match mtp.step35.as_ref() {
2591 Some(s) => (
2592 s.n_rot,
2593 s.rope_base,
2594 if s.swa { None } else {
2595 self.step35_aux.as_ref().and_then(|a| a.rope_freqs(e))
2596 },
2597 ),
2598 None => (geometry.n_rot as usize, geometry.rope_base, None),
2599 };
2600 #[cfg(debug_assertions)]
2601 if let Some(ff) = ff {
2602 crate::debug_assert_tensor_stream_device(ff, &e.stream(),
2603 "mtp_kv_fill.rope_freqs");
2604 }
2605 match ff {
2606 Some(f) => e.rope_neox_ff(&mut k, &pos_d, head_dim, rope_dims, n_head_kv, t,
2607 rope_base, 1.0, f)?,
2608 None => e.rope_neox(&mut k, &pos_d, head_dim, rope_dims, n_head_kv, t,
2609 rope_base, 1.0)?,
2610 }
2611
2612 let kv = &mut scratch.kv;
2613 // Match the trunk prime contract: a chunk may need the aligned window immediately before
2614 // its first row, so preserve that prefix when the physical tail rebases at wrap.
2615 let retain_from = kv
2616 .ring
2617 .as_ref()
2618 .map(|ring| pos0.saturating_sub(ring.window() - 1) & !31usize)
2619 .unwrap_or(0);
2620 let write_row = e.prepare_kv_append(kv, retain_from, t)?;
2621 for i in 0..t {
2622 let k_row = k.slice(i * kv.kv_dim_k..(i + 1) * kv.kv_dim_k);
2623 let v_row = v.slice(i * kv.kv_dim_v..(i + 1) * kv.kv_dim_v);
2624 e.append_kv_quantized_view(
2625 &k_row,
2626 &v_row,
2627 &mut kv.k,
2628 &mut kv.v,
2629 write_row + i,
2630 kv.kv_dim_k,
2631 kv.kv_dim_v,
2632 kv.k_tok_bytes,
2633 kv.v_tok_bytes,
2634 false,
2635 )?;
2636 }
2637 kv.len = pos0 + t;
2638 e.set_i32_one(&mut kv.len_d, kv.len as i32)?;
2639 Ok(())
2640 }
2641
2642 /// CAPTURE body for the GRAPH DRAFT (stage 2 of graph-grade spec): ONE MTP head forward with
2643 /// every varying input device-resident —
2644 /// - token id from the persistent `tok_d` (the previous replay's in-graph argmax wrote it,
2645 /// so the chain feeds itself; the host reads the same 4 bytes for the draft list),
2646 /// - h_seed from the persistent `h_seed_d` (h_nextn is copied BACK into it at the end),
2647 /// - rope pos from the persistent `pos_d` counter (inc'd in-graph),
2648 /// - scratch KV slot/bound from `scratch.kv.len_d` (see mtp_full_attn_dc).
2649 /// The p-min confidence lands in the persistent `p_d` iff `with_prob` (env is fixed per run).
2650 /// Same kernels, same dispatch as the eager mtp_head_forward_dev chain -> same draft tokens
2651 /// (exactness never depends on drafts — the verify arbitrates — but acceptance parity does).
2652 /// `with_head=false` captures the HEAD-LESS twin for the pseudo-seed replay (2026-07-03):
2653 /// the pseudo pass only needs h_nextn (op 10) + the scratch append — the lm_head read
2654 /// (~1.06ms q6_K on the 9B), argmax and prob are dead weight there. h_nextn's inputs are
2655 /// untouched, so the seed value is identical; round-start resets overwrite tok_d/p_d anyway.
2656 /// `sampled_cap` = Some((ctr_d, perturb_d, q_out_d, seed, temp)) captures the SAMPLED twin
2657 /// (step 3 of the sampled-spec arc): head logits are retained in the persistent `q_out_d`
2658 /// (host D2Ds them to the round's q slot after each replay), the DEVICE event counter is
2659 /// bumped in-graph, and the argmax reads GUMBEL-PERTURBED logits — one categorical draw per
2660 /// replay, bit-identical to the eager arm's gumbel_perturb at the same (seed, sctr, temp).
2661 /// seed/temp are capture-time constants (fixed per generate call, like p_min).
2662 #[allow(clippy::too_many_arguments)]
2663 fn mtp_head_forward_cap(
2664 &self,
2665 e: &Engine,
2666 mtp: &MtpHead,
2667 tok_d: &mut CudaSlice<u32>,
2668 pos_d: &mut CudaSlice<i32>,
2669 h_seed_d: &mut CudaSlice<f32>,
2670 p_d: &mut CudaSlice<f32>,
2671 scratch: &mut MtpScratch,
2672 with_prob: bool,
2673 with_head: bool,
2674 embd_gpu: &CudaSlice<u8>,
2675 embd_qt: i32,
2676 embd_rb: usize,
2677 d_vocab: usize,
2678 sampled_cap: Option<(
2679 &mut CudaSlice<u32>,
2680 &mut CudaSlice<f32>,
2681 &mut CudaSlice<f32>,
2682 u64,
2683 f32,
2684 )>,
2685 stream_pack: Option<(&mut CudaSlice<u32>, usize, Option<&CudaSlice<u32>>)>,
2686 // DRAFT-SIDE GRAMMAR MASK (lane/draft-mask): (packed draft-vocab allowed-set buffer,
2687 // word count). Captured as ONE mask_logits_f32 node between the head matmul and the
2688 // in-graph argmax; the buffer address is baked, its CONTENTS are re-uploaded by the
2689 // host before every replay (the decode.rs graph-mask pattern). All-ones contents = a
2690 // no-op ban, so a position the grammar cannot constrain costs one pass over the row.
2691 mask_cap: Option<(&CudaSlice<u32>, usize)>,
2692 ) -> Result<(), Box<dyn std::error::Error>> {
2693 let cfg = &self.cfg;
2694 let n_embd = cfg.n_embd as usize;
2695 // step35 REFUSAL (deliberate, named): the graph body's attention is `mtp_full_attn_dc`,
2696 // whose device-counter key bound always starts at row 0 — it cannot express this block's
2697 // SWA view offset, so a captured chain would silently attend OUTSIDE the window once the
2698 // persistent scratch passes 512 rows. Nothing is lost today: `graph_draft` also requires
2699 // `trunk_dense`, and step35's trunk is a 288-expert MoE, so the eager chain
2700 // (`mtp_head_forward_dev` -> `mtp_step35_attn`) is the served path. Returning Err (not a
2701 // panic) is what the two capture sites and the round-stream capture already handle by
2702 // degrading to eager / stream-off.
2703 if mtp.step35.is_some() {
2704 return Err("step35 has no captured draft chain (fa_decode_dc cannot express the MTP \
2705 block's SWA view offset; same root cause as the dc decode refusal) — the \
2706 eager draft chain serves this arch".into());
2707 }
2708 // student inner width (see mtp_head_forward_dev) — interface dims stay n_embd.
2709 let di = mtp.geom.as_ref().map(|g| g.d_inner).unwrap_or(n_embd);
2710 let eps = cfg.rms_eps;
2711 let e_emb = e.embed_gather_device(embd_gpu, tok_d, n_embd, embd_qt, embd_rb)?;
2712 let mut e_norm = e.zeros(n_embd)?;
2713 e.rms_norm(&e_emb, mtp.enorm.float_data(), &mut e_norm, n_embd, 1, eps)?;
2714 let mut h_norm = e.zeros(n_embd)?;
2715 e.rms_norm(
2716 &*h_seed_d,
2717 mtp.hnorm.float_data(),
2718 &mut h_norm,
2719 n_embd,
2720 1,
2721 eps,
2722 )?;
2723 let mut concat = e.zeros(2 * n_embd)?;
2724 e.copy_into(&mut concat, 0, &e_norm, n_embd)?;
2725 e.copy_into(&mut concat, n_embd, &h_norm, n_embd)?;
2726 let inp_sa = e.matmul(&mtp.eh_proj, &concat, 1)?;
2727 let mut a_norm = e.zeros(di)?;
2728 e.rms_norm(&inp_sa, mtp.attn_norm.float_data(), &mut a_norm, di, 1, eps)?;
2729 let attn_out = match &mtp.mixer {
2730 Mixer::Full(fa) => {
2731 self.mtp_full_attn_dc(e, fa, &a_norm, pos_d, scratch, mtp.geom.as_ref())?
2732 }
2733 Mixer::Linear(_) => {
2734 panic!("MTP block is full-attn in qwen35; linear MTP not supported")
2735 }
2736 Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
2737 };
2738 let mut x1 = e.zeros(di)?;
2739 e.add(&inp_sa, &attn_out, &mut x1, di)?;
2740 let mut z = e.zeros(di)?;
2741 e.rms_norm(&x1, mtp.post_attn_norm.float_data(), &mut z, di, 1, eps)?;
2742 let ffn_out = match &mtp.ffn {
2743 crate::hybrid::Ffn::Dense {
2744 ffn_gate,
2745 ffn_up,
2746 ffn_down,
2747 } => {
2748 let n_ff = ffn_gate.out_features();
2749 let (gate, up) = if e.uses_q8_1_fast(ffn_gate) && e.uses_q8_1_fast(ffn_up) {
2750 let (zq, zd) = e.quantize_q8_1(&z, 1, di)?;
2751 (
2752 e.matmul_pre(ffn_gate, &zq, &zd, &z, 1)?,
2753 e.matmul_pre(ffn_up, &zq, &zd, &z, 1)?,
2754 )
2755 } else {
2756 (e.matmul(ffn_gate, &z, 1)?, e.matmul(ffn_up, &z, 1)?)
2757 };
2758 let mut act = e.zeros(n_ff)?;
2759 Self::ffn_act(e, &self.cfg, &gate, &up, &mut act, n_ff)?;
2760 e.matmul(ffn_down, &act, 1)?
2761 }
2762 // ROUND-STREAM: the 35B NextN block carries a MoE FFN. With RESIDENT experts the
2763 // dev path is pure device launches (device top-k + rows kernels, ZERO-DtoH by
2764 // design) — capture-legal. Non-resident (SLRU-lock) stays rejected: the capture
2765 // error arm degrades the caller to eager/stream-off.
2766 crate::hybrid::Ffn::Moe(m) if m.dev_exps.is_some() => {
2767 self.moe_ffn_il(e, m, &z, 1, u16::MAX)?
2768 }
2769 crate::hybrid::Ffn::Moe(_) => {
2770 return Err("graph draft requires a Dense (or resident-MoE) MTP FFN".into())
2771 }
2772 };
2773 let mut h_inner = e.zeros(di)?;
2774 e.add(&x1, &ffn_out, &mut h_inner, di)?;
2775 // student: up-project back to n_embd (carrier + head input; see mtp_head_forward_dev).
2776 let h_nextn = match mtp.geom.as_ref() {
2777 Some(g) => e.matmul(&g.out_up, &h_inner, 1)?,
2778 None => h_inner,
2779 };
2780 // MEMRA_SPEC_HPOST needs final_h even head-less (it IS the next seed under that convention).
2781 let final_h = if with_head || spec_hpost() {
2782 let final_norm = mtp.shared_head_norm.as_ref().unwrap_or(&self.output_norm);
2783 let mut fh = e.zeros(n_embd)?;
2784 e.rms_norm(&h_nextn, final_norm.float_data(), &mut fh, n_embd, 1, eps)?;
2785 Some(fh)
2786 } else {
2787 None
2788 };
2789 if with_head {
2790 let head = mtp.shared_head_head.as_ref().unwrap_or(&self.output);
2791 let mut logits = e.matmul(head, final_h.as_ref().unwrap(), 1)?;
2792 // DRAFT-SIDE GRAMMAR MASK: ban the grammar-illegal draft ids IN the captured chain,
2793 // before the argmax — proposals become legal by construction. Contents-only
2794 // per-replay upload keeps the capture valid.
2795 if let Some((mask_d, mw)) = mask_cap {
2796 e.mask_logits_col(&mut logits, mask_d, 0, d_vocab, mw)?;
2797 }
2798 if let Some((ctr_d, perturb_d, q_out_d, seed, temp)) = sampled_cap {
2799 // SAMPLED chain: retain q (raw head logits -> persistent q_out_d; the matmul's
2800 // own buffer is pool-recycled after the capture body returns, so it can't be the
2801 // retention target), bump the device event counter, gumbel-perturb reading it,
2802 // and argmax the PERTURBED logits into tok_d — the in-graph categorical draw.
2803 e.copy_into(q_out_d, 0, &logits, d_vocab)?;
2804 e.sctr_inc(ctr_d)?;
2805 e.gumbel_perturb_ctr(&logits, perturb_d, d_vocab, seed, ctr_d, temp)?;
2806 e.argmax_token_device_into(perturb_d, tok_d, d_vocab)?;
2807 // p-min prob = the head's RAW softmax confidence in the SAMPLED pick — same
2808 // semantics as the eager sampled arm's prob_of_token_device(dl_d, tok_d).
2809 if with_prob {
2810 e.prob_of_token_device_into(&logits, tok_d, p_d, d_vocab)?;
2811 }
2812 } else {
2813 // draft token -> persistent tok_d (next replay's embed reads it; host reads the 4 bytes).
2814 e.argmax_token_device_into(&logits, tok_d, d_vocab)?;
2815 // p-min under a draft mask reads the MASKED row: confidence relative to the
2816 // grammar-LEGAL alternatives (illegal ids leave the softmax denominator), which
2817 // is the right semantics for "does the drafter know what comes next here" and
2818 // the same row the pick came from. Draft-quality only — verify arbitrates.
2819 if with_prob {
2820 e.prob_of_token_device_into(&logits, tok_d, p_d, d_vocab)?;
2821 }
2822 }
2823 }
2824 // ROUND-STREAM K-chain: pack (tok, p) into slot j, then remap tok through d2t so the
2825 // NEXT chained body's embed reads the TARGET id — zero host involvement per step.
2826 if let Some((out, slot, d2t)) = stream_pack {
2827 e.pack_tok_p(tok_d, p_d, out, slot)?;
2828 if let Some(map) = d2t {
2829 e.tok_map_u32(tok_d, map)?;
2830 }
2831 }
2832 // Next draft step's h_seed: pre-norm h_nextn (default) or post-norm final_h (HPOST).
2833 if spec_hpost() {
2834 e.copy_into(h_seed_d, 0, final_h.as_ref().unwrap(), n_embd)?;
2835 } else {
2836 e.copy_into(h_seed_d, 0, &h_nextn, n_embd)?;
2837 }
2838 // advance the draft rope position in-graph.
2839 e.inc_seqlen(pos_d)?;
2840 Ok(())
2841 }
2842
2843 /// Batched target verify forward over `tokens` at positions `pos0..pos0+T` (§D.3, T=K+1).
2844 /// Returns ALL T logit columns (host f32, [T*n_vocab]); appends T cols to every full-attn KV
2845 /// and advances every linear-attn recur state by T steps (the recur steps are SEQUENTIAL T=1).
2846 /// Advances `cache.pos` by T.
2847 pub fn decode_step_t(&self, e: &Engine, tokens: &[u32], pos0: usize, cache: &mut Cache)
2848 -> Result<Vec<f32>, Box<dyn std::error::Error>> {
2849 if self.is_gemma4_e4b() {
2850 return Ok(self.gemma4_e4b_decode_step_t_h(e, tokens, pos0, cache)?.0);
2851 }
2852 if self.cfg.gemma4.is_some() {
2853 return self.gemma4_decode_step_t(e, tokens, pos0, cache);
2854 }
2855 Ok(self.decode_step_t_h(e, tokens, pos0, cache)?.0)
2856 }
2857
2858 /// Like `decode_step_t` but ALSO returns the LAST column's pre-output_norm hidden (h_seed for
2859 /// the next draft round). This lets partial-accept replay run as ONE batched T=(n_acc+1) forward
2860 /// (single weight read) instead of n_acc+1 separate T=1 decode_steps (n_acc+1 weight reads).
2861 /// At batch=1 decode is bandwidth-bound, so batching the replay is THE MTP profitability lever.
2862 pub fn decode_step_t_h(
2863 &self,
2864 e: &Engine,
2865 tokens: &[u32],
2866 pos0: usize,
2867 cache: &mut Cache,
2868 ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
2869 self.decode_step_t_h_emb(e, tokens, pos0, cache, None)
2870 }
2871
2872 /// Like `decode_step_t_h` with an optional RESIDENT embed table (spec hot loop): device
2873 /// gather instead of host dequant + [T, n_embd] f32 htod. Bit-identical rows.
2874 pub fn decode_step_t_h_emb(
2875 &self,
2876 e: &Engine,
2877 tokens: &[u32],
2878 pos0: usize,
2879 cache: &mut Cache,
2880 embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
2881 ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
2882 let (logits_d, h_seed) = self.decode_step_t_h_emb_dev(e, tokens, pos0, cache, embd_dev)?;
2883 Ok((e.dtoh(&logits_d)?, h_seed))
2884 }
2885
2886 /// DEVICE-LOGITS verify forward (spec device-argmax lever): identical kernel chain to
2887 /// `decode_step_t_h_emb` but returns the [T, n_vocab] logits ON DEVICE — the accept walk
2888 /// argmaxes each column on-device and reads back ONE [T] u32 instead of dtoh'ing the full
2889 /// T x n_vocab f32 block (~1-4 MB + T host argmaxes, every round). Kernel dispatch is
2890 /// UNCHANGED (same decode-exact kernels); only the post-logits transfer moves.
2891 pub fn decode_step_t_h_emb_dev(
2892 &self,
2893 e: &Engine,
2894 tokens: &[u32],
2895 pos0: usize,
2896 cache: &mut Cache,
2897 embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
2898 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
2899 let n_embd = self.cfg.n_embd as usize;
2900 let t = tokens.len();
2901 let (logits, x) = self.decode_step_t_core(e, tokens, pos0, cache, embd_dev, None)?;
2902 // h_seed for the next round = LAST column's pre-output_norm hidden ([n_embd]).
2903 let mut hs = vbuf(e, n_embd)?; // fully written by copy_view_into below
2904 e.copy_view_into(&mut hs, 0, &x.slice((t - 1) * n_embd..t * n_embd), n_embd)?;
2905 Ok((logits, hs))
2906 }
2907
2908 /// CORE verify forward: the `decode_step_t_h_emb_dev` kernel chain, returning the FULL
2909 /// pre-output_norm hidden stack x ([T, n_embd], any column extractable) and optionally
2910 /// filling a `VerifyCkpt` (retained per-layer state-rebuild inputs) for the REPLAY-FREE
2911 /// partial accept. `ckpt: None` => byte-for-byte the old behavior (the ckpt writes are pure
2912 /// retains/copies — they never change what any kernel computes).
2913 fn decode_step_t_core(
2914 &self,
2915 e: &Engine,
2916 tokens: &[u32],
2917 pos0: usize,
2918 cache: &mut Cache,
2919 embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
2920 mut ckpt: Option<&mut VerifyCkpt>,
2921 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
2922 self.decode_step_t_core_stream(
2923 e, tokens, pos0, cache, embd_dev, ckpt.take(), None, None,
2924 )
2925 }
2926
2927 /// Increment-0 two-session PP seam: release the peer after this lane's stage-0 boundary TX.
2928 /// The two independent sessions keep their own cache/checkpoint state; only issue order moves.
2929 fn decode_step_t_core_pipelined(
2930 &self,
2931 e: &Engine,
2932 tokens: &[u32],
2933 pos0: usize,
2934 cache: &mut Cache,
2935 embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
2936 mut ckpt: Option<&mut VerifyCkpt>,
2937 pipe: &SpecPipeLane,
2938 round: usize,
2939 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
2940 let fence = crate::pp::pp_cuts(self.layers.len())
2941 .ok_or("two-session speculative pipeline requires a PP stage cut")?;
2942 if crate::pp::pp2_streams_off() || !crate::pp::spec_pp_on() {
2943 return Err("two-session speculative pipeline requires the PP verify split".into());
2944 }
2945 let interval_fence = pipe.stage0_begin(round)?;
2946 let ticket = self.verify_stage0_issue(
2947 e,
2948 tokens,
2949 pos0,
2950 cache,
2951 embd_dev,
2952 ckpt.as_deref_mut(),
2953 None,
2954 &fence,
2955 Some(interval_fence),
2956 pipe.trace(round),
2957 )?;
2958 pipe.stage0_end(round);
2959 pipe.stage1_begin(round)?;
2960 let result = self.verify_stage1_finish(e, ticket, cache, ckpt, None, &fence, true)?;
2961 pipe.verify_end(round);
2962 Ok(result)
2963 }
2964
2965 /// ROUND-STREAM stage (c) 4: `stream` = (device verify tokens [t], device pos counter) —
2966 /// when Some, rope positions come from pos_iota over the counter, the embed gathers the
2967 /// device tokens, and full_attn_verify routes appends/FA through the _dc twins reading the
2968 /// SAME counter (every layer's kvl.len == cache.pos, one counter drives all three). The
2969 /// host `tokens`/`pos0` args still size buffers (t is FIXED K+1 in stream mode).
2970 #[allow(clippy::too_many_arguments)]
2971 fn decode_step_t_core_stream(
2972 &self,
2973 e: &Engine,
2974 tokens: &[u32],
2975 pos0: usize,
2976 cache: &mut Cache,
2977 embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
2978 mut ckpt: Option<&mut VerifyCkpt>,
2979 stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
2980 pp_pipe: Option<bool>,
2981 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
2982 // PP DOOR (lane/pp2-spec 2026-08-06): the verify trunk now takes its OWN stage split,
2983 // exactly as the eager and batched steps do. This is the single funnel every verify
2984 // forward reaches (decode_step_t / _h / _h_emb / _h_emb_dev / _core all land here), so
2985 // wiring it here wires the whole spec surface — the draft/accept/commit machinery above
2986 // is untouched.
2987 //
2988 // History: pp2-hardening (2026-08-06) made this funnel FAIL CLOSED, because its trunk
2989 // walk was unsplit on one stream and a sharded cross-device placement peer-read every
2990 // remote layer's weights on every spec round (measured 13.9-28x on the batched twin).
2991 // The refusal below survives to cover the residue — MEMRA_SPEC_PP=0, MEMRA_PP_STREAMS=0,
2992 // or a placement whose PpNRt fails to build — so a config that would still walk the
2993 // whole trunk on one stream refuses instead of regressing 28x.
2994 if let Some(fence) = crate::pp::pp_cuts(self.layers.len()) {
2995 if !crate::pp::pp2_streams_off() && crate::pp::spec_pp_on() {
2996 return self.decode_step_t_core_ppn(
2997 e, tokens, pos0, cache, embd_dev, ckpt.take(), stream, &fence, pp_pipe,
2998 );
2999 }
3000 }
3001 crate::pp::refuse_unsplit_if_remote(
3002 "decode_step_t (spec verify)",
3003 "drop MEMRA_SPEC_PP=0 / MEMRA_PP_STREAMS=0 so the verify trunk takes its OWN stage \
3004 split (decode_step_t_core_ppn); or run spec on one device",
3005 )?;
3006 let cfg = &self.cfg;
3007 let n_embd = cfg.n_embd as usize;
3008 let eps = cfg.rms_eps;
3009 let t = tokens.len();
3010 let pos_d = match stream {
3011 Some((_, ctr)) => {
3012 let mut p = e.alloc_uninit::<i32>(t)?;
3013 e.pos_iota(ctr, &mut p, t)?;
3014 p
3015 }
3016 None => {
3017 let pos_vec: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
3018 e.htod_i32(&pos_vec)?
3019 }
3020 };
3021
3022 // embed T tokens -> [T, n_embd] token-major (device gather on the spec hot loop)
3023 let x = match (stream, embd_dev) {
3024 (Some((vtok, _)), Some((g, qt, rb))) => {
3025 e.embed_gather_device_td(g, vtok, t, n_embd, qt, rb)?
3026 }
3027 (None, Some((g, qt, rb))) => e.embed_gather_device_t(g, tokens, n_embd, qt, rb)?,
3028 _ => e.htod(&self.embd.gather(n_embd, tokens))?,
3029 };
3030
3031 // TRUNK WALK: layers [0, n_layers) through the SAME range-scoped subgraph the PP-N
3032 // stage split calls per stage (`verify_layers`) — one code path, so the split cannot
3033 // drift from the unsplit dispatch mirroring. lane/pp2-spec 2026-08-06.
3034 let x = self.verify_layers(
3035 e, x, 0, self.layers.len(), &pos_d, pos0, t, cache, ckpt.take(), stream,
3036 )?;
3037
3038 let mut hn = vbuf(e, t * n_embd)?;
3039 let serving_head = self.cfg.step35.is_some()
3040 || matches!(self.cfg.arch, memra_gguf::config::Arch::Qwen35Moe);
3041 let logits = if serving_head {
3042 // Step35 and Qwen35-MoE serving use one batched numeric class at every live width,
3043 // including B=1. Keep the verify head in that same class; other generic families
3044 // retain the decode-exact head that their run-spec contract pins.
3045 e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
3046 e.matmul(&self.output, &hn, t)?
3047 } else {
3048 e.rms_norm_decode(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
3049 e.matmul_decode_exact(&self.output, &hn, t)?
3050 };
3051 // stream: the device pos counter owns position; host mirror reconciles at drain.
3052 if stream.is_none() {
3053 cache.pos += t;
3054 }
3055 // Hidden stack for seeds/refresh-fills: pre-norm x (default) or post-norm hn (HPOST).
3056 Ok((logits, if spec_hpost() { hn } else { x }))
3057 }
3058
3059 /// THE VERIFY TRUNK OVER PP-N (lane/pp2-spec 2026-08-06): `decode_step_t_core_stream`'s walk
3060 /// as N stage subgraphs, each on its own engine/stream (and, under `MEMRA_PP_DEVICES`, its own
3061 /// device), with a `[T, n_embd]` boundary transfer between them. T = K+1 (the verify batch),
3062 /// so this is the batched-boundary shape the pp2-batch lane's grow-only slots already handle
3063 /// (`tx(b, x, t*n_embd)`; the slot grows to the high-water T and the transport moves exactly
3064 /// the payload).
3065 ///
3066 /// Structure is `decode_step_batch_ppn`'s, which is `decode_step_h_ppn`'s. FOUR THINGS ARE
3067 /// PER-STAGE and each for a measured reason (see `decode_step_batch_ppn`'s header for the
3068 /// receipts):
3069 ///
3070 /// 1. THE ENGINE (`rt.engine(s, e)`) — `Engine` owns lazily-grown stable-pointer scratch
3071 /// (`fa_part_pool`, `fa_vf16_scratch`, `argmax_partials`) that is single-stream-safe BY
3072 /// DESIGN. Two stage streams through one Engine is the 2026-08-02 shared-scratch race
3073 /// (35% flake, nondeterministic all-logits divergence). `PpNRt::build` gives every stage
3074 /// s>0 its own Engine even on the primary device; honouring it here is what scopes the
3075 /// pools. The verify path allocates MORE of that scratch than eager decode does (FA at
3076 /// m=T, and the per-layer `GdnStash` retains), so this is load-bearing, not inherited.
3077 ///
3078 /// 2. `pos_d` — each stage uploads its OWN copy of the T rope positions on ITS stream, so the
3079 /// buffer is allocated, consumed and freed on one stream. In `stream` mode that means each
3080 /// stage runs its own `pos_iota` over the SHARED device counter (`pos_ctr`): the counter is
3081 /// read-only during the forward (the round's `inc`/`copy_add` happen outside it), so every
3082 /// stage derives the identical iota, and each stage's own output buffer is stream-local.
3083 ///
3084 /// 3. THE EMBED lives with stage 0 (`self.embd` / `embd_gpu` are host/primary-side; the
3085 /// sharded loader leaves the table with stage 0 by construction).
3086 ///
3087 /// 4. THE HEAD (`output_norm` + `output`) runs on the LAST stage — the sharded loader uploaded
3088 /// both through that stage's engine (`hybrid.rs`: `e_head = layer_engine(e, n_trunk,
3089 /// n_trunk-1)`), so reading them anywhere else is a peer read of the biggest tensor in the
3090 /// model, every round.
3091 ///
3092 /// WHAT STAYS ON THE PRIMARY, deliberately: the returned logits and hidden stack `x`. Both are
3093 /// last-stage-allocated device buffers, and every consumer (the device argmax walk, the accept
3094 /// kernels, `spec_seed_gather`, the ckpt rebuild in `commit_verified_prefix`) reads them
3095 /// through the primary context by UVA — the same read the batched serving epilogue's
3096 /// `last_logits_dev` park does. Those consumers are per-round O(T x n_vocab) and O(n_embd),
3097 /// not per-layer, so they are not the 28x class; splitting them is a separate lane.
3098 ///
3099 /// The MTP HEAD (draft side) is NOT split: it is one block, it lives wherever the loader put
3100 /// it (`load_mtp` uses the primary engine), and it is ~1-2 GB against the trunk's tens. Draft
3101 /// placement is measured, not assumed — see `research/pp2-spec-20260806`.
3102 ///
3103 /// EXACTNESS: PP-N adds ZERO deviation. Each stage runs the SAME kernels on the SAME bytes in
3104 /// the same order via the SAME `verify_layers` the unsplit body calls; the only change is
3105 /// where the residual is materialized, and the boundary is a straight f32 copy (dtod
3106 /// same-device / `cudaMemcpyPeerAsync` cross-device, no conversion). So the split MUST be
3107 /// BIT-IDENTICAL to the unsplit verify at the same T, in both placement orders. Gate:
3108 /// `decode-batch-gate --mode ppspec`. Acceptance counts are a DERIVED consequence — greedy
3109 /// accept argmaxes these logits, so bit-identical logits force identical accept walks; the
3110 /// `run-spec` K=1..8 arm checks that end-to-end rather than trusting the implication.
3111 #[allow(clippy::too_many_arguments)]
3112 fn decode_step_t_core_ppn(
3113 &self,
3114 e: &Engine,
3115 tokens: &[u32],
3116 pos0: usize,
3117 cache: &mut Cache,
3118 embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
3119 mut ckpt: Option<&mut VerifyCkpt>,
3120 stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
3121 fence: &[usize],
3122 pp_pipe: Option<bool>,
3123 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
3124 let ticket = self.verify_stage0_issue(
3125 e,
3126 tokens,
3127 pos0,
3128 cache,
3129 embd_dev,
3130 ckpt.as_deref_mut(),
3131 stream,
3132 fence,
3133 pp_pipe,
3134 None,
3135 )?;
3136 self.verify_stage1_finish(e, ticket, cache, ckpt, stream, fence, true)
3137 }
3138
3139 /// Enqueue embed, stage 0, and the first boundary TX, then return the actual boundary slot.
3140 /// The ordinary PP verify wrapper calls `verify_stage1_finish` immediately after this return.
3141 #[allow(clippy::too_many_arguments)]
3142 fn verify_stage0_issue(
3143 &self,
3144 e: &Engine,
3145 tokens: &[u32],
3146 pos0: usize,
3147 cache: &mut Cache,
3148 embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
3149 mut ckpt: Option<&mut VerifyCkpt>,
3150 stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
3151 fence: &[usize],
3152 pp_pipe: Option<bool>,
3153 trace: Option<SpecPipeTraceCtx>,
3154 ) -> Result<VerifyBoundaryTicket, Box<dyn std::error::Error>> {
3155 assert!(
3156 !self.is_gemma4_e4b() && self.cfg.gemma4.is_none(),
3157 "decode_step_t_core_ppn covers the hybrid non-gemma4 verify trunk only \
3158 (the gemma4 arms have their own decode_step_t twins)"
3159 );
3160 if crate::pp::pp_host_bounce_active() && (stream.is_some() || embd_dev.is_some()) {
3161 return Err(
3162 "decode_step_t_core_ppn: refused with MEMRA_PP_HOST_BOUNCE=1 — the trunk \
3163 boundary itself is host-staged, but device-resident verify still peer-reads \
3164 primary-device token/position/embedding buffers from stage 0. Run plain PP \
3165 serving on this host class; spec requires local per-stage inputs first."
3166 .into(),
3167 );
3168 }
3169 let rt = crate::pp::PpNRt::get(e)?;
3170 let n_st = fence.len() - 1;
3171 assert_eq!(
3172 rt.n_stages(), n_st,
3173 "PpNRt stage count {} != fence stages {n_st}", rt.n_stages()
3174 );
3175 let n_embd = self.cfg.n_embd as usize;
3176 let t = tokens.len();
3177 let payload = t * n_embd;
3178 if pp_pipe.is_some() {
3179 assert_eq!(n_st, 2, "spec pipeline requires exactly two PP stages");
3180 }
3181 // One-shot lane diagnostic: force natural PP-2 boundaries to completion so the server
3182 // log can price stage 0, the peer hop, the RX copy, and stage 1 + head separately without
3183 // nsys. The ordinary path keeps every enqueue asynchronous. N>2 is deliberately excluded:
3184 // the report below names exactly two stages and must never imply it measured middle ones.
3185 let pp_anatomy = n_st == 2
3186 && std::env::var("MEMRA_SPEC_PP_ANATOMY").as_deref() == Ok("1");
3187 let pp_started = std::time::Instant::now();
3188 let (mut reverse_ms, mut stage0_ms, mut tx_ms) = (0.0f64, 0.0f64, 0.0f64);
3189 // The CALLER's ambient stream, captured BEFORE any `rt.enter()` pushes a stage stream:
3190 // this body returns DEVICE-RESIDENT buffers (the device-argmax accept walk's contract),
3191 // so the exit needs the same publication the boundaries get — see `PpNRt::publish_to`.
3192 // Taken here, not at the end, because inside the last-stage scope `e.stream()` IS the
3193 // stage stream and the wait would self-order into a no-op.
3194 let caller_stream = e.stream();
3195 // #87 ROOT-CAUSE FENCE (lane/pp2spec-crash): the PREVIOUS round's stage-allocated
3196 // outputs (logits/hidden/ckpt stashes) freed stream-ordered on the STAGE streams while
3197 // the primary stream still holds queued reads of them — with event tracking elided,
3198 // nothing stops the pool from reusing those blocks for THIS round's stage allocations,
3199 // whose writes then race the queued reads (measured: 13/4096-NaN random-bits garbage in
3200 // the spec round seed; the full anatomy is on `PpNRt::fence_stages_behind`). Order every
3201 // stage stream behind the caller before enqueueing new stage work.
3202 let reverse_started = std::time::Instant::now();
3203 if pp_pipe != Some(false) {
3204 rt.fence_stages_behind(&caller_stream)?;
3205 }
3206 if pp_pipe == Some(true) {
3207 // Both session verifies must alternate boundary slots even when the ordinary
3208 // decode overlap experiment is off. Prewarm before A's stage 0 so B cannot grow
3209 // slot 1 by synchronizing the RX stream while A's stage 1 is in flight.
3210 rt.prepare_overlap_slots(0, payload)?;
3211 }
3212 if pp_anatomy {
3213 // Drain the reverse-publication dependency before timing stage 0 itself. At c=1 this
3214 // prices any primary-stream rollback/refresh tail inherited from the prior round.
3215 for s in 0..n_st {
3216 let _st = rt.enter(s);
3217 rt.engine(s, e).stream().synchronize()?;
3218 }
3219 reverse_ms = reverse_started.elapsed().as_secs_f64() * 1e3;
3220 }
3221
3222 // Per-stage rope positions: in host mode the same [T] iota each stage uploads itself; in
3223 // stream mode each stage's own `pos_iota` over the shared read-only device counter.
3224 let stage_pos = |es: &Engine| -> Result<CudaSlice<i32>, Box<dyn std::error::Error>> {
3225 match stream {
3226 Some((_, ctr)) => {
3227 let mut p = es.alloc_uninit::<i32>(t)?;
3228 es.pos_iota(ctr, &mut p, t)?;
3229 Ok(p)
3230 }
3231 None => {
3232 let pos_vec: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
3233 es.htod_i32(&pos_vec)
3234 }
3235 }
3236 };
3237
3238 // ---- STAGE 0: embed (the table lives with stage 0) + layers [0, fence[1]) + TX ----
3239 let slot = {
3240 let _st0 = rt.enter(0);
3241 let e0 = rt.engine(0, e);
3242 enqueue_spec_pipe_trace_marker(&e0.stream(), trace.as_ref(), "S0", "start", None)?;
3243 let stage0_started = std::time::Instant::now();
3244 let pos_d = stage_pos(e0)?;
3245 let x = match (stream, embd_dev) {
3246 (Some((vtok, _)), Some((g, qt, rb))) => {
3247 e0.embed_gather_device_td(g, vtok, t, n_embd, qt, rb)?
3248 }
3249 (None, Some((g, qt, rb))) => e0.embed_gather_device_t(g, tokens, n_embd, qt, rb)?,
3250 _ => e0.htod(&self.embd.gather(n_embd, tokens))?,
3251 };
3252 let x = self.verify_layers(
3253 e0, x, fence[0], fence[1], &pos_d, pos0, t, cache, ckpt.as_deref_mut(), stream,
3254 )?;
3255 if pp_anatomy {
3256 e0.stream().synchronize()?;
3257 stage0_ms = stage0_started.elapsed().as_secs_f64() * 1e3;
3258 }
3259 let tx_started = std::time::Instant::now();
3260 let slot = if pp_pipe.is_some() {
3261 rt.tx_pipelined(0, &x, payload)?
3262 } else {
3263 rt.tx(0, &x, payload)?
3264 };
3265 enqueue_spec_pipe_trace_marker(
3266 &e0.stream(),
3267 trace.as_ref(),
3268 "S0",
3269 "end",
3270 Some(slot),
3271 )?;
3272 if pp_anatomy {
3273 e0.stream().synchronize()?;
3274 tx_ms = tx_started.elapsed().as_secs_f64() * 1e3;
3275 }
3276 slot
3277 // x + pos_d drop here: freed stream-ordered on stage-0's stream after use.
3278 };
3279
3280 Ok(VerifyBoundaryTicket {
3281 rt,
3282 caller_stream,
3283 slot,
3284 pos0,
3285 t,
3286 payload,
3287 n_st,
3288 pipelined: pp_pipe.is_some(),
3289 pp_anatomy,
3290 pp_started,
3291 reverse_ms,
3292 stage0_ms,
3293 tx_ms,
3294 trace,
3295 })
3296 }
3297
3298 /// Consume a stage-0 boundary ticket and enqueue the remaining PP stages plus the head.
3299 /// On PP-2 this is exactly stage 1; PP-N keeps its pre-existing middle-stage walk here.
3300 #[allow(clippy::too_many_arguments)]
3301 fn verify_stage1_finish(
3302 &self,
3303 e: &Engine,
3304 ticket: VerifyBoundaryTicket,
3305 cache: &mut Cache,
3306 mut ckpt: Option<&mut VerifyCkpt>,
3307 stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
3308 fence: &[usize],
3309 publish_to_caller: bool,
3310 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
3311 let VerifyBoundaryTicket {
3312 rt,
3313 caller_stream,
3314 slot,
3315 pos0,
3316 t,
3317 payload,
3318 n_st,
3319 pipelined,
3320 pp_anatomy,
3321 pp_started,
3322 reverse_ms,
3323 stage0_ms,
3324 tx_ms,
3325 trace,
3326 } = ticket;
3327 let n_embd = self.cfg.n_embd as usize;
3328 let eps = self.cfg.rms_eps;
3329 let mut slot = slot;
3330 let (mut rx_ms, mut stage1_ms) = (0.0f64, 0.0f64);
3331 let stage_pos = |es: &Engine| -> Result<CudaSlice<i32>, Box<dyn std::error::Error>> {
3332 match stream {
3333 Some((_, ctr)) => {
3334 let mut p = es.alloc_uninit::<i32>(t)?;
3335 es.pos_iota(ctr, &mut p, t)?;
3336 Ok(p)
3337 }
3338 None => {
3339 let pos_vec: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
3340 es.htod_i32(&pos_vec)
3341 }
3342 }
3343 };
3344
3345 // ---- MIDDLE STAGES: RX boundary s-1 -> range -> TX boundary s ----
3346 for s in 1..n_st - 1 {
3347 let _st = rt.enter(s);
3348 let es = rt.engine(s, e);
3349 let pos_d = stage_pos(es)?;
3350 let x = rt.rx(s - 1, slot, payload)?;
3351 let x = self.verify_layers(
3352 es, x, fence[s], fence[s + 1], &pos_d, pos0, t, cache,
3353 ckpt.as_deref_mut(), stream,
3354 )?;
3355 slot = if pipelined {
3356 rt.tx_pipelined(s, &x, payload)?
3357 } else {
3358 rt.tx(s, &x, payload)?
3359 };
3360 }
3361
3362 // ---- LAST STAGE: RX + final range + output_norm + lm head ----
3363 let _stl = rt.enter(n_st - 1);
3364 let el = rt.engine(n_st - 1, e);
3365 let pos_d = stage_pos(el)?;
3366 let rx_started = std::time::Instant::now();
3367 let x = rt.rx(n_st - 2, slot, payload)?;
3368 if pp_anatomy {
3369 el.stream().synchronize()?;
3370 rx_ms = rx_started.elapsed().as_secs_f64() * 1e3;
3371 }
3372 enqueue_spec_pipe_trace_marker(
3373 &el.stream(),
3374 trace.as_ref(),
3375 "S1",
3376 "start",
3377 Some(slot),
3378 )?;
3379 let stage1_started = std::time::Instant::now();
3380 let x = self.verify_layers(
3381 el, x, fence[n_st - 1], fence[n_st], &pos_d, pos0, t, cache,
3382 ckpt.as_deref_mut(), stream,
3383 )?;
3384
3385 let mut hn = vbuf(el, payload)?;
3386 let logits = if self.cfg.step35.is_some() {
3387 // The PP Step35 serving path uses rms_norm + matmul for B=1 as well as B>1.
3388 // Verify must not switch numeric class merely because the same session speculates.
3389 el.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
3390 el.matmul(&self.output, &hn, t)?
3391 } else {
3392 el.rms_norm_decode(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
3393 el.matmul_decode_exact(&self.output, &hn, t)?
3394 };
3395 enqueue_spec_pipe_trace_marker(
3396 &el.stream(),
3397 trace.as_ref(),
3398 "S1",
3399 "end",
3400 Some(slot),
3401 )?;
3402 if pp_anatomy {
3403 el.stream().synchronize()?;
3404 stage1_ms = stage1_started.elapsed().as_secs_f64() * 1e3;
3405 }
3406 // EXIT PUBLICATION: both returned buffers are still being produced on the last stage's
3407 // stream. Order the caller's stream behind that work before the buffers escape this
3408 // scope (the 2026-08-06 same-device ppspec find: without it the caller's primary-stream
3409 // consumer read unwritten logits — nondeterministic, one-device-only, and it poisoned
3410 // the following arm's KV in the same process).
3411 if publish_to_caller {
3412 rt.publish_to(n_st - 1, &caller_stream)?;
3413 }
3414 if pp_anatomy {
3415 if publish_to_caller {
3416 caller_stream.synchronize()?;
3417 }
3418 eprintln!(
3419 "[spec-pp-anatomy] t={t} reverse={reverse_ms:.3}ms stage0={stage0_ms:.3}ms \
3420 tx={tx_ms:.3}ms rx={rx_ms:.3}ms stage1-head={stage1_ms:.3}ms total={:.3}ms",
3421 pp_started.elapsed().as_secs_f64() * 1e3,
3422 );
3423 }
3424 // stream: the device pos counter owns position; host mirror reconciles at drain.
3425 if stream.is_none() {
3426 cache.pos += t;
3427 }
3428 Ok((logits, if spec_hpost() { hn } else { x }))
3429 }
3430
3431 /// Step3.5/Step3.7 verify trunk in the serving batched numeric class.
3432 ///
3433 /// `step35_decode_batch_layers` is now authoritative at every live serving width, including
3434 /// B=1 (lane/cx-b1fix). The older verify walk deliberately mirrored the eager T=1 class:
3435 /// it replayed `step35_decode_attn` per row and used the eager/decode-exact FFN dispatch.
3436 /// Those classes are individually stable, but a near-tie prompt can choose different greedy
3437 /// bytes when a request moves from batched plain serving into speculative verify. Run the
3438 /// same authoritative B=1 stage subgraph for each verify row here. Rows still advance
3439 /// layer-by-layer, so every layer sees the preceding verify rows in its attention cache while
3440 /// every norm/projection/FFN uses exactly the live serving dispatch.
3441 #[allow(clippy::too_many_arguments)]
3442 fn step35_verify_batch_layers(
3443 &self,
3444 e: &Engine,
3445 mut x: CudaSlice<f32>,
3446 lo: usize,
3447 hi: usize,
3448 pos0: usize,
3449 t: usize,
3450 cache: &mut Cache,
3451 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3452 let n_embd = self.cfg.n_embd as usize;
3453 self.cfg.step35.as_ref().ok_or("step35 verify batch requires step35 cfg")?;
3454 let mut ph_last = std::time::Instant::now();
3455 for il in lo..hi {
3456 let mut next = e.uninit(t * n_embd)?;
3457 for r in 0..t {
3458 let mut row = e.uninit(n_embd)?;
3459 e.dtod_copy_view(&x.slice(r * n_embd..(r + 1) * n_embd), &mut row)?;
3460 // The caller owns this verify's position. During controller overlap, cache.pos
3461 // still describes generation N while this stage-0 walk belongs to N+1.
3462 let row_pos = e.htod_i32(&[(pos0 + r) as i32])?;
3463 let mut one = [&mut *cache];
3464 let out = self.step35_decode_batch_layers(
3465 e,
3466 row,
3467 &mut one,
3468 &row_pos,
3469 il,
3470 il + 1,
3471 &mut ph_last,
3472 )?;
3473 e.dtod_copy_into(&out, &mut next, r * n_embd)?;
3474 }
3475 x = next;
3476 }
3477 Ok(x)
3478 }
3479
3480 /// Qwen35-MoE verify trunk in the live serving numeric class.
3481 ///
3482 /// Serving intentionally keeps this architecture in the generic batched program even at
3483 /// B=1. The older verify walk used its own mirrored dispatch and can flip near-tie argmaxes.
3484 /// Replay each verify row through the authoritative serving layer body while preserving the
3485 /// single-session autoregressive cache order.
3486 #[allow(clippy::too_many_arguments)]
3487 fn qwen35_moe_verify_batch_layers(
3488 &self,
3489 e: &Engine,
3490 mut x: CudaSlice<f32>,
3491 lo: usize,
3492 hi: usize,
3493 pos0: usize,
3494 t: usize,
3495 cache: &mut Cache,
3496 mut ckpt: Option<&mut VerifyCkpt>,
3497 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3498 let n_embd = self.cfg.n_embd as usize;
3499 let saved_pos = cache.pos;
3500 let mut ph_last = std::time::Instant::now();
3501 for il in lo..hi {
3502 let mut next = e.uninit(t * n_embd)?;
3503 let mut col_states: Option<Vec<(CudaSlice<f32>, CudaSlice<f32>)>> =
3504 if ckpt.is_some()
3505 && t >= 2
3506 && matches!(self.layers[il].mixer, Mixer::Linear(_))
3507 {
3508 Some(Vec::with_capacity(t - 1))
3509 } else {
3510 None
3511 };
3512 for r in 0..t {
3513 cache.pos = pos0 + r;
3514 let mut row = e.uninit(n_embd)?;
3515 e.dtod_copy_view(&x.slice(r * n_embd..(r + 1) * n_embd), &mut row)?;
3516 let row_pos = e.htod_i32(&[(pos0 + r) as i32])?;
3517 let mut one = [&mut *cache];
3518 let ctx = self.batch_layer_ctx(e, &one, il, il + 1)?;
3519 let out = match self.decode_batch_layers(
3520 e,
3521 row,
3522 &mut one,
3523 &ctx,
3524 &row_pos,
3525 &mut ph_last,
3526 ) {
3527 Ok(out) => out,
3528 Err(error) => {
3529 cache.pos = saved_pos;
3530 return Err(error);
3531 }
3532 };
3533 e.dtod_copy_into(&out, &mut next, r * n_embd)?;
3534 if r + 1 < t {
3535 if let Some(states) = col_states.as_mut() {
3536 let recur = cache.recur[il]
3537 .as_ref()
3538 .ok_or("Qwen35-MoE linear verify layer has no recurrent state")?;
3539 states.push((
3540 e.clone_dtod(&recur.conv_state)?,
3541 e.clone_dtod(&recur.ssm_state)?,
3542 ));
3543 }
3544 }
3545 }
3546 if let (Some(checkpoint), Some(states)) = (ckpt.as_deref_mut(), col_states) {
3547 checkpoint.cols[il] = Some(states);
3548 }
3549 x = next;
3550 }
3551 cache.pos = saved_pos;
3552 Ok(x)
3553 }
3554
3555 /// PP-N STAGE SUBGRAPH of the verify trunk: layers `[lo, hi)` of `decode_step_t_core_stream`'s
3556 /// walk, verbatim. Enters with a MATERIALIZED `[T, n_embd]` residual (no pending fusion pair
3557 /// carried in from outside the range) and exits with the range's final residual materialized
3558 /// (the trailing add executed) — exactly the `decode_layers_eager(lo, hi)` contract, T rows
3559 /// instead of one.
3560 ///
3561 /// EXTRACTED (lane/pp2-spec 2026-08-06) rather than duplicated: `decode_step_t_core_stream` IS
3562 /// the single funnel every verify forward reaches, and its per-layer dispatch MIRRORING (norm
3563 /// fusion per layer, the t>=3/spec_m2 batched-linear window, the fused-q8 FFN chain, the
3564 /// decode-exact projections) is what makes verify bit-identical to eager decode. A second copy
3565 /// for the split arm is how those mirrors drift apart on the next lever. The unsplit body now
3566 /// calls this with `(0, n_layers)`, so the whole-trunk path and every stage range run the SAME
3567 /// code — there is no "split version" of the verify math.
3568 ///
3569 /// Bit-identity of a cut rests on the same kernel-check-pinned identity the eager arm's cut
3570 /// does — `add_rms_norm_q8_1 == add then rms_norm_q8_1` at nrows=T — because the ONLY thing a
3571 /// fence changes is that the cross-layer fusion carry breaks at `hi-1` and is re-materialized
3572 /// as an explicit `add`. `decode-batch-gate --mode ppspec` verifies end-to-end on real weights.
3573 #[allow(clippy::too_many_arguments)]
3574 fn verify_layers(
3575 &self,
3576 e: &Engine,
3577 mut x: CudaSlice<f32>,
3578 lo: usize,
3579 hi: usize,
3580 pos_d: &CudaSlice<i32>,
3581 pos0: usize,
3582 t: usize,
3583 cache: &mut Cache,
3584 mut ckpt: Option<&mut VerifyCkpt>,
3585 stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
3586 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3587 if self.cfg.step35.is_some() {
3588 if stream.is_some() {
3589 return Err("step35 has no ROUND-STREAM verify arm (the device-counter _dc twins \
3590 cannot express the SWA offset KV view)".into());
3591 }
3592 return self.step35_verify_batch_layers(e, x, lo, hi, pos0, t, cache);
3593 }
3594 if matches!(self.cfg.arch, memra_gguf::config::Arch::Qwen35Moe) {
3595 if stream.is_some() {
3596 return Err("Qwen35-MoE serving-class verify has no ROUND-STREAM arm".into());
3597 }
3598 return self.qwen35_moe_verify_batch_layers(
3599 e,
3600 x,
3601 lo,
3602 hi,
3603 pos0,
3604 t,
3605 cache,
3606 ckpt.take(),
3607 );
3608 }
3609 let n_embd = self.cfg.n_embd as usize;
3610 let eps = self.cfg.rms_eps;
3611 // CROSS-LAYER ADD+NORM FUSION (lane/vt-fixes fix 2, mirroring decode_step_h's
3612 // launch-arc form): layer il's post-FFN residual add (x2 = x1 + ffn_out) and layer
3613 // il+1's attn_norm(+quantize) are consecutive row-wise ops — ONE add_rms_norm_q8_1
3614 // launch at nrows=t does all three (bit-identity pinned by the T-row kernel-check
3615 // arms). Carry the un-added (x1, ffn_out) pair; the fused launch materializes x2 (the
3616 // residual the next layer needs) as its `res` output. Falls back to the separate add
3617 // when the next layer is off the fused-q8 path.
3618 let mut pending: Option<(CudaSlice<f32>, CudaSlice<f32>)> = None;
3619 for il in lo..hi {
3620 let layer = &self.layers[il];
3621 // DISPATCH-MIRRORED attn-input RMSNorm (FP-order lesson #8): eager decode fuses the
3622 // 1024-thread rms_norm_q8_1 ONLY when every mixer projection is q8_1-fast; layers with
3623 // Float projections (ssm_beta/ssm_alpha on layers 1/2/4 of the 9B NVFP4 GGUF) take the
3624 // UNFUSED 256-thread rms_norm. The verify norm must mirror that PER-LAYER choice —
3625 // blockDim changes the sum-of-squares reduce order, and the ULP shift amplifies through
3626 // the GDN recurrence into argmax flips (measured: 9B text prompt, 1 ULP at layer 2 ->
3627 // 2.3e-1 logit maxdiff at the head -> K=1..8 divergence at a 0.03-margin token).
3628 let mixer_fast = self.mixer_in_q8_1_fast(e, &layer.mixer);
3629 let norm_fused = std::env::var("MEMRA_NO_FUSE_NORMQ").is_err() && mixer_fast;
3630 // BATCHED EPILOGUE RE-FUSE (lane/vt-fixes fix 2, 2026-08-03): when the norm is
3631 // dispatch-fused AND every consumer of `h` reads only its q8_1 form (Full mixer:
3632 // projections only; Linear mixer: the batched arm — the per-column fallback needs
3633 // f32 h), emit the attn-input norm DIRECTLY as q8_1 via `rms_norm_q8_1` at nrows=t
3634 // (row-indexed kernel — the T-row launch is the per-row m=1 program, kernel-check
3635 // pins bit-identity vs rms_norm_decode -> quantize_q8_1). Kills the standalone
3636 // quantize launch(es) + the f32 h HBM round-trip that decode never pays.
3637 // step35 (Full mixer) is the third case that needs f32 `h`: its verify arm is a
3638 // per-ROW replay of the eager decode mixer, whose `pre_q` contract is a single row —
3639 // a T-row q8_1 pair cannot be handed to it, and re-deriving per-row q8_1 from the f32
3640 // rows is exactly the dispatch being mirrored. Keep step35 on the unfused arm.
3641 let lin_q8_only = match &layer.mixer {
3642 Mixer::Linear(la) => {
3643 (t >= 3 || (t == 2 && spec_m2())) && e.uses_q8_1_fast(&la.ssm_out)
3644 }
3645 Mixer::Full(_) if self.cfg.step35.is_some() => false,
3646 _ => true,
3647 };
3648 // NOTE decode.rs's take()-first lesson: take the pending pair BEFORE branching so
3649 // a non-fused layer still performs the residual add.
3650 let taken = pending.take();
3651 let (h, h_q8) = if norm_fused && lin_q8_only {
3652 let pair = match taken {
3653 // fused add + attn_norm + q8_1: ONE launch resolves the carried residual
3654 // AND emits this layer's mixer input pre-quantized. res -> x2 (= new x).
3655 Some((x1p, f1p)) => {
3656 let mut x2 = vbuf(e, t * n_embd)?; // fully written (res output)
3657 let p = e.add_rms_norm_q8_1(
3658 &x1p, &f1p, layer.attn_norm.float_data(), &mut x2, n_embd, t, eps,
3659 )?;
3660 x = x2;
3661 p
3662 }
3663 None => e.rms_norm_q8_1(&x, layer.attn_norm.float_data(), n_embd, t, eps)?,
3664 };
3665 (e.zeros(0)?, Some(pair)) // h unused on this path (q8-only consumers)
3666 } else {
3667 if let Some((x1p, f1p)) = taken {
3668 let mut x2 = vbuf(e, t * n_embd)?; // fully written by add
3669 e.add(&x1p, &f1p, &mut x2, t * n_embd)?;
3670 x = x2;
3671 }
3672 let mut h = vbuf(e, t * n_embd)?; // fully written by either rms_norm arm
3673 if norm_fused {
3674 e.rms_norm_decode(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
3675 } else {
3676 e.rms_norm(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
3677 }
3678 (h, None)
3679 };
3680 let h_q8_ref = h_q8.as_ref().map(|(q, d)| (q, d));
3681
3682 let mixed = match &layer.mixer {
3683 Mixer::Full(fa) => {
3684 self.full_attn_verify(e, fa, &h, h_q8_ref, pos_d, t, cache, il,
3685 stream.map(|(_, c)| c))?
3686 }
3687 Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
3688 Mixer::Linear(la) => {
3689 // BATCHED linear verify (2026-07-03, the MTP-profit lever): one T-token pass —
3690 // batched projections (weight read ONCE, hits the m=2-4 weight-resident matvec),
3691 // carried-state conv (ssm_conv1d_tm_state), GDN prep on the prefill kernels, and
3692 // ONE gdn_scan whose internal sequential t-loop is the SAME recurrence as T
3693 // chained T=1 steps (bit-identical). Falls back to the sequential per-column
3694 // chain when T < d_conv-1 (conv ring update needs T >= pad) — or when ANY
3695 // projection is off the q8_1 fast path: matmul_decode_exact would route a Float
3696 // tensor to cuBLAS at m=t (different FP accumulation than eager's per-token
3697 // GEMV), so mixed-dtype layers stay on the eager-identical per-column chain.
3698 // MEMRA_SPEC_M2 (lane/spec-m2): the t==2 batch rides the same arm — the conv
3699 // wrapper handles t<pad with a pure-copy ring rebuild; see spec_m2() header.
3700 if (t >= 3 || (t == 2 && spec_m2()))
3701 && mixer_fast
3702 && e.uses_q8_1_fast(&la.ssm_out)
3703 {
3704 let want = ckpt.is_some();
3705 let (out, stash) =
3706 self.linear_attn_verify_t(e, la, &h, h_q8_ref, t, cache, il, want)?;
3707 if let (Some(ck), Some(st)) = (ckpt.as_deref_mut(), stash) {
3708 ck.gdn[il] = Some(st);
3709 }
3710 out
3711 } else {
3712 let mut out = vbuf(e, t * n_embd)?; // every col written by copy_into
3713 let mut col_states: Option<Vec<(CudaSlice<f32>, CudaSlice<f32>)>> =
3714 if ckpt.is_some() && t >= 2 {
3715 Some(Vec::with_capacity(t - 1))
3716 } else {
3717 None
3718 };
3719 for col in 0..t {
3720 let mut h_col = vbuf(e, n_embd)?; // fully written by copy_view_into
3721 let src = h.slice(col * n_embd..(col + 1) * n_embd);
3722 e.copy_view_into(&mut h_col, 0, &src, n_embd)?;
3723 let m_col = self.linear_attn_decode(e, la, &h_col, cache, il)?;
3724 e.copy_into(&mut out, col * n_embd, &m_col, n_embd)?;
3725 // REPLAY-FREE ckpt: clone the chain's ACTUAL state after this column
3726 // (pure dtod — cannot change any computed value). Last column skipped:
3727 // rebuild targets are j <= t-1 columns.
3728 if let Some(cs) = col_states.as_mut() {
3729 if col + 1 < t {
3730 let rl = cache.recur[il].as_ref().unwrap();
3731 cs.push((
3732 e.clone_dtod(&rl.conv_state)?,
3733 e.clone_dtod(&rl.ssm_state)?,
3734 ));
3735 }
3736 }
3737 }
3738 if let (Some(ck), Some(cs)) = (ckpt.as_deref_mut(), col_states) {
3739 // ReplaySSM-assessment instrumentation (2026-07-30): the
3740 // per-column clones are the only true state snapshots left in
3741 // the verify (the batched path stashes INPUTS and replays).
3742 if std::env::var("MEMRA_SPEC_STATS").as_deref() == Ok("1") {
3743 static ONCE: std::sync::Once = std::sync::Once::new();
3744 let bytes: usize = cs.iter()
3745 .map(|(c, s)| (c.len() + s.len()) * 4).sum();
3746 ONCE.call_once(|| eprintln!(
3747 "[verify-ckpt] per-column layer il={il}: {} clones, {:.2} MB/layer/round",
3748 cs.len(), bytes as f64 / 1e6));
3749 }
3750 ck.cols[il] = Some(cs);
3751 }
3752 out
3753 }
3754 }
3755 };
3756
3757 // DISPATCH-MIRRORED post-attn norm: eager residual_norm_ffn fuses add+norm+quant
3758 // (1024-thread add_rms_norm_q8_1) only for Dense FFNs whose gate+up are q8_1-fast;
3759 // otherwise (and for MoE) it runs the 256-thread fused add_rms_norm. Mirror per layer.
3760 let ffn_fuse = match &layer.ffn {
3761 crate::hybrid::Ffn::Dense {
3762 ffn_gate, ffn_up, ..
3763 } => {
3764 std::env::var("MEMRA_NO_FUSE_NORMQ").is_err()
3765 && e.uses_q8_1_fast(ffn_gate)
3766 && e.uses_q8_1_fast(ffn_up)
3767 }
3768 crate::hybrid::Ffn::Moe(_) => false,
3769 };
3770 // BATCHED EPILOGUE RE-FUSE (lane/vt-fixes fix 2): on the ffn_fuse path (Dense,
3771 // gate+up q8_1-fast, non-M3) the FFN input is emitted DIRECTLY as q8_1 by ONE
3772 // add_rms_norm_q8_1 launch at nrows=t (row-indexed kernel: the T-row launch is the
3773 // per-row m=1 program; kernel-check pins bit-identity vs the unfused
3774 // add_f32 -> rms_norm_decode -> quantize_q8_1 chain at T=2/4/5/8) — replacing the
3775 // add + rms_norm_decode launches AND the dual/singles' internal re-quantize.
3776 // M3's swigluoai must keep the f32 chain (the fused SwiGLU epilogue encodes plain
3777 // SiLU), mirroring residual_norm_ffn's m3 guard on the decode path.
3778 // step35: same guard per LAYER. A dense FFN's clamp is the SHEXP array (upstream's
3779 // one build_ffn serves dense + shared expert, llama-graph.cpp:1751), and verify MUST
3780 // mirror decode's dispatch or spec self-consistency fails.
3781 let dense_lim = self.cfg.clamp_shexp_at(il as u32);
3782 let fuse_q8 = ffn_fuse && self.cfg.m3.is_none() && dense_lim.is_none();
3783 let mut x1 = vbuf(e, t * n_embd)?; // fully written by add / add_rms_norm*
3784 let mut z = e.zeros(0)?; // replaced below on the unfused arms
3785 let z_q8 = if fuse_q8 {
3786 Some(e.add_rms_norm_q8_1(
3787 &x,
3788 &mixed,
3789 layer.post_attn_norm.float_data(),
3790 &mut x1,
3791 n_embd,
3792 t,
3793 eps,
3794 )?)
3795 } else {
3796 let mut zf = vbuf(e, t * n_embd)?; // fully written by rms_norm_decode / add_rms_norm
3797 if ffn_fuse {
3798 e.add(&x, &mixed, &mut x1, t * n_embd)?;
3799 e.rms_norm_decode(
3800 &x1,
3801 layer.post_attn_norm.float_data(),
3802 &mut zf,
3803 n_embd,
3804 t,
3805 eps,
3806 )?;
3807 } else {
3808 e.add_rms_norm(
3809 &x,
3810 &mixed,
3811 layer.post_attn_norm.float_data(),
3812 &mut x1,
3813 &mut zf,
3814 n_embd,
3815 t,
3816 eps,
3817 )?;
3818 }
3819 z = zf;
3820 None
3821 };
3822 // DECODE-EXACT FFN projections: force MMVQ for gate/up/down at any T to match the
3823 // T=1 decode FP accumulation order. At T>=5 the generic matmul/matmul_pre falls to dp4a
3824 // (128-thread, different FP sum order). At T=2-4 the batched MMVQ is already bit-identical.
3825 let ffn_out = match &layer.ffn {
3826 crate::hybrid::Ffn::Dense {
3827 ffn_gate,
3828 ffn_up,
3829 ffn_down,
3830 } => {
3831 let n_ff = ffn_gate.out_features();
3832 if let Some((zq, zd)) = z_q8.as_ref() {
3833 // FUSED CHAIN (fix 2): pre-quantized z feeds the projections; the SwiGLU
3834 // epilogue emits act pre-quantized for ffn_down (silu_mul_scaled_q8_1,
3835 // bit-identical to silu_mul + quantize — kernel-check-pinned) with the
3836 // NVFP4 macro-scales folded (deferred-scale dual: y*s inline == the
3837 // scale_inplace store, value-exact) — the exact m=1 decode epilogue
3838 // structure at nrows=t.
3839 let pair = match e.matmul_decode_exact_dual_pre(ffn_gate, ffn_up, zq, zd, t)? {
3840 Some(((g, gs), (u, us))) => Some((g, gs, u, us)),
3841 None => None,
3842 };
3843 let (gate, gs, up, us) = match pair {
3844 Some(x4) => x4,
3845 None => (
3846 e.matmul_decode_exact_pre(ffn_gate, zq, zd, t)?,
3847 1.0, // scale already applied inside _pre
3848 e.matmul_decode_exact_pre(ffn_up, zq, zd, t)?,
3849 1.0,
3850 ),
3851 };
3852 if e.uses_q8_1_fast(ffn_down) {
3853 let (aq, ad) = e.silu_mul_scaled_q8_1(&gate, &up, gs, us, t * n_ff)?;
3854 e.matmul_decode_exact_pre(ffn_down, &aq, &ad, t)?
3855 } else {
3856 let mut act = vbuf(e, t * n_ff)?;
3857 e.silu_mul_scaled(&gate, &up, gs, us, &mut act, t * n_ff)?;
3858 e.matmul_decode_exact(ffn_down, &act, t)?
3859 }
3860 } else {
3861 // UNFUSED (pre-fix) chain — MoE-adjacent/M3/off-fast layers, unchanged.
3862 // DUAL gate+up batched twin (lane/verify-economics, 2026-08-02): one launch
3863 // for the pair at t=2..8 — bit-identical per (tensor,token,row) to the two
3864 // singles (kernel-check pins bitwise; MEMRA_SPEC_DUAL_T=0 reverts). None
3865 // (non-NVFP4 / t outside the tier / seam off) -> the two singles, unchanged.
3866 let (gate, up) = match e.matmul_decode_exact_dual(ffn_gate, ffn_up, &z, t)? {
3867 Some(pair) => pair,
3868 None => (
3869 e.matmul_decode_exact(ffn_gate, &z, t)?,
3870 e.matmul_decode_exact(ffn_up, &z, t)?,
3871 ),
3872 };
3873 let mut act = vbuf(e, t * n_ff)?; // fully written by ffn_act_lim
3874 Self::ffn_act_lim(e, &self.cfg, &gate, &up, 1.0, 1.0, dense_lim,
3875 &mut act, t * n_ff)?;
3876 e.matmul_decode_exact(ffn_down, &act, t)?
3877 }
3878 }
3879 crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il(e, m, &z, t, il as u16)?,
3880 };
3881 // CROSS-LAYER fusion: defer this layer's post-FFN residual add — the next layer's
3882 // fused-q8 attn norm folds it in (add_rms_norm_q8_1 == add; rms_norm; quantize,
3883 // kernel-check-pinned at nrows=T). Non-fused next layers add explicitly above.
3884 pending = Some((x1, ffn_out));
3885 }
3886 // RANGE's final add (no next norm INSIDE the range to fuse with; for the
3887 // whole-trunk call that is the last layer, whose next norm is output_norm — f32-out).
3888 if let Some((x1p, f1p)) = pending.take() {
3889 let mut x2 = vbuf(e, t * n_embd)?; // fully written by add
3890 e.add(&x1p, &f1p, &mut x2, t * n_embd)?;
3891 x = x2;
3892 }
3893 Ok(x)
3894 }
3895 /// BATCHED linear-attn verify (T=K+1): the whole layer in ~10 launches instead of T x the
3896 /// T=1 decode chain (T x ~12 launches + T weight reads of the four projections). The GDN
3897 /// recurrence itself is inherently sequential — gdn_scan_s128 runs its internal t-loop with
3898 /// the SAME per-token math as chained T=1 calls (bit-identical state evolution); everything
3899 /// around it (projections, conv, prep, gated norm, out-proj) batches. Advances conv ring +
3900 /// ssm state exactly like T sequential decode steps.
3901 /// `want_stash`: additionally RETAIN the gdn-scan inputs (pure buffer keep-alives, zero extra
3902 /// kernels) so a partial accept can rebuild the state after any column prefix (REPLAY-FREE).
3903 #[allow(clippy::too_many_arguments)]
3904 fn linear_attn_verify_t(
3905 &self,
3906 e: &Engine,
3907 la: &LinearAttnLayer,
3908 h: &CudaSlice<f32>,
3909 h_q8: Option<(&CudaSlice<i8>, &CudaSlice<f32>)>,
3910 t: usize,
3911 cache: &mut Cache,
3912 il: usize,
3913 want_stash: bool,
3914 ) -> Result<(CudaSlice<f32>, Option<GdnStash>), Box<dyn std::error::Error>> {
3915 let cfg = &self.cfg;
3916 let ssm = cfg.ssm.as_ref().unwrap();
3917 let d_state = ssm.state_size as usize;
3918 let num_k = ssm.group_count as usize;
3919 let num_v = ssm.time_step_rank as usize;
3920 let d_conv = ssm.conv_kernel as usize;
3921 let key_dim = d_state * num_k;
3922 let conv_dim = key_dim * 2 + d_state * num_v;
3923 let eps = cfg.rms_eps;
3924 let scale = 1.0 / (d_state as f32).sqrt();
3925
3926 // DECODE-EXACT projections: matmul_decode_exact forces the MMVQ (warp-per-row, 32-thread)
3927 // accumulation order for EVERY m, matching the T=1 decode path bit-for-bit. The generic
3928 // `matmul` at m>=5 falls to dp4a (128-thread, two-level reduce) which has a different FP
3929 // sum order — ULP differences propagate through gdn_scan and flip argmax on the 27B.
3930 // Q8 TRUNK-FUSION at T=1 (35B: wqkv+wqkv_gate both Q8_0): one fused2 launch, bit-identical
3931 // per (tensor,row) to the two m=1 MMVQ dispatches below — decode-exact contract holds.
3932 // VERIFY-TIER TRUNK FUSION (MEMRA_SPEC_FUSED_T, t=2-4): quantize h ONCE for every
3933 // fused-eligible same-input Q8_0 pair of this layer (35B wqkv+wqkv_gate; 9B
3934 // ssm_beta+ssm_alpha) — each fused2 batched launch then replaces two decode-exact
3935 // calls (each of which re-quantizes the same h + runs its own _b2/_b4 launch).
3936 // Bit-identical per (tensor,token,row) — see spec_fused_t().
3937 // BATCHED EPILOGUE RE-FUSE (lane/vt-fixes fix 2): `h_q8` = the attn-input norm emitted
3938 // directly as q8_1 by the caller's fused rms_norm_q8_1 (bit-identical to the unfused
3939 // chain, kernel-check-pinned). When present it REPLACES the standalone quantize below
3940 // and feeds every projection; the caller guaranteed all four input projections are
3941 // q8_1-fast. When absent, the old shared-quantize (fused-t window) stands.
3942 let h_q8_t = if h_q8.is_none()
3943 && spec_fused_t()
3944 && (2..=4).contains(&t)
3945 && ((e.uses_q8_1_fast(&la.wqkv) && e.uses_q8_1_fast(&la.wqkv_gate))
3946 || (e.uses_q8_1_fast(&la.ssm_beta) && e.uses_q8_1_fast(&la.ssm_alpha)))
3947 {
3948 Some(e.quantize_q8_1(h, t, cfg.n_embd as usize)?)
3949 } else {
3950 None
3951 };
3952 // one view: the caller's fused-norm q8 or this fn's own shared quantize.
3953 let hq8_any: Option<(&CudaSlice<i8>, &CudaSlice<f32>)> =
3954 h_q8.or(h_q8_t.as_ref().map(|(q, d)| (q, d)));
3955 let (qkv_mixed, z) = {
3956 let mut fused = None;
3957 if t == 1 && e.uses_q8_1_fast(&la.wqkv) && e.uses_q8_1_fast(&la.wqkv_gate) {
3958 let (hq, hd) = e.quantize_q8_1(h, 1, cfg.n_embd as usize)?;
3959 fused = e.matmul_q8_fused2(&la.wqkv, &la.wqkv_gate, &hq, &hd)?;
3960 } else if let Some((hq, hd)) = hq8_any {
3961 if spec_fused_t() && (2..=4).contains(&t) {
3962 fused = e.matmul_q8_fused2_t(&la.wqkv, &la.wqkv_gate, hq, hd, t)?;
3963 }
3964 }
3965 match (fused, hq8_any) {
3966 (Some(pair), _) => pair,
3967 (None, Some((hq, hd))) if h_q8.is_some() => (
3968 e.matmul_decode_exact_pre(&la.wqkv, hq, hd, t)?,
3969 e.matmul_decode_exact_pre(&la.wqkv_gate, hq, hd, t)?,
3970 ),
3971 (None, _) => (
3972 e.matmul_decode_exact(&la.wqkv, h, t)?,
3973 e.matmul_decode_exact(&la.wqkv_gate, h, t)?,
3974 ),
3975 }
3976 };
3977 // beta+alpha DUAL at T=1 (75% of p3 rounds run T=1 verify — p-min chain cuts): the dual
3978 // mr2 kernel is bit-identical per element to the m=1 MMVQ matmul_decode_exact dispatches
3979 // (same warp-per-row body, blockIdx.y picks the weight), so the decode-exact contract
3980 // holds; the run-spec battery is the arbiter. T>1 keeps the per-tensor decode-exact path.
3981 let (beta_raw, alpha) = if t == 1 {
3982 let (hq, hd) = e.quantize_q8_1(h, 1, cfg.n_embd as usize)?;
3983 match e.matmul_pre_dual_noscale(&la.ssm_beta, &la.ssm_alpha, &hq, &hd, 1)? {
3984 Some(((mut b, bs), (mut a, as_))) => {
3985 if bs != 1.0 {
3986 e.scale_inplace(&mut b, bs, la.ssm_beta.out_features())?;
3987 }
3988 if as_ != 1.0 {
3989 e.scale_inplace(&mut a, as_, la.ssm_alpha.out_features())?;
3990 }
3991 (b, a)
3992 }
3993 // Q8_0 fused2 twin (9B stores beta/alpha as Q8_0): DISPATCH-MIRRORS the eager
3994 // decode's beta_alpha closure — the fused body is qmatvec_q8_0_mmvq verbatim,
3995 // bit-identical per row (kernel-check rel=0.00e0 gate), so decode==verify holds.
3996 None => match e.matmul_q8_fused2(&la.ssm_beta, &la.ssm_alpha, &hq, &hd)? {
3997 Some((b, a)) => (b, a),
3998 None => (
3999 e.matmul_decode_exact(&la.ssm_beta, h, 1)?,
4000 e.matmul_decode_exact(&la.ssm_alpha, h, 1)?,
4001 ),
4002 },
4003 }
4004 } else {
4005 // fused-t twin (9B stores beta/alpha as Q8_0): same shared-quantize + one launch
4006 // contract as the wqkv pair above; 35B beta/alpha are Float -> None -> fallback.
4007 let mut nvfp4_fused = None;
4008 let mut q8_fused = None;
4009 if let Some((hq, hd)) = hq8_any {
4010 if t == 3 && std::env::var("MEMRA_NVFP4_AUX_DUAL").as_deref() != Ok("0") {
4011 nvfp4_fused = e.matmul_decode_exact_dual_pre(
4012 &la.ssm_beta,
4013 &la.ssm_alpha,
4014 hq,
4015 hd,
4016 t,
4017 )?;
4018 if nvfp4_fused.is_some() && std::env::var("MEMRA_DEBUG").is_ok() {
4019 static ONCE: std::sync::Once = std::sync::Once::new();
4020 ONCE.call_once(|| eprintln!(
4021 "[memra] NVFP4 beta+alpha batched aux dual ENGAGED (t={t})"
4022 ));
4023 }
4024 }
4025 if nvfp4_fused.is_none() && spec_fused_t() && (2..=4).contains(&t) {
4026 q8_fused = e.matmul_q8_fused2_t(&la.ssm_beta, &la.ssm_alpha, hq, hd, t)?;
4027 }
4028 }
4029 if let Some(((mut b, bs), (mut a, as_))) = nvfp4_fused {
4030 if bs != 1.0 {
4031 e.scale_inplace(&mut b, bs, t * la.ssm_beta.out_features())?;
4032 }
4033 if as_ != 1.0 {
4034 e.scale_inplace(&mut a, as_, t * la.ssm_alpha.out_features())?;
4035 }
4036 (b, a)
4037 } else if let Some(pair) = q8_fused {
4038 pair
4039 } else { match hq8_any {
4040 Some((hq, hd)) if h_q8.is_some() => (
4041 e.matmul_decode_exact_pre(&la.ssm_beta, hq, hd, t)?,
4042 e.matmul_decode_exact_pre(&la.ssm_alpha, hq, hd, t)?,
4043 ),
4044 _ => (
4045 e.matmul_decode_exact(&la.ssm_beta, h, t)?,
4046 e.matmul_decode_exact(&la.ssm_alpha, h, t)?,
4047 ),
4048 }}
4049 };
4050
4051 // conv with CARRIED state + ring roll (T >= pad rides the input-column update kernel;
4052 // T < pad — the MEMRA_SPEC_M2 t=2 arm — rolls via the pure-copy ring rebuild).
4053 let rl = cache.recur[il].as_mut().unwrap();
4054 let mut conv_out = e.uninit(conv_dim * t)?;
4055 e.ssm_conv1d_tm_state(
4056 &qkv_mixed,
4057 &mut rl.conv_state,
4058 la.ssm_conv1d.float_data(),
4059 &mut conv_out,
4060 conv_dim,
4061 t,
4062 d_conv,
4063 )?;
4064
4065 // GDN prep via the prefill kernels (repack + L2 + sigmoid + glog), T-wide.
4066 let mut q_g = e.uninit(d_state * num_v * t)?;
4067 let mut k_g = e.uninit(d_state * num_v * t)?;
4068 let mut v_g = e.uninit(d_state * num_v * t)?;
4069 e.qkv_to_gdn_repack(
4070 &conv_out, &mut q_g, &mut k_g, &mut v_g, d_state, num_v, num_k, key_dim, t,
4071 )?;
4072 let mut q_l2 = e.uninit(d_state * num_v * t)?;
4073 e.l2_norm_decode(&q_g, &mut q_l2, d_state, num_v * t, eps)?;
4074 let mut k_l2 = e.uninit(d_state * num_v * t)?;
4075 e.l2_norm_decode(&k_g, &mut k_l2, d_state, num_v * t, eps)?;
4076 let mut beta = e.uninit(t * num_v)?;
4077 e.sigmoid(&beta_raw, &mut beta, t * num_v)?;
4078 let mut g_log = e.uninit(t * num_v)?;
4079 e.gdn_glog(
4080 &alpha,
4081 la.ssm_dt.float_data(),
4082 la.ssm_a.float_data(),
4083 &mut g_log,
4084 num_v,
4085 t,
4086 )?;
4087
4088 // ONE gdn_scan over T tokens from the carried state (internal sequential loop ==
4089 // T chained T=1 steps). Ping-pong the resident buffers like eager decode.
4090 let mut o = e.uninit(d_state * num_v * t)?;
4091 {
4092 let crate::cache::RecurLayer {
4093 ssm_state,
4094 ssm_state_alt,
4095 ..
4096 } = rl;
4097 e.gdn_scan_s128(
4098 &q_l2,
4099 &k_l2,
4100 &v_g,
4101 &g_log,
4102 &beta,
4103 ssm_state,
4104 ssm_state_alt,
4105 &mut o,
4106 num_v,
4107 t,
4108 scale,
4109 )?;
4110 }
4111 std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
4112
4113 // gated RMSNorm + out projection, T-wide. FUSED-QUANTIZE ARM (lane/vt-fixes fix 2,
4114 // mirroring the T=1 decode's launch-arc form): when ssm_out rides the q8_1 fast path,
4115 // emit q8_1 straight from the gated norm at nrows=num_v*t (row-indexed kernel, the
4116 // T-wide launch is the per-row program; kernel-check pins bit-identity vs
4117 // gated_rmsnorm -> quantize_q8_1 at T=1 and T=5) and feed the decode-exact dispatch
4118 // pre-quantized — one launch replaces norm + quantize. Fallback = the f32 chain.
4119 let out = if e.uses_q8_1_fast(&la.ssm_out) {
4120 let (gq, gd) =
4121 e.gated_rmsnorm_q8_1(&o, la.ssm_norm.float_data(), &z, d_state, num_v * t, eps)?;
4122 e.matmul_decode_exact_pre(&la.ssm_out, &gq, &gd, t)?
4123 } else {
4124 let mut gn = e.uninit(d_state * num_v * t)?;
4125 e.gated_rmsnorm(
4126 &o,
4127 la.ssm_norm.float_data(),
4128 &z,
4129 &mut gn,
4130 d_state,
4131 num_v * t,
4132 eps,
4133 )?;
4134 // DECODE-EXACT out-projection: same MMVQ path as the T=1 decode (ssm_out at m>=5
4135 // would fall to dp4a with a different FP reduction order — same class of bug as
4136 // the input projs).
4137 e.matmul_decode_exact(&la.ssm_out, &gn, t)?
4138 };
4139 let stash = if want_stash {
4140 Some(GdnStash {
4141 qkv_mixed,
4142 q_l2,
4143 k_l2,
4144 v_g,
4145 g_log,
4146 beta,
4147 })
4148 } else {
4149 None
4150 };
4151 Ok((out, stash))
4152 }
4153
4154 /// REPLAY-FREE partial-accept commit (2026-07-03): make the cache state == "committed through
4155 /// the first `j` verify columns" WITHOUT the legacy rollback + duplicate trunk replay.
4156 /// - Full-attn KV: truncate len to snapshot + j. The verify's appended rows for those columns
4157 /// are bit-identical to what an eager T=1 chain writes (the decode-exact contract the
4158 /// verify-probe gates), so keeping them == replaying them.
4159 /// - Linear layers, batched path: rebuild the conv ring by PURE COPIES (ring holds raw input
4160 /// columns) and the ssm state by a prefix re-run of the SAME gdn_scan kernel (t=j) from the
4161 /// snapshot state over the stash's identical inputs — the kernel's t-loop carries state in
4162 /// registers and writes it once at the end, so iterations 0..j-1 are independent of T:
4163 /// bit-identical to the verify's own state after j tokens == the eager chain state.
4164 /// - Linear layers, per-column path: restore the cloned actual state after column j-1.
4165 /// Caller guarantees 1 <= j <= t-1 (j==0 rounds take the legacy rollback; j==t is full accept).
4166 fn commit_verified_prefix(
4167 &self,
4168 e: &Engine,
4169 cache: &mut Cache,
4170 snap: &crate::cache::CacheSnapshot,
4171 ckpt: &VerifyCkpt,
4172 j: usize,
4173 kv_lens_done: bool,
4174 dev_j: Option<(&CudaSlice<u32>, usize, usize)>,
4175 ) -> Result<(), Box<dyn std::error::Error>> {
4176 let cfg = &self.cfg;
4177 let ssm = cfg.ssm.as_ref().unwrap();
4178 let d_state = ssm.state_size as usize;
4179 let num_k = ssm.group_count as usize;
4180 let num_v = ssm.time_step_rank as usize;
4181 let d_conv = ssm.conv_kernel as usize;
4182 let conv_dim = d_state * num_k * 2 + d_state * num_v;
4183 let scale = 1.0 / (d_state as f32).sqrt();
4184 for il in 0..self.layers.len() {
4185 if let (Some(kvl), Some(saved)) = (cache.kv[il].as_mut(), snap.kv_len[il]) {
4186 kvl.len = saved + j;
4187 // devacc 3a: spec_rollback_kv already wrote len_d on-device (same value).
4188 if !kv_lens_done {
4189 e.set_i32_one(&mut kvl.len_d, kvl.len as i32)?;
4190 }
4191 }
4192 if let Some(rl) = cache.recur[il].as_mut() {
4193 if let Some(st) = &ckpt.gdn[il] {
4194 let ring_old = snap.conv[il].as_ref().expect("snapshot missing conv");
4195 let state_in = snap.ssm[il].as_ref().expect("snapshot missing ssm");
4196 if let Some((acc, base, t_v)) = dev_j {
4197 // 3b: j read on-device (_dc twins, same bodies; full accept early-exits).
4198 e.ssm_conv_ring_rebuild_dc(
4199 &st.qkv_mixed,
4200 ring_old,
4201 &mut rl.conv_state,
4202 conv_dim,
4203 acc,
4204 base,
4205 t_v,
4206 d_conv,
4207 )?;
4208 let mut o = e.uninit(d_state * num_v * j.max(1))?;
4209 e.gdn_scan_s128_dc(
4210 &st.q_l2,
4211 &st.k_l2,
4212 &st.v_g,
4213 &st.g_log,
4214 &st.beta,
4215 state_in,
4216 &mut rl.ssm_state,
4217 &mut o,
4218 num_v,
4219 acc,
4220 base,
4221 t_v,
4222 scale,
4223 )?;
4224 } else {
4225 e.ssm_conv_ring_rebuild(
4226 &st.qkv_mixed,
4227 ring_old,
4228 &mut rl.conv_state,
4229 conv_dim,
4230 j,
4231 d_conv,
4232 )?;
4233 let mut o = e.uninit(d_state * num_v * j)?; // scan output, discarded
4234 e.gdn_scan_s128(
4235 &st.q_l2,
4236 &st.k_l2,
4237 &st.v_g,
4238 &st.g_log,
4239 &st.beta,
4240 state_in,
4241 &mut rl.ssm_state,
4242 &mut o,
4243 num_v,
4244 j,
4245 scale,
4246 )?;
4247 }
4248 } else if let Some(cols) = &ckpt.cols[il] {
4249 let (c, s) = &cols[j - 1];
4250 e.copy_into(&mut rl.conv_state, 0, c, c.len())?;
4251 e.copy_into(&mut rl.ssm_state, 0, s, s.len())?;
4252 } else {
4253 return Err(
4254 "commit_verified_prefix: verify ckpt missing for linear layer".into(),
4255 );
4256 }
4257 }
4258 }
4259 cache.pos = snap.pos + j;
4260 Ok(())
4261 }
4262
4263 /// ROUND-STREAM: recur restore with device-j (the _dc twins; full accept early-exits
4264 /// in-kernel). Requires the batched-linear stash on every linear layer (stream gate).
4265 fn commit_verified_prefix_stream(
4266 &self,
4267 e: &Engine,
4268 cache: &mut Cache,
4269 snap: &crate::cache::CacheSnapshot,
4270 ckpt: &VerifyCkpt,
4271 acc: &CudaSlice<u32>,
4272 base: usize,
4273 t_v: usize,
4274 ) -> Result<(), Box<dyn std::error::Error>> {
4275 let cfg = &self.cfg;
4276 let ssm = cfg.ssm.as_ref().unwrap();
4277 let d_state = ssm.state_size as usize;
4278 let num_k = ssm.group_count as usize;
4279 let num_v = ssm.time_step_rank as usize;
4280 let d_conv = ssm.conv_kernel as usize;
4281 let conv_dim = d_state * num_k * 2 + d_state * num_v;
4282 let scale = 1.0 / (d_state as f32).sqrt();
4283 for il in 0..self.layers.len() {
4284 if let Some(rl) = cache.recur[il].as_mut() {
4285 let st = ckpt.gdn[il]
4286 .as_ref()
4287 .ok_or("stream restore: batched-linear stash missing")?;
4288 let ring_old = snap.conv[il].as_ref().expect("snapshot missing conv");
4289 let state_in = snap.ssm[il].as_ref().expect("snapshot missing ssm");
4290 e.ssm_conv_ring_rebuild_dc(
4291 &st.qkv_mixed,
4292 ring_old,
4293 &mut rl.conv_state,
4294 conv_dim,
4295 acc,
4296 base,
4297 t_v,
4298 d_conv,
4299 )?;
4300 let mut o = e.uninit(d_state * num_v * t_v)?;
4301 e.gdn_scan_s128_dc(
4302 &st.q_l2,
4303 &st.k_l2,
4304 &st.v_g,
4305 &st.g_log,
4306 &st.beta,
4307 state_in,
4308 &mut rl.ssm_state,
4309 &mut o,
4310 num_v,
4311 acc,
4312 base,
4313 t_v,
4314 scale,
4315 )?;
4316 }
4317 }
4318 Ok(())
4319 }
4320
4321 /// EAGLE3 aux-capturing verify forward over `tokens` (T) — mirrors `decode_step_t_h` exactly
4322 /// (same KV append, same causal verify, same recur advance) but ALSO clones the aux residual-
4323 /// stream hiddens (blocks in `aux_layers`) for TWO columns: the LAST column (always) and the
4324 /// optional `pred_col` (the EAGLE seed = bonus's predecessor). Returns
4325 /// (all_T_logits host, last_col_aux, pred_col_aux?). Used by the EAGLE3 orchestrator's commit.
4326 pub fn decode_step_t_aux2(
4327 &self,
4328 e: &Engine,
4329 tokens: &[u32],
4330 pos0: usize,
4331 cache: &mut Cache,
4332 aux_layers: &[usize],
4333 pred_col: Option<usize>,
4334 ) -> Result<
4335 (Vec<f32>, Vec<CudaSlice<f32>>, Option<Vec<CudaSlice<f32>>>),
4336 Box<dyn std::error::Error>,
4337 > {
4338 let cfg = &self.cfg;
4339 let n_embd = cfg.n_embd as usize;
4340 let eps = cfg.rms_eps;
4341 let t = tokens.len();
4342 let pos_vec: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
4343 let pos_d = e.htod_i32(&pos_vec)?;
4344 let mut x = e.htod(&self.embd.gather(n_embd, tokens))?;
4345 let mut aux_last: Vec<CudaSlice<f32>> = Vec::with_capacity(aux_layers.len());
4346 let mut aux_pred: Vec<CudaSlice<f32>> = Vec::new();
4347 let want_pred = pred_col.is_some();
4348
4349 for (il, layer) in self.layers.iter().enumerate() {
4350 // DISPATCH-MIRRORED norms (FP-order lesson #8) — see decode_step_t_h_emb.
4351 let mixer_fast = self.mixer_in_q8_1_fast(e, &layer.mixer);
4352 let norm_fused = std::env::var("MEMRA_NO_FUSE_NORMQ").is_err() && mixer_fast;
4353 let mut h = vbuf(e, t * n_embd)?; // fully written by either rms_norm arm
4354 if norm_fused {
4355 e.rms_norm_decode(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
4356 } else {
4357 e.rms_norm(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
4358 }
4359 let mixed = match &layer.mixer {
4360 Mixer::Full(fa) => {
4361 self.full_attn_verify(e, fa, &h, None, &pos_d, t, cache, il, None)?
4362 }
4363 Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
4364 Mixer::Linear(la) => {
4365 let mut out = e.zeros(t * n_embd)?;
4366 for col in 0..t {
4367 let mut h_col = e.zeros(n_embd)?;
4368 let src = h.slice(col * n_embd..(col + 1) * n_embd);
4369 e.copy_view_into(&mut h_col, 0, &src, n_embd)?;
4370 let m_col = self.linear_attn_decode(e, la, &h_col, cache, il)?;
4371 e.copy_into(&mut out, col * n_embd, &m_col, n_embd)?;
4372 }
4373 out
4374 }
4375 };
4376 let ffn_fuse = match &layer.ffn {
4377 crate::hybrid::Ffn::Dense {
4378 ffn_gate, ffn_up, ..
4379 } => {
4380 std::env::var("MEMRA_NO_FUSE_NORMQ").is_err()
4381 && e.uses_q8_1_fast(ffn_gate)
4382 && e.uses_q8_1_fast(ffn_up)
4383 }
4384 crate::hybrid::Ffn::Moe(_) => false,
4385 };
4386 let mut x1 = vbuf(e, t * n_embd)?; // fully written by add / add_rms_norm
4387 let mut z = vbuf(e, t * n_embd)?; // fully written by rms_norm_decode / add_rms_norm
4388 if ffn_fuse {
4389 e.add(&x, &mixed, &mut x1, t * n_embd)?;
4390 e.rms_norm_decode(
4391 &x1,
4392 layer.post_attn_norm.float_data(),
4393 &mut z,
4394 n_embd,
4395 t,
4396 eps,
4397 )?;
4398 } else {
4399 e.add_rms_norm(
4400 &x,
4401 &mixed,
4402 layer.post_attn_norm.float_data(),
4403 &mut x1,
4404 &mut z,
4405 n_embd,
4406 t,
4407 eps,
4408 )?;
4409 }
4410 let ffn_out = match &layer.ffn {
4411 crate::hybrid::Ffn::Dense {
4412 ffn_gate,
4413 ffn_up,
4414 ffn_down,
4415 } => {
4416 let n_ff = ffn_gate.out_features();
4417 let gate = e.matmul_decode_exact(ffn_gate, &z, t)?;
4418 let up = e.matmul_decode_exact(ffn_up, &z, t)?;
4419 let mut act = vbuf(e, t * n_ff)?; // fully written by ffn_act_lim
4420 // dense FFN clamp = the SHEXP array (upstream build_ffn serves both).
4421 Self::ffn_act_lim(e, &self.cfg, &gate, &up, 1.0, 1.0,
4422 self.cfg.clamp_shexp_at(il as u32), &mut act, t * n_ff)?;
4423 e.matmul_decode_exact(ffn_down, &act, t)?
4424 }
4425 crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il(e, m, &z, t, il as u16)?,
4426 };
4427 let mut x2 = vbuf(e, t * n_embd)?; // fully written by add
4428 e.add(&x1, &ffn_out, &mut x2, t * n_embd)?;
4429 if aux_layers.contains(&il) {
4430 let mut a = e.zeros(n_embd)?;
4431 e.copy_view_into(&mut a, 0, &x2.slice((t - 1) * n_embd..t * n_embd), n_embd)?;
4432 aux_last.push(a);
4433 if let Some(pc) = pred_col {
4434 let mut ap = e.zeros(n_embd)?;
4435 e.copy_view_into(
4436 &mut ap,
4437 0,
4438 &x2.slice(pc * n_embd..(pc + 1) * n_embd),
4439 n_embd,
4440 )?;
4441 aux_pred.push(ap);
4442 }
4443 }
4444 x = x2;
4445 }
4446 let mut hn = vbuf(e, t * n_embd)?; // fully written by rms_norm_decode
4447 e.rms_norm_decode(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
4448 let logits = e.matmul_decode_exact(&self.output, &hn, t)?;
4449 let host = e.dtoh(&logits)?;
4450 cache.pos += t;
4451 Ok((
4452 host,
4453 aux_last,
4454 if want_pred { Some(aux_pred) } else { None },
4455 ))
4456 }
4457
4458 /// step35 SPEC-VERIFY attention over T query tokens — a per-row REPLAY of the eager
4459 /// `step35_decode_attn`.
4460 ///
4461 /// WHY A REPLAY AND NOT A BATCHED TWIN. The verify's whole job is to be bit-identical to what
4462 /// the eager decode would have computed for the same tokens; that is what makes greedy spec
4463 /// decode exact (run-spec asserts token identity for K=1..8). Every other verify arm in this
4464 /// file earns that identity by carefully mirroring dispatch (`matmul_decode_exact` to force
4465 /// MMVQ at any m, per-layer `ffn_fuse` mirroring, per-row `fa_decode` key bounds). step35
4466 /// stacks FOUR more per-layer degrees of freedom on top of that — per-layer `n_head`
4467 /// (64 full / 96 SWA), per-layer rotary width (64 full / 128 SWA), per-layer rope base, and a
4468 /// SEPARATE `attn_gate` tensor whose projection shares the attn-normed input — and its SWA
4469 /// layers attend through a token-OFFSET view whose offset is a function of the ABSOLUTE
4470 /// position of each query row. A batched twin would have to reproduce all of that AND the
4471 /// per-row offset in one launch; the offset alone rules out the existing rows kernels (they
4472 /// take one `base_len`, not a per-row offset).
4473 ///
4474 /// So this arm calls the eager path itself, once per row, on the same cache. Identity is then
4475 /// true BY CONSTRUCTION rather than by mirroring: row r runs exactly the kernel sequence that
4476 /// eager decode step r runs (same projections, same q8_1 fusion decision, same append, same
4477 /// view arithmetic, same `fa_decode_kvmod`, same gate), because it IS that code. Cost: T x the
4478 /// eager decode mixer instead of one batched pass — the same trade the generic arm's `else`
4479 /// per-row loop already accepts when `fa_rows_eligible` says no. Correctness first; a batched
4480 /// step35 twin is a perf lane's job and must be gated against this arm.
4481 ///
4482 /// The `h_q8` pre-quantized pair from the caller's fused norm is NOT forwarded: it is a
4483 /// T-row buffer and `step35_decode_attn`'s `pre_q` contract is one row. Instead each row's
4484 /// f32 `h` slice is handed over and the callee re-derives its own q8_1 exactly as eager decode
4485 /// does (`quantize_q8_1(h, 1, n_embd)`) — which is the dispatch being mirrored. Callers that
4486 /// took the fused arm therefore MUST still pass a live `h`; `step35_verify` asserts that.
4487 #[allow(clippy::too_many_arguments)]
4488 fn step35_verify(
4489 &self,
4490 e: &Engine,
4491 fa: &FullAttnLayer,
4492 h: &CudaSlice<f32>,
4493 h_q8: Option<(&CudaSlice<i8>, &CudaSlice<f32>)>,
4494 t: usize,
4495 cache: &mut Cache,
4496 il: usize,
4497 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4498 let n_embd = self.cfg.n_embd as usize;
4499 // The fused (h-less) attn-norm arm hands `h` as a zero-length placeholder. This arm needs
4500 // the f32 rows, so the caller must not take that lever for step35 — enforced at the call
4501 // site by the `Mixer::Full(_) if self.cfg.step35.is_some() => false` arm of
4502 // `lin_q8_only` in `decode_step_t_core_stream`, and asserted here so a future caller
4503 // cannot regress it into silently reading an empty buffer.
4504 assert_eq!(
4505 h.len(),
4506 t * n_embd,
4507 "step35_verify needs the f32 attn-normed rows ([t*n_embd]); the caller took the \
4508 fused q8-only norm arm (h_q8={}) — step35 must stay on the unfused arm",
4509 h_q8.is_some()
4510 );
4511 // ROW WIDTH IS n_embd, NOT n_head*head_dim: `step35_decode_attn` returns the mixer output
4512 // AFTER `wo`, so a row is [n_embd] — the same contract the generic arm's
4513 // `matmul_decode_exact(&fa.wo, &attn_g, t)` return has. Sizing this buffer from the
4514 // per-layer head geometry (8192 on full-attn, 12288 on SWA) instead overran the row on the
4515 // FIRST copy and panicked inside `copy_into`'s `CudaView::slice` unwrap
4516 // (raw/mtp-bt-20260806T212127Z.log frames 12-13).
4517 let mut out = vbuf(e, t * n_embd)?; // each row fully written by the copy below
4518 for r in 0..t {
4519 // Absolute position of this query row. `cache.pos` is the committed length at round
4520 // start and every row before r has already been appended by this loop, so the r-th
4521 // verify token sits at cache.pos + r — the same position eager decode would give it.
4522 let pos_d = e.htod_i32(&[(cache.pos + r) as i32])?;
4523 let mut h_row = vbuf(e, n_embd)?; // fully written by copy_view_into
4524 e.copy_view_into(&mut h_row, 0, &h.slice(r * n_embd..(r + 1) * n_embd), n_embd)?;
4525 // THE eager decode mixer: appends this row's K/V at kvl.len, advances it, then
4526 // attends over the (SWA-offset) view. Post-`wo`, same contract as this fn returns.
4527 let o = self.step35_decode_attn(e, fa, il, &h_row, None, &pos_d, cache)?;
4528 debug_assert_eq!(o.len(), n_embd, "step35_decode_attn returns post-wo [n_embd]");
4529 e.copy_into(&mut out, r * n_embd, &o, n_embd)?;
4530 }
4531 Ok(out)
4532 }
4533
4534 /// Full-attention mixer over T query tokens with a GROWING resident KV (verify path, §D.3).
4535 /// Appends the T new K/V columns to cache.kv[il] then attends causally over [0..len) via
4536 /// fa_prefill. Token-major [T, kv_dim] projection layout == cache row layout (single copy).
4537 #[allow(clippy::too_many_arguments)]
4538 fn full_attn_verify(
4539 &self,
4540 e: &Engine,
4541 fa: &FullAttnLayer,
4542 h: &CudaSlice<f32>,
4543 h_q8: Option<(&CudaSlice<i8>, &CudaSlice<f32>)>,
4544 pos_d: &CudaSlice<i32>,
4545 t: usize,
4546 cache: &mut Cache,
4547 il: usize,
4548 stream_ctr: Option<&CudaSlice<i32>>,
4549 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4550 // step35: the generic geometry below is wrong for this arch (per-layer n_head, partial
4551 // per-layer rope, the SWA offset view, and a SEPARATE head-wise gate tensor), so it takes
4552 // its own arm. A verify that silently computes different attention than decode defeats the
4553 // whole self-consistency gate, so the arm is a per-row REPLAY of `step35_decode_attn`
4554 // rather than a batched twin — see `step35_verify` for why that is the exactness-correct
4555 // shape and not laziness.
4556 if self.cfg.step35.is_some() {
4557 if stream_ctr.is_some() {
4558 return Err("step35 has no ROUND-STREAM verify arm (the device-counter _dc twins \
4559 cannot express the SWA offset KV view; same root cause as the dc \
4560 decode refusal) — run spec without the stream arm".into());
4561 }
4562 return self.step35_verify(e, fa, h, h_q8, t, cache, il);
4563 }
4564 let cfg = &self.cfg;
4565 let geometry = cfg.full_attention_geometry_at(il as u32);
4566 let n_head = geometry.n_head as usize;
4567 let n_head_kv = geometry.n_head_kv as usize;
4568 let head_dim = geometry.head_dim_k as usize;
4569 let eps = cfg.rms_eps;
4570 let scale = geometry.attention_scale();
4571 let n_embd = cfg.n_embd as usize;
4572
4573 // DECODE-EXACT Q/K/V projections: matmul_decode_exact forces the MMVQ (warp-per-row) path
4574 // for every m, matching the T=1 decode's FP accumulation order. matmul_pre at m>=5 would
4575 // fall to dp4a (128-thread, two-level reduce) with a different FP sum order.
4576 // Q8 TRUNK-FUSION at T=1: DISPATCH-MIRRORS the eager decode's fused3 (bit-identical body).
4577 // BATCHED EPILOGUE RE-FUSE (lane/vt-fixes fix 2): `h_q8` = the attn-input norm's q8_1
4578 // form emitted by the fused rms_norm_q8_1 (bit-identical to rms_norm_decode ->
4579 // quantize_q8_1, kernel-check-pinned). When present (caller checked mixer q8_1-fast),
4580 // every projection consumes it — `h` may be a zero-len placeholder and must not be read.
4581 let (qf, mut k, v) = {
4582 let mut fused = None;
4583 let qkv_fast = e.uses_q8_1_fast(&fa.wq)
4584 && e.uses_q8_1_fast(&fa.wk)
4585 && e.uses_q8_1_fast(&fa.wv);
4586 if t == 1 && qkv_fast {
4587 let (hq_o, hd_o);
4588 let (hq, hd): (&CudaSlice<i8>, &CudaSlice<f32>) = match h_q8 {
4589 Some(p) => p,
4590 None => {
4591 (hq_o, hd_o) = e.quantize_q8_1(h, 1, n_embd)?;
4592 (&hq_o, &hd_o)
4593 }
4594 };
4595 fused = e.matmul_q8_fused3(&fa.wq, &fa.wk, &fa.wv, hq, hd)?;
4596 } else if spec_fused_t() && (2..=4).contains(&t) && qkv_fast {
4597 // VERIFY-TIER TRUNK FUSION (MEMRA_SPEC_FUSED_T): one shared quantize + one
4598 // fused3 batched launch replaces three decode-exact calls (3 re-quantizes of
4599 // the same h + 3 _b2/_b4 launches). Bit-identical per (tensor,token,row).
4600 let (hq_o, hd_o);
4601 let (hq, hd): (&CudaSlice<i8>, &CudaSlice<f32>) = match h_q8 {
4602 Some(p) => p,
4603 None => {
4604 (hq_o, hd_o) = e.quantize_q8_1(h, t, n_embd)?;
4605 (&hq_o, &hd_o)
4606 }
4607 };
4608 fused = e.matmul_q8_fused3_t(&fa.wq, &fa.wk, &fa.wv, hq, hd, t)?;
4609 }
4610 match (fused, h_q8) {
4611 (Some(triple), _) => triple,
4612 // shared pre-quantized activation (q8_1-fast guaranteed by the caller): the
4613 // decode-exact dispatch consumes (hq, hd) instead of re-quantizing 3x.
4614 (None, Some((hq, hd))) if qkv_fast => (
4615 e.matmul_decode_exact_pre(&fa.wq, hq, hd, t)?,
4616 e.matmul_decode_exact_pre(&fa.wk, hq, hd, t)?,
4617 e.matmul_decode_exact_pre(&fa.wv, hq, hd, t)?,
4618 ),
4619 (None, _) => (
4620 e.matmul_decode_exact(&fa.wq, h, t)?,
4621 e.matmul_decode_exact(&fa.wk, h, t)?,
4622 e.matmul_decode_exact(&fa.wv, h, t)?,
4623 ),
4624 }
4625 };
4626 // M3/Hy3 have no attention output gate — wq out is exactly q; skip the split.
4627 let gated = geometry.attention_gate
4628 == memra_gguf::config::AttentionGateKind::FusedQ;
4629 let (mut q, gate) = if gated {
4630 let mut q = vbuf(e, t * n_head * head_dim)?; // fully written by q_gate_split
4631 let mut gate = vbuf(e, t * n_head * head_dim)?; // fully written by q_gate_split
4632 e.q_gate_split(&qf, &mut q, &mut gate, head_dim, n_head, t)?;
4633 (q, Some(gate))
4634 } else {
4635 (qf, None)
4636 };
4637
4638 let mut qn = vbuf(e, t * n_head * head_dim)?; // fully written by rms_norm
4639 e.rms_norm(
4640 &q,
4641 fa.q_norm.float_data(),
4642 &mut qn,
4643 head_dim,
4644 n_head * t,
4645 eps,
4646 )?;
4647 q = qn;
4648 let mut kn = vbuf(e, t * n_head_kv * head_dim)?; // fully written by rms_norm
4649 e.rms_norm(
4650 &k,
4651 fa.k_norm.float_data(),
4652 &mut kn,
4653 head_dim,
4654 n_head_kv * t,
4655 eps,
4656 )?;
4657 k = kn;
4658 let rope_dims = geometry.n_rot as usize;
4659 e.rope_neox(
4660 &mut q,
4661 pos_d,
4662 head_dim,
4663 rope_dims,
4664 n_head,
4665 t,
4666 geometry.rope_base,
4667 1.0,
4668 )?;
4669 e.rope_neox(
4670 &mut k,
4671 pos_d,
4672 head_dim,
4673 rope_dims,
4674 n_head_kv,
4675 t,
4676 geometry.rope_base,
4677 1.0,
4678 )?;
4679
4680 // append T new K/V columns to the resident QUANTIZED cache. k/v are token-major [T, kv_dim]
4681 // f32; append-quantize each of the T token rows into the byte cache (q8_0 K / q5_1 V).
4682 let kvl = cache.kv[il].as_mut().unwrap();
4683 let (kv_dim_k, kv_dim_v, ktb, vtb) =
4684 (kvl.kv_dim_k, kvl.kv_dim_v, kvl.k_tok_bytes, kvl.v_tok_bytes);
4685 if let Some(ctr) = stream_ctr {
4686 // stream: ONE batched append at the device counter (rows kernel = the per-view warp
4687 // math on a (block, token) grid, documented byte-identical); host len is a stale
4688 // LOWER BOUND under pre-issue (drain reconciles it).
4689 e.append_kv_quantized_rows_dc(
4690 &k,
4691 &v,
4692 &mut kvl.k,
4693 &mut kvl.v,
4694 ctr,
4695 t,
4696 kv_dim_k,
4697 kv_dim_v,
4698 ktb,
4699 vtb,
4700 crate::Engine::kv_fp8_on(),
4701 )?;
4702 } else {
4703 for i in 0..t {
4704 let k_row = k.slice(i * kv_dim_k..(i + 1) * kv_dim_k);
4705 let v_row = v.slice(i * kv_dim_v..(i + 1) * kv_dim_v);
4706 e.append_kv_quantized_view(
4707 &k_row,
4708 &v_row,
4709 &mut kvl.k,
4710 &mut kvl.v,
4711 kvl.len + i,
4712 kv_dim_k,
4713 kv_dim_v,
4714 ktb,
4715 vtb,
4716 crate::Engine::kv_fp8_on(),
4717 )?;
4718 }
4719 kvl.len += t;
4720 }
4721
4722 // BIT-IDENTICAL VERIFY ATTENTION (spec-exactness fix): the FP accumulation order must be
4723 // byte-for-byte identical to the eager decode path. fa_prefill uses a different tile size
4724 // (BLOCK_Q=64, BK=32) and online-softmax structure than fa_decode's split-K + combine,
4725 // which changes FP summation order and can flip argmax at tight logit margins. Query row r
4726 // attends to keys [0..base_len+r+1) — each successive row sees one more key (the causal
4727 // property). This matches eager: decode appends k at len, then fa_decode sees t_kv = len+1
4728 // keys. The verify appends all T tokens first but bounds the key range per row.
4729 //
4730 // MULTI-ROW FUSED PATH (the long-ctx spec fix, 2026-07-03): when every row takes the vec
4731 // kernel (base_len+1 >= FA_VEC_MIN_TKV), ONE fa_decode_rows launch executes the exact
4732 // per-row program for all T rows (grid.z = row, per-row n_splits from the same
4733 // fa_split_keys formula) — replacing T x (2 launches + 2 dtod copies + 5 partial allocs)
4734 // and multiplying resident CTAs by T on a latency-bound kernel. Bit-identical per row by
4735 // construction; kernel-check pins rows-vs-loop byte identity, run-spec is the end gate.
4736 // Short ctx (any row below the vec crossover) and MEMRA_NO_FA_VEC/MEMRA_FA_ROWS_OFF keep the
4737 // per-row loop (whose fa_decode picks scalar/vec per row exactly like eager decode).
4738 let mut attn = vbuf(e, t * n_head * head_dim)?; // fully written by every FA arm below
4739 let base_len = kvl.len - t; // KV len BEFORE this round's T tokens were appended
4740 // T=1 INCLUDED (2026-07-05): p-min cuts the draft to 1 in ~75% of rounds on hard
4741 // (agentic) content — the old t>1 gate sent those rounds to the per-row loop (262us/row
4742 // + q-row copy + per-row allocs vs 93us/row through the fused kernel at grid.z=1, same
4743 // program). nsys accounting: 1088 of 1456 verify FA launches were T=1 escapees.
4744 // LEAN T=1 ARM (MEMRA_SPEC_LEAN, close35): at t==1, q IS one row and fa_decode on it is
4745 // the EXACT eager decode dispatch (vec_q_v2 + combine_f32; the rows pair measured +50us
4746 // at m=1). Byte-identical: kernel-check pins rows-vs-loop identity, and the per-row loop
4747 // at t=1 is fa_decode on the same q with zero-offset copies. Gates arbitrate.
4748 if let Some(ctr) = stream_ctr {
4749 // STREAM ARM: causal base from the device counter; host kvl.len is a stale lower
4750 // bound used only for the split-sizing upper bound (+64 slack covers M pre-issued
4751 // rounds at K<=8). Views span the bound; per-row limits derive in-kernel.
4752 let upper = kvl.len + t + 64;
4753 let k_view = e.view_u8(&kvl.k, (upper.min(cache.max_ctx)) * ktb);
4754 let v_view = e.view_u8(&kvl.v, (upper.min(cache.max_ctx)) * vtb);
4755 e.fa_decode_rows_dc(
4756 &q,
4757 &k_view,
4758 &v_view,
4759 &mut attn,
4760 head_dim,
4761 n_head,
4762 n_head_kv,
4763 ctr,
4764 upper.min(cache.max_ctx),
4765 t,
4766 scale,
4767 ktb,
4768 vtb,
4769 0,
4770 false,
4771 )?;
4772 } else if spec_lean() && t == 1 {
4773 let t_kv = base_len + 1;
4774 let k_view = e.view_u8(&kvl.k, t_kv * ktb);
4775 let v_view = e.view_u8(&kvl.v, t_kv * vtb);
4776 e.fa_decode_kvmod(
4777 &q,
4778 &k_view,
4779 &v_view,
4780 &mut attn,
4781 head_dim,
4782 n_head,
4783 n_head_kv,
4784 t_kv,
4785 scale,
4786 ktb,
4787 vtb,
4788 crate::Engine::kv_fp8_on(),
4789 )?;
4790 } else if e.fa_rows_eligible(base_len, head_dim) {
4791 let k_view = e.view_u8(&kvl.k, (base_len + t) * ktb);
4792 let v_view = e.view_u8(&kvl.v, (base_len + t) * vtb);
4793 e.fa_decode_rows(
4794 &q,
4795 &k_view,
4796 &v_view,
4797 &mut attn,
4798 head_dim,
4799 n_head,
4800 n_head_kv,
4801 base_len,
4802 t,
4803 scale,
4804 ktb,
4805 vtb,
4806 None,
4807 false,
4808 crate::Engine::kv_fp8_on(),
4809 None,
4810 )?;
4811 } else {
4812 for r in 0..t {
4813 let t_kv_r = base_len + r + 1; // this row sees keys [0..t_kv_r)
4814 let k_view_r = e.view_u8(&kvl.k, t_kv_r * ktb);
4815 let v_view_r = e.view_u8(&kvl.v, t_kv_r * vtb);
4816 // copy q row into an owned buffer (fa_decode takes &CudaSlice, not CudaView)
4817 let mut q_row = vbuf(e, n_head * head_dim)?; // fully written by copy_view_into
4818 let q_src = q.slice(r * n_head * head_dim..(r + 1) * n_head * head_dim);
4819 e.copy_view_into(&mut q_row, 0, &q_src, n_head * head_dim)?;
4820 let mut attn_row = vbuf(e, n_head * head_dim)?; // fully written by fa_decode
4821 e.fa_decode_kvmod(
4822 &q_row,
4823 &k_view_r,
4824 &v_view_r,
4825 &mut attn_row,
4826 head_dim,
4827 n_head,
4828 n_head_kv,
4829 t_kv_r,
4830 scale,
4831 ktb,
4832 vtb,
4833 crate::Engine::kv_fp8_on(),
4834 )?;
4835 e.copy_into(
4836 &mut attn,
4837 r * n_head * head_dim,
4838 &attn_row,
4839 n_head * head_dim,
4840 )?;
4841 }
4842 }
4843
4844 let attn_g = match &gate {
4845 Some(gate) => {
4846 let mut gsig = vbuf(e, t * n_head * head_dim)?; // fully written by sigmoid
4847 e.sigmoid(gate, &mut gsig, t * n_head * head_dim)?;
4848 let mut ag = vbuf(e, t * n_head * head_dim)?; // fully written by mul
4849 e.mul(&attn, &gsig, &mut ag, t * n_head * head_dim)?;
4850 ag
4851 }
4852 None => attn,
4853 };
4854 // DECODE-EXACT wo projection: at m>=5 (K=4+ with pending) the generic matmul would use dp4a
4855 // (128-thread, different FP sum order than MMVQ). Force MMVQ for bit-identity with decode.
4856 Ok(e.matmul_decode_exact(&fa.wo, &attn_g, t)?)
4857 }
4858
4859 /// Context-linear bytes for a plain serving session's trunk cache.
4860 pub fn plain_session_kv_bytes_per_token(&self) -> usize {
4861 crate::cache::cache_bytes_per_token(&self.cfg)
4862 }
4863
4864 /// `(logical bytes/token, ring-capped bytes/token, ring row cap)` for exact admission.
4865 pub fn plain_session_kv_shape(&self) -> (usize, usize, usize) {
4866 (
4867 self.plain_session_kv_bytes_per_token(),
4868 crate::cache::cache_ring_bytes_per_token(&self.cfg),
4869 crate::cache::cache_ring_row_cap(&self.cfg),
4870 )
4871 }
4872
4873 /// Context-linear bytes for a speculative serving session: trunk cache plus persistent MTP
4874 /// scratch. With no MTP head this equals the plain coefficient.
4875 pub fn spec_session_kv_bytes_per_token(&self) -> usize {
4876 let scratch = self
4877 .mtp
4878 .as_ref()
4879 .map(|mtp| {
4880 let (_, _, k, v) = mtp_scratch_layout(&self.cfg, mtp.geom.as_ref());
4881 k + v
4882 })
4883 .unwrap_or(0);
4884 self.plain_session_kv_bytes_per_token()
4885 .saturating_add(scratch)
4886 }
4887
4888 /// Spec twin of [`HybridModel::plain_session_kv_shape`]; Step35's persistent MTP scratch is
4889 /// capped by the same SWA ring rows as the trunk.
4890 pub fn spec_session_kv_shape(&self) -> (usize, usize, usize) {
4891 let total = self.spec_session_kv_bytes_per_token();
4892 let (_, mut ring, rows) = self.plain_session_kv_shape();
4893 if rows > 0 {
4894 ring = ring.saturating_add(
4895 self.mtp
4896 .as_ref()
4897 .map(|mtp| {
4898 let (_, _, k, v) = mtp_scratch_layout(&self.cfg, mtp.geom.as_ref());
4899 k + v
4900 })
4901 .unwrap_or(0),
4902 );
4903 }
4904 (total, ring, rows)
4905 }
4906
4907 /// Greedy MTP speculative decode (§B). Token-identical to `generate(prompt, max_new)` but uses
4908 /// the NextN head to draft K tokens then verifies them in one batched target forward.
4909 /// Returns (generated tokens, total_drafted, total_accepted) so the caller can report
4910 /// acceptance rate. `k` = draft length per round.
4911 ///
4912 /// GRAPH DRAFT (stage 2 of graph-grade spec): when the model is all-Dense and the MTP head is
4913 /// Dense (no MoE host readbacks), the fixed-shape T=1 MTP forward is CUDA-graph-captured ONCE
4914 /// and replayed per draft step — the ~40 eager launches per drafted token collapse into one
4915 /// graph dispatch; only the 4-byte token id (and 4-byte p-min confidence) round-trip per step.
4916 /// Event tracking is disabled for the whole call (generate_graph pattern) so every buffer the
4917 /// captured graph references is event-free; the spec loop is strictly single-stream.
4918 /// MEMRA_SPEC_NOGRAPH=1 forces the eager draft chain.
4919 /// SAMPLED mode (MEMRA_SPEC_TEMP>0) has its OWN capture (gumbel-perturbed in-graph argmax,
4920 /// device Philox event counter, persistent q retention) — graph-vs-eager sampled streams are
4921 /// bit-identical for the same (seed, prompt, K, temp); see the sampled-graph setup in
4922 /// generate_spec_inner2.
4923 /// Multi-turn session: trunk cache + MTP draft scratch persist across generate calls, so
4924 /// turn N+1 primes ONLY its new suffix (the 124k-conversation daily pattern — re-priming a
4925 /// 32k history costs ~54s; a suffix prime costs seconds). APPEND-ONLY by construction: the
4926 /// hybrid linear-attn states are in-place (no position index), so a session can extend but
4927 /// never rewind — `committed` is the exact token list whose state the caches hold (includes
4928 /// any overshoot tokens past max_new; the caller renders from `committed`, not its own echo).
4929 pub fn new_session(
4930 &self,
4931 e: &Engine,
4932 max_ctx: usize,
4933 ) -> Result<SpecSession, Box<dyn std::error::Error>> {
4934 Ok(SpecSession {
4935 // STAGE-OWNED KV (lane/pp2-spec 2026-08-06): `pp::new_cache`, not `Cache::new`. This
4936 // is the SERVING spec-session path, and with the ppN door open across two cards a
4937 // primary-homed cache makes every remote stage peer-read its OWN KV on every verify
4938 // round — the wrong-card class already fixed on the two batched serving paths
4939 // (worker.rs 2483 / 2837). With the door shut `new_cache` IS `Cache::new` (same
4940 // branch, same allocations), so single-device behavior is byte-unchanged.
4941 cache: crate::pp::new_cache(e, &self.cfg, max_ctx)?,
4942 scratch: MtpScratch::new(
4943 e,
4944 &self.cfg,
4945 max_ctx,
4946 self.mtp.as_ref().and_then(|m| m.geom.as_ref()),
4947 )?,
4948 committed: Vec::new(),
4949 last_h: None,
4950 next_pred: None,
4951 sctr: 0,
4952 uctr: 0,
4953 draft_ctx: None,
4954 pending_tok: None,
4955 turn_ckpt: None,
4956 telem: SpecTelemetryCounters::default(),
4957 })
4958 }
4959
4960 /// Forced-gate exact state comparison. This intentionally reads the real live prefixes from
4961 /// their owning PP devices: matching emitted ids alone would miss a stale `len_d`, recurrent
4962 /// snapshot, or draft-KV row that only corrupts the following round.
4963 pub fn optipipe_compare_session_state(
4964 &self,
4965 e: &Engine,
4966 reference: &SpecSession,
4967 candidate: &SpecSession,
4968 ) -> Result<OptiForkStateIdentity, Box<dyn std::error::Error>> {
4969 fn fail(what: &str) -> Box<dyn std::error::Error> {
4970 format!("optipipe state mismatch: {what}").into()
4971 }
4972 fn same_f32(a: &[f32], b: &[f32]) -> bool {
4973 a.len() == b.len()
4974 && a.iter().zip(b).all(|(x, y)| x.to_bits() == y.to_bits())
4975 }
4976 fn compare_layers(
4977 es: &Engine,
4978 range: std::ops::Range<usize>,
4979 reference: &SpecSession,
4980 candidate: &SpecSession,
4981 report: &mut OptiForkStateIdentity,
4982 ) -> Result<(), Box<dyn std::error::Error>> {
4983 for il in range {
4984 match (&reference.cache.kv[il], &candidate.cache.kv[il]) {
4985 (Some(a), Some(b)) => {
4986 if a.len != b.len {
4987 return Err(fail(&format!("layer {il} host KV len {} != {}", a.len, b.len)));
4988 }
4989 let ad = es.dtoh_i32(&a.len_d)?;
4990 let bd = es.dtoh_i32(&b.len_d)?;
4991 if ad != bd || ad.first().copied() != Some(a.len as i32) {
4992 return Err(fail(&format!(
4993 "layer {il} device KV len {ad:?} != {bd:?} (host={})",
4994 a.len,
4995 )));
4996 }
4997 let kb = a.len * a.k_tok_bytes;
4998 let vb = a.len * a.v_tok_bytes;
4999 if kb > 0 {
5000 let ak = es.dtoh_u8_view(&a.k.slice(0..kb))?;
5001 let bk = es.dtoh_u8_view(&b.k.slice(0..kb))?;
5002 if ak != bk {
5003 let at = ak.iter().zip(&bk).position(|(x, y)| x != y).unwrap();
5004 return Err(fail(&format!(
5005 "layer {il} K bytes at byte {at} row {} offset {}: {} != {}",
5006 at / a.k_tok_bytes,
5007 at % a.k_tok_bytes,
5008 ak[at],
5009 bk[at],
5010 )));
5011 }
5012 }
5013 if vb > 0 {
5014 let av = es.dtoh_u8_view(&a.v.slice(0..vb))?;
5015 let bv = es.dtoh_u8_view(&b.v.slice(0..vb))?;
5016 if av != bv {
5017 let at = av.iter().zip(&bv).position(|(x, y)| x != y).unwrap();
5018 return Err(fail(&format!(
5019 "layer {il} V bytes at byte {at} row {} offset {}: {} != {}",
5020 at / a.v_tok_bytes,
5021 at % a.v_tok_bytes,
5022 av[at],
5023 bv[at],
5024 )));
5025 }
5026 }
5027 report.trunk_kv_bytes += kb + vb;
5028 }
5029 (None, None) => {}
5030 _ => return Err(fail(&format!("layer {il} KV presence"))),
5031 }
5032 match (&reference.cache.recur[il], &candidate.cache.recur[il]) {
5033 (Some(a), Some(b)) => {
5034 let ac = es.dtoh(&a.conv_state)?;
5035 let bc = es.dtoh(&b.conv_state)?;
5036 if !same_f32(&ac, &bc) {
5037 return Err(fail(&format!("layer {il} conv state")));
5038 }
5039 let as_ = es.dtoh(&a.ssm_state)?;
5040 let bs = es.dtoh(&b.ssm_state)?;
5041 if !same_f32(&as_, &bs) {
5042 return Err(fail(&format!("layer {il} SSM state")));
5043 }
5044 report.recurrent_bytes += (ac.len() + as_.len()) * 4;
5045 }
5046 (None, None) => {}
5047 _ => return Err(fail(&format!("layer {il} recurrent presence"))),
5048 }
5049 }
5050 Ok(())
5051 }
5052
5053 if reference.committed != candidate.committed {
5054 return Err(fail("committed token ids"));
5055 }
5056 if reference.cache.pos != candidate.cache.pos
5057 || reference.cache.max_ctx != candidate.cache.max_ctx
5058 {
5059 return Err(fail("cache pos/capacity"));
5060 }
5061 if reference.pending_tok != candidate.pending_tok
5062 || reference.next_pred != candidate.next_pred
5063 || reference.sctr != candidate.sctr
5064 || reference.uctr != candidate.uctr
5065 {
5066 return Err(fail("pending/prediction/counter tail"));
5067 }
5068
5069 let mut report = OptiForkStateIdentity::default();
5070 if let Some(fence) = crate::pp::pp_cuts(self.layers.len()) {
5071 let rt = crate::pp::PpNRt::get(e)?;
5072 for stage in 0..rt.n_stages() {
5073 let _scope = rt.enter(stage);
5074 compare_layers(
5075 rt.engine(stage, e),
5076 fence[stage]..fence[stage + 1],
5077 reference,
5078 candidate,
5079 &mut report,
5080 )?;
5081 }
5082 } else {
5083 compare_layers(
5084 e,
5085 0..self.layers.len(),
5086 reference,
5087 candidate,
5088 &mut report,
5089 )?;
5090 }
5091
5092 let (a, b) = (&reference.scratch.kv, &candidate.scratch.kv);
5093 if a.len != b.len || e.dtoh_i32(&a.len_d)? != e.dtoh_i32(&b.len_d)? {
5094 return Err(fail("draft scratch length"));
5095 }
5096 let kb = a.len * a.k_tok_bytes;
5097 let vb = a.len * a.v_tok_bytes;
5098 if kb > 0
5099 && e.dtoh_u8_view(&a.k.slice(0..kb))? != e.dtoh_u8_view(&b.k.slice(0..kb))?
5100 {
5101 return Err(fail("draft scratch K bytes"));
5102 }
5103 if vb > 0
5104 && e.dtoh_u8_view(&a.v.slice(0..vb))? != e.dtoh_u8_view(&b.v.slice(0..vb))?
5105 {
5106 return Err(fail("draft scratch V bytes"));
5107 }
5108 report.scratch_kv_bytes = kb + vb;
5109
5110 match (&reference.last_h, &candidate.last_h) {
5111 (Some(a), Some(b)) => {
5112 let ah = e.dtoh(a)?;
5113 let bh = e.dtoh(b)?;
5114 if !same_f32(&ah, &bh) {
5115 return Err(fail("last hidden/seed bytes"));
5116 }
5117 report.hidden_bytes = ah.len() * 4;
5118 }
5119 (None, None) => {}
5120 _ => return Err(fail("last hidden/seed presence")),
5121 }
5122 Ok(report)
5123 }
5124
5125 /// SESSION-AFFINITY REWIND (lane/session-affinity, 2026-08-05): roll `sess` back to its
5126 /// retained prompt-end checkpoint, so a request whose prompt matches
5127 /// `committed[..rewind_pos()]` exactly can resume there and prime only its own delta.
5128 ///
5129 /// EXACTNESS. After this returns, the session is byte-for-byte the state it was in AT that
5130 /// boundary: full-attn KV truncated to it (append-only, position-addressed), GDN conv/ssm
5131 /// restored from the device copy taken there, draft scratch length reset, `committed`
5132 /// truncated, `last_h` = the boundary's predecessor anchor. That is precisely the state a
5133 /// fresh prime of `committed[..pos]` would have produced, so the following suffix prime and
5134 /// every burst after it are identical to a cold run of the same token stream — the
5135 /// committed-tokens-authoritative contract.
5136 ///
5137 /// `next_pred` and `pending_tok` are CLEARED: both describe generation past the boundary,
5138 /// which the rewind discards. The caller therefore must supply a non-empty suffix (a
5139 /// rewound session cannot serve an empty-suffix continuation burst — there is nothing to
5140 /// continue). The persistent draft graph survives: it bakes only session-stable pointers
5141 /// (the scratch KV, the resident embedding), none of which the rewind moves.
5142 ///
5143 /// The checkpoint is CONSUMED (`turn_ckpt` taken): its snapshot buffers are freed here, and
5144 /// this turn's own prime installs a fresh one at the new prompt end. Returns the position
5145 /// rewound to, or `None` when the session holds no checkpoint (caller: full re-prime).
5146 pub fn spec_rewind_to_checkpoint(
5147 &self,
5148 e: &Engine,
5149 sess: &mut SpecSession,
5150 ) -> Result<Option<usize>, Box<dyn std::error::Error>> {
5151 if sess
5152 .turn_ckpt
5153 .as_ref()
5154 .is_some_and(|ckpt| {
5155 !sess.cache.can_rollback(&ckpt.snap, 0)
5156 || !sess.scratch.can_rewind_to(ckpt.pos)
5157 })
5158 {
5159 return Err("SWA ring rewind checkpoint has been lapped; full re-prime required".into());
5160 }
5161 let Some(ckpt) = sess.turn_ckpt.take() else {
5162 return Ok(None);
5163 };
5164 assert!(
5165 ckpt.pos <= sess.committed.len(),
5166 "checkpoint past committed ({} > {})",
5167 ckpt.pos,
5168 sess.committed.len()
5169 );
5170 // Restore through each layer's owning engine. A single primary-engine rollback is not
5171 // sufficient when the serving cache is stage-owned under cross-device PP.
5172 crate::pp::restore_cache_checkpoint(e, &self.cfg, None, &mut sess.cache, &ckpt.snap)?;
5173 debug_assert_eq!(sess.cache.pos, ckpt.pos, "rollback landed off the checkpoint");
5174 sess.scratch.set_len(e, ckpt.pos)?;
5175 sess.committed.truncate(ckpt.pos);
5176 sess.last_h = Some(ckpt.last_h);
5177 sess.next_pred = None;
5178 sess.pending_tok = None;
5179 Ok(Some(ckpt.pos))
5180 }
5181
5182 /// Grow a parked speculative session to `target_cap` and rewind it to its retained turn
5183 /// checkpoint without re-priming the checkpoint prefix.
5184 ///
5185 /// The trunk cache is restored exactly like a plain grown cache: append-only full-attention
5186 /// KV rows come from the parked cache, while recurrent state comes from the checkpoint's
5187 /// owned snapshot. The MTP scratch is also context-linear and its rows below the checkpoint
5188 /// remain authoritative, so they are copied into a fresh larger scratch before its length is
5189 /// truncated. Pointer-baking draft graphs are dropped and recaptured on the next burst.
5190 ///
5191 /// All fallible work completes before `sess` is mutated. A failed allocation or copy leaves
5192 /// the parked session intact, allowing the caller one reclaim-and-retry attempt.
5193 pub fn spec_grow_and_rewind_to_checkpoint(
5194 &self,
5195 e: &Engine,
5196 sess: &mut SpecSession,
5197 target_cap: usize,
5198 ) -> Result<Option<usize>, Box<dyn std::error::Error>> {
5199 if target_cap <= sess.cache.max_ctx {
5200 return self.spec_rewind_to_checkpoint(e, sess);
5201 }
5202 let Some(ckpt) = sess.turn_ckpt.as_ref() else {
5203 return Ok(None);
5204 };
5205 if ckpt.pos == 0 || ckpt.pos > sess.committed.len() {
5206 return Err(format!(
5207 "checkpoint pos {} outside committed length {}",
5208 ckpt.pos,
5209 sess.committed.len(),
5210 )
5211 .into());
5212 }
5213 if ckpt.pos > target_cap {
5214 return Err(format!(
5215 "checkpoint pos {} exceeds grown capacity {target_cap}",
5216 ckpt.pos,
5217 )
5218 .into());
5219 }
5220
5221 let mut grown_cache = crate::pp::new_cache(e, &self.cfg, target_cap)?;
5222 let mut grown_scratch = MtpScratch::new(
5223 e,
5224 &self.cfg,
5225 target_cap,
5226 self.mtp.as_ref().and_then(|m| m.geom.as_ref()),
5227 )?;
5228 crate::pp::restore_cache_checkpoint(
5229 e,
5230 &self.cfg,
5231 Some(&sess.cache),
5232 &mut grown_cache,
5233 &ckpt.snap,
5234 )?;
5235
5236 let src = &sess.scratch.kv;
5237 let dst = &mut grown_scratch.kv;
5238 if ckpt.pos > src.len
5239 || src.kv_dim_k != dst.kv_dim_k
5240 || src.kv_dim_v != dst.kv_dim_v
5241 || src.k_tok_bytes != dst.k_tok_bytes
5242 || src.v_tok_bytes != dst.v_tok_bytes
5243 {
5244 return Err(format!(
5245 "checkpoint draft layout mismatch (pos {}, source len {})",
5246 ckpt.pos, src.len,
5247 )
5248 .into());
5249 }
5250 let kb = ckpt.pos * src.k_tok_bytes;
5251 let vb = ckpt.pos * src.v_tok_bytes;
5252 if kb > 0 {
5253 e.copy_u8_into(&mut dst.k, 0, &src.k, kb)?;
5254 }
5255 if vb > 0 {
5256 e.copy_u8_into(&mut dst.v, 0, &src.v, vb)?;
5257 }
5258 grown_scratch.set_len(e, ckpt.pos)?;
5259 // The old scratch is dropped immediately after publication below. Bound its D2D reads
5260 // first; growth happens once per rewritten turn, outside the decode hot loop.
5261 e.stream().synchronize()?;
5262
5263 let ckpt = sess
5264 .turn_ckpt
5265 .take()
5266 .expect("checkpoint remained present through transactional grow");
5267 let pos = ckpt.pos;
5268 sess.cache = grown_cache;
5269 sess.scratch = grown_scratch;
5270 sess.committed.truncate(pos);
5271 sess.last_h = Some(ckpt.last_h);
5272 sess.next_pred = None;
5273 sess.pending_tok = None;
5274 sess.draft_ctx = None;
5275 debug_assert_eq!(sess.cache.pos, pos, "grown rewind landed off checkpoint");
5276 debug_assert_eq!(sess.scratch.kv.len, pos, "grown draft rewind landed off checkpoint");
5277 Ok(Some(pos))
5278 }
5279
5280 /// Commit a carried pending bonus (see SpecSession::pending_tok): one T=1 trunk pass
5281 /// (its logits' argmax becomes next_pred) + the draft-KV fill at the carried anchor —
5282 /// byte-identical to the pre-carry session tail. Required before a non-empty-suffix
5283 /// prime, a sampled turn, or parking a session for pool reuse. No-op without a pending.
5284 pub fn spec_flush_pending(
5285 &self,
5286 e: &Engine,
5287 sess: &mut SpecSession,
5288 ) -> Result<(), Box<dyn std::error::Error>> {
5289 let Some(b) = sess.pending_tok.take() else {
5290 return Ok(());
5291 };
5292 let mtp = self.mtp.as_ref().expect("pending carry requires an MTP head");
5293 let n_embd = self.cfg.n_embd as usize;
5294 let (embd_qt, embd_rb) = self.embd.qt_and_row_bytes(n_embd);
5295 let embd_gpu = if spec_host_embd() {
5296 None
5297 } else {
5298 Some(
5299 self.embd_gpu
5300 .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload")),
5301 )
5302 };
5303 let embd_dev = embd_gpu.map(|g| (g, embd_qt, embd_rb));
5304 let pos_b = sess.cache.pos;
5305 sess.scratch.set_len(e, pos_b)?;
5306 let (lg_b, hb) = self.spec_target_step_h(e, b, &mut sess.cache)?;
5307 sess.next_pred = Some(argmax(&lg_b) as u32);
5308 let anchor = sess
5309 .last_h
5310 .as_ref()
5311 .expect("pending carry requires last_h (the predecessor-row anchor)");
5312 self.mtp_kv_fill(e, mtp, &[b], anchor, pos_b, &mut sess.scratch, embd_dev)?;
5313 sess.last_h = Some(hb);
5314 sess.committed.push(b);
5315 Ok(())
5316 }
5317
5318 /// Solo target feed used only at speculative round boundaries. Step35 serving made its
5319 /// staged batched B=1 graph authoritative, so a speculative session must enter and leave
5320 /// rounds through that same graph. Other model families keep their eager T=1 contract.
5321 fn spec_target_step_h(
5322 &self,
5323 e: &Engine,
5324 token: u32,
5325 cache: &mut Cache,
5326 ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
5327 if self.cfg.step35.is_none()
5328 && !matches!(self.cfg.arch, memra_gguf::config::Arch::Qwen35Moe)
5329 {
5330 return self.decode_step_h(e, token, cache);
5331 }
5332 let pos0 = cache.pos;
5333 let (logits, hidden) = self.decode_step_t_core(e, &[token], pos0, cache, None, None)?;
5334 Ok((e.dtoh(&logits)?, hidden))
5335 }
5336
5337 /// Reduced-matrix admission for increment 1. This deliberately does not change the PP-2
5338 /// serving policy: the worker calls it only after `MEMRA_SPEC_PIPE=1` and an explicit spec
5339 /// session already exist.
5340 pub fn spec_pipe_available(&self, e: &Engine) -> bool {
5341 if std::env::var("MEMRA_SPEC_PIPE").as_deref() != Ok("1")
5342 || !spec_devacc()
5343 || std::env::var("MEMRA_SPEC_REPLAY").is_ok()
5344 || spec_stream()
5345 || std::env::var("MEMRA_SPEC_ADAPT").as_deref() == Ok("1")
5346 || std::env::var("MEMRA_SPEC_PMIN0").as_deref() == Ok("1")
5347 || std::env::var("MEMRA_SPEC_PP_ANATOMY").as_deref() == Ok("1")
5348 || std::env::var("MEMRA_SPEC_PMIN")
5349 .ok()
5350 .and_then(|v| v.parse::<f32>().ok())
5351 .unwrap_or(0.0) > 0.0
5352 || self.is_gemma4_e4b()
5353 || self.cfg.gemma4.is_some()
5354 || self.mtp.is_none()
5355 {
5356 return false;
5357 }
5358 let Some(cuts) = crate::pp::pp_cuts(self.layers.len()) else {
5359 return false;
5360 };
5361 if cuts.len() != 3 || crate::pp::pp2_streams_off() || !crate::pp::spec_pp_on() {
5362 return false;
5363 }
5364 crate::pp::PpNRt::get(e)
5365 .map(|rt| rt.n_stages() == 2 && rt.cross_device())
5366 .unwrap_or(false)
5367 }
5368
5369 /// Two warm greedy continuation bursts over one PP-2 interval coordinator. The two existing
5370 /// `generate_spec_inner2` call stacks own all per-session round locals; only phase issue order
5371 /// changes. No callback is accepted in increment 1 — the worker publishes each completed burst.
5372 #[allow(clippy::too_many_arguments)]
5373 pub fn generate_spec_session_pair(
5374 &self,
5375 e: &Engine,
5376 sess_a: &mut SpecSession,
5377 max_new_a: usize,
5378 k_a: usize,
5379 sess_b: &mut SpecSession,
5380 max_new_b: usize,
5381 k_b: usize,
5382 ) -> Result<
5383 ((Vec<u32>, usize, usize), (Vec<u32>, usize, usize)),
5384 Box<dyn std::error::Error>,
5385 > {
5386 if !self.spec_pipe_available(e) {
5387 return Err("two-session speculative pipeline is outside its reduced matrix".into());
5388 }
5389 if max_new_a == 0 || max_new_b == 0 || k_a == 0 || k_b == 0 {
5390 return Err("two-session speculative pipeline requires non-empty positive-K bursts".into());
5391 }
5392 for sess in [&*sess_a, &*sess_b] {
5393 if sess.committed.is_empty()
5394 || sess.last_h.is_none()
5395 || (sess.next_pred.is_none() && sess.pending_tok.is_none())
5396 {
5397 return Err("two-session speculative pipeline requires warm continuations".into());
5398 }
5399 }
5400
5401 let mtp_dense = self
5402 .mtp
5403 .as_ref()
5404 .map(|m| matches!(m.ffn, crate::hybrid::Ffn::Dense { .. }))
5405 .unwrap_or(false);
5406 let trunk_dense = self
5407 .layers
5408 .iter()
5409 .all(|l| matches!(l.ffn, crate::hybrid::Ffn::Dense { .. }));
5410 let graph_ok = std::env::var("MEMRA_SPEC_NOGRAPH").is_err()
5411 && !spec_host_embd()
5412 && mtp_dense
5413 && trunk_dense
5414 && !crate::model::full_prec_enabled();
5415 let graph_a = graph_ok && k_a + 2 < 96;
5416 let graph_b = graph_ok && k_b + 2 < 96;
5417 let was_tracking = e.ctx().is_event_tracking();
5418 if (graph_a || graph_b) && was_tracking {
5419 unsafe {
5420 e.ctx().disable_event_tracking();
5421 }
5422 }
5423
5424 static LOGGED: std::sync::Once = std::sync::Once::new();
5425 LOGGED.call_once(|| {
5426 eprintln!("[spec-pipe] two-session PP-2 continuation pipeline engaged");
5427 });
5428 let sync = std::sync::Arc::new(SpecPipeSync::new());
5429 let lane_a = SpecPipeLane { sync: sync.clone(), lane: 0 };
5430 let lane_b = SpecPipeLane { sync, lane: 1 };
5431 let mut sess_b_ptr = SpecPipeSessionPtr(sess_b as *mut SpecSession);
5432 let (result_a, result_b) = std::thread::scope(|scope| {
5433 let b = scope.spawn(move || {
5434 let mut finish = SpecPipeFinish::new(&lane_b);
5435 let sess_b = unsafe { sess_b_ptr.get_mut() };
5436 let result = e
5437 .ctx()
5438 .bind_to_thread()
5439 .map_err(|err| err.to_string())
5440 .and_then(|_| {
5441 self.generate_spec_inner2(
5442 e,
5443 &[],
5444 max_new_b,
5445 k_b,
5446 graph_b,
5447 Some(sess_b),
5448 None,
5449 None,
5450 None,
5451 None,
5452 Some(&lane_b),
5453 )
5454 .map_err(|err| err.to_string())
5455 });
5456 finish.close(result.is_err());
5457 result
5458 });
5459 let mut finish = SpecPipeFinish::new(&lane_a);
5460 let result_a = self.generate_spec_inner2(
5461 e,
5462 &[],
5463 max_new_a,
5464 k_a,
5465 graph_a,
5466 Some(sess_a),
5467 None,
5468 None,
5469 None,
5470 None,
5471 Some(&lane_a),
5472 );
5473 finish.close(result_a.is_err());
5474 let result_b = b
5475 .join()
5476 .map_err(|_| "paired speculative session B panicked".to_string())
5477 .and_then(|r| r);
5478 (result_a, result_b)
5479 });
5480
5481 if (graph_a || graph_b) && was_tracking {
5482 unsafe {
5483 e.ctx().enable_event_tracking();
5484 }
5485 }
5486 let result_a = result_a?;
5487 let result_b = result_b.map_err(|err| -> Box<dyn std::error::Error> { err.into() })?;
5488 Ok((result_a, result_b))
5489 }
5490
5491 /// One spec-decode turn on a live session. `suffix` = the NEW tokens only (turn N+1's user
5492 /// message rendered through the chat template continuation). Returns (new tokens emitted,
5493 /// drafted, accepted); session.committed grows by suffix + emitted.
5494 pub fn generate_spec_session(
5495 &self,
5496 e: &Engine,
5497 sess: &mut SpecSession,
5498 suffix: &[u32],
5499 max_new: usize,
5500 k: usize,
5501 ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
5502 self.generate_spec_session_sampled(e, sess, suffix, max_new, k, None, None)
5503 }
5504
5505 /// Serve-path sampled spec: routes the burst through the rejection-sampling verify with
5506 /// per-SESSION Philox continuity (sess.sctr/uctr). None = env-driven (CLI) or greedy.
5507 /// Filters (top-k/p/min-p) apply SYMMETRICALLY to draft q and verify p — distribution-exact
5508 /// for the filtered target (feat/filtered-spec).
5509 ///
5510 /// `on_commit` (sse-cadence, 2026-08-05): called with each newly-emitted slice of the
5511 /// output — once right after the prime's first token, then once per round commit — so a
5512 /// streaming caller can flush text at round cadence instead of once per burst. The slices
5513 /// are disjoint, in order, and concatenate to exactly the returned token vec. Emission-
5514 /// timing only: token bytes, session state, and exactness are untouched.
5515 ///
5516 /// The returned bool is a CONTINUE-VERDICT (admission yield, 2026-08-06): `false` ends
5517 /// the burst at the current round boundary, exactly as if `max_new` had been reached —
5518 /// the caller's scheduler regains control without waiting the burst out. Burst size is
5519 /// content-neutral (spec-levers battery), so an early exit moves WHEN the burst returns,
5520 /// never what tokens say. The slice may be EMPTY (a poll-only boundary — round-stream
5521 /// drains and the defensive tail flush can land with nothing new committed).
5522 #[allow(clippy::too_many_arguments)]
5523 pub fn generate_spec_session_sampled(
5524 &self,
5525 e: &Engine,
5526 sess: &mut SpecSession,
5527 suffix: &[u32],
5528 max_new: usize,
5529 k: usize,
5530 sampling: Option<SpecSampling>,
5531 on_commit: Option<&mut dyn FnMut(&[u32]) -> bool>,
5532 ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
5533 self.generate_spec_session_sampled_prime_split(
5534 e, sess, suffix, max_new, k, sampling, None, on_commit,
5535 )
5536 }
5537
5538 /// Serve-only cold-prime segmentation twin. `prime_split` is the same stable boundary the
5539 /// plain worker would honor before entering its sub-floor tokenwise tail; warm continuations
5540 /// pass `None` and stay on the existing zero-prime path.
5541 #[allow(clippy::too_many_arguments)]
5542 pub fn generate_spec_session_sampled_prime_split(
5543 &self,
5544 e: &Engine,
5545 sess: &mut SpecSession,
5546 suffix: &[u32],
5547 max_new: usize,
5548 k: usize,
5549 sampling: Option<SpecSampling>,
5550 prime_split: Option<usize>,
5551 on_commit: Option<&mut dyn FnMut(&[u32]) -> bool>,
5552 ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
5553 self.generate_spec_session_constrained_prime_split(
5554 e, sess, suffix, max_new, k, sampling, None, prime_split, on_commit,
5555 )
5556 }
5557
5558 /// `generate_spec_session_sampled` + GRAMMAR (constrained decoding, 2026-08-03): the
5559 /// hook truncates acceptance at the first grammar-illegal token AFTER the exactness
5560 /// verify (grammar is an extra rejection rule, ordering like the batched-verify twins)
5561 /// and replaces an illegal bonus with the MASKED argmax of the target's own verify
5562 /// column — token-identical to constrained plain greedy decode. GREEDY only (the
5563 /// worker routes sampled constrained to plain decode). Acceptance under tight grammars
5564 /// may drop (drafter is unconstrained); that is measured, not hidden.
5565 #[allow(clippy::too_many_arguments)]
5566 pub fn generate_spec_session_constrained(
5567 &self,
5568 e: &Engine,
5569 sess: &mut SpecSession,
5570 suffix: &[u32],
5571 max_new: usize,
5572 k: usize,
5573 sampling: Option<SpecSampling>,
5574 constraint: Option<&mut dyn SpecConstraint>,
5575 on_commit: Option<&mut dyn FnMut(&[u32]) -> bool>,
5576 ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
5577 self.generate_spec_session_constrained_prime_split(
5578 e, sess, suffix, max_new, k, sampling, constraint, None, on_commit,
5579 )
5580 }
5581
5582 #[allow(clippy::too_many_arguments)]
5583 pub fn generate_spec_session_constrained_prime_split(
5584 &self,
5585 e: &Engine,
5586 sess: &mut SpecSession,
5587 suffix: &[u32],
5588 max_new: usize,
5589 k: usize,
5590 sampling: Option<SpecSampling>,
5591 constraint: Option<&mut dyn SpecConstraint>,
5592 prime_split: Option<usize>,
5593 on_commit: Option<&mut dyn FnMut(&[u32]) -> bool>,
5594 ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
5595 if constraint.is_some() && sampling.is_some_and(|s| s.temp > 0.0) {
5596 return Err("constrained spec decode is greedy-only (worker routes sampled \
5597 constrained to plain decode)".into());
5598 }
5599 // PENDING-CARRY entry flush: a carried bonus precedes any new suffix in the sequence,
5600 // so it must commit BEFORE the suffix primes; the sampled path doesn't carry (its
5601 // round-0 accept needs the commit pass's logits). Empty-suffix greedy bursts — the
5602 // serve continuation case — consume the carry in-loop with zero solo passes.
5603 if sess.pending_tok.is_some()
5604 && (!suffix.is_empty() || sampling.map_or(false, |s| s.temp > 0.0))
5605 {
5606 self.spec_flush_pending(e, sess)?;
5607 }
5608 let mtp_dense = self
5609 .mtp
5610 .as_ref()
5611 .map(|m| matches!(m.ffn, crate::hybrid::Ffn::Dense { .. }))
5612 .unwrap_or(false);
5613 let trunk_dense = self
5614 .layers
5615 .iter()
5616 .all(|l| matches!(l.ffn, crate::hybrid::Ffn::Dense { .. }));
5617 // FULL_PREC forces the EAGER draft: the graph capture would enclose cuBLASLt f32 GEMV
5618 // (the FloatBf16 else-branches) and a bf16_to_f32 dequant alloc — neither is stream-capture
5619 // safe. Eager rides matmul/matmul_decode_exact, which dequant FloatBf16 on use. (§item 2.)
5620 let graph_draft = std::env::var("MEMRA_SPEC_NOGRAPH").is_err()
5621 && !spec_host_embd()
5622 && mtp_dense
5623 && trunk_dense
5624 && k + 2 < 96
5625 && !crate::model::full_prec_enabled();
5626 let was_tracking = e.ctx().is_event_tracking();
5627 if graph_draft && was_tracking {
5628 unsafe {
5629 e.ctx().disable_event_tracking();
5630 }
5631 }
5632 let r = self.generate_spec_inner2(
5633 e, suffix, max_new, k, graph_draft, Some(sess), sampling, constraint, on_commit,
5634 prime_split, None,
5635 );
5636 if graph_draft && was_tracking {
5637 unsafe {
5638 e.ctx().enable_event_tracking();
5639 }
5640 }
5641 let (out, d, a) = r?;
5642 Ok((out, d, a))
5643 }
5644
5645 pub fn generate_spec(
5646 &self,
5647 e: &Engine,
5648 prompt: &[u32],
5649 max_new: usize,
5650 k: usize,
5651 ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
5652 let mtp_dense = self
5653 .mtp
5654 .as_ref()
5655 .map(|m| matches!(m.ffn, crate::hybrid::Ffn::Dense { .. }))
5656 .unwrap_or(false);
5657 let trunk_dense = self
5658 .layers
5659 .iter()
5660 .all(|l| matches!(l.ffn, crate::hybrid::Ffn::Dense { .. }));
5661 // FULL_PREC forces eager (see generate_spec_session note): CUDA graph capture cannot
5662 // enclose cuBLASLt f32 GEMV or the bf16_to_f32 dequant alloc the FloatBf16 path needs.
5663 let graph_draft = std::env::var("MEMRA_SPEC_NOGRAPH").is_err()
5664 && !spec_host_embd()
5665 && mtp_dense
5666 && trunk_dense
5667 && k + 2 < 96
5668 && !crate::model::full_prec_enabled();
5669 if !graph_draft {
5670 return self.generate_spec_inner2(
5671 e, prompt, max_new, k, false, None, None, None, None, None, None,
5672 );
5673 }
5674 let was_tracking = e.ctx().is_event_tracking();
5675 if was_tracking {
5676 unsafe {
5677 e.ctx().disable_event_tracking();
5678 }
5679 }
5680 let r = self.generate_spec_inner2(
5681 e, prompt, max_new, k, true, None, None, None, None, None, None,
5682 );
5683 if was_tracking {
5684 unsafe {
5685 e.ctx().enable_event_tracking();
5686 }
5687 }
5688 r
5689 }
5690
5691 fn generate_spec_inner2(
5692 &self,
5693 e: &Engine,
5694 prompt: &[u32],
5695 max_new: usize,
5696 k: usize,
5697 graph_draft: bool,
5698 mut sess: Option<&mut SpecSession>,
5699 sampling: Option<SpecSampling>,
5700 mut constraint: Option<&mut dyn SpecConstraint>,
5701 mut on_commit: Option<&mut dyn FnMut(&[u32]) -> bool>,
5702 prime_split: Option<usize>,
5703 pipe: Option<&SpecPipeLane>,
5704 ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
5705 assert!(k >= 1, "k must be >= 1");
5706 if let Some(p) = pipe {
5707 p.setup_begin()?;
5708 }
5709 // sse-cadence flush cursor: everything in out[..flushed] has been handed to on_commit.
5710 let mut flushed = 0usize;
5711 // admission yield (2026-08-06): on_commit's continue-verdict; false = end the burst
5712 // at the next round boundary (same exit as max_new reached — the session tail runs).
5713 // Initialized by the unconditional post-prime flush below.
5714 let mut keep_going;
5715 let mtp = self
5716 .mtp
5717 .as_ref()
5718 .expect("generate_spec requires an MTP head (nextn_predict_layers>0)");
5719 let n_vocab = self.output.out_features();
5720 // FR-Spec: the draft head may be TRIMMED (fewer rows than n_vocab); the draft argmax runs
5721 // over the draft vocab and the winning index maps through d2t to a TARGET token id.
5722 // Everything downstream (verify/accept/commit) sees target ids only — exactness unchanged.
5723 let d_vocab = mtp
5724 .shared_head_head
5725 .as_ref()
5726 .unwrap_or(&self.output)
5727 .out_features();
5728 let n_embd = self.cfg.n_embd as usize;
5729 // SESSION MODE: reuse the live cache/scratch, prime only the suffix. `base` = tokens
5730 // already committed (their state is in the caches); 0 = fresh single-shot call.
5731 let session_mode = sess.is_some();
5732 let max_ctx = match sess.as_ref() {
5733 Some(s) => s.cache.max_ctx,
5734 None => prompt.len() + max_new + k + 8,
5735 };
5736 let mut own_cache;
5737 let mut own_scratch;
5738 let (
5739 cache,
5740 scratch,
5741 mut sess_tail,
5742 mut sess_draft_slot,
5743 mut sess_pending_slot,
5744 sess_ckpt_slot,
5745 sess_telem,
5746 ): (
5747 &mut Cache,
5748 &mut MtpScratch,
5749 Option<(
5750 &mut Vec<u32>,
5751 &mut Option<CudaSlice<f32>>,
5752 &mut Option<u32>,
5753 &mut u32,
5754 &mut u32,
5755 )>,
5756 Option<&mut Option<DraftGraphCtx>>,
5757 Option<&mut Option<u32>>,
5758 Option<&mut Option<SpecCheckpoint>>,
5759 Option<&SpecTelemetryCounters>,
5760 ) = match sess.take() {
5761 Some(sr) => {
5762 let SpecSession {
5763 cache,
5764 scratch,
5765 committed,
5766 last_h,
5767 next_pred,
5768 sctr: s_sctr,
5769 uctr: s_uctr,
5770 draft_ctx,
5771 pending_tok,
5772 turn_ckpt,
5773 telem,
5774 } = sr;
5775 (
5776 cache,
5777 scratch,
5778 Some((committed, last_h, next_pred, s_sctr, s_uctr)),
5779 Some(draft_ctx),
5780 Some(pending_tok),
5781 Some(turn_ckpt),
5782 Some(telem),
5783 )
5784 }
5785 None => {
5786 // STAGE-OWNED KV (lane/pp2-spec 2026-08-06) — see `new_session`. Door shut =
5787 // `Cache::new` verbatim.
5788 own_cache = crate::pp::new_cache(e, &self.cfg, max_ctx)?;
5789 // Persistent scratch = max_ctx rows (~2KB/token quantized).
5790 own_scratch = MtpScratch::new(
5791 e,
5792 &self.cfg,
5793 max_ctx,
5794 self.mtp.as_ref().and_then(|m| m.geom.as_ref()),
5795 )?;
5796 (&mut own_cache, &mut own_scratch, None, None, None, None, None)
5797 }
5798 };
5799 let base = cache.pos;
5800 // PENDING-CARRY consume (2026-08-01): a carried bonus reaches here only on the
5801 // empty-suffix GREEDY continuation path (generate_spec_session_sampled flushed every
5802 // other case). It enters the round loop as round-0's pending — verify col 0 — exactly
5803 // like a mid-burst full-accept boundary: no init feed, no tail commit pass.
5804 let carried_pending: Option<u32> = sess_pending_slot.as_mut().and_then(|s| s.take());
5805 // PERSISTENT DRAFT KV (the only mode since 2026-07-08 — the legacy round-local scratch,
5806 // MEMRA_SPEC_KVLOCAL, measured -35 acceptance pts on the 27B p3 sweep and was removed;
5807 // acceptance-only — exactness is verify's job either way).
5808 // HIDDEN-PAIRING CONVENTION (DEFAULT = predecessor-row, 2026-07-04 — the 27B acceptance
5809 // unlock, +16pts): the MTP head is TRAINED on rows pairing token x_p with the trunk
5810 // hidden of its PREDECESSOR h_{p-1} (the reference engine's mtp_update shifts the target
5811 // hiddens right by one; its draft step 0 feeds (id_last, TRUE hidden of the row id_last
5812 // was sampled from)). memra's historical convention paired SAME-ROW (x_p, h_p) in the fill
5813 // and seeded chain step 0 through an extra MTP pass on a duplicated token (the
5814 // pseudo-seed) — measured 27B p2 K=3 acceptance 0.569 vs 0.731, p3 0.445 vs 0.63+, and
5815 // the chain steps j>=1 were already predecessor-shaped, so ONLY the fill + step-0 seed
5816 // move. The fill shifts by one and the chain seeds from the predecessor's true hidden
5817 // DIRECTLY (vh_seed / vx[j-1]) — the pseudo pass disappears (one MTP-block pass saved
5818 // per round on top of the acceptance win). Draft-quality-only: exactness stays the
5819 // verify's job either way. (The legacy same-row pairing seam, MEMRA_SPEC_HSAME, and its
5820 // pseudo-seed passes were removed 2026-07-08 — predecessor pairing won by +16 acc pts;
5821 // the legacy round-local scratch, MEMRA_SPEC_KVLOCAL, went with it.)
5822 // REPLAY-FREE PARTIAL ACCEPT (default, 2026-07-03): partial rounds keep the verify's own
5823 // bit-identical committed-prefix state (KV truncate + recur rebuild from the VerifyCkpt)
5824 // and leave the bonus PENDING — no duplicate trunk pass (profiled ~0.54 extra full weight
5825 // reads/round at long ctx). MEMRA_SPEC_REPLAY=1 restores the legacy rollback+replay (A/B
5826 // + fallback seam).
5827 // Qwen35-MoE stays on the correctness reference path until its retained verify-state
5828 // commit is proven equivalent to sequential serving on the long-prompt gate. Replaying
5829 // every accepted round through the serving-class verifier is slower, but prevents a
5830 // numerically exact verify result from carrying a drifted recurrent cache into the next
5831 // round.
5832 let spec_replay = std::env::var("MEMRA_SPEC_REPLAY").is_ok()
5833 || matches!(self.cfg.arch, memra_gguf::config::Arch::Qwen35Moe);
5834 if constraint.is_some() && spec_replay {
5835 return Err("constrained spec decode does not support MEMRA_SPEC_REPLAY=1 \
5836 (legacy replay commits an unmasked bonus)".into());
5837 }
5838 // TRUE-HIDDEN REFRESH (default in persistent-draft-KV mode): every round overwrites the
5839 // committed positions' scratch entries from the verify's exact hiddens (mtp_kv_fill batch)
5840 // instead of keeping chain-approximate entries. MEMRA_SPEC_NOREFRESH=1 = legacy (A/B seam).
5841 let refresh = std::env::var("MEMRA_SPEC_NOREFRESH").is_err();
5842
5843 // prime: BATCHED cache prime (prime_cache — the measured #1 e2e gap: tokenwise primed at
5844 // ~102/38 tok/s vs the engine's ~2000-5900 tok/s batched prefill). prime_cache returns the
5845 // full pre-output_norm hidden stack [T, n_embd], which IS prompt_h (the persistent-draft-KV
5846 // mtp_kv_fill input) — no per-token collection needed. Prompts below PRIME_MIN_T, and
5847 // MEMRA_PRIME_TOKENWISE=1, and frozen Hy3 CPU/GPU expert splits take the tokenwise
5848 // decode_step_h loop. The latter avoids transient GPU staging of the spilled expert bank.
5849 // EMPTY-SUFFIX CONTINUATION (serve bursts): a session turn with NO new tokens resumes
5850 // generation exactly where the last turn stopped — no prime at all. The stashed
5851 // `next_pred` plays prime_logits' argmax role (it IS the argmax of the logits after
5852 // committed.last()); `last_h` seeds the predecessor pairing below. Fresh calls and
5853 // non-empty suffixes take the normal path.
5854 let continuation = prompt.is_empty();
5855 if continuation {
5856 assert!(session_mode, "empty prompt requires a session");
5857 assert!(
5858 sess_tail
5859 .as_ref()
5860 .map_or(false, |(c, lh, np, _, _)| !c.is_empty()
5861 && lh.is_some()
5862 && (np.is_some() || carried_pending.is_some())),
5863 "empty-suffix continuation needs a primed session (committed + last_h + next_pred|pending)"
5864 );
5865 }
5866 let mut prime_logits;
5867 let mut prompt_h: Option<CudaSlice<f32>> = None;
5868 let t_prime = std::time::Instant::now();
5869 let batched_prime = !continuation
5870 && prompt.len() >= crate::hybrid_forward::PRIME_MIN_T
5871 && std::env::var("MEMRA_PRIME_TOKENWISE").is_err()
5872 && !e.frozen_cpu_experts_prefer_tokenwise_prime();
5873 let prime_split = prime_split.filter(|&split| split > 0 && split < prompt.len());
5874 if prime_split.is_some() && (continuation || base != 0) {
5875 return Err("spec prime split is cold-session-only".into());
5876 }
5877 if continuation {
5878 prime_logits = Vec::new();
5879 } else if let Some(split) = prime_split {
5880 if split < crate::hybrid_forward::PRIME_MIN_T {
5881 return Err(format!(
5882 "spec prime split {split} is below PRIME_MIN_T {}",
5883 crate::hybrid_forward::PRIME_MIN_T,
5884 ).into());
5885 }
5886 // Mirror the plain worker's affinity boundary exactly. The prefix is a request-level
5887 // prime (`queued_after` keeps Step35 arm selection independent of this stop); a tail
5888 // below PRIME_MIN_T then takes the same eager tokenwise continuation as prefill_tick.
5889 // Retain every hidden row so the draft scratch fill remains one coherent prompt.
5890 let mut h_all = e.uninit(prompt.len() * n_embd)?;
5891 let (l, _, h_prefix) =
5892 self.prime_cache(e, &prompt[..split], &mut *cache, prompt.len() - split)?;
5893 e.copy_into(&mut h_all, 0, &h_prefix, split * n_embd)?;
5894 prime_logits = l;
5895 let tail = &prompt[split..];
5896 if tail.len() >= crate::hybrid_forward::PRIME_MIN_T
5897 && std::env::var("MEMRA_PRIME_TOKENWISE").is_err()
5898 && !e.frozen_cpu_experts_prefer_tokenwise_prime()
5899 {
5900 let (l, _, h_tail) = self.prime_cache(e, tail, &mut *cache, 0)?;
5901 e.copy_into(&mut h_all, split * n_embd, &h_tail, tail.len() * n_embd)?;
5902 prime_logits = l;
5903 } else {
5904 for (i, &tok) in tail.iter().enumerate() {
5905 let (l, h) = self.decode_step_h(e, tok, &mut *cache)?;
5906 e.copy_into(&mut h_all, (split + i) * n_embd, &h, n_embd)?;
5907 prime_logits = l;
5908 }
5909 }
5910 if std::env::var("MEMRA_SPEC_STATS").as_deref() == Ok("1") {
5911 eprintln!("[spec-prime] affinity split={split} tail={}", tail.len());
5912 }
5913 prompt_h = Some(h_all);
5914 } else if batched_prime {
5915 let (l, _h_seed, hiddens) = self.prime_cache(e, prompt, &mut *cache, 0)?;
5916 prime_logits = l;
5917 prompt_h = Some(hiddens);
5918 } else {
5919 prime_logits = Vec::new();
5920 prompt_h = Some(e.uninit(prompt.len() * n_embd)?);
5921 for (i, &tok) in prompt.iter().enumerate() {
5922 let (l, h) = self.spec_target_step_h(e, tok, &mut *cache)?;
5923 if let Some(ph) = prompt_h.as_mut() {
5924 e.copy_into(ph, i * n_embd, &h, n_embd)?;
5925 }
5926 prime_logits = l;
5927 }
5928 }
5929 e.stream().synchronize()?;
5930 // Harness timing contract (see crate::PRIME_NANOS): gen-only throughput without the
5931 // prime-subtraction hack.
5932 crate::PRIME_NANOS.store(
5933 t_prime.elapsed().as_nanos() as u64,
5934 std::sync::atomic::Ordering::Relaxed,
5935 );
5936
5937 let (embd_qt, embd_rb) = self.embd.qt_and_row_bytes(n_embd);
5938 // Resident table is fastest when it fits. Large spill deployments can preserve that HBM
5939 // for expert-cache slots and gather only the exact rows needed by MTP/verify from host.
5940 let host_embd = spec_host_embd();
5941 let embd_gpu = if host_embd {
5942 None
5943 } else {
5944 Some(
5945 self.embd_gpu
5946 .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload")),
5947 )
5948 };
5949 let embd_dev = embd_gpu.map(|g| (g, embd_qt, embd_rb));
5950 if host_embd {
5951 eprintln!(
5952 "[spec] host-row embedding: {} bytes kept off HBM",
5953 self.embd.raw.len()
5954 );
5955 }
5956 let mut out: Vec<u32> = Vec::with_capacity(max_new);
5957 let mut total_drafted = 0usize;
5958 let mut total_accepted = 0usize;
5959
5960 // First generated token = argmax of the prompt's last logits (== greedy's first token).
5961 // Emit it, then FEED it to establish the loop invariant below.
5962 // PENDING-CARRY: the carried bonus was already emitted by the LAST burst — it becomes
5963 // last_token WITHOUT re-emission, and round 0 consumes it as pending (no init feed).
5964 // CONSTRAINED entry rules: the first emitted token is the MASKED argmax of the
5965 // prompt's last logits (plain constrained-greedy identity); a continuation without
5966 // a carried pending would emit an UNMASKED stashed next_pred — refused loudly (the
5967 // worker never resumes constrained sessions from the pool, so this cannot fire).
5968 if let Some(c) = constraint.as_deref_mut() {
5969 if continuation && carried_pending.is_none() {
5970 return Err("constrained spec continuation requires a carried pending \
5971 (pool resume is unconstrained-only)".into());
5972 }
5973 if !continuation {
5974 c.mask_logits(&mut prime_logits)
5975 .map_err(|e2| format!("constraint: {e2}"))?;
5976 }
5977 }
5978 let mut last_token = if let Some(b) = carried_pending {
5979 b
5980 } else if continuation {
5981 sess_tail.as_ref().unwrap().2.unwrap()
5982 } else {
5983 argmax(&prime_logits) as u32
5984 };
5985 if carried_pending.is_none() {
5986 out.push(last_token);
5987 // grammar advances with every emitted token (carried pendings were consumed
5988 // by the burst that emitted them).
5989 if let Some(c) = constraint.as_deref_mut() {
5990 c.consume(last_token).map_err(|e2| format!("constraint: {e2}"))?;
5991 }
5992 }
5993 if continuation {
5994 // draft-KV invariant: entries [0..base) are the session's exact fills; truncate any
5995 // overhang so the chain's first append lands at slot base (== committed.len()).
5996 scratch.set_len(e, base)?;
5997 }
5998 // sse-cadence: hand the caller every not-yet-flushed token (disjoint in-order slices
5999 // concatenating to the full `out`). Called after the prime's first token and after each
6000 // round commit — emission timing only, token bytes untouched. The slice may be EMPTY
6001 // (poll-only boundary: zero-round folds commit nothing new); returns the caller's
6002 // continue-verdict (admission yield, 2026-08-06) — false ends the burst at this round.
6003 fn flush_commit(
6004 cb: &mut Option<&mut dyn FnMut(&[u32]) -> bool>,
6005 out: &[u32],
6006 flushed: &mut usize,
6007 ) -> bool {
6008 if let Some(f) = cb.as_mut() {
6009 let keep = f(&out[*flushed..]);
6010 *flushed = out.len();
6011 keep
6012 } else {
6013 true
6014 }
6015 }
6016 keep_going = flush_commit(&mut on_commit, &out, &mut flushed);
6017 // INVARIANT at loop top: `last_token` is the most-recently-committed/emitted token, its
6018 // KV+recur state IS in `cache` (cache.pos = position right AFTER last_token), `last_pred`
6019 // is the greedy ARGMAX of the logits that predict the token FOLLOWING last_token, and
6020 // `h_seed` = last_token's pre-output_norm hidden. Establish it by feeding last_token once
6021 // (mirrors plain greedy). DEVICE-ARGMAX lever: the accept walk only ever consumes the
6022 // argmax of those logits — never the full vector — so a host u32 replaces the Vec<f32>.
6023 // --- SAMPLED SPEC (MEMRA_SPEC_TEMP>0, research/sampled-spec-impl-map.md): rejection-
6024 // sampling verify (Leviathan/Chen) — accept draft x at u < p(x)/q(x), resample from
6025 // norm(max(0,p-q)) on reject, bonus sampled from p on full accept. Counter-based Philox
6026 // everywhere (seed, event) -> reproducible. temp==0/unset = the greedy path, untouched.
6027 let sp = sampling.unwrap_or_else(|| SpecSampling {
6028 temp: std::env::var("MEMRA_SPEC_TEMP")
6029 .ok()
6030 .and_then(|v| v.parse().ok())
6031 .unwrap_or(0.0),
6032 seed: std::env::var("MEMRA_SEED")
6033 .ok()
6034 .and_then(|v| v.parse().ok())
6035 .unwrap_or(42),
6036 top_k: std::env::var("MEMRA_TOP_K")
6037 .ok()
6038 .and_then(|v| v.parse().ok())
6039 .unwrap_or(0),
6040 top_p: std::env::var("MEMRA_TOP_P")
6041 .ok()
6042 .and_then(|v| v.parse().ok())
6043 .unwrap_or(1.0),
6044 min_p: std::env::var("MEMRA_MIN_P")
6045 .ok()
6046 .and_then(|v| v.parse().ok())
6047 .unwrap_or(0.0),
6048 penalty_last_n: std::env::var("MEMRA_PENALTY_LAST_N")
6049 .ok()
6050 .and_then(|v| v.parse().ok())
6051 .unwrap_or(0),
6052 penalty_repeat: std::env::var("MEMRA_PENALTY_REPEAT")
6053 .ok()
6054 .and_then(|v| v.parse().ok())
6055 .unwrap_or(1.0),
6056 penalty_freq: std::env::var("MEMRA_PENALTY_FREQ")
6057 .ok()
6058 .and_then(|v| v.parse().ok())
6059 .unwrap_or(0.0),
6060 penalty_present: std::env::var("MEMRA_PENALTY_PRESENT")
6061 .ok()
6062 .and_then(|v| v.parse().ok())
6063 .unwrap_or(0.0),
6064 });
6065 let (sp_temp, sp_seed) = (sp.temp, sp.seed);
6066 let sampled = sp_temp > 0.0;
6067 // Trimmed heads: q lives on the trimmed vocab; accept gathers use the TRIMMED index and
6068 // the residual scatters q into target-id space (q=-inf off-trim — the head cannot propose
6069 // those, so their residual mass is p(x), correct by construction).
6070 let d2t_dev: Option<CudaSlice<u32>> = if sampled || crate::spec::spec_stream() {
6071 match &mtp.d2t {
6072 Some(map) => Some(e.htod_u32_v(map)?),
6073 None => None,
6074 }
6075 } else {
6076 None
6077 };
6078 let mut q_full_buf: Option<CudaSlice<f32>> = None;
6079 // Counters resume from the session (burst continuity: randomness must never repeat
6080 // across generate_spec_session calls); one-shot callers start at (0,0). Read through
6081 // sess_tail — `sess` was take()n into it above, so sess.as_ref() here is always None.
6082 let mut sctr: u32 = sess_tail.as_ref().map(|(_, _, _, s, _)| **s).unwrap_or(0);
6083 let mut uctr: u32 = sess_tail.as_ref().map(|(_, _, _, _, u)| **u).unwrap_or(0);
6084 // host Philox4x32-10 (mirrors spec_sample.cu; independent stream via ctr_lo tag)
6085 let host_u01 = |seed: u64, ctr: u32| -> f32 {
6086 let (m0, m1) = (0xD2511F53u32, 0xCD9E8D57u32);
6087 let (mut c0, mut c1, mut c2, mut c3) = (0xFFFF_FFFEu32, ctr, 0u32, 0u32);
6088 let (mut k0, mut k1) = ((seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
6089 for _ in 0..10 {
6090 let (h0, l0) = (((m0 as u64 * c0 as u64) >> 32) as u32, m0.wrapping_mul(c0));
6091 let (h1, l1) = (((m1 as u64 * c2 as u64) >> 32) as u32, m1.wrapping_mul(c2));
6092 let (n0, n1, n2, n3) = (h1 ^ c1 ^ k0, l1, h0 ^ c3 ^ k1, l0);
6093 c0 = n0;
6094 c1 = n1;
6095 c2 = n2;
6096 c3 = n3;
6097 k0 = k0.wrapping_add(0x9E3779B9);
6098 k1 = k1.wrapping_add(0xBB67AE85);
6099 }
6100 (c0 as f32 + 1.0) * (1.0 / 4294967296.0)
6101 };
6102 let mut draft_logits: Vec<CudaSlice<f32>> = Vec::new(); // retained head logits (q), per slot
6103 let mut draft_stats: Vec<(f32, f32, f32)> = Vec::new(); // (row_max, th_e, z_e) per slot
6104 let mut perturb_buf: Option<CudaSlice<f32>> = None; // gumbel scratch (max(n_vocab,d_vocab))
6105 let mut sample_tok = e.alloc_u32_zeroed(1)?; // residual/bonus sample out
6106 let mut col_buf: Option<CudaSlice<f32>> = None; // materialized verify column
6107 // Penalties (v2.1): applied to COPIES of q rows and p columns symmetrically (exactness
6108 // for the penalized+filtered target). History = generated tokens, host-tracked window.
6109 let pen_on = sampled
6110 && sp.penalty_last_n > 0
6111 && (sp.penalty_repeat != 1.0 || sp.penalty_freq != 0.0 || sp.penalty_present != 0.0);
6112 let mut pen_hist: Vec<u32> = if pen_on {
6113 prompt.iter().rev().take(64).rev().cloned().collect() // llama-parity: history spans prompt tail too
6114 } else {
6115 Vec::new()
6116 };
6117 let mut pen_hist_d: Option<CudaSlice<u32>> = None;
6118 let mut pcol_buf: Option<CudaSlice<f32>> = None; // penalized p-column scratch
6119 // MEMRA_SPEC_SETUP_TRACE=1 (diagnostics): per-call wall decomposition of the burst
6120 // SETUP + TAIL segments (the round loop's internals are MEMRA_SPEC_PHASE's job) —
6121 // built to pin the serve per-burst fixed cost (research/spec-serving-20260801).
6122 let setup_trace = std::env::var("MEMRA_SPEC_SETUP_TRACE").as_deref() == Ok("1");
6123 let t_ent = std::time::Instant::now();
6124
6125 // SESSION-AFFINITY TURN CHECKPOINT (lane/session-affinity, 2026-08-05): capture the
6126 // PROMPT-END boundary state so a LATER turn can rewind here and re-prime only its own
6127 // delta instead of the whole conversation. See `SpecCheckpoint` for why this boundary is
6128 // the one that matters (a history-rewriting client mutates what the session GENERATED,
6129 // so the next turn's prompt agrees with this one up to exactly here).
6130 //
6131 // WHERE — AND WHY THIS EXACT LINE. Right after the trunk prime, BEFORE the init feed
6132 // (`decode_step_h(last_token)`) and before round 0: the last instant at which the caches
6133 // hold exactly `base + prompt.len()` rows and nothing generated.
6134 //
6135 // This was WRONG in the first cut of this lane: the capture sat after the draft-KV fill,
6136 // which is also after the init feed, so `cache.pos` was `base + prompt.len() + 1` — the
6137 // boundary included the FIRST GENERATED TOKEN. That token is the first thing inside the
6138 // `<think>` block the client strips, so every later turn's diff diverged exactly one
6139 // token below the checkpoint and affinity declined 100% of the time. Measured on the
6140 // owner regime: "history diverged at 12233 of checkpoint 12234". The off-by-one made the
6141 // whole mechanism inert while looking, from the outside, like a working
6142 // correctness-declines-safely path — hence the decline log carries the offsets.
6143 //
6144 // The full-attn planes are `len`-truncatable so the snapshot copies only the GDN conv/ssm
6145 // state (the reason a spec session could not rewind before). The draft scratch needs no
6146 // copy: rows below the boundary are rewritten by the next turn's own fill.
6147 //
6148 // WHEN: non-empty prime only. An empty-suffix continuation burst adds no prompt boundary
6149 // (its "prompt end" IS the previous checkpoint's, already held), so it keeps the existing
6150 // checkpoint rather than replacing it with a strictly worse one.
6151 //
6152 // FAILURE IS SILENT BY DESIGN: on a VRAM-tight rig the snapshot alloc can fail. That
6153 // costs the NEXT turn its rewind (it re-primes fully, today's behavior) and must never
6154 // fail the burst that is already running — so the error is swallowed, loud only under
6155 // MEMRA_DEBUG_SPEC.
6156 if let Some(slot) = sess_ckpt_slot {
6157 if !continuation {
6158 let pos = cache.pos;
6159 debug_assert_eq!(
6160 pos,
6161 base + prompt.len(),
6162 "turn checkpoint must sit at the prompt end, before the init feed"
6163 );
6164 let anchor: Result<CudaSlice<f32>, Box<dyn std::error::Error>> =
6165 if let Some(ph) = &prompt_h {
6166 // hidden of the LAST primed row = the predecessor anchor at this
6167 // boundary (exactly what a fresh prime of committed[..pos] leaves in
6168 // last_h, and what the next prime's fill reads for its first row).
6169 let np = prompt.len();
6170 e.uninit(n_embd).and_then(|mut a| {
6171 e.copy_view_into(
6172 &mut a,
6173 0,
6174 &ph.slice((np - 1) * n_embd..np * n_embd),
6175 n_embd,
6176 )?;
6177 Ok(a)
6178 })
6179 } else {
6180 Err("no prompt hiddens".into())
6181 };
6182 match (cache.snapshot(e), anchor) {
6183 (Ok(snap), Ok(last_h)) => {
6184 *slot = Some(SpecCheckpoint { snap, pos, last_h });
6185 }
6186 (s, a) => {
6187 *slot = None; // a stale checkpoint would rewind to the WRONG boundary
6188 if std::env::var("MEMRA_DEBUG_SPEC").is_ok() {
6189 let err = s.err().map(|e| e.to_string())
6190 .or_else(|| a.err().map(|e| e.to_string()))
6191 .unwrap_or_default();
6192 eprintln!("[spec] turn checkpoint skipped ({err}); \
6193 next turn re-primes in full");
6194 }
6195 }
6196 }
6197 }
6198 }
6199 // INIT FEED — skipped on a pending carry: last_token (the carried bonus) is NOT in the
6200 // caches and must NOT be fed solo; round 0's batched verify commits it as col 0. Its
6201 // seed/anchor hidden is the carried last_h (copied below); last_pred is dead in the
6202 // pending path (t_pred reads verify col 0 — the accept walk overwrites it).
6203 let mut last_pred = 0u32;
6204 let mut last_col_logits: Option<CudaSlice<f32>> = None;
6205 // CONSTRAINED: the init feed's logits back the (n_acc==0, base==0) masked-argmax
6206 // recompute in the grammar-truncation walk — retained host-side, round 0 only.
6207 let mut init_logits_host: Option<Vec<f32>> = None;
6208 let h_seed0: CudaSlice<f32> = if carried_pending.is_none() {
6209 let (init_logits, h) = self.spec_target_step_h(e, last_token, &mut *cache)?;
6210 last_pred = argmax(&init_logits) as u32;
6211 if constraint.is_some() {
6212 init_logits_host = Some(init_logits.clone());
6213 }
6214 // sampled mode: p-distribution after last_token, for the j==0/base==0 accept test.
6215 if sampled {
6216 last_col_logits = Some(e.htod(&init_logits)?);
6217 }
6218 h
6219 } else {
6220 // predecessor-row anchor: hidden of the last COMMITTED row (the carry contract).
6221 let lh = sess_tail
6222 .as_ref()
6223 .unwrap()
6224 .1
6225 .as_ref()
6226 .expect("pending carry requires last_h");
6227 e.clone_dtod(lh)?
6228 };
6229 let t_init = t_ent.elapsed();
6230 let mut last_col_stats: Option<(f32, f32, f32)> = None;
6231 // PERSISTENT h_seed buffer (allocated BEFORE any graph capture so no captured scratch can
6232 // alias it): every path that updates the round seed copies INTO it — no per-round allocs,
6233 // stable pointer for the graph-draft round-start copy.
6234 let mut h_seed_buf = e.clone_dtod(&h_seed0)?;
6235 // Predecessor-pairing trackers: `fill_prev` = trunk hidden AT the last COMMITTED row (the
6236 // predecessor of the next verify's col 0 — the reference's carried pending-h analogue;
6237 // also the predecessor-row hidden for the round-0 legacy-replay seed). At round 0 that
6238 // row is last_token's own (h_seed0). The chain step-0 seed under the pairing default =
6239 // hidden of the row BEFORE last_token = the prompt's last row at round 0 (h_seed_buf
6240 // overwritten below).
6241 let mut fill_prev = e.clone_dtod(&h_seed0)?;
6242 {
6243 if let Some(ph) = &prompt_h {
6244 let np = prompt.len();
6245 e.copy_view_into(
6246 &mut h_seed_buf,
6247 0,
6248 &ph.slice((np - 1) * n_embd..np * n_embd),
6249 n_embd,
6250 )?;
6251 } else if continuation {
6252 if let Some((_, lh, _, _, _)) = sess_tail.as_ref() {
6253 if let Some(lh) = lh.as_ref() {
6254 e.copy_into(&mut h_seed_buf, 0, lh, n_embd)?;
6255 }
6256 }
6257 }
6258 }
6259 // Persistent device prediction slots for the accept walk (max k+1 verify columns).
6260 let mut preds_d = e.alloc_u32_zeroed(k + 2)?;
6261
6262 let debug_spec = std::env::var("MEMRA_DEBUG_SPEC").is_ok();
6263 let fork_mode = OptiForkGateMode::configured();
6264 // MEMRA_SPEC_STATS=1: per-slot accept histogram + draft-length histogram, printed once at
6265 // the end. Metric normalization vs the reference engine: BOTH engines count
6266 // accepted/drafted where the chain stopped at p-min and the sub-threshold token is
6267 // discarded uncounted — per-slot decay + chain-length mix are the extra dimensions.
6268 let spec_stats = std::env::var("MEMRA_SPEC_STATS").is_ok();
6269 let mut st_drafted = vec![0usize; k];
6270 let mut st_accepted = vec![0usize; k];
6271 let mut st_len_hist = vec![0usize; k + 1];
6272 let mut st_full = 0usize;
6273 // P-MIN CONFIDENCE GATE (MEMRA_SPEC_PMIN, the serve script's --spec-draft-p-min mechanism):
6274 // stop the draft chain early when the head's softmax confidence in its own pick drops
6275 // below p_min. Hoisted above the loop: the graph capture bakes the prob kernels iff on.
6276 static PMIN: std::sync::OnceLock<f32> = std::sync::OnceLock::new();
6277 let p_min = *PMIN.get_or_init(|| {
6278 std::env::var("MEMRA_SPEC_PMIN")
6279 .ok()
6280 .and_then(|v| v.parse().ok())
6281 .unwrap_or(0.0)
6282 });
6283 // ZERO-DRAFT ROUNDS (MEMRA_SPEC_PMIN0=1, vendored from llama.cpp's draft gating): let the
6284 // p-min gate apply at j==0 too, so a low-confidence round drafts NOTHING and the verify
6285 // batch is just the pending bonus (m=1 = a plain decode step). llama's 35B win rides
6286 // exactly this — draft acceptance 76% at mean len 2.5 because unpredictable stretches
6287 // never pay draft+verify overhead. Only legal when a pending bonus exists (an empty
6288 // verify batch is not); the j==0 exemption stays for pending-less rounds.
6289 let pmin0 = std::env::var("MEMRA_SPEC_PMIN0")
6290 .map(|v| v == "1")
6291 .unwrap_or(false);
6292
6293 // --- GRAPH DRAFT setup: persistent I/O buffers + ONE capture (2 warmups inside). The
6294 // warmups mutate scratch len_d / pos / tok / seed — all reset at every round start, so the
6295 // only restore needed is the scratch counter. Capture failure (e.g. a non-capturable
6296 // cuBLAS path in an exotic head) falls back to the eager draft chain.
6297 // PER-SESSION PERSISTENCE (2026-08-01): session calls reuse the DraftGraphCtx parked on
6298 // the SpecSession — the capture (2 warmup head forwards + instantiate) ran ONCE at the
6299 // session's first burst, not per burst (measured ~16ms/burst fixed cost on H100 q27,
6300 // research/spec-serving-20260801). Reuse is pointer-exact: the graph bakes the session's
6301 // own scratch KV (never realloc'd), the model's resident embedding, the OnceLock p_min,
6302 // and the g_* buffers carried in the ctx — replay dispatch is identical to a fresh
6303 // capture, so draft tokens are bit-identical (drafts never decide exactness anyway; the
6304 // verify arbitrates). Single-shot calls (sess=None) build a fresh ctx and drop it.
6305 let mut dctx: DraftGraphCtx = match sess_draft_slot.as_mut().and_then(|s| s.take()) {
6306 Some(c) => c,
6307 None => DraftGraphCtx::new(e, n_embd, if sampled { d_vocab } else { 1 })?,
6308 };
6309 // A session that ran greedy bursts first sized g_q/g_perturb at 1; a sampled resume
6310 // needs d_vocab. Realloc is legal exactly while graph_s is None (nothing baked them).
6311 if sampled && dctx.g_q.len() < d_vocab {
6312 dctx.g_q = e.zeros(d_vocab)?;
6313 dctx.g_perturb = e.zeros(d_vocab)?;
6314 }
6315 // DRAFT-SIDE GRAMMAR MASK (lane/draft-mask, 2026-08-04): the drafter samples the
6316 // grammar's legal set, so proposals are legal BY CONSTRUCTION and the verify-side
6317 // truncation (the correctness backstop) stops cutting every tight-schema round.
6318 // The mask is one node inside the captured draft chain — presence is a CAPTURE-TIME
6319 // shape, so a parked graph of the other shape is dropped and recaptured.
6320 let dmask_on = constraint.as_deref().is_some_and(|c| c.draft_mask_enabled());
6321 let dmask_words = if dmask_on { d_vocab.div_ceil(32) } else { 0 };
6322 if dmask_on && dctx.g_dmask.len() < dmask_words {
6323 dctx.g_dmask = e.alloc_u32_zeroed(dmask_words)?;
6324 dctx.graph = None; // the old capture baked the old (or no) mask pointer
6325 dctx.failed.clear_greedy();
6326 dctx.keeper.clear();
6327 }
6328 if dctx.graph.is_some() && dctx.graph_masked != dmask_on {
6329 dctx.graph = None;
6330 dctx.failed.clear_greedy();
6331 dctx.keeper.clear();
6332 }
6333 if graph_draft && !sampled && dctx.graph.is_none() && !dctx.failed.greedy_failed() {
6334 let DraftGraphCtx { g_tok, g_pos, g_seed, g_p, g_dmask, .. } = &mut dctx;
6335 // capture-time contents: ALL-ONES (ban nothing). A replay only ever runs after the
6336 // host uploads the position's real words, so the warmups stay grammar-free.
6337 if dmask_on {
6338 e.htod_u32_into(g_dmask, &vec![u32::MAX; dmask_words])?;
6339 }
6340 let g_dmask_ro: &CudaSlice<u32> = &*g_dmask;
6341 // CAPTURE-RETAIN (#68 fix): the warmup transients' pool addresses are baked into the
6342 // captured graph; the keeper pins them for the graph's lifetime. capture_graph (non-
6343 // retained) freed them at exit — safe for one-shot generate_spec (nothing else touches
6344 // the pool between replays) but WRONG for sessions: burst-boundary prime/fill/commit
6345 // passes (and, in serve, other sessions) recycle those addresses and the replay then
6346 // clobbers live buffers — the ST serve-spec corruption (research/serve-st-20260803).
6347 let cap_res = e.capture_graph_retained(|e| {
6348 self.mtp_head_forward_cap(
6349 e,
6350 mtp,
6351 g_tok,
6352 g_pos,
6353 g_seed,
6354 g_p,
6355 &mut *scratch,
6356 p_min > 0.0 || fork_mode == OptiForkGateMode::Controller,
6357 true,
6358 embd_gpu.expect("graph draft requires resident embedding"),
6359 embd_qt,
6360 embd_rb,
6361 d_vocab,
6362 None,
6363 None,
6364 if dmask_on { Some((g_dmask_ro, dmask_words)) } else { None },
6365 )
6366 });
6367 match cap_res {
6368 Ok((g, keep)) => {
6369 scratch.set_len(e, base)?;
6370 dctx.graph = Some(g);
6371 dctx.graph_masked = dmask_on;
6372 dctx.keeper = keep;
6373 }
6374 Err(err) => {
6375 scratch.set_len(e, base)?;
6376 // LOUD flip (audit Q2): a dropped draft graph is a coverage loss, never
6377 // silent. Once per flip — mark returns None on an already-failed ctx.
6378 if let Some(line) = dctx.failed.mark_greedy(&err.to_string()) {
6379 eprintln!("{line}");
6380 }
6381 }
6382 }
6383 }
6384 // --- SAMPLED GRAPH DRAFT setup (step 3 of the sampled-spec arc): a SECOND capture, own
6385 // graph object, built only when sampled && graph-eligible — the greedy capture above is
6386 // untouched (and skipped when sampled: its graph would never be launched). Same head
6387 // forward, but the in-graph argmax reads GUMBEL-PERTURBED logits; the Philox event
6388 // counter lives in the persistent device g_ctr (bumped in-graph, host-seeded from sctr
6389 // once per round); the raw head logits land in the persistent g_q for the host's
6390 // per-replay async D2D into the round's q slot (q_slots, K x d_vocab, allocated once).
6391 // seed/temp are capture-time constants — baked into graph_s, so a pool-resumed request
6392 // with a different (seed, temp, k) drops the parked sampled graph and recaptures.
6393 // COST OF THE FRESH-SEED SERVE DEFAULT (dogfood F4, 2026-08-04): omitting `seed` on a
6394 // serve request now draws fresh per-request entropy (it used to default to a pinned 0),
6395 // so a seed-omitting request that RESUMES a parked spec session finds an s_key baked
6396 // with the PREVIOUS request's seed and pays one recapture. Bounded, and it does not
6397 // reopen the ~16ms/burst regression the persistent ctx exists to fix: a session's seed
6398 // is fixed for its whole lifetime (worker.rs reads s.sampler.seed() per burst), so
6399 // this compare misses at most ONCE per resumed request — the first burst recaptures
6400 // and every later burst in that request replays. A client that wants the parked graph
6401 // AND reproducibility supplies an explicit `seed`, honored exactly, which keeps s_key
6402 // stable across its whole conversation.
6403 // COMPOSITION RULE (fspec x gsd merge): the in-graph chain samples from the RAW
6404 // softmax — it can hold neither per-row filter stats nor the varying penalty history.
6405 // The sampled graph therefore engages only in the PURE-TEMP regime; filters/penalties
6406 // force the eager draft (which computes stats/penalties per row).
6407 let pure_temp = sp.top_k == 0 && sp.top_p >= 1.0 && sp.min_p <= 0.0 && !pen_on;
6408 let s_key = (sp_seed, sp_temp.to_bits(), k);
6409 if sampled && dctx.s_key.is_some_and(|old| old != s_key) {
6410 dctx.graph_s = None;
6411 dctx.failed.clear_sampled();
6412 dctx.s_key = None;
6413 dctx.q_slots.clear();
6414 dctx.keeper_s.clear();
6415 }
6416 if graph_draft && sampled && pure_temp && dctx.graph_s.is_none()
6417 && !dctx.failed.sampled_failed()
6418 {
6419 let DraftGraphCtx { g_tok, g_pos, g_seed, g_p, g_ctr, g_perturb, g_q, .. } = &mut dctx;
6420 // CAPTURE-RETAIN (#68 fix): same keeper contract as the greedy capture above.
6421 let cap_res = e.capture_graph_retained(|e| {
6422 self.mtp_head_forward_cap(
6423 e,
6424 mtp,
6425 g_tok,
6426 g_pos,
6427 g_seed,
6428 g_p,
6429 &mut *scratch,
6430 p_min > 0.0,
6431 true,
6432 embd_gpu.expect("graph draft requires resident embedding"),
6433 embd_qt,
6434 embd_rb,
6435 d_vocab,
6436 Some((g_ctr, g_perturb, g_q, sp_seed, sp_temp)),
6437 None,
6438 None, // constrained spec is greedy-only — sampled never carries a hook
6439 )
6440 });
6441 match cap_res {
6442 Ok((g, keep)) => {
6443 scratch.set_len(e, base)?;
6444 for _ in 0..k {
6445 dctx.q_slots.push(e.zeros(d_vocab)?);
6446 }
6447 dctx.graph_s = Some(g);
6448 dctx.s_key = Some(s_key);
6449 dctx.keeper_s = keep;
6450 }
6451 Err(err) => {
6452 scratch.set_len(e, base)?;
6453 // LOUD flip (audit Q2): same contract as the greedy capture above.
6454 if let Some(line) = dctx.failed.mark_sampled(&err.to_string()) {
6455 eprintln!("{line}");
6456 }
6457 }
6458 }
6459 }
6460 let t_cap = t_ent.elapsed();
6461 // PERSISTENT DRAFT KV: fill the MTP block's K/V for every prompt position from the exact
6462 // trunk hiddens collected during prime — ONE batched K/V-only pass (overwrites any
6463 // capture-warmup garbage; capture left len at 0). last_token (the init feed) needs no
6464 // fill: the first chain step processes it and appends its entry at slot prompt.len().
6465 if let Some(ph) = &prompt_h {
6466 // SESSION: rows [0..base) are the previous turns' exact fills (refresh overwrote them
6467 // with true verify hiddens) — truncate any draft overhang, fill ONLY the suffix at
6468 // global positions [base..base+tp). Fresh call: base==0, identical to before.
6469 scratch.set_len(e, base)?;
6470 // CHUNKED FILL (long-ctx OOM fix, 2026-07-05): mtp_kv_fill's transients scale with its
6471 // T (concat = T*2*n_embd*4B — 1.5GB at 40k) and its concat loop is 2*T launches. The
6472 // fill is a pure sequential append, so chunking is exact: each chunk appends its rows
6473 // at pos0=base+start with the identical per-row math. Same knob as the trunk prime.
6474 let tp = prompt.len();
6475 let fill_chunk: usize = if crate::cache::swa_ring_on() {
6476 crate::hybrid_forward::prime_chunk_tokens(tp, self.layers.len())
6477 } else {
6478 // Preserve the flag-OFF schedule byte-for-byte, including the legacy zero value
6479 // meaning one monolithic fill.
6480 std::env::var("MEMRA_PRIME_CHUNK")
6481 .ok()
6482 .and_then(|v| v.parse().ok())
6483 .unwrap_or(4096)
6484 };
6485 let fill_chunk = if fill_chunk == 0 { tp } else { fill_chunk };
6486 let mut start = 0usize;
6487 while start < tp {
6488 let end = (start + fill_chunk).min(tp);
6489 let tc = end - start;
6490 {
6491 // PREDECESSOR pairing: row i gets h[i-1]; global row 0 a zeros row (the
6492 // reference engine's initial pending-h is zeroed too); a session turn's row 0
6493 // gets the PREVIOUS turn's last committed hidden (sess.last_h). Per chunk:
6494 // rows start..end read h[start-1..end-1] — one dtod into a chunk buffer.
6495 let mut phs = e.zeros(tc * n_embd)?;
6496 let (src_lo, dst_off) = if start == 0 {
6497 (0, n_embd)
6498 } else {
6499 ((start - 1) * n_embd, 0)
6500 };
6501 let n_copy = if start == 0 {
6502 (tc - 1) * n_embd
6503 } else {
6504 tc * n_embd
6505 };
6506 if start == 0 {
6507 if let Some((_, lh, _, _, _)) = sess_tail.as_ref() {
6508 if let Some(lh) = lh.as_ref() {
6509 e.copy_into(&mut phs, 0, lh, n_embd)?;
6510 }
6511 }
6512 }
6513 if n_copy > 0 {
6514 e.copy_view_into(
6515 &mut phs,
6516 dst_off,
6517 &ph.slice(src_lo..src_lo + n_copy),
6518 n_copy,
6519 )?;
6520 }
6521 self.mtp_kv_fill(
6522 e,
6523 mtp,
6524 &prompt[start..end],
6525 &phs,
6526 base + start,
6527 &mut *scratch,
6528 embd_dev,
6529 )?;
6530 }
6531 start = end;
6532 }
6533 }
6534 // MEMRA_PROFILE_SPEC=2: profiler capture starts HERE — after the prime, so an
6535 // `nsys -c cudaProfilerApi` capture contains ONLY the round loop (draft/verify/commit).
6536 // (=1 brackets the whole call in run_spec.rs, prime included.)
6537 if std::env::var("MEMRA_PROFILE_SPEC").as_deref() == Ok("2") {
6538 unsafe extern "C" {
6539 fn cudaProfilerStart() -> i32;
6540 }
6541 unsafe {
6542 cudaProfilerStart();
6543 }
6544 }
6545 // ROUND-STREAM stage (c) 4 (MEMRA_SPEC_STREAM=1, experimental): pre-issued M-round
6546 // bursts with ZERO per-round host readbacks — the accept/seed/rollback/ring kernels
6547 // consume each other's device outputs; the host drains the ring every M rounds. v1
6548 // constraints: greedy, !spec_replay, single-shot, batched-linear layers, no refresh
6549 // fills (acceptance effect A/B-arbitrated), enters from round 1 (pending guaranteed).
6550 // NOTE: not gated on the caller's graph_draft (its trunk_dense conjunct turns the 35B
6551 // MoE off) — the stream capture encloses ONLY the dense MTP head; the head-dense /
6552 // full-prec / k gates are re-derived here and a failed capture degrades to stream-off.
6553 let stream_on = crate::spec::spec_stream()
6554 && !sampled
6555 && !spec_replay
6556 && constraint.is_none()
6557 && !session_mode
6558 && embd_gpu.is_some()
6559 && !crate::model::full_prec_enabled()
6560 && k + 2 < 96;
6561 let mut stream_graph: Option<cudarc::driver::CudaGraph> = None;
6562 let mut g_tokp2k = e.alloc_u32_zeroed(2 * k.max(1))?;
6563 if stream_on {
6564 let cap = e.capture_graph(|e| {
6565 for j in 0..k.max(1) {
6566 self.mtp_head_forward_cap(
6567 e,
6568 mtp,
6569 &mut dctx.g_tok,
6570 &mut dctx.g_pos,
6571 &mut dctx.g_seed,
6572 &mut dctx.g_p,
6573 &mut *scratch,
6574 true,
6575 true,
6576 embd_gpu.expect("round stream requires resident embedding"),
6577 embd_qt,
6578 embd_rb,
6579 d_vocab,
6580 None,
6581 Some((&mut g_tokp2k, j, d2t_dev.as_ref())),
6582 None, // round-stream requires constraint.is_none() (see stream_on)
6583 )?;
6584 }
6585 Ok(())
6586 });
6587 match cap {
6588 Ok(g) => {
6589 scratch.set_len(e, 0)?;
6590 stream_graph = Some(g);
6591 }
6592 Err(err) => {
6593 scratch.set_len(e, 0)?;
6594 if debug_spec {
6595 eprintln!("[spec] stream-graph capture failed ({err}); stream off");
6596 }
6597 }
6598 }
6599 }
6600 let stream_active = stream_on && stream_graph.is_some();
6601 if debug_spec {
6602 eprintln!("[spec] stream_on={stream_on} env={} samp={sampled} dg={} captured={} active={stream_active} session={session_mode} replay={spec_replay}",
6603 crate::spec::spec_stream(), dctx.graph.is_some(), stream_graph.is_some());
6604 }
6605 let t_v_s = k + 1;
6606 // ROUND-STREAM buffers + ptr tables now live in the model-generic round_stream
6607 // module (extracted 2026-07-12; the gemma burst reuses them).
6608 let sb = crate::round_stream::StreamBufs::new(e, k, crate::spec::spec_stream_m())?;
6609 let crate::round_stream::StreamBufs {
6610 mut vtok_d,
6611 mut brk_d,
6612 mut pend_d,
6613 last_pred_d,
6614 mut pos_ctr,
6615 mut pos_start_d,
6616 mut ring_d,
6617 acc_d: mut stream_acc,
6618 m_rounds,
6619 k: _,
6620 } = sb;
6621 let stream_ptrs: Option<CudaSlice<u64>> = if stream_active {
6622 Some(crate::round_stream::kv_len_ptr_table(
6623 e,
6624 cache,
6625 Some(&pos_ctr),
6626 )?)
6627 } else {
6628 None
6629 };
6630
6631 let t_fill = t_ent.elapsed();
6632 let mut round = 0usize;
6633 // ADAPTIVE DRAFT LENGTH (MEMRA_SPEC_ADAPT=1, opt-in — the gemma_spec accepted-run law,
6634 // ported 2026-08-01): next round's draft depth = last round's accepted run + 1, clamped
6635 // to [floor(pos), k_cap] — a miss shrinks the next draft to the miss point + 1,
6636 // full-accept streaks re-deepen one step per round. NOT the 2026-07-07 acceptance-EMA
6637 // (that arm measured an HONEST LOSS to static per-class optima — 115.0/85.8/73.4 vs
6638 // 121.6/92.7/75.6, EMA lag — and was removed 2026-07-08; rig5090.jsonl has the record).
6639 // The gemma law has no lag class: it reacts within one round, and was worth +7-20% on
6640 // the gemma cells at unchanged exactness (2026-07-10 flip; floor sweep 2026-07-25;
6641 // position key 2026-07-26). Signal = n_acc from the round's EXISTING accept readback —
6642 // zero new syncs; the draft graph is a SINGLE-STEP capture replayed per drafted token,
6643 // so a per-round depth needs no re-capture (unlike gemma's whole-chain graphs). qwen's
6644 // in-round p-min cut already shortens chains mid-round, so gemma's one-round-late p-min
6645 // fold into kc is unnecessary here — the accepted-run law sees the cut via n_acc.
6646 // Exactness is the verify's job at ANY depth (same contract as p-min variable rounds).
6647 // DEFAULT OFF on the qwen path until its cells gate a flip (gemma's is default-on).
6648 // MEASURED 2026-08-01 (H100 GPU-3, interleaved x3, NGEN=256, same-invocation plain
6649 // denominators; research/qwen-adaptive-k-20260801/): REFUTED on the tuned qwen configs.
6650 // q27 K=3+HPOST+PMIN=0.3: short +0.8% (noise; law ~idles, len_hist identical), board
6651 // -1.9%, agentic -0.5%; board PMIN=0 -2.1% (not p-min shadowing — the law itself);
6652 // floor=1 -2.8% (gemma's floor-collapse, reproduced). q35 K=2 board: -6.4% (52/136
6653 // rounds shrink to depth 1; no depth to reclaim at K=2). The gemma direction DOES
6654 // appear at untuned depth-K — q27 K=6 floor=4 +1.5% over fixed K=6 — but stays -3.7%
6655 // below fixed K=3: same verdict class as the retired EMA arm (honest loss to static
6656 // per-class optima). Acceptance-rate rises under the law while tokens/round falls —
6657 // it buys accept-% by adding rounds, and a round's fixed draft+verify cost wins.
6658 // K=1..8 self-consistency PASS both models with the law ON (exactness held).
6659 let adapt = std::env::var("MEMRA_SPEC_ADAPT").as_deref() == Ok("1");
6660 // floor: per-model default keyed on n_embd (gemma's tiering — models with an expensive
6661 // verify keep deep drafts after a miss); MEMRA_SPEC_ADAPT_FLOOR pins it everywhere.
6662 let adapt_floor_env: Option<usize> = std::env::var("MEMRA_SPEC_ADAPT_FLOOR")
6663 .ok()
6664 .and_then(|v| v.parse().ok());
6665 let adapt_floor_default: usize = if self.cfg.n_embd as usize >= 3500 {
6666 4
6667 } else if self.cfg.n_embd as usize >= 2500 {
6668 2
6669 } else {
6670 1
6671 };
6672 let adapt_floor: usize = adapt_floor_env.unwrap_or(adapt_floor_default);
6673 // position key: past floor_ctx a HIGH floor (>=4) relaxes to 1 — forced-deep drafts
6674 // turn net-negative at depth (gemma 31B d1736 evidence); MEMRA_SPEC_FLOOR_CTX moves
6675 // the boundary, an explicit MEMRA_SPEC_ADAPT_FLOOR pins the floor everywhere.
6676 let floor_ctx: usize = std::env::var("MEMRA_SPEC_FLOOR_CTX")
6677 .ok()
6678 .and_then(|v| v.parse().ok())
6679 .unwrap_or(1024);
6680 let floor_at = |pos: usize| -> usize {
6681 if adapt_floor_env.is_some() || pos < floor_ctx {
6682 adapt_floor
6683 } else if adapt_floor >= 4 {
6684 1
6685 } else {
6686 adapt_floor
6687 }
6688 };
6689 // cap: MEMRA_SPEC_CAPMAX (gemma semantics, default 7). Binds only under adapt — the
6690 // fixed-K default path is untouched by this whole block.
6691 let cap_max: usize = std::env::var("MEMRA_SPEC_CAPMAX")
6692 .ok()
6693 .and_then(|v| v.parse().ok())
6694 .unwrap_or(7);
6695 let k_cap = k.min(cap_max).max(1);
6696 let mut kc = k_cap;
6697 let mut opti_fork: Option<OptiForkState> = None;
6698 let mut fork_snapshot: Option<crate::cache::CacheSnapshot> = None;
6699 if fork_mode != OptiForkGateMode::Disabled {
6700 let fence = crate::pp::pp_cuts(self.layers.len());
6701 let refusal = if !session_mode {
6702 Some("not-session")
6703 } else if k != 1 || adapt {
6704 Some("requires-fixed-k1")
6705 } else if sampled || constraint.is_some() || spec_replay {
6706 Some("sampled-constrained-or-replay")
6707 } else if pipe.is_some() {
6708 Some("two-session-pipeline")
6709 } else if !spec_devacc() {
6710 Some("requires-device-accept")
6711 } else if stream_active || crate::spec::spec_stream() {
6712 Some("round-stream")
6713 } else if crate::cache::swa_ring_on() || cache.has_swa_ring() {
6714 Some("swa-ring")
6715 } else if crate::pp::pp_host_bounce_active() {
6716 Some("host-bounce")
6717 } else if fork_mode == OptiForkGateMode::Controller
6718 && cache.recur.iter().any(Option::is_some)
6719 {
6720 Some("controller-requires-zero-recurrent-state")
6721 } else if fence.as_ref().is_none_or(|f| f.len() != 3) {
6722 Some("requires-pp2")
6723 } else {
6724 None
6725 };
6726 if let Some(reason) = refusal {
6727 OPTI_FORK_REFUSALS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
6728 eprintln!("[opti-fork] refused reason={reason}");
6729 } else {
6730 let fence = fence.expect("validated PP-2 fence");
6731 let rt = crate::pp::PpNRt::get(e)?;
6732 let primary_stage0 = rt.engine(0, e).ctx().ordinal() == e.ctx().ordinal();
6733 let primary_stage1 = rt.engine(1, e).ctx().ordinal() == e.ctx().ordinal();
6734 let primary_supported = primary_stage0
6735 || (fork_mode == OptiForkGateMode::Controller && primary_stage1);
6736 if !rt.cross_device() || !primary_supported {
6737 OPTI_FORK_REFUSALS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
6738 eprintln!(
6739 "[opti-fork] refused reason=requires-supported-primary-cross-device"
6740 );
6741 } else {
6742 // Both recurrent snapshots and both seed generations are allocated before
6743 // the first fork, each through its owning PP stage. Allocation failure
6744 // therefore happens before any optimistic state mutation can occur.
6745 let current_snapshot = opti_snapshot_stage_owned(e, cache, rt, &fence)?;
6746 let alternate_snapshot = opti_snapshot_stage_owned(e, cache, rt, &fence)?;
6747 let fork = OptiForkState::new(
6748 e,
6749 cache,
6750 fork_mode,
6751 alternate_snapshot,
6752 &h_seed_buf,
6753 &fill_prev,
6754 rt,
6755 fence[1],
6756 self.layers.len(),
6757 )?;
6758 eprintln!(
6759 "[opti-fork] armed mode={fork_mode:?} snapshots=2 seeds=2 split={} \
6760 payload_dev0={} payload_dev1={} q_threshold={:.3}",
6761 fence[1],
6762 fork.logical_payload_bytes[0],
6763 fork.logical_payload_bytes[1],
6764 fork.controller.map_or(0.0, |policy| policy.threshold),
6765 );
6766 fork_snapshot = Some(current_snapshot);
6767 opti_fork = Some(fork);
6768 }
6769 }
6770 }
6771 // Persistent snapshot buffers are allocated once and refreshed in place. The fork arm
6772 // uses stage-owned snapshots; refused/disabled arms retain the existing generic helper.
6773 let mut snap = match fork_snapshot {
6774 Some(snapshot) => snapshot,
6775 None => cache.snapshot(e)?,
6776 };
6777 let mut carried_opti: Option<OptiControllerTicket> = None;
6778 // ROUND-STREAM stage (b) 3a: device table of per-layer kvl.len_d pointers (stable — the
6779 // cache never reallocates len_d; see cache.rs "stable pointer" note). 0 = no KV layer.
6780 let kv_len_ptrs: Option<CudaSlice<u64>> = if spec_devacc() && !spec_replay {
6781 Some(crate::round_stream::kv_len_ptr_table(e, cache, None)?)
6782 } else {
6783 None
6784 };
6785 // BONUS FOLD (2026-07-04): after a FULL accept the bonus token is NOT committed with a
6786 // separate T=1 trunk pass (a full weight read per round). It stays PENDING and rides as
6787 // column 0 of the NEXT round's verify batch. Under predecessor pairing the next chain
6788 // seeds from the bonus's predecessor's TRUE verify hidden (free — no extra
6789 // pass of any kind). Verify still
6790 // checks every emitted token against the target -> exactness holds by construction; only
6791 // DRAFT QUALITY can shift, which the acceptance numbers arbitrate.
6792 // bonus emitted but not yet committed to cache. A carried pending (see SpecSession::
6793 // pending_tok) enters round 0 directly — the burst boundary becomes a plain round edge.
6794 let mut pending: Option<u32> = carried_pending;
6795 // MEMRA_SPEC_PHASE=1: per-round wall decomposition (draft / verify / accept+commit) —
6796 // no tracing, no extra syncs (each phase is naturally sync-bounded: draft readbacks,
6797 // the verify accept readback). Printed once at loop end via spec-stats.
6798 let anatomy_on = std::env::var("MEMRA_SPEC_PP_ANATOMY").as_deref() == Ok("1");
6799 let phase_on = anatomy_on
6800 || std::env::var("MEMRA_SPEC_PHASE").as_deref() == Ok("1");
6801 // DRAFT-MASK receipt (lane/draft-mask): speculative-clone wall + rounds, printed with
6802 // spec-stats. The clone is the one cost the design adds per round — measured, not assumed.
6803 let (mut dm_clone_ns, mut dm_rounds) = (0u128, 0usize);
6804 // grammar-truncation counters: how many rounds the verify-side cut fired and how many
6805 // already-verified tokens it threw away. THIS is the quantity draft masking targets.
6806 let (mut dm_cuts, mut dm_cut_tokens) = (0usize, 0usize);
6807 let (mut ph_draft, mut ph_verify, mut ph_rest) = (0f64, 0f64, 0f64);
6808 let mut ph_wait = 0f64;
6809 let mut ph_commit = 0f64;
6810 let mut ph_t = std::time::Instant::now();
6811 let mut ph_mark = |acc: &mut f64, on: bool| {
6812 if on {
6813 let now = std::time::Instant::now();
6814 *acc += (now - ph_t).as_secs_f64();
6815 ph_t = now;
6816 }
6817 };
6818 if let Some(p) = pipe {
6819 p.setup_end();
6820 }
6821 while keep_going && out.len() < max_new {
6822 // ROUND-STREAM BURST: from round 1 (pending guaranteed by every non-replay arm),
6823 // issue M rounds with zero readbacks, then drain the ring + reconcile mirrors.
6824 if let (true, Some(sg), Some(ptrs)) = (
6825 stream_active && round >= 1 && pending.is_some(),
6826 &stream_graph,
6827 &stream_ptrs,
6828 ) {
6829 if debug_spec {
6830 static ONCE: std::sync::Once = std::sync::Once::new();
6831 ONCE.call_once(|| {
6832 eprintln!("[memra] ROUND-STREAM burst engaged (M={m_rounds} k={k})")
6833 });
6834 }
6835 e.set_i32_one(&mut pos_ctr, cache.pos as i32)?;
6836 e.set_u32_one(&mut pend_d, pending.unwrap())?;
6837 e.set_u32_one(&mut ring_d, 0)?; // ring count = 0 (writes element 0)
6838 for _mi in 0..m_rounds {
6839 e.i32_copy_add(&pos_ctr, &mut pos_start_d, 0)?;
6840 cache.snapshot_into(e, &mut snap)?; // device D2Ds, stream-ordered
6841 e.i32_copy_add(&pos_ctr, &mut scratch.kv.len_d, 0)?; // draft-KV rollback
6842 e.i32_copy_add(&pos_ctr, &mut dctx.g_pos, 1)?; // rope pos = pos + base
6843 e.u32_copy(&pend_d, &mut dctx.g_tok)?;
6844 e.copy_into(&mut dctx.g_seed, 0, &h_seed_buf, n_embd)?;
6845 sg.launch()?;
6846 e.spec_assemble_verify(
6847 &g_tokp2k,
6848 &pend_d,
6849 d2t_dev.as_ref(),
6850 &mut vtok_d,
6851 &mut brk_d,
6852 p_min,
6853 k,
6854 pmin0,
6855 )?;
6856 let mut ck = VerifyCkpt::new(self.layers.len());
6857 let dummy = vec![0u32; t_v_s];
6858 let (tl_d, vx) = self.decode_step_t_core_stream(
6859 e,
6860 &dummy,
6861 0,
6862 &mut *cache,
6863 embd_dev,
6864 Some(&mut ck),
6865 Some((&vtok_d, &pos_ctr)),
6866 None,
6867 )?;
6868 for j in 0..t_v_s {
6869 e.argmax_token_device_col(&tl_d, j, n_vocab, &mut preds_d, j)?;
6870 }
6871 e.spec_accept_greedy_dc(
6872 &preds_d,
6873 &vtok_d,
6874 &last_pred_d,
6875 &brk_d,
6876 &mut stream_acc,
6877 )?;
6878 e.spec_seed_gather(&vx, &fill_prev, &stream_acc, &mut h_seed_buf, 1, n_embd)?;
6879 e.copy_into(&mut fill_prev, 0, &h_seed_buf, n_embd)?;
6880 self.commit_verified_prefix_stream(
6881 e,
6882 &mut *cache,
6883 &snap,
6884 &ck,
6885 &stream_acc,
6886 1,
6887 t_v_s,
6888 )?;
6889 e.spec_rollback_stream(
6890 ptrs,
6891 &pos_start_d,
6892 &stream_acc,
6893 1,
6894 self.layers.len() + 1,
6895 )?;
6896 e.spec_ring_commit(&vtok_d, &stream_acc, &brk_d, &mut ring_d, &mut pend_d)?;
6897 }
6898 e.stream().synchronize()?;
6899 let ring_h = e.dtoh_u32(&ring_d)?;
6900 let cnt = ring_h[0] as usize;
6901 for i in 0..cnt {
6902 if out.len() < max_new {
6903 out.push(ring_h[1 + i]);
6904 }
6905 }
6906 let pos_h = e.dtoh_i32(&pos_ctr)?[0] as usize;
6907 for il in 0..self.layers.len() {
6908 if let Some(kvl) = cache.kv[il].as_mut() {
6909 kvl.len = pos_h;
6910 }
6911 }
6912 cache.pos = pos_h;
6913 scratch.kv.len = pos_h;
6914 pending = Some(ring_h[cnt]); // last drained token = the live bonus
6915 last_token = ring_h[cnt];
6916 total_drafted += k * m_rounds; // upper bound (p-min breaks uncounted)
6917 total_accepted += cnt.saturating_sub(m_rounds);
6918 if let Some(t) = sess_telem {
6919 // totals only — the burst's per-round accept counts stayed on device
6920 // (that is the point of the round-stream arm). pos_* untouched.
6921 t.record_totals(
6922 m_rounds,
6923 k * m_rounds,
6924 cnt.saturating_sub(m_rounds),
6925 );
6926 }
6927 round += m_rounds;
6928 // sse-cadence: the drained ring is committed — flush it at burst-drain cadence.
6929 keep_going = flush_commit(&mut on_commit, &out, &mut flushed);
6930 continue;
6931 }
6932 let pipe_draft = match pipe {
6933 Some(p) => Some(p.draft_begin(round)?),
6934 None => None,
6935 };
6936 let pos = cache.pos; // #tokens committed (EXCLUDES a pending bonus)
6937 let mut current_opti = carried_opti.take();
6938 let mut fork_generation = if current_opti.is_none() && pending.is_some() {
6939 match opti_fork.as_mut() {
6940 Some(fork) if fork.mode.is_forced() => Some(fork.reserve(&mut snap)?),
6941 None => None,
6942 Some(_) => None,
6943 }
6944 } else {
6945 None
6946 };
6947 if current_opti.is_none() {
6948 if let Some(fork) = opti_fork.as_ref() {
6949 opti_snapshot_stage_owned_into(e, cache, fork.rt, &fork.fence, &mut snap)?;
6950 } else {
6951 cache.snapshot_into(e, &mut snap)?;
6952 }
6953 } else if snap.pos != pos {
6954 return Err(format!(
6955 "optipipe carried snapshot pos {} != current pos {pos}", snap.pos
6956 )
6957 .into());
6958 } // §C: snapshot BEFORE draft+verify (already retained for a carried successor)
6959 ph_mark(&mut ph_rest, phase_on);
6960
6961 // --- 1. DRAFT k tokens with the NextN head (autoregressive, T=1 each) ---
6962 // p-min semantics (both paths): stop the chain early when the head's confidence in
6963 // its own pick drops below p_min — the just-drafted token is DISCARDED, but its
6964 // scratch append stands (identical to the eager chain's ordering). j==0 always drafts.
6965 let base0 = if pending.is_some() { 1usize } else { 0usize };
6966 // fixed draft length by default; MEMRA_SPEC_ADAPT=1 drafts at last round's
6967 // accepted run + 1 (the gemma law — see the setup block above the loop).
6968 let k_this = if adapt { kc } else { k };
6969 let mut draft: Vec<u32> = Vec::with_capacity(k);
6970 let mut draft_idx: Vec<u32> = Vec::with_capacity(k); // trimmed-vocab ids (== draft when untrimmed)
6971 let mut controller_draft_prob: Option<f32> = None;
6972 let mut controller_eager_state: Option<(u32, CudaSlice<f32>)> = None;
6973 if let Some(ticket) = current_opti.as_mut() {
6974 let carried_pending = pending.ok_or("optipipe carried successor lost pending")?;
6975 if ticket.verify_tokens[0] != carried_pending {
6976 return Err(format!(
6977 "optipipe carried pending mismatch: ticket={} live={carried_pending}",
6978 ticket.verify_tokens[0],
6979 )
6980 .into());
6981 }
6982 draft.push(ticket.verify_tokens[1]);
6983 controller_draft_prob = Some(ticket.draft_prob);
6984 controller_eager_state = ticket
6985 .take_eager_seed()
6986 .map(|seed| (ticket.verify_tokens[1], seed));
6987 } else {
6988 // Round-start draft-KV sync (BOTH paths). Persistent: truncate/align to the committed
6989 // history — slots 0..P hold entries for the tokens before last_token@P (P = pos +
6990 // base0 - 1); this single set_len IS the draft-side rollback (drops last round's
6991 // rejected drafts and p-min extras via the len mechanism).
6992 scratch.set_len(e, pos + base0 - 1)?;
6993 if pen_on {
6994 let w0 = pen_hist.len().saturating_sub(sp.penalty_last_n);
6995 pen_hist_d = Some(e.htod_u32_v(&pen_hist[w0..])?);
6996 }
6997 if sampled {
6998 draft_logits.clear();
6999 draft_stats.clear();
7000 }
7001 // DRAFT-SIDE GRAMMAR MASK: clone the committed grammar state ONCE per round; each
7002 // position's mask is computed on that clone and advanced by the PROPOSED token. The
7003 // real state moves only on emission (verify's job), so the emitted stream is
7004 // unchanged — the mask only removes tokens the verify would have truncated anyway.
7005 let mut dmask_live = dmask_on;
7006 if dmask_live {
7007 let t_c = std::time::Instant::now();
7008 constraint
7009 .as_deref_mut()
7010 .unwrap()
7011 .draft_begin()
7012 .map_err(|e2| format!("constraint: {e2}"))?;
7013 dm_clone_ns += t_c.elapsed().as_nanos();
7014 dm_rounds += 1;
7015 }
7016 if let (false, Some(gr)) = (sampled || pen_on, &dctx.graph) {
7017 // GRAPH DRAFT: one dispatch per drafted token. The chain feeds itself on-device
7018 // (in-graph argmax -> tok_d -> next replay's embed; h_nextn -> h_seed_d; pos_d
7019 // inc'd in-graph); the host only reads 4B token (+4B p) and decides the break.
7020 e.set_i32_one(&mut dctx.g_pos, (pos + base0) as i32)?;
7021 e.set_u32_one(&mut dctx.g_tok, last_token)?;
7022 e.copy_into(&mut dctx.g_seed, 0, &h_seed_buf, n_embd)?;
7023 for j in 0..k_this {
7024 // per-position mask upload (contents only — the graph's baked pointer is
7025 // dctx.g_dmask). All-ones once masking goes dead mid-chain, so the captured
7026 // mask node degrades to a no-op ban instead of needing a second graph.
7027 if dmask_live
7028 && !upload_draft_mask(
7029 e,
7030 constraint.as_deref_mut().unwrap(),
7031 &mut dctx.g_dmask,
7032 mtp.d2t.as_ref(),
7033 d_vocab,
7034 dmask_words,
7035 )?
7036 {
7037 // no draft-vocab row is grammar-legal here (a trimmed FR-Spec head can
7038 // genuinely miss the legal set): neutralize the captured mask node and
7039 // finish the chain UNMASKED — exactly pre-lane behaviour, never worse.
7040 e.htod_u32_into(&mut dctx.g_dmask, &vec![u32::MAX; dmask_words])?;
7041 dmask_live = false;
7042 }
7043 gr.launch()?;
7044 scratch.kv.len += 1; // host mirror (len_d advanced in-graph)
7045 let idx = e.dtoh_u32_one(&dctx.g_tok)?;
7046 // #87 SENTINEL TRAP: an all-NaN head-logits row leaves the device argmax's
7047 // init sentinel (0x7FFFFFFF) in g_tok — feeding it onward dereferences
7048 // embed_row(sentinel) = table + ~4.6TB (never mapped) inside the NEXT graph
7049 // replay's embed node, and the MMU fault kills the CUDA context for the
7050 // whole process (research/pp2spec-crash-20260807: 3 coredumps, byte-exact
7051 // VA arithmetic). Refuse loudly instead; the diagnostics name the first-NaN
7052 // buffer (g_seed = the verify-side handoff vs head-side compute).
7053 if (idx as usize) >= d_vocab {
7054 // g_seed is SELF-FED (the replay writes h_nextn back into it), so it
7055 // reads as the head's OUTPUT at j; h_seed_buf is the round's INPUT
7056 // seed, untouched since the round-start copy — the pair discriminates
7057 // "seed arrived poisoned" from "head forward produced NaN".
7058 let seed_h = e.dtoh(&dctx.g_seed)?;
7059 let seed_nan = seed_h.iter().filter(|v| v.is_nan()).count();
7060 let in_h = e.dtoh(&h_seed_buf)?;
7061 let in_nan = in_h.iter().filter(|v| v.is_nan()).count();
7062 return Err(format!(
7063 "draft(graph) argmax sentinel 0x{idx:08x} >= d_vocab {d_vocab} at \
7064 round {round} j={j} pos={pos}: head-out NaN {seed_nan}/{n_embd}, \
7065 round-input-seed NaN {in_nan}/{n_embd} — refusing to dereference \
7066 the embed row (#87 trap)"
7067 )
7068 .into());
7069 }
7070 // trimmed draft vocab -> target token id (identity when no d2t map)
7071 let d = match &mtp.d2t {
7072 Some(map) => map[idx as usize],
7073 None => idx,
7074 };
7075 let draft_p = if p_min > 0.0
7076 || opti_fork.as_ref().is_some_and(|fork| fork.controller.is_some())
7077 {
7078 Some(e.dtoh(&dctx.g_p)?[0])
7079 } else {
7080 None
7081 };
7082 if j == 0 {
7083 controller_draft_prob = draft_p;
7084 }
7085 if let Some(p) = draft_p.filter(|_| p_min > 0.0) {
7086 if p < p_min && (j > 0 || (pmin0 && base0 == 1)) {
7087 break;
7088 }
7089 }
7090 draft.push(d);
7091 // with a trimmed head the NEXT embed must read the TARGET id, not the draft
7092 // index the argmax wrote — patch the persistent token buffer (4B htod).
7093 if d != idx {
7094 e.set_u32_one(&mut dctx.g_tok, d)?;
7095 }
7096 // advance the SPECULATIVE state with the proposal; a dead chain drops to
7097 // unmasked drafting for the remaining positions (verify still arbitrates).
7098 // speculative advance; a chain the grammar can no longer follow (EOS
7099 // proposed) ends here. The captured mask node always runs, so a dead chain
7100 // leaves the buffer NEUTRAL (all-ones = ban nothing) before it exits.
7101 if dmask_live
7102 && !constraint
7103 .as_deref_mut()
7104 .unwrap()
7105 .draft_advance(d)
7106 .map_err(|e2| format!("constraint: {e2}"))?
7107 {
7108 e.htod_u32_into(&mut dctx.g_dmask, &vec![u32::MAX; dmask_words])?;
7109 break;
7110 }
7111 }
7112 } else if let (true, Some(gr)) = (sampled, &dctx.graph_s) {
7113 // SAMPLED GRAPH DRAFT: one replay per drafted token — head forward + gumbel +
7114 // argmax in ONE dispatch; the host reads 4B token (+4B p), D2Ds q into slot j,
7115 // and decides the break. Event-counter continuity: g_ctr is host-seeded to
7116 // sctr-1 ONCE per round (outside the graph); the in-graph bump runs BEFORE the
7117 // perturb, so replay j consumes counter sctr+j — exactly the eager arm's Philox
7118 // stream. Host sctr advances in lockstep (computed, no readback needed).
7119 e.set_i32_one(&mut dctx.g_pos, (pos + base0) as i32)?;
7120 e.set_u32_one(&mut dctx.g_tok, last_token)?;
7121 e.copy_into(&mut dctx.g_seed, 0, &h_seed_buf, n_embd)?;
7122 e.set_u32_one(&mut dctx.g_ctr, sctr.wrapping_sub(1))?;
7123 for j in 0..k_this {
7124 gr.launch()?;
7125 scratch.kv.len += 1; // host mirror (len_d advanced in-graph)
7126 sctr += 1; // mirrors the in-graph g_ctr bump (eager parity:
7127 // counts the p-min-discarded token too)
7128 // q retention: ONE async D2D of the persistent head-logits buffer into this
7129 // round's slot j (stream-ordered after the replay, before the next one).
7130 e.copy_into(&mut dctx.q_slots[j], 0, &dctx.g_q, d_vocab)?;
7131 let idx = e.dtoh_u32_one(&dctx.g_tok)?;
7132 // #87 SENTINEL TRAP (see the greedy graph arm above).
7133 if (idx as usize) >= d_vocab {
7134 let seed_h = e.dtoh(&dctx.g_seed)?;
7135 let seed_nan = seed_h.iter().filter(|v| v.is_nan()).count();
7136 return Err(format!(
7137 "draft(graph-sampled) argmax sentinel 0x{idx:08x} >= d_vocab \
7138 {d_vocab} at round {round} j={j} pos={pos}: round-seed NaN \
7139 {seed_nan}/{n_embd} — refusing to dereference the embed row \
7140 (#87 trap)"
7141 )
7142 .into());
7143 }
7144 let d = match &mtp.d2t {
7145 Some(map) => map[idx as usize],
7146 None => idx,
7147 };
7148 draft_idx.push(idx);
7149 if p_min > 0.0 {
7150 let p = e.dtoh(&dctx.g_p)?[0];
7151 if p < p_min && (j > 0 || (pmin0 && base0 == 1)) {
7152 break;
7153 }
7154 }
7155 draft.push(d);
7156 // trimmed head: the NEXT embed must read the TARGET id (see the greedy arm).
7157 if d != idx {
7158 e.set_u32_one(&mut dctx.g_tok, d)?;
7159 }
7160 }
7161 // uniform accept path: fill draft_stats per used slot (pure-temp regime — the
7162 // stats degenerate to th=0 / full-Z; one filter_stats launch per slot, tiny).
7163 for j in 0..draft.len().max(draft_idx.len()) {
7164 let rows0 = e.htod_i32(&[0])?;
7165 let (mut th_d, mut z_d, mut mx_d) = (e.zeros(1)?, e.zeros(1)?, e.zeros(1)?);
7166 e.filter_stats(
7167 &dctx.q_slots[j],
7168 d_vocab,
7169 &rows0,
7170 &mut th_d,
7171 &mut z_d,
7172 &mut mx_d,
7173 d_vocab,
7174 1,
7175 sp_temp,
7176 sp.top_k,
7177 sp.top_p,
7178 sp.min_p,
7179 )?;
7180 draft_stats.push((e.dtoh(&mx_d)?[0], e.dtoh(&th_d)?[0], e.dtoh(&z_d)?[0]));
7181 }
7182 } else {
7183 // EAGER DRAFT (fallback: MoE head/trunk, huge k, MEMRA_SPEC_NOGRAPH, capture fail).
7184 let mut e_tok = last_token;
7185 let mut d_seed = e.clone_dtod(&h_seed_buf)?;
7186 for j in 0..k_this {
7187 // GPU-ARGMAX DRAFT (2026-07-03): device logits + device argmax + 4-byte token
7188 // read instead of the ~600KB full-vocab dtoh + host argmax per draft token.
7189 let mtp_pos = pos + base0 + j;
7190 // draft-side grammar mask (eager twin of the graph arm's in-graph node).
7191 // A position with no legal draft-vocab row drops to unmasked drafting for
7192 // the rest of the chain (pre-lane behaviour; verify still arbitrates).
7193 if dmask_live {
7194 dmask_live = upload_draft_mask(
7195 e,
7196 constraint.as_deref_mut().unwrap(),
7197 &mut dctx.g_dmask,
7198 mtp.d2t.as_ref(),
7199 d_vocab,
7200 dmask_words,
7201 )?;
7202 }
7203 let (dl_d, h_nextn) = self.mtp_head_forward_dev(
7204 e,
7205 mtp,
7206 e_tok,
7207 &d_seed,
7208 &mut *scratch,
7209 mtp_pos,
7210 embd_dev,
7211 if dmask_live { Some((&dctx.g_dmask, dmask_words)) } else { None },
7212 )?;
7213 let tok_d = if sampled {
7214 // FILTERED Gumbel-max: stats -> masked perturb -> argmax = one draw from
7215 // the filtered softmax (filters off => th=0, exact v1 semantics).
7216 if perturb_buf.is_none() {
7217 perturb_buf = Some(e.zeros(d_vocab.max(n_vocab))?);
7218 }
7219 let mut q_row = e.clone_dtod(&dl_d)?; // retained q (penalized when on)
7220 if pen_on {
7221 let h = pen_hist_d.as_ref().unwrap();
7222 let nh = h.len();
7223 e.penalize_logits(
7224 &mut q_row,
7225 h,
7226 nh,
7227 sp.penalty_repeat,
7228 sp.penalty_freq,
7229 sp.penalty_present,
7230 d_vocab,
7231 )?;
7232 }
7233 let rows0 = e.htod_i32(&[0])?;
7234 let (mut th_d, mut z_d, mut mx_d) = (e.zeros(1)?, e.zeros(1)?, e.zeros(1)?);
7235 e.filter_stats(
7236 &q_row, d_vocab, &rows0, &mut th_d, &mut z_d, &mut mx_d, d_vocab, 1,
7237 sp_temp, sp.top_k, sp.top_p, sp.min_p,
7238 )?;
7239 let (th, z, mx) = (e.dtoh(&th_d)?[0], e.dtoh(&z_d)?[0], e.dtoh(&mx_d)?[0]);
7240 let pb = perturb_buf.as_mut().unwrap();
7241 e.gumbel_perturb_filtered(
7242 &q_row, pb, d_vocab, sp_seed, sctr, sp_temp, mx, th,
7243 )?;
7244 sctr += 1;
7245 draft_logits.push(q_row);
7246 draft_stats.push((mx, th, z));
7247 e.argmax_token_device(pb, d_vocab)?
7248 } else {
7249 e.argmax_token_device(&dl_d, d_vocab)?
7250 };
7251 let idx = e.dtoh_u32_one(&tok_d)?;
7252 // #87 SENTINEL TRAP (eager twin — see the graph arm). Extra diagnostics
7253 // here because the eager chain's operands are all readable: dl_d (the head
7254 // logits row) and d_seed (this step's h_seed) name the first-NaN buffer.
7255 if (idx as usize) >= d_vocab {
7256 let dl_h = e.dtoh(&dl_d)?;
7257 let dl_nan = dl_h.iter().filter(|v| v.is_nan()).count();
7258 let seed_h = e.dtoh(&d_seed)?;
7259 let seed_nan = seed_h.iter().filter(|v| v.is_nan()).count();
7260 return Err(format!(
7261 "draft(eager) argmax sentinel 0x{idx:08x} >= d_vocab {d_vocab} at \
7262 round {round} j={j} pos={pos}: head-logits NaN {dl_nan}/{d_vocab}, \
7263 step-seed NaN {seed_nan}/{n_embd} — refusing to dereference the \
7264 embed row (#87 trap)"
7265 )
7266 .into());
7267 }
7268 let d = match &mtp.d2t {
7269 Some(map) => map[idx as usize],
7270 None => idx,
7271 };
7272 if sampled {
7273 draft_idx.push(idx);
7274 }
7275 let draft_p = if p_min > 0.0
7276 || opti_fork.as_ref().is_some_and(|fork| fork.controller.is_some())
7277 {
7278 let p_d = e.prob_of_token_device(&dl_d, &tok_d, d_vocab)?;
7279 Some(e.dtoh(&p_d)?[0])
7280 } else {
7281 None
7282 };
7283 if j == 0 {
7284 controller_draft_prob = draft_p;
7285 }
7286 if let Some(p) = draft_p.filter(|_| p_min > 0.0) {
7287 if p < p_min && (j > 0 || (pmin0 && base0 == 1)) {
7288 break;
7289 }
7290 }
7291 draft.push(d);
7292 e_tok = d;
7293 d_seed = h_nextn;
7294 // speculative advance; a chain the grammar can no longer follow (EOS
7295 // proposed) ends here — the prefix already proposed still rides verify.
7296 if dmask_live
7297 && !constraint
7298 .as_deref_mut()
7299 .unwrap()
7300 .draft_advance(d)
7301 .map_err(|e2| format!("constraint: {e2}"))?
7302 {
7303 break;
7304 }
7305 }
7306 if opti_fork.as_ref().is_some_and(|fork| fork.controller.is_some()) {
7307 controller_eager_state = Some((e_tok, d_seed));
7308 }
7309 }
7310 }
7311 let k_round = draft.len();
7312 if let Some(p) = pipe {
7313 p.draft_end(round);
7314 }
7315 drop(pipe_draft);
7316
7317 ph_mark(&mut ph_draft, phase_on);
7318 // --- 2. VERIFY: one batched target forward. With a pending bonus, it rides as col 0
7319 // (committing its KV/recur inside the SAME weight read); drafts follow. ---
7320 let verify_tokens: Vec<u32> = match pending {
7321 Some(b) => {
7322 let mut v = Vec::with_capacity(k_round + 1);
7323 v.push(b);
7324 v.extend_from_slice(&draft);
7325 v
7326 }
7327 None => draft.clone(),
7328 };
7329 let base = if pending.is_some() { 1 } else { 0 };
7330 // ckpt (REPLAY-FREE partial accept): retain per-layer state-rebuild inputs alongside
7331 // the verify. Pure buffer keep-alives + dtod clones — kernel work is unchanged.
7332 let mut ckpt = if let Some(ticket) = current_opti.as_mut() {
7333 Some(ticket.take_ckpt())
7334 } else if spec_replay {
7335 None
7336 } else {
7337 Some(VerifyCkpt::new(self.layers.len()))
7338 };
7339 let controller_can_probe = base == 1
7340 && k_round == 1
7341 && out.len().saturating_add(2) < max_new
7342 && controller_draft_prob.is_some()
7343 && opti_fork
7344 .as_ref()
7345 .and_then(|fork| fork.controller.as_ref())
7346 .is_some_and(|policy| !policy.breaker_tripped);
7347 let mut successor_attempt: Option<OptiControllerTicket> = None;
7348 let mut rejected_probe: Option<(f32, u32)> = None;
7349 let mut controller_prepared: Option<OptiControllerPrepared> = None;
7350 if controller_can_probe {
7351 // Prepare d2/q and, on admission, d3 before either current verify half is
7352 // issued. N stage 0 can then be followed immediately by N+1 stage 0; once N's
7353 // boundary fires, those dev0 launches overlap N stage 1 on dev1. Preparing on
7354 // the primary stream after N stage 1 would serialize the supposed pipeline.
7355 let eager_pos = scratch.kv.len + 1;
7356 let (optimistic_pending, pending_probability) = self.opti_controller_draft_step(
7357 e,
7358 mtp,
7359 &mut dctx,
7360 &mut *scratch,
7361 d_vocab,
7362 &mut controller_eager_state,
7363 eager_pos,
7364 embd_dev,
7365 )?;
7366 let first_probability = controller_draft_prob
7367 .ok_or("optipipe controller probe lost first-token probability")?;
7368 let q_proxy = first_probability * pending_probability;
7369 OPTI_GATE_CHECKS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
7370 OPTI_SHADOW_DRAFT_TOKENS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
7371 let admitted = opti_fork
7372 .as_ref()
7373 .and_then(|fork| fork.controller.as_ref())
7374 .ok_or("optipipe controller policy disappeared")?
7375 .admit(q_proxy);
7376 if admitted {
7377 OPTI_GATE_ADMITS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
7378 let eager_pos = scratch.kv.len + 1;
7379 let (optimistic_draft, optimistic_draft_probability) =
7380 self.opti_controller_draft_step(
7381 e,
7382 mtp,
7383 &mut dctx,
7384 &mut *scratch,
7385 d_vocab,
7386 &mut controller_eager_state,
7387 eager_pos,
7388 embd_dev,
7389 )?;
7390 OPTI_SHADOW_DRAFT_TOKENS
7391 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
7392 let eager_seed = controller_eager_state.take().map(|(token, seed)| {
7393 debug_assert_eq!(token, optimistic_draft);
7394 seed
7395 });
7396 controller_prepared = Some(OptiControllerPrepared {
7397 verify_tokens: [optimistic_pending, optimistic_draft],
7398 draft_prob: optimistic_draft_probability,
7399 eager_seed,
7400 q_proxy,
7401 scratch_len: scratch.kv.len,
7402 });
7403 } else {
7404 OPTI_GATE_REJECTS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
7405 OPTI_WASTED_DRAFT_TOKENS
7406 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
7407 rejected_probe = Some((q_proxy, optimistic_pending));
7408 eprintln!(
7409 "[opti-controller] reject q={q_proxy:.6} threshold={:.3}",
7410 opti_fork
7411 .as_ref()
7412 .and_then(|fork| fork.controller.as_ref())
7413 .expect("controller policy")
7414 .threshold,
7415 );
7416 }
7417 }
7418 let fork_attempt = match fork_generation.take() {
7419 Some(generation) if base == 1 && k_round == 1 => Some(generation),
7420 Some(generation) => {
7421 opti_fork
7422 .as_mut()
7423 .expect("fork generation without fork state")
7424 .retire(generation)?;
7425 None
7426 }
7427 None => None,
7428 };
7429 let (tlogits_d, vx) = if let Some(p) = pipe {
7430 self.decode_step_t_core_pipelined(
7431 e,
7432 &verify_tokens,
7433 pos,
7434 &mut *cache,
7435 embd_dev,
7436 ckpt.as_mut(),
7437 p,
7438 round,
7439 )?
7440 } else if controller_can_probe {
7441 let fence = opti_fork
7442 .as_ref()
7443 .ok_or("optipipe controller probe lost fork state")?
7444 .fence;
7445 let boundary = match current_opti.as_mut() {
7446 Some(ticket) => ticket.take_boundary(),
7447 None => self.verify_stage0_issue(
7448 e,
7449 &verify_tokens,
7450 pos,
7451 &mut *cache,
7452 embd_dev,
7453 ckpt.as_mut(),
7454 None,
7455 &fence,
7456 Some(true),
7457 None,
7458 )?,
7459 };
7460 if let Some(prepared) = controller_prepared.take() {
7461 let generation = {
7462 let fork = opti_fork
7463 .as_mut()
7464 .ok_or("optipipe controller admission lost fork state")?;
7465 let generation = fork.reserve_successor()?;
7466 let rt = fork.rt;
7467 let snapshot_fence = fork.fence;
7468 opti_snapshot_one_stage_owned_into(
7469 e,
7470 cache,
7471 rt,
7472 &snapshot_fence,
7473 0,
7474 fork.successor_snapshot_mut(),
7475 )?;
7476 generation
7477 };
7478 let mut successor_ckpt = VerifyCkpt::new(self.layers.len());
7479 let successor_boundary = self.verify_stage0_issue(
7480 e,
7481 &prepared.verify_tokens,
7482 pos + verify_tokens.len(),
7483 &mut *cache,
7484 embd_dev,
7485 Some(&mut successor_ckpt),
7486 None,
7487 &fence,
7488 Some(false),
7489 None,
7490 )?;
7491 OPTI_FORK_ATTEMPTS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
7492 let fork = opti_fork
7493 .as_ref()
7494 .ok_or("optipipe controller ticket lost fork state")?;
7495 successor_attempt = Some(fork.controller_ticket(
7496 generation,
7497 successor_boundary,
7498 successor_ckpt,
7499 prepared.verify_tokens,
7500 prepared.draft_prob,
7501 prepared.eager_seed,
7502 prepared.q_proxy,
7503 prepared.scratch_len,
7504 ));
7505 eprintln!(
7506 "[opti-controller] issue generation={} q={:.6} threshold={:.3} \
7507 verify={:?}",
7508 generation.id,
7509 prepared.q_proxy,
7510 fork.controller.expect("controller policy").threshold,
7511 prepared.verify_tokens,
7512 );
7513 }
7514 let result = self.verify_stage1_finish(
7515 e,
7516 boundary,
7517 &mut *cache,
7518 ckpt.as_mut(),
7519 None,
7520 &fence,
7521 successor_attempt.is_none(),
7522 )?;
7523 if let Some(ticket) = current_opti.as_mut() {
7524 ticket.settle();
7525 }
7526 if successor_attempt.is_some() {
7527 let fork = opti_fork
7528 .as_mut()
7529 .ok_or("optipipe successor snapshot lost fork state")?;
7530 let rt = fork.rt;
7531 let snapshot_fence = fork.fence;
7532 opti_snapshot_one_stage_owned_into(
7533 e,
7534 cache,
7535 rt,
7536 &snapshot_fence,
7537 1,
7538 fork.successor_snapshot_mut(),
7539 )?;
7540 // Publish N only after both independent successor-state queues are complete.
7541 fork.rt.publish_to(1, &e.stream())?;
7542 }
7543 result
7544 } else if let Some(ticket) = current_opti.as_mut() {
7545 let fork = opti_fork
7546 .as_mut()
7547 .ok_or("optipipe carried controller ticket lost fork state")?;
7548 let boundary = ticket.take_boundary();
7549 let result = self.verify_stage1_finish(
7550 e,
7551 boundary,
7552 &mut *cache,
7553 ckpt.as_mut(),
7554 None,
7555 &fork.fence,
7556 true,
7557 )?;
7558 ticket.settle();
7559 result
7560 } else if let Some(generation) = fork_attempt {
7561 let fork = opti_fork.as_mut().expect("fork generation without fork state");
7562 fork.capture_seed(
7563 e,
7564 generation,
7565 &h_seed_buf,
7566 &fill_prev,
7567 scratch.kv.len,
7568 )?;
7569 let action = fork.mode.action(generation.id);
7570 let boundary = self.verify_stage0_issue(
7571 e,
7572 &verify_tokens,
7573 pos,
7574 &mut *cache,
7575 embd_dev,
7576 ckpt.as_mut(),
7577 None,
7578 &fork.fence,
7579 Some(true),
7580 None,
7581 )?;
7582 OPTI_FORK_ATTEMPTS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
7583 let mut ticket = fork.ticket(generation, boundary);
7584 if action == OptiForkAction::Abort {
7585 return Err(format!(
7586 "optipipe forced abort with generation {} stage0 in flight",
7587 generation.id,
7588 )
7589 .into());
7590 }
7591 fork.reconcile(
7592 e,
7593 &mut *cache,
7594 &mut *scratch,
7595 &snap,
7596 &mut h_seed_buf,
7597 &mut fill_prev,
7598 generation,
7599 action,
7600 verify_tokens[0],
7601 )?;
7602 let result = if action == OptiForkAction::Hit {
7603 let boundary = ticket.take_boundary();
7604 self.verify_stage1_finish(
7605 e,
7606 boundary,
7607 &mut *cache,
7608 ckpt.as_mut(),
7609 None,
7610 &fork.fence,
7611 true,
7612 )?
7613 } else {
7614 // The optimistic boundary slot has no reader. Re-run the unchanged serial
7615 // verify only after E_restart published the restored stage-0 state.
7616 self.decode_step_t_core(
7617 e,
7618 &verify_tokens,
7619 pos,
7620 &mut *cache,
7621 embd_dev,
7622 ckpt.as_mut(),
7623 )?
7624 };
7625 ticket.settle();
7626 debug_assert_eq!(ticket.generation, generation);
7627 fork.retire(generation)?;
7628 result
7629 } else {
7630 self.decode_step_t_core(
7631 e,
7632 &verify_tokens,
7633 pos,
7634 &mut *cache,
7635 embd_dev,
7636 ckpt.as_mut(),
7637 )?
7638 };
7639 let pipe_accept = match pipe {
7640 Some(p) => Some(p.accept_begin(round)?),
7641 None => None,
7642 };
7643
7644 ph_mark(&mut ph_verify, phase_on);
7645 // --- 3. GREEDY ACCEPT (walk prefix, stop at first mismatch) ---
7646 // DEVICE-ARGMAX ACCEPT: argmax every verify column ON DEVICE (same 2-pass kernels +
7647 // smallest-index tie-break as host argmax, argmax_gate-validated) and read back ONE
7648 // [T] u32 — replaces the T x n_vocab f32 dtoh + T host argmaxes per round.
7649 // t_pred[j] = target's greedy prediction for the slot after draft[j-1] (j>=1) or after
7650 // last_token (j==0). With a pending bonus, col 0 IS the prediction after last_token
7651 // (== the bonus), so every index shifts by `base` and last_pred is unused.
7652 let t_v = verify_tokens.len();
7653 let mut preds: Vec<u32> = Vec::new();
7654 if !sampled {
7655 for j in 0..t_v {
7656 e.argmax_token_device_col(&tlogits_d, j, n_vocab, &mut preds_d, j)?;
7657 }
7658 preds = e.dtoh_u32(&preds_d)?; // <- the verify-GPU wait lands here
7659 // #87 SENTINEL TRAP, verify side: a sentinel pred becomes the round's bonus =
7660 // next round's last_token = the next chain's embed lookup. Catch it at the
7661 // source with the column named — an all-NaN VERIFY column implicates the
7662 // stage-split trunk (decode_step_t_core_ppn), not the draft head.
7663 if let Some(bad) = preds[..t_v].iter().position(|&p| (p as usize) >= n_vocab) {
7664 let col = &tlogits_d.slice(bad * n_vocab..(bad + 1) * n_vocab);
7665 let mut probe = e.zeros(n_vocab)?;
7666 e.copy_view_into(&mut probe, 0, col, n_vocab)?;
7667 let col_h = e.dtoh(&probe)?;
7668 let col_nan = col_h.iter().filter(|v| v.is_nan()).count();
7669 return Err(format!(
7670 "verify argmax sentinel 0x{:08x} >= n_vocab {n_vocab} at round {round} \
7671 col {bad}/{t_v} pos={pos}: verify-logits col NaN {col_nan}/{n_vocab} \
7672 — the stage-split verify produced a poisoned column (#87 trap)",
7673 preds[bad]
7674 )
7675 .into());
7676 }
7677 }
7678 ph_mark(&mut ph_wait, phase_on);
7679 let t_pred = |j: usize| -> u32 {
7680 if j == 0 && base == 0 {
7681 last_pred
7682 } else {
7683 preds[base + j - 1]
7684 }
7685 };
7686 let mut devacc_seeded = false;
7687 let mut devacc_acc: Option<CudaSlice<u32>> = None;
7688 let (n_acc, bonus) = if !sampled {
7689 // ROUND-STREAM stage (a) (MEMRA_SPEC_DEVACC=1 opt-in): the walk runs ON DEVICE
7690 // (spec_accept_greedy, verbatim rule) and the host reads back 8B (n_acc, bonus)
7691 // instead of the [T] preds. Same sync count — machinery for stages (b)/(c),
7692 // gated on token identity vs the host walk (the arms below are bit-equal rules).
7693 if crate::spec::spec_devacc() && k_round > 0 && !spec_replay
7694 && constraint.is_none() {
7695 let draft_d = e.htod_u32_v(&draft)?;
7696 let mut acc_out = e.alloc_u32_zeroed(2)?;
7697 e.spec_accept_greedy(
7698 &preds_d,
7699 &draft_d,
7700 last_pred,
7701 base,
7702 k_round,
7703 &mut acc_out,
7704 )?;
7705 devacc_acc = Some(acc_out.clone());
7706 // stage (b): next-round seed gathered ON DEVICE from acc_out before the host
7707 // ever reads n_acc (j=base+n_acc -> vx col j-1; j==0 -> fill_prev). The three
7708 // non-replay commit arms skip their host-offset seed copies (guarded below);
7709 // the legacy spec_replay arm keeps its own rx-based seeding (excluded here).
7710 // NOTE: fill_prev is NOT updated here — the commit arms' TRUE-HIDDEN
7711 // REFRESH reads the OLD fill_prev (predecessor of this round's verify batch);
7712 // the update lands after the arms (devacc_seeded guard below).
7713 e.spec_seed_gather(&vx, &fill_prev, &acc_out, &mut h_seed_buf, base, n_embd)?;
7714 // 3a: KV lens roll back on device (len = saved + base + n_acc, all arms'
7715 // unified rule; full accept rewrites the verify-left value). Host mirrors
7716 // update after the readback; commit_verified_prefix skips its len_d writes.
7717 if let Some(successor) = successor_attempt.as_ref() {
7718 opti_fork
7719 .as_mut()
7720 .ok_or("optipipe successor reconcile lost fork state")?
7721 .queue_actual_reconcile(
7722 e,
7723 &snap,
7724 &acc_out,
7725 successor.verify_tokens[0],
7726 base,
7727 )?;
7728 } else if let Some(ptrs) = &kv_len_ptrs {
7729 let saved: Vec<i32> = (0..self.layers.len())
7730 .map(|il| snap.kv_len[il].map(|v| v as i32).unwrap_or(0))
7731 .collect();
7732 let saved_d = e.htod_i32(&saved)?;
7733 e.spec_rollback_kv(ptrs, &saved_d, &acc_out, base, self.layers.len())?;
7734 }
7735 devacc_seeded = true;
7736 let ab = e.dtoh_u32(&acc_out)?;
7737 (ab[0] as usize, ab[1])
7738 } else {
7739 let mut n_acc = 0usize;
7740 for j in 0..k_round {
7741 if t_pred(j) == draft[j] {
7742 n_acc += 1;
7743 } else {
7744 break;
7745 }
7746 }
7747 // bonus = target's own token at the first non-accepted slot. n_acc in 0..=k; t_pred
7748 // is defined for j in 0..=k (j==0 -> last_logits, j>=1 -> col j-1, last col = k-1).
7749 (n_acc, t_pred(n_acc))
7750 }
7751 } else {
7752 // --- SAMPLED ACCEPT (rejection sampling): u_j < p_j(x_j)/q_j(x_j) walk ---
7753 if col_buf.is_none() {
7754 col_buf = Some(e.zeros(n_vocab)?);
7755 }
7756 // FILTERED p_j: per-verify-col stats (one batched filter_stats call), then the
7757 // filtered gather. j==0&&base==0 reads last_col (its own stats row appended).
7758 let mut pj = vec![0f32; k_round.max(1)];
7759 let mut col_stats: Vec<(f32, f32, f32)> = Vec::new(); // (max, th, z) per verify col used
7760 if k_round > 0 {
7761 let mut ids: Vec<u32> = Vec::new();
7762 let mut rows: Vec<i32> = Vec::new();
7763 for j in 0..k_round {
7764 if j > 0 || base == 1 {
7765 ids.push(draft[j]);
7766 rows.push((base + j) as i32 - 1);
7767 }
7768 }
7769 if !ids.is_empty() {
7770 let nr = rows.len();
7771 // penalties: materialize the used columns into one contiguous penalized
7772 // buffer (rows remapped 0..nr) so stats+gathers see the penalized p.
7773 // penalties: materialize used columns contiguously, penalize all rows in
7774 // one launch, and point stats+gathers at the penalized buffer (rows 0..nr).
7775 let p_rows: Vec<i32> = if pen_on {
7776 (0..nr as i32).collect()
7777 } else {
7778 rows.clone()
7779 };
7780 if pen_on {
7781 if pcol_buf.as_ref().map(|b| b.len()).unwrap_or(0) < nr * n_vocab {
7782 pcol_buf = Some(e.zeros(nr * n_vocab)?);
7783 }
7784 let pc = pcol_buf.as_mut().unwrap();
7785 for (i2, &r) in rows.iter().enumerate() {
7786 let c = r as usize;
7787 e.copy_view_into(
7788 pc,
7789 i2 * n_vocab,
7790 &tlogits_d.slice(c * n_vocab..(c + 1) * n_vocab),
7791 n_vocab,
7792 )?;
7793 }
7794 let h = pen_hist_d.as_ref().unwrap();
7795 let nh = h.len();
7796 e.penalize_logits_rows(
7797 pc,
7798 h,
7799 nh,
7800 sp.penalty_repeat,
7801 sp.penalty_freq,
7802 sp.penalty_present,
7803 n_vocab,
7804 nr,
7805 )?;
7806 }
7807 let p_src: &CudaSlice<f32> = if pen_on {
7808 pcol_buf.as_ref().unwrap()
7809 } else {
7810 &tlogits_d
7811 };
7812 let rowsd = e.htod_i32(&p_rows)?;
7813 let (mut th_d, mut z_d, mut mx_d) =
7814 (e.zeros(nr)?, e.zeros(nr)?, e.zeros(nr)?);
7815 e.filter_stats(
7816 p_src, n_vocab, &rowsd, &mut th_d, &mut z_d, &mut mx_d, n_vocab, nr,
7817 sp_temp, sp.top_k, sp.top_p, sp.min_p,
7818 )?;
7819 let idsd = e.htod_u32_v(&ids)?;
7820 let mut outd = e.zeros(nr)?;
7821 e.softmax_gather_filtered(
7822 p_src, n_vocab, &idsd, &rowsd, &th_d, &z_d, &mut outd, n_vocab, nr,
7823 sp_temp,
7824 )?;
7825 let outv = e.dtoh(&outd)?;
7826 let (thv, zv, mxv) = (e.dtoh(&th_d)?, e.dtoh(&z_d)?, e.dtoh(&mx_d)?);
7827 let mut oi = 0usize;
7828 for j in 0..k_round {
7829 if j > 0 || base == 1 {
7830 pj[j] = outv[oi];
7831 oi += 1;
7832 }
7833 }
7834 col_stats = (0..nr).map(|i| (mxv[i], thv[i], zv[i])).collect();
7835 }
7836 if base == 0 {
7837 let lc: &CudaSlice<f32> = if pen_on {
7838 if col_buf.is_none() {
7839 col_buf = Some(e.zeros(n_vocab)?);
7840 }
7841 let cb = col_buf.as_mut().unwrap();
7842 e.copy_into(
7843 cb,
7844 0,
7845 last_col_logits
7846 .as_ref()
7847 .expect("sampled: last_col_logits unset"),
7848 n_vocab,
7849 )?;
7850 let h = pen_hist_d.as_ref().unwrap();
7851 let nh = h.len();
7852 e.penalize_logits(
7853 cb,
7854 h,
7855 nh,
7856 sp.penalty_repeat,
7857 sp.penalty_freq,
7858 sp.penalty_present,
7859 n_vocab,
7860 )?;
7861 col_buf.as_ref().unwrap()
7862 } else {
7863 last_col_logits
7864 .as_ref()
7865 .expect("sampled: last_col_logits unset")
7866 };
7867 let rows0 = e.htod_i32(&[0])?;
7868 let (mut th_d, mut z_d, mut mx_d) = (e.zeros(1)?, e.zeros(1)?, e.zeros(1)?);
7869 e.filter_stats(
7870 lc, n_vocab, &rows0, &mut th_d, &mut z_d, &mut mx_d, n_vocab, 1,
7871 sp_temp, sp.top_k, sp.top_p, sp.min_p,
7872 )?;
7873 let idsd = e.htod_u32_v(&[draft[0]])?;
7874 let mut outd = e.zeros(1)?;
7875 e.softmax_gather_filtered(
7876 lc, n_vocab, &idsd, &rows0, &th_d, &z_d, &mut outd, n_vocab, 1, sp_temp,
7877 )?;
7878 pj[0] = e.dtoh(&outd)?[0];
7879 last_col_stats =
7880 Some((e.dtoh(&mx_d)?[0], e.dtoh(&th_d)?[0], e.dtoh(&z_d)?[0]));
7881 }
7882 }
7883 // q source: the graph arm retained the head logits in the persistent q_slots;
7884 // the eager arm in per-round draft_logits clones. Same raw-logit values either way.
7885 // FILTERED q_j: stats from draft_stats (eager pushes in-chain; the graph arm
7886 // computes them post-replay — graph engages only filter/penalty-free, so the
7887 // stats degenerate to th=0/full-Z there, keeping ONE accept path).
7888 let q_bufs: &[CudaSlice<f32>] = if dctx.graph_s.is_some() {
7889 &dctx.q_slots
7890 } else {
7891 &draft_logits
7892 };
7893 let mut n_acc = 0usize;
7894 for j in 0..k_round {
7895 let (qmx, qth, qz) = draft_stats[j];
7896 let idsd = e.htod_u32_v(&[draft_idx[j]])?;
7897 let rowsd = e.htod_i32(&[0])?;
7898 let thd = e.htod(&[qth])?;
7899 let zd = e.htod(&[qz])?;
7900 let _ = qmx;
7901 let mut outd = e.zeros(1)?;
7902 e.softmax_gather_filtered(
7903 &q_bufs[j], d_vocab, &idsd, &rowsd, &thd, &zd, &mut outd, d_vocab, 1,
7904 sp_temp,
7905 )?;
7906 let qj = e.dtoh(&outd)?[0];
7907 let u = host_u01(sp_seed, uctr);
7908 uctr += 1;
7909 if (u as f64) * (qj as f64) < pj[j] as f64 {
7910 n_acc += 1;
7911 } else {
7912 break;
7913 }
7914 }
7915 let bonus = if n_acc == k_round {
7916 // FULL ACCEPT: bonus ~ FILTERED softmax at the last verify column.
7917 let col = base + k_round - 1;
7918 let cb = col_buf.as_mut().unwrap();
7919 e.copy_view_into(
7920 cb,
7921 0,
7922 &tlogits_d.slice(col * n_vocab..(col + 1) * n_vocab),
7923 n_vocab,
7924 )?;
7925 if pen_on {
7926 let h = pen_hist_d.as_ref().unwrap();
7927 let nh = h.len();
7928 e.penalize_logits(
7929 cb,
7930 h,
7931 nh,
7932 sp.penalty_repeat,
7933 sp.penalty_freq,
7934 sp.penalty_present,
7935 n_vocab,
7936 )?;
7937 }
7938 if perturb_buf.is_none() {
7939 perturb_buf = Some(e.zeros(d_vocab.max(n_vocab))?);
7940 }
7941 // STATS MUST COME FROM THIS COLUMN (bug fix 2026-08-05, lane/sampler-
7942 // truncation-fix; receipts research/sampfix-20260805/). The old code reused
7943 // `col_stats.last()` here, which is ALWAYS the wrong row: the gathered set
7944 // covers verify columns 0..=(base+k_round-2) (rows pushed as base+j-1), while
7945 // the full-accept bonus samples column base+k_round-1 — exactly ONE PAST the
7946 // last gathered column, in both base arms. `th` is a threshold in e-units of
7947 // its OWN row's max, so feeding a neighbour's (row_max, th) into
7948 // gumbel_perturb_filtered mis-scales every e0 = exp((x-row_max)/T). When the
7949 // donor column's peak is higher by more than T*ln(1/th), EVERY id fails
7950 // `e0 >= th`, the whole perturbed row becomes -3.4e38, and the 2-pass argmax
7951 // falls through to its smallest-index tie-break => token id 0 ("!") spliced
7952 // mid-word. Fragility is ordered by how large th is: min_p pins th = min_p
7953 // (0.05 => trigger at delta > 2.4 at T=0.8, fires constantly), top_p's
7954 // mass-boundary th is smaller, top_k's k-th-largest th smaller still — which
7955 // is why the head-to-head matrix saw min_p and top_p corrupt while top_k-only
7956 // stayed clean. The pure-temp default regime is immune (th == 0 masks nothing,
7957 // and row_max is unused once nothing is masked), so this fix is a byte-level
7958 // no-op for the untruncated serve default. One extra one-block filter_stats
7959 // per full-accept round is the whole cost.
7960 let (mx, th) = {
7961 let rows0 = e.htod_i32(&[0])?;
7962 let (mut th_d, mut z_d, mut mx_d) = (e.zeros(1)?, e.zeros(1)?, e.zeros(1)?);
7963 let cb0 = col_buf.as_ref().unwrap();
7964 e.filter_stats(
7965 cb0, n_vocab, &rows0, &mut th_d, &mut z_d, &mut mx_d, n_vocab, 1,
7966 sp_temp, sp.top_k, sp.top_p, sp.min_p,
7967 )?;
7968 (e.dtoh(&mx_d)?[0], e.dtoh(&th_d)?[0])
7969 };
7970 let pb = perturb_buf.as_mut().unwrap();
7971 let cb2 = col_buf.as_ref().unwrap();
7972 e.gumbel_perturb_filtered(cb2, pb, n_vocab, sp_seed, sctr, sp_temp, mx, th)?;
7973 sctr += 1;
7974 let td = e.argmax_token_device(pb, n_vocab)?;
7975 e.dtoh_u32_one(&td)?
7976 } else {
7977 // REJECT at n_acc: bonus ~ norm(max(0, softmax_T(p) - softmax_T(q))).
7978 let cb = col_buf.as_mut().unwrap();
7979 if n_acc > 0 || base == 1 {
7980 let col = base + n_acc - 1;
7981 e.copy_view_into(
7982 cb,
7983 0,
7984 &tlogits_d.slice(col * n_vocab..(col + 1) * n_vocab),
7985 n_vocab,
7986 )?;
7987 } else {
7988 let lc = last_col_logits.as_ref().unwrap();
7989 e.copy_into(cb, 0, lc, n_vocab)?;
7990 }
7991 if pen_on {
7992 let h = pen_hist_d.as_ref().unwrap();
7993 let nh = h.len();
7994 e.penalize_logits(
7995 cb,
7996 h,
7997 nh,
7998 sp.penalty_repeat,
7999 sp.penalty_freq,
8000 sp.penalty_present,
8001 n_vocab,
8002 )?;
8003 }
8004 let cb2 = col_buf.as_ref().unwrap();
8005 let sc = sctr;
8006 sctr += 1;
8007 // p-stats for the reject column: from col_stats when the col was gathered,
8008 // else (j==0&&base==0) from last_col_stats.
8009 let p_stats = if n_acc > 0 || base == 1 {
8010 // col index within the gathered set == number of gathered cols before n_acc
8011 let gi = if base == 1 { n_acc } else { n_acc - 1 };
8012 col_stats.get(gi).copied().unwrap_or_else(|| {
8013 (0.0, 0.0, 1.0) // unreachable: gathered cols always cover the reject slot
8014 })
8015 } else {
8016 last_col_stats.expect("sampled: last_col_stats unset at reject")
8017 };
8018 let q_stats = draft_stats[n_acc];
8019 if let Some(map) = &d2t_dev {
8020 if q_full_buf.is_none() {
8021 q_full_buf = Some(e.zeros(n_vocab)?);
8022 }
8023 let qf = q_full_buf.as_mut().unwrap();
8024 e.scatter_trim_logits(&q_bufs[n_acc], map, qf, d_vocab, n_vocab)?;
8025 let qf2 = q_full_buf.as_ref().unwrap();
8026 e.residual_sample_filtered(
8027 cb2,
8028 Some(qf2),
8029 n_vocab,
8030 sp_temp,
8031 sp_seed,
8032 sc,
8033 p_stats,
8034 q_stats,
8035 &mut sample_tok,
8036 )?;
8037 } else {
8038 e.residual_sample_filtered(
8039 cb2,
8040 Some(&q_bufs[n_acc]),
8041 n_vocab,
8042 sp_temp,
8043 sp_seed,
8044 sc,
8045 p_stats,
8046 q_stats,
8047 &mut sample_tok,
8048 )?;
8049 }
8050 e.dtoh_u32(&sample_tok)?[0]
8051 };
8052 (n_acc, bonus)
8053 };
8054 // --- 3b. GRAMMAR TRUNCATION (constrained spec, 2026-08-03): the grammar is
8055 // an extra rejection rule AFTER the exactness verify (the batched-verify-twins
8056 // ordering). Walk the accepted drafts through the grammar in commit order; the
8057 // first illegal token truncates acceptance at its slot, and that slot's emission
8058 // is recomputed as the MASKED argmax of the target's own verify column — token-
8059 // identical to constrained plain greedy decode (an unmasked argmax that is
8060 // grammar-legal IS the masked argmax: masking only removes competitors). The
8061 // column D2H (~1MB) is paid only when a cut fires — the tight-grammar cost,
8062 // measured in acceptance numbers, never hidden.
8063 let (n_acc, bonus) = match constraint.as_deref_mut() {
8064 None => (n_acc, bonus),
8065 Some(c) => {
8066 fn ce(e2: String) -> Box<dyn std::error::Error> {
8067 format!("constraint: {e2}").into()
8068 }
8069 let mut na = n_acc;
8070 let mut cut = false;
8071 for (j, &d) in draft.iter().enumerate().take(n_acc) {
8072 if c.is_allowed(d).map_err(ce)? {
8073 c.consume(d).map_err(ce)?;
8074 } else {
8075 na = j;
8076 cut = true;
8077 dm_cut_tokens += n_acc - j;
8078 break;
8079 }
8080 }
8081 if cut {
8082 dm_cuts += 1;
8083 }
8084 let mut bo = bonus;
8085 if cut || !c.is_allowed(bo).map_err(ce)? {
8086 let mut row = if na == 0 && base == 0 {
8087 init_logits_host.clone()
8088 .ok_or("constraint: init logits missing (round-0 cut)")?
8089 } else {
8090 e.dtoh_view(&tlogits_d.slice(
8091 (base + na - 1) * n_vocab..(base + na) * n_vocab))?
8092 };
8093 c.mask_logits(&mut row).map_err(ce)?;
8094 bo = argmax(&row) as u32;
8095 }
8096 c.consume(bo).map_err(ce)?;
8097 (na, bo)
8098 }
8099 };
8100 let mut successor_valid = false;
8101 if let Some((q_proxy, expected_d2)) = rejected_probe {
8102 let v_n = n_acc == 1 && bonus == expected_d2;
8103 eprintln!(
8104 "[opti-controller] shadow q={q_proxy:.6} admitted=false v_n={v_n} \
8105 expected_d2={expected_d2} n_acc={n_acc} bonus={bonus}",
8106 );
8107 }
8108 if let Some(successor) = successor_attempt.as_ref() {
8109 successor_valid = n_acc == 1 && bonus == successor.verify_tokens[0];
8110 let generation = successor.generation;
8111 let q_proxy = successor.q_proxy;
8112 let expected_pending = successor.verify_tokens[0];
8113 let resolution_ms = successor.issued_at.elapsed().as_secs_f64() * 1e3;
8114 let fork = opti_fork
8115 .as_mut()
8116 .ok_or("optipipe successor resolution lost fork state")?;
8117 fork.finish_actual_reconcile(
8118 e,
8119 &mut *cache,
8120 &snap,
8121 n_acc,
8122 base,
8123 successor_valid,
8124 )?;
8125 if successor_valid {
8126 OPTI_FORK_HITS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
8127 } else {
8128 OPTI_FORK_MISSES.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
8129 OPTI_RECONCILES.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
8130 OPTI_WASTED_DRAFT_TOKENS
8131 .fetch_add(2, std::sync::atomic::Ordering::Relaxed);
8132 }
8133 let breaker_tripped = fork
8134 .controller
8135 .as_mut()
8136 .expect("controller policy")
8137 .resolve(successor_valid);
8138 if breaker_tripped {
8139 OPTI_BREAKER_TRIPS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
8140 }
8141 eprintln!(
8142 "[opti-controller] resolve generation={} hit={} q={q_proxy:.6} \
8143 expected_pending={expected_pending} n_acc={n_acc} bonus={bonus} \
8144 resolution_ms={resolution_ms:.3} reconcile={} breaker={}",
8145 generation.id,
8146 successor_valid,
8147 !successor_valid,
8148 breaker_tripped,
8149 );
8150 if !successor_valid {
8151 let mut successor = successor_attempt
8152 .take()
8153 .expect("controller successor disappeared on miss");
8154 successor.settle();
8155 fork.retire(generation)?;
8156 }
8157 }
8158 total_drafted += k_round;
8159 total_accepted += n_acc;
8160 if let Some(t) = sess_telem {
8161 // Greedy, rejection-sampling, and grammar truncation all converge here after
8162 // the accept decision is already on host. Fixed-size relaxed atomics only.
8163 t.record_round(k_round, n_acc);
8164 }
8165 if spec_stats {
8166 st_len_hist[k_round] += 1;
8167 for j in 0..k_round {
8168 st_drafted[j] += 1;
8169 }
8170 for j in 0..n_acc {
8171 st_accepted[j] += 1;
8172 }
8173 if n_acc == k_round {
8174 st_full += 1;
8175 }
8176 }
8177
8178 if debug_spec {
8179 eprintln!("[R{round}] pos={pos} out_len={} last_tok={last_token} draft={draft:?} n_acc={n_acc} bonus={bonus} t_pred0={}", out.len(), t_pred(0));
8180 }
8181
8182 // --- 4. COMMIT: draft[0..n_acc] then bonus (n_acc + 1 tokens) ---
8183 let commit_started = std::time::Instant::now();
8184 // SESSION MODE: every accepted column is already in the CACHE — `out` must carry all
8185 // of them (overshoot past max_new included) or `committed` under-counts the cache rows
8186 // and the next turn's continuation seeds one token off (gate-caught 2026-07-05). The
8187 // single-shot path keeps the cap (its caller truncates + drops the cache anyway).
8188 for j in 0..n_acc {
8189 if !session_mode && out.len() >= max_new {
8190 break;
8191 }
8192 out.push(draft[j]);
8193 }
8194 if pen_on {
8195 pen_hist.extend_from_slice(&draft[0..n_acc]);
8196 pen_hist.push(bonus);
8197 }
8198 let bonus_emitted = session_mode || out.len() < max_new;
8199 if bonus_emitted {
8200 out.push(bonus);
8201 }
8202 last_token = bonus;
8203
8204 // --- 5. ROLLBACK + advance (§C) ---
8205 if n_acc == k_round && !spec_replay {
8206 // FULL ACCEPT, BONUS FOLD: all verify columns (pending? + drafts) are committed in
8207 // cache; the NEW bonus stays PENDING for the next round's verify batch — NO extra
8208 // T=1 trunk pass. The next draft chain seeds from the MTP block's h_nextn at the
8209 // bonus position: one MTP-block pass (~1/33 trunk cost) replaces the trunk read.
8210 // last_pred is dead in the pending path (t_pred reads verify col 0).
8211 //
8212 // PERSISTENT DRAFT KV, full-accept fill: the chain covered last_token +
8213 // draft[0..k_round-2] as INPUTS (slots P..P'-2); draft[k_round-1] (slot P'-1) was
8214 // only ever an output, so its entry is MISSING. Fill it from vh_seed — its EXACT
8215 // trunk hidden (the last verify column). set_len first: a p-min break may have
8216 // left one extra chain append at that slot. Partial accepts need NO fill (the
8217 // chain already covered every accepted position; round-start set_len truncates).
8218 let mut vh_seed = e.zeros(n_embd)?;
8219 e.copy_view_into(
8220 &mut vh_seed,
8221 0,
8222 &vx.slice((t_v - 1) * n_embd..t_v * n_embd),
8223 n_embd,
8224 )?;
8225 if refresh {
8226 // TRUE-HIDDEN REFRESH (2026-07-03, the HANDOVER-listed acceptance lever):
8227 // overwrite ALL committed positions' scratch entries with K/V from their EXACT
8228 // verify hiddens — the reference engine's mtp_update fills from true hiddens;
8229 // the full stack (vx) is already resident from the verify. Replaces both the
8230 // chain-approximate entries AND the old last-token-only fill. Acceptance-only
8231 // (draft attention quality); exactness stays the verify's job.
8232 scratch.set_len(e, pos)?;
8233 // PREDECESSOR pairing: row i gets vx[i-1]; row 0 the carried fill_prev
8234 // (hidden of the last committed row before this verify batch).
8235 let mut vxs = e.zeros(t_v * n_embd)?;
8236 e.copy_into(&mut vxs, 0, &fill_prev, n_embd)?;
8237 if t_v > 1 {
8238 e.copy_view_into(
8239 &mut vxs,
8240 n_embd,
8241 &vx.slice(0..(t_v - 1) * n_embd),
8242 (t_v - 1) * n_embd,
8243 )?;
8244 }
8245 self.mtp_kv_fill(e, mtp, &verify_tokens, &vxs, pos, &mut *scratch, embd_dev)?;
8246 } else {
8247 scratch.set_len(e, pos + base + k_round - 1)?;
8248 // predecessor of the last draft = verify col t_v-2 (or fill_prev at t_v==1)
8249 let mut hp = e.zeros(n_embd)?;
8250 if t_v >= 2 {
8251 e.copy_view_into(
8252 &mut hp,
8253 0,
8254 &vx.slice((t_v - 2) * n_embd..(t_v - 1) * n_embd),
8255 n_embd,
8256 )?;
8257 } else {
8258 e.copy_into(&mut hp, 0, &fill_prev, n_embd)?;
8259 }
8260 self.mtp_kv_fill(
8261 e,
8262 mtp,
8263 &[draft[k_round - 1]],
8264 &hp,
8265 pos + base + k_round - 1,
8266 &mut *scratch,
8267 embd_dev,
8268 )?;
8269 }
8270 // REFERENCE SEEDING: no pseudo pass — the next chain's step 0 IS the
8271 // reference's (id_last, h_prev) draft row; it appends the bonus's scratch
8272 // entry itself. Seed = TRUE hidden of the bonus's predecessor (last verify
8273 // col). Saves one MTP-block pass per round on top of the pairing fix.
8274 if !devacc_seeded {
8275 e.copy_into(&mut h_seed_buf, 0, &vh_seed, n_embd)?;
8276 e.copy_into(&mut fill_prev, 0, &vh_seed, n_embd)?;
8277 }
8278 pending = Some(bonus);
8279 if debug_spec {
8280 eprintln!(" -> FULL ACCEPT (bonus pending, prev-h seed)");
8281 }
8282 } else if !spec_replay && base + n_acc >= 1 {
8283 // PARTIAL ACCEPT, REPLAY-FREE (2026-07-03 — the profiled #1 long-ctx spec cost):
8284 // the verify's first j = base+n_acc columns ARE the committed sequence, computed
8285 // bit-identically to eager (decode-exact contract) — so KEEP them: KV truncates to
8286 // pos+j, recurrent state rebuilds from the VerifyCkpt (same-kernel gdn prefix
8287 // re-run / pure state-clone restore), and the bonus stays PENDING exactly like the
8288 // full-accept path — the legacy duplicate trunk replay is gone. The next chain
8289 // seeds from the MTP pseudo-hidden of the bonus, whose seed = the TRUE verify
8290 // hidden of its predecessor (col j-1) — same one-hop pseudo structure as full
8291 // accept (never compounds: the next verify recomputes true hiddens for all
8292 // committed columns).
8293 let j = base + n_acc;
8294 self.commit_verified_prefix(
8295 e,
8296 &mut *cache,
8297 &snap,
8298 ckpt.as_ref().unwrap(),
8299 j,
8300 devacc_seeded,
8301 if devacc_seeded {
8302 devacc_acc.as_ref().map(|a| (a, base, t_v))
8303 } else {
8304 None
8305 },
8306 )?;
8307 let mut seed = e.zeros(n_embd)?;
8308 e.copy_view_into(
8309 &mut seed,
8310 0,
8311 &vx.slice((j - 1) * n_embd..j * n_embd),
8312 n_embd,
8313 )?;
8314 // Draft scratch: TRUE-HIDDEN REFRESH of the committed prefix (see the full-accept
8315 // branch); without it the chain entries stand and only the tail truncates. Either
8316 // way len ends at pos+j so the pseudo append lands at the bonus's slot pos+j
8317 // (persistent mode), rope pos+j+1 (chain convention).
8318 if refresh {
8319 scratch.set_len(e, pos)?;
8320 let mut vxs = e.zeros(j * n_embd)?;
8321 e.copy_into(&mut vxs, 0, &fill_prev, n_embd)?;
8322 if j > 1 {
8323 e.copy_view_into(
8324 &mut vxs,
8325 n_embd,
8326 &vx.slice(0..(j - 1) * n_embd),
8327 (j - 1) * n_embd,
8328 )?;
8329 }
8330 self.mtp_kv_fill(
8331 e,
8332 mtp,
8333 &verify_tokens[0..j],
8334 &vxs,
8335 pos,
8336 &mut *scratch,
8337 embd_dev,
8338 )?;
8339 } else {
8340 scratch.set_len(e, pos + j)?;
8341 }
8342 // REFERENCE SEEDING (see the full-accept branch): seed = TRUE hidden of the
8343 // bonus's predecessor (verify col j-1); no pseudo pass.
8344 if !devacc_seeded {
8345 e.copy_into(&mut h_seed_buf, 0, &seed, n_embd)?;
8346 e.copy_into(&mut fill_prev, 0, &seed, n_embd)?;
8347 }
8348 pending = Some(bonus);
8349 if debug_spec {
8350 eprintln!(" -> PARTIAL(replay-free j={j}, bonus pending, prev-h seed)");
8351 }
8352 } else if !spec_replay {
8353 // ZERO ROUND FOLD (2026-07-10, verify-cost target #3): base+n_acc == 0 — a
8354 // pending-less round where nothing was accepted (PMIN0 zero-draft chains after a
8355 // replay/commit, or plain 0-accept rounds at round 0). The old path replayed
8356 // [bonus] through a FULL m=1 trunk+head forward (the 489us full-vocab head pass
8357 // measured at ~0.75/round on PMIN0 configs). Instead: restore the pre-round
8358 // snapshot and let the bonus ride the NEXT round's verify as col 0 — the existing
8359 // base=1 pending machinery, bit-identical by the decode-exact verify contract.
8360 // Seed: the bonus's predecessor is the last COMMITTED token, whose hidden
8361 // fill_prev already carries (same seeding as the 1-token-replay case it replaces).
8362 cache.rollback(e, &snap, 0)?;
8363 scratch.set_len(e, pos)?;
8364 e.copy_into(&mut h_seed_buf, 0, &fill_prev, n_embd)?;
8365 pending = Some(bonus);
8366 if debug_spec {
8367 eprintln!(" -> ZERO-ROUND FOLD (bonus pending, fill_prev seed)");
8368 }
8369 } else {
8370 // PARTIAL ACCEPT, LEGACY REPLAY (seam MEMRA_SPEC_REPLAY=1 — or j==0: nothing of
8371 // this round survives, only possible before the first pending exists, ~round 0):
8372 // restore EVERYTHING to the pre-round snapshot (KV truncate to pos + recur
8373 // restore), then replay the committed prefix pending? ++ draft[0..n_acc] ++
8374 // [bonus] as ONE batched T forward — single weight read, bit-identical to greedy
8375 // (the verify-all-columns path is the same math). Commits the bonus with a TRUE
8376 // trunk hidden.
8377 cache.rollback(e, &snap, 0)?; // accept_len=0: KV len = pos, recur = snapshot
8378 let mut replay: Vec<u32> = Vec::with_capacity(base + n_acc + 1);
8379 if let Some(b) = pending.take() {
8380 replay.push(b);
8381 }
8382 replay.extend_from_slice(&draft[0..n_acc]);
8383 replay.push(bonus);
8384 // Full-stack forward (decode_step_t_core = decode_step_t_h_emb_dev's body):
8385 // Predecessor pairing seeds from the PREDECESSOR row (col len-2) — the same-row path takes the
8386 // last col exactly as before (byte-identical to the old _h_emb_dev call).
8387 let (rl_d, rx) =
8388 if matches!(self.cfg.arch, memra_gguf::config::Arch::Qwen35Moe) {
8389 let mut logits = Vec::with_capacity(replay.len() * n_vocab);
8390 let mut hidden = e.uninit(replay.len() * n_embd)?;
8391 for (row, &token) in replay.iter().enumerate() {
8392 let (row_logits, row_hidden) =
8393 self.spec_target_step_h(e, token, &mut *cache)?;
8394 logits.extend_from_slice(&row_logits);
8395 e.dtod_copy_into(&row_hidden, &mut hidden, row * n_embd)?;
8396 }
8397 (e.htod(&logits)?, hidden)
8398 } else {
8399 self.decode_step_t_core(e, &replay, pos, &mut *cache, embd_dev, None)?
8400 };
8401 // last_pred = argmax of the LAST column's logits (predicts the token after `bonus`)
8402 // — device argmax + one 4-byte read instead of the full-vocab column dtoh.
8403 e.argmax_token_device_col(&rl_d, replay.len() - 1, n_vocab, &mut preds_d, 0)?;
8404 last_pred = e.dtoh_u32(&preds_d)?[0];
8405 if sampled {
8406 let lr0 = replay.len();
8407 let lc = last_col_logits
8408 .as_mut()
8409 .expect("sampled: last_col_logits unset");
8410 e.copy_view_into(
8411 lc,
8412 0,
8413 &rl_d.slice((lr0 - 1) * n_vocab..lr0 * n_vocab),
8414 n_vocab,
8415 )?;
8416 }
8417 let lr = replay.len();
8418 if lr >= 2 {
8419 e.copy_view_into(
8420 &mut h_seed_buf,
8421 0,
8422 &rx.slice((lr - 2) * n_embd..(lr - 1) * n_embd),
8423 n_embd,
8424 )?;
8425 } else {
8426 // 1-token replay (round-0 miss): the bonus's predecessor is the OLD
8427 // last_token, whose own-row hidden fill_prev still holds.
8428 e.copy_into(&mut h_seed_buf, 0, &fill_prev, n_embd)?;
8429 }
8430 // the bonus is COMMITTED here — it becomes the last committed row.
8431 let mut rh_last = e.zeros(n_embd)?;
8432 e.copy_view_into(
8433 &mut rh_last,
8434 0,
8435 &rx.slice((lr - 1) * n_embd..lr * n_embd),
8436 n_embd,
8437 )?;
8438 e.copy_into(&mut fill_prev, 0, &rh_last, n_embd)?;
8439 if debug_spec {
8440 eprintln!(" -> PARTIAL(replay={replay:?}), next_pred={last_pred}");
8441 }
8442 }
8443 if devacc_seeded {
8444 // stage (b) epilogue: fill_prev takes the gathered seed AFTER the refresh fills
8445 // consumed the old value (both slots carry the same value in every non-replay arm).
8446 e.copy_into(&mut fill_prev, 0, &h_seed_buf, n_embd)?;
8447 }
8448 if successor_valid {
8449 let optimistic_scratch_len = successor_attempt
8450 .as_ref()
8451 .expect("valid controller successor disappeared")
8452 .scratch_len;
8453 // The normal current-round commit refreshed/truncated the logical scratch tail.
8454 // Its optimistic successor row was already written physically, so restoring only
8455 // the retained logical length makes that row live for the carried round.
8456 scratch.set_len(e, optimistic_scratch_len)?;
8457 }
8458 if let Some(current) = current_opti.take() {
8459 opti_fork
8460 .as_mut()
8461 .ok_or("optipipe current retirement lost fork state")?
8462 .retire(current.generation)?;
8463 }
8464 if successor_valid {
8465 let successor = successor_attempt
8466 .take()
8467 .expect("valid controller successor disappeared before promotion");
8468 let generation = successor.generation;
8469 opti_fork
8470 .as_mut()
8471 .ok_or("optipipe successor promotion lost fork state")?
8472 .promote_successor_snapshot(&mut snap, generation);
8473 carried_opti = Some(successor);
8474 }
8475 if anatomy_on {
8476 // Commit/rollback is normally asynchronous on the primary/head stream. Bound it
8477 // only for this diagnostic so it does not disappear into the following draft's
8478 // first token readback.
8479 e.stream().synchronize()?;
8480 ph_commit += commit_started.elapsed().as_secs_f64();
8481 }
8482 // adaptive-K update (host math, zero syncs): next round drafts accepted-run + 1,
8483 // clamped to [floor(pos), k_cap]. cache.pos is post-rollback here (the round's
8484 // final position — the floor's position key reads the committed depth). Burst
8485 // rounds (`continue` above) draft the captured fixed depth and skip this, exactly
8486 // like gemma's burst arm.
8487 if adapt {
8488 let fl_now = floor_at(cache.pos);
8489 kc = (n_acc + 1).clamp(fl_now.min(k_cap), k_cap);
8490 }
8491 ph_mark(&mut ph_rest, phase_on);
8492 if let Some(p) = pipe {
8493 p.accept_end(round);
8494 }
8495 drop(pipe_accept);
8496 round += 1;
8497 // sse-cadence: this round's accepted drafts + bonus are committed (out is
8498 // append-only past step 4) — flush at round cadence.
8499 keep_going = flush_commit(&mut on_commit, &out, &mut flushed);
8500 }
8501 if let Some(mut ticket) = carried_opti.take() {
8502 opti_fork
8503 .as_mut()
8504 .ok_or("optipipe tail drain lost fork state")?
8505 .cancel_controller_ticket(e, &mut *cache, &mut *scratch, &snap, &mut ticket)?;
8506 }
8507 // sse-cadence: nothing below appends to `out`; flush any remainder (defensive).
8508 // (verdict ignored — the burst is over either way; the session tail runs unchanged.)
8509 let _ = flush_commit(&mut on_commit, &out, &mut flushed);
8510
8511 if spec_stats {
8512 let per_slot: Vec<String> = (0..k)
8513 .map(|j| {
8514 if st_drafted[j] > 0 {
8515 format!(
8516 "{}/{}={:.3}",
8517 st_accepted[j],
8518 st_drafted[j],
8519 st_accepted[j] as f64 / st_drafted[j] as f64
8520 )
8521 } else {
8522 "0/0".into()
8523 }
8524 })
8525 .collect();
8526 let acc = if total_drafted > 0 {
8527 total_accepted as f64 / total_drafted as f64
8528 } else {
8529 0.0
8530 };
8531 eprintln!(
8532 "[spec-stats] rounds={round} full_accept={st_full} len_hist={st_len_hist:?} \
8533 per_slot=[{}] total={total_accepted}/{total_drafted}={acc:.3} \
8534 tok_per_round={:.3}",
8535 per_slot.join(" "),
8536 (total_accepted + round) as f64 / round.max(1) as f64
8537 );
8538 }
8539 if constraint.is_some() {
8540 eprintln!(
8541 "[draft-mask] mask_rounds={dm_rounds} clone_total={:.3}ms \
8542 clone_per_round={:.4}ms gram_cuts={dm_cuts}/{round} cut_tokens={dm_cut_tokens}",
8543 dm_clone_ns as f64 / 1e6,
8544 dm_clone_ns as f64 / 1e6 / dm_rounds.max(1) as f64
8545 );
8546 }
8547 if phase_on {
8548 let tot = ph_draft + ph_verify + ph_wait + ph_rest;
8549 eprintln!("[spec-phase] draft={:.1}ms ({:.1}%) verify-issue={:.1}ms ({:.1}%) verify-wait={:.1}ms ({:.1}%) commit-host={:.1}ms ({:.1}%) rounds={round}",
8550 ph_draft * 1e3, ph_draft / tot * 100.0,
8551 ph_verify * 1e3, ph_verify / tot * 100.0,
8552 ph_wait * 1e3, ph_wait / tot * 100.0,
8553 ph_rest * 1e3, ph_rest / tot * 100.0);
8554 }
8555 if anatomy_on {
8556 let rounds_f = round.max(1) as f64;
8557 let other = (ph_rest - ph_commit).max(0.0);
8558 eprintln!(
8559 "[spec-anatomy] per-round draft={:.3}ms pp-verify={:.3}ms \
8560 verify-accept={:.3}ms commit-rollback={:.3}ms other={:.3}ms rounds={round}",
8561 ph_draft * 1e3 / rounds_f,
8562 ph_verify * 1e3 / rounds_f,
8563 ph_wait * 1e3 / rounds_f,
8564 ph_commit * 1e3 / rounds_f,
8565 other * 1e3 / rounds_f,
8566 );
8567 }
8568 let _pipe_tail = pipe.map(|p| p.primary());
8569 // SESSION TAIL: leave the session in the exact invariant the next turn's suffix prime
8570 // expects — every row in `committed` has trunk KV/recur state AND an exact draft-KV row.
8571 // Park the draft-graph ctx back on the session (the serve-burst fixed-cost fix): the next
8572 // burst replays instead of recapturing. Error paths (`?` above) drop it — recaptured then.
8573 if let Some(slot) = sess_draft_slot.take() {
8574 *slot = Some(dctx);
8575 }
8576 let t_rounds = t_ent.elapsed();
8577 if let Some((committed, last_h, next_pred_slot, sctr_slot, uctr_slot)) = sess_tail.take() {
8578 *sctr_slot = sctr;
8579 *uctr_slot = uctr;
8580 *next_pred_slot = Some(last_pred);
8581 let mut stashed_pending = false;
8582 if let Some(b) = pending.take() {
8583 if !sampled {
8584 // PENDING-CARRY (2026-08-01): stash the bonus on the session instead of
8585 // committing it with a solo T=1 pass — the next empty-suffix greedy burst
8586 // consumes it as round-0 verify col 0 (a plain round edge; the old tail
8587 // commit + next burst's init feed were 11.6+11.5ms solo trunk passes per
8588 // burst on H100 q27, [spec-setup] trace). b stays in `out` (emitted) but
8589 // OUT of `committed` (cache rows == committed); the consuming call
8590 // prepends it once its verify commits the row. next_pred is unknowable
8591 // without the commit pass — None; callers gate on pending_tok too.
8592 debug_assert_eq!(out.last(), Some(&b), "pending must be the last emitted");
8593 if let Some(slot) = sess_pending_slot.take() {
8594 *slot = Some(b);
8595 }
8596 *next_pred_slot = None;
8597 // fill_prev = hidden of the last COMMITTED row (b's predecessor) — the
8598 // exact chain-seed/fill anchor the consuming burst (or a flush) needs.
8599 *last_h = Some(e.clone_dtod(&fill_prev)?);
8600 stashed_pending = true;
8601 } else {
8602 // SAMPLED tail (unchanged): commit the bonus (one T=1 pass) + draft fill —
8603 // the sampled round-0 accept needs this pass's logits (last_col_logits).
8604 let pos_b = cache.pos;
8605 scratch.set_len(e, pos_b)?;
8606 let (lg_b, hb) = self.spec_target_step_h(e, b, &mut *cache)?;
8607 // after a FULL-accept exit `last_pred` is STALE (it predicted the bonus
8608 // itself — the prediction AFTER the bonus never materialized; it would have
8609 // been the next round's verify col 0). The commit's logits ARE that
8610 // prediction.
8611 *next_pred_slot = Some(argmax(&lg_b) as u32);
8612 self.mtp_kv_fill(e, mtp, &[b], &fill_prev, pos_b, &mut *scratch, embd_dev)?;
8613 *last_h = Some(hb);
8614 }
8615 } else {
8616 // fill_prev tracks the hidden of the last COMMITTED row throughout the loop.
8617 *last_h = Some(e.clone_dtod(&fill_prev)?);
8618 }
8619 committed.extend_from_slice(prompt);
8620 if let Some(cb) = carried_pending {
8621 // the consumed carry's cache row landed in round 0's verify (every pending
8622 // round commits col 0) — it joins `committed` here, in sequence order.
8623 committed.push(cb);
8624 }
8625 if stashed_pending {
8626 // ZERO-EMIT BURST (2026-08-06 c=8 serve panic, pre-existing since b4aea184):
8627 // `out.len() - 1` underflowed on an EMPTY `out` — "range end index
8628 // 18446744073709551615 out of range for slice of length 0", killing the
8629 // memra-gpu-worker and failing 31 of 32 concurrent requests with "worker closed
8630 // stream". Reachable because `pending` starts as `carried_pending` (a bonus
8631 // stashed by the PREVIOUS burst) while `out` starts empty, and the carry is
8632 // deliberately NOT pushed to `out` (line ~3239: the burst that emitted it already
8633 // did). So a burst that stashes a pending without emitting anything of its own —
8634 // the round loop exits before a push, e.g. the ring drain's `out.len() < max_new`
8635 // guard skipping every token under a tight budget — arrives here with
8636 // out.len() == 0 and stashed_pending == true.
8637 //
8638 // The invariant is unchanged: `committed` gets every emitted token EXCEPT the
8639 // stashed bonus. With nothing emitted, that is nothing — and the carry pushed
8640 // just above is already accounted. Saturating, not a min/assert: an empty `out`
8641 // here is a legitimate burst shape, not a corrupt state.
8642 let emitted = out.len().saturating_sub(1);
8643 committed.extend_from_slice(&out[..emitted]);
8644 } else {
8645 committed.extend_from_slice(&out); // FULL out incl. overshoot — all committed
8646 }
8647 debug_assert_eq!(
8648 cache.pos,
8649 committed.len(),
8650 "session invariant: cache rows == committed tokens"
8651 );
8652 if setup_trace {
8653 e.stream().synchronize()?; // bound the async tail fill in the trace
8654 let t_tail = t_ent.elapsed();
8655 eprintln!(
8656 "[spec-setup] init={:.2}ms cap={:.2}ms fill={:.2}ms rounds={:.2}ms tail={:.2}ms total={:.2}ms out={} cont={}",
8657 t_init.as_secs_f64() * 1e3,
8658 (t_cap - t_init).as_secs_f64() * 1e3,
8659 (t_fill - t_cap).as_secs_f64() * 1e3,
8660 (t_rounds - t_fill).as_secs_f64() * 1e3,
8661 (t_tail - t_rounds).as_secs_f64() * 1e3,
8662 t_tail.as_secs_f64() * 1e3,
8663 out.len(),
8664 continuation
8665 );
8666 }
8667 return Ok((out, total_drafted, total_accepted));
8668 }
8669 out.truncate(max_new);
8670 Ok((out, total_drafted, total_accepted))
8671 }
8672
8673 /// Anchor-bounded DSpark target extraction. The trunk sees the exact generated token tape;
8674 /// only requested hidden rows and target-logit rows cross PCIe. An anchor token at p pairs
8675 /// with the pre-output-norm h[p-1] carrier, exactly as the existing replay/NextN path does.
8676 pub fn extract_dspark_anchors(
8677 &self,
8678 e: &Engine,
8679 tokens: &[u32],
8680 anchor_positions: &[usize],
8681 gamma: usize,
8682 top_k: usize,
8683 chunk: usize,
8684 temperature: f32,
8685 ) -> Result<Vec<DsparkAnchorRecord>, Box<dyn std::error::Error>> {
8686 if tokens.len() < gamma + 2 || gamma == 0 || chunk < 2 {
8687 return Err("DSpark extraction token tape/gamma/chunk is invalid".into());
8688 }
8689 if anchor_positions.windows(2).any(|pair| pair[0] >= pair[1]) {
8690 return Err("DSpark anchor positions must be sorted and unique".into());
8691 }
8692 for &position in anchor_positions {
8693 if position == 0 || position + gamma >= tokens.len() {
8694 return Err(format!(
8695 "DSpark anchor {position} has no predecessor or cannot cover gamma={gamma} in {} tokens",
8696 tokens.len()
8697 )
8698 .into());
8699 }
8700 }
8701
8702 let n_vocab = self.output.out_features();
8703 let n_embd = self.cfg.n_embd as usize;
8704 let mut cache = crate::pp::new_cache(e, &self.cfg, tokens.len() + gamma + 8)?;
8705 let (embd_qt, embd_rb) = self.embd.qt_and_row_bytes(n_embd);
8706 let embd_gpu = if spec_host_embd() {
8707 None
8708 } else {
8709 Some(
8710 self.embd_gpu
8711 .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload")),
8712 )
8713 };
8714 let embd_dev = embd_gpu.map(|gpu| (gpu, embd_qt, embd_rb));
8715
8716 struct PendingRecord {
8717 position: usize,
8718 hidden: Option<Vec<f32>>,
8719 tokens: Vec<u32>,
8720 target_top_ids: Vec<Option<Vec<u32>>>,
8721 target_top_logits: Vec<Option<Vec<f32>>>,
8722 target_top_probs: Vec<Option<Vec<f32>>>,
8723 target_tail_probs: Vec<Option<f32>>,
8724 }
8725
8726 let mut pending: Vec<PendingRecord> = anchor_positions
8727 .iter()
8728 .map(|&position| PendingRecord {
8729 position,
8730 hidden: None,
8731 tokens: tokens[position..=position + gamma].to_vec(),
8732 target_top_ids: vec![None; gamma],
8733 target_top_logits: vec![None; gamma],
8734 target_top_probs: vec![None; gamma],
8735 target_tail_probs: vec![None; gamma],
8736 })
8737 .collect();
8738
8739 let mut start = 0usize;
8740 while start < tokens.len() {
8741 let end = (start + chunk).min(tokens.len());
8742 let chunk_tokens = &tokens[start..end];
8743 let (target_logits, hidden_rows) =
8744 self.decode_step_t_core(e, chunk_tokens, start, &mut cache, embd_dev, None)?;
8745 for record in &mut pending {
8746 let hidden_position = record.position - 1;
8747 if hidden_position >= start && hidden_position < end {
8748 let local = hidden_position - start;
8749 record.hidden = Some(e.dtoh_view(
8750 &hidden_rows.slice(local * n_embd..(local + 1) * n_embd),
8751 )?);
8752 }
8753 for slot in 0..gamma {
8754 let target_row = record.position + slot;
8755 if target_row < start || target_row >= end {
8756 continue;
8757 }
8758 let local = target_row - start;
8759 let logits = e.dtoh_view(
8760 &target_logits.slice(local * n_vocab..(local + 1) * n_vocab),
8761 )?;
8762 let (ids, top_logits, probs, tail) =
8763 dspark_sparse_softmax_topk(&logits, top_k, temperature)?;
8764 record.target_top_ids[slot] = Some(ids);
8765 record.target_top_logits[slot] = Some(top_logits);
8766 record.target_top_probs[slot] = Some(probs);
8767 record.target_tail_probs[slot] = Some(tail);
8768 }
8769 }
8770 start = end;
8771 }
8772
8773 pending
8774 .into_iter()
8775 .map(|record| {
8776 let hidden = record
8777 .hidden
8778 .ok_or_else(|| format!("missing DSpark hidden at {}", record.position))?;
8779 let target_top_ids =
8780 flatten_dspark_rows(record.target_top_ids, record.position, "target ids")?;
8781 let target_top_logits = flatten_dspark_rows(
8782 record.target_top_logits,
8783 record.position,
8784 "target logits",
8785 )?;
8786 let target_top_probs = flatten_dspark_rows(
8787 record.target_top_probs,
8788 record.position,
8789 "target probs",
8790 )?;
8791 let target_tail_probs = record
8792 .target_tail_probs
8793 .into_iter()
8794 .enumerate()
8795 .map(|(slot, value)| {
8796 value.ok_or_else(|| {
8797 format!("missing DSpark tail at {} slot {slot}", record.position)
8798 })
8799 })
8800 .collect::<Result<Vec<_>, _>>()?;
8801 Ok(DsparkAnchorRecord {
8802 position: record.position,
8803 hidden,
8804 tokens: record.tokens,
8805 target_top_ids,
8806 target_top_logits,
8807 target_top_probs,
8808 target_tail_probs,
8809 })
8810 })
8811 .collect()
8812 }
8813
8814 /// TEACHER-FORCED REPLAY ACCEPTANCE (hqmtp MTP-heal protocol): walk a FIXED token
8815 /// sequence and, at sampled positions, compare the MTP head's K-token draft chain against
8816 /// the trunk's own teacher-forced greedy predictions. Nothing is generated — the context is
8817 /// the corpus text itself, so (a) degenerate self-generated loops cannot inflate acceptance
8818 /// and (b) two arms (bf16 ceiling vs NVFP4) score on IDENTICAL contexts, isolating the
8819 /// quant-induced head/hidden-state mismatch from text drift.
8820 ///
8821 /// Per eval position p (context = tokens[0..=p], predecessor pairing as in spec decode):
8822 /// draft_j = chain token j from (tokens[p], h_{p-1}), then its own drafts — the exact
8823 /// eager spec-decode chain (same mtp_head_forward_dev, same rope positions).
8824 /// target_j = teacher-forced greedy pick for position p+1+j (argmax of the trunk logits
8825 /// at forced context tokens[0..p+j]). For j==0 this equals live spec
8826 /// acceptance; for j>=1 live verify would condition on the drafts, here it
8827 /// conditions on the corpus — deterministic and arm-comparable by design.
8828 ///
8829 /// Returns (rows, bg): one (p, drafts[k], targets[k]) row per eval position (ascending p),
8830 /// plus the full teacher-forced greedy track bg (bg[i] = greedy pick for position i, i>=1)
8831 /// so harnesses can cross-check runs (e.g. different chunk sizes must give identical bg).
8832 ///
8833 /// `hdump`: when Some, every position's pre-output_norm trunk hidden (the exact rows the
8834 /// draft-KV fill pairs from) streams to the file as little-endian f32 [t_total, n_embd] —
8835 /// the head-distillation extraction (hqmtp): the ENGINE is the source of truth for trunk
8836 /// hiddens (HF torch reproductions of the hybrid trunk measured only ~0.5 greedy
8837 /// agreement vs this path — not usable as a training-data source).
8838 pub fn replay_acceptance(
8839 &self,
8840 e: &Engine,
8841 tokens: &[u32],
8842 k: usize,
8843 stride: usize,
8844 chunk: usize,
8845 mut hdump: Option<&mut std::fs::File>,
8846 ) -> Result<(Vec<(usize, Vec<u32>, Vec<u32>)>, Vec<u32>), Box<dyn std::error::Error>> {
8847 assert!(k >= 1 && stride >= 1 && chunk >= 2);
8848 let mtp = self
8849 .mtp
8850 .as_ref()
8851 .expect("replay_acceptance requires an MTP head");
8852 let n_vocab = self.output.out_features();
8853 let d_vocab = mtp
8854 .shared_head_head
8855 .as_ref()
8856 .unwrap_or(&self.output)
8857 .out_features();
8858 let n_embd = self.cfg.n_embd as usize;
8859 let t_total = tokens.len();
8860 assert!(t_total >= 8, "corpus too short ({t_total} tokens)");
8861 // STAGE-OWNED KV (lane/pp2-spec 2026-08-06) — see `new_session`. Door shut = `Cache::new`.
8862 let mut cache = crate::pp::new_cache(e, &self.cfg, t_total + k + 8)?;
8863 let mut scratch = MtpScratch::new(
8864 e,
8865 &self.cfg,
8866 t_total + k + 8,
8867 self.mtp.as_ref().and_then(|m| m.geom.as_ref()),
8868 )?;
8869 let (embd_qt, embd_rb) = self.embd.qt_and_row_bytes(n_embd);
8870 let embd_gpu = if spec_host_embd() {
8871 None
8872 } else {
8873 Some(
8874 self.embd_gpu
8875 .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload")),
8876 )
8877 };
8878 let embd_dev = embd_gpu.map(|g| (g, embd_qt, embd_rb));
8879
8880 // bg[i] = the trunk's greedy pick for position i under the forced context (i >= 1).
8881 let mut bg: Vec<u32> = vec![0; t_total + 1];
8882 let mut rows: Vec<(usize, Vec<u32>, Vec<u32>)> = Vec::new();
8883 let mut prev_last_h = e.zeros(n_embd)?; // predecessor hidden entering the chunk
8884 let mut seed_buf = e.zeros(n_embd)?;
8885 let mut preds_d = e.alloc_u32_zeroed(chunk)?;
8886 let nll_on = std::env::var("MEMRA_REPLAY_NLL").as_deref() == Ok("1");
8887 let (mut nll_sum, mut nll_cnt) = (0f64, 0u64);
8888 let mut s = 0usize;
8889 while s < t_total {
8890 let cend = (s + chunk).min(t_total);
8891 let tc = cend - s;
8892 let ch = &tokens[s..cend];
8893 // 1. forced trunk pass — verify path (decode-exact contract): all-column logits +
8894 // the chunk's true hiddens.
8895 let (tl_d, vx) = self.decode_step_t_core(e, ch, s, &mut cache, embd_dev, None)?;
8896 for j in 0..tc {
8897 e.argmax_token_device_col(&tl_d, j, n_vocab, &mut preds_d, j)?;
8898 }
8899 let preds = e.dtoh_u32(&preds_d)?;
8900 for j in 0..tc {
8901 bg[s + j + 1] = preds[j];
8902 }
8903 // MEMRA_REPLAY_NLL=1: teacher-forced NLL/perplexity over the same forced pass — the
8904 // checkpoint-quality metric (position j's logits score the GOLD next token).
8905 if nll_on {
8906 let jmax = if cend < t_total { tc } else { tc - 1 }; // last pos has no gold next
8907 if jmax > 0 {
8908 let ids: Vec<u32> = (0..jmax).map(|j| tokens[s + j + 1]).collect();
8909 let rows: Vec<i32> = (0..jmax as i32).collect();
8910 let idsd = e.htod_u32_v(&ids)?;
8911 let rowsd = e.htod_i32(&rows)?;
8912 let mut outd = e.zeros(jmax)?;
8913 e.softmax_gather(&tl_d, n_vocab, &idsd, &rowsd, &mut outd, n_vocab, jmax, 1.0)?;
8914 for pr in e.dtoh(&outd)? {
8915 nll_sum += -((pr.max(1e-30)) as f64).ln();
8916 nll_cnt += 1;
8917 }
8918 }
8919 }
8920 if let Some(f) = hdump.as_deref_mut() {
8921 use std::io::Write;
8922 let host: Vec<f32> = e.dtoh(&vx)?;
8923 // bf16 round-to-nearest-even — f32 doubled the disk bill at bulk
8924 // extraction scale (20M tokens x 4096 = 320GB f32 vs 160GB bf16).
8925 let mut bytes = Vec::with_capacity(tc * n_embd * 2);
8926 for v in &host[..tc * n_embd] {
8927 let b = v.to_bits();
8928 let r = b.wrapping_add(0x7FFF + ((b >> 16) & 1));
8929 bytes.extend_from_slice(&((r >> 16) as u16).to_le_bytes());
8930 }
8931 f.write_all(&bytes)?;
8932 }
8933 // CHAINLESS extraction (stride > corpus, the bulk-hdump mode): no chunk ever
8934 // drafts, so the draft-KV fills are pure waste — skip them (2 MTP-block passes
8935 // per token saved; the forced trunk pass + hdump is all the mode needs).
8936 let chainless = stride > t_total;
8937 if chainless {
8938 e.copy_view_into(
8939 &mut prev_last_h,
8940 0,
8941 &vx.slice((tc - 1) * n_embd..tc * n_embd),
8942 n_embd,
8943 )?;
8944 s = cend;
8945 continue;
8946 }
8947 // 2. TRUE predecessor-paired draft-KV fill for the chunk (row i carries h_{i-1};
8948 // row s reads the previous chunk's last true hidden, zeros at corpus start).
8949 let mut vxs = e.zeros(tc * n_embd)?;
8950 e.copy_into(&mut vxs, 0, &prev_last_h, n_embd)?;
8951 if tc > 1 {
8952 e.copy_view_into(
8953 &mut vxs,
8954 n_embd,
8955 &vx.slice(0..(tc - 1) * n_embd),
8956 (tc - 1) * n_embd,
8957 )?;
8958 }
8959 scratch.set_len(e, s)?;
8960 self.mtp_kv_fill(e, mtp, ch, &vxs, s, &mut scratch, embd_dev)?;
8961 // 3. draft chains at sampled positions, DESCENDING: a chain reads only slots
8962 // [0..p) (true fills) and appends at >= p; the next (smaller-p) chain's set_len
8963 // truncates those approximate appends before they can ever be read.
8964 let ps: Vec<usize> = (s..cend)
8965 .filter(|p| *p >= 1 && *p % stride == 0 && *p + k <= t_total)
8966 .collect();
8967 for &p in ps.iter().rev() {
8968 scratch.set_len(e, p)?;
8969 if p == s {
8970 e.copy_into(&mut seed_buf, 0, &prev_last_h, n_embd)?;
8971 } else {
8972 e.copy_view_into(
8973 &mut seed_buf,
8974 0,
8975 &vx.slice((p - 1 - s) * n_embd..(p - s) * n_embd),
8976 n_embd,
8977 )?;
8978 }
8979 let mut e_tok = tokens[p];
8980 let mut d_seed = e.clone_dtod(&seed_buf)?;
8981 let mut drafts: Vec<u32> = Vec::with_capacity(k);
8982 for j in 0..k {
8983 let (dl_d, h_nextn) = self.mtp_head_forward_dev(
8984 e,
8985 mtp,
8986 e_tok,
8987 &d_seed,
8988 &mut scratch,
8989 p + 1 + j,
8990 embd_dev,
8991 None, // acceptance-oracle walk: no grammar
8992 )?;
8993 let tok_d = e.argmax_token_device(&dl_d, d_vocab)?;
8994 let idx = e.dtoh_u32_one(&tok_d)?;
8995 let d = match &mtp.d2t {
8996 Some(map) => map[idx as usize],
8997 None => idx,
8998 };
8999 drafts.push(d);
9000 e_tok = d;
9001 d_seed = h_nextn;
9002 }
9003 // targets may live in a LATER chunk's bg — resolved after the walk.
9004 rows.push((p, drafts, Vec::new()));
9005 }
9006 // 4. restore TRUE entries for the whole chunk (the next chunk's chains and fills
9007 // expect scratch.len == cend with exact rows).
9008 scratch.set_len(e, s)?;
9009 self.mtp_kv_fill(e, mtp, ch, &vxs, s, &mut scratch, embd_dev)?;
9010 e.copy_view_into(
9011 &mut prev_last_h,
9012 0,
9013 &vx.slice((tc - 1) * n_embd..tc * n_embd),
9014 n_embd,
9015 )?;
9016 s = cend;
9017 }
9018 for (p, drafts, targets) in rows.iter_mut() {
9019 for j in 0..drafts.len() {
9020 targets.push(bg[*p + 1 + j]);
9021 }
9022 }
9023 rows.sort_by_key(|r| r.0);
9024 if nll_cnt > 0 {
9025 let mean = nll_sum / nll_cnt as f64;
9026 println!(
9027 "[replay-nll] tokens={nll_cnt} nll/token={mean:.5} ppl={:.4}",
9028 mean.exp()
9029 );
9030 }
9031 Ok((rows, bg))
9032 }
9033}
9034
9035#[cfg(test)]
9036mod dspark_sparse_tests {
9037 use super::dspark_sparse_softmax_topk;
9038
9039 #[test]
9040 fn topk_keeps_full_softmax_mass_and_stable_ties() {
9041 let logits = [1.0f32, 3.0, 3.0, -2.0];
9042 let (ids, top_logits, probs, tail) =
9043 dspark_sparse_softmax_topk(&logits, 2, 1.0).unwrap();
9044 assert_eq!(ids, vec![1, 2]);
9045 assert_eq!(top_logits, vec![3.0, 3.0]);
9046 let denominator = logits.iter().map(|value| (value - 3.0).exp()).sum::<f32>();
9047 let expected = 1.0 / denominator;
9048 assert!((probs[0] - expected).abs() < 1.0e-6);
9049 assert!((probs[1] - expected).abs() < 1.0e-6);
9050 assert!((tail - (1.0 - 2.0 * expected)).abs() < 1.0e-6);
9051 assert!((probs.iter().sum::<f32>() + tail - 1.0).abs() < 1.0e-6);
9052 }
9053}
9054
9055#[cfg(test)]
9056mod telem_tests {
9057 use super::{SpecTelemetry, SpecTelemetryCounters, SPEC_TELEM_POS};
9058
9059 #[test]
9060 fn synthetic_accept_masks_produce_tau_and_position_histogram() {
9061 let counters = SpecTelemetryCounters::default();
9062 for mask in [
9063 [true, true, true],
9064 [true, true, false],
9065 [true, false, false],
9066 [false, false, false],
9067 ] {
9068 let accepted = mask.iter().take_while(|&&value| value).count();
9069 counters.record_round(mask.len(), accepted);
9070 }
9071
9072 let snapshot = counters.snapshot();
9073 assert_eq!((snapshot.rounds, snapshot.drafted, snapshot.accepted), (4, 12, 6));
9074 assert_eq!(&snapshot.pos_drafted[..3], &[4, 4, 4]);
9075 assert_eq!(&snapshot.pos_accepted[..3], &[3, 2, 1]);
9076 assert_eq!(snapshot.tau(), 1.5);
9077 assert_eq!(snapshot.pos_drafted[3..], [0; SPEC_TELEM_POS - 3]);
9078 assert_eq!(snapshot.pos_accepted[3..], [0; SPEC_TELEM_POS - 3]);
9079 }
9080
9081 /// The worker's per-burst pattern: stash, accumulate, diff — the delta must isolate
9082 /// exactly the burst's contribution (pool-resumed sessions carry prior requests' counts).
9083 #[test]
9084 fn delta_isolates_burst_contribution() {
9085 let mut t = SpecTelemetry::default();
9086 // "previous request": 2 rounds of k=3, accepts 3 then 1.
9087 for (kr, na) in [(3usize, 3usize), (3, 1)] {
9088 t.rounds += 1;
9089 t.drafted += kr as u64;
9090 t.accepted += na as u64;
9091 for j in 0..kr { t.pos_drafted[j] += 1; }
9092 for j in 0..na { t.pos_accepted[j] += 1; }
9093 }
9094 let before = t;
9095 // "this burst": 1 round k=3, accepts 2.
9096 t.rounds += 1;
9097 t.drafted += 3;
9098 t.accepted += 2;
9099 for j in 0..3 { t.pos_drafted[j] += 1; }
9100 for j in 0..2 { t.pos_accepted[j] += 1; }
9101 let d = t.delta_since(&before);
9102 assert_eq!((d.rounds, d.drafted, d.accepted), (1, 3, 2));
9103 assert_eq!(&d.pos_drafted[..3], &[1, 1, 1]);
9104 assert_eq!(&d.pos_accepted[..3], &[1, 1, 0]);
9105 assert_eq!(d.pos_drafted[3..], [0; SPEC_TELEM_POS - 3]);
9106 }
9107
9108 /// merge(delta) then merge(delta2) equals accumulating both — the per-model /metrics
9109 /// aggregation invariant.
9110 #[test]
9111 fn merge_accumulates_fieldwise() {
9112 let mut agg = SpecTelemetry::default();
9113 let mut d1 = SpecTelemetry { rounds: 2, drafted: 6, accepted: 4, ..Default::default() };
9114 d1.pos_drafted[0] = 2;
9115 d1.pos_accepted[0] = 2;
9116 let mut d2 = SpecTelemetry { rounds: 1, drafted: 3, accepted: 1, ..Default::default() };
9117 d2.pos_drafted[0] = 1;
9118 d2.pos_accepted[0] = 1;
9119 d2.pos_drafted[1] = 1;
9120 agg.merge(&d1);
9121 agg.merge(&d2);
9122 assert_eq!((agg.rounds, agg.drafted, agg.accepted), (3, 9, 5));
9123 assert_eq!(agg.pos_drafted[0], 3);
9124 assert_eq!(agg.pos_accepted[0], 3);
9125 assert_eq!(agg.pos_drafted[1], 1);
9126 assert_eq!(agg.pos_accepted[1], 0);
9127 }
9128
9129 /// Wrong-snapshot diff saturates to zero instead of wrapping — the counters feed a
9130 /// public metrics surface and must never publish a u64-wrapped garbage value.
9131 #[test]
9132 fn delta_saturates_never_wraps() {
9133 let small = SpecTelemetry { rounds: 1, drafted: 2, accepted: 1, ..Default::default() };
9134 let big = SpecTelemetry { rounds: 5, drafted: 15, accepted: 9, ..Default::default() };
9135 let d = small.delta_since(&big);
9136 assert_eq!((d.rounds, d.drafted, d.accepted), (0, 0, 0));
9137 }
9138}
9139
9140#[cfg(test)]
9141mod opti_fork_tests {
9142 use super::{
9143 OptiControllerPolicy, OptiForkAction, OptiForkGateMode, OptiForkGenerationTracker,
9144 };
9145
9146 #[test]
9147 fn controller_threshold_and_three_miss_breaker_are_exact() {
9148 let mut policy = OptiControllerPolicy {
9149 threshold: 0.7,
9150 consecutive_misses: 0,
9151 breaker_tripped: false,
9152 };
9153 assert!(!policy.admit(0.699_999));
9154 assert!(policy.admit(0.7));
9155 assert!(!policy.resolve(false));
9156 assert!(!policy.resolve(false));
9157 assert!(policy.resolve(false));
9158 assert!(policy.breaker_tripped);
9159 assert!(!policy.admit(1.0));
9160 assert!(!policy.resolve(true), "a resolved hit cannot re-arm a tripped request");
9161 assert!(policy.breaker_tripped);
9162 }
9163
9164 #[test]
9165 fn zero_threshold_is_the_true_unconditional_measurement_arm() {
9166 let mut policy = OptiControllerPolicy {
9167 threshold: 0.0,
9168 consecutive_misses: 0,
9169 breaker_tripped: false,
9170 };
9171 for _ in 0..16 {
9172 assert!(policy.admit(0.0));
9173 assert!(!policy.resolve(false));
9174 }
9175 for invalid in [f32::NAN, f32::INFINITY, -0.01, 1.01] {
9176 assert!(!policy.admit(invalid), "invalid q proxy must fail closed: {invalid}");
9177 }
9178 assert!(!policy.breaker_tripped);
9179 assert_eq!(policy.consecutive_misses, 0);
9180 }
9181
9182 #[test]
9183 fn alternating_mode_flips_by_generation_not_round_parity() {
9184 assert_eq!(OptiForkGateMode::Alternate.action(0), OptiForkAction::Hit);
9185 assert_eq!(OptiForkGateMode::Alternate.action(1), OptiForkAction::Miss);
9186 assert_eq!(OptiForkGateMode::Alternate.action(8), OptiForkAction::Hit);
9187 assert_eq!(OptiForkGateMode::Alternate.action(9), OptiForkAction::Miss);
9188 }
9189
9190 #[test]
9191 fn live_generation_cannot_be_overwritten() {
9192 let mut tracker = OptiForkGenerationTracker::default();
9193 let g0 = tracker.reserve().unwrap();
9194 let g1 = tracker.reserve().unwrap();
9195 let err = tracker.reserve().unwrap_err().to_string();
9196 assert!(err.contains("still owns generation 0"), "unexpected error: {err}");
9197 tracker.retire(g0).unwrap();
9198 let g2 = tracker.reserve().unwrap();
9199 assert_eq!((g2.id, g2.slot), (2, 0));
9200 tracker.retire(g1).unwrap();
9201 tracker.retire(g2).unwrap();
9202 }
9203
9204 #[test]
9205 fn teardown_rejects_a_stale_generation_tag() {
9206 let mut tracker = OptiForkGenerationTracker::default();
9207 let g0 = tracker.reserve().unwrap();
9208 tracker.retire(g0).unwrap();
9209 let err = tracker.retire(g0).unwrap_err().to_string();
9210 assert!(err.contains("teardown mismatch"), "unexpected error: {err}");
9211 }
9212}
9213
9214#[cfg(test)]
9215mod draft_graph_fallback_tests {
9216 use super::DraftGraphFallback;
9217
9218 /// The Q2 contract, part (a): a fallback flip is LOUD — exactly once per flip.
9219 #[test]
9220 fn flip_is_loud_once_and_memoized_after() {
9221 let mut f = DraftGraphFallback::default();
9222 let line = f.mark_greedy("out of memory").expect("first flip must return the warn line");
9223 assert!(line.contains("WARN"), "flip line must be warn-level: {line}");
9224 assert!(line.contains("out of memory"), "flip line must carry the reason: {line}");
9225 assert!(f.greedy_failed());
9226 // re-marking an already-failed graph is the memoization: quiet, still failed.
9227 assert!(f.mark_greedy("out of memory").is_none());
9228 assert!(f.greedy_failed());
9229 // the two graphs' flags are independent (greedy flip leaves sampled capturable).
9230 assert!(!f.sampled_failed());
9231 let line_s = f.mark_sampled("capture unsupported").expect("sampled flip is its own flip");
9232 assert!(line_s.contains("sampled"), "sampled flip names itself: {line_s}");
9233 assert!(f.mark_sampled("capture unsupported").is_none());
9234 }
9235
9236 /// The Q2 contract, part (b): resume-from-pool RESETS both flags (fresh capture chance),
9237 /// and says so exactly when there was something to reset.
9238 #[test]
9239 fn reset_on_resume_clears_flags_and_logs_once() {
9240 let mut f = DraftGraphFallback::default();
9241 // clean session: resume is silent, nothing to reset.
9242 assert!(f.reset_on_resume().is_none());
9243 f.mark_greedy("oom").unwrap();
9244 f.mark_sampled("oom").unwrap();
9245 let note = f.reset_on_resume().expect("a set flag must produce the reset note");
9246 assert!(note.contains("greedy+sampled"), "note names what was reset: {note}");
9247 assert!(!f.greedy_failed() && !f.sampled_failed(), "both flags cleared");
9248 // and the NEXT failure after a reset is a fresh flip — loud again.
9249 assert!(f.mark_greedy("oom again").is_some());
9250 let note2 = f.reset_on_resume().expect("greedy-only reset");
9251 assert!(note2.contains("(greedy)"), "single-flag note: {note2}");
9252 }
9253
9254 /// Shape-change clears (dmask realloc / mask-shape mismatch / s_key change) stay silent —
9255 /// they precede a fresh capture attempt whose own failure re-flips loudly.
9256 #[test]
9257 fn shape_change_clears_are_silent() {
9258 let mut f = DraftGraphFallback::default();
9259 f.mark_greedy("oom").unwrap();
9260 f.clear_greedy();
9261 assert!(!f.greedy_failed());
9262 f.mark_sampled("oom").unwrap();
9263 f.clear_sampled();
9264 assert!(!f.sampled_failed());
9265 // after a silent clear there is nothing left for resume to report.
9266 assert!(f.reset_on_resume().is_none());
9267 }
9268}