memra_engine/spec.rs
1//! Qwen3.5 MTP (NextN) greedy speculative decode (research/mtp/MTP-PLAN.md §A/§B/§C/§D).
2//!
3//! Greedy spec decode is MATHEMATICALLY EXACT: the accepted+bonus token stream is token-for-token
4//! identical to plain greedy `generate`. This module provides:
5//! - `mtp_head_forward` (§A, T=1): one NextN draft-token forward.
6//! - `decode_step_t` (§D.3, T=K+1): batched target verify forward, all-column logits.
7//! - `generate_spec` (§B): the draft/verify/accept/rollback orchestrator.
8//! Cache snapshot/rollback lives in cache.rs (§D.4). The MTP head uses its OWN scratch KV (§D.6),
9//! PERSISTENT over the committed sequence (see `MtpScratch`).
10
11use crate::Engine;
12use crate::cache::{Cache, KvLayer};
13use crate::forward::argmax;
14use crate::hybrid::{FullAttnLayer, HybridModel, LinearAttnLayer, Mixer, MtpHead};
15use cudarc::driver::CudaSlice;
16use std::sync::atomic::{AtomicU64, Ordering};
17
18/// Parse the documented `MEMRA_SPEC_REPLAY=1` rollback seam.
19///
20/// Keep this shared with serving admission so `=0` cannot select replay in one
21/// layer while another layer treats it as disabled.
22pub fn spec_replay_env_on(value: Option<&str>) -> bool {
23 value == Some("1")
24}
25
26pub fn spec_replay_env_enabled() -> bool {
27 let value = std::env::var("MEMRA_SPEC_REPLAY").ok();
28 spec_replay_env_on(value.as_deref())
29}
30
31/// One compact, anchor-bounded DSpark supervision record. `tokens[0]` is the anchor at p and
32/// `hidden` is its predecessor carrier h[p-1], matching the live NextN/DSpark pairing. Target
33/// rows p..p+gamma-1 score tokens p+1..p+gamma. They are the full-target softmax's top-k
34/// entries; `target_tail_probs[j]` is the probability mass outside those rows. All flattened
35/// target arrays are `[gamma, top_k]` in row-major order.
36pub struct DsparkAnchorRecord {
37 pub position: usize,
38 pub hidden: Vec<f32>,
39 pub tokens: Vec<u32>,
40 pub target_top_ids: Vec<u32>,
41 pub target_top_logits: Vec<f32>,
42 pub target_top_probs: Vec<f32>,
43 pub target_tail_probs: Vec<f32>,
44}
45
46fn dspark_sparse_softmax_topk(
47 logits: &[f32],
48 top_k: usize,
49 temperature: f32,
50) -> Result<(Vec<u32>, Vec<f32>, Vec<f32>, f32), Box<dyn std::error::Error>> {
51 if logits.is_empty() || top_k == 0 || top_k > logits.len() || temperature <= 0.0 {
52 return Err("invalid DSpark sparse-softmax shape or temperature".into());
53 }
54 if logits.iter().any(|value| !value.is_finite()) {
55 return Err("DSpark target logits contain a non-finite value".into());
56 }
57 let mut ranked: Vec<(u32, f32)> = logits
58 .iter()
59 .copied()
60 .enumerate()
61 .map(|(index, value)| (index as u32, value))
62 .collect();
63 let compare = |left: &(u32, f32), right: &(u32, f32)| {
64 right.1.total_cmp(&left.1).then(left.0.cmp(&right.0))
65 };
66 ranked.select_nth_unstable_by(top_k - 1, compare);
67 ranked[..top_k].sort_unstable_by(compare);
68
69 let max_logit = logits.iter().copied().fold(f32::NEG_INFINITY, f32::max);
70 let inv_temperature = 1.0f64 / temperature as f64;
71 let denominator: f64 = logits
72 .iter()
73 .map(|value| (((*value - max_logit) as f64) * inv_temperature).exp())
74 .sum();
75 let ids: Vec<u32> = ranked[..top_k].iter().map(|(index, _)| *index).collect();
76 let top_logits: Vec<f32> = ranked[..top_k].iter().map(|(_, value)| *value).collect();
77 let top_probs: Vec<f32> = top_logits
78 .iter()
79 .map(|value| ((((value - max_logit) as f64) * inv_temperature).exp() / denominator) as f32)
80 .collect();
81 let top_mass: f64 = top_probs.iter().map(|value| *value as f64).sum();
82 let tail = (1.0f64 - top_mass).clamp(0.0, 1.0) as f32;
83 Ok((ids, top_logits, top_probs, tail))
84}
85
86fn flatten_dspark_rows<T>(
87 rows: Vec<Option<Vec<T>>>,
88 position: usize,
89 label: &str,
90) -> Result<Vec<T>, Box<dyn std::error::Error>> {
91 let mut flattened = Vec::new();
92 for (slot, row) in rows.into_iter().enumerate() {
93 flattened.extend(
94 row.ok_or_else(|| format!("missing DSpark {label} at {position} slot {slot}"))?,
95 );
96 }
97 Ok(flattened)
98}
99
100/// H-SEED CONVENTION (MEMRA_SPEC_HPOST=1): feed the MTP head the POST-norm hidden — trunk rows
101/// hand over `output_norm(x)` and the draft chain recurrence hands over `shared_head_norm(h_nextn)`
102/// (= final_h) — matching the reference engines: llama.cpp #24025 ("qwen35: use post-norm hidden
103/// state for MTP", t_h_nextn is taken AFTER the final norm in both trunk and MTP graphs) and
104/// SGLang's qwen3_5_mtp (spec_info.hidden_states = the target model's post-norm output). memra's
105/// historical convention (default, MTP-PLAN §A) is PRE-norm x. Draft-quality-only: exactness is
106/// the verify's job either way; acceptance arbitrates. OnceLock: read once, hot-loop safe.
107pub(crate) fn spec_hpost() -> bool {
108 static H: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
109 *H.get_or_init(|| {
110 std::env::var("MEMRA_SPEC_HPOST")
111 .map(|v| v != "0")
112 .unwrap_or(false)
113 })
114}
115
116/// LEAN VERIFY (default ON since 2026-07-08; MEMRA_SPEC_LEAN=0 reverts — close35 lane): the verify m-scaling
117/// probe + nsys diff showed the verify t-path pays ~1.0ms/call at m=1 over eager decode on the
118/// 35B, and the kernels are NOT the cause (dev-MoE identical, kernel-time delta only +179us).
119/// The overhead is (a) ~250 extra cuMemsetD8Async/call from `e.zeros()` on buffers every kernel
120/// fully overwrites (~0.9ms host issue + ~0.35ms GPU) and (b) the t=1 FA rows dispatch (rows_v2 +
121/// combine_rows, +50us vs the eager fa_decode pair). This flag switches (a) fully-overwritten
122/// verify buffers to `e.uninit` (identical bytes: every element is written before read) and
123/// (b) t==1 verify FA to the eager `fa_decode` entry (byte-identical: kernel-check pins the
124/// rows-vs-loop identity and the per-row loop at t=1 IS fa_decode on the same q). Gates arbitrate.
125pub(crate) fn spec_lean() -> bool {
126 static L: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
127 // DEFAULT ON since 2026-07-08 (MEMRA_SPEC_LEAN=0 reverts): bit-identical (buffers fully
128 // overwritten; gates green incl maxdiff-identical run-gen) and measured +2.4% e2e p3 /
129 // +1.5% p2 at the daily 35B config. m=1 verify now costs eager-decode parity.
130 *L.get_or_init(|| {
131 std::env::var("MEMRA_SPEC_LEAN")
132 .map(|v| v != "0")
133 .unwrap_or(true)
134 })
135}
136
137/// SMALL-M BATCHED VERIFY (default ON since 2026-07-09; MEMRA_SPEC_M2=0 reverts — lane/spec-m2): extend the
138/// batched linear-attn verify arm down to t=2 and batch the MoE dev token loop over a
139/// grid.z=token axis at every verify t. The close35 m-scaling probe put the m=2 verify tier at
140/// x1.54 of m=1 (llama x1.14); the per-column linear chain (t<3) and the serial MoE dev token
141/// loop are the two launch-structure causes. Both changes are LAUNCH-STRUCTURE ONLY:
142/// (a) the batched conv's t<pad ring update is pure copies (ssm_conv_ring_rebuild from a cloned
143/// ring — the ring stores raw input columns); every arithmetic kernel is the same one the
144/// t>=3 arm already runs (matmul_decode_exact bit-identical at m=2-4, gdn_scan's internal
145/// t-loop == chained T=1 steps);
146/// (b) the MoE dev-rows twins run the serial loop's per-token warp program with tok-offset
147/// pointers (same sel/w/aq/ad bytes, same dot order, same slot-ordered FMA chain).
148/// Gates arbitrate: run-spec K=1..8 self-consistency (35B+9B), kernel-check, run-gen argmax.
149pub(crate) fn spec_m2() -> bool {
150 static M: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
151 // DEFAULT ON since 2026-07-09 (MEMRA_SPEC_M2=0 reverts): launch-structure only — t=2
152 // batched linear arm (ring-roll copies, zero new FP order) + MoE dev-rows kernels
153 // (grid.z=token, 4 launches/layer at any verify t). Acceptance bit-identical at every K;
154 // 35B p2 +3.4% / p3 +3.6%; the profitable-K plateau widens (new optimum K=3 at 223).
155 *M.get_or_init(|| {
156 std::env::var("MEMRA_SPEC_M2")
157 .map(|v| v != "0")
158 .unwrap_or(true)
159 })
160}
161pub(crate) fn spec_stream() -> bool {
162 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
163 *ON.get_or_init(|| std::env::var("MEMRA_SPEC_STREAM").as_deref() == Ok("1"))
164}
165pub(crate) fn spec_stream_m() -> usize {
166 static M: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
167 *M.get_or_init(|| {
168 std::env::var("MEMRA_SPEC_STREAM_M")
169 .ok()
170 .and_then(|v| v.parse().ok())
171 .unwrap_or(4)
172 })
173}
174pub(crate) fn spec_devacc() -> bool {
175 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
176 *ON.get_or_init(|| std::env::var("MEMRA_SPEC_DEVACC").as_deref() == Ok("1"))
177}
178/// Engine-bundle slice 2 (DSF-ROUNDCOST-20260820 §1.1 host/device round trips + §2 rows 2-3),
179/// DEFAULT ON (`MEMRA_DSPARK_DEFER_READBACK=0` reverts): the dspark round's draft-chain DtoH
180/// is DEFERRED past verify dispatch and merged with the verify-argmax readback into ONE host
181/// sync (2 blocking DtoH/round -> 1). Verify embeds DEVICE tokens (`chain_d`) through the
182/// resident embed table — `embed_gather_u32_t`, bit-identical rows to the host gather by its
183/// own pinned contract. The host therefore dispatches snap + the whole verify while the DRAFT
184/// is still executing, instead of blocking ~1.7 ms on the chain and letting the device drain.
185/// Ladder arm only: the confidence policies size vt from a pre-verify head readback (their
186/// chain readback merges into that same sync instead). Exactness unchanged BY CONSTRUCTION —
187/// same tokens, same kernels, same order; E2E + accept-bank gates arbitrate.
188pub(crate) fn dspark_defer_readback_on() -> bool {
189 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
190 *ON.get_or_init(|| {
191 std::env::var("MEMRA_DSPARK_DEFER_READBACK")
192 .map(|v| v != "0")
193 .unwrap_or(true)
194 })
195}
196/// Engine-bundle slice 1 (DSF-ROUNDCOST-20260820 §1.1, lane/dspark-engine-bundle-20260820),
197/// DEFAULT ON (`MEMRA_STATE_COPY_BATCH=0` reverts): batch the dspark round's GDN state
198/// snapshot and partial-accept restore into single `copy_batch_uniform_f32` launches
199/// instead of ~2 memcpy dispatches (+2 alloc_zeros on the snap side) per linear layer per
200/// round — measured 0.67 ms/round snap + 0.25 ms/round commit of pure dispatch on the q38
201/// route. Launch-structure only: bytes, buffers and stream order are unchanged, so
202/// acceptance and streams stay bit-identical (E2E-gated on the B1 packs).
203pub(crate) fn state_copy_batch_on() -> bool {
204 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
205 *ON.get_or_init(|| {
206 std::env::var("MEMRA_STATE_COPY_BATCH")
207 .map(|v| v != "0")
208 .unwrap_or(true)
209 })
210}
211/// Engine-bundle slice 3 + fa-execupdate slice 4c (DSF-ROUNDCOST-20260820 §5 rank 1),
212/// DEFAULT OFF — `MEMRA_DSPARK_VERIFY_GRAPH=1` opts in: per-(segment, vt) CUDA graphs
213/// for the LINEAR-layer runs, plus the full-verify single graph per (vt, rung) when a
214/// round's rows all ride one seqs rung — see [`DsparkVerifyGraphs`]. Requires the
215/// slice-2 deferred path (device tokens); the eager walk is the byte-identical fallback.
216///
217/// MEASURED disposition (box6 card0, agentic pack, 2026-08-20, both slices): exactness
218/// holds everywhere (ALL EXACT, accept lines byte-match the banks, ckpt-gate oracle
219/// green over the graph + slab-commit paths). Slice-3's AUTO_FREE launch-scan limiter
220/// (25.6 us x 16 launches ≈ 0.41 ms/round) is FIXED — the captured bodies' alloc nodes
221/// are balanced by in-graph frees (census 84/84 per segment, 1776/1776 full) so graphs
222/// instantiate USE_NODE_PRIORITY and the scan is gone. What remains at gate scale:
223/// segment graphs +0.1 tok/s over the batched-rows default (114.4 vs 114.3 x5
224/// interleaved — the linear launch overhead was only ~0.1 ms); the FULL-verify graph is
225/// NET NEGATIVE at gate scale (110.6 vs 114.2: ~14-21 (vt, rung) captures/process at
226/// 2 full-walk executions + ~2.9k-node instantiate each eat far more than the ~0.2-0.3
227/// ms/round of remaining launch overhead). The orchestration ceiling of §1.3 is spent —
228/// the fa/append recovery landed DEFAULT-ON as the batched rows arm
229/// (`dspark_fa_rows_on`), not as a graph. The serve-lifetime cell (DSF-ROUNDCOST §9,
230/// nj-ws-solo) measured the amortization: crossover K≈33 requests, steady −0.246
231/// ms/round, −1.25% session wall over 240 requests — and the graphs-serve lane wired
232/// the door into the session arm (`dspark_spec_session_burst`) as a model-owned
233/// capture pool shared across sessions. Stays opt-in pending the owner's default-ON
234/// ratification on the serve-surface battery.
235pub(crate) fn dspark_verify_graph_on() -> bool {
236 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
237 *ON.get_or_init(|| std::env::var("MEMRA_DSPARK_VERIFY_GRAPH").as_deref() == Ok("1"))
238}
239/// MTP-ROUTE verify graphs, DEFAULT ON for the GDN+MoE family since 2026-08-23
240/// (`MEMRA_SPEC_VERIFY_GRAPH=0` is the kill switch, `=1` opts other families in).
241///
242/// The slice-4c capture already lived inside `qwen35_verify_tparallel` and said so in its own
243/// comment — "stream rides the qwen35moe burst, graphs ride the dspark route" — with no caller
244/// on this route. The MTP spec round is that caller.
245///
246/// WHY it is worth a default (receipts: `research/orndecode-20260822/VGRAPH.md`). With
247/// `MEMRA_SPEC_PHASE=1` this route's round reads verify-ISSUE 44-58% and verify-WAIT **0.0%**:
248/// the host is never waiting for the device, it is spending its own time launching the trunk.
249/// Replay collapses that into one graph launch and the phase all but disappears (55-62 ms ->
250/// 8-10 ms per burst).
251///
252/// MEASURED, two host generations, forced ON/OFF, balanced 4+4 boots in both orders:
253/// * current-generation host (9950X, the serving class): OFF 266.0-266.5, ON 318.8-319.5
254/// tok/s — **+19.7%**, no overlap, sub-1% spread per arm; per-round 6.9 -> 5.7 ms.
255/// * Zen 3 host: +3-9% (that rig's own clock drift is wider than the effect, so the ratio
256/// comes from per-round phase totals, which are internal to each boot).
257/// The ON arm lands at ~320 tok/s on BOTH hosts while OFF tracks host speed — the arm moves
258/// the round off the host and onto the device, which is the whole point.
259///
260/// EXACTNESS is structural (same kernels, same order) and gated anyway: a fixed-seed SAMPLED
261/// completion hashes identically ON vs OFF **and across both hosts** (`08941d5bb9762b21`),
262/// greedy seed-pinned likewise, `run-spec` K=1..8 PASS on both arms with identical acceptance
263/// at every K, kernel-check ALL GREEN.
264///
265/// SCOPE, deliberately narrow: default ON only where it was measured — the GatedDeltaNet +
266/// MoE family (`vgraph_family_default`). Qwen3.8-27B is GDN + DENSE mlp and would otherwise
267/// inherit this default unmeasured, which is the family-by-family law this repo keeps; it can
268/// opt in with `=1` once it has its own interleave. Also never armed together with
269/// ROUND-STREAM, and a round wider than the pool declines it for the eager walk.
270pub(crate) fn spec_verify_graph_env() -> Option<bool> {
271 static ON: std::sync::OnceLock<Option<bool>> = std::sync::OnceLock::new();
272 *ON.get_or_init(
273 || match std::env::var("MEMRA_SPEC_VERIFY_GRAPH").as_deref() {
274 Ok("1") => Some(true),
275 Ok("0") => Some(false),
276 _ => None,
277 },
278 )
279}
280/// SERVE-ROUTE twin of [`dspark_verify_graph_on`], DEFAULT ON — owner-ratified
281/// 2026-08-22 on the §10 serve-lifetime battery (DSF-ROUNDCOST-20260820 §10.3:
282/// crossover K=36–43, steady −0.357 ms/round, session wall −1.55..−1.65%, byte-exact
283/// 240/240 ×3 pairs, pool bounded at 8,852 MiB under `MEMRA_DSPARK_VG_MAX`). The env
284/// stays as the kill-switch: `MEMRA_DSPARK_VERIFY_GRAPH=0` restores the eager walk
285/// (byte-identical body); `MEMRA_DSPARK_VG_MAX=0` is the finer freeze valve. The BIN
286/// arm keeps its own opt-in default (`dspark_verify_graph_on`): at gate scale the
287/// capture toll is never repaid (§8 measured disposition — 14–21 captures over a
288/// 256-token run vs the serve session's thousands of rounds), and the two
289/// instruments must keep their own measured dispositions rather than share one flag.
290pub(crate) fn dspark_verify_graph_serve_on() -> bool {
291 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
292 *ON.get_or_init(|| std::env::var("MEMRA_DSPARK_VERIFY_GRAPH").as_deref() != Ok("0"))
293}
294/// Capture-count ceiling for the dspark verify-graph pool (graphs-serve lane) — the
295/// pool's memory policy STATED instead of silently unbounded. The keyspace is
296/// intrinsically finite — segment keys (run_start, vt) ≤ 16 runs x 7 windows, full
297/// keys (vt, rung, hi) ≤ 7 windows x the split-rung ladder (8 rungs at 32k ctx), ~168
298/// on the q38 export — so the default (256) never engages there; the knob is the
299/// safety valve for a future export with a wider ladder. At the ceiling the pool
300/// FREEZES: existing keys keep replaying, rounds needing a new capture run the eager
301/// walk byte-identically (round-atomic — a partial refusal would mix slab- and
302/// cols-stashed layers inside one commit). No eviction by design: destroying a live
303/// exec graph re-opens the stale-address class the indirect tables exist to close,
304/// and the bounded keyspace makes reclaim worthless.
305pub(crate) fn dspark_vg_cap() -> usize {
306 static CAP: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
307 *CAP.get_or_init(|| {
308 std::env::var("MEMRA_DSPARK_VG_MAX")
309 .ok()
310 .and_then(|v| v.parse().ok())
311 .unwrap_or(256)
312 })
313}
314/// Engine-bundle slice 4 (fa-execupdate lane, DSF-ROUNDCOST-20260820 §6 close: "the
315/// residual gap lives in the FULL-ATTENTION per-row section"), DEFAULT ON —
316/// `MEMRA_DSPARK_FA_ROWS=0` reverts to the per-row loop: when every row of a verify
317/// round takes the v4-seqs arm on ONE `fa_split_keys` rung (the straddle law, evaluated
318/// at the round's first and last t_kv — both eligibility gates are intervals in t_kv),
319/// the qwen35 t-parallel verify's per-row KV-append + fa-decode loop collapses into the
320/// z-batched serving twins: ONE `append_quantize_kv_q8_0_q5_1_seqs` + ONE
321/// `fa_decode_vec_q_seqs_v4` + ONE combine per full-attention layer, replacing
322/// T x (4 dtod row copies + append + 3 memsets + main + combine) launches. Bytes are
323/// pinned by the batched-tick increment-2 kernel-check (seqs-vs-per-seq-loop bit
324/// identity: per-row T_kv derives in-kernel from pos_seq[z]; splits >= ns_eff write the
325/// empty partial the combine never reads, so the shared n_splits_max stride changes no
326/// bytes) and re-gated e2e by this lane's battery.
327pub(crate) fn dspark_fa_rows_on() -> bool {
328 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
329 *ON.get_or_init(|| {
330 std::env::var("MEMRA_DSPARK_FA_ROWS")
331 .map(|v| v != "0")
332 .unwrap_or(true)
333 })
334}
335
336/// `t_pred0` for the `MEMRA_DEBUG_SPEC` per-round print, sampled-safe.
337///
338/// `generate_spec_inner2` fills its `preds` vector ONLY on the greedy path (`if !sampled`), and
339/// the per-round debug print was the sole consumer in the sampled arm: `t_pred(0)` survives round
340/// 0 (`base == 0` returns `last_pred`) and from round 1 (`base == 1`, a pending bonus) indexes an
341/// EMPTY vector — `index out of bounds: the len is 0 but the index is 0`, in the GPU worker
342/// thread, which then respawns and reloads weights while the request dies. So any sampled spec
343/// request longer than one round used to kill the worker whenever `MEMRA_DEBUG_SPEC` was set:
344/// the flag crashed precisely the regime it exists to investigate.
345///
346/// Fixed at the print site, not inside the closure, so the greedy accept walk keeps its strict
347/// indexing (an out-of-range pred there is a real bug and must still be loud).
348fn debug_t_pred0(sampled: bool, base: usize, last_pred: u32, preds: &[u32]) -> String {
349 if base == 0 {
350 return last_pred.to_string();
351 }
352 match preds.get(base - 1) {
353 Some(p) => p.to_string(),
354 // sampled: the greedy per-column argmax was never run for this round.
355 None => {
356 debug_assert!(
357 sampled,
358 "greedy spec: preds[{}] missing at base {base}",
359 base - 1
360 );
361 "n/a".to_string()
362 }
363 }
364}
365
366/// `MEMRA_SKEY_PROBE=1` — sampled-draft-graph key probe (lane/graph-s-key-exactness-20260819).
367///
368/// Reports, per burst and per round, which draft chain the sampled arm chose and under which
369/// filter regime, plus the ONE observable that separates a legal filtered draft from a stale
370/// pure-temp graph replayed under filters: an accept test whose gathered `q` is exactly 0.
371/// A draft token sampled from the FILTERED softmax can never gather q=0 (it was drawn from the
372/// kept set), so `q=0` in the verify means the draft came from a distribution the verify does
373/// not believe in — and `u * 0 < p` then accepts it unconditionally.
374///
375/// Its own env var, deliberately NOT `MEMRA_DEBUG_SPEC`: that flag panicked the GPU worker on
376/// any sampled spec request past round 0 until this lane fixed it (§2 of the bank note).
377pub(crate) fn skey_probe() -> bool {
378 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
379 *ON.get_or_init(|| std::env::var("MEMRA_SKEY_PROBE").as_deref() == Ok("1"))
380}
381
382/// GRAMMAR HOOK for constrained spec decode (lane/constrained-full, 2026-08-03). The engine
383/// stays llguidance-agnostic: the server adapts its per-session grammar state behind this
384/// trait. CONTRACT (the verify-side truncation rule — token-identical to constrained plain
385/// greedy decode): the exactness walk runs UNMASKED first; the hook then (a) truncates
386/// acceptance at the first grammar-illegal accepted token, and (b) when the truncation fired
387/// or the bonus is illegal, the engine recomputes that slot as the MASKED argmax of the
388/// target's own verify column (an unmasked argmax that is grammar-legal IS the masked argmax
389/// — masking only removes tokens — so the common case pays nothing). `consume` advances the
390/// state with each EMITTED token in order; EOS handling is the implementor's job (skip).
391pub trait SpecConstraint {
392 /// -inf the current state's banned ids on a HOST logits row (prompt-tail / init-feed
393 /// masked argmax).
394 fn mask_logits(&mut self, logits: &mut [f32]) -> Result<(), String>;
395 /// Packed 32-bit bitset words of the CURRENT state's allowed set (device-mask form).
396 fn mask_words(&mut self) -> Result<Vec<u32>, String>;
397 /// Is `tok` consumable in the CURRENT state?
398 fn is_allowed(&mut self, tok: u32) -> Result<bool, String>;
399 /// Advance the state with an emitted token.
400 fn consume(&mut self, tok: u32) -> Result<(), String>;
401
402 // --- DRAFT-SIDE MASKING (lane/draft-mask, 2026-08-04) ---
403 // The drafter proposed grammar-illegal tokens under tight schemas, so verify-side
404 // truncation cut nearly every round (measured acceptance 0.467-0.513 tight vs 0.62-0.82
405 // loose, research/constrained-full-20260803). These three methods let the engine mask the
406 // DRAFT model's own sampling with the grammar's legal set, so proposals are legal by
407 // construction. The state they walk is a SPECULATIVE CLONE of the session matcher — the
408 // real state is advanced only by `consume` (emitted tokens), so verify-side truncation
409 // stays the correctness backstop and the emitted stream is unchanged by construction
410 // (an accepted draft is the target's unmasked argmax AND grammar-legal, hence the masked
411 // argmax; a cut slot is recomputed as the masked argmax either way).
412 // Default impls = feature OFF (pre-lane behaviour: unmasked drafts).
413
414 /// Is draft-side masking available on this hook? Probed ONCE per burst, before the draft
415 /// graph is captured (the mask is an in-graph node — its presence is a capture-time shape).
416 fn draft_mask_enabled(&self) -> bool {
417 false
418 }
419 /// Start a draft chain: clone the CURRENT (committed) grammar state into the speculative
420 /// slot. Called once per spec round, before the first draft position.
421 fn draft_begin(&mut self) -> Result<(), String> {
422 Ok(())
423 }
424 /// Packed 32-bit bitset words of the SPECULATIVE state's allowed set (target-vocab ids),
425 /// for the draft position about to be sampled. `None` = draft masking off (no-op).
426 fn draft_mask_words(&mut self) -> Result<Option<Vec<u32>>, String> {
427 Ok(None)
428 }
429 /// Advance the SPECULATIVE state with a PROPOSED draft token. `false` = the chain cannot
430 /// continue (EOS proposed, or an unmasked position proposed something illegal) — the
431 /// engine stops drafting; the token already pushed still goes through verify.
432 fn draft_advance(&mut self, _tok: u32) -> Result<bool, String> {
433 Ok(false)
434 }
435}
436
437/// DRAFT-MASK UPLOAD (lane/draft-mask): pull the speculative state's allowed set (TARGET-id
438/// space) from the hook, project it into the DRAFT head's vocab space, and upload it into the
439/// stable device buffer the draft chain reads. Returns false when the chain must stop drafting:
440/// the hook handed out no mask, or NO draft-vocab row is grammar-legal at this position (a
441/// trimmed FR-Spec head genuinely cannot propose a legal token there — masking it would leave
442/// a fully-banned row whose argmax is meaningless, so the round drafts fewer tokens and the
443/// verify emits the masked argmax as usual).
444fn upload_draft_mask(
445 e: &Engine,
446 c: &mut dyn SpecConstraint,
447 dst: &mut CudaSlice<u32>,
448 d2t: Option<&Vec<u32>>,
449 d_vocab: usize,
450 words: usize,
451) -> Result<bool, Box<dyn std::error::Error>> {
452 let Some(tw) = c
453 .draft_mask_words()
454 .map_err(|e2| format!("constraint: {e2}"))?
455 else {
456 return Ok(false);
457 };
458 let bit = |t: usize| -> bool {
459 let w = t >> 5;
460 w < tw.len() && (tw[w] >> (t & 31)) & 1 == 1
461 };
462 let mut buf = vec![0u32; words];
463 match d2t {
464 // TRIMMED draft head: row i proposes target id d2t[i] — permute the mask accordingly.
465 Some(map) => {
466 for (i, &t) in map.iter().enumerate().take(d_vocab) {
467 if bit(t as usize) {
468 buf[i >> 5] |= 1u32 << (i & 31);
469 }
470 }
471 }
472 // UNTRIMMED: draft ids ARE target ids; the packed words transfer verbatim (a short
473 // mask leaves the padded tail zeroed == banned, same rule as constrained::apply_mask).
474 None => {
475 let n = tw.len().min(words);
476 buf[..n].copy_from_slice(&tw[..n]);
477 }
478 }
479 if buf.iter().all(|w| *w == 0) {
480 return Ok(false);
481 }
482 e.htod_u32_into(dst, &buf)?;
483 Ok(true)
484}
485
486/// Keep the full token-embedding table in host memory and upload only the rows needed by each
487/// MTP/verify step. This is an exact memory-capacity seam for very large BF16 vocab tables: host
488/// gather expands the same source bits to f32, and only O(T*n_embd) bytes cross PCIe per step.
489/// CUDA-graph/round-stream draft paths require device token ids and therefore stay disabled.
490pub(crate) fn spec_host_embd() -> bool {
491 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
492 *ON.get_or_init(|| std::env::var("MEMRA_SPEC_HOST_EMBD").as_deref() == Ok("1"))
493}
494
495/// VERIFY-TIER TRUNK LAUNCH-FUSION (default ON since 2026-07-09; MEMRA_SPEC_FUSED_T=0 reverts — lane/close35b): extend
496/// the t=1 fused2/fused3 Q8_0 trunk launches to the batched verify tier (t=2-4, the K=1..3
497/// verify shapes). At t>1 the trunk pairs/triples (35B wqkv+wqkv_gate, wq/wk/wv,
498/// gate_shexp+up_shexp) each run a separate `matmul_decode_exact` — one q8_1 re-quantize of the
499/// SAME activation plus one _b2/_b4 launch per tensor. The fused twins share ONE quantize and
500/// ONE launch per group; per (tensor,token,row) the kernel body is q8_0_mmvq_batched verbatim
501/// with the identical row mapping -> BIT-IDENTICAL by construction (kernel-check pins it,
502/// run-spec K=1..8 + acceptance identity arbitrate e2e).
503pub(crate) fn spec_fused_t() -> bool {
504 static F: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
505 // DEFAULT ON since 2026-07-09 (MEMRA_SPEC_FUSED_T=0 reverts): verify t=2-4 trunk launch-fusion
506 // (fused2/fused3 Q8_0 batched twins, bit-identical by construction — m=1 block-offset split on
507 // the batched body). m=2 marginal token 2117->1762us; 35B daily: p3 +3.7% (crosses llama), p2 +5%.
508 *F.get_or_init(|| {
509 std::env::var("MEMRA_SPEC_FUSED_T")
510 .map(|v| v != "0")
511 .unwrap_or(true)
512 })
513}
514
515/// zeros/uninit switch for verify-path buffers that are FULLY OVERWRITTEN before any read.
516/// Only call this on such buffers — the lean contract is "identical bytes by construction".
517fn vbuf(e: &Engine, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
518 if spec_lean() { e.uninit(n) } else { e.zeros(n) }
519}
520
521/// Scratch KV for the MTP block (one full-attn layer).
522///
523/// PERSISTENT MODE (default, 2026-07-03 — the acceptance lever): sized cap = max_ctx and kept in
524/// sync with the COMMITTED sequence — slot p holds the MTP block's K/V for committed token p
525/// (roped p+1, the chain's rope convention), so the draft chain's self-attention sees the FULL
526/// committed history instead of only the current round's 1..K+1 chain tokens (the reference
527/// engine's "mtp_update" design). Entries come from two sources:
528/// - chain appends: accepted positions KEEP their chain-computed entries (embedding exact,
529/// hidden chain-approximate — the reference engine accepts the same);
530/// - `mtp_kv_fill` batches: prompt positions + the last-draft position on full accept, computed
531/// from EXACT trunk hiddens (K/V-only MTP-block pass, no attention/FFN/lm_head).
532/// Rejected drafts / p-min extras / pseudo-seed appends are all discarded by the round-start
533/// `set_len` truncation (the KvLayer len mechanism — §C rollback for the draft side).
534/// Multi-turn spec-decode session (2026-07-05): trunk Cache + persistent MTP draft scratch +
535/// the committed token list, alive across generate_spec_session calls. Turn N+1 primes ONLY its
536/// suffix (chunked continuation prime over the quantized past) and mtp_kv_fill's its suffix rows,
537/// then the round loop runs unchanged. `last_h` carries the pre-output_norm hidden of the last
538/// committed row across turns (the predecessor-pairing seed + fill anchor).
539/// Per-request sampling config for the sampled-spec serve path.
540#[derive(Clone, Copy, Debug)]
541pub struct SpecSampling {
542 pub temp: f32,
543 pub seed: u64,
544 pub top_k: i32, // 0 = off
545 pub top_p: f32, // 1.0 = off
546 pub min_p: f32, // 0.0 = off
547 pub penalty_last_n: usize, // 0 = penalties off
548 pub penalty_repeat: f32,
549 pub penalty_freq: f32,
550 pub penalty_present: f32,
551}
552
553impl SpecSampling {
554 /// Non-identity penalties requested — THE `pen_on` predicate (one definition; the
555 /// same group-off rule `SamplerIdentity::of` canonicalizes: a window with neutral
556 /// coefficients is penalties-absent). Both spec routes and the dspark accept walk
557 /// key their penalty arms off this.
558 pub fn pen_on(&self) -> bool {
559 self.penalty_last_n > 0
560 && (self.penalty_repeat != 1.0
561 || self.penalty_freq != 0.0
562 || self.penalty_present != 0.0)
563 }
564}
565
566/// Host Philox4x32-10 uniform in (0,1) — mirrors spec_sample.cu's `philox4`/`u01` with the
567/// ctr_lo tag 0xFFFF_FFFE, so the host accept-test stream never collides with any device
568/// sampling event (device Gumbel uses (i>>2, stream_pos); device residual uses 0xFFFF_FFFD).
569/// One value per (seed, ctr) EVENT; callers own the counter discipline. Extracted verbatim
570/// from generate_spec_inner2's closure for the dspark sampled-admission walk (the two paths
571/// MUST consume the identical stream construction — two ad-hoc Philox copies drifting apart
572/// is a distributional bug, not a style problem).
573pub(crate) fn host_u01(seed: u64, ctr: u32) -> f32 {
574 let (m0, m1) = (0xD2511F53u32, 0xCD9E8D57u32);
575 let (mut c0, mut c1, mut c2, mut c3) = (0xFFFF_FFFEu32, ctr, 0u32, 0u32);
576 let (mut k0, mut k1) = ((seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
577 for _ in 0..10 {
578 let (h0, l0) = (((m0 as u64 * c0 as u64) >> 32) as u32, m0.wrapping_mul(c0));
579 let (h1, l1) = (((m1 as u64 * c2 as u64) >> 32) as u32, m1.wrapping_mul(c2));
580 let (n0, n1, n2, n3) = (h1 ^ c1 ^ k0, l1, h0 ^ c3 ^ k1, l0);
581 c0 = n0;
582 c1 = n1;
583 c2 = n2;
584 c3 = n3;
585 k0 = k0.wrapping_add(0x9E3779B9);
586 k1 = k1.wrapping_add(0xBB67AE85);
587 }
588 (c0 as f32 + 1.0) * (1.0 / 4294967296.0)
589}
590
591/// Tracked draft positions for [`SpecTelemetry`] (serve K defaults to 3; the run-spec gate
592/// sweeps K=1..8, and MEMRA_SPEC_CAPMAX defaults to 7 — 8 covers every tuned config).
593pub const SPEC_TELEM_POS: usize = 8;
594
595/// Always-on per-draft-position acceptance telemetry (lane/accept-telemetry, 2026-08-05 —
596/// the llama.cpp #26389 / vLLM spec-decode counter schema, upstream-sweeps 2026-08-05).
597/// Lives on the [`SpecSession`] and accumulates across bursts; the serve worker diffs a
598/// stashed copy per burst for its per-model /metrics aggregation and per-request usage.
599/// Same normalization as the `[spec-stats]` line: p-min-discarded chain tokens are counted
600/// in NEITHER drafted nor accepted.
601#[derive(Clone, Copy, Default, Debug)]
602pub struct SpecTelemetry {
603 /// verify rounds completed (a round-stream burst counts each of its M rounds).
604 pub rounds: u64,
605 /// tokens drafted / accepted across all rounds.
606 pub drafted: u64,
607 pub accepted: u64,
608 /// how often draft position j (0-based within a round's chain) was offered / accepted.
609 /// Positions >= SPEC_TELEM_POS are untracked (totals still count them). The opt-in
610 /// round-stream arm (MEMRA_SPEC_STREAM=1) reads back only totals, so under it these
611 /// arrays cover the standard-path rounds only and their sums may undercount the totals.
612 pub pos_drafted: [u64; SPEC_TELEM_POS],
613 pub pos_accepted: [u64; SPEC_TELEM_POS],
614}
615
616impl SpecTelemetry {
617 /// Fieldwise `self - prev` — the worker's per-burst delta off a copy stashed before the
618 /// burst call. Saturating: a caller diffing against the wrong snapshot gets zeros, not
619 /// a wrapped counter.
620 pub fn delta_since(&self, prev: &SpecTelemetry) -> SpecTelemetry {
621 let mut d = SpecTelemetry {
622 rounds: self.rounds.saturating_sub(prev.rounds),
623 drafted: self.drafted.saturating_sub(prev.drafted),
624 accepted: self.accepted.saturating_sub(prev.accepted),
625 ..Default::default()
626 };
627 for j in 0..SPEC_TELEM_POS {
628 d.pos_drafted[j] = self.pos_drafted[j].saturating_sub(prev.pos_drafted[j]);
629 d.pos_accepted[j] = self.pos_accepted[j].saturating_sub(prev.pos_accepted[j]);
630 }
631 d
632 }
633 /// Fieldwise `self += d` — the worker's per-model aggregation.
634 pub fn merge(&mut self, d: &SpecTelemetry) {
635 self.rounds += d.rounds;
636 self.drafted += d.drafted;
637 self.accepted += d.accepted;
638 for j in 0..SPEC_TELEM_POS {
639 self.pos_drafted[j] += d.pos_drafted[j];
640 self.pos_accepted[j] += d.pos_accepted[j];
641 }
642 }
643
644 /// Mean accepted draft-prefix length per verify round (tau).
645 pub fn tau(&self) -> f64 {
646 if self.rounds > 0 {
647 self.accepted as f64 / self.rounds as f64
648 } else {
649 0.0
650 }
651 }
652}
653
654/// Session-lifetime atomic acceptance counters. The verifier records only after the greedy or
655/// rejection-sampling walk has resolved on the host, so these relaxed increments add no GPU
656/// launch, synchronization, allocation, or ordering dependency to the numeric path.
657struct SpecTelemetryCounters {
658 rounds: AtomicU64,
659 drafted: AtomicU64,
660 accepted: AtomicU64,
661 pos_drafted: [AtomicU64; SPEC_TELEM_POS],
662 pos_accepted: [AtomicU64; SPEC_TELEM_POS],
663}
664
665impl Default for SpecTelemetryCounters {
666 fn default() -> Self {
667 Self {
668 rounds: AtomicU64::new(0),
669 drafted: AtomicU64::new(0),
670 accepted: AtomicU64::new(0),
671 pos_drafted: std::array::from_fn(|_| AtomicU64::new(0)),
672 pos_accepted: std::array::from_fn(|_| AtomicU64::new(0)),
673 }
674 }
675}
676
677impl SpecTelemetryCounters {
678 fn record_round(&self, drafted: usize, accepted: usize) {
679 debug_assert!(accepted <= drafted);
680 self.rounds.fetch_add(1, Ordering::Relaxed);
681 self.drafted.fetch_add(drafted as u64, Ordering::Relaxed);
682 self.accepted.fetch_add(accepted as u64, Ordering::Relaxed);
683 for counter in self.pos_drafted.iter().take(drafted) {
684 counter.fetch_add(1, Ordering::Relaxed);
685 }
686 for counter in self.pos_accepted.iter().take(accepted) {
687 counter.fetch_add(1, Ordering::Relaxed);
688 }
689 }
690
691 /// Round-stream keeps each round's accept length on device; retain exact scalar totals while
692 /// leaving the per-position arrays untouched, matching the pre-existing telemetry contract.
693 fn record_totals(&self, rounds: usize, drafted: usize, accepted: usize) {
694 self.rounds.fetch_add(rounds as u64, Ordering::Relaxed);
695 self.drafted.fetch_add(drafted as u64, Ordering::Relaxed);
696 self.accepted.fetch_add(accepted as u64, Ordering::Relaxed);
697 }
698
699 fn snapshot(&self) -> SpecTelemetry {
700 SpecTelemetry {
701 rounds: self.rounds.load(Ordering::Relaxed),
702 drafted: self.drafted.load(Ordering::Relaxed),
703 accepted: self.accepted.load(Ordering::Relaxed),
704 pos_drafted: std::array::from_fn(|j| self.pos_drafted[j].load(Ordering::Relaxed)),
705 pos_accepted: std::array::from_fn(|j| self.pos_accepted[j].load(Ordering::Relaxed)),
706 }
707 }
708}
709
710pub struct SpecSession {
711 pub(crate) cache: Cache,
712 pub(crate) scratch: MtpScratch,
713 /// Every token whose state the caches hold, in order (prompt turns + generated), INCLUDING
714 /// overshoot: spec commits accepted drafts past max_new; those rows are in the caches, so the
715 /// session must count them. Callers render output from this, not from their own echo.
716 pub committed: Vec<u32>,
717 /// Pre-output_norm hidden of the LAST committed row (device). None before the first turn.
718 pub(crate) last_h: Option<CudaSlice<f32>>,
719 /// Greedy argmax predicting the token AFTER committed.last() (from the last turn's final
720 /// logits). Fuels empty-suffix continuation bursts (serve): the next turn emits this token
721 /// first, feeds it, and the round loop resumes without any prime. None before the first turn.
722 pub next_pred: Option<u32>,
723 /// SAMPLED-SPEC stream continuity across bursts: Philox event counters persist here so a
724 /// session's randomness never repeats between generate_spec_session calls. (0,0) at admit.
725 pub sctr: u32,
726 pub uctr: u32,
727 /// PERSISTENT DRAFT-GRAPH CONTEXT (2026-08-01, the serve-burst fixed-cost fix): the captured
728 /// draft graph(s) + every device I/O buffer they bake, carried ACROSS generate_spec_session
729 /// calls. Before this, every serve burst re-captured the draft graph (2 warmup forwards +
730 /// instantiate) — measured ~16ms/burst on H100 q27 (MEMRA_SPEC_BURST sweep,
731 /// research/spec-serving-20260801). None before the first turn; error paths drop it
732 /// (next burst recaptures — serve retires errored sessions anyway).
733 pub(crate) draft_ctx: Option<DraftGraphCtx>,
734 /// PENDING-CARRY across bursts (2026-08-01, the serve burst-boundary fix): the bonus token
735 /// emitted by the last round but NOT committed to the caches. The old tail committed it with
736 /// a solo T=1 trunk pass (+ draft fill), and the next burst's setup fed the stashed next_pred
737 /// with ANOTHER solo pass — 2x ~11.5ms/burst measured on H100 q27 ([spec-setup] trace).
738 /// Carrying it lets the next empty-suffix greedy burst consume it as round-0 verify col 0,
739 /// exactly like a mid-burst full-accept boundary (no solo passes). INVARIANT: when set,
740 /// `committed` (== cache rows) EXCLUDES this token although it was already emitted in the
741 /// last burst's output, and `last_h` holds the hidden of the last COMMITTED row (its
742 /// predecessor — the chain-seed/fill anchor). `next_pred` is None (unknown without the
743 /// commit pass). Non-empty-suffix or sampled turns must flush first (spec_flush_pending);
744 /// generate_spec_session_sampled does this at entry, and serve parks only flushed sessions.
745 pub pending_tok: Option<u32>,
746 /// SESSION-AFFINITY TURN CHECKPOINT (lane/session-affinity, 2026-08-05): the state at this
747 /// turn's PROMPT-END boundary, retained so a later turn can REWIND here. See
748 /// [`SpecCheckpoint`]. Refreshed by every non-empty prime; None until the first one, and on
749 /// a rig too tight to hold it (a failed capture is silent — resume just isn't available).
750 pub(crate) turn_ckpt: Option<SpecCheckpoint>,
751 /// Session-lifetime acceptance telemetry. Relaxed atomics update at the host-side round
752 /// accounting the loop already does — no syncs, no allocation. NOTE a
753 /// pool-resumed session carries the PREVIOUS requests' counts; per-request consumers
754 /// diff with [`SpecTelemetry::delta_since`] around each burst.
755 telem: SpecTelemetryCounters,
756 /// PREFIX-CACHE publication request (lane/spec-prefix-cache): worker sets this to the
757 /// miss-LCP boundary before a cold burst; the prime captures at exactly that split (it must
758 /// coincide with the burst's `prime_split` or no capture happens). One-shot: consumed by the
759 /// prime, result lands in `boundary_captures`.
760 pub capture_at: Option<usize>,
761 /// The captures the last prime produced (see [`SpecBoundaryCapture`]). Worker drains them
762 /// post-burst to assemble prefix entries. A failed capture is silent, like `turn_ckpt` —
763 /// publication just isn't available for that request. Plural since
764 /// lane/frspec-multiturn-cache (2026-08-21): a cold burst can capture BOTH the miss-LCP
765 /// split (the shared-prefix class) and the stable pre-generation boundary (the
766 /// next-turn re-render class) — one entry per stop, exactly the boundary set the plain
767 /// prefill tick publishes/checkpoints.
768 pub boundary_captures: Vec<SpecBoundaryCapture>,
769 /// STABLE-BOUNDARY TURN CHECKPOINT REQUEST (lane/frspec-multiturn-cache, 2026-08-21): the
770 /// ABSOLUTE committed-length position the next non-empty prime should capture `turn_ckpt`
771 /// at, instead of prompt-end. The worker sets it to the STABLE PRE-GENERATION boundary
772 /// (`plain_checkpoint_boundary` — before the live generation header the client rewrites),
773 /// porting the 2026-08-09 plain-tier fix: a prompt-end spec checkpoint includes the
774 /// template's live assistant-generation header (`<|im_start|>assistant\n<think>\n`), which
775 /// the NEXT turn's re-render replaces, so `affinity_match` diverged a couple tokens below
776 /// the checkpoint and the spec pool declined 100% of multi-turn agent traffic (measured:
777 /// `spec-affinity: declined (history diverged at 6811 of checkpoint 6813)`,
778 /// research/multiturn-cache-20260821 B4). One-shot, `capture_at` convention; None = legacy
779 /// prompt-end capture.
780 pub ckpt_at: Option<usize>,
781}
782impl SpecSession {
783 /// Context capacity of the session's caches (the server's ContextFull guard).
784 pub fn cache_max_ctx(&self) -> usize {
785 self.cache.max_ctx
786 }
787 /// Read access to the live trunk cache (lane/spec-prefix-cache): the worker slices
788 /// full-attn KV rows `[0..capture.pos)` out of it when publishing a boundary capture —
789 /// those rows are append-only for the session's lifetime (rollbacks never truncate below
790 /// the prime boundary), so no copy was taken at prime time.
791 pub fn cache_ref(&self) -> &Cache {
792 &self.cache
793 }
794 /// Read access to the persistent draft-scratch plane (lane/spec-on-cache-hit): the
795 /// worker slices rows `[0..capture.pos)` when publishing a boundary capture, exactly
796 /// like the trunk KV — draft rows below the prompt end are append-only for the
797 /// session's lifetime (the prime fill wrote them once; rollbacks reset `len_d` to the
798 /// committed length, never below the prime boundary, and the true-hidden refresh
799 /// rewrites generated positions only). Returns `(k, v, k_tok_bytes, v_tok_bytes)`.
800 /// None when the scratch is ring-backed (Step35 SWA — physical rows are not
801 /// prefix-addressable; the prefix cache already refuses that class end to end).
802 pub fn draft_plane_ref(&self) -> Option<(&CudaSlice<u8>, &CudaSlice<u8>, usize, usize)> {
803 if self.scratch.kv.ring.is_some() {
804 return None;
805 }
806 Some((
807 &self.scratch.kv.k,
808 &self.scratch.kv.v,
809 self.scratch.kv.k_tok_bytes,
810 self.scratch.kv.v_tok_bytes,
811 ))
812 }
813 /// Snapshot the session's process-local acceptance counters for per-burst diffing.
814 pub fn telemetry(&self) -> SpecTelemetry {
815 self.telem.snapshot()
816 }
817 /// Committed position this session can REWIND to (its retained prompt-end boundary), if any.
818 /// A request whose prompt matches `committed[..pos]` exactly can resume from here — see
819 /// `spec_rewind_to_checkpoint`.
820 pub fn rewind_pos(&self) -> Option<usize> {
821 self.turn_ckpt.as_ref().map(|c| c.pos)
822 }
823 /// Whether every ring-backed trunk/draft row needed by the retained checkpoint is resident.
824 pub fn rewind_is_resident(&self) -> bool {
825 self.turn_ckpt.as_ref().is_some_and(|ckpt| {
826 self.cache.can_rollback(&ckpt.snap, 0) && self.scratch.can_rewind_to(ckpt.pos)
827 })
828 }
829 /// Is this session in the DEMOTION-READY shape (see [`SpecSession::into_demoted`])?
830 /// `false` means a carried pending must be flushed first (`spec_flush_pending`), or the
831 /// session has never run a turn and has no prediction to hand over.
832 pub fn demote_ready(&self) -> bool {
833 self.pending_tok.is_none() && self.next_pred.is_some()
834 }
835 /// Does this session hold a carried pending bonus (flush required before a handoff/park)?
836 pub fn has_pending(&self) -> bool {
837 self.pending_tok.is_some()
838 }
839 /// Committed row count == cache rows (the session invariant), for the caller's own
840 /// `fed`-length cross-check at a handoff boundary.
841 pub fn committed_len(&self) -> usize {
842 self.committed.len()
843 }
844 /// DEMOTION HANDOFF (lane/spec-gate, 2026-08-07): consume this session and hand its trunk
845 /// cache + next-token prediction to the plain batched-decode path.
846 ///
847 /// WHY THIS IS EXACT (greedy). The invariant at a burst boundary is `cache.pos ==
848 /// committed.len()`: every committed row has trunk KV + recurrent state, exactly as a plain
849 /// tokenwise prime of the same `committed` sequence would have left it (that is the
850 /// session-tail contract, and the same property `spec_rewind_to_checkpoint` and the reuse
851 /// pool already rely on). `next_pred` is the argmax of the verify's logits for the LAST
852 /// committed row — and verify-column logits are bit-identical to plain decode's logits at
853 /// that position, because `matmul_decode_exact` bit-identity IS the basis of the greedy
854 /// accept walk. So handing (cache, next_pred) to the batched path continues the stream from
855 /// a state indistinguishable from one the batched path produced itself: the batched tick
856 /// emits `next_pred`, feeds it into this same cache, and decodes on.
857 ///
858 /// `None` when the session is not in the handoff shape — a carried pending (its bonus row is
859 /// NOT in the cache, so `spec_flush_pending` must commit it first) or no `next_pred` yet
860 /// (never bursted). Callers must not force it: a half-committed cache handed to the batched
861 /// path would silently skip a token.
862 ///
863 /// The MTP draft scratch, the persistent draft-graph context and the turn checkpoint are
864 /// DROPPED here (freeing their VRAM): the batched path never drafts, and this handoff is
865 /// one-way by design — there is no cheap symmetric re-promotion (rebuilding the draft KV
866 /// would mean an `mtp_kv_fill` over the whole committed history).
867 pub fn into_demoted(self) -> Option<(Cache, u32)> {
868 if self.pending_tok.is_some() {
869 return None;
870 }
871 let np = self.next_pred?;
872 debug_assert_eq!(
873 self.cache.pos,
874 self.committed.len(),
875 "demotion handoff: cache rows != committed tokens"
876 );
877 Some((self.cache, np))
878 }
879 /// Pool-resume hook (audit Q2): clear the parked draft-graph failure memoization so a
880 /// NEW request resuming this session gets one fresh capture chance — a transient-pressure
881 /// capture failure must not persist for the pool's whole lifetime (the TRT #16072 class).
882 /// Logs once iff a flag was actually set; a no-fallback resume is silent and free.
883 pub fn reset_graph_fallback_on_resume(&mut self) {
884 if let Some(line) = self
885 .draft_ctx
886 .as_mut()
887 .and_then(|c| c.failed.reset_on_resume())
888 {
889 eprintln!("{line}");
890 }
891 }
892}
893
894/// A session's PROMPT-END boundary state, the rewind target for session-affinity resume.
895///
896/// WHY THIS BOUNDARY, AND WHY IT IS THE ONLY ONE WORTH KEEPING. The rewrite class this lane
897/// exists for (a client that strips `<think>` blocks out of prior assistant turns) mutates the
898/// text the session GENERATED, never the prompt it was given. So turn N's prompt agrees with
899/// turn N-1's committed tokens up to almost exactly where turn N-1's generation began — the
900/// prompt-end boundary. Keeping a checkpoint there means the next turn re-primes only its own
901/// delta (the rewritten answer + the new user turn) instead of the whole conversation.
902///
903/// WHAT IT MUST HOLD. Full-attn KV is append-only and position-addressed, so rewinding it is a
904/// `len` truncation (no data). Linear-attn (GDN) conv/ssm state is mutated IN PLACE with no
905/// position index, so it must be a real device COPY — that copy is the entire reason a spec
906/// session could not previously rewind. The MTP draft scratch needs no copy either: its rows
907/// below the boundary were written by this turn's fill and are never revisited (the per-round
908/// true-hidden refresh only rewrites the CURRENT burst's committed positions), so rewinding it
909/// is also just a `len` reset. `last_h` is the hidden of the last row below the boundary — the
910/// predecessor-pairing anchor the next prime's fill reads for its first row.
911///
912/// COST: one `Cache::snapshot` per TURN, on a code path that already takes one per ROUND.
913pub(crate) struct SpecCheckpoint {
914 snap: crate::cache::CacheSnapshot,
915 /// Committed length at the boundary (== cache.pos there, the session invariant).
916 pos: usize,
917 /// Pre-output_norm hidden of row `pos - 1`.
918 last_h: CudaSlice<f32>,
919}
920
921/// PREFIX-CACHE BOUNDARY CAPTURE (lane/spec-prefix-cache, 2026-08-14): the state a spec session
922/// records at its cold-prime split so the WORKER can publish a cross-request prefix entry —
923/// the commit-gated-publication port (research/cache-spec-design-20260814/PORT-PLAN.md item 1).
924/// Only the pieces that are DESTROYED by continuing the prime need copies here: the in-place
925/// GDN conv/ssm states (via `Cache::snapshot`, same mechanism as [`SpecCheckpoint`]) and the
926/// boundary logits. Full-attn KV rows `[0..pos)` and draft-scratch rows `[0..pos)` are
927/// append-only for the session's lifetime (rollbacks never truncate below the prime boundary),
928/// so the worker slices those from the live caches post-burst instead of copying at prime time.
929pub struct SpecBoundaryCapture {
930 pub snap: crate::cache::CacheSnapshot,
931 /// Token boundary (== cache.pos at capture; == the worker's miss-LCP split).
932 pub pos: usize,
933 /// Full-vocab logits after the prefix prime — the entry's boundary logits.
934 pub logits: Vec<f32>,
935 /// Pre-output_norm trunk hidden of row `pos - 1` (lane/spec-on-cache-hit): the
936 /// predecessor-pairing anchor a RESTORED spec session's first suffix-fill row reads
937 /// (the `SpecSession::last_h` convention). Empty = unavailable (capture stays valid;
938 /// the fill's zeros row-0 fallback covers it at a bounded acceptance cost).
939 pub last_h: Vec<f32>,
940}
941
942/// D2H one hidden row out of a `[T, n_embd]` prime hidden stack — the boundary anchor a
943/// spec boundary capture carries for later restored-session fills. Failure is silent
944/// (`turn_ckpt` convention): the capture publishes without an anchor.
945fn capture_boundary_hidden(
946 e: &Engine,
947 h_rows: &CudaSlice<f32>,
948 pos: usize,
949 n_embd: usize,
950) -> Vec<f32> {
951 if pos == 0 || h_rows.len() < pos * n_embd {
952 return Vec::new();
953 }
954 let Ok(mut row) = e.uninit(n_embd) else {
955 return Vec::new();
956 };
957 if e.copy_view_into(
958 &mut row,
959 0,
960 &h_rows.slice((pos - 1) * n_embd..pos * n_embd),
961 n_embd,
962 )
963 .is_err()
964 {
965 return Vec::new();
966 }
967 e.dtoh(&row).unwrap_or_default()
968}
969
970/// ROLLBACK DOOR for sampled BOUNDARY tokens (lane/sampled-spec-quality, 2026-08-19).
971/// Default ON: the token a burst emits at its own boundary is drawn from the request's
972/// sampler. `MEMRA_SPEC_SAMPLED_BOUNDARY=0` restores the pre-lane posture (an ARGMAX at
973/// every boundary) without touching greedy, which is byte-unaffected either way.
974pub fn spec_sampled_boundary_on() -> bool {
975 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
976 *ON.get_or_init(|| std::env::var("MEMRA_SPEC_SAMPLED_BOUNDARY").as_deref() != Ok("0"))
977}
978
979/// ROLLBACK DOOR for SESSION-SPANNING penalty history (lane/sampled-spec-quality).
980/// Default ON: `pen_hist` is seeded from the session's committed tail, so repetition /
981/// frequency / presence penalties see the whole stream. `MEMRA_SPEC_PEN_SESSION=0`
982/// restores the pre-lane posture (each burst restarts the window from its own prompt
983/// slice, i.e. from NOTHING on a continuation burst) — and with the door shut the worker
984/// must keep refusing penalized sampled prefix-cache restores, because the restored
985/// session's continuation burst is handed no prompt slice at all.
986pub fn spec_pen_session_on() -> bool {
987 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
988 *ON.get_or_init(|| std::env::var("MEMRA_SPEC_PEN_SESSION").as_deref() != Ok("0"))
989}
990
991/// ROLLBACK DOOR for extended-entry publication from a RESTORED session
992/// (lane/sampled-spec-quality, Item 3). Default ON: a converted prefix-cache hit that fed a
993/// suffix captures its own prompt-end boundary so the NEXT turn can hit a longer prefix.
994/// `MEMRA_SPEC_RESTORE_REPUBLISH=0` restores the pre-lane posture (a namespace learns exactly
995/// one boundary and never advances it). Whole-entry semantics only — the boundary is the
996/// restored session's own prompt end, so `entry_pos != fed_len` still refuses on the way in.
997pub fn spec_restore_republish_on() -> bool {
998 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
999 *ON.get_or_init(|| std::env::var("MEMRA_SPEC_RESTORE_REPUBLISH").as_deref() != Ok("0"))
1000}
1001
1002/// Diagnostics: name every boundary token on stderr (`MEMRA_SPEC_BOUNDARY_TRACE=1`), with
1003/// the argmax the pre-lane code would have emitted from the same row. This is how the
1004/// lane MEASURES the boundary rate and the deviation rate instead of estimating them.
1005fn spec_boundary_trace() -> bool {
1006 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1007 *ON.get_or_init(|| std::env::var("MEMRA_SPEC_BOUNDARY_TRACE").as_deref() == Ok("1"))
1008}
1009
1010/// llama-parity floor for the penalty window when the request does not ask for a bigger
1011/// one (`repeat_last_n` default). The serve API arms `penalty_last_n = usize::MAX` for any
1012/// non-identity penalty, so this floor only matters to explicit small windows and to the
1013/// CLI env path.
1014const PEN_WINDOW_FLOOR: usize = 64;
1015
1016/// CEILING on the penalty window, and it is a COST bound, not a semantic preference.
1017/// `penalize_logits_f32` (cu/spec_sample.cu) dedups on device by having thread `i` scan
1018/// `hist[0..i]`, so a pass is O(n_hist²) and it runs ~3x per verify round (the q rows, the
1019/// p column, the bonus column). The serve API arms `penalty_last_n = usize::MAX` for ANY
1020/// non-identity penalty — "the whole context", llama's `repeat_last_n = -1` — so an
1021/// uncapped session window would put a 128k-token history through that kernel: ~1.7e10
1022/// comparisons per pass, tens of ms per round, i.e. penalties would silently destroy decode
1023/// throughput on exactly the long-context requests that most want them. 8192 keeps a pass
1024/// at ~7e7 comparisons (tens of microseconds) while still being **128x wider than the
1025/// pre-lane effective window** (64 prompt-tail tokens + whatever the current burst had
1026/// generated). A request that genuinely needs a window beyond this wants host-side dedup +
1027/// counts through a new kernel signature — a follow-up lane, named here rather than hidden.
1028/// `pub` since lane/dspark-penalized-sampled-20260821: the dspark route's accept walk and
1029/// the dspark_sample_gate binary trim their uploads with the SAME cap — a second constant
1030/// is a second thing to drift.
1031pub const PEN_WINDOW_MAX: usize = 8192;
1032
1033/// Seed a penalty window over the SESSION, not the burst (lane/sampled-spec-quality,
1034/// Item 2). The window is the last `max(penalty_last_n, 64)` tokens of
1035/// `session_committed ++ burst_prompt` — for a cold turn-1 burst (`session_committed`
1036/// empty, default `penalty_last_n`) that is byte-identically the pre-lane
1037/// `prompt.iter().rev().take(64).rev()`; for a continuation burst it is the stream the
1038/// client actually asked us to penalize, where the pre-lane code had NOTHING.
1039/// `pub` since lane/dspark-penalized-sampled-20260821: the dspark route seeds its session
1040/// window through the SAME function (one definition of "the window" across both spec
1041/// routes and the gate binary's trunk-only reference arm).
1042pub fn pen_window_seed(
1043 session_committed: &[u32],
1044 burst_prompt: &[u32],
1045 penalty_last_n: usize,
1046) -> Vec<u32> {
1047 let win = penalty_last_n.clamp(PEN_WINDOW_FLOOR, PEN_WINDOW_MAX);
1048 let take_prompt = burst_prompt.len().min(win);
1049 let take_sess = (win - take_prompt).min(session_committed.len());
1050 let mut hist = Vec::with_capacity(take_sess + take_prompt);
1051 hist.extend_from_slice(&session_committed[session_committed.len() - take_sess..]);
1052 hist.extend_from_slice(&burst_prompt[burst_prompt.len() - take_prompt..]);
1053 hist
1054}
1055
1056/// Draw a BOUNDARY token from the target distribution the request asked for
1057/// (lane/sampled-spec-quality, Item 1) — the fix for "sampled spec emits an ARGMAX token at
1058/// every burst boundary".
1059///
1060/// WHY THIS EXISTS. A spec burst's first emitted token is not produced by the accept walk:
1061/// it comes off a logits row that already exists (the prime's last row on a cold burst; the
1062/// row after the last committed token on a continuation burst; the prefix-cache entry's
1063/// boundary row on a restored one). Pre-lane that token was `argmax` in BOTH sampling
1064/// regimes, so a sampled stream took a greedy token once per burst — measured, not
1065/// estimated, in research/spec-cache-20260818/SAMPLED-QUALITY.md. At temperature > 0 the
1066/// customer asked for a sampled token, so this draws one.
1067///
1068/// THE PROGRAM IS THE FULL-ACCEPT BONUS'S PROGRAM, deliberately: penalize the row (over the
1069/// session's window), take this row's OWN filter stats (the sampfix-20260805 law — stats
1070/// from a neighbour row mis-scale every `e0` and can wipe the row to token 0), gumbel-perturb
1071/// with the session's Philox stream at `*sctr`, argmax the perturbed row. Reusing the bonus's
1072/// composition means `sample_check`'s distributional oracle covers this draw too, and the
1073/// boundary token is drawn from the same filtered/penalized `p` the accept walk targets.
1074///
1075/// THE STREAM IS THE SESSION'S, NOT A FRESH ONE. `sctr` is the caller's live counter and is
1076/// advanced by exactly one, so a boundary draw consumes the next value in the same Philox
1077/// stream the accept walk uses — never a second, independently seeded stream (which would be
1078/// a new distributional bug: two streams from one seed correlate wherever their counters
1079/// collide). That also makes a restored session's boundary draw at `sctr == 0` bit-identical
1080/// to the cold session's own first draw from the same logits row, which is what preserves the
1081/// sampled-hit lane's per-seed hit==cold byte identity.
1082#[allow(clippy::too_many_arguments)]
1083pub fn sample_boundary_token_dev(
1084 e: &Engine,
1085 logits: &CudaSlice<f32>,
1086 n_vocab: usize,
1087 sp: &SpecSampling,
1088 pen_hist: &[u32],
1089 sctr: &mut u32,
1090 site: &str,
1091) -> Result<u32, Box<dyn std::error::Error>> {
1092 debug_assert!(
1093 sp.temp > 0.0,
1094 "boundary sampling is the sampled regime only"
1095 );
1096 // Own copy: penalize_logits mutates in place and the caller's row is live state
1097 // (prime_logits back the constrained recompute; last_col_logits backs round 0's accept).
1098 let mut col = e.zeros(n_vocab)?;
1099 e.copy_into(&mut col, 0, logits, n_vocab)?;
1100 let pen_on = sp.penalty_last_n > 0
1101 && (sp.penalty_repeat != 1.0 || sp.penalty_freq != 0.0 || sp.penalty_present != 0.0);
1102 if pen_on && !pen_hist.is_empty() {
1103 // window trim mirrors the round loop's own upload (`pen_hist[w0..]`), cap included.
1104 let w0 = pen_hist
1105 .len()
1106 .saturating_sub(sp.penalty_last_n.min(PEN_WINDOW_MAX));
1107 let hist = &pen_hist[w0..];
1108 let hd = e.htod_u32_v(hist)?;
1109 e.penalize_logits(
1110 &mut col,
1111 &hd,
1112 hist.len(),
1113 sp.penalty_repeat,
1114 sp.penalty_freq,
1115 sp.penalty_present,
1116 n_vocab,
1117 )?;
1118 }
1119 let rows0 = e.htod_i32(&[0])?;
1120 let (mut th_d, mut z_d, mut mx_d) = (e.zeros(1)?, e.zeros(1)?, e.zeros(1)?);
1121 e.filter_stats(
1122 &col, n_vocab, &rows0, &mut th_d, &mut z_d, &mut mx_d, n_vocab, 1, sp.temp, sp.top_k,
1123 sp.top_p, sp.min_p,
1124 )?;
1125 let (th, mx) = (e.dtoh(&th_d)?[0], e.dtoh(&mx_d)?[0]);
1126 let mut perturb = e.zeros(n_vocab)?;
1127 e.gumbel_perturb_filtered(&col, &mut perturb, n_vocab, sp.seed, *sctr, sp.temp, mx, th)?;
1128 *sctr = sctr.wrapping_add(1);
1129 let td = e.argmax_token_device(&perturb, n_vocab)?;
1130 let tok = e.dtoh_u32_one(&td)?;
1131 if spec_boundary_trace() {
1132 // the pre-lane token, from the SAME row, so the deviation rate is measurable.
1133 let raw = e.argmax_token_device(logits, n_vocab)?;
1134 let greedy = e.dtoh_u32_one(&raw)?;
1135 eprintln!(
1136 "[spec-boundary] site={site} sampled={tok} argmax={greedy} \
1137 deviates={} temp={} sctr={}",
1138 (tok != greedy) as u8,
1139 sp.temp,
1140 sctr.wrapping_sub(1),
1141 );
1142 }
1143 Ok(tok)
1144}
1145
1146/// Host-row twin of [`sample_boundary_token_dev`] (the prime / feed / entry rows arrive as
1147/// host `Vec<f32>`).
1148#[allow(clippy::too_many_arguments)]
1149pub fn sample_boundary_token(
1150 e: &Engine,
1151 logits: &[f32],
1152 sp: &SpecSampling,
1153 pen_hist: &[u32],
1154 sctr: &mut u32,
1155 site: &str,
1156) -> Result<u32, Box<dyn std::error::Error>> {
1157 let n_vocab = logits.len();
1158 let d = e.htod(logits)?;
1159 sample_boundary_token_dev(e, &d, n_vocab, sp, pen_hist, sctr, site)
1160}
1161
1162struct SpecPipeTraceClock {
1163 pair: usize,
1164 started: std::time::Instant,
1165}
1166
1167#[derive(Clone)]
1168struct SpecPipeTraceCtx {
1169 clock: std::sync::Arc<SpecPipeTraceClock>,
1170 round: usize,
1171 lane: usize,
1172}
1173
1174struct SpecPipeTraceMarker {
1175 trace: SpecPipeTraceCtx,
1176 phase: &'static str,
1177 edge: &'static str,
1178 slot: Option<usize>,
1179}
1180
1181unsafe extern "C" fn spec_pipe_trace_marker(raw: *mut std::ffi::c_void) {
1182 let marker = unsafe { Box::from_raw(raw.cast::<SpecPipeTraceMarker>()) };
1183 let lane = if marker.trace.lane == 0 { "A" } else { "B" };
1184 let slot = marker
1185 .slot
1186 .map(|v| v.to_string())
1187 .unwrap_or_else(|| "-".into());
1188 let t_ms = marker.trace.clock.started.elapsed().as_secs_f64() * 1e3;
1189 use std::io::Write as _;
1190 let stderr = std::io::stderr();
1191 let mut stderr = stderr.lock();
1192 let _ = writeln!(
1193 stderr,
1194 "[spec-pipe-timeline] pair={} round={} lane={lane} phase={} edge={} \
1195 slot={slot} t_ms={t_ms:.3}",
1196 marker.trace.clock.pair, marker.trace.round, marker.phase, marker.edge,
1197 );
1198}
1199
1200fn enqueue_spec_pipe_trace_marker(
1201 stream: &cudarc::driver::CudaStream,
1202 trace: Option<&SpecPipeTraceCtx>,
1203 phase: &'static str,
1204 edge: &'static str,
1205 slot: Option<usize>,
1206) -> Result<(), Box<dyn std::error::Error>> {
1207 let Some(trace) = trace else {
1208 return Ok(());
1209 };
1210 let marker = Box::new(SpecPipeTraceMarker {
1211 trace: trace.clone(),
1212 phase,
1213 edge,
1214 slot,
1215 });
1216 let raw = Box::into_raw(marker);
1217 let result = unsafe {
1218 cudarc::driver::result::stream::launch_host_function(
1219 stream.cu_stream(),
1220 spec_pipe_trace_marker,
1221 raw.cast(),
1222 )
1223 };
1224 if let Err(err) = result {
1225 unsafe {
1226 drop(Box::from_raw(raw));
1227 }
1228 return Err(err.into());
1229 }
1230 Ok(())
1231}
1232
1233#[derive(Default)]
1234struct SpecPipeProgress {
1235 setup_done: [bool; 2],
1236 draft_done: [usize; 2],
1237 stage0_done: [usize; 2],
1238 verify_done: [usize; 2],
1239 accept_done: [usize; 2],
1240 finished: [bool; 2],
1241 aborted: bool,
1242}
1243
1244/// Host-side issue coordinator for the reduced two-session speculative pipeline. Each session
1245/// keeps its existing call stack and round locals; this object only orders phase entry. The
1246/// primary mutex spans whole draft/accept/tail issue regions so Engine's single-stream scratch
1247/// cannot be interleaved by the two host threads.
1248struct SpecPipeSync {
1249 progress: std::sync::Mutex<SpecPipeProgress>,
1250 changed: std::sync::Condvar,
1251 primary: std::sync::Mutex<()>,
1252 trace: Option<std::sync::Arc<SpecPipeTraceClock>>,
1253}
1254
1255impl SpecPipeSync {
1256 fn new() -> Self {
1257 static TRACE_PAIR: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
1258 let trace = (std::env::var("MEMRA_SPEC_PIPE_TRACE").as_deref() == Ok("1")).then(|| {
1259 std::sync::Arc::new(SpecPipeTraceClock {
1260 pair: TRACE_PAIR.fetch_add(1, std::sync::atomic::Ordering::Relaxed) + 1,
1261 started: std::time::Instant::now(),
1262 })
1263 });
1264 Self {
1265 progress: std::sync::Mutex::new(SpecPipeProgress::default()),
1266 changed: std::sync::Condvar::new(),
1267 primary: std::sync::Mutex::new(()),
1268 trace,
1269 }
1270 }
1271}
1272
1273#[derive(Clone)]
1274struct SpecPipeLane {
1275 sync: std::sync::Arc<SpecPipeSync>,
1276 lane: usize,
1277}
1278
1279impl SpecPipeLane {
1280 fn peer(&self) -> usize {
1281 1 - self.lane
1282 }
1283
1284 fn aborted() -> Box<dyn std::error::Error> {
1285 "paired speculative peer aborted".into()
1286 }
1287
1288 fn trace(&self, round: usize) -> Option<SpecPipeTraceCtx> {
1289 self.sync.trace.as_ref().map(|clock| SpecPipeTraceCtx {
1290 clock: clock.clone(),
1291 round,
1292 lane: self.lane,
1293 })
1294 }
1295
1296 fn setup_begin(&self) -> Result<(), Box<dyn std::error::Error>> {
1297 let mut p = self.sync.progress.lock().unwrap();
1298 while !p.aborted && self.lane == 1 && !p.setup_done[0] && !p.finished[0] {
1299 p = self.sync.changed.wait(p).unwrap();
1300 }
1301 if p.aborted {
1302 Err(Self::aborted())
1303 } else {
1304 Ok(())
1305 }
1306 }
1307
1308 fn setup_end(&self) {
1309 let mut p = self.sync.progress.lock().unwrap();
1310 p.setup_done[self.lane] = true;
1311 self.sync.changed.notify_all();
1312 }
1313
1314 fn draft_begin(
1315 &self,
1316 round: usize,
1317 ) -> Result<std::sync::MutexGuard<'_, ()>, Box<dyn std::error::Error>> {
1318 let peer = self.peer();
1319 let mut p = self.sync.progress.lock().unwrap();
1320 loop {
1321 if p.aborted {
1322 return Err(Self::aborted());
1323 }
1324 let setup_ready =
1325 (p.setup_done[0] || p.finished[0]) && (p.setup_done[1] || p.finished[1]);
1326 let prior_ready = p.accept_done[self.lane] >= round
1327 && (p.accept_done[peer] >= round || p.finished[peer]);
1328 let turn_ready = if self.lane == 0 {
1329 true
1330 } else {
1331 p.draft_done[0] > round || p.finished[0]
1332 };
1333 if setup_ready && prior_ready && turn_ready {
1334 break;
1335 }
1336 p = self.sync.changed.wait(p).unwrap();
1337 }
1338 drop(p);
1339 Ok(self.sync.primary.lock().unwrap())
1340 }
1341
1342 fn draft_end(&self, round: usize) {
1343 let mut p = self.sync.progress.lock().unwrap();
1344 p.draft_done[self.lane] = round + 1;
1345 self.sync.changed.notify_all();
1346 }
1347
1348 /// Admit stage 0 and return whether this lane owns the interval's one reverse fence.
1349 /// Lane B releases as soon as lane A has issued its boundary TX, not after A's full body.
1350 fn stage0_begin(&self, round: usize) -> Result<bool, Box<dyn std::error::Error>> {
1351 let peer = self.peer();
1352 let mut p = self.sync.progress.lock().unwrap();
1353 loop {
1354 if p.aborted {
1355 return Err(Self::aborted());
1356 }
1357 let ready = if self.lane == 0 {
1358 p.draft_done[0] > round && (p.draft_done[1] > round || p.finished[1])
1359 } else {
1360 p.draft_done[1] > round && (p.stage0_done[0] > round || p.finished[0])
1361 };
1362 if ready {
1363 return Ok(self.lane == 0 || p.finished[peer]);
1364 }
1365 p = self.sync.changed.wait(p).unwrap();
1366 }
1367 }
1368
1369 fn stage0_end(&self, round: usize) {
1370 let mut p = self.sync.progress.lock().unwrap();
1371 p.stage0_done[self.lane] = round + 1;
1372 self.sync.changed.notify_all();
1373 }
1374
1375 /// Stage 1 is single-owner per engine. A proceeds immediately after its own ticket; B waits
1376 /// for A's full stage1/head issue so only A.S1 and B.S0 can overlap.
1377 fn stage1_begin(&self, round: usize) -> Result<(), Box<dyn std::error::Error>> {
1378 let mut p = self.sync.progress.lock().unwrap();
1379 while !p.aborted
1380 && !(p.stage0_done[self.lane] > round
1381 && (self.lane == 0 || p.verify_done[0] > round || p.finished[0]))
1382 {
1383 p = self.sync.changed.wait(p).unwrap();
1384 }
1385 if p.aborted {
1386 Err(Self::aborted())
1387 } else {
1388 Ok(())
1389 }
1390 }
1391
1392 fn verify_end(&self, round: usize) {
1393 let mut p = self.sync.progress.lock().unwrap();
1394 p.verify_done[self.lane] = round + 1;
1395 self.sync.changed.notify_all();
1396 }
1397
1398 fn accept_begin(
1399 &self,
1400 round: usize,
1401 ) -> Result<std::sync::MutexGuard<'_, ()>, Box<dyn std::error::Error>> {
1402 let mut p = self.sync.progress.lock().unwrap();
1403 loop {
1404 if p.aborted {
1405 return Err(Self::aborted());
1406 }
1407 let ready = if self.lane == 0 {
1408 p.verify_done[0] > round && (p.verify_done[1] > round || p.finished[1])
1409 } else {
1410 p.verify_done[1] > round && (p.accept_done[0] > round || p.finished[0])
1411 };
1412 if ready {
1413 break;
1414 }
1415 p = self.sync.changed.wait(p).unwrap();
1416 }
1417 drop(p);
1418 Ok(self.sync.primary.lock().unwrap())
1419 }
1420
1421 fn accept_end(&self, round: usize) {
1422 let mut p = self.sync.progress.lock().unwrap();
1423 p.accept_done[self.lane] = round + 1;
1424 self.sync.changed.notify_all();
1425 }
1426
1427 fn primary(&self) -> std::sync::MutexGuard<'_, ()> {
1428 self.sync.primary.lock().unwrap()
1429 }
1430
1431 fn finish(&self, failed: bool) {
1432 let mut p = self.sync.progress.lock().unwrap();
1433 p.finished[self.lane] = true;
1434 p.aborted |= failed;
1435 self.sync.changed.notify_all();
1436 }
1437}
1438
1439struct SpecPipeFinish<'a> {
1440 lane: &'a SpecPipeLane,
1441 closed: bool,
1442}
1443
1444impl<'a> SpecPipeFinish<'a> {
1445 fn new(lane: &'a SpecPipeLane) -> Self {
1446 Self {
1447 lane,
1448 closed: false,
1449 }
1450 }
1451
1452 fn close(&mut self, failed: bool) {
1453 self.lane.finish(failed);
1454 self.closed = true;
1455 }
1456}
1457
1458impl Drop for SpecPipeFinish<'_> {
1459 fn drop(&mut self) {
1460 if !self.closed {
1461 self.lane.finish(true);
1462 }
1463 }
1464}
1465
1466/// Scoped transfer of one exclusively-borrowed session to the second host issue thread.
1467/// `CudaGraph` is not marked Send by cudarc because its raw driver handles carry no automatic
1468/// trait. CUDA driver graph handles are context-scoped rather than OS-thread-affine; the caller
1469/// binds that context before touching the session, joins before returning, and never aliases the
1470/// pointer. Keep this exception local to the experimental pair call instead of marking the public
1471/// session type Send.
1472struct SpecPipeSessionPtr(*mut SpecSession);
1473
1474unsafe impl Send for SpecPipeSessionPtr {}
1475
1476impl SpecPipeSessionPtr {
1477 unsafe fn get_mut(&mut self) -> &mut SpecSession {
1478 unsafe { &mut *self.0 }
1479 }
1480}
1481
1482/// Per-session persistent draft-graph context: the captured CUDA graph(s) plus the device
1483/// buffers whose POINTERS the capture bakes. Reuse legality: the greedy capture bakes only
1484/// session-stable pointers (the session's own MtpScratch KV — allocated once, never realloc'd;
1485/// the model's resident embedding; the process-wide OnceLock p_min) and the g_* buffers held
1486/// HERE — so one capture serves the session's whole lifetime. The sampled capture additionally
1487/// bakes (seed, temp) as capture-time constants and needs k q-slots — keyed by `s_key`, dropped
1488/// and recaptured when a pool-resumed request changes them. `*_failed` memoizes a failed capture
1489/// so the eager fallback doesn't pay a doomed capture attempt every burst.
1490/// Capture identity of the parked SAMPLED draft graph (`DraftGraphCtx::graph_s`).
1491///
1492/// EXACTNESS, not perf (lane/graph-s-key-exactness-20260819; receipts
1493/// `research/spec-cache-20260818/GRAPH-S-KEY.md`). Two classes of field live here, both
1494/// load-bearing:
1495///
1496/// - **Baked constants.** `seed` and `temp` are capture-time constants INSIDE the graph and `k`
1497/// sizes the q slots its replays write. A resumed request changing any of them must recapture.
1498/// This is all the key used to carry.
1499/// - **Regime fields.** `top_k`/`top_p`/`min_p`/`pen_on` are not baked, but they decide whether
1500/// the captured graph is a legal draft chain AT ALL. The in-graph draw is one gumbel-max over
1501/// the RAW softmax (`gumbel_perturb_ctr`, unfiltered by construction), while the verify builds
1502/// the accept test's `q` from `filter_stats(q_slots, top_k, top_p, min_p)`. If those disagree
1503/// the accept test evaluates a distribution the draft was never sampled from: a draft token
1504/// below the filter threshold gathers `q = 0` (`softmax_gather_filtered_f32`,
1505/// `cu/spec_sample.cu`) and `u * 0 < p` accepts it UNCONDITIONALLY.
1506///
1507/// Omitting the regime fields was reachable — not through the prefix-cache spec restore (that
1508/// path is greedy-only, `memra-server` `spec_restore_convertible`), but through WHOLE-SESSION
1509/// spec reuse: a parked `SpecSession` carries this `DraftGraphCtx`, and the pool-resume probe
1510/// applies no sampler predicate at all. Turn 1 pure-temp parks a graph; turn 2 of the same
1511/// conversation, same explicit seed and temperature, adds `top_p`/`top_k` and inherits it.
1512#[derive(Clone, Copy, PartialEq, Eq, Debug)]
1513pub(crate) struct SampledGraphKey {
1514 seed: u64,
1515 temp_bits: u32,
1516 k: usize,
1517 top_k: i32,
1518 top_p_bits: u32,
1519 min_p_bits: u32,
1520 pen_on: bool,
1521}
1522
1523impl SampledGraphKey {
1524 pub(crate) fn new(
1525 seed: u64,
1526 temp: f32,
1527 k: usize,
1528 top_k: i32,
1529 top_p: f32,
1530 min_p: f32,
1531 pen_on: bool,
1532 ) -> Self {
1533 SampledGraphKey {
1534 seed,
1535 temp_bits: temp.to_bits(),
1536 k,
1537 top_k,
1538 top_p_bits: top_p.to_bits(),
1539 min_p_bits: min_p.to_bits(),
1540 pen_on,
1541 }
1542 }
1543
1544 /// The one regime the in-graph sampled chain may stand in for the eager one: nothing but
1545 /// temperature shapes `q`. Computed FROM THE KEY so the capture guard, the launch guard and
1546 /// the key can never drift apart (they were three separate expressions before this lane, and
1547 /// the launch site simply forgot to ask).
1548 pub(crate) fn pure_temp(&self) -> bool {
1549 self.top_k == 0
1550 && f32::from_bits(self.top_p_bits) >= 1.0
1551 && f32::from_bits(self.min_p_bits) <= 0.0
1552 && !self.pen_on
1553 }
1554}
1555
1556pub(crate) struct DraftGraphCtx {
1557 g_tok: CudaSlice<u32>,
1558 g_pos: CudaSlice<i32>,
1559 g_seed: CudaSlice<f32>,
1560 g_p: CudaSlice<f32>,
1561 g_ctr: CudaSlice<u32>,
1562 g_q: CudaSlice<f32>,
1563 g_perturb: CudaSlice<f32>,
1564 q_slots: Vec<CudaSlice<f32>>,
1565 /// DRAFT-SIDE GRAMMAR MASK (lane/draft-mask): packed allowed-set words over the DRAFT
1566 /// head's vocab, at a STABLE address so the captured draft graph's mask node reads the
1567 /// per-position contents the host re-uploads before each replay (the graph-promote
1568 /// pattern from decode.rs). Empty unless the session drafts under a grammar.
1569 g_dmask: CudaSlice<u32>,
1570 /// was `graph` captured WITH the mask node? A parked graph of the wrong shape is dropped.
1571 graph_masked: bool,
1572 graph: Option<cudarc::driver::CudaGraph>,
1573 graph_s: Option<cudarc::driver::CudaGraph>,
1574 /// Failed-capture memoization for both graphs — LOUD on flip, cleared on pool resume
1575 /// (audit Q2, the TRT #16072 silent-permanent-coverage-loss class).
1576 failed: DraftGraphFallback,
1577 /// Capture identity of `graph_s` — see [`SampledGraphKey`]. `None` iff no sampled graph is
1578 /// parked; a request whose key differs drops the parked graph (and its q slots/keeper).
1579 s_key: Option<SampledGraphKey>,
1580 /// CAPTURE-RETAIN keepers (#68 root cause, 2026-08-04): the warmup-run transients whose
1581 /// pool addresses the captured graph(s) bake. Without these, the transients return to the
1582 /// pool at capture-body exit and later work (burst-boundary prime/fill/commit passes, or a
1583 /// co-served session in the worker) reuses those addresses — the persisted graph's replay
1584 /// then reads/writes live unrelated buffers (exactness corruption, first seen as the ST
1585 /// serve-spec 4B graph-arm corruption; one-shot CLI calls never re-shuffled the pool, which
1586 /// is why run-spec K=1..8 passed on the same checkpoint). Same fix class as
1587 /// capture_graph_retained's gemma/decode.rs sites — hold as long as the graph replays.
1588 keeper: Vec<Box<dyn std::any::Any + Send>>,
1589 keeper_s: Vec<Box<dyn std::any::Any + Send>>,
1590}
1591
1592/// Failed-capture memoization for the two draft graphs (audit Q2, 2026-08-05 — the
1593/// TRT #16072 trap class: pressure-triggered, silent, long-lived coverage loss).
1594///
1595/// Three contracts:
1596/// - LOUD FLIP: `mark_*` returns the warn line exactly on the false→true transition
1597/// (returned, not printed, so the once-per-flip contract is unit-testable); the caller
1598/// `eprintln!`s it UNCONDITIONALLY — a dropped draft graph is never silent. Re-marking
1599/// an already-failed graph returns None (the per-burst memoization that keeps the eager
1600/// fallback from paying a doomed capture attempt every burst).
1601/// - RESET ON RESUME: `reset_on_resume` clears both flags — a parked session resumed by a
1602/// NEW request gets one fresh capture chance instead of carrying a transient-pressure
1603/// failure for the pool's whole lifetime. Returns the note line only when a flag was
1604/// actually set (quiet on the common clean-resume path).
1605/// - Shape-change clears (`clear_*`) stay silent, exactly as before: they precede a fresh
1606/// capture attempt whose own failure would re-flip loudly.
1607#[derive(Default)]
1608pub(crate) struct DraftGraphFallback {
1609 greedy: bool,
1610 sampled: bool,
1611}
1612impl DraftGraphFallback {
1613 fn mark_greedy(&mut self, reason: &str) -> Option<String> {
1614 if self.greedy {
1615 return None;
1616 }
1617 self.greedy = true;
1618 Some(format!(
1619 "[spec] WARN: draft-graph capture failed ({reason}); eager fallback until session resume"
1620 ))
1621 }
1622 fn mark_sampled(&mut self, reason: &str) -> Option<String> {
1623 if self.sampled {
1624 return None;
1625 }
1626 self.sampled = true;
1627 Some(format!(
1628 "[spec] WARN: sampled draft-graph capture failed ({reason}); eager fallback until session resume"
1629 ))
1630 }
1631 fn greedy_failed(&self) -> bool {
1632 self.greedy
1633 }
1634 fn sampled_failed(&self) -> bool {
1635 self.sampled
1636 }
1637 fn clear_greedy(&mut self) {
1638 self.greedy = false;
1639 }
1640 fn clear_sampled(&mut self) {
1641 self.sampled = false;
1642 }
1643 /// Pool-resume reset: both graphs get a fresh capture chance. Some(note) iff any flag
1644 /// was set (so clean resumes stay quiet).
1645 pub(crate) fn reset_on_resume(&mut self) -> Option<String> {
1646 if !self.greedy && !self.sampled {
1647 return None;
1648 }
1649 let which = match (self.greedy, self.sampled) {
1650 (true, true) => "greedy+sampled",
1651 (true, false) => "greedy",
1652 _ => "sampled",
1653 };
1654 self.greedy = false;
1655 self.sampled = false;
1656 Some(format!(
1657 "[spec] draft-graph fallback reset on session resume ({which}); recapture eligible"
1658 ))
1659 }
1660}
1661
1662impl DraftGraphCtx {
1663 fn new(e: &Engine, n_embd: usize, qlen: usize) -> Result<Self, Box<dyn std::error::Error>> {
1664 Ok(DraftGraphCtx {
1665 g_tok: e.alloc_u32_zeroed(1)?,
1666 g_pos: e.htod_i32(&[0])?,
1667 g_seed: e.zeros(n_embd)?,
1668 g_p: e.zeros(1)?,
1669 g_ctr: e.alloc_u32_zeroed(1)?,
1670 g_q: e.zeros(qlen)?,
1671 g_perturb: e.zeros(qlen)?,
1672 q_slots: Vec::new(),
1673 g_dmask: e.alloc_u32_zeroed(1)?,
1674 graph_masked: false,
1675 graph: None,
1676 graph_s: None,
1677 failed: DraftGraphFallback::default(),
1678 s_key: None,
1679 keeper: Vec::new(),
1680 keeper_s: Vec::new(),
1681 })
1682 }
1683}
1684
1685pub(crate) struct MtpScratch {
1686 kv: KvLayer,
1687 /// Logical row capacity. On the graph/DC draft path it also doubles as fa_decode_dc's
1688 /// bucket_max: n_splits is sized from it ONCE, so the graph captured at round 0 stays valid
1689 /// for every later t_kv. Step35 refuses that path and may back this logical extent with the
1690 /// smaller host-indexed SWA ring instead.
1691 cap: usize,
1692 extra: Vec<MtpScratchPlane>,
1693}
1694
1695struct MtpScratchPlane {
1696 kv: KvLayer,
1697 cap: usize,
1698}
1699
1700fn mtp_scratch_layout(
1701 cfg: &memra_gguf::config::ModelConfig,
1702 geom: Option<&crate::hybrid::DraftGeom>,
1703) -> (usize, usize, usize, usize) {
1704 // Student draft heads carry fewer KV heads (head_dim unchanged) -> smaller scratch rows.
1705 let n_head_kv = geom.map(|g| g.n_head_kv).unwrap_or(cfg.n_head_kv as usize);
1706 let head_dim_k = cfg.head_dim_k as usize;
1707 let head_dim_v = cfg.head_dim_v as usize;
1708 assert!(
1709 head_dim_k % 32 == 0 && head_dim_v % 32 == 0,
1710 "KVQUANT requires head_dim%32==0 (MTP scratch)"
1711 );
1712 let kv_dim_k = head_dim_k * n_head_kv;
1713 let kv_dim_v = head_dim_v * n_head_kv;
1714 // The fp8-KV arm deliberately does not reach the draft scratch; keep the exact format
1715 // policy shared with `MtpScratch::new` so admission scales the same allocation.
1716 let (kbb, vbb) = crate::kv_blk_bytes();
1717 let k_tok_bytes = (kv_dim_k / 32) * kbb;
1718 let v_tok_bytes = (kv_dim_v / 32) * vbb;
1719 (kv_dim_k, kv_dim_v, k_tok_bytes, v_tok_bytes)
1720}
1721
1722fn mtp_chain_head_index(step: usize, head_count: usize) -> usize {
1723 assert!(head_count > 0, "MTP chain requires at least one head");
1724 step % head_count
1725}
1726
1727impl MtpScratch {
1728 fn alloc_plane(
1729 e: &Engine,
1730 cfg: &memra_gguf::config::ModelConfig,
1731 plan: &memra_gguf::model_plan::ModelPlan,
1732 cap: usize,
1733 geom: Option<&crate::hybrid::DraftGeom>,
1734 ) -> Result<MtpScratchPlane, Box<dyn std::error::Error>> {
1735 let (kv_dim_k, kv_dim_v, k_tok_bytes, v_tok_bytes) = mtp_scratch_layout(cfg, geom);
1736 let ring = if crate::cache::swa_ring_on()
1737 && crate::plan_backend::decode_batch_program(plan)
1738 == crate::plan_backend::DecodeBatchProgram::SlidingGatedMoe
1739 {
1740 let window = plan
1741 .layers
1742 .iter()
1743 .find_map(|layer| match layer.attention {
1744 memra_gguf::model_plan::AttentionPlan::SlidingWindow { window, .. } => {
1745 Some(window as usize)
1746 }
1747 _ => None,
1748 })
1749 .ok_or("sliding-gated-MoE draft scratch has no sliding-window layer")?;
1750 Some(crate::cache::KvRing::new(
1751 crate::cache::swa_ring_rows(window, cap),
1752 window,
1753 ))
1754 } else {
1755 None
1756 };
1757 let alloc_rows = ring.as_ref().map(crate::cache::KvRing::rows).unwrap_or(cap);
1758 Ok(MtpScratchPlane {
1759 kv: KvLayer {
1760 k: e.alloc_u8(alloc_rows * k_tok_bytes)?,
1761 v: e.alloc_u8(alloc_rows * v_tok_bytes)?,
1762 kv_dim_k,
1763 kv_dim_v,
1764 k_tok_bytes,
1765 v_tok_bytes,
1766 len: 0,
1767 ring,
1768 len_d: e.htod_i32(&[0])?,
1769 },
1770 cap,
1771 })
1772 }
1773
1774 fn new(
1775 e: &Engine,
1776 cfg: &memra_gguf::config::ModelConfig,
1777 plan: &memra_gguf::model_plan::ModelPlan,
1778 cap: usize,
1779 geom: Option<&crate::hybrid::DraftGeom>,
1780 ) -> Result<Self, Box<dyn std::error::Error>> {
1781 // env-selected KV formats (default 34/24). The fp8-KV arm (MEMRA_KV_FP8) deliberately
1782 // does NOT reach the draft scratch: fp8 drafts drifted acceptance 69-88% -> 46%
1783 // (2026-07-12 A/B); the scratch is tiny, so it keeps baseline q8_0/q5_1 numerics
1784 // while the TRUNK cache carries the fp8 depth win. Scratch append/fa pass g=false.
1785 let primary = Self::alloc_plane(e, cfg, plan, cap, geom)?;
1786 Ok(MtpScratch {
1787 kv: primary.kv,
1788 cap: primary.cap,
1789 extra: Vec::new(),
1790 })
1791 }
1792
1793 fn push_plane(
1794 &mut self,
1795 e: &Engine,
1796 cfg: &memra_gguf::config::ModelConfig,
1797 plan: &memra_gguf::model_plan::ModelPlan,
1798 geom: Option<&crate::hybrid::DraftGeom>,
1799 ) -> Result<(), Box<dyn std::error::Error>> {
1800 self.extra
1801 .push(Self::alloc_plane(e, cfg, plan, self.cap, geom)?);
1802 Ok(())
1803 }
1804
1805 fn plane_count(&self) -> usize {
1806 1 + self.extra.len()
1807 }
1808
1809 fn plane(&self, index: usize) -> (&KvLayer, usize) {
1810 if index == 0 {
1811 (&self.kv, self.cap)
1812 } else {
1813 let plane = &self.extra[index - 1];
1814 (&plane.kv, plane.cap)
1815 }
1816 }
1817
1818 fn plane_mut(&mut self, index: usize) -> (&mut KvLayer, usize) {
1819 if index == 0 {
1820 (&mut self.kv, self.cap)
1821 } else {
1822 let plane = &mut self.extra[index - 1];
1823 (&mut plane.kv, plane.cap)
1824 }
1825 }
1826
1827 fn set_plane_len(
1828 &mut self,
1829 e: &Engine,
1830 index: usize,
1831 n: usize,
1832 ) -> Result<(), Box<dyn std::error::Error>> {
1833 let (kv, _) = self.plane_mut(index);
1834 if kv.ring.as_ref().is_some_and(|ring| !ring.can_rewind_to(n)) {
1835 return Err("SWA ring MTP checkpoint has been lapped; full re-prime required".into());
1836 }
1837 kv.len = n;
1838 e.set_i32_one(&mut kv.len_d, n as i32)
1839 }
1840
1841 /// Set BOTH length counters: the host mirror AND the device len_d the captured append/fa read
1842 /// (a 4-byte in-place htod — the counter pointer is baked into the graph, never realloc'd).
1843 /// This is the ONLY truncation/rollback mechanism the persistent draft KV needs.
1844 fn set_len(&mut self, e: &Engine, n: usize) -> Result<(), Box<dyn std::error::Error>> {
1845 if !self.can_rewind_to(n) {
1846 return Err("SWA ring MTP checkpoint has been lapped; full re-prime required".into());
1847 }
1848 for index in 0..self.plane_count() {
1849 self.set_plane_len(e, index, n)?;
1850 }
1851 Ok(())
1852 }
1853
1854 fn can_rewind_to(&self, n: usize) -> bool {
1855 (0..self.plane_count()).all(|index| {
1856 self.plane(index)
1857 .0
1858 .ring
1859 .as_ref()
1860 .is_none_or(|ring| ring.can_rewind_to(n))
1861 })
1862 }
1863}
1864
1865/// Retained verify intermediates for the REPLAY-FREE partial accept (2026-07-03, the profiled
1866/// #1 spec cost at long ctx: the partial-accept replay was a DUPLICATE trunk pass — ~0.54 extra
1867/// full weight reads per round — recomputing columns the verify had already produced
1868/// bit-identically). Holds, per linear layer, everything needed to rebuild its recurrent state
1869/// to "after the first j verify columns" WITHOUT re-running the trunk:
1870/// - BATCHED-path layers (`gdn`): the exact token-major inputs the round's ONE gdn_scan
1871/// consumed. A prefix re-run of the SAME kernel (t=j) from the snapshot state is bit-identical
1872/// to the first j iterations of the verify's scan — the kernel's t-loop carries state in
1873/// registers and iteration t never depends on T. `qkv_mixed` (the conv input) feeds the
1874/// pure-copy ring rebuild.
1875/// - PER-COLUMN-path layers (`cols`): dtod clones of (conv_state, ssm_state) taken after each
1876/// column 0..t-2 — pure copies of the actual chain states (the last column is never a rebuild
1877/// target: j <= t-1).
1878/// Full-attn layers need nothing: their verify KV rows are bit-identical to eager's (the
1879/// decode-exact contract; verify-probe pins it), so rollback = len truncation.
1880struct GdnStash {
1881 qkv_mixed: CudaSlice<f32>, // [t, conv_dim] token-major (conv input)
1882 q_l2: CudaSlice<f32>,
1883 k_l2: CudaSlice<f32>,
1884 v_g: CudaSlice<f32>, // [t, num_v, d_state]
1885 g_log: CudaSlice<f32>,
1886 beta: CudaSlice<f32>, // [t, num_v]
1887}
1888pub(crate) struct VerifyCkpt {
1889 gdn: Vec<Option<GdnStash>>, // [n_layer], Some iff batched linear path ran
1890 cols: Vec<Option<Vec<(CudaSlice<f32>, CudaSlice<f32>)>>>, // [n_layer][col] = (conv, ssm) after col
1891}
1892/// Opaque handle for the dspark round (dflash.rs) — VerifyCkpt stays spec-private.
1893pub(crate) struct DsparkVerifyCkpt(VerifyCkpt);
1894
1895/// Engine-bundle slice 3 (DSF-ROUNDCOST-20260820 §2 row 4 / §5 rank 1): bucketed CUDA
1896/// graphs for the dspark verify's LINEAR-layer segments. The measured verify is ~2,800
1897/// eager launches whose residual cost is DEVICE-side per-launch overhead (slice 2 proved
1898/// host dispatch is not the binder: fully-deferred dispatch bought ~0 wall). The 48 GDN
1899/// layers between full-attention layers are shape-static given vt — no positions, no
1900/// t_kv, state addressed through pointer tables — so runs of them capture per
1901/// (segment, vt) and replay as ONE graph launch each. Full-attention layers stay eager
1902/// (their per-row append/fa arm picks are t_kv-driven — the exec-update extension).
1903///
1904/// Per round out-of-graph: one pointer-table refresh (gdn ping-pong moves the canonical
1905/// handles), one input-staging copy per segment, host parity bookkeeping. Captured via
1906/// `capture_graph_retained` (2 warmups + capture, keeper retains warmup transients so
1907/// pool addresses stay stable); the warmups EXECUTE, so segment conv/ssm state is saved
1908/// before and restored after — the graph's first real launch starts from the exact
1909/// pre-round state. The ckpt column stash rides persistent slabs (written inside the
1910/// graph as memcpy nodes); commit reads them via `dspark_commit_prefix_slab`.
1911/// `MEMRA_DSPARK_VERIFY_GRAPH=0` reverts to the eager walk (byte-identical body).
1912pub(crate) struct DsparkVerifyGraphs {
1913 /// Linear-attention layer indices ascending; `lin_pos[il]` = index into the vecs.
1914 lin: Vec<usize>,
1915 lin_pos: std::collections::HashMap<usize, usize>,
1916 /// [n_lin x 6] pointer table (conv, s0, s1, conv, s1, s0 per layer), refreshed per
1917 /// verify from the live handles; layer il's slice starts at lin_pos[il]*6.
1918 table_all: CudaSlice<u64>,
1919 host_table: Vec<u64>,
1920 /// Persistent per-layer ckpt stash slabs: row r of the verify at slab offset
1921 /// r*words. Shared by every (segment, vt) bucket — one verify runs at a time.
1922 stash_conv: Vec<CudaSlice<f32>>,
1923 stash_ssm: Vec<CudaSlice<f32>>,
1924 conv_words: usize,
1925 ssm_words: usize,
1926 /// Per-vt input/output staging (stable addresses the graphs bake).
1927 stage: std::collections::HashMap<usize, (CudaSlice<f32>, CudaSlice<f32>)>,
1928 /// Per-vt dflash tap-sink buffers — the captured segments bake the tap dst address,
1929 /// so the sink buffer must live (and persist) with the graphs, not with the round.
1930 pub(crate) tap_bufs: std::collections::HashMap<usize, CudaSlice<f32>>,
1931 graphs: std::collections::HashMap<(usize, usize), DsparkSegGraph>,
1932 /// Warmup-corruption guard scratch: pre-capture conv/ssm of every linear layer
1933 /// (sized n_lin — the slice-4c full-verify warmups execute the whole walk).
1934 save_conv: CudaSlice<f32>,
1935 save_ssm: CudaSlice<f32>,
1936 max_run: usize,
1937 n_embd: usize,
1938 /// Set by the verify walk: this round's linear ckpt lives in the slabs (the caller
1939 /// commits through `dspark_commit_prefix_slab` instead of the cols arm).
1940 pub(crate) round_slab: bool,
1941 // ---- slice 4c: full-verify single graph per (vt, rung) ----
1942 /// Full-attention layer indices ascending; `fa_pos[il]` = index into the vec.
1943 fa: Vec<usize>,
1944 fa_pos: std::collections::HashMap<usize, usize>,
1945 /// [n_fa x 2 x t_cap] interleaved (k,v) base-pointer pairs, refreshed per verify;
1946 /// layer il's slice starts at `fa_pos[il] * 2 * t_cap` (the seqs twins read pairs
1947 /// [2z], z < t <= t_cap, so one t_cap-sized table serves every vt).
1948 fa_table: CudaSlice<u64>,
1949 fa_host_table: Vec<u64>,
1950 t_cap: usize,
1951 /// Per-vt position staging for the captured bodies — contents refreshed per round
1952 /// (rope reads row r; the seqs twins derive append slot and T_kv per z from it).
1953 pos_stage: std::collections::HashMap<usize, CudaSlice<i32>>,
1954 /// Full-verify graphs keyed (vt, rung_end, hi).
1955 full: std::collections::HashMap<(usize, usize, usize), DsparkSegGraph>,
1956 /// Largest n with every layer in [0, n) linear or full-attention (walk coverage).
1957 covered: usize,
1958 /// Every layer in [0, n) is linear or full-attention (no MLA/unknown mixers) — the
1959 /// full-verify capture walks all of them.
1960 walk_uniform: bool,
1961}
1962
1963struct DsparkSegGraph {
1964 graph: cudarc::driver::CudaGraph,
1965 _keeper: Vec<Box<dyn std::any::Any + Send>>,
1966}
1967
1968/// Per-call arguments of [`HybridModel::qwen35_tparallel_fa_layer`] — one struct so the
1969/// eager walk and the slice-4c captured full-verify graphs hand the SAME body its two
1970/// modes without a second copy of the math.
1971pub(crate) struct FaLayerArgs<'a> {
1972 /// [T] per-row positions (device): rope reads them row-indexed; the seqs twins read
1973 /// them per-z (append slot = pos, T_kv = pos + 1).
1974 pub pos_d: &'a CudaSlice<i32>,
1975 /// Verify-level lazy per-row 1-element position buffers — only the per-row fallback
1976 /// arm builds/uses them (graph mode refuses that arm).
1977 pub pos_rows: &'a mut Option<Vec<CudaSlice<i32>>>,
1978 pub pos0: usize,
1979 pub seqs_append: bool,
1980 pub batch_fa_on: bool,
1981 /// Some((kv pointer table, offset-in-u64s, rung_end)) = captured-graph mode.
1982 pub graph_cap: Option<(&'a CudaSlice<u64>, usize, usize)>,
1983 /// ROUND-STREAM (lane/draftcost-moe, v0.100 train merge): Some((token stream, device
1984 /// round counter)) routes the FA attend through the dc rows kernels and the Linear
1985 /// mixer through `linear_attn_verify_t` (the stream arms the old inline body carried).
1986 /// Never armed together with `graph_cap` (the verify-level merge guard refuses).
1987 pub stream: Option<(&'a CudaSlice<u32>, &'a CudaSlice<i32>)>,
1988 /// VerifyCkpt for the stream-Linear arm's GdnStash install; None in graph mode and
1989 /// for FA layers that never touch it.
1990 pub ckpt: Option<&'a mut VerifyCkpt>,
1991}
1992
1993// SAFETY: `CudaGraph` is not marked Send by cudarc because its raw driver handles carry
1994// no automatic trait; CUDA driver graph handles are context-scoped rather than
1995// OS-thread-affine (the SpecPipeSessionPtr precedent above). The ctx lives in
1996// `HybridModel::dspark_vgraphs` behind a Mutex and every touch happens on the engine's
1997// single decode-stream thread.
1998unsafe impl Send for DsparkVerifyGraphs {}
1999
2000impl DsparkVerifyGraphs {
2001 /// Build for this cache's shape. None when there are no linear layers, sizes are
2002 /// non-uniform, or the trunk keeps a gemma4 config (never on the qwen35 family).
2003 pub(crate) fn new(
2004 e: &Engine,
2005 cache: &Cache,
2006 t_max: usize,
2007 n_embd: usize,
2008 ) -> Result<Option<Self>, Box<dyn std::error::Error>> {
2009 let lin: Vec<usize> = (0..cache.recur.len())
2010 .filter(|&il| cache.recur[il].is_some())
2011 .collect();
2012 if lin.is_empty() || t_max < 2 {
2013 return Ok(None);
2014 }
2015 let first = cache.recur[lin[0]].as_ref().unwrap();
2016 let (conv_words, ssm_words) = (first.conv_state.len(), first.ssm_state.len());
2017 for &il in &lin {
2018 let rl = cache.recur[il].as_ref().unwrap();
2019 if rl.conv_state.len() != conv_words || rl.ssm_state.len() != ssm_words {
2020 return Ok(None);
2021 }
2022 }
2023 let n = lin.len();
2024 let mut lin_pos = std::collections::HashMap::with_capacity(n);
2025 for (k, &il) in lin.iter().enumerate() {
2026 lin_pos.insert(il, k);
2027 }
2028 // longest run of consecutive linear layers (save-scratch sizing)
2029 let mut max_run = 1usize;
2030 let mut run = 1usize;
2031 for w in lin.windows(2) {
2032 if w[1] == w[0] + 1 {
2033 run += 1;
2034 max_run = max_run.max(run);
2035 } else {
2036 run = 1;
2037 }
2038 }
2039 let rows = t_max - 1;
2040 let mut stash_conv = Vec::with_capacity(n);
2041 let mut stash_ssm = Vec::with_capacity(n);
2042 for _ in 0..n {
2043 stash_conv.push(e.uninit(rows * conv_words)?);
2044 stash_ssm.push(e.uninit(rows * ssm_words)?);
2045 }
2046 let host_table = vec![0u64; n * 6];
2047 let table_all = e.htod_u64(&host_table)?;
2048 // slice 4c: full-attention census for the full-verify graphs.
2049 let fa: Vec<usize> = (0..cache.kv.len())
2050 .filter(|&il| cache.kv[il].is_some())
2051 .collect();
2052 let mut fa_pos = std::collections::HashMap::with_capacity(fa.len());
2053 for (k, &il) in fa.iter().enumerate() {
2054 fa_pos.insert(il, k);
2055 }
2056 let n_layers = cache.kv.len().max(cache.recur.len());
2057 // exactly one of (linear state, kv cache) per layer — no MLA/unknown mixers.
2058 let walk_uniform = (0..n_layers).all(|il| {
2059 cache.recur.get(il).is_some_and(|r| r.is_some())
2060 != cache.kv.get(il).is_some_and(|k| k.is_some())
2061 });
2062 // Contiguous covered prefix: the largest n such that every layer in [0, n) is
2063 // linear or full-attention. The TRUNK walk is [0, layers.len()) and the cache
2064 // vecs can carry EXTRA state slots past it (the q38 export keeps the MTP head
2065 // layer's kv at the tail — hi == lin+fa never held, the s4c battery's zero
2066 // 'full' captures). The full-graph guard is walk coverage, not slot arithmetic.
2067 let covered = (0..n_layers)
2068 .take_while(|il| lin_pos.contains_key(il) || fa_pos.contains_key(il))
2069 .count();
2070 let t_cap = t_max;
2071 let fa_host_table = vec![0u64; fa.len() * 2 * t_cap];
2072 let fa_table = e.htod_u64(&fa_host_table)?;
2073 Ok(Some(Self {
2074 lin,
2075 lin_pos,
2076 table_all,
2077 host_table,
2078 stash_conv,
2079 stash_ssm,
2080 conv_words,
2081 ssm_words,
2082 stage: std::collections::HashMap::new(),
2083 tap_bufs: std::collections::HashMap::new(),
2084 graphs: std::collections::HashMap::new(),
2085 save_conv: e.uninit(n * conv_words)?,
2086 save_ssm: e.uninit(n * ssm_words)?,
2087 max_run,
2088 n_embd,
2089 round_slab: false,
2090 fa,
2091 fa_pos,
2092 fa_table,
2093 fa_host_table,
2094 t_cap,
2095 pos_stage: std::collections::HashMap::new(),
2096 full: std::collections::HashMap::new(),
2097 covered,
2098 walk_uniform,
2099 }))
2100 }
2101
2102 /// Rebuild the pointer tables from the live handles (once per verify — the gdn
2103 /// ping-pong swaps the canonical/alt handles between rounds; a fresh generation's
2104 /// cache buffers land at new addresses; a stale table would read the wrong state).
2105 pub(crate) fn refresh_tables(
2106 &mut self,
2107 e: &Engine,
2108 cache: &Cache,
2109 ) -> Result<(), Box<dyn std::error::Error>> {
2110 use cudarc::driver::DevicePtr;
2111 {
2112 let s = &e.gpu.stream();
2113 for (k, &il) in self.lin.iter().enumerate() {
2114 let rl = cache.recur[il].as_ref().unwrap();
2115 let (pc, _g0) = rl.conv_state.device_ptr(s);
2116 let (p0, _g1) = rl.ssm_state.device_ptr(s);
2117 let (p1, _g2) = rl.ssm_state_alt.device_ptr(s);
2118 let o = k * 6;
2119 self.host_table[o] = pc as u64;
2120 self.host_table[o + 1] = p0 as u64;
2121 self.host_table[o + 2] = p1 as u64;
2122 self.host_table[o + 3] = pc as u64;
2123 self.host_table[o + 4] = p1 as u64;
2124 self.host_table[o + 5] = p0 as u64;
2125 }
2126 for (k, &il) in self.fa.iter().enumerate() {
2127 let kvl = cache.kv[il].as_ref().unwrap();
2128 let (pk, _g0) = kvl.k.device_ptr(s);
2129 let (pv, _g1) = kvl.v.device_ptr(s);
2130 let o = k * 2 * self.t_cap;
2131 for z in 0..self.t_cap {
2132 self.fa_host_table[o + 2 * z] = pk as u64;
2133 self.fa_host_table[o + 2 * z + 1] = pv as u64;
2134 }
2135 }
2136 }
2137 e.htod_u64_into(&self.host_table, &mut self.table_all)?;
2138 if !self.fa_host_table.is_empty() {
2139 e.htod_u64_into(&self.fa_host_table, &mut self.fa_table)?;
2140 }
2141 Ok(())
2142 }
2143
2144 /// Slice 4c eligibility: Some(rung_end) when this round can replay (or capture) a
2145 /// full-verify graph — the whole walk [lo, hi) is covered, every layer is linear or
2146 /// full-attention, and ALL of the round's per-row t_kv values take the v4-seqs arm
2147 /// on ONE `fa_split_keys` ladder step that the rung also sits on (the straddle law;
2148 /// both gates are t_kv intervals, so ends-inside means all-inside). The rung is the
2149 /// round's next power of two — grid/partial sizing only (`n_splits_max` is pure
2150 /// stride; splits >= ns_eff write the empty partial the combine never reads), so one
2151 /// captured graph is bit-identical for every round the rung covers.
2152 #[allow(clippy::too_many_arguments)]
2153 pub(crate) fn full_rung(
2154 &self,
2155 model: &crate::hybrid::HybridModel,
2156 cache: &Cache,
2157 lo: usize,
2158 hi: usize,
2159 t: usize,
2160 seqs_arms_on: bool,
2161 ) -> Option<usize> {
2162 if std::env::var("MEMRA_DSPARK_FULLG_DEBUG").as_deref() == Ok("1") {
2163 static ONCE: std::sync::Once = std::sync::Once::new();
2164 let len0 = self
2165 .fa
2166 .first()
2167 .and_then(|&il| cache.kv[il].as_ref())
2168 .map(|k| k.len);
2169 ONCE.call_once(|| {
2170 eprintln!(
2171 "[fullg-debug] walk_uniform={} covered={} seqs_arms_on={} fa_rows_on={} t={} lo={} hi={} lin={} fa={} t_cap={} len0={:?}",
2172 self.walk_uniform, self.covered, seqs_arms_on, dspark_fa_rows_on(), t, lo, hi,
2173 self.lin.len(), self.fa.len(), self.t_cap, len0
2174 );
2175 });
2176 }
2177 if !self.walk_uniform
2178 || !seqs_arms_on
2179 || !dspark_fa_rows_on()
2180 || t < 2
2181 || lo != 0
2182 || hi > self.covered
2183 || t > self.t_cap
2184 || self.fa.is_empty()
2185 {
2186 return None;
2187 }
2188 let cfg = &model.cfg;
2189 let head_dim_global = cfg.head_dim_k as usize;
2190 let nkv = cfg.n_head_kv as usize;
2191 let kvl0 = cache.kv[self.fa[0]].as_ref().unwrap();
2192 // the z-batched twins read stacked rows at the cache's kv dims — must equal the
2193 // projection stride (the body's guard, hoisted so ineligible models fall back
2194 // instead of refusing mid-capture).
2195 let geom = cfg.full_attention_geometry_at(self.fa[0] as u32);
2196 let kv_dim = geom.n_head_kv as usize * geom.head_dim_k as usize;
2197 if kvl0.kv_dim_k != kv_dim || kvl0.kv_dim_v != kv_dim {
2198 return None;
2199 }
2200 let len0 = kvl0.len;
2201 let (t_kv_first, t_kv_last) = (len0 + 1, len0 + t);
2202 if !crate::fa_seqs_eligible(t_kv_first, head_dim_global)
2203 || !crate::fa_seqs_eligible(t_kv_last, head_dim_global)
2204 || crate::fa_split_keys(t_kv_first, nkv) != crate::fa_split_keys(t_kv_last, nkv)
2205 {
2206 return None;
2207 }
2208 let rung = t_kv_last.next_power_of_two().max(256);
2209 if crate::fa_split_keys(rung, nkv) != crate::fa_split_keys(t_kv_last, nkv) {
2210 return None;
2211 }
2212 Some(rung)
2213 }
2214
2215 /// Run the WHOLE verify walk [lo, hi) as one captured graph at (vt=t, rung): stage
2216 /// the residual + refresh the per-vt position staging, capture on first encounter
2217 /// (2 executing warmups bracketed by a full linear-state save/restore; KV warmup
2218 /// appends write the exact slots the replay writes — idempotent), launch, then apply
2219 /// the host bookkeeping the captured body skipped (per-linear-layer parity swap for
2220 /// odd t, per-fa-layer len bump). Returns the fresh residual.
2221 #[allow(clippy::too_many_arguments)]
2222 pub(crate) fn run_full(
2223 &mut self,
2224 model: &crate::hybrid::HybridModel,
2225 e: &Engine,
2226 lo: usize,
2227 hi: usize,
2228 x: &CudaSlice<f32>,
2229 t: usize,
2230 pos0: usize,
2231 rung: usize,
2232 cache: &mut Cache,
2233 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2234 let n_embd = self.n_embd;
2235 if !self.stage.contains_key(&t) {
2236 let xin = e.uninit(t * n_embd)?;
2237 let xout = e.uninit(t * n_embd)?;
2238 self.stage.insert(t, (xin, xout));
2239 }
2240 if !self.pos_stage.contains_key(&t) {
2241 self.pos_stage.insert(t, e.htod_i32(&vec![0i32; t])?);
2242 }
2243 // Per-round refresh: position contents + input staging (both addresses are baked
2244 // by the captured bodies; only their CONTENTS change round to round).
2245 {
2246 let pos_host: Vec<i32> = (0..t).map(|r| (pos0 + r) as i32).collect();
2247 let pb = self.pos_stage.get_mut(&t).unwrap();
2248 e.htod_i32_into(pb, &pos_host)?;
2249 let (xin, _) = self.stage.get_mut(&t).unwrap();
2250 e.copy_into(xin, 0, x, t * n_embd)?;
2251 }
2252 let key = (t, rung, hi);
2253 if !self.full.contains_key(&key) {
2254 // The warmups EXECUTE the whole walk on live state — save every linear
2255 // layer's conv + canonical ssm first, restore after (KV needs no restore:
2256 // graph mode never bumps host lens and the appends write this round's own
2257 // slots).
2258 for (k, &il) in self.lin.iter().enumerate() {
2259 let rl = cache.recur[il].as_ref().unwrap();
2260 e.copy_into(
2261 &mut self.save_conv,
2262 k * self.conv_words,
2263 &rl.conv_state,
2264 self.conv_words,
2265 )?;
2266 e.copy_into(
2267 &mut self.save_ssm,
2268 k * self.ssm_words,
2269 &rl.ssm_state,
2270 self.ssm_words,
2271 )?;
2272 }
2273 let (graph, keeper) = {
2274 let table_all = &self.table_all;
2275 let lin_pos = &self.lin_pos;
2276 let fa_pos = &self.fa_pos;
2277 let fa_table = &self.fa_table;
2278 let t_cap = self.t_cap;
2279 let stash_conv = &mut self.stash_conv;
2280 let stash_ssm = &mut self.stash_ssm;
2281 let pos_d: &CudaSlice<i32> = &self.pos_stage[&t];
2282 let (xin, xout) = self
2283 .stage
2284 .get_mut(&t)
2285 .map(|(a, b)| (&*a, b))
2286 .expect("stage bucket created above");
2287 let cache_ref: &mut Cache = cache;
2288 let iflag = if std::env::var("MEMRA_DSPARK_VG_AUTOFREE").as_deref() == Ok("1") {
2289 cudarc::driver::sys::CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH
2290 } else {
2291 cudarc::driver::sys::CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_USE_NODE_PRIORITY
2292 };
2293 e.capture_graph_retained_flags(iflag, move |e| {
2294 let mut xc: Option<CudaSlice<f32>> = None;
2295 for il in lo..hi {
2296 let xr: &CudaSlice<f32> = xc.as_ref().unwrap_or(xin);
2297 let nx = if let Some(&k) = lin_pos.get(&il) {
2298 model.qwen35_tparallel_linear_layer(
2299 e,
2300 il,
2301 xr,
2302 t,
2303 cache_ref,
2304 None,
2305 Some((&mut stash_conv[k], &mut stash_ssm[k])),
2306 Some((table_all, k * 6)),
2307 )?
2308 } else if let Some(&kf) = fa_pos.get(&il) {
2309 let mut no_rows: Option<Vec<CudaSlice<i32>>> = None;
2310 model.qwen35_tparallel_fa_layer(
2311 e,
2312 il,
2313 xr,
2314 t,
2315 cache_ref,
2316 FaLayerArgs {
2317 pos_d,
2318 pos_rows: &mut no_rows,
2319 pos0,
2320 seqs_append: true,
2321 batch_fa_on: true,
2322 graph_cap: Some((fa_table, kf * 2 * t_cap, rung)),
2323 stream: None,
2324 ckpt: None,
2325 },
2326 )?
2327 } else {
2328 return Err(format!(
2329 "run_full: layer {il} is neither linear nor full-attention"
2330 )
2331 .into());
2332 };
2333 xc = Some(nx);
2334 }
2335 e.copy_into(xout, 0, xc.as_ref().unwrap(), t * n_embd)?;
2336 Ok(())
2337 })?
2338 };
2339 // Undo the net host parity motion of the 3 body runs (each run swaps iff t
2340 // is odd -> 3 runs = net one swap), then restore the device state the
2341 // warmups consumed (walk scope only — layers past hi never executed). The
2342 // launch below then behaves exactly like one run.
2343 if t % 2 == 1 {
2344 for &il in &self.lin {
2345 if il < lo || il >= hi {
2346 continue;
2347 }
2348 let rl = cache.recur[il].as_mut().unwrap();
2349 std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
2350 }
2351 }
2352 for (k, &il) in self.lin.iter().enumerate() {
2353 if il < lo || il >= hi {
2354 continue;
2355 }
2356 let rl = cache.recur[il].as_mut().unwrap();
2357 let (cw, sw) = (self.conv_words, self.ssm_words);
2358 {
2359 let sv = e.view(&self.save_conv, self.lin.len() * cw);
2360 let win = sv.slice(k * cw..(k + 1) * cw);
2361 e.copy_view_into(&mut rl.conv_state, 0, &win, cw)?;
2362 }
2363 {
2364 let sv = e.view(&self.save_ssm, self.lin.len() * sw);
2365 let win = sv.slice(k * sw..(k + 1) * sw);
2366 e.copy_view_into(&mut rl.ssm_state, 0, &win, sw)?;
2367 }
2368 }
2369 if std::env::var("MEMRA_GRAPH_CENSUS").as_deref() == Ok("1") {
2370 if let Ok(c) = crate::graph_update::node_census(&graph) {
2371 eprintln!("[dspark-vg-census] full vt={t} rung={rung} {c:?}");
2372 }
2373 }
2374 self.full.insert(
2375 key,
2376 DsparkSegGraph {
2377 graph,
2378 _keeper: keeper,
2379 },
2380 );
2381 }
2382 self.full[&key].graph.launch()?;
2383 // Host bookkeeping for the replayed body (captured host code does not re-run):
2384 // gdn parity swap per linear layer (t odd), kv len bump per fa layer — scoped
2385 // to the WALK [lo, hi): the cache can carry extra state slots past it (the MTP
2386 // head layer's kv) that the walk never touches.
2387 if t % 2 == 1 {
2388 for &il in &self.lin {
2389 if il < lo || il >= hi {
2390 continue;
2391 }
2392 let rl = cache.recur[il].as_mut().unwrap();
2393 std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
2394 }
2395 }
2396 for &il in &self.fa {
2397 if il < lo || il >= hi {
2398 continue;
2399 }
2400 cache.kv[il].as_mut().unwrap().len += t;
2401 }
2402 let (_, xout) = self.stage.get(&t).unwrap();
2403 let mut out = e.uninit(t * n_embd)?;
2404 e.copy_into(&mut out, 0, xout, t * n_embd)?;
2405 Ok(out)
2406 }
2407
2408 /// Run layers [start, end) (all linear) as one captured graph at this vt: stage the
2409 /// residual into the bucket's x_in, capture on first encounter (2 executing warmups
2410 /// bracketed by a segment state save/restore), launch, then apply the host parity
2411 /// bookkeeping the captured body would have done. Returns the fresh residual.
2412 #[allow(clippy::too_many_arguments)]
2413 fn run_segment(
2414 &mut self,
2415 model: &crate::hybrid::HybridModel,
2416 e: &Engine,
2417 start: usize,
2418 end: usize,
2419 x: &CudaSlice<f32>,
2420 t: usize,
2421 cache: &mut Cache,
2422 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2423 let n_embd = self.n_embd;
2424 debug_assert!(end - start <= self.max_run);
2425 if !self.stage.contains_key(&t) {
2426 let xin = e.uninit(t * n_embd)?;
2427 let xout = e.uninit(t * n_embd)?;
2428 self.stage.insert(t, (xin, xout));
2429 }
2430 // Stage the residual at the bucket's baked input address.
2431 {
2432 let (xin, _) = self.stage.get_mut(&t).unwrap();
2433 e.copy_into(xin, 0, x, t * n_embd)?;
2434 }
2435 let key = (start, t);
2436 if !self.graphs.contains_key(&key) {
2437 // The 2 warmups EXECUTE the segment on live state — save conv + the canonical
2438 // ssm of every segment layer first, restore after, so the graph's first real
2439 // launch starts from the exact pre-round state (bytes gated e2e).
2440 for (k, il) in (start..end).enumerate() {
2441 let rl = cache.recur[il].as_ref().unwrap();
2442 e.copy_into(
2443 &mut self.save_conv,
2444 k * self.conv_words,
2445 &rl.conv_state,
2446 self.conv_words,
2447 )?;
2448 e.copy_into(
2449 &mut self.save_ssm,
2450 k * self.ssm_words,
2451 &rl.ssm_state,
2452 self.ssm_words,
2453 )?;
2454 }
2455 let (graph, keeper) = {
2456 let table_all = &self.table_all;
2457 let lin_pos = &self.lin_pos;
2458 let stash_conv = &mut self.stash_conv;
2459 let stash_ssm = &mut self.stash_ssm;
2460 let (xin, xout) = self
2461 .stage
2462 .get_mut(&t)
2463 .map(|(a, b)| (&*a, b))
2464 .expect("stage bucket created above");
2465 let cache_ref: &mut Cache = cache;
2466 // Slice 4 (fa-execupdate lane): USE_NODE_PRIORITY instead of
2467 // AUTO_FREE_ON_LAUNCH. The slice-3 measured limiter was AUTO_FREE's
2468 // launch-time mem-pool scan — 25.6 us per cuGraphLaunch x 16 segments
2469 // = ~0.41 ms/round, most of the eager-launch savings. The captured
2470 // body's cuMemAllocAsync transients are BALANCED by in-graph frees
2471 // (every transient drops inside the capture region — the generic
2472 // capture path's census precedent, 1589/1589), so AUTO_FREE has
2473 // nothing to reclaim and the graph is legal to instantiate without
2474 // it; PRIORITY is the flag the gemma slotted door ships for exactly
2475 // this reason (both alternatives drop the scan; UPLOAD via
2476 // cuGraphInstantiateWithFlags is WithParams-only and refused).
2477 // MEMRA_DSPARK_VG_AUTOFREE=1 reverts; MEMRA_GRAPH_CENSUS=1 prints
2478 // the node census at capture (the ALLOC==FREE receipt).
2479 let iflag = if std::env::var("MEMRA_DSPARK_VG_AUTOFREE").as_deref() == Ok("1") {
2480 cudarc::driver::sys::CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH
2481 } else {
2482 cudarc::driver::sys::CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_USE_NODE_PRIORITY
2483 };
2484 e.capture_graph_retained_flags(iflag, move |e| {
2485 let mut xc: Option<CudaSlice<f32>> = None;
2486 for il in start..end {
2487 let k = lin_pos[&il];
2488 let xr: &CudaSlice<f32> = xc.as_ref().unwrap_or(xin);
2489 let nx = model.qwen35_tparallel_linear_layer(
2490 e,
2491 il,
2492 xr,
2493 t,
2494 cache_ref,
2495 None,
2496 Some((&mut stash_conv[k], &mut stash_ssm[k])),
2497 Some((table_all, k * 6)),
2498 )?;
2499 xc = Some(nx);
2500 }
2501 e.copy_into(xout, 0, xc.as_ref().unwrap(), t * n_embd)?;
2502 Ok(())
2503 })?
2504 };
2505 // Undo the net host parity motion of the 3 body runs (each run swaps iff t
2506 // is odd -> 3 runs = net one swap), then restore the device state the
2507 // warmups consumed. The launch below then behaves exactly like one run.
2508 if t % 2 == 1 {
2509 for il in start..end {
2510 let rl = cache.recur[il].as_mut().unwrap();
2511 std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
2512 }
2513 }
2514 for (k, il) in (start..end).enumerate() {
2515 let rl = cache.recur[il].as_mut().unwrap();
2516 let (cw, sw) = (self.conv_words, self.ssm_words);
2517 {
2518 let sv = e.view(&self.save_conv, self.lin.len() * cw);
2519 let win = sv.slice(k * cw..(k + 1) * cw);
2520 e.copy_view_into(&mut rl.conv_state, 0, &win, cw)?;
2521 }
2522 {
2523 let sv = e.view(&self.save_ssm, self.lin.len() * sw);
2524 let win = sv.slice(k * sw..(k + 1) * sw);
2525 e.copy_view_into(&mut rl.ssm_state, 0, &win, sw)?;
2526 }
2527 }
2528 if std::env::var("MEMRA_GRAPH_CENSUS").as_deref() == Ok("1") {
2529 if let Ok(c) = crate::graph_update::node_census(&graph) {
2530 eprintln!("[dspark-vg-census] seg={start}..{end} vt={t} {c:?}");
2531 }
2532 }
2533 self.graphs.insert(
2534 key,
2535 DsparkSegGraph {
2536 graph,
2537 _keeper: keeper,
2538 },
2539 );
2540 }
2541 self.graphs[&key].graph.launch()?;
2542 // Host parity bookkeeping for the replayed body (the captured host swaps do not
2543 // re-run at replay).
2544 if t % 2 == 1 {
2545 for il in start..end {
2546 let rl = cache.recur[il].as_mut().unwrap();
2547 std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
2548 }
2549 }
2550 let (_, xout) = self.stage.get(&t).unwrap();
2551 let mut out = e.uninit(t * n_embd)?;
2552 e.copy_into(&mut out, 0, xout, t * n_embd)?;
2553 Ok(out)
2554 }
2555
2556 /// Pool freeze check (`dspark_vg_cap`): below the ceiling new keys may capture.
2557 fn can_capture(&self) -> bool {
2558 self.graphs.len() + self.full.len() < dspark_vg_cap()
2559 }
2560
2561 /// Round-atomic segment-door readiness: TRUE when this round's walk can ride the
2562 /// per-(segment, vt) graphs without a NEW capture past the pool ceiling — every
2563 /// linear run in [lo, hi) already has its (run_start, t) key, or capture is still
2564 /// allowed. FALSE sends the WHOLE round down the eager cols-ckpt walk: a partial
2565 /// refusal would stash some layers in the ctx slabs and others in the round's cols
2566 /// while one commit reads only one of them.
2567 pub(crate) fn segments_ready(
2568 &self,
2569 model: &crate::hybrid::HybridModel,
2570 lo: usize,
2571 hi: usize,
2572 t: usize,
2573 ) -> bool {
2574 if self.can_capture() {
2575 return true;
2576 }
2577 let mut il = lo;
2578 while il < hi {
2579 if matches!(model.layers[il].mixer, Mixer::Linear(_)) {
2580 let start = il;
2581 while il < hi && matches!(model.layers[il].mixer, Mixer::Linear(_)) {
2582 il += 1;
2583 }
2584 if !self.graphs.contains_key(&(start, t)) {
2585 return false;
2586 }
2587 } else {
2588 il += 1;
2589 }
2590 }
2591 true
2592 }
2593
2594 /// Widest verify window this pool was built for. A caller whose round exceeds it must
2595 /// take the eager walk: the stash slabs hold `t_capacity() - 1` column rows, and slicing
2596 /// past them is a panic rather than a refusal.
2597 pub(crate) fn t_capacity(&self) -> usize {
2598 self.t_cap
2599 }
2600
2601 /// Slab row (conv, ssm) device pointers + lengths for the commit restore of column
2602 /// `row` (0-based) of layer `il`. None for non-linear layers.
2603 pub(crate) fn slab_row(
2604 &self,
2605 e: &Engine,
2606 il: usize,
2607 row: usize,
2608 ) -> Option<(u64, u64, usize, usize)> {
2609 use cudarc::driver::DevicePtr;
2610 let k = *self.lin_pos.get(&il)?;
2611 let s = &e.gpu.stream();
2612 let (pc, _g0) = self.stash_conv[k].device_ptr(s);
2613 let (ps, _g1) = self.stash_ssm[k].device_ptr(s);
2614 Some((
2615 pc as u64 + (row * self.conv_words * 4) as u64,
2616 ps as u64 + (row * self.ssm_words * 4) as u64,
2617 self.conv_words,
2618 self.ssm_words,
2619 ))
2620 }
2621}
2622
2623impl VerifyCkpt {
2624 fn new(n_layer: usize) -> Self {
2625 VerifyCkpt {
2626 gdn: (0..n_layer).map(|_| None).collect(),
2627 cols: (0..n_layer).map(|_| None).collect(),
2628 }
2629 }
2630}
2631
2632/// The stage-0/TX half of one PP verify. The boundary slot is the ownership token: stage 1
2633/// consumes exactly the slot selected by `tx()` / `tx_pipelined()`, never a slot inferred from
2634/// a logical round number.
2635struct VerifyBoundaryTicket {
2636 rt: &'static crate::pp::PpNRt,
2637 caller_stream: std::sync::Arc<cudarc::driver::CudaStream>,
2638 slot: usize,
2639 pos0: usize,
2640 t: usize,
2641 payload: usize,
2642 n_st: usize,
2643 pipelined: bool,
2644 pp_anatomy: bool,
2645 pp_started: std::time::Instant,
2646 reverse_ms: f64,
2647 stage0_ms: f64,
2648 tx_ms: f64,
2649 trace: Option<SpecPipeTraceCtx>,
2650}
2651
2652/// Explicit OPTIPIPE diagnostic control. Forced modes are set only by `optipipe-gate`; the
2653/// increment-2 controller can also be armed by the server's fresh-process research door.
2654#[derive(Clone, Copy, Debug, PartialEq, Eq)]
2655pub enum OptiForkGateMode {
2656 Disabled,
2657 Hit,
2658 Miss,
2659 Alternate,
2660 Abort,
2661 Controller,
2662}
2663
2664static OPTI_FORK_GATE_MODE: std::sync::atomic::AtomicU8 = std::sync::atomic::AtomicU8::new(0);
2665static OPTI_CONTROLLER_THRESHOLD: std::sync::atomic::AtomicU32 =
2666 std::sync::atomic::AtomicU32::new(0);
2667static OPTI_FORK_ATTEMPTS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
2668static OPTI_FORK_HITS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
2669static OPTI_FORK_MISSES: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
2670static OPTI_FORK_ABORT_DRAINS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
2671static OPTI_FORK_REFUSALS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
2672static OPTI_GATE_CHECKS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
2673static OPTI_GATE_ADMITS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
2674static OPTI_GATE_REJECTS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
2675static OPTI_RECONCILES: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
2676static OPTI_WASTED_DRAFT_TOKENS: std::sync::atomic::AtomicU64 =
2677 std::sync::atomic::AtomicU64::new(0);
2678static OPTI_SHADOW_DRAFT_TOKENS: std::sync::atomic::AtomicU64 =
2679 std::sync::atomic::AtomicU64::new(0);
2680static OPTI_BREAKER_TRIPS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
2681
2682impl OptiForkGateMode {
2683 fn code(self) -> u8 {
2684 match self {
2685 Self::Disabled => 0,
2686 Self::Hit => 1,
2687 Self::Miss => 2,
2688 Self::Alternate => 3,
2689 Self::Abort => 4,
2690 Self::Controller => 5,
2691 }
2692 }
2693
2694 fn configured() -> Self {
2695 match OPTI_FORK_GATE_MODE.load(std::sync::atomic::Ordering::Relaxed) {
2696 1 => Self::Hit,
2697 2 => Self::Miss,
2698 3 => Self::Alternate,
2699 4 => Self::Abort,
2700 5 => Self::Controller,
2701 _ => Self::Disabled,
2702 }
2703 }
2704
2705 fn action(self, generation: u64) -> OptiForkAction {
2706 match self {
2707 Self::Hit => OptiForkAction::Hit,
2708 Self::Miss => OptiForkAction::Miss,
2709 Self::Alternate if generation & 1 == 0 => OptiForkAction::Hit,
2710 Self::Alternate => OptiForkAction::Miss,
2711 Self::Abort => OptiForkAction::Abort,
2712 Self::Disabled | Self::Controller => {
2713 unreachable!("non-forced mode cannot choose a forced fork action")
2714 }
2715 }
2716 }
2717
2718 fn is_forced(self) -> bool {
2719 matches!(self, Self::Hit | Self::Miss | Self::Alternate | Self::Abort)
2720 }
2721}
2722
2723/// Arm or disarm the forced harness. Serving uses only `set_optipipe_controller_threshold`.
2724pub fn set_optipipe_gate_mode(mode: OptiForkGateMode) {
2725 OPTI_FORK_GATE_MODE.store(mode.code(), std::sync::atomic::Ordering::Relaxed);
2726}
2727
2728/// Arm the increment-2 diagnostic controller. The threshold applies to the uncalibrated
2729/// two-token draft-probability product. Serving can call this only through its explicit
2730/// fresh-process research door; the absent-door default remains byte-for-byte disabled.
2731pub fn set_optipipe_controller_threshold(threshold: f32) {
2732 assert!(threshold.is_finite() && (0.0..=1.0).contains(&threshold));
2733 OPTI_CONTROLLER_THRESHOLD.store(threshold.to_bits(), std::sync::atomic::Ordering::Relaxed);
2734 set_optipipe_gate_mode(OptiForkGateMode::Controller);
2735}
2736
2737#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
2738pub struct OptiForkGateStats {
2739 pub attempts: u64,
2740 pub hits: u64,
2741 pub misses: u64,
2742 pub abort_drains: u64,
2743 pub refusals: u64,
2744 pub gate_checks: u64,
2745 pub gate_admits: u64,
2746 pub gate_rejects: u64,
2747 pub reconciles: u64,
2748 pub wasted_draft_tokens: u64,
2749 pub shadow_draft_tokens: u64,
2750 pub breaker_trips: u64,
2751}
2752
2753#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
2754pub struct OptiForkStateIdentity {
2755 pub trunk_kv_bytes: usize,
2756 pub recurrent_bytes: usize,
2757 pub scratch_kv_bytes: usize,
2758 pub hidden_bytes: usize,
2759}
2760
2761pub fn reset_optipipe_gate_stats() {
2762 for counter in [
2763 &OPTI_FORK_ATTEMPTS,
2764 &OPTI_FORK_HITS,
2765 &OPTI_FORK_MISSES,
2766 &OPTI_FORK_ABORT_DRAINS,
2767 &OPTI_FORK_REFUSALS,
2768 &OPTI_GATE_CHECKS,
2769 &OPTI_GATE_ADMITS,
2770 &OPTI_GATE_REJECTS,
2771 &OPTI_RECONCILES,
2772 &OPTI_WASTED_DRAFT_TOKENS,
2773 &OPTI_SHADOW_DRAFT_TOKENS,
2774 &OPTI_BREAKER_TRIPS,
2775 ] {
2776 counter.store(0, std::sync::atomic::Ordering::Relaxed);
2777 }
2778}
2779
2780pub fn optipipe_gate_stats() -> OptiForkGateStats {
2781 let load = |v: &std::sync::atomic::AtomicU64| v.load(std::sync::atomic::Ordering::Relaxed);
2782 OptiForkGateStats {
2783 attempts: load(&OPTI_FORK_ATTEMPTS),
2784 hits: load(&OPTI_FORK_HITS),
2785 misses: load(&OPTI_FORK_MISSES),
2786 abort_drains: load(&OPTI_FORK_ABORT_DRAINS),
2787 refusals: load(&OPTI_FORK_REFUSALS),
2788 gate_checks: load(&OPTI_GATE_CHECKS),
2789 gate_admits: load(&OPTI_GATE_ADMITS),
2790 gate_rejects: load(&OPTI_GATE_REJECTS),
2791 reconciles: load(&OPTI_RECONCILES),
2792 wasted_draft_tokens: load(&OPTI_WASTED_DRAFT_TOKENS),
2793 shadow_draft_tokens: load(&OPTI_SHADOW_DRAFT_TOKENS),
2794 breaker_trips: load(&OPTI_BREAKER_TRIPS),
2795 }
2796}
2797
2798#[derive(Clone, Copy, Debug)]
2799struct OptiControllerPolicy {
2800 threshold: f32,
2801 consecutive_misses: u8,
2802 breaker_tripped: bool,
2803}
2804
2805impl OptiControllerPolicy {
2806 fn configured() -> Self {
2807 Self {
2808 threshold: f32::from_bits(
2809 OPTI_CONTROLLER_THRESHOLD.load(std::sync::atomic::Ordering::Relaxed),
2810 ),
2811 consecutive_misses: 0,
2812 breaker_tripped: false,
2813 }
2814 }
2815
2816 fn admit(&self, q_proxy: f32) -> bool {
2817 q_proxy.is_finite()
2818 && (0.0..=1.0).contains(&q_proxy)
2819 && (self.threshold == 0.0 || (!self.breaker_tripped && q_proxy >= self.threshold))
2820 }
2821
2822 /// Returns true exactly when this resolution newly trips the three-miss breaker.
2823 fn resolve(&mut self, hit: bool) -> bool {
2824 // q*=0 is the lane's explicit unconditional measurement arm. Its purpose is to price
2825 // every optimistic opportunity, so the safety breaker is measured separately and must
2826 // not silently turn this arm into "three attempts then serial".
2827 if self.threshold == 0.0 {
2828 self.consecutive_misses = 0;
2829 return false;
2830 }
2831 if hit {
2832 self.consecutive_misses = 0;
2833 return false;
2834 }
2835 self.consecutive_misses = self.consecutive_misses.saturating_add(1);
2836 if !self.breaker_tripped && self.consecutive_misses >= 3 {
2837 self.breaker_tripped = true;
2838 return true;
2839 }
2840 false
2841 }
2842}
2843
2844#[derive(Clone, Copy, Debug, PartialEq, Eq)]
2845enum OptiForkAction {
2846 Hit,
2847 Miss,
2848 Abort,
2849}
2850
2851#[derive(Clone, Copy, Debug, PartialEq, Eq)]
2852struct OptiForkGeneration {
2853 id: u64,
2854 slot: usize,
2855}
2856
2857#[derive(Default)]
2858struct OptiForkGenerationTracker {
2859 next: u64,
2860 live: [Option<u64>; 2],
2861}
2862
2863impl OptiForkGenerationTracker {
2864 fn reserve(&mut self) -> Result<OptiForkGeneration, Box<dyn std::error::Error>> {
2865 let generation = OptiForkGeneration {
2866 id: self.next,
2867 slot: (self.next & 1) as usize,
2868 };
2869 if let Some(live) = self.live[generation.slot] {
2870 return Err(format!(
2871 "optipipe snapshot slot {} still owns generation {live}; refusing to overwrite it",
2872 generation.slot,
2873 )
2874 .into());
2875 }
2876 self.next += 1;
2877 self.live[generation.slot] = Some(generation.id);
2878 Ok(generation)
2879 }
2880
2881 fn retire(&mut self, generation: OptiForkGeneration) -> Result<(), Box<dyn std::error::Error>> {
2882 match self.live[generation.slot] {
2883 Some(id) if id == generation.id => {
2884 self.live[generation.slot] = None;
2885 Ok(())
2886 }
2887 other => Err(format!(
2888 "optipipe generation teardown mismatch: ticket={} slot={} live={other:?}",
2889 generation.id, generation.slot,
2890 )
2891 .into()),
2892 }
2893 }
2894}
2895
2896struct OptiForkSeedGeneration {
2897 h_seed: CudaSlice<f32>,
2898 fill_prev: CudaSlice<f32>,
2899 scratch_len: usize,
2900}
2901
2902/// Allocate or refresh one full checkpoint through the engine that owns each PP stage. The
2903/// generic cache helper accepts one device and therefore cannot copy GDN state split across
2904/// devices. KV lengths and position stay host metadata; only recurrent buffers need stage-local
2905/// device ownership.
2906fn opti_snapshot_stage_owned(
2907 e: &Engine,
2908 cache: &Cache,
2909 rt: &'static crate::pp::PpNRt,
2910 fence: &[usize],
2911) -> Result<crate::cache::CacheSnapshot, Box<dyn std::error::Error>> {
2912 let n = cache.kv.len();
2913 let mut snapshot = crate::cache::CacheSnapshot {
2914 kv_len: vec![None; n],
2915 tp_kv_len: vec![None; n],
2916 conv: (0..n).map(|_| None).collect(),
2917 ssm: (0..n).map(|_| None).collect(),
2918 pos: cache.pos,
2919 };
2920 opti_snapshot_stage_owned_into(e, cache, rt, fence, &mut snapshot)?;
2921 Ok(snapshot)
2922}
2923
2924fn opti_snapshot_stage_owned_into(
2925 e: &Engine,
2926 cache: &Cache,
2927 rt: &'static crate::pp::PpNRt,
2928 fence: &[usize],
2929 snapshot: &mut crate::cache::CacheSnapshot,
2930) -> Result<(), Box<dyn std::error::Error>> {
2931 if fence.len() != rt.n_stages() + 1
2932 || snapshot.kv_len.len() != cache.kv.len()
2933 || snapshot.tp_kv_len.len() != cache.tp_kv.len()
2934 {
2935 return Err("optipipe stage-owned snapshot shape mismatch".into());
2936 }
2937 for stage in 0..rt.n_stages() {
2938 opti_snapshot_one_stage_owned_into(e, cache, rt, fence, stage, snapshot)?;
2939 }
2940 snapshot.pos = cache.pos;
2941 Ok(())
2942}
2943
2944/// Refresh one PP stage of a checkpoint. Increment 2 uses this split form so stage 0's
2945/// optimistic post-N state is captured before N+1 stage 0 is queued, while stage 1's matching
2946/// post-N state is captured only after N stage 1 is enqueued. Calling the all-stage helper at
2947/// either point would capture one side of the fork at the wrong generation.
2948fn opti_snapshot_one_stage_owned_into(
2949 e: &Engine,
2950 cache: &Cache,
2951 rt: &'static crate::pp::PpNRt,
2952 fence: &[usize],
2953 stage: usize,
2954 snapshot: &mut crate::cache::CacheSnapshot,
2955) -> Result<(), Box<dyn std::error::Error>> {
2956 if fence.len() != rt.n_stages() + 1
2957 || snapshot.kv_len.len() != cache.kv.len()
2958 || snapshot.tp_kv_len.len() != cache.tp_kv.len()
2959 || stage >= rt.n_stages()
2960 {
2961 return Err("optipipe single-stage snapshot shape mismatch".into());
2962 }
2963 let _scope = rt.enter(stage);
2964 let owner = rt.engine(stage, e);
2965 for il in fence[stage]..fence[stage + 1] {
2966 snapshot.kv_len[il] = cache.kv[il].as_ref().map(|kv| kv.len);
2967 snapshot.tp_kv_len[il] = cache.tp_kv[il]
2968 .as_ref()
2969 .map(crate::tp::ResidentTpKvCache::committed_len);
2970 match &cache.recur[il] {
2971 Some(recur) => {
2972 match snapshot.conv[il].as_mut() {
2973 Some(dst) => {
2974 owner.copy_into(dst, 0, &recur.conv_state, recur.conv_state.len())?
2975 }
2976 None => snapshot.conv[il] = Some(owner.clone_dtod(&recur.conv_state)?),
2977 }
2978 match snapshot.ssm[il].as_mut() {
2979 Some(dst) => {
2980 owner.copy_into(dst, 0, &recur.ssm_state, recur.ssm_state.len())?
2981 }
2982 None => snapshot.ssm[il] = Some(owner.clone_dtod(&recur.ssm_state)?),
2983 }
2984 }
2985 None if snapshot.conv[il].is_some() || snapshot.ssm[il].is_some() => {
2986 return Err(
2987 format!("optipipe stage-owned snapshot layer {il} changed shape").into(),
2988 );
2989 }
2990 None => {}
2991 }
2992 }
2993 snapshot.pos = cache.pos;
2994 Ok(())
2995}
2996
2997/// Increment-1 persistent fork state. Exactly two snapshot/seed slots alternate; a live ticket
2998/// names its generation and keeps teardown fail-closed. Only stage 0 is allowed to mutate before
2999/// resolve, so the reconcile tables and conditional restores are stage-local.
3000struct OptiForkState {
3001 mode: OptiForkGateMode,
3002 controller: Option<OptiControllerPolicy>,
3003 generations: OptiForkGenerationTracker,
3004 active_snapshot_slot: usize,
3005 alternate_snapshot: crate::cache::CacheSnapshot,
3006 seeds: [OptiForkSeedGeneration; 2],
3007 rt: &'static crate::pp::PpNRt,
3008 fence: [usize; 3],
3009 split: usize,
3010 len_ptrs: CudaSlice<u64>,
3011 saved_lens: CudaSlice<i32>,
3012 forced_acc: CudaSlice<u32>,
3013 valid: CudaSlice<u32>,
3014 stage0_stream: std::sync::Arc<cudarc::driver::CudaStream>,
3015 logical_payload_bytes: [usize; 2],
3016}
3017
3018struct OptiForkTicket {
3019 generation: OptiForkGeneration,
3020 boundary: Option<VerifyBoundaryTicket>,
3021 drain: std::sync::Arc<cudarc::driver::CudaStream>,
3022 settled: bool,
3023}
3024
3025struct OptiControllerTicket {
3026 generation: OptiForkGeneration,
3027 boundary: Option<VerifyBoundaryTicket>,
3028 ckpt: Option<VerifyCkpt>,
3029 verify_tokens: [u32; 2],
3030 draft_prob: f32,
3031 eager_seed: Option<CudaSlice<f32>>,
3032 q_proxy: f32,
3033 scratch_len: usize,
3034 issued_at: std::time::Instant,
3035 drain: std::sync::Arc<cudarc::driver::CudaStream>,
3036 settled: bool,
3037}
3038
3039struct OptiControllerPrepared {
3040 verify_tokens: [u32; 2],
3041 draft_prob: f32,
3042 eager_seed: Option<CudaSlice<f32>>,
3043 q_proxy: f32,
3044 scratch_len: usize,
3045}
3046
3047impl OptiControllerTicket {
3048 fn take_boundary(&mut self) -> VerifyBoundaryTicket {
3049 self.boundary
3050 .take()
3051 .expect("controller boundary ticket already consumed")
3052 }
3053
3054 fn take_ckpt(&mut self) -> VerifyCkpt {
3055 self.ckpt
3056 .take()
3057 .expect("controller verify checkpoint already consumed")
3058 }
3059
3060 fn take_eager_seed(&mut self) -> Option<CudaSlice<f32>> {
3061 self.eager_seed.take()
3062 }
3063
3064 fn settle(&mut self) {
3065 self.settled = true;
3066 }
3067}
3068
3069impl Drop for OptiControllerTicket {
3070 fn drop(&mut self) {
3071 if !self.settled {
3072 let _ = self.drain.synchronize();
3073 OPTI_FORK_ABORT_DRAINS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
3074 }
3075 }
3076}
3077
3078impl OptiForkTicket {
3079 fn take_boundary(&mut self) -> VerifyBoundaryTicket {
3080 self.boundary
3081 .take()
3082 .expect("fork ticket boundary already consumed")
3083 }
3084
3085 fn settle(&mut self) {
3086 self.settled = true;
3087 }
3088}
3089
3090impl Drop for OptiForkTicket {
3091 fn drop(&mut self) {
3092 if !self.settled {
3093 let _ = self.drain.synchronize();
3094 OPTI_FORK_ABORT_DRAINS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
3095 }
3096 }
3097}
3098
3099impl OptiForkState {
3100 #[allow(clippy::too_many_arguments)]
3101 fn new(
3102 e: &Engine,
3103 cache: &Cache,
3104 mode: OptiForkGateMode,
3105 alternate_snapshot: crate::cache::CacheSnapshot,
3106 h_seed: &CudaSlice<f32>,
3107 fill_prev: &CudaSlice<f32>,
3108 rt: &'static crate::pp::PpNRt,
3109 split: usize,
3110 n_layer: usize,
3111 ) -> Result<Self, Box<dyn std::error::Error>> {
3112 let fence = [0, split, n_layer];
3113 let mut logical_payload_bytes = [0usize; 2];
3114 for stage in 0..2 {
3115 for il in fence[stage]..fence[stage + 1] {
3116 logical_payload_bytes[stage] += alternate_snapshot.conv[il]
3117 .as_ref()
3118 .map_or(0, |v| v.len() * std::mem::size_of::<f32>());
3119 logical_payload_bytes[stage] += alternate_snapshot.ssm[il]
3120 .as_ref()
3121 .map_or(0, |v| v.len() * std::mem::size_of::<f32>());
3122 }
3123 }
3124 let seeds = [
3125 OptiForkSeedGeneration {
3126 h_seed: e.clone_dtod(h_seed)?,
3127 fill_prev: e.clone_dtod(fill_prev)?,
3128 scratch_len: 0,
3129 },
3130 OptiForkSeedGeneration {
3131 h_seed: e.clone_dtod(h_seed)?,
3132 fill_prev: e.clone_dtod(fill_prev)?,
3133 scratch_len: 0,
3134 },
3135 ];
3136 let (len_ptrs, saved_lens, forced_acc, valid, stage0_stream) = {
3137 let _stage = rt.enter(0);
3138 let e0 = rt.engine(0, e);
3139 (
3140 crate::round_stream::kv_len_ptr_table_range(e0, cache, 0..split, None)?,
3141 e0.htod_i32(&vec![0; split])?,
3142 e0.alloc_u32_zeroed(2)?,
3143 e0.alloc_u32_zeroed(1)?,
3144 e0.stream(),
3145 )
3146 };
3147 logical_payload_bytes[0] += seeds
3148 .iter()
3149 .map(|seed| (seed.h_seed.len() + seed.fill_prev.len()) * std::mem::size_of::<f32>())
3150 .sum::<usize>();
3151 logical_payload_bytes[0] += len_ptrs.len() * std::mem::size_of::<u64>()
3152 + saved_lens.len() * std::mem::size_of::<i32>()
3153 + forced_acc.len() * std::mem::size_of::<u32>()
3154 + valid.len() * std::mem::size_of::<u32>();
3155 Ok(Self {
3156 mode,
3157 controller: (mode == OptiForkGateMode::Controller)
3158 .then(OptiControllerPolicy::configured),
3159 generations: OptiForkGenerationTracker::default(),
3160 active_snapshot_slot: 0,
3161 alternate_snapshot,
3162 seeds,
3163 rt,
3164 fence,
3165 split,
3166 len_ptrs,
3167 saved_lens,
3168 forced_acc,
3169 valid,
3170 stage0_stream,
3171 logical_payload_bytes,
3172 })
3173 }
3174
3175 fn reserve(
3176 &mut self,
3177 current_snapshot: &mut crate::cache::CacheSnapshot,
3178 ) -> Result<OptiForkGeneration, Box<dyn std::error::Error>> {
3179 let generation = self.generations.reserve()?;
3180 if generation.slot != self.active_snapshot_slot {
3181 std::mem::swap(current_snapshot, &mut self.alternate_snapshot);
3182 self.active_snapshot_slot = generation.slot;
3183 }
3184 Ok(generation)
3185 }
3186
3187 fn capture_seed(
3188 &mut self,
3189 e: &Engine,
3190 generation: OptiForkGeneration,
3191 h_seed: &CudaSlice<f32>,
3192 fill_prev: &CudaSlice<f32>,
3193 scratch_len: usize,
3194 ) -> Result<(), Box<dyn std::error::Error>> {
3195 let seed = &mut self.seeds[generation.slot];
3196 e.copy_into(&mut seed.h_seed, 0, h_seed, h_seed.len())?;
3197 e.copy_into(&mut seed.fill_prev, 0, fill_prev, fill_prev.len())?;
3198 seed.scratch_len = scratch_len;
3199 Ok(())
3200 }
3201
3202 fn ticket(
3203 &self,
3204 generation: OptiForkGeneration,
3205 boundary: VerifyBoundaryTicket,
3206 ) -> OptiForkTicket {
3207 OptiForkTicket {
3208 generation,
3209 boundary: Some(boundary),
3210 drain: self.stage0_stream.clone(),
3211 settled: false,
3212 }
3213 }
3214
3215 #[allow(clippy::too_many_arguments)]
3216 fn controller_ticket(
3217 &self,
3218 generation: OptiForkGeneration,
3219 boundary: VerifyBoundaryTicket,
3220 ckpt: VerifyCkpt,
3221 verify_tokens: [u32; 2],
3222 draft_prob: f32,
3223 eager_seed: Option<CudaSlice<f32>>,
3224 q_proxy: f32,
3225 scratch_len: usize,
3226 ) -> OptiControllerTicket {
3227 OptiControllerTicket {
3228 generation,
3229 boundary: Some(boundary),
3230 ckpt: Some(ckpt),
3231 verify_tokens,
3232 draft_prob,
3233 eager_seed,
3234 q_proxy,
3235 scratch_len,
3236 issued_at: std::time::Instant::now(),
3237 drain: self.stage0_stream.clone(),
3238 settled: false,
3239 }
3240 }
3241
3242 fn reserve_successor(&mut self) -> Result<OptiForkGeneration, Box<dyn std::error::Error>> {
3243 self.generations.reserve()
3244 }
3245
3246 fn successor_snapshot_mut(&mut self) -> &mut crate::cache::CacheSnapshot {
3247 &mut self.alternate_snapshot
3248 }
3249
3250 fn promote_successor_snapshot(
3251 &mut self,
3252 current_snapshot: &mut crate::cache::CacheSnapshot,
3253 generation: OptiForkGeneration,
3254 ) {
3255 std::mem::swap(current_snapshot, &mut self.alternate_snapshot);
3256 self.active_snapshot_slot = generation.slot;
3257 }
3258
3259 fn queue_actual_reconcile(
3260 &mut self,
3261 e: &Engine,
3262 snapshot: &crate::cache::CacheSnapshot,
3263 acc: &CudaSlice<u32>,
3264 optimistic_pending: u32,
3265 base: usize,
3266 ) -> Result<(), Box<dyn std::error::Error>> {
3267 let saved: Vec<i32> = (0..self.split)
3268 .map(|il| snapshot.kv_len[il].map(|v| v as i32).unwrap_or(0))
3269 .collect();
3270 // Serving keeps the caller/accept walk on the head (stage-1) device. Record the accept
3271 // decision point there and append a wait to stage 0 after its optimistic successor/TX;
3272 // the validity/reconcile kernels must never peer-read acc before it is written. The
3273 // increment-1 harness uses primary stage 0, where stream order already provides this.
3274 if self.rt.engine(0, e).ctx().ordinal() != e.ctx().ordinal() {
3275 self.rt.fence_stages_behind(&e.stream())?;
3276 }
3277 let _stage = self.rt.enter(0);
3278 let e0 = self.rt.engine(0, e);
3279 e0.htod_i32_into(&mut self.saved_lens, &saved)?;
3280 e0.spec_fork_valid(acc, optimistic_pending, &mut self.valid)?;
3281 e0.spec_fork_reconcile_kv(
3282 &self.len_ptrs,
3283 &self.saved_lens,
3284 acc,
3285 &self.valid,
3286 base,
3287 self.split,
3288 )
3289 }
3290
3291 fn finish_actual_reconcile(
3292 &mut self,
3293 e: &Engine,
3294 cache: &mut Cache,
3295 snapshot: &crate::cache::CacheSnapshot,
3296 n_acc: usize,
3297 base: usize,
3298 hit: bool,
3299 ) -> Result<(), Box<dyn std::error::Error>> {
3300 if hit {
3301 return Ok(());
3302 }
3303 let len_delta = base + n_acc;
3304 for il in 0..self.split {
3305 if let (Some(kv), Some(saved)) = (cache.kv[il].as_mut(), snapshot.kv_len[il]) {
3306 kv.len = saved + len_delta;
3307 }
3308 }
3309 {
3310 let _stage = self.rt.enter(1);
3311 let e1 = self.rt.engine(1, e);
3312 for il in self.split..self.fence[2] {
3313 if let (Some(kv), Some(saved)) = (cache.kv[il].as_mut(), snapshot.kv_len[il]) {
3314 kv.len = saved + len_delta;
3315 e1.set_i32_one(&mut kv.len_d, kv.len as i32)?;
3316 }
3317 }
3318 }
3319 self.rt.publish_to(0, &e.stream())?;
3320 Ok(())
3321 }
3322
3323 fn cancel_controller_ticket(
3324 &mut self,
3325 e: &Engine,
3326 cache: &mut Cache,
3327 scratch: &mut MtpScratch,
3328 snapshot: &crate::cache::CacheSnapshot,
3329 ticket: &mut OptiControllerTicket,
3330 ) -> Result<(), Box<dyn std::error::Error>> {
3331 {
3332 let _stage = self.rt.enter(0);
3333 let e0 = self.rt.engine(0, e);
3334 for il in 0..self.split {
3335 if let (Some(kv), Some(saved)) = (cache.kv[il].as_mut(), snapshot.kv_len[il]) {
3336 kv.len = saved;
3337 e0.set_i32_one(&mut kv.len_d, saved as i32)?;
3338 }
3339 }
3340 }
3341 scratch.set_len(e, snapshot.pos)?;
3342 ticket.settle();
3343 self.generations.retire(ticket.generation)?;
3344 OPTI_FORK_ABORT_DRAINS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
3345 OPTI_WASTED_DRAFT_TOKENS.fetch_add(2, std::sync::atomic::Ordering::Relaxed);
3346 eprintln!(
3347 "[opti-controller] tail-drain generation={} slot={}",
3348 ticket.generation.id, ticket.generation.slot,
3349 );
3350 Ok(())
3351 }
3352
3353 #[allow(clippy::too_many_arguments)]
3354 fn reconcile(
3355 &mut self,
3356 e: &Engine,
3357 cache: &mut Cache,
3358 scratch: &mut MtpScratch,
3359 snapshot: &crate::cache::CacheSnapshot,
3360 h_seed: &mut CudaSlice<f32>,
3361 fill_prev: &mut CudaSlice<f32>,
3362 generation: OptiForkGeneration,
3363 action: OptiForkAction,
3364 optimistic_pending: u32,
3365 ) -> Result<(), Box<dyn std::error::Error>> {
3366 debug_assert!(action != OptiForkAction::Abort);
3367 let miss_started = std::time::Instant::now();
3368 let keep = action == OptiForkAction::Hit;
3369 let saved: Vec<i32> = (0..self.split)
3370 .map(|il| snapshot.kv_len[il].map(|v| v as i32).unwrap_or(0))
3371 .collect();
3372 let seed = &self.seeds[generation.slot];
3373 {
3374 let _stage = self.rt.enter(0);
3375 let e0 = self.rt.engine(0, e);
3376 e0.htod_i32_into(&mut self.saved_lens, &saved)?;
3377 let forced = if keep {
3378 [1u32, optimistic_pending]
3379 } else {
3380 [0u32, optimistic_pending]
3381 };
3382 e0.htod_u32_into(&mut self.forced_acc, &forced)?;
3383 e0.spec_fork_valid(&self.forced_acc, optimistic_pending, &mut self.valid)?;
3384 e0.spec_fork_reconcile_kv(
3385 &self.len_ptrs,
3386 &self.saved_lens,
3387 &self.forced_acc,
3388 &self.valid,
3389 0,
3390 self.split,
3391 )?;
3392 for il in 0..self.split {
3393 if let Some(recur) = cache.recur[il].as_mut() {
3394 let conv = snapshot.conv[il]
3395 .as_ref()
3396 .ok_or("optipipe stage0 snapshot missing conv state")?;
3397 let ssm = snapshot.ssm[il]
3398 .as_ref()
3399 .ok_or("optipipe stage0 snapshot missing ssm state")?;
3400 e0.spec_fork_restore_f32(conv, &mut recur.conv_state, &self.valid)?;
3401 e0.spec_fork_restore_f32(ssm, &mut recur.ssm_state, &self.valid)?;
3402 }
3403 }
3404 e0.spec_fork_restore_f32(&seed.h_seed, h_seed, &self.valid)?;
3405 e0.spec_fork_restore_f32(&seed.fill_prev, fill_prev, &self.valid)?;
3406 }
3407
3408 if keep {
3409 OPTI_FORK_HITS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
3410 return Ok(());
3411 }
3412
3413 for il in 0..self.split {
3414 if let (Some(kv), Some(saved)) = (cache.kv[il].as_mut(), snapshot.kv_len[il]) {
3415 kv.len = saved;
3416 }
3417 }
3418 scratch.set_len(e, seed.scratch_len)?;
3419 // Targeted E_restart: publish only stage 0's reconcile to the caller, then bound the
3420 // forced diagnostic so the retained number is the actual miss cost, not enqueue time.
3421 let caller = e.stream();
3422 self.rt.publish_to(0, &caller)?;
3423 caller.synchronize()?;
3424 let miss_ms = miss_started.elapsed().as_secs_f64() * 1e3;
3425 eprintln!(
3426 "[opti-fork-reconcile] generation={} slot={} miss_ms={miss_ms:.3}",
3427 generation.id, generation.slot,
3428 );
3429 OPTI_FORK_MISSES.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
3430 Ok(())
3431 }
3432
3433 fn retire(&mut self, generation: OptiForkGeneration) -> Result<(), Box<dyn std::error::Error>> {
3434 self.generations.retire(generation)
3435 }
3436}
3437
3438fn rewind_tp_kv_verified_prefix(
3439 tp_kv: &mut [Option<crate::tp::ResidentTpKvCache>],
3440 saved_lens: &[Option<usize>],
3441 accepted: usize,
3442) -> Result<(), Box<dyn std::error::Error>> {
3443 if tp_kv.len() != saved_lens.len() {
3444 return Err("spec TP KV snapshot shape mismatch".into());
3445 }
3446 for (layer, (cache, saved)) in tp_kv.iter_mut().zip(saved_lens).enumerate() {
3447 match (cache.as_mut(), *saved) {
3448 (Some(cache), Some(saved)) => {
3449 let target = saved
3450 .checked_add(accepted)
3451 .ok_or("spec TP KV committed length overflow")?;
3452 cache.rewind_to(target)?;
3453 }
3454 (None, None) => {}
3455 _ => {
3456 return Err(
3457 format!("spec TP KV layer {layer} changed shape since its snapshot").into(),
3458 );
3459 }
3460 }
3461 }
3462 Ok(())
3463}
3464
3465impl HybridModel {
3466 fn mtp_head_count(&self) -> usize {
3467 usize::from(self.mtp.is_some()) + self.mtp_extra.len()
3468 }
3469
3470 fn mtp_head_at(&self, index: usize) -> &MtpHead {
3471 if index == 0 {
3472 self.mtp.as_ref().expect("MTP head 0 is unavailable")
3473 } else {
3474 &self.mtp_extra[index - 1]
3475 }
3476 }
3477
3478 fn new_mtp_scratch(
3479 &self,
3480 e: &Engine,
3481 cap: usize,
3482 ) -> Result<MtpScratch, Box<dyn std::error::Error>> {
3483 let mut scratch = MtpScratch::new(
3484 e,
3485 &self.cfg,
3486 &self.plan,
3487 cap,
3488 self.mtp.as_ref().and_then(|head| head.geom.as_ref()),
3489 )?;
3490 for head in &self.mtp_extra {
3491 scratch.push_plane(e, &self.cfg, &self.plan, head.geom.as_ref())?;
3492 }
3493 Ok(scratch)
3494 }
3495
3496 fn opti_graph_draft_step(
3497 &self,
3498 e: &Engine,
3499 mtp: &MtpHead,
3500 dctx: &mut DraftGraphCtx,
3501 scratch: &mut MtpScratch,
3502 d_vocab: usize,
3503 ) -> Result<(u32, f32), Box<dyn std::error::Error>> {
3504 dctx.graph
3505 .as_ref()
3506 .ok_or("optipipe controller requires the greedy draft graph")?
3507 .launch()?;
3508 scratch.kv.len += 1;
3509 let idx = e.dtoh_u32_one(&dctx.g_tok)?;
3510 if (idx as usize) >= d_vocab {
3511 return Err(
3512 format!("optipipe draft argmax sentinel 0x{idx:08x} >= d_vocab {d_vocab}").into(),
3513 );
3514 }
3515 let probability = e.dtoh(&dctx.g_p)?[0];
3516 if !probability.is_finite() || !(0.0..=1.0).contains(&probability) {
3517 return Err(format!("optipipe draft probability is invalid: {probability}").into());
3518 }
3519 let token = match &mtp.d2t {
3520 Some(map) => map[idx as usize],
3521 None => idx,
3522 };
3523 if token != idx {
3524 e.set_u32_one(&mut dctx.g_tok, token)?;
3525 }
3526 Ok((token, probability))
3527 }
3528
3529 #[allow(clippy::too_many_arguments)]
3530 fn opti_controller_draft_step(
3531 &self,
3532 e: &Engine,
3533 mtp: &MtpHead,
3534 dctx: &mut DraftGraphCtx,
3535 scratch: &mut MtpScratch,
3536 d_vocab: usize,
3537 eager_state: &mut Option<(u32, CudaSlice<f32>)>,
3538 eager_pos: usize,
3539 embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
3540 ) -> Result<(u32, f32), Box<dyn std::error::Error>> {
3541 if dctx.graph.is_some() {
3542 return self.opti_graph_draft_step(e, mtp, dctx, scratch, d_vocab);
3543 }
3544 let (input_token, input_seed) = eager_state
3545 .take()
3546 .ok_or("optipipe eager continuation seed is unavailable")?;
3547 let (logits, next_seed) = self.mtp_head_forward_dev(
3548 e,
3549 mtp,
3550 input_token,
3551 &input_seed,
3552 scratch,
3553 eager_pos,
3554 embd_dev,
3555 None,
3556 )?;
3557 let token_d = e.argmax_token_device(&logits, d_vocab)?;
3558 let idx = e.dtoh_u32_one(&token_d)?;
3559 if (idx as usize) >= d_vocab {
3560 return Err(format!(
3561 "optipipe eager draft argmax sentinel 0x{idx:08x} >= d_vocab {d_vocab}"
3562 )
3563 .into());
3564 }
3565 let probability_d = e.prob_of_token_device(&logits, &token_d, d_vocab)?;
3566 let probability = e.dtoh(&probability_d)?[0];
3567 if !probability.is_finite() || !(0.0..=1.0).contains(&probability) {
3568 return Err(
3569 format!("optipipe eager draft probability is invalid: {probability}").into(),
3570 );
3571 }
3572 let token = match &mtp.d2t {
3573 Some(map) => map[idx as usize],
3574 None => idx,
3575 };
3576 *eager_state = Some((token, next_seed));
3577 Ok((token, probability))
3578 }
3579
3580 /// NextN head forward for ONE draft token (§A ops 1-13, T=1).
3581 /// Inputs: `e_tok` = the token to predict FROM (last committed / previous draft); `h_seed` =
3582 /// the trunk's pre-output_norm hidden of that token (§A op 2 input). `mtp_pos` = absolute
3583 /// position of the token being predicted from. Returns (draft_logits[n_vocab] host, h_nextn dev).
3584 /// `h_nextn` (§A op 10) becomes `h_seed` for the next autoregressive draft step.
3585 /// Device-resident: returns draft logits ON DEVICE (no [n_vocab] dtoh). The greedy draft
3586 /// loop only needs argmax — paired with `argmax_token_device` this cuts the ~600KB logits
3587 /// transfer + host argmax per draft token from the K-token draft chain.
3588 #[allow(clippy::too_many_arguments)]
3589 fn mtp_head_forward_dev(
3590 &self,
3591 e: &Engine,
3592 mtp: &MtpHead,
3593 e_tok: u32,
3594 h_seed: &CudaSlice<f32>,
3595 scratch: &mut MtpScratch,
3596 mtp_pos: usize,
3597 embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
3598 mask: Option<(&CudaSlice<u32>, usize)>,
3599 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
3600 self.mtp_head_forward_dev_at(e, mtp, e_tok, h_seed, scratch, 0, mtp_pos, embd_dev, mask)
3601 }
3602
3603 #[allow(clippy::too_many_arguments)]
3604 fn mtp_head_forward_dev_at(
3605 &self,
3606 e: &Engine,
3607 mtp: &MtpHead,
3608 e_tok: u32,
3609 h_seed: &CudaSlice<f32>,
3610 scratch: &mut MtpScratch,
3611 scratch_index: usize,
3612 mtp_pos: usize,
3613 embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
3614 // DRAFT-SIDE GRAMMAR MASK (lane/draft-mask): (packed draft-vocab allowed set, words).
3615 // Applied to the head logits BEFORE they are returned, so every consumer (argmax,
3616 // gumbel draw, p-min prob) sees the grammar-legal row. None = unmasked (pre-lane).
3617 mask: Option<(&CudaSlice<u32>, usize)>,
3618 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
3619 // MEMRA_SPEC_ANATOMY=1 — eager-step phase timers (diagnostic only). Phase boundaries
3620 // sync the stream, so absolute time inflates; the BREAKDOWN is the signal. Cumulative
3621 // summary on stderr every 128 steps: glue (embed..attn_norm), attn, ffn, head.
3622 use std::sync::atomic::{AtomicU64, Ordering::Relaxed};
3623 static ANAT_NS: [AtomicU64; 5] = [
3624 AtomicU64::new(0),
3625 AtomicU64::new(0),
3626 AtomicU64::new(0),
3627 AtomicU64::new(0),
3628 AtomicU64::new(0),
3629 ];
3630 static ANAT_STEPS: AtomicU64 = AtomicU64::new(0);
3631 let anat = {
3632 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
3633 *ON.get_or_init(|| std::env::var("MEMRA_SPEC_ANATOMY").as_deref() == Ok("1"))
3634 };
3635 if anat {
3636 e.stream().synchronize()?; // drain prior queue so phase 0 starts clean
3637 }
3638 let t_all = std::time::Instant::now();
3639 let mut t_ph = std::time::Instant::now();
3640 let mut anat_mark = |i: usize,
3641 e: &Engine,
3642 t: &mut std::time::Instant|
3643 -> Result<(), Box<dyn std::error::Error>> {
3644 if anat {
3645 e.stream().synchronize()?;
3646 ANAT_NS[i].fetch_add(t.elapsed().as_nanos() as u64, Relaxed);
3647 *t = std::time::Instant::now();
3648 }
3649 Ok(())
3650 };
3651 let cfg = &self.cfg;
3652 let n_embd = cfg.n_embd as usize;
3653 // Distilled-student geometry: the block runs at the INNER width `di` (eh_proj out /
3654 // attn / ffn); the n_embd interface (embed, norms in, carrier out, head in) is unchanged.
3655 let di = mtp.geom.as_ref().map(|g| g.d_inner).unwrap_or(n_embd);
3656 let eps = cfg.rms_eps;
3657 let pos_d = e.htod_i32(&[mtp_pos as i32])?;
3658
3659 // op A: a resident table transfers one 4B token id. The exact host-row capacity path
3660 // expands this one row on CPU and transfers n_embd f32 values instead.
3661 let e_emb = match embd_dev {
3662 Some((g, qt, rb)) => e.embed_gather_device_t(g, &[e_tok], n_embd, qt, rb)?,
3663 None => e.htod(&self.embd.gather(n_embd, &[e_tok]))?,
3664 };
3665
3666 // op 1/2: e_norm = RMSNorm(e, enorm); h_norm = RMSNorm(h_seed, hnorm)
3667 let mut e_norm = e.zeros(n_embd)?;
3668 e.rms_norm(&e_emb, mtp.enorm.float_data(), &mut e_norm, n_embd, 1, eps)?;
3669 let mut h_norm = e.zeros(n_embd)?;
3670 e.rms_norm(h_seed, mtp.hnorm.float_data(), &mut h_norm, n_embd, 1, eps)?;
3671
3672 // op 3: concat = [e_norm ; h_norm] -> [2*n_embd], e_norm in [0,n_embd), h_norm in [n_embd,2n_embd)
3673 let mut concat = e.zeros(2 * n_embd)?;
3674 e.copy_into(&mut concat, 0, &e_norm, n_embd)?;
3675 e.copy_into(&mut concat, n_embd, &h_norm, n_embd)?;
3676
3677 // op 4: inpSA = eh_proj @ concat (eh_proj [2*n_embd, n_embd]) -> [n_embd]
3678 let inp_sa = e.matmul(&mtp.eh_proj, &concat, 1)?;
3679
3680 // op 5: a_norm = RMSNorm(inpSA, attn_norm)
3681 let mut a_norm = e.zeros(di)?;
3682 e.rms_norm(&inp_sa, mtp.attn_norm.float_data(), &mut a_norm, di, 1, eps)?;
3683 anat_mark(0, e, &mut t_ph)?;
3684
3685 // op 6: attention on the scratch KV. SAME dc launcher as the graph path (bucket_max =
3686 // scratch.cap, length from the device len_d) so eager drafts match graph drafts
3687 // bit-for-bit at any t_kv (the parity gate). Host len mirrored here (the dc append
3688 // advances only the device counter).
3689 let attn_out = match (&mtp.mixer, mtp.step35.as_ref()) {
3690 // step35 MTP block: PER-LAYER geometry + a separate head-wise gate + an SWA window,
3691 // none of which the dc launcher can express (see `mtp_step35_attn`). Host-len arm.
3692 // Advances BOTH the host len and the device counter itself (unlike the dc arm,
3693 // whose host-side mirror the caller does).
3694 (Mixer::Full(fa), Some(g)) => {
3695 self.mtp_step35_attn(e, fa, g, &a_norm, &pos_d, scratch, scratch_index)?
3696 }
3697 (Mixer::Full(fa), None) => {
3698 let out = self.mtp_full_attn_dc(
3699 e,
3700 fa,
3701 &a_norm,
3702 &pos_d,
3703 scratch,
3704 scratch_index,
3705 mtp.geom.as_ref(),
3706 )?;
3707 scratch.plane_mut(scratch_index).0.len += 1;
3708 out
3709 }
3710 (Mixer::Linear(_), _) => {
3711 panic!("MTP block is full-attn in qwen35; linear MTP not supported")
3712 }
3713 (Mixer::Mla(_), _) => crate::hybrid::mla_forward_unimplemented(),
3714 };
3715 anat_mark(1, e, &mut t_ph)?;
3716
3717 // op 7: x1 = inpSA + attn_out
3718 let mut x1 = e.zeros(di)?;
3719 e.add(&inp_sa, &attn_out, &mut x1, di)?;
3720
3721 // op 8: z = RMSNorm(x1, post_attn_norm) (pre-FFN norm)
3722 let mut z = e.zeros(di)?;
3723 e.rms_norm(&x1, mtp.post_attn_norm.float_data(), &mut z, di, 1, eps)?;
3724
3725 // op 9: FFN (Dense or MoE) — same as the trunk decode FFN
3726 let ffn_out = match &mtp.ffn {
3727 crate::hybrid::Ffn::Dense {
3728 ffn_gate,
3729 ffn_up,
3730 ffn_down,
3731 } => {
3732 let n_ff = ffn_gate.out_features();
3733 let (gate, up) = if e.uses_q8_1_fast(ffn_gate) && e.uses_q8_1_fast(ffn_up) {
3734 let (zq, zd) = e.quantize_q8_1(&z, 1, di)?;
3735 (
3736 e.matmul_pre(ffn_gate, &zq, &zd, &z, 1)?,
3737 e.matmul_pre(ffn_up, &zq, &zd, &z, 1)?,
3738 )
3739 } else {
3740 (e.matmul(ffn_gate, &z, 1)?, e.matmul(ffn_up, &z, 1)?)
3741 };
3742 let mut act = e.zeros(n_ff)?;
3743 // step35: a DENSE FFN reads the per-layer SHEXP clamp (upstream's one `build_ffn`
3744 // serves the dense MLP and the shared expert off `swiglu_clamp_shexp` —
3745 // llama-graph.cpp:1751), resolved for the MTP block's OWN index. Every other arch
3746 // passes None, which is `ffn_act`'s dispatch verbatim.
3747 Self::ffn_act_lim(
3748 e,
3749 &self.cfg,
3750 &gate,
3751 &up,
3752 1.0,
3753 1.0,
3754 mtp.step35.as_ref().and_then(|s| s.clamp_shexp),
3755 &mut act,
3756 n_ff,
3757 )?;
3758 e.matmul(ffn_down, &act, 1)?
3759 }
3760 // MTP head is a distinct block — key its experts under a separate layer index (u16::MAX)
3761 // so they never alias trunk layer 0's cache keys.
3762 crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il(e, m, &z, 1, u16::MAX)?,
3763 };
3764 anat_mark(2, e, &mut t_ph)?;
3765
3766 // op 10: h_nextn = x1 + ffn_out (at di)
3767 let mut h_inner = e.zeros(di)?;
3768 e.add(&x1, &ffn_out, &mut h_inner, di)?;
3769
3770 // op 10.5 (student): up-project the inner hidden back to n_embd — training semantics:
3771 // the chain carrier AND the head input are out_up(h_inner) (pre-final-norm).
3772 let h_nextn = match mtp.geom.as_ref() {
3773 Some(g) => e.matmul(&g.out_up, &h_inner, 1)?,
3774 None => h_inner,
3775 };
3776
3777 // op 11: final = RMSNorm(h_nextn, shared_head_norm OR output_norm)
3778 let final_norm = mtp.shared_head_norm.as_ref().unwrap_or(&self.output_norm);
3779 let mut final_h = e.zeros(n_embd)?;
3780 e.rms_norm(
3781 &h_nextn,
3782 final_norm.float_data(),
3783 &mut final_h,
3784 n_embd,
3785 1,
3786 eps,
3787 )?;
3788
3789 // op 12: draft_logits = (shared_head_head OR output) @ final — stays ON DEVICE.
3790 let head = mtp.shared_head_head.as_ref().unwrap_or(&self.output);
3791 let mut logits = e.matmul(head, &final_h, 1)?;
3792 // op 12b (lane/draft-mask): grammar mask over the DRAFT vocab, applied here so the
3793 // caller's argmax / gumbel draw / p-min prob all read the grammar-legal row.
3794 if let Some((mask_d, mw)) = mask {
3795 let d_vocab = head.out_features();
3796 e.mask_logits_col(&mut logits, mask_d, 0, d_vocab, mw)?;
3797 }
3798 anat_mark(3, e, &mut t_ph)?;
3799 if anat {
3800 ANAT_NS[4].fetch_add(t_all.elapsed().as_nanos() as u64, Relaxed);
3801 let n = ANAT_STEPS.fetch_add(1, Relaxed) + 1;
3802 if n % 128 == 0 {
3803 let us = |i: usize| ANAT_NS[i].load(Relaxed) / n / 1000;
3804 eprintln!(
3805 "[spec-anatomy] steps={n} avg us/step: glue={} attn={} ffn={} head={} total={}",
3806 us(0),
3807 us(1),
3808 us(2),
3809 us(3),
3810 us(4)
3811 );
3812 }
3813 }
3814 // Chain recurrence hand-over: pre-norm h_nextn (default) or post-norm final_h
3815 // (MEMRA_SPEC_HPOST — llama.cpp #24025's t_h_nextn is taken AFTER the head norm).
3816 Ok((logits, if spec_hpost() { final_h } else { h_nextn }))
3817 }
3818
3819 #[allow(clippy::too_many_arguments)]
3820 fn mtp_chain_forward_dev(
3821 &self,
3822 e: &Engine,
3823 tokens: &[u32],
3824 seeds: &[CudaSlice<f32>],
3825 scratch: &mut MtpScratch,
3826 committed_scratch_len: usize,
3827 embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
3828 mask: Option<(&CudaSlice<u32>, usize)>,
3829 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
3830 if tokens.is_empty() || tokens.len() != seeds.len() {
3831 return Err("multi-head MTP prefix tokens/seeds are malformed".into());
3832 }
3833 let index = mtp_chain_head_index(tokens.len() - 1, self.mtp_head_count());
3834 let head = self.mtp_head_at(index);
3835 scratch.set_plane_len(e, index, committed_scratch_len)?;
3836
3837 let mut last = None;
3838 for row in 0..tokens.len() {
3839 let is_last = row + 1 == tokens.len();
3840 last = Some(self.mtp_head_forward_dev_at(
3841 e,
3842 head,
3843 tokens[row],
3844 &seeds[row],
3845 scratch,
3846 index,
3847 committed_scratch_len + row + 1,
3848 embd_dev,
3849 if is_last { mask } else { None },
3850 )?);
3851 }
3852 Ok(last.expect("non-empty MTP prefix produced no row"))
3853 }
3854
3855 /// step35 MTP-block attention, T=1, on the scratch KV — the EAGER-ONLY twin of
3856 /// `mtp_full_attn_dc`. Three things force a separate arm rather than a geometry parameter on
3857 /// the dc path, and all three are properties of this arch's MTP block:
3858 ///
3859 /// 1. **The SWA window.** Block 45 is an SWA-type block (`sliding_window_pattern[45]=true`,
3860 /// window 512). Windowed decode in memra is a token-aligned VIEW OFFSET into the quantized
3861 /// cache (the gemma4 R6 / `step35_decode_attn` pattern: keys carry absolute rope and the
3862 /// mask is purely positional, so one query at `len-1` attending the last `win` rows IS the
3863 /// windowed result). `fa_decode_dc` takes the key count from a DEVICE counter and always
3864 /// starts at row 0 — it cannot express a nonzero offset, so a windowed dc arm would need a
3865 /// new kernel. That is deliberately not built here: see the CUDA-graph note below.
3866 /// 2. **Per-layer head count.** 96 q heads over 8 KV (GQA 12) at this block, vs the trunk's 64
3867 /// on its full-attn layers. The trunk cfg's `n_head` scalar is the MAX over layers, and the
3868 /// trunk ARTIFACT's per-layer arrays stop at index 44 — so the count must come from the
3869 /// resolved `Step35MtpGeom`, never from `cfg`.
3870 /// 3. **The separate head-wise gate.** `blk.45.attn_gate.weight [n_embd, 96]` produces one
3871 /// sigmoid scalar per head (broadcast over head_dim) — `attn_head_gate`, not the qwen35
3872 /// fused-into-wq `q_gate_split` form the dc arm handles.
3873 ///
3874 /// WHY EAGER-ONLY IS NOT A GAP TODAY: `mtp_head_forward_cap` refuses step35 heads
3875 /// explicitly (the SWA-window refusal below), so the graph draft never engages for step35
3876 /// models regardless of eligibility — the eager chain IS the served path. `mtp_head_forward_cap` refuses step35 explicitly rather
3877 /// than silently capturing a window-less (wrong past `win` draft rows) graph.
3878 ///
3879 /// Unlike the dc arm this advances BOTH the host `kv.len` and the device counter, so the
3880 /// caller must not mirror.
3881 fn mtp_step35_attn(
3882 &self,
3883 e: &Engine,
3884 fa: &FullAttnLayer,
3885 g: &crate::hybrid::Step35MtpGeom,
3886 h: &CudaSlice<f32>,
3887 pos_d: &CudaSlice<i32>,
3888 scratch: &mut MtpScratch,
3889 scratch_index: usize,
3890 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3891 let (nh, nkv, hd) = (g.n_head, g.n_head_kv, self.cfg.head_dim_k as usize);
3892 let eps = self.cfg.rms_eps;
3893 let scale = 1.0 / (hd as f32).sqrt(); // step35.cpp:255 kq_scale
3894 let n_embd = self.cfg.n_embd as usize;
3895 let gw = fa
3896 .attn_gate
3897 .as_ref()
3898 .ok_or("step35 MTP block is missing attn_gate.weight (head-wise attention gate)")?;
3899
3900 let (q0, k0, v0, gt) = if e.uses_q8_1_fast(&fa.wq)
3901 && e.uses_q8_1_fast(&fa.wk)
3902 && e.uses_q8_1_fast(&fa.wv)
3903 && e.uses_q8_1_fast(gw)
3904 {
3905 let (hq, hdq) = e.quantize_q8_1(h, 1, n_embd)?;
3906 let (a, b, c) = match e.matmul_q8_fused3(&fa.wq, &fa.wk, &fa.wv, &hq, &hdq)? {
3907 Some(t3) => t3,
3908 None => (
3909 e.matmul_pre(&fa.wq, &hq, &hdq, h, 1)?,
3910 e.matmul_pre(&fa.wk, &hq, &hdq, h, 1)?,
3911 e.matmul_pre(&fa.wv, &hq, &hdq, h, 1)?,
3912 ),
3913 };
3914 (a, b, c, e.matmul_pre(gw, &hq, &hdq, h, 1)?)
3915 } else {
3916 (
3917 e.matmul(&fa.wq, h, 1)?,
3918 e.matmul(&fa.wk, h, 1)?,
3919 e.matmul(&fa.wv, h, 1)?,
3920 e.matmul(gw, h, 1)?,
3921 )
3922 };
3923
3924 let mut q = e.uninit(nh * hd)?;
3925 e.rms_norm(&q0, fa.q_norm.float_data(), &mut q, hd, nh, eps)?;
3926 let mut k = e.uninit(nkv * hd)?;
3927 e.rms_norm(&k0, fa.k_norm.float_data(), &mut k, hd, nkv, eps)?;
3928 // `rope_freqs.weight` (llama3 factors) applies to the FULL-attn layers ONLY; SWA passes
3929 // null (llama-hparams / step35.cpp). Block 45 is SWA, so `ff` is None there — but read
3930 // the resolved flag, not the constant, so an all-full sibling stays correct.
3931 let ff = if g.swa {
3932 None
3933 } else {
3934 self.step35_aux.as_ref().and_then(|a| a.rope_freqs(e))
3935 };
3936 #[cfg(debug_assertions)]
3937 if let Some(ff) = ff {
3938 crate::debug_assert_tensor_stream_device(ff, &e.stream(), "mtp_step35_attn.rope_freqs");
3939 }
3940 e.rope_neox2(
3941 &mut q,
3942 &mut k,
3943 pos_d,
3944 hd,
3945 g.n_rot,
3946 nh,
3947 nkv,
3948 1,
3949 g.rope_base,
3950 1.0,
3951 ff,
3952 )?;
3953
3954 // Append at the HOST slot, then re-stamp the device counter: the eager chain has the
3955 // length on the host anyway, and the windowed view below needs it there to compute the
3956 // offset. The device counter is kept in lockstep so `mtp_kv_fill`'s `set_i32_one` and any
3957 // dc-family consumer of this scratch still agree.
3958 let (kv, scratch_cap) = scratch.plane_mut(scratch_index);
3959 assert!(
3960 kv.len < scratch_cap,
3961 "step35 MTP scratch overflow ({} >= {})",
3962 kv.len,
3963 scratch_cap
3964 );
3965 let next_len = kv.len + 1;
3966 let (off, t_kv) = if g.swa && next_len > g.window {
3967 (next_len - g.window, g.window)
3968 } else {
3969 (0, next_len)
3970 };
3971 let write_row = e.prepare_kv_append(kv, off & !31usize, 1)?;
3972 e.append_kv_quantized(
3973 &k,
3974 &v0,
3975 &mut kv.k,
3976 &mut kv.v,
3977 write_row,
3978 kv.kv_dim_k,
3979 kv.kv_dim_v,
3980 kv.k_tok_bytes,
3981 kv.v_tok_bytes,
3982 false,
3983 )?;
3984 kv.len = next_len;
3985 e.set_i32_one(&mut kv.len_d, kv.len as i32)?;
3986 // SWA view offset (see note 1). The draft chain is short (k+2 rows), but the scratch is
3987 // PERSISTENT across rounds — `mtp_kv_fill` leaves one row per committed token behind, so
3988 // `kv.len` tracks absolute position and crosses 512 in any real generation. The window is
3989 // therefore live, not theoretical.
3990 let physical = kv.physical_rows(off, off + t_kv)?;
3991 let k_view = e.view_u8_range(
3992 &kv.k,
3993 physical.start * kv.k_tok_bytes,
3994 physical.end * kv.k_tok_bytes,
3995 );
3996 let v_view = e.view_u8_range(
3997 &kv.v,
3998 physical.start * kv.v_tok_bytes,
3999 physical.end * kv.v_tok_bytes,
4000 );
4001 let mut attn = e.uninit(nh * hd)?;
4002 e.fa_decode_kvmod(
4003 &q,
4004 &k_view,
4005 &v_view,
4006 &mut attn,
4007 hd,
4008 nh,
4009 nkv,
4010 t_kv,
4011 scale,
4012 kv.k_tok_bytes,
4013 kv.v_tok_bytes,
4014 false,
4015 )?;
4016
4017 let mut ag = e.uninit(nh * hd)?;
4018 e.attn_head_gate(&attn, >, &mut ag, None, hd, nh, 1)?;
4019 Ok(e.matmul(&fa.wo, &ag, 1)?)
4020 }
4021
4022 /// MTP-block full attention, T=1, on the scratch KV (BOTH draft paths — eager and graph):
4023 /// the scratch write slot and the attention bound come from `scratch.kv.len_d` (device i32[1])
4024 /// so the launch args are FIXED across draft steps — ONE captured graph serves the whole
4025 /// chain, and replays keep seeing KV growth through the device counter (no recapture).
4026 /// Geometry contract: n_splits is sized from `scratch.cap` (the persistent capacity); splits
4027 /// whose key range lies beyond the device t_kv exit empty and the shared combine skips them
4028 /// (fa_decode_dc bit-correct-for-any-t_kv<=bucket_max contract). The eager path uses the SAME
4029 /// launcher with the SAME bucket_max -> identical dispatch -> bit-identical draft tokens (the
4030 /// graph-vs-eager parity gate). Host len is NOT advanced here (graph contract); callers mirror.
4031 fn mtp_full_attn_dc(
4032 &self,
4033 e: &Engine,
4034 fa: &FullAttnLayer,
4035 h: &CudaSlice<f32>,
4036 pos_d: &CudaSlice<i32>,
4037 scratch: &mut MtpScratch,
4038 scratch_index: usize,
4039 geom: Option<&crate::hybrid::DraftGeom>,
4040 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4041 let cfg = &self.cfg;
4042 let mtp_il = cfg.n_layer.saturating_sub(cfg.nextn_predict_layers);
4043 let geometry = cfg.full_attention_geometry_at(mtp_il);
4044 let n_head = geom.map(|g| g.n_head).unwrap_or(geometry.n_head as usize);
4045 let n_head_kv = geom
4046 .map(|g| g.n_head_kv)
4047 .unwrap_or(geometry.n_head_kv as usize);
4048 let head_dim = geometry.head_dim_k as usize;
4049 let eps = cfg.rms_eps;
4050 let scale = geometry.attention_scale();
4051 let n_embd = geom.map(|g| g.d_inner).unwrap_or(cfg.n_embd as usize);
4052 let bucket_max = scratch.plane(scratch_index).1;
4053
4054 let (qf, mut k, v) =
4055 if e.uses_q8_1_fast(&fa.wq) && e.uses_q8_1_fast(&fa.wk) && e.uses_q8_1_fast(&fa.wv) {
4056 let (hq, hd) = e.quantize_q8_1(h, 1, n_embd)?;
4057 (
4058 e.matmul_pre(&fa.wq, &hq, &hd, h, 1)?,
4059 e.matmul_pre(&fa.wk, &hq, &hd, h, 1)?,
4060 e.matmul_pre(&fa.wv, &hq, &hd, h, 1)?,
4061 )
4062 } else {
4063 (
4064 e.matmul(&fa.wq, h, 1)?,
4065 e.matmul(&fa.wk, h, 1)?,
4066 e.matmul(&fa.wv, h, 1)?,
4067 )
4068 };
4069 // M3/Hy3 have no attention output gate — wq out is exactly q; skip the split.
4070 let gated = geometry.attention_gate == memra_gguf::config::AttentionGateKind::FusedQ;
4071 let (mut q, gate) = if gated {
4072 let mut q = e.zeros(n_head * head_dim)?;
4073 let mut gate = e.zeros(n_head * head_dim)?;
4074 e.q_gate_split(&qf, &mut q, &mut gate, head_dim, n_head, 1)?;
4075 (q, Some(gate))
4076 } else {
4077 (qf, None)
4078 };
4079
4080 let mut qn = e.zeros(n_head * head_dim)?;
4081 e.rms_norm(&q, fa.q_norm.float_data(), &mut qn, head_dim, n_head, eps)?;
4082 q = qn;
4083 let mut kn = e.zeros(n_head_kv * head_dim)?;
4084 e.rms_norm(
4085 &k,
4086 fa.k_norm.float_data(),
4087 &mut kn,
4088 head_dim,
4089 n_head_kv,
4090 eps,
4091 )?;
4092 k = kn;
4093 let rope_dims = geometry.n_rot as usize;
4094 e.rope_neox(
4095 &mut q,
4096 pos_d,
4097 head_dim,
4098 rope_dims,
4099 n_head,
4100 1,
4101 geometry.rope_base,
4102 1.0,
4103 )?;
4104 e.rope_neox(
4105 &mut k,
4106 pos_d,
4107 head_dim,
4108 rope_dims,
4109 n_head_kv,
4110 1,
4111 geometry.rope_base,
4112 1.0,
4113 )?;
4114
4115 let kv = scratch.plane_mut(scratch_index).0;
4116 // append at the DEVICE slot (kv.len_d == old len), then advance the counter in-graph.
4117 e.append_kv_quantized_dc(
4118 &k,
4119 &v,
4120 &mut kv.k,
4121 &mut kv.v,
4122 &kv.len_d,
4123 kv.kv_dim_k,
4124 kv.kv_dim_v,
4125 kv.k_tok_bytes,
4126 kv.v_tok_bytes,
4127 false,
4128 )?;
4129 e.inc_seqlen(&mut kv.len_d)?;
4130 // full-buffer views (any in-round t_kv stays in range on replay); the kernel bounds the
4131 // key range from the device counter.
4132 let k_view = e.view_u8(&kv.k, kv.k.len());
4133 let v_view = e.view_u8(&kv.v, kv.v.len());
4134 let (ktb, vtb) = (kv.k_tok_bytes, kv.v_tok_bytes);
4135 let mut attn = e.zeros(n_head * head_dim)?;
4136 e.fa_decode_dc(
4137 &q, &k_view, &v_view, &mut attn, head_dim, n_head, n_head_kv, &kv.len_d, bucket_max,
4138 scale, ktb, vtb, false,
4139 )?;
4140
4141 let attn_g = match &gate {
4142 Some(gate) => {
4143 let mut gsig = e.zeros(n_head * head_dim)?;
4144 e.sigmoid(gate, &mut gsig, n_head * head_dim)?;
4145 let mut ag = e.zeros(n_head * head_dim)?;
4146 e.mul(&attn, &gsig, &mut ag, n_head * head_dim)?;
4147 ag
4148 }
4149 None => attn,
4150 };
4151 Ok(e.matmul(&fa.wo, &attn_g, 1)?)
4152 }
4153
4154 /// PERSISTENT-DRAFT-KV fill (the reference engine's "mtp_update" analogue): compute the MTP
4155 /// block's K/V for `tokens` (committed tokens at positions pos0..pos0+T) from their EXACT
4156 /// trunk hiddens `h` ([T, n_embd] token-major, pre-output_norm) and append at slots pos0.. of
4157 /// the scratch KV. K/V-ONLY — ops A/1-5 plus the K-side of op 6 (wk/wv + k_norm + rope +
4158 /// quantized append); no wq/attention/FFN/lm_head, so per-token cost ~= eh_proj + wk/wv (a
4159 /// small fraction of one trunk layer), T-batched. Rope follows the chain convention
4160 /// rope(token@p) = p+1. Runs at round boundaries OUTSIDE the captured graph in BOTH draft
4161 /// modes -> draft parity by construction. Caller must have scratch.kv.len == pos0.
4162 #[allow(clippy::too_many_arguments)]
4163 fn mtp_kv_fill_at(
4164 &self,
4165 e: &Engine,
4166 mtp: &MtpHead,
4167 tokens: &[u32],
4168 h: &CudaSlice<f32>,
4169 pos0: usize,
4170 scratch: &mut MtpScratch,
4171 scratch_index: usize,
4172 embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
4173 ) -> Result<(), Box<dyn std::error::Error>> {
4174 let cfg = &self.cfg;
4175 let n_embd = cfg.n_embd as usize;
4176 let eps = cfg.rms_eps;
4177 let t = tokens.len();
4178 let (scratch_kv, scratch_cap) = scratch.plane(scratch_index);
4179 assert_eq!(scratch_kv.len, pos0, "mtp_kv_fill: append slot mismatch");
4180 assert!(pos0 + t <= scratch_cap, "mtp_kv_fill: scratch overflow");
4181 let Mixer::Full(fa) = &mtp.mixer else {
4182 panic!("MTP block is full-attn in qwen35; linear MTP not supported")
4183 };
4184 let pos_vec: Vec<i32> = (0..t).map(|i| (pos0 + i + 1) as i32).collect();
4185 let pos_d = e.htod_i32(&pos_vec)?;
4186
4187 // ops A/1/2: embed + the two input norms, T-wide.
4188 let e_emb = match embd_dev {
4189 Some((g, qt, rb)) => e.embed_gather_device_t(g, tokens, n_embd, qt, rb)?,
4190 None => e.htod(&self.embd.gather(n_embd, tokens))?,
4191 };
4192 let mut e_norm = e.zeros(t * n_embd)?;
4193 e.rms_norm(&e_emb, mtp.enorm.float_data(), &mut e_norm, n_embd, t, eps)?;
4194 let mut h_norm = e.zeros(t * n_embd)?;
4195 e.rms_norm(h, mtp.hnorm.float_data(), &mut h_norm, n_embd, t, eps)?;
4196
4197 // op 3: per-row [e_norm ; h_norm] concat, token-major [T, 2*n_embd].
4198 let mut concat = e.zeros(t * 2 * n_embd)?;
4199 for i in 0..t {
4200 e.copy_view_into(
4201 &mut concat,
4202 i * 2 * n_embd,
4203 &e_norm.slice(i * n_embd..(i + 1) * n_embd),
4204 n_embd,
4205 )?;
4206 e.copy_view_into(
4207 &mut concat,
4208 i * 2 * n_embd + n_embd,
4209 &h_norm.slice(i * n_embd..(i + 1) * n_embd),
4210 n_embd,
4211 )?;
4212 }
4213
4214 // ops 4/5: eh_proj + attn_norm, T-wide (at the student inner width when geom is set).
4215 let di = mtp.geom.as_ref().map(|g| g.d_inner).unwrap_or(n_embd);
4216 let inp_sa = e.matmul(&mtp.eh_proj, &concat, t)?;
4217 let mut a_norm = e.zeros(t * di)?;
4218 e.rms_norm(&inp_sa, mtp.attn_norm.float_data(), &mut a_norm, di, t, eps)?;
4219
4220 // op 6 (K/V half): wk/wv + k_norm + rope + per-row quantized append. No wq/attention —
4221 // the fill only has to leave correct K/V rows behind for later chains to attend over.
4222 let n_head_kv = mtp
4223 .geom
4224 .as_ref()
4225 .map(|g| g.n_head_kv)
4226 .or(mtp.step35.as_ref().map(|s| s.n_head_kv))
4227 .unwrap_or_else(|| {
4228 let mtp_il = cfg.n_layer.saturating_sub(cfg.nextn_predict_layers);
4229 cfg.full_attention_geometry_at(mtp_il).n_head_kv as usize
4230 });
4231 let mtp_il = cfg.n_layer.saturating_sub(cfg.nextn_predict_layers);
4232 let geometry = cfg.full_attention_geometry_at(mtp_il);
4233 let head_dim = geometry.head_dim_k as usize;
4234 let mut k = e.matmul(&fa.wk, &a_norm, t)?;
4235 let v = e.matmul(&fa.wv, &a_norm, t)?;
4236 let mut kn = e.zeros(t * n_head_kv * head_dim)?;
4237 e.rms_norm(
4238 &k,
4239 fa.k_norm.float_data(),
4240 &mut kn,
4241 head_dim,
4242 n_head_kv * t,
4243 eps,
4244 )?;
4245 k = kn;
4246 // step35: rotary width AND base are per-layer, and the MTP block's values come from the
4247 // resolved `Step35MtpGeom` — NOT from `cfg.rope_dim_count`/`cfg.rope_freq_base`, which
4248 // carry the arch defaults (128 / 5e6, i.e. the FULL-attn layers' base). Getting this wrong
4249 // writes K rows the attention arm then re-derives at a different theta: correct-looking
4250 // output with dead acceptance, invisible to the exactness gates.
4251 let (rope_dims, rope_base, ff) = match mtp.step35.as_ref() {
4252 Some(s) => (
4253 s.n_rot,
4254 s.rope_base,
4255 if s.swa {
4256 None
4257 } else {
4258 self.step35_aux.as_ref().and_then(|a| a.rope_freqs(e))
4259 },
4260 ),
4261 None => (geometry.n_rot as usize, geometry.rope_base, None),
4262 };
4263 #[cfg(debug_assertions)]
4264 if let Some(ff) = ff {
4265 crate::debug_assert_tensor_stream_device(ff, &e.stream(), "mtp_kv_fill.rope_freqs");
4266 }
4267 match ff {
4268 Some(f) => e.rope_neox_ff(
4269 &mut k, &pos_d, head_dim, rope_dims, n_head_kv, t, rope_base, 1.0, f,
4270 )?,
4271 None => e.rope_neox(
4272 &mut k, &pos_d, head_dim, rope_dims, n_head_kv, t, rope_base, 1.0,
4273 )?,
4274 }
4275
4276 let kv = scratch.plane_mut(scratch_index).0;
4277 // Match the trunk prime contract: a chunk may need the aligned window immediately before
4278 // its first row, so preserve that prefix when the physical tail rebases at wrap.
4279 let retain_from = kv
4280 .ring
4281 .as_ref()
4282 .map(|ring| pos0.saturating_sub(ring.window() - 1) & !31usize)
4283 .unwrap_or(0);
4284 let write_row = e.prepare_kv_append(kv, retain_from, t)?;
4285 for i in 0..t {
4286 let k_row = k.slice(i * kv.kv_dim_k..(i + 1) * kv.kv_dim_k);
4287 let v_row = v.slice(i * kv.kv_dim_v..(i + 1) * kv.kv_dim_v);
4288 e.append_kv_quantized_view(
4289 &k_row,
4290 &v_row,
4291 &mut kv.k,
4292 &mut kv.v,
4293 write_row + i,
4294 kv.kv_dim_k,
4295 kv.kv_dim_v,
4296 kv.k_tok_bytes,
4297 kv.v_tok_bytes,
4298 false,
4299 )?;
4300 }
4301 kv.len = pos0 + t;
4302 e.set_i32_one(&mut kv.len_d, kv.len as i32)?;
4303 Ok(())
4304 }
4305
4306 #[allow(clippy::too_many_arguments)]
4307 fn mtp_kv_fill_all(
4308 &self,
4309 e: &Engine,
4310 tokens: &[u32],
4311 h: &CudaSlice<f32>,
4312 pos0: usize,
4313 scratch: &mut MtpScratch,
4314 embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
4315 ) -> Result<(), Box<dyn std::error::Error>> {
4316 debug_assert_eq!(self.mtp_head_count(), scratch.plane_count());
4317 for index in 0..self.mtp_head_count() {
4318 self.mtp_kv_fill_at(
4319 e,
4320 self.mtp_head_at(index),
4321 tokens,
4322 h,
4323 pos0,
4324 scratch,
4325 index,
4326 embd_dev,
4327 )?;
4328 }
4329 Ok(())
4330 }
4331
4332 /// CAPTURE body for the GRAPH DRAFT (stage 2 of graph-grade spec): ONE MTP head forward with
4333 /// every varying input device-resident —
4334 /// - token id from the persistent `tok_d` (the previous replay's in-graph argmax wrote it,
4335 /// so the chain feeds itself; the host reads the same 4 bytes for the draft list),
4336 /// - h_seed from the persistent `h_seed_d` (h_nextn is copied BACK into it at the end),
4337 /// - rope pos from the persistent `pos_d` counter (inc'd in-graph),
4338 /// - scratch KV slot/bound from `scratch.kv.len_d` (see mtp_full_attn_dc).
4339 /// The p-min confidence lands in the persistent `p_d` iff `with_prob` (env is fixed per run).
4340 /// Same kernels, same dispatch as the eager mtp_head_forward_dev chain -> same draft tokens
4341 /// (exactness never depends on drafts — the verify arbitrates — but acceptance parity does).
4342 /// `with_head=false` captures the HEAD-LESS twin for the pseudo-seed replay (2026-07-03):
4343 /// the pseudo pass only needs h_nextn (op 10) + the scratch append — the lm_head read
4344 /// (~1.06ms q6_K on the 9B), argmax and prob are dead weight there. h_nextn's inputs are
4345 /// untouched, so the seed value is identical; round-start resets overwrite tok_d/p_d anyway.
4346 /// `sampled_cap` = Some((ctr_d, perturb_d, q_out_d, seed, temp)) captures the SAMPLED twin
4347 /// (step 3 of the sampled-spec arc): head logits are retained in the persistent `q_out_d`
4348 /// (host D2Ds them to the round's q slot after each replay), the DEVICE event counter is
4349 /// bumped in-graph, and the argmax reads GUMBEL-PERTURBED logits — one categorical draw per
4350 /// replay, bit-identical to the eager arm's gumbel_perturb at the same (seed, sctr, temp).
4351 /// seed/temp are capture-time constants (fixed per generate call, like p_min).
4352 #[allow(clippy::too_many_arguments)]
4353 fn mtp_head_forward_cap(
4354 &self,
4355 e: &Engine,
4356 mtp: &MtpHead,
4357 tok_d: &mut CudaSlice<u32>,
4358 pos_d: &mut CudaSlice<i32>,
4359 h_seed_d: &mut CudaSlice<f32>,
4360 p_d: &mut CudaSlice<f32>,
4361 scratch: &mut MtpScratch,
4362 with_prob: bool,
4363 with_head: bool,
4364 embd_gpu: &CudaSlice<u8>,
4365 embd_qt: i32,
4366 embd_rb: usize,
4367 d_vocab: usize,
4368 sampled_cap: Option<(
4369 &mut CudaSlice<u32>,
4370 &mut CudaSlice<f32>,
4371 &mut CudaSlice<f32>,
4372 u64,
4373 f32,
4374 )>,
4375 stream_pack: Option<(&mut CudaSlice<u32>, usize, Option<&CudaSlice<u32>>)>,
4376 // DRAFT-SIDE GRAMMAR MASK (lane/draft-mask): (packed draft-vocab allowed-set buffer,
4377 // word count). Captured as ONE mask_logits_f32 node between the head matmul and the
4378 // in-graph argmax; the buffer address is baked, its CONTENTS are re-uploaded by the
4379 // host before every replay (the decode.rs graph-mask pattern). All-ones contents = a
4380 // no-op ban, so a position the grammar cannot constrain costs one pass over the row.
4381 mask_cap: Option<(&CudaSlice<u32>, usize)>,
4382 ) -> Result<(), Box<dyn std::error::Error>> {
4383 let cfg = &self.cfg;
4384 let n_embd = cfg.n_embd as usize;
4385 // step35 REFUSAL (deliberate, named): the graph body's attention is `mtp_full_attn_dc`,
4386 // whose device-counter key bound always starts at row 0 — it cannot express this block's
4387 // SWA view offset, so a captured chain would silently attend OUTSIDE the window once the
4388 // persistent scratch passes 512 rows. Nothing is lost today: `mtp_head_forward_cap`
4389 // refuses step35 heads explicitly (SWA refusal), so the eager chain
4390 // (`mtp_head_forward_dev` -> `mtp_step35_attn`) is the served path. Returning Err (not a
4391 // panic) is what the two capture sites and the round-stream capture already handle by
4392 // degrading to eager / stream-off.
4393 if mtp.step35.is_some() {
4394 return Err(
4395 "step35 has no captured draft chain (fa_decode_dc cannot express the MTP \
4396 block's SWA view offset; same root cause as the dc decode refusal) — the \
4397 eager draft chain serves this arch"
4398 .into(),
4399 );
4400 }
4401 // student inner width (see mtp_head_forward_dev) — interface dims stay n_embd.
4402 let di = mtp.geom.as_ref().map(|g| g.d_inner).unwrap_or(n_embd);
4403 let eps = cfg.rms_eps;
4404 let e_emb = e.embed_gather_device(embd_gpu, tok_d, n_embd, embd_qt, embd_rb)?;
4405 let mut e_norm = e.zeros(n_embd)?;
4406 e.rms_norm(&e_emb, mtp.enorm.float_data(), &mut e_norm, n_embd, 1, eps)?;
4407 let mut h_norm = e.zeros(n_embd)?;
4408 e.rms_norm(
4409 &*h_seed_d,
4410 mtp.hnorm.float_data(),
4411 &mut h_norm,
4412 n_embd,
4413 1,
4414 eps,
4415 )?;
4416 let mut concat = e.zeros(2 * n_embd)?;
4417 e.copy_into(&mut concat, 0, &e_norm, n_embd)?;
4418 e.copy_into(&mut concat, n_embd, &h_norm, n_embd)?;
4419 let inp_sa = e.matmul(&mtp.eh_proj, &concat, 1)?;
4420 let mut a_norm = e.zeros(di)?;
4421 e.rms_norm(&inp_sa, mtp.attn_norm.float_data(), &mut a_norm, di, 1, eps)?;
4422 let attn_out = match &mtp.mixer {
4423 Mixer::Full(fa) => {
4424 self.mtp_full_attn_dc(e, fa, &a_norm, pos_d, scratch, 0, mtp.geom.as_ref())?
4425 }
4426 Mixer::Linear(_) => {
4427 panic!("MTP block is full-attn in qwen35; linear MTP not supported")
4428 }
4429 Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
4430 };
4431 let mut x1 = e.zeros(di)?;
4432 e.add(&inp_sa, &attn_out, &mut x1, di)?;
4433 let mut z = e.zeros(di)?;
4434 e.rms_norm(&x1, mtp.post_attn_norm.float_data(), &mut z, di, 1, eps)?;
4435 let ffn_out = match &mtp.ffn {
4436 crate::hybrid::Ffn::Dense {
4437 ffn_gate,
4438 ffn_up,
4439 ffn_down,
4440 } => {
4441 let n_ff = ffn_gate.out_features();
4442 let (gate, up) = if e.uses_q8_1_fast(ffn_gate) && e.uses_q8_1_fast(ffn_up) {
4443 let (zq, zd) = e.quantize_q8_1(&z, 1, di)?;
4444 (
4445 e.matmul_pre(ffn_gate, &zq, &zd, &z, 1)?,
4446 e.matmul_pre(ffn_up, &zq, &zd, &z, 1)?,
4447 )
4448 } else {
4449 (e.matmul(ffn_gate, &z, 1)?, e.matmul(ffn_up, &z, 1)?)
4450 };
4451 let mut act = e.zeros(n_ff)?;
4452 Self::ffn_act(e, &self.cfg, &gate, &up, &mut act, n_ff)?;
4453 e.matmul(ffn_down, &act, 1)?
4454 }
4455 // ROUND-STREAM: the 35B NextN block carries a MoE FFN. With RESIDENT experts the
4456 // dev path is pure device launches (device top-k + rows kernels, ZERO-DtoH by
4457 // design) — capture-legal. Non-resident (SLRU-lock) stays rejected: the capture
4458 // error arm degrades the caller to eager/stream-off.
4459 crate::hybrid::Ffn::Moe(m) if m.dev_exps.is_some() => {
4460 self.moe_ffn_il(e, m, &z, 1, u16::MAX)?
4461 }
4462 crate::hybrid::Ffn::Moe(_) => {
4463 return Err("graph draft requires a Dense (or resident-MoE) MTP FFN".into());
4464 }
4465 };
4466 let mut h_inner = e.zeros(di)?;
4467 e.add(&x1, &ffn_out, &mut h_inner, di)?;
4468 // student: up-project back to n_embd (carrier + head input; see mtp_head_forward_dev).
4469 let h_nextn = match mtp.geom.as_ref() {
4470 Some(g) => e.matmul(&g.out_up, &h_inner, 1)?,
4471 None => h_inner,
4472 };
4473 // MEMRA_SPEC_HPOST needs final_h even head-less (it IS the next seed under that convention).
4474 let final_h = if with_head || spec_hpost() {
4475 let final_norm = mtp.shared_head_norm.as_ref().unwrap_or(&self.output_norm);
4476 let mut fh = e.zeros(n_embd)?;
4477 e.rms_norm(&h_nextn, final_norm.float_data(), &mut fh, n_embd, 1, eps)?;
4478 Some(fh)
4479 } else {
4480 None
4481 };
4482 if with_head {
4483 let head = mtp.shared_head_head.as_ref().unwrap_or(&self.output);
4484 let mut logits = e.matmul(head, final_h.as_ref().unwrap(), 1)?;
4485 // DRAFT-SIDE GRAMMAR MASK: ban the grammar-illegal draft ids IN the captured chain,
4486 // before the argmax — proposals become legal by construction. Contents-only
4487 // per-replay upload keeps the capture valid.
4488 if let Some((mask_d, mw)) = mask_cap {
4489 e.mask_logits_col(&mut logits, mask_d, 0, d_vocab, mw)?;
4490 }
4491 if let Some((ctr_d, perturb_d, q_out_d, seed, temp)) = sampled_cap {
4492 // SAMPLED chain: retain q (raw head logits -> persistent q_out_d; the matmul's
4493 // own buffer is pool-recycled after the capture body returns, so it can't be the
4494 // retention target), bump the device event counter, gumbel-perturb reading it,
4495 // and argmax the PERTURBED logits into tok_d — the in-graph categorical draw.
4496 e.copy_into(q_out_d, 0, &logits, d_vocab)?;
4497 e.sctr_inc(ctr_d)?;
4498 e.gumbel_perturb_ctr(&logits, perturb_d, d_vocab, seed, ctr_d, temp)?;
4499 e.argmax_token_device_into(perturb_d, tok_d, d_vocab)?;
4500 // p-min prob = the head's RAW softmax confidence in the SAMPLED pick — same
4501 // semantics as the eager sampled arm's prob_of_token_device(dl_d, tok_d).
4502 if with_prob {
4503 e.prob_of_token_device_into(&logits, tok_d, p_d, d_vocab)?;
4504 }
4505 } else {
4506 // draft token -> persistent tok_d (next replay's embed reads it; host reads the 4 bytes).
4507 e.argmax_token_device_into(&logits, tok_d, d_vocab)?;
4508 // p-min under a draft mask reads the MASKED row: confidence relative to the
4509 // grammar-LEGAL alternatives (illegal ids leave the softmax denominator), which
4510 // is the right semantics for "does the drafter know what comes next here" and
4511 // the same row the pick came from. Draft-quality only — verify arbitrates.
4512 if with_prob {
4513 e.prob_of_token_device_into(&logits, tok_d, p_d, d_vocab)?;
4514 }
4515 }
4516 }
4517 // ROUND-STREAM K-chain: pack (tok, p) into slot j, then remap tok through d2t so the
4518 // NEXT chained body's embed reads the TARGET id — zero host involvement per step.
4519 if let Some((out, slot, d2t)) = stream_pack {
4520 e.pack_tok_p(tok_d, p_d, out, slot)?;
4521 if let Some(map) = d2t {
4522 e.tok_map_u32(tok_d, map)?;
4523 }
4524 }
4525 // Next draft step's h_seed: pre-norm h_nextn (default) or post-norm final_h (HPOST).
4526 if spec_hpost() {
4527 e.copy_into(h_seed_d, 0, final_h.as_ref().unwrap(), n_embd)?;
4528 } else {
4529 e.copy_into(h_seed_d, 0, &h_nextn, n_embd)?;
4530 }
4531 // advance the draft rope position in-graph.
4532 e.inc_seqlen(pos_d)?;
4533 Ok(())
4534 }
4535
4536 /// Batched target verify forward over `tokens` at positions `pos0..pos0+T` (§D.3, T=K+1).
4537 /// Returns ALL T logit columns (host f32, [T*n_vocab]); appends T cols to every full-attn KV
4538 /// and advances every linear-attn recur state by T steps (the recur steps are SEQUENTIAL T=1).
4539 /// Advances `cache.pos` by T.
4540 pub fn decode_step_t(
4541 &self,
4542 e: &Engine,
4543 tokens: &[u32],
4544 pos0: usize,
4545 cache: &mut Cache,
4546 ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
4547 if self.is_gemma4_e4b() {
4548 return Ok(self.gemma4_e4b_decode_step_t_h(e, tokens, pos0, cache)?.0);
4549 }
4550 if self.gemma_batch_program() {
4551 return self.gemma4_decode_step_t(e, tokens, pos0, cache);
4552 }
4553 Ok(self.decode_step_t_h(e, tokens, pos0, cache)?.0)
4554 }
4555
4556 /// Like `decode_step_t` but ALSO returns the LAST column's pre-output_norm hidden (h_seed for
4557 /// the next draft round). This lets partial-accept replay run as ONE batched T=(n_acc+1) forward
4558 /// (single weight read) instead of n_acc+1 separate T=1 decode_steps (n_acc+1 weight reads).
4559 /// At batch=1 decode is bandwidth-bound, so batching the replay is THE MTP profitability lever.
4560 pub fn decode_step_t_h(
4561 &self,
4562 e: &Engine,
4563 tokens: &[u32],
4564 pos0: usize,
4565 cache: &mut Cache,
4566 ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4567 self.decode_step_t_h_emb(e, tokens, pos0, cache, None)
4568 }
4569
4570 /// Like `decode_step_t_h` with an optional RESIDENT embed table (spec hot loop): device
4571 /// gather instead of host dequant + [T, n_embd] f32 htod. Bit-identical rows.
4572 pub fn decode_step_t_h_emb(
4573 &self,
4574 e: &Engine,
4575 tokens: &[u32],
4576 pos0: usize,
4577 cache: &mut Cache,
4578 embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
4579 ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4580 let (logits_d, h_seed) = self.decode_step_t_h_emb_dev(e, tokens, pos0, cache, embd_dev)?;
4581 Ok((e.dtoh(&logits_d)?, h_seed))
4582 }
4583
4584 /// DEVICE-LOGITS verify forward (spec device-argmax lever): identical kernel chain to
4585 /// `decode_step_t_h_emb` but returns the [T, n_vocab] logits ON DEVICE — the accept walk
4586 /// argmaxes each column on-device and reads back ONE [T] u32 instead of dtoh'ing the full
4587 /// T x n_vocab f32 block (~1-4 MB + T host argmaxes, every round). Kernel dispatch is
4588 /// UNCHANGED (same decode-exact kernels); only the post-logits transfer moves.
4589 pub fn decode_step_t_h_emb_dev(
4590 &self,
4591 e: &Engine,
4592 tokens: &[u32],
4593 pos0: usize,
4594 cache: &mut Cache,
4595 embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
4596 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4597 let n_embd = self.cfg.n_embd as usize;
4598 let t = tokens.len();
4599 let (logits, x) = self.decode_step_t_core(e, tokens, pos0, cache, embd_dev, None)?;
4600 // h_seed for the next round = LAST column's pre-output_norm hidden ([n_embd]).
4601 let mut hs = vbuf(e, n_embd)?; // fully written by copy_view_into below
4602 e.copy_view_into(&mut hs, 0, &x.slice((t - 1) * n_embd..t * n_embd), n_embd)?;
4603 Ok((logits, hs))
4604 }
4605
4606 /// CORE verify forward: the `decode_step_t_h_emb_dev` kernel chain, returning the FULL
4607 /// pre-output_norm hidden stack x ([T, n_embd], any column extractable) and optionally
4608 /// filling a `VerifyCkpt` (retained per-layer state-rebuild inputs) for the REPLAY-FREE
4609 /// partial accept. `ckpt: None` => byte-for-byte the old behavior (the ckpt writes are pure
4610 /// retains/copies — they never change what any kernel computes).
4611 fn decode_step_t_core(
4612 &self,
4613 e: &Engine,
4614 tokens: &[u32],
4615 pos0: usize,
4616 cache: &mut Cache,
4617 embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
4618 mut ckpt: Option<&mut VerifyCkpt>,
4619 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4620 self.decode_step_t_core_stream(
4621 e,
4622 tokens,
4623 pos0,
4624 cache,
4625 embd_dev,
4626 ckpt.take(),
4627 None,
4628 None,
4629 None,
4630 None,
4631 )
4632 }
4633
4634 /// [`Self::decode_step_t_core`] with the MTP route's verify-graph pool armed
4635 /// (`MEMRA_SPEC_VERIFY_GRAPH`). `graphs: None` reproduces `decode_step_t_core`
4636 /// argument-for-argument, so the eager walk stays the byte-identical fallback.
4637 fn decode_step_t_core_vg(
4638 &self,
4639 e: &Engine,
4640 tokens: &[u32],
4641 pos0: usize,
4642 cache: &mut Cache,
4643 embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
4644 mut ckpt: Option<&mut VerifyCkpt>,
4645 graphs: Option<&mut DsparkVerifyGraphs>,
4646 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4647 self.decode_step_t_core_stream(
4648 e,
4649 tokens,
4650 pos0,
4651 cache,
4652 embd_dev,
4653 ckpt.take(),
4654 None,
4655 None,
4656 None,
4657 graphs,
4658 )
4659 }
4660
4661 /// Increment-0 two-session PP seam: release the peer after this lane's stage-0 boundary TX.
4662 /// The two independent sessions keep their own cache/checkpoint state; only issue order moves.
4663 fn decode_step_t_core_pipelined(
4664 &self,
4665 e: &Engine,
4666 tokens: &[u32],
4667 pos0: usize,
4668 cache: &mut Cache,
4669 embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
4670 mut ckpt: Option<&mut VerifyCkpt>,
4671 pipe: &SpecPipeLane,
4672 round: usize,
4673 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4674 let fence = crate::pp::pp_cuts(self.layers.len())
4675 .ok_or("two-session speculative pipeline requires a PP stage cut")?;
4676 if crate::pp::pp2_streams_off() || !crate::pp::spec_pp_on() {
4677 return Err("two-session speculative pipeline requires the PP verify split".into());
4678 }
4679 let interval_fence = pipe.stage0_begin(round)?;
4680 let ticket = self.verify_stage0_issue(
4681 e,
4682 tokens,
4683 pos0,
4684 cache,
4685 embd_dev,
4686 ckpt.as_deref_mut(),
4687 None,
4688 &fence,
4689 Some(interval_fence),
4690 pipe.trace(round),
4691 )?;
4692 pipe.stage0_end(round);
4693 pipe.stage1_begin(round)?;
4694 let result = self.verify_stage1_finish(e, ticket, cache, ckpt, None, &fence, true)?;
4695 pipe.verify_end(round);
4696 Ok(result)
4697 }
4698
4699 /// ROUND-STREAM stage (c) 4: `stream` = (device verify tokens [t], device pos counter) —
4700 /// when Some, rope positions come from pos_iota over the counter, the embed gathers the
4701 /// device tokens, and full_attn_verify routes appends/FA through the _dc twins reading the
4702 /// SAME counter (every layer's kvl.len == cache.pos, one counter drives all three). The
4703 /// host `tokens`/`pos0` args still size buffers (t is FIXED K+1 in stream mode).
4704 /// `vtok_dev` (engine-bundle slice 2): device verify tokens for the EMBED only —
4705 /// unlike `stream` mode it changes nothing else (host pos iota, host-len KV appends).
4706 /// `tokens` then only sizes buffers (the dummy-slice pattern the round-stream arm uses).
4707 #[allow(clippy::too_many_arguments)]
4708 fn decode_step_t_core_stream(
4709 &self,
4710 e: &Engine,
4711 tokens: &[u32],
4712 pos0: usize,
4713 cache: &mut Cache,
4714 embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
4715 mut ckpt: Option<&mut VerifyCkpt>,
4716 stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
4717 pp_pipe: Option<bool>,
4718 vtok_dev: Option<&CudaSlice<u32>>,
4719 graphs: Option<&mut DsparkVerifyGraphs>,
4720 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4721 // PP DOOR (lane/pp2-spec 2026-08-06): the verify trunk now takes its OWN stage split,
4722 // exactly as the eager and batched steps do. This is the single funnel every verify
4723 // forward reaches (decode_step_t / _h / _h_emb / _h_emb_dev / _core all land here), so
4724 // wiring it here wires the whole spec surface — the draft/accept/commit machinery above
4725 // is untouched.
4726 //
4727 // History: pp2-hardening (2026-08-06) made this funnel FAIL CLOSED, because its trunk
4728 // walk was unsplit on one stream and a sharded cross-device placement peer-read every
4729 // remote layer's weights on every spec round (measured 13.9-28x on the batched twin).
4730 // The refusal below survives to cover the residue — MEMRA_SPEC_PP=0, MEMRA_PP_STREAMS=0,
4731 // or a placement whose PpNRt fails to build — so a config that would still walk the
4732 // whole trunk on one stream refuses instead of regressing 28x.
4733 if let Some(fence) = crate::pp::pp_cuts(self.layers.len()) {
4734 if !crate::pp::pp2_streams_off() && crate::pp::spec_pp_on() {
4735 if vtok_dev.is_some() {
4736 return Err(
4737 "device-token dspark verify (slice-2 deferred readback) has no PP \
4738 stage-split arm; set MEMRA_DSPARK_DEFER_READBACK=0 or run the dspark \
4739 route on one device"
4740 .into(),
4741 );
4742 }
4743 return self.decode_step_t_core_ppn(
4744 e,
4745 tokens,
4746 pos0,
4747 cache,
4748 embd_dev,
4749 ckpt.take(),
4750 stream,
4751 &fence,
4752 pp_pipe,
4753 );
4754 }
4755 }
4756 crate::pp::refuse_unsplit_if_remote(
4757 "decode_step_t (spec verify)",
4758 "drop MEMRA_SPEC_PP=0 / MEMRA_PP_STREAMS=0 so the verify trunk takes its OWN stage \
4759 split (decode_step_t_core_ppn); or run spec on one device",
4760 )?;
4761 let cfg = &self.cfg;
4762 let n_embd = cfg.n_embd as usize;
4763 let eps = cfg.rms_eps;
4764 let t = tokens.len();
4765 let pos_d = match stream {
4766 Some((_, ctr)) => {
4767 let mut p = e.alloc_uninit::<i32>(t)?;
4768 e.pos_iota(ctr, &mut p, t)?;
4769 p
4770 }
4771 None => {
4772 let pos_vec: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
4773 e.htod_i32(&pos_vec)?
4774 }
4775 };
4776
4777 // embed T tokens -> [T, n_embd] token-major (device gather on the spec hot loop)
4778 let x = match (stream, embd_dev) {
4779 (Some((vtok, _)), Some((g, qt, rb))) => {
4780 e.embed_gather_device_td(g, vtok, t, n_embd, qt, rb)?
4781 }
4782 (None, Some((g, qt, rb))) => match vtok_dev {
4783 // slice 2: device verify tokens, same embed_gather_u32_t kernel —
4784 // bit-identical rows to the host-token arm (same per-dtype deq).
4785 Some(vt_d) => e.embed_gather_device_td(g, vt_d, t, n_embd, qt, rb)?,
4786 None => e.embed_gather_device_t(g, tokens, n_embd, qt, rb)?,
4787 },
4788 _ => {
4789 assert!(
4790 vtok_dev.is_none(),
4791 "device-token verify requires the resident embed table (embd_dev)"
4792 );
4793 e.htod(&self.embd.gather(n_embd, tokens))?
4794 }
4795 };
4796
4797 // TRUNK WALK: layers [0, n_layers) through the SAME range-scoped subgraph the PP-N
4798 // stage split calls per stage (`verify_layers`) — one code path, so the split cannot
4799 // drift from the unsplit dispatch mirroring. lane/pp2-spec 2026-08-06.
4800 let x = self.verify_layers(
4801 e,
4802 x,
4803 0,
4804 self.layers.len(),
4805 &pos_d,
4806 pos0,
4807 t,
4808 cache,
4809 ckpt.take(),
4810 stream,
4811 graphs,
4812 )?;
4813
4814 let mut hn = vbuf(e, t * n_embd)?;
4815 // Stage-A door: with the serving-class row-outer verify walk, the TAIL must be the
4816 // t=1 decode program per row too (rms_norm t=1 + the single-row bf16 head — the
4817 // split head's concat is receipted bit-identical to it). The batched cuBLASLt head
4818 // is a different ULP class and flips near-tie argmaxes off the greedy tape.
4819 let eager_tail = self.sliding_gated_moe_batch_program()
4820 && std::env::var("MEMRA_SPEC_VERIFY_EAGER").as_deref() == Ok("1");
4821 if eager_tail {
4822 let n_vocab = self.cfg.n_vocab as usize;
4823 let mut logits = vbuf(e, t * n_vocab)?;
4824 for r in 0..t {
4825 let mut row = e.uninit(n_embd)?;
4826 e.dtod_copy_view(&x.slice(r * n_embd..(r + 1) * n_embd), &mut row)?;
4827 let mut hr = e.uninit(n_embd)?;
4828 e.rms_norm(&row, self.output_norm.float_data(), &mut hr, n_embd, 1, eps)?;
4829 let lr = e.matmul(&self.output, &hr, 1)?;
4830 e.dtod_copy_into(&lr, &mut logits, r * n_vocab)?;
4831 e.dtod_copy_into(&hr, &mut hn, r * n_embd)?;
4832 }
4833 if stream.is_none() {
4834 cache.pos += t;
4835 }
4836 return Ok((logits, if spec_hpost() { hn } else { x }));
4837 }
4838 let serving_head =
4839 self.sliding_gated_moe_batch_program() || self.batched_serving_numeric_class();
4840 let logits = if serving_head {
4841 // Step35 and the qwen35 family (MoE 2026-08-14 AM, dense-hybrid same day PM — the
4842 // Q3.8 bring-up reproduced the identical near-tie class on dense: eager-class verify
4843 // vs batched-class live serving, ULP drift amplified through the GDN recurrence)
4844 // serve one batched numeric class at every live width, including B=1. Keep the
4845 // verify head in that same class; other generic families retain the decode-exact
4846 // head that their run-spec contract pins.
4847 e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
4848 e.matmul(&self.output, &hn, t)?
4849 } else {
4850 e.rms_norm_decode(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
4851 e.matmul_decode_exact(&self.output, &hn, t)?
4852 };
4853 // stream: the device pos counter owns position; host mirror reconciles at drain.
4854 if stream.is_none() {
4855 cache.pos += t;
4856 }
4857 // Hidden stack for seeds/refresh-fills: pre-norm x (default) or post-norm hn (HPOST).
4858 Ok((logits, if spec_hpost() { hn } else { x }))
4859 }
4860
4861 /// THE VERIFY TRUNK OVER PP-N (lane/pp2-spec 2026-08-06): `decode_step_t_core_stream`'s walk
4862 /// as N stage subgraphs, each on its own engine/stream (and, under `MEMRA_PP_DEVICES`, its own
4863 /// device), with a `[T, n_embd]` boundary transfer between them. T = K+1 (the verify batch),
4864 /// so this is the batched-boundary shape the pp2-batch lane's grow-only slots already handle
4865 /// (`tx(b, x, t*n_embd)`; the slot grows to the high-water T and the transport moves exactly
4866 /// the payload).
4867 ///
4868 /// Structure is `decode_step_batch_ppn`'s, which is `decode_step_h_ppn`'s. FOUR THINGS ARE
4869 /// PER-STAGE and each for a measured reason (see `decode_step_batch_ppn`'s header for the
4870 /// receipts):
4871 ///
4872 /// 1. THE ENGINE (`rt.engine(s, e)`) — `Engine` owns lazily-grown stable-pointer scratch
4873 /// (`fa_part_pool`, `fa_vf16_scratch`, `argmax_partials`) that is single-stream-safe BY
4874 /// DESIGN. Two stage streams through one Engine is the 2026-08-02 shared-scratch race
4875 /// (35% flake, nondeterministic all-logits divergence). `PpNRt::build` gives every stage
4876 /// s>0 its own Engine even on the primary device; honouring it here is what scopes the
4877 /// pools. The verify path allocates MORE of that scratch than eager decode does (FA at
4878 /// m=T, and the per-layer `GdnStash` retains), so this is load-bearing, not inherited.
4879 ///
4880 /// 2. `pos_d` — each stage uploads its OWN copy of the T rope positions on ITS stream, so the
4881 /// buffer is allocated, consumed and freed on one stream. In `stream` mode that means each
4882 /// stage runs its own `pos_iota` over the SHARED device counter (`pos_ctr`): the counter is
4883 /// read-only during the forward (the round's `inc`/`copy_add` happen outside it), so every
4884 /// stage derives the identical iota, and each stage's own output buffer is stream-local.
4885 ///
4886 /// 3. THE EMBED lives with stage 0 (`self.embd` / `embd_gpu` are host/primary-side; the
4887 /// sharded loader leaves the table with stage 0 by construction).
4888 ///
4889 /// 4. THE HEAD (`output_norm` + `output`) runs on the LAST stage — the sharded loader uploaded
4890 /// both through that stage's engine (`hybrid.rs`: `e_head = layer_engine(e, n_trunk,
4891 /// n_trunk-1)`), so reading them anywhere else is a peer read of the biggest tensor in the
4892 /// model, every round.
4893 ///
4894 /// WHAT STAYS ON THE PRIMARY, deliberately: the returned logits and hidden stack `x`. Both are
4895 /// last-stage-allocated device buffers, and every consumer (the device argmax walk, the accept
4896 /// kernels, `spec_seed_gather`, the ckpt rebuild in `commit_verified_prefix`) reads them
4897 /// through the primary context by UVA — the same read the batched serving epilogue's
4898 /// `last_logits_dev` park does. Those consumers are per-round O(T x n_vocab) and O(n_embd),
4899 /// not per-layer, so they are not the 28x class; splitting them is a separate lane.
4900 ///
4901 /// The MTP HEAD (draft side) is NOT split: it is one block, it lives wherever the loader put
4902 /// it (`load_mtp` uses the primary engine), and it is ~1-2 GB against the trunk's tens. Draft
4903 /// placement is measured, not assumed — see `research/pp2-spec-20260806`.
4904 ///
4905 /// EXACTNESS: PP-N adds ZERO deviation. Each stage runs the SAME kernels on the SAME bytes in
4906 /// the same order via the SAME `verify_layers` the unsplit body calls; the only change is
4907 /// where the residual is materialized, and the boundary is a straight f32 copy (dtod
4908 /// same-device / `cudaMemcpyPeerAsync` cross-device, no conversion). So the split MUST be
4909 /// BIT-IDENTICAL to the unsplit verify at the same T, in both placement orders. Gate:
4910 /// `decode-batch-gate --mode ppspec`. Acceptance counts are a DERIVED consequence — greedy
4911 /// accept argmaxes these logits, so bit-identical logits force identical accept walks; the
4912 /// `run-spec` K=1..8 arm checks that end-to-end rather than trusting the implication.
4913 #[allow(clippy::too_many_arguments)]
4914 fn decode_step_t_core_ppn(
4915 &self,
4916 e: &Engine,
4917 tokens: &[u32],
4918 pos0: usize,
4919 cache: &mut Cache,
4920 embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
4921 mut ckpt: Option<&mut VerifyCkpt>,
4922 stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
4923 fence: &[usize],
4924 pp_pipe: Option<bool>,
4925 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4926 let ticket = self.verify_stage0_issue(
4927 e,
4928 tokens,
4929 pos0,
4930 cache,
4931 embd_dev,
4932 ckpt.as_deref_mut(),
4933 stream,
4934 fence,
4935 pp_pipe,
4936 None,
4937 )?;
4938 self.verify_stage1_finish(e, ticket, cache, ckpt, stream, fence, true)
4939 }
4940
4941 /// Enqueue embed, stage 0, and the first boundary TX, then return the actual boundary slot.
4942 /// The ordinary PP verify wrapper calls `verify_stage1_finish` immediately after this return.
4943 #[allow(clippy::too_many_arguments)]
4944 fn verify_stage0_issue(
4945 &self,
4946 e: &Engine,
4947 tokens: &[u32],
4948 pos0: usize,
4949 cache: &mut Cache,
4950 embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
4951 mut ckpt: Option<&mut VerifyCkpt>,
4952 stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
4953 fence: &[usize],
4954 pp_pipe: Option<bool>,
4955 trace: Option<SpecPipeTraceCtx>,
4956 ) -> Result<VerifyBoundaryTicket, Box<dyn std::error::Error>> {
4957 assert!(
4958 !self.is_gemma4_e4b() && !self.gemma_batch_program(),
4959 "decode_step_t_core_ppn covers the hybrid non-gemma4 verify trunk only \
4960 (the gemma4 arms have their own decode_step_t twins)"
4961 );
4962 if crate::pp::pp_host_bounce_active() && (stream.is_some() || embd_dev.is_some()) {
4963 return Err(
4964 "decode_step_t_core_ppn: refused with MEMRA_PP_HOST_BOUNCE=1 — the trunk \
4965 boundary itself is host-staged, but device-resident verify still peer-reads \
4966 primary-device token/position/embedding buffers from stage 0. Run plain PP \
4967 serving on this host class; spec requires local per-stage inputs first."
4968 .into(),
4969 );
4970 }
4971 let rt = crate::pp::PpNRt::get(e)?;
4972 let n_st = fence.len() - 1;
4973 assert_eq!(
4974 rt.n_stages(),
4975 n_st,
4976 "PpNRt stage count {} != fence stages {n_st}",
4977 rt.n_stages()
4978 );
4979 let n_embd = self.cfg.n_embd as usize;
4980 let t = tokens.len();
4981 let payload = t * n_embd;
4982 if pp_pipe.is_some() {
4983 assert_eq!(n_st, 2, "spec pipeline requires exactly two PP stages");
4984 }
4985 // One-shot lane diagnostic: force natural PP-2 boundaries to completion so the server
4986 // log can price stage 0, the peer hop, the RX copy, and stage 1 + head separately without
4987 // nsys. The ordinary path keeps every enqueue asynchronous. N>2 is deliberately excluded:
4988 // the report below names exactly two stages and must never imply it measured middle ones.
4989 let pp_anatomy = n_st == 2 && std::env::var("MEMRA_SPEC_PP_ANATOMY").as_deref() == Ok("1");
4990 let pp_started = std::time::Instant::now();
4991 let (mut reverse_ms, mut stage0_ms, mut tx_ms) = (0.0f64, 0.0f64, 0.0f64);
4992 // The CALLER's ambient stream, captured BEFORE any `rt.enter()` pushes a stage stream:
4993 // this body returns DEVICE-RESIDENT buffers (the device-argmax accept walk's contract),
4994 // so the exit needs the same publication the boundaries get — see `PpNRt::publish_to`.
4995 // Taken here, not at the end, because inside the last-stage scope `e.stream()` IS the
4996 // stage stream and the wait would self-order into a no-op.
4997 let caller_stream = e.stream();
4998 // #87 ROOT-CAUSE FENCE (lane/pp2spec-crash): the PREVIOUS round's stage-allocated
4999 // outputs (logits/hidden/ckpt stashes) freed stream-ordered on the STAGE streams while
5000 // the primary stream still holds queued reads of them — with event tracking elided,
5001 // nothing stops the pool from reusing those blocks for THIS round's stage allocations,
5002 // whose writes then race the queued reads (measured: 13/4096-NaN random-bits garbage in
5003 // the spec round seed; the full anatomy is on `PpNRt::fence_stages_behind`). Order every
5004 // stage stream behind the caller before enqueueing new stage work.
5005 let reverse_started = std::time::Instant::now();
5006 if pp_pipe != Some(false) {
5007 rt.fence_stages_behind(&caller_stream)?;
5008 }
5009 if pp_pipe == Some(true) {
5010 // Both session verifies must alternate boundary slots even when the ordinary
5011 // decode overlap experiment is off. Prewarm before A's stage 0 so B cannot grow
5012 // slot 1 by synchronizing the RX stream while A's stage 1 is in flight.
5013 rt.prepare_overlap_slots(0, payload)?;
5014 }
5015 if pp_anatomy {
5016 // Drain the reverse-publication dependency before timing stage 0 itself. At c=1 this
5017 // prices any primary-stream rollback/refresh tail inherited from the prior round.
5018 for s in 0..n_st {
5019 let _st = rt.enter(s);
5020 rt.engine(s, e).stream().synchronize()?;
5021 }
5022 reverse_ms = reverse_started.elapsed().as_secs_f64() * 1e3;
5023 }
5024
5025 // Per-stage rope positions: in host mode the same [T] iota each stage uploads itself; in
5026 // stream mode each stage's own `pos_iota` over the shared read-only device counter.
5027 let stage_pos = |es: &Engine| -> Result<CudaSlice<i32>, Box<dyn std::error::Error>> {
5028 match stream {
5029 Some((_, ctr)) => {
5030 let mut p = es.alloc_uninit::<i32>(t)?;
5031 es.pos_iota(ctr, &mut p, t)?;
5032 Ok(p)
5033 }
5034 None => {
5035 let pos_vec: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
5036 es.htod_i32(&pos_vec)
5037 }
5038 }
5039 };
5040
5041 // ---- STAGE 0: embed (the table lives with stage 0) + layers [0, fence[1]) + TX ----
5042 let slot = {
5043 let _st0 = rt.enter(0);
5044 let e0 = rt.engine(0, e);
5045 enqueue_spec_pipe_trace_marker(&e0.stream(), trace.as_ref(), "S0", "start", None)?;
5046 let stage0_started = std::time::Instant::now();
5047 let pos_d = stage_pos(e0)?;
5048 let x = match (stream, embd_dev) {
5049 (Some((vtok, _)), Some((g, qt, rb))) => {
5050 e0.embed_gather_device_td(g, vtok, t, n_embd, qt, rb)?
5051 }
5052 (None, Some((g, qt, rb))) => e0.embed_gather_device_t(g, tokens, n_embd, qt, rb)?,
5053 _ => e0.htod(&self.embd.gather(n_embd, tokens))?,
5054 };
5055 let x = self.verify_layers(
5056 e0,
5057 x,
5058 fence[0],
5059 fence[1],
5060 &pos_d,
5061 pos0,
5062 t,
5063 cache,
5064 ckpt.as_deref_mut(),
5065 stream,
5066 None,
5067 )?;
5068 if pp_anatomy {
5069 e0.stream().synchronize()?;
5070 stage0_ms = stage0_started.elapsed().as_secs_f64() * 1e3;
5071 }
5072 let tx_started = std::time::Instant::now();
5073 let slot = if pp_pipe.is_some() {
5074 rt.tx_pipelined(0, &x, payload)?
5075 } else {
5076 rt.tx(0, &x, payload)?
5077 };
5078 enqueue_spec_pipe_trace_marker(&e0.stream(), trace.as_ref(), "S0", "end", Some(slot))?;
5079 if pp_anatomy {
5080 e0.stream().synchronize()?;
5081 tx_ms = tx_started.elapsed().as_secs_f64() * 1e3;
5082 }
5083 slot
5084 // x + pos_d drop here: freed stream-ordered on stage-0's stream after use.
5085 };
5086
5087 Ok(VerifyBoundaryTicket {
5088 rt,
5089 caller_stream,
5090 slot,
5091 pos0,
5092 t,
5093 payload,
5094 n_st,
5095 pipelined: pp_pipe.is_some(),
5096 pp_anatomy,
5097 pp_started,
5098 reverse_ms,
5099 stage0_ms,
5100 tx_ms,
5101 trace,
5102 })
5103 }
5104
5105 /// Consume a stage-0 boundary ticket and enqueue the remaining PP stages plus the head.
5106 /// On PP-2 this is exactly stage 1; PP-N keeps its pre-existing middle-stage walk here.
5107 #[allow(clippy::too_many_arguments)]
5108 fn verify_stage1_finish(
5109 &self,
5110 e: &Engine,
5111 ticket: VerifyBoundaryTicket,
5112 cache: &mut Cache,
5113 mut ckpt: Option<&mut VerifyCkpt>,
5114 stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
5115 fence: &[usize],
5116 publish_to_caller: bool,
5117 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
5118 let VerifyBoundaryTicket {
5119 rt,
5120 caller_stream,
5121 slot,
5122 pos0,
5123 t,
5124 payload,
5125 n_st,
5126 pipelined,
5127 pp_anatomy,
5128 pp_started,
5129 reverse_ms,
5130 stage0_ms,
5131 tx_ms,
5132 trace,
5133 } = ticket;
5134 let n_embd = self.cfg.n_embd as usize;
5135 let eps = self.cfg.rms_eps;
5136 let mut slot = slot;
5137 let (mut rx_ms, mut stage1_ms) = (0.0f64, 0.0f64);
5138 let stage_pos = |es: &Engine| -> Result<CudaSlice<i32>, Box<dyn std::error::Error>> {
5139 match stream {
5140 Some((_, ctr)) => {
5141 let mut p = es.alloc_uninit::<i32>(t)?;
5142 es.pos_iota(ctr, &mut p, t)?;
5143 Ok(p)
5144 }
5145 None => {
5146 let pos_vec: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
5147 es.htod_i32(&pos_vec)
5148 }
5149 }
5150 };
5151
5152 // ---- MIDDLE STAGES: RX boundary s-1 -> range -> TX boundary s ----
5153 for s in 1..n_st - 1 {
5154 let _st = rt.enter(s);
5155 let es = rt.engine(s, e);
5156 let pos_d = stage_pos(es)?;
5157 let x = rt.rx(s - 1, slot, payload)?;
5158 let x = self.verify_layers(
5159 es,
5160 x,
5161 fence[s],
5162 fence[s + 1],
5163 &pos_d,
5164 pos0,
5165 t,
5166 cache,
5167 ckpt.as_deref_mut(),
5168 stream,
5169 None,
5170 )?;
5171 slot = if pipelined {
5172 rt.tx_pipelined(s, &x, payload)?
5173 } else {
5174 rt.tx(s, &x, payload)?
5175 };
5176 }
5177
5178 // ---- LAST STAGE: RX + final range + output_norm + lm head ----
5179 let _stl = rt.enter(n_st - 1);
5180 let el = rt.engine(n_st - 1, e);
5181 let pos_d = stage_pos(el)?;
5182 let rx_started = std::time::Instant::now();
5183 let x = rt.rx(n_st - 2, slot, payload)?;
5184 if pp_anatomy {
5185 el.stream().synchronize()?;
5186 rx_ms = rx_started.elapsed().as_secs_f64() * 1e3;
5187 }
5188 enqueue_spec_pipe_trace_marker(&el.stream(), trace.as_ref(), "S1", "start", Some(slot))?;
5189 let stage1_started = std::time::Instant::now();
5190 let x = self.verify_layers(
5191 el,
5192 x,
5193 fence[n_st - 1],
5194 fence[n_st],
5195 &pos_d,
5196 pos0,
5197 t,
5198 cache,
5199 ckpt.as_deref_mut(),
5200 stream,
5201 None,
5202 )?;
5203
5204 let mut hn = vbuf(el, payload)?;
5205 let logits = if self.sliding_gated_moe_batch_program() {
5206 // The PP Step35 serving path uses rms_norm + matmul for B=1 as well as B>1.
5207 // Verify must not switch numeric class merely because the same session speculates.
5208 el.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
5209 el.matmul(&self.output, &hn, t)?
5210 } else {
5211 el.rms_norm_decode(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
5212 el.matmul_decode_exact(&self.output, &hn, t)?
5213 };
5214 enqueue_spec_pipe_trace_marker(&el.stream(), trace.as_ref(), "S1", "end", Some(slot))?;
5215 if pp_anatomy {
5216 el.stream().synchronize()?;
5217 stage1_ms = stage1_started.elapsed().as_secs_f64() * 1e3;
5218 }
5219 // EXIT PUBLICATION: both returned buffers are still being produced on the last stage's
5220 // stream. Order the caller's stream behind that work before the buffers escape this
5221 // scope (the 2026-08-06 same-device ppspec find: without it the caller's primary-stream
5222 // consumer read unwritten logits — nondeterministic, one-device-only, and it poisoned
5223 // the following arm's KV in the same process).
5224 if publish_to_caller {
5225 rt.publish_to(n_st - 1, &caller_stream)?;
5226 }
5227 if pp_anatomy {
5228 if publish_to_caller {
5229 caller_stream.synchronize()?;
5230 }
5231 eprintln!(
5232 "[spec-pp-anatomy] t={t} reverse={reverse_ms:.3}ms stage0={stage0_ms:.3}ms \
5233 tx={tx_ms:.3}ms rx={rx_ms:.3}ms stage1-head={stage1_ms:.3}ms total={:.3}ms",
5234 pp_started.elapsed().as_secs_f64() * 1e3,
5235 );
5236 }
5237 // stream: the device pos counter owns position; host mirror reconciles at drain.
5238 if stream.is_none() {
5239 cache.pos += t;
5240 }
5241 Ok((logits, if spec_hpost() { hn } else { x }))
5242 }
5243
5244 /// Step3.5/Step3.7 verify trunk in the serving batched numeric class.
5245 ///
5246 /// `step35_decode_batch_layers` is now authoritative at every live serving width, including
5247 /// B=1 (lane/cx-b1fix). The older verify walk deliberately mirrored the eager T=1 class:
5248 /// it replayed `step35_decode_attn` per row and used the eager/decode-exact FFN dispatch.
5249 /// Those classes are individually stable, but a near-tie prompt can choose different greedy
5250 /// bytes when a request moves from batched plain serving into speculative verify. Run the
5251 /// same authoritative B=1 stage subgraph for each verify row here. Rows still advance
5252 /// layer-by-layer, so every layer sees the preceding verify rows in its attention cache while
5253 /// every norm/projection/FFN uses exactly the live serving dispatch.
5254 #[allow(clippy::too_many_arguments)]
5255 fn step35_verify_batch_layers(
5256 &self,
5257 e: &Engine,
5258 mut x: CudaSlice<f32>,
5259 lo: usize,
5260 hi: usize,
5261 pos0: usize,
5262 t: usize,
5263 cache: &mut Cache,
5264 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5265 let n_embd = self.cfg.n_embd as usize;
5266 if !self.uses_sliding_gated_moe_program() {
5267 return Err(
5268 "serving-class verify requires sliding-gated-MoE canonical operations".into(),
5269 );
5270 }
5271 // SERVING-CLASS VERIFY (MEMRA_SPEC_VERIFY_EAGER=1, step37 MTP bring-up): each verify
5272 // column rides decode_layers_eager — the EXACT t=1 program live serving runs (all TP2
5273 // doors) — row-outer, so row r's appends land before row r+1 attends: bit-equal to
5274 // plain greedy by construction. Only the unsplit full-range walk qualifies; PP splits
5275 // and the tap path keep the batch-layer class.
5276 static VE: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
5277 let eager_verify = *VE
5278 .get_or_init(|| std::env::var("MEMRA_SPEC_VERIFY_EAGER").as_deref() == Ok("1"))
5279 && lo == 0
5280 && hi == self.layers.len();
5281 if eager_verify {
5282 // T-COLUMN LAYER-OUTER WALK (MEMRA_SPEC_VERIFY_TCOL=1): per layer, one t-grid
5283 // attn norm + ONE weight-amortized QKV(+gate) over all T columns, then each
5284 // column runs the UNMODIFIED t=1 attention program via the col-select door and
5285 // the ordinary residual/FFN body. Values per column are bit-equal to the
5286 // row-outer walk: rms over the materialized residual == the fused add+norm
5287 // (kernel_check identity), the tcol kernel's per-column FP order == the t=1
5288 // kernel, and every downstream op IS the t=1 program.
5289 static TCOL: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
5290 let tcol =
5291 *TCOL.get_or_init(|| std::env::var("MEMRA_SPEC_VERIFY_TCOL").as_deref() == Ok("1"));
5292 if tcol && t >= 2 && t <= 8 {
5293 // MEMRA_TCOL_PROF=1: synchronized per-segment wall profile of the walk
5294 // (norm+QKV precompute / per-col attention / per-col residual+FFN). The
5295 // syncs serialize the stream, so the split is for TARGETING amortization
5296 // work only — never a perf claim.
5297 static PROF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
5298 let prof =
5299 *PROF.get_or_init(|| std::env::var("MEMRA_TCOL_PROF").as_deref() == Ok("1"));
5300 let mut prof_ms = [0f64; 3];
5301 let eps = self.cfg.rms_eps;
5302 let mut x_t = x;
5303 let mut h_t = e.uninit(t * n_embd)?;
5304 let mut h_row = e.uninit(n_embd)?; // real row: the non-dcw fallback reads it
5305 // Per-column pos buffers hoisted out of the layer loop (a per-col-per-layer
5306 // pageable htod was an in-stream engine turnaround x t x 45).
5307 let mut pos_rows = Vec::with_capacity(t);
5308 for r in 0..t {
5309 pos_rows.push(e.htod_i32(&[(pos0 + r) as i32])?);
5310 }
5311 let mut ok = true;
5312 // MEMRA_TCOL_OPROJ=1: defer each column's o_proj — the finish seam
5313 // stashes `gated` instead of joining per column; one b4_tcol per rank +
5314 // one slab join produce every column's `mixed` after the attention pass.
5315 // Bit-exact per column (t=1 b4 program per column; elementwise join).
5316 // MEMRA_TCOL_FFN=1 (implies the o_proj defer): when every column of a
5317 // MoE layer deferred, the residual norm runs as one t-grid launch
5318 // (per-row program == t=1) and the FFN as ONE two-column device-routed
5319 // sweep + per-column shexp — the two columns' expert weights dedup
5320 // through L2 instead of reading HBM twice.
5321 static FFN2: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
5322 let ffn_batch =
5323 *FFN2.get_or_init(|| std::env::var("MEMRA_TCOL_FFN").as_deref() == Ok("1"));
5324 let oproj_batch = crate::tp::tcol_oproj_on() || ffn_batch;
5325 let mut mixed_row = e.uninit(n_embd)?;
5326 for il in lo..hi {
5327 let layer = &self.layers[il];
5328 let mut seg = std::time::Instant::now();
5329 e.rms_norm(&x_t, layer.attn_norm.float_data(), &mut h_t, n_embd, t, eps)?;
5330 if !self.step35_verify_qkv_precompute(e, il, &h_t, t)? {
5331 ok = false;
5332 break;
5333 }
5334 if prof {
5335 e.stream().synchronize()?;
5336 prof_ms[0] += seg.elapsed().as_secs_f64() * 1e3;
5337 seg = std::time::Instant::now();
5338 }
5339 let mut next = e.uninit(t * n_embd)?;
5340 // Columns whose o_proj was deferred (their FFN runs after the join).
5341 // A NON-deferred column's FFN must run INSIDE the column loop: the
5342 // oproj-tail handoff is a single cell that the same column's
5343 // residual_norm_ffn consumes before the next column's finish.
5344 let mut deferred: Vec<usize> = Vec::new();
5345 let mut ffn_col =
5346 |r: usize,
5347 mixed: &CudaSlice<f32>,
5348 next: &mut CudaSlice<f32>|
5349 -> Result<(), Box<dyn std::error::Error>> {
5350 let mut x_row = e.uninit(n_embd)?;
5351 e.dtod_copy_view(&x_t.slice(r * n_embd..(r + 1) * n_embd), &mut x_row)?;
5352 let (x1, ffn_out) =
5353 self.residual_norm_ffn(e, layer, &x_row, mixed, n_embd, il, eps)?;
5354 let mut x2 = e.uninit(n_embd)?;
5355 e.add(&x1, &ffn_out, &mut x2, n_embd)?;
5356 e.dtod_copy_into(&x2, next, r * n_embd)?;
5357 Ok(())
5358 };
5359 for r in 0..t {
5360 e.dtod_copy_view(&h_t.slice(r * n_embd..(r + 1) * n_embd), &mut h_row)?;
5361 let row_pos = &pos_rows[r];
5362 crate::tp::set_verify_tcol(Some(r));
5363 if oproj_batch {
5364 crate::tp::set_tcol_oproj_defer(Some(r));
5365 }
5366 let mixed = match &layer.mixer {
5367 crate::hybrid::Mixer::Full(fa) => {
5368 self.full_attn_decode(e, fa, &h_row, row_pos, pos0 + r, cache, il)
5369 }
5370 _ => Err("step35 verify expects full attention".into()),
5371 };
5372 crate::tp::set_verify_tcol(None);
5373 crate::tp::set_tcol_oproj_defer(None);
5374 let mixed = mixed?;
5375 if oproj_batch && crate::tp::take_tcol_oproj_stashed() {
5376 deferred.push(r);
5377 } else {
5378 ffn_col(r, &mixed, &mut next)?;
5379 }
5380 }
5381 if prof {
5382 e.stream().synchronize()?;
5383 prof_ms[1] += seg.elapsed().as_secs_f64() * 1e3;
5384 seg = std::time::Instant::now();
5385 }
5386 if !deferred.is_empty() {
5387 let mixed_t = self.step35_verify_oproj_tcol(e, il, t)?;
5388 let o_out = mixed_t.len() / t;
5389 // Batched t=2 residual+MoE: one t-grid add_rms_norm (per-row
5390 // program == t=1; bit-identical to the oproj-tail join per the
5391 // M2 verbatim-program contract) feeding the two-column routed
5392 // sweep. Ineligible layers (dense FFN, non-nvfp4) fall through
5393 // to the per-column body.
5394 let mut batched = false;
5395 if ffn_batch && t == 2 && deferred.len() == t && o_out == n_embd {
5396 let mut x1_t = e.uninit(t * n_embd)?;
5397 let mut z_t = e.uninit(t * n_embd)?;
5398 e.add_rms_norm(
5399 &x_t,
5400 &mixed_t,
5401 layer.post_attn_norm.float_data(),
5402 &mut x1_t,
5403 &mut z_t,
5404 n_embd,
5405 t,
5406 eps,
5407 )?;
5408 if let Some(ffn_t) = self.step35_verify_moe_t2(e, il, &z_t)? {
5409 let mut x2_t = e.uninit(t * n_embd)?;
5410 e.add(&x1_t, &ffn_t, &mut x2_t, t * n_embd)?;
5411 next = x2_t;
5412 batched = true;
5413 }
5414 }
5415 if !batched {
5416 for &r in &deferred {
5417 e.dtod_copy_view(
5418 &mixed_t.slice(r * o_out..(r + 1) * o_out),
5419 &mut mixed_row,
5420 )?;
5421 ffn_col(r, &mixed_row, &mut next)?;
5422 }
5423 }
5424 }
5425 if prof {
5426 e.stream().synchronize()?;
5427 prof_ms[2] += seg.elapsed().as_secs_f64() * 1e3;
5428 }
5429 drop(ffn_col);
5430 x_t = next;
5431 }
5432 if prof {
5433 eprintln!(
5434 "[tcol-prof] t={t} norm+qkv={:.3}ms attn={:.3}ms ffn={:.3}ms",
5435 prof_ms[0], prof_ms[1], prof_ms[2]
5436 );
5437 }
5438 if ok {
5439 return Ok(x_t);
5440 }
5441 // fall through to the row-outer walk on ineligible layers
5442 x = x_t;
5443 }
5444 let mut next = e.uninit(t * n_embd)?;
5445 for r in 0..t {
5446 let mut row = e.uninit(n_embd)?;
5447 e.dtod_copy_view(&x.slice(r * n_embd..(r + 1) * n_embd), &mut row)?;
5448 let row_pos = e.htod_i32(&[(pos0 + r) as i32])?;
5449 let out = self.decode_layers_eager(e, row, lo, hi, &row_pos, pos0 + r, cache)?;
5450 e.dtod_copy_into(&out, &mut next, r * n_embd)?;
5451 }
5452 // dflash taps are NOT produced on this arm (they need per-layer hiddens the
5453 // row-outer walk does not materialize); the door is a step37 MTP bring-up
5454 // surface where taps are unused.
5455 return Ok(next);
5456 }
5457 let mut ph_last = std::time::Instant::now();
5458 for il in lo..hi {
5459 let mut next = e.uninit(t * n_embd)?;
5460 for r in 0..t {
5461 let mut row = e.uninit(n_embd)?;
5462 e.dtod_copy_view(&x.slice(r * n_embd..(r + 1) * n_embd), &mut row)?;
5463 // The caller owns this verify's position. During controller overlap, cache.pos
5464 // still describes generation N while this stage-0 walk belongs to N+1.
5465 let row_pos = e.htod_i32(&[(pos0 + r) as i32])?;
5466 let mut one = [&mut *cache];
5467 let out = self.step35_decode_batch_layers(
5468 e,
5469 row,
5470 &mut one,
5471 &[(pos0 + r) as i32],
5472 &row_pos,
5473 il,
5474 il + 1,
5475 &mut ph_last,
5476 )?;
5477 e.dtod_copy_into(&out, &mut next, r * n_embd)?;
5478 }
5479 self.dflash_tap(e, cache, il, &next, t)?;
5480 x = next;
5481 }
5482 Ok(x)
5483 }
5484
5485 /// DSpark drafter verify (lane/dspark-q38-recover): one t-row forward through the
5486 /// SERVING-CLASS verify funnel (`decode_step_t_core_stream` — the same numeric class
5487 /// MTP verify rides, GDN state advanced in place), returning per-row argmax tokens.
5488 /// Advances `cache.pos += t`; the caller owns snapshot/rollback (block acceptance is
5489 /// prefix-keep, not all-or-nothing).
5490 pub(crate) fn dspark_verify_t_am(
5491 &self,
5492 e: &Engine,
5493 tokens: &[u32],
5494 pos0: usize,
5495 cache: &mut Cache,
5496 ) -> Result<Vec<u32>, Box<dyn std::error::Error>> {
5497 let (logits, _hn) = self.decode_step_t_core_stream(
5498 e, tokens, pos0, cache, None, None, None, None, None, None,
5499 )?;
5500 let t = tokens.len();
5501 let v = self.output.out_features();
5502 let mut am_d = e.stream().alloc_zeros::<u32>(t)?;
5503 for r in 0..t {
5504 e.argmax_token_device_col(&logits, r, v, &mut am_d, r)?;
5505 }
5506 Ok(e.dtoh_u32(&am_d)?)
5507 }
5508
5509 /// DSpark verify returning the RAW verify logits [t, n_vocab] (device-resident) instead
5510 /// of per-row argmaxes — the sampled-admission arm's input (rejection-sampling accept
5511 /// gathers filtered p from these columns; lane/dspark-sampled-admission-20260820). Same
5512 /// forward as `dspark_verify_t_am`; the greedy arm keeps its argmax wrapper untouched.
5513 pub(crate) fn dspark_verify_t_logits(
5514 &self,
5515 e: &Engine,
5516 tokens: &[u32],
5517 pos0: usize,
5518 cache: &mut Cache,
5519 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5520 let (logits, _hn) = self.decode_step_t_core_stream(
5521 e, tokens, pos0, cache, None, None, None, None, None, None,
5522 )?;
5523 Ok(logits)
5524 }
5525
5526 /// DSpark verify with the MTP column-stash armed: identical forward to
5527 /// `dspark_verify_t_am`, but fills a `VerifyCkpt` so a partial accept can restore
5528 /// column state directly (`dspark_commit_prefix`) instead of snapshot-replay.
5529 /// The ckpt type is opaque outside spec.rs (newtype) — dflash.rs threads it through.
5530 pub(crate) fn dspark_verify_t_am_ckpt(
5531 &self,
5532 e: &Engine,
5533 tokens: &[u32],
5534 pos0: usize,
5535 cache: &mut Cache,
5536 ) -> Result<(Vec<u32>, DsparkVerifyCkpt), Box<dyn std::error::Error>> {
5537 let mut ck = VerifyCkpt::new(self.layers.len());
5538 let (logits, _hn) = self.decode_step_t_core_stream(
5539 e,
5540 tokens,
5541 pos0,
5542 cache,
5543 None,
5544 Some(&mut ck),
5545 None,
5546 None,
5547 None,
5548 None,
5549 )?;
5550 let t = tokens.len();
5551 let v = self.output.out_features();
5552 let mut am_d = e.stream().alloc_zeros::<u32>(t)?;
5553 for r in 0..t {
5554 e.argmax_token_device_col(&logits, r, v, &mut am_d, r)?;
5555 }
5556 Ok((e.dtoh_u32(&am_d)?, DsparkVerifyCkpt(ck)))
5557 }
5558
5559 /// Engine-bundle slice 2: `dspark_verify_t_am_ckpt` with DEVICE tokens and NO readback.
5560 /// The verify tokens are the round's `chain_d` (cand layout: [anchor, drafts...]); the
5561 /// embed gathers its first `t` entries on-device (`embed_gather_u32_t` — bit-identical
5562 /// rows to the host arm), so the host never blocks on the draft chain before dispatching
5563 /// verify. Returns the device per-row argmax buffer; the caller merges its readback with
5564 /// the chain's into ONE sync. Forward, ckpt fill and argmax walk are `_ckpt` verbatim.
5565 pub(crate) fn dspark_verify_t_am_ckpt_dev(
5566 &self,
5567 e: &Engine,
5568 vtok: &CudaSlice<u32>,
5569 t: usize,
5570 pos0: usize,
5571 cache: &mut Cache,
5572 embd_dev: (&CudaSlice<u8>, i32, usize),
5573 graphs: Option<&mut DsparkVerifyGraphs>,
5574 ) -> Result<(CudaSlice<u32>, DsparkVerifyCkpt), Box<dyn std::error::Error>> {
5575 debug_assert!(
5576 vtok.len() >= t,
5577 "verify window exceeds the device token buffer"
5578 );
5579 // The slab flag is a per-round statement: clear it here so a verify that never
5580 // reaches the graphs door (rowwise env, a non-tparallel arm) cannot leave a
5581 // stale `true` steering the commit at slabs the round never wrote.
5582 let mut graphs = graphs;
5583 if let Some(g) = graphs.as_deref_mut() {
5584 g.round_slab = false;
5585 }
5586 let mut ck = VerifyCkpt::new(self.layers.len());
5587 // Dummy host tokens size the funnel; the embed reads `vtok` (the round-stream
5588 // arm's established pattern — spec.rs stream-mode verify does the same).
5589 let dummy = vec![0u32; t];
5590 let (logits, _hn) = self.decode_step_t_core_stream(
5591 e,
5592 &dummy,
5593 pos0,
5594 cache,
5595 Some(embd_dev),
5596 Some(&mut ck),
5597 None,
5598 None,
5599 Some(vtok),
5600 graphs,
5601 )?;
5602 let v = self.output.out_features();
5603 let mut am_d = e.stream().alloc_zeros::<u32>(t)?;
5604 for r in 0..t {
5605 e.argmax_token_device_col(&logits, r, v, &mut am_d, r)?;
5606 }
5607 Ok((am_d, DsparkVerifyCkpt(ck)))
5608 }
5609
5610 /// Ckpt-armed twin of [`Self::dspark_verify_t_logits`] (sampled-admission arm).
5611 pub(crate) fn dspark_verify_t_logits_ckpt(
5612 &self,
5613 e: &Engine,
5614 tokens: &[u32],
5615 pos0: usize,
5616 cache: &mut Cache,
5617 ) -> Result<(CudaSlice<f32>, DsparkVerifyCkpt), Box<dyn std::error::Error>> {
5618 let mut ck = VerifyCkpt::new(self.layers.len());
5619 let (logits, _hn) = self.decode_step_t_core_stream(
5620 e,
5621 tokens,
5622 pos0,
5623 cache,
5624 None,
5625 Some(&mut ck),
5626 None,
5627 None,
5628 None,
5629 None,
5630 )?;
5631 Ok((logits, DsparkVerifyCkpt(ck)))
5632 }
5633
5634 /// Restore the round to `keep` accepted columns from the verify stash: KV lens and
5635 /// pos from the pre-verify snapshot + keep, GDN conv/ssm from the stashed column
5636 /// state — no replay forward. The exact `commit_verified_prefix` the MTP path ships.
5637 pub(crate) fn dspark_commit_prefix(
5638 &self,
5639 e: &Engine,
5640 cache: &mut Cache,
5641 snap: &crate::cache::CacheSnapshot,
5642 ckpt: &DsparkVerifyCkpt,
5643 keep: usize,
5644 ) -> Result<(), Box<dyn std::error::Error>> {
5645 self.commit_verified_prefix(e, cache, snap, &ckpt.0, keep, false, None)
5646 }
5647
5648 /// Slice-3 commit twin: restore to `keep` accepted columns when the round's linear
5649 /// column stash lives in the graphs ctx's persistent slabs (`DsparkVerifyGraphs`) —
5650 /// the cols arm's exact semantics (KV lens + pos from the snapshot, GDN conv/ssm
5651 /// from the stash of column keep-1), slab-addressed and batched into two copy
5652 /// launches. `MEMRA_STATE_COPY_BATCH=0` falls back to per-layer view copies.
5653 pub(crate) fn dspark_commit_prefix_slab(
5654 &self,
5655 e: &Engine,
5656 cache: &mut Cache,
5657 snap: &crate::cache::CacheSnapshot,
5658 ctx: &DsparkVerifyGraphs,
5659 keep: usize,
5660 ) -> Result<(), Box<dyn std::error::Error>> {
5661 use cudarc::driver::DevicePtr;
5662 debug_assert!(keep >= 1, "keep==0 rounds take the legacy rollback");
5663 let mut conv_src: Vec<u64> = Vec::new();
5664 let mut ssm_src: Vec<u64> = Vec::new();
5665 let mut conv_dst: Vec<u64> = Vec::new();
5666 let mut ssm_dst: Vec<u64> = Vec::new();
5667 for il in 0..self.layers.len() {
5668 if let (Some(kvl), Some(saved)) = (cache.kv[il].as_mut(), snap.kv_len[il]) {
5669 kvl.len = saved + keep;
5670 e.set_i32_one(&mut kvl.len_d, kvl.len as i32)?;
5671 }
5672 if let Some(rl) = cache.recur[il].as_ref() {
5673 let (pc, ps, _cw, _sw) = ctx
5674 .slab_row(e, il, keep - 1)
5675 .ok_or("slab commit: linear layer missing from the graphs ctx")?;
5676 conv_src.push(pc);
5677 ssm_src.push(ps);
5678 let st = &e.gpu.stream();
5679 let (dc, _g0) = rl.conv_state.device_ptr(st);
5680 let (ds, _g1) = rl.ssm_state.device_ptr(st);
5681 conv_dst.push(dc as u64);
5682 ssm_dst.push(ds as u64);
5683 }
5684 }
5685 let n = conv_src.len();
5686 if n > 0 {
5687 if state_copy_batch_on() {
5688 let mut tt = vec![0u64; 2 * n];
5689 tt[..n].copy_from_slice(&conv_src);
5690 tt[n..].copy_from_slice(&conv_dst);
5691 let ct = e.htod_u64(&tt)?;
5692 tt[..n].copy_from_slice(&ssm_src);
5693 tt[n..].copy_from_slice(&ssm_dst);
5694 let st = e.htod_u64(&tt)?;
5695 e.copy_batch_uniform_f32(&ct, n, ctx.conv_words)?;
5696 e.copy_batch_uniform_f32(&st, n, ctx.ssm_words)?;
5697 } else {
5698 let (cw, sw) = (ctx.conv_words, ctx.ssm_words);
5699 let row = keep - 1;
5700 for il in 0..self.layers.len() {
5701 let Some(rl) = cache.recur[il].as_mut() else {
5702 continue;
5703 };
5704 let k = ctx.lin_pos[&il];
5705 {
5706 let sv = e.view(&ctx.stash_conv[k], (row + 1) * cw);
5707 let win = sv.slice(row * cw..(row + 1) * cw);
5708 e.copy_view_into(&mut rl.conv_state, 0, &win, cw)?;
5709 }
5710 {
5711 let sv = e.view(&ctx.stash_ssm[k], (row + 1) * sw);
5712 let win = sv.slice(row * sw..(row + 1) * sw);
5713 e.copy_view_into(&mut rl.ssm_state, 0, &win, sw)?;
5714 }
5715 }
5716 }
5717 }
5718 cache.pos = snap.pos + keep;
5719 Ok(())
5720 }
5721
5722 /// Qwen35-family verify trunk in the live serving numeric class.
5723 ///
5724 /// Serving intentionally keeps this architecture in the generic batched program even at
5725 /// B=1. The older verify walk used its own mirrored dispatch and can flip near-tie argmaxes.
5726 ///
5727 /// Two arms, one numeric class:
5728 /// - DENSE GDN (`DenseMlp`, t<=16): `qwen35_verify_tparallel` — the weight ops (norms,
5729 /// projections, FFN) hoist to m=T through the exact-tier batched kernels whose per-row
5730 /// program IS the m=1 program (`matmul_pre == fused2 per (tensor,row); _bN mmvq per-row
5731 /// == m=1` — decode_batch.rs v2 note), while the state ops (conv ring, gdn scan, KV
5732 /// append, fa decode) stay a per-row loop running the b_n=1 serving kernels with each
5733 /// row's own t_kv-driven arm pick (the straddle law: every row executes the exact
5734 /// program its isolated serving step would). One weight read per layer per round
5735 /// instead of T — this is what makes MTP profitable in the exact class (the per-row
5736 /// walk measured verify(K+1) ~= (K+1) plain steps: 69 -> 44 tok/s served, 2026-08-15).
5737 /// - MoE / t>16 / `MEMRA_SPEC_VERIFY_ROWWISE=1`: the per-row replay of the authoritative
5738 /// serving layer body, preserving single-session autoregressive cache order (the
5739 /// correctness reference; also the rollback seam for the t-parallel arm).
5740 ///
5741 /// Bit-identity of the t-parallel arm vs the rowwise arm is gated by spec-serve-gate
5742 /// (zero differing logits at T=1..4, K arms) + the 8-prompt ON/OFF canary before ship.
5743 #[allow(clippy::too_many_arguments)]
5744 fn qwen35_verify_batch_layers(
5745 &self,
5746 e: &Engine,
5747 x: CudaSlice<f32>,
5748 lo: usize,
5749 hi: usize,
5750 pos0: usize,
5751 t: usize,
5752 cache: &mut Cache,
5753 ckpt: Option<&mut VerifyCkpt>,
5754 stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
5755 graphs: Option<&mut DsparkVerifyGraphs>,
5756 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5757 // Qwen35Moe admitted 2026-08-20 (lane/draftcost-moe): the t-parallel arm already
5758 // carries the MoE FFN (`moe_ffn_il_zq8` at m=T) and the GDN per-row state loop; the
5759 // arch fence was a qualification gate, not a mechanism gap. Measured disease on the
5760 // 35B-A3B class: rowwise verify ~= 5.6 ms per drafted token (one full trunk step
5761 // each) — the same (K+1)-plain-steps wall the dense admission fixed on 2026-08-15.
5762 // Rollback seam unchanged: MEMRA_SPEC_VERIFY_ROWWISE=1.
5763 let rowwise = std::env::var("MEMRA_SPEC_VERIFY_ROWWISE").as_deref() == Ok("1")
5764 || !self.batched_serving_numeric_class()
5765 || t > 16;
5766 if rowwise {
5767 if stream.is_some() {
5768 // rowwise replays per row with host cache.pos — irreconcilable with a
5769 // device position counter. Burst callers must keep t <= 16 and the
5770 // ROWWISE env unset; refusing beats silently mispositioned rows.
5771 return Err("qwen35 rowwise verify has no ROUND-STREAM arm \
5772 (t > 16 or MEMRA_SPEC_VERIFY_ROWWISE=1)"
5773 .into());
5774 }
5775 self.qwen35_verify_rowwise(e, x, lo, hi, pos0, t, cache, ckpt)
5776 } else {
5777 self.qwen35_verify_tparallel(e, x, lo, hi, pos0, t, cache, ckpt, stream, graphs)
5778 }
5779 }
5780
5781 /// The per-row correctness reference: replay each verify row through the authoritative
5782 /// serving layer body (`decode_batch_layers` at b_n=1). T full weight reads per layer.
5783 #[allow(clippy::too_many_arguments)]
5784 fn qwen35_verify_rowwise(
5785 &self,
5786 e: &Engine,
5787 mut x: CudaSlice<f32>,
5788 lo: usize,
5789 hi: usize,
5790 pos0: usize,
5791 t: usize,
5792 cache: &mut Cache,
5793 mut ckpt: Option<&mut VerifyCkpt>,
5794 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5795 let n_embd = self.cfg.n_embd as usize;
5796 let saved_pos = cache.pos;
5797 let mut ph_last = std::time::Instant::now();
5798 for il in lo..hi {
5799 let mut next = e.uninit(t * n_embd)?;
5800 let mut col_states: Option<Vec<(CudaSlice<f32>, CudaSlice<f32>)>> =
5801 if ckpt.is_some() && t >= 2 && matches!(self.layers[il].mixer, Mixer::Linear(_)) {
5802 Some(Vec::with_capacity(t - 1))
5803 } else {
5804 None
5805 };
5806 for r in 0..t {
5807 cache.pos = pos0 + r;
5808 let mut row = e.uninit(n_embd)?;
5809 e.dtod_copy_view(&x.slice(r * n_embd..(r + 1) * n_embd), &mut row)?;
5810 let row_pos = e.htod_i32(&[(pos0 + r) as i32])?;
5811 let mut one = [&mut *cache];
5812 let ctx = self.batch_layer_ctx(e, &one, il, il + 1)?;
5813 let out = match self.decode_batch_layers(
5814 e,
5815 row,
5816 &mut one,
5817 &ctx,
5818 &row_pos,
5819 &mut ph_last,
5820 ) {
5821 Ok(out) => out,
5822 Err(error) => {
5823 cache.pos = saved_pos;
5824 return Err(error);
5825 }
5826 };
5827 e.dtod_copy_into(&out, &mut next, r * n_embd)?;
5828 if r + 1 < t {
5829 if let Some(states) = col_states.as_mut() {
5830 let recur = cache.recur[il]
5831 .as_ref()
5832 .ok_or("Qwen35-MoE linear verify layer has no recurrent state")?;
5833 states.push((
5834 e.clone_dtod(&recur.conv_state)?,
5835 e.clone_dtod(&recur.ssm_state)?,
5836 ));
5837 }
5838 }
5839 }
5840 if let (Some(checkpoint), Some(states)) = (ckpt.as_deref_mut(), col_states) {
5841 checkpoint.cols[il] = Some(states);
5842 }
5843 x = next;
5844 }
5845 cache.pos = saved_pos;
5846 Ok(x)
5847 }
5848
5849 /// T-PARALLEL VERIFY IN THE SERVING NUMERIC CLASS (lane/tparallel-verify, 2026-08-15).
5850 ///
5851 /// The weight ops run ONCE per layer at m=T; the state ops run per row through the same
5852 /// b_n=1 serving kernels the rowwise replay uses. Per-row bit-identity rests on the two
5853 /// pins the serving batch tier already carries:
5854 /// * `matmul_pre` / `_bN` mmvq: per-row program == m=1 program (decode_batch.rs v2 note,
5855 /// kernel-check pinned) — so a [T, n_embd] projection row equals the row projected
5856 /// alone;
5857 /// * row-indexed norms/elementwise (`rms_norm`, `quantize_q8_1`, `add_rms_norm`,
5858 /// `gated_rmsnorm[_q8_1]`, `silu_mul`, `rope_neox` with per-row positions): the T-row
5859 /// launch is the per-row program (same pin the generic verify's fused norms rely on).
5860 /// The sequential dependencies keep their exact serving order: the conv ring / gdn scan
5861 /// chain state row -> row through the `_b` kernels at b_n=1 (ping-pong via a 6-entry
5862 /// alternating pointer table, host handles swapped per row so VerifyCkpt clones the
5863 /// canonical state exactly as the rowwise arm does), and each row's KV append + fa decode
5864 /// picks its arm from ITS OWN t_kv (append: format-only; fa: `fa_seqs_eligible` + its own
5865 /// `fa_split_keys` rung at b_n=1) — the straddle law per row, so every row executes the
5866 /// program its isolated B=1 serving step would.
5867 ///
5868 /// Cost: 1 weight read per layer per round + T state micro-launches, vs the rowwise arm's
5869 /// T weight reads. Gated bit-identical vs the rowwise arm by spec-serve-gate + canary.
5870 #[allow(clippy::too_many_arguments)]
5871 fn qwen35_verify_tparallel(
5872 &self,
5873 e: &Engine,
5874 mut x: CudaSlice<f32>,
5875 lo: usize,
5876 hi: usize,
5877 pos0: usize,
5878 t: usize,
5879 cache: &mut Cache,
5880 mut ckpt: Option<&mut VerifyCkpt>,
5881 stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
5882 mut graphs: Option<&mut DsparkVerifyGraphs>,
5883 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5884 let seqs_append =
5885 std::env::var("MEMRA_BATCH_APPEND").as_deref() != Ok("0") && !Engine::kv_fp8_on();
5886 let batch_fa_on = std::env::var("MEMRA_BATCH_FA").as_deref() != Ok("0");
5887
5888 // Merge guard (v0.98 train, re-affirmed on the v0.100 train over slice 4c): the
5889 // ROUND-STREAM arm (lane/draftcost-moe, device position counter) and the dspark
5890 // verify graphs (engine-bundle slice 3 / trunk slice 4c) have no common caller —
5891 // stream rides the qwen35moe burst, graphs ride the dspark route. If a future
5892 // caller arms both, refuse loudly instead of silently dropping the graphs ctx
5893 // (the stream linear arm takes linear_attn_verify_t, not the graphed segment or
5894 // full-verify bodies).
5895 if stream.is_some() && graphs.is_some() {
5896 return Err(
5897 "qwen35 tparallel verify: ROUND-STREAM and dspark verify graphs \
5898 cannot arm together"
5899 .into(),
5900 );
5901 }
5902 // Engine-bundle slice 3 + slice 4c: with a graphs ctx armed, pointer tables are
5903 // refreshed once per verify (the gdn ping-pong moves handles; a fresh generation
5904 // moves the kv caches). Then:
5905 // - slice 4c: when the WHOLE round rides one seqs rung (every row batchable, one
5906 // split-ladder step, rung covers the round), the ENTIRE walk replays as ONE
5907 // full-verify graph per (vt, rung) — linear layers through the shared
5908 // `qwen35_tparallel_linear_layer` body, full-attention layers through the
5909 // shared `qwen35_tparallel_fa_layer` body in graph mode.
5910 // - fallback (straddle rounds, below the vec floor, partial walks): runs of
5911 // consecutive LINEAR layers replay the slice-3 per-(segment, vt) graphs and
5912 // the full-attention layers run eager (batched rows when eligible).
5913 if let Some(g) = graphs.as_deref_mut() {
5914 g.refresh_tables(e, cache)?;
5915 g.round_slab = false;
5916 if let Some(rung) = g.full_rung(self, cache, lo, hi, t, seqs_append && batch_fa_on) {
5917 // Pool ceiling (dspark_vg_cap): an existing key always replays; a NEW
5918 // full capture past the ceiling falls through to the segment/eager arms.
5919 if g.full.contains_key(&(t, rung, hi)) || g.can_capture() {
5920 let out = g.run_full(self, e, lo, hi, &x, t, pos0, rung, cache)?;
5921 g.round_slab = true;
5922 return Ok(out);
5923 }
5924 }
5925 // Round-atomic ceiling check for the segment door: if any linear run in this
5926 // walk would need a NEW capture past the ceiling, the whole round runs the
5927 // eager cols-ckpt walk (mixing slab- and cols-stashed layers in one round
5928 // would corrupt the commit).
5929 if !g.segments_ready(self, lo, hi, t) {
5930 graphs = None;
5931 }
5932 }
5933 // STREAM (2b, lane/draftcost-moe): positions come from the device round counter
5934 // (pos_iota / i32_copy_add) so a burst round needs no host position knowledge.
5935 let pos_d = match stream {
5936 Some((_, ctr)) => {
5937 let mut p = e.alloc_uninit::<i32>(t)?;
5938 e.pos_iota(ctr, &mut p, t)?;
5939 p
5940 }
5941 None => {
5942 let pos_host: Vec<i32> = (0..t).map(|r| (pos0 + r) as i32).collect();
5943 e.htod_i32(&pos_host)?
5944 }
5945 };
5946 // Per-row 1-element position buffers, built ONCE per verify (the append/fa wrappers
5947 // take owned pos slices; building these inside the layer x row loops cost 16xT H2Ds).
5948 // LAZY since slice 4: the batched fa/append arm never touches them — they are built
5949 // on the first per-row fallback layer only (stream-aware there; the stream FA arm
5950 // rides the dc rows kernels and never reaches the fallback).
5951 let mut pos_rows: Option<Vec<CudaSlice<i32>>> = None;
5952 let mut il = lo;
5953 while il < hi {
5954 if graphs.is_some() && matches!(self.layers[il].mixer, Mixer::Linear(_)) {
5955 let mut end = il;
5956 while end < hi && matches!(self.layers[end].mixer, Mixer::Linear(_)) {
5957 end += 1;
5958 }
5959 let g = graphs.as_deref_mut().expect("checked above");
5960 x = g.run_segment(self, e, il, end, &x, t, cache)?;
5961 g.round_slab = true;
5962 il = end;
5963 continue;
5964 }
5965 let layer = &self.layers[il];
5966 if stream.is_none() && matches!(layer.mixer, Mixer::Linear(_)) {
5967 // Eager linear layer (no graphs ctx): the shared body, legacy cols-ckpt arm.
5968 // Under ROUND-STREAM the linear layers ride the fa-body match's stream arm
5969 // below (linear_attn_verify_t — the stream COMMIT needs its GdnStash).
5970 x = self.qwen35_tparallel_linear_layer(
5971 e,
5972 il,
5973 &x,
5974 t,
5975 cache,
5976 ckpt.as_deref_mut(),
5977 None,
5978 None,
5979 )?;
5980 il += 1;
5981 continue;
5982 }
5983 // Full-attention (or stream-Linear, or MLA-refusing) layer: the extracted
5984 // shared body — eager arm (fresh per-verify pos/table, exact t_kv sizing,
5985 // in-body len bump). The slice-4c captured full-verify graphs run the SAME
5986 // body in graph mode; under ROUND-STREAM the body's dc-rows / GDN stream arms
5987 // run (lane/draftcost-moe).
5988 x = self.qwen35_tparallel_fa_layer(
5989 e,
5990 il,
5991 &x,
5992 t,
5993 cache,
5994 FaLayerArgs {
5995 pos_d: &pos_d,
5996 pos_rows: &mut pos_rows,
5997 pos0,
5998 seqs_append,
5999 batch_fa_on,
6000 graph_cap: None,
6001 stream,
6002 ckpt: ckpt.as_deref_mut(),
6003 },
6004 )?;
6005 il += 1;
6006 }
6007 Ok(x)
6008 }
6009
6010 /// SHARED dense-FFN body for the qwen35 t-parallel layers (trunk-kernels slice B) —
6011 /// ONE copy for the fa and linear layer bodies (the verify_layers extraction lesson).
6012 /// Dual arm (MEMRA_TK_FFN_DUAL, default on): gate+up in ONE dual launch from the
6013 /// pre-quantized activation with macro-scales DEFERRED into the fused SwiGLU+q8_1
6014 /// epilogue, then ffn_down from the fused (aq, ad) — the q27 verify chain verbatim.
6015 /// Every door is the bit-identical proven one: `matmul_decode_exact_dual_pre` (per
6016 /// (tensor,token,row) == the two singles), `silu_mul_scaled_q8_1` (y*s inline == the
6017 /// scale_inplace store, value-exact; fused quantize == quantize_q8_1 bytes),
6018 /// `matmul_decode_exact_pre` (dispatch mirror of the singles' q8_1-fast tail).
6019 /// Dual-refused (t outside 2..=7, non-NVFP4, layout mismatch) or seam off -> the
6020 /// original singles chain, byte-for-byte.
6021 #[allow(clippy::too_many_arguments)]
6022 fn qwen35_tparallel_dense_ffn(
6023 &self,
6024 e: &Engine,
6025 ffn_gate: &crate::model::GpuTensor,
6026 ffn_up: &crate::model::GpuTensor,
6027 ffn_down: &crate::model::GpuTensor,
6028 zn: &CudaSlice<f32>,
6029 t: usize,
6030 n_embd: usize,
6031 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6032 let n_ff = ffn_gate.out_features();
6033 let (zq, zd) = e.quantize_q8_1(zn, t, n_embd)?;
6034 if Engine::tk_ffn_dual_on() {
6035 if let Some(((g, gs), (u, us))) =
6036 e.matmul_decode_exact_dual_pre(ffn_gate, ffn_up, &zq, &zd, t)?
6037 {
6038 if e.uses_q8_1_fast(ffn_down) {
6039 let (aq, ad) = e.silu_mul_scaled_q8_1(&g, &u, gs, us, t * n_ff)?;
6040 return e.matmul_decode_exact_pre(ffn_down, &aq, &ad, t);
6041 }
6042 let mut act = e.uninit(t * n_ff)?;
6043 e.silu_mul_scaled(&g, &u, gs, us, &mut act, t * n_ff)?;
6044 let (aq, ad) = e.quantize_q8_1(&act, t, n_ff)?;
6045 return e.matmul_pre(ffn_down, &aq, &ad, &act, t);
6046 }
6047 }
6048 // v1 singles chain (seam off or dual-refused) — the pre-slice-B body verbatim.
6049 let g = e.matmul_pre(ffn_gate, &zq, &zd, zn, t)?;
6050 let u = e.matmul_pre(ffn_up, &zq, &zd, zn, t)?;
6051 let mut act = e.uninit(t * n_ff)?;
6052 e.silu_mul(&g, &u, &mut act, t * n_ff)?;
6053 let (aq, ad) = e.quantize_q8_1(&act, t, n_ff)?;
6054 e.matmul_pre(ffn_down, &aq, &ad, &act, t)
6055 }
6056
6057 /// ONE t-parallel FULL-ATTENTION layer (attn_norm + fa mixer + post_attn_norm + FFN +
6058 /// tap) — extracted from the walk exactly like `qwen35_tparallel_linear_layer` so the
6059 /// eager walk and the slice-4c captured full-verify graphs execute the SAME body (a
6060 /// second copy is how dispatch mirrors drift — the verify_layers extraction lesson).
6061 ///
6062 /// `args.graph_cap = Some((table, off, rung_end))` is the captured-graph mode:
6063 /// - kv base-pointer pairs come from the ctx-owned persistent table at `off` (a fresh
6064 /// generation's cache lands at new addresses that only the per-verify table refresh
6065 /// knows — the slice-3 baked-address lesson);
6066 /// - the seqs twins size partials/grid at `rung_end` and pin `split_keys` to the
6067 /// rung's ladder value: `n_splits_max` is pure stride, splits >= ns_eff write the
6068 /// EMPTY partial the combine never reads, and every per-row T_kv derives in-kernel
6069 /// from `pos_seq[z]` — so one captured launch replays bit-identically for every
6070 /// round whose rows all sit inside the rung;
6071 /// - the host len bump moves to the replay caller (captured host code does not
6072 /// re-run at replay).
6073 /// Graph mode REFUSES any round the batched arm cannot take: the per-row fallback
6074 /// host-branches on t_kv and must never be captured.
6075 #[allow(clippy::too_many_arguments)]
6076 fn qwen35_tparallel_fa_layer(
6077 &self,
6078 e: &Engine,
6079 il: usize,
6080 x: &CudaSlice<f32>,
6081 t: usize,
6082 cache: &mut Cache,
6083 args: FaLayerArgs<'_>,
6084 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6085 use cudarc::driver::DevicePtr;
6086 let cfg = &self.cfg;
6087 let n_embd = cfg.n_embd as usize;
6088 let eps = cfg.rms_eps;
6089 let head_dim_global = cfg.head_dim_k as usize;
6090 let layer = &self.layers[il];
6091 let FaLayerArgs {
6092 pos_d,
6093 pos_rows,
6094 pos0,
6095 seqs_append,
6096 batch_fa_on,
6097 graph_cap,
6098 stream,
6099 mut ckpt,
6100 } = args;
6101
6102 // ---- attn_norm + q8_1 quantize at m=T (row-indexed == per-row) ----
6103 let anorm = layer.attn_norm.float_data();
6104 let mut xn = e.uninit(t * n_embd)?;
6105 e.rms_norm(x, anorm, &mut xn, n_embd, t, eps)?;
6106 let (hq, hd) = e.quantize_q8_1(&xn, t, n_embd)?;
6107
6108 let mixed: CudaSlice<f32> = match &layer.mixer {
6109 Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
6110 // STREAM ARM (2b, lane/draftcost-moe): under a device position counter the
6111 // per-row serving-kernel chain cannot run (host state swaps keyed on host
6112 // row index are fine, but the stream COMMIT needs the GdnStash for its _dc
6113 // rebuild — the per-row chain only produces per-column clones). GDN rides
6114 // `linear_attn_verify_t`: batched q8_1-class projections, stash-producing,
6115 // and its one-scan recurrence is pinned bit-identical to T chained T=1
6116 // steps (its header + kernel-check). Position-independent, so no counter
6117 // plumbing is needed. Guards mirror the generic call site exactly.
6118 Mixer::Linear(la) if stream.is_some() => {
6119 if !(t >= 3 || (t == 2 && spec_m2()))
6120 || !self.mixer_in_q8_1_fast(e, &layer.mixer)
6121 || !e.uses_q8_1_fast(&la.ssm_out)
6122 {
6123 return Err("qwen35 stream verify: GDN batched arm requires t>=3 \
6124 (or MEMRA_SPEC_M2 at t=2) and q8_1-fast projections"
6125 .into());
6126 }
6127 let want = ckpt.is_some();
6128 let (out, stash) =
6129 self.linear_attn_verify_t(e, la, &xn, Some((&hq, &hd)), t, cache, il, want)?;
6130 if let (Some(ck), Some(st)) = (ckpt.as_deref_mut(), stash) {
6131 ck.gdn[il] = Some(st);
6132 }
6133 out
6134 }
6135 Mixer::Linear(_) => {
6136 unreachable!("linear layers ride qwen35_tparallel_linear_layer")
6137 }
6138 Mixer::Full(fa) => {
6139 let geometry = cfg.full_attention_geometry_at(il as u32);
6140 let n_head = geometry.n_head as usize;
6141 let n_head_kv = geometry.n_head_kv as usize;
6142 let head_dim = geometry.head_dim_k as usize;
6143 let rope_dims = geometry.n_rot as usize;
6144 let rope_base = geometry.rope_base;
6145 let scale = geometry.attention_scale();
6146 // Batched projections: one weight read serves all T rows.
6147 // GROUP-3 twin (trunk-kernels slice D): q/k/v in ONE launch — the group4
6148 // kernel with n3=0, bit-identical per (tensor, token, row) to the three
6149 // singles; refused or MEMRA_TK_FA_GROUP=0 -> singles byte-for-byte.
6150 let (qf, mut k, v) = match e.matmul_decode_exact_group3_pre(
6151 [&fa.wq, &fa.wk, &fa.wv],
6152 &hq,
6153 &hd,
6154 t,
6155 )? {
6156 Some(mut g3) => {
6157 let v = g3.pop().unwrap();
6158 let k = g3.pop().unwrap();
6159 let qf = g3.pop().unwrap();
6160 (qf, k, v)
6161 }
6162 None => (
6163 e.matmul_pre(&fa.wq, &hq, &hd, &xn, t)?,
6164 e.matmul_pre(&fa.wk, &hq, &hd, &xn, t)?,
6165 e.matmul_pre(&fa.wv, &hq, &hd, &xn, t)?,
6166 ),
6167 };
6168 let gated =
6169 geometry.attention_gate == memra_gguf::config::AttentionGateKind::FusedQ;
6170 let (mut q, gate) = if gated {
6171 let mut qs = e.uninit(t * n_head * head_dim)?;
6172 let mut gs = e.uninit(t * n_head * head_dim)?;
6173 e.q_gate_split(&qf, &mut qs, &mut gs, head_dim, n_head, t)?;
6174 (qs, Some(gs))
6175 } else {
6176 (qf, None)
6177 };
6178 let mut qn = e.uninit(t * n_head * head_dim)?;
6179 e.rms_norm(
6180 &q,
6181 fa.q_norm.float_data(),
6182 &mut qn,
6183 head_dim,
6184 t * n_head,
6185 eps,
6186 )?;
6187 q = qn;
6188 let mut kn = e.uninit(t * n_head_kv * head_dim)?;
6189 e.rms_norm(
6190 &k,
6191 fa.k_norm.float_data(),
6192 &mut kn,
6193 head_dim,
6194 t * n_head_kv,
6195 eps,
6196 )?;
6197 k = kn;
6198 e.rope_neox(
6199 &mut q, pos_d, head_dim, rope_dims, n_head, t, rope_base, 1.0,
6200 )?;
6201 e.rope_neox(
6202 &mut k, pos_d, head_dim, rope_dims, n_head_kv, t, rope_base, 1.0,
6203 )?;
6204
6205 // Per-row append + attend: row r sees rows 0..r in KV (causal within the
6206 // draft), each through the b_n=1 serving kernels at its own t_kv.
6207 let q_dim = n_head * head_dim;
6208 let kv_dim = n_head_kv * head_dim;
6209 let mut attn = e.uninit(t * q_dim)?;
6210 let (kdk, kdv, ktb, vtb, len0, kv_local) = {
6211 let kvl = cache.kv[il].as_ref().unwrap();
6212 // [2T] interleaved k,v base pointers: entry pair z serves row z of
6213 // the batched twins; the per-row fallback reads pair 0 (same cache
6214 // for every row of one layer). Graph mode reads the ctx table.
6215 let local: Option<CudaSlice<u64>> = match graph_cap {
6216 Some(_) => None,
6217 None => {
6218 let s = &e.gpu.stream();
6219 let (pk, _g) = kvl.k.device_ptr(s);
6220 let (pv, _g2) = kvl.v.device_ptr(s);
6221 let mut tbl = Vec::with_capacity(2 * t);
6222 for _ in 0..t {
6223 tbl.push(pk as u64);
6224 tbl.push(pv as u64);
6225 }
6226 Some(e.htod_u64(&tbl)?)
6227 }
6228 };
6229 (
6230 kvl.kv_dim_k,
6231 kvl.kv_dim_v,
6232 kvl.k_tok_bytes,
6233 kvl.v_tok_bytes,
6234 kvl.len,
6235 local,
6236 )
6237 };
6238 let (kv_tbl, kv_off): (&CudaSlice<u64>, usize) = match graph_cap {
6239 Some((tb, off, _)) => (tb, off),
6240 None => (kv_local.as_ref().expect("built above"), 0),
6241 };
6242 // Slice 4 (fa/append rows — see dspark_fa_rows_on): the whole per-row
6243 // section batches into the z-batched serving twins when every row of
6244 // this round takes the v4-seqs arm on ONE fa_split_keys rung. Both
6245 // guards are evaluated at the round's FIRST and LAST t_kv — the
6246 // eligibility window (vec floor .. v4 max) and each split-ladder rung
6247 // are intervals in t_kv, so ends-inside means all-inside (the straddle
6248 // law). Appending all T rows before any attend is read-equivalent to
6249 // the interleaved order: row r's walk reads keys 0..len0+r only, and
6250 // rows > r land at slots it never touches; every written cache row is
6251 // the per-token appender's exact warp program (kernel-check pinned).
6252 let t_kv_first = len0 + 1;
6253 let t_kv_last = len0 + t;
6254 let rows_batched = t >= 2
6255 && seqs_append
6256 && batch_fa_on
6257 && dspark_fa_rows_on()
6258 // the z-batched twins read stacked rows at the CACHE's kv dims;
6259 // the projection stack is [T, n_head_kv*head_dim] — they must be
6260 // the same stride or row z misaligns (true for this family; the
6261 // guard keeps any asymmetric-kv model on the per-row loop).
6262 && kdk == kv_dim
6263 && kdv == kv_dim
6264 && crate::fa_seqs_eligible(t_kv_first, head_dim_global)
6265 && crate::fa_seqs_eligible(t_kv_last, head_dim_global)
6266 && crate::fa_split_keys(t_kv_first, cfg.n_head_kv as usize)
6267 == crate::fa_split_keys(t_kv_last, cfg.n_head_kv as usize);
6268 // Sizing: eager = exact round bound; graph mode = the rung end (stride +
6269 // grid only — bytes proven equal above). Capture-time invariants refuse
6270 // loudly rather than bake a divergent body.
6271 let (size_kv_max, sp) = match graph_cap {
6272 Some((_, _, rung)) => {
6273 if !rows_batched {
6274 return Err(format!(
6275 "fa graph capture: layer {il} round is not batchable \
6276 (t_kv {t_kv_first}..{t_kv_last}) — the per-row fallback \
6277 must never be captured"
6278 )
6279 .into());
6280 }
6281 let sp_r = crate::fa_split_keys(rung, cfg.n_head_kv as usize);
6282 if t_kv_last > rung
6283 || sp_r != crate::fa_split_keys(t_kv_last, cfg.n_head_kv as usize)
6284 {
6285 return Err(format!(
6286 "fa graph capture: rung {rung} does not cover round \
6287 t_kv {t_kv_first}..{t_kv_last} on one split ladder step"
6288 )
6289 .into());
6290 }
6291 (rung, sp_r)
6292 }
6293 None => (
6294 t_kv_last,
6295 crate::fa_split_keys(t_kv_last, cfg.n_head_kv as usize),
6296 ),
6297 };
6298 if let Some((_, ctr)) = stream {
6299 // STREAM ARM (2b): one batched dc append + the multi-row dc attention
6300 // — the generic stream arm's exact shape (rows kernels are pinned
6301 // byte-identical to the per-row programs by kernel-check). Host len
6302 // stays a stale lower bound; the burst drain reconciles it.
6303 let kvl = cache.kv[il].as_mut().unwrap();
6304 e.append_kv_quantized_rows_dc(
6305 &k,
6306 &v,
6307 &mut kvl.k,
6308 &mut kvl.v,
6309 ctr,
6310 t,
6311 kdk,
6312 kdv,
6313 ktb,
6314 vtb,
6315 Engine::kv_fp8_on(),
6316 )?;
6317 let upper = (kvl.len + t + 64).min(cache.max_ctx);
6318 let k_view = e.view_u8(&kvl.k, upper * ktb);
6319 let v_view = e.view_u8(&kvl.v, upper * vtb);
6320 e.fa_decode_rows_dc(
6321 &q, &k_view, &v_view, &mut attn, head_dim, n_head, n_head_kv, ctr, upper,
6322 t, scale, ktb, vtb, 0, false,
6323 )?;
6324 } else if rows_batched {
6325 e.append_kv_quantized_seqs(
6326 &k,
6327 &v,
6328 &kv_tbl.slice(kv_off..kv_off + 2 * t),
6329 pos_d,
6330 t,
6331 kdk,
6332 kdv,
6333 ktb,
6334 vtb,
6335 )?;
6336 if graph_cap.is_none() {
6337 cache.kv[il].as_mut().unwrap().len += t;
6338 }
6339 e.fa_decode_batch_seqs_v4(
6340 &q,
6341 &kv_tbl.slice(kv_off..kv_off + 2 * t),
6342 pos_d,
6343 &mut attn,
6344 head_dim,
6345 n_head,
6346 n_head_kv,
6347 t,
6348 size_kv_max,
6349 scale,
6350 sp,
6351 ktb,
6352 vtb,
6353 )?;
6354 } else {
6355 if pos_rows.is_none() {
6356 // Stream-aware for symmetry with pos_d (the stream FA arm rides
6357 // the dc rows kernels above and never reaches this fallback).
6358 *pos_rows = Some(match stream {
6359 Some((_, ctr)) => (0..t)
6360 .map(|r| {
6361 let mut b = e.alloc_uninit::<i32>(1)?;
6362 e.i32_copy_add(ctr, &mut b, r as i32)?;
6363 Ok(b)
6364 })
6365 .collect::<Result<_, Box<dyn std::error::Error>>>()?,
6366 None => (0..t)
6367 .map(|r| e.htod_i32(&[(pos0 + r) as i32]))
6368 .collect::<Result<_, _>>()?,
6369 });
6370 }
6371 let pos_rows = pos_rows.as_ref().unwrap();
6372 for r in 0..t {
6373 // Owned per-row scratch: the b_n=1 kernels take packed batch buffers
6374 // whose row 0 is this row (arithmetic-free materialization copies,
6375 // same as decode's per-seq fallback arm).
6376 let mut k_row = e.uninit(kv_dim)?;
6377 e.dtod_copy_view(&k.slice(r * kv_dim..(r + 1) * kv_dim), &mut k_row)?;
6378 let mut v_row = e.uninit(kv_dim)?;
6379 e.dtod_copy_view(&v.slice(r * kv_dim..(r + 1) * kv_dim), &mut v_row)?;
6380 let pos_row = &pos_rows[r];
6381 let kvl = cache.kv[il].as_mut().unwrap();
6382 if seqs_append {
6383 e.append_kv_quantized_seqs(
6384 &k_row,
6385 &v_row,
6386 &kv_tbl.slice(kv_off..kv_off + 2),
6387 pos_row,
6388 1,
6389 kdk,
6390 kdv,
6391 ktb,
6392 vtb,
6393 )?;
6394 kvl.len += 1;
6395 } else {
6396 e.append_kv_quantized_view(
6397 &k_row.slice(0..kv_dim),
6398 &v_row.slice(0..kv_dim),
6399 &mut kvl.k,
6400 &mut kvl.v,
6401 kvl.len,
6402 kvl.kv_dim_k,
6403 kvl.kv_dim_v,
6404 kvl.k_tok_bytes,
6405 kvl.v_tok_bytes,
6406 Engine::kv_fp8_on(),
6407 )?;
6408 kvl.len += 1;
6409 }
6410 let t_kv = kvl.len;
6411 let mut q_row = e.uninit(q_dim)?;
6412 e.dtod_copy_view(&q.slice(r * q_dim..(r + 1) * q_dim), &mut q_row)?;
6413 let mut a_row = e.uninit(q_dim)?;
6414 if batch_fa_on && crate::fa_seqs_eligible(t_kv, head_dim_global) {
6415 let sp0_r = crate::fa_split_keys(t_kv, cfg.n_head_kv as usize);
6416 e.fa_decode_batch_seqs_v4(
6417 &q_row,
6418 &kv_tbl.slice(kv_off..kv_off + 2),
6419 pos_row,
6420 &mut a_row,
6421 head_dim,
6422 n_head,
6423 n_head_kv,
6424 1,
6425 t_kv,
6426 scale,
6427 sp0_r,
6428 ktb,
6429 vtb,
6430 )?;
6431 } else {
6432 let k_view = e.view_u8(&kvl.k, t_kv * kvl.k_tok_bytes);
6433 let v_view = e.view_u8(&kvl.v, t_kv * kvl.v_tok_bytes);
6434 let mut a_view = a_row.slice_mut(0..q_dim);
6435 e.fa_decode_kvmod_view(
6436 &q_row.slice(0..q_dim),
6437 &k_view,
6438 &v_view,
6439 &mut a_view,
6440 head_dim,
6441 n_head,
6442 n_head_kv,
6443 t_kv,
6444 scale,
6445 kvl.k_tok_bytes,
6446 kvl.v_tok_bytes,
6447 Engine::kv_fp8_on(),
6448 )?;
6449 }
6450 e.dtod_copy_into(&a_row, &mut attn, r * q_dim)?;
6451 }
6452 }
6453
6454 // Output gate (element-wise) + o-proj at m=T.
6455 let attn_g = match &gate {
6456 Some(g) => {
6457 let n = t * q_dim;
6458 let mut gsig = e.uninit(n)?;
6459 e.sigmoid(g, &mut gsig, n)?;
6460 let mut ag = e.uninit(n)?;
6461 e.mul(&attn, &gsig, &mut ag, n)?;
6462 ag
6463 }
6464 None => attn,
6465 };
6466 e.matmul(&fa.wo, &attn_g, t)?
6467 }
6468 };
6469
6470 // ---- residual add + post_attn_norm + FFN at m=T (serving dispatch verbatim) ----
6471 let pnorm = layer.post_attn_norm.float_data();
6472 let mut x1 = e.uninit(t * n_embd)?;
6473 let mut zn = e.uninit(t * n_embd)?;
6474 e.add_rms_norm(x, &mixed, pnorm, &mut x1, &mut zn, n_embd, t, eps)?;
6475 let ffn_out = match &layer.ffn {
6476 crate::hybrid::Ffn::Dense {
6477 ffn_gate,
6478 ffn_up,
6479 ffn_down,
6480 } => {
6481 assert!(
6482 self.cfg.m3.is_none(),
6483 "qwen35 t-parallel verify: M3 swigluoai FFN not yet batched"
6484 );
6485 self.qwen35_tparallel_dense_ffn(e, ffn_gate, ffn_up, ffn_down, &zn, t, n_embd)?
6486 }
6487 crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il_zq8(e, m, &zn, None, t, il as u16)?,
6488 };
6489 let mut x2 = e.uninit(t * n_embd)?;
6490 e.add(&x1, &ffn_out, &mut x2, t * n_embd)?;
6491 // dspark drafter tap (no-op when no sink armed): post-layer residual verify rows
6492 self.dflash_tap(e, cache, il, &x2, t)?;
6493 Ok(x2)
6494 }
6495
6496 /// ONE t-parallel LINEAR layer (attn_norm + gdn mixer + post_attn_norm + FFN + tap) —
6497 /// the exact body the old in-loop Linear arm ran, extracted so the eager walk and the
6498 /// slice-3 captured segments execute the SAME code (a second copy is how dispatch
6499 /// mirrors drift — the verify_layers extraction lesson). Two deliberate changes, both
6500 /// bit-identical by construction:
6501 /// - the gdn ping-pong host swap moves from per-row to ONE end-of-body swap (t odd):
6502 /// the device sequence is driven entirely by the 6-entry pointer table, which
6503 /// already encodes both parities; the ckpt stash reads name row r's out buffer
6504 /// directly (r even -> alt handle, odd -> canonical) — the same physical bytes the
6505 /// legacy post-swap clone read.
6506 /// - `stash` (slice-3 ctx): persistent per-layer slabs written by copy_into instead of
6507 /// per-row clone_dtod allocs — same bytes, capture-legal (no per-round host objects).
6508 /// `table_src` = (persistent pointer table, offset) when the ctx owns the tables;
6509 /// None builds the per-verify table exactly as before.
6510 #[allow(clippy::too_many_arguments)]
6511 fn qwen35_tparallel_linear_layer(
6512 &self,
6513 e: &Engine,
6514 il: usize,
6515 x: &CudaSlice<f32>,
6516 t: usize,
6517 cache: &mut Cache,
6518 mut ckpt: Option<&mut VerifyCkpt>,
6519 stash: Option<(&mut CudaSlice<f32>, &mut CudaSlice<f32>)>,
6520 table_src: Option<(&CudaSlice<u64>, usize)>,
6521 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6522 use cudarc::driver::DevicePtr;
6523 let cfg = &self.cfg;
6524 let n_embd = cfg.n_embd as usize;
6525 let eps = cfg.rms_eps;
6526 let layer = &self.layers[il];
6527 let Mixer::Linear(la) = &layer.mixer else {
6528 return Err("qwen35_tparallel_linear_layer on a non-linear layer".into());
6529 };
6530 // ---- attn_norm + q8_1 quantize at m=T (row-indexed == per-row) ----
6531 let anorm = layer.attn_norm.float_data();
6532 let mut xn = e.uninit(t * n_embd)?;
6533 e.rms_norm(x, anorm, &mut xn, n_embd, t, eps)?;
6534 let (hq, hd) = e.quantize_q8_1(&xn, t, n_embd)?;
6535
6536 let geometry = la.geometry;
6537 let d_state = geometry.key_head_dim as usize;
6538 let num_k = geometry.key_heads as usize;
6539 let num_v = geometry.value_heads as usize;
6540 let d_conv = geometry.conv_kernel as usize;
6541 let key_dim = d_state * num_k;
6542 let value_dim = geometry.value_head_dim as usize * num_v;
6543 let conv_dim = key_dim * 2 + value_dim;
6544 let gdn_scale = 1.0 / (d_state as f32).sqrt();
6545
6546 // ---- batched projections: one weight read for all T rows ----
6547 // GROUP-4 twin (trunk-kernels slice C): the whole 4-tuple in ONE launch, bit-identical
6548 // per (tensor, token, row) to the four singles; refused (layout/tier) or
6549 // MEMRA_TK_GDN_GROUP=0 -> the singles chain byte-for-byte.
6550 let (qkv_mixed, z, beta_raw, alpha) = match e.matmul_decode_exact_group4_pre(
6551 [&la.wqkv, &la.wqkv_gate, &la.ssm_beta, &la.ssm_alpha],
6552 &hq,
6553 &hd,
6554 t,
6555 )? {
6556 Some(mut g4) => {
6557 let alpha = g4.pop().unwrap();
6558 let beta_raw = g4.pop().unwrap();
6559 let z = g4.pop().unwrap();
6560 let qkv_mixed = g4.pop().unwrap();
6561 (qkv_mixed, z, beta_raw, alpha)
6562 }
6563 None => (
6564 e.matmul_pre(&la.wqkv, &hq, &hd, &xn, t)?,
6565 e.matmul_pre(&la.wqkv_gate, &hq, &hd, &xn, t)?,
6566 e.matmul_pre(&la.ssm_beta, &hq, &hd, &xn, t)?,
6567 e.matmul_pre(&la.ssm_alpha, &hq, &hd, &xn, t)?,
6568 ),
6569 };
6570 let beta_w = la.ssm_beta.out_features();
6571 let alpha_w = la.ssm_alpha.out_features();
6572 let qkv_w = la.wqkv.out_features();
6573
6574 // ---- per-row state chain through the b_n=1 serving kernels ----
6575 // 6-entry alternating pointer table expresses the ping-pong without a rebuild per
6576 // row: even rows scan s0 -> s1, odd rows s1 -> s0.
6577 let table_local: Option<CudaSlice<u64>> = match table_src {
6578 Some(_) => None,
6579 None => {
6580 let rl = cache.recur[il].as_ref().unwrap();
6581 let s = &e.gpu.stream();
6582 let (pc, _g0) = rl.conv_state.device_ptr(s);
6583 let (p0, _g1) = rl.ssm_state.device_ptr(s);
6584 let (p1, _g2) = rl.ssm_state_alt.device_ptr(s);
6585 Some(e.htod_u64(&[
6586 pc as u64, p0 as u64, p1 as u64, pc as u64, p1 as u64, p0 as u64,
6587 ])?)
6588 }
6589 };
6590 let (table, toff): (&CudaSlice<u64>, usize) = match table_src {
6591 Some((tb, off)) => (tb, off),
6592 None => (table_local.as_ref().unwrap(), 0),
6593 };
6594 let mut o_all = e.uninit(t * value_dim)?;
6595 let mut col_states: Option<Vec<(CudaSlice<f32>, CudaSlice<f32>)>> =
6596 if ckpt.is_some() && stash.is_none() && t >= 2 {
6597 Some(Vec::with_capacity(t - 1))
6598 } else {
6599 None
6600 };
6601 let mut stash = stash;
6602 // Per-row scratch reused across rows (uninit is cheap but not free at
6603 // 48 layers x T rows); row inputs/outputs pass as VIEWS into the packed
6604 // [T, ...] buffers — zero arithmetic-free copies in this loop.
6605 let mut conv_out = e.uninit(conv_dim)?;
6606 let mut q_l2 = e.uninit(value_dim)?;
6607 let mut k_l2 = e.uninit(value_dim)?;
6608 let mut v_gd = e.uninit(value_dim)?;
6609 let mut beta_b = e.uninit(num_v)?;
6610 let mut g_log = e.uninit(num_v)?;
6611 for r in 0..t {
6612 let base = toff + if r % 2 == 0 { 0 } else { 3 };
6613 let conv_view = table.slice(base..base + 1);
6614 let in_view = table.slice(base + 1..base + 2);
6615 let out_view = table.slice(base + 2..base + 3);
6616 e.ssm_conv1d_fused_decode_b_view(
6617 &qkv_mixed.slice(r * qkv_w..(r + 1) * qkv_w),
6618 &conv_view,
6619 la.ssm_conv1d.float_data(),
6620 &mut conv_out,
6621 conv_dim,
6622 d_conv,
6623 1,
6624 )?;
6625 e.gdn_prep_decode_b_view(
6626 &conv_out,
6627 &beta_raw.slice(r * beta_w..(r + 1) * beta_w),
6628 &alpha.slice(r * alpha_w..(r + 1) * alpha_w),
6629 la.ssm_dt.float_data(),
6630 la.ssm_a.float_data(),
6631 &mut q_l2,
6632 &mut k_l2,
6633 &mut v_gd,
6634 &mut beta_b,
6635 &mut g_log,
6636 d_state,
6637 num_v,
6638 num_k,
6639 key_dim,
6640 eps,
6641 conv_dim,
6642 1,
6643 )?;
6644 let mut o_row = o_all.slice_mut(r * value_dim..(r + 1) * value_dim);
6645 e.gdn_scan_s128_batched_view(
6646 &q_l2, &k_l2, &v_gd, &g_log, &beta_b, &in_view, &out_view, &mut o_row, num_v, 1,
6647 gdn_scale,
6648 )?;
6649 if r + 1 < t {
6650 // Row r's out buffer: even rows write s1 (the alt handle — no swaps ran),
6651 // odd rows write s0 — the same physical state the legacy post-swap
6652 // canonical clone read.
6653 let rl = cache.recur[il]
6654 .as_ref()
6655 .ok_or("qwen35 linear verify layer has no recurrent state")?;
6656 let ssm_src = if r % 2 == 0 {
6657 &rl.ssm_state_alt
6658 } else {
6659 &rl.ssm_state
6660 };
6661 match stash.as_mut() {
6662 Some((conv_slab, ssm_slab)) => {
6663 // BOTH stash reads go through the pointer table at run time: the
6664 // ssm handles ping-pong between rounds, and the ctx (with its
6665 // captured graphs) outlives the Cache — a fresh generation's
6666 // conv/ssm buffers land at new addresses that only the per-round
6667 // table refresh knows. A baked direct copy would read freed
6668 // memory (parity was the slice-3 smoke divergence; cache
6669 // lifetime is the cross-generation twin).
6670 e.copy_indirect_src_f32(
6671 &conv_view,
6672 conv_slab,
6673 r * conv_dim * (d_conv - 1),
6674 conv_dim * (d_conv - 1),
6675 )?;
6676 // The ssm handles PING-PONG between rounds: a captured direct
6677 // copy would bake the capture-time physical buffer and read the
6678 // wrong parity after any odd-vt round (the slice-3 smoke
6679 // divergence). Read the src address from row r's OUT table
6680 // entry at run time — the same entry the scan just wrote.
6681 e.copy_indirect_src_f32(
6682 &out_view,
6683 ssm_slab,
6684 r * d_state * d_state * num_v,
6685 d_state * d_state * num_v,
6686 )?;
6687 }
6688 None => {
6689 if let Some(states) = col_states.as_mut() {
6690 states.push((e.clone_dtod(&rl.conv_state)?, e.clone_dtod(ssm_src)?));
6691 }
6692 }
6693 }
6694 }
6695 }
6696 // ONE end-of-body parity swap (t odd) — the legacy loop swapped per row; the net
6697 // handle motion is identical and the device sequence never read the handles.
6698 if t % 2 == 1 {
6699 let rl = cache.recur[il].as_mut().unwrap();
6700 std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
6701 }
6702 if let (Some(checkpoint), Some(states)) = (ckpt.as_deref_mut(), col_states) {
6703 checkpoint.cols[il] = Some(states);
6704 }
6705
6706 // ---- batched gated norm + out-projection at m=T ----
6707 let mixed = if e.uses_q8_1_fast(&la.ssm_out) {
6708 let (gq, gd) = e.gated_rmsnorm_q8_1(
6709 &o_all,
6710 la.ssm_norm.float_data(),
6711 &z,
6712 d_state,
6713 t * num_v,
6714 eps,
6715 )?;
6716 let g0 = e.zeros(0)?;
6717 e.matmul_pre(&la.ssm_out, &gq, &gd, &g0, t)?
6718 } else {
6719 let mut gn = e.uninit(t * value_dim)?;
6720 e.gated_rmsnorm(
6721 &o_all,
6722 la.ssm_norm.float_data(),
6723 &z,
6724 &mut gn,
6725 d_state,
6726 t * num_v,
6727 eps,
6728 )?;
6729 e.matmul(&la.ssm_out, &gn, t)?
6730 };
6731
6732 // ---- residual add + post_attn_norm + FFN at m=T (serving dispatch verbatim) ----
6733 let pnorm = layer.post_attn_norm.float_data();
6734 let mut x1 = e.uninit(t * n_embd)?;
6735 let mut zn = e.uninit(t * n_embd)?;
6736 e.add_rms_norm(x, &mixed, pnorm, &mut x1, &mut zn, n_embd, t, eps)?;
6737 let ffn_out = match &layer.ffn {
6738 crate::hybrid::Ffn::Dense {
6739 ffn_gate,
6740 ffn_up,
6741 ffn_down,
6742 } => {
6743 assert!(
6744 self.cfg.m3.is_none(),
6745 "qwen35 t-parallel verify: M3 swigluoai FFN not yet batched"
6746 );
6747 self.qwen35_tparallel_dense_ffn(e, ffn_gate, ffn_up, ffn_down, &zn, t, n_embd)?
6748 }
6749 crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il_zq8(e, m, &zn, None, t, il as u16)?,
6750 };
6751 let mut x2 = e.uninit(t * n_embd)?;
6752 e.add(&x1, &ffn_out, &mut x2, t * n_embd)?;
6753 // dspark drafter tap (no-op when no sink armed): post-layer residual verify rows
6754 self.dflash_tap(e, cache, il, &x2, t)?;
6755 Ok(x2)
6756 }
6757
6758 /// PP-N STAGE SUBGRAPH of the verify trunk: layers `[lo, hi)` of `decode_step_t_core_stream`'s
6759 /// walk, verbatim. Enters with a MATERIALIZED `[T, n_embd]` residual (no pending fusion pair
6760 /// carried in from outside the range) and exits with the range's final residual materialized
6761 /// (the trailing add executed) — exactly the `decode_layers_eager(lo, hi)` contract, T rows
6762 /// instead of one.
6763 ///
6764 /// EXTRACTED (lane/pp2-spec 2026-08-06) rather than duplicated: `decode_step_t_core_stream` IS
6765 /// the single funnel every verify forward reaches, and its per-layer dispatch MIRRORING (norm
6766 /// fusion per layer, the t>=3/spec_m2 batched-linear window, the fused-q8 FFN chain, the
6767 /// decode-exact projections) is what makes verify bit-identical to eager decode. A second copy
6768 /// for the split arm is how those mirrors drift apart on the next lever. The unsplit body now
6769 /// calls this with `(0, n_layers)`, so the whole-trunk path and every stage range run the SAME
6770 /// code — there is no "split version" of the verify math.
6771 ///
6772 /// Bit-identity of a cut rests on the same kernel-check-pinned identity the eager arm's cut
6773 /// does — `add_rms_norm_q8_1 == add then rms_norm_q8_1` at nrows=T — because the ONLY thing a
6774 /// fence changes is that the cross-layer fusion carry breaks at `hi-1` and is re-materialized
6775 /// as an explicit `add`. `decode-batch-gate --mode ppspec` verifies end-to-end on real weights.
6776 #[allow(clippy::too_many_arguments)]
6777 fn verify_layers(
6778 &self,
6779 e: &Engine,
6780 mut x: CudaSlice<f32>,
6781 lo: usize,
6782 hi: usize,
6783 pos_d: &CudaSlice<i32>,
6784 pos0: usize,
6785 t: usize,
6786 cache: &mut Cache,
6787 mut ckpt: Option<&mut VerifyCkpt>,
6788 stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
6789 graphs: Option<&mut DsparkVerifyGraphs>,
6790 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6791 if self.sliding_gated_moe_batch_program() {
6792 if stream.is_some() {
6793 return Err(
6794 "step35 has no ROUND-STREAM verify arm (the device-counter _dc twins \
6795 cannot express the SWA offset KV view)"
6796 .into(),
6797 );
6798 }
6799 return self.step35_verify_batch_layers(e, x, lo, hi, pos0, t, cache);
6800 }
6801 if self.batched_serving_numeric_class() {
6802 return self.qwen35_verify_batch_layers(
6803 e,
6804 x,
6805 lo,
6806 hi,
6807 pos0,
6808 t,
6809 cache,
6810 ckpt.take(),
6811 stream,
6812 graphs,
6813 );
6814 }
6815 let n_embd = self.cfg.n_embd as usize;
6816 let eps = self.cfg.rms_eps;
6817 // CROSS-LAYER ADD+NORM FUSION (lane/vt-fixes fix 2, mirroring decode_step_h's
6818 // launch-arc form): layer il's post-FFN residual add (x2 = x1 + ffn_out) and layer
6819 // il+1's attn_norm(+quantize) are consecutive row-wise ops — ONE add_rms_norm_q8_1
6820 // launch at nrows=t does all three (bit-identity pinned by the T-row kernel-check
6821 // arms). Carry the un-added (x1, ffn_out) pair; the fused launch materializes x2 (the
6822 // residual the next layer needs) as its `res` output. Falls back to the separate add
6823 // when the next layer is off the fused-q8 path.
6824 let mut pending: Option<(CudaSlice<f32>, CudaSlice<f32>)> = None;
6825 for il in lo..hi {
6826 let layer = &self.layers[il];
6827 // DISPATCH-MIRRORED attn-input RMSNorm (FP-order lesson #8): eager decode fuses the
6828 // 1024-thread rms_norm_q8_1 ONLY when every mixer projection is q8_1-fast; layers with
6829 // Float projections (ssm_beta/ssm_alpha on layers 1/2/4 of the 9B NVFP4 GGUF) take the
6830 // UNFUSED 256-thread rms_norm. The verify norm must mirror that PER-LAYER choice —
6831 // blockDim changes the sum-of-squares reduce order, and the ULP shift amplifies through
6832 // the GDN recurrence into argmax flips (measured: 9B text prompt, 1 ULP at layer 2 ->
6833 // 2.3e-1 logit maxdiff at the head -> K=1..8 divergence at a 0.03-margin token).
6834 let mixer_fast = self.mixer_in_q8_1_fast(e, &layer.mixer);
6835 let norm_fused = std::env::var("MEMRA_NO_FUSE_NORMQ").is_err() && mixer_fast;
6836 // BATCHED EPILOGUE RE-FUSE (lane/vt-fixes fix 2, 2026-08-03): when the norm is
6837 // dispatch-fused AND every consumer of `h` reads only its q8_1 form (Full mixer:
6838 // projections only; Linear mixer: the batched arm — the per-column fallback needs
6839 // f32 h), emit the attn-input norm DIRECTLY as q8_1 via `rms_norm_q8_1` at nrows=t
6840 // (row-indexed kernel — the T-row launch is the per-row m=1 program, kernel-check
6841 // pins bit-identity vs rms_norm_decode -> quantize_q8_1). Kills the standalone
6842 // quantize launch(es) + the f32 h HBM round-trip that decode never pays.
6843 // step35 (Full mixer) is the third case that needs f32 `h`: its verify arm is a
6844 // per-ROW replay of the eager decode mixer, whose `pre_q` contract is a single row —
6845 // a T-row q8_1 pair cannot be handed to it, and re-deriving per-row q8_1 from the f32
6846 // rows is exactly the dispatch being mirrored. Keep step35 on the unfused arm.
6847 let lin_q8_only = match &layer.mixer {
6848 Mixer::Linear(la) => {
6849 (t >= 3 || (t == 2 && spec_m2())) && e.uses_q8_1_fast(&la.ssm_out)
6850 }
6851 Mixer::Full(_) if self.sliding_gated_moe_batch_program() => false,
6852 _ => true,
6853 };
6854 // NOTE decode.rs's take()-first lesson: take the pending pair BEFORE branching so
6855 // a non-fused layer still performs the residual add.
6856 let taken = pending.take();
6857 let (h, h_q8) = if norm_fused && lin_q8_only {
6858 let pair = match taken {
6859 // fused add + attn_norm + q8_1: ONE launch resolves the carried residual
6860 // AND emits this layer's mixer input pre-quantized. res -> x2 (= new x).
6861 Some((x1p, f1p)) => {
6862 let mut x2 = vbuf(e, t * n_embd)?; // fully written (res output)
6863 let p = e.add_rms_norm_q8_1(
6864 &x1p,
6865 &f1p,
6866 layer.attn_norm.float_data(),
6867 &mut x2,
6868 n_embd,
6869 t,
6870 eps,
6871 )?;
6872 x = x2;
6873 p
6874 }
6875 None => e.rms_norm_q8_1(&x, layer.attn_norm.float_data(), n_embd, t, eps)?,
6876 };
6877 (e.zeros(0)?, Some(pair)) // h unused on this path (q8-only consumers)
6878 } else {
6879 if let Some((x1p, f1p)) = taken {
6880 let mut x2 = vbuf(e, t * n_embd)?; // fully written by add
6881 e.add(&x1p, &f1p, &mut x2, t * n_embd)?;
6882 x = x2;
6883 }
6884 let mut h = vbuf(e, t * n_embd)?; // fully written by either rms_norm arm
6885 if norm_fused {
6886 e.rms_norm_decode(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
6887 } else {
6888 e.rms_norm(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
6889 }
6890 (h, None)
6891 };
6892 let h_q8_ref = h_q8.as_ref().map(|(q, d)| (q, d));
6893
6894 let mixed = match &layer.mixer {
6895 Mixer::Full(fa) => self.full_attn_verify(
6896 e,
6897 fa,
6898 &h,
6899 h_q8_ref,
6900 pos_d,
6901 t,
6902 cache,
6903 il,
6904 stream.map(|(_, c)| c),
6905 )?,
6906 Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
6907 Mixer::Linear(la) => {
6908 // BATCHED linear verify (2026-07-03, the MTP-profit lever): one T-token pass —
6909 // batched projections (weight read ONCE, hits the m=2-4 weight-resident matvec),
6910 // carried-state conv (ssm_conv1d_tm_state), GDN prep on the prefill kernels, and
6911 // ONE gdn_scan whose internal sequential t-loop is the SAME recurrence as T
6912 // chained T=1 steps (bit-identical). Falls back to the sequential per-column
6913 // chain when T < d_conv-1 (conv ring update needs T >= pad) — or when ANY
6914 // projection is off the q8_1 fast path: matmul_decode_exact would route a Float
6915 // tensor to cuBLAS at m=t (different FP accumulation than eager's per-token
6916 // GEMV), so mixed-dtype layers stay on the eager-identical per-column chain.
6917 // MEMRA_SPEC_M2 (lane/spec-m2): the t==2 batch rides the same arm — the conv
6918 // wrapper handles t<pad with a pure-copy ring rebuild; see spec_m2() header.
6919 if (t >= 3 || (t == 2 && spec_m2()))
6920 && mixer_fast
6921 && e.uses_q8_1_fast(&la.ssm_out)
6922 {
6923 let want = ckpt.is_some();
6924 let (out, stash) =
6925 self.linear_attn_verify_t(e, la, &h, h_q8_ref, t, cache, il, want)?;
6926 if let (Some(ck), Some(st)) = (ckpt.as_deref_mut(), stash) {
6927 ck.gdn[il] = Some(st);
6928 }
6929 out
6930 } else {
6931 let mut out = vbuf(e, t * n_embd)?; // every col written by copy_into
6932 let mut col_states: Option<Vec<(CudaSlice<f32>, CudaSlice<f32>)>> =
6933 if ckpt.is_some() && t >= 2 {
6934 Some(Vec::with_capacity(t - 1))
6935 } else {
6936 None
6937 };
6938 for col in 0..t {
6939 let mut h_col = vbuf(e, n_embd)?; // fully written by copy_view_into
6940 let src = h.slice(col * n_embd..(col + 1) * n_embd);
6941 e.copy_view_into(&mut h_col, 0, &src, n_embd)?;
6942 let m_col = self.linear_attn_decode(e, la, &h_col, cache, il)?;
6943 e.copy_into(&mut out, col * n_embd, &m_col, n_embd)?;
6944 // REPLAY-FREE ckpt: clone the chain's ACTUAL state after this column
6945 // (pure dtod — cannot change any computed value). Last column skipped:
6946 // rebuild targets are j <= t-1 columns.
6947 if let Some(cs) = col_states.as_mut() {
6948 if col + 1 < t {
6949 let rl = cache.recur[il].as_ref().unwrap();
6950 cs.push((
6951 e.clone_dtod(&rl.conv_state)?,
6952 e.clone_dtod(&rl.ssm_state)?,
6953 ));
6954 }
6955 }
6956 }
6957 if let (Some(ck), Some(cs)) = (ckpt.as_deref_mut(), col_states) {
6958 // ReplaySSM-assessment instrumentation (2026-07-30): the
6959 // per-column clones are the only true state snapshots left in
6960 // the verify (the batched path stashes INPUTS and replays).
6961 if std::env::var("MEMRA_SPEC_STATS").as_deref() == Ok("1") {
6962 static ONCE: std::sync::Once = std::sync::Once::new();
6963 let bytes: usize =
6964 cs.iter().map(|(c, s)| (c.len() + s.len()) * 4).sum();
6965 ONCE.call_once(|| eprintln!(
6966 "[verify-ckpt] per-column layer il={il}: {} clones, {:.2} MB/layer/round",
6967 cs.len(), bytes as f64 / 1e6));
6968 }
6969 ck.cols[il] = Some(cs);
6970 }
6971 out
6972 }
6973 }
6974 };
6975
6976 // DISPATCH-MIRRORED post-attn norm: eager residual_norm_ffn fuses add+norm+quant
6977 // (1024-thread add_rms_norm_q8_1) only for Dense FFNs whose gate+up are q8_1-fast;
6978 // otherwise (and for MoE) it runs the 256-thread fused add_rms_norm. Mirror per layer.
6979 let ffn_fuse = match &layer.ffn {
6980 crate::hybrid::Ffn::Dense {
6981 ffn_gate, ffn_up, ..
6982 } => {
6983 std::env::var("MEMRA_NO_FUSE_NORMQ").is_err()
6984 && e.uses_q8_1_fast(ffn_gate)
6985 && e.uses_q8_1_fast(ffn_up)
6986 }
6987 crate::hybrid::Ffn::Moe(_) => false,
6988 };
6989 // BATCHED EPILOGUE RE-FUSE (lane/vt-fixes fix 2): on the ffn_fuse path (Dense,
6990 // gate+up q8_1-fast, non-M3) the FFN input is emitted DIRECTLY as q8_1 by ONE
6991 // add_rms_norm_q8_1 launch at nrows=t (row-indexed kernel: the T-row launch is the
6992 // per-row m=1 program; kernel-check pins bit-identity vs the unfused
6993 // add_f32 -> rms_norm_decode -> quantize_q8_1 chain at T=2/4/5/8) — replacing the
6994 // add + rms_norm_decode launches AND the dual/singles' internal re-quantize.
6995 // M3's swigluoai must keep the f32 chain (the fused SwiGLU epilogue encodes plain
6996 // SiLU), mirroring residual_norm_ffn's m3 guard on the decode path.
6997 // step35: same guard per LAYER. A dense FFN's clamp is the SHEXP array (upstream's
6998 // one build_ffn serves dense + shared expert, llama-graph.cpp:1751), and verify MUST
6999 // mirror decode's dispatch or spec self-consistency fails.
7000 let dense_lim = self.cfg.clamp_shexp_at(il as u32);
7001 let fuse_q8 = ffn_fuse && self.cfg.m3.is_none() && dense_lim.is_none();
7002 let mut x1 = vbuf(e, t * n_embd)?; // fully written by add / add_rms_norm*
7003 let mut z = e.zeros(0)?; // replaced below on the unfused arms
7004 let z_q8 = if fuse_q8 {
7005 Some(e.add_rms_norm_q8_1(
7006 &x,
7007 &mixed,
7008 layer.post_attn_norm.float_data(),
7009 &mut x1,
7010 n_embd,
7011 t,
7012 eps,
7013 )?)
7014 } else {
7015 let mut zf = vbuf(e, t * n_embd)?; // fully written by rms_norm_decode / add_rms_norm
7016 if ffn_fuse {
7017 e.add(&x, &mixed, &mut x1, t * n_embd)?;
7018 e.rms_norm_decode(
7019 &x1,
7020 layer.post_attn_norm.float_data(),
7021 &mut zf,
7022 n_embd,
7023 t,
7024 eps,
7025 )?;
7026 } else {
7027 e.add_rms_norm(
7028 &x,
7029 &mixed,
7030 layer.post_attn_norm.float_data(),
7031 &mut x1,
7032 &mut zf,
7033 n_embd,
7034 t,
7035 eps,
7036 )?;
7037 }
7038 z = zf;
7039 None
7040 };
7041 // DECODE-EXACT FFN projections: force MMVQ for gate/up/down at any T to match the
7042 // T=1 decode FP accumulation order. At T>=5 the generic matmul/matmul_pre falls to dp4a
7043 // (128-thread, different FP sum order). At T=2-4 the batched MMVQ is already bit-identical.
7044 let ffn_out = match &layer.ffn {
7045 crate::hybrid::Ffn::Dense {
7046 ffn_gate,
7047 ffn_up,
7048 ffn_down,
7049 } => {
7050 let n_ff = ffn_gate.out_features();
7051 if let Some((zq, zd)) = z_q8.as_ref() {
7052 // FUSED CHAIN (fix 2): pre-quantized z feeds the projections; the SwiGLU
7053 // epilogue emits act pre-quantized for ffn_down (silu_mul_scaled_q8_1,
7054 // bit-identical to silu_mul + quantize — kernel-check-pinned) with the
7055 // NVFP4 macro-scales folded (deferred-scale dual: y*s inline == the
7056 // scale_inplace store, value-exact) — the exact m=1 decode epilogue
7057 // structure at nrows=t.
7058 let pair =
7059 match e.matmul_decode_exact_dual_pre(ffn_gate, ffn_up, zq, zd, t)? {
7060 Some(((g, gs), (u, us))) => Some((g, gs, u, us)),
7061 None => None,
7062 };
7063 let (gate, gs, up, us) = match pair {
7064 Some(x4) => x4,
7065 None => (
7066 e.matmul_decode_exact_pre(ffn_gate, zq, zd, t)?,
7067 1.0, // scale already applied inside _pre
7068 e.matmul_decode_exact_pre(ffn_up, zq, zd, t)?,
7069 1.0,
7070 ),
7071 };
7072 if e.uses_q8_1_fast(ffn_down) {
7073 let (aq, ad) = e.silu_mul_scaled_q8_1(&gate, &up, gs, us, t * n_ff)?;
7074 e.matmul_decode_exact_pre(ffn_down, &aq, &ad, t)?
7075 } else {
7076 let mut act = vbuf(e, t * n_ff)?;
7077 e.silu_mul_scaled(&gate, &up, gs, us, &mut act, t * n_ff)?;
7078 e.matmul_decode_exact(ffn_down, &act, t)?
7079 }
7080 } else {
7081 // UNFUSED (pre-fix) chain — MoE-adjacent/M3/off-fast layers, unchanged.
7082 // DUAL gate+up batched twin (lane/verify-economics, 2026-08-02): one launch
7083 // for the pair at t=2..8 — bit-identical per (tensor,token,row) to the two
7084 // singles (kernel-check pins bitwise; MEMRA_SPEC_DUAL_T=0 reverts). None
7085 // (non-NVFP4 / t outside the tier / seam off) -> the two singles, unchanged.
7086 let (gate, up) =
7087 match e.matmul_decode_exact_dual(ffn_gate, ffn_up, &z, t)? {
7088 Some(pair) => pair,
7089 None => (
7090 e.matmul_decode_exact(ffn_gate, &z, t)?,
7091 e.matmul_decode_exact(ffn_up, &z, t)?,
7092 ),
7093 };
7094 let mut act = vbuf(e, t * n_ff)?; // fully written by ffn_act_lim
7095 Self::ffn_act_lim(
7096 e,
7097 &self.cfg,
7098 &gate,
7099 &up,
7100 1.0,
7101 1.0,
7102 dense_lim,
7103 &mut act,
7104 t * n_ff,
7105 )?;
7106 e.matmul_decode_exact(ffn_down, &act, t)?
7107 }
7108 }
7109 crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il(e, m, &z, t, il as u16)?,
7110 };
7111 // CROSS-LAYER fusion: defer this layer's post-FFN residual add — the next layer's
7112 // fused-q8 attn norm folds it in (add_rms_norm_q8_1 == add; rms_norm; quantize,
7113 // kernel-check-pinned at nrows=T). Non-fused next layers add explicitly above.
7114 pending = Some((x1, ffn_out));
7115 }
7116 // RANGE's final add (no next norm INSIDE the range to fuse with; for the
7117 // whole-trunk call that is the last layer, whose next norm is output_norm — f32-out).
7118 if let Some((x1p, f1p)) = pending.take() {
7119 let mut x2 = vbuf(e, t * n_embd)?; // fully written by add
7120 e.add(&x1p, &f1p, &mut x2, t * n_embd)?;
7121 x = x2;
7122 }
7123 Ok(x)
7124 }
7125 /// BATCHED linear-attn verify (T=K+1): the whole layer in ~10 launches instead of T x the
7126 /// T=1 decode chain (T x ~12 launches + T weight reads of the four projections). The GDN
7127 /// recurrence itself is inherently sequential — gdn_scan_s128 runs its internal t-loop with
7128 /// the SAME per-token math as chained T=1 calls (bit-identical state evolution); everything
7129 /// around it (projections, conv, prep, gated norm, out-proj) batches. Advances conv ring +
7130 /// ssm state exactly like T sequential decode steps.
7131 /// `want_stash`: additionally RETAIN the gdn-scan inputs (pure buffer keep-alives, zero extra
7132 /// kernels) so a partial accept can rebuild the state after any column prefix (REPLAY-FREE).
7133 #[allow(clippy::too_many_arguments)]
7134 fn linear_attn_verify_t(
7135 &self,
7136 e: &Engine,
7137 la: &LinearAttnLayer,
7138 h: &CudaSlice<f32>,
7139 h_q8: Option<(&CudaSlice<i8>, &CudaSlice<f32>)>,
7140 t: usize,
7141 cache: &mut Cache,
7142 il: usize,
7143 want_stash: bool,
7144 ) -> Result<(CudaSlice<f32>, Option<GdnStash>), Box<dyn std::error::Error>> {
7145 let cfg = &self.cfg;
7146 let geometry = la.geometry;
7147 let d_state = geometry.key_head_dim as usize;
7148 let num_k = geometry.key_heads as usize;
7149 let num_v = geometry.value_heads as usize;
7150 let d_conv = geometry.conv_kernel as usize;
7151 let key_dim = d_state * num_k;
7152 let conv_dim = key_dim * 2 + geometry.value_head_dim as usize * num_v;
7153 let eps = cfg.rms_eps;
7154 let scale = 1.0 / (d_state as f32).sqrt();
7155
7156 // DECODE-EXACT projections: matmul_decode_exact forces the MMVQ (warp-per-row, 32-thread)
7157 // accumulation order for EVERY m, matching the T=1 decode path bit-for-bit. The generic
7158 // `matmul` at m>=5 falls to dp4a (128-thread, two-level reduce) which has a different FP
7159 // sum order — ULP differences propagate through gdn_scan and flip argmax on the 27B.
7160 // Q8 TRUNK-FUSION at T=1 (35B: wqkv+wqkv_gate both Q8_0): one fused2 launch, bit-identical
7161 // per (tensor,row) to the two m=1 MMVQ dispatches below — decode-exact contract holds.
7162 // VERIFY-TIER TRUNK FUSION (MEMRA_SPEC_FUSED_T, t=2-4): quantize h ONCE for every
7163 // fused-eligible same-input Q8_0 pair of this layer (35B wqkv+wqkv_gate; 9B
7164 // ssm_beta+ssm_alpha) — each fused2 batched launch then replaces two decode-exact
7165 // calls (each of which re-quantizes the same h + runs its own _b2/_b4 launch).
7166 // Bit-identical per (tensor,token,row) — see spec_fused_t().
7167 // BATCHED EPILOGUE RE-FUSE (lane/vt-fixes fix 2): `h_q8` = the attn-input norm emitted
7168 // directly as q8_1 by the caller's fused rms_norm_q8_1 (bit-identical to the unfused
7169 // chain, kernel-check-pinned). When present it REPLACES the standalone quantize below
7170 // and feeds every projection; the caller guaranteed all four input projections are
7171 // q8_1-fast. When absent, the old shared-quantize (fused-t window) stands.
7172 let h_q8_t = if h_q8.is_none()
7173 && spec_fused_t()
7174 && (2..=4).contains(&t)
7175 && ((e.uses_q8_1_fast(&la.wqkv) && e.uses_q8_1_fast(&la.wqkv_gate))
7176 || (e.uses_q8_1_fast(&la.ssm_beta) && e.uses_q8_1_fast(&la.ssm_alpha)))
7177 {
7178 Some(e.quantize_q8_1(h, t, cfg.n_embd as usize)?)
7179 } else {
7180 None
7181 };
7182 // one view: the caller's fused-norm q8 or this fn's own shared quantize.
7183 let hq8_any: Option<(&CudaSlice<i8>, &CudaSlice<f32>)> =
7184 h_q8.or(h_q8_t.as_ref().map(|(q, d)| (q, d)));
7185 let (qkv_mixed, z) = {
7186 let mut fused = None;
7187 if t == 1 && e.uses_q8_1_fast(&la.wqkv) && e.uses_q8_1_fast(&la.wqkv_gate) {
7188 let (hq, hd) = e.quantize_q8_1(h, 1, cfg.n_embd as usize)?;
7189 fused = e.matmul_q8_fused2(&la.wqkv, &la.wqkv_gate, &hq, &hd)?;
7190 } else if let Some((hq, hd)) = hq8_any {
7191 if spec_fused_t() && (2..=4).contains(&t) {
7192 fused = e.matmul_q8_fused2_t(&la.wqkv, &la.wqkv_gate, hq, hd, t)?;
7193 }
7194 }
7195 match (fused, hq8_any) {
7196 (Some(pair), _) => pair,
7197 (None, Some((hq, hd))) if h_q8.is_some() => (
7198 e.matmul_decode_exact_pre(&la.wqkv, hq, hd, t)?,
7199 e.matmul_decode_exact_pre(&la.wqkv_gate, hq, hd, t)?,
7200 ),
7201 (None, _) => (
7202 e.matmul_decode_exact(&la.wqkv, h, t)?,
7203 e.matmul_decode_exact(&la.wqkv_gate, h, t)?,
7204 ),
7205 }
7206 };
7207 // beta+alpha DUAL at T=1 (75% of p3 rounds run T=1 verify — p-min chain cuts): the dual
7208 // mr2 kernel is bit-identical per element to the m=1 MMVQ matmul_decode_exact dispatches
7209 // (same warp-per-row body, blockIdx.y picks the weight), so the decode-exact contract
7210 // holds; the run-spec battery is the arbiter. T>1 keeps the per-tensor decode-exact path.
7211 let (beta_raw, alpha) = if t == 1 {
7212 let (hq, hd) = e.quantize_q8_1(h, 1, cfg.n_embd as usize)?;
7213 match e.matmul_pre_dual_noscale(&la.ssm_beta, &la.ssm_alpha, &hq, &hd, 1)? {
7214 Some(((mut b, bs), (mut a, as_))) => {
7215 if bs != 1.0 {
7216 e.scale_inplace(&mut b, bs, la.ssm_beta.out_features())?;
7217 }
7218 if as_ != 1.0 {
7219 e.scale_inplace(&mut a, as_, la.ssm_alpha.out_features())?;
7220 }
7221 (b, a)
7222 }
7223 // Q8_0 fused2 twin (9B stores beta/alpha as Q8_0): DISPATCH-MIRRORS the eager
7224 // decode's beta_alpha closure — the fused body is qmatvec_q8_0_mmvq verbatim,
7225 // bit-identical per row (kernel-check rel=0.00e0 gate), so decode==verify holds.
7226 None => match e.matmul_q8_fused2(&la.ssm_beta, &la.ssm_alpha, &hq, &hd)? {
7227 Some((b, a)) => (b, a),
7228 None => (
7229 e.matmul_decode_exact(&la.ssm_beta, h, 1)?,
7230 e.matmul_decode_exact(&la.ssm_alpha, h, 1)?,
7231 ),
7232 },
7233 }
7234 } else {
7235 // fused-t twin (9B stores beta/alpha as Q8_0): same shared-quantize + one launch
7236 // contract as the wqkv pair above; 35B beta/alpha are Float -> None -> fallback.
7237 let mut nvfp4_fused = None;
7238 let mut q8_fused = None;
7239 if let Some((hq, hd)) = hq8_any {
7240 if t == 3 && std::env::var("MEMRA_NVFP4_AUX_DUAL").as_deref() != Ok("0") {
7241 nvfp4_fused =
7242 e.matmul_decode_exact_dual_pre(&la.ssm_beta, &la.ssm_alpha, hq, hd, t)?;
7243 if nvfp4_fused.is_some() && std::env::var("MEMRA_DEBUG").is_ok() {
7244 static ONCE: std::sync::Once = std::sync::Once::new();
7245 ONCE.call_once(|| {
7246 eprintln!("[memra] NVFP4 beta+alpha batched aux dual ENGAGED (t={t})")
7247 });
7248 }
7249 }
7250 if nvfp4_fused.is_none() && spec_fused_t() && (2..=4).contains(&t) {
7251 q8_fused = e.matmul_q8_fused2_t(&la.ssm_beta, &la.ssm_alpha, hq, hd, t)?;
7252 }
7253 }
7254 if let Some(((mut b, bs), (mut a, as_))) = nvfp4_fused {
7255 if bs != 1.0 {
7256 e.scale_inplace(&mut b, bs, t * la.ssm_beta.out_features())?;
7257 }
7258 if as_ != 1.0 {
7259 e.scale_inplace(&mut a, as_, t * la.ssm_alpha.out_features())?;
7260 }
7261 (b, a)
7262 } else if let Some(pair) = q8_fused {
7263 pair
7264 } else {
7265 match hq8_any {
7266 Some((hq, hd)) if h_q8.is_some() => (
7267 e.matmul_decode_exact_pre(&la.ssm_beta, hq, hd, t)?,
7268 e.matmul_decode_exact_pre(&la.ssm_alpha, hq, hd, t)?,
7269 ),
7270 _ => (
7271 e.matmul_decode_exact(&la.ssm_beta, h, t)?,
7272 e.matmul_decode_exact(&la.ssm_alpha, h, t)?,
7273 ),
7274 }
7275 }
7276 };
7277
7278 // conv with CARRIED state + ring roll (T >= pad rides the input-column update kernel;
7279 // T < pad — the MEMRA_SPEC_M2 t=2 arm — rolls via the pure-copy ring rebuild).
7280 let rl = cache.recur[il].as_mut().unwrap();
7281 let mut conv_out = e.uninit(conv_dim * t)?;
7282 e.ssm_conv1d_tm_state(
7283 &qkv_mixed,
7284 &mut rl.conv_state,
7285 la.ssm_conv1d.float_data(),
7286 &mut conv_out,
7287 conv_dim,
7288 t,
7289 d_conv,
7290 )?;
7291
7292 // GDN prep via the prefill kernels (repack + L2 + sigmoid + glog), T-wide.
7293 let mut q_g = e.uninit(d_state * num_v * t)?;
7294 let mut k_g = e.uninit(d_state * num_v * t)?;
7295 let mut v_g = e.uninit(d_state * num_v * t)?;
7296 e.qkv_to_gdn_repack(
7297 &conv_out, &mut q_g, &mut k_g, &mut v_g, d_state, num_v, num_k, key_dim, t,
7298 )?;
7299 let mut q_l2 = e.uninit(d_state * num_v * t)?;
7300 e.l2_norm_decode(&q_g, &mut q_l2, d_state, num_v * t, eps)?;
7301 let mut k_l2 = e.uninit(d_state * num_v * t)?;
7302 e.l2_norm_decode(&k_g, &mut k_l2, d_state, num_v * t, eps)?;
7303 let mut beta = e.uninit(t * num_v)?;
7304 e.sigmoid(&beta_raw, &mut beta, t * num_v)?;
7305 let mut g_log = e.uninit(t * num_v)?;
7306 e.gdn_glog(
7307 &alpha,
7308 la.ssm_dt.float_data(),
7309 la.ssm_a.float_data(),
7310 &mut g_log,
7311 num_v,
7312 t,
7313 )?;
7314
7315 // ONE gdn_scan over T tokens from the carried state (internal sequential loop ==
7316 // T chained T=1 steps). Ping-pong the resident buffers like eager decode.
7317 let mut o = e.uninit(d_state * num_v * t)?;
7318 {
7319 let crate::cache::RecurLayer {
7320 ssm_state,
7321 ssm_state_alt,
7322 ..
7323 } = rl;
7324 e.gdn_scan_s128(
7325 &q_l2,
7326 &k_l2,
7327 &v_g,
7328 &g_log,
7329 &beta,
7330 ssm_state,
7331 ssm_state_alt,
7332 &mut o,
7333 num_v,
7334 t,
7335 scale,
7336 )?;
7337 }
7338 std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
7339
7340 // gated RMSNorm + out projection, T-wide. FUSED-QUANTIZE ARM (lane/vt-fixes fix 2,
7341 // mirroring the T=1 decode's launch-arc form): when ssm_out rides the q8_1 fast path,
7342 // emit q8_1 straight from the gated norm at nrows=num_v*t (row-indexed kernel, the
7343 // T-wide launch is the per-row program; kernel-check pins bit-identity vs
7344 // gated_rmsnorm -> quantize_q8_1 at T=1 and T=5) and feed the decode-exact dispatch
7345 // pre-quantized — one launch replaces norm + quantize. Fallback = the f32 chain.
7346 let out = if e.uses_q8_1_fast(&la.ssm_out) {
7347 let (gq, gd) =
7348 e.gated_rmsnorm_q8_1(&o, la.ssm_norm.float_data(), &z, d_state, num_v * t, eps)?;
7349 e.matmul_decode_exact_pre(&la.ssm_out, &gq, &gd, t)?
7350 } else {
7351 let mut gn = e.uninit(d_state * num_v * t)?;
7352 e.gated_rmsnorm(
7353 &o,
7354 la.ssm_norm.float_data(),
7355 &z,
7356 &mut gn,
7357 d_state,
7358 num_v * t,
7359 eps,
7360 )?;
7361 // DECODE-EXACT out-projection: same MMVQ path as the T=1 decode (ssm_out at m>=5
7362 // would fall to dp4a with a different FP reduction order — same class of bug as
7363 // the input projs).
7364 e.matmul_decode_exact(&la.ssm_out, &gn, t)?
7365 };
7366 let stash = if want_stash {
7367 Some(GdnStash {
7368 qkv_mixed,
7369 q_l2,
7370 k_l2,
7371 v_g,
7372 g_log,
7373 beta,
7374 })
7375 } else {
7376 None
7377 };
7378 Ok((out, stash))
7379 }
7380
7381 /// REPLAY-FREE partial-accept commit (2026-07-03): make the cache state == "committed through
7382 /// the first `j` verify columns" WITHOUT the legacy rollback + duplicate trunk replay.
7383 /// - Full-attn KV: truncate both the owning-stage shadow and every TP rank to snapshot + j.
7384 /// The verify's appended rows for those columns are bit-identical to what an eager T=1
7385 /// chain writes (the decode-exact contract the verify-probe gates), so keeping them ==
7386 /// replaying them.
7387 /// - Linear layers, batched path: rebuild the conv ring by PURE COPIES (ring holds raw input
7388 /// columns) and the ssm state by a prefix re-run of the SAME gdn_scan kernel (t=j) from the
7389 /// snapshot state over the stash's identical inputs — the kernel's t-loop carries state in
7390 /// registers and writes it once at the end, so iterations 0..j-1 are independent of T:
7391 /// bit-identical to the verify's own state after j tokens == the eager chain state.
7392 /// - Linear layers, per-column path: restore the cloned actual state after column j-1.
7393 /// Caller guarantees 1 <= j <= t-1 (j==0 rounds take the legacy rollback; j==t is full accept).
7394 fn commit_verified_prefix(
7395 &self,
7396 e: &Engine,
7397 cache: &mut Cache,
7398 snap: &crate::cache::CacheSnapshot,
7399 ckpt: &VerifyCkpt,
7400 j: usize,
7401 kv_lens_done: bool,
7402 dev_j: Option<(&CudaSlice<u32>, usize, usize)>,
7403 ) -> Result<(), Box<dyn std::error::Error>> {
7404 // GDN geometry derives lazily inside recurrent-layer arms. Full-attention plans carry no
7405 // recurrent state and must never be forced through a synthetic SSM geometry.
7406 // Engine-bundle slice 1 (DSF-ROUNDCOST-20260820 §1.1): the per-column-arm restores
7407 // are 2 tiny D2D copies per linear layer (~96 dispatches/partial round on the q38
7408 // route). When every cols-arm layer shares uniform state sizes (single ssm cfg —
7409 // always true today), batch them into two `copy_batch_uniform_f32` launches. Bytes,
7410 // buffers and stream order are identical to the per-layer memcpy sequence; the
7411 // kernel-rebuild (gdn-stash) arm below is untouched. MEMRA_STATE_COPY_BATCH=0 reverts.
7412 let mut batched_cols = false;
7413 if state_copy_batch_on() && dev_j.is_none() {
7414 use cudarc::driver::DevicePtr;
7415 let s = &e.gpu.stream();
7416 let mut conv_pairs: Vec<(u64, u64)> = Vec::new();
7417 let mut ssm_pairs: Vec<(u64, u64)> = Vec::new();
7418 let (mut conv_words, mut ssm_words) = (0usize, 0usize);
7419 let mut uniform = true;
7420 for il in 0..self.layers.len() {
7421 let Some(rl) = cache.recur[il].as_ref() else {
7422 continue;
7423 };
7424 if ckpt.gdn[il].is_some() {
7425 continue; // kernel-rebuild arm restores below, per layer
7426 }
7427 let Some(cols) = &ckpt.cols[il] else {
7428 continue; // missing-ckpt error surfaces in the main loop
7429 };
7430 let (c, st) = &cols[j - 1];
7431 if conv_pairs.is_empty() {
7432 conv_words = c.len();
7433 ssm_words = st.len();
7434 } else if c.len() != conv_words || st.len() != ssm_words {
7435 uniform = false;
7436 break;
7437 }
7438 let (pc, _g0) = c.device_ptr(s);
7439 let (dc, _g1) = rl.conv_state.device_ptr(s);
7440 let (ps, _g2) = st.device_ptr(s);
7441 let (ds, _g3) = rl.ssm_state.device_ptr(s);
7442 conv_pairs.push((pc as u64, dc as u64));
7443 ssm_pairs.push((ps as u64, ds as u64));
7444 }
7445 if uniform && !conv_pairs.is_empty() {
7446 let n = conv_pairs.len();
7447 let mut t = vec![0u64; 2 * n];
7448 for (k, &(src, dst)) in conv_pairs.iter().enumerate() {
7449 t[k] = src;
7450 t[n + k] = dst;
7451 }
7452 let conv_t = e.htod_u64(&t)?;
7453 for (k, &(src, dst)) in ssm_pairs.iter().enumerate() {
7454 t[k] = src;
7455 t[n + k] = dst;
7456 }
7457 let ssm_t = e.htod_u64(&t)?;
7458 e.copy_batch_uniform_f32(&conv_t, n, conv_words)?;
7459 e.copy_batch_uniform_f32(&ssm_t, n, ssm_words)?;
7460 batched_cols = true;
7461 }
7462 }
7463 rewind_tp_kv_verified_prefix(&mut cache.tp_kv, &snap.tp_kv_len, j)?;
7464 for il in 0..self.layers.len() {
7465 if let (Some(kvl), Some(saved)) = (cache.kv[il].as_mut(), snap.kv_len[il]) {
7466 kvl.len = saved + j;
7467 // devacc 3a: spec_rollback_kv already wrote len_d on-device (same value).
7468 if !kv_lens_done {
7469 e.set_i32_one(&mut kvl.len_d, kvl.len as i32)?;
7470 }
7471 }
7472 if let Some(rl) = cache.recur[il].as_mut() {
7473 let Mixer::Linear(linear) = &self.layers[il].mixer else {
7474 return Err(format!("recurrent cache layer {il} has no GDN plan").into());
7475 };
7476 let geometry = linear.geometry;
7477 let d_state = geometry.key_head_dim as usize;
7478 let num_k = geometry.key_heads as usize;
7479 let num_v = geometry.value_heads as usize;
7480 let d_conv = geometry.conv_kernel as usize;
7481 let conv_dim = d_state * num_k * 2 + geometry.value_head_dim as usize * num_v;
7482 let scale = 1.0 / (d_state as f32).sqrt();
7483 if let Some(st) = &ckpt.gdn[il] {
7484 let ring_old = snap.conv[il].as_ref().expect("snapshot missing conv");
7485 let state_in = snap.ssm[il].as_ref().expect("snapshot missing ssm");
7486 if let Some((acc, base, t_v)) = dev_j {
7487 // 3b: j read on-device (_dc twins, same bodies; full accept early-exits).
7488 e.ssm_conv_ring_rebuild_dc(
7489 &st.qkv_mixed,
7490 ring_old,
7491 &mut rl.conv_state,
7492 conv_dim,
7493 acc,
7494 base,
7495 t_v,
7496 d_conv,
7497 )?;
7498 let mut o = e.uninit(d_state * num_v * j.max(1))?;
7499 e.gdn_scan_s128_dc(
7500 &st.q_l2,
7501 &st.k_l2,
7502 &st.v_g,
7503 &st.g_log,
7504 &st.beta,
7505 state_in,
7506 &mut rl.ssm_state,
7507 &mut o,
7508 num_v,
7509 acc,
7510 base,
7511 t_v,
7512 scale,
7513 )?;
7514 } else {
7515 e.ssm_conv_ring_rebuild(
7516 &st.qkv_mixed,
7517 ring_old,
7518 &mut rl.conv_state,
7519 conv_dim,
7520 j,
7521 d_conv,
7522 )?;
7523 let mut o = e.uninit(d_state * num_v * j)?; // scan output, discarded
7524 e.gdn_scan_s128(
7525 &st.q_l2,
7526 &st.k_l2,
7527 &st.v_g,
7528 &st.g_log,
7529 &st.beta,
7530 state_in,
7531 &mut rl.ssm_state,
7532 &mut o,
7533 num_v,
7534 j,
7535 scale,
7536 )?;
7537 }
7538 } else if let Some(cols) = &ckpt.cols[il] {
7539 if !batched_cols {
7540 let (c, s) = &cols[j - 1];
7541 e.copy_into(&mut rl.conv_state, 0, c, c.len())?;
7542 e.copy_into(&mut rl.ssm_state, 0, s, s.len())?;
7543 }
7544 } else {
7545 return Err(
7546 "commit_verified_prefix: verify ckpt missing for linear layer".into(),
7547 );
7548 }
7549 }
7550 }
7551 cache.pos = snap.pos + j;
7552 Ok(())
7553 }
7554
7555 /// ROUND-STREAM: recur restore with device-j (the _dc twins; full accept early-exits
7556 /// in-kernel). Requires the batched-linear stash on every linear layer (stream gate).
7557 fn commit_verified_prefix_stream(
7558 &self,
7559 e: &Engine,
7560 cache: &mut Cache,
7561 snap: &crate::cache::CacheSnapshot,
7562 ckpt: &VerifyCkpt,
7563 acc: &CudaSlice<u32>,
7564 base: usize,
7565 t_v: usize,
7566 ) -> Result<(), Box<dyn std::error::Error>> {
7567 for il in 0..self.layers.len() {
7568 if let Some(rl) = cache.recur[il].as_mut() {
7569 let Mixer::Linear(linear) = &self.layers[il].mixer else {
7570 return Err(format!("recurrent cache layer {il} has no GDN plan").into());
7571 };
7572 let geometry = linear.geometry;
7573 let d_state = geometry.key_head_dim as usize;
7574 let num_k = geometry.key_heads as usize;
7575 let num_v = geometry.value_heads as usize;
7576 let d_conv = geometry.conv_kernel as usize;
7577 let conv_dim = d_state * num_k * 2 + geometry.value_head_dim as usize * num_v;
7578 let scale = 1.0 / (d_state as f32).sqrt();
7579 let st = ckpt.gdn[il]
7580 .as_ref()
7581 .ok_or("stream restore: batched-linear stash missing")?;
7582 let ring_old = snap.conv[il].as_ref().expect("snapshot missing conv");
7583 let state_in = snap.ssm[il].as_ref().expect("snapshot missing ssm");
7584 e.ssm_conv_ring_rebuild_dc(
7585 &st.qkv_mixed,
7586 ring_old,
7587 &mut rl.conv_state,
7588 conv_dim,
7589 acc,
7590 base,
7591 t_v,
7592 d_conv,
7593 )?;
7594 let mut o = e.uninit(d_state * num_v * t_v)?;
7595 e.gdn_scan_s128_dc(
7596 &st.q_l2,
7597 &st.k_l2,
7598 &st.v_g,
7599 &st.g_log,
7600 &st.beta,
7601 state_in,
7602 &mut rl.ssm_state,
7603 &mut o,
7604 num_v,
7605 acc,
7606 base,
7607 t_v,
7608 scale,
7609 )?;
7610 }
7611 }
7612 Ok(())
7613 }
7614
7615 /// EAGLE3 aux-capturing verify forward over `tokens` (T) — mirrors `decode_step_t_h` exactly
7616 /// (same KV append, same causal verify, same recur advance) but ALSO clones the aux residual-
7617 /// stream hiddens (blocks in `aux_layers`) for TWO columns: the LAST column (always) and the
7618 /// optional `pred_col` (the EAGLE seed = bonus's predecessor). Returns
7619 /// (all_T_logits host, last_col_aux, pred_col_aux?). Used by the EAGLE3 orchestrator's commit.
7620 pub fn decode_step_t_aux2(
7621 &self,
7622 e: &Engine,
7623 tokens: &[u32],
7624 pos0: usize,
7625 cache: &mut Cache,
7626 aux_layers: &[usize],
7627 pred_col: Option<usize>,
7628 ) -> Result<
7629 (Vec<f32>, Vec<CudaSlice<f32>>, Option<Vec<CudaSlice<f32>>>),
7630 Box<dyn std::error::Error>,
7631 > {
7632 let cfg = &self.cfg;
7633 let n_embd = cfg.n_embd as usize;
7634 let eps = cfg.rms_eps;
7635 let t = tokens.len();
7636 let pos_vec: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
7637 let pos_d = e.htod_i32(&pos_vec)?;
7638 let mut x = e.htod(&self.embd.gather(n_embd, tokens))?;
7639 let mut aux_last: Vec<CudaSlice<f32>> = Vec::with_capacity(aux_layers.len());
7640 let mut aux_pred: Vec<CudaSlice<f32>> = Vec::new();
7641 let want_pred = pred_col.is_some();
7642
7643 for (il, layer) in self.layers.iter().enumerate() {
7644 // DISPATCH-MIRRORED norms (FP-order lesson #8) — see decode_step_t_h_emb.
7645 let mixer_fast = self.mixer_in_q8_1_fast(e, &layer.mixer);
7646 let norm_fused = std::env::var("MEMRA_NO_FUSE_NORMQ").is_err() && mixer_fast;
7647 let mut h = vbuf(e, t * n_embd)?; // fully written by either rms_norm arm
7648 if norm_fused {
7649 e.rms_norm_decode(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
7650 } else {
7651 e.rms_norm(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
7652 }
7653 let mixed = match &layer.mixer {
7654 Mixer::Full(fa) => {
7655 self.full_attn_verify(e, fa, &h, None, &pos_d, t, cache, il, None)?
7656 }
7657 Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
7658 Mixer::Linear(la) => {
7659 let mut out = e.zeros(t * n_embd)?;
7660 for col in 0..t {
7661 let mut h_col = e.zeros(n_embd)?;
7662 let src = h.slice(col * n_embd..(col + 1) * n_embd);
7663 e.copy_view_into(&mut h_col, 0, &src, n_embd)?;
7664 let m_col = self.linear_attn_decode(e, la, &h_col, cache, il)?;
7665 e.copy_into(&mut out, col * n_embd, &m_col, n_embd)?;
7666 }
7667 out
7668 }
7669 };
7670 let ffn_fuse = match &layer.ffn {
7671 crate::hybrid::Ffn::Dense {
7672 ffn_gate, ffn_up, ..
7673 } => {
7674 std::env::var("MEMRA_NO_FUSE_NORMQ").is_err()
7675 && e.uses_q8_1_fast(ffn_gate)
7676 && e.uses_q8_1_fast(ffn_up)
7677 }
7678 crate::hybrid::Ffn::Moe(_) => false,
7679 };
7680 let mut x1 = vbuf(e, t * n_embd)?; // fully written by add / add_rms_norm
7681 let mut z = vbuf(e, t * n_embd)?; // fully written by rms_norm_decode / add_rms_norm
7682 if ffn_fuse {
7683 e.add(&x, &mixed, &mut x1, t * n_embd)?;
7684 e.rms_norm_decode(
7685 &x1,
7686 layer.post_attn_norm.float_data(),
7687 &mut z,
7688 n_embd,
7689 t,
7690 eps,
7691 )?;
7692 } else {
7693 e.add_rms_norm(
7694 &x,
7695 &mixed,
7696 layer.post_attn_norm.float_data(),
7697 &mut x1,
7698 &mut z,
7699 n_embd,
7700 t,
7701 eps,
7702 )?;
7703 }
7704 let ffn_out = match &layer.ffn {
7705 crate::hybrid::Ffn::Dense {
7706 ffn_gate,
7707 ffn_up,
7708 ffn_down,
7709 } => {
7710 let n_ff = ffn_gate.out_features();
7711 let gate = e.matmul_decode_exact(ffn_gate, &z, t)?;
7712 let up = e.matmul_decode_exact(ffn_up, &z, t)?;
7713 let mut act = vbuf(e, t * n_ff)?; // fully written by ffn_act_lim
7714 // dense FFN clamp = the SHEXP array (upstream build_ffn serves both).
7715 Self::ffn_act_lim(
7716 e,
7717 &self.cfg,
7718 &gate,
7719 &up,
7720 1.0,
7721 1.0,
7722 self.cfg.clamp_shexp_at(il as u32),
7723 &mut act,
7724 t * n_ff,
7725 )?;
7726 e.matmul_decode_exact(ffn_down, &act, t)?
7727 }
7728 crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il(e, m, &z, t, il as u16)?,
7729 };
7730 let mut x2 = vbuf(e, t * n_embd)?; // fully written by add
7731 e.add(&x1, &ffn_out, &mut x2, t * n_embd)?;
7732 if aux_layers.contains(&il) {
7733 let mut a = e.zeros(n_embd)?;
7734 e.copy_view_into(&mut a, 0, &x2.slice((t - 1) * n_embd..t * n_embd), n_embd)?;
7735 aux_last.push(a);
7736 if let Some(pc) = pred_col {
7737 let mut ap = e.zeros(n_embd)?;
7738 e.copy_view_into(
7739 &mut ap,
7740 0,
7741 &x2.slice(pc * n_embd..(pc + 1) * n_embd),
7742 n_embd,
7743 )?;
7744 aux_pred.push(ap);
7745 }
7746 }
7747 x = x2;
7748 }
7749 let mut hn = vbuf(e, t * n_embd)?; // fully written by rms_norm_decode
7750 e.rms_norm_decode(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
7751 let logits = e.matmul_decode_exact(&self.output, &hn, t)?;
7752 let host = e.dtoh(&logits)?;
7753 cache.pos += t;
7754 Ok((
7755 host,
7756 aux_last,
7757 if want_pred { Some(aux_pred) } else { None },
7758 ))
7759 }
7760
7761 /// step35 SPEC-VERIFY attention over T query tokens — a per-row REPLAY of the eager
7762 /// `step35_decode_attn`.
7763 ///
7764 /// WHY A REPLAY AND NOT A BATCHED TWIN. The verify's whole job is to be bit-identical to what
7765 /// the eager decode would have computed for the same tokens; that is what makes greedy spec
7766 /// decode exact (run-spec asserts token identity for K=1..8). Every other verify arm in this
7767 /// file earns that identity by carefully mirroring dispatch (`matmul_decode_exact` to force
7768 /// MMVQ at any m, per-layer `ffn_fuse` mirroring, per-row `fa_decode` key bounds). step35
7769 /// stacks FOUR more per-layer degrees of freedom on top of that — per-layer `n_head`
7770 /// (64 full / 96 SWA), per-layer rotary width (64 full / 128 SWA), per-layer rope base, and a
7771 /// SEPARATE `attn_gate` tensor whose projection shares the attn-normed input — and its SWA
7772 /// layers attend through a token-OFFSET view whose offset is a function of the ABSOLUTE
7773 /// position of each query row. A batched twin would have to reproduce all of that AND the
7774 /// per-row offset in one launch; the offset alone rules out the existing rows kernels (they
7775 /// take one `base_len`, not a per-row offset).
7776 ///
7777 /// So this arm calls the eager path itself, once per row, on the same cache. Identity is then
7778 /// true BY CONSTRUCTION rather than by mirroring: row r runs exactly the kernel sequence that
7779 /// eager decode step r runs (same projections, same q8_1 fusion decision, same append, same
7780 /// view arithmetic, same `fa_decode_kvmod`, same gate), because it IS that code. Cost: T x the
7781 /// eager decode mixer instead of one batched pass — the same trade the generic arm's `else`
7782 /// per-row loop already accepts when `fa_rows_eligible` says no. Correctness first; a batched
7783 /// step35 twin is a perf lane's job and must be gated against this arm.
7784 ///
7785 /// The `h_q8` pre-quantized pair from the caller's fused norm is NOT forwarded: it is a
7786 /// T-row buffer and `step35_decode_attn`'s `pre_q` contract is one row. Instead each row's
7787 /// f32 `h` slice is handed over and the callee re-derives its own q8_1 exactly as eager decode
7788 /// does (`quantize_q8_1(h, 1, n_embd)`) — which is the dispatch being mirrored. Callers that
7789 /// took the fused arm therefore MUST still pass a live `h`; `step35_verify` asserts that.
7790 #[allow(clippy::too_many_arguments)]
7791 fn step35_verify(
7792 &self,
7793 e: &Engine,
7794 fa: &FullAttnLayer,
7795 h: &CudaSlice<f32>,
7796 h_q8: Option<(&CudaSlice<i8>, &CudaSlice<f32>)>,
7797 t: usize,
7798 cache: &mut Cache,
7799 il: usize,
7800 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7801 let n_embd = self.cfg.n_embd as usize;
7802 // The fused (h-less) attn-norm arm hands `h` as a zero-length placeholder. This arm needs
7803 // the f32 rows, so the caller must not take that lever for step35 — enforced at the call
7804 // site by the sliding-gated-MoE `Mixer::Full(_) => false` arm of
7805 // `lin_q8_only` in `decode_step_t_core_stream`, and asserted here so a future caller
7806 // cannot regress it into silently reading an empty buffer.
7807 assert_eq!(
7808 h.len(),
7809 t * n_embd,
7810 "step35_verify needs the f32 attn-normed rows ([t*n_embd]); the caller took the \
7811 fused q8-only norm arm (h_q8={}) — step35 must stay on the unfused arm",
7812 h_q8.is_some()
7813 );
7814 // ROW WIDTH IS n_embd, NOT n_head*head_dim: `step35_decode_attn` returns the mixer output
7815 // AFTER `wo`, so a row is [n_embd] — the same contract the generic arm's
7816 // `matmul_decode_exact(&fa.wo, &attn_g, t)` return has. Sizing this buffer from the
7817 // per-layer head geometry (8192 on full-attn, 12288 on SWA) instead overran the row on the
7818 // FIRST copy and panicked inside `copy_into`'s `CudaView::slice` unwrap
7819 // (raw/mtp-bt-20260806T212127Z.log frames 12-13).
7820 let mut out = vbuf(e, t * n_embd)?; // each row fully written by the copy below
7821 for r in 0..t {
7822 // Absolute position of this query row. `cache.pos` is the committed length at round
7823 // start and every row before r has already been appended by this loop, so the r-th
7824 // verify token sits at cache.pos + r — the same position eager decode would give it.
7825 let pos_d = e.htod_i32(&[(cache.pos + r) as i32])?;
7826 let mut h_row = vbuf(e, n_embd)?; // fully written by copy_view_into
7827 e.copy_view_into(
7828 &mut h_row,
7829 0,
7830 &h.slice(r * n_embd..(r + 1) * n_embd),
7831 n_embd,
7832 )?;
7833 // THE eager decode mixer: appends this row's K/V at kvl.len, advances it, then
7834 // attends over the (SWA-offset) view. Post-`wo`, same contract as this fn returns.
7835 let o = self.step35_decode_attn(e, fa, il, &h_row, None, &pos_d, cache)?;
7836 debug_assert_eq!(
7837 o.len(),
7838 n_embd,
7839 "step35_decode_attn returns post-wo [n_embd]"
7840 );
7841 e.copy_into(&mut out, r * n_embd, &o, n_embd)?;
7842 }
7843 Ok(out)
7844 }
7845
7846 /// Full-attention mixer over T query tokens with a GROWING resident KV (verify path, §D.3).
7847 /// Appends the T new K/V columns to cache.kv[il] then attends causally over [0..len) via
7848 /// fa_prefill. Token-major [T, kv_dim] projection layout == cache row layout (single copy).
7849 #[allow(clippy::too_many_arguments)]
7850 fn full_attn_verify(
7851 &self,
7852 e: &Engine,
7853 fa: &FullAttnLayer,
7854 h: &CudaSlice<f32>,
7855 h_q8: Option<(&CudaSlice<i8>, &CudaSlice<f32>)>,
7856 pos_d: &CudaSlice<i32>,
7857 t: usize,
7858 cache: &mut Cache,
7859 il: usize,
7860 stream_ctr: Option<&CudaSlice<i32>>,
7861 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7862 // step35: the generic geometry below is wrong for this arch (per-layer n_head, partial
7863 // per-layer rope, the SWA offset view, and a SEPARATE head-wise gate tensor), so it takes
7864 // its own arm. A verify that silently computes different attention than decode defeats the
7865 // whole self-consistency gate, so the arm is a per-row REPLAY of `step35_decode_attn`
7866 // rather than a batched twin — see `step35_verify` for why that is the exactness-correct
7867 // shape and not laziness.
7868 if self.sliding_gated_moe_batch_program() {
7869 if stream_ctr.is_some() {
7870 return Err(
7871 "step35 has no ROUND-STREAM verify arm (the device-counter _dc twins \
7872 cannot express the SWA offset KV view; same root cause as the dc \
7873 decode refusal) — run spec without the stream arm"
7874 .into(),
7875 );
7876 }
7877 return self.step35_verify(e, fa, h, h_q8, t, cache, il);
7878 }
7879 let cfg = &self.cfg;
7880 let geometry = cfg.full_attention_geometry_at(il as u32);
7881 let n_head = geometry.n_head as usize;
7882 let n_head_kv = geometry.n_head_kv as usize;
7883 let head_dim = geometry.head_dim_k as usize;
7884 let eps = cfg.rms_eps;
7885 let scale = geometry.attention_scale();
7886 let n_embd = cfg.n_embd as usize;
7887
7888 // DECODE-EXACT Q/K/V projections: matmul_decode_exact forces the MMVQ (warp-per-row) path
7889 // for every m, matching the T=1 decode's FP accumulation order. matmul_pre at m>=5 would
7890 // fall to dp4a (128-thread, two-level reduce) with a different FP sum order.
7891 // Q8 TRUNK-FUSION at T=1: DISPATCH-MIRRORS the eager decode's fused3 (bit-identical body).
7892 // BATCHED EPILOGUE RE-FUSE (lane/vt-fixes fix 2): `h_q8` = the attn-input norm's q8_1
7893 // form emitted by the fused rms_norm_q8_1 (bit-identical to rms_norm_decode ->
7894 // quantize_q8_1, kernel-check-pinned). When present (caller checked mixer q8_1-fast),
7895 // every projection consumes it — `h` may be a zero-len placeholder and must not be read.
7896 let (qf, mut k, v) = {
7897 let mut fused = None;
7898 let qkv_fast =
7899 e.uses_q8_1_fast(&fa.wq) && e.uses_q8_1_fast(&fa.wk) && e.uses_q8_1_fast(&fa.wv);
7900 if t == 1 && qkv_fast {
7901 let (hq_o, hd_o);
7902 let (hq, hd): (&CudaSlice<i8>, &CudaSlice<f32>) = match h_q8 {
7903 Some(p) => p,
7904 None => {
7905 (hq_o, hd_o) = e.quantize_q8_1(h, 1, n_embd)?;
7906 (&hq_o, &hd_o)
7907 }
7908 };
7909 fused = e.matmul_q8_fused3(&fa.wq, &fa.wk, &fa.wv, hq, hd)?;
7910 } else if spec_fused_t() && (2..=4).contains(&t) && qkv_fast {
7911 // VERIFY-TIER TRUNK FUSION (MEMRA_SPEC_FUSED_T): one shared quantize + one
7912 // fused3 batched launch replaces three decode-exact calls (3 re-quantizes of
7913 // the same h + 3 _b2/_b4 launches). Bit-identical per (tensor,token,row).
7914 let (hq_o, hd_o);
7915 let (hq, hd): (&CudaSlice<i8>, &CudaSlice<f32>) = match h_q8 {
7916 Some(p) => p,
7917 None => {
7918 (hq_o, hd_o) = e.quantize_q8_1(h, t, n_embd)?;
7919 (&hq_o, &hd_o)
7920 }
7921 };
7922 fused = e.matmul_q8_fused3_t(&fa.wq, &fa.wk, &fa.wv, hq, hd, t)?;
7923 }
7924 match (fused, h_q8) {
7925 (Some(triple), _) => triple,
7926 // shared pre-quantized activation (q8_1-fast guaranteed by the caller): the
7927 // decode-exact dispatch consumes (hq, hd) instead of re-quantizing 3x.
7928 (None, Some((hq, hd))) if qkv_fast => (
7929 e.matmul_decode_exact_pre(&fa.wq, hq, hd, t)?,
7930 e.matmul_decode_exact_pre(&fa.wk, hq, hd, t)?,
7931 e.matmul_decode_exact_pre(&fa.wv, hq, hd, t)?,
7932 ),
7933 (None, _) => (
7934 e.matmul_decode_exact(&fa.wq, h, t)?,
7935 e.matmul_decode_exact(&fa.wk, h, t)?,
7936 e.matmul_decode_exact(&fa.wv, h, t)?,
7937 ),
7938 }
7939 };
7940 // M3/Hy3 have no attention output gate — wq out is exactly q; skip the split.
7941 let gated = geometry.attention_gate == memra_gguf::config::AttentionGateKind::FusedQ;
7942 let (mut q, gate) = if gated {
7943 let mut q = vbuf(e, t * n_head * head_dim)?; // fully written by q_gate_split
7944 let mut gate = vbuf(e, t * n_head * head_dim)?; // fully written by q_gate_split
7945 e.q_gate_split(&qf, &mut q, &mut gate, head_dim, n_head, t)?;
7946 (q, Some(gate))
7947 } else {
7948 (qf, None)
7949 };
7950
7951 let mut qn = vbuf(e, t * n_head * head_dim)?; // fully written by rms_norm
7952 e.rms_norm(
7953 &q,
7954 fa.q_norm.float_data(),
7955 &mut qn,
7956 head_dim,
7957 n_head * t,
7958 eps,
7959 )?;
7960 q = qn;
7961 let mut kn = vbuf(e, t * n_head_kv * head_dim)?; // fully written by rms_norm
7962 e.rms_norm(
7963 &k,
7964 fa.k_norm.float_data(),
7965 &mut kn,
7966 head_dim,
7967 n_head_kv * t,
7968 eps,
7969 )?;
7970 k = kn;
7971 let rope_dims = geometry.n_rot as usize;
7972 e.rope_neox(
7973 &mut q,
7974 pos_d,
7975 head_dim,
7976 rope_dims,
7977 n_head,
7978 t,
7979 geometry.rope_base,
7980 1.0,
7981 )?;
7982 e.rope_neox(
7983 &mut k,
7984 pos_d,
7985 head_dim,
7986 rope_dims,
7987 n_head_kv,
7988 t,
7989 geometry.rope_base,
7990 1.0,
7991 )?;
7992
7993 // append T new K/V columns to the resident QUANTIZED cache. k/v are token-major [T, kv_dim]
7994 // f32; append-quantize each of the T token rows into the byte cache (q8_0 K / q5_1 V).
7995 let kvl = cache.kv[il].as_mut().unwrap();
7996 let (kv_dim_k, kv_dim_v, ktb, vtb) =
7997 (kvl.kv_dim_k, kvl.kv_dim_v, kvl.k_tok_bytes, kvl.v_tok_bytes);
7998 if let Some(ctr) = stream_ctr {
7999 // stream: ONE batched append at the device counter (rows kernel = the per-view warp
8000 // math on a (block, token) grid, documented byte-identical); host len is a stale
8001 // LOWER BOUND under pre-issue (drain reconciles it).
8002 e.append_kv_quantized_rows_dc(
8003 &k,
8004 &v,
8005 &mut kvl.k,
8006 &mut kvl.v,
8007 ctr,
8008 t,
8009 kv_dim_k,
8010 kv_dim_v,
8011 ktb,
8012 vtb,
8013 crate::Engine::kv_fp8_on(),
8014 )?;
8015 } else {
8016 for i in 0..t {
8017 let k_row = k.slice(i * kv_dim_k..(i + 1) * kv_dim_k);
8018 let v_row = v.slice(i * kv_dim_v..(i + 1) * kv_dim_v);
8019 e.append_kv_quantized_view(
8020 &k_row,
8021 &v_row,
8022 &mut kvl.k,
8023 &mut kvl.v,
8024 kvl.len + i,
8025 kv_dim_k,
8026 kv_dim_v,
8027 ktb,
8028 vtb,
8029 crate::Engine::kv_fp8_on(),
8030 )?;
8031 }
8032 kvl.len += t;
8033 }
8034
8035 // BIT-IDENTICAL VERIFY ATTENTION (spec-exactness fix): the FP accumulation order must be
8036 // byte-for-byte identical to the eager decode path. fa_prefill uses a different tile size
8037 // (BLOCK_Q=64, BK=32) and online-softmax structure than fa_decode's split-K + combine,
8038 // which changes FP summation order and can flip argmax at tight logit margins. Query row r
8039 // attends to keys [0..base_len+r+1) — each successive row sees one more key (the causal
8040 // property). This matches eager: decode appends k at len, then fa_decode sees t_kv = len+1
8041 // keys. The verify appends all T tokens first but bounds the key range per row.
8042 //
8043 // MULTI-ROW FUSED PATH (the long-ctx spec fix, 2026-07-03): when every row takes the vec
8044 // kernel (base_len+1 >= FA_VEC_MIN_TKV), ONE fa_decode_rows launch executes the exact
8045 // per-row program for all T rows (grid.z = row, per-row n_splits from the same
8046 // fa_split_keys formula) — replacing T x (2 launches + 2 dtod copies + 5 partial allocs)
8047 // and multiplying resident CTAs by T on a latency-bound kernel. Bit-identical per row by
8048 // construction; kernel-check pins rows-vs-loop byte identity, run-spec is the end gate.
8049 // Short ctx (any row below the vec crossover) and MEMRA_NO_FA_VEC/MEMRA_FA_ROWS_OFF keep the
8050 // per-row loop (whose fa_decode picks scalar/vec per row exactly like eager decode).
8051 let mut attn = vbuf(e, t * n_head * head_dim)?; // fully written by every FA arm below
8052 let base_len = kvl.len - t; // KV len BEFORE this round's T tokens were appended
8053 // T=1 INCLUDED (2026-07-05): p-min cuts the draft to 1 in ~75% of rounds on hard
8054 // (agentic) content — the old t>1 gate sent those rounds to the per-row loop (262us/row
8055 // + q-row copy + per-row allocs vs 93us/row through the fused kernel at grid.z=1, same
8056 // program). nsys accounting: 1088 of 1456 verify FA launches were T=1 escapees.
8057 // LEAN T=1 ARM (MEMRA_SPEC_LEAN, close35): at t==1, q IS one row and fa_decode on it is
8058 // the EXACT eager decode dispatch (vec_q_v2 + combine_f32; the rows pair measured +50us
8059 // at m=1). Byte-identical: kernel-check pins rows-vs-loop identity, and the per-row loop
8060 // at t=1 is fa_decode on the same q with zero-offset copies. Gates arbitrate.
8061 if let Some(ctr) = stream_ctr {
8062 // STREAM ARM: causal base from the device counter; host kvl.len is a stale lower
8063 // bound used only for the split-sizing upper bound (+64 slack covers M pre-issued
8064 // rounds at K<=8). Views span the bound; per-row limits derive in-kernel.
8065 let upper = kvl.len + t + 64;
8066 let k_view = e.view_u8(&kvl.k, (upper.min(cache.max_ctx)) * ktb);
8067 let v_view = e.view_u8(&kvl.v, (upper.min(cache.max_ctx)) * vtb);
8068 e.fa_decode_rows_dc(
8069 &q,
8070 &k_view,
8071 &v_view,
8072 &mut attn,
8073 head_dim,
8074 n_head,
8075 n_head_kv,
8076 ctr,
8077 upper.min(cache.max_ctx),
8078 t,
8079 scale,
8080 ktb,
8081 vtb,
8082 0,
8083 false,
8084 )?;
8085 } else if spec_lean() && t == 1 {
8086 let t_kv = base_len + 1;
8087 let k_view = e.view_u8(&kvl.k, t_kv * ktb);
8088 let v_view = e.view_u8(&kvl.v, t_kv * vtb);
8089 e.fa_decode_kvmod(
8090 &q,
8091 &k_view,
8092 &v_view,
8093 &mut attn,
8094 head_dim,
8095 n_head,
8096 n_head_kv,
8097 t_kv,
8098 scale,
8099 ktb,
8100 vtb,
8101 crate::Engine::kv_fp8_on(),
8102 )?;
8103 } else if e.fa_rows_eligible(base_len, head_dim) {
8104 let k_view = e.view_u8(&kvl.k, (base_len + t) * ktb);
8105 let v_view = e.view_u8(&kvl.v, (base_len + t) * vtb);
8106 e.fa_decode_rows(
8107 &q,
8108 &k_view,
8109 &v_view,
8110 &mut attn,
8111 head_dim,
8112 n_head,
8113 n_head_kv,
8114 base_len,
8115 t,
8116 scale,
8117 ktb,
8118 vtb,
8119 None,
8120 false,
8121 crate::Engine::kv_fp8_on(),
8122 None,
8123 )?;
8124 } else {
8125 for r in 0..t {
8126 let t_kv_r = base_len + r + 1; // this row sees keys [0..t_kv_r)
8127 let k_view_r = e.view_u8(&kvl.k, t_kv_r * ktb);
8128 let v_view_r = e.view_u8(&kvl.v, t_kv_r * vtb);
8129 // copy q row into an owned buffer (fa_decode takes &CudaSlice, not CudaView)
8130 let mut q_row = vbuf(e, n_head * head_dim)?; // fully written by copy_view_into
8131 let q_src = q.slice(r * n_head * head_dim..(r + 1) * n_head * head_dim);
8132 e.copy_view_into(&mut q_row, 0, &q_src, n_head * head_dim)?;
8133 let mut attn_row = vbuf(e, n_head * head_dim)?; // fully written by fa_decode
8134 e.fa_decode_kvmod(
8135 &q_row,
8136 &k_view_r,
8137 &v_view_r,
8138 &mut attn_row,
8139 head_dim,
8140 n_head,
8141 n_head_kv,
8142 t_kv_r,
8143 scale,
8144 ktb,
8145 vtb,
8146 crate::Engine::kv_fp8_on(),
8147 )?;
8148 e.copy_into(
8149 &mut attn,
8150 r * n_head * head_dim,
8151 &attn_row,
8152 n_head * head_dim,
8153 )?;
8154 }
8155 }
8156
8157 let attn_g = match &gate {
8158 Some(gate) => {
8159 let mut gsig = vbuf(e, t * n_head * head_dim)?; // fully written by sigmoid
8160 e.sigmoid(gate, &mut gsig, t * n_head * head_dim)?;
8161 let mut ag = vbuf(e, t * n_head * head_dim)?; // fully written by mul
8162 e.mul(&attn, &gsig, &mut ag, t * n_head * head_dim)?;
8163 ag
8164 }
8165 None => attn,
8166 };
8167 // DECODE-EXACT wo projection: at m>=5 (K=4+ with pending) the generic matmul would use dp4a
8168 // (128-thread, different FP sum order than MMVQ). Force MMVQ for bit-identity with decode.
8169 Ok(e.matmul_decode_exact(&fa.wo, &attn_g, t)?)
8170 }
8171
8172 /// Context-linear bytes for a plain serving session's trunk cache.
8173 pub fn plain_session_kv_bytes_per_token(&self) -> usize {
8174 crate::cache::cache_bytes_per_token_for_plan(
8175 &self.cfg,
8176 &self.plan,
8177 0,
8178 self.plan.layers.len(),
8179 )
8180 }
8181
8182 /// `(logical bytes/token, ring-capped bytes/token, ring row cap)` for exact admission.
8183 pub fn plain_session_kv_shape(&self) -> (usize, usize, usize) {
8184 (
8185 self.plain_session_kv_bytes_per_token(),
8186 crate::cache::cache_ring_bytes_per_token_for_plan(
8187 &self.cfg,
8188 &self.plan,
8189 0,
8190 self.plan.layers.len(),
8191 ),
8192 crate::cache::cache_ring_row_cap_for_plan(&self.plan),
8193 )
8194 }
8195
8196 /// Context-linear bytes for a speculative serving session: trunk cache plus persistent MTP
8197 /// scratch. With no MTP head this equals the plain coefficient.
8198 pub fn spec_session_kv_bytes_per_token(&self) -> usize {
8199 let scratch = self
8200 .mtp
8201 .iter()
8202 .chain(self.mtp_extra.iter())
8203 .map(|mtp| {
8204 let (_, _, k, v) = mtp_scratch_layout(&self.cfg, mtp.geom.as_ref());
8205 k + v
8206 })
8207 .sum::<usize>();
8208 self.plain_session_kv_bytes_per_token()
8209 .saturating_add(scratch)
8210 }
8211
8212 /// Spec twin of [`HybridModel::plain_session_kv_shape`]; Step35's persistent MTP scratch is
8213 /// capped by the same SWA ring rows as the trunk.
8214 pub fn spec_session_kv_shape(&self) -> (usize, usize, usize) {
8215 let total = self.spec_session_kv_bytes_per_token();
8216 let (_, mut ring, rows) = self.plain_session_kv_shape();
8217 if rows > 0 {
8218 ring = ring.saturating_add(
8219 self.mtp
8220 .iter()
8221 .chain(self.mtp_extra.iter())
8222 .map(|mtp| {
8223 let (_, _, k, v) = mtp_scratch_layout(&self.cfg, mtp.geom.as_ref());
8224 k + v
8225 })
8226 .sum::<usize>(),
8227 );
8228 }
8229 (total, ring, rows)
8230 }
8231
8232 /// Greedy MTP speculative decode (§B). Token-identical to `generate(prompt, max_new)` but uses
8233 /// the NextN head to draft K tokens then verifies them in one batched target forward.
8234 /// Returns (generated tokens, total_drafted, total_accepted) so the caller can report
8235 /// acceptance rate. `k` = draft length per round.
8236 ///
8237 /// GRAPH DRAFT (stage 2 of graph-grade spec): when the model is all-Dense and the MTP head is
8238 /// Dense (no MoE host readbacks), the fixed-shape T=1 MTP forward is CUDA-graph-captured ONCE
8239 /// and replayed per draft step — the ~40 eager launches per drafted token collapse into one
8240 /// graph dispatch; only the 4-byte token id (and 4-byte p-min confidence) round-trip per step.
8241 /// Event tracking is disabled for the whole call (generate_graph pattern) so every buffer the
8242 /// captured graph references is event-free; the spec loop is strictly single-stream.
8243 /// MEMRA_SPEC_NOGRAPH=1 forces the eager draft chain.
8244 /// SAMPLED mode (MEMRA_SPEC_TEMP>0) has its OWN capture (gumbel-perturbed in-graph argmax,
8245 /// device Philox event counter, persistent q retention) — graph-vs-eager sampled streams are
8246 /// bit-identical for the same (seed, prompt, K, temp); see the sampled-graph setup in
8247 /// generate_spec_inner2.
8248 /// Multi-turn session: trunk cache + MTP draft scratch persist across generate calls, so
8249 /// turn N+1 primes ONLY its new suffix (the 124k-conversation daily pattern — re-priming a
8250 /// 32k history costs ~54s; a suffix prime costs seconds). APPEND-ONLY by construction: the
8251 /// hybrid linear-attn states are in-place (no position index), so a session can extend but
8252 /// never rewind — `committed` is the exact token list whose state the caches hold (includes
8253 /// any overshoot tokens past max_new; the caller renders from `committed`, not its own echo).
8254 pub fn new_session(
8255 &self,
8256 e: &Engine,
8257 max_ctx: usize,
8258 ) -> Result<SpecSession, Box<dyn std::error::Error>> {
8259 Ok(SpecSession {
8260 // STAGE-OWNED KV (lane/pp2-spec 2026-08-06): `pp::new_cache`, not `Cache::new`. This
8261 // is the SERVING spec-session path, and with the ppN door open across two cards a
8262 // primary-homed cache makes every remote stage peer-read its OWN KV on every verify
8263 // round — the wrong-card class already fixed on the two batched serving paths
8264 // (worker.rs 2483 / 2837). With the door shut `new_cache` IS `Cache::new` (same
8265 // branch, same allocations), so single-device behavior is byte-unchanged.
8266 cache: crate::pp::new_cache_planned(e, &self.cfg, &self.plan, max_ctx)?,
8267 scratch: self.new_mtp_scratch(e, max_ctx)?,
8268 committed: Vec::new(),
8269 last_h: None,
8270 next_pred: None,
8271 sctr: 0,
8272 uctr: 0,
8273 draft_ctx: None,
8274 pending_tok: None,
8275 turn_ckpt: None,
8276 telem: SpecTelemetryCounters::default(),
8277 capture_at: None,
8278 boundary_captures: Vec::new(),
8279 ckpt_at: None,
8280 })
8281 }
8282
8283 /// SPEC-ON-CACHE-HIT restore (lane/spec-on-cache-hit, 2026-08-18 — PORT-PLAN item 3,
8284 /// research/cache-spec-design-20260814, scoped to WHOLE-ENTRY restores only): build a
8285 /// SpecSession around a trunk cache the worker already restored from a prefix-cache
8286 /// entry, re-installing the entry's published draft plane as the MTP scratch rows
8287 /// `[0..prefix.len())` and the entry's boundary hidden as `last_h`, then feeding the
8288 /// prompt SUFFIX here — through EXACTLY the plain path's program selection — so the
8289 /// worker always receives a fully-warm continuation session (committed = whole
8290 /// prompt, `next_pred` + `last_h` set; caller sets `next_pred` from the entry's
8291 /// boundary logits on the empty-suffix shape).
8292 ///
8293 /// PROGRAM LAW (the splitiso two-programs class, learned AGAIN in this lane's own
8294 /// gate): the identity target for a converted hit is the PLAIN hit serving the same
8295 /// request, and plain feeds a carried suffix via eager `decode_step` below
8296 /// PRIME_MIN_T and via `prime_cache` at/above it (prefill_tick's arms). The generate
8297 /// path's tokenwise arm routes qwen35-class through the BATCHED T=1 program
8298 /// (`spec_target_step_h`) instead — ULP-different suffix rows, and the gate measured
8299 /// the near-tie flip at generated token ~8 (research/spec-cache-20260818, qwen r3).
8300 /// So the suffix is fed HERE, mirroring prefill_tick arm-for-arm, not handed to the
8301 /// burst prime.
8302 ///
8303 /// SEED RULE (both sampling regimes; lane/sampled-hit-spec 2026-08-19, sampled draw
8304 /// added by lane/sampled-spec-quality 2026-08-19). The boundary token is produced by
8305 /// EXACTLY the rule the cold burst entry applies to its own first token from the same
8306 /// logits row: `argmax` when greedy, and a `sample_boundary_token` draw at Philox
8307 /// counter 0 when sampled. Both shapes are covered — the entry's boundary logits on a
8308 /// full-cover (empty-suffix) hit, this feed's own boundary logits on a suffix hit.
8309 /// That is what keeps a restored session seed-identical to a cold one PER SEED: the
8310 /// cold session draws from the identical row at counter 0 and then runs its rounds from
8311 /// counter 1, so the restored session admits with `sctr = 1` after its own draw.
8312 /// The WORKER owns the one refusal this constructor cannot see — a constrained request.
8313 /// (The penalized-sampled refusal was LIFTED once the burst's penalty window learned to
8314 /// span the session: `committed` here is the WHOLE prompt, so the restored session's
8315 /// window is the cold session's window. It comes back if `MEMRA_SPEC_PEN_SESSION=0`.)
8316 ///
8317 /// NOT the rolled-back partial-restore hazard: the caller restores at exactly the
8318 /// entry's captured endpoint (`e.pos`) through the shipping whole-entry path;
8319 /// mid-entry (`at < e.pos`) trunk restores stay behind MEMRA_PREFIX_PARTIAL_RESTORE
8320 /// and are never routed here.
8321 ///
8322 /// Failure contract: `Err((Some(cache), why))` before any trunk mutation — the
8323 /// worker rebuilds the plain carrier and the hit serves plain, byte-unchanged.
8324 /// `Err((None, why))` after the suffix feed began — the carrier is part-fed and
8325 /// UNUSABLE; the worker serves the request cold-plain (correct, slower) and the
8326 /// entry stays published for the next request.
8327 #[allow(clippy::too_many_arguments)]
8328 pub fn spec_session_from_restored(
8329 &self,
8330 e: &Engine,
8331 mut cache: Cache,
8332 prefix: Vec<u32>,
8333 suffix: &[u32],
8334 draft_k: &CudaSlice<u8>,
8335 draft_v: &CudaSlice<u8>,
8336 draft_k_tok_bytes: usize,
8337 draft_v_tok_bytes: usize,
8338 draft_len: usize,
8339 last_h: &[f32],
8340 // The ENTRY's boundary logits row (the full-cover shape's seed source). May be empty
8341 // when a suffix follows — the feed's own logits are the boundary then.
8342 boundary_logits: &[f32],
8343 // The request's sampler, or None for greedy. Owned here so the seed rule lives in
8344 // ONE place instead of being half-applied by the worker.
8345 sampling: Option<SpecSampling>,
8346 require_anchor: bool,
8347 max_ctx: usize,
8348 // STABLE-BOUNDARY REPUBLICATION (lane/frspec-multiturn-cache, 2026-08-21): ABSOLUTE
8349 // prompt position to split the suffix feed at and capture the extended-entry
8350 // publication + this session's `turn_ckpt` — the worker's stable pre-generation
8351 // boundary (`plain_checkpoint_boundary`). None = legacy prompt-end republication.
8352 // WHY: the prompt-end capture below includes the template's live generation header
8353 // (`<|im_start|>assistant\n<think>\n`), which the next turn's re-render replaces, so
8354 // for a hybrid (whole-entry restores only) every extended entry's last ~2 tokens
8355 // diverged from every future prompt and the hit boundary FROZE at the first
8356 // lcp-split entry forever (measured: cached 6811 of 38228 by turn 8, B4).
8357 republish_at: Option<usize>,
8358 ) -> Result<SpecSession, (Option<Cache>, String)> {
8359 let pos = prefix.len();
8360 let fail = |cache: Cache, msg: String| -> Result<SpecSession, (Option<Cache>, String)> {
8361 Err((Some(cache), msg))
8362 };
8363 if self.mtp.is_none() {
8364 return fail(cache, "no MTP head attached (nothing to draft with)".into());
8365 }
8366 if pos == 0 {
8367 return fail(cache, "empty committed prefix".into());
8368 }
8369 if cache.pos != pos {
8370 let msg = format!(
8371 "restored cache pos {} != restored prefix len {pos}",
8372 cache.pos
8373 );
8374 return fail(cache, msg);
8375 }
8376 if draft_len != pos {
8377 return fail(
8378 cache,
8379 format!("draft plane len {draft_len} != restored prefix len {pos}"),
8380 );
8381 }
8382 if pos + suffix.len() >= max_ctx {
8383 return fail(
8384 cache,
8385 format!(
8386 "prompt {} + suffix would not leave generation room in ctx {max_ctx}",
8387 pos + suffix.len(),
8388 ),
8389 );
8390 }
8391 let mut scratch = match MtpScratch::new(
8392 e,
8393 &self.cfg,
8394 &self.plan,
8395 max_ctx,
8396 self.mtp.as_ref().and_then(|m| m.geom.as_ref()),
8397 ) {
8398 Ok(s) => s,
8399 Err(err) => return fail(cache, format!("draft scratch alloc failed: {err}")),
8400 };
8401 if scratch.kv.ring.is_some() {
8402 return fail(
8403 cache,
8404 "ring-backed draft scratch (Step35 SWA) cannot take a flat prefix restore".into(),
8405 );
8406 }
8407 if scratch.kv.k_tok_bytes != draft_k_tok_bytes
8408 || scratch.kv.v_tok_bytes != draft_v_tok_bytes
8409 {
8410 return fail(
8411 cache,
8412 format!(
8413 "draft plane layout {draft_k_tok_bytes}/{draft_v_tok_bytes} != scratch \
8414 {}/{} bytes/token (stale entry across a format change)",
8415 scratch.kv.k_tok_bytes, scratch.kv.v_tok_bytes,
8416 ),
8417 );
8418 }
8419 if pos > scratch.cap {
8420 return fail(
8421 cache,
8422 format!(
8423 "draft plane rows {pos} exceed scratch capacity {}",
8424 scratch.cap
8425 ),
8426 );
8427 }
8428 let kb = pos * draft_k_tok_bytes;
8429 let vb = pos * draft_v_tok_bytes;
8430 if draft_k.len() < kb || draft_v.len() < vb {
8431 return fail(
8432 cache,
8433 format!(
8434 "truncated draft plane: K {} < {kb} or V {} < {vb} bytes",
8435 draft_k.len(),
8436 draft_v.len(),
8437 ),
8438 );
8439 }
8440 if kb > 0 {
8441 if let Err(err) = e.copy_u8_into(&mut scratch.kv.k, 0, draft_k, kb) {
8442 return fail(cache, format!("draft K restore copy failed: {err}"));
8443 }
8444 }
8445 if vb > 0 {
8446 if let Err(err) = e.copy_u8_into(&mut scratch.kv.v, 0, draft_v, vb) {
8447 return fail(cache, format!("draft V restore copy failed: {err}"));
8448 }
8449 }
8450 if let Err(err) = scratch.set_len(e, pos) {
8451 return fail(cache, format!("draft scratch len set failed: {err}"));
8452 }
8453 let mut last_h_dev = if last_h.len() == self.cfg.n_embd as usize {
8454 // anchor upload failure is acceptance-only when a suffix feed follows (fill
8455 // row-0 falls back to zeros) but FATAL for an empty-suffix continuation (the
8456 // burst entry asserts committed + last_h + next_pred) — the caller says which.
8457 e.htod(last_h).ok()
8458 } else {
8459 None
8460 };
8461 if require_anchor && last_h_dev.is_none() {
8462 return fail(
8463 cache,
8464 "empty-suffix continuation requires the entry's boundary hidden anchor".into(),
8465 );
8466 }
8467 let mut committed = prefix;
8468 // Set on BOTH shapes below (suffix-fed and full-cover) — never left None, which is
8469 // what the empty-suffix continuation assert in the burst entry requires.
8470 let next_pred;
8471 // Philox: (0,0) at admit exactly like a fresh session; a sampled boundary draw below
8472 // consumes counter 0 and leaves 1, which is the state a cold session reaches after
8473 // drawing its own first token from the same row.
8474 let mut sctr = 0u32;
8475 let sampled = sampling.is_some_and(|s| s.temp > 0.0) && spec_sampled_boundary_on();
8476 // Penalty window for the boundary draw: the last `penalty_last_n` tokens of the WHOLE
8477 // prompt, which is what the cold session's own burst sees (Item 2's window). Built
8478 // after the suffix joins `committed` below.
8479 let mut boundary_captures: Vec<SpecBoundaryCapture> = Vec::new();
8480 let mut restored_turn_ckpt: Option<SpecCheckpoint> = None;
8481 if !suffix.is_empty() {
8482 // ---- SUFFIX FEED, mirroring prefill_tick's program selection exactly ----
8483 // From here on the trunk cache mutates: failures return Err((None, _)) and
8484 // the worker serves the request cold-plain instead of reusing the carrier.
8485 let dirty =
8486 |msg: String| -> Result<SpecSession, (Option<Cache>, String)> { Err((None, msg)) };
8487 let n_embd = self.cfg.n_embd as usize;
8488 let t = suffix.len();
8489 let mut h_rows = match e.uninit(t * n_embd) {
8490 Ok(b) => b,
8491 Err(err) => return fail(cache, format!("suffix hidden buffer alloc: {err}")),
8492 };
8493 // STABLE-BOUNDARY split (see `republish_at`): feed stops at the boundary so the
8494 // in-place GDN conv/ssm state can be snapshotted there — the only moment it
8495 // exists (the cold prime-split law). suffix-relative; None = one-segment legacy.
8496 let b_rel = republish_at
8497 .and_then(|abs| abs.checked_sub(pos))
8498 .filter(|&r| r > 0 && r < t);
8499 let mut feed_logits = Vec::new();
8500 let tokenwise_env = std::env::var("MEMRA_PRIME_TOKENWISE").is_ok()
8501 || e.frozen_cpu_experts_prefer_tokenwise_prime();
8502 let mut fed = 0usize;
8503 for seg_end in [b_rel, Some(t)].into_iter().flatten() {
8504 if seg_end <= fed {
8505 continue;
8506 }
8507 let seg = &suffix[fed..seg_end];
8508 let batched = seg.len() >= crate::hybrid_forward::PRIME_MIN_T && !tokenwise_env;
8509 if batched {
8510 // prefill_tick's prime arm: request-level prime_cache call; tokens still
8511 // queued after this segment ride `queued_after` so Step35 arm selection
8512 // stays keyed to the request's end (tick-seg law).
8513 match self.prime_cache(e, seg, &mut cache, t - seg_end) {
8514 Ok((l, _h_seed, hiddens)) => {
8515 if let Err(err) =
8516 e.copy_into(&mut h_rows, fed * n_embd, &hiddens, seg.len() * n_embd)
8517 {
8518 return dirty(format!("suffix hidden copy: {err}"));
8519 }
8520 feed_logits = l;
8521 }
8522 Err(err) => return dirty(format!("suffix prime failed: {err}")),
8523 }
8524 } else {
8525 // prefill_tick's tokenwise arm: eager decode_step, one token at a time.
8526 for (i, &tok) in seg.iter().enumerate() {
8527 match self.decode_step_h(e, tok, &mut cache) {
8528 Ok((l, h)) => {
8529 if let Err(err) =
8530 e.copy_into(&mut h_rows, (fed + i) * n_embd, &h, n_embd)
8531 {
8532 return dirty(format!("suffix hidden copy: {err}"));
8533 }
8534 feed_logits = l;
8535 }
8536 Err(err) => return dirty(format!("suffix decode_step failed: {err}")),
8537 }
8538 }
8539 }
8540 fed = seg_end;
8541 if Some(seg_end) == b_rel {
8542 // The stable pre-generation boundary: capture the extended-entry
8543 // publication AND this session's own turn checkpoint here instead of at
8544 // prompt-end (both would otherwise carry the volatile live-header tail
8545 // the next re-render replaces). Failure silent, turn_ckpt convention.
8546 debug_assert_eq!(
8547 cache.pos,
8548 pos + seg_end,
8549 "stable-boundary capture off the feed split"
8550 );
8551 if spec_restore_republish_on() {
8552 if let Ok(snap) = cache.snapshot(e) {
8553 boundary_captures.push(SpecBoundaryCapture {
8554 snap,
8555 pos: pos + seg_end,
8556 logits: feed_logits.clone(),
8557 last_h: capture_boundary_hidden(e, &h_rows, seg_end, n_embd),
8558 });
8559 }
8560 }
8561 let anchor: Result<CudaSlice<f32>, Box<dyn std::error::Error>> =
8562 e.uninit(n_embd).and_then(|mut a| {
8563 e.copy_view_into(
8564 &mut a,
8565 0,
8566 &h_rows.slice((seg_end - 1) * n_embd..seg_end * n_embd),
8567 n_embd,
8568 )?;
8569 Ok(a)
8570 });
8571 if let (Ok(snap), Ok(last_h)) = (cache.snapshot(e), anchor) {
8572 restored_turn_ckpt = Some(SpecCheckpoint {
8573 snap,
8574 pos: pos + seg_end,
8575 last_h,
8576 });
8577 }
8578 }
8579 }
8580 // Draft-scratch fill for the suffix rows, predecessor-paired: row `pos` reads
8581 // the entry's boundary anchor (zeros fallback — acceptance-only), row `pos+i`
8582 // reads h_rows[i-1]. Chunked like the generate path's fill (transients scale
8583 // with T). Fill failures are acceptance-only — truncate to the restored rows
8584 // and continue; the burst's own set_len keeps the invariant.
8585 let mtp = self.mtp.as_ref().expect("mtp checked above");
8586 let (embd_qt, embd_rb) = self.embd.qt_and_row_bytes(n_embd);
8587 let embd_gpu = if spec_host_embd() {
8588 None
8589 } else {
8590 Some(
8591 self.embd_gpu
8592 .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload")),
8593 )
8594 };
8595 let embd_dev = embd_gpu.map(|g| (g, embd_qt, embd_rb));
8596 let fill_chunk = 4096usize;
8597 let mut filled = true;
8598 let mut start = 0usize;
8599 'fill: while start < t {
8600 let end = (start + fill_chunk).min(t);
8601 let tc = end - start;
8602 let Ok(mut phs) = e.zeros(tc * n_embd) else {
8603 filled = false;
8604 break 'fill;
8605 };
8606 let (src_lo, dst_off, n_copy) = if start == 0 {
8607 (0, n_embd, (tc - 1) * n_embd)
8608 } else {
8609 ((start - 1) * n_embd, 0, tc * n_embd)
8610 };
8611 if start == 0 {
8612 if let Some(lh) = last_h_dev.as_ref() {
8613 if e.copy_into(&mut phs, 0, lh, n_embd).is_err() {
8614 filled = false;
8615 break 'fill;
8616 }
8617 }
8618 }
8619 if n_copy > 0
8620 && e.copy_view_into(
8621 &mut phs,
8622 dst_off,
8623 &h_rows.slice(src_lo..src_lo + n_copy),
8624 n_copy,
8625 )
8626 .is_err()
8627 {
8628 filled = false;
8629 break 'fill;
8630 }
8631 if self
8632 .mtp_kv_fill_all(
8633 e,
8634 &suffix[start..end],
8635 &phs,
8636 pos + start,
8637 &mut scratch,
8638 embd_dev,
8639 )
8640 .is_err()
8641 {
8642 filled = false;
8643 break 'fill;
8644 }
8645 start = end;
8646 }
8647 if !filled {
8648 // acceptance-only: drafts over missing suffix rows are cheap and wrong,
8649 // so keep only the restored rows resident and let verify arbitrate.
8650 if let Err(err) = scratch.set_len(e, pos) {
8651 return dirty(format!("scratch truncation after failed fill: {err}"));
8652 }
8653 }
8654 // EXTENDED-ENTRY PUBLICATION (lane/sampled-spec-quality, Item 3 — the fix for
8655 // "a restored spec session never publishes an extended entry", SAMPLED-HIT.md
8656 // finding (d)). Pre-lane, publication was armed only for COLD sessions
8657 // (`spec_resumed == 0` in the worker) and both engine capture sites require a
8658 // non-continuation burst — but a converted hit's first burst IS a continuation,
8659 // so a growing conversation learned exactly ONE boundary and turn 3 could never
8660 // hit a longer prefix than turn 2 did.
8661 //
8662 // WHERE, and why it is safe here: `cache.pos == prefix + suffix` at this exact
8663 // line — the trunk is primed over the whole prompt, nothing is generated, and the
8664 // draft plane rows [0..prompt) are filled just above. That is a complete
8665 // whole-entry boundary (`pos == fed_len`), the same shape the cold seed capture
8666 // publishes; the worker's existing publication sweep picks it up because it is
8667 // keyed on non-empty `boundary_captures` and is sampler- and resume-independent.
8668 // NOT the partial-restore hazard: the boundary is this session's own prompt END,
8669 // never mid-entry, so `entry_pos != fed_len` still refuses on the way back in.
8670 // Failure is SILENT by design (the turn_ckpt / boundary-capture convention):
8671 // publication is an optimization, never a correctness dependency.
8672 //
8673 // SUPERSEDED WHEN `republish_at` FIRED (lane/frspec-multiturn-cache): a prompt-end
8674 // entry's tail is the live generation header the next re-render replaces, so on a
8675 // hybrid (whole-entry restores) it can never serve the conversation's next turn —
8676 // the stable-boundary capture above IS this publication, minus the poisoned tail.
8677 if spec_restore_republish_on() && boundary_captures.is_empty() {
8678 debug_assert_eq!(
8679 cache.pos,
8680 pos + t,
8681 "extended-entry capture must sit at the restored session's prompt end",
8682 );
8683 if let Ok(snap) = cache.snapshot(e) {
8684 boundary_captures.push(SpecBoundaryCapture {
8685 snap,
8686 pos: pos + t,
8687 logits: feed_logits.clone(),
8688 last_h: capture_boundary_hidden(e, &h_rows, t, n_embd),
8689 });
8690 }
8691 }
8692 // continuation seed: the feed's boundary logits ARE the plain path's boundary
8693 // logits (same program), so greedy's argmax here is plain's first emitted token,
8694 // and the sampled draw is the cold sampled session's own first token.
8695 next_pred = Some(if sampled {
8696 let sp = sampling.expect("sampled implies a sampler");
8697 // `committed` is still the restored prefix here; the suffix joins it below —
8698 // so this is the last-N window over the WHOLE prompt, exactly the cold
8699 // session's own window at its first token.
8700 let hist = pen_window_seed(&committed, suffix, sp.penalty_last_n);
8701 match sample_boundary_token(
8702 e,
8703 &feed_logits,
8704 &sp,
8705 &hist,
8706 &mut sctr,
8707 "restore-suffix-feed",
8708 ) {
8709 Ok(t) => t,
8710 // the trunk is already fed: hand nothing back, the worker serves the
8711 // request cold-plain. Never fall back to an argmax — that would put a
8712 // greedy token in a sampled stream to save a slow path.
8713 Err(err) => {
8714 return dirty(format!("boundary token draw failed: {err}"));
8715 }
8716 }
8717 } else {
8718 argmax(&feed_logits) as u32
8719 });
8720 let mut lh = match e.uninit(n_embd) {
8721 Ok(b) => b,
8722 Err(err) => return dirty(format!("boundary hidden alloc: {err}")),
8723 };
8724 if let Err(err) = e.copy_view_into(
8725 &mut lh,
8726 0,
8727 &h_rows.slice((t - 1) * n_embd..t * n_embd),
8728 n_embd,
8729 ) {
8730 return dirty(format!("boundary hidden copy: {err}"));
8731 }
8732 last_h_dev = Some(lh);
8733 committed.extend_from_slice(suffix);
8734 } else {
8735 // FULL-COVER shape (empty suffix — the identical-repeat / agent-loop shape): the
8736 // ENTRY's boundary logits are the boundary row, and this is the token the cold
8737 // session emits from that same row. Owned here rather than in the worker so the
8738 // sampled draw cannot be half-applied on one shape (the worker used to argmax it).
8739 if boundary_logits.is_empty() {
8740 return fail(
8741 cache,
8742 "full-cover restore without the entry's boundary logits".into(),
8743 );
8744 }
8745 next_pred = Some(if sampled {
8746 let sp = sampling.expect("sampled implies a sampler");
8747 let hist = pen_window_seed(&committed, &[], sp.penalty_last_n);
8748 match sample_boundary_token(
8749 e,
8750 boundary_logits,
8751 &sp,
8752 &hist,
8753 &mut sctr,
8754 "restore-full-cover",
8755 ) {
8756 Ok(t) => t,
8757 // nothing has been mutated on this shape — hand the carrier back and let
8758 // the hit serve PLAIN (the banked pre-lane path).
8759 Err(err) => {
8760 return fail(cache, format!("boundary token draw failed: {err}"));
8761 }
8762 }
8763 } else {
8764 argmax(boundary_logits) as u32
8765 });
8766 }
8767 Ok(SpecSession {
8768 cache,
8769 scratch,
8770 committed,
8771 last_h: last_h_dev,
8772 next_pred,
8773 sctr,
8774 uctr: 0,
8775 draft_ctx: None,
8776 pending_tok: None,
8777 // Stable-boundary capture from the split feed above (None on the legacy shape):
8778 // a restored session previously parked WITHOUT a checkpoint, so the next turn's
8779 // affinity probe declined ("no turn checkpoint retained") and the conversation
8780 // fell back to the frozen prefix entry forever.
8781 turn_ckpt: restored_turn_ckpt,
8782 telem: SpecTelemetryCounters::default(),
8783 capture_at: None,
8784 boundary_captures,
8785 ckpt_at: None,
8786 })
8787 }
8788
8789 /// Forced-gate exact state comparison. This intentionally reads the real live prefixes from
8790 /// their owning PP devices: matching emitted ids alone would miss a stale `len_d`, recurrent
8791 /// snapshot, or draft-KV row that only corrupts the following round.
8792 pub fn optipipe_compare_session_state(
8793 &self,
8794 e: &Engine,
8795 reference: &SpecSession,
8796 candidate: &SpecSession,
8797 ) -> Result<OptiForkStateIdentity, Box<dyn std::error::Error>> {
8798 fn fail(what: &str) -> Box<dyn std::error::Error> {
8799 format!("optipipe state mismatch: {what}").into()
8800 }
8801 fn same_f32(a: &[f32], b: &[f32]) -> bool {
8802 a.len() == b.len() && a.iter().zip(b).all(|(x, y)| x.to_bits() == y.to_bits())
8803 }
8804 fn compare_layers(
8805 es: &Engine,
8806 range: std::ops::Range<usize>,
8807 reference: &SpecSession,
8808 candidate: &SpecSession,
8809 report: &mut OptiForkStateIdentity,
8810 ) -> Result<(), Box<dyn std::error::Error>> {
8811 for il in range {
8812 match (&reference.cache.kv[il], &candidate.cache.kv[il]) {
8813 (Some(a), Some(b)) => {
8814 if a.len != b.len {
8815 return Err(fail(&format!(
8816 "layer {il} host KV len {} != {}",
8817 a.len, b.len
8818 )));
8819 }
8820 let ad = es.dtoh_i32(&a.len_d)?;
8821 let bd = es.dtoh_i32(&b.len_d)?;
8822 if ad != bd || ad.first().copied() != Some(a.len as i32) {
8823 return Err(fail(&format!(
8824 "layer {il} device KV len {ad:?} != {bd:?} (host={})",
8825 a.len,
8826 )));
8827 }
8828 let kb = a.len * a.k_tok_bytes;
8829 let vb = a.len * a.v_tok_bytes;
8830 if kb > 0 {
8831 let ak = es.dtoh_u8_view(&a.k.slice(0..kb))?;
8832 let bk = es.dtoh_u8_view(&b.k.slice(0..kb))?;
8833 if ak != bk {
8834 let at = ak.iter().zip(&bk).position(|(x, y)| x != y).unwrap();
8835 return Err(fail(&format!(
8836 "layer {il} K bytes at byte {at} row {} offset {}: {} != {}",
8837 at / a.k_tok_bytes,
8838 at % a.k_tok_bytes,
8839 ak[at],
8840 bk[at],
8841 )));
8842 }
8843 }
8844 if vb > 0 {
8845 let av = es.dtoh_u8_view(&a.v.slice(0..vb))?;
8846 let bv = es.dtoh_u8_view(&b.v.slice(0..vb))?;
8847 if av != bv {
8848 let at = av.iter().zip(&bv).position(|(x, y)| x != y).unwrap();
8849 return Err(fail(&format!(
8850 "layer {il} V bytes at byte {at} row {} offset {}: {} != {}",
8851 at / a.v_tok_bytes,
8852 at % a.v_tok_bytes,
8853 av[at],
8854 bv[at],
8855 )));
8856 }
8857 }
8858 report.trunk_kv_bytes += kb + vb;
8859 }
8860 (None, None) => {}
8861 _ => return Err(fail(&format!("layer {il} KV presence"))),
8862 }
8863 match (&reference.cache.recur[il], &candidate.cache.recur[il]) {
8864 (Some(a), Some(b)) => {
8865 let ac = es.dtoh(&a.conv_state)?;
8866 let bc = es.dtoh(&b.conv_state)?;
8867 if !same_f32(&ac, &bc) {
8868 return Err(fail(&format!("layer {il} conv state")));
8869 }
8870 let as_ = es.dtoh(&a.ssm_state)?;
8871 let bs = es.dtoh(&b.ssm_state)?;
8872 if !same_f32(&as_, &bs) {
8873 return Err(fail(&format!("layer {il} SSM state")));
8874 }
8875 report.recurrent_bytes += (ac.len() + as_.len()) * 4;
8876 }
8877 (None, None) => {}
8878 _ => return Err(fail(&format!("layer {il} recurrent presence"))),
8879 }
8880 }
8881 Ok(())
8882 }
8883
8884 if reference.committed != candidate.committed {
8885 return Err(fail("committed token ids"));
8886 }
8887 if reference.cache.pos != candidate.cache.pos
8888 || reference.cache.max_ctx != candidate.cache.max_ctx
8889 {
8890 return Err(fail("cache pos/capacity"));
8891 }
8892 if reference.pending_tok != candidate.pending_tok
8893 || reference.next_pred != candidate.next_pred
8894 || reference.sctr != candidate.sctr
8895 || reference.uctr != candidate.uctr
8896 {
8897 return Err(fail("pending/prediction/counter tail"));
8898 }
8899
8900 let mut report = OptiForkStateIdentity::default();
8901 if let Some(fence) = crate::pp::pp_cuts(self.layers.len()) {
8902 let rt = crate::pp::PpNRt::get(e)?;
8903 for stage in 0..rt.n_stages() {
8904 let _scope = rt.enter(stage);
8905 compare_layers(
8906 rt.engine(stage, e),
8907 fence[stage]..fence[stage + 1],
8908 reference,
8909 candidate,
8910 &mut report,
8911 )?;
8912 }
8913 } else {
8914 compare_layers(e, 0..self.layers.len(), reference, candidate, &mut report)?;
8915 }
8916
8917 if reference.scratch.plane_count() != candidate.scratch.plane_count() {
8918 return Err(fail("draft scratch plane count"));
8919 }
8920 for index in 0..reference.scratch.plane_count() {
8921 let (a, _) = reference.scratch.plane(index);
8922 let (b, _) = candidate.scratch.plane(index);
8923 if a.len != b.len
8924 || a.kv_dim_k != b.kv_dim_k
8925 || a.kv_dim_v != b.kv_dim_v
8926 || a.k_tok_bytes != b.k_tok_bytes
8927 || a.v_tok_bytes != b.v_tok_bytes
8928 || e.dtoh_i32(&a.len_d)? != e.dtoh_i32(&b.len_d)?
8929 {
8930 return Err(fail(&format!("draft scratch plane {index} length/layout")));
8931 }
8932 let kb = a.len * a.k_tok_bytes;
8933 let vb = a.len * a.v_tok_bytes;
8934 if kb > 0 && e.dtoh_u8_view(&a.k.slice(0..kb))? != e.dtoh_u8_view(&b.k.slice(0..kb))? {
8935 return Err(fail(&format!("draft scratch plane {index} K bytes")));
8936 }
8937 if vb > 0 && e.dtoh_u8_view(&a.v.slice(0..vb))? != e.dtoh_u8_view(&b.v.slice(0..vb))? {
8938 return Err(fail(&format!("draft scratch plane {index} V bytes")));
8939 }
8940 report.scratch_kv_bytes += kb + vb;
8941 }
8942
8943 match (&reference.last_h, &candidate.last_h) {
8944 (Some(a), Some(b)) => {
8945 let ah = e.dtoh(a)?;
8946 let bh = e.dtoh(b)?;
8947 if !same_f32(&ah, &bh) {
8948 return Err(fail("last hidden/seed bytes"));
8949 }
8950 report.hidden_bytes = ah.len() * 4;
8951 }
8952 (None, None) => {}
8953 _ => return Err(fail("last hidden/seed presence")),
8954 }
8955 Ok(report)
8956 }
8957
8958 /// SESSION-AFFINITY REWIND (lane/session-affinity, 2026-08-05): roll `sess` back to its
8959 /// retained prompt-end checkpoint, so a request whose prompt matches
8960 /// `committed[..rewind_pos()]` exactly can resume there and prime only its own delta.
8961 ///
8962 /// EXACTNESS. After this returns, the session is byte-for-byte the state it was in AT that
8963 /// boundary: full-attn KV truncated to it (append-only, position-addressed), GDN conv/ssm
8964 /// restored from the device copy taken there, draft scratch length reset, `committed`
8965 /// truncated, `last_h` = the boundary's predecessor anchor. That is precisely the state a
8966 /// fresh prime of `committed[..pos]` would have produced, so the following suffix prime and
8967 /// every burst after it are identical to a cold run of the same token stream — the
8968 /// committed-tokens-authoritative contract.
8969 ///
8970 /// `next_pred` and `pending_tok` are CLEARED: both describe generation past the boundary,
8971 /// which the rewind discards. The caller therefore must supply a non-empty suffix (a
8972 /// rewound session cannot serve an empty-suffix continuation burst — there is nothing to
8973 /// continue). The persistent draft graph survives: it bakes only session-stable pointers
8974 /// (the scratch KV, the resident embedding), none of which the rewind moves.
8975 ///
8976 /// The checkpoint is CONSUMED (`turn_ckpt` taken): its snapshot buffers are freed here, and
8977 /// this turn's own prime installs a fresh one at the new prompt end. Returns the position
8978 /// rewound to, or `None` when the session holds no checkpoint (caller: full re-prime).
8979 pub fn spec_rewind_to_checkpoint(
8980 &self,
8981 e: &Engine,
8982 sess: &mut SpecSession,
8983 ) -> Result<Option<usize>, Box<dyn std::error::Error>> {
8984 if sess.turn_ckpt.as_ref().is_some_and(|ckpt| {
8985 !sess.cache.can_rollback(&ckpt.snap, 0) || !sess.scratch.can_rewind_to(ckpt.pos)
8986 }) {
8987 return Err(
8988 "SWA ring rewind checkpoint has been lapped; full re-prime required".into(),
8989 );
8990 }
8991 let Some(ckpt) = sess.turn_ckpt.take() else {
8992 return Ok(None);
8993 };
8994 assert!(
8995 ckpt.pos <= sess.committed.len(),
8996 "checkpoint past committed ({} > {})",
8997 ckpt.pos,
8998 sess.committed.len()
8999 );
9000 // Restore through each layer's owning engine. A single primary-engine rollback is not
9001 // sufficient when the serving cache is stage-owned under cross-device PP.
9002 crate::pp::restore_cache_checkpoint(e, self, None, &mut sess.cache, &ckpt.snap)?;
9003 debug_assert_eq!(
9004 sess.cache.pos, ckpt.pos,
9005 "rollback landed off the checkpoint"
9006 );
9007 sess.scratch.set_len(e, ckpt.pos)?;
9008 sess.committed.truncate(ckpt.pos);
9009 sess.last_h = Some(ckpt.last_h);
9010 sess.next_pred = None;
9011 sess.pending_tok = None;
9012 Ok(Some(ckpt.pos))
9013 }
9014
9015 /// Grow a parked speculative session to `target_cap` and rewind it to its retained turn
9016 /// checkpoint without re-priming the checkpoint prefix.
9017 ///
9018 /// The trunk cache is restored exactly like a plain grown cache: append-only full-attention
9019 /// KV rows come from the parked cache, while recurrent state comes from the checkpoint's
9020 /// owned snapshot. The MTP scratch is also context-linear and its rows below the checkpoint
9021 /// remain authoritative, so they are copied into a fresh larger scratch before its length is
9022 /// truncated. Pointer-baking draft graphs are dropped and recaptured on the next burst.
9023 ///
9024 /// All fallible work completes before `sess` is mutated. A failed allocation or copy leaves
9025 /// the parked session intact, allowing the caller one reclaim-and-retry attempt.
9026 pub fn spec_grow_and_rewind_to_checkpoint(
9027 &self,
9028 e: &Engine,
9029 sess: &mut SpecSession,
9030 target_cap: usize,
9031 ) -> Result<Option<usize>, Box<dyn std::error::Error>> {
9032 if target_cap <= sess.cache.max_ctx {
9033 return self.spec_rewind_to_checkpoint(e, sess);
9034 }
9035 let Some(ckpt) = sess.turn_ckpt.as_ref() else {
9036 return Ok(None);
9037 };
9038 if ckpt.pos == 0 || ckpt.pos > sess.committed.len() {
9039 return Err(format!(
9040 "checkpoint pos {} outside committed length {}",
9041 ckpt.pos,
9042 sess.committed.len(),
9043 )
9044 .into());
9045 }
9046 if ckpt.pos > target_cap {
9047 return Err(format!(
9048 "checkpoint pos {} exceeds grown capacity {target_cap}",
9049 ckpt.pos,
9050 )
9051 .into());
9052 }
9053
9054 let mut grown_cache = crate::pp::new_cache_planned(e, &self.cfg, &self.plan, target_cap)?;
9055 let mut grown_scratch = self.new_mtp_scratch(e, target_cap)?;
9056 crate::pp::restore_cache_checkpoint(
9057 e,
9058 self,
9059 Some(&sess.cache),
9060 &mut grown_cache,
9061 &ckpt.snap,
9062 )?;
9063
9064 if sess.scratch.plane_count() != grown_scratch.plane_count() {
9065 return Err("checkpoint draft plane count mismatch".into());
9066 }
9067 for index in 0..sess.scratch.plane_count() {
9068 let (src, _) = sess.scratch.plane(index);
9069 let (dst, _) = grown_scratch.plane_mut(index);
9070 if ckpt.pos > src.len
9071 || src.kv_dim_k != dst.kv_dim_k
9072 || src.kv_dim_v != dst.kv_dim_v
9073 || src.k_tok_bytes != dst.k_tok_bytes
9074 || src.v_tok_bytes != dst.v_tok_bytes
9075 {
9076 return Err(format!(
9077 "checkpoint draft plane {index} layout mismatch (pos {}, source len {})",
9078 ckpt.pos, src.len,
9079 )
9080 .into());
9081 }
9082 let kb = ckpt.pos * src.k_tok_bytes;
9083 let vb = ckpt.pos * src.v_tok_bytes;
9084 if kb > 0 {
9085 e.copy_u8_into(&mut dst.k, 0, &src.k, kb)?;
9086 }
9087 if vb > 0 {
9088 e.copy_u8_into(&mut dst.v, 0, &src.v, vb)?;
9089 }
9090 }
9091 grown_scratch.set_len(e, ckpt.pos)?;
9092 // The old scratch is dropped immediately after publication below. Bound its D2D reads
9093 // first; growth happens once per rewritten turn, outside the decode hot loop.
9094 e.stream().synchronize()?;
9095
9096 let ckpt = sess
9097 .turn_ckpt
9098 .take()
9099 .expect("checkpoint remained present through transactional grow");
9100 let pos = ckpt.pos;
9101 sess.cache = grown_cache;
9102 sess.scratch = grown_scratch;
9103 sess.committed.truncate(pos);
9104 sess.last_h = Some(ckpt.last_h);
9105 sess.next_pred = None;
9106 sess.pending_tok = None;
9107 sess.draft_ctx = None;
9108 debug_assert_eq!(sess.cache.pos, pos, "grown rewind landed off checkpoint");
9109 debug_assert!(
9110 (0..sess.scratch.plane_count()).all(|index| sess.scratch.plane(index).0.len == pos),
9111 "grown draft rewind landed off checkpoint"
9112 );
9113 Ok(Some(pos))
9114 }
9115
9116 /// Commit a carried pending bonus (see SpecSession::pending_tok): one T=1 trunk pass
9117 /// (its logits' argmax becomes next_pred) + the draft-KV fill at the carried anchor —
9118 /// byte-identical to the pre-carry session tail. Required before a non-empty-suffix
9119 /// prime, a sampled turn, or parking a session for pool reuse. No-op without a pending.
9120 /// `sampling` is the sampler of the request that will CONSUME the resulting `next_pred`
9121 /// (lane/sampled-spec-quality): this is a boundary site like any other, so a sampled
9122 /// consumer must get a DRAWN token, not an argmax. Pass `None` from the park/demote
9123 /// callers — a pending only ever exists on the GREEDY tail, and the consumer of a
9124 /// park-time flush is a future request whose sampler is not knowable here (residual
9125 /// named at the pool-resume probe in worker.rs and in SAMPLED-QUALITY.md).
9126 pub fn spec_flush_pending(
9127 &self,
9128 e: &Engine,
9129 sess: &mut SpecSession,
9130 sampling: Option<SpecSampling>,
9131 ) -> Result<(), Box<dyn std::error::Error>> {
9132 let Some(b) = sess.pending_tok.take() else {
9133 return Ok(());
9134 };
9135 if self.mtp.is_none() {
9136 return Err("pending carry requires an MTP head".into());
9137 }
9138 let n_embd = self.cfg.n_embd as usize;
9139 let (embd_qt, embd_rb) = self.embd.qt_and_row_bytes(n_embd);
9140 let embd_gpu = if spec_host_embd() {
9141 None
9142 } else {
9143 Some(
9144 self.embd_gpu
9145 .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload")),
9146 )
9147 };
9148 let embd_dev = embd_gpu.map(|g| (g, embd_qt, embd_rb));
9149 let pos_b = sess.cache.pos;
9150 sess.scratch.set_len(e, pos_b)?;
9151 let (lg_b, hb) = self.spec_target_step_h(e, b, &mut sess.cache)?;
9152 sess.next_pred = Some(match sampling {
9153 Some(sp) if sp.temp > 0.0 && spec_sampled_boundary_on() => {
9154 // window includes `b` itself: it is committed by this pass, and the pre-lane
9155 // code never counted a boundary token in the penalty history at all.
9156 let hist = pen_window_seed(&sess.committed, &[b], sp.penalty_last_n);
9157 sample_boundary_token(e, &lg_b, &sp, &hist, &mut sess.sctr, "flush-pending")?
9158 }
9159 _ => argmax(&lg_b) as u32,
9160 });
9161 let anchor = sess
9162 .last_h
9163 .as_ref()
9164 .expect("pending carry requires last_h (the predecessor-row anchor)");
9165 self.mtp_kv_fill_all(e, &[b], anchor, pos_b, &mut sess.scratch, embd_dev)?;
9166 sess.last_h = Some(hb);
9167 sess.committed.push(b);
9168 Ok(())
9169 }
9170
9171 /// Solo target feed used only at speculative round boundaries. Step35 serving made its
9172 /// staged batched B=1 graph authoritative, so a speculative session must enter and leave
9173 /// rounds through that same graph. Other model families keep their eager T=1 contract.
9174 fn spec_target_step_h(
9175 &self,
9176 e: &Engine,
9177 token: u32,
9178 cache: &mut Cache,
9179 ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9180 if !self.sliding_gated_moe_batch_program() && !self.batched_serving_numeric_class() {
9181 return self.decode_step_h(e, token, cache);
9182 }
9183 let pos0 = cache.pos;
9184 let (logits, hidden) = self.decode_step_t_core(e, &[token], pos0, cache, None, None)?;
9185 Ok((e.dtoh(&logits)?, hidden))
9186 }
9187
9188 /// The archs whose LIVE B=1 serving runs the generic BATCHED numeric class (decode_step_batch
9189 /// walk + batched head), so their spec verify must run the SAME class. MoE learned this
9190 /// 2026-08-14 AM (4b777ccc5); the dense hybrid reproduced the identical near-tie flip class
9191 /// the same day on Qwen3.8-27B — eager-class verify logits drift from batched-class serving
9192 /// logits ("1 ULP at layer 2 → 2.3e-1 logit maxdiff at the head"), and the GDN recurrence
9193 /// carries the drift until a near-tie flips deep in generation. One predicate so the five
9194 /// dispatch sites cannot drift apart again.
9195 /// Draft-graph head admissibility (lane/draftcost-moe, 2026-08-20): the capture body
9196 /// (`mtp_head_forward_cap`) supports Dense heads AND resident-MoE heads
9197 /// (`Ffn::Moe(m) if m.dev_exps.is_some()`); non-resident MoE still refuses inside the
9198 /// capture and the caller falls back to the eager chain by design. Trunk FFN class is
9199 /// irrelevant — the graph body is the HEAD forward only. One predicate for all three
9200 /// eligibility sites so they cannot drift (the serving numeric-class lesson).
9201 fn mtp_graph_capturable(&self) -> bool {
9202 self.mtp
9203 .as_ref()
9204 .map(|m| match &m.ffn {
9205 crate::hybrid::Ffn::Dense { .. } => true,
9206 crate::hybrid::Ffn::Moe(mo) => mo.dev_exps.is_some(),
9207 })
9208 .unwrap_or(false)
9209 }
9210
9211 fn batched_serving_numeric_class(&self) -> bool {
9212 self.plan
9213 .trunk_operations()
9214 .contains(&memra_gguf::model_plan::OperationKind::GatedDeltaNet)
9215 }
9216
9217 /// The family the MTP verify-graph default was measured on: GatedDeltaNet state layers
9218 /// (a `recur` mixer) together with a routed-MoE FFN — Ornith-1.5-35B-A3B and its kin. The
9219 /// server-side twin of this test is `model_forces_spec_replay` (GatedDeltaNet + MoeMlp);
9220 /// keeping the engine's own version structural rather than name-based means a new
9221 /// checkpoint of the same shape inherits the default, and a different shape does not.
9222 fn vgraph_family_default(&self) -> bool {
9223 let has_linear = self
9224 .layers
9225 .iter()
9226 .any(|l| matches!(l.mixer, Mixer::Linear(_)));
9227 let has_moe = self
9228 .layers
9229 .iter()
9230 .any(|l| matches!(l.ffn, crate::hybrid::Ffn::Moe(_)));
9231 has_linear && has_moe
9232 }
9233
9234 fn sliding_gated_moe_batch_program(&self) -> bool {
9235 self.uses_sliding_gated_moe_program()
9236 }
9237
9238 fn gemma_batch_program(&self) -> bool {
9239 self.uses_gemma_program()
9240 }
9241
9242 /// Reduced-matrix admission for increment 1. This deliberately does not change the PP-2
9243 /// serving policy: the worker calls it only after `MEMRA_SPEC_PIPE=1` and an explicit spec
9244 /// session already exist.
9245 pub fn spec_pipe_available(&self, e: &Engine) -> bool {
9246 if std::env::var("MEMRA_SPEC_PIPE").as_deref() != Ok("1")
9247 || !spec_devacc()
9248 || spec_replay_env_enabled()
9249 || spec_stream()
9250 || std::env::var("MEMRA_SPEC_ADAPT").as_deref() == Ok("1")
9251 || std::env::var("MEMRA_SPEC_PMIN0").as_deref() == Ok("1")
9252 || std::env::var("MEMRA_SPEC_PP_ANATOMY").as_deref() == Ok("1")
9253 || std::env::var("MEMRA_SPEC_PMIN")
9254 .ok()
9255 .and_then(|v| v.parse::<f32>().ok())
9256 .unwrap_or(0.0)
9257 > 0.0
9258 || self.is_gemma4_e4b()
9259 || self.gemma_batch_program()
9260 || self.mtp.is_none()
9261 || !self.mtp_extra.is_empty()
9262 {
9263 return false;
9264 }
9265 let Some(cuts) = crate::pp::pp_cuts(self.layers.len()) else {
9266 return false;
9267 };
9268 if cuts.len() != 3 || crate::pp::pp2_streams_off() || !crate::pp::spec_pp_on() {
9269 return false;
9270 }
9271 crate::pp::PpNRt::get(e)
9272 .map(|rt| rt.n_stages() == 2 && rt.cross_device())
9273 .unwrap_or(false)
9274 }
9275
9276 /// Two warm greedy continuation bursts over one PP-2 interval coordinator. The two existing
9277 /// `generate_spec_inner2` call stacks own all per-session round locals; only phase issue order
9278 /// changes. No callback is accepted in increment 1 — the worker publishes each completed burst.
9279 #[allow(clippy::too_many_arguments)]
9280 pub fn generate_spec_session_pair(
9281 &self,
9282 e: &Engine,
9283 sess_a: &mut SpecSession,
9284 max_new_a: usize,
9285 k_a: usize,
9286 sess_b: &mut SpecSession,
9287 max_new_b: usize,
9288 k_b: usize,
9289 ) -> Result<((Vec<u32>, usize, usize), (Vec<u32>, usize, usize)), Box<dyn std::error::Error>>
9290 {
9291 if !self.spec_pipe_available(e) {
9292 return Err("two-session speculative pipeline is outside its reduced matrix".into());
9293 }
9294 if max_new_a == 0 || max_new_b == 0 || k_a == 0 || k_b == 0 {
9295 return Err(
9296 "two-session speculative pipeline requires non-empty positive-K bursts".into(),
9297 );
9298 }
9299 for sess in [&*sess_a, &*sess_b] {
9300 if sess.committed.is_empty()
9301 || sess.last_h.is_none()
9302 || (sess.next_pred.is_none() && sess.pending_tok.is_none())
9303 {
9304 return Err("two-session speculative pipeline requires warm continuations".into());
9305 }
9306 }
9307
9308 let graph_ok = std::env::var("MEMRA_SPEC_NOGRAPH").is_err()
9309 && !spec_host_embd()
9310 && self.mtp_graph_capturable()
9311 && self.mtp_extra.is_empty()
9312 && !crate::model::full_prec_enabled();
9313 let graph_a = graph_ok && k_a + 2 < 96;
9314 let graph_b = graph_ok && k_b + 2 < 96;
9315 let was_tracking = e.ctx().is_event_tracking();
9316 if (graph_a || graph_b) && was_tracking {
9317 unsafe {
9318 e.ctx().disable_event_tracking();
9319 }
9320 }
9321
9322 static LOGGED: std::sync::Once = std::sync::Once::new();
9323 LOGGED.call_once(|| {
9324 eprintln!("[spec-pipe] two-session PP-2 continuation pipeline engaged");
9325 });
9326 let sync = std::sync::Arc::new(SpecPipeSync::new());
9327 let lane_a = SpecPipeLane {
9328 sync: sync.clone(),
9329 lane: 0,
9330 };
9331 let lane_b = SpecPipeLane { sync, lane: 1 };
9332 let mut sess_b_ptr = SpecPipeSessionPtr(sess_b as *mut SpecSession);
9333 let (result_a, result_b) = std::thread::scope(|scope| {
9334 let b = scope.spawn(move || {
9335 let mut finish = SpecPipeFinish::new(&lane_b);
9336 let sess_b = unsafe { sess_b_ptr.get_mut() };
9337 let result = e
9338 .ctx()
9339 .bind_to_thread()
9340 .map_err(|err| err.to_string())
9341 .and_then(|_| {
9342 self.generate_spec_inner2(
9343 e,
9344 &[],
9345 max_new_b,
9346 k_b,
9347 graph_b,
9348 Some(sess_b),
9349 None,
9350 None,
9351 None,
9352 None,
9353 Some(&lane_b),
9354 )
9355 .map_err(|err| err.to_string())
9356 });
9357 finish.close(result.is_err());
9358 result
9359 });
9360 let mut finish = SpecPipeFinish::new(&lane_a);
9361 let result_a = self.generate_spec_inner2(
9362 e,
9363 &[],
9364 max_new_a,
9365 k_a,
9366 graph_a,
9367 Some(sess_a),
9368 None,
9369 None,
9370 None,
9371 None,
9372 Some(&lane_a),
9373 );
9374 finish.close(result_a.is_err());
9375 let result_b = b
9376 .join()
9377 .map_err(|_| "paired speculative session B panicked".to_string())
9378 .and_then(|r| r);
9379 (result_a, result_b)
9380 });
9381
9382 if (graph_a || graph_b) && was_tracking {
9383 unsafe {
9384 e.ctx().enable_event_tracking();
9385 }
9386 }
9387 let result_a = result_a?;
9388 let result_b = result_b.map_err(|err| -> Box<dyn std::error::Error> { err.into() })?;
9389 Ok((result_a, result_b))
9390 }
9391
9392 /// One spec-decode turn on a live session. `suffix` = the NEW tokens only (turn N+1's user
9393 /// message rendered through the chat template continuation). Returns (new tokens emitted,
9394 /// drafted, accepted); session.committed grows by suffix + emitted.
9395 pub fn generate_spec_session(
9396 &self,
9397 e: &Engine,
9398 sess: &mut SpecSession,
9399 suffix: &[u32],
9400 max_new: usize,
9401 k: usize,
9402 ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
9403 self.generate_spec_session_sampled(e, sess, suffix, max_new, k, None, None)
9404 }
9405
9406 /// Serve-path sampled spec: routes the burst through the rejection-sampling verify with
9407 /// per-SESSION Philox continuity (sess.sctr/uctr). None = env-driven (CLI) or greedy.
9408 /// Filters (top-k/p/min-p) apply SYMMETRICALLY to draft q and verify p — distribution-exact
9409 /// for the filtered target (feat/filtered-spec).
9410 ///
9411 /// `on_commit` (sse-cadence, 2026-08-05): called with each newly-emitted slice of the
9412 /// output — once right after the prime's first token, then once per round commit — so a
9413 /// streaming caller can flush text at round cadence instead of once per burst. The slices
9414 /// are disjoint, in order, and concatenate to exactly the returned token vec. Emission-
9415 /// timing only: token bytes, session state, and exactness are untouched.
9416 ///
9417 /// The returned bool is a CONTINUE-VERDICT (admission yield, 2026-08-06): `false` ends
9418 /// the burst at the current round boundary, exactly as if `max_new` had been reached —
9419 /// the caller's scheduler regains control without waiting the burst out. Burst size is
9420 /// content-neutral (spec-levers battery), so an early exit moves WHEN the burst returns,
9421 /// never what tokens say. The slice may be EMPTY (a poll-only boundary — round-stream
9422 /// drains and the defensive tail flush can land with nothing new committed).
9423 #[allow(clippy::too_many_arguments)]
9424 pub fn generate_spec_session_sampled(
9425 &self,
9426 e: &Engine,
9427 sess: &mut SpecSession,
9428 suffix: &[u32],
9429 max_new: usize,
9430 k: usize,
9431 sampling: Option<SpecSampling>,
9432 on_commit: Option<&mut dyn FnMut(&[u32]) -> bool>,
9433 ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
9434 self.generate_spec_session_sampled_prime_split(
9435 e, sess, suffix, max_new, k, sampling, None, on_commit,
9436 )
9437 }
9438
9439 /// Serve-only cold-prime segmentation twin. `prime_split` is the same stable boundary the
9440 /// plain worker would honor before entering its sub-floor tokenwise tail; warm continuations
9441 /// pass `None` and stay on the existing zero-prime path.
9442 #[allow(clippy::too_many_arguments)]
9443 pub fn generate_spec_session_sampled_prime_split(
9444 &self,
9445 e: &Engine,
9446 sess: &mut SpecSession,
9447 suffix: &[u32],
9448 max_new: usize,
9449 k: usize,
9450 sampling: Option<SpecSampling>,
9451 prime_split: Option<usize>,
9452 on_commit: Option<&mut dyn FnMut(&[u32]) -> bool>,
9453 ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
9454 self.generate_spec_session_constrained_prime_split(
9455 e,
9456 sess,
9457 suffix,
9458 max_new,
9459 k,
9460 sampling,
9461 None,
9462 prime_split,
9463 on_commit,
9464 )
9465 }
9466
9467 /// `generate_spec_session_sampled` + GRAMMAR (constrained decoding, 2026-08-03): the
9468 /// hook truncates acceptance at the first grammar-illegal token AFTER the exactness
9469 /// verify (grammar is an extra rejection rule, ordering like the batched-verify twins)
9470 /// and replaces an illegal bonus with the MASKED argmax of the target's own verify
9471 /// column — token-identical to constrained plain greedy decode. GREEDY only (the
9472 /// worker routes sampled constrained to plain decode). Acceptance under tight grammars
9473 /// may drop (drafter is unconstrained); that is measured, not hidden.
9474 #[allow(clippy::too_many_arguments)]
9475 pub fn generate_spec_session_constrained(
9476 &self,
9477 e: &Engine,
9478 sess: &mut SpecSession,
9479 suffix: &[u32],
9480 max_new: usize,
9481 k: usize,
9482 sampling: Option<SpecSampling>,
9483 constraint: Option<&mut dyn SpecConstraint>,
9484 on_commit: Option<&mut dyn FnMut(&[u32]) -> bool>,
9485 ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
9486 self.generate_spec_session_constrained_prime_split(
9487 e, sess, suffix, max_new, k, sampling, constraint, None, on_commit,
9488 )
9489 }
9490
9491 #[allow(clippy::too_many_arguments)]
9492 pub fn generate_spec_session_constrained_prime_split(
9493 &self,
9494 e: &Engine,
9495 sess: &mut SpecSession,
9496 suffix: &[u32],
9497 max_new: usize,
9498 k: usize,
9499 sampling: Option<SpecSampling>,
9500 constraint: Option<&mut dyn SpecConstraint>,
9501 prime_split: Option<usize>,
9502 on_commit: Option<&mut dyn FnMut(&[u32]) -> bool>,
9503 ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
9504 if constraint.is_some() && sampling.is_some_and(|s| s.temp > 0.0) {
9505 return Err(
9506 "constrained spec decode is greedy-only (worker routes sampled \
9507 constrained to plain decode)"
9508 .into(),
9509 );
9510 }
9511 // PENDING-CARRY entry flush: a carried bonus precedes any new suffix in the sequence,
9512 // so it must commit BEFORE the suffix primes; the sampled path doesn't carry (its
9513 // round-0 accept needs the commit pass's logits). Empty-suffix greedy bursts — the
9514 // serve continuation case — consume the carry in-loop with zero solo passes.
9515 if sess.pending_tok.is_some()
9516 && (!suffix.is_empty() || sampling.map_or(false, |s| s.temp > 0.0))
9517 {
9518 self.spec_flush_pending(e, sess, sampling)?;
9519 }
9520
9521 // FULL_PREC forces the EAGER draft: the graph capture would enclose cuBLASLt f32 GEMV
9522 // (the FloatBf16 else-branches) and a bf16_to_f32 dequant alloc — neither is stream-capture
9523 // safe. Eager rides matmul/matmul_decode_exact, which dequant FloatBf16 on use. (§item 2.)
9524 let graph_draft = std::env::var("MEMRA_SPEC_NOGRAPH").is_err()
9525 && !spec_host_embd()
9526 && self.mtp_graph_capturable()
9527 && self.mtp_extra.is_empty()
9528 && k + 2 < 96
9529 && !crate::model::full_prec_enabled();
9530 let was_tracking = e.ctx().is_event_tracking();
9531 if graph_draft && was_tracking {
9532 unsafe {
9533 e.ctx().disable_event_tracking();
9534 }
9535 }
9536 let r = self.generate_spec_inner2(
9537 e,
9538 suffix,
9539 max_new,
9540 k,
9541 graph_draft,
9542 Some(sess),
9543 sampling,
9544 constraint,
9545 on_commit,
9546 prime_split,
9547 None,
9548 );
9549 if graph_draft && was_tracking {
9550 unsafe {
9551 e.ctx().enable_event_tracking();
9552 }
9553 }
9554 let (out, d, a) = r?;
9555 Ok((out, d, a))
9556 }
9557
9558 pub fn generate_spec(
9559 &self,
9560 e: &Engine,
9561 prompt: &[u32],
9562 max_new: usize,
9563 k: usize,
9564 ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
9565 if crate::pp::pp_cuts(self.layers.len()).is_some()
9566 && !self.rewrite_allowed(memra_gguf::execution_manifest::RewriteSurface::Pipeline)
9567 {
9568 return Err("pipeline rewrite is not qualified for speculative decode".into());
9569 }
9570 if !self.rewrite_allowed(memra_gguf::execution_manifest::RewriteSurface::MtpSpec) {
9571 return Err("speculative rewrite is not qualified for this ModelPlan".into());
9572 }
9573 // FULL_PREC forces eager (see generate_spec_session note): CUDA graph capture cannot
9574 // enclose cuBLASLt f32 GEMV or the bf16_to_f32 dequant alloc the FloatBf16 path needs.
9575 let graph_draft = std::env::var("MEMRA_SPEC_NOGRAPH").is_err()
9576 && !spec_host_embd()
9577 && self.mtp_graph_capturable()
9578 && self.mtp_extra.is_empty()
9579 && k + 2 < 96
9580 && !crate::model::full_prec_enabled();
9581 if !graph_draft {
9582 return self.generate_spec_inner2(
9583 e, prompt, max_new, k, false, None, None, None, None, None, None,
9584 );
9585 }
9586 let was_tracking = e.ctx().is_event_tracking();
9587 if was_tracking {
9588 unsafe {
9589 e.ctx().disable_event_tracking();
9590 }
9591 }
9592 let r = self.generate_spec_inner2(
9593 e, prompt, max_new, k, true, None, None, None, None, None, None,
9594 );
9595 if was_tracking {
9596 unsafe {
9597 e.ctx().enable_event_tracking();
9598 }
9599 }
9600 r
9601 }
9602
9603 fn generate_spec_inner2(
9604 &self,
9605 e: &Engine,
9606 prompt: &[u32],
9607 max_new: usize,
9608 k: usize,
9609 graph_draft: bool,
9610 mut sess: Option<&mut SpecSession>,
9611 sampling: Option<SpecSampling>,
9612 mut constraint: Option<&mut dyn SpecConstraint>,
9613 mut on_commit: Option<&mut dyn FnMut(&[u32]) -> bool>,
9614 prime_split: Option<usize>,
9615 pipe: Option<&SpecPipeLane>,
9616 ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
9617 assert!(k >= 1, "k must be >= 1");
9618 if let Some(p) = pipe {
9619 p.setup_begin()?;
9620 }
9621 // sse-cadence flush cursor: everything in out[..flushed] has been handed to on_commit.
9622 let mut flushed = 0usize;
9623 // admission yield (2026-08-06): on_commit's continue-verdict; false = end the burst
9624 // at the next round boundary (same exit as max_new reached — the session tail runs).
9625 // Initialized by the unconditional post-prime flush below.
9626 let mut keep_going;
9627 let mtp = self
9628 .mtp
9629 .as_ref()
9630 .expect("generate_spec requires an MTP head (nextn_predict_layers>0)");
9631 let n_vocab = self.output.out_features();
9632 // FR-Spec: the draft head may be TRIMMED (fewer rows than n_vocab); the draft argmax runs
9633 // over the draft vocab and the winning index maps through d2t to a TARGET token id.
9634 // Everything downstream (verify/accept/commit) sees target ids only — exactness unchanged.
9635 let d_vocab = mtp
9636 .shared_head_head
9637 .as_ref()
9638 .unwrap_or(&self.output)
9639 .out_features();
9640 if !self.mtp_extra.is_empty() {
9641 if self.plan.draft_source != memra_gguf::model_plan::DraftSourcePlan::Embedded
9642 || self.plan.mtp_blocks.len() != self.mtp_head_count()
9643 || mtp.d2t.is_some()
9644 {
9645 return Err(
9646 "multi-head MTP requires one embedded canonical block per loaded head".into(),
9647 );
9648 }
9649 for (offset, head) in self.mtp_extra.iter().enumerate() {
9650 if head.d2t.is_some()
9651 || head
9652 .shared_head_head
9653 .as_ref()
9654 .unwrap_or(&self.output)
9655 .out_features()
9656 != d_vocab
9657 {
9658 return Err(format!(
9659 "embedded MTP head {} has incompatible draft vocabulary",
9660 offset + 1
9661 )
9662 .into());
9663 }
9664 }
9665 eprintln!(
9666 "[mtp-chain] heads={} policy=step-modulo prefix-replay kv=per-head",
9667 self.mtp_head_count()
9668 );
9669 }
9670 let n_embd = self.cfg.n_embd as usize;
9671 // SESSION MODE: reuse the live cache/scratch, prime only the suffix. `base` = tokens
9672 // already committed (their state is in the caches); 0 = fresh single-shot call.
9673 let session_mode = sess.is_some();
9674 let max_ctx = match sess.as_ref() {
9675 Some(s) => s.cache.max_ctx,
9676 None => prompt.len() + max_new + k + 8,
9677 };
9678 let mut own_cache;
9679 let mut own_scratch;
9680 // PREFIX-CACHE capture request threaded out of the session (lane/spec-prefix-cache):
9681 // (requested split, destination list). Single-shot per burst; fresh calls have none.
9682 let mut sess_capture: Option<(Option<usize>, &mut Vec<SpecBoundaryCapture>)> = None;
9683 // STABLE-BOUNDARY turn-checkpoint request (lane/frspec-multiturn-cache): ABSOLUTE
9684 // committed-length position; consumed one-shot like `capture_at`. None = legacy
9685 // prompt-end capture below.
9686 let mut ckpt_req: Option<usize> = None;
9687 let (
9688 cache,
9689 scratch,
9690 mut sess_tail,
9691 mut sess_draft_slot,
9692 mut sess_pending_slot,
9693 sess_ckpt_slot,
9694 sess_telem,
9695 ): (
9696 &mut Cache,
9697 &mut MtpScratch,
9698 Option<(
9699 &mut Vec<u32>,
9700 &mut Option<CudaSlice<f32>>,
9701 &mut Option<u32>,
9702 &mut u32,
9703 &mut u32,
9704 )>,
9705 Option<&mut Option<DraftGraphCtx>>,
9706 Option<&mut Option<u32>>,
9707 Option<&mut Option<SpecCheckpoint>>,
9708 Option<&SpecTelemetryCounters>,
9709 ) = match sess.take() {
9710 Some(sr) => {
9711 let SpecSession {
9712 cache,
9713 scratch,
9714 committed,
9715 last_h,
9716 next_pred,
9717 sctr: s_sctr,
9718 uctr: s_uctr,
9719 draft_ctx,
9720 pending_tok,
9721 turn_ckpt,
9722 telem,
9723 capture_at,
9724 boundary_captures,
9725 ckpt_at,
9726 } = sr;
9727 sess_capture = Some((capture_at.take(), boundary_captures));
9728 ckpt_req = ckpt_at.take();
9729 (
9730 cache,
9731 scratch,
9732 Some((committed, last_h, next_pred, s_sctr, s_uctr)),
9733 Some(draft_ctx),
9734 Some(pending_tok),
9735 Some(turn_ckpt),
9736 Some(telem),
9737 )
9738 }
9739 None => {
9740 // STAGE-OWNED KV (lane/pp2-spec 2026-08-06) — see `new_session`. Door shut =
9741 // `Cache::new` verbatim.
9742 own_cache = crate::pp::new_cache_planned(e, &self.cfg, &self.plan, max_ctx)?;
9743 // Persistent scratch = max_ctx rows (~2KB/token quantized).
9744 own_scratch = self.new_mtp_scratch(e, max_ctx)?;
9745 (
9746 &mut own_cache,
9747 &mut own_scratch,
9748 None,
9749 None,
9750 None,
9751 None,
9752 None,
9753 )
9754 }
9755 };
9756 if scratch.plane_count() != self.mtp_head_count() {
9757 return Err(format!(
9758 "MTP scratch/head count mismatch ({}/{})",
9759 scratch.plane_count(),
9760 self.mtp_head_count()
9761 )
9762 .into());
9763 }
9764 let base = cache.pos;
9765 // PENDING-CARRY consume (2026-08-01): a carried bonus reaches here only on the
9766 // empty-suffix GREEDY continuation path (generate_spec_session_sampled flushed every
9767 // other case). It enters the round loop as round-0's pending — verify col 0 — exactly
9768 // like a mid-burst full-accept boundary: no init feed, no tail commit pass.
9769 let carried_pending: Option<u32> = sess_pending_slot.as_mut().and_then(|s| s.take());
9770 // PERSISTENT DRAFT KV (the only mode since 2026-07-08 — the legacy round-local scratch,
9771 // MEMRA_SPEC_KVLOCAL, measured -35 acceptance pts on the 27B p3 sweep and was removed;
9772 // acceptance-only — exactness is verify's job either way).
9773 // HIDDEN-PAIRING CONVENTION (DEFAULT = predecessor-row, 2026-07-04 — the 27B acceptance
9774 // unlock, +16pts): the MTP head is TRAINED on rows pairing token x_p with the trunk
9775 // hidden of its PREDECESSOR h_{p-1} (the reference engine's mtp_update shifts the target
9776 // hiddens right by one; its draft step 0 feeds (id_last, TRUE hidden of the row id_last
9777 // was sampled from)). memra's historical convention paired SAME-ROW (x_p, h_p) in the fill
9778 // and seeded chain step 0 through an extra MTP pass on a duplicated token (the
9779 // pseudo-seed) — measured 27B p2 K=3 acceptance 0.569 vs 0.731, p3 0.445 vs 0.63+, and
9780 // the chain steps j>=1 were already predecessor-shaped, so ONLY the fill + step-0 seed
9781 // move. The fill shifts by one and the chain seeds from the predecessor's true hidden
9782 // DIRECTLY (vh_seed / vx[j-1]) — the pseudo pass disappears (one MTP-block pass saved
9783 // per round on top of the acceptance win). Draft-quality-only: exactness stays the
9784 // verify's job either way. (The legacy same-row pairing seam, MEMRA_SPEC_HSAME, and its
9785 // pseudo-seed passes were removed 2026-07-08 — predecessor pairing won by +16 acc pts;
9786 // the legacy round-local scratch, MEMRA_SPEC_KVLOCAL, went with it.)
9787 // REPLAY-FREE PARTIAL ACCEPT (default, 2026-07-03): partial rounds keep the verify's own
9788 // bit-identical committed-prefix state (KV truncate + recur rebuild from the VerifyCkpt)
9789 // and leave the bonus PENDING — no duplicate trunk pass (profiled ~0.54 extra full weight
9790 // reads/round at long ctx). MEMRA_SPEC_REPLAY=1 restores the legacy rollback+replay (A/B
9791 // + fallback seam).
9792 // Qwen35-MoE replay pin LIFTED (lane/draftcost-moe, 2026-08-20). The pin's stated
9793 // bar — the retained verify-state commit proven equivalent to sequential serving —
9794 // was waiting on this arch running the serving batched verify class, which the
9795 // t-parallel admission (this lane, increment 1) provided: the VerifyCkpt the
9796 // replay-free commit consumes is now produced by the SAME serving-class verify that
9797 // qualified dense qwen35 on 2026-08-15 (where the per-round duplicate replay
9798 // measured 69 -> 30 tok/s). Qualification receipts (run-spec K=1..8 both arms,
9799 // 8-prompt replay-vs-replay-free canary, long-prompt cell):
9800 // research/draftcost-moe-20260820/RECEIPTS.md. MEMRA_SPEC_REPLAY=1 stays the
9801 // rollback + A/B seam.
9802 let spec_replay = spec_replay_env_enabled();
9803 if constraint.is_some() && spec_replay {
9804 return Err(
9805 "constrained spec decode does not support MEMRA_SPEC_REPLAY=1 \
9806 (legacy replay commits an unmasked bonus)"
9807 .into(),
9808 );
9809 }
9810 // TRUE-HIDDEN REFRESH (default in persistent-draft-KV mode): every round overwrites the
9811 // committed positions' scratch entries from the verify's exact hiddens (mtp_kv_fill batch)
9812 // instead of keeping chain-approximate entries. MEMRA_SPEC_NOREFRESH=1 = legacy (A/B seam).
9813 let refresh = std::env::var("MEMRA_SPEC_NOREFRESH").is_err();
9814 if !refresh && !self.mtp_extra.is_empty() {
9815 return Err("multi-head MTP requires exact accepted-prefix refresh".into());
9816 }
9817
9818 // prime: BATCHED cache prime (prime_cache — the measured #1 e2e gap: tokenwise primed at
9819 // ~102/38 tok/s vs the engine's ~2000-5900 tok/s batched prefill). prime_cache returns the
9820 // full pre-output_norm hidden stack [T, n_embd], which IS prompt_h (the persistent-draft-KV
9821 // mtp_kv_fill input) — no per-token collection needed. Prompts below PRIME_MIN_T, and
9822 // MEMRA_PRIME_TOKENWISE=1, and frozen Hy3 CPU/GPU expert splits take the tokenwise
9823 // decode_step_h loop. The latter avoids transient GPU staging of the spilled expert bank.
9824 // EMPTY-SUFFIX CONTINUATION (serve bursts): a session turn with NO new tokens resumes
9825 // generation exactly where the last turn stopped — no prime at all. The stashed
9826 // `next_pred` plays prime_logits' role: it is the token produced from the logits after
9827 // committed.last() by the same rule this entry applies to a cold prime's last row —
9828 // an argmax when greedy, a `sample_boundary_token` draw when sampled (the burst tail,
9829 // or `spec_session_from_restored` for a converted prefix-cache hit, did the drawing
9830 // where the sampler and the session's Philox counters were live). `last_h` seeds the
9831 // predecessor pairing below. Fresh calls and non-empty suffixes take the normal path.
9832 let continuation = prompt.is_empty();
9833 if continuation {
9834 assert!(session_mode, "empty prompt requires a session");
9835 assert!(
9836 sess_tail
9837 .as_ref()
9838 .map_or(false, |(c, lh, np, _, _)| !c.is_empty()
9839 && lh.is_some()
9840 && (np.is_some() || carried_pending.is_some())),
9841 "empty-suffix continuation needs a primed session (committed + last_h + next_pred|pending)"
9842 );
9843 }
9844 let mut prime_logits;
9845 let mut prompt_h: Option<CudaSlice<f32>> = None;
9846 let t_prime = std::time::Instant::now();
9847 let batched_prime = !continuation
9848 && prompt.len() >= crate::hybrid_forward::PRIME_MIN_T
9849 && std::env::var("MEMRA_PRIME_TOKENWISE").is_err()
9850 && !e.frozen_cpu_experts_prefer_tokenwise_prime();
9851 let prime_split = prime_split.filter(|&split| split > 0 && split < prompt.len());
9852 if prime_split.is_some() && continuation {
9853 return Err("spec prime split requires a non-empty prime".into());
9854 }
9855 // STABLE-BOUNDARY TURN CHECKPOINT stop (lane/frspec-multiturn-cache, 2026-08-21):
9856 // the worker's `ckpt_at` request, ABSOLUTE -> prompt-relative. On WARM bursts
9857 // (base != 0, an affinity-rewound or pool-resumed session priming its own delta)
9858 // this is the only stop; on COLD bursts it usually coincides with `prime_split`
9859 // (both are the plain tier's stable pre-generation boundary). A boundary the prime
9860 // cannot honor (outside this prime's range) silently drops the capture — the
9861 // turn_ckpt convention: the next turn re-primes in full, never a wrong resume.
9862 let ckpt_rel = if continuation {
9863 None
9864 } else {
9865 ckpt_req
9866 .and_then(|abs| abs.checked_sub(base))
9867 .filter(|&r| r > 0 && r < prompt.len())
9868 };
9869 // Prime stops, ordered: each is a boundary the prime halts at so the in-place GDN
9870 // conv/ssm state can be snapshotted there (the only moment it exists). One stop =
9871 // the legacy single-split program, byte-for-byte.
9872 let mut stops: Vec<usize> = Vec::new();
9873 for b in [prime_split, ckpt_rel].into_iter().flatten() {
9874 if !stops.contains(&b) {
9875 stops.push(b);
9876 }
9877 }
9878 stops.sort_unstable();
9879 // Captured at the ckpt stop, installed into the session slot post-prime (replacing
9880 // the legacy prompt-end capture). Some(None) = capture attempted and failed -> the
9881 // slot is cleared (a stale checkpoint would rewind to the WRONG boundary).
9882 let mut ckpt_early: Option<Option<SpecCheckpoint>> = None;
9883 if continuation {
9884 prime_logits = Vec::new();
9885 } else if !stops.is_empty() {
9886 if let Some(&first) = stops.first() {
9887 if prime_split == Some(first) && first < crate::hybrid_forward::PRIME_MIN_T {
9888 return Err(format!(
9889 "spec prime split {first} is below PRIME_MIN_T {}",
9890 crate::hybrid_forward::PRIME_MIN_T,
9891 )
9892 .into());
9893 }
9894 }
9895 // Mirror the plain worker's boundary stops exactly. Each segment is a
9896 // request-level prime (`queued_after` keeps Step35 arm selection independent of
9897 // the stops — tick-seg law); a segment below PRIME_MIN_T (and the final tail
9898 // under MEMRA_PRIME_TOKENWISE) takes the same eager tokenwise continuation as
9899 // prefill_tick. Retain every hidden row so the draft scratch fill remains one
9900 // coherent prompt.
9901 let mut h_all = e.uninit(prompt.len() * n_embd)?;
9902 prime_logits = Vec::new();
9903 let mut prev = 0usize;
9904 for seg_end in stops.iter().copied().chain(std::iter::once(prompt.len())) {
9905 if seg_end <= prev {
9906 continue;
9907 }
9908 let seg = &prompt[prev..seg_end];
9909 let is_final = seg_end == prompt.len();
9910 let batched_seg = seg.len() >= crate::hybrid_forward::PRIME_MIN_T
9911 && (!is_final
9912 || (std::env::var("MEMRA_PRIME_TOKENWISE").is_err()
9913 && !e.frozen_cpu_experts_prefer_tokenwise_prime()));
9914 if batched_seg {
9915 let (l, _, h_seg) =
9916 self.prime_cache(e, seg, &mut *cache, prompt.len() - seg_end)?;
9917 e.copy_into(&mut h_all, prev * n_embd, &h_seg, seg.len() * n_embd)?;
9918 prime_logits = l;
9919 } else {
9920 for (i, &tok) in seg.iter().enumerate() {
9921 let (l, h) = self.decode_step_h(e, tok, &mut *cache)?;
9922 e.copy_into(&mut h_all, (prev + i) * n_embd, &h, n_embd)?;
9923 prime_logits = l;
9924 }
9925 }
9926 prev = seg_end;
9927 if is_final {
9928 break;
9929 }
9930 debug_assert_eq!(cache.pos, base + seg_end, "prime stop landed off boundary");
9931 // PREFIX-CACHE BOUNDARY CAPTURE (lane/spec-prefix-cache): the GDN conv/ssm
9932 // states are about to be advanced in place by the next segment, so this is
9933 // the ONLY moment the boundary's recurrent state exists. Capture iff the
9934 // worker requested exactly this stop (cold sessions only — `capture_at` is
9935 // never armed warm). A failed snapshot is silent (turn_ckpt convention) —
9936 // publication is an optimization, never a correctness dependency.
9937 if base == 0 {
9938 if let Some((requested, slot)) = sess_capture.as_mut() {
9939 // Publish at the requested miss-LCP stop (the shared-prefix class)
9940 // AND at the stable-boundary stop (the next-turn re-render class,
9941 // lane/frspec-multiturn-cache) — the same boundary set the plain
9942 // prefill tick learns. Without the second entry, the turn after a
9943 // cold re-park could only hit the OLDER lcp entry (the measured
9944 // one-turn transient: t3 restored 607 of 24122 while the plain arm
9945 // rewound to 15222). Dedupe is the worker sweep's has_key.
9946 if *requested == Some(seg_end) || ckpt_rel == Some(seg_end) {
9947 if let Ok(snap) = cache.snapshot(e) {
9948 slot.push(SpecBoundaryCapture {
9949 snap,
9950 pos: seg_end,
9951 logits: prime_logits.clone(),
9952 // rows [0..seg_end) of h_all are primed — the following
9953 // segments append, never overwrite.
9954 last_h: capture_boundary_hidden(e, &h_all, seg_end, n_embd),
9955 });
9956 }
9957 }
9958 }
9959 }
9960 // SESSION-AFFINITY TURN CHECKPOINT at the STABLE boundary (see `ckpt_at`):
9961 // same snapshot mechanics, installed post-prime in place of the prompt-end
9962 // capture the re-render class always diverged below.
9963 if ckpt_rel == Some(seg_end) {
9964 let anchor: Result<CudaSlice<f32>, Box<dyn std::error::Error>> =
9965 e.uninit(n_embd).and_then(|mut a| {
9966 e.copy_view_into(
9967 &mut a,
9968 0,
9969 &h_all.slice((seg_end - 1) * n_embd..seg_end * n_embd),
9970 n_embd,
9971 )?;
9972 Ok(a)
9973 });
9974 ckpt_early = Some(match (cache.snapshot(e), anchor) {
9975 (Ok(snap), Ok(last_h)) => Some(SpecCheckpoint {
9976 snap,
9977 pos: base + seg_end,
9978 last_h,
9979 }),
9980 _ => None,
9981 });
9982 }
9983 }
9984 if std::env::var("MEMRA_SPEC_STATS").as_deref() == Ok("1") {
9985 eprintln!(
9986 "[spec-prime] stops={stops:?} tail={}",
9987 prompt.len() - stops.last().copied().unwrap_or(0)
9988 );
9989 }
9990 prompt_h = Some(h_all);
9991 } else if batched_prime {
9992 let (l, _h_seed, hiddens) = self.prime_cache(e, prompt, &mut *cache, 0)?;
9993 prime_logits = l;
9994 prompt_h = Some(hiddens);
9995 } else {
9996 prime_logits = Vec::new();
9997 prompt_h = Some(e.uninit(prompt.len() * n_embd)?);
9998 for (i, &tok) in prompt.iter().enumerate() {
9999 let (l, h) = self.spec_target_step_h(e, tok, &mut *cache)?;
10000 if let Some(ph) = prompt_h.as_mut() {
10001 e.copy_into(ph, i * n_embd, &h, n_embd)?;
10002 }
10003 prime_logits = l;
10004 }
10005 }
10006 e.stream().synchronize()?;
10007 // PREFIX-CACHE SEED CAPTURE (lane/spec-prefix-cache): boundary == prompt end (the seed
10008 // case — no shared-prefix split, publish the whole prompt). The prime just finished, so
10009 // cache.pos == base + prompt.len() and the recurrent state IS the boundary state;
10010 // prime_logits are the boundary logits. Cold sessions only (base == 0) — same law as
10011 // prime_split. The mid-prompt capture above already consumed the request if it matched.
10012 if !continuation && base == 0 {
10013 if let Some((requested, slot)) = sess_capture.as_mut() {
10014 if *requested == Some(prompt.len()) && slot.is_empty() {
10015 debug_assert_eq!(cache.pos, prompt.len(), "seed capture off prompt end");
10016 if let Ok(snap) = cache.snapshot(e) {
10017 slot.push(SpecBoundaryCapture {
10018 snap,
10019 pos: prompt.len(),
10020 logits: prime_logits.clone(),
10021 last_h: prompt_h
10022 .as_ref()
10023 .map(|ph| capture_boundary_hidden(e, ph, prompt.len(), n_embd))
10024 .unwrap_or_default(),
10025 });
10026 }
10027 }
10028 }
10029 }
10030 // Harness timing contract (see crate::PRIME_NANOS): gen-only throughput without the
10031 // prime-subtraction hack.
10032 crate::PRIME_NANOS.store(
10033 t_prime.elapsed().as_nanos() as u64,
10034 std::sync::atomic::Ordering::Relaxed,
10035 );
10036
10037 let (embd_qt, embd_rb) = self.embd.qt_and_row_bytes(n_embd);
10038 // Resident table is fastest when it fits. Large spill deployments can preserve that HBM
10039 // for expert-cache slots and gather only the exact rows needed by MTP/verify from host.
10040 let host_embd = spec_host_embd();
10041 let embd_gpu = if host_embd {
10042 None
10043 } else {
10044 Some(
10045 self.embd_gpu
10046 .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload")),
10047 )
10048 };
10049 let embd_dev = embd_gpu.map(|g| (g, embd_qt, embd_rb));
10050 if host_embd {
10051 eprintln!(
10052 "[spec] host-row embedding: {} bytes kept off HBM",
10053 self.embd.raw.len()
10054 );
10055 }
10056 let mut out: Vec<u32> = Vec::with_capacity(max_new);
10057 let mut total_drafted = 0usize;
10058 let mut total_accepted = 0usize;
10059
10060 // --- SAMPLER FIRST (lane/sampled-spec-quality, 2026-08-19) ---
10061 // The sampler config, the session's Philox counters and the penalty window are parsed
10062 // HERE, above the boundary-token selection, because the boundary token must be drawn
10063 // from the sampler the request asked for. Pre-lane this block sat ~50 lines BELOW the
10064 // selection, which is the whole mechanical reason the boundary token was an argmax:
10065 // the sampler state was not in scope yet. Nothing here depends on the round loop, so
10066 // moving it up is a pure reordering for greedy (`sampled == false` ⇒ every branch
10067 // below takes the argmax path it always took).
10068 // --- SAMPLED SPEC (MEMRA_SPEC_TEMP>0, research/sampled-spec-impl-map.md): rejection-
10069 // sampling verify (Leviathan/Chen) — accept draft x at u < p(x)/q(x), resample from
10070 // norm(max(0,p-q)) on reject, bonus sampled from p on full accept. Counter-based Philox
10071 // everywhere (seed, event) -> reproducible. temp==0/unset = the greedy path, untouched.
10072 let sp = sampling.unwrap_or_else(|| SpecSampling {
10073 temp: std::env::var("MEMRA_SPEC_TEMP")
10074 .ok()
10075 .and_then(|v| v.parse().ok())
10076 .unwrap_or(0.0),
10077 seed: std::env::var("MEMRA_SEED")
10078 .ok()
10079 .and_then(|v| v.parse().ok())
10080 .unwrap_or(42),
10081 top_k: std::env::var("MEMRA_TOP_K")
10082 .ok()
10083 .and_then(|v| v.parse().ok())
10084 .unwrap_or(0),
10085 top_p: std::env::var("MEMRA_TOP_P")
10086 .ok()
10087 .and_then(|v| v.parse().ok())
10088 .unwrap_or(1.0),
10089 min_p: std::env::var("MEMRA_MIN_P")
10090 .ok()
10091 .and_then(|v| v.parse().ok())
10092 .unwrap_or(0.0),
10093 penalty_last_n: std::env::var("MEMRA_PENALTY_LAST_N")
10094 .ok()
10095 .and_then(|v| v.parse().ok())
10096 .unwrap_or(0),
10097 penalty_repeat: std::env::var("MEMRA_PENALTY_REPEAT")
10098 .ok()
10099 .and_then(|v| v.parse().ok())
10100 .unwrap_or(1.0),
10101 penalty_freq: std::env::var("MEMRA_PENALTY_FREQ")
10102 .ok()
10103 .and_then(|v| v.parse().ok())
10104 .unwrap_or(0.0),
10105 penalty_present: std::env::var("MEMRA_PENALTY_PRESENT")
10106 .ok()
10107 .and_then(|v| v.parse().ok())
10108 .unwrap_or(0.0),
10109 });
10110 let (sp_temp, sp_seed) = (sp.temp, sp.seed);
10111 let sampled = sp_temp > 0.0;
10112 // Counters resume from the session (burst continuity: randomness must never repeat
10113 // across generate_spec_session calls); one-shot callers start at (0,0). Read through
10114 // sess_tail — `sess` was take()n into it above, so sess.as_ref() here is always None.
10115 let mut sctr: u32 = sess_tail.as_ref().map(|(_, _, _, s, _)| **s).unwrap_or(0);
10116 let mut uctr: u32 = sess_tail.as_ref().map(|(_, _, _, _, u)| **u).unwrap_or(0);
10117 // Penalties (v2.1): applied to COPIES of q rows and p columns symmetrically (exactness
10118 // for the penalized+filtered target). History = generated tokens, host-tracked window.
10119 let pen_on = sampled
10120 && sp.penalty_last_n > 0
10121 && (sp.penalty_repeat != 1.0 || sp.penalty_freq != 0.0 || sp.penalty_present != 0.0);
10122 // SESSION-SPANNING PENALTY WINDOW (Item 2). Pre-lane this was
10123 // `prompt.iter().rev().take(64).rev()` — the BURST's suffix slice — so a continuation
10124 // burst (the majority of a stream's tokens, and ALL of a converted cache hit's) started
10125 // with an EMPTY penalty history and the client's repetition/frequency/presence penalties
10126 // silently reset at every burst boundary. The window now spans `committed ++ prompt`,
10127 // which is what the API contract says and what the plain sampler's own `history` does.
10128 // Byte-identical to the pre-lane seed for a cold turn-1 burst at the default window.
10129 let mut pen_hist: Vec<u32> = if pen_on {
10130 let sess_hist: &[u32] = if spec_pen_session_on() {
10131 sess_tail
10132 .as_ref()
10133 .map(|(c, ..)| c.as_slice())
10134 .unwrap_or(&[])
10135 } else {
10136 &[] // MEMRA_SPEC_PEN_SESSION=0: pre-lane burst-local window
10137 };
10138 pen_window_seed(sess_hist, prompt, sp.penalty_last_n)
10139 } else {
10140 Vec::new()
10141 };
10142 // First generated token = the BOUNDARY token: greedy takes the argmax of the prompt's
10143 // last logits (== greedy's first token, byte-contract); SAMPLED draws it from the
10144 // request's own filtered/penalized target through the session's Philox stream
10145 // (`sample_boundary_token`, lane/sampled-spec-quality Item 1 — pre-lane this was an
10146 // argmax in both regimes, so ~1 token per burst of a sampled stream was greedy).
10147 // Emit it, then FEED it to establish the loop invariant below.
10148 // PENDING-CARRY: the carried bonus was already emitted by the LAST burst — it becomes
10149 // last_token WITHOUT re-emission, and round 0 consumes it as pending (no init feed).
10150 // CONSTRAINED entry rules: the first emitted token is the MASKED argmax of the
10151 // prompt's last logits (plain constrained-greedy identity); a continuation without
10152 // a carried pending would emit an UNMASKED stashed next_pred — refused loudly (the
10153 // worker never resumes constrained sessions from the pool, so this cannot fire).
10154 if let Some(c) = constraint.as_deref_mut() {
10155 if continuation && carried_pending.is_none() {
10156 return Err("constrained spec continuation requires a carried pending \
10157 (pool resume is unconstrained-only)"
10158 .into());
10159 }
10160 if !continuation {
10161 c.mask_logits(&mut prime_logits)
10162 .map_err(|e2| format!("constraint: {e2}"))?;
10163 }
10164 }
10165 let mut last_token = if let Some(b) = carried_pending {
10166 b
10167 } else if continuation {
10168 // A continuation's boundary token was DRAWN by the burst that stashed it (the
10169 // session tail below), or by `spec_session_from_restored` for a converted
10170 // prefix-cache hit — in both cases from the correct logits row with this same
10171 // session's Philox stream, which is why it can be consumed here as-is.
10172 sess_tail.as_ref().unwrap().2.unwrap()
10173 } else if sampled && constraint.is_none() && spec_sampled_boundary_on() {
10174 sample_boundary_token(e, &prime_logits, &sp, &pen_hist, &mut sctr, "cold-prime")?
10175 } else {
10176 // greedy (byte contract), the rollback door, or constrained (masked-argmax
10177 // identity — the worker routes sampled+constrained to the plain path, and this
10178 // function refuses the combination outright above).
10179 argmax(&prime_logits) as u32
10180 };
10181 if pen_on {
10182 // The boundary token is a GENERATED token: the plain sampler `accept()`s every
10183 // emitted token into its penalty history, and pre-lane the burst's first token
10184 // was invisible to penalties forever (never pushed, and never in `committed`
10185 // until this burst's tail). Covers the carry/continuation seeds too — neither is
10186 // in `committed` yet.
10187 pen_hist.push(last_token);
10188 }
10189 if carried_pending.is_none() {
10190 out.push(last_token);
10191 // grammar advances with every emitted token (carried pendings were consumed
10192 // by the burst that emitted them).
10193 if let Some(c) = constraint.as_deref_mut() {
10194 c.consume(last_token)
10195 .map_err(|e2| format!("constraint: {e2}"))?;
10196 }
10197 }
10198 if continuation {
10199 // draft-KV invariant: entries [0..base) are the session's exact fills; truncate any
10200 // overhang so the chain's first append lands at slot base (== committed.len()).
10201 scratch.set_len(e, base)?;
10202 }
10203 // sse-cadence: hand the caller every not-yet-flushed token (disjoint in-order slices
10204 // concatenating to the full `out`). Called after the prime's first token and after each
10205 // round commit — emission timing only, token bytes untouched. The slice may be EMPTY
10206 // (poll-only boundary: zero-round folds commit nothing new); returns the caller's
10207 // continue-verdict (admission yield, 2026-08-06) — false ends the burst at this round.
10208 fn flush_commit(
10209 cb: &mut Option<&mut dyn FnMut(&[u32]) -> bool>,
10210 out: &[u32],
10211 flushed: &mut usize,
10212 ) -> bool {
10213 if let Some(f) = cb.as_mut() {
10214 let keep = f(&out[*flushed..]);
10215 *flushed = out.len();
10216 keep
10217 } else {
10218 true
10219 }
10220 }
10221 keep_going = flush_commit(&mut on_commit, &out, &mut flushed);
10222 // INVARIANT at loop top: `last_token` is the most-recently-committed/emitted token, its
10223 // KV+recur state IS in `cache` (cache.pos = position right AFTER last_token), `last_pred`
10224 // is the greedy ARGMAX of the logits that predict the token FOLLOWING last_token, and
10225 // `h_seed` = last_token's pre-output_norm hidden. Establish it by feeding last_token once
10226 // (mirrors plain greedy). DEVICE-ARGMAX lever: the accept walk only ever consumes the
10227 // argmax of those logits — never the full vector — so a host u32 replaces the Vec<f32>.
10228 // Trimmed heads: q lives on the trimmed vocab; accept gathers use the TRIMMED index and
10229 // the residual scatters q into target-id space (q=-inf off-trim — the head cannot propose
10230 // those, so their residual mass is p(x), correct by construction).
10231 let d2t_dev: Option<CudaSlice<u32>> = if sampled || crate::spec::spec_stream() {
10232 match &mtp.d2t {
10233 Some(map) => Some(e.htod_u32_v(map)?),
10234 None => None,
10235 }
10236 } else {
10237 None
10238 };
10239 let mut q_full_buf: Option<CudaSlice<f32>> = None;
10240 // host Philox4x32-10 accept-test uniforms: module fn `host_u01` (shared with the
10241 // dspark sampled-admission walk); byte-identical to the closure it replaces.
10242 let mut draft_logits: Vec<CudaSlice<f32>> = Vec::new(); // retained head logits (q), per slot
10243 let mut draft_stats: Vec<(f32, f32, f32)> = Vec::new(); // (row_max, th_e, z_e) per slot
10244 let mut perturb_buf: Option<CudaSlice<f32>> = None; // gumbel scratch (max(n_vocab,d_vocab))
10245 let mut sample_tok = e.alloc_u32_zeroed(1)?; // residual/bonus sample out
10246 let mut col_buf: Option<CudaSlice<f32>> = None; // materialized verify column
10247 let mut pen_hist_d: Option<CudaSlice<u32>> = None;
10248 let mut pcol_buf: Option<CudaSlice<f32>> = None; // penalized p-column scratch
10249 // MEMRA_SPEC_SETUP_TRACE=1 (diagnostics): per-call wall decomposition of the burst
10250 // SETUP + TAIL segments (the round loop's internals are MEMRA_SPEC_PHASE's job) —
10251 // built to pin the serve per-burst fixed cost (research/spec-serving-20260801).
10252 let setup_trace = std::env::var("MEMRA_SPEC_SETUP_TRACE").as_deref() == Ok("1");
10253 let t_ent = std::time::Instant::now();
10254
10255 // SESSION-AFFINITY TURN CHECKPOINT (lane/session-affinity, 2026-08-05): capture the
10256 // PROMPT-END boundary state so a LATER turn can rewind here and re-prime only its own
10257 // delta instead of the whole conversation. See `SpecCheckpoint` for why this boundary is
10258 // the one that matters (a history-rewriting client mutates what the session GENERATED,
10259 // so the next turn's prompt agrees with this one up to exactly here).
10260 //
10261 // WHERE — AND WHY THIS EXACT LINE. Right after the trunk prime, BEFORE the init feed
10262 // (`decode_step_h(last_token)`) and before round 0: the last instant at which the caches
10263 // hold exactly `base + prompt.len()` rows and nothing generated.
10264 //
10265 // This was WRONG in the first cut of this lane: the capture sat after the draft-KV fill,
10266 // which is also after the init feed, so `cache.pos` was `base + prompt.len() + 1` — the
10267 // boundary included the FIRST GENERATED TOKEN. That token is the first thing inside the
10268 // `<think>` block the client strips, so every later turn's diff diverged exactly one
10269 // token below the checkpoint and affinity declined 100% of the time. Measured on the
10270 // owner regime: "history diverged at 12233 of checkpoint 12234". The off-by-one made the
10271 // whole mechanism inert while looking, from the outside, like a working
10272 // correctness-declines-safely path — hence the decline log carries the offsets.
10273 //
10274 // The full-attn planes are `len`-truncatable so the snapshot copies only the GDN conv/ssm
10275 // state (the reason a spec session could not rewind before). The draft scratch needs no
10276 // copy: rows below the boundary are rewritten by the next turn's own fill.
10277 //
10278 // WHEN: non-empty prime only. An empty-suffix continuation burst adds no prompt boundary
10279 // (its "prompt end" IS the previous checkpoint's, already held), so it keeps the existing
10280 // checkpoint rather than replacing it with a strictly worse one.
10281 //
10282 // FAILURE IS SILENT BY DESIGN: on a VRAM-tight rig the snapshot alloc can fail. That
10283 // costs the NEXT turn its rewind (it re-primes fully, today's behavior) and must never
10284 // fail the burst that is already running — so the error is swallowed, loud only under
10285 // MEMRA_DEBUG_SPEC.
10286 //
10287 // STABLE-BOUNDARY OVERRIDE (lane/frspec-multiturn-cache, 2026-08-21): the prompt-end
10288 // posture above was DISPROVED for the think-posture template class — the prompt's own
10289 // tail is the live generation header (`<|im_start|>assistant\n<think>\n`) that the
10290 // next turn's re-render replaces, so the diff diverged a couple tokens BELOW the
10291 // checkpoint and affinity declined 100% of multi-turn agent traffic (the same class
10292 // the plain tier fixed on 2026-08-09 via `plain_checkpoint_boundary`; the port to the
10293 // spec tier is this lane). When the worker armed `ckpt_at`, the capture happened at
10294 // that stop inside the prime above (`ckpt_early`) and is installed here instead;
10295 // capture-attempted-but-failed clears the slot exactly like the legacy arm.
10296 if let Some(slot) = sess_ckpt_slot {
10297 if let Some(early) = ckpt_early {
10298 if early.is_none() && std::env::var("MEMRA_DEBUG_SPEC").is_ok() {
10299 eprintln!(
10300 "[spec] stable-boundary turn checkpoint skipped; \
10301 next turn re-primes in full"
10302 );
10303 }
10304 *slot = early;
10305 } else if !continuation {
10306 let pos = cache.pos;
10307 debug_assert_eq!(
10308 pos,
10309 base + prompt.len(),
10310 "turn checkpoint must sit at the prompt end, before the init feed"
10311 );
10312 let anchor: Result<CudaSlice<f32>, Box<dyn std::error::Error>> =
10313 if let Some(ph) = &prompt_h {
10314 // hidden of the LAST primed row = the predecessor anchor at this
10315 // boundary (exactly what a fresh prime of committed[..pos] leaves in
10316 // last_h, and what the next prime's fill reads for its first row).
10317 let np = prompt.len();
10318 e.uninit(n_embd).and_then(|mut a| {
10319 e.copy_view_into(
10320 &mut a,
10321 0,
10322 &ph.slice((np - 1) * n_embd..np * n_embd),
10323 n_embd,
10324 )?;
10325 Ok(a)
10326 })
10327 } else {
10328 Err("no prompt hiddens".into())
10329 };
10330 match (cache.snapshot(e), anchor) {
10331 (Ok(snap), Ok(last_h)) => {
10332 *slot = Some(SpecCheckpoint { snap, pos, last_h });
10333 }
10334 (s, a) => {
10335 *slot = None; // a stale checkpoint would rewind to the WRONG boundary
10336 if std::env::var("MEMRA_DEBUG_SPEC").is_ok() {
10337 let err = s
10338 .err()
10339 .map(|e| e.to_string())
10340 .or_else(|| a.err().map(|e| e.to_string()))
10341 .unwrap_or_default();
10342 eprintln!(
10343 "[spec] turn checkpoint skipped ({err}); \
10344 next turn re-primes in full"
10345 );
10346 }
10347 }
10348 }
10349 }
10350 }
10351 // INIT FEED — skipped on a pending carry: last_token (the carried bonus) is NOT in the
10352 // caches and must NOT be fed solo; round 0's batched verify commits it as col 0. Its
10353 // seed/anchor hidden is the carried last_h (copied below); last_pred is dead in the
10354 // pending path (t_pred reads verify col 0 — the accept walk overwrites it).
10355 let mut last_pred = 0u32;
10356 let mut last_col_logits: Option<CudaSlice<f32>> = None;
10357 // CONSTRAINED: the init feed's logits back the (n_acc==0, base==0) masked-argmax
10358 // recompute in the grammar-truncation walk — retained host-side, round 0 only.
10359 let mut init_logits_host: Option<Vec<f32>> = None;
10360 let h_seed0: CudaSlice<f32> = if carried_pending.is_none() {
10361 let (init_logits, h) = self.spec_target_step_h(e, last_token, &mut *cache)?;
10362 last_pred = argmax(&init_logits) as u32;
10363 if constraint.is_some() {
10364 init_logits_host = Some(init_logits.clone());
10365 }
10366 // sampled mode: p-distribution after last_token, for the j==0/base==0 accept test.
10367 if sampled {
10368 last_col_logits = Some(e.htod(&init_logits)?);
10369 }
10370 h
10371 } else {
10372 // predecessor-row anchor: hidden of the last COMMITTED row (the carry contract).
10373 let lh = sess_tail
10374 .as_ref()
10375 .unwrap()
10376 .1
10377 .as_ref()
10378 .expect("pending carry requires last_h");
10379 e.clone_dtod(lh)?
10380 };
10381 let t_init = t_ent.elapsed();
10382 let mut last_col_stats: Option<(f32, f32, f32)> = None;
10383 // PERSISTENT h_seed buffer (allocated BEFORE any graph capture so no captured scratch can
10384 // alias it): every path that updates the round seed copies INTO it — no per-round allocs,
10385 // stable pointer for the graph-draft round-start copy.
10386 let mut h_seed_buf = e.clone_dtod(&h_seed0)?;
10387 // Predecessor-pairing trackers: `fill_prev` = trunk hidden AT the last COMMITTED row (the
10388 // predecessor of the next verify's col 0 — the reference's carried pending-h analogue;
10389 // also the predecessor-row hidden for the round-0 legacy-replay seed). At round 0 that
10390 // row is last_token's own (h_seed0). The chain step-0 seed under the pairing default =
10391 // hidden of the row BEFORE last_token = the prompt's last row at round 0 (h_seed_buf
10392 // overwritten below).
10393 let mut fill_prev = e.clone_dtod(&h_seed0)?;
10394 {
10395 if let Some(ph) = &prompt_h {
10396 let np = prompt.len();
10397 e.copy_view_into(
10398 &mut h_seed_buf,
10399 0,
10400 &ph.slice((np - 1) * n_embd..np * n_embd),
10401 n_embd,
10402 )?;
10403 } else if continuation {
10404 if let Some((_, lh, _, _, _)) = sess_tail.as_ref() {
10405 if let Some(lh) = lh.as_ref() {
10406 e.copy_into(&mut h_seed_buf, 0, lh, n_embd)?;
10407 }
10408 }
10409 }
10410 }
10411 // Persistent device prediction slots for the accept walk (max k+1 verify columns).
10412 let mut preds_d = e.alloc_u32_zeroed(k + 2)?;
10413
10414 let debug_spec = std::env::var("MEMRA_DEBUG_SPEC").is_ok();
10415 let fork_mode = OptiForkGateMode::configured();
10416 // MEMRA_SPEC_STATS=1: per-slot accept histogram + draft-length histogram, printed once at
10417 // the end. Metric normalization vs the reference engine: BOTH engines count
10418 // accepted/drafted where the chain stopped at p-min and the sub-threshold token is
10419 // discarded uncounted — per-slot decay + chain-length mix are the extra dimensions.
10420 let spec_stats = std::env::var("MEMRA_SPEC_STATS").is_ok();
10421 let mut st_drafted = vec![0usize; k];
10422 let mut st_accepted = vec![0usize; k];
10423 let mut st_len_hist = vec![0usize; k + 1];
10424 let mut st_full = 0usize;
10425 // P-MIN CONFIDENCE GATE (MEMRA_SPEC_PMIN, the serve script's --spec-draft-p-min mechanism):
10426 // stop the draft chain early when the head's softmax confidence in its own pick drops
10427 // below p_min. Hoisted above the loop: the graph capture bakes the prob kernels iff on.
10428 static PMIN: std::sync::OnceLock<f32> = std::sync::OnceLock::new();
10429 let p_min = *PMIN.get_or_init(|| {
10430 std::env::var("MEMRA_SPEC_PMIN")
10431 .ok()
10432 .and_then(|v| v.parse().ok())
10433 .unwrap_or(0.0)
10434 });
10435 // ZERO-DRAFT ROUNDS (MEMRA_SPEC_PMIN0=1, vendored from llama.cpp's draft gating): let the
10436 // p-min gate apply at j==0 too, so a low-confidence round drafts NOTHING and the verify
10437 // batch is just the pending bonus (m=1 = a plain decode step). llama's 35B win rides
10438 // exactly this — draft acceptance 76% at mean len 2.5 because unpredictable stretches
10439 // never pay draft+verify overhead. Only legal when a pending bonus exists (an empty
10440 // verify batch is not); the j==0 exemption stays for pending-less rounds.
10441 let pmin0 = std::env::var("MEMRA_SPEC_PMIN0")
10442 .map(|v| v == "1")
10443 .unwrap_or(false);
10444
10445 // --- GRAPH DRAFT setup: persistent I/O buffers + ONE capture (2 warmups inside). The
10446 // warmups mutate scratch len_d / pos / tok / seed — all reset at every round start, so the
10447 // only restore needed is the scratch counter. Capture failure (e.g. a non-capturable
10448 // cuBLAS path in an exotic head) falls back to the eager draft chain.
10449 // PER-SESSION PERSISTENCE (2026-08-01): session calls reuse the DraftGraphCtx parked on
10450 // the SpecSession — the capture (2 warmup head forwards + instantiate) ran ONCE at the
10451 // session's first burst, not per burst (measured ~16ms/burst fixed cost on H100 q27,
10452 // research/spec-serving-20260801). Reuse is pointer-exact: the graph bakes the session's
10453 // own scratch KV (never realloc'd), the model's resident embedding, the OnceLock p_min,
10454 // and the g_* buffers carried in the ctx — replay dispatch is identical to a fresh
10455 // capture, so draft tokens are bit-identical (drafts never decide exactness anyway; the
10456 // verify arbitrates). Single-shot calls (sess=None) build a fresh ctx and drop it.
10457 let mut dctx: DraftGraphCtx = match sess_draft_slot.as_mut().and_then(|s| s.take()) {
10458 Some(c) => c,
10459 None => DraftGraphCtx::new(e, n_embd, if sampled { d_vocab } else { 1 })?,
10460 };
10461 // A session that ran greedy bursts first sized g_q/g_perturb at 1; a sampled resume
10462 // needs d_vocab. Realloc is legal exactly while graph_s is None (nothing baked them).
10463 if sampled && dctx.g_q.len() < d_vocab {
10464 dctx.g_q = e.zeros(d_vocab)?;
10465 dctx.g_perturb = e.zeros(d_vocab)?;
10466 }
10467 // DRAFT-SIDE GRAMMAR MASK (lane/draft-mask, 2026-08-04): the drafter samples the
10468 // grammar's legal set, so proposals are legal BY CONSTRUCTION and the verify-side
10469 // truncation (the correctness backstop) stops cutting every tight-schema round.
10470 // The mask is one node inside the captured draft chain — presence is a CAPTURE-TIME
10471 // shape, so a parked graph of the other shape is dropped and recaptured.
10472 let dmask_on = constraint
10473 .as_deref()
10474 .is_some_and(|c| c.draft_mask_enabled());
10475 let dmask_words = if dmask_on { d_vocab.div_ceil(32) } else { 0 };
10476 if dmask_on && dctx.g_dmask.len() < dmask_words {
10477 dctx.g_dmask = e.alloc_u32_zeroed(dmask_words)?;
10478 dctx.graph = None; // the old capture baked the old (or no) mask pointer
10479 dctx.failed.clear_greedy();
10480 dctx.keeper.clear();
10481 }
10482 if dctx.graph.is_some() && dctx.graph_masked != dmask_on {
10483 dctx.graph = None;
10484 dctx.failed.clear_greedy();
10485 dctx.keeper.clear();
10486 }
10487 if graph_draft && !sampled && dctx.graph.is_none() && !dctx.failed.greedy_failed() {
10488 let DraftGraphCtx {
10489 g_tok,
10490 g_pos,
10491 g_seed,
10492 g_p,
10493 g_dmask,
10494 ..
10495 } = &mut dctx;
10496 // capture-time contents: ALL-ONES (ban nothing). A replay only ever runs after the
10497 // host uploads the position's real words, so the warmups stay grammar-free.
10498 if dmask_on {
10499 e.htod_u32_into(g_dmask, &vec![u32::MAX; dmask_words])?;
10500 }
10501 let g_dmask_ro: &CudaSlice<u32> = &*g_dmask;
10502 // CAPTURE-RETAIN (#68 fix): the warmup transients' pool addresses are baked into the
10503 // captured graph; the keeper pins them for the graph's lifetime. capture_graph (non-
10504 // retained) freed them at exit — safe for one-shot generate_spec (nothing else touches
10505 // the pool between replays) but WRONG for sessions: burst-boundary prime/fill/commit
10506 // passes (and, in serve, other sessions) recycle those addresses and the replay then
10507 // clobbers live buffers — the ST serve-spec corruption (research/serve-st-20260803).
10508 let cap_res = e.capture_graph_retained(|e| {
10509 self.mtp_head_forward_cap(
10510 e,
10511 mtp,
10512 g_tok,
10513 g_pos,
10514 g_seed,
10515 g_p,
10516 &mut *scratch,
10517 p_min > 0.0 || fork_mode == OptiForkGateMode::Controller,
10518 true,
10519 embd_gpu.expect("graph draft requires resident embedding"),
10520 embd_qt,
10521 embd_rb,
10522 d_vocab,
10523 None,
10524 None,
10525 if dmask_on {
10526 Some((g_dmask_ro, dmask_words))
10527 } else {
10528 None
10529 },
10530 )
10531 });
10532 match cap_res {
10533 Ok((g, keep)) => {
10534 scratch.set_len(e, base)?;
10535 dctx.graph = Some(g);
10536 dctx.graph_masked = dmask_on;
10537 dctx.keeper = keep;
10538 }
10539 Err(err) => {
10540 scratch.set_len(e, base)?;
10541 // LOUD flip (audit Q2): a dropped draft graph is a coverage loss, never
10542 // silent. Once per flip — mark returns None on an already-failed ctx.
10543 if let Some(line) = dctx.failed.mark_greedy(&err.to_string()) {
10544 eprintln!("{line}");
10545 }
10546 }
10547 }
10548 }
10549 // --- SAMPLED GRAPH DRAFT setup (step 3 of the sampled-spec arc): a SECOND capture, own
10550 // graph object, built only when sampled && graph-eligible — the greedy capture above is
10551 // untouched (and skipped when sampled: its graph would never be launched). Same head
10552 // forward, but the in-graph argmax reads GUMBEL-PERTURBED logits; the Philox event
10553 // counter lives in the persistent device g_ctr (bumped in-graph, host-seeded from sctr
10554 // once per round); the raw head logits land in the persistent g_q for the host's
10555 // per-replay async D2D into the round's q slot (q_slots, K x d_vocab, allocated once).
10556 // seed/temp are capture-time constants — baked into graph_s, so a pool-resumed request
10557 // with a different (seed, temp, k) drops the parked sampled graph and recaptures.
10558 // COST OF THE FRESH-SEED SERVE DEFAULT (dogfood F4, 2026-08-04): omitting `seed` on a
10559 // serve request now draws fresh per-request entropy (it used to default to a pinned 0),
10560 // so a seed-omitting request that RESUMES a parked spec session finds an s_key baked
10561 // with the PREVIOUS request's seed and pays one recapture. Bounded, and it does not
10562 // reopen the ~16ms/burst regression the persistent ctx exists to fix: a session's seed
10563 // is fixed for its whole lifetime (worker.rs reads s.sampler.seed() per burst), so
10564 // this compare misses at most ONCE per resumed request — the first burst recaptures
10565 // and every later burst in that request replays. A client that wants the parked graph
10566 // AND reproducibility supplies an explicit `seed`, honored exactly, which keeps s_key
10567 // stable across its whole conversation.
10568 // COMPOSITION RULE (fspec x gsd merge): the in-graph chain samples from the RAW
10569 // softmax — it can hold neither per-row filter stats nor the varying penalty history.
10570 // The sampled graph therefore engages only in the PURE-TEMP regime; filters/penalties
10571 // force the eager draft (which computes stats/penalties per row).
10572 // KEY THE WHOLE REGIME, not just the baked constants (lane/graph-s-key-exactness-
10573 // 20260819). `s_key` used to be `(seed, temp, k)`; the filters and penalties were left
10574 // out, so a filtered request resuming a session that parked a PURE-TEMP graph kept it —
10575 // and the launch site never re-asked `pure_temp`. See [`SampledGraphKey`] for what that
10576 // costs (an unconditional accept of out-of-head draft tokens, i.e. an exactness bug on
10577 // the request shape the vendor-default flip makes the majority).
10578 let s_key = SampledGraphKey::new(sp_seed, sp_temp, k, sp.top_k, sp.top_p, sp.min_p, pen_on);
10579 let pure_temp = s_key.pure_temp();
10580 if sampled && dctx.s_key.is_some_and(|old| old != s_key) {
10581 dctx.graph_s = None;
10582 dctx.failed.clear_sampled();
10583 dctx.s_key = None;
10584 dctx.q_slots.clear();
10585 dctx.keeper_s.clear();
10586 }
10587 if graph_draft
10588 && sampled
10589 && pure_temp
10590 && dctx.graph_s.is_none()
10591 && !dctx.failed.sampled_failed()
10592 {
10593 let DraftGraphCtx {
10594 g_tok,
10595 g_pos,
10596 g_seed,
10597 g_p,
10598 g_ctr,
10599 g_perturb,
10600 g_q,
10601 ..
10602 } = &mut dctx;
10603 // CAPTURE-RETAIN (#68 fix): same keeper contract as the greedy capture above.
10604 let cap_res = e.capture_graph_retained(|e| {
10605 self.mtp_head_forward_cap(
10606 e,
10607 mtp,
10608 g_tok,
10609 g_pos,
10610 g_seed,
10611 g_p,
10612 &mut *scratch,
10613 p_min > 0.0,
10614 true,
10615 embd_gpu.expect("graph draft requires resident embedding"),
10616 embd_qt,
10617 embd_rb,
10618 d_vocab,
10619 Some((g_ctr, g_perturb, g_q, sp_seed, sp_temp)),
10620 None,
10621 None, // constrained spec is greedy-only — sampled never carries a hook
10622 )
10623 });
10624 match cap_res {
10625 Ok((g, keep)) => {
10626 scratch.set_len(e, base)?;
10627 for _ in 0..k {
10628 dctx.q_slots.push(e.zeros(d_vocab)?);
10629 }
10630 dctx.graph_s = Some(g);
10631 dctx.s_key = Some(s_key);
10632 dctx.keeper_s = keep;
10633 }
10634 Err(err) => {
10635 scratch.set_len(e, base)?;
10636 // LOUD flip (audit Q2): same contract as the greedy capture above.
10637 if let Some(line) = dctx.failed.mark_sampled(&err.to_string()) {
10638 eprintln!("{line}");
10639 }
10640 }
10641 }
10642 }
10643 // ---- EXACTNESS GUARD, the enforceable half (lane/graph-s-key-exactness-20260819) ----
10644 // With the filters and penalties in `s_key`, a graph that SURVIVED the drop above was
10645 // captured under this request's exact regime, and capture requires `pure_temp` — so a
10646 // parked graph implies `pure_temp`. That implication is the whole exactness argument for
10647 // the graph arm, so it is asserted here rather than assumed: a future change that widens
10648 // the capture condition, narrows the key, or copies a `DraftGraphCtx` across regimes
10649 // fails LOUDLY at this line instead of silently drafting from the raw softmax while the
10650 // verify applies filter stats. Release builds refuse the graph (drop it, draft eager)
10651 // rather than launching it; the launch site re-tests `pure_temp` independently.
10652 if sampled && !pure_temp && dctx.graph_s.is_some() {
10653 debug_assert!(
10654 false,
10655 "sampled draft graph parked under {:?} survived into a FILTERED request \
10656 (top_k={} top_p={} min_p={} pen_on={}): the in-graph chain draws from the RAW \
10657 softmax, so the verify's filtered q would test a distribution the draft was \
10658 never sampled from",
10659 dctx.s_key, sp.top_k, sp.top_p, sp.min_p, pen_on,
10660 );
10661 eprintln!(
10662 "[spec] BUG: dropping a parked sampled draft graph that outlived its capture \
10663 regime (s_key={:?}, request top_k={} top_p={} min_p={} pen_on={}); drafting \
10664 EAGER — the key must carry every field that shapes q",
10665 dctx.s_key, sp.top_k, sp.top_p, sp.min_p, pen_on,
10666 );
10667 dctx.graph_s = None;
10668 dctx.s_key = None;
10669 dctx.q_slots.clear();
10670 dctx.keeper_s.clear();
10671 }
10672 // SKEY PROBE (MEMRA_SKEY_PROBE=1): the burst-entry facts the reachability question turns
10673 // on — is this request sampled, is it in the pure-temp regime the sampled graph is only
10674 // legal in, and is a graph PARKED from an earlier request of the same session? The launch
10675 // arms below print which chain actually ran, so the probe never restates the condition.
10676 if skey_probe() {
10677 eprintln!(
10678 "[skey] burst sampled={} pure_temp={} temp={} top_k={} top_p={} min_p={} \
10679 pen_on={} k={} graph_draft={} graph_s_parked={} s_key_parked={:?}",
10680 sampled as u8,
10681 pure_temp as u8,
10682 sp_temp,
10683 sp.top_k,
10684 sp.top_p,
10685 sp.min_p,
10686 pen_on as u8,
10687 k,
10688 graph_draft as u8,
10689 dctx.graph_s.is_some() as u8,
10690 dctx.s_key,
10691 );
10692 }
10693 let t_cap = t_ent.elapsed();
10694 // PERSISTENT DRAFT KV: fill the MTP block's K/V for every prompt position from the exact
10695 // trunk hiddens collected during prime — ONE batched K/V-only pass (overwrites any
10696 // capture-warmup garbage; capture left len at 0). last_token (the init feed) needs no
10697 // fill: the first chain step processes it and appends its entry at slot prompt.len().
10698 if let Some(ph) = &prompt_h {
10699 // SESSION: rows [0..base) are the previous turns' exact fills (refresh overwrote them
10700 // with true verify hiddens) — truncate any draft overhang, fill ONLY the suffix at
10701 // global positions [base..base+tp). Fresh call: base==0, identical to before.
10702 scratch.set_len(e, base)?;
10703 // CHUNKED FILL (long-ctx OOM fix, 2026-07-05): mtp_kv_fill's transients scale with its
10704 // T (concat = T*2*n_embd*4B — 1.5GB at 40k) and its concat loop is 2*T launches. The
10705 // fill is a pure sequential append, so chunking is exact: each chunk appends its rows
10706 // at pos0=base+start with the identical per-row math. Same knob as the trunk prime.
10707 let tp = prompt.len();
10708 let fill_chunk: usize = if crate::cache::swa_ring_on() {
10709 crate::hybrid_forward::prime_chunk_tokens(tp, self.layers.len())
10710 } else {
10711 // Preserve the flag-OFF schedule byte-for-byte, including the legacy zero value
10712 // meaning one monolithic fill.
10713 std::env::var("MEMRA_PRIME_CHUNK")
10714 .ok()
10715 .and_then(|v| v.parse().ok())
10716 .unwrap_or(4096)
10717 };
10718 let fill_chunk = if fill_chunk == 0 { tp } else { fill_chunk };
10719 let mut start = 0usize;
10720 while start < tp {
10721 let end = (start + fill_chunk).min(tp);
10722 let tc = end - start;
10723 {
10724 // PREDECESSOR pairing: row i gets h[i-1]; global row 0 a zeros row (the
10725 // reference engine's initial pending-h is zeroed too); a session turn's row 0
10726 // gets the PREVIOUS turn's last committed hidden (sess.last_h). Per chunk:
10727 // rows start..end read h[start-1..end-1] — one dtod into a chunk buffer.
10728 let mut phs = e.zeros(tc * n_embd)?;
10729 let (src_lo, dst_off) = if start == 0 {
10730 (0, n_embd)
10731 } else {
10732 ((start - 1) * n_embd, 0)
10733 };
10734 let n_copy = if start == 0 {
10735 (tc - 1) * n_embd
10736 } else {
10737 tc * n_embd
10738 };
10739 if start == 0 {
10740 if let Some((_, lh, _, _, _)) = sess_tail.as_ref() {
10741 if let Some(lh) = lh.as_ref() {
10742 e.copy_into(&mut phs, 0, lh, n_embd)?;
10743 }
10744 }
10745 }
10746 if n_copy > 0 {
10747 e.copy_view_into(
10748 &mut phs,
10749 dst_off,
10750 &ph.slice(src_lo..src_lo + n_copy),
10751 n_copy,
10752 )?;
10753 }
10754 self.mtp_kv_fill_all(
10755 e,
10756 &prompt[start..end],
10757 &phs,
10758 base + start,
10759 &mut *scratch,
10760 embd_dev,
10761 )?;
10762 }
10763 start = end;
10764 }
10765 }
10766 // MEMRA_PROFILE_SPEC=2: profiler capture starts HERE — after the prime, so an
10767 // `nsys -c cudaProfilerApi` capture contains ONLY the round loop (draft/verify/commit).
10768 // (=1 brackets the whole call in run_spec.rs, prime included.)
10769 if std::env::var("MEMRA_PROFILE_SPEC").as_deref() == Ok("2") {
10770 unsafe extern "C" {
10771 fn cudaProfilerStart() -> i32;
10772 }
10773 unsafe {
10774 cudaProfilerStart();
10775 }
10776 }
10777 // ROUND-STREAM stage (c) 4 (MEMRA_SPEC_STREAM=1, experimental): pre-issued M-round
10778 // bursts with ZERO per-round host readbacks — the accept/seed/rollback/ring kernels
10779 // consume each other's device outputs; the host drains the ring every M rounds. v1
10780 // constraints: greedy, !spec_replay, single-shot, batched-linear layers, no refresh
10781 // fills (acceptance effect A/B-arbitrated), enters from round 1 (pending guaranteed).
10782 // NOTE: not gated on the caller's graph_draft (its trunk_dense conjunct turns the 35B
10783 // MoE off) — the stream capture encloses ONLY the dense MTP head; the head-dense /
10784 // full-prec / k gates are re-derived here and a failed capture degrades to stream-off.
10785 let stream_on = crate::spec::spec_stream()
10786 && !sampled
10787 && !spec_replay
10788 && self.mtp_extra.is_empty()
10789 && constraint.is_none()
10790 && !session_mode
10791 && embd_gpu.is_some()
10792 && !crate::model::full_prec_enabled()
10793 && k + 2 < 96;
10794 let mut stream_graph: Option<cudarc::driver::CudaGraph> = None;
10795 let mut g_tokp2k = e.alloc_u32_zeroed(2 * k.max(1))?;
10796 if stream_on {
10797 let cap = e.capture_graph(|e| {
10798 for j in 0..k.max(1) {
10799 self.mtp_head_forward_cap(
10800 e,
10801 mtp,
10802 &mut dctx.g_tok,
10803 &mut dctx.g_pos,
10804 &mut dctx.g_seed,
10805 &mut dctx.g_p,
10806 &mut *scratch,
10807 true,
10808 true,
10809 embd_gpu.expect("round stream requires resident embedding"),
10810 embd_qt,
10811 embd_rb,
10812 d_vocab,
10813 None,
10814 Some((&mut g_tokp2k, j, d2t_dev.as_ref())),
10815 None, // round-stream requires constraint.is_none() (see stream_on)
10816 )?;
10817 }
10818 Ok(())
10819 });
10820 match cap {
10821 Ok(g) => {
10822 scratch.set_len(e, 0)?;
10823 stream_graph = Some(g);
10824 }
10825 Err(err) => {
10826 scratch.set_len(e, 0)?;
10827 if debug_spec {
10828 eprintln!("[spec] stream-graph capture failed ({err}); stream off");
10829 }
10830 }
10831 }
10832 }
10833 let stream_active = stream_on && stream_graph.is_some();
10834 if debug_spec {
10835 eprintln!(
10836 "[spec] stream_on={stream_on} env={} samp={sampled} dg={} captured={} active={stream_active} session={session_mode} replay={spec_replay}",
10837 crate::spec::spec_stream(),
10838 dctx.graph.is_some(),
10839 stream_graph.is_some()
10840 );
10841 }
10842 let t_v_s = k + 1;
10843 // ROUND-STREAM buffers + ptr tables now live in the model-generic round_stream
10844 // module (extracted 2026-07-12; the gemma burst reuses them).
10845 let sb = crate::round_stream::StreamBufs::new(e, k, crate::spec::spec_stream_m())?;
10846 let crate::round_stream::StreamBufs {
10847 mut vtok_d,
10848 mut brk_d,
10849 mut pend_d,
10850 last_pred_d,
10851 mut pos_ctr,
10852 mut pos_start_d,
10853 mut ring_d,
10854 acc_d: mut stream_acc,
10855 m_rounds,
10856 k: _,
10857 } = sb;
10858 let stream_ptrs: Option<CudaSlice<u64>> = if stream_active {
10859 Some(crate::round_stream::kv_len_ptr_table(
10860 e,
10861 cache,
10862 Some(&pos_ctr),
10863 )?)
10864 } else {
10865 None
10866 };
10867
10868 let t_fill = t_ent.elapsed();
10869 let mut round = 0usize;
10870 // ADAPTIVE DRAFT LENGTH (MEMRA_SPEC_ADAPT=1, opt-in — the gemma_spec accepted-run law,
10871 // ported 2026-08-01): next round's draft depth = last round's accepted run + 1, clamped
10872 // to [floor(pos), k_cap] — a miss shrinks the next draft to the miss point + 1,
10873 // full-accept streaks re-deepen one step per round. NOT the 2026-07-07 acceptance-EMA
10874 // (that arm measured an HONEST LOSS to static per-class optima — 115.0/85.8/73.4 vs
10875 // 121.6/92.7/75.6, EMA lag — and was removed 2026-07-08; rig5090.jsonl has the record).
10876 // The gemma law has no lag class: it reacts within one round, and was worth +7-20% on
10877 // the gemma cells at unchanged exactness (2026-07-10 flip; floor sweep 2026-07-25;
10878 // position key 2026-07-26). Signal = n_acc from the round's EXISTING accept readback —
10879 // zero new syncs; the draft graph is a SINGLE-STEP capture replayed per drafted token,
10880 // so a per-round depth needs no re-capture (unlike gemma's whole-chain graphs). qwen's
10881 // in-round p-min cut already shortens chains mid-round, so gemma's one-round-late p-min
10882 // fold into kc is unnecessary here — the accepted-run law sees the cut via n_acc.
10883 // Exactness is the verify's job at ANY depth (same contract as p-min variable rounds).
10884 // DEFAULT OFF on the qwen path until its cells gate a flip (gemma's is default-on).
10885 // MEASURED 2026-08-01 (H100 GPU-3, interleaved x3, NGEN=256, same-invocation plain
10886 // denominators; research/qwen-adaptive-k-20260801/): REFUTED on the tuned qwen configs.
10887 // q27 K=3+HPOST+PMIN=0.3: short +0.8% (noise; law ~idles, len_hist identical), board
10888 // -1.9%, agentic -0.5%; board PMIN=0 -2.1% (not p-min shadowing — the law itself);
10889 // floor=1 -2.8% (gemma's floor-collapse, reproduced). q35 K=2 board: -6.4% (52/136
10890 // rounds shrink to depth 1; no depth to reclaim at K=2). The gemma direction DOES
10891 // appear at untuned depth-K — q27 K=6 floor=4 +1.5% over fixed K=6 — but stays -3.7%
10892 // below fixed K=3: same verdict class as the retired EMA arm (honest loss to static
10893 // per-class optima). Acceptance-rate rises under the law while tokens/round falls —
10894 // it buys accept-% by adding rounds, and a round's fixed draft+verify cost wins.
10895 // K=1..8 self-consistency PASS both models with the law ON (exactness held).
10896 let adapt = std::env::var("MEMRA_SPEC_ADAPT").as_deref() == Ok("1");
10897 // floor: per-model default keyed on n_embd (gemma's tiering — models with an expensive
10898 // verify keep deep drafts after a miss); MEMRA_SPEC_ADAPT_FLOOR pins it everywhere.
10899 let adapt_floor_env: Option<usize> = std::env::var("MEMRA_SPEC_ADAPT_FLOOR")
10900 .ok()
10901 .and_then(|v| v.parse().ok());
10902 let adapt_floor_default: usize = if self.cfg.n_embd as usize >= 3500 {
10903 4
10904 } else if self.cfg.n_embd as usize >= 2500 {
10905 2
10906 } else {
10907 1
10908 };
10909 let adapt_floor: usize = adapt_floor_env.unwrap_or(adapt_floor_default);
10910 // position key: past floor_ctx a HIGH floor (>=4) relaxes to 1 — forced-deep drafts
10911 // turn net-negative at depth (gemma 31B d1736 evidence); MEMRA_SPEC_FLOOR_CTX moves
10912 // the boundary, an explicit MEMRA_SPEC_ADAPT_FLOOR pins the floor everywhere.
10913 let floor_ctx: usize = std::env::var("MEMRA_SPEC_FLOOR_CTX")
10914 .ok()
10915 .and_then(|v| v.parse().ok())
10916 .unwrap_or(1024);
10917 let floor_at = |pos: usize| -> usize {
10918 if adapt_floor_env.is_some() || pos < floor_ctx {
10919 adapt_floor
10920 } else if adapt_floor >= 4 {
10921 1
10922 } else {
10923 adapt_floor
10924 }
10925 };
10926 // cap: MEMRA_SPEC_CAPMAX (gemma semantics, default 7). Binds only under adapt — the
10927 // fixed-K default path is untouched by this whole block.
10928 let cap_max: usize = std::env::var("MEMRA_SPEC_CAPMAX")
10929 .ok()
10930 .and_then(|v| v.parse().ok())
10931 .unwrap_or(7);
10932 let k_cap = k.min(cap_max).max(1);
10933 let mut kc = k_cap;
10934 let mut opti_fork: Option<OptiForkState> = None;
10935 let mut fork_snapshot: Option<crate::cache::CacheSnapshot> = None;
10936 if fork_mode != OptiForkGateMode::Disabled {
10937 let fence = crate::pp::pp_cuts(self.layers.len());
10938 let refusal = if !session_mode {
10939 Some("not-session")
10940 } else if k != 1 || adapt {
10941 Some("requires-fixed-k1")
10942 } else if sampled || constraint.is_some() || spec_replay {
10943 Some("sampled-constrained-or-replay")
10944 } else if pipe.is_some() {
10945 Some("two-session-pipeline")
10946 } else if !spec_devacc() {
10947 Some("requires-device-accept")
10948 } else if stream_active || crate::spec::spec_stream() {
10949 Some("round-stream")
10950 } else if !self.mtp_extra.is_empty() {
10951 Some("multi-head-mtp")
10952 } else if crate::cache::swa_ring_on() || cache.has_swa_ring() {
10953 Some("swa-ring")
10954 } else if crate::pp::pp_host_bounce_active() {
10955 Some("host-bounce")
10956 } else if fork_mode == OptiForkGateMode::Controller
10957 && cache.recur.iter().any(Option::is_some)
10958 {
10959 Some("controller-requires-zero-recurrent-state")
10960 } else if fence.as_ref().is_none_or(|f| f.len() != 3) {
10961 Some("requires-pp2")
10962 } else {
10963 None
10964 };
10965 if let Some(reason) = refusal {
10966 OPTI_FORK_REFUSALS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
10967 eprintln!("[opti-fork] refused reason={reason}");
10968 } else {
10969 let fence = fence.expect("validated PP-2 fence");
10970 let rt = crate::pp::PpNRt::get(e)?;
10971 let primary_stage0 = rt.engine(0, e).ctx().ordinal() == e.ctx().ordinal();
10972 let primary_stage1 = rt.engine(1, e).ctx().ordinal() == e.ctx().ordinal();
10973 let primary_supported =
10974 primary_stage0 || (fork_mode == OptiForkGateMode::Controller && primary_stage1);
10975 if !rt.cross_device() || !primary_supported {
10976 OPTI_FORK_REFUSALS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
10977 eprintln!("[opti-fork] refused reason=requires-supported-primary-cross-device");
10978 } else {
10979 // Both recurrent snapshots and both seed generations are allocated before
10980 // the first fork, each through its owning PP stage. Allocation failure
10981 // therefore happens before any optimistic state mutation can occur.
10982 let current_snapshot = opti_snapshot_stage_owned(e, cache, rt, &fence)?;
10983 let alternate_snapshot = opti_snapshot_stage_owned(e, cache, rt, &fence)?;
10984 let fork = OptiForkState::new(
10985 e,
10986 cache,
10987 fork_mode,
10988 alternate_snapshot,
10989 &h_seed_buf,
10990 &fill_prev,
10991 rt,
10992 fence[1],
10993 self.layers.len(),
10994 )?;
10995 eprintln!(
10996 "[opti-fork] armed mode={fork_mode:?} snapshots=2 seeds=2 split={} \
10997 payload_dev0={} payload_dev1={} q_threshold={:.3}",
10998 fence[1],
10999 fork.logical_payload_bytes[0],
11000 fork.logical_payload_bytes[1],
11001 fork.controller.map_or(0.0, |policy| policy.threshold),
11002 );
11003 fork_snapshot = Some(current_snapshot);
11004 opti_fork = Some(fork);
11005 }
11006 }
11007 }
11008 // Persistent snapshot buffers are allocated once and refreshed in place. The fork arm
11009 // uses stage-owned snapshots; refused/disabled arms retain the existing generic helper.
11010 let mut snap = match fork_snapshot {
11011 Some(snapshot) => snapshot,
11012 None => cache.snapshot(e)?,
11013 };
11014 let mut carried_opti: Option<OptiControllerTicket> = None;
11015 // ROUND-STREAM stage (b) 3a: device table of per-layer kvl.len_d pointers (stable — the
11016 // cache never reallocates len_d; see cache.rs "stable pointer" note). 0 = no KV layer.
11017 let kv_len_ptrs: Option<CudaSlice<u64>> = if spec_devacc() && !spec_replay {
11018 Some(crate::round_stream::kv_len_ptr_table(e, cache, None)?)
11019 } else {
11020 None
11021 };
11022 // BONUS FOLD (2026-07-04): after a FULL accept the bonus token is NOT committed with a
11023 // separate T=1 trunk pass (a full weight read per round). It stays PENDING and rides as
11024 // column 0 of the NEXT round's verify batch. Under predecessor pairing the next chain
11025 // seeds from the bonus's predecessor's TRUE verify hidden (free — no extra
11026 // pass of any kind). Verify still
11027 // checks every emitted token against the target -> exactness holds by construction; only
11028 // DRAFT QUALITY can shift, which the acceptance numbers arbitrate.
11029 // bonus emitted but not yet committed to cache. A carried pending (see SpecSession::
11030 // pending_tok) enters round 0 directly — the burst boundary becomes a plain round edge.
11031 let mut pending: Option<u32> = carried_pending;
11032 // MEMRA_SPEC_PHASE=1: per-round wall decomposition (draft / verify / accept+commit) —
11033 // no tracing, no extra syncs (each phase is naturally sync-bounded: draft readbacks,
11034 // the verify accept readback). Printed once at loop end via spec-stats.
11035 let anatomy_on = std::env::var("MEMRA_SPEC_PP_ANATOMY").as_deref() == Ok("1");
11036 let phase_on = anatomy_on || std::env::var("MEMRA_SPEC_PHASE").as_deref() == Ok("1");
11037 // DRAFT-MASK receipt (lane/draft-mask): speculative-clone wall + rounds, printed with
11038 // spec-stats. The clone is the one cost the design adds per round — measured, not assumed.
11039 let (mut dm_clone_ns, mut dm_rounds) = (0u128, 0usize);
11040 // grammar-truncation counters: how many rounds the verify-side cut fired and how many
11041 // already-verified tokens it threw away. THIS is the quantity draft masking targets.
11042 let (mut dm_cuts, mut dm_cut_tokens) = (0usize, 0usize);
11043 let (mut ph_draft, mut ph_verify, mut ph_rest) = (0f64, 0f64, 0f64);
11044 let mut ph_wait = 0f64;
11045 let mut ph_commit = 0f64;
11046 let mut ph_t = std::time::Instant::now();
11047 let mut ph_mark = |acc: &mut f64, on: bool| {
11048 if on {
11049 let now = std::time::Instant::now();
11050 *acc += (now - ph_t).as_secs_f64();
11051 ph_t = now;
11052 }
11053 };
11054 // MTP-ROUTE VERIFY GRAPHS (`MEMRA_SPEC_VERIFY_GRAPH`, see the flag doc): the
11055 // model-owned capture pool, locked for the whole burst exactly as the dspark serve
11056 // arm holds it — the slab stash is live verify -> commit inside a round, and the
11057 // worker drives rounds from one scheduler thread. PERSISTENT across generations on
11058 // the model (rebuilding per call re-captures the pool per prompt, which is the
11059 // measured way to lose more than the launches cost); the captured bodies are
11060 // cache-independent, every state read going through per-round refreshed pointer
11061 // tables. None = the eager walk, byte-identical.
11062 //
11063 // Never armed together with ROUND-STREAM: the tparallel verify refuses that pair
11064 // loudly, and `stream_active` owns the burst arm above, so the door stays shut
11065 // whenever the stream is live rather than relying on that refusal.
11066 // The lock is taken ONLY when the door is armed: with the flag off this whole block
11067 // is inert, so the default path cannot serialize two spec generations behind a mutex
11068 // it never reads.
11069 let vg_armed =
11070 crate::spec::spec_verify_graph_env().unwrap_or_else(|| self.vgraph_family_default());
11071 let mut vg_guard = if vg_armed && !stream_active {
11072 let mut g = self.dspark_vgraphs.lock().unwrap();
11073 if g.is_none() {
11074 // Size by the WIDEST verify this run can present, which is k+1 and NOT
11075 // k_cap+1: the sampled arm's own window is `t_v_s = k + 1`, so a pool built
11076 // from a smaller adaptive cap gets sliced past its stash rows (a `slice_mut`
11077 // panic in the sampled ON arm, measured before this line said k+1).
11078 let vt_cap = (k.max(k_cap) + 1).max(2);
11079 *g = DsparkVerifyGraphs::new(e, cache, vt_cap, n_embd)?;
11080 if g.is_some() {
11081 // Engagement receipt (the dead-arm lesson): prove the door is LIVE rather
11082 // than trusting that a flag set means a pool built.
11083 eprintln!("[spec-vg] MTP verify-graph pool ENGAGED (vt_cap={vt_cap})");
11084 } else {
11085 eprintln!(
11086 "[spec-vg] MTP verify-graph pool declined (no linear layers, \
11087 non-uniform state, or vt_cap < 2) — eager walk"
11088 );
11089 }
11090 }
11091 Some(g)
11092 } else {
11093 None
11094 };
11095 // Capacity fail-safe: a round wider than the pool was built for must take the eager
11096 // walk, not slice the stash past its rows. The sizing above already covers every
11097 // round this run can present; this keeps a future caller (or a k that grows behind
11098 // the pool's back) on the byte-identical fallback instead of a panic.
11099 let vg_t_cap = vg_guard
11100 .as_ref()
11101 .and_then(|g| g.as_ref())
11102 .map(|g| g.t_capacity())
11103 .unwrap_or(0);
11104 if let Some(p) = pipe {
11105 p.setup_end();
11106 }
11107 while keep_going && out.len() < max_new {
11108 // ROUND-STREAM BURST: from round 1 (pending guaranteed by every non-replay arm),
11109 // issue M rounds with zero readbacks, then drain the ring + reconcile mirrors.
11110 if let (true, Some(sg), Some(ptrs)) = (
11111 stream_active && round >= 1 && pending.is_some(),
11112 &stream_graph,
11113 &stream_ptrs,
11114 ) {
11115 if debug_spec {
11116 static ONCE: std::sync::Once = std::sync::Once::new();
11117 ONCE.call_once(|| {
11118 eprintln!("[memra] ROUND-STREAM burst engaged (M={m_rounds} k={k})")
11119 });
11120 }
11121 e.set_i32_one(&mut pos_ctr, cache.pos as i32)?;
11122 e.set_u32_one(&mut pend_d, pending.unwrap())?;
11123 e.set_u32_one(&mut ring_d, 0)?; // ring count = 0 (writes element 0)
11124 for _mi in 0..m_rounds {
11125 e.i32_copy_add(&pos_ctr, &mut pos_start_d, 0)?;
11126 cache.snapshot_into(e, &mut snap)?; // device D2Ds, stream-ordered
11127 e.i32_copy_add(&pos_ctr, &mut scratch.kv.len_d, 0)?; // draft-KV rollback
11128 e.i32_copy_add(&pos_ctr, &mut dctx.g_pos, 1)?; // rope pos = pos + base
11129 e.u32_copy(&pend_d, &mut dctx.g_tok)?;
11130 e.copy_into(&mut dctx.g_seed, 0, &h_seed_buf, n_embd)?;
11131 sg.launch()?;
11132 e.spec_assemble_verify(
11133 &g_tokp2k,
11134 &pend_d,
11135 d2t_dev.as_ref(),
11136 &mut vtok_d,
11137 &mut brk_d,
11138 p_min,
11139 k,
11140 pmin0,
11141 )?;
11142 let mut ck = VerifyCkpt::new(self.layers.len());
11143 let dummy = vec![0u32; t_v_s];
11144 let (tl_d, vx) = self.decode_step_t_core_stream(
11145 e,
11146 &dummy,
11147 0,
11148 &mut *cache,
11149 embd_dev,
11150 Some(&mut ck),
11151 Some((&vtok_d, &pos_ctr)),
11152 None,
11153 None,
11154 None,
11155 )?;
11156 for j in 0..t_v_s {
11157 e.argmax_token_device_col(&tl_d, j, n_vocab, &mut preds_d, j)?;
11158 }
11159 e.spec_accept_greedy_dc(
11160 &preds_d,
11161 &vtok_d,
11162 &last_pred_d,
11163 &brk_d,
11164 &mut stream_acc,
11165 )?;
11166 e.spec_seed_gather(&vx, &fill_prev, &stream_acc, &mut h_seed_buf, 1, n_embd)?;
11167 e.copy_into(&mut fill_prev, 0, &h_seed_buf, n_embd)?;
11168 self.commit_verified_prefix_stream(
11169 e,
11170 &mut *cache,
11171 &snap,
11172 &ck,
11173 &stream_acc,
11174 1,
11175 t_v_s,
11176 )?;
11177 e.spec_rollback_stream(
11178 ptrs,
11179 &pos_start_d,
11180 &stream_acc,
11181 1,
11182 self.layers.len() + 1,
11183 )?;
11184 e.spec_ring_commit(&vtok_d, &stream_acc, &brk_d, &mut ring_d, &mut pend_d)?;
11185 }
11186 e.stream().synchronize()?;
11187 let ring_h = e.dtoh_u32(&ring_d)?;
11188 let cnt = ring_h[0] as usize;
11189 for i in 0..cnt {
11190 if out.len() < max_new {
11191 out.push(ring_h[1 + i]);
11192 }
11193 }
11194 let pos_h = e.dtoh_i32(&pos_ctr)?[0] as usize;
11195 for il in 0..self.layers.len() {
11196 if let Some(kvl) = cache.kv[il].as_mut() {
11197 kvl.len = pos_h;
11198 }
11199 }
11200 cache.pos = pos_h;
11201 scratch.kv.len = pos_h;
11202 pending = Some(ring_h[cnt]); // last drained token = the live bonus
11203 last_token = ring_h[cnt];
11204 total_drafted += k * m_rounds; // upper bound (p-min breaks uncounted)
11205 total_accepted += cnt.saturating_sub(m_rounds);
11206 if let Some(t) = sess_telem {
11207 // totals only — the burst's per-round accept counts stayed on device
11208 // (that is the point of the round-stream arm). pos_* untouched.
11209 t.record_totals(m_rounds, k * m_rounds, cnt.saturating_sub(m_rounds));
11210 }
11211 round += m_rounds;
11212 // sse-cadence: the drained ring is committed — flush it at burst-drain cadence.
11213 keep_going = flush_commit(&mut on_commit, &out, &mut flushed);
11214 continue;
11215 }
11216 let pipe_draft = match pipe {
11217 Some(p) => Some(p.draft_begin(round)?),
11218 None => None,
11219 };
11220 let pos = cache.pos; // #tokens committed (EXCLUDES a pending bonus)
11221 let mut current_opti = carried_opti.take();
11222 let mut fork_generation = if current_opti.is_none() && pending.is_some() {
11223 match opti_fork.as_mut() {
11224 Some(fork) if fork.mode.is_forced() => Some(fork.reserve(&mut snap)?),
11225 None => None,
11226 Some(_) => None,
11227 }
11228 } else {
11229 None
11230 };
11231 if current_opti.is_none() {
11232 if let Some(fork) = opti_fork.as_ref() {
11233 opti_snapshot_stage_owned_into(e, cache, fork.rt, &fork.fence, &mut snap)?;
11234 } else {
11235 cache.snapshot_into(e, &mut snap)?;
11236 }
11237 } else if snap.pos != pos {
11238 return Err(format!(
11239 "optipipe carried snapshot pos {} != current pos {pos}",
11240 snap.pos
11241 )
11242 .into());
11243 } // §C: snapshot BEFORE draft+verify (already retained for a carried successor)
11244 ph_mark(&mut ph_rest, phase_on);
11245
11246 // --- 1. DRAFT k tokens with the NextN head (autoregressive, T=1 each) ---
11247 // p-min semantics (both paths): stop the chain early when the head's confidence in
11248 // its own pick drops below p_min — the just-drafted token is DISCARDED, but its
11249 // scratch append stands (identical to the eager chain's ordering). j==0 always drafts.
11250 let base0 = if pending.is_some() { 1usize } else { 0usize };
11251 // fixed draft length by default; MEMRA_SPEC_ADAPT=1 drafts at last round's
11252 // accepted run + 1 (the gemma law — see the setup block above the loop).
11253 let k_this = if adapt { kc } else { k };
11254 let mut draft: Vec<u32> = Vec::with_capacity(k);
11255 let mut draft_idx: Vec<u32> = Vec::with_capacity(k); // trimmed-vocab ids (== draft when untrimmed)
11256 let mut controller_draft_prob: Option<f32> = None;
11257 let mut controller_eager_state: Option<(u32, CudaSlice<f32>)> = None;
11258 if let Some(ticket) = current_opti.as_mut() {
11259 let carried_pending = pending.ok_or("optipipe carried successor lost pending")?;
11260 if ticket.verify_tokens[0] != carried_pending {
11261 return Err(format!(
11262 "optipipe carried pending mismatch: ticket={} live={carried_pending}",
11263 ticket.verify_tokens[0],
11264 )
11265 .into());
11266 }
11267 draft.push(ticket.verify_tokens[1]);
11268 controller_draft_prob = Some(ticket.draft_prob);
11269 controller_eager_state = ticket
11270 .take_eager_seed()
11271 .map(|seed| (ticket.verify_tokens[1], seed));
11272 } else {
11273 // Round-start draft-KV sync (BOTH paths). Persistent: truncate/align to the committed
11274 // history — slots 0..P hold entries for the tokens before last_token@P (P = pos +
11275 // base0 - 1); this single set_len IS the draft-side rollback (drops last round's
11276 // rejected drafts and p-min extras via the len mechanism).
11277 scratch.set_len(e, pos + base0 - 1)?;
11278 if pen_on {
11279 // PEN_WINDOW_MAX also bounds the per-round upload and the O(n_hist^2)
11280 // device dedup: `penalty_last_n` is usize::MAX for any serve request with
11281 // a penalty, so without the cap this grew with the whole session.
11282 let win = sp.penalty_last_n.min(PEN_WINDOW_MAX);
11283 let w0 = pen_hist.len().saturating_sub(win);
11284 pen_hist_d = Some(e.htod_u32_v(&pen_hist[w0..])?);
11285 }
11286 if sampled {
11287 draft_logits.clear();
11288 draft_stats.clear();
11289 }
11290 // DRAFT-SIDE GRAMMAR MASK: clone the committed grammar state ONCE per round; each
11291 // position's mask is computed on that clone and advanced by the PROPOSED token. The
11292 // real state moves only on emission (verify's job), so the emitted stream is
11293 // unchanged — the mask only removes tokens the verify would have truncated anyway.
11294 let mut dmask_live = dmask_on;
11295 if dmask_live {
11296 let t_c = std::time::Instant::now();
11297 constraint
11298 .as_deref_mut()
11299 .unwrap()
11300 .draft_begin()
11301 .map_err(|e2| format!("constraint: {e2}"))?;
11302 dm_clone_ns += t_c.elapsed().as_nanos();
11303 dm_rounds += 1;
11304 }
11305 if let (false, Some(gr)) = (sampled || pen_on, &dctx.graph) {
11306 // GRAPH DRAFT: one dispatch per drafted token. The chain feeds itself on-device
11307 // (in-graph argmax -> tok_d -> next replay's embed; h_nextn -> h_seed_d; pos_d
11308 // inc'd in-graph); the host only reads 4B token (+4B p) and decides the break.
11309 e.set_i32_one(&mut dctx.g_pos, (pos + base0) as i32)?;
11310 e.set_u32_one(&mut dctx.g_tok, last_token)?;
11311 e.copy_into(&mut dctx.g_seed, 0, &h_seed_buf, n_embd)?;
11312 for j in 0..k_this {
11313 // per-position mask upload (contents only — the graph's baked pointer is
11314 // dctx.g_dmask). All-ones once masking goes dead mid-chain, so the captured
11315 // mask node degrades to a no-op ban instead of needing a second graph.
11316 if dmask_live
11317 && !upload_draft_mask(
11318 e,
11319 constraint.as_deref_mut().unwrap(),
11320 &mut dctx.g_dmask,
11321 mtp.d2t.as_ref(),
11322 d_vocab,
11323 dmask_words,
11324 )?
11325 {
11326 // no draft-vocab row is grammar-legal here (a trimmed FR-Spec head can
11327 // genuinely miss the legal set): neutralize the captured mask node and
11328 // finish the chain UNMASKED — exactly pre-lane behaviour, never worse.
11329 e.htod_u32_into(&mut dctx.g_dmask, &vec![u32::MAX; dmask_words])?;
11330 dmask_live = false;
11331 }
11332 gr.launch()?;
11333 scratch.kv.len += 1; // host mirror (len_d advanced in-graph)
11334 let idx = e.dtoh_u32_one(&dctx.g_tok)?;
11335 // #87 SENTINEL TRAP: an all-NaN head-logits row leaves the device argmax's
11336 // init sentinel (0x7FFFFFFF) in g_tok — feeding it onward dereferences
11337 // embed_row(sentinel) = table + ~4.6TB (never mapped) inside the NEXT graph
11338 // replay's embed node, and the MMU fault kills the CUDA context for the
11339 // whole process (research/pp2spec-crash-20260807: 3 coredumps, byte-exact
11340 // VA arithmetic). Refuse loudly instead; the diagnostics name the first-NaN
11341 // buffer (g_seed = the verify-side handoff vs head-side compute).
11342 if (idx as usize) >= d_vocab {
11343 // g_seed is SELF-FED (the replay writes h_nextn back into it), so it
11344 // reads as the head's OUTPUT at j; h_seed_buf is the round's INPUT
11345 // seed, untouched since the round-start copy — the pair discriminates
11346 // "seed arrived poisoned" from "head forward produced NaN".
11347 let seed_h = e.dtoh(&dctx.g_seed)?;
11348 let seed_nan = seed_h.iter().filter(|v| v.is_nan()).count();
11349 let in_h = e.dtoh(&h_seed_buf)?;
11350 let in_nan = in_h.iter().filter(|v| v.is_nan()).count();
11351 return Err(format!(
11352 "draft(graph) argmax sentinel 0x{idx:08x} >= d_vocab {d_vocab} at \
11353 round {round} j={j} pos={pos}: head-out NaN {seed_nan}/{n_embd}, \
11354 round-input-seed NaN {in_nan}/{n_embd} — refusing to dereference \
11355 the embed row (#87 trap)"
11356 )
11357 .into());
11358 }
11359 // trimmed draft vocab -> target token id (identity when no d2t map)
11360 let d = match &mtp.d2t {
11361 Some(map) => map[idx as usize],
11362 None => idx,
11363 };
11364 let draft_p = if p_min > 0.0
11365 || opti_fork
11366 .as_ref()
11367 .is_some_and(|fork| fork.controller.is_some())
11368 {
11369 Some(e.dtoh(&dctx.g_p)?[0])
11370 } else {
11371 None
11372 };
11373 if j == 0 {
11374 controller_draft_prob = draft_p;
11375 }
11376 if let Some(p) = draft_p.filter(|_| p_min > 0.0) {
11377 if p < p_min && (j > 0 || (pmin0 && base0 == 1)) {
11378 break;
11379 }
11380 }
11381 draft.push(d);
11382 // with a trimmed head the NEXT embed must read the TARGET id, not the draft
11383 // index the argmax wrote — patch the persistent token buffer (4B htod).
11384 if d != idx {
11385 e.set_u32_one(&mut dctx.g_tok, d)?;
11386 }
11387 // advance the SPECULATIVE state with the proposal; a dead chain drops to
11388 // unmasked drafting for the remaining positions (verify still arbitrates).
11389 // speculative advance; a chain the grammar can no longer follow (EOS
11390 // proposed) ends here. The captured mask node always runs, so a dead chain
11391 // leaves the buffer NEUTRAL (all-ones = ban nothing) before it exits.
11392 if dmask_live
11393 && !constraint
11394 .as_deref_mut()
11395 .unwrap()
11396 .draft_advance(d)
11397 .map_err(|e2| format!("constraint: {e2}"))?
11398 {
11399 e.htod_u32_into(&mut dctx.g_dmask, &vec![u32::MAX; dmask_words])?;
11400 break;
11401 }
11402 }
11403 // PURE-TEMP RE-TEST (lane/graph-s-key-exactness-20260819): the sampled graph is
11404 // legal ONLY in the regime it was captured in. The condition used to read
11405 // `(sampled, &dctx.graph_s)` and trusted `s_key` to have dropped anything else —
11406 // which it could not, because the key omitted the filters. Both halves are now
11407 // enforced: the key drops a stale graph, and this site refuses to launch one.
11408 } else if let (true, Some(gr)) = (sampled && pure_temp, &dctx.graph_s) {
11409 if skey_probe() {
11410 eprintln!(
11411 "[skey] chain=graph_s round={round} pure_temp={} top_k={} \
11412 top_p={} min_p={} s_key_parked={:?}",
11413 pure_temp as u8, sp.top_k, sp.top_p, sp.min_p, dctx.s_key,
11414 );
11415 }
11416 // SAMPLED GRAPH DRAFT: one replay per drafted token — head forward + gumbel +
11417 // argmax in ONE dispatch; the host reads 4B token (+4B p), D2Ds q into slot j,
11418 // and decides the break. Event-counter continuity: g_ctr is host-seeded to
11419 // sctr-1 ONCE per round (outside the graph); the in-graph bump runs BEFORE the
11420 // perturb, so replay j consumes counter sctr+j — exactly the eager arm's Philox
11421 // stream. Host sctr advances in lockstep (computed, no readback needed).
11422 e.set_i32_one(&mut dctx.g_pos, (pos + base0) as i32)?;
11423 e.set_u32_one(&mut dctx.g_tok, last_token)?;
11424 e.copy_into(&mut dctx.g_seed, 0, &h_seed_buf, n_embd)?;
11425 e.set_u32_one(&mut dctx.g_ctr, sctr.wrapping_sub(1))?;
11426 for j in 0..k_this {
11427 gr.launch()?;
11428 scratch.kv.len += 1; // host mirror (len_d advanced in-graph)
11429 sctr += 1; // mirrors the in-graph g_ctr bump (eager parity:
11430 // counts the p-min-discarded token too)
11431 // q retention: ONE async D2D of the persistent head-logits buffer into this
11432 // round's slot j (stream-ordered after the replay, before the next one).
11433 e.copy_into(&mut dctx.q_slots[j], 0, &dctx.g_q, d_vocab)?;
11434 let idx = e.dtoh_u32_one(&dctx.g_tok)?;
11435 // #87 SENTINEL TRAP (see the greedy graph arm above).
11436 if (idx as usize) >= d_vocab {
11437 let seed_h = e.dtoh(&dctx.g_seed)?;
11438 let seed_nan = seed_h.iter().filter(|v| v.is_nan()).count();
11439 return Err(format!(
11440 "draft(graph-sampled) argmax sentinel 0x{idx:08x} >= d_vocab \
11441 {d_vocab} at round {round} j={j} pos={pos}: round-seed NaN \
11442 {seed_nan}/{n_embd} — refusing to dereference the embed row \
11443 (#87 trap)"
11444 )
11445 .into());
11446 }
11447 let d = match &mtp.d2t {
11448 Some(map) => map[idx as usize],
11449 None => idx,
11450 };
11451 draft_idx.push(idx);
11452 if p_min > 0.0 {
11453 let p = e.dtoh(&dctx.g_p)?[0];
11454 if p < p_min && (j > 0 || (pmin0 && base0 == 1)) {
11455 break;
11456 }
11457 }
11458 draft.push(d);
11459 // trimmed head: the NEXT embed must read the TARGET id (see the greedy arm).
11460 if d != idx {
11461 e.set_u32_one(&mut dctx.g_tok, d)?;
11462 }
11463 }
11464 // uniform accept path: fill draft_stats per used slot (pure-temp regime — the
11465 // stats degenerate to th=0 / full-Z; one filter_stats launch per slot, tiny).
11466 for j in 0..draft.len().max(draft_idx.len()) {
11467 let rows0 = e.htod_i32(&[0])?;
11468 let (mut th_d, mut z_d, mut mx_d) = (e.zeros(1)?, e.zeros(1)?, e.zeros(1)?);
11469 e.filter_stats(
11470 &dctx.q_slots[j],
11471 d_vocab,
11472 &rows0,
11473 &mut th_d,
11474 &mut z_d,
11475 &mut mx_d,
11476 d_vocab,
11477 1,
11478 sp_temp,
11479 sp.top_k,
11480 sp.top_p,
11481 sp.min_p,
11482 )?;
11483 draft_stats.push((e.dtoh(&mx_d)?[0], e.dtoh(&th_d)?[0], e.dtoh(&z_d)?[0]));
11484 }
11485 } else {
11486 if skey_probe() && sampled {
11487 eprintln!(
11488 "[skey] chain=eager round={round} pure_temp={} top_k={} \
11489 top_p={} min_p={} s_key_parked={:?}",
11490 pure_temp as u8, sp.top_k, sp.top_p, sp.min_p, dctx.s_key,
11491 );
11492 }
11493 // EAGER DRAFT (fallback: MoE head/trunk, huge k, MEMRA_SPEC_NOGRAPH, capture fail).
11494 let chain_heads = !self.mtp_extra.is_empty();
11495 let mut e_tok = last_token;
11496 let mut d_seed = e.clone_dtod(&h_seed_buf)?;
11497 let mut chain_tokens = if chain_heads {
11498 vec![last_token]
11499 } else {
11500 Vec::new()
11501 };
11502 let mut chain_seeds = if chain_heads {
11503 vec![e.clone_dtod(&h_seed_buf)?]
11504 } else {
11505 Vec::new()
11506 };
11507 for j in 0..k_this {
11508 // GPU-ARGMAX DRAFT (2026-07-03): device logits + device argmax + 4-byte token
11509 // read instead of the ~600KB full-vocab dtoh + host argmax per draft token.
11510 let mtp_pos = pos + base0 + j;
11511 // draft-side grammar mask (eager twin of the graph arm's in-graph node).
11512 // A position with no legal draft-vocab row drops to unmasked drafting for
11513 // the rest of the chain (pre-lane behaviour; verify still arbitrates).
11514 if dmask_live {
11515 dmask_live = upload_draft_mask(
11516 e,
11517 constraint.as_deref_mut().unwrap(),
11518 &mut dctx.g_dmask,
11519 mtp.d2t.as_ref(),
11520 d_vocab,
11521 dmask_words,
11522 )?;
11523 }
11524 let mask = if dmask_live {
11525 Some((&dctx.g_dmask, dmask_words))
11526 } else {
11527 None
11528 };
11529 let (dl_d, h_nextn) = if chain_heads {
11530 if debug_spec {
11531 eprintln!(
11532 "[mtp-chain-step] round={round} j={j} head={} replay_rows={}",
11533 mtp_chain_head_index(j, self.mtp_head_count()),
11534 chain_tokens.len(),
11535 );
11536 }
11537 self.mtp_chain_forward_dev(
11538 e,
11539 &chain_tokens,
11540 &chain_seeds,
11541 &mut *scratch,
11542 pos + base0 - 1,
11543 embd_dev,
11544 mask,
11545 )?
11546 } else {
11547 self.mtp_head_forward_dev(
11548 e,
11549 mtp,
11550 e_tok,
11551 &d_seed,
11552 &mut *scratch,
11553 mtp_pos,
11554 embd_dev,
11555 mask,
11556 )?
11557 };
11558 let tok_d = if sampled {
11559 // FILTERED Gumbel-max: stats -> masked perturb -> argmax = one draw from
11560 // the filtered softmax (filters off => th=0, exact v1 semantics).
11561 if perturb_buf.is_none() {
11562 perturb_buf = Some(e.zeros(d_vocab.max(n_vocab))?);
11563 }
11564 let mut q_row = e.clone_dtod(&dl_d)?; // retained q (penalized when on)
11565 if pen_on {
11566 let h = pen_hist_d.as_ref().unwrap();
11567 let nh = h.len();
11568 e.penalize_logits(
11569 &mut q_row,
11570 h,
11571 nh,
11572 sp.penalty_repeat,
11573 sp.penalty_freq,
11574 sp.penalty_present,
11575 d_vocab,
11576 )?;
11577 }
11578 let rows0 = e.htod_i32(&[0])?;
11579 let (mut th_d, mut z_d, mut mx_d) =
11580 (e.zeros(1)?, e.zeros(1)?, e.zeros(1)?);
11581 e.filter_stats(
11582 &q_row, d_vocab, &rows0, &mut th_d, &mut z_d, &mut mx_d, d_vocab,
11583 1, sp_temp, sp.top_k, sp.top_p, sp.min_p,
11584 )?;
11585 let (th, z, mx) =
11586 (e.dtoh(&th_d)?[0], e.dtoh(&z_d)?[0], e.dtoh(&mx_d)?[0]);
11587 let pb = perturb_buf.as_mut().unwrap();
11588 e.gumbel_perturb_filtered(
11589 &q_row, pb, d_vocab, sp_seed, sctr, sp_temp, mx, th,
11590 )?;
11591 sctr += 1;
11592 draft_logits.push(q_row);
11593 draft_stats.push((mx, th, z));
11594 e.argmax_token_device(pb, d_vocab)?
11595 } else {
11596 e.argmax_token_device(&dl_d, d_vocab)?
11597 };
11598 let idx = e.dtoh_u32_one(&tok_d)?;
11599 // #87 SENTINEL TRAP (eager twin — see the graph arm). Extra diagnostics
11600 // here because the eager chain's operands are all readable: dl_d (the head
11601 // logits row) and d_seed (this step's h_seed) name the first-NaN buffer.
11602 if (idx as usize) >= d_vocab {
11603 let dl_h = e.dtoh(&dl_d)?;
11604 let dl_nan = dl_h.iter().filter(|v| v.is_nan()).count();
11605 let seed_h = if chain_heads {
11606 e.dtoh(chain_seeds.last().unwrap())?
11607 } else {
11608 e.dtoh(&d_seed)?
11609 };
11610 let seed_nan = seed_h.iter().filter(|v| v.is_nan()).count();
11611 return Err(format!(
11612 "draft(eager) argmax sentinel 0x{idx:08x} >= d_vocab {d_vocab} at \
11613 round {round} j={j} pos={pos}: head-logits NaN {dl_nan}/{d_vocab}, \
11614 step-seed NaN {seed_nan}/{n_embd} — refusing to dereference the \
11615 embed row (#87 trap)"
11616 )
11617 .into());
11618 }
11619 let d = match &mtp.d2t {
11620 Some(map) => map[idx as usize],
11621 None => idx,
11622 };
11623 if sampled {
11624 draft_idx.push(idx);
11625 }
11626 let draft_p = if p_min > 0.0
11627 || opti_fork
11628 .as_ref()
11629 .is_some_and(|fork| fork.controller.is_some())
11630 {
11631 let p_d = e.prob_of_token_device(&dl_d, &tok_d, d_vocab)?;
11632 Some(e.dtoh(&p_d)?[0])
11633 } else {
11634 None
11635 };
11636 if j == 0 {
11637 controller_draft_prob = draft_p;
11638 }
11639 if let Some(p) = draft_p.filter(|_| p_min > 0.0) {
11640 if p < p_min && (j > 0 || (pmin0 && base0 == 1)) {
11641 break;
11642 }
11643 }
11644 draft.push(d);
11645 if chain_heads {
11646 chain_tokens.push(d);
11647 chain_seeds.push(h_nextn);
11648 } else {
11649 e_tok = d;
11650 d_seed = h_nextn;
11651 }
11652 // speculative advance; a chain the grammar can no longer follow (EOS
11653 // proposed) ends here — the prefix already proposed still rides verify.
11654 if dmask_live
11655 && !constraint
11656 .as_deref_mut()
11657 .unwrap()
11658 .draft_advance(d)
11659 .map_err(|e2| format!("constraint: {e2}"))?
11660 {
11661 break;
11662 }
11663 }
11664 if !chain_heads
11665 && opti_fork
11666 .as_ref()
11667 .is_some_and(|fork| fork.controller.is_some())
11668 {
11669 controller_eager_state = Some((e_tok, d_seed));
11670 }
11671 }
11672 }
11673 let k_round = draft.len();
11674 if let Some(p) = pipe {
11675 p.draft_end(round);
11676 }
11677 drop(pipe_draft);
11678
11679 ph_mark(&mut ph_draft, phase_on);
11680 // --- 2. VERIFY: one batched target forward. With a pending bonus, it rides as col 0
11681 // (committing its KV/recur inside the SAME weight read); drafts follow. ---
11682 let verify_tokens: Vec<u32> = match pending {
11683 Some(b) => {
11684 let mut v = Vec::with_capacity(k_round + 1);
11685 v.push(b);
11686 v.extend_from_slice(&draft);
11687 v
11688 }
11689 None => draft.clone(),
11690 };
11691 let base = if pending.is_some() { 1 } else { 0 };
11692 // ckpt (REPLAY-FREE partial accept): retain per-layer state-rebuild inputs alongside
11693 // the verify. Pure buffer keep-alives + dtod clones — kernel work is unchanged.
11694 let mut ckpt = if let Some(ticket) = current_opti.as_mut() {
11695 Some(ticket.take_ckpt())
11696 } else if spec_replay {
11697 None
11698 } else {
11699 Some(VerifyCkpt::new(self.layers.len()))
11700 };
11701 let controller_can_probe = base == 1
11702 && k_round == 1
11703 && out.len().saturating_add(2) < max_new
11704 && controller_draft_prob.is_some()
11705 && opti_fork
11706 .as_ref()
11707 .and_then(|fork| fork.controller.as_ref())
11708 .is_some_and(|policy| !policy.breaker_tripped);
11709 let mut successor_attempt: Option<OptiControllerTicket> = None;
11710 let mut rejected_probe: Option<(f32, u32)> = None;
11711 let mut controller_prepared: Option<OptiControllerPrepared> = None;
11712 if controller_can_probe {
11713 // Prepare d2/q and, on admission, d3 before either current verify half is
11714 // issued. N stage 0 can then be followed immediately by N+1 stage 0; once N's
11715 // boundary fires, those dev0 launches overlap N stage 1 on dev1. Preparing on
11716 // the primary stream after N stage 1 would serialize the supposed pipeline.
11717 let eager_pos = scratch.kv.len + 1;
11718 let (optimistic_pending, pending_probability) = self.opti_controller_draft_step(
11719 e,
11720 mtp,
11721 &mut dctx,
11722 &mut *scratch,
11723 d_vocab,
11724 &mut controller_eager_state,
11725 eager_pos,
11726 embd_dev,
11727 )?;
11728 let first_probability = controller_draft_prob
11729 .ok_or("optipipe controller probe lost first-token probability")?;
11730 let q_proxy = first_probability * pending_probability;
11731 OPTI_GATE_CHECKS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
11732 OPTI_SHADOW_DRAFT_TOKENS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
11733 let admitted = opti_fork
11734 .as_ref()
11735 .and_then(|fork| fork.controller.as_ref())
11736 .ok_or("optipipe controller policy disappeared")?
11737 .admit(q_proxy);
11738 if admitted {
11739 OPTI_GATE_ADMITS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
11740 let eager_pos = scratch.kv.len + 1;
11741 let (optimistic_draft, optimistic_draft_probability) = self
11742 .opti_controller_draft_step(
11743 e,
11744 mtp,
11745 &mut dctx,
11746 &mut *scratch,
11747 d_vocab,
11748 &mut controller_eager_state,
11749 eager_pos,
11750 embd_dev,
11751 )?;
11752 OPTI_SHADOW_DRAFT_TOKENS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
11753 let eager_seed = controller_eager_state.take().map(|(token, seed)| {
11754 debug_assert_eq!(token, optimistic_draft);
11755 seed
11756 });
11757 controller_prepared = Some(OptiControllerPrepared {
11758 verify_tokens: [optimistic_pending, optimistic_draft],
11759 draft_prob: optimistic_draft_probability,
11760 eager_seed,
11761 q_proxy,
11762 scratch_len: scratch.kv.len,
11763 });
11764 } else {
11765 OPTI_GATE_REJECTS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
11766 OPTI_WASTED_DRAFT_TOKENS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
11767 rejected_probe = Some((q_proxy, optimistic_pending));
11768 eprintln!(
11769 "[opti-controller] reject q={q_proxy:.6} threshold={:.3}",
11770 opti_fork
11771 .as_ref()
11772 .and_then(|fork| fork.controller.as_ref())
11773 .expect("controller policy")
11774 .threshold,
11775 );
11776 }
11777 }
11778 let fork_attempt = match fork_generation.take() {
11779 Some(generation) if base == 1 && k_round == 1 => Some(generation),
11780 Some(generation) => {
11781 opti_fork
11782 .as_mut()
11783 .expect("fork generation without fork state")
11784 .retire(generation)?;
11785 None
11786 }
11787 None => None,
11788 };
11789 let (tlogits_d, vx) = if let Some(p) = pipe {
11790 self.decode_step_t_core_pipelined(
11791 e,
11792 &verify_tokens,
11793 pos,
11794 &mut *cache,
11795 embd_dev,
11796 ckpt.as_mut(),
11797 p,
11798 round,
11799 )?
11800 } else if controller_can_probe {
11801 let fence = opti_fork
11802 .as_ref()
11803 .ok_or("optipipe controller probe lost fork state")?
11804 .fence;
11805 let boundary = match current_opti.as_mut() {
11806 Some(ticket) => ticket.take_boundary(),
11807 None => self.verify_stage0_issue(
11808 e,
11809 &verify_tokens,
11810 pos,
11811 &mut *cache,
11812 embd_dev,
11813 ckpt.as_mut(),
11814 None,
11815 &fence,
11816 Some(true),
11817 None,
11818 )?,
11819 };
11820 if let Some(prepared) = controller_prepared.take() {
11821 let generation = {
11822 let fork = opti_fork
11823 .as_mut()
11824 .ok_or("optipipe controller admission lost fork state")?;
11825 let generation = fork.reserve_successor()?;
11826 let rt = fork.rt;
11827 let snapshot_fence = fork.fence;
11828 opti_snapshot_one_stage_owned_into(
11829 e,
11830 cache,
11831 rt,
11832 &snapshot_fence,
11833 0,
11834 fork.successor_snapshot_mut(),
11835 )?;
11836 generation
11837 };
11838 let mut successor_ckpt = VerifyCkpt::new(self.layers.len());
11839 let successor_boundary = self.verify_stage0_issue(
11840 e,
11841 &prepared.verify_tokens,
11842 pos + verify_tokens.len(),
11843 &mut *cache,
11844 embd_dev,
11845 Some(&mut successor_ckpt),
11846 None,
11847 &fence,
11848 Some(false),
11849 None,
11850 )?;
11851 OPTI_FORK_ATTEMPTS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
11852 let fork = opti_fork
11853 .as_ref()
11854 .ok_or("optipipe controller ticket lost fork state")?;
11855 successor_attempt = Some(fork.controller_ticket(
11856 generation,
11857 successor_boundary,
11858 successor_ckpt,
11859 prepared.verify_tokens,
11860 prepared.draft_prob,
11861 prepared.eager_seed,
11862 prepared.q_proxy,
11863 prepared.scratch_len,
11864 ));
11865 eprintln!(
11866 "[opti-controller] issue generation={} q={:.6} threshold={:.3} \
11867 verify={:?}",
11868 generation.id,
11869 prepared.q_proxy,
11870 fork.controller.expect("controller policy").threshold,
11871 prepared.verify_tokens,
11872 );
11873 }
11874 let result = self.verify_stage1_finish(
11875 e,
11876 boundary,
11877 &mut *cache,
11878 ckpt.as_mut(),
11879 None,
11880 &fence,
11881 successor_attempt.is_none(),
11882 )?;
11883 if let Some(ticket) = current_opti.as_mut() {
11884 ticket.settle();
11885 }
11886 if successor_attempt.is_some() {
11887 let fork = opti_fork
11888 .as_mut()
11889 .ok_or("optipipe successor snapshot lost fork state")?;
11890 let rt = fork.rt;
11891 let snapshot_fence = fork.fence;
11892 opti_snapshot_one_stage_owned_into(
11893 e,
11894 cache,
11895 rt,
11896 &snapshot_fence,
11897 1,
11898 fork.successor_snapshot_mut(),
11899 )?;
11900 // Publish N only after both independent successor-state queues are complete.
11901 fork.rt.publish_to(1, &e.stream())?;
11902 }
11903 result
11904 } else if let Some(ticket) = current_opti.as_mut() {
11905 let fork = opti_fork
11906 .as_mut()
11907 .ok_or("optipipe carried controller ticket lost fork state")?;
11908 let boundary = ticket.take_boundary();
11909 let result = self.verify_stage1_finish(
11910 e,
11911 boundary,
11912 &mut *cache,
11913 ckpt.as_mut(),
11914 None,
11915 &fork.fence,
11916 true,
11917 )?;
11918 ticket.settle();
11919 result
11920 } else if let Some(generation) = fork_attempt {
11921 let fork = opti_fork
11922 .as_mut()
11923 .expect("fork generation without fork state");
11924 fork.capture_seed(e, generation, &h_seed_buf, &fill_prev, scratch.kv.len)?;
11925 let action = fork.mode.action(generation.id);
11926 let boundary = self.verify_stage0_issue(
11927 e,
11928 &verify_tokens,
11929 pos,
11930 &mut *cache,
11931 embd_dev,
11932 ckpt.as_mut(),
11933 None,
11934 &fork.fence,
11935 Some(true),
11936 None,
11937 )?;
11938 OPTI_FORK_ATTEMPTS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
11939 let mut ticket = fork.ticket(generation, boundary);
11940 if action == OptiForkAction::Abort {
11941 return Err(format!(
11942 "optipipe forced abort with generation {} stage0 in flight",
11943 generation.id,
11944 )
11945 .into());
11946 }
11947 fork.reconcile(
11948 e,
11949 &mut *cache,
11950 &mut *scratch,
11951 &snap,
11952 &mut h_seed_buf,
11953 &mut fill_prev,
11954 generation,
11955 action,
11956 verify_tokens[0],
11957 )?;
11958 let result = if action == OptiForkAction::Hit {
11959 let boundary = ticket.take_boundary();
11960 self.verify_stage1_finish(
11961 e,
11962 boundary,
11963 &mut *cache,
11964 ckpt.as_mut(),
11965 None,
11966 &fork.fence,
11967 true,
11968 )?
11969 } else {
11970 // The optimistic boundary slot has no reader. Re-run the unchanged serial
11971 // verify only after E_restart published the restored stage-0 state.
11972 self.decode_step_t_core(
11973 e,
11974 &verify_tokens,
11975 pos,
11976 &mut *cache,
11977 embd_dev,
11978 ckpt.as_mut(),
11979 )?
11980 };
11981 ticket.settle();
11982 debug_assert_eq!(ticket.generation, generation);
11983 fork.retire(generation)?;
11984 result
11985 } else {
11986 // The serial verify every non-fork round takes — the MTP route's
11987 // verify-graph door. The pool is None unless MEMRA_SPEC_VERIFY_GRAPH armed
11988 // a pool above, and then the walk replays the captured trunk instead of
11989 // re-issuing it launch by launch.
11990 let vg_round = if verify_tokens.len() <= vg_t_cap {
11991 vg_guard.as_mut().and_then(|g| g.as_mut())
11992 } else {
11993 if let Some(g) = vg_guard.as_mut().and_then(|g| g.as_mut()) {
11994 // The commit reads this flag to pick its arm; a round that declines
11995 // the pool must not inherit a stale `true` from the round before it.
11996 g.round_slab = false;
11997 }
11998 None
11999 };
12000 self.decode_step_t_core_vg(
12001 e,
12002 &verify_tokens,
12003 pos,
12004 &mut *cache,
12005 embd_dev,
12006 ckpt.as_mut(),
12007 vg_round,
12008 )?
12009 };
12010 let pipe_accept = match pipe {
12011 Some(p) => Some(p.accept_begin(round)?),
12012 None => None,
12013 };
12014
12015 ph_mark(&mut ph_verify, phase_on);
12016 // --- 3. GREEDY ACCEPT (walk prefix, stop at first mismatch) ---
12017 // DEVICE-ARGMAX ACCEPT: argmax every verify column ON DEVICE (same 2-pass kernels +
12018 // smallest-index tie-break as host argmax, argmax_gate-validated) and read back ONE
12019 // [T] u32 — replaces the T x n_vocab f32 dtoh + T host argmaxes per round.
12020 // t_pred[j] = target's greedy prediction for the slot after draft[j-1] (j>=1) or after
12021 // last_token (j==0). With a pending bonus, col 0 IS the prediction after last_token
12022 // (== the bonus), so every index shifts by `base` and last_pred is unused.
12023 let t_v = verify_tokens.len();
12024 let mut preds: Vec<u32> = Vec::new();
12025 if !sampled {
12026 for j in 0..t_v {
12027 e.argmax_token_device_col(&tlogits_d, j, n_vocab, &mut preds_d, j)?;
12028 }
12029 preds = e.dtoh_u32(&preds_d)?; // <- the verify-GPU wait lands here
12030 // #87 SENTINEL TRAP, verify side: a sentinel pred becomes the round's bonus =
12031 // next round's last_token = the next chain's embed lookup. Catch it at the
12032 // source with the column named — an all-NaN VERIFY column implicates the
12033 // stage-split trunk (decode_step_t_core_ppn), not the draft head.
12034 if let Some(bad) = preds[..t_v].iter().position(|&p| (p as usize) >= n_vocab) {
12035 let col = &tlogits_d.slice(bad * n_vocab..(bad + 1) * n_vocab);
12036 let mut probe = e.zeros(n_vocab)?;
12037 e.copy_view_into(&mut probe, 0, col, n_vocab)?;
12038 let col_h = e.dtoh(&probe)?;
12039 let col_nan = col_h.iter().filter(|v| v.is_nan()).count();
12040 return Err(format!(
12041 "verify argmax sentinel 0x{:08x} >= n_vocab {n_vocab} at round {round} \
12042 col {bad}/{t_v} pos={pos}: verify-logits col NaN {col_nan}/{n_vocab} \
12043 — the stage-split verify produced a poisoned column (#87 trap)",
12044 preds[bad]
12045 )
12046 .into());
12047 }
12048 }
12049 ph_mark(&mut ph_wait, phase_on);
12050 let t_pred = |j: usize| -> u32 {
12051 if j == 0 && base == 0 {
12052 last_pred
12053 } else {
12054 // GREEDY-ONLY: `preds` is filled under `if !sampled` above. The debug print
12055 // used to call this from the sampled arm and panicked the worker; it now goes
12056 // through `debug_t_pred0`. Keep the strict index here — in the greedy walk an
12057 // out-of-range pred is a real bug, not something to paper over.
12058 debug_assert!(
12059 !sampled,
12060 "t_pred is greedy-only: `preds` is empty in the sampled arm"
12061 );
12062 preds[base + j - 1]
12063 }
12064 };
12065 let mut devacc_seeded = false;
12066 let mut devacc_acc: Option<CudaSlice<u32>> = None;
12067 let (n_acc, bonus) = if !sampled {
12068 // ROUND-STREAM stage (a) (MEMRA_SPEC_DEVACC=1 opt-in): the walk runs ON DEVICE
12069 // (spec_accept_greedy, verbatim rule) and the host reads back 8B (n_acc, bonus)
12070 // instead of the [T] preds. Same sync count — machinery for stages (b)/(c),
12071 // gated on token identity vs the host walk (the arms below are bit-equal rules).
12072 if crate::spec::spec_devacc() && k_round > 0 && !spec_replay && constraint.is_none()
12073 {
12074 let draft_d = e.htod_u32_v(&draft)?;
12075 let mut acc_out = e.alloc_u32_zeroed(2)?;
12076 e.spec_accept_greedy(
12077 &preds_d,
12078 &draft_d,
12079 last_pred,
12080 base,
12081 k_round,
12082 &mut acc_out,
12083 )?;
12084 devacc_acc = Some(acc_out.clone());
12085 // stage (b): next-round seed gathered ON DEVICE from acc_out before the host
12086 // ever reads n_acc (j=base+n_acc -> vx col j-1; j==0 -> fill_prev). The three
12087 // non-replay commit arms skip their host-offset seed copies (guarded below);
12088 // the legacy spec_replay arm keeps its own rx-based seeding (excluded here).
12089 // NOTE: fill_prev is NOT updated here — the commit arms' TRUE-HIDDEN
12090 // REFRESH reads the OLD fill_prev (predecessor of this round's verify batch);
12091 // the update lands after the arms (devacc_seeded guard below).
12092 e.spec_seed_gather(&vx, &fill_prev, &acc_out, &mut h_seed_buf, base, n_embd)?;
12093 // 3a: KV lens roll back on device (len = saved + base + n_acc, all arms'
12094 // unified rule; full accept rewrites the verify-left value). Host mirrors
12095 // update after the readback; commit_verified_prefix skips its len_d writes.
12096 if let Some(successor) = successor_attempt.as_ref() {
12097 opti_fork
12098 .as_mut()
12099 .ok_or("optipipe successor reconcile lost fork state")?
12100 .queue_actual_reconcile(
12101 e,
12102 &snap,
12103 &acc_out,
12104 successor.verify_tokens[0],
12105 base,
12106 )?;
12107 } else if let Some(ptrs) = &kv_len_ptrs {
12108 let saved: Vec<i32> = (0..self.layers.len())
12109 .map(|il| snap.kv_len[il].map(|v| v as i32).unwrap_or(0))
12110 .collect();
12111 let saved_d = e.htod_i32(&saved)?;
12112 e.spec_rollback_kv(ptrs, &saved_d, &acc_out, base, self.layers.len())?;
12113 }
12114 devacc_seeded = true;
12115 let ab = e.dtoh_u32(&acc_out)?;
12116 (ab[0] as usize, ab[1])
12117 } else {
12118 let mut n_acc = 0usize;
12119 for j in 0..k_round {
12120 if t_pred(j) == draft[j] {
12121 n_acc += 1;
12122 } else {
12123 break;
12124 }
12125 }
12126 // bonus = target's own token at the first non-accepted slot. n_acc in 0..=k; t_pred
12127 // is defined for j in 0..=k (j==0 -> last_logits, j>=1 -> col j-1, last col = k-1).
12128 (n_acc, t_pred(n_acc))
12129 }
12130 } else {
12131 // --- SAMPLED ACCEPT (rejection sampling): u_j < p_j(x_j)/q_j(x_j) walk ---
12132 if col_buf.is_none() {
12133 col_buf = Some(e.zeros(n_vocab)?);
12134 }
12135 // FILTERED p_j: per-verify-col stats (one batched filter_stats call), then the
12136 // filtered gather. j==0&&base==0 reads last_col (its own stats row appended).
12137 let mut pj = vec![0f32; k_round.max(1)];
12138 let mut col_stats: Vec<(f32, f32, f32)> = Vec::new(); // (max, th, z) per verify col used
12139 if k_round > 0 {
12140 let mut ids: Vec<u32> = Vec::new();
12141 let mut rows: Vec<i32> = Vec::new();
12142 for j in 0..k_round {
12143 if j > 0 || base == 1 {
12144 ids.push(draft[j]);
12145 rows.push((base + j) as i32 - 1);
12146 }
12147 }
12148 if !ids.is_empty() {
12149 let nr = rows.len();
12150 // penalties: materialize the used columns into one contiguous penalized
12151 // buffer (rows remapped 0..nr) so stats+gathers see the penalized p.
12152 // penalties: materialize used columns contiguously, penalize all rows in
12153 // one launch, and point stats+gathers at the penalized buffer (rows 0..nr).
12154 let p_rows: Vec<i32> = if pen_on {
12155 (0..nr as i32).collect()
12156 } else {
12157 rows.clone()
12158 };
12159 if pen_on {
12160 if pcol_buf.as_ref().map(|b| b.len()).unwrap_or(0) < nr * n_vocab {
12161 pcol_buf = Some(e.zeros(nr * n_vocab)?);
12162 }
12163 let pc = pcol_buf.as_mut().unwrap();
12164 for (i2, &r) in rows.iter().enumerate() {
12165 let c = r as usize;
12166 e.copy_view_into(
12167 pc,
12168 i2 * n_vocab,
12169 &tlogits_d.slice(c * n_vocab..(c + 1) * n_vocab),
12170 n_vocab,
12171 )?;
12172 }
12173 let h = pen_hist_d.as_ref().unwrap();
12174 let nh = h.len();
12175 e.penalize_logits_rows(
12176 pc,
12177 h,
12178 nh,
12179 sp.penalty_repeat,
12180 sp.penalty_freq,
12181 sp.penalty_present,
12182 n_vocab,
12183 nr,
12184 )?;
12185 }
12186 let p_src: &CudaSlice<f32> = if pen_on {
12187 pcol_buf.as_ref().unwrap()
12188 } else {
12189 &tlogits_d
12190 };
12191 let rowsd = e.htod_i32(&p_rows)?;
12192 let (mut th_d, mut z_d, mut mx_d) =
12193 (e.zeros(nr)?, e.zeros(nr)?, e.zeros(nr)?);
12194 e.filter_stats(
12195 p_src, n_vocab, &rowsd, &mut th_d, &mut z_d, &mut mx_d, n_vocab, nr,
12196 sp_temp, sp.top_k, sp.top_p, sp.min_p,
12197 )?;
12198 let idsd = e.htod_u32_v(&ids)?;
12199 let mut outd = e.zeros(nr)?;
12200 e.softmax_gather_filtered(
12201 p_src, n_vocab, &idsd, &rowsd, &th_d, &z_d, &mut outd, n_vocab, nr,
12202 sp_temp,
12203 )?;
12204 let outv = e.dtoh(&outd)?;
12205 let (thv, zv, mxv) = (e.dtoh(&th_d)?, e.dtoh(&z_d)?, e.dtoh(&mx_d)?);
12206 let mut oi = 0usize;
12207 for j in 0..k_round {
12208 if j > 0 || base == 1 {
12209 pj[j] = outv[oi];
12210 oi += 1;
12211 }
12212 }
12213 col_stats = (0..nr).map(|i| (mxv[i], thv[i], zv[i])).collect();
12214 }
12215 if base == 0 {
12216 let lc: &CudaSlice<f32> = if pen_on {
12217 if col_buf.is_none() {
12218 col_buf = Some(e.zeros(n_vocab)?);
12219 }
12220 let cb = col_buf.as_mut().unwrap();
12221 e.copy_into(
12222 cb,
12223 0,
12224 last_col_logits
12225 .as_ref()
12226 .expect("sampled: last_col_logits unset"),
12227 n_vocab,
12228 )?;
12229 let h = pen_hist_d.as_ref().unwrap();
12230 let nh = h.len();
12231 e.penalize_logits(
12232 cb,
12233 h,
12234 nh,
12235 sp.penalty_repeat,
12236 sp.penalty_freq,
12237 sp.penalty_present,
12238 n_vocab,
12239 )?;
12240 col_buf.as_ref().unwrap()
12241 } else {
12242 last_col_logits
12243 .as_ref()
12244 .expect("sampled: last_col_logits unset")
12245 };
12246 let rows0 = e.htod_i32(&[0])?;
12247 let (mut th_d, mut z_d, mut mx_d) = (e.zeros(1)?, e.zeros(1)?, e.zeros(1)?);
12248 e.filter_stats(
12249 lc, n_vocab, &rows0, &mut th_d, &mut z_d, &mut mx_d, n_vocab, 1,
12250 sp_temp, sp.top_k, sp.top_p, sp.min_p,
12251 )?;
12252 let idsd = e.htod_u32_v(&[draft[0]])?;
12253 let mut outd = e.zeros(1)?;
12254 e.softmax_gather_filtered(
12255 lc, n_vocab, &idsd, &rows0, &th_d, &z_d, &mut outd, n_vocab, 1, sp_temp,
12256 )?;
12257 pj[0] = e.dtoh(&outd)?[0];
12258 last_col_stats =
12259 Some((e.dtoh(&mx_d)?[0], e.dtoh(&th_d)?[0], e.dtoh(&z_d)?[0]));
12260 }
12261 }
12262 // q source: the graph arm retained the head logits in the persistent q_slots;
12263 // the eager arm in per-round draft_logits clones. Same raw-logit values either way.
12264 // FILTERED q_j: stats from draft_stats (eager pushes in-chain; the graph arm
12265 // computes them post-replay — graph engages only filter/penalty-free, so the
12266 // stats degenerate to th=0/full-Z there, keeping ONE accept path).
12267 let q_bufs: &[CudaSlice<f32>] = if dctx.graph_s.is_some() {
12268 &dctx.q_slots
12269 } else {
12270 &draft_logits
12271 };
12272 let mut n_acc = 0usize;
12273 for j in 0..k_round {
12274 let (qmx, qth, qz) = draft_stats[j];
12275 let idsd = e.htod_u32_v(&[draft_idx[j]])?;
12276 let rowsd = e.htod_i32(&[0])?;
12277 let thd = e.htod(&[qth])?;
12278 let zd = e.htod(&[qz])?;
12279 let _ = qmx;
12280 let mut outd = e.zeros(1)?;
12281 e.softmax_gather_filtered(
12282 &q_bufs[j], d_vocab, &idsd, &rowsd, &thd, &zd, &mut outd, d_vocab, 1,
12283 sp_temp,
12284 )?;
12285 let qj = e.dtoh(&outd)?[0];
12286 let u = host_u01(sp_seed, uctr);
12287 uctr += 1;
12288 let accept = (u as f64) * (qj as f64) < pj[j] as f64;
12289 // SKEY PROBE: q == 0 for the token the draft actually proposed is the
12290 // exactness signature (see `skey_probe`). Impossible when the draft was
12291 // drawn from the same filtered distribution the verify reconstructs here;
12292 // `u * 0 < p` makes it an UNCONDITIONAL accept whenever p > 0.
12293 if skey_probe() && qj == 0.0 {
12294 eprintln!(
12295 "[skey] EXACTNESS q=0 round={round} j={j} draft_tok={} \
12296 draft_idx={} p={:e} u={u} accepted={} th_z={:?}",
12297 draft[j], draft_idx[j], pj[j], accept as u8, draft_stats[j],
12298 );
12299 }
12300 if accept {
12301 n_acc += 1;
12302 } else {
12303 break;
12304 }
12305 }
12306 let bonus = if n_acc == k_round {
12307 // FULL ACCEPT: bonus ~ FILTERED softmax at the last verify column.
12308 let col = base + k_round - 1;
12309 let cb = col_buf.as_mut().unwrap();
12310 e.copy_view_into(
12311 cb,
12312 0,
12313 &tlogits_d.slice(col * n_vocab..(col + 1) * n_vocab),
12314 n_vocab,
12315 )?;
12316 if pen_on {
12317 let h = pen_hist_d.as_ref().unwrap();
12318 let nh = h.len();
12319 e.penalize_logits(
12320 cb,
12321 h,
12322 nh,
12323 sp.penalty_repeat,
12324 sp.penalty_freq,
12325 sp.penalty_present,
12326 n_vocab,
12327 )?;
12328 }
12329 if perturb_buf.is_none() {
12330 perturb_buf = Some(e.zeros(d_vocab.max(n_vocab))?);
12331 }
12332 // STATS MUST COME FROM THIS COLUMN (bug fix 2026-08-05, lane/sampler-
12333 // truncation-fix; receipts research/sampfix-20260805/). The old code reused
12334 // `col_stats.last()` here, which is ALWAYS the wrong row: the gathered set
12335 // covers verify columns 0..=(base+k_round-2) (rows pushed as base+j-1), while
12336 // the full-accept bonus samples column base+k_round-1 — exactly ONE PAST the
12337 // last gathered column, in both base arms. `th` is a threshold in e-units of
12338 // its OWN row's max, so feeding a neighbour's (row_max, th) into
12339 // gumbel_perturb_filtered mis-scales every e0 = exp((x-row_max)/T). When the
12340 // donor column's peak is higher by more than T*ln(1/th), EVERY id fails
12341 // `e0 >= th`, the whole perturbed row becomes -3.4e38, and the 2-pass argmax
12342 // falls through to its smallest-index tie-break => token id 0 ("!") spliced
12343 // mid-word. Fragility is ordered by how large th is: min_p pins th = min_p
12344 // (0.05 => trigger at delta > 2.4 at T=0.8, fires constantly), top_p's
12345 // mass-boundary th is smaller, top_k's k-th-largest th smaller still — which
12346 // is why the head-to-head matrix saw min_p and top_p corrupt while top_k-only
12347 // stayed clean. The pure-temp default regime is immune (th == 0 masks nothing,
12348 // and row_max is unused once nothing is masked), so this fix is a byte-level
12349 // no-op for the untruncated serve default. One extra one-block filter_stats
12350 // per full-accept round is the whole cost.
12351 let (mx, th) = {
12352 let rows0 = e.htod_i32(&[0])?;
12353 let (mut th_d, mut z_d, mut mx_d) = (e.zeros(1)?, e.zeros(1)?, e.zeros(1)?);
12354 let cb0 = col_buf.as_ref().unwrap();
12355 e.filter_stats(
12356 cb0, n_vocab, &rows0, &mut th_d, &mut z_d, &mut mx_d, n_vocab, 1,
12357 sp_temp, sp.top_k, sp.top_p, sp.min_p,
12358 )?;
12359 (e.dtoh(&mx_d)?[0], e.dtoh(&th_d)?[0])
12360 };
12361 let pb = perturb_buf.as_mut().unwrap();
12362 let cb2 = col_buf.as_ref().unwrap();
12363 e.gumbel_perturb_filtered(cb2, pb, n_vocab, sp_seed, sctr, sp_temp, mx, th)?;
12364 sctr += 1;
12365 let td = e.argmax_token_device(pb, n_vocab)?;
12366 e.dtoh_u32_one(&td)?
12367 } else {
12368 // REJECT at n_acc: bonus ~ norm(max(0, softmax_T(p) - softmax_T(q))).
12369 let cb = col_buf.as_mut().unwrap();
12370 if n_acc > 0 || base == 1 {
12371 let col = base + n_acc - 1;
12372 e.copy_view_into(
12373 cb,
12374 0,
12375 &tlogits_d.slice(col * n_vocab..(col + 1) * n_vocab),
12376 n_vocab,
12377 )?;
12378 } else {
12379 let lc = last_col_logits.as_ref().unwrap();
12380 e.copy_into(cb, 0, lc, n_vocab)?;
12381 }
12382 if pen_on {
12383 let h = pen_hist_d.as_ref().unwrap();
12384 let nh = h.len();
12385 e.penalize_logits(
12386 cb,
12387 h,
12388 nh,
12389 sp.penalty_repeat,
12390 sp.penalty_freq,
12391 sp.penalty_present,
12392 n_vocab,
12393 )?;
12394 }
12395 let cb2 = col_buf.as_ref().unwrap();
12396 let sc = sctr;
12397 sctr += 1;
12398 // p-stats for the reject column: from col_stats when the col was gathered,
12399 // else (j==0&&base==0) from last_col_stats.
12400 let p_stats = if n_acc > 0 || base == 1 {
12401 // col index within the gathered set == number of gathered cols before n_acc
12402 let gi = if base == 1 { n_acc } else { n_acc - 1 };
12403 col_stats.get(gi).copied().unwrap_or_else(|| {
12404 (0.0, 0.0, 1.0) // unreachable: gathered cols always cover the reject slot
12405 })
12406 } else {
12407 last_col_stats.expect("sampled: last_col_stats unset at reject")
12408 };
12409 let q_stats = draft_stats[n_acc];
12410 if let Some(map) = &d2t_dev {
12411 if q_full_buf.is_none() {
12412 q_full_buf = Some(e.zeros(n_vocab)?);
12413 }
12414 let qf = q_full_buf.as_mut().unwrap();
12415 e.scatter_trim_logits(&q_bufs[n_acc], map, qf, d_vocab, n_vocab)?;
12416 let qf2 = q_full_buf.as_ref().unwrap();
12417 e.residual_sample_filtered(
12418 cb2,
12419 Some(qf2),
12420 n_vocab,
12421 sp_temp,
12422 sp_seed,
12423 sc,
12424 p_stats,
12425 q_stats,
12426 &mut sample_tok,
12427 )?;
12428 } else {
12429 e.residual_sample_filtered(
12430 cb2,
12431 Some(&q_bufs[n_acc]),
12432 n_vocab,
12433 sp_temp,
12434 sp_seed,
12435 sc,
12436 p_stats,
12437 q_stats,
12438 &mut sample_tok,
12439 )?;
12440 }
12441 e.dtoh_u32(&sample_tok)?[0]
12442 };
12443 (n_acc, bonus)
12444 };
12445 // --- 3b. GRAMMAR TRUNCATION (constrained spec, 2026-08-03): the grammar is
12446 // an extra rejection rule AFTER the exactness verify (the batched-verify-twins
12447 // ordering). Walk the accepted drafts through the grammar in commit order; the
12448 // first illegal token truncates acceptance at its slot, and that slot's emission
12449 // is recomputed as the MASKED argmax of the target's own verify column — token-
12450 // identical to constrained plain greedy decode (an unmasked argmax that is
12451 // grammar-legal IS the masked argmax: masking only removes competitors). The
12452 // column D2H (~1MB) is paid only when a cut fires — the tight-grammar cost,
12453 // measured in acceptance numbers, never hidden.
12454 let (n_acc, bonus) = match constraint.as_deref_mut() {
12455 None => (n_acc, bonus),
12456 Some(c) => {
12457 fn ce(e2: String) -> Box<dyn std::error::Error> {
12458 format!("constraint: {e2}").into()
12459 }
12460 let mut na = n_acc;
12461 let mut cut = false;
12462 for (j, &d) in draft.iter().enumerate().take(n_acc) {
12463 if c.is_allowed(d).map_err(ce)? {
12464 c.consume(d).map_err(ce)?;
12465 } else {
12466 na = j;
12467 cut = true;
12468 dm_cut_tokens += n_acc - j;
12469 break;
12470 }
12471 }
12472 if cut {
12473 dm_cuts += 1;
12474 }
12475 let mut bo = bonus;
12476 if cut || !c.is_allowed(bo).map_err(ce)? {
12477 let mut row = if na == 0 && base == 0 {
12478 init_logits_host
12479 .clone()
12480 .ok_or("constraint: init logits missing (round-0 cut)")?
12481 } else {
12482 e.dtoh_view(
12483 &tlogits_d.slice((base + na - 1) * n_vocab..(base + na) * n_vocab),
12484 )?
12485 };
12486 c.mask_logits(&mut row).map_err(ce)?;
12487 bo = argmax(&row) as u32;
12488 }
12489 c.consume(bo).map_err(ce)?;
12490 (na, bo)
12491 }
12492 };
12493 let mut successor_valid = false;
12494 if let Some((q_proxy, expected_d2)) = rejected_probe {
12495 let v_n = n_acc == 1 && bonus == expected_d2;
12496 eprintln!(
12497 "[opti-controller] shadow q={q_proxy:.6} admitted=false v_n={v_n} \
12498 expected_d2={expected_d2} n_acc={n_acc} bonus={bonus}",
12499 );
12500 }
12501 if let Some(successor) = successor_attempt.as_ref() {
12502 successor_valid = n_acc == 1 && bonus == successor.verify_tokens[0];
12503 let generation = successor.generation;
12504 let q_proxy = successor.q_proxy;
12505 let expected_pending = successor.verify_tokens[0];
12506 let resolution_ms = successor.issued_at.elapsed().as_secs_f64() * 1e3;
12507 let fork = opti_fork
12508 .as_mut()
12509 .ok_or("optipipe successor resolution lost fork state")?;
12510 fork.finish_actual_reconcile(e, &mut *cache, &snap, n_acc, base, successor_valid)?;
12511 if successor_valid {
12512 OPTI_FORK_HITS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
12513 } else {
12514 OPTI_FORK_MISSES.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
12515 OPTI_RECONCILES.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
12516 OPTI_WASTED_DRAFT_TOKENS.fetch_add(2, std::sync::atomic::Ordering::Relaxed);
12517 }
12518 let breaker_tripped = fork
12519 .controller
12520 .as_mut()
12521 .expect("controller policy")
12522 .resolve(successor_valid);
12523 if breaker_tripped {
12524 OPTI_BREAKER_TRIPS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
12525 }
12526 eprintln!(
12527 "[opti-controller] resolve generation={} hit={} q={q_proxy:.6} \
12528 expected_pending={expected_pending} n_acc={n_acc} bonus={bonus} \
12529 resolution_ms={resolution_ms:.3} reconcile={} breaker={}",
12530 generation.id, successor_valid, !successor_valid, breaker_tripped,
12531 );
12532 if !successor_valid {
12533 let mut successor = successor_attempt
12534 .take()
12535 .expect("controller successor disappeared on miss");
12536 successor.settle();
12537 fork.retire(generation)?;
12538 }
12539 }
12540 total_drafted += k_round;
12541 total_accepted += n_acc;
12542 if let Some(t) = sess_telem {
12543 // Greedy, rejection-sampling, and grammar truncation all converge here after
12544 // the accept decision is already on host. Fixed-size relaxed atomics only.
12545 t.record_round(k_round, n_acc);
12546 }
12547 if spec_stats {
12548 st_len_hist[k_round] += 1;
12549 for j in 0..k_round {
12550 st_drafted[j] += 1;
12551 }
12552 for j in 0..n_acc {
12553 st_accepted[j] += 1;
12554 }
12555 if n_acc == k_round {
12556 st_full += 1;
12557 }
12558 }
12559
12560 if debug_spec {
12561 eprintln!(
12562 "[R{round}] pos={pos} out_len={} last_tok={last_token} draft={draft:?} n_acc={n_acc} bonus={bonus} t_pred0={}",
12563 out.len(),
12564 // NOT `t_pred(0)`: `preds` is filled only under `if !sampled` above, so on a
12565 // sampled request round >= 1 (base == 1) indexed an EMPTY vector and PANICKED
12566 // the GPU worker thread — a debug flag that killed the exact regime you would
12567 // set it to investigate. See `debug_t_pred0`.
12568 debug_t_pred0(sampled, base, last_pred, &preds)
12569 );
12570 }
12571
12572 // --- 4. COMMIT: draft[0..n_acc] then bonus (n_acc + 1 tokens) ---
12573 let commit_started = std::time::Instant::now();
12574 // SESSION MODE: every accepted column is already in the CACHE — `out` must carry all
12575 // of them (overshoot past max_new included) or `committed` under-counts the cache rows
12576 // and the next turn's continuation seeds one token off (gate-caught 2026-07-05). The
12577 // single-shot path keeps the cap (its caller truncates + drops the cache anyway).
12578 for j in 0..n_acc {
12579 if !session_mode && out.len() >= max_new {
12580 break;
12581 }
12582 out.push(draft[j]);
12583 }
12584 if pen_on {
12585 pen_hist.extend_from_slice(&draft[0..n_acc]);
12586 pen_hist.push(bonus);
12587 }
12588 let bonus_emitted = session_mode || out.len() < max_new;
12589 if bonus_emitted {
12590 out.push(bonus);
12591 }
12592 last_token = bonus;
12593
12594 // --- 5. ROLLBACK + advance (§C) ---
12595 if n_acc == k_round && !spec_replay {
12596 // FULL ACCEPT, BONUS FOLD: all verify columns (pending? + drafts) are committed in
12597 // cache; the NEW bonus stays PENDING for the next round's verify batch — NO extra
12598 // T=1 trunk pass. The next draft chain seeds from the MTP block's h_nextn at the
12599 // bonus position: one MTP-block pass (~1/33 trunk cost) replaces the trunk read.
12600 // last_pred is dead in the pending path (t_pred reads verify col 0).
12601 //
12602 // PERSISTENT DRAFT KV, full-accept fill: the chain covered last_token +
12603 // draft[0..k_round-2] as INPUTS (slots P..P'-2); draft[k_round-1] (slot P'-1) was
12604 // only ever an output, so its entry is MISSING. Fill it from vh_seed — its EXACT
12605 // trunk hidden (the last verify column). set_len first: a p-min break may have
12606 // left one extra chain append at that slot. Partial accepts need NO fill (the
12607 // chain already covered every accepted position; round-start set_len truncates).
12608 let mut vh_seed = e.zeros(n_embd)?;
12609 e.copy_view_into(
12610 &mut vh_seed,
12611 0,
12612 &vx.slice((t_v - 1) * n_embd..t_v * n_embd),
12613 n_embd,
12614 )?;
12615 if refresh {
12616 // TRUE-HIDDEN REFRESH (2026-07-03, the HANDOVER-listed acceptance lever):
12617 // overwrite ALL committed positions' scratch entries with K/V from their EXACT
12618 // verify hiddens — the reference engine's mtp_update fills from true hiddens;
12619 // the full stack (vx) is already resident from the verify. Replaces both the
12620 // chain-approximate entries AND the old last-token-only fill. Acceptance-only
12621 // (draft attention quality); exactness stays the verify's job.
12622 scratch.set_len(e, pos)?;
12623 // PREDECESSOR pairing: row i gets vx[i-1]; row 0 the carried fill_prev
12624 // (hidden of the last committed row before this verify batch).
12625 let mut vxs = e.zeros(t_v * n_embd)?;
12626 e.copy_into(&mut vxs, 0, &fill_prev, n_embd)?;
12627 if t_v > 1 {
12628 e.copy_view_into(
12629 &mut vxs,
12630 n_embd,
12631 &vx.slice(0..(t_v - 1) * n_embd),
12632 (t_v - 1) * n_embd,
12633 )?;
12634 }
12635 self.mtp_kv_fill_all(e, &verify_tokens, &vxs, pos, &mut *scratch, embd_dev)?;
12636 } else {
12637 scratch.set_len(e, pos + base + k_round - 1)?;
12638 // predecessor of the last draft = verify col t_v-2 (or fill_prev at t_v==1)
12639 let mut hp = e.zeros(n_embd)?;
12640 if t_v >= 2 {
12641 e.copy_view_into(
12642 &mut hp,
12643 0,
12644 &vx.slice((t_v - 2) * n_embd..(t_v - 1) * n_embd),
12645 n_embd,
12646 )?;
12647 } else {
12648 e.copy_into(&mut hp, 0, &fill_prev, n_embd)?;
12649 }
12650 self.mtp_kv_fill_all(
12651 e,
12652 &[draft[k_round - 1]],
12653 &hp,
12654 pos + base + k_round - 1,
12655 &mut *scratch,
12656 embd_dev,
12657 )?;
12658 }
12659 // REFERENCE SEEDING: no pseudo pass — the next chain's step 0 IS the
12660 // reference's (id_last, h_prev) draft row; it appends the bonus's scratch
12661 // entry itself. Seed = TRUE hidden of the bonus's predecessor (last verify
12662 // col). Saves one MTP-block pass per round on top of the pairing fix.
12663 if !devacc_seeded {
12664 e.copy_into(&mut h_seed_buf, 0, &vh_seed, n_embd)?;
12665 e.copy_into(&mut fill_prev, 0, &vh_seed, n_embd)?;
12666 }
12667 pending = Some(bonus);
12668 if debug_spec {
12669 eprintln!(" -> FULL ACCEPT (bonus pending, prev-h seed)");
12670 }
12671 } else if !spec_replay && base + n_acc >= 1 {
12672 // PARTIAL ACCEPT, REPLAY-FREE (2026-07-03 — the profiled #1 long-ctx spec cost):
12673 // the verify's first j = base+n_acc columns ARE the committed sequence, computed
12674 // bit-identically to eager (decode-exact contract) — so KEEP them: KV truncates to
12675 // pos+j, recurrent state rebuilds from the VerifyCkpt (same-kernel gdn prefix
12676 // re-run / pure state-clone restore), and the bonus stays PENDING exactly like the
12677 // full-accept path — the legacy duplicate trunk replay is gone. The next chain
12678 // seeds from the MTP pseudo-hidden of the bonus, whose seed = the TRUE verify
12679 // hidden of its predecessor (col j-1) — same one-hop pseudo structure as full
12680 // accept (never compounds: the next verify recomputes true hiddens for all
12681 // committed columns).
12682 let j = base + n_acc;
12683 // VERIFY-GRAPH SLAB COMMIT: when the captured trunk ran, the linear layers'
12684 // column stash was written into the graphs ctx's persistent slabs as in-graph
12685 // memcpy nodes, NOT into the per-column VerifyCkpt the cols arm reads — so the
12686 // commit must take the slab twin (same semantics, slab-addressed sources). The
12687 // ctx states which of the two this round produced via `round_slab`; trusting the
12688 // flag rather than the env keeps a round that fell back to the eager walk (a
12689 // capture that declined, a t the pool never captured) on the cols arm.
12690 let slab_commit = vg_guard
12691 .as_ref()
12692 .and_then(|g| g.as_ref())
12693 .map(|g| g.round_slab)
12694 .unwrap_or(false);
12695 if slab_commit {
12696 self.dspark_commit_prefix_slab(
12697 e,
12698 &mut *cache,
12699 &snap,
12700 vg_guard
12701 .as_ref()
12702 .and_then(|g| g.as_ref())
12703 .expect("slab_commit implies a graphs ctx"),
12704 j,
12705 )?;
12706 } else {
12707 self.commit_verified_prefix(
12708 e,
12709 &mut *cache,
12710 &snap,
12711 ckpt.as_ref().unwrap(),
12712 j,
12713 devacc_seeded,
12714 if devacc_seeded {
12715 devacc_acc.as_ref().map(|a| (a, base, t_v))
12716 } else {
12717 None
12718 },
12719 )?;
12720 }
12721 let mut seed = e.zeros(n_embd)?;
12722 e.copy_view_into(
12723 &mut seed,
12724 0,
12725 &vx.slice((j - 1) * n_embd..j * n_embd),
12726 n_embd,
12727 )?;
12728 // Draft scratch: TRUE-HIDDEN REFRESH of the committed prefix (see the full-accept
12729 // branch); without it the chain entries stand and only the tail truncates. Either
12730 // way len ends at pos+j so the pseudo append lands at the bonus's slot pos+j
12731 // (persistent mode), rope pos+j+1 (chain convention).
12732 if refresh {
12733 scratch.set_len(e, pos)?;
12734 let mut vxs = e.zeros(j * n_embd)?;
12735 e.copy_into(&mut vxs, 0, &fill_prev, n_embd)?;
12736 if j > 1 {
12737 e.copy_view_into(
12738 &mut vxs,
12739 n_embd,
12740 &vx.slice(0..(j - 1) * n_embd),
12741 (j - 1) * n_embd,
12742 )?;
12743 }
12744 self.mtp_kv_fill_all(
12745 e,
12746 &verify_tokens[0..j],
12747 &vxs,
12748 pos,
12749 &mut *scratch,
12750 embd_dev,
12751 )?;
12752 } else {
12753 scratch.set_len(e, pos + j)?;
12754 }
12755 // REFERENCE SEEDING (see the full-accept branch): seed = TRUE hidden of the
12756 // bonus's predecessor (verify col j-1); no pseudo pass.
12757 if !devacc_seeded {
12758 e.copy_into(&mut h_seed_buf, 0, &seed, n_embd)?;
12759 e.copy_into(&mut fill_prev, 0, &seed, n_embd)?;
12760 }
12761 pending = Some(bonus);
12762 if debug_spec {
12763 eprintln!(" -> PARTIAL(replay-free j={j}, bonus pending, prev-h seed)");
12764 }
12765 } else if !spec_replay {
12766 // ZERO ROUND FOLD (2026-07-10, verify-cost target #3): base+n_acc == 0 — a
12767 // pending-less round where nothing was accepted (PMIN0 zero-draft chains after a
12768 // replay/commit, or plain 0-accept rounds at round 0). The old path replayed
12769 // [bonus] through a FULL m=1 trunk+head forward (the 489us full-vocab head pass
12770 // measured at ~0.75/round on PMIN0 configs). Instead: restore the pre-round
12771 // snapshot and let the bonus ride the NEXT round's verify as col 0 — the existing
12772 // base=1 pending machinery, bit-identical by the decode-exact verify contract.
12773 // Seed: the bonus's predecessor is the last COMMITTED token, whose hidden
12774 // fill_prev already carries (same seeding as the 1-token-replay case it replaces).
12775 cache.rollback(e, &snap, 0)?;
12776 scratch.set_len(e, pos)?;
12777 e.copy_into(&mut h_seed_buf, 0, &fill_prev, n_embd)?;
12778 pending = Some(bonus);
12779 if debug_spec {
12780 eprintln!(" -> ZERO-ROUND FOLD (bonus pending, fill_prev seed)");
12781 }
12782 } else {
12783 // PARTIAL ACCEPT, LEGACY REPLAY (seam MEMRA_SPEC_REPLAY=1 — or j==0: nothing of
12784 // this round survives, only possible before the first pending exists, ~round 0):
12785 // restore EVERYTHING to the pre-round snapshot (KV truncate to pos + recur
12786 // restore), then replay the committed prefix pending? ++ draft[0..n_acc] ++
12787 // [bonus] as ONE batched T forward — single weight read, bit-identical to greedy
12788 // (the verify-all-columns path is the same math). Commits the bonus with a TRUE
12789 // trunk hidden.
12790 cache.rollback(e, &snap, 0)?; // accept_len=0: KV len = pos, recur = snapshot
12791 let mut replay: Vec<u32> = Vec::with_capacity(base + n_acc + 1);
12792 if let Some(b) = pending.take() {
12793 replay.push(b);
12794 }
12795 replay.extend_from_slice(&draft[0..n_acc]);
12796 replay.push(bonus);
12797 // Full-stack forward (decode_step_t_core = decode_step_t_h_emb_dev's body):
12798 // Predecessor pairing seeds from the PREDECESSOR row (col len-2) — the same-row path takes the
12799 // last col exactly as before (byte-identical to the old _h_emb_dev call).
12800 let (rl_d, rx) = if self.batched_serving_numeric_class() {
12801 let mut logits = Vec::with_capacity(replay.len() * n_vocab);
12802 let mut hidden = e.uninit(replay.len() * n_embd)?;
12803 for (row, &token) in replay.iter().enumerate() {
12804 let (row_logits, row_hidden) =
12805 self.spec_target_step_h(e, token, &mut *cache)?;
12806 logits.extend_from_slice(&row_logits);
12807 e.dtod_copy_into(&row_hidden, &mut hidden, row * n_embd)?;
12808 }
12809 (e.htod(&logits)?, hidden)
12810 } else {
12811 self.decode_step_t_core(e, &replay, pos, &mut *cache, embd_dev, None)?
12812 };
12813 // last_pred = argmax of the LAST column's logits (predicts the token after `bonus`)
12814 // — device argmax + one 4-byte read instead of the full-vocab column dtoh.
12815 e.argmax_token_device_col(&rl_d, replay.len() - 1, n_vocab, &mut preds_d, 0)?;
12816 last_pred = e.dtoh_u32(&preds_d)?[0];
12817 if sampled {
12818 let lr0 = replay.len();
12819 let lc = last_col_logits
12820 .as_mut()
12821 .expect("sampled: last_col_logits unset");
12822 e.copy_view_into(
12823 lc,
12824 0,
12825 &rl_d.slice((lr0 - 1) * n_vocab..lr0 * n_vocab),
12826 n_vocab,
12827 )?;
12828 }
12829 let lr = replay.len();
12830 if lr >= 2 {
12831 e.copy_view_into(
12832 &mut h_seed_buf,
12833 0,
12834 &rx.slice((lr - 2) * n_embd..(lr - 1) * n_embd),
12835 n_embd,
12836 )?;
12837 } else {
12838 // 1-token replay (round-0 miss): the bonus's predecessor is the OLD
12839 // last_token, whose own-row hidden fill_prev still holds.
12840 e.copy_into(&mut h_seed_buf, 0, &fill_prev, n_embd)?;
12841 }
12842 // the bonus is COMMITTED here — it becomes the last committed row.
12843 let mut rh_last = e.zeros(n_embd)?;
12844 e.copy_view_into(
12845 &mut rh_last,
12846 0,
12847 &rx.slice((lr - 1) * n_embd..lr * n_embd),
12848 n_embd,
12849 )?;
12850 e.copy_into(&mut fill_prev, 0, &rh_last, n_embd)?;
12851 if debug_spec {
12852 eprintln!(" -> PARTIAL(replay={replay:?}), next_pred={last_pred}");
12853 }
12854 }
12855 if devacc_seeded {
12856 // stage (b) epilogue: fill_prev takes the gathered seed AFTER the refresh fills
12857 // consumed the old value (both slots carry the same value in every non-replay arm).
12858 e.copy_into(&mut fill_prev, 0, &h_seed_buf, n_embd)?;
12859 }
12860 if successor_valid {
12861 let optimistic_scratch_len = successor_attempt
12862 .as_ref()
12863 .expect("valid controller successor disappeared")
12864 .scratch_len;
12865 // The normal current-round commit refreshed/truncated the logical scratch tail.
12866 // Its optimistic successor row was already written physically, so restoring only
12867 // the retained logical length makes that row live for the carried round.
12868 scratch.set_len(e, optimistic_scratch_len)?;
12869 }
12870 if let Some(current) = current_opti.take() {
12871 opti_fork
12872 .as_mut()
12873 .ok_or("optipipe current retirement lost fork state")?
12874 .retire(current.generation)?;
12875 }
12876 if successor_valid {
12877 let successor = successor_attempt
12878 .take()
12879 .expect("valid controller successor disappeared before promotion");
12880 let generation = successor.generation;
12881 opti_fork
12882 .as_mut()
12883 .ok_or("optipipe successor promotion lost fork state")?
12884 .promote_successor_snapshot(&mut snap, generation);
12885 carried_opti = Some(successor);
12886 }
12887 if anatomy_on {
12888 // Commit/rollback is normally asynchronous on the primary/head stream. Bound it
12889 // only for this diagnostic so it does not disappear into the following draft's
12890 // first token readback.
12891 e.stream().synchronize()?;
12892 ph_commit += commit_started.elapsed().as_secs_f64();
12893 }
12894 // adaptive-K update (host math, zero syncs): next round drafts accepted-run + 1,
12895 // clamped to [floor(pos), k_cap]. cache.pos is post-rollback here (the round's
12896 // final position — the floor's position key reads the committed depth). Burst
12897 // rounds (`continue` above) draft the captured fixed depth and skip this, exactly
12898 // like gemma's burst arm.
12899 if adapt {
12900 let fl_now = floor_at(cache.pos);
12901 kc = (n_acc + 1).clamp(fl_now.min(k_cap), k_cap);
12902 }
12903 ph_mark(&mut ph_rest, phase_on);
12904 if let Some(p) = pipe {
12905 p.accept_end(round);
12906 }
12907 drop(pipe_accept);
12908 round += 1;
12909 // sse-cadence: this round's accepted drafts + bonus are committed (out is
12910 // append-only past step 4) — flush at round cadence.
12911 keep_going = flush_commit(&mut on_commit, &out, &mut flushed);
12912 }
12913 if let Some(mut ticket) = carried_opti.take() {
12914 opti_fork
12915 .as_mut()
12916 .ok_or("optipipe tail drain lost fork state")?
12917 .cancel_controller_ticket(e, &mut *cache, &mut *scratch, &snap, &mut ticket)?;
12918 }
12919 // sse-cadence: nothing below appends to `out`; flush any remainder (defensive).
12920 // (verdict ignored — the burst is over either way; the session tail runs unchanged.)
12921 let _ = flush_commit(&mut on_commit, &out, &mut flushed);
12922
12923 if spec_stats {
12924 let per_slot: Vec<String> = (0..k)
12925 .map(|j| {
12926 if st_drafted[j] > 0 {
12927 format!(
12928 "{}/{}={:.3}",
12929 st_accepted[j],
12930 st_drafted[j],
12931 st_accepted[j] as f64 / st_drafted[j] as f64
12932 )
12933 } else {
12934 "0/0".into()
12935 }
12936 })
12937 .collect();
12938 let acc = if total_drafted > 0 {
12939 total_accepted as f64 / total_drafted as f64
12940 } else {
12941 0.0
12942 };
12943 eprintln!(
12944 "[spec-stats] rounds={round} full_accept={st_full} len_hist={st_len_hist:?} \
12945 per_slot=[{}] total={total_accepted}/{total_drafted}={acc:.3} \
12946 tok_per_round={:.3}",
12947 per_slot.join(" "),
12948 (total_accepted + round) as f64 / round.max(1) as f64
12949 );
12950 }
12951 if constraint.is_some() {
12952 eprintln!(
12953 "[draft-mask] mask_rounds={dm_rounds} clone_total={:.3}ms \
12954 clone_per_round={:.4}ms gram_cuts={dm_cuts}/{round} cut_tokens={dm_cut_tokens}",
12955 dm_clone_ns as f64 / 1e6,
12956 dm_clone_ns as f64 / 1e6 / dm_rounds.max(1) as f64
12957 );
12958 }
12959 if phase_on {
12960 let tot = ph_draft + ph_verify + ph_wait + ph_rest;
12961 eprintln!(
12962 "[spec-phase] draft={:.1}ms ({:.1}%) verify-issue={:.1}ms ({:.1}%) verify-wait={:.1}ms ({:.1}%) commit-host={:.1}ms ({:.1}%) rounds={round}",
12963 ph_draft * 1e3,
12964 ph_draft / tot * 100.0,
12965 ph_verify * 1e3,
12966 ph_verify / tot * 100.0,
12967 ph_wait * 1e3,
12968 ph_wait / tot * 100.0,
12969 ph_rest * 1e3,
12970 ph_rest / tot * 100.0
12971 );
12972 }
12973 if anatomy_on {
12974 let rounds_f = round.max(1) as f64;
12975 let other = (ph_rest - ph_commit).max(0.0);
12976 eprintln!(
12977 "[spec-anatomy] per-round draft={:.3}ms pp-verify={:.3}ms \
12978 verify-accept={:.3}ms commit-rollback={:.3}ms other={:.3}ms rounds={round}",
12979 ph_draft * 1e3 / rounds_f,
12980 ph_verify * 1e3 / rounds_f,
12981 ph_wait * 1e3 / rounds_f,
12982 ph_commit * 1e3 / rounds_f,
12983 other * 1e3 / rounds_f,
12984 );
12985 }
12986 let _pipe_tail = pipe.map(|p| p.primary());
12987 // SESSION TAIL: leave the session in the exact invariant the next turn's suffix prime
12988 // expects — every row in `committed` has trunk KV/recur state AND an exact draft-KV row.
12989 // Park the draft-graph ctx back on the session (the serve-burst fixed-cost fix): the next
12990 // burst replays instead of recapturing. Error paths (`?` above) drop it — recaptured then.
12991 if let Some(slot) = sess_draft_slot.take() {
12992 *slot = Some(dctx);
12993 }
12994 let t_rounds = t_ent.elapsed();
12995 if let Some((committed, last_h, next_pred_slot, sctr_slot, uctr_slot)) = sess_tail.take() {
12996 // NEXT BURST'S BOUNDARY TOKEN (lane/sampled-spec-quality, Item 1). Greedy stashes
12997 // the argmax `last_pred` exactly as before (byte contract). SAMPLED draws the token
12998 // HERE, where the sampler, the session Philox counters and the penalty window are
12999 // all live and the boundary logits row still exists — that is the "make the state
13000 // available" half of the fix; the consuming burst then just emits it. `sctr` is
13001 // written to the session BELOW the draws so the advance is never lost.
13002 *next_pred_slot = Some(last_pred);
13003 let sample_boundary = sampled && constraint.is_none() && spec_sampled_boundary_on();
13004 let mut stashed_pending = false;
13005 if let Some(b) = pending.take() {
13006 if !sampled {
13007 // PENDING-CARRY (2026-08-01): stash the bonus on the session instead of
13008 // committing it with a solo T=1 pass — the next empty-suffix greedy burst
13009 // consumes it as round-0 verify col 0 (a plain round edge; the old tail
13010 // commit + next burst's init feed were 11.6+11.5ms solo trunk passes per
13011 // burst on H100 q27, [spec-setup] trace). b stays in `out` (emitted) but
13012 // OUT of `committed` (cache rows == committed); the consuming call
13013 // prepends it once its verify commits the row. next_pred is unknowable
13014 // without the commit pass — None; callers gate on pending_tok too.
13015 debug_assert_eq!(out.last(), Some(&b), "pending must be the last emitted");
13016 if let Some(slot) = sess_pending_slot.take() {
13017 *slot = Some(b);
13018 }
13019 *next_pred_slot = None;
13020 // fill_prev = hidden of the last COMMITTED row (b's predecessor) — the
13021 // exact chain-seed/fill anchor the consuming burst (or a flush) needs.
13022 *last_h = Some(e.clone_dtod(&fill_prev)?);
13023 stashed_pending = true;
13024 } else {
13025 // SAMPLED tail (unchanged): commit the bonus (one T=1 pass) + draft fill —
13026 // the sampled round-0 accept needs this pass's logits (last_col_logits).
13027 let pos_b = cache.pos;
13028 scratch.set_len(e, pos_b)?;
13029 let (lg_b, hb) = self.spec_target_step_h(e, b, &mut *cache)?;
13030 // after a FULL-accept exit `last_pred` is STALE (it predicted the bonus
13031 // itself — the prediction AFTER the bonus never materialized; it would have
13032 // been the next round's verify col 0). The commit's logits ARE that
13033 // prediction — so they are also the row the next burst's boundary token
13034 // comes off, and (lane/sampled-spec-quality) it is DRAWN from them here.
13035 *next_pred_slot = Some(if sample_boundary {
13036 sample_boundary_token(
13037 e,
13038 &lg_b,
13039 &sp,
13040 &pen_hist,
13041 &mut sctr,
13042 "burst-tail-commit",
13043 )?
13044 } else {
13045 argmax(&lg_b) as u32
13046 });
13047 self.mtp_kv_fill_all(e, &[b], &fill_prev, pos_b, &mut *scratch, embd_dev)?;
13048 *last_h = Some(hb);
13049 }
13050 } else {
13051 // fill_prev tracks the hidden of the last COMMITTED row throughout the loop.
13052 *last_h = Some(e.clone_dtod(&fill_prev)?);
13053 if sample_boundary {
13054 // No pending to commit, so the boundary row is the one `last_pred` was
13055 // argmaxed from and the sampled path keeps it on device: the init feed's
13056 // logits when the burst ran zero rounds, else the legacy-replay path's
13057 // last verify column (both predict the token AFTER the last committed
13058 // row). It is retained precisely because round 0's accept test needs it,
13059 // so the draw costs no extra D2H of the [n_vocab] row.
13060 match last_col_logits.as_ref() {
13061 Some(lc) => {
13062 *next_pred_slot = Some(sample_boundary_token_dev(
13063 e,
13064 lc,
13065 n_vocab,
13066 &sp,
13067 &pen_hist,
13068 &mut sctr,
13069 "burst-tail-nopending",
13070 )?);
13071 }
13072 // NAME THE FALLBACK (house standard): unreachable today — a sampled
13073 // burst always feeds or replays, so the row exists — but if it ever
13074 // is, the stream takes a greedy token and SAYS so rather than
13075 // silently regressing to the pre-lane behaviour.
13076 None => eprintln!(
13077 "[spec-boundary] sampled tail kept the ARGMAX boundary token \
13078 (reason: no retained boundary logits row)"
13079 ),
13080 }
13081 }
13082 }
13083 *sctr_slot = sctr;
13084 *uctr_slot = uctr;
13085 committed.extend_from_slice(prompt);
13086 if let Some(cb) = carried_pending {
13087 // the consumed carry's cache row landed in round 0's verify (every pending
13088 // round commits col 0) — it joins `committed` here, in sequence order.
13089 committed.push(cb);
13090 }
13091 if stashed_pending {
13092 // ZERO-EMIT BURST (2026-08-06 c=8 serve panic, pre-existing since b4aea184):
13093 // `out.len() - 1` underflowed on an EMPTY `out` — "range end index
13094 // 18446744073709551615 out of range for slice of length 0", killing the
13095 // memra-gpu-worker and failing 31 of 32 concurrent requests with "worker closed
13096 // stream". Reachable because `pending` starts as `carried_pending` (a bonus
13097 // stashed by the PREVIOUS burst) while `out` starts empty, and the carry is
13098 // deliberately NOT pushed to `out` (line ~3239: the burst that emitted it already
13099 // did). So a burst that stashes a pending without emitting anything of its own —
13100 // the round loop exits before a push, e.g. the ring drain's `out.len() < max_new`
13101 // guard skipping every token under a tight budget — arrives here with
13102 // out.len() == 0 and stashed_pending == true.
13103 //
13104 // The invariant is unchanged: `committed` gets every emitted token EXCEPT the
13105 // stashed bonus. With nothing emitted, that is nothing — and the carry pushed
13106 // just above is already accounted. Saturating, not a min/assert: an empty `out`
13107 // here is a legitimate burst shape, not a corrupt state.
13108 let emitted = out.len().saturating_sub(1);
13109 committed.extend_from_slice(&out[..emitted]);
13110 } else {
13111 committed.extend_from_slice(&out); // FULL out incl. overshoot — all committed
13112 }
13113 debug_assert_eq!(
13114 cache.pos,
13115 committed.len(),
13116 "session invariant: cache rows == committed tokens"
13117 );
13118 if setup_trace {
13119 e.stream().synchronize()?; // bound the async tail fill in the trace
13120 let t_tail = t_ent.elapsed();
13121 eprintln!(
13122 "[spec-setup] init={:.2}ms cap={:.2}ms fill={:.2}ms rounds={:.2}ms tail={:.2}ms total={:.2}ms out={} cont={}",
13123 t_init.as_secs_f64() * 1e3,
13124 (t_cap - t_init).as_secs_f64() * 1e3,
13125 (t_fill - t_cap).as_secs_f64() * 1e3,
13126 (t_rounds - t_fill).as_secs_f64() * 1e3,
13127 (t_tail - t_rounds).as_secs_f64() * 1e3,
13128 t_tail.as_secs_f64() * 1e3,
13129 out.len(),
13130 continuation
13131 );
13132 }
13133 return Ok((out, total_drafted, total_accepted));
13134 }
13135 out.truncate(max_new);
13136 Ok((out, total_drafted, total_accepted))
13137 }
13138
13139 /// Anchor-bounded DSpark target extraction. The trunk sees the exact generated token tape;
13140 /// only requested hidden rows and target-logit rows cross PCIe. An anchor token at p pairs
13141 /// with the pre-output-norm h[p-1] carrier, exactly as the existing replay/NextN path does.
13142 pub fn extract_dspark_anchors(
13143 &self,
13144 e: &Engine,
13145 tokens: &[u32],
13146 anchor_positions: &[usize],
13147 gamma: usize,
13148 top_k: usize,
13149 chunk: usize,
13150 temperature: f32,
13151 ) -> Result<Vec<DsparkAnchorRecord>, Box<dyn std::error::Error>> {
13152 if tokens.len() < gamma + 2 || gamma == 0 || chunk < 2 {
13153 return Err("DSpark extraction token tape/gamma/chunk is invalid".into());
13154 }
13155 if anchor_positions.windows(2).any(|pair| pair[0] >= pair[1]) {
13156 return Err("DSpark anchor positions must be sorted and unique".into());
13157 }
13158 for &position in anchor_positions {
13159 if position == 0 || position + gamma >= tokens.len() {
13160 return Err(format!(
13161 "DSpark anchor {position} has no predecessor or cannot cover gamma={gamma} in {} tokens",
13162 tokens.len()
13163 )
13164 .into());
13165 }
13166 }
13167
13168 let n_vocab = self.output.out_features();
13169 let n_embd = self.cfg.n_embd as usize;
13170 let mut cache =
13171 crate::pp::new_cache_planned(e, &self.cfg, &self.plan, tokens.len() + gamma + 8)?;
13172 let (embd_qt, embd_rb) = self.embd.qt_and_row_bytes(n_embd);
13173 let embd_gpu = if spec_host_embd() {
13174 None
13175 } else {
13176 Some(
13177 self.embd_gpu
13178 .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload")),
13179 )
13180 };
13181 let embd_dev = embd_gpu.map(|gpu| (gpu, embd_qt, embd_rb));
13182
13183 struct PendingRecord {
13184 position: usize,
13185 hidden: Option<Vec<f32>>,
13186 tokens: Vec<u32>,
13187 target_top_ids: Vec<Option<Vec<u32>>>,
13188 target_top_logits: Vec<Option<Vec<f32>>>,
13189 target_top_probs: Vec<Option<Vec<f32>>>,
13190 target_tail_probs: Vec<Option<f32>>,
13191 }
13192
13193 let mut pending: Vec<PendingRecord> = anchor_positions
13194 .iter()
13195 .map(|&position| PendingRecord {
13196 position,
13197 hidden: None,
13198 tokens: tokens[position..=position + gamma].to_vec(),
13199 target_top_ids: vec![None; gamma],
13200 target_top_logits: vec![None; gamma],
13201 target_top_probs: vec![None; gamma],
13202 target_tail_probs: vec![None; gamma],
13203 })
13204 .collect();
13205
13206 let mut start = 0usize;
13207 while start < tokens.len() {
13208 let end = (start + chunk).min(tokens.len());
13209 let chunk_tokens = &tokens[start..end];
13210 let (target_logits, hidden_rows) =
13211 self.decode_step_t_core(e, chunk_tokens, start, &mut cache, embd_dev, None)?;
13212 for record in &mut pending {
13213 let hidden_position = record.position - 1;
13214 if hidden_position >= start && hidden_position < end {
13215 let local = hidden_position - start;
13216 record.hidden = Some(
13217 e.dtoh_view(&hidden_rows.slice(local * n_embd..(local + 1) * n_embd))?,
13218 );
13219 }
13220 for slot in 0..gamma {
13221 let target_row = record.position + slot;
13222 if target_row < start || target_row >= end {
13223 continue;
13224 }
13225 let local = target_row - start;
13226 let logits =
13227 e.dtoh_view(&target_logits.slice(local * n_vocab..(local + 1) * n_vocab))?;
13228 let (ids, top_logits, probs, tail) =
13229 dspark_sparse_softmax_topk(&logits, top_k, temperature)?;
13230 record.target_top_ids[slot] = Some(ids);
13231 record.target_top_logits[slot] = Some(top_logits);
13232 record.target_top_probs[slot] = Some(probs);
13233 record.target_tail_probs[slot] = Some(tail);
13234 }
13235 }
13236 start = end;
13237 }
13238
13239 pending
13240 .into_iter()
13241 .map(|record| {
13242 let hidden = record
13243 .hidden
13244 .ok_or_else(|| format!("missing DSpark hidden at {}", record.position))?;
13245 let target_top_ids =
13246 flatten_dspark_rows(record.target_top_ids, record.position, "target ids")?;
13247 let target_top_logits = flatten_dspark_rows(
13248 record.target_top_logits,
13249 record.position,
13250 "target logits",
13251 )?;
13252 let target_top_probs =
13253 flatten_dspark_rows(record.target_top_probs, record.position, "target probs")?;
13254 let target_tail_probs = record
13255 .target_tail_probs
13256 .into_iter()
13257 .enumerate()
13258 .map(|(slot, value)| {
13259 value.ok_or_else(|| {
13260 format!("missing DSpark tail at {} slot {slot}", record.position)
13261 })
13262 })
13263 .collect::<Result<Vec<_>, _>>()?;
13264 Ok(DsparkAnchorRecord {
13265 position: record.position,
13266 hidden,
13267 tokens: record.tokens,
13268 target_top_ids,
13269 target_top_logits,
13270 target_top_probs,
13271 target_tail_probs,
13272 })
13273 })
13274 .collect()
13275 }
13276
13277 /// TEACHER-FORCED REPLAY ACCEPTANCE (hqmtp MTP-heal protocol): walk a FIXED token
13278 /// sequence and, at sampled positions, compare the MTP head's K-token draft chain against
13279 /// the trunk's own teacher-forced greedy predictions. Nothing is generated — the context is
13280 /// the corpus text itself, so (a) degenerate self-generated loops cannot inflate acceptance
13281 /// and (b) two arms (bf16 ceiling vs NVFP4) score on IDENTICAL contexts, isolating the
13282 /// quant-induced head/hidden-state mismatch from text drift.
13283 ///
13284 /// Per eval position p (context = tokens[0..=p], predecessor pairing as in spec decode):
13285 /// draft_j = chain token j from (tokens[p], h_{p-1}), then its own drafts — the exact
13286 /// eager spec-decode chain (same mtp_head_forward_dev, same rope positions).
13287 /// target_j = teacher-forced greedy pick for position p+1+j (argmax of the trunk logits
13288 /// at forced context tokens[0..p+j]). For j==0 this equals live spec
13289 /// acceptance; for j>=1 live verify would condition on the drafts, here it
13290 /// conditions on the corpus — deterministic and arm-comparable by design.
13291 ///
13292 /// Returns (rows, bg): one (p, drafts[k], targets[k]) row per eval position (ascending p),
13293 /// plus the full teacher-forced greedy track bg (bg[i] = greedy pick for position i, i>=1)
13294 /// so harnesses can cross-check runs (e.g. different chunk sizes must give identical bg).
13295 ///
13296 /// `hdump`: when Some, every position's pre-output_norm trunk hidden (the exact rows the
13297 /// draft-KV fill pairs from) streams to the file as little-endian f32 [t_total, n_embd] —
13298 /// the head-distillation extraction (hqmtp): the ENGINE is the source of truth for trunk
13299 /// hiddens (HF torch reproductions of the hybrid trunk measured only ~0.5 greedy
13300 /// agreement vs this path — not usable as a training-data source).
13301 pub fn replay_acceptance(
13302 &self,
13303 e: &Engine,
13304 tokens: &[u32],
13305 k: usize,
13306 stride: usize,
13307 chunk: usize,
13308 mut hdump: Option<&mut std::fs::File>,
13309 ) -> Result<(Vec<(usize, Vec<u32>, Vec<u32>)>, Vec<u32>), Box<dyn std::error::Error>> {
13310 assert!(k >= 1 && stride >= 1 && chunk >= 2);
13311 let mtp = self
13312 .mtp
13313 .as_ref()
13314 .expect("replay_acceptance requires an MTP head");
13315 let n_vocab = self.output.out_features();
13316 let d_vocab = mtp
13317 .shared_head_head
13318 .as_ref()
13319 .unwrap_or(&self.output)
13320 .out_features();
13321 let n_embd = self.cfg.n_embd as usize;
13322 let t_total = tokens.len();
13323 assert!(t_total >= 8, "corpus too short ({t_total} tokens)");
13324 // STAGE-OWNED KV (lane/pp2-spec 2026-08-06) — see `new_session`. Door shut = `Cache::new`.
13325 let mut cache = crate::pp::new_cache_planned(e, &self.cfg, &self.plan, t_total + k + 8)?;
13326 let mut scratch = self.new_mtp_scratch(e, t_total + k + 8)?;
13327 let (embd_qt, embd_rb) = self.embd.qt_and_row_bytes(n_embd);
13328 let embd_gpu = if spec_host_embd() {
13329 None
13330 } else {
13331 Some(
13332 self.embd_gpu
13333 .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload")),
13334 )
13335 };
13336 let embd_dev = embd_gpu.map(|g| (g, embd_qt, embd_rb));
13337
13338 // bg[i] = the trunk's greedy pick for position i under the forced context (i >= 1).
13339 let mut bg: Vec<u32> = vec![0; t_total + 1];
13340 let mut rows: Vec<(usize, Vec<u32>, Vec<u32>)> = Vec::new();
13341 let mut prev_last_h = e.zeros(n_embd)?; // predecessor hidden entering the chunk
13342 let mut seed_buf = e.zeros(n_embd)?;
13343 let mut preds_d = e.alloc_u32_zeroed(chunk)?;
13344 let nll_on = std::env::var("MEMRA_REPLAY_NLL").as_deref() == Ok("1");
13345 let (mut nll_sum, mut nll_cnt) = (0f64, 0u64);
13346 let mut s = 0usize;
13347 while s < t_total {
13348 let cend = (s + chunk).min(t_total);
13349 let tc = cend - s;
13350 let ch = &tokens[s..cend];
13351 // 1. forced trunk pass — verify path (decode-exact contract): all-column logits +
13352 // the chunk's true hiddens.
13353 let (tl_d, vx) = self.decode_step_t_core(e, ch, s, &mut cache, embd_dev, None)?;
13354 for j in 0..tc {
13355 e.argmax_token_device_col(&tl_d, j, n_vocab, &mut preds_d, j)?;
13356 }
13357 let preds = e.dtoh_u32(&preds_d)?;
13358 for j in 0..tc {
13359 bg[s + j + 1] = preds[j];
13360 }
13361 // MEMRA_REPLAY_NLL=1: teacher-forced NLL/perplexity over the same forced pass — the
13362 // checkpoint-quality metric (position j's logits score the GOLD next token).
13363 if nll_on {
13364 let jmax = if cend < t_total { tc } else { tc - 1 }; // last pos has no gold next
13365 if jmax > 0 {
13366 let ids: Vec<u32> = (0..jmax).map(|j| tokens[s + j + 1]).collect();
13367 let rows: Vec<i32> = (0..jmax as i32).collect();
13368 let idsd = e.htod_u32_v(&ids)?;
13369 let rowsd = e.htod_i32(&rows)?;
13370 let mut outd = e.zeros(jmax)?;
13371 e.softmax_gather(&tl_d, n_vocab, &idsd, &rowsd, &mut outd, n_vocab, jmax, 1.0)?;
13372 for pr in e.dtoh(&outd)? {
13373 nll_sum += -((pr.max(1e-30)) as f64).ln();
13374 nll_cnt += 1;
13375 }
13376 }
13377 }
13378 if let Some(f) = hdump.as_deref_mut() {
13379 use std::io::Write;
13380 let host: Vec<f32> = e.dtoh(&vx)?;
13381 // bf16 round-to-nearest-even — f32 doubled the disk bill at bulk
13382 // extraction scale (20M tokens x 4096 = 320GB f32 vs 160GB bf16).
13383 let mut bytes = Vec::with_capacity(tc * n_embd * 2);
13384 for v in &host[..tc * n_embd] {
13385 let b = v.to_bits();
13386 let r = b.wrapping_add(0x7FFF + ((b >> 16) & 1));
13387 bytes.extend_from_slice(&((r >> 16) as u16).to_le_bytes());
13388 }
13389 f.write_all(&bytes)?;
13390 }
13391 // CHAINLESS extraction (stride > corpus, the bulk-hdump mode): no chunk ever
13392 // drafts, so the draft-KV fills are pure waste — skip them (2 MTP-block passes
13393 // per token saved; the forced trunk pass + hdump is all the mode needs).
13394 let chainless = stride > t_total;
13395 if chainless {
13396 e.copy_view_into(
13397 &mut prev_last_h,
13398 0,
13399 &vx.slice((tc - 1) * n_embd..tc * n_embd),
13400 n_embd,
13401 )?;
13402 s = cend;
13403 continue;
13404 }
13405 // 2. TRUE predecessor-paired draft-KV fill for the chunk (row i carries h_{i-1};
13406 // row s reads the previous chunk's last true hidden, zeros at corpus start).
13407 let mut vxs = e.zeros(tc * n_embd)?;
13408 e.copy_into(&mut vxs, 0, &prev_last_h, n_embd)?;
13409 if tc > 1 {
13410 e.copy_view_into(
13411 &mut vxs,
13412 n_embd,
13413 &vx.slice(0..(tc - 1) * n_embd),
13414 (tc - 1) * n_embd,
13415 )?;
13416 }
13417 scratch.set_len(e, s)?;
13418 self.mtp_kv_fill_all(e, ch, &vxs, s, &mut scratch, embd_dev)?;
13419 // 3. draft chains at sampled positions, DESCENDING: a chain reads only slots
13420 // [0..p) (true fills) and appends at >= p; the next (smaller-p) chain's set_len
13421 // truncates those approximate appends before they can ever be read.
13422 let ps: Vec<usize> = (s..cend)
13423 .filter(|p| *p >= 1 && *p % stride == 0 && *p + k <= t_total)
13424 .collect();
13425 for &p in ps.iter().rev() {
13426 scratch.set_len(e, p)?;
13427 if p == s {
13428 e.copy_into(&mut seed_buf, 0, &prev_last_h, n_embd)?;
13429 } else {
13430 e.copy_view_into(
13431 &mut seed_buf,
13432 0,
13433 &vx.slice((p - 1 - s) * n_embd..(p - s) * n_embd),
13434 n_embd,
13435 )?;
13436 }
13437 let mut e_tok = tokens[p];
13438 let mut d_seed = e.clone_dtod(&seed_buf)?;
13439 let chain_heads = !self.mtp_extra.is_empty();
13440 let mut chain_tokens = if chain_heads {
13441 vec![tokens[p]]
13442 } else {
13443 Vec::new()
13444 };
13445 let mut chain_seeds = if chain_heads {
13446 vec![e.clone_dtod(&seed_buf)?]
13447 } else {
13448 Vec::new()
13449 };
13450 let mut drafts: Vec<u32> = Vec::with_capacity(k);
13451 for j in 0..k {
13452 let (dl_d, h_nextn) = if chain_heads {
13453 self.mtp_chain_forward_dev(
13454 e,
13455 &chain_tokens,
13456 &chain_seeds,
13457 &mut scratch,
13458 p,
13459 embd_dev,
13460 None,
13461 )?
13462 } else {
13463 self.mtp_head_forward_dev(
13464 e,
13465 mtp,
13466 e_tok,
13467 &d_seed,
13468 &mut scratch,
13469 p + 1 + j,
13470 embd_dev,
13471 None,
13472 )?
13473 };
13474 let tok_d = e.argmax_token_device(&dl_d, d_vocab)?;
13475 let idx = e.dtoh_u32_one(&tok_d)?;
13476 let d = match &mtp.d2t {
13477 Some(map) => map[idx as usize],
13478 None => idx,
13479 };
13480 drafts.push(d);
13481 if chain_heads {
13482 chain_tokens.push(d);
13483 chain_seeds.push(h_nextn);
13484 } else {
13485 e_tok = d;
13486 d_seed = h_nextn;
13487 }
13488 }
13489 // targets may live in a LATER chunk's bg — resolved after the walk.
13490 rows.push((p, drafts, Vec::new()));
13491 }
13492 // 4. restore TRUE entries for the whole chunk (the next chunk's chains and fills
13493 // expect scratch.len == cend with exact rows).
13494 scratch.set_len(e, s)?;
13495 self.mtp_kv_fill_all(e, ch, &vxs, s, &mut scratch, embd_dev)?;
13496 e.copy_view_into(
13497 &mut prev_last_h,
13498 0,
13499 &vx.slice((tc - 1) * n_embd..tc * n_embd),
13500 n_embd,
13501 )?;
13502 s = cend;
13503 }
13504 for (p, drafts, targets) in rows.iter_mut() {
13505 for j in 0..drafts.len() {
13506 targets.push(bg[*p + 1 + j]);
13507 }
13508 }
13509 rows.sort_by_key(|r| r.0);
13510 if nll_cnt > 0 {
13511 let mean = nll_sum / nll_cnt as f64;
13512 println!(
13513 "[replay-nll] tokens={nll_cnt} nll/token={mean:.5} ppl={:.4}",
13514 mean.exp()
13515 );
13516 }
13517 Ok((rows, bg))
13518 }
13519}
13520
13521#[cfg(test)]
13522mod mtp_chain_tests {
13523 use super::mtp_chain_head_index;
13524
13525 #[test]
13526 fn embedded_step_heads_cycle_in_declared_order() {
13527 let actual: Vec<usize> = (0..8).map(|step| mtp_chain_head_index(step, 3)).collect();
13528 assert_eq!(actual, [0, 1, 2, 0, 1, 2, 0, 1]);
13529 }
13530
13531 #[test]
13532 fn standalone_draft_remains_single_head() {
13533 assert!((0..8).all(|step| mtp_chain_head_index(step, 1) == 0));
13534 }
13535}
13536
13537#[cfg(test)]
13538mod tp_verified_prefix_tests {
13539 use super::rewind_tp_kv_verified_prefix;
13540 use crate::tp::ResidentTpKvCache;
13541
13542 fn cache_with_committed_len(committed: usize) -> ResidentTpKvCache {
13543 let mut cache = ResidentTpKvCache::new(Vec::new(), 1, 1, 1, 1, 8);
13544 let transaction = cache.begin_transaction().unwrap();
13545 let target = cache.append_target(transaction, committed).unwrap();
13546 cache.publish_append(transaction, target).unwrap();
13547 let target = cache.commit_target(transaction, committed).unwrap();
13548 cache.publish_finalize(transaction, target).unwrap();
13549 cache
13550 }
13551
13552 #[test]
13553 fn replay_free_prefix_rewinds_tp_visibility_to_snapshot_plus_accepts() {
13554 let mut layers = vec![Some(cache_with_committed_len(5)), None];
13555 rewind_tp_kv_verified_prefix(&mut layers, &[Some(2), None], 1).unwrap();
13556 let cache = layers[0].as_ref().unwrap();
13557 assert_eq!(cache.committed_len(), 3);
13558 assert_eq!(cache.staged_len(), 3);
13559 }
13560
13561 #[test]
13562 fn replay_free_prefix_rejects_a_changed_tp_cache_shape() {
13563 let mut layers = vec![Some(cache_with_committed_len(1))];
13564 let error = rewind_tp_kv_verified_prefix(&mut layers, &[None], 1)
13565 .unwrap_err()
13566 .to_string();
13567 assert!(error.contains("changed shape"), "unexpected error: {error}");
13568 }
13569}
13570
13571#[cfg(test)]
13572mod dspark_sparse_tests {
13573 use super::dspark_sparse_softmax_topk;
13574
13575 #[test]
13576 fn topk_keeps_full_softmax_mass_and_stable_ties() {
13577 let logits = [1.0f32, 3.0, 3.0, -2.0];
13578 let (ids, top_logits, probs, tail) = dspark_sparse_softmax_topk(&logits, 2, 1.0).unwrap();
13579 assert_eq!(ids, vec![1, 2]);
13580 assert_eq!(top_logits, vec![3.0, 3.0]);
13581 let denominator = logits.iter().map(|value| (value - 3.0).exp()).sum::<f32>();
13582 let expected = 1.0 / denominator;
13583 assert!((probs[0] - expected).abs() < 1.0e-6);
13584 assert!((probs[1] - expected).abs() < 1.0e-6);
13585 assert!((tail - (1.0 - 2.0 * expected)).abs() < 1.0e-6);
13586 assert!((probs.iter().sum::<f32>() + tail - 1.0).abs() < 1.0e-6);
13587 }
13588}
13589
13590#[cfg(test)]
13591mod spec_replay_env_tests {
13592 use super::spec_replay_env_on;
13593
13594 #[test]
13595 fn replay_requires_literal_one() {
13596 assert!(!spec_replay_env_on(None));
13597 assert!(!spec_replay_env_on(Some("")));
13598 assert!(!spec_replay_env_on(Some("0")));
13599 assert!(!spec_replay_env_on(Some("true")));
13600 assert!(!spec_replay_env_on(Some("2")));
13601 assert!(spec_replay_env_on(Some("1")));
13602 }
13603}
13604
13605#[cfg(test)]
13606mod telem_tests {
13607 use super::{SPEC_TELEM_POS, SpecTelemetry, SpecTelemetryCounters};
13608
13609 #[test]
13610 fn synthetic_accept_masks_produce_tau_and_position_histogram() {
13611 let counters = SpecTelemetryCounters::default();
13612 for mask in [
13613 [true, true, true],
13614 [true, true, false],
13615 [true, false, false],
13616 [false, false, false],
13617 ] {
13618 let accepted = mask.iter().take_while(|&&value| value).count();
13619 counters.record_round(mask.len(), accepted);
13620 }
13621
13622 let snapshot = counters.snapshot();
13623 assert_eq!(
13624 (snapshot.rounds, snapshot.drafted, snapshot.accepted),
13625 (4, 12, 6)
13626 );
13627 assert_eq!(&snapshot.pos_drafted[..3], &[4, 4, 4]);
13628 assert_eq!(&snapshot.pos_accepted[..3], &[3, 2, 1]);
13629 assert_eq!(snapshot.tau(), 1.5);
13630 assert_eq!(snapshot.pos_drafted[3..], [0; SPEC_TELEM_POS - 3]);
13631 assert_eq!(snapshot.pos_accepted[3..], [0; SPEC_TELEM_POS - 3]);
13632 }
13633
13634 /// The worker's per-burst pattern: stash, accumulate, diff — the delta must isolate
13635 /// exactly the burst's contribution (pool-resumed sessions carry prior requests' counts).
13636 #[test]
13637 fn delta_isolates_burst_contribution() {
13638 let mut t = SpecTelemetry::default();
13639 // "previous request": 2 rounds of k=3, accepts 3 then 1.
13640 for (kr, na) in [(3usize, 3usize), (3, 1)] {
13641 t.rounds += 1;
13642 t.drafted += kr as u64;
13643 t.accepted += na as u64;
13644 for j in 0..kr {
13645 t.pos_drafted[j] += 1;
13646 }
13647 for j in 0..na {
13648 t.pos_accepted[j] += 1;
13649 }
13650 }
13651 let before = t;
13652 // "this burst": 1 round k=3, accepts 2.
13653 t.rounds += 1;
13654 t.drafted += 3;
13655 t.accepted += 2;
13656 for j in 0..3 {
13657 t.pos_drafted[j] += 1;
13658 }
13659 for j in 0..2 {
13660 t.pos_accepted[j] += 1;
13661 }
13662 let d = t.delta_since(&before);
13663 assert_eq!((d.rounds, d.drafted, d.accepted), (1, 3, 2));
13664 assert_eq!(&d.pos_drafted[..3], &[1, 1, 1]);
13665 assert_eq!(&d.pos_accepted[..3], &[1, 1, 0]);
13666 assert_eq!(d.pos_drafted[3..], [0; SPEC_TELEM_POS - 3]);
13667 }
13668
13669 /// merge(delta) then merge(delta2) equals accumulating both — the per-model /metrics
13670 /// aggregation invariant.
13671 #[test]
13672 fn merge_accumulates_fieldwise() {
13673 let mut agg = SpecTelemetry::default();
13674 let mut d1 = SpecTelemetry {
13675 rounds: 2,
13676 drafted: 6,
13677 accepted: 4,
13678 ..Default::default()
13679 };
13680 d1.pos_drafted[0] = 2;
13681 d1.pos_accepted[0] = 2;
13682 let mut d2 = SpecTelemetry {
13683 rounds: 1,
13684 drafted: 3,
13685 accepted: 1,
13686 ..Default::default()
13687 };
13688 d2.pos_drafted[0] = 1;
13689 d2.pos_accepted[0] = 1;
13690 d2.pos_drafted[1] = 1;
13691 agg.merge(&d1);
13692 agg.merge(&d2);
13693 assert_eq!((agg.rounds, agg.drafted, agg.accepted), (3, 9, 5));
13694 assert_eq!(agg.pos_drafted[0], 3);
13695 assert_eq!(agg.pos_accepted[0], 3);
13696 assert_eq!(agg.pos_drafted[1], 1);
13697 assert_eq!(agg.pos_accepted[1], 0);
13698 }
13699
13700 /// Wrong-snapshot diff saturates to zero instead of wrapping — the counters feed a
13701 /// public metrics surface and must never publish a u64-wrapped garbage value.
13702 #[test]
13703 fn delta_saturates_never_wraps() {
13704 let small = SpecTelemetry {
13705 rounds: 1,
13706 drafted: 2,
13707 accepted: 1,
13708 ..Default::default()
13709 };
13710 let big = SpecTelemetry {
13711 rounds: 5,
13712 drafted: 15,
13713 accepted: 9,
13714 ..Default::default()
13715 };
13716 let d = small.delta_since(&big);
13717 assert_eq!((d.rounds, d.drafted, d.accepted), (0, 0, 0));
13718 }
13719}
13720
13721#[cfg(test)]
13722mod opti_fork_tests {
13723 use super::{
13724 OptiControllerPolicy, OptiForkAction, OptiForkGateMode, OptiForkGenerationTracker,
13725 };
13726
13727 #[test]
13728 fn controller_threshold_and_three_miss_breaker_are_exact() {
13729 let mut policy = OptiControllerPolicy {
13730 threshold: 0.7,
13731 consecutive_misses: 0,
13732 breaker_tripped: false,
13733 };
13734 assert!(!policy.admit(0.699_999));
13735 assert!(policy.admit(0.7));
13736 assert!(!policy.resolve(false));
13737 assert!(!policy.resolve(false));
13738 assert!(policy.resolve(false));
13739 assert!(policy.breaker_tripped);
13740 assert!(!policy.admit(1.0));
13741 assert!(
13742 !policy.resolve(true),
13743 "a resolved hit cannot re-arm a tripped request"
13744 );
13745 assert!(policy.breaker_tripped);
13746 }
13747
13748 #[test]
13749 fn zero_threshold_is_the_true_unconditional_measurement_arm() {
13750 let mut policy = OptiControllerPolicy {
13751 threshold: 0.0,
13752 consecutive_misses: 0,
13753 breaker_tripped: false,
13754 };
13755 for _ in 0..16 {
13756 assert!(policy.admit(0.0));
13757 assert!(!policy.resolve(false));
13758 }
13759 for invalid in [f32::NAN, f32::INFINITY, -0.01, 1.01] {
13760 assert!(
13761 !policy.admit(invalid),
13762 "invalid q proxy must fail closed: {invalid}"
13763 );
13764 }
13765 assert!(!policy.breaker_tripped);
13766 assert_eq!(policy.consecutive_misses, 0);
13767 }
13768
13769 #[test]
13770 fn alternating_mode_flips_by_generation_not_round_parity() {
13771 assert_eq!(OptiForkGateMode::Alternate.action(0), OptiForkAction::Hit);
13772 assert_eq!(OptiForkGateMode::Alternate.action(1), OptiForkAction::Miss);
13773 assert_eq!(OptiForkGateMode::Alternate.action(8), OptiForkAction::Hit);
13774 assert_eq!(OptiForkGateMode::Alternate.action(9), OptiForkAction::Miss);
13775 }
13776
13777 #[test]
13778 fn live_generation_cannot_be_overwritten() {
13779 let mut tracker = OptiForkGenerationTracker::default();
13780 let g0 = tracker.reserve().unwrap();
13781 let g1 = tracker.reserve().unwrap();
13782 let err = tracker.reserve().unwrap_err().to_string();
13783 assert!(
13784 err.contains("still owns generation 0"),
13785 "unexpected error: {err}"
13786 );
13787 tracker.retire(g0).unwrap();
13788 let g2 = tracker.reserve().unwrap();
13789 assert_eq!((g2.id, g2.slot), (2, 0));
13790 tracker.retire(g1).unwrap();
13791 tracker.retire(g2).unwrap();
13792 }
13793
13794 #[test]
13795 fn teardown_rejects_a_stale_generation_tag() {
13796 let mut tracker = OptiForkGenerationTracker::default();
13797 let g0 = tracker.reserve().unwrap();
13798 tracker.retire(g0).unwrap();
13799 let err = tracker.retire(g0).unwrap_err().to_string();
13800 assert!(err.contains("teardown mismatch"), "unexpected error: {err}");
13801 }
13802}
13803
13804#[cfg(test)]
13805mod draft_graph_fallback_tests {
13806 use super::DraftGraphFallback;
13807
13808 /// The Q2 contract, part (a): a fallback flip is LOUD — exactly once per flip.
13809 #[test]
13810 fn flip_is_loud_once_and_memoized_after() {
13811 let mut f = DraftGraphFallback::default();
13812 let line = f
13813 .mark_greedy("out of memory")
13814 .expect("first flip must return the warn line");
13815 assert!(
13816 line.contains("WARN"),
13817 "flip line must be warn-level: {line}"
13818 );
13819 assert!(
13820 line.contains("out of memory"),
13821 "flip line must carry the reason: {line}"
13822 );
13823 assert!(f.greedy_failed());
13824 // re-marking an already-failed graph is the memoization: quiet, still failed.
13825 assert!(f.mark_greedy("out of memory").is_none());
13826 assert!(f.greedy_failed());
13827 // the two graphs' flags are independent (greedy flip leaves sampled capturable).
13828 assert!(!f.sampled_failed());
13829 let line_s = f
13830 .mark_sampled("capture unsupported")
13831 .expect("sampled flip is its own flip");
13832 assert!(
13833 line_s.contains("sampled"),
13834 "sampled flip names itself: {line_s}"
13835 );
13836 assert!(f.mark_sampled("capture unsupported").is_none());
13837 }
13838
13839 /// The Q2 contract, part (b): resume-from-pool RESETS both flags (fresh capture chance),
13840 /// and says so exactly when there was something to reset.
13841 #[test]
13842 fn reset_on_resume_clears_flags_and_logs_once() {
13843 let mut f = DraftGraphFallback::default();
13844 // clean session: resume is silent, nothing to reset.
13845 assert!(f.reset_on_resume().is_none());
13846 f.mark_greedy("oom").unwrap();
13847 f.mark_sampled("oom").unwrap();
13848 let note = f
13849 .reset_on_resume()
13850 .expect("a set flag must produce the reset note");
13851 assert!(
13852 note.contains("greedy+sampled"),
13853 "note names what was reset: {note}"
13854 );
13855 assert!(
13856 !f.greedy_failed() && !f.sampled_failed(),
13857 "both flags cleared"
13858 );
13859 // and the NEXT failure after a reset is a fresh flip — loud again.
13860 assert!(f.mark_greedy("oom again").is_some());
13861 let note2 = f.reset_on_resume().expect("greedy-only reset");
13862 assert!(note2.contains("(greedy)"), "single-flag note: {note2}");
13863 }
13864
13865 /// Shape-change clears (dmask realloc / mask-shape mismatch / s_key change) stay silent —
13866 /// they precede a fresh capture attempt whose own failure re-flips loudly.
13867 #[test]
13868 fn shape_change_clears_are_silent() {
13869 let mut f = DraftGraphFallback::default();
13870 f.mark_greedy("oom").unwrap();
13871 f.clear_greedy();
13872 assert!(!f.greedy_failed());
13873 f.mark_sampled("oom").unwrap();
13874 f.clear_sampled();
13875 assert!(!f.sampled_failed());
13876 // after a silent clear there is nothing left for resume to report.
13877 assert!(f.reset_on_resume().is_none());
13878 }
13879}
13880
13881/// SAMPLED DRAFT-GRAPH KEY (lane/graph-s-key-exactness-20260819).
13882///
13883/// These are the CPU teeth for an exactness bug whose live reproduction needs a GPU, a trunk, a
13884/// drafter and a two-turn session: the key itself. Every test below fails against the pre-fix key
13885/// `(seed, temp.to_bits(), k)` — `legacy_key` restates it so the collision is explicit rather
13886/// than remembered.
13887#[cfg(test)]
13888mod sampled_graph_key_tests {
13889 use super::{SampledGraphKey, debug_t_pred0};
13890
13891 /// The pre-fix key, verbatim: `let s_key = (sp_seed, sp_temp.to_bits(), k);`
13892 fn legacy_key(k: &SampledGraphKey) -> (u64, u32, usize) {
13893 (k.seed, k.temp_bits, k.k)
13894 }
13895
13896 fn pure_temp_key() -> SampledGraphKey {
13897 // temperature 1.0, filters off — today's serve default, the shape that parks a graph.
13898 SampledGraphKey::new(12345, 1.0, 3, 0, 1.0, 0.0, false)
13899 }
13900
13901 /// THE COLLISION. Two requests that differ ONLY in the truncation filters shared one key, so
13902 /// a parked pure-temp graph survived into a filtered request and the launch site launched it.
13903 #[test]
13904 fn vendor_filters_change_the_key() {
13905 let parked = pure_temp_key();
13906 // qwen3.8 generation_config.json — what the vendor-default flip makes the default shape.
13907 let vendor = SampledGraphKey::new(12345, 1.0, 3, 20, 0.95, 0.0, false);
13908 assert_eq!(
13909 legacy_key(&parked),
13910 legacy_key(&vendor),
13911 "pre-fix key collided: this is the bug, and the reason a test asserts on it",
13912 );
13913 assert_ne!(parked, vendor, "post-fix key must separate the two regimes");
13914 assert!(parked.pure_temp());
13915 assert!(!vendor.pure_temp());
13916 }
13917
13918 /// Each distribution-shaping field alone is enough to drop the parked graph.
13919 #[test]
13920 fn every_filter_field_is_keyed() {
13921 let base = pure_temp_key();
13922 for (what, other) in [
13923 (
13924 "top_k",
13925 SampledGraphKey::new(12345, 1.0, 3, 20, 1.0, 0.0, false),
13926 ),
13927 (
13928 "top_p",
13929 SampledGraphKey::new(12345, 1.0, 3, 0, 0.95, 0.0, false),
13930 ),
13931 (
13932 "min_p",
13933 SampledGraphKey::new(12345, 1.0, 3, 0, 1.0, 0.05, false),
13934 ),
13935 (
13936 "penalties",
13937 SampledGraphKey::new(12345, 1.0, 3, 0, 1.0, 0.0, true),
13938 ),
13939 ] {
13940 assert_ne!(base, other, "{what} must be part of the key");
13941 assert!(!other.pure_temp(), "{what} leaves the pure-temp regime");
13942 assert_eq!(
13943 legacy_key(&base),
13944 legacy_key(&other),
13945 "{what} was invisible to the pre-fix key",
13946 );
13947 }
13948 }
13949
13950 /// The baked constants stay keyed (this half was always right — regression cover for it).
13951 #[test]
13952 fn baked_constants_stay_keyed() {
13953 let base = pure_temp_key();
13954 assert_ne!(
13955 base,
13956 SampledGraphKey::new(999, 1.0, 3, 0, 1.0, 0.0, false),
13957 "seed"
13958 );
13959 assert_ne!(
13960 base,
13961 SampledGraphKey::new(12345, 0.7, 3, 0, 1.0, 0.0, false),
13962 "temp"
13963 );
13964 assert_ne!(
13965 base,
13966 SampledGraphKey::new(12345, 1.0, 4, 0, 1.0, 0.0, false),
13967 "k"
13968 );
13969 // bitwise on temperature: 0.7f32 vs the same value re-derived must NOT differ.
13970 assert_eq!(
13971 SampledGraphKey::new(1, 0.7, 3, 0, 1.0, 0.0, false),
13972 SampledGraphKey::new(1, 7.0 / 10.0, 3, 0, 1.0, 0.0, false),
13973 );
13974 }
13975
13976 /// THE LOAD-BEARING HALF OF THE SEED DECISION (lane/session-resume-sampler-predicate-
13977 /// 20260820). The whole-session resume predicate deliberately does NOT compare `seed`: an
13978 /// omitted serve `seed` draws fresh per-request entropy, so comparing it would refuse every
13979 /// seed-omitting sampled conversation. That is only sound because the one piece of parked state
13980 /// that BAKES the seed — this graph — is re-keyed on it, so a seed change drops and recaptures.
13981 ///
13982 /// This test is the other end of that argument, asserted here rather than remembered in a
13983 /// comment: if a future change dropped `seed` from the key, the resume predicate's exclusion
13984 /// would silently become the unsound thing it is documented not to be.
13985 /// (Paired with `seed_alone_does_not_refuse` in `memra-sampling`.)
13986 #[test]
13987 fn seed_alone_still_rekeys_the_draft_graph() {
13988 let parked = pure_temp_key();
13989 let reseeded = SampledGraphKey::new(999, 1.0, 3, 0, 1.0, 0.0, false);
13990 assert_ne!(
13991 parked, reseeded,
13992 "a seed-only change MUST drop the parked sampled graph — the resume predicate's \
13993 decision not to compare seed rests on exactly this",
13994 );
13995 // Same regime on both sides: the drop is a recapture, not a fall to the eager chain
13996 // because of a filter difference.
13997 assert!(parked.pure_temp() && reseeded.pure_temp());
13998 }
13999
14000 /// `pure_temp()` is the capture guard's predicate, computed from the key so the two cannot
14001 /// drift. The equality below is the invariant the launch-site guard asserts: identical keys
14002 /// agree on the regime, so a graph that survives the drop is legal to launch.
14003 #[test]
14004 fn equal_keys_agree_on_the_regime() {
14005 let a = SampledGraphKey::new(7, 0.8, 3, 20, 0.95, 0.0, false);
14006 let b = SampledGraphKey::new(7, 0.8, 3, 20, 0.95, 0.0, false);
14007 assert_eq!(a, b);
14008 assert_eq!(a.pure_temp(), b.pure_temp());
14009 // top_p slightly above 1.0 (a client sending 1.0 exactly, or an operator default) is
14010 // still the unfiltered regime, matching the original `sp.top_p >= 1.0` test.
14011 assert!(SampledGraphKey::new(7, 0.8, 3, 0, 1.0, 0.0, false).pure_temp());
14012 assert!(SampledGraphKey::new(7, 0.8, 3, 0, 1.5, -1.0, false).pure_temp());
14013 }
14014
14015 /// MEMRA_DEBUG_SPEC on a SAMPLED spec request past round 0: the print must render without
14016 /// indexing the empty greedy `preds` vector (it panicked the GPU worker before this lane).
14017 #[test]
14018 fn debug_print_survives_the_sampled_arm() {
14019 // round >= 1 with a pending bonus == base 1, sampled == `preds` empty.
14020 assert_eq!(debug_t_pred0(true, 1, 4242, &[]), "n/a");
14021 assert_eq!(debug_t_pred0(true, 2, 4242, &[]), "n/a");
14022 // round 0 without a pending bonus still reports last_pred, in both arms.
14023 assert_eq!(debug_t_pred0(true, 0, 4242, &[]), "4242");
14024 assert_eq!(debug_t_pred0(false, 0, 4242, &[7, 8]), "4242");
14025 // greedy keeps the real prediction it always printed.
14026 assert_eq!(debug_t_pred0(false, 1, 4242, &[7, 8]), "7");
14027 assert_eq!(debug_t_pred0(false, 2, 4242, &[7, 8]), "8");
14028 }
14029}