memra_engine/spec.rs
1//! Qwen3.5 MTP (NextN) greedy speculative decode (research/mtp/MTP-PLAN.md §A/§B/§C/§D).
2//!
3//! Greedy spec decode is MATHEMATICALLY EXACT: the accepted+bonus token stream is token-for-token
4//! identical to plain greedy `generate`. This module provides:
5//! - `mtp_head_forward` (§A, T=1): one NextN draft-token forward.
6//! - `decode_step_t` (§D.3, T=K+1): batched target verify forward, all-column logits.
7//! - `generate_spec` (§B): the draft/verify/accept/rollback orchestrator.
8//! Cache snapshot/rollback lives in cache.rs (§D.4). The MTP head uses its OWN scratch KV (§D.6),
9//! PERSISTENT over the committed sequence (see `MtpScratch`).
10
11use crate::Engine;
12use crate::cache::{Cache, KvLayer};
13use crate::forward::argmax;
14use crate::hybrid::{FullAttnLayer, HybridModel, LinearAttnLayer, Mixer, MtpHead};
15use cudarc::driver::CudaSlice;
16use std::sync::atomic::{AtomicU64, Ordering};
17
18/// Parse the documented `MEMRA_SPEC_REPLAY=1` rollback seam.
19///
20/// Keep this shared with serving admission so `=0` cannot select replay in one
21/// layer while another layer treats it as disabled.
22pub fn spec_replay_env_on(value: Option<&str>) -> bool {
23 value == Some("1")
24}
25
26pub fn spec_replay_env_enabled() -> bool {
27 let value = std::env::var("MEMRA_SPEC_REPLAY").ok();
28 spec_replay_env_on(value.as_deref())
29}
30
31/// One compact, anchor-bounded DSpark supervision record. `tokens[0]` is the anchor at p and
32/// `hidden` is its predecessor carrier h[p-1], matching the live NextN/DSpark pairing. Target
33/// rows p..p+gamma-1 score tokens p+1..p+gamma. They are the full-target softmax's top-k
34/// entries; `target_tail_probs[j]` is the probability mass outside those rows. All flattened
35/// target arrays are `[gamma, top_k]` in row-major order.
36pub struct DsparkAnchorRecord {
37 pub position: usize,
38 pub hidden: Vec<f32>,
39 pub tokens: Vec<u32>,
40 pub target_top_ids: Vec<u32>,
41 pub target_top_logits: Vec<f32>,
42 pub target_top_probs: Vec<f32>,
43 pub target_tail_probs: Vec<f32>,
44}
45
46fn dspark_sparse_softmax_topk(
47 logits: &[f32],
48 top_k: usize,
49 temperature: f32,
50) -> Result<(Vec<u32>, Vec<f32>, Vec<f32>, f32), Box<dyn std::error::Error>> {
51 if logits.is_empty() || top_k == 0 || top_k > logits.len() || temperature <= 0.0 {
52 return Err("invalid DSpark sparse-softmax shape or temperature".into());
53 }
54 if logits.iter().any(|value| !value.is_finite()) {
55 return Err("DSpark target logits contain a non-finite value".into());
56 }
57 let mut ranked: Vec<(u32, f32)> = logits
58 .iter()
59 .copied()
60 .enumerate()
61 .map(|(index, value)| (index as u32, value))
62 .collect();
63 let compare = |left: &(u32, f32), right: &(u32, f32)| {
64 right.1.total_cmp(&left.1).then(left.0.cmp(&right.0))
65 };
66 ranked.select_nth_unstable_by(top_k - 1, compare);
67 ranked[..top_k].sort_unstable_by(compare);
68
69 let max_logit = logits.iter().copied().fold(f32::NEG_INFINITY, f32::max);
70 let inv_temperature = 1.0f64 / temperature as f64;
71 let denominator: f64 = logits
72 .iter()
73 .map(|value| (((*value - max_logit) as f64) * inv_temperature).exp())
74 .sum();
75 let ids: Vec<u32> = ranked[..top_k].iter().map(|(index, _)| *index).collect();
76 let top_logits: Vec<f32> = ranked[..top_k].iter().map(|(_, value)| *value).collect();
77 let top_probs: Vec<f32> = top_logits
78 .iter()
79 .map(|value| ((((value - max_logit) as f64) * inv_temperature).exp() / denominator) as f32)
80 .collect();
81 let top_mass: f64 = top_probs.iter().map(|value| *value as f64).sum();
82 let tail = (1.0f64 - top_mass).clamp(0.0, 1.0) as f32;
83 Ok((ids, top_logits, top_probs, tail))
84}
85
86fn flatten_dspark_rows<T>(
87 rows: Vec<Option<Vec<T>>>,
88 position: usize,
89 label: &str,
90) -> Result<Vec<T>, Box<dyn std::error::Error>> {
91 let mut flattened = Vec::new();
92 for (slot, row) in rows.into_iter().enumerate() {
93 flattened.extend(
94 row.ok_or_else(|| format!("missing DSpark {label} at {position} slot {slot}"))?,
95 );
96 }
97 Ok(flattened)
98}
99
100/// H-SEED CONVENTION (MEMRA_SPEC_HPOST=1): feed the MTP head the POST-norm hidden — trunk rows
101/// hand over `output_norm(x)` and the draft chain recurrence hands over `shared_head_norm(h_nextn)`
102/// (= final_h) — matching the reference engines: llama.cpp #24025 ("qwen35: use post-norm hidden
103/// state for MTP", t_h_nextn is taken AFTER the final norm in both trunk and MTP graphs) and
104/// SGLang's qwen3_5_mtp (spec_info.hidden_states = the target model's post-norm output). memra's
105/// historical convention (default, MTP-PLAN §A) is PRE-norm x. Draft-quality-only: exactness is
106/// the verify's job either way; acceptance arbitrates. OnceLock: read once, hot-loop safe.
107pub(crate) fn spec_hpost() -> bool {
108 static H: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
109 *H.get_or_init(|| {
110 std::env::var("MEMRA_SPEC_HPOST")
111 .map(|v| v != "0")
112 .unwrap_or(false)
113 })
114}
115
116/// LEAN VERIFY (default ON since 2026-07-08; MEMRA_SPEC_LEAN=0 reverts — close35 lane): the verify m-scaling
117/// probe + nsys diff showed the verify t-path pays ~1.0ms/call at m=1 over eager decode on the
118/// 35B, and the kernels are NOT the cause (dev-MoE identical, kernel-time delta only +179us).
119/// The overhead is (a) ~250 extra cuMemsetD8Async/call from `e.zeros()` on buffers every kernel
120/// fully overwrites (~0.9ms host issue + ~0.35ms GPU) and (b) the t=1 FA rows dispatch (rows_v2 +
121/// combine_rows, +50us vs the eager fa_decode pair). This flag switches (a) fully-overwritten
122/// verify buffers to `e.uninit` (identical bytes: every element is written before read) and
123/// (b) t==1 verify FA to the eager `fa_decode` entry (byte-identical: kernel-check pins the
124/// rows-vs-loop identity and the per-row loop at t=1 IS fa_decode on the same q). Gates arbitrate.
125pub(crate) fn spec_lean() -> bool {
126 static L: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
127 // DEFAULT ON since 2026-07-08 (MEMRA_SPEC_LEAN=0 reverts): bit-identical (buffers fully
128 // overwritten; gates green incl maxdiff-identical run-gen) and measured +2.4% e2e p3 /
129 // +1.5% p2 at the daily 35B config. m=1 verify now costs eager-decode parity.
130 *L.get_or_init(|| {
131 std::env::var("MEMRA_SPEC_LEAN")
132 .map(|v| v != "0")
133 .unwrap_or(true)
134 })
135}
136
137/// SMALL-M BATCHED VERIFY (default ON since 2026-07-09; MEMRA_SPEC_M2=0 reverts — lane/spec-m2): extend the
138/// batched linear-attn verify arm down to t=2 and batch the MoE dev token loop over a
139/// grid.z=token axis at every verify t. The close35 m-scaling probe put the m=2 verify tier at
140/// x1.54 of m=1 (llama x1.14); the per-column linear chain (t<3) and the serial MoE dev token
141/// loop are the two launch-structure causes. Both changes are LAUNCH-STRUCTURE ONLY:
142/// (a) the batched conv's t<pad ring update is pure copies (ssm_conv_ring_rebuild from a cloned
143/// ring — the ring stores raw input columns); every arithmetic kernel is the same one the
144/// t>=3 arm already runs (matmul_decode_exact bit-identical at m=2-4, gdn_scan's internal
145/// t-loop == chained T=1 steps);
146/// (b) the MoE dev-rows twins run the serial loop's per-token warp program with tok-offset
147/// pointers (same sel/w/aq/ad bytes, same dot order, same slot-ordered FMA chain).
148/// Gates arbitrate: run-spec K=1..8 self-consistency (35B+9B), kernel-check, run-gen argmax.
149pub(crate) fn spec_m2() -> bool {
150 static M: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
151 // DEFAULT ON since 2026-07-09 (MEMRA_SPEC_M2=0 reverts): launch-structure only — t=2
152 // batched linear arm (ring-roll copies, zero new FP order) + MoE dev-rows kernels
153 // (grid.z=token, 4 launches/layer at any verify t). Acceptance bit-identical at every K;
154 // 35B p2 +3.4% / p3 +3.6%; the profitable-K plateau widens (new optimum K=3 at 223).
155 *M.get_or_init(|| {
156 std::env::var("MEMRA_SPEC_M2")
157 .map(|v| v != "0")
158 .unwrap_or(true)
159 })
160}
161pub(crate) fn spec_stream() -> bool {
162 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
163 *ON.get_or_init(|| std::env::var("MEMRA_SPEC_STREAM").as_deref() == Ok("1"))
164}
165pub(crate) fn spec_stream_m() -> usize {
166 static M: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
167 *M.get_or_init(|| {
168 std::env::var("MEMRA_SPEC_STREAM_M")
169 .ok()
170 .and_then(|v| v.parse().ok())
171 .unwrap_or(4)
172 })
173}
174pub(crate) fn spec_devacc() -> bool {
175 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
176 *ON.get_or_init(|| std::env::var("MEMRA_SPEC_DEVACC").as_deref() == Ok("1"))
177}
178
179/// GRAMMAR HOOK for constrained spec decode (lane/constrained-full, 2026-08-03). The engine
180/// stays llguidance-agnostic: the server adapts its per-session grammar state behind this
181/// trait. CONTRACT (the verify-side truncation rule — token-identical to constrained plain
182/// greedy decode): the exactness walk runs UNMASKED first; the hook then (a) truncates
183/// acceptance at the first grammar-illegal accepted token, and (b) when the truncation fired
184/// or the bonus is illegal, the engine recomputes that slot as the MASKED argmax of the
185/// target's own verify column (an unmasked argmax that is grammar-legal IS the masked argmax
186/// — masking only removes tokens — so the common case pays nothing). `consume` advances the
187/// state with each EMITTED token in order; EOS handling is the implementor's job (skip).
188pub trait SpecConstraint {
189 /// -inf the current state's banned ids on a HOST logits row (prompt-tail / init-feed
190 /// masked argmax).
191 fn mask_logits(&mut self, logits: &mut [f32]) -> Result<(), String>;
192 /// Packed 32-bit bitset words of the CURRENT state's allowed set (device-mask form).
193 fn mask_words(&mut self) -> Result<Vec<u32>, String>;
194 /// Is `tok` consumable in the CURRENT state?
195 fn is_allowed(&mut self, tok: u32) -> Result<bool, String>;
196 /// Advance the state with an emitted token.
197 fn consume(&mut self, tok: u32) -> Result<(), String>;
198
199 // --- DRAFT-SIDE MASKING (lane/draft-mask, 2026-08-04) ---
200 // The drafter proposed grammar-illegal tokens under tight schemas, so verify-side
201 // truncation cut nearly every round (measured acceptance 0.467-0.513 tight vs 0.62-0.82
202 // loose, research/constrained-full-20260803). These three methods let the engine mask the
203 // DRAFT model's own sampling with the grammar's legal set, so proposals are legal by
204 // construction. The state they walk is a SPECULATIVE CLONE of the session matcher — the
205 // real state is advanced only by `consume` (emitted tokens), so verify-side truncation
206 // stays the correctness backstop and the emitted stream is unchanged by construction
207 // (an accepted draft is the target's unmasked argmax AND grammar-legal, hence the masked
208 // argmax; a cut slot is recomputed as the masked argmax either way).
209 // Default impls = feature OFF (pre-lane behaviour: unmasked drafts).
210
211 /// Is draft-side masking available on this hook? Probed ONCE per burst, before the draft
212 /// graph is captured (the mask is an in-graph node — its presence is a capture-time shape).
213 fn draft_mask_enabled(&self) -> bool {
214 false
215 }
216 /// Start a draft chain: clone the CURRENT (committed) grammar state into the speculative
217 /// slot. Called once per spec round, before the first draft position.
218 fn draft_begin(&mut self) -> Result<(), String> {
219 Ok(())
220 }
221 /// Packed 32-bit bitset words of the SPECULATIVE state's allowed set (target-vocab ids),
222 /// for the draft position about to be sampled. `None` = draft masking off (no-op).
223 fn draft_mask_words(&mut self) -> Result<Option<Vec<u32>>, String> {
224 Ok(None)
225 }
226 /// Advance the SPECULATIVE state with a PROPOSED draft token. `false` = the chain cannot
227 /// continue (EOS proposed, or an unmasked position proposed something illegal) — the
228 /// engine stops drafting; the token already pushed still goes through verify.
229 fn draft_advance(&mut self, _tok: u32) -> Result<bool, String> {
230 Ok(false)
231 }
232}
233
234/// DRAFT-MASK UPLOAD (lane/draft-mask): pull the speculative state's allowed set (TARGET-id
235/// space) from the hook, project it into the DRAFT head's vocab space, and upload it into the
236/// stable device buffer the draft chain reads. Returns false when the chain must stop drafting:
237/// the hook handed out no mask, or NO draft-vocab row is grammar-legal at this position (a
238/// trimmed FR-Spec head genuinely cannot propose a legal token there — masking it would leave
239/// a fully-banned row whose argmax is meaningless, so the round drafts fewer tokens and the
240/// verify emits the masked argmax as usual).
241fn upload_draft_mask(
242 e: &Engine,
243 c: &mut dyn SpecConstraint,
244 dst: &mut CudaSlice<u32>,
245 d2t: Option<&Vec<u32>>,
246 d_vocab: usize,
247 words: usize,
248) -> Result<bool, Box<dyn std::error::Error>> {
249 let Some(tw) = c
250 .draft_mask_words()
251 .map_err(|e2| format!("constraint: {e2}"))?
252 else {
253 return Ok(false);
254 };
255 let bit = |t: usize| -> bool {
256 let w = t >> 5;
257 w < tw.len() && (tw[w] >> (t & 31)) & 1 == 1
258 };
259 let mut buf = vec![0u32; words];
260 match d2t {
261 // TRIMMED draft head: row i proposes target id d2t[i] — permute the mask accordingly.
262 Some(map) => {
263 for (i, &t) in map.iter().enumerate().take(d_vocab) {
264 if bit(t as usize) {
265 buf[i >> 5] |= 1u32 << (i & 31);
266 }
267 }
268 }
269 // UNTRIMMED: draft ids ARE target ids; the packed words transfer verbatim (a short
270 // mask leaves the padded tail zeroed == banned, same rule as constrained::apply_mask).
271 None => {
272 let n = tw.len().min(words);
273 buf[..n].copy_from_slice(&tw[..n]);
274 }
275 }
276 if buf.iter().all(|w| *w == 0) {
277 return Ok(false);
278 }
279 e.htod_u32_into(dst, &buf)?;
280 Ok(true)
281}
282
283/// Keep the full token-embedding table in host memory and upload only the rows needed by each
284/// MTP/verify step. This is an exact memory-capacity seam for very large BF16 vocab tables: host
285/// gather expands the same source bits to f32, and only O(T*n_embd) bytes cross PCIe per step.
286/// CUDA-graph/round-stream draft paths require device token ids and therefore stay disabled.
287pub(crate) fn spec_host_embd() -> bool {
288 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
289 *ON.get_or_init(|| std::env::var("MEMRA_SPEC_HOST_EMBD").as_deref() == Ok("1"))
290}
291
292/// VERIFY-TIER TRUNK LAUNCH-FUSION (default ON since 2026-07-09; MEMRA_SPEC_FUSED_T=0 reverts — lane/close35b): extend
293/// the t=1 fused2/fused3 Q8_0 trunk launches to the batched verify tier (t=2-4, the K=1..3
294/// verify shapes). At t>1 the trunk pairs/triples (35B wqkv+wqkv_gate, wq/wk/wv,
295/// gate_shexp+up_shexp) each run a separate `matmul_decode_exact` — one q8_1 re-quantize of the
296/// SAME activation plus one _b2/_b4 launch per tensor. The fused twins share ONE quantize and
297/// ONE launch per group; per (tensor,token,row) the kernel body is q8_0_mmvq_batched verbatim
298/// with the identical row mapping -> BIT-IDENTICAL by construction (kernel-check pins it,
299/// run-spec K=1..8 + acceptance identity arbitrate e2e).
300pub(crate) fn spec_fused_t() -> bool {
301 static F: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
302 // DEFAULT ON since 2026-07-09 (MEMRA_SPEC_FUSED_T=0 reverts): verify t=2-4 trunk launch-fusion
303 // (fused2/fused3 Q8_0 batched twins, bit-identical by construction — m=1 block-offset split on
304 // the batched body). m=2 marginal token 2117->1762us; 35B daily: p3 +3.7% (crosses llama), p2 +5%.
305 *F.get_or_init(|| {
306 std::env::var("MEMRA_SPEC_FUSED_T")
307 .map(|v| v != "0")
308 .unwrap_or(true)
309 })
310}
311
312/// zeros/uninit switch for verify-path buffers that are FULLY OVERWRITTEN before any read.
313/// Only call this on such buffers — the lean contract is "identical bytes by construction".
314fn vbuf(e: &Engine, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
315 if spec_lean() { e.uninit(n) } else { e.zeros(n) }
316}
317
318/// Scratch KV for the MTP block (one full-attn layer).
319///
320/// PERSISTENT MODE (default, 2026-07-03 — the acceptance lever): sized cap = max_ctx and kept in
321/// sync with the COMMITTED sequence — slot p holds the MTP block's K/V for committed token p
322/// (roped p+1, the chain's rope convention), so the draft chain's self-attention sees the FULL
323/// committed history instead of only the current round's 1..K+1 chain tokens (the reference
324/// engine's "mtp_update" design). Entries come from two sources:
325/// - chain appends: accepted positions KEEP their chain-computed entries (embedding exact,
326/// hidden chain-approximate — the reference engine accepts the same);
327/// - `mtp_kv_fill` batches: prompt positions + the last-draft position on full accept, computed
328/// from EXACT trunk hiddens (K/V-only MTP-block pass, no attention/FFN/lm_head).
329/// Rejected drafts / p-min extras / pseudo-seed appends are all discarded by the round-start
330/// `set_len` truncation (the KvLayer len mechanism — §C rollback for the draft side).
331/// Multi-turn spec-decode session (2026-07-05): trunk Cache + persistent MTP draft scratch +
332/// the committed token list, alive across generate_spec_session calls. Turn N+1 primes ONLY its
333/// suffix (chunked continuation prime over the quantized past) and mtp_kv_fill's its suffix rows,
334/// then the round loop runs unchanged. `last_h` carries the pre-output_norm hidden of the last
335/// committed row across turns (the predecessor-pairing seed + fill anchor).
336/// Per-request sampling config for the sampled-spec serve path.
337#[derive(Clone, Copy, Debug)]
338pub struct SpecSampling {
339 pub temp: f32,
340 pub seed: u64,
341 pub top_k: i32, // 0 = off
342 pub top_p: f32, // 1.0 = off
343 pub min_p: f32, // 0.0 = off
344 pub penalty_last_n: usize, // 0 = penalties off
345 pub penalty_repeat: f32,
346 pub penalty_freq: f32,
347 pub penalty_present: f32,
348}
349
350/// Tracked draft positions for [`SpecTelemetry`] (serve K defaults to 3; the run-spec gate
351/// sweeps K=1..8, and MEMRA_SPEC_CAPMAX defaults to 7 — 8 covers every tuned config).
352pub const SPEC_TELEM_POS: usize = 8;
353
354/// Always-on per-draft-position acceptance telemetry (lane/accept-telemetry, 2026-08-05 —
355/// the llama.cpp #26389 / vLLM spec-decode counter schema, upstream-sweeps 2026-08-05).
356/// Lives on the [`SpecSession`] and accumulates across bursts; the serve worker diffs a
357/// stashed copy per burst for its per-model /metrics aggregation and per-request usage.
358/// Same normalization as the `[spec-stats]` line: p-min-discarded chain tokens are counted
359/// in NEITHER drafted nor accepted.
360#[derive(Clone, Copy, Default, Debug)]
361pub struct SpecTelemetry {
362 /// verify rounds completed (a round-stream burst counts each of its M rounds).
363 pub rounds: u64,
364 /// tokens drafted / accepted across all rounds.
365 pub drafted: u64,
366 pub accepted: u64,
367 /// how often draft position j (0-based within a round's chain) was offered / accepted.
368 /// Positions >= SPEC_TELEM_POS are untracked (totals still count them). The opt-in
369 /// round-stream arm (MEMRA_SPEC_STREAM=1) reads back only totals, so under it these
370 /// arrays cover the standard-path rounds only and their sums may undercount the totals.
371 pub pos_drafted: [u64; SPEC_TELEM_POS],
372 pub pos_accepted: [u64; SPEC_TELEM_POS],
373}
374
375impl SpecTelemetry {
376 /// Fieldwise `self - prev` — the worker's per-burst delta off a copy stashed before the
377 /// burst call. Saturating: a caller diffing against the wrong snapshot gets zeros, not
378 /// a wrapped counter.
379 pub fn delta_since(&self, prev: &SpecTelemetry) -> SpecTelemetry {
380 let mut d = SpecTelemetry {
381 rounds: self.rounds.saturating_sub(prev.rounds),
382 drafted: self.drafted.saturating_sub(prev.drafted),
383 accepted: self.accepted.saturating_sub(prev.accepted),
384 ..Default::default()
385 };
386 for j in 0..SPEC_TELEM_POS {
387 d.pos_drafted[j] = self.pos_drafted[j].saturating_sub(prev.pos_drafted[j]);
388 d.pos_accepted[j] = self.pos_accepted[j].saturating_sub(prev.pos_accepted[j]);
389 }
390 d
391 }
392 /// Fieldwise `self += d` — the worker's per-model aggregation.
393 pub fn merge(&mut self, d: &SpecTelemetry) {
394 self.rounds += d.rounds;
395 self.drafted += d.drafted;
396 self.accepted += d.accepted;
397 for j in 0..SPEC_TELEM_POS {
398 self.pos_drafted[j] += d.pos_drafted[j];
399 self.pos_accepted[j] += d.pos_accepted[j];
400 }
401 }
402
403 /// Mean accepted draft-prefix length per verify round (tau).
404 pub fn tau(&self) -> f64 {
405 if self.rounds > 0 {
406 self.accepted as f64 / self.rounds as f64
407 } else {
408 0.0
409 }
410 }
411}
412
413/// Session-lifetime atomic acceptance counters. The verifier records only after the greedy or
414/// rejection-sampling walk has resolved on the host, so these relaxed increments add no GPU
415/// launch, synchronization, allocation, or ordering dependency to the numeric path.
416struct SpecTelemetryCounters {
417 rounds: AtomicU64,
418 drafted: AtomicU64,
419 accepted: AtomicU64,
420 pos_drafted: [AtomicU64; SPEC_TELEM_POS],
421 pos_accepted: [AtomicU64; SPEC_TELEM_POS],
422}
423
424impl Default for SpecTelemetryCounters {
425 fn default() -> Self {
426 Self {
427 rounds: AtomicU64::new(0),
428 drafted: AtomicU64::new(0),
429 accepted: AtomicU64::new(0),
430 pos_drafted: std::array::from_fn(|_| AtomicU64::new(0)),
431 pos_accepted: std::array::from_fn(|_| AtomicU64::new(0)),
432 }
433 }
434}
435
436impl SpecTelemetryCounters {
437 fn record_round(&self, drafted: usize, accepted: usize) {
438 debug_assert!(accepted <= drafted);
439 self.rounds.fetch_add(1, Ordering::Relaxed);
440 self.drafted.fetch_add(drafted as u64, Ordering::Relaxed);
441 self.accepted.fetch_add(accepted as u64, Ordering::Relaxed);
442 for counter in self.pos_drafted.iter().take(drafted) {
443 counter.fetch_add(1, Ordering::Relaxed);
444 }
445 for counter in self.pos_accepted.iter().take(accepted) {
446 counter.fetch_add(1, Ordering::Relaxed);
447 }
448 }
449
450 /// Round-stream keeps each round's accept length on device; retain exact scalar totals while
451 /// leaving the per-position arrays untouched, matching the pre-existing telemetry contract.
452 fn record_totals(&self, rounds: usize, drafted: usize, accepted: usize) {
453 self.rounds.fetch_add(rounds as u64, Ordering::Relaxed);
454 self.drafted.fetch_add(drafted as u64, Ordering::Relaxed);
455 self.accepted.fetch_add(accepted as u64, Ordering::Relaxed);
456 }
457
458 fn snapshot(&self) -> SpecTelemetry {
459 SpecTelemetry {
460 rounds: self.rounds.load(Ordering::Relaxed),
461 drafted: self.drafted.load(Ordering::Relaxed),
462 accepted: self.accepted.load(Ordering::Relaxed),
463 pos_drafted: std::array::from_fn(|j| self.pos_drafted[j].load(Ordering::Relaxed)),
464 pos_accepted: std::array::from_fn(|j| self.pos_accepted[j].load(Ordering::Relaxed)),
465 }
466 }
467}
468
469pub struct SpecSession {
470 pub(crate) cache: Cache,
471 pub(crate) scratch: MtpScratch,
472 /// Every token whose state the caches hold, in order (prompt turns + generated), INCLUDING
473 /// overshoot: spec commits accepted drafts past max_new; those rows are in the caches, so the
474 /// session must count them. Callers render output from this, not from their own echo.
475 pub committed: Vec<u32>,
476 /// Pre-output_norm hidden of the LAST committed row (device). None before the first turn.
477 pub(crate) last_h: Option<CudaSlice<f32>>,
478 /// Greedy argmax predicting the token AFTER committed.last() (from the last turn's final
479 /// logits). Fuels empty-suffix continuation bursts (serve): the next turn emits this token
480 /// first, feeds it, and the round loop resumes without any prime. None before the first turn.
481 pub next_pred: Option<u32>,
482 /// SAMPLED-SPEC stream continuity across bursts: Philox event counters persist here so a
483 /// session's randomness never repeats between generate_spec_session calls. (0,0) at admit.
484 pub sctr: u32,
485 pub uctr: u32,
486 /// PERSISTENT DRAFT-GRAPH CONTEXT (2026-08-01, the serve-burst fixed-cost fix): the captured
487 /// draft graph(s) + every device I/O buffer they bake, carried ACROSS generate_spec_session
488 /// calls. Before this, every serve burst re-captured the draft graph (2 warmup forwards +
489 /// instantiate) — measured ~16ms/burst on H100 q27 (MEMRA_SPEC_BURST sweep,
490 /// research/spec-serving-20260801). None before the first turn; error paths drop it
491 /// (next burst recaptures — serve retires errored sessions anyway).
492 pub(crate) draft_ctx: Option<DraftGraphCtx>,
493 /// PENDING-CARRY across bursts (2026-08-01, the serve burst-boundary fix): the bonus token
494 /// emitted by the last round but NOT committed to the caches. The old tail committed it with
495 /// a solo T=1 trunk pass (+ draft fill), and the next burst's setup fed the stashed next_pred
496 /// with ANOTHER solo pass — 2x ~11.5ms/burst measured on H100 q27 ([spec-setup] trace).
497 /// Carrying it lets the next empty-suffix greedy burst consume it as round-0 verify col 0,
498 /// exactly like a mid-burst full-accept boundary (no solo passes). INVARIANT: when set,
499 /// `committed` (== cache rows) EXCLUDES this token although it was already emitted in the
500 /// last burst's output, and `last_h` holds the hidden of the last COMMITTED row (its
501 /// predecessor — the chain-seed/fill anchor). `next_pred` is None (unknown without the
502 /// commit pass). Non-empty-suffix or sampled turns must flush first (spec_flush_pending);
503 /// generate_spec_session_sampled does this at entry, and serve parks only flushed sessions.
504 pub pending_tok: Option<u32>,
505 /// SESSION-AFFINITY TURN CHECKPOINT (lane/session-affinity, 2026-08-05): the state at this
506 /// turn's PROMPT-END boundary, retained so a later turn can REWIND here. See
507 /// [`SpecCheckpoint`]. Refreshed by every non-empty prime; None until the first one, and on
508 /// a rig too tight to hold it (a failed capture is silent — resume just isn't available).
509 pub(crate) turn_ckpt: Option<SpecCheckpoint>,
510 /// Session-lifetime acceptance telemetry. Relaxed atomics update at the host-side round
511 /// accounting the loop already does — no syncs, no allocation. NOTE a
512 /// pool-resumed session carries the PREVIOUS requests' counts; per-request consumers
513 /// diff with [`SpecTelemetry::delta_since`] around each burst.
514 telem: SpecTelemetryCounters,
515 /// PREFIX-CACHE publication request (lane/spec-prefix-cache): worker sets this to the
516 /// miss-LCP boundary before a cold burst; the prime captures at exactly that split (it must
517 /// coincide with the burst's `prime_split` or no capture happens). One-shot: consumed by the
518 /// prime, result lands in `boundary_capture`.
519 pub capture_at: Option<usize>,
520 /// The capture the last cold prime produced (see [`SpecBoundaryCapture`]). Worker takes it
521 /// post-burst to assemble the prefix entry. A failed capture is silent, like `turn_ckpt` —
522 /// publication just isn't available for that request.
523 pub boundary_capture: Option<SpecBoundaryCapture>,
524}
525impl SpecSession {
526 /// Context capacity of the session's caches (the server's ContextFull guard).
527 pub fn cache_max_ctx(&self) -> usize {
528 self.cache.max_ctx
529 }
530 /// Read access to the live trunk cache (lane/spec-prefix-cache): the worker slices
531 /// full-attn KV rows `[0..capture.pos)` out of it when publishing a boundary capture —
532 /// those rows are append-only for the session's lifetime (rollbacks never truncate below
533 /// the prime boundary), so no copy was taken at prime time.
534 pub fn cache_ref(&self) -> &Cache {
535 &self.cache
536 }
537 /// Snapshot the session's process-local acceptance counters for per-burst diffing.
538 pub fn telemetry(&self) -> SpecTelemetry {
539 self.telem.snapshot()
540 }
541 /// Committed position this session can REWIND to (its retained prompt-end boundary), if any.
542 /// A request whose prompt matches `committed[..pos]` exactly can resume from here — see
543 /// `spec_rewind_to_checkpoint`.
544 pub fn rewind_pos(&self) -> Option<usize> {
545 self.turn_ckpt.as_ref().map(|c| c.pos)
546 }
547 /// Whether every ring-backed trunk/draft row needed by the retained checkpoint is resident.
548 pub fn rewind_is_resident(&self) -> bool {
549 self.turn_ckpt.as_ref().is_some_and(|ckpt| {
550 self.cache.can_rollback(&ckpt.snap, 0) && self.scratch.can_rewind_to(ckpt.pos)
551 })
552 }
553 /// Is this session in the DEMOTION-READY shape (see [`SpecSession::into_demoted`])?
554 /// `false` means a carried pending must be flushed first (`spec_flush_pending`), or the
555 /// session has never run a turn and has no prediction to hand over.
556 pub fn demote_ready(&self) -> bool {
557 self.pending_tok.is_none() && self.next_pred.is_some()
558 }
559 /// Does this session hold a carried pending bonus (flush required before a handoff/park)?
560 pub fn has_pending(&self) -> bool {
561 self.pending_tok.is_some()
562 }
563 /// Committed row count == cache rows (the session invariant), for the caller's own
564 /// `fed`-length cross-check at a handoff boundary.
565 pub fn committed_len(&self) -> usize {
566 self.committed.len()
567 }
568 /// DEMOTION HANDOFF (lane/spec-gate, 2026-08-07): consume this session and hand its trunk
569 /// cache + next-token prediction to the plain batched-decode path.
570 ///
571 /// WHY THIS IS EXACT (greedy). The invariant at a burst boundary is `cache.pos ==
572 /// committed.len()`: every committed row has trunk KV + recurrent state, exactly as a plain
573 /// tokenwise prime of the same `committed` sequence would have left it (that is the
574 /// session-tail contract, and the same property `spec_rewind_to_checkpoint` and the reuse
575 /// pool already rely on). `next_pred` is the argmax of the verify's logits for the LAST
576 /// committed row — and verify-column logits are bit-identical to plain decode's logits at
577 /// that position, because `matmul_decode_exact` bit-identity IS the basis of the greedy
578 /// accept walk. So handing (cache, next_pred) to the batched path continues the stream from
579 /// a state indistinguishable from one the batched path produced itself: the batched tick
580 /// emits `next_pred`, feeds it into this same cache, and decodes on.
581 ///
582 /// `None` when the session is not in the handoff shape — a carried pending (its bonus row is
583 /// NOT in the cache, so `spec_flush_pending` must commit it first) or no `next_pred` yet
584 /// (never bursted). Callers must not force it: a half-committed cache handed to the batched
585 /// path would silently skip a token.
586 ///
587 /// The MTP draft scratch, the persistent draft-graph context and the turn checkpoint are
588 /// DROPPED here (freeing their VRAM): the batched path never drafts, and this handoff is
589 /// one-way by design — there is no cheap symmetric re-promotion (rebuilding the draft KV
590 /// would mean an `mtp_kv_fill` over the whole committed history).
591 pub fn into_demoted(self) -> Option<(Cache, u32)> {
592 if self.pending_tok.is_some() {
593 return None;
594 }
595 let np = self.next_pred?;
596 debug_assert_eq!(
597 self.cache.pos,
598 self.committed.len(),
599 "demotion handoff: cache rows != committed tokens"
600 );
601 Some((self.cache, np))
602 }
603 /// Pool-resume hook (audit Q2): clear the parked draft-graph failure memoization so a
604 /// NEW request resuming this session gets one fresh capture chance — a transient-pressure
605 /// capture failure must not persist for the pool's whole lifetime (the TRT #16072 class).
606 /// Logs once iff a flag was actually set; a no-fallback resume is silent and free.
607 pub fn reset_graph_fallback_on_resume(&mut self) {
608 if let Some(line) = self
609 .draft_ctx
610 .as_mut()
611 .and_then(|c| c.failed.reset_on_resume())
612 {
613 eprintln!("{line}");
614 }
615 }
616}
617
618/// A session's PROMPT-END boundary state, the rewind target for session-affinity resume.
619///
620/// WHY THIS BOUNDARY, AND WHY IT IS THE ONLY ONE WORTH KEEPING. The rewrite class this lane
621/// exists for (a client that strips `<think>` blocks out of prior assistant turns) mutates the
622/// text the session GENERATED, never the prompt it was given. So turn N's prompt agrees with
623/// turn N-1's committed tokens up to almost exactly where turn N-1's generation began — the
624/// prompt-end boundary. Keeping a checkpoint there means the next turn re-primes only its own
625/// delta (the rewritten answer + the new user turn) instead of the whole conversation.
626///
627/// WHAT IT MUST HOLD. Full-attn KV is append-only and position-addressed, so rewinding it is a
628/// `len` truncation (no data). Linear-attn (GDN) conv/ssm state is mutated IN PLACE with no
629/// position index, so it must be a real device COPY — that copy is the entire reason a spec
630/// session could not previously rewind. The MTP draft scratch needs no copy either: its rows
631/// below the boundary were written by this turn's fill and are never revisited (the per-round
632/// true-hidden refresh only rewrites the CURRENT burst's committed positions), so rewinding it
633/// is also just a `len` reset. `last_h` is the hidden of the last row below the boundary — the
634/// predecessor-pairing anchor the next prime's fill reads for its first row.
635///
636/// COST: one `Cache::snapshot` per TURN, on a code path that already takes one per ROUND.
637pub(crate) struct SpecCheckpoint {
638 snap: crate::cache::CacheSnapshot,
639 /// Committed length at the boundary (== cache.pos there, the session invariant).
640 pos: usize,
641 /// Pre-output_norm hidden of row `pos - 1`.
642 last_h: CudaSlice<f32>,
643}
644
645/// PREFIX-CACHE BOUNDARY CAPTURE (lane/spec-prefix-cache, 2026-08-14): the state a spec session
646/// records at its cold-prime split so the WORKER can publish a cross-request prefix entry —
647/// the commit-gated-publication port (research/cache-spec-design-20260814/PORT-PLAN.md item 1).
648/// Only the pieces that are DESTROYED by continuing the prime need copies here: the in-place
649/// GDN conv/ssm states (via `Cache::snapshot`, same mechanism as [`SpecCheckpoint`]) and the
650/// boundary logits. Full-attn KV rows `[0..pos)` and draft-scratch rows `[0..pos)` are
651/// append-only for the session's lifetime (rollbacks never truncate below the prime boundary),
652/// so the worker slices those from the live caches post-burst instead of copying at prime time.
653pub struct SpecBoundaryCapture {
654 pub snap: crate::cache::CacheSnapshot,
655 /// Token boundary (== cache.pos at capture; == the worker's miss-LCP split).
656 pub pos: usize,
657 /// Full-vocab logits after the prefix prime — the entry's boundary logits.
658 pub logits: Vec<f32>,
659}
660
661struct SpecPipeTraceClock {
662 pair: usize,
663 started: std::time::Instant,
664}
665
666#[derive(Clone)]
667struct SpecPipeTraceCtx {
668 clock: std::sync::Arc<SpecPipeTraceClock>,
669 round: usize,
670 lane: usize,
671}
672
673struct SpecPipeTraceMarker {
674 trace: SpecPipeTraceCtx,
675 phase: &'static str,
676 edge: &'static str,
677 slot: Option<usize>,
678}
679
680unsafe extern "C" fn spec_pipe_trace_marker(raw: *mut std::ffi::c_void) {
681 let marker = unsafe { Box::from_raw(raw.cast::<SpecPipeTraceMarker>()) };
682 let lane = if marker.trace.lane == 0 { "A" } else { "B" };
683 let slot = marker
684 .slot
685 .map(|v| v.to_string())
686 .unwrap_or_else(|| "-".into());
687 let t_ms = marker.trace.clock.started.elapsed().as_secs_f64() * 1e3;
688 use std::io::Write as _;
689 let stderr = std::io::stderr();
690 let mut stderr = stderr.lock();
691 let _ = writeln!(
692 stderr,
693 "[spec-pipe-timeline] pair={} round={} lane={lane} phase={} edge={} \
694 slot={slot} t_ms={t_ms:.3}",
695 marker.trace.clock.pair, marker.trace.round, marker.phase, marker.edge,
696 );
697}
698
699fn enqueue_spec_pipe_trace_marker(
700 stream: &cudarc::driver::CudaStream,
701 trace: Option<&SpecPipeTraceCtx>,
702 phase: &'static str,
703 edge: &'static str,
704 slot: Option<usize>,
705) -> Result<(), Box<dyn std::error::Error>> {
706 let Some(trace) = trace else {
707 return Ok(());
708 };
709 let marker = Box::new(SpecPipeTraceMarker {
710 trace: trace.clone(),
711 phase,
712 edge,
713 slot,
714 });
715 let raw = Box::into_raw(marker);
716 let result = unsafe {
717 cudarc::driver::result::stream::launch_host_function(
718 stream.cu_stream(),
719 spec_pipe_trace_marker,
720 raw.cast(),
721 )
722 };
723 if let Err(err) = result {
724 unsafe {
725 drop(Box::from_raw(raw));
726 }
727 return Err(err.into());
728 }
729 Ok(())
730}
731
732#[derive(Default)]
733struct SpecPipeProgress {
734 setup_done: [bool; 2],
735 draft_done: [usize; 2],
736 stage0_done: [usize; 2],
737 verify_done: [usize; 2],
738 accept_done: [usize; 2],
739 finished: [bool; 2],
740 aborted: bool,
741}
742
743/// Host-side issue coordinator for the reduced two-session speculative pipeline. Each session
744/// keeps its existing call stack and round locals; this object only orders phase entry. The
745/// primary mutex spans whole draft/accept/tail issue regions so Engine's single-stream scratch
746/// cannot be interleaved by the two host threads.
747struct SpecPipeSync {
748 progress: std::sync::Mutex<SpecPipeProgress>,
749 changed: std::sync::Condvar,
750 primary: std::sync::Mutex<()>,
751 trace: Option<std::sync::Arc<SpecPipeTraceClock>>,
752}
753
754impl SpecPipeSync {
755 fn new() -> Self {
756 static TRACE_PAIR: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
757 let trace = (std::env::var("MEMRA_SPEC_PIPE_TRACE").as_deref() == Ok("1")).then(|| {
758 std::sync::Arc::new(SpecPipeTraceClock {
759 pair: TRACE_PAIR.fetch_add(1, std::sync::atomic::Ordering::Relaxed) + 1,
760 started: std::time::Instant::now(),
761 })
762 });
763 Self {
764 progress: std::sync::Mutex::new(SpecPipeProgress::default()),
765 changed: std::sync::Condvar::new(),
766 primary: std::sync::Mutex::new(()),
767 trace,
768 }
769 }
770}
771
772#[derive(Clone)]
773struct SpecPipeLane {
774 sync: std::sync::Arc<SpecPipeSync>,
775 lane: usize,
776}
777
778impl SpecPipeLane {
779 fn peer(&self) -> usize {
780 1 - self.lane
781 }
782
783 fn aborted() -> Box<dyn std::error::Error> {
784 "paired speculative peer aborted".into()
785 }
786
787 fn trace(&self, round: usize) -> Option<SpecPipeTraceCtx> {
788 self.sync.trace.as_ref().map(|clock| SpecPipeTraceCtx {
789 clock: clock.clone(),
790 round,
791 lane: self.lane,
792 })
793 }
794
795 fn setup_begin(&self) -> Result<(), Box<dyn std::error::Error>> {
796 let mut p = self.sync.progress.lock().unwrap();
797 while !p.aborted && self.lane == 1 && !p.setup_done[0] && !p.finished[0] {
798 p = self.sync.changed.wait(p).unwrap();
799 }
800 if p.aborted {
801 Err(Self::aborted())
802 } else {
803 Ok(())
804 }
805 }
806
807 fn setup_end(&self) {
808 let mut p = self.sync.progress.lock().unwrap();
809 p.setup_done[self.lane] = true;
810 self.sync.changed.notify_all();
811 }
812
813 fn draft_begin(
814 &self,
815 round: usize,
816 ) -> Result<std::sync::MutexGuard<'_, ()>, Box<dyn std::error::Error>> {
817 let peer = self.peer();
818 let mut p = self.sync.progress.lock().unwrap();
819 loop {
820 if p.aborted {
821 return Err(Self::aborted());
822 }
823 let setup_ready =
824 (p.setup_done[0] || p.finished[0]) && (p.setup_done[1] || p.finished[1]);
825 let prior_ready = p.accept_done[self.lane] >= round
826 && (p.accept_done[peer] >= round || p.finished[peer]);
827 let turn_ready = if self.lane == 0 {
828 true
829 } else {
830 p.draft_done[0] > round || p.finished[0]
831 };
832 if setup_ready && prior_ready && turn_ready {
833 break;
834 }
835 p = self.sync.changed.wait(p).unwrap();
836 }
837 drop(p);
838 Ok(self.sync.primary.lock().unwrap())
839 }
840
841 fn draft_end(&self, round: usize) {
842 let mut p = self.sync.progress.lock().unwrap();
843 p.draft_done[self.lane] = round + 1;
844 self.sync.changed.notify_all();
845 }
846
847 /// Admit stage 0 and return whether this lane owns the interval's one reverse fence.
848 /// Lane B releases as soon as lane A has issued its boundary TX, not after A's full body.
849 fn stage0_begin(&self, round: usize) -> Result<bool, Box<dyn std::error::Error>> {
850 let peer = self.peer();
851 let mut p = self.sync.progress.lock().unwrap();
852 loop {
853 if p.aborted {
854 return Err(Self::aborted());
855 }
856 let ready = if self.lane == 0 {
857 p.draft_done[0] > round && (p.draft_done[1] > round || p.finished[1])
858 } else {
859 p.draft_done[1] > round && (p.stage0_done[0] > round || p.finished[0])
860 };
861 if ready {
862 return Ok(self.lane == 0 || p.finished[peer]);
863 }
864 p = self.sync.changed.wait(p).unwrap();
865 }
866 }
867
868 fn stage0_end(&self, round: usize) {
869 let mut p = self.sync.progress.lock().unwrap();
870 p.stage0_done[self.lane] = round + 1;
871 self.sync.changed.notify_all();
872 }
873
874 /// Stage 1 is single-owner per engine. A proceeds immediately after its own ticket; B waits
875 /// for A's full stage1/head issue so only A.S1 and B.S0 can overlap.
876 fn stage1_begin(&self, round: usize) -> Result<(), Box<dyn std::error::Error>> {
877 let mut p = self.sync.progress.lock().unwrap();
878 while !p.aborted
879 && !(p.stage0_done[self.lane] > round
880 && (self.lane == 0 || p.verify_done[0] > round || p.finished[0]))
881 {
882 p = self.sync.changed.wait(p).unwrap();
883 }
884 if p.aborted {
885 Err(Self::aborted())
886 } else {
887 Ok(())
888 }
889 }
890
891 fn verify_end(&self, round: usize) {
892 let mut p = self.sync.progress.lock().unwrap();
893 p.verify_done[self.lane] = round + 1;
894 self.sync.changed.notify_all();
895 }
896
897 fn accept_begin(
898 &self,
899 round: usize,
900 ) -> Result<std::sync::MutexGuard<'_, ()>, Box<dyn std::error::Error>> {
901 let mut p = self.sync.progress.lock().unwrap();
902 loop {
903 if p.aborted {
904 return Err(Self::aborted());
905 }
906 let ready = if self.lane == 0 {
907 p.verify_done[0] > round && (p.verify_done[1] > round || p.finished[1])
908 } else {
909 p.verify_done[1] > round && (p.accept_done[0] > round || p.finished[0])
910 };
911 if ready {
912 break;
913 }
914 p = self.sync.changed.wait(p).unwrap();
915 }
916 drop(p);
917 Ok(self.sync.primary.lock().unwrap())
918 }
919
920 fn accept_end(&self, round: usize) {
921 let mut p = self.sync.progress.lock().unwrap();
922 p.accept_done[self.lane] = round + 1;
923 self.sync.changed.notify_all();
924 }
925
926 fn primary(&self) -> std::sync::MutexGuard<'_, ()> {
927 self.sync.primary.lock().unwrap()
928 }
929
930 fn finish(&self, failed: bool) {
931 let mut p = self.sync.progress.lock().unwrap();
932 p.finished[self.lane] = true;
933 p.aborted |= failed;
934 self.sync.changed.notify_all();
935 }
936}
937
938struct SpecPipeFinish<'a> {
939 lane: &'a SpecPipeLane,
940 closed: bool,
941}
942
943impl<'a> SpecPipeFinish<'a> {
944 fn new(lane: &'a SpecPipeLane) -> Self {
945 Self {
946 lane,
947 closed: false,
948 }
949 }
950
951 fn close(&mut self, failed: bool) {
952 self.lane.finish(failed);
953 self.closed = true;
954 }
955}
956
957impl Drop for SpecPipeFinish<'_> {
958 fn drop(&mut self) {
959 if !self.closed {
960 self.lane.finish(true);
961 }
962 }
963}
964
965/// Scoped transfer of one exclusively-borrowed session to the second host issue thread.
966/// `CudaGraph` is not marked Send by cudarc because its raw driver handles carry no automatic
967/// trait. CUDA driver graph handles are context-scoped rather than OS-thread-affine; the caller
968/// binds that context before touching the session, joins before returning, and never aliases the
969/// pointer. Keep this exception local to the experimental pair call instead of marking the public
970/// session type Send.
971struct SpecPipeSessionPtr(*mut SpecSession);
972
973unsafe impl Send for SpecPipeSessionPtr {}
974
975impl SpecPipeSessionPtr {
976 unsafe fn get_mut(&mut self) -> &mut SpecSession {
977 unsafe { &mut *self.0 }
978 }
979}
980
981/// Per-session persistent draft-graph context: the captured CUDA graph(s) plus the device
982/// buffers whose POINTERS the capture bakes. Reuse legality: the greedy capture bakes only
983/// session-stable pointers (the session's own MtpScratch KV — allocated once, never realloc'd;
984/// the model's resident embedding; the process-wide OnceLock p_min) and the g_* buffers held
985/// HERE — so one capture serves the session's whole lifetime. The sampled capture additionally
986/// bakes (seed, temp) as capture-time constants and needs k q-slots — keyed by `s_key`, dropped
987/// and recaptured when a pool-resumed request changes them. `*_failed` memoizes a failed capture
988/// so the eager fallback doesn't pay a doomed capture attempt every burst.
989pub(crate) struct DraftGraphCtx {
990 g_tok: CudaSlice<u32>,
991 g_pos: CudaSlice<i32>,
992 g_seed: CudaSlice<f32>,
993 g_p: CudaSlice<f32>,
994 g_ctr: CudaSlice<u32>,
995 g_q: CudaSlice<f32>,
996 g_perturb: CudaSlice<f32>,
997 q_slots: Vec<CudaSlice<f32>>,
998 /// DRAFT-SIDE GRAMMAR MASK (lane/draft-mask): packed allowed-set words over the DRAFT
999 /// head's vocab, at a STABLE address so the captured draft graph's mask node reads the
1000 /// per-position contents the host re-uploads before each replay (the graph-promote
1001 /// pattern from decode.rs). Empty unless the session drafts under a grammar.
1002 g_dmask: CudaSlice<u32>,
1003 /// was `graph` captured WITH the mask node? A parked graph of the wrong shape is dropped.
1004 graph_masked: bool,
1005 graph: Option<cudarc::driver::CudaGraph>,
1006 graph_s: Option<cudarc::driver::CudaGraph>,
1007 /// Failed-capture memoization for both graphs — LOUD on flip, cleared on pool resume
1008 /// (audit Q2, the TRT #16072 silent-permanent-coverage-loss class).
1009 failed: DraftGraphFallback,
1010 /// (seed, temp.to_bits(), k) baked into graph_s at its capture.
1011 s_key: Option<(u64, u32, usize)>,
1012 /// CAPTURE-RETAIN keepers (#68 root cause, 2026-08-04): the warmup-run transients whose
1013 /// pool addresses the captured graph(s) bake. Without these, the transients return to the
1014 /// pool at capture-body exit and later work (burst-boundary prime/fill/commit passes, or a
1015 /// co-served session in the worker) reuses those addresses — the persisted graph's replay
1016 /// then reads/writes live unrelated buffers (exactness corruption, first seen as the ST
1017 /// serve-spec 4B graph-arm corruption; one-shot CLI calls never re-shuffled the pool, which
1018 /// is why run-spec K=1..8 passed on the same checkpoint). Same fix class as
1019 /// capture_graph_retained's gemma/decode.rs sites — hold as long as the graph replays.
1020 keeper: Vec<Box<dyn std::any::Any + Send>>,
1021 keeper_s: Vec<Box<dyn std::any::Any + Send>>,
1022}
1023
1024/// Failed-capture memoization for the two draft graphs (audit Q2, 2026-08-05 — the
1025/// TRT #16072 trap class: pressure-triggered, silent, long-lived coverage loss).
1026///
1027/// Three contracts:
1028/// - LOUD FLIP: `mark_*` returns the warn line exactly on the false→true transition
1029/// (returned, not printed, so the once-per-flip contract is unit-testable); the caller
1030/// `eprintln!`s it UNCONDITIONALLY — a dropped draft graph is never silent. Re-marking
1031/// an already-failed graph returns None (the per-burst memoization that keeps the eager
1032/// fallback from paying a doomed capture attempt every burst).
1033/// - RESET ON RESUME: `reset_on_resume` clears both flags — a parked session resumed by a
1034/// NEW request gets one fresh capture chance instead of carrying a transient-pressure
1035/// failure for the pool's whole lifetime. Returns the note line only when a flag was
1036/// actually set (quiet on the common clean-resume path).
1037/// - Shape-change clears (`clear_*`) stay silent, exactly as before: they precede a fresh
1038/// capture attempt whose own failure would re-flip loudly.
1039#[derive(Default)]
1040pub(crate) struct DraftGraphFallback {
1041 greedy: bool,
1042 sampled: bool,
1043}
1044impl DraftGraphFallback {
1045 fn mark_greedy(&mut self, reason: &str) -> Option<String> {
1046 if self.greedy {
1047 return None;
1048 }
1049 self.greedy = true;
1050 Some(format!(
1051 "[spec] WARN: draft-graph capture failed ({reason}); eager fallback until session resume"
1052 ))
1053 }
1054 fn mark_sampled(&mut self, reason: &str) -> Option<String> {
1055 if self.sampled {
1056 return None;
1057 }
1058 self.sampled = true;
1059 Some(format!(
1060 "[spec] WARN: sampled draft-graph capture failed ({reason}); eager fallback until session resume"
1061 ))
1062 }
1063 fn greedy_failed(&self) -> bool {
1064 self.greedy
1065 }
1066 fn sampled_failed(&self) -> bool {
1067 self.sampled
1068 }
1069 fn clear_greedy(&mut self) {
1070 self.greedy = false;
1071 }
1072 fn clear_sampled(&mut self) {
1073 self.sampled = false;
1074 }
1075 /// Pool-resume reset: both graphs get a fresh capture chance. Some(note) iff any flag
1076 /// was set (so clean resumes stay quiet).
1077 pub(crate) fn reset_on_resume(&mut self) -> Option<String> {
1078 if !self.greedy && !self.sampled {
1079 return None;
1080 }
1081 let which = match (self.greedy, self.sampled) {
1082 (true, true) => "greedy+sampled",
1083 (true, false) => "greedy",
1084 _ => "sampled",
1085 };
1086 self.greedy = false;
1087 self.sampled = false;
1088 Some(format!(
1089 "[spec] draft-graph fallback reset on session resume ({which}); recapture eligible"
1090 ))
1091 }
1092}
1093
1094impl DraftGraphCtx {
1095 fn new(e: &Engine, n_embd: usize, qlen: usize) -> Result<Self, Box<dyn std::error::Error>> {
1096 Ok(DraftGraphCtx {
1097 g_tok: e.alloc_u32_zeroed(1)?,
1098 g_pos: e.htod_i32(&[0])?,
1099 g_seed: e.zeros(n_embd)?,
1100 g_p: e.zeros(1)?,
1101 g_ctr: e.alloc_u32_zeroed(1)?,
1102 g_q: e.zeros(qlen)?,
1103 g_perturb: e.zeros(qlen)?,
1104 q_slots: Vec::new(),
1105 g_dmask: e.alloc_u32_zeroed(1)?,
1106 graph_masked: false,
1107 graph: None,
1108 graph_s: None,
1109 failed: DraftGraphFallback::default(),
1110 s_key: None,
1111 keeper: Vec::new(),
1112 keeper_s: Vec::new(),
1113 })
1114 }
1115}
1116
1117pub(crate) struct MtpScratch {
1118 kv: KvLayer,
1119 /// Logical row capacity. On the graph/DC draft path it also doubles as fa_decode_dc's
1120 /// bucket_max: n_splits is sized from it ONCE, so the graph captured at round 0 stays valid
1121 /// for every later t_kv. Step35 refuses that path and may back this logical extent with the
1122 /// smaller host-indexed SWA ring instead.
1123 cap: usize,
1124}
1125
1126fn mtp_scratch_layout(
1127 cfg: &memra_gguf::config::ModelConfig,
1128 geom: Option<&crate::hybrid::DraftGeom>,
1129) -> (usize, usize, usize, usize) {
1130 // Student draft heads carry fewer KV heads (head_dim unchanged) -> smaller scratch rows.
1131 let n_head_kv = geom.map(|g| g.n_head_kv).unwrap_or(cfg.n_head_kv as usize);
1132 let head_dim_k = cfg.head_dim_k as usize;
1133 let head_dim_v = cfg.head_dim_v as usize;
1134 assert!(
1135 head_dim_k % 32 == 0 && head_dim_v % 32 == 0,
1136 "KVQUANT requires head_dim%32==0 (MTP scratch)"
1137 );
1138 let kv_dim_k = head_dim_k * n_head_kv;
1139 let kv_dim_v = head_dim_v * n_head_kv;
1140 // The fp8-KV arm deliberately does not reach the draft scratch; keep the exact format
1141 // policy shared with `MtpScratch::new` so admission scales the same allocation.
1142 let (kbb, vbb) = crate::kv_blk_bytes();
1143 let k_tok_bytes = (kv_dim_k / 32) * kbb;
1144 let v_tok_bytes = (kv_dim_v / 32) * vbb;
1145 (kv_dim_k, kv_dim_v, k_tok_bytes, v_tok_bytes)
1146}
1147
1148impl MtpScratch {
1149 fn new(
1150 e: &Engine,
1151 cfg: &memra_gguf::config::ModelConfig,
1152 cap: usize,
1153 geom: Option<&crate::hybrid::DraftGeom>,
1154 ) -> Result<Self, Box<dyn std::error::Error>> {
1155 // env-selected KV formats (default 34/24). The fp8-KV arm (MEMRA_KV_FP8) deliberately
1156 // does NOT reach the draft scratch: fp8 drafts drifted acceptance 69-88% -> 46%
1157 // (2026-07-12 A/B); the scratch is tiny, so it keeps baseline q8_0/q5_1 numerics
1158 // while the TRUNK cache carries the fp8 depth win. Scratch append/fa pass g=false.
1159 let (kv_dim_k, kv_dim_v, k_tok_bytes, v_tok_bytes) = mtp_scratch_layout(cfg, geom);
1160 let ring = if crate::cache::swa_ring_on() && cfg.arch.is_step35() {
1161 let window = cfg.step35.as_ref().unwrap().sliding_window as usize;
1162 Some(crate::cache::KvRing::new(
1163 crate::cache::swa_ring_rows(window, cap),
1164 window,
1165 ))
1166 } else {
1167 None
1168 };
1169 let alloc_rows = ring.as_ref().map(crate::cache::KvRing::rows).unwrap_or(cap);
1170 Ok(MtpScratch {
1171 kv: KvLayer {
1172 k: e.alloc_u8(alloc_rows * k_tok_bytes)?,
1173 v: e.alloc_u8(alloc_rows * v_tok_bytes)?,
1174 kv_dim_k,
1175 kv_dim_v,
1176 k_tok_bytes,
1177 v_tok_bytes,
1178 len: 0,
1179 ring,
1180 len_d: e.htod_i32(&[0])?,
1181 },
1182 cap,
1183 })
1184 }
1185 /// Set BOTH length counters: the host mirror AND the device len_d the captured append/fa read
1186 /// (a 4-byte in-place htod — the counter pointer is baked into the graph, never realloc'd).
1187 /// This is the ONLY truncation/rollback mechanism the persistent draft KV needs.
1188 fn set_len(&mut self, e: &Engine, n: usize) -> Result<(), Box<dyn std::error::Error>> {
1189 if self
1190 .kv
1191 .ring
1192 .as_ref()
1193 .is_some_and(|ring| !ring.can_rewind_to(n))
1194 {
1195 return Err("SWA ring MTP checkpoint has been lapped; full re-prime required".into());
1196 }
1197 self.kv.len = n;
1198 e.set_i32_one(&mut self.kv.len_d, n as i32)
1199 }
1200
1201 fn can_rewind_to(&self, n: usize) -> bool {
1202 self.kv
1203 .ring
1204 .as_ref()
1205 .is_none_or(|ring| ring.can_rewind_to(n))
1206 }
1207}
1208
1209/// Retained verify intermediates for the REPLAY-FREE partial accept (2026-07-03, the profiled
1210/// #1 spec cost at long ctx: the partial-accept replay was a DUPLICATE trunk pass — ~0.54 extra
1211/// full weight reads per round — recomputing columns the verify had already produced
1212/// bit-identically). Holds, per linear layer, everything needed to rebuild its recurrent state
1213/// to "after the first j verify columns" WITHOUT re-running the trunk:
1214/// - BATCHED-path layers (`gdn`): the exact token-major inputs the round's ONE gdn_scan
1215/// consumed. A prefix re-run of the SAME kernel (t=j) from the snapshot state is bit-identical
1216/// to the first j iterations of the verify's scan — the kernel's t-loop carries state in
1217/// registers and iteration t never depends on T. `qkv_mixed` (the conv input) feeds the
1218/// pure-copy ring rebuild.
1219/// - PER-COLUMN-path layers (`cols`): dtod clones of (conv_state, ssm_state) taken after each
1220/// column 0..t-2 — pure copies of the actual chain states (the last column is never a rebuild
1221/// target: j <= t-1).
1222/// Full-attn layers need nothing: their verify KV rows are bit-identical to eager's (the
1223/// decode-exact contract; verify-probe pins it), so rollback = len truncation.
1224struct GdnStash {
1225 qkv_mixed: CudaSlice<f32>, // [t, conv_dim] token-major (conv input)
1226 q_l2: CudaSlice<f32>,
1227 k_l2: CudaSlice<f32>,
1228 v_g: CudaSlice<f32>, // [t, num_v, d_state]
1229 g_log: CudaSlice<f32>,
1230 beta: CudaSlice<f32>, // [t, num_v]
1231}
1232struct VerifyCkpt {
1233 gdn: Vec<Option<GdnStash>>, // [n_layer], Some iff batched linear path ran
1234 cols: Vec<Option<Vec<(CudaSlice<f32>, CudaSlice<f32>)>>>, // [n_layer][col] = (conv, ssm) after col
1235}
1236impl VerifyCkpt {
1237 fn new(n_layer: usize) -> Self {
1238 VerifyCkpt {
1239 gdn: (0..n_layer).map(|_| None).collect(),
1240 cols: (0..n_layer).map(|_| None).collect(),
1241 }
1242 }
1243}
1244
1245/// The stage-0/TX half of one PP verify. The boundary slot is the ownership token: stage 1
1246/// consumes exactly the slot selected by `tx()` / `tx_pipelined()`, never a slot inferred from
1247/// a logical round number.
1248struct VerifyBoundaryTicket {
1249 rt: &'static crate::pp::PpNRt,
1250 caller_stream: std::sync::Arc<cudarc::driver::CudaStream>,
1251 slot: usize,
1252 pos0: usize,
1253 t: usize,
1254 payload: usize,
1255 n_st: usize,
1256 pipelined: bool,
1257 pp_anatomy: bool,
1258 pp_started: std::time::Instant,
1259 reverse_ms: f64,
1260 stage0_ms: f64,
1261 tx_ms: f64,
1262 trace: Option<SpecPipeTraceCtx>,
1263}
1264
1265/// Explicit OPTIPIPE diagnostic control. Forced modes are set only by `optipipe-gate`; the
1266/// increment-2 controller can also be armed by the server's fresh-process research door.
1267#[derive(Clone, Copy, Debug, PartialEq, Eq)]
1268pub enum OptiForkGateMode {
1269 Disabled,
1270 Hit,
1271 Miss,
1272 Alternate,
1273 Abort,
1274 Controller,
1275}
1276
1277static OPTI_FORK_GATE_MODE: std::sync::atomic::AtomicU8 = std::sync::atomic::AtomicU8::new(0);
1278static OPTI_CONTROLLER_THRESHOLD: std::sync::atomic::AtomicU32 =
1279 std::sync::atomic::AtomicU32::new(0);
1280static OPTI_FORK_ATTEMPTS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
1281static OPTI_FORK_HITS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
1282static OPTI_FORK_MISSES: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
1283static OPTI_FORK_ABORT_DRAINS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
1284static OPTI_FORK_REFUSALS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
1285static OPTI_GATE_CHECKS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
1286static OPTI_GATE_ADMITS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
1287static OPTI_GATE_REJECTS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
1288static OPTI_RECONCILES: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
1289static OPTI_WASTED_DRAFT_TOKENS: std::sync::atomic::AtomicU64 =
1290 std::sync::atomic::AtomicU64::new(0);
1291static OPTI_SHADOW_DRAFT_TOKENS: std::sync::atomic::AtomicU64 =
1292 std::sync::atomic::AtomicU64::new(0);
1293static OPTI_BREAKER_TRIPS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
1294
1295impl OptiForkGateMode {
1296 fn code(self) -> u8 {
1297 match self {
1298 Self::Disabled => 0,
1299 Self::Hit => 1,
1300 Self::Miss => 2,
1301 Self::Alternate => 3,
1302 Self::Abort => 4,
1303 Self::Controller => 5,
1304 }
1305 }
1306
1307 fn configured() -> Self {
1308 match OPTI_FORK_GATE_MODE.load(std::sync::atomic::Ordering::Relaxed) {
1309 1 => Self::Hit,
1310 2 => Self::Miss,
1311 3 => Self::Alternate,
1312 4 => Self::Abort,
1313 5 => Self::Controller,
1314 _ => Self::Disabled,
1315 }
1316 }
1317
1318 fn action(self, generation: u64) -> OptiForkAction {
1319 match self {
1320 Self::Hit => OptiForkAction::Hit,
1321 Self::Miss => OptiForkAction::Miss,
1322 Self::Alternate if generation & 1 == 0 => OptiForkAction::Hit,
1323 Self::Alternate => OptiForkAction::Miss,
1324 Self::Abort => OptiForkAction::Abort,
1325 Self::Disabled | Self::Controller => {
1326 unreachable!("non-forced mode cannot choose a forced fork action")
1327 }
1328 }
1329 }
1330
1331 fn is_forced(self) -> bool {
1332 matches!(self, Self::Hit | Self::Miss | Self::Alternate | Self::Abort)
1333 }
1334}
1335
1336/// Arm or disarm the forced harness. Serving uses only `set_optipipe_controller_threshold`.
1337pub fn set_optipipe_gate_mode(mode: OptiForkGateMode) {
1338 OPTI_FORK_GATE_MODE.store(mode.code(), std::sync::atomic::Ordering::Relaxed);
1339}
1340
1341/// Arm the increment-2 diagnostic controller. The threshold applies to the uncalibrated
1342/// two-token draft-probability product. Serving can call this only through its explicit
1343/// fresh-process research door; the absent-door default remains byte-for-byte disabled.
1344pub fn set_optipipe_controller_threshold(threshold: f32) {
1345 assert!(threshold.is_finite() && (0.0..=1.0).contains(&threshold));
1346 OPTI_CONTROLLER_THRESHOLD.store(threshold.to_bits(), std::sync::atomic::Ordering::Relaxed);
1347 set_optipipe_gate_mode(OptiForkGateMode::Controller);
1348}
1349
1350#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
1351pub struct OptiForkGateStats {
1352 pub attempts: u64,
1353 pub hits: u64,
1354 pub misses: u64,
1355 pub abort_drains: u64,
1356 pub refusals: u64,
1357 pub gate_checks: u64,
1358 pub gate_admits: u64,
1359 pub gate_rejects: u64,
1360 pub reconciles: u64,
1361 pub wasted_draft_tokens: u64,
1362 pub shadow_draft_tokens: u64,
1363 pub breaker_trips: u64,
1364}
1365
1366#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
1367pub struct OptiForkStateIdentity {
1368 pub trunk_kv_bytes: usize,
1369 pub recurrent_bytes: usize,
1370 pub scratch_kv_bytes: usize,
1371 pub hidden_bytes: usize,
1372}
1373
1374pub fn reset_optipipe_gate_stats() {
1375 for counter in [
1376 &OPTI_FORK_ATTEMPTS,
1377 &OPTI_FORK_HITS,
1378 &OPTI_FORK_MISSES,
1379 &OPTI_FORK_ABORT_DRAINS,
1380 &OPTI_FORK_REFUSALS,
1381 &OPTI_GATE_CHECKS,
1382 &OPTI_GATE_ADMITS,
1383 &OPTI_GATE_REJECTS,
1384 &OPTI_RECONCILES,
1385 &OPTI_WASTED_DRAFT_TOKENS,
1386 &OPTI_SHADOW_DRAFT_TOKENS,
1387 &OPTI_BREAKER_TRIPS,
1388 ] {
1389 counter.store(0, std::sync::atomic::Ordering::Relaxed);
1390 }
1391}
1392
1393pub fn optipipe_gate_stats() -> OptiForkGateStats {
1394 let load = |v: &std::sync::atomic::AtomicU64| v.load(std::sync::atomic::Ordering::Relaxed);
1395 OptiForkGateStats {
1396 attempts: load(&OPTI_FORK_ATTEMPTS),
1397 hits: load(&OPTI_FORK_HITS),
1398 misses: load(&OPTI_FORK_MISSES),
1399 abort_drains: load(&OPTI_FORK_ABORT_DRAINS),
1400 refusals: load(&OPTI_FORK_REFUSALS),
1401 gate_checks: load(&OPTI_GATE_CHECKS),
1402 gate_admits: load(&OPTI_GATE_ADMITS),
1403 gate_rejects: load(&OPTI_GATE_REJECTS),
1404 reconciles: load(&OPTI_RECONCILES),
1405 wasted_draft_tokens: load(&OPTI_WASTED_DRAFT_TOKENS),
1406 shadow_draft_tokens: load(&OPTI_SHADOW_DRAFT_TOKENS),
1407 breaker_trips: load(&OPTI_BREAKER_TRIPS),
1408 }
1409}
1410
1411#[derive(Clone, Copy, Debug)]
1412struct OptiControllerPolicy {
1413 threshold: f32,
1414 consecutive_misses: u8,
1415 breaker_tripped: bool,
1416}
1417
1418impl OptiControllerPolicy {
1419 fn configured() -> Self {
1420 Self {
1421 threshold: f32::from_bits(
1422 OPTI_CONTROLLER_THRESHOLD.load(std::sync::atomic::Ordering::Relaxed),
1423 ),
1424 consecutive_misses: 0,
1425 breaker_tripped: false,
1426 }
1427 }
1428
1429 fn admit(&self, q_proxy: f32) -> bool {
1430 q_proxy.is_finite()
1431 && (0.0..=1.0).contains(&q_proxy)
1432 && (self.threshold == 0.0 || (!self.breaker_tripped && q_proxy >= self.threshold))
1433 }
1434
1435 /// Returns true exactly when this resolution newly trips the three-miss breaker.
1436 fn resolve(&mut self, hit: bool) -> bool {
1437 // q*=0 is the lane's explicit unconditional measurement arm. Its purpose is to price
1438 // every optimistic opportunity, so the safety breaker is measured separately and must
1439 // not silently turn this arm into "three attempts then serial".
1440 if self.threshold == 0.0 {
1441 self.consecutive_misses = 0;
1442 return false;
1443 }
1444 if hit {
1445 self.consecutive_misses = 0;
1446 return false;
1447 }
1448 self.consecutive_misses = self.consecutive_misses.saturating_add(1);
1449 if !self.breaker_tripped && self.consecutive_misses >= 3 {
1450 self.breaker_tripped = true;
1451 return true;
1452 }
1453 false
1454 }
1455}
1456
1457#[derive(Clone, Copy, Debug, PartialEq, Eq)]
1458enum OptiForkAction {
1459 Hit,
1460 Miss,
1461 Abort,
1462}
1463
1464#[derive(Clone, Copy, Debug, PartialEq, Eq)]
1465struct OptiForkGeneration {
1466 id: u64,
1467 slot: usize,
1468}
1469
1470#[derive(Default)]
1471struct OptiForkGenerationTracker {
1472 next: u64,
1473 live: [Option<u64>; 2],
1474}
1475
1476impl OptiForkGenerationTracker {
1477 fn reserve(&mut self) -> Result<OptiForkGeneration, Box<dyn std::error::Error>> {
1478 let generation = OptiForkGeneration {
1479 id: self.next,
1480 slot: (self.next & 1) as usize,
1481 };
1482 if let Some(live) = self.live[generation.slot] {
1483 return Err(format!(
1484 "optipipe snapshot slot {} still owns generation {live}; refusing to overwrite it",
1485 generation.slot,
1486 )
1487 .into());
1488 }
1489 self.next += 1;
1490 self.live[generation.slot] = Some(generation.id);
1491 Ok(generation)
1492 }
1493
1494 fn retire(&mut self, generation: OptiForkGeneration) -> Result<(), Box<dyn std::error::Error>> {
1495 match self.live[generation.slot] {
1496 Some(id) if id == generation.id => {
1497 self.live[generation.slot] = None;
1498 Ok(())
1499 }
1500 other => Err(format!(
1501 "optipipe generation teardown mismatch: ticket={} slot={} live={other:?}",
1502 generation.id, generation.slot,
1503 )
1504 .into()),
1505 }
1506 }
1507}
1508
1509struct OptiForkSeedGeneration {
1510 h_seed: CudaSlice<f32>,
1511 fill_prev: CudaSlice<f32>,
1512 scratch_len: usize,
1513}
1514
1515/// Allocate or refresh one full checkpoint through the engine that owns each PP stage. The
1516/// generic cache helper accepts one device and therefore cannot copy GDN state split across
1517/// devices. KV lengths and position stay host metadata; only recurrent buffers need stage-local
1518/// device ownership.
1519fn opti_snapshot_stage_owned(
1520 e: &Engine,
1521 cache: &Cache,
1522 rt: &'static crate::pp::PpNRt,
1523 fence: &[usize],
1524) -> Result<crate::cache::CacheSnapshot, Box<dyn std::error::Error>> {
1525 let n = cache.kv.len();
1526 let mut snapshot = crate::cache::CacheSnapshot {
1527 kv_len: vec![None; n],
1528 conv: (0..n).map(|_| None).collect(),
1529 ssm: (0..n).map(|_| None).collect(),
1530 pos: cache.pos,
1531 };
1532 opti_snapshot_stage_owned_into(e, cache, rt, fence, &mut snapshot)?;
1533 Ok(snapshot)
1534}
1535
1536fn opti_snapshot_stage_owned_into(
1537 e: &Engine,
1538 cache: &Cache,
1539 rt: &'static crate::pp::PpNRt,
1540 fence: &[usize],
1541 snapshot: &mut crate::cache::CacheSnapshot,
1542) -> Result<(), Box<dyn std::error::Error>> {
1543 if fence.len() != rt.n_stages() + 1 || snapshot.kv_len.len() != cache.kv.len() {
1544 return Err("optipipe stage-owned snapshot shape mismatch".into());
1545 }
1546 for stage in 0..rt.n_stages() {
1547 opti_snapshot_one_stage_owned_into(e, cache, rt, fence, stage, snapshot)?;
1548 }
1549 snapshot.pos = cache.pos;
1550 Ok(())
1551}
1552
1553/// Refresh one PP stage of a checkpoint. Increment 2 uses this split form so stage 0's
1554/// optimistic post-N state is captured before N+1 stage 0 is queued, while stage 1's matching
1555/// post-N state is captured only after N stage 1 is enqueued. Calling the all-stage helper at
1556/// either point would capture one side of the fork at the wrong generation.
1557fn opti_snapshot_one_stage_owned_into(
1558 e: &Engine,
1559 cache: &Cache,
1560 rt: &'static crate::pp::PpNRt,
1561 fence: &[usize],
1562 stage: usize,
1563 snapshot: &mut crate::cache::CacheSnapshot,
1564) -> Result<(), Box<dyn std::error::Error>> {
1565 if fence.len() != rt.n_stages() + 1
1566 || snapshot.kv_len.len() != cache.kv.len()
1567 || stage >= rt.n_stages()
1568 {
1569 return Err("optipipe single-stage snapshot shape mismatch".into());
1570 }
1571 let _scope = rt.enter(stage);
1572 let owner = rt.engine(stage, e);
1573 for il in fence[stage]..fence[stage + 1] {
1574 snapshot.kv_len[il] = cache.kv[il].as_ref().map(|kv| kv.len);
1575 match &cache.recur[il] {
1576 Some(recur) => {
1577 match snapshot.conv[il].as_mut() {
1578 Some(dst) => {
1579 owner.copy_into(dst, 0, &recur.conv_state, recur.conv_state.len())?
1580 }
1581 None => snapshot.conv[il] = Some(owner.clone_dtod(&recur.conv_state)?),
1582 }
1583 match snapshot.ssm[il].as_mut() {
1584 Some(dst) => {
1585 owner.copy_into(dst, 0, &recur.ssm_state, recur.ssm_state.len())?
1586 }
1587 None => snapshot.ssm[il] = Some(owner.clone_dtod(&recur.ssm_state)?),
1588 }
1589 }
1590 None if snapshot.conv[il].is_some() || snapshot.ssm[il].is_some() => {
1591 return Err(
1592 format!("optipipe stage-owned snapshot layer {il} changed shape").into(),
1593 );
1594 }
1595 None => {}
1596 }
1597 }
1598 snapshot.pos = cache.pos;
1599 Ok(())
1600}
1601
1602/// Increment-1 persistent fork state. Exactly two snapshot/seed slots alternate; a live ticket
1603/// names its generation and keeps teardown fail-closed. Only stage 0 is allowed to mutate before
1604/// resolve, so the reconcile tables and conditional restores are stage-local.
1605struct OptiForkState {
1606 mode: OptiForkGateMode,
1607 controller: Option<OptiControllerPolicy>,
1608 generations: OptiForkGenerationTracker,
1609 active_snapshot_slot: usize,
1610 alternate_snapshot: crate::cache::CacheSnapshot,
1611 seeds: [OptiForkSeedGeneration; 2],
1612 rt: &'static crate::pp::PpNRt,
1613 fence: [usize; 3],
1614 split: usize,
1615 len_ptrs: CudaSlice<u64>,
1616 saved_lens: CudaSlice<i32>,
1617 forced_acc: CudaSlice<u32>,
1618 valid: CudaSlice<u32>,
1619 stage0_stream: std::sync::Arc<cudarc::driver::CudaStream>,
1620 logical_payload_bytes: [usize; 2],
1621}
1622
1623struct OptiForkTicket {
1624 generation: OptiForkGeneration,
1625 boundary: Option<VerifyBoundaryTicket>,
1626 drain: std::sync::Arc<cudarc::driver::CudaStream>,
1627 settled: bool,
1628}
1629
1630struct OptiControllerTicket {
1631 generation: OptiForkGeneration,
1632 boundary: Option<VerifyBoundaryTicket>,
1633 ckpt: Option<VerifyCkpt>,
1634 verify_tokens: [u32; 2],
1635 draft_prob: f32,
1636 eager_seed: Option<CudaSlice<f32>>,
1637 q_proxy: f32,
1638 scratch_len: usize,
1639 issued_at: std::time::Instant,
1640 drain: std::sync::Arc<cudarc::driver::CudaStream>,
1641 settled: bool,
1642}
1643
1644struct OptiControllerPrepared {
1645 verify_tokens: [u32; 2],
1646 draft_prob: f32,
1647 eager_seed: Option<CudaSlice<f32>>,
1648 q_proxy: f32,
1649 scratch_len: usize,
1650}
1651
1652impl OptiControllerTicket {
1653 fn take_boundary(&mut self) -> VerifyBoundaryTicket {
1654 self.boundary
1655 .take()
1656 .expect("controller boundary ticket already consumed")
1657 }
1658
1659 fn take_ckpt(&mut self) -> VerifyCkpt {
1660 self.ckpt
1661 .take()
1662 .expect("controller verify checkpoint already consumed")
1663 }
1664
1665 fn take_eager_seed(&mut self) -> Option<CudaSlice<f32>> {
1666 self.eager_seed.take()
1667 }
1668
1669 fn settle(&mut self) {
1670 self.settled = true;
1671 }
1672}
1673
1674impl Drop for OptiControllerTicket {
1675 fn drop(&mut self) {
1676 if !self.settled {
1677 let _ = self.drain.synchronize();
1678 OPTI_FORK_ABORT_DRAINS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
1679 }
1680 }
1681}
1682
1683impl OptiForkTicket {
1684 fn take_boundary(&mut self) -> VerifyBoundaryTicket {
1685 self.boundary
1686 .take()
1687 .expect("fork ticket boundary already consumed")
1688 }
1689
1690 fn settle(&mut self) {
1691 self.settled = true;
1692 }
1693}
1694
1695impl Drop for OptiForkTicket {
1696 fn drop(&mut self) {
1697 if !self.settled {
1698 let _ = self.drain.synchronize();
1699 OPTI_FORK_ABORT_DRAINS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
1700 }
1701 }
1702}
1703
1704impl OptiForkState {
1705 #[allow(clippy::too_many_arguments)]
1706 fn new(
1707 e: &Engine,
1708 cache: &Cache,
1709 mode: OptiForkGateMode,
1710 alternate_snapshot: crate::cache::CacheSnapshot,
1711 h_seed: &CudaSlice<f32>,
1712 fill_prev: &CudaSlice<f32>,
1713 rt: &'static crate::pp::PpNRt,
1714 split: usize,
1715 n_layer: usize,
1716 ) -> Result<Self, Box<dyn std::error::Error>> {
1717 let fence = [0, split, n_layer];
1718 let mut logical_payload_bytes = [0usize; 2];
1719 for stage in 0..2 {
1720 for il in fence[stage]..fence[stage + 1] {
1721 logical_payload_bytes[stage] += alternate_snapshot.conv[il]
1722 .as_ref()
1723 .map_or(0, |v| v.len() * std::mem::size_of::<f32>());
1724 logical_payload_bytes[stage] += alternate_snapshot.ssm[il]
1725 .as_ref()
1726 .map_or(0, |v| v.len() * std::mem::size_of::<f32>());
1727 }
1728 }
1729 let seeds = [
1730 OptiForkSeedGeneration {
1731 h_seed: e.clone_dtod(h_seed)?,
1732 fill_prev: e.clone_dtod(fill_prev)?,
1733 scratch_len: 0,
1734 },
1735 OptiForkSeedGeneration {
1736 h_seed: e.clone_dtod(h_seed)?,
1737 fill_prev: e.clone_dtod(fill_prev)?,
1738 scratch_len: 0,
1739 },
1740 ];
1741 let (len_ptrs, saved_lens, forced_acc, valid, stage0_stream) = {
1742 let _stage = rt.enter(0);
1743 let e0 = rt.engine(0, e);
1744 (
1745 crate::round_stream::kv_len_ptr_table_range(e0, cache, 0..split, None)?,
1746 e0.htod_i32(&vec![0; split])?,
1747 e0.alloc_u32_zeroed(2)?,
1748 e0.alloc_u32_zeroed(1)?,
1749 e0.stream(),
1750 )
1751 };
1752 logical_payload_bytes[0] += seeds
1753 .iter()
1754 .map(|seed| (seed.h_seed.len() + seed.fill_prev.len()) * std::mem::size_of::<f32>())
1755 .sum::<usize>();
1756 logical_payload_bytes[0] += len_ptrs.len() * std::mem::size_of::<u64>()
1757 + saved_lens.len() * std::mem::size_of::<i32>()
1758 + forced_acc.len() * std::mem::size_of::<u32>()
1759 + valid.len() * std::mem::size_of::<u32>();
1760 Ok(Self {
1761 mode,
1762 controller: (mode == OptiForkGateMode::Controller)
1763 .then(OptiControllerPolicy::configured),
1764 generations: OptiForkGenerationTracker::default(),
1765 active_snapshot_slot: 0,
1766 alternate_snapshot,
1767 seeds,
1768 rt,
1769 fence,
1770 split,
1771 len_ptrs,
1772 saved_lens,
1773 forced_acc,
1774 valid,
1775 stage0_stream,
1776 logical_payload_bytes,
1777 })
1778 }
1779
1780 fn reserve(
1781 &mut self,
1782 current_snapshot: &mut crate::cache::CacheSnapshot,
1783 ) -> Result<OptiForkGeneration, Box<dyn std::error::Error>> {
1784 let generation = self.generations.reserve()?;
1785 if generation.slot != self.active_snapshot_slot {
1786 std::mem::swap(current_snapshot, &mut self.alternate_snapshot);
1787 self.active_snapshot_slot = generation.slot;
1788 }
1789 Ok(generation)
1790 }
1791
1792 fn capture_seed(
1793 &mut self,
1794 e: &Engine,
1795 generation: OptiForkGeneration,
1796 h_seed: &CudaSlice<f32>,
1797 fill_prev: &CudaSlice<f32>,
1798 scratch_len: usize,
1799 ) -> Result<(), Box<dyn std::error::Error>> {
1800 let seed = &mut self.seeds[generation.slot];
1801 e.copy_into(&mut seed.h_seed, 0, h_seed, h_seed.len())?;
1802 e.copy_into(&mut seed.fill_prev, 0, fill_prev, fill_prev.len())?;
1803 seed.scratch_len = scratch_len;
1804 Ok(())
1805 }
1806
1807 fn ticket(
1808 &self,
1809 generation: OptiForkGeneration,
1810 boundary: VerifyBoundaryTicket,
1811 ) -> OptiForkTicket {
1812 OptiForkTicket {
1813 generation,
1814 boundary: Some(boundary),
1815 drain: self.stage0_stream.clone(),
1816 settled: false,
1817 }
1818 }
1819
1820 #[allow(clippy::too_many_arguments)]
1821 fn controller_ticket(
1822 &self,
1823 generation: OptiForkGeneration,
1824 boundary: VerifyBoundaryTicket,
1825 ckpt: VerifyCkpt,
1826 verify_tokens: [u32; 2],
1827 draft_prob: f32,
1828 eager_seed: Option<CudaSlice<f32>>,
1829 q_proxy: f32,
1830 scratch_len: usize,
1831 ) -> OptiControllerTicket {
1832 OptiControllerTicket {
1833 generation,
1834 boundary: Some(boundary),
1835 ckpt: Some(ckpt),
1836 verify_tokens,
1837 draft_prob,
1838 eager_seed,
1839 q_proxy,
1840 scratch_len,
1841 issued_at: std::time::Instant::now(),
1842 drain: self.stage0_stream.clone(),
1843 settled: false,
1844 }
1845 }
1846
1847 fn reserve_successor(&mut self) -> Result<OptiForkGeneration, Box<dyn std::error::Error>> {
1848 self.generations.reserve()
1849 }
1850
1851 fn successor_snapshot_mut(&mut self) -> &mut crate::cache::CacheSnapshot {
1852 &mut self.alternate_snapshot
1853 }
1854
1855 fn promote_successor_snapshot(
1856 &mut self,
1857 current_snapshot: &mut crate::cache::CacheSnapshot,
1858 generation: OptiForkGeneration,
1859 ) {
1860 std::mem::swap(current_snapshot, &mut self.alternate_snapshot);
1861 self.active_snapshot_slot = generation.slot;
1862 }
1863
1864 fn queue_actual_reconcile(
1865 &mut self,
1866 e: &Engine,
1867 snapshot: &crate::cache::CacheSnapshot,
1868 acc: &CudaSlice<u32>,
1869 optimistic_pending: u32,
1870 base: usize,
1871 ) -> Result<(), Box<dyn std::error::Error>> {
1872 let saved: Vec<i32> = (0..self.split)
1873 .map(|il| snapshot.kv_len[il].map(|v| v as i32).unwrap_or(0))
1874 .collect();
1875 // Serving keeps the caller/accept walk on the head (stage-1) device. Record the accept
1876 // decision point there and append a wait to stage 0 after its optimistic successor/TX;
1877 // the validity/reconcile kernels must never peer-read acc before it is written. The
1878 // increment-1 harness uses primary stage 0, where stream order already provides this.
1879 if self.rt.engine(0, e).ctx().ordinal() != e.ctx().ordinal() {
1880 self.rt.fence_stages_behind(&e.stream())?;
1881 }
1882 let _stage = self.rt.enter(0);
1883 let e0 = self.rt.engine(0, e);
1884 e0.htod_i32_into(&mut self.saved_lens, &saved)?;
1885 e0.spec_fork_valid(acc, optimistic_pending, &mut self.valid)?;
1886 e0.spec_fork_reconcile_kv(
1887 &self.len_ptrs,
1888 &self.saved_lens,
1889 acc,
1890 &self.valid,
1891 base,
1892 self.split,
1893 )
1894 }
1895
1896 fn finish_actual_reconcile(
1897 &mut self,
1898 e: &Engine,
1899 cache: &mut Cache,
1900 snapshot: &crate::cache::CacheSnapshot,
1901 n_acc: usize,
1902 base: usize,
1903 hit: bool,
1904 ) -> Result<(), Box<dyn std::error::Error>> {
1905 if hit {
1906 return Ok(());
1907 }
1908 let len_delta = base + n_acc;
1909 for il in 0..self.split {
1910 if let (Some(kv), Some(saved)) = (cache.kv[il].as_mut(), snapshot.kv_len[il]) {
1911 kv.len = saved + len_delta;
1912 }
1913 }
1914 {
1915 let _stage = self.rt.enter(1);
1916 let e1 = self.rt.engine(1, e);
1917 for il in self.split..self.fence[2] {
1918 if let (Some(kv), Some(saved)) = (cache.kv[il].as_mut(), snapshot.kv_len[il]) {
1919 kv.len = saved + len_delta;
1920 e1.set_i32_one(&mut kv.len_d, kv.len as i32)?;
1921 }
1922 }
1923 }
1924 self.rt.publish_to(0, &e.stream())?;
1925 Ok(())
1926 }
1927
1928 fn cancel_controller_ticket(
1929 &mut self,
1930 e: &Engine,
1931 cache: &mut Cache,
1932 scratch: &mut MtpScratch,
1933 snapshot: &crate::cache::CacheSnapshot,
1934 ticket: &mut OptiControllerTicket,
1935 ) -> Result<(), Box<dyn std::error::Error>> {
1936 {
1937 let _stage = self.rt.enter(0);
1938 let e0 = self.rt.engine(0, e);
1939 for il in 0..self.split {
1940 if let (Some(kv), Some(saved)) = (cache.kv[il].as_mut(), snapshot.kv_len[il]) {
1941 kv.len = saved;
1942 e0.set_i32_one(&mut kv.len_d, saved as i32)?;
1943 }
1944 }
1945 }
1946 scratch.set_len(e, snapshot.pos)?;
1947 ticket.settle();
1948 self.generations.retire(ticket.generation)?;
1949 OPTI_FORK_ABORT_DRAINS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
1950 OPTI_WASTED_DRAFT_TOKENS.fetch_add(2, std::sync::atomic::Ordering::Relaxed);
1951 eprintln!(
1952 "[opti-controller] tail-drain generation={} slot={}",
1953 ticket.generation.id, ticket.generation.slot,
1954 );
1955 Ok(())
1956 }
1957
1958 #[allow(clippy::too_many_arguments)]
1959 fn reconcile(
1960 &mut self,
1961 e: &Engine,
1962 cache: &mut Cache,
1963 scratch: &mut MtpScratch,
1964 snapshot: &crate::cache::CacheSnapshot,
1965 h_seed: &mut CudaSlice<f32>,
1966 fill_prev: &mut CudaSlice<f32>,
1967 generation: OptiForkGeneration,
1968 action: OptiForkAction,
1969 optimistic_pending: u32,
1970 ) -> Result<(), Box<dyn std::error::Error>> {
1971 debug_assert!(action != OptiForkAction::Abort);
1972 let miss_started = std::time::Instant::now();
1973 let keep = action == OptiForkAction::Hit;
1974 let saved: Vec<i32> = (0..self.split)
1975 .map(|il| snapshot.kv_len[il].map(|v| v as i32).unwrap_or(0))
1976 .collect();
1977 let seed = &self.seeds[generation.slot];
1978 {
1979 let _stage = self.rt.enter(0);
1980 let e0 = self.rt.engine(0, e);
1981 e0.htod_i32_into(&mut self.saved_lens, &saved)?;
1982 let forced = if keep {
1983 [1u32, optimistic_pending]
1984 } else {
1985 [0u32, optimistic_pending]
1986 };
1987 e0.htod_u32_into(&mut self.forced_acc, &forced)?;
1988 e0.spec_fork_valid(&self.forced_acc, optimistic_pending, &mut self.valid)?;
1989 e0.spec_fork_reconcile_kv(
1990 &self.len_ptrs,
1991 &self.saved_lens,
1992 &self.forced_acc,
1993 &self.valid,
1994 0,
1995 self.split,
1996 )?;
1997 for il in 0..self.split {
1998 if let Some(recur) = cache.recur[il].as_mut() {
1999 let conv = snapshot.conv[il]
2000 .as_ref()
2001 .ok_or("optipipe stage0 snapshot missing conv state")?;
2002 let ssm = snapshot.ssm[il]
2003 .as_ref()
2004 .ok_or("optipipe stage0 snapshot missing ssm state")?;
2005 e0.spec_fork_restore_f32(conv, &mut recur.conv_state, &self.valid)?;
2006 e0.spec_fork_restore_f32(ssm, &mut recur.ssm_state, &self.valid)?;
2007 }
2008 }
2009 e0.spec_fork_restore_f32(&seed.h_seed, h_seed, &self.valid)?;
2010 e0.spec_fork_restore_f32(&seed.fill_prev, fill_prev, &self.valid)?;
2011 }
2012
2013 if keep {
2014 OPTI_FORK_HITS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2015 return Ok(());
2016 }
2017
2018 for il in 0..self.split {
2019 if let (Some(kv), Some(saved)) = (cache.kv[il].as_mut(), snapshot.kv_len[il]) {
2020 kv.len = saved;
2021 }
2022 }
2023 scratch.set_len(e, seed.scratch_len)?;
2024 // Targeted E_restart: publish only stage 0's reconcile to the caller, then bound the
2025 // forced diagnostic so the retained number is the actual miss cost, not enqueue time.
2026 let caller = e.stream();
2027 self.rt.publish_to(0, &caller)?;
2028 caller.synchronize()?;
2029 let miss_ms = miss_started.elapsed().as_secs_f64() * 1e3;
2030 eprintln!(
2031 "[opti-fork-reconcile] generation={} slot={} miss_ms={miss_ms:.3}",
2032 generation.id, generation.slot,
2033 );
2034 OPTI_FORK_MISSES.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2035 Ok(())
2036 }
2037
2038 fn retire(&mut self, generation: OptiForkGeneration) -> Result<(), Box<dyn std::error::Error>> {
2039 self.generations.retire(generation)
2040 }
2041}
2042
2043impl HybridModel {
2044 fn opti_graph_draft_step(
2045 &self,
2046 e: &Engine,
2047 mtp: &MtpHead,
2048 dctx: &mut DraftGraphCtx,
2049 scratch: &mut MtpScratch,
2050 d_vocab: usize,
2051 ) -> Result<(u32, f32), Box<dyn std::error::Error>> {
2052 dctx.graph
2053 .as_ref()
2054 .ok_or("optipipe controller requires the greedy draft graph")?
2055 .launch()?;
2056 scratch.kv.len += 1;
2057 let idx = e.dtoh_u32_one(&dctx.g_tok)?;
2058 if (idx as usize) >= d_vocab {
2059 return Err(
2060 format!("optipipe draft argmax sentinel 0x{idx:08x} >= d_vocab {d_vocab}").into(),
2061 );
2062 }
2063 let probability = e.dtoh(&dctx.g_p)?[0];
2064 if !probability.is_finite() || !(0.0..=1.0).contains(&probability) {
2065 return Err(format!("optipipe draft probability is invalid: {probability}").into());
2066 }
2067 let token = match &mtp.d2t {
2068 Some(map) => map[idx as usize],
2069 None => idx,
2070 };
2071 if token != idx {
2072 e.set_u32_one(&mut dctx.g_tok, token)?;
2073 }
2074 Ok((token, probability))
2075 }
2076
2077 #[allow(clippy::too_many_arguments)]
2078 fn opti_controller_draft_step(
2079 &self,
2080 e: &Engine,
2081 mtp: &MtpHead,
2082 dctx: &mut DraftGraphCtx,
2083 scratch: &mut MtpScratch,
2084 d_vocab: usize,
2085 eager_state: &mut Option<(u32, CudaSlice<f32>)>,
2086 eager_pos: usize,
2087 embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
2088 ) -> Result<(u32, f32), Box<dyn std::error::Error>> {
2089 if dctx.graph.is_some() {
2090 return self.opti_graph_draft_step(e, mtp, dctx, scratch, d_vocab);
2091 }
2092 let (input_token, input_seed) = eager_state
2093 .take()
2094 .ok_or("optipipe eager continuation seed is unavailable")?;
2095 let (logits, next_seed) = self.mtp_head_forward_dev(
2096 e,
2097 mtp,
2098 input_token,
2099 &input_seed,
2100 scratch,
2101 eager_pos,
2102 embd_dev,
2103 None,
2104 )?;
2105 let token_d = e.argmax_token_device(&logits, d_vocab)?;
2106 let idx = e.dtoh_u32_one(&token_d)?;
2107 if (idx as usize) >= d_vocab {
2108 return Err(format!(
2109 "optipipe eager draft argmax sentinel 0x{idx:08x} >= d_vocab {d_vocab}"
2110 )
2111 .into());
2112 }
2113 let probability_d = e.prob_of_token_device(&logits, &token_d, d_vocab)?;
2114 let probability = e.dtoh(&probability_d)?[0];
2115 if !probability.is_finite() || !(0.0..=1.0).contains(&probability) {
2116 return Err(
2117 format!("optipipe eager draft probability is invalid: {probability}").into(),
2118 );
2119 }
2120 let token = match &mtp.d2t {
2121 Some(map) => map[idx as usize],
2122 None => idx,
2123 };
2124 *eager_state = Some((token, next_seed));
2125 Ok((token, probability))
2126 }
2127
2128 /// NextN head forward for ONE draft token (§A ops 1-13, T=1).
2129 /// Inputs: `e_tok` = the token to predict FROM (last committed / previous draft); `h_seed` =
2130 /// the trunk's pre-output_norm hidden of that token (§A op 2 input). `mtp_pos` = absolute
2131 /// position of the token being predicted from. Returns (draft_logits[n_vocab] host, h_nextn dev).
2132 /// `h_nextn` (§A op 10) becomes `h_seed` for the next autoregressive draft step.
2133 /// Device-resident: returns draft logits ON DEVICE (no [n_vocab] dtoh). The greedy draft
2134 /// loop only needs argmax — paired with `argmax_token_device` this cuts the ~600KB logits
2135 /// transfer + host argmax per draft token from the K-token draft chain.
2136 #[allow(clippy::too_many_arguments)]
2137 fn mtp_head_forward_dev(
2138 &self,
2139 e: &Engine,
2140 mtp: &MtpHead,
2141 e_tok: u32,
2142 h_seed: &CudaSlice<f32>,
2143 scratch: &mut MtpScratch,
2144 mtp_pos: usize,
2145 embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
2146 // DRAFT-SIDE GRAMMAR MASK (lane/draft-mask): (packed draft-vocab allowed set, words).
2147 // Applied to the head logits BEFORE they are returned, so every consumer (argmax,
2148 // gumbel draw, p-min prob) sees the grammar-legal row. None = unmasked (pre-lane).
2149 mask: Option<(&CudaSlice<u32>, usize)>,
2150 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
2151 let cfg = &self.cfg;
2152 let n_embd = cfg.n_embd as usize;
2153 // Distilled-student geometry: the block runs at the INNER width `di` (eh_proj out /
2154 // attn / ffn); the n_embd interface (embed, norms in, carrier out, head in) is unchanged.
2155 let di = mtp.geom.as_ref().map(|g| g.d_inner).unwrap_or(n_embd);
2156 let eps = cfg.rms_eps;
2157 let pos_d = e.htod_i32(&[mtp_pos as i32])?;
2158
2159 // op A: a resident table transfers one 4B token id. The exact host-row capacity path
2160 // expands this one row on CPU and transfers n_embd f32 values instead.
2161 let e_emb = match embd_dev {
2162 Some((g, qt, rb)) => e.embed_gather_device_t(g, &[e_tok], n_embd, qt, rb)?,
2163 None => e.htod(&self.embd.gather(n_embd, &[e_tok]))?,
2164 };
2165
2166 // op 1/2: e_norm = RMSNorm(e, enorm); h_norm = RMSNorm(h_seed, hnorm)
2167 let mut e_norm = e.zeros(n_embd)?;
2168 e.rms_norm(&e_emb, mtp.enorm.float_data(), &mut e_norm, n_embd, 1, eps)?;
2169 let mut h_norm = e.zeros(n_embd)?;
2170 e.rms_norm(h_seed, mtp.hnorm.float_data(), &mut h_norm, n_embd, 1, eps)?;
2171
2172 // op 3: concat = [e_norm ; h_norm] -> [2*n_embd], e_norm in [0,n_embd), h_norm in [n_embd,2n_embd)
2173 let mut concat = e.zeros(2 * n_embd)?;
2174 e.copy_into(&mut concat, 0, &e_norm, n_embd)?;
2175 e.copy_into(&mut concat, n_embd, &h_norm, n_embd)?;
2176
2177 // op 4: inpSA = eh_proj @ concat (eh_proj [2*n_embd, n_embd]) -> [n_embd]
2178 let inp_sa = e.matmul(&mtp.eh_proj, &concat, 1)?;
2179
2180 // op 5: a_norm = RMSNorm(inpSA, attn_norm)
2181 let mut a_norm = e.zeros(di)?;
2182 e.rms_norm(&inp_sa, mtp.attn_norm.float_data(), &mut a_norm, di, 1, eps)?;
2183
2184 // op 6: attention on the scratch KV. SAME dc launcher as the graph path (bucket_max =
2185 // scratch.cap, length from the device len_d) so eager drafts match graph drafts
2186 // bit-for-bit at any t_kv (the parity gate). Host len mirrored here (the dc append
2187 // advances only the device counter).
2188 let attn_out = match (&mtp.mixer, mtp.step35.as_ref()) {
2189 // step35 MTP block: PER-LAYER geometry + a separate head-wise gate + an SWA window,
2190 // none of which the dc launcher can express (see `mtp_step35_attn`). Host-len arm.
2191 // Advances BOTH the host len and the device counter itself (unlike the dc arm,
2192 // whose host-side mirror the caller does).
2193 (Mixer::Full(fa), Some(g)) => {
2194 self.mtp_step35_attn(e, fa, g, &a_norm, &pos_d, scratch)?
2195 }
2196 (Mixer::Full(fa), None) => {
2197 let out =
2198 self.mtp_full_attn_dc(e, fa, &a_norm, &pos_d, scratch, mtp.geom.as_ref())?;
2199 scratch.kv.len += 1;
2200 out
2201 }
2202 (Mixer::Linear(_), _) => {
2203 panic!("MTP block is full-attn in qwen35; linear MTP not supported")
2204 }
2205 (Mixer::Mla(_), _) => crate::hybrid::mla_forward_unimplemented(),
2206 };
2207
2208 // op 7: x1 = inpSA + attn_out
2209 let mut x1 = e.zeros(di)?;
2210 e.add(&inp_sa, &attn_out, &mut x1, di)?;
2211
2212 // op 8: z = RMSNorm(x1, post_attn_norm) (pre-FFN norm)
2213 let mut z = e.zeros(di)?;
2214 e.rms_norm(&x1, mtp.post_attn_norm.float_data(), &mut z, di, 1, eps)?;
2215
2216 // op 9: FFN (Dense or MoE) — same as the trunk decode FFN
2217 let ffn_out = match &mtp.ffn {
2218 crate::hybrid::Ffn::Dense {
2219 ffn_gate,
2220 ffn_up,
2221 ffn_down,
2222 } => {
2223 let n_ff = ffn_gate.out_features();
2224 let (gate, up) = if e.uses_q8_1_fast(ffn_gate) && e.uses_q8_1_fast(ffn_up) {
2225 let (zq, zd) = e.quantize_q8_1(&z, 1, di)?;
2226 (
2227 e.matmul_pre(ffn_gate, &zq, &zd, &z, 1)?,
2228 e.matmul_pre(ffn_up, &zq, &zd, &z, 1)?,
2229 )
2230 } else {
2231 (e.matmul(ffn_gate, &z, 1)?, e.matmul(ffn_up, &z, 1)?)
2232 };
2233 let mut act = e.zeros(n_ff)?;
2234 // step35: a DENSE FFN reads the per-layer SHEXP clamp (upstream's one `build_ffn`
2235 // serves the dense MLP and the shared expert off `swiglu_clamp_shexp` —
2236 // llama-graph.cpp:1751), resolved for the MTP block's OWN index. Every other arch
2237 // passes None, which is `ffn_act`'s dispatch verbatim.
2238 Self::ffn_act_lim(
2239 e,
2240 &self.cfg,
2241 &gate,
2242 &up,
2243 1.0,
2244 1.0,
2245 mtp.step35.as_ref().and_then(|s| s.clamp_shexp),
2246 &mut act,
2247 n_ff,
2248 )?;
2249 e.matmul(ffn_down, &act, 1)?
2250 }
2251 // MTP head is a distinct block — key its experts under a separate layer index (u16::MAX)
2252 // so they never alias trunk layer 0's cache keys.
2253 crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il(e, m, &z, 1, u16::MAX)?,
2254 };
2255
2256 // op 10: h_nextn = x1 + ffn_out (at di)
2257 let mut h_inner = e.zeros(di)?;
2258 e.add(&x1, &ffn_out, &mut h_inner, di)?;
2259
2260 // op 10.5 (student): up-project the inner hidden back to n_embd — training semantics:
2261 // the chain carrier AND the head input are out_up(h_inner) (pre-final-norm).
2262 let h_nextn = match mtp.geom.as_ref() {
2263 Some(g) => e.matmul(&g.out_up, &h_inner, 1)?,
2264 None => h_inner,
2265 };
2266
2267 // op 11: final = RMSNorm(h_nextn, shared_head_norm OR output_norm)
2268 let final_norm = mtp.shared_head_norm.as_ref().unwrap_or(&self.output_norm);
2269 let mut final_h = e.zeros(n_embd)?;
2270 e.rms_norm(
2271 &h_nextn,
2272 final_norm.float_data(),
2273 &mut final_h,
2274 n_embd,
2275 1,
2276 eps,
2277 )?;
2278
2279 // op 12: draft_logits = (shared_head_head OR output) @ final — stays ON DEVICE.
2280 let head = mtp.shared_head_head.as_ref().unwrap_or(&self.output);
2281 let mut logits = e.matmul(head, &final_h, 1)?;
2282 // op 12b (lane/draft-mask): grammar mask over the DRAFT vocab, applied here so the
2283 // caller's argmax / gumbel draw / p-min prob all read the grammar-legal row.
2284 if let Some((mask_d, mw)) = mask {
2285 let d_vocab = head.out_features();
2286 e.mask_logits_col(&mut logits, mask_d, 0, d_vocab, mw)?;
2287 }
2288 // Chain recurrence hand-over: pre-norm h_nextn (default) or post-norm final_h
2289 // (MEMRA_SPEC_HPOST — llama.cpp #24025's t_h_nextn is taken AFTER the head norm).
2290 Ok((logits, if spec_hpost() { final_h } else { h_nextn }))
2291 }
2292
2293 /// step35 MTP-block attention, T=1, on the scratch KV — the EAGER-ONLY twin of
2294 /// `mtp_full_attn_dc`. Three things force a separate arm rather than a geometry parameter on
2295 /// the dc path, and all three are properties of this arch's MTP block:
2296 ///
2297 /// 1. **The SWA window.** Block 45 is an SWA-type block (`sliding_window_pattern[45]=true`,
2298 /// window 512). Windowed decode in memra is a token-aligned VIEW OFFSET into the quantized
2299 /// cache (the gemma4 R6 / `step35_decode_attn` pattern: keys carry absolute rope and the
2300 /// mask is purely positional, so one query at `len-1` attending the last `win` rows IS the
2301 /// windowed result). `fa_decode_dc` takes the key count from a DEVICE counter and always
2302 /// starts at row 0 — it cannot express a nonzero offset, so a windowed dc arm would need a
2303 /// new kernel. That is deliberately not built here: see the CUDA-graph note below.
2304 /// 2. **Per-layer head count.** 96 q heads over 8 KV (GQA 12) at this block, vs the trunk's 64
2305 /// on its full-attn layers. The trunk cfg's `n_head` scalar is the MAX over layers, and the
2306 /// trunk ARTIFACT's per-layer arrays stop at index 44 — so the count must come from the
2307 /// resolved `Step35MtpGeom`, never from `cfg`.
2308 /// 3. **The separate head-wise gate.** `blk.45.attn_gate.weight [n_embd, 96]` produces one
2309 /// sigmoid scalar per head (broadcast over head_dim) — `attn_head_gate`, not the qwen35
2310 /// fused-into-wq `q_gate_split` form the dc arm handles.
2311 ///
2312 /// WHY EAGER-ONLY IS NOT A GAP TODAY: `graph_draft` requires `trunk_dense`, and this SKU's
2313 /// trunk is a 288-expert MoE, so the graph draft is already off for every step35 model — the
2314 /// eager chain IS the served path. `mtp_head_forward_cap` refuses step35 explicitly rather
2315 /// than silently capturing a window-less (wrong past `win` draft rows) graph.
2316 ///
2317 /// Unlike the dc arm this advances BOTH the host `kv.len` and the device counter, so the
2318 /// caller must not mirror.
2319 fn mtp_step35_attn(
2320 &self,
2321 e: &Engine,
2322 fa: &FullAttnLayer,
2323 g: &crate::hybrid::Step35MtpGeom,
2324 h: &CudaSlice<f32>,
2325 pos_d: &CudaSlice<i32>,
2326 scratch: &mut MtpScratch,
2327 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2328 let (nh, nkv, hd) = (g.n_head, g.n_head_kv, self.cfg.head_dim_k as usize);
2329 let eps = self.cfg.rms_eps;
2330 let scale = 1.0 / (hd as f32).sqrt(); // step35.cpp:255 kq_scale
2331 let n_embd = self.cfg.n_embd as usize;
2332 let gw = fa
2333 .attn_gate
2334 .as_ref()
2335 .ok_or("step35 MTP block is missing attn_gate.weight (head-wise attention gate)")?;
2336
2337 let (q0, k0, v0, gt) = if e.uses_q8_1_fast(&fa.wq)
2338 && e.uses_q8_1_fast(&fa.wk)
2339 && e.uses_q8_1_fast(&fa.wv)
2340 && e.uses_q8_1_fast(gw)
2341 {
2342 let (hq, hdq) = e.quantize_q8_1(h, 1, n_embd)?;
2343 let (a, b, c) = match e.matmul_q8_fused3(&fa.wq, &fa.wk, &fa.wv, &hq, &hdq)? {
2344 Some(t3) => t3,
2345 None => (
2346 e.matmul_pre(&fa.wq, &hq, &hdq, h, 1)?,
2347 e.matmul_pre(&fa.wk, &hq, &hdq, h, 1)?,
2348 e.matmul_pre(&fa.wv, &hq, &hdq, h, 1)?,
2349 ),
2350 };
2351 (a, b, c, e.matmul_pre(gw, &hq, &hdq, h, 1)?)
2352 } else {
2353 (
2354 e.matmul(&fa.wq, h, 1)?,
2355 e.matmul(&fa.wk, h, 1)?,
2356 e.matmul(&fa.wv, h, 1)?,
2357 e.matmul(gw, h, 1)?,
2358 )
2359 };
2360
2361 let mut q = e.uninit(nh * hd)?;
2362 e.rms_norm(&q0, fa.q_norm.float_data(), &mut q, hd, nh, eps)?;
2363 let mut k = e.uninit(nkv * hd)?;
2364 e.rms_norm(&k0, fa.k_norm.float_data(), &mut k, hd, nkv, eps)?;
2365 // `rope_freqs.weight` (llama3 factors) applies to the FULL-attn layers ONLY; SWA passes
2366 // null (llama-hparams / step35.cpp). Block 45 is SWA, so `ff` is None there — but read
2367 // the resolved flag, not the constant, so an all-full sibling stays correct.
2368 let ff = if g.swa {
2369 None
2370 } else {
2371 self.step35_aux.as_ref().and_then(|a| a.rope_freqs(e))
2372 };
2373 #[cfg(debug_assertions)]
2374 if let Some(ff) = ff {
2375 crate::debug_assert_tensor_stream_device(ff, &e.stream(), "mtp_step35_attn.rope_freqs");
2376 }
2377 e.rope_neox2(
2378 &mut q,
2379 &mut k,
2380 pos_d,
2381 hd,
2382 g.n_rot,
2383 nh,
2384 nkv,
2385 1,
2386 g.rope_base,
2387 1.0,
2388 ff,
2389 )?;
2390
2391 // Append at the HOST slot, then re-stamp the device counter: the eager chain has the
2392 // length on the host anyway, and the windowed view below needs it there to compute the
2393 // offset. The device counter is kept in lockstep so `mtp_kv_fill`'s `set_i32_one` and any
2394 // dc-family consumer of this scratch still agree.
2395 let kv = &mut scratch.kv;
2396 assert!(
2397 kv.len < scratch.cap,
2398 "step35 MTP scratch overflow ({} >= {})",
2399 kv.len,
2400 scratch.cap
2401 );
2402 let next_len = kv.len + 1;
2403 let (off, t_kv) = if g.swa && next_len > g.window {
2404 (next_len - g.window, g.window)
2405 } else {
2406 (0, next_len)
2407 };
2408 let write_row = e.prepare_kv_append(kv, off & !31usize, 1)?;
2409 e.append_kv_quantized(
2410 &k,
2411 &v0,
2412 &mut kv.k,
2413 &mut kv.v,
2414 write_row,
2415 kv.kv_dim_k,
2416 kv.kv_dim_v,
2417 kv.k_tok_bytes,
2418 kv.v_tok_bytes,
2419 false,
2420 )?;
2421 kv.len = next_len;
2422 e.set_i32_one(&mut kv.len_d, kv.len as i32)?;
2423 // SWA view offset (see note 1). The draft chain is short (k+2 rows), but the scratch is
2424 // PERSISTENT across rounds — `mtp_kv_fill` leaves one row per committed token behind, so
2425 // `kv.len` tracks absolute position and crosses 512 in any real generation. The window is
2426 // therefore live, not theoretical.
2427 let physical = kv.physical_rows(off, off + t_kv)?;
2428 let k_view = e.view_u8_range(
2429 &kv.k,
2430 physical.start * kv.k_tok_bytes,
2431 physical.end * kv.k_tok_bytes,
2432 );
2433 let v_view = e.view_u8_range(
2434 &kv.v,
2435 physical.start * kv.v_tok_bytes,
2436 physical.end * kv.v_tok_bytes,
2437 );
2438 let mut attn = e.uninit(nh * hd)?;
2439 e.fa_decode_kvmod(
2440 &q,
2441 &k_view,
2442 &v_view,
2443 &mut attn,
2444 hd,
2445 nh,
2446 nkv,
2447 t_kv,
2448 scale,
2449 kv.k_tok_bytes,
2450 kv.v_tok_bytes,
2451 false,
2452 )?;
2453
2454 let mut ag = e.uninit(nh * hd)?;
2455 e.attn_head_gate(&attn, >, &mut ag, None, hd, nh, 1)?;
2456 Ok(e.matmul(&fa.wo, &ag, 1)?)
2457 }
2458
2459 /// MTP-block full attention, T=1, on the scratch KV (BOTH draft paths — eager and graph):
2460 /// the scratch write slot and the attention bound come from `scratch.kv.len_d` (device i32[1])
2461 /// so the launch args are FIXED across draft steps — ONE captured graph serves the whole
2462 /// chain, and replays keep seeing KV growth through the device counter (no recapture).
2463 /// Geometry contract: n_splits is sized from `scratch.cap` (the persistent capacity); splits
2464 /// whose key range lies beyond the device t_kv exit empty and the shared combine skips them
2465 /// (fa_decode_dc bit-correct-for-any-t_kv<=bucket_max contract). The eager path uses the SAME
2466 /// launcher with the SAME bucket_max -> identical dispatch -> bit-identical draft tokens (the
2467 /// graph-vs-eager parity gate). Host len is NOT advanced here (graph contract); callers mirror.
2468 fn mtp_full_attn_dc(
2469 &self,
2470 e: &Engine,
2471 fa: &FullAttnLayer,
2472 h: &CudaSlice<f32>,
2473 pos_d: &CudaSlice<i32>,
2474 scratch: &mut MtpScratch,
2475 geom: Option<&crate::hybrid::DraftGeom>,
2476 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2477 let cfg = &self.cfg;
2478 let mtp_il = cfg.n_layer.saturating_sub(cfg.nextn_predict_layers);
2479 let geometry = cfg.full_attention_geometry_at(mtp_il);
2480 let n_head = geom.map(|g| g.n_head).unwrap_or(geometry.n_head as usize);
2481 let n_head_kv = geom
2482 .map(|g| g.n_head_kv)
2483 .unwrap_or(geometry.n_head_kv as usize);
2484 let head_dim = geometry.head_dim_k as usize;
2485 let eps = cfg.rms_eps;
2486 let scale = geometry.attention_scale();
2487 let n_embd = geom.map(|g| g.d_inner).unwrap_or(cfg.n_embd as usize);
2488 let bucket_max = scratch.cap; // < 96 guaranteed by the graph_draft eligibility gate
2489
2490 let (qf, mut k, v) =
2491 if e.uses_q8_1_fast(&fa.wq) && e.uses_q8_1_fast(&fa.wk) && e.uses_q8_1_fast(&fa.wv) {
2492 let (hq, hd) = e.quantize_q8_1(h, 1, n_embd)?;
2493 (
2494 e.matmul_pre(&fa.wq, &hq, &hd, h, 1)?,
2495 e.matmul_pre(&fa.wk, &hq, &hd, h, 1)?,
2496 e.matmul_pre(&fa.wv, &hq, &hd, h, 1)?,
2497 )
2498 } else {
2499 (
2500 e.matmul(&fa.wq, h, 1)?,
2501 e.matmul(&fa.wk, h, 1)?,
2502 e.matmul(&fa.wv, h, 1)?,
2503 )
2504 };
2505 // M3/Hy3 have no attention output gate — wq out is exactly q; skip the split.
2506 let gated = geometry.attention_gate == memra_gguf::config::AttentionGateKind::FusedQ;
2507 let (mut q, gate) = if gated {
2508 let mut q = e.zeros(n_head * head_dim)?;
2509 let mut gate = e.zeros(n_head * head_dim)?;
2510 e.q_gate_split(&qf, &mut q, &mut gate, head_dim, n_head, 1)?;
2511 (q, Some(gate))
2512 } else {
2513 (qf, None)
2514 };
2515
2516 let mut qn = e.zeros(n_head * head_dim)?;
2517 e.rms_norm(&q, fa.q_norm.float_data(), &mut qn, head_dim, n_head, eps)?;
2518 q = qn;
2519 let mut kn = e.zeros(n_head_kv * head_dim)?;
2520 e.rms_norm(
2521 &k,
2522 fa.k_norm.float_data(),
2523 &mut kn,
2524 head_dim,
2525 n_head_kv,
2526 eps,
2527 )?;
2528 k = kn;
2529 let rope_dims = geometry.n_rot as usize;
2530 e.rope_neox(
2531 &mut q,
2532 pos_d,
2533 head_dim,
2534 rope_dims,
2535 n_head,
2536 1,
2537 geometry.rope_base,
2538 1.0,
2539 )?;
2540 e.rope_neox(
2541 &mut k,
2542 pos_d,
2543 head_dim,
2544 rope_dims,
2545 n_head_kv,
2546 1,
2547 geometry.rope_base,
2548 1.0,
2549 )?;
2550
2551 let kv = &mut scratch.kv;
2552 // append at the DEVICE slot (kv.len_d == old len), then advance the counter in-graph.
2553 e.append_kv_quantized_dc(
2554 &k,
2555 &v,
2556 &mut kv.k,
2557 &mut kv.v,
2558 &kv.len_d,
2559 kv.kv_dim_k,
2560 kv.kv_dim_v,
2561 kv.k_tok_bytes,
2562 kv.v_tok_bytes,
2563 false,
2564 )?;
2565 e.inc_seqlen(&mut kv.len_d)?;
2566 // full-buffer views (any in-round t_kv stays in range on replay); the kernel bounds the
2567 // key range from the device counter.
2568 let k_view = e.view_u8(&kv.k, kv.k.len());
2569 let v_view = e.view_u8(&kv.v, kv.v.len());
2570 let (ktb, vtb) = (kv.k_tok_bytes, kv.v_tok_bytes);
2571 let mut attn = e.zeros(n_head * head_dim)?;
2572 e.fa_decode_dc(
2573 &q, &k_view, &v_view, &mut attn, head_dim, n_head, n_head_kv, &kv.len_d, bucket_max,
2574 scale, ktb, vtb, false,
2575 )?;
2576
2577 let attn_g = match &gate {
2578 Some(gate) => {
2579 let mut gsig = e.zeros(n_head * head_dim)?;
2580 e.sigmoid(gate, &mut gsig, n_head * head_dim)?;
2581 let mut ag = e.zeros(n_head * head_dim)?;
2582 e.mul(&attn, &gsig, &mut ag, n_head * head_dim)?;
2583 ag
2584 }
2585 None => attn,
2586 };
2587 Ok(e.matmul(&fa.wo, &attn_g, 1)?)
2588 }
2589
2590 /// PERSISTENT-DRAFT-KV fill (the reference engine's "mtp_update" analogue): compute the MTP
2591 /// block's K/V for `tokens` (committed tokens at positions pos0..pos0+T) from their EXACT
2592 /// trunk hiddens `h` ([T, n_embd] token-major, pre-output_norm) and append at slots pos0.. of
2593 /// the scratch KV. K/V-ONLY — ops A/1-5 plus the K-side of op 6 (wk/wv + k_norm + rope +
2594 /// quantized append); no wq/attention/FFN/lm_head, so per-token cost ~= eh_proj + wk/wv (a
2595 /// small fraction of one trunk layer), T-batched. Rope follows the chain convention
2596 /// rope(token@p) = p+1. Runs at round boundaries OUTSIDE the captured graph in BOTH draft
2597 /// modes -> draft parity by construction. Caller must have scratch.kv.len == pos0.
2598 #[allow(clippy::too_many_arguments)]
2599 fn mtp_kv_fill(
2600 &self,
2601 e: &Engine,
2602 mtp: &MtpHead,
2603 tokens: &[u32],
2604 h: &CudaSlice<f32>,
2605 pos0: usize,
2606 scratch: &mut MtpScratch,
2607 embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
2608 ) -> Result<(), Box<dyn std::error::Error>> {
2609 let cfg = &self.cfg;
2610 let n_embd = cfg.n_embd as usize;
2611 let eps = cfg.rms_eps;
2612 let t = tokens.len();
2613 assert_eq!(scratch.kv.len, pos0, "mtp_kv_fill: append slot mismatch");
2614 assert!(pos0 + t <= scratch.cap, "mtp_kv_fill: scratch overflow");
2615 let Mixer::Full(fa) = &mtp.mixer else {
2616 panic!("MTP block is full-attn in qwen35; linear MTP not supported")
2617 };
2618 let pos_vec: Vec<i32> = (0..t).map(|i| (pos0 + i + 1) as i32).collect();
2619 let pos_d = e.htod_i32(&pos_vec)?;
2620
2621 // ops A/1/2: embed + the two input norms, T-wide.
2622 let e_emb = match embd_dev {
2623 Some((g, qt, rb)) => e.embed_gather_device_t(g, tokens, n_embd, qt, rb)?,
2624 None => e.htod(&self.embd.gather(n_embd, tokens))?,
2625 };
2626 let mut e_norm = e.zeros(t * n_embd)?;
2627 e.rms_norm(&e_emb, mtp.enorm.float_data(), &mut e_norm, n_embd, t, eps)?;
2628 let mut h_norm = e.zeros(t * n_embd)?;
2629 e.rms_norm(h, mtp.hnorm.float_data(), &mut h_norm, n_embd, t, eps)?;
2630
2631 // op 3: per-row [e_norm ; h_norm] concat, token-major [T, 2*n_embd].
2632 let mut concat = e.zeros(t * 2 * n_embd)?;
2633 for i in 0..t {
2634 e.copy_view_into(
2635 &mut concat,
2636 i * 2 * n_embd,
2637 &e_norm.slice(i * n_embd..(i + 1) * n_embd),
2638 n_embd,
2639 )?;
2640 e.copy_view_into(
2641 &mut concat,
2642 i * 2 * n_embd + n_embd,
2643 &h_norm.slice(i * n_embd..(i + 1) * n_embd),
2644 n_embd,
2645 )?;
2646 }
2647
2648 // ops 4/5: eh_proj + attn_norm, T-wide (at the student inner width when geom is set).
2649 let di = mtp.geom.as_ref().map(|g| g.d_inner).unwrap_or(n_embd);
2650 let inp_sa = e.matmul(&mtp.eh_proj, &concat, t)?;
2651 let mut a_norm = e.zeros(t * di)?;
2652 e.rms_norm(&inp_sa, mtp.attn_norm.float_data(), &mut a_norm, di, t, eps)?;
2653
2654 // op 6 (K/V half): wk/wv + k_norm + rope + per-row quantized append. No wq/attention —
2655 // the fill only has to leave correct K/V rows behind for later chains to attend over.
2656 let n_head_kv = mtp
2657 .geom
2658 .as_ref()
2659 .map(|g| g.n_head_kv)
2660 .or(mtp.step35.as_ref().map(|s| s.n_head_kv))
2661 .unwrap_or_else(|| {
2662 let mtp_il = cfg.n_layer.saturating_sub(cfg.nextn_predict_layers);
2663 cfg.full_attention_geometry_at(mtp_il).n_head_kv as usize
2664 });
2665 let mtp_il = cfg.n_layer.saturating_sub(cfg.nextn_predict_layers);
2666 let geometry = cfg.full_attention_geometry_at(mtp_il);
2667 let head_dim = geometry.head_dim_k as usize;
2668 let mut k = e.matmul(&fa.wk, &a_norm, t)?;
2669 let v = e.matmul(&fa.wv, &a_norm, t)?;
2670 let mut kn = e.zeros(t * n_head_kv * head_dim)?;
2671 e.rms_norm(
2672 &k,
2673 fa.k_norm.float_data(),
2674 &mut kn,
2675 head_dim,
2676 n_head_kv * t,
2677 eps,
2678 )?;
2679 k = kn;
2680 // step35: rotary width AND base are per-layer, and the MTP block's values come from the
2681 // resolved `Step35MtpGeom` — NOT from `cfg.rope_dim_count`/`cfg.rope_freq_base`, which
2682 // carry the arch defaults (128 / 5e6, i.e. the FULL-attn layers' base). Getting this wrong
2683 // writes K rows the attention arm then re-derives at a different theta: correct-looking
2684 // output with dead acceptance, invisible to the exactness gates.
2685 let (rope_dims, rope_base, ff) = match mtp.step35.as_ref() {
2686 Some(s) => (
2687 s.n_rot,
2688 s.rope_base,
2689 if s.swa {
2690 None
2691 } else {
2692 self.step35_aux.as_ref().and_then(|a| a.rope_freqs(e))
2693 },
2694 ),
2695 None => (geometry.n_rot as usize, geometry.rope_base, None),
2696 };
2697 #[cfg(debug_assertions)]
2698 if let Some(ff) = ff {
2699 crate::debug_assert_tensor_stream_device(ff, &e.stream(), "mtp_kv_fill.rope_freqs");
2700 }
2701 match ff {
2702 Some(f) => e.rope_neox_ff(
2703 &mut k, &pos_d, head_dim, rope_dims, n_head_kv, t, rope_base, 1.0, f,
2704 )?,
2705 None => e.rope_neox(
2706 &mut k, &pos_d, head_dim, rope_dims, n_head_kv, t, rope_base, 1.0,
2707 )?,
2708 }
2709
2710 let kv = &mut scratch.kv;
2711 // Match the trunk prime contract: a chunk may need the aligned window immediately before
2712 // its first row, so preserve that prefix when the physical tail rebases at wrap.
2713 let retain_from = kv
2714 .ring
2715 .as_ref()
2716 .map(|ring| pos0.saturating_sub(ring.window() - 1) & !31usize)
2717 .unwrap_or(0);
2718 let write_row = e.prepare_kv_append(kv, retain_from, t)?;
2719 for i in 0..t {
2720 let k_row = k.slice(i * kv.kv_dim_k..(i + 1) * kv.kv_dim_k);
2721 let v_row = v.slice(i * kv.kv_dim_v..(i + 1) * kv.kv_dim_v);
2722 e.append_kv_quantized_view(
2723 &k_row,
2724 &v_row,
2725 &mut kv.k,
2726 &mut kv.v,
2727 write_row + i,
2728 kv.kv_dim_k,
2729 kv.kv_dim_v,
2730 kv.k_tok_bytes,
2731 kv.v_tok_bytes,
2732 false,
2733 )?;
2734 }
2735 kv.len = pos0 + t;
2736 e.set_i32_one(&mut kv.len_d, kv.len as i32)?;
2737 Ok(())
2738 }
2739
2740 /// CAPTURE body for the GRAPH DRAFT (stage 2 of graph-grade spec): ONE MTP head forward with
2741 /// every varying input device-resident —
2742 /// - token id from the persistent `tok_d` (the previous replay's in-graph argmax wrote it,
2743 /// so the chain feeds itself; the host reads the same 4 bytes for the draft list),
2744 /// - h_seed from the persistent `h_seed_d` (h_nextn is copied BACK into it at the end),
2745 /// - rope pos from the persistent `pos_d` counter (inc'd in-graph),
2746 /// - scratch KV slot/bound from `scratch.kv.len_d` (see mtp_full_attn_dc).
2747 /// The p-min confidence lands in the persistent `p_d` iff `with_prob` (env is fixed per run).
2748 /// Same kernels, same dispatch as the eager mtp_head_forward_dev chain -> same draft tokens
2749 /// (exactness never depends on drafts — the verify arbitrates — but acceptance parity does).
2750 /// `with_head=false` captures the HEAD-LESS twin for the pseudo-seed replay (2026-07-03):
2751 /// the pseudo pass only needs h_nextn (op 10) + the scratch append — the lm_head read
2752 /// (~1.06ms q6_K on the 9B), argmax and prob are dead weight there. h_nextn's inputs are
2753 /// untouched, so the seed value is identical; round-start resets overwrite tok_d/p_d anyway.
2754 /// `sampled_cap` = Some((ctr_d, perturb_d, q_out_d, seed, temp)) captures the SAMPLED twin
2755 /// (step 3 of the sampled-spec arc): head logits are retained in the persistent `q_out_d`
2756 /// (host D2Ds them to the round's q slot after each replay), the DEVICE event counter is
2757 /// bumped in-graph, and the argmax reads GUMBEL-PERTURBED logits — one categorical draw per
2758 /// replay, bit-identical to the eager arm's gumbel_perturb at the same (seed, sctr, temp).
2759 /// seed/temp are capture-time constants (fixed per generate call, like p_min).
2760 #[allow(clippy::too_many_arguments)]
2761 fn mtp_head_forward_cap(
2762 &self,
2763 e: &Engine,
2764 mtp: &MtpHead,
2765 tok_d: &mut CudaSlice<u32>,
2766 pos_d: &mut CudaSlice<i32>,
2767 h_seed_d: &mut CudaSlice<f32>,
2768 p_d: &mut CudaSlice<f32>,
2769 scratch: &mut MtpScratch,
2770 with_prob: bool,
2771 with_head: bool,
2772 embd_gpu: &CudaSlice<u8>,
2773 embd_qt: i32,
2774 embd_rb: usize,
2775 d_vocab: usize,
2776 sampled_cap: Option<(
2777 &mut CudaSlice<u32>,
2778 &mut CudaSlice<f32>,
2779 &mut CudaSlice<f32>,
2780 u64,
2781 f32,
2782 )>,
2783 stream_pack: Option<(&mut CudaSlice<u32>, usize, Option<&CudaSlice<u32>>)>,
2784 // DRAFT-SIDE GRAMMAR MASK (lane/draft-mask): (packed draft-vocab allowed-set buffer,
2785 // word count). Captured as ONE mask_logits_f32 node between the head matmul and the
2786 // in-graph argmax; the buffer address is baked, its CONTENTS are re-uploaded by the
2787 // host before every replay (the decode.rs graph-mask pattern). All-ones contents = a
2788 // no-op ban, so a position the grammar cannot constrain costs one pass over the row.
2789 mask_cap: Option<(&CudaSlice<u32>, usize)>,
2790 ) -> Result<(), Box<dyn std::error::Error>> {
2791 let cfg = &self.cfg;
2792 let n_embd = cfg.n_embd as usize;
2793 // step35 REFUSAL (deliberate, named): the graph body's attention is `mtp_full_attn_dc`,
2794 // whose device-counter key bound always starts at row 0 — it cannot express this block's
2795 // SWA view offset, so a captured chain would silently attend OUTSIDE the window once the
2796 // persistent scratch passes 512 rows. Nothing is lost today: `graph_draft` also requires
2797 // `trunk_dense`, and step35's trunk is a 288-expert MoE, so the eager chain
2798 // (`mtp_head_forward_dev` -> `mtp_step35_attn`) is the served path. Returning Err (not a
2799 // panic) is what the two capture sites and the round-stream capture already handle by
2800 // degrading to eager / stream-off.
2801 if mtp.step35.is_some() {
2802 return Err(
2803 "step35 has no captured draft chain (fa_decode_dc cannot express the MTP \
2804 block's SWA view offset; same root cause as the dc decode refusal) — the \
2805 eager draft chain serves this arch"
2806 .into(),
2807 );
2808 }
2809 // student inner width (see mtp_head_forward_dev) — interface dims stay n_embd.
2810 let di = mtp.geom.as_ref().map(|g| g.d_inner).unwrap_or(n_embd);
2811 let eps = cfg.rms_eps;
2812 let e_emb = e.embed_gather_device(embd_gpu, tok_d, n_embd, embd_qt, embd_rb)?;
2813 let mut e_norm = e.zeros(n_embd)?;
2814 e.rms_norm(&e_emb, mtp.enorm.float_data(), &mut e_norm, n_embd, 1, eps)?;
2815 let mut h_norm = e.zeros(n_embd)?;
2816 e.rms_norm(
2817 &*h_seed_d,
2818 mtp.hnorm.float_data(),
2819 &mut h_norm,
2820 n_embd,
2821 1,
2822 eps,
2823 )?;
2824 let mut concat = e.zeros(2 * n_embd)?;
2825 e.copy_into(&mut concat, 0, &e_norm, n_embd)?;
2826 e.copy_into(&mut concat, n_embd, &h_norm, n_embd)?;
2827 let inp_sa = e.matmul(&mtp.eh_proj, &concat, 1)?;
2828 let mut a_norm = e.zeros(di)?;
2829 e.rms_norm(&inp_sa, mtp.attn_norm.float_data(), &mut a_norm, di, 1, eps)?;
2830 let attn_out = match &mtp.mixer {
2831 Mixer::Full(fa) => {
2832 self.mtp_full_attn_dc(e, fa, &a_norm, pos_d, scratch, mtp.geom.as_ref())?
2833 }
2834 Mixer::Linear(_) => {
2835 panic!("MTP block is full-attn in qwen35; linear MTP not supported")
2836 }
2837 Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
2838 };
2839 let mut x1 = e.zeros(di)?;
2840 e.add(&inp_sa, &attn_out, &mut x1, di)?;
2841 let mut z = e.zeros(di)?;
2842 e.rms_norm(&x1, mtp.post_attn_norm.float_data(), &mut z, di, 1, eps)?;
2843 let ffn_out = match &mtp.ffn {
2844 crate::hybrid::Ffn::Dense {
2845 ffn_gate,
2846 ffn_up,
2847 ffn_down,
2848 } => {
2849 let n_ff = ffn_gate.out_features();
2850 let (gate, up) = if e.uses_q8_1_fast(ffn_gate) && e.uses_q8_1_fast(ffn_up) {
2851 let (zq, zd) = e.quantize_q8_1(&z, 1, di)?;
2852 (
2853 e.matmul_pre(ffn_gate, &zq, &zd, &z, 1)?,
2854 e.matmul_pre(ffn_up, &zq, &zd, &z, 1)?,
2855 )
2856 } else {
2857 (e.matmul(ffn_gate, &z, 1)?, e.matmul(ffn_up, &z, 1)?)
2858 };
2859 let mut act = e.zeros(n_ff)?;
2860 Self::ffn_act(e, &self.cfg, &gate, &up, &mut act, n_ff)?;
2861 e.matmul(ffn_down, &act, 1)?
2862 }
2863 // ROUND-STREAM: the 35B NextN block carries a MoE FFN. With RESIDENT experts the
2864 // dev path is pure device launches (device top-k + rows kernels, ZERO-DtoH by
2865 // design) — capture-legal. Non-resident (SLRU-lock) stays rejected: the capture
2866 // error arm degrades the caller to eager/stream-off.
2867 crate::hybrid::Ffn::Moe(m) if m.dev_exps.is_some() => {
2868 self.moe_ffn_il(e, m, &z, 1, u16::MAX)?
2869 }
2870 crate::hybrid::Ffn::Moe(_) => {
2871 return Err("graph draft requires a Dense (or resident-MoE) MTP FFN".into());
2872 }
2873 };
2874 let mut h_inner = e.zeros(di)?;
2875 e.add(&x1, &ffn_out, &mut h_inner, di)?;
2876 // student: up-project back to n_embd (carrier + head input; see mtp_head_forward_dev).
2877 let h_nextn = match mtp.geom.as_ref() {
2878 Some(g) => e.matmul(&g.out_up, &h_inner, 1)?,
2879 None => h_inner,
2880 };
2881 // MEMRA_SPEC_HPOST needs final_h even head-less (it IS the next seed under that convention).
2882 let final_h = if with_head || spec_hpost() {
2883 let final_norm = mtp.shared_head_norm.as_ref().unwrap_or(&self.output_norm);
2884 let mut fh = e.zeros(n_embd)?;
2885 e.rms_norm(&h_nextn, final_norm.float_data(), &mut fh, n_embd, 1, eps)?;
2886 Some(fh)
2887 } else {
2888 None
2889 };
2890 if with_head {
2891 let head = mtp.shared_head_head.as_ref().unwrap_or(&self.output);
2892 let mut logits = e.matmul(head, final_h.as_ref().unwrap(), 1)?;
2893 // DRAFT-SIDE GRAMMAR MASK: ban the grammar-illegal draft ids IN the captured chain,
2894 // before the argmax — proposals become legal by construction. Contents-only
2895 // per-replay upload keeps the capture valid.
2896 if let Some((mask_d, mw)) = mask_cap {
2897 e.mask_logits_col(&mut logits, mask_d, 0, d_vocab, mw)?;
2898 }
2899 if let Some((ctr_d, perturb_d, q_out_d, seed, temp)) = sampled_cap {
2900 // SAMPLED chain: retain q (raw head logits -> persistent q_out_d; the matmul's
2901 // own buffer is pool-recycled after the capture body returns, so it can't be the
2902 // retention target), bump the device event counter, gumbel-perturb reading it,
2903 // and argmax the PERTURBED logits into tok_d — the in-graph categorical draw.
2904 e.copy_into(q_out_d, 0, &logits, d_vocab)?;
2905 e.sctr_inc(ctr_d)?;
2906 e.gumbel_perturb_ctr(&logits, perturb_d, d_vocab, seed, ctr_d, temp)?;
2907 e.argmax_token_device_into(perturb_d, tok_d, d_vocab)?;
2908 // p-min prob = the head's RAW softmax confidence in the SAMPLED pick — same
2909 // semantics as the eager sampled arm's prob_of_token_device(dl_d, tok_d).
2910 if with_prob {
2911 e.prob_of_token_device_into(&logits, tok_d, p_d, d_vocab)?;
2912 }
2913 } else {
2914 // draft token -> persistent tok_d (next replay's embed reads it; host reads the 4 bytes).
2915 e.argmax_token_device_into(&logits, tok_d, d_vocab)?;
2916 // p-min under a draft mask reads the MASKED row: confidence relative to the
2917 // grammar-LEGAL alternatives (illegal ids leave the softmax denominator), which
2918 // is the right semantics for "does the drafter know what comes next here" and
2919 // the same row the pick came from. Draft-quality only — verify arbitrates.
2920 if with_prob {
2921 e.prob_of_token_device_into(&logits, tok_d, p_d, d_vocab)?;
2922 }
2923 }
2924 }
2925 // ROUND-STREAM K-chain: pack (tok, p) into slot j, then remap tok through d2t so the
2926 // NEXT chained body's embed reads the TARGET id — zero host involvement per step.
2927 if let Some((out, slot, d2t)) = stream_pack {
2928 e.pack_tok_p(tok_d, p_d, out, slot)?;
2929 if let Some(map) = d2t {
2930 e.tok_map_u32(tok_d, map)?;
2931 }
2932 }
2933 // Next draft step's h_seed: pre-norm h_nextn (default) or post-norm final_h (HPOST).
2934 if spec_hpost() {
2935 e.copy_into(h_seed_d, 0, final_h.as_ref().unwrap(), n_embd)?;
2936 } else {
2937 e.copy_into(h_seed_d, 0, &h_nextn, n_embd)?;
2938 }
2939 // advance the draft rope position in-graph.
2940 e.inc_seqlen(pos_d)?;
2941 Ok(())
2942 }
2943
2944 /// Batched target verify forward over `tokens` at positions `pos0..pos0+T` (§D.3, T=K+1).
2945 /// Returns ALL T logit columns (host f32, [T*n_vocab]); appends T cols to every full-attn KV
2946 /// and advances every linear-attn recur state by T steps (the recur steps are SEQUENTIAL T=1).
2947 /// Advances `cache.pos` by T.
2948 pub fn decode_step_t(
2949 &self,
2950 e: &Engine,
2951 tokens: &[u32],
2952 pos0: usize,
2953 cache: &mut Cache,
2954 ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
2955 if self.is_gemma4_e4b() {
2956 return Ok(self.gemma4_e4b_decode_step_t_h(e, tokens, pos0, cache)?.0);
2957 }
2958 if self.cfg.gemma4.is_some() {
2959 return self.gemma4_decode_step_t(e, tokens, pos0, cache);
2960 }
2961 Ok(self.decode_step_t_h(e, tokens, pos0, cache)?.0)
2962 }
2963
2964 /// Like `decode_step_t` but ALSO returns the LAST column's pre-output_norm hidden (h_seed for
2965 /// the next draft round). This lets partial-accept replay run as ONE batched T=(n_acc+1) forward
2966 /// (single weight read) instead of n_acc+1 separate T=1 decode_steps (n_acc+1 weight reads).
2967 /// At batch=1 decode is bandwidth-bound, so batching the replay is THE MTP profitability lever.
2968 pub fn decode_step_t_h(
2969 &self,
2970 e: &Engine,
2971 tokens: &[u32],
2972 pos0: usize,
2973 cache: &mut Cache,
2974 ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
2975 self.decode_step_t_h_emb(e, tokens, pos0, cache, None)
2976 }
2977
2978 /// Like `decode_step_t_h` with an optional RESIDENT embed table (spec hot loop): device
2979 /// gather instead of host dequant + [T, n_embd] f32 htod. Bit-identical rows.
2980 pub fn decode_step_t_h_emb(
2981 &self,
2982 e: &Engine,
2983 tokens: &[u32],
2984 pos0: usize,
2985 cache: &mut Cache,
2986 embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
2987 ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
2988 let (logits_d, h_seed) = self.decode_step_t_h_emb_dev(e, tokens, pos0, cache, embd_dev)?;
2989 Ok((e.dtoh(&logits_d)?, h_seed))
2990 }
2991
2992 /// DEVICE-LOGITS verify forward (spec device-argmax lever): identical kernel chain to
2993 /// `decode_step_t_h_emb` but returns the [T, n_vocab] logits ON DEVICE — the accept walk
2994 /// argmaxes each column on-device and reads back ONE [T] u32 instead of dtoh'ing the full
2995 /// T x n_vocab f32 block (~1-4 MB + T host argmaxes, every round). Kernel dispatch is
2996 /// UNCHANGED (same decode-exact kernels); only the post-logits transfer moves.
2997 pub fn decode_step_t_h_emb_dev(
2998 &self,
2999 e: &Engine,
3000 tokens: &[u32],
3001 pos0: usize,
3002 cache: &mut Cache,
3003 embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
3004 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
3005 let n_embd = self.cfg.n_embd as usize;
3006 let t = tokens.len();
3007 let (logits, x) = self.decode_step_t_core(e, tokens, pos0, cache, embd_dev, None)?;
3008 // h_seed for the next round = LAST column's pre-output_norm hidden ([n_embd]).
3009 let mut hs = vbuf(e, n_embd)?; // fully written by copy_view_into below
3010 e.copy_view_into(&mut hs, 0, &x.slice((t - 1) * n_embd..t * n_embd), n_embd)?;
3011 Ok((logits, hs))
3012 }
3013
3014 /// CORE verify forward: the `decode_step_t_h_emb_dev` kernel chain, returning the FULL
3015 /// pre-output_norm hidden stack x ([T, n_embd], any column extractable) and optionally
3016 /// filling a `VerifyCkpt` (retained per-layer state-rebuild inputs) for the REPLAY-FREE
3017 /// partial accept. `ckpt: None` => byte-for-byte the old behavior (the ckpt writes are pure
3018 /// retains/copies — they never change what any kernel computes).
3019 fn decode_step_t_core(
3020 &self,
3021 e: &Engine,
3022 tokens: &[u32],
3023 pos0: usize,
3024 cache: &mut Cache,
3025 embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
3026 mut ckpt: Option<&mut VerifyCkpt>,
3027 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
3028 self.decode_step_t_core_stream(e, tokens, pos0, cache, embd_dev, ckpt.take(), None, None)
3029 }
3030
3031 /// Increment-0 two-session PP seam: release the peer after this lane's stage-0 boundary TX.
3032 /// The two independent sessions keep their own cache/checkpoint state; only issue order moves.
3033 fn decode_step_t_core_pipelined(
3034 &self,
3035 e: &Engine,
3036 tokens: &[u32],
3037 pos0: usize,
3038 cache: &mut Cache,
3039 embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
3040 mut ckpt: Option<&mut VerifyCkpt>,
3041 pipe: &SpecPipeLane,
3042 round: usize,
3043 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
3044 let fence = crate::pp::pp_cuts(self.layers.len())
3045 .ok_or("two-session speculative pipeline requires a PP stage cut")?;
3046 if crate::pp::pp2_streams_off() || !crate::pp::spec_pp_on() {
3047 return Err("two-session speculative pipeline requires the PP verify split".into());
3048 }
3049 let interval_fence = pipe.stage0_begin(round)?;
3050 let ticket = self.verify_stage0_issue(
3051 e,
3052 tokens,
3053 pos0,
3054 cache,
3055 embd_dev,
3056 ckpt.as_deref_mut(),
3057 None,
3058 &fence,
3059 Some(interval_fence),
3060 pipe.trace(round),
3061 )?;
3062 pipe.stage0_end(round);
3063 pipe.stage1_begin(round)?;
3064 let result = self.verify_stage1_finish(e, ticket, cache, ckpt, None, &fence, true)?;
3065 pipe.verify_end(round);
3066 Ok(result)
3067 }
3068
3069 /// ROUND-STREAM stage (c) 4: `stream` = (device verify tokens [t], device pos counter) —
3070 /// when Some, rope positions come from pos_iota over the counter, the embed gathers the
3071 /// device tokens, and full_attn_verify routes appends/FA through the _dc twins reading the
3072 /// SAME counter (every layer's kvl.len == cache.pos, one counter drives all three). The
3073 /// host `tokens`/`pos0` args still size buffers (t is FIXED K+1 in stream mode).
3074 #[allow(clippy::too_many_arguments)]
3075 fn decode_step_t_core_stream(
3076 &self,
3077 e: &Engine,
3078 tokens: &[u32],
3079 pos0: usize,
3080 cache: &mut Cache,
3081 embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
3082 mut ckpt: Option<&mut VerifyCkpt>,
3083 stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
3084 pp_pipe: Option<bool>,
3085 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
3086 // PP DOOR (lane/pp2-spec 2026-08-06): the verify trunk now takes its OWN stage split,
3087 // exactly as the eager and batched steps do. This is the single funnel every verify
3088 // forward reaches (decode_step_t / _h / _h_emb / _h_emb_dev / _core all land here), so
3089 // wiring it here wires the whole spec surface — the draft/accept/commit machinery above
3090 // is untouched.
3091 //
3092 // History: pp2-hardening (2026-08-06) made this funnel FAIL CLOSED, because its trunk
3093 // walk was unsplit on one stream and a sharded cross-device placement peer-read every
3094 // remote layer's weights on every spec round (measured 13.9-28x on the batched twin).
3095 // The refusal below survives to cover the residue — MEMRA_SPEC_PP=0, MEMRA_PP_STREAMS=0,
3096 // or a placement whose PpNRt fails to build — so a config that would still walk the
3097 // whole trunk on one stream refuses instead of regressing 28x.
3098 if let Some(fence) = crate::pp::pp_cuts(self.layers.len()) {
3099 if !crate::pp::pp2_streams_off() && crate::pp::spec_pp_on() {
3100 return self.decode_step_t_core_ppn(
3101 e,
3102 tokens,
3103 pos0,
3104 cache,
3105 embd_dev,
3106 ckpt.take(),
3107 stream,
3108 &fence,
3109 pp_pipe,
3110 );
3111 }
3112 }
3113 crate::pp::refuse_unsplit_if_remote(
3114 "decode_step_t (spec verify)",
3115 "drop MEMRA_SPEC_PP=0 / MEMRA_PP_STREAMS=0 so the verify trunk takes its OWN stage \
3116 split (decode_step_t_core_ppn); or run spec on one device",
3117 )?;
3118 let cfg = &self.cfg;
3119 let n_embd = cfg.n_embd as usize;
3120 let eps = cfg.rms_eps;
3121 let t = tokens.len();
3122 let pos_d = match stream {
3123 Some((_, ctr)) => {
3124 let mut p = e.alloc_uninit::<i32>(t)?;
3125 e.pos_iota(ctr, &mut p, t)?;
3126 p
3127 }
3128 None => {
3129 let pos_vec: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
3130 e.htod_i32(&pos_vec)?
3131 }
3132 };
3133
3134 // embed T tokens -> [T, n_embd] token-major (device gather on the spec hot loop)
3135 let x = match (stream, embd_dev) {
3136 (Some((vtok, _)), Some((g, qt, rb))) => {
3137 e.embed_gather_device_td(g, vtok, t, n_embd, qt, rb)?
3138 }
3139 (None, Some((g, qt, rb))) => e.embed_gather_device_t(g, tokens, n_embd, qt, rb)?,
3140 _ => e.htod(&self.embd.gather(n_embd, tokens))?,
3141 };
3142
3143 // TRUNK WALK: layers [0, n_layers) through the SAME range-scoped subgraph the PP-N
3144 // stage split calls per stage (`verify_layers`) — one code path, so the split cannot
3145 // drift from the unsplit dispatch mirroring. lane/pp2-spec 2026-08-06.
3146 let x = self.verify_layers(
3147 e,
3148 x,
3149 0,
3150 self.layers.len(),
3151 &pos_d,
3152 pos0,
3153 t,
3154 cache,
3155 ckpt.take(),
3156 stream,
3157 )?;
3158
3159 let mut hn = vbuf(e, t * n_embd)?;
3160 let serving_head = self.cfg.step35.is_some() || self.qwen35_serving_class();
3161 let logits = if serving_head {
3162 // Step35 and the qwen35 family (MoE 2026-08-14 AM, dense-hybrid same day PM — the
3163 // Q3.8 bring-up reproduced the identical near-tie class on dense: eager-class verify
3164 // vs batched-class live serving, ULP drift amplified through the GDN recurrence)
3165 // serve one batched numeric class at every live width, including B=1. Keep the
3166 // verify head in that same class; other generic families retain the decode-exact
3167 // head that their run-spec contract pins.
3168 e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
3169 e.matmul(&self.output, &hn, t)?
3170 } else {
3171 e.rms_norm_decode(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
3172 e.matmul_decode_exact(&self.output, &hn, t)?
3173 };
3174 // stream: the device pos counter owns position; host mirror reconciles at drain.
3175 if stream.is_none() {
3176 cache.pos += t;
3177 }
3178 // Hidden stack for seeds/refresh-fills: pre-norm x (default) or post-norm hn (HPOST).
3179 Ok((logits, if spec_hpost() { hn } else { x }))
3180 }
3181
3182 /// THE VERIFY TRUNK OVER PP-N (lane/pp2-spec 2026-08-06): `decode_step_t_core_stream`'s walk
3183 /// as N stage subgraphs, each on its own engine/stream (and, under `MEMRA_PP_DEVICES`, its own
3184 /// device), with a `[T, n_embd]` boundary transfer between them. T = K+1 (the verify batch),
3185 /// so this is the batched-boundary shape the pp2-batch lane's grow-only slots already handle
3186 /// (`tx(b, x, t*n_embd)`; the slot grows to the high-water T and the transport moves exactly
3187 /// the payload).
3188 ///
3189 /// Structure is `decode_step_batch_ppn`'s, which is `decode_step_h_ppn`'s. FOUR THINGS ARE
3190 /// PER-STAGE and each for a measured reason (see `decode_step_batch_ppn`'s header for the
3191 /// receipts):
3192 ///
3193 /// 1. THE ENGINE (`rt.engine(s, e)`) — `Engine` owns lazily-grown stable-pointer scratch
3194 /// (`fa_part_pool`, `fa_vf16_scratch`, `argmax_partials`) that is single-stream-safe BY
3195 /// DESIGN. Two stage streams through one Engine is the 2026-08-02 shared-scratch race
3196 /// (35% flake, nondeterministic all-logits divergence). `PpNRt::build` gives every stage
3197 /// s>0 its own Engine even on the primary device; honouring it here is what scopes the
3198 /// pools. The verify path allocates MORE of that scratch than eager decode does (FA at
3199 /// m=T, and the per-layer `GdnStash` retains), so this is load-bearing, not inherited.
3200 ///
3201 /// 2. `pos_d` — each stage uploads its OWN copy of the T rope positions on ITS stream, so the
3202 /// buffer is allocated, consumed and freed on one stream. In `stream` mode that means each
3203 /// stage runs its own `pos_iota` over the SHARED device counter (`pos_ctr`): the counter is
3204 /// read-only during the forward (the round's `inc`/`copy_add` happen outside it), so every
3205 /// stage derives the identical iota, and each stage's own output buffer is stream-local.
3206 ///
3207 /// 3. THE EMBED lives with stage 0 (`self.embd` / `embd_gpu` are host/primary-side; the
3208 /// sharded loader leaves the table with stage 0 by construction).
3209 ///
3210 /// 4. THE HEAD (`output_norm` + `output`) runs on the LAST stage — the sharded loader uploaded
3211 /// both through that stage's engine (`hybrid.rs`: `e_head = layer_engine(e, n_trunk,
3212 /// n_trunk-1)`), so reading them anywhere else is a peer read of the biggest tensor in the
3213 /// model, every round.
3214 ///
3215 /// WHAT STAYS ON THE PRIMARY, deliberately: the returned logits and hidden stack `x`. Both are
3216 /// last-stage-allocated device buffers, and every consumer (the device argmax walk, the accept
3217 /// kernels, `spec_seed_gather`, the ckpt rebuild in `commit_verified_prefix`) reads them
3218 /// through the primary context by UVA — the same read the batched serving epilogue's
3219 /// `last_logits_dev` park does. Those consumers are per-round O(T x n_vocab) and O(n_embd),
3220 /// not per-layer, so they are not the 28x class; splitting them is a separate lane.
3221 ///
3222 /// The MTP HEAD (draft side) is NOT split: it is one block, it lives wherever the loader put
3223 /// it (`load_mtp` uses the primary engine), and it is ~1-2 GB against the trunk's tens. Draft
3224 /// placement is measured, not assumed — see `research/pp2-spec-20260806`.
3225 ///
3226 /// EXACTNESS: PP-N adds ZERO deviation. Each stage runs the SAME kernels on the SAME bytes in
3227 /// the same order via the SAME `verify_layers` the unsplit body calls; the only change is
3228 /// where the residual is materialized, and the boundary is a straight f32 copy (dtod
3229 /// same-device / `cudaMemcpyPeerAsync` cross-device, no conversion). So the split MUST be
3230 /// BIT-IDENTICAL to the unsplit verify at the same T, in both placement orders. Gate:
3231 /// `decode-batch-gate --mode ppspec`. Acceptance counts are a DERIVED consequence — greedy
3232 /// accept argmaxes these logits, so bit-identical logits force identical accept walks; the
3233 /// `run-spec` K=1..8 arm checks that end-to-end rather than trusting the implication.
3234 #[allow(clippy::too_many_arguments)]
3235 fn decode_step_t_core_ppn(
3236 &self,
3237 e: &Engine,
3238 tokens: &[u32],
3239 pos0: usize,
3240 cache: &mut Cache,
3241 embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
3242 mut ckpt: Option<&mut VerifyCkpt>,
3243 stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
3244 fence: &[usize],
3245 pp_pipe: Option<bool>,
3246 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
3247 let ticket = self.verify_stage0_issue(
3248 e,
3249 tokens,
3250 pos0,
3251 cache,
3252 embd_dev,
3253 ckpt.as_deref_mut(),
3254 stream,
3255 fence,
3256 pp_pipe,
3257 None,
3258 )?;
3259 self.verify_stage1_finish(e, ticket, cache, ckpt, stream, fence, true)
3260 }
3261
3262 /// Enqueue embed, stage 0, and the first boundary TX, then return the actual boundary slot.
3263 /// The ordinary PP verify wrapper calls `verify_stage1_finish` immediately after this return.
3264 #[allow(clippy::too_many_arguments)]
3265 fn verify_stage0_issue(
3266 &self,
3267 e: &Engine,
3268 tokens: &[u32],
3269 pos0: usize,
3270 cache: &mut Cache,
3271 embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
3272 mut ckpt: Option<&mut VerifyCkpt>,
3273 stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
3274 fence: &[usize],
3275 pp_pipe: Option<bool>,
3276 trace: Option<SpecPipeTraceCtx>,
3277 ) -> Result<VerifyBoundaryTicket, Box<dyn std::error::Error>> {
3278 assert!(
3279 !self.is_gemma4_e4b() && self.cfg.gemma4.is_none(),
3280 "decode_step_t_core_ppn covers the hybrid non-gemma4 verify trunk only \
3281 (the gemma4 arms have their own decode_step_t twins)"
3282 );
3283 if crate::pp::pp_host_bounce_active() && (stream.is_some() || embd_dev.is_some()) {
3284 return Err(
3285 "decode_step_t_core_ppn: refused with MEMRA_PP_HOST_BOUNCE=1 — the trunk \
3286 boundary itself is host-staged, but device-resident verify still peer-reads \
3287 primary-device token/position/embedding buffers from stage 0. Run plain PP \
3288 serving on this host class; spec requires local per-stage inputs first."
3289 .into(),
3290 );
3291 }
3292 let rt = crate::pp::PpNRt::get(e)?;
3293 let n_st = fence.len() - 1;
3294 assert_eq!(
3295 rt.n_stages(),
3296 n_st,
3297 "PpNRt stage count {} != fence stages {n_st}",
3298 rt.n_stages()
3299 );
3300 let n_embd = self.cfg.n_embd as usize;
3301 let t = tokens.len();
3302 let payload = t * n_embd;
3303 if pp_pipe.is_some() {
3304 assert_eq!(n_st, 2, "spec pipeline requires exactly two PP stages");
3305 }
3306 // One-shot lane diagnostic: force natural PP-2 boundaries to completion so the server
3307 // log can price stage 0, the peer hop, the RX copy, and stage 1 + head separately without
3308 // nsys. The ordinary path keeps every enqueue asynchronous. N>2 is deliberately excluded:
3309 // the report below names exactly two stages and must never imply it measured middle ones.
3310 let pp_anatomy = n_st == 2 && std::env::var("MEMRA_SPEC_PP_ANATOMY").as_deref() == Ok("1");
3311 let pp_started = std::time::Instant::now();
3312 let (mut reverse_ms, mut stage0_ms, mut tx_ms) = (0.0f64, 0.0f64, 0.0f64);
3313 // The CALLER's ambient stream, captured BEFORE any `rt.enter()` pushes a stage stream:
3314 // this body returns DEVICE-RESIDENT buffers (the device-argmax accept walk's contract),
3315 // so the exit needs the same publication the boundaries get — see `PpNRt::publish_to`.
3316 // Taken here, not at the end, because inside the last-stage scope `e.stream()` IS the
3317 // stage stream and the wait would self-order into a no-op.
3318 let caller_stream = e.stream();
3319 // #87 ROOT-CAUSE FENCE (lane/pp2spec-crash): the PREVIOUS round's stage-allocated
3320 // outputs (logits/hidden/ckpt stashes) freed stream-ordered on the STAGE streams while
3321 // the primary stream still holds queued reads of them — with event tracking elided,
3322 // nothing stops the pool from reusing those blocks for THIS round's stage allocations,
3323 // whose writes then race the queued reads (measured: 13/4096-NaN random-bits garbage in
3324 // the spec round seed; the full anatomy is on `PpNRt::fence_stages_behind`). Order every
3325 // stage stream behind the caller before enqueueing new stage work.
3326 let reverse_started = std::time::Instant::now();
3327 if pp_pipe != Some(false) {
3328 rt.fence_stages_behind(&caller_stream)?;
3329 }
3330 if pp_pipe == Some(true) {
3331 // Both session verifies must alternate boundary slots even when the ordinary
3332 // decode overlap experiment is off. Prewarm before A's stage 0 so B cannot grow
3333 // slot 1 by synchronizing the RX stream while A's stage 1 is in flight.
3334 rt.prepare_overlap_slots(0, payload)?;
3335 }
3336 if pp_anatomy {
3337 // Drain the reverse-publication dependency before timing stage 0 itself. At c=1 this
3338 // prices any primary-stream rollback/refresh tail inherited from the prior round.
3339 for s in 0..n_st {
3340 let _st = rt.enter(s);
3341 rt.engine(s, e).stream().synchronize()?;
3342 }
3343 reverse_ms = reverse_started.elapsed().as_secs_f64() * 1e3;
3344 }
3345
3346 // Per-stage rope positions: in host mode the same [T] iota each stage uploads itself; in
3347 // stream mode each stage's own `pos_iota` over the shared read-only device counter.
3348 let stage_pos = |es: &Engine| -> Result<CudaSlice<i32>, Box<dyn std::error::Error>> {
3349 match stream {
3350 Some((_, ctr)) => {
3351 let mut p = es.alloc_uninit::<i32>(t)?;
3352 es.pos_iota(ctr, &mut p, t)?;
3353 Ok(p)
3354 }
3355 None => {
3356 let pos_vec: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
3357 es.htod_i32(&pos_vec)
3358 }
3359 }
3360 };
3361
3362 // ---- STAGE 0: embed (the table lives with stage 0) + layers [0, fence[1]) + TX ----
3363 let slot = {
3364 let _st0 = rt.enter(0);
3365 let e0 = rt.engine(0, e);
3366 enqueue_spec_pipe_trace_marker(&e0.stream(), trace.as_ref(), "S0", "start", None)?;
3367 let stage0_started = std::time::Instant::now();
3368 let pos_d = stage_pos(e0)?;
3369 let x = match (stream, embd_dev) {
3370 (Some((vtok, _)), Some((g, qt, rb))) => {
3371 e0.embed_gather_device_td(g, vtok, t, n_embd, qt, rb)?
3372 }
3373 (None, Some((g, qt, rb))) => e0.embed_gather_device_t(g, tokens, n_embd, qt, rb)?,
3374 _ => e0.htod(&self.embd.gather(n_embd, tokens))?,
3375 };
3376 let x = self.verify_layers(
3377 e0,
3378 x,
3379 fence[0],
3380 fence[1],
3381 &pos_d,
3382 pos0,
3383 t,
3384 cache,
3385 ckpt.as_deref_mut(),
3386 stream,
3387 )?;
3388 if pp_anatomy {
3389 e0.stream().synchronize()?;
3390 stage0_ms = stage0_started.elapsed().as_secs_f64() * 1e3;
3391 }
3392 let tx_started = std::time::Instant::now();
3393 let slot = if pp_pipe.is_some() {
3394 rt.tx_pipelined(0, &x, payload)?
3395 } else {
3396 rt.tx(0, &x, payload)?
3397 };
3398 enqueue_spec_pipe_trace_marker(&e0.stream(), trace.as_ref(), "S0", "end", Some(slot))?;
3399 if pp_anatomy {
3400 e0.stream().synchronize()?;
3401 tx_ms = tx_started.elapsed().as_secs_f64() * 1e3;
3402 }
3403 slot
3404 // x + pos_d drop here: freed stream-ordered on stage-0's stream after use.
3405 };
3406
3407 Ok(VerifyBoundaryTicket {
3408 rt,
3409 caller_stream,
3410 slot,
3411 pos0,
3412 t,
3413 payload,
3414 n_st,
3415 pipelined: pp_pipe.is_some(),
3416 pp_anatomy,
3417 pp_started,
3418 reverse_ms,
3419 stage0_ms,
3420 tx_ms,
3421 trace,
3422 })
3423 }
3424
3425 /// Consume a stage-0 boundary ticket and enqueue the remaining PP stages plus the head.
3426 /// On PP-2 this is exactly stage 1; PP-N keeps its pre-existing middle-stage walk here.
3427 #[allow(clippy::too_many_arguments)]
3428 fn verify_stage1_finish(
3429 &self,
3430 e: &Engine,
3431 ticket: VerifyBoundaryTicket,
3432 cache: &mut Cache,
3433 mut ckpt: Option<&mut VerifyCkpt>,
3434 stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
3435 fence: &[usize],
3436 publish_to_caller: bool,
3437 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
3438 let VerifyBoundaryTicket {
3439 rt,
3440 caller_stream,
3441 slot,
3442 pos0,
3443 t,
3444 payload,
3445 n_st,
3446 pipelined,
3447 pp_anatomy,
3448 pp_started,
3449 reverse_ms,
3450 stage0_ms,
3451 tx_ms,
3452 trace,
3453 } = ticket;
3454 let n_embd = self.cfg.n_embd as usize;
3455 let eps = self.cfg.rms_eps;
3456 let mut slot = slot;
3457 let (mut rx_ms, mut stage1_ms) = (0.0f64, 0.0f64);
3458 let stage_pos = |es: &Engine| -> Result<CudaSlice<i32>, Box<dyn std::error::Error>> {
3459 match stream {
3460 Some((_, ctr)) => {
3461 let mut p = es.alloc_uninit::<i32>(t)?;
3462 es.pos_iota(ctr, &mut p, t)?;
3463 Ok(p)
3464 }
3465 None => {
3466 let pos_vec: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
3467 es.htod_i32(&pos_vec)
3468 }
3469 }
3470 };
3471
3472 // ---- MIDDLE STAGES: RX boundary s-1 -> range -> TX boundary s ----
3473 for s in 1..n_st - 1 {
3474 let _st = rt.enter(s);
3475 let es = rt.engine(s, e);
3476 let pos_d = stage_pos(es)?;
3477 let x = rt.rx(s - 1, slot, payload)?;
3478 let x = self.verify_layers(
3479 es,
3480 x,
3481 fence[s],
3482 fence[s + 1],
3483 &pos_d,
3484 pos0,
3485 t,
3486 cache,
3487 ckpt.as_deref_mut(),
3488 stream,
3489 )?;
3490 slot = if pipelined {
3491 rt.tx_pipelined(s, &x, payload)?
3492 } else {
3493 rt.tx(s, &x, payload)?
3494 };
3495 }
3496
3497 // ---- LAST STAGE: RX + final range + output_norm + lm head ----
3498 let _stl = rt.enter(n_st - 1);
3499 let el = rt.engine(n_st - 1, e);
3500 let pos_d = stage_pos(el)?;
3501 let rx_started = std::time::Instant::now();
3502 let x = rt.rx(n_st - 2, slot, payload)?;
3503 if pp_anatomy {
3504 el.stream().synchronize()?;
3505 rx_ms = rx_started.elapsed().as_secs_f64() * 1e3;
3506 }
3507 enqueue_spec_pipe_trace_marker(&el.stream(), trace.as_ref(), "S1", "start", Some(slot))?;
3508 let stage1_started = std::time::Instant::now();
3509 let x = self.verify_layers(
3510 el,
3511 x,
3512 fence[n_st - 1],
3513 fence[n_st],
3514 &pos_d,
3515 pos0,
3516 t,
3517 cache,
3518 ckpt.as_deref_mut(),
3519 stream,
3520 )?;
3521
3522 let mut hn = vbuf(el, payload)?;
3523 let logits = if self.cfg.step35.is_some() {
3524 // The PP Step35 serving path uses rms_norm + matmul for B=1 as well as B>1.
3525 // Verify must not switch numeric class merely because the same session speculates.
3526 el.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
3527 el.matmul(&self.output, &hn, t)?
3528 } else {
3529 el.rms_norm_decode(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
3530 el.matmul_decode_exact(&self.output, &hn, t)?
3531 };
3532 enqueue_spec_pipe_trace_marker(&el.stream(), trace.as_ref(), "S1", "end", Some(slot))?;
3533 if pp_anatomy {
3534 el.stream().synchronize()?;
3535 stage1_ms = stage1_started.elapsed().as_secs_f64() * 1e3;
3536 }
3537 // EXIT PUBLICATION: both returned buffers are still being produced on the last stage's
3538 // stream. Order the caller's stream behind that work before the buffers escape this
3539 // scope (the 2026-08-06 same-device ppspec find: without it the caller's primary-stream
3540 // consumer read unwritten logits — nondeterministic, one-device-only, and it poisoned
3541 // the following arm's KV in the same process).
3542 if publish_to_caller {
3543 rt.publish_to(n_st - 1, &caller_stream)?;
3544 }
3545 if pp_anatomy {
3546 if publish_to_caller {
3547 caller_stream.synchronize()?;
3548 }
3549 eprintln!(
3550 "[spec-pp-anatomy] t={t} reverse={reverse_ms:.3}ms stage0={stage0_ms:.3}ms \
3551 tx={tx_ms:.3}ms rx={rx_ms:.3}ms stage1-head={stage1_ms:.3}ms total={:.3}ms",
3552 pp_started.elapsed().as_secs_f64() * 1e3,
3553 );
3554 }
3555 // stream: the device pos counter owns position; host mirror reconciles at drain.
3556 if stream.is_none() {
3557 cache.pos += t;
3558 }
3559 Ok((logits, if spec_hpost() { hn } else { x }))
3560 }
3561
3562 /// Step3.5/Step3.7 verify trunk in the serving batched numeric class.
3563 ///
3564 /// `step35_decode_batch_layers` is now authoritative at every live serving width, including
3565 /// B=1 (lane/cx-b1fix). The older verify walk deliberately mirrored the eager T=1 class:
3566 /// it replayed `step35_decode_attn` per row and used the eager/decode-exact FFN dispatch.
3567 /// Those classes are individually stable, but a near-tie prompt can choose different greedy
3568 /// bytes when a request moves from batched plain serving into speculative verify. Run the
3569 /// same authoritative B=1 stage subgraph for each verify row here. Rows still advance
3570 /// layer-by-layer, so every layer sees the preceding verify rows in its attention cache while
3571 /// every norm/projection/FFN uses exactly the live serving dispatch.
3572 #[allow(clippy::too_many_arguments)]
3573 fn step35_verify_batch_layers(
3574 &self,
3575 e: &Engine,
3576 mut x: CudaSlice<f32>,
3577 lo: usize,
3578 hi: usize,
3579 pos0: usize,
3580 t: usize,
3581 cache: &mut Cache,
3582 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3583 let n_embd = self.cfg.n_embd as usize;
3584 self.cfg
3585 .step35
3586 .as_ref()
3587 .ok_or("step35 verify batch requires step35 cfg")?;
3588 let mut ph_last = std::time::Instant::now();
3589 for il in lo..hi {
3590 let mut next = e.uninit(t * n_embd)?;
3591 for r in 0..t {
3592 let mut row = e.uninit(n_embd)?;
3593 e.dtod_copy_view(&x.slice(r * n_embd..(r + 1) * n_embd), &mut row)?;
3594 // The caller owns this verify's position. During controller overlap, cache.pos
3595 // still describes generation N while this stage-0 walk belongs to N+1.
3596 let row_pos = e.htod_i32(&[(pos0 + r) as i32])?;
3597 let mut one = [&mut *cache];
3598 let out = self.step35_decode_batch_layers(
3599 e,
3600 row,
3601 &mut one,
3602 &row_pos,
3603 il,
3604 il + 1,
3605 &mut ph_last,
3606 )?;
3607 e.dtod_copy_into(&out, &mut next, r * n_embd)?;
3608 }
3609 x = next;
3610 }
3611 Ok(x)
3612 }
3613
3614 /// Qwen35-family verify trunk in the live serving numeric class.
3615 ///
3616 /// Serving intentionally keeps this architecture in the generic batched program even at
3617 /// B=1. The older verify walk used its own mirrored dispatch and can flip near-tie argmaxes.
3618 ///
3619 /// Two arms, one numeric class:
3620 /// - DENSE (`Arch::Qwen35`, t<=16): `qwen35_verify_tparallel` — the weight ops (norms,
3621 /// projections, FFN) hoist to m=T through the exact-tier batched kernels whose per-row
3622 /// program IS the m=1 program (`matmul_pre == fused2 per (tensor,row); _bN mmvq per-row
3623 /// == m=1` — decode_batch.rs v2 note), while the state ops (conv ring, gdn scan, KV
3624 /// append, fa decode) stay a per-row loop running the b_n=1 serving kernels with each
3625 /// row's own t_kv-driven arm pick (the straddle law: every row executes the exact
3626 /// program its isolated serving step would). One weight read per layer per round
3627 /// instead of T — this is what makes MTP profitable in the exact class (the per-row
3628 /// walk measured verify(K+1) ~= (K+1) plain steps: 69 -> 44 tok/s served, 2026-08-15).
3629 /// - MoE / t>16 / `MEMRA_SPEC_VERIFY_ROWWISE=1`: the per-row replay of the authoritative
3630 /// serving layer body, preserving single-session autoregressive cache order (the
3631 /// correctness reference; also the rollback seam for the t-parallel arm).
3632 ///
3633 /// Bit-identity of the t-parallel arm vs the rowwise arm is gated by spec-serve-gate
3634 /// (zero differing logits at T=1..4, K arms) + the 8-prompt ON/OFF canary before ship.
3635 #[allow(clippy::too_many_arguments)]
3636 fn qwen35_verify_batch_layers(
3637 &self,
3638 e: &Engine,
3639 x: CudaSlice<f32>,
3640 lo: usize,
3641 hi: usize,
3642 pos0: usize,
3643 t: usize,
3644 cache: &mut Cache,
3645 ckpt: Option<&mut VerifyCkpt>,
3646 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3647 let rowwise = std::env::var("MEMRA_SPEC_VERIFY_ROWWISE").as_deref() == Ok("1")
3648 || !matches!(self.cfg.arch, memra_gguf::config::Arch::Qwen35)
3649 || t > 16;
3650 if rowwise {
3651 self.qwen35_verify_rowwise(e, x, lo, hi, pos0, t, cache, ckpt)
3652 } else {
3653 self.qwen35_verify_tparallel(e, x, lo, hi, pos0, t, cache, ckpt)
3654 }
3655 }
3656
3657 /// The per-row correctness reference: replay each verify row through the authoritative
3658 /// serving layer body (`decode_batch_layers` at b_n=1). T full weight reads per layer.
3659 #[allow(clippy::too_many_arguments)]
3660 fn qwen35_verify_rowwise(
3661 &self,
3662 e: &Engine,
3663 mut x: CudaSlice<f32>,
3664 lo: usize,
3665 hi: usize,
3666 pos0: usize,
3667 t: usize,
3668 cache: &mut Cache,
3669 mut ckpt: Option<&mut VerifyCkpt>,
3670 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3671 let n_embd = self.cfg.n_embd as usize;
3672 let saved_pos = cache.pos;
3673 let mut ph_last = std::time::Instant::now();
3674 for il in lo..hi {
3675 let mut next = e.uninit(t * n_embd)?;
3676 let mut col_states: Option<Vec<(CudaSlice<f32>, CudaSlice<f32>)>> =
3677 if ckpt.is_some() && t >= 2 && matches!(self.layers[il].mixer, Mixer::Linear(_)) {
3678 Some(Vec::with_capacity(t - 1))
3679 } else {
3680 None
3681 };
3682 for r in 0..t {
3683 cache.pos = pos0 + r;
3684 let mut row = e.uninit(n_embd)?;
3685 e.dtod_copy_view(&x.slice(r * n_embd..(r + 1) * n_embd), &mut row)?;
3686 let row_pos = e.htod_i32(&[(pos0 + r) as i32])?;
3687 let mut one = [&mut *cache];
3688 let ctx = self.batch_layer_ctx(e, &one, il, il + 1)?;
3689 let out = match self.decode_batch_layers(
3690 e,
3691 row,
3692 &mut one,
3693 &ctx,
3694 &row_pos,
3695 &mut ph_last,
3696 ) {
3697 Ok(out) => out,
3698 Err(error) => {
3699 cache.pos = saved_pos;
3700 return Err(error);
3701 }
3702 };
3703 e.dtod_copy_into(&out, &mut next, r * n_embd)?;
3704 if r + 1 < t {
3705 if let Some(states) = col_states.as_mut() {
3706 let recur = cache.recur[il]
3707 .as_ref()
3708 .ok_or("Qwen35-MoE linear verify layer has no recurrent state")?;
3709 states.push((
3710 e.clone_dtod(&recur.conv_state)?,
3711 e.clone_dtod(&recur.ssm_state)?,
3712 ));
3713 }
3714 }
3715 }
3716 if let (Some(checkpoint), Some(states)) = (ckpt.as_deref_mut(), col_states) {
3717 checkpoint.cols[il] = Some(states);
3718 }
3719 x = next;
3720 }
3721 cache.pos = saved_pos;
3722 Ok(x)
3723 }
3724
3725 /// T-PARALLEL VERIFY IN THE SERVING NUMERIC CLASS (lane/tparallel-verify, 2026-08-15).
3726 ///
3727 /// The weight ops run ONCE per layer at m=T; the state ops run per row through the same
3728 /// b_n=1 serving kernels the rowwise replay uses. Per-row bit-identity rests on the two
3729 /// pins the serving batch tier already carries:
3730 /// * `matmul_pre` / `_bN` mmvq: per-row program == m=1 program (decode_batch.rs v2 note,
3731 /// kernel-check pinned) — so a [T, n_embd] projection row equals the row projected
3732 /// alone;
3733 /// * row-indexed norms/elementwise (`rms_norm`, `quantize_q8_1`, `add_rms_norm`,
3734 /// `gated_rmsnorm[_q8_1]`, `silu_mul`, `rope_neox` with per-row positions): the T-row
3735 /// launch is the per-row program (same pin the generic verify's fused norms rely on).
3736 /// The sequential dependencies keep their exact serving order: the conv ring / gdn scan
3737 /// chain state row -> row through the `_b` kernels at b_n=1 (ping-pong via a 6-entry
3738 /// alternating pointer table, host handles swapped per row so VerifyCkpt clones the
3739 /// canonical state exactly as the rowwise arm does), and each row's KV append + fa decode
3740 /// picks its arm from ITS OWN t_kv (append: format-only; fa: `fa_seqs_eligible` + its own
3741 /// `fa_split_keys` rung at b_n=1) — the straddle law per row, so every row executes the
3742 /// program its isolated B=1 serving step would.
3743 ///
3744 /// Cost: 1 weight read per layer per round + T state micro-launches, vs the rowwise arm's
3745 /// T weight reads. Gated bit-identical vs the rowwise arm by spec-serve-gate + canary.
3746 #[allow(clippy::too_many_arguments)]
3747 fn qwen35_verify_tparallel(
3748 &self,
3749 e: &Engine,
3750 mut x: CudaSlice<f32>,
3751 lo: usize,
3752 hi: usize,
3753 pos0: usize,
3754 t: usize,
3755 cache: &mut Cache,
3756 mut ckpt: Option<&mut VerifyCkpt>,
3757 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3758 use cudarc::driver::DevicePtr;
3759 let cfg = &self.cfg;
3760 let n_embd = cfg.n_embd as usize;
3761 let eps = cfg.rms_eps;
3762 let head_dim_global = cfg.head_dim_k as usize;
3763 let pos_host: Vec<i32> = (0..t).map(|r| (pos0 + r) as i32).collect();
3764 let pos_d = e.htod_i32(&pos_host)?;
3765 let seqs_append =
3766 std::env::var("MEMRA_BATCH_APPEND").as_deref() != Ok("0") && !Engine::kv_fp8_on();
3767 let batch_fa_on = std::env::var("MEMRA_BATCH_FA").as_deref() != Ok("0");
3768
3769 for il in lo..hi {
3770 let layer = &self.layers[il];
3771 // ---- attn_norm + q8_1 quantize at m=T (row-indexed == per-row) ----
3772 let anorm = layer.attn_norm.float_data();
3773 let mut xn = e.uninit(t * n_embd)?;
3774 e.rms_norm(&x, anorm, &mut xn, n_embd, t, eps)?;
3775 let (hq, hd) = e.quantize_q8_1(&xn, t, n_embd)?;
3776
3777 let mixed: CudaSlice<f32> = match &layer.mixer {
3778 Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
3779 Mixer::Full(fa) => {
3780 let geometry = cfg.full_attention_geometry_at(il as u32);
3781 let n_head = geometry.n_head as usize;
3782 let n_head_kv = geometry.n_head_kv as usize;
3783 let head_dim = geometry.head_dim_k as usize;
3784 let rope_dims = geometry.n_rot as usize;
3785 let rope_base = geometry.rope_base;
3786 let scale = geometry.attention_scale();
3787 // Batched projections: one weight read serves all T rows.
3788 let qf = e.matmul_pre(&fa.wq, &hq, &hd, &xn, t)?;
3789 let mut k = e.matmul_pre(&fa.wk, &hq, &hd, &xn, t)?;
3790 let v = e.matmul_pre(&fa.wv, &hq, &hd, &xn, t)?;
3791 let gated =
3792 geometry.attention_gate == memra_gguf::config::AttentionGateKind::FusedQ;
3793 let (mut q, gate) = if gated {
3794 let mut qs = e.uninit(t * n_head * head_dim)?;
3795 let mut gs = e.uninit(t * n_head * head_dim)?;
3796 e.q_gate_split(&qf, &mut qs, &mut gs, head_dim, n_head, t)?;
3797 (qs, Some(gs))
3798 } else {
3799 (qf, None)
3800 };
3801 let mut qn = e.uninit(t * n_head * head_dim)?;
3802 e.rms_norm(
3803 &q,
3804 fa.q_norm.float_data(),
3805 &mut qn,
3806 head_dim,
3807 t * n_head,
3808 eps,
3809 )?;
3810 q = qn;
3811 let mut kn = e.uninit(t * n_head_kv * head_dim)?;
3812 e.rms_norm(
3813 &k,
3814 fa.k_norm.float_data(),
3815 &mut kn,
3816 head_dim,
3817 t * n_head_kv,
3818 eps,
3819 )?;
3820 k = kn;
3821 e.rope_neox(
3822 &mut q, &pos_d, head_dim, rope_dims, n_head, t, rope_base, 1.0,
3823 )?;
3824 e.rope_neox(
3825 &mut k, &pos_d, head_dim, rope_dims, n_head_kv, t, rope_base, 1.0,
3826 )?;
3827
3828 // Per-row append + attend: row r sees rows 0..r in KV (causal within the
3829 // draft), each through the b_n=1 serving kernels at its own t_kv.
3830 let q_dim = n_head * head_dim;
3831 let kv_dim = n_head_kv * head_dim;
3832 let mut attn = e.uninit(t * q_dim)?;
3833 let (kdk, kdv, ktb, vtb, kv_view) = {
3834 let kvl = cache.kv[il].as_ref().unwrap();
3835 let s = &e.gpu.stream();
3836 let (pk, _g) = kvl.k.device_ptr(s);
3837 let (pv, _g2) = kvl.v.device_ptr(s);
3838 (
3839 kvl.kv_dim_k,
3840 kvl.kv_dim_v,
3841 kvl.k_tok_bytes,
3842 kvl.v_tok_bytes,
3843 e.htod_u64(&[pk as u64, pv as u64])?,
3844 )
3845 };
3846 for r in 0..t {
3847 // Owned per-row scratch: the b_n=1 kernels take packed batch buffers
3848 // whose row 0 is this row (arithmetic-free materialization copies,
3849 // same as decode's per-seq fallback arm).
3850 let mut k_row = e.uninit(kv_dim)?;
3851 e.dtod_copy_view(&k.slice(r * kv_dim..(r + 1) * kv_dim), &mut k_row)?;
3852 let mut v_row = e.uninit(kv_dim)?;
3853 e.dtod_copy_view(&v.slice(r * kv_dim..(r + 1) * kv_dim), &mut v_row)?;
3854 let pos_row = e.htod_i32(&[(pos0 + r) as i32])?;
3855 let kvl = cache.kv[il].as_mut().unwrap();
3856 if seqs_append {
3857 e.append_kv_quantized_seqs(
3858 &k_row,
3859 &v_row,
3860 &kv_view.slice(0..2),
3861 &pos_row,
3862 1,
3863 kdk,
3864 kdv,
3865 ktb,
3866 vtb,
3867 )?;
3868 kvl.len += 1;
3869 } else {
3870 e.append_kv_quantized_view(
3871 &k_row.slice(0..kv_dim),
3872 &v_row.slice(0..kv_dim),
3873 &mut kvl.k,
3874 &mut kvl.v,
3875 kvl.len,
3876 kvl.kv_dim_k,
3877 kvl.kv_dim_v,
3878 kvl.k_tok_bytes,
3879 kvl.v_tok_bytes,
3880 Engine::kv_fp8_on(),
3881 )?;
3882 kvl.len += 1;
3883 }
3884 let t_kv = kvl.len;
3885 let mut q_row = e.uninit(q_dim)?;
3886 e.dtod_copy_view(&q.slice(r * q_dim..(r + 1) * q_dim), &mut q_row)?;
3887 let mut a_row = e.uninit(q_dim)?;
3888 if batch_fa_on && crate::fa_seqs_eligible(t_kv, head_dim_global) {
3889 let sp0_r = crate::fa_split_keys(t_kv, cfg.n_head_kv as usize);
3890 e.fa_decode_batch_seqs_v4(
3891 &q_row,
3892 &kv_view.slice(0..2),
3893 &pos_row,
3894 &mut a_row,
3895 head_dim,
3896 n_head,
3897 n_head_kv,
3898 1,
3899 t_kv,
3900 scale,
3901 sp0_r,
3902 ktb,
3903 vtb,
3904 )?;
3905 } else {
3906 let k_view = e.view_u8(&kvl.k, t_kv * kvl.k_tok_bytes);
3907 let v_view = e.view_u8(&kvl.v, t_kv * kvl.v_tok_bytes);
3908 let mut a_view = a_row.slice_mut(0..q_dim);
3909 e.fa_decode_kvmod_view(
3910 &q_row.slice(0..q_dim),
3911 &k_view,
3912 &v_view,
3913 &mut a_view,
3914 head_dim,
3915 n_head,
3916 n_head_kv,
3917 t_kv,
3918 scale,
3919 kvl.k_tok_bytes,
3920 kvl.v_tok_bytes,
3921 Engine::kv_fp8_on(),
3922 )?;
3923 }
3924 e.dtod_copy_into(&a_row, &mut attn, r * q_dim)?;
3925 }
3926
3927 // Output gate (element-wise) + o-proj at m=T.
3928 let attn_g = match &gate {
3929 Some(g) => {
3930 let n = t * q_dim;
3931 let mut gsig = e.uninit(n)?;
3932 e.sigmoid(g, &mut gsig, n)?;
3933 let mut ag = e.uninit(n)?;
3934 e.mul(&attn, &gsig, &mut ag, n)?;
3935 ag
3936 }
3937 None => attn,
3938 };
3939 e.matmul(&fa.wo, &attn_g, t)?
3940 }
3941 Mixer::Linear(la) => {
3942 let ssm = cfg.ssm.as_ref().expect("linear mixer requires ssm cfg");
3943 let d_state = ssm.state_size as usize;
3944 let num_k = ssm.group_count as usize;
3945 let num_v = ssm.time_step_rank as usize;
3946 let d_conv = ssm.conv_kernel as usize;
3947 let key_dim = d_state * num_k;
3948 let value_dim = d_state * num_v;
3949 let conv_dim = key_dim * 2 + value_dim;
3950 let gdn_scale = 1.0 / (d_state as f32).sqrt();
3951
3952 // ---- batched projections: one weight read for all T rows ----
3953 let qkv_mixed = e.matmul_pre(&la.wqkv, &hq, &hd, &xn, t)?;
3954 let z = e.matmul_pre(&la.wqkv_gate, &hq, &hd, &xn, t)?;
3955 let beta_raw = e.matmul_pre(&la.ssm_beta, &hq, &hd, &xn, t)?;
3956 let alpha = e.matmul_pre(&la.ssm_alpha, &hq, &hd, &xn, t)?;
3957 let beta_w = la.ssm_beta.out_features();
3958 let alpha_w = la.ssm_alpha.out_features();
3959 let qkv_w = la.wqkv.out_features();
3960
3961 // ---- per-row state chain through the b_n=1 serving kernels ----
3962 // 6-entry alternating pointer table expresses the ping-pong without a
3963 // rebuild per row: even rows scan s0 -> s1, odd rows s1 -> s0. Host
3964 // handles swap per row so ckpt clones the canonical state (and the
3965 // post-verify canonical handle matches the last write), exactly as the
3966 // rowwise arm leaves them.
3967 let table = {
3968 let rl = cache.recur[il].as_ref().unwrap();
3969 let s = &e.gpu.stream();
3970 let (pc, _g0) = rl.conv_state.device_ptr(s);
3971 let (p0, _g1) = rl.ssm_state.device_ptr(s);
3972 let (p1, _g2) = rl.ssm_state_alt.device_ptr(s);
3973 e.htod_u64(&[
3974 pc as u64, p0 as u64, p1 as u64, pc as u64, p1 as u64, p0 as u64,
3975 ])?
3976 };
3977 let mut o_all = e.uninit(t * value_dim)?;
3978 let mut col_states: Option<Vec<(CudaSlice<f32>, CudaSlice<f32>)>> =
3979 if ckpt.is_some() && t >= 2 {
3980 Some(Vec::with_capacity(t - 1))
3981 } else {
3982 None
3983 };
3984 for r in 0..t {
3985 let base = if r % 2 == 0 { 0 } else { 3 };
3986 let conv_view = table.slice(base..base + 1);
3987 let in_view = table.slice(base + 1..base + 2);
3988 let out_view = table.slice(base + 2..base + 3);
3989 let mut qkv_row = e.uninit(qkv_w)?;
3990 e.dtod_copy_view(
3991 &qkv_mixed.slice(r * qkv_w..(r + 1) * qkv_w),
3992 &mut qkv_row,
3993 )?;
3994 let mut conv_out = e.uninit(conv_dim)?;
3995 e.ssm_conv1d_fused_decode_b(
3996 &qkv_row,
3997 &conv_view,
3998 la.ssm_conv1d.float_data(),
3999 &mut conv_out,
4000 conv_dim,
4001 d_conv,
4002 1,
4003 )?;
4004 let mut beta_row = e.uninit(beta_w)?;
4005 e.dtod_copy_view(
4006 &beta_raw.slice(r * beta_w..(r + 1) * beta_w),
4007 &mut beta_row,
4008 )?;
4009 let mut alpha_row = e.uninit(alpha_w)?;
4010 e.dtod_copy_view(
4011 &alpha.slice(r * alpha_w..(r + 1) * alpha_w),
4012 &mut alpha_row,
4013 )?;
4014 let mut q_l2 = e.uninit(value_dim)?;
4015 let mut k_l2 = e.uninit(value_dim)?;
4016 let mut v_gd = e.uninit(value_dim)?;
4017 let mut beta_b = e.uninit(num_v)?;
4018 let mut g_log = e.uninit(num_v)?;
4019 e.gdn_prep_decode_b(
4020 &conv_out,
4021 &beta_row,
4022 &alpha_row,
4023 la.ssm_dt.float_data(),
4024 la.ssm_a.float_data(),
4025 &mut q_l2,
4026 &mut k_l2,
4027 &mut v_gd,
4028 &mut beta_b,
4029 &mut g_log,
4030 d_state,
4031 num_v,
4032 num_k,
4033 key_dim,
4034 eps,
4035 conv_dim,
4036 1,
4037 )?;
4038 let mut o_row = e.uninit(value_dim)?;
4039 e.gdn_scan_s128_batched(
4040 &q_l2, &k_l2, &v_gd, &g_log, &beta_b, &in_view, &out_view, &mut o_row,
4041 num_v, 1, gdn_scale,
4042 )?;
4043 e.dtod_copy_into(&o_row, &mut o_all, r * value_dim)?;
4044 {
4045 let rl = cache.recur[il].as_mut().unwrap();
4046 std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
4047 }
4048 if r + 1 < t {
4049 if let Some(states) = col_states.as_mut() {
4050 let recur = cache.recur[il]
4051 .as_ref()
4052 .ok_or("qwen35 linear verify layer has no recurrent state")?;
4053 states.push((
4054 e.clone_dtod(&recur.conv_state)?,
4055 e.clone_dtod(&recur.ssm_state)?,
4056 ));
4057 }
4058 }
4059 }
4060 if let (Some(checkpoint), Some(states)) = (ckpt.as_deref_mut(), col_states) {
4061 checkpoint.cols[il] = Some(states);
4062 }
4063
4064 // ---- batched gated norm + out-projection at m=T ----
4065 if e.uses_q8_1_fast(&la.ssm_out) {
4066 let (gq, gd) = e.gated_rmsnorm_q8_1(
4067 &o_all,
4068 la.ssm_norm.float_data(),
4069 &z,
4070 d_state,
4071 t * num_v,
4072 eps,
4073 )?;
4074 let g0 = e.zeros(0)?;
4075 e.matmul_pre(&la.ssm_out, &gq, &gd, &g0, t)?
4076 } else {
4077 let mut gn = e.uninit(t * value_dim)?;
4078 e.gated_rmsnorm(
4079 &o_all,
4080 la.ssm_norm.float_data(),
4081 &z,
4082 &mut gn,
4083 d_state,
4084 t * num_v,
4085 eps,
4086 )?;
4087 e.matmul(&la.ssm_out, &gn, t)?
4088 }
4089 }
4090 };
4091
4092 // ---- residual add + post_attn_norm + FFN at m=T (serving dispatch verbatim) ----
4093 let pnorm = layer.post_attn_norm.float_data();
4094 let mut x1 = e.uninit(t * n_embd)?;
4095 let mut zn = e.uninit(t * n_embd)?;
4096 e.add_rms_norm(&x, &mixed, pnorm, &mut x1, &mut zn, n_embd, t, eps)?;
4097 let ffn_out = match &layer.ffn {
4098 crate::hybrid::Ffn::Dense {
4099 ffn_gate,
4100 ffn_up,
4101 ffn_down,
4102 } => {
4103 assert!(
4104 self.cfg.m3.is_none(),
4105 "qwen35 t-parallel verify: M3 swigluoai FFN not yet batched"
4106 );
4107 let n_ff = ffn_gate.out_features();
4108 let (zq, zd) = e.quantize_q8_1(&zn, t, n_embd)?;
4109 let g = e.matmul_pre(ffn_gate, &zq, &zd, &zn, t)?;
4110 let u = e.matmul_pre(ffn_up, &zq, &zd, &zn, t)?;
4111 let mut act = e.uninit(t * n_ff)?;
4112 e.silu_mul(&g, &u, &mut act, t * n_ff)?;
4113 let (aq, ad) = e.quantize_q8_1(&act, t, n_ff)?;
4114 e.matmul_pre(ffn_down, &aq, &ad, &act, t)?
4115 }
4116 crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il_zq8(e, m, &zn, None, t, il as u16)?,
4117 };
4118 let mut x2 = e.uninit(t * n_embd)?;
4119 e.add(&x1, &ffn_out, &mut x2, t * n_embd)?;
4120 x = x2;
4121 }
4122 Ok(x)
4123 }
4124
4125 /// PP-N STAGE SUBGRAPH of the verify trunk: layers `[lo, hi)` of `decode_step_t_core_stream`'s
4126 /// walk, verbatim. Enters with a MATERIALIZED `[T, n_embd]` residual (no pending fusion pair
4127 /// carried in from outside the range) and exits with the range's final residual materialized
4128 /// (the trailing add executed) — exactly the `decode_layers_eager(lo, hi)` contract, T rows
4129 /// instead of one.
4130 ///
4131 /// EXTRACTED (lane/pp2-spec 2026-08-06) rather than duplicated: `decode_step_t_core_stream` IS
4132 /// the single funnel every verify forward reaches, and its per-layer dispatch MIRRORING (norm
4133 /// fusion per layer, the t>=3/spec_m2 batched-linear window, the fused-q8 FFN chain, the
4134 /// decode-exact projections) is what makes verify bit-identical to eager decode. A second copy
4135 /// for the split arm is how those mirrors drift apart on the next lever. The unsplit body now
4136 /// calls this with `(0, n_layers)`, so the whole-trunk path and every stage range run the SAME
4137 /// code — there is no "split version" of the verify math.
4138 ///
4139 /// Bit-identity of a cut rests on the same kernel-check-pinned identity the eager arm's cut
4140 /// does — `add_rms_norm_q8_1 == add then rms_norm_q8_1` at nrows=T — because the ONLY thing a
4141 /// fence changes is that the cross-layer fusion carry breaks at `hi-1` and is re-materialized
4142 /// as an explicit `add`. `decode-batch-gate --mode ppspec` verifies end-to-end on real weights.
4143 #[allow(clippy::too_many_arguments)]
4144 fn verify_layers(
4145 &self,
4146 e: &Engine,
4147 mut x: CudaSlice<f32>,
4148 lo: usize,
4149 hi: usize,
4150 pos_d: &CudaSlice<i32>,
4151 pos0: usize,
4152 t: usize,
4153 cache: &mut Cache,
4154 mut ckpt: Option<&mut VerifyCkpt>,
4155 stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
4156 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4157 if self.cfg.step35.is_some() {
4158 if stream.is_some() {
4159 return Err(
4160 "step35 has no ROUND-STREAM verify arm (the device-counter _dc twins \
4161 cannot express the SWA offset KV view)"
4162 .into(),
4163 );
4164 }
4165 return self.step35_verify_batch_layers(e, x, lo, hi, pos0, t, cache);
4166 }
4167 if self.qwen35_serving_class() {
4168 if stream.is_some() {
4169 return Err("qwen35-family serving-class verify has no ROUND-STREAM arm".into());
4170 }
4171 return self.qwen35_verify_batch_layers(e, x, lo, hi, pos0, t, cache, ckpt.take());
4172 }
4173 let n_embd = self.cfg.n_embd as usize;
4174 let eps = self.cfg.rms_eps;
4175 // CROSS-LAYER ADD+NORM FUSION (lane/vt-fixes fix 2, mirroring decode_step_h's
4176 // launch-arc form): layer il's post-FFN residual add (x2 = x1 + ffn_out) and layer
4177 // il+1's attn_norm(+quantize) are consecutive row-wise ops — ONE add_rms_norm_q8_1
4178 // launch at nrows=t does all three (bit-identity pinned by the T-row kernel-check
4179 // arms). Carry the un-added (x1, ffn_out) pair; the fused launch materializes x2 (the
4180 // residual the next layer needs) as its `res` output. Falls back to the separate add
4181 // when the next layer is off the fused-q8 path.
4182 let mut pending: Option<(CudaSlice<f32>, CudaSlice<f32>)> = None;
4183 for il in lo..hi {
4184 let layer = &self.layers[il];
4185 // DISPATCH-MIRRORED attn-input RMSNorm (FP-order lesson #8): eager decode fuses the
4186 // 1024-thread rms_norm_q8_1 ONLY when every mixer projection is q8_1-fast; layers with
4187 // Float projections (ssm_beta/ssm_alpha on layers 1/2/4 of the 9B NVFP4 GGUF) take the
4188 // UNFUSED 256-thread rms_norm. The verify norm must mirror that PER-LAYER choice —
4189 // blockDim changes the sum-of-squares reduce order, and the ULP shift amplifies through
4190 // the GDN recurrence into argmax flips (measured: 9B text prompt, 1 ULP at layer 2 ->
4191 // 2.3e-1 logit maxdiff at the head -> K=1..8 divergence at a 0.03-margin token).
4192 let mixer_fast = self.mixer_in_q8_1_fast(e, &layer.mixer);
4193 let norm_fused = std::env::var("MEMRA_NO_FUSE_NORMQ").is_err() && mixer_fast;
4194 // BATCHED EPILOGUE RE-FUSE (lane/vt-fixes fix 2, 2026-08-03): when the norm is
4195 // dispatch-fused AND every consumer of `h` reads only its q8_1 form (Full mixer:
4196 // projections only; Linear mixer: the batched arm — the per-column fallback needs
4197 // f32 h), emit the attn-input norm DIRECTLY as q8_1 via `rms_norm_q8_1` at nrows=t
4198 // (row-indexed kernel — the T-row launch is the per-row m=1 program, kernel-check
4199 // pins bit-identity vs rms_norm_decode -> quantize_q8_1). Kills the standalone
4200 // quantize launch(es) + the f32 h HBM round-trip that decode never pays.
4201 // step35 (Full mixer) is the third case that needs f32 `h`: its verify arm is a
4202 // per-ROW replay of the eager decode mixer, whose `pre_q` contract is a single row —
4203 // a T-row q8_1 pair cannot be handed to it, and re-deriving per-row q8_1 from the f32
4204 // rows is exactly the dispatch being mirrored. Keep step35 on the unfused arm.
4205 let lin_q8_only = match &layer.mixer {
4206 Mixer::Linear(la) => {
4207 (t >= 3 || (t == 2 && spec_m2())) && e.uses_q8_1_fast(&la.ssm_out)
4208 }
4209 Mixer::Full(_) if self.cfg.step35.is_some() => false,
4210 _ => true,
4211 };
4212 // NOTE decode.rs's take()-first lesson: take the pending pair BEFORE branching so
4213 // a non-fused layer still performs the residual add.
4214 let taken = pending.take();
4215 let (h, h_q8) = if norm_fused && lin_q8_only {
4216 let pair = match taken {
4217 // fused add + attn_norm + q8_1: ONE launch resolves the carried residual
4218 // AND emits this layer's mixer input pre-quantized. res -> x2 (= new x).
4219 Some((x1p, f1p)) => {
4220 let mut x2 = vbuf(e, t * n_embd)?; // fully written (res output)
4221 let p = e.add_rms_norm_q8_1(
4222 &x1p,
4223 &f1p,
4224 layer.attn_norm.float_data(),
4225 &mut x2,
4226 n_embd,
4227 t,
4228 eps,
4229 )?;
4230 x = x2;
4231 p
4232 }
4233 None => e.rms_norm_q8_1(&x, layer.attn_norm.float_data(), n_embd, t, eps)?,
4234 };
4235 (e.zeros(0)?, Some(pair)) // h unused on this path (q8-only consumers)
4236 } else {
4237 if let Some((x1p, f1p)) = taken {
4238 let mut x2 = vbuf(e, t * n_embd)?; // fully written by add
4239 e.add(&x1p, &f1p, &mut x2, t * n_embd)?;
4240 x = x2;
4241 }
4242 let mut h = vbuf(e, t * n_embd)?; // fully written by either rms_norm arm
4243 if norm_fused {
4244 e.rms_norm_decode(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
4245 } else {
4246 e.rms_norm(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
4247 }
4248 (h, None)
4249 };
4250 let h_q8_ref = h_q8.as_ref().map(|(q, d)| (q, d));
4251
4252 let mixed = match &layer.mixer {
4253 Mixer::Full(fa) => self.full_attn_verify(
4254 e,
4255 fa,
4256 &h,
4257 h_q8_ref,
4258 pos_d,
4259 t,
4260 cache,
4261 il,
4262 stream.map(|(_, c)| c),
4263 )?,
4264 Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
4265 Mixer::Linear(la) => {
4266 // BATCHED linear verify (2026-07-03, the MTP-profit lever): one T-token pass —
4267 // batched projections (weight read ONCE, hits the m=2-4 weight-resident matvec),
4268 // carried-state conv (ssm_conv1d_tm_state), GDN prep on the prefill kernels, and
4269 // ONE gdn_scan whose internal sequential t-loop is the SAME recurrence as T
4270 // chained T=1 steps (bit-identical). Falls back to the sequential per-column
4271 // chain when T < d_conv-1 (conv ring update needs T >= pad) — or when ANY
4272 // projection is off the q8_1 fast path: matmul_decode_exact would route a Float
4273 // tensor to cuBLAS at m=t (different FP accumulation than eager's per-token
4274 // GEMV), so mixed-dtype layers stay on the eager-identical per-column chain.
4275 // MEMRA_SPEC_M2 (lane/spec-m2): the t==2 batch rides the same arm — the conv
4276 // wrapper handles t<pad with a pure-copy ring rebuild; see spec_m2() header.
4277 if (t >= 3 || (t == 2 && spec_m2()))
4278 && mixer_fast
4279 && e.uses_q8_1_fast(&la.ssm_out)
4280 {
4281 let want = ckpt.is_some();
4282 let (out, stash) =
4283 self.linear_attn_verify_t(e, la, &h, h_q8_ref, t, cache, il, want)?;
4284 if let (Some(ck), Some(st)) = (ckpt.as_deref_mut(), stash) {
4285 ck.gdn[il] = Some(st);
4286 }
4287 out
4288 } else {
4289 let mut out = vbuf(e, t * n_embd)?; // every col written by copy_into
4290 let mut col_states: Option<Vec<(CudaSlice<f32>, CudaSlice<f32>)>> =
4291 if ckpt.is_some() && t >= 2 {
4292 Some(Vec::with_capacity(t - 1))
4293 } else {
4294 None
4295 };
4296 for col in 0..t {
4297 let mut h_col = vbuf(e, n_embd)?; // fully written by copy_view_into
4298 let src = h.slice(col * n_embd..(col + 1) * n_embd);
4299 e.copy_view_into(&mut h_col, 0, &src, n_embd)?;
4300 let m_col = self.linear_attn_decode(e, la, &h_col, cache, il)?;
4301 e.copy_into(&mut out, col * n_embd, &m_col, n_embd)?;
4302 // REPLAY-FREE ckpt: clone the chain's ACTUAL state after this column
4303 // (pure dtod — cannot change any computed value). Last column skipped:
4304 // rebuild targets are j <= t-1 columns.
4305 if let Some(cs) = col_states.as_mut() {
4306 if col + 1 < t {
4307 let rl = cache.recur[il].as_ref().unwrap();
4308 cs.push((
4309 e.clone_dtod(&rl.conv_state)?,
4310 e.clone_dtod(&rl.ssm_state)?,
4311 ));
4312 }
4313 }
4314 }
4315 if let (Some(ck), Some(cs)) = (ckpt.as_deref_mut(), col_states) {
4316 // ReplaySSM-assessment instrumentation (2026-07-30): the
4317 // per-column clones are the only true state snapshots left in
4318 // the verify (the batched path stashes INPUTS and replays).
4319 if std::env::var("MEMRA_SPEC_STATS").as_deref() == Ok("1") {
4320 static ONCE: std::sync::Once = std::sync::Once::new();
4321 let bytes: usize =
4322 cs.iter().map(|(c, s)| (c.len() + s.len()) * 4).sum();
4323 ONCE.call_once(|| eprintln!(
4324 "[verify-ckpt] per-column layer il={il}: {} clones, {:.2} MB/layer/round",
4325 cs.len(), bytes as f64 / 1e6));
4326 }
4327 ck.cols[il] = Some(cs);
4328 }
4329 out
4330 }
4331 }
4332 };
4333
4334 // DISPATCH-MIRRORED post-attn norm: eager residual_norm_ffn fuses add+norm+quant
4335 // (1024-thread add_rms_norm_q8_1) only for Dense FFNs whose gate+up are q8_1-fast;
4336 // otherwise (and for MoE) it runs the 256-thread fused add_rms_norm. Mirror per layer.
4337 let ffn_fuse = match &layer.ffn {
4338 crate::hybrid::Ffn::Dense {
4339 ffn_gate, ffn_up, ..
4340 } => {
4341 std::env::var("MEMRA_NO_FUSE_NORMQ").is_err()
4342 && e.uses_q8_1_fast(ffn_gate)
4343 && e.uses_q8_1_fast(ffn_up)
4344 }
4345 crate::hybrid::Ffn::Moe(_) => false,
4346 };
4347 // BATCHED EPILOGUE RE-FUSE (lane/vt-fixes fix 2): on the ffn_fuse path (Dense,
4348 // gate+up q8_1-fast, non-M3) the FFN input is emitted DIRECTLY as q8_1 by ONE
4349 // add_rms_norm_q8_1 launch at nrows=t (row-indexed kernel: the T-row launch is the
4350 // per-row m=1 program; kernel-check pins bit-identity vs the unfused
4351 // add_f32 -> rms_norm_decode -> quantize_q8_1 chain at T=2/4/5/8) — replacing the
4352 // add + rms_norm_decode launches AND the dual/singles' internal re-quantize.
4353 // M3's swigluoai must keep the f32 chain (the fused SwiGLU epilogue encodes plain
4354 // SiLU), mirroring residual_norm_ffn's m3 guard on the decode path.
4355 // step35: same guard per LAYER. A dense FFN's clamp is the SHEXP array (upstream's
4356 // one build_ffn serves dense + shared expert, llama-graph.cpp:1751), and verify MUST
4357 // mirror decode's dispatch or spec self-consistency fails.
4358 let dense_lim = self.cfg.clamp_shexp_at(il as u32);
4359 let fuse_q8 = ffn_fuse && self.cfg.m3.is_none() && dense_lim.is_none();
4360 let mut x1 = vbuf(e, t * n_embd)?; // fully written by add / add_rms_norm*
4361 let mut z = e.zeros(0)?; // replaced below on the unfused arms
4362 let z_q8 = if fuse_q8 {
4363 Some(e.add_rms_norm_q8_1(
4364 &x,
4365 &mixed,
4366 layer.post_attn_norm.float_data(),
4367 &mut x1,
4368 n_embd,
4369 t,
4370 eps,
4371 )?)
4372 } else {
4373 let mut zf = vbuf(e, t * n_embd)?; // fully written by rms_norm_decode / add_rms_norm
4374 if ffn_fuse {
4375 e.add(&x, &mixed, &mut x1, t * n_embd)?;
4376 e.rms_norm_decode(
4377 &x1,
4378 layer.post_attn_norm.float_data(),
4379 &mut zf,
4380 n_embd,
4381 t,
4382 eps,
4383 )?;
4384 } else {
4385 e.add_rms_norm(
4386 &x,
4387 &mixed,
4388 layer.post_attn_norm.float_data(),
4389 &mut x1,
4390 &mut zf,
4391 n_embd,
4392 t,
4393 eps,
4394 )?;
4395 }
4396 z = zf;
4397 None
4398 };
4399 // DECODE-EXACT FFN projections: force MMVQ for gate/up/down at any T to match the
4400 // T=1 decode FP accumulation order. At T>=5 the generic matmul/matmul_pre falls to dp4a
4401 // (128-thread, different FP sum order). At T=2-4 the batched MMVQ is already bit-identical.
4402 let ffn_out = match &layer.ffn {
4403 crate::hybrid::Ffn::Dense {
4404 ffn_gate,
4405 ffn_up,
4406 ffn_down,
4407 } => {
4408 let n_ff = ffn_gate.out_features();
4409 if let Some((zq, zd)) = z_q8.as_ref() {
4410 // FUSED CHAIN (fix 2): pre-quantized z feeds the projections; the SwiGLU
4411 // epilogue emits act pre-quantized for ffn_down (silu_mul_scaled_q8_1,
4412 // bit-identical to silu_mul + quantize — kernel-check-pinned) with the
4413 // NVFP4 macro-scales folded (deferred-scale dual: y*s inline == the
4414 // scale_inplace store, value-exact) — the exact m=1 decode epilogue
4415 // structure at nrows=t.
4416 let pair =
4417 match e.matmul_decode_exact_dual_pre(ffn_gate, ffn_up, zq, zd, t)? {
4418 Some(((g, gs), (u, us))) => Some((g, gs, u, us)),
4419 None => None,
4420 };
4421 let (gate, gs, up, us) = match pair {
4422 Some(x4) => x4,
4423 None => (
4424 e.matmul_decode_exact_pre(ffn_gate, zq, zd, t)?,
4425 1.0, // scale already applied inside _pre
4426 e.matmul_decode_exact_pre(ffn_up, zq, zd, t)?,
4427 1.0,
4428 ),
4429 };
4430 if e.uses_q8_1_fast(ffn_down) {
4431 let (aq, ad) = e.silu_mul_scaled_q8_1(&gate, &up, gs, us, t * n_ff)?;
4432 e.matmul_decode_exact_pre(ffn_down, &aq, &ad, t)?
4433 } else {
4434 let mut act = vbuf(e, t * n_ff)?;
4435 e.silu_mul_scaled(&gate, &up, gs, us, &mut act, t * n_ff)?;
4436 e.matmul_decode_exact(ffn_down, &act, t)?
4437 }
4438 } else {
4439 // UNFUSED (pre-fix) chain — MoE-adjacent/M3/off-fast layers, unchanged.
4440 // DUAL gate+up batched twin (lane/verify-economics, 2026-08-02): one launch
4441 // for the pair at t=2..8 — bit-identical per (tensor,token,row) to the two
4442 // singles (kernel-check pins bitwise; MEMRA_SPEC_DUAL_T=0 reverts). None
4443 // (non-NVFP4 / t outside the tier / seam off) -> the two singles, unchanged.
4444 let (gate, up) =
4445 match e.matmul_decode_exact_dual(ffn_gate, ffn_up, &z, t)? {
4446 Some(pair) => pair,
4447 None => (
4448 e.matmul_decode_exact(ffn_gate, &z, t)?,
4449 e.matmul_decode_exact(ffn_up, &z, t)?,
4450 ),
4451 };
4452 let mut act = vbuf(e, t * n_ff)?; // fully written by ffn_act_lim
4453 Self::ffn_act_lim(
4454 e,
4455 &self.cfg,
4456 &gate,
4457 &up,
4458 1.0,
4459 1.0,
4460 dense_lim,
4461 &mut act,
4462 t * n_ff,
4463 )?;
4464 e.matmul_decode_exact(ffn_down, &act, t)?
4465 }
4466 }
4467 crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il(e, m, &z, t, il as u16)?,
4468 };
4469 // CROSS-LAYER fusion: defer this layer's post-FFN residual add — the next layer's
4470 // fused-q8 attn norm folds it in (add_rms_norm_q8_1 == add; rms_norm; quantize,
4471 // kernel-check-pinned at nrows=T). Non-fused next layers add explicitly above.
4472 pending = Some((x1, ffn_out));
4473 }
4474 // RANGE's final add (no next norm INSIDE the range to fuse with; for the
4475 // whole-trunk call that is the last layer, whose next norm is output_norm — f32-out).
4476 if let Some((x1p, f1p)) = pending.take() {
4477 let mut x2 = vbuf(e, t * n_embd)?; // fully written by add
4478 e.add(&x1p, &f1p, &mut x2, t * n_embd)?;
4479 x = x2;
4480 }
4481 Ok(x)
4482 }
4483 /// BATCHED linear-attn verify (T=K+1): the whole layer in ~10 launches instead of T x the
4484 /// T=1 decode chain (T x ~12 launches + T weight reads of the four projections). The GDN
4485 /// recurrence itself is inherently sequential — gdn_scan_s128 runs its internal t-loop with
4486 /// the SAME per-token math as chained T=1 calls (bit-identical state evolution); everything
4487 /// around it (projections, conv, prep, gated norm, out-proj) batches. Advances conv ring +
4488 /// ssm state exactly like T sequential decode steps.
4489 /// `want_stash`: additionally RETAIN the gdn-scan inputs (pure buffer keep-alives, zero extra
4490 /// kernels) so a partial accept can rebuild the state after any column prefix (REPLAY-FREE).
4491 #[allow(clippy::too_many_arguments)]
4492 fn linear_attn_verify_t(
4493 &self,
4494 e: &Engine,
4495 la: &LinearAttnLayer,
4496 h: &CudaSlice<f32>,
4497 h_q8: Option<(&CudaSlice<i8>, &CudaSlice<f32>)>,
4498 t: usize,
4499 cache: &mut Cache,
4500 il: usize,
4501 want_stash: bool,
4502 ) -> Result<(CudaSlice<f32>, Option<GdnStash>), Box<dyn std::error::Error>> {
4503 let cfg = &self.cfg;
4504 let ssm = cfg.ssm.as_ref().unwrap();
4505 let d_state = ssm.state_size as usize;
4506 let num_k = ssm.group_count as usize;
4507 let num_v = ssm.time_step_rank as usize;
4508 let d_conv = ssm.conv_kernel as usize;
4509 let key_dim = d_state * num_k;
4510 let conv_dim = key_dim * 2 + d_state * num_v;
4511 let eps = cfg.rms_eps;
4512 let scale = 1.0 / (d_state as f32).sqrt();
4513
4514 // DECODE-EXACT projections: matmul_decode_exact forces the MMVQ (warp-per-row, 32-thread)
4515 // accumulation order for EVERY m, matching the T=1 decode path bit-for-bit. The generic
4516 // `matmul` at m>=5 falls to dp4a (128-thread, two-level reduce) which has a different FP
4517 // sum order — ULP differences propagate through gdn_scan and flip argmax on the 27B.
4518 // Q8 TRUNK-FUSION at T=1 (35B: wqkv+wqkv_gate both Q8_0): one fused2 launch, bit-identical
4519 // per (tensor,row) to the two m=1 MMVQ dispatches below — decode-exact contract holds.
4520 // VERIFY-TIER TRUNK FUSION (MEMRA_SPEC_FUSED_T, t=2-4): quantize h ONCE for every
4521 // fused-eligible same-input Q8_0 pair of this layer (35B wqkv+wqkv_gate; 9B
4522 // ssm_beta+ssm_alpha) — each fused2 batched launch then replaces two decode-exact
4523 // calls (each of which re-quantizes the same h + runs its own _b2/_b4 launch).
4524 // Bit-identical per (tensor,token,row) — see spec_fused_t().
4525 // BATCHED EPILOGUE RE-FUSE (lane/vt-fixes fix 2): `h_q8` = the attn-input norm emitted
4526 // directly as q8_1 by the caller's fused rms_norm_q8_1 (bit-identical to the unfused
4527 // chain, kernel-check-pinned). When present it REPLACES the standalone quantize below
4528 // and feeds every projection; the caller guaranteed all four input projections are
4529 // q8_1-fast. When absent, the old shared-quantize (fused-t window) stands.
4530 let h_q8_t = if h_q8.is_none()
4531 && spec_fused_t()
4532 && (2..=4).contains(&t)
4533 && ((e.uses_q8_1_fast(&la.wqkv) && e.uses_q8_1_fast(&la.wqkv_gate))
4534 || (e.uses_q8_1_fast(&la.ssm_beta) && e.uses_q8_1_fast(&la.ssm_alpha)))
4535 {
4536 Some(e.quantize_q8_1(h, t, cfg.n_embd as usize)?)
4537 } else {
4538 None
4539 };
4540 // one view: the caller's fused-norm q8 or this fn's own shared quantize.
4541 let hq8_any: Option<(&CudaSlice<i8>, &CudaSlice<f32>)> =
4542 h_q8.or(h_q8_t.as_ref().map(|(q, d)| (q, d)));
4543 let (qkv_mixed, z) = {
4544 let mut fused = None;
4545 if t == 1 && e.uses_q8_1_fast(&la.wqkv) && e.uses_q8_1_fast(&la.wqkv_gate) {
4546 let (hq, hd) = e.quantize_q8_1(h, 1, cfg.n_embd as usize)?;
4547 fused = e.matmul_q8_fused2(&la.wqkv, &la.wqkv_gate, &hq, &hd)?;
4548 } else if let Some((hq, hd)) = hq8_any {
4549 if spec_fused_t() && (2..=4).contains(&t) {
4550 fused = e.matmul_q8_fused2_t(&la.wqkv, &la.wqkv_gate, hq, hd, t)?;
4551 }
4552 }
4553 match (fused, hq8_any) {
4554 (Some(pair), _) => pair,
4555 (None, Some((hq, hd))) if h_q8.is_some() => (
4556 e.matmul_decode_exact_pre(&la.wqkv, hq, hd, t)?,
4557 e.matmul_decode_exact_pre(&la.wqkv_gate, hq, hd, t)?,
4558 ),
4559 (None, _) => (
4560 e.matmul_decode_exact(&la.wqkv, h, t)?,
4561 e.matmul_decode_exact(&la.wqkv_gate, h, t)?,
4562 ),
4563 }
4564 };
4565 // beta+alpha DUAL at T=1 (75% of p3 rounds run T=1 verify — p-min chain cuts): the dual
4566 // mr2 kernel is bit-identical per element to the m=1 MMVQ matmul_decode_exact dispatches
4567 // (same warp-per-row body, blockIdx.y picks the weight), so the decode-exact contract
4568 // holds; the run-spec battery is the arbiter. T>1 keeps the per-tensor decode-exact path.
4569 let (beta_raw, alpha) = if t == 1 {
4570 let (hq, hd) = e.quantize_q8_1(h, 1, cfg.n_embd as usize)?;
4571 match e.matmul_pre_dual_noscale(&la.ssm_beta, &la.ssm_alpha, &hq, &hd, 1)? {
4572 Some(((mut b, bs), (mut a, as_))) => {
4573 if bs != 1.0 {
4574 e.scale_inplace(&mut b, bs, la.ssm_beta.out_features())?;
4575 }
4576 if as_ != 1.0 {
4577 e.scale_inplace(&mut a, as_, la.ssm_alpha.out_features())?;
4578 }
4579 (b, a)
4580 }
4581 // Q8_0 fused2 twin (9B stores beta/alpha as Q8_0): DISPATCH-MIRRORS the eager
4582 // decode's beta_alpha closure — the fused body is qmatvec_q8_0_mmvq verbatim,
4583 // bit-identical per row (kernel-check rel=0.00e0 gate), so decode==verify holds.
4584 None => match e.matmul_q8_fused2(&la.ssm_beta, &la.ssm_alpha, &hq, &hd)? {
4585 Some((b, a)) => (b, a),
4586 None => (
4587 e.matmul_decode_exact(&la.ssm_beta, h, 1)?,
4588 e.matmul_decode_exact(&la.ssm_alpha, h, 1)?,
4589 ),
4590 },
4591 }
4592 } else {
4593 // fused-t twin (9B stores beta/alpha as Q8_0): same shared-quantize + one launch
4594 // contract as the wqkv pair above; 35B beta/alpha are Float -> None -> fallback.
4595 let mut nvfp4_fused = None;
4596 let mut q8_fused = None;
4597 if let Some((hq, hd)) = hq8_any {
4598 if t == 3 && std::env::var("MEMRA_NVFP4_AUX_DUAL").as_deref() != Ok("0") {
4599 nvfp4_fused =
4600 e.matmul_decode_exact_dual_pre(&la.ssm_beta, &la.ssm_alpha, hq, hd, t)?;
4601 if nvfp4_fused.is_some() && std::env::var("MEMRA_DEBUG").is_ok() {
4602 static ONCE: std::sync::Once = std::sync::Once::new();
4603 ONCE.call_once(|| {
4604 eprintln!("[memra] NVFP4 beta+alpha batched aux dual ENGAGED (t={t})")
4605 });
4606 }
4607 }
4608 if nvfp4_fused.is_none() && spec_fused_t() && (2..=4).contains(&t) {
4609 q8_fused = e.matmul_q8_fused2_t(&la.ssm_beta, &la.ssm_alpha, hq, hd, t)?;
4610 }
4611 }
4612 if let Some(((mut b, bs), (mut a, as_))) = nvfp4_fused {
4613 if bs != 1.0 {
4614 e.scale_inplace(&mut b, bs, t * la.ssm_beta.out_features())?;
4615 }
4616 if as_ != 1.0 {
4617 e.scale_inplace(&mut a, as_, t * la.ssm_alpha.out_features())?;
4618 }
4619 (b, a)
4620 } else if let Some(pair) = q8_fused {
4621 pair
4622 } else {
4623 match hq8_any {
4624 Some((hq, hd)) if h_q8.is_some() => (
4625 e.matmul_decode_exact_pre(&la.ssm_beta, hq, hd, t)?,
4626 e.matmul_decode_exact_pre(&la.ssm_alpha, hq, hd, t)?,
4627 ),
4628 _ => (
4629 e.matmul_decode_exact(&la.ssm_beta, h, t)?,
4630 e.matmul_decode_exact(&la.ssm_alpha, h, t)?,
4631 ),
4632 }
4633 }
4634 };
4635
4636 // conv with CARRIED state + ring roll (T >= pad rides the input-column update kernel;
4637 // T < pad — the MEMRA_SPEC_M2 t=2 arm — rolls via the pure-copy ring rebuild).
4638 let rl = cache.recur[il].as_mut().unwrap();
4639 let mut conv_out = e.uninit(conv_dim * t)?;
4640 e.ssm_conv1d_tm_state(
4641 &qkv_mixed,
4642 &mut rl.conv_state,
4643 la.ssm_conv1d.float_data(),
4644 &mut conv_out,
4645 conv_dim,
4646 t,
4647 d_conv,
4648 )?;
4649
4650 // GDN prep via the prefill kernels (repack + L2 + sigmoid + glog), T-wide.
4651 let mut q_g = e.uninit(d_state * num_v * t)?;
4652 let mut k_g = e.uninit(d_state * num_v * t)?;
4653 let mut v_g = e.uninit(d_state * num_v * t)?;
4654 e.qkv_to_gdn_repack(
4655 &conv_out, &mut q_g, &mut k_g, &mut v_g, d_state, num_v, num_k, key_dim, t,
4656 )?;
4657 let mut q_l2 = e.uninit(d_state * num_v * t)?;
4658 e.l2_norm_decode(&q_g, &mut q_l2, d_state, num_v * t, eps)?;
4659 let mut k_l2 = e.uninit(d_state * num_v * t)?;
4660 e.l2_norm_decode(&k_g, &mut k_l2, d_state, num_v * t, eps)?;
4661 let mut beta = e.uninit(t * num_v)?;
4662 e.sigmoid(&beta_raw, &mut beta, t * num_v)?;
4663 let mut g_log = e.uninit(t * num_v)?;
4664 e.gdn_glog(
4665 &alpha,
4666 la.ssm_dt.float_data(),
4667 la.ssm_a.float_data(),
4668 &mut g_log,
4669 num_v,
4670 t,
4671 )?;
4672
4673 // ONE gdn_scan over T tokens from the carried state (internal sequential loop ==
4674 // T chained T=1 steps). Ping-pong the resident buffers like eager decode.
4675 let mut o = e.uninit(d_state * num_v * t)?;
4676 {
4677 let crate::cache::RecurLayer {
4678 ssm_state,
4679 ssm_state_alt,
4680 ..
4681 } = rl;
4682 e.gdn_scan_s128(
4683 &q_l2,
4684 &k_l2,
4685 &v_g,
4686 &g_log,
4687 &beta,
4688 ssm_state,
4689 ssm_state_alt,
4690 &mut o,
4691 num_v,
4692 t,
4693 scale,
4694 )?;
4695 }
4696 std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
4697
4698 // gated RMSNorm + out projection, T-wide. FUSED-QUANTIZE ARM (lane/vt-fixes fix 2,
4699 // mirroring the T=1 decode's launch-arc form): when ssm_out rides the q8_1 fast path,
4700 // emit q8_1 straight from the gated norm at nrows=num_v*t (row-indexed kernel, the
4701 // T-wide launch is the per-row program; kernel-check pins bit-identity vs
4702 // gated_rmsnorm -> quantize_q8_1 at T=1 and T=5) and feed the decode-exact dispatch
4703 // pre-quantized — one launch replaces norm + quantize. Fallback = the f32 chain.
4704 let out = if e.uses_q8_1_fast(&la.ssm_out) {
4705 let (gq, gd) =
4706 e.gated_rmsnorm_q8_1(&o, la.ssm_norm.float_data(), &z, d_state, num_v * t, eps)?;
4707 e.matmul_decode_exact_pre(&la.ssm_out, &gq, &gd, t)?
4708 } else {
4709 let mut gn = e.uninit(d_state * num_v * t)?;
4710 e.gated_rmsnorm(
4711 &o,
4712 la.ssm_norm.float_data(),
4713 &z,
4714 &mut gn,
4715 d_state,
4716 num_v * t,
4717 eps,
4718 )?;
4719 // DECODE-EXACT out-projection: same MMVQ path as the T=1 decode (ssm_out at m>=5
4720 // would fall to dp4a with a different FP reduction order — same class of bug as
4721 // the input projs).
4722 e.matmul_decode_exact(&la.ssm_out, &gn, t)?
4723 };
4724 let stash = if want_stash {
4725 Some(GdnStash {
4726 qkv_mixed,
4727 q_l2,
4728 k_l2,
4729 v_g,
4730 g_log,
4731 beta,
4732 })
4733 } else {
4734 None
4735 };
4736 Ok((out, stash))
4737 }
4738
4739 /// REPLAY-FREE partial-accept commit (2026-07-03): make the cache state == "committed through
4740 /// the first `j` verify columns" WITHOUT the legacy rollback + duplicate trunk replay.
4741 /// - Full-attn KV: truncate len to snapshot + j. The verify's appended rows for those columns
4742 /// are bit-identical to what an eager T=1 chain writes (the decode-exact contract the
4743 /// verify-probe gates), so keeping them == replaying them.
4744 /// - Linear layers, batched path: rebuild the conv ring by PURE COPIES (ring holds raw input
4745 /// columns) and the ssm state by a prefix re-run of the SAME gdn_scan kernel (t=j) from the
4746 /// snapshot state over the stash's identical inputs — the kernel's t-loop carries state in
4747 /// registers and writes it once at the end, so iterations 0..j-1 are independent of T:
4748 /// bit-identical to the verify's own state after j tokens == the eager chain state.
4749 /// - Linear layers, per-column path: restore the cloned actual state after column j-1.
4750 /// Caller guarantees 1 <= j <= t-1 (j==0 rounds take the legacy rollback; j==t is full accept).
4751 fn commit_verified_prefix(
4752 &self,
4753 e: &Engine,
4754 cache: &mut Cache,
4755 snap: &crate::cache::CacheSnapshot,
4756 ckpt: &VerifyCkpt,
4757 j: usize,
4758 kv_lens_done: bool,
4759 dev_j: Option<(&CudaSlice<u32>, usize, usize)>,
4760 ) -> Result<(), Box<dyn std::error::Error>> {
4761 let cfg = &self.cfg;
4762 let ssm = cfg.ssm.as_ref().unwrap();
4763 let d_state = ssm.state_size as usize;
4764 let num_k = ssm.group_count as usize;
4765 let num_v = ssm.time_step_rank as usize;
4766 let d_conv = ssm.conv_kernel as usize;
4767 let conv_dim = d_state * num_k * 2 + d_state * num_v;
4768 let scale = 1.0 / (d_state as f32).sqrt();
4769 for il in 0..self.layers.len() {
4770 if let (Some(kvl), Some(saved)) = (cache.kv[il].as_mut(), snap.kv_len[il]) {
4771 kvl.len = saved + j;
4772 // devacc 3a: spec_rollback_kv already wrote len_d on-device (same value).
4773 if !kv_lens_done {
4774 e.set_i32_one(&mut kvl.len_d, kvl.len as i32)?;
4775 }
4776 }
4777 if let Some(rl) = cache.recur[il].as_mut() {
4778 if let Some(st) = &ckpt.gdn[il] {
4779 let ring_old = snap.conv[il].as_ref().expect("snapshot missing conv");
4780 let state_in = snap.ssm[il].as_ref().expect("snapshot missing ssm");
4781 if let Some((acc, base, t_v)) = dev_j {
4782 // 3b: j read on-device (_dc twins, same bodies; full accept early-exits).
4783 e.ssm_conv_ring_rebuild_dc(
4784 &st.qkv_mixed,
4785 ring_old,
4786 &mut rl.conv_state,
4787 conv_dim,
4788 acc,
4789 base,
4790 t_v,
4791 d_conv,
4792 )?;
4793 let mut o = e.uninit(d_state * num_v * j.max(1))?;
4794 e.gdn_scan_s128_dc(
4795 &st.q_l2,
4796 &st.k_l2,
4797 &st.v_g,
4798 &st.g_log,
4799 &st.beta,
4800 state_in,
4801 &mut rl.ssm_state,
4802 &mut o,
4803 num_v,
4804 acc,
4805 base,
4806 t_v,
4807 scale,
4808 )?;
4809 } else {
4810 e.ssm_conv_ring_rebuild(
4811 &st.qkv_mixed,
4812 ring_old,
4813 &mut rl.conv_state,
4814 conv_dim,
4815 j,
4816 d_conv,
4817 )?;
4818 let mut o = e.uninit(d_state * num_v * j)?; // scan output, discarded
4819 e.gdn_scan_s128(
4820 &st.q_l2,
4821 &st.k_l2,
4822 &st.v_g,
4823 &st.g_log,
4824 &st.beta,
4825 state_in,
4826 &mut rl.ssm_state,
4827 &mut o,
4828 num_v,
4829 j,
4830 scale,
4831 )?;
4832 }
4833 } else if let Some(cols) = &ckpt.cols[il] {
4834 let (c, s) = &cols[j - 1];
4835 e.copy_into(&mut rl.conv_state, 0, c, c.len())?;
4836 e.copy_into(&mut rl.ssm_state, 0, s, s.len())?;
4837 } else {
4838 return Err(
4839 "commit_verified_prefix: verify ckpt missing for linear layer".into(),
4840 );
4841 }
4842 }
4843 }
4844 cache.pos = snap.pos + j;
4845 Ok(())
4846 }
4847
4848 /// ROUND-STREAM: recur restore with device-j (the _dc twins; full accept early-exits
4849 /// in-kernel). Requires the batched-linear stash on every linear layer (stream gate).
4850 fn commit_verified_prefix_stream(
4851 &self,
4852 e: &Engine,
4853 cache: &mut Cache,
4854 snap: &crate::cache::CacheSnapshot,
4855 ckpt: &VerifyCkpt,
4856 acc: &CudaSlice<u32>,
4857 base: usize,
4858 t_v: usize,
4859 ) -> Result<(), Box<dyn std::error::Error>> {
4860 let cfg = &self.cfg;
4861 let ssm = cfg.ssm.as_ref().unwrap();
4862 let d_state = ssm.state_size as usize;
4863 let num_k = ssm.group_count as usize;
4864 let num_v = ssm.time_step_rank as usize;
4865 let d_conv = ssm.conv_kernel as usize;
4866 let conv_dim = d_state * num_k * 2 + d_state * num_v;
4867 let scale = 1.0 / (d_state as f32).sqrt();
4868 for il in 0..self.layers.len() {
4869 if let Some(rl) = cache.recur[il].as_mut() {
4870 let st = ckpt.gdn[il]
4871 .as_ref()
4872 .ok_or("stream restore: batched-linear stash missing")?;
4873 let ring_old = snap.conv[il].as_ref().expect("snapshot missing conv");
4874 let state_in = snap.ssm[il].as_ref().expect("snapshot missing ssm");
4875 e.ssm_conv_ring_rebuild_dc(
4876 &st.qkv_mixed,
4877 ring_old,
4878 &mut rl.conv_state,
4879 conv_dim,
4880 acc,
4881 base,
4882 t_v,
4883 d_conv,
4884 )?;
4885 let mut o = e.uninit(d_state * num_v * t_v)?;
4886 e.gdn_scan_s128_dc(
4887 &st.q_l2,
4888 &st.k_l2,
4889 &st.v_g,
4890 &st.g_log,
4891 &st.beta,
4892 state_in,
4893 &mut rl.ssm_state,
4894 &mut o,
4895 num_v,
4896 acc,
4897 base,
4898 t_v,
4899 scale,
4900 )?;
4901 }
4902 }
4903 Ok(())
4904 }
4905
4906 /// EAGLE3 aux-capturing verify forward over `tokens` (T) — mirrors `decode_step_t_h` exactly
4907 /// (same KV append, same causal verify, same recur advance) but ALSO clones the aux residual-
4908 /// stream hiddens (blocks in `aux_layers`) for TWO columns: the LAST column (always) and the
4909 /// optional `pred_col` (the EAGLE seed = bonus's predecessor). Returns
4910 /// (all_T_logits host, last_col_aux, pred_col_aux?). Used by the EAGLE3 orchestrator's commit.
4911 pub fn decode_step_t_aux2(
4912 &self,
4913 e: &Engine,
4914 tokens: &[u32],
4915 pos0: usize,
4916 cache: &mut Cache,
4917 aux_layers: &[usize],
4918 pred_col: Option<usize>,
4919 ) -> Result<
4920 (Vec<f32>, Vec<CudaSlice<f32>>, Option<Vec<CudaSlice<f32>>>),
4921 Box<dyn std::error::Error>,
4922 > {
4923 let cfg = &self.cfg;
4924 let n_embd = cfg.n_embd as usize;
4925 let eps = cfg.rms_eps;
4926 let t = tokens.len();
4927 let pos_vec: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
4928 let pos_d = e.htod_i32(&pos_vec)?;
4929 let mut x = e.htod(&self.embd.gather(n_embd, tokens))?;
4930 let mut aux_last: Vec<CudaSlice<f32>> = Vec::with_capacity(aux_layers.len());
4931 let mut aux_pred: Vec<CudaSlice<f32>> = Vec::new();
4932 let want_pred = pred_col.is_some();
4933
4934 for (il, layer) in self.layers.iter().enumerate() {
4935 // DISPATCH-MIRRORED norms (FP-order lesson #8) — see decode_step_t_h_emb.
4936 let mixer_fast = self.mixer_in_q8_1_fast(e, &layer.mixer);
4937 let norm_fused = std::env::var("MEMRA_NO_FUSE_NORMQ").is_err() && mixer_fast;
4938 let mut h = vbuf(e, t * n_embd)?; // fully written by either rms_norm arm
4939 if norm_fused {
4940 e.rms_norm_decode(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
4941 } else {
4942 e.rms_norm(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
4943 }
4944 let mixed = match &layer.mixer {
4945 Mixer::Full(fa) => {
4946 self.full_attn_verify(e, fa, &h, None, &pos_d, t, cache, il, None)?
4947 }
4948 Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
4949 Mixer::Linear(la) => {
4950 let mut out = e.zeros(t * n_embd)?;
4951 for col in 0..t {
4952 let mut h_col = e.zeros(n_embd)?;
4953 let src = h.slice(col * n_embd..(col + 1) * n_embd);
4954 e.copy_view_into(&mut h_col, 0, &src, n_embd)?;
4955 let m_col = self.linear_attn_decode(e, la, &h_col, cache, il)?;
4956 e.copy_into(&mut out, col * n_embd, &m_col, n_embd)?;
4957 }
4958 out
4959 }
4960 };
4961 let ffn_fuse = match &layer.ffn {
4962 crate::hybrid::Ffn::Dense {
4963 ffn_gate, ffn_up, ..
4964 } => {
4965 std::env::var("MEMRA_NO_FUSE_NORMQ").is_err()
4966 && e.uses_q8_1_fast(ffn_gate)
4967 && e.uses_q8_1_fast(ffn_up)
4968 }
4969 crate::hybrid::Ffn::Moe(_) => false,
4970 };
4971 let mut x1 = vbuf(e, t * n_embd)?; // fully written by add / add_rms_norm
4972 let mut z = vbuf(e, t * n_embd)?; // fully written by rms_norm_decode / add_rms_norm
4973 if ffn_fuse {
4974 e.add(&x, &mixed, &mut x1, t * n_embd)?;
4975 e.rms_norm_decode(
4976 &x1,
4977 layer.post_attn_norm.float_data(),
4978 &mut z,
4979 n_embd,
4980 t,
4981 eps,
4982 )?;
4983 } else {
4984 e.add_rms_norm(
4985 &x,
4986 &mixed,
4987 layer.post_attn_norm.float_data(),
4988 &mut x1,
4989 &mut z,
4990 n_embd,
4991 t,
4992 eps,
4993 )?;
4994 }
4995 let ffn_out = match &layer.ffn {
4996 crate::hybrid::Ffn::Dense {
4997 ffn_gate,
4998 ffn_up,
4999 ffn_down,
5000 } => {
5001 let n_ff = ffn_gate.out_features();
5002 let gate = e.matmul_decode_exact(ffn_gate, &z, t)?;
5003 let up = e.matmul_decode_exact(ffn_up, &z, t)?;
5004 let mut act = vbuf(e, t * n_ff)?; // fully written by ffn_act_lim
5005 // dense FFN clamp = the SHEXP array (upstream build_ffn serves both).
5006 Self::ffn_act_lim(
5007 e,
5008 &self.cfg,
5009 &gate,
5010 &up,
5011 1.0,
5012 1.0,
5013 self.cfg.clamp_shexp_at(il as u32),
5014 &mut act,
5015 t * n_ff,
5016 )?;
5017 e.matmul_decode_exact(ffn_down, &act, t)?
5018 }
5019 crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il(e, m, &z, t, il as u16)?,
5020 };
5021 let mut x2 = vbuf(e, t * n_embd)?; // fully written by add
5022 e.add(&x1, &ffn_out, &mut x2, t * n_embd)?;
5023 if aux_layers.contains(&il) {
5024 let mut a = e.zeros(n_embd)?;
5025 e.copy_view_into(&mut a, 0, &x2.slice((t - 1) * n_embd..t * n_embd), n_embd)?;
5026 aux_last.push(a);
5027 if let Some(pc) = pred_col {
5028 let mut ap = e.zeros(n_embd)?;
5029 e.copy_view_into(
5030 &mut ap,
5031 0,
5032 &x2.slice(pc * n_embd..(pc + 1) * n_embd),
5033 n_embd,
5034 )?;
5035 aux_pred.push(ap);
5036 }
5037 }
5038 x = x2;
5039 }
5040 let mut hn = vbuf(e, t * n_embd)?; // fully written by rms_norm_decode
5041 e.rms_norm_decode(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
5042 let logits = e.matmul_decode_exact(&self.output, &hn, t)?;
5043 let host = e.dtoh(&logits)?;
5044 cache.pos += t;
5045 Ok((
5046 host,
5047 aux_last,
5048 if want_pred { Some(aux_pred) } else { None },
5049 ))
5050 }
5051
5052 /// step35 SPEC-VERIFY attention over T query tokens — a per-row REPLAY of the eager
5053 /// `step35_decode_attn`.
5054 ///
5055 /// WHY A REPLAY AND NOT A BATCHED TWIN. The verify's whole job is to be bit-identical to what
5056 /// the eager decode would have computed for the same tokens; that is what makes greedy spec
5057 /// decode exact (run-spec asserts token identity for K=1..8). Every other verify arm in this
5058 /// file earns that identity by carefully mirroring dispatch (`matmul_decode_exact` to force
5059 /// MMVQ at any m, per-layer `ffn_fuse` mirroring, per-row `fa_decode` key bounds). step35
5060 /// stacks FOUR more per-layer degrees of freedom on top of that — per-layer `n_head`
5061 /// (64 full / 96 SWA), per-layer rotary width (64 full / 128 SWA), per-layer rope base, and a
5062 /// SEPARATE `attn_gate` tensor whose projection shares the attn-normed input — and its SWA
5063 /// layers attend through a token-OFFSET view whose offset is a function of the ABSOLUTE
5064 /// position of each query row. A batched twin would have to reproduce all of that AND the
5065 /// per-row offset in one launch; the offset alone rules out the existing rows kernels (they
5066 /// take one `base_len`, not a per-row offset).
5067 ///
5068 /// So this arm calls the eager path itself, once per row, on the same cache. Identity is then
5069 /// true BY CONSTRUCTION rather than by mirroring: row r runs exactly the kernel sequence that
5070 /// eager decode step r runs (same projections, same q8_1 fusion decision, same append, same
5071 /// view arithmetic, same `fa_decode_kvmod`, same gate), because it IS that code. Cost: T x the
5072 /// eager decode mixer instead of one batched pass — the same trade the generic arm's `else`
5073 /// per-row loop already accepts when `fa_rows_eligible` says no. Correctness first; a batched
5074 /// step35 twin is a perf lane's job and must be gated against this arm.
5075 ///
5076 /// The `h_q8` pre-quantized pair from the caller's fused norm is NOT forwarded: it is a
5077 /// T-row buffer and `step35_decode_attn`'s `pre_q` contract is one row. Instead each row's
5078 /// f32 `h` slice is handed over and the callee re-derives its own q8_1 exactly as eager decode
5079 /// does (`quantize_q8_1(h, 1, n_embd)`) — which is the dispatch being mirrored. Callers that
5080 /// took the fused arm therefore MUST still pass a live `h`; `step35_verify` asserts that.
5081 #[allow(clippy::too_many_arguments)]
5082 fn step35_verify(
5083 &self,
5084 e: &Engine,
5085 fa: &FullAttnLayer,
5086 h: &CudaSlice<f32>,
5087 h_q8: Option<(&CudaSlice<i8>, &CudaSlice<f32>)>,
5088 t: usize,
5089 cache: &mut Cache,
5090 il: usize,
5091 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5092 let n_embd = self.cfg.n_embd as usize;
5093 // The fused (h-less) attn-norm arm hands `h` as a zero-length placeholder. This arm needs
5094 // the f32 rows, so the caller must not take that lever for step35 — enforced at the call
5095 // site by the `Mixer::Full(_) if self.cfg.step35.is_some() => false` arm of
5096 // `lin_q8_only` in `decode_step_t_core_stream`, and asserted here so a future caller
5097 // cannot regress it into silently reading an empty buffer.
5098 assert_eq!(
5099 h.len(),
5100 t * n_embd,
5101 "step35_verify needs the f32 attn-normed rows ([t*n_embd]); the caller took the \
5102 fused q8-only norm arm (h_q8={}) — step35 must stay on the unfused arm",
5103 h_q8.is_some()
5104 );
5105 // ROW WIDTH IS n_embd, NOT n_head*head_dim: `step35_decode_attn` returns the mixer output
5106 // AFTER `wo`, so a row is [n_embd] — the same contract the generic arm's
5107 // `matmul_decode_exact(&fa.wo, &attn_g, t)` return has. Sizing this buffer from the
5108 // per-layer head geometry (8192 on full-attn, 12288 on SWA) instead overran the row on the
5109 // FIRST copy and panicked inside `copy_into`'s `CudaView::slice` unwrap
5110 // (raw/mtp-bt-20260806T212127Z.log frames 12-13).
5111 let mut out = vbuf(e, t * n_embd)?; // each row fully written by the copy below
5112 for r in 0..t {
5113 // Absolute position of this query row. `cache.pos` is the committed length at round
5114 // start and every row before r has already been appended by this loop, so the r-th
5115 // verify token sits at cache.pos + r — the same position eager decode would give it.
5116 let pos_d = e.htod_i32(&[(cache.pos + r) as i32])?;
5117 let mut h_row = vbuf(e, n_embd)?; // fully written by copy_view_into
5118 e.copy_view_into(
5119 &mut h_row,
5120 0,
5121 &h.slice(r * n_embd..(r + 1) * n_embd),
5122 n_embd,
5123 )?;
5124 // THE eager decode mixer: appends this row's K/V at kvl.len, advances it, then
5125 // attends over the (SWA-offset) view. Post-`wo`, same contract as this fn returns.
5126 let o = self.step35_decode_attn(e, fa, il, &h_row, None, &pos_d, cache)?;
5127 debug_assert_eq!(
5128 o.len(),
5129 n_embd,
5130 "step35_decode_attn returns post-wo [n_embd]"
5131 );
5132 e.copy_into(&mut out, r * n_embd, &o, n_embd)?;
5133 }
5134 Ok(out)
5135 }
5136
5137 /// Full-attention mixer over T query tokens with a GROWING resident KV (verify path, §D.3).
5138 /// Appends the T new K/V columns to cache.kv[il] then attends causally over [0..len) via
5139 /// fa_prefill. Token-major [T, kv_dim] projection layout == cache row layout (single copy).
5140 #[allow(clippy::too_many_arguments)]
5141 fn full_attn_verify(
5142 &self,
5143 e: &Engine,
5144 fa: &FullAttnLayer,
5145 h: &CudaSlice<f32>,
5146 h_q8: Option<(&CudaSlice<i8>, &CudaSlice<f32>)>,
5147 pos_d: &CudaSlice<i32>,
5148 t: usize,
5149 cache: &mut Cache,
5150 il: usize,
5151 stream_ctr: Option<&CudaSlice<i32>>,
5152 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5153 // step35: the generic geometry below is wrong for this arch (per-layer n_head, partial
5154 // per-layer rope, the SWA offset view, and a SEPARATE head-wise gate tensor), so it takes
5155 // its own arm. A verify that silently computes different attention than decode defeats the
5156 // whole self-consistency gate, so the arm is a per-row REPLAY of `step35_decode_attn`
5157 // rather than a batched twin — see `step35_verify` for why that is the exactness-correct
5158 // shape and not laziness.
5159 if self.cfg.step35.is_some() {
5160 if stream_ctr.is_some() {
5161 return Err(
5162 "step35 has no ROUND-STREAM verify arm (the device-counter _dc twins \
5163 cannot express the SWA offset KV view; same root cause as the dc \
5164 decode refusal) — run spec without the stream arm"
5165 .into(),
5166 );
5167 }
5168 return self.step35_verify(e, fa, h, h_q8, t, cache, il);
5169 }
5170 let cfg = &self.cfg;
5171 let geometry = cfg.full_attention_geometry_at(il as u32);
5172 let n_head = geometry.n_head as usize;
5173 let n_head_kv = geometry.n_head_kv as usize;
5174 let head_dim = geometry.head_dim_k as usize;
5175 let eps = cfg.rms_eps;
5176 let scale = geometry.attention_scale();
5177 let n_embd = cfg.n_embd as usize;
5178
5179 // DECODE-EXACT Q/K/V projections: matmul_decode_exact forces the MMVQ (warp-per-row) path
5180 // for every m, matching the T=1 decode's FP accumulation order. matmul_pre at m>=5 would
5181 // fall to dp4a (128-thread, two-level reduce) with a different FP sum order.
5182 // Q8 TRUNK-FUSION at T=1: DISPATCH-MIRRORS the eager decode's fused3 (bit-identical body).
5183 // BATCHED EPILOGUE RE-FUSE (lane/vt-fixes fix 2): `h_q8` = the attn-input norm's q8_1
5184 // form emitted by the fused rms_norm_q8_1 (bit-identical to rms_norm_decode ->
5185 // quantize_q8_1, kernel-check-pinned). When present (caller checked mixer q8_1-fast),
5186 // every projection consumes it — `h` may be a zero-len placeholder and must not be read.
5187 let (qf, mut k, v) = {
5188 let mut fused = None;
5189 let qkv_fast =
5190 e.uses_q8_1_fast(&fa.wq) && e.uses_q8_1_fast(&fa.wk) && e.uses_q8_1_fast(&fa.wv);
5191 if t == 1 && qkv_fast {
5192 let (hq_o, hd_o);
5193 let (hq, hd): (&CudaSlice<i8>, &CudaSlice<f32>) = match h_q8 {
5194 Some(p) => p,
5195 None => {
5196 (hq_o, hd_o) = e.quantize_q8_1(h, 1, n_embd)?;
5197 (&hq_o, &hd_o)
5198 }
5199 };
5200 fused = e.matmul_q8_fused3(&fa.wq, &fa.wk, &fa.wv, hq, hd)?;
5201 } else if spec_fused_t() && (2..=4).contains(&t) && qkv_fast {
5202 // VERIFY-TIER TRUNK FUSION (MEMRA_SPEC_FUSED_T): one shared quantize + one
5203 // fused3 batched launch replaces three decode-exact calls (3 re-quantizes of
5204 // the same h + 3 _b2/_b4 launches). Bit-identical per (tensor,token,row).
5205 let (hq_o, hd_o);
5206 let (hq, hd): (&CudaSlice<i8>, &CudaSlice<f32>) = match h_q8 {
5207 Some(p) => p,
5208 None => {
5209 (hq_o, hd_o) = e.quantize_q8_1(h, t, n_embd)?;
5210 (&hq_o, &hd_o)
5211 }
5212 };
5213 fused = e.matmul_q8_fused3_t(&fa.wq, &fa.wk, &fa.wv, hq, hd, t)?;
5214 }
5215 match (fused, h_q8) {
5216 (Some(triple), _) => triple,
5217 // shared pre-quantized activation (q8_1-fast guaranteed by the caller): the
5218 // decode-exact dispatch consumes (hq, hd) instead of re-quantizing 3x.
5219 (None, Some((hq, hd))) if qkv_fast => (
5220 e.matmul_decode_exact_pre(&fa.wq, hq, hd, t)?,
5221 e.matmul_decode_exact_pre(&fa.wk, hq, hd, t)?,
5222 e.matmul_decode_exact_pre(&fa.wv, hq, hd, t)?,
5223 ),
5224 (None, _) => (
5225 e.matmul_decode_exact(&fa.wq, h, t)?,
5226 e.matmul_decode_exact(&fa.wk, h, t)?,
5227 e.matmul_decode_exact(&fa.wv, h, t)?,
5228 ),
5229 }
5230 };
5231 // M3/Hy3 have no attention output gate — wq out is exactly q; skip the split.
5232 let gated = geometry.attention_gate == memra_gguf::config::AttentionGateKind::FusedQ;
5233 let (mut q, gate) = if gated {
5234 let mut q = vbuf(e, t * n_head * head_dim)?; // fully written by q_gate_split
5235 let mut gate = vbuf(e, t * n_head * head_dim)?; // fully written by q_gate_split
5236 e.q_gate_split(&qf, &mut q, &mut gate, head_dim, n_head, t)?;
5237 (q, Some(gate))
5238 } else {
5239 (qf, None)
5240 };
5241
5242 let mut qn = vbuf(e, t * n_head * head_dim)?; // fully written by rms_norm
5243 e.rms_norm(
5244 &q,
5245 fa.q_norm.float_data(),
5246 &mut qn,
5247 head_dim,
5248 n_head * t,
5249 eps,
5250 )?;
5251 q = qn;
5252 let mut kn = vbuf(e, t * n_head_kv * head_dim)?; // fully written by rms_norm
5253 e.rms_norm(
5254 &k,
5255 fa.k_norm.float_data(),
5256 &mut kn,
5257 head_dim,
5258 n_head_kv * t,
5259 eps,
5260 )?;
5261 k = kn;
5262 let rope_dims = geometry.n_rot as usize;
5263 e.rope_neox(
5264 &mut q,
5265 pos_d,
5266 head_dim,
5267 rope_dims,
5268 n_head,
5269 t,
5270 geometry.rope_base,
5271 1.0,
5272 )?;
5273 e.rope_neox(
5274 &mut k,
5275 pos_d,
5276 head_dim,
5277 rope_dims,
5278 n_head_kv,
5279 t,
5280 geometry.rope_base,
5281 1.0,
5282 )?;
5283
5284 // append T new K/V columns to the resident QUANTIZED cache. k/v are token-major [T, kv_dim]
5285 // f32; append-quantize each of the T token rows into the byte cache (q8_0 K / q5_1 V).
5286 let kvl = cache.kv[il].as_mut().unwrap();
5287 let (kv_dim_k, kv_dim_v, ktb, vtb) =
5288 (kvl.kv_dim_k, kvl.kv_dim_v, kvl.k_tok_bytes, kvl.v_tok_bytes);
5289 if let Some(ctr) = stream_ctr {
5290 // stream: ONE batched append at the device counter (rows kernel = the per-view warp
5291 // math on a (block, token) grid, documented byte-identical); host len is a stale
5292 // LOWER BOUND under pre-issue (drain reconciles it).
5293 e.append_kv_quantized_rows_dc(
5294 &k,
5295 &v,
5296 &mut kvl.k,
5297 &mut kvl.v,
5298 ctr,
5299 t,
5300 kv_dim_k,
5301 kv_dim_v,
5302 ktb,
5303 vtb,
5304 crate::Engine::kv_fp8_on(),
5305 )?;
5306 } else {
5307 for i in 0..t {
5308 let k_row = k.slice(i * kv_dim_k..(i + 1) * kv_dim_k);
5309 let v_row = v.slice(i * kv_dim_v..(i + 1) * kv_dim_v);
5310 e.append_kv_quantized_view(
5311 &k_row,
5312 &v_row,
5313 &mut kvl.k,
5314 &mut kvl.v,
5315 kvl.len + i,
5316 kv_dim_k,
5317 kv_dim_v,
5318 ktb,
5319 vtb,
5320 crate::Engine::kv_fp8_on(),
5321 )?;
5322 }
5323 kvl.len += t;
5324 }
5325
5326 // BIT-IDENTICAL VERIFY ATTENTION (spec-exactness fix): the FP accumulation order must be
5327 // byte-for-byte identical to the eager decode path. fa_prefill uses a different tile size
5328 // (BLOCK_Q=64, BK=32) and online-softmax structure than fa_decode's split-K + combine,
5329 // which changes FP summation order and can flip argmax at tight logit margins. Query row r
5330 // attends to keys [0..base_len+r+1) — each successive row sees one more key (the causal
5331 // property). This matches eager: decode appends k at len, then fa_decode sees t_kv = len+1
5332 // keys. The verify appends all T tokens first but bounds the key range per row.
5333 //
5334 // MULTI-ROW FUSED PATH (the long-ctx spec fix, 2026-07-03): when every row takes the vec
5335 // kernel (base_len+1 >= FA_VEC_MIN_TKV), ONE fa_decode_rows launch executes the exact
5336 // per-row program for all T rows (grid.z = row, per-row n_splits from the same
5337 // fa_split_keys formula) — replacing T x (2 launches + 2 dtod copies + 5 partial allocs)
5338 // and multiplying resident CTAs by T on a latency-bound kernel. Bit-identical per row by
5339 // construction; kernel-check pins rows-vs-loop byte identity, run-spec is the end gate.
5340 // Short ctx (any row below the vec crossover) and MEMRA_NO_FA_VEC/MEMRA_FA_ROWS_OFF keep the
5341 // per-row loop (whose fa_decode picks scalar/vec per row exactly like eager decode).
5342 let mut attn = vbuf(e, t * n_head * head_dim)?; // fully written by every FA arm below
5343 let base_len = kvl.len - t; // KV len BEFORE this round's T tokens were appended
5344 // T=1 INCLUDED (2026-07-05): p-min cuts the draft to 1 in ~75% of rounds on hard
5345 // (agentic) content — the old t>1 gate sent those rounds to the per-row loop (262us/row
5346 // + q-row copy + per-row allocs vs 93us/row through the fused kernel at grid.z=1, same
5347 // program). nsys accounting: 1088 of 1456 verify FA launches were T=1 escapees.
5348 // LEAN T=1 ARM (MEMRA_SPEC_LEAN, close35): at t==1, q IS one row and fa_decode on it is
5349 // the EXACT eager decode dispatch (vec_q_v2 + combine_f32; the rows pair measured +50us
5350 // at m=1). Byte-identical: kernel-check pins rows-vs-loop identity, and the per-row loop
5351 // at t=1 is fa_decode on the same q with zero-offset copies. Gates arbitrate.
5352 if let Some(ctr) = stream_ctr {
5353 // STREAM ARM: causal base from the device counter; host kvl.len is a stale lower
5354 // bound used only for the split-sizing upper bound (+64 slack covers M pre-issued
5355 // rounds at K<=8). Views span the bound; per-row limits derive in-kernel.
5356 let upper = kvl.len + t + 64;
5357 let k_view = e.view_u8(&kvl.k, (upper.min(cache.max_ctx)) * ktb);
5358 let v_view = e.view_u8(&kvl.v, (upper.min(cache.max_ctx)) * vtb);
5359 e.fa_decode_rows_dc(
5360 &q,
5361 &k_view,
5362 &v_view,
5363 &mut attn,
5364 head_dim,
5365 n_head,
5366 n_head_kv,
5367 ctr,
5368 upper.min(cache.max_ctx),
5369 t,
5370 scale,
5371 ktb,
5372 vtb,
5373 0,
5374 false,
5375 )?;
5376 } else if spec_lean() && t == 1 {
5377 let t_kv = base_len + 1;
5378 let k_view = e.view_u8(&kvl.k, t_kv * ktb);
5379 let v_view = e.view_u8(&kvl.v, t_kv * vtb);
5380 e.fa_decode_kvmod(
5381 &q,
5382 &k_view,
5383 &v_view,
5384 &mut attn,
5385 head_dim,
5386 n_head,
5387 n_head_kv,
5388 t_kv,
5389 scale,
5390 ktb,
5391 vtb,
5392 crate::Engine::kv_fp8_on(),
5393 )?;
5394 } else if e.fa_rows_eligible(base_len, head_dim) {
5395 let k_view = e.view_u8(&kvl.k, (base_len + t) * ktb);
5396 let v_view = e.view_u8(&kvl.v, (base_len + t) * vtb);
5397 e.fa_decode_rows(
5398 &q,
5399 &k_view,
5400 &v_view,
5401 &mut attn,
5402 head_dim,
5403 n_head,
5404 n_head_kv,
5405 base_len,
5406 t,
5407 scale,
5408 ktb,
5409 vtb,
5410 None,
5411 false,
5412 crate::Engine::kv_fp8_on(),
5413 None,
5414 )?;
5415 } else {
5416 for r in 0..t {
5417 let t_kv_r = base_len + r + 1; // this row sees keys [0..t_kv_r)
5418 let k_view_r = e.view_u8(&kvl.k, t_kv_r * ktb);
5419 let v_view_r = e.view_u8(&kvl.v, t_kv_r * vtb);
5420 // copy q row into an owned buffer (fa_decode takes &CudaSlice, not CudaView)
5421 let mut q_row = vbuf(e, n_head * head_dim)?; // fully written by copy_view_into
5422 let q_src = q.slice(r * n_head * head_dim..(r + 1) * n_head * head_dim);
5423 e.copy_view_into(&mut q_row, 0, &q_src, n_head * head_dim)?;
5424 let mut attn_row = vbuf(e, n_head * head_dim)?; // fully written by fa_decode
5425 e.fa_decode_kvmod(
5426 &q_row,
5427 &k_view_r,
5428 &v_view_r,
5429 &mut attn_row,
5430 head_dim,
5431 n_head,
5432 n_head_kv,
5433 t_kv_r,
5434 scale,
5435 ktb,
5436 vtb,
5437 crate::Engine::kv_fp8_on(),
5438 )?;
5439 e.copy_into(
5440 &mut attn,
5441 r * n_head * head_dim,
5442 &attn_row,
5443 n_head * head_dim,
5444 )?;
5445 }
5446 }
5447
5448 let attn_g = match &gate {
5449 Some(gate) => {
5450 let mut gsig = vbuf(e, t * n_head * head_dim)?; // fully written by sigmoid
5451 e.sigmoid(gate, &mut gsig, t * n_head * head_dim)?;
5452 let mut ag = vbuf(e, t * n_head * head_dim)?; // fully written by mul
5453 e.mul(&attn, &gsig, &mut ag, t * n_head * head_dim)?;
5454 ag
5455 }
5456 None => attn,
5457 };
5458 // DECODE-EXACT wo projection: at m>=5 (K=4+ with pending) the generic matmul would use dp4a
5459 // (128-thread, different FP sum order than MMVQ). Force MMVQ for bit-identity with decode.
5460 Ok(e.matmul_decode_exact(&fa.wo, &attn_g, t)?)
5461 }
5462
5463 /// Context-linear bytes for a plain serving session's trunk cache.
5464 pub fn plain_session_kv_bytes_per_token(&self) -> usize {
5465 crate::cache::cache_bytes_per_token(&self.cfg)
5466 }
5467
5468 /// `(logical bytes/token, ring-capped bytes/token, ring row cap)` for exact admission.
5469 pub fn plain_session_kv_shape(&self) -> (usize, usize, usize) {
5470 (
5471 self.plain_session_kv_bytes_per_token(),
5472 crate::cache::cache_ring_bytes_per_token(&self.cfg),
5473 crate::cache::cache_ring_row_cap(&self.cfg),
5474 )
5475 }
5476
5477 /// Context-linear bytes for a speculative serving session: trunk cache plus persistent MTP
5478 /// scratch. With no MTP head this equals the plain coefficient.
5479 pub fn spec_session_kv_bytes_per_token(&self) -> usize {
5480 let scratch = self
5481 .mtp
5482 .as_ref()
5483 .map(|mtp| {
5484 let (_, _, k, v) = mtp_scratch_layout(&self.cfg, mtp.geom.as_ref());
5485 k + v
5486 })
5487 .unwrap_or(0);
5488 self.plain_session_kv_bytes_per_token()
5489 .saturating_add(scratch)
5490 }
5491
5492 /// Spec twin of [`HybridModel::plain_session_kv_shape`]; Step35's persistent MTP scratch is
5493 /// capped by the same SWA ring rows as the trunk.
5494 pub fn spec_session_kv_shape(&self) -> (usize, usize, usize) {
5495 let total = self.spec_session_kv_bytes_per_token();
5496 let (_, mut ring, rows) = self.plain_session_kv_shape();
5497 if rows > 0 {
5498 ring = ring.saturating_add(
5499 self.mtp
5500 .as_ref()
5501 .map(|mtp| {
5502 let (_, _, k, v) = mtp_scratch_layout(&self.cfg, mtp.geom.as_ref());
5503 k + v
5504 })
5505 .unwrap_or(0),
5506 );
5507 }
5508 (total, ring, rows)
5509 }
5510
5511 /// Greedy MTP speculative decode (§B). Token-identical to `generate(prompt, max_new)` but uses
5512 /// the NextN head to draft K tokens then verifies them in one batched target forward.
5513 /// Returns (generated tokens, total_drafted, total_accepted) so the caller can report
5514 /// acceptance rate. `k` = draft length per round.
5515 ///
5516 /// GRAPH DRAFT (stage 2 of graph-grade spec): when the model is all-Dense and the MTP head is
5517 /// Dense (no MoE host readbacks), the fixed-shape T=1 MTP forward is CUDA-graph-captured ONCE
5518 /// and replayed per draft step — the ~40 eager launches per drafted token collapse into one
5519 /// graph dispatch; only the 4-byte token id (and 4-byte p-min confidence) round-trip per step.
5520 /// Event tracking is disabled for the whole call (generate_graph pattern) so every buffer the
5521 /// captured graph references is event-free; the spec loop is strictly single-stream.
5522 /// MEMRA_SPEC_NOGRAPH=1 forces the eager draft chain.
5523 /// SAMPLED mode (MEMRA_SPEC_TEMP>0) has its OWN capture (gumbel-perturbed in-graph argmax,
5524 /// device Philox event counter, persistent q retention) — graph-vs-eager sampled streams are
5525 /// bit-identical for the same (seed, prompt, K, temp); see the sampled-graph setup in
5526 /// generate_spec_inner2.
5527 /// Multi-turn session: trunk cache + MTP draft scratch persist across generate calls, so
5528 /// turn N+1 primes ONLY its new suffix (the 124k-conversation daily pattern — re-priming a
5529 /// 32k history costs ~54s; a suffix prime costs seconds). APPEND-ONLY by construction: the
5530 /// hybrid linear-attn states are in-place (no position index), so a session can extend but
5531 /// never rewind — `committed` is the exact token list whose state the caches hold (includes
5532 /// any overshoot tokens past max_new; the caller renders from `committed`, not its own echo).
5533 pub fn new_session(
5534 &self,
5535 e: &Engine,
5536 max_ctx: usize,
5537 ) -> Result<SpecSession, Box<dyn std::error::Error>> {
5538 Ok(SpecSession {
5539 // STAGE-OWNED KV (lane/pp2-spec 2026-08-06): `pp::new_cache`, not `Cache::new`. This
5540 // is the SERVING spec-session path, and with the ppN door open across two cards a
5541 // primary-homed cache makes every remote stage peer-read its OWN KV on every verify
5542 // round — the wrong-card class already fixed on the two batched serving paths
5543 // (worker.rs 2483 / 2837). With the door shut `new_cache` IS `Cache::new` (same
5544 // branch, same allocations), so single-device behavior is byte-unchanged.
5545 cache: crate::pp::new_cache(e, &self.cfg, max_ctx)?,
5546 scratch: MtpScratch::new(
5547 e,
5548 &self.cfg,
5549 max_ctx,
5550 self.mtp.as_ref().and_then(|m| m.geom.as_ref()),
5551 )?,
5552 committed: Vec::new(),
5553 last_h: None,
5554 next_pred: None,
5555 sctr: 0,
5556 uctr: 0,
5557 draft_ctx: None,
5558 pending_tok: None,
5559 turn_ckpt: None,
5560 telem: SpecTelemetryCounters::default(),
5561 capture_at: None,
5562 boundary_capture: None,
5563 })
5564 }
5565
5566 /// Forced-gate exact state comparison. This intentionally reads the real live prefixes from
5567 /// their owning PP devices: matching emitted ids alone would miss a stale `len_d`, recurrent
5568 /// snapshot, or draft-KV row that only corrupts the following round.
5569 pub fn optipipe_compare_session_state(
5570 &self,
5571 e: &Engine,
5572 reference: &SpecSession,
5573 candidate: &SpecSession,
5574 ) -> Result<OptiForkStateIdentity, Box<dyn std::error::Error>> {
5575 fn fail(what: &str) -> Box<dyn std::error::Error> {
5576 format!("optipipe state mismatch: {what}").into()
5577 }
5578 fn same_f32(a: &[f32], b: &[f32]) -> bool {
5579 a.len() == b.len() && a.iter().zip(b).all(|(x, y)| x.to_bits() == y.to_bits())
5580 }
5581 fn compare_layers(
5582 es: &Engine,
5583 range: std::ops::Range<usize>,
5584 reference: &SpecSession,
5585 candidate: &SpecSession,
5586 report: &mut OptiForkStateIdentity,
5587 ) -> Result<(), Box<dyn std::error::Error>> {
5588 for il in range {
5589 match (&reference.cache.kv[il], &candidate.cache.kv[il]) {
5590 (Some(a), Some(b)) => {
5591 if a.len != b.len {
5592 return Err(fail(&format!(
5593 "layer {il} host KV len {} != {}",
5594 a.len, b.len
5595 )));
5596 }
5597 let ad = es.dtoh_i32(&a.len_d)?;
5598 let bd = es.dtoh_i32(&b.len_d)?;
5599 if ad != bd || ad.first().copied() != Some(a.len as i32) {
5600 return Err(fail(&format!(
5601 "layer {il} device KV len {ad:?} != {bd:?} (host={})",
5602 a.len,
5603 )));
5604 }
5605 let kb = a.len * a.k_tok_bytes;
5606 let vb = a.len * a.v_tok_bytes;
5607 if kb > 0 {
5608 let ak = es.dtoh_u8_view(&a.k.slice(0..kb))?;
5609 let bk = es.dtoh_u8_view(&b.k.slice(0..kb))?;
5610 if ak != bk {
5611 let at = ak.iter().zip(&bk).position(|(x, y)| x != y).unwrap();
5612 return Err(fail(&format!(
5613 "layer {il} K bytes at byte {at} row {} offset {}: {} != {}",
5614 at / a.k_tok_bytes,
5615 at % a.k_tok_bytes,
5616 ak[at],
5617 bk[at],
5618 )));
5619 }
5620 }
5621 if vb > 0 {
5622 let av = es.dtoh_u8_view(&a.v.slice(0..vb))?;
5623 let bv = es.dtoh_u8_view(&b.v.slice(0..vb))?;
5624 if av != bv {
5625 let at = av.iter().zip(&bv).position(|(x, y)| x != y).unwrap();
5626 return Err(fail(&format!(
5627 "layer {il} V bytes at byte {at} row {} offset {}: {} != {}",
5628 at / a.v_tok_bytes,
5629 at % a.v_tok_bytes,
5630 av[at],
5631 bv[at],
5632 )));
5633 }
5634 }
5635 report.trunk_kv_bytes += kb + vb;
5636 }
5637 (None, None) => {}
5638 _ => return Err(fail(&format!("layer {il} KV presence"))),
5639 }
5640 match (&reference.cache.recur[il], &candidate.cache.recur[il]) {
5641 (Some(a), Some(b)) => {
5642 let ac = es.dtoh(&a.conv_state)?;
5643 let bc = es.dtoh(&b.conv_state)?;
5644 if !same_f32(&ac, &bc) {
5645 return Err(fail(&format!("layer {il} conv state")));
5646 }
5647 let as_ = es.dtoh(&a.ssm_state)?;
5648 let bs = es.dtoh(&b.ssm_state)?;
5649 if !same_f32(&as_, &bs) {
5650 return Err(fail(&format!("layer {il} SSM state")));
5651 }
5652 report.recurrent_bytes += (ac.len() + as_.len()) * 4;
5653 }
5654 (None, None) => {}
5655 _ => return Err(fail(&format!("layer {il} recurrent presence"))),
5656 }
5657 }
5658 Ok(())
5659 }
5660
5661 if reference.committed != candidate.committed {
5662 return Err(fail("committed token ids"));
5663 }
5664 if reference.cache.pos != candidate.cache.pos
5665 || reference.cache.max_ctx != candidate.cache.max_ctx
5666 {
5667 return Err(fail("cache pos/capacity"));
5668 }
5669 if reference.pending_tok != candidate.pending_tok
5670 || reference.next_pred != candidate.next_pred
5671 || reference.sctr != candidate.sctr
5672 || reference.uctr != candidate.uctr
5673 {
5674 return Err(fail("pending/prediction/counter tail"));
5675 }
5676
5677 let mut report = OptiForkStateIdentity::default();
5678 if let Some(fence) = crate::pp::pp_cuts(self.layers.len()) {
5679 let rt = crate::pp::PpNRt::get(e)?;
5680 for stage in 0..rt.n_stages() {
5681 let _scope = rt.enter(stage);
5682 compare_layers(
5683 rt.engine(stage, e),
5684 fence[stage]..fence[stage + 1],
5685 reference,
5686 candidate,
5687 &mut report,
5688 )?;
5689 }
5690 } else {
5691 compare_layers(e, 0..self.layers.len(), reference, candidate, &mut report)?;
5692 }
5693
5694 let (a, b) = (&reference.scratch.kv, &candidate.scratch.kv);
5695 if a.len != b.len || e.dtoh_i32(&a.len_d)? != e.dtoh_i32(&b.len_d)? {
5696 return Err(fail("draft scratch length"));
5697 }
5698 let kb = a.len * a.k_tok_bytes;
5699 let vb = a.len * a.v_tok_bytes;
5700 if kb > 0 && e.dtoh_u8_view(&a.k.slice(0..kb))? != e.dtoh_u8_view(&b.k.slice(0..kb))? {
5701 return Err(fail("draft scratch K bytes"));
5702 }
5703 if vb > 0 && e.dtoh_u8_view(&a.v.slice(0..vb))? != e.dtoh_u8_view(&b.v.slice(0..vb))? {
5704 return Err(fail("draft scratch V bytes"));
5705 }
5706 report.scratch_kv_bytes = kb + vb;
5707
5708 match (&reference.last_h, &candidate.last_h) {
5709 (Some(a), Some(b)) => {
5710 let ah = e.dtoh(a)?;
5711 let bh = e.dtoh(b)?;
5712 if !same_f32(&ah, &bh) {
5713 return Err(fail("last hidden/seed bytes"));
5714 }
5715 report.hidden_bytes = ah.len() * 4;
5716 }
5717 (None, None) => {}
5718 _ => return Err(fail("last hidden/seed presence")),
5719 }
5720 Ok(report)
5721 }
5722
5723 /// SESSION-AFFINITY REWIND (lane/session-affinity, 2026-08-05): roll `sess` back to its
5724 /// retained prompt-end checkpoint, so a request whose prompt matches
5725 /// `committed[..rewind_pos()]` exactly can resume there and prime only its own delta.
5726 ///
5727 /// EXACTNESS. After this returns, the session is byte-for-byte the state it was in AT that
5728 /// boundary: full-attn KV truncated to it (append-only, position-addressed), GDN conv/ssm
5729 /// restored from the device copy taken there, draft scratch length reset, `committed`
5730 /// truncated, `last_h` = the boundary's predecessor anchor. That is precisely the state a
5731 /// fresh prime of `committed[..pos]` would have produced, so the following suffix prime and
5732 /// every burst after it are identical to a cold run of the same token stream — the
5733 /// committed-tokens-authoritative contract.
5734 ///
5735 /// `next_pred` and `pending_tok` are CLEARED: both describe generation past the boundary,
5736 /// which the rewind discards. The caller therefore must supply a non-empty suffix (a
5737 /// rewound session cannot serve an empty-suffix continuation burst — there is nothing to
5738 /// continue). The persistent draft graph survives: it bakes only session-stable pointers
5739 /// (the scratch KV, the resident embedding), none of which the rewind moves.
5740 ///
5741 /// The checkpoint is CONSUMED (`turn_ckpt` taken): its snapshot buffers are freed here, and
5742 /// this turn's own prime installs a fresh one at the new prompt end. Returns the position
5743 /// rewound to, or `None` when the session holds no checkpoint (caller: full re-prime).
5744 pub fn spec_rewind_to_checkpoint(
5745 &self,
5746 e: &Engine,
5747 sess: &mut SpecSession,
5748 ) -> Result<Option<usize>, Box<dyn std::error::Error>> {
5749 if sess.turn_ckpt.as_ref().is_some_and(|ckpt| {
5750 !sess.cache.can_rollback(&ckpt.snap, 0) || !sess.scratch.can_rewind_to(ckpt.pos)
5751 }) {
5752 return Err(
5753 "SWA ring rewind checkpoint has been lapped; full re-prime required".into(),
5754 );
5755 }
5756 let Some(ckpt) = sess.turn_ckpt.take() else {
5757 return Ok(None);
5758 };
5759 assert!(
5760 ckpt.pos <= sess.committed.len(),
5761 "checkpoint past committed ({} > {})",
5762 ckpt.pos,
5763 sess.committed.len()
5764 );
5765 // Restore through each layer's owning engine. A single primary-engine rollback is not
5766 // sufficient when the serving cache is stage-owned under cross-device PP.
5767 crate::pp::restore_cache_checkpoint(e, &self.cfg, None, &mut sess.cache, &ckpt.snap)?;
5768 debug_assert_eq!(
5769 sess.cache.pos, ckpt.pos,
5770 "rollback landed off the checkpoint"
5771 );
5772 sess.scratch.set_len(e, ckpt.pos)?;
5773 sess.committed.truncate(ckpt.pos);
5774 sess.last_h = Some(ckpt.last_h);
5775 sess.next_pred = None;
5776 sess.pending_tok = None;
5777 Ok(Some(ckpt.pos))
5778 }
5779
5780 /// Grow a parked speculative session to `target_cap` and rewind it to its retained turn
5781 /// checkpoint without re-priming the checkpoint prefix.
5782 ///
5783 /// The trunk cache is restored exactly like a plain grown cache: append-only full-attention
5784 /// KV rows come from the parked cache, while recurrent state comes from the checkpoint's
5785 /// owned snapshot. The MTP scratch is also context-linear and its rows below the checkpoint
5786 /// remain authoritative, so they are copied into a fresh larger scratch before its length is
5787 /// truncated. Pointer-baking draft graphs are dropped and recaptured on the next burst.
5788 ///
5789 /// All fallible work completes before `sess` is mutated. A failed allocation or copy leaves
5790 /// the parked session intact, allowing the caller one reclaim-and-retry attempt.
5791 pub fn spec_grow_and_rewind_to_checkpoint(
5792 &self,
5793 e: &Engine,
5794 sess: &mut SpecSession,
5795 target_cap: usize,
5796 ) -> Result<Option<usize>, Box<dyn std::error::Error>> {
5797 if target_cap <= sess.cache.max_ctx {
5798 return self.spec_rewind_to_checkpoint(e, sess);
5799 }
5800 let Some(ckpt) = sess.turn_ckpt.as_ref() else {
5801 return Ok(None);
5802 };
5803 if ckpt.pos == 0 || ckpt.pos > sess.committed.len() {
5804 return Err(format!(
5805 "checkpoint pos {} outside committed length {}",
5806 ckpt.pos,
5807 sess.committed.len(),
5808 )
5809 .into());
5810 }
5811 if ckpt.pos > target_cap {
5812 return Err(format!(
5813 "checkpoint pos {} exceeds grown capacity {target_cap}",
5814 ckpt.pos,
5815 )
5816 .into());
5817 }
5818
5819 let mut grown_cache = crate::pp::new_cache(e, &self.cfg, target_cap)?;
5820 let mut grown_scratch = MtpScratch::new(
5821 e,
5822 &self.cfg,
5823 target_cap,
5824 self.mtp.as_ref().and_then(|m| m.geom.as_ref()),
5825 )?;
5826 crate::pp::restore_cache_checkpoint(
5827 e,
5828 &self.cfg,
5829 Some(&sess.cache),
5830 &mut grown_cache,
5831 &ckpt.snap,
5832 )?;
5833
5834 let src = &sess.scratch.kv;
5835 let dst = &mut grown_scratch.kv;
5836 if ckpt.pos > src.len
5837 || src.kv_dim_k != dst.kv_dim_k
5838 || src.kv_dim_v != dst.kv_dim_v
5839 || src.k_tok_bytes != dst.k_tok_bytes
5840 || src.v_tok_bytes != dst.v_tok_bytes
5841 {
5842 return Err(format!(
5843 "checkpoint draft layout mismatch (pos {}, source len {})",
5844 ckpt.pos, src.len,
5845 )
5846 .into());
5847 }
5848 let kb = ckpt.pos * src.k_tok_bytes;
5849 let vb = ckpt.pos * src.v_tok_bytes;
5850 if kb > 0 {
5851 e.copy_u8_into(&mut dst.k, 0, &src.k, kb)?;
5852 }
5853 if vb > 0 {
5854 e.copy_u8_into(&mut dst.v, 0, &src.v, vb)?;
5855 }
5856 grown_scratch.set_len(e, ckpt.pos)?;
5857 // The old scratch is dropped immediately after publication below. Bound its D2D reads
5858 // first; growth happens once per rewritten turn, outside the decode hot loop.
5859 e.stream().synchronize()?;
5860
5861 let ckpt = sess
5862 .turn_ckpt
5863 .take()
5864 .expect("checkpoint remained present through transactional grow");
5865 let pos = ckpt.pos;
5866 sess.cache = grown_cache;
5867 sess.scratch = grown_scratch;
5868 sess.committed.truncate(pos);
5869 sess.last_h = Some(ckpt.last_h);
5870 sess.next_pred = None;
5871 sess.pending_tok = None;
5872 sess.draft_ctx = None;
5873 debug_assert_eq!(sess.cache.pos, pos, "grown rewind landed off checkpoint");
5874 debug_assert_eq!(
5875 sess.scratch.kv.len, pos,
5876 "grown draft rewind landed off checkpoint"
5877 );
5878 Ok(Some(pos))
5879 }
5880
5881 /// Commit a carried pending bonus (see SpecSession::pending_tok): one T=1 trunk pass
5882 /// (its logits' argmax becomes next_pred) + the draft-KV fill at the carried anchor —
5883 /// byte-identical to the pre-carry session tail. Required before a non-empty-suffix
5884 /// prime, a sampled turn, or parking a session for pool reuse. No-op without a pending.
5885 pub fn spec_flush_pending(
5886 &self,
5887 e: &Engine,
5888 sess: &mut SpecSession,
5889 ) -> Result<(), Box<dyn std::error::Error>> {
5890 let Some(b) = sess.pending_tok.take() else {
5891 return Ok(());
5892 };
5893 let mtp = self
5894 .mtp
5895 .as_ref()
5896 .expect("pending carry requires an MTP head");
5897 let n_embd = self.cfg.n_embd as usize;
5898 let (embd_qt, embd_rb) = self.embd.qt_and_row_bytes(n_embd);
5899 let embd_gpu = if spec_host_embd() {
5900 None
5901 } else {
5902 Some(
5903 self.embd_gpu
5904 .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload")),
5905 )
5906 };
5907 let embd_dev = embd_gpu.map(|g| (g, embd_qt, embd_rb));
5908 let pos_b = sess.cache.pos;
5909 sess.scratch.set_len(e, pos_b)?;
5910 let (lg_b, hb) = self.spec_target_step_h(e, b, &mut sess.cache)?;
5911 sess.next_pred = Some(argmax(&lg_b) as u32);
5912 let anchor = sess
5913 .last_h
5914 .as_ref()
5915 .expect("pending carry requires last_h (the predecessor-row anchor)");
5916 self.mtp_kv_fill(e, mtp, &[b], anchor, pos_b, &mut sess.scratch, embd_dev)?;
5917 sess.last_h = Some(hb);
5918 sess.committed.push(b);
5919 Ok(())
5920 }
5921
5922 /// Solo target feed used only at speculative round boundaries. Step35 serving made its
5923 /// staged batched B=1 graph authoritative, so a speculative session must enter and leave
5924 /// rounds through that same graph. Other model families keep their eager T=1 contract.
5925 fn spec_target_step_h(
5926 &self,
5927 e: &Engine,
5928 token: u32,
5929 cache: &mut Cache,
5930 ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
5931 if self.cfg.step35.is_none() && !self.qwen35_serving_class() {
5932 return self.decode_step_h(e, token, cache);
5933 }
5934 let pos0 = cache.pos;
5935 let (logits, hidden) = self.decode_step_t_core(e, &[token], pos0, cache, None, None)?;
5936 Ok((e.dtoh(&logits)?, hidden))
5937 }
5938
5939 /// The archs whose LIVE B=1 serving runs the generic BATCHED numeric class (decode_step_batch
5940 /// walk + batched head), so their spec verify must run the SAME class. MoE learned this
5941 /// 2026-08-14 AM (4b777ccc5); the dense hybrid reproduced the identical near-tie flip class
5942 /// the same day on Qwen3.8-27B — eager-class verify logits drift from batched-class serving
5943 /// logits ("1 ULP at layer 2 → 2.3e-1 logit maxdiff at the head"), and the GDN recurrence
5944 /// carries the drift until a near-tie flips deep in generation. One predicate so the five
5945 /// dispatch sites cannot drift apart again.
5946 fn qwen35_serving_class(&self) -> bool {
5947 matches!(
5948 self.cfg.arch,
5949 memra_gguf::config::Arch::Qwen35 | memra_gguf::config::Arch::Qwen35Moe
5950 )
5951 }
5952
5953 /// Reduced-matrix admission for increment 1. This deliberately does not change the PP-2
5954 /// serving policy: the worker calls it only after `MEMRA_SPEC_PIPE=1` and an explicit spec
5955 /// session already exist.
5956 pub fn spec_pipe_available(&self, e: &Engine) -> bool {
5957 if std::env::var("MEMRA_SPEC_PIPE").as_deref() != Ok("1")
5958 || !spec_devacc()
5959 || spec_replay_env_enabled()
5960 || spec_stream()
5961 || std::env::var("MEMRA_SPEC_ADAPT").as_deref() == Ok("1")
5962 || std::env::var("MEMRA_SPEC_PMIN0").as_deref() == Ok("1")
5963 || std::env::var("MEMRA_SPEC_PP_ANATOMY").as_deref() == Ok("1")
5964 || std::env::var("MEMRA_SPEC_PMIN")
5965 .ok()
5966 .and_then(|v| v.parse::<f32>().ok())
5967 .unwrap_or(0.0)
5968 > 0.0
5969 || self.is_gemma4_e4b()
5970 || self.cfg.gemma4.is_some()
5971 || self.mtp.is_none()
5972 {
5973 return false;
5974 }
5975 let Some(cuts) = crate::pp::pp_cuts(self.layers.len()) else {
5976 return false;
5977 };
5978 if cuts.len() != 3 || crate::pp::pp2_streams_off() || !crate::pp::spec_pp_on() {
5979 return false;
5980 }
5981 crate::pp::PpNRt::get(e)
5982 .map(|rt| rt.n_stages() == 2 && rt.cross_device())
5983 .unwrap_or(false)
5984 }
5985
5986 /// Two warm greedy continuation bursts over one PP-2 interval coordinator. The two existing
5987 /// `generate_spec_inner2` call stacks own all per-session round locals; only phase issue order
5988 /// changes. No callback is accepted in increment 1 — the worker publishes each completed burst.
5989 #[allow(clippy::too_many_arguments)]
5990 pub fn generate_spec_session_pair(
5991 &self,
5992 e: &Engine,
5993 sess_a: &mut SpecSession,
5994 max_new_a: usize,
5995 k_a: usize,
5996 sess_b: &mut SpecSession,
5997 max_new_b: usize,
5998 k_b: usize,
5999 ) -> Result<((Vec<u32>, usize, usize), (Vec<u32>, usize, usize)), Box<dyn std::error::Error>>
6000 {
6001 if !self.spec_pipe_available(e) {
6002 return Err("two-session speculative pipeline is outside its reduced matrix".into());
6003 }
6004 if max_new_a == 0 || max_new_b == 0 || k_a == 0 || k_b == 0 {
6005 return Err(
6006 "two-session speculative pipeline requires non-empty positive-K bursts".into(),
6007 );
6008 }
6009 for sess in [&*sess_a, &*sess_b] {
6010 if sess.committed.is_empty()
6011 || sess.last_h.is_none()
6012 || (sess.next_pred.is_none() && sess.pending_tok.is_none())
6013 {
6014 return Err("two-session speculative pipeline requires warm continuations".into());
6015 }
6016 }
6017
6018 let mtp_dense = self
6019 .mtp
6020 .as_ref()
6021 .map(|m| matches!(m.ffn, crate::hybrid::Ffn::Dense { .. }))
6022 .unwrap_or(false);
6023 let trunk_dense = self
6024 .layers
6025 .iter()
6026 .all(|l| matches!(l.ffn, crate::hybrid::Ffn::Dense { .. }));
6027 let graph_ok = std::env::var("MEMRA_SPEC_NOGRAPH").is_err()
6028 && !spec_host_embd()
6029 && mtp_dense
6030 && trunk_dense
6031 && !crate::model::full_prec_enabled();
6032 let graph_a = graph_ok && k_a + 2 < 96;
6033 let graph_b = graph_ok && k_b + 2 < 96;
6034 let was_tracking = e.ctx().is_event_tracking();
6035 if (graph_a || graph_b) && was_tracking {
6036 unsafe {
6037 e.ctx().disable_event_tracking();
6038 }
6039 }
6040
6041 static LOGGED: std::sync::Once = std::sync::Once::new();
6042 LOGGED.call_once(|| {
6043 eprintln!("[spec-pipe] two-session PP-2 continuation pipeline engaged");
6044 });
6045 let sync = std::sync::Arc::new(SpecPipeSync::new());
6046 let lane_a = SpecPipeLane {
6047 sync: sync.clone(),
6048 lane: 0,
6049 };
6050 let lane_b = SpecPipeLane { sync, lane: 1 };
6051 let mut sess_b_ptr = SpecPipeSessionPtr(sess_b as *mut SpecSession);
6052 let (result_a, result_b) = std::thread::scope(|scope| {
6053 let b = scope.spawn(move || {
6054 let mut finish = SpecPipeFinish::new(&lane_b);
6055 let sess_b = unsafe { sess_b_ptr.get_mut() };
6056 let result = e
6057 .ctx()
6058 .bind_to_thread()
6059 .map_err(|err| err.to_string())
6060 .and_then(|_| {
6061 self.generate_spec_inner2(
6062 e,
6063 &[],
6064 max_new_b,
6065 k_b,
6066 graph_b,
6067 Some(sess_b),
6068 None,
6069 None,
6070 None,
6071 None,
6072 Some(&lane_b),
6073 )
6074 .map_err(|err| err.to_string())
6075 });
6076 finish.close(result.is_err());
6077 result
6078 });
6079 let mut finish = SpecPipeFinish::new(&lane_a);
6080 let result_a = self.generate_spec_inner2(
6081 e,
6082 &[],
6083 max_new_a,
6084 k_a,
6085 graph_a,
6086 Some(sess_a),
6087 None,
6088 None,
6089 None,
6090 None,
6091 Some(&lane_a),
6092 );
6093 finish.close(result_a.is_err());
6094 let result_b = b
6095 .join()
6096 .map_err(|_| "paired speculative session B panicked".to_string())
6097 .and_then(|r| r);
6098 (result_a, result_b)
6099 });
6100
6101 if (graph_a || graph_b) && was_tracking {
6102 unsafe {
6103 e.ctx().enable_event_tracking();
6104 }
6105 }
6106 let result_a = result_a?;
6107 let result_b = result_b.map_err(|err| -> Box<dyn std::error::Error> { err.into() })?;
6108 Ok((result_a, result_b))
6109 }
6110
6111 /// One spec-decode turn on a live session. `suffix` = the NEW tokens only (turn N+1's user
6112 /// message rendered through the chat template continuation). Returns (new tokens emitted,
6113 /// drafted, accepted); session.committed grows by suffix + emitted.
6114 pub fn generate_spec_session(
6115 &self,
6116 e: &Engine,
6117 sess: &mut SpecSession,
6118 suffix: &[u32],
6119 max_new: usize,
6120 k: usize,
6121 ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
6122 self.generate_spec_session_sampled(e, sess, suffix, max_new, k, None, None)
6123 }
6124
6125 /// Serve-path sampled spec: routes the burst through the rejection-sampling verify with
6126 /// per-SESSION Philox continuity (sess.sctr/uctr). None = env-driven (CLI) or greedy.
6127 /// Filters (top-k/p/min-p) apply SYMMETRICALLY to draft q and verify p — distribution-exact
6128 /// for the filtered target (feat/filtered-spec).
6129 ///
6130 /// `on_commit` (sse-cadence, 2026-08-05): called with each newly-emitted slice of the
6131 /// output — once right after the prime's first token, then once per round commit — so a
6132 /// streaming caller can flush text at round cadence instead of once per burst. The slices
6133 /// are disjoint, in order, and concatenate to exactly the returned token vec. Emission-
6134 /// timing only: token bytes, session state, and exactness are untouched.
6135 ///
6136 /// The returned bool is a CONTINUE-VERDICT (admission yield, 2026-08-06): `false` ends
6137 /// the burst at the current round boundary, exactly as if `max_new` had been reached —
6138 /// the caller's scheduler regains control without waiting the burst out. Burst size is
6139 /// content-neutral (spec-levers battery), so an early exit moves WHEN the burst returns,
6140 /// never what tokens say. The slice may be EMPTY (a poll-only boundary — round-stream
6141 /// drains and the defensive tail flush can land with nothing new committed).
6142 #[allow(clippy::too_many_arguments)]
6143 pub fn generate_spec_session_sampled(
6144 &self,
6145 e: &Engine,
6146 sess: &mut SpecSession,
6147 suffix: &[u32],
6148 max_new: usize,
6149 k: usize,
6150 sampling: Option<SpecSampling>,
6151 on_commit: Option<&mut dyn FnMut(&[u32]) -> bool>,
6152 ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
6153 self.generate_spec_session_sampled_prime_split(
6154 e, sess, suffix, max_new, k, sampling, None, on_commit,
6155 )
6156 }
6157
6158 /// Serve-only cold-prime segmentation twin. `prime_split` is the same stable boundary the
6159 /// plain worker would honor before entering its sub-floor tokenwise tail; warm continuations
6160 /// pass `None` and stay on the existing zero-prime path.
6161 #[allow(clippy::too_many_arguments)]
6162 pub fn generate_spec_session_sampled_prime_split(
6163 &self,
6164 e: &Engine,
6165 sess: &mut SpecSession,
6166 suffix: &[u32],
6167 max_new: usize,
6168 k: usize,
6169 sampling: Option<SpecSampling>,
6170 prime_split: Option<usize>,
6171 on_commit: Option<&mut dyn FnMut(&[u32]) -> bool>,
6172 ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
6173 self.generate_spec_session_constrained_prime_split(
6174 e,
6175 sess,
6176 suffix,
6177 max_new,
6178 k,
6179 sampling,
6180 None,
6181 prime_split,
6182 on_commit,
6183 )
6184 }
6185
6186 /// `generate_spec_session_sampled` + GRAMMAR (constrained decoding, 2026-08-03): the
6187 /// hook truncates acceptance at the first grammar-illegal token AFTER the exactness
6188 /// verify (grammar is an extra rejection rule, ordering like the batched-verify twins)
6189 /// and replaces an illegal bonus with the MASKED argmax of the target's own verify
6190 /// column — token-identical to constrained plain greedy decode. GREEDY only (the
6191 /// worker routes sampled constrained to plain decode). Acceptance under tight grammars
6192 /// may drop (drafter is unconstrained); that is measured, not hidden.
6193 #[allow(clippy::too_many_arguments)]
6194 pub fn generate_spec_session_constrained(
6195 &self,
6196 e: &Engine,
6197 sess: &mut SpecSession,
6198 suffix: &[u32],
6199 max_new: usize,
6200 k: usize,
6201 sampling: Option<SpecSampling>,
6202 constraint: Option<&mut dyn SpecConstraint>,
6203 on_commit: Option<&mut dyn FnMut(&[u32]) -> bool>,
6204 ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
6205 self.generate_spec_session_constrained_prime_split(
6206 e, sess, suffix, max_new, k, sampling, constraint, None, on_commit,
6207 )
6208 }
6209
6210 #[allow(clippy::too_many_arguments)]
6211 pub fn generate_spec_session_constrained_prime_split(
6212 &self,
6213 e: &Engine,
6214 sess: &mut SpecSession,
6215 suffix: &[u32],
6216 max_new: usize,
6217 k: usize,
6218 sampling: Option<SpecSampling>,
6219 constraint: Option<&mut dyn SpecConstraint>,
6220 prime_split: Option<usize>,
6221 on_commit: Option<&mut dyn FnMut(&[u32]) -> bool>,
6222 ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
6223 if constraint.is_some() && sampling.is_some_and(|s| s.temp > 0.0) {
6224 return Err(
6225 "constrained spec decode is greedy-only (worker routes sampled \
6226 constrained to plain decode)"
6227 .into(),
6228 );
6229 }
6230 // PENDING-CARRY entry flush: a carried bonus precedes any new suffix in the sequence,
6231 // so it must commit BEFORE the suffix primes; the sampled path doesn't carry (its
6232 // round-0 accept needs the commit pass's logits). Empty-suffix greedy bursts — the
6233 // serve continuation case — consume the carry in-loop with zero solo passes.
6234 if sess.pending_tok.is_some()
6235 && (!suffix.is_empty() || sampling.map_or(false, |s| s.temp > 0.0))
6236 {
6237 self.spec_flush_pending(e, sess)?;
6238 }
6239 let mtp_dense = self
6240 .mtp
6241 .as_ref()
6242 .map(|m| matches!(m.ffn, crate::hybrid::Ffn::Dense { .. }))
6243 .unwrap_or(false);
6244 let trunk_dense = self
6245 .layers
6246 .iter()
6247 .all(|l| matches!(l.ffn, crate::hybrid::Ffn::Dense { .. }));
6248 // FULL_PREC forces the EAGER draft: the graph capture would enclose cuBLASLt f32 GEMV
6249 // (the FloatBf16 else-branches) and a bf16_to_f32 dequant alloc — neither is stream-capture
6250 // safe. Eager rides matmul/matmul_decode_exact, which dequant FloatBf16 on use. (§item 2.)
6251 let graph_draft = std::env::var("MEMRA_SPEC_NOGRAPH").is_err()
6252 && !spec_host_embd()
6253 && mtp_dense
6254 && trunk_dense
6255 && k + 2 < 96
6256 && !crate::model::full_prec_enabled();
6257 let was_tracking = e.ctx().is_event_tracking();
6258 if graph_draft && was_tracking {
6259 unsafe {
6260 e.ctx().disable_event_tracking();
6261 }
6262 }
6263 let r = self.generate_spec_inner2(
6264 e,
6265 suffix,
6266 max_new,
6267 k,
6268 graph_draft,
6269 Some(sess),
6270 sampling,
6271 constraint,
6272 on_commit,
6273 prime_split,
6274 None,
6275 );
6276 if graph_draft && was_tracking {
6277 unsafe {
6278 e.ctx().enable_event_tracking();
6279 }
6280 }
6281 let (out, d, a) = r?;
6282 Ok((out, d, a))
6283 }
6284
6285 pub fn generate_spec(
6286 &self,
6287 e: &Engine,
6288 prompt: &[u32],
6289 max_new: usize,
6290 k: usize,
6291 ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
6292 let mtp_dense = self
6293 .mtp
6294 .as_ref()
6295 .map(|m| matches!(m.ffn, crate::hybrid::Ffn::Dense { .. }))
6296 .unwrap_or(false);
6297 let trunk_dense = self
6298 .layers
6299 .iter()
6300 .all(|l| matches!(l.ffn, crate::hybrid::Ffn::Dense { .. }));
6301 // FULL_PREC forces eager (see generate_spec_session note): CUDA graph capture cannot
6302 // enclose cuBLASLt f32 GEMV or the bf16_to_f32 dequant alloc the FloatBf16 path needs.
6303 let graph_draft = std::env::var("MEMRA_SPEC_NOGRAPH").is_err()
6304 && !spec_host_embd()
6305 && mtp_dense
6306 && trunk_dense
6307 && k + 2 < 96
6308 && !crate::model::full_prec_enabled();
6309 if !graph_draft {
6310 return self.generate_spec_inner2(
6311 e, prompt, max_new, k, false, None, None, None, None, None, None,
6312 );
6313 }
6314 let was_tracking = e.ctx().is_event_tracking();
6315 if was_tracking {
6316 unsafe {
6317 e.ctx().disable_event_tracking();
6318 }
6319 }
6320 let r = self.generate_spec_inner2(
6321 e, prompt, max_new, k, true, None, None, None, None, None, None,
6322 );
6323 if was_tracking {
6324 unsafe {
6325 e.ctx().enable_event_tracking();
6326 }
6327 }
6328 r
6329 }
6330
6331 fn generate_spec_inner2(
6332 &self,
6333 e: &Engine,
6334 prompt: &[u32],
6335 max_new: usize,
6336 k: usize,
6337 graph_draft: bool,
6338 mut sess: Option<&mut SpecSession>,
6339 sampling: Option<SpecSampling>,
6340 mut constraint: Option<&mut dyn SpecConstraint>,
6341 mut on_commit: Option<&mut dyn FnMut(&[u32]) -> bool>,
6342 prime_split: Option<usize>,
6343 pipe: Option<&SpecPipeLane>,
6344 ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
6345 assert!(k >= 1, "k must be >= 1");
6346 if let Some(p) = pipe {
6347 p.setup_begin()?;
6348 }
6349 // sse-cadence flush cursor: everything in out[..flushed] has been handed to on_commit.
6350 let mut flushed = 0usize;
6351 // admission yield (2026-08-06): on_commit's continue-verdict; false = end the burst
6352 // at the next round boundary (same exit as max_new reached — the session tail runs).
6353 // Initialized by the unconditional post-prime flush below.
6354 let mut keep_going;
6355 let mtp = self
6356 .mtp
6357 .as_ref()
6358 .expect("generate_spec requires an MTP head (nextn_predict_layers>0)");
6359 let n_vocab = self.output.out_features();
6360 // FR-Spec: the draft head may be TRIMMED (fewer rows than n_vocab); the draft argmax runs
6361 // over the draft vocab and the winning index maps through d2t to a TARGET token id.
6362 // Everything downstream (verify/accept/commit) sees target ids only — exactness unchanged.
6363 let d_vocab = mtp
6364 .shared_head_head
6365 .as_ref()
6366 .unwrap_or(&self.output)
6367 .out_features();
6368 let n_embd = self.cfg.n_embd as usize;
6369 // SESSION MODE: reuse the live cache/scratch, prime only the suffix. `base` = tokens
6370 // already committed (their state is in the caches); 0 = fresh single-shot call.
6371 let session_mode = sess.is_some();
6372 let max_ctx = match sess.as_ref() {
6373 Some(s) => s.cache.max_ctx,
6374 None => prompt.len() + max_new + k + 8,
6375 };
6376 let mut own_cache;
6377 let mut own_scratch;
6378 // PREFIX-CACHE capture request threaded out of the session (lane/spec-prefix-cache):
6379 // (requested split, destination slot). Single-shot per burst; fresh calls have none.
6380 let mut sess_capture: Option<(Option<usize>, &mut Option<SpecBoundaryCapture>)> = None;
6381 let (
6382 cache,
6383 scratch,
6384 mut sess_tail,
6385 mut sess_draft_slot,
6386 mut sess_pending_slot,
6387 sess_ckpt_slot,
6388 sess_telem,
6389 ): (
6390 &mut Cache,
6391 &mut MtpScratch,
6392 Option<(
6393 &mut Vec<u32>,
6394 &mut Option<CudaSlice<f32>>,
6395 &mut Option<u32>,
6396 &mut u32,
6397 &mut u32,
6398 )>,
6399 Option<&mut Option<DraftGraphCtx>>,
6400 Option<&mut Option<u32>>,
6401 Option<&mut Option<SpecCheckpoint>>,
6402 Option<&SpecTelemetryCounters>,
6403 ) = match sess.take() {
6404 Some(sr) => {
6405 let SpecSession {
6406 cache,
6407 scratch,
6408 committed,
6409 last_h,
6410 next_pred,
6411 sctr: s_sctr,
6412 uctr: s_uctr,
6413 draft_ctx,
6414 pending_tok,
6415 turn_ckpt,
6416 telem,
6417 capture_at,
6418 boundary_capture,
6419 } = sr;
6420 sess_capture = Some((capture_at.take(), boundary_capture));
6421 (
6422 cache,
6423 scratch,
6424 Some((committed, last_h, next_pred, s_sctr, s_uctr)),
6425 Some(draft_ctx),
6426 Some(pending_tok),
6427 Some(turn_ckpt),
6428 Some(telem),
6429 )
6430 }
6431 None => {
6432 // STAGE-OWNED KV (lane/pp2-spec 2026-08-06) — see `new_session`. Door shut =
6433 // `Cache::new` verbatim.
6434 own_cache = crate::pp::new_cache(e, &self.cfg, max_ctx)?;
6435 // Persistent scratch = max_ctx rows (~2KB/token quantized).
6436 own_scratch = MtpScratch::new(
6437 e,
6438 &self.cfg,
6439 max_ctx,
6440 self.mtp.as_ref().and_then(|m| m.geom.as_ref()),
6441 )?;
6442 (
6443 &mut own_cache,
6444 &mut own_scratch,
6445 None,
6446 None,
6447 None,
6448 None,
6449 None,
6450 )
6451 }
6452 };
6453 let base = cache.pos;
6454 // PENDING-CARRY consume (2026-08-01): a carried bonus reaches here only on the
6455 // empty-suffix GREEDY continuation path (generate_spec_session_sampled flushed every
6456 // other case). It enters the round loop as round-0's pending — verify col 0 — exactly
6457 // like a mid-burst full-accept boundary: no init feed, no tail commit pass.
6458 let carried_pending: Option<u32> = sess_pending_slot.as_mut().and_then(|s| s.take());
6459 // PERSISTENT DRAFT KV (the only mode since 2026-07-08 — the legacy round-local scratch,
6460 // MEMRA_SPEC_KVLOCAL, measured -35 acceptance pts on the 27B p3 sweep and was removed;
6461 // acceptance-only — exactness is verify's job either way).
6462 // HIDDEN-PAIRING CONVENTION (DEFAULT = predecessor-row, 2026-07-04 — the 27B acceptance
6463 // unlock, +16pts): the MTP head is TRAINED on rows pairing token x_p with the trunk
6464 // hidden of its PREDECESSOR h_{p-1} (the reference engine's mtp_update shifts the target
6465 // hiddens right by one; its draft step 0 feeds (id_last, TRUE hidden of the row id_last
6466 // was sampled from)). memra's historical convention paired SAME-ROW (x_p, h_p) in the fill
6467 // and seeded chain step 0 through an extra MTP pass on a duplicated token (the
6468 // pseudo-seed) — measured 27B p2 K=3 acceptance 0.569 vs 0.731, p3 0.445 vs 0.63+, and
6469 // the chain steps j>=1 were already predecessor-shaped, so ONLY the fill + step-0 seed
6470 // move. The fill shifts by one and the chain seeds from the predecessor's true hidden
6471 // DIRECTLY (vh_seed / vx[j-1]) — the pseudo pass disappears (one MTP-block pass saved
6472 // per round on top of the acceptance win). Draft-quality-only: exactness stays the
6473 // verify's job either way. (The legacy same-row pairing seam, MEMRA_SPEC_HSAME, and its
6474 // pseudo-seed passes were removed 2026-07-08 — predecessor pairing won by +16 acc pts;
6475 // the legacy round-local scratch, MEMRA_SPEC_KVLOCAL, went with it.)
6476 // REPLAY-FREE PARTIAL ACCEPT (default, 2026-07-03): partial rounds keep the verify's own
6477 // bit-identical committed-prefix state (KV truncate + recur rebuild from the VerifyCkpt)
6478 // and leave the bonus PENDING — no duplicate trunk pass (profiled ~0.54 extra full weight
6479 // reads/round at long ctx). MEMRA_SPEC_REPLAY=1 restores the legacy rollback+replay (A/B
6480 // + fallback seam).
6481 // Qwen35-MoE stays on the correctness reference path until its retained verify-state
6482 // commit is proven equivalent to sequential serving on the long-prompt gate. Replaying
6483 // every accepted round through the serving-class verifier is slower, but prevents a
6484 // numerically exact verify result from carrying a drifted recurrent cache into the next
6485 // round. DENSE qwen35 runs replay-free: its verify already executes the serving batched
6486 // class (qwen35_verify_batch_layers), and the serving-class replay loop below steps
6487 // per-row T=1 (replay.len() full weight reads/round — measured 69 -> 30 tok/s on
6488 // Qwen3.8-27B, 2026-08-15); the replay-free VerifyCkpt commit is gated bit-identical by
6489 // the spec-serve battery before release.
6490 let spec_replay = spec_replay_env_enabled()
6491 || matches!(self.cfg.arch, memra_gguf::config::Arch::Qwen35Moe);
6492 if constraint.is_some() && spec_replay {
6493 return Err(
6494 "constrained spec decode does not support MEMRA_SPEC_REPLAY=1 \
6495 (legacy replay commits an unmasked bonus)"
6496 .into(),
6497 );
6498 }
6499 // TRUE-HIDDEN REFRESH (default in persistent-draft-KV mode): every round overwrites the
6500 // committed positions' scratch entries from the verify's exact hiddens (mtp_kv_fill batch)
6501 // instead of keeping chain-approximate entries. MEMRA_SPEC_NOREFRESH=1 = legacy (A/B seam).
6502 let refresh = std::env::var("MEMRA_SPEC_NOREFRESH").is_err();
6503
6504 // prime: BATCHED cache prime (prime_cache — the measured #1 e2e gap: tokenwise primed at
6505 // ~102/38 tok/s vs the engine's ~2000-5900 tok/s batched prefill). prime_cache returns the
6506 // full pre-output_norm hidden stack [T, n_embd], which IS prompt_h (the persistent-draft-KV
6507 // mtp_kv_fill input) — no per-token collection needed. Prompts below PRIME_MIN_T, and
6508 // MEMRA_PRIME_TOKENWISE=1, and frozen Hy3 CPU/GPU expert splits take the tokenwise
6509 // decode_step_h loop. The latter avoids transient GPU staging of the spilled expert bank.
6510 // EMPTY-SUFFIX CONTINUATION (serve bursts): a session turn with NO new tokens resumes
6511 // generation exactly where the last turn stopped — no prime at all. The stashed
6512 // `next_pred` plays prime_logits' argmax role (it IS the argmax of the logits after
6513 // committed.last()); `last_h` seeds the predecessor pairing below. Fresh calls and
6514 // non-empty suffixes take the normal path.
6515 let continuation = prompt.is_empty();
6516 if continuation {
6517 assert!(session_mode, "empty prompt requires a session");
6518 assert!(
6519 sess_tail
6520 .as_ref()
6521 .map_or(false, |(c, lh, np, _, _)| !c.is_empty()
6522 && lh.is_some()
6523 && (np.is_some() || carried_pending.is_some())),
6524 "empty-suffix continuation needs a primed session (committed + last_h + next_pred|pending)"
6525 );
6526 }
6527 let mut prime_logits;
6528 let mut prompt_h: Option<CudaSlice<f32>> = None;
6529 let t_prime = std::time::Instant::now();
6530 let batched_prime = !continuation
6531 && prompt.len() >= crate::hybrid_forward::PRIME_MIN_T
6532 && std::env::var("MEMRA_PRIME_TOKENWISE").is_err()
6533 && !e.frozen_cpu_experts_prefer_tokenwise_prime();
6534 let prime_split = prime_split.filter(|&split| split > 0 && split < prompt.len());
6535 if prime_split.is_some() && (continuation || base != 0) {
6536 return Err("spec prime split is cold-session-only".into());
6537 }
6538 if continuation {
6539 prime_logits = Vec::new();
6540 } else if let Some(split) = prime_split {
6541 if split < crate::hybrid_forward::PRIME_MIN_T {
6542 return Err(format!(
6543 "spec prime split {split} is below PRIME_MIN_T {}",
6544 crate::hybrid_forward::PRIME_MIN_T,
6545 )
6546 .into());
6547 }
6548 // Mirror the plain worker's affinity boundary exactly. The prefix is a request-level
6549 // prime (`queued_after` keeps Step35 arm selection independent of this stop); a tail
6550 // below PRIME_MIN_T then takes the same eager tokenwise continuation as prefill_tick.
6551 // Retain every hidden row so the draft scratch fill remains one coherent prompt.
6552 let mut h_all = e.uninit(prompt.len() * n_embd)?;
6553 let (l, _, h_prefix) =
6554 self.prime_cache(e, &prompt[..split], &mut *cache, prompt.len() - split)?;
6555 e.copy_into(&mut h_all, 0, &h_prefix, split * n_embd)?;
6556 prime_logits = l;
6557 // PREFIX-CACHE BOUNDARY CAPTURE (lane/spec-prefix-cache): the GDN conv/ssm states
6558 // are about to be advanced in place by the tail prime, so this is the ONLY moment
6559 // the boundary's recurrent state exists. Capture iff the worker requested exactly
6560 // this split. cache.pos == split here (the prefix prime just finished). A failed
6561 // snapshot is silent (turn_ckpt convention) — publication is an optimization,
6562 // never a correctness dependency.
6563 if let Some((requested, slot)) = sess_capture.as_mut() {
6564 if *requested == Some(split) {
6565 debug_assert_eq!(cache.pos, split, "boundary capture off the prime split");
6566 if let Ok(snap) = cache.snapshot(e) {
6567 **slot = Some(SpecBoundaryCapture {
6568 snap,
6569 pos: split,
6570 logits: prime_logits.clone(),
6571 });
6572 }
6573 }
6574 }
6575 let tail = &prompt[split..];
6576 if tail.len() >= crate::hybrid_forward::PRIME_MIN_T
6577 && std::env::var("MEMRA_PRIME_TOKENWISE").is_err()
6578 && !e.frozen_cpu_experts_prefer_tokenwise_prime()
6579 {
6580 let (l, _, h_tail) = self.prime_cache(e, tail, &mut *cache, 0)?;
6581 e.copy_into(&mut h_all, split * n_embd, &h_tail, tail.len() * n_embd)?;
6582 prime_logits = l;
6583 } else {
6584 for (i, &tok) in tail.iter().enumerate() {
6585 let (l, h) = self.decode_step_h(e, tok, &mut *cache)?;
6586 e.copy_into(&mut h_all, (split + i) * n_embd, &h, n_embd)?;
6587 prime_logits = l;
6588 }
6589 }
6590 if std::env::var("MEMRA_SPEC_STATS").as_deref() == Ok("1") {
6591 eprintln!("[spec-prime] affinity split={split} tail={}", tail.len());
6592 }
6593 prompt_h = Some(h_all);
6594 } else if batched_prime {
6595 let (l, _h_seed, hiddens) = self.prime_cache(e, prompt, &mut *cache, 0)?;
6596 prime_logits = l;
6597 prompt_h = Some(hiddens);
6598 } else {
6599 prime_logits = Vec::new();
6600 prompt_h = Some(e.uninit(prompt.len() * n_embd)?);
6601 for (i, &tok) in prompt.iter().enumerate() {
6602 let (l, h) = self.spec_target_step_h(e, tok, &mut *cache)?;
6603 if let Some(ph) = prompt_h.as_mut() {
6604 e.copy_into(ph, i * n_embd, &h, n_embd)?;
6605 }
6606 prime_logits = l;
6607 }
6608 }
6609 e.stream().synchronize()?;
6610 // PREFIX-CACHE SEED CAPTURE (lane/spec-prefix-cache): boundary == prompt end (the seed
6611 // case — no shared-prefix split, publish the whole prompt). The prime just finished, so
6612 // cache.pos == base + prompt.len() and the recurrent state IS the boundary state;
6613 // prime_logits are the boundary logits. Cold sessions only (base == 0) — same law as
6614 // prime_split. The mid-prompt capture above already consumed the request if it matched.
6615 if !continuation && base == 0 {
6616 if let Some((requested, slot)) = sess_capture.as_mut() {
6617 if *requested == Some(prompt.len()) && slot.is_none() {
6618 debug_assert_eq!(cache.pos, prompt.len(), "seed capture off prompt end");
6619 if let Ok(snap) = cache.snapshot(e) {
6620 **slot = Some(SpecBoundaryCapture {
6621 snap,
6622 pos: prompt.len(),
6623 logits: prime_logits.clone(),
6624 });
6625 }
6626 }
6627 }
6628 }
6629 // Harness timing contract (see crate::PRIME_NANOS): gen-only throughput without the
6630 // prime-subtraction hack.
6631 crate::PRIME_NANOS.store(
6632 t_prime.elapsed().as_nanos() as u64,
6633 std::sync::atomic::Ordering::Relaxed,
6634 );
6635
6636 let (embd_qt, embd_rb) = self.embd.qt_and_row_bytes(n_embd);
6637 // Resident table is fastest when it fits. Large spill deployments can preserve that HBM
6638 // for expert-cache slots and gather only the exact rows needed by MTP/verify from host.
6639 let host_embd = spec_host_embd();
6640 let embd_gpu = if host_embd {
6641 None
6642 } else {
6643 Some(
6644 self.embd_gpu
6645 .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload")),
6646 )
6647 };
6648 let embd_dev = embd_gpu.map(|g| (g, embd_qt, embd_rb));
6649 if host_embd {
6650 eprintln!(
6651 "[spec] host-row embedding: {} bytes kept off HBM",
6652 self.embd.raw.len()
6653 );
6654 }
6655 let mut out: Vec<u32> = Vec::with_capacity(max_new);
6656 let mut total_drafted = 0usize;
6657 let mut total_accepted = 0usize;
6658
6659 // First generated token = argmax of the prompt's last logits (== greedy's first token).
6660 // Emit it, then FEED it to establish the loop invariant below.
6661 // PENDING-CARRY: the carried bonus was already emitted by the LAST burst — it becomes
6662 // last_token WITHOUT re-emission, and round 0 consumes it as pending (no init feed).
6663 // CONSTRAINED entry rules: the first emitted token is the MASKED argmax of the
6664 // prompt's last logits (plain constrained-greedy identity); a continuation without
6665 // a carried pending would emit an UNMASKED stashed next_pred — refused loudly (the
6666 // worker never resumes constrained sessions from the pool, so this cannot fire).
6667 if let Some(c) = constraint.as_deref_mut() {
6668 if continuation && carried_pending.is_none() {
6669 return Err("constrained spec continuation requires a carried pending \
6670 (pool resume is unconstrained-only)"
6671 .into());
6672 }
6673 if !continuation {
6674 c.mask_logits(&mut prime_logits)
6675 .map_err(|e2| format!("constraint: {e2}"))?;
6676 }
6677 }
6678 let mut last_token = if let Some(b) = carried_pending {
6679 b
6680 } else if continuation {
6681 sess_tail.as_ref().unwrap().2.unwrap()
6682 } else {
6683 argmax(&prime_logits) as u32
6684 };
6685 if carried_pending.is_none() {
6686 out.push(last_token);
6687 // grammar advances with every emitted token (carried pendings were consumed
6688 // by the burst that emitted them).
6689 if let Some(c) = constraint.as_deref_mut() {
6690 c.consume(last_token)
6691 .map_err(|e2| format!("constraint: {e2}"))?;
6692 }
6693 }
6694 if continuation {
6695 // draft-KV invariant: entries [0..base) are the session's exact fills; truncate any
6696 // overhang so the chain's first append lands at slot base (== committed.len()).
6697 scratch.set_len(e, base)?;
6698 }
6699 // sse-cadence: hand the caller every not-yet-flushed token (disjoint in-order slices
6700 // concatenating to the full `out`). Called after the prime's first token and after each
6701 // round commit — emission timing only, token bytes untouched. The slice may be EMPTY
6702 // (poll-only boundary: zero-round folds commit nothing new); returns the caller's
6703 // continue-verdict (admission yield, 2026-08-06) — false ends the burst at this round.
6704 fn flush_commit(
6705 cb: &mut Option<&mut dyn FnMut(&[u32]) -> bool>,
6706 out: &[u32],
6707 flushed: &mut usize,
6708 ) -> bool {
6709 if let Some(f) = cb.as_mut() {
6710 let keep = f(&out[*flushed..]);
6711 *flushed = out.len();
6712 keep
6713 } else {
6714 true
6715 }
6716 }
6717 keep_going = flush_commit(&mut on_commit, &out, &mut flushed);
6718 // INVARIANT at loop top: `last_token` is the most-recently-committed/emitted token, its
6719 // KV+recur state IS in `cache` (cache.pos = position right AFTER last_token), `last_pred`
6720 // is the greedy ARGMAX of the logits that predict the token FOLLOWING last_token, and
6721 // `h_seed` = last_token's pre-output_norm hidden. Establish it by feeding last_token once
6722 // (mirrors plain greedy). DEVICE-ARGMAX lever: the accept walk only ever consumes the
6723 // argmax of those logits — never the full vector — so a host u32 replaces the Vec<f32>.
6724 // --- SAMPLED SPEC (MEMRA_SPEC_TEMP>0, research/sampled-spec-impl-map.md): rejection-
6725 // sampling verify (Leviathan/Chen) — accept draft x at u < p(x)/q(x), resample from
6726 // norm(max(0,p-q)) on reject, bonus sampled from p on full accept. Counter-based Philox
6727 // everywhere (seed, event) -> reproducible. temp==0/unset = the greedy path, untouched.
6728 let sp = sampling.unwrap_or_else(|| SpecSampling {
6729 temp: std::env::var("MEMRA_SPEC_TEMP")
6730 .ok()
6731 .and_then(|v| v.parse().ok())
6732 .unwrap_or(0.0),
6733 seed: std::env::var("MEMRA_SEED")
6734 .ok()
6735 .and_then(|v| v.parse().ok())
6736 .unwrap_or(42),
6737 top_k: std::env::var("MEMRA_TOP_K")
6738 .ok()
6739 .and_then(|v| v.parse().ok())
6740 .unwrap_or(0),
6741 top_p: std::env::var("MEMRA_TOP_P")
6742 .ok()
6743 .and_then(|v| v.parse().ok())
6744 .unwrap_or(1.0),
6745 min_p: std::env::var("MEMRA_MIN_P")
6746 .ok()
6747 .and_then(|v| v.parse().ok())
6748 .unwrap_or(0.0),
6749 penalty_last_n: std::env::var("MEMRA_PENALTY_LAST_N")
6750 .ok()
6751 .and_then(|v| v.parse().ok())
6752 .unwrap_or(0),
6753 penalty_repeat: std::env::var("MEMRA_PENALTY_REPEAT")
6754 .ok()
6755 .and_then(|v| v.parse().ok())
6756 .unwrap_or(1.0),
6757 penalty_freq: std::env::var("MEMRA_PENALTY_FREQ")
6758 .ok()
6759 .and_then(|v| v.parse().ok())
6760 .unwrap_or(0.0),
6761 penalty_present: std::env::var("MEMRA_PENALTY_PRESENT")
6762 .ok()
6763 .and_then(|v| v.parse().ok())
6764 .unwrap_or(0.0),
6765 });
6766 let (sp_temp, sp_seed) = (sp.temp, sp.seed);
6767 let sampled = sp_temp > 0.0;
6768 // Trimmed heads: q lives on the trimmed vocab; accept gathers use the TRIMMED index and
6769 // the residual scatters q into target-id space (q=-inf off-trim — the head cannot propose
6770 // those, so their residual mass is p(x), correct by construction).
6771 let d2t_dev: Option<CudaSlice<u32>> = if sampled || crate::spec::spec_stream() {
6772 match &mtp.d2t {
6773 Some(map) => Some(e.htod_u32_v(map)?),
6774 None => None,
6775 }
6776 } else {
6777 None
6778 };
6779 let mut q_full_buf: Option<CudaSlice<f32>> = None;
6780 // Counters resume from the session (burst continuity: randomness must never repeat
6781 // across generate_spec_session calls); one-shot callers start at (0,0). Read through
6782 // sess_tail — `sess` was take()n into it above, so sess.as_ref() here is always None.
6783 let mut sctr: u32 = sess_tail.as_ref().map(|(_, _, _, s, _)| **s).unwrap_or(0);
6784 let mut uctr: u32 = sess_tail.as_ref().map(|(_, _, _, _, u)| **u).unwrap_or(0);
6785 // host Philox4x32-10 (mirrors spec_sample.cu; independent stream via ctr_lo tag)
6786 let host_u01 = |seed: u64, ctr: u32| -> f32 {
6787 let (m0, m1) = (0xD2511F53u32, 0xCD9E8D57u32);
6788 let (mut c0, mut c1, mut c2, mut c3) = (0xFFFF_FFFEu32, ctr, 0u32, 0u32);
6789 let (mut k0, mut k1) = ((seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
6790 for _ in 0..10 {
6791 let (h0, l0) = (((m0 as u64 * c0 as u64) >> 32) as u32, m0.wrapping_mul(c0));
6792 let (h1, l1) = (((m1 as u64 * c2 as u64) >> 32) as u32, m1.wrapping_mul(c2));
6793 let (n0, n1, n2, n3) = (h1 ^ c1 ^ k0, l1, h0 ^ c3 ^ k1, l0);
6794 c0 = n0;
6795 c1 = n1;
6796 c2 = n2;
6797 c3 = n3;
6798 k0 = k0.wrapping_add(0x9E3779B9);
6799 k1 = k1.wrapping_add(0xBB67AE85);
6800 }
6801 (c0 as f32 + 1.0) * (1.0 / 4294967296.0)
6802 };
6803 let mut draft_logits: Vec<CudaSlice<f32>> = Vec::new(); // retained head logits (q), per slot
6804 let mut draft_stats: Vec<(f32, f32, f32)> = Vec::new(); // (row_max, th_e, z_e) per slot
6805 let mut perturb_buf: Option<CudaSlice<f32>> = None; // gumbel scratch (max(n_vocab,d_vocab))
6806 let mut sample_tok = e.alloc_u32_zeroed(1)?; // residual/bonus sample out
6807 let mut col_buf: Option<CudaSlice<f32>> = None; // materialized verify column
6808 // Penalties (v2.1): applied to COPIES of q rows and p columns symmetrically (exactness
6809 // for the penalized+filtered target). History = generated tokens, host-tracked window.
6810 let pen_on = sampled
6811 && sp.penalty_last_n > 0
6812 && (sp.penalty_repeat != 1.0 || sp.penalty_freq != 0.0 || sp.penalty_present != 0.0);
6813 let mut pen_hist: Vec<u32> = if pen_on {
6814 prompt.iter().rev().take(64).rev().cloned().collect() // llama-parity: history spans prompt tail too
6815 } else {
6816 Vec::new()
6817 };
6818 let mut pen_hist_d: Option<CudaSlice<u32>> = None;
6819 let mut pcol_buf: Option<CudaSlice<f32>> = None; // penalized p-column scratch
6820 // MEMRA_SPEC_SETUP_TRACE=1 (diagnostics): per-call wall decomposition of the burst
6821 // SETUP + TAIL segments (the round loop's internals are MEMRA_SPEC_PHASE's job) —
6822 // built to pin the serve per-burst fixed cost (research/spec-serving-20260801).
6823 let setup_trace = std::env::var("MEMRA_SPEC_SETUP_TRACE").as_deref() == Ok("1");
6824 let t_ent = std::time::Instant::now();
6825
6826 // SESSION-AFFINITY TURN CHECKPOINT (lane/session-affinity, 2026-08-05): capture the
6827 // PROMPT-END boundary state so a LATER turn can rewind here and re-prime only its own
6828 // delta instead of the whole conversation. See `SpecCheckpoint` for why this boundary is
6829 // the one that matters (a history-rewriting client mutates what the session GENERATED,
6830 // so the next turn's prompt agrees with this one up to exactly here).
6831 //
6832 // WHERE — AND WHY THIS EXACT LINE. Right after the trunk prime, BEFORE the init feed
6833 // (`decode_step_h(last_token)`) and before round 0: the last instant at which the caches
6834 // hold exactly `base + prompt.len()` rows and nothing generated.
6835 //
6836 // This was WRONG in the first cut of this lane: the capture sat after the draft-KV fill,
6837 // which is also after the init feed, so `cache.pos` was `base + prompt.len() + 1` — the
6838 // boundary included the FIRST GENERATED TOKEN. That token is the first thing inside the
6839 // `<think>` block the client strips, so every later turn's diff diverged exactly one
6840 // token below the checkpoint and affinity declined 100% of the time. Measured on the
6841 // owner regime: "history diverged at 12233 of checkpoint 12234". The off-by-one made the
6842 // whole mechanism inert while looking, from the outside, like a working
6843 // correctness-declines-safely path — hence the decline log carries the offsets.
6844 //
6845 // The full-attn planes are `len`-truncatable so the snapshot copies only the GDN conv/ssm
6846 // state (the reason a spec session could not rewind before). The draft scratch needs no
6847 // copy: rows below the boundary are rewritten by the next turn's own fill.
6848 //
6849 // WHEN: non-empty prime only. An empty-suffix continuation burst adds no prompt boundary
6850 // (its "prompt end" IS the previous checkpoint's, already held), so it keeps the existing
6851 // checkpoint rather than replacing it with a strictly worse one.
6852 //
6853 // FAILURE IS SILENT BY DESIGN: on a VRAM-tight rig the snapshot alloc can fail. That
6854 // costs the NEXT turn its rewind (it re-primes fully, today's behavior) and must never
6855 // fail the burst that is already running — so the error is swallowed, loud only under
6856 // MEMRA_DEBUG_SPEC.
6857 if let Some(slot) = sess_ckpt_slot {
6858 if !continuation {
6859 let pos = cache.pos;
6860 debug_assert_eq!(
6861 pos,
6862 base + prompt.len(),
6863 "turn checkpoint must sit at the prompt end, before the init feed"
6864 );
6865 let anchor: Result<CudaSlice<f32>, Box<dyn std::error::Error>> =
6866 if let Some(ph) = &prompt_h {
6867 // hidden of the LAST primed row = the predecessor anchor at this
6868 // boundary (exactly what a fresh prime of committed[..pos] leaves in
6869 // last_h, and what the next prime's fill reads for its first row).
6870 let np = prompt.len();
6871 e.uninit(n_embd).and_then(|mut a| {
6872 e.copy_view_into(
6873 &mut a,
6874 0,
6875 &ph.slice((np - 1) * n_embd..np * n_embd),
6876 n_embd,
6877 )?;
6878 Ok(a)
6879 })
6880 } else {
6881 Err("no prompt hiddens".into())
6882 };
6883 match (cache.snapshot(e), anchor) {
6884 (Ok(snap), Ok(last_h)) => {
6885 *slot = Some(SpecCheckpoint { snap, pos, last_h });
6886 }
6887 (s, a) => {
6888 *slot = None; // a stale checkpoint would rewind to the WRONG boundary
6889 if std::env::var("MEMRA_DEBUG_SPEC").is_ok() {
6890 let err = s
6891 .err()
6892 .map(|e| e.to_string())
6893 .or_else(|| a.err().map(|e| e.to_string()))
6894 .unwrap_or_default();
6895 eprintln!(
6896 "[spec] turn checkpoint skipped ({err}); \
6897 next turn re-primes in full"
6898 );
6899 }
6900 }
6901 }
6902 }
6903 }
6904 // INIT FEED — skipped on a pending carry: last_token (the carried bonus) is NOT in the
6905 // caches and must NOT be fed solo; round 0's batched verify commits it as col 0. Its
6906 // seed/anchor hidden is the carried last_h (copied below); last_pred is dead in the
6907 // pending path (t_pred reads verify col 0 — the accept walk overwrites it).
6908 let mut last_pred = 0u32;
6909 let mut last_col_logits: Option<CudaSlice<f32>> = None;
6910 // CONSTRAINED: the init feed's logits back the (n_acc==0, base==0) masked-argmax
6911 // recompute in the grammar-truncation walk — retained host-side, round 0 only.
6912 let mut init_logits_host: Option<Vec<f32>> = None;
6913 let h_seed0: CudaSlice<f32> = if carried_pending.is_none() {
6914 let (init_logits, h) = self.spec_target_step_h(e, last_token, &mut *cache)?;
6915 last_pred = argmax(&init_logits) as u32;
6916 if constraint.is_some() {
6917 init_logits_host = Some(init_logits.clone());
6918 }
6919 // sampled mode: p-distribution after last_token, for the j==0/base==0 accept test.
6920 if sampled {
6921 last_col_logits = Some(e.htod(&init_logits)?);
6922 }
6923 h
6924 } else {
6925 // predecessor-row anchor: hidden of the last COMMITTED row (the carry contract).
6926 let lh = sess_tail
6927 .as_ref()
6928 .unwrap()
6929 .1
6930 .as_ref()
6931 .expect("pending carry requires last_h");
6932 e.clone_dtod(lh)?
6933 };
6934 let t_init = t_ent.elapsed();
6935 let mut last_col_stats: Option<(f32, f32, f32)> = None;
6936 // PERSISTENT h_seed buffer (allocated BEFORE any graph capture so no captured scratch can
6937 // alias it): every path that updates the round seed copies INTO it — no per-round allocs,
6938 // stable pointer for the graph-draft round-start copy.
6939 let mut h_seed_buf = e.clone_dtod(&h_seed0)?;
6940 // Predecessor-pairing trackers: `fill_prev` = trunk hidden AT the last COMMITTED row (the
6941 // predecessor of the next verify's col 0 — the reference's carried pending-h analogue;
6942 // also the predecessor-row hidden for the round-0 legacy-replay seed). At round 0 that
6943 // row is last_token's own (h_seed0). The chain step-0 seed under the pairing default =
6944 // hidden of the row BEFORE last_token = the prompt's last row at round 0 (h_seed_buf
6945 // overwritten below).
6946 let mut fill_prev = e.clone_dtod(&h_seed0)?;
6947 {
6948 if let Some(ph) = &prompt_h {
6949 let np = prompt.len();
6950 e.copy_view_into(
6951 &mut h_seed_buf,
6952 0,
6953 &ph.slice((np - 1) * n_embd..np * n_embd),
6954 n_embd,
6955 )?;
6956 } else if continuation {
6957 if let Some((_, lh, _, _, _)) = sess_tail.as_ref() {
6958 if let Some(lh) = lh.as_ref() {
6959 e.copy_into(&mut h_seed_buf, 0, lh, n_embd)?;
6960 }
6961 }
6962 }
6963 }
6964 // Persistent device prediction slots for the accept walk (max k+1 verify columns).
6965 let mut preds_d = e.alloc_u32_zeroed(k + 2)?;
6966
6967 let debug_spec = std::env::var("MEMRA_DEBUG_SPEC").is_ok();
6968 let fork_mode = OptiForkGateMode::configured();
6969 // MEMRA_SPEC_STATS=1: per-slot accept histogram + draft-length histogram, printed once at
6970 // the end. Metric normalization vs the reference engine: BOTH engines count
6971 // accepted/drafted where the chain stopped at p-min and the sub-threshold token is
6972 // discarded uncounted — per-slot decay + chain-length mix are the extra dimensions.
6973 let spec_stats = std::env::var("MEMRA_SPEC_STATS").is_ok();
6974 let mut st_drafted = vec![0usize; k];
6975 let mut st_accepted = vec![0usize; k];
6976 let mut st_len_hist = vec![0usize; k + 1];
6977 let mut st_full = 0usize;
6978 // P-MIN CONFIDENCE GATE (MEMRA_SPEC_PMIN, the serve script's --spec-draft-p-min mechanism):
6979 // stop the draft chain early when the head's softmax confidence in its own pick drops
6980 // below p_min. Hoisted above the loop: the graph capture bakes the prob kernels iff on.
6981 static PMIN: std::sync::OnceLock<f32> = std::sync::OnceLock::new();
6982 let p_min = *PMIN.get_or_init(|| {
6983 std::env::var("MEMRA_SPEC_PMIN")
6984 .ok()
6985 .and_then(|v| v.parse().ok())
6986 .unwrap_or(0.0)
6987 });
6988 // ZERO-DRAFT ROUNDS (MEMRA_SPEC_PMIN0=1, vendored from llama.cpp's draft gating): let the
6989 // p-min gate apply at j==0 too, so a low-confidence round drafts NOTHING and the verify
6990 // batch is just the pending bonus (m=1 = a plain decode step). llama's 35B win rides
6991 // exactly this — draft acceptance 76% at mean len 2.5 because unpredictable stretches
6992 // never pay draft+verify overhead. Only legal when a pending bonus exists (an empty
6993 // verify batch is not); the j==0 exemption stays for pending-less rounds.
6994 let pmin0 = std::env::var("MEMRA_SPEC_PMIN0")
6995 .map(|v| v == "1")
6996 .unwrap_or(false);
6997
6998 // --- GRAPH DRAFT setup: persistent I/O buffers + ONE capture (2 warmups inside). The
6999 // warmups mutate scratch len_d / pos / tok / seed — all reset at every round start, so the
7000 // only restore needed is the scratch counter. Capture failure (e.g. a non-capturable
7001 // cuBLAS path in an exotic head) falls back to the eager draft chain.
7002 // PER-SESSION PERSISTENCE (2026-08-01): session calls reuse the DraftGraphCtx parked on
7003 // the SpecSession — the capture (2 warmup head forwards + instantiate) ran ONCE at the
7004 // session's first burst, not per burst (measured ~16ms/burst fixed cost on H100 q27,
7005 // research/spec-serving-20260801). Reuse is pointer-exact: the graph bakes the session's
7006 // own scratch KV (never realloc'd), the model's resident embedding, the OnceLock p_min,
7007 // and the g_* buffers carried in the ctx — replay dispatch is identical to a fresh
7008 // capture, so draft tokens are bit-identical (drafts never decide exactness anyway; the
7009 // verify arbitrates). Single-shot calls (sess=None) build a fresh ctx and drop it.
7010 let mut dctx: DraftGraphCtx = match sess_draft_slot.as_mut().and_then(|s| s.take()) {
7011 Some(c) => c,
7012 None => DraftGraphCtx::new(e, n_embd, if sampled { d_vocab } else { 1 })?,
7013 };
7014 // A session that ran greedy bursts first sized g_q/g_perturb at 1; a sampled resume
7015 // needs d_vocab. Realloc is legal exactly while graph_s is None (nothing baked them).
7016 if sampled && dctx.g_q.len() < d_vocab {
7017 dctx.g_q = e.zeros(d_vocab)?;
7018 dctx.g_perturb = e.zeros(d_vocab)?;
7019 }
7020 // DRAFT-SIDE GRAMMAR MASK (lane/draft-mask, 2026-08-04): the drafter samples the
7021 // grammar's legal set, so proposals are legal BY CONSTRUCTION and the verify-side
7022 // truncation (the correctness backstop) stops cutting every tight-schema round.
7023 // The mask is one node inside the captured draft chain — presence is a CAPTURE-TIME
7024 // shape, so a parked graph of the other shape is dropped and recaptured.
7025 let dmask_on = constraint
7026 .as_deref()
7027 .is_some_and(|c| c.draft_mask_enabled());
7028 let dmask_words = if dmask_on { d_vocab.div_ceil(32) } else { 0 };
7029 if dmask_on && dctx.g_dmask.len() < dmask_words {
7030 dctx.g_dmask = e.alloc_u32_zeroed(dmask_words)?;
7031 dctx.graph = None; // the old capture baked the old (or no) mask pointer
7032 dctx.failed.clear_greedy();
7033 dctx.keeper.clear();
7034 }
7035 if dctx.graph.is_some() && dctx.graph_masked != dmask_on {
7036 dctx.graph = None;
7037 dctx.failed.clear_greedy();
7038 dctx.keeper.clear();
7039 }
7040 if graph_draft && !sampled && dctx.graph.is_none() && !dctx.failed.greedy_failed() {
7041 let DraftGraphCtx {
7042 g_tok,
7043 g_pos,
7044 g_seed,
7045 g_p,
7046 g_dmask,
7047 ..
7048 } = &mut dctx;
7049 // capture-time contents: ALL-ONES (ban nothing). A replay only ever runs after the
7050 // host uploads the position's real words, so the warmups stay grammar-free.
7051 if dmask_on {
7052 e.htod_u32_into(g_dmask, &vec![u32::MAX; dmask_words])?;
7053 }
7054 let g_dmask_ro: &CudaSlice<u32> = &*g_dmask;
7055 // CAPTURE-RETAIN (#68 fix): the warmup transients' pool addresses are baked into the
7056 // captured graph; the keeper pins them for the graph's lifetime. capture_graph (non-
7057 // retained) freed them at exit — safe for one-shot generate_spec (nothing else touches
7058 // the pool between replays) but WRONG for sessions: burst-boundary prime/fill/commit
7059 // passes (and, in serve, other sessions) recycle those addresses and the replay then
7060 // clobbers live buffers — the ST serve-spec corruption (research/serve-st-20260803).
7061 let cap_res = e.capture_graph_retained(|e| {
7062 self.mtp_head_forward_cap(
7063 e,
7064 mtp,
7065 g_tok,
7066 g_pos,
7067 g_seed,
7068 g_p,
7069 &mut *scratch,
7070 p_min > 0.0 || fork_mode == OptiForkGateMode::Controller,
7071 true,
7072 embd_gpu.expect("graph draft requires resident embedding"),
7073 embd_qt,
7074 embd_rb,
7075 d_vocab,
7076 None,
7077 None,
7078 if dmask_on {
7079 Some((g_dmask_ro, dmask_words))
7080 } else {
7081 None
7082 },
7083 )
7084 });
7085 match cap_res {
7086 Ok((g, keep)) => {
7087 scratch.set_len(e, base)?;
7088 dctx.graph = Some(g);
7089 dctx.graph_masked = dmask_on;
7090 dctx.keeper = keep;
7091 }
7092 Err(err) => {
7093 scratch.set_len(e, base)?;
7094 // LOUD flip (audit Q2): a dropped draft graph is a coverage loss, never
7095 // silent. Once per flip — mark returns None on an already-failed ctx.
7096 if let Some(line) = dctx.failed.mark_greedy(&err.to_string()) {
7097 eprintln!("{line}");
7098 }
7099 }
7100 }
7101 }
7102 // --- SAMPLED GRAPH DRAFT setup (step 3 of the sampled-spec arc): a SECOND capture, own
7103 // graph object, built only when sampled && graph-eligible — the greedy capture above is
7104 // untouched (and skipped when sampled: its graph would never be launched). Same head
7105 // forward, but the in-graph argmax reads GUMBEL-PERTURBED logits; the Philox event
7106 // counter lives in the persistent device g_ctr (bumped in-graph, host-seeded from sctr
7107 // once per round); the raw head logits land in the persistent g_q for the host's
7108 // per-replay async D2D into the round's q slot (q_slots, K x d_vocab, allocated once).
7109 // seed/temp are capture-time constants — baked into graph_s, so a pool-resumed request
7110 // with a different (seed, temp, k) drops the parked sampled graph and recaptures.
7111 // COST OF THE FRESH-SEED SERVE DEFAULT (dogfood F4, 2026-08-04): omitting `seed` on a
7112 // serve request now draws fresh per-request entropy (it used to default to a pinned 0),
7113 // so a seed-omitting request that RESUMES a parked spec session finds an s_key baked
7114 // with the PREVIOUS request's seed and pays one recapture. Bounded, and it does not
7115 // reopen the ~16ms/burst regression the persistent ctx exists to fix: a session's seed
7116 // is fixed for its whole lifetime (worker.rs reads s.sampler.seed() per burst), so
7117 // this compare misses at most ONCE per resumed request — the first burst recaptures
7118 // and every later burst in that request replays. A client that wants the parked graph
7119 // AND reproducibility supplies an explicit `seed`, honored exactly, which keeps s_key
7120 // stable across its whole conversation.
7121 // COMPOSITION RULE (fspec x gsd merge): the in-graph chain samples from the RAW
7122 // softmax — it can hold neither per-row filter stats nor the varying penalty history.
7123 // The sampled graph therefore engages only in the PURE-TEMP regime; filters/penalties
7124 // force the eager draft (which computes stats/penalties per row).
7125 let pure_temp = sp.top_k == 0 && sp.top_p >= 1.0 && sp.min_p <= 0.0 && !pen_on;
7126 let s_key = (sp_seed, sp_temp.to_bits(), k);
7127 if sampled && dctx.s_key.is_some_and(|old| old != s_key) {
7128 dctx.graph_s = None;
7129 dctx.failed.clear_sampled();
7130 dctx.s_key = None;
7131 dctx.q_slots.clear();
7132 dctx.keeper_s.clear();
7133 }
7134 if graph_draft
7135 && sampled
7136 && pure_temp
7137 && dctx.graph_s.is_none()
7138 && !dctx.failed.sampled_failed()
7139 {
7140 let DraftGraphCtx {
7141 g_tok,
7142 g_pos,
7143 g_seed,
7144 g_p,
7145 g_ctr,
7146 g_perturb,
7147 g_q,
7148 ..
7149 } = &mut dctx;
7150 // CAPTURE-RETAIN (#68 fix): same keeper contract as the greedy capture above.
7151 let cap_res = e.capture_graph_retained(|e| {
7152 self.mtp_head_forward_cap(
7153 e,
7154 mtp,
7155 g_tok,
7156 g_pos,
7157 g_seed,
7158 g_p,
7159 &mut *scratch,
7160 p_min > 0.0,
7161 true,
7162 embd_gpu.expect("graph draft requires resident embedding"),
7163 embd_qt,
7164 embd_rb,
7165 d_vocab,
7166 Some((g_ctr, g_perturb, g_q, sp_seed, sp_temp)),
7167 None,
7168 None, // constrained spec is greedy-only — sampled never carries a hook
7169 )
7170 });
7171 match cap_res {
7172 Ok((g, keep)) => {
7173 scratch.set_len(e, base)?;
7174 for _ in 0..k {
7175 dctx.q_slots.push(e.zeros(d_vocab)?);
7176 }
7177 dctx.graph_s = Some(g);
7178 dctx.s_key = Some(s_key);
7179 dctx.keeper_s = keep;
7180 }
7181 Err(err) => {
7182 scratch.set_len(e, base)?;
7183 // LOUD flip (audit Q2): same contract as the greedy capture above.
7184 if let Some(line) = dctx.failed.mark_sampled(&err.to_string()) {
7185 eprintln!("{line}");
7186 }
7187 }
7188 }
7189 }
7190 let t_cap = t_ent.elapsed();
7191 // PERSISTENT DRAFT KV: fill the MTP block's K/V for every prompt position from the exact
7192 // trunk hiddens collected during prime — ONE batched K/V-only pass (overwrites any
7193 // capture-warmup garbage; capture left len at 0). last_token (the init feed) needs no
7194 // fill: the first chain step processes it and appends its entry at slot prompt.len().
7195 if let Some(ph) = &prompt_h {
7196 // SESSION: rows [0..base) are the previous turns' exact fills (refresh overwrote them
7197 // with true verify hiddens) — truncate any draft overhang, fill ONLY the suffix at
7198 // global positions [base..base+tp). Fresh call: base==0, identical to before.
7199 scratch.set_len(e, base)?;
7200 // CHUNKED FILL (long-ctx OOM fix, 2026-07-05): mtp_kv_fill's transients scale with its
7201 // T (concat = T*2*n_embd*4B — 1.5GB at 40k) and its concat loop is 2*T launches. The
7202 // fill is a pure sequential append, so chunking is exact: each chunk appends its rows
7203 // at pos0=base+start with the identical per-row math. Same knob as the trunk prime.
7204 let tp = prompt.len();
7205 let fill_chunk: usize = if crate::cache::swa_ring_on() {
7206 crate::hybrid_forward::prime_chunk_tokens(tp, self.layers.len())
7207 } else {
7208 // Preserve the flag-OFF schedule byte-for-byte, including the legacy zero value
7209 // meaning one monolithic fill.
7210 std::env::var("MEMRA_PRIME_CHUNK")
7211 .ok()
7212 .and_then(|v| v.parse().ok())
7213 .unwrap_or(4096)
7214 };
7215 let fill_chunk = if fill_chunk == 0 { tp } else { fill_chunk };
7216 let mut start = 0usize;
7217 while start < tp {
7218 let end = (start + fill_chunk).min(tp);
7219 let tc = end - start;
7220 {
7221 // PREDECESSOR pairing: row i gets h[i-1]; global row 0 a zeros row (the
7222 // reference engine's initial pending-h is zeroed too); a session turn's row 0
7223 // gets the PREVIOUS turn's last committed hidden (sess.last_h). Per chunk:
7224 // rows start..end read h[start-1..end-1] — one dtod into a chunk buffer.
7225 let mut phs = e.zeros(tc * n_embd)?;
7226 let (src_lo, dst_off) = if start == 0 {
7227 (0, n_embd)
7228 } else {
7229 ((start - 1) * n_embd, 0)
7230 };
7231 let n_copy = if start == 0 {
7232 (tc - 1) * n_embd
7233 } else {
7234 tc * n_embd
7235 };
7236 if start == 0 {
7237 if let Some((_, lh, _, _, _)) = sess_tail.as_ref() {
7238 if let Some(lh) = lh.as_ref() {
7239 e.copy_into(&mut phs, 0, lh, n_embd)?;
7240 }
7241 }
7242 }
7243 if n_copy > 0 {
7244 e.copy_view_into(
7245 &mut phs,
7246 dst_off,
7247 &ph.slice(src_lo..src_lo + n_copy),
7248 n_copy,
7249 )?;
7250 }
7251 self.mtp_kv_fill(
7252 e,
7253 mtp,
7254 &prompt[start..end],
7255 &phs,
7256 base + start,
7257 &mut *scratch,
7258 embd_dev,
7259 )?;
7260 }
7261 start = end;
7262 }
7263 }
7264 // MEMRA_PROFILE_SPEC=2: profiler capture starts HERE — after the prime, so an
7265 // `nsys -c cudaProfilerApi` capture contains ONLY the round loop (draft/verify/commit).
7266 // (=1 brackets the whole call in run_spec.rs, prime included.)
7267 if std::env::var("MEMRA_PROFILE_SPEC").as_deref() == Ok("2") {
7268 unsafe extern "C" {
7269 fn cudaProfilerStart() -> i32;
7270 }
7271 unsafe {
7272 cudaProfilerStart();
7273 }
7274 }
7275 // ROUND-STREAM stage (c) 4 (MEMRA_SPEC_STREAM=1, experimental): pre-issued M-round
7276 // bursts with ZERO per-round host readbacks — the accept/seed/rollback/ring kernels
7277 // consume each other's device outputs; the host drains the ring every M rounds. v1
7278 // constraints: greedy, !spec_replay, single-shot, batched-linear layers, no refresh
7279 // fills (acceptance effect A/B-arbitrated), enters from round 1 (pending guaranteed).
7280 // NOTE: not gated on the caller's graph_draft (its trunk_dense conjunct turns the 35B
7281 // MoE off) — the stream capture encloses ONLY the dense MTP head; the head-dense /
7282 // full-prec / k gates are re-derived here and a failed capture degrades to stream-off.
7283 let stream_on = crate::spec::spec_stream()
7284 && !sampled
7285 && !spec_replay
7286 && constraint.is_none()
7287 && !session_mode
7288 && embd_gpu.is_some()
7289 && !crate::model::full_prec_enabled()
7290 && k + 2 < 96;
7291 let mut stream_graph: Option<cudarc::driver::CudaGraph> = None;
7292 let mut g_tokp2k = e.alloc_u32_zeroed(2 * k.max(1))?;
7293 if stream_on {
7294 let cap = e.capture_graph(|e| {
7295 for j in 0..k.max(1) {
7296 self.mtp_head_forward_cap(
7297 e,
7298 mtp,
7299 &mut dctx.g_tok,
7300 &mut dctx.g_pos,
7301 &mut dctx.g_seed,
7302 &mut dctx.g_p,
7303 &mut *scratch,
7304 true,
7305 true,
7306 embd_gpu.expect("round stream requires resident embedding"),
7307 embd_qt,
7308 embd_rb,
7309 d_vocab,
7310 None,
7311 Some((&mut g_tokp2k, j, d2t_dev.as_ref())),
7312 None, // round-stream requires constraint.is_none() (see stream_on)
7313 )?;
7314 }
7315 Ok(())
7316 });
7317 match cap {
7318 Ok(g) => {
7319 scratch.set_len(e, 0)?;
7320 stream_graph = Some(g);
7321 }
7322 Err(err) => {
7323 scratch.set_len(e, 0)?;
7324 if debug_spec {
7325 eprintln!("[spec] stream-graph capture failed ({err}); stream off");
7326 }
7327 }
7328 }
7329 }
7330 let stream_active = stream_on && stream_graph.is_some();
7331 if debug_spec {
7332 eprintln!(
7333 "[spec] stream_on={stream_on} env={} samp={sampled} dg={} captured={} active={stream_active} session={session_mode} replay={spec_replay}",
7334 crate::spec::spec_stream(),
7335 dctx.graph.is_some(),
7336 stream_graph.is_some()
7337 );
7338 }
7339 let t_v_s = k + 1;
7340 // ROUND-STREAM buffers + ptr tables now live in the model-generic round_stream
7341 // module (extracted 2026-07-12; the gemma burst reuses them).
7342 let sb = crate::round_stream::StreamBufs::new(e, k, crate::spec::spec_stream_m())?;
7343 let crate::round_stream::StreamBufs {
7344 mut vtok_d,
7345 mut brk_d,
7346 mut pend_d,
7347 last_pred_d,
7348 mut pos_ctr,
7349 mut pos_start_d,
7350 mut ring_d,
7351 acc_d: mut stream_acc,
7352 m_rounds,
7353 k: _,
7354 } = sb;
7355 let stream_ptrs: Option<CudaSlice<u64>> = if stream_active {
7356 Some(crate::round_stream::kv_len_ptr_table(
7357 e,
7358 cache,
7359 Some(&pos_ctr),
7360 )?)
7361 } else {
7362 None
7363 };
7364
7365 let t_fill = t_ent.elapsed();
7366 let mut round = 0usize;
7367 // ADAPTIVE DRAFT LENGTH (MEMRA_SPEC_ADAPT=1, opt-in — the gemma_spec accepted-run law,
7368 // ported 2026-08-01): next round's draft depth = last round's accepted run + 1, clamped
7369 // to [floor(pos), k_cap] — a miss shrinks the next draft to the miss point + 1,
7370 // full-accept streaks re-deepen one step per round. NOT the 2026-07-07 acceptance-EMA
7371 // (that arm measured an HONEST LOSS to static per-class optima — 115.0/85.8/73.4 vs
7372 // 121.6/92.7/75.6, EMA lag — and was removed 2026-07-08; rig5090.jsonl has the record).
7373 // The gemma law has no lag class: it reacts within one round, and was worth +7-20% on
7374 // the gemma cells at unchanged exactness (2026-07-10 flip; floor sweep 2026-07-25;
7375 // position key 2026-07-26). Signal = n_acc from the round's EXISTING accept readback —
7376 // zero new syncs; the draft graph is a SINGLE-STEP capture replayed per drafted token,
7377 // so a per-round depth needs no re-capture (unlike gemma's whole-chain graphs). qwen's
7378 // in-round p-min cut already shortens chains mid-round, so gemma's one-round-late p-min
7379 // fold into kc is unnecessary here — the accepted-run law sees the cut via n_acc.
7380 // Exactness is the verify's job at ANY depth (same contract as p-min variable rounds).
7381 // DEFAULT OFF on the qwen path until its cells gate a flip (gemma's is default-on).
7382 // MEASURED 2026-08-01 (H100 GPU-3, interleaved x3, NGEN=256, same-invocation plain
7383 // denominators; research/qwen-adaptive-k-20260801/): REFUTED on the tuned qwen configs.
7384 // q27 K=3+HPOST+PMIN=0.3: short +0.8% (noise; law ~idles, len_hist identical), board
7385 // -1.9%, agentic -0.5%; board PMIN=0 -2.1% (not p-min shadowing — the law itself);
7386 // floor=1 -2.8% (gemma's floor-collapse, reproduced). q35 K=2 board: -6.4% (52/136
7387 // rounds shrink to depth 1; no depth to reclaim at K=2). The gemma direction DOES
7388 // appear at untuned depth-K — q27 K=6 floor=4 +1.5% over fixed K=6 — but stays -3.7%
7389 // below fixed K=3: same verdict class as the retired EMA arm (honest loss to static
7390 // per-class optima). Acceptance-rate rises under the law while tokens/round falls —
7391 // it buys accept-% by adding rounds, and a round's fixed draft+verify cost wins.
7392 // K=1..8 self-consistency PASS both models with the law ON (exactness held).
7393 let adapt = std::env::var("MEMRA_SPEC_ADAPT").as_deref() == Ok("1");
7394 // floor: per-model default keyed on n_embd (gemma's tiering — models with an expensive
7395 // verify keep deep drafts after a miss); MEMRA_SPEC_ADAPT_FLOOR pins it everywhere.
7396 let adapt_floor_env: Option<usize> = std::env::var("MEMRA_SPEC_ADAPT_FLOOR")
7397 .ok()
7398 .and_then(|v| v.parse().ok());
7399 let adapt_floor_default: usize = if self.cfg.n_embd as usize >= 3500 {
7400 4
7401 } else if self.cfg.n_embd as usize >= 2500 {
7402 2
7403 } else {
7404 1
7405 };
7406 let adapt_floor: usize = adapt_floor_env.unwrap_or(adapt_floor_default);
7407 // position key: past floor_ctx a HIGH floor (>=4) relaxes to 1 — forced-deep drafts
7408 // turn net-negative at depth (gemma 31B d1736 evidence); MEMRA_SPEC_FLOOR_CTX moves
7409 // the boundary, an explicit MEMRA_SPEC_ADAPT_FLOOR pins the floor everywhere.
7410 let floor_ctx: usize = std::env::var("MEMRA_SPEC_FLOOR_CTX")
7411 .ok()
7412 .and_then(|v| v.parse().ok())
7413 .unwrap_or(1024);
7414 let floor_at = |pos: usize| -> usize {
7415 if adapt_floor_env.is_some() || pos < floor_ctx {
7416 adapt_floor
7417 } else if adapt_floor >= 4 {
7418 1
7419 } else {
7420 adapt_floor
7421 }
7422 };
7423 // cap: MEMRA_SPEC_CAPMAX (gemma semantics, default 7). Binds only under adapt — the
7424 // fixed-K default path is untouched by this whole block.
7425 let cap_max: usize = std::env::var("MEMRA_SPEC_CAPMAX")
7426 .ok()
7427 .and_then(|v| v.parse().ok())
7428 .unwrap_or(7);
7429 let k_cap = k.min(cap_max).max(1);
7430 let mut kc = k_cap;
7431 let mut opti_fork: Option<OptiForkState> = None;
7432 let mut fork_snapshot: Option<crate::cache::CacheSnapshot> = None;
7433 if fork_mode != OptiForkGateMode::Disabled {
7434 let fence = crate::pp::pp_cuts(self.layers.len());
7435 let refusal = if !session_mode {
7436 Some("not-session")
7437 } else if k != 1 || adapt {
7438 Some("requires-fixed-k1")
7439 } else if sampled || constraint.is_some() || spec_replay {
7440 Some("sampled-constrained-or-replay")
7441 } else if pipe.is_some() {
7442 Some("two-session-pipeline")
7443 } else if !spec_devacc() {
7444 Some("requires-device-accept")
7445 } else if stream_active || crate::spec::spec_stream() {
7446 Some("round-stream")
7447 } else if crate::cache::swa_ring_on() || cache.has_swa_ring() {
7448 Some("swa-ring")
7449 } else if crate::pp::pp_host_bounce_active() {
7450 Some("host-bounce")
7451 } else if fork_mode == OptiForkGateMode::Controller
7452 && cache.recur.iter().any(Option::is_some)
7453 {
7454 Some("controller-requires-zero-recurrent-state")
7455 } else if fence.as_ref().is_none_or(|f| f.len() != 3) {
7456 Some("requires-pp2")
7457 } else {
7458 None
7459 };
7460 if let Some(reason) = refusal {
7461 OPTI_FORK_REFUSALS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
7462 eprintln!("[opti-fork] refused reason={reason}");
7463 } else {
7464 let fence = fence.expect("validated PP-2 fence");
7465 let rt = crate::pp::PpNRt::get(e)?;
7466 let primary_stage0 = rt.engine(0, e).ctx().ordinal() == e.ctx().ordinal();
7467 let primary_stage1 = rt.engine(1, e).ctx().ordinal() == e.ctx().ordinal();
7468 let primary_supported =
7469 primary_stage0 || (fork_mode == OptiForkGateMode::Controller && primary_stage1);
7470 if !rt.cross_device() || !primary_supported {
7471 OPTI_FORK_REFUSALS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
7472 eprintln!("[opti-fork] refused reason=requires-supported-primary-cross-device");
7473 } else {
7474 // Both recurrent snapshots and both seed generations are allocated before
7475 // the first fork, each through its owning PP stage. Allocation failure
7476 // therefore happens before any optimistic state mutation can occur.
7477 let current_snapshot = opti_snapshot_stage_owned(e, cache, rt, &fence)?;
7478 let alternate_snapshot = opti_snapshot_stage_owned(e, cache, rt, &fence)?;
7479 let fork = OptiForkState::new(
7480 e,
7481 cache,
7482 fork_mode,
7483 alternate_snapshot,
7484 &h_seed_buf,
7485 &fill_prev,
7486 rt,
7487 fence[1],
7488 self.layers.len(),
7489 )?;
7490 eprintln!(
7491 "[opti-fork] armed mode={fork_mode:?} snapshots=2 seeds=2 split={} \
7492 payload_dev0={} payload_dev1={} q_threshold={:.3}",
7493 fence[1],
7494 fork.logical_payload_bytes[0],
7495 fork.logical_payload_bytes[1],
7496 fork.controller.map_or(0.0, |policy| policy.threshold),
7497 );
7498 fork_snapshot = Some(current_snapshot);
7499 opti_fork = Some(fork);
7500 }
7501 }
7502 }
7503 // Persistent snapshot buffers are allocated once and refreshed in place. The fork arm
7504 // uses stage-owned snapshots; refused/disabled arms retain the existing generic helper.
7505 let mut snap = match fork_snapshot {
7506 Some(snapshot) => snapshot,
7507 None => cache.snapshot(e)?,
7508 };
7509 let mut carried_opti: Option<OptiControllerTicket> = None;
7510 // ROUND-STREAM stage (b) 3a: device table of per-layer kvl.len_d pointers (stable — the
7511 // cache never reallocates len_d; see cache.rs "stable pointer" note). 0 = no KV layer.
7512 let kv_len_ptrs: Option<CudaSlice<u64>> = if spec_devacc() && !spec_replay {
7513 Some(crate::round_stream::kv_len_ptr_table(e, cache, None)?)
7514 } else {
7515 None
7516 };
7517 // BONUS FOLD (2026-07-04): after a FULL accept the bonus token is NOT committed with a
7518 // separate T=1 trunk pass (a full weight read per round). It stays PENDING and rides as
7519 // column 0 of the NEXT round's verify batch. Under predecessor pairing the next chain
7520 // seeds from the bonus's predecessor's TRUE verify hidden (free — no extra
7521 // pass of any kind). Verify still
7522 // checks every emitted token against the target -> exactness holds by construction; only
7523 // DRAFT QUALITY can shift, which the acceptance numbers arbitrate.
7524 // bonus emitted but not yet committed to cache. A carried pending (see SpecSession::
7525 // pending_tok) enters round 0 directly — the burst boundary becomes a plain round edge.
7526 let mut pending: Option<u32> = carried_pending;
7527 // MEMRA_SPEC_PHASE=1: per-round wall decomposition (draft / verify / accept+commit) —
7528 // no tracing, no extra syncs (each phase is naturally sync-bounded: draft readbacks,
7529 // the verify accept readback). Printed once at loop end via spec-stats.
7530 let anatomy_on = std::env::var("MEMRA_SPEC_PP_ANATOMY").as_deref() == Ok("1");
7531 let phase_on = anatomy_on || std::env::var("MEMRA_SPEC_PHASE").as_deref() == Ok("1");
7532 // DRAFT-MASK receipt (lane/draft-mask): speculative-clone wall + rounds, printed with
7533 // spec-stats. The clone is the one cost the design adds per round — measured, not assumed.
7534 let (mut dm_clone_ns, mut dm_rounds) = (0u128, 0usize);
7535 // grammar-truncation counters: how many rounds the verify-side cut fired and how many
7536 // already-verified tokens it threw away. THIS is the quantity draft masking targets.
7537 let (mut dm_cuts, mut dm_cut_tokens) = (0usize, 0usize);
7538 let (mut ph_draft, mut ph_verify, mut ph_rest) = (0f64, 0f64, 0f64);
7539 let mut ph_wait = 0f64;
7540 let mut ph_commit = 0f64;
7541 let mut ph_t = std::time::Instant::now();
7542 let mut ph_mark = |acc: &mut f64, on: bool| {
7543 if on {
7544 let now = std::time::Instant::now();
7545 *acc += (now - ph_t).as_secs_f64();
7546 ph_t = now;
7547 }
7548 };
7549 if let Some(p) = pipe {
7550 p.setup_end();
7551 }
7552 while keep_going && out.len() < max_new {
7553 // ROUND-STREAM BURST: from round 1 (pending guaranteed by every non-replay arm),
7554 // issue M rounds with zero readbacks, then drain the ring + reconcile mirrors.
7555 if let (true, Some(sg), Some(ptrs)) = (
7556 stream_active && round >= 1 && pending.is_some(),
7557 &stream_graph,
7558 &stream_ptrs,
7559 ) {
7560 if debug_spec {
7561 static ONCE: std::sync::Once = std::sync::Once::new();
7562 ONCE.call_once(|| {
7563 eprintln!("[memra] ROUND-STREAM burst engaged (M={m_rounds} k={k})")
7564 });
7565 }
7566 e.set_i32_one(&mut pos_ctr, cache.pos as i32)?;
7567 e.set_u32_one(&mut pend_d, pending.unwrap())?;
7568 e.set_u32_one(&mut ring_d, 0)?; // ring count = 0 (writes element 0)
7569 for _mi in 0..m_rounds {
7570 e.i32_copy_add(&pos_ctr, &mut pos_start_d, 0)?;
7571 cache.snapshot_into(e, &mut snap)?; // device D2Ds, stream-ordered
7572 e.i32_copy_add(&pos_ctr, &mut scratch.kv.len_d, 0)?; // draft-KV rollback
7573 e.i32_copy_add(&pos_ctr, &mut dctx.g_pos, 1)?; // rope pos = pos + base
7574 e.u32_copy(&pend_d, &mut dctx.g_tok)?;
7575 e.copy_into(&mut dctx.g_seed, 0, &h_seed_buf, n_embd)?;
7576 sg.launch()?;
7577 e.spec_assemble_verify(
7578 &g_tokp2k,
7579 &pend_d,
7580 d2t_dev.as_ref(),
7581 &mut vtok_d,
7582 &mut brk_d,
7583 p_min,
7584 k,
7585 pmin0,
7586 )?;
7587 let mut ck = VerifyCkpt::new(self.layers.len());
7588 let dummy = vec![0u32; t_v_s];
7589 let (tl_d, vx) = self.decode_step_t_core_stream(
7590 e,
7591 &dummy,
7592 0,
7593 &mut *cache,
7594 embd_dev,
7595 Some(&mut ck),
7596 Some((&vtok_d, &pos_ctr)),
7597 None,
7598 )?;
7599 for j in 0..t_v_s {
7600 e.argmax_token_device_col(&tl_d, j, n_vocab, &mut preds_d, j)?;
7601 }
7602 e.spec_accept_greedy_dc(
7603 &preds_d,
7604 &vtok_d,
7605 &last_pred_d,
7606 &brk_d,
7607 &mut stream_acc,
7608 )?;
7609 e.spec_seed_gather(&vx, &fill_prev, &stream_acc, &mut h_seed_buf, 1, n_embd)?;
7610 e.copy_into(&mut fill_prev, 0, &h_seed_buf, n_embd)?;
7611 self.commit_verified_prefix_stream(
7612 e,
7613 &mut *cache,
7614 &snap,
7615 &ck,
7616 &stream_acc,
7617 1,
7618 t_v_s,
7619 )?;
7620 e.spec_rollback_stream(
7621 ptrs,
7622 &pos_start_d,
7623 &stream_acc,
7624 1,
7625 self.layers.len() + 1,
7626 )?;
7627 e.spec_ring_commit(&vtok_d, &stream_acc, &brk_d, &mut ring_d, &mut pend_d)?;
7628 }
7629 e.stream().synchronize()?;
7630 let ring_h = e.dtoh_u32(&ring_d)?;
7631 let cnt = ring_h[0] as usize;
7632 for i in 0..cnt {
7633 if out.len() < max_new {
7634 out.push(ring_h[1 + i]);
7635 }
7636 }
7637 let pos_h = e.dtoh_i32(&pos_ctr)?[0] as usize;
7638 for il in 0..self.layers.len() {
7639 if let Some(kvl) = cache.kv[il].as_mut() {
7640 kvl.len = pos_h;
7641 }
7642 }
7643 cache.pos = pos_h;
7644 scratch.kv.len = pos_h;
7645 pending = Some(ring_h[cnt]); // last drained token = the live bonus
7646 last_token = ring_h[cnt];
7647 total_drafted += k * m_rounds; // upper bound (p-min breaks uncounted)
7648 total_accepted += cnt.saturating_sub(m_rounds);
7649 if let Some(t) = sess_telem {
7650 // totals only — the burst's per-round accept counts stayed on device
7651 // (that is the point of the round-stream arm). pos_* untouched.
7652 t.record_totals(m_rounds, k * m_rounds, cnt.saturating_sub(m_rounds));
7653 }
7654 round += m_rounds;
7655 // sse-cadence: the drained ring is committed — flush it at burst-drain cadence.
7656 keep_going = flush_commit(&mut on_commit, &out, &mut flushed);
7657 continue;
7658 }
7659 let pipe_draft = match pipe {
7660 Some(p) => Some(p.draft_begin(round)?),
7661 None => None,
7662 };
7663 let pos = cache.pos; // #tokens committed (EXCLUDES a pending bonus)
7664 let mut current_opti = carried_opti.take();
7665 let mut fork_generation = if current_opti.is_none() && pending.is_some() {
7666 match opti_fork.as_mut() {
7667 Some(fork) if fork.mode.is_forced() => Some(fork.reserve(&mut snap)?),
7668 None => None,
7669 Some(_) => None,
7670 }
7671 } else {
7672 None
7673 };
7674 if current_opti.is_none() {
7675 if let Some(fork) = opti_fork.as_ref() {
7676 opti_snapshot_stage_owned_into(e, cache, fork.rt, &fork.fence, &mut snap)?;
7677 } else {
7678 cache.snapshot_into(e, &mut snap)?;
7679 }
7680 } else if snap.pos != pos {
7681 return Err(format!(
7682 "optipipe carried snapshot pos {} != current pos {pos}",
7683 snap.pos
7684 )
7685 .into());
7686 } // §C: snapshot BEFORE draft+verify (already retained for a carried successor)
7687 ph_mark(&mut ph_rest, phase_on);
7688
7689 // --- 1. DRAFT k tokens with the NextN head (autoregressive, T=1 each) ---
7690 // p-min semantics (both paths): stop the chain early when the head's confidence in
7691 // its own pick drops below p_min — the just-drafted token is DISCARDED, but its
7692 // scratch append stands (identical to the eager chain's ordering). j==0 always drafts.
7693 let base0 = if pending.is_some() { 1usize } else { 0usize };
7694 // fixed draft length by default; MEMRA_SPEC_ADAPT=1 drafts at last round's
7695 // accepted run + 1 (the gemma law — see the setup block above the loop).
7696 let k_this = if adapt { kc } else { k };
7697 let mut draft: Vec<u32> = Vec::with_capacity(k);
7698 let mut draft_idx: Vec<u32> = Vec::with_capacity(k); // trimmed-vocab ids (== draft when untrimmed)
7699 let mut controller_draft_prob: Option<f32> = None;
7700 let mut controller_eager_state: Option<(u32, CudaSlice<f32>)> = None;
7701 if let Some(ticket) = current_opti.as_mut() {
7702 let carried_pending = pending.ok_or("optipipe carried successor lost pending")?;
7703 if ticket.verify_tokens[0] != carried_pending {
7704 return Err(format!(
7705 "optipipe carried pending mismatch: ticket={} live={carried_pending}",
7706 ticket.verify_tokens[0],
7707 )
7708 .into());
7709 }
7710 draft.push(ticket.verify_tokens[1]);
7711 controller_draft_prob = Some(ticket.draft_prob);
7712 controller_eager_state = ticket
7713 .take_eager_seed()
7714 .map(|seed| (ticket.verify_tokens[1], seed));
7715 } else {
7716 // Round-start draft-KV sync (BOTH paths). Persistent: truncate/align to the committed
7717 // history — slots 0..P hold entries for the tokens before last_token@P (P = pos +
7718 // base0 - 1); this single set_len IS the draft-side rollback (drops last round's
7719 // rejected drafts and p-min extras via the len mechanism).
7720 scratch.set_len(e, pos + base0 - 1)?;
7721 if pen_on {
7722 let w0 = pen_hist.len().saturating_sub(sp.penalty_last_n);
7723 pen_hist_d = Some(e.htod_u32_v(&pen_hist[w0..])?);
7724 }
7725 if sampled {
7726 draft_logits.clear();
7727 draft_stats.clear();
7728 }
7729 // DRAFT-SIDE GRAMMAR MASK: clone the committed grammar state ONCE per round; each
7730 // position's mask is computed on that clone and advanced by the PROPOSED token. The
7731 // real state moves only on emission (verify's job), so the emitted stream is
7732 // unchanged — the mask only removes tokens the verify would have truncated anyway.
7733 let mut dmask_live = dmask_on;
7734 if dmask_live {
7735 let t_c = std::time::Instant::now();
7736 constraint
7737 .as_deref_mut()
7738 .unwrap()
7739 .draft_begin()
7740 .map_err(|e2| format!("constraint: {e2}"))?;
7741 dm_clone_ns += t_c.elapsed().as_nanos();
7742 dm_rounds += 1;
7743 }
7744 if let (false, Some(gr)) = (sampled || pen_on, &dctx.graph) {
7745 // GRAPH DRAFT: one dispatch per drafted token. The chain feeds itself on-device
7746 // (in-graph argmax -> tok_d -> next replay's embed; h_nextn -> h_seed_d; pos_d
7747 // inc'd in-graph); the host only reads 4B token (+4B p) and decides the break.
7748 e.set_i32_one(&mut dctx.g_pos, (pos + base0) as i32)?;
7749 e.set_u32_one(&mut dctx.g_tok, last_token)?;
7750 e.copy_into(&mut dctx.g_seed, 0, &h_seed_buf, n_embd)?;
7751 for j in 0..k_this {
7752 // per-position mask upload (contents only — the graph's baked pointer is
7753 // dctx.g_dmask). All-ones once masking goes dead mid-chain, so the captured
7754 // mask node degrades to a no-op ban instead of needing a second graph.
7755 if dmask_live
7756 && !upload_draft_mask(
7757 e,
7758 constraint.as_deref_mut().unwrap(),
7759 &mut dctx.g_dmask,
7760 mtp.d2t.as_ref(),
7761 d_vocab,
7762 dmask_words,
7763 )?
7764 {
7765 // no draft-vocab row is grammar-legal here (a trimmed FR-Spec head can
7766 // genuinely miss the legal set): neutralize the captured mask node and
7767 // finish the chain UNMASKED — exactly pre-lane behaviour, never worse.
7768 e.htod_u32_into(&mut dctx.g_dmask, &vec![u32::MAX; dmask_words])?;
7769 dmask_live = false;
7770 }
7771 gr.launch()?;
7772 scratch.kv.len += 1; // host mirror (len_d advanced in-graph)
7773 let idx = e.dtoh_u32_one(&dctx.g_tok)?;
7774 // #87 SENTINEL TRAP: an all-NaN head-logits row leaves the device argmax's
7775 // init sentinel (0x7FFFFFFF) in g_tok — feeding it onward dereferences
7776 // embed_row(sentinel) = table + ~4.6TB (never mapped) inside the NEXT graph
7777 // replay's embed node, and the MMU fault kills the CUDA context for the
7778 // whole process (research/pp2spec-crash-20260807: 3 coredumps, byte-exact
7779 // VA arithmetic). Refuse loudly instead; the diagnostics name the first-NaN
7780 // buffer (g_seed = the verify-side handoff vs head-side compute).
7781 if (idx as usize) >= d_vocab {
7782 // g_seed is SELF-FED (the replay writes h_nextn back into it), so it
7783 // reads as the head's OUTPUT at j; h_seed_buf is the round's INPUT
7784 // seed, untouched since the round-start copy — the pair discriminates
7785 // "seed arrived poisoned" from "head forward produced NaN".
7786 let seed_h = e.dtoh(&dctx.g_seed)?;
7787 let seed_nan = seed_h.iter().filter(|v| v.is_nan()).count();
7788 let in_h = e.dtoh(&h_seed_buf)?;
7789 let in_nan = in_h.iter().filter(|v| v.is_nan()).count();
7790 return Err(format!(
7791 "draft(graph) argmax sentinel 0x{idx:08x} >= d_vocab {d_vocab} at \
7792 round {round} j={j} pos={pos}: head-out NaN {seed_nan}/{n_embd}, \
7793 round-input-seed NaN {in_nan}/{n_embd} — refusing to dereference \
7794 the embed row (#87 trap)"
7795 )
7796 .into());
7797 }
7798 // trimmed draft vocab -> target token id (identity when no d2t map)
7799 let d = match &mtp.d2t {
7800 Some(map) => map[idx as usize],
7801 None => idx,
7802 };
7803 let draft_p = if p_min > 0.0
7804 || opti_fork
7805 .as_ref()
7806 .is_some_and(|fork| fork.controller.is_some())
7807 {
7808 Some(e.dtoh(&dctx.g_p)?[0])
7809 } else {
7810 None
7811 };
7812 if j == 0 {
7813 controller_draft_prob = draft_p;
7814 }
7815 if let Some(p) = draft_p.filter(|_| p_min > 0.0) {
7816 if p < p_min && (j > 0 || (pmin0 && base0 == 1)) {
7817 break;
7818 }
7819 }
7820 draft.push(d);
7821 // with a trimmed head the NEXT embed must read the TARGET id, not the draft
7822 // index the argmax wrote — patch the persistent token buffer (4B htod).
7823 if d != idx {
7824 e.set_u32_one(&mut dctx.g_tok, d)?;
7825 }
7826 // advance the SPECULATIVE state with the proposal; a dead chain drops to
7827 // unmasked drafting for the remaining positions (verify still arbitrates).
7828 // speculative advance; a chain the grammar can no longer follow (EOS
7829 // proposed) ends here. The captured mask node always runs, so a dead chain
7830 // leaves the buffer NEUTRAL (all-ones = ban nothing) before it exits.
7831 if dmask_live
7832 && !constraint
7833 .as_deref_mut()
7834 .unwrap()
7835 .draft_advance(d)
7836 .map_err(|e2| format!("constraint: {e2}"))?
7837 {
7838 e.htod_u32_into(&mut dctx.g_dmask, &vec![u32::MAX; dmask_words])?;
7839 break;
7840 }
7841 }
7842 } else if let (true, Some(gr)) = (sampled, &dctx.graph_s) {
7843 // SAMPLED GRAPH DRAFT: one replay per drafted token — head forward + gumbel +
7844 // argmax in ONE dispatch; the host reads 4B token (+4B p), D2Ds q into slot j,
7845 // and decides the break. Event-counter continuity: g_ctr is host-seeded to
7846 // sctr-1 ONCE per round (outside the graph); the in-graph bump runs BEFORE the
7847 // perturb, so replay j consumes counter sctr+j — exactly the eager arm's Philox
7848 // stream. Host sctr advances in lockstep (computed, no readback needed).
7849 e.set_i32_one(&mut dctx.g_pos, (pos + base0) as i32)?;
7850 e.set_u32_one(&mut dctx.g_tok, last_token)?;
7851 e.copy_into(&mut dctx.g_seed, 0, &h_seed_buf, n_embd)?;
7852 e.set_u32_one(&mut dctx.g_ctr, sctr.wrapping_sub(1))?;
7853 for j in 0..k_this {
7854 gr.launch()?;
7855 scratch.kv.len += 1; // host mirror (len_d advanced in-graph)
7856 sctr += 1; // mirrors the in-graph g_ctr bump (eager parity:
7857 // counts the p-min-discarded token too)
7858 // q retention: ONE async D2D of the persistent head-logits buffer into this
7859 // round's slot j (stream-ordered after the replay, before the next one).
7860 e.copy_into(&mut dctx.q_slots[j], 0, &dctx.g_q, d_vocab)?;
7861 let idx = e.dtoh_u32_one(&dctx.g_tok)?;
7862 // #87 SENTINEL TRAP (see the greedy graph arm above).
7863 if (idx as usize) >= d_vocab {
7864 let seed_h = e.dtoh(&dctx.g_seed)?;
7865 let seed_nan = seed_h.iter().filter(|v| v.is_nan()).count();
7866 return Err(format!(
7867 "draft(graph-sampled) argmax sentinel 0x{idx:08x} >= d_vocab \
7868 {d_vocab} at round {round} j={j} pos={pos}: round-seed NaN \
7869 {seed_nan}/{n_embd} — refusing to dereference the embed row \
7870 (#87 trap)"
7871 )
7872 .into());
7873 }
7874 let d = match &mtp.d2t {
7875 Some(map) => map[idx as usize],
7876 None => idx,
7877 };
7878 draft_idx.push(idx);
7879 if p_min > 0.0 {
7880 let p = e.dtoh(&dctx.g_p)?[0];
7881 if p < p_min && (j > 0 || (pmin0 && base0 == 1)) {
7882 break;
7883 }
7884 }
7885 draft.push(d);
7886 // trimmed head: the NEXT embed must read the TARGET id (see the greedy arm).
7887 if d != idx {
7888 e.set_u32_one(&mut dctx.g_tok, d)?;
7889 }
7890 }
7891 // uniform accept path: fill draft_stats per used slot (pure-temp regime — the
7892 // stats degenerate to th=0 / full-Z; one filter_stats launch per slot, tiny).
7893 for j in 0..draft.len().max(draft_idx.len()) {
7894 let rows0 = e.htod_i32(&[0])?;
7895 let (mut th_d, mut z_d, mut mx_d) = (e.zeros(1)?, e.zeros(1)?, e.zeros(1)?);
7896 e.filter_stats(
7897 &dctx.q_slots[j],
7898 d_vocab,
7899 &rows0,
7900 &mut th_d,
7901 &mut z_d,
7902 &mut mx_d,
7903 d_vocab,
7904 1,
7905 sp_temp,
7906 sp.top_k,
7907 sp.top_p,
7908 sp.min_p,
7909 )?;
7910 draft_stats.push((e.dtoh(&mx_d)?[0], e.dtoh(&th_d)?[0], e.dtoh(&z_d)?[0]));
7911 }
7912 } else {
7913 // EAGER DRAFT (fallback: MoE head/trunk, huge k, MEMRA_SPEC_NOGRAPH, capture fail).
7914 let mut e_tok = last_token;
7915 let mut d_seed = e.clone_dtod(&h_seed_buf)?;
7916 for j in 0..k_this {
7917 // GPU-ARGMAX DRAFT (2026-07-03): device logits + device argmax + 4-byte token
7918 // read instead of the ~600KB full-vocab dtoh + host argmax per draft token.
7919 let mtp_pos = pos + base0 + j;
7920 // draft-side grammar mask (eager twin of the graph arm's in-graph node).
7921 // A position with no legal draft-vocab row drops to unmasked drafting for
7922 // the rest of the chain (pre-lane behaviour; verify still arbitrates).
7923 if dmask_live {
7924 dmask_live = upload_draft_mask(
7925 e,
7926 constraint.as_deref_mut().unwrap(),
7927 &mut dctx.g_dmask,
7928 mtp.d2t.as_ref(),
7929 d_vocab,
7930 dmask_words,
7931 )?;
7932 }
7933 let (dl_d, h_nextn) = self.mtp_head_forward_dev(
7934 e,
7935 mtp,
7936 e_tok,
7937 &d_seed,
7938 &mut *scratch,
7939 mtp_pos,
7940 embd_dev,
7941 if dmask_live {
7942 Some((&dctx.g_dmask, dmask_words))
7943 } else {
7944 None
7945 },
7946 )?;
7947 let tok_d = if sampled {
7948 // FILTERED Gumbel-max: stats -> masked perturb -> argmax = one draw from
7949 // the filtered softmax (filters off => th=0, exact v1 semantics).
7950 if perturb_buf.is_none() {
7951 perturb_buf = Some(e.zeros(d_vocab.max(n_vocab))?);
7952 }
7953 let mut q_row = e.clone_dtod(&dl_d)?; // retained q (penalized when on)
7954 if pen_on {
7955 let h = pen_hist_d.as_ref().unwrap();
7956 let nh = h.len();
7957 e.penalize_logits(
7958 &mut q_row,
7959 h,
7960 nh,
7961 sp.penalty_repeat,
7962 sp.penalty_freq,
7963 sp.penalty_present,
7964 d_vocab,
7965 )?;
7966 }
7967 let rows0 = e.htod_i32(&[0])?;
7968 let (mut th_d, mut z_d, mut mx_d) =
7969 (e.zeros(1)?, e.zeros(1)?, e.zeros(1)?);
7970 e.filter_stats(
7971 &q_row, d_vocab, &rows0, &mut th_d, &mut z_d, &mut mx_d, d_vocab,
7972 1, sp_temp, sp.top_k, sp.top_p, sp.min_p,
7973 )?;
7974 let (th, z, mx) =
7975 (e.dtoh(&th_d)?[0], e.dtoh(&z_d)?[0], e.dtoh(&mx_d)?[0]);
7976 let pb = perturb_buf.as_mut().unwrap();
7977 e.gumbel_perturb_filtered(
7978 &q_row, pb, d_vocab, sp_seed, sctr, sp_temp, mx, th,
7979 )?;
7980 sctr += 1;
7981 draft_logits.push(q_row);
7982 draft_stats.push((mx, th, z));
7983 e.argmax_token_device(pb, d_vocab)?
7984 } else {
7985 e.argmax_token_device(&dl_d, d_vocab)?
7986 };
7987 let idx = e.dtoh_u32_one(&tok_d)?;
7988 // #87 SENTINEL TRAP (eager twin — see the graph arm). Extra diagnostics
7989 // here because the eager chain's operands are all readable: dl_d (the head
7990 // logits row) and d_seed (this step's h_seed) name the first-NaN buffer.
7991 if (idx as usize) >= d_vocab {
7992 let dl_h = e.dtoh(&dl_d)?;
7993 let dl_nan = dl_h.iter().filter(|v| v.is_nan()).count();
7994 let seed_h = e.dtoh(&d_seed)?;
7995 let seed_nan = seed_h.iter().filter(|v| v.is_nan()).count();
7996 return Err(format!(
7997 "draft(eager) argmax sentinel 0x{idx:08x} >= d_vocab {d_vocab} at \
7998 round {round} j={j} pos={pos}: head-logits NaN {dl_nan}/{d_vocab}, \
7999 step-seed NaN {seed_nan}/{n_embd} — refusing to dereference the \
8000 embed row (#87 trap)"
8001 )
8002 .into());
8003 }
8004 let d = match &mtp.d2t {
8005 Some(map) => map[idx as usize],
8006 None => idx,
8007 };
8008 if sampled {
8009 draft_idx.push(idx);
8010 }
8011 let draft_p = if p_min > 0.0
8012 || opti_fork
8013 .as_ref()
8014 .is_some_and(|fork| fork.controller.is_some())
8015 {
8016 let p_d = e.prob_of_token_device(&dl_d, &tok_d, d_vocab)?;
8017 Some(e.dtoh(&p_d)?[0])
8018 } else {
8019 None
8020 };
8021 if j == 0 {
8022 controller_draft_prob = draft_p;
8023 }
8024 if let Some(p) = draft_p.filter(|_| p_min > 0.0) {
8025 if p < p_min && (j > 0 || (pmin0 && base0 == 1)) {
8026 break;
8027 }
8028 }
8029 draft.push(d);
8030 e_tok = d;
8031 d_seed = h_nextn;
8032 // speculative advance; a chain the grammar can no longer follow (EOS
8033 // proposed) ends here — the prefix already proposed still rides verify.
8034 if dmask_live
8035 && !constraint
8036 .as_deref_mut()
8037 .unwrap()
8038 .draft_advance(d)
8039 .map_err(|e2| format!("constraint: {e2}"))?
8040 {
8041 break;
8042 }
8043 }
8044 if opti_fork
8045 .as_ref()
8046 .is_some_and(|fork| fork.controller.is_some())
8047 {
8048 controller_eager_state = Some((e_tok, d_seed));
8049 }
8050 }
8051 }
8052 let k_round = draft.len();
8053 if let Some(p) = pipe {
8054 p.draft_end(round);
8055 }
8056 drop(pipe_draft);
8057
8058 ph_mark(&mut ph_draft, phase_on);
8059 // --- 2. VERIFY: one batched target forward. With a pending bonus, it rides as col 0
8060 // (committing its KV/recur inside the SAME weight read); drafts follow. ---
8061 let verify_tokens: Vec<u32> = match pending {
8062 Some(b) => {
8063 let mut v = Vec::with_capacity(k_round + 1);
8064 v.push(b);
8065 v.extend_from_slice(&draft);
8066 v
8067 }
8068 None => draft.clone(),
8069 };
8070 let base = if pending.is_some() { 1 } else { 0 };
8071 // ckpt (REPLAY-FREE partial accept): retain per-layer state-rebuild inputs alongside
8072 // the verify. Pure buffer keep-alives + dtod clones — kernel work is unchanged.
8073 let mut ckpt = if let Some(ticket) = current_opti.as_mut() {
8074 Some(ticket.take_ckpt())
8075 } else if spec_replay {
8076 None
8077 } else {
8078 Some(VerifyCkpt::new(self.layers.len()))
8079 };
8080 let controller_can_probe = base == 1
8081 && k_round == 1
8082 && out.len().saturating_add(2) < max_new
8083 && controller_draft_prob.is_some()
8084 && opti_fork
8085 .as_ref()
8086 .and_then(|fork| fork.controller.as_ref())
8087 .is_some_and(|policy| !policy.breaker_tripped);
8088 let mut successor_attempt: Option<OptiControllerTicket> = None;
8089 let mut rejected_probe: Option<(f32, u32)> = None;
8090 let mut controller_prepared: Option<OptiControllerPrepared> = None;
8091 if controller_can_probe {
8092 // Prepare d2/q and, on admission, d3 before either current verify half is
8093 // issued. N stage 0 can then be followed immediately by N+1 stage 0; once N's
8094 // boundary fires, those dev0 launches overlap N stage 1 on dev1. Preparing on
8095 // the primary stream after N stage 1 would serialize the supposed pipeline.
8096 let eager_pos = scratch.kv.len + 1;
8097 let (optimistic_pending, pending_probability) = self.opti_controller_draft_step(
8098 e,
8099 mtp,
8100 &mut dctx,
8101 &mut *scratch,
8102 d_vocab,
8103 &mut controller_eager_state,
8104 eager_pos,
8105 embd_dev,
8106 )?;
8107 let first_probability = controller_draft_prob
8108 .ok_or("optipipe controller probe lost first-token probability")?;
8109 let q_proxy = first_probability * pending_probability;
8110 OPTI_GATE_CHECKS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
8111 OPTI_SHADOW_DRAFT_TOKENS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
8112 let admitted = opti_fork
8113 .as_ref()
8114 .and_then(|fork| fork.controller.as_ref())
8115 .ok_or("optipipe controller policy disappeared")?
8116 .admit(q_proxy);
8117 if admitted {
8118 OPTI_GATE_ADMITS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
8119 let eager_pos = scratch.kv.len + 1;
8120 let (optimistic_draft, optimistic_draft_probability) = self
8121 .opti_controller_draft_step(
8122 e,
8123 mtp,
8124 &mut dctx,
8125 &mut *scratch,
8126 d_vocab,
8127 &mut controller_eager_state,
8128 eager_pos,
8129 embd_dev,
8130 )?;
8131 OPTI_SHADOW_DRAFT_TOKENS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
8132 let eager_seed = controller_eager_state.take().map(|(token, seed)| {
8133 debug_assert_eq!(token, optimistic_draft);
8134 seed
8135 });
8136 controller_prepared = Some(OptiControllerPrepared {
8137 verify_tokens: [optimistic_pending, optimistic_draft],
8138 draft_prob: optimistic_draft_probability,
8139 eager_seed,
8140 q_proxy,
8141 scratch_len: scratch.kv.len,
8142 });
8143 } else {
8144 OPTI_GATE_REJECTS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
8145 OPTI_WASTED_DRAFT_TOKENS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
8146 rejected_probe = Some((q_proxy, optimistic_pending));
8147 eprintln!(
8148 "[opti-controller] reject q={q_proxy:.6} threshold={:.3}",
8149 opti_fork
8150 .as_ref()
8151 .and_then(|fork| fork.controller.as_ref())
8152 .expect("controller policy")
8153 .threshold,
8154 );
8155 }
8156 }
8157 let fork_attempt = match fork_generation.take() {
8158 Some(generation) if base == 1 && k_round == 1 => Some(generation),
8159 Some(generation) => {
8160 opti_fork
8161 .as_mut()
8162 .expect("fork generation without fork state")
8163 .retire(generation)?;
8164 None
8165 }
8166 None => None,
8167 };
8168 let (tlogits_d, vx) = if let Some(p) = pipe {
8169 self.decode_step_t_core_pipelined(
8170 e,
8171 &verify_tokens,
8172 pos,
8173 &mut *cache,
8174 embd_dev,
8175 ckpt.as_mut(),
8176 p,
8177 round,
8178 )?
8179 } else if controller_can_probe {
8180 let fence = opti_fork
8181 .as_ref()
8182 .ok_or("optipipe controller probe lost fork state")?
8183 .fence;
8184 let boundary = match current_opti.as_mut() {
8185 Some(ticket) => ticket.take_boundary(),
8186 None => self.verify_stage0_issue(
8187 e,
8188 &verify_tokens,
8189 pos,
8190 &mut *cache,
8191 embd_dev,
8192 ckpt.as_mut(),
8193 None,
8194 &fence,
8195 Some(true),
8196 None,
8197 )?,
8198 };
8199 if let Some(prepared) = controller_prepared.take() {
8200 let generation = {
8201 let fork = opti_fork
8202 .as_mut()
8203 .ok_or("optipipe controller admission lost fork state")?;
8204 let generation = fork.reserve_successor()?;
8205 let rt = fork.rt;
8206 let snapshot_fence = fork.fence;
8207 opti_snapshot_one_stage_owned_into(
8208 e,
8209 cache,
8210 rt,
8211 &snapshot_fence,
8212 0,
8213 fork.successor_snapshot_mut(),
8214 )?;
8215 generation
8216 };
8217 let mut successor_ckpt = VerifyCkpt::new(self.layers.len());
8218 let successor_boundary = self.verify_stage0_issue(
8219 e,
8220 &prepared.verify_tokens,
8221 pos + verify_tokens.len(),
8222 &mut *cache,
8223 embd_dev,
8224 Some(&mut successor_ckpt),
8225 None,
8226 &fence,
8227 Some(false),
8228 None,
8229 )?;
8230 OPTI_FORK_ATTEMPTS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
8231 let fork = opti_fork
8232 .as_ref()
8233 .ok_or("optipipe controller ticket lost fork state")?;
8234 successor_attempt = Some(fork.controller_ticket(
8235 generation,
8236 successor_boundary,
8237 successor_ckpt,
8238 prepared.verify_tokens,
8239 prepared.draft_prob,
8240 prepared.eager_seed,
8241 prepared.q_proxy,
8242 prepared.scratch_len,
8243 ));
8244 eprintln!(
8245 "[opti-controller] issue generation={} q={:.6} threshold={:.3} \
8246 verify={:?}",
8247 generation.id,
8248 prepared.q_proxy,
8249 fork.controller.expect("controller policy").threshold,
8250 prepared.verify_tokens,
8251 );
8252 }
8253 let result = self.verify_stage1_finish(
8254 e,
8255 boundary,
8256 &mut *cache,
8257 ckpt.as_mut(),
8258 None,
8259 &fence,
8260 successor_attempt.is_none(),
8261 )?;
8262 if let Some(ticket) = current_opti.as_mut() {
8263 ticket.settle();
8264 }
8265 if successor_attempt.is_some() {
8266 let fork = opti_fork
8267 .as_mut()
8268 .ok_or("optipipe successor snapshot lost fork state")?;
8269 let rt = fork.rt;
8270 let snapshot_fence = fork.fence;
8271 opti_snapshot_one_stage_owned_into(
8272 e,
8273 cache,
8274 rt,
8275 &snapshot_fence,
8276 1,
8277 fork.successor_snapshot_mut(),
8278 )?;
8279 // Publish N only after both independent successor-state queues are complete.
8280 fork.rt.publish_to(1, &e.stream())?;
8281 }
8282 result
8283 } else if let Some(ticket) = current_opti.as_mut() {
8284 let fork = opti_fork
8285 .as_mut()
8286 .ok_or("optipipe carried controller ticket lost fork state")?;
8287 let boundary = ticket.take_boundary();
8288 let result = self.verify_stage1_finish(
8289 e,
8290 boundary,
8291 &mut *cache,
8292 ckpt.as_mut(),
8293 None,
8294 &fork.fence,
8295 true,
8296 )?;
8297 ticket.settle();
8298 result
8299 } else if let Some(generation) = fork_attempt {
8300 let fork = opti_fork
8301 .as_mut()
8302 .expect("fork generation without fork state");
8303 fork.capture_seed(e, generation, &h_seed_buf, &fill_prev, scratch.kv.len)?;
8304 let action = fork.mode.action(generation.id);
8305 let boundary = self.verify_stage0_issue(
8306 e,
8307 &verify_tokens,
8308 pos,
8309 &mut *cache,
8310 embd_dev,
8311 ckpt.as_mut(),
8312 None,
8313 &fork.fence,
8314 Some(true),
8315 None,
8316 )?;
8317 OPTI_FORK_ATTEMPTS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
8318 let mut ticket = fork.ticket(generation, boundary);
8319 if action == OptiForkAction::Abort {
8320 return Err(format!(
8321 "optipipe forced abort with generation {} stage0 in flight",
8322 generation.id,
8323 )
8324 .into());
8325 }
8326 fork.reconcile(
8327 e,
8328 &mut *cache,
8329 &mut *scratch,
8330 &snap,
8331 &mut h_seed_buf,
8332 &mut fill_prev,
8333 generation,
8334 action,
8335 verify_tokens[0],
8336 )?;
8337 let result = if action == OptiForkAction::Hit {
8338 let boundary = ticket.take_boundary();
8339 self.verify_stage1_finish(
8340 e,
8341 boundary,
8342 &mut *cache,
8343 ckpt.as_mut(),
8344 None,
8345 &fork.fence,
8346 true,
8347 )?
8348 } else {
8349 // The optimistic boundary slot has no reader. Re-run the unchanged serial
8350 // verify only after E_restart published the restored stage-0 state.
8351 self.decode_step_t_core(
8352 e,
8353 &verify_tokens,
8354 pos,
8355 &mut *cache,
8356 embd_dev,
8357 ckpt.as_mut(),
8358 )?
8359 };
8360 ticket.settle();
8361 debug_assert_eq!(ticket.generation, generation);
8362 fork.retire(generation)?;
8363 result
8364 } else {
8365 self.decode_step_t_core(
8366 e,
8367 &verify_tokens,
8368 pos,
8369 &mut *cache,
8370 embd_dev,
8371 ckpt.as_mut(),
8372 )?
8373 };
8374 let pipe_accept = match pipe {
8375 Some(p) => Some(p.accept_begin(round)?),
8376 None => None,
8377 };
8378
8379 ph_mark(&mut ph_verify, phase_on);
8380 // --- 3. GREEDY ACCEPT (walk prefix, stop at first mismatch) ---
8381 // DEVICE-ARGMAX ACCEPT: argmax every verify column ON DEVICE (same 2-pass kernels +
8382 // smallest-index tie-break as host argmax, argmax_gate-validated) and read back ONE
8383 // [T] u32 — replaces the T x n_vocab f32 dtoh + T host argmaxes per round.
8384 // t_pred[j] = target's greedy prediction for the slot after draft[j-1] (j>=1) or after
8385 // last_token (j==0). With a pending bonus, col 0 IS the prediction after last_token
8386 // (== the bonus), so every index shifts by `base` and last_pred is unused.
8387 let t_v = verify_tokens.len();
8388 let mut preds: Vec<u32> = Vec::new();
8389 if !sampled {
8390 for j in 0..t_v {
8391 e.argmax_token_device_col(&tlogits_d, j, n_vocab, &mut preds_d, j)?;
8392 }
8393 preds = e.dtoh_u32(&preds_d)?; // <- the verify-GPU wait lands here
8394 // #87 SENTINEL TRAP, verify side: a sentinel pred becomes the round's bonus =
8395 // next round's last_token = the next chain's embed lookup. Catch it at the
8396 // source with the column named — an all-NaN VERIFY column implicates the
8397 // stage-split trunk (decode_step_t_core_ppn), not the draft head.
8398 if let Some(bad) = preds[..t_v].iter().position(|&p| (p as usize) >= n_vocab) {
8399 let col = &tlogits_d.slice(bad * n_vocab..(bad + 1) * n_vocab);
8400 let mut probe = e.zeros(n_vocab)?;
8401 e.copy_view_into(&mut probe, 0, col, n_vocab)?;
8402 let col_h = e.dtoh(&probe)?;
8403 let col_nan = col_h.iter().filter(|v| v.is_nan()).count();
8404 return Err(format!(
8405 "verify argmax sentinel 0x{:08x} >= n_vocab {n_vocab} at round {round} \
8406 col {bad}/{t_v} pos={pos}: verify-logits col NaN {col_nan}/{n_vocab} \
8407 — the stage-split verify produced a poisoned column (#87 trap)",
8408 preds[bad]
8409 )
8410 .into());
8411 }
8412 }
8413 ph_mark(&mut ph_wait, phase_on);
8414 let t_pred = |j: usize| -> u32 {
8415 if j == 0 && base == 0 {
8416 last_pred
8417 } else {
8418 preds[base + j - 1]
8419 }
8420 };
8421 let mut devacc_seeded = false;
8422 let mut devacc_acc: Option<CudaSlice<u32>> = None;
8423 let (n_acc, bonus) = if !sampled {
8424 // ROUND-STREAM stage (a) (MEMRA_SPEC_DEVACC=1 opt-in): the walk runs ON DEVICE
8425 // (spec_accept_greedy, verbatim rule) and the host reads back 8B (n_acc, bonus)
8426 // instead of the [T] preds. Same sync count — machinery for stages (b)/(c),
8427 // gated on token identity vs the host walk (the arms below are bit-equal rules).
8428 if crate::spec::spec_devacc() && k_round > 0 && !spec_replay && constraint.is_none()
8429 {
8430 let draft_d = e.htod_u32_v(&draft)?;
8431 let mut acc_out = e.alloc_u32_zeroed(2)?;
8432 e.spec_accept_greedy(
8433 &preds_d,
8434 &draft_d,
8435 last_pred,
8436 base,
8437 k_round,
8438 &mut acc_out,
8439 )?;
8440 devacc_acc = Some(acc_out.clone());
8441 // stage (b): next-round seed gathered ON DEVICE from acc_out before the host
8442 // ever reads n_acc (j=base+n_acc -> vx col j-1; j==0 -> fill_prev). The three
8443 // non-replay commit arms skip their host-offset seed copies (guarded below);
8444 // the legacy spec_replay arm keeps its own rx-based seeding (excluded here).
8445 // NOTE: fill_prev is NOT updated here — the commit arms' TRUE-HIDDEN
8446 // REFRESH reads the OLD fill_prev (predecessor of this round's verify batch);
8447 // the update lands after the arms (devacc_seeded guard below).
8448 e.spec_seed_gather(&vx, &fill_prev, &acc_out, &mut h_seed_buf, base, n_embd)?;
8449 // 3a: KV lens roll back on device (len = saved + base + n_acc, all arms'
8450 // unified rule; full accept rewrites the verify-left value). Host mirrors
8451 // update after the readback; commit_verified_prefix skips its len_d writes.
8452 if let Some(successor) = successor_attempt.as_ref() {
8453 opti_fork
8454 .as_mut()
8455 .ok_or("optipipe successor reconcile lost fork state")?
8456 .queue_actual_reconcile(
8457 e,
8458 &snap,
8459 &acc_out,
8460 successor.verify_tokens[0],
8461 base,
8462 )?;
8463 } else if let Some(ptrs) = &kv_len_ptrs {
8464 let saved: Vec<i32> = (0..self.layers.len())
8465 .map(|il| snap.kv_len[il].map(|v| v as i32).unwrap_or(0))
8466 .collect();
8467 let saved_d = e.htod_i32(&saved)?;
8468 e.spec_rollback_kv(ptrs, &saved_d, &acc_out, base, self.layers.len())?;
8469 }
8470 devacc_seeded = true;
8471 let ab = e.dtoh_u32(&acc_out)?;
8472 (ab[0] as usize, ab[1])
8473 } else {
8474 let mut n_acc = 0usize;
8475 for j in 0..k_round {
8476 if t_pred(j) == draft[j] {
8477 n_acc += 1;
8478 } else {
8479 break;
8480 }
8481 }
8482 // bonus = target's own token at the first non-accepted slot. n_acc in 0..=k; t_pred
8483 // is defined for j in 0..=k (j==0 -> last_logits, j>=1 -> col j-1, last col = k-1).
8484 (n_acc, t_pred(n_acc))
8485 }
8486 } else {
8487 // --- SAMPLED ACCEPT (rejection sampling): u_j < p_j(x_j)/q_j(x_j) walk ---
8488 if col_buf.is_none() {
8489 col_buf = Some(e.zeros(n_vocab)?);
8490 }
8491 // FILTERED p_j: per-verify-col stats (one batched filter_stats call), then the
8492 // filtered gather. j==0&&base==0 reads last_col (its own stats row appended).
8493 let mut pj = vec![0f32; k_round.max(1)];
8494 let mut col_stats: Vec<(f32, f32, f32)> = Vec::new(); // (max, th, z) per verify col used
8495 if k_round > 0 {
8496 let mut ids: Vec<u32> = Vec::new();
8497 let mut rows: Vec<i32> = Vec::new();
8498 for j in 0..k_round {
8499 if j > 0 || base == 1 {
8500 ids.push(draft[j]);
8501 rows.push((base + j) as i32 - 1);
8502 }
8503 }
8504 if !ids.is_empty() {
8505 let nr = rows.len();
8506 // penalties: materialize the used columns into one contiguous penalized
8507 // buffer (rows remapped 0..nr) so stats+gathers see the penalized p.
8508 // penalties: materialize used columns contiguously, penalize all rows in
8509 // one launch, and point stats+gathers at the penalized buffer (rows 0..nr).
8510 let p_rows: Vec<i32> = if pen_on {
8511 (0..nr as i32).collect()
8512 } else {
8513 rows.clone()
8514 };
8515 if pen_on {
8516 if pcol_buf.as_ref().map(|b| b.len()).unwrap_or(0) < nr * n_vocab {
8517 pcol_buf = Some(e.zeros(nr * n_vocab)?);
8518 }
8519 let pc = pcol_buf.as_mut().unwrap();
8520 for (i2, &r) in rows.iter().enumerate() {
8521 let c = r as usize;
8522 e.copy_view_into(
8523 pc,
8524 i2 * n_vocab,
8525 &tlogits_d.slice(c * n_vocab..(c + 1) * n_vocab),
8526 n_vocab,
8527 )?;
8528 }
8529 let h = pen_hist_d.as_ref().unwrap();
8530 let nh = h.len();
8531 e.penalize_logits_rows(
8532 pc,
8533 h,
8534 nh,
8535 sp.penalty_repeat,
8536 sp.penalty_freq,
8537 sp.penalty_present,
8538 n_vocab,
8539 nr,
8540 )?;
8541 }
8542 let p_src: &CudaSlice<f32> = if pen_on {
8543 pcol_buf.as_ref().unwrap()
8544 } else {
8545 &tlogits_d
8546 };
8547 let rowsd = e.htod_i32(&p_rows)?;
8548 let (mut th_d, mut z_d, mut mx_d) =
8549 (e.zeros(nr)?, e.zeros(nr)?, e.zeros(nr)?);
8550 e.filter_stats(
8551 p_src, n_vocab, &rowsd, &mut th_d, &mut z_d, &mut mx_d, n_vocab, nr,
8552 sp_temp, sp.top_k, sp.top_p, sp.min_p,
8553 )?;
8554 let idsd = e.htod_u32_v(&ids)?;
8555 let mut outd = e.zeros(nr)?;
8556 e.softmax_gather_filtered(
8557 p_src, n_vocab, &idsd, &rowsd, &th_d, &z_d, &mut outd, n_vocab, nr,
8558 sp_temp,
8559 )?;
8560 let outv = e.dtoh(&outd)?;
8561 let (thv, zv, mxv) = (e.dtoh(&th_d)?, e.dtoh(&z_d)?, e.dtoh(&mx_d)?);
8562 let mut oi = 0usize;
8563 for j in 0..k_round {
8564 if j > 0 || base == 1 {
8565 pj[j] = outv[oi];
8566 oi += 1;
8567 }
8568 }
8569 col_stats = (0..nr).map(|i| (mxv[i], thv[i], zv[i])).collect();
8570 }
8571 if base == 0 {
8572 let lc: &CudaSlice<f32> = if pen_on {
8573 if col_buf.is_none() {
8574 col_buf = Some(e.zeros(n_vocab)?);
8575 }
8576 let cb = col_buf.as_mut().unwrap();
8577 e.copy_into(
8578 cb,
8579 0,
8580 last_col_logits
8581 .as_ref()
8582 .expect("sampled: last_col_logits unset"),
8583 n_vocab,
8584 )?;
8585 let h = pen_hist_d.as_ref().unwrap();
8586 let nh = h.len();
8587 e.penalize_logits(
8588 cb,
8589 h,
8590 nh,
8591 sp.penalty_repeat,
8592 sp.penalty_freq,
8593 sp.penalty_present,
8594 n_vocab,
8595 )?;
8596 col_buf.as_ref().unwrap()
8597 } else {
8598 last_col_logits
8599 .as_ref()
8600 .expect("sampled: last_col_logits unset")
8601 };
8602 let rows0 = e.htod_i32(&[0])?;
8603 let (mut th_d, mut z_d, mut mx_d) = (e.zeros(1)?, e.zeros(1)?, e.zeros(1)?);
8604 e.filter_stats(
8605 lc, n_vocab, &rows0, &mut th_d, &mut z_d, &mut mx_d, n_vocab, 1,
8606 sp_temp, sp.top_k, sp.top_p, sp.min_p,
8607 )?;
8608 let idsd = e.htod_u32_v(&[draft[0]])?;
8609 let mut outd = e.zeros(1)?;
8610 e.softmax_gather_filtered(
8611 lc, n_vocab, &idsd, &rows0, &th_d, &z_d, &mut outd, n_vocab, 1, sp_temp,
8612 )?;
8613 pj[0] = e.dtoh(&outd)?[0];
8614 last_col_stats =
8615 Some((e.dtoh(&mx_d)?[0], e.dtoh(&th_d)?[0], e.dtoh(&z_d)?[0]));
8616 }
8617 }
8618 // q source: the graph arm retained the head logits in the persistent q_slots;
8619 // the eager arm in per-round draft_logits clones. Same raw-logit values either way.
8620 // FILTERED q_j: stats from draft_stats (eager pushes in-chain; the graph arm
8621 // computes them post-replay — graph engages only filter/penalty-free, so the
8622 // stats degenerate to th=0/full-Z there, keeping ONE accept path).
8623 let q_bufs: &[CudaSlice<f32>] = if dctx.graph_s.is_some() {
8624 &dctx.q_slots
8625 } else {
8626 &draft_logits
8627 };
8628 let mut n_acc = 0usize;
8629 for j in 0..k_round {
8630 let (qmx, qth, qz) = draft_stats[j];
8631 let idsd = e.htod_u32_v(&[draft_idx[j]])?;
8632 let rowsd = e.htod_i32(&[0])?;
8633 let thd = e.htod(&[qth])?;
8634 let zd = e.htod(&[qz])?;
8635 let _ = qmx;
8636 let mut outd = e.zeros(1)?;
8637 e.softmax_gather_filtered(
8638 &q_bufs[j], d_vocab, &idsd, &rowsd, &thd, &zd, &mut outd, d_vocab, 1,
8639 sp_temp,
8640 )?;
8641 let qj = e.dtoh(&outd)?[0];
8642 let u = host_u01(sp_seed, uctr);
8643 uctr += 1;
8644 if (u as f64) * (qj as f64) < pj[j] as f64 {
8645 n_acc += 1;
8646 } else {
8647 break;
8648 }
8649 }
8650 let bonus = if n_acc == k_round {
8651 // FULL ACCEPT: bonus ~ FILTERED softmax at the last verify column.
8652 let col = base + k_round - 1;
8653 let cb = col_buf.as_mut().unwrap();
8654 e.copy_view_into(
8655 cb,
8656 0,
8657 &tlogits_d.slice(col * n_vocab..(col + 1) * n_vocab),
8658 n_vocab,
8659 )?;
8660 if pen_on {
8661 let h = pen_hist_d.as_ref().unwrap();
8662 let nh = h.len();
8663 e.penalize_logits(
8664 cb,
8665 h,
8666 nh,
8667 sp.penalty_repeat,
8668 sp.penalty_freq,
8669 sp.penalty_present,
8670 n_vocab,
8671 )?;
8672 }
8673 if perturb_buf.is_none() {
8674 perturb_buf = Some(e.zeros(d_vocab.max(n_vocab))?);
8675 }
8676 // STATS MUST COME FROM THIS COLUMN (bug fix 2026-08-05, lane/sampler-
8677 // truncation-fix; receipts research/sampfix-20260805/). The old code reused
8678 // `col_stats.last()` here, which is ALWAYS the wrong row: the gathered set
8679 // covers verify columns 0..=(base+k_round-2) (rows pushed as base+j-1), while
8680 // the full-accept bonus samples column base+k_round-1 — exactly ONE PAST the
8681 // last gathered column, in both base arms. `th` is a threshold in e-units of
8682 // its OWN row's max, so feeding a neighbour's (row_max, th) into
8683 // gumbel_perturb_filtered mis-scales every e0 = exp((x-row_max)/T). When the
8684 // donor column's peak is higher by more than T*ln(1/th), EVERY id fails
8685 // `e0 >= th`, the whole perturbed row becomes -3.4e38, and the 2-pass argmax
8686 // falls through to its smallest-index tie-break => token id 0 ("!") spliced
8687 // mid-word. Fragility is ordered by how large th is: min_p pins th = min_p
8688 // (0.05 => trigger at delta > 2.4 at T=0.8, fires constantly), top_p's
8689 // mass-boundary th is smaller, top_k's k-th-largest th smaller still — which
8690 // is why the head-to-head matrix saw min_p and top_p corrupt while top_k-only
8691 // stayed clean. The pure-temp default regime is immune (th == 0 masks nothing,
8692 // and row_max is unused once nothing is masked), so this fix is a byte-level
8693 // no-op for the untruncated serve default. One extra one-block filter_stats
8694 // per full-accept round is the whole cost.
8695 let (mx, th) = {
8696 let rows0 = e.htod_i32(&[0])?;
8697 let (mut th_d, mut z_d, mut mx_d) = (e.zeros(1)?, e.zeros(1)?, e.zeros(1)?);
8698 let cb0 = col_buf.as_ref().unwrap();
8699 e.filter_stats(
8700 cb0, n_vocab, &rows0, &mut th_d, &mut z_d, &mut mx_d, n_vocab, 1,
8701 sp_temp, sp.top_k, sp.top_p, sp.min_p,
8702 )?;
8703 (e.dtoh(&mx_d)?[0], e.dtoh(&th_d)?[0])
8704 };
8705 let pb = perturb_buf.as_mut().unwrap();
8706 let cb2 = col_buf.as_ref().unwrap();
8707 e.gumbel_perturb_filtered(cb2, pb, n_vocab, sp_seed, sctr, sp_temp, mx, th)?;
8708 sctr += 1;
8709 let td = e.argmax_token_device(pb, n_vocab)?;
8710 e.dtoh_u32_one(&td)?
8711 } else {
8712 // REJECT at n_acc: bonus ~ norm(max(0, softmax_T(p) - softmax_T(q))).
8713 let cb = col_buf.as_mut().unwrap();
8714 if n_acc > 0 || base == 1 {
8715 let col = base + n_acc - 1;
8716 e.copy_view_into(
8717 cb,
8718 0,
8719 &tlogits_d.slice(col * n_vocab..(col + 1) * n_vocab),
8720 n_vocab,
8721 )?;
8722 } else {
8723 let lc = last_col_logits.as_ref().unwrap();
8724 e.copy_into(cb, 0, lc, n_vocab)?;
8725 }
8726 if pen_on {
8727 let h = pen_hist_d.as_ref().unwrap();
8728 let nh = h.len();
8729 e.penalize_logits(
8730 cb,
8731 h,
8732 nh,
8733 sp.penalty_repeat,
8734 sp.penalty_freq,
8735 sp.penalty_present,
8736 n_vocab,
8737 )?;
8738 }
8739 let cb2 = col_buf.as_ref().unwrap();
8740 let sc = sctr;
8741 sctr += 1;
8742 // p-stats for the reject column: from col_stats when the col was gathered,
8743 // else (j==0&&base==0) from last_col_stats.
8744 let p_stats = if n_acc > 0 || base == 1 {
8745 // col index within the gathered set == number of gathered cols before n_acc
8746 let gi = if base == 1 { n_acc } else { n_acc - 1 };
8747 col_stats.get(gi).copied().unwrap_or_else(|| {
8748 (0.0, 0.0, 1.0) // unreachable: gathered cols always cover the reject slot
8749 })
8750 } else {
8751 last_col_stats.expect("sampled: last_col_stats unset at reject")
8752 };
8753 let q_stats = draft_stats[n_acc];
8754 if let Some(map) = &d2t_dev {
8755 if q_full_buf.is_none() {
8756 q_full_buf = Some(e.zeros(n_vocab)?);
8757 }
8758 let qf = q_full_buf.as_mut().unwrap();
8759 e.scatter_trim_logits(&q_bufs[n_acc], map, qf, d_vocab, n_vocab)?;
8760 let qf2 = q_full_buf.as_ref().unwrap();
8761 e.residual_sample_filtered(
8762 cb2,
8763 Some(qf2),
8764 n_vocab,
8765 sp_temp,
8766 sp_seed,
8767 sc,
8768 p_stats,
8769 q_stats,
8770 &mut sample_tok,
8771 )?;
8772 } else {
8773 e.residual_sample_filtered(
8774 cb2,
8775 Some(&q_bufs[n_acc]),
8776 n_vocab,
8777 sp_temp,
8778 sp_seed,
8779 sc,
8780 p_stats,
8781 q_stats,
8782 &mut sample_tok,
8783 )?;
8784 }
8785 e.dtoh_u32(&sample_tok)?[0]
8786 };
8787 (n_acc, bonus)
8788 };
8789 // --- 3b. GRAMMAR TRUNCATION (constrained spec, 2026-08-03): the grammar is
8790 // an extra rejection rule AFTER the exactness verify (the batched-verify-twins
8791 // ordering). Walk the accepted drafts through the grammar in commit order; the
8792 // first illegal token truncates acceptance at its slot, and that slot's emission
8793 // is recomputed as the MASKED argmax of the target's own verify column — token-
8794 // identical to constrained plain greedy decode (an unmasked argmax that is
8795 // grammar-legal IS the masked argmax: masking only removes competitors). The
8796 // column D2H (~1MB) is paid only when a cut fires — the tight-grammar cost,
8797 // measured in acceptance numbers, never hidden.
8798 let (n_acc, bonus) = match constraint.as_deref_mut() {
8799 None => (n_acc, bonus),
8800 Some(c) => {
8801 fn ce(e2: String) -> Box<dyn std::error::Error> {
8802 format!("constraint: {e2}").into()
8803 }
8804 let mut na = n_acc;
8805 let mut cut = false;
8806 for (j, &d) in draft.iter().enumerate().take(n_acc) {
8807 if c.is_allowed(d).map_err(ce)? {
8808 c.consume(d).map_err(ce)?;
8809 } else {
8810 na = j;
8811 cut = true;
8812 dm_cut_tokens += n_acc - j;
8813 break;
8814 }
8815 }
8816 if cut {
8817 dm_cuts += 1;
8818 }
8819 let mut bo = bonus;
8820 if cut || !c.is_allowed(bo).map_err(ce)? {
8821 let mut row = if na == 0 && base == 0 {
8822 init_logits_host
8823 .clone()
8824 .ok_or("constraint: init logits missing (round-0 cut)")?
8825 } else {
8826 e.dtoh_view(
8827 &tlogits_d.slice((base + na - 1) * n_vocab..(base + na) * n_vocab),
8828 )?
8829 };
8830 c.mask_logits(&mut row).map_err(ce)?;
8831 bo = argmax(&row) as u32;
8832 }
8833 c.consume(bo).map_err(ce)?;
8834 (na, bo)
8835 }
8836 };
8837 let mut successor_valid = false;
8838 if let Some((q_proxy, expected_d2)) = rejected_probe {
8839 let v_n = n_acc == 1 && bonus == expected_d2;
8840 eprintln!(
8841 "[opti-controller] shadow q={q_proxy:.6} admitted=false v_n={v_n} \
8842 expected_d2={expected_d2} n_acc={n_acc} bonus={bonus}",
8843 );
8844 }
8845 if let Some(successor) = successor_attempt.as_ref() {
8846 successor_valid = n_acc == 1 && bonus == successor.verify_tokens[0];
8847 let generation = successor.generation;
8848 let q_proxy = successor.q_proxy;
8849 let expected_pending = successor.verify_tokens[0];
8850 let resolution_ms = successor.issued_at.elapsed().as_secs_f64() * 1e3;
8851 let fork = opti_fork
8852 .as_mut()
8853 .ok_or("optipipe successor resolution lost fork state")?;
8854 fork.finish_actual_reconcile(e, &mut *cache, &snap, n_acc, base, successor_valid)?;
8855 if successor_valid {
8856 OPTI_FORK_HITS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
8857 } else {
8858 OPTI_FORK_MISSES.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
8859 OPTI_RECONCILES.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
8860 OPTI_WASTED_DRAFT_TOKENS.fetch_add(2, std::sync::atomic::Ordering::Relaxed);
8861 }
8862 let breaker_tripped = fork
8863 .controller
8864 .as_mut()
8865 .expect("controller policy")
8866 .resolve(successor_valid);
8867 if breaker_tripped {
8868 OPTI_BREAKER_TRIPS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
8869 }
8870 eprintln!(
8871 "[opti-controller] resolve generation={} hit={} q={q_proxy:.6} \
8872 expected_pending={expected_pending} n_acc={n_acc} bonus={bonus} \
8873 resolution_ms={resolution_ms:.3} reconcile={} breaker={}",
8874 generation.id, successor_valid, !successor_valid, breaker_tripped,
8875 );
8876 if !successor_valid {
8877 let mut successor = successor_attempt
8878 .take()
8879 .expect("controller successor disappeared on miss");
8880 successor.settle();
8881 fork.retire(generation)?;
8882 }
8883 }
8884 total_drafted += k_round;
8885 total_accepted += n_acc;
8886 if let Some(t) = sess_telem {
8887 // Greedy, rejection-sampling, and grammar truncation all converge here after
8888 // the accept decision is already on host. Fixed-size relaxed atomics only.
8889 t.record_round(k_round, n_acc);
8890 }
8891 if spec_stats {
8892 st_len_hist[k_round] += 1;
8893 for j in 0..k_round {
8894 st_drafted[j] += 1;
8895 }
8896 for j in 0..n_acc {
8897 st_accepted[j] += 1;
8898 }
8899 if n_acc == k_round {
8900 st_full += 1;
8901 }
8902 }
8903
8904 if debug_spec {
8905 eprintln!(
8906 "[R{round}] pos={pos} out_len={} last_tok={last_token} draft={draft:?} n_acc={n_acc} bonus={bonus} t_pred0={}",
8907 out.len(),
8908 t_pred(0)
8909 );
8910 }
8911
8912 // --- 4. COMMIT: draft[0..n_acc] then bonus (n_acc + 1 tokens) ---
8913 let commit_started = std::time::Instant::now();
8914 // SESSION MODE: every accepted column is already in the CACHE — `out` must carry all
8915 // of them (overshoot past max_new included) or `committed` under-counts the cache rows
8916 // and the next turn's continuation seeds one token off (gate-caught 2026-07-05). The
8917 // single-shot path keeps the cap (its caller truncates + drops the cache anyway).
8918 for j in 0..n_acc {
8919 if !session_mode && out.len() >= max_new {
8920 break;
8921 }
8922 out.push(draft[j]);
8923 }
8924 if pen_on {
8925 pen_hist.extend_from_slice(&draft[0..n_acc]);
8926 pen_hist.push(bonus);
8927 }
8928 let bonus_emitted = session_mode || out.len() < max_new;
8929 if bonus_emitted {
8930 out.push(bonus);
8931 }
8932 last_token = bonus;
8933
8934 // --- 5. ROLLBACK + advance (§C) ---
8935 if n_acc == k_round && !spec_replay {
8936 // FULL ACCEPT, BONUS FOLD: all verify columns (pending? + drafts) are committed in
8937 // cache; the NEW bonus stays PENDING for the next round's verify batch — NO extra
8938 // T=1 trunk pass. The next draft chain seeds from the MTP block's h_nextn at the
8939 // bonus position: one MTP-block pass (~1/33 trunk cost) replaces the trunk read.
8940 // last_pred is dead in the pending path (t_pred reads verify col 0).
8941 //
8942 // PERSISTENT DRAFT KV, full-accept fill: the chain covered last_token +
8943 // draft[0..k_round-2] as INPUTS (slots P..P'-2); draft[k_round-1] (slot P'-1) was
8944 // only ever an output, so its entry is MISSING. Fill it from vh_seed — its EXACT
8945 // trunk hidden (the last verify column). set_len first: a p-min break may have
8946 // left one extra chain append at that slot. Partial accepts need NO fill (the
8947 // chain already covered every accepted position; round-start set_len truncates).
8948 let mut vh_seed = e.zeros(n_embd)?;
8949 e.copy_view_into(
8950 &mut vh_seed,
8951 0,
8952 &vx.slice((t_v - 1) * n_embd..t_v * n_embd),
8953 n_embd,
8954 )?;
8955 if refresh {
8956 // TRUE-HIDDEN REFRESH (2026-07-03, the HANDOVER-listed acceptance lever):
8957 // overwrite ALL committed positions' scratch entries with K/V from their EXACT
8958 // verify hiddens — the reference engine's mtp_update fills from true hiddens;
8959 // the full stack (vx) is already resident from the verify. Replaces both the
8960 // chain-approximate entries AND the old last-token-only fill. Acceptance-only
8961 // (draft attention quality); exactness stays the verify's job.
8962 scratch.set_len(e, pos)?;
8963 // PREDECESSOR pairing: row i gets vx[i-1]; row 0 the carried fill_prev
8964 // (hidden of the last committed row before this verify batch).
8965 let mut vxs = e.zeros(t_v * n_embd)?;
8966 e.copy_into(&mut vxs, 0, &fill_prev, n_embd)?;
8967 if t_v > 1 {
8968 e.copy_view_into(
8969 &mut vxs,
8970 n_embd,
8971 &vx.slice(0..(t_v - 1) * n_embd),
8972 (t_v - 1) * n_embd,
8973 )?;
8974 }
8975 self.mtp_kv_fill(e, mtp, &verify_tokens, &vxs, pos, &mut *scratch, embd_dev)?;
8976 } else {
8977 scratch.set_len(e, pos + base + k_round - 1)?;
8978 // predecessor of the last draft = verify col t_v-2 (or fill_prev at t_v==1)
8979 let mut hp = e.zeros(n_embd)?;
8980 if t_v >= 2 {
8981 e.copy_view_into(
8982 &mut hp,
8983 0,
8984 &vx.slice((t_v - 2) * n_embd..(t_v - 1) * n_embd),
8985 n_embd,
8986 )?;
8987 } else {
8988 e.copy_into(&mut hp, 0, &fill_prev, n_embd)?;
8989 }
8990 self.mtp_kv_fill(
8991 e,
8992 mtp,
8993 &[draft[k_round - 1]],
8994 &hp,
8995 pos + base + k_round - 1,
8996 &mut *scratch,
8997 embd_dev,
8998 )?;
8999 }
9000 // REFERENCE SEEDING: no pseudo pass — the next chain's step 0 IS the
9001 // reference's (id_last, h_prev) draft row; it appends the bonus's scratch
9002 // entry itself. Seed = TRUE hidden of the bonus's predecessor (last verify
9003 // col). Saves one MTP-block pass per round on top of the pairing fix.
9004 if !devacc_seeded {
9005 e.copy_into(&mut h_seed_buf, 0, &vh_seed, n_embd)?;
9006 e.copy_into(&mut fill_prev, 0, &vh_seed, n_embd)?;
9007 }
9008 pending = Some(bonus);
9009 if debug_spec {
9010 eprintln!(" -> FULL ACCEPT (bonus pending, prev-h seed)");
9011 }
9012 } else if !spec_replay && base + n_acc >= 1 {
9013 // PARTIAL ACCEPT, REPLAY-FREE (2026-07-03 — the profiled #1 long-ctx spec cost):
9014 // the verify's first j = base+n_acc columns ARE the committed sequence, computed
9015 // bit-identically to eager (decode-exact contract) — so KEEP them: KV truncates to
9016 // pos+j, recurrent state rebuilds from the VerifyCkpt (same-kernel gdn prefix
9017 // re-run / pure state-clone restore), and the bonus stays PENDING exactly like the
9018 // full-accept path — the legacy duplicate trunk replay is gone. The next chain
9019 // seeds from the MTP pseudo-hidden of the bonus, whose seed = the TRUE verify
9020 // hidden of its predecessor (col j-1) — same one-hop pseudo structure as full
9021 // accept (never compounds: the next verify recomputes true hiddens for all
9022 // committed columns).
9023 let j = base + n_acc;
9024 self.commit_verified_prefix(
9025 e,
9026 &mut *cache,
9027 &snap,
9028 ckpt.as_ref().unwrap(),
9029 j,
9030 devacc_seeded,
9031 if devacc_seeded {
9032 devacc_acc.as_ref().map(|a| (a, base, t_v))
9033 } else {
9034 None
9035 },
9036 )?;
9037 let mut seed = e.zeros(n_embd)?;
9038 e.copy_view_into(
9039 &mut seed,
9040 0,
9041 &vx.slice((j - 1) * n_embd..j * n_embd),
9042 n_embd,
9043 )?;
9044 // Draft scratch: TRUE-HIDDEN REFRESH of the committed prefix (see the full-accept
9045 // branch); without it the chain entries stand and only the tail truncates. Either
9046 // way len ends at pos+j so the pseudo append lands at the bonus's slot pos+j
9047 // (persistent mode), rope pos+j+1 (chain convention).
9048 if refresh {
9049 scratch.set_len(e, pos)?;
9050 let mut vxs = e.zeros(j * n_embd)?;
9051 e.copy_into(&mut vxs, 0, &fill_prev, n_embd)?;
9052 if j > 1 {
9053 e.copy_view_into(
9054 &mut vxs,
9055 n_embd,
9056 &vx.slice(0..(j - 1) * n_embd),
9057 (j - 1) * n_embd,
9058 )?;
9059 }
9060 self.mtp_kv_fill(
9061 e,
9062 mtp,
9063 &verify_tokens[0..j],
9064 &vxs,
9065 pos,
9066 &mut *scratch,
9067 embd_dev,
9068 )?;
9069 } else {
9070 scratch.set_len(e, pos + j)?;
9071 }
9072 // REFERENCE SEEDING (see the full-accept branch): seed = TRUE hidden of the
9073 // bonus's predecessor (verify col j-1); no pseudo pass.
9074 if !devacc_seeded {
9075 e.copy_into(&mut h_seed_buf, 0, &seed, n_embd)?;
9076 e.copy_into(&mut fill_prev, 0, &seed, n_embd)?;
9077 }
9078 pending = Some(bonus);
9079 if debug_spec {
9080 eprintln!(" -> PARTIAL(replay-free j={j}, bonus pending, prev-h seed)");
9081 }
9082 } else if !spec_replay {
9083 // ZERO ROUND FOLD (2026-07-10, verify-cost target #3): base+n_acc == 0 — a
9084 // pending-less round where nothing was accepted (PMIN0 zero-draft chains after a
9085 // replay/commit, or plain 0-accept rounds at round 0). The old path replayed
9086 // [bonus] through a FULL m=1 trunk+head forward (the 489us full-vocab head pass
9087 // measured at ~0.75/round on PMIN0 configs). Instead: restore the pre-round
9088 // snapshot and let the bonus ride the NEXT round's verify as col 0 — the existing
9089 // base=1 pending machinery, bit-identical by the decode-exact verify contract.
9090 // Seed: the bonus's predecessor is the last COMMITTED token, whose hidden
9091 // fill_prev already carries (same seeding as the 1-token-replay case it replaces).
9092 cache.rollback(e, &snap, 0)?;
9093 scratch.set_len(e, pos)?;
9094 e.copy_into(&mut h_seed_buf, 0, &fill_prev, n_embd)?;
9095 pending = Some(bonus);
9096 if debug_spec {
9097 eprintln!(" -> ZERO-ROUND FOLD (bonus pending, fill_prev seed)");
9098 }
9099 } else {
9100 // PARTIAL ACCEPT, LEGACY REPLAY (seam MEMRA_SPEC_REPLAY=1 — or j==0: nothing of
9101 // this round survives, only possible before the first pending exists, ~round 0):
9102 // restore EVERYTHING to the pre-round snapshot (KV truncate to pos + recur
9103 // restore), then replay the committed prefix pending? ++ draft[0..n_acc] ++
9104 // [bonus] as ONE batched T forward — single weight read, bit-identical to greedy
9105 // (the verify-all-columns path is the same math). Commits the bonus with a TRUE
9106 // trunk hidden.
9107 cache.rollback(e, &snap, 0)?; // accept_len=0: KV len = pos, recur = snapshot
9108 let mut replay: Vec<u32> = Vec::with_capacity(base + n_acc + 1);
9109 if let Some(b) = pending.take() {
9110 replay.push(b);
9111 }
9112 replay.extend_from_slice(&draft[0..n_acc]);
9113 replay.push(bonus);
9114 // Full-stack forward (decode_step_t_core = decode_step_t_h_emb_dev's body):
9115 // Predecessor pairing seeds from the PREDECESSOR row (col len-2) — the same-row path takes the
9116 // last col exactly as before (byte-identical to the old _h_emb_dev call).
9117 let (rl_d, rx) = if self.qwen35_serving_class() {
9118 let mut logits = Vec::with_capacity(replay.len() * n_vocab);
9119 let mut hidden = e.uninit(replay.len() * n_embd)?;
9120 for (row, &token) in replay.iter().enumerate() {
9121 let (row_logits, row_hidden) =
9122 self.spec_target_step_h(e, token, &mut *cache)?;
9123 logits.extend_from_slice(&row_logits);
9124 e.dtod_copy_into(&row_hidden, &mut hidden, row * n_embd)?;
9125 }
9126 (e.htod(&logits)?, hidden)
9127 } else {
9128 self.decode_step_t_core(e, &replay, pos, &mut *cache, embd_dev, None)?
9129 };
9130 // last_pred = argmax of the LAST column's logits (predicts the token after `bonus`)
9131 // — device argmax + one 4-byte read instead of the full-vocab column dtoh.
9132 e.argmax_token_device_col(&rl_d, replay.len() - 1, n_vocab, &mut preds_d, 0)?;
9133 last_pred = e.dtoh_u32(&preds_d)?[0];
9134 if sampled {
9135 let lr0 = replay.len();
9136 let lc = last_col_logits
9137 .as_mut()
9138 .expect("sampled: last_col_logits unset");
9139 e.copy_view_into(
9140 lc,
9141 0,
9142 &rl_d.slice((lr0 - 1) * n_vocab..lr0 * n_vocab),
9143 n_vocab,
9144 )?;
9145 }
9146 let lr = replay.len();
9147 if lr >= 2 {
9148 e.copy_view_into(
9149 &mut h_seed_buf,
9150 0,
9151 &rx.slice((lr - 2) * n_embd..(lr - 1) * n_embd),
9152 n_embd,
9153 )?;
9154 } else {
9155 // 1-token replay (round-0 miss): the bonus's predecessor is the OLD
9156 // last_token, whose own-row hidden fill_prev still holds.
9157 e.copy_into(&mut h_seed_buf, 0, &fill_prev, n_embd)?;
9158 }
9159 // the bonus is COMMITTED here — it becomes the last committed row.
9160 let mut rh_last = e.zeros(n_embd)?;
9161 e.copy_view_into(
9162 &mut rh_last,
9163 0,
9164 &rx.slice((lr - 1) * n_embd..lr * n_embd),
9165 n_embd,
9166 )?;
9167 e.copy_into(&mut fill_prev, 0, &rh_last, n_embd)?;
9168 if debug_spec {
9169 eprintln!(" -> PARTIAL(replay={replay:?}), next_pred={last_pred}");
9170 }
9171 }
9172 if devacc_seeded {
9173 // stage (b) epilogue: fill_prev takes the gathered seed AFTER the refresh fills
9174 // consumed the old value (both slots carry the same value in every non-replay arm).
9175 e.copy_into(&mut fill_prev, 0, &h_seed_buf, n_embd)?;
9176 }
9177 if successor_valid {
9178 let optimistic_scratch_len = successor_attempt
9179 .as_ref()
9180 .expect("valid controller successor disappeared")
9181 .scratch_len;
9182 // The normal current-round commit refreshed/truncated the logical scratch tail.
9183 // Its optimistic successor row was already written physically, so restoring only
9184 // the retained logical length makes that row live for the carried round.
9185 scratch.set_len(e, optimistic_scratch_len)?;
9186 }
9187 if let Some(current) = current_opti.take() {
9188 opti_fork
9189 .as_mut()
9190 .ok_or("optipipe current retirement lost fork state")?
9191 .retire(current.generation)?;
9192 }
9193 if successor_valid {
9194 let successor = successor_attempt
9195 .take()
9196 .expect("valid controller successor disappeared before promotion");
9197 let generation = successor.generation;
9198 opti_fork
9199 .as_mut()
9200 .ok_or("optipipe successor promotion lost fork state")?
9201 .promote_successor_snapshot(&mut snap, generation);
9202 carried_opti = Some(successor);
9203 }
9204 if anatomy_on {
9205 // Commit/rollback is normally asynchronous on the primary/head stream. Bound it
9206 // only for this diagnostic so it does not disappear into the following draft's
9207 // first token readback.
9208 e.stream().synchronize()?;
9209 ph_commit += commit_started.elapsed().as_secs_f64();
9210 }
9211 // adaptive-K update (host math, zero syncs): next round drafts accepted-run + 1,
9212 // clamped to [floor(pos), k_cap]. cache.pos is post-rollback here (the round's
9213 // final position — the floor's position key reads the committed depth). Burst
9214 // rounds (`continue` above) draft the captured fixed depth and skip this, exactly
9215 // like gemma's burst arm.
9216 if adapt {
9217 let fl_now = floor_at(cache.pos);
9218 kc = (n_acc + 1).clamp(fl_now.min(k_cap), k_cap);
9219 }
9220 ph_mark(&mut ph_rest, phase_on);
9221 if let Some(p) = pipe {
9222 p.accept_end(round);
9223 }
9224 drop(pipe_accept);
9225 round += 1;
9226 // sse-cadence: this round's accepted drafts + bonus are committed (out is
9227 // append-only past step 4) — flush at round cadence.
9228 keep_going = flush_commit(&mut on_commit, &out, &mut flushed);
9229 }
9230 if let Some(mut ticket) = carried_opti.take() {
9231 opti_fork
9232 .as_mut()
9233 .ok_or("optipipe tail drain lost fork state")?
9234 .cancel_controller_ticket(e, &mut *cache, &mut *scratch, &snap, &mut ticket)?;
9235 }
9236 // sse-cadence: nothing below appends to `out`; flush any remainder (defensive).
9237 // (verdict ignored — the burst is over either way; the session tail runs unchanged.)
9238 let _ = flush_commit(&mut on_commit, &out, &mut flushed);
9239
9240 if spec_stats {
9241 let per_slot: Vec<String> = (0..k)
9242 .map(|j| {
9243 if st_drafted[j] > 0 {
9244 format!(
9245 "{}/{}={:.3}",
9246 st_accepted[j],
9247 st_drafted[j],
9248 st_accepted[j] as f64 / st_drafted[j] as f64
9249 )
9250 } else {
9251 "0/0".into()
9252 }
9253 })
9254 .collect();
9255 let acc = if total_drafted > 0 {
9256 total_accepted as f64 / total_drafted as f64
9257 } else {
9258 0.0
9259 };
9260 eprintln!(
9261 "[spec-stats] rounds={round} full_accept={st_full} len_hist={st_len_hist:?} \
9262 per_slot=[{}] total={total_accepted}/{total_drafted}={acc:.3} \
9263 tok_per_round={:.3}",
9264 per_slot.join(" "),
9265 (total_accepted + round) as f64 / round.max(1) as f64
9266 );
9267 }
9268 if constraint.is_some() {
9269 eprintln!(
9270 "[draft-mask] mask_rounds={dm_rounds} clone_total={:.3}ms \
9271 clone_per_round={:.4}ms gram_cuts={dm_cuts}/{round} cut_tokens={dm_cut_tokens}",
9272 dm_clone_ns as f64 / 1e6,
9273 dm_clone_ns as f64 / 1e6 / dm_rounds.max(1) as f64
9274 );
9275 }
9276 if phase_on {
9277 let tot = ph_draft + ph_verify + ph_wait + ph_rest;
9278 eprintln!(
9279 "[spec-phase] draft={:.1}ms ({:.1}%) verify-issue={:.1}ms ({:.1}%) verify-wait={:.1}ms ({:.1}%) commit-host={:.1}ms ({:.1}%) rounds={round}",
9280 ph_draft * 1e3,
9281 ph_draft / tot * 100.0,
9282 ph_verify * 1e3,
9283 ph_verify / tot * 100.0,
9284 ph_wait * 1e3,
9285 ph_wait / tot * 100.0,
9286 ph_rest * 1e3,
9287 ph_rest / tot * 100.0
9288 );
9289 }
9290 if anatomy_on {
9291 let rounds_f = round.max(1) as f64;
9292 let other = (ph_rest - ph_commit).max(0.0);
9293 eprintln!(
9294 "[spec-anatomy] per-round draft={:.3}ms pp-verify={:.3}ms \
9295 verify-accept={:.3}ms commit-rollback={:.3}ms other={:.3}ms rounds={round}",
9296 ph_draft * 1e3 / rounds_f,
9297 ph_verify * 1e3 / rounds_f,
9298 ph_wait * 1e3 / rounds_f,
9299 ph_commit * 1e3 / rounds_f,
9300 other * 1e3 / rounds_f,
9301 );
9302 }
9303 let _pipe_tail = pipe.map(|p| p.primary());
9304 // SESSION TAIL: leave the session in the exact invariant the next turn's suffix prime
9305 // expects — every row in `committed` has trunk KV/recur state AND an exact draft-KV row.
9306 // Park the draft-graph ctx back on the session (the serve-burst fixed-cost fix): the next
9307 // burst replays instead of recapturing. Error paths (`?` above) drop it — recaptured then.
9308 if let Some(slot) = sess_draft_slot.take() {
9309 *slot = Some(dctx);
9310 }
9311 let t_rounds = t_ent.elapsed();
9312 if let Some((committed, last_h, next_pred_slot, sctr_slot, uctr_slot)) = sess_tail.take() {
9313 *sctr_slot = sctr;
9314 *uctr_slot = uctr;
9315 *next_pred_slot = Some(last_pred);
9316 let mut stashed_pending = false;
9317 if let Some(b) = pending.take() {
9318 if !sampled {
9319 // PENDING-CARRY (2026-08-01): stash the bonus on the session instead of
9320 // committing it with a solo T=1 pass — the next empty-suffix greedy burst
9321 // consumes it as round-0 verify col 0 (a plain round edge; the old tail
9322 // commit + next burst's init feed were 11.6+11.5ms solo trunk passes per
9323 // burst on H100 q27, [spec-setup] trace). b stays in `out` (emitted) but
9324 // OUT of `committed` (cache rows == committed); the consuming call
9325 // prepends it once its verify commits the row. next_pred is unknowable
9326 // without the commit pass — None; callers gate on pending_tok too.
9327 debug_assert_eq!(out.last(), Some(&b), "pending must be the last emitted");
9328 if let Some(slot) = sess_pending_slot.take() {
9329 *slot = Some(b);
9330 }
9331 *next_pred_slot = None;
9332 // fill_prev = hidden of the last COMMITTED row (b's predecessor) — the
9333 // exact chain-seed/fill anchor the consuming burst (or a flush) needs.
9334 *last_h = Some(e.clone_dtod(&fill_prev)?);
9335 stashed_pending = true;
9336 } else {
9337 // SAMPLED tail (unchanged): commit the bonus (one T=1 pass) + draft fill —
9338 // the sampled round-0 accept needs this pass's logits (last_col_logits).
9339 let pos_b = cache.pos;
9340 scratch.set_len(e, pos_b)?;
9341 let (lg_b, hb) = self.spec_target_step_h(e, b, &mut *cache)?;
9342 // after a FULL-accept exit `last_pred` is STALE (it predicted the bonus
9343 // itself — the prediction AFTER the bonus never materialized; it would have
9344 // been the next round's verify col 0). The commit's logits ARE that
9345 // prediction.
9346 *next_pred_slot = Some(argmax(&lg_b) as u32);
9347 self.mtp_kv_fill(e, mtp, &[b], &fill_prev, pos_b, &mut *scratch, embd_dev)?;
9348 *last_h = Some(hb);
9349 }
9350 } else {
9351 // fill_prev tracks the hidden of the last COMMITTED row throughout the loop.
9352 *last_h = Some(e.clone_dtod(&fill_prev)?);
9353 }
9354 committed.extend_from_slice(prompt);
9355 if let Some(cb) = carried_pending {
9356 // the consumed carry's cache row landed in round 0's verify (every pending
9357 // round commits col 0) — it joins `committed` here, in sequence order.
9358 committed.push(cb);
9359 }
9360 if stashed_pending {
9361 // ZERO-EMIT BURST (2026-08-06 c=8 serve panic, pre-existing since b4aea184):
9362 // `out.len() - 1` underflowed on an EMPTY `out` — "range end index
9363 // 18446744073709551615 out of range for slice of length 0", killing the
9364 // memra-gpu-worker and failing 31 of 32 concurrent requests with "worker closed
9365 // stream". Reachable because `pending` starts as `carried_pending` (a bonus
9366 // stashed by the PREVIOUS burst) while `out` starts empty, and the carry is
9367 // deliberately NOT pushed to `out` (line ~3239: the burst that emitted it already
9368 // did). So a burst that stashes a pending without emitting anything of its own —
9369 // the round loop exits before a push, e.g. the ring drain's `out.len() < max_new`
9370 // guard skipping every token under a tight budget — arrives here with
9371 // out.len() == 0 and stashed_pending == true.
9372 //
9373 // The invariant is unchanged: `committed` gets every emitted token EXCEPT the
9374 // stashed bonus. With nothing emitted, that is nothing — and the carry pushed
9375 // just above is already accounted. Saturating, not a min/assert: an empty `out`
9376 // here is a legitimate burst shape, not a corrupt state.
9377 let emitted = out.len().saturating_sub(1);
9378 committed.extend_from_slice(&out[..emitted]);
9379 } else {
9380 committed.extend_from_slice(&out); // FULL out incl. overshoot — all committed
9381 }
9382 debug_assert_eq!(
9383 cache.pos,
9384 committed.len(),
9385 "session invariant: cache rows == committed tokens"
9386 );
9387 if setup_trace {
9388 e.stream().synchronize()?; // bound the async tail fill in the trace
9389 let t_tail = t_ent.elapsed();
9390 eprintln!(
9391 "[spec-setup] init={:.2}ms cap={:.2}ms fill={:.2}ms rounds={:.2}ms tail={:.2}ms total={:.2}ms out={} cont={}",
9392 t_init.as_secs_f64() * 1e3,
9393 (t_cap - t_init).as_secs_f64() * 1e3,
9394 (t_fill - t_cap).as_secs_f64() * 1e3,
9395 (t_rounds - t_fill).as_secs_f64() * 1e3,
9396 (t_tail - t_rounds).as_secs_f64() * 1e3,
9397 t_tail.as_secs_f64() * 1e3,
9398 out.len(),
9399 continuation
9400 );
9401 }
9402 return Ok((out, total_drafted, total_accepted));
9403 }
9404 out.truncate(max_new);
9405 Ok((out, total_drafted, total_accepted))
9406 }
9407
9408 /// Anchor-bounded DSpark target extraction. The trunk sees the exact generated token tape;
9409 /// only requested hidden rows and target-logit rows cross PCIe. An anchor token at p pairs
9410 /// with the pre-output-norm h[p-1] carrier, exactly as the existing replay/NextN path does.
9411 pub fn extract_dspark_anchors(
9412 &self,
9413 e: &Engine,
9414 tokens: &[u32],
9415 anchor_positions: &[usize],
9416 gamma: usize,
9417 top_k: usize,
9418 chunk: usize,
9419 temperature: f32,
9420 ) -> Result<Vec<DsparkAnchorRecord>, Box<dyn std::error::Error>> {
9421 if tokens.len() < gamma + 2 || gamma == 0 || chunk < 2 {
9422 return Err("DSpark extraction token tape/gamma/chunk is invalid".into());
9423 }
9424 if anchor_positions.windows(2).any(|pair| pair[0] >= pair[1]) {
9425 return Err("DSpark anchor positions must be sorted and unique".into());
9426 }
9427 for &position in anchor_positions {
9428 if position == 0 || position + gamma >= tokens.len() {
9429 return Err(format!(
9430 "DSpark anchor {position} has no predecessor or cannot cover gamma={gamma} in {} tokens",
9431 tokens.len()
9432 )
9433 .into());
9434 }
9435 }
9436
9437 let n_vocab = self.output.out_features();
9438 let n_embd = self.cfg.n_embd as usize;
9439 let mut cache = crate::pp::new_cache(e, &self.cfg, tokens.len() + gamma + 8)?;
9440 let (embd_qt, embd_rb) = self.embd.qt_and_row_bytes(n_embd);
9441 let embd_gpu = if spec_host_embd() {
9442 None
9443 } else {
9444 Some(
9445 self.embd_gpu
9446 .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload")),
9447 )
9448 };
9449 let embd_dev = embd_gpu.map(|gpu| (gpu, embd_qt, embd_rb));
9450
9451 struct PendingRecord {
9452 position: usize,
9453 hidden: Option<Vec<f32>>,
9454 tokens: Vec<u32>,
9455 target_top_ids: Vec<Option<Vec<u32>>>,
9456 target_top_logits: Vec<Option<Vec<f32>>>,
9457 target_top_probs: Vec<Option<Vec<f32>>>,
9458 target_tail_probs: Vec<Option<f32>>,
9459 }
9460
9461 let mut pending: Vec<PendingRecord> = anchor_positions
9462 .iter()
9463 .map(|&position| PendingRecord {
9464 position,
9465 hidden: None,
9466 tokens: tokens[position..=position + gamma].to_vec(),
9467 target_top_ids: vec![None; gamma],
9468 target_top_logits: vec![None; gamma],
9469 target_top_probs: vec![None; gamma],
9470 target_tail_probs: vec![None; gamma],
9471 })
9472 .collect();
9473
9474 let mut start = 0usize;
9475 while start < tokens.len() {
9476 let end = (start + chunk).min(tokens.len());
9477 let chunk_tokens = &tokens[start..end];
9478 let (target_logits, hidden_rows) =
9479 self.decode_step_t_core(e, chunk_tokens, start, &mut cache, embd_dev, None)?;
9480 for record in &mut pending {
9481 let hidden_position = record.position - 1;
9482 if hidden_position >= start && hidden_position < end {
9483 let local = hidden_position - start;
9484 record.hidden = Some(
9485 e.dtoh_view(&hidden_rows.slice(local * n_embd..(local + 1) * n_embd))?,
9486 );
9487 }
9488 for slot in 0..gamma {
9489 let target_row = record.position + slot;
9490 if target_row < start || target_row >= end {
9491 continue;
9492 }
9493 let local = target_row - start;
9494 let logits =
9495 e.dtoh_view(&target_logits.slice(local * n_vocab..(local + 1) * n_vocab))?;
9496 let (ids, top_logits, probs, tail) =
9497 dspark_sparse_softmax_topk(&logits, top_k, temperature)?;
9498 record.target_top_ids[slot] = Some(ids);
9499 record.target_top_logits[slot] = Some(top_logits);
9500 record.target_top_probs[slot] = Some(probs);
9501 record.target_tail_probs[slot] = Some(tail);
9502 }
9503 }
9504 start = end;
9505 }
9506
9507 pending
9508 .into_iter()
9509 .map(|record| {
9510 let hidden = record
9511 .hidden
9512 .ok_or_else(|| format!("missing DSpark hidden at {}", record.position))?;
9513 let target_top_ids =
9514 flatten_dspark_rows(record.target_top_ids, record.position, "target ids")?;
9515 let target_top_logits = flatten_dspark_rows(
9516 record.target_top_logits,
9517 record.position,
9518 "target logits",
9519 )?;
9520 let target_top_probs =
9521 flatten_dspark_rows(record.target_top_probs, record.position, "target probs")?;
9522 let target_tail_probs = record
9523 .target_tail_probs
9524 .into_iter()
9525 .enumerate()
9526 .map(|(slot, value)| {
9527 value.ok_or_else(|| {
9528 format!("missing DSpark tail at {} slot {slot}", record.position)
9529 })
9530 })
9531 .collect::<Result<Vec<_>, _>>()?;
9532 Ok(DsparkAnchorRecord {
9533 position: record.position,
9534 hidden,
9535 tokens: record.tokens,
9536 target_top_ids,
9537 target_top_logits,
9538 target_top_probs,
9539 target_tail_probs,
9540 })
9541 })
9542 .collect()
9543 }
9544
9545 /// TEACHER-FORCED REPLAY ACCEPTANCE (hqmtp MTP-heal protocol): walk a FIXED token
9546 /// sequence and, at sampled positions, compare the MTP head's K-token draft chain against
9547 /// the trunk's own teacher-forced greedy predictions. Nothing is generated — the context is
9548 /// the corpus text itself, so (a) degenerate self-generated loops cannot inflate acceptance
9549 /// and (b) two arms (bf16 ceiling vs NVFP4) score on IDENTICAL contexts, isolating the
9550 /// quant-induced head/hidden-state mismatch from text drift.
9551 ///
9552 /// Per eval position p (context = tokens[0..=p], predecessor pairing as in spec decode):
9553 /// draft_j = chain token j from (tokens[p], h_{p-1}), then its own drafts — the exact
9554 /// eager spec-decode chain (same mtp_head_forward_dev, same rope positions).
9555 /// target_j = teacher-forced greedy pick for position p+1+j (argmax of the trunk logits
9556 /// at forced context tokens[0..p+j]). For j==0 this equals live spec
9557 /// acceptance; for j>=1 live verify would condition on the drafts, here it
9558 /// conditions on the corpus — deterministic and arm-comparable by design.
9559 ///
9560 /// Returns (rows, bg): one (p, drafts[k], targets[k]) row per eval position (ascending p),
9561 /// plus the full teacher-forced greedy track bg (bg[i] = greedy pick for position i, i>=1)
9562 /// so harnesses can cross-check runs (e.g. different chunk sizes must give identical bg).
9563 ///
9564 /// `hdump`: when Some, every position's pre-output_norm trunk hidden (the exact rows the
9565 /// draft-KV fill pairs from) streams to the file as little-endian f32 [t_total, n_embd] —
9566 /// the head-distillation extraction (hqmtp): the ENGINE is the source of truth for trunk
9567 /// hiddens (HF torch reproductions of the hybrid trunk measured only ~0.5 greedy
9568 /// agreement vs this path — not usable as a training-data source).
9569 pub fn replay_acceptance(
9570 &self,
9571 e: &Engine,
9572 tokens: &[u32],
9573 k: usize,
9574 stride: usize,
9575 chunk: usize,
9576 mut hdump: Option<&mut std::fs::File>,
9577 ) -> Result<(Vec<(usize, Vec<u32>, Vec<u32>)>, Vec<u32>), Box<dyn std::error::Error>> {
9578 assert!(k >= 1 && stride >= 1 && chunk >= 2);
9579 let mtp = self
9580 .mtp
9581 .as_ref()
9582 .expect("replay_acceptance requires an MTP head");
9583 let n_vocab = self.output.out_features();
9584 let d_vocab = mtp
9585 .shared_head_head
9586 .as_ref()
9587 .unwrap_or(&self.output)
9588 .out_features();
9589 let n_embd = self.cfg.n_embd as usize;
9590 let t_total = tokens.len();
9591 assert!(t_total >= 8, "corpus too short ({t_total} tokens)");
9592 // STAGE-OWNED KV (lane/pp2-spec 2026-08-06) — see `new_session`. Door shut = `Cache::new`.
9593 let mut cache = crate::pp::new_cache(e, &self.cfg, t_total + k + 8)?;
9594 let mut scratch = MtpScratch::new(
9595 e,
9596 &self.cfg,
9597 t_total + k + 8,
9598 self.mtp.as_ref().and_then(|m| m.geom.as_ref()),
9599 )?;
9600 let (embd_qt, embd_rb) = self.embd.qt_and_row_bytes(n_embd);
9601 let embd_gpu = if spec_host_embd() {
9602 None
9603 } else {
9604 Some(
9605 self.embd_gpu
9606 .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload")),
9607 )
9608 };
9609 let embd_dev = embd_gpu.map(|g| (g, embd_qt, embd_rb));
9610
9611 // bg[i] = the trunk's greedy pick for position i under the forced context (i >= 1).
9612 let mut bg: Vec<u32> = vec![0; t_total + 1];
9613 let mut rows: Vec<(usize, Vec<u32>, Vec<u32>)> = Vec::new();
9614 let mut prev_last_h = e.zeros(n_embd)?; // predecessor hidden entering the chunk
9615 let mut seed_buf = e.zeros(n_embd)?;
9616 let mut preds_d = e.alloc_u32_zeroed(chunk)?;
9617 let nll_on = std::env::var("MEMRA_REPLAY_NLL").as_deref() == Ok("1");
9618 let (mut nll_sum, mut nll_cnt) = (0f64, 0u64);
9619 let mut s = 0usize;
9620 while s < t_total {
9621 let cend = (s + chunk).min(t_total);
9622 let tc = cend - s;
9623 let ch = &tokens[s..cend];
9624 // 1. forced trunk pass — verify path (decode-exact contract): all-column logits +
9625 // the chunk's true hiddens.
9626 let (tl_d, vx) = self.decode_step_t_core(e, ch, s, &mut cache, embd_dev, None)?;
9627 for j in 0..tc {
9628 e.argmax_token_device_col(&tl_d, j, n_vocab, &mut preds_d, j)?;
9629 }
9630 let preds = e.dtoh_u32(&preds_d)?;
9631 for j in 0..tc {
9632 bg[s + j + 1] = preds[j];
9633 }
9634 // MEMRA_REPLAY_NLL=1: teacher-forced NLL/perplexity over the same forced pass — the
9635 // checkpoint-quality metric (position j's logits score the GOLD next token).
9636 if nll_on {
9637 let jmax = if cend < t_total { tc } else { tc - 1 }; // last pos has no gold next
9638 if jmax > 0 {
9639 let ids: Vec<u32> = (0..jmax).map(|j| tokens[s + j + 1]).collect();
9640 let rows: Vec<i32> = (0..jmax as i32).collect();
9641 let idsd = e.htod_u32_v(&ids)?;
9642 let rowsd = e.htod_i32(&rows)?;
9643 let mut outd = e.zeros(jmax)?;
9644 e.softmax_gather(&tl_d, n_vocab, &idsd, &rowsd, &mut outd, n_vocab, jmax, 1.0)?;
9645 for pr in e.dtoh(&outd)? {
9646 nll_sum += -((pr.max(1e-30)) as f64).ln();
9647 nll_cnt += 1;
9648 }
9649 }
9650 }
9651 if let Some(f) = hdump.as_deref_mut() {
9652 use std::io::Write;
9653 let host: Vec<f32> = e.dtoh(&vx)?;
9654 // bf16 round-to-nearest-even — f32 doubled the disk bill at bulk
9655 // extraction scale (20M tokens x 4096 = 320GB f32 vs 160GB bf16).
9656 let mut bytes = Vec::with_capacity(tc * n_embd * 2);
9657 for v in &host[..tc * n_embd] {
9658 let b = v.to_bits();
9659 let r = b.wrapping_add(0x7FFF + ((b >> 16) & 1));
9660 bytes.extend_from_slice(&((r >> 16) as u16).to_le_bytes());
9661 }
9662 f.write_all(&bytes)?;
9663 }
9664 // CHAINLESS extraction (stride > corpus, the bulk-hdump mode): no chunk ever
9665 // drafts, so the draft-KV fills are pure waste — skip them (2 MTP-block passes
9666 // per token saved; the forced trunk pass + hdump is all the mode needs).
9667 let chainless = stride > t_total;
9668 if chainless {
9669 e.copy_view_into(
9670 &mut prev_last_h,
9671 0,
9672 &vx.slice((tc - 1) * n_embd..tc * n_embd),
9673 n_embd,
9674 )?;
9675 s = cend;
9676 continue;
9677 }
9678 // 2. TRUE predecessor-paired draft-KV fill for the chunk (row i carries h_{i-1};
9679 // row s reads the previous chunk's last true hidden, zeros at corpus start).
9680 let mut vxs = e.zeros(tc * n_embd)?;
9681 e.copy_into(&mut vxs, 0, &prev_last_h, n_embd)?;
9682 if tc > 1 {
9683 e.copy_view_into(
9684 &mut vxs,
9685 n_embd,
9686 &vx.slice(0..(tc - 1) * n_embd),
9687 (tc - 1) * n_embd,
9688 )?;
9689 }
9690 scratch.set_len(e, s)?;
9691 self.mtp_kv_fill(e, mtp, ch, &vxs, s, &mut scratch, embd_dev)?;
9692 // 3. draft chains at sampled positions, DESCENDING: a chain reads only slots
9693 // [0..p) (true fills) and appends at >= p; the next (smaller-p) chain's set_len
9694 // truncates those approximate appends before they can ever be read.
9695 let ps: Vec<usize> = (s..cend)
9696 .filter(|p| *p >= 1 && *p % stride == 0 && *p + k <= t_total)
9697 .collect();
9698 for &p in ps.iter().rev() {
9699 scratch.set_len(e, p)?;
9700 if p == s {
9701 e.copy_into(&mut seed_buf, 0, &prev_last_h, n_embd)?;
9702 } else {
9703 e.copy_view_into(
9704 &mut seed_buf,
9705 0,
9706 &vx.slice((p - 1 - s) * n_embd..(p - s) * n_embd),
9707 n_embd,
9708 )?;
9709 }
9710 let mut e_tok = tokens[p];
9711 let mut d_seed = e.clone_dtod(&seed_buf)?;
9712 let mut drafts: Vec<u32> = Vec::with_capacity(k);
9713 for j in 0..k {
9714 let (dl_d, h_nextn) = self.mtp_head_forward_dev(
9715 e,
9716 mtp,
9717 e_tok,
9718 &d_seed,
9719 &mut scratch,
9720 p + 1 + j,
9721 embd_dev,
9722 None, // acceptance-oracle walk: no grammar
9723 )?;
9724 let tok_d = e.argmax_token_device(&dl_d, d_vocab)?;
9725 let idx = e.dtoh_u32_one(&tok_d)?;
9726 let d = match &mtp.d2t {
9727 Some(map) => map[idx as usize],
9728 None => idx,
9729 };
9730 drafts.push(d);
9731 e_tok = d;
9732 d_seed = h_nextn;
9733 }
9734 // targets may live in a LATER chunk's bg — resolved after the walk.
9735 rows.push((p, drafts, Vec::new()));
9736 }
9737 // 4. restore TRUE entries for the whole chunk (the next chunk's chains and fills
9738 // expect scratch.len == cend with exact rows).
9739 scratch.set_len(e, s)?;
9740 self.mtp_kv_fill(e, mtp, ch, &vxs, s, &mut scratch, embd_dev)?;
9741 e.copy_view_into(
9742 &mut prev_last_h,
9743 0,
9744 &vx.slice((tc - 1) * n_embd..tc * n_embd),
9745 n_embd,
9746 )?;
9747 s = cend;
9748 }
9749 for (p, drafts, targets) in rows.iter_mut() {
9750 for j in 0..drafts.len() {
9751 targets.push(bg[*p + 1 + j]);
9752 }
9753 }
9754 rows.sort_by_key(|r| r.0);
9755 if nll_cnt > 0 {
9756 let mean = nll_sum / nll_cnt as f64;
9757 println!(
9758 "[replay-nll] tokens={nll_cnt} nll/token={mean:.5} ppl={:.4}",
9759 mean.exp()
9760 );
9761 }
9762 Ok((rows, bg))
9763 }
9764}
9765
9766#[cfg(test)]
9767mod dspark_sparse_tests {
9768 use super::dspark_sparse_softmax_topk;
9769
9770 #[test]
9771 fn topk_keeps_full_softmax_mass_and_stable_ties() {
9772 let logits = [1.0f32, 3.0, 3.0, -2.0];
9773 let (ids, top_logits, probs, tail) = dspark_sparse_softmax_topk(&logits, 2, 1.0).unwrap();
9774 assert_eq!(ids, vec![1, 2]);
9775 assert_eq!(top_logits, vec![3.0, 3.0]);
9776 let denominator = logits.iter().map(|value| (value - 3.0).exp()).sum::<f32>();
9777 let expected = 1.0 / denominator;
9778 assert!((probs[0] - expected).abs() < 1.0e-6);
9779 assert!((probs[1] - expected).abs() < 1.0e-6);
9780 assert!((tail - (1.0 - 2.0 * expected)).abs() < 1.0e-6);
9781 assert!((probs.iter().sum::<f32>() + tail - 1.0).abs() < 1.0e-6);
9782 }
9783}
9784
9785#[cfg(test)]
9786mod spec_replay_env_tests {
9787 use super::spec_replay_env_on;
9788
9789 #[test]
9790 fn replay_requires_literal_one() {
9791 assert!(!spec_replay_env_on(None));
9792 assert!(!spec_replay_env_on(Some("")));
9793 assert!(!spec_replay_env_on(Some("0")));
9794 assert!(!spec_replay_env_on(Some("true")));
9795 assert!(!spec_replay_env_on(Some("2")));
9796 assert!(spec_replay_env_on(Some("1")));
9797 }
9798}
9799
9800#[cfg(test)]
9801mod telem_tests {
9802 use super::{SPEC_TELEM_POS, SpecTelemetry, SpecTelemetryCounters};
9803
9804 #[test]
9805 fn synthetic_accept_masks_produce_tau_and_position_histogram() {
9806 let counters = SpecTelemetryCounters::default();
9807 for mask in [
9808 [true, true, true],
9809 [true, true, false],
9810 [true, false, false],
9811 [false, false, false],
9812 ] {
9813 let accepted = mask.iter().take_while(|&&value| value).count();
9814 counters.record_round(mask.len(), accepted);
9815 }
9816
9817 let snapshot = counters.snapshot();
9818 assert_eq!(
9819 (snapshot.rounds, snapshot.drafted, snapshot.accepted),
9820 (4, 12, 6)
9821 );
9822 assert_eq!(&snapshot.pos_drafted[..3], &[4, 4, 4]);
9823 assert_eq!(&snapshot.pos_accepted[..3], &[3, 2, 1]);
9824 assert_eq!(snapshot.tau(), 1.5);
9825 assert_eq!(snapshot.pos_drafted[3..], [0; SPEC_TELEM_POS - 3]);
9826 assert_eq!(snapshot.pos_accepted[3..], [0; SPEC_TELEM_POS - 3]);
9827 }
9828
9829 /// The worker's per-burst pattern: stash, accumulate, diff — the delta must isolate
9830 /// exactly the burst's contribution (pool-resumed sessions carry prior requests' counts).
9831 #[test]
9832 fn delta_isolates_burst_contribution() {
9833 let mut t = SpecTelemetry::default();
9834 // "previous request": 2 rounds of k=3, accepts 3 then 1.
9835 for (kr, na) in [(3usize, 3usize), (3, 1)] {
9836 t.rounds += 1;
9837 t.drafted += kr as u64;
9838 t.accepted += na as u64;
9839 for j in 0..kr {
9840 t.pos_drafted[j] += 1;
9841 }
9842 for j in 0..na {
9843 t.pos_accepted[j] += 1;
9844 }
9845 }
9846 let before = t;
9847 // "this burst": 1 round k=3, accepts 2.
9848 t.rounds += 1;
9849 t.drafted += 3;
9850 t.accepted += 2;
9851 for j in 0..3 {
9852 t.pos_drafted[j] += 1;
9853 }
9854 for j in 0..2 {
9855 t.pos_accepted[j] += 1;
9856 }
9857 let d = t.delta_since(&before);
9858 assert_eq!((d.rounds, d.drafted, d.accepted), (1, 3, 2));
9859 assert_eq!(&d.pos_drafted[..3], &[1, 1, 1]);
9860 assert_eq!(&d.pos_accepted[..3], &[1, 1, 0]);
9861 assert_eq!(d.pos_drafted[3..], [0; SPEC_TELEM_POS - 3]);
9862 }
9863
9864 /// merge(delta) then merge(delta2) equals accumulating both — the per-model /metrics
9865 /// aggregation invariant.
9866 #[test]
9867 fn merge_accumulates_fieldwise() {
9868 let mut agg = SpecTelemetry::default();
9869 let mut d1 = SpecTelemetry {
9870 rounds: 2,
9871 drafted: 6,
9872 accepted: 4,
9873 ..Default::default()
9874 };
9875 d1.pos_drafted[0] = 2;
9876 d1.pos_accepted[0] = 2;
9877 let mut d2 = SpecTelemetry {
9878 rounds: 1,
9879 drafted: 3,
9880 accepted: 1,
9881 ..Default::default()
9882 };
9883 d2.pos_drafted[0] = 1;
9884 d2.pos_accepted[0] = 1;
9885 d2.pos_drafted[1] = 1;
9886 agg.merge(&d1);
9887 agg.merge(&d2);
9888 assert_eq!((agg.rounds, agg.drafted, agg.accepted), (3, 9, 5));
9889 assert_eq!(agg.pos_drafted[0], 3);
9890 assert_eq!(agg.pos_accepted[0], 3);
9891 assert_eq!(agg.pos_drafted[1], 1);
9892 assert_eq!(agg.pos_accepted[1], 0);
9893 }
9894
9895 /// Wrong-snapshot diff saturates to zero instead of wrapping — the counters feed a
9896 /// public metrics surface and must never publish a u64-wrapped garbage value.
9897 #[test]
9898 fn delta_saturates_never_wraps() {
9899 let small = SpecTelemetry {
9900 rounds: 1,
9901 drafted: 2,
9902 accepted: 1,
9903 ..Default::default()
9904 };
9905 let big = SpecTelemetry {
9906 rounds: 5,
9907 drafted: 15,
9908 accepted: 9,
9909 ..Default::default()
9910 };
9911 let d = small.delta_since(&big);
9912 assert_eq!((d.rounds, d.drafted, d.accepted), (0, 0, 0));
9913 }
9914}
9915
9916#[cfg(test)]
9917mod opti_fork_tests {
9918 use super::{
9919 OptiControllerPolicy, OptiForkAction, OptiForkGateMode, OptiForkGenerationTracker,
9920 };
9921
9922 #[test]
9923 fn controller_threshold_and_three_miss_breaker_are_exact() {
9924 let mut policy = OptiControllerPolicy {
9925 threshold: 0.7,
9926 consecutive_misses: 0,
9927 breaker_tripped: false,
9928 };
9929 assert!(!policy.admit(0.699_999));
9930 assert!(policy.admit(0.7));
9931 assert!(!policy.resolve(false));
9932 assert!(!policy.resolve(false));
9933 assert!(policy.resolve(false));
9934 assert!(policy.breaker_tripped);
9935 assert!(!policy.admit(1.0));
9936 assert!(
9937 !policy.resolve(true),
9938 "a resolved hit cannot re-arm a tripped request"
9939 );
9940 assert!(policy.breaker_tripped);
9941 }
9942
9943 #[test]
9944 fn zero_threshold_is_the_true_unconditional_measurement_arm() {
9945 let mut policy = OptiControllerPolicy {
9946 threshold: 0.0,
9947 consecutive_misses: 0,
9948 breaker_tripped: false,
9949 };
9950 for _ in 0..16 {
9951 assert!(policy.admit(0.0));
9952 assert!(!policy.resolve(false));
9953 }
9954 for invalid in [f32::NAN, f32::INFINITY, -0.01, 1.01] {
9955 assert!(
9956 !policy.admit(invalid),
9957 "invalid q proxy must fail closed: {invalid}"
9958 );
9959 }
9960 assert!(!policy.breaker_tripped);
9961 assert_eq!(policy.consecutive_misses, 0);
9962 }
9963
9964 #[test]
9965 fn alternating_mode_flips_by_generation_not_round_parity() {
9966 assert_eq!(OptiForkGateMode::Alternate.action(0), OptiForkAction::Hit);
9967 assert_eq!(OptiForkGateMode::Alternate.action(1), OptiForkAction::Miss);
9968 assert_eq!(OptiForkGateMode::Alternate.action(8), OptiForkAction::Hit);
9969 assert_eq!(OptiForkGateMode::Alternate.action(9), OptiForkAction::Miss);
9970 }
9971
9972 #[test]
9973 fn live_generation_cannot_be_overwritten() {
9974 let mut tracker = OptiForkGenerationTracker::default();
9975 let g0 = tracker.reserve().unwrap();
9976 let g1 = tracker.reserve().unwrap();
9977 let err = tracker.reserve().unwrap_err().to_string();
9978 assert!(
9979 err.contains("still owns generation 0"),
9980 "unexpected error: {err}"
9981 );
9982 tracker.retire(g0).unwrap();
9983 let g2 = tracker.reserve().unwrap();
9984 assert_eq!((g2.id, g2.slot), (2, 0));
9985 tracker.retire(g1).unwrap();
9986 tracker.retire(g2).unwrap();
9987 }
9988
9989 #[test]
9990 fn teardown_rejects_a_stale_generation_tag() {
9991 let mut tracker = OptiForkGenerationTracker::default();
9992 let g0 = tracker.reserve().unwrap();
9993 tracker.retire(g0).unwrap();
9994 let err = tracker.retire(g0).unwrap_err().to_string();
9995 assert!(err.contains("teardown mismatch"), "unexpected error: {err}");
9996 }
9997}
9998
9999#[cfg(test)]
10000mod draft_graph_fallback_tests {
10001 use super::DraftGraphFallback;
10002
10003 /// The Q2 contract, part (a): a fallback flip is LOUD — exactly once per flip.
10004 #[test]
10005 fn flip_is_loud_once_and_memoized_after() {
10006 let mut f = DraftGraphFallback::default();
10007 let line = f
10008 .mark_greedy("out of memory")
10009 .expect("first flip must return the warn line");
10010 assert!(
10011 line.contains("WARN"),
10012 "flip line must be warn-level: {line}"
10013 );
10014 assert!(
10015 line.contains("out of memory"),
10016 "flip line must carry the reason: {line}"
10017 );
10018 assert!(f.greedy_failed());
10019 // re-marking an already-failed graph is the memoization: quiet, still failed.
10020 assert!(f.mark_greedy("out of memory").is_none());
10021 assert!(f.greedy_failed());
10022 // the two graphs' flags are independent (greedy flip leaves sampled capturable).
10023 assert!(!f.sampled_failed());
10024 let line_s = f
10025 .mark_sampled("capture unsupported")
10026 .expect("sampled flip is its own flip");
10027 assert!(
10028 line_s.contains("sampled"),
10029 "sampled flip names itself: {line_s}"
10030 );
10031 assert!(f.mark_sampled("capture unsupported").is_none());
10032 }
10033
10034 /// The Q2 contract, part (b): resume-from-pool RESETS both flags (fresh capture chance),
10035 /// and says so exactly when there was something to reset.
10036 #[test]
10037 fn reset_on_resume_clears_flags_and_logs_once() {
10038 let mut f = DraftGraphFallback::default();
10039 // clean session: resume is silent, nothing to reset.
10040 assert!(f.reset_on_resume().is_none());
10041 f.mark_greedy("oom").unwrap();
10042 f.mark_sampled("oom").unwrap();
10043 let note = f
10044 .reset_on_resume()
10045 .expect("a set flag must produce the reset note");
10046 assert!(
10047 note.contains("greedy+sampled"),
10048 "note names what was reset: {note}"
10049 );
10050 assert!(
10051 !f.greedy_failed() && !f.sampled_failed(),
10052 "both flags cleared"
10053 );
10054 // and the NEXT failure after a reset is a fresh flip — loud again.
10055 assert!(f.mark_greedy("oom again").is_some());
10056 let note2 = f.reset_on_resume().expect("greedy-only reset");
10057 assert!(note2.contains("(greedy)"), "single-flag note: {note2}");
10058 }
10059
10060 /// Shape-change clears (dmask realloc / mask-shape mismatch / s_key change) stay silent —
10061 /// they precede a fresh capture attempt whose own failure re-flips loudly.
10062 #[test]
10063 fn shape_change_clears_are_silent() {
10064 let mut f = DraftGraphFallback::default();
10065 f.mark_greedy("oom").unwrap();
10066 f.clear_greedy();
10067 assert!(!f.greedy_failed());
10068 f.mark_sampled("oom").unwrap();
10069 f.clear_sampled();
10070 assert!(!f.sampled_failed());
10071 // after a silent clear there is nothing left for resume to report.
10072 assert!(f.reset_on_resume().is_none());
10073 }
10074}