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
315/// PROJECTED REMAINING GROWTH of the verify-graph pool, in bytes (lane/hermes-perf-fixes,
316/// 2026-08-23 — the admission accounting the "pool dwarfs spec admission reserve" finding
317/// asks for). The pool was measured at 8,852 MiB at storm-complete on the q38 export while
318/// admission's transient floor (`SPEC_SHRINK_RESERVE`) is 1.5 GiB and never charged for it:
319/// sessions admitted while the pool is cold overcommit VRAM the pool WILL hold, because the
320/// pool grows monotonically (no eviction by design) and is model-owned across sessions.
321///
322/// SELF-MEASURING, no per-model constant (generic-model law — the 8,852 MiB is a q38 number
323/// and proves nothing about another export): the debt is remaining capture slots x the
324/// MARGINAL bytes a capture adds to this device's graph mem pool.
325///
326/// MARGINAL, NOT MEAN — measured correction (box9 on-box receipt, 2026-08-23). The first
327/// version of this used the mean (`reserved / captures`) and the live serve log showed why
328/// that is wrong: with the pool's reservation flat at ~33.6 MiB across captures 1..3, the
329/// mean-based debt printed **8,556 MB, then 4,261, then 2,830** — it extrapolated capture
330/// #1's ONE-TIME shared allocation (staging buffers, stash slabs, pointer tables: sized
331/// once per pool, shared by every key) across all 256 slots. An 8.5 GB phantom reserve at
332/// boot can refuse admissions that would have fit, which is a worse defect than the
333/// under-charge this accounting exists to remove. The marginal reading prices what an
334/// ADDITIONAL key actually costs: two observations `(captures, reserved)` give
335/// `(r1 - r0) / (c1 - c0)`, which is ~0 on an export whose pool does not grow per key and
336/// tracks real growth on one that does.
337///
338/// BOOTSTRAP (only one observation so far, so growth is unmeasurable): reserve one more
339/// pool's worth — `min(remaining x mean, reserved)`. "We have measured `reserved` bytes for
340/// `captures` keys; until growth is measurable, assume at most a doubling" is fail-safe in
341/// the same direction as the old rule without the 255x extrapolation.
342///
343/// Before the FIRST capture the debt is 0 (a single capture lands well inside the existing
344/// 1.5 GiB floor). `cap` is the intrinsic freeze ceiling (`MEMRA_DSPARK_VG_MAX`; =0 freeze
345/// valve => the pool cannot grow => debt 0); at or past the cap the pool FREEZES, so the
346/// debt is 0 there too.
347pub fn dspark_vg_debt_projection(
348 captures: usize,
349 cap: usize,
350 reserved_bytes: usize,
351 prev: Option<(usize, usize)>,
352) -> usize {
353 if captures == 0 || cap == 0 {
354 return 0;
355 }
356 let remaining = cap.saturating_sub(captures);
357 if remaining == 0 {
358 return 0;
359 }
360 match prev {
361 // marginal growth between two observations of the same pool
362 Some((c0, r0)) if captures > c0 => {
363 let marginal = reserved_bytes.saturating_sub(r0) / (captures - c0);
364 remaining.saturating_mul(marginal)
365 }
366 // bootstrap: at most one more pool's worth
367 _ => remaining
368 .saturating_mul(reserved_bytes / captures)
369 .min(reserved_bytes),
370 }
371}
372/// Engine-bundle slice 4 (fa-execupdate lane, DSF-ROUNDCOST-20260820 §6 close: "the
373/// residual gap lives in the FULL-ATTENTION per-row section"), DEFAULT ON —
374/// `MEMRA_DSPARK_FA_ROWS=0` reverts to the per-row loop: when every row of a verify
375/// round takes the v4-seqs arm on ONE `fa_split_keys` rung (the straddle law, evaluated
376/// at the round's first and last t_kv — both eligibility gates are intervals in t_kv),
377/// the qwen35 t-parallel verify's per-row KV-append + fa-decode loop collapses into the
378/// z-batched serving twins: ONE `append_quantize_kv_q8_0_q5_1_seqs` + ONE
379/// `fa_decode_vec_q_seqs_v4` + ONE combine per full-attention layer, replacing
380/// T x (4 dtod row copies + append + 3 memsets + main + combine) launches. Bytes are
381/// pinned by the batched-tick increment-2 kernel-check (seqs-vs-per-seq-loop bit
382/// identity: per-row T_kv derives in-kernel from pos_seq[z]; splits >= ns_eff write the
383/// empty partial the combine never reads, so the shared n_splits_max stride changes no
384/// bytes) and re-gated e2e by this lane's battery.
385pub(crate) fn dspark_fa_rows_on() -> bool {
386 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
387 *ON.get_or_init(|| {
388 std::env::var("MEMRA_DSPARK_FA_ROWS")
389 .map(|v| v != "0")
390 .unwrap_or(true)
391 })
392}
393
394/// `t_pred0` for the `MEMRA_DEBUG_SPEC` per-round print, sampled-safe.
395///
396/// `generate_spec_inner2` fills its `preds` vector ONLY on the greedy path (`if !sampled`), and
397/// the per-round debug print was the sole consumer in the sampled arm: `t_pred(0)` survives round
398/// 0 (`base == 0` returns `last_pred`) and from round 1 (`base == 1`, a pending bonus) indexes an
399/// EMPTY vector — `index out of bounds: the len is 0 but the index is 0`, in the GPU worker
400/// thread, which then respawns and reloads weights while the request dies. So any sampled spec
401/// request longer than one round used to kill the worker whenever `MEMRA_DEBUG_SPEC` was set:
402/// the flag crashed precisely the regime it exists to investigate.
403///
404/// Fixed at the print site, not inside the closure, so the greedy accept walk keeps its strict
405/// indexing (an out-of-range pred there is a real bug and must still be loud).
406fn debug_t_pred0(sampled: bool, base: usize, last_pred: u32, preds: &[u32]) -> String {
407 if base == 0 {
408 return last_pred.to_string();
409 }
410 match preds.get(base - 1) {
411 Some(p) => p.to_string(),
412 // sampled: the greedy per-column argmax was never run for this round.
413 None => {
414 debug_assert!(
415 sampled,
416 "greedy spec: preds[{}] missing at base {base}",
417 base - 1
418 );
419 "n/a".to_string()
420 }
421 }
422}
423
424/// `MEMRA_SKEY_PROBE=1` — sampled-draft-graph key probe (lane/graph-s-key-exactness-20260819).
425///
426/// Reports, per burst and per round, which draft chain the sampled arm chose and under which
427/// filter regime, plus the ONE observable that separates a legal filtered draft from a stale
428/// pure-temp graph replayed under filters: an accept test whose gathered `q` is exactly 0.
429/// A draft token sampled from the FILTERED softmax can never gather q=0 (it was drawn from the
430/// kept set), so `q=0` in the verify means the draft came from a distribution the verify does
431/// not believe in — and `u * 0 < p` then accepts it unconditionally.
432///
433/// Its own env var, deliberately NOT `MEMRA_DEBUG_SPEC`: that flag panicked the GPU worker on
434/// any sampled spec request past round 0 until this lane fixed it (§2 of the bank note).
435pub(crate) fn skey_probe() -> bool {
436 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
437 *ON.get_or_init(|| std::env::var("MEMRA_SKEY_PROBE").as_deref() == Ok("1"))
438}
439
440/// GRAMMAR HOOK for constrained spec decode (lane/constrained-full, 2026-08-03). The engine
441/// stays llguidance-agnostic: the server adapts its per-session grammar state behind this
442/// trait. CONTRACT (the verify-side truncation rule — token-identical to constrained plain
443/// greedy decode): the exactness walk runs UNMASKED first; the hook then (a) truncates
444/// acceptance at the first grammar-illegal accepted token, and (b) when the truncation fired
445/// or the bonus is illegal, the engine recomputes that slot as the MASKED argmax of the
446/// target's own verify column (an unmasked argmax that is grammar-legal IS the masked argmax
447/// — masking only removes tokens — so the common case pays nothing). `consume` advances the
448/// state with each EMITTED token in order; EOS handling is the implementor's job (skip).
449pub trait SpecConstraint {
450 /// -inf the current state's banned ids on a HOST logits row (prompt-tail / init-feed
451 /// masked argmax).
452 fn mask_logits(&mut self, logits: &mut [f32]) -> Result<(), String>;
453 /// Packed 32-bit bitset words of the CURRENT state's allowed set (device-mask form).
454 fn mask_words(&mut self) -> Result<Vec<u32>, String>;
455 /// Is `tok` consumable in the CURRENT state?
456 fn is_allowed(&mut self, tok: u32) -> Result<bool, String>;
457 /// Advance the state with an emitted token.
458 fn consume(&mut self, tok: u32) -> Result<(), String>;
459
460 // --- DRAFT-SIDE MASKING (lane/draft-mask, 2026-08-04) ---
461 // The drafter proposed grammar-illegal tokens under tight schemas, so verify-side
462 // truncation cut nearly every round (measured acceptance 0.467-0.513 tight vs 0.62-0.82
463 // loose, research/constrained-full-20260803). These three methods let the engine mask the
464 // DRAFT model's own sampling with the grammar's legal set, so proposals are legal by
465 // construction. The state they walk is a SPECULATIVE CLONE of the session matcher — the
466 // real state is advanced only by `consume` (emitted tokens), so verify-side truncation
467 // stays the correctness backstop and the emitted stream is unchanged by construction
468 // (an accepted draft is the target's unmasked argmax AND grammar-legal, hence the masked
469 // argmax; a cut slot is recomputed as the masked argmax either way).
470 // Default impls = feature OFF (pre-lane behaviour: unmasked drafts).
471
472 /// Is draft-side masking available on this hook? Probed ONCE per burst, before the draft
473 /// graph is captured (the mask is an in-graph node — its presence is a capture-time shape).
474 fn draft_mask_enabled(&self) -> bool {
475 false
476 }
477 /// Start a draft chain: clone the CURRENT (committed) grammar state into the speculative
478 /// slot. Called once per spec round, before the first draft position.
479 fn draft_begin(&mut self) -> Result<(), String> {
480 Ok(())
481 }
482 /// Packed 32-bit bitset words of the SPECULATIVE state's allowed set (target-vocab ids),
483 /// for the draft position about to be sampled. `None` = draft masking off (no-op).
484 fn draft_mask_words(&mut self) -> Result<Option<Vec<u32>>, String> {
485 Ok(None)
486 }
487 /// Advance the SPECULATIVE state with a PROPOSED draft token. `false` = the chain cannot
488 /// continue (EOS proposed, or an unmasked position proposed something illegal) — the
489 /// engine stops drafting; the token already pushed still goes through verify.
490 fn draft_advance(&mut self, _tok: u32) -> Result<bool, String> {
491 Ok(false)
492 }
493}
494
495/// DRAFT-MASK UPLOAD (lane/draft-mask): pull the speculative state's allowed set (TARGET-id
496/// space) from the hook, project it into the DRAFT head's vocab space, and upload it into the
497/// stable device buffer the draft chain reads. Returns false when the chain must stop drafting:
498/// the hook handed out no mask, or NO draft-vocab row is grammar-legal at this position (a
499/// trimmed FR-Spec head genuinely cannot propose a legal token there — masking it would leave
500/// a fully-banned row whose argmax is meaningless, so the round drafts fewer tokens and the
501/// verify emits the masked argmax as usual).
502fn upload_draft_mask(
503 e: &Engine,
504 c: &mut dyn SpecConstraint,
505 dst: &mut CudaSlice<u32>,
506 d2t: Option<&Vec<u32>>,
507 d_vocab: usize,
508 words: usize,
509) -> Result<bool, Box<dyn std::error::Error>> {
510 let Some(tw) = c
511 .draft_mask_words()
512 .map_err(|e2| format!("constraint: {e2}"))?
513 else {
514 return Ok(false);
515 };
516 let bit = |t: usize| -> bool {
517 let w = t >> 5;
518 w < tw.len() && (tw[w] >> (t & 31)) & 1 == 1
519 };
520 let mut buf = vec![0u32; words];
521 match d2t {
522 // TRIMMED draft head: row i proposes target id d2t[i] — permute the mask accordingly.
523 Some(map) => {
524 for (i, &t) in map.iter().enumerate().take(d_vocab) {
525 if bit(t as usize) {
526 buf[i >> 5] |= 1u32 << (i & 31);
527 }
528 }
529 }
530 // UNTRIMMED: draft ids ARE target ids; the packed words transfer verbatim (a short
531 // mask leaves the padded tail zeroed == banned, same rule as constrained::apply_mask).
532 None => {
533 let n = tw.len().min(words);
534 buf[..n].copy_from_slice(&tw[..n]);
535 }
536 }
537 if buf.iter().all(|w| *w == 0) {
538 return Ok(false);
539 }
540 e.htod_u32_into(dst, &buf)?;
541 Ok(true)
542}
543
544/// Keep the full token-embedding table in host memory and upload only the rows needed by each
545/// MTP/verify step. This is an exact memory-capacity seam for very large BF16 vocab tables: host
546/// gather expands the same source bits to f32, and only O(T*n_embd) bytes cross PCIe per step.
547/// CUDA-graph/round-stream draft paths require device token ids and therefore stay disabled.
548pub(crate) fn spec_host_embd() -> bool {
549 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
550 *ON.get_or_init(|| std::env::var("MEMRA_SPEC_HOST_EMBD").as_deref() == Ok("1"))
551}
552
553/// VERIFY-TIER TRUNK LAUNCH-FUSION (default ON since 2026-07-09; MEMRA_SPEC_FUSED_T=0 reverts — lane/close35b): extend
554/// the t=1 fused2/fused3 Q8_0 trunk launches to the batched verify tier (t=2-4, the K=1..3
555/// verify shapes). At t>1 the trunk pairs/triples (35B wqkv+wqkv_gate, wq/wk/wv,
556/// gate_shexp+up_shexp) each run a separate `matmul_decode_exact` — one q8_1 re-quantize of the
557/// SAME activation plus one _b2/_b4 launch per tensor. The fused twins share ONE quantize and
558/// ONE launch per group; per (tensor,token,row) the kernel body is q8_0_mmvq_batched verbatim
559/// with the identical row mapping -> BIT-IDENTICAL by construction (kernel-check pins it,
560/// run-spec K=1..8 + acceptance identity arbitrate e2e).
561pub(crate) fn spec_fused_t() -> bool {
562 static F: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
563 // DEFAULT ON since 2026-07-09 (MEMRA_SPEC_FUSED_T=0 reverts): verify t=2-4 trunk launch-fusion
564 // (fused2/fused3 Q8_0 batched twins, bit-identical by construction — m=1 block-offset split on
565 // the batched body). m=2 marginal token 2117->1762us; 35B daily: p3 +3.7% (crosses llama), p2 +5%.
566 *F.get_or_init(|| {
567 std::env::var("MEMRA_SPEC_FUSED_T")
568 .map(|v| v != "0")
569 .unwrap_or(true)
570 })
571}
572
573/// zeros/uninit switch for verify-path buffers that are FULLY OVERWRITTEN before any read.
574/// Only call this on such buffers — the lean contract is "identical bytes by construction".
575fn vbuf(e: &Engine, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
576 if spec_lean() { e.uninit(n) } else { e.zeros(n) }
577}
578
579/// Scratch KV for the MTP block (one full-attn layer).
580///
581/// PERSISTENT MODE (default, 2026-07-03 — the acceptance lever): sized cap = max_ctx and kept in
582/// sync with the COMMITTED sequence — slot p holds the MTP block's K/V for committed token p
583/// (roped p+1, the chain's rope convention), so the draft chain's self-attention sees the FULL
584/// committed history instead of only the current round's 1..K+1 chain tokens (the reference
585/// engine's "mtp_update" design). Entries come from two sources:
586/// - chain appends: accepted positions KEEP their chain-computed entries (embedding exact,
587/// hidden chain-approximate — the reference engine accepts the same);
588/// - `mtp_kv_fill` batches: prompt positions + the last-draft position on full accept, computed
589/// from EXACT trunk hiddens (K/V-only MTP-block pass, no attention/FFN/lm_head).
590/// Rejected drafts / p-min extras / pseudo-seed appends are all discarded by the round-start
591/// `set_len` truncation (the KvLayer len mechanism — §C rollback for the draft side).
592/// Multi-turn spec-decode session (2026-07-05): trunk Cache + persistent MTP draft scratch +
593/// the committed token list, alive across generate_spec_session calls. Turn N+1 primes ONLY its
594/// suffix (chunked continuation prime over the quantized past) and mtp_kv_fill's its suffix rows,
595/// then the round loop runs unchanged. `last_h` carries the pre-output_norm hidden of the last
596/// committed row across turns (the predecessor-pairing seed + fill anchor).
597/// Per-request sampling config for the sampled-spec serve path.
598#[derive(Clone, Copy, Debug)]
599pub struct SpecSampling {
600 pub temp: f32,
601 pub seed: u64,
602 pub top_k: i32, // 0 = off
603 pub top_p: f32, // 1.0 = off
604 pub min_p: f32, // 0.0 = off
605 pub penalty_last_n: usize, // 0 = penalties off
606 pub penalty_repeat: f32,
607 pub penalty_freq: f32,
608 pub penalty_present: f32,
609}
610
611impl SpecSampling {
612 /// Non-identity penalties requested — THE `pen_on` predicate (one definition; the
613 /// same group-off rule `SamplerIdentity::of` canonicalizes: a window with neutral
614 /// coefficients is penalties-absent). Both spec routes and the dspark accept walk
615 /// key their penalty arms off this.
616 pub fn pen_on(&self) -> bool {
617 self.penalty_last_n > 0
618 && (self.penalty_repeat != 1.0
619 || self.penalty_freq != 0.0
620 || self.penalty_present != 0.0)
621 }
622}
623
624/// Host Philox4x32-10 uniform in (0,1) — mirrors spec_sample.cu's `philox4`/`u01` with the
625/// ctr_lo tag 0xFFFF_FFFE, so the host accept-test stream never collides with any device
626/// sampling event (device Gumbel uses (i>>2, stream_pos); device residual uses 0xFFFF_FFFD).
627/// One value per (seed, ctr) EVENT; callers own the counter discipline. Extracted verbatim
628/// from generate_spec_inner2's closure for the dspark sampled-admission walk (the two paths
629/// MUST consume the identical stream construction — two ad-hoc Philox copies drifting apart
630/// is a distributional bug, not a style problem).
631pub(crate) fn host_u01(seed: u64, ctr: u32) -> f32 {
632 let (m0, m1) = (0xD2511F53u32, 0xCD9E8D57u32);
633 let (mut c0, mut c1, mut c2, mut c3) = (0xFFFF_FFFEu32, ctr, 0u32, 0u32);
634 let (mut k0, mut k1) = ((seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
635 for _ in 0..10 {
636 let (h0, l0) = (((m0 as u64 * c0 as u64) >> 32) as u32, m0.wrapping_mul(c0));
637 let (h1, l1) = (((m1 as u64 * c2 as u64) >> 32) as u32, m1.wrapping_mul(c2));
638 let (n0, n1, n2, n3) = (h1 ^ c1 ^ k0, l1, h0 ^ c3 ^ k1, l0);
639 c0 = n0;
640 c1 = n1;
641 c2 = n2;
642 c3 = n3;
643 k0 = k0.wrapping_add(0x9E3779B9);
644 k1 = k1.wrapping_add(0xBB67AE85);
645 }
646 (c0 as f32 + 1.0) * (1.0 / 4294967296.0)
647}
648
649/// Tracked draft positions for [`SpecTelemetry`] (serve K defaults to 3; the run-spec gate
650/// sweeps K=1..8, and MEMRA_SPEC_CAPMAX defaults to 7 — 8 covers every tuned config).
651pub const SPEC_TELEM_POS: usize = 8;
652
653/// Always-on per-draft-position acceptance telemetry (lane/accept-telemetry, 2026-08-05 —
654/// the llama.cpp #26389 / vLLM spec-decode counter schema, upstream-sweeps 2026-08-05).
655/// Lives on the [`SpecSession`] and accumulates across bursts; the serve worker diffs a
656/// stashed copy per burst for its per-model /metrics aggregation and per-request usage.
657/// Same normalization as the `[spec-stats]` line: p-min-discarded chain tokens are counted
658/// in NEITHER drafted nor accepted.
659#[derive(Clone, Copy, Default, Debug)]
660pub struct SpecTelemetry {
661 /// verify rounds completed (a round-stream burst counts each of its M rounds).
662 pub rounds: u64,
663 /// tokens drafted / accepted across all rounds.
664 pub drafted: u64,
665 pub accepted: u64,
666 /// how often draft position j (0-based within a round's chain) was offered / accepted.
667 /// Positions >= SPEC_TELEM_POS are untracked (totals still count them). The opt-in
668 /// round-stream arm (MEMRA_SPEC_STREAM=1) reads back only totals, so under it these
669 /// arrays cover the standard-path rounds only and their sums may undercount the totals.
670 pub pos_drafted: [u64; SPEC_TELEM_POS],
671 pub pos_accepted: [u64; SPEC_TELEM_POS],
672}
673
674impl SpecTelemetry {
675 /// Fieldwise `self - prev` — the worker's per-burst delta off a copy stashed before the
676 /// burst call. Saturating: a caller diffing against the wrong snapshot gets zeros, not
677 /// a wrapped counter.
678 pub fn delta_since(&self, prev: &SpecTelemetry) -> SpecTelemetry {
679 let mut d = SpecTelemetry {
680 rounds: self.rounds.saturating_sub(prev.rounds),
681 drafted: self.drafted.saturating_sub(prev.drafted),
682 accepted: self.accepted.saturating_sub(prev.accepted),
683 ..Default::default()
684 };
685 for j in 0..SPEC_TELEM_POS {
686 d.pos_drafted[j] = self.pos_drafted[j].saturating_sub(prev.pos_drafted[j]);
687 d.pos_accepted[j] = self.pos_accepted[j].saturating_sub(prev.pos_accepted[j]);
688 }
689 d
690 }
691 /// Fieldwise `self += d` — the worker's per-model aggregation.
692 pub fn merge(&mut self, d: &SpecTelemetry) {
693 self.rounds += d.rounds;
694 self.drafted += d.drafted;
695 self.accepted += d.accepted;
696 for j in 0..SPEC_TELEM_POS {
697 self.pos_drafted[j] += d.pos_drafted[j];
698 self.pos_accepted[j] += d.pos_accepted[j];
699 }
700 }
701
702 /// Mean accepted draft-prefix length per verify round (tau).
703 pub fn tau(&self) -> f64 {
704 if self.rounds > 0 {
705 self.accepted as f64 / self.rounds as f64
706 } else {
707 0.0
708 }
709 }
710}
711
712/// Session-lifetime atomic acceptance counters. The verifier records only after the greedy or
713/// rejection-sampling walk has resolved on the host, so these relaxed increments add no GPU
714/// launch, synchronization, allocation, or ordering dependency to the numeric path.
715struct SpecTelemetryCounters {
716 rounds: AtomicU64,
717 drafted: AtomicU64,
718 accepted: AtomicU64,
719 pos_drafted: [AtomicU64; SPEC_TELEM_POS],
720 pos_accepted: [AtomicU64; SPEC_TELEM_POS],
721}
722
723impl Default for SpecTelemetryCounters {
724 fn default() -> Self {
725 Self {
726 rounds: AtomicU64::new(0),
727 drafted: AtomicU64::new(0),
728 accepted: AtomicU64::new(0),
729 pos_drafted: std::array::from_fn(|_| AtomicU64::new(0)),
730 pos_accepted: std::array::from_fn(|_| AtomicU64::new(0)),
731 }
732 }
733}
734
735impl SpecTelemetryCounters {
736 fn record_round(&self, drafted: usize, accepted: usize) {
737 debug_assert!(accepted <= drafted);
738 self.rounds.fetch_add(1, Ordering::Relaxed);
739 self.drafted.fetch_add(drafted as u64, Ordering::Relaxed);
740 self.accepted.fetch_add(accepted as u64, Ordering::Relaxed);
741 for counter in self.pos_drafted.iter().take(drafted) {
742 counter.fetch_add(1, Ordering::Relaxed);
743 }
744 for counter in self.pos_accepted.iter().take(accepted) {
745 counter.fetch_add(1, Ordering::Relaxed);
746 }
747 }
748
749 /// Round-stream keeps each round's accept length on device; retain exact scalar totals while
750 /// leaving the per-position arrays untouched, matching the pre-existing telemetry contract.
751 fn record_totals(&self, rounds: usize, drafted: usize, accepted: usize) {
752 self.rounds.fetch_add(rounds as u64, Ordering::Relaxed);
753 self.drafted.fetch_add(drafted as u64, Ordering::Relaxed);
754 self.accepted.fetch_add(accepted as u64, Ordering::Relaxed);
755 }
756
757 fn snapshot(&self) -> SpecTelemetry {
758 SpecTelemetry {
759 rounds: self.rounds.load(Ordering::Relaxed),
760 drafted: self.drafted.load(Ordering::Relaxed),
761 accepted: self.accepted.load(Ordering::Relaxed),
762 pos_drafted: std::array::from_fn(|j| self.pos_drafted[j].load(Ordering::Relaxed)),
763 pos_accepted: std::array::from_fn(|j| self.pos_accepted[j].load(Ordering::Relaxed)),
764 }
765 }
766}
767
768pub struct SpecSession {
769 pub(crate) cache: Cache,
770 pub(crate) scratch: MtpScratch,
771 /// Every token whose state the caches hold, in order (prompt turns + generated), INCLUDING
772 /// overshoot: spec commits accepted drafts past max_new; those rows are in the caches, so the
773 /// session must count them. Callers render output from this, not from their own echo.
774 pub committed: Vec<u32>,
775 /// Pre-output_norm hidden of the LAST committed row (device). None before the first turn.
776 pub(crate) last_h: Option<CudaSlice<f32>>,
777 /// Greedy argmax predicting the token AFTER committed.last() (from the last turn's final
778 /// logits). Fuels empty-suffix continuation bursts (serve): the next turn emits this token
779 /// first, feeds it, and the round loop resumes without any prime. None before the first turn.
780 pub next_pred: Option<u32>,
781 /// SAMPLED-SPEC stream continuity across bursts: Philox event counters persist here so a
782 /// session's randomness never repeats between generate_spec_session calls. (0,0) at admit.
783 pub sctr: u32,
784 pub uctr: u32,
785 /// PERSISTENT DRAFT-GRAPH CONTEXT (2026-08-01, the serve-burst fixed-cost fix): the captured
786 /// draft graph(s) + every device I/O buffer they bake, carried ACROSS generate_spec_session
787 /// calls. Before this, every serve burst re-captured the draft graph (2 warmup forwards +
788 /// instantiate) — measured ~16ms/burst on H100 q27 (MEMRA_SPEC_BURST sweep,
789 /// research/spec-serving-20260801). None before the first turn; error paths drop it
790 /// (next burst recaptures — serve retires errored sessions anyway).
791 pub(crate) draft_ctx: Option<DraftGraphCtx>,
792 /// PENDING-CARRY across bursts (2026-08-01, the serve burst-boundary fix): the bonus token
793 /// emitted by the last round but NOT committed to the caches. The old tail committed it with
794 /// a solo T=1 trunk pass (+ draft fill), and the next burst's setup fed the stashed next_pred
795 /// with ANOTHER solo pass — 2x ~11.5ms/burst measured on H100 q27 ([spec-setup] trace).
796 /// Carrying it lets the next empty-suffix greedy burst consume it as round-0 verify col 0,
797 /// exactly like a mid-burst full-accept boundary (no solo passes). INVARIANT: when set,
798 /// `committed` (== cache rows) EXCLUDES this token although it was already emitted in the
799 /// last burst's output, and `last_h` holds the hidden of the last COMMITTED row (its
800 /// predecessor — the chain-seed/fill anchor). `next_pred` is None (unknown without the
801 /// commit pass). Non-empty-suffix or sampled turns must flush first (spec_flush_pending);
802 /// generate_spec_session_sampled does this at entry, and serve parks only flushed sessions.
803 pub pending_tok: Option<u32>,
804 /// SESSION-AFFINITY TURN CHECKPOINT (lane/session-affinity, 2026-08-05): the state at this
805 /// turn's PROMPT-END boundary, retained so a later turn can REWIND here. See
806 /// [`SpecCheckpoint`]. Refreshed by every non-empty prime; None until the first one, and on
807 /// a rig too tight to hold it (a failed capture is silent — resume just isn't available).
808 pub(crate) turn_ckpt: Option<SpecCheckpoint>,
809 /// Session-lifetime acceptance telemetry. Relaxed atomics update at the host-side round
810 /// accounting the loop already does — no syncs, no allocation. NOTE a
811 /// pool-resumed session carries the PREVIOUS requests' counts; per-request consumers
812 /// diff with [`SpecTelemetry::delta_since`] around each burst.
813 telem: SpecTelemetryCounters,
814 /// PREFIX-CACHE publication request (lane/spec-prefix-cache): worker sets this to the
815 /// miss-LCP boundary before a cold burst; the prime captures at exactly that split (it must
816 /// coincide with the burst's `prime_split` or no capture happens). One-shot: consumed by the
817 /// prime, result lands in `boundary_captures`.
818 pub capture_at: Option<usize>,
819 /// The captures the last prime produced (see [`SpecBoundaryCapture`]). Worker drains them
820 /// post-burst to assemble prefix entries. A failed capture is silent, like `turn_ckpt` —
821 /// publication just isn't available for that request. Plural since
822 /// lane/frspec-multiturn-cache (2026-08-21): a cold burst can capture BOTH the miss-LCP
823 /// split (the shared-prefix class) and the stable pre-generation boundary (the
824 /// next-turn re-render class) — one entry per stop, exactly the boundary set the plain
825 /// prefill tick publishes/checkpoints.
826 pub boundary_captures: Vec<SpecBoundaryCapture>,
827 /// STABLE-BOUNDARY TURN CHECKPOINT REQUEST (lane/frspec-multiturn-cache, 2026-08-21): the
828 /// ABSOLUTE committed-length position the next non-empty prime should capture `turn_ckpt`
829 /// at, instead of prompt-end. The worker sets it to the STABLE PRE-GENERATION boundary
830 /// (`plain_checkpoint_boundary` — before the live generation header the client rewrites),
831 /// porting the 2026-08-09 plain-tier fix: a prompt-end spec checkpoint includes the
832 /// template's live assistant-generation header (`<|im_start|>assistant\n<think>\n`), which
833 /// the NEXT turn's re-render replaces, so `affinity_match` diverged a couple tokens below
834 /// the checkpoint and the spec pool declined 100% of multi-turn agent traffic (measured:
835 /// `spec-affinity: declined (history diverged at 6811 of checkpoint 6813)`,
836 /// research/multiturn-cache-20260821 B4). One-shot, `capture_at` convention; None = legacy
837 /// prompt-end capture.
838 pub ckpt_at: Option<usize>,
839}
840impl SpecSession {
841 /// Context capacity of the session's caches (the server's ContextFull guard).
842 pub fn cache_max_ctx(&self) -> usize {
843 self.cache.max_ctx
844 }
845 /// Read access to the live trunk cache (lane/spec-prefix-cache): the worker slices
846 /// full-attn KV rows `[0..capture.pos)` out of it when publishing a boundary capture —
847 /// those rows are append-only for the session's lifetime (rollbacks never truncate below
848 /// the prime boundary), so no copy was taken at prime time.
849 pub fn cache_ref(&self) -> &Cache {
850 &self.cache
851 }
852 /// Read access to the persistent draft-scratch plane (lane/spec-on-cache-hit): the
853 /// worker slices rows `[0..capture.pos)` when publishing a boundary capture, exactly
854 /// like the trunk KV — draft rows below the prompt end are append-only for the
855 /// session's lifetime (the prime fill wrote them once; rollbacks reset `len_d` to the
856 /// committed length, never below the prime boundary, and the true-hidden refresh
857 /// rewrites generated positions only). Returns `(k, v, k_tok_bytes, v_tok_bytes)`.
858 /// None when the scratch is ring-backed (Step35 SWA — physical rows are not
859 /// prefix-addressable; the prefix cache already refuses that class end to end).
860 pub fn draft_plane_ref(&self) -> Option<(&CudaSlice<u8>, &CudaSlice<u8>, usize, usize)> {
861 if self.scratch.kv.ring.is_some() {
862 return None;
863 }
864 Some((
865 &self.scratch.kv.k,
866 &self.scratch.kv.v,
867 self.scratch.kv.k_tok_bytes,
868 self.scratch.kv.v_tok_bytes,
869 ))
870 }
871 /// Snapshot the session's process-local acceptance counters for per-burst diffing.
872 pub fn telemetry(&self) -> SpecTelemetry {
873 self.telem.snapshot()
874 }
875 /// Committed position this session can REWIND to (its retained prompt-end boundary), if any.
876 /// A request whose prompt matches `committed[..pos]` exactly can resume from here — see
877 /// `spec_rewind_to_checkpoint`.
878 pub fn rewind_pos(&self) -> Option<usize> {
879 self.turn_ckpt.as_ref().map(|c| c.pos)
880 }
881 /// Whether every ring-backed trunk/draft row needed by the retained checkpoint is resident.
882 pub fn rewind_is_resident(&self) -> bool {
883 self.turn_ckpt.as_ref().is_some_and(|ckpt| {
884 self.cache.can_rollback(&ckpt.snap, 0) && self.scratch.can_rewind_to(ckpt.pos)
885 })
886 }
887 /// Is this session in the DEMOTION-READY shape (see [`SpecSession::into_demoted`])?
888 /// `false` means a carried pending must be flushed first (`spec_flush_pending`), or the
889 /// session has never run a turn and has no prediction to hand over.
890 pub fn demote_ready(&self) -> bool {
891 self.pending_tok.is_none() && self.next_pred.is_some()
892 }
893 /// Does this session hold a carried pending bonus (flush required before a handoff/park)?
894 pub fn has_pending(&self) -> bool {
895 self.pending_tok.is_some()
896 }
897 /// Committed row count == cache rows (the session invariant), for the caller's own
898 /// `fed`-length cross-check at a handoff boundary.
899 pub fn committed_len(&self) -> usize {
900 self.committed.len()
901 }
902 /// DEMOTION HANDOFF (lane/spec-gate, 2026-08-07): consume this session and hand its trunk
903 /// cache + next-token prediction to the plain batched-decode path.
904 ///
905 /// WHY THIS IS EXACT (greedy). The invariant at a burst boundary is `cache.pos ==
906 /// committed.len()`: every committed row has trunk KV + recurrent state, exactly as a plain
907 /// tokenwise prime of the same `committed` sequence would have left it (that is the
908 /// session-tail contract, and the same property `spec_rewind_to_checkpoint` and the reuse
909 /// pool already rely on). `next_pred` is the argmax of the verify's logits for the LAST
910 /// committed row — and verify-column logits are bit-identical to plain decode's logits at
911 /// that position, because `matmul_decode_exact` bit-identity IS the basis of the greedy
912 /// accept walk. So handing (cache, next_pred) to the batched path continues the stream from
913 /// a state indistinguishable from one the batched path produced itself: the batched tick
914 /// emits `next_pred`, feeds it into this same cache, and decodes on.
915 ///
916 /// `None` when the session is not in the handoff shape — a carried pending (its bonus row is
917 /// NOT in the cache, so `spec_flush_pending` must commit it first) or no `next_pred` yet
918 /// (never bursted). Callers must not force it: a half-committed cache handed to the batched
919 /// path would silently skip a token.
920 ///
921 /// The MTP draft scratch, the persistent draft-graph context and the turn checkpoint are
922 /// DROPPED here (freeing their VRAM): the batched path never drafts, and this handoff is
923 /// one-way by design — there is no cheap symmetric re-promotion (rebuilding the draft KV
924 /// would mean an `mtp_kv_fill` over the whole committed history).
925 pub fn into_demoted(self) -> Option<(Cache, u32)> {
926 if self.pending_tok.is_some() {
927 return None;
928 }
929 let np = self.next_pred?;
930 debug_assert_eq!(
931 self.cache.pos,
932 self.committed.len(),
933 "demotion handoff: cache rows != committed tokens"
934 );
935 Some((self.cache, np))
936 }
937 /// Pool-resume hook (audit Q2): clear the parked draft-graph failure memoization so a
938 /// NEW request resuming this session gets one fresh capture chance — a transient-pressure
939 /// capture failure must not persist for the pool's whole lifetime (the TRT #16072 class).
940 /// Logs once iff a flag was actually set; a no-fallback resume is silent and free.
941 pub fn reset_graph_fallback_on_resume(&mut self) {
942 if let Some(line) = self
943 .draft_ctx
944 .as_mut()
945 .and_then(|c| c.failed.reset_on_resume())
946 {
947 eprintln!("{line}");
948 }
949 }
950}
951
952/// A session's PROMPT-END boundary state, the rewind target for session-affinity resume.
953///
954/// WHY THIS BOUNDARY, AND WHY IT IS THE ONLY ONE WORTH KEEPING. The rewrite class this lane
955/// exists for (a client that strips `<think>` blocks out of prior assistant turns) mutates the
956/// text the session GENERATED, never the prompt it was given. So turn N's prompt agrees with
957/// turn N-1's committed tokens up to almost exactly where turn N-1's generation began — the
958/// prompt-end boundary. Keeping a checkpoint there means the next turn re-primes only its own
959/// delta (the rewritten answer + the new user turn) instead of the whole conversation.
960///
961/// WHAT IT MUST HOLD. Full-attn KV is append-only and position-addressed, so rewinding it is a
962/// `len` truncation (no data). Linear-attn (GDN) conv/ssm state is mutated IN PLACE with no
963/// position index, so it must be a real device COPY — that copy is the entire reason a spec
964/// session could not previously rewind. The MTP draft scratch needs no copy either: its rows
965/// below the boundary were written by this turn's fill and are never revisited (the per-round
966/// true-hidden refresh only rewrites the CURRENT burst's committed positions), so rewinding it
967/// is also just a `len` reset. `last_h` is the hidden of the last row below the boundary — the
968/// predecessor-pairing anchor the next prime's fill reads for its first row.
969///
970/// COST: one `Cache::snapshot` per TURN, on a code path that already takes one per ROUND.
971pub(crate) struct SpecCheckpoint {
972 snap: crate::cache::CacheSnapshot,
973 /// Committed length at the boundary (== cache.pos there, the session invariant).
974 pos: usize,
975 /// Pre-output_norm hidden of row `pos - 1`.
976 last_h: CudaSlice<f32>,
977}
978
979/// PREFIX-CACHE BOUNDARY CAPTURE (lane/spec-prefix-cache, 2026-08-14): the state a spec session
980/// records at its cold-prime split so the WORKER can publish a cross-request prefix entry —
981/// the commit-gated-publication port (research/cache-spec-design-20260814/PORT-PLAN.md item 1).
982/// Only the pieces that are DESTROYED by continuing the prime need copies here: the in-place
983/// GDN conv/ssm states (via `Cache::snapshot`, same mechanism as [`SpecCheckpoint`]) and the
984/// boundary logits. Full-attn KV rows `[0..pos)` and draft-scratch rows `[0..pos)` are
985/// append-only for the session's lifetime (rollbacks never truncate below the prime boundary),
986/// so the worker slices those from the live caches post-burst instead of copying at prime time.
987pub struct SpecBoundaryCapture {
988 pub snap: crate::cache::CacheSnapshot,
989 /// Token boundary (== cache.pos at capture; == the worker's miss-LCP split).
990 pub pos: usize,
991 /// Full-vocab logits after the prefix prime — the entry's boundary logits.
992 pub logits: Vec<f32>,
993 /// Pre-output_norm trunk hidden of row `pos - 1` (lane/spec-on-cache-hit): the
994 /// predecessor-pairing anchor a RESTORED spec session's first suffix-fill row reads
995 /// (the `SpecSession::last_h` convention). Empty = unavailable (capture stays valid;
996 /// the fill's zeros row-0 fallback covers it at a bounded acceptance cost).
997 pub last_h: Vec<f32>,
998}
999
1000/// D2H one hidden row out of a `[T, n_embd]` prime hidden stack — the boundary anchor a
1001/// spec boundary capture carries for later restored-session fills. Failure is silent
1002/// (`turn_ckpt` convention): the capture publishes without an anchor.
1003fn capture_boundary_hidden(
1004 e: &Engine,
1005 h_rows: &CudaSlice<f32>,
1006 pos: usize,
1007 n_embd: usize,
1008) -> Vec<f32> {
1009 if pos == 0 || h_rows.len() < pos * n_embd {
1010 return Vec::new();
1011 }
1012 let Ok(mut row) = e.uninit(n_embd) else {
1013 return Vec::new();
1014 };
1015 if e.copy_view_into(
1016 &mut row,
1017 0,
1018 &h_rows.slice((pos - 1) * n_embd..pos * n_embd),
1019 n_embd,
1020 )
1021 .is_err()
1022 {
1023 return Vec::new();
1024 }
1025 e.dtoh(&row).unwrap_or_default()
1026}
1027
1028/// ROLLBACK DOOR for sampled BOUNDARY tokens (lane/sampled-spec-quality, 2026-08-19).
1029/// Default ON: the token a burst emits at its own boundary is drawn from the request's
1030/// sampler. `MEMRA_SPEC_SAMPLED_BOUNDARY=0` restores the pre-lane posture (an ARGMAX at
1031/// every boundary) without touching greedy, which is byte-unaffected either way.
1032pub fn spec_sampled_boundary_on() -> bool {
1033 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1034 *ON.get_or_init(|| std::env::var("MEMRA_SPEC_SAMPLED_BOUNDARY").as_deref() != Ok("0"))
1035}
1036
1037/// ROLLBACK DOOR for SESSION-SPANNING penalty history (lane/sampled-spec-quality).
1038/// Default ON: `pen_hist` is seeded from the session's committed tail, so repetition /
1039/// frequency / presence penalties see the whole stream. `MEMRA_SPEC_PEN_SESSION=0`
1040/// restores the pre-lane posture (each burst restarts the window from its own prompt
1041/// slice, i.e. from NOTHING on a continuation burst) — and with the door shut the worker
1042/// must keep refusing penalized sampled prefix-cache restores, because the restored
1043/// session's continuation burst is handed no prompt slice at all.
1044pub fn spec_pen_session_on() -> bool {
1045 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1046 *ON.get_or_init(|| std::env::var("MEMRA_SPEC_PEN_SESSION").as_deref() != Ok("0"))
1047}
1048
1049/// ROLLBACK DOOR for extended-entry publication from a RESTORED session
1050/// (lane/sampled-spec-quality, Item 3). Default ON: a converted prefix-cache hit that fed a
1051/// suffix captures its own prompt-end boundary so the NEXT turn can hit a longer prefix.
1052/// `MEMRA_SPEC_RESTORE_REPUBLISH=0` restores the pre-lane posture (a namespace learns exactly
1053/// one boundary and never advances it). Whole-entry semantics only — the boundary is the
1054/// restored session's own prompt end, so `entry_pos != fed_len` still refuses on the way in.
1055pub fn spec_restore_republish_on() -> bool {
1056 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1057 *ON.get_or_init(|| std::env::var("MEMRA_SPEC_RESTORE_REPUBLISH").as_deref() != Ok("0"))
1058}
1059
1060/// Diagnostics: name every boundary token on stderr (`MEMRA_SPEC_BOUNDARY_TRACE=1`), with
1061/// the argmax the pre-lane code would have emitted from the same row. This is how the
1062/// lane MEASURES the boundary rate and the deviation rate instead of estimating them.
1063fn spec_boundary_trace() -> bool {
1064 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1065 *ON.get_or_init(|| std::env::var("MEMRA_SPEC_BOUNDARY_TRACE").as_deref() == Ok("1"))
1066}
1067
1068/// llama-parity floor for the penalty window when the request does not ask for a bigger
1069/// one (`repeat_last_n` default). The serve API arms `penalty_last_n = usize::MAX` for any
1070/// non-identity penalty, so this floor only matters to explicit small windows and to the
1071/// CLI env path.
1072const PEN_WINDOW_FLOOR: usize = 64;
1073
1074/// CEILING on the penalty window, and it is a COST bound, not a semantic preference.
1075/// `penalize_logits_f32` (cu/spec_sample.cu) dedups on device by having thread `i` scan
1076/// `hist[0..i]`, so a pass is O(n_hist²) and it runs ~3x per verify round (the q rows, the
1077/// p column, the bonus column). The serve API arms `penalty_last_n = usize::MAX` for ANY
1078/// non-identity penalty — "the whole context", llama's `repeat_last_n = -1` — so an
1079/// uncapped session window would put a 128k-token history through that kernel: ~1.7e10
1080/// comparisons per pass, tens of ms per round, i.e. penalties would silently destroy decode
1081/// throughput on exactly the long-context requests that most want them. 8192 keeps a pass
1082/// at ~7e7 comparisons (tens of microseconds) while still being **128x wider than the
1083/// pre-lane effective window** (64 prompt-tail tokens + whatever the current burst had
1084/// generated). A request that genuinely needs a window beyond this wants host-side dedup +
1085/// counts through a new kernel signature — a follow-up lane, named here rather than hidden.
1086/// `pub` since lane/dspark-penalized-sampled-20260821: the dspark route's accept walk and
1087/// the dspark_sample_gate binary trim their uploads with the SAME cap — a second constant
1088/// is a second thing to drift.
1089pub const PEN_WINDOW_MAX: usize = 8192;
1090
1091/// Seed a penalty window over the SESSION, not the burst (lane/sampled-spec-quality,
1092/// Item 2). The window is the last `max(penalty_last_n, 64)` tokens of
1093/// `session_committed ++ burst_prompt` — for a cold turn-1 burst (`session_committed`
1094/// empty, default `penalty_last_n`) that is byte-identically the pre-lane
1095/// `prompt.iter().rev().take(64).rev()`; for a continuation burst it is the stream the
1096/// client actually asked us to penalize, where the pre-lane code had NOTHING.
1097/// `pub` since lane/dspark-penalized-sampled-20260821: the dspark route seeds its session
1098/// window through the SAME function (one definition of "the window" across both spec
1099/// routes and the gate binary's trunk-only reference arm).
1100pub fn pen_window_seed(
1101 session_committed: &[u32],
1102 burst_prompt: &[u32],
1103 penalty_last_n: usize,
1104) -> Vec<u32> {
1105 let win = penalty_last_n.clamp(PEN_WINDOW_FLOOR, PEN_WINDOW_MAX);
1106 let take_prompt = burst_prompt.len().min(win);
1107 let take_sess = (win - take_prompt).min(session_committed.len());
1108 let mut hist = Vec::with_capacity(take_sess + take_prompt);
1109 hist.extend_from_slice(&session_committed[session_committed.len() - take_sess..]);
1110 hist.extend_from_slice(&burst_prompt[burst_prompt.len() - take_prompt..]);
1111 hist
1112}
1113
1114/// Draw a BOUNDARY token from the target distribution the request asked for
1115/// (lane/sampled-spec-quality, Item 1) — the fix for "sampled spec emits an ARGMAX token at
1116/// every burst boundary".
1117///
1118/// WHY THIS EXISTS. A spec burst's first emitted token is not produced by the accept walk:
1119/// it comes off a logits row that already exists (the prime's last row on a cold burst; the
1120/// row after the last committed token on a continuation burst; the prefix-cache entry's
1121/// boundary row on a restored one). Pre-lane that token was `argmax` in BOTH sampling
1122/// regimes, so a sampled stream took a greedy token once per burst — measured, not
1123/// estimated, in research/spec-cache-20260818/SAMPLED-QUALITY.md. At temperature > 0 the
1124/// customer asked for a sampled token, so this draws one.
1125///
1126/// THE PROGRAM IS THE FULL-ACCEPT BONUS'S PROGRAM, deliberately: penalize the row (over the
1127/// session's window), take this row's OWN filter stats (the sampfix-20260805 law — stats
1128/// from a neighbour row mis-scale every `e0` and can wipe the row to token 0), gumbel-perturb
1129/// with the session's Philox stream at `*sctr`, argmax the perturbed row. Reusing the bonus's
1130/// composition means `sample_check`'s distributional oracle covers this draw too, and the
1131/// boundary token is drawn from the same filtered/penalized `p` the accept walk targets.
1132///
1133/// THE STREAM IS THE SESSION'S, NOT A FRESH ONE. `sctr` is the caller's live counter and is
1134/// advanced by exactly one, so a boundary draw consumes the next value in the same Philox
1135/// stream the accept walk uses — never a second, independently seeded stream (which would be
1136/// a new distributional bug: two streams from one seed correlate wherever their counters
1137/// collide). That also makes a restored session's boundary draw at `sctr == 0` bit-identical
1138/// to the cold session's own first draw from the same logits row, which is what preserves the
1139/// sampled-hit lane's per-seed hit==cold byte identity.
1140#[allow(clippy::too_many_arguments)]
1141pub fn sample_boundary_token_dev(
1142 e: &Engine,
1143 logits: &CudaSlice<f32>,
1144 n_vocab: usize,
1145 sp: &SpecSampling,
1146 pen_hist: &[u32],
1147 sctr: &mut u32,
1148 site: &str,
1149) -> Result<u32, Box<dyn std::error::Error>> {
1150 debug_assert!(
1151 sp.temp > 0.0,
1152 "boundary sampling is the sampled regime only"
1153 );
1154 // Own copy: penalize_logits mutates in place and the caller's row is live state
1155 // (prime_logits back the constrained recompute; last_col_logits backs round 0's accept).
1156 let mut col = e.zeros(n_vocab)?;
1157 e.copy_into(&mut col, 0, logits, n_vocab)?;
1158 let pen_on = sp.penalty_last_n > 0
1159 && (sp.penalty_repeat != 1.0 || sp.penalty_freq != 0.0 || sp.penalty_present != 0.0);
1160 if pen_on && !pen_hist.is_empty() {
1161 // window trim mirrors the round loop's own upload (`pen_hist[w0..]`), cap included.
1162 let w0 = pen_hist
1163 .len()
1164 .saturating_sub(sp.penalty_last_n.min(PEN_WINDOW_MAX));
1165 let hist = &pen_hist[w0..];
1166 let hd = e.htod_u32_v(hist)?;
1167 e.penalize_logits(
1168 &mut col,
1169 &hd,
1170 hist.len(),
1171 sp.penalty_repeat,
1172 sp.penalty_freq,
1173 sp.penalty_present,
1174 n_vocab,
1175 )?;
1176 }
1177 let rows0 = e.htod_i32(&[0])?;
1178 let (mut th_d, mut z_d, mut mx_d) = (e.zeros(1)?, e.zeros(1)?, e.zeros(1)?);
1179 e.filter_stats(
1180 &col, n_vocab, &rows0, &mut th_d, &mut z_d, &mut mx_d, n_vocab, 1, sp.temp, sp.top_k,
1181 sp.top_p, sp.min_p,
1182 )?;
1183 let (th, mx) = (e.dtoh(&th_d)?[0], e.dtoh(&mx_d)?[0]);
1184 let mut perturb = e.zeros(n_vocab)?;
1185 e.gumbel_perturb_filtered(&col, &mut perturb, n_vocab, sp.seed, *sctr, sp.temp, mx, th)?;
1186 *sctr = sctr.wrapping_add(1);
1187 let td = e.argmax_token_device(&perturb, n_vocab)?;
1188 let tok = e.dtoh_u32_one(&td)?;
1189 if spec_boundary_trace() {
1190 // the pre-lane token, from the SAME row, so the deviation rate is measurable.
1191 let raw = e.argmax_token_device(logits, n_vocab)?;
1192 let greedy = e.dtoh_u32_one(&raw)?;
1193 eprintln!(
1194 "[spec-boundary] site={site} sampled={tok} argmax={greedy} \
1195 deviates={} temp={} sctr={}",
1196 (tok != greedy) as u8,
1197 sp.temp,
1198 sctr.wrapping_sub(1),
1199 );
1200 }
1201 Ok(tok)
1202}
1203
1204/// Host-row twin of [`sample_boundary_token_dev`] (the prime / feed / entry rows arrive as
1205/// host `Vec<f32>`).
1206#[allow(clippy::too_many_arguments)]
1207pub fn sample_boundary_token(
1208 e: &Engine,
1209 logits: &[f32],
1210 sp: &SpecSampling,
1211 pen_hist: &[u32],
1212 sctr: &mut u32,
1213 site: &str,
1214) -> Result<u32, Box<dyn std::error::Error>> {
1215 let n_vocab = logits.len();
1216 let d = e.htod(logits)?;
1217 sample_boundary_token_dev(e, &d, n_vocab, sp, pen_hist, sctr, site)
1218}
1219
1220struct SpecPipeTraceClock {
1221 pair: usize,
1222 started: std::time::Instant,
1223}
1224
1225#[derive(Clone)]
1226struct SpecPipeTraceCtx {
1227 clock: std::sync::Arc<SpecPipeTraceClock>,
1228 round: usize,
1229 lane: usize,
1230}
1231
1232struct SpecPipeTraceMarker {
1233 trace: SpecPipeTraceCtx,
1234 phase: &'static str,
1235 edge: &'static str,
1236 slot: Option<usize>,
1237}
1238
1239unsafe extern "C" fn spec_pipe_trace_marker(raw: *mut std::ffi::c_void) {
1240 let marker = unsafe { Box::from_raw(raw.cast::<SpecPipeTraceMarker>()) };
1241 let lane = if marker.trace.lane == 0 { "A" } else { "B" };
1242 let slot = marker
1243 .slot
1244 .map(|v| v.to_string())
1245 .unwrap_or_else(|| "-".into());
1246 let t_ms = marker.trace.clock.started.elapsed().as_secs_f64() * 1e3;
1247 use std::io::Write as _;
1248 let stderr = std::io::stderr();
1249 let mut stderr = stderr.lock();
1250 let _ = writeln!(
1251 stderr,
1252 "[spec-pipe-timeline] pair={} round={} lane={lane} phase={} edge={} \
1253 slot={slot} t_ms={t_ms:.3}",
1254 marker.trace.clock.pair, marker.trace.round, marker.phase, marker.edge,
1255 );
1256}
1257
1258fn enqueue_spec_pipe_trace_marker(
1259 stream: &cudarc::driver::CudaStream,
1260 trace: Option<&SpecPipeTraceCtx>,
1261 phase: &'static str,
1262 edge: &'static str,
1263 slot: Option<usize>,
1264) -> Result<(), Box<dyn std::error::Error>> {
1265 let Some(trace) = trace else {
1266 return Ok(());
1267 };
1268 let marker = Box::new(SpecPipeTraceMarker {
1269 trace: trace.clone(),
1270 phase,
1271 edge,
1272 slot,
1273 });
1274 let raw = Box::into_raw(marker);
1275 let result = unsafe {
1276 cudarc::driver::result::stream::launch_host_function(
1277 stream.cu_stream(),
1278 spec_pipe_trace_marker,
1279 raw.cast(),
1280 )
1281 };
1282 if let Err(err) = result {
1283 unsafe {
1284 drop(Box::from_raw(raw));
1285 }
1286 return Err(err.into());
1287 }
1288 Ok(())
1289}
1290
1291#[derive(Default)]
1292struct SpecPipeProgress {
1293 setup_done: [bool; 2],
1294 draft_done: [usize; 2],
1295 stage0_done: [usize; 2],
1296 verify_done: [usize; 2],
1297 accept_done: [usize; 2],
1298 finished: [bool; 2],
1299 aborted: bool,
1300}
1301
1302/// Host-side issue coordinator for the reduced two-session speculative pipeline. Each session
1303/// keeps its existing call stack and round locals; this object only orders phase entry. The
1304/// primary mutex spans whole draft/accept/tail issue regions so Engine's single-stream scratch
1305/// cannot be interleaved by the two host threads.
1306struct SpecPipeSync {
1307 progress: std::sync::Mutex<SpecPipeProgress>,
1308 changed: std::sync::Condvar,
1309 primary: std::sync::Mutex<()>,
1310 trace: Option<std::sync::Arc<SpecPipeTraceClock>>,
1311}
1312
1313impl SpecPipeSync {
1314 fn new() -> Self {
1315 static TRACE_PAIR: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
1316 let trace = (std::env::var("MEMRA_SPEC_PIPE_TRACE").as_deref() == Ok("1")).then(|| {
1317 std::sync::Arc::new(SpecPipeTraceClock {
1318 pair: TRACE_PAIR.fetch_add(1, std::sync::atomic::Ordering::Relaxed) + 1,
1319 started: std::time::Instant::now(),
1320 })
1321 });
1322 Self {
1323 progress: std::sync::Mutex::new(SpecPipeProgress::default()),
1324 changed: std::sync::Condvar::new(),
1325 primary: std::sync::Mutex::new(()),
1326 trace,
1327 }
1328 }
1329}
1330
1331#[derive(Clone)]
1332struct SpecPipeLane {
1333 sync: std::sync::Arc<SpecPipeSync>,
1334 lane: usize,
1335}
1336
1337impl SpecPipeLane {
1338 fn peer(&self) -> usize {
1339 1 - self.lane
1340 }
1341
1342 fn aborted() -> Box<dyn std::error::Error> {
1343 "paired speculative peer aborted".into()
1344 }
1345
1346 fn trace(&self, round: usize) -> Option<SpecPipeTraceCtx> {
1347 self.sync.trace.as_ref().map(|clock| SpecPipeTraceCtx {
1348 clock: clock.clone(),
1349 round,
1350 lane: self.lane,
1351 })
1352 }
1353
1354 fn setup_begin(&self) -> Result<(), Box<dyn std::error::Error>> {
1355 let mut p = self.sync.progress.lock().unwrap();
1356 while !p.aborted && self.lane == 1 && !p.setup_done[0] && !p.finished[0] {
1357 p = self.sync.changed.wait(p).unwrap();
1358 }
1359 if p.aborted {
1360 Err(Self::aborted())
1361 } else {
1362 Ok(())
1363 }
1364 }
1365
1366 fn setup_end(&self) {
1367 let mut p = self.sync.progress.lock().unwrap();
1368 p.setup_done[self.lane] = true;
1369 self.sync.changed.notify_all();
1370 }
1371
1372 fn draft_begin(
1373 &self,
1374 round: usize,
1375 ) -> Result<std::sync::MutexGuard<'_, ()>, Box<dyn std::error::Error>> {
1376 let peer = self.peer();
1377 let mut p = self.sync.progress.lock().unwrap();
1378 loop {
1379 if p.aborted {
1380 return Err(Self::aborted());
1381 }
1382 let setup_ready =
1383 (p.setup_done[0] || p.finished[0]) && (p.setup_done[1] || p.finished[1]);
1384 let prior_ready = p.accept_done[self.lane] >= round
1385 && (p.accept_done[peer] >= round || p.finished[peer]);
1386 let turn_ready = if self.lane == 0 {
1387 true
1388 } else {
1389 p.draft_done[0] > round || p.finished[0]
1390 };
1391 if setup_ready && prior_ready && turn_ready {
1392 break;
1393 }
1394 p = self.sync.changed.wait(p).unwrap();
1395 }
1396 drop(p);
1397 Ok(self.sync.primary.lock().unwrap())
1398 }
1399
1400 fn draft_end(&self, round: usize) {
1401 let mut p = self.sync.progress.lock().unwrap();
1402 p.draft_done[self.lane] = round + 1;
1403 self.sync.changed.notify_all();
1404 }
1405
1406 /// Admit stage 0 and return whether this lane owns the interval's one reverse fence.
1407 /// Lane B releases as soon as lane A has issued its boundary TX, not after A's full body.
1408 fn stage0_begin(&self, round: usize) -> Result<bool, Box<dyn std::error::Error>> {
1409 let peer = self.peer();
1410 let mut p = self.sync.progress.lock().unwrap();
1411 loop {
1412 if p.aborted {
1413 return Err(Self::aborted());
1414 }
1415 let ready = if self.lane == 0 {
1416 p.draft_done[0] > round && (p.draft_done[1] > round || p.finished[1])
1417 } else {
1418 p.draft_done[1] > round && (p.stage0_done[0] > round || p.finished[0])
1419 };
1420 if ready {
1421 return Ok(self.lane == 0 || p.finished[peer]);
1422 }
1423 p = self.sync.changed.wait(p).unwrap();
1424 }
1425 }
1426
1427 fn stage0_end(&self, round: usize) {
1428 let mut p = self.sync.progress.lock().unwrap();
1429 p.stage0_done[self.lane] = round + 1;
1430 self.sync.changed.notify_all();
1431 }
1432
1433 /// Stage 1 is single-owner per engine. A proceeds immediately after its own ticket; B waits
1434 /// for A's full stage1/head issue so only A.S1 and B.S0 can overlap.
1435 fn stage1_begin(&self, round: usize) -> Result<(), Box<dyn std::error::Error>> {
1436 let mut p = self.sync.progress.lock().unwrap();
1437 while !p.aborted
1438 && !(p.stage0_done[self.lane] > round
1439 && (self.lane == 0 || p.verify_done[0] > round || p.finished[0]))
1440 {
1441 p = self.sync.changed.wait(p).unwrap();
1442 }
1443 if p.aborted {
1444 Err(Self::aborted())
1445 } else {
1446 Ok(())
1447 }
1448 }
1449
1450 fn verify_end(&self, round: usize) {
1451 let mut p = self.sync.progress.lock().unwrap();
1452 p.verify_done[self.lane] = round + 1;
1453 self.sync.changed.notify_all();
1454 }
1455
1456 fn accept_begin(
1457 &self,
1458 round: usize,
1459 ) -> Result<std::sync::MutexGuard<'_, ()>, Box<dyn std::error::Error>> {
1460 let mut p = self.sync.progress.lock().unwrap();
1461 loop {
1462 if p.aborted {
1463 return Err(Self::aborted());
1464 }
1465 let ready = if self.lane == 0 {
1466 p.verify_done[0] > round && (p.verify_done[1] > round || p.finished[1])
1467 } else {
1468 p.verify_done[1] > round && (p.accept_done[0] > round || p.finished[0])
1469 };
1470 if ready {
1471 break;
1472 }
1473 p = self.sync.changed.wait(p).unwrap();
1474 }
1475 drop(p);
1476 Ok(self.sync.primary.lock().unwrap())
1477 }
1478
1479 fn accept_end(&self, round: usize) {
1480 let mut p = self.sync.progress.lock().unwrap();
1481 p.accept_done[self.lane] = round + 1;
1482 self.sync.changed.notify_all();
1483 }
1484
1485 fn primary(&self) -> std::sync::MutexGuard<'_, ()> {
1486 self.sync.primary.lock().unwrap()
1487 }
1488
1489 fn finish(&self, failed: bool) {
1490 let mut p = self.sync.progress.lock().unwrap();
1491 p.finished[self.lane] = true;
1492 p.aborted |= failed;
1493 self.sync.changed.notify_all();
1494 }
1495}
1496
1497struct SpecPipeFinish<'a> {
1498 lane: &'a SpecPipeLane,
1499 closed: bool,
1500}
1501
1502impl<'a> SpecPipeFinish<'a> {
1503 fn new(lane: &'a SpecPipeLane) -> Self {
1504 Self {
1505 lane,
1506 closed: false,
1507 }
1508 }
1509
1510 fn close(&mut self, failed: bool) {
1511 self.lane.finish(failed);
1512 self.closed = true;
1513 }
1514}
1515
1516impl Drop for SpecPipeFinish<'_> {
1517 fn drop(&mut self) {
1518 if !self.closed {
1519 self.lane.finish(true);
1520 }
1521 }
1522}
1523
1524/// Scoped transfer of one exclusively-borrowed session to the second host issue thread.
1525/// `CudaGraph` is not marked Send by cudarc because its raw driver handles carry no automatic
1526/// trait. CUDA driver graph handles are context-scoped rather than OS-thread-affine; the caller
1527/// binds that context before touching the session, joins before returning, and never aliases the
1528/// pointer. Keep this exception local to the experimental pair call instead of marking the public
1529/// session type Send.
1530struct SpecPipeSessionPtr(*mut SpecSession);
1531
1532unsafe impl Send for SpecPipeSessionPtr {}
1533
1534impl SpecPipeSessionPtr {
1535 unsafe fn get_mut(&mut self) -> &mut SpecSession {
1536 unsafe { &mut *self.0 }
1537 }
1538}
1539
1540/// Per-session persistent draft-graph context: the captured CUDA graph(s) plus the device
1541/// buffers whose POINTERS the capture bakes. Reuse legality: the greedy capture bakes only
1542/// session-stable pointers (the session's own MtpScratch KV — allocated once, never realloc'd;
1543/// the model's resident embedding; the process-wide OnceLock p_min) and the g_* buffers held
1544/// HERE — so one capture serves the session's whole lifetime. The sampled capture additionally
1545/// bakes (seed, temp) as capture-time constants and needs k q-slots — keyed by `s_key`, dropped
1546/// and recaptured when a pool-resumed request changes them. `*_failed` memoizes a failed capture
1547/// so the eager fallback doesn't pay a doomed capture attempt every burst.
1548/// Capture identity of the parked SAMPLED draft graph (`DraftGraphCtx::graph_s`).
1549///
1550/// EXACTNESS, not perf (lane/graph-s-key-exactness-20260819; receipts
1551/// `research/spec-cache-20260818/GRAPH-S-KEY.md`). Two classes of field live here, both
1552/// load-bearing:
1553///
1554/// - **Baked constants.** `seed` and `temp` are capture-time constants INSIDE the graph and `k`
1555/// sizes the q slots its replays write. A resumed request changing any of them must recapture.
1556/// This is all the key used to carry.
1557/// - **Regime fields.** `top_k`/`top_p`/`min_p`/`pen_on` are not baked, but they decide whether
1558/// the captured graph is a legal draft chain AT ALL. The in-graph draw is one gumbel-max over
1559/// the RAW softmax (`gumbel_perturb_ctr`, unfiltered by construction), while the verify builds
1560/// the accept test's `q` from `filter_stats(q_slots, top_k, top_p, min_p)`. If those disagree
1561/// the accept test evaluates a distribution the draft was never sampled from: a draft token
1562/// below the filter threshold gathers `q = 0` (`softmax_gather_filtered_f32`,
1563/// `cu/spec_sample.cu`) and `u * 0 < p` accepts it UNCONDITIONALLY.
1564///
1565/// Omitting the regime fields was reachable — not through the prefix-cache spec restore (that
1566/// path is greedy-only, `memra-server` `spec_restore_convertible`), but through WHOLE-SESSION
1567/// spec reuse: a parked `SpecSession` carries this `DraftGraphCtx`, and the pool-resume probe
1568/// applies no sampler predicate at all. Turn 1 pure-temp parks a graph; turn 2 of the same
1569/// conversation, same explicit seed and temperature, adds `top_p`/`top_k` and inherits it.
1570#[derive(Clone, Copy, PartialEq, Eq, Debug)]
1571pub(crate) struct SampledGraphKey {
1572 seed: u64,
1573 temp_bits: u32,
1574 k: usize,
1575 top_k: i32,
1576 top_p_bits: u32,
1577 min_p_bits: u32,
1578 pen_on: bool,
1579}
1580
1581impl SampledGraphKey {
1582 pub(crate) fn new(
1583 seed: u64,
1584 temp: f32,
1585 k: usize,
1586 top_k: i32,
1587 top_p: f32,
1588 min_p: f32,
1589 pen_on: bool,
1590 ) -> Self {
1591 SampledGraphKey {
1592 seed,
1593 temp_bits: temp.to_bits(),
1594 k,
1595 top_k,
1596 top_p_bits: top_p.to_bits(),
1597 min_p_bits: min_p.to_bits(),
1598 pen_on,
1599 }
1600 }
1601
1602 /// The one regime the in-graph sampled chain may stand in for the eager one: nothing but
1603 /// temperature shapes `q`. Computed FROM THE KEY so the capture guard, the launch guard and
1604 /// the key can never drift apart (they were three separate expressions before this lane, and
1605 /// the launch site simply forgot to ask).
1606 pub(crate) fn pure_temp(&self) -> bool {
1607 self.top_k == 0
1608 && f32::from_bits(self.top_p_bits) >= 1.0
1609 && f32::from_bits(self.min_p_bits) <= 0.0
1610 && !self.pen_on
1611 }
1612}
1613
1614pub(crate) struct DraftGraphCtx {
1615 g_tok: CudaSlice<u32>,
1616 g_pos: CudaSlice<i32>,
1617 g_seed: CudaSlice<f32>,
1618 g_p: CudaSlice<f32>,
1619 g_ctr: CudaSlice<u32>,
1620 g_q: CudaSlice<f32>,
1621 g_perturb: CudaSlice<f32>,
1622 q_slots: Vec<CudaSlice<f32>>,
1623 /// DRAFT-SIDE GRAMMAR MASK (lane/draft-mask): packed allowed-set words over the DRAFT
1624 /// head's vocab, at a STABLE address so the captured draft graph's mask node reads the
1625 /// per-position contents the host re-uploads before each replay (the graph-promote
1626 /// pattern from decode.rs). Empty unless the session drafts under a grammar.
1627 g_dmask: CudaSlice<u32>,
1628 /// was `graph` captured WITH the mask node? A parked graph of the wrong shape is dropped.
1629 graph_masked: bool,
1630 graph: Option<cudarc::driver::CudaGraph>,
1631 graph_s: Option<cudarc::driver::CudaGraph>,
1632 /// Failed-capture memoization for both graphs — LOUD on flip, cleared on pool resume
1633 /// (audit Q2, the TRT #16072 silent-permanent-coverage-loss class).
1634 failed: DraftGraphFallback,
1635 /// Capture identity of `graph_s` — see [`SampledGraphKey`]. `None` iff no sampled graph is
1636 /// parked; a request whose key differs drops the parked graph (and its q slots/keeper).
1637 s_key: Option<SampledGraphKey>,
1638 /// CAPTURE-RETAIN keepers (#68 root cause, 2026-08-04): the warmup-run transients whose
1639 /// pool addresses the captured graph(s) bake. Without these, the transients return to the
1640 /// pool at capture-body exit and later work (burst-boundary prime/fill/commit passes, or a
1641 /// co-served session in the worker) reuses those addresses — the persisted graph's replay
1642 /// then reads/writes live unrelated buffers (exactness corruption, first seen as the ST
1643 /// serve-spec 4B graph-arm corruption; one-shot CLI calls never re-shuffled the pool, which
1644 /// is why run-spec K=1..8 passed on the same checkpoint). Same fix class as
1645 /// capture_graph_retained's gemma/decode.rs sites — hold as long as the graph replays.
1646 keeper: Vec<Box<dyn std::any::Any + Send>>,
1647 keeper_s: Vec<Box<dyn std::any::Any + Send>>,
1648}
1649
1650/// Failed-capture memoization for the two draft graphs (audit Q2, 2026-08-05 — the
1651/// TRT #16072 trap class: pressure-triggered, silent, long-lived coverage loss).
1652///
1653/// Three contracts:
1654/// - LOUD FLIP: `mark_*` returns the warn line exactly on the false→true transition
1655/// (returned, not printed, so the once-per-flip contract is unit-testable); the caller
1656/// `eprintln!`s it UNCONDITIONALLY — a dropped draft graph is never silent. Re-marking
1657/// an already-failed graph returns None (the per-burst memoization that keeps the eager
1658/// fallback from paying a doomed capture attempt every burst).
1659/// - RESET ON RESUME: `reset_on_resume` clears both flags — a parked session resumed by a
1660/// NEW request gets one fresh capture chance instead of carrying a transient-pressure
1661/// failure for the pool's whole lifetime. Returns the note line only when a flag was
1662/// actually set (quiet on the common clean-resume path).
1663/// - Shape-change clears (`clear_*`) stay silent, exactly as before: they precede a fresh
1664/// capture attempt whose own failure would re-flip loudly.
1665#[derive(Default)]
1666pub(crate) struct DraftGraphFallback {
1667 greedy: bool,
1668 sampled: bool,
1669}
1670impl DraftGraphFallback {
1671 fn mark_greedy(&mut self, reason: &str) -> Option<String> {
1672 if self.greedy {
1673 return None;
1674 }
1675 self.greedy = true;
1676 Some(format!(
1677 "[spec] WARN: draft-graph capture failed ({reason}); eager fallback until session resume"
1678 ))
1679 }
1680 fn mark_sampled(&mut self, reason: &str) -> Option<String> {
1681 if self.sampled {
1682 return None;
1683 }
1684 self.sampled = true;
1685 Some(format!(
1686 "[spec] WARN: sampled draft-graph capture failed ({reason}); eager fallback until session resume"
1687 ))
1688 }
1689 fn greedy_failed(&self) -> bool {
1690 self.greedy
1691 }
1692 fn sampled_failed(&self) -> bool {
1693 self.sampled
1694 }
1695 fn clear_greedy(&mut self) {
1696 self.greedy = false;
1697 }
1698 fn clear_sampled(&mut self) {
1699 self.sampled = false;
1700 }
1701 /// Pool-resume reset: both graphs get a fresh capture chance. Some(note) iff any flag
1702 /// was set (so clean resumes stay quiet).
1703 pub(crate) fn reset_on_resume(&mut self) -> Option<String> {
1704 if !self.greedy && !self.sampled {
1705 return None;
1706 }
1707 let which = match (self.greedy, self.sampled) {
1708 (true, true) => "greedy+sampled",
1709 (true, false) => "greedy",
1710 _ => "sampled",
1711 };
1712 self.greedy = false;
1713 self.sampled = false;
1714 Some(format!(
1715 "[spec] draft-graph fallback reset on session resume ({which}); recapture eligible"
1716 ))
1717 }
1718}
1719
1720impl DraftGraphCtx {
1721 fn new(e: &Engine, n_embd: usize, qlen: usize) -> Result<Self, Box<dyn std::error::Error>> {
1722 Ok(DraftGraphCtx {
1723 g_tok: e.alloc_u32_zeroed(1)?,
1724 g_pos: e.htod_i32(&[0])?,
1725 g_seed: e.zeros(n_embd)?,
1726 g_p: e.zeros(1)?,
1727 g_ctr: e.alloc_u32_zeroed(1)?,
1728 g_q: e.zeros(qlen)?,
1729 g_perturb: e.zeros(qlen)?,
1730 q_slots: Vec::new(),
1731 g_dmask: e.alloc_u32_zeroed(1)?,
1732 graph_masked: false,
1733 graph: None,
1734 graph_s: None,
1735 failed: DraftGraphFallback::default(),
1736 s_key: None,
1737 keeper: Vec::new(),
1738 keeper_s: Vec::new(),
1739 })
1740 }
1741}
1742
1743pub(crate) struct MtpScratch {
1744 kv: KvLayer,
1745 /// Logical row capacity. On the graph/DC draft path it also doubles as fa_decode_dc's
1746 /// bucket_max: n_splits is sized from it ONCE, so the graph captured at round 0 stays valid
1747 /// for every later t_kv. Step35 refuses that path and may back this logical extent with the
1748 /// smaller host-indexed SWA ring instead.
1749 cap: usize,
1750 extra: Vec<MtpScratchPlane>,
1751}
1752
1753struct MtpScratchPlane {
1754 kv: KvLayer,
1755 cap: usize,
1756}
1757
1758fn mtp_scratch_layout(
1759 cfg: &memra_gguf::config::ModelConfig,
1760 geom: Option<&crate::hybrid::DraftGeom>,
1761) -> (usize, usize, usize, usize) {
1762 // Student draft heads carry fewer KV heads (head_dim unchanged) -> smaller scratch rows.
1763 let n_head_kv = geom.map(|g| g.n_head_kv).unwrap_or(cfg.n_head_kv as usize);
1764 let head_dim_k = cfg.head_dim_k as usize;
1765 let head_dim_v = cfg.head_dim_v as usize;
1766 assert!(
1767 head_dim_k % 32 == 0 && head_dim_v % 32 == 0,
1768 "KVQUANT requires head_dim%32==0 (MTP scratch)"
1769 );
1770 let kv_dim_k = head_dim_k * n_head_kv;
1771 let kv_dim_v = head_dim_v * n_head_kv;
1772 // The fp8-KV arm deliberately does not reach the draft scratch; keep the exact format
1773 // policy shared with `MtpScratch::new` so admission scales the same allocation.
1774 let (kbb, vbb) = crate::kv_blk_bytes();
1775 let k_tok_bytes = (kv_dim_k / 32) * kbb;
1776 let v_tok_bytes = (kv_dim_v / 32) * vbb;
1777 (kv_dim_k, kv_dim_v, k_tok_bytes, v_tok_bytes)
1778}
1779
1780fn mtp_chain_head_index(step: usize, head_count: usize) -> usize {
1781 assert!(head_count > 0, "MTP chain requires at least one head");
1782 step % head_count
1783}
1784
1785impl MtpScratch {
1786 fn alloc_plane(
1787 e: &Engine,
1788 cfg: &memra_gguf::config::ModelConfig,
1789 plan: &memra_gguf::model_plan::ModelPlan,
1790 cap: usize,
1791 geom: Option<&crate::hybrid::DraftGeom>,
1792 ) -> Result<MtpScratchPlane, Box<dyn std::error::Error>> {
1793 let (kv_dim_k, kv_dim_v, k_tok_bytes, v_tok_bytes) = mtp_scratch_layout(cfg, geom);
1794 let ring = if crate::cache::swa_ring_on()
1795 && crate::plan_backend::decode_batch_program(plan)
1796 == crate::plan_backend::DecodeBatchProgram::SlidingGatedMoe
1797 {
1798 let window = plan
1799 .layers
1800 .iter()
1801 .find_map(|layer| match layer.attention {
1802 memra_gguf::model_plan::AttentionPlan::SlidingWindow { window, .. } => {
1803 Some(window as usize)
1804 }
1805 _ => None,
1806 })
1807 .ok_or("sliding-gated-MoE draft scratch has no sliding-window layer")?;
1808 Some(crate::cache::KvRing::new(
1809 crate::cache::swa_ring_rows(window, cap),
1810 window,
1811 ))
1812 } else {
1813 None
1814 };
1815 let alloc_rows = ring.as_ref().map(crate::cache::KvRing::rows).unwrap_or(cap);
1816 Ok(MtpScratchPlane {
1817 kv: KvLayer {
1818 k: e.alloc_u8(alloc_rows * k_tok_bytes)?,
1819 v: e.alloc_u8(alloc_rows * v_tok_bytes)?,
1820 kv_dim_k,
1821 kv_dim_v,
1822 k_tok_bytes,
1823 v_tok_bytes,
1824 len: 0,
1825 ring,
1826 len_d: e.htod_i32(&[0])?,
1827 },
1828 cap,
1829 })
1830 }
1831
1832 fn new(
1833 e: &Engine,
1834 cfg: &memra_gguf::config::ModelConfig,
1835 plan: &memra_gguf::model_plan::ModelPlan,
1836 cap: usize,
1837 geom: Option<&crate::hybrid::DraftGeom>,
1838 ) -> Result<Self, Box<dyn std::error::Error>> {
1839 // env-selected KV formats (default 34/24). The fp8-KV arm (MEMRA_KV_FP8) deliberately
1840 // does NOT reach the draft scratch: fp8 drafts drifted acceptance 69-88% -> 46%
1841 // (2026-07-12 A/B); the scratch is tiny, so it keeps baseline q8_0/q5_1 numerics
1842 // while the TRUNK cache carries the fp8 depth win. Scratch append/fa pass g=false.
1843 let primary = Self::alloc_plane(e, cfg, plan, cap, geom)?;
1844 Ok(MtpScratch {
1845 kv: primary.kv,
1846 cap: primary.cap,
1847 extra: Vec::new(),
1848 })
1849 }
1850
1851 fn push_plane(
1852 &mut self,
1853 e: &Engine,
1854 cfg: &memra_gguf::config::ModelConfig,
1855 plan: &memra_gguf::model_plan::ModelPlan,
1856 geom: Option<&crate::hybrid::DraftGeom>,
1857 ) -> Result<(), Box<dyn std::error::Error>> {
1858 self.extra
1859 .push(Self::alloc_plane(e, cfg, plan, self.cap, geom)?);
1860 Ok(())
1861 }
1862
1863 fn plane_count(&self) -> usize {
1864 1 + self.extra.len()
1865 }
1866
1867 fn plane(&self, index: usize) -> (&KvLayer, usize) {
1868 if index == 0 {
1869 (&self.kv, self.cap)
1870 } else {
1871 let plane = &self.extra[index - 1];
1872 (&plane.kv, plane.cap)
1873 }
1874 }
1875
1876 fn plane_mut(&mut self, index: usize) -> (&mut KvLayer, usize) {
1877 if index == 0 {
1878 (&mut self.kv, self.cap)
1879 } else {
1880 let plane = &mut self.extra[index - 1];
1881 (&mut plane.kv, plane.cap)
1882 }
1883 }
1884
1885 fn set_plane_len(
1886 &mut self,
1887 e: &Engine,
1888 index: usize,
1889 n: usize,
1890 ) -> Result<(), Box<dyn std::error::Error>> {
1891 let (kv, _) = self.plane_mut(index);
1892 if kv.ring.as_ref().is_some_and(|ring| !ring.can_rewind_to(n)) {
1893 return Err("SWA ring MTP checkpoint has been lapped; full re-prime required".into());
1894 }
1895 kv.len = n;
1896 e.set_i32_one(&mut kv.len_d, n as i32)
1897 }
1898
1899 /// Set BOTH length counters: the host mirror AND the device len_d the captured append/fa read
1900 /// (a 4-byte in-place htod — the counter pointer is baked into the graph, never realloc'd).
1901 /// This is the ONLY truncation/rollback mechanism the persistent draft KV needs.
1902 fn set_len(&mut self, e: &Engine, n: usize) -> Result<(), Box<dyn std::error::Error>> {
1903 if !self.can_rewind_to(n) {
1904 return Err("SWA ring MTP checkpoint has been lapped; full re-prime required".into());
1905 }
1906 for index in 0..self.plane_count() {
1907 self.set_plane_len(e, index, n)?;
1908 }
1909 Ok(())
1910 }
1911
1912 fn can_rewind_to(&self, n: usize) -> bool {
1913 (0..self.plane_count()).all(|index| {
1914 self.plane(index)
1915 .0
1916 .ring
1917 .as_ref()
1918 .is_none_or(|ring| ring.can_rewind_to(n))
1919 })
1920 }
1921}
1922
1923/// Retained verify intermediates for the REPLAY-FREE partial accept (2026-07-03, the profiled
1924/// #1 spec cost at long ctx: the partial-accept replay was a DUPLICATE trunk pass — ~0.54 extra
1925/// full weight reads per round — recomputing columns the verify had already produced
1926/// bit-identically). Holds, per linear layer, everything needed to rebuild its recurrent state
1927/// to "after the first j verify columns" WITHOUT re-running the trunk:
1928/// - BATCHED-path layers (`gdn`): the exact token-major inputs the round's ONE gdn_scan
1929/// consumed. A prefix re-run of the SAME kernel (t=j) from the snapshot state is bit-identical
1930/// to the first j iterations of the verify's scan — the kernel's t-loop carries state in
1931/// registers and iteration t never depends on T. `qkv_mixed` (the conv input) feeds the
1932/// pure-copy ring rebuild.
1933/// - PER-COLUMN-path layers (`cols`): dtod clones of (conv_state, ssm_state) taken after each
1934/// column 0..t-2 — pure copies of the actual chain states (the last column is never a rebuild
1935/// target: j <= t-1).
1936/// Full-attn layers need nothing: their verify KV rows are bit-identical to eager's (the
1937/// decode-exact contract; verify-probe pins it), so rollback = len truncation.
1938struct GdnStash {
1939 qkv_mixed: CudaSlice<f32>, // [t, conv_dim] token-major (conv input)
1940 q_l2: CudaSlice<f32>,
1941 k_l2: CudaSlice<f32>,
1942 v_g: CudaSlice<f32>, // [t, num_v, d_state]
1943 g_log: CudaSlice<f32>,
1944 beta: CudaSlice<f32>, // [t, num_v]
1945}
1946pub(crate) struct VerifyCkpt {
1947 gdn: Vec<Option<GdnStash>>, // [n_layer], Some iff batched linear path ran
1948 cols: Vec<Option<Vec<(CudaSlice<f32>, CudaSlice<f32>)>>>, // [n_layer][col] = (conv, ssm) after col
1949}
1950/// Opaque handle for the dspark round (dflash.rs) — VerifyCkpt stays spec-private.
1951pub(crate) struct DsparkVerifyCkpt(VerifyCkpt);
1952
1953/// Engine-bundle slice 3 (DSF-ROUNDCOST-20260820 §2 row 4 / §5 rank 1): bucketed CUDA
1954/// graphs for the dspark verify's LINEAR-layer segments. The measured verify is ~2,800
1955/// eager launches whose residual cost is DEVICE-side per-launch overhead (slice 2 proved
1956/// host dispatch is not the binder: fully-deferred dispatch bought ~0 wall). The 48 GDN
1957/// layers between full-attention layers are shape-static given vt — no positions, no
1958/// t_kv, state addressed through pointer tables — so runs of them capture per
1959/// (segment, vt) and replay as ONE graph launch each. Full-attention layers stay eager
1960/// (their per-row append/fa arm picks are t_kv-driven — the exec-update extension).
1961///
1962/// Per round out-of-graph: one pointer-table refresh (gdn ping-pong moves the canonical
1963/// handles), one input-staging copy per segment, host parity bookkeeping. Captured via
1964/// `capture_graph_retained` (2 warmups + capture, keeper retains warmup transients so
1965/// pool addresses stay stable); the warmups EXECUTE, so segment conv/ssm state is saved
1966/// before and restored after — the graph's first real launch starts from the exact
1967/// pre-round state. The ckpt column stash rides persistent slabs (written inside the
1968/// graph as memcpy nodes); commit reads them via `dspark_commit_prefix_slab`.
1969/// `MEMRA_DSPARK_VERIFY_GRAPH=0` reverts to the eager walk (byte-identical body).
1970pub(crate) struct DsparkVerifyGraphs {
1971 /// Linear-attention layer indices ascending; `lin_pos[il]` = index into the vecs.
1972 lin: Vec<usize>,
1973 lin_pos: std::collections::HashMap<usize, usize>,
1974 /// [n_lin x 6] pointer table (conv, s0, s1, conv, s1, s0 per layer), refreshed per
1975 /// verify from the live handles; layer il's slice starts at lin_pos[il]*6.
1976 table_all: CudaSlice<u64>,
1977 host_table: Vec<u64>,
1978 /// Persistent per-layer ckpt stash slabs: row r of the verify at slab offset
1979 /// r*words. Shared by every (segment, vt) bucket — one verify runs at a time.
1980 stash_conv: Vec<CudaSlice<f32>>,
1981 stash_ssm: Vec<CudaSlice<f32>>,
1982 conv_words: usize,
1983 ssm_words: usize,
1984 /// Per-vt input/output staging (stable addresses the graphs bake).
1985 stage: std::collections::HashMap<usize, (CudaSlice<f32>, CudaSlice<f32>)>,
1986 /// Per-vt dflash tap-sink buffers — the captured segments bake the tap dst address,
1987 /// so the sink buffer must live (and persist) with the graphs, not with the round.
1988 pub(crate) tap_bufs: std::collections::HashMap<usize, CudaSlice<f32>>,
1989 graphs: std::collections::HashMap<(usize, usize), DsparkSegGraph>,
1990 /// Warmup-corruption guard scratch: pre-capture conv/ssm of every linear layer
1991 /// (sized n_lin — the slice-4c full-verify warmups execute the whole walk).
1992 save_conv: CudaSlice<f32>,
1993 save_ssm: CudaSlice<f32>,
1994 max_run: usize,
1995 n_embd: usize,
1996 /// Set by the verify walk: this round's linear ckpt lives in the slabs (the caller
1997 /// commits through `dspark_commit_prefix_slab` instead of the cols arm).
1998 pub(crate) round_slab: bool,
1999 // ---- slice 4c: full-verify single graph per (vt, rung) ----
2000 /// Full-attention layer indices ascending; `fa_pos[il]` = index into the vec.
2001 fa: Vec<usize>,
2002 fa_pos: std::collections::HashMap<usize, usize>,
2003 /// [n_fa x 2 x t_cap] interleaved (k,v) base-pointer pairs, refreshed per verify;
2004 /// layer il's slice starts at `fa_pos[il] * 2 * t_cap` (the seqs twins read pairs
2005 /// [2z], z < t <= t_cap, so one t_cap-sized table serves every vt).
2006 fa_table: CudaSlice<u64>,
2007 fa_host_table: Vec<u64>,
2008 t_cap: usize,
2009 /// Per-vt position staging for the captured bodies — contents refreshed per round
2010 /// (rope reads row r; the seqs twins derive append slot and T_kv per z from it).
2011 pos_stage: std::collections::HashMap<usize, CudaSlice<i32>>,
2012 /// Full-verify graphs keyed (vt, rung_end, hi).
2013 full: std::collections::HashMap<(usize, usize, usize), DsparkSegGraph>,
2014 /// Largest n with every layer in [0, n) linear or full-attention (walk coverage).
2015 covered: usize,
2016 /// Every layer in [0, n) is linear or full-attention (no MLA/unknown mixers) — the
2017 /// full-verify capture walks all of them.
2018 walk_uniform: bool,
2019 /// Last `(captures, device graph-mem reserved bytes)` reading taken by
2020 /// `HybridModel::dspark_vg_admission_debt` — the two-point base of the MARGINAL debt
2021 /// projection (see `dspark_vg_debt_projection`; a mean-based reading extrapolated the
2022 /// pool's one-time shared allocation and reserved 8.5 GB of phantom VRAM).
2023 debt_obs: Option<(usize, usize)>,
2024}
2025
2026struct DsparkSegGraph {
2027 graph: cudarc::driver::CudaGraph,
2028 _keeper: Vec<Box<dyn std::any::Any + Send>>,
2029}
2030
2031/// Per-call arguments of [`HybridModel::qwen35_tparallel_fa_layer`] — one struct so the
2032/// eager walk and the slice-4c captured full-verify graphs hand the SAME body its two
2033/// modes without a second copy of the math.
2034pub(crate) struct FaLayerArgs<'a> {
2035 /// [T] per-row positions (device): rope reads them row-indexed; the seqs twins read
2036 /// them per-z (append slot = pos, T_kv = pos + 1).
2037 pub pos_d: &'a CudaSlice<i32>,
2038 /// Verify-level lazy per-row 1-element position buffers — only the per-row fallback
2039 /// arm builds/uses them (graph mode refuses that arm).
2040 pub pos_rows: &'a mut Option<Vec<CudaSlice<i32>>>,
2041 pub pos0: usize,
2042 pub seqs_append: bool,
2043 pub batch_fa_on: bool,
2044 /// Some((kv pointer table, offset-in-u64s, rung_end)) = captured-graph mode.
2045 pub graph_cap: Option<(&'a CudaSlice<u64>, usize, usize)>,
2046 /// ROUND-STREAM (lane/draftcost-moe, v0.100 train merge): Some((token stream, device
2047 /// round counter)) routes the FA attend through the dc rows kernels and the Linear
2048 /// mixer through `linear_attn_verify_t` (the stream arms the old inline body carried).
2049 /// Never armed together with `graph_cap` (the verify-level merge guard refuses).
2050 pub stream: Option<(&'a CudaSlice<u32>, &'a CudaSlice<i32>)>,
2051 /// VerifyCkpt for the stream-Linear arm's GdnStash install; None in graph mode and
2052 /// for FA layers that never touch it.
2053 pub ckpt: Option<&'a mut VerifyCkpt>,
2054}
2055
2056// SAFETY: `CudaGraph` is not marked Send by cudarc because its raw driver handles carry
2057// no automatic trait; CUDA driver graph handles are context-scoped rather than
2058// OS-thread-affine (the SpecPipeSessionPtr precedent above). The ctx lives in
2059// `HybridModel::dspark_vgraphs` behind a Mutex and every touch happens on the engine's
2060// single decode-stream thread.
2061unsafe impl Send for DsparkVerifyGraphs {}
2062
2063impl DsparkVerifyGraphs {
2064 /// Live capture count (segment + full graphs) — the denominator of
2065 /// [`dspark_vg_debt_projection`]'s observed bytes/capture mean.
2066 pub(crate) fn captures(&self) -> usize {
2067 self.graphs.len() + self.full.len()
2068 }
2069
2070 /// Take the marginal-growth debt reading and record this observation for the next one.
2071 /// Called under the pool mutex by `HybridModel::dspark_vg_admission_debt`.
2072 pub(crate) fn admission_debt(&mut self, reserved_bytes: usize) -> usize {
2073 let captures = self.captures();
2074 let debt =
2075 dspark_vg_debt_projection(captures, dspark_vg_cap(), reserved_bytes, self.debt_obs);
2076 if captures > 0 {
2077 match self.debt_obs {
2078 Some((c0, _)) if captures <= c0 => {}
2079 _ => self.debt_obs = Some((captures, reserved_bytes)),
2080 }
2081 }
2082 debt
2083 }
2084
2085 /// Build for this cache's shape. None when there are no linear layers, sizes are
2086 /// non-uniform, or the trunk keeps a gemma4 config (never on the qwen35 family).
2087 pub(crate) fn new(
2088 e: &Engine,
2089 cache: &Cache,
2090 t_max: usize,
2091 n_embd: usize,
2092 ) -> Result<Option<Self>, Box<dyn std::error::Error>> {
2093 let lin: Vec<usize> = (0..cache.recur.len())
2094 .filter(|&il| cache.recur[il].is_some())
2095 .collect();
2096 if lin.is_empty() || t_max < 2 {
2097 return Ok(None);
2098 }
2099 let first = cache.recur[lin[0]].as_ref().unwrap();
2100 let (conv_words, ssm_words) = (first.conv_state.len(), first.ssm_state.len());
2101 for &il in &lin {
2102 let rl = cache.recur[il].as_ref().unwrap();
2103 if rl.conv_state.len() != conv_words || rl.ssm_state.len() != ssm_words {
2104 return Ok(None);
2105 }
2106 }
2107 let n = lin.len();
2108 let mut lin_pos = std::collections::HashMap::with_capacity(n);
2109 for (k, &il) in lin.iter().enumerate() {
2110 lin_pos.insert(il, k);
2111 }
2112 // longest run of consecutive linear layers (save-scratch sizing)
2113 let mut max_run = 1usize;
2114 let mut run = 1usize;
2115 for w in lin.windows(2) {
2116 if w[1] == w[0] + 1 {
2117 run += 1;
2118 max_run = max_run.max(run);
2119 } else {
2120 run = 1;
2121 }
2122 }
2123 let rows = t_max - 1;
2124 let mut stash_conv = Vec::with_capacity(n);
2125 let mut stash_ssm = Vec::with_capacity(n);
2126 for _ in 0..n {
2127 stash_conv.push(e.uninit(rows * conv_words)?);
2128 stash_ssm.push(e.uninit(rows * ssm_words)?);
2129 }
2130 let host_table = vec![0u64; n * 6];
2131 let table_all = e.htod_u64(&host_table)?;
2132 // slice 4c: full-attention census for the full-verify graphs.
2133 let fa: Vec<usize> = (0..cache.kv.len())
2134 .filter(|&il| cache.kv[il].is_some())
2135 .collect();
2136 let mut fa_pos = std::collections::HashMap::with_capacity(fa.len());
2137 for (k, &il) in fa.iter().enumerate() {
2138 fa_pos.insert(il, k);
2139 }
2140 let n_layers = cache.kv.len().max(cache.recur.len());
2141 // exactly one of (linear state, kv cache) per layer — no MLA/unknown mixers.
2142 let walk_uniform = (0..n_layers).all(|il| {
2143 cache.recur.get(il).is_some_and(|r| r.is_some())
2144 != cache.kv.get(il).is_some_and(|k| k.is_some())
2145 });
2146 // Contiguous covered prefix: the largest n such that every layer in [0, n) is
2147 // linear or full-attention. The TRUNK walk is [0, layers.len()) and the cache
2148 // vecs can carry EXTRA state slots past it (the q38 export keeps the MTP head
2149 // layer's kv at the tail — hi == lin+fa never held, the s4c battery's zero
2150 // 'full' captures). The full-graph guard is walk coverage, not slot arithmetic.
2151 let covered = (0..n_layers)
2152 .take_while(|il| lin_pos.contains_key(il) || fa_pos.contains_key(il))
2153 .count();
2154 let t_cap = t_max;
2155 let fa_host_table = vec![0u64; fa.len() * 2 * t_cap];
2156 let fa_table = e.htod_u64(&fa_host_table)?;
2157 Ok(Some(Self {
2158 lin,
2159 lin_pos,
2160 table_all,
2161 host_table,
2162 stash_conv,
2163 stash_ssm,
2164 conv_words,
2165 ssm_words,
2166 stage: std::collections::HashMap::new(),
2167 tap_bufs: std::collections::HashMap::new(),
2168 graphs: std::collections::HashMap::new(),
2169 save_conv: e.uninit(n * conv_words)?,
2170 save_ssm: e.uninit(n * ssm_words)?,
2171 max_run,
2172 n_embd,
2173 round_slab: false,
2174 fa,
2175 fa_pos,
2176 fa_table,
2177 fa_host_table,
2178 t_cap,
2179 pos_stage: std::collections::HashMap::new(),
2180 full: std::collections::HashMap::new(),
2181 covered,
2182 walk_uniform,
2183 debt_obs: None,
2184 }))
2185 }
2186
2187 /// Rebuild the pointer tables from the live handles (once per verify — the gdn
2188 /// ping-pong swaps the canonical/alt handles between rounds; a fresh generation's
2189 /// cache buffers land at new addresses; a stale table would read the wrong state).
2190 pub(crate) fn refresh_tables(
2191 &mut self,
2192 e: &Engine,
2193 cache: &Cache,
2194 ) -> Result<(), Box<dyn std::error::Error>> {
2195 use cudarc::driver::DevicePtr;
2196 {
2197 let s = &e.gpu.stream();
2198 for (k, &il) in self.lin.iter().enumerate() {
2199 let rl = cache.recur[il].as_ref().unwrap();
2200 let (pc, _g0) = rl.conv_state.device_ptr(s);
2201 let (p0, _g1) = rl.ssm_state.device_ptr(s);
2202 let (p1, _g2) = rl.ssm_state_alt.device_ptr(s);
2203 let o = k * 6;
2204 self.host_table[o] = pc as u64;
2205 self.host_table[o + 1] = p0 as u64;
2206 self.host_table[o + 2] = p1 as u64;
2207 self.host_table[o + 3] = pc as u64;
2208 self.host_table[o + 4] = p1 as u64;
2209 self.host_table[o + 5] = p0 as u64;
2210 }
2211 for (k, &il) in self.fa.iter().enumerate() {
2212 let kvl = cache.kv[il].as_ref().unwrap();
2213 let (pk, _g0) = kvl.k.device_ptr(s);
2214 let (pv, _g1) = kvl.v.device_ptr(s);
2215 let o = k * 2 * self.t_cap;
2216 for z in 0..self.t_cap {
2217 self.fa_host_table[o + 2 * z] = pk as u64;
2218 self.fa_host_table[o + 2 * z + 1] = pv as u64;
2219 }
2220 }
2221 }
2222 e.htod_u64_into(&self.host_table, &mut self.table_all)?;
2223 if !self.fa_host_table.is_empty() {
2224 e.htod_u64_into(&self.fa_host_table, &mut self.fa_table)?;
2225 }
2226 Ok(())
2227 }
2228
2229 /// Slice 4c eligibility: Some(rung_end) when this round can replay (or capture) a
2230 /// full-verify graph — the whole walk [lo, hi) is covered, every layer is linear or
2231 /// full-attention, and ALL of the round's per-row t_kv values take the v4-seqs arm
2232 /// on ONE `fa_split_keys` ladder step that the rung also sits on (the straddle law;
2233 /// both gates are t_kv intervals, so ends-inside means all-inside). The rung is the
2234 /// round's next power of two — grid/partial sizing only (`n_splits_max` is pure
2235 /// stride; splits >= ns_eff write the empty partial the combine never reads), so one
2236 /// captured graph is bit-identical for every round the rung covers.
2237 #[allow(clippy::too_many_arguments)]
2238 pub(crate) fn full_rung(
2239 &self,
2240 model: &crate::hybrid::HybridModel,
2241 cache: &Cache,
2242 lo: usize,
2243 hi: usize,
2244 t: usize,
2245 seqs_arms_on: bool,
2246 ) -> Option<usize> {
2247 if std::env::var("MEMRA_DSPARK_FULLG_DEBUG").as_deref() == Ok("1") {
2248 static ONCE: std::sync::Once = std::sync::Once::new();
2249 let len0 = self
2250 .fa
2251 .first()
2252 .and_then(|&il| cache.kv[il].as_ref())
2253 .map(|k| k.len);
2254 ONCE.call_once(|| {
2255 eprintln!(
2256 "[fullg-debug] walk_uniform={} covered={} seqs_arms_on={} fa_rows_on={} t={} lo={} hi={} lin={} fa={} t_cap={} len0={:?}",
2257 self.walk_uniform, self.covered, seqs_arms_on, dspark_fa_rows_on(), t, lo, hi,
2258 self.lin.len(), self.fa.len(), self.t_cap, len0
2259 );
2260 });
2261 }
2262 if !self.walk_uniform
2263 || !seqs_arms_on
2264 || !dspark_fa_rows_on()
2265 || t < 2
2266 || lo != 0
2267 || hi > self.covered
2268 || t > self.t_cap
2269 || self.fa.is_empty()
2270 {
2271 return None;
2272 }
2273 let cfg = &model.cfg;
2274 let head_dim_global = cfg.head_dim_k as usize;
2275 let nkv = cfg.n_head_kv as usize;
2276 let kvl0 = cache.kv[self.fa[0]].as_ref().unwrap();
2277 // the z-batched twins read stacked rows at the cache's kv dims — must equal the
2278 // projection stride (the body's guard, hoisted so ineligible models fall back
2279 // instead of refusing mid-capture).
2280 let geom = cfg.full_attention_geometry_at(self.fa[0] as u32);
2281 let kv_dim = geom.n_head_kv as usize * geom.head_dim_k as usize;
2282 if kvl0.kv_dim_k != kv_dim || kvl0.kv_dim_v != kv_dim {
2283 return None;
2284 }
2285 let len0 = kvl0.len;
2286 let (t_kv_first, t_kv_last) = (len0 + 1, len0 + t);
2287 if !crate::fa_seqs_eligible(t_kv_first, head_dim_global)
2288 || !crate::fa_seqs_eligible(t_kv_last, head_dim_global)
2289 || crate::fa_split_keys(t_kv_first, nkv) != crate::fa_split_keys(t_kv_last, nkv)
2290 {
2291 return None;
2292 }
2293 let rung = t_kv_last.next_power_of_two().max(256);
2294 if crate::fa_split_keys(rung, nkv) != crate::fa_split_keys(t_kv_last, nkv) {
2295 return None;
2296 }
2297 Some(rung)
2298 }
2299
2300 /// Run the WHOLE verify walk [lo, hi) as one captured graph at (vt=t, rung): stage
2301 /// the residual + refresh the per-vt position staging, capture on first encounter
2302 /// (2 executing warmups bracketed by a full linear-state save/restore; KV warmup
2303 /// appends write the exact slots the replay writes — idempotent), launch, then apply
2304 /// the host bookkeeping the captured body skipped (per-linear-layer parity swap for
2305 /// odd t, per-fa-layer len bump). Returns the fresh residual.
2306 #[allow(clippy::too_many_arguments)]
2307 pub(crate) fn run_full(
2308 &mut self,
2309 model: &crate::hybrid::HybridModel,
2310 e: &Engine,
2311 lo: usize,
2312 hi: usize,
2313 x: &CudaSlice<f32>,
2314 t: usize,
2315 pos0: usize,
2316 rung: usize,
2317 cache: &mut Cache,
2318 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2319 let n_embd = self.n_embd;
2320 if !self.stage.contains_key(&t) {
2321 let xin = e.uninit(t * n_embd)?;
2322 let xout = e.uninit(t * n_embd)?;
2323 self.stage.insert(t, (xin, xout));
2324 }
2325 if !self.pos_stage.contains_key(&t) {
2326 self.pos_stage.insert(t, e.htod_i32(&vec![0i32; t])?);
2327 }
2328 // Per-round refresh: position contents + input staging (both addresses are baked
2329 // by the captured bodies; only their CONTENTS change round to round).
2330 {
2331 let pos_host: Vec<i32> = (0..t).map(|r| (pos0 + r) as i32).collect();
2332 let pb = self.pos_stage.get_mut(&t).unwrap();
2333 e.htod_i32_into(pb, &pos_host)?;
2334 let (xin, _) = self.stage.get_mut(&t).unwrap();
2335 e.copy_into(xin, 0, x, t * n_embd)?;
2336 }
2337 let key = (t, rung, hi);
2338 if !self.full.contains_key(&key) {
2339 // The warmups EXECUTE the whole walk on live state — save every linear
2340 // layer's conv + canonical ssm first, restore after (KV needs no restore:
2341 // graph mode never bumps host lens and the appends write this round's own
2342 // slots).
2343 for (k, &il) in self.lin.iter().enumerate() {
2344 let rl = cache.recur[il].as_ref().unwrap();
2345 e.copy_into(
2346 &mut self.save_conv,
2347 k * self.conv_words,
2348 &rl.conv_state,
2349 self.conv_words,
2350 )?;
2351 e.copy_into(
2352 &mut self.save_ssm,
2353 k * self.ssm_words,
2354 &rl.ssm_state,
2355 self.ssm_words,
2356 )?;
2357 }
2358 let (graph, keeper) = {
2359 let table_all = &self.table_all;
2360 let lin_pos = &self.lin_pos;
2361 let fa_pos = &self.fa_pos;
2362 let fa_table = &self.fa_table;
2363 let t_cap = self.t_cap;
2364 let stash_conv = &mut self.stash_conv;
2365 let stash_ssm = &mut self.stash_ssm;
2366 let pos_d: &CudaSlice<i32> = &self.pos_stage[&t];
2367 let (xin, xout) = self
2368 .stage
2369 .get_mut(&t)
2370 .map(|(a, b)| (&*a, b))
2371 .expect("stage bucket created above");
2372 let cache_ref: &mut Cache = cache;
2373 let iflag = if std::env::var("MEMRA_DSPARK_VG_AUTOFREE").as_deref() == Ok("1") {
2374 cudarc::driver::sys::CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH
2375 } else {
2376 cudarc::driver::sys::CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_USE_NODE_PRIORITY
2377 };
2378 e.capture_graph_retained_flags(iflag, move |e| {
2379 let mut xc: Option<CudaSlice<f32>> = None;
2380 for il in lo..hi {
2381 let xr: &CudaSlice<f32> = xc.as_ref().unwrap_or(xin);
2382 let nx = if let Some(&k) = lin_pos.get(&il) {
2383 model.qwen35_tparallel_linear_layer(
2384 e,
2385 il,
2386 xr,
2387 t,
2388 cache_ref,
2389 None,
2390 Some((&mut stash_conv[k], &mut stash_ssm[k])),
2391 Some((table_all, k * 6)),
2392 )?
2393 } else if let Some(&kf) = fa_pos.get(&il) {
2394 let mut no_rows: Option<Vec<CudaSlice<i32>>> = None;
2395 model.qwen35_tparallel_fa_layer(
2396 e,
2397 il,
2398 xr,
2399 t,
2400 cache_ref,
2401 FaLayerArgs {
2402 pos_d,
2403 pos_rows: &mut no_rows,
2404 pos0,
2405 seqs_append: true,
2406 batch_fa_on: true,
2407 graph_cap: Some((fa_table, kf * 2 * t_cap, rung)),
2408 stream: None,
2409 ckpt: None,
2410 },
2411 )?
2412 } else {
2413 return Err(format!(
2414 "run_full: layer {il} is neither linear nor full-attention"
2415 )
2416 .into());
2417 };
2418 xc = Some(nx);
2419 }
2420 e.copy_into(xout, 0, xc.as_ref().unwrap(), t * n_embd)?;
2421 Ok(())
2422 })?
2423 };
2424 // Undo the net host parity motion of the 3 body runs (each run swaps iff t
2425 // is odd -> 3 runs = net one swap), then restore the device state the
2426 // warmups consumed (walk scope only — layers past hi never executed). The
2427 // launch below then behaves exactly like one run.
2428 if t % 2 == 1 {
2429 for &il in &self.lin {
2430 if il < lo || il >= hi {
2431 continue;
2432 }
2433 let rl = cache.recur[il].as_mut().unwrap();
2434 std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
2435 }
2436 }
2437 for (k, &il) in self.lin.iter().enumerate() {
2438 if il < lo || il >= hi {
2439 continue;
2440 }
2441 let rl = cache.recur[il].as_mut().unwrap();
2442 let (cw, sw) = (self.conv_words, self.ssm_words);
2443 {
2444 let sv = e.view(&self.save_conv, self.lin.len() * cw);
2445 let win = sv.slice(k * cw..(k + 1) * cw);
2446 e.copy_view_into(&mut rl.conv_state, 0, &win, cw)?;
2447 }
2448 {
2449 let sv = e.view(&self.save_ssm, self.lin.len() * sw);
2450 let win = sv.slice(k * sw..(k + 1) * sw);
2451 e.copy_view_into(&mut rl.ssm_state, 0, &win, sw)?;
2452 }
2453 }
2454 if std::env::var("MEMRA_GRAPH_CENSUS").as_deref() == Ok("1") {
2455 if let Ok(c) = crate::graph_update::node_census(&graph) {
2456 eprintln!("[dspark-vg-census] full vt={t} rung={rung} {c:?}");
2457 }
2458 }
2459 self.full.insert(
2460 key,
2461 DsparkSegGraph {
2462 graph,
2463 _keeper: keeper,
2464 },
2465 );
2466 }
2467 self.full[&key].graph.launch()?;
2468 // Host bookkeeping for the replayed body (captured host code does not re-run):
2469 // gdn parity swap per linear layer (t odd), kv len bump per fa layer — scoped
2470 // to the WALK [lo, hi): the cache can carry extra state slots past it (the MTP
2471 // head layer's kv) that the walk never touches.
2472 if t % 2 == 1 {
2473 for &il in &self.lin {
2474 if il < lo || il >= hi {
2475 continue;
2476 }
2477 let rl = cache.recur[il].as_mut().unwrap();
2478 std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
2479 }
2480 }
2481 for &il in &self.fa {
2482 if il < lo || il >= hi {
2483 continue;
2484 }
2485 cache.kv[il].as_mut().unwrap().len += t;
2486 }
2487 let (_, xout) = self.stage.get(&t).unwrap();
2488 let mut out = e.uninit(t * n_embd)?;
2489 e.copy_into(&mut out, 0, xout, t * n_embd)?;
2490 Ok(out)
2491 }
2492
2493 /// Run layers [start, end) (all linear) as one captured graph at this vt: stage the
2494 /// residual into the bucket's x_in, capture on first encounter (2 executing warmups
2495 /// bracketed by a segment state save/restore), launch, then apply the host parity
2496 /// bookkeeping the captured body would have done. Returns the fresh residual.
2497 #[allow(clippy::too_many_arguments)]
2498 fn run_segment(
2499 &mut self,
2500 model: &crate::hybrid::HybridModel,
2501 e: &Engine,
2502 start: usize,
2503 end: usize,
2504 x: &CudaSlice<f32>,
2505 t: usize,
2506 cache: &mut Cache,
2507 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2508 let n_embd = self.n_embd;
2509 debug_assert!(end - start <= self.max_run);
2510 if !self.stage.contains_key(&t) {
2511 let xin = e.uninit(t * n_embd)?;
2512 let xout = e.uninit(t * n_embd)?;
2513 self.stage.insert(t, (xin, xout));
2514 }
2515 // Stage the residual at the bucket's baked input address.
2516 {
2517 let (xin, _) = self.stage.get_mut(&t).unwrap();
2518 e.copy_into(xin, 0, x, t * n_embd)?;
2519 }
2520 let key = (start, t);
2521 if !self.graphs.contains_key(&key) {
2522 // The 2 warmups EXECUTE the segment on live state — save conv + the canonical
2523 // ssm of every segment layer first, restore after, so the graph's first real
2524 // launch starts from the exact pre-round state (bytes gated e2e).
2525 for (k, il) in (start..end).enumerate() {
2526 let rl = cache.recur[il].as_ref().unwrap();
2527 e.copy_into(
2528 &mut self.save_conv,
2529 k * self.conv_words,
2530 &rl.conv_state,
2531 self.conv_words,
2532 )?;
2533 e.copy_into(
2534 &mut self.save_ssm,
2535 k * self.ssm_words,
2536 &rl.ssm_state,
2537 self.ssm_words,
2538 )?;
2539 }
2540 let (graph, keeper) = {
2541 let table_all = &self.table_all;
2542 let lin_pos = &self.lin_pos;
2543 let stash_conv = &mut self.stash_conv;
2544 let stash_ssm = &mut self.stash_ssm;
2545 let (xin, xout) = self
2546 .stage
2547 .get_mut(&t)
2548 .map(|(a, b)| (&*a, b))
2549 .expect("stage bucket created above");
2550 let cache_ref: &mut Cache = cache;
2551 // Slice 4 (fa-execupdate lane): USE_NODE_PRIORITY instead of
2552 // AUTO_FREE_ON_LAUNCH. The slice-3 measured limiter was AUTO_FREE's
2553 // launch-time mem-pool scan — 25.6 us per cuGraphLaunch x 16 segments
2554 // = ~0.41 ms/round, most of the eager-launch savings. The captured
2555 // body's cuMemAllocAsync transients are BALANCED by in-graph frees
2556 // (every transient drops inside the capture region — the generic
2557 // capture path's census precedent, 1589/1589), so AUTO_FREE has
2558 // nothing to reclaim and the graph is legal to instantiate without
2559 // it; PRIORITY is the flag the gemma slotted door ships for exactly
2560 // this reason (both alternatives drop the scan; UPLOAD via
2561 // cuGraphInstantiateWithFlags is WithParams-only and refused).
2562 // MEMRA_DSPARK_VG_AUTOFREE=1 reverts; MEMRA_GRAPH_CENSUS=1 prints
2563 // the node census at capture (the ALLOC==FREE receipt).
2564 let iflag = if std::env::var("MEMRA_DSPARK_VG_AUTOFREE").as_deref() == Ok("1") {
2565 cudarc::driver::sys::CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH
2566 } else {
2567 cudarc::driver::sys::CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_USE_NODE_PRIORITY
2568 };
2569 e.capture_graph_retained_flags(iflag, move |e| {
2570 let mut xc: Option<CudaSlice<f32>> = None;
2571 for il in start..end {
2572 let k = lin_pos[&il];
2573 let xr: &CudaSlice<f32> = xc.as_ref().unwrap_or(xin);
2574 let nx = model.qwen35_tparallel_linear_layer(
2575 e,
2576 il,
2577 xr,
2578 t,
2579 cache_ref,
2580 None,
2581 Some((&mut stash_conv[k], &mut stash_ssm[k])),
2582 Some((table_all, k * 6)),
2583 )?;
2584 xc = Some(nx);
2585 }
2586 e.copy_into(xout, 0, xc.as_ref().unwrap(), t * n_embd)?;
2587 Ok(())
2588 })?
2589 };
2590 // Undo the net host parity motion of the 3 body runs (each run swaps iff t
2591 // is odd -> 3 runs = net one swap), then restore the device state the
2592 // warmups consumed. The launch below then behaves exactly like one run.
2593 if t % 2 == 1 {
2594 for il in start..end {
2595 let rl = cache.recur[il].as_mut().unwrap();
2596 std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
2597 }
2598 }
2599 for (k, il) in (start..end).enumerate() {
2600 let rl = cache.recur[il].as_mut().unwrap();
2601 let (cw, sw) = (self.conv_words, self.ssm_words);
2602 {
2603 let sv = e.view(&self.save_conv, self.lin.len() * cw);
2604 let win = sv.slice(k * cw..(k + 1) * cw);
2605 e.copy_view_into(&mut rl.conv_state, 0, &win, cw)?;
2606 }
2607 {
2608 let sv = e.view(&self.save_ssm, self.lin.len() * sw);
2609 let win = sv.slice(k * sw..(k + 1) * sw);
2610 e.copy_view_into(&mut rl.ssm_state, 0, &win, sw)?;
2611 }
2612 }
2613 if std::env::var("MEMRA_GRAPH_CENSUS").as_deref() == Ok("1") {
2614 if let Ok(c) = crate::graph_update::node_census(&graph) {
2615 eprintln!("[dspark-vg-census] seg={start}..{end} vt={t} {c:?}");
2616 }
2617 }
2618 self.graphs.insert(
2619 key,
2620 DsparkSegGraph {
2621 graph,
2622 _keeper: keeper,
2623 },
2624 );
2625 }
2626 self.graphs[&key].graph.launch()?;
2627 // Host parity bookkeeping for the replayed body (the captured host swaps do not
2628 // re-run at replay).
2629 if t % 2 == 1 {
2630 for il in start..end {
2631 let rl = cache.recur[il].as_mut().unwrap();
2632 std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
2633 }
2634 }
2635 let (_, xout) = self.stage.get(&t).unwrap();
2636 let mut out = e.uninit(t * n_embd)?;
2637 e.copy_into(&mut out, 0, xout, t * n_embd)?;
2638 Ok(out)
2639 }
2640
2641 /// Pool freeze check (`dspark_vg_cap`): below the ceiling new keys may capture.
2642 fn can_capture(&self) -> bool {
2643 self.graphs.len() + self.full.len() < dspark_vg_cap()
2644 }
2645
2646 /// Round-atomic segment-door readiness: TRUE when this round's walk can ride the
2647 /// per-(segment, vt) graphs without a NEW capture past the pool ceiling — every
2648 /// linear run in [lo, hi) already has its (run_start, t) key, or capture is still
2649 /// allowed. FALSE sends the WHOLE round down the eager cols-ckpt walk: a partial
2650 /// refusal would stash some layers in the ctx slabs and others in the round's cols
2651 /// while one commit reads only one of them.
2652 pub(crate) fn segments_ready(
2653 &self,
2654 model: &crate::hybrid::HybridModel,
2655 lo: usize,
2656 hi: usize,
2657 t: usize,
2658 ) -> bool {
2659 if self.can_capture() {
2660 return true;
2661 }
2662 let mut il = lo;
2663 while il < hi {
2664 if matches!(model.layers[il].mixer, Mixer::Linear(_)) {
2665 let start = il;
2666 while il < hi && matches!(model.layers[il].mixer, Mixer::Linear(_)) {
2667 il += 1;
2668 }
2669 if !self.graphs.contains_key(&(start, t)) {
2670 return false;
2671 }
2672 } else {
2673 il += 1;
2674 }
2675 }
2676 true
2677 }
2678
2679 /// Widest verify window this pool was built for. A caller whose round exceeds it must
2680 /// take the eager walk: the stash slabs hold `t_capacity() - 1` column rows, and slicing
2681 /// past them is a panic rather than a refusal.
2682 pub(crate) fn t_capacity(&self) -> usize {
2683 self.t_cap
2684 }
2685
2686 /// Slab row (conv, ssm) device pointers + lengths for the commit restore of column
2687 /// `row` (0-based) of layer `il`. None for non-linear layers.
2688 pub(crate) fn slab_row(
2689 &self,
2690 e: &Engine,
2691 il: usize,
2692 row: usize,
2693 ) -> Option<(u64, u64, usize, usize)> {
2694 use cudarc::driver::DevicePtr;
2695 let k = *self.lin_pos.get(&il)?;
2696 let s = &e.gpu.stream();
2697 let (pc, _g0) = self.stash_conv[k].device_ptr(s);
2698 let (ps, _g1) = self.stash_ssm[k].device_ptr(s);
2699 Some((
2700 pc as u64 + (row * self.conv_words * 4) as u64,
2701 ps as u64 + (row * self.ssm_words * 4) as u64,
2702 self.conv_words,
2703 self.ssm_words,
2704 ))
2705 }
2706}
2707
2708impl VerifyCkpt {
2709 fn new(n_layer: usize) -> Self {
2710 VerifyCkpt {
2711 gdn: (0..n_layer).map(|_| None).collect(),
2712 cols: (0..n_layer).map(|_| None).collect(),
2713 }
2714 }
2715}
2716
2717/// The stage-0/TX half of one PP verify. The boundary slot is the ownership token: stage 1
2718/// consumes exactly the slot selected by `tx()` / `tx_pipelined()`, never a slot inferred from
2719/// a logical round number.
2720struct VerifyBoundaryTicket {
2721 rt: &'static crate::pp::PpNRt,
2722 caller_stream: std::sync::Arc<cudarc::driver::CudaStream>,
2723 slot: usize,
2724 pos0: usize,
2725 t: usize,
2726 payload: usize,
2727 n_st: usize,
2728 pipelined: bool,
2729 pp_anatomy: bool,
2730 pp_started: std::time::Instant,
2731 reverse_ms: f64,
2732 stage0_ms: f64,
2733 tx_ms: f64,
2734 trace: Option<SpecPipeTraceCtx>,
2735}
2736
2737/// Explicit OPTIPIPE diagnostic control. Forced modes are set only by `optipipe-gate`; the
2738/// increment-2 controller can also be armed by the server's fresh-process research door.
2739#[derive(Clone, Copy, Debug, PartialEq, Eq)]
2740pub enum OptiForkGateMode {
2741 Disabled,
2742 Hit,
2743 Miss,
2744 Alternate,
2745 Abort,
2746 Controller,
2747}
2748
2749static OPTI_FORK_GATE_MODE: std::sync::atomic::AtomicU8 = std::sync::atomic::AtomicU8::new(0);
2750static OPTI_CONTROLLER_THRESHOLD: std::sync::atomic::AtomicU32 =
2751 std::sync::atomic::AtomicU32::new(0);
2752static OPTI_FORK_ATTEMPTS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
2753static OPTI_FORK_HITS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
2754static OPTI_FORK_MISSES: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
2755static OPTI_FORK_ABORT_DRAINS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
2756static OPTI_FORK_REFUSALS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
2757static OPTI_GATE_CHECKS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
2758static OPTI_GATE_ADMITS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
2759static OPTI_GATE_REJECTS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
2760static OPTI_RECONCILES: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
2761static OPTI_WASTED_DRAFT_TOKENS: std::sync::atomic::AtomicU64 =
2762 std::sync::atomic::AtomicU64::new(0);
2763static OPTI_SHADOW_DRAFT_TOKENS: std::sync::atomic::AtomicU64 =
2764 std::sync::atomic::AtomicU64::new(0);
2765static OPTI_BREAKER_TRIPS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
2766
2767impl OptiForkGateMode {
2768 fn code(self) -> u8 {
2769 match self {
2770 Self::Disabled => 0,
2771 Self::Hit => 1,
2772 Self::Miss => 2,
2773 Self::Alternate => 3,
2774 Self::Abort => 4,
2775 Self::Controller => 5,
2776 }
2777 }
2778
2779 fn configured() -> Self {
2780 match OPTI_FORK_GATE_MODE.load(std::sync::atomic::Ordering::Relaxed) {
2781 1 => Self::Hit,
2782 2 => Self::Miss,
2783 3 => Self::Alternate,
2784 4 => Self::Abort,
2785 5 => Self::Controller,
2786 _ => Self::Disabled,
2787 }
2788 }
2789
2790 fn action(self, generation: u64) -> OptiForkAction {
2791 match self {
2792 Self::Hit => OptiForkAction::Hit,
2793 Self::Miss => OptiForkAction::Miss,
2794 Self::Alternate if generation & 1 == 0 => OptiForkAction::Hit,
2795 Self::Alternate => OptiForkAction::Miss,
2796 Self::Abort => OptiForkAction::Abort,
2797 Self::Disabled | Self::Controller => {
2798 unreachable!("non-forced mode cannot choose a forced fork action")
2799 }
2800 }
2801 }
2802
2803 fn is_forced(self) -> bool {
2804 matches!(self, Self::Hit | Self::Miss | Self::Alternate | Self::Abort)
2805 }
2806}
2807
2808/// Arm or disarm the forced harness. Serving uses only `set_optipipe_controller_threshold`.
2809pub fn set_optipipe_gate_mode(mode: OptiForkGateMode) {
2810 OPTI_FORK_GATE_MODE.store(mode.code(), std::sync::atomic::Ordering::Relaxed);
2811}
2812
2813/// Arm the increment-2 diagnostic controller. The threshold applies to the uncalibrated
2814/// two-token draft-probability product. Serving can call this only through its explicit
2815/// fresh-process research door; the absent-door default remains byte-for-byte disabled.
2816pub fn set_optipipe_controller_threshold(threshold: f32) {
2817 assert!(threshold.is_finite() && (0.0..=1.0).contains(&threshold));
2818 OPTI_CONTROLLER_THRESHOLD.store(threshold.to_bits(), std::sync::atomic::Ordering::Relaxed);
2819 set_optipipe_gate_mode(OptiForkGateMode::Controller);
2820}
2821
2822#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
2823pub struct OptiForkGateStats {
2824 pub attempts: u64,
2825 pub hits: u64,
2826 pub misses: u64,
2827 pub abort_drains: u64,
2828 pub refusals: u64,
2829 pub gate_checks: u64,
2830 pub gate_admits: u64,
2831 pub gate_rejects: u64,
2832 pub reconciles: u64,
2833 pub wasted_draft_tokens: u64,
2834 pub shadow_draft_tokens: u64,
2835 pub breaker_trips: u64,
2836}
2837
2838#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
2839pub struct OptiForkStateIdentity {
2840 pub trunk_kv_bytes: usize,
2841 pub recurrent_bytes: usize,
2842 pub scratch_kv_bytes: usize,
2843 pub hidden_bytes: usize,
2844}
2845
2846pub fn reset_optipipe_gate_stats() {
2847 for counter in [
2848 &OPTI_FORK_ATTEMPTS,
2849 &OPTI_FORK_HITS,
2850 &OPTI_FORK_MISSES,
2851 &OPTI_FORK_ABORT_DRAINS,
2852 &OPTI_FORK_REFUSALS,
2853 &OPTI_GATE_CHECKS,
2854 &OPTI_GATE_ADMITS,
2855 &OPTI_GATE_REJECTS,
2856 &OPTI_RECONCILES,
2857 &OPTI_WASTED_DRAFT_TOKENS,
2858 &OPTI_SHADOW_DRAFT_TOKENS,
2859 &OPTI_BREAKER_TRIPS,
2860 ] {
2861 counter.store(0, std::sync::atomic::Ordering::Relaxed);
2862 }
2863}
2864
2865pub fn optipipe_gate_stats() -> OptiForkGateStats {
2866 let load = |v: &std::sync::atomic::AtomicU64| v.load(std::sync::atomic::Ordering::Relaxed);
2867 OptiForkGateStats {
2868 attempts: load(&OPTI_FORK_ATTEMPTS),
2869 hits: load(&OPTI_FORK_HITS),
2870 misses: load(&OPTI_FORK_MISSES),
2871 abort_drains: load(&OPTI_FORK_ABORT_DRAINS),
2872 refusals: load(&OPTI_FORK_REFUSALS),
2873 gate_checks: load(&OPTI_GATE_CHECKS),
2874 gate_admits: load(&OPTI_GATE_ADMITS),
2875 gate_rejects: load(&OPTI_GATE_REJECTS),
2876 reconciles: load(&OPTI_RECONCILES),
2877 wasted_draft_tokens: load(&OPTI_WASTED_DRAFT_TOKENS),
2878 shadow_draft_tokens: load(&OPTI_SHADOW_DRAFT_TOKENS),
2879 breaker_trips: load(&OPTI_BREAKER_TRIPS),
2880 }
2881}
2882
2883#[derive(Clone, Copy, Debug)]
2884struct OptiControllerPolicy {
2885 threshold: f32,
2886 consecutive_misses: u8,
2887 breaker_tripped: bool,
2888}
2889
2890impl OptiControllerPolicy {
2891 fn configured() -> Self {
2892 Self {
2893 threshold: f32::from_bits(
2894 OPTI_CONTROLLER_THRESHOLD.load(std::sync::atomic::Ordering::Relaxed),
2895 ),
2896 consecutive_misses: 0,
2897 breaker_tripped: false,
2898 }
2899 }
2900
2901 fn admit(&self, q_proxy: f32) -> bool {
2902 q_proxy.is_finite()
2903 && (0.0..=1.0).contains(&q_proxy)
2904 && (self.threshold == 0.0 || (!self.breaker_tripped && q_proxy >= self.threshold))
2905 }
2906
2907 /// Returns true exactly when this resolution newly trips the three-miss breaker.
2908 fn resolve(&mut self, hit: bool) -> bool {
2909 // q*=0 is the lane's explicit unconditional measurement arm. Its purpose is to price
2910 // every optimistic opportunity, so the safety breaker is measured separately and must
2911 // not silently turn this arm into "three attempts then serial".
2912 if self.threshold == 0.0 {
2913 self.consecutive_misses = 0;
2914 return false;
2915 }
2916 if hit {
2917 self.consecutive_misses = 0;
2918 return false;
2919 }
2920 self.consecutive_misses = self.consecutive_misses.saturating_add(1);
2921 if !self.breaker_tripped && self.consecutive_misses >= 3 {
2922 self.breaker_tripped = true;
2923 return true;
2924 }
2925 false
2926 }
2927}
2928
2929#[derive(Clone, Copy, Debug, PartialEq, Eq)]
2930enum OptiForkAction {
2931 Hit,
2932 Miss,
2933 Abort,
2934}
2935
2936#[derive(Clone, Copy, Debug, PartialEq, Eq)]
2937struct OptiForkGeneration {
2938 id: u64,
2939 slot: usize,
2940}
2941
2942#[derive(Default)]
2943struct OptiForkGenerationTracker {
2944 next: u64,
2945 live: [Option<u64>; 2],
2946}
2947
2948impl OptiForkGenerationTracker {
2949 fn reserve(&mut self) -> Result<OptiForkGeneration, Box<dyn std::error::Error>> {
2950 let generation = OptiForkGeneration {
2951 id: self.next,
2952 slot: (self.next & 1) as usize,
2953 };
2954 if let Some(live) = self.live[generation.slot] {
2955 return Err(format!(
2956 "optipipe snapshot slot {} still owns generation {live}; refusing to overwrite it",
2957 generation.slot,
2958 )
2959 .into());
2960 }
2961 self.next += 1;
2962 self.live[generation.slot] = Some(generation.id);
2963 Ok(generation)
2964 }
2965
2966 fn retire(&mut self, generation: OptiForkGeneration) -> Result<(), Box<dyn std::error::Error>> {
2967 match self.live[generation.slot] {
2968 Some(id) if id == generation.id => {
2969 self.live[generation.slot] = None;
2970 Ok(())
2971 }
2972 other => Err(format!(
2973 "optipipe generation teardown mismatch: ticket={} slot={} live={other:?}",
2974 generation.id, generation.slot,
2975 )
2976 .into()),
2977 }
2978 }
2979}
2980
2981struct OptiForkSeedGeneration {
2982 h_seed: CudaSlice<f32>,
2983 fill_prev: CudaSlice<f32>,
2984 scratch_len: usize,
2985}
2986
2987/// Allocate or refresh one full checkpoint through the engine that owns each PP stage. The
2988/// generic cache helper accepts one device and therefore cannot copy GDN state split across
2989/// devices. KV lengths and position stay host metadata; only recurrent buffers need stage-local
2990/// device ownership.
2991fn opti_snapshot_stage_owned(
2992 e: &Engine,
2993 cache: &Cache,
2994 rt: &'static crate::pp::PpNRt,
2995 fence: &[usize],
2996) -> Result<crate::cache::CacheSnapshot, Box<dyn std::error::Error>> {
2997 let n = cache.kv.len();
2998 let mut snapshot = crate::cache::CacheSnapshot {
2999 kv_len: vec![None; n],
3000 tp_kv_len: vec![None; n],
3001 conv: (0..n).map(|_| None).collect(),
3002 ssm: (0..n).map(|_| None).collect(),
3003 pos: cache.pos,
3004 };
3005 opti_snapshot_stage_owned_into(e, cache, rt, fence, &mut snapshot)?;
3006 Ok(snapshot)
3007}
3008
3009fn opti_snapshot_stage_owned_into(
3010 e: &Engine,
3011 cache: &Cache,
3012 rt: &'static crate::pp::PpNRt,
3013 fence: &[usize],
3014 snapshot: &mut crate::cache::CacheSnapshot,
3015) -> Result<(), Box<dyn std::error::Error>> {
3016 if fence.len() != rt.n_stages() + 1
3017 || snapshot.kv_len.len() != cache.kv.len()
3018 || snapshot.tp_kv_len.len() != cache.tp_kv.len()
3019 {
3020 return Err("optipipe stage-owned snapshot shape mismatch".into());
3021 }
3022 for stage in 0..rt.n_stages() {
3023 opti_snapshot_one_stage_owned_into(e, cache, rt, fence, stage, snapshot)?;
3024 }
3025 snapshot.pos = cache.pos;
3026 Ok(())
3027}
3028
3029/// Refresh one PP stage of a checkpoint. Increment 2 uses this split form so stage 0's
3030/// optimistic post-N state is captured before N+1 stage 0 is queued, while stage 1's matching
3031/// post-N state is captured only after N stage 1 is enqueued. Calling the all-stage helper at
3032/// either point would capture one side of the fork at the wrong generation.
3033fn opti_snapshot_one_stage_owned_into(
3034 e: &Engine,
3035 cache: &Cache,
3036 rt: &'static crate::pp::PpNRt,
3037 fence: &[usize],
3038 stage: usize,
3039 snapshot: &mut crate::cache::CacheSnapshot,
3040) -> Result<(), Box<dyn std::error::Error>> {
3041 if fence.len() != rt.n_stages() + 1
3042 || snapshot.kv_len.len() != cache.kv.len()
3043 || snapshot.tp_kv_len.len() != cache.tp_kv.len()
3044 || stage >= rt.n_stages()
3045 {
3046 return Err("optipipe single-stage snapshot shape mismatch".into());
3047 }
3048 let _scope = rt.enter(stage);
3049 let owner = rt.engine(stage, e);
3050 for il in fence[stage]..fence[stage + 1] {
3051 snapshot.kv_len[il] = cache.kv[il].as_ref().map(|kv| kv.len);
3052 snapshot.tp_kv_len[il] = cache.tp_kv[il]
3053 .as_ref()
3054 .map(crate::tp::ResidentTpKvCache::committed_len);
3055 match &cache.recur[il] {
3056 Some(recur) => {
3057 match snapshot.conv[il].as_mut() {
3058 Some(dst) => {
3059 owner.copy_into(dst, 0, &recur.conv_state, recur.conv_state.len())?
3060 }
3061 None => snapshot.conv[il] = Some(owner.clone_dtod(&recur.conv_state)?),
3062 }
3063 match snapshot.ssm[il].as_mut() {
3064 Some(dst) => {
3065 owner.copy_into(dst, 0, &recur.ssm_state, recur.ssm_state.len())?
3066 }
3067 None => snapshot.ssm[il] = Some(owner.clone_dtod(&recur.ssm_state)?),
3068 }
3069 }
3070 None if snapshot.conv[il].is_some() || snapshot.ssm[il].is_some() => {
3071 return Err(
3072 format!("optipipe stage-owned snapshot layer {il} changed shape").into(),
3073 );
3074 }
3075 None => {}
3076 }
3077 }
3078 snapshot.pos = cache.pos;
3079 Ok(())
3080}
3081
3082/// Increment-1 persistent fork state. Exactly two snapshot/seed slots alternate; a live ticket
3083/// names its generation and keeps teardown fail-closed. Only stage 0 is allowed to mutate before
3084/// resolve, so the reconcile tables and conditional restores are stage-local.
3085struct OptiForkState {
3086 mode: OptiForkGateMode,
3087 controller: Option<OptiControllerPolicy>,
3088 generations: OptiForkGenerationTracker,
3089 active_snapshot_slot: usize,
3090 alternate_snapshot: crate::cache::CacheSnapshot,
3091 seeds: [OptiForkSeedGeneration; 2],
3092 rt: &'static crate::pp::PpNRt,
3093 fence: [usize; 3],
3094 split: usize,
3095 len_ptrs: CudaSlice<u64>,
3096 saved_lens: CudaSlice<i32>,
3097 forced_acc: CudaSlice<u32>,
3098 valid: CudaSlice<u32>,
3099 stage0_stream: std::sync::Arc<cudarc::driver::CudaStream>,
3100 logical_payload_bytes: [usize; 2],
3101}
3102
3103struct OptiForkTicket {
3104 generation: OptiForkGeneration,
3105 boundary: Option<VerifyBoundaryTicket>,
3106 drain: std::sync::Arc<cudarc::driver::CudaStream>,
3107 settled: bool,
3108}
3109
3110struct OptiControllerTicket {
3111 generation: OptiForkGeneration,
3112 boundary: Option<VerifyBoundaryTicket>,
3113 ckpt: Option<VerifyCkpt>,
3114 verify_tokens: [u32; 2],
3115 draft_prob: f32,
3116 eager_seed: Option<CudaSlice<f32>>,
3117 q_proxy: f32,
3118 scratch_len: usize,
3119 issued_at: std::time::Instant,
3120 drain: std::sync::Arc<cudarc::driver::CudaStream>,
3121 settled: bool,
3122}
3123
3124struct OptiControllerPrepared {
3125 verify_tokens: [u32; 2],
3126 draft_prob: f32,
3127 eager_seed: Option<CudaSlice<f32>>,
3128 q_proxy: f32,
3129 scratch_len: usize,
3130}
3131
3132impl OptiControllerTicket {
3133 fn take_boundary(&mut self) -> VerifyBoundaryTicket {
3134 self.boundary
3135 .take()
3136 .expect("controller boundary ticket already consumed")
3137 }
3138
3139 fn take_ckpt(&mut self) -> VerifyCkpt {
3140 self.ckpt
3141 .take()
3142 .expect("controller verify checkpoint already consumed")
3143 }
3144
3145 fn take_eager_seed(&mut self) -> Option<CudaSlice<f32>> {
3146 self.eager_seed.take()
3147 }
3148
3149 fn settle(&mut self) {
3150 self.settled = true;
3151 }
3152}
3153
3154impl Drop for OptiControllerTicket {
3155 fn drop(&mut self) {
3156 if !self.settled {
3157 let _ = self.drain.synchronize();
3158 OPTI_FORK_ABORT_DRAINS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
3159 }
3160 }
3161}
3162
3163impl OptiForkTicket {
3164 fn take_boundary(&mut self) -> VerifyBoundaryTicket {
3165 self.boundary
3166 .take()
3167 .expect("fork ticket boundary already consumed")
3168 }
3169
3170 fn settle(&mut self) {
3171 self.settled = true;
3172 }
3173}
3174
3175impl Drop for OptiForkTicket {
3176 fn drop(&mut self) {
3177 if !self.settled {
3178 let _ = self.drain.synchronize();
3179 OPTI_FORK_ABORT_DRAINS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
3180 }
3181 }
3182}
3183
3184impl OptiForkState {
3185 #[allow(clippy::too_many_arguments)]
3186 fn new(
3187 e: &Engine,
3188 cache: &Cache,
3189 mode: OptiForkGateMode,
3190 alternate_snapshot: crate::cache::CacheSnapshot,
3191 h_seed: &CudaSlice<f32>,
3192 fill_prev: &CudaSlice<f32>,
3193 rt: &'static crate::pp::PpNRt,
3194 split: usize,
3195 n_layer: usize,
3196 ) -> Result<Self, Box<dyn std::error::Error>> {
3197 let fence = [0, split, n_layer];
3198 let mut logical_payload_bytes = [0usize; 2];
3199 for stage in 0..2 {
3200 for il in fence[stage]..fence[stage + 1] {
3201 logical_payload_bytes[stage] += alternate_snapshot.conv[il]
3202 .as_ref()
3203 .map_or(0, |v| v.len() * std::mem::size_of::<f32>());
3204 logical_payload_bytes[stage] += alternate_snapshot.ssm[il]
3205 .as_ref()
3206 .map_or(0, |v| v.len() * std::mem::size_of::<f32>());
3207 }
3208 }
3209 let seeds = [
3210 OptiForkSeedGeneration {
3211 h_seed: e.clone_dtod(h_seed)?,
3212 fill_prev: e.clone_dtod(fill_prev)?,
3213 scratch_len: 0,
3214 },
3215 OptiForkSeedGeneration {
3216 h_seed: e.clone_dtod(h_seed)?,
3217 fill_prev: e.clone_dtod(fill_prev)?,
3218 scratch_len: 0,
3219 },
3220 ];
3221 let (len_ptrs, saved_lens, forced_acc, valid, stage0_stream) = {
3222 let _stage = rt.enter(0);
3223 let e0 = rt.engine(0, e);
3224 (
3225 crate::round_stream::kv_len_ptr_table_range(e0, cache, 0..split, None)?,
3226 e0.htod_i32(&vec![0; split])?,
3227 e0.alloc_u32_zeroed(2)?,
3228 e0.alloc_u32_zeroed(1)?,
3229 e0.stream(),
3230 )
3231 };
3232 logical_payload_bytes[0] += seeds
3233 .iter()
3234 .map(|seed| (seed.h_seed.len() + seed.fill_prev.len()) * std::mem::size_of::<f32>())
3235 .sum::<usize>();
3236 logical_payload_bytes[0] += len_ptrs.len() * std::mem::size_of::<u64>()
3237 + saved_lens.len() * std::mem::size_of::<i32>()
3238 + forced_acc.len() * std::mem::size_of::<u32>()
3239 + valid.len() * std::mem::size_of::<u32>();
3240 Ok(Self {
3241 mode,
3242 controller: (mode == OptiForkGateMode::Controller)
3243 .then(OptiControllerPolicy::configured),
3244 generations: OptiForkGenerationTracker::default(),
3245 active_snapshot_slot: 0,
3246 alternate_snapshot,
3247 seeds,
3248 rt,
3249 fence,
3250 split,
3251 len_ptrs,
3252 saved_lens,
3253 forced_acc,
3254 valid,
3255 stage0_stream,
3256 logical_payload_bytes,
3257 })
3258 }
3259
3260 fn reserve(
3261 &mut self,
3262 current_snapshot: &mut crate::cache::CacheSnapshot,
3263 ) -> Result<OptiForkGeneration, Box<dyn std::error::Error>> {
3264 let generation = self.generations.reserve()?;
3265 if generation.slot != self.active_snapshot_slot {
3266 std::mem::swap(current_snapshot, &mut self.alternate_snapshot);
3267 self.active_snapshot_slot = generation.slot;
3268 }
3269 Ok(generation)
3270 }
3271
3272 fn capture_seed(
3273 &mut self,
3274 e: &Engine,
3275 generation: OptiForkGeneration,
3276 h_seed: &CudaSlice<f32>,
3277 fill_prev: &CudaSlice<f32>,
3278 scratch_len: usize,
3279 ) -> Result<(), Box<dyn std::error::Error>> {
3280 let seed = &mut self.seeds[generation.slot];
3281 e.copy_into(&mut seed.h_seed, 0, h_seed, h_seed.len())?;
3282 e.copy_into(&mut seed.fill_prev, 0, fill_prev, fill_prev.len())?;
3283 seed.scratch_len = scratch_len;
3284 Ok(())
3285 }
3286
3287 fn ticket(
3288 &self,
3289 generation: OptiForkGeneration,
3290 boundary: VerifyBoundaryTicket,
3291 ) -> OptiForkTicket {
3292 OptiForkTicket {
3293 generation,
3294 boundary: Some(boundary),
3295 drain: self.stage0_stream.clone(),
3296 settled: false,
3297 }
3298 }
3299
3300 #[allow(clippy::too_many_arguments)]
3301 fn controller_ticket(
3302 &self,
3303 generation: OptiForkGeneration,
3304 boundary: VerifyBoundaryTicket,
3305 ckpt: VerifyCkpt,
3306 verify_tokens: [u32; 2],
3307 draft_prob: f32,
3308 eager_seed: Option<CudaSlice<f32>>,
3309 q_proxy: f32,
3310 scratch_len: usize,
3311 ) -> OptiControllerTicket {
3312 OptiControllerTicket {
3313 generation,
3314 boundary: Some(boundary),
3315 ckpt: Some(ckpt),
3316 verify_tokens,
3317 draft_prob,
3318 eager_seed,
3319 q_proxy,
3320 scratch_len,
3321 issued_at: std::time::Instant::now(),
3322 drain: self.stage0_stream.clone(),
3323 settled: false,
3324 }
3325 }
3326
3327 fn reserve_successor(&mut self) -> Result<OptiForkGeneration, Box<dyn std::error::Error>> {
3328 self.generations.reserve()
3329 }
3330
3331 fn successor_snapshot_mut(&mut self) -> &mut crate::cache::CacheSnapshot {
3332 &mut self.alternate_snapshot
3333 }
3334
3335 fn promote_successor_snapshot(
3336 &mut self,
3337 current_snapshot: &mut crate::cache::CacheSnapshot,
3338 generation: OptiForkGeneration,
3339 ) {
3340 std::mem::swap(current_snapshot, &mut self.alternate_snapshot);
3341 self.active_snapshot_slot = generation.slot;
3342 }
3343
3344 fn queue_actual_reconcile(
3345 &mut self,
3346 e: &Engine,
3347 snapshot: &crate::cache::CacheSnapshot,
3348 acc: &CudaSlice<u32>,
3349 optimistic_pending: u32,
3350 base: usize,
3351 ) -> Result<(), Box<dyn std::error::Error>> {
3352 let saved: Vec<i32> = (0..self.split)
3353 .map(|il| snapshot.kv_len[il].map(|v| v as i32).unwrap_or(0))
3354 .collect();
3355 // Serving keeps the caller/accept walk on the head (stage-1) device. Record the accept
3356 // decision point there and append a wait to stage 0 after its optimistic successor/TX;
3357 // the validity/reconcile kernels must never peer-read acc before it is written. The
3358 // increment-1 harness uses primary stage 0, where stream order already provides this.
3359 if self.rt.engine(0, e).ctx().ordinal() != e.ctx().ordinal() {
3360 self.rt.fence_stages_behind(&e.stream())?;
3361 }
3362 let _stage = self.rt.enter(0);
3363 let e0 = self.rt.engine(0, e);
3364 e0.htod_i32_into(&mut self.saved_lens, &saved)?;
3365 e0.spec_fork_valid(acc, optimistic_pending, &mut self.valid)?;
3366 e0.spec_fork_reconcile_kv(
3367 &self.len_ptrs,
3368 &self.saved_lens,
3369 acc,
3370 &self.valid,
3371 base,
3372 self.split,
3373 )
3374 }
3375
3376 fn finish_actual_reconcile(
3377 &mut self,
3378 e: &Engine,
3379 cache: &mut Cache,
3380 snapshot: &crate::cache::CacheSnapshot,
3381 n_acc: usize,
3382 base: usize,
3383 hit: bool,
3384 ) -> Result<(), Box<dyn std::error::Error>> {
3385 if hit {
3386 return Ok(());
3387 }
3388 let len_delta = base + n_acc;
3389 for il in 0..self.split {
3390 if let (Some(kv), Some(saved)) = (cache.kv[il].as_mut(), snapshot.kv_len[il]) {
3391 kv.len = saved + len_delta;
3392 }
3393 }
3394 {
3395 let _stage = self.rt.enter(1);
3396 let e1 = self.rt.engine(1, e);
3397 for il in self.split..self.fence[2] {
3398 if let (Some(kv), Some(saved)) = (cache.kv[il].as_mut(), snapshot.kv_len[il]) {
3399 kv.len = saved + len_delta;
3400 e1.set_i32_one(&mut kv.len_d, kv.len as i32)?;
3401 }
3402 }
3403 }
3404 self.rt.publish_to(0, &e.stream())?;
3405 Ok(())
3406 }
3407
3408 fn cancel_controller_ticket(
3409 &mut self,
3410 e: &Engine,
3411 cache: &mut Cache,
3412 scratch: &mut MtpScratch,
3413 snapshot: &crate::cache::CacheSnapshot,
3414 ticket: &mut OptiControllerTicket,
3415 ) -> Result<(), Box<dyn std::error::Error>> {
3416 {
3417 let _stage = self.rt.enter(0);
3418 let e0 = self.rt.engine(0, e);
3419 for il in 0..self.split {
3420 if let (Some(kv), Some(saved)) = (cache.kv[il].as_mut(), snapshot.kv_len[il]) {
3421 kv.len = saved;
3422 e0.set_i32_one(&mut kv.len_d, saved as i32)?;
3423 }
3424 }
3425 }
3426 scratch.set_len(e, snapshot.pos)?;
3427 ticket.settle();
3428 self.generations.retire(ticket.generation)?;
3429 OPTI_FORK_ABORT_DRAINS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
3430 OPTI_WASTED_DRAFT_TOKENS.fetch_add(2, std::sync::atomic::Ordering::Relaxed);
3431 eprintln!(
3432 "[opti-controller] tail-drain generation={} slot={}",
3433 ticket.generation.id, ticket.generation.slot,
3434 );
3435 Ok(())
3436 }
3437
3438 #[allow(clippy::too_many_arguments)]
3439 fn reconcile(
3440 &mut self,
3441 e: &Engine,
3442 cache: &mut Cache,
3443 scratch: &mut MtpScratch,
3444 snapshot: &crate::cache::CacheSnapshot,
3445 h_seed: &mut CudaSlice<f32>,
3446 fill_prev: &mut CudaSlice<f32>,
3447 generation: OptiForkGeneration,
3448 action: OptiForkAction,
3449 optimistic_pending: u32,
3450 ) -> Result<(), Box<dyn std::error::Error>> {
3451 debug_assert!(action != OptiForkAction::Abort);
3452 let miss_started = std::time::Instant::now();
3453 let keep = action == OptiForkAction::Hit;
3454 let saved: Vec<i32> = (0..self.split)
3455 .map(|il| snapshot.kv_len[il].map(|v| v as i32).unwrap_or(0))
3456 .collect();
3457 let seed = &self.seeds[generation.slot];
3458 {
3459 let _stage = self.rt.enter(0);
3460 let e0 = self.rt.engine(0, e);
3461 e0.htod_i32_into(&mut self.saved_lens, &saved)?;
3462 let forced = if keep {
3463 [1u32, optimistic_pending]
3464 } else {
3465 [0u32, optimistic_pending]
3466 };
3467 e0.htod_u32_into(&mut self.forced_acc, &forced)?;
3468 e0.spec_fork_valid(&self.forced_acc, optimistic_pending, &mut self.valid)?;
3469 e0.spec_fork_reconcile_kv(
3470 &self.len_ptrs,
3471 &self.saved_lens,
3472 &self.forced_acc,
3473 &self.valid,
3474 0,
3475 self.split,
3476 )?;
3477 for il in 0..self.split {
3478 if let Some(recur) = cache.recur[il].as_mut() {
3479 let conv = snapshot.conv[il]
3480 .as_ref()
3481 .ok_or("optipipe stage0 snapshot missing conv state")?;
3482 let ssm = snapshot.ssm[il]
3483 .as_ref()
3484 .ok_or("optipipe stage0 snapshot missing ssm state")?;
3485 e0.spec_fork_restore_f32(conv, &mut recur.conv_state, &self.valid)?;
3486 e0.spec_fork_restore_f32(ssm, &mut recur.ssm_state, &self.valid)?;
3487 }
3488 }
3489 e0.spec_fork_restore_f32(&seed.h_seed, h_seed, &self.valid)?;
3490 e0.spec_fork_restore_f32(&seed.fill_prev, fill_prev, &self.valid)?;
3491 }
3492
3493 if keep {
3494 OPTI_FORK_HITS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
3495 return Ok(());
3496 }
3497
3498 for il in 0..self.split {
3499 if let (Some(kv), Some(saved)) = (cache.kv[il].as_mut(), snapshot.kv_len[il]) {
3500 kv.len = saved;
3501 }
3502 }
3503 scratch.set_len(e, seed.scratch_len)?;
3504 // Targeted E_restart: publish only stage 0's reconcile to the caller, then bound the
3505 // forced diagnostic so the retained number is the actual miss cost, not enqueue time.
3506 let caller = e.stream();
3507 self.rt.publish_to(0, &caller)?;
3508 caller.synchronize()?;
3509 let miss_ms = miss_started.elapsed().as_secs_f64() * 1e3;
3510 eprintln!(
3511 "[opti-fork-reconcile] generation={} slot={} miss_ms={miss_ms:.3}",
3512 generation.id, generation.slot,
3513 );
3514 OPTI_FORK_MISSES.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
3515 Ok(())
3516 }
3517
3518 fn retire(&mut self, generation: OptiForkGeneration) -> Result<(), Box<dyn std::error::Error>> {
3519 self.generations.retire(generation)
3520 }
3521}
3522
3523fn rewind_tp_kv_verified_prefix(
3524 tp_kv: &mut [Option<crate::tp::ResidentTpKvCache>],
3525 saved_lens: &[Option<usize>],
3526 accepted: usize,
3527) -> Result<(), Box<dyn std::error::Error>> {
3528 if tp_kv.len() != saved_lens.len() {
3529 return Err("spec TP KV snapshot shape mismatch".into());
3530 }
3531 for (layer, (cache, saved)) in tp_kv.iter_mut().zip(saved_lens).enumerate() {
3532 match (cache.as_mut(), *saved) {
3533 (Some(cache), Some(saved)) => {
3534 let target = saved
3535 .checked_add(accepted)
3536 .ok_or("spec TP KV committed length overflow")?;
3537 cache.rewind_to(target)?;
3538 }
3539 (None, None) => {}
3540 _ => {
3541 return Err(
3542 format!("spec TP KV layer {layer} changed shape since its snapshot").into(),
3543 );
3544 }
3545 }
3546 }
3547 Ok(())
3548}
3549
3550impl HybridModel {
3551 fn mtp_head_count(&self) -> usize {
3552 usize::from(self.mtp.is_some()) + self.mtp_extra.len()
3553 }
3554
3555 fn mtp_head_at(&self, index: usize) -> &MtpHead {
3556 if index == 0 {
3557 self.mtp.as_ref().expect("MTP head 0 is unavailable")
3558 } else {
3559 &self.mtp_extra[index - 1]
3560 }
3561 }
3562
3563 fn new_mtp_scratch(
3564 &self,
3565 e: &Engine,
3566 cap: usize,
3567 ) -> Result<MtpScratch, Box<dyn std::error::Error>> {
3568 let mut scratch = MtpScratch::new(
3569 e,
3570 &self.cfg,
3571 &self.plan,
3572 cap,
3573 self.mtp.as_ref().and_then(|head| head.geom.as_ref()),
3574 )?;
3575 for head in &self.mtp_extra {
3576 scratch.push_plane(e, &self.cfg, &self.plan, head.geom.as_ref())?;
3577 }
3578 Ok(scratch)
3579 }
3580
3581 fn opti_graph_draft_step(
3582 &self,
3583 e: &Engine,
3584 mtp: &MtpHead,
3585 dctx: &mut DraftGraphCtx,
3586 scratch: &mut MtpScratch,
3587 d_vocab: usize,
3588 ) -> Result<(u32, f32), Box<dyn std::error::Error>> {
3589 dctx.graph
3590 .as_ref()
3591 .ok_or("optipipe controller requires the greedy draft graph")?
3592 .launch()?;
3593 scratch.kv.len += 1;
3594 let idx = e.dtoh_u32_one(&dctx.g_tok)?;
3595 if (idx as usize) >= d_vocab {
3596 return Err(
3597 format!("optipipe draft argmax sentinel 0x{idx:08x} >= d_vocab {d_vocab}").into(),
3598 );
3599 }
3600 let probability = e.dtoh(&dctx.g_p)?[0];
3601 if !probability.is_finite() || !(0.0..=1.0).contains(&probability) {
3602 return Err(format!("optipipe draft probability is invalid: {probability}").into());
3603 }
3604 let token = match &mtp.d2t {
3605 Some(map) => map[idx as usize],
3606 None => idx,
3607 };
3608 if token != idx {
3609 e.set_u32_one(&mut dctx.g_tok, token)?;
3610 }
3611 Ok((token, probability))
3612 }
3613
3614 #[allow(clippy::too_many_arguments)]
3615 fn opti_controller_draft_step(
3616 &self,
3617 e: &Engine,
3618 mtp: &MtpHead,
3619 dctx: &mut DraftGraphCtx,
3620 scratch: &mut MtpScratch,
3621 d_vocab: usize,
3622 eager_state: &mut Option<(u32, CudaSlice<f32>)>,
3623 eager_pos: usize,
3624 embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
3625 ) -> Result<(u32, f32), Box<dyn std::error::Error>> {
3626 if dctx.graph.is_some() {
3627 return self.opti_graph_draft_step(e, mtp, dctx, scratch, d_vocab);
3628 }
3629 let (input_token, input_seed) = eager_state
3630 .take()
3631 .ok_or("optipipe eager continuation seed is unavailable")?;
3632 let (logits, next_seed) = self.mtp_head_forward_dev(
3633 e,
3634 mtp,
3635 input_token,
3636 &input_seed,
3637 scratch,
3638 eager_pos,
3639 embd_dev,
3640 None,
3641 )?;
3642 let token_d = e.argmax_token_device(&logits, d_vocab)?;
3643 let idx = e.dtoh_u32_one(&token_d)?;
3644 if (idx as usize) >= d_vocab {
3645 return Err(format!(
3646 "optipipe eager draft argmax sentinel 0x{idx:08x} >= d_vocab {d_vocab}"
3647 )
3648 .into());
3649 }
3650 let probability_d = e.prob_of_token_device(&logits, &token_d, d_vocab)?;
3651 let probability = e.dtoh(&probability_d)?[0];
3652 if !probability.is_finite() || !(0.0..=1.0).contains(&probability) {
3653 return Err(
3654 format!("optipipe eager draft probability is invalid: {probability}").into(),
3655 );
3656 }
3657 let token = match &mtp.d2t {
3658 Some(map) => map[idx as usize],
3659 None => idx,
3660 };
3661 *eager_state = Some((token, next_seed));
3662 Ok((token, probability))
3663 }
3664
3665 /// NextN head forward for ONE draft token (§A ops 1-13, T=1).
3666 /// Inputs: `e_tok` = the token to predict FROM (last committed / previous draft); `h_seed` =
3667 /// the trunk's pre-output_norm hidden of that token (§A op 2 input). `mtp_pos` = absolute
3668 /// position of the token being predicted from. Returns (draft_logits[n_vocab] host, h_nextn dev).
3669 /// `h_nextn` (§A op 10) becomes `h_seed` for the next autoregressive draft step.
3670 /// Device-resident: returns draft logits ON DEVICE (no [n_vocab] dtoh). The greedy draft
3671 /// loop only needs argmax — paired with `argmax_token_device` this cuts the ~600KB logits
3672 /// transfer + host argmax per draft token from the K-token draft chain.
3673 #[allow(clippy::too_many_arguments)]
3674 fn mtp_head_forward_dev(
3675 &self,
3676 e: &Engine,
3677 mtp: &MtpHead,
3678 e_tok: u32,
3679 h_seed: &CudaSlice<f32>,
3680 scratch: &mut MtpScratch,
3681 mtp_pos: usize,
3682 embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
3683 mask: Option<(&CudaSlice<u32>, usize)>,
3684 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
3685 self.mtp_head_forward_dev_at(e, mtp, e_tok, h_seed, scratch, 0, mtp_pos, embd_dev, mask)
3686 }
3687
3688 #[allow(clippy::too_many_arguments)]
3689 fn mtp_head_forward_dev_at(
3690 &self,
3691 e: &Engine,
3692 mtp: &MtpHead,
3693 e_tok: u32,
3694 h_seed: &CudaSlice<f32>,
3695 scratch: &mut MtpScratch,
3696 scratch_index: usize,
3697 mtp_pos: usize,
3698 embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
3699 // DRAFT-SIDE GRAMMAR MASK (lane/draft-mask): (packed draft-vocab allowed set, words).
3700 // Applied to the head logits BEFORE they are returned, so every consumer (argmax,
3701 // gumbel draw, p-min prob) sees the grammar-legal row. None = unmasked (pre-lane).
3702 mask: Option<(&CudaSlice<u32>, usize)>,
3703 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
3704 // MEMRA_SPEC_ANATOMY=1 — eager-step phase timers (diagnostic only). Phase boundaries
3705 // sync the stream, so absolute time inflates; the BREAKDOWN is the signal. Cumulative
3706 // summary on stderr every 128 steps: glue (embed..attn_norm), attn, ffn, head.
3707 use std::sync::atomic::{AtomicU64, Ordering::Relaxed};
3708 static ANAT_NS: [AtomicU64; 5] = [
3709 AtomicU64::new(0),
3710 AtomicU64::new(0),
3711 AtomicU64::new(0),
3712 AtomicU64::new(0),
3713 AtomicU64::new(0),
3714 ];
3715 static ANAT_STEPS: AtomicU64 = AtomicU64::new(0);
3716 let anat = {
3717 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
3718 *ON.get_or_init(|| std::env::var("MEMRA_SPEC_ANATOMY").as_deref() == Ok("1"))
3719 };
3720 if anat {
3721 e.stream().synchronize()?; // drain prior queue so phase 0 starts clean
3722 }
3723 let t_all = std::time::Instant::now();
3724 let mut t_ph = std::time::Instant::now();
3725 let mut anat_mark = |i: usize,
3726 e: &Engine,
3727 t: &mut std::time::Instant|
3728 -> Result<(), Box<dyn std::error::Error>> {
3729 if anat {
3730 e.stream().synchronize()?;
3731 ANAT_NS[i].fetch_add(t.elapsed().as_nanos() as u64, Relaxed);
3732 *t = std::time::Instant::now();
3733 }
3734 Ok(())
3735 };
3736 let cfg = &self.cfg;
3737 let n_embd = cfg.n_embd as usize;
3738 // Distilled-student geometry: the block runs at the INNER width `di` (eh_proj out /
3739 // attn / ffn); the n_embd interface (embed, norms in, carrier out, head in) is unchanged.
3740 let di = mtp.geom.as_ref().map(|g| g.d_inner).unwrap_or(n_embd);
3741 let eps = cfg.rms_eps;
3742 let pos_d = e.htod_i32(&[mtp_pos as i32])?;
3743
3744 // op A: a resident table transfers one 4B token id. The exact host-row capacity path
3745 // expands this one row on CPU and transfers n_embd f32 values instead.
3746 let e_emb = match embd_dev {
3747 Some((g, qt, rb)) => e.embed_gather_device_t(g, &[e_tok], n_embd, qt, rb)?,
3748 None => e.htod(&self.embd.gather(n_embd, &[e_tok]))?,
3749 };
3750
3751 // op 1/2: e_norm = RMSNorm(e, enorm); h_norm = RMSNorm(h_seed, hnorm)
3752 let mut e_norm = e.zeros(n_embd)?;
3753 e.rms_norm(&e_emb, mtp.enorm.float_data(), &mut e_norm, n_embd, 1, eps)?;
3754 let mut h_norm = e.zeros(n_embd)?;
3755 e.rms_norm(h_seed, mtp.hnorm.float_data(), &mut h_norm, n_embd, 1, eps)?;
3756
3757 // op 3: concat = [e_norm ; h_norm] -> [2*n_embd], e_norm in [0,n_embd), h_norm in [n_embd,2n_embd)
3758 let mut concat = e.zeros(2 * n_embd)?;
3759 e.copy_into(&mut concat, 0, &e_norm, n_embd)?;
3760 e.copy_into(&mut concat, n_embd, &h_norm, n_embd)?;
3761
3762 // op 4: inpSA = eh_proj @ concat (eh_proj [2*n_embd, n_embd]) -> [n_embd]
3763 let inp_sa = e.matmul(&mtp.eh_proj, &concat, 1)?;
3764
3765 // op 5: a_norm = RMSNorm(inpSA, attn_norm)
3766 let mut a_norm = e.zeros(di)?;
3767 e.rms_norm(&inp_sa, mtp.attn_norm.float_data(), &mut a_norm, di, 1, eps)?;
3768 anat_mark(0, e, &mut t_ph)?;
3769
3770 // op 6: attention on the scratch KV. SAME dc launcher as the graph path (bucket_max =
3771 // scratch.cap, length from the device len_d) so eager drafts match graph drafts
3772 // bit-for-bit at any t_kv (the parity gate). Host len mirrored here (the dc append
3773 // advances only the device counter).
3774 let attn_out = match (&mtp.mixer, mtp.step35.as_ref()) {
3775 // step35 MTP block: PER-LAYER geometry + a separate head-wise gate + an SWA window,
3776 // none of which the dc launcher can express (see `mtp_step35_attn`). Host-len arm.
3777 // Advances BOTH the host len and the device counter itself (unlike the dc arm,
3778 // whose host-side mirror the caller does).
3779 (Mixer::Full(fa), Some(g)) => {
3780 self.mtp_step35_attn(e, fa, g, &a_norm, &pos_d, scratch, scratch_index)?
3781 }
3782 (Mixer::Full(fa), None) => {
3783 let out = self.mtp_full_attn_dc(
3784 e,
3785 fa,
3786 &a_norm,
3787 &pos_d,
3788 scratch,
3789 scratch_index,
3790 mtp.geom.as_ref(),
3791 )?;
3792 scratch.plane_mut(scratch_index).0.len += 1;
3793 out
3794 }
3795 (Mixer::Linear(_), _) => {
3796 panic!("MTP block is full-attn in qwen35; linear MTP not supported")
3797 }
3798 (Mixer::Mla(_), _) => crate::hybrid::mla_forward_unimplemented(),
3799 };
3800 anat_mark(1, e, &mut t_ph)?;
3801
3802 // op 7: x1 = inpSA + attn_out
3803 let mut x1 = e.zeros(di)?;
3804 e.add(&inp_sa, &attn_out, &mut x1, di)?;
3805
3806 // op 8: z = RMSNorm(x1, post_attn_norm) (pre-FFN norm)
3807 let mut z = e.zeros(di)?;
3808 e.rms_norm(&x1, mtp.post_attn_norm.float_data(), &mut z, di, 1, eps)?;
3809
3810 // op 9: FFN (Dense or MoE) — same as the trunk decode FFN
3811 let ffn_out = match &mtp.ffn {
3812 crate::hybrid::Ffn::Dense {
3813 ffn_gate,
3814 ffn_up,
3815 ffn_down,
3816 } => {
3817 let n_ff = ffn_gate.out_features();
3818 let (gate, up) = if e.uses_q8_1_fast(ffn_gate) && e.uses_q8_1_fast(ffn_up) {
3819 let (zq, zd) = e.quantize_q8_1(&z, 1, di)?;
3820 (
3821 e.matmul_pre(ffn_gate, &zq, &zd, &z, 1)?,
3822 e.matmul_pre(ffn_up, &zq, &zd, &z, 1)?,
3823 )
3824 } else {
3825 (e.matmul(ffn_gate, &z, 1)?, e.matmul(ffn_up, &z, 1)?)
3826 };
3827 let mut act = e.zeros(n_ff)?;
3828 // step35: a DENSE FFN reads the per-layer SHEXP clamp (upstream's one `build_ffn`
3829 // serves the dense MLP and the shared expert off `swiglu_clamp_shexp` —
3830 // llama-graph.cpp:1751), resolved for the MTP block's OWN index. Every other arch
3831 // passes None, which is `ffn_act`'s dispatch verbatim.
3832 Self::ffn_act_lim(
3833 e,
3834 &self.cfg,
3835 &gate,
3836 &up,
3837 1.0,
3838 1.0,
3839 mtp.step35.as_ref().and_then(|s| s.clamp_shexp),
3840 &mut act,
3841 n_ff,
3842 )?;
3843 e.matmul(ffn_down, &act, 1)?
3844 }
3845 // MTP head is a distinct block — key its experts under a separate layer index (u16::MAX)
3846 // so they never alias trunk layer 0's cache keys.
3847 crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il(e, m, &z, 1, u16::MAX)?,
3848 };
3849 anat_mark(2, e, &mut t_ph)?;
3850
3851 // op 10: h_nextn = x1 + ffn_out (at di)
3852 let mut h_inner = e.zeros(di)?;
3853 e.add(&x1, &ffn_out, &mut h_inner, di)?;
3854
3855 // op 10.5 (student): up-project the inner hidden back to n_embd — training semantics:
3856 // the chain carrier AND the head input are out_up(h_inner) (pre-final-norm).
3857 let h_nextn = match mtp.geom.as_ref() {
3858 Some(g) => e.matmul(&g.out_up, &h_inner, 1)?,
3859 None => h_inner,
3860 };
3861
3862 // op 11: final = RMSNorm(h_nextn, shared_head_norm OR output_norm)
3863 let final_norm = mtp.shared_head_norm.as_ref().unwrap_or(&self.output_norm);
3864 let mut final_h = e.zeros(n_embd)?;
3865 e.rms_norm(
3866 &h_nextn,
3867 final_norm.float_data(),
3868 &mut final_h,
3869 n_embd,
3870 1,
3871 eps,
3872 )?;
3873
3874 // op 12: draft_logits = (shared_head_head OR output) @ final — stays ON DEVICE.
3875 let head = mtp.shared_head_head.as_ref().unwrap_or(&self.output);
3876 let mut logits = e.matmul(head, &final_h, 1)?;
3877 // op 12b (lane/draft-mask): grammar mask over the DRAFT vocab, applied here so the
3878 // caller's argmax / gumbel draw / p-min prob all read the grammar-legal row.
3879 if let Some((mask_d, mw)) = mask {
3880 let d_vocab = head.out_features();
3881 e.mask_logits_col(&mut logits, mask_d, 0, d_vocab, mw)?;
3882 }
3883 anat_mark(3, e, &mut t_ph)?;
3884 if anat {
3885 ANAT_NS[4].fetch_add(t_all.elapsed().as_nanos() as u64, Relaxed);
3886 let n = ANAT_STEPS.fetch_add(1, Relaxed) + 1;
3887 if n % 128 == 0 {
3888 let us = |i: usize| ANAT_NS[i].load(Relaxed) / n / 1000;
3889 eprintln!(
3890 "[spec-anatomy] steps={n} avg us/step: glue={} attn={} ffn={} head={} total={}",
3891 us(0),
3892 us(1),
3893 us(2),
3894 us(3),
3895 us(4)
3896 );
3897 }
3898 }
3899 // Chain recurrence hand-over: pre-norm h_nextn (default) or post-norm final_h
3900 // (MEMRA_SPEC_HPOST — llama.cpp #24025's t_h_nextn is taken AFTER the head norm).
3901 Ok((logits, if spec_hpost() { final_h } else { h_nextn }))
3902 }
3903
3904 #[allow(clippy::too_many_arguments)]
3905 fn mtp_chain_forward_dev(
3906 &self,
3907 e: &Engine,
3908 tokens: &[u32],
3909 seeds: &[CudaSlice<f32>],
3910 scratch: &mut MtpScratch,
3911 committed_scratch_len: usize,
3912 embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
3913 mask: Option<(&CudaSlice<u32>, usize)>,
3914 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
3915 if tokens.is_empty() || tokens.len() != seeds.len() {
3916 return Err("multi-head MTP prefix tokens/seeds are malformed".into());
3917 }
3918 let index = mtp_chain_head_index(tokens.len() - 1, self.mtp_head_count());
3919 let head = self.mtp_head_at(index);
3920 scratch.set_plane_len(e, index, committed_scratch_len)?;
3921
3922 let mut last = None;
3923 for row in 0..tokens.len() {
3924 let is_last = row + 1 == tokens.len();
3925 last = Some(self.mtp_head_forward_dev_at(
3926 e,
3927 head,
3928 tokens[row],
3929 &seeds[row],
3930 scratch,
3931 index,
3932 committed_scratch_len + row + 1,
3933 embd_dev,
3934 if is_last { mask } else { None },
3935 )?);
3936 }
3937 Ok(last.expect("non-empty MTP prefix produced no row"))
3938 }
3939
3940 /// step35 MTP-block attention, T=1, on the scratch KV — the EAGER-ONLY twin of
3941 /// `mtp_full_attn_dc`. Three things force a separate arm rather than a geometry parameter on
3942 /// the dc path, and all three are properties of this arch's MTP block:
3943 ///
3944 /// 1. **The SWA window.** Block 45 is an SWA-type block (`sliding_window_pattern[45]=true`,
3945 /// window 512). Windowed decode in memra is a token-aligned VIEW OFFSET into the quantized
3946 /// cache (the gemma4 R6 / `step35_decode_attn` pattern: keys carry absolute rope and the
3947 /// mask is purely positional, so one query at `len-1` attending the last `win` rows IS the
3948 /// windowed result). `fa_decode_dc` takes the key count from a DEVICE counter and always
3949 /// starts at row 0 — it cannot express a nonzero offset, so a windowed dc arm would need a
3950 /// new kernel. That is deliberately not built here: see the CUDA-graph note below.
3951 /// 2. **Per-layer head count.** 96 q heads over 8 KV (GQA 12) at this block, vs the trunk's 64
3952 /// on its full-attn layers. The trunk cfg's `n_head` scalar is the MAX over layers, and the
3953 /// trunk ARTIFACT's per-layer arrays stop at index 44 — so the count must come from the
3954 /// resolved `Step35MtpGeom`, never from `cfg`.
3955 /// 3. **The separate head-wise gate.** `blk.45.attn_gate.weight [n_embd, 96]` produces one
3956 /// sigmoid scalar per head (broadcast over head_dim) — `attn_head_gate`, not the qwen35
3957 /// fused-into-wq `q_gate_split` form the dc arm handles.
3958 ///
3959 /// WHY EAGER-ONLY IS NOT A GAP TODAY: `mtp_head_forward_cap` refuses step35 heads
3960 /// explicitly (the SWA-window refusal below), so the graph draft never engages for step35
3961 /// models regardless of eligibility — the eager chain IS the served path. `mtp_head_forward_cap` refuses step35 explicitly rather
3962 /// than silently capturing a window-less (wrong past `win` draft rows) graph.
3963 ///
3964 /// Unlike the dc arm this advances BOTH the host `kv.len` and the device counter, so the
3965 /// caller must not mirror.
3966 fn mtp_step35_attn(
3967 &self,
3968 e: &Engine,
3969 fa: &FullAttnLayer,
3970 g: &crate::hybrid::Step35MtpGeom,
3971 h: &CudaSlice<f32>,
3972 pos_d: &CudaSlice<i32>,
3973 scratch: &mut MtpScratch,
3974 scratch_index: usize,
3975 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3976 let (nh, nkv, hd) = (g.n_head, g.n_head_kv, self.cfg.head_dim_k as usize);
3977 let eps = self.cfg.rms_eps;
3978 let scale = 1.0 / (hd as f32).sqrt(); // step35.cpp:255 kq_scale
3979 let n_embd = self.cfg.n_embd as usize;
3980 let gw = fa
3981 .attn_gate
3982 .as_ref()
3983 .ok_or("step35 MTP block is missing attn_gate.weight (head-wise attention gate)")?;
3984
3985 let (q0, k0, v0, gt) = if e.uses_q8_1_fast(&fa.wq)
3986 && e.uses_q8_1_fast(&fa.wk)
3987 && e.uses_q8_1_fast(&fa.wv)
3988 && e.uses_q8_1_fast(gw)
3989 {
3990 let (hq, hdq) = e.quantize_q8_1(h, 1, n_embd)?;
3991 let (a, b, c) = match e.matmul_q8_fused3(&fa.wq, &fa.wk, &fa.wv, &hq, &hdq)? {
3992 Some(t3) => t3,
3993 None => (
3994 e.matmul_pre(&fa.wq, &hq, &hdq, h, 1)?,
3995 e.matmul_pre(&fa.wk, &hq, &hdq, h, 1)?,
3996 e.matmul_pre(&fa.wv, &hq, &hdq, h, 1)?,
3997 ),
3998 };
3999 (a, b, c, e.matmul_pre(gw, &hq, &hdq, h, 1)?)
4000 } else {
4001 (
4002 e.matmul(&fa.wq, h, 1)?,
4003 e.matmul(&fa.wk, h, 1)?,
4004 e.matmul(&fa.wv, h, 1)?,
4005 e.matmul(gw, h, 1)?,
4006 )
4007 };
4008
4009 let mut q = e.uninit(nh * hd)?;
4010 e.rms_norm(&q0, fa.q_norm.float_data(), &mut q, hd, nh, eps)?;
4011 let mut k = e.uninit(nkv * hd)?;
4012 e.rms_norm(&k0, fa.k_norm.float_data(), &mut k, hd, nkv, eps)?;
4013 // `rope_freqs.weight` (llama3 factors) applies to the FULL-attn layers ONLY; SWA passes
4014 // null (llama-hparams / step35.cpp). Block 45 is SWA, so `ff` is None there — but read
4015 // the resolved flag, not the constant, so an all-full sibling stays correct.
4016 let ff = if g.swa {
4017 None
4018 } else {
4019 self.step35_aux.as_ref().and_then(|a| a.rope_freqs(e))
4020 };
4021 #[cfg(debug_assertions)]
4022 if let Some(ff) = ff {
4023 crate::debug_assert_tensor_stream_device(ff, &e.stream(), "mtp_step35_attn.rope_freqs");
4024 }
4025 e.rope_neox2(
4026 &mut q,
4027 &mut k,
4028 pos_d,
4029 hd,
4030 g.n_rot,
4031 nh,
4032 nkv,
4033 1,
4034 g.rope_base,
4035 1.0,
4036 ff,
4037 )?;
4038
4039 // Append at the HOST slot, then re-stamp the device counter: the eager chain has the
4040 // length on the host anyway, and the windowed view below needs it there to compute the
4041 // offset. The device counter is kept in lockstep so `mtp_kv_fill`'s `set_i32_one` and any
4042 // dc-family consumer of this scratch still agree.
4043 let (kv, scratch_cap) = scratch.plane_mut(scratch_index);
4044 assert!(
4045 kv.len < scratch_cap,
4046 "step35 MTP scratch overflow ({} >= {})",
4047 kv.len,
4048 scratch_cap
4049 );
4050 let next_len = kv.len + 1;
4051 let (off, t_kv) = if g.swa && next_len > g.window {
4052 (next_len - g.window, g.window)
4053 } else {
4054 (0, next_len)
4055 };
4056 let write_row = e.prepare_kv_append(kv, off & !31usize, 1)?;
4057 e.append_kv_quantized(
4058 &k,
4059 &v0,
4060 &mut kv.k,
4061 &mut kv.v,
4062 write_row,
4063 kv.kv_dim_k,
4064 kv.kv_dim_v,
4065 kv.k_tok_bytes,
4066 kv.v_tok_bytes,
4067 false,
4068 )?;
4069 kv.len = next_len;
4070 e.set_i32_one(&mut kv.len_d, kv.len as i32)?;
4071 // SWA view offset (see note 1). The draft chain is short (k+2 rows), but the scratch is
4072 // PERSISTENT across rounds — `mtp_kv_fill` leaves one row per committed token behind, so
4073 // `kv.len` tracks absolute position and crosses 512 in any real generation. The window is
4074 // therefore live, not theoretical.
4075 let physical = kv.physical_rows(off, off + t_kv)?;
4076 let k_view = e.view_u8_range(
4077 &kv.k,
4078 physical.start * kv.k_tok_bytes,
4079 physical.end * kv.k_tok_bytes,
4080 );
4081 let v_view = e.view_u8_range(
4082 &kv.v,
4083 physical.start * kv.v_tok_bytes,
4084 physical.end * kv.v_tok_bytes,
4085 );
4086 let mut attn = e.uninit(nh * hd)?;
4087 e.fa_decode_kvmod(
4088 &q,
4089 &k_view,
4090 &v_view,
4091 &mut attn,
4092 hd,
4093 nh,
4094 nkv,
4095 t_kv,
4096 scale,
4097 kv.k_tok_bytes,
4098 kv.v_tok_bytes,
4099 false,
4100 )?;
4101
4102 let mut ag = e.uninit(nh * hd)?;
4103 e.attn_head_gate(&attn, >, &mut ag, None, hd, nh, 1)?;
4104 Ok(e.matmul(&fa.wo, &ag, 1)?)
4105 }
4106
4107 /// MTP-block full attention, T=1, on the scratch KV (BOTH draft paths — eager and graph):
4108 /// the scratch write slot and the attention bound come from `scratch.kv.len_d` (device i32[1])
4109 /// so the launch args are FIXED across draft steps — ONE captured graph serves the whole
4110 /// chain, and replays keep seeing KV growth through the device counter (no recapture).
4111 /// Geometry contract: n_splits is sized from `scratch.cap` (the persistent capacity); splits
4112 /// whose key range lies beyond the device t_kv exit empty and the shared combine skips them
4113 /// (fa_decode_dc bit-correct-for-any-t_kv<=bucket_max contract). The eager path uses the SAME
4114 /// launcher with the SAME bucket_max -> identical dispatch -> bit-identical draft tokens (the
4115 /// graph-vs-eager parity gate). Host len is NOT advanced here (graph contract); callers mirror.
4116 fn mtp_full_attn_dc(
4117 &self,
4118 e: &Engine,
4119 fa: &FullAttnLayer,
4120 h: &CudaSlice<f32>,
4121 pos_d: &CudaSlice<i32>,
4122 scratch: &mut MtpScratch,
4123 scratch_index: usize,
4124 geom: Option<&crate::hybrid::DraftGeom>,
4125 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4126 let cfg = &self.cfg;
4127 let mtp_il = cfg.n_layer.saturating_sub(cfg.nextn_predict_layers);
4128 let geometry = cfg.full_attention_geometry_at(mtp_il);
4129 let n_head = geom.map(|g| g.n_head).unwrap_or(geometry.n_head as usize);
4130 let n_head_kv = geom
4131 .map(|g| g.n_head_kv)
4132 .unwrap_or(geometry.n_head_kv as usize);
4133 let head_dim = geometry.head_dim_k as usize;
4134 let eps = cfg.rms_eps;
4135 let scale = geometry.attention_scale();
4136 let n_embd = geom.map(|g| g.d_inner).unwrap_or(cfg.n_embd as usize);
4137 let bucket_max = scratch.plane(scratch_index).1;
4138
4139 let (qf, mut k, v) =
4140 if e.uses_q8_1_fast(&fa.wq) && e.uses_q8_1_fast(&fa.wk) && e.uses_q8_1_fast(&fa.wv) {
4141 let (hq, hd) = e.quantize_q8_1(h, 1, n_embd)?;
4142 (
4143 e.matmul_pre(&fa.wq, &hq, &hd, h, 1)?,
4144 e.matmul_pre(&fa.wk, &hq, &hd, h, 1)?,
4145 e.matmul_pre(&fa.wv, &hq, &hd, h, 1)?,
4146 )
4147 } else {
4148 (
4149 e.matmul(&fa.wq, h, 1)?,
4150 e.matmul(&fa.wk, h, 1)?,
4151 e.matmul(&fa.wv, h, 1)?,
4152 )
4153 };
4154 // M3/Hy3 have no attention output gate — wq out is exactly q; skip the split.
4155 let gated = geometry.attention_gate == memra_gguf::config::AttentionGateKind::FusedQ;
4156 let (mut q, gate) = if gated {
4157 let mut q = e.zeros(n_head * head_dim)?;
4158 let mut gate = e.zeros(n_head * head_dim)?;
4159 e.q_gate_split(&qf, &mut q, &mut gate, head_dim, n_head, 1)?;
4160 (q, Some(gate))
4161 } else {
4162 (qf, None)
4163 };
4164
4165 let mut qn = e.zeros(n_head * head_dim)?;
4166 e.rms_norm(&q, fa.q_norm.float_data(), &mut qn, head_dim, n_head, eps)?;
4167 q = qn;
4168 let mut kn = e.zeros(n_head_kv * head_dim)?;
4169 e.rms_norm(
4170 &k,
4171 fa.k_norm.float_data(),
4172 &mut kn,
4173 head_dim,
4174 n_head_kv,
4175 eps,
4176 )?;
4177 k = kn;
4178 let rope_dims = geometry.n_rot as usize;
4179 e.rope_neox(
4180 &mut q,
4181 pos_d,
4182 head_dim,
4183 rope_dims,
4184 n_head,
4185 1,
4186 geometry.rope_base,
4187 1.0,
4188 )?;
4189 e.rope_neox(
4190 &mut k,
4191 pos_d,
4192 head_dim,
4193 rope_dims,
4194 n_head_kv,
4195 1,
4196 geometry.rope_base,
4197 1.0,
4198 )?;
4199
4200 let kv = scratch.plane_mut(scratch_index).0;
4201 // append at the DEVICE slot (kv.len_d == old len), then advance the counter in-graph.
4202 e.append_kv_quantized_dc(
4203 &k,
4204 &v,
4205 &mut kv.k,
4206 &mut kv.v,
4207 &kv.len_d,
4208 kv.kv_dim_k,
4209 kv.kv_dim_v,
4210 kv.k_tok_bytes,
4211 kv.v_tok_bytes,
4212 false,
4213 )?;
4214 e.inc_seqlen(&mut kv.len_d)?;
4215 // full-buffer views (any in-round t_kv stays in range on replay); the kernel bounds the
4216 // key range from the device counter.
4217 let k_view = e.view_u8(&kv.k, kv.k.len());
4218 let v_view = e.view_u8(&kv.v, kv.v.len());
4219 let (ktb, vtb) = (kv.k_tok_bytes, kv.v_tok_bytes);
4220 let mut attn = e.zeros(n_head * head_dim)?;
4221 e.fa_decode_dc(
4222 &q, &k_view, &v_view, &mut attn, head_dim, n_head, n_head_kv, &kv.len_d, bucket_max,
4223 scale, ktb, vtb, false,
4224 )?;
4225
4226 let attn_g = match &gate {
4227 Some(gate) => {
4228 let mut gsig = e.zeros(n_head * head_dim)?;
4229 e.sigmoid(gate, &mut gsig, n_head * head_dim)?;
4230 let mut ag = e.zeros(n_head * head_dim)?;
4231 e.mul(&attn, &gsig, &mut ag, n_head * head_dim)?;
4232 ag
4233 }
4234 None => attn,
4235 };
4236 Ok(e.matmul(&fa.wo, &attn_g, 1)?)
4237 }
4238
4239 /// PERSISTENT-DRAFT-KV fill (the reference engine's "mtp_update" analogue): compute the MTP
4240 /// block's K/V for `tokens` (committed tokens at positions pos0..pos0+T) from their EXACT
4241 /// trunk hiddens `h` ([T, n_embd] token-major, pre-output_norm) and append at slots pos0.. of
4242 /// the scratch KV. K/V-ONLY — ops A/1-5 plus the K-side of op 6 (wk/wv + k_norm + rope +
4243 /// quantized append); no wq/attention/FFN/lm_head, so per-token cost ~= eh_proj + wk/wv (a
4244 /// small fraction of one trunk layer), T-batched. Rope follows the chain convention
4245 /// rope(token@p) = p+1. Runs at round boundaries OUTSIDE the captured graph in BOTH draft
4246 /// modes -> draft parity by construction. Caller must have scratch.kv.len == pos0.
4247 #[allow(clippy::too_many_arguments)]
4248 fn mtp_kv_fill_at(
4249 &self,
4250 e: &Engine,
4251 mtp: &MtpHead,
4252 tokens: &[u32],
4253 h: &CudaSlice<f32>,
4254 pos0: usize,
4255 scratch: &mut MtpScratch,
4256 scratch_index: usize,
4257 embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
4258 ) -> Result<(), Box<dyn std::error::Error>> {
4259 let cfg = &self.cfg;
4260 let n_embd = cfg.n_embd as usize;
4261 let eps = cfg.rms_eps;
4262 let t = tokens.len();
4263 let (scratch_kv, scratch_cap) = scratch.plane(scratch_index);
4264 assert_eq!(scratch_kv.len, pos0, "mtp_kv_fill: append slot mismatch");
4265 assert!(pos0 + t <= scratch_cap, "mtp_kv_fill: scratch overflow");
4266 let Mixer::Full(fa) = &mtp.mixer else {
4267 panic!("MTP block is full-attn in qwen35; linear MTP not supported")
4268 };
4269 let pos_vec: Vec<i32> = (0..t).map(|i| (pos0 + i + 1) as i32).collect();
4270 let pos_d = e.htod_i32(&pos_vec)?;
4271
4272 // ops A/1/2: embed + the two input norms, T-wide.
4273 let e_emb = match embd_dev {
4274 Some((g, qt, rb)) => e.embed_gather_device_t(g, tokens, n_embd, qt, rb)?,
4275 None => e.htod(&self.embd.gather(n_embd, tokens))?,
4276 };
4277 let mut e_norm = e.zeros(t * n_embd)?;
4278 e.rms_norm(&e_emb, mtp.enorm.float_data(), &mut e_norm, n_embd, t, eps)?;
4279 let mut h_norm = e.zeros(t * n_embd)?;
4280 e.rms_norm(h, mtp.hnorm.float_data(), &mut h_norm, n_embd, t, eps)?;
4281
4282 // op 3: per-row [e_norm ; h_norm] concat, token-major [T, 2*n_embd].
4283 let mut concat = e.zeros(t * 2 * n_embd)?;
4284 for i in 0..t {
4285 e.copy_view_into(
4286 &mut concat,
4287 i * 2 * n_embd,
4288 &e_norm.slice(i * n_embd..(i + 1) * n_embd),
4289 n_embd,
4290 )?;
4291 e.copy_view_into(
4292 &mut concat,
4293 i * 2 * n_embd + n_embd,
4294 &h_norm.slice(i * n_embd..(i + 1) * n_embd),
4295 n_embd,
4296 )?;
4297 }
4298
4299 // ops 4/5: eh_proj + attn_norm, T-wide (at the student inner width when geom is set).
4300 let di = mtp.geom.as_ref().map(|g| g.d_inner).unwrap_or(n_embd);
4301 let inp_sa = e.matmul(&mtp.eh_proj, &concat, t)?;
4302 let mut a_norm = e.zeros(t * di)?;
4303 e.rms_norm(&inp_sa, mtp.attn_norm.float_data(), &mut a_norm, di, t, eps)?;
4304
4305 // op 6 (K/V half): wk/wv + k_norm + rope + per-row quantized append. No wq/attention —
4306 // the fill only has to leave correct K/V rows behind for later chains to attend over.
4307 let n_head_kv = mtp
4308 .geom
4309 .as_ref()
4310 .map(|g| g.n_head_kv)
4311 .or(mtp.step35.as_ref().map(|s| s.n_head_kv))
4312 .unwrap_or_else(|| {
4313 let mtp_il = cfg.n_layer.saturating_sub(cfg.nextn_predict_layers);
4314 cfg.full_attention_geometry_at(mtp_il).n_head_kv as usize
4315 });
4316 let mtp_il = cfg.n_layer.saturating_sub(cfg.nextn_predict_layers);
4317 let geometry = cfg.full_attention_geometry_at(mtp_il);
4318 let head_dim = geometry.head_dim_k as usize;
4319 let mut k = e.matmul(&fa.wk, &a_norm, t)?;
4320 let v = e.matmul(&fa.wv, &a_norm, t)?;
4321 let mut kn = e.zeros(t * n_head_kv * head_dim)?;
4322 e.rms_norm(
4323 &k,
4324 fa.k_norm.float_data(),
4325 &mut kn,
4326 head_dim,
4327 n_head_kv * t,
4328 eps,
4329 )?;
4330 k = kn;
4331 // step35: rotary width AND base are per-layer, and the MTP block's values come from the
4332 // resolved `Step35MtpGeom` — NOT from `cfg.rope_dim_count`/`cfg.rope_freq_base`, which
4333 // carry the arch defaults (128 / 5e6, i.e. the FULL-attn layers' base). Getting this wrong
4334 // writes K rows the attention arm then re-derives at a different theta: correct-looking
4335 // output with dead acceptance, invisible to the exactness gates.
4336 let (rope_dims, rope_base, ff) = match mtp.step35.as_ref() {
4337 Some(s) => (
4338 s.n_rot,
4339 s.rope_base,
4340 if s.swa {
4341 None
4342 } else {
4343 self.step35_aux.as_ref().and_then(|a| a.rope_freqs(e))
4344 },
4345 ),
4346 None => (geometry.n_rot as usize, geometry.rope_base, None),
4347 };
4348 #[cfg(debug_assertions)]
4349 if let Some(ff) = ff {
4350 crate::debug_assert_tensor_stream_device(ff, &e.stream(), "mtp_kv_fill.rope_freqs");
4351 }
4352 match ff {
4353 Some(f) => e.rope_neox_ff(
4354 &mut k, &pos_d, head_dim, rope_dims, n_head_kv, t, rope_base, 1.0, f,
4355 )?,
4356 None => e.rope_neox(
4357 &mut k, &pos_d, head_dim, rope_dims, n_head_kv, t, rope_base, 1.0,
4358 )?,
4359 }
4360
4361 let kv = scratch.plane_mut(scratch_index).0;
4362 // Match the trunk prime contract: a chunk may need the aligned window immediately before
4363 // its first row, so preserve that prefix when the physical tail rebases at wrap.
4364 let retain_from = kv
4365 .ring
4366 .as_ref()
4367 .map(|ring| pos0.saturating_sub(ring.window() - 1) & !31usize)
4368 .unwrap_or(0);
4369 let write_row = e.prepare_kv_append(kv, retain_from, t)?;
4370 for i in 0..t {
4371 let k_row = k.slice(i * kv.kv_dim_k..(i + 1) * kv.kv_dim_k);
4372 let v_row = v.slice(i * kv.kv_dim_v..(i + 1) * kv.kv_dim_v);
4373 e.append_kv_quantized_view(
4374 &k_row,
4375 &v_row,
4376 &mut kv.k,
4377 &mut kv.v,
4378 write_row + i,
4379 kv.kv_dim_k,
4380 kv.kv_dim_v,
4381 kv.k_tok_bytes,
4382 kv.v_tok_bytes,
4383 false,
4384 )?;
4385 }
4386 kv.len = pos0 + t;
4387 e.set_i32_one(&mut kv.len_d, kv.len as i32)?;
4388 Ok(())
4389 }
4390
4391 #[allow(clippy::too_many_arguments)]
4392 fn mtp_kv_fill_all(
4393 &self,
4394 e: &Engine,
4395 tokens: &[u32],
4396 h: &CudaSlice<f32>,
4397 pos0: usize,
4398 scratch: &mut MtpScratch,
4399 embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
4400 ) -> Result<(), Box<dyn std::error::Error>> {
4401 debug_assert_eq!(self.mtp_head_count(), scratch.plane_count());
4402 for index in 0..self.mtp_head_count() {
4403 self.mtp_kv_fill_at(
4404 e,
4405 self.mtp_head_at(index),
4406 tokens,
4407 h,
4408 pos0,
4409 scratch,
4410 index,
4411 embd_dev,
4412 )?;
4413 }
4414 Ok(())
4415 }
4416
4417 /// CAPTURE body for the GRAPH DRAFT (stage 2 of graph-grade spec): ONE MTP head forward with
4418 /// every varying input device-resident —
4419 /// - token id from the persistent `tok_d` (the previous replay's in-graph argmax wrote it,
4420 /// so the chain feeds itself; the host reads the same 4 bytes for the draft list),
4421 /// - h_seed from the persistent `h_seed_d` (h_nextn is copied BACK into it at the end),
4422 /// - rope pos from the persistent `pos_d` counter (inc'd in-graph),
4423 /// - scratch KV slot/bound from `scratch.kv.len_d` (see mtp_full_attn_dc).
4424 /// The p-min confidence lands in the persistent `p_d` iff `with_prob` (env is fixed per run).
4425 /// Same kernels, same dispatch as the eager mtp_head_forward_dev chain -> same draft tokens
4426 /// (exactness never depends on drafts — the verify arbitrates — but acceptance parity does).
4427 /// `with_head=false` captures the HEAD-LESS twin for the pseudo-seed replay (2026-07-03):
4428 /// the pseudo pass only needs h_nextn (op 10) + the scratch append — the lm_head read
4429 /// (~1.06ms q6_K on the 9B), argmax and prob are dead weight there. h_nextn's inputs are
4430 /// untouched, so the seed value is identical; round-start resets overwrite tok_d/p_d anyway.
4431 /// `sampled_cap` = Some((ctr_d, perturb_d, q_out_d, seed, temp)) captures the SAMPLED twin
4432 /// (step 3 of the sampled-spec arc): head logits are retained in the persistent `q_out_d`
4433 /// (host D2Ds them to the round's q slot after each replay), the DEVICE event counter is
4434 /// bumped in-graph, and the argmax reads GUMBEL-PERTURBED logits — one categorical draw per
4435 /// replay, bit-identical to the eager arm's gumbel_perturb at the same (seed, sctr, temp).
4436 /// seed/temp are capture-time constants (fixed per generate call, like p_min).
4437 #[allow(clippy::too_many_arguments)]
4438 fn mtp_head_forward_cap(
4439 &self,
4440 e: &Engine,
4441 mtp: &MtpHead,
4442 tok_d: &mut CudaSlice<u32>,
4443 pos_d: &mut CudaSlice<i32>,
4444 h_seed_d: &mut CudaSlice<f32>,
4445 p_d: &mut CudaSlice<f32>,
4446 scratch: &mut MtpScratch,
4447 with_prob: bool,
4448 with_head: bool,
4449 embd_gpu: &CudaSlice<u8>,
4450 embd_qt: i32,
4451 embd_rb: usize,
4452 d_vocab: usize,
4453 sampled_cap: Option<(
4454 &mut CudaSlice<u32>,
4455 &mut CudaSlice<f32>,
4456 &mut CudaSlice<f32>,
4457 u64,
4458 f32,
4459 )>,
4460 stream_pack: Option<(&mut CudaSlice<u32>, usize, Option<&CudaSlice<u32>>)>,
4461 // DRAFT-SIDE GRAMMAR MASK (lane/draft-mask): (packed draft-vocab allowed-set buffer,
4462 // word count). Captured as ONE mask_logits_f32 node between the head matmul and the
4463 // in-graph argmax; the buffer address is baked, its CONTENTS are re-uploaded by the
4464 // host before every replay (the decode.rs graph-mask pattern). All-ones contents = a
4465 // no-op ban, so a position the grammar cannot constrain costs one pass over the row.
4466 mask_cap: Option<(&CudaSlice<u32>, usize)>,
4467 ) -> Result<(), Box<dyn std::error::Error>> {
4468 let cfg = &self.cfg;
4469 let n_embd = cfg.n_embd as usize;
4470 // step35 REFUSAL (deliberate, named): the graph body's attention is `mtp_full_attn_dc`,
4471 // whose device-counter key bound always starts at row 0 — it cannot express this block's
4472 // SWA view offset, so a captured chain would silently attend OUTSIDE the window once the
4473 // persistent scratch passes 512 rows. Nothing is lost today: `mtp_head_forward_cap`
4474 // refuses step35 heads explicitly (SWA refusal), so the eager chain
4475 // (`mtp_head_forward_dev` -> `mtp_step35_attn`) is the served path. Returning Err (not a
4476 // panic) is what the two capture sites and the round-stream capture already handle by
4477 // degrading to eager / stream-off.
4478 if mtp.step35.is_some() {
4479 return Err(
4480 "step35 has no captured draft chain (fa_decode_dc cannot express the MTP \
4481 block's SWA view offset; same root cause as the dc decode refusal) — the \
4482 eager draft chain serves this arch"
4483 .into(),
4484 );
4485 }
4486 // student inner width (see mtp_head_forward_dev) — interface dims stay n_embd.
4487 let di = mtp.geom.as_ref().map(|g| g.d_inner).unwrap_or(n_embd);
4488 let eps = cfg.rms_eps;
4489 let e_emb = e.embed_gather_device(embd_gpu, tok_d, n_embd, embd_qt, embd_rb)?;
4490 let mut e_norm = e.zeros(n_embd)?;
4491 e.rms_norm(&e_emb, mtp.enorm.float_data(), &mut e_norm, n_embd, 1, eps)?;
4492 let mut h_norm = e.zeros(n_embd)?;
4493 e.rms_norm(
4494 &*h_seed_d,
4495 mtp.hnorm.float_data(),
4496 &mut h_norm,
4497 n_embd,
4498 1,
4499 eps,
4500 )?;
4501 let mut concat = e.zeros(2 * n_embd)?;
4502 e.copy_into(&mut concat, 0, &e_norm, n_embd)?;
4503 e.copy_into(&mut concat, n_embd, &h_norm, n_embd)?;
4504 let inp_sa = e.matmul(&mtp.eh_proj, &concat, 1)?;
4505 let mut a_norm = e.zeros(di)?;
4506 e.rms_norm(&inp_sa, mtp.attn_norm.float_data(), &mut a_norm, di, 1, eps)?;
4507 let attn_out = match &mtp.mixer {
4508 Mixer::Full(fa) => {
4509 self.mtp_full_attn_dc(e, fa, &a_norm, pos_d, scratch, 0, mtp.geom.as_ref())?
4510 }
4511 Mixer::Linear(_) => {
4512 panic!("MTP block is full-attn in qwen35; linear MTP not supported")
4513 }
4514 Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
4515 };
4516 let mut x1 = e.zeros(di)?;
4517 e.add(&inp_sa, &attn_out, &mut x1, di)?;
4518 let mut z = e.zeros(di)?;
4519 e.rms_norm(&x1, mtp.post_attn_norm.float_data(), &mut z, di, 1, eps)?;
4520 let ffn_out = match &mtp.ffn {
4521 crate::hybrid::Ffn::Dense {
4522 ffn_gate,
4523 ffn_up,
4524 ffn_down,
4525 } => {
4526 let n_ff = ffn_gate.out_features();
4527 let (gate, up) = if e.uses_q8_1_fast(ffn_gate) && e.uses_q8_1_fast(ffn_up) {
4528 let (zq, zd) = e.quantize_q8_1(&z, 1, di)?;
4529 (
4530 e.matmul_pre(ffn_gate, &zq, &zd, &z, 1)?,
4531 e.matmul_pre(ffn_up, &zq, &zd, &z, 1)?,
4532 )
4533 } else {
4534 (e.matmul(ffn_gate, &z, 1)?, e.matmul(ffn_up, &z, 1)?)
4535 };
4536 let mut act = e.zeros(n_ff)?;
4537 Self::ffn_act(e, &self.cfg, &gate, &up, &mut act, n_ff)?;
4538 e.matmul(ffn_down, &act, 1)?
4539 }
4540 // ROUND-STREAM: the 35B NextN block carries a MoE FFN. With RESIDENT experts the
4541 // dev path is pure device launches (device top-k + rows kernels, ZERO-DtoH by
4542 // design) — capture-legal. Non-resident (SLRU-lock) stays rejected: the capture
4543 // error arm degrades the caller to eager/stream-off.
4544 crate::hybrid::Ffn::Moe(m) if m.dev_exps.is_some() => {
4545 self.moe_ffn_il(e, m, &z, 1, u16::MAX)?
4546 }
4547 crate::hybrid::Ffn::Moe(_) => {
4548 return Err("graph draft requires a Dense (or resident-MoE) MTP FFN".into());
4549 }
4550 };
4551 let mut h_inner = e.zeros(di)?;
4552 e.add(&x1, &ffn_out, &mut h_inner, di)?;
4553 // student: up-project back to n_embd (carrier + head input; see mtp_head_forward_dev).
4554 let h_nextn = match mtp.geom.as_ref() {
4555 Some(g) => e.matmul(&g.out_up, &h_inner, 1)?,
4556 None => h_inner,
4557 };
4558 // MEMRA_SPEC_HPOST needs final_h even head-less (it IS the next seed under that convention).
4559 let final_h = if with_head || spec_hpost() {
4560 let final_norm = mtp.shared_head_norm.as_ref().unwrap_or(&self.output_norm);
4561 let mut fh = e.zeros(n_embd)?;
4562 e.rms_norm(&h_nextn, final_norm.float_data(), &mut fh, n_embd, 1, eps)?;
4563 Some(fh)
4564 } else {
4565 None
4566 };
4567 if with_head {
4568 let head = mtp.shared_head_head.as_ref().unwrap_or(&self.output);
4569 let mut logits = e.matmul(head, final_h.as_ref().unwrap(), 1)?;
4570 // DRAFT-SIDE GRAMMAR MASK: ban the grammar-illegal draft ids IN the captured chain,
4571 // before the argmax — proposals become legal by construction. Contents-only
4572 // per-replay upload keeps the capture valid.
4573 if let Some((mask_d, mw)) = mask_cap {
4574 e.mask_logits_col(&mut logits, mask_d, 0, d_vocab, mw)?;
4575 }
4576 if let Some((ctr_d, perturb_d, q_out_d, seed, temp)) = sampled_cap {
4577 // SAMPLED chain: retain q (raw head logits -> persistent q_out_d; the matmul's
4578 // own buffer is pool-recycled after the capture body returns, so it can't be the
4579 // retention target), bump the device event counter, gumbel-perturb reading it,
4580 // and argmax the PERTURBED logits into tok_d — the in-graph categorical draw.
4581 e.copy_into(q_out_d, 0, &logits, d_vocab)?;
4582 e.sctr_inc(ctr_d)?;
4583 e.gumbel_perturb_ctr(&logits, perturb_d, d_vocab, seed, ctr_d, temp)?;
4584 e.argmax_token_device_into(perturb_d, tok_d, d_vocab)?;
4585 // p-min prob = the head's RAW softmax confidence in the SAMPLED pick — same
4586 // semantics as the eager sampled arm's prob_of_token_device(dl_d, tok_d).
4587 if with_prob {
4588 e.prob_of_token_device_into(&logits, tok_d, p_d, d_vocab)?;
4589 }
4590 } else {
4591 // draft token -> persistent tok_d (next replay's embed reads it; host reads the 4 bytes).
4592 e.argmax_token_device_into(&logits, tok_d, d_vocab)?;
4593 // p-min under a draft mask reads the MASKED row: confidence relative to the
4594 // grammar-LEGAL alternatives (illegal ids leave the softmax denominator), which
4595 // is the right semantics for "does the drafter know what comes next here" and
4596 // the same row the pick came from. Draft-quality only — verify arbitrates.
4597 if with_prob {
4598 e.prob_of_token_device_into(&logits, tok_d, p_d, d_vocab)?;
4599 }
4600 }
4601 }
4602 // ROUND-STREAM K-chain: pack (tok, p) into slot j, then remap tok through d2t so the
4603 // NEXT chained body's embed reads the TARGET id — zero host involvement per step.
4604 if let Some((out, slot, d2t)) = stream_pack {
4605 e.pack_tok_p(tok_d, p_d, out, slot)?;
4606 if let Some(map) = d2t {
4607 e.tok_map_u32(tok_d, map)?;
4608 }
4609 }
4610 // Next draft step's h_seed: pre-norm h_nextn (default) or post-norm final_h (HPOST).
4611 if spec_hpost() {
4612 e.copy_into(h_seed_d, 0, final_h.as_ref().unwrap(), n_embd)?;
4613 } else {
4614 e.copy_into(h_seed_d, 0, &h_nextn, n_embd)?;
4615 }
4616 // advance the draft rope position in-graph.
4617 e.inc_seqlen(pos_d)?;
4618 Ok(())
4619 }
4620
4621 /// Batched target verify forward over `tokens` at positions `pos0..pos0+T` (§D.3, T=K+1).
4622 /// Returns ALL T logit columns (host f32, [T*n_vocab]); appends T cols to every full-attn KV
4623 /// and advances every linear-attn recur state by T steps (the recur steps are SEQUENTIAL T=1).
4624 /// Advances `cache.pos` by T.
4625 pub fn decode_step_t(
4626 &self,
4627 e: &Engine,
4628 tokens: &[u32],
4629 pos0: usize,
4630 cache: &mut Cache,
4631 ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
4632 if self.is_gemma4_e4b() {
4633 return Ok(self.gemma4_e4b_decode_step_t_h(e, tokens, pos0, cache)?.0);
4634 }
4635 if self.gemma_batch_program() {
4636 return self.gemma4_decode_step_t(e, tokens, pos0, cache);
4637 }
4638 Ok(self.decode_step_t_h(e, tokens, pos0, cache)?.0)
4639 }
4640
4641 /// Like `decode_step_t` but ALSO returns the LAST column's pre-output_norm hidden (h_seed for
4642 /// the next draft round). This lets partial-accept replay run as ONE batched T=(n_acc+1) forward
4643 /// (single weight read) instead of n_acc+1 separate T=1 decode_steps (n_acc+1 weight reads).
4644 /// At batch=1 decode is bandwidth-bound, so batching the replay is THE MTP profitability lever.
4645 pub fn decode_step_t_h(
4646 &self,
4647 e: &Engine,
4648 tokens: &[u32],
4649 pos0: usize,
4650 cache: &mut Cache,
4651 ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4652 self.decode_step_t_h_emb(e, tokens, pos0, cache, None)
4653 }
4654
4655 /// Like `decode_step_t_h` with an optional RESIDENT embed table (spec hot loop): device
4656 /// gather instead of host dequant + [T, n_embd] f32 htod. Bit-identical rows.
4657 pub fn decode_step_t_h_emb(
4658 &self,
4659 e: &Engine,
4660 tokens: &[u32],
4661 pos0: usize,
4662 cache: &mut Cache,
4663 embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
4664 ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4665 let (logits_d, h_seed) = self.decode_step_t_h_emb_dev(e, tokens, pos0, cache, embd_dev)?;
4666 Ok((e.dtoh(&logits_d)?, h_seed))
4667 }
4668
4669 /// DEVICE-LOGITS verify forward (spec device-argmax lever): identical kernel chain to
4670 /// `decode_step_t_h_emb` but returns the [T, n_vocab] logits ON DEVICE — the accept walk
4671 /// argmaxes each column on-device and reads back ONE [T] u32 instead of dtoh'ing the full
4672 /// T x n_vocab f32 block (~1-4 MB + T host argmaxes, every round). Kernel dispatch is
4673 /// UNCHANGED (same decode-exact kernels); only the post-logits transfer moves.
4674 pub fn decode_step_t_h_emb_dev(
4675 &self,
4676 e: &Engine,
4677 tokens: &[u32],
4678 pos0: usize,
4679 cache: &mut Cache,
4680 embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
4681 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4682 let n_embd = self.cfg.n_embd as usize;
4683 let t = tokens.len();
4684 let (logits, x) = self.decode_step_t_core(e, tokens, pos0, cache, embd_dev, None)?;
4685 // h_seed for the next round = LAST column's pre-output_norm hidden ([n_embd]).
4686 let mut hs = vbuf(e, n_embd)?; // fully written by copy_view_into below
4687 e.copy_view_into(&mut hs, 0, &x.slice((t - 1) * n_embd..t * n_embd), n_embd)?;
4688 Ok((logits, hs))
4689 }
4690
4691 /// CORE verify forward: the `decode_step_t_h_emb_dev` kernel chain, returning the FULL
4692 /// pre-output_norm hidden stack x ([T, n_embd], any column extractable) and optionally
4693 /// filling a `VerifyCkpt` (retained per-layer state-rebuild inputs) for the REPLAY-FREE
4694 /// partial accept. `ckpt: None` => byte-for-byte the old behavior (the ckpt writes are pure
4695 /// retains/copies — they never change what any kernel computes).
4696 fn decode_step_t_core(
4697 &self,
4698 e: &Engine,
4699 tokens: &[u32],
4700 pos0: usize,
4701 cache: &mut Cache,
4702 embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
4703 mut ckpt: Option<&mut VerifyCkpt>,
4704 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4705 self.decode_step_t_core_stream(
4706 e,
4707 tokens,
4708 pos0,
4709 cache,
4710 embd_dev,
4711 ckpt.take(),
4712 None,
4713 None,
4714 None,
4715 None,
4716 )
4717 }
4718
4719 /// [`Self::decode_step_t_core`] with the MTP route's verify-graph pool armed
4720 /// (`MEMRA_SPEC_VERIFY_GRAPH`). `graphs: None` reproduces `decode_step_t_core`
4721 /// argument-for-argument, so the eager walk stays the byte-identical fallback.
4722 fn decode_step_t_core_vg(
4723 &self,
4724 e: &Engine,
4725 tokens: &[u32],
4726 pos0: usize,
4727 cache: &mut Cache,
4728 embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
4729 mut ckpt: Option<&mut VerifyCkpt>,
4730 graphs: Option<&mut DsparkVerifyGraphs>,
4731 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4732 self.decode_step_t_core_stream(
4733 e,
4734 tokens,
4735 pos0,
4736 cache,
4737 embd_dev,
4738 ckpt.take(),
4739 None,
4740 None,
4741 None,
4742 graphs,
4743 )
4744 }
4745
4746 /// Increment-0 two-session PP seam: release the peer after this lane's stage-0 boundary TX.
4747 /// The two independent sessions keep their own cache/checkpoint state; only issue order moves.
4748 fn decode_step_t_core_pipelined(
4749 &self,
4750 e: &Engine,
4751 tokens: &[u32],
4752 pos0: usize,
4753 cache: &mut Cache,
4754 embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
4755 mut ckpt: Option<&mut VerifyCkpt>,
4756 pipe: &SpecPipeLane,
4757 round: usize,
4758 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4759 let fence = crate::pp::pp_cuts(self.layers.len())
4760 .ok_or("two-session speculative pipeline requires a PP stage cut")?;
4761 if crate::pp::pp2_streams_off() || !crate::pp::spec_pp_on() {
4762 return Err("two-session speculative pipeline requires the PP verify split".into());
4763 }
4764 let interval_fence = pipe.stage0_begin(round)?;
4765 let ticket = self.verify_stage0_issue(
4766 e,
4767 tokens,
4768 pos0,
4769 cache,
4770 embd_dev,
4771 ckpt.as_deref_mut(),
4772 None,
4773 &fence,
4774 Some(interval_fence),
4775 pipe.trace(round),
4776 )?;
4777 pipe.stage0_end(round);
4778 pipe.stage1_begin(round)?;
4779 let result = self.verify_stage1_finish(e, ticket, cache, ckpt, None, &fence, true)?;
4780 pipe.verify_end(round);
4781 Ok(result)
4782 }
4783
4784 /// ROUND-STREAM stage (c) 4: `stream` = (device verify tokens [t], device pos counter) —
4785 /// when Some, rope positions come from pos_iota over the counter, the embed gathers the
4786 /// device tokens, and full_attn_verify routes appends/FA through the _dc twins reading the
4787 /// SAME counter (every layer's kvl.len == cache.pos, one counter drives all three). The
4788 /// host `tokens`/`pos0` args still size buffers (t is FIXED K+1 in stream mode).
4789 /// `vtok_dev` (engine-bundle slice 2): device verify tokens for the EMBED only —
4790 /// unlike `stream` mode it changes nothing else (host pos iota, host-len KV appends).
4791 /// `tokens` then only sizes buffers (the dummy-slice pattern the round-stream arm uses).
4792 #[allow(clippy::too_many_arguments)]
4793 fn decode_step_t_core_stream(
4794 &self,
4795 e: &Engine,
4796 tokens: &[u32],
4797 pos0: usize,
4798 cache: &mut Cache,
4799 embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
4800 mut ckpt: Option<&mut VerifyCkpt>,
4801 stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
4802 pp_pipe: Option<bool>,
4803 vtok_dev: Option<&CudaSlice<u32>>,
4804 graphs: Option<&mut DsparkVerifyGraphs>,
4805 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4806 // PP DOOR (lane/pp2-spec 2026-08-06): the verify trunk now takes its OWN stage split,
4807 // exactly as the eager and batched steps do. This is the single funnel every verify
4808 // forward reaches (decode_step_t / _h / _h_emb / _h_emb_dev / _core all land here), so
4809 // wiring it here wires the whole spec surface — the draft/accept/commit machinery above
4810 // is untouched.
4811 //
4812 // History: pp2-hardening (2026-08-06) made this funnel FAIL CLOSED, because its trunk
4813 // walk was unsplit on one stream and a sharded cross-device placement peer-read every
4814 // remote layer's weights on every spec round (measured 13.9-28x on the batched twin).
4815 // The refusal below survives to cover the residue — MEMRA_SPEC_PP=0, MEMRA_PP_STREAMS=0,
4816 // or a placement whose PpNRt fails to build — so a config that would still walk the
4817 // whole trunk on one stream refuses instead of regressing 28x.
4818 if let Some(fence) = crate::pp::pp_cuts(self.layers.len()) {
4819 if !crate::pp::pp2_streams_off() && crate::pp::spec_pp_on() {
4820 if vtok_dev.is_some() {
4821 return Err(
4822 "device-token dspark verify (slice-2 deferred readback) has no PP \
4823 stage-split arm; set MEMRA_DSPARK_DEFER_READBACK=0 or run the dspark \
4824 route on one device"
4825 .into(),
4826 );
4827 }
4828 return self.decode_step_t_core_ppn(
4829 e,
4830 tokens,
4831 pos0,
4832 cache,
4833 embd_dev,
4834 ckpt.take(),
4835 stream,
4836 &fence,
4837 pp_pipe,
4838 );
4839 }
4840 }
4841 crate::pp::refuse_unsplit_if_remote(
4842 "decode_step_t (spec verify)",
4843 "drop MEMRA_SPEC_PP=0 / MEMRA_PP_STREAMS=0 so the verify trunk takes its OWN stage \
4844 split (decode_step_t_core_ppn); or run spec on one device",
4845 )?;
4846 let cfg = &self.cfg;
4847 let n_embd = cfg.n_embd as usize;
4848 let eps = cfg.rms_eps;
4849 let t = tokens.len();
4850 let pos_d = match stream {
4851 Some((_, ctr)) => {
4852 let mut p = e.alloc_uninit::<i32>(t)?;
4853 e.pos_iota(ctr, &mut p, t)?;
4854 p
4855 }
4856 None => {
4857 let pos_vec: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
4858 e.htod_i32(&pos_vec)?
4859 }
4860 };
4861
4862 // embed T tokens -> [T, n_embd] token-major (device gather on the spec hot loop)
4863 let x = match (stream, embd_dev) {
4864 (Some((vtok, _)), Some((g, qt, rb))) => {
4865 e.embed_gather_device_td(g, vtok, t, n_embd, qt, rb)?
4866 }
4867 (None, Some((g, qt, rb))) => match vtok_dev {
4868 // slice 2: device verify tokens, same embed_gather_u32_t kernel —
4869 // bit-identical rows to the host-token arm (same per-dtype deq).
4870 Some(vt_d) => e.embed_gather_device_td(g, vt_d, t, n_embd, qt, rb)?,
4871 None => e.embed_gather_device_t(g, tokens, n_embd, qt, rb)?,
4872 },
4873 _ => {
4874 assert!(
4875 vtok_dev.is_none(),
4876 "device-token verify requires the resident embed table (embd_dev)"
4877 );
4878 e.htod(&self.embd.gather(n_embd, tokens))?
4879 }
4880 };
4881
4882 // TRUNK WALK: layers [0, n_layers) through the SAME range-scoped subgraph the PP-N
4883 // stage split calls per stage (`verify_layers`) — one code path, so the split cannot
4884 // drift from the unsplit dispatch mirroring. lane/pp2-spec 2026-08-06.
4885 let x = self.verify_layers(
4886 e,
4887 x,
4888 0,
4889 self.layers.len(),
4890 &pos_d,
4891 pos0,
4892 t,
4893 cache,
4894 ckpt.take(),
4895 stream,
4896 graphs,
4897 )?;
4898
4899 let mut hn = vbuf(e, t * n_embd)?;
4900 // Stage-A door: with the serving-class row-outer verify walk, the TAIL must be the
4901 // t=1 decode program per row too (rms_norm t=1 + the single-row bf16 head — the
4902 // split head's concat is receipted bit-identical to it). The batched cuBLASLt head
4903 // is a different ULP class and flips near-tie argmaxes off the greedy tape.
4904 let eager_tail = self.sliding_gated_moe_batch_program()
4905 && std::env::var("MEMRA_SPEC_VERIFY_EAGER").as_deref() == Ok("1");
4906 if eager_tail {
4907 let n_vocab = self.cfg.n_vocab as usize;
4908 let mut logits = vbuf(e, t * n_vocab)?;
4909 for r in 0..t {
4910 let mut row = e.uninit(n_embd)?;
4911 e.dtod_copy_view(&x.slice(r * n_embd..(r + 1) * n_embd), &mut row)?;
4912 let mut hr = e.uninit(n_embd)?;
4913 e.rms_norm(&row, self.output_norm.float_data(), &mut hr, n_embd, 1, eps)?;
4914 let lr = e.matmul(&self.output, &hr, 1)?;
4915 e.dtod_copy_into(&lr, &mut logits, r * n_vocab)?;
4916 e.dtod_copy_into(&hr, &mut hn, r * n_embd)?;
4917 }
4918 if stream.is_none() {
4919 cache.pos += t;
4920 }
4921 return Ok((logits, if spec_hpost() { hn } else { x }));
4922 }
4923 let serving_head =
4924 self.sliding_gated_moe_batch_program() || self.batched_serving_numeric_class();
4925 let logits = if serving_head {
4926 // Step35 and the qwen35 family (MoE 2026-08-14 AM, dense-hybrid same day PM — the
4927 // Q3.8 bring-up reproduced the identical near-tie class on dense: eager-class verify
4928 // vs batched-class live serving, ULP drift amplified through the GDN recurrence)
4929 // serve one batched numeric class at every live width, including B=1. Keep the
4930 // verify head in that same class; other generic families retain the decode-exact
4931 // head that their run-spec contract pins.
4932 e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
4933 e.matmul(&self.output, &hn, t)?
4934 } else {
4935 e.rms_norm_decode(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
4936 e.matmul_decode_exact(&self.output, &hn, t)?
4937 };
4938 // stream: the device pos counter owns position; host mirror reconciles at drain.
4939 if stream.is_none() {
4940 cache.pos += t;
4941 }
4942 // Hidden stack for seeds/refresh-fills: pre-norm x (default) or post-norm hn (HPOST).
4943 Ok((logits, if spec_hpost() { hn } else { x }))
4944 }
4945
4946 /// THE VERIFY TRUNK OVER PP-N (lane/pp2-spec 2026-08-06): `decode_step_t_core_stream`'s walk
4947 /// as N stage subgraphs, each on its own engine/stream (and, under `MEMRA_PP_DEVICES`, its own
4948 /// device), with a `[T, n_embd]` boundary transfer between them. T = K+1 (the verify batch),
4949 /// so this is the batched-boundary shape the pp2-batch lane's grow-only slots already handle
4950 /// (`tx(b, x, t*n_embd)`; the slot grows to the high-water T and the transport moves exactly
4951 /// the payload).
4952 ///
4953 /// Structure is `decode_step_batch_ppn`'s, which is `decode_step_h_ppn`'s. FOUR THINGS ARE
4954 /// PER-STAGE and each for a measured reason (see `decode_step_batch_ppn`'s header for the
4955 /// receipts):
4956 ///
4957 /// 1. THE ENGINE (`rt.engine(s, e)`) — `Engine` owns lazily-grown stable-pointer scratch
4958 /// (`fa_part_pool`, `fa_vf16_scratch`, `argmax_partials`) that is single-stream-safe BY
4959 /// DESIGN. Two stage streams through one Engine is the 2026-08-02 shared-scratch race
4960 /// (35% flake, nondeterministic all-logits divergence). `PpNRt::build` gives every stage
4961 /// s>0 its own Engine even on the primary device; honouring it here is what scopes the
4962 /// pools. The verify path allocates MORE of that scratch than eager decode does (FA at
4963 /// m=T, and the per-layer `GdnStash` retains), so this is load-bearing, not inherited.
4964 ///
4965 /// 2. `pos_d` — each stage uploads its OWN copy of the T rope positions on ITS stream, so the
4966 /// buffer is allocated, consumed and freed on one stream. In `stream` mode that means each
4967 /// stage runs its own `pos_iota` over the SHARED device counter (`pos_ctr`): the counter is
4968 /// read-only during the forward (the round's `inc`/`copy_add` happen outside it), so every
4969 /// stage derives the identical iota, and each stage's own output buffer is stream-local.
4970 ///
4971 /// 3. THE EMBED lives with stage 0 (`self.embd` / `embd_gpu` are host/primary-side; the
4972 /// sharded loader leaves the table with stage 0 by construction).
4973 ///
4974 /// 4. THE HEAD (`output_norm` + `output`) runs on the LAST stage — the sharded loader uploaded
4975 /// both through that stage's engine (`hybrid.rs`: `e_head = layer_engine(e, n_trunk,
4976 /// n_trunk-1)`), so reading them anywhere else is a peer read of the biggest tensor in the
4977 /// model, every round.
4978 ///
4979 /// WHAT STAYS ON THE PRIMARY, deliberately: the returned logits and hidden stack `x`. Both are
4980 /// last-stage-allocated device buffers, and every consumer (the device argmax walk, the accept
4981 /// kernels, `spec_seed_gather`, the ckpt rebuild in `commit_verified_prefix`) reads them
4982 /// through the primary context by UVA — the same read the batched serving epilogue's
4983 /// `last_logits_dev` park does. Those consumers are per-round O(T x n_vocab) and O(n_embd),
4984 /// not per-layer, so they are not the 28x class; splitting them is a separate lane.
4985 ///
4986 /// The MTP HEAD (draft side) is NOT split: it is one block, it lives wherever the loader put
4987 /// it (`load_mtp` uses the primary engine), and it is ~1-2 GB against the trunk's tens. Draft
4988 /// placement is measured, not assumed — see `research/pp2-spec-20260806`.
4989 ///
4990 /// EXACTNESS: PP-N adds ZERO deviation. Each stage runs the SAME kernels on the SAME bytes in
4991 /// the same order via the SAME `verify_layers` the unsplit body calls; the only change is
4992 /// where the residual is materialized, and the boundary is a straight f32 copy (dtod
4993 /// same-device / `cudaMemcpyPeerAsync` cross-device, no conversion). So the split MUST be
4994 /// BIT-IDENTICAL to the unsplit verify at the same T, in both placement orders. Gate:
4995 /// `decode-batch-gate --mode ppspec`. Acceptance counts are a DERIVED consequence — greedy
4996 /// accept argmaxes these logits, so bit-identical logits force identical accept walks; the
4997 /// `run-spec` K=1..8 arm checks that end-to-end rather than trusting the implication.
4998 #[allow(clippy::too_many_arguments)]
4999 fn decode_step_t_core_ppn(
5000 &self,
5001 e: &Engine,
5002 tokens: &[u32],
5003 pos0: usize,
5004 cache: &mut Cache,
5005 embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
5006 mut ckpt: Option<&mut VerifyCkpt>,
5007 stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
5008 fence: &[usize],
5009 pp_pipe: Option<bool>,
5010 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
5011 let ticket = self.verify_stage0_issue(
5012 e,
5013 tokens,
5014 pos0,
5015 cache,
5016 embd_dev,
5017 ckpt.as_deref_mut(),
5018 stream,
5019 fence,
5020 pp_pipe,
5021 None,
5022 )?;
5023 self.verify_stage1_finish(e, ticket, cache, ckpt, stream, fence, true)
5024 }
5025
5026 /// Enqueue embed, stage 0, and the first boundary TX, then return the actual boundary slot.
5027 /// The ordinary PP verify wrapper calls `verify_stage1_finish` immediately after this return.
5028 #[allow(clippy::too_many_arguments)]
5029 fn verify_stage0_issue(
5030 &self,
5031 e: &Engine,
5032 tokens: &[u32],
5033 pos0: usize,
5034 cache: &mut Cache,
5035 embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
5036 mut ckpt: Option<&mut VerifyCkpt>,
5037 stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
5038 fence: &[usize],
5039 pp_pipe: Option<bool>,
5040 trace: Option<SpecPipeTraceCtx>,
5041 ) -> Result<VerifyBoundaryTicket, Box<dyn std::error::Error>> {
5042 assert!(
5043 !self.is_gemma4_e4b() && !self.gemma_batch_program(),
5044 "decode_step_t_core_ppn covers the hybrid non-gemma4 verify trunk only \
5045 (the gemma4 arms have their own decode_step_t twins)"
5046 );
5047 if crate::pp::pp_host_bounce_active() && (stream.is_some() || embd_dev.is_some()) {
5048 return Err(
5049 "decode_step_t_core_ppn: refused with MEMRA_PP_HOST_BOUNCE=1 — the trunk \
5050 boundary itself is host-staged, but device-resident verify still peer-reads \
5051 primary-device token/position/embedding buffers from stage 0. Run plain PP \
5052 serving on this host class; spec requires local per-stage inputs first."
5053 .into(),
5054 );
5055 }
5056 let rt = crate::pp::PpNRt::get(e)?;
5057 let n_st = fence.len() - 1;
5058 assert_eq!(
5059 rt.n_stages(),
5060 n_st,
5061 "PpNRt stage count {} != fence stages {n_st}",
5062 rt.n_stages()
5063 );
5064 let n_embd = self.cfg.n_embd as usize;
5065 let t = tokens.len();
5066 let payload = t * n_embd;
5067 if pp_pipe.is_some() {
5068 assert_eq!(n_st, 2, "spec pipeline requires exactly two PP stages");
5069 }
5070 // One-shot lane diagnostic: force natural PP-2 boundaries to completion so the server
5071 // log can price stage 0, the peer hop, the RX copy, and stage 1 + head separately without
5072 // nsys. The ordinary path keeps every enqueue asynchronous. N>2 is deliberately excluded:
5073 // the report below names exactly two stages and must never imply it measured middle ones.
5074 let pp_anatomy = n_st == 2 && std::env::var("MEMRA_SPEC_PP_ANATOMY").as_deref() == Ok("1");
5075 let pp_started = std::time::Instant::now();
5076 let (mut reverse_ms, mut stage0_ms, mut tx_ms) = (0.0f64, 0.0f64, 0.0f64);
5077 // The CALLER's ambient stream, captured BEFORE any `rt.enter()` pushes a stage stream:
5078 // this body returns DEVICE-RESIDENT buffers (the device-argmax accept walk's contract),
5079 // so the exit needs the same publication the boundaries get — see `PpNRt::publish_to`.
5080 // Taken here, not at the end, because inside the last-stage scope `e.stream()` IS the
5081 // stage stream and the wait would self-order into a no-op.
5082 let caller_stream = e.stream();
5083 // #87 ROOT-CAUSE FENCE (lane/pp2spec-crash): the PREVIOUS round's stage-allocated
5084 // outputs (logits/hidden/ckpt stashes) freed stream-ordered on the STAGE streams while
5085 // the primary stream still holds queued reads of them — with event tracking elided,
5086 // nothing stops the pool from reusing those blocks for THIS round's stage allocations,
5087 // whose writes then race the queued reads (measured: 13/4096-NaN random-bits garbage in
5088 // the spec round seed; the full anatomy is on `PpNRt::fence_stages_behind`). Order every
5089 // stage stream behind the caller before enqueueing new stage work.
5090 let reverse_started = std::time::Instant::now();
5091 if pp_pipe != Some(false) {
5092 rt.fence_stages_behind(&caller_stream)?;
5093 }
5094 if pp_pipe == Some(true) {
5095 // Both session verifies must alternate boundary slots even when the ordinary
5096 // decode overlap experiment is off. Prewarm before A's stage 0 so B cannot grow
5097 // slot 1 by synchronizing the RX stream while A's stage 1 is in flight.
5098 rt.prepare_overlap_slots(0, payload)?;
5099 }
5100 if pp_anatomy {
5101 // Drain the reverse-publication dependency before timing stage 0 itself. At c=1 this
5102 // prices any primary-stream rollback/refresh tail inherited from the prior round.
5103 for s in 0..n_st {
5104 let _st = rt.enter(s);
5105 rt.engine(s, e).stream().synchronize()?;
5106 }
5107 reverse_ms = reverse_started.elapsed().as_secs_f64() * 1e3;
5108 }
5109
5110 // Per-stage rope positions: in host mode the same [T] iota each stage uploads itself; in
5111 // stream mode each stage's own `pos_iota` over the shared read-only device counter.
5112 let stage_pos = |es: &Engine| -> Result<CudaSlice<i32>, Box<dyn std::error::Error>> {
5113 match stream {
5114 Some((_, ctr)) => {
5115 let mut p = es.alloc_uninit::<i32>(t)?;
5116 es.pos_iota(ctr, &mut p, t)?;
5117 Ok(p)
5118 }
5119 None => {
5120 let pos_vec: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
5121 es.htod_i32(&pos_vec)
5122 }
5123 }
5124 };
5125
5126 // ---- STAGE 0: embed (the table lives with stage 0) + layers [0, fence[1]) + TX ----
5127 let slot = {
5128 let _st0 = rt.enter(0);
5129 let e0 = rt.engine(0, e);
5130 enqueue_spec_pipe_trace_marker(&e0.stream(), trace.as_ref(), "S0", "start", None)?;
5131 let stage0_started = std::time::Instant::now();
5132 let pos_d = stage_pos(e0)?;
5133 let x = match (stream, embd_dev) {
5134 (Some((vtok, _)), Some((g, qt, rb))) => {
5135 e0.embed_gather_device_td(g, vtok, t, n_embd, qt, rb)?
5136 }
5137 (None, Some((g, qt, rb))) => e0.embed_gather_device_t(g, tokens, n_embd, qt, rb)?,
5138 _ => e0.htod(&self.embd.gather(n_embd, tokens))?,
5139 };
5140 let x = self.verify_layers(
5141 e0,
5142 x,
5143 fence[0],
5144 fence[1],
5145 &pos_d,
5146 pos0,
5147 t,
5148 cache,
5149 ckpt.as_deref_mut(),
5150 stream,
5151 None,
5152 )?;
5153 if pp_anatomy {
5154 e0.stream().synchronize()?;
5155 stage0_ms = stage0_started.elapsed().as_secs_f64() * 1e3;
5156 }
5157 let tx_started = std::time::Instant::now();
5158 let slot = if pp_pipe.is_some() {
5159 rt.tx_pipelined(0, &x, payload)?
5160 } else {
5161 rt.tx(0, &x, payload)?
5162 };
5163 enqueue_spec_pipe_trace_marker(&e0.stream(), trace.as_ref(), "S0", "end", Some(slot))?;
5164 if pp_anatomy {
5165 e0.stream().synchronize()?;
5166 tx_ms = tx_started.elapsed().as_secs_f64() * 1e3;
5167 }
5168 slot
5169 // x + pos_d drop here: freed stream-ordered on stage-0's stream after use.
5170 };
5171
5172 Ok(VerifyBoundaryTicket {
5173 rt,
5174 caller_stream,
5175 slot,
5176 pos0,
5177 t,
5178 payload,
5179 n_st,
5180 pipelined: pp_pipe.is_some(),
5181 pp_anatomy,
5182 pp_started,
5183 reverse_ms,
5184 stage0_ms,
5185 tx_ms,
5186 trace,
5187 })
5188 }
5189
5190 /// Consume a stage-0 boundary ticket and enqueue the remaining PP stages plus the head.
5191 /// On PP-2 this is exactly stage 1; PP-N keeps its pre-existing middle-stage walk here.
5192 #[allow(clippy::too_many_arguments)]
5193 fn verify_stage1_finish(
5194 &self,
5195 e: &Engine,
5196 ticket: VerifyBoundaryTicket,
5197 cache: &mut Cache,
5198 mut ckpt: Option<&mut VerifyCkpt>,
5199 stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
5200 fence: &[usize],
5201 publish_to_caller: bool,
5202 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
5203 let VerifyBoundaryTicket {
5204 rt,
5205 caller_stream,
5206 slot,
5207 pos0,
5208 t,
5209 payload,
5210 n_st,
5211 pipelined,
5212 pp_anatomy,
5213 pp_started,
5214 reverse_ms,
5215 stage0_ms,
5216 tx_ms,
5217 trace,
5218 } = ticket;
5219 let n_embd = self.cfg.n_embd as usize;
5220 let eps = self.cfg.rms_eps;
5221 let mut slot = slot;
5222 let (mut rx_ms, mut stage1_ms) = (0.0f64, 0.0f64);
5223 let stage_pos = |es: &Engine| -> Result<CudaSlice<i32>, Box<dyn std::error::Error>> {
5224 match stream {
5225 Some((_, ctr)) => {
5226 let mut p = es.alloc_uninit::<i32>(t)?;
5227 es.pos_iota(ctr, &mut p, t)?;
5228 Ok(p)
5229 }
5230 None => {
5231 let pos_vec: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
5232 es.htod_i32(&pos_vec)
5233 }
5234 }
5235 };
5236
5237 // ---- MIDDLE STAGES: RX boundary s-1 -> range -> TX boundary s ----
5238 for s in 1..n_st - 1 {
5239 let _st = rt.enter(s);
5240 let es = rt.engine(s, e);
5241 let pos_d = stage_pos(es)?;
5242 let x = rt.rx(s - 1, slot, payload)?;
5243 let x = self.verify_layers(
5244 es,
5245 x,
5246 fence[s],
5247 fence[s + 1],
5248 &pos_d,
5249 pos0,
5250 t,
5251 cache,
5252 ckpt.as_deref_mut(),
5253 stream,
5254 None,
5255 )?;
5256 slot = if pipelined {
5257 rt.tx_pipelined(s, &x, payload)?
5258 } else {
5259 rt.tx(s, &x, payload)?
5260 };
5261 }
5262
5263 // ---- LAST STAGE: RX + final range + output_norm + lm head ----
5264 let _stl = rt.enter(n_st - 1);
5265 let el = rt.engine(n_st - 1, e);
5266 let pos_d = stage_pos(el)?;
5267 let rx_started = std::time::Instant::now();
5268 let x = rt.rx(n_st - 2, slot, payload)?;
5269 if pp_anatomy {
5270 el.stream().synchronize()?;
5271 rx_ms = rx_started.elapsed().as_secs_f64() * 1e3;
5272 }
5273 enqueue_spec_pipe_trace_marker(&el.stream(), trace.as_ref(), "S1", "start", Some(slot))?;
5274 let stage1_started = std::time::Instant::now();
5275 let x = self.verify_layers(
5276 el,
5277 x,
5278 fence[n_st - 1],
5279 fence[n_st],
5280 &pos_d,
5281 pos0,
5282 t,
5283 cache,
5284 ckpt.as_deref_mut(),
5285 stream,
5286 None,
5287 )?;
5288
5289 let mut hn = vbuf(el, payload)?;
5290 let logits = if self.sliding_gated_moe_batch_program() {
5291 // The PP Step35 serving path uses rms_norm + matmul for B=1 as well as B>1.
5292 // Verify must not switch numeric class merely because the same session speculates.
5293 el.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
5294 el.matmul(&self.output, &hn, t)?
5295 } else {
5296 el.rms_norm_decode(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
5297 el.matmul_decode_exact(&self.output, &hn, t)?
5298 };
5299 enqueue_spec_pipe_trace_marker(&el.stream(), trace.as_ref(), "S1", "end", Some(slot))?;
5300 if pp_anatomy {
5301 el.stream().synchronize()?;
5302 stage1_ms = stage1_started.elapsed().as_secs_f64() * 1e3;
5303 }
5304 // EXIT PUBLICATION: both returned buffers are still being produced on the last stage's
5305 // stream. Order the caller's stream behind that work before the buffers escape this
5306 // scope (the 2026-08-06 same-device ppspec find: without it the caller's primary-stream
5307 // consumer read unwritten logits — nondeterministic, one-device-only, and it poisoned
5308 // the following arm's KV in the same process).
5309 if publish_to_caller {
5310 rt.publish_to(n_st - 1, &caller_stream)?;
5311 }
5312 if pp_anatomy {
5313 if publish_to_caller {
5314 caller_stream.synchronize()?;
5315 }
5316 eprintln!(
5317 "[spec-pp-anatomy] t={t} reverse={reverse_ms:.3}ms stage0={stage0_ms:.3}ms \
5318 tx={tx_ms:.3}ms rx={rx_ms:.3}ms stage1-head={stage1_ms:.3}ms total={:.3}ms",
5319 pp_started.elapsed().as_secs_f64() * 1e3,
5320 );
5321 }
5322 // stream: the device pos counter owns position; host mirror reconciles at drain.
5323 if stream.is_none() {
5324 cache.pos += t;
5325 }
5326 Ok((logits, if spec_hpost() { hn } else { x }))
5327 }
5328
5329 /// Step3.5/Step3.7 verify trunk in the serving batched numeric class.
5330 ///
5331 /// `step35_decode_batch_layers` is now authoritative at every live serving width, including
5332 /// B=1 (lane/cx-b1fix). The older verify walk deliberately mirrored the eager T=1 class:
5333 /// it replayed `step35_decode_attn` per row and used the eager/decode-exact FFN dispatch.
5334 /// Those classes are individually stable, but a near-tie prompt can choose different greedy
5335 /// bytes when a request moves from batched plain serving into speculative verify. Run the
5336 /// same authoritative B=1 stage subgraph for each verify row here. Rows still advance
5337 /// layer-by-layer, so every layer sees the preceding verify rows in its attention cache while
5338 /// every norm/projection/FFN uses exactly the live serving dispatch.
5339 #[allow(clippy::too_many_arguments)]
5340 /// PRIME-BY-T-ROWS (MEMRA_PRIME_TROWS=1): prefill the prompt through the same-session
5341 /// t-row walk in 32-row chunks — every row runs the t=1 decode program bit-for-bit
5342 /// (the TOKENWISE-prime ORACLE class), so this door is exact against the exactness
5343 /// reference while replacing the host-canonical per-token prime. Requires the walk
5344 /// doors (MEMRA_SPEC_VERIFY_EAGER/TCOL); returns the prime contract trio.
5345 #[allow(clippy::type_complexity)]
5346 pub(crate) fn step35_prime_trows(
5347 &self,
5348 e: &Engine,
5349 tokens: &[u32],
5350 cache: &mut Cache,
5351 ) -> Result<Option<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
5352 {
5353 let dbg = std::env::var("MEMRA_SPEC_FA2_DEBUG").as_deref() == Ok("1");
5354 if std::env::var("MEMRA_PRIME_TROWS").as_deref() != Ok("1") {
5355 return Ok(None);
5356 }
5357 if !self.uses_sliding_gated_moe_program()
5358 || cache.pos != 0
5359 || cache.dflash_taps.is_some()
5360 || std::env::var("MEMRA_SPEC_VERIFY_EAGER").as_deref() != Ok("1")
5361 || std::env::var("MEMRA_SPEC_VERIFY_TCOL").as_deref() != Ok("1")
5362 {
5363 if dbg {
5364 eprintln!(
5365 "[prime-trows] refuse: program={} pos={} taps={} eager={:?} tcol={:?}",
5366 self.uses_sliding_gated_moe_program(),
5367 cache.pos,
5368 cache.dflash_taps.is_some(),
5369 std::env::var("MEMRA_SPEC_VERIFY_EAGER").ok(),
5370 std::env::var("MEMRA_SPEC_VERIFY_TCOL").ok()
5371 );
5372 }
5373 return Ok(None);
5374 }
5375 let n_embd = self.cfg.n_embd as usize;
5376 let n_layers = self.layers.len();
5377 let t_total = tokens.len();
5378 let Some(embd_gpu) = self.embd_gpu_try(e) else {
5379 if dbg {
5380 eprintln!("[prime-trows] refuse: no device embed table");
5381 }
5382 return Ok(None);
5383 };
5384 let embd_qtype = match self.embd.ggml_type {
5385 memra_gguf::GgmlType::BF16 => crate::QT_BF16,
5386 memra_gguf::GgmlType::Q8_0 => crate::QT_Q8_0,
5387 other => {
5388 if dbg {
5389 eprintln!("[prime-trows] refuse: embed dtype {other:?}");
5390 }
5391 return Ok(None);
5392 }
5393 };
5394 let embd_row_bytes = self.embd.raw.len() / self.cfg.n_vocab as usize;
5395 // Chunk plan: 32-row chunks; a 1-token tail folds into the previous chunk
5396 // (the walk floor is t >= 2).
5397 let mut bounds = Vec::new();
5398 let mut start = 0usize;
5399 while start < t_total {
5400 let mut end = (start + 32).min(t_total);
5401 if t_total - end == 1 {
5402 end -= 1;
5403 }
5404 bounds.push((start, end));
5405 start = end;
5406 }
5407 if bounds.iter().any(|(a, b)| b - a < 2) {
5408 return Ok(None); // degenerate short prompt keeps the ordinary prime
5409 }
5410 let mut hiddens = e.uninit(t_total * n_embd)?;
5411 let mut last: Option<CudaSlice<f32>> = None;
5412 for &(a, b) in &bounds {
5413 let tc = b - a;
5414 let tok_d = e.stream().clone_htod(&tokens[a..b])?;
5415 let x =
5416 e.embed_gather_device_td(embd_gpu, &tok_d, tc, n_embd, embd_qtype, embd_row_bytes)?;
5417 let out = self.step35_verify_batch_layers(e, x, 0, n_layers, a, tc, cache)?;
5418 e.copy_into(&mut hiddens, a * n_embd, &out, tc * n_embd)?;
5419 if b == t_total {
5420 let mut h = e.uninit(n_embd)?;
5421 e.dtod_copy_view(&out.slice((tc - 1) * n_embd..tc * n_embd), &mut h)?;
5422 last = Some(h);
5423 }
5424 }
5425 let h_seed = last.expect("last chunk produced the seed row");
5426 let mut hn = e.uninit(n_embd)?;
5427 e.rms_norm_decode(
5428 &h_seed,
5429 self.output_norm.float_data(),
5430 &mut hn,
5431 n_embd,
5432 1,
5433 self.cfg.rms_eps,
5434 )?;
5435 let logits_d = e.matmul_decode_exact(&self.output, &hn, 1)?;
5436 let logits = e.dtoh(&logits_d)?;
5437 cache.pos = t_total;
5438 Ok(Some((logits, h_seed, hiddens)))
5439 }
5440
5441 fn step35_verify_batch_layers(
5442 &self,
5443 e: &Engine,
5444 mut x: CudaSlice<f32>,
5445 lo: usize,
5446 hi: usize,
5447 pos0: usize,
5448 t: usize,
5449 cache: &mut Cache,
5450 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5451 let n_embd = self.cfg.n_embd as usize;
5452 if !self.uses_sliding_gated_moe_program() {
5453 return Err(
5454 "serving-class verify requires sliding-gated-MoE canonical operations".into(),
5455 );
5456 }
5457 // SERVING-CLASS VERIFY (MEMRA_SPEC_VERIFY_EAGER=1, step37 MTP bring-up): each verify
5458 // column rides decode_layers_eager — the EXACT t=1 program live serving runs (all TP2
5459 // doors) — row-outer, so row r's appends land before row r+1 attends: bit-equal to
5460 // plain greedy by construction. Only the unsplit full-range walk qualifies; PP splits
5461 // and the tap path keep the batch-layer class.
5462 static VE: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
5463 let eager_verify = *VE
5464 .get_or_init(|| std::env::var("MEMRA_SPEC_VERIFY_EAGER").as_deref() == Ok("1"))
5465 && lo == 0
5466 && hi == self.layers.len();
5467 if eager_verify {
5468 // T-COLUMN LAYER-OUTER WALK (MEMRA_SPEC_VERIFY_TCOL=1): per layer, one t-grid
5469 // attn norm + ONE weight-amortized QKV(+gate) over all T columns, then each
5470 // column runs the UNMODIFIED t=1 attention program via the col-select door and
5471 // the ordinary residual/FFN body. Values per column are bit-equal to the
5472 // row-outer walk: rms over the materialized residual == the fused add+norm
5473 // (kernel_check identity), the tcol kernel's per-column FP order == the t=1
5474 // kernel, and every downstream op IS the t=1 program.
5475 static TCOL: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
5476 let tcol =
5477 *TCOL.get_or_init(|| std::env::var("MEMRA_SPEC_VERIFY_TCOL").as_deref() == Ok("1"));
5478 // T > 32 (prefill-class): run the SAME walk in 32-row chunks — each chunk's
5479 // rows are the t=1 program bit-for-bit and the rope pass advances the cache,
5480 // so a chunked call is value-identical to the row-outer loop it replaces.
5481 static TROWS_PREFILL: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
5482 let trows_prefill = *TROWS_PREFILL
5483 .get_or_init(|| std::env::var("MEMRA_PRIME_TROWS").as_deref() == Ok("1"));
5484 // MEMRA_PRIME_TROWS_T=<w>: chunk width, default 8 = the REAL cap of this walk.
5485 // The workspace slabs go to 32 rows, but `matvec_bf16_qkvg_tcol_into` refuses
5486 // t > 8 (compile-time-T twins exist for 2/4/8 only; the runtime-t kernel spills
5487 // its accumulators to local memory), so a wider chunk fails the request with
5488 // "matvec_bf16_qkvg_tcol geometry" — which is exactly how the first server-path
5489 // TROWS arm died. Measured at 193 tokens: w=8 2.459 s, w=4 2.574 s.
5490 static TROWS_W: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
5491 let trows_w = *TROWS_W.get_or_init(|| {
5492 std::env::var("MEMRA_PRIME_TROWS_T")
5493 .ok()
5494 .and_then(|v| v.parse::<usize>().ok())
5495 .filter(|w| (2..=8).contains(w))
5496 .unwrap_or(8)
5497 });
5498 if tcol && trows_prefill && t > trows_w {
5499 // One-time engagement receipt: without it a prefill gate cannot tell a
5500 // chunked walk from the row-outer fallback it is supposed to replace
5501 // (the first PRIME_TROWS gate passed vacuously on exactly that).
5502 static SEEN: std::sync::atomic::AtomicBool =
5503 std::sync::atomic::AtomicBool::new(false);
5504 if !SEEN.swap(true, std::sync::atomic::Ordering::Relaxed) {
5505 eprintln!(
5506 "[prime-trows] ENGAGED t={t} width={trows_w} chunks={} layers={}..{}",
5507 t.div_ceil(trows_w),
5508 lo,
5509 hi
5510 );
5511 }
5512 let mut out = e.uninit(t * n_embd)?;
5513 let mut start = 0usize;
5514 while start < t {
5515 let mut end = (start + trows_w).min(t);
5516 if t - end == 1 {
5517 end -= 1;
5518 }
5519 let tc = end - start;
5520 let mut xc = e.uninit(tc * n_embd)?;
5521 e.dtod_copy_view(&x.slice(start * n_embd..end * n_embd), &mut xc)?;
5522 let oc =
5523 self.step35_verify_batch_layers(e, xc, lo, hi, pos0 + start, tc, cache)?;
5524 e.copy_into(&mut out, start * n_embd, &oc, tc * n_embd)?;
5525 start = end;
5526 }
5527 return Ok(out);
5528 }
5529 if tcol && t >= 2 && t <= 32 {
5530 // MEMRA_TCOL_PROF=1: synchronized per-segment wall profile of the walk
5531 // (norm+QKV precompute / per-col attention / per-col residual+FFN). The
5532 // syncs serialize the stream, so the split is for TARGETING amortization
5533 // work only — never a perf claim.
5534 static PROF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
5535 let prof =
5536 *PROF.get_or_init(|| std::env::var("MEMRA_TCOL_PROF").as_deref() == Ok("1"));
5537 let mut prof_ms = [0f64; 3];
5538 let eps = self.cfg.rms_eps;
5539 let mut x_t = x;
5540 let mut h_t = e.uninit(t * n_embd)?;
5541 let mut h_row = e.uninit(n_embd)?; // real row: the non-dcw fallback reads it
5542 // Per-column pos buffers hoisted out of the layer loop (a per-col-per-layer
5543 // pageable htod was an in-stream engine turnaround x t x 45).
5544 let mut pos_rows = Vec::with_capacity(t);
5545 for r in 0..t {
5546 pos_rows.push(e.htod_i32(&[(pos0 + r) as i32])?);
5547 }
5548 let mut ok = true;
5549 // MEMRA_TCOL_OPROJ=1: defer each column's o_proj — the finish seam
5550 // stashes `gated` instead of joining per column; one b4_tcol per rank +
5551 // one slab join produce every column's `mixed` after the attention pass.
5552 // Bit-exact per column (t=1 b4 program per column; elementwise join).
5553 // MEMRA_TCOL_FFN=1 (implies the o_proj defer): when every column of a
5554 // MoE layer deferred, the residual norm runs as one t-grid launch
5555 // (per-row program == t=1) and the FFN as ONE two-column device-routed
5556 // sweep + per-column shexp — the two columns' expert weights dedup
5557 // through L2 instead of reading HBM twice.
5558 static FFN2: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
5559 let ffn_batch =
5560 *FFN2.get_or_init(|| std::env::var("MEMRA_TCOL_FFN").as_deref() == Ok("1"));
5561 let oproj_batch = crate::tp::tcol_oproj_on() || ffn_batch;
5562 // MEMRA_SPEC_FA2=1 (T=2 only): eligible layers defer BOTH columns' fa —
5563 // the per-column pass norms/ropes/appends and stashes q+gate, then one
5564 // shared-KV fa_decode_dcw2 per rank + the o_proj join produce the
5565 // [2, o_out] mixed slab. The precheck runs before arming (stashing is
5566 // unrecoverable); ineligible/boundary layers run the ordinary program.
5567 let fa2 = crate::tp::spec_fa2_on() && t <= 32;
5568 let mut mixed_row = e.uninit(n_embd)?;
5569 let mut pos_staged = false;
5570 for il in lo..hi {
5571 let layer = &self.layers[il];
5572 let fa2_layer = fa2 && self.step35_fa_rows_precheck(cache, il, pos0, t)?;
5573 let mut seg = std::time::Instant::now();
5574 e.rms_norm(&x_t, layer.attn_norm.float_data(), &mut h_t, n_embd, t, eps)?;
5575 if !self.step35_verify_qkv_precompute(e, il, &h_t, t)? {
5576 ok = false;
5577 break;
5578 }
5579 // FULL t-row attention pass (rope/append + fa + combine + o_proj in
5580 // 3 launches/rank): same-session rows, slot = len-base+r, one len
5581 // advance by t. Host cache bookkeeping mirrors the per-column tail.
5582 if fa2_layer {
5583 if let Some(mixed_t) =
5584 self.step35_verify_rope_fa_pass(e, il, cache, pos0, t, !pos_staged)?
5585 {
5586 pos_staged = true;
5587 {
5588 let tp_kv = cache.tp_kv[il]
5589 .as_mut()
5590 .expect("precheck verified the distributed cache");
5591 let transaction = tp_kv.begin_transaction()?;
5592 let crate::hybrid::Mixer::Full(fa) = &layer.mixer else {
5593 return Err("verify rope pass expects full attention".into());
5594 };
5595 let tp = fa
5596 .step_tp_qkv
5597 .as_ref()
5598 .ok_or("verify rope pass lost its TP state")?;
5599 let empty: [CudaSlice<f32>; 0] = [];
5600 tp.runtime.append_tp_kv_transaction_inner(
5601 tp_kv,
5602 transaction,
5603 &empty,
5604 &empty,
5605 t,
5606 true,
5607 )?;
5608 tp.runtime.commit_tp_kv_transaction_external(
5609 tp_kv,
5610 transaction,
5611 t,
5612 )?;
5613 if let Some(local) = cache.kv[il].as_mut() {
5614 local.len = pos0 + t;
5615 if !crate::tp::len_mirror_lazy_on() {
5616 e.set_i32_one(&mut local.len_d, local.len as i32)?;
5617 }
5618 }
5619 }
5620 if prof {
5621 e.stream().synchronize()?;
5622 prof_ms[1] += seg.elapsed().as_secs_f64() * 1e3;
5623 seg = std::time::Instant::now();
5624 }
5625 let o_out = mixed_t.len() / t;
5626 let mut next = e.uninit(t * n_embd)?;
5627 let mut batched = false;
5628 if ffn_batch && o_out == n_embd {
5629 let mut x1_t = e.uninit(t * n_embd)?;
5630 let mut z_t = e.uninit(t * n_embd)?;
5631 e.add_rms_norm(
5632 &x_t,
5633 &mixed_t,
5634 layer.post_attn_norm.float_data(),
5635 &mut x1_t,
5636 &mut z_t,
5637 n_embd,
5638 t,
5639 eps,
5640 )?;
5641 if let Some(ffn_t) = self.step35_verify_moe_tn(e, il, &z_t, t)? {
5642 let mut x2_t = e.uninit(t * n_embd)?;
5643 e.add(&x1_t, &ffn_t, &mut x2_t, t * n_embd)?;
5644 next = x2_t;
5645 batched = true;
5646 }
5647 }
5648 if !batched {
5649 for r in 0..t {
5650 e.dtod_copy_view(
5651 &mixed_t.slice(r * o_out..(r + 1) * o_out),
5652 &mut mixed_row,
5653 )?;
5654 let mut x_row = e.uninit(n_embd)?;
5655 e.dtod_copy_view(
5656 &x_t.slice(r * n_embd..(r + 1) * n_embd),
5657 &mut x_row,
5658 )?;
5659 let (x1, ffn_out) = self.residual_norm_ffn(
5660 e, layer, &x_row, &mixed_row, n_embd, il, eps,
5661 )?;
5662 let mut x2 = e.uninit(n_embd)?;
5663 e.add(&x1, &ffn_out, &mut x2, n_embd)?;
5664 e.dtod_copy_into(&x2, &mut next, r * n_embd)?;
5665 }
5666 }
5667 if prof {
5668 e.stream().synchronize()?;
5669 prof_ms[2] += seg.elapsed().as_secs_f64() * 1e3;
5670 }
5671 x_t = next;
5672 continue;
5673 }
5674 }
5675 if prof {
5676 e.stream().synchronize()?;
5677 prof_ms[0] += seg.elapsed().as_secs_f64() * 1e3;
5678 seg = std::time::Instant::now();
5679 }
5680 let mut next = e.uninit(t * n_embd)?;
5681 // Columns whose o_proj was deferred (their FFN runs after the join).
5682 // A NON-deferred column's FFN must run INSIDE the column loop: the
5683 // oproj-tail handoff is a single cell that the same column's
5684 // residual_norm_ffn consumes before the next column's finish.
5685 let mut deferred: Vec<usize> = Vec::new();
5686 let mut fa2_deferred: Vec<usize> = Vec::new();
5687 let mut ffn_col =
5688 |r: usize,
5689 mixed: &CudaSlice<f32>,
5690 next: &mut CudaSlice<f32>|
5691 -> Result<(), Box<dyn std::error::Error>> {
5692 let mut x_row = e.uninit(n_embd)?;
5693 e.dtod_copy_view(&x_t.slice(r * n_embd..(r + 1) * n_embd), &mut x_row)?;
5694 let (x1, ffn_out) =
5695 self.residual_norm_ffn(e, layer, &x_row, mixed, n_embd, il, eps)?;
5696 let mut x2 = e.uninit(n_embd)?;
5697 e.add(&x1, &ffn_out, &mut x2, n_embd)?;
5698 e.dtod_copy_into(&x2, next, r * n_embd)?;
5699 Ok(())
5700 };
5701 for r in 0..t {
5702 e.dtod_copy_view(&h_t.slice(r * n_embd..(r + 1) * n_embd), &mut h_row)?;
5703 let row_pos = &pos_rows[r];
5704 crate::tp::set_verify_tcol(Some(r));
5705 if fa2_layer {
5706 crate::tp::set_spec_fa2_defer(Some(r));
5707 } else if oproj_batch {
5708 crate::tp::set_tcol_oproj_defer(Some(r));
5709 }
5710 let mixed = match &layer.mixer {
5711 crate::hybrid::Mixer::Full(fa) => {
5712 self.full_attn_decode(e, fa, &h_row, row_pos, pos0 + r, cache, il)
5713 }
5714 _ => Err("step35 verify expects full attention".into()),
5715 };
5716 crate::tp::set_verify_tcol(None);
5717 crate::tp::set_spec_fa2_defer(None);
5718 crate::tp::set_tcol_oproj_defer(None);
5719 let mixed = mixed?;
5720 if fa2_layer && crate::tp::take_spec_fa2_stashed() {
5721 fa2_deferred.push(r);
5722 } else if oproj_batch && crate::tp::take_tcol_oproj_stashed() {
5723 deferred.push(r);
5724 } else {
5725 ffn_col(r, &mixed, &mut next)?;
5726 }
5727 }
5728 if !fa2_deferred.is_empty() && fa2_deferred.len() != t {
5729 // The precheck guarantees both columns stash or neither; a strict
5730 // subset means a column's output was never produced anywhere.
5731 return Err("spec fa2 stash engaged for a subset of columns".into());
5732 }
5733 if prof {
5734 e.stream().synchronize()?;
5735 prof_ms[1] += seg.elapsed().as_secs_f64() * 1e3;
5736 seg = std::time::Instant::now();
5737 }
5738 if !fa2_deferred.is_empty() {
5739 deferred = fa2_deferred;
5740 }
5741 if !deferred.is_empty() {
5742 let mixed_t = if fa2_layer {
5743 self.step35_verify_fa_rows_join(e, il, cache, pos0, t)?
5744 } else {
5745 self.step35_verify_oproj_tcol(e, il, t)?
5746 };
5747 let o_out = mixed_t.len() / t;
5748 // Batched t=2 residual+MoE: one t-grid add_rms_norm (per-row
5749 // program == t=1; bit-identical to the oproj-tail join per the
5750 // M2 verbatim-program contract) feeding the two-column routed
5751 // sweep. Ineligible layers (dense FFN, non-nvfp4) fall through
5752 // to the per-column body.
5753 let mut batched = false;
5754 if ffn_batch && deferred.len() == t && o_out == n_embd {
5755 let mut x1_t = e.uninit(t * n_embd)?;
5756 let mut z_t = e.uninit(t * n_embd)?;
5757 e.add_rms_norm(
5758 &x_t,
5759 &mixed_t,
5760 layer.post_attn_norm.float_data(),
5761 &mut x1_t,
5762 &mut z_t,
5763 n_embd,
5764 t,
5765 eps,
5766 )?;
5767 if let Some(ffn_t) = self.step35_verify_moe_tn(e, il, &z_t, t)? {
5768 let mut x2_t = e.uninit(t * n_embd)?;
5769 e.add(&x1_t, &ffn_t, &mut x2_t, t * n_embd)?;
5770 next = x2_t;
5771 batched = true;
5772 }
5773 }
5774 if !batched {
5775 for &r in &deferred {
5776 e.dtod_copy_view(
5777 &mixed_t.slice(r * o_out..(r + 1) * o_out),
5778 &mut mixed_row,
5779 )?;
5780 ffn_col(r, &mixed_row, &mut next)?;
5781 }
5782 }
5783 }
5784 if prof {
5785 e.stream().synchronize()?;
5786 prof_ms[2] += seg.elapsed().as_secs_f64() * 1e3;
5787 }
5788 drop(ffn_col);
5789 x_t = next;
5790 }
5791 if prof {
5792 eprintln!(
5793 "[tcol-prof] t={t} norm+qkv={:.3}ms attn={:.3}ms ffn={:.3}ms",
5794 prof_ms[0], prof_ms[1], prof_ms[2]
5795 );
5796 }
5797 if ok {
5798 return Ok(x_t);
5799 }
5800 // fall through to the row-outer walk on ineligible layers
5801 x = x_t;
5802 }
5803 let mut next = e.uninit(t * n_embd)?;
5804 for r in 0..t {
5805 let mut row = e.uninit(n_embd)?;
5806 e.dtod_copy_view(&x.slice(r * n_embd..(r + 1) * n_embd), &mut row)?;
5807 let row_pos = e.htod_i32(&[(pos0 + r) as i32])?;
5808 let out = self.decode_layers_eager(e, row, lo, hi, &row_pos, pos0 + r, cache)?;
5809 e.dtod_copy_into(&out, &mut next, r * n_embd)?;
5810 }
5811 // dflash taps are NOT produced on this arm (they need per-layer hiddens the
5812 // row-outer walk does not materialize); the door is a step37 MTP bring-up
5813 // surface where taps are unused.
5814 return Ok(next);
5815 }
5816 let mut ph_last = std::time::Instant::now();
5817 for il in lo..hi {
5818 let mut next = e.uninit(t * n_embd)?;
5819 for r in 0..t {
5820 let mut row = e.uninit(n_embd)?;
5821 e.dtod_copy_view(&x.slice(r * n_embd..(r + 1) * n_embd), &mut row)?;
5822 // The caller owns this verify's position. During controller overlap, cache.pos
5823 // still describes generation N while this stage-0 walk belongs to N+1.
5824 let row_pos = e.htod_i32(&[(pos0 + r) as i32])?;
5825 let mut one = [&mut *cache];
5826 let out = self.step35_decode_batch_layers(
5827 e,
5828 row,
5829 &mut one,
5830 &[(pos0 + r) as i32],
5831 &row_pos,
5832 il,
5833 il + 1,
5834 &mut ph_last,
5835 )?;
5836 e.dtod_copy_into(&out, &mut next, r * n_embd)?;
5837 }
5838 self.dflash_tap(e, cache, il, &next, t)?;
5839 x = next;
5840 }
5841 Ok(x)
5842 }
5843
5844 /// DSpark drafter verify (lane/dspark-q38-recover): one t-row forward through the
5845 /// SERVING-CLASS verify funnel (`decode_step_t_core_stream` — the same numeric class
5846 /// MTP verify rides, GDN state advanced in place), returning per-row argmax tokens.
5847 /// Advances `cache.pos += t`; the caller owns snapshot/rollback (block acceptance is
5848 /// prefix-keep, not all-or-nothing).
5849 pub(crate) fn dspark_verify_t_am(
5850 &self,
5851 e: &Engine,
5852 tokens: &[u32],
5853 pos0: usize,
5854 cache: &mut Cache,
5855 ) -> Result<Vec<u32>, Box<dyn std::error::Error>> {
5856 let (logits, _hn) = self.decode_step_t_core_stream(
5857 e, tokens, pos0, cache, None, None, None, None, None, None,
5858 )?;
5859 let t = tokens.len();
5860 let v = self.output.out_features();
5861 let mut am_d = e.stream().alloc_zeros::<u32>(t)?;
5862 for r in 0..t {
5863 e.argmax_token_device_col(&logits, r, v, &mut am_d, r)?;
5864 }
5865 Ok(e.dtoh_u32(&am_d)?)
5866 }
5867
5868 /// DSpark verify returning the RAW verify logits [t, n_vocab] (device-resident) instead
5869 /// of per-row argmaxes — the sampled-admission arm's input (rejection-sampling accept
5870 /// gathers filtered p from these columns; lane/dspark-sampled-admission-20260820). Same
5871 /// forward as `dspark_verify_t_am`; the greedy arm keeps its argmax wrapper untouched.
5872 pub(crate) fn dspark_verify_t_logits(
5873 &self,
5874 e: &Engine,
5875 tokens: &[u32],
5876 pos0: usize,
5877 cache: &mut Cache,
5878 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5879 let (logits, _hn) = self.decode_step_t_core_stream(
5880 e, tokens, pos0, cache, None, None, None, None, None, None,
5881 )?;
5882 Ok(logits)
5883 }
5884
5885 /// DSpark verify with the MTP column-stash armed: identical forward to
5886 /// `dspark_verify_t_am`, but fills a `VerifyCkpt` so a partial accept can restore
5887 /// column state directly (`dspark_commit_prefix`) instead of snapshot-replay.
5888 /// The ckpt type is opaque outside spec.rs (newtype) — dflash.rs threads it through.
5889 pub(crate) fn dspark_verify_t_am_ckpt(
5890 &self,
5891 e: &Engine,
5892 tokens: &[u32],
5893 pos0: usize,
5894 cache: &mut Cache,
5895 ) -> Result<(Vec<u32>, DsparkVerifyCkpt), Box<dyn std::error::Error>> {
5896 let mut ck = VerifyCkpt::new(self.layers.len());
5897 let (logits, _hn) = self.decode_step_t_core_stream(
5898 e,
5899 tokens,
5900 pos0,
5901 cache,
5902 None,
5903 Some(&mut ck),
5904 None,
5905 None,
5906 None,
5907 None,
5908 )?;
5909 let t = tokens.len();
5910 let v = self.output.out_features();
5911 let mut am_d = e.stream().alloc_zeros::<u32>(t)?;
5912 for r in 0..t {
5913 e.argmax_token_device_col(&logits, r, v, &mut am_d, r)?;
5914 }
5915 Ok((e.dtoh_u32(&am_d)?, DsparkVerifyCkpt(ck)))
5916 }
5917
5918 /// Engine-bundle slice 2: `dspark_verify_t_am_ckpt` with DEVICE tokens and NO readback.
5919 /// The verify tokens are the round's `chain_d` (cand layout: [anchor, drafts...]); the
5920 /// embed gathers its first `t` entries on-device (`embed_gather_u32_t` — bit-identical
5921 /// rows to the host arm), so the host never blocks on the draft chain before dispatching
5922 /// verify. Returns the device per-row argmax buffer; the caller merges its readback with
5923 /// the chain's into ONE sync. Forward, ckpt fill and argmax walk are `_ckpt` verbatim.
5924 pub(crate) fn dspark_verify_t_am_ckpt_dev(
5925 &self,
5926 e: &Engine,
5927 vtok: &CudaSlice<u32>,
5928 t: usize,
5929 pos0: usize,
5930 cache: &mut Cache,
5931 embd_dev: (&CudaSlice<u8>, i32, usize),
5932 graphs: Option<&mut DsparkVerifyGraphs>,
5933 ) -> Result<(CudaSlice<u32>, DsparkVerifyCkpt), Box<dyn std::error::Error>> {
5934 debug_assert!(
5935 vtok.len() >= t,
5936 "verify window exceeds the device token buffer"
5937 );
5938 // The slab flag is a per-round statement: clear it here so a verify that never
5939 // reaches the graphs door (rowwise env, a non-tparallel arm) cannot leave a
5940 // stale `true` steering the commit at slabs the round never wrote.
5941 let mut graphs = graphs;
5942 if let Some(g) = graphs.as_deref_mut() {
5943 g.round_slab = false;
5944 }
5945 let mut ck = VerifyCkpt::new(self.layers.len());
5946 // Dummy host tokens size the funnel; the embed reads `vtok` (the round-stream
5947 // arm's established pattern — spec.rs stream-mode verify does the same).
5948 let dummy = vec![0u32; t];
5949 let (logits, _hn) = self.decode_step_t_core_stream(
5950 e,
5951 &dummy,
5952 pos0,
5953 cache,
5954 Some(embd_dev),
5955 Some(&mut ck),
5956 None,
5957 None,
5958 Some(vtok),
5959 graphs,
5960 )?;
5961 let v = self.output.out_features();
5962 let mut am_d = e.stream().alloc_zeros::<u32>(t)?;
5963 for r in 0..t {
5964 e.argmax_token_device_col(&logits, r, v, &mut am_d, r)?;
5965 }
5966 Ok((am_d, DsparkVerifyCkpt(ck)))
5967 }
5968
5969 /// Ckpt-armed twin of [`Self::dspark_verify_t_logits`] (sampled-admission arm).
5970 pub(crate) fn dspark_verify_t_logits_ckpt(
5971 &self,
5972 e: &Engine,
5973 tokens: &[u32],
5974 pos0: usize,
5975 cache: &mut Cache,
5976 ) -> Result<(CudaSlice<f32>, DsparkVerifyCkpt), Box<dyn std::error::Error>> {
5977 let mut ck = VerifyCkpt::new(self.layers.len());
5978 let (logits, _hn) = self.decode_step_t_core_stream(
5979 e,
5980 tokens,
5981 pos0,
5982 cache,
5983 None,
5984 Some(&mut ck),
5985 None,
5986 None,
5987 None,
5988 None,
5989 )?;
5990 Ok((logits, DsparkVerifyCkpt(ck)))
5991 }
5992
5993 /// Restore the round to `keep` accepted columns from the verify stash: KV lens and
5994 /// pos from the pre-verify snapshot + keep, GDN conv/ssm from the stashed column
5995 /// state — no replay forward. The exact `commit_verified_prefix` the MTP path ships.
5996 pub(crate) fn dspark_commit_prefix(
5997 &self,
5998 e: &Engine,
5999 cache: &mut Cache,
6000 snap: &crate::cache::CacheSnapshot,
6001 ckpt: &DsparkVerifyCkpt,
6002 keep: usize,
6003 ) -> Result<(), Box<dyn std::error::Error>> {
6004 self.commit_verified_prefix(e, cache, snap, &ckpt.0, keep, false, None)
6005 }
6006
6007 /// Slice-3 commit twin: restore to `keep` accepted columns when the round's linear
6008 /// column stash lives in the graphs ctx's persistent slabs (`DsparkVerifyGraphs`) —
6009 /// the cols arm's exact semantics (KV lens + pos from the snapshot, GDN conv/ssm
6010 /// from the stash of column keep-1), slab-addressed and batched into two copy
6011 /// launches. `MEMRA_STATE_COPY_BATCH=0` falls back to per-layer view copies.
6012 pub(crate) fn dspark_commit_prefix_slab(
6013 &self,
6014 e: &Engine,
6015 cache: &mut Cache,
6016 snap: &crate::cache::CacheSnapshot,
6017 ctx: &DsparkVerifyGraphs,
6018 keep: usize,
6019 ) -> Result<(), Box<dyn std::error::Error>> {
6020 use cudarc::driver::DevicePtr;
6021 debug_assert!(keep >= 1, "keep==0 rounds take the legacy rollback");
6022 let mut conv_src: Vec<u64> = Vec::new();
6023 let mut ssm_src: Vec<u64> = Vec::new();
6024 let mut conv_dst: Vec<u64> = Vec::new();
6025 let mut ssm_dst: Vec<u64> = Vec::new();
6026 for il in 0..self.layers.len() {
6027 if let (Some(kvl), Some(saved)) = (cache.kv[il].as_mut(), snap.kv_len[il]) {
6028 kvl.len = saved + keep;
6029 e.set_i32_one(&mut kvl.len_d, kvl.len as i32)?;
6030 }
6031 if let Some(rl) = cache.recur[il].as_ref() {
6032 let (pc, ps, _cw, _sw) = ctx
6033 .slab_row(e, il, keep - 1)
6034 .ok_or("slab commit: linear layer missing from the graphs ctx")?;
6035 conv_src.push(pc);
6036 ssm_src.push(ps);
6037 let st = &e.gpu.stream();
6038 let (dc, _g0) = rl.conv_state.device_ptr(st);
6039 let (ds, _g1) = rl.ssm_state.device_ptr(st);
6040 conv_dst.push(dc as u64);
6041 ssm_dst.push(ds as u64);
6042 }
6043 }
6044 let n = conv_src.len();
6045 if n > 0 {
6046 if state_copy_batch_on() {
6047 let mut tt = vec![0u64; 2 * n];
6048 tt[..n].copy_from_slice(&conv_src);
6049 tt[n..].copy_from_slice(&conv_dst);
6050 let ct = e.htod_u64(&tt)?;
6051 tt[..n].copy_from_slice(&ssm_src);
6052 tt[n..].copy_from_slice(&ssm_dst);
6053 let st = e.htod_u64(&tt)?;
6054 e.copy_batch_uniform_f32(&ct, n, ctx.conv_words)?;
6055 e.copy_batch_uniform_f32(&st, n, ctx.ssm_words)?;
6056 } else {
6057 let (cw, sw) = (ctx.conv_words, ctx.ssm_words);
6058 let row = keep - 1;
6059 for il in 0..self.layers.len() {
6060 let Some(rl) = cache.recur[il].as_mut() else {
6061 continue;
6062 };
6063 let k = ctx.lin_pos[&il];
6064 {
6065 let sv = e.view(&ctx.stash_conv[k], (row + 1) * cw);
6066 let win = sv.slice(row * cw..(row + 1) * cw);
6067 e.copy_view_into(&mut rl.conv_state, 0, &win, cw)?;
6068 }
6069 {
6070 let sv = e.view(&ctx.stash_ssm[k], (row + 1) * sw);
6071 let win = sv.slice(row * sw..(row + 1) * sw);
6072 e.copy_view_into(&mut rl.ssm_state, 0, &win, sw)?;
6073 }
6074 }
6075 }
6076 }
6077 cache.pos = snap.pos + keep;
6078 Ok(())
6079 }
6080
6081 /// Qwen35-family verify trunk in the live serving numeric class.
6082 ///
6083 /// Serving intentionally keeps this architecture in the generic batched program even at
6084 /// B=1. The older verify walk used its own mirrored dispatch and can flip near-tie argmaxes.
6085 ///
6086 /// Two arms, one numeric class:
6087 /// - DENSE GDN (`DenseMlp`, t<=16): `qwen35_verify_tparallel` — the weight ops (norms,
6088 /// projections, FFN) hoist to m=T through the exact-tier batched kernels whose per-row
6089 /// program IS the m=1 program (`matmul_pre == fused2 per (tensor,row); _bN mmvq per-row
6090 /// == m=1` — decode_batch.rs v2 note), while the state ops (conv ring, gdn scan, KV
6091 /// append, fa decode) stay a per-row loop running the b_n=1 serving kernels with each
6092 /// row's own t_kv-driven arm pick (the straddle law: every row executes the exact
6093 /// program its isolated serving step would). One weight read per layer per round
6094 /// instead of T — this is what makes MTP profitable in the exact class (the per-row
6095 /// walk measured verify(K+1) ~= (K+1) plain steps: 69 -> 44 tok/s served, 2026-08-15).
6096 /// - MoE / t>16 / `MEMRA_SPEC_VERIFY_ROWWISE=1`: the per-row replay of the authoritative
6097 /// serving layer body, preserving single-session autoregressive cache order (the
6098 /// correctness reference; also the rollback seam for the t-parallel arm).
6099 ///
6100 /// Bit-identity of the t-parallel arm vs the rowwise arm is gated by spec-serve-gate
6101 /// (zero differing logits at T=1..4, K arms) + the 8-prompt ON/OFF canary before ship.
6102 #[allow(clippy::too_many_arguments)]
6103 fn qwen35_verify_batch_layers(
6104 &self,
6105 e: &Engine,
6106 x: CudaSlice<f32>,
6107 lo: usize,
6108 hi: usize,
6109 pos0: usize,
6110 t: usize,
6111 cache: &mut Cache,
6112 ckpt: Option<&mut VerifyCkpt>,
6113 stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
6114 graphs: Option<&mut DsparkVerifyGraphs>,
6115 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6116 // Qwen35Moe admitted 2026-08-20 (lane/draftcost-moe): the t-parallel arm already
6117 // carries the MoE FFN (`moe_ffn_il_zq8` at m=T) and the GDN per-row state loop; the
6118 // arch fence was a qualification gate, not a mechanism gap. Measured disease on the
6119 // 35B-A3B class: rowwise verify ~= 5.6 ms per drafted token (one full trunk step
6120 // each) — the same (K+1)-plain-steps wall the dense admission fixed on 2026-08-15.
6121 // Rollback seam unchanged: MEMRA_SPEC_VERIFY_ROWWISE=1.
6122 let rowwise = std::env::var("MEMRA_SPEC_VERIFY_ROWWISE").as_deref() == Ok("1")
6123 || !self.batched_serving_numeric_class()
6124 || t > 16;
6125 if rowwise {
6126 if stream.is_some() {
6127 // rowwise replays per row with host cache.pos — irreconcilable with a
6128 // device position counter. Burst callers must keep t <= 16 and the
6129 // ROWWISE env unset; refusing beats silently mispositioned rows.
6130 return Err("qwen35 rowwise verify has no ROUND-STREAM arm \
6131 (t > 16 or MEMRA_SPEC_VERIFY_ROWWISE=1)"
6132 .into());
6133 }
6134 self.qwen35_verify_rowwise(e, x, lo, hi, pos0, t, cache, ckpt)
6135 } else {
6136 self.qwen35_verify_tparallel(e, x, lo, hi, pos0, t, cache, ckpt, stream, graphs)
6137 }
6138 }
6139
6140 /// The per-row correctness reference: replay each verify row through the authoritative
6141 /// serving layer body (`decode_batch_layers` at b_n=1). T full weight reads per layer.
6142 #[allow(clippy::too_many_arguments)]
6143 fn qwen35_verify_rowwise(
6144 &self,
6145 e: &Engine,
6146 mut x: CudaSlice<f32>,
6147 lo: usize,
6148 hi: usize,
6149 pos0: usize,
6150 t: usize,
6151 cache: &mut Cache,
6152 mut ckpt: Option<&mut VerifyCkpt>,
6153 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6154 let n_embd = self.cfg.n_embd as usize;
6155 let saved_pos = cache.pos;
6156 let mut ph_last = std::time::Instant::now();
6157 for il in lo..hi {
6158 let mut next = e.uninit(t * n_embd)?;
6159 let mut col_states: Option<Vec<(CudaSlice<f32>, CudaSlice<f32>)>> =
6160 if ckpt.is_some() && t >= 2 && matches!(self.layers[il].mixer, Mixer::Linear(_)) {
6161 Some(Vec::with_capacity(t - 1))
6162 } else {
6163 None
6164 };
6165 for r in 0..t {
6166 cache.pos = pos0 + r;
6167 let mut row = e.uninit(n_embd)?;
6168 e.dtod_copy_view(&x.slice(r * n_embd..(r + 1) * n_embd), &mut row)?;
6169 let row_pos = e.htod_i32(&[(pos0 + r) as i32])?;
6170 let mut one = [&mut *cache];
6171 let ctx = self.batch_layer_ctx(e, &one, il, il + 1)?;
6172 let out = match self.decode_batch_layers(
6173 e,
6174 row,
6175 &mut one,
6176 &ctx,
6177 &row_pos,
6178 &mut ph_last,
6179 ) {
6180 Ok(out) => out,
6181 Err(error) => {
6182 cache.pos = saved_pos;
6183 return Err(error);
6184 }
6185 };
6186 e.dtod_copy_into(&out, &mut next, r * n_embd)?;
6187 if r + 1 < t {
6188 if let Some(states) = col_states.as_mut() {
6189 let recur = cache.recur[il]
6190 .as_ref()
6191 .ok_or("Qwen35-MoE linear verify layer has no recurrent state")?;
6192 states.push((
6193 e.clone_dtod(&recur.conv_state)?,
6194 e.clone_dtod(&recur.ssm_state)?,
6195 ));
6196 }
6197 }
6198 }
6199 if let (Some(checkpoint), Some(states)) = (ckpt.as_deref_mut(), col_states) {
6200 checkpoint.cols[il] = Some(states);
6201 }
6202 x = next;
6203 }
6204 cache.pos = saved_pos;
6205 Ok(x)
6206 }
6207
6208 /// T-PARALLEL VERIFY IN THE SERVING NUMERIC CLASS (lane/tparallel-verify, 2026-08-15).
6209 ///
6210 /// The weight ops run ONCE per layer at m=T; the state ops run per row through the same
6211 /// b_n=1 serving kernels the rowwise replay uses. Per-row bit-identity rests on the two
6212 /// pins the serving batch tier already carries:
6213 /// * `matmul_pre` / `_bN` mmvq: per-row program == m=1 program (decode_batch.rs v2 note,
6214 /// kernel-check pinned) — so a [T, n_embd] projection row equals the row projected
6215 /// alone;
6216 /// * row-indexed norms/elementwise (`rms_norm`, `quantize_q8_1`, `add_rms_norm`,
6217 /// `gated_rmsnorm[_q8_1]`, `silu_mul`, `rope_neox` with per-row positions): the T-row
6218 /// launch is the per-row program (same pin the generic verify's fused norms rely on).
6219 /// The sequential dependencies keep their exact serving order: the conv ring / gdn scan
6220 /// chain state row -> row through the `_b` kernels at b_n=1 (ping-pong via a 6-entry
6221 /// alternating pointer table, host handles swapped per row so VerifyCkpt clones the
6222 /// canonical state exactly as the rowwise arm does), and each row's KV append + fa decode
6223 /// picks its arm from ITS OWN t_kv (append: format-only; fa: `fa_seqs_eligible` + its own
6224 /// `fa_split_keys` rung at b_n=1) — the straddle law per row, so every row executes the
6225 /// program its isolated B=1 serving step would.
6226 ///
6227 /// Cost: 1 weight read per layer per round + T state micro-launches, vs the rowwise arm's
6228 /// T weight reads. Gated bit-identical vs the rowwise arm by spec-serve-gate + canary.
6229 #[allow(clippy::too_many_arguments)]
6230 fn qwen35_verify_tparallel(
6231 &self,
6232 e: &Engine,
6233 mut x: CudaSlice<f32>,
6234 lo: usize,
6235 hi: usize,
6236 pos0: usize,
6237 t: usize,
6238 cache: &mut Cache,
6239 mut ckpt: Option<&mut VerifyCkpt>,
6240 stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
6241 mut graphs: Option<&mut DsparkVerifyGraphs>,
6242 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6243 let seqs_append =
6244 std::env::var("MEMRA_BATCH_APPEND").as_deref() != Ok("0") && !Engine::kv_fp8_on();
6245 let batch_fa_on = std::env::var("MEMRA_BATCH_FA").as_deref() != Ok("0");
6246
6247 // Merge guard (v0.98 train, re-affirmed on the v0.100 train over slice 4c): the
6248 // ROUND-STREAM arm (lane/draftcost-moe, device position counter) and the dspark
6249 // verify graphs (engine-bundle slice 3 / trunk slice 4c) have no common caller —
6250 // stream rides the qwen35moe burst, graphs ride the dspark route. If a future
6251 // caller arms both, refuse loudly instead of silently dropping the graphs ctx
6252 // (the stream linear arm takes linear_attn_verify_t, not the graphed segment or
6253 // full-verify bodies).
6254 if stream.is_some() && graphs.is_some() {
6255 return Err(
6256 "qwen35 tparallel verify: ROUND-STREAM and dspark verify graphs \
6257 cannot arm together"
6258 .into(),
6259 );
6260 }
6261 // Engine-bundle slice 3 + slice 4c: with a graphs ctx armed, pointer tables are
6262 // refreshed once per verify (the gdn ping-pong moves handles; a fresh generation
6263 // moves the kv caches). Then:
6264 // - slice 4c: when the WHOLE round rides one seqs rung (every row batchable, one
6265 // split-ladder step, rung covers the round), the ENTIRE walk replays as ONE
6266 // full-verify graph per (vt, rung) — linear layers through the shared
6267 // `qwen35_tparallel_linear_layer` body, full-attention layers through the
6268 // shared `qwen35_tparallel_fa_layer` body in graph mode.
6269 // - fallback (straddle rounds, below the vec floor, partial walks): runs of
6270 // consecutive LINEAR layers replay the slice-3 per-(segment, vt) graphs and
6271 // the full-attention layers run eager (batched rows when eligible).
6272 if let Some(g) = graphs.as_deref_mut() {
6273 g.refresh_tables(e, cache)?;
6274 g.round_slab = false;
6275 if let Some(rung) = g.full_rung(self, cache, lo, hi, t, seqs_append && batch_fa_on) {
6276 // Pool ceiling (dspark_vg_cap): an existing key always replays; a NEW
6277 // full capture past the ceiling falls through to the segment/eager arms.
6278 if g.full.contains_key(&(t, rung, hi)) || g.can_capture() {
6279 let out = g.run_full(self, e, lo, hi, &x, t, pos0, rung, cache)?;
6280 g.round_slab = true;
6281 return Ok(out);
6282 }
6283 }
6284 // Round-atomic ceiling check for the segment door: if any linear run in this
6285 // walk would need a NEW capture past the ceiling, the whole round runs the
6286 // eager cols-ckpt walk (mixing slab- and cols-stashed layers in one round
6287 // would corrupt the commit).
6288 if !g.segments_ready(self, lo, hi, t) {
6289 graphs = None;
6290 }
6291 }
6292 // STREAM (2b, lane/draftcost-moe): positions come from the device round counter
6293 // (pos_iota / i32_copy_add) so a burst round needs no host position knowledge.
6294 let pos_d = match stream {
6295 Some((_, ctr)) => {
6296 let mut p = e.alloc_uninit::<i32>(t)?;
6297 e.pos_iota(ctr, &mut p, t)?;
6298 p
6299 }
6300 None => {
6301 let pos_host: Vec<i32> = (0..t).map(|r| (pos0 + r) as i32).collect();
6302 e.htod_i32(&pos_host)?
6303 }
6304 };
6305 // Per-row 1-element position buffers, built ONCE per verify (the append/fa wrappers
6306 // take owned pos slices; building these inside the layer x row loops cost 16xT H2Ds).
6307 // LAZY since slice 4: the batched fa/append arm never touches them — they are built
6308 // on the first per-row fallback layer only (stream-aware there; the stream FA arm
6309 // rides the dc rows kernels and never reaches the fallback).
6310 let mut pos_rows: Option<Vec<CudaSlice<i32>>> = None;
6311 let mut il = lo;
6312 while il < hi {
6313 if graphs.is_some() && matches!(self.layers[il].mixer, Mixer::Linear(_)) {
6314 let mut end = il;
6315 while end < hi && matches!(self.layers[end].mixer, Mixer::Linear(_)) {
6316 end += 1;
6317 }
6318 let g = graphs.as_deref_mut().expect("checked above");
6319 x = g.run_segment(self, e, il, end, &x, t, cache)?;
6320 g.round_slab = true;
6321 il = end;
6322 continue;
6323 }
6324 let layer = &self.layers[il];
6325 if stream.is_none() && matches!(layer.mixer, Mixer::Linear(_)) {
6326 // Eager linear layer (no graphs ctx): the shared body, legacy cols-ckpt arm.
6327 // Under ROUND-STREAM the linear layers ride the fa-body match's stream arm
6328 // below (linear_attn_verify_t — the stream COMMIT needs its GdnStash).
6329 x = self.qwen35_tparallel_linear_layer(
6330 e,
6331 il,
6332 &x,
6333 t,
6334 cache,
6335 ckpt.as_deref_mut(),
6336 None,
6337 None,
6338 )?;
6339 il += 1;
6340 continue;
6341 }
6342 // Full-attention (or stream-Linear, or MLA-refusing) layer: the extracted
6343 // shared body — eager arm (fresh per-verify pos/table, exact t_kv sizing,
6344 // in-body len bump). The slice-4c captured full-verify graphs run the SAME
6345 // body in graph mode; under ROUND-STREAM the body's dc-rows / GDN stream arms
6346 // run (lane/draftcost-moe).
6347 x = self.qwen35_tparallel_fa_layer(
6348 e,
6349 il,
6350 &x,
6351 t,
6352 cache,
6353 FaLayerArgs {
6354 pos_d: &pos_d,
6355 pos_rows: &mut pos_rows,
6356 pos0,
6357 seqs_append,
6358 batch_fa_on,
6359 graph_cap: None,
6360 stream,
6361 ckpt: ckpt.as_deref_mut(),
6362 },
6363 )?;
6364 il += 1;
6365 }
6366 Ok(x)
6367 }
6368
6369 /// SHARED dense-FFN body for the qwen35 t-parallel layers (trunk-kernels slice B) —
6370 /// ONE copy for the fa and linear layer bodies (the verify_layers extraction lesson).
6371 /// Dual arm (MEMRA_TK_FFN_DUAL, default on): gate+up in ONE dual launch from the
6372 /// pre-quantized activation with macro-scales DEFERRED into the fused SwiGLU+q8_1
6373 /// epilogue, then ffn_down from the fused (aq, ad) — the q27 verify chain verbatim.
6374 /// Every door is the bit-identical proven one: `matmul_decode_exact_dual_pre` (per
6375 /// (tensor,token,row) == the two singles), `silu_mul_scaled_q8_1` (y*s inline == the
6376 /// scale_inplace store, value-exact; fused quantize == quantize_q8_1 bytes),
6377 /// `matmul_decode_exact_pre` (dispatch mirror of the singles' q8_1-fast tail).
6378 /// Dual-refused (t outside 2..=7, non-NVFP4, layout mismatch) or seam off -> the
6379 /// original singles chain, byte-for-byte.
6380 #[allow(clippy::too_many_arguments)]
6381 fn qwen35_tparallel_dense_ffn(
6382 &self,
6383 e: &Engine,
6384 ffn_gate: &crate::model::GpuTensor,
6385 ffn_up: &crate::model::GpuTensor,
6386 ffn_down: &crate::model::GpuTensor,
6387 zn: &CudaSlice<f32>,
6388 t: usize,
6389 n_embd: usize,
6390 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6391 let n_ff = ffn_gate.out_features();
6392 let (zq, zd) = e.quantize_q8_1(zn, t, n_embd)?;
6393 if Engine::tk_ffn_dual_on() {
6394 if let Some(((g, gs), (u, us))) =
6395 e.matmul_decode_exact_dual_pre(ffn_gate, ffn_up, &zq, &zd, t)?
6396 {
6397 if e.uses_q8_1_fast(ffn_down) {
6398 let (aq, ad) = e.silu_mul_scaled_q8_1(&g, &u, gs, us, t * n_ff)?;
6399 return e.matmul_decode_exact_pre(ffn_down, &aq, &ad, t);
6400 }
6401 let mut act = e.uninit(t * n_ff)?;
6402 e.silu_mul_scaled(&g, &u, gs, us, &mut act, t * n_ff)?;
6403 let (aq, ad) = e.quantize_q8_1(&act, t, n_ff)?;
6404 return e.matmul_pre(ffn_down, &aq, &ad, &act, t);
6405 }
6406 }
6407 // v1 singles chain (seam off or dual-refused) — the pre-slice-B body verbatim.
6408 let g = e.matmul_pre(ffn_gate, &zq, &zd, zn, t)?;
6409 let u = e.matmul_pre(ffn_up, &zq, &zd, zn, t)?;
6410 let mut act = e.uninit(t * n_ff)?;
6411 e.silu_mul(&g, &u, &mut act, t * n_ff)?;
6412 let (aq, ad) = e.quantize_q8_1(&act, t, n_ff)?;
6413 e.matmul_pre(ffn_down, &aq, &ad, &act, t)
6414 }
6415
6416 /// ONE t-parallel FULL-ATTENTION layer (attn_norm + fa mixer + post_attn_norm + FFN +
6417 /// tap) — extracted from the walk exactly like `qwen35_tparallel_linear_layer` so the
6418 /// eager walk and the slice-4c captured full-verify graphs execute the SAME body (a
6419 /// second copy is how dispatch mirrors drift — the verify_layers extraction lesson).
6420 ///
6421 /// `args.graph_cap = Some((table, off, rung_end))` is the captured-graph mode:
6422 /// - kv base-pointer pairs come from the ctx-owned persistent table at `off` (a fresh
6423 /// generation's cache lands at new addresses that only the per-verify table refresh
6424 /// knows — the slice-3 baked-address lesson);
6425 /// - the seqs twins size partials/grid at `rung_end` and pin `split_keys` to the
6426 /// rung's ladder value: `n_splits_max` is pure stride, splits >= ns_eff write the
6427 /// EMPTY partial the combine never reads, and every per-row T_kv derives in-kernel
6428 /// from `pos_seq[z]` — so one captured launch replays bit-identically for every
6429 /// round whose rows all sit inside the rung;
6430 /// - the host len bump moves to the replay caller (captured host code does not
6431 /// re-run at replay).
6432 /// Graph mode REFUSES any round the batched arm cannot take: the per-row fallback
6433 /// host-branches on t_kv and must never be captured.
6434 #[allow(clippy::too_many_arguments)]
6435 fn qwen35_tparallel_fa_layer(
6436 &self,
6437 e: &Engine,
6438 il: usize,
6439 x: &CudaSlice<f32>,
6440 t: usize,
6441 cache: &mut Cache,
6442 args: FaLayerArgs<'_>,
6443 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6444 use cudarc::driver::DevicePtr;
6445 let cfg = &self.cfg;
6446 let n_embd = cfg.n_embd as usize;
6447 let eps = cfg.rms_eps;
6448 let head_dim_global = cfg.head_dim_k as usize;
6449 let layer = &self.layers[il];
6450 let FaLayerArgs {
6451 pos_d,
6452 pos_rows,
6453 pos0,
6454 seqs_append,
6455 batch_fa_on,
6456 graph_cap,
6457 stream,
6458 mut ckpt,
6459 } = args;
6460
6461 // ---- attn_norm + q8_1 quantize at m=T (row-indexed == per-row) ----
6462 let anorm = layer.attn_norm.float_data();
6463 let mut xn = e.uninit(t * n_embd)?;
6464 e.rms_norm(x, anorm, &mut xn, n_embd, t, eps)?;
6465 let (hq, hd) = e.quantize_q8_1(&xn, t, n_embd)?;
6466
6467 let mixed: CudaSlice<f32> = match &layer.mixer {
6468 Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
6469 // STREAM ARM (2b, lane/draftcost-moe): under a device position counter the
6470 // per-row serving-kernel chain cannot run (host state swaps keyed on host
6471 // row index are fine, but the stream COMMIT needs the GdnStash for its _dc
6472 // rebuild — the per-row chain only produces per-column clones). GDN rides
6473 // `linear_attn_verify_t`: batched q8_1-class projections, stash-producing,
6474 // and its one-scan recurrence is pinned bit-identical to T chained T=1
6475 // steps (its header + kernel-check). Position-independent, so no counter
6476 // plumbing is needed. Guards mirror the generic call site exactly.
6477 Mixer::Linear(la) if stream.is_some() => {
6478 if !(t >= 3 || (t == 2 && spec_m2()))
6479 || !self.mixer_in_q8_1_fast(e, &layer.mixer)
6480 || !e.uses_q8_1_fast(&la.ssm_out)
6481 {
6482 return Err("qwen35 stream verify: GDN batched arm requires t>=3 \
6483 (or MEMRA_SPEC_M2 at t=2) and q8_1-fast projections"
6484 .into());
6485 }
6486 let want = ckpt.is_some();
6487 let (out, stash) =
6488 self.linear_attn_verify_t(e, la, &xn, Some((&hq, &hd)), t, cache, il, want)?;
6489 if let (Some(ck), Some(st)) = (ckpt.as_deref_mut(), stash) {
6490 ck.gdn[il] = Some(st);
6491 }
6492 out
6493 }
6494 Mixer::Linear(_) => {
6495 unreachable!("linear layers ride qwen35_tparallel_linear_layer")
6496 }
6497 Mixer::Full(fa) => {
6498 let geometry = cfg.full_attention_geometry_at(il as u32);
6499 let n_head = geometry.n_head as usize;
6500 let n_head_kv = geometry.n_head_kv as usize;
6501 let head_dim = geometry.head_dim_k as usize;
6502 let rope_dims = geometry.n_rot as usize;
6503 let rope_base = geometry.rope_base;
6504 let scale = geometry.attention_scale();
6505 // Batched projections: one weight read serves all T rows.
6506 // GROUP-3 twin (trunk-kernels slice D): q/k/v in ONE launch — the group4
6507 // kernel with n3=0, bit-identical per (tensor, token, row) to the three
6508 // singles; refused or MEMRA_TK_FA_GROUP=0 -> singles byte-for-byte.
6509 let (qf, mut k, v) = match e.matmul_decode_exact_group3_pre(
6510 [&fa.wq, &fa.wk, &fa.wv],
6511 &hq,
6512 &hd,
6513 t,
6514 )? {
6515 Some(mut g3) => {
6516 let v = g3.pop().unwrap();
6517 let k = g3.pop().unwrap();
6518 let qf = g3.pop().unwrap();
6519 (qf, k, v)
6520 }
6521 None => (
6522 e.matmul_pre(&fa.wq, &hq, &hd, &xn, t)?,
6523 e.matmul_pre(&fa.wk, &hq, &hd, &xn, t)?,
6524 e.matmul_pre(&fa.wv, &hq, &hd, &xn, t)?,
6525 ),
6526 };
6527 let gated =
6528 geometry.attention_gate == memra_gguf::config::AttentionGateKind::FusedQ;
6529 let (mut q, gate) = if gated {
6530 let mut qs = e.uninit(t * n_head * head_dim)?;
6531 let mut gs = e.uninit(t * n_head * head_dim)?;
6532 e.q_gate_split(&qf, &mut qs, &mut gs, head_dim, n_head, t)?;
6533 (qs, Some(gs))
6534 } else {
6535 (qf, None)
6536 };
6537 let mut qn = e.uninit(t * n_head * head_dim)?;
6538 e.rms_norm(
6539 &q,
6540 fa.q_norm.float_data(),
6541 &mut qn,
6542 head_dim,
6543 t * n_head,
6544 eps,
6545 )?;
6546 q = qn;
6547 let mut kn = e.uninit(t * n_head_kv * head_dim)?;
6548 e.rms_norm(
6549 &k,
6550 fa.k_norm.float_data(),
6551 &mut kn,
6552 head_dim,
6553 t * n_head_kv,
6554 eps,
6555 )?;
6556 k = kn;
6557 e.rope_neox(
6558 &mut q, pos_d, head_dim, rope_dims, n_head, t, rope_base, 1.0,
6559 )?;
6560 e.rope_neox(
6561 &mut k, pos_d, head_dim, rope_dims, n_head_kv, t, rope_base, 1.0,
6562 )?;
6563
6564 // Per-row append + attend: row r sees rows 0..r in KV (causal within the
6565 // draft), each through the b_n=1 serving kernels at its own t_kv.
6566 let q_dim = n_head * head_dim;
6567 let kv_dim = n_head_kv * head_dim;
6568 let mut attn = e.uninit(t * q_dim)?;
6569 let (kdk, kdv, ktb, vtb, len0, kv_local) = {
6570 let kvl = cache.kv[il].as_ref().unwrap();
6571 // [2T] interleaved k,v base pointers: entry pair z serves row z of
6572 // the batched twins; the per-row fallback reads pair 0 (same cache
6573 // for every row of one layer). Graph mode reads the ctx table.
6574 let local: Option<CudaSlice<u64>> = match graph_cap {
6575 Some(_) => None,
6576 None => {
6577 let s = &e.gpu.stream();
6578 let (pk, _g) = kvl.k.device_ptr(s);
6579 let (pv, _g2) = kvl.v.device_ptr(s);
6580 let mut tbl = Vec::with_capacity(2 * t);
6581 for _ in 0..t {
6582 tbl.push(pk as u64);
6583 tbl.push(pv as u64);
6584 }
6585 Some(e.htod_u64(&tbl)?)
6586 }
6587 };
6588 (
6589 kvl.kv_dim_k,
6590 kvl.kv_dim_v,
6591 kvl.k_tok_bytes,
6592 kvl.v_tok_bytes,
6593 kvl.len,
6594 local,
6595 )
6596 };
6597 let (kv_tbl, kv_off): (&CudaSlice<u64>, usize) = match graph_cap {
6598 Some((tb, off, _)) => (tb, off),
6599 None => (kv_local.as_ref().expect("built above"), 0),
6600 };
6601 // Slice 4 (fa/append rows — see dspark_fa_rows_on): the whole per-row
6602 // section batches into the z-batched serving twins when every row of
6603 // this round takes the v4-seqs arm on ONE fa_split_keys rung. Both
6604 // guards are evaluated at the round's FIRST and LAST t_kv — the
6605 // eligibility window (vec floor .. v4 max) and each split-ladder rung
6606 // are intervals in t_kv, so ends-inside means all-inside (the straddle
6607 // law). Appending all T rows before any attend is read-equivalent to
6608 // the interleaved order: row r's walk reads keys 0..len0+r only, and
6609 // rows > r land at slots it never touches; every written cache row is
6610 // the per-token appender's exact warp program (kernel-check pinned).
6611 let t_kv_first = len0 + 1;
6612 let t_kv_last = len0 + t;
6613 let rows_batched = t >= 2
6614 && seqs_append
6615 && batch_fa_on
6616 && dspark_fa_rows_on()
6617 // the z-batched twins read stacked rows at the CACHE's kv dims;
6618 // the projection stack is [T, n_head_kv*head_dim] — they must be
6619 // the same stride or row z misaligns (true for this family; the
6620 // guard keeps any asymmetric-kv model on the per-row loop).
6621 && kdk == kv_dim
6622 && kdv == kv_dim
6623 && crate::fa_seqs_eligible(t_kv_first, head_dim_global)
6624 && crate::fa_seqs_eligible(t_kv_last, head_dim_global)
6625 && crate::fa_split_keys(t_kv_first, cfg.n_head_kv as usize)
6626 == crate::fa_split_keys(t_kv_last, cfg.n_head_kv as usize);
6627 // Sizing: eager = exact round bound; graph mode = the rung end (stride +
6628 // grid only — bytes proven equal above). Capture-time invariants refuse
6629 // loudly rather than bake a divergent body.
6630 let (size_kv_max, sp) = match graph_cap {
6631 Some((_, _, rung)) => {
6632 if !rows_batched {
6633 return Err(format!(
6634 "fa graph capture: layer {il} round is not batchable \
6635 (t_kv {t_kv_first}..{t_kv_last}) — the per-row fallback \
6636 must never be captured"
6637 )
6638 .into());
6639 }
6640 let sp_r = crate::fa_split_keys(rung, cfg.n_head_kv as usize);
6641 if t_kv_last > rung
6642 || sp_r != crate::fa_split_keys(t_kv_last, cfg.n_head_kv as usize)
6643 {
6644 return Err(format!(
6645 "fa graph capture: rung {rung} does not cover round \
6646 t_kv {t_kv_first}..{t_kv_last} on one split ladder step"
6647 )
6648 .into());
6649 }
6650 (rung, sp_r)
6651 }
6652 None => (
6653 t_kv_last,
6654 crate::fa_split_keys(t_kv_last, cfg.n_head_kv as usize),
6655 ),
6656 };
6657 if let Some((_, ctr)) = stream {
6658 // STREAM ARM (2b): one batched dc append + the multi-row dc attention
6659 // — the generic stream arm's exact shape (rows kernels are pinned
6660 // byte-identical to the per-row programs by kernel-check). Host len
6661 // stays a stale lower bound; the burst drain reconciles it.
6662 let kvl = cache.kv[il].as_mut().unwrap();
6663 e.append_kv_quantized_rows_dc(
6664 &k,
6665 &v,
6666 &mut kvl.k,
6667 &mut kvl.v,
6668 ctr,
6669 t,
6670 kdk,
6671 kdv,
6672 ktb,
6673 vtb,
6674 Engine::kv_fp8_on(),
6675 )?;
6676 let upper = (kvl.len + t + 64).min(cache.max_ctx);
6677 let k_view = e.view_u8(&kvl.k, upper * ktb);
6678 let v_view = e.view_u8(&kvl.v, upper * vtb);
6679 e.fa_decode_rows_dc(
6680 &q, &k_view, &v_view, &mut attn, head_dim, n_head, n_head_kv, ctr, upper,
6681 t, scale, ktb, vtb, 0, false,
6682 )?;
6683 } else if rows_batched {
6684 e.append_kv_quantized_seqs(
6685 &k,
6686 &v,
6687 &kv_tbl.slice(kv_off..kv_off + 2 * t),
6688 pos_d,
6689 t,
6690 kdk,
6691 kdv,
6692 ktb,
6693 vtb,
6694 )?;
6695 if graph_cap.is_none() {
6696 cache.kv[il].as_mut().unwrap().len += t;
6697 }
6698 e.fa_decode_batch_seqs_v4(
6699 &q,
6700 &kv_tbl.slice(kv_off..kv_off + 2 * t),
6701 pos_d,
6702 &mut attn,
6703 head_dim,
6704 n_head,
6705 n_head_kv,
6706 t,
6707 size_kv_max,
6708 scale,
6709 sp,
6710 ktb,
6711 vtb,
6712 )?;
6713 } else {
6714 if pos_rows.is_none() {
6715 // Stream-aware for symmetry with pos_d (the stream FA arm rides
6716 // the dc rows kernels above and never reaches this fallback).
6717 *pos_rows = Some(match stream {
6718 Some((_, ctr)) => (0..t)
6719 .map(|r| {
6720 let mut b = e.alloc_uninit::<i32>(1)?;
6721 e.i32_copy_add(ctr, &mut b, r as i32)?;
6722 Ok(b)
6723 })
6724 .collect::<Result<_, Box<dyn std::error::Error>>>()?,
6725 None => (0..t)
6726 .map(|r| e.htod_i32(&[(pos0 + r) as i32]))
6727 .collect::<Result<_, _>>()?,
6728 });
6729 }
6730 let pos_rows = pos_rows.as_ref().unwrap();
6731 for r in 0..t {
6732 // Owned per-row scratch: the b_n=1 kernels take packed batch buffers
6733 // whose row 0 is this row (arithmetic-free materialization copies,
6734 // same as decode's per-seq fallback arm).
6735 let mut k_row = e.uninit(kv_dim)?;
6736 e.dtod_copy_view(&k.slice(r * kv_dim..(r + 1) * kv_dim), &mut k_row)?;
6737 let mut v_row = e.uninit(kv_dim)?;
6738 e.dtod_copy_view(&v.slice(r * kv_dim..(r + 1) * kv_dim), &mut v_row)?;
6739 let pos_row = &pos_rows[r];
6740 let kvl = cache.kv[il].as_mut().unwrap();
6741 if seqs_append {
6742 e.append_kv_quantized_seqs(
6743 &k_row,
6744 &v_row,
6745 &kv_tbl.slice(kv_off..kv_off + 2),
6746 pos_row,
6747 1,
6748 kdk,
6749 kdv,
6750 ktb,
6751 vtb,
6752 )?;
6753 kvl.len += 1;
6754 } else {
6755 e.append_kv_quantized_view(
6756 &k_row.slice(0..kv_dim),
6757 &v_row.slice(0..kv_dim),
6758 &mut kvl.k,
6759 &mut kvl.v,
6760 kvl.len,
6761 kvl.kv_dim_k,
6762 kvl.kv_dim_v,
6763 kvl.k_tok_bytes,
6764 kvl.v_tok_bytes,
6765 Engine::kv_fp8_on(),
6766 )?;
6767 kvl.len += 1;
6768 }
6769 let t_kv = kvl.len;
6770 let mut q_row = e.uninit(q_dim)?;
6771 e.dtod_copy_view(&q.slice(r * q_dim..(r + 1) * q_dim), &mut q_row)?;
6772 let mut a_row = e.uninit(q_dim)?;
6773 if batch_fa_on && crate::fa_seqs_eligible(t_kv, head_dim_global) {
6774 let sp0_r = crate::fa_split_keys(t_kv, cfg.n_head_kv as usize);
6775 e.fa_decode_batch_seqs_v4(
6776 &q_row,
6777 &kv_tbl.slice(kv_off..kv_off + 2),
6778 pos_row,
6779 &mut a_row,
6780 head_dim,
6781 n_head,
6782 n_head_kv,
6783 1,
6784 t_kv,
6785 scale,
6786 sp0_r,
6787 ktb,
6788 vtb,
6789 )?;
6790 } else {
6791 let k_view = e.view_u8(&kvl.k, t_kv * kvl.k_tok_bytes);
6792 let v_view = e.view_u8(&kvl.v, t_kv * kvl.v_tok_bytes);
6793 let mut a_view = a_row.slice_mut(0..q_dim);
6794 e.fa_decode_kvmod_view(
6795 &q_row.slice(0..q_dim),
6796 &k_view,
6797 &v_view,
6798 &mut a_view,
6799 head_dim,
6800 n_head,
6801 n_head_kv,
6802 t_kv,
6803 scale,
6804 kvl.k_tok_bytes,
6805 kvl.v_tok_bytes,
6806 Engine::kv_fp8_on(),
6807 )?;
6808 }
6809 e.dtod_copy_into(&a_row, &mut attn, r * q_dim)?;
6810 }
6811 }
6812
6813 // Output gate (element-wise) + o-proj at m=T.
6814 let attn_g = match &gate {
6815 Some(g) => {
6816 let n = t * q_dim;
6817 let mut gsig = e.uninit(n)?;
6818 e.sigmoid(g, &mut gsig, n)?;
6819 let mut ag = e.uninit(n)?;
6820 e.mul(&attn, &gsig, &mut ag, n)?;
6821 ag
6822 }
6823 None => attn,
6824 };
6825 e.matmul(&fa.wo, &attn_g, t)?
6826 }
6827 };
6828
6829 // ---- residual add + post_attn_norm + FFN at m=T (serving dispatch verbatim) ----
6830 let pnorm = layer.post_attn_norm.float_data();
6831 let mut x1 = e.uninit(t * n_embd)?;
6832 let mut zn = e.uninit(t * n_embd)?;
6833 e.add_rms_norm(x, &mixed, pnorm, &mut x1, &mut zn, n_embd, t, eps)?;
6834 let ffn_out = match &layer.ffn {
6835 crate::hybrid::Ffn::Dense {
6836 ffn_gate,
6837 ffn_up,
6838 ffn_down,
6839 } => {
6840 assert!(
6841 self.cfg.m3.is_none(),
6842 "qwen35 t-parallel verify: M3 swigluoai FFN not yet batched"
6843 );
6844 self.qwen35_tparallel_dense_ffn(e, ffn_gate, ffn_up, ffn_down, &zn, t, n_embd)?
6845 }
6846 crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il_zq8(e, m, &zn, None, t, il as u16)?,
6847 };
6848 let mut x2 = e.uninit(t * n_embd)?;
6849 e.add(&x1, &ffn_out, &mut x2, t * n_embd)?;
6850 // dspark drafter tap (no-op when no sink armed): post-layer residual verify rows
6851 self.dflash_tap(e, cache, il, &x2, t)?;
6852 Ok(x2)
6853 }
6854
6855 /// ONE t-parallel LINEAR layer (attn_norm + gdn mixer + post_attn_norm + FFN + tap) —
6856 /// the exact body the old in-loop Linear arm ran, extracted so the eager walk and the
6857 /// slice-3 captured segments execute the SAME code (a second copy is how dispatch
6858 /// mirrors drift — the verify_layers extraction lesson). Two deliberate changes, both
6859 /// bit-identical by construction:
6860 /// - the gdn ping-pong host swap moves from per-row to ONE end-of-body swap (t odd):
6861 /// the device sequence is driven entirely by the 6-entry pointer table, which
6862 /// already encodes both parities; the ckpt stash reads name row r's out buffer
6863 /// directly (r even -> alt handle, odd -> canonical) — the same physical bytes the
6864 /// legacy post-swap clone read.
6865 /// - `stash` (slice-3 ctx): persistent per-layer slabs written by copy_into instead of
6866 /// per-row clone_dtod allocs — same bytes, capture-legal (no per-round host objects).
6867 /// `table_src` = (persistent pointer table, offset) when the ctx owns the tables;
6868 /// None builds the per-verify table exactly as before.
6869 #[allow(clippy::too_many_arguments)]
6870 fn qwen35_tparallel_linear_layer(
6871 &self,
6872 e: &Engine,
6873 il: usize,
6874 x: &CudaSlice<f32>,
6875 t: usize,
6876 cache: &mut Cache,
6877 mut ckpt: Option<&mut VerifyCkpt>,
6878 stash: Option<(&mut CudaSlice<f32>, &mut CudaSlice<f32>)>,
6879 table_src: Option<(&CudaSlice<u64>, usize)>,
6880 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6881 use cudarc::driver::DevicePtr;
6882 let cfg = &self.cfg;
6883 let n_embd = cfg.n_embd as usize;
6884 let eps = cfg.rms_eps;
6885 let layer = &self.layers[il];
6886 let Mixer::Linear(la) = &layer.mixer else {
6887 return Err("qwen35_tparallel_linear_layer on a non-linear layer".into());
6888 };
6889 // ---- attn_norm + q8_1 quantize at m=T (row-indexed == per-row) ----
6890 let anorm = layer.attn_norm.float_data();
6891 let mut xn = e.uninit(t * n_embd)?;
6892 e.rms_norm(x, anorm, &mut xn, n_embd, t, eps)?;
6893 let (hq, hd) = e.quantize_q8_1(&xn, t, n_embd)?;
6894
6895 let geometry = la.geometry;
6896 let d_state = geometry.key_head_dim as usize;
6897 let num_k = geometry.key_heads as usize;
6898 let num_v = geometry.value_heads as usize;
6899 let d_conv = geometry.conv_kernel as usize;
6900 let key_dim = d_state * num_k;
6901 let value_dim = geometry.value_head_dim as usize * num_v;
6902 let conv_dim = key_dim * 2 + value_dim;
6903 let gdn_scale = 1.0 / (d_state as f32).sqrt();
6904
6905 // ---- batched projections: one weight read for all T rows ----
6906 // GROUP-4 twin (trunk-kernels slice C): the whole 4-tuple in ONE launch, bit-identical
6907 // per (tensor, token, row) to the four singles; refused (layout/tier) or
6908 // MEMRA_TK_GDN_GROUP=0 -> the singles chain byte-for-byte.
6909 let (qkv_mixed, z, beta_raw, alpha) = match e.matmul_decode_exact_group4_pre(
6910 [&la.wqkv, &la.wqkv_gate, &la.ssm_beta, &la.ssm_alpha],
6911 &hq,
6912 &hd,
6913 t,
6914 )? {
6915 Some(mut g4) => {
6916 let alpha = g4.pop().unwrap();
6917 let beta_raw = g4.pop().unwrap();
6918 let z = g4.pop().unwrap();
6919 let qkv_mixed = g4.pop().unwrap();
6920 (qkv_mixed, z, beta_raw, alpha)
6921 }
6922 None => (
6923 e.matmul_pre(&la.wqkv, &hq, &hd, &xn, t)?,
6924 e.matmul_pre(&la.wqkv_gate, &hq, &hd, &xn, t)?,
6925 e.matmul_pre(&la.ssm_beta, &hq, &hd, &xn, t)?,
6926 e.matmul_pre(&la.ssm_alpha, &hq, &hd, &xn, t)?,
6927 ),
6928 };
6929 let beta_w = la.ssm_beta.out_features();
6930 let alpha_w = la.ssm_alpha.out_features();
6931 let qkv_w = la.wqkv.out_features();
6932
6933 // ---- per-row state chain through the b_n=1 serving kernels ----
6934 // 6-entry alternating pointer table expresses the ping-pong without a rebuild per
6935 // row: even rows scan s0 -> s1, odd rows s1 -> s0.
6936 let table_local: Option<CudaSlice<u64>> = match table_src {
6937 Some(_) => None,
6938 None => {
6939 let rl = cache.recur[il].as_ref().unwrap();
6940 let s = &e.gpu.stream();
6941 let (pc, _g0) = rl.conv_state.device_ptr(s);
6942 let (p0, _g1) = rl.ssm_state.device_ptr(s);
6943 let (p1, _g2) = rl.ssm_state_alt.device_ptr(s);
6944 Some(e.htod_u64(&[
6945 pc as u64, p0 as u64, p1 as u64, pc as u64, p1 as u64, p0 as u64,
6946 ])?)
6947 }
6948 };
6949 let (table, toff): (&CudaSlice<u64>, usize) = match table_src {
6950 Some((tb, off)) => (tb, off),
6951 None => (table_local.as_ref().unwrap(), 0),
6952 };
6953 let mut o_all = e.uninit(t * value_dim)?;
6954 let mut col_states: Option<Vec<(CudaSlice<f32>, CudaSlice<f32>)>> =
6955 if ckpt.is_some() && stash.is_none() && t >= 2 {
6956 Some(Vec::with_capacity(t - 1))
6957 } else {
6958 None
6959 };
6960 let mut stash = stash;
6961 // Per-row scratch reused across rows (uninit is cheap but not free at
6962 // 48 layers x T rows); row inputs/outputs pass as VIEWS into the packed
6963 // [T, ...] buffers — zero arithmetic-free copies in this loop.
6964 let mut conv_out = e.uninit(conv_dim)?;
6965 let mut q_l2 = e.uninit(value_dim)?;
6966 let mut k_l2 = e.uninit(value_dim)?;
6967 let mut v_gd = e.uninit(value_dim)?;
6968 let mut beta_b = e.uninit(num_v)?;
6969 let mut g_log = e.uninit(num_v)?;
6970 for r in 0..t {
6971 let base = toff + if r % 2 == 0 { 0 } else { 3 };
6972 let conv_view = table.slice(base..base + 1);
6973 let in_view = table.slice(base + 1..base + 2);
6974 let out_view = table.slice(base + 2..base + 3);
6975 e.ssm_conv1d_fused_decode_b_view(
6976 &qkv_mixed.slice(r * qkv_w..(r + 1) * qkv_w),
6977 &conv_view,
6978 la.ssm_conv1d.float_data(),
6979 &mut conv_out,
6980 conv_dim,
6981 d_conv,
6982 1,
6983 )?;
6984 e.gdn_prep_decode_b_view(
6985 &conv_out,
6986 &beta_raw.slice(r * beta_w..(r + 1) * beta_w),
6987 &alpha.slice(r * alpha_w..(r + 1) * alpha_w),
6988 la.ssm_dt.float_data(),
6989 la.ssm_a.float_data(),
6990 &mut q_l2,
6991 &mut k_l2,
6992 &mut v_gd,
6993 &mut beta_b,
6994 &mut g_log,
6995 d_state,
6996 num_v,
6997 num_k,
6998 key_dim,
6999 eps,
7000 conv_dim,
7001 1,
7002 )?;
7003 let mut o_row = o_all.slice_mut(r * value_dim..(r + 1) * value_dim);
7004 e.gdn_scan_s128_batched_view(
7005 &q_l2, &k_l2, &v_gd, &g_log, &beta_b, &in_view, &out_view, &mut o_row, num_v, 1,
7006 gdn_scale,
7007 )?;
7008 if r + 1 < t {
7009 // Row r's out buffer: even rows write s1 (the alt handle — no swaps ran),
7010 // odd rows write s0 — the same physical state the legacy post-swap
7011 // canonical clone read.
7012 let rl = cache.recur[il]
7013 .as_ref()
7014 .ok_or("qwen35 linear verify layer has no recurrent state")?;
7015 let ssm_src = if r % 2 == 0 {
7016 &rl.ssm_state_alt
7017 } else {
7018 &rl.ssm_state
7019 };
7020 match stash.as_mut() {
7021 Some((conv_slab, ssm_slab)) => {
7022 // BOTH stash reads go through the pointer table at run time: the
7023 // ssm handles ping-pong between rounds, and the ctx (with its
7024 // captured graphs) outlives the Cache — a fresh generation's
7025 // conv/ssm buffers land at new addresses that only the per-round
7026 // table refresh knows. A baked direct copy would read freed
7027 // memory (parity was the slice-3 smoke divergence; cache
7028 // lifetime is the cross-generation twin).
7029 e.copy_indirect_src_f32(
7030 &conv_view,
7031 conv_slab,
7032 r * conv_dim * (d_conv - 1),
7033 conv_dim * (d_conv - 1),
7034 )?;
7035 // The ssm handles PING-PONG between rounds: a captured direct
7036 // copy would bake the capture-time physical buffer and read the
7037 // wrong parity after any odd-vt round (the slice-3 smoke
7038 // divergence). Read the src address from row r's OUT table
7039 // entry at run time — the same entry the scan just wrote.
7040 e.copy_indirect_src_f32(
7041 &out_view,
7042 ssm_slab,
7043 r * d_state * d_state * num_v,
7044 d_state * d_state * num_v,
7045 )?;
7046 }
7047 None => {
7048 if let Some(states) = col_states.as_mut() {
7049 states.push((e.clone_dtod(&rl.conv_state)?, e.clone_dtod(ssm_src)?));
7050 }
7051 }
7052 }
7053 }
7054 }
7055 // ONE end-of-body parity swap (t odd) — the legacy loop swapped per row; the net
7056 // handle motion is identical and the device sequence never read the handles.
7057 if t % 2 == 1 {
7058 let rl = cache.recur[il].as_mut().unwrap();
7059 std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
7060 }
7061 if let (Some(checkpoint), Some(states)) = (ckpt.as_deref_mut(), col_states) {
7062 checkpoint.cols[il] = Some(states);
7063 }
7064
7065 // ---- batched gated norm + out-projection at m=T ----
7066 let mixed = if e.uses_q8_1_fast(&la.ssm_out) {
7067 let (gq, gd) = e.gated_rmsnorm_q8_1(
7068 &o_all,
7069 la.ssm_norm.float_data(),
7070 &z,
7071 d_state,
7072 t * num_v,
7073 eps,
7074 )?;
7075 let g0 = e.zeros(0)?;
7076 e.matmul_pre(&la.ssm_out, &gq, &gd, &g0, t)?
7077 } else {
7078 let mut gn = e.uninit(t * value_dim)?;
7079 e.gated_rmsnorm(
7080 &o_all,
7081 la.ssm_norm.float_data(),
7082 &z,
7083 &mut gn,
7084 d_state,
7085 t * num_v,
7086 eps,
7087 )?;
7088 e.matmul(&la.ssm_out, &gn, t)?
7089 };
7090
7091 // ---- residual add + post_attn_norm + FFN at m=T (serving dispatch verbatim) ----
7092 let pnorm = layer.post_attn_norm.float_data();
7093 let mut x1 = e.uninit(t * n_embd)?;
7094 let mut zn = e.uninit(t * n_embd)?;
7095 e.add_rms_norm(x, &mixed, pnorm, &mut x1, &mut zn, n_embd, t, eps)?;
7096 let ffn_out = match &layer.ffn {
7097 crate::hybrid::Ffn::Dense {
7098 ffn_gate,
7099 ffn_up,
7100 ffn_down,
7101 } => {
7102 assert!(
7103 self.cfg.m3.is_none(),
7104 "qwen35 t-parallel verify: M3 swigluoai FFN not yet batched"
7105 );
7106 self.qwen35_tparallel_dense_ffn(e, ffn_gate, ffn_up, ffn_down, &zn, t, n_embd)?
7107 }
7108 crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il_zq8(e, m, &zn, None, t, il as u16)?,
7109 };
7110 let mut x2 = e.uninit(t * n_embd)?;
7111 e.add(&x1, &ffn_out, &mut x2, t * n_embd)?;
7112 // dspark drafter tap (no-op when no sink armed): post-layer residual verify rows
7113 self.dflash_tap(e, cache, il, &x2, t)?;
7114 Ok(x2)
7115 }
7116
7117 /// PP-N STAGE SUBGRAPH of the verify trunk: layers `[lo, hi)` of `decode_step_t_core_stream`'s
7118 /// walk, verbatim. Enters with a MATERIALIZED `[T, n_embd]` residual (no pending fusion pair
7119 /// carried in from outside the range) and exits with the range's final residual materialized
7120 /// (the trailing add executed) — exactly the `decode_layers_eager(lo, hi)` contract, T rows
7121 /// instead of one.
7122 ///
7123 /// EXTRACTED (lane/pp2-spec 2026-08-06) rather than duplicated: `decode_step_t_core_stream` IS
7124 /// the single funnel every verify forward reaches, and its per-layer dispatch MIRRORING (norm
7125 /// fusion per layer, the t>=3/spec_m2 batched-linear window, the fused-q8 FFN chain, the
7126 /// decode-exact projections) is what makes verify bit-identical to eager decode. A second copy
7127 /// for the split arm is how those mirrors drift apart on the next lever. The unsplit body now
7128 /// calls this with `(0, n_layers)`, so the whole-trunk path and every stage range run the SAME
7129 /// code — there is no "split version" of the verify math.
7130 ///
7131 /// Bit-identity of a cut rests on the same kernel-check-pinned identity the eager arm's cut
7132 /// does — `add_rms_norm_q8_1 == add then rms_norm_q8_1` at nrows=T — because the ONLY thing a
7133 /// fence changes is that the cross-layer fusion carry breaks at `hi-1` and is re-materialized
7134 /// as an explicit `add`. `decode-batch-gate --mode ppspec` verifies end-to-end on real weights.
7135 #[allow(clippy::too_many_arguments)]
7136 fn verify_layers(
7137 &self,
7138 e: &Engine,
7139 mut x: CudaSlice<f32>,
7140 lo: usize,
7141 hi: usize,
7142 pos_d: &CudaSlice<i32>,
7143 pos0: usize,
7144 t: usize,
7145 cache: &mut Cache,
7146 mut ckpt: Option<&mut VerifyCkpt>,
7147 stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
7148 graphs: Option<&mut DsparkVerifyGraphs>,
7149 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7150 if self.sliding_gated_moe_batch_program() {
7151 if stream.is_some() {
7152 return Err(
7153 "step35 has no ROUND-STREAM verify arm (the device-counter _dc twins \
7154 cannot express the SWA offset KV view)"
7155 .into(),
7156 );
7157 }
7158 return self.step35_verify_batch_layers(e, x, lo, hi, pos0, t, cache);
7159 }
7160 if self.batched_serving_numeric_class() {
7161 return self.qwen35_verify_batch_layers(
7162 e,
7163 x,
7164 lo,
7165 hi,
7166 pos0,
7167 t,
7168 cache,
7169 ckpt.take(),
7170 stream,
7171 graphs,
7172 );
7173 }
7174 let n_embd = self.cfg.n_embd as usize;
7175 let eps = self.cfg.rms_eps;
7176 // CROSS-LAYER ADD+NORM FUSION (lane/vt-fixes fix 2, mirroring decode_step_h's
7177 // launch-arc form): layer il's post-FFN residual add (x2 = x1 + ffn_out) and layer
7178 // il+1's attn_norm(+quantize) are consecutive row-wise ops — ONE add_rms_norm_q8_1
7179 // launch at nrows=t does all three (bit-identity pinned by the T-row kernel-check
7180 // arms). Carry the un-added (x1, ffn_out) pair; the fused launch materializes x2 (the
7181 // residual the next layer needs) as its `res` output. Falls back to the separate add
7182 // when the next layer is off the fused-q8 path.
7183 let mut pending: Option<(CudaSlice<f32>, CudaSlice<f32>)> = None;
7184 for il in lo..hi {
7185 let layer = &self.layers[il];
7186 // DISPATCH-MIRRORED attn-input RMSNorm (FP-order lesson #8): eager decode fuses the
7187 // 1024-thread rms_norm_q8_1 ONLY when every mixer projection is q8_1-fast; layers with
7188 // Float projections (ssm_beta/ssm_alpha on layers 1/2/4 of the 9B NVFP4 GGUF) take the
7189 // UNFUSED 256-thread rms_norm. The verify norm must mirror that PER-LAYER choice —
7190 // blockDim changes the sum-of-squares reduce order, and the ULP shift amplifies through
7191 // the GDN recurrence into argmax flips (measured: 9B text prompt, 1 ULP at layer 2 ->
7192 // 2.3e-1 logit maxdiff at the head -> K=1..8 divergence at a 0.03-margin token).
7193 let mixer_fast = self.mixer_in_q8_1_fast(e, &layer.mixer);
7194 let norm_fused = std::env::var("MEMRA_NO_FUSE_NORMQ").is_err() && mixer_fast;
7195 // BATCHED EPILOGUE RE-FUSE (lane/vt-fixes fix 2, 2026-08-03): when the norm is
7196 // dispatch-fused AND every consumer of `h` reads only its q8_1 form (Full mixer:
7197 // projections only; Linear mixer: the batched arm — the per-column fallback needs
7198 // f32 h), emit the attn-input norm DIRECTLY as q8_1 via `rms_norm_q8_1` at nrows=t
7199 // (row-indexed kernel — the T-row launch is the per-row m=1 program, kernel-check
7200 // pins bit-identity vs rms_norm_decode -> quantize_q8_1). Kills the standalone
7201 // quantize launch(es) + the f32 h HBM round-trip that decode never pays.
7202 // step35 (Full mixer) is the third case that needs f32 `h`: its verify arm is a
7203 // per-ROW replay of the eager decode mixer, whose `pre_q` contract is a single row —
7204 // a T-row q8_1 pair cannot be handed to it, and re-deriving per-row q8_1 from the f32
7205 // rows is exactly the dispatch being mirrored. Keep step35 on the unfused arm.
7206 let lin_q8_only = match &layer.mixer {
7207 Mixer::Linear(la) => {
7208 (t >= 3 || (t == 2 && spec_m2())) && e.uses_q8_1_fast(&la.ssm_out)
7209 }
7210 Mixer::Full(_) if self.sliding_gated_moe_batch_program() => false,
7211 _ => true,
7212 };
7213 // NOTE decode.rs's take()-first lesson: take the pending pair BEFORE branching so
7214 // a non-fused layer still performs the residual add.
7215 let taken = pending.take();
7216 let (h, h_q8) = if norm_fused && lin_q8_only {
7217 let pair = match taken {
7218 // fused add + attn_norm + q8_1: ONE launch resolves the carried residual
7219 // AND emits this layer's mixer input pre-quantized. res -> x2 (= new x).
7220 Some((x1p, f1p)) => {
7221 let mut x2 = vbuf(e, t * n_embd)?; // fully written (res output)
7222 let p = e.add_rms_norm_q8_1(
7223 &x1p,
7224 &f1p,
7225 layer.attn_norm.float_data(),
7226 &mut x2,
7227 n_embd,
7228 t,
7229 eps,
7230 )?;
7231 x = x2;
7232 p
7233 }
7234 None => e.rms_norm_q8_1(&x, layer.attn_norm.float_data(), n_embd, t, eps)?,
7235 };
7236 (e.zeros(0)?, Some(pair)) // h unused on this path (q8-only consumers)
7237 } else {
7238 if let Some((x1p, f1p)) = taken {
7239 let mut x2 = vbuf(e, t * n_embd)?; // fully written by add
7240 e.add(&x1p, &f1p, &mut x2, t * n_embd)?;
7241 x = x2;
7242 }
7243 let mut h = vbuf(e, t * n_embd)?; // fully written by either rms_norm arm
7244 if norm_fused {
7245 e.rms_norm_decode(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
7246 } else {
7247 e.rms_norm(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
7248 }
7249 (h, None)
7250 };
7251 let h_q8_ref = h_q8.as_ref().map(|(q, d)| (q, d));
7252
7253 let mixed = match &layer.mixer {
7254 Mixer::Full(fa) => self.full_attn_verify(
7255 e,
7256 fa,
7257 &h,
7258 h_q8_ref,
7259 pos_d,
7260 t,
7261 cache,
7262 il,
7263 stream.map(|(_, c)| c),
7264 )?,
7265 Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
7266 Mixer::Linear(la) => {
7267 // BATCHED linear verify (2026-07-03, the MTP-profit lever): one T-token pass —
7268 // batched projections (weight read ONCE, hits the m=2-4 weight-resident matvec),
7269 // carried-state conv (ssm_conv1d_tm_state), GDN prep on the prefill kernels, and
7270 // ONE gdn_scan whose internal sequential t-loop is the SAME recurrence as T
7271 // chained T=1 steps (bit-identical). Falls back to the sequential per-column
7272 // chain when T < d_conv-1 (conv ring update needs T >= pad) — or when ANY
7273 // projection is off the q8_1 fast path: matmul_decode_exact would route a Float
7274 // tensor to cuBLAS at m=t (different FP accumulation than eager's per-token
7275 // GEMV), so mixed-dtype layers stay on the eager-identical per-column chain.
7276 // MEMRA_SPEC_M2 (lane/spec-m2): the t==2 batch rides the same arm — the conv
7277 // wrapper handles t<pad with a pure-copy ring rebuild; see spec_m2() header.
7278 if (t >= 3 || (t == 2 && spec_m2()))
7279 && mixer_fast
7280 && e.uses_q8_1_fast(&la.ssm_out)
7281 {
7282 let want = ckpt.is_some();
7283 let (out, stash) =
7284 self.linear_attn_verify_t(e, la, &h, h_q8_ref, t, cache, il, want)?;
7285 if let (Some(ck), Some(st)) = (ckpt.as_deref_mut(), stash) {
7286 ck.gdn[il] = Some(st);
7287 }
7288 out
7289 } else {
7290 let mut out = vbuf(e, t * n_embd)?; // every col written by copy_into
7291 let mut col_states: Option<Vec<(CudaSlice<f32>, CudaSlice<f32>)>> =
7292 if ckpt.is_some() && t >= 2 {
7293 Some(Vec::with_capacity(t - 1))
7294 } else {
7295 None
7296 };
7297 for col in 0..t {
7298 let mut h_col = vbuf(e, n_embd)?; // fully written by copy_view_into
7299 let src = h.slice(col * n_embd..(col + 1) * n_embd);
7300 e.copy_view_into(&mut h_col, 0, &src, n_embd)?;
7301 let m_col = self.linear_attn_decode(e, la, &h_col, cache, il)?;
7302 e.copy_into(&mut out, col * n_embd, &m_col, n_embd)?;
7303 // REPLAY-FREE ckpt: clone the chain's ACTUAL state after this column
7304 // (pure dtod — cannot change any computed value). Last column skipped:
7305 // rebuild targets are j <= t-1 columns.
7306 if let Some(cs) = col_states.as_mut() {
7307 if col + 1 < t {
7308 let rl = cache.recur[il].as_ref().unwrap();
7309 cs.push((
7310 e.clone_dtod(&rl.conv_state)?,
7311 e.clone_dtod(&rl.ssm_state)?,
7312 ));
7313 }
7314 }
7315 }
7316 if let (Some(ck), Some(cs)) = (ckpt.as_deref_mut(), col_states) {
7317 // ReplaySSM-assessment instrumentation (2026-07-30): the
7318 // per-column clones are the only true state snapshots left in
7319 // the verify (the batched path stashes INPUTS and replays).
7320 if std::env::var("MEMRA_SPEC_STATS").as_deref() == Ok("1") {
7321 static ONCE: std::sync::Once = std::sync::Once::new();
7322 let bytes: usize =
7323 cs.iter().map(|(c, s)| (c.len() + s.len()) * 4).sum();
7324 ONCE.call_once(|| eprintln!(
7325 "[verify-ckpt] per-column layer il={il}: {} clones, {:.2} MB/layer/round",
7326 cs.len(), bytes as f64 / 1e6));
7327 }
7328 ck.cols[il] = Some(cs);
7329 }
7330 out
7331 }
7332 }
7333 };
7334
7335 // DISPATCH-MIRRORED post-attn norm: eager residual_norm_ffn fuses add+norm+quant
7336 // (1024-thread add_rms_norm_q8_1) only for Dense FFNs whose gate+up are q8_1-fast;
7337 // otherwise (and for MoE) it runs the 256-thread fused add_rms_norm. Mirror per layer.
7338 let ffn_fuse = match &layer.ffn {
7339 crate::hybrid::Ffn::Dense {
7340 ffn_gate, ffn_up, ..
7341 } => {
7342 std::env::var("MEMRA_NO_FUSE_NORMQ").is_err()
7343 && e.uses_q8_1_fast(ffn_gate)
7344 && e.uses_q8_1_fast(ffn_up)
7345 }
7346 crate::hybrid::Ffn::Moe(_) => false,
7347 };
7348 // BATCHED EPILOGUE RE-FUSE (lane/vt-fixes fix 2): on the ffn_fuse path (Dense,
7349 // gate+up q8_1-fast, non-M3) the FFN input is emitted DIRECTLY as q8_1 by ONE
7350 // add_rms_norm_q8_1 launch at nrows=t (row-indexed kernel: the T-row launch is the
7351 // per-row m=1 program; kernel-check pins bit-identity vs the unfused
7352 // add_f32 -> rms_norm_decode -> quantize_q8_1 chain at T=2/4/5/8) — replacing the
7353 // add + rms_norm_decode launches AND the dual/singles' internal re-quantize.
7354 // M3's swigluoai must keep the f32 chain (the fused SwiGLU epilogue encodes plain
7355 // SiLU), mirroring residual_norm_ffn's m3 guard on the decode path.
7356 // step35: same guard per LAYER. A dense FFN's clamp is the SHEXP array (upstream's
7357 // one build_ffn serves dense + shared expert, llama-graph.cpp:1751), and verify MUST
7358 // mirror decode's dispatch or spec self-consistency fails.
7359 let dense_lim = self.cfg.clamp_shexp_at(il as u32);
7360 let fuse_q8 = ffn_fuse && self.cfg.m3.is_none() && dense_lim.is_none();
7361 let mut x1 = vbuf(e, t * n_embd)?; // fully written by add / add_rms_norm*
7362 let mut z = e.zeros(0)?; // replaced below on the unfused arms
7363 let z_q8 = if fuse_q8 {
7364 Some(e.add_rms_norm_q8_1(
7365 &x,
7366 &mixed,
7367 layer.post_attn_norm.float_data(),
7368 &mut x1,
7369 n_embd,
7370 t,
7371 eps,
7372 )?)
7373 } else {
7374 let mut zf = vbuf(e, t * n_embd)?; // fully written by rms_norm_decode / add_rms_norm
7375 if ffn_fuse {
7376 e.add(&x, &mixed, &mut x1, t * n_embd)?;
7377 e.rms_norm_decode(
7378 &x1,
7379 layer.post_attn_norm.float_data(),
7380 &mut zf,
7381 n_embd,
7382 t,
7383 eps,
7384 )?;
7385 } else {
7386 e.add_rms_norm(
7387 &x,
7388 &mixed,
7389 layer.post_attn_norm.float_data(),
7390 &mut x1,
7391 &mut zf,
7392 n_embd,
7393 t,
7394 eps,
7395 )?;
7396 }
7397 z = zf;
7398 None
7399 };
7400 // DECODE-EXACT FFN projections: force MMVQ for gate/up/down at any T to match the
7401 // T=1 decode FP accumulation order. At T>=5 the generic matmul/matmul_pre falls to dp4a
7402 // (128-thread, different FP sum order). At T=2-4 the batched MMVQ is already bit-identical.
7403 let ffn_out = match &layer.ffn {
7404 crate::hybrid::Ffn::Dense {
7405 ffn_gate,
7406 ffn_up,
7407 ffn_down,
7408 } => {
7409 let n_ff = ffn_gate.out_features();
7410 if let Some((zq, zd)) = z_q8.as_ref() {
7411 // FUSED CHAIN (fix 2): pre-quantized z feeds the projections; the SwiGLU
7412 // epilogue emits act pre-quantized for ffn_down (silu_mul_scaled_q8_1,
7413 // bit-identical to silu_mul + quantize — kernel-check-pinned) with the
7414 // NVFP4 macro-scales folded (deferred-scale dual: y*s inline == the
7415 // scale_inplace store, value-exact) — the exact m=1 decode epilogue
7416 // structure at nrows=t.
7417 let pair =
7418 match e.matmul_decode_exact_dual_pre(ffn_gate, ffn_up, zq, zd, t)? {
7419 Some(((g, gs), (u, us))) => Some((g, gs, u, us)),
7420 None => None,
7421 };
7422 let (gate, gs, up, us) = match pair {
7423 Some(x4) => x4,
7424 None => (
7425 e.matmul_decode_exact_pre(ffn_gate, zq, zd, t)?,
7426 1.0, // scale already applied inside _pre
7427 e.matmul_decode_exact_pre(ffn_up, zq, zd, t)?,
7428 1.0,
7429 ),
7430 };
7431 if e.uses_q8_1_fast(ffn_down) {
7432 let (aq, ad) = e.silu_mul_scaled_q8_1(&gate, &up, gs, us, t * n_ff)?;
7433 e.matmul_decode_exact_pre(ffn_down, &aq, &ad, t)?
7434 } else {
7435 let mut act = vbuf(e, t * n_ff)?;
7436 e.silu_mul_scaled(&gate, &up, gs, us, &mut act, t * n_ff)?;
7437 e.matmul_decode_exact(ffn_down, &act, t)?
7438 }
7439 } else {
7440 // UNFUSED (pre-fix) chain — MoE-adjacent/M3/off-fast layers, unchanged.
7441 // DUAL gate+up batched twin (lane/verify-economics, 2026-08-02): one launch
7442 // for the pair at t=2..8 — bit-identical per (tensor,token,row) to the two
7443 // singles (kernel-check pins bitwise; MEMRA_SPEC_DUAL_T=0 reverts). None
7444 // (non-NVFP4 / t outside the tier / seam off) -> the two singles, unchanged.
7445 let (gate, up) =
7446 match e.matmul_decode_exact_dual(ffn_gate, ffn_up, &z, t)? {
7447 Some(pair) => pair,
7448 None => (
7449 e.matmul_decode_exact(ffn_gate, &z, t)?,
7450 e.matmul_decode_exact(ffn_up, &z, t)?,
7451 ),
7452 };
7453 let mut act = vbuf(e, t * n_ff)?; // fully written by ffn_act_lim
7454 Self::ffn_act_lim(
7455 e,
7456 &self.cfg,
7457 &gate,
7458 &up,
7459 1.0,
7460 1.0,
7461 dense_lim,
7462 &mut act,
7463 t * n_ff,
7464 )?;
7465 e.matmul_decode_exact(ffn_down, &act, t)?
7466 }
7467 }
7468 crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il(e, m, &z, t, il as u16)?,
7469 };
7470 // CROSS-LAYER fusion: defer this layer's post-FFN residual add — the next layer's
7471 // fused-q8 attn norm folds it in (add_rms_norm_q8_1 == add; rms_norm; quantize,
7472 // kernel-check-pinned at nrows=T). Non-fused next layers add explicitly above.
7473 pending = Some((x1, ffn_out));
7474 }
7475 // RANGE's final add (no next norm INSIDE the range to fuse with; for the
7476 // whole-trunk call that is the last layer, whose next norm is output_norm — f32-out).
7477 if let Some((x1p, f1p)) = pending.take() {
7478 let mut x2 = vbuf(e, t * n_embd)?; // fully written by add
7479 e.add(&x1p, &f1p, &mut x2, t * n_embd)?;
7480 x = x2;
7481 }
7482 Ok(x)
7483 }
7484 /// BATCHED linear-attn verify (T=K+1): the whole layer in ~10 launches instead of T x the
7485 /// T=1 decode chain (T x ~12 launches + T weight reads of the four projections). The GDN
7486 /// recurrence itself is inherently sequential — gdn_scan_s128 runs its internal t-loop with
7487 /// the SAME per-token math as chained T=1 calls (bit-identical state evolution); everything
7488 /// around it (projections, conv, prep, gated norm, out-proj) batches. Advances conv ring +
7489 /// ssm state exactly like T sequential decode steps.
7490 /// `want_stash`: additionally RETAIN the gdn-scan inputs (pure buffer keep-alives, zero extra
7491 /// kernels) so a partial accept can rebuild the state after any column prefix (REPLAY-FREE).
7492 #[allow(clippy::too_many_arguments)]
7493 fn linear_attn_verify_t(
7494 &self,
7495 e: &Engine,
7496 la: &LinearAttnLayer,
7497 h: &CudaSlice<f32>,
7498 h_q8: Option<(&CudaSlice<i8>, &CudaSlice<f32>)>,
7499 t: usize,
7500 cache: &mut Cache,
7501 il: usize,
7502 want_stash: bool,
7503 ) -> Result<(CudaSlice<f32>, Option<GdnStash>), Box<dyn std::error::Error>> {
7504 let cfg = &self.cfg;
7505 let geometry = la.geometry;
7506 let d_state = geometry.key_head_dim as usize;
7507 let num_k = geometry.key_heads as usize;
7508 let num_v = geometry.value_heads as usize;
7509 let d_conv = geometry.conv_kernel as usize;
7510 let key_dim = d_state * num_k;
7511 let conv_dim = key_dim * 2 + geometry.value_head_dim as usize * num_v;
7512 let eps = cfg.rms_eps;
7513 let scale = 1.0 / (d_state as f32).sqrt();
7514
7515 // DECODE-EXACT projections: matmul_decode_exact forces the MMVQ (warp-per-row, 32-thread)
7516 // accumulation order for EVERY m, matching the T=1 decode path bit-for-bit. The generic
7517 // `matmul` at m>=5 falls to dp4a (128-thread, two-level reduce) which has a different FP
7518 // sum order — ULP differences propagate through gdn_scan and flip argmax on the 27B.
7519 // Q8 TRUNK-FUSION at T=1 (35B: wqkv+wqkv_gate both Q8_0): one fused2 launch, bit-identical
7520 // per (tensor,row) to the two m=1 MMVQ dispatches below — decode-exact contract holds.
7521 // VERIFY-TIER TRUNK FUSION (MEMRA_SPEC_FUSED_T, t=2-4): quantize h ONCE for every
7522 // fused-eligible same-input Q8_0 pair of this layer (35B wqkv+wqkv_gate; 9B
7523 // ssm_beta+ssm_alpha) — each fused2 batched launch then replaces two decode-exact
7524 // calls (each of which re-quantizes the same h + runs its own _b2/_b4 launch).
7525 // Bit-identical per (tensor,token,row) — see spec_fused_t().
7526 // BATCHED EPILOGUE RE-FUSE (lane/vt-fixes fix 2): `h_q8` = the attn-input norm emitted
7527 // directly as q8_1 by the caller's fused rms_norm_q8_1 (bit-identical to the unfused
7528 // chain, kernel-check-pinned). When present it REPLACES the standalone quantize below
7529 // and feeds every projection; the caller guaranteed all four input projections are
7530 // q8_1-fast. When absent, the old shared-quantize (fused-t window) stands.
7531 let h_q8_t = if h_q8.is_none()
7532 && spec_fused_t()
7533 && (2..=4).contains(&t)
7534 && ((e.uses_q8_1_fast(&la.wqkv) && e.uses_q8_1_fast(&la.wqkv_gate))
7535 || (e.uses_q8_1_fast(&la.ssm_beta) && e.uses_q8_1_fast(&la.ssm_alpha)))
7536 {
7537 Some(e.quantize_q8_1(h, t, cfg.n_embd as usize)?)
7538 } else {
7539 None
7540 };
7541 // one view: the caller's fused-norm q8 or this fn's own shared quantize.
7542 let hq8_any: Option<(&CudaSlice<i8>, &CudaSlice<f32>)> =
7543 h_q8.or(h_q8_t.as_ref().map(|(q, d)| (q, d)));
7544 let (qkv_mixed, z) = {
7545 let mut fused = None;
7546 if t == 1 && e.uses_q8_1_fast(&la.wqkv) && e.uses_q8_1_fast(&la.wqkv_gate) {
7547 let (hq, hd) = e.quantize_q8_1(h, 1, cfg.n_embd as usize)?;
7548 fused = e.matmul_q8_fused2(&la.wqkv, &la.wqkv_gate, &hq, &hd)?;
7549 } else if let Some((hq, hd)) = hq8_any {
7550 if spec_fused_t() && (2..=4).contains(&t) {
7551 fused = e.matmul_q8_fused2_t(&la.wqkv, &la.wqkv_gate, hq, hd, t)?;
7552 }
7553 }
7554 match (fused, hq8_any) {
7555 (Some(pair), _) => pair,
7556 (None, Some((hq, hd))) if h_q8.is_some() => (
7557 e.matmul_decode_exact_pre(&la.wqkv, hq, hd, t)?,
7558 e.matmul_decode_exact_pre(&la.wqkv_gate, hq, hd, t)?,
7559 ),
7560 (None, _) => (
7561 e.matmul_decode_exact(&la.wqkv, h, t)?,
7562 e.matmul_decode_exact(&la.wqkv_gate, h, t)?,
7563 ),
7564 }
7565 };
7566 // beta+alpha DUAL at T=1 (75% of p3 rounds run T=1 verify — p-min chain cuts): the dual
7567 // mr2 kernel is bit-identical per element to the m=1 MMVQ matmul_decode_exact dispatches
7568 // (same warp-per-row body, blockIdx.y picks the weight), so the decode-exact contract
7569 // holds; the run-spec battery is the arbiter. T>1 keeps the per-tensor decode-exact path.
7570 let (beta_raw, alpha) = if t == 1 {
7571 let (hq, hd) = e.quantize_q8_1(h, 1, cfg.n_embd as usize)?;
7572 match e.matmul_pre_dual_noscale(&la.ssm_beta, &la.ssm_alpha, &hq, &hd, 1)? {
7573 Some(((mut b, bs), (mut a, as_))) => {
7574 if bs != 1.0 {
7575 e.scale_inplace(&mut b, bs, la.ssm_beta.out_features())?;
7576 }
7577 if as_ != 1.0 {
7578 e.scale_inplace(&mut a, as_, la.ssm_alpha.out_features())?;
7579 }
7580 (b, a)
7581 }
7582 // Q8_0 fused2 twin (9B stores beta/alpha as Q8_0): DISPATCH-MIRRORS the eager
7583 // decode's beta_alpha closure — the fused body is qmatvec_q8_0_mmvq verbatim,
7584 // bit-identical per row (kernel-check rel=0.00e0 gate), so decode==verify holds.
7585 None => match e.matmul_q8_fused2(&la.ssm_beta, &la.ssm_alpha, &hq, &hd)? {
7586 Some((b, a)) => (b, a),
7587 None => (
7588 e.matmul_decode_exact(&la.ssm_beta, h, 1)?,
7589 e.matmul_decode_exact(&la.ssm_alpha, h, 1)?,
7590 ),
7591 },
7592 }
7593 } else {
7594 // fused-t twin (9B stores beta/alpha as Q8_0): same shared-quantize + one launch
7595 // contract as the wqkv pair above; 35B beta/alpha are Float -> None -> fallback.
7596 let mut nvfp4_fused = None;
7597 let mut q8_fused = None;
7598 if let Some((hq, hd)) = hq8_any {
7599 if t == 3 && std::env::var("MEMRA_NVFP4_AUX_DUAL").as_deref() != Ok("0") {
7600 nvfp4_fused =
7601 e.matmul_decode_exact_dual_pre(&la.ssm_beta, &la.ssm_alpha, hq, hd, t)?;
7602 if nvfp4_fused.is_some() && std::env::var("MEMRA_DEBUG").is_ok() {
7603 static ONCE: std::sync::Once = std::sync::Once::new();
7604 ONCE.call_once(|| {
7605 eprintln!("[memra] NVFP4 beta+alpha batched aux dual ENGAGED (t={t})")
7606 });
7607 }
7608 }
7609 if nvfp4_fused.is_none() && spec_fused_t() && (2..=4).contains(&t) {
7610 q8_fused = e.matmul_q8_fused2_t(&la.ssm_beta, &la.ssm_alpha, hq, hd, t)?;
7611 }
7612 }
7613 if let Some(((mut b, bs), (mut a, as_))) = nvfp4_fused {
7614 if bs != 1.0 {
7615 e.scale_inplace(&mut b, bs, t * la.ssm_beta.out_features())?;
7616 }
7617 if as_ != 1.0 {
7618 e.scale_inplace(&mut a, as_, t * la.ssm_alpha.out_features())?;
7619 }
7620 (b, a)
7621 } else if let Some(pair) = q8_fused {
7622 pair
7623 } else {
7624 match hq8_any {
7625 Some((hq, hd)) if h_q8.is_some() => (
7626 e.matmul_decode_exact_pre(&la.ssm_beta, hq, hd, t)?,
7627 e.matmul_decode_exact_pre(&la.ssm_alpha, hq, hd, t)?,
7628 ),
7629 _ => (
7630 e.matmul_decode_exact(&la.ssm_beta, h, t)?,
7631 e.matmul_decode_exact(&la.ssm_alpha, h, t)?,
7632 ),
7633 }
7634 }
7635 };
7636
7637 // conv with CARRIED state + ring roll (T >= pad rides the input-column update kernel;
7638 // T < pad — the MEMRA_SPEC_M2 t=2 arm — rolls via the pure-copy ring rebuild).
7639 let rl = cache.recur[il].as_mut().unwrap();
7640 let mut conv_out = e.uninit(conv_dim * t)?;
7641 e.ssm_conv1d_tm_state(
7642 &qkv_mixed,
7643 &mut rl.conv_state,
7644 la.ssm_conv1d.float_data(),
7645 &mut conv_out,
7646 conv_dim,
7647 t,
7648 d_conv,
7649 )?;
7650
7651 // GDN prep via the prefill kernels (repack + L2 + sigmoid + glog), T-wide.
7652 let mut q_g = e.uninit(d_state * num_v * t)?;
7653 let mut k_g = e.uninit(d_state * num_v * t)?;
7654 let mut v_g = e.uninit(d_state * num_v * t)?;
7655 e.qkv_to_gdn_repack(
7656 &conv_out, &mut q_g, &mut k_g, &mut v_g, d_state, num_v, num_k, key_dim, t,
7657 )?;
7658 let mut q_l2 = e.uninit(d_state * num_v * t)?;
7659 e.l2_norm_decode(&q_g, &mut q_l2, d_state, num_v * t, eps)?;
7660 let mut k_l2 = e.uninit(d_state * num_v * t)?;
7661 e.l2_norm_decode(&k_g, &mut k_l2, d_state, num_v * t, eps)?;
7662 let mut beta = e.uninit(t * num_v)?;
7663 e.sigmoid(&beta_raw, &mut beta, t * num_v)?;
7664 let mut g_log = e.uninit(t * num_v)?;
7665 e.gdn_glog(
7666 &alpha,
7667 la.ssm_dt.float_data(),
7668 la.ssm_a.float_data(),
7669 &mut g_log,
7670 num_v,
7671 t,
7672 )?;
7673
7674 // ONE gdn_scan over T tokens from the carried state (internal sequential loop ==
7675 // T chained T=1 steps). Ping-pong the resident buffers like eager decode.
7676 let mut o = e.uninit(d_state * num_v * t)?;
7677 {
7678 let crate::cache::RecurLayer {
7679 ssm_state,
7680 ssm_state_alt,
7681 ..
7682 } = rl;
7683 e.gdn_scan_s128(
7684 &q_l2,
7685 &k_l2,
7686 &v_g,
7687 &g_log,
7688 &beta,
7689 ssm_state,
7690 ssm_state_alt,
7691 &mut o,
7692 num_v,
7693 t,
7694 scale,
7695 )?;
7696 }
7697 std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
7698
7699 // gated RMSNorm + out projection, T-wide. FUSED-QUANTIZE ARM (lane/vt-fixes fix 2,
7700 // mirroring the T=1 decode's launch-arc form): when ssm_out rides the q8_1 fast path,
7701 // emit q8_1 straight from the gated norm at nrows=num_v*t (row-indexed kernel, the
7702 // T-wide launch is the per-row program; kernel-check pins bit-identity vs
7703 // gated_rmsnorm -> quantize_q8_1 at T=1 and T=5) and feed the decode-exact dispatch
7704 // pre-quantized — one launch replaces norm + quantize. Fallback = the f32 chain.
7705 let out = if e.uses_q8_1_fast(&la.ssm_out) {
7706 let (gq, gd) =
7707 e.gated_rmsnorm_q8_1(&o, la.ssm_norm.float_data(), &z, d_state, num_v * t, eps)?;
7708 e.matmul_decode_exact_pre(&la.ssm_out, &gq, &gd, t)?
7709 } else {
7710 let mut gn = e.uninit(d_state * num_v * t)?;
7711 e.gated_rmsnorm(
7712 &o,
7713 la.ssm_norm.float_data(),
7714 &z,
7715 &mut gn,
7716 d_state,
7717 num_v * t,
7718 eps,
7719 )?;
7720 // DECODE-EXACT out-projection: same MMVQ path as the T=1 decode (ssm_out at m>=5
7721 // would fall to dp4a with a different FP reduction order — same class of bug as
7722 // the input projs).
7723 e.matmul_decode_exact(&la.ssm_out, &gn, t)?
7724 };
7725 let stash = if want_stash {
7726 Some(GdnStash {
7727 qkv_mixed,
7728 q_l2,
7729 k_l2,
7730 v_g,
7731 g_log,
7732 beta,
7733 })
7734 } else {
7735 None
7736 };
7737 Ok((out, stash))
7738 }
7739
7740 /// REPLAY-FREE partial-accept commit (2026-07-03): make the cache state == "committed through
7741 /// the first `j` verify columns" WITHOUT the legacy rollback + duplicate trunk replay.
7742 /// - Full-attn KV: truncate both the owning-stage shadow and every TP rank to snapshot + j.
7743 /// The verify's appended rows for those columns are bit-identical to what an eager T=1
7744 /// chain writes (the decode-exact contract the verify-probe gates), so keeping them ==
7745 /// replaying them.
7746 /// - Linear layers, batched path: rebuild the conv ring by PURE COPIES (ring holds raw input
7747 /// columns) and the ssm state by a prefix re-run of the SAME gdn_scan kernel (t=j) from the
7748 /// snapshot state over the stash's identical inputs — the kernel's t-loop carries state in
7749 /// registers and writes it once at the end, so iterations 0..j-1 are independent of T:
7750 /// bit-identical to the verify's own state after j tokens == the eager chain state.
7751 /// - Linear layers, per-column path: restore the cloned actual state after column j-1.
7752 /// Caller guarantees 1 <= j <= t-1 (j==0 rounds take the legacy rollback; j==t is full accept).
7753 fn commit_verified_prefix(
7754 &self,
7755 e: &Engine,
7756 cache: &mut Cache,
7757 snap: &crate::cache::CacheSnapshot,
7758 ckpt: &VerifyCkpt,
7759 j: usize,
7760 kv_lens_done: bool,
7761 dev_j: Option<(&CudaSlice<u32>, usize, usize)>,
7762 ) -> Result<(), Box<dyn std::error::Error>> {
7763 // GDN geometry derives lazily inside recurrent-layer arms. Full-attention plans carry no
7764 // recurrent state and must never be forced through a synthetic SSM geometry.
7765 // Engine-bundle slice 1 (DSF-ROUNDCOST-20260820 §1.1): the per-column-arm restores
7766 // are 2 tiny D2D copies per linear layer (~96 dispatches/partial round on the q38
7767 // route). When every cols-arm layer shares uniform state sizes (single ssm cfg —
7768 // always true today), batch them into two `copy_batch_uniform_f32` launches. Bytes,
7769 // buffers and stream order are identical to the per-layer memcpy sequence; the
7770 // kernel-rebuild (gdn-stash) arm below is untouched. MEMRA_STATE_COPY_BATCH=0 reverts.
7771 let mut batched_cols = false;
7772 if state_copy_batch_on() && dev_j.is_none() {
7773 use cudarc::driver::DevicePtr;
7774 let s = &e.gpu.stream();
7775 let mut conv_pairs: Vec<(u64, u64)> = Vec::new();
7776 let mut ssm_pairs: Vec<(u64, u64)> = Vec::new();
7777 let (mut conv_words, mut ssm_words) = (0usize, 0usize);
7778 let mut uniform = true;
7779 for il in 0..self.layers.len() {
7780 let Some(rl) = cache.recur[il].as_ref() else {
7781 continue;
7782 };
7783 if ckpt.gdn[il].is_some() {
7784 continue; // kernel-rebuild arm restores below, per layer
7785 }
7786 let Some(cols) = &ckpt.cols[il] else {
7787 continue; // missing-ckpt error surfaces in the main loop
7788 };
7789 let (c, st) = &cols[j - 1];
7790 if conv_pairs.is_empty() {
7791 conv_words = c.len();
7792 ssm_words = st.len();
7793 } else if c.len() != conv_words || st.len() != ssm_words {
7794 uniform = false;
7795 break;
7796 }
7797 let (pc, _g0) = c.device_ptr(s);
7798 let (dc, _g1) = rl.conv_state.device_ptr(s);
7799 let (ps, _g2) = st.device_ptr(s);
7800 let (ds, _g3) = rl.ssm_state.device_ptr(s);
7801 conv_pairs.push((pc as u64, dc as u64));
7802 ssm_pairs.push((ps as u64, ds as u64));
7803 }
7804 if uniform && !conv_pairs.is_empty() {
7805 let n = conv_pairs.len();
7806 let mut t = vec![0u64; 2 * n];
7807 for (k, &(src, dst)) in conv_pairs.iter().enumerate() {
7808 t[k] = src;
7809 t[n + k] = dst;
7810 }
7811 let conv_t = e.htod_u64(&t)?;
7812 for (k, &(src, dst)) in ssm_pairs.iter().enumerate() {
7813 t[k] = src;
7814 t[n + k] = dst;
7815 }
7816 let ssm_t = e.htod_u64(&t)?;
7817 e.copy_batch_uniform_f32(&conv_t, n, conv_words)?;
7818 e.copy_batch_uniform_f32(&ssm_t, n, ssm_words)?;
7819 batched_cols = true;
7820 }
7821 }
7822 rewind_tp_kv_verified_prefix(&mut cache.tp_kv, &snap.tp_kv_len, j)?;
7823 for il in 0..self.layers.len() {
7824 if let (Some(kvl), Some(saved)) = (cache.kv[il].as_mut(), snap.kv_len[il]) {
7825 kvl.len = saved + j;
7826 // devacc 3a: spec_rollback_kv already wrote len_d on-device (same value).
7827 if !kv_lens_done {
7828 e.set_i32_one(&mut kvl.len_d, kvl.len as i32)?;
7829 }
7830 }
7831 if let Some(rl) = cache.recur[il].as_mut() {
7832 let Mixer::Linear(linear) = &self.layers[il].mixer else {
7833 return Err(format!("recurrent cache layer {il} has no GDN plan").into());
7834 };
7835 let geometry = linear.geometry;
7836 let d_state = geometry.key_head_dim as usize;
7837 let num_k = geometry.key_heads as usize;
7838 let num_v = geometry.value_heads as usize;
7839 let d_conv = geometry.conv_kernel as usize;
7840 let conv_dim = d_state * num_k * 2 + geometry.value_head_dim as usize * num_v;
7841 let scale = 1.0 / (d_state as f32).sqrt();
7842 if let Some(st) = &ckpt.gdn[il] {
7843 let ring_old = snap.conv[il].as_ref().expect("snapshot missing conv");
7844 let state_in = snap.ssm[il].as_ref().expect("snapshot missing ssm");
7845 if let Some((acc, base, t_v)) = dev_j {
7846 // 3b: j read on-device (_dc twins, same bodies; full accept early-exits).
7847 e.ssm_conv_ring_rebuild_dc(
7848 &st.qkv_mixed,
7849 ring_old,
7850 &mut rl.conv_state,
7851 conv_dim,
7852 acc,
7853 base,
7854 t_v,
7855 d_conv,
7856 )?;
7857 let mut o = e.uninit(d_state * num_v * j.max(1))?;
7858 e.gdn_scan_s128_dc(
7859 &st.q_l2,
7860 &st.k_l2,
7861 &st.v_g,
7862 &st.g_log,
7863 &st.beta,
7864 state_in,
7865 &mut rl.ssm_state,
7866 &mut o,
7867 num_v,
7868 acc,
7869 base,
7870 t_v,
7871 scale,
7872 )?;
7873 } else {
7874 e.ssm_conv_ring_rebuild(
7875 &st.qkv_mixed,
7876 ring_old,
7877 &mut rl.conv_state,
7878 conv_dim,
7879 j,
7880 d_conv,
7881 )?;
7882 let mut o = e.uninit(d_state * num_v * j)?; // scan output, discarded
7883 e.gdn_scan_s128(
7884 &st.q_l2,
7885 &st.k_l2,
7886 &st.v_g,
7887 &st.g_log,
7888 &st.beta,
7889 state_in,
7890 &mut rl.ssm_state,
7891 &mut o,
7892 num_v,
7893 j,
7894 scale,
7895 )?;
7896 }
7897 } else if let Some(cols) = &ckpt.cols[il] {
7898 if !batched_cols {
7899 let (c, s) = &cols[j - 1];
7900 e.copy_into(&mut rl.conv_state, 0, c, c.len())?;
7901 e.copy_into(&mut rl.ssm_state, 0, s, s.len())?;
7902 }
7903 } else {
7904 return Err(
7905 "commit_verified_prefix: verify ckpt missing for linear layer".into(),
7906 );
7907 }
7908 }
7909 }
7910 cache.pos = snap.pos + j;
7911 Ok(())
7912 }
7913
7914 /// ROUND-STREAM: recur restore with device-j (the _dc twins; full accept early-exits
7915 /// in-kernel). Requires the batched-linear stash on every linear layer (stream gate).
7916 fn commit_verified_prefix_stream(
7917 &self,
7918 e: &Engine,
7919 cache: &mut Cache,
7920 snap: &crate::cache::CacheSnapshot,
7921 ckpt: &VerifyCkpt,
7922 acc: &CudaSlice<u32>,
7923 base: usize,
7924 t_v: usize,
7925 ) -> Result<(), Box<dyn std::error::Error>> {
7926 for il in 0..self.layers.len() {
7927 if let Some(rl) = cache.recur[il].as_mut() {
7928 let Mixer::Linear(linear) = &self.layers[il].mixer else {
7929 return Err(format!("recurrent cache layer {il} has no GDN plan").into());
7930 };
7931 let geometry = linear.geometry;
7932 let d_state = geometry.key_head_dim as usize;
7933 let num_k = geometry.key_heads as usize;
7934 let num_v = geometry.value_heads as usize;
7935 let d_conv = geometry.conv_kernel as usize;
7936 let conv_dim = d_state * num_k * 2 + geometry.value_head_dim as usize * num_v;
7937 let scale = 1.0 / (d_state as f32).sqrt();
7938 let st = ckpt.gdn[il]
7939 .as_ref()
7940 .ok_or("stream restore: batched-linear stash missing")?;
7941 let ring_old = snap.conv[il].as_ref().expect("snapshot missing conv");
7942 let state_in = snap.ssm[il].as_ref().expect("snapshot missing ssm");
7943 e.ssm_conv_ring_rebuild_dc(
7944 &st.qkv_mixed,
7945 ring_old,
7946 &mut rl.conv_state,
7947 conv_dim,
7948 acc,
7949 base,
7950 t_v,
7951 d_conv,
7952 )?;
7953 let mut o = e.uninit(d_state * num_v * t_v)?;
7954 e.gdn_scan_s128_dc(
7955 &st.q_l2,
7956 &st.k_l2,
7957 &st.v_g,
7958 &st.g_log,
7959 &st.beta,
7960 state_in,
7961 &mut rl.ssm_state,
7962 &mut o,
7963 num_v,
7964 acc,
7965 base,
7966 t_v,
7967 scale,
7968 )?;
7969 }
7970 }
7971 Ok(())
7972 }
7973
7974 /// EAGLE3 aux-capturing verify forward over `tokens` (T) — mirrors `decode_step_t_h` exactly
7975 /// (same KV append, same causal verify, same recur advance) but ALSO clones the aux residual-
7976 /// stream hiddens (blocks in `aux_layers`) for TWO columns: the LAST column (always) and the
7977 /// optional `pred_col` (the EAGLE seed = bonus's predecessor). Returns
7978 /// (all_T_logits host, last_col_aux, pred_col_aux?). Used by the EAGLE3 orchestrator's commit.
7979 pub fn decode_step_t_aux2(
7980 &self,
7981 e: &Engine,
7982 tokens: &[u32],
7983 pos0: usize,
7984 cache: &mut Cache,
7985 aux_layers: &[usize],
7986 pred_col: Option<usize>,
7987 ) -> Result<
7988 (Vec<f32>, Vec<CudaSlice<f32>>, Option<Vec<CudaSlice<f32>>>),
7989 Box<dyn std::error::Error>,
7990 > {
7991 let cfg = &self.cfg;
7992 let n_embd = cfg.n_embd as usize;
7993 let eps = cfg.rms_eps;
7994 let t = tokens.len();
7995 let pos_vec: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
7996 let pos_d = e.htod_i32(&pos_vec)?;
7997 let mut x = e.htod(&self.embd.gather(n_embd, tokens))?;
7998 let mut aux_last: Vec<CudaSlice<f32>> = Vec::with_capacity(aux_layers.len());
7999 let mut aux_pred: Vec<CudaSlice<f32>> = Vec::new();
8000 let want_pred = pred_col.is_some();
8001
8002 for (il, layer) in self.layers.iter().enumerate() {
8003 // DISPATCH-MIRRORED norms (FP-order lesson #8) — see decode_step_t_h_emb.
8004 let mixer_fast = self.mixer_in_q8_1_fast(e, &layer.mixer);
8005 let norm_fused = std::env::var("MEMRA_NO_FUSE_NORMQ").is_err() && mixer_fast;
8006 let mut h = vbuf(e, t * n_embd)?; // fully written by either rms_norm arm
8007 if norm_fused {
8008 e.rms_norm_decode(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
8009 } else {
8010 e.rms_norm(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
8011 }
8012 let mixed = match &layer.mixer {
8013 Mixer::Full(fa) => {
8014 self.full_attn_verify(e, fa, &h, None, &pos_d, t, cache, il, None)?
8015 }
8016 Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
8017 Mixer::Linear(la) => {
8018 let mut out = e.zeros(t * n_embd)?;
8019 for col in 0..t {
8020 let mut h_col = e.zeros(n_embd)?;
8021 let src = h.slice(col * n_embd..(col + 1) * n_embd);
8022 e.copy_view_into(&mut h_col, 0, &src, n_embd)?;
8023 let m_col = self.linear_attn_decode(e, la, &h_col, cache, il)?;
8024 e.copy_into(&mut out, col * n_embd, &m_col, n_embd)?;
8025 }
8026 out
8027 }
8028 };
8029 let ffn_fuse = match &layer.ffn {
8030 crate::hybrid::Ffn::Dense {
8031 ffn_gate, ffn_up, ..
8032 } => {
8033 std::env::var("MEMRA_NO_FUSE_NORMQ").is_err()
8034 && e.uses_q8_1_fast(ffn_gate)
8035 && e.uses_q8_1_fast(ffn_up)
8036 }
8037 crate::hybrid::Ffn::Moe(_) => false,
8038 };
8039 let mut x1 = vbuf(e, t * n_embd)?; // fully written by add / add_rms_norm
8040 let mut z = vbuf(e, t * n_embd)?; // fully written by rms_norm_decode / add_rms_norm
8041 if ffn_fuse {
8042 e.add(&x, &mixed, &mut x1, t * n_embd)?;
8043 e.rms_norm_decode(
8044 &x1,
8045 layer.post_attn_norm.float_data(),
8046 &mut z,
8047 n_embd,
8048 t,
8049 eps,
8050 )?;
8051 } else {
8052 e.add_rms_norm(
8053 &x,
8054 &mixed,
8055 layer.post_attn_norm.float_data(),
8056 &mut x1,
8057 &mut z,
8058 n_embd,
8059 t,
8060 eps,
8061 )?;
8062 }
8063 let ffn_out = match &layer.ffn {
8064 crate::hybrid::Ffn::Dense {
8065 ffn_gate,
8066 ffn_up,
8067 ffn_down,
8068 } => {
8069 let n_ff = ffn_gate.out_features();
8070 let gate = e.matmul_decode_exact(ffn_gate, &z, t)?;
8071 let up = e.matmul_decode_exact(ffn_up, &z, t)?;
8072 let mut act = vbuf(e, t * n_ff)?; // fully written by ffn_act_lim
8073 // dense FFN clamp = the SHEXP array (upstream build_ffn serves both).
8074 Self::ffn_act_lim(
8075 e,
8076 &self.cfg,
8077 &gate,
8078 &up,
8079 1.0,
8080 1.0,
8081 self.cfg.clamp_shexp_at(il as u32),
8082 &mut act,
8083 t * n_ff,
8084 )?;
8085 e.matmul_decode_exact(ffn_down, &act, t)?
8086 }
8087 crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il(e, m, &z, t, il as u16)?,
8088 };
8089 let mut x2 = vbuf(e, t * n_embd)?; // fully written by add
8090 e.add(&x1, &ffn_out, &mut x2, t * n_embd)?;
8091 if aux_layers.contains(&il) {
8092 let mut a = e.zeros(n_embd)?;
8093 e.copy_view_into(&mut a, 0, &x2.slice((t - 1) * n_embd..t * n_embd), n_embd)?;
8094 aux_last.push(a);
8095 if let Some(pc) = pred_col {
8096 let mut ap = e.zeros(n_embd)?;
8097 e.copy_view_into(
8098 &mut ap,
8099 0,
8100 &x2.slice(pc * n_embd..(pc + 1) * n_embd),
8101 n_embd,
8102 )?;
8103 aux_pred.push(ap);
8104 }
8105 }
8106 x = x2;
8107 }
8108 let mut hn = vbuf(e, t * n_embd)?; // fully written by rms_norm_decode
8109 e.rms_norm_decode(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
8110 let logits = e.matmul_decode_exact(&self.output, &hn, t)?;
8111 let host = e.dtoh(&logits)?;
8112 cache.pos += t;
8113 Ok((
8114 host,
8115 aux_last,
8116 if want_pred { Some(aux_pred) } else { None },
8117 ))
8118 }
8119
8120 /// step35 SPEC-VERIFY attention over T query tokens — a per-row REPLAY of the eager
8121 /// `step35_decode_attn`.
8122 ///
8123 /// WHY A REPLAY AND NOT A BATCHED TWIN. The verify's whole job is to be bit-identical to what
8124 /// the eager decode would have computed for the same tokens; that is what makes greedy spec
8125 /// decode exact (run-spec asserts token identity for K=1..8). Every other verify arm in this
8126 /// file earns that identity by carefully mirroring dispatch (`matmul_decode_exact` to force
8127 /// MMVQ at any m, per-layer `ffn_fuse` mirroring, per-row `fa_decode` key bounds). step35
8128 /// stacks FOUR more per-layer degrees of freedom on top of that — per-layer `n_head`
8129 /// (64 full / 96 SWA), per-layer rotary width (64 full / 128 SWA), per-layer rope base, and a
8130 /// SEPARATE `attn_gate` tensor whose projection shares the attn-normed input — and its SWA
8131 /// layers attend through a token-OFFSET view whose offset is a function of the ABSOLUTE
8132 /// position of each query row. A batched twin would have to reproduce all of that AND the
8133 /// per-row offset in one launch; the offset alone rules out the existing rows kernels (they
8134 /// take one `base_len`, not a per-row offset).
8135 ///
8136 /// So this arm calls the eager path itself, once per row, on the same cache. Identity is then
8137 /// true BY CONSTRUCTION rather than by mirroring: row r runs exactly the kernel sequence that
8138 /// eager decode step r runs (same projections, same q8_1 fusion decision, same append, same
8139 /// view arithmetic, same `fa_decode_kvmod`, same gate), because it IS that code. Cost: T x the
8140 /// eager decode mixer instead of one batched pass — the same trade the generic arm's `else`
8141 /// per-row loop already accepts when `fa_rows_eligible` says no. Correctness first; a batched
8142 /// step35 twin is a perf lane's job and must be gated against this arm.
8143 ///
8144 /// The `h_q8` pre-quantized pair from the caller's fused norm is NOT forwarded: it is a
8145 /// T-row buffer and `step35_decode_attn`'s `pre_q` contract is one row. Instead each row's
8146 /// f32 `h` slice is handed over and the callee re-derives its own q8_1 exactly as eager decode
8147 /// does (`quantize_q8_1(h, 1, n_embd)`) — which is the dispatch being mirrored. Callers that
8148 /// took the fused arm therefore MUST still pass a live `h`; `step35_verify` asserts that.
8149 #[allow(clippy::too_many_arguments)]
8150 fn step35_verify(
8151 &self,
8152 e: &Engine,
8153 fa: &FullAttnLayer,
8154 h: &CudaSlice<f32>,
8155 h_q8: Option<(&CudaSlice<i8>, &CudaSlice<f32>)>,
8156 t: usize,
8157 cache: &mut Cache,
8158 il: usize,
8159 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8160 let n_embd = self.cfg.n_embd as usize;
8161 // The fused (h-less) attn-norm arm hands `h` as a zero-length placeholder. This arm needs
8162 // the f32 rows, so the caller must not take that lever for step35 — enforced at the call
8163 // site by the sliding-gated-MoE `Mixer::Full(_) => false` arm of
8164 // `lin_q8_only` in `decode_step_t_core_stream`, and asserted here so a future caller
8165 // cannot regress it into silently reading an empty buffer.
8166 assert_eq!(
8167 h.len(),
8168 t * n_embd,
8169 "step35_verify needs the f32 attn-normed rows ([t*n_embd]); the caller took the \
8170 fused q8-only norm arm (h_q8={}) — step35 must stay on the unfused arm",
8171 h_q8.is_some()
8172 );
8173 // ROW WIDTH IS n_embd, NOT n_head*head_dim: `step35_decode_attn` returns the mixer output
8174 // AFTER `wo`, so a row is [n_embd] — the same contract the generic arm's
8175 // `matmul_decode_exact(&fa.wo, &attn_g, t)` return has. Sizing this buffer from the
8176 // per-layer head geometry (8192 on full-attn, 12288 on SWA) instead overran the row on the
8177 // FIRST copy and panicked inside `copy_into`'s `CudaView::slice` unwrap
8178 // (raw/mtp-bt-20260806T212127Z.log frames 12-13).
8179 let mut out = vbuf(e, t * n_embd)?; // each row fully written by the copy below
8180 for r in 0..t {
8181 // Absolute position of this query row. `cache.pos` is the committed length at round
8182 // start and every row before r has already been appended by this loop, so the r-th
8183 // verify token sits at cache.pos + r — the same position eager decode would give it.
8184 let pos_d = e.htod_i32(&[(cache.pos + r) as i32])?;
8185 let mut h_row = vbuf(e, n_embd)?; // fully written by copy_view_into
8186 e.copy_view_into(
8187 &mut h_row,
8188 0,
8189 &h.slice(r * n_embd..(r + 1) * n_embd),
8190 n_embd,
8191 )?;
8192 // THE eager decode mixer: appends this row's K/V at kvl.len, advances it, then
8193 // attends over the (SWA-offset) view. Post-`wo`, same contract as this fn returns.
8194 let o = self.step35_decode_attn(e, fa, il, &h_row, None, &pos_d, cache)?;
8195 debug_assert_eq!(
8196 o.len(),
8197 n_embd,
8198 "step35_decode_attn returns post-wo [n_embd]"
8199 );
8200 e.copy_into(&mut out, r * n_embd, &o, n_embd)?;
8201 }
8202 Ok(out)
8203 }
8204
8205 /// Full-attention mixer over T query tokens with a GROWING resident KV (verify path, §D.3).
8206 /// Appends the T new K/V columns to cache.kv[il] then attends causally over [0..len) via
8207 /// fa_prefill. Token-major [T, kv_dim] projection layout == cache row layout (single copy).
8208 #[allow(clippy::too_many_arguments)]
8209 fn full_attn_verify(
8210 &self,
8211 e: &Engine,
8212 fa: &FullAttnLayer,
8213 h: &CudaSlice<f32>,
8214 h_q8: Option<(&CudaSlice<i8>, &CudaSlice<f32>)>,
8215 pos_d: &CudaSlice<i32>,
8216 t: usize,
8217 cache: &mut Cache,
8218 il: usize,
8219 stream_ctr: Option<&CudaSlice<i32>>,
8220 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8221 // step35: the generic geometry below is wrong for this arch (per-layer n_head, partial
8222 // per-layer rope, the SWA offset view, and a SEPARATE head-wise gate tensor), so it takes
8223 // its own arm. A verify that silently computes different attention than decode defeats the
8224 // whole self-consistency gate, so the arm is a per-row REPLAY of `step35_decode_attn`
8225 // rather than a batched twin — see `step35_verify` for why that is the exactness-correct
8226 // shape and not laziness.
8227 if self.sliding_gated_moe_batch_program() {
8228 if stream_ctr.is_some() {
8229 return Err(
8230 "step35 has no ROUND-STREAM verify arm (the device-counter _dc twins \
8231 cannot express the SWA offset KV view; same root cause as the dc \
8232 decode refusal) — run spec without the stream arm"
8233 .into(),
8234 );
8235 }
8236 return self.step35_verify(e, fa, h, h_q8, t, cache, il);
8237 }
8238 let cfg = &self.cfg;
8239 let geometry = cfg.full_attention_geometry_at(il as u32);
8240 let n_head = geometry.n_head as usize;
8241 let n_head_kv = geometry.n_head_kv as usize;
8242 let head_dim = geometry.head_dim_k as usize;
8243 let eps = cfg.rms_eps;
8244 let scale = geometry.attention_scale();
8245 let n_embd = cfg.n_embd as usize;
8246
8247 // DECODE-EXACT Q/K/V projections: matmul_decode_exact forces the MMVQ (warp-per-row) path
8248 // for every m, matching the T=1 decode's FP accumulation order. matmul_pre at m>=5 would
8249 // fall to dp4a (128-thread, two-level reduce) with a different FP sum order.
8250 // Q8 TRUNK-FUSION at T=1: DISPATCH-MIRRORS the eager decode's fused3 (bit-identical body).
8251 // BATCHED EPILOGUE RE-FUSE (lane/vt-fixes fix 2): `h_q8` = the attn-input norm's q8_1
8252 // form emitted by the fused rms_norm_q8_1 (bit-identical to rms_norm_decode ->
8253 // quantize_q8_1, kernel-check-pinned). When present (caller checked mixer q8_1-fast),
8254 // every projection consumes it — `h` may be a zero-len placeholder and must not be read.
8255 let (qf, mut k, v) = {
8256 let mut fused = None;
8257 let qkv_fast =
8258 e.uses_q8_1_fast(&fa.wq) && e.uses_q8_1_fast(&fa.wk) && e.uses_q8_1_fast(&fa.wv);
8259 if t == 1 && qkv_fast {
8260 let (hq_o, hd_o);
8261 let (hq, hd): (&CudaSlice<i8>, &CudaSlice<f32>) = match h_q8 {
8262 Some(p) => p,
8263 None => {
8264 (hq_o, hd_o) = e.quantize_q8_1(h, 1, n_embd)?;
8265 (&hq_o, &hd_o)
8266 }
8267 };
8268 fused = e.matmul_q8_fused3(&fa.wq, &fa.wk, &fa.wv, hq, hd)?;
8269 } else if spec_fused_t() && (2..=4).contains(&t) && qkv_fast {
8270 // VERIFY-TIER TRUNK FUSION (MEMRA_SPEC_FUSED_T): one shared quantize + one
8271 // fused3 batched launch replaces three decode-exact calls (3 re-quantizes of
8272 // the same h + 3 _b2/_b4 launches). Bit-identical per (tensor,token,row).
8273 let (hq_o, hd_o);
8274 let (hq, hd): (&CudaSlice<i8>, &CudaSlice<f32>) = match h_q8 {
8275 Some(p) => p,
8276 None => {
8277 (hq_o, hd_o) = e.quantize_q8_1(h, t, n_embd)?;
8278 (&hq_o, &hd_o)
8279 }
8280 };
8281 fused = e.matmul_q8_fused3_t(&fa.wq, &fa.wk, &fa.wv, hq, hd, t)?;
8282 }
8283 match (fused, h_q8) {
8284 (Some(triple), _) => triple,
8285 // shared pre-quantized activation (q8_1-fast guaranteed by the caller): the
8286 // decode-exact dispatch consumes (hq, hd) instead of re-quantizing 3x.
8287 (None, Some((hq, hd))) if qkv_fast => (
8288 e.matmul_decode_exact_pre(&fa.wq, hq, hd, t)?,
8289 e.matmul_decode_exact_pre(&fa.wk, hq, hd, t)?,
8290 e.matmul_decode_exact_pre(&fa.wv, hq, hd, t)?,
8291 ),
8292 (None, _) => (
8293 e.matmul_decode_exact(&fa.wq, h, t)?,
8294 e.matmul_decode_exact(&fa.wk, h, t)?,
8295 e.matmul_decode_exact(&fa.wv, h, t)?,
8296 ),
8297 }
8298 };
8299 // M3/Hy3 have no attention output gate — wq out is exactly q; skip the split.
8300 let gated = geometry.attention_gate == memra_gguf::config::AttentionGateKind::FusedQ;
8301 let (mut q, gate) = if gated {
8302 let mut q = vbuf(e, t * n_head * head_dim)?; // fully written by q_gate_split
8303 let mut gate = vbuf(e, t * n_head * head_dim)?; // fully written by q_gate_split
8304 e.q_gate_split(&qf, &mut q, &mut gate, head_dim, n_head, t)?;
8305 (q, Some(gate))
8306 } else {
8307 (qf, None)
8308 };
8309
8310 let mut qn = vbuf(e, t * n_head * head_dim)?; // fully written by rms_norm
8311 e.rms_norm(
8312 &q,
8313 fa.q_norm.float_data(),
8314 &mut qn,
8315 head_dim,
8316 n_head * t,
8317 eps,
8318 )?;
8319 q = qn;
8320 let mut kn = vbuf(e, t * n_head_kv * head_dim)?; // fully written by rms_norm
8321 e.rms_norm(
8322 &k,
8323 fa.k_norm.float_data(),
8324 &mut kn,
8325 head_dim,
8326 n_head_kv * t,
8327 eps,
8328 )?;
8329 k = kn;
8330 let rope_dims = geometry.n_rot as usize;
8331 e.rope_neox(
8332 &mut q,
8333 pos_d,
8334 head_dim,
8335 rope_dims,
8336 n_head,
8337 t,
8338 geometry.rope_base,
8339 1.0,
8340 )?;
8341 e.rope_neox(
8342 &mut k,
8343 pos_d,
8344 head_dim,
8345 rope_dims,
8346 n_head_kv,
8347 t,
8348 geometry.rope_base,
8349 1.0,
8350 )?;
8351
8352 // append T new K/V columns to the resident QUANTIZED cache. k/v are token-major [T, kv_dim]
8353 // f32; append-quantize each of the T token rows into the byte cache (q8_0 K / q5_1 V).
8354 let kvl = cache.kv[il].as_mut().unwrap();
8355 let (kv_dim_k, kv_dim_v, ktb, vtb) =
8356 (kvl.kv_dim_k, kvl.kv_dim_v, kvl.k_tok_bytes, kvl.v_tok_bytes);
8357 if let Some(ctr) = stream_ctr {
8358 // stream: ONE batched append at the device counter (rows kernel = the per-view warp
8359 // math on a (block, token) grid, documented byte-identical); host len is a stale
8360 // LOWER BOUND under pre-issue (drain reconciles it).
8361 e.append_kv_quantized_rows_dc(
8362 &k,
8363 &v,
8364 &mut kvl.k,
8365 &mut kvl.v,
8366 ctr,
8367 t,
8368 kv_dim_k,
8369 kv_dim_v,
8370 ktb,
8371 vtb,
8372 crate::Engine::kv_fp8_on(),
8373 )?;
8374 } else {
8375 for i in 0..t {
8376 let k_row = k.slice(i * kv_dim_k..(i + 1) * kv_dim_k);
8377 let v_row = v.slice(i * kv_dim_v..(i + 1) * kv_dim_v);
8378 e.append_kv_quantized_view(
8379 &k_row,
8380 &v_row,
8381 &mut kvl.k,
8382 &mut kvl.v,
8383 kvl.len + i,
8384 kv_dim_k,
8385 kv_dim_v,
8386 ktb,
8387 vtb,
8388 crate::Engine::kv_fp8_on(),
8389 )?;
8390 }
8391 kvl.len += t;
8392 }
8393
8394 // BIT-IDENTICAL VERIFY ATTENTION (spec-exactness fix): the FP accumulation order must be
8395 // byte-for-byte identical to the eager decode path. fa_prefill uses a different tile size
8396 // (BLOCK_Q=64, BK=32) and online-softmax structure than fa_decode's split-K + combine,
8397 // which changes FP summation order and can flip argmax at tight logit margins. Query row r
8398 // attends to keys [0..base_len+r+1) — each successive row sees one more key (the causal
8399 // property). This matches eager: decode appends k at len, then fa_decode sees t_kv = len+1
8400 // keys. The verify appends all T tokens first but bounds the key range per row.
8401 //
8402 // MULTI-ROW FUSED PATH (the long-ctx spec fix, 2026-07-03): when every row takes the vec
8403 // kernel (base_len+1 >= FA_VEC_MIN_TKV), ONE fa_decode_rows launch executes the exact
8404 // per-row program for all T rows (grid.z = row, per-row n_splits from the same
8405 // fa_split_keys formula) — replacing T x (2 launches + 2 dtod copies + 5 partial allocs)
8406 // and multiplying resident CTAs by T on a latency-bound kernel. Bit-identical per row by
8407 // construction; kernel-check pins rows-vs-loop byte identity, run-spec is the end gate.
8408 // Short ctx (any row below the vec crossover) and MEMRA_NO_FA_VEC/MEMRA_FA_ROWS_OFF keep the
8409 // per-row loop (whose fa_decode picks scalar/vec per row exactly like eager decode).
8410 let mut attn = vbuf(e, t * n_head * head_dim)?; // fully written by every FA arm below
8411 let base_len = kvl.len - t; // KV len BEFORE this round's T tokens were appended
8412 // T=1 INCLUDED (2026-07-05): p-min cuts the draft to 1 in ~75% of rounds on hard
8413 // (agentic) content — the old t>1 gate sent those rounds to the per-row loop (262us/row
8414 // + q-row copy + per-row allocs vs 93us/row through the fused kernel at grid.z=1, same
8415 // program). nsys accounting: 1088 of 1456 verify FA launches were T=1 escapees.
8416 // LEAN T=1 ARM (MEMRA_SPEC_LEAN, close35): at t==1, q IS one row and fa_decode on it is
8417 // the EXACT eager decode dispatch (vec_q_v2 + combine_f32; the rows pair measured +50us
8418 // at m=1). Byte-identical: kernel-check pins rows-vs-loop identity, and the per-row loop
8419 // at t=1 is fa_decode on the same q with zero-offset copies. Gates arbitrate.
8420 if let Some(ctr) = stream_ctr {
8421 // STREAM ARM: causal base from the device counter; host kvl.len is a stale lower
8422 // bound used only for the split-sizing upper bound (+64 slack covers M pre-issued
8423 // rounds at K<=8). Views span the bound; per-row limits derive in-kernel.
8424 let upper = kvl.len + t + 64;
8425 let k_view = e.view_u8(&kvl.k, (upper.min(cache.max_ctx)) * ktb);
8426 let v_view = e.view_u8(&kvl.v, (upper.min(cache.max_ctx)) * vtb);
8427 e.fa_decode_rows_dc(
8428 &q,
8429 &k_view,
8430 &v_view,
8431 &mut attn,
8432 head_dim,
8433 n_head,
8434 n_head_kv,
8435 ctr,
8436 upper.min(cache.max_ctx),
8437 t,
8438 scale,
8439 ktb,
8440 vtb,
8441 0,
8442 false,
8443 )?;
8444 } else if spec_lean() && t == 1 {
8445 let t_kv = base_len + 1;
8446 let k_view = e.view_u8(&kvl.k, t_kv * ktb);
8447 let v_view = e.view_u8(&kvl.v, t_kv * vtb);
8448 e.fa_decode_kvmod(
8449 &q,
8450 &k_view,
8451 &v_view,
8452 &mut attn,
8453 head_dim,
8454 n_head,
8455 n_head_kv,
8456 t_kv,
8457 scale,
8458 ktb,
8459 vtb,
8460 crate::Engine::kv_fp8_on(),
8461 )?;
8462 } else if e.fa_rows_eligible(base_len, head_dim) {
8463 let k_view = e.view_u8(&kvl.k, (base_len + t) * ktb);
8464 let v_view = e.view_u8(&kvl.v, (base_len + t) * vtb);
8465 e.fa_decode_rows(
8466 &q,
8467 &k_view,
8468 &v_view,
8469 &mut attn,
8470 head_dim,
8471 n_head,
8472 n_head_kv,
8473 base_len,
8474 t,
8475 scale,
8476 ktb,
8477 vtb,
8478 None,
8479 false,
8480 crate::Engine::kv_fp8_on(),
8481 None,
8482 )?;
8483 } else {
8484 for r in 0..t {
8485 let t_kv_r = base_len + r + 1; // this row sees keys [0..t_kv_r)
8486 let k_view_r = e.view_u8(&kvl.k, t_kv_r * ktb);
8487 let v_view_r = e.view_u8(&kvl.v, t_kv_r * vtb);
8488 // copy q row into an owned buffer (fa_decode takes &CudaSlice, not CudaView)
8489 let mut q_row = vbuf(e, n_head * head_dim)?; // fully written by copy_view_into
8490 let q_src = q.slice(r * n_head * head_dim..(r + 1) * n_head * head_dim);
8491 e.copy_view_into(&mut q_row, 0, &q_src, n_head * head_dim)?;
8492 let mut attn_row = vbuf(e, n_head * head_dim)?; // fully written by fa_decode
8493 e.fa_decode_kvmod(
8494 &q_row,
8495 &k_view_r,
8496 &v_view_r,
8497 &mut attn_row,
8498 head_dim,
8499 n_head,
8500 n_head_kv,
8501 t_kv_r,
8502 scale,
8503 ktb,
8504 vtb,
8505 crate::Engine::kv_fp8_on(),
8506 )?;
8507 e.copy_into(
8508 &mut attn,
8509 r * n_head * head_dim,
8510 &attn_row,
8511 n_head * head_dim,
8512 )?;
8513 }
8514 }
8515
8516 let attn_g = match &gate {
8517 Some(gate) => {
8518 let mut gsig = vbuf(e, t * n_head * head_dim)?; // fully written by sigmoid
8519 e.sigmoid(gate, &mut gsig, t * n_head * head_dim)?;
8520 let mut ag = vbuf(e, t * n_head * head_dim)?; // fully written by mul
8521 e.mul(&attn, &gsig, &mut ag, t * n_head * head_dim)?;
8522 ag
8523 }
8524 None => attn,
8525 };
8526 // DECODE-EXACT wo projection: at m>=5 (K=4+ with pending) the generic matmul would use dp4a
8527 // (128-thread, different FP sum order than MMVQ). Force MMVQ for bit-identity with decode.
8528 Ok(e.matmul_decode_exact(&fa.wo, &attn_g, t)?)
8529 }
8530
8531 /// Context-linear bytes for a plain serving session's trunk cache.
8532 pub fn plain_session_kv_bytes_per_token(&self) -> usize {
8533 crate::cache::cache_bytes_per_token_for_plan(
8534 &self.cfg,
8535 &self.plan,
8536 0,
8537 self.plan.layers.len(),
8538 )
8539 }
8540
8541 /// `(logical bytes/token, ring-capped bytes/token, ring row cap)` for exact admission.
8542 pub fn plain_session_kv_shape(&self) -> (usize, usize, usize) {
8543 (
8544 self.plain_session_kv_bytes_per_token(),
8545 crate::cache::cache_ring_bytes_per_token_for_plan(
8546 &self.cfg,
8547 &self.plan,
8548 0,
8549 self.plan.layers.len(),
8550 ),
8551 crate::cache::cache_ring_row_cap_for_plan(&self.plan),
8552 )
8553 }
8554
8555 /// Context-linear bytes for a speculative serving session: trunk cache plus persistent MTP
8556 /// scratch. With no MTP head this equals the plain coefficient.
8557 pub fn spec_session_kv_bytes_per_token(&self) -> usize {
8558 let scratch = self
8559 .mtp
8560 .iter()
8561 .chain(self.mtp_extra.iter())
8562 .map(|mtp| {
8563 let (_, _, k, v) = mtp_scratch_layout(&self.cfg, mtp.geom.as_ref());
8564 k + v
8565 })
8566 .sum::<usize>();
8567 self.plain_session_kv_bytes_per_token()
8568 .saturating_add(scratch)
8569 }
8570
8571 /// Spec twin of [`HybridModel::plain_session_kv_shape`]; Step35's persistent MTP scratch is
8572 /// capped by the same SWA ring rows as the trunk.
8573 pub fn spec_session_kv_shape(&self) -> (usize, usize, usize) {
8574 let total = self.spec_session_kv_bytes_per_token();
8575 let (_, mut ring, rows) = self.plain_session_kv_shape();
8576 if rows > 0 {
8577 ring = ring.saturating_add(
8578 self.mtp
8579 .iter()
8580 .chain(self.mtp_extra.iter())
8581 .map(|mtp| {
8582 let (_, _, k, v) = mtp_scratch_layout(&self.cfg, mtp.geom.as_ref());
8583 k + v
8584 })
8585 .sum::<usize>(),
8586 );
8587 }
8588 (total, ring, rows)
8589 }
8590
8591 /// Greedy MTP speculative decode (§B). Token-identical to `generate(prompt, max_new)` but uses
8592 /// the NextN head to draft K tokens then verifies them in one batched target forward.
8593 /// Returns (generated tokens, total_drafted, total_accepted) so the caller can report
8594 /// acceptance rate. `k` = draft length per round.
8595 ///
8596 /// GRAPH DRAFT (stage 2 of graph-grade spec): when the model is all-Dense and the MTP head is
8597 /// Dense (no MoE host readbacks), the fixed-shape T=1 MTP forward is CUDA-graph-captured ONCE
8598 /// and replayed per draft step — the ~40 eager launches per drafted token collapse into one
8599 /// graph dispatch; only the 4-byte token id (and 4-byte p-min confidence) round-trip per step.
8600 /// Event tracking is disabled for the whole call (generate_graph pattern) so every buffer the
8601 /// captured graph references is event-free; the spec loop is strictly single-stream.
8602 /// MEMRA_SPEC_NOGRAPH=1 forces the eager draft chain.
8603 /// SAMPLED mode (MEMRA_SPEC_TEMP>0) has its OWN capture (gumbel-perturbed in-graph argmax,
8604 /// device Philox event counter, persistent q retention) — graph-vs-eager sampled streams are
8605 /// bit-identical for the same (seed, prompt, K, temp); see the sampled-graph setup in
8606 /// generate_spec_inner2.
8607 /// Multi-turn session: trunk cache + MTP draft scratch persist across generate calls, so
8608 /// turn N+1 primes ONLY its new suffix (the 124k-conversation daily pattern — re-priming a
8609 /// 32k history costs ~54s; a suffix prime costs seconds). APPEND-ONLY by construction: the
8610 /// hybrid linear-attn states are in-place (no position index), so a session can extend but
8611 /// never rewind — `committed` is the exact token list whose state the caches hold (includes
8612 /// any overshoot tokens past max_new; the caller renders from `committed`, not its own echo).
8613 pub fn new_session(
8614 &self,
8615 e: &Engine,
8616 max_ctx: usize,
8617 ) -> Result<SpecSession, Box<dyn std::error::Error>> {
8618 Ok(SpecSession {
8619 // STAGE-OWNED KV (lane/pp2-spec 2026-08-06): `pp::new_cache`, not `Cache::new`. This
8620 // is the SERVING spec-session path, and with the ppN door open across two cards a
8621 // primary-homed cache makes every remote stage peer-read its OWN KV on every verify
8622 // round — the wrong-card class already fixed on the two batched serving paths
8623 // (worker.rs 2483 / 2837). With the door shut `new_cache` IS `Cache::new` (same
8624 // branch, same allocations), so single-device behavior is byte-unchanged.
8625 cache: crate::pp::new_cache_planned(e, &self.cfg, &self.plan, max_ctx)?,
8626 scratch: self.new_mtp_scratch(e, max_ctx)?,
8627 committed: Vec::new(),
8628 last_h: None,
8629 next_pred: None,
8630 sctr: 0,
8631 uctr: 0,
8632 draft_ctx: None,
8633 pending_tok: None,
8634 turn_ckpt: None,
8635 telem: SpecTelemetryCounters::default(),
8636 capture_at: None,
8637 boundary_captures: Vec::new(),
8638 ckpt_at: None,
8639 })
8640 }
8641
8642 /// SPEC-ON-CACHE-HIT restore (lane/spec-on-cache-hit, 2026-08-18 — PORT-PLAN item 3,
8643 /// research/cache-spec-design-20260814, scoped to WHOLE-ENTRY restores only): build a
8644 /// SpecSession around a trunk cache the worker already restored from a prefix-cache
8645 /// entry, re-installing the entry's published draft plane as the MTP scratch rows
8646 /// `[0..prefix.len())` and the entry's boundary hidden as `last_h`, then feeding the
8647 /// prompt SUFFIX here — through EXACTLY the plain path's program selection — so the
8648 /// worker always receives a fully-warm continuation session (committed = whole
8649 /// prompt, `next_pred` + `last_h` set; caller sets `next_pred` from the entry's
8650 /// boundary logits on the empty-suffix shape).
8651 ///
8652 /// PROGRAM LAW (the splitiso two-programs class, learned AGAIN in this lane's own
8653 /// gate): the identity target for a converted hit is the PLAIN hit serving the same
8654 /// request, and plain feeds a carried suffix via eager `decode_step` below
8655 /// PRIME_MIN_T and via `prime_cache` at/above it (prefill_tick's arms). The generate
8656 /// path's tokenwise arm routes qwen35-class through the BATCHED T=1 program
8657 /// (`spec_target_step_h`) instead — ULP-different suffix rows, and the gate measured
8658 /// the near-tie flip at generated token ~8 (research/spec-cache-20260818, qwen r3).
8659 /// So the suffix is fed HERE, mirroring prefill_tick arm-for-arm, not handed to the
8660 /// burst prime.
8661 ///
8662 /// SEED RULE (both sampling regimes; lane/sampled-hit-spec 2026-08-19, sampled draw
8663 /// added by lane/sampled-spec-quality 2026-08-19). The boundary token is produced by
8664 /// EXACTLY the rule the cold burst entry applies to its own first token from the same
8665 /// logits row: `argmax` when greedy, and a `sample_boundary_token` draw at Philox
8666 /// counter 0 when sampled. Both shapes are covered — the entry's boundary logits on a
8667 /// full-cover (empty-suffix) hit, this feed's own boundary logits on a suffix hit.
8668 /// That is what keeps a restored session seed-identical to a cold one PER SEED: the
8669 /// cold session draws from the identical row at counter 0 and then runs its rounds from
8670 /// counter 1, so the restored session admits with `sctr = 1` after its own draw.
8671 /// The WORKER owns the one refusal this constructor cannot see — a constrained request.
8672 /// (The penalized-sampled refusal was LIFTED once the burst's penalty window learned to
8673 /// span the session: `committed` here is the WHOLE prompt, so the restored session's
8674 /// window is the cold session's window. It comes back if `MEMRA_SPEC_PEN_SESSION=0`.)
8675 ///
8676 /// NOT the rolled-back partial-restore hazard: the caller restores at exactly the
8677 /// entry's captured endpoint (`e.pos`) through the shipping whole-entry path;
8678 /// mid-entry (`at < e.pos`) trunk restores stay behind MEMRA_PREFIX_PARTIAL_RESTORE
8679 /// and are never routed here.
8680 ///
8681 /// Failure contract: `Err((Some(cache), why))` before any trunk mutation — the
8682 /// worker rebuilds the plain carrier and the hit serves plain, byte-unchanged.
8683 /// `Err((None, why))` after the suffix feed began — the carrier is part-fed and
8684 /// UNUSABLE; the worker serves the request cold-plain (correct, slower) and the
8685 /// entry stays published for the next request.
8686 #[allow(clippy::too_many_arguments)]
8687 pub fn spec_session_from_restored(
8688 &self,
8689 e: &Engine,
8690 mut cache: Cache,
8691 prefix: Vec<u32>,
8692 suffix: &[u32],
8693 draft_k: &CudaSlice<u8>,
8694 draft_v: &CudaSlice<u8>,
8695 draft_k_tok_bytes: usize,
8696 draft_v_tok_bytes: usize,
8697 draft_len: usize,
8698 last_h: &[f32],
8699 // The ENTRY's boundary logits row (the full-cover shape's seed source). May be empty
8700 // when a suffix follows — the feed's own logits are the boundary then.
8701 boundary_logits: &[f32],
8702 // The request's sampler, or None for greedy. Owned here so the seed rule lives in
8703 // ONE place instead of being half-applied by the worker.
8704 sampling: Option<SpecSampling>,
8705 require_anchor: bool,
8706 max_ctx: usize,
8707 // STABLE-BOUNDARY REPUBLICATION (lane/frspec-multiturn-cache, 2026-08-21): ABSOLUTE
8708 // prompt position to split the suffix feed at and capture the extended-entry
8709 // publication + this session's `turn_ckpt` — the worker's stable pre-generation
8710 // boundary (`plain_checkpoint_boundary`). None = legacy prompt-end republication.
8711 // WHY: the prompt-end capture below includes the template's live generation header
8712 // (`<|im_start|>assistant\n<think>\n`), which the next turn's re-render replaces, so
8713 // for a hybrid (whole-entry restores only) every extended entry's last ~2 tokens
8714 // diverged from every future prompt and the hit boundary FROZE at the first
8715 // lcp-split entry forever (measured: cached 6811 of 38228 by turn 8, B4).
8716 republish_at: Option<usize>,
8717 ) -> Result<SpecSession, (Option<Cache>, String)> {
8718 let pos = prefix.len();
8719 let fail = |cache: Cache, msg: String| -> Result<SpecSession, (Option<Cache>, String)> {
8720 Err((Some(cache), msg))
8721 };
8722 if self.mtp.is_none() {
8723 return fail(cache, "no MTP head attached (nothing to draft with)".into());
8724 }
8725 if pos == 0 {
8726 return fail(cache, "empty committed prefix".into());
8727 }
8728 if cache.pos != pos {
8729 let msg = format!(
8730 "restored cache pos {} != restored prefix len {pos}",
8731 cache.pos
8732 );
8733 return fail(cache, msg);
8734 }
8735 if draft_len != pos {
8736 return fail(
8737 cache,
8738 format!("draft plane len {draft_len} != restored prefix len {pos}"),
8739 );
8740 }
8741 if pos + suffix.len() >= max_ctx {
8742 return fail(
8743 cache,
8744 format!(
8745 "prompt {} + suffix would not leave generation room in ctx {max_ctx}",
8746 pos + suffix.len(),
8747 ),
8748 );
8749 }
8750 let mut scratch = match MtpScratch::new(
8751 e,
8752 &self.cfg,
8753 &self.plan,
8754 max_ctx,
8755 self.mtp.as_ref().and_then(|m| m.geom.as_ref()),
8756 ) {
8757 Ok(s) => s,
8758 Err(err) => return fail(cache, format!("draft scratch alloc failed: {err}")),
8759 };
8760 if scratch.kv.ring.is_some() {
8761 return fail(
8762 cache,
8763 "ring-backed draft scratch (Step35 SWA) cannot take a flat prefix restore".into(),
8764 );
8765 }
8766 if scratch.kv.k_tok_bytes != draft_k_tok_bytes
8767 || scratch.kv.v_tok_bytes != draft_v_tok_bytes
8768 {
8769 return fail(
8770 cache,
8771 format!(
8772 "draft plane layout {draft_k_tok_bytes}/{draft_v_tok_bytes} != scratch \
8773 {}/{} bytes/token (stale entry across a format change)",
8774 scratch.kv.k_tok_bytes, scratch.kv.v_tok_bytes,
8775 ),
8776 );
8777 }
8778 if pos > scratch.cap {
8779 return fail(
8780 cache,
8781 format!(
8782 "draft plane rows {pos} exceed scratch capacity {}",
8783 scratch.cap
8784 ),
8785 );
8786 }
8787 let kb = pos * draft_k_tok_bytes;
8788 let vb = pos * draft_v_tok_bytes;
8789 if draft_k.len() < kb || draft_v.len() < vb {
8790 return fail(
8791 cache,
8792 format!(
8793 "truncated draft plane: K {} < {kb} or V {} < {vb} bytes",
8794 draft_k.len(),
8795 draft_v.len(),
8796 ),
8797 );
8798 }
8799 if kb > 0 {
8800 if let Err(err) = e.copy_u8_into(&mut scratch.kv.k, 0, draft_k, kb) {
8801 return fail(cache, format!("draft K restore copy failed: {err}"));
8802 }
8803 }
8804 if vb > 0 {
8805 if let Err(err) = e.copy_u8_into(&mut scratch.kv.v, 0, draft_v, vb) {
8806 return fail(cache, format!("draft V restore copy failed: {err}"));
8807 }
8808 }
8809 if let Err(err) = scratch.set_len(e, pos) {
8810 return fail(cache, format!("draft scratch len set failed: {err}"));
8811 }
8812 let mut last_h_dev = if last_h.len() == self.cfg.n_embd as usize {
8813 // anchor upload failure is acceptance-only when a suffix feed follows (fill
8814 // row-0 falls back to zeros) but FATAL for an empty-suffix continuation (the
8815 // burst entry asserts committed + last_h + next_pred) — the caller says which.
8816 e.htod(last_h).ok()
8817 } else {
8818 None
8819 };
8820 if require_anchor && last_h_dev.is_none() {
8821 return fail(
8822 cache,
8823 "empty-suffix continuation requires the entry's boundary hidden anchor".into(),
8824 );
8825 }
8826 let mut committed = prefix;
8827 // Set on BOTH shapes below (suffix-fed and full-cover) — never left None, which is
8828 // what the empty-suffix continuation assert in the burst entry requires.
8829 let next_pred;
8830 // Philox: (0,0) at admit exactly like a fresh session; a sampled boundary draw below
8831 // consumes counter 0 and leaves 1, which is the state a cold session reaches after
8832 // drawing its own first token from the same row.
8833 let mut sctr = 0u32;
8834 let sampled = sampling.is_some_and(|s| s.temp > 0.0) && spec_sampled_boundary_on();
8835 // Penalty window for the boundary draw: the last `penalty_last_n` tokens of the WHOLE
8836 // prompt, which is what the cold session's own burst sees (Item 2's window). Built
8837 // after the suffix joins `committed` below.
8838 let mut boundary_captures: Vec<SpecBoundaryCapture> = Vec::new();
8839 let mut restored_turn_ckpt: Option<SpecCheckpoint> = None;
8840 if !suffix.is_empty() {
8841 // ---- SUFFIX FEED, mirroring prefill_tick's program selection exactly ----
8842 // From here on the trunk cache mutates: failures return Err((None, _)) and
8843 // the worker serves the request cold-plain instead of reusing the carrier.
8844 let dirty =
8845 |msg: String| -> Result<SpecSession, (Option<Cache>, String)> { Err((None, msg)) };
8846 let n_embd = self.cfg.n_embd as usize;
8847 let t = suffix.len();
8848 let mut h_rows = match e.uninit(t * n_embd) {
8849 Ok(b) => b,
8850 Err(err) => return fail(cache, format!("suffix hidden buffer alloc: {err}")),
8851 };
8852 // STABLE-BOUNDARY split (see `republish_at`): feed stops at the boundary so the
8853 // in-place GDN conv/ssm state can be snapshotted there — the only moment it
8854 // exists (the cold prime-split law). suffix-relative; None = one-segment legacy.
8855 let b_rel = republish_at
8856 .and_then(|abs| abs.checked_sub(pos))
8857 .filter(|&r| r > 0 && r < t);
8858 let mut feed_logits = Vec::new();
8859 let tokenwise_env = std::env::var("MEMRA_PRIME_TOKENWISE").is_ok()
8860 || e.frozen_cpu_experts_prefer_tokenwise_prime();
8861 let mut fed = 0usize;
8862 for seg_end in [b_rel, Some(t)].into_iter().flatten() {
8863 if seg_end <= fed {
8864 continue;
8865 }
8866 let seg = &suffix[fed..seg_end];
8867 let batched = seg.len() >= crate::hybrid_forward::PRIME_MIN_T && !tokenwise_env;
8868 if batched {
8869 // prefill_tick's prime arm: request-level prime_cache call; tokens still
8870 // queued after this segment ride `queued_after` so Step35 arm selection
8871 // stays keyed to the request's end (tick-seg law).
8872 match self.prime_cache(e, seg, &mut cache, t - seg_end) {
8873 Ok((l, _h_seed, hiddens)) => {
8874 if let Err(err) =
8875 e.copy_into(&mut h_rows, fed * n_embd, &hiddens, seg.len() * n_embd)
8876 {
8877 return dirty(format!("suffix hidden copy: {err}"));
8878 }
8879 feed_logits = l;
8880 }
8881 Err(err) => return dirty(format!("suffix prime failed: {err}")),
8882 }
8883 } else {
8884 // prefill_tick's tokenwise arm: eager decode_step, one token at a time.
8885 for (i, &tok) in seg.iter().enumerate() {
8886 match self.decode_step_h(e, tok, &mut cache) {
8887 Ok((l, h)) => {
8888 if let Err(err) =
8889 e.copy_into(&mut h_rows, (fed + i) * n_embd, &h, n_embd)
8890 {
8891 return dirty(format!("suffix hidden copy: {err}"));
8892 }
8893 feed_logits = l;
8894 }
8895 Err(err) => return dirty(format!("suffix decode_step failed: {err}")),
8896 }
8897 }
8898 }
8899 fed = seg_end;
8900 if Some(seg_end) == b_rel {
8901 // The stable pre-generation boundary: capture the extended-entry
8902 // publication AND this session's own turn checkpoint here instead of at
8903 // prompt-end (both would otherwise carry the volatile live-header tail
8904 // the next re-render replaces). Failure silent, turn_ckpt convention.
8905 debug_assert_eq!(
8906 cache.pos,
8907 pos + seg_end,
8908 "stable-boundary capture off the feed split"
8909 );
8910 if spec_restore_republish_on() {
8911 if let Ok(snap) = cache.snapshot(e) {
8912 boundary_captures.push(SpecBoundaryCapture {
8913 snap,
8914 pos: pos + seg_end,
8915 logits: feed_logits.clone(),
8916 last_h: capture_boundary_hidden(e, &h_rows, seg_end, n_embd),
8917 });
8918 }
8919 }
8920 let anchor: Result<CudaSlice<f32>, Box<dyn std::error::Error>> =
8921 e.uninit(n_embd).and_then(|mut a| {
8922 e.copy_view_into(
8923 &mut a,
8924 0,
8925 &h_rows.slice((seg_end - 1) * n_embd..seg_end * n_embd),
8926 n_embd,
8927 )?;
8928 Ok(a)
8929 });
8930 if let (Ok(snap), Ok(last_h)) = (cache.snapshot(e), anchor) {
8931 restored_turn_ckpt = Some(SpecCheckpoint {
8932 snap,
8933 pos: pos + seg_end,
8934 last_h,
8935 });
8936 }
8937 }
8938 }
8939 // Draft-scratch fill for the suffix rows, predecessor-paired: row `pos` reads
8940 // the entry's boundary anchor (zeros fallback — acceptance-only), row `pos+i`
8941 // reads h_rows[i-1]. Chunked like the generate path's fill (transients scale
8942 // with T). Fill failures are acceptance-only — truncate to the restored rows
8943 // and continue; the burst's own set_len keeps the invariant.
8944 let mtp = self.mtp.as_ref().expect("mtp checked above");
8945 let (embd_qt, embd_rb) = self.embd.qt_and_row_bytes(n_embd);
8946 let embd_gpu = if spec_host_embd() {
8947 None
8948 } else {
8949 Some(
8950 self.embd_gpu
8951 .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload")),
8952 )
8953 };
8954 let embd_dev = embd_gpu.map(|g| (g, embd_qt, embd_rb));
8955 let fill_chunk = 4096usize;
8956 let mut filled = true;
8957 let mut start = 0usize;
8958 'fill: while start < t {
8959 let end = (start + fill_chunk).min(t);
8960 let tc = end - start;
8961 let Ok(mut phs) = e.zeros(tc * n_embd) else {
8962 filled = false;
8963 break 'fill;
8964 };
8965 let (src_lo, dst_off, n_copy) = if start == 0 {
8966 (0, n_embd, (tc - 1) * n_embd)
8967 } else {
8968 ((start - 1) * n_embd, 0, tc * n_embd)
8969 };
8970 if start == 0 {
8971 if let Some(lh) = last_h_dev.as_ref() {
8972 if e.copy_into(&mut phs, 0, lh, n_embd).is_err() {
8973 filled = false;
8974 break 'fill;
8975 }
8976 }
8977 }
8978 if n_copy > 0
8979 && e.copy_view_into(
8980 &mut phs,
8981 dst_off,
8982 &h_rows.slice(src_lo..src_lo + n_copy),
8983 n_copy,
8984 )
8985 .is_err()
8986 {
8987 filled = false;
8988 break 'fill;
8989 }
8990 if self
8991 .mtp_kv_fill_all(
8992 e,
8993 &suffix[start..end],
8994 &phs,
8995 pos + start,
8996 &mut scratch,
8997 embd_dev,
8998 )
8999 .is_err()
9000 {
9001 filled = false;
9002 break 'fill;
9003 }
9004 start = end;
9005 }
9006 if !filled {
9007 // acceptance-only: drafts over missing suffix rows are cheap and wrong,
9008 // so keep only the restored rows resident and let verify arbitrate.
9009 if let Err(err) = scratch.set_len(e, pos) {
9010 return dirty(format!("scratch truncation after failed fill: {err}"));
9011 }
9012 }
9013 // EXTENDED-ENTRY PUBLICATION (lane/sampled-spec-quality, Item 3 — the fix for
9014 // "a restored spec session never publishes an extended entry", SAMPLED-HIT.md
9015 // finding (d)). Pre-lane, publication was armed only for COLD sessions
9016 // (`spec_resumed == 0` in the worker) and both engine capture sites require a
9017 // non-continuation burst — but a converted hit's first burst IS a continuation,
9018 // so a growing conversation learned exactly ONE boundary and turn 3 could never
9019 // hit a longer prefix than turn 2 did.
9020 //
9021 // WHERE, and why it is safe here: `cache.pos == prefix + suffix` at this exact
9022 // line — the trunk is primed over the whole prompt, nothing is generated, and the
9023 // draft plane rows [0..prompt) are filled just above. That is a complete
9024 // whole-entry boundary (`pos == fed_len`), the same shape the cold seed capture
9025 // publishes; the worker's existing publication sweep picks it up because it is
9026 // keyed on non-empty `boundary_captures` and is sampler- and resume-independent.
9027 // NOT the partial-restore hazard: the boundary is this session's own prompt END,
9028 // never mid-entry, so `entry_pos != fed_len` still refuses on the way back in.
9029 // Failure is SILENT by design (the turn_ckpt / boundary-capture convention):
9030 // publication is an optimization, never a correctness dependency.
9031 //
9032 // SUPERSEDED WHEN `republish_at` FIRED (lane/frspec-multiturn-cache): a prompt-end
9033 // entry's tail is the live generation header the next re-render replaces, so on a
9034 // hybrid (whole-entry restores) it can never serve the conversation's next turn —
9035 // the stable-boundary capture above IS this publication, minus the poisoned tail.
9036 if spec_restore_republish_on() && boundary_captures.is_empty() {
9037 debug_assert_eq!(
9038 cache.pos,
9039 pos + t,
9040 "extended-entry capture must sit at the restored session's prompt end",
9041 );
9042 if let Ok(snap) = cache.snapshot(e) {
9043 boundary_captures.push(SpecBoundaryCapture {
9044 snap,
9045 pos: pos + t,
9046 logits: feed_logits.clone(),
9047 last_h: capture_boundary_hidden(e, &h_rows, t, n_embd),
9048 });
9049 }
9050 }
9051 // continuation seed: the feed's boundary logits ARE the plain path's boundary
9052 // logits (same program), so greedy's argmax here is plain's first emitted token,
9053 // and the sampled draw is the cold sampled session's own first token.
9054 next_pred = Some(if sampled {
9055 let sp = sampling.expect("sampled implies a sampler");
9056 // `committed` is still the restored prefix here; the suffix joins it below —
9057 // so this is the last-N window over the WHOLE prompt, exactly the cold
9058 // session's own window at its first token.
9059 let hist = pen_window_seed(&committed, suffix, sp.penalty_last_n);
9060 match sample_boundary_token(
9061 e,
9062 &feed_logits,
9063 &sp,
9064 &hist,
9065 &mut sctr,
9066 "restore-suffix-feed",
9067 ) {
9068 Ok(t) => t,
9069 // the trunk is already fed: hand nothing back, the worker serves the
9070 // request cold-plain. Never fall back to an argmax — that would put a
9071 // greedy token in a sampled stream to save a slow path.
9072 Err(err) => {
9073 return dirty(format!("boundary token draw failed: {err}"));
9074 }
9075 }
9076 } else {
9077 argmax(&feed_logits) as u32
9078 });
9079 let mut lh = match e.uninit(n_embd) {
9080 Ok(b) => b,
9081 Err(err) => return dirty(format!("boundary hidden alloc: {err}")),
9082 };
9083 if let Err(err) = e.copy_view_into(
9084 &mut lh,
9085 0,
9086 &h_rows.slice((t - 1) * n_embd..t * n_embd),
9087 n_embd,
9088 ) {
9089 return dirty(format!("boundary hidden copy: {err}"));
9090 }
9091 last_h_dev = Some(lh);
9092 committed.extend_from_slice(suffix);
9093 } else {
9094 // FULL-COVER shape (empty suffix — the identical-repeat / agent-loop shape): the
9095 // ENTRY's boundary logits are the boundary row, and this is the token the cold
9096 // session emits from that same row. Owned here rather than in the worker so the
9097 // sampled draw cannot be half-applied on one shape (the worker used to argmax it).
9098 if boundary_logits.is_empty() {
9099 return fail(
9100 cache,
9101 "full-cover restore without the entry's boundary logits".into(),
9102 );
9103 }
9104 next_pred = Some(if sampled {
9105 let sp = sampling.expect("sampled implies a sampler");
9106 let hist = pen_window_seed(&committed, &[], sp.penalty_last_n);
9107 match sample_boundary_token(
9108 e,
9109 boundary_logits,
9110 &sp,
9111 &hist,
9112 &mut sctr,
9113 "restore-full-cover",
9114 ) {
9115 Ok(t) => t,
9116 // nothing has been mutated on this shape — hand the carrier back and let
9117 // the hit serve PLAIN (the banked pre-lane path).
9118 Err(err) => {
9119 return fail(cache, format!("boundary token draw failed: {err}"));
9120 }
9121 }
9122 } else {
9123 argmax(boundary_logits) as u32
9124 });
9125 }
9126 Ok(SpecSession {
9127 cache,
9128 scratch,
9129 committed,
9130 last_h: last_h_dev,
9131 next_pred,
9132 sctr,
9133 uctr: 0,
9134 draft_ctx: None,
9135 pending_tok: None,
9136 // Stable-boundary capture from the split feed above (None on the legacy shape):
9137 // a restored session previously parked WITHOUT a checkpoint, so the next turn's
9138 // affinity probe declined ("no turn checkpoint retained") and the conversation
9139 // fell back to the frozen prefix entry forever.
9140 turn_ckpt: restored_turn_ckpt,
9141 telem: SpecTelemetryCounters::default(),
9142 capture_at: None,
9143 boundary_captures,
9144 ckpt_at: None,
9145 })
9146 }
9147
9148 /// Forced-gate exact state comparison. This intentionally reads the real live prefixes from
9149 /// their owning PP devices: matching emitted ids alone would miss a stale `len_d`, recurrent
9150 /// snapshot, or draft-KV row that only corrupts the following round.
9151 pub fn optipipe_compare_session_state(
9152 &self,
9153 e: &Engine,
9154 reference: &SpecSession,
9155 candidate: &SpecSession,
9156 ) -> Result<OptiForkStateIdentity, Box<dyn std::error::Error>> {
9157 fn fail(what: &str) -> Box<dyn std::error::Error> {
9158 format!("optipipe state mismatch: {what}").into()
9159 }
9160 fn same_f32(a: &[f32], b: &[f32]) -> bool {
9161 a.len() == b.len() && a.iter().zip(b).all(|(x, y)| x.to_bits() == y.to_bits())
9162 }
9163 fn compare_layers(
9164 es: &Engine,
9165 range: std::ops::Range<usize>,
9166 reference: &SpecSession,
9167 candidate: &SpecSession,
9168 report: &mut OptiForkStateIdentity,
9169 ) -> Result<(), Box<dyn std::error::Error>> {
9170 for il in range {
9171 match (&reference.cache.kv[il], &candidate.cache.kv[il]) {
9172 (Some(a), Some(b)) => {
9173 if a.len != b.len {
9174 return Err(fail(&format!(
9175 "layer {il} host KV len {} != {}",
9176 a.len, b.len
9177 )));
9178 }
9179 let ad = es.dtoh_i32(&a.len_d)?;
9180 let bd = es.dtoh_i32(&b.len_d)?;
9181 if ad != bd || ad.first().copied() != Some(a.len as i32) {
9182 return Err(fail(&format!(
9183 "layer {il} device KV len {ad:?} != {bd:?} (host={})",
9184 a.len,
9185 )));
9186 }
9187 let kb = a.len * a.k_tok_bytes;
9188 let vb = a.len * a.v_tok_bytes;
9189 if kb > 0 {
9190 let ak = es.dtoh_u8_view(&a.k.slice(0..kb))?;
9191 let bk = es.dtoh_u8_view(&b.k.slice(0..kb))?;
9192 if ak != bk {
9193 let at = ak.iter().zip(&bk).position(|(x, y)| x != y).unwrap();
9194 return Err(fail(&format!(
9195 "layer {il} K bytes at byte {at} row {} offset {}: {} != {}",
9196 at / a.k_tok_bytes,
9197 at % a.k_tok_bytes,
9198 ak[at],
9199 bk[at],
9200 )));
9201 }
9202 }
9203 if vb > 0 {
9204 let av = es.dtoh_u8_view(&a.v.slice(0..vb))?;
9205 let bv = es.dtoh_u8_view(&b.v.slice(0..vb))?;
9206 if av != bv {
9207 let at = av.iter().zip(&bv).position(|(x, y)| x != y).unwrap();
9208 return Err(fail(&format!(
9209 "layer {il} V bytes at byte {at} row {} offset {}: {} != {}",
9210 at / a.v_tok_bytes,
9211 at % a.v_tok_bytes,
9212 av[at],
9213 bv[at],
9214 )));
9215 }
9216 }
9217 report.trunk_kv_bytes += kb + vb;
9218 }
9219 (None, None) => {}
9220 _ => return Err(fail(&format!("layer {il} KV presence"))),
9221 }
9222 match (&reference.cache.recur[il], &candidate.cache.recur[il]) {
9223 (Some(a), Some(b)) => {
9224 let ac = es.dtoh(&a.conv_state)?;
9225 let bc = es.dtoh(&b.conv_state)?;
9226 if !same_f32(&ac, &bc) {
9227 return Err(fail(&format!("layer {il} conv state")));
9228 }
9229 let as_ = es.dtoh(&a.ssm_state)?;
9230 let bs = es.dtoh(&b.ssm_state)?;
9231 if !same_f32(&as_, &bs) {
9232 return Err(fail(&format!("layer {il} SSM state")));
9233 }
9234 report.recurrent_bytes += (ac.len() + as_.len()) * 4;
9235 }
9236 (None, None) => {}
9237 _ => return Err(fail(&format!("layer {il} recurrent presence"))),
9238 }
9239 }
9240 Ok(())
9241 }
9242
9243 if reference.committed != candidate.committed {
9244 return Err(fail("committed token ids"));
9245 }
9246 if reference.cache.pos != candidate.cache.pos
9247 || reference.cache.max_ctx != candidate.cache.max_ctx
9248 {
9249 return Err(fail("cache pos/capacity"));
9250 }
9251 if reference.pending_tok != candidate.pending_tok
9252 || reference.next_pred != candidate.next_pred
9253 || reference.sctr != candidate.sctr
9254 || reference.uctr != candidate.uctr
9255 {
9256 return Err(fail("pending/prediction/counter tail"));
9257 }
9258
9259 let mut report = OptiForkStateIdentity::default();
9260 if let Some(fence) = crate::pp::pp_cuts(self.layers.len()) {
9261 let rt = crate::pp::PpNRt::get(e)?;
9262 for stage in 0..rt.n_stages() {
9263 let _scope = rt.enter(stage);
9264 compare_layers(
9265 rt.engine(stage, e),
9266 fence[stage]..fence[stage + 1],
9267 reference,
9268 candidate,
9269 &mut report,
9270 )?;
9271 }
9272 } else {
9273 compare_layers(e, 0..self.layers.len(), reference, candidate, &mut report)?;
9274 }
9275
9276 if reference.scratch.plane_count() != candidate.scratch.plane_count() {
9277 return Err(fail("draft scratch plane count"));
9278 }
9279 for index in 0..reference.scratch.plane_count() {
9280 let (a, _) = reference.scratch.plane(index);
9281 let (b, _) = candidate.scratch.plane(index);
9282 if a.len != b.len
9283 || a.kv_dim_k != b.kv_dim_k
9284 || a.kv_dim_v != b.kv_dim_v
9285 || a.k_tok_bytes != b.k_tok_bytes
9286 || a.v_tok_bytes != b.v_tok_bytes
9287 || e.dtoh_i32(&a.len_d)? != e.dtoh_i32(&b.len_d)?
9288 {
9289 return Err(fail(&format!("draft scratch plane {index} length/layout")));
9290 }
9291 let kb = a.len * a.k_tok_bytes;
9292 let vb = a.len * a.v_tok_bytes;
9293 if kb > 0 && e.dtoh_u8_view(&a.k.slice(0..kb))? != e.dtoh_u8_view(&b.k.slice(0..kb))? {
9294 return Err(fail(&format!("draft scratch plane {index} K bytes")));
9295 }
9296 if vb > 0 && e.dtoh_u8_view(&a.v.slice(0..vb))? != e.dtoh_u8_view(&b.v.slice(0..vb))? {
9297 return Err(fail(&format!("draft scratch plane {index} V bytes")));
9298 }
9299 report.scratch_kv_bytes += kb + vb;
9300 }
9301
9302 match (&reference.last_h, &candidate.last_h) {
9303 (Some(a), Some(b)) => {
9304 let ah = e.dtoh(a)?;
9305 let bh = e.dtoh(b)?;
9306 if !same_f32(&ah, &bh) {
9307 return Err(fail("last hidden/seed bytes"));
9308 }
9309 report.hidden_bytes = ah.len() * 4;
9310 }
9311 (None, None) => {}
9312 _ => return Err(fail("last hidden/seed presence")),
9313 }
9314 Ok(report)
9315 }
9316
9317 /// SESSION-AFFINITY REWIND (lane/session-affinity, 2026-08-05): roll `sess` back to its
9318 /// retained prompt-end checkpoint, so a request whose prompt matches
9319 /// `committed[..rewind_pos()]` exactly can resume there and prime only its own delta.
9320 ///
9321 /// EXACTNESS. After this returns, the session is byte-for-byte the state it was in AT that
9322 /// boundary: full-attn KV truncated to it (append-only, position-addressed), GDN conv/ssm
9323 /// restored from the device copy taken there, draft scratch length reset, `committed`
9324 /// truncated, `last_h` = the boundary's predecessor anchor. That is precisely the state a
9325 /// fresh prime of `committed[..pos]` would have produced, so the following suffix prime and
9326 /// every burst after it are identical to a cold run of the same token stream — the
9327 /// committed-tokens-authoritative contract.
9328 ///
9329 /// `next_pred` and `pending_tok` are CLEARED: both describe generation past the boundary,
9330 /// which the rewind discards. The caller therefore must supply a non-empty suffix (a
9331 /// rewound session cannot serve an empty-suffix continuation burst — there is nothing to
9332 /// continue). The persistent draft graph survives: it bakes only session-stable pointers
9333 /// (the scratch KV, the resident embedding), none of which the rewind moves.
9334 ///
9335 /// The checkpoint is CONSUMED (`turn_ckpt` taken): its snapshot buffers are freed here, and
9336 /// this turn's own prime installs a fresh one at the new prompt end. Returns the position
9337 /// rewound to, or `None` when the session holds no checkpoint (caller: full re-prime).
9338 pub fn spec_rewind_to_checkpoint(
9339 &self,
9340 e: &Engine,
9341 sess: &mut SpecSession,
9342 ) -> Result<Option<usize>, Box<dyn std::error::Error>> {
9343 if sess.turn_ckpt.as_ref().is_some_and(|ckpt| {
9344 !sess.cache.can_rollback(&ckpt.snap, 0) || !sess.scratch.can_rewind_to(ckpt.pos)
9345 }) {
9346 return Err(
9347 "SWA ring rewind checkpoint has been lapped; full re-prime required".into(),
9348 );
9349 }
9350 let Some(ckpt) = sess.turn_ckpt.take() else {
9351 return Ok(None);
9352 };
9353 assert!(
9354 ckpt.pos <= sess.committed.len(),
9355 "checkpoint past committed ({} > {})",
9356 ckpt.pos,
9357 sess.committed.len()
9358 );
9359 // Restore through each layer's owning engine. A single primary-engine rollback is not
9360 // sufficient when the serving cache is stage-owned under cross-device PP.
9361 crate::pp::restore_cache_checkpoint(e, self, None, &mut sess.cache, &ckpt.snap)?;
9362 debug_assert_eq!(
9363 sess.cache.pos, ckpt.pos,
9364 "rollback landed off the checkpoint"
9365 );
9366 sess.scratch.set_len(e, ckpt.pos)?;
9367 sess.committed.truncate(ckpt.pos);
9368 sess.last_h = Some(ckpt.last_h);
9369 sess.next_pred = None;
9370 sess.pending_tok = None;
9371 Ok(Some(ckpt.pos))
9372 }
9373
9374 /// Grow a parked speculative session to `target_cap` and rewind it to its retained turn
9375 /// checkpoint without re-priming the checkpoint prefix.
9376 ///
9377 /// The trunk cache is restored exactly like a plain grown cache: append-only full-attention
9378 /// KV rows come from the parked cache, while recurrent state comes from the checkpoint's
9379 /// owned snapshot. The MTP scratch is also context-linear and its rows below the checkpoint
9380 /// remain authoritative, so they are copied into a fresh larger scratch before its length is
9381 /// truncated. Pointer-baking draft graphs are dropped and recaptured on the next burst.
9382 ///
9383 /// All fallible work completes before `sess` is mutated. A failed allocation or copy leaves
9384 /// the parked session intact, allowing the caller one reclaim-and-retry attempt.
9385 pub fn spec_grow_and_rewind_to_checkpoint(
9386 &self,
9387 e: &Engine,
9388 sess: &mut SpecSession,
9389 target_cap: usize,
9390 ) -> Result<Option<usize>, Box<dyn std::error::Error>> {
9391 if target_cap <= sess.cache.max_ctx {
9392 return self.spec_rewind_to_checkpoint(e, sess);
9393 }
9394 let Some(ckpt) = sess.turn_ckpt.as_ref() else {
9395 return Ok(None);
9396 };
9397 if ckpt.pos == 0 || ckpt.pos > sess.committed.len() {
9398 return Err(format!(
9399 "checkpoint pos {} outside committed length {}",
9400 ckpt.pos,
9401 sess.committed.len(),
9402 )
9403 .into());
9404 }
9405 if ckpt.pos > target_cap {
9406 return Err(format!(
9407 "checkpoint pos {} exceeds grown capacity {target_cap}",
9408 ckpt.pos,
9409 )
9410 .into());
9411 }
9412
9413 let mut grown_cache = crate::pp::new_cache_planned(e, &self.cfg, &self.plan, target_cap)?;
9414 let mut grown_scratch = self.new_mtp_scratch(e, target_cap)?;
9415 crate::pp::restore_cache_checkpoint(
9416 e,
9417 self,
9418 Some(&sess.cache),
9419 &mut grown_cache,
9420 &ckpt.snap,
9421 )?;
9422
9423 if sess.scratch.plane_count() != grown_scratch.plane_count() {
9424 return Err("checkpoint draft plane count mismatch".into());
9425 }
9426 for index in 0..sess.scratch.plane_count() {
9427 let (src, _) = sess.scratch.plane(index);
9428 let (dst, _) = grown_scratch.plane_mut(index);
9429 if ckpt.pos > src.len
9430 || src.kv_dim_k != dst.kv_dim_k
9431 || src.kv_dim_v != dst.kv_dim_v
9432 || src.k_tok_bytes != dst.k_tok_bytes
9433 || src.v_tok_bytes != dst.v_tok_bytes
9434 {
9435 return Err(format!(
9436 "checkpoint draft plane {index} layout mismatch (pos {}, source len {})",
9437 ckpt.pos, src.len,
9438 )
9439 .into());
9440 }
9441 let kb = ckpt.pos * src.k_tok_bytes;
9442 let vb = ckpt.pos * src.v_tok_bytes;
9443 if kb > 0 {
9444 e.copy_u8_into(&mut dst.k, 0, &src.k, kb)?;
9445 }
9446 if vb > 0 {
9447 e.copy_u8_into(&mut dst.v, 0, &src.v, vb)?;
9448 }
9449 }
9450 grown_scratch.set_len(e, ckpt.pos)?;
9451 // The old scratch is dropped immediately after publication below. Bound its D2D reads
9452 // first; growth happens once per rewritten turn, outside the decode hot loop.
9453 e.stream().synchronize()?;
9454
9455 let ckpt = sess
9456 .turn_ckpt
9457 .take()
9458 .expect("checkpoint remained present through transactional grow");
9459 let pos = ckpt.pos;
9460 sess.cache = grown_cache;
9461 sess.scratch = grown_scratch;
9462 sess.committed.truncate(pos);
9463 sess.last_h = Some(ckpt.last_h);
9464 sess.next_pred = None;
9465 sess.pending_tok = None;
9466 sess.draft_ctx = None;
9467 debug_assert_eq!(sess.cache.pos, pos, "grown rewind landed off checkpoint");
9468 debug_assert!(
9469 (0..sess.scratch.plane_count()).all(|index| sess.scratch.plane(index).0.len == pos),
9470 "grown draft rewind landed off checkpoint"
9471 );
9472 Ok(Some(pos))
9473 }
9474
9475 /// Commit a carried pending bonus (see SpecSession::pending_tok): one T=1 trunk pass
9476 /// (its logits' argmax becomes next_pred) + the draft-KV fill at the carried anchor —
9477 /// byte-identical to the pre-carry session tail. Required before a non-empty-suffix
9478 /// prime, a sampled turn, or parking a session for pool reuse. No-op without a pending.
9479 /// `sampling` is the sampler of the request that will CONSUME the resulting `next_pred`
9480 /// (lane/sampled-spec-quality): this is a boundary site like any other, so a sampled
9481 /// consumer must get a DRAWN token, not an argmax. Pass `None` from the park/demote
9482 /// callers — a pending only ever exists on the GREEDY tail, and the consumer of a
9483 /// park-time flush is a future request whose sampler is not knowable here (residual
9484 /// named at the pool-resume probe in worker.rs and in SAMPLED-QUALITY.md).
9485 pub fn spec_flush_pending(
9486 &self,
9487 e: &Engine,
9488 sess: &mut SpecSession,
9489 sampling: Option<SpecSampling>,
9490 ) -> Result<(), Box<dyn std::error::Error>> {
9491 let Some(b) = sess.pending_tok.take() else {
9492 return Ok(());
9493 };
9494 if self.mtp.is_none() {
9495 return Err("pending carry requires an MTP head".into());
9496 }
9497 let n_embd = self.cfg.n_embd as usize;
9498 let (embd_qt, embd_rb) = self.embd.qt_and_row_bytes(n_embd);
9499 let embd_gpu = if spec_host_embd() {
9500 None
9501 } else {
9502 Some(
9503 self.embd_gpu
9504 .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload")),
9505 )
9506 };
9507 let embd_dev = embd_gpu.map(|g| (g, embd_qt, embd_rb));
9508 let pos_b = sess.cache.pos;
9509 sess.scratch.set_len(e, pos_b)?;
9510 let (lg_b, hb) = self.spec_target_step_h(e, b, &mut sess.cache)?;
9511 sess.next_pred = Some(match sampling {
9512 Some(sp) if sp.temp > 0.0 && spec_sampled_boundary_on() => {
9513 // window includes `b` itself: it is committed by this pass, and the pre-lane
9514 // code never counted a boundary token in the penalty history at all.
9515 let hist = pen_window_seed(&sess.committed, &[b], sp.penalty_last_n);
9516 sample_boundary_token(e, &lg_b, &sp, &hist, &mut sess.sctr, "flush-pending")?
9517 }
9518 _ => argmax(&lg_b) as u32,
9519 });
9520 let anchor = sess
9521 .last_h
9522 .as_ref()
9523 .expect("pending carry requires last_h (the predecessor-row anchor)");
9524 self.mtp_kv_fill_all(e, &[b], anchor, pos_b, &mut sess.scratch, embd_dev)?;
9525 sess.last_h = Some(hb);
9526 sess.committed.push(b);
9527 Ok(())
9528 }
9529
9530 /// Solo target feed used only at speculative round boundaries. Step35 serving made its
9531 /// staged batched B=1 graph authoritative, so a speculative session must enter and leave
9532 /// rounds through that same graph. Other model families keep their eager T=1 contract.
9533 fn spec_target_step_h(
9534 &self,
9535 e: &Engine,
9536 token: u32,
9537 cache: &mut Cache,
9538 ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9539 if !self.sliding_gated_moe_batch_program() && !self.batched_serving_numeric_class() {
9540 return self.decode_step_h(e, token, cache);
9541 }
9542 let pos0 = cache.pos;
9543 let (logits, hidden) = self.decode_step_t_core(e, &[token], pos0, cache, None, None)?;
9544 Ok((e.dtoh(&logits)?, hidden))
9545 }
9546
9547 /// The archs whose LIVE B=1 serving runs the generic BATCHED numeric class (decode_step_batch
9548 /// walk + batched head), so their spec verify must run the SAME class. MoE learned this
9549 /// 2026-08-14 AM (4b777ccc5); the dense hybrid reproduced the identical near-tie flip class
9550 /// the same day on Qwen3.8-27B — eager-class verify logits drift from batched-class serving
9551 /// logits ("1 ULP at layer 2 → 2.3e-1 logit maxdiff at the head"), and the GDN recurrence
9552 /// carries the drift until a near-tie flips deep in generation. One predicate so the five
9553 /// dispatch sites cannot drift apart again.
9554 /// Draft-graph head admissibility (lane/draftcost-moe, 2026-08-20): the capture body
9555 /// (`mtp_head_forward_cap`) supports Dense heads AND resident-MoE heads
9556 /// (`Ffn::Moe(m) if m.dev_exps.is_some()`); non-resident MoE still refuses inside the
9557 /// capture and the caller falls back to the eager chain by design. Trunk FFN class is
9558 /// irrelevant — the graph body is the HEAD forward only. One predicate for all three
9559 /// eligibility sites so they cannot drift (the serving numeric-class lesson).
9560 fn mtp_graph_capturable(&self) -> bool {
9561 self.mtp
9562 .as_ref()
9563 .map(|m| match &m.ffn {
9564 crate::hybrid::Ffn::Dense { .. } => true,
9565 crate::hybrid::Ffn::Moe(mo) => mo.dev_exps.is_some(),
9566 })
9567 .unwrap_or(false)
9568 }
9569
9570 fn batched_serving_numeric_class(&self) -> bool {
9571 self.plan
9572 .trunk_operations()
9573 .contains(&memra_gguf::model_plan::OperationKind::GatedDeltaNet)
9574 }
9575
9576 /// The family the MTP verify-graph default was measured on: GatedDeltaNet state layers
9577 /// (a `recur` mixer) together with a routed-MoE FFN — Ornith-1.5-35B-A3B and its kin. The
9578 /// server-side twin of this test is `model_forces_spec_replay` (GatedDeltaNet + MoeMlp);
9579 /// keeping the engine's own version structural rather than name-based means a new
9580 /// checkpoint of the same shape inherits the default, and a different shape does not.
9581 fn vgraph_family_default(&self) -> bool {
9582 let has_linear = self
9583 .layers
9584 .iter()
9585 .any(|l| matches!(l.mixer, Mixer::Linear(_)));
9586 let has_moe = self
9587 .layers
9588 .iter()
9589 .any(|l| matches!(l.ffn, crate::hybrid::Ffn::Moe(_)));
9590 has_linear && has_moe
9591 }
9592
9593 fn sliding_gated_moe_batch_program(&self) -> bool {
9594 self.uses_sliding_gated_moe_program()
9595 }
9596
9597 fn gemma_batch_program(&self) -> bool {
9598 self.uses_gemma_program()
9599 }
9600
9601 /// Reduced-matrix admission for increment 1. This deliberately does not change the PP-2
9602 /// serving policy: the worker calls it only after `MEMRA_SPEC_PIPE=1` and an explicit spec
9603 /// session already exist.
9604 pub fn spec_pipe_available(&self, e: &Engine) -> bool {
9605 if std::env::var("MEMRA_SPEC_PIPE").as_deref() != Ok("1")
9606 || !spec_devacc()
9607 || spec_replay_env_enabled()
9608 || spec_stream()
9609 || std::env::var("MEMRA_SPEC_ADAPT").as_deref() == Ok("1")
9610 || std::env::var("MEMRA_SPEC_PMIN0").as_deref() == Ok("1")
9611 || std::env::var("MEMRA_SPEC_PP_ANATOMY").as_deref() == Ok("1")
9612 || std::env::var("MEMRA_SPEC_PMIN")
9613 .ok()
9614 .and_then(|v| v.parse::<f32>().ok())
9615 .unwrap_or(0.0)
9616 > 0.0
9617 || self.is_gemma4_e4b()
9618 || self.gemma_batch_program()
9619 || self.mtp.is_none()
9620 || !self.mtp_extra.is_empty()
9621 {
9622 return false;
9623 }
9624 let Some(cuts) = crate::pp::pp_cuts(self.layers.len()) else {
9625 return false;
9626 };
9627 if cuts.len() != 3 || crate::pp::pp2_streams_off() || !crate::pp::spec_pp_on() {
9628 return false;
9629 }
9630 crate::pp::PpNRt::get(e)
9631 .map(|rt| rt.n_stages() == 2 && rt.cross_device())
9632 .unwrap_or(false)
9633 }
9634
9635 /// Two warm greedy continuation bursts over one PP-2 interval coordinator. The two existing
9636 /// `generate_spec_inner2` call stacks own all per-session round locals; only phase issue order
9637 /// changes. No callback is accepted in increment 1 — the worker publishes each completed burst.
9638 #[allow(clippy::too_many_arguments)]
9639 pub fn generate_spec_session_pair(
9640 &self,
9641 e: &Engine,
9642 sess_a: &mut SpecSession,
9643 max_new_a: usize,
9644 k_a: usize,
9645 sess_b: &mut SpecSession,
9646 max_new_b: usize,
9647 k_b: usize,
9648 ) -> Result<((Vec<u32>, usize, usize), (Vec<u32>, usize, usize)), Box<dyn std::error::Error>>
9649 {
9650 if !self.spec_pipe_available(e) {
9651 return Err("two-session speculative pipeline is outside its reduced matrix".into());
9652 }
9653 if max_new_a == 0 || max_new_b == 0 || k_a == 0 || k_b == 0 {
9654 return Err(
9655 "two-session speculative pipeline requires non-empty positive-K bursts".into(),
9656 );
9657 }
9658 for sess in [&*sess_a, &*sess_b] {
9659 if sess.committed.is_empty()
9660 || sess.last_h.is_none()
9661 || (sess.next_pred.is_none() && sess.pending_tok.is_none())
9662 {
9663 return Err("two-session speculative pipeline requires warm continuations".into());
9664 }
9665 }
9666
9667 let graph_ok = std::env::var("MEMRA_SPEC_NOGRAPH").is_err()
9668 && !spec_host_embd()
9669 && self.mtp_graph_capturable()
9670 && self.mtp_extra.is_empty()
9671 && !crate::model::full_prec_enabled();
9672 let graph_a = graph_ok && k_a + 2 < 96;
9673 let graph_b = graph_ok && k_b + 2 < 96;
9674 let was_tracking = e.ctx().is_event_tracking();
9675 if (graph_a || graph_b) && was_tracking {
9676 unsafe {
9677 e.ctx().disable_event_tracking();
9678 }
9679 }
9680
9681 static LOGGED: std::sync::Once = std::sync::Once::new();
9682 LOGGED.call_once(|| {
9683 eprintln!("[spec-pipe] two-session PP-2 continuation pipeline engaged");
9684 });
9685 let sync = std::sync::Arc::new(SpecPipeSync::new());
9686 let lane_a = SpecPipeLane {
9687 sync: sync.clone(),
9688 lane: 0,
9689 };
9690 let lane_b = SpecPipeLane { sync, lane: 1 };
9691 let mut sess_b_ptr = SpecPipeSessionPtr(sess_b as *mut SpecSession);
9692 let (result_a, result_b) = std::thread::scope(|scope| {
9693 let b = scope.spawn(move || {
9694 let mut finish = SpecPipeFinish::new(&lane_b);
9695 let sess_b = unsafe { sess_b_ptr.get_mut() };
9696 let result = e
9697 .ctx()
9698 .bind_to_thread()
9699 .map_err(|err| err.to_string())
9700 .and_then(|_| {
9701 self.generate_spec_inner2(
9702 e,
9703 &[],
9704 max_new_b,
9705 k_b,
9706 graph_b,
9707 Some(sess_b),
9708 None,
9709 None,
9710 None,
9711 None,
9712 Some(&lane_b),
9713 )
9714 .map_err(|err| err.to_string())
9715 });
9716 finish.close(result.is_err());
9717 result
9718 });
9719 let mut finish = SpecPipeFinish::new(&lane_a);
9720 let result_a = self.generate_spec_inner2(
9721 e,
9722 &[],
9723 max_new_a,
9724 k_a,
9725 graph_a,
9726 Some(sess_a),
9727 None,
9728 None,
9729 None,
9730 None,
9731 Some(&lane_a),
9732 );
9733 finish.close(result_a.is_err());
9734 let result_b = b
9735 .join()
9736 .map_err(|_| "paired speculative session B panicked".to_string())
9737 .and_then(|r| r);
9738 (result_a, result_b)
9739 });
9740
9741 if (graph_a || graph_b) && was_tracking {
9742 unsafe {
9743 e.ctx().enable_event_tracking();
9744 }
9745 }
9746 let result_a = result_a?;
9747 let result_b = result_b.map_err(|err| -> Box<dyn std::error::Error> { err.into() })?;
9748 Ok((result_a, result_b))
9749 }
9750
9751 /// One spec-decode turn on a live session. `suffix` = the NEW tokens only (turn N+1's user
9752 /// message rendered through the chat template continuation). Returns (new tokens emitted,
9753 /// drafted, accepted); session.committed grows by suffix + emitted.
9754 pub fn generate_spec_session(
9755 &self,
9756 e: &Engine,
9757 sess: &mut SpecSession,
9758 suffix: &[u32],
9759 max_new: usize,
9760 k: usize,
9761 ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
9762 self.generate_spec_session_sampled(e, sess, suffix, max_new, k, None, None)
9763 }
9764
9765 /// Serve-path sampled spec: routes the burst through the rejection-sampling verify with
9766 /// per-SESSION Philox continuity (sess.sctr/uctr). None = env-driven (CLI) or greedy.
9767 /// Filters (top-k/p/min-p) apply SYMMETRICALLY to draft q and verify p — distribution-exact
9768 /// for the filtered target (feat/filtered-spec).
9769 ///
9770 /// `on_commit` (sse-cadence, 2026-08-05): called with each newly-emitted slice of the
9771 /// output — once right after the prime's first token, then once per round commit — so a
9772 /// streaming caller can flush text at round cadence instead of once per burst. The slices
9773 /// are disjoint, in order, and concatenate to exactly the returned token vec. Emission-
9774 /// timing only: token bytes, session state, and exactness are untouched.
9775 ///
9776 /// The returned bool is a CONTINUE-VERDICT (admission yield, 2026-08-06): `false` ends
9777 /// the burst at the current round boundary, exactly as if `max_new` had been reached —
9778 /// the caller's scheduler regains control without waiting the burst out. Burst size is
9779 /// content-neutral (spec-levers battery), so an early exit moves WHEN the burst returns,
9780 /// never what tokens say. The slice may be EMPTY (a poll-only boundary — round-stream
9781 /// drains and the defensive tail flush can land with nothing new committed).
9782 #[allow(clippy::too_many_arguments)]
9783 pub fn generate_spec_session_sampled(
9784 &self,
9785 e: &Engine,
9786 sess: &mut SpecSession,
9787 suffix: &[u32],
9788 max_new: usize,
9789 k: usize,
9790 sampling: Option<SpecSampling>,
9791 on_commit: Option<&mut dyn FnMut(&[u32]) -> bool>,
9792 ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
9793 self.generate_spec_session_sampled_prime_split(
9794 e, sess, suffix, max_new, k, sampling, None, on_commit,
9795 )
9796 }
9797
9798 /// Serve-only cold-prime segmentation twin. `prime_split` is the same stable boundary the
9799 /// plain worker would honor before entering its sub-floor tokenwise tail; warm continuations
9800 /// pass `None` and stay on the existing zero-prime path.
9801 #[allow(clippy::too_many_arguments)]
9802 pub fn generate_spec_session_sampled_prime_split(
9803 &self,
9804 e: &Engine,
9805 sess: &mut SpecSession,
9806 suffix: &[u32],
9807 max_new: usize,
9808 k: usize,
9809 sampling: Option<SpecSampling>,
9810 prime_split: Option<usize>,
9811 on_commit: Option<&mut dyn FnMut(&[u32]) -> bool>,
9812 ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
9813 self.generate_spec_session_constrained_prime_split(
9814 e,
9815 sess,
9816 suffix,
9817 max_new,
9818 k,
9819 sampling,
9820 None,
9821 prime_split,
9822 on_commit,
9823 )
9824 }
9825
9826 /// `generate_spec_session_sampled` + GRAMMAR (constrained decoding, 2026-08-03): the
9827 /// hook truncates acceptance at the first grammar-illegal token AFTER the exactness
9828 /// verify (grammar is an extra rejection rule, ordering like the batched-verify twins)
9829 /// and replaces an illegal bonus with the MASKED argmax of the target's own verify
9830 /// column — token-identical to constrained plain greedy decode. GREEDY only (the
9831 /// worker routes sampled constrained to plain decode). Acceptance under tight grammars
9832 /// may drop (drafter is unconstrained); that is measured, not hidden.
9833 #[allow(clippy::too_many_arguments)]
9834 pub fn generate_spec_session_constrained(
9835 &self,
9836 e: &Engine,
9837 sess: &mut SpecSession,
9838 suffix: &[u32],
9839 max_new: usize,
9840 k: usize,
9841 sampling: Option<SpecSampling>,
9842 constraint: Option<&mut dyn SpecConstraint>,
9843 on_commit: Option<&mut dyn FnMut(&[u32]) -> bool>,
9844 ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
9845 self.generate_spec_session_constrained_prime_split(
9846 e, sess, suffix, max_new, k, sampling, constraint, None, on_commit,
9847 )
9848 }
9849
9850 #[allow(clippy::too_many_arguments)]
9851 pub fn generate_spec_session_constrained_prime_split(
9852 &self,
9853 e: &Engine,
9854 sess: &mut SpecSession,
9855 suffix: &[u32],
9856 max_new: usize,
9857 k: usize,
9858 sampling: Option<SpecSampling>,
9859 constraint: Option<&mut dyn SpecConstraint>,
9860 prime_split: Option<usize>,
9861 on_commit: Option<&mut dyn FnMut(&[u32]) -> bool>,
9862 ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
9863 if constraint.is_some() && sampling.is_some_and(|s| s.temp > 0.0) {
9864 return Err(
9865 "constrained spec decode is greedy-only (worker routes sampled \
9866 constrained to plain decode)"
9867 .into(),
9868 );
9869 }
9870 // PENDING-CARRY entry flush: a carried bonus precedes any new suffix in the sequence,
9871 // so it must commit BEFORE the suffix primes; the sampled path doesn't carry (its
9872 // round-0 accept needs the commit pass's logits). Empty-suffix greedy bursts — the
9873 // serve continuation case — consume the carry in-loop with zero solo passes.
9874 if sess.pending_tok.is_some()
9875 && (!suffix.is_empty() || sampling.map_or(false, |s| s.temp > 0.0))
9876 {
9877 self.spec_flush_pending(e, sess, sampling)?;
9878 }
9879
9880 // FULL_PREC forces the EAGER draft: the graph capture would enclose cuBLASLt f32 GEMV
9881 // (the FloatBf16 else-branches) and a bf16_to_f32 dequant alloc — neither is stream-capture
9882 // safe. Eager rides matmul/matmul_decode_exact, which dequant FloatBf16 on use. (§item 2.)
9883 let graph_draft = std::env::var("MEMRA_SPEC_NOGRAPH").is_err()
9884 && !spec_host_embd()
9885 && self.mtp_graph_capturable()
9886 && self.mtp_extra.is_empty()
9887 && k + 2 < 96
9888 && !crate::model::full_prec_enabled();
9889 let was_tracking = e.ctx().is_event_tracking();
9890 if graph_draft && was_tracking {
9891 unsafe {
9892 e.ctx().disable_event_tracking();
9893 }
9894 }
9895 let r = self.generate_spec_inner2(
9896 e,
9897 suffix,
9898 max_new,
9899 k,
9900 graph_draft,
9901 Some(sess),
9902 sampling,
9903 constraint,
9904 on_commit,
9905 prime_split,
9906 None,
9907 );
9908 if graph_draft && was_tracking {
9909 unsafe {
9910 e.ctx().enable_event_tracking();
9911 }
9912 }
9913 let (out, d, a) = r?;
9914 Ok((out, d, a))
9915 }
9916
9917 pub fn generate_spec(
9918 &self,
9919 e: &Engine,
9920 prompt: &[u32],
9921 max_new: usize,
9922 k: usize,
9923 ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
9924 if crate::pp::pp_cuts(self.layers.len()).is_some()
9925 && !self.rewrite_allowed(memra_gguf::execution_manifest::RewriteSurface::Pipeline)
9926 {
9927 return Err("pipeline rewrite is not qualified for speculative decode".into());
9928 }
9929 if !self.rewrite_allowed(memra_gguf::execution_manifest::RewriteSurface::MtpSpec) {
9930 return Err("speculative rewrite is not qualified for this ModelPlan".into());
9931 }
9932 // FULL_PREC forces eager (see generate_spec_session note): CUDA graph capture cannot
9933 // enclose cuBLASLt f32 GEMV or the bf16_to_f32 dequant alloc the FloatBf16 path needs.
9934 let graph_draft = std::env::var("MEMRA_SPEC_NOGRAPH").is_err()
9935 && !spec_host_embd()
9936 && self.mtp_graph_capturable()
9937 && self.mtp_extra.is_empty()
9938 && k + 2 < 96
9939 && !crate::model::full_prec_enabled();
9940 if !graph_draft {
9941 return self.generate_spec_inner2(
9942 e, prompt, max_new, k, false, None, None, None, None, None, None,
9943 );
9944 }
9945 let was_tracking = e.ctx().is_event_tracking();
9946 if was_tracking {
9947 unsafe {
9948 e.ctx().disable_event_tracking();
9949 }
9950 }
9951 let r = self.generate_spec_inner2(
9952 e, prompt, max_new, k, true, None, None, None, None, None, None,
9953 );
9954 if was_tracking {
9955 unsafe {
9956 e.ctx().enable_event_tracking();
9957 }
9958 }
9959 r
9960 }
9961
9962 fn generate_spec_inner2(
9963 &self,
9964 e: &Engine,
9965 prompt: &[u32],
9966 max_new: usize,
9967 k: usize,
9968 graph_draft: bool,
9969 mut sess: Option<&mut SpecSession>,
9970 sampling: Option<SpecSampling>,
9971 mut constraint: Option<&mut dyn SpecConstraint>,
9972 mut on_commit: Option<&mut dyn FnMut(&[u32]) -> bool>,
9973 prime_split: Option<usize>,
9974 pipe: Option<&SpecPipeLane>,
9975 ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
9976 assert!(k >= 1, "k must be >= 1");
9977 if let Some(p) = pipe {
9978 p.setup_begin()?;
9979 }
9980 // sse-cadence flush cursor: everything in out[..flushed] has been handed to on_commit.
9981 let mut flushed = 0usize;
9982 // admission yield (2026-08-06): on_commit's continue-verdict; false = end the burst
9983 // at the next round boundary (same exit as max_new reached — the session tail runs).
9984 // Initialized by the unconditional post-prime flush below.
9985 let mut keep_going;
9986 let mtp = self
9987 .mtp
9988 .as_ref()
9989 .expect("generate_spec requires an MTP head (nextn_predict_layers>0)");
9990 let n_vocab = self.output.out_features();
9991 // FR-Spec: the draft head may be TRIMMED (fewer rows than n_vocab); the draft argmax runs
9992 // over the draft vocab and the winning index maps through d2t to a TARGET token id.
9993 // Everything downstream (verify/accept/commit) sees target ids only — exactness unchanged.
9994 let d_vocab = mtp
9995 .shared_head_head
9996 .as_ref()
9997 .unwrap_or(&self.output)
9998 .out_features();
9999 if !self.mtp_extra.is_empty() {
10000 if self.plan.draft_source != memra_gguf::model_plan::DraftSourcePlan::Embedded
10001 || self.plan.mtp_blocks.len() != self.mtp_head_count()
10002 || mtp.d2t.is_some()
10003 {
10004 return Err(
10005 "multi-head MTP requires one embedded canonical block per loaded head".into(),
10006 );
10007 }
10008 for (offset, head) in self.mtp_extra.iter().enumerate() {
10009 if head.d2t.is_some()
10010 || head
10011 .shared_head_head
10012 .as_ref()
10013 .unwrap_or(&self.output)
10014 .out_features()
10015 != d_vocab
10016 {
10017 return Err(format!(
10018 "embedded MTP head {} has incompatible draft vocabulary",
10019 offset + 1
10020 )
10021 .into());
10022 }
10023 }
10024 eprintln!(
10025 "[mtp-chain] heads={} policy=step-modulo prefix-replay kv=per-head",
10026 self.mtp_head_count()
10027 );
10028 }
10029 let n_embd = self.cfg.n_embd as usize;
10030 // SESSION MODE: reuse the live cache/scratch, prime only the suffix. `base` = tokens
10031 // already committed (their state is in the caches); 0 = fresh single-shot call.
10032 let session_mode = sess.is_some();
10033 let max_ctx = match sess.as_ref() {
10034 Some(s) => s.cache.max_ctx,
10035 None => prompt.len() + max_new + k + 8,
10036 };
10037 let mut own_cache;
10038 let mut own_scratch;
10039 // PREFIX-CACHE capture request threaded out of the session (lane/spec-prefix-cache):
10040 // (requested split, destination list). Single-shot per burst; fresh calls have none.
10041 let mut sess_capture: Option<(Option<usize>, &mut Vec<SpecBoundaryCapture>)> = None;
10042 // STABLE-BOUNDARY turn-checkpoint request (lane/frspec-multiturn-cache): ABSOLUTE
10043 // committed-length position; consumed one-shot like `capture_at`. None = legacy
10044 // prompt-end capture below.
10045 let mut ckpt_req: Option<usize> = None;
10046 let (
10047 cache,
10048 scratch,
10049 mut sess_tail,
10050 mut sess_draft_slot,
10051 mut sess_pending_slot,
10052 sess_ckpt_slot,
10053 sess_telem,
10054 ): (
10055 &mut Cache,
10056 &mut MtpScratch,
10057 Option<(
10058 &mut Vec<u32>,
10059 &mut Option<CudaSlice<f32>>,
10060 &mut Option<u32>,
10061 &mut u32,
10062 &mut u32,
10063 )>,
10064 Option<&mut Option<DraftGraphCtx>>,
10065 Option<&mut Option<u32>>,
10066 Option<&mut Option<SpecCheckpoint>>,
10067 Option<&SpecTelemetryCounters>,
10068 ) = match sess.take() {
10069 Some(sr) => {
10070 let SpecSession {
10071 cache,
10072 scratch,
10073 committed,
10074 last_h,
10075 next_pred,
10076 sctr: s_sctr,
10077 uctr: s_uctr,
10078 draft_ctx,
10079 pending_tok,
10080 turn_ckpt,
10081 telem,
10082 capture_at,
10083 boundary_captures,
10084 ckpt_at,
10085 } = sr;
10086 sess_capture = Some((capture_at.take(), boundary_captures));
10087 ckpt_req = ckpt_at.take();
10088 (
10089 cache,
10090 scratch,
10091 Some((committed, last_h, next_pred, s_sctr, s_uctr)),
10092 Some(draft_ctx),
10093 Some(pending_tok),
10094 Some(turn_ckpt),
10095 Some(telem),
10096 )
10097 }
10098 None => {
10099 // STAGE-OWNED KV (lane/pp2-spec 2026-08-06) — see `new_session`. Door shut =
10100 // `Cache::new` verbatim.
10101 own_cache = crate::pp::new_cache_planned(e, &self.cfg, &self.plan, max_ctx)?;
10102 // Persistent scratch = max_ctx rows (~2KB/token quantized).
10103 own_scratch = self.new_mtp_scratch(e, max_ctx)?;
10104 (
10105 &mut own_cache,
10106 &mut own_scratch,
10107 None,
10108 None,
10109 None,
10110 None,
10111 None,
10112 )
10113 }
10114 };
10115 if scratch.plane_count() != self.mtp_head_count() {
10116 return Err(format!(
10117 "MTP scratch/head count mismatch ({}/{})",
10118 scratch.plane_count(),
10119 self.mtp_head_count()
10120 )
10121 .into());
10122 }
10123 let base = cache.pos;
10124 // PENDING-CARRY consume (2026-08-01): a carried bonus reaches here only on the
10125 // empty-suffix GREEDY continuation path (generate_spec_session_sampled flushed every
10126 // other case). It enters the round loop as round-0's pending — verify col 0 — exactly
10127 // like a mid-burst full-accept boundary: no init feed, no tail commit pass.
10128 let carried_pending: Option<u32> = sess_pending_slot.as_mut().and_then(|s| s.take());
10129 // PERSISTENT DRAFT KV (the only mode since 2026-07-08 — the legacy round-local scratch,
10130 // MEMRA_SPEC_KVLOCAL, measured -35 acceptance pts on the 27B p3 sweep and was removed;
10131 // acceptance-only — exactness is verify's job either way).
10132 // HIDDEN-PAIRING CONVENTION (DEFAULT = predecessor-row, 2026-07-04 — the 27B acceptance
10133 // unlock, +16pts): the MTP head is TRAINED on rows pairing token x_p with the trunk
10134 // hidden of its PREDECESSOR h_{p-1} (the reference engine's mtp_update shifts the target
10135 // hiddens right by one; its draft step 0 feeds (id_last, TRUE hidden of the row id_last
10136 // was sampled from)). memra's historical convention paired SAME-ROW (x_p, h_p) in the fill
10137 // and seeded chain step 0 through an extra MTP pass on a duplicated token (the
10138 // pseudo-seed) — measured 27B p2 K=3 acceptance 0.569 vs 0.731, p3 0.445 vs 0.63+, and
10139 // the chain steps j>=1 were already predecessor-shaped, so ONLY the fill + step-0 seed
10140 // move. The fill shifts by one and the chain seeds from the predecessor's true hidden
10141 // DIRECTLY (vh_seed / vx[j-1]) — the pseudo pass disappears (one MTP-block pass saved
10142 // per round on top of the acceptance win). Draft-quality-only: exactness stays the
10143 // verify's job either way. (The legacy same-row pairing seam, MEMRA_SPEC_HSAME, and its
10144 // pseudo-seed passes were removed 2026-07-08 — predecessor pairing won by +16 acc pts;
10145 // the legacy round-local scratch, MEMRA_SPEC_KVLOCAL, went with it.)
10146 // REPLAY-FREE PARTIAL ACCEPT (default, 2026-07-03): partial rounds keep the verify's own
10147 // bit-identical committed-prefix state (KV truncate + recur rebuild from the VerifyCkpt)
10148 // and leave the bonus PENDING — no duplicate trunk pass (profiled ~0.54 extra full weight
10149 // reads/round at long ctx). MEMRA_SPEC_REPLAY=1 restores the legacy rollback+replay (A/B
10150 // + fallback seam).
10151 // Qwen35-MoE replay pin LIFTED (lane/draftcost-moe, 2026-08-20). The pin's stated
10152 // bar — the retained verify-state commit proven equivalent to sequential serving —
10153 // was waiting on this arch running the serving batched verify class, which the
10154 // t-parallel admission (this lane, increment 1) provided: the VerifyCkpt the
10155 // replay-free commit consumes is now produced by the SAME serving-class verify that
10156 // qualified dense qwen35 on 2026-08-15 (where the per-round duplicate replay
10157 // measured 69 -> 30 tok/s). Qualification receipts (run-spec K=1..8 both arms,
10158 // 8-prompt replay-vs-replay-free canary, long-prompt cell):
10159 // research/draftcost-moe-20260820/RECEIPTS.md. MEMRA_SPEC_REPLAY=1 stays the
10160 // rollback + A/B seam.
10161 let spec_replay = spec_replay_env_enabled();
10162 if constraint.is_some() && spec_replay {
10163 return Err(
10164 "constrained spec decode does not support MEMRA_SPEC_REPLAY=1 \
10165 (legacy replay commits an unmasked bonus)"
10166 .into(),
10167 );
10168 }
10169 // TRUE-HIDDEN REFRESH (default in persistent-draft-KV mode): every round overwrites the
10170 // committed positions' scratch entries from the verify's exact hiddens (mtp_kv_fill batch)
10171 // instead of keeping chain-approximate entries. MEMRA_SPEC_NOREFRESH=1 = legacy (A/B seam).
10172 let refresh = std::env::var("MEMRA_SPEC_NOREFRESH").is_err();
10173 if !refresh && !self.mtp_extra.is_empty() {
10174 return Err("multi-head MTP requires exact accepted-prefix refresh".into());
10175 }
10176
10177 // prime: BATCHED cache prime (prime_cache — the measured #1 e2e gap: tokenwise primed at
10178 // ~102/38 tok/s vs the engine's ~2000-5900 tok/s batched prefill). prime_cache returns the
10179 // full pre-output_norm hidden stack [T, n_embd], which IS prompt_h (the persistent-draft-KV
10180 // mtp_kv_fill input) — no per-token collection needed. Prompts below PRIME_MIN_T, and
10181 // MEMRA_PRIME_TOKENWISE=1, and frozen Hy3 CPU/GPU expert splits take the tokenwise
10182 // decode_step_h loop. The latter avoids transient GPU staging of the spilled expert bank.
10183 // EMPTY-SUFFIX CONTINUATION (serve bursts): a session turn with NO new tokens resumes
10184 // generation exactly where the last turn stopped — no prime at all. The stashed
10185 // `next_pred` plays prime_logits' role: it is the token produced from the logits after
10186 // committed.last() by the same rule this entry applies to a cold prime's last row —
10187 // an argmax when greedy, a `sample_boundary_token` draw when sampled (the burst tail,
10188 // or `spec_session_from_restored` for a converted prefix-cache hit, did the drawing
10189 // where the sampler and the session's Philox counters were live). `last_h` seeds the
10190 // predecessor pairing below. Fresh calls and non-empty suffixes take the normal path.
10191 let continuation = prompt.is_empty();
10192 if continuation {
10193 assert!(session_mode, "empty prompt requires a session");
10194 assert!(
10195 sess_tail
10196 .as_ref()
10197 .map_or(false, |(c, lh, np, _, _)| !c.is_empty()
10198 && lh.is_some()
10199 && (np.is_some() || carried_pending.is_some())),
10200 "empty-suffix continuation needs a primed session (committed + last_h + next_pred|pending)"
10201 );
10202 }
10203 let mut prime_logits;
10204 let mut prompt_h: Option<CudaSlice<f32>> = None;
10205 let t_prime = std::time::Instant::now();
10206 let batched_prime = !continuation
10207 && prompt.len() >= crate::hybrid_forward::PRIME_MIN_T
10208 && std::env::var("MEMRA_PRIME_TOKENWISE").is_err()
10209 && !e.frozen_cpu_experts_prefer_tokenwise_prime();
10210 let prime_split = prime_split.filter(|&split| split > 0 && split < prompt.len());
10211 if prime_split.is_some() && continuation {
10212 return Err("spec prime split requires a non-empty prime".into());
10213 }
10214 // STABLE-BOUNDARY TURN CHECKPOINT stop (lane/frspec-multiturn-cache, 2026-08-21):
10215 // the worker's `ckpt_at` request, ABSOLUTE -> prompt-relative. On WARM bursts
10216 // (base != 0, an affinity-rewound or pool-resumed session priming its own delta)
10217 // this is the only stop; on COLD bursts it usually coincides with `prime_split`
10218 // (both are the plain tier's stable pre-generation boundary). A boundary the prime
10219 // cannot honor (outside this prime's range) silently drops the capture — the
10220 // turn_ckpt convention: the next turn re-primes in full, never a wrong resume.
10221 let ckpt_rel = if continuation {
10222 None
10223 } else {
10224 ckpt_req
10225 .and_then(|abs| abs.checked_sub(base))
10226 .filter(|&r| r > 0 && r < prompt.len())
10227 };
10228 // Prime stops, ordered: each is a boundary the prime halts at so the in-place GDN
10229 // conv/ssm state can be snapshotted there (the only moment it exists). One stop =
10230 // the legacy single-split program, byte-for-byte.
10231 let mut stops: Vec<usize> = Vec::new();
10232 for b in [prime_split, ckpt_rel].into_iter().flatten() {
10233 if !stops.contains(&b) {
10234 stops.push(b);
10235 }
10236 }
10237 stops.sort_unstable();
10238 // Captured at the ckpt stop, installed into the session slot post-prime (replacing
10239 // the legacy prompt-end capture). Some(None) = capture attempted and failed -> the
10240 // slot is cleared (a stale checkpoint would rewind to the WRONG boundary).
10241 let mut ckpt_early: Option<Option<SpecCheckpoint>> = None;
10242 if continuation {
10243 prime_logits = Vec::new();
10244 } else if !stops.is_empty() {
10245 if let Some(&first) = stops.first() {
10246 if prime_split == Some(first) && first < crate::hybrid_forward::PRIME_MIN_T {
10247 return Err(format!(
10248 "spec prime split {first} is below PRIME_MIN_T {}",
10249 crate::hybrid_forward::PRIME_MIN_T,
10250 )
10251 .into());
10252 }
10253 }
10254 // Mirror the plain worker's boundary stops exactly. Each segment is a
10255 // request-level prime (`queued_after` keeps Step35 arm selection independent of
10256 // the stops — tick-seg law); a segment below PRIME_MIN_T (and the final tail
10257 // under MEMRA_PRIME_TOKENWISE) takes the same eager tokenwise continuation as
10258 // prefill_tick. Retain every hidden row so the draft scratch fill remains one
10259 // coherent prompt.
10260 let mut h_all = e.uninit(prompt.len() * n_embd)?;
10261 prime_logits = Vec::new();
10262 let mut prev = 0usize;
10263 for seg_end in stops.iter().copied().chain(std::iter::once(prompt.len())) {
10264 if seg_end <= prev {
10265 continue;
10266 }
10267 let seg = &prompt[prev..seg_end];
10268 let is_final = seg_end == prompt.len();
10269 let batched_seg = seg.len() >= crate::hybrid_forward::PRIME_MIN_T
10270 && (!is_final
10271 || (std::env::var("MEMRA_PRIME_TOKENWISE").is_err()
10272 && !e.frozen_cpu_experts_prefer_tokenwise_prime()));
10273 if batched_seg {
10274 let (l, _, h_seg) =
10275 self.prime_cache(e, seg, &mut *cache, prompt.len() - seg_end)?;
10276 e.copy_into(&mut h_all, prev * n_embd, &h_seg, seg.len() * n_embd)?;
10277 prime_logits = l;
10278 } else {
10279 for (i, &tok) in seg.iter().enumerate() {
10280 let (l, h) = self.decode_step_h(e, tok, &mut *cache)?;
10281 e.copy_into(&mut h_all, (prev + i) * n_embd, &h, n_embd)?;
10282 prime_logits = l;
10283 }
10284 }
10285 prev = seg_end;
10286 if is_final {
10287 break;
10288 }
10289 debug_assert_eq!(cache.pos, base + seg_end, "prime stop landed off boundary");
10290 // PREFIX-CACHE BOUNDARY CAPTURE (lane/spec-prefix-cache): the GDN conv/ssm
10291 // states are about to be advanced in place by the next segment, so this is
10292 // the ONLY moment the boundary's recurrent state exists. Capture iff the
10293 // worker requested exactly this stop (cold sessions only — `capture_at` is
10294 // never armed warm). A failed snapshot is silent (turn_ckpt convention) —
10295 // publication is an optimization, never a correctness dependency.
10296 if base == 0 {
10297 if let Some((requested, slot)) = sess_capture.as_mut() {
10298 // Publish at the requested miss-LCP stop (the shared-prefix class)
10299 // AND at the stable-boundary stop (the next-turn re-render class,
10300 // lane/frspec-multiturn-cache) — the same boundary set the plain
10301 // prefill tick learns. Without the second entry, the turn after a
10302 // cold re-park could only hit the OLDER lcp entry (the measured
10303 // one-turn transient: t3 restored 607 of 24122 while the plain arm
10304 // rewound to 15222). Dedupe is the worker sweep's has_key.
10305 if *requested == Some(seg_end) || ckpt_rel == Some(seg_end) {
10306 if let Ok(snap) = cache.snapshot(e) {
10307 slot.push(SpecBoundaryCapture {
10308 snap,
10309 pos: seg_end,
10310 logits: prime_logits.clone(),
10311 // rows [0..seg_end) of h_all are primed — the following
10312 // segments append, never overwrite.
10313 last_h: capture_boundary_hidden(e, &h_all, seg_end, n_embd),
10314 });
10315 }
10316 }
10317 }
10318 }
10319 // SESSION-AFFINITY TURN CHECKPOINT at the STABLE boundary (see `ckpt_at`):
10320 // same snapshot mechanics, installed post-prime in place of the prompt-end
10321 // capture the re-render class always diverged below.
10322 if ckpt_rel == Some(seg_end) {
10323 let anchor: Result<CudaSlice<f32>, Box<dyn std::error::Error>> =
10324 e.uninit(n_embd).and_then(|mut a| {
10325 e.copy_view_into(
10326 &mut a,
10327 0,
10328 &h_all.slice((seg_end - 1) * n_embd..seg_end * n_embd),
10329 n_embd,
10330 )?;
10331 Ok(a)
10332 });
10333 ckpt_early = Some(match (cache.snapshot(e), anchor) {
10334 (Ok(snap), Ok(last_h)) => Some(SpecCheckpoint {
10335 snap,
10336 pos: base + seg_end,
10337 last_h,
10338 }),
10339 _ => None,
10340 });
10341 }
10342 }
10343 if std::env::var("MEMRA_SPEC_STATS").as_deref() == Ok("1") {
10344 eprintln!(
10345 "[spec-prime] stops={stops:?} tail={}",
10346 prompt.len() - stops.last().copied().unwrap_or(0)
10347 );
10348 }
10349 prompt_h = Some(h_all);
10350 } else if batched_prime {
10351 let (l, _h_seed, hiddens) = self.prime_cache(e, prompt, &mut *cache, 0)?;
10352 prime_logits = l;
10353 prompt_h = Some(hiddens);
10354 } else {
10355 prime_logits = Vec::new();
10356 prompt_h = Some(e.uninit(prompt.len() * n_embd)?);
10357 for (i, &tok) in prompt.iter().enumerate() {
10358 let (l, h) = self.spec_target_step_h(e, tok, &mut *cache)?;
10359 if let Some(ph) = prompt_h.as_mut() {
10360 e.copy_into(ph, i * n_embd, &h, n_embd)?;
10361 }
10362 prime_logits = l;
10363 }
10364 }
10365 e.stream().synchronize()?;
10366 // PREFIX-CACHE SEED CAPTURE (lane/spec-prefix-cache): boundary == prompt end (the seed
10367 // case — no shared-prefix split, publish the whole prompt). The prime just finished, so
10368 // cache.pos == base + prompt.len() and the recurrent state IS the boundary state;
10369 // prime_logits are the boundary logits. Cold sessions only (base == 0) — same law as
10370 // prime_split. The mid-prompt capture above already consumed the request if it matched.
10371 if !continuation && base == 0 {
10372 if let Some((requested, slot)) = sess_capture.as_mut() {
10373 if *requested == Some(prompt.len()) && slot.is_empty() {
10374 debug_assert_eq!(cache.pos, prompt.len(), "seed capture off prompt end");
10375 if let Ok(snap) = cache.snapshot(e) {
10376 slot.push(SpecBoundaryCapture {
10377 snap,
10378 pos: prompt.len(),
10379 logits: prime_logits.clone(),
10380 last_h: prompt_h
10381 .as_ref()
10382 .map(|ph| capture_boundary_hidden(e, ph, prompt.len(), n_embd))
10383 .unwrap_or_default(),
10384 });
10385 }
10386 }
10387 }
10388 }
10389 // Harness timing contract (see crate::PRIME_NANOS): gen-only throughput without the
10390 // prime-subtraction hack.
10391 crate::PRIME_NANOS.store(
10392 t_prime.elapsed().as_nanos() as u64,
10393 std::sync::atomic::Ordering::Relaxed,
10394 );
10395
10396 let (embd_qt, embd_rb) = self.embd.qt_and_row_bytes(n_embd);
10397 // Resident table is fastest when it fits. Large spill deployments can preserve that HBM
10398 // for expert-cache slots and gather only the exact rows needed by MTP/verify from host.
10399 let host_embd = spec_host_embd();
10400 let embd_gpu = if host_embd {
10401 None
10402 } else {
10403 Some(
10404 self.embd_gpu
10405 .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload")),
10406 )
10407 };
10408 let embd_dev = embd_gpu.map(|g| (g, embd_qt, embd_rb));
10409 if host_embd {
10410 eprintln!(
10411 "[spec] host-row embedding: {} bytes kept off HBM",
10412 self.embd.raw.len()
10413 );
10414 }
10415 let mut out: Vec<u32> = Vec::with_capacity(max_new);
10416 let mut total_drafted = 0usize;
10417 let mut total_accepted = 0usize;
10418
10419 // --- SAMPLER FIRST (lane/sampled-spec-quality, 2026-08-19) ---
10420 // The sampler config, the session's Philox counters and the penalty window are parsed
10421 // HERE, above the boundary-token selection, because the boundary token must be drawn
10422 // from the sampler the request asked for. Pre-lane this block sat ~50 lines BELOW the
10423 // selection, which is the whole mechanical reason the boundary token was an argmax:
10424 // the sampler state was not in scope yet. Nothing here depends on the round loop, so
10425 // moving it up is a pure reordering for greedy (`sampled == false` ⇒ every branch
10426 // below takes the argmax path it always took).
10427 // --- SAMPLED SPEC (MEMRA_SPEC_TEMP>0, research/sampled-spec-impl-map.md): rejection-
10428 // sampling verify (Leviathan/Chen) — accept draft x at u < p(x)/q(x), resample from
10429 // norm(max(0,p-q)) on reject, bonus sampled from p on full accept. Counter-based Philox
10430 // everywhere (seed, event) -> reproducible. temp==0/unset = the greedy path, untouched.
10431 let sp = sampling.unwrap_or_else(|| SpecSampling {
10432 temp: std::env::var("MEMRA_SPEC_TEMP")
10433 .ok()
10434 .and_then(|v| v.parse().ok())
10435 .unwrap_or(0.0),
10436 seed: std::env::var("MEMRA_SEED")
10437 .ok()
10438 .and_then(|v| v.parse().ok())
10439 .unwrap_or(42),
10440 top_k: std::env::var("MEMRA_TOP_K")
10441 .ok()
10442 .and_then(|v| v.parse().ok())
10443 .unwrap_or(0),
10444 top_p: std::env::var("MEMRA_TOP_P")
10445 .ok()
10446 .and_then(|v| v.parse().ok())
10447 .unwrap_or(1.0),
10448 min_p: std::env::var("MEMRA_MIN_P")
10449 .ok()
10450 .and_then(|v| v.parse().ok())
10451 .unwrap_or(0.0),
10452 penalty_last_n: std::env::var("MEMRA_PENALTY_LAST_N")
10453 .ok()
10454 .and_then(|v| v.parse().ok())
10455 .unwrap_or(0),
10456 penalty_repeat: std::env::var("MEMRA_PENALTY_REPEAT")
10457 .ok()
10458 .and_then(|v| v.parse().ok())
10459 .unwrap_or(1.0),
10460 penalty_freq: std::env::var("MEMRA_PENALTY_FREQ")
10461 .ok()
10462 .and_then(|v| v.parse().ok())
10463 .unwrap_or(0.0),
10464 penalty_present: std::env::var("MEMRA_PENALTY_PRESENT")
10465 .ok()
10466 .and_then(|v| v.parse().ok())
10467 .unwrap_or(0.0),
10468 });
10469 let (sp_temp, sp_seed) = (sp.temp, sp.seed);
10470 let sampled = sp_temp > 0.0;
10471 // Counters resume from the session (burst continuity: randomness must never repeat
10472 // across generate_spec_session calls); one-shot callers start at (0,0). Read through
10473 // sess_tail — `sess` was take()n into it above, so sess.as_ref() here is always None.
10474 let mut sctr: u32 = sess_tail.as_ref().map(|(_, _, _, s, _)| **s).unwrap_or(0);
10475 let mut uctr: u32 = sess_tail.as_ref().map(|(_, _, _, _, u)| **u).unwrap_or(0);
10476 // Penalties (v2.1): applied to COPIES of q rows and p columns symmetrically (exactness
10477 // for the penalized+filtered target). History = generated tokens, host-tracked window.
10478 let pen_on = sampled
10479 && sp.penalty_last_n > 0
10480 && (sp.penalty_repeat != 1.0 || sp.penalty_freq != 0.0 || sp.penalty_present != 0.0);
10481 // SESSION-SPANNING PENALTY WINDOW (Item 2). Pre-lane this was
10482 // `prompt.iter().rev().take(64).rev()` — the BURST's suffix slice — so a continuation
10483 // burst (the majority of a stream's tokens, and ALL of a converted cache hit's) started
10484 // with an EMPTY penalty history and the client's repetition/frequency/presence penalties
10485 // silently reset at every burst boundary. The window now spans `committed ++ prompt`,
10486 // which is what the API contract says and what the plain sampler's own `history` does.
10487 // Byte-identical to the pre-lane seed for a cold turn-1 burst at the default window.
10488 let mut pen_hist: Vec<u32> = if pen_on {
10489 let sess_hist: &[u32] = if spec_pen_session_on() {
10490 sess_tail
10491 .as_ref()
10492 .map(|(c, ..)| c.as_slice())
10493 .unwrap_or(&[])
10494 } else {
10495 &[] // MEMRA_SPEC_PEN_SESSION=0: pre-lane burst-local window
10496 };
10497 pen_window_seed(sess_hist, prompt, sp.penalty_last_n)
10498 } else {
10499 Vec::new()
10500 };
10501 // First generated token = the BOUNDARY token: greedy takes the argmax of the prompt's
10502 // last logits (== greedy's first token, byte-contract); SAMPLED draws it from the
10503 // request's own filtered/penalized target through the session's Philox stream
10504 // (`sample_boundary_token`, lane/sampled-spec-quality Item 1 — pre-lane this was an
10505 // argmax in both regimes, so ~1 token per burst of a sampled stream was greedy).
10506 // Emit it, then FEED it to establish the loop invariant below.
10507 // PENDING-CARRY: the carried bonus was already emitted by the LAST burst — it becomes
10508 // last_token WITHOUT re-emission, and round 0 consumes it as pending (no init feed).
10509 // CONSTRAINED entry rules: the first emitted token is the MASKED argmax of the
10510 // prompt's last logits (plain constrained-greedy identity); a continuation without
10511 // a carried pending would emit an UNMASKED stashed next_pred — refused loudly (the
10512 // worker never resumes constrained sessions from the pool, so this cannot fire).
10513 if let Some(c) = constraint.as_deref_mut() {
10514 if continuation && carried_pending.is_none() {
10515 return Err("constrained spec continuation requires a carried pending \
10516 (pool resume is unconstrained-only)"
10517 .into());
10518 }
10519 if !continuation {
10520 c.mask_logits(&mut prime_logits)
10521 .map_err(|e2| format!("constraint: {e2}"))?;
10522 }
10523 }
10524 let mut last_token = if let Some(b) = carried_pending {
10525 b
10526 } else if continuation {
10527 // A continuation's boundary token was DRAWN by the burst that stashed it (the
10528 // session tail below), or by `spec_session_from_restored` for a converted
10529 // prefix-cache hit — in both cases from the correct logits row with this same
10530 // session's Philox stream, which is why it can be consumed here as-is.
10531 sess_tail.as_ref().unwrap().2.unwrap()
10532 } else if sampled && constraint.is_none() && spec_sampled_boundary_on() {
10533 sample_boundary_token(e, &prime_logits, &sp, &pen_hist, &mut sctr, "cold-prime")?
10534 } else {
10535 // greedy (byte contract), the rollback door, or constrained (masked-argmax
10536 // identity — the worker routes sampled+constrained to the plain path, and this
10537 // function refuses the combination outright above).
10538 argmax(&prime_logits) as u32
10539 };
10540 if pen_on {
10541 // The boundary token is a GENERATED token: the plain sampler `accept()`s every
10542 // emitted token into its penalty history, and pre-lane the burst's first token
10543 // was invisible to penalties forever (never pushed, and never in `committed`
10544 // until this burst's tail). Covers the carry/continuation seeds too — neither is
10545 // in `committed` yet.
10546 pen_hist.push(last_token);
10547 }
10548 if carried_pending.is_none() {
10549 out.push(last_token);
10550 // grammar advances with every emitted token (carried pendings were consumed
10551 // by the burst that emitted them).
10552 if let Some(c) = constraint.as_deref_mut() {
10553 c.consume(last_token)
10554 .map_err(|e2| format!("constraint: {e2}"))?;
10555 }
10556 }
10557 if continuation {
10558 // draft-KV invariant: entries [0..base) are the session's exact fills; truncate any
10559 // overhang so the chain's first append lands at slot base (== committed.len()).
10560 scratch.set_len(e, base)?;
10561 }
10562 // sse-cadence: hand the caller every not-yet-flushed token (disjoint in-order slices
10563 // concatenating to the full `out`). Called after the prime's first token and after each
10564 // round commit — emission timing only, token bytes untouched. The slice may be EMPTY
10565 // (poll-only boundary: zero-round folds commit nothing new); returns the caller's
10566 // continue-verdict (admission yield, 2026-08-06) — false ends the burst at this round.
10567 fn flush_commit(
10568 cb: &mut Option<&mut dyn FnMut(&[u32]) -> bool>,
10569 out: &[u32],
10570 flushed: &mut usize,
10571 ) -> bool {
10572 if let Some(f) = cb.as_mut() {
10573 let keep = f(&out[*flushed..]);
10574 *flushed = out.len();
10575 keep
10576 } else {
10577 true
10578 }
10579 }
10580 keep_going = flush_commit(&mut on_commit, &out, &mut flushed);
10581 // INVARIANT at loop top: `last_token` is the most-recently-committed/emitted token, its
10582 // KV+recur state IS in `cache` (cache.pos = position right AFTER last_token), `last_pred`
10583 // is the greedy ARGMAX of the logits that predict the token FOLLOWING last_token, and
10584 // `h_seed` = last_token's pre-output_norm hidden. Establish it by feeding last_token once
10585 // (mirrors plain greedy). DEVICE-ARGMAX lever: the accept walk only ever consumes the
10586 // argmax of those logits — never the full vector — so a host u32 replaces the Vec<f32>.
10587 // Trimmed heads: q lives on the trimmed vocab; accept gathers use the TRIMMED index and
10588 // the residual scatters q into target-id space (q=-inf off-trim — the head cannot propose
10589 // those, so their residual mass is p(x), correct by construction).
10590 let d2t_dev: Option<CudaSlice<u32>> = if sampled || crate::spec::spec_stream() {
10591 match &mtp.d2t {
10592 Some(map) => Some(e.htod_u32_v(map)?),
10593 None => None,
10594 }
10595 } else {
10596 None
10597 };
10598 let mut q_full_buf: Option<CudaSlice<f32>> = None;
10599 // host Philox4x32-10 accept-test uniforms: module fn `host_u01` (shared with the
10600 // dspark sampled-admission walk); byte-identical to the closure it replaces.
10601 let mut draft_logits: Vec<CudaSlice<f32>> = Vec::new(); // retained head logits (q), per slot
10602 let mut draft_stats: Vec<(f32, f32, f32)> = Vec::new(); // (row_max, th_e, z_e) per slot
10603 let mut perturb_buf: Option<CudaSlice<f32>> = None; // gumbel scratch (max(n_vocab,d_vocab))
10604 let mut sample_tok = e.alloc_u32_zeroed(1)?; // residual/bonus sample out
10605 let mut col_buf: Option<CudaSlice<f32>> = None; // materialized verify column
10606 let mut pen_hist_d: Option<CudaSlice<u32>> = None;
10607 let mut pcol_buf: Option<CudaSlice<f32>> = None; // penalized p-column scratch
10608 // MEMRA_SPEC_SETUP_TRACE=1 (diagnostics): per-call wall decomposition of the burst
10609 // SETUP + TAIL segments (the round loop's internals are MEMRA_SPEC_PHASE's job) —
10610 // built to pin the serve per-burst fixed cost (research/spec-serving-20260801).
10611 let setup_trace = std::env::var("MEMRA_SPEC_SETUP_TRACE").as_deref() == Ok("1");
10612 let t_ent = std::time::Instant::now();
10613
10614 // SESSION-AFFINITY TURN CHECKPOINT (lane/session-affinity, 2026-08-05): capture the
10615 // PROMPT-END boundary state so a LATER turn can rewind here and re-prime only its own
10616 // delta instead of the whole conversation. See `SpecCheckpoint` for why this boundary is
10617 // the one that matters (a history-rewriting client mutates what the session GENERATED,
10618 // so the next turn's prompt agrees with this one up to exactly here).
10619 //
10620 // WHERE — AND WHY THIS EXACT LINE. Right after the trunk prime, BEFORE the init feed
10621 // (`decode_step_h(last_token)`) and before round 0: the last instant at which the caches
10622 // hold exactly `base + prompt.len()` rows and nothing generated.
10623 //
10624 // This was WRONG in the first cut of this lane: the capture sat after the draft-KV fill,
10625 // which is also after the init feed, so `cache.pos` was `base + prompt.len() + 1` — the
10626 // boundary included the FIRST GENERATED TOKEN. That token is the first thing inside the
10627 // `<think>` block the client strips, so every later turn's diff diverged exactly one
10628 // token below the checkpoint and affinity declined 100% of the time. Measured on the
10629 // owner regime: "history diverged at 12233 of checkpoint 12234". The off-by-one made the
10630 // whole mechanism inert while looking, from the outside, like a working
10631 // correctness-declines-safely path — hence the decline log carries the offsets.
10632 //
10633 // The full-attn planes are `len`-truncatable so the snapshot copies only the GDN conv/ssm
10634 // state (the reason a spec session could not rewind before). The draft scratch needs no
10635 // copy: rows below the boundary are rewritten by the next turn's own fill.
10636 //
10637 // WHEN: non-empty prime only. An empty-suffix continuation burst adds no prompt boundary
10638 // (its "prompt end" IS the previous checkpoint's, already held), so it keeps the existing
10639 // checkpoint rather than replacing it with a strictly worse one.
10640 //
10641 // FAILURE IS SILENT BY DESIGN: on a VRAM-tight rig the snapshot alloc can fail. That
10642 // costs the NEXT turn its rewind (it re-primes fully, today's behavior) and must never
10643 // fail the burst that is already running — so the error is swallowed, loud only under
10644 // MEMRA_DEBUG_SPEC.
10645 //
10646 // STABLE-BOUNDARY OVERRIDE (lane/frspec-multiturn-cache, 2026-08-21): the prompt-end
10647 // posture above was DISPROVED for the think-posture template class — the prompt's own
10648 // tail is the live generation header (`<|im_start|>assistant\n<think>\n`) that the
10649 // next turn's re-render replaces, so the diff diverged a couple tokens BELOW the
10650 // checkpoint and affinity declined 100% of multi-turn agent traffic (the same class
10651 // the plain tier fixed on 2026-08-09 via `plain_checkpoint_boundary`; the port to the
10652 // spec tier is this lane). When the worker armed `ckpt_at`, the capture happened at
10653 // that stop inside the prime above (`ckpt_early`) and is installed here instead;
10654 // capture-attempted-but-failed clears the slot exactly like the legacy arm.
10655 if let Some(slot) = sess_ckpt_slot {
10656 if let Some(early) = ckpt_early {
10657 if early.is_none() && std::env::var("MEMRA_DEBUG_SPEC").is_ok() {
10658 eprintln!(
10659 "[spec] stable-boundary turn checkpoint skipped; \
10660 next turn re-primes in full"
10661 );
10662 }
10663 *slot = early;
10664 } else if !continuation {
10665 let pos = cache.pos;
10666 debug_assert_eq!(
10667 pos,
10668 base + prompt.len(),
10669 "turn checkpoint must sit at the prompt end, before the init feed"
10670 );
10671 let anchor: Result<CudaSlice<f32>, Box<dyn std::error::Error>> =
10672 if let Some(ph) = &prompt_h {
10673 // hidden of the LAST primed row = the predecessor anchor at this
10674 // boundary (exactly what a fresh prime of committed[..pos] leaves in
10675 // last_h, and what the next prime's fill reads for its first row).
10676 let np = prompt.len();
10677 e.uninit(n_embd).and_then(|mut a| {
10678 e.copy_view_into(
10679 &mut a,
10680 0,
10681 &ph.slice((np - 1) * n_embd..np * n_embd),
10682 n_embd,
10683 )?;
10684 Ok(a)
10685 })
10686 } else {
10687 Err("no prompt hiddens".into())
10688 };
10689 match (cache.snapshot(e), anchor) {
10690 (Ok(snap), Ok(last_h)) => {
10691 *slot = Some(SpecCheckpoint { snap, pos, last_h });
10692 }
10693 (s, a) => {
10694 *slot = None; // a stale checkpoint would rewind to the WRONG boundary
10695 if std::env::var("MEMRA_DEBUG_SPEC").is_ok() {
10696 let err = s
10697 .err()
10698 .map(|e| e.to_string())
10699 .or_else(|| a.err().map(|e| e.to_string()))
10700 .unwrap_or_default();
10701 eprintln!(
10702 "[spec] turn checkpoint skipped ({err}); \
10703 next turn re-primes in full"
10704 );
10705 }
10706 }
10707 }
10708 }
10709 }
10710 // INIT FEED — skipped on a pending carry: last_token (the carried bonus) is NOT in the
10711 // caches and must NOT be fed solo; round 0's batched verify commits it as col 0. Its
10712 // seed/anchor hidden is the carried last_h (copied below); last_pred is dead in the
10713 // pending path (t_pred reads verify col 0 — the accept walk overwrites it).
10714 let mut last_pred = 0u32;
10715 let mut last_col_logits: Option<CudaSlice<f32>> = None;
10716 // CONSTRAINED: the init feed's logits back the (n_acc==0, base==0) masked-argmax
10717 // recompute in the grammar-truncation walk — retained host-side, round 0 only.
10718 let mut init_logits_host: Option<Vec<f32>> = None;
10719 let h_seed0: CudaSlice<f32> = if carried_pending.is_none() {
10720 let (init_logits, h) = self.spec_target_step_h(e, last_token, &mut *cache)?;
10721 last_pred = argmax(&init_logits) as u32;
10722 if constraint.is_some() {
10723 init_logits_host = Some(init_logits.clone());
10724 }
10725 // sampled mode: p-distribution after last_token, for the j==0/base==0 accept test.
10726 if sampled {
10727 last_col_logits = Some(e.htod(&init_logits)?);
10728 }
10729 h
10730 } else {
10731 // predecessor-row anchor: hidden of the last COMMITTED row (the carry contract).
10732 let lh = sess_tail
10733 .as_ref()
10734 .unwrap()
10735 .1
10736 .as_ref()
10737 .expect("pending carry requires last_h");
10738 e.clone_dtod(lh)?
10739 };
10740 let t_init = t_ent.elapsed();
10741 let mut last_col_stats: Option<(f32, f32, f32)> = None;
10742 // PERSISTENT h_seed buffer (allocated BEFORE any graph capture so no captured scratch can
10743 // alias it): every path that updates the round seed copies INTO it — no per-round allocs,
10744 // stable pointer for the graph-draft round-start copy.
10745 let mut h_seed_buf = e.clone_dtod(&h_seed0)?;
10746 // Predecessor-pairing trackers: `fill_prev` = trunk hidden AT the last COMMITTED row (the
10747 // predecessor of the next verify's col 0 — the reference's carried pending-h analogue;
10748 // also the predecessor-row hidden for the round-0 legacy-replay seed). At round 0 that
10749 // row is last_token's own (h_seed0). The chain step-0 seed under the pairing default =
10750 // hidden of the row BEFORE last_token = the prompt's last row at round 0 (h_seed_buf
10751 // overwritten below).
10752 let mut fill_prev = e.clone_dtod(&h_seed0)?;
10753 {
10754 if let Some(ph) = &prompt_h {
10755 let np = prompt.len();
10756 e.copy_view_into(
10757 &mut h_seed_buf,
10758 0,
10759 &ph.slice((np - 1) * n_embd..np * n_embd),
10760 n_embd,
10761 )?;
10762 } else if continuation {
10763 if let Some((_, lh, _, _, _)) = sess_tail.as_ref() {
10764 if let Some(lh) = lh.as_ref() {
10765 e.copy_into(&mut h_seed_buf, 0, lh, n_embd)?;
10766 }
10767 }
10768 }
10769 }
10770 // Persistent device prediction slots for the accept walk (max k+1 verify columns).
10771 let mut preds_d = e.alloc_u32_zeroed(k + 2)?;
10772
10773 let debug_spec = std::env::var("MEMRA_DEBUG_SPEC").is_ok();
10774 let fork_mode = OptiForkGateMode::configured();
10775 // MEMRA_SPEC_STATS=1: per-slot accept histogram + draft-length histogram, printed once at
10776 // the end. Metric normalization vs the reference engine: BOTH engines count
10777 // accepted/drafted where the chain stopped at p-min and the sub-threshold token is
10778 // discarded uncounted — per-slot decay + chain-length mix are the extra dimensions.
10779 let spec_stats = std::env::var("MEMRA_SPEC_STATS").is_ok();
10780 let mut st_drafted = vec![0usize; k];
10781 let mut st_accepted = vec![0usize; k];
10782 let mut st_len_hist = vec![0usize; k + 1];
10783 let mut st_full = 0usize;
10784 // P-MIN CONFIDENCE GATE (MEMRA_SPEC_PMIN, the serve script's --spec-draft-p-min mechanism):
10785 // stop the draft chain early when the head's softmax confidence in its own pick drops
10786 // below p_min. Hoisted above the loop: the graph capture bakes the prob kernels iff on.
10787 static PMIN: std::sync::OnceLock<f32> = std::sync::OnceLock::new();
10788 let p_min = *PMIN.get_or_init(|| {
10789 std::env::var("MEMRA_SPEC_PMIN")
10790 .ok()
10791 .and_then(|v| v.parse().ok())
10792 .unwrap_or(0.0)
10793 });
10794 // ZERO-DRAFT ROUNDS (MEMRA_SPEC_PMIN0=1, vendored from llama.cpp's draft gating): let the
10795 // p-min gate apply at j==0 too, so a low-confidence round drafts NOTHING and the verify
10796 // batch is just the pending bonus (m=1 = a plain decode step). llama's 35B win rides
10797 // exactly this — draft acceptance 76% at mean len 2.5 because unpredictable stretches
10798 // never pay draft+verify overhead. Only legal when a pending bonus exists (an empty
10799 // verify batch is not); the j==0 exemption stays for pending-less rounds.
10800 let pmin0 = std::env::var("MEMRA_SPEC_PMIN0")
10801 .map(|v| v == "1")
10802 .unwrap_or(false);
10803
10804 // --- GRAPH DRAFT setup: persistent I/O buffers + ONE capture (2 warmups inside). The
10805 // warmups mutate scratch len_d / pos / tok / seed — all reset at every round start, so the
10806 // only restore needed is the scratch counter. Capture failure (e.g. a non-capturable
10807 // cuBLAS path in an exotic head) falls back to the eager draft chain.
10808 // PER-SESSION PERSISTENCE (2026-08-01): session calls reuse the DraftGraphCtx parked on
10809 // the SpecSession — the capture (2 warmup head forwards + instantiate) ran ONCE at the
10810 // session's first burst, not per burst (measured ~16ms/burst fixed cost on H100 q27,
10811 // research/spec-serving-20260801). Reuse is pointer-exact: the graph bakes the session's
10812 // own scratch KV (never realloc'd), the model's resident embedding, the OnceLock p_min,
10813 // and the g_* buffers carried in the ctx — replay dispatch is identical to a fresh
10814 // capture, so draft tokens are bit-identical (drafts never decide exactness anyway; the
10815 // verify arbitrates). Single-shot calls (sess=None) build a fresh ctx and drop it.
10816 let mut dctx: DraftGraphCtx = match sess_draft_slot.as_mut().and_then(|s| s.take()) {
10817 Some(c) => c,
10818 None => DraftGraphCtx::new(e, n_embd, if sampled { d_vocab } else { 1 })?,
10819 };
10820 // A session that ran greedy bursts first sized g_q/g_perturb at 1; a sampled resume
10821 // needs d_vocab. Realloc is legal exactly while graph_s is None (nothing baked them).
10822 if sampled && dctx.g_q.len() < d_vocab {
10823 dctx.g_q = e.zeros(d_vocab)?;
10824 dctx.g_perturb = e.zeros(d_vocab)?;
10825 }
10826 // DRAFT-SIDE GRAMMAR MASK (lane/draft-mask, 2026-08-04): the drafter samples the
10827 // grammar's legal set, so proposals are legal BY CONSTRUCTION and the verify-side
10828 // truncation (the correctness backstop) stops cutting every tight-schema round.
10829 // The mask is one node inside the captured draft chain — presence is a CAPTURE-TIME
10830 // shape, so a parked graph of the other shape is dropped and recaptured.
10831 let dmask_on = constraint
10832 .as_deref()
10833 .is_some_and(|c| c.draft_mask_enabled());
10834 let dmask_words = if dmask_on { d_vocab.div_ceil(32) } else { 0 };
10835 if dmask_on && dctx.g_dmask.len() < dmask_words {
10836 dctx.g_dmask = e.alloc_u32_zeroed(dmask_words)?;
10837 dctx.graph = None; // the old capture baked the old (or no) mask pointer
10838 dctx.failed.clear_greedy();
10839 dctx.keeper.clear();
10840 }
10841 if dctx.graph.is_some() && dctx.graph_masked != dmask_on {
10842 dctx.graph = None;
10843 dctx.failed.clear_greedy();
10844 dctx.keeper.clear();
10845 }
10846 if graph_draft && !sampled && dctx.graph.is_none() && !dctx.failed.greedy_failed() {
10847 let DraftGraphCtx {
10848 g_tok,
10849 g_pos,
10850 g_seed,
10851 g_p,
10852 g_dmask,
10853 ..
10854 } = &mut dctx;
10855 // capture-time contents: ALL-ONES (ban nothing). A replay only ever runs after the
10856 // host uploads the position's real words, so the warmups stay grammar-free.
10857 if dmask_on {
10858 e.htod_u32_into(g_dmask, &vec![u32::MAX; dmask_words])?;
10859 }
10860 let g_dmask_ro: &CudaSlice<u32> = &*g_dmask;
10861 // CAPTURE-RETAIN (#68 fix): the warmup transients' pool addresses are baked into the
10862 // captured graph; the keeper pins them for the graph's lifetime. capture_graph (non-
10863 // retained) freed them at exit — safe for one-shot generate_spec (nothing else touches
10864 // the pool between replays) but WRONG for sessions: burst-boundary prime/fill/commit
10865 // passes (and, in serve, other sessions) recycle those addresses and the replay then
10866 // clobbers live buffers — the ST serve-spec corruption (research/serve-st-20260803).
10867 let cap_res = e.capture_graph_retained(|e| {
10868 self.mtp_head_forward_cap(
10869 e,
10870 mtp,
10871 g_tok,
10872 g_pos,
10873 g_seed,
10874 g_p,
10875 &mut *scratch,
10876 p_min > 0.0 || fork_mode == OptiForkGateMode::Controller,
10877 true,
10878 embd_gpu.expect("graph draft requires resident embedding"),
10879 embd_qt,
10880 embd_rb,
10881 d_vocab,
10882 None,
10883 None,
10884 if dmask_on {
10885 Some((g_dmask_ro, dmask_words))
10886 } else {
10887 None
10888 },
10889 )
10890 });
10891 match cap_res {
10892 Ok((g, keep)) => {
10893 scratch.set_len(e, base)?;
10894 dctx.graph = Some(g);
10895 dctx.graph_masked = dmask_on;
10896 dctx.keeper = keep;
10897 }
10898 Err(err) => {
10899 scratch.set_len(e, base)?;
10900 // LOUD flip (audit Q2): a dropped draft graph is a coverage loss, never
10901 // silent. Once per flip — mark returns None on an already-failed ctx.
10902 if let Some(line) = dctx.failed.mark_greedy(&err.to_string()) {
10903 eprintln!("{line}");
10904 }
10905 }
10906 }
10907 }
10908 // --- SAMPLED GRAPH DRAFT setup (step 3 of the sampled-spec arc): a SECOND capture, own
10909 // graph object, built only when sampled && graph-eligible — the greedy capture above is
10910 // untouched (and skipped when sampled: its graph would never be launched). Same head
10911 // forward, but the in-graph argmax reads GUMBEL-PERTURBED logits; the Philox event
10912 // counter lives in the persistent device g_ctr (bumped in-graph, host-seeded from sctr
10913 // once per round); the raw head logits land in the persistent g_q for the host's
10914 // per-replay async D2D into the round's q slot (q_slots, K x d_vocab, allocated once).
10915 // seed/temp are capture-time constants — baked into graph_s, so a pool-resumed request
10916 // with a different (seed, temp, k) drops the parked sampled graph and recaptures.
10917 // COST OF THE FRESH-SEED SERVE DEFAULT (dogfood F4, 2026-08-04): omitting `seed` on a
10918 // serve request now draws fresh per-request entropy (it used to default to a pinned 0),
10919 // so a seed-omitting request that RESUMES a parked spec session finds an s_key baked
10920 // with the PREVIOUS request's seed and pays one recapture. Bounded, and it does not
10921 // reopen the ~16ms/burst regression the persistent ctx exists to fix: a session's seed
10922 // is fixed for its whole lifetime (worker.rs reads s.sampler.seed() per burst), so
10923 // this compare misses at most ONCE per resumed request — the first burst recaptures
10924 // and every later burst in that request replays. A client that wants the parked graph
10925 // AND reproducibility supplies an explicit `seed`, honored exactly, which keeps s_key
10926 // stable across its whole conversation.
10927 // COMPOSITION RULE (fspec x gsd merge): the in-graph chain samples from the RAW
10928 // softmax — it can hold neither per-row filter stats nor the varying penalty history.
10929 // The sampled graph therefore engages only in the PURE-TEMP regime; filters/penalties
10930 // force the eager draft (which computes stats/penalties per row).
10931 // KEY THE WHOLE REGIME, not just the baked constants (lane/graph-s-key-exactness-
10932 // 20260819). `s_key` used to be `(seed, temp, k)`; the filters and penalties were left
10933 // out, so a filtered request resuming a session that parked a PURE-TEMP graph kept it —
10934 // and the launch site never re-asked `pure_temp`. See [`SampledGraphKey`] for what that
10935 // costs (an unconditional accept of out-of-head draft tokens, i.e. an exactness bug on
10936 // the request shape the vendor-default flip makes the majority).
10937 let s_key = SampledGraphKey::new(sp_seed, sp_temp, k, sp.top_k, sp.top_p, sp.min_p, pen_on);
10938 let pure_temp = s_key.pure_temp();
10939 if sampled && dctx.s_key.is_some_and(|old| old != s_key) {
10940 dctx.graph_s = None;
10941 dctx.failed.clear_sampled();
10942 dctx.s_key = None;
10943 dctx.q_slots.clear();
10944 dctx.keeper_s.clear();
10945 }
10946 if graph_draft
10947 && sampled
10948 && pure_temp
10949 && dctx.graph_s.is_none()
10950 && !dctx.failed.sampled_failed()
10951 {
10952 let DraftGraphCtx {
10953 g_tok,
10954 g_pos,
10955 g_seed,
10956 g_p,
10957 g_ctr,
10958 g_perturb,
10959 g_q,
10960 ..
10961 } = &mut dctx;
10962 // CAPTURE-RETAIN (#68 fix): same keeper contract as the greedy capture above.
10963 let cap_res = e.capture_graph_retained(|e| {
10964 self.mtp_head_forward_cap(
10965 e,
10966 mtp,
10967 g_tok,
10968 g_pos,
10969 g_seed,
10970 g_p,
10971 &mut *scratch,
10972 p_min > 0.0,
10973 true,
10974 embd_gpu.expect("graph draft requires resident embedding"),
10975 embd_qt,
10976 embd_rb,
10977 d_vocab,
10978 Some((g_ctr, g_perturb, g_q, sp_seed, sp_temp)),
10979 None,
10980 None, // constrained spec is greedy-only — sampled never carries a hook
10981 )
10982 });
10983 match cap_res {
10984 Ok((g, keep)) => {
10985 scratch.set_len(e, base)?;
10986 for _ in 0..k {
10987 dctx.q_slots.push(e.zeros(d_vocab)?);
10988 }
10989 dctx.graph_s = Some(g);
10990 dctx.s_key = Some(s_key);
10991 dctx.keeper_s = keep;
10992 }
10993 Err(err) => {
10994 scratch.set_len(e, base)?;
10995 // LOUD flip (audit Q2): same contract as the greedy capture above.
10996 if let Some(line) = dctx.failed.mark_sampled(&err.to_string()) {
10997 eprintln!("{line}");
10998 }
10999 }
11000 }
11001 }
11002 // ---- EXACTNESS GUARD, the enforceable half (lane/graph-s-key-exactness-20260819) ----
11003 // With the filters and penalties in `s_key`, a graph that SURVIVED the drop above was
11004 // captured under this request's exact regime, and capture requires `pure_temp` — so a
11005 // parked graph implies `pure_temp`. That implication is the whole exactness argument for
11006 // the graph arm, so it is asserted here rather than assumed: a future change that widens
11007 // the capture condition, narrows the key, or copies a `DraftGraphCtx` across regimes
11008 // fails LOUDLY at this line instead of silently drafting from the raw softmax while the
11009 // verify applies filter stats. Release builds refuse the graph (drop it, draft eager)
11010 // rather than launching it; the launch site re-tests `pure_temp` independently.
11011 if sampled && !pure_temp && dctx.graph_s.is_some() {
11012 debug_assert!(
11013 false,
11014 "sampled draft graph parked under {:?} survived into a FILTERED request \
11015 (top_k={} top_p={} min_p={} pen_on={}): the in-graph chain draws from the RAW \
11016 softmax, so the verify's filtered q would test a distribution the draft was \
11017 never sampled from",
11018 dctx.s_key, sp.top_k, sp.top_p, sp.min_p, pen_on,
11019 );
11020 eprintln!(
11021 "[spec] BUG: dropping a parked sampled draft graph that outlived its capture \
11022 regime (s_key={:?}, request top_k={} top_p={} min_p={} pen_on={}); drafting \
11023 EAGER — the key must carry every field that shapes q",
11024 dctx.s_key, sp.top_k, sp.top_p, sp.min_p, pen_on,
11025 );
11026 dctx.graph_s = None;
11027 dctx.s_key = None;
11028 dctx.q_slots.clear();
11029 dctx.keeper_s.clear();
11030 }
11031 // SKEY PROBE (MEMRA_SKEY_PROBE=1): the burst-entry facts the reachability question turns
11032 // on — is this request sampled, is it in the pure-temp regime the sampled graph is only
11033 // legal in, and is a graph PARKED from an earlier request of the same session? The launch
11034 // arms below print which chain actually ran, so the probe never restates the condition.
11035 if skey_probe() {
11036 eprintln!(
11037 "[skey] burst sampled={} pure_temp={} temp={} top_k={} top_p={} min_p={} \
11038 pen_on={} k={} graph_draft={} graph_s_parked={} s_key_parked={:?}",
11039 sampled as u8,
11040 pure_temp as u8,
11041 sp_temp,
11042 sp.top_k,
11043 sp.top_p,
11044 sp.min_p,
11045 pen_on as u8,
11046 k,
11047 graph_draft as u8,
11048 dctx.graph_s.is_some() as u8,
11049 dctx.s_key,
11050 );
11051 }
11052 let t_cap = t_ent.elapsed();
11053 // PERSISTENT DRAFT KV: fill the MTP block's K/V for every prompt position from the exact
11054 // trunk hiddens collected during prime — ONE batched K/V-only pass (overwrites any
11055 // capture-warmup garbage; capture left len at 0). last_token (the init feed) needs no
11056 // fill: the first chain step processes it and appends its entry at slot prompt.len().
11057 if let Some(ph) = &prompt_h {
11058 // SESSION: rows [0..base) are the previous turns' exact fills (refresh overwrote them
11059 // with true verify hiddens) — truncate any draft overhang, fill ONLY the suffix at
11060 // global positions [base..base+tp). Fresh call: base==0, identical to before.
11061 scratch.set_len(e, base)?;
11062 // CHUNKED FILL (long-ctx OOM fix, 2026-07-05): mtp_kv_fill's transients scale with its
11063 // T (concat = T*2*n_embd*4B — 1.5GB at 40k) and its concat loop is 2*T launches. The
11064 // fill is a pure sequential append, so chunking is exact: each chunk appends its rows
11065 // at pos0=base+start with the identical per-row math. Same knob as the trunk prime.
11066 let tp = prompt.len();
11067 let fill_chunk: usize = if crate::cache::swa_ring_on() {
11068 crate::hybrid_forward::prime_chunk_tokens(tp, self.layers.len())
11069 } else {
11070 // Preserve the flag-OFF schedule byte-for-byte, including the legacy zero value
11071 // meaning one monolithic fill.
11072 std::env::var("MEMRA_PRIME_CHUNK")
11073 .ok()
11074 .and_then(|v| v.parse().ok())
11075 .unwrap_or(4096)
11076 };
11077 let fill_chunk = if fill_chunk == 0 { tp } else { fill_chunk };
11078 let mut start = 0usize;
11079 while start < tp {
11080 let end = (start + fill_chunk).min(tp);
11081 let tc = end - start;
11082 {
11083 // PREDECESSOR pairing: row i gets h[i-1]; global row 0 a zeros row (the
11084 // reference engine's initial pending-h is zeroed too); a session turn's row 0
11085 // gets the PREVIOUS turn's last committed hidden (sess.last_h). Per chunk:
11086 // rows start..end read h[start-1..end-1] — one dtod into a chunk buffer.
11087 let mut phs = e.zeros(tc * n_embd)?;
11088 let (src_lo, dst_off) = if start == 0 {
11089 (0, n_embd)
11090 } else {
11091 ((start - 1) * n_embd, 0)
11092 };
11093 let n_copy = if start == 0 {
11094 (tc - 1) * n_embd
11095 } else {
11096 tc * n_embd
11097 };
11098 if start == 0 {
11099 if let Some((_, lh, _, _, _)) = sess_tail.as_ref() {
11100 if let Some(lh) = lh.as_ref() {
11101 e.copy_into(&mut phs, 0, lh, n_embd)?;
11102 }
11103 }
11104 }
11105 if n_copy > 0 {
11106 e.copy_view_into(
11107 &mut phs,
11108 dst_off,
11109 &ph.slice(src_lo..src_lo + n_copy),
11110 n_copy,
11111 )?;
11112 }
11113 self.mtp_kv_fill_all(
11114 e,
11115 &prompt[start..end],
11116 &phs,
11117 base + start,
11118 &mut *scratch,
11119 embd_dev,
11120 )?;
11121 }
11122 start = end;
11123 }
11124 }
11125 // MEMRA_PROFILE_SPEC=2: profiler capture starts HERE — after the prime, so an
11126 // `nsys -c cudaProfilerApi` capture contains ONLY the round loop (draft/verify/commit).
11127 // (=1 brackets the whole call in run_spec.rs, prime included.)
11128 if std::env::var("MEMRA_PROFILE_SPEC").as_deref() == Ok("2") {
11129 unsafe extern "C" {
11130 fn cudaProfilerStart() -> i32;
11131 }
11132 unsafe {
11133 cudaProfilerStart();
11134 }
11135 }
11136 // ROUND-STREAM stage (c) 4 (MEMRA_SPEC_STREAM=1, experimental): pre-issued M-round
11137 // bursts with ZERO per-round host readbacks — the accept/seed/rollback/ring kernels
11138 // consume each other's device outputs; the host drains the ring every M rounds. v1
11139 // constraints: greedy, !spec_replay, single-shot, batched-linear layers, no refresh
11140 // fills (acceptance effect A/B-arbitrated), enters from round 1 (pending guaranteed).
11141 // NOTE: not gated on the caller's graph_draft (its trunk_dense conjunct turns the 35B
11142 // MoE off) — the stream capture encloses ONLY the dense MTP head; the head-dense /
11143 // full-prec / k gates are re-derived here and a failed capture degrades to stream-off.
11144 let stream_on = crate::spec::spec_stream()
11145 && !sampled
11146 && !spec_replay
11147 && self.mtp_extra.is_empty()
11148 && constraint.is_none()
11149 && !session_mode
11150 && embd_gpu.is_some()
11151 && !crate::model::full_prec_enabled()
11152 && k + 2 < 96;
11153 let mut stream_graph: Option<cudarc::driver::CudaGraph> = None;
11154 let mut g_tokp2k = e.alloc_u32_zeroed(2 * k.max(1))?;
11155 if stream_on {
11156 let cap = e.capture_graph(|e| {
11157 for j in 0..k.max(1) {
11158 self.mtp_head_forward_cap(
11159 e,
11160 mtp,
11161 &mut dctx.g_tok,
11162 &mut dctx.g_pos,
11163 &mut dctx.g_seed,
11164 &mut dctx.g_p,
11165 &mut *scratch,
11166 true,
11167 true,
11168 embd_gpu.expect("round stream requires resident embedding"),
11169 embd_qt,
11170 embd_rb,
11171 d_vocab,
11172 None,
11173 Some((&mut g_tokp2k, j, d2t_dev.as_ref())),
11174 None, // round-stream requires constraint.is_none() (see stream_on)
11175 )?;
11176 }
11177 Ok(())
11178 });
11179 match cap {
11180 Ok(g) => {
11181 scratch.set_len(e, 0)?;
11182 stream_graph = Some(g);
11183 }
11184 Err(err) => {
11185 scratch.set_len(e, 0)?;
11186 if debug_spec {
11187 eprintln!("[spec] stream-graph capture failed ({err}); stream off");
11188 }
11189 }
11190 }
11191 }
11192 let stream_active = stream_on && stream_graph.is_some();
11193 if debug_spec {
11194 eprintln!(
11195 "[spec] stream_on={stream_on} env={} samp={sampled} dg={} captured={} active={stream_active} session={session_mode} replay={spec_replay}",
11196 crate::spec::spec_stream(),
11197 dctx.graph.is_some(),
11198 stream_graph.is_some()
11199 );
11200 }
11201 let t_v_s = k + 1;
11202 // ROUND-STREAM buffers + ptr tables now live in the model-generic round_stream
11203 // module (extracted 2026-07-12; the gemma burst reuses them).
11204 let sb = crate::round_stream::StreamBufs::new(e, k, crate::spec::spec_stream_m())?;
11205 let crate::round_stream::StreamBufs {
11206 mut vtok_d,
11207 mut brk_d,
11208 mut pend_d,
11209 last_pred_d,
11210 mut pos_ctr,
11211 mut pos_start_d,
11212 mut ring_d,
11213 acc_d: mut stream_acc,
11214 m_rounds,
11215 k: _,
11216 } = sb;
11217 let stream_ptrs: Option<CudaSlice<u64>> = if stream_active {
11218 Some(crate::round_stream::kv_len_ptr_table(
11219 e,
11220 cache,
11221 Some(&pos_ctr),
11222 )?)
11223 } else {
11224 None
11225 };
11226
11227 let t_fill = t_ent.elapsed();
11228 let mut round = 0usize;
11229 // ADAPTIVE DRAFT LENGTH (MEMRA_SPEC_ADAPT=1, opt-in — the gemma_spec accepted-run law,
11230 // ported 2026-08-01): next round's draft depth = last round's accepted run + 1, clamped
11231 // to [floor(pos), k_cap] — a miss shrinks the next draft to the miss point + 1,
11232 // full-accept streaks re-deepen one step per round. NOT the 2026-07-07 acceptance-EMA
11233 // (that arm measured an HONEST LOSS to static per-class optima — 115.0/85.8/73.4 vs
11234 // 121.6/92.7/75.6, EMA lag — and was removed 2026-07-08; rig5090.jsonl has the record).
11235 // The gemma law has no lag class: it reacts within one round, and was worth +7-20% on
11236 // the gemma cells at unchanged exactness (2026-07-10 flip; floor sweep 2026-07-25;
11237 // position key 2026-07-26). Signal = n_acc from the round's EXISTING accept readback —
11238 // zero new syncs; the draft graph is a SINGLE-STEP capture replayed per drafted token,
11239 // so a per-round depth needs no re-capture (unlike gemma's whole-chain graphs). qwen's
11240 // in-round p-min cut already shortens chains mid-round, so gemma's one-round-late p-min
11241 // fold into kc is unnecessary here — the accepted-run law sees the cut via n_acc.
11242 // Exactness is the verify's job at ANY depth (same contract as p-min variable rounds).
11243 // DEFAULT OFF on the qwen path until its cells gate a flip (gemma's is default-on).
11244 // MEASURED 2026-08-01 (H100 GPU-3, interleaved x3, NGEN=256, same-invocation plain
11245 // denominators; research/qwen-adaptive-k-20260801/): REFUTED on the tuned qwen configs.
11246 // q27 K=3+HPOST+PMIN=0.3: short +0.8% (noise; law ~idles, len_hist identical), board
11247 // -1.9%, agentic -0.5%; board PMIN=0 -2.1% (not p-min shadowing — the law itself);
11248 // floor=1 -2.8% (gemma's floor-collapse, reproduced). q35 K=2 board: -6.4% (52/136
11249 // rounds shrink to depth 1; no depth to reclaim at K=2). The gemma direction DOES
11250 // appear at untuned depth-K — q27 K=6 floor=4 +1.5% over fixed K=6 — but stays -3.7%
11251 // below fixed K=3: same verdict class as the retired EMA arm (honest loss to static
11252 // per-class optima). Acceptance-rate rises under the law while tokens/round falls —
11253 // it buys accept-% by adding rounds, and a round's fixed draft+verify cost wins.
11254 // K=1..8 self-consistency PASS both models with the law ON (exactness held).
11255 let adapt = std::env::var("MEMRA_SPEC_ADAPT").as_deref() == Ok("1");
11256 // floor: per-model default keyed on n_embd (gemma's tiering — models with an expensive
11257 // verify keep deep drafts after a miss); MEMRA_SPEC_ADAPT_FLOOR pins it everywhere.
11258 let adapt_floor_env: Option<usize> = std::env::var("MEMRA_SPEC_ADAPT_FLOOR")
11259 .ok()
11260 .and_then(|v| v.parse().ok());
11261 let adapt_floor_default: usize = if self.cfg.n_embd as usize >= 3500 {
11262 4
11263 } else if self.cfg.n_embd as usize >= 2500 {
11264 2
11265 } else {
11266 1
11267 };
11268 let adapt_floor: usize = adapt_floor_env.unwrap_or(adapt_floor_default);
11269 // position key: past floor_ctx a HIGH floor (>=4) relaxes to 1 — forced-deep drafts
11270 // turn net-negative at depth (gemma 31B d1736 evidence); MEMRA_SPEC_FLOOR_CTX moves
11271 // the boundary, an explicit MEMRA_SPEC_ADAPT_FLOOR pins the floor everywhere.
11272 let floor_ctx: usize = std::env::var("MEMRA_SPEC_FLOOR_CTX")
11273 .ok()
11274 .and_then(|v| v.parse().ok())
11275 .unwrap_or(1024);
11276 let floor_at = |pos: usize| -> usize {
11277 if adapt_floor_env.is_some() || pos < floor_ctx {
11278 adapt_floor
11279 } else if adapt_floor >= 4 {
11280 1
11281 } else {
11282 adapt_floor
11283 }
11284 };
11285 // cap: MEMRA_SPEC_CAPMAX (gemma semantics, default 7). Binds only under adapt — the
11286 // fixed-K default path is untouched by this whole block.
11287 let cap_max: usize = std::env::var("MEMRA_SPEC_CAPMAX")
11288 .ok()
11289 .and_then(|v| v.parse().ok())
11290 .unwrap_or(7);
11291 let k_cap = k.min(cap_max).max(1);
11292 let mut kc = k_cap;
11293 let mut opti_fork: Option<OptiForkState> = None;
11294 let mut fork_snapshot: Option<crate::cache::CacheSnapshot> = None;
11295 if fork_mode != OptiForkGateMode::Disabled {
11296 let fence = crate::pp::pp_cuts(self.layers.len());
11297 let refusal = if !session_mode {
11298 Some("not-session")
11299 } else if k != 1 || adapt {
11300 Some("requires-fixed-k1")
11301 } else if sampled || constraint.is_some() || spec_replay {
11302 Some("sampled-constrained-or-replay")
11303 } else if pipe.is_some() {
11304 Some("two-session-pipeline")
11305 } else if !spec_devacc() {
11306 Some("requires-device-accept")
11307 } else if stream_active || crate::spec::spec_stream() {
11308 Some("round-stream")
11309 } else if !self.mtp_extra.is_empty() {
11310 Some("multi-head-mtp")
11311 } else if crate::cache::swa_ring_on() || cache.has_swa_ring() {
11312 Some("swa-ring")
11313 } else if crate::pp::pp_host_bounce_active() {
11314 Some("host-bounce")
11315 } else if fork_mode == OptiForkGateMode::Controller
11316 && cache.recur.iter().any(Option::is_some)
11317 {
11318 Some("controller-requires-zero-recurrent-state")
11319 } else if fence.as_ref().is_none_or(|f| f.len() != 3) {
11320 Some("requires-pp2")
11321 } else {
11322 None
11323 };
11324 if let Some(reason) = refusal {
11325 OPTI_FORK_REFUSALS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
11326 eprintln!("[opti-fork] refused reason={reason}");
11327 } else {
11328 let fence = fence.expect("validated PP-2 fence");
11329 let rt = crate::pp::PpNRt::get(e)?;
11330 let primary_stage0 = rt.engine(0, e).ctx().ordinal() == e.ctx().ordinal();
11331 let primary_stage1 = rt.engine(1, e).ctx().ordinal() == e.ctx().ordinal();
11332 let primary_supported =
11333 primary_stage0 || (fork_mode == OptiForkGateMode::Controller && primary_stage1);
11334 if !rt.cross_device() || !primary_supported {
11335 OPTI_FORK_REFUSALS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
11336 eprintln!("[opti-fork] refused reason=requires-supported-primary-cross-device");
11337 } else {
11338 // Both recurrent snapshots and both seed generations are allocated before
11339 // the first fork, each through its owning PP stage. Allocation failure
11340 // therefore happens before any optimistic state mutation can occur.
11341 let current_snapshot = opti_snapshot_stage_owned(e, cache, rt, &fence)?;
11342 let alternate_snapshot = opti_snapshot_stage_owned(e, cache, rt, &fence)?;
11343 let fork = OptiForkState::new(
11344 e,
11345 cache,
11346 fork_mode,
11347 alternate_snapshot,
11348 &h_seed_buf,
11349 &fill_prev,
11350 rt,
11351 fence[1],
11352 self.layers.len(),
11353 )?;
11354 eprintln!(
11355 "[opti-fork] armed mode={fork_mode:?} snapshots=2 seeds=2 split={} \
11356 payload_dev0={} payload_dev1={} q_threshold={:.3}",
11357 fence[1],
11358 fork.logical_payload_bytes[0],
11359 fork.logical_payload_bytes[1],
11360 fork.controller.map_or(0.0, |policy| policy.threshold),
11361 );
11362 fork_snapshot = Some(current_snapshot);
11363 opti_fork = Some(fork);
11364 }
11365 }
11366 }
11367 // Persistent snapshot buffers are allocated once and refreshed in place. The fork arm
11368 // uses stage-owned snapshots; refused/disabled arms retain the existing generic helper.
11369 let mut snap = match fork_snapshot {
11370 Some(snapshot) => snapshot,
11371 None => cache.snapshot(e)?,
11372 };
11373 let mut carried_opti: Option<OptiControllerTicket> = None;
11374 // ROUND-STREAM stage (b) 3a: device table of per-layer kvl.len_d pointers (stable — the
11375 // cache never reallocates len_d; see cache.rs "stable pointer" note). 0 = no KV layer.
11376 let kv_len_ptrs: Option<CudaSlice<u64>> = if spec_devacc() && !spec_replay {
11377 Some(crate::round_stream::kv_len_ptr_table(e, cache, None)?)
11378 } else {
11379 None
11380 };
11381 // BONUS FOLD (2026-07-04): after a FULL accept the bonus token is NOT committed with a
11382 // separate T=1 trunk pass (a full weight read per round). It stays PENDING and rides as
11383 // column 0 of the NEXT round's verify batch. Under predecessor pairing the next chain
11384 // seeds from the bonus's predecessor's TRUE verify hidden (free — no extra
11385 // pass of any kind). Verify still
11386 // checks every emitted token against the target -> exactness holds by construction; only
11387 // DRAFT QUALITY can shift, which the acceptance numbers arbitrate.
11388 // bonus emitted but not yet committed to cache. A carried pending (see SpecSession::
11389 // pending_tok) enters round 0 directly — the burst boundary becomes a plain round edge.
11390 let mut pending: Option<u32> = carried_pending;
11391 // MEMRA_SPEC_PHASE=1: per-round wall decomposition (draft / verify / accept+commit) —
11392 // no tracing, no extra syncs (each phase is naturally sync-bounded: draft readbacks,
11393 // the verify accept readback). Printed once at loop end via spec-stats.
11394 let anatomy_on = std::env::var("MEMRA_SPEC_PP_ANATOMY").as_deref() == Ok("1");
11395 let phase_on = anatomy_on || std::env::var("MEMRA_SPEC_PHASE").as_deref() == Ok("1");
11396 // DRAFT-MASK receipt (lane/draft-mask): speculative-clone wall + rounds, printed with
11397 // spec-stats. The clone is the one cost the design adds per round — measured, not assumed.
11398 let (mut dm_clone_ns, mut dm_rounds) = (0u128, 0usize);
11399 // grammar-truncation counters: how many rounds the verify-side cut fired and how many
11400 // already-verified tokens it threw away. THIS is the quantity draft masking targets.
11401 let (mut dm_cuts, mut dm_cut_tokens) = (0usize, 0usize);
11402 let (mut ph_draft, mut ph_verify, mut ph_rest) = (0f64, 0f64, 0f64);
11403 let mut ph_wait = 0f64;
11404 let mut ph_commit = 0f64;
11405 let mut ph_t = std::time::Instant::now();
11406 let mut ph_mark = |acc: &mut f64, on: bool| {
11407 if on {
11408 let now = std::time::Instant::now();
11409 *acc += (now - ph_t).as_secs_f64();
11410 ph_t = now;
11411 }
11412 };
11413 // MTP-ROUTE VERIFY GRAPHS (`MEMRA_SPEC_VERIFY_GRAPH`, see the flag doc): the
11414 // model-owned capture pool, locked for the whole burst exactly as the dspark serve
11415 // arm holds it — the slab stash is live verify -> commit inside a round, and the
11416 // worker drives rounds from one scheduler thread. PERSISTENT across generations on
11417 // the model (rebuilding per call re-captures the pool per prompt, which is the
11418 // measured way to lose more than the launches cost); the captured bodies are
11419 // cache-independent, every state read going through per-round refreshed pointer
11420 // tables. None = the eager walk, byte-identical.
11421 //
11422 // Never armed together with ROUND-STREAM: the tparallel verify refuses that pair
11423 // loudly, and `stream_active` owns the burst arm above, so the door stays shut
11424 // whenever the stream is live rather than relying on that refusal.
11425 // The lock is taken ONLY when the door is armed: with the flag off this whole block
11426 // is inert, so the default path cannot serialize two spec generations behind a mutex
11427 // it never reads.
11428 let vg_armed =
11429 crate::spec::spec_verify_graph_env().unwrap_or_else(|| self.vgraph_family_default());
11430 let mut vg_guard = if vg_armed && !stream_active {
11431 let mut g = self.dspark_vgraphs.lock().unwrap();
11432 if g.is_none() {
11433 // Size by the WIDEST verify this run can present, which is k+1 and NOT
11434 // k_cap+1: the sampled arm's own window is `t_v_s = k + 1`, so a pool built
11435 // from a smaller adaptive cap gets sliced past its stash rows (a `slice_mut`
11436 // panic in the sampled ON arm, measured before this line said k+1).
11437 let vt_cap = (k.max(k_cap) + 1).max(2);
11438 *g = DsparkVerifyGraphs::new(e, cache, vt_cap, n_embd)?;
11439 if g.is_some() {
11440 // Engagement receipt (the dead-arm lesson): prove the door is LIVE rather
11441 // than trusting that a flag set means a pool built.
11442 eprintln!("[spec-vg] MTP verify-graph pool ENGAGED (vt_cap={vt_cap})");
11443 } else {
11444 eprintln!(
11445 "[spec-vg] MTP verify-graph pool declined (no linear layers, \
11446 non-uniform state, or vt_cap < 2) — eager walk"
11447 );
11448 }
11449 }
11450 Some(g)
11451 } else {
11452 None
11453 };
11454 // Capacity fail-safe: a round wider than the pool was built for must take the eager
11455 // walk, not slice the stash past its rows. The sizing above already covers every
11456 // round this run can present; this keeps a future caller (or a k that grows behind
11457 // the pool's back) on the byte-identical fallback instead of a panic.
11458 let vg_t_cap = vg_guard
11459 .as_ref()
11460 .and_then(|g| g.as_ref())
11461 .map(|g| g.t_capacity())
11462 .unwrap_or(0);
11463 if let Some(p) = pipe {
11464 p.setup_end();
11465 }
11466 while keep_going && out.len() < max_new {
11467 // ROUND-STREAM BURST: from round 1 (pending guaranteed by every non-replay arm),
11468 // issue M rounds with zero readbacks, then drain the ring + reconcile mirrors.
11469 if let (true, Some(sg), Some(ptrs)) = (
11470 stream_active && round >= 1 && pending.is_some(),
11471 &stream_graph,
11472 &stream_ptrs,
11473 ) {
11474 if debug_spec {
11475 static ONCE: std::sync::Once = std::sync::Once::new();
11476 ONCE.call_once(|| {
11477 eprintln!("[memra] ROUND-STREAM burst engaged (M={m_rounds} k={k})")
11478 });
11479 }
11480 e.set_i32_one(&mut pos_ctr, cache.pos as i32)?;
11481 e.set_u32_one(&mut pend_d, pending.unwrap())?;
11482 e.set_u32_one(&mut ring_d, 0)?; // ring count = 0 (writes element 0)
11483 for _mi in 0..m_rounds {
11484 e.i32_copy_add(&pos_ctr, &mut pos_start_d, 0)?;
11485 cache.snapshot_into(e, &mut snap)?; // device D2Ds, stream-ordered
11486 e.i32_copy_add(&pos_ctr, &mut scratch.kv.len_d, 0)?; // draft-KV rollback
11487 e.i32_copy_add(&pos_ctr, &mut dctx.g_pos, 1)?; // rope pos = pos + base
11488 e.u32_copy(&pend_d, &mut dctx.g_tok)?;
11489 e.copy_into(&mut dctx.g_seed, 0, &h_seed_buf, n_embd)?;
11490 sg.launch()?;
11491 e.spec_assemble_verify(
11492 &g_tokp2k,
11493 &pend_d,
11494 d2t_dev.as_ref(),
11495 &mut vtok_d,
11496 &mut brk_d,
11497 p_min,
11498 k,
11499 pmin0,
11500 )?;
11501 let mut ck = VerifyCkpt::new(self.layers.len());
11502 let dummy = vec![0u32; t_v_s];
11503 let (tl_d, vx) = self.decode_step_t_core_stream(
11504 e,
11505 &dummy,
11506 0,
11507 &mut *cache,
11508 embd_dev,
11509 Some(&mut ck),
11510 Some((&vtok_d, &pos_ctr)),
11511 None,
11512 None,
11513 None,
11514 )?;
11515 for j in 0..t_v_s {
11516 e.argmax_token_device_col(&tl_d, j, n_vocab, &mut preds_d, j)?;
11517 }
11518 e.spec_accept_greedy_dc(
11519 &preds_d,
11520 &vtok_d,
11521 &last_pred_d,
11522 &brk_d,
11523 &mut stream_acc,
11524 )?;
11525 e.spec_seed_gather(&vx, &fill_prev, &stream_acc, &mut h_seed_buf, 1, n_embd)?;
11526 e.copy_into(&mut fill_prev, 0, &h_seed_buf, n_embd)?;
11527 self.commit_verified_prefix_stream(
11528 e,
11529 &mut *cache,
11530 &snap,
11531 &ck,
11532 &stream_acc,
11533 1,
11534 t_v_s,
11535 )?;
11536 e.spec_rollback_stream(
11537 ptrs,
11538 &pos_start_d,
11539 &stream_acc,
11540 1,
11541 self.layers.len() + 1,
11542 )?;
11543 e.spec_ring_commit(&vtok_d, &stream_acc, &brk_d, &mut ring_d, &mut pend_d)?;
11544 }
11545 e.stream().synchronize()?;
11546 let ring_h = e.dtoh_u32(&ring_d)?;
11547 let cnt = ring_h[0] as usize;
11548 for i in 0..cnt {
11549 if out.len() < max_new {
11550 out.push(ring_h[1 + i]);
11551 }
11552 }
11553 let pos_h = e.dtoh_i32(&pos_ctr)?[0] as usize;
11554 for il in 0..self.layers.len() {
11555 if let Some(kvl) = cache.kv[il].as_mut() {
11556 kvl.len = pos_h;
11557 }
11558 }
11559 cache.pos = pos_h;
11560 scratch.kv.len = pos_h;
11561 pending = Some(ring_h[cnt]); // last drained token = the live bonus
11562 last_token = ring_h[cnt];
11563 total_drafted += k * m_rounds; // upper bound (p-min breaks uncounted)
11564 total_accepted += cnt.saturating_sub(m_rounds);
11565 if let Some(t) = sess_telem {
11566 // totals only — the burst's per-round accept counts stayed on device
11567 // (that is the point of the round-stream arm). pos_* untouched.
11568 t.record_totals(m_rounds, k * m_rounds, cnt.saturating_sub(m_rounds));
11569 }
11570 round += m_rounds;
11571 // sse-cadence: the drained ring is committed — flush it at burst-drain cadence.
11572 keep_going = flush_commit(&mut on_commit, &out, &mut flushed);
11573 continue;
11574 }
11575 let pipe_draft = match pipe {
11576 Some(p) => Some(p.draft_begin(round)?),
11577 None => None,
11578 };
11579 let pos = cache.pos; // #tokens committed (EXCLUDES a pending bonus)
11580 let mut current_opti = carried_opti.take();
11581 let mut fork_generation = if current_opti.is_none() && pending.is_some() {
11582 match opti_fork.as_mut() {
11583 Some(fork) if fork.mode.is_forced() => Some(fork.reserve(&mut snap)?),
11584 None => None,
11585 Some(_) => None,
11586 }
11587 } else {
11588 None
11589 };
11590 if current_opti.is_none() {
11591 if let Some(fork) = opti_fork.as_ref() {
11592 opti_snapshot_stage_owned_into(e, cache, fork.rt, &fork.fence, &mut snap)?;
11593 } else {
11594 cache.snapshot_into(e, &mut snap)?;
11595 }
11596 } else if snap.pos != pos {
11597 return Err(format!(
11598 "optipipe carried snapshot pos {} != current pos {pos}",
11599 snap.pos
11600 )
11601 .into());
11602 } // §C: snapshot BEFORE draft+verify (already retained for a carried successor)
11603 ph_mark(&mut ph_rest, phase_on);
11604
11605 // --- 1. DRAFT k tokens with the NextN head (autoregressive, T=1 each) ---
11606 // p-min semantics (both paths): stop the chain early when the head's confidence in
11607 // its own pick drops below p_min — the just-drafted token is DISCARDED, but its
11608 // scratch append stands (identical to the eager chain's ordering). j==0 always drafts.
11609 let base0 = if pending.is_some() { 1usize } else { 0usize };
11610 // fixed draft length by default; MEMRA_SPEC_ADAPT=1 drafts at last round's
11611 // accepted run + 1 (the gemma law — see the setup block above the loop).
11612 let k_this = if adapt { kc } else { k };
11613 let mut draft: Vec<u32> = Vec::with_capacity(k);
11614 let mut draft_idx: Vec<u32> = Vec::with_capacity(k); // trimmed-vocab ids (== draft when untrimmed)
11615 let mut controller_draft_prob: Option<f32> = None;
11616 let mut controller_eager_state: Option<(u32, CudaSlice<f32>)> = None;
11617 if let Some(ticket) = current_opti.as_mut() {
11618 let carried_pending = pending.ok_or("optipipe carried successor lost pending")?;
11619 if ticket.verify_tokens[0] != carried_pending {
11620 return Err(format!(
11621 "optipipe carried pending mismatch: ticket={} live={carried_pending}",
11622 ticket.verify_tokens[0],
11623 )
11624 .into());
11625 }
11626 draft.push(ticket.verify_tokens[1]);
11627 controller_draft_prob = Some(ticket.draft_prob);
11628 controller_eager_state = ticket
11629 .take_eager_seed()
11630 .map(|seed| (ticket.verify_tokens[1], seed));
11631 } else {
11632 // Round-start draft-KV sync (BOTH paths). Persistent: truncate/align to the committed
11633 // history — slots 0..P hold entries for the tokens before last_token@P (P = pos +
11634 // base0 - 1); this single set_len IS the draft-side rollback (drops last round's
11635 // rejected drafts and p-min extras via the len mechanism).
11636 scratch.set_len(e, pos + base0 - 1)?;
11637 if pen_on {
11638 // PEN_WINDOW_MAX also bounds the per-round upload and the O(n_hist^2)
11639 // device dedup: `penalty_last_n` is usize::MAX for any serve request with
11640 // a penalty, so without the cap this grew with the whole session.
11641 let win = sp.penalty_last_n.min(PEN_WINDOW_MAX);
11642 let w0 = pen_hist.len().saturating_sub(win);
11643 pen_hist_d = Some(e.htod_u32_v(&pen_hist[w0..])?);
11644 }
11645 if sampled {
11646 draft_logits.clear();
11647 draft_stats.clear();
11648 }
11649 // DRAFT-SIDE GRAMMAR MASK: clone the committed grammar state ONCE per round; each
11650 // position's mask is computed on that clone and advanced by the PROPOSED token. The
11651 // real state moves only on emission (verify's job), so the emitted stream is
11652 // unchanged — the mask only removes tokens the verify would have truncated anyway.
11653 let mut dmask_live = dmask_on;
11654 if dmask_live {
11655 let t_c = std::time::Instant::now();
11656 constraint
11657 .as_deref_mut()
11658 .unwrap()
11659 .draft_begin()
11660 .map_err(|e2| format!("constraint: {e2}"))?;
11661 dm_clone_ns += t_c.elapsed().as_nanos();
11662 dm_rounds += 1;
11663 }
11664 if let (false, Some(gr)) = (sampled || pen_on, &dctx.graph) {
11665 // GRAPH DRAFT: one dispatch per drafted token. The chain feeds itself on-device
11666 // (in-graph argmax -> tok_d -> next replay's embed; h_nextn -> h_seed_d; pos_d
11667 // inc'd in-graph); the host only reads 4B token (+4B p) and decides the break.
11668 e.set_i32_one(&mut dctx.g_pos, (pos + base0) as i32)?;
11669 e.set_u32_one(&mut dctx.g_tok, last_token)?;
11670 e.copy_into(&mut dctx.g_seed, 0, &h_seed_buf, n_embd)?;
11671 for j in 0..k_this {
11672 // per-position mask upload (contents only — the graph's baked pointer is
11673 // dctx.g_dmask). All-ones once masking goes dead mid-chain, so the captured
11674 // mask node degrades to a no-op ban instead of needing a second graph.
11675 if dmask_live
11676 && !upload_draft_mask(
11677 e,
11678 constraint.as_deref_mut().unwrap(),
11679 &mut dctx.g_dmask,
11680 mtp.d2t.as_ref(),
11681 d_vocab,
11682 dmask_words,
11683 )?
11684 {
11685 // no draft-vocab row is grammar-legal here (a trimmed FR-Spec head can
11686 // genuinely miss the legal set): neutralize the captured mask node and
11687 // finish the chain UNMASKED — exactly pre-lane behaviour, never worse.
11688 e.htod_u32_into(&mut dctx.g_dmask, &vec![u32::MAX; dmask_words])?;
11689 dmask_live = false;
11690 }
11691 gr.launch()?;
11692 scratch.kv.len += 1; // host mirror (len_d advanced in-graph)
11693 let idx = e.dtoh_u32_one(&dctx.g_tok)?;
11694 // #87 SENTINEL TRAP: an all-NaN head-logits row leaves the device argmax's
11695 // init sentinel (0x7FFFFFFF) in g_tok — feeding it onward dereferences
11696 // embed_row(sentinel) = table + ~4.6TB (never mapped) inside the NEXT graph
11697 // replay's embed node, and the MMU fault kills the CUDA context for the
11698 // whole process (research/pp2spec-crash-20260807: 3 coredumps, byte-exact
11699 // VA arithmetic). Refuse loudly instead; the diagnostics name the first-NaN
11700 // buffer (g_seed = the verify-side handoff vs head-side compute).
11701 if (idx as usize) >= d_vocab {
11702 // g_seed is SELF-FED (the replay writes h_nextn back into it), so it
11703 // reads as the head's OUTPUT at j; h_seed_buf is the round's INPUT
11704 // seed, untouched since the round-start copy — the pair discriminates
11705 // "seed arrived poisoned" from "head forward produced NaN".
11706 let seed_h = e.dtoh(&dctx.g_seed)?;
11707 let seed_nan = seed_h.iter().filter(|v| v.is_nan()).count();
11708 let in_h = e.dtoh(&h_seed_buf)?;
11709 let in_nan = in_h.iter().filter(|v| v.is_nan()).count();
11710 return Err(format!(
11711 "draft(graph) argmax sentinel 0x{idx:08x} >= d_vocab {d_vocab} at \
11712 round {round} j={j} pos={pos}: head-out NaN {seed_nan}/{n_embd}, \
11713 round-input-seed NaN {in_nan}/{n_embd} — refusing to dereference \
11714 the embed row (#87 trap)"
11715 )
11716 .into());
11717 }
11718 // trimmed draft vocab -> target token id (identity when no d2t map)
11719 let d = match &mtp.d2t {
11720 Some(map) => map[idx as usize],
11721 None => idx,
11722 };
11723 let draft_p = if p_min > 0.0
11724 || opti_fork
11725 .as_ref()
11726 .is_some_and(|fork| fork.controller.is_some())
11727 {
11728 Some(e.dtoh(&dctx.g_p)?[0])
11729 } else {
11730 None
11731 };
11732 if j == 0 {
11733 controller_draft_prob = draft_p;
11734 }
11735 if let Some(p) = draft_p.filter(|_| p_min > 0.0) {
11736 if p < p_min && (j > 0 || (pmin0 && base0 == 1)) {
11737 break;
11738 }
11739 }
11740 draft.push(d);
11741 // with a trimmed head the NEXT embed must read the TARGET id, not the draft
11742 // index the argmax wrote — patch the persistent token buffer (4B htod).
11743 if d != idx {
11744 e.set_u32_one(&mut dctx.g_tok, d)?;
11745 }
11746 // advance the SPECULATIVE state with the proposal; a dead chain drops to
11747 // unmasked drafting for the remaining positions (verify still arbitrates).
11748 // speculative advance; a chain the grammar can no longer follow (EOS
11749 // proposed) ends here. The captured mask node always runs, so a dead chain
11750 // leaves the buffer NEUTRAL (all-ones = ban nothing) before it exits.
11751 if dmask_live
11752 && !constraint
11753 .as_deref_mut()
11754 .unwrap()
11755 .draft_advance(d)
11756 .map_err(|e2| format!("constraint: {e2}"))?
11757 {
11758 e.htod_u32_into(&mut dctx.g_dmask, &vec![u32::MAX; dmask_words])?;
11759 break;
11760 }
11761 }
11762 // PURE-TEMP RE-TEST (lane/graph-s-key-exactness-20260819): the sampled graph is
11763 // legal ONLY in the regime it was captured in. The condition used to read
11764 // `(sampled, &dctx.graph_s)` and trusted `s_key` to have dropped anything else —
11765 // which it could not, because the key omitted the filters. Both halves are now
11766 // enforced: the key drops a stale graph, and this site refuses to launch one.
11767 } else if let (true, Some(gr)) = (sampled && pure_temp, &dctx.graph_s) {
11768 if skey_probe() {
11769 eprintln!(
11770 "[skey] chain=graph_s round={round} pure_temp={} top_k={} \
11771 top_p={} min_p={} s_key_parked={:?}",
11772 pure_temp as u8, sp.top_k, sp.top_p, sp.min_p, dctx.s_key,
11773 );
11774 }
11775 // SAMPLED GRAPH DRAFT: one replay per drafted token — head forward + gumbel +
11776 // argmax in ONE dispatch; the host reads 4B token (+4B p), D2Ds q into slot j,
11777 // and decides the break. Event-counter continuity: g_ctr is host-seeded to
11778 // sctr-1 ONCE per round (outside the graph); the in-graph bump runs BEFORE the
11779 // perturb, so replay j consumes counter sctr+j — exactly the eager arm's Philox
11780 // stream. Host sctr advances in lockstep (computed, no readback needed).
11781 e.set_i32_one(&mut dctx.g_pos, (pos + base0) as i32)?;
11782 e.set_u32_one(&mut dctx.g_tok, last_token)?;
11783 e.copy_into(&mut dctx.g_seed, 0, &h_seed_buf, n_embd)?;
11784 e.set_u32_one(&mut dctx.g_ctr, sctr.wrapping_sub(1))?;
11785 for j in 0..k_this {
11786 gr.launch()?;
11787 scratch.kv.len += 1; // host mirror (len_d advanced in-graph)
11788 sctr += 1; // mirrors the in-graph g_ctr bump (eager parity:
11789 // counts the p-min-discarded token too)
11790 // q retention: ONE async D2D of the persistent head-logits buffer into this
11791 // round's slot j (stream-ordered after the replay, before the next one).
11792 e.copy_into(&mut dctx.q_slots[j], 0, &dctx.g_q, d_vocab)?;
11793 let idx = e.dtoh_u32_one(&dctx.g_tok)?;
11794 // #87 SENTINEL TRAP (see the greedy graph arm above).
11795 if (idx as usize) >= d_vocab {
11796 let seed_h = e.dtoh(&dctx.g_seed)?;
11797 let seed_nan = seed_h.iter().filter(|v| v.is_nan()).count();
11798 return Err(format!(
11799 "draft(graph-sampled) argmax sentinel 0x{idx:08x} >= d_vocab \
11800 {d_vocab} at round {round} j={j} pos={pos}: round-seed NaN \
11801 {seed_nan}/{n_embd} — refusing to dereference the embed row \
11802 (#87 trap)"
11803 )
11804 .into());
11805 }
11806 let d = match &mtp.d2t {
11807 Some(map) => map[idx as usize],
11808 None => idx,
11809 };
11810 draft_idx.push(idx);
11811 if p_min > 0.0 {
11812 let p = e.dtoh(&dctx.g_p)?[0];
11813 if p < p_min && (j > 0 || (pmin0 && base0 == 1)) {
11814 break;
11815 }
11816 }
11817 draft.push(d);
11818 // trimmed head: the NEXT embed must read the TARGET id (see the greedy arm).
11819 if d != idx {
11820 e.set_u32_one(&mut dctx.g_tok, d)?;
11821 }
11822 }
11823 // uniform accept path: fill draft_stats per used slot (pure-temp regime — the
11824 // stats degenerate to th=0 / full-Z; one filter_stats launch per slot, tiny).
11825 for j in 0..draft.len().max(draft_idx.len()) {
11826 let rows0 = e.htod_i32(&[0])?;
11827 let (mut th_d, mut z_d, mut mx_d) = (e.zeros(1)?, e.zeros(1)?, e.zeros(1)?);
11828 e.filter_stats(
11829 &dctx.q_slots[j],
11830 d_vocab,
11831 &rows0,
11832 &mut th_d,
11833 &mut z_d,
11834 &mut mx_d,
11835 d_vocab,
11836 1,
11837 sp_temp,
11838 sp.top_k,
11839 sp.top_p,
11840 sp.min_p,
11841 )?;
11842 draft_stats.push((e.dtoh(&mx_d)?[0], e.dtoh(&th_d)?[0], e.dtoh(&z_d)?[0]));
11843 }
11844 } else {
11845 if skey_probe() && sampled {
11846 eprintln!(
11847 "[skey] chain=eager round={round} pure_temp={} top_k={} \
11848 top_p={} min_p={} s_key_parked={:?}",
11849 pure_temp as u8, sp.top_k, sp.top_p, sp.min_p, dctx.s_key,
11850 );
11851 }
11852 // EAGER DRAFT (fallback: MoE head/trunk, huge k, MEMRA_SPEC_NOGRAPH, capture fail).
11853 let chain_heads = !self.mtp_extra.is_empty();
11854 let mut e_tok = last_token;
11855 let mut d_seed = e.clone_dtod(&h_seed_buf)?;
11856 let mut chain_tokens = if chain_heads {
11857 vec![last_token]
11858 } else {
11859 Vec::new()
11860 };
11861 let mut chain_seeds = if chain_heads {
11862 vec![e.clone_dtod(&h_seed_buf)?]
11863 } else {
11864 Vec::new()
11865 };
11866 for j in 0..k_this {
11867 // GPU-ARGMAX DRAFT (2026-07-03): device logits + device argmax + 4-byte token
11868 // read instead of the ~600KB full-vocab dtoh + host argmax per draft token.
11869 let mtp_pos = pos + base0 + j;
11870 // draft-side grammar mask (eager twin of the graph arm's in-graph node).
11871 // A position with no legal draft-vocab row drops to unmasked drafting for
11872 // the rest of the chain (pre-lane behaviour; verify still arbitrates).
11873 if dmask_live {
11874 dmask_live = upload_draft_mask(
11875 e,
11876 constraint.as_deref_mut().unwrap(),
11877 &mut dctx.g_dmask,
11878 mtp.d2t.as_ref(),
11879 d_vocab,
11880 dmask_words,
11881 )?;
11882 }
11883 let mask = if dmask_live {
11884 Some((&dctx.g_dmask, dmask_words))
11885 } else {
11886 None
11887 };
11888 let (dl_d, h_nextn) = if chain_heads {
11889 if debug_spec {
11890 eprintln!(
11891 "[mtp-chain-step] round={round} j={j} head={} replay_rows={}",
11892 mtp_chain_head_index(j, self.mtp_head_count()),
11893 chain_tokens.len(),
11894 );
11895 }
11896 self.mtp_chain_forward_dev(
11897 e,
11898 &chain_tokens,
11899 &chain_seeds,
11900 &mut *scratch,
11901 pos + base0 - 1,
11902 embd_dev,
11903 mask,
11904 )?
11905 } else {
11906 self.mtp_head_forward_dev(
11907 e,
11908 mtp,
11909 e_tok,
11910 &d_seed,
11911 &mut *scratch,
11912 mtp_pos,
11913 embd_dev,
11914 mask,
11915 )?
11916 };
11917 let tok_d = if sampled {
11918 // FILTERED Gumbel-max: stats -> masked perturb -> argmax = one draw from
11919 // the filtered softmax (filters off => th=0, exact v1 semantics).
11920 if perturb_buf.is_none() {
11921 perturb_buf = Some(e.zeros(d_vocab.max(n_vocab))?);
11922 }
11923 let mut q_row = e.clone_dtod(&dl_d)?; // retained q (penalized when on)
11924 if pen_on {
11925 let h = pen_hist_d.as_ref().unwrap();
11926 let nh = h.len();
11927 e.penalize_logits(
11928 &mut q_row,
11929 h,
11930 nh,
11931 sp.penalty_repeat,
11932 sp.penalty_freq,
11933 sp.penalty_present,
11934 d_vocab,
11935 )?;
11936 }
11937 let rows0 = e.htod_i32(&[0])?;
11938 let (mut th_d, mut z_d, mut mx_d) =
11939 (e.zeros(1)?, e.zeros(1)?, e.zeros(1)?);
11940 e.filter_stats(
11941 &q_row, d_vocab, &rows0, &mut th_d, &mut z_d, &mut mx_d, d_vocab,
11942 1, sp_temp, sp.top_k, sp.top_p, sp.min_p,
11943 )?;
11944 let (th, z, mx) =
11945 (e.dtoh(&th_d)?[0], e.dtoh(&z_d)?[0], e.dtoh(&mx_d)?[0]);
11946 let pb = perturb_buf.as_mut().unwrap();
11947 e.gumbel_perturb_filtered(
11948 &q_row, pb, d_vocab, sp_seed, sctr, sp_temp, mx, th,
11949 )?;
11950 sctr += 1;
11951 draft_logits.push(q_row);
11952 draft_stats.push((mx, th, z));
11953 e.argmax_token_device(pb, d_vocab)?
11954 } else {
11955 e.argmax_token_device(&dl_d, d_vocab)?
11956 };
11957 let idx = e.dtoh_u32_one(&tok_d)?;
11958 // #87 SENTINEL TRAP (eager twin — see the graph arm). Extra diagnostics
11959 // here because the eager chain's operands are all readable: dl_d (the head
11960 // logits row) and d_seed (this step's h_seed) name the first-NaN buffer.
11961 if (idx as usize) >= d_vocab {
11962 let dl_h = e.dtoh(&dl_d)?;
11963 let dl_nan = dl_h.iter().filter(|v| v.is_nan()).count();
11964 let seed_h = if chain_heads {
11965 e.dtoh(chain_seeds.last().unwrap())?
11966 } else {
11967 e.dtoh(&d_seed)?
11968 };
11969 let seed_nan = seed_h.iter().filter(|v| v.is_nan()).count();
11970 return Err(format!(
11971 "draft(eager) argmax sentinel 0x{idx:08x} >= d_vocab {d_vocab} at \
11972 round {round} j={j} pos={pos}: head-logits NaN {dl_nan}/{d_vocab}, \
11973 step-seed NaN {seed_nan}/{n_embd} — refusing to dereference the \
11974 embed row (#87 trap)"
11975 )
11976 .into());
11977 }
11978 let d = match &mtp.d2t {
11979 Some(map) => map[idx as usize],
11980 None => idx,
11981 };
11982 if sampled {
11983 draft_idx.push(idx);
11984 }
11985 let draft_p = if p_min > 0.0
11986 || opti_fork
11987 .as_ref()
11988 .is_some_and(|fork| fork.controller.is_some())
11989 {
11990 let p_d = e.prob_of_token_device(&dl_d, &tok_d, d_vocab)?;
11991 Some(e.dtoh(&p_d)?[0])
11992 } else {
11993 None
11994 };
11995 if j == 0 {
11996 controller_draft_prob = draft_p;
11997 }
11998 if let Some(p) = draft_p.filter(|_| p_min > 0.0) {
11999 if p < p_min && (j > 0 || (pmin0 && base0 == 1)) {
12000 break;
12001 }
12002 }
12003 draft.push(d);
12004 if chain_heads {
12005 chain_tokens.push(d);
12006 chain_seeds.push(h_nextn);
12007 } else {
12008 e_tok = d;
12009 d_seed = h_nextn;
12010 }
12011 // speculative advance; a chain the grammar can no longer follow (EOS
12012 // proposed) ends here — the prefix already proposed still rides verify.
12013 if dmask_live
12014 && !constraint
12015 .as_deref_mut()
12016 .unwrap()
12017 .draft_advance(d)
12018 .map_err(|e2| format!("constraint: {e2}"))?
12019 {
12020 break;
12021 }
12022 }
12023 if !chain_heads
12024 && opti_fork
12025 .as_ref()
12026 .is_some_and(|fork| fork.controller.is_some())
12027 {
12028 controller_eager_state = Some((e_tok, d_seed));
12029 }
12030 }
12031 }
12032 let k_round = draft.len();
12033 if let Some(p) = pipe {
12034 p.draft_end(round);
12035 }
12036 drop(pipe_draft);
12037
12038 ph_mark(&mut ph_draft, phase_on);
12039 // --- 2. VERIFY: one batched target forward. With a pending bonus, it rides as col 0
12040 // (committing its KV/recur inside the SAME weight read); drafts follow. ---
12041 let verify_tokens: Vec<u32> = match pending {
12042 Some(b) => {
12043 let mut v = Vec::with_capacity(k_round + 1);
12044 v.push(b);
12045 v.extend_from_slice(&draft);
12046 v
12047 }
12048 None => draft.clone(),
12049 };
12050 let base = if pending.is_some() { 1 } else { 0 };
12051 // ckpt (REPLAY-FREE partial accept): retain per-layer state-rebuild inputs alongside
12052 // the verify. Pure buffer keep-alives + dtod clones — kernel work is unchanged.
12053 let mut ckpt = if let Some(ticket) = current_opti.as_mut() {
12054 Some(ticket.take_ckpt())
12055 } else if spec_replay {
12056 None
12057 } else {
12058 Some(VerifyCkpt::new(self.layers.len()))
12059 };
12060 let controller_can_probe = base == 1
12061 && k_round == 1
12062 && out.len().saturating_add(2) < max_new
12063 && controller_draft_prob.is_some()
12064 && opti_fork
12065 .as_ref()
12066 .and_then(|fork| fork.controller.as_ref())
12067 .is_some_and(|policy| !policy.breaker_tripped);
12068 let mut successor_attempt: Option<OptiControllerTicket> = None;
12069 let mut rejected_probe: Option<(f32, u32)> = None;
12070 let mut controller_prepared: Option<OptiControllerPrepared> = None;
12071 if controller_can_probe {
12072 // Prepare d2/q and, on admission, d3 before either current verify half is
12073 // issued. N stage 0 can then be followed immediately by N+1 stage 0; once N's
12074 // boundary fires, those dev0 launches overlap N stage 1 on dev1. Preparing on
12075 // the primary stream after N stage 1 would serialize the supposed pipeline.
12076 let eager_pos = scratch.kv.len + 1;
12077 let (optimistic_pending, pending_probability) = self.opti_controller_draft_step(
12078 e,
12079 mtp,
12080 &mut dctx,
12081 &mut *scratch,
12082 d_vocab,
12083 &mut controller_eager_state,
12084 eager_pos,
12085 embd_dev,
12086 )?;
12087 let first_probability = controller_draft_prob
12088 .ok_or("optipipe controller probe lost first-token probability")?;
12089 let q_proxy = first_probability * pending_probability;
12090 OPTI_GATE_CHECKS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
12091 OPTI_SHADOW_DRAFT_TOKENS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
12092 let admitted = opti_fork
12093 .as_ref()
12094 .and_then(|fork| fork.controller.as_ref())
12095 .ok_or("optipipe controller policy disappeared")?
12096 .admit(q_proxy);
12097 if admitted {
12098 OPTI_GATE_ADMITS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
12099 let eager_pos = scratch.kv.len + 1;
12100 let (optimistic_draft, optimistic_draft_probability) = self
12101 .opti_controller_draft_step(
12102 e,
12103 mtp,
12104 &mut dctx,
12105 &mut *scratch,
12106 d_vocab,
12107 &mut controller_eager_state,
12108 eager_pos,
12109 embd_dev,
12110 )?;
12111 OPTI_SHADOW_DRAFT_TOKENS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
12112 let eager_seed = controller_eager_state.take().map(|(token, seed)| {
12113 debug_assert_eq!(token, optimistic_draft);
12114 seed
12115 });
12116 controller_prepared = Some(OptiControllerPrepared {
12117 verify_tokens: [optimistic_pending, optimistic_draft],
12118 draft_prob: optimistic_draft_probability,
12119 eager_seed,
12120 q_proxy,
12121 scratch_len: scratch.kv.len,
12122 });
12123 } else {
12124 OPTI_GATE_REJECTS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
12125 OPTI_WASTED_DRAFT_TOKENS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
12126 rejected_probe = Some((q_proxy, optimistic_pending));
12127 eprintln!(
12128 "[opti-controller] reject q={q_proxy:.6} threshold={:.3}",
12129 opti_fork
12130 .as_ref()
12131 .and_then(|fork| fork.controller.as_ref())
12132 .expect("controller policy")
12133 .threshold,
12134 );
12135 }
12136 }
12137 let fork_attempt = match fork_generation.take() {
12138 Some(generation) if base == 1 && k_round == 1 => Some(generation),
12139 Some(generation) => {
12140 opti_fork
12141 .as_mut()
12142 .expect("fork generation without fork state")
12143 .retire(generation)?;
12144 None
12145 }
12146 None => None,
12147 };
12148 let (tlogits_d, vx) = if let Some(p) = pipe {
12149 self.decode_step_t_core_pipelined(
12150 e,
12151 &verify_tokens,
12152 pos,
12153 &mut *cache,
12154 embd_dev,
12155 ckpt.as_mut(),
12156 p,
12157 round,
12158 )?
12159 } else if controller_can_probe {
12160 let fence = opti_fork
12161 .as_ref()
12162 .ok_or("optipipe controller probe lost fork state")?
12163 .fence;
12164 let boundary = match current_opti.as_mut() {
12165 Some(ticket) => ticket.take_boundary(),
12166 None => self.verify_stage0_issue(
12167 e,
12168 &verify_tokens,
12169 pos,
12170 &mut *cache,
12171 embd_dev,
12172 ckpt.as_mut(),
12173 None,
12174 &fence,
12175 Some(true),
12176 None,
12177 )?,
12178 };
12179 if let Some(prepared) = controller_prepared.take() {
12180 let generation = {
12181 let fork = opti_fork
12182 .as_mut()
12183 .ok_or("optipipe controller admission lost fork state")?;
12184 let generation = fork.reserve_successor()?;
12185 let rt = fork.rt;
12186 let snapshot_fence = fork.fence;
12187 opti_snapshot_one_stage_owned_into(
12188 e,
12189 cache,
12190 rt,
12191 &snapshot_fence,
12192 0,
12193 fork.successor_snapshot_mut(),
12194 )?;
12195 generation
12196 };
12197 let mut successor_ckpt = VerifyCkpt::new(self.layers.len());
12198 let successor_boundary = self.verify_stage0_issue(
12199 e,
12200 &prepared.verify_tokens,
12201 pos + verify_tokens.len(),
12202 &mut *cache,
12203 embd_dev,
12204 Some(&mut successor_ckpt),
12205 None,
12206 &fence,
12207 Some(false),
12208 None,
12209 )?;
12210 OPTI_FORK_ATTEMPTS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
12211 let fork = opti_fork
12212 .as_ref()
12213 .ok_or("optipipe controller ticket lost fork state")?;
12214 successor_attempt = Some(fork.controller_ticket(
12215 generation,
12216 successor_boundary,
12217 successor_ckpt,
12218 prepared.verify_tokens,
12219 prepared.draft_prob,
12220 prepared.eager_seed,
12221 prepared.q_proxy,
12222 prepared.scratch_len,
12223 ));
12224 eprintln!(
12225 "[opti-controller] issue generation={} q={:.6} threshold={:.3} \
12226 verify={:?}",
12227 generation.id,
12228 prepared.q_proxy,
12229 fork.controller.expect("controller policy").threshold,
12230 prepared.verify_tokens,
12231 );
12232 }
12233 let result = self.verify_stage1_finish(
12234 e,
12235 boundary,
12236 &mut *cache,
12237 ckpt.as_mut(),
12238 None,
12239 &fence,
12240 successor_attempt.is_none(),
12241 )?;
12242 if let Some(ticket) = current_opti.as_mut() {
12243 ticket.settle();
12244 }
12245 if successor_attempt.is_some() {
12246 let fork = opti_fork
12247 .as_mut()
12248 .ok_or("optipipe successor snapshot lost fork state")?;
12249 let rt = fork.rt;
12250 let snapshot_fence = fork.fence;
12251 opti_snapshot_one_stage_owned_into(
12252 e,
12253 cache,
12254 rt,
12255 &snapshot_fence,
12256 1,
12257 fork.successor_snapshot_mut(),
12258 )?;
12259 // Publish N only after both independent successor-state queues are complete.
12260 fork.rt.publish_to(1, &e.stream())?;
12261 }
12262 result
12263 } else if let Some(ticket) = current_opti.as_mut() {
12264 let fork = opti_fork
12265 .as_mut()
12266 .ok_or("optipipe carried controller ticket lost fork state")?;
12267 let boundary = ticket.take_boundary();
12268 let result = self.verify_stage1_finish(
12269 e,
12270 boundary,
12271 &mut *cache,
12272 ckpt.as_mut(),
12273 None,
12274 &fork.fence,
12275 true,
12276 )?;
12277 ticket.settle();
12278 result
12279 } else if let Some(generation) = fork_attempt {
12280 let fork = opti_fork
12281 .as_mut()
12282 .expect("fork generation without fork state");
12283 fork.capture_seed(e, generation, &h_seed_buf, &fill_prev, scratch.kv.len)?;
12284 let action = fork.mode.action(generation.id);
12285 let boundary = self.verify_stage0_issue(
12286 e,
12287 &verify_tokens,
12288 pos,
12289 &mut *cache,
12290 embd_dev,
12291 ckpt.as_mut(),
12292 None,
12293 &fork.fence,
12294 Some(true),
12295 None,
12296 )?;
12297 OPTI_FORK_ATTEMPTS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
12298 let mut ticket = fork.ticket(generation, boundary);
12299 if action == OptiForkAction::Abort {
12300 return Err(format!(
12301 "optipipe forced abort with generation {} stage0 in flight",
12302 generation.id,
12303 )
12304 .into());
12305 }
12306 fork.reconcile(
12307 e,
12308 &mut *cache,
12309 &mut *scratch,
12310 &snap,
12311 &mut h_seed_buf,
12312 &mut fill_prev,
12313 generation,
12314 action,
12315 verify_tokens[0],
12316 )?;
12317 let result = if action == OptiForkAction::Hit {
12318 let boundary = ticket.take_boundary();
12319 self.verify_stage1_finish(
12320 e,
12321 boundary,
12322 &mut *cache,
12323 ckpt.as_mut(),
12324 None,
12325 &fork.fence,
12326 true,
12327 )?
12328 } else {
12329 // The optimistic boundary slot has no reader. Re-run the unchanged serial
12330 // verify only after E_restart published the restored stage-0 state.
12331 self.decode_step_t_core(
12332 e,
12333 &verify_tokens,
12334 pos,
12335 &mut *cache,
12336 embd_dev,
12337 ckpt.as_mut(),
12338 )?
12339 };
12340 ticket.settle();
12341 debug_assert_eq!(ticket.generation, generation);
12342 fork.retire(generation)?;
12343 result
12344 } else {
12345 // The serial verify every non-fork round takes — the MTP route's
12346 // verify-graph door. The pool is None unless MEMRA_SPEC_VERIFY_GRAPH armed
12347 // a pool above, and then the walk replays the captured trunk instead of
12348 // re-issuing it launch by launch.
12349 let vg_round = if verify_tokens.len() <= vg_t_cap {
12350 vg_guard.as_mut().and_then(|g| g.as_mut())
12351 } else {
12352 if let Some(g) = vg_guard.as_mut().and_then(|g| g.as_mut()) {
12353 // The commit reads this flag to pick its arm; a round that declines
12354 // the pool must not inherit a stale `true` from the round before it.
12355 g.round_slab = false;
12356 }
12357 None
12358 };
12359 self.decode_step_t_core_vg(
12360 e,
12361 &verify_tokens,
12362 pos,
12363 &mut *cache,
12364 embd_dev,
12365 ckpt.as_mut(),
12366 vg_round,
12367 )?
12368 };
12369 let pipe_accept = match pipe {
12370 Some(p) => Some(p.accept_begin(round)?),
12371 None => None,
12372 };
12373
12374 ph_mark(&mut ph_verify, phase_on);
12375 // --- 3. GREEDY ACCEPT (walk prefix, stop at first mismatch) ---
12376 // DEVICE-ARGMAX ACCEPT: argmax every verify column ON DEVICE (same 2-pass kernels +
12377 // smallest-index tie-break as host argmax, argmax_gate-validated) and read back ONE
12378 // [T] u32 — replaces the T x n_vocab f32 dtoh + T host argmaxes per round.
12379 // t_pred[j] = target's greedy prediction for the slot after draft[j-1] (j>=1) or after
12380 // last_token (j==0). With a pending bonus, col 0 IS the prediction after last_token
12381 // (== the bonus), so every index shifts by `base` and last_pred is unused.
12382 let t_v = verify_tokens.len();
12383 let mut preds: Vec<u32> = Vec::new();
12384 if !sampled {
12385 for j in 0..t_v {
12386 e.argmax_token_device_col(&tlogits_d, j, n_vocab, &mut preds_d, j)?;
12387 }
12388 preds = e.dtoh_u32(&preds_d)?; // <- the verify-GPU wait lands here
12389 // #87 SENTINEL TRAP, verify side: a sentinel pred becomes the round's bonus =
12390 // next round's last_token = the next chain's embed lookup. Catch it at the
12391 // source with the column named — an all-NaN VERIFY column implicates the
12392 // stage-split trunk (decode_step_t_core_ppn), not the draft head.
12393 if let Some(bad) = preds[..t_v].iter().position(|&p| (p as usize) >= n_vocab) {
12394 let col = &tlogits_d.slice(bad * n_vocab..(bad + 1) * n_vocab);
12395 let mut probe = e.zeros(n_vocab)?;
12396 e.copy_view_into(&mut probe, 0, col, n_vocab)?;
12397 let col_h = e.dtoh(&probe)?;
12398 let col_nan = col_h.iter().filter(|v| v.is_nan()).count();
12399 return Err(format!(
12400 "verify argmax sentinel 0x{:08x} >= n_vocab {n_vocab} at round {round} \
12401 col {bad}/{t_v} pos={pos}: verify-logits col NaN {col_nan}/{n_vocab} \
12402 — the stage-split verify produced a poisoned column (#87 trap)",
12403 preds[bad]
12404 )
12405 .into());
12406 }
12407 }
12408 ph_mark(&mut ph_wait, phase_on);
12409 let t_pred = |j: usize| -> u32 {
12410 if j == 0 && base == 0 {
12411 last_pred
12412 } else {
12413 // GREEDY-ONLY: `preds` is filled under `if !sampled` above. The debug print
12414 // used to call this from the sampled arm and panicked the worker; it now goes
12415 // through `debug_t_pred0`. Keep the strict index here — in the greedy walk an
12416 // out-of-range pred is a real bug, not something to paper over.
12417 debug_assert!(
12418 !sampled,
12419 "t_pred is greedy-only: `preds` is empty in the sampled arm"
12420 );
12421 preds[base + j - 1]
12422 }
12423 };
12424 let mut devacc_seeded = false;
12425 let mut devacc_acc: Option<CudaSlice<u32>> = None;
12426 let (n_acc, bonus) = if !sampled {
12427 // ROUND-STREAM stage (a) (MEMRA_SPEC_DEVACC=1 opt-in): the walk runs ON DEVICE
12428 // (spec_accept_greedy, verbatim rule) and the host reads back 8B (n_acc, bonus)
12429 // instead of the [T] preds. Same sync count — machinery for stages (b)/(c),
12430 // gated on token identity vs the host walk (the arms below are bit-equal rules).
12431 if crate::spec::spec_devacc() && k_round > 0 && !spec_replay && constraint.is_none()
12432 {
12433 let draft_d = e.htod_u32_v(&draft)?;
12434 let mut acc_out = e.alloc_u32_zeroed(2)?;
12435 e.spec_accept_greedy(
12436 &preds_d,
12437 &draft_d,
12438 last_pred,
12439 base,
12440 k_round,
12441 &mut acc_out,
12442 )?;
12443 devacc_acc = Some(acc_out.clone());
12444 // stage (b): next-round seed gathered ON DEVICE from acc_out before the host
12445 // ever reads n_acc (j=base+n_acc -> vx col j-1; j==0 -> fill_prev). The three
12446 // non-replay commit arms skip their host-offset seed copies (guarded below);
12447 // the legacy spec_replay arm keeps its own rx-based seeding (excluded here).
12448 // NOTE: fill_prev is NOT updated here — the commit arms' TRUE-HIDDEN
12449 // REFRESH reads the OLD fill_prev (predecessor of this round's verify batch);
12450 // the update lands after the arms (devacc_seeded guard below).
12451 e.spec_seed_gather(&vx, &fill_prev, &acc_out, &mut h_seed_buf, base, n_embd)?;
12452 // 3a: KV lens roll back on device (len = saved + base + n_acc, all arms'
12453 // unified rule; full accept rewrites the verify-left value). Host mirrors
12454 // update after the readback; commit_verified_prefix skips its len_d writes.
12455 if let Some(successor) = successor_attempt.as_ref() {
12456 opti_fork
12457 .as_mut()
12458 .ok_or("optipipe successor reconcile lost fork state")?
12459 .queue_actual_reconcile(
12460 e,
12461 &snap,
12462 &acc_out,
12463 successor.verify_tokens[0],
12464 base,
12465 )?;
12466 } else if let Some(ptrs) = &kv_len_ptrs {
12467 let saved: Vec<i32> = (0..self.layers.len())
12468 .map(|il| snap.kv_len[il].map(|v| v as i32).unwrap_or(0))
12469 .collect();
12470 let saved_d = e.htod_i32(&saved)?;
12471 e.spec_rollback_kv(ptrs, &saved_d, &acc_out, base, self.layers.len())?;
12472 }
12473 devacc_seeded = true;
12474 let ab = e.dtoh_u32(&acc_out)?;
12475 (ab[0] as usize, ab[1])
12476 } else {
12477 let mut n_acc = 0usize;
12478 for j in 0..k_round {
12479 if t_pred(j) == draft[j] {
12480 n_acc += 1;
12481 } else {
12482 break;
12483 }
12484 }
12485 // bonus = target's own token at the first non-accepted slot. n_acc in 0..=k; t_pred
12486 // is defined for j in 0..=k (j==0 -> last_logits, j>=1 -> col j-1, last col = k-1).
12487 (n_acc, t_pred(n_acc))
12488 }
12489 } else {
12490 // --- SAMPLED ACCEPT (rejection sampling): u_j < p_j(x_j)/q_j(x_j) walk ---
12491 if col_buf.is_none() {
12492 col_buf = Some(e.zeros(n_vocab)?);
12493 }
12494 // FILTERED p_j: per-verify-col stats (one batched filter_stats call), then the
12495 // filtered gather. j==0&&base==0 reads last_col (its own stats row appended).
12496 let mut pj = vec![0f32; k_round.max(1)];
12497 let mut col_stats: Vec<(f32, f32, f32)> = Vec::new(); // (max, th, z) per verify col used
12498 if k_round > 0 {
12499 let mut ids: Vec<u32> = Vec::new();
12500 let mut rows: Vec<i32> = Vec::new();
12501 for j in 0..k_round {
12502 if j > 0 || base == 1 {
12503 ids.push(draft[j]);
12504 rows.push((base + j) as i32 - 1);
12505 }
12506 }
12507 if !ids.is_empty() {
12508 let nr = rows.len();
12509 // penalties: materialize the used columns into one contiguous penalized
12510 // buffer (rows remapped 0..nr) so stats+gathers see the penalized p.
12511 // penalties: materialize used columns contiguously, penalize all rows in
12512 // one launch, and point stats+gathers at the penalized buffer (rows 0..nr).
12513 let p_rows: Vec<i32> = if pen_on {
12514 (0..nr as i32).collect()
12515 } else {
12516 rows.clone()
12517 };
12518 if pen_on {
12519 if pcol_buf.as_ref().map(|b| b.len()).unwrap_or(0) < nr * n_vocab {
12520 pcol_buf = Some(e.zeros(nr * n_vocab)?);
12521 }
12522 let pc = pcol_buf.as_mut().unwrap();
12523 for (i2, &r) in rows.iter().enumerate() {
12524 let c = r as usize;
12525 e.copy_view_into(
12526 pc,
12527 i2 * n_vocab,
12528 &tlogits_d.slice(c * n_vocab..(c + 1) * n_vocab),
12529 n_vocab,
12530 )?;
12531 }
12532 let h = pen_hist_d.as_ref().unwrap();
12533 let nh = h.len();
12534 e.penalize_logits_rows(
12535 pc,
12536 h,
12537 nh,
12538 sp.penalty_repeat,
12539 sp.penalty_freq,
12540 sp.penalty_present,
12541 n_vocab,
12542 nr,
12543 )?;
12544 }
12545 let p_src: &CudaSlice<f32> = if pen_on {
12546 pcol_buf.as_ref().unwrap()
12547 } else {
12548 &tlogits_d
12549 };
12550 let rowsd = e.htod_i32(&p_rows)?;
12551 let (mut th_d, mut z_d, mut mx_d) =
12552 (e.zeros(nr)?, e.zeros(nr)?, e.zeros(nr)?);
12553 e.filter_stats(
12554 p_src, n_vocab, &rowsd, &mut th_d, &mut z_d, &mut mx_d, n_vocab, nr,
12555 sp_temp, sp.top_k, sp.top_p, sp.min_p,
12556 )?;
12557 let idsd = e.htod_u32_v(&ids)?;
12558 let mut outd = e.zeros(nr)?;
12559 e.softmax_gather_filtered(
12560 p_src, n_vocab, &idsd, &rowsd, &th_d, &z_d, &mut outd, n_vocab, nr,
12561 sp_temp,
12562 )?;
12563 let outv = e.dtoh(&outd)?;
12564 let (thv, zv, mxv) = (e.dtoh(&th_d)?, e.dtoh(&z_d)?, e.dtoh(&mx_d)?);
12565 let mut oi = 0usize;
12566 for j in 0..k_round {
12567 if j > 0 || base == 1 {
12568 pj[j] = outv[oi];
12569 oi += 1;
12570 }
12571 }
12572 col_stats = (0..nr).map(|i| (mxv[i], thv[i], zv[i])).collect();
12573 }
12574 if base == 0 {
12575 let lc: &CudaSlice<f32> = if pen_on {
12576 if col_buf.is_none() {
12577 col_buf = Some(e.zeros(n_vocab)?);
12578 }
12579 let cb = col_buf.as_mut().unwrap();
12580 e.copy_into(
12581 cb,
12582 0,
12583 last_col_logits
12584 .as_ref()
12585 .expect("sampled: last_col_logits unset"),
12586 n_vocab,
12587 )?;
12588 let h = pen_hist_d.as_ref().unwrap();
12589 let nh = h.len();
12590 e.penalize_logits(
12591 cb,
12592 h,
12593 nh,
12594 sp.penalty_repeat,
12595 sp.penalty_freq,
12596 sp.penalty_present,
12597 n_vocab,
12598 )?;
12599 col_buf.as_ref().unwrap()
12600 } else {
12601 last_col_logits
12602 .as_ref()
12603 .expect("sampled: last_col_logits unset")
12604 };
12605 let rows0 = e.htod_i32(&[0])?;
12606 let (mut th_d, mut z_d, mut mx_d) = (e.zeros(1)?, e.zeros(1)?, e.zeros(1)?);
12607 e.filter_stats(
12608 lc, n_vocab, &rows0, &mut th_d, &mut z_d, &mut mx_d, n_vocab, 1,
12609 sp_temp, sp.top_k, sp.top_p, sp.min_p,
12610 )?;
12611 let idsd = e.htod_u32_v(&[draft[0]])?;
12612 let mut outd = e.zeros(1)?;
12613 e.softmax_gather_filtered(
12614 lc, n_vocab, &idsd, &rows0, &th_d, &z_d, &mut outd, n_vocab, 1, sp_temp,
12615 )?;
12616 pj[0] = e.dtoh(&outd)?[0];
12617 last_col_stats =
12618 Some((e.dtoh(&mx_d)?[0], e.dtoh(&th_d)?[0], e.dtoh(&z_d)?[0]));
12619 }
12620 }
12621 // q source: the graph arm retained the head logits in the persistent q_slots;
12622 // the eager arm in per-round draft_logits clones. Same raw-logit values either way.
12623 // FILTERED q_j: stats from draft_stats (eager pushes in-chain; the graph arm
12624 // computes them post-replay — graph engages only filter/penalty-free, so the
12625 // stats degenerate to th=0/full-Z there, keeping ONE accept path).
12626 let q_bufs: &[CudaSlice<f32>] = if dctx.graph_s.is_some() {
12627 &dctx.q_slots
12628 } else {
12629 &draft_logits
12630 };
12631 let mut n_acc = 0usize;
12632 for j in 0..k_round {
12633 let (qmx, qth, qz) = draft_stats[j];
12634 let idsd = e.htod_u32_v(&[draft_idx[j]])?;
12635 let rowsd = e.htod_i32(&[0])?;
12636 let thd = e.htod(&[qth])?;
12637 let zd = e.htod(&[qz])?;
12638 let _ = qmx;
12639 let mut outd = e.zeros(1)?;
12640 e.softmax_gather_filtered(
12641 &q_bufs[j], d_vocab, &idsd, &rowsd, &thd, &zd, &mut outd, d_vocab, 1,
12642 sp_temp,
12643 )?;
12644 let qj = e.dtoh(&outd)?[0];
12645 let u = host_u01(sp_seed, uctr);
12646 uctr += 1;
12647 let accept = (u as f64) * (qj as f64) < pj[j] as f64;
12648 // SKEY PROBE: q == 0 for the token the draft actually proposed is the
12649 // exactness signature (see `skey_probe`). Impossible when the draft was
12650 // drawn from the same filtered distribution the verify reconstructs here;
12651 // `u * 0 < p` makes it an UNCONDITIONAL accept whenever p > 0.
12652 if skey_probe() && qj == 0.0 {
12653 eprintln!(
12654 "[skey] EXACTNESS q=0 round={round} j={j} draft_tok={} \
12655 draft_idx={} p={:e} u={u} accepted={} th_z={:?}",
12656 draft[j], draft_idx[j], pj[j], accept as u8, draft_stats[j],
12657 );
12658 }
12659 if accept {
12660 n_acc += 1;
12661 } else {
12662 break;
12663 }
12664 }
12665 let bonus = if n_acc == k_round {
12666 // FULL ACCEPT: bonus ~ FILTERED softmax at the last verify column.
12667 let col = base + k_round - 1;
12668 let cb = col_buf.as_mut().unwrap();
12669 e.copy_view_into(
12670 cb,
12671 0,
12672 &tlogits_d.slice(col * n_vocab..(col + 1) * n_vocab),
12673 n_vocab,
12674 )?;
12675 if pen_on {
12676 let h = pen_hist_d.as_ref().unwrap();
12677 let nh = h.len();
12678 e.penalize_logits(
12679 cb,
12680 h,
12681 nh,
12682 sp.penalty_repeat,
12683 sp.penalty_freq,
12684 sp.penalty_present,
12685 n_vocab,
12686 )?;
12687 }
12688 if perturb_buf.is_none() {
12689 perturb_buf = Some(e.zeros(d_vocab.max(n_vocab))?);
12690 }
12691 // STATS MUST COME FROM THIS COLUMN (bug fix 2026-08-05, lane/sampler-
12692 // truncation-fix; receipts research/sampfix-20260805/). The old code reused
12693 // `col_stats.last()` here, which is ALWAYS the wrong row: the gathered set
12694 // covers verify columns 0..=(base+k_round-2) (rows pushed as base+j-1), while
12695 // the full-accept bonus samples column base+k_round-1 — exactly ONE PAST the
12696 // last gathered column, in both base arms. `th` is a threshold in e-units of
12697 // its OWN row's max, so feeding a neighbour's (row_max, th) into
12698 // gumbel_perturb_filtered mis-scales every e0 = exp((x-row_max)/T). When the
12699 // donor column's peak is higher by more than T*ln(1/th), EVERY id fails
12700 // `e0 >= th`, the whole perturbed row becomes -3.4e38, and the 2-pass argmax
12701 // falls through to its smallest-index tie-break => token id 0 ("!") spliced
12702 // mid-word. Fragility is ordered by how large th is: min_p pins th = min_p
12703 // (0.05 => trigger at delta > 2.4 at T=0.8, fires constantly), top_p's
12704 // mass-boundary th is smaller, top_k's k-th-largest th smaller still — which
12705 // is why the head-to-head matrix saw min_p and top_p corrupt while top_k-only
12706 // stayed clean. The pure-temp default regime is immune (th == 0 masks nothing,
12707 // and row_max is unused once nothing is masked), so this fix is a byte-level
12708 // no-op for the untruncated serve default. One extra one-block filter_stats
12709 // per full-accept round is the whole cost.
12710 let (mx, th) = {
12711 let rows0 = e.htod_i32(&[0])?;
12712 let (mut th_d, mut z_d, mut mx_d) = (e.zeros(1)?, e.zeros(1)?, e.zeros(1)?);
12713 let cb0 = col_buf.as_ref().unwrap();
12714 e.filter_stats(
12715 cb0, n_vocab, &rows0, &mut th_d, &mut z_d, &mut mx_d, n_vocab, 1,
12716 sp_temp, sp.top_k, sp.top_p, sp.min_p,
12717 )?;
12718 (e.dtoh(&mx_d)?[0], e.dtoh(&th_d)?[0])
12719 };
12720 let pb = perturb_buf.as_mut().unwrap();
12721 let cb2 = col_buf.as_ref().unwrap();
12722 e.gumbel_perturb_filtered(cb2, pb, n_vocab, sp_seed, sctr, sp_temp, mx, th)?;
12723 sctr += 1;
12724 let td = e.argmax_token_device(pb, n_vocab)?;
12725 e.dtoh_u32_one(&td)?
12726 } else {
12727 // REJECT at n_acc: bonus ~ norm(max(0, softmax_T(p) - softmax_T(q))).
12728 let cb = col_buf.as_mut().unwrap();
12729 if n_acc > 0 || base == 1 {
12730 let col = base + n_acc - 1;
12731 e.copy_view_into(
12732 cb,
12733 0,
12734 &tlogits_d.slice(col * n_vocab..(col + 1) * n_vocab),
12735 n_vocab,
12736 )?;
12737 } else {
12738 let lc = last_col_logits.as_ref().unwrap();
12739 e.copy_into(cb, 0, lc, n_vocab)?;
12740 }
12741 if pen_on {
12742 let h = pen_hist_d.as_ref().unwrap();
12743 let nh = h.len();
12744 e.penalize_logits(
12745 cb,
12746 h,
12747 nh,
12748 sp.penalty_repeat,
12749 sp.penalty_freq,
12750 sp.penalty_present,
12751 n_vocab,
12752 )?;
12753 }
12754 let cb2 = col_buf.as_ref().unwrap();
12755 let sc = sctr;
12756 sctr += 1;
12757 // p-stats for the reject column: from col_stats when the col was gathered,
12758 // else (j==0&&base==0) from last_col_stats.
12759 let p_stats = if n_acc > 0 || base == 1 {
12760 // col index within the gathered set == number of gathered cols before n_acc
12761 let gi = if base == 1 { n_acc } else { n_acc - 1 };
12762 col_stats.get(gi).copied().unwrap_or_else(|| {
12763 (0.0, 0.0, 1.0) // unreachable: gathered cols always cover the reject slot
12764 })
12765 } else {
12766 last_col_stats.expect("sampled: last_col_stats unset at reject")
12767 };
12768 let q_stats = draft_stats[n_acc];
12769 if let Some(map) = &d2t_dev {
12770 if q_full_buf.is_none() {
12771 q_full_buf = Some(e.zeros(n_vocab)?);
12772 }
12773 let qf = q_full_buf.as_mut().unwrap();
12774 e.scatter_trim_logits(&q_bufs[n_acc], map, qf, d_vocab, n_vocab)?;
12775 let qf2 = q_full_buf.as_ref().unwrap();
12776 e.residual_sample_filtered(
12777 cb2,
12778 Some(qf2),
12779 n_vocab,
12780 sp_temp,
12781 sp_seed,
12782 sc,
12783 p_stats,
12784 q_stats,
12785 &mut sample_tok,
12786 )?;
12787 } else {
12788 e.residual_sample_filtered(
12789 cb2,
12790 Some(&q_bufs[n_acc]),
12791 n_vocab,
12792 sp_temp,
12793 sp_seed,
12794 sc,
12795 p_stats,
12796 q_stats,
12797 &mut sample_tok,
12798 )?;
12799 }
12800 e.dtoh_u32(&sample_tok)?[0]
12801 };
12802 (n_acc, bonus)
12803 };
12804 // --- 3b. GRAMMAR TRUNCATION (constrained spec, 2026-08-03): the grammar is
12805 // an extra rejection rule AFTER the exactness verify (the batched-verify-twins
12806 // ordering). Walk the accepted drafts through the grammar in commit order; the
12807 // first illegal token truncates acceptance at its slot, and that slot's emission
12808 // is recomputed as the MASKED argmax of the target's own verify column — token-
12809 // identical to constrained plain greedy decode (an unmasked argmax that is
12810 // grammar-legal IS the masked argmax: masking only removes competitors). The
12811 // column D2H (~1MB) is paid only when a cut fires — the tight-grammar cost,
12812 // measured in acceptance numbers, never hidden.
12813 let (n_acc, bonus) = match constraint.as_deref_mut() {
12814 None => (n_acc, bonus),
12815 Some(c) => {
12816 fn ce(e2: String) -> Box<dyn std::error::Error> {
12817 format!("constraint: {e2}").into()
12818 }
12819 let mut na = n_acc;
12820 let mut cut = false;
12821 for (j, &d) in draft.iter().enumerate().take(n_acc) {
12822 if c.is_allowed(d).map_err(ce)? {
12823 c.consume(d).map_err(ce)?;
12824 } else {
12825 na = j;
12826 cut = true;
12827 dm_cut_tokens += n_acc - j;
12828 break;
12829 }
12830 }
12831 if cut {
12832 dm_cuts += 1;
12833 }
12834 let mut bo = bonus;
12835 if cut || !c.is_allowed(bo).map_err(ce)? {
12836 let mut row = if na == 0 && base == 0 {
12837 init_logits_host
12838 .clone()
12839 .ok_or("constraint: init logits missing (round-0 cut)")?
12840 } else {
12841 e.dtoh_view(
12842 &tlogits_d.slice((base + na - 1) * n_vocab..(base + na) * n_vocab),
12843 )?
12844 };
12845 c.mask_logits(&mut row).map_err(ce)?;
12846 bo = argmax(&row) as u32;
12847 }
12848 c.consume(bo).map_err(ce)?;
12849 (na, bo)
12850 }
12851 };
12852 let mut successor_valid = false;
12853 if let Some((q_proxy, expected_d2)) = rejected_probe {
12854 let v_n = n_acc == 1 && bonus == expected_d2;
12855 eprintln!(
12856 "[opti-controller] shadow q={q_proxy:.6} admitted=false v_n={v_n} \
12857 expected_d2={expected_d2} n_acc={n_acc} bonus={bonus}",
12858 );
12859 }
12860 if let Some(successor) = successor_attempt.as_ref() {
12861 successor_valid = n_acc == 1 && bonus == successor.verify_tokens[0];
12862 let generation = successor.generation;
12863 let q_proxy = successor.q_proxy;
12864 let expected_pending = successor.verify_tokens[0];
12865 let resolution_ms = successor.issued_at.elapsed().as_secs_f64() * 1e3;
12866 let fork = opti_fork
12867 .as_mut()
12868 .ok_or("optipipe successor resolution lost fork state")?;
12869 fork.finish_actual_reconcile(e, &mut *cache, &snap, n_acc, base, successor_valid)?;
12870 if successor_valid {
12871 OPTI_FORK_HITS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
12872 } else {
12873 OPTI_FORK_MISSES.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
12874 OPTI_RECONCILES.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
12875 OPTI_WASTED_DRAFT_TOKENS.fetch_add(2, std::sync::atomic::Ordering::Relaxed);
12876 }
12877 let breaker_tripped = fork
12878 .controller
12879 .as_mut()
12880 .expect("controller policy")
12881 .resolve(successor_valid);
12882 if breaker_tripped {
12883 OPTI_BREAKER_TRIPS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
12884 }
12885 eprintln!(
12886 "[opti-controller] resolve generation={} hit={} q={q_proxy:.6} \
12887 expected_pending={expected_pending} n_acc={n_acc} bonus={bonus} \
12888 resolution_ms={resolution_ms:.3} reconcile={} breaker={}",
12889 generation.id, successor_valid, !successor_valid, breaker_tripped,
12890 );
12891 if !successor_valid {
12892 let mut successor = successor_attempt
12893 .take()
12894 .expect("controller successor disappeared on miss");
12895 successor.settle();
12896 fork.retire(generation)?;
12897 }
12898 }
12899 total_drafted += k_round;
12900 total_accepted += n_acc;
12901 if let Some(t) = sess_telem {
12902 // Greedy, rejection-sampling, and grammar truncation all converge here after
12903 // the accept decision is already on host. Fixed-size relaxed atomics only.
12904 t.record_round(k_round, n_acc);
12905 }
12906 if spec_stats {
12907 st_len_hist[k_round] += 1;
12908 for j in 0..k_round {
12909 st_drafted[j] += 1;
12910 }
12911 for j in 0..n_acc {
12912 st_accepted[j] += 1;
12913 }
12914 if n_acc == k_round {
12915 st_full += 1;
12916 }
12917 }
12918
12919 if debug_spec {
12920 eprintln!(
12921 "[R{round}] pos={pos} out_len={} last_tok={last_token} draft={draft:?} n_acc={n_acc} bonus={bonus} t_pred0={}",
12922 out.len(),
12923 // NOT `t_pred(0)`: `preds` is filled only under `if !sampled` above, so on a
12924 // sampled request round >= 1 (base == 1) indexed an EMPTY vector and PANICKED
12925 // the GPU worker thread — a debug flag that killed the exact regime you would
12926 // set it to investigate. See `debug_t_pred0`.
12927 debug_t_pred0(sampled, base, last_pred, &preds)
12928 );
12929 }
12930
12931 // --- 4. COMMIT: draft[0..n_acc] then bonus (n_acc + 1 tokens) ---
12932 let commit_started = std::time::Instant::now();
12933 // SESSION MODE: every accepted column is already in the CACHE — `out` must carry all
12934 // of them (overshoot past max_new included) or `committed` under-counts the cache rows
12935 // and the next turn's continuation seeds one token off (gate-caught 2026-07-05). The
12936 // single-shot path keeps the cap (its caller truncates + drops the cache anyway).
12937 for j in 0..n_acc {
12938 if !session_mode && out.len() >= max_new {
12939 break;
12940 }
12941 out.push(draft[j]);
12942 }
12943 if pen_on {
12944 pen_hist.extend_from_slice(&draft[0..n_acc]);
12945 pen_hist.push(bonus);
12946 }
12947 let bonus_emitted = session_mode || out.len() < max_new;
12948 if bonus_emitted {
12949 out.push(bonus);
12950 }
12951 last_token = bonus;
12952
12953 // --- 5. ROLLBACK + advance (§C) ---
12954 if n_acc == k_round && !spec_replay {
12955 // FULL ACCEPT, BONUS FOLD: all verify columns (pending? + drafts) are committed in
12956 // cache; the NEW bonus stays PENDING for the next round's verify batch — NO extra
12957 // T=1 trunk pass. The next draft chain seeds from the MTP block's h_nextn at the
12958 // bonus position: one MTP-block pass (~1/33 trunk cost) replaces the trunk read.
12959 // last_pred is dead in the pending path (t_pred reads verify col 0).
12960 //
12961 // PERSISTENT DRAFT KV, full-accept fill: the chain covered last_token +
12962 // draft[0..k_round-2] as INPUTS (slots P..P'-2); draft[k_round-1] (slot P'-1) was
12963 // only ever an output, so its entry is MISSING. Fill it from vh_seed — its EXACT
12964 // trunk hidden (the last verify column). set_len first: a p-min break may have
12965 // left one extra chain append at that slot. Partial accepts need NO fill (the
12966 // chain already covered every accepted position; round-start set_len truncates).
12967 let mut vh_seed = e.zeros(n_embd)?;
12968 e.copy_view_into(
12969 &mut vh_seed,
12970 0,
12971 &vx.slice((t_v - 1) * n_embd..t_v * n_embd),
12972 n_embd,
12973 )?;
12974 if refresh {
12975 // TRUE-HIDDEN REFRESH (2026-07-03, the HANDOVER-listed acceptance lever):
12976 // overwrite ALL committed positions' scratch entries with K/V from their EXACT
12977 // verify hiddens — the reference engine's mtp_update fills from true hiddens;
12978 // the full stack (vx) is already resident from the verify. Replaces both the
12979 // chain-approximate entries AND the old last-token-only fill. Acceptance-only
12980 // (draft attention quality); exactness stays the verify's job.
12981 scratch.set_len(e, pos)?;
12982 // PREDECESSOR pairing: row i gets vx[i-1]; row 0 the carried fill_prev
12983 // (hidden of the last committed row before this verify batch).
12984 let mut vxs = e.zeros(t_v * n_embd)?;
12985 e.copy_into(&mut vxs, 0, &fill_prev, n_embd)?;
12986 if t_v > 1 {
12987 e.copy_view_into(
12988 &mut vxs,
12989 n_embd,
12990 &vx.slice(0..(t_v - 1) * n_embd),
12991 (t_v - 1) * n_embd,
12992 )?;
12993 }
12994 self.mtp_kv_fill_all(e, &verify_tokens, &vxs, pos, &mut *scratch, embd_dev)?;
12995 } else {
12996 scratch.set_len(e, pos + base + k_round - 1)?;
12997 // predecessor of the last draft = verify col t_v-2 (or fill_prev at t_v==1)
12998 let mut hp = e.zeros(n_embd)?;
12999 if t_v >= 2 {
13000 e.copy_view_into(
13001 &mut hp,
13002 0,
13003 &vx.slice((t_v - 2) * n_embd..(t_v - 1) * n_embd),
13004 n_embd,
13005 )?;
13006 } else {
13007 e.copy_into(&mut hp, 0, &fill_prev, n_embd)?;
13008 }
13009 self.mtp_kv_fill_all(
13010 e,
13011 &[draft[k_round - 1]],
13012 &hp,
13013 pos + base + k_round - 1,
13014 &mut *scratch,
13015 embd_dev,
13016 )?;
13017 }
13018 // REFERENCE SEEDING: no pseudo pass — the next chain's step 0 IS the
13019 // reference's (id_last, h_prev) draft row; it appends the bonus's scratch
13020 // entry itself. Seed = TRUE hidden of the bonus's predecessor (last verify
13021 // col). Saves one MTP-block pass per round on top of the pairing fix.
13022 if !devacc_seeded {
13023 e.copy_into(&mut h_seed_buf, 0, &vh_seed, n_embd)?;
13024 e.copy_into(&mut fill_prev, 0, &vh_seed, n_embd)?;
13025 }
13026 pending = Some(bonus);
13027 if debug_spec {
13028 eprintln!(" -> FULL ACCEPT (bonus pending, prev-h seed)");
13029 }
13030 } else if !spec_replay && base + n_acc >= 1 {
13031 // PARTIAL ACCEPT, REPLAY-FREE (2026-07-03 — the profiled #1 long-ctx spec cost):
13032 // the verify's first j = base+n_acc columns ARE the committed sequence, computed
13033 // bit-identically to eager (decode-exact contract) — so KEEP them: KV truncates to
13034 // pos+j, recurrent state rebuilds from the VerifyCkpt (same-kernel gdn prefix
13035 // re-run / pure state-clone restore), and the bonus stays PENDING exactly like the
13036 // full-accept path — the legacy duplicate trunk replay is gone. The next chain
13037 // seeds from the MTP pseudo-hidden of the bonus, whose seed = the TRUE verify
13038 // hidden of its predecessor (col j-1) — same one-hop pseudo structure as full
13039 // accept (never compounds: the next verify recomputes true hiddens for all
13040 // committed columns).
13041 let j = base + n_acc;
13042 // VERIFY-GRAPH SLAB COMMIT: when the captured trunk ran, the linear layers'
13043 // column stash was written into the graphs ctx's persistent slabs as in-graph
13044 // memcpy nodes, NOT into the per-column VerifyCkpt the cols arm reads — so the
13045 // commit must take the slab twin (same semantics, slab-addressed sources). The
13046 // ctx states which of the two this round produced via `round_slab`; trusting the
13047 // flag rather than the env keeps a round that fell back to the eager walk (a
13048 // capture that declined, a t the pool never captured) on the cols arm.
13049 let slab_commit = vg_guard
13050 .as_ref()
13051 .and_then(|g| g.as_ref())
13052 .map(|g| g.round_slab)
13053 .unwrap_or(false);
13054 if slab_commit {
13055 self.dspark_commit_prefix_slab(
13056 e,
13057 &mut *cache,
13058 &snap,
13059 vg_guard
13060 .as_ref()
13061 .and_then(|g| g.as_ref())
13062 .expect("slab_commit implies a graphs ctx"),
13063 j,
13064 )?;
13065 } else {
13066 self.commit_verified_prefix(
13067 e,
13068 &mut *cache,
13069 &snap,
13070 ckpt.as_ref().unwrap(),
13071 j,
13072 devacc_seeded,
13073 if devacc_seeded {
13074 devacc_acc.as_ref().map(|a| (a, base, t_v))
13075 } else {
13076 None
13077 },
13078 )?;
13079 }
13080 let mut seed = e.zeros(n_embd)?;
13081 e.copy_view_into(
13082 &mut seed,
13083 0,
13084 &vx.slice((j - 1) * n_embd..j * n_embd),
13085 n_embd,
13086 )?;
13087 // Draft scratch: TRUE-HIDDEN REFRESH of the committed prefix (see the full-accept
13088 // branch); without it the chain entries stand and only the tail truncates. Either
13089 // way len ends at pos+j so the pseudo append lands at the bonus's slot pos+j
13090 // (persistent mode), rope pos+j+1 (chain convention).
13091 if refresh {
13092 scratch.set_len(e, pos)?;
13093 let mut vxs = e.zeros(j * n_embd)?;
13094 e.copy_into(&mut vxs, 0, &fill_prev, n_embd)?;
13095 if j > 1 {
13096 e.copy_view_into(
13097 &mut vxs,
13098 n_embd,
13099 &vx.slice(0..(j - 1) * n_embd),
13100 (j - 1) * n_embd,
13101 )?;
13102 }
13103 self.mtp_kv_fill_all(
13104 e,
13105 &verify_tokens[0..j],
13106 &vxs,
13107 pos,
13108 &mut *scratch,
13109 embd_dev,
13110 )?;
13111 } else {
13112 scratch.set_len(e, pos + j)?;
13113 }
13114 // REFERENCE SEEDING (see the full-accept branch): seed = TRUE hidden of the
13115 // bonus's predecessor (verify col j-1); no pseudo pass.
13116 if !devacc_seeded {
13117 e.copy_into(&mut h_seed_buf, 0, &seed, n_embd)?;
13118 e.copy_into(&mut fill_prev, 0, &seed, n_embd)?;
13119 }
13120 pending = Some(bonus);
13121 if debug_spec {
13122 eprintln!(" -> PARTIAL(replay-free j={j}, bonus pending, prev-h seed)");
13123 }
13124 } else if !spec_replay {
13125 // ZERO ROUND FOLD (2026-07-10, verify-cost target #3): base+n_acc == 0 — a
13126 // pending-less round where nothing was accepted (PMIN0 zero-draft chains after a
13127 // replay/commit, or plain 0-accept rounds at round 0). The old path replayed
13128 // [bonus] through a FULL m=1 trunk+head forward (the 489us full-vocab head pass
13129 // measured at ~0.75/round on PMIN0 configs). Instead: restore the pre-round
13130 // snapshot and let the bonus ride the NEXT round's verify as col 0 — the existing
13131 // base=1 pending machinery, bit-identical by the decode-exact verify contract.
13132 // Seed: the bonus's predecessor is the last COMMITTED token, whose hidden
13133 // fill_prev already carries (same seeding as the 1-token-replay case it replaces).
13134 cache.rollback(e, &snap, 0)?;
13135 scratch.set_len(e, pos)?;
13136 e.copy_into(&mut h_seed_buf, 0, &fill_prev, n_embd)?;
13137 pending = Some(bonus);
13138 if debug_spec {
13139 eprintln!(" -> ZERO-ROUND FOLD (bonus pending, fill_prev seed)");
13140 }
13141 } else {
13142 // PARTIAL ACCEPT, LEGACY REPLAY (seam MEMRA_SPEC_REPLAY=1 — or j==0: nothing of
13143 // this round survives, only possible before the first pending exists, ~round 0):
13144 // restore EVERYTHING to the pre-round snapshot (KV truncate to pos + recur
13145 // restore), then replay the committed prefix pending? ++ draft[0..n_acc] ++
13146 // [bonus] as ONE batched T forward — single weight read, bit-identical to greedy
13147 // (the verify-all-columns path is the same math). Commits the bonus with a TRUE
13148 // trunk hidden.
13149 cache.rollback(e, &snap, 0)?; // accept_len=0: KV len = pos, recur = snapshot
13150 let mut replay: Vec<u32> = Vec::with_capacity(base + n_acc + 1);
13151 if let Some(b) = pending.take() {
13152 replay.push(b);
13153 }
13154 replay.extend_from_slice(&draft[0..n_acc]);
13155 replay.push(bonus);
13156 // Full-stack forward (decode_step_t_core = decode_step_t_h_emb_dev's body):
13157 // Predecessor pairing seeds from the PREDECESSOR row (col len-2) — the same-row path takes the
13158 // last col exactly as before (byte-identical to the old _h_emb_dev call).
13159 let (rl_d, rx) = if self.batched_serving_numeric_class() {
13160 let mut logits = Vec::with_capacity(replay.len() * n_vocab);
13161 let mut hidden = e.uninit(replay.len() * n_embd)?;
13162 for (row, &token) in replay.iter().enumerate() {
13163 let (row_logits, row_hidden) =
13164 self.spec_target_step_h(e, token, &mut *cache)?;
13165 logits.extend_from_slice(&row_logits);
13166 e.dtod_copy_into(&row_hidden, &mut hidden, row * n_embd)?;
13167 }
13168 (e.htod(&logits)?, hidden)
13169 } else {
13170 self.decode_step_t_core(e, &replay, pos, &mut *cache, embd_dev, None)?
13171 };
13172 // last_pred = argmax of the LAST column's logits (predicts the token after `bonus`)
13173 // — device argmax + one 4-byte read instead of the full-vocab column dtoh.
13174 e.argmax_token_device_col(&rl_d, replay.len() - 1, n_vocab, &mut preds_d, 0)?;
13175 last_pred = e.dtoh_u32(&preds_d)?[0];
13176 if sampled {
13177 let lr0 = replay.len();
13178 let lc = last_col_logits
13179 .as_mut()
13180 .expect("sampled: last_col_logits unset");
13181 e.copy_view_into(
13182 lc,
13183 0,
13184 &rl_d.slice((lr0 - 1) * n_vocab..lr0 * n_vocab),
13185 n_vocab,
13186 )?;
13187 }
13188 let lr = replay.len();
13189 if lr >= 2 {
13190 e.copy_view_into(
13191 &mut h_seed_buf,
13192 0,
13193 &rx.slice((lr - 2) * n_embd..(lr - 1) * n_embd),
13194 n_embd,
13195 )?;
13196 } else {
13197 // 1-token replay (round-0 miss): the bonus's predecessor is the OLD
13198 // last_token, whose own-row hidden fill_prev still holds.
13199 e.copy_into(&mut h_seed_buf, 0, &fill_prev, n_embd)?;
13200 }
13201 // the bonus is COMMITTED here — it becomes the last committed row.
13202 let mut rh_last = e.zeros(n_embd)?;
13203 e.copy_view_into(
13204 &mut rh_last,
13205 0,
13206 &rx.slice((lr - 1) * n_embd..lr * n_embd),
13207 n_embd,
13208 )?;
13209 e.copy_into(&mut fill_prev, 0, &rh_last, n_embd)?;
13210 if debug_spec {
13211 eprintln!(" -> PARTIAL(replay={replay:?}), next_pred={last_pred}");
13212 }
13213 }
13214 if devacc_seeded {
13215 // stage (b) epilogue: fill_prev takes the gathered seed AFTER the refresh fills
13216 // consumed the old value (both slots carry the same value in every non-replay arm).
13217 e.copy_into(&mut fill_prev, 0, &h_seed_buf, n_embd)?;
13218 }
13219 if successor_valid {
13220 let optimistic_scratch_len = successor_attempt
13221 .as_ref()
13222 .expect("valid controller successor disappeared")
13223 .scratch_len;
13224 // The normal current-round commit refreshed/truncated the logical scratch tail.
13225 // Its optimistic successor row was already written physically, so restoring only
13226 // the retained logical length makes that row live for the carried round.
13227 scratch.set_len(e, optimistic_scratch_len)?;
13228 }
13229 if let Some(current) = current_opti.take() {
13230 opti_fork
13231 .as_mut()
13232 .ok_or("optipipe current retirement lost fork state")?
13233 .retire(current.generation)?;
13234 }
13235 if successor_valid {
13236 let successor = successor_attempt
13237 .take()
13238 .expect("valid controller successor disappeared before promotion");
13239 let generation = successor.generation;
13240 opti_fork
13241 .as_mut()
13242 .ok_or("optipipe successor promotion lost fork state")?
13243 .promote_successor_snapshot(&mut snap, generation);
13244 carried_opti = Some(successor);
13245 }
13246 if anatomy_on {
13247 // Commit/rollback is normally asynchronous on the primary/head stream. Bound it
13248 // only for this diagnostic so it does not disappear into the following draft's
13249 // first token readback.
13250 e.stream().synchronize()?;
13251 ph_commit += commit_started.elapsed().as_secs_f64();
13252 }
13253 // adaptive-K update (host math, zero syncs): next round drafts accepted-run + 1,
13254 // clamped to [floor(pos), k_cap]. cache.pos is post-rollback here (the round's
13255 // final position — the floor's position key reads the committed depth). Burst
13256 // rounds (`continue` above) draft the captured fixed depth and skip this, exactly
13257 // like gemma's burst arm.
13258 if adapt {
13259 let fl_now = floor_at(cache.pos);
13260 kc = (n_acc + 1).clamp(fl_now.min(k_cap), k_cap);
13261 }
13262 ph_mark(&mut ph_rest, phase_on);
13263 if let Some(p) = pipe {
13264 p.accept_end(round);
13265 }
13266 drop(pipe_accept);
13267 round += 1;
13268 // sse-cadence: this round's accepted drafts + bonus are committed (out is
13269 // append-only past step 4) — flush at round cadence.
13270 keep_going = flush_commit(&mut on_commit, &out, &mut flushed);
13271 }
13272 if let Some(mut ticket) = carried_opti.take() {
13273 opti_fork
13274 .as_mut()
13275 .ok_or("optipipe tail drain lost fork state")?
13276 .cancel_controller_ticket(e, &mut *cache, &mut *scratch, &snap, &mut ticket)?;
13277 }
13278 // sse-cadence: nothing below appends to `out`; flush any remainder (defensive).
13279 // (verdict ignored — the burst is over either way; the session tail runs unchanged.)
13280 let _ = flush_commit(&mut on_commit, &out, &mut flushed);
13281
13282 if spec_stats {
13283 let per_slot: Vec<String> = (0..k)
13284 .map(|j| {
13285 if st_drafted[j] > 0 {
13286 format!(
13287 "{}/{}={:.3}",
13288 st_accepted[j],
13289 st_drafted[j],
13290 st_accepted[j] as f64 / st_drafted[j] as f64
13291 )
13292 } else {
13293 "0/0".into()
13294 }
13295 })
13296 .collect();
13297 let acc = if total_drafted > 0 {
13298 total_accepted as f64 / total_drafted as f64
13299 } else {
13300 0.0
13301 };
13302 eprintln!(
13303 "[spec-stats] rounds={round} full_accept={st_full} len_hist={st_len_hist:?} \
13304 per_slot=[{}] total={total_accepted}/{total_drafted}={acc:.3} \
13305 tok_per_round={:.3}",
13306 per_slot.join(" "),
13307 (total_accepted + round) as f64 / round.max(1) as f64
13308 );
13309 }
13310 if constraint.is_some() {
13311 eprintln!(
13312 "[draft-mask] mask_rounds={dm_rounds} clone_total={:.3}ms \
13313 clone_per_round={:.4}ms gram_cuts={dm_cuts}/{round} cut_tokens={dm_cut_tokens}",
13314 dm_clone_ns as f64 / 1e6,
13315 dm_clone_ns as f64 / 1e6 / dm_rounds.max(1) as f64
13316 );
13317 }
13318 if phase_on {
13319 let tot = ph_draft + ph_verify + ph_wait + ph_rest;
13320 eprintln!(
13321 "[spec-phase] draft={:.1}ms ({:.1}%) verify-issue={:.1}ms ({:.1}%) verify-wait={:.1}ms ({:.1}%) commit-host={:.1}ms ({:.1}%) rounds={round}",
13322 ph_draft * 1e3,
13323 ph_draft / tot * 100.0,
13324 ph_verify * 1e3,
13325 ph_verify / tot * 100.0,
13326 ph_wait * 1e3,
13327 ph_wait / tot * 100.0,
13328 ph_rest * 1e3,
13329 ph_rest / tot * 100.0
13330 );
13331 }
13332 if anatomy_on {
13333 let rounds_f = round.max(1) as f64;
13334 let other = (ph_rest - ph_commit).max(0.0);
13335 eprintln!(
13336 "[spec-anatomy] per-round draft={:.3}ms pp-verify={:.3}ms \
13337 verify-accept={:.3}ms commit-rollback={:.3}ms other={:.3}ms rounds={round}",
13338 ph_draft * 1e3 / rounds_f,
13339 ph_verify * 1e3 / rounds_f,
13340 ph_wait * 1e3 / rounds_f,
13341 ph_commit * 1e3 / rounds_f,
13342 other * 1e3 / rounds_f,
13343 );
13344 }
13345 let _pipe_tail = pipe.map(|p| p.primary());
13346 // SESSION TAIL: leave the session in the exact invariant the next turn's suffix prime
13347 // expects — every row in `committed` has trunk KV/recur state AND an exact draft-KV row.
13348 // Park the draft-graph ctx back on the session (the serve-burst fixed-cost fix): the next
13349 // burst replays instead of recapturing. Error paths (`?` above) drop it — recaptured then.
13350 if let Some(slot) = sess_draft_slot.take() {
13351 *slot = Some(dctx);
13352 }
13353 let t_rounds = t_ent.elapsed();
13354 if let Some((committed, last_h, next_pred_slot, sctr_slot, uctr_slot)) = sess_tail.take() {
13355 // NEXT BURST'S BOUNDARY TOKEN (lane/sampled-spec-quality, Item 1). Greedy stashes
13356 // the argmax `last_pred` exactly as before (byte contract). SAMPLED draws the token
13357 // HERE, where the sampler, the session Philox counters and the penalty window are
13358 // all live and the boundary logits row still exists — that is the "make the state
13359 // available" half of the fix; the consuming burst then just emits it. `sctr` is
13360 // written to the session BELOW the draws so the advance is never lost.
13361 *next_pred_slot = Some(last_pred);
13362 let sample_boundary = sampled && constraint.is_none() && spec_sampled_boundary_on();
13363 let mut stashed_pending = false;
13364 if let Some(b) = pending.take() {
13365 if !sampled {
13366 // PENDING-CARRY (2026-08-01): stash the bonus on the session instead of
13367 // committing it with a solo T=1 pass — the next empty-suffix greedy burst
13368 // consumes it as round-0 verify col 0 (a plain round edge; the old tail
13369 // commit + next burst's init feed were 11.6+11.5ms solo trunk passes per
13370 // burst on H100 q27, [spec-setup] trace). b stays in `out` (emitted) but
13371 // OUT of `committed` (cache rows == committed); the consuming call
13372 // prepends it once its verify commits the row. next_pred is unknowable
13373 // without the commit pass — None; callers gate on pending_tok too.
13374 debug_assert_eq!(out.last(), Some(&b), "pending must be the last emitted");
13375 if let Some(slot) = sess_pending_slot.take() {
13376 *slot = Some(b);
13377 }
13378 *next_pred_slot = None;
13379 // fill_prev = hidden of the last COMMITTED row (b's predecessor) — the
13380 // exact chain-seed/fill anchor the consuming burst (or a flush) needs.
13381 *last_h = Some(e.clone_dtod(&fill_prev)?);
13382 stashed_pending = true;
13383 } else {
13384 // SAMPLED tail (unchanged): commit the bonus (one T=1 pass) + draft fill —
13385 // the sampled round-0 accept needs this pass's logits (last_col_logits).
13386 let pos_b = cache.pos;
13387 scratch.set_len(e, pos_b)?;
13388 let (lg_b, hb) = self.spec_target_step_h(e, b, &mut *cache)?;
13389 // after a FULL-accept exit `last_pred` is STALE (it predicted the bonus
13390 // itself — the prediction AFTER the bonus never materialized; it would have
13391 // been the next round's verify col 0). The commit's logits ARE that
13392 // prediction — so they are also the row the next burst's boundary token
13393 // comes off, and (lane/sampled-spec-quality) it is DRAWN from them here.
13394 *next_pred_slot = Some(if sample_boundary {
13395 sample_boundary_token(
13396 e,
13397 &lg_b,
13398 &sp,
13399 &pen_hist,
13400 &mut sctr,
13401 "burst-tail-commit",
13402 )?
13403 } else {
13404 argmax(&lg_b) as u32
13405 });
13406 self.mtp_kv_fill_all(e, &[b], &fill_prev, pos_b, &mut *scratch, embd_dev)?;
13407 *last_h = Some(hb);
13408 }
13409 } else {
13410 // fill_prev tracks the hidden of the last COMMITTED row throughout the loop.
13411 *last_h = Some(e.clone_dtod(&fill_prev)?);
13412 if sample_boundary {
13413 // No pending to commit, so the boundary row is the one `last_pred` was
13414 // argmaxed from and the sampled path keeps it on device: the init feed's
13415 // logits when the burst ran zero rounds, else the legacy-replay path's
13416 // last verify column (both predict the token AFTER the last committed
13417 // row). It is retained precisely because round 0's accept test needs it,
13418 // so the draw costs no extra D2H of the [n_vocab] row.
13419 match last_col_logits.as_ref() {
13420 Some(lc) => {
13421 *next_pred_slot = Some(sample_boundary_token_dev(
13422 e,
13423 lc,
13424 n_vocab,
13425 &sp,
13426 &pen_hist,
13427 &mut sctr,
13428 "burst-tail-nopending",
13429 )?);
13430 }
13431 // NAME THE FALLBACK (house standard): unreachable today — a sampled
13432 // burst always feeds or replays, so the row exists — but if it ever
13433 // is, the stream takes a greedy token and SAYS so rather than
13434 // silently regressing to the pre-lane behaviour.
13435 None => eprintln!(
13436 "[spec-boundary] sampled tail kept the ARGMAX boundary token \
13437 (reason: no retained boundary logits row)"
13438 ),
13439 }
13440 }
13441 }
13442 *sctr_slot = sctr;
13443 *uctr_slot = uctr;
13444 committed.extend_from_slice(prompt);
13445 if let Some(cb) = carried_pending {
13446 // the consumed carry's cache row landed in round 0's verify (every pending
13447 // round commits col 0) — it joins `committed` here, in sequence order.
13448 committed.push(cb);
13449 }
13450 if stashed_pending {
13451 // ZERO-EMIT BURST (2026-08-06 c=8 serve panic, pre-existing since b4aea184):
13452 // `out.len() - 1` underflowed on an EMPTY `out` — "range end index
13453 // 18446744073709551615 out of range for slice of length 0", killing the
13454 // memra-gpu-worker and failing 31 of 32 concurrent requests with "worker closed
13455 // stream". Reachable because `pending` starts as `carried_pending` (a bonus
13456 // stashed by the PREVIOUS burst) while `out` starts empty, and the carry is
13457 // deliberately NOT pushed to `out` (line ~3239: the burst that emitted it already
13458 // did). So a burst that stashes a pending without emitting anything of its own —
13459 // the round loop exits before a push, e.g. the ring drain's `out.len() < max_new`
13460 // guard skipping every token under a tight budget — arrives here with
13461 // out.len() == 0 and stashed_pending == true.
13462 //
13463 // The invariant is unchanged: `committed` gets every emitted token EXCEPT the
13464 // stashed bonus. With nothing emitted, that is nothing — and the carry pushed
13465 // just above is already accounted. Saturating, not a min/assert: an empty `out`
13466 // here is a legitimate burst shape, not a corrupt state.
13467 let emitted = out.len().saturating_sub(1);
13468 committed.extend_from_slice(&out[..emitted]);
13469 } else {
13470 committed.extend_from_slice(&out); // FULL out incl. overshoot — all committed
13471 }
13472 debug_assert_eq!(
13473 cache.pos,
13474 committed.len(),
13475 "session invariant: cache rows == committed tokens"
13476 );
13477 if setup_trace {
13478 e.stream().synchronize()?; // bound the async tail fill in the trace
13479 let t_tail = t_ent.elapsed();
13480 eprintln!(
13481 "[spec-setup] init={:.2}ms cap={:.2}ms fill={:.2}ms rounds={:.2}ms tail={:.2}ms total={:.2}ms out={} cont={}",
13482 t_init.as_secs_f64() * 1e3,
13483 (t_cap - t_init).as_secs_f64() * 1e3,
13484 (t_fill - t_cap).as_secs_f64() * 1e3,
13485 (t_rounds - t_fill).as_secs_f64() * 1e3,
13486 (t_tail - t_rounds).as_secs_f64() * 1e3,
13487 t_tail.as_secs_f64() * 1e3,
13488 out.len(),
13489 continuation
13490 );
13491 }
13492 return Ok((out, total_drafted, total_accepted));
13493 }
13494 out.truncate(max_new);
13495 Ok((out, total_drafted, total_accepted))
13496 }
13497
13498 /// Anchor-bounded DSpark target extraction. The trunk sees the exact generated token tape;
13499 /// only requested hidden rows and target-logit rows cross PCIe. An anchor token at p pairs
13500 /// with the pre-output-norm h[p-1] carrier, exactly as the existing replay/NextN path does.
13501 pub fn extract_dspark_anchors(
13502 &self,
13503 e: &Engine,
13504 tokens: &[u32],
13505 anchor_positions: &[usize],
13506 gamma: usize,
13507 top_k: usize,
13508 chunk: usize,
13509 temperature: f32,
13510 ) -> Result<Vec<DsparkAnchorRecord>, Box<dyn std::error::Error>> {
13511 if tokens.len() < gamma + 2 || gamma == 0 || chunk < 2 {
13512 return Err("DSpark extraction token tape/gamma/chunk is invalid".into());
13513 }
13514 if anchor_positions.windows(2).any(|pair| pair[0] >= pair[1]) {
13515 return Err("DSpark anchor positions must be sorted and unique".into());
13516 }
13517 for &position in anchor_positions {
13518 if position == 0 || position + gamma >= tokens.len() {
13519 return Err(format!(
13520 "DSpark anchor {position} has no predecessor or cannot cover gamma={gamma} in {} tokens",
13521 tokens.len()
13522 )
13523 .into());
13524 }
13525 }
13526
13527 let n_vocab = self.output.out_features();
13528 let n_embd = self.cfg.n_embd as usize;
13529 let mut cache =
13530 crate::pp::new_cache_planned(e, &self.cfg, &self.plan, tokens.len() + gamma + 8)?;
13531 let (embd_qt, embd_rb) = self.embd.qt_and_row_bytes(n_embd);
13532 let embd_gpu = if spec_host_embd() {
13533 None
13534 } else {
13535 Some(
13536 self.embd_gpu
13537 .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload")),
13538 )
13539 };
13540 let embd_dev = embd_gpu.map(|gpu| (gpu, embd_qt, embd_rb));
13541
13542 struct PendingRecord {
13543 position: usize,
13544 hidden: Option<Vec<f32>>,
13545 tokens: Vec<u32>,
13546 target_top_ids: Vec<Option<Vec<u32>>>,
13547 target_top_logits: Vec<Option<Vec<f32>>>,
13548 target_top_probs: Vec<Option<Vec<f32>>>,
13549 target_tail_probs: Vec<Option<f32>>,
13550 }
13551
13552 let mut pending: Vec<PendingRecord> = anchor_positions
13553 .iter()
13554 .map(|&position| PendingRecord {
13555 position,
13556 hidden: None,
13557 tokens: tokens[position..=position + gamma].to_vec(),
13558 target_top_ids: vec![None; gamma],
13559 target_top_logits: vec![None; gamma],
13560 target_top_probs: vec![None; gamma],
13561 target_tail_probs: vec![None; gamma],
13562 })
13563 .collect();
13564
13565 let mut start = 0usize;
13566 while start < tokens.len() {
13567 let end = (start + chunk).min(tokens.len());
13568 let chunk_tokens = &tokens[start..end];
13569 let (target_logits, hidden_rows) =
13570 self.decode_step_t_core(e, chunk_tokens, start, &mut cache, embd_dev, None)?;
13571 for record in &mut pending {
13572 let hidden_position = record.position - 1;
13573 if hidden_position >= start && hidden_position < end {
13574 let local = hidden_position - start;
13575 record.hidden = Some(
13576 e.dtoh_view(&hidden_rows.slice(local * n_embd..(local + 1) * n_embd))?,
13577 );
13578 }
13579 for slot in 0..gamma {
13580 let target_row = record.position + slot;
13581 if target_row < start || target_row >= end {
13582 continue;
13583 }
13584 let local = target_row - start;
13585 let logits =
13586 e.dtoh_view(&target_logits.slice(local * n_vocab..(local + 1) * n_vocab))?;
13587 let (ids, top_logits, probs, tail) =
13588 dspark_sparse_softmax_topk(&logits, top_k, temperature)?;
13589 record.target_top_ids[slot] = Some(ids);
13590 record.target_top_logits[slot] = Some(top_logits);
13591 record.target_top_probs[slot] = Some(probs);
13592 record.target_tail_probs[slot] = Some(tail);
13593 }
13594 }
13595 start = end;
13596 }
13597
13598 pending
13599 .into_iter()
13600 .map(|record| {
13601 let hidden = record
13602 .hidden
13603 .ok_or_else(|| format!("missing DSpark hidden at {}", record.position))?;
13604 let target_top_ids =
13605 flatten_dspark_rows(record.target_top_ids, record.position, "target ids")?;
13606 let target_top_logits = flatten_dspark_rows(
13607 record.target_top_logits,
13608 record.position,
13609 "target logits",
13610 )?;
13611 let target_top_probs =
13612 flatten_dspark_rows(record.target_top_probs, record.position, "target probs")?;
13613 let target_tail_probs = record
13614 .target_tail_probs
13615 .into_iter()
13616 .enumerate()
13617 .map(|(slot, value)| {
13618 value.ok_or_else(|| {
13619 format!("missing DSpark tail at {} slot {slot}", record.position)
13620 })
13621 })
13622 .collect::<Result<Vec<_>, _>>()?;
13623 Ok(DsparkAnchorRecord {
13624 position: record.position,
13625 hidden,
13626 tokens: record.tokens,
13627 target_top_ids,
13628 target_top_logits,
13629 target_top_probs,
13630 target_tail_probs,
13631 })
13632 })
13633 .collect()
13634 }
13635
13636 /// TEACHER-FORCED REPLAY ACCEPTANCE (hqmtp MTP-heal protocol): walk a FIXED token
13637 /// sequence and, at sampled positions, compare the MTP head's K-token draft chain against
13638 /// the trunk's own teacher-forced greedy predictions. Nothing is generated — the context is
13639 /// the corpus text itself, so (a) degenerate self-generated loops cannot inflate acceptance
13640 /// and (b) two arms (bf16 ceiling vs NVFP4) score on IDENTICAL contexts, isolating the
13641 /// quant-induced head/hidden-state mismatch from text drift.
13642 ///
13643 /// Per eval position p (context = tokens[0..=p], predecessor pairing as in spec decode):
13644 /// draft_j = chain token j from (tokens[p], h_{p-1}), then its own drafts — the exact
13645 /// eager spec-decode chain (same mtp_head_forward_dev, same rope positions).
13646 /// target_j = teacher-forced greedy pick for position p+1+j (argmax of the trunk logits
13647 /// at forced context tokens[0..p+j]). For j==0 this equals live spec
13648 /// acceptance; for j>=1 live verify would condition on the drafts, here it
13649 /// conditions on the corpus — deterministic and arm-comparable by design.
13650 ///
13651 /// Returns (rows, bg): one (p, drafts[k], targets[k]) row per eval position (ascending p),
13652 /// plus the full teacher-forced greedy track bg (bg[i] = greedy pick for position i, i>=1)
13653 /// so harnesses can cross-check runs (e.g. different chunk sizes must give identical bg).
13654 ///
13655 /// `hdump`: when Some, every position's pre-output_norm trunk hidden (the exact rows the
13656 /// draft-KV fill pairs from) streams to the file as little-endian f32 [t_total, n_embd] —
13657 /// the head-distillation extraction (hqmtp): the ENGINE is the source of truth for trunk
13658 /// hiddens (HF torch reproductions of the hybrid trunk measured only ~0.5 greedy
13659 /// agreement vs this path — not usable as a training-data source).
13660 pub fn replay_acceptance(
13661 &self,
13662 e: &Engine,
13663 tokens: &[u32],
13664 k: usize,
13665 stride: usize,
13666 chunk: usize,
13667 mut hdump: Option<&mut std::fs::File>,
13668 ) -> Result<(Vec<(usize, Vec<u32>, Vec<u32>)>, Vec<u32>), Box<dyn std::error::Error>> {
13669 assert!(k >= 1 && stride >= 1 && chunk >= 2);
13670 let mtp = self
13671 .mtp
13672 .as_ref()
13673 .expect("replay_acceptance requires an MTP head");
13674 let n_vocab = self.output.out_features();
13675 let d_vocab = mtp
13676 .shared_head_head
13677 .as_ref()
13678 .unwrap_or(&self.output)
13679 .out_features();
13680 let n_embd = self.cfg.n_embd as usize;
13681 let t_total = tokens.len();
13682 assert!(t_total >= 8, "corpus too short ({t_total} tokens)");
13683 // STAGE-OWNED KV (lane/pp2-spec 2026-08-06) — see `new_session`. Door shut = `Cache::new`.
13684 let mut cache = crate::pp::new_cache_planned(e, &self.cfg, &self.plan, t_total + k + 8)?;
13685 let mut scratch = self.new_mtp_scratch(e, t_total + k + 8)?;
13686 let (embd_qt, embd_rb) = self.embd.qt_and_row_bytes(n_embd);
13687 let embd_gpu = if spec_host_embd() {
13688 None
13689 } else {
13690 Some(
13691 self.embd_gpu
13692 .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload")),
13693 )
13694 };
13695 let embd_dev = embd_gpu.map(|g| (g, embd_qt, embd_rb));
13696
13697 // bg[i] = the trunk's greedy pick for position i under the forced context (i >= 1).
13698 let mut bg: Vec<u32> = vec![0; t_total + 1];
13699 let mut rows: Vec<(usize, Vec<u32>, Vec<u32>)> = Vec::new();
13700 let mut prev_last_h = e.zeros(n_embd)?; // predecessor hidden entering the chunk
13701 let mut seed_buf = e.zeros(n_embd)?;
13702 let mut preds_d = e.alloc_u32_zeroed(chunk)?;
13703 let nll_on = std::env::var("MEMRA_REPLAY_NLL").as_deref() == Ok("1");
13704 let (mut nll_sum, mut nll_cnt) = (0f64, 0u64);
13705 let mut s = 0usize;
13706 while s < t_total {
13707 let cend = (s + chunk).min(t_total);
13708 let tc = cend - s;
13709 let ch = &tokens[s..cend];
13710 // 1. forced trunk pass — verify path (decode-exact contract): all-column logits +
13711 // the chunk's true hiddens.
13712 let (tl_d, vx) = self.decode_step_t_core(e, ch, s, &mut cache, embd_dev, None)?;
13713 for j in 0..tc {
13714 e.argmax_token_device_col(&tl_d, j, n_vocab, &mut preds_d, j)?;
13715 }
13716 let preds = e.dtoh_u32(&preds_d)?;
13717 for j in 0..tc {
13718 bg[s + j + 1] = preds[j];
13719 }
13720 // MEMRA_REPLAY_NLL=1: teacher-forced NLL/perplexity over the same forced pass — the
13721 // checkpoint-quality metric (position j's logits score the GOLD next token).
13722 if nll_on {
13723 let jmax = if cend < t_total { tc } else { tc - 1 }; // last pos has no gold next
13724 if jmax > 0 {
13725 let ids: Vec<u32> = (0..jmax).map(|j| tokens[s + j + 1]).collect();
13726 let rows: Vec<i32> = (0..jmax as i32).collect();
13727 let idsd = e.htod_u32_v(&ids)?;
13728 let rowsd = e.htod_i32(&rows)?;
13729 let mut outd = e.zeros(jmax)?;
13730 e.softmax_gather(&tl_d, n_vocab, &idsd, &rowsd, &mut outd, n_vocab, jmax, 1.0)?;
13731 for pr in e.dtoh(&outd)? {
13732 nll_sum += -((pr.max(1e-30)) as f64).ln();
13733 nll_cnt += 1;
13734 }
13735 }
13736 }
13737 if let Some(f) = hdump.as_deref_mut() {
13738 use std::io::Write;
13739 let host: Vec<f32> = e.dtoh(&vx)?;
13740 // bf16 round-to-nearest-even — f32 doubled the disk bill at bulk
13741 // extraction scale (20M tokens x 4096 = 320GB f32 vs 160GB bf16).
13742 let mut bytes = Vec::with_capacity(tc * n_embd * 2);
13743 for v in &host[..tc * n_embd] {
13744 let b = v.to_bits();
13745 let r = b.wrapping_add(0x7FFF + ((b >> 16) & 1));
13746 bytes.extend_from_slice(&((r >> 16) as u16).to_le_bytes());
13747 }
13748 f.write_all(&bytes)?;
13749 }
13750 // CHAINLESS extraction (stride > corpus, the bulk-hdump mode): no chunk ever
13751 // drafts, so the draft-KV fills are pure waste — skip them (2 MTP-block passes
13752 // per token saved; the forced trunk pass + hdump is all the mode needs).
13753 let chainless = stride > t_total;
13754 if chainless {
13755 e.copy_view_into(
13756 &mut prev_last_h,
13757 0,
13758 &vx.slice((tc - 1) * n_embd..tc * n_embd),
13759 n_embd,
13760 )?;
13761 s = cend;
13762 continue;
13763 }
13764 // 2. TRUE predecessor-paired draft-KV fill for the chunk (row i carries h_{i-1};
13765 // row s reads the previous chunk's last true hidden, zeros at corpus start).
13766 let mut vxs = e.zeros(tc * n_embd)?;
13767 e.copy_into(&mut vxs, 0, &prev_last_h, n_embd)?;
13768 if tc > 1 {
13769 e.copy_view_into(
13770 &mut vxs,
13771 n_embd,
13772 &vx.slice(0..(tc - 1) * n_embd),
13773 (tc - 1) * n_embd,
13774 )?;
13775 }
13776 scratch.set_len(e, s)?;
13777 self.mtp_kv_fill_all(e, ch, &vxs, s, &mut scratch, embd_dev)?;
13778 // 3. draft chains at sampled positions, DESCENDING: a chain reads only slots
13779 // [0..p) (true fills) and appends at >= p; the next (smaller-p) chain's set_len
13780 // truncates those approximate appends before they can ever be read.
13781 let ps: Vec<usize> = (s..cend)
13782 .filter(|p| *p >= 1 && *p % stride == 0 && *p + k <= t_total)
13783 .collect();
13784 for &p in ps.iter().rev() {
13785 scratch.set_len(e, p)?;
13786 if p == s {
13787 e.copy_into(&mut seed_buf, 0, &prev_last_h, n_embd)?;
13788 } else {
13789 e.copy_view_into(
13790 &mut seed_buf,
13791 0,
13792 &vx.slice((p - 1 - s) * n_embd..(p - s) * n_embd),
13793 n_embd,
13794 )?;
13795 }
13796 let mut e_tok = tokens[p];
13797 let mut d_seed = e.clone_dtod(&seed_buf)?;
13798 let chain_heads = !self.mtp_extra.is_empty();
13799 let mut chain_tokens = if chain_heads {
13800 vec![tokens[p]]
13801 } else {
13802 Vec::new()
13803 };
13804 let mut chain_seeds = if chain_heads {
13805 vec![e.clone_dtod(&seed_buf)?]
13806 } else {
13807 Vec::new()
13808 };
13809 let mut drafts: Vec<u32> = Vec::with_capacity(k);
13810 for j in 0..k {
13811 let (dl_d, h_nextn) = if chain_heads {
13812 self.mtp_chain_forward_dev(
13813 e,
13814 &chain_tokens,
13815 &chain_seeds,
13816 &mut scratch,
13817 p,
13818 embd_dev,
13819 None,
13820 )?
13821 } else {
13822 self.mtp_head_forward_dev(
13823 e,
13824 mtp,
13825 e_tok,
13826 &d_seed,
13827 &mut scratch,
13828 p + 1 + j,
13829 embd_dev,
13830 None,
13831 )?
13832 };
13833 let tok_d = e.argmax_token_device(&dl_d, d_vocab)?;
13834 let idx = e.dtoh_u32_one(&tok_d)?;
13835 let d = match &mtp.d2t {
13836 Some(map) => map[idx as usize],
13837 None => idx,
13838 };
13839 drafts.push(d);
13840 if chain_heads {
13841 chain_tokens.push(d);
13842 chain_seeds.push(h_nextn);
13843 } else {
13844 e_tok = d;
13845 d_seed = h_nextn;
13846 }
13847 }
13848 // targets may live in a LATER chunk's bg — resolved after the walk.
13849 rows.push((p, drafts, Vec::new()));
13850 }
13851 // 4. restore TRUE entries for the whole chunk (the next chunk's chains and fills
13852 // expect scratch.len == cend with exact rows).
13853 scratch.set_len(e, s)?;
13854 self.mtp_kv_fill_all(e, ch, &vxs, s, &mut scratch, embd_dev)?;
13855 e.copy_view_into(
13856 &mut prev_last_h,
13857 0,
13858 &vx.slice((tc - 1) * n_embd..tc * n_embd),
13859 n_embd,
13860 )?;
13861 s = cend;
13862 }
13863 for (p, drafts, targets) in rows.iter_mut() {
13864 for j in 0..drafts.len() {
13865 targets.push(bg[*p + 1 + j]);
13866 }
13867 }
13868 rows.sort_by_key(|r| r.0);
13869 if nll_cnt > 0 {
13870 let mean = nll_sum / nll_cnt as f64;
13871 println!(
13872 "[replay-nll] tokens={nll_cnt} nll/token={mean:.5} ppl={:.4}",
13873 mean.exp()
13874 );
13875 }
13876 Ok((rows, bg))
13877 }
13878}
13879
13880#[cfg(test)]
13881mod vg_debt_tests {
13882 use super::dspark_vg_debt_projection;
13883
13884 /// TOOTH for the verify-graph admission accounting: the pool's projected remaining
13885 /// growth must be charged (pre-fix, admission charged 0 for a pool measured at
13886 /// 8,852 MiB), the projection must price the MARGINAL cost of one more key rather than
13887 /// extrapolating the pool's one-time shared allocation, and the doors that make growth
13888 /// impossible must zero the debt.
13889 #[test]
13890 fn vg_debt_projects_remaining_growth_and_respects_the_freeze_valves() {
13891 const MIB: usize = 1 << 20;
13892 let d = dspark_vg_debt_projection;
13893 // cold pool: nothing observed, one capture fits inside SPEC_SHRINK_RESERVE.
13894 assert_eq!(d(0, 256, 0, None), 0);
13895 // freeze valve MEMRA_DSPARK_VG_MAX=0: the pool cannot grow.
13896 assert_eq!(d(10, 0, 500 * MIB, None), 0);
13897 // saturated pool: at/past the cap the pool FREEZES, nothing left to reserve.
13898 assert_eq!(d(256, 256, 8852 * MIB, None), 0);
13899 assert_eq!(d(300, 256, 8852 * MIB, None), 0);
13900
13901 // BOOTSTRAP (one observation, growth unmeasurable): at most one more pool's worth.
13902 // The pre-fix mean rule extrapolated 255x here — the measured 8.5 GB phantom.
13903 assert_eq!(d(1, 256, 33 * MIB, None), 33 * MIB);
13904
13905 // MARGINAL, flat pool (the box9 receipt: reserved stayed ~33.6 MiB across captures
13906 // 1..3, so an additional key costs ~nothing and the debt must collapse to ~0 —
13907 // NOT the 8,556/4,261/2,830 MB the mean rule printed).
13908 assert_eq!(d(3, 256, 33 * MIB, Some((1, 33 * MIB))), 0);
13909
13910 // MARGINAL, genuinely growing pool: 40 MiB per new key over 2 keys, 250 slots left.
13911 let debt = d(6, 256, 273 * MIB, Some((4, 193 * MIB)));
13912 assert_eq!(debt, 250 * (40 * MIB));
13913 assert!(
13914 debt > 3 * (1536 * MIB),
13915 "real growth must dwarf SPEC_SHRINK_RESERVE"
13916 );
13917
13918 // a shrinking/recycled reading never becomes a negative charge.
13919 assert_eq!(d(6, 256, 10 * MIB, Some((4, 99 * MIB))), 0);
13920 // a stale observation at the same capture count falls back to bootstrap.
13921 assert_eq!(d(4, 256, 80 * MIB, Some((4, 80 * MIB))), 80 * MIB);
13922 }
13923}
13924
13925#[cfg(test)]
13926mod mtp_chain_tests {
13927 use super::mtp_chain_head_index;
13928
13929 #[test]
13930 fn embedded_step_heads_cycle_in_declared_order() {
13931 let actual: Vec<usize> = (0..8).map(|step| mtp_chain_head_index(step, 3)).collect();
13932 assert_eq!(actual, [0, 1, 2, 0, 1, 2, 0, 1]);
13933 }
13934
13935 #[test]
13936 fn standalone_draft_remains_single_head() {
13937 assert!((0..8).all(|step| mtp_chain_head_index(step, 1) == 0));
13938 }
13939}
13940
13941#[cfg(test)]
13942mod tp_verified_prefix_tests {
13943 use super::rewind_tp_kv_verified_prefix;
13944 use crate::tp::ResidentTpKvCache;
13945
13946 fn cache_with_committed_len(committed: usize) -> ResidentTpKvCache {
13947 let mut cache = ResidentTpKvCache::new(Vec::new(), 1, 1, 1, 1, 8);
13948 let transaction = cache.begin_transaction().unwrap();
13949 let target = cache.append_target(transaction, committed).unwrap();
13950 cache.publish_append(transaction, target).unwrap();
13951 let target = cache.commit_target(transaction, committed).unwrap();
13952 cache.publish_finalize(transaction, target).unwrap();
13953 cache
13954 }
13955
13956 #[test]
13957 fn replay_free_prefix_rewinds_tp_visibility_to_snapshot_plus_accepts() {
13958 let mut layers = vec![Some(cache_with_committed_len(5)), None];
13959 rewind_tp_kv_verified_prefix(&mut layers, &[Some(2), None], 1).unwrap();
13960 let cache = layers[0].as_ref().unwrap();
13961 assert_eq!(cache.committed_len(), 3);
13962 assert_eq!(cache.staged_len(), 3);
13963 }
13964
13965 #[test]
13966 fn replay_free_prefix_rejects_a_changed_tp_cache_shape() {
13967 let mut layers = vec![Some(cache_with_committed_len(1))];
13968 let error = rewind_tp_kv_verified_prefix(&mut layers, &[None], 1)
13969 .unwrap_err()
13970 .to_string();
13971 assert!(error.contains("changed shape"), "unexpected error: {error}");
13972 }
13973}
13974
13975#[cfg(test)]
13976mod dspark_sparse_tests {
13977 use super::dspark_sparse_softmax_topk;
13978
13979 #[test]
13980 fn topk_keeps_full_softmax_mass_and_stable_ties() {
13981 let logits = [1.0f32, 3.0, 3.0, -2.0];
13982 let (ids, top_logits, probs, tail) = dspark_sparse_softmax_topk(&logits, 2, 1.0).unwrap();
13983 assert_eq!(ids, vec![1, 2]);
13984 assert_eq!(top_logits, vec![3.0, 3.0]);
13985 let denominator = logits.iter().map(|value| (value - 3.0).exp()).sum::<f32>();
13986 let expected = 1.0 / denominator;
13987 assert!((probs[0] - expected).abs() < 1.0e-6);
13988 assert!((probs[1] - expected).abs() < 1.0e-6);
13989 assert!((tail - (1.0 - 2.0 * expected)).abs() < 1.0e-6);
13990 assert!((probs.iter().sum::<f32>() + tail - 1.0).abs() < 1.0e-6);
13991 }
13992}
13993
13994#[cfg(test)]
13995mod spec_replay_env_tests {
13996 use super::spec_replay_env_on;
13997
13998 #[test]
13999 fn replay_requires_literal_one() {
14000 assert!(!spec_replay_env_on(None));
14001 assert!(!spec_replay_env_on(Some("")));
14002 assert!(!spec_replay_env_on(Some("0")));
14003 assert!(!spec_replay_env_on(Some("true")));
14004 assert!(!spec_replay_env_on(Some("2")));
14005 assert!(spec_replay_env_on(Some("1")));
14006 }
14007}
14008
14009#[cfg(test)]
14010mod telem_tests {
14011 use super::{SPEC_TELEM_POS, SpecTelemetry, SpecTelemetryCounters};
14012
14013 #[test]
14014 fn synthetic_accept_masks_produce_tau_and_position_histogram() {
14015 let counters = SpecTelemetryCounters::default();
14016 for mask in [
14017 [true, true, true],
14018 [true, true, false],
14019 [true, false, false],
14020 [false, false, false],
14021 ] {
14022 let accepted = mask.iter().take_while(|&&value| value).count();
14023 counters.record_round(mask.len(), accepted);
14024 }
14025
14026 let snapshot = counters.snapshot();
14027 assert_eq!(
14028 (snapshot.rounds, snapshot.drafted, snapshot.accepted),
14029 (4, 12, 6)
14030 );
14031 assert_eq!(&snapshot.pos_drafted[..3], &[4, 4, 4]);
14032 assert_eq!(&snapshot.pos_accepted[..3], &[3, 2, 1]);
14033 assert_eq!(snapshot.tau(), 1.5);
14034 assert_eq!(snapshot.pos_drafted[3..], [0; SPEC_TELEM_POS - 3]);
14035 assert_eq!(snapshot.pos_accepted[3..], [0; SPEC_TELEM_POS - 3]);
14036 }
14037
14038 /// The worker's per-burst pattern: stash, accumulate, diff — the delta must isolate
14039 /// exactly the burst's contribution (pool-resumed sessions carry prior requests' counts).
14040 #[test]
14041 fn delta_isolates_burst_contribution() {
14042 let mut t = SpecTelemetry::default();
14043 // "previous request": 2 rounds of k=3, accepts 3 then 1.
14044 for (kr, na) in [(3usize, 3usize), (3, 1)] {
14045 t.rounds += 1;
14046 t.drafted += kr as u64;
14047 t.accepted += na as u64;
14048 for j in 0..kr {
14049 t.pos_drafted[j] += 1;
14050 }
14051 for j in 0..na {
14052 t.pos_accepted[j] += 1;
14053 }
14054 }
14055 let before = t;
14056 // "this burst": 1 round k=3, accepts 2.
14057 t.rounds += 1;
14058 t.drafted += 3;
14059 t.accepted += 2;
14060 for j in 0..3 {
14061 t.pos_drafted[j] += 1;
14062 }
14063 for j in 0..2 {
14064 t.pos_accepted[j] += 1;
14065 }
14066 let d = t.delta_since(&before);
14067 assert_eq!((d.rounds, d.drafted, d.accepted), (1, 3, 2));
14068 assert_eq!(&d.pos_drafted[..3], &[1, 1, 1]);
14069 assert_eq!(&d.pos_accepted[..3], &[1, 1, 0]);
14070 assert_eq!(d.pos_drafted[3..], [0; SPEC_TELEM_POS - 3]);
14071 }
14072
14073 /// merge(delta) then merge(delta2) equals accumulating both — the per-model /metrics
14074 /// aggregation invariant.
14075 #[test]
14076 fn merge_accumulates_fieldwise() {
14077 let mut agg = SpecTelemetry::default();
14078 let mut d1 = SpecTelemetry {
14079 rounds: 2,
14080 drafted: 6,
14081 accepted: 4,
14082 ..Default::default()
14083 };
14084 d1.pos_drafted[0] = 2;
14085 d1.pos_accepted[0] = 2;
14086 let mut d2 = SpecTelemetry {
14087 rounds: 1,
14088 drafted: 3,
14089 accepted: 1,
14090 ..Default::default()
14091 };
14092 d2.pos_drafted[0] = 1;
14093 d2.pos_accepted[0] = 1;
14094 d2.pos_drafted[1] = 1;
14095 agg.merge(&d1);
14096 agg.merge(&d2);
14097 assert_eq!((agg.rounds, agg.drafted, agg.accepted), (3, 9, 5));
14098 assert_eq!(agg.pos_drafted[0], 3);
14099 assert_eq!(agg.pos_accepted[0], 3);
14100 assert_eq!(agg.pos_drafted[1], 1);
14101 assert_eq!(agg.pos_accepted[1], 0);
14102 }
14103
14104 /// Wrong-snapshot diff saturates to zero instead of wrapping — the counters feed a
14105 /// public metrics surface and must never publish a u64-wrapped garbage value.
14106 #[test]
14107 fn delta_saturates_never_wraps() {
14108 let small = SpecTelemetry {
14109 rounds: 1,
14110 drafted: 2,
14111 accepted: 1,
14112 ..Default::default()
14113 };
14114 let big = SpecTelemetry {
14115 rounds: 5,
14116 drafted: 15,
14117 accepted: 9,
14118 ..Default::default()
14119 };
14120 let d = small.delta_since(&big);
14121 assert_eq!((d.rounds, d.drafted, d.accepted), (0, 0, 0));
14122 }
14123}
14124
14125#[cfg(test)]
14126mod opti_fork_tests {
14127 use super::{
14128 OptiControllerPolicy, OptiForkAction, OptiForkGateMode, OptiForkGenerationTracker,
14129 };
14130
14131 #[test]
14132 fn controller_threshold_and_three_miss_breaker_are_exact() {
14133 let mut policy = OptiControllerPolicy {
14134 threshold: 0.7,
14135 consecutive_misses: 0,
14136 breaker_tripped: false,
14137 };
14138 assert!(!policy.admit(0.699_999));
14139 assert!(policy.admit(0.7));
14140 assert!(!policy.resolve(false));
14141 assert!(!policy.resolve(false));
14142 assert!(policy.resolve(false));
14143 assert!(policy.breaker_tripped);
14144 assert!(!policy.admit(1.0));
14145 assert!(
14146 !policy.resolve(true),
14147 "a resolved hit cannot re-arm a tripped request"
14148 );
14149 assert!(policy.breaker_tripped);
14150 }
14151
14152 #[test]
14153 fn zero_threshold_is_the_true_unconditional_measurement_arm() {
14154 let mut policy = OptiControllerPolicy {
14155 threshold: 0.0,
14156 consecutive_misses: 0,
14157 breaker_tripped: false,
14158 };
14159 for _ in 0..16 {
14160 assert!(policy.admit(0.0));
14161 assert!(!policy.resolve(false));
14162 }
14163 for invalid in [f32::NAN, f32::INFINITY, -0.01, 1.01] {
14164 assert!(
14165 !policy.admit(invalid),
14166 "invalid q proxy must fail closed: {invalid}"
14167 );
14168 }
14169 assert!(!policy.breaker_tripped);
14170 assert_eq!(policy.consecutive_misses, 0);
14171 }
14172
14173 #[test]
14174 fn alternating_mode_flips_by_generation_not_round_parity() {
14175 assert_eq!(OptiForkGateMode::Alternate.action(0), OptiForkAction::Hit);
14176 assert_eq!(OptiForkGateMode::Alternate.action(1), OptiForkAction::Miss);
14177 assert_eq!(OptiForkGateMode::Alternate.action(8), OptiForkAction::Hit);
14178 assert_eq!(OptiForkGateMode::Alternate.action(9), OptiForkAction::Miss);
14179 }
14180
14181 #[test]
14182 fn live_generation_cannot_be_overwritten() {
14183 let mut tracker = OptiForkGenerationTracker::default();
14184 let g0 = tracker.reserve().unwrap();
14185 let g1 = tracker.reserve().unwrap();
14186 let err = tracker.reserve().unwrap_err().to_string();
14187 assert!(
14188 err.contains("still owns generation 0"),
14189 "unexpected error: {err}"
14190 );
14191 tracker.retire(g0).unwrap();
14192 let g2 = tracker.reserve().unwrap();
14193 assert_eq!((g2.id, g2.slot), (2, 0));
14194 tracker.retire(g1).unwrap();
14195 tracker.retire(g2).unwrap();
14196 }
14197
14198 #[test]
14199 fn teardown_rejects_a_stale_generation_tag() {
14200 let mut tracker = OptiForkGenerationTracker::default();
14201 let g0 = tracker.reserve().unwrap();
14202 tracker.retire(g0).unwrap();
14203 let err = tracker.retire(g0).unwrap_err().to_string();
14204 assert!(err.contains("teardown mismatch"), "unexpected error: {err}");
14205 }
14206}
14207
14208#[cfg(test)]
14209mod draft_graph_fallback_tests {
14210 use super::DraftGraphFallback;
14211
14212 /// The Q2 contract, part (a): a fallback flip is LOUD — exactly once per flip.
14213 #[test]
14214 fn flip_is_loud_once_and_memoized_after() {
14215 let mut f = DraftGraphFallback::default();
14216 let line = f
14217 .mark_greedy("out of memory")
14218 .expect("first flip must return the warn line");
14219 assert!(
14220 line.contains("WARN"),
14221 "flip line must be warn-level: {line}"
14222 );
14223 assert!(
14224 line.contains("out of memory"),
14225 "flip line must carry the reason: {line}"
14226 );
14227 assert!(f.greedy_failed());
14228 // re-marking an already-failed graph is the memoization: quiet, still failed.
14229 assert!(f.mark_greedy("out of memory").is_none());
14230 assert!(f.greedy_failed());
14231 // the two graphs' flags are independent (greedy flip leaves sampled capturable).
14232 assert!(!f.sampled_failed());
14233 let line_s = f
14234 .mark_sampled("capture unsupported")
14235 .expect("sampled flip is its own flip");
14236 assert!(
14237 line_s.contains("sampled"),
14238 "sampled flip names itself: {line_s}"
14239 );
14240 assert!(f.mark_sampled("capture unsupported").is_none());
14241 }
14242
14243 /// The Q2 contract, part (b): resume-from-pool RESETS both flags (fresh capture chance),
14244 /// and says so exactly when there was something to reset.
14245 #[test]
14246 fn reset_on_resume_clears_flags_and_logs_once() {
14247 let mut f = DraftGraphFallback::default();
14248 // clean session: resume is silent, nothing to reset.
14249 assert!(f.reset_on_resume().is_none());
14250 f.mark_greedy("oom").unwrap();
14251 f.mark_sampled("oom").unwrap();
14252 let note = f
14253 .reset_on_resume()
14254 .expect("a set flag must produce the reset note");
14255 assert!(
14256 note.contains("greedy+sampled"),
14257 "note names what was reset: {note}"
14258 );
14259 assert!(
14260 !f.greedy_failed() && !f.sampled_failed(),
14261 "both flags cleared"
14262 );
14263 // and the NEXT failure after a reset is a fresh flip — loud again.
14264 assert!(f.mark_greedy("oom again").is_some());
14265 let note2 = f.reset_on_resume().expect("greedy-only reset");
14266 assert!(note2.contains("(greedy)"), "single-flag note: {note2}");
14267 }
14268
14269 /// Shape-change clears (dmask realloc / mask-shape mismatch / s_key change) stay silent —
14270 /// they precede a fresh capture attempt whose own failure re-flips loudly.
14271 #[test]
14272 fn shape_change_clears_are_silent() {
14273 let mut f = DraftGraphFallback::default();
14274 f.mark_greedy("oom").unwrap();
14275 f.clear_greedy();
14276 assert!(!f.greedy_failed());
14277 f.mark_sampled("oom").unwrap();
14278 f.clear_sampled();
14279 assert!(!f.sampled_failed());
14280 // after a silent clear there is nothing left for resume to report.
14281 assert!(f.reset_on_resume().is_none());
14282 }
14283}
14284
14285/// SAMPLED DRAFT-GRAPH KEY (lane/graph-s-key-exactness-20260819).
14286///
14287/// These are the CPU teeth for an exactness bug whose live reproduction needs a GPU, a trunk, a
14288/// drafter and a two-turn session: the key itself. Every test below fails against the pre-fix key
14289/// `(seed, temp.to_bits(), k)` — `legacy_key` restates it so the collision is explicit rather
14290/// than remembered.
14291#[cfg(test)]
14292mod sampled_graph_key_tests {
14293 use super::{SampledGraphKey, debug_t_pred0};
14294
14295 /// The pre-fix key, verbatim: `let s_key = (sp_seed, sp_temp.to_bits(), k);`
14296 fn legacy_key(k: &SampledGraphKey) -> (u64, u32, usize) {
14297 (k.seed, k.temp_bits, k.k)
14298 }
14299
14300 fn pure_temp_key() -> SampledGraphKey {
14301 // temperature 1.0, filters off — today's serve default, the shape that parks a graph.
14302 SampledGraphKey::new(12345, 1.0, 3, 0, 1.0, 0.0, false)
14303 }
14304
14305 /// THE COLLISION. Two requests that differ ONLY in the truncation filters shared one key, so
14306 /// a parked pure-temp graph survived into a filtered request and the launch site launched it.
14307 #[test]
14308 fn vendor_filters_change_the_key() {
14309 let parked = pure_temp_key();
14310 // qwen3.8 generation_config.json — what the vendor-default flip makes the default shape.
14311 let vendor = SampledGraphKey::new(12345, 1.0, 3, 20, 0.95, 0.0, false);
14312 assert_eq!(
14313 legacy_key(&parked),
14314 legacy_key(&vendor),
14315 "pre-fix key collided: this is the bug, and the reason a test asserts on it",
14316 );
14317 assert_ne!(parked, vendor, "post-fix key must separate the two regimes");
14318 assert!(parked.pure_temp());
14319 assert!(!vendor.pure_temp());
14320 }
14321
14322 /// Each distribution-shaping field alone is enough to drop the parked graph.
14323 #[test]
14324 fn every_filter_field_is_keyed() {
14325 let base = pure_temp_key();
14326 for (what, other) in [
14327 (
14328 "top_k",
14329 SampledGraphKey::new(12345, 1.0, 3, 20, 1.0, 0.0, false),
14330 ),
14331 (
14332 "top_p",
14333 SampledGraphKey::new(12345, 1.0, 3, 0, 0.95, 0.0, false),
14334 ),
14335 (
14336 "min_p",
14337 SampledGraphKey::new(12345, 1.0, 3, 0, 1.0, 0.05, false),
14338 ),
14339 (
14340 "penalties",
14341 SampledGraphKey::new(12345, 1.0, 3, 0, 1.0, 0.0, true),
14342 ),
14343 ] {
14344 assert_ne!(base, other, "{what} must be part of the key");
14345 assert!(!other.pure_temp(), "{what} leaves the pure-temp regime");
14346 assert_eq!(
14347 legacy_key(&base),
14348 legacy_key(&other),
14349 "{what} was invisible to the pre-fix key",
14350 );
14351 }
14352 }
14353
14354 /// The baked constants stay keyed (this half was always right — regression cover for it).
14355 #[test]
14356 fn baked_constants_stay_keyed() {
14357 let base = pure_temp_key();
14358 assert_ne!(
14359 base,
14360 SampledGraphKey::new(999, 1.0, 3, 0, 1.0, 0.0, false),
14361 "seed"
14362 );
14363 assert_ne!(
14364 base,
14365 SampledGraphKey::new(12345, 0.7, 3, 0, 1.0, 0.0, false),
14366 "temp"
14367 );
14368 assert_ne!(
14369 base,
14370 SampledGraphKey::new(12345, 1.0, 4, 0, 1.0, 0.0, false),
14371 "k"
14372 );
14373 // bitwise on temperature: 0.7f32 vs the same value re-derived must NOT differ.
14374 assert_eq!(
14375 SampledGraphKey::new(1, 0.7, 3, 0, 1.0, 0.0, false),
14376 SampledGraphKey::new(1, 7.0 / 10.0, 3, 0, 1.0, 0.0, false),
14377 );
14378 }
14379
14380 /// THE LOAD-BEARING HALF OF THE SEED DECISION (lane/session-resume-sampler-predicate-
14381 /// 20260820). The whole-session resume predicate deliberately does NOT compare `seed`: an
14382 /// omitted serve `seed` draws fresh per-request entropy, so comparing it would refuse every
14383 /// seed-omitting sampled conversation. That is only sound because the one piece of parked state
14384 /// that BAKES the seed — this graph — is re-keyed on it, so a seed change drops and recaptures.
14385 ///
14386 /// This test is the other end of that argument, asserted here rather than remembered in a
14387 /// comment: if a future change dropped `seed` from the key, the resume predicate's exclusion
14388 /// would silently become the unsound thing it is documented not to be.
14389 /// (Paired with `seed_alone_does_not_refuse` in `memra-sampling`.)
14390 #[test]
14391 fn seed_alone_still_rekeys_the_draft_graph() {
14392 let parked = pure_temp_key();
14393 let reseeded = SampledGraphKey::new(999, 1.0, 3, 0, 1.0, 0.0, false);
14394 assert_ne!(
14395 parked, reseeded,
14396 "a seed-only change MUST drop the parked sampled graph — the resume predicate's \
14397 decision not to compare seed rests on exactly this",
14398 );
14399 // Same regime on both sides: the drop is a recapture, not a fall to the eager chain
14400 // because of a filter difference.
14401 assert!(parked.pure_temp() && reseeded.pure_temp());
14402 }
14403
14404 /// `pure_temp()` is the capture guard's predicate, computed from the key so the two cannot
14405 /// drift. The equality below is the invariant the launch-site guard asserts: identical keys
14406 /// agree on the regime, so a graph that survives the drop is legal to launch.
14407 #[test]
14408 fn equal_keys_agree_on_the_regime() {
14409 let a = SampledGraphKey::new(7, 0.8, 3, 20, 0.95, 0.0, false);
14410 let b = SampledGraphKey::new(7, 0.8, 3, 20, 0.95, 0.0, false);
14411 assert_eq!(a, b);
14412 assert_eq!(a.pure_temp(), b.pure_temp());
14413 // top_p slightly above 1.0 (a client sending 1.0 exactly, or an operator default) is
14414 // still the unfiltered regime, matching the original `sp.top_p >= 1.0` test.
14415 assert!(SampledGraphKey::new(7, 0.8, 3, 0, 1.0, 0.0, false).pure_temp());
14416 assert!(SampledGraphKey::new(7, 0.8, 3, 0, 1.5, -1.0, false).pure_temp());
14417 }
14418
14419 /// MEMRA_DEBUG_SPEC on a SAMPLED spec request past round 0: the print must render without
14420 /// indexing the empty greedy `preds` vector (it panicked the GPU worker before this lane).
14421 #[test]
14422 fn debug_print_survives_the_sampled_arm() {
14423 // round >= 1 with a pending bonus == base 1, sampled == `preds` empty.
14424 assert_eq!(debug_t_pred0(true, 1, 4242, &[]), "n/a");
14425 assert_eq!(debug_t_pred0(true, 2, 4242, &[]), "n/a");
14426 // round 0 without a pending bonus still reports last_pred, in both arms.
14427 assert_eq!(debug_t_pred0(true, 0, 4242, &[]), "4242");
14428 assert_eq!(debug_t_pred0(false, 0, 4242, &[7, 8]), "4242");
14429 // greedy keeps the real prediction it always printed.
14430 assert_eq!(debug_t_pred0(false, 1, 4242, &[7, 8]), "7");
14431 assert_eq!(debug_t_pred0(false, 2, 4242, &[7, 8]), "8");
14432 }
14433}