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 memra_gguf::config::SwigluClamp;
17use std::sync::atomic::{AtomicU64, Ordering};
18
19/// Parse the documented `MEMRA_SPEC_REPLAY=1` rollback seam.
20///
21/// Keep this shared with serving admission so `=0` cannot select replay in one
22/// layer while another layer treats it as disabled.
23pub fn spec_replay_env_on(value: Option<&str>) -> bool {
24 value == Some("1")
25}
26
27pub fn spec_replay_env_enabled() -> bool {
28 let value = std::env::var("MEMRA_SPEC_REPLAY").ok();
29 spec_replay_env_on(value.as_deref())
30}
31
32/// step35 dcw draft-chain door (lane/step37-draft-graph-20260829). ON routes the step35 MTP
33/// block's draft attention through the WINDOWED device-counter family
34/// (`append_kv_quantized_dcw` + `fa_decode_dcw`, the step TP graph arc's kernels), which
35/// derives the SWA view entirely from device state (len_d, base_d, window): exactly the view
36/// offset the old capture refusal said `fa_decode_dc` could not express. BOTH draft modes
37/// switch together: eager and captured run the ONE launcher at the ONE bucket
38/// (min(cap, window)), so graph-vs-eager draft parity holds by construction (the
39/// `mtp_full_attn_dc` precedent).
40///
41/// DEFAULT ON since lane/step37-draft-graph-serving-20260830: the 20260829 lane shipped it
42/// OFF because it enabled nothing at the shipping head count (capture was structurally
43/// unreachable at heads=3); with the multi-head chain capture and the in-graph filtered
44/// sampler landed, this door is the kernel prerequisite for the captured chain on the
45/// QUALIFIED serving shape, and the exactness battery (greedy K=1..8 identity, per-K
46/// acceptance identity, seeded sampled twins) banks on the ON arm. Rollback seam:
47/// MEMRA_STEP35_DRAFT_DCW=0 restores the host-len eager arm (`mtp_step35_attn`) plus the
48/// named capture refusal, byte-for-byte the pre-lane serving; no state survives restart.
49fn step35_draft_dcw_on() -> bool {
50 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
51 *ON.get_or_init(|| std::env::var("MEMRA_STEP35_DRAFT_DCW").as_deref() != Ok("0"))
52}
53
54/// Multi-head MTP draft-chain capture door (lane/step37-draft-graph-serving-20260830,
55/// default ON — receipts in the lane RESULTS). ON lets the step-modulo prefix-replay chain
56/// (`mtp_extra` non-empty, the step37 3-head shipping shape) capture per-head single-row
57/// CUDA graphs and replay them in the exact eager launch order; the chain POLICY (head
58/// selection, prefix length, seed history) stays host-side, so graph-vs-eager drafts are
59/// bit-identical by construction. A failed capture degrades LOUDLY to the eager chain (the
60/// draft-graph WARN contract). OFF (=0) keeps the eager chain as the only multi-head path —
61/// the pre-lane serving byte-for-byte. Single-head capture is untouched by this door.
62fn mtp_chain_graph_on() -> bool {
63 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
64 *ON.get_or_init(|| std::env::var("MEMRA_MTP_CHAIN_GRAPH").as_deref() != Ok("0"))
65}
66
67/// In-graph FILTERED sampled draft door (lane/step37-draft-graph-serving-20260830, default
68/// ON — receipts in the lane RESULTS). ON widens the sampled draft-graph capture from the
69/// pure-temp regime to every truncation-filtered regime (top_k / top_p / min_p): the capture
70/// body runs `filter_stats` + `gumbel_perturb_filtered_ctr` IN-GRAPH, so the draft draws
71/// from the SAME filtered distribution the verify's accept test reconstructs (the
72/// graph-s-key exactness law, now satisfied inside the graph instead of by refusing it).
73/// Penalties stay eager either way (the history varies per round and cannot be baked).
74/// The pure-temp capture body is UNTOUCHED by this door (byte-identical to the pre-lane
75/// graph). OFF (=0) restores the pure-temp-only capture guard: filtered requests draft
76/// eager, byte-for-byte the pre-lane behavior.
77fn spec_graph_filtered_on() -> bool {
78 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
79 *ON.get_or_init(|| std::env::var("MEMRA_SPEC_GRAPH_FILTERED").as_deref() != Ok("0"))
80}
81
82fn parse_prime_trows_width(value: Option<&str>) -> Result<usize, String> {
83 let Some(raw) = value else {
84 return Ok(8);
85 };
86 let width = raw
87 .parse::<usize>()
88 .map_err(|_| format!("MEMRA_PRIME_TROWS_T must be an integer in 2..=8, got {raw:?}"))?;
89 if !(2..=8).contains(&width) {
90 return Err(format!("MEMRA_PRIME_TROWS_T must be in 2..=8, got {width}"));
91 }
92 Ok(width)
93}
94
95#[cfg(test)]
96mod prime_trows_width_tests {
97 #[test]
98 fn width_defaults_to_eight_and_refuses_invalid_operator_values() {
99 assert_eq!(super::parse_prime_trows_width(None), Ok(8));
100 assert_eq!(super::parse_prime_trows_width(Some("2")), Ok(2));
101 assert_eq!(super::parse_prime_trows_width(Some("8")), Ok(8));
102 for invalid in ["", "1", "9", "32", "wide"] {
103 let err = super::parse_prime_trows_width(Some(invalid)).unwrap_err();
104 assert!(err.contains("MEMRA_PRIME_TROWS_T"), "{err}");
105 assert!(err.contains("2..=8"), "{err}");
106 }
107 }
108}
109
110/// One compact, anchor-bounded DSpark supervision record. `tokens[0]` is the anchor at p and
111/// `hidden` is its predecessor carrier h[p-1], matching the live NextN/DSpark pairing. Target
112/// rows p..p+gamma-1 score tokens p+1..p+gamma. They are the full-target softmax's top-k
113/// entries; `target_tail_probs[j]` is the probability mass outside those rows. All flattened
114/// target arrays are `[gamma, top_k]` in row-major order.
115pub struct DsparkAnchorRecord {
116 pub position: usize,
117 pub hidden: Vec<f32>,
118 pub tokens: Vec<u32>,
119 pub target_top_ids: Vec<u32>,
120 pub target_top_logits: Vec<f32>,
121 pub target_top_probs: Vec<f32>,
122 pub target_tail_probs: Vec<f32>,
123}
124
125#[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
126fn dspark_sparse_softmax_topk(
127 logits: &[f32],
128 top_k: usize,
129 temperature: f32,
130) -> Result<(Vec<u32>, Vec<f32>, Vec<f32>, f32), Box<dyn std::error::Error>> {
131 if logits.is_empty() || top_k == 0 || top_k > logits.len() || temperature <= 0.0 {
132 return Err("invalid DSpark sparse-softmax shape or temperature".into());
133 }
134 if logits.iter().any(|value| !value.is_finite()) {
135 return Err("DSpark target logits contain a non-finite value".into());
136 }
137 let mut ranked: Vec<(u32, f32)> = logits
138 .iter()
139 .copied()
140 .enumerate()
141 .map(|(index, value)| (index as u32, value))
142 .collect();
143 let compare = |left: &(u32, f32), right: &(u32, f32)| {
144 right.1.total_cmp(&left.1).then(left.0.cmp(&right.0))
145 };
146 ranked.select_nth_unstable_by(top_k - 1, compare);
147 ranked[..top_k].sort_unstable_by(compare);
148
149 let max_logit = logits.iter().copied().fold(f32::NEG_INFINITY, f32::max);
150 let inv_temperature = 1.0f64 / temperature as f64;
151 let denominator: f64 = logits
152 .iter()
153 .map(|value| (((*value - max_logit) as f64) * inv_temperature).exp())
154 .sum();
155 let ids: Vec<u32> = ranked[..top_k].iter().map(|(index, _)| *index).collect();
156 let top_logits: Vec<f32> = ranked[..top_k].iter().map(|(_, value)| *value).collect();
157 let top_probs: Vec<f32> = top_logits
158 .iter()
159 .map(|value| ((((value - max_logit) as f64) * inv_temperature).exp() / denominator) as f32)
160 .collect();
161 let top_mass: f64 = top_probs.iter().map(|value| *value as f64).sum();
162 let tail = (1.0f64 - top_mass).clamp(0.0, 1.0) as f32;
163 Ok((ids, top_logits, top_probs, tail))
164}
165
166fn flatten_dspark_rows<T>(
167 rows: Vec<Option<Vec<T>>>,
168 position: usize,
169 label: &str,
170) -> Result<Vec<T>, Box<dyn std::error::Error>> {
171 let mut flattened = Vec::new();
172 for (slot, row) in rows.into_iter().enumerate() {
173 flattened.extend(
174 row.ok_or_else(|| format!("missing DSpark {label} at {position} slot {slot}"))?,
175 );
176 }
177 Ok(flattened)
178}
179
180/// H-SEED CONVENTION (MEMRA_SPEC_HPOST=1): feed the MTP head the POST-norm hidden — trunk rows
181/// hand over `output_norm(x)` and the draft chain recurrence hands over `shared_head_norm(h_nextn)`
182/// (= final_h) — matching the reference engines: llama.cpp #24025 ("qwen35: use post-norm hidden
183/// state for MTP", t_h_nextn is taken AFTER the final norm in both trunk and MTP graphs) and
184/// SGLang's qwen3_5_mtp (spec_info.hidden_states = the target model's post-norm output). memra's
185/// historical convention (default, MTP-PLAN §A) is PRE-norm x. Draft-quality-only: exactness is
186/// the verify's job either way; acceptance arbitrates. OnceLock: read once, hot-loop safe.
187/// `MEMRA_SPEC_HEAD_ROWS=1` — batch the verify tail's LM head over its t columns instead of running
188/// it at m=1 once per column. See the call site in `decode_step_t_core_stream` for why the batched
189/// form is the same per-row arithmetic (the bf16/q8 rows twins, not cuBLASLt) and what it costs
190/// today: the head is re-streamed t times per verify pass. Default off until the byte tape says so.
191pub(crate) fn head_rows_on() -> bool {
192 static ENV: std::sync::OnceLock<Option<bool>> = std::sync::OnceLock::new();
193 crate::step37_door(&ENV, "MEMRA_SPEC_HEAD_ROWS")
194}
195
196/// The serving walk's own doors, tri-stated the same way (owner flip 2026-08-27): env forces,
197/// unset takes the step37 family default. Call sites are the t-row verify walk itself.
198pub(crate) fn spec_verify_eager_on() -> bool {
199 static ENV: std::sync::OnceLock<Option<bool>> = std::sync::OnceLock::new();
200 crate::step37_door(&ENV, "MEMRA_SPEC_VERIFY_EAGER")
201}
202
203pub(crate) fn spec_verify_tcol_on() -> bool {
204 static ENV: std::sync::OnceLock<Option<bool>> = std::sync::OnceLock::new();
205 crate::step37_door(&ENV, "MEMRA_SPEC_VERIFY_TCOL")
206}
207
208/// NOT family-armed (2026-08-27): the walk's prime leaves its sub-32 TAIL chunk out of the
209/// DISTRIBUTED kv, so the server refuses before decode with "cache lengths diverged
210/// local=N distributed=floor(N/32)*32" for every prompt whose token count is not a multiple of
211/// 32 — i.e. nearly all real traffic. Isolated on the server route: defaults ERR (local=445
212/// distributed=416), MEMRA_PRIME_TROWS=0 OK. It was default-OFF before the 2026-08-27 flip and
213/// goes back to opt-in until the tail append is fixed and gated ON THE SERVER ROUTE, not just
214/// run-gen (run-gen calls decode_step_t on the whole prompt and never exercises this path — the
215/// reason a run-gen-only receipt could not see it). The GEMM prime supersedes it on this route.
216pub(crate) fn prime_trows_on() -> bool {
217 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
218 *ON.get_or_init(|| std::env::var("MEMRA_PRIME_TROWS").as_deref() == Ok("1"))
219}
220
221pub(crate) fn tcol_ffn_on() -> bool {
222 static ENV: std::sync::OnceLock<Option<bool>> = std::sync::OnceLock::new();
223 crate::step37_door(&ENV, "MEMRA_TCOL_FFN")
224}
225
226pub(crate) fn spec_hpost() -> bool {
227 static H: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
228 *H.get_or_init(|| {
229 std::env::var("MEMRA_SPEC_HPOST")
230 .map(|v| v != "0")
231 .unwrap_or(false)
232 })
233}
234
235/// LEAN VERIFY (default ON since 2026-07-08; MEMRA_SPEC_LEAN=0 reverts — close35 lane): the verify m-scaling
236/// probe + nsys diff showed the verify t-path pays ~1.0ms/call at m=1 over eager decode on the
237/// 35B, and the kernels are NOT the cause (dev-MoE identical, kernel-time delta only +179us).
238/// The overhead is (a) ~250 extra cuMemsetD8Async/call from `e.zeros()` on buffers every kernel
239/// fully overwrites (~0.9ms host issue + ~0.35ms GPU) and (b) the t=1 FA rows dispatch (rows_v2 +
240/// combine_rows, +50us vs the eager fa_decode pair). This flag switches (a) fully-overwritten
241/// verify buffers to `e.uninit` (identical bytes: every element is written before read) and
242/// (b) t==1 verify FA to the eager `fa_decode` entry (byte-identical: kernel-check pins the
243/// rows-vs-loop identity and the per-row loop at t=1 IS fa_decode on the same q). Gates arbitrate.
244pub(crate) fn spec_lean() -> bool {
245 static L: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
246 // DEFAULT ON since 2026-07-08 (MEMRA_SPEC_LEAN=0 reverts): bit-identical (buffers fully
247 // overwritten; gates green incl maxdiff-identical run-gen) and measured +2.4% e2e p3 /
248 // +1.5% p2 at the daily 35B config. m=1 verify now costs eager-decode parity.
249 *L.get_or_init(|| {
250 std::env::var("MEMRA_SPEC_LEAN")
251 .map(|v| v != "0")
252 .unwrap_or(true)
253 })
254}
255
256/// SMALL-M BATCHED VERIFY (default ON since 2026-07-09; MEMRA_SPEC_M2=0 reverts — lane/spec-m2): extend the
257/// batched linear-attn verify arm down to t=2 and batch the MoE dev token loop over a
258/// grid.z=token axis at every verify t. The close35 m-scaling probe put the m=2 verify tier at
259/// x1.54 of m=1 (llama x1.14); the per-column linear chain (t<3) and the serial MoE dev token
260/// loop are the two launch-structure causes. Both changes are LAUNCH-STRUCTURE ONLY:
261/// (a) the batched conv's t<pad ring update is pure copies (ssm_conv_ring_rebuild from a cloned
262/// ring — the ring stores raw input columns); every arithmetic kernel is the same one the
263/// t>=3 arm already runs (matmul_decode_exact bit-identical at m=2-4, gdn_scan's internal
264/// t-loop == chained T=1 steps);
265/// (b) the MoE dev-rows twins run the serial loop's per-token warp program with tok-offset
266/// pointers (same sel/w/aq/ad bytes, same dot order, same slot-ordered FMA chain).
267/// Gates arbitrate: run-spec K=1..8 self-consistency (35B+9B), kernel-check, run-gen argmax.
268pub(crate) fn spec_m2() -> bool {
269 static M: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
270 // DEFAULT ON since 2026-07-09 (MEMRA_SPEC_M2=0 reverts): launch-structure only — t=2
271 // batched linear arm (ring-roll copies, zero new FP order) + MoE dev-rows kernels
272 // (grid.z=token, 4 launches/layer at any verify t). Acceptance bit-identical at every K;
273 // 35B p2 +3.4% / p3 +3.6%; the profitable-K plateau widens (new optimum K=3 at 223).
274 *M.get_or_init(|| {
275 std::env::var("MEMRA_SPEC_M2")
276 .map(|v| v != "0")
277 .unwrap_or(true)
278 })
279}
280pub(crate) fn spec_stream() -> bool {
281 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
282 *ON.get_or_init(|| std::env::var("MEMRA_SPEC_STREAM").as_deref() == Ok("1"))
283}
284pub(crate) fn spec_stream_m() -> usize {
285 static M: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
286 *M.get_or_init(|| {
287 std::env::var("MEMRA_SPEC_STREAM_M")
288 .ok()
289 .and_then(|v| v.parse().ok())
290 .unwrap_or(4)
291 })
292}
293/// Engine-bundle slice 3 + fa-execupdate slice 4c (DSF-ROUNDCOST-20260820 §5 rank 1),
294/// DEFAULT OFF — `MEMRA_DSPARK_VERIFY_GRAPH=1` opts in: per-(segment, vt) CUDA graphs
295/// for the LINEAR-layer runs, plus the full-verify single graph per (vt, rung) when a
296/// round's rows all ride one seqs rung — see [`DsparkVerifyGraphs`]. Requires the
297/// slice-2 deferred path (device tokens); the eager walk is the byte-identical fallback.
298///
299/// MEASURED disposition (box6 card0, agentic pack, 2026-08-20, both slices): exactness
300/// holds everywhere (ALL EXACT, accept lines byte-match the banks, ckpt-gate oracle
301/// green over the graph + slab-commit paths). Slice-3's AUTO_FREE launch-scan limiter
302/// (25.6 us x 16 launches ≈ 0.41 ms/round) is FIXED — the captured bodies' alloc nodes
303/// are balanced by in-graph frees (census 84/84 per segment, 1776/1776 full) so graphs
304/// instantiate USE_NODE_PRIORITY and the scan is gone. What remains at gate scale:
305/// segment graphs +0.1 tok/s over the batched-rows default (114.4 vs 114.3 x5
306/// interleaved — the linear launch overhead was only ~0.1 ms); the FULL-verify graph is
307/// NET NEGATIVE at gate scale (110.6 vs 114.2: ~14-21 (vt, rung) captures/process at
308/// 2 full-walk executions + ~2.9k-node instantiate each eat far more than the ~0.2-0.3
309/// ms/round of remaining launch overhead). The orchestration ceiling of §1.3 is spent —
310/// the fa/append recovery landed DEFAULT-ON as the batched rows arm
311/// (`dspark_fa_rows_on`), not as a graph. The serve-lifetime cell (DSF-ROUNDCOST §9,
312/// nj-ws-solo) measured the amortization: crossover K≈33 requests, steady −0.246
313/// ms/round, −1.25% session wall over 240 requests — and the graphs-serve lane wired
314/// the door into the session arm (`dspark_spec_session_burst`) as a model-owned
315/// capture pool shared across sessions. Stays opt-in pending the owner's default-ON
316/// ratification on the serve-surface battery.
317pub(crate) fn dspark_verify_graph_on() -> bool {
318 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
319 *ON.get_or_init(|| std::env::var("MEMRA_DSPARK_VERIFY_GRAPH").as_deref() == Ok("1"))
320}
321/// MTP-ROUTE verify graphs, DEFAULT ON for the GDN+MoE family since 2026-08-23
322/// (`MEMRA_SPEC_VERIFY_GRAPH=0` is the kill switch, `=1` opts other families in).
323///
324/// The slice-4c capture already lived inside `qwen35_verify_tparallel` and said so in its own
325/// comment — "stream rides the qwen35moe burst, graphs ride the dspark route" — with no caller
326/// on this route. The MTP spec round is that caller.
327///
328/// WHY it is worth a default (receipts: `research/orndecode-20260822/VGRAPH.md`). With
329/// `MEMRA_SPEC_PHASE=1` this route's round reads verify-ISSUE 44-58% and verify-WAIT **0.0%**:
330/// the host is never waiting for the device, it is spending its own time launching the trunk.
331/// Replay collapses that into one graph launch and the phase all but disappears (55-62 ms ->
332/// 8-10 ms per burst).
333///
334/// MEASURED, two host generations, forced ON/OFF, balanced 4+4 boots in both orders:
335/// * current-generation host (9950X, the serving class): OFF 266.0-266.5, ON 318.8-319.5
336/// tok/s — **+19.7%**, no overlap, sub-1% spread per arm; per-round 6.9 -> 5.7 ms.
337/// * Zen 3 host: +3-9% (that rig's own clock drift is wider than the effect, so the ratio
338/// comes from per-round phase totals, which are internal to each boot).
339/// The ON arm lands at ~320 tok/s on BOTH hosts while OFF tracks host speed — the arm moves
340/// the round off the host and onto the device, which is the whole point.
341///
342/// EXACTNESS is structural (same kernels, same order) and gated anyway: a fixed-seed SAMPLED
343/// completion hashes identically ON vs OFF **and across both hosts** (`08941d5bb9762b21`),
344/// greedy seed-pinned likewise, `run-spec` K=1..8 PASS on both arms with identical acceptance
345/// at every K, kernel-check ALL GREEN.
346///
347/// SCOPE, deliberately narrow: default ON only where it was measured — the GatedDeltaNet +
348/// MoE family (`vgraph_family_default`). Qwen3.8-27B is GDN + DENSE mlp and would otherwise
349/// inherit this default unmeasured, which is the family-by-family law this repo keeps; it can
350/// opt in with `=1` once it has its own interleave. Also never armed together with
351/// ROUND-STREAM, and a round wider than the pool declines it for the eager walk.
352pub(crate) fn spec_verify_graph_env() -> Option<bool> {
353 static ON: std::sync::OnceLock<Option<bool>> = std::sync::OnceLock::new();
354 *ON.get_or_init(
355 || match std::env::var("MEMRA_SPEC_VERIFY_GRAPH").as_deref() {
356 Ok("1") => Some(true),
357 Ok("0") => Some(false),
358 _ => None,
359 },
360 )
361}
362/// SERVE-ROUTE twin of [`dspark_verify_graph_on`], DEFAULT ON — owner-ratified
363/// 2026-08-22 on the §10 serve-lifetime battery (DSF-ROUNDCOST-20260820 §10.3:
364/// crossover K=36–43, steady −0.357 ms/round, session wall −1.55..−1.65%, byte-exact
365/// 240/240 ×3 pairs, pool bounded at 8,852 MiB under `MEMRA_DSPARK_VG_MAX`). The env
366/// stays as the kill-switch: `MEMRA_DSPARK_VERIFY_GRAPH=0` restores the eager walk
367/// (byte-identical body); `MEMRA_DSPARK_VG_MAX=0` is the finer freeze valve. The BIN
368/// arm keeps its own opt-in default (`dspark_verify_graph_on`): at gate scale the
369/// capture toll is never repaid (§8 measured disposition — 14–21 captures over a
370/// 256-token run vs the serve session's thousands of rounds), and the two
371/// instruments must keep their own measured dispositions rather than share one flag.
372pub(crate) fn dspark_verify_graph_serve_on() -> bool {
373 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
374 *ON.get_or_init(|| std::env::var("MEMRA_DSPARK_VERIFY_GRAPH").as_deref() != Ok("0"))
375}
376/// Capture-count ceiling for the dspark verify-graph pool (graphs-serve lane) — the
377/// pool's memory policy STATED instead of silently unbounded. The keyspace is
378/// intrinsically finite — segment keys (run_start, vt) ≤ 16 runs x 7 windows, full
379/// keys (vt, rung, hi) ≤ 7 windows x the split-rung ladder (8 rungs at 32k ctx), ~168
380/// on the q38 export — so the default (256) never engages there; the knob is the
381/// safety valve for a future export with a wider ladder. At the ceiling the pool
382/// FREEZES: existing keys keep replaying, rounds needing a new capture run the eager
383/// walk byte-identically (round-atomic — a partial refusal would mix slab- and
384/// cols-stashed layers inside one commit). No eviction by design: destroying a live
385/// exec graph re-opens the stale-address class the indirect tables exist to close,
386/// and the bounded keyspace makes reclaim worthless.
387pub(crate) fn dspark_vg_cap() -> usize {
388 static CAP: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
389 *CAP.get_or_init(|| {
390 std::env::var("MEMRA_DSPARK_VG_MAX")
391 .ok()
392 .and_then(|v| v.parse().ok())
393 .unwrap_or(256)
394 })
395}
396
397/// PROJECTED REMAINING GROWTH of the verify-graph pool, in bytes (lane/hermes-perf-fixes,
398/// 2026-08-23 — the admission accounting the "pool dwarfs spec admission reserve" finding
399/// asks for). The pool was measured at 8,852 MiB at storm-complete on the q38 export while
400/// admission's transient floor (`SPEC_SHRINK_RESERVE`) is 1.5 GiB and never charged for it:
401/// sessions admitted while the pool is cold overcommit VRAM the pool WILL hold, because the
402/// pool grows monotonically (no eviction by design) and is model-owned across sessions.
403///
404/// SELF-MEASURING, no per-model constant (generic-model law — the 8,852 MiB is a q38 number
405/// and proves nothing about another export): the debt is remaining capture slots x the
406/// MARGINAL bytes a capture adds to this device's graph mem pool.
407///
408/// MARGINAL, NOT MEAN — measured correction (box9 on-box receipt, 2026-08-23). The first
409/// version of this used the mean (`reserved / captures`) and the live serve log showed why
410/// that is wrong: with the pool's reservation flat at ~33.6 MiB across captures 1..3, the
411/// mean-based debt printed **8,556 MB, then 4,261, then 2,830** — it extrapolated capture
412/// #1's ONE-TIME shared allocation (staging buffers, stash slabs, pointer tables: sized
413/// once per pool, shared by every key) across all 256 slots. An 8.5 GB phantom reserve at
414/// boot can refuse admissions that would have fit, which is a worse defect than the
415/// under-charge this accounting exists to remove. The marginal reading prices what an
416/// ADDITIONAL key actually costs: two observations `(captures, reserved)` give
417/// `(r1 - r0) / (c1 - c0)`, which is ~0 on an export whose pool does not grow per key and
418/// tracks real growth on one that does.
419///
420/// BOOTSTRAP (only one observation so far, so growth is unmeasurable): reserve one more
421/// pool's worth — `min(remaining x mean, reserved)`. "We have measured `reserved` bytes for
422/// `captures` keys; until growth is measurable, assume at most a doubling" is fail-safe in
423/// the same direction as the old rule without the 255x extrapolation.
424///
425/// Before the FIRST capture the debt is 0 (a single capture lands well inside the existing
426/// 1.5 GiB floor). `cap` is the intrinsic freeze ceiling (`MEMRA_DSPARK_VG_MAX`; =0 freeze
427/// valve => the pool cannot grow => debt 0); at or past the cap the pool FREEZES, so the
428/// debt is 0 there too.
429pub fn dspark_vg_debt_projection(
430 captures: usize,
431 cap: usize,
432 reserved_bytes: usize,
433 prev: Option<(usize, usize)>,
434) -> usize {
435 if captures == 0 || cap == 0 {
436 return 0;
437 }
438 let remaining = cap.saturating_sub(captures);
439 if remaining == 0 {
440 return 0;
441 }
442 match prev {
443 // marginal growth between two observations of the same pool
444 Some((c0, r0)) if captures > c0 => {
445 let marginal = reserved_bytes.saturating_sub(r0) / (captures - c0);
446 remaining.saturating_mul(marginal)
447 }
448 // bootstrap: at most one more pool's worth
449 _ => remaining
450 .saturating_mul(reserved_bytes / captures)
451 .min(reserved_bytes),
452 }
453}
454/// PRE-CAPTURE VRAM RESERVE CHECK door (lane/step37-vram-admission-20260830), DEFAULT ON.
455/// A draft-graph capture attempt on a tight card used to be try-and-fail: the 2 warmup
456/// forwards + instantiate grew the pool to the edge BEFORE the OOM surfaced, and the
457/// "eager fallback" then ran on a card the failed attempt had just exhausted (the owner's
458/// single-session second-prompt OOM: capture WARN followed by 28 step-OOM engine errors,
459/// device at 5 MiB free). With the gate ON, a capture is attempted only when the device's
460/// effective free (driver free + async-pool cached) covers the capture's expected appetite
461/// PLUS a post-capture safety floor — otherwise the session falls back to eager EARLY,
462/// with headroom intact, through the same LOUD once-per-flip WARN. `=0` restores
463/// try-and-fail (diagnostics door; the trim-on-OOM recovery below stays active either way).
464pub fn spec_capture_gate_on() -> bool {
465 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
466 *ON.get_or_init(|| std::env::var("MEMRA_SPEC_CAPTURE_GATE").as_deref() != Ok("0"))
467}
468
469/// Post-capture safety floor the reserve check keeps free ON TOP of the capture's own
470/// appetite: the same measured constant class as the admission transient floor
471/// (capture arenas + verify activations — the admit-oom control fit). A capture that
472/// would leave less than this behind is not worth its eager-coverage risk.
473pub(crate) const CAPTURE_HEADROOM_FLOOR: usize = 1536 << 20;
474
475/// Pure verdict half of the pre-capture reserve check (unit-testable): given the device's
476/// driver-free and pool-cached bytes and the capture's expected `need`, returns
477/// `Some((required, effective))` when the capture must be REFUSED, `None` when it fits.
478pub(crate) fn capture_headroom_verdict(
479 driver_free: usize,
480 pool_cached: usize,
481 need: usize,
482 floor: usize,
483) -> Option<(usize, usize)> {
484 let effective = driver_free.saturating_add(pool_cached);
485 let required = need.saturating_add(floor);
486 (effective < required).then_some((required, effective))
487}
488
489/// Expected device appetite of a draft-graph capture attempt when no measurement exists
490/// yet (bootstrap only — the model-owned high-water gauge takes over after the first
491/// observed capture). Deliberately conservative and shape-derived, never a per-family
492/// constant: per (head, mode) capture the two warmups + capture each walk one head
493/// forward whose dominant transients are a handful of `n_embd` rows and one `d_vocab`
494/// logits row, retained by the keeper; the sampled tail additionally parks
495/// `k` q-slots + perturb/q buffers of `d_vocab` each.
496pub(crate) fn draft_capture_bootstrap_estimate(
497 heads: usize,
498 k: usize,
499 d_vocab: usize,
500 n_embd: usize,
501) -> usize {
502 let per_capture = 3usize // 2 warmups + capture body, each retaining its transients
503 .saturating_mul(d_vocab.saturating_add(8 * n_embd))
504 .saturating_mul(4)
505 .max(32 << 20); // instantiate + driver-side graph backing per capture, floor
506 let captures = heads.max(1).saturating_mul(2); // interior + last per head
507 let sampled_slots = (k.saturating_add(2))
508 .saturating_mul(d_vocab)
509 .saturating_mul(4);
510 captures
511 .saturating_mul(per_capture)
512 .saturating_add(sampled_slots)
513 .max(64 << 20)
514}
515
516/// OOM predicate for capture-failure recovery (engine-side twin of the worker's
517/// `is_cuda_oom` — the same quoted-text contract).
518pub(crate) fn capture_err_is_oom(reason: &str) -> bool {
519 reason.contains("CUDA_ERROR_OUT_OF_MEMORY") || reason.contains("out of memory")
520}
521
522/// Impure half of the pre-capture reserve check: reads the device, trims the async pool
523/// when the driver alone is short but cached blocks would cover it (graph instantiate and
524/// cuBLAS workspaces allocate from the DRIVER, not from our pool — a pool sitting on freed
525/// blocks starves them), and returns the refusal reason line when the capture must not be
526/// attempted. `None` = go ahead.
527pub(crate) fn capture_headroom_refusal(e: &Engine, need: usize) -> Option<String> {
528 let Ok((driver_free, _total)) = e.ctx().mem_get_info() else {
529 return None; // unreadable device: keep the historical try-and-fail behavior
530 };
531 let pool_cached = e.pool_cached_bytes();
532 // A capture may take AT MOST HALF the discretionary headroom: required =
533 // 2x appetite + two floors (owner's contract: "fall back to eager EARLY with headroom
534 // intact"). Measured escalation on the owner-shape cells: one floor of slack let the
535 // capture walk the card to the edge and the burst step-OOM'd immediately; two floors
536 // still allowed a capture whose session then OOM'd on its own admission-charged work,
537 // because the capture had consumed the memory the charge was counting on. Requiring
538 // the appetite TWICE means the card retains a whole capture's worth of room after the
539 // capture lands - enough for the session's charged classes and its peers' bursts. The
540 // capture is an optimization worth ~2-3 ms of TTFT (draft-graph lane receipts); at the
541 // margin it is never worth an OOM incident.
542 let floor = CAPTURE_HEADROOM_FLOOR.saturating_mul(2);
543 let required_need = need.saturating_mul(2);
544 let required = required_need.saturating_add(floor);
545 match capture_headroom_verdict(driver_free, pool_cached, required_need, floor) {
546 Some((required, effective)) => Some(format!(
547 "insufficient VRAM headroom for capture: effective free {}MB (driver {}MB + pool-cached \
548 {}MB) < required {}MB (2x appetite {}MB + floor {}MB); capture skipped pre-attempt",
549 effective / (1 << 20),
550 driver_free / (1 << 20),
551 pool_cached / (1 << 20),
552 required / (1 << 20),
553 need / (1 << 20),
554 floor / (1 << 20),
555 )),
556 None => {
557 if driver_free < required && pool_cached > 0 {
558 let trimmed = e.pool_trim_to_zero();
559 if trimmed > 0 {
560 eprintln!(
561 "[spec] pre-capture pool trim: released {}MB cached back to the driver \
562 (driver free {}MB < required {}MB; instantiate allocates from the driver)",
563 trimmed / (1 << 20),
564 driver_free / (1 << 20),
565 required / (1 << 20),
566 );
567 }
568 }
569 None
570 }
571 }
572}
573
574/// GRAPH-LAUNCH HEADROOM FLOOR (lane/step37-vram-admission-20260830, defect 3 root
575/// cause): `cuGraphLaunch` SEGFAULTS inside libcuda (offset +0x27c87f, a null internal
576/// dereference at address 0x60) when a captured graph is dispatched into a
577/// driver-exhausted card — reproduced on this lane's box with core dumps on BOTH the
578/// pre-lane and lane binaries (multi-active step-OOM squeeze; the crashing thread sits in
579/// `CudaGraph::launch` inside `generate_spec_inner2`). The eager arms fail RECOVERABLY on
580/// the same card (a quoted CUDA OOM the park path handles), so below this driver-free
581/// floor every graph arm yields to eager for the round. A named constant, not a knob: the
582/// winning value is the default and the guard exists to make a driver segfault
583/// unreachable, not to tune anything.
584pub(crate) const GRAPH_LAUNCH_MIN_FREE: usize = 256 << 20;
585
586/// Per-round guard for the floor above. Read failure keeps serving (never a false
587/// refusal from an unreadable device); one `mem_get_info` (~microseconds) per ~25ms round.
588pub(crate) fn graph_launch_headroom_ok(e: &Engine) -> bool {
589 match e.ctx().mem_get_info() {
590 Ok((free, _total)) => free >= GRAPH_LAUNCH_MIN_FREE,
591 Err(_) => true,
592 }
593}
594
595/// One grep-stable suspension line per ROUTE (each call site holds its own
596/// process-lifetime `Once`): every captured-graph launch route below the floor names
597/// itself in the tag while keeping the same `graph replay suspended:` key the step37
598/// admission lane's squeeze cell greps for. The spec-round guard keeps its original
599/// per-generation `[spec]` line; the sweep routes (graph-launch-guard-sweep lane,
600/// 2026-08-31) note once per process — presence is what the gates assert, and a
601/// suspended round is otherwise byte-identical to its eager twin.
602pub(crate) fn graph_replay_suspended_note(route: &str) {
603 eprintln!(
604 "[{route}] graph replay suspended: driver free below the {}MB launch floor \
605 (eager arms serve; cuGraphLaunch segfaults into an exhausted card)",
606 GRAPH_LAUNCH_MIN_FREE / (1 << 20)
607 );
608}
609
610/// Engine-bundle slice 4 (fa-execupdate lane, DSF-ROUNDCOST-20260820 §6 close: "the
611/// residual gap lives in the FULL-ATTENTION per-row section"), DEFAULT ON —
612/// `MEMRA_DSPARK_FA_ROWS=0` reverts to the per-row loop: when every row of a verify
613/// round takes the v4-seqs arm on ONE `fa_split_keys` rung (the straddle law, evaluated
614/// at the round's first and last t_kv — both eligibility gates are intervals in t_kv),
615/// the qwen35 t-parallel verify's per-row KV-append + fa-decode loop collapses into the
616/// z-batched serving twins: ONE `append_quantize_kv_q8_0_q5_1_seqs` + ONE
617/// `fa_decode_vec_q_seqs_v4` + ONE combine per full-attention layer, replacing
618/// T x (4 dtod row copies + append + 3 memsets + main + combine) launches. Bytes are
619/// pinned by the batched-tick increment-2 kernel-check (seqs-vs-per-seq-loop bit
620/// identity: per-row T_kv derives in-kernel from pos_seq[z]; splits >= ns_eff write the
621/// empty partial the combine never reads, so the shared n_splits_max stride changes no
622/// bytes) and re-gated e2e by this lane's battery.
623pub(crate) fn dspark_fa_rows_on() -> bool {
624 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
625 *ON.get_or_init(|| {
626 std::env::var("MEMRA_DSPARK_FA_ROWS")
627 .map(|v| v != "0")
628 .unwrap_or(true)
629 })
630}
631
632/// `t_pred0` for the `MEMRA_DEBUG_SPEC` per-round print, sampled-safe.
633///
634/// `generate_spec_inner2` fills its `preds` vector ONLY on the greedy path (`if !sampled`), and
635/// the per-round debug print was the sole consumer in the sampled arm: `t_pred(0)` survives round
636/// 0 (`base == 0` returns `last_pred`) and from round 1 (`base == 1`, a pending bonus) indexes an
637/// EMPTY vector — `index out of bounds: the len is 0 but the index is 0`, in the GPU worker
638/// thread, which then respawns and reloads weights while the request dies. So any sampled spec
639/// request longer than one round used to kill the worker whenever `MEMRA_DEBUG_SPEC` was set:
640/// the flag crashed precisely the regime it exists to investigate.
641///
642/// Fixed at the print site, not inside the closure, so the greedy accept walk keeps its strict
643/// indexing (an out-of-range pred there is a real bug and must still be loud).
644fn debug_t_pred0(sampled: bool, base: usize, last_pred: u32, preds: &[u32]) -> String {
645 if base == 0 {
646 return last_pred.to_string();
647 }
648 match preds.get(base - 1) {
649 Some(p) => p.to_string(),
650 // sampled: the greedy per-column argmax was never run for this round.
651 None => {
652 debug_assert!(
653 sampled,
654 "greedy spec: preds[{}] missing at base {base}",
655 base - 1
656 );
657 "n/a".to_string()
658 }
659 }
660}
661
662/// `MEMRA_SKEY_PROBE=1` — sampled-draft-graph key probe (lane/graph-s-key-exactness-20260819).
663///
664/// Reports, per burst and per round, which draft chain the sampled arm chose and under which
665/// filter regime, plus the ONE observable that separates a legal filtered draft from a stale
666/// pure-temp graph replayed under filters: an accept test whose gathered `q` is exactly 0.
667/// A draft token sampled from the FILTERED softmax can never gather q=0 (it was drawn from the
668/// kept set), so `q=0` in the verify means the draft came from a distribution the verify does
669/// not believe in — and `u * 0 < p` then accepts it unconditionally.
670///
671/// Its own env var, deliberately NOT `MEMRA_DEBUG_SPEC`: that flag panicked the GPU worker on
672/// any sampled spec request past round 0 until this lane fixed it (§2 of the bank note).
673pub(crate) fn skey_probe() -> bool {
674 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
675 *ON.get_or_init(|| std::env::var("MEMRA_SKEY_PROBE").as_deref() == Ok("1"))
676}
677
678/// GRAMMAR HOOK for constrained spec decode (lane/constrained-full, 2026-08-03). The engine
679/// stays llguidance-agnostic: the server adapts its per-session grammar state behind this
680/// trait. CONTRACT (the verify-side truncation rule — token-identical to constrained plain
681/// greedy decode): the exactness walk runs UNMASKED first; the hook then (a) truncates
682/// acceptance at the first grammar-illegal accepted token, and (b) when the truncation fired
683/// or the bonus is illegal, the engine recomputes that slot as the MASKED argmax of the
684/// target's own verify column (an unmasked argmax that is grammar-legal IS the masked argmax
685/// — masking only removes tokens — so the common case pays nothing). `consume` advances the
686/// state with each EMITTED token in order; EOS handling is the implementor's job (skip).
687pub trait SpecConstraint {
688 /// -inf the current state's banned ids on a HOST logits row (prompt-tail / init-feed
689 /// masked argmax).
690 fn mask_logits(&mut self, logits: &mut [f32]) -> Result<(), String>;
691 /// Packed 32-bit bitset words of the CURRENT state's allowed set (device-mask form).
692 fn mask_words(&mut self) -> Result<Vec<u32>, String>;
693 /// Is `tok` consumable in the CURRENT state?
694 fn is_allowed(&mut self, tok: u32) -> Result<bool, String>;
695 /// Advance the state with an emitted token.
696 fn consume(&mut self, tok: u32) -> Result<(), String>;
697
698 // --- DRAFT-SIDE MASKING (lane/draft-mask, 2026-08-04) ---
699 // The drafter proposed grammar-illegal tokens under tight schemas, so verify-side
700 // truncation cut nearly every round (measured acceptance 0.467-0.513 tight vs 0.62-0.82
701 // loose, research/constrained-full-20260803). These three methods let the engine mask the
702 // DRAFT model's own sampling with the grammar's legal set, so proposals are legal by
703 // construction. The state they walk is a SPECULATIVE CLONE of the session matcher — the
704 // real state is advanced only by `consume` (emitted tokens), so verify-side truncation
705 // stays the correctness backstop and the emitted stream is unchanged by construction
706 // (an accepted draft is the target's unmasked argmax AND grammar-legal, hence the masked
707 // argmax; a cut slot is recomputed as the masked argmax either way).
708 // Default impls = feature OFF (pre-lane behaviour: unmasked drafts).
709
710 /// Is draft-side masking available on this hook? Probed ONCE per burst, before the draft
711 /// graph is captured (the mask is an in-graph node — its presence is a capture-time shape).
712 fn draft_mask_enabled(&self) -> bool {
713 false
714 }
715 /// Start a draft chain: clone the CURRENT (committed) grammar state into the speculative
716 /// slot. Called once per spec round, before the first draft position.
717 fn draft_begin(&mut self) -> Result<(), String> {
718 Ok(())
719 }
720 /// Packed 32-bit bitset words of the SPECULATIVE state's allowed set (target-vocab ids),
721 /// for the draft position about to be sampled. `None` = draft masking off (no-op).
722 fn draft_mask_words(&mut self) -> Result<Option<Vec<u32>>, String> {
723 Ok(None)
724 }
725 /// Advance the SPECULATIVE state with a PROPOSED draft token. `false` = the chain cannot
726 /// continue (EOS proposed, or an unmasked position proposed something illegal) — the
727 /// engine stops drafting; the token already pushed still goes through verify.
728 fn draft_advance(&mut self, _tok: u32) -> Result<bool, String> {
729 Ok(false)
730 }
731}
732
733/// DRAFT-MASK UPLOAD (lane/draft-mask): pull the speculative state's allowed set (TARGET-id
734/// space) from the hook, project it into the DRAFT head's vocab space, and upload it into the
735/// stable device buffer the draft chain reads. Returns false when the chain must stop drafting:
736/// the hook handed out no mask, or NO draft-vocab row is grammar-legal at this position (a
737/// trimmed FR-Spec head genuinely cannot propose a legal token there — masking it would leave
738/// a fully-banned row whose argmax is meaningless, so the round drafts fewer tokens and the
739/// verify emits the masked argmax as usual).
740fn upload_draft_mask(
741 e: &Engine,
742 c: &mut dyn SpecConstraint,
743 dst: &mut CudaSlice<u32>,
744 d2t: Option<&Vec<u32>>,
745 d_vocab: usize,
746 words: usize,
747) -> Result<bool, Box<dyn std::error::Error>> {
748 let Some(tw) = c
749 .draft_mask_words()
750 .map_err(|e2| format!("constraint: {e2}"))?
751 else {
752 return Ok(false);
753 };
754 let bit = |t: usize| -> bool {
755 let w = t >> 5;
756 w < tw.len() && (tw[w] >> (t & 31)) & 1 == 1
757 };
758 let mut buf = vec![0u32; words];
759 match d2t {
760 // TRIMMED draft head: row i proposes target id d2t[i] — permute the mask accordingly.
761 Some(map) => {
762 for (i, &t) in map.iter().enumerate().take(d_vocab) {
763 if bit(t as usize) {
764 buf[i >> 5] |= 1u32 << (i & 31);
765 }
766 }
767 }
768 // UNTRIMMED: draft ids ARE target ids; the packed words transfer verbatim (a short
769 // mask leaves the padded tail zeroed == banned, same rule as constrained::apply_mask).
770 None => {
771 let n = tw.len().min(words);
772 buf[..n].copy_from_slice(&tw[..n]);
773 }
774 }
775 if buf.iter().all(|w| *w == 0) {
776 return Ok(false);
777 }
778 e.htod_u32_into(dst, &buf)?;
779 Ok(true)
780}
781
782/// Keep the full token-embedding table in host memory and upload only the rows needed by each
783/// MTP/verify step. This is an exact memory-capacity seam for very large BF16 vocab tables: host
784/// gather expands the same source bits to f32, and only O(T*n_embd) bytes cross PCIe per step.
785/// CUDA-graph/round-stream draft paths require device token ids and therefore stay disabled.
786pub(crate) fn spec_host_embd() -> bool {
787 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
788 *ON.get_or_init(|| std::env::var("MEMRA_SPEC_HOST_EMBD").as_deref() == Ok("1"))
789}
790
791/// VERIFY-TIER TRUNK LAUNCH-FUSION (default ON since 2026-07-09; MEMRA_SPEC_FUSED_T=0 reverts — lane/close35b): extend
792/// the t=1 fused2/fused3 Q8_0 trunk launches to the batched verify tier (t=2-4, the K=1..3
793/// verify shapes). At t>1 the trunk pairs/triples (35B wqkv+wqkv_gate, wq/wk/wv,
794/// gate_shexp+up_shexp) each run a separate `matmul_decode_exact` — one q8_1 re-quantize of the
795/// SAME activation plus one _b2/_b4 launch per tensor. The fused twins share ONE quantize and
796/// ONE launch per group; per (tensor,token,row) the kernel body is q8_0_mmvq_batched verbatim
797/// with the identical row mapping -> BIT-IDENTICAL by construction (kernel-check pins it,
798/// run-spec K=1..8 + acceptance identity arbitrate e2e).
799pub(crate) fn spec_fused_t() -> bool {
800 static F: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
801 // DEFAULT ON since 2026-07-09 (MEMRA_SPEC_FUSED_T=0 reverts): verify t=2-4 trunk launch-fusion
802 // (fused2/fused3 Q8_0 batched twins, bit-identical by construction — m=1 block-offset split on
803 // the batched body). m=2 marginal token 2117->1762us; 35B daily: p3 +3.7% (crosses llama), p2 +5%.
804 *F.get_or_init(|| {
805 std::env::var("MEMRA_SPEC_FUSED_T")
806 .map(|v| v != "0")
807 .unwrap_or(true)
808 })
809}
810
811/// zeros/uninit switch for verify-path buffers that are FULLY OVERWRITTEN before any read.
812/// Only call this on such buffers — the lean contract is "identical bytes by construction".
813/// TOKEN-ID GUARD for every id that reaches an embed gather (#87 family).
814///
815/// A device argmax seeds its running index with 0x7FFFFFFF and replaces it only through
816/// comparisons, all of which are FALSE against NaN. An all-NaN logits row therefore returns
817/// the sentinel, and the next thing done with a token id is `embed_row(id)` — table +
818/// ~4.6 TB, never mapped, an MMU fault that kills the CUDA context for the whole process
819/// (research/pp2spec-crash-20260807). The draft chain and the GREEDY verify walk already
820/// trap this; the SAMPLED verify bonus, the boundary sampler and the replay arm's last_pred
821/// did not, which is why the recoverable fault on the greedy instrument is a TERMINAL one on
822/// the vendor-default sampled shape we actually serve.
823pub(crate) fn guard_vocab_token(
824 tok: u32,
825 n_vocab: usize,
826 what: &str,
827) -> Result<u32, Box<dyn std::error::Error>> {
828 if (tok as usize) >= n_vocab {
829 return Err(format!(
830 "{what}: token id 0x{tok:08x} >= n_vocab {n_vocab} — an all-NaN logits row left \
831 the device argmax's init sentinel in place; refusing to dereference the embed \
832 row (#87 trap)"
833 )
834 .into());
835 }
836 Ok(tok)
837}
838
839/// SPEC NaN-ORIGIN SCAN (`MEMRA_SPEC_NAN_SCAN=1`, DEFAULT OFF, diagnostic only).
840///
841/// The `#87` trap reports an all-NaN VERIFY logits column, which says the poison reached the
842/// head but not where it entered. With the scan armed the verify walk syncs and reads back
843/// every layer's output, so the FIRST layer whose residual carries a NaN names itself with the
844/// round's row and position. Off by default and never on a serving path: it costs one host
845/// sync + one `t*n_embd` D2H per layer, and the syncs change scheduling (so a run that stops
846/// reproducing under the scan is itself a datum, not an all-clear).
847///
848/// Rollback seam: unset `MEMRA_SPEC_NAN_SCAN` (or set it to 0). Every call site is behind
849/// `spec_nan_scan()`, so the default path keeps the exact launch sequence it had.
850pub(crate) fn spec_nan_scan() -> bool {
851 spec_nan_scan_level() > 0
852}
853
854/// `MEMRA_SPEC_NAN_SCAN` as a LEVEL, not a boolean. `1` scans each layer's residual, which
855/// names the layer. `2` also scans INSIDE the t-column layer body — the per-column attention
856/// output, the deferred-column o-proj/fa2 join, the post-attention norm and the routed-MoE
857/// output — because "layer 20 poisons row 0" does not say whether the attention or the routed
858/// MoE produced it, and those are different bugs with different fixes.
859pub(crate) fn spec_nan_scan_level() -> u8 {
860 static LVL: std::sync::OnceLock<u8> = std::sync::OnceLock::new();
861 *LVL.get_or_init(|| match std::env::var("MEMRA_SPEC_NAN_SCAN").as_deref() {
862 Ok("1") => 1,
863 Ok("2") => 2,
864 _ => 0,
865 })
866}
867
868/// Read back `[rows, cols]` and fail with the first NaN's coordinates. `what` names the
869/// producer (layer index, walk arm) so the error line is the localization.
870/// VERIFY-ARM RECEIPT (rides `MEMRA_SPEC_NAN_SCAN>=1`, bounded to 200 lines).
871///
872/// Names, per trunk layer, WHICH attention arm the t-column walk actually took. This exists
873/// because the level-1 residual scan below sat only on the non-fused tail: the fused
874/// rope+append+fa arm ends in `continue`, so every layer that fused was NEVER SCANNED and
875/// silently read as "clean". A poisoned residual therefore first reported at the next
876/// non-fused layer, which is how "layer 20 creates the poison" could be true of the scan and
877/// false of the engine. Also carries the row-table lookup counter, so "the fused path never
878/// ran" is distinguishable from "it ran and was innocent".
879/// KV-PLANE SCAN (`MEMRA_KV_PLANE_SCAN=1`, DEFAULT OFF, diagnostic only).
880///
881/// Reads back the STAGED rows of a layer's distributed K/V planes and reports the first row
882/// whose quantization scale is not finite. No kernel required: q8_0 blocks are
883/// `[half d][32 x i8]` and q5_1 blocks carry `half d` then `half m`, so the fp16 scale at the
884/// head of each block is host-checkable straight out of the byte plane.
885///
886/// It exists because the level-2 bad-row bitmap says EVERY verify row is non-finite at a
887/// global-attention layer's join, and row r attends a strict superset of row r-1's keys: that
888/// implicates the shared KV history those rows walk, not per-column staging. "The attention
889/// output is NaN" and "the KV history it attends is already NaN" are different bugs with
890/// different owners, and nothing measured so far separates them. A first-corrupt-row index
891/// also dates the corruption against the prime/decode boundary.
892///
893/// Bounded hard: only layers whose geometry has NO window (the global planes), only the first
894/// `MEMRA_KV_PLANE_SCAN_ROUNDS` verify rounds of a process (default 2), and it copies only
895/// `[0, staged_len)`, which is ~1.6 MB at the 1480-token repro rather than the 262144-row
896/// provision. It still syncs per layer, so it is never a serving or a measured-perf arm.
897pub(crate) fn kv_plane_scan_on() -> bool {
898 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
899 *ON.get_or_init(|| std::env::var("MEMRA_KV_PLANE_SCAN").as_deref() == Ok("1"))
900}
901
902fn kv_plane_scan_rounds() -> usize {
903 static R: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
904 *R.get_or_init(|| {
905 std::env::var("MEMRA_KV_PLANE_SCAN_ROUNDS")
906 .ok()
907 .and_then(|v| v.parse().ok())
908 .unwrap_or(2)
909 })
910}
911
912/// First non-finite fp16 block scale in `bytes`, as (block index, raw u16), scanning one
913/// scale every `stride` bytes. Returns None when every block scale is finite.
914fn first_bad_scale(bytes: &[u8], stride: usize) -> Option<(usize, u16)> {
915 if stride == 0 {
916 return None;
917 }
918 for (i, blk) in bytes.chunks_exact(stride).enumerate() {
919 let raw = u16::from_le_bytes([blk[0], blk[1]]);
920 if half_is_non_finite(raw) {
921 return Some((i, raw));
922 }
923 }
924 None
925}
926
927/// IEEE binary16: exponent all ones is Inf or NaN, whatever the mantissa says.
928fn half_is_non_finite(raw: u16) -> bool {
929 (raw & 0x7C00) == 0x7C00
930}
931
932/// Scan one layer's staged K/V planes for a non-finite quantization scale. Returns the
933/// receipt line, or None when the layer is out of scope or every scale is finite.
934pub(crate) fn scan_kv_plane(
935 e: &crate::Engine,
936 distributed: &memra_kv::ResidentTpKvCache,
937 il: usize,
938 pos0: usize,
939) -> Result<(), Box<dyn std::error::Error>> {
940 // One "round" is one pos0, not one layer: the walk visits 45 layers per verify. The
941 // default of 2 rounds is for a fault that shows up immediately; the step37 repro does not
942 // fire until rep 3 or later, i.e. round ~60 of the process, so that arm MUST raise
943 // MEMRA_KV_PLANE_SCAN_ROUNDS or it will scan only the two rounds that were never going to
944 // be poisoned and report a clean history it never looked at.
945 static ROUNDS: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
946 static LAST_POS: std::sync::atomic::AtomicUsize =
947 std::sync::atomic::AtomicUsize::new(usize::MAX);
948 if LAST_POS.swap(pos0, std::sync::atomic::Ordering::Relaxed) != pos0 {
949 ROUNDS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
950 }
951 if ROUNDS.load(std::sync::atomic::Ordering::Relaxed) > kv_plane_scan_rounds() {
952 return Ok(());
953 }
954 let staged = distributed.staged_len();
955 if staged == 0 {
956 return Ok(());
957 }
958 // ENGAGEMENT RECEIPT. This scan prints only on corruption, so `kvbad=0` in a cell would
959 // read the same whether the history was clean or the scan never ran once. Bounded so a
960 // 45-layer walk cannot flood the log.
961 static SEEN: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
962 let seen = SEEN.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
963 let (ktb, vtb) = (distributed.k_tok_bytes(), distributed.v_tok_bytes());
964 if seen < 4 {
965 eprintln!(
966 "[kv-plane] engaged #{seen} layer {il} pos0={pos0} staged={staged} \
967 ktok={ktb} vtok={vtb} (scan armed; a corrupt plane prints its own line)"
968 );
969 }
970 for rank in 0..distributed.ranks().len() {
971 let Some(rc) = distributed.rank(rank) else {
972 continue;
973 };
974 // q8_0 K blocks are [half d][32 x i8] = 34B; q5_1 V blocks lead with half d then half m.
975 let kbytes = e.dtoh_u8_view(&rc.k().slice(0..staged * ktb))?;
976 let vbytes = e.dtoh_u8_view(&rc.v().slice(0..staged * vtb))?;
977 let kbad = first_bad_scale(&kbytes, 34);
978 let vbad = first_bad_scale(&vbytes, 24);
979 if kbad.is_some() || vbad.is_some() {
980 let row = |b: Option<(usize, u16)>, tok: usize| {
981 b.map(|(i, raw)| format!("blk {i} (row {}) raw={raw:#06x}", i * 34 / tok.max(1)))
982 .unwrap_or_else(|| "clean".into())
983 };
984 eprintln!(
985 "[kv-plane] layer {il} rank {rank} pos0={pos0} staged={staged} K={} V={} - the attended KV history is ALREADY non-finite, so a non-finite attention output here is a symptom and not the origin",
986 row(kbad, ktb),
987 row(vbad, vtb)
988 );
989 return Ok(());
990 }
991 }
992 Ok(())
993}
994
995pub(crate) fn verify_arm_receipt(
996 arm: &str,
997 il: usize,
998 pos0: usize,
999 t: usize,
1000 staged: Option<usize>,
1001) {
1002 static N: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
1003 if N.fetch_add(1, std::sync::atomic::Ordering::Relaxed) >= 200 {
1004 return;
1005 }
1006 eprintln!(
1007 "[verify-arm] layer {il} arm={arm} pos0={pos0} t={t} staged_len={} rows_tab_lookups={}",
1008 staged.map(|v| v as i64).unwrap_or(-1),
1009 crate::tp::ROWS_TAB_ENGAGED.load(std::sync::atomic::Ordering::Relaxed)
1010 );
1011}
1012
1013pub(crate) fn nan_scan_rows(
1014 e: &Engine,
1015 buf: &CudaSlice<f32>,
1016 rows: usize,
1017 cols: usize,
1018 what: &str,
1019) -> Result<(), Box<dyn std::error::Error>> {
1020 // The readback is also the ATTRIBUTION point for an asynchronous fault: a
1021 // CUDA_ERROR_ILLEGAL_ADDRESS raised by any launch since the previous scan surfaces on this
1022 // sync, and the bare DriverError names nothing. Wrapping it with `what` turns "the process
1023 // died somewhere" into "it died at or before this layer, on this row, at this position".
1024 let host = e.dtoh(buf).map_err(|err| -> Box<dyn std::error::Error> {
1025 format!(
1026 "spec nan-scan: sync at {what} FAILED: {err} — the fault is at or before \
1027 this point in the walk"
1028 )
1029 .into()
1030 })?;
1031 if host.len() < rows * cols {
1032 return Err(format!(
1033 "nan-scan {what}: buffer holds {} < {rows}x{cols}",
1034 host.len()
1035 )
1036 .into());
1037 }
1038 // SCAN EVERY ROW BEFORE REPORTING. A first-hit return says "row 0 is bad" and leaves the
1039 // other rows UNEXAMINED, which is exactly the bit that discriminates the two mechanisms: in
1040 // the t-column verify, row 0 attends keys [0..p+1) and row 1 attends [0..p+2), a strict
1041 // superset, so poison in the SHARED KV history must appear in BOTH rows, while poison in
1042 // per-column staging can appear in one. Report the whole map.
1043 let mut per_row: Vec<usize> = Vec::with_capacity(rows);
1044 let mut first_bad: Option<(usize, usize)> = None;
1045 for r in 0..rows {
1046 let row = &host[r * cols..(r + 1) * cols];
1047 let bad = row.iter().filter(|v| !v.is_finite()).count();
1048 per_row.push(bad);
1049 if bad > 0 && first_bad.is_none() {
1050 first_bad = Some((r, row.iter().position(|v| !v.is_finite()).unwrap_or(0)));
1051 }
1052 }
1053 if let Some((r0, c0)) = first_bad {
1054 let map: String = per_row
1055 .iter()
1056 .map(|&b| if b == 0 { '.' } else { 'X' })
1057 .collect();
1058 return Err(format!(
1059 "spec nan-scan: {what} produced non-finite values — rows[{rows}] map={map} \
1060 counts={per_row:?} of {cols} each; first at row {r0} element {c0}. Both rows bad \
1061 implicates shared state (the KV history this layer reads); one row bad implicates \
1062 per-column staging."
1063 )
1064 .into());
1065 }
1066 Ok(())
1067}
1068
1069fn vbuf(e: &Engine, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
1070 if spec_lean() { e.uninit(n) } else { e.zeros(n) }
1071}
1072
1073/// Scratch KV for the MTP block (one full-attn layer).
1074///
1075/// PERSISTENT MODE (default, 2026-07-03 — the acceptance lever): sized cap = max_ctx and kept in
1076/// sync with the COMMITTED sequence — slot p holds the MTP block's K/V for committed token p
1077/// (roped p+1, the chain's rope convention), so the draft chain's self-attention sees the FULL
1078/// committed history instead of only the current round's 1..K+1 chain tokens (the reference
1079/// engine's "mtp_update" design). Entries come from two sources:
1080/// - chain appends: accepted positions KEEP their chain-computed entries (embedding exact,
1081/// hidden chain-approximate — the reference engine accepts the same);
1082/// - `mtp_kv_fill` batches: prompt positions + the last-draft position on full accept, computed
1083/// from EXACT trunk hiddens (K/V-only MTP-block pass, no attention/FFN/lm_head).
1084/// Rejected drafts / p-min extras / pseudo-seed appends are all discarded by the round-start
1085/// `set_len` truncation (the KvLayer len mechanism — §C rollback for the draft side).
1086/// Multi-turn spec-decode session (2026-07-05): trunk Cache + persistent MTP draft scratch +
1087/// the committed token list, alive across generate_spec_session calls. Turn N+1 primes ONLY its
1088/// suffix (chunked continuation prime over the quantized past) and mtp_kv_fill's its suffix rows,
1089/// then the round loop runs unchanged. `last_h` carries the pre-output_norm hidden of the last
1090/// committed row across turns (the predecessor-pairing seed + fill anchor).
1091/// Per-request sampling config for the sampled-spec serve path.
1092#[derive(Clone, Copy, Debug)]
1093pub struct SpecSampling {
1094 pub temp: f32,
1095 pub seed: u64,
1096 pub top_k: i32, // 0 = off
1097 pub top_p: f32, // 1.0 = off
1098 pub min_p: f32, // 0.0 = off
1099 pub penalty_last_n: usize, // 0 = penalties off
1100 pub penalty_repeat: f32,
1101 pub penalty_freq: f32,
1102 pub penalty_present: f32,
1103}
1104
1105impl SpecSampling {
1106 /// Non-identity penalties requested — THE `pen_on` predicate (one definition; the
1107 /// same group-off rule `SamplerIdentity::of` canonicalizes: a window with neutral
1108 /// coefficients is penalties-absent). Both spec routes and the dspark accept walk
1109 /// key their penalty arms off this.
1110 pub fn pen_on(&self) -> bool {
1111 self.penalty_last_n > 0
1112 && (self.penalty_repeat != 1.0
1113 || self.penalty_freq != 0.0
1114 || self.penalty_present != 0.0)
1115 }
1116}
1117
1118/// Which draft source a spec session is pinned to. The ENGINE-LEVEL half of
1119/// `DraftSourcePlan` (memra-gguf `model_plan.rs`, always general): the plan states what the
1120/// model DECLARES, this states what actually LOADED and therefore what the session runs.
1121/// Pinned at session creation for the session's lifetime.
1122///
1123/// Family-agnostic on purpose (lane/glm5-extract2, the DraftSource seam): glm5 is today's
1124/// consumer with NativeMtp | Dflash2; the hy3/qwen-next spec lanes select through the same
1125/// three-way law instead of re-deriving it. What each family still owns is the per-session
1126/// STATE behind the kind (see `dflash.rs`'s seam note for why that half is not a trait yet).
1127#[derive(Clone, Copy, Debug, PartialEq, Eq)]
1128pub enum DraftSourceKind {
1129 /// The model's own embedded NextN/MTP head.
1130 NativeMtp,
1131 /// A separately loaded DFlash2 block-diffusion drafter
1132 /// ([`crate::dflash::DflashDrafter`]).
1133 Dflash2,
1134}
1135
1136/// The uniform draft-source selection law. Pure — no env, no engine, no family types — so it
1137/// is CPU-gateable and so every spec family answers "which source" the same way.
1138///
1139/// THE LAW, in precedence order:
1140/// 1. A LOADED DFlash2 drafter IS the source. The operator asked for it by name (a set
1141/// drafter flag that cannot load is already a loud boot failure, never a silent
1142/// fallback), and the family's embedded head is deliberately NOT loaded for this source —
1143/// it is a full trunk layer of VRAM.
1144/// 2. Otherwise the embedded head, and only when the PLAN declares an embedded source: a
1145/// loaded head under a plan that does not declare `Embedded` is a load-path bug, not a
1146/// draft source, and it is refused by name rather than drafted from.
1147/// 3. Otherwise there is no draft source and speculative decode must refuse before drafting.
1148pub fn resolve_draft_source_kind(
1149 plan: memra_gguf::model_plan::DraftSourcePlan,
1150 embedded_head_loaded: bool,
1151 dflash_loaded: bool,
1152) -> Result<DraftSourceKind, String> {
1153 use memra_gguf::model_plan::DraftSourcePlan as P;
1154 if dflash_loaded {
1155 return Ok(DraftSourceKind::Dflash2);
1156 }
1157 if embedded_head_loaded {
1158 if plan != P::Embedded {
1159 return Err(format!(
1160 "an embedded draft head is loaded but the ModelPlan declares \
1161 draft_source={plan:?} — refused rather than drafting from a head the plan \
1162 does not claim"
1163 ));
1164 }
1165 return Ok(DraftSourceKind::NativeMtp);
1166 }
1167 Err(format!(
1168 "no draft source loaded (ModelPlan declares draft_source={plan:?}): speculative \
1169 decode has nothing to draft from"
1170 ))
1171}
1172
1173#[cfg(test)]
1174mod draft_source_kind_tests {
1175 use super::{DraftSourceKind, resolve_draft_source_kind};
1176 use memra_gguf::model_plan::DraftSourcePlan as P;
1177
1178 #[test]
1179 fn a_loaded_drafter_wins_over_a_co_loaded_embedded_head() {
1180 // The operator asked for the drafter BY NAME (a set drafter flag that cannot load is
1181 // already a loud boot failure), so it takes precedence under every plan value —
1182 // including ExternalArtifact, which is what a pack declares when the draft weights
1183 // are not in the model file.
1184 for plan in [P::Embedded, P::ExternalArtifact, P::None] {
1185 assert_eq!(
1186 resolve_draft_source_kind(plan, true, true).unwrap(),
1187 DraftSourceKind::Dflash2,
1188 "plan {plan:?}: a loaded drafter must win"
1189 );
1190 assert_eq!(
1191 resolve_draft_source_kind(plan, false, true).unwrap(),
1192 DraftSourceKind::Dflash2
1193 );
1194 }
1195 }
1196
1197 #[test]
1198 fn the_embedded_head_is_the_source_only_under_a_plan_that_claims_it() {
1199 assert_eq!(
1200 resolve_draft_source_kind(P::Embedded, true, false).unwrap(),
1201 DraftSourceKind::NativeMtp
1202 );
1203 // A head loaded under a plan that does not declare Embedded is a LOAD-PATH BUG, not a
1204 // draft source. Unreachable on glm5 today (its pack hardcodes Embedded and the head
1205 // only loads under it) — which is exactly why it is pinned here: an unreachable
1206 // refusal with no arm is an untested refusal, and the next family is the one that
1207 // makes it reachable.
1208 for plan in [P::ExternalArtifact, P::None] {
1209 let err = resolve_draft_source_kind(plan, true, false)
1210 .expect_err("a head under a non-Embedded plan must refuse");
1211 assert!(err.contains("does not claim"), "{err}");
1212 assert!(err.contains(&format!("{plan:?}")), "{err}");
1213 }
1214 }
1215
1216 #[test]
1217 fn nothing_loaded_refuses_before_drafting_and_names_the_plan() {
1218 for plan in [P::Embedded, P::ExternalArtifact, P::None] {
1219 let err =
1220 resolve_draft_source_kind(plan, false, false).expect_err("no source must refuse");
1221 assert!(err.contains("no draft source loaded"), "{err}");
1222 assert!(err.contains(&format!("{plan:?}")), "{err}");
1223 }
1224 }
1225}
1226
1227/// `MEMRA_SPEC_PMIN` break semantics over per-slot draft confidences (the chain break this
1228/// module's drafting loops apply inline: `p < p_min && (j > 0 || pmin0)`): keep the longest
1229/// prefix whose every slot clears `p_min`; slot 0 survives a miss unless PMIN0 arms
1230/// zero-draft rounds. Prefix truncation is forced by the accept rule anyway (a kept slot
1231/// after a dropped one could never commit — the dspark confidence-slot argument). Pure so
1232/// the rule is CPU-gateable; the SHARED K-policy surface every spec family consumes
1233/// (hoisted from the glm5 loop, lane/glm5-extract-general).
1234pub fn spec_conf_keep(q: &[f32], p_min: f32, pmin0: bool) -> usize {
1235 if p_min <= 0.0 {
1236 return q.len();
1237 }
1238 let mut kept = 0usize;
1239 for (j, &qj) in q.iter().enumerate() {
1240 if qj < p_min && (j > 0 || pmin0) {
1241 break;
1242 }
1243 kept += 1;
1244 }
1245 kept
1246}
1247
1248/// Host Philox4x32-10 uniform in (0,1) — mirrors spec_sample.cu's `philox4`/`u01` with the
1249/// ctr_lo tag 0xFFFF_FFFE, so the host accept-test stream never collides with any device
1250/// sampling event (device Gumbel uses (i>>2, stream_pos); device residual uses 0xFFFF_FFFD).
1251/// One value per (seed, ctr) EVENT; callers own the counter discipline. Extracted verbatim
1252/// from generate_spec_inner2's closure for the dspark sampled-admission walk (the two paths
1253/// MUST consume the identical stream construction — two ad-hoc Philox copies drifting apart
1254/// is a distributional bug, not a style problem).
1255pub(crate) fn host_u01(seed: u64, ctr: u32) -> f32 {
1256 let (m0, m1) = (0xD2511F53u32, 0xCD9E8D57u32);
1257 let (mut c0, mut c1, mut c2, mut c3) = (0xFFFF_FFFEu32, ctr, 0u32, 0u32);
1258 let (mut k0, mut k1) = ((seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
1259 for _ in 0..10 {
1260 let (h0, l0) = (((m0 as u64 * c0 as u64) >> 32) as u32, m0.wrapping_mul(c0));
1261 let (h1, l1) = (((m1 as u64 * c2 as u64) >> 32) as u32, m1.wrapping_mul(c2));
1262 let (n0, n1, n2, n3) = (h1 ^ c1 ^ k0, l1, h0 ^ c3 ^ k1, l0);
1263 c0 = n0;
1264 c1 = n1;
1265 c2 = n2;
1266 c3 = n3;
1267 k0 = k0.wrapping_add(0x9E3779B9);
1268 k1 = k1.wrapping_add(0xBB67AE85);
1269 }
1270 (c0 as f32 + 1.0) * (1.0 / 4294967296.0)
1271}
1272
1273/// Tracked draft positions for [`SpecTelemetry`] (serve K defaults to 3; the run-spec gate
1274/// sweeps K=1..8, and MEMRA_SPEC_CAPMAX defaults to 7 — 8 covers every tuned config).
1275pub const SPEC_TELEM_POS: usize = 8;
1276
1277/// Always-on per-draft-position acceptance telemetry (lane/accept-telemetry, 2026-08-05 —
1278/// the llama.cpp #26389 / vLLM spec-decode counter schema, upstream-sweeps 2026-08-05).
1279/// Lives on the [`SpecSession`] and accumulates across bursts; the serve worker diffs a
1280/// stashed copy per burst for its per-model /metrics aggregation and per-request usage.
1281/// Same normalization as the `[spec-stats]` line: p-min-discarded chain tokens are counted
1282/// in NEITHER drafted nor accepted.
1283#[derive(Clone, Copy, Default, Debug)]
1284pub struct SpecTelemetry {
1285 /// verify rounds completed (a round-stream burst counts each of its M rounds).
1286 pub rounds: u64,
1287 /// tokens drafted / accepted across all rounds.
1288 pub drafted: u64,
1289 pub accepted: u64,
1290 /// how often draft position j (0-based within a round's chain) was offered / accepted.
1291 /// Positions >= SPEC_TELEM_POS are untracked (totals still count them). The opt-in
1292 /// round-stream arm (MEMRA_SPEC_STREAM=1) reads back only totals, so under it these
1293 /// arrays cover the standard-path rounds only and their sums may undercount the totals.
1294 pub pos_drafted: [u64; SPEC_TELEM_POS],
1295 pub pos_accepted: [u64; SPEC_TELEM_POS],
1296}
1297
1298impl SpecTelemetry {
1299 /// Fieldwise `self - prev` — the worker's per-burst delta off a copy stashed before the
1300 /// burst call. Saturating: a caller diffing against the wrong snapshot gets zeros, not
1301 /// a wrapped counter.
1302 pub fn delta_since(&self, prev: &SpecTelemetry) -> SpecTelemetry {
1303 let mut d = SpecTelemetry {
1304 rounds: self.rounds.saturating_sub(prev.rounds),
1305 drafted: self.drafted.saturating_sub(prev.drafted),
1306 accepted: self.accepted.saturating_sub(prev.accepted),
1307 ..Default::default()
1308 };
1309 for j in 0..SPEC_TELEM_POS {
1310 d.pos_drafted[j] = self.pos_drafted[j].saturating_sub(prev.pos_drafted[j]);
1311 d.pos_accepted[j] = self.pos_accepted[j].saturating_sub(prev.pos_accepted[j]);
1312 }
1313 d
1314 }
1315 /// Fieldwise `self += d` — the worker's per-model aggregation.
1316 pub fn merge(&mut self, d: &SpecTelemetry) {
1317 self.rounds += d.rounds;
1318 self.drafted += d.drafted;
1319 self.accepted += d.accepted;
1320 for j in 0..SPEC_TELEM_POS {
1321 self.pos_drafted[j] += d.pos_drafted[j];
1322 self.pos_accepted[j] += d.pos_accepted[j];
1323 }
1324 }
1325
1326 /// Mean accepted draft-prefix length per verify round (tau).
1327 pub fn tau(&self) -> f64 {
1328 if self.rounds > 0 {
1329 self.accepted as f64 / self.rounds as f64
1330 } else {
1331 0.0
1332 }
1333 }
1334}
1335
1336/// Session-lifetime atomic acceptance counters. The verifier records only after the greedy or
1337/// rejection-sampling walk has resolved on the host, so these relaxed increments add no GPU
1338/// launch, synchronization, allocation, or ordering dependency to the numeric path.
1339struct SpecTelemetryCounters {
1340 rounds: AtomicU64,
1341 drafted: AtomicU64,
1342 accepted: AtomicU64,
1343 pos_drafted: [AtomicU64; SPEC_TELEM_POS],
1344 pos_accepted: [AtomicU64; SPEC_TELEM_POS],
1345}
1346
1347impl Default for SpecTelemetryCounters {
1348 fn default() -> Self {
1349 Self {
1350 rounds: AtomicU64::new(0),
1351 drafted: AtomicU64::new(0),
1352 accepted: AtomicU64::new(0),
1353 pos_drafted: std::array::from_fn(|_| AtomicU64::new(0)),
1354 pos_accepted: std::array::from_fn(|_| AtomicU64::new(0)),
1355 }
1356 }
1357}
1358
1359impl SpecTelemetryCounters {
1360 fn record_round(&self, drafted: usize, accepted: usize) {
1361 debug_assert!(accepted <= drafted);
1362 self.rounds.fetch_add(1, Ordering::Relaxed);
1363 self.drafted.fetch_add(drafted as u64, Ordering::Relaxed);
1364 self.accepted.fetch_add(accepted as u64, Ordering::Relaxed);
1365 for counter in self.pos_drafted.iter().take(drafted) {
1366 counter.fetch_add(1, Ordering::Relaxed);
1367 }
1368 for counter in self.pos_accepted.iter().take(accepted) {
1369 counter.fetch_add(1, Ordering::Relaxed);
1370 }
1371 }
1372
1373 /// Round-stream keeps each round's accept length on device; retain exact scalar totals while
1374 /// leaving the per-position arrays untouched, matching the pre-existing telemetry contract.
1375 fn record_totals(&self, rounds: usize, drafted: usize, accepted: usize) {
1376 self.rounds.fetch_add(rounds as u64, Ordering::Relaxed);
1377 self.drafted.fetch_add(drafted as u64, Ordering::Relaxed);
1378 self.accepted.fetch_add(accepted as u64, Ordering::Relaxed);
1379 }
1380
1381 fn snapshot(&self) -> SpecTelemetry {
1382 SpecTelemetry {
1383 rounds: self.rounds.load(Ordering::Relaxed),
1384 drafted: self.drafted.load(Ordering::Relaxed),
1385 accepted: self.accepted.load(Ordering::Relaxed),
1386 pos_drafted: std::array::from_fn(|j| self.pos_drafted[j].load(Ordering::Relaxed)),
1387 pos_accepted: std::array::from_fn(|j| self.pos_accepted[j].load(Ordering::Relaxed)),
1388 }
1389 }
1390}
1391
1392pub struct SpecSession {
1393 pub(crate) cache: Cache,
1394 pub(crate) scratch: MtpScratch,
1395 /// Every token whose state the caches hold, in order (prompt turns + generated), INCLUDING
1396 /// overshoot: spec commits accepted drafts past max_new; those rows are in the caches, so the
1397 /// session must count them. Callers render output from this, not from their own echo.
1398 pub committed: Vec<u32>,
1399 /// Pre-output_norm hidden of the LAST committed row (device). None before the first turn.
1400 pub(crate) last_h: Option<CudaSlice<f32>>,
1401 /// Greedy argmax predicting the token AFTER committed.last() (from the last turn's final
1402 /// logits). Fuels empty-suffix continuation bursts (serve): the next turn emits this token
1403 /// first, feeds it, and the round loop resumes without any prime. None before the first turn.
1404 pub next_pred: Option<u32>,
1405 /// SAMPLED-SPEC stream continuity across bursts: Philox event counters persist here so a
1406 /// session's randomness never repeats between generate_spec_session calls. (0,0) at admit.
1407 pub sctr: u32,
1408 pub uctr: u32,
1409 /// PERSISTENT DRAFT-GRAPH CONTEXT (2026-08-01, the serve-burst fixed-cost fix): the captured
1410 /// draft graph(s) + every device I/O buffer they bake, carried ACROSS generate_spec_session
1411 /// calls. Before this, every serve burst re-captured the draft graph (2 warmup forwards +
1412 /// instantiate) — measured ~16ms/burst on H100 q27 (MEMRA_SPEC_BURST sweep,
1413 /// research/spec-serving-20260801). None before the first turn; error paths drop it
1414 /// (next burst recaptures — serve retires errored sessions anyway).
1415 pub(crate) draft_ctx: Option<DraftGraphCtx>,
1416 /// PENDING-CARRY across bursts (2026-08-01, the serve burst-boundary fix): the bonus token
1417 /// emitted by the last round but NOT committed to the caches. The old tail committed it with
1418 /// a solo T=1 trunk pass (+ draft fill), and the next burst's setup fed the stashed next_pred
1419 /// with ANOTHER solo pass — 2x ~11.5ms/burst measured on H100 q27 ([spec-setup] trace).
1420 /// Carrying it lets the next empty-suffix greedy burst consume it as round-0 verify col 0,
1421 /// exactly like a mid-burst full-accept boundary (no solo passes). INVARIANT: when set,
1422 /// `committed` (== cache rows) EXCLUDES this token although it was already emitted in the
1423 /// last burst's output, and `last_h` holds the hidden of the last COMMITTED row (its
1424 /// predecessor — the chain-seed/fill anchor). `next_pred` is None (unknown without the
1425 /// commit pass). Non-empty-suffix or sampled turns must flush first (spec_flush_pending);
1426 /// generate_spec_session_sampled does this at entry, and serve parks only flushed sessions.
1427 pub pending_tok: Option<u32>,
1428 /// SESSION-AFFINITY TURN CHECKPOINT (lane/session-affinity, 2026-08-05): the state at this
1429 /// turn's PROMPT-END boundary, retained so a later turn can REWIND here. See
1430 /// [`SpecCheckpoint`]. Refreshed by every non-empty prime; None until the first one, and on
1431 /// a rig too tight to hold it (a failed capture is silent — resume just isn't available).
1432 pub(crate) turn_ckpt: Option<SpecCheckpoint>,
1433 /// Session-lifetime acceptance telemetry. Relaxed atomics update at the host-side round
1434 /// accounting the loop already does — no syncs, no allocation. NOTE a
1435 /// pool-resumed session carries the PREVIOUS requests' counts; per-request consumers
1436 /// diff with [`SpecTelemetry::delta_since`] around each burst.
1437 telem: SpecTelemetryCounters,
1438 /// PREFIX-CACHE publication request (lane/spec-prefix-cache): worker sets this to the
1439 /// miss-LCP boundary before a cold burst; the prime captures at exactly that split (it must
1440 /// coincide with the burst's `prime_split` or no capture happens). One-shot: consumed by the
1441 /// prime, result lands in `boundary_captures`.
1442 pub capture_at: Option<usize>,
1443 /// The captures the last prime produced (see [`SpecBoundaryCapture`]). Worker drains them
1444 /// post-burst to assemble prefix entries. A failed capture is silent, like `turn_ckpt` —
1445 /// publication just isn't available for that request. Plural since
1446 /// lane/frspec-multiturn-cache (2026-08-21): a cold burst can capture BOTH the miss-LCP
1447 /// split (the shared-prefix class) and the stable pre-generation boundary (the
1448 /// next-turn re-render class) — one entry per stop, exactly the boundary set the plain
1449 /// prefill tick publishes/checkpoints.
1450 pub boundary_captures: Vec<SpecBoundaryCapture>,
1451 /// STABLE-BOUNDARY TURN CHECKPOINT REQUEST (lane/frspec-multiturn-cache, 2026-08-21): the
1452 /// ABSOLUTE committed-length position the next non-empty prime should capture `turn_ckpt`
1453 /// at, instead of prompt-end. The worker sets it to the STABLE PRE-GENERATION boundary
1454 /// (`plain_checkpoint_boundary` — before the live generation header the client rewrites),
1455 /// porting the 2026-08-09 plain-tier fix: a prompt-end spec checkpoint includes the
1456 /// template's live assistant-generation header (`<|im_start|>assistant\n<think>\n`), which
1457 /// the NEXT turn's re-render replaces, so `affinity_match` diverged a couple tokens below
1458 /// the checkpoint and the spec pool declined 100% of multi-turn agent traffic (measured:
1459 /// `spec-affinity: declined (history diverged at 6811 of checkpoint 6813)`,
1460 /// research/multiturn-cache-20260821 B4). One-shot, `capture_at` convention; None = legacy
1461 /// prompt-end capture.
1462 pub ckpt_at: Option<usize>,
1463 /// FAIL-SAFE (lane/step37-vram-admission-20260830, external-review corroboration): set
1464 /// by the worker on a session serving a step-OOM park REPLAY. The burst entry pre-marks
1465 /// the draft-graph fallback so the replay never re-enters the capture path — the capture
1466 /// appetite is part of what drove the card to the OOM, and a replay that recaptures
1467 /// re-runs the incident. If the eager replay still cannot fit, the bounded retry budget
1468 /// exhausts into the honest recoverable Overloaded error instead of looping.
1469 pub capture_disabled: bool,
1470}
1471impl SpecSession {
1472 /// Context capacity of the session's caches (the server's ContextFull guard).
1473 pub fn cache_max_ctx(&self) -> usize {
1474 self.cache.max_ctx
1475 }
1476 /// Read access to the live trunk cache (lane/spec-prefix-cache): the worker slices
1477 /// full-attn KV rows `[0..capture.pos)` out of it when publishing a boundary capture —
1478 /// those rows are append-only for the session's lifetime (rollbacks never truncate below
1479 /// the prime boundary), so no copy was taken at prime time.
1480 pub fn cache_ref(&self) -> &Cache {
1481 &self.cache
1482 }
1483 /// Read access to the persistent draft-scratch plane (lane/spec-on-cache-hit): the
1484 /// worker slices rows `[0..capture.pos)` when publishing a boundary capture, exactly
1485 /// like the trunk KV — draft rows below the prompt end are append-only for the
1486 /// session's lifetime (the prime fill wrote them once; rollbacks reset `len_d` to the
1487 /// committed length, never below the prime boundary, and the true-hidden refresh
1488 /// rewrites generated positions only). Returns `(k, v, k_tok_bytes, v_tok_bytes)`.
1489 /// None when the scratch is ring-backed (Step35 SWA — physical rows are not
1490 /// prefix-addressable; the prefix cache already refuses that class end to end).
1491 pub fn draft_plane_ref(&self) -> Option<(&CudaSlice<u8>, &CudaSlice<u8>, usize, usize)> {
1492 if self.scratch.kv.ring.is_some() {
1493 return None;
1494 }
1495 Some((
1496 &self.scratch.kv.k,
1497 &self.scratch.kv.v,
1498 self.scratch.kv.k_tok_bytes,
1499 self.scratch.kv.v_tok_bytes,
1500 ))
1501 }
1502 /// Snapshot the session's process-local acceptance counters for per-burst diffing.
1503 pub fn telemetry(&self) -> SpecTelemetry {
1504 self.telem.snapshot()
1505 }
1506 /// Committed position this session can REWIND to (its retained prompt-end boundary), if any.
1507 /// A request whose prompt matches `committed[..pos]` exactly can resume from here — see
1508 /// `spec_rewind_to_checkpoint`.
1509 pub fn rewind_pos(&self) -> Option<usize> {
1510 self.turn_ckpt.as_ref().map(|c| c.pos)
1511 }
1512 /// Whether every ring-backed trunk/draft row needed by the retained checkpoint is resident.
1513 pub fn rewind_is_resident(&self) -> bool {
1514 self.turn_ckpt.as_ref().is_some_and(|ckpt| {
1515 self.cache.can_rollback(&ckpt.snap, 0) && self.scratch.can_rewind_to(ckpt.pos)
1516 })
1517 }
1518 /// Is this session in the DEMOTION-READY shape (see [`SpecSession::into_demoted`])?
1519 /// `false` means a carried pending must be flushed first (`spec_flush_pending`), or the
1520 /// session has never run a turn and has no prediction to hand over.
1521 pub fn demote_ready(&self) -> bool {
1522 self.pending_tok.is_none() && self.next_pred.is_some()
1523 }
1524 /// Does this session hold a carried pending bonus (flush required before a handoff/park)?
1525 pub fn has_pending(&self) -> bool {
1526 self.pending_tok.is_some()
1527 }
1528 /// Committed row count == cache rows (the session invariant), for the caller's own
1529 /// `fed`-length cross-check at a handoff boundary.
1530 pub fn committed_len(&self) -> usize {
1531 self.committed.len()
1532 }
1533 /// DEMOTION HANDOFF (lane/spec-gate, 2026-08-07): consume this session and hand its trunk
1534 /// cache + next-token prediction to the plain batched-decode path.
1535 ///
1536 /// WHY THIS IS EXACT (greedy). The invariant at a burst boundary is `cache.pos ==
1537 /// committed.len()`: every committed row has trunk KV + recurrent state, exactly as a plain
1538 /// tokenwise prime of the same `committed` sequence would have left it (that is the
1539 /// session-tail contract, and the same property `spec_rewind_to_checkpoint` and the reuse
1540 /// pool already rely on). `next_pred` is the argmax of the verify's logits for the LAST
1541 /// committed row — and verify-column logits are bit-identical to plain decode's logits at
1542 /// that position, because `matmul_decode_exact` bit-identity IS the basis of the greedy
1543 /// accept walk. So handing (cache, next_pred) to the batched path continues the stream from
1544 /// a state indistinguishable from one the batched path produced itself: the batched tick
1545 /// emits `next_pred`, feeds it into this same cache, and decodes on.
1546 ///
1547 /// `None` when the session is not in the handoff shape — a carried pending (its bonus row is
1548 /// NOT in the cache, so `spec_flush_pending` must commit it first) or no `next_pred` yet
1549 /// (never bursted). Callers must not force it: a half-committed cache handed to the batched
1550 /// path would silently skip a token.
1551 ///
1552 /// The MTP draft scratch, the persistent draft-graph context and the turn checkpoint are
1553 /// DROPPED here (freeing their VRAM): the batched path never drafts, and this handoff is
1554 /// one-way by design — there is no cheap symmetric re-promotion (rebuilding the draft KV
1555 /// would mean an `mtp_kv_fill` over the whole committed history).
1556 pub fn into_demoted(self) -> Option<(Cache, u32)> {
1557 if self.pending_tok.is_some() || self.cache.tainted {
1558 return None;
1559 }
1560 let np = self.next_pred?;
1561 debug_assert_eq!(
1562 self.cache.pos,
1563 self.committed.len(),
1564 "demotion handoff: cache rows != committed tokens"
1565 );
1566 Some((self.cache, np))
1567 }
1568 /// Pool-resume hook (audit Q2): clear the parked draft-graph failure memoization so a
1569 /// NEW request resuming this session gets one fresh capture chance — a transient-pressure
1570 /// capture failure must not persist for the pool's whole lifetime (the TRT #16072 class).
1571 /// Logs once iff a flag was actually set; a no-fallback resume is silent and free.
1572 pub fn reset_graph_fallback_on_resume(&mut self) {
1573 if let Some(line) = self
1574 .draft_ctx
1575 .as_mut()
1576 .and_then(|c| c.failed.reset_on_resume())
1577 {
1578 eprintln!("{line}");
1579 }
1580 }
1581}
1582
1583/// A session's PROMPT-END boundary state, the rewind target for session-affinity resume.
1584///
1585/// WHY THIS BOUNDARY, AND WHY IT IS THE ONLY ONE WORTH KEEPING. The rewrite class this lane
1586/// exists for (a client that strips `<think>` blocks out of prior assistant turns) mutates the
1587/// text the session GENERATED, never the prompt it was given. So turn N's prompt agrees with
1588/// turn N-1's committed tokens up to almost exactly where turn N-1's generation began — the
1589/// prompt-end boundary. Keeping a checkpoint there means the next turn re-primes only its own
1590/// delta (the rewritten answer + the new user turn) instead of the whole conversation.
1591///
1592/// WHAT IT MUST HOLD. Full-attn KV is append-only and position-addressed, so rewinding it is a
1593/// `len` truncation (no data). Linear-attn (GDN) conv/ssm state is mutated IN PLACE with no
1594/// position index, so it must be a real device COPY — that copy is the entire reason a spec
1595/// session could not previously rewind. The MTP draft scratch needs no copy either: its rows
1596/// below the boundary were written by this turn's fill and are never revisited (the per-round
1597/// true-hidden refresh only rewrites the CURRENT burst's committed positions), so rewinding it
1598/// is also just a `len` reset. `last_h` is the hidden of the last row below the boundary — the
1599/// predecessor-pairing anchor the next prime's fill reads for its first row.
1600///
1601/// COST: one `Cache::snapshot` per TURN, on a code path that already takes one per ROUND.
1602pub(crate) struct SpecCheckpoint {
1603 snap: crate::cache::CacheSnapshot,
1604 /// Committed length at the boundary (== cache.pos there, the session invariant).
1605 pos: usize,
1606 /// Pre-output_norm hidden of row `pos - 1`.
1607 last_h: CudaSlice<f32>,
1608}
1609
1610/// PREFIX-CACHE BOUNDARY CAPTURE (lane/spec-prefix-cache, 2026-08-14): the state a spec session
1611/// records at its cold-prime split so the WORKER can publish a cross-request prefix entry —
1612/// the commit-gated-publication port (research/cache-spec-design-20260814/PORT-PLAN.md item 1).
1613/// Only the pieces that are DESTROYED by continuing the prime need copies here: the in-place
1614/// GDN conv/ssm states (via `Cache::snapshot`, same mechanism as [`SpecCheckpoint`]) and the
1615/// boundary logits. Full-attn KV rows `[0..pos)` and draft-scratch rows `[0..pos)` are
1616/// append-only for the session's lifetime (rollbacks never truncate below the prime boundary),
1617/// so the worker slices those from the live caches post-burst instead of copying at prime time.
1618pub struct SpecBoundaryCapture {
1619 pub snap: crate::cache::CacheSnapshot,
1620 /// Token boundary (== cache.pos at capture; == the worker's miss-LCP split).
1621 pub pos: usize,
1622 /// Full-vocab logits after the prefix prime — the entry's boundary logits.
1623 pub logits: Vec<f32>,
1624 /// Pre-output_norm trunk hidden of row `pos - 1` (lane/spec-on-cache-hit): the
1625 /// predecessor-pairing anchor a RESTORED spec session's first suffix-fill row reads
1626 /// (the `SpecSession::last_h` convention). Empty = unavailable (capture stays valid;
1627 /// the fill's zeros row-0 fallback covers it at a bounded acceptance cost).
1628 pub last_h: Vec<f32>,
1629 /// Per-layer latent boundary tails (lane/glm5-prefix-latent2, 2026-09-01): the
1630 /// generation-destroyed slice of each MLA/DSA layer's boundary state, captured eagerly
1631 /// so the worker's DEFERRED publication can slice the append-only planes from the live
1632 /// cache (`LatentKvLayer::snapshot_plane_at`). EMPTY on every two-plane model — the
1633 /// pre-field captures are byte-identical; a latent-bearing cache with an EMPTY vec here
1634 /// keeps the publisher's loud refusal (the fail-closed door stays shut).
1635 pub latent_tails: Vec<Option<crate::cache::LatentTailCapture>>,
1636}
1637
1638/// D2H one hidden row out of a `[T, n_embd]` prime hidden stack — the boundary anchor a
1639/// spec boundary capture carries for later restored-session fills. Failure is silent
1640/// (`turn_ckpt` convention): the capture publishes without an anchor.
1641pub(crate) fn capture_boundary_hidden(
1642 e: &Engine,
1643 h_rows: &CudaSlice<f32>,
1644 pos: usize,
1645 n_embd: usize,
1646) -> Vec<f32> {
1647 if pos == 0 || h_rows.len() < pos * n_embd {
1648 return Vec::new();
1649 }
1650 let Ok(mut row) = e.uninit(n_embd) else {
1651 return Vec::new();
1652 };
1653 if e.copy_view_into(
1654 &mut row,
1655 0,
1656 &h_rows.slice((pos - 1) * n_embd..pos * n_embd),
1657 n_embd,
1658 )
1659 .is_err()
1660 {
1661 return Vec::new();
1662 }
1663 e.dtoh(&row).unwrap_or_default()
1664}
1665
1666/// ROLLBACK DOOR for sampled BOUNDARY tokens (lane/sampled-spec-quality, 2026-08-19).
1667/// Default ON: the token a burst emits at its own boundary is drawn from the request's
1668/// sampler. `MEMRA_SPEC_SAMPLED_BOUNDARY=0` restores the pre-lane posture (an ARGMAX at
1669/// every boundary) without touching greedy, which is byte-unaffected either way.
1670pub fn spec_sampled_boundary_on() -> bool {
1671 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1672 *ON.get_or_init(|| std::env::var("MEMRA_SPEC_SAMPLED_BOUNDARY").as_deref() != Ok("0"))
1673}
1674
1675/// ROLLBACK DOOR for SESSION-SPANNING penalty history (lane/sampled-spec-quality).
1676/// Default ON: `pen_hist` is seeded from the session's committed tail, so repetition /
1677/// frequency / presence penalties see the whole stream. `MEMRA_SPEC_PEN_SESSION=0`
1678/// restores the pre-lane posture (each burst restarts the window from its own prompt
1679/// slice, i.e. from NOTHING on a continuation burst) — and with the door shut the worker
1680/// must keep refusing penalized sampled prefix-cache restores, because the restored
1681/// session's continuation burst is handed no prompt slice at all.
1682pub fn spec_pen_session_on() -> bool {
1683 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1684 *ON.get_or_init(|| std::env::var("MEMRA_SPEC_PEN_SESSION").as_deref() != Ok("0"))
1685}
1686
1687/// ROLLBACK DOOR for extended-entry publication from a RESTORED session
1688/// (lane/sampled-spec-quality, Item 3). Default ON: a converted prefix-cache hit that fed a
1689/// suffix captures its own prompt-end boundary so the NEXT turn can hit a longer prefix.
1690/// `MEMRA_SPEC_RESTORE_REPUBLISH=0` restores the pre-lane posture (a namespace learns exactly
1691/// one boundary and never advances it). Whole-entry semantics only — the boundary is the
1692/// restored session's own prompt end, so `entry_pos != fed_len` still refuses on the way in.
1693pub fn spec_restore_republish_on() -> bool {
1694 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1695 *ON.get_or_init(|| std::env::var("MEMRA_SPEC_RESTORE_REPUBLISH").as_deref() != Ok("0"))
1696}
1697
1698/// Diagnostics: name every boundary token on stderr (`MEMRA_SPEC_BOUNDARY_TRACE=1`), with
1699/// the argmax the pre-lane code would have emitted from the same row. This is how the
1700/// lane MEASURES the boundary rate and the deviation rate instead of estimating them.
1701fn spec_boundary_trace() -> bool {
1702 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1703 *ON.get_or_init(|| std::env::var("MEMRA_SPEC_BOUNDARY_TRACE").as_deref() == Ok("1"))
1704}
1705
1706/// llama-parity floor for the penalty window when the request does not ask for a bigger
1707/// one (`repeat_last_n` default). The serve API arms `penalty_last_n = PEN_WINDOW_MAX` for any
1708/// non-identity penalty, so this floor only matters to explicit small windows and to the
1709/// CLI env path.
1710const PEN_WINDOW_FLOOR: usize = 64;
1711
1712/// CEILING on the penalty window, and it is a COST bound, not a semantic preference.
1713/// `penalize_logits_f32` (cu/spec_sample.cu) dedups on device by having thread `i` scan
1714/// `hist[0..i]`, so a pass is O(n_hist²) and it runs ~3x per verify round (the q rows, the
1715/// p column, the bonus column). The serve API uses this same bound for every non-identity
1716/// penalty so host/plain, sparse-device, and speculative sampling cannot change logits on
1717/// admission demotion. An uncapped 128k-token history would put ~1.7e10
1718/// comparisons per pass, tens of ms per round, i.e. penalties would silently destroy decode
1719/// throughput on exactly the long-context requests that most want them. 8192 keeps a pass
1720/// at ~7e7 comparisons (tens of microseconds) while still being **128x wider than the
1721/// pre-lane effective window** (64 prompt-tail tokens + whatever the current burst had
1722/// generated). A request that genuinely needs a window beyond this wants host-side dedup +
1723/// counts through a new kernel signature — a follow-up lane, named here rather than hidden.
1724/// `pub` since lane/dspark-penalized-sampled-20260821: the dspark route's accept walk and
1725/// the dspark_sample_gate binary trim their uploads with the SAME cap — a second constant
1726/// is a second thing to drift.
1727pub const PEN_WINDOW_MAX: usize = 8192;
1728
1729/// Seed a penalty window over the SESSION, not the burst (lane/sampled-spec-quality,
1730/// Item 2). The window is the last `max(penalty_last_n, 64)` tokens of
1731/// `session_committed ++ burst_prompt` — for a cold turn-1 burst (`session_committed`
1732/// empty, default `penalty_last_n`) that is byte-identically the pre-lane
1733/// `prompt.iter().rev().take(64).rev()`; for a continuation burst it is the stream the
1734/// client actually asked us to penalize, where the pre-lane code had NOTHING.
1735/// `pub` since lane/dspark-penalized-sampled-20260821: the dspark route seeds its session
1736/// window through the SAME function (one definition of "the window" across both spec
1737/// routes and the gate binary's trunk-only reference arm).
1738pub fn pen_window_seed(
1739 session_committed: &[u32],
1740 burst_prompt: &[u32],
1741 penalty_last_n: usize,
1742) -> Vec<u32> {
1743 let win = penalty_last_n.clamp(PEN_WINDOW_FLOOR, PEN_WINDOW_MAX);
1744 let take_prompt = burst_prompt.len().min(win);
1745 let take_sess = (win - take_prompt).min(session_committed.len());
1746 let mut hist = Vec::with_capacity(take_sess + take_prompt);
1747 hist.extend_from_slice(&session_committed[session_committed.len() - take_sess..]);
1748 hist.extend_from_slice(&burst_prompt[burst_prompt.len() - take_prompt..]);
1749 hist
1750}
1751
1752/// Draw a BOUNDARY token from the target distribution the request asked for
1753/// (lane/sampled-spec-quality, Item 1) — the fix for "sampled spec emits an ARGMAX token at
1754/// every burst boundary".
1755///
1756/// WHY THIS EXISTS. A spec burst's first emitted token is not produced by the accept walk:
1757/// it comes off a logits row that already exists (the prime's last row on a cold burst; the
1758/// row after the last committed token on a continuation burst; the prefix-cache entry's
1759/// boundary row on a restored one). Pre-lane that token was `argmax` in BOTH sampling
1760/// regimes, so a sampled stream took a greedy token once per burst — measured, not
1761/// estimated, in research/spec-cache-20260818/SAMPLED-QUALITY.md. At temperature > 0 the
1762/// customer asked for a sampled token, so this draws one.
1763///
1764/// THE PROGRAM IS THE FULL-ACCEPT BONUS'S PROGRAM, deliberately: penalize the row (over the
1765/// session's window), take this row's OWN filter stats (the sampfix-20260805 law — stats
1766/// from a neighbour row mis-scale every `e0` and can wipe the row to token 0), gumbel-perturb
1767/// with the session's Philox stream at `*sctr`, argmax the perturbed row. Reusing the bonus's
1768/// composition means `sample_check`'s distributional oracle covers this draw too, and the
1769/// boundary token is drawn from the same filtered/penalized `p` the accept walk targets.
1770///
1771/// THE STREAM IS THE SESSION'S, NOT A FRESH ONE. `sctr` is the caller's live counter and is
1772/// advanced by exactly one, so a boundary draw consumes the next value in the same Philox
1773/// stream the accept walk uses — never a second, independently seeded stream (which would be
1774/// a new distributional bug: two streams from one seed correlate wherever their counters
1775/// collide). That also makes a restored session's boundary draw at `sctr == 0` bit-identical
1776/// to the cold session's own first draw from the same logits row, which is what preserves the
1777/// sampled-hit lane's per-seed hit==cold byte identity.
1778#[allow(clippy::too_many_arguments)]
1779pub fn sample_boundary_token_dev(
1780 e: &Engine,
1781 logits: &CudaSlice<f32>,
1782 n_vocab: usize,
1783 sp: &SpecSampling,
1784 pen_hist: &[u32],
1785 sctr: &mut u32,
1786 site: &str,
1787) -> Result<u32, Box<dyn std::error::Error>> {
1788 debug_assert!(
1789 sp.temp > 0.0,
1790 "boundary sampling is the sampled regime only"
1791 );
1792 // Own copy: penalize_logits mutates in place and the caller's row is live state
1793 // (prime_logits back the constrained recompute; last_col_logits backs round 0's accept).
1794 let mut col = e.zeros(n_vocab)?;
1795 e.copy_into(&mut col, 0, logits, n_vocab)?;
1796 let pen_on = sp.penalty_last_n > 0
1797 && (sp.penalty_repeat != 1.0 || sp.penalty_freq != 0.0 || sp.penalty_present != 0.0);
1798 if pen_on && !pen_hist.is_empty() {
1799 // window trim mirrors the round loop's own upload (`pen_hist[w0..]`), cap included.
1800 let w0 = pen_hist
1801 .len()
1802 .saturating_sub(sp.penalty_last_n.min(PEN_WINDOW_MAX));
1803 let hist = &pen_hist[w0..];
1804 let hd = e.htod_u32_v(hist)?;
1805 e.penalize_logits(
1806 &mut col,
1807 &hd,
1808 hist.len(),
1809 sp.penalty_repeat,
1810 sp.penalty_freq,
1811 sp.penalty_present,
1812 n_vocab,
1813 )?;
1814 }
1815 let rows0 = e.htod_i32(&[0])?;
1816 let (mut th_d, mut z_d, mut mx_d) = (e.zeros(1)?, e.zeros(1)?, e.zeros(1)?);
1817 e.filter_stats(
1818 &col, n_vocab, &rows0, &mut th_d, &mut z_d, &mut mx_d, n_vocab, 1, sp.temp, sp.top_k,
1819 sp.top_p, sp.min_p,
1820 )?;
1821 let (th, mx) = (e.dtoh(&th_d)?[0], e.dtoh(&mx_d)?[0]);
1822 let mut perturb = e.zeros(n_vocab)?;
1823 e.gumbel_perturb_filtered(&col, &mut perturb, n_vocab, sp.seed, *sctr, sp.temp, mx, th)?;
1824 *sctr = sctr.wrapping_add(1);
1825 let td = e.argmax_token_device(&perturb, n_vocab)?;
1826 let tok = guard_vocab_token(
1827 e.dtoh_u32_one(&td)?,
1828 n_vocab,
1829 &format!("sampled boundary token (site={site})"),
1830 )?;
1831 if spec_boundary_trace() {
1832 // the pre-lane token, from the SAME row, so the deviation rate is measurable.
1833 let raw = e.argmax_token_device(logits, n_vocab)?;
1834 let greedy = e.dtoh_u32_one(&raw)?;
1835 eprintln!(
1836 "[spec-boundary] site={site} sampled={tok} argmax={greedy} \
1837 deviates={} temp={} sctr={}",
1838 (tok != greedy) as u8,
1839 sp.temp,
1840 sctr.wrapping_sub(1),
1841 );
1842 }
1843 Ok(tok)
1844}
1845
1846/// Host-row twin of [`sample_boundary_token_dev`] (the prime / feed / entry rows arrive as
1847/// host `Vec<f32>`).
1848#[allow(clippy::too_many_arguments)]
1849pub fn sample_boundary_token(
1850 e: &Engine,
1851 logits: &[f32],
1852 sp: &SpecSampling,
1853 pen_hist: &[u32],
1854 sctr: &mut u32,
1855 site: &str,
1856) -> Result<u32, Box<dyn std::error::Error>> {
1857 let n_vocab = logits.len();
1858 let d = e.htod(logits)?;
1859 sample_boundary_token_dev(e, &d, n_vocab, sp, pen_hist, sctr, site)
1860}
1861
1862#[derive(Clone, Copy, PartialEq, Eq, Debug)]
1863pub(crate) struct SampledGraphKey {
1864 seed: u64,
1865 temp_bits: u32,
1866 k: usize,
1867 top_k: i32,
1868 top_p_bits: u32,
1869 min_p_bits: u32,
1870 pen_on: bool,
1871}
1872
1873impl SampledGraphKey {
1874 pub(crate) fn new(
1875 seed: u64,
1876 temp: f32,
1877 k: usize,
1878 top_k: i32,
1879 top_p: f32,
1880 min_p: f32,
1881 pen_on: bool,
1882 ) -> Self {
1883 SampledGraphKey {
1884 seed,
1885 temp_bits: temp.to_bits(),
1886 k,
1887 top_k,
1888 top_p_bits: top_p.to_bits(),
1889 min_p_bits: min_p.to_bits(),
1890 pen_on,
1891 }
1892 }
1893
1894 /// The one regime the PURE-TEMP in-graph sampled chain may stand in for the eager one:
1895 /// nothing but temperature shapes `q`. Computed FROM THE KEY so the capture guard, the
1896 /// launch guard and the key can never drift apart (they were three separate expressions
1897 /// before this lane, and the launch site simply forgot to ask).
1898 pub(crate) fn pure_temp(&self) -> bool {
1899 self.top_k == 0
1900 && f32::from_bits(self.top_p_bits) >= 1.0
1901 && f32::from_bits(self.min_p_bits) <= 0.0
1902 && !self.pen_on
1903 }
1904
1905 /// Truncation filters active — the capture body needs the IN-GRAPH filter nodes
1906 /// (`filter_stats` + `gumbel_perturb_filtered_ctr`) so the draft draws from the same
1907 /// filtered distribution the accept test reconstructs. Meaningful only when
1908 /// `graph_capturable`; penalties never reach a capture body.
1909 pub(crate) fn filtered(&self) -> bool {
1910 !self.pure_temp()
1911 }
1912
1913 /// May the sampled draft graph be CAPTURED (and a parked one LAUNCHED) for this regime?
1914 /// Pure-temp always; filtered regimes when the filtered-capture door is on
1915 /// (lane/step37-draft-graph-serving-20260830); penalties never — the per-round history
1916 /// cannot be baked into a graph, and composing a raw-softmax (or stale-history) draw
1917 /// with a penalized accept test is the unconditional-accept exactness bug. Computed FROM
1918 /// THE KEY for the same no-drift reason as `pure_temp`.
1919 pub(crate) fn graph_capturable(&self) -> bool {
1920 !self.pen_on && (self.pure_temp() || spec_graph_filtered_on())
1921 }
1922}
1923
1924/// Per-head captured graphs for the MULTI-HEAD MTP draft chain (step-modulo prefix-replay,
1925/// lane/step37-draft-graph-serving-20260830). The chain POLICY — which head serves step j,
1926/// how long the replayed prefix is, which stored seed feeds row r — stays HOST-SIDE in the
1927/// launch loop, exactly `mtp_chain_forward_dev`'s order; the graphs capture ONE head-row
1928/// forward each, on the head's OWN scratch plane:
1929/// - `interior[i]`: head i, `with_head=false` — KV append + carrier only. Interior rows'
1930/// logits are dead in the eager chain too (`mtp_chain_forward_dev` keeps only the last
1931/// row), so skipping the head matmul changes no consumed byte and removes the eager
1932/// chain's per-replay-row full-vocab matmul.
1933/// - `last[i]`: head i, `with_head=true` + the mode's tail (greedy argmax, or the sampled
1934/// gumbel draw — filtered in-graph when the request carries filters).
1935///
1936/// One `DraftChainGraphs` per MODE (greedy vs sampled), owning its keeper: dropping the
1937/// sampled chain on an s_key change never invalidates the greedy one.
1938struct DraftChainGraphs {
1939 interior: Vec<cudarc::driver::CudaGraph>,
1940 last: Vec<cudarc::driver::CudaGraph>,
1941 /// Never read: exists to OWN the captured graphs' backing buffers for as long as the
1942 /// graphs replay (the capture-retain law; same class as `DsparkSegGraph::_keeper`).
1943 _keeper: Vec<Box<dyn std::any::Any + Send>>,
1944}
1945
1946/// Sampled-tail capture pack for `mtp_head_forward_cap`: the persistent buffers and baked
1947/// constants of the in-graph categorical draw. `filt: None` = the PURE-TEMP body (gumbel
1948/// over the raw softmax), byte-identical to the pre-lane capture; `Some` adds the in-graph
1949/// truncation filter (`filter_stats` + `gumbel_perturb_filtered_ctr`) so the draft draws
1950/// from the same filtered distribution the accept test reconstructs
1951/// (lane/step37-draft-graph-serving-20260830).
1952struct SampledCapArgs<'a> {
1953 ctr: &'a mut CudaSlice<u32>,
1954 perturb: &'a mut CudaSlice<f32>,
1955 q_out: &'a mut CudaSlice<f32>,
1956 seed: u64,
1957 temp: f32,
1958 filt: Option<SampledCapFilter<'a>>,
1959}
1960
1961/// In-graph truncation-filter nodes: the stat slots `filter_stats` fills and the perturb
1962/// reads, plus the filter constants baked into the capture (they live in `s_key`, so a
1963/// request whose filters differ drops the parked graph before this ever goes stale).
1964struct SampledCapFilter<'a> {
1965 rows0: &'a CudaSlice<i32>,
1966 th: &'a mut CudaSlice<f32>,
1967 z: &'a mut CudaSlice<f32>,
1968 mx: &'a mut CudaSlice<f32>,
1969 top_k: i32,
1970 top_p: f32,
1971 min_p: f32,
1972}
1973
1974pub(crate) struct DraftGraphCtx {
1975 g_tok: CudaSlice<u32>,
1976 g_pos: CudaSlice<i32>,
1977 g_seed: CudaSlice<f32>,
1978 g_p: CudaSlice<f32>,
1979 g_ctr: CudaSlice<u32>,
1980 g_q: CudaSlice<f32>,
1981 g_perturb: CudaSlice<f32>,
1982 /// IN-GRAPH filter-stat slots (filtered sampled capture): `filter_stats` writes
1983 /// (th, z, mx) here inside the graph; `gumbel_perturb_filtered_ctr` reads (mx, th) from
1984 /// the same slots. Persistent so the baked pointers survive replays. `g_rows0` is the
1985 /// constant row-index-0 the single-row `filter_stats` launch reads (a captured memcpy
1986 /// source must not be a host temporary).
1987 g_rows0: CudaSlice<i32>,
1988 g_th: CudaSlice<f32>,
1989 g_z: CudaSlice<f32>,
1990 g_mx: CudaSlice<f32>,
1991 q_slots: Vec<CudaSlice<f32>>,
1992 /// DRAFT-SIDE GRAMMAR MASK (lane/draft-mask): packed allowed-set words over the DRAFT
1993 /// head's vocab, at a STABLE address so the captured draft graph's mask node reads the
1994 /// per-position contents the host re-uploads before each replay (the graph-promote
1995 /// pattern from decode.rs). Empty unless the session drafts under a grammar.
1996 g_dmask: CudaSlice<u32>,
1997 /// was `graph` captured WITH the mask node? A parked graph of the wrong shape is dropped.
1998 /// Covers the multi-head `chain` too (single-head and chain are mutually exclusive for a
1999 /// given model, so one flag serves whichever is active).
2000 graph_masked: bool,
2001 graph: Option<cudarc::driver::CudaGraph>,
2002 graph_s: Option<cudarc::driver::CudaGraph>,
2003 /// Multi-head chain graphs (see [`DraftChainGraphs`]): greedy and sampled chains, the
2004 /// chain twins of `graph` / `graph_s`. `chain_s`'s capture identity is `s_key` (shared
2005 /// with `graph_s` — a session is either single-head or chain, never both), and it obeys
2006 /// the same drop rules (key mismatch, penalty regime, mask-shape change).
2007 chain: Option<DraftChainGraphs>,
2008 chain_s: Option<DraftChainGraphs>,
2009 /// Failed-capture memoization for both graphs — LOUD on flip, cleared on pool resume
2010 /// (audit Q2, the TRT #16072 silent-permanent-coverage-loss class).
2011 failed: DraftGraphFallback,
2012 /// Capture identity of `graph_s` — see [`SampledGraphKey`]. `None` iff no sampled graph is
2013 /// parked; a request whose key differs drops the parked graph (and its q slots/keeper).
2014 s_key: Option<SampledGraphKey>,
2015 /// CAPTURE-RETAIN keepers (#68 root cause, 2026-08-04): the warmup-run transients whose
2016 /// pool addresses the captured graph(s) bake. Without these, the transients return to the
2017 /// pool at capture-body exit and later work (burst-boundary prime/fill/commit passes, or a
2018 /// co-served session in the worker) reuses those addresses — the persisted graph's replay
2019 /// then reads/writes live unrelated buffers (exactness corruption, first seen as the ST
2020 /// serve-spec 4B graph-arm corruption; one-shot CLI calls never re-shuffled the pool, which
2021 /// is why run-spec K=1..8 passed on the same checkpoint). Same fix class as
2022 /// capture_graph_retained's gemma/decode.rs sites — hold as long as the graph replays.
2023 keeper: Vec<Box<dyn std::any::Any + Send>>,
2024 keeper_s: Vec<Box<dyn std::any::Any + Send>>,
2025}
2026
2027/// Failed-capture memoization for the two draft graphs (audit Q2, 2026-08-05 — the
2028/// TRT #16072 trap class: pressure-triggered, silent, long-lived coverage loss).
2029///
2030/// Three contracts:
2031/// - LOUD FLIP: `mark_*` returns the warn line exactly on the false→true transition
2032/// (returned, not printed, so the once-per-flip contract is unit-testable); the caller
2033/// `eprintln!`s it UNCONDITIONALLY — a dropped draft graph is never silent. Re-marking
2034/// an already-failed graph returns None (the per-burst memoization that keeps the eager
2035/// fallback from paying a doomed capture attempt every burst).
2036/// - RESET ON RESUME: `reset_on_resume` clears both flags — a parked session resumed by a
2037/// NEW request gets one fresh capture chance instead of carrying a transient-pressure
2038/// failure for the pool's whole lifetime. Returns the note line only when a flag was
2039/// actually set (quiet on the common clean-resume path).
2040/// - Shape-change clears (`clear_*`) stay silent, exactly as before: they precede a fresh
2041/// capture attempt whose own failure would re-flip loudly.
2042#[derive(Default)]
2043pub(crate) struct DraftGraphFallback {
2044 greedy: bool,
2045 sampled: bool,
2046}
2047impl DraftGraphFallback {
2048 fn mark_greedy(&mut self, reason: &str) -> Option<String> {
2049 if self.greedy {
2050 return None;
2051 }
2052 self.greedy = true;
2053 Some(format!(
2054 "[spec] WARN: draft-graph capture failed ({reason}); eager fallback until session resume"
2055 ))
2056 }
2057 fn mark_sampled(&mut self, reason: &str) -> Option<String> {
2058 if self.sampled {
2059 return None;
2060 }
2061 self.sampled = true;
2062 Some(format!(
2063 "[spec] WARN: sampled draft-graph capture failed ({reason}); eager fallback until session resume"
2064 ))
2065 }
2066 fn greedy_failed(&self) -> bool {
2067 self.greedy
2068 }
2069 fn sampled_failed(&self) -> bool {
2070 self.sampled
2071 }
2072 fn clear_greedy(&mut self) {
2073 self.greedy = false;
2074 }
2075 fn clear_sampled(&mut self) {
2076 self.sampled = false;
2077 }
2078 /// Pool-resume reset: both graphs get a fresh capture chance. Some(note) iff any flag
2079 /// was set (so clean resumes stay quiet).
2080 pub(crate) fn reset_on_resume(&mut self) -> Option<String> {
2081 if !self.greedy && !self.sampled {
2082 return None;
2083 }
2084 let which = match (self.greedy, self.sampled) {
2085 (true, true) => "greedy+sampled",
2086 (true, false) => "greedy",
2087 _ => "sampled",
2088 };
2089 self.greedy = false;
2090 self.sampled = false;
2091 Some(format!(
2092 "[spec] draft-graph fallback reset on session resume ({which}); recapture eligible"
2093 ))
2094 }
2095}
2096
2097impl DraftGraphCtx {
2098 fn new(e: &Engine, n_embd: usize, qlen: usize) -> Result<Self, Box<dyn std::error::Error>> {
2099 Ok(DraftGraphCtx {
2100 g_tok: e.alloc_u32_zeroed(1)?,
2101 g_pos: e.htod_i32(&[0])?,
2102 g_seed: e.zeros(n_embd)?,
2103 g_p: e.zeros(1)?,
2104 g_ctr: e.alloc_u32_zeroed(1)?,
2105 g_q: e.zeros(qlen)?,
2106 g_perturb: e.zeros(qlen)?,
2107 g_rows0: e.htod_i32(&[0])?,
2108 g_th: e.zeros(1)?,
2109 g_z: e.zeros(1)?,
2110 g_mx: e.zeros(1)?,
2111 q_slots: Vec::new(),
2112 g_dmask: e.alloc_u32_zeroed(1)?,
2113 graph_masked: false,
2114 graph: None,
2115 graph_s: None,
2116 chain: None,
2117 chain_s: None,
2118 failed: DraftGraphFallback::default(),
2119 s_key: None,
2120 keeper: Vec::new(),
2121 keeper_s: Vec::new(),
2122 })
2123 }
2124}
2125
2126pub(crate) struct MtpScratch {
2127 kv: KvLayer,
2128 /// Logical row capacity. On the graph/DC draft path it also doubles as fa_decode_dc's
2129 /// bucket_max: n_splits is sized from it ONCE, so the graph captured at round 0 stays valid
2130 /// for every later t_kv. Step35 refuses that path and may back this logical extent with the
2131 /// smaller host-indexed SWA ring instead.
2132 cap: usize,
2133 extra: Vec<MtpScratchPlane>,
2134}
2135
2136struct MtpScratchPlane {
2137 kv: KvLayer,
2138 cap: usize,
2139}
2140
2141fn mtp_scratch_layout(
2142 cfg: &memra_gguf::config::ModelConfig,
2143 geom: Option<&crate::hybrid::DraftGeom>,
2144) -> (usize, usize, usize, usize) {
2145 // Student draft heads carry fewer KV heads (head_dim unchanged) -> smaller scratch rows.
2146 let n_head_kv = geom.map(|g| g.n_head_kv).unwrap_or(cfg.n_head_kv as usize);
2147 let head_dim_k = cfg.head_dim_k as usize;
2148 let head_dim_v = cfg.head_dim_v as usize;
2149 assert!(
2150 head_dim_k.is_multiple_of(32) && head_dim_v.is_multiple_of(32),
2151 "KVQUANT requires head_dim%32==0 (MTP scratch)"
2152 );
2153 let kv_dim_k = head_dim_k * n_head_kv;
2154 let kv_dim_v = head_dim_v * n_head_kv;
2155 // The fp8-KV arm deliberately does not reach the draft scratch; keep the exact format
2156 // policy shared with `MtpScratch::new` so admission scales the same allocation.
2157 let (kbb, vbb) = crate::kv_blk_bytes();
2158 let k_tok_bytes = (kv_dim_k / 32) * kbb;
2159 let v_tok_bytes = (kv_dim_v / 32) * vbb;
2160 (kv_dim_k, kv_dim_v, k_tok_bytes, v_tok_bytes)
2161}
2162
2163fn mtp_chain_head_index(step: usize, head_count: usize) -> usize {
2164 assert!(head_count > 0, "MTP chain requires at least one head");
2165 step % head_count
2166}
2167
2168impl MtpScratch {
2169 fn alloc_plane(
2170 e: &Engine,
2171 cfg: &memra_gguf::config::ModelConfig,
2172 plan: &memra_gguf::model_plan::ModelPlan,
2173 cap: usize,
2174 geom: Option<&crate::hybrid::DraftGeom>,
2175 ) -> Result<MtpScratchPlane, Box<dyn std::error::Error>> {
2176 let (kv_dim_k, kv_dim_v, k_tok_bytes, v_tok_bytes) = mtp_scratch_layout(cfg, geom);
2177 let ring = if crate::cache::swa_ring_on()
2178 && crate::plan_backend::decode_batch_program(plan)
2179 == crate::plan_backend::DecodeBatchProgram::SlidingGatedMoe
2180 {
2181 let window = plan
2182 .layers
2183 .iter()
2184 .find_map(|layer| match layer.attention {
2185 memra_gguf::model_plan::AttentionPlan::SlidingWindow { window, .. } => {
2186 Some(window as usize)
2187 }
2188 _ => None,
2189 })
2190 .ok_or("sliding-gated-MoE draft scratch has no sliding-window layer")?;
2191 Some(crate::cache::KvRing::new(
2192 crate::cache::swa_ring_rows(window, cap),
2193 window,
2194 ))
2195 } else {
2196 None
2197 };
2198 let alloc_rows = ring.as_ref().map(crate::cache::KvRing::rows).unwrap_or(cap);
2199 // Ring-backed planes arm the device base mirror for the dcw draft arm (see
2200 // KvLayer::base_d): the captured chain derives its physical rows from
2201 // (len_d, base_d, window) with zero per-token node updates.
2202 let base_d = match ring.as_ref() {
2203 Some(_) => Some(e.htod_i32(&[0])?),
2204 None => None,
2205 };
2206 Ok(MtpScratchPlane {
2207 kv: KvLayer {
2208 k: e.alloc_u8(alloc_rows * k_tok_bytes)?,
2209 v: e.alloc_u8(alloc_rows * v_tok_bytes)?,
2210 kv_dim_k,
2211 kv_dim_v,
2212 k_tok_bytes,
2213 v_tok_bytes,
2214 len: 0,
2215 ring,
2216 len_d: e.htod_i32(&[0])?,
2217 base_d,
2218 },
2219 cap,
2220 })
2221 }
2222
2223 fn new(
2224 e: &Engine,
2225 cfg: &memra_gguf::config::ModelConfig,
2226 plan: &memra_gguf::model_plan::ModelPlan,
2227 cap: usize,
2228 geom: Option<&crate::hybrid::DraftGeom>,
2229 ) -> Result<Self, Box<dyn std::error::Error>> {
2230 // Trunk KV block bytes (34/24, q8_0/q5_1). The scratch keeps the baseline numerics and
2231 // its append/fa pass g=false (the e4m3 arms drifted draft acceptance 69-88% -> 46% in the
2232 // 2026-07-12 A/B, and the trunk fp8 door itself was removed 2026-09-05).
2233 let primary = Self::alloc_plane(e, cfg, plan, cap, geom)?;
2234 Ok(MtpScratch {
2235 kv: primary.kv,
2236 cap: primary.cap,
2237 extra: Vec::new(),
2238 })
2239 }
2240
2241 fn push_plane(
2242 &mut self,
2243 e: &Engine,
2244 cfg: &memra_gguf::config::ModelConfig,
2245 plan: &memra_gguf::model_plan::ModelPlan,
2246 geom: Option<&crate::hybrid::DraftGeom>,
2247 ) -> Result<(), Box<dyn std::error::Error>> {
2248 self.extra
2249 .push(Self::alloc_plane(e, cfg, plan, self.cap, geom)?);
2250 Ok(())
2251 }
2252
2253 fn plane_count(&self) -> usize {
2254 1 + self.extra.len()
2255 }
2256
2257 fn plane(&self, index: usize) -> (&KvLayer, usize) {
2258 if index == 0 {
2259 (&self.kv, self.cap)
2260 } else {
2261 let plane = &self.extra[index - 1];
2262 (&plane.kv, plane.cap)
2263 }
2264 }
2265
2266 fn plane_mut(&mut self, index: usize) -> (&mut KvLayer, usize) {
2267 if index == 0 {
2268 (&mut self.kv, self.cap)
2269 } else {
2270 let plane = &mut self.extra[index - 1];
2271 (&mut plane.kv, plane.cap)
2272 }
2273 }
2274
2275 // #[track_caller]: set_len/set_plane_len have eight call sites (checkpoint restore, spec
2276 // rollback, session grow, seed replay ...) and the lap failure needs to say WHICH one, not
2277 // just that a rewind was refused.
2278 #[track_caller]
2279 fn set_plane_len(
2280 &mut self,
2281 e: &Engine,
2282 index: usize,
2283 n: usize,
2284 ) -> Result<(), Box<dyn std::error::Error>> {
2285 let caller = std::panic::Location::caller();
2286 let (kv, cap) = self.plane_mut(index);
2287 if let Some(ring) = kv.ring.as_ref()
2288 && !ring.can_rewind_to(n)
2289 {
2290 // NAME THE NUMBERS (2026-08-28). This error is a step37 serving blocker on the
2291 // vendor-default shape and it fires from more than one call path with more than
2292 // one trigger: a long generation walks the checkpoint out of the ring, but a
2293 // ~4.5k-token prompt also fails within 5 s of prime, which accumulation cannot
2294 // explain. A bare message forced two rounds of guessing; the operands make each
2295 // trigger name itself.
2296 let raw = n.saturating_sub(ring.window().saturating_sub(1));
2297 return Err(format!(
2298 "SWA ring MTP checkpoint has been lapped; full re-prime required (plane={index} rewind_to={n} window={} base={} rows={} cap={cap} needed_view_start={} < base, called from {caller})",
2299 ring.window(),
2300 ring.base(),
2301 ring.rows(),
2302 raw & !31usize,
2303 )
2304 .into());
2305 }
2306 kv.len = n;
2307 e.set_i32_one(&mut kv.len_d, n as i32)
2308 }
2309
2310 /// Set BOTH length counters: the host mirror AND the device len_d the captured append/fa read
2311 /// (a 4-byte in-place htod — the counter pointer is baked into the graph, never realloc'd).
2312 /// This is the ONLY truncation/rollback mechanism the persistent draft KV needs.
2313 #[track_caller]
2314 fn set_len(&mut self, e: &Engine, n: usize) -> Result<(), Box<dyn std::error::Error>> {
2315 let caller = std::panic::Location::caller();
2316 if !self.can_rewind_to(n) {
2317 // set_plane_len re-checks and reports the operands; call it so the failure carries
2318 // which plane refused and why, instead of this bare aggregate.
2319 for index in 0..self.plane_count() {
2320 self.set_plane_len(e, index, n)?;
2321 }
2322 return Err(format!(
2323 "SWA ring MTP checkpoint has been lapped; full re-prime required (aggregate rewind_to={n}, no single plane reported, called from {caller})"
2324 )
2325 .into());
2326 }
2327 for index in 0..self.plane_count() {
2328 self.set_plane_len(e, index, n)?;
2329 }
2330 Ok(())
2331 }
2332
2333 fn can_rewind_to(&self, n: usize) -> bool {
2334 (0..self.plane_count()).all(|index| {
2335 self.plane(index)
2336 .0
2337 .ring
2338 .as_ref()
2339 .is_none_or(|ring| ring.can_rewind_to(n))
2340 })
2341 }
2342
2343 /// Pre-arm ring headroom for `rows` upcoming DEVICE-COUNTER appends (the dcw draft arm):
2344 /// a captured chain cannot rebase mid-replay, so any rebase the coming appends could need
2345 /// happens HERE, host-side, before the capture warmups or the round's replays (the rebase
2346 /// arm of `prepare_kv_append` also refreshes the plane's `base_d` device mirror). No-op on
2347 /// flat planes and when the ring already has room; `len` is untouched either way.
2348 fn ensure_dcw_headroom(
2349 &mut self,
2350 e: &Engine,
2351 rows: usize,
2352 ) -> Result<(), Box<dyn std::error::Error>> {
2353 for index in 0..self.plane_count() {
2354 let (kv, _) = self.plane_mut(index);
2355 let Some(ring) = kv.ring.as_ref() else {
2356 continue;
2357 };
2358 let retain = memra_kv::swa_retain_from(kv.len, ring.window(), ring.base());
2359 e.prepare_kv_append(kv, retain, rows)?;
2360 }
2361 Ok(())
2362 }
2363}
2364
2365/// Retained verify intermediates for the REPLAY-FREE partial accept (2026-07-03, the profiled
2366/// #1 spec cost at long ctx: the partial-accept replay was a DUPLICATE trunk pass — ~0.54 extra
2367/// full weight reads per round — recomputing columns the verify had already produced
2368/// bit-identically). Holds, per linear layer, everything needed to rebuild its recurrent state
2369/// to "after the first j verify columns" WITHOUT re-running the trunk:
2370/// - BATCHED-path layers (`gdn`): the exact token-major inputs the round's ONE gdn_scan
2371/// consumed. A prefix re-run of the SAME kernel (t=j) from the snapshot state is bit-identical
2372/// to the first j iterations of the verify's scan — the kernel's t-loop carries state in
2373/// registers and iteration t never depends on T. `qkv_mixed` (the conv input) feeds the
2374/// pure-copy ring rebuild.
2375/// - PER-COLUMN-path layers (`cols`): dtod clones of (conv_state, ssm_state) taken after each
2376/// column 0..t-2 — pure copies of the actual chain states (the last column is never a rebuild
2377/// target: j <= t-1).
2378/// Full-attn layers need nothing: their verify KV rows are bit-identical to eager's (the
2379/// decode-exact contract; verify-probe pins it), so rollback = len truncation.
2380struct GdnStash {
2381 qkv_mixed: CudaSlice<f32>, // [t, conv_dim] token-major (conv input)
2382 q_l2: CudaSlice<f32>,
2383 k_l2: CudaSlice<f32>,
2384 v_g: CudaSlice<f32>, // [t, num_v, d_state]
2385 g_log: CudaSlice<f32>,
2386 beta: CudaSlice<f32>, // [t, num_v]
2387}
2388pub(crate) struct VerifyCkpt {
2389 gdn: Vec<Option<GdnStash>>, // [n_layer], Some iff batched linear path ran
2390 #[allow(clippy::type_complexity)]
2391 // allow: one-shot composite type; naming it would hide the shape that matters at the call site
2392 cols: Vec<Option<Vec<(CudaSlice<f32>, CudaSlice<f32>)>>>, // [n_layer][col] = (conv, ssm) after col
2393}
2394/// Opaque handle for the dspark round (dflash.rs) — VerifyCkpt stays spec-private.
2395pub(crate) struct DsparkVerifyCkpt(VerifyCkpt);
2396
2397/// Engine-bundle slice 3 (DSF-ROUNDCOST-20260820 §2 row 4 / §5 rank 1): bucketed CUDA
2398/// graphs for the dspark verify's LINEAR-layer segments. The measured verify is ~2,800
2399/// eager launches whose residual cost is DEVICE-side per-launch overhead (slice 2 proved
2400/// host dispatch is not the binder: fully-deferred dispatch bought ~0 wall). The 48 GDN
2401/// layers between full-attention layers are shape-static given vt — no positions, no
2402/// t_kv, state addressed through pointer tables — so runs of them capture per
2403/// (segment, vt) and replay as ONE graph launch each. Full-attention layers stay eager
2404/// (their per-row append/fa arm picks are t_kv-driven — the exec-update extension).
2405///
2406/// Per round out-of-graph: one pointer-table refresh (gdn ping-pong moves the canonical
2407/// handles), one input-staging copy per segment, host parity bookkeeping. Captured via
2408/// `capture_graph_retained` (2 warmups + capture, keeper retains warmup transients so
2409/// pool addresses stay stable); the warmups EXECUTE, so segment conv/ssm state is saved
2410/// before and restored after — the graph's first real launch starts from the exact
2411/// pre-round state. The ckpt column stash rides persistent slabs (written inside the
2412/// graph as memcpy nodes); commit reads them via `dspark_commit_prefix_slab`.
2413/// `MEMRA_DSPARK_VERIFY_GRAPH=0` reverts to the eager walk (byte-identical body).
2414pub(crate) struct DsparkVerifyGraphs {
2415 /// Linear-attention layer indices ascending; `lin_pos[il]` = index into the vecs.
2416 lin: Vec<usize>,
2417 lin_pos: std::collections::HashMap<usize, usize>,
2418 /// [n_lin x 6] pointer table (conv, s0, s1, conv, s1, s0 per layer), refreshed per
2419 /// verify from the live handles; layer il's slice starts at lin_pos[il]*6.
2420 table_all: CudaSlice<u64>,
2421 host_table: Vec<u64>,
2422 /// Persistent per-layer ckpt stash slabs: row r of the verify at slab offset
2423 /// r*words. Shared by every (segment, vt) bucket — one verify runs at a time.
2424 stash_conv: Vec<CudaSlice<f32>>,
2425 stash_ssm: Vec<CudaSlice<f32>>,
2426 conv_words: usize,
2427 ssm_words: usize,
2428 /// Per-vt input/output staging (stable addresses the graphs bake).
2429 stage: std::collections::HashMap<usize, (CudaSlice<f32>, CudaSlice<f32>)>,
2430 /// Per-vt dflash tap-sink buffers — the captured segments bake the tap dst address,
2431 /// so the sink buffer must live (and persist) with the graphs, not with the round.
2432 pub(crate) tap_bufs: std::collections::HashMap<usize, CudaSlice<f32>>,
2433 graphs: std::collections::HashMap<(usize, usize), DsparkSegGraph>,
2434 /// Warmup-corruption guard scratch: pre-capture conv/ssm of every linear layer
2435 /// (sized n_lin — the slice-4c full-verify warmups execute the whole walk).
2436 save_conv: CudaSlice<f32>,
2437 save_ssm: CudaSlice<f32>,
2438 max_run: usize,
2439 n_embd: usize,
2440 /// Set by the verify walk: this round's linear ckpt lives in the slabs (the caller
2441 /// commits through `dspark_commit_prefix_slab` instead of the cols arm).
2442 pub(crate) round_slab: bool,
2443 // ---- slice 4c: full-verify single graph per (vt, rung) ----
2444 /// Full-attention layer indices ascending; `fa_pos[il]` = index into the vec.
2445 fa: Vec<usize>,
2446 fa_pos: std::collections::HashMap<usize, usize>,
2447 /// [n_fa x 2 x t_cap] interleaved (k,v) base-pointer pairs, refreshed per verify;
2448 /// layer il's slice starts at `fa_pos[il] * 2 * t_cap` (the seqs twins read pairs
2449 /// [2z], z < t <= t_cap, so one t_cap-sized table serves every vt).
2450 fa_table: CudaSlice<u64>,
2451 fa_host_table: Vec<u64>,
2452 t_cap: usize,
2453 /// Per-vt position staging for the captured bodies — contents refreshed per round
2454 /// (rope reads row r; the seqs twins derive append slot and T_kv per z from it).
2455 pos_stage: std::collections::HashMap<usize, CudaSlice<i32>>,
2456 /// Full-verify graphs keyed (vt, rung_end, hi).
2457 full: std::collections::HashMap<(usize, usize, usize), DsparkSegGraph>,
2458 /// Largest n with every layer in [0, n) linear or full-attention (walk coverage).
2459 covered: usize,
2460 /// Every layer in [0, n) is linear or full-attention (no MLA/unknown mixers) — the
2461 /// full-verify capture walks all of them.
2462 walk_uniform: bool,
2463 /// Last `(captures, device graph-mem reserved bytes)` reading taken by
2464 /// `HybridModel::dspark_vg_admission_debt` — the two-point base of the MARGINAL debt
2465 /// projection (see `dspark_vg_debt_projection`; a mean-based reading extrapolated the
2466 /// pool's one-time shared allocation and reserved 8.5 GB of phantom VRAM).
2467 debt_obs: Option<(usize, usize)>,
2468}
2469
2470struct DsparkSegGraph {
2471 graph: cudarc::driver::CudaGraph,
2472 _keeper: Vec<Box<dyn std::any::Any + Send>>,
2473}
2474
2475/// Per-call arguments of [`HybridModel::qwen35_tparallel_fa_layer`] — one struct so the
2476/// eager walk and the slice-4c captured full-verify graphs hand the SAME body its two
2477/// modes without a second copy of the math.
2478pub(crate) struct FaLayerArgs<'a> {
2479 /// [T] per-row positions (device): rope reads them row-indexed; the seqs twins read
2480 /// them per-z (append slot = pos, T_kv = pos + 1).
2481 pub pos_d: &'a CudaSlice<i32>,
2482 /// Verify-level lazy per-row 1-element position buffers — only the per-row fallback
2483 /// arm builds/uses them (graph mode refuses that arm).
2484 pub pos_rows: &'a mut Option<Vec<CudaSlice<i32>>>,
2485 pub pos0: usize,
2486 /// Some((kv pointer table, offset-in-u64s, rung_end)) = captured-graph mode.
2487 pub graph_cap: Option<(&'a CudaSlice<u64>, usize, usize)>,
2488 /// ROUND-STREAM (lane/draftcost-moe, v0.100 train merge): Some((token stream, device
2489 /// round counter)) routes the FA attend through the dc rows kernels and the Linear
2490 /// mixer through `linear_attn_verify_t` (the stream arms the old inline body carried).
2491 /// Never armed together with `graph_cap` (the verify-level merge guard refuses).
2492 pub stream: Option<(&'a CudaSlice<u32>, &'a CudaSlice<i32>)>,
2493 /// VerifyCkpt for the stream-Linear arm's GdnStash install; None in graph mode and
2494 /// for FA layers that never touch it.
2495 pub ckpt: Option<&'a mut VerifyCkpt>,
2496}
2497
2498// SAFETY: `CudaGraph` is not marked Send by cudarc because its raw driver handles carry
2499// no automatic trait; CUDA driver graph handles are context-scoped rather than
2500// OS-thread-affine. The ctx lives in
2501// `HybridModel::dspark_vgraphs` behind a Mutex and every touch happens on the engine's
2502// single decode-stream thread.
2503unsafe impl Send for DsparkVerifyGraphs {}
2504
2505impl DsparkVerifyGraphs {
2506 /// Live capture count (segment + full graphs) — the denominator of
2507 /// [`dspark_vg_debt_projection`]'s observed bytes/capture mean.
2508 pub(crate) fn captures(&self) -> usize {
2509 self.graphs.len() + self.full.len()
2510 }
2511
2512 /// Take the marginal-growth debt reading and record this observation for the next one.
2513 /// Called under the pool mutex by `HybridModel::dspark_vg_admission_debt`.
2514 pub(crate) fn admission_debt(&mut self, reserved_bytes: usize) -> usize {
2515 let captures = self.captures();
2516 let debt =
2517 dspark_vg_debt_projection(captures, dspark_vg_cap(), reserved_bytes, self.debt_obs);
2518 if captures > 0 {
2519 match self.debt_obs {
2520 Some((c0, _)) if captures <= c0 => {}
2521 _ => self.debt_obs = Some((captures, reserved_bytes)),
2522 }
2523 }
2524 debt
2525 }
2526
2527 /// Build for this cache's shape. None when there are no linear layers, sizes are
2528 /// non-uniform, or the trunk keeps a gemma4 config (never on the qwen35 family).
2529 pub(crate) fn new(
2530 e: &Engine,
2531 cache: &Cache,
2532 t_max: usize,
2533 n_embd: usize,
2534 ) -> Result<Option<Self>, Box<dyn std::error::Error>> {
2535 let lin: Vec<usize> = (0..cache.recur.len())
2536 .filter(|&il| cache.recur[il].is_some())
2537 .collect();
2538 if lin.is_empty() || t_max < 2 {
2539 return Ok(None);
2540 }
2541 let first = cache.recur[lin[0]].as_ref().unwrap();
2542 let (conv_words, ssm_words) = (first.conv_state.len(), first.ssm_state.len());
2543 for &il in &lin {
2544 let rl = cache.recur[il].as_ref().unwrap();
2545 if rl.conv_state.len() != conv_words || rl.ssm_state.len() != ssm_words {
2546 return Ok(None);
2547 }
2548 }
2549 let n = lin.len();
2550 let mut lin_pos = std::collections::HashMap::with_capacity(n);
2551 for (k, &il) in lin.iter().enumerate() {
2552 lin_pos.insert(il, k);
2553 }
2554 // longest run of consecutive linear layers (save-scratch sizing)
2555 let mut max_run = 1usize;
2556 let mut run = 1usize;
2557 for w in lin.windows(2) {
2558 if w[1] == w[0] + 1 {
2559 run += 1;
2560 max_run = max_run.max(run);
2561 } else {
2562 run = 1;
2563 }
2564 }
2565 let rows = t_max - 1;
2566 let mut stash_conv = Vec::with_capacity(n);
2567 let mut stash_ssm = Vec::with_capacity(n);
2568 for _ in 0..n {
2569 stash_conv.push(e.uninit(rows * conv_words)?);
2570 stash_ssm.push(e.uninit(rows * ssm_words)?);
2571 }
2572 let host_table = vec![0u64; n * 6];
2573 let table_all = e.htod_u64(&host_table)?;
2574 // slice 4c: full-attention census for the full-verify graphs.
2575 let fa: Vec<usize> = (0..cache.kv.len())
2576 .filter(|&il| cache.kv[il].is_some())
2577 .collect();
2578 let mut fa_pos = std::collections::HashMap::with_capacity(fa.len());
2579 for (k, &il) in fa.iter().enumerate() {
2580 fa_pos.insert(il, k);
2581 }
2582 let n_layers = cache.kv.len().max(cache.recur.len());
2583 // exactly one of (linear state, kv cache) per layer — no MLA/unknown mixers.
2584 let walk_uniform = (0..n_layers).all(|il| {
2585 cache.recur.get(il).is_some_and(|r| r.is_some())
2586 != cache.kv.get(il).is_some_and(|k| k.is_some())
2587 });
2588 // Contiguous covered prefix: the largest n such that every layer in [0, n) is
2589 // linear or full-attention. The TRUNK walk is [0, layers.len()) and the cache
2590 // vecs can carry EXTRA state slots past it (the q38 export keeps the MTP head
2591 // layer's kv at the tail — hi == lin+fa never held, the s4c battery's zero
2592 // 'full' captures). The full-graph guard is walk coverage, not slot arithmetic.
2593 let covered = (0..n_layers)
2594 .take_while(|il| lin_pos.contains_key(il) || fa_pos.contains_key(il))
2595 .count();
2596 let t_cap = t_max;
2597 let fa_host_table = vec![0u64; fa.len() * 2 * t_cap];
2598 let fa_table = e.htod_u64(&fa_host_table)?;
2599 Ok(Some(Self {
2600 lin,
2601 lin_pos,
2602 table_all,
2603 host_table,
2604 stash_conv,
2605 stash_ssm,
2606 conv_words,
2607 ssm_words,
2608 stage: std::collections::HashMap::new(),
2609 tap_bufs: std::collections::HashMap::new(),
2610 graphs: std::collections::HashMap::new(),
2611 save_conv: e.uninit(n * conv_words)?,
2612 save_ssm: e.uninit(n * ssm_words)?,
2613 max_run,
2614 n_embd,
2615 round_slab: false,
2616 fa,
2617 fa_pos,
2618 fa_table,
2619 fa_host_table,
2620 t_cap,
2621 pos_stage: std::collections::HashMap::new(),
2622 full: std::collections::HashMap::new(),
2623 covered,
2624 walk_uniform,
2625 debt_obs: None,
2626 }))
2627 }
2628
2629 /// Rebuild the pointer tables from the live handles (once per verify — the gdn
2630 /// ping-pong swaps the canonical/alt handles between rounds; a fresh generation's
2631 /// cache buffers land at new addresses; a stale table would read the wrong state).
2632 pub(crate) fn refresh_tables(
2633 &mut self,
2634 e: &Engine,
2635 cache: &Cache,
2636 ) -> Result<(), Box<dyn std::error::Error>> {
2637 use cudarc::driver::DevicePtr;
2638 {
2639 let s = &e.gpu.stream();
2640 for (k, &il) in self.lin.iter().enumerate() {
2641 let rl = cache.recur[il].as_ref().unwrap();
2642 let (pc, _g0) = rl.conv_state.device_ptr(s);
2643 let (p0, _g1) = rl.ssm_state.device_ptr(s);
2644 let (p1, _g2) = rl.ssm_state_alt.device_ptr(s);
2645 let o = k * 6;
2646 self.host_table[o] = pc;
2647 self.host_table[o + 1] = p0;
2648 self.host_table[o + 2] = p1;
2649 self.host_table[o + 3] = pc;
2650 self.host_table[o + 4] = p1;
2651 self.host_table[o + 5] = p0;
2652 }
2653 for (k, &il) in self.fa.iter().enumerate() {
2654 let kvl = cache.kv[il].as_ref().unwrap();
2655 let (pk, _g0) = kvl.k.device_ptr(s);
2656 let (pv, _g1) = kvl.v.device_ptr(s);
2657 let o = k * 2 * self.t_cap;
2658 for z in 0..self.t_cap {
2659 self.fa_host_table[o + 2 * z] = pk;
2660 self.fa_host_table[o + 2 * z + 1] = pv;
2661 }
2662 }
2663 }
2664 e.htod_u64_into(&self.host_table, &mut self.table_all)?;
2665 if !self.fa_host_table.is_empty() {
2666 e.htod_u64_into(&self.fa_host_table, &mut self.fa_table)?;
2667 }
2668 Ok(())
2669 }
2670
2671 /// Slice 4c eligibility: Some(rung_end) when this round can replay (or capture) a
2672 /// full-verify graph — the whole walk [lo, hi) is covered, every layer is linear or
2673 /// full-attention, and ALL of the round's per-row t_kv values take the v4-seqs arm
2674 /// on ONE `fa_split_keys` ladder step that the rung also sits on (the straddle law;
2675 /// both gates are t_kv intervals, so ends-inside means all-inside). The rung is the
2676 /// round's next power of two — grid/partial sizing only (`n_splits_max` is pure
2677 /// stride; splits >= ns_eff write the empty partial the combine never reads), so one
2678 /// captured graph is bit-identical for every round the rung covers.
2679 #[allow(clippy::too_many_arguments)]
2680 pub(crate) fn full_rung(
2681 &self,
2682 model: &crate::hybrid::HybridModel,
2683 cache: &Cache,
2684 lo: usize,
2685 hi: usize,
2686 t: usize,
2687 ) -> Option<usize> {
2688 if std::env::var("MEMRA_DSPARK_FULLG_DEBUG").as_deref() == Ok("1") {
2689 static ONCE: std::sync::Once = std::sync::Once::new();
2690 let len0 = self
2691 .fa
2692 .first()
2693 .and_then(|&il| cache.kv[il].as_ref())
2694 .map(|k| k.len);
2695 ONCE.call_once(|| {
2696 eprintln!(
2697 "[fullg-debug] walk_uniform={} covered={} fa_rows_on={} t={} lo={} hi={} lin={} fa={} t_cap={} len0={:?}",
2698 self.walk_uniform, self.covered, dspark_fa_rows_on(), t, lo, hi,
2699 self.lin.len(), self.fa.len(), self.t_cap, len0
2700 );
2701 });
2702 }
2703 if !self.walk_uniform
2704 || !dspark_fa_rows_on()
2705 || t < 2
2706 || lo != 0
2707 || hi > self.covered
2708 || t > self.t_cap
2709 || self.fa.is_empty()
2710 {
2711 return None;
2712 }
2713 let cfg = &model.cfg;
2714 let head_dim_global = cfg.head_dim_k as usize;
2715 let nkv = cfg.n_head_kv as usize;
2716 let kvl0 = cache.kv[self.fa[0]].as_ref().unwrap();
2717 // the z-batched twins read stacked rows at the cache's kv dims — must equal the
2718 // projection stride (the body's guard, hoisted so ineligible models fall back
2719 // instead of refusing mid-capture).
2720 let geom = cfg.full_attention_geometry_at(self.fa[0] as u32);
2721 let kv_dim = geom.n_head_kv as usize * geom.head_dim_k as usize;
2722 if kvl0.kv_dim_k != kv_dim || kvl0.kv_dim_v != kv_dim {
2723 return None;
2724 }
2725 let len0 = kvl0.len;
2726 let (t_kv_first, t_kv_last) = (len0 + 1, len0 + t);
2727 if !crate::fa_seqs_eligible(t_kv_first, head_dim_global)
2728 || !crate::fa_seqs_eligible(t_kv_last, head_dim_global)
2729 || crate::fa_split_keys(t_kv_first, nkv) != crate::fa_split_keys(t_kv_last, nkv)
2730 {
2731 return None;
2732 }
2733 let rung = t_kv_last.next_power_of_two().max(256);
2734 if crate::fa_split_keys(rung, nkv) != crate::fa_split_keys(t_kv_last, nkv) {
2735 return None;
2736 }
2737 Some(rung)
2738 }
2739
2740 /// Run the WHOLE verify walk [lo, hi) as one captured graph at (vt=t, rung): stage
2741 /// the residual + refresh the per-vt position staging, capture on first encounter
2742 /// (2 executing warmups bracketed by a full linear-state save/restore; KV warmup
2743 /// appends write the exact slots the replay writes — idempotent), launch, then apply
2744 /// the host bookkeeping the captured body skipped (per-linear-layer parity swap for
2745 /// odd t, per-fa-layer len bump). Returns the fresh residual.
2746 #[allow(clippy::too_many_arguments)]
2747 #[allow(clippy::map_entry)] // allow: the init body is fallible (`?`); Entry::or_insert_with cannot propagate errors
2748 pub(crate) fn run_full(
2749 &mut self,
2750 model: &crate::hybrid::HybridModel,
2751 e: &Engine,
2752 lo: usize,
2753 hi: usize,
2754 x: &CudaSlice<f32>,
2755 t: usize,
2756 pos0: usize,
2757 rung: usize,
2758 cache: &mut Cache,
2759 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2760 let n_embd = self.n_embd;
2761 if !self.stage.contains_key(&t) {
2762 let xin = e.uninit(t * n_embd)?;
2763 let xout = e.uninit(t * n_embd)?;
2764 self.stage.insert(t, (xin, xout));
2765 }
2766 if !self.pos_stage.contains_key(&t) {
2767 self.pos_stage.insert(t, e.htod_i32(&vec![0i32; t])?);
2768 }
2769 // Per-round refresh: position contents + input staging (both addresses are baked
2770 // by the captured bodies; only their CONTENTS change round to round).
2771 {
2772 let pos_host: Vec<i32> = (0..t).map(|r| (pos0 + r) as i32).collect();
2773 let pb = self.pos_stage.get_mut(&t).unwrap();
2774 e.htod_i32_into(pb, &pos_host)?;
2775 let (xin, _) = self.stage.get_mut(&t).unwrap();
2776 e.copy_into(xin, 0, x, t * n_embd)?;
2777 }
2778 let key = (t, rung, hi);
2779 if !self.full.contains_key(&key) {
2780 // The warmups EXECUTE the whole walk on live state — save every linear
2781 // layer's conv + canonical ssm first, restore after (KV needs no restore:
2782 // graph mode never bumps host lens and the appends write this round's own
2783 // slots).
2784 for (k, &il) in self.lin.iter().enumerate() {
2785 let rl = cache.recur[il].as_ref().unwrap();
2786 e.copy_into(
2787 &mut self.save_conv,
2788 k * self.conv_words,
2789 &rl.conv_state,
2790 self.conv_words,
2791 )?;
2792 e.copy_into(
2793 &mut self.save_ssm,
2794 k * self.ssm_words,
2795 &rl.ssm_state,
2796 self.ssm_words,
2797 )?;
2798 }
2799 let (graph, keeper) = {
2800 let table_all = &self.table_all;
2801 let lin_pos = &self.lin_pos;
2802 let fa_pos = &self.fa_pos;
2803 let fa_table = &self.fa_table;
2804 let t_cap = self.t_cap;
2805 let stash_conv = &mut self.stash_conv;
2806 let stash_ssm = &mut self.stash_ssm;
2807 let pos_d: &CudaSlice<i32> = &self.pos_stage[&t];
2808 let (xin, xout) = self
2809 .stage
2810 .get_mut(&t)
2811 .map(|(a, b)| (&*a, b))
2812 .expect("stage bucket created above");
2813 let cache_ref: &mut Cache = cache;
2814 let iflag = if std::env::var("MEMRA_DSPARK_VG_AUTOFREE").as_deref() == Ok("1") {
2815 cudarc::driver::sys::CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH
2816 } else {
2817 cudarc::driver::sys::CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_USE_NODE_PRIORITY
2818 };
2819 e.capture_graph_retained_flags(iflag, move |e| {
2820 let mut xc: Option<CudaSlice<f32>> = None;
2821 for il in lo..hi {
2822 let xr: &CudaSlice<f32> = xc.as_ref().unwrap_or(xin);
2823 let nx = if let Some(&k) = lin_pos.get(&il) {
2824 model.qwen35_tparallel_linear_layer(
2825 e,
2826 il,
2827 xr,
2828 t,
2829 cache_ref,
2830 None,
2831 Some((&mut stash_conv[k], &mut stash_ssm[k])),
2832 Some((table_all, k * 6)),
2833 )?
2834 } else if let Some(&kf) = fa_pos.get(&il) {
2835 let mut no_rows: Option<Vec<CudaSlice<i32>>> = None;
2836 model.qwen35_tparallel_fa_layer(
2837 e,
2838 il,
2839 xr,
2840 t,
2841 cache_ref,
2842 FaLayerArgs {
2843 pos_d,
2844 pos_rows: &mut no_rows,
2845 pos0,
2846 graph_cap: Some((fa_table, kf * 2 * t_cap, rung)),
2847 stream: None,
2848 ckpt: None,
2849 },
2850 )?
2851 } else {
2852 return Err(format!(
2853 "run_full: layer {il} is neither linear nor full-attention"
2854 )
2855 .into());
2856 };
2857 xc = Some(nx);
2858 }
2859 e.copy_into(xout, 0, xc.as_ref().unwrap(), t * n_embd)?;
2860 Ok(())
2861 })?
2862 };
2863 // Undo the net host parity motion of the 3 body runs (each run swaps iff t
2864 // is odd -> 3 runs = net one swap), then restore the device state the
2865 // warmups consumed (walk scope only — layers past hi never executed). The
2866 // launch below then behaves exactly like one run.
2867 if t % 2 == 1 {
2868 for &il in &self.lin {
2869 if il < lo || il >= hi {
2870 continue;
2871 }
2872 let rl = cache.recur[il].as_mut().unwrap();
2873 std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
2874 }
2875 }
2876 for (k, &il) in self.lin.iter().enumerate() {
2877 if il < lo || il >= hi {
2878 continue;
2879 }
2880 let rl = cache.recur[il].as_mut().unwrap();
2881 let (cw, sw) = (self.conv_words, self.ssm_words);
2882 {
2883 let sv = e.view(&self.save_conv, self.lin.len() * cw);
2884 let win = sv.slice(k * cw..(k + 1) * cw);
2885 e.copy_view_into(&mut rl.conv_state, 0, &win, cw)?;
2886 }
2887 {
2888 let sv = e.view(&self.save_ssm, self.lin.len() * sw);
2889 let win = sv.slice(k * sw..(k + 1) * sw);
2890 e.copy_view_into(&mut rl.ssm_state, 0, &win, sw)?;
2891 }
2892 }
2893 if std::env::var("MEMRA_GRAPH_CENSUS").as_deref() == Ok("1")
2894 && let Ok(c) = crate::graph_update::node_census(&graph)
2895 {
2896 eprintln!("[dspark-vg-census] full vt={t} rung={rung} {c:?}");
2897 }
2898 self.full.insert(
2899 key,
2900 DsparkSegGraph {
2901 graph,
2902 _keeper: keeper,
2903 },
2904 );
2905 }
2906 self.full[&key].graph.launch()?;
2907 // Host bookkeeping for the replayed body (captured host code does not re-run):
2908 // gdn parity swap per linear layer (t odd), kv len bump per fa layer — scoped
2909 // to the WALK [lo, hi): the cache can carry extra state slots past it (the MTP
2910 // head layer's kv) that the walk never touches.
2911 if t % 2 == 1 {
2912 for &il in &self.lin {
2913 if il < lo || il >= hi {
2914 continue;
2915 }
2916 let rl = cache.recur[il].as_mut().unwrap();
2917 std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
2918 }
2919 }
2920 for &il in &self.fa {
2921 if il < lo || il >= hi {
2922 continue;
2923 }
2924 cache.kv[il].as_mut().unwrap().len += t;
2925 }
2926 let (_, xout) = self.stage.get(&t).unwrap();
2927 let mut out = e.uninit(t * n_embd)?;
2928 e.copy_into(&mut out, 0, xout, t * n_embd)?;
2929 Ok(out)
2930 }
2931
2932 /// Run layers [start, end) (all linear) as one captured graph at this vt: stage the
2933 /// residual into the bucket's x_in, capture on first encounter (2 executing warmups
2934 /// bracketed by a segment state save/restore), launch, then apply the host parity
2935 /// bookkeeping the captured body would have done. Returns the fresh residual.
2936 #[allow(clippy::too_many_arguments)]
2937 #[allow(clippy::map_entry)] // allow: the init body is fallible (`?`); Entry::or_insert_with cannot propagate errors
2938 fn run_segment(
2939 &mut self,
2940 model: &crate::hybrid::HybridModel,
2941 e: &Engine,
2942 start: usize,
2943 end: usize,
2944 x: &CudaSlice<f32>,
2945 t: usize,
2946 cache: &mut Cache,
2947 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2948 let n_embd = self.n_embd;
2949 debug_assert!(end - start <= self.max_run);
2950 if !self.stage.contains_key(&t) {
2951 let xin = e.uninit(t * n_embd)?;
2952 let xout = e.uninit(t * n_embd)?;
2953 self.stage.insert(t, (xin, xout));
2954 }
2955 // Stage the residual at the bucket's baked input address.
2956 {
2957 let (xin, _) = self.stage.get_mut(&t).unwrap();
2958 e.copy_into(xin, 0, x, t * n_embd)?;
2959 }
2960 let key = (start, t);
2961 if !self.graphs.contains_key(&key) {
2962 // The 2 warmups EXECUTE the segment on live state — save conv + the canonical
2963 // ssm of every segment layer first, restore after, so the graph's first real
2964 // launch starts from the exact pre-round state (bytes gated e2e).
2965 for (k, il) in (start..end).enumerate() {
2966 let rl = cache.recur[il].as_ref().unwrap();
2967 e.copy_into(
2968 &mut self.save_conv,
2969 k * self.conv_words,
2970 &rl.conv_state,
2971 self.conv_words,
2972 )?;
2973 e.copy_into(
2974 &mut self.save_ssm,
2975 k * self.ssm_words,
2976 &rl.ssm_state,
2977 self.ssm_words,
2978 )?;
2979 }
2980 let (graph, keeper) = {
2981 let table_all = &self.table_all;
2982 let lin_pos = &self.lin_pos;
2983 let stash_conv = &mut self.stash_conv;
2984 let stash_ssm = &mut self.stash_ssm;
2985 let (xin, xout) = self
2986 .stage
2987 .get_mut(&t)
2988 .map(|(a, b)| (&*a, b))
2989 .expect("stage bucket created above");
2990 let cache_ref: &mut Cache = cache;
2991 // Slice 4 (fa-execupdate lane): USE_NODE_PRIORITY instead of
2992 // AUTO_FREE_ON_LAUNCH. The slice-3 measured limiter was AUTO_FREE's
2993 // launch-time mem-pool scan — 25.6 us per cuGraphLaunch x 16 segments
2994 // = ~0.41 ms/round, most of the eager-launch savings. The captured
2995 // body's cuMemAllocAsync transients are BALANCED by in-graph frees
2996 // (every transient drops inside the capture region — the generic
2997 // capture path's census precedent, 1589/1589), so AUTO_FREE has
2998 // nothing to reclaim and the graph is legal to instantiate without
2999 // it; PRIORITY is the flag the gemma slotted door ships for exactly
3000 // this reason (both alternatives drop the scan; UPLOAD via
3001 // cuGraphInstantiateWithFlags is WithParams-only and refused).
3002 // MEMRA_DSPARK_VG_AUTOFREE=1 reverts; MEMRA_GRAPH_CENSUS=1 prints
3003 // the node census at capture (the ALLOC==FREE receipt).
3004 let iflag = if std::env::var("MEMRA_DSPARK_VG_AUTOFREE").as_deref() == Ok("1") {
3005 cudarc::driver::sys::CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH
3006 } else {
3007 cudarc::driver::sys::CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_USE_NODE_PRIORITY
3008 };
3009 e.capture_graph_retained_flags(iflag, move |e| {
3010 let mut xc: Option<CudaSlice<f32>> = None;
3011 for il in start..end {
3012 let k = lin_pos[&il];
3013 let xr: &CudaSlice<f32> = xc.as_ref().unwrap_or(xin);
3014 let nx = model.qwen35_tparallel_linear_layer(
3015 e,
3016 il,
3017 xr,
3018 t,
3019 cache_ref,
3020 None,
3021 Some((&mut stash_conv[k], &mut stash_ssm[k])),
3022 Some((table_all, k * 6)),
3023 )?;
3024 xc = Some(nx);
3025 }
3026 e.copy_into(xout, 0, xc.as_ref().unwrap(), t * n_embd)?;
3027 Ok(())
3028 })?
3029 };
3030 // Undo the net host parity motion of the 3 body runs (each run swaps iff t
3031 // is odd -> 3 runs = net one swap), then restore the device state the
3032 // warmups consumed. The launch below then behaves exactly like one run.
3033 if t % 2 == 1 {
3034 for il in start..end {
3035 let rl = cache.recur[il].as_mut().unwrap();
3036 std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
3037 }
3038 }
3039 for (k, il) in (start..end).enumerate() {
3040 let rl = cache.recur[il].as_mut().unwrap();
3041 let (cw, sw) = (self.conv_words, self.ssm_words);
3042 {
3043 let sv = e.view(&self.save_conv, self.lin.len() * cw);
3044 let win = sv.slice(k * cw..(k + 1) * cw);
3045 e.copy_view_into(&mut rl.conv_state, 0, &win, cw)?;
3046 }
3047 {
3048 let sv = e.view(&self.save_ssm, self.lin.len() * sw);
3049 let win = sv.slice(k * sw..(k + 1) * sw);
3050 e.copy_view_into(&mut rl.ssm_state, 0, &win, sw)?;
3051 }
3052 }
3053 if std::env::var("MEMRA_GRAPH_CENSUS").as_deref() == Ok("1")
3054 && let Ok(c) = crate::graph_update::node_census(&graph)
3055 {
3056 eprintln!("[dspark-vg-census] seg={start}..{end} vt={t} {c:?}");
3057 }
3058 self.graphs.insert(
3059 key,
3060 DsparkSegGraph {
3061 graph,
3062 _keeper: keeper,
3063 },
3064 );
3065 }
3066 self.graphs[&key].graph.launch()?;
3067 // Host parity bookkeeping for the replayed body (the captured host swaps do not
3068 // re-run at replay).
3069 if t % 2 == 1 {
3070 for il in start..end {
3071 let rl = cache.recur[il].as_mut().unwrap();
3072 std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
3073 }
3074 }
3075 let (_, xout) = self.stage.get(&t).unwrap();
3076 let mut out = e.uninit(t * n_embd)?;
3077 e.copy_into(&mut out, 0, xout, t * n_embd)?;
3078 Ok(out)
3079 }
3080
3081 /// Pool freeze check (`dspark_vg_cap`): below the ceiling new keys may capture.
3082 fn can_capture(&self) -> bool {
3083 self.graphs.len() + self.full.len() < dspark_vg_cap()
3084 }
3085
3086 /// Round-atomic segment-door readiness: TRUE when this round's walk can ride the
3087 /// per-(segment, vt) graphs without a NEW capture past the pool ceiling — every
3088 /// linear run in [lo, hi) already has its (run_start, t) key, or capture is still
3089 /// allowed. FALSE sends the WHOLE round down the eager cols-ckpt walk: a partial
3090 /// refusal would stash some layers in the ctx slabs and others in the round's cols
3091 /// while one commit reads only one of them.
3092 pub(crate) fn segments_ready(
3093 &self,
3094 model: &crate::hybrid::HybridModel,
3095 lo: usize,
3096 hi: usize,
3097 t: usize,
3098 ) -> bool {
3099 if self.can_capture() {
3100 return true;
3101 }
3102 let mut il = lo;
3103 while il < hi {
3104 if matches!(model.layers[il].mixer, Mixer::Linear(_)) {
3105 let start = il;
3106 while il < hi && matches!(model.layers[il].mixer, Mixer::Linear(_)) {
3107 il += 1;
3108 }
3109 if !self.graphs.contains_key(&(start, t)) {
3110 return false;
3111 }
3112 } else {
3113 il += 1;
3114 }
3115 }
3116 true
3117 }
3118
3119 /// Widest verify window this pool was built for. A caller whose round exceeds it must
3120 /// take the eager walk: the stash slabs hold `t_capacity() - 1` column rows, and slicing
3121 /// past them is a panic rather than a refusal.
3122 pub(crate) fn t_capacity(&self) -> usize {
3123 self.t_cap
3124 }
3125
3126 /// Slab row (conv, ssm) device pointers + lengths for the commit restore of column
3127 /// `row` (0-based) of layer `il`. None for non-linear layers.
3128 pub(crate) fn slab_row(
3129 &self,
3130 e: &Engine,
3131 il: usize,
3132 row: usize,
3133 ) -> Option<(u64, u64, usize, usize)> {
3134 use cudarc::driver::DevicePtr;
3135 let k = *self.lin_pos.get(&il)?;
3136 let s = &e.gpu.stream();
3137 let (pc, _g0) = self.stash_conv[k].device_ptr(s);
3138 let (ps, _g1) = self.stash_ssm[k].device_ptr(s);
3139 Some((
3140 pc + (row * self.conv_words * 4) as u64,
3141 ps + (row * self.ssm_words * 4) as u64,
3142 self.conv_words,
3143 self.ssm_words,
3144 ))
3145 }
3146}
3147
3148impl VerifyCkpt {
3149 fn new(n_layer: usize) -> Self {
3150 VerifyCkpt {
3151 gdn: (0..n_layer).map(|_| None).collect(),
3152 cols: (0..n_layer).map(|_| None).collect(),
3153 }
3154 }
3155}
3156
3157/// The stage-0/TX half of one PP verify. The boundary slot is the ownership token: stage 1
3158/// consumes exactly the slot selected by `tx()`, never a slot inferred from
3159/// a logical round number.
3160struct VerifyBoundaryTicket {
3161 rt: &'static crate::pp::PpNRt,
3162 caller_stream: std::sync::Arc<cudarc::driver::CudaStream>,
3163 slot: usize,
3164 pos0: usize,
3165 t: usize,
3166 payload: usize,
3167 n_st: usize,
3168 pp_anatomy: bool,
3169 pp_started: std::time::Instant,
3170 reverse_ms: f64,
3171 stage0_ms: f64,
3172 tx_ms: f64,
3173 _walk_owner: crate::pp::PpWalkLease,
3174}
3175
3176/// MEMRA_SPEC_ROUND_PROF counters: whole-round wall, so the round can be weighed against the
3177/// draft-step ([spec-anatomy]) and verify-walk ([tcol-prof]) splits we already print.
3178static ROUND_PROF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
3179static ROUND_MS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
3180static ROUND_N: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
3181
3182fn validate_tp_kv_snapshot_shape(
3183 tp_kv: &[Option<crate::tp::ResidentTpKvCache>],
3184 saved_lens: &[Option<usize>],
3185) -> Result<(), Box<dyn std::error::Error>> {
3186 if tp_kv.len() != saved_lens.len() {
3187 return Err("spec TP KV snapshot shape mismatch".into());
3188 }
3189 for (layer, (cache, saved)) in tp_kv.iter().zip(saved_lens).enumerate() {
3190 if cache.is_some() != saved.is_some() {
3191 return Err(
3192 format!("spec TP KV layer {layer} changed shape since its snapshot").into(),
3193 );
3194 }
3195 }
3196 Ok(())
3197}
3198
3199impl HybridModel {
3200 /// memra#128: the canonical-to-rank byte copy that `5e0fffb97` added to
3201 /// `restore_step_tp_kv_verified_prefix`. OFF by default (written decision, docs/FLAGS.md
3202 /// `MEMRA_STEP_TP_KV_RESTORE`): on step37 NVFP4 TP2 the only shapes that pass the
3203 /// production acceptance gate are the engine before the copy (arm A) and the copy skipped
3204 /// on every step-TP layer (arm F, byte-identical answers to A); the copy on any layer
3205 /// spliced answers or shifted decode (darklanes research/memra128-bisect-20260903).
3206 /// `1` re-enables the copy for ordinary-commit layers; on-device-written layers
3207 /// (`rows_external`) are skipped either way, their canonical bytes are stale.
3208 fn step_tp_kv_restore_copy_on() -> bool {
3209 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
3210 *ON.get_or_init(|| std::env::var("MEMRA_STEP_TP_KV_RESTORE").ok().as_deref() == Some("1"))
3211 }
3212
3213 fn restore_step_tp_kv_verified_prefix(
3214 &self,
3215 e: &Engine,
3216 cache: &mut Cache,
3217 snap: &crate::cache::CacheSnapshot,
3218 accepted: usize,
3219 // memra#128: what an externally-written (dcw / fa2) layer needs from this call.
3220 // PARTIAL accept (commit_verified_prefix): 5e0fffb97 replaced the standalone
3221 // rewind_tp_kv_verified_prefix with this restore, so the length shrink to
3222 // saved+accepted must still happen here - without it E ran with distributed=259
3223 // against local=257. FULL accept: before 5e0fffb97 nothing touched the
3224 // distributed length there and it was right (arm A passed); rewinding to
3225 // saved+t_v shrinks it by one and the next verify's SWA ring view falls off the
3226 // end ("view [5148,5152) is outside resident [0,5151)", arm E2).
3227 rewind_external: bool,
3228 ) -> Result<(), Box<dyn std::error::Error>> {
3229 validate_tp_kv_snapshot_shape(&cache.tp_kv, &snap.tp_kv_len)?;
3230 e.stream().synchronize()?;
3231 {
3232 let (local_layers, distributed_layers) = (&cache.kv, &mut cache.tp_kv);
3233 let stream = e.gpu.stream();
3234 let mut runtime: Option<std::sync::Arc<crate::tp::TpE4m3HostBounce>> = None;
3235 let mut uniform_runtime = true;
3236 let mut batch = Vec::new();
3237 for (il, (distributed_slot, local_slot)) in distributed_layers
3238 .iter_mut()
3239 .zip(local_layers.iter())
3240 .enumerate()
3241 {
3242 let (Some(distributed), Some(saved)) =
3243 (distributed_slot.as_mut(), snap.tp_kv_len[il])
3244 else {
3245 continue;
3246 };
3247 // memra#128: on the dcw / fa2 verify path the rank rows for [saved, target)
3248 // were written on-device and the canonical cache holds only stale bytes for
3249 // them (the verify bumps `local.len` and writes nothing). Copying those over
3250 // the correct rank rows is exactly what spliced two requests' answers
3251 // together on step-3.7-flash. The rewind already kept the right rows; skip.
3252 let target = saved
3253 .checked_add(accepted)
3254 .ok_or("spec TP KV batch restore length overflow")?;
3255 if !Self::step_tp_kv_restore_copy_on() || distributed.rows_external() {
3256 // Rank rows already right (written on-device). Length: see the
3257 // `rewind_external` note on the signature.
3258 if rewind_external {
3259 distributed.rewind_to(target)?;
3260 }
3261 continue;
3262 }
3263 let local = local_slot
3264 .as_ref()
3265 .ok_or_else(|| format!("spec TP KV layer {il} lost its owning cache"))?;
3266 if local.len < target {
3267 return Err(format!(
3268 "spec TP KV layer {il} local length {} precedes restore target {target}",
3269 local.len
3270 )
3271 .into());
3272 }
3273 let physical = local.physical_rows(saved, target)?;
3274 if physical.len() != accepted {
3275 return Err(format!(
3276 "spec TP KV layer {il} restore [{saved},{target}) is not contiguous"
3277 )
3278 .into());
3279 }
3280 let Mixer::Full(fa) = &self.layers[il].mixer else {
3281 return Err(format!("spec TP KV layer {il} is not full attention").into());
3282 };
3283 let tp = fa
3284 .step_tp_qkv
3285 .as_ref()
3286 .ok_or_else(|| format!("spec TP KV layer {il} lost its TP runtime"))?;
3287 if let Some(first) = runtime.as_ref() {
3288 if !std::sync::Arc::ptr_eq(first, &tp.runtime) {
3289 uniform_runtime = false;
3290 break;
3291 }
3292 } else {
3293 runtime = Some(tp.runtime.clone());
3294 }
3295 use cudarc::driver::DevicePtr;
3296 let (k_base, _k_guard) = local.k.device_ptr(&stream);
3297 let (v_base, _v_guard) = local.v.device_ptr(&stream);
3298 batch.push(crate::tp::TpKvVerifiedLayer {
3299 cache: distributed,
3300 start: saved,
3301 logical_len: target,
3302 source_k_raw: k_base + (physical.start * local.k_tok_bytes) as u64,
3303 source_v_raw: v_base + (physical.start * local.v_tok_bytes) as u64,
3304 source_k_tok_bytes: local.k_tok_bytes,
3305 source_v_tok_bytes: local.v_tok_bytes,
3306 });
3307 }
3308 if uniform_runtime
3309 && let Some(runtime) = runtime
3310 && runtime.restore_tp_kv_layers_from_device(&mut batch)?
3311 {
3312 return Ok(());
3313 }
3314 }
3315 for il in 0..self.layers.len() {
3316 let (Some(distributed), Some(saved)) = (cache.tp_kv[il].as_mut(), snap.tp_kv_len[il])
3317 else {
3318 continue;
3319 };
3320 let target = saved
3321 .checked_add(accepted)
3322 .ok_or("spec TP KV restore length overflow")?;
3323 if !Self::step_tp_kv_restore_copy_on() || distributed.rows_external() {
3324 if rewind_external {
3325 distributed.rewind_to(target)?;
3326 }
3327 continue;
3328 }
3329 let local = cache.kv[il]
3330 .as_ref()
3331 .ok_or_else(|| format!("spec TP KV layer {il} lost its owning cache"))?;
3332 if local.len < target {
3333 return Err(format!(
3334 "spec TP KV layer {il} local length {} precedes restore target {target}",
3335 local.len
3336 )
3337 .into());
3338 }
3339 let physical = local.physical_rows(saved, target)?;
3340 if physical.len() != accepted {
3341 return Err(format!(
3342 "spec TP KV layer {il} restore [{saved},{target}) is not contiguous"
3343 )
3344 .into());
3345 }
3346 use cudarc::driver::DevicePtr;
3347 let stream = e.gpu.stream();
3348 let (k_base, _k_guard) = local.k.device_ptr(&stream);
3349 let (v_base, _v_guard) = local.v.device_ptr(&stream);
3350 let k_raw = k_base + (physical.start * local.k_tok_bytes) as u64;
3351 let v_raw = v_base + (physical.start * local.v_tok_bytes) as u64;
3352 let Mixer::Full(fa) = &self.layers[il].mixer else {
3353 return Err(format!("spec TP KV layer {il} is not full attention").into());
3354 };
3355 let tp = fa
3356 .step_tp_qkv
3357 .as_ref()
3358 .ok_or_else(|| format!("spec TP KV layer {il} lost its TP runtime"))?;
3359 tp.runtime.restore_tp_kv_rows_from_device(
3360 distributed,
3361 saved,
3362 target,
3363 k_raw,
3364 v_raw,
3365 local.k_tok_bytes,
3366 local.v_tok_bytes,
3367 )?;
3368 }
3369 Ok(())
3370 }
3371
3372 fn mtp_head_count(&self) -> usize {
3373 usize::from(self.mtp.is_some()) + self.mtp_extra.len()
3374 }
3375
3376 fn mtp_head_at(&self, index: usize) -> &MtpHead {
3377 if index == 0 {
3378 self.mtp.as_ref().expect("MTP head 0 is unavailable")
3379 } else {
3380 &self.mtp_extra[index - 1]
3381 }
3382 }
3383
3384 fn new_mtp_scratch(
3385 &self,
3386 e: &Engine,
3387 cap: usize,
3388 ) -> Result<MtpScratch, Box<dyn std::error::Error>> {
3389 let mut scratch = MtpScratch::new(
3390 e,
3391 &self.cfg,
3392 &self.plan,
3393 cap,
3394 self.mtp.as_ref().and_then(|head| head.geom.as_ref()),
3395 )?;
3396 for head in &self.mtp_extra {
3397 scratch.push_plane(e, &self.cfg, &self.plan, head.geom.as_ref())?;
3398 }
3399 Ok(scratch)
3400 }
3401
3402 /// NextN head forward for ONE draft token (§A ops 1-13, T=1).
3403 /// Inputs: `e_tok` = the token to predict FROM (last committed / previous draft); `h_seed` =
3404 /// the trunk's pre-output_norm hidden of that token (§A op 2 input). `mtp_pos` = absolute
3405 /// position of the token being predicted from. Returns (draft_logits[n_vocab] host, h_nextn dev).
3406 /// `h_nextn` (§A op 10) becomes `h_seed` for the next autoregressive draft step.
3407 /// Device-resident: returns draft logits ON DEVICE (no [n_vocab] dtoh). The greedy draft
3408 /// loop only needs argmax — paired with `argmax_token_device` this cuts the ~600KB logits
3409 /// transfer + host argmax per draft token from the K-token draft chain.
3410 #[allow(clippy::too_many_arguments)]
3411 fn mtp_head_forward_dev(
3412 &self,
3413 e: &Engine,
3414 mtp: &MtpHead,
3415 e_tok: u32,
3416 h_seed: &CudaSlice<f32>,
3417 scratch: &mut MtpScratch,
3418 mtp_pos: usize,
3419 embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
3420 mask: Option<(&CudaSlice<u32>, usize)>,
3421 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
3422 self.mtp_head_forward_dev_at(e, mtp, e_tok, h_seed, scratch, 0, mtp_pos, embd_dev, mask)
3423 }
3424
3425 #[allow(clippy::too_many_arguments)]
3426 fn mtp_head_forward_dev_at(
3427 &self,
3428 e: &Engine,
3429 mtp: &MtpHead,
3430 e_tok: u32,
3431 h_seed: &CudaSlice<f32>,
3432 scratch: &mut MtpScratch,
3433 scratch_index: usize,
3434 mtp_pos: usize,
3435 embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
3436 // DRAFT-SIDE GRAMMAR MASK (lane/draft-mask): (packed draft-vocab allowed set, words).
3437 // Applied to the head logits BEFORE they are returned, so every consumer (argmax,
3438 // gumbel draw, p-min prob) sees the grammar-legal row. None = unmasked (pre-lane).
3439 mask: Option<(&CudaSlice<u32>, usize)>,
3440 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
3441 // MEMRA_SPEC_ANATOMY=1 — eager-step phase timers (diagnostic only). Phase boundaries
3442 // sync the stream, so absolute time inflates; the BREAKDOWN is the signal. Cumulative
3443 // summary on stderr every 128 steps: glue (embed..attn_norm), attn, ffn, head.
3444 use std::sync::atomic::{AtomicU64, Ordering::Relaxed};
3445 static ANAT_NS: [AtomicU64; 5] = [
3446 AtomicU64::new(0),
3447 AtomicU64::new(0),
3448 AtomicU64::new(0),
3449 AtomicU64::new(0),
3450 AtomicU64::new(0),
3451 ];
3452 static ANAT_STEPS: AtomicU64 = AtomicU64::new(0);
3453 let anat = {
3454 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
3455 *ON.get_or_init(|| std::env::var("MEMRA_SPEC_ANATOMY").as_deref() == Ok("1"))
3456 };
3457 if anat {
3458 e.stream().synchronize()?; // drain prior queue so phase 0 starts clean
3459 }
3460 let t_all = std::time::Instant::now();
3461 let mut t_ph = std::time::Instant::now();
3462 let anat_mark = |i: usize,
3463 e: &Engine,
3464 t: &mut std::time::Instant|
3465 -> Result<(), Box<dyn std::error::Error>> {
3466 if anat {
3467 e.stream().synchronize()?;
3468 ANAT_NS[i].fetch_add(t.elapsed().as_nanos() as u64, Relaxed);
3469 *t = std::time::Instant::now();
3470 }
3471 Ok(())
3472 };
3473 let cfg = &self.cfg;
3474 let n_embd = cfg.n_embd as usize;
3475 // Distilled-student geometry: the block runs at the INNER width `di` (eh_proj out /
3476 // attn / ffn); the n_embd interface (embed, norms in, carrier out, head in) is unchanged.
3477 let di = mtp.geom.as_ref().map(|g| g.d_inner).unwrap_or(n_embd);
3478 let eps = cfg.rms_eps;
3479 let pos_d = e.htod_i32(&[mtp_pos as i32])?;
3480
3481 // op A: a resident table transfers one 4B token id. The exact host-row capacity path
3482 // expands this one row on CPU and transfers n_embd f32 values instead.
3483 let e_emb = match embd_dev {
3484 Some((g, qt, rb)) => e.embed_gather_device_t(g, &[e_tok], n_embd, qt, rb)?,
3485 None => e.htod(&self.embd.try_gather(n_embd, &[e_tok])?)?,
3486 };
3487
3488 // op 1/2: e_norm = RMSNorm(e, enorm); h_norm = RMSNorm(h_seed, hnorm)
3489 let mut e_norm = e.zeros(n_embd)?;
3490 e.rms_norm(&e_emb, mtp.enorm.float_data(), &mut e_norm, n_embd, 1, eps)?;
3491 let mut h_norm = e.zeros(n_embd)?;
3492 e.rms_norm(h_seed, mtp.hnorm.float_data(), &mut h_norm, n_embd, 1, eps)?;
3493
3494 // op 3: concat = [e_norm ; h_norm] -> [2*n_embd], e_norm in [0,n_embd), h_norm in [n_embd,2n_embd)
3495 let mut concat = e.zeros(2 * n_embd)?;
3496 e.copy_into(&mut concat, 0, &e_norm, n_embd)?;
3497 e.copy_into(&mut concat, n_embd, &h_norm, n_embd)?;
3498
3499 // op 4: inpSA = eh_proj @ concat (eh_proj [2*n_embd, n_embd]) -> [n_embd]
3500 let inp_sa = e.matmul(&mtp.eh_proj, &concat, 1)?;
3501
3502 // op 5: a_norm = RMSNorm(inpSA, attn_norm)
3503 let mut a_norm = e.zeros(di)?;
3504 e.rms_norm(&inp_sa, mtp.attn_norm.float_data(), &mut a_norm, di, 1, eps)?;
3505 anat_mark(0, e, &mut t_ph)?;
3506
3507 // op 6: attention on the scratch KV. SAME dc launcher as the graph path (bucket_max =
3508 // scratch.cap, length from the device len_d) so eager drafts match graph drafts
3509 // bit-for-bit at any t_kv (the parity gate). Host len mirrored here (the dc append
3510 // advances only the device counter).
3511 let attn_out = match (&mtp.mixer, mtp.step35.as_ref()) {
3512 // step35 MTP block, dcw door armed: the SAME windowed device-counter launcher as
3513 // the captured chain (draft parity by construction). Per-step ring headroom runs
3514 // HERE (eager is host-len work, a rebase is legal); host len mirrored like the
3515 // plain dc arm below.
3516 (Mixer::Full(fa), Some(g))
3517 if self.step35_dcw_eligible(g, scratch.plane(scratch_index).1) =>
3518 {
3519 {
3520 let (kv, _) = scratch.plane_mut(scratch_index);
3521 let retain = match kv.ring.as_ref() {
3522 Some(ring) => memra_kv::swa_retain_from(kv.len, ring.window(), ring.base()),
3523 None => 0,
3524 };
3525 e.prepare_kv_append(kv, retain, 1)?;
3526 }
3527 let out =
3528 self.mtp_step35_attn_dcw(e, fa, g, &a_norm, &pos_d, scratch, scratch_index)?;
3529 scratch.plane_mut(scratch_index).0.len += 1;
3530 out
3531 }
3532 // step35 MTP block, door off (MEMRA_STEP35_DRAFT_DCW=0 rollback) or class-
3533 // ineligible: PER-LAYER geometry + a separate head-wise gate + an SWA window,
3534 // none of which the plain dc launcher can express (see `mtp_step35_attn`).
3535 // Host-len arm. Advances BOTH the
3536 // host len and the device counter itself (unlike the dc arm, whose host-side
3537 // mirror the caller does).
3538 (Mixer::Full(fa), Some(g)) => {
3539 self.mtp_step35_attn(e, fa, g, &a_norm, &pos_d, scratch, scratch_index)?
3540 }
3541 (Mixer::Full(fa), None) => {
3542 let out = self.mtp_full_attn_dc(
3543 e,
3544 fa,
3545 &a_norm,
3546 &pos_d,
3547 scratch,
3548 scratch_index,
3549 mtp.geom.as_ref(),
3550 )?;
3551 scratch.plane_mut(scratch_index).0.len += 1;
3552 out
3553 }
3554 (Mixer::Linear(_), _) => {
3555 panic!("MTP block is full-attn in qwen35; linear MTP not supported")
3556 }
3557 (Mixer::Mla(_), _) => crate::hybrid::mla_path_unimplemented("MTP head forward"),
3558 (Mixer::Kda(_), _) => crate::hybrid::kda_path_unimplemented("MTP head forward"),
3559 };
3560 anat_mark(1, e, &mut t_ph)?;
3561
3562 // op 7: x1 = inpSA + attn_out
3563 let mut x1 = e.zeros(di)?;
3564 e.add(&inp_sa, &attn_out, &mut x1, di)?;
3565
3566 // op 8: z = RMSNorm(x1, post_attn_norm) (pre-FFN norm)
3567 let mut z = e.zeros(di)?;
3568 e.rms_norm(&x1, mtp.post_attn_norm.float_data(), &mut z, di, 1, eps)?;
3569
3570 // op 9: FFN (Dense or MoE) — same as the trunk decode FFN
3571 let ffn_out = match &mtp.ffn {
3572 crate::hybrid::Ffn::Dense {
3573 ffn_gate,
3574 ffn_up,
3575 ffn_down,
3576 } => {
3577 let n_ff = ffn_gate.out_features();
3578 let (gate, up) = if e.uses_q8_1_fast(ffn_gate) && e.uses_q8_1_fast(ffn_up) {
3579 let (zq, zd) = e.quantize_q8_1(&z, 1, di)?;
3580 (
3581 e.matmul_pre(ffn_gate, &zq, &zd, &z, 1)?,
3582 e.matmul_pre(ffn_up, &zq, &zd, &z, 1)?,
3583 )
3584 } else {
3585 (e.matmul(ffn_gate, &z, 1)?, e.matmul(ffn_up, &z, 1)?)
3586 };
3587 let mut act = e.zeros(n_ff)?;
3588 // step35: a DENSE FFN reads the per-layer SHEXP clamp (upstream's one `build_ffn`
3589 // serves the dense MLP and the shared expert off `swiglu_clamp_shexp` —
3590 // llama-graph.cpp:1751), resolved for the MTP block's OWN index. Every other arch
3591 // passes None, which is `ffn_act`'s dispatch verbatim.
3592 Self::ffn_act_lim(
3593 e,
3594 &self.cfg,
3595 &gate,
3596 &up,
3597 1.0,
3598 1.0,
3599 mtp.step35
3600 .as_ref()
3601 .and_then(|s| s.clamp_shexp)
3602 .map(SwigluClamp::Post),
3603 &mut act,
3604 n_ff,
3605 )?;
3606 e.matmul(ffn_down, &act, 1)?
3607 }
3608 // MTP head is a distinct block — key its experts under a separate layer index (u16::MAX)
3609 // so they never alias trunk layer 0's cache keys.
3610 crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il(e, m, &z, 1, u16::MAX)?,
3611 };
3612 anat_mark(2, e, &mut t_ph)?;
3613
3614 // op 10: h_nextn = x1 + ffn_out (at di)
3615 let mut h_inner = e.zeros(di)?;
3616 e.add(&x1, &ffn_out, &mut h_inner, di)?;
3617
3618 // op 10.5 (student): up-project the inner hidden back to n_embd — training semantics:
3619 // the chain carrier AND the head input are out_up(h_inner) (pre-final-norm).
3620 let h_nextn = match mtp.geom.as_ref() {
3621 Some(g) => e.matmul(&g.out_up, &h_inner, 1)?,
3622 None => h_inner,
3623 };
3624
3625 // op 11: final = RMSNorm(h_nextn, shared_head_norm OR output_norm)
3626 let final_norm = mtp.shared_head_norm.as_ref().unwrap_or(&self.output_norm);
3627 let mut final_h = e.zeros(n_embd)?;
3628 e.rms_norm(
3629 &h_nextn,
3630 final_norm.float_data(),
3631 &mut final_h,
3632 n_embd,
3633 1,
3634 eps,
3635 )?;
3636
3637 // op 12: draft_logits = (shared_head_head OR output) @ final — stays ON DEVICE.
3638 let head = mtp.shared_head_head.as_ref().unwrap_or(&self.output);
3639 let mut logits = e.matmul(head, &final_h, 1)?;
3640 // op 12b (lane/draft-mask): grammar mask over the DRAFT vocab, applied here so the
3641 // caller's argmax / gumbel draw / p-min prob all read the grammar-legal row.
3642 if let Some((mask_d, mw)) = mask {
3643 let d_vocab = head.out_features();
3644 e.mask_logits_col(&mut logits, mask_d, 0, d_vocab, mw)?;
3645 }
3646 anat_mark(3, e, &mut t_ph)?;
3647 if anat {
3648 ANAT_NS[4].fetch_add(t_all.elapsed().as_nanos() as u64, Relaxed);
3649 let n = ANAT_STEPS.fetch_add(1, Relaxed) + 1;
3650 if n.is_multiple_of(128) {
3651 let us = |i: usize| ANAT_NS[i].load(Relaxed) / n / 1000;
3652 eprintln!(
3653 "[spec-anatomy] steps={n} avg us/step: glue={} attn={} ffn={} head={} total={}",
3654 us(0),
3655 us(1),
3656 us(2),
3657 us(3),
3658 us(4)
3659 );
3660 }
3661 }
3662 // Chain recurrence hand-over: pre-norm h_nextn (default) or post-norm final_h
3663 // (MEMRA_SPEC_HPOST — llama.cpp #24025's t_h_nextn is taken AFTER the head norm).
3664 Ok((logits, if spec_hpost() { final_h } else { h_nextn }))
3665 }
3666
3667 /// One NextN/MTP draft step for an **MLA-mixer** MTP block (glm5_next class: MLA + own
3668 /// k-pool indexer + MoE, serial residual — the NextN layer carries no hc_* tensors), on
3669 /// the model `Cache`'s own MTP latent plane rather than the full-attn `MtpScratch` the
3670 /// qwen35/step35 chain uses. Gate: `glm5_mtp_head_gpu` (engine vs `memra_reference`
3671 /// `execute_mtp`, teacher-forced walk, eh_proj-transpose and h_seed-off-by-one red arms).
3672 ///
3673 /// The interface, stated precisely for the verify arc:
3674 /// - `h_seed`: `[n_embd]` f32 device — the trunk's COLLAPSED PRE-output_norm hidden of
3675 /// the position whose next token is being drafted (MTP-PLAN §A; exactly what
3676 /// `prime_cache`/`decode_step` return for hc models). `MEMRA_SPEC_HPOST` flips both
3677 /// this producer and the returned carrier to the post-norm variant, same as the dev path.
3678 /// - `e_tok`: the token at the seeded position's SUCCESSOR — the token the trunk just
3679 /// sampled/accepted (reference oracle pairing: `fused[i] = eh_proj([enorm(embed(ids[i]));
3680 /// hnorm(trunk_hidden[i])])`, i.e. this call with `e_tok = ids[i]`, `h_seed = h[i]`,
3681 /// `mtp_pos = i` reproduces the reference's row `i`).
3682 /// - `mtp_pos`: the absolute position this step appends to the MTP block's latent plane;
3683 /// must equal that plane's current length (the plane advances by ONE row per call inside
3684 /// `mla_attn_cached`; rollback on rejection = the verify arc's latent-plane len reset).
3685 /// - returns `(draft_logits [n_vocab], carrier [n_embd])` on device. glm5_next ships no
3686 /// private MTP head, so the logits ride the trunk `lm_head` (full vocab, no d2t).
3687 pub fn mtp_head_forward_mla_cached(
3688 &self,
3689 e: &Engine,
3690 depth: usize,
3691 e_tok: u32,
3692 h_seed: &CudaSlice<f32>,
3693 cache: &mut Cache,
3694 mtp_pos: usize,
3695 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
3696 if depth >= self.mtp_head_count() {
3697 return Err(format!(
3698 "MTP depth {depth} out of range: {} embedded head(s) loaded \
3699 (is MEMRA_GLM5_MTP=1 set for a glm5_next model?)",
3700 self.mtp_head_count()
3701 )
3702 .into());
3703 }
3704 let mtp = self.mtp_head_at(depth);
3705 let block = self
3706 .plan
3707 .mtp_blocks
3708 .get(depth)
3709 .ok_or_else(|| format!("ModelPlan declares no MTP block at depth {depth}"))?;
3710 let il = block.layer.index as usize;
3711 let Mixer::Mla(mla) = &mtp.mixer else {
3712 return Err(
3713 "mtp_head_forward_mla_cached serves MLA-mixer MTP blocks only; full-attn \
3714 blocks take mtp_head_forward_dev's scratch path"
3715 .into(),
3716 );
3717 };
3718 if matches!(mtp.ffn, crate::hybrid::Ffn::Dense { .. }) {
3719 return Err(
3720 "MLA-mixer MTP block with a Dense FFN has no gated arm yet (glm5_next and \
3721 glm-dsa NextN blocks are MoE); refusing rather than running ungated math"
3722 .into(),
3723 );
3724 }
3725 let plane_len = cache
3726 .latent
3727 .get(il)
3728 .and_then(|plane| plane.as_ref())
3729 .map(|plane| plane.len)
3730 .ok_or_else(|| {
3731 format!(
3732 "MTP block layer {il} has no latent cache plane — the Cache must be \
3733 built from a plan whose mtp_blocks declare StatePlan::LatentKvCache"
3734 )
3735 })?;
3736 if mtp_pos != plane_len {
3737 return Err(format!(
3738 "MTP draft position {mtp_pos} != the MTP latent plane's length {plane_len} — \
3739 the plane advances one row per draft step and rolls back by len reset; a \
3740 skipped or repeated position would attend the wrong horizon"
3741 )
3742 .into());
3743 }
3744
3745 let cfg = &self.cfg;
3746 let n_embd = cfg.n_embd as usize;
3747 let eps = cfg.rms_eps;
3748 let pos_d = e.htod_i32(&[mtp_pos as i32])?;
3749
3750 // Same op chain as `mtp_head_forward_dev_at` (ops 1-12), same kernels — only the
3751 // attention arm differs: `mla_attn_cached` on the plan's own MTP plane instead of
3752 // `mtp_full_attn_dc` on the MtpScratch.
3753 let e_emb = e.htod(&self.embd.try_gather(n_embd, &[e_tok])?)?;
3754 let mut e_norm = e.zeros(n_embd)?;
3755 e.rms_norm(&e_emb, mtp.enorm.float_data(), &mut e_norm, n_embd, 1, eps)?;
3756 let mut h_norm = e.zeros(n_embd)?;
3757 e.rms_norm(h_seed, mtp.hnorm.float_data(), &mut h_norm, n_embd, 1, eps)?;
3758
3759 let mut concat = e.zeros(2 * n_embd)?;
3760 e.copy_into(&mut concat, 0, &e_norm, n_embd)?;
3761 e.copy_into(&mut concat, n_embd, &h_norm, n_embd)?;
3762 let inp_sa = e.matmul(&mtp.eh_proj, &concat, 1)?;
3763
3764 let mut a_norm = e.zeros(n_embd)?;
3765 e.rms_norm(
3766 &inp_sa,
3767 mtp.attn_norm.float_data(),
3768 &mut a_norm,
3769 n_embd,
3770 1,
3771 eps,
3772 )?;
3773 let attn_out = self.mla_attn_cached(e, mla, &a_norm, &pos_d, 1, il, cache)?;
3774
3775 let mut x1 = e.zeros(n_embd)?;
3776 e.add(&inp_sa, &attn_out, &mut x1, n_embd)?;
3777 let mut z = e.zeros(n_embd)?;
3778 e.rms_norm(&x1, mtp.post_attn_norm.float_data(), &mut z, n_embd, 1, eps)?;
3779 let ffn_out = match &mtp.ffn {
3780 // Distinct block — key its experts off the trunk layers' cache keys (dev-path rule).
3781 crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il(e, m, &z, 1, u16::MAX)?,
3782 crate::hybrid::Ffn::Dense { .. } => unreachable!("refused above"),
3783 };
3784 let mut h_nextn = e.zeros(n_embd)?;
3785 e.add(&x1, &ffn_out, &mut h_nextn, n_embd)?;
3786
3787 let final_norm = mtp.shared_head_norm.as_ref().unwrap_or(&self.output_norm);
3788 let mut final_h = e.zeros(n_embd)?;
3789 e.rms_norm(
3790 &h_nextn,
3791 final_norm.float_data(),
3792 &mut final_h,
3793 n_embd,
3794 1,
3795 eps,
3796 )?;
3797 let head = mtp.shared_head_head.as_ref().unwrap_or(&self.output);
3798 let logits = e.matmul(head, &final_h, 1)?;
3799 Ok((logits, if spec_hpost() { final_h } else { h_nextn }))
3800 }
3801
3802 #[allow(clippy::too_many_arguments)]
3803 fn mtp_chain_forward_dev(
3804 &self,
3805 e: &Engine,
3806 tokens: &[u32],
3807 seeds: &[CudaSlice<f32>],
3808 scratch: &mut MtpScratch,
3809 committed_scratch_len: usize,
3810 embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
3811 mask: Option<(&CudaSlice<u32>, usize)>,
3812 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
3813 if tokens.is_empty() || tokens.len() != seeds.len() {
3814 return Err("multi-head MTP prefix tokens/seeds are malformed".into());
3815 }
3816 let index = mtp_chain_head_index(tokens.len() - 1, self.mtp_head_count());
3817 let head = self.mtp_head_at(index);
3818 scratch.set_plane_len(e, index, committed_scratch_len)?;
3819
3820 let mut last = None;
3821 for row in 0..tokens.len() {
3822 let is_last = row + 1 == tokens.len();
3823 last = Some(self.mtp_head_forward_dev_at(
3824 e,
3825 head,
3826 tokens[row],
3827 &seeds[row],
3828 scratch,
3829 index,
3830 committed_scratch_len + row + 1,
3831 embd_dev,
3832 if is_last { mask } else { None },
3833 )?);
3834 }
3835 Ok(last.expect("non-empty MTP prefix produced no row"))
3836 }
3837
3838 /// step35 MTP-block attention, T=1, on the scratch KV — the EAGER-ONLY twin of
3839 /// `mtp_full_attn_dc`. Three things force a separate arm rather than a geometry parameter on
3840 /// the dc path, and all three are properties of this arch's MTP block:
3841 ///
3842 /// 1. **The SWA window.** Block 45 is an SWA-type block (`sliding_window_pattern[45]=true`,
3843 /// window 512). Windowed decode in memra is a token-aligned VIEW OFFSET into the quantized
3844 /// cache (the gemma4 R6 / `step35_decode_attn` pattern: keys carry absolute rope and the
3845 /// mask is purely positional, so one query at `len-1` attending the last `win` rows IS the
3846 /// windowed result). `fa_decode_dc` takes the key count from a DEVICE counter and always
3847 /// starts at row 0 — it cannot express a nonzero offset. The windowed dc arm is
3848 /// `mtp_step35_attn_dcw` (`fa_decode_dcw`, doored via MEMRA_STEP35_DRAFT_DCW —
3849 /// default ON since lane/step37-draft-graph-serving-20260830); this host-len arm is
3850 /// the =0 rollback and the class-ineligibility fallback.
3851 /// 2. **Per-layer head count.** 96 q heads over 8 KV (GQA 12) at this block, vs the trunk's 64
3852 /// on its full-attn layers. The trunk cfg's `n_head` scalar is the MAX over layers, and the
3853 /// trunk ARTIFACT's per-layer arrays stop at index 44 — so the count must come from the
3854 /// resolved `Step35MtpGeom`, never from `cfg`.
3855 /// 3. **The separate head-wise gate.** `blk.45.attn_gate.weight [n_embd, 96]` produces one
3856 /// sigmoid scalar per head (broadcast over head_dim) — `attn_head_gate`, not the qwen35
3857 /// fused-into-wq `q_gate_split` form the dc arm handles.
3858 ///
3859 /// DOOR STATE: with MEMRA_STEP35_DRAFT_DCW=0 (or a sub-eligible kernel class),
3860 /// `mtp_head_forward_cap` refuses step35 heads explicitly (rather than silently capturing
3861 /// a window-less, wrong-past-`win` graph) and this eager chain IS the served path. With
3862 /// the door armed (the default), BOTH draft modes run the `mtp_step35_attn_dcw` twin
3863 /// instead of this arm.
3864 ///
3865 /// Unlike the dc arm this advances BOTH the host `kv.len` and the device counter, so the
3866 /// caller must not mirror.
3867 #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
3868 fn mtp_step35_attn(
3869 &self,
3870 e: &Engine,
3871 fa: &FullAttnLayer,
3872 g: &crate::hybrid::Step35MtpGeom,
3873 h: &CudaSlice<f32>,
3874 pos_d: &CudaSlice<i32>,
3875 scratch: &mut MtpScratch,
3876 scratch_index: usize,
3877 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3878 let (nh, nkv, hd) = (g.n_head, g.n_head_kv, self.cfg.head_dim_k as usize);
3879 // MTP-GEOM RECEIPT, once per process, on the SERVED draft path. Slot-0 acceptance is
3880 // 0.725 here against 0.994 for vLLM MTP3 on the same checkpoint family and card class, and
3881 // the first three explanations for that gap were all wrong: head assignment (step-modulo
3882 // is index 0 at K=1, correct), MEMRA_SPEC_HPOST (identical 84/116 both arms), and this
3883 // block's geometry. Geometry was the one that could have failed SILENTLY — a wrong window
3884 // makes the draft attend the whole context instead of Step-3.7's 512, stays fluent, and
3885 // shows up only as acceptance — so it gets a standing receipt rather than another reading
3886 // of the source. Prints the resolved Step35MtpGeom the served path actually runs on;
3887 // `full_attention_geometry_at`'s missing-row fallback (window: None) does NOT reach here.
3888 {
3889 static ONCE: std::sync::OnceLock<()> = std::sync::OnceLock::new();
3890 ONCE.get_or_init(|| {
3891 eprintln!(
3892 "[mtp-geom] arm=eager block={} swa={} window={} n_head={nh} n_head_kv={nkv} \
3893 head_dim_k={hd} n_rot={} rope_base={} clamp_shexp={:?}",
3894 g.il, g.swa, g.window, g.n_rot, g.rope_base, g.clamp_shexp,
3895 );
3896 });
3897 }
3898 let eps = self.cfg.rms_eps;
3899 let scale = 1.0 / (hd as f32).sqrt(); // step35.cpp:255 kq_scale
3900 let n_embd = self.cfg.n_embd as usize;
3901 let gw = fa
3902 .attn_gate
3903 .as_ref()
3904 .ok_or("step35 MTP block is missing attn_gate.weight (head-wise attention gate)")?;
3905
3906 let (q0, k0, v0, gt) = if e.uses_q8_1_fast(&fa.wq)
3907 && e.uses_q8_1_fast(&fa.wk)
3908 && e.uses_q8_1_fast(&fa.wv)
3909 && e.uses_q8_1_fast(gw)
3910 {
3911 let (hq, hdq) = e.quantize_q8_1(h, 1, n_embd)?;
3912 let (a, b, c) = match e.matmul_q8_fused3(&fa.wq, &fa.wk, &fa.wv, &hq, &hdq)? {
3913 Some(t3) => t3,
3914 None => (
3915 e.matmul_pre(&fa.wq, &hq, &hdq, h, 1)?,
3916 e.matmul_pre(&fa.wk, &hq, &hdq, h, 1)?,
3917 e.matmul_pre(&fa.wv, &hq, &hdq, h, 1)?,
3918 ),
3919 };
3920 (a, b, c, e.matmul_pre(gw, &hq, &hdq, h, 1)?)
3921 } else {
3922 (
3923 e.matmul(&fa.wq, h, 1)?,
3924 e.matmul(&fa.wk, h, 1)?,
3925 e.matmul(&fa.wv, h, 1)?,
3926 e.matmul(gw, h, 1)?,
3927 )
3928 };
3929
3930 let mut q = e.uninit(nh * hd)?;
3931 e.rms_norm(&q0, fa.q_norm.float_data(), &mut q, hd, nh, eps)?;
3932 let mut k = e.uninit(nkv * hd)?;
3933 e.rms_norm(&k0, fa.k_norm.float_data(), &mut k, hd, nkv, eps)?;
3934 // `rope_freqs.weight` (llama3 factors) applies to the FULL-attn layers ONLY; SWA passes
3935 // null (llama-hparams / step35.cpp). Block 45 is SWA, so `ff` is None there — but read
3936 // the resolved flag, not the constant, so an all-full sibling stays correct.
3937 let ff = if g.swa {
3938 None
3939 } else {
3940 self.step35_aux.as_ref().and_then(|a| a.rope_freqs(e))
3941 };
3942 #[cfg(debug_assertions)]
3943 if let Some(ff) = ff {
3944 crate::debug_assert_tensor_stream_device(ff, &e.stream(), "mtp_step35_attn.rope_freqs");
3945 }
3946 e.rope_neox2(
3947 &mut q,
3948 &mut k,
3949 pos_d,
3950 hd,
3951 g.n_rot,
3952 nh,
3953 nkv,
3954 1,
3955 g.rope_base,
3956 1.0,
3957 ff,
3958 )?;
3959
3960 // Append at the HOST slot, then re-stamp the device counter: the eager chain has the
3961 // length on the host anyway, and the windowed view below needs it there to compute the
3962 // offset. The device counter is kept in lockstep so `mtp_kv_fill`'s `set_i32_one` and any
3963 // dc-family consumer of this scratch still agree.
3964 let (kv, scratch_cap) = scratch.plane_mut(scratch_index);
3965 assert!(
3966 kv.len < scratch_cap,
3967 "step35 MTP scratch overflow ({} >= {})",
3968 kv.len,
3969 scratch_cap
3970 );
3971 let next_len = kv.len + 1;
3972 let (off, t_kv) = if g.swa && next_len > g.window {
3973 (next_len - g.window, g.window)
3974 } else {
3975 (0, next_len)
3976 };
3977 // `off`/`t_kv` stay the ATTENTION view; the retain is a separate, lower bound so the
3978 // rewind that follows this append is still resident. THIS is the only site that rebases
3979 // this plane (MEMRA_KV_REBASE_TRACE, one run: 1 rebase, all from here), so it is the site
3980 // that decides `base` for everyone.
3981 let retain_from = match kv.ring.as_ref() {
3982 Some(ring) => memra_kv::swa_retain_from(kv.len, ring.window(), ring.base()),
3983 None => off & !31usize,
3984 };
3985 let write_row = e.prepare_kv_append(kv, retain_from, 1)?;
3986 e.append_kv_quantized(
3987 &k,
3988 &v0,
3989 &mut kv.k,
3990 &mut kv.v,
3991 write_row,
3992 kv.kv_dim_k,
3993 kv.kv_dim_v,
3994 kv.k_tok_bytes,
3995 kv.v_tok_bytes,
3996 false,
3997 )?;
3998 kv.len = next_len;
3999 e.set_i32_one(&mut kv.len_d, kv.len as i32)?;
4000 // SWA view offset (see note 1). The draft chain is short (k+2 rows), but the scratch is
4001 // PERSISTENT across rounds — `mtp_kv_fill` leaves one row per committed token behind, so
4002 // `kv.len` tracks absolute position and crosses 512 in any real generation. The window is
4003 // therefore live, not theoretical.
4004 let physical = kv.physical_rows(off, off + t_kv)?;
4005 let k_view = e.view_u8_range(
4006 &kv.k,
4007 physical.start * kv.k_tok_bytes,
4008 physical.end * kv.k_tok_bytes,
4009 );
4010 let v_view = e.view_u8_range(
4011 &kv.v,
4012 physical.start * kv.v_tok_bytes,
4013 physical.end * kv.v_tok_bytes,
4014 );
4015 let mut attn = e.uninit(nh * hd)?;
4016 e.fa_decode_kvmod(
4017 &q,
4018 &k_view,
4019 &v_view,
4020 &mut attn,
4021 hd,
4022 nh,
4023 nkv,
4024 t_kv,
4025 scale,
4026 kv.k_tok_bytes,
4027 kv.v_tok_bytes,
4028 false,
4029 )?;
4030
4031 let mut ag = e.uninit(nh * hd)?;
4032 e.attn_head_gate(&attn, >, &mut ag, None, hd, nh, 1)?;
4033 e.matmul(&fa.wo, &ag, 1)
4034 }
4035
4036 /// The dcw draft arm's kernel-class precondition, mirrored from `fa_decode_dcw`'s own
4037 /// refusal plus the v3 walk's format contract (`fa_v3_active`), so the DEV dispatch can
4038 /// never pick an arm the launcher would refuse mid-chain (the eager chain has no graceful
4039 /// fallback point) and the CAP site refuses with the named reason instead.
4040 ///
4041 /// `cap` = the SESSION's scratch-plane row capacity: the launcher's vec gate reads
4042 /// `bucket_max = min(window, cap)`, so a SMALL session (tiny prompt + tiny max_tokens,
4043 /// e.g. a max_tokens=8 probe: cap ~62 < the 96 vec floor) is OUTSIDE the dcw domain even
4044 /// though the WINDOW clears the floor. Mirroring the window alone shipped exactly that
4045 /// hole when the door default flipped ON (2026-08-30, vision-cell receipt: sampled
4046 /// capture WARN + `[engine-error] fa_decode_dcw supports the default v3-vec class only`
4047 /// hard-failing the burst — the eager dcw arm has no graceful fallback point). Sub-floor
4048 /// sessions now take the host-len kvmod arm, byte-for-byte the door-off serving.
4049 fn step35_dcw_eligible(&self, g: &crate::hybrid::Step35MtpGeom, cap: usize) -> bool {
4050 let hd = self.cfg.head_dim_k as usize;
4051 step35_draft_dcw_on()
4052 && g.swa
4053 && g.window.min(cap) >= crate::fa_vec_min_tkv()
4054 && std::env::var("MEMRA_NO_FA_VEC").is_err()
4055 && crate::fa_v3_active(hd)
4056 && hd <= 256
4057 && hd.is_multiple_of(32)
4058 }
4059
4060 /// step35 MTP-block attention, T=1, on the scratch KV: the WINDOWED DEVICE-COUNTER twin
4061 /// of `mtp_step35_attn`, serving BOTH draft paths when `step35_draft_dcw_on`. Write slot,
4062 /// key bound and SWA view offset all derive from device state (`len_d`, `base_d` written
4063 /// only at host-side rebases, and the block's `window`), so ONE captured graph serves the
4064 /// whole chain and replays see KV growth through the counter: the `mtp_full_attn_dc`
4065 /// contract plus the view offset the plain `_dc` kernel could not express (the old
4066 /// capture-refusal root cause). The three step35 properties stay per-geom exactly as in
4067 /// the eager twin: nh/nkv from `Step35MtpGeom`, the separate head-wise gate
4068 /// (`attn_head_gate`), per-layer rope width/base with SWA passing null freqs.
4069 ///
4070 /// bucket_max = min(cap, window): the windowed view never exceeds `window` rows, so the
4071 /// capture-time grid stays valid for every replayed len, and the kernel derives ns_eff
4072 /// from the LIVE T_kv at the fixed split_keys (one-partition law). Both arms call THIS
4073 /// launcher at THIS bucket, so eager and captured drafts are bit-identical by
4074 /// construction; vs the retired-by-flag `mtp_step35_attn` the only numeric-class deltas
4075 /// are the sub-vec-floor region (t_kv < 96: kvmod ran scalar, dcw stays vec) and any
4076 /// live-len split-ladder rung below the bucket's, both draft-side only (the verify
4077 /// arbitrates emitted bytes; acceptance is gated by the battery).
4078 ///
4079 /// Host len is NOT advanced here (graph contract); callers mirror. The EAGER caller runs
4080 /// `prepare_kv_append` per step (ring headroom, rebase legal there); the CAPTURED path
4081 /// pre-arms headroom at capture time and round start (`MtpScratch::ensure_dcw_headroom`)
4082 /// because a rebase is host work no captured chain may contain.
4083 #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the capture/call contract; bundling into a struct is a refactor, not a lint fix
4084 fn mtp_step35_attn_dcw(
4085 &self,
4086 e: &Engine,
4087 fa: &FullAttnLayer,
4088 g: &crate::hybrid::Step35MtpGeom,
4089 h: &CudaSlice<f32>,
4090 pos_d: &CudaSlice<i32>,
4091 scratch: &mut MtpScratch,
4092 scratch_index: usize,
4093 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4094 let (nh, nkv, hd) = (g.n_head, g.n_head_kv, self.cfg.head_dim_k as usize);
4095 // MTP-GEOM RECEIPT (dcw twin of the `mtp_step35_attn` receipt): once per process,
4096 // naming the arm, so a serving log proves WHICH draft attention program ran (the
4097 // engagement receipt for the flag door, both directions).
4098 {
4099 static ONCE: std::sync::OnceLock<()> = std::sync::OnceLock::new();
4100 ONCE.get_or_init(|| {
4101 eprintln!(
4102 "[mtp-geom] arm=dcw block={} swa={} window={} n_head={nh} n_head_kv={nkv} \
4103 head_dim_k={hd} n_rot={} rope_base={} clamp_shexp={:?}",
4104 g.il, g.swa, g.window, g.n_rot, g.rope_base, g.clamp_shexp,
4105 );
4106 });
4107 }
4108 let eps = self.cfg.rms_eps;
4109 let scale = 1.0 / (hd as f32).sqrt(); // step35.cpp:255 kq_scale
4110 let n_embd = self.cfg.n_embd as usize;
4111 let gw = fa
4112 .attn_gate
4113 .as_ref()
4114 .ok_or("step35 MTP block is missing attn_gate.weight (head-wise attention gate)")?;
4115
4116 let (q0, k0, v0, gt) = if e.uses_q8_1_fast(&fa.wq)
4117 && e.uses_q8_1_fast(&fa.wk)
4118 && e.uses_q8_1_fast(&fa.wv)
4119 && e.uses_q8_1_fast(gw)
4120 {
4121 let (hq, hdq) = e.quantize_q8_1(h, 1, n_embd)?;
4122 let (a, b, c) = match e.matmul_q8_fused3(&fa.wq, &fa.wk, &fa.wv, &hq, &hdq)? {
4123 Some(t3) => t3,
4124 None => (
4125 e.matmul_pre(&fa.wq, &hq, &hdq, h, 1)?,
4126 e.matmul_pre(&fa.wk, &hq, &hdq, h, 1)?,
4127 e.matmul_pre(&fa.wv, &hq, &hdq, h, 1)?,
4128 ),
4129 };
4130 (a, b, c, e.matmul_pre(gw, &hq, &hdq, h, 1)?)
4131 } else {
4132 (
4133 e.matmul(&fa.wq, h, 1)?,
4134 e.matmul(&fa.wk, h, 1)?,
4135 e.matmul(&fa.wv, h, 1)?,
4136 e.matmul(gw, h, 1)?,
4137 )
4138 };
4139
4140 let mut q = e.zeros(nh * hd)?;
4141 e.rms_norm(&q0, fa.q_norm.float_data(), &mut q, hd, nh, eps)?;
4142 let mut k = e.zeros(nkv * hd)?;
4143 e.rms_norm(&k0, fa.k_norm.float_data(), &mut k, hd, nkv, eps)?;
4144 // rope_freqs (llama3 factors) apply to the FULL-attn layers ONLY; SWA passes null
4145 // (the eager twin's rule, resolved from the flag, not the constant).
4146 let ff = if g.swa {
4147 None
4148 } else {
4149 self.step35_aux.as_ref().and_then(|a| a.rope_freqs(e))
4150 };
4151 #[cfg(debug_assertions)]
4152 if let Some(ff) = ff {
4153 crate::debug_assert_tensor_stream_device(
4154 ff,
4155 &e.stream(),
4156 "mtp_step35_attn_dcw.rope_freqs",
4157 );
4158 }
4159 e.rope_neox2(
4160 &mut q,
4161 &mut k,
4162 pos_d,
4163 hd,
4164 g.n_rot,
4165 nh,
4166 nkv,
4167 1,
4168 g.rope_base,
4169 1.0,
4170 ff,
4171 )?;
4172
4173 let (kv, cap) = scratch.plane_mut(scratch_index);
4174 // Append at the DEVICE slot's PHYSICAL row (len_d - base_d), then advance the counter
4175 // in-graph. Physical room is the callers' headroom contract (see the fn doc).
4176 e.append_kv_quantized_dcw(
4177 &k,
4178 &v0,
4179 &mut kv.k,
4180 &mut kv.v,
4181 &kv.len_d,
4182 kv.base_d.as_ref(),
4183 kv.kv_dim_k,
4184 kv.kv_dim_v,
4185 kv.k_tok_bytes,
4186 kv.v_tok_bytes,
4187 )?;
4188 e.inc_seqlen(&mut kv.len_d)?;
4189 // Full-buffer views (any in-round physical row stays in range under the headroom
4190 // contract); the kernel bounds and offsets the key range from (len_d, base_d, window).
4191 let k_view = e.view_u8(&kv.k, kv.k.len());
4192 let v_view = e.view_u8(&kv.v, kv.v.len());
4193 let bucket = g.window.min(cap);
4194 let mut attn = e.zeros(nh * hd)?;
4195 e.fa_decode_dcw(
4196 &q,
4197 &k_view,
4198 &v_view,
4199 &mut attn,
4200 hd,
4201 nh,
4202 nkv,
4203 &kv.len_d,
4204 kv.base_d.as_ref(),
4205 if g.swa { g.window } else { 0 },
4206 bucket,
4207 scale,
4208 kv.k_tok_bytes,
4209 kv.v_tok_bytes,
4210 None,
4211 )?;
4212
4213 let mut ag = e.zeros(nh * hd)?;
4214 e.attn_head_gate(&attn, >, &mut ag, None, hd, nh, 1)?;
4215 e.matmul(&fa.wo, &ag, 1)
4216 }
4217
4218 /// MTP-block full attention, T=1, on the scratch KV (BOTH draft paths — eager and graph):
4219 /// the scratch write slot and the attention bound come from `scratch.kv.len_d` (device i32[1])
4220 /// so the launch args are FIXED across draft steps — ONE captured graph serves the whole
4221 /// chain, and replays keep seeing KV growth through the device counter (no recapture).
4222 /// Geometry contract: n_splits is sized from `scratch.cap` (the persistent capacity); splits
4223 /// whose key range lies beyond the device t_kv exit empty and the shared combine skips them
4224 /// (fa_decode_dc bit-correct-for-any-t_kv<=bucket_max contract). The eager path uses the SAME
4225 /// launcher with the SAME bucket_max -> identical dispatch -> bit-identical draft tokens (the
4226 /// graph-vs-eager parity gate). Host len is NOT advanced here (graph contract); callers mirror.
4227 #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
4228 fn mtp_full_attn_dc(
4229 &self,
4230 e: &Engine,
4231 fa: &FullAttnLayer,
4232 h: &CudaSlice<f32>,
4233 pos_d: &CudaSlice<i32>,
4234 scratch: &mut MtpScratch,
4235 scratch_index: usize,
4236 geom: Option<&crate::hybrid::DraftGeom>,
4237 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4238 let cfg = &self.cfg;
4239 let mtp_il = cfg.n_layer.saturating_sub(cfg.nextn_predict_layers);
4240 let geometry = cfg.full_attention_geometry_at(mtp_il);
4241 let n_head = geom.map(|g| g.n_head).unwrap_or(geometry.n_head as usize);
4242 let n_head_kv = geom
4243 .map(|g| g.n_head_kv)
4244 .unwrap_or(geometry.n_head_kv as usize);
4245 let head_dim = geometry.head_dim_k as usize;
4246 let eps = cfg.rms_eps;
4247 let scale = geometry.attention_scale();
4248 let n_embd = geom.map(|g| g.d_inner).unwrap_or(cfg.n_embd as usize);
4249 let bucket_max = scratch.plane(scratch_index).1;
4250
4251 let (qf, mut k, v) =
4252 if e.uses_q8_1_fast(&fa.wq) && e.uses_q8_1_fast(&fa.wk) && e.uses_q8_1_fast(&fa.wv) {
4253 let (hq, hd) = e.quantize_q8_1(h, 1, n_embd)?;
4254 (
4255 e.matmul_pre(&fa.wq, &hq, &hd, h, 1)?,
4256 e.matmul_pre(&fa.wk, &hq, &hd, h, 1)?,
4257 e.matmul_pre(&fa.wv, &hq, &hd, h, 1)?,
4258 )
4259 } else {
4260 (
4261 e.matmul(&fa.wq, h, 1)?,
4262 e.matmul(&fa.wk, h, 1)?,
4263 e.matmul(&fa.wv, h, 1)?,
4264 )
4265 };
4266 // M3/Hy3 have no attention output gate — wq out is exactly q; skip the split.
4267 let gated = geometry.attention_gate == memra_gguf::config::AttentionGateKind::FusedQ;
4268 let (mut q, gate) = if gated {
4269 let mut q = e.zeros(n_head * head_dim)?;
4270 let mut gate = e.zeros(n_head * head_dim)?;
4271 e.q_gate_split(&qf, &mut q, &mut gate, head_dim, n_head, 1)?;
4272 (q, Some(gate))
4273 } else {
4274 (qf, None)
4275 };
4276
4277 let mut qn = e.zeros(n_head * head_dim)?;
4278 e.rms_norm(&q, fa.q_norm.float_data(), &mut qn, head_dim, n_head, eps)?;
4279 q = qn;
4280 let mut kn = e.zeros(n_head_kv * head_dim)?;
4281 e.rms_norm(
4282 &k,
4283 fa.k_norm.float_data(),
4284 &mut kn,
4285 head_dim,
4286 n_head_kv,
4287 eps,
4288 )?;
4289 k = kn;
4290 let rope_dims = geometry.n_rot as usize;
4291 e.rope_neox(
4292 &mut q,
4293 pos_d,
4294 head_dim,
4295 rope_dims,
4296 n_head,
4297 1,
4298 geometry.rope_base,
4299 1.0,
4300 )?;
4301 e.rope_neox(
4302 &mut k,
4303 pos_d,
4304 head_dim,
4305 rope_dims,
4306 n_head_kv,
4307 1,
4308 geometry.rope_base,
4309 1.0,
4310 )?;
4311
4312 let kv = scratch.plane_mut(scratch_index).0;
4313 // append at the DEVICE slot (kv.len_d == old len), then advance the counter in-graph.
4314 e.append_kv_quantized_dc(
4315 &k,
4316 &v,
4317 &mut kv.k,
4318 &mut kv.v,
4319 &kv.len_d,
4320 kv.kv_dim_k,
4321 kv.kv_dim_v,
4322 kv.k_tok_bytes,
4323 kv.v_tok_bytes,
4324 false,
4325 )?;
4326 e.inc_seqlen(&mut kv.len_d)?;
4327 // full-buffer views (any in-round t_kv stays in range on replay); the kernel bounds the
4328 // key range from the device counter.
4329 let k_view = e.view_u8(&kv.k, kv.k.len());
4330 let v_view = e.view_u8(&kv.v, kv.v.len());
4331 let (ktb, vtb) = (kv.k_tok_bytes, kv.v_tok_bytes);
4332 let mut attn = e.zeros(n_head * head_dim)?;
4333 e.fa_decode_dc(
4334 &q, &k_view, &v_view, &mut attn, head_dim, n_head, n_head_kv, &kv.len_d, bucket_max,
4335 scale, ktb, vtb, false,
4336 )?;
4337
4338 let attn_g = match &gate {
4339 Some(gate) => {
4340 let mut gsig = e.zeros(n_head * head_dim)?;
4341 e.sigmoid(gate, &mut gsig, n_head * head_dim)?;
4342 let mut ag = e.zeros(n_head * head_dim)?;
4343 e.mul(&attn, &gsig, &mut ag, n_head * head_dim)?;
4344 ag
4345 }
4346 None => attn,
4347 };
4348 e.matmul(&fa.wo, &attn_g, 1)
4349 }
4350
4351 /// PERSISTENT-DRAFT-KV fill (the reference engine's "mtp_update" analogue): compute the MTP
4352 /// block's K/V for `tokens` (committed tokens at positions pos0..pos0+T) from their EXACT
4353 /// trunk hiddens `h` ([T, n_embd] token-major, pre-output_norm) and append at slots pos0.. of
4354 /// the scratch KV. K/V-ONLY — ops A/1-5 plus the K-side of op 6 (wk/wv + k_norm + rope +
4355 /// quantized append); no wq/attention/FFN/lm_head, so per-token cost ~= eh_proj + wk/wv (a
4356 /// small fraction of one trunk layer), T-batched. Rope follows the chain convention
4357 /// rope(token@p) = p+1. Runs at round boundaries OUTSIDE the captured graph in BOTH draft
4358 /// modes -> draft parity by construction. Caller must have scratch.kv.len == pos0.
4359 #[allow(clippy::too_many_arguments)]
4360 fn mtp_kv_fill_at(
4361 &self,
4362 e: &Engine,
4363 mtp: &MtpHead,
4364 tokens: &[u32],
4365 h: &CudaSlice<f32>,
4366 pos0: usize,
4367 scratch: &mut MtpScratch,
4368 scratch_index: usize,
4369 embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
4370 ) -> Result<(), Box<dyn std::error::Error>> {
4371 let cfg = &self.cfg;
4372 let n_embd = cfg.n_embd as usize;
4373 let eps = cfg.rms_eps;
4374 let t = tokens.len();
4375 let (scratch_kv, scratch_cap) = scratch.plane(scratch_index);
4376 assert_eq!(scratch_kv.len, pos0, "mtp_kv_fill: append slot mismatch");
4377 assert!(pos0 + t <= scratch_cap, "mtp_kv_fill: scratch overflow");
4378 let Mixer::Full(fa) = &mtp.mixer else {
4379 panic!("MTP block is full-attn in qwen35; linear MTP not supported")
4380 };
4381 let pos_vec: Vec<i32> = (0..t).map(|i| (pos0 + i + 1) as i32).collect();
4382 let pos_d = e.htod_i32(&pos_vec)?;
4383
4384 // ops A/1/2: embed + the two input norms, T-wide.
4385 let e_emb = match embd_dev {
4386 Some((g, qt, rb)) => e.embed_gather_device_t(g, tokens, n_embd, qt, rb)?,
4387 None => e.htod(&self.embd.try_gather(n_embd, tokens)?)?,
4388 };
4389 let mut e_norm = e.zeros(t * n_embd)?;
4390 e.rms_norm(&e_emb, mtp.enorm.float_data(), &mut e_norm, n_embd, t, eps)?;
4391 let mut h_norm = e.zeros(t * n_embd)?;
4392 e.rms_norm(h, mtp.hnorm.float_data(), &mut h_norm, n_embd, t, eps)?;
4393
4394 // op 3: per-row [e_norm ; h_norm] concat, token-major [T, 2*n_embd].
4395 let mut concat = e.zeros(t * 2 * n_embd)?;
4396 for i in 0..t {
4397 e.copy_view_into(
4398 &mut concat,
4399 i * 2 * n_embd,
4400 &e_norm.slice(i * n_embd..(i + 1) * n_embd),
4401 n_embd,
4402 )?;
4403 e.copy_view_into(
4404 &mut concat,
4405 i * 2 * n_embd + n_embd,
4406 &h_norm.slice(i * n_embd..(i + 1) * n_embd),
4407 n_embd,
4408 )?;
4409 }
4410
4411 // ops 4/5: eh_proj + attn_norm, T-wide (at the student inner width when geom is set).
4412 let di = mtp.geom.as_ref().map(|g| g.d_inner).unwrap_or(n_embd);
4413 let inp_sa = e.matmul(&mtp.eh_proj, &concat, t)?;
4414 let mut a_norm = e.zeros(t * di)?;
4415 e.rms_norm(&inp_sa, mtp.attn_norm.float_data(), &mut a_norm, di, t, eps)?;
4416
4417 // op 6 (K/V half): wk/wv + k_norm + rope + per-row quantized append. No wq/attention —
4418 // the fill only has to leave correct K/V rows behind for later chains to attend over.
4419 let n_head_kv = mtp
4420 .geom
4421 .as_ref()
4422 .map(|g| g.n_head_kv)
4423 .or(mtp.step35.as_ref().map(|s| s.n_head_kv))
4424 .unwrap_or_else(|| {
4425 let mtp_il = cfg.n_layer.saturating_sub(cfg.nextn_predict_layers);
4426 cfg.full_attention_geometry_at(mtp_il).n_head_kv as usize
4427 });
4428 let mtp_il = cfg.n_layer.saturating_sub(cfg.nextn_predict_layers);
4429 let geometry = cfg.full_attention_geometry_at(mtp_il);
4430 let head_dim = geometry.head_dim_k as usize;
4431 let mut k = e.matmul(&fa.wk, &a_norm, t)?;
4432 let v = e.matmul(&fa.wv, &a_norm, t)?;
4433 let mut kn = e.zeros(t * n_head_kv * head_dim)?;
4434 e.rms_norm(
4435 &k,
4436 fa.k_norm.float_data(),
4437 &mut kn,
4438 head_dim,
4439 n_head_kv * t,
4440 eps,
4441 )?;
4442 k = kn;
4443 // step35: rotary width AND base are per-layer, and the MTP block's values come from the
4444 // resolved `Step35MtpGeom` — NOT from `cfg.rope_dim_count`/`cfg.rope_freq_base`, which
4445 // carry the arch defaults (128 / 5e6, i.e. the FULL-attn layers' base). Getting this wrong
4446 // writes K rows the attention arm then re-derives at a different theta: correct-looking
4447 // output with dead acceptance, invisible to the exactness gates.
4448 let (rope_dims, rope_base, ff) = match mtp.step35.as_ref() {
4449 Some(s) => (
4450 s.n_rot,
4451 s.rope_base,
4452 if s.swa {
4453 None
4454 } else {
4455 self.step35_aux.as_ref().and_then(|a| a.rope_freqs(e))
4456 },
4457 ),
4458 None => (geometry.n_rot as usize, geometry.rope_base, None),
4459 };
4460 #[cfg(debug_assertions)]
4461 if let Some(ff) = ff {
4462 crate::debug_assert_tensor_stream_device(ff, &e.stream(), "mtp_kv_fill.rope_freqs");
4463 }
4464 match ff {
4465 Some(f) => e.rope_neox_ff(
4466 &mut k, &pos_d, head_dim, rope_dims, n_head_kv, t, rope_base, 1.0, f,
4467 )?,
4468 None => e.rope_neox(
4469 &mut k, &pos_d, head_dim, rope_dims, n_head_kv, t, rope_base, 1.0,
4470 )?,
4471 }
4472
4473 let kv = scratch.plane_mut(scratch_index).0;
4474 // Match the trunk prime contract: a chunk may need the aligned window immediately before
4475 // its first row, so preserve that prefix when the physical tail rebases at wrap.
4476 let retain_from = kv
4477 .ring
4478 .as_ref()
4479 .map(|ring| memra_kv::swa_retain_from(pos0, ring.window(), ring.base()))
4480 .unwrap_or(0);
4481 let write_row = e.prepare_kv_append(kv, retain_from, t)?;
4482 for i in 0..t {
4483 let k_row = k.slice(i * kv.kv_dim_k..(i + 1) * kv.kv_dim_k);
4484 let v_row = v.slice(i * kv.kv_dim_v..(i + 1) * kv.kv_dim_v);
4485 e.append_kv_quantized_view(
4486 &k_row,
4487 &v_row,
4488 &mut kv.k,
4489 &mut kv.v,
4490 write_row + i,
4491 kv.kv_dim_k,
4492 kv.kv_dim_v,
4493 kv.k_tok_bytes,
4494 kv.v_tok_bytes,
4495 false,
4496 )?;
4497 }
4498 kv.len = pos0 + t;
4499 e.set_i32_one(&mut kv.len_d, kv.len as i32)?;
4500 Ok(())
4501 }
4502
4503 #[allow(clippy::too_many_arguments)]
4504 fn mtp_kv_fill_all(
4505 &self,
4506 e: &Engine,
4507 tokens: &[u32],
4508 h: &CudaSlice<f32>,
4509 pos0: usize,
4510 scratch: &mut MtpScratch,
4511 embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
4512 ) -> Result<(), Box<dyn std::error::Error>> {
4513 debug_assert_eq!(self.mtp_head_count(), scratch.plane_count());
4514 for index in 0..self.mtp_head_count() {
4515 self.mtp_kv_fill_at(
4516 e,
4517 self.mtp_head_at(index),
4518 tokens,
4519 h,
4520 pos0,
4521 scratch,
4522 index,
4523 embd_dev,
4524 )?;
4525 }
4526 Ok(())
4527 }
4528
4529 /// CAPTURE body for the GRAPH DRAFT (stage 2 of graph-grade spec): ONE MTP head forward with
4530 /// every varying input device-resident —
4531 /// - token id from the persistent `tok_d` (the previous replay's in-graph argmax wrote it,
4532 /// so the chain feeds itself; the host reads the same 4 bytes for the draft list),
4533 /// - h_seed from the persistent `h_seed_d` (h_nextn is copied BACK into it at the end),
4534 /// - rope pos from the persistent `pos_d` counter (inc'd in-graph),
4535 /// - scratch KV slot/bound from `scratch.kv.len_d` (see mtp_full_attn_dc).
4536 /// The p-min confidence lands in the persistent `p_d` iff `with_prob` (env is fixed per run).
4537 /// Same kernels, same dispatch as the eager mtp_head_forward_dev chain -> same draft tokens
4538 /// (exactness never depends on drafts — the verify arbitrates — but acceptance parity does).
4539 /// `with_head=false` captures the HEAD-LESS twin for the pseudo-seed replay (2026-07-03):
4540 /// the pseudo pass only needs h_nextn (op 10) + the scratch append — the lm_head read
4541 /// (~1.06ms q6_K on the 9B), argmax and prob are dead weight there. h_nextn's inputs are
4542 /// untouched, so the seed value is identical; round-start resets overwrite tok_d/p_d anyway.
4543 /// `sampled_cap` = Some((ctr_d, perturb_d, q_out_d, seed, temp)) captures the SAMPLED twin
4544 /// (step 3 of the sampled-spec arc): head logits are retained in the persistent `q_out_d`
4545 /// (host D2Ds them to the round's q slot after each replay), the DEVICE event counter is
4546 /// bumped in-graph, and the argmax reads GUMBEL-PERTURBED logits — one categorical draw per
4547 /// replay, bit-identical to the eager arm's gumbel_perturb at the same (seed, sctr, temp).
4548 /// seed/temp are capture-time constants (fixed per generate call, like p_min).
4549 #[allow(clippy::too_many_arguments)]
4550 #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
4551 fn mtp_head_forward_cap(
4552 &self,
4553 e: &Engine,
4554 mtp: &MtpHead,
4555 tok_d: &mut CudaSlice<u32>,
4556 pos_d: &mut CudaSlice<i32>,
4557 h_seed_d: &mut CudaSlice<f32>,
4558 p_d: &mut CudaSlice<f32>,
4559 scratch: &mut MtpScratch,
4560 // Which scratch plane this head appends to / attends over: 0 for the single-head
4561 // chain (every pre-lane caller), the head's own plane index for the multi-head
4562 // chain graphs (each head owns one plane — `mtp_chain_forward_dev`'s contract).
4563 scratch_index: usize,
4564 with_prob: bool,
4565 with_head: bool,
4566 embd_gpu: &CudaSlice<u8>,
4567 embd_qt: i32,
4568 embd_rb: usize,
4569 d_vocab: usize,
4570 sampled_cap: Option<SampledCapArgs<'_>>,
4571 stream_pack: Option<(&mut CudaSlice<u32>, usize, Option<&CudaSlice<u32>>)>,
4572 // DRAFT-SIDE GRAMMAR MASK (lane/draft-mask): (packed draft-vocab allowed-set buffer,
4573 // word count). Captured as ONE mask_logits_f32 node between the head matmul and the
4574 // in-graph argmax; the buffer address is baked, its CONTENTS are re-uploaded by the
4575 // host before every replay (the decode.rs graph-mask pattern). All-ones contents = a
4576 // no-op ban, so a position the grammar cannot constrain costs one pass over the row.
4577 mask_cap: Option<(&CudaSlice<u32>, usize)>,
4578 ) -> Result<(), Box<dyn std::error::Error>> {
4579 let cfg = &self.cfg;
4580 let n_embd = cfg.n_embd as usize;
4581 // step35: capturable through the WINDOWED device-counter arm (`mtp_step35_attn_dcw`)
4582 // once the dcw door is armed and the v3-vec class is live. Without the door this stays
4583 // the deliberate, named refusal: the plain `_dc` attention's key bound always starts at
4584 // row 0, cannot express this block's SWA view offset, and a captured chain would
4585 // silently attend OUTSIDE the window once the persistent scratch passes 512 rows.
4586 // Returning Err (not a panic) is what the capture sites already handle by degrading to
4587 // the eager chain (`mtp_head_forward_dev` -> `mtp_step35_attn`).
4588 // ROUND-STREAM stays refused EITHER WAY: the stream VERIFY has no step35 twin (see the
4589 // step35_verify refusal), so a stream capture that succeeded here would only move the
4590 // failure from capture time (graceful stream-off) to serve time (a failed round).
4591 if let Some(g) = mtp.step35.as_ref() {
4592 if stream_pack.is_some() {
4593 return Err(
4594 "step35 has no ROUND-STREAM draft arm (the stream verify has no step35 \
4595 twin); stream off"
4596 .into(),
4597 );
4598 }
4599 if !self.step35_dcw_eligible(g, scratch.plane(scratch_index).1) {
4600 return Err(format!(
4601 "step35 has no captured draft chain (fa_decode_dc cannot express the MTP \
4602 block's SWA view offset; the windowed dcw capture needs \
4603 MEMRA_STEP35_DRAFT_DCW armed [default ON, =0 disarms] and the v3-vec \
4604 class live at bucket=min(window {}, scratch cap {})) - the eager draft \
4605 chain serves this shape",
4606 g.window,
4607 scratch.plane(scratch_index).1,
4608 )
4609 .into());
4610 }
4611 }
4612 // student inner width (see mtp_head_forward_dev) — interface dims stay n_embd.
4613 let di = mtp.geom.as_ref().map(|g| g.d_inner).unwrap_or(n_embd);
4614 let eps = cfg.rms_eps;
4615 let e_emb = e.embed_gather_device(embd_gpu, tok_d, n_embd, embd_qt, embd_rb)?;
4616 let mut e_norm = e.zeros(n_embd)?;
4617 e.rms_norm(&e_emb, mtp.enorm.float_data(), &mut e_norm, n_embd, 1, eps)?;
4618 let mut h_norm = e.zeros(n_embd)?;
4619 e.rms_norm(
4620 &*h_seed_d,
4621 mtp.hnorm.float_data(),
4622 &mut h_norm,
4623 n_embd,
4624 1,
4625 eps,
4626 )?;
4627 let mut concat = e.zeros(2 * n_embd)?;
4628 e.copy_into(&mut concat, 0, &e_norm, n_embd)?;
4629 e.copy_into(&mut concat, n_embd, &h_norm, n_embd)?;
4630 let inp_sa = e.matmul(&mtp.eh_proj, &concat, 1)?;
4631 let mut a_norm = e.zeros(di)?;
4632 e.rms_norm(&inp_sa, mtp.attn_norm.float_data(), &mut a_norm, di, 1, eps)?;
4633 let attn_out = match (&mtp.mixer, mtp.step35.as_ref()) {
4634 // step35 (eligibility already enforced by the refusal above): the windowed dcw
4635 // arm, the SAME launcher the eager dev arm runs when the door is armed. No host
4636 // work here (this is the capture body); headroom is the callers' pre-arm.
4637 (Mixer::Full(fa), Some(g)) => {
4638 self.mtp_step35_attn_dcw(e, fa, g, &a_norm, pos_d, scratch, scratch_index)?
4639 }
4640 (Mixer::Full(fa), None) => self.mtp_full_attn_dc(
4641 e,
4642 fa,
4643 &a_norm,
4644 pos_d,
4645 scratch,
4646 scratch_index,
4647 mtp.geom.as_ref(),
4648 )?,
4649 (Mixer::Linear(_), _) => {
4650 panic!("MTP block is full-attn in qwen35; linear MTP not supported")
4651 }
4652 (Mixer::Mla(_), _) => {
4653 crate::hybrid::mla_path_unimplemented("captured MTP head forward")
4654 }
4655 (Mixer::Kda(_), _) => {
4656 crate::hybrid::kda_path_unimplemented("captured MTP head forward")
4657 }
4658 };
4659 let mut x1 = e.zeros(di)?;
4660 e.add(&inp_sa, &attn_out, &mut x1, di)?;
4661 let mut z = e.zeros(di)?;
4662 e.rms_norm(&x1, mtp.post_attn_norm.float_data(), &mut z, di, 1, eps)?;
4663 let ffn_out = match &mtp.ffn {
4664 crate::hybrid::Ffn::Dense {
4665 ffn_gate,
4666 ffn_up,
4667 ffn_down,
4668 } => {
4669 let n_ff = ffn_gate.out_features();
4670 let (gate, up) = if e.uses_q8_1_fast(ffn_gate) && e.uses_q8_1_fast(ffn_up) {
4671 let (zq, zd) = e.quantize_q8_1(&z, 1, di)?;
4672 (
4673 e.matmul_pre(ffn_gate, &zq, &zd, &z, 1)?,
4674 e.matmul_pre(ffn_up, &zq, &zd, &z, 1)?,
4675 )
4676 } else {
4677 (e.matmul(ffn_gate, &z, 1)?, e.matmul(ffn_up, &z, 1)?)
4678 };
4679 let mut act = e.zeros(n_ff)?;
4680 // step35: the dense FFN reads the per-layer SHEXP clamp, resolved for the MTP
4681 // block's own index (the mtp_head_forward_dev rule; None for every other arch,
4682 // which is `ffn_act`'s dispatch verbatim). The eager and captured chains must
4683 // run the ONE activation program.
4684 Self::ffn_act_lim(
4685 e,
4686 &self.cfg,
4687 &gate,
4688 &up,
4689 1.0,
4690 1.0,
4691 mtp.step35
4692 .as_ref()
4693 .and_then(|s| s.clamp_shexp)
4694 .map(SwigluClamp::Post),
4695 &mut act,
4696 n_ff,
4697 )?;
4698 e.matmul(ffn_down, &act, 1)?
4699 }
4700 // ROUND-STREAM: a softmax-routed resident MoE takes the zero-D2H device router +
4701 // expert program and is capture-legal. Sigmoid-routed MoE (Hy3/M3/Step) still
4702 // selects through the host-visible sigmoid router; capturing that stream sync
4703 // invalidates CUDA capture, so it stays on the eager draft chain even when every
4704 // expert is resident. Non-resident (SLRU-lock) is likewise rejected.
4705 crate::hybrid::Ffn::Moe(m)
4706 if m.dev_exps.is_some() && self.cfg.sigmoid_router().is_none() =>
4707 {
4708 self.moe_ffn_il(e, m, &z, 1, u16::MAX)?
4709 }
4710 crate::hybrid::Ffn::Moe(_) => {
4711 return Err(
4712 "graph draft requires a Dense or device-routed resident-MoE MTP FFN".into(),
4713 );
4714 }
4715 };
4716 let mut h_inner = e.zeros(di)?;
4717 e.add(&x1, &ffn_out, &mut h_inner, di)?;
4718 // student: up-project back to n_embd (carrier + head input; see mtp_head_forward_dev).
4719 let h_nextn = match mtp.geom.as_ref() {
4720 Some(g) => e.matmul(&g.out_up, &h_inner, 1)?,
4721 None => h_inner,
4722 };
4723 // MEMRA_SPEC_HPOST needs final_h even head-less (it IS the next seed under that convention).
4724 let final_h = if with_head || spec_hpost() {
4725 let final_norm = mtp.shared_head_norm.as_ref().unwrap_or(&self.output_norm);
4726 let mut fh = e.zeros(n_embd)?;
4727 e.rms_norm(&h_nextn, final_norm.float_data(), &mut fh, n_embd, 1, eps)?;
4728 Some(fh)
4729 } else {
4730 None
4731 };
4732 if with_head {
4733 let head = mtp.shared_head_head.as_ref().unwrap_or(&self.output);
4734 let mut logits = e.matmul(head, final_h.as_ref().unwrap(), 1)?;
4735 // DRAFT-SIDE GRAMMAR MASK: ban the grammar-illegal draft ids IN the captured chain,
4736 // before the argmax — proposals become legal by construction. Contents-only
4737 // per-replay upload keeps the capture valid.
4738 if let Some((mask_d, mw)) = mask_cap {
4739 e.mask_logits_col(&mut logits, mask_d, 0, d_vocab, mw)?;
4740 }
4741 if let Some(SampledCapArgs {
4742 ctr: ctr_d,
4743 perturb: perturb_d,
4744 q_out: q_out_d,
4745 seed,
4746 temp,
4747 filt,
4748 }) = sampled_cap
4749 {
4750 // SAMPLED chain: retain q (raw head logits -> persistent q_out_d; the matmul's
4751 // own buffer is pool-recycled after the capture body returns, so it can't be the
4752 // retention target), bump the device event counter, gumbel-perturb reading it,
4753 // and argmax the PERTURBED logits into tok_d — the in-graph categorical draw.
4754 e.copy_into(q_out_d, 0, &logits, d_vocab)?;
4755 e.sctr_inc(ctr_d)?;
4756 match filt {
4757 // PURE-TEMP: gumbel over the raw softmax — byte-identical to the
4758 // pre-lane capture body.
4759 None => e.gumbel_perturb_ctr(&logits, perturb_d, d_vocab, seed, ctr_d, temp)?,
4760 // FILTERED (lane/step37-draft-graph-serving-20260830): the SAME
4761 // filter_stats program the eager arm and the accept path run (the
4762 // wrapper's coop/plain choice is deployment-keyed, never per-call), then
4763 // the device-stat/device-counter perturb twin — the draft draws from the
4764 // exact filtered distribution the verify gathers `q` from. q was
4765 // retained ABOVE, pre-perturb, so the accept path's post-replay stats
4766 // recompute (same kernel, same bits) reconstructs these th/z exactly.
4767 Some(f) => {
4768 e.filter_stats(
4769 &logits, d_vocab, f.rows0, f.th, f.z, f.mx, d_vocab, 1, temp, f.top_k,
4770 f.top_p, f.min_p,
4771 )?;
4772 e.gumbel_perturb_filtered_ctr(
4773 &logits, perturb_d, d_vocab, seed, ctr_d, temp, f.mx, f.th,
4774 )?;
4775 }
4776 }
4777 e.argmax_token_device_into(perturb_d, tok_d, d_vocab)?;
4778 // p-min prob = the head's RAW softmax confidence in the SAMPLED pick — same
4779 // semantics as the eager sampled arm's prob_of_token_device(dl_d, tok_d).
4780 if with_prob {
4781 e.prob_of_token_device_into(&logits, tok_d, p_d, d_vocab)?;
4782 }
4783 } else {
4784 // draft token -> persistent tok_d (next replay's embed reads it; host reads the 4 bytes).
4785 e.argmax_token_device_into(&logits, tok_d, d_vocab)?;
4786 // p-min under a draft mask reads the MASKED row: confidence relative to the
4787 // grammar-LEGAL alternatives (illegal ids leave the softmax denominator), which
4788 // is the right semantics for "does the drafter know what comes next here" and
4789 // the same row the pick came from. Draft-quality only — verify arbitrates.
4790 if with_prob {
4791 e.prob_of_token_device_into(&logits, tok_d, p_d, d_vocab)?;
4792 }
4793 }
4794 }
4795 // ROUND-STREAM K-chain: pack (tok, p) into slot j, then remap tok through d2t so the
4796 // NEXT chained body's embed reads the TARGET id — zero host involvement per step.
4797 if let Some((out, slot, d2t)) = stream_pack {
4798 e.pack_tok_p(tok_d, p_d, out, slot)?;
4799 if let Some(map) = d2t {
4800 e.tok_map_u32(tok_d, map)?;
4801 }
4802 }
4803 // Next draft step's h_seed: pre-norm h_nextn (default) or post-norm final_h (HPOST).
4804 if spec_hpost() {
4805 e.copy_into(h_seed_d, 0, final_h.as_ref().unwrap(), n_embd)?;
4806 } else {
4807 e.copy_into(h_seed_d, 0, &h_nextn, n_embd)?;
4808 }
4809 // advance the draft rope position in-graph.
4810 e.inc_seqlen(pos_d)?;
4811 Ok(())
4812 }
4813
4814 /// Batched target verify forward over `tokens` at positions `pos0..pos0+T` (§D.3, T=K+1).
4815 /// Returns ALL T logit columns (host f32, [T*n_vocab]); appends T cols to every full-attn KV
4816 /// and advances every linear-attn recur state by T steps (the recur steps are SEQUENTIAL T=1).
4817 /// Advances `cache.pos` by T.
4818 pub fn decode_step_t(
4819 &self,
4820 e: &Engine,
4821 tokens: &[u32],
4822 pos0: usize,
4823 cache: &mut Cache,
4824 ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
4825 if self.is_gemma4_e4b() {
4826 return Ok(self.gemma4_e4b_decode_step_t_h(e, tokens, pos0, cache)?.0);
4827 }
4828 if self.gemma_batch_program() {
4829 return self.gemma4_decode_step_t(e, tokens, pos0, cache);
4830 }
4831 Ok(self.decode_step_t_h(e, tokens, pos0, cache)?.0)
4832 }
4833
4834 /// Like `decode_step_t` but ALSO returns the LAST column's pre-output_norm hidden (h_seed for
4835 /// the next draft round). This lets partial-accept replay run as ONE batched T=(n_acc+1) forward
4836 /// (single weight read) instead of n_acc+1 separate T=1 decode_steps (n_acc+1 weight reads).
4837 /// At batch=1 decode is bandwidth-bound, so batching the replay is THE MTP profitability lever.
4838 pub fn decode_step_t_h(
4839 &self,
4840 e: &Engine,
4841 tokens: &[u32],
4842 pos0: usize,
4843 cache: &mut Cache,
4844 ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4845 self.decode_step_t_h_emb(e, tokens, pos0, cache, None)
4846 }
4847
4848 /// Like `decode_step_t_h` with an optional RESIDENT embed table (spec hot loop): device
4849 /// gather instead of host dequant + [T, n_embd] f32 htod. Bit-identical rows.
4850 pub fn decode_step_t_h_emb(
4851 &self,
4852 e: &Engine,
4853 tokens: &[u32],
4854 pos0: usize,
4855 cache: &mut Cache,
4856 embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
4857 ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4858 let (logits_d, h_seed) = self.decode_step_t_h_emb_dev(e, tokens, pos0, cache, embd_dev)?;
4859 Ok((e.dtoh(&logits_d)?, h_seed))
4860 }
4861
4862 /// DEVICE-LOGITS verify forward (spec device-argmax lever): identical kernel chain to
4863 /// `decode_step_t_h_emb` but returns the [T, n_vocab] logits ON DEVICE — the accept walk
4864 /// argmaxes each column on-device and reads back ONE [T] u32 instead of dtoh'ing the full
4865 /// T x n_vocab f32 block (~1-4 MB + T host argmaxes, every round). Kernel dispatch is
4866 /// UNCHANGED (same decode-exact kernels); only the post-logits transfer moves.
4867 pub fn decode_step_t_h_emb_dev(
4868 &self,
4869 e: &Engine,
4870 tokens: &[u32],
4871 pos0: usize,
4872 cache: &mut Cache,
4873 embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
4874 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4875 cache.ensure_usable("decode_step_t")?;
4876 let n_embd = self.cfg.n_embd as usize;
4877 let t = tokens.len();
4878 let (logits, x) = self.decode_step_t_core(e, tokens, pos0, cache, embd_dev, None)?;
4879 // h_seed for the next round = LAST column's pre-output_norm hidden ([n_embd]).
4880 let mut hs = vbuf(e, n_embd)?; // fully written by copy_view_into below
4881 e.copy_view_into(&mut hs, 0, &x.slice((t - 1) * n_embd..t * n_embd), n_embd)?;
4882 Ok((logits, hs))
4883 }
4884
4885 /// CORE verify forward: the `decode_step_t_h_emb_dev` kernel chain, returning the FULL
4886 /// pre-output_norm hidden stack x ([T, n_embd], any column extractable) and optionally
4887 /// filling a `VerifyCkpt` (retained per-layer state-rebuild inputs) for the REPLAY-FREE
4888 /// partial accept. `ckpt: None` => byte-for-byte the old behavior (the ckpt writes are pure
4889 /// retains/copies — they never change what any kernel computes).
4890 fn decode_step_t_core(
4891 &self,
4892 e: &Engine,
4893 tokens: &[u32],
4894 pos0: usize,
4895 cache: &mut Cache,
4896 embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
4897 mut ckpt: Option<&mut VerifyCkpt>,
4898 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4899 self.decode_step_t_core_stream(
4900 e,
4901 tokens,
4902 pos0,
4903 cache,
4904 embd_dev,
4905 ckpt.take(),
4906 None,
4907 None,
4908 None,
4909 )
4910 }
4911
4912 /// [`Self::decode_step_t_core`] with the MTP route's verify-graph pool armed
4913 /// (`MEMRA_SPEC_VERIFY_GRAPH`). `graphs: None` reproduces `decode_step_t_core`
4914 /// argument-for-argument, so the eager walk stays the byte-identical fallback.
4915 #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
4916 fn decode_step_t_core_vg(
4917 &self,
4918 e: &Engine,
4919 tokens: &[u32],
4920 pos0: usize,
4921 cache: &mut Cache,
4922 embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
4923 mut ckpt: Option<&mut VerifyCkpt>,
4924 graphs: Option<&mut DsparkVerifyGraphs>,
4925 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4926 self.decode_step_t_core_stream(
4927 e,
4928 tokens,
4929 pos0,
4930 cache,
4931 embd_dev,
4932 ckpt.take(),
4933 None,
4934 None,
4935 graphs,
4936 )
4937 }
4938
4939 /// ROUND-STREAM stage (c) 4: `stream` = (device verify tokens [t], device pos counter) —
4940 /// when Some, rope positions come from pos_iota over the counter, the embed gathers the
4941 /// device tokens, and full_attn_verify routes appends/FA through the _dc twins reading the
4942 /// SAME counter (every layer's kvl.len == cache.pos, one counter drives all three). The
4943 /// host `tokens`/`pos0` args still size buffers (t is FIXED K+1 in stream mode).
4944 /// `vtok_dev` (engine-bundle slice 2): device verify tokens for the EMBED only —
4945 /// unlike `stream` mode it changes nothing else (host pos iota, host-len KV appends).
4946 /// `tokens` then only sizes buffers (the dummy-slice pattern the round-stream arm uses).
4947 #[allow(clippy::too_many_arguments)]
4948 fn decode_step_t_core_stream(
4949 &self,
4950 e: &Engine,
4951 tokens: &[u32],
4952 pos0: usize,
4953 cache: &mut Cache,
4954 embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
4955 mut ckpt: Option<&mut VerifyCkpt>,
4956 stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
4957 vtok_dev: Option<&CudaSlice<u32>>,
4958 graphs: Option<&mut DsparkVerifyGraphs>,
4959 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4960 // PP DOOR (lane/pp2-spec 2026-08-06): the verify trunk now takes its OWN stage split,
4961 // exactly as the eager and batched steps do. This is the single funnel every verify
4962 // forward reaches (decode_step_t / _h / _h_emb / _h_emb_dev / _core all land here), so
4963 // wiring it here wires the whole spec surface — the draft/accept/commit machinery above
4964 // is untouched.
4965 //
4966 // History: pp2-hardening (2026-08-06) made this funnel FAIL CLOSED, because its trunk
4967 // walk was unsplit on one stream and a sharded cross-device placement peer-read every
4968 // remote layer's weights on every spec round (measured 13.9-28x on the batched twin).
4969 // The refusal below survives to cover the residue — MEMRA_SPEC_PP=0, MEMRA_PP_STREAMS=0,
4970 // or a placement whose PpNRt fails to build — so a config that would still walk the
4971 // whole trunk on one stream refuses instead of regressing 28x.
4972 if let Some(fence) = crate::pp::pp_cuts(self.layers.len())
4973 && !crate::pp::pp2_streams_off()
4974 && crate::pp::spec_pp_on()
4975 {
4976 if vtok_dev.is_some() {
4977 return Err(
4978 "device-token dspark verify (slice-2 deferred readback) has no PP \
4979 stage-split arm; run the dspark route on one device"
4980 .into(),
4981 );
4982 }
4983 return self.decode_step_t_core_ppn(
4984 e,
4985 tokens,
4986 pos0,
4987 cache,
4988 embd_dev,
4989 ckpt.take(),
4990 stream,
4991 &fence,
4992 );
4993 }
4994 crate::pp::refuse_unsplit_if_remote(
4995 "decode_step_t (spec verify)",
4996 "drop MEMRA_SPEC_PP=0 / MEMRA_PP_STREAMS=0 so the verify trunk takes its OWN stage \
4997 split (decode_step_t_core_ppn); or run spec on one device",
4998 )?;
4999 let cfg = &self.cfg;
5000 let n_embd = cfg.n_embd as usize;
5001 let eps = cfg.rms_eps;
5002 let t = tokens.len();
5003 let pos_d = match stream {
5004 Some((_, ctr)) => {
5005 let mut p = e.alloc_uninit::<i32>(t)?;
5006 e.pos_iota(ctr, &mut p, t)?;
5007 p
5008 }
5009 None => {
5010 let pos_vec: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
5011 e.htod_i32(&pos_vec)?
5012 }
5013 };
5014
5015 // embed T tokens -> [T, n_embd] token-major (device gather on the spec hot loop)
5016 let x = match (stream, embd_dev) {
5017 (Some((vtok, _)), Some((g, qt, rb))) => {
5018 e.embed_gather_device_td(g, vtok, t, n_embd, qt, rb)?
5019 }
5020 (None, Some((g, qt, rb))) => match vtok_dev {
5021 // slice 2: device verify tokens, same embed_gather_u32_t kernel —
5022 // bit-identical rows to the host-token arm (same per-dtype deq).
5023 Some(vt_d) => e.embed_gather_device_td(g, vt_d, t, n_embd, qt, rb)?,
5024 None => e.embed_gather_device_t(g, tokens, n_embd, qt, rb)?,
5025 },
5026 _ => {
5027 assert!(
5028 vtok_dev.is_none(),
5029 "device-token verify requires the resident embed table (embd_dev)"
5030 );
5031 e.htod(&self.embd.try_gather(n_embd, tokens)?)?
5032 }
5033 };
5034
5035 // TRUNK WALK: layers [0, n_layers) through the SAME range-scoped subgraph the PP-N
5036 // stage split calls per stage (`verify_layers`) — one code path, so the split cannot
5037 // drift from the unsplit dispatch mirroring. lane/pp2-spec 2026-08-06.
5038 let x = self.verify_layers(
5039 e,
5040 x,
5041 0,
5042 self.layers.len(),
5043 &pos_d,
5044 pos0,
5045 t,
5046 cache,
5047 ckpt.take(),
5048 stream,
5049 graphs,
5050 )?;
5051 if spec_nan_scan() {
5052 nan_scan_rows(e, &x, t, n_embd, &format!("verify trunk exit pos0={pos0}"))?;
5053 }
5054
5055 let mut hn = vbuf(e, t * n_embd)?;
5056 // Stage-A door: with the serving-class row-outer verify walk, the TAIL must be the
5057 // t=1 decode program per row too (rms_norm t=1 + the single-row bf16 head — the
5058 // split head's concat is receipted bit-identical to it). The batched cuBLASLt head
5059 // is a different ULP class and flips near-tie argmaxes off the greedy tape.
5060 let eager_tail = self.sliding_gated_moe_batch_program() && spec_verify_eager_on();
5061 if eager_tail {
5062 let n_vocab = self.cfg.n_vocab as usize;
5063 // MEMRA_SPEC_HEAD_ROWS=1 — THE VERIFY TAIL'S REDUNDANT HEAD READ.
5064 //
5065 // The loop below runs the head at m=1 once PER COLUMN, so the LM head's weights are
5066 // streamed t times per verify pass. On step37 that head is ~0.49 GiB per card after the
5067 // rank split, ~1.07 ms of pure re-read at t=2 and worse at every wider t — which is a
5068 // large part of why the fixed K ladder LOSES (K=1 81.2 > K=2 73.1 > K=3 62.7 tok/s).
5069 //
5070 // The loop's justification is the comment above: the batched cuBLASLt head is a
5071 // different ULP class and flips near-tie argmaxes off the greedy tape. That is true of
5072 // cuBLASLt and it does NOT apply here, because a FloatBf16 head at 1..=32 rows never
5073 // reaches cuBLASLt: `matmul` routes it to `matvec_bf16_rows_into` (lib.rs:12248), whose
5074 // own doc says `matvec_bf16_f32acc_x4_rows` "runs the t=1 decode head program PER ROW
5075 // (identical dot + reduce), so decode/verify tiers keep the t=1 numeric class". Under
5076 // the W8 doors both widths route to the q8 mirror instead, and the t-column mirror is
5077 // documented "bit-identical to t single-row calls". So the batched form is the SAME
5078 // arithmetic per row on both paths, with one weight read instead of t.
5079 //
5080 // rms_norm is row-wise, so norm(t) is per-row identical to t x norm(1) by construction.
5081 //
5082 // DEFAULT OFF for exactly one turn of the crank: "bit-identical by two documented
5083 // claims" is still an argument. The greedy byte tape decides, and the door flips only
5084 // once the tape is a receipt.
5085 if head_rows_on() {
5086 e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
5087 let logits = e.matmul(&self.output, &hn, t)?;
5088 if stream.is_none() {
5089 cache.pos += t;
5090 }
5091 return Ok((logits, if spec_hpost() { hn } else { x }));
5092 }
5093 let mut logits = vbuf(e, t * n_vocab)?;
5094 for r in 0..t {
5095 let mut row = e.uninit(n_embd)?;
5096 e.dtod_copy_view(&x.slice(r * n_embd..(r + 1) * n_embd), &mut row)?;
5097 let mut hr = e.uninit(n_embd)?;
5098 e.rms_norm(&row, self.output_norm.float_data(), &mut hr, n_embd, 1, eps)?;
5099 let lr = e.matmul(&self.output, &hr, 1)?;
5100 e.dtod_copy_into(&lr, &mut logits, r * n_vocab)?;
5101 e.dtod_copy_into(&hr, &mut hn, r * n_embd)?;
5102 }
5103 if stream.is_none() {
5104 cache.pos += t;
5105 }
5106 return Ok((logits, if spec_hpost() { hn } else { x }));
5107 }
5108 let serving_head =
5109 self.sliding_gated_moe_batch_program() || self.batched_serving_numeric_class();
5110 let logits = if serving_head {
5111 // Step35 and the qwen35 family (MoE 2026-08-14 AM, dense-hybrid same day PM — the
5112 // Q3.8 bring-up reproduced the identical near-tie class on dense: eager-class verify
5113 // vs batched-class live serving, ULP drift amplified through the GDN recurrence)
5114 // serve one batched numeric class at every live width, including B=1. Keep the
5115 // verify head in that same class; other generic families retain the decode-exact
5116 // head that their run-spec contract pins.
5117 e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
5118 e.matmul(&self.output, &hn, t)?
5119 } else {
5120 e.rms_norm_decode(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
5121 e.matmul_decode_exact(&self.output, &hn, t)?
5122 };
5123 // stream: the device pos counter owns position; host mirror reconciles at drain.
5124 if stream.is_none() {
5125 cache.pos += t;
5126 }
5127 // Hidden stack for seeds/refresh-fills: pre-norm x (default) or post-norm hn (HPOST).
5128 Ok((logits, if spec_hpost() { hn } else { x }))
5129 }
5130
5131 /// THE VERIFY TRUNK OVER PP-N (lane/pp2-spec 2026-08-06): `decode_step_t_core_stream`'s walk
5132 /// as N stage subgraphs, each on its own engine/stream (and, under `MEMRA_PP_DEVICES`, its own
5133 /// device), with a `[T, n_embd]` boundary transfer between them. T = K+1 (the verify batch),
5134 /// so this is the batched-boundary shape the pp2-batch lane's grow-only slots already handle
5135 /// (`tx(b, x, t*n_embd)`; the slot grows to the high-water T and the transport moves exactly
5136 /// the payload).
5137 ///
5138 /// Structure is `decode_step_batch_ppn`'s, which is `decode_step_h_ppn`'s. FOUR THINGS ARE
5139 /// PER-STAGE and each for a measured reason (see `decode_step_batch_ppn`'s header for the
5140 /// receipts):
5141 ///
5142 /// 1. THE ENGINE (`rt.engine(s, e)`) — `Engine` owns lazily-grown stable-pointer scratch
5143 /// (`fa_part_pool`, `fa_vf16_scratch`, `argmax_partials`) that is single-stream-safe BY
5144 /// DESIGN. Two stage streams through one Engine is the 2026-08-02 shared-scratch race
5145 /// (35% flake, nondeterministic all-logits divergence). `PpNRt::build` gives every stage
5146 /// s>0 its own Engine even on the primary device; honouring it here is what scopes the
5147 /// pools. The verify path allocates MORE of that scratch than eager decode does (FA at
5148 /// m=T, and the per-layer `GdnStash` retains), so this is load-bearing, not inherited.
5149 ///
5150 /// 2. `pos_d` — each stage uploads its OWN copy of the T rope positions on ITS stream, so the
5151 /// buffer is allocated, consumed and freed on one stream. In `stream` mode that means each
5152 /// stage runs its own `pos_iota` over the SHARED device counter (`pos_ctr`): the counter is
5153 /// read-only during the forward (the round's `inc`/`copy_add` happen outside it), so every
5154 /// stage derives the identical iota, and each stage's own output buffer is stream-local.
5155 ///
5156 /// 3. THE EMBED lives with stage 0 (`self.embd` / `embd_gpu` are host/primary-side; the
5157 /// sharded loader leaves the table with stage 0 by construction).
5158 ///
5159 /// 4. THE HEAD (`output_norm` + `output`) runs on the LAST stage — the sharded loader uploaded
5160 /// both through that stage's engine (`hybrid.rs`: `e_head = layer_engine(e, n_trunk,
5161 /// n_trunk-1)`), so reading them anywhere else is a peer read of the biggest tensor in the
5162 /// model, every round.
5163 ///
5164 /// WHAT STAYS ON THE PRIMARY, deliberately: the returned logits and hidden stack `x`. Both are
5165 /// last-stage-allocated device buffers, and every consumer (the device argmax walk, the accept
5166 /// kernels, `spec_seed_gather`, the ckpt rebuild in `commit_verified_prefix`) reads them
5167 /// through the primary context by UVA — the same read the batched serving epilogue's
5168 /// `last_logits_dev` park does. Those consumers are per-round O(T x n_vocab) and O(n_embd),
5169 /// not per-layer, so they are not the 28x class; splitting them is a separate lane.
5170 ///
5171 /// The MTP HEAD (draft side) is NOT split: it is one block, it lives wherever the loader put
5172 /// it (`load_mtp` uses the primary engine), and it is ~1-2 GB against the trunk's tens. Draft
5173 /// placement is measured, not assumed — see `research/pp2-spec-20260806`.
5174 ///
5175 /// EXACTNESS: PP-N adds ZERO deviation. Each stage runs the SAME kernels on the SAME bytes in
5176 /// the same order via the SAME `verify_layers` the unsplit body calls; the only change is
5177 /// where the residual is materialized, and the boundary is a straight f32 copy (dtod
5178 /// same-device / `cudaMemcpyPeerAsync` cross-device, no conversion). So the split MUST be
5179 /// BIT-IDENTICAL to the unsplit verify at the same T, in both placement orders. Gate:
5180 /// `decode-batch-gate --mode ppspec`. Acceptance counts are a DERIVED consequence — greedy
5181 /// accept argmaxes these logits, so bit-identical logits force identical accept walks; the
5182 /// `run-spec` K=1..8 arm checks that end-to-end rather than trusting the implication.
5183 #[allow(clippy::too_many_arguments)]
5184 fn decode_step_t_core_ppn(
5185 &self,
5186 e: &Engine,
5187 tokens: &[u32],
5188 pos0: usize,
5189 cache: &mut Cache,
5190 embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
5191 mut ckpt: Option<&mut VerifyCkpt>,
5192 stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
5193 fence: &[usize],
5194 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
5195 let ticket = self.verify_stage0_issue(
5196 e,
5197 tokens,
5198 pos0,
5199 cache,
5200 embd_dev,
5201 ckpt.as_deref_mut(),
5202 stream,
5203 fence,
5204 )?;
5205 self.verify_stage1_finish(e, ticket, cache, ckpt, stream, fence, true)
5206 }
5207
5208 /// Enqueue embed, stage 0, and the first boundary TX, then return the actual boundary slot.
5209 /// The ordinary PP verify wrapper calls `verify_stage1_finish` immediately after this return.
5210 #[allow(clippy::too_many_arguments)]
5211 fn verify_stage0_issue(
5212 &self,
5213 e: &Engine,
5214 tokens: &[u32],
5215 pos0: usize,
5216 cache: &mut Cache,
5217 embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
5218 ckpt: Option<&mut VerifyCkpt>,
5219 stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
5220 fence: &[usize],
5221 ) -> Result<VerifyBoundaryTicket, Box<dyn std::error::Error>> {
5222 assert!(
5223 !self.is_gemma4_e4b() && !self.gemma_batch_program(),
5224 "decode_step_t_core_ppn covers the hybrid non-gemma4 verify trunk only \
5225 (the gemma4 arms have their own decode_step_t twins)"
5226 );
5227 if crate::pp::pp_host_bounce_active() && (stream.is_some() || embd_dev.is_some()) {
5228 return Err(
5229 "decode_step_t_core_ppn: refused with MEMRA_PP_HOST_BOUNCE=1 — the trunk \
5230 boundary itself is host-staged, but device-resident verify still peer-reads \
5231 primary-device token/position/embedding buffers from stage 0. Run plain PP \
5232 serving on this host class; spec requires local per-stage inputs first."
5233 .into(),
5234 );
5235 }
5236 let rt = crate::pp::PpNRt::get(e)?;
5237 // Pipelined callers do not bypass ownership: their explicit coordinator borrow makes
5238 // this acquire clone the same active generation. Ordinary callers acquire a fresh lease.
5239 let walk_owner = rt.acquire_walk("verify_stage0_issue")?;
5240 let n_st = fence.len() - 1;
5241 assert_eq!(
5242 rt.n_stages(),
5243 n_st,
5244 "PpNRt stage count {} != fence stages {n_st}",
5245 rt.n_stages()
5246 );
5247 let n_embd = self.cfg.n_embd as usize;
5248 let t = tokens.len();
5249 let payload = t * n_embd;
5250 // One-shot lane diagnostic: force natural PP-2 boundaries to completion so the server
5251 // log can price stage 0, the peer hop, the RX copy, and stage 1 + head separately without
5252 // nsys. The ordinary path keeps every enqueue asynchronous. N>2 is deliberately excluded:
5253 // the report below names exactly two stages and must never imply it measured middle ones.
5254 let pp_anatomy = n_st == 2 && std::env::var("MEMRA_SPEC_PP_ANATOMY").as_deref() == Ok("1");
5255 let pp_started = std::time::Instant::now();
5256 let (mut reverse_ms, mut stage0_ms, mut tx_ms) = (0.0f64, 0.0f64, 0.0f64);
5257 // The CALLER's ambient stream, captured BEFORE any `rt.enter()` pushes a stage stream:
5258 // this body returns DEVICE-RESIDENT buffers (the device-argmax accept walk's contract),
5259 // so the exit needs the same publication the boundaries get — see `PpNRt::publish_to`.
5260 // Taken here, not at the end, because inside the last-stage scope `e.stream()` IS the
5261 // stage stream and the wait would self-order into a no-op.
5262 let caller_stream = e.stream();
5263 // #87 ROOT-CAUSE FENCE (lane/pp2spec-crash): the PREVIOUS round's stage-allocated
5264 // outputs (logits/hidden/ckpt stashes) freed stream-ordered on the STAGE streams while
5265 // the primary stream still holds queued reads of them — with event tracking elided,
5266 // nothing stops the pool from reusing those blocks for THIS round's stage allocations,
5267 // whose writes then race the queued reads (measured: 13/4096-NaN random-bits garbage in
5268 // the spec round seed; the full anatomy is on `PpNRt::fence_stages_behind`). Order every
5269 // stage stream behind the caller before enqueueing new stage work.
5270 let reverse_started = std::time::Instant::now();
5271 rt.fence_stages_behind(&caller_stream)?;
5272 if pp_anatomy {
5273 // Drain the reverse-publication dependency before timing stage 0 itself. At c=1 this
5274 // prices any primary-stream rollback/refresh tail inherited from the prior round.
5275 for s in 0..n_st {
5276 let _st = rt.enter(s);
5277 rt.engine(s, e).stream().synchronize()?;
5278 }
5279 reverse_ms = reverse_started.elapsed().as_secs_f64() * 1e3;
5280 }
5281
5282 // Per-stage rope positions: in host mode the same [T] iota each stage uploads itself; in
5283 // stream mode each stage's own `pos_iota` over the shared read-only device counter.
5284 let stage_pos = |es: &Engine| -> Result<CudaSlice<i32>, Box<dyn std::error::Error>> {
5285 match stream {
5286 Some((_, ctr)) => {
5287 let mut p = es.alloc_uninit::<i32>(t)?;
5288 es.pos_iota(ctr, &mut p, t)?;
5289 Ok(p)
5290 }
5291 None => {
5292 let pos_vec: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
5293 es.htod_i32(&pos_vec)
5294 }
5295 }
5296 };
5297
5298 // ---- STAGE 0: embed (the table lives with stage 0) + layers [0, fence[1]) + TX ----
5299 let slot = {
5300 let _st0 = rt.enter(0);
5301 let e0 = rt.engine(0, e);
5302 let stage0_started = std::time::Instant::now();
5303 let pos_d = stage_pos(e0)?;
5304 let x = match (stream, embd_dev) {
5305 (Some((vtok, _)), Some((g, qt, rb))) => {
5306 e0.embed_gather_device_td(g, vtok, t, n_embd, qt, rb)?
5307 }
5308 (None, Some((g, qt, rb))) => e0.embed_gather_device_t(g, tokens, n_embd, qt, rb)?,
5309 _ => e0.htod(&self.embd.try_gather(n_embd, tokens)?)?,
5310 };
5311 let x = self.verify_layers(
5312 e0, x, fence[0], fence[1], &pos_d, pos0, t, cache, ckpt, stream, None,
5313 )?;
5314 if pp_anatomy {
5315 e0.stream().synchronize()?;
5316 stage0_ms = stage0_started.elapsed().as_secs_f64() * 1e3;
5317 }
5318 let tx_started = std::time::Instant::now();
5319 let slot = rt.tx(0, &x, payload)?;
5320 if pp_anatomy {
5321 e0.stream().synchronize()?;
5322 tx_ms = tx_started.elapsed().as_secs_f64() * 1e3;
5323 }
5324 slot
5325 // x + pos_d drop here: freed stream-ordered on stage-0's stream after use.
5326 };
5327
5328 Ok(VerifyBoundaryTicket {
5329 rt,
5330 caller_stream,
5331 slot,
5332 pos0,
5333 t,
5334 payload,
5335 n_st,
5336 pp_anatomy,
5337 pp_started,
5338 reverse_ms,
5339 stage0_ms,
5340 tx_ms,
5341 _walk_owner: walk_owner,
5342 })
5343 }
5344
5345 /// Consume a stage-0 boundary ticket and enqueue the remaining PP stages plus the head.
5346 /// On PP-2 this is exactly stage 1; PP-N keeps its pre-existing middle-stage walk here.
5347 #[allow(clippy::too_many_arguments)]
5348 fn verify_stage1_finish(
5349 &self,
5350 e: &Engine,
5351 ticket: VerifyBoundaryTicket,
5352 cache: &mut Cache,
5353 mut ckpt: Option<&mut VerifyCkpt>,
5354 stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
5355 fence: &[usize],
5356 publish_to_caller: bool,
5357 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
5358 let VerifyBoundaryTicket {
5359 rt,
5360 caller_stream,
5361 slot,
5362 pos0,
5363 t,
5364 payload,
5365 n_st,
5366 pp_anatomy,
5367 pp_started,
5368 reverse_ms,
5369 stage0_ms,
5370 tx_ms,
5371 _walk_owner,
5372 } = ticket;
5373 let n_embd = self.cfg.n_embd as usize;
5374 let eps = self.cfg.rms_eps;
5375 let mut slot = slot;
5376 let (mut rx_ms, mut stage1_ms) = (0.0f64, 0.0f64);
5377 let stage_pos = |es: &Engine| -> Result<CudaSlice<i32>, Box<dyn std::error::Error>> {
5378 match stream {
5379 Some((_, ctr)) => {
5380 let mut p = es.alloc_uninit::<i32>(t)?;
5381 es.pos_iota(ctr, &mut p, t)?;
5382 Ok(p)
5383 }
5384 None => {
5385 let pos_vec: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
5386 es.htod_i32(&pos_vec)
5387 }
5388 }
5389 };
5390
5391 // ---- MIDDLE STAGES: RX boundary s-1 -> range -> TX boundary s ----
5392 for s in 1..n_st - 1 {
5393 let _st = rt.enter(s);
5394 let es = rt.engine(s, e);
5395 let pos_d = stage_pos(es)?;
5396 let x = rt.rx(s - 1, slot, payload)?;
5397 let x = self.verify_layers(
5398 es,
5399 x,
5400 fence[s],
5401 fence[s + 1],
5402 &pos_d,
5403 pos0,
5404 t,
5405 cache,
5406 ckpt.as_deref_mut(),
5407 stream,
5408 None,
5409 )?;
5410 slot = rt.tx(s, &x, payload)?;
5411 }
5412
5413 // ---- LAST STAGE: RX + final range + output_norm + lm head ----
5414 let _stl = rt.enter(n_st - 1);
5415 let el = rt.engine(n_st - 1, e);
5416 let pos_d = stage_pos(el)?;
5417 let rx_started = std::time::Instant::now();
5418 let x = rt.rx(n_st - 2, slot, payload)?;
5419 if pp_anatomy {
5420 el.stream().synchronize()?;
5421 rx_ms = rx_started.elapsed().as_secs_f64() * 1e3;
5422 }
5423 let stage1_started = std::time::Instant::now();
5424 let x = self.verify_layers(
5425 el,
5426 x,
5427 fence[n_st - 1],
5428 fence[n_st],
5429 &pos_d,
5430 pos0,
5431 t,
5432 cache,
5433 ckpt,
5434 stream,
5435 None,
5436 )?;
5437
5438 let mut hn = vbuf(el, payload)?;
5439 let logits = if self.sliding_gated_moe_batch_program() {
5440 // The PP Step35 serving path uses rms_norm + matmul for B=1 as well as B>1.
5441 // Verify must not switch numeric class merely because the same session speculates.
5442 el.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
5443 el.matmul(&self.output, &hn, t)?
5444 } else {
5445 el.rms_norm_decode(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
5446 el.matmul_decode_exact(&self.output, &hn, t)?
5447 };
5448 if pp_anatomy {
5449 el.stream().synchronize()?;
5450 stage1_ms = stage1_started.elapsed().as_secs_f64() * 1e3;
5451 }
5452 // EXIT PUBLICATION: both returned buffers are still being produced on the last stage's
5453 // stream. Order the caller's stream behind that work before the buffers escape this
5454 // scope (the 2026-08-06 same-device ppspec find: without it the caller's primary-stream
5455 // consumer read unwritten logits — nondeterministic, one-device-only, and it poisoned
5456 // the following arm's KV in the same process).
5457 if publish_to_caller {
5458 rt.publish_to(n_st - 1, &caller_stream)?;
5459 }
5460 if pp_anatomy {
5461 if publish_to_caller {
5462 caller_stream.synchronize()?;
5463 }
5464 eprintln!(
5465 "[spec-pp-anatomy] t={t} reverse={reverse_ms:.3}ms stage0={stage0_ms:.3}ms \
5466 tx={tx_ms:.3}ms rx={rx_ms:.3}ms stage1-head={stage1_ms:.3}ms total={:.3}ms",
5467 pp_started.elapsed().as_secs_f64() * 1e3,
5468 );
5469 }
5470 // stream: the device pos counter owns position; host mirror reconciles at drain.
5471 if stream.is_none() {
5472 cache.pos += t;
5473 }
5474 Ok((logits, if spec_hpost() { hn } else { x }))
5475 }
5476
5477 /// Step3.5/Step3.7 verify trunk in the serving batched numeric class.
5478 ///
5479 /// `step35_decode_batch_layers` is now authoritative at every live serving width, including
5480 /// B=1 (lane/cx-b1fix). The older verify walk deliberately mirrored the eager T=1 class:
5481 /// it replayed `step35_decode_attn` per row and used the eager/decode-exact FFN dispatch.
5482 /// Those classes are individually stable, but a near-tie prompt can choose different greedy
5483 /// bytes when a request moves from batched plain serving into speculative verify. Run the
5484 /// same authoritative B=1 stage subgraph for each verify row here. Rows still advance
5485 /// layer-by-layer, so every layer sees the preceding verify rows in its attention cache while
5486 /// every norm/projection/FFN uses exactly the live serving dispatch.
5487 #[allow(clippy::too_many_arguments)]
5488 /// PRIME-BY-T-ROWS (MEMRA_PRIME_TROWS=1): prefill the prompt through the same-session
5489 /// t-row walk in 32-row chunks — every row runs the t=1 decode program bit-for-bit
5490 /// (the TOKENWISE-prime ORACLE class), so this door is exact against the exactness
5491 /// reference while replacing the host-canonical per-token prime. Requires the walk
5492 /// doors (MEMRA_SPEC_VERIFY_EAGER/TCOL); returns the prime contract trio.
5493 #[allow(clippy::type_complexity)]
5494 pub(crate) fn step35_prime_trows(
5495 &self,
5496 e: &Engine,
5497 tokens: &[u32],
5498 cache: &mut Cache,
5499 ) -> Result<Option<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
5500 {
5501 let dbg = std::env::var("MEMRA_SPEC_FA2_DEBUG").as_deref() == Ok("1");
5502 if !prime_trows_on() {
5503 return Ok(None);
5504 }
5505 if !self.uses_sliding_gated_moe_program()
5506 || cache.pos != 0
5507 || cache.dflash_taps.is_some()
5508 || !spec_verify_eager_on()
5509 || !spec_verify_tcol_on()
5510 {
5511 if dbg {
5512 eprintln!(
5513 "[prime-trows] refuse: program={} pos={} taps={} eager={:?} tcol={:?}",
5514 self.uses_sliding_gated_moe_program(),
5515 cache.pos,
5516 cache.dflash_taps.is_some(),
5517 std::env::var("MEMRA_SPEC_VERIFY_EAGER").ok(),
5518 std::env::var("MEMRA_SPEC_VERIFY_TCOL").ok()
5519 );
5520 }
5521 return Ok(None);
5522 }
5523 let n_embd = self.cfg.n_embd as usize;
5524 let n_layers = self.layers.len();
5525 let t_total = tokens.len();
5526 let Some(embd_gpu) = self.embd_gpu_try(e) else {
5527 if dbg {
5528 eprintln!("[prime-trows] refuse: no device embed table");
5529 }
5530 return Ok(None);
5531 };
5532 let embd_qtype = match self.embd.ggml_type {
5533 memra_gguf::GgmlType::BF16 => crate::QT_BF16,
5534 memra_gguf::GgmlType::Q8_0 => crate::QT_Q8_0,
5535 other => {
5536 if dbg {
5537 eprintln!("[prime-trows] refuse: embed dtype {other:?}");
5538 }
5539 return Ok(None);
5540 }
5541 };
5542 let embd_row_bytes = self.embd.raw.len() / self.cfg.n_vocab as usize;
5543 // Chunk plan: 32-row chunks; a 1-token tail folds into the previous chunk
5544 // (the walk floor is t >= 2).
5545 let mut bounds = Vec::new();
5546 let mut start = 0usize;
5547 while start < t_total {
5548 let mut end = (start + 32).min(t_total);
5549 if t_total - end == 1 {
5550 end -= 1;
5551 }
5552 bounds.push((start, end));
5553 start = end;
5554 }
5555 if bounds.iter().any(|(a, b)| b - a < 2) {
5556 return Ok(None); // degenerate short prompt keeps the ordinary prime
5557 }
5558 let mut hiddens = e.uninit(t_total * n_embd)?;
5559 let mut last: Option<CudaSlice<f32>> = None;
5560 for &(a, b) in &bounds {
5561 let tc = b - a;
5562 let tok_d = e.stream().clone_htod(&tokens[a..b])?;
5563 let x =
5564 e.embed_gather_device_td(embd_gpu, &tok_d, tc, n_embd, embd_qtype, embd_row_bytes)?;
5565 let out = self.step35_verify_batch_layers(e, x, 0, n_layers, a, tc, cache)?;
5566 e.copy_into(&mut hiddens, a * n_embd, &out, tc * n_embd)?;
5567 if b == t_total {
5568 let mut h = e.uninit(n_embd)?;
5569 e.dtod_copy_view(&out.slice((tc - 1) * n_embd..tc * n_embd), &mut h)?;
5570 last = Some(h);
5571 }
5572 }
5573 let h_seed = last.expect("last chunk produced the seed row");
5574 let mut hn = e.uninit(n_embd)?;
5575 e.rms_norm_decode(
5576 &h_seed,
5577 self.output_norm.float_data(),
5578 &mut hn,
5579 n_embd,
5580 1,
5581 self.cfg.rms_eps,
5582 )?;
5583 let logits_d = e.matmul_decode_exact(&self.output, &hn, 1)?;
5584 let logits = e.dtoh(&logits_d)?;
5585 cache.pos = t_total;
5586 Ok(Some((logits, h_seed, hiddens)))
5587 }
5588
5589 #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
5590 fn step35_verify_batch_layers(
5591 &self,
5592 e: &Engine,
5593 mut x: CudaSlice<f32>,
5594 lo: usize,
5595 hi: usize,
5596 pos0: usize,
5597 t: usize,
5598 cache: &mut Cache,
5599 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5600 let n_embd = self.cfg.n_embd as usize;
5601 if !self.uses_sliding_gated_moe_program() {
5602 return Err(
5603 "serving-class verify requires sliding-gated-MoE canonical operations".into(),
5604 );
5605 }
5606 // SERVING-CLASS VERIFY (MEMRA_SPEC_VERIFY_EAGER=1, step37 MTP bring-up): each verify
5607 // column rides decode_layers_eager — the EXACT t=1 program live serving runs (all TP2
5608 // doors) — row-outer, so row r's appends land before row r+1 attends: bit-equal to
5609 // plain greedy by construction. Only the unsplit full-range walk qualifies; PP splits
5610 // and the tap path keep the batch-layer class.
5611 static VE: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
5612 let eager_verify =
5613 *VE.get_or_init(spec_verify_eager_on) && lo == 0 && hi == self.layers.len();
5614 if eager_verify {
5615 // T-COLUMN LAYER-OUTER WALK (MEMRA_SPEC_VERIFY_TCOL=1): per layer, one t-grid
5616 // attn norm + ONE weight-amortized QKV(+gate) over all T columns, then each
5617 // column runs the UNMODIFIED t=1 attention program via the col-select door and
5618 // the ordinary residual/FFN body. Values per column are bit-equal to the
5619 // row-outer walk: rms over the materialized residual == the fused add+norm
5620 // (kernel_check identity), the tcol kernel's per-column FP order == the t=1
5621 // kernel, and every downstream op IS the t=1 program.
5622 static TCOL: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
5623 let tcol = *TCOL.get_or_init(spec_verify_tcol_on);
5624 // T > 32 (prefill-class): run the SAME walk in 32-row chunks — each chunk's
5625 // rows are the t=1 program bit-for-bit and the rope pass advances the cache,
5626 // so a chunked call is value-identical to the row-outer loop it replaces.
5627 static TROWS_PREFILL: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
5628 // MEMRA_STEP_GEMM_PRIME outranks the walk: with the grouped GEMM prime armed, the
5629 // t-row walk defers so the batch path (GEMM trunk + grouped MoE) takes the prompt —
5630 // flag precedence between two existing doors, not a new flag. Without this, both
5631 // doors ON meant the walk still won and the GEMM prime needed PRIME_TROWS=0 by hand.
5632 let trows_prefill =
5633 *TROWS_PREFILL.get_or_init(|| prime_trows_on() && !crate::step_gemm_prime_on());
5634 // MEMRA_PRIME_TROWS_T=<w>: chunk width, default 8 = the REAL cap of this walk.
5635 // The workspace slabs go to 32 rows, but `matvec_bf16_qkvg_tcol_into` refuses
5636 // t > 8 (compile-time-T twins exist for 2/4/8 only; the runtime-t kernel spills
5637 // its accumulators to local memory), so a wider chunk fails the request with
5638 // "matvec_bf16_qkvg_tcol geometry" — which is exactly how the first server-path
5639 // TROWS arm died. Measured at 193 tokens: w=8 2.459 s, w=4 2.574 s.
5640 static TROWS_W: std::sync::OnceLock<Result<usize, String>> = std::sync::OnceLock::new();
5641 let trows_w = match TROWS_W.get_or_init(|| {
5642 let value = std::env::var("MEMRA_PRIME_TROWS_T").ok();
5643 parse_prime_trows_width(value.as_deref())
5644 }) {
5645 Ok(width) => *width,
5646 Err(err) => return Err(err.clone().into()),
5647 };
5648 if tcol && trows_prefill && t > trows_w {
5649 // One-time engagement receipt: without it a prefill gate cannot tell a
5650 // chunked walk from the row-outer fallback it is supposed to replace
5651 // (the first PRIME_TROWS gate passed vacuously on exactly that).
5652 static SEEN: std::sync::atomic::AtomicBool =
5653 std::sync::atomic::AtomicBool::new(false);
5654 if !SEEN.swap(true, std::sync::atomic::Ordering::Relaxed) {
5655 eprintln!(
5656 "[prime-trows] ENGAGED t={t} width={trows_w} chunks={} layers={}..{}",
5657 t.div_ceil(trows_w),
5658 lo,
5659 hi
5660 );
5661 }
5662 let mut out = e.uninit(t * n_embd)?;
5663 let mut start = 0usize;
5664 while start < t {
5665 let mut end = (start + trows_w).min(t);
5666 if t - end == 1 {
5667 end -= 1;
5668 }
5669 let tc = end - start;
5670 let mut xc = e.uninit(tc * n_embd)?;
5671 e.dtod_copy_view(&x.slice(start * n_embd..end * n_embd), &mut xc)?;
5672 let oc =
5673 self.step35_verify_batch_layers(e, xc, lo, hi, pos0 + start, tc, cache)?;
5674 e.copy_into(&mut out, start * n_embd, &oc, tc * n_embd)?;
5675 start = end;
5676 }
5677 return Ok(out);
5678 }
5679 if tcol && (2..=32).contains(&t) {
5680 // MEMRA_TCOL_PROF=1: synchronized per-segment wall profile of the walk
5681 // (norm+QKV precompute / per-col attention / per-col residual+FFN). The
5682 // syncs serialize the stream, so the split is for TARGETING amortization
5683 // work only — never a perf claim.
5684 static PROF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
5685 let prof =
5686 *PROF.get_or_init(|| std::env::var("MEMRA_TCOL_PROF").as_deref() == Ok("1"));
5687 let mut prof_ms = [0f64; 3];
5688 let eps = self.cfg.rms_eps;
5689 let mut x_t = x;
5690 let mut h_t = e.uninit(t * n_embd)?;
5691 let mut h_row = e.uninit(n_embd)?; // real row: the non-dcw fallback reads it
5692 // Per-column pos buffers hoisted out of the layer loop (a per-col-per-layer
5693 // pageable htod was an in-stream engine turnaround x t x 45).
5694 let mut pos_rows = Vec::with_capacity(t);
5695 for r in 0..t {
5696 pos_rows.push(e.htod_i32(&[(pos0 + r) as i32])?);
5697 }
5698 let mut ok = true;
5699 // MEMRA_TCOL_OPROJ=1: defer each column's o_proj — the finish seam
5700 // stashes `gated` instead of joining per column; one b4_tcol per rank +
5701 // one slab join produce every column's `mixed` after the attention pass.
5702 // Bit-exact per column (t=1 b4 program per column; elementwise join).
5703 // MEMRA_TCOL_FFN=1: today this only IMPLIES the o_proj defer above. Its
5704 // named feature, the two-column device-routed FFN sweep, rode the
5705 // slot-major v2 TP banks and was REMOVED with the MEMRA_NVFP4_BANK_V2 door
5706 // (2026-08-29, research/step37-bankv2-removal-20260829): the v2 layout
5707 // changed generated text in serving. The flag itself stays because it is
5708 // family-armed in the step37 serving defaults and killing it here would
5709 // silently drop the o_proj defer from the qualified serving shape.
5710 static FFN2: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
5711 let ffn_batch = *FFN2.get_or_init(tcol_ffn_on);
5712 let oproj_batch = crate::tp::tcol_oproj_on() || ffn_batch;
5713 // MEMRA_SPEC_FA2=1: eligible layers defer every verify column's fa —
5714 // the per-column pass norms/ropes/appends and stashes q+gate, then one
5715 // fa_decode_vec_q_v3_dcw_rows per rank attends every stashed row and the
5716 // o_proj join produces the mixed slab. Ineligible/boundary layers run the
5717 // ordinary program.
5718 let fa2 = crate::tp::spec_fa2_on() && t <= 32;
5719 let mut mixed_row = e.uninit(n_embd)?;
5720 let mut pos_staged = false;
5721 for il in lo..hi {
5722 let layer = &self.layers[il];
5723 // BEFORE this layer touches its planes: is the history it is about to
5724 // attend already poisoned? Global (non-ring) layers only, which are the
5725 // ones the level-2 bitmap implicates.
5726 if kv_plane_scan_on()
5727 && self.step35_geom(il).window.is_none()
5728 && let Some(distributed) = cache.tp_kv[il].as_ref()
5729 {
5730 scan_kv_plane(e, distributed, il, pos0)?;
5731 }
5732 let fa2_layer = fa2 && self.step35_fa_rows_precheck(cache, il, pos0, t)?;
5733 let mut seg = std::time::Instant::now();
5734 e.rms_norm(&x_t, layer.attn_norm.float_data(), &mut h_t, n_embd, t, eps)?;
5735 if !self.step35_verify_qkv_precompute(e, il, &h_t, t)? {
5736 ok = false;
5737 break;
5738 }
5739 // FULL t-row attention pass (rope/append + fa + combine + o_proj in
5740 // 3 launches/rank): same-session rows, slot = len-base+r, one len
5741 // advance by t. Ring rebase happens during append BEFORE the fused
5742 // kernel so device base_d and memory are already rebased for rows.
5743 // Host cache bookkeeping mirrors the per-column tail.
5744 let mut mixed_t_opt: Option<CudaSlice<f32>> = None;
5745 if fa2_layer {
5746 let crate::hybrid::Mixer::Full(fa) = &layer.mixer else {
5747 return Err("verify rope pass expects full attention".into());
5748 };
5749 let tp = fa
5750 .step_tp_qkv
5751 .as_ref()
5752 .ok_or("verify rope pass lost its TP state")?;
5753 let empty: [CudaSlice<f32>; 0] = [];
5754 let transaction = {
5755 let tp_kv = cache.tp_kv[il]
5756 .as_mut()
5757 .expect("precheck verified the distributed cache");
5758 let tx = tp_kv.begin_transaction()?;
5759 tp.runtime.append_tp_kv_transaction_inner(
5760 tp_kv, tx, &empty, &empty, t, true,
5761 )?;
5762 tx
5763 };
5764 match self.step35_verify_rope_fa_pass(e, il, cache, pos0, t, !pos_staged) {
5765 Ok(Some(mixed_t)) => {
5766 let tp_kv = cache.tp_kv[il]
5767 .as_mut()
5768 .expect("precheck verified the distributed cache");
5769 tp.runtime.commit_tp_kv_transaction_external(
5770 tp_kv,
5771 transaction,
5772 t,
5773 )?;
5774 if let Some(local) = cache.kv[il].as_mut() {
5775 local.len = pos0 + t;
5776 if !crate::tp::len_mirror_lazy_on() {
5777 e.set_i32_one(&mut local.len_d, local.len as i32)?;
5778 }
5779 if let (Some(ring), Some(tp_base)) =
5780 (local.ring.as_mut(), tp_kv.ring_base())
5781 && ring.base() != tp_base
5782 {
5783 ring.apply_rebase(tp_base);
5784 if let Some(base_d) = local.base_d.as_mut() {
5785 e.set_i32_one(base_d, tp_base as i32)?;
5786 }
5787 }
5788 }
5789 pos_staged = true;
5790 mixed_t_opt = Some(mixed_t);
5791 }
5792 Ok(None) => {
5793 let tp_kv = cache.tp_kv[il]
5794 .as_mut()
5795 .expect("precheck verified the distributed cache");
5796 tp.runtime.rollback_tp_kv_transaction(tp_kv, transaction)?;
5797 }
5798 Err(err) => {
5799 if let Some(tp_kv) = cache.tp_kv[il].as_mut() {
5800 let _ =
5801 tp.runtime.rollback_tp_kv_transaction(tp_kv, transaction);
5802 }
5803 return Err(err);
5804 }
5805 }
5806 }
5807 if let Some(mixed_t) = mixed_t_opt {
5808 if prof {
5809 e.stream().synchronize()?;
5810 prof_ms[1] += seg.elapsed().as_secs_f64() * 1e3;
5811 seg = std::time::Instant::now();
5812 }
5813 let o_out = mixed_t.len() / t;
5814 let mut next = e.uninit(t * n_embd)?;
5815 {
5816 for r in 0..t {
5817 e.dtod_copy_view(
5818 &mixed_t.slice(r * o_out..(r + 1) * o_out),
5819 &mut mixed_row,
5820 )?;
5821 let mut x_row = e.uninit(n_embd)?;
5822 e.dtod_copy_view(
5823 &x_t.slice(r * n_embd..(r + 1) * n_embd),
5824 &mut x_row,
5825 )?;
5826 let (x1, ffn_out) = self.residual_norm_ffn(
5827 e, layer, &x_row, &mixed_row, n_embd, il, eps,
5828 )?;
5829 let mut x2 = e.uninit(n_embd)?;
5830 e.add(&x1, &ffn_out, &mut x2, n_embd)?;
5831 e.dtod_copy_into(&x2, &mut next, r * n_embd)?;
5832 }
5833 }
5834 if prof {
5835 e.stream().synchronize()?;
5836 prof_ms[2] += seg.elapsed().as_secs_f64() * 1e3;
5837 }
5838 x_t = next;
5839 if spec_nan_scan() {
5840 // The scan MUST sit on this arm too. It used to live only on
5841 // the non-fused tail, so a fused layer's poison was first
5842 // reported by the next non-fused layer.
5843 verify_arm_receipt(
5844 "fused",
5845 il,
5846 pos0,
5847 t,
5848 cache.tp_kv[il].as_ref().map(|d| d.staged_len()),
5849 );
5850 nan_scan_rows(
5851 e,
5852 &x_t,
5853 t,
5854 n_embd,
5855 &format!("tcol layer {il} pos0={pos0} arm=fused"),
5856 )?;
5857 }
5858 continue;
5859 }
5860 if prof {
5861 e.stream().synchronize()?;
5862 prof_ms[0] += seg.elapsed().as_secs_f64() * 1e3;
5863 seg = std::time::Instant::now();
5864 }
5865 let mut next = e.uninit(t * n_embd)?;
5866 // Columns whose o_proj was deferred (their FFN runs after the join).
5867 // A NON-deferred column's FFN must run INSIDE the column loop: the
5868 // oproj-tail handoff is a single cell that the same column's
5869 // residual_norm_ffn consumes before the next column's finish.
5870 let mut deferred: Vec<usize> = Vec::new();
5871 let mut fa2_deferred: Vec<usize> = Vec::new();
5872 let ffn_col = |r: usize,
5873 mixed: &CudaSlice<f32>,
5874 next: &mut CudaSlice<f32>|
5875 -> Result<(), Box<dyn std::error::Error>> {
5876 let mut x_row = e.uninit(n_embd)?;
5877 e.dtod_copy_view(&x_t.slice(r * n_embd..(r + 1) * n_embd), &mut x_row)?;
5878 let (x1, ffn_out) =
5879 self.residual_norm_ffn(e, layer, &x_row, mixed, n_embd, il, eps)?;
5880 if spec_nan_scan_level() >= 2 {
5881 nan_scan_rows(
5882 e,
5883 &ffn_out,
5884 1,
5885 n_embd,
5886 &format!("tcol layer {il} col {r} per-column FFN out"),
5887 )?;
5888 }
5889 let mut x2 = e.uninit(n_embd)?;
5890 e.add(&x1, &ffn_out, &mut x2, n_embd)?;
5891 e.dtod_copy_into(&x2, next, r * n_embd)?;
5892 Ok(())
5893 };
5894 #[allow(clippy::needless_range_loop)]
5895 // allow: the explicit index loop keeps the offset arithmetic visible and aligned with the device-side indexing
5896 for r in 0..t {
5897 e.dtod_copy_view(&h_t.slice(r * n_embd..(r + 1) * n_embd), &mut h_row)?;
5898 let row_pos = &pos_rows[r];
5899 crate::tp::set_verify_tcol(Some(r));
5900 if fa2_layer {
5901 crate::tp::set_spec_fa2_defer(Some(r));
5902 } else if oproj_batch {
5903 crate::tp::set_tcol_oproj_defer(Some(r));
5904 }
5905 let mixed = match &layer.mixer {
5906 crate::hybrid::Mixer::Full(fa) => {
5907 self.full_attn_decode(e, fa, &h_row, row_pos, pos0 + r, cache, il)
5908 }
5909 _ => Err("step35 verify expects full attention".into()),
5910 };
5911 crate::tp::set_verify_tcol(None);
5912 crate::tp::set_spec_fa2_defer(None);
5913 crate::tp::set_tcol_oproj_defer(None);
5914 let mixed = mixed?;
5915 if fa2_layer && crate::tp::take_spec_fa2_stashed() {
5916 fa2_deferred.push(r);
5917 } else if oproj_batch && crate::tp::take_tcol_oproj_stashed() {
5918 deferred.push(r);
5919 } else {
5920 if spec_nan_scan_level() >= 2 {
5921 let cols = mixed.len();
5922 nan_scan_rows(
5923 e,
5924 &mixed,
5925 1,
5926 cols,
5927 &format!("tcol layer {il} col {r} per-column ATTN out"),
5928 )?;
5929 }
5930 ffn_col(r, &mixed, &mut next)?;
5931 }
5932 }
5933 if !fa2_deferred.is_empty() && fa2_deferred.len() != t {
5934 // The precheck guarantees both columns stash or neither; a strict
5935 // subset means a column's output was never produced anywhere.
5936 return Err("spec fa2 stash engaged for a subset of columns".into());
5937 }
5938 if prof {
5939 e.stream().synchronize()?;
5940 prof_ms[1] += seg.elapsed().as_secs_f64() * 1e3;
5941 seg = std::time::Instant::now();
5942 }
5943 if !fa2_deferred.is_empty() {
5944 deferred = fa2_deferred;
5945 }
5946 if !deferred.is_empty() {
5947 let mixed_t = if fa2_layer {
5948 self.step35_verify_fa_rows_join(e, il, cache, pos0, t)?
5949 } else {
5950 self.step35_verify_oproj_tcol(e, il, t)?
5951 };
5952 let o_out = mixed_t.len() / t;
5953 if spec_nan_scan_level() >= 2 {
5954 nan_scan_rows(
5955 e,
5956 &mixed_t,
5957 t,
5958 o_out,
5959 &format!("tcol layer {il} JOINED attn over deferred cols"),
5960 )?;
5961 }
5962 // Batched t=2 residual+MoE: one t-grid add_rms_norm (per-row
5963 // program == t=1; bit-identical to the oproj-tail join per the
5964 // M2 verbatim-program contract) feeding the two-column routed
5965 // sweep. Ineligible layers (dense FFN, non-nvfp4) fall through
5966 // to the per-column body.
5967 {
5968 for &r in &deferred {
5969 e.dtod_copy_view(
5970 &mixed_t.slice(r * o_out..(r + 1) * o_out),
5971 &mut mixed_row,
5972 )?;
5973 ffn_col(r, &mixed_row, &mut next)?;
5974 }
5975 }
5976 }
5977 if prof {
5978 e.stream().synchronize()?;
5979 prof_ms[2] += seg.elapsed().as_secs_f64() * 1e3;
5980 }
5981 x_t = next;
5982 if spec_nan_scan() {
5983 verify_arm_receipt(
5984 if fa2_layer { "join" } else { "percol" },
5985 il,
5986 pos0,
5987 t,
5988 cache.tp_kv[il].as_ref().map(|d| d.staged_len()),
5989 );
5990 nan_scan_rows(
5991 e,
5992 &x_t,
5993 t,
5994 n_embd,
5995 &format!(
5996 "tcol layer {il} pos0={pos0} arm={}",
5997 if fa2_layer { "join" } else { "percol" }
5998 ),
5999 )?;
6000 }
6001 }
6002 if prof {
6003 eprintln!(
6004 "[tcol-prof] t={t} norm+qkv={:.3}ms attn={:.3}ms ffn={:.3}ms",
6005 prof_ms[0], prof_ms[1], prof_ms[2]
6006 );
6007 }
6008 if ok {
6009 return Ok(x_t);
6010 }
6011 // fall through to the row-outer walk on ineligible layers
6012 x = x_t;
6013 }
6014 let mut next = e.uninit(t * n_embd)?;
6015 let scan = spec_nan_scan();
6016 for r in 0..t {
6017 let mut row = e.uninit(n_embd)?;
6018 e.dtod_copy_view(&x.slice(r * n_embd..(r + 1) * n_embd), &mut row)?;
6019 let row_pos = e.htod_i32(&[(pos0 + r) as i32])?;
6020 let out = if scan {
6021 // Diagnostic arm: the same range walked one layer at a time so the first
6022 // poisoned layer names itself. `decode_layers_eager(lo, hi)` is range-scoped
6023 // and executes its trailing residual add, so a per-layer chain is the same
6024 // program with the cross-layer add+norm fusion unrolled.
6025 nan_scan_rows(
6026 e,
6027 &row,
6028 1,
6029 n_embd,
6030 &format!("embed row r={r} pos={}", pos0 + r),
6031 )?;
6032 let mut acc = row;
6033 for il in lo..hi {
6034 acc = self.decode_layers_eager(
6035 e,
6036 acc,
6037 il,
6038 il + 1,
6039 &row_pos,
6040 pos0 + r,
6041 cache,
6042 )?;
6043 nan_scan_rows(
6044 e,
6045 &acc,
6046 1,
6047 n_embd,
6048 &format!("row-outer layer {il} r={r} pos={}", pos0 + r),
6049 )?;
6050 }
6051 acc
6052 } else {
6053 self.decode_layers_eager(e, row, lo, hi, &row_pos, pos0 + r, cache)?
6054 };
6055 e.dtod_copy_into(&out, &mut next, r * n_embd)?;
6056 }
6057 // dflash taps are NOT produced on this arm (they need per-layer hiddens the
6058 // row-outer walk does not materialize); the door is a step37 MTP bring-up
6059 // surface where taps are unused.
6060 return Ok(next);
6061 }
6062 let mut ph_last = std::time::Instant::now();
6063 for il in lo..hi {
6064 let mut next = e.uninit(t * n_embd)?;
6065 for r in 0..t {
6066 let mut row = e.uninit(n_embd)?;
6067 e.dtod_copy_view(&x.slice(r * n_embd..(r + 1) * n_embd), &mut row)?;
6068 // The caller owns this verify's position. During controller overlap, cache.pos
6069 // still describes generation N while this stage-0 walk belongs to N+1.
6070 let row_pos = e.htod_i32(&[(pos0 + r) as i32])?;
6071 let mut one = [&mut *cache];
6072 let out = self.step35_decode_batch_layers(
6073 e,
6074 row,
6075 &mut one,
6076 &[(pos0 + r) as i32],
6077 &row_pos,
6078 il,
6079 il + 1,
6080 &mut ph_last,
6081 )?;
6082 e.dtod_copy_into(&out, &mut next, r * n_embd)?;
6083 }
6084 self.dflash_tap(e, cache, il, &next, t)?;
6085 x = next;
6086 if spec_nan_scan() {
6087 nan_scan_rows(e, &x, t, n_embd, &format!("batch-layer {il} pos0={pos0}"))?;
6088 }
6089 }
6090 Ok(x)
6091 }
6092
6093 /// DSpark drafter verify (lane/dspark-q38-recover): one t-row forward through the
6094 /// SERVING-CLASS verify funnel (`decode_step_t_core_stream` — the same numeric class
6095 /// MTP verify rides, GDN state advanced in place), returning per-row argmax tokens.
6096 /// Advances `cache.pos += t`; the caller owns snapshot/rollback (block acceptance is
6097 /// prefix-keep, not all-or-nothing).
6098 pub(crate) fn dspark_verify_t_am(
6099 &self,
6100 e: &Engine,
6101 tokens: &[u32],
6102 pos0: usize,
6103 cache: &mut Cache,
6104 ) -> Result<Vec<u32>, Box<dyn std::error::Error>> {
6105 let (logits, _hn) =
6106 self.decode_step_t_core_stream(e, tokens, pos0, cache, None, None, None, None, None)?;
6107 let t = tokens.len();
6108 let v = self.output.out_features();
6109 let mut am_d = e.stream().alloc_zeros::<u32>(t)?;
6110 for r in 0..t {
6111 e.argmax_token_device_col(&logits, r, v, &mut am_d, r)?;
6112 }
6113 e.dtoh_u32(&am_d)
6114 }
6115
6116 /// DSpark verify returning the RAW verify logits [t, n_vocab] (device-resident) instead
6117 /// of per-row argmaxes — the sampled-admission arm's input (rejection-sampling accept
6118 /// gathers filtered p from these columns; lane/dspark-sampled-admission-20260820). Same
6119 /// forward as `dspark_verify_t_am`; the greedy arm keeps its argmax wrapper untouched.
6120 pub(crate) fn dspark_verify_t_logits(
6121 &self,
6122 e: &Engine,
6123 tokens: &[u32],
6124 pos0: usize,
6125 cache: &mut Cache,
6126 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6127 let (logits, _hn) =
6128 self.decode_step_t_core_stream(e, tokens, pos0, cache, None, None, None, None, None)?;
6129 Ok(logits)
6130 }
6131
6132 /// DSpark verify with the MTP column-stash armed: identical forward to
6133 /// `dspark_verify_t_am`, but fills a `VerifyCkpt` so a partial accept can restore
6134 /// column state directly (`dspark_commit_prefix`) instead of snapshot-replay.
6135 /// The ckpt type is opaque outside spec.rs (newtype) — dflash.rs threads it through.
6136 pub(crate) fn dspark_verify_t_am_ckpt(
6137 &self,
6138 e: &Engine,
6139 tokens: &[u32],
6140 pos0: usize,
6141 cache: &mut Cache,
6142 ) -> Result<(Vec<u32>, DsparkVerifyCkpt), Box<dyn std::error::Error>> {
6143 let mut ck = VerifyCkpt::new(self.layers.len());
6144 let (logits, _hn) = self.decode_step_t_core_stream(
6145 e,
6146 tokens,
6147 pos0,
6148 cache,
6149 None,
6150 Some(&mut ck),
6151 None,
6152 None,
6153 None,
6154 )?;
6155 let t = tokens.len();
6156 let v = self.output.out_features();
6157 let mut am_d = e.stream().alloc_zeros::<u32>(t)?;
6158 for r in 0..t {
6159 e.argmax_token_device_col(&logits, r, v, &mut am_d, r)?;
6160 }
6161 Ok((e.dtoh_u32(&am_d)?, DsparkVerifyCkpt(ck)))
6162 }
6163
6164 /// Engine-bundle slice 2: `dspark_verify_t_am_ckpt` with DEVICE tokens and NO readback.
6165 /// The verify tokens are the round's `chain_d` (cand layout: [anchor, drafts...]); the
6166 /// embed gathers its first `t` entries on-device (`embed_gather_u32_t` — bit-identical
6167 /// rows to the host arm), so the host never blocks on the draft chain before dispatching
6168 /// verify. Returns the device per-row argmax buffer; the caller merges its readback with
6169 /// the chain's into ONE sync. Forward, ckpt fill and argmax walk are `_ckpt` verbatim.
6170 #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
6171 pub(crate) fn dspark_verify_t_am_ckpt_dev(
6172 &self,
6173 e: &Engine,
6174 vtok: &CudaSlice<u32>,
6175 t: usize,
6176 pos0: usize,
6177 cache: &mut Cache,
6178 embd_dev: (&CudaSlice<u8>, i32, usize),
6179 graphs: Option<&mut DsparkVerifyGraphs>,
6180 ) -> Result<(CudaSlice<u32>, DsparkVerifyCkpt), Box<dyn std::error::Error>> {
6181 debug_assert!(
6182 vtok.len() >= t,
6183 "verify window exceeds the device token buffer"
6184 );
6185 // The slab flag is a per-round statement: clear it here so a verify that never
6186 // reaches the graphs door (rowwise env, a non-tparallel arm) cannot leave a
6187 // stale `true` steering the commit at slabs the round never wrote.
6188 let mut graphs = graphs;
6189 if let Some(g) = graphs.as_deref_mut() {
6190 g.round_slab = false;
6191 }
6192 let mut ck = VerifyCkpt::new(self.layers.len());
6193 // Dummy host tokens size the funnel; the embed reads `vtok` (the round-stream
6194 // arm's established pattern — spec.rs stream-mode verify does the same).
6195 let dummy = vec![0u32; t];
6196 let (logits, _hn) = self.decode_step_t_core_stream(
6197 e,
6198 &dummy,
6199 pos0,
6200 cache,
6201 Some(embd_dev),
6202 Some(&mut ck),
6203 None,
6204 Some(vtok),
6205 graphs,
6206 )?;
6207 let v = self.output.out_features();
6208 let mut am_d = e.stream().alloc_zeros::<u32>(t)?;
6209 for r in 0..t {
6210 e.argmax_token_device_col(&logits, r, v, &mut am_d, r)?;
6211 }
6212 Ok((am_d, DsparkVerifyCkpt(ck)))
6213 }
6214
6215 /// Ckpt-armed twin of [`Self::dspark_verify_t_logits`] (sampled-admission arm).
6216 pub(crate) fn dspark_verify_t_logits_ckpt(
6217 &self,
6218 e: &Engine,
6219 tokens: &[u32],
6220 pos0: usize,
6221 cache: &mut Cache,
6222 ) -> Result<(CudaSlice<f32>, DsparkVerifyCkpt), Box<dyn std::error::Error>> {
6223 let mut ck = VerifyCkpt::new(self.layers.len());
6224 let (logits, _hn) = self.decode_step_t_core_stream(
6225 e,
6226 tokens,
6227 pos0,
6228 cache,
6229 None,
6230 Some(&mut ck),
6231 None,
6232 None,
6233 None,
6234 )?;
6235 Ok((logits, DsparkVerifyCkpt(ck)))
6236 }
6237
6238 /// Restore the round to `keep` accepted columns from the verify stash: KV lens and
6239 /// pos from the pre-verify snapshot + keep, GDN conv/ssm from the stashed column
6240 /// state — no replay forward. The exact `commit_verified_prefix` the MTP path ships.
6241 pub(crate) fn dspark_commit_prefix(
6242 &self,
6243 e: &Engine,
6244 cache: &mut Cache,
6245 snap: &crate::cache::CacheSnapshot,
6246 ckpt: &DsparkVerifyCkpt,
6247 keep: usize,
6248 ) -> Result<(), Box<dyn std::error::Error>> {
6249 self.commit_verified_prefix(e, cache, snap, &ckpt.0, keep)
6250 }
6251
6252 /// Slice-3 commit twin: restore to `keep` accepted columns when the round's linear
6253 /// column stash lives in the graphs ctx's persistent slabs (`DsparkVerifyGraphs`) —
6254 /// the cols arm's exact semantics (KV lens + pos from the snapshot, GDN conv/ssm
6255 /// from the stash of column keep-1), slab-addressed and batched into two copy
6256 /// launches.
6257 pub(crate) fn dspark_commit_prefix_slab(
6258 &self,
6259 e: &Engine,
6260 cache: &mut Cache,
6261 snap: &crate::cache::CacheSnapshot,
6262 ctx: &DsparkVerifyGraphs,
6263 keep: usize,
6264 ) -> Result<(), Box<dyn std::error::Error>> {
6265 use cudarc::driver::DevicePtr;
6266 debug_assert!(keep >= 1, "keep==0 rounds take the legacy rollback");
6267 let mut conv_src: Vec<u64> = Vec::new();
6268 let mut ssm_src: Vec<u64> = Vec::new();
6269 let mut conv_dst: Vec<u64> = Vec::new();
6270 let mut ssm_dst: Vec<u64> = Vec::new();
6271 for il in 0..self.layers.len() {
6272 if let (Some(kvl), Some(saved)) = (cache.kv[il].as_mut(), snap.kv_len[il]) {
6273 kvl.len = saved + keep;
6274 e.set_i32_one(&mut kvl.len_d, kvl.len as i32)?;
6275 }
6276 if let Some(rl) = cache.recur[il].as_ref() {
6277 let (pc, ps, _cw, _sw) = ctx
6278 .slab_row(e, il, keep - 1)
6279 .ok_or("slab commit: linear layer missing from the graphs ctx")?;
6280 conv_src.push(pc);
6281 ssm_src.push(ps);
6282 let st = &e.gpu.stream();
6283 let (dc, _g0) = rl.conv_state.device_ptr(st);
6284 let (ds, _g1) = rl.ssm_state.device_ptr(st);
6285 conv_dst.push(dc);
6286 ssm_dst.push(ds);
6287 }
6288 }
6289 let n = conv_src.len();
6290 if n > 0 {
6291 let mut tt = vec![0u64; 2 * n];
6292 tt[..n].copy_from_slice(&conv_src);
6293 tt[n..].copy_from_slice(&conv_dst);
6294 let ct = e.htod_u64(&tt)?;
6295 tt[..n].copy_from_slice(&ssm_src);
6296 tt[n..].copy_from_slice(&ssm_dst);
6297 let st = e.htod_u64(&tt)?;
6298 e.copy_batch_uniform_f32(&ct, n, ctx.conv_words)?;
6299 e.copy_batch_uniform_f32(&st, n, ctx.ssm_words)?;
6300 }
6301 cache.pos = snap.pos + keep;
6302 Ok(())
6303 }
6304
6305 /// Qwen35-family verify trunk in the live serving numeric class.
6306 ///
6307 /// Serving intentionally keeps this architecture in the generic batched program even at
6308 /// B=1. The older verify walk used its own mirrored dispatch and can flip near-tie argmaxes.
6309 ///
6310 /// Two arms, one numeric class:
6311 /// - DENSE GDN (`DenseMlp`, t<=16): `qwen35_verify_tparallel` — the weight ops (norms,
6312 /// projections, FFN) hoist to m=T through the exact-tier batched kernels whose per-row
6313 /// program IS the m=1 program (`matmul_pre == fused2 per (tensor,row); _bN mmvq per-row
6314 /// == m=1` — decode_batch.rs v2 note), while the state ops (conv ring, gdn scan, KV
6315 /// append, fa decode) stay a per-row loop running the b_n=1 serving kernels with each
6316 /// row's own t_kv-driven arm pick (the straddle law: every row executes the exact
6317 /// program its isolated serving step would). One weight read per layer per round
6318 /// instead of T — this is what makes MTP profitable in the exact class (the per-row
6319 /// walk measured verify(K+1) ~= (K+1) plain steps: 69 -> 44 tok/s served, 2026-08-15).
6320 /// - MoE / t>16 / `MEMRA_SPEC_VERIFY_ROWWISE=1`: the per-row replay of the authoritative
6321 /// serving layer body, preserving single-session autoregressive cache order (the
6322 /// correctness reference; also the rollback seam for the t-parallel arm).
6323 ///
6324 /// Bit-identity of the t-parallel arm vs the rowwise arm is gated by spec-serve-gate
6325 /// (zero differing logits at T=1..4, K arms) + the 8-prompt ON/OFF canary before ship.
6326 #[allow(clippy::too_many_arguments)]
6327 fn qwen35_verify_batch_layers(
6328 &self,
6329 e: &Engine,
6330 x: CudaSlice<f32>,
6331 lo: usize,
6332 hi: usize,
6333 pos0: usize,
6334 t: usize,
6335 cache: &mut Cache,
6336 ckpt: Option<&mut VerifyCkpt>,
6337 stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
6338 graphs: Option<&mut DsparkVerifyGraphs>,
6339 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6340 // Qwen35Moe admitted 2026-08-20 (lane/draftcost-moe): the t-parallel arm already
6341 // carries the MoE FFN (`moe_ffn_il_zq8` at m=T) and the GDN per-row state loop; the
6342 // arch fence was a qualification gate, not a mechanism gap. Measured disease on the
6343 // 35B-A3B class: rowwise verify ~= 5.6 ms per drafted token (one full trunk step
6344 // each) — the same (K+1)-plain-steps wall the dense admission fixed on 2026-08-15.
6345 // Rollback seam unchanged: MEMRA_SPEC_VERIFY_ROWWISE=1.
6346 let rowwise = std::env::var("MEMRA_SPEC_VERIFY_ROWWISE").as_deref() == Ok("1")
6347 || !self.batched_serving_numeric_class()
6348 || t > 16;
6349 if rowwise {
6350 if stream.is_some() {
6351 // rowwise replays per row with host cache.pos — irreconcilable with a
6352 // device position counter. Burst callers must keep t <= 16 and the
6353 // ROWWISE env unset; refusing beats silently mispositioned rows.
6354 return Err("qwen35 rowwise verify has no ROUND-STREAM arm \
6355 (t > 16 or MEMRA_SPEC_VERIFY_ROWWISE=1)"
6356 .into());
6357 }
6358 self.qwen35_verify_rowwise(e, x, lo, hi, pos0, t, cache, ckpt)
6359 } else {
6360 self.qwen35_verify_tparallel(e, x, lo, hi, pos0, t, cache, ckpt, stream, graphs)
6361 }
6362 }
6363
6364 /// The per-row correctness reference: replay each verify row through the authoritative
6365 /// serving layer body (`decode_batch_layers` at b_n=1). T full weight reads per layer.
6366 #[allow(clippy::too_many_arguments)]
6367 fn qwen35_verify_rowwise(
6368 &self,
6369 e: &Engine,
6370 mut x: CudaSlice<f32>,
6371 lo: usize,
6372 hi: usize,
6373 pos0: usize,
6374 t: usize,
6375 cache: &mut Cache,
6376 mut ckpt: Option<&mut VerifyCkpt>,
6377 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6378 let n_embd = self.cfg.n_embd as usize;
6379 let saved_pos = cache.pos;
6380 let mut ph_last = std::time::Instant::now();
6381 for il in lo..hi {
6382 let mut next = e.uninit(t * n_embd)?;
6383 let mut col_states: Option<Vec<(CudaSlice<f32>, CudaSlice<f32>)>> =
6384 if ckpt.is_some() && t >= 2 && matches!(self.layers[il].mixer, Mixer::Linear(_)) {
6385 Some(Vec::with_capacity(t - 1))
6386 } else {
6387 None
6388 };
6389 for r in 0..t {
6390 cache.pos = pos0 + r;
6391 let mut row = e.uninit(n_embd)?;
6392 e.dtod_copy_view(&x.slice(r * n_embd..(r + 1) * n_embd), &mut row)?;
6393 let row_pos = e.htod_i32(&[(pos0 + r) as i32])?;
6394 let mut one = [&mut *cache];
6395 let ctx = self.batch_layer_ctx(e, &one, il, il + 1)?;
6396 let out = match self.decode_batch_layers(
6397 e,
6398 row,
6399 &mut one,
6400 &ctx,
6401 &row_pos,
6402 &mut ph_last,
6403 ) {
6404 Ok(out) => out,
6405 Err(error) => {
6406 cache.pos = saved_pos;
6407 return Err(error);
6408 }
6409 };
6410 e.dtod_copy_into(&out, &mut next, r * n_embd)?;
6411 if r + 1 < t
6412 && let Some(states) = col_states.as_mut()
6413 {
6414 let recur = cache.recur[il]
6415 .as_ref()
6416 .ok_or("Qwen35-MoE linear verify layer has no recurrent state")?;
6417 states.push((
6418 e.clone_dtod(&recur.conv_state)?,
6419 e.clone_dtod(&recur.ssm_state)?,
6420 ));
6421 }
6422 }
6423 if let (Some(checkpoint), Some(states)) = (ckpt.as_deref_mut(), col_states) {
6424 checkpoint.cols[il] = Some(states);
6425 }
6426 x = next;
6427 }
6428 cache.pos = saved_pos;
6429 Ok(x)
6430 }
6431
6432 /// T-PARALLEL VERIFY IN THE SERVING NUMERIC CLASS (lane/tparallel-verify, 2026-08-15).
6433 ///
6434 /// The weight ops run ONCE per layer at m=T; the state ops run per row through the same
6435 /// b_n=1 serving kernels the rowwise replay uses. Per-row bit-identity rests on the two
6436 /// pins the serving batch tier already carries:
6437 /// * `matmul_pre` / `_bN` mmvq: per-row program == m=1 program (decode_batch.rs v2 note,
6438 /// kernel-check pinned) — so a [T, n_embd] projection row equals the row projected
6439 /// alone;
6440 /// * row-indexed norms/elementwise (`rms_norm`, `quantize_q8_1`, `add_rms_norm`,
6441 /// `gated_rmsnorm[_q8_1]`, `silu_mul`, `rope_neox` with per-row positions): the T-row
6442 /// launch is the per-row program (same pin the generic verify's fused norms rely on).
6443 /// The sequential dependencies keep their exact serving order: the conv ring / gdn scan
6444 /// chain state row -> row through the `_b` kernels at b_n=1 (ping-pong via a 6-entry
6445 /// alternating pointer table, host handles swapped per row so VerifyCkpt clones the
6446 /// canonical state exactly as the rowwise arm does), and each row's KV append + fa decode
6447 /// picks its arm from ITS OWN t_kv (append: format-only; fa: `fa_seqs_eligible` + its own
6448 /// `fa_split_keys` rung at b_n=1) — the straddle law per row, so every row executes the
6449 /// program its isolated B=1 serving step would.
6450 ///
6451 /// Cost: 1 weight read per layer per round + T state micro-launches, vs the rowwise arm's
6452 /// T weight reads. Gated bit-identical vs the rowwise arm by spec-serve-gate + canary.
6453 #[allow(clippy::too_many_arguments)]
6454 fn qwen35_verify_tparallel(
6455 &self,
6456 e: &Engine,
6457 mut x: CudaSlice<f32>,
6458 lo: usize,
6459 hi: usize,
6460 pos0: usize,
6461 t: usize,
6462 cache: &mut Cache,
6463 mut ckpt: Option<&mut VerifyCkpt>,
6464 stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
6465 mut graphs: Option<&mut DsparkVerifyGraphs>,
6466 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6467 // Merge guard (v0.98 train, re-affirmed on the v0.100 train over slice 4c): the
6468 // ROUND-STREAM arm (lane/draftcost-moe, device position counter) and the dspark
6469 // verify graphs (engine-bundle slice 3 / trunk slice 4c) have no common caller —
6470 // stream rides the qwen35moe burst, graphs ride the dspark route. If a future
6471 // caller arms both, refuse loudly instead of silently dropping the graphs ctx
6472 // (the stream linear arm takes linear_attn_verify_t, not the graphed segment or
6473 // full-verify bodies).
6474 if stream.is_some() && graphs.is_some() {
6475 return Err(
6476 "qwen35 tparallel verify: ROUND-STREAM and dspark verify graphs \
6477 cannot arm together"
6478 .into(),
6479 );
6480 }
6481 // Engine-bundle slice 3 + slice 4c: with a graphs ctx armed, pointer tables are
6482 // refreshed once per verify (the gdn ping-pong moves handles; a fresh generation
6483 // moves the kv caches). Then:
6484 // - slice 4c: when the WHOLE round rides one seqs rung (every row batchable, one
6485 // split-ladder step, rung covers the round), the ENTIRE walk replays as ONE
6486 // full-verify graph per (vt, rung) — linear layers through the shared
6487 // `qwen35_tparallel_linear_layer` body, full-attention layers through the
6488 // shared `qwen35_tparallel_fa_layer` body in graph mode.
6489 // - fallback (straddle rounds, below the vec floor, partial walks): runs of
6490 // consecutive LINEAR layers replay the slice-3 per-(segment, vt) graphs and
6491 // the full-attention layers run eager (batched rows when eligible).
6492 //
6493 // GRAPH-LAUNCH HEADROOM GUARD (see GRAPH_LAUNCH_MIN_FREE): the dspark verify
6494 // graphs replay through this walk from THREE callers — the MTP spec round's vg
6495 // door (already dropped per round by `graph_round_ok` before it gets here), the
6496 // dspark one-shot, and the dspark SERVE round (default ON since v0.108). Below
6497 // the driver-free floor the WHOLE round takes the byte-identical eager
6498 // cols-ckpt walk — the same drop-the-ctx fallback the pool ceiling already
6499 // takes — instead of feeding cuGraphLaunch a card it segfaults on.
6500 if let Some(g) = graphs.as_deref_mut()
6501 && !graph_launch_headroom_ok(e)
6502 {
6503 g.round_slab = false;
6504 graphs = None;
6505 static NOTED: std::sync::Once = std::sync::Once::new();
6506 NOTED.call_once(|| graph_replay_suspended_note("dspark-vg"));
6507 }
6508 if let Some(g) = graphs.as_deref_mut() {
6509 g.refresh_tables(e, cache)?;
6510 g.round_slab = false;
6511 if let Some(rung) = g.full_rung(self, cache, lo, hi, t) {
6512 // Pool ceiling (dspark_vg_cap): an existing key always replays; a NEW
6513 // full capture past the ceiling falls through to the segment/eager arms.
6514 if g.full.contains_key(&(t, rung, hi)) || g.can_capture() {
6515 let out = g.run_full(self, e, lo, hi, &x, t, pos0, rung, cache)?;
6516 g.round_slab = true;
6517 return Ok(out);
6518 }
6519 }
6520 // Round-atomic ceiling check for the segment door: if any linear run in this
6521 // walk would need a NEW capture past the ceiling, the whole round runs the
6522 // eager cols-ckpt walk (mixing slab- and cols-stashed layers in one round
6523 // would corrupt the commit).
6524 if !g.segments_ready(self, lo, hi, t) {
6525 graphs = None;
6526 }
6527 }
6528 // STREAM (2b, lane/draftcost-moe): positions come from the device round counter
6529 // (pos_iota / i32_copy_add) so a burst round needs no host position knowledge.
6530 let pos_d = match stream {
6531 Some((_, ctr)) => {
6532 let mut p = e.alloc_uninit::<i32>(t)?;
6533 e.pos_iota(ctr, &mut p, t)?;
6534 p
6535 }
6536 None => {
6537 let pos_host: Vec<i32> = (0..t).map(|r| (pos0 + r) as i32).collect();
6538 e.htod_i32(&pos_host)?
6539 }
6540 };
6541 // Per-row 1-element position buffers, built ONCE per verify (the append/fa wrappers
6542 // take owned pos slices; building these inside the layer x row loops cost 16xT H2Ds).
6543 // LAZY since slice 4: the batched fa/append arm never touches them — they are built
6544 // on the first per-row fallback layer only (stream-aware there; the stream FA arm
6545 // rides the dc rows kernels and never reaches the fallback).
6546 let mut pos_rows: Option<Vec<CudaSlice<i32>>> = None;
6547 let mut il = lo;
6548 while il < hi {
6549 if graphs.is_some() && matches!(self.layers[il].mixer, Mixer::Linear(_)) {
6550 let mut end = il;
6551 while end < hi && matches!(self.layers[end].mixer, Mixer::Linear(_)) {
6552 end += 1;
6553 }
6554 let g = graphs.as_deref_mut().expect("checked above");
6555 x = g.run_segment(self, e, il, end, &x, t, cache)?;
6556 g.round_slab = true;
6557 il = end;
6558 continue;
6559 }
6560 let layer = &self.layers[il];
6561 if stream.is_none() && matches!(layer.mixer, Mixer::Linear(_)) {
6562 // Eager linear layer (no graphs ctx): the shared body, legacy cols-ckpt arm.
6563 // Under ROUND-STREAM the linear layers ride the fa-body match's stream arm
6564 // below (linear_attn_verify_t — the stream COMMIT needs its GdnStash).
6565 x = self.qwen35_tparallel_linear_layer(
6566 e,
6567 il,
6568 &x,
6569 t,
6570 cache,
6571 ckpt.as_deref_mut(),
6572 None,
6573 None,
6574 )?;
6575 il += 1;
6576 continue;
6577 }
6578 // Full-attention (or stream-Linear, or MLA-refusing) layer: the extracted
6579 // shared body — eager arm (fresh per-verify pos/table, exact t_kv sizing,
6580 // in-body len bump). The slice-4c captured full-verify graphs run the SAME
6581 // body in graph mode; under ROUND-STREAM the body's dc-rows / GDN stream arms
6582 // run (lane/draftcost-moe).
6583 x = self.qwen35_tparallel_fa_layer(
6584 e,
6585 il,
6586 &x,
6587 t,
6588 cache,
6589 FaLayerArgs {
6590 pos_d: &pos_d,
6591 pos_rows: &mut pos_rows,
6592 pos0,
6593 graph_cap: None,
6594 stream,
6595 ckpt: ckpt.as_deref_mut(),
6596 },
6597 )?;
6598 il += 1;
6599 }
6600 Ok(x)
6601 }
6602
6603 /// SHARED dense-FFN body for the qwen35 t-parallel layers (trunk-kernels slice B) —
6604 /// ONE copy for the fa and linear layer bodies (the verify_layers extraction lesson).
6605 /// Dual arm (MEMRA_TK_FFN_DUAL, default on): gate+up in ONE dual launch from the
6606 /// pre-quantized activation with macro-scales DEFERRED into the fused SwiGLU+q8_1
6607 /// epilogue, then ffn_down from the fused (aq, ad) — the q27 verify chain verbatim.
6608 /// Every door is the bit-identical proven one: `matmul_decode_exact_dual_pre` (per
6609 /// (tensor,token,row) == the two singles), `silu_mul_scaled_q8_1` (y*s inline == the
6610 /// scale_inplace store, value-exact; fused quantize == quantize_q8_1 bytes),
6611 /// `matmul_decode_exact_pre` (dispatch mirror of the singles' q8_1-fast tail).
6612 /// Dual-refused (t outside 2..=7, non-NVFP4, layout mismatch) or seam off -> the
6613 /// original singles chain, byte-for-byte.
6614 #[allow(clippy::too_many_arguments)]
6615 fn qwen35_tparallel_dense_ffn(
6616 &self,
6617 e: &Engine,
6618 ffn_gate: &crate::model::GpuTensor,
6619 ffn_up: &crate::model::GpuTensor,
6620 ffn_down: &crate::model::GpuTensor,
6621 zn: &CudaSlice<f32>,
6622 t: usize,
6623 n_embd: usize,
6624 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6625 let n_ff = ffn_gate.out_features();
6626 let (zq, zd) = e.quantize_q8_1(zn, t, n_embd)?;
6627 if Engine::tk_ffn_dual_on()
6628 && let Some(((g, gs), (u, us))) =
6629 e.matmul_decode_exact_dual_pre(ffn_gate, ffn_up, &zq, &zd, t)?
6630 {
6631 if e.uses_q8_1_fast(ffn_down) {
6632 let (aq, ad) = e.silu_mul_scaled_q8_1(&g, &u, gs, us, t * n_ff)?;
6633 return e.matmul_decode_exact_pre(ffn_down, &aq, &ad, t);
6634 }
6635 let mut act = e.uninit(t * n_ff)?;
6636 e.silu_mul_scaled(&g, &u, gs, us, &mut act, t * n_ff)?;
6637 let (aq, ad) = e.quantize_q8_1(&act, t, n_ff)?;
6638 return e.matmul_pre(ffn_down, &aq, &ad, &act, t);
6639 }
6640 // v1 singles chain (seam off or dual-refused) — the pre-slice-B body verbatim.
6641 let g = e.matmul_pre(ffn_gate, &zq, &zd, zn, t)?;
6642 let u = e.matmul_pre(ffn_up, &zq, &zd, zn, t)?;
6643 let mut act = e.uninit(t * n_ff)?;
6644 e.silu_mul(&g, &u, &mut act, t * n_ff)?;
6645 let (aq, ad) = e.quantize_q8_1(&act, t, n_ff)?;
6646 e.matmul_pre(ffn_down, &aq, &ad, &act, t)
6647 }
6648
6649 /// ONE t-parallel FULL-ATTENTION layer (attn_norm + fa mixer + post_attn_norm + FFN +
6650 /// tap) — extracted from the walk exactly like `qwen35_tparallel_linear_layer` so the
6651 /// eager walk and the slice-4c captured full-verify graphs execute the SAME body (a
6652 /// second copy is how dispatch mirrors drift — the verify_layers extraction lesson).
6653 ///
6654 /// `args.graph_cap = Some((table, off, rung_end))` is the captured-graph mode:
6655 /// - kv base-pointer pairs come from the ctx-owned persistent table at `off` (a fresh
6656 /// generation's cache lands at new addresses that only the per-verify table refresh
6657 /// knows — the slice-3 baked-address lesson);
6658 /// - the seqs twins size partials/grid at `rung_end` and pin `split_keys` to the
6659 /// rung's ladder value: `n_splits_max` is pure stride, splits >= ns_eff write the
6660 /// EMPTY partial the combine never reads, and every per-row T_kv derives in-kernel
6661 /// from `pos_seq[z]` — so one captured launch replays bit-identically for every
6662 /// round whose rows all sit inside the rung;
6663 /// - the host len bump moves to the replay caller (captured host code does not
6664 /// re-run at replay).
6665 /// Graph mode REFUSES any round the batched arm cannot take: the per-row fallback
6666 /// host-branches on t_kv and must never be captured.
6667 #[allow(clippy::too_many_arguments)]
6668 fn qwen35_tparallel_fa_layer(
6669 &self,
6670 e: &Engine,
6671 il: usize,
6672 x: &CudaSlice<f32>,
6673 t: usize,
6674 cache: &mut Cache,
6675 args: FaLayerArgs<'_>,
6676 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6677 use cudarc::driver::DevicePtr;
6678 let cfg = &self.cfg;
6679 let n_embd = cfg.n_embd as usize;
6680 let eps = cfg.rms_eps;
6681 let head_dim_global = cfg.head_dim_k as usize;
6682 let layer = &self.layers[il];
6683 let FaLayerArgs {
6684 pos_d,
6685 pos_rows,
6686 pos0,
6687 graph_cap,
6688 stream,
6689 ckpt,
6690 } = args;
6691
6692 // ---- attn_norm + q8_1 quantize at m=T (row-indexed == per-row) ----
6693 let anorm = layer.attn_norm.float_data();
6694 let mut xn = e.uninit(t * n_embd)?;
6695 e.rms_norm(x, anorm, &mut xn, n_embd, t, eps)?;
6696 let (hq, hd) = e.quantize_q8_1(&xn, t, n_embd)?;
6697
6698 let mixed: CudaSlice<f32> = match &layer.mixer {
6699 Mixer::Mla(_) => crate::hybrid::mla_path_unimplemented("tensor-parallel attention"),
6700 Mixer::Kda(_) => crate::hybrid::kda_path_unimplemented("T-parallel attention"),
6701 // STREAM ARM (2b, lane/draftcost-moe): under a device position counter the
6702 // per-row serving-kernel chain cannot run (host state swaps keyed on host
6703 // row index are fine, but the stream COMMIT needs the GdnStash for its _dc
6704 // rebuild — the per-row chain only produces per-column clones). GDN rides
6705 // `linear_attn_verify_t`: batched q8_1-class projections, stash-producing,
6706 // and its one-scan recurrence is pinned bit-identical to T chained T=1
6707 // steps (its header + kernel-check). Position-independent, so no counter
6708 // plumbing is needed. Guards mirror the generic call site exactly.
6709 Mixer::Linear(la) if stream.is_some() => {
6710 if !(t >= 3 || (t == 2 && spec_m2()))
6711 || !self.mixer_in_q8_1_fast(e, &layer.mixer)
6712 || !e.uses_q8_1_fast(&la.ssm_out)
6713 {
6714 return Err("qwen35 stream verify: GDN batched arm requires t>=3 \
6715 (or MEMRA_SPEC_M2 at t=2) and q8_1-fast projections"
6716 .into());
6717 }
6718 let want = ckpt.is_some();
6719 let (out, stash) =
6720 self.linear_attn_verify_t(e, la, &xn, Some((&hq, &hd)), t, cache, il, want)?;
6721 if let (Some(ck), Some(st)) = (ckpt, stash) {
6722 ck.gdn[il] = Some(st);
6723 }
6724 out
6725 }
6726 Mixer::Linear(_) => {
6727 unreachable!("linear layers ride qwen35_tparallel_linear_layer")
6728 }
6729 Mixer::Full(fa) => {
6730 let geometry = cfg.full_attention_geometry_at(il as u32);
6731 let n_head = geometry.n_head as usize;
6732 let n_head_kv = geometry.n_head_kv as usize;
6733 let head_dim = geometry.head_dim_k as usize;
6734 let rope_dims = geometry.n_rot as usize;
6735 let rope_base = geometry.rope_base;
6736 let scale = geometry.attention_scale();
6737 // Batched projections: one weight read serves all T rows.
6738 // GROUP-3 twin (trunk-kernels slice D): q/k/v in ONE launch — the group4
6739 // kernel with n3=0, bit-identical per (tensor, token, row) to the three
6740 // singles; refused or MEMRA_TK_FA_GROUP=0 -> singles byte-for-byte.
6741 let (qf, mut k, v) = match e.matmul_decode_exact_group3_pre(
6742 [&fa.wq, &fa.wk, &fa.wv],
6743 &hq,
6744 &hd,
6745 t,
6746 )? {
6747 Some(mut g3) => {
6748 let v = g3.pop().unwrap();
6749 let k = g3.pop().unwrap();
6750 let qf = g3.pop().unwrap();
6751 (qf, k, v)
6752 }
6753 None => (
6754 e.matmul_pre(&fa.wq, &hq, &hd, &xn, t)?,
6755 e.matmul_pre(&fa.wk, &hq, &hd, &xn, t)?,
6756 e.matmul_pre(&fa.wv, &hq, &hd, &xn, t)?,
6757 ),
6758 };
6759 let gated =
6760 geometry.attention_gate == memra_gguf::config::AttentionGateKind::FusedQ;
6761 let (mut q, gate) = if gated {
6762 let mut qs = e.uninit(t * n_head * head_dim)?;
6763 let mut gs = e.uninit(t * n_head * head_dim)?;
6764 e.q_gate_split(&qf, &mut qs, &mut gs, head_dim, n_head, t)?;
6765 (qs, Some(gs))
6766 } else {
6767 (qf, None)
6768 };
6769 let mut qn = e.uninit(t * n_head * head_dim)?;
6770 e.rms_norm(
6771 &q,
6772 fa.q_norm.float_data(),
6773 &mut qn,
6774 head_dim,
6775 t * n_head,
6776 eps,
6777 )?;
6778 q = qn;
6779 let mut kn = e.uninit(t * n_head_kv * head_dim)?;
6780 e.rms_norm(
6781 &k,
6782 fa.k_norm.float_data(),
6783 &mut kn,
6784 head_dim,
6785 t * n_head_kv,
6786 eps,
6787 )?;
6788 k = kn;
6789 e.rope_neox(
6790 &mut q, pos_d, head_dim, rope_dims, n_head, t, rope_base, 1.0,
6791 )?;
6792 e.rope_neox(
6793 &mut k, pos_d, head_dim, rope_dims, n_head_kv, t, rope_base, 1.0,
6794 )?;
6795
6796 // Per-row append + attend: row r sees rows 0..r in KV (causal within the
6797 // draft), each through the b_n=1 serving kernels at its own t_kv.
6798 let q_dim = n_head * head_dim;
6799 let kv_dim = n_head_kv * head_dim;
6800 let mut attn = e.uninit(t * q_dim)?;
6801 let (kdk, kdv, ktb, vtb, len0, kv_local) = {
6802 let kvl = cache.kv[il].as_ref().unwrap();
6803 // [2T] interleaved k,v base pointers: entry pair z serves row z of
6804 // the batched twins; the per-row fallback reads pair 0 (same cache
6805 // for every row of one layer). Graph mode reads the ctx table.
6806 let local: Option<CudaSlice<u64>> = match graph_cap {
6807 Some(_) => None,
6808 None => {
6809 let s = &e.gpu.stream();
6810 let (pk, _g) = kvl.k.device_ptr(s);
6811 let (pv, _g2) = kvl.v.device_ptr(s);
6812 let mut tbl = Vec::with_capacity(2 * t);
6813 for _ in 0..t {
6814 tbl.push(pk);
6815 tbl.push(pv);
6816 }
6817 Some(e.htod_u64(&tbl)?)
6818 }
6819 };
6820 (
6821 kvl.kv_dim_k,
6822 kvl.kv_dim_v,
6823 kvl.k_tok_bytes,
6824 kvl.v_tok_bytes,
6825 kvl.len,
6826 local,
6827 )
6828 };
6829 let (kv_tbl, kv_off): (&CudaSlice<u64>, usize) = match graph_cap {
6830 Some((tb, off, _)) => (tb, off),
6831 None => (kv_local.as_ref().expect("built above"), 0),
6832 };
6833 // Slice 4 (fa/append rows — see dspark_fa_rows_on): the whole per-row
6834 // section batches into the z-batched serving twins when every row of
6835 // this round takes the v4-seqs arm on ONE fa_split_keys rung. Both
6836 // guards are evaluated at the round's FIRST and LAST t_kv — the
6837 // eligibility window (vec floor .. v4 max) and each split-ladder rung
6838 // are intervals in t_kv, so ends-inside means all-inside (the straddle
6839 // law). Appending all T rows before any attend is read-equivalent to
6840 // the interleaved order: row r's walk reads keys 0..len0+r only, and
6841 // rows > r land at slots it never touches; every written cache row is
6842 // the per-token appender's exact warp program (kernel-check pinned).
6843 let t_kv_first = len0 + 1;
6844 let t_kv_last = len0 + t;
6845 let rows_batched = t >= 2
6846 && dspark_fa_rows_on()
6847 // the z-batched twins read stacked rows at the CACHE's kv dims;
6848 // the projection stack is [T, n_head_kv*head_dim] — they must be
6849 // the same stride or row z misaligns (true for this family; the
6850 // guard keeps any asymmetric-kv model on the per-row loop).
6851 && kdk == kv_dim
6852 && kdv == kv_dim
6853 && crate::fa_seqs_eligible(t_kv_first, head_dim_global)
6854 && crate::fa_seqs_eligible(t_kv_last, head_dim_global)
6855 && crate::fa_split_keys(t_kv_first, cfg.n_head_kv as usize)
6856 == crate::fa_split_keys(t_kv_last, cfg.n_head_kv as usize);
6857 // Sizing: eager = exact round bound; graph mode = the rung end (stride +
6858 // grid only — bytes proven equal above). Capture-time invariants refuse
6859 // loudly rather than bake a divergent body.
6860 let (size_kv_max, sp) = match graph_cap {
6861 Some((_, _, rung)) => {
6862 if !rows_batched {
6863 return Err(format!(
6864 "fa graph capture: layer {il} round is not batchable \
6865 (t_kv {t_kv_first}..{t_kv_last}) — the per-row fallback \
6866 must never be captured"
6867 )
6868 .into());
6869 }
6870 let sp_r = crate::fa_split_keys(rung, cfg.n_head_kv as usize);
6871 if t_kv_last > rung
6872 || sp_r != crate::fa_split_keys(t_kv_last, cfg.n_head_kv as usize)
6873 {
6874 return Err(format!(
6875 "fa graph capture: rung {rung} does not cover round \
6876 t_kv {t_kv_first}..{t_kv_last} on one split ladder step"
6877 )
6878 .into());
6879 }
6880 (rung, sp_r)
6881 }
6882 None => (
6883 t_kv_last,
6884 crate::fa_split_keys(t_kv_last, cfg.n_head_kv as usize),
6885 ),
6886 };
6887 if let Some((_, ctr)) = stream {
6888 // STREAM ARM (2b): one batched dc append + the multi-row dc attention
6889 // — the generic stream arm's exact shape (rows kernels are pinned
6890 // byte-identical to the per-row programs by kernel-check). Host len
6891 // stays a stale lower bound; the burst drain reconciles it.
6892 let kvl = cache.kv[il].as_mut().unwrap();
6893 e.append_kv_quantized_rows_dc(
6894 &k, &v, &mut kvl.k, &mut kvl.v, ctr, t, kdk, kdv, ktb, vtb, false,
6895 )?;
6896 let upper = (kvl.len + t + 64).min(cache.max_ctx);
6897 let k_view = e.view_u8(&kvl.k, upper * ktb);
6898 let v_view = e.view_u8(&kvl.v, upper * vtb);
6899 e.fa_decode_rows_dc(
6900 &q, &k_view, &v_view, &mut attn, head_dim, n_head, n_head_kv, ctr, upper,
6901 t, scale, ktb, vtb, 0, false,
6902 )?;
6903 } else if rows_batched {
6904 e.append_kv_quantized_seqs(
6905 &k,
6906 &v,
6907 &kv_tbl.slice(kv_off..kv_off + 2 * t),
6908 pos_d,
6909 t,
6910 kdk,
6911 kdv,
6912 ktb,
6913 vtb,
6914 )?;
6915 if graph_cap.is_none() {
6916 cache.kv[il].as_mut().unwrap().len += t;
6917 }
6918 e.fa_decode_batch_seqs_v4(
6919 &q,
6920 &kv_tbl.slice(kv_off..kv_off + 2 * t),
6921 pos_d,
6922 &mut attn,
6923 head_dim,
6924 n_head,
6925 n_head_kv,
6926 t,
6927 size_kv_max,
6928 scale,
6929 sp,
6930 ktb,
6931 vtb,
6932 )?;
6933 } else {
6934 if pos_rows.is_none() {
6935 // Stream-aware for symmetry with pos_d (the stream FA arm rides
6936 // the dc rows kernels above and never reaches this fallback).
6937 *pos_rows = Some(match stream {
6938 Some((_, ctr)) => (0..t)
6939 .map(|r| {
6940 let mut b = e.alloc_uninit::<i32>(1)?;
6941 e.i32_copy_add(ctr, &mut b, r as i32)?;
6942 Ok(b)
6943 })
6944 .collect::<Result<_, Box<dyn std::error::Error>>>()?,
6945 None => (0..t)
6946 .map(|r| e.htod_i32(&[(pos0 + r) as i32]))
6947 .collect::<Result<_, _>>()?,
6948 });
6949 }
6950 let pos_rows = pos_rows.as_ref().unwrap();
6951 #[allow(clippy::needless_range_loop)]
6952 // allow: the explicit index loop keeps the offset arithmetic visible and aligned with the device-side indexing
6953 for r in 0..t {
6954 // Owned per-row scratch: the b_n=1 kernels take packed batch buffers
6955 // whose row 0 is this row (arithmetic-free materialization copies,
6956 // same as decode's per-seq fallback arm).
6957 let mut k_row = e.uninit(kv_dim)?;
6958 e.dtod_copy_view(&k.slice(r * kv_dim..(r + 1) * kv_dim), &mut k_row)?;
6959 let mut v_row = e.uninit(kv_dim)?;
6960 e.dtod_copy_view(&v.slice(r * kv_dim..(r + 1) * kv_dim), &mut v_row)?;
6961 let pos_row = &pos_rows[r];
6962 let kvl = cache.kv[il].as_mut().unwrap();
6963 e.append_kv_quantized_seqs(
6964 &k_row,
6965 &v_row,
6966 &kv_tbl.slice(kv_off..kv_off + 2),
6967 pos_row,
6968 1,
6969 kdk,
6970 kdv,
6971 ktb,
6972 vtb,
6973 )?;
6974 kvl.len += 1;
6975 let t_kv = kvl.len;
6976 let mut q_row = e.uninit(q_dim)?;
6977 e.dtod_copy_view(&q.slice(r * q_dim..(r + 1) * q_dim), &mut q_row)?;
6978 let mut a_row = e.uninit(q_dim)?;
6979 if crate::fa_seqs_eligible(t_kv, head_dim_global) {
6980 let sp0_r = crate::fa_split_keys(t_kv, cfg.n_head_kv as usize);
6981 e.fa_decode_batch_seqs_v4(
6982 &q_row,
6983 &kv_tbl.slice(kv_off..kv_off + 2),
6984 pos_row,
6985 &mut a_row,
6986 head_dim,
6987 n_head,
6988 n_head_kv,
6989 1,
6990 t_kv,
6991 scale,
6992 sp0_r,
6993 ktb,
6994 vtb,
6995 )?;
6996 } else {
6997 let k_view = e.view_u8(&kvl.k, t_kv * kvl.k_tok_bytes);
6998 let v_view = e.view_u8(&kvl.v, t_kv * kvl.v_tok_bytes);
6999 let mut a_view = a_row.slice_mut(0..q_dim);
7000 e.fa_decode_kvmod_view(
7001 &q_row.slice(0..q_dim),
7002 &k_view,
7003 &v_view,
7004 &mut a_view,
7005 head_dim,
7006 n_head,
7007 n_head_kv,
7008 t_kv,
7009 scale,
7010 kvl.k_tok_bytes,
7011 kvl.v_tok_bytes,
7012 false,
7013 )?;
7014 }
7015 e.dtod_copy_into(&a_row, &mut attn, r * q_dim)?;
7016 }
7017 }
7018
7019 // Output gate (element-wise) + o-proj at m=T.
7020 let attn_g = match &gate {
7021 Some(g) => {
7022 let n = t * q_dim;
7023 let mut gsig = e.uninit(n)?;
7024 e.sigmoid(g, &mut gsig, n)?;
7025 let mut ag = e.uninit(n)?;
7026 e.mul(&attn, &gsig, &mut ag, n)?;
7027 ag
7028 }
7029 None => attn,
7030 };
7031 e.matmul(&fa.wo, &attn_g, t)?
7032 }
7033 };
7034
7035 // ---- residual add + post_attn_norm + FFN at m=T (serving dispatch verbatim) ----
7036 let pnorm = layer.post_attn_norm.float_data();
7037 let mut x1 = e.uninit(t * n_embd)?;
7038 let mut zn = e.uninit(t * n_embd)?;
7039 e.add_rms_norm(x, &mixed, pnorm, &mut x1, &mut zn, n_embd, t, eps)?;
7040 let ffn_out = match &layer.ffn {
7041 crate::hybrid::Ffn::Dense {
7042 ffn_gate,
7043 ffn_up,
7044 ffn_down,
7045 } => {
7046 assert!(
7047 self.cfg.m3.is_none(),
7048 "qwen35 t-parallel verify: M3 swigluoai FFN not yet batched"
7049 );
7050 self.qwen35_tparallel_dense_ffn(e, ffn_gate, ffn_up, ffn_down, &zn, t, n_embd)?
7051 }
7052 crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il_zq8(e, m, &zn, None, t, il as u16)?,
7053 };
7054 let mut x2 = e.uninit(t * n_embd)?;
7055 e.add(&x1, &ffn_out, &mut x2, t * n_embd)?;
7056 // dspark drafter tap (no-op when no sink armed): post-layer residual verify rows
7057 self.dflash_tap(e, cache, il, &x2, t)?;
7058 Ok(x2)
7059 }
7060
7061 /// ONE t-parallel LINEAR layer (attn_norm + gdn mixer + post_attn_norm + FFN + tap) —
7062 /// the exact body the old in-loop Linear arm ran, extracted so the eager walk and the
7063 /// slice-3 captured segments execute the SAME code (a second copy is how dispatch
7064 /// mirrors drift — the verify_layers extraction lesson). Two deliberate changes, both
7065 /// bit-identical by construction:
7066 /// - the gdn ping-pong host swap moves from per-row to ONE end-of-body swap (t odd):
7067 /// the device sequence is driven entirely by the 6-entry pointer table, which
7068 /// already encodes both parities; the ckpt stash reads name row r's out buffer
7069 /// directly (r even -> alt handle, odd -> canonical) — the same physical bytes the
7070 /// legacy post-swap clone read.
7071 /// - `stash` (slice-3 ctx): persistent per-layer slabs written by copy_into instead of
7072 /// per-row clone_dtod allocs — same bytes, capture-legal (no per-round host objects).
7073 /// `table_src` = (persistent pointer table, offset) when the ctx owns the tables;
7074 /// None builds the per-verify table exactly as before.
7075 #[allow(clippy::too_many_arguments)]
7076 fn qwen35_tparallel_linear_layer(
7077 &self,
7078 e: &Engine,
7079 il: usize,
7080 x: &CudaSlice<f32>,
7081 t: usize,
7082 cache: &mut Cache,
7083 ckpt: Option<&mut VerifyCkpt>,
7084 stash: Option<(&mut CudaSlice<f32>, &mut CudaSlice<f32>)>,
7085 table_src: Option<(&CudaSlice<u64>, usize)>,
7086 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7087 use cudarc::driver::DevicePtr;
7088 let cfg = &self.cfg;
7089 let n_embd = cfg.n_embd as usize;
7090 let eps = cfg.rms_eps;
7091 let layer = &self.layers[il];
7092 let Mixer::Linear(la) = &layer.mixer else {
7093 return Err("qwen35_tparallel_linear_layer on a non-linear layer".into());
7094 };
7095 // ---- attn_norm + q8_1 quantize at m=T (row-indexed == per-row) ----
7096 let anorm = layer.attn_norm.float_data();
7097 let mut xn = e.uninit(t * n_embd)?;
7098 e.rms_norm(x, anorm, &mut xn, n_embd, t, eps)?;
7099 let (hq, hd) = e.quantize_q8_1(&xn, t, n_embd)?;
7100
7101 let geometry = la.geometry;
7102 let d_state = geometry.key_head_dim as usize;
7103 let num_k = geometry.key_heads as usize;
7104 let num_v = geometry.value_heads as usize;
7105 let d_conv = geometry.conv_kernel as usize;
7106 let key_dim = d_state * num_k;
7107 let value_dim = geometry.value_head_dim as usize * num_v;
7108 let conv_dim = key_dim * 2 + value_dim;
7109 let gdn_scale = 1.0 / (d_state as f32).sqrt();
7110
7111 // ---- batched projections: one weight read for all T rows ----
7112 // GROUP-4 twin (trunk-kernels slice C): the whole 4-tuple in ONE launch, bit-identical
7113 // per (tensor, token, row) to the four singles; refused (layout/tier) or
7114 // MEMRA_TK_GDN_GROUP=0 -> the singles chain byte-for-byte.
7115 let (qkv_mixed, z, beta_raw, alpha) = match e.matmul_decode_exact_group4_pre(
7116 [&la.wqkv, &la.wqkv_gate, &la.ssm_beta, &la.ssm_alpha],
7117 &hq,
7118 &hd,
7119 t,
7120 )? {
7121 Some(mut g4) => {
7122 let alpha = g4.pop().unwrap();
7123 let beta_raw = g4.pop().unwrap();
7124 let z = g4.pop().unwrap();
7125 let qkv_mixed = g4.pop().unwrap();
7126 (qkv_mixed, z, beta_raw, alpha)
7127 }
7128 None => (
7129 e.matmul_pre(&la.wqkv, &hq, &hd, &xn, t)?,
7130 e.matmul_pre(&la.wqkv_gate, &hq, &hd, &xn, t)?,
7131 e.matmul_pre(&la.ssm_beta, &hq, &hd, &xn, t)?,
7132 e.matmul_pre(&la.ssm_alpha, &hq, &hd, &xn, t)?,
7133 ),
7134 };
7135 let beta_w = la.ssm_beta.out_features();
7136 let alpha_w = la.ssm_alpha.out_features();
7137 let qkv_w = la.wqkv.out_features();
7138
7139 // ---- per-row state chain through the b_n=1 serving kernels ----
7140 // 6-entry alternating pointer table expresses the ping-pong without a rebuild per
7141 // row: even rows scan s0 -> s1, odd rows s1 -> s0.
7142 let table_local: Option<CudaSlice<u64>> = match table_src {
7143 Some(_) => None,
7144 None => {
7145 let rl = cache.recur[il].as_ref().unwrap();
7146 let s = &e.gpu.stream();
7147 let (pc, _g0) = rl.conv_state.device_ptr(s);
7148 let (p0, _g1) = rl.ssm_state.device_ptr(s);
7149 let (p1, _g2) = rl.ssm_state_alt.device_ptr(s);
7150 Some(e.htod_u64(&[pc, p0, p1, pc, p1, p0])?)
7151 }
7152 };
7153 let (table, toff): (&CudaSlice<u64>, usize) = match table_src {
7154 Some((tb, off)) => (tb, off),
7155 None => (table_local.as_ref().unwrap(), 0),
7156 };
7157 let mut o_all = e.uninit(t * value_dim)?;
7158 let mut col_states: Option<Vec<(CudaSlice<f32>, CudaSlice<f32>)>> =
7159 if ckpt.is_some() && stash.is_none() && t >= 2 {
7160 Some(Vec::with_capacity(t - 1))
7161 } else {
7162 None
7163 };
7164 let mut stash = stash;
7165 // Per-row scratch reused across rows (uninit is cheap but not free at
7166 // 48 layers x T rows); row inputs/outputs pass as VIEWS into the packed
7167 // [T, ...] buffers — zero arithmetic-free copies in this loop.
7168 let mut conv_out = e.uninit(conv_dim)?;
7169 let mut q_l2 = e.uninit(value_dim)?;
7170 let mut k_l2 = e.uninit(value_dim)?;
7171 let mut v_gd = e.uninit(value_dim)?;
7172 let mut beta_b = e.uninit(num_v)?;
7173 let mut g_log = e.uninit(num_v)?;
7174 for r in 0..t {
7175 let base = toff + if r % 2 == 0 { 0 } else { 3 };
7176 let conv_view = table.slice(base..base + 1);
7177 let in_view = table.slice(base + 1..base + 2);
7178 let out_view = table.slice(base + 2..base + 3);
7179 e.ssm_conv1d_fused_decode_b_view(
7180 &qkv_mixed.slice(r * qkv_w..(r + 1) * qkv_w),
7181 &conv_view,
7182 la.ssm_conv1d.float_data(),
7183 &mut conv_out,
7184 conv_dim,
7185 d_conv,
7186 1,
7187 )?;
7188 e.gdn_prep_decode_b_view(
7189 &conv_out,
7190 &beta_raw.slice(r * beta_w..(r + 1) * beta_w),
7191 &alpha.slice(r * alpha_w..(r + 1) * alpha_w),
7192 la.ssm_dt.float_data(),
7193 la.ssm_a.float_data(),
7194 &mut q_l2,
7195 &mut k_l2,
7196 &mut v_gd,
7197 &mut beta_b,
7198 &mut g_log,
7199 d_state,
7200 num_v,
7201 num_k,
7202 key_dim,
7203 eps,
7204 conv_dim,
7205 1,
7206 )?;
7207 let mut o_row = o_all.slice_mut(r * value_dim..(r + 1) * value_dim);
7208 e.gdn_scan_s128_batched_view(
7209 &q_l2, &k_l2, &v_gd, &g_log, &beta_b, &in_view, &out_view, &mut o_row, num_v, 1,
7210 gdn_scale,
7211 )?;
7212 if r + 1 < t {
7213 // Row r's out buffer: even rows write s1 (the alt handle — no swaps ran),
7214 // odd rows write s0 — the same physical state the legacy post-swap
7215 // canonical clone read.
7216 let rl = cache.recur[il]
7217 .as_ref()
7218 .ok_or("qwen35 linear verify layer has no recurrent state")?;
7219 let ssm_src = if r % 2 == 0 {
7220 &rl.ssm_state_alt
7221 } else {
7222 &rl.ssm_state
7223 };
7224 match stash.as_mut() {
7225 Some((conv_slab, ssm_slab)) => {
7226 // BOTH stash reads go through the pointer table at run time: the
7227 // ssm handles ping-pong between rounds, and the ctx (with its
7228 // captured graphs) outlives the Cache — a fresh generation's
7229 // conv/ssm buffers land at new addresses that only the per-round
7230 // table refresh knows. A baked direct copy would read freed
7231 // memory (parity was the slice-3 smoke divergence; cache
7232 // lifetime is the cross-generation twin).
7233 e.copy_indirect_src_f32(
7234 &conv_view,
7235 conv_slab,
7236 r * conv_dim * (d_conv - 1),
7237 conv_dim * (d_conv - 1),
7238 )?;
7239 // The ssm handles PING-PONG between rounds: a captured direct
7240 // copy would bake the capture-time physical buffer and read the
7241 // wrong parity after any odd-vt round (the slice-3 smoke
7242 // divergence). Read the src address from row r's OUT table
7243 // entry at run time — the same entry the scan just wrote.
7244 e.copy_indirect_src_f32(
7245 &out_view,
7246 ssm_slab,
7247 r * d_state * d_state * num_v,
7248 d_state * d_state * num_v,
7249 )?;
7250 }
7251 None => {
7252 if let Some(states) = col_states.as_mut() {
7253 states.push((e.clone_dtod(&rl.conv_state)?, e.clone_dtod(ssm_src)?));
7254 }
7255 }
7256 }
7257 }
7258 }
7259 // ONE end-of-body parity swap (t odd) — the legacy loop swapped per row; the net
7260 // handle motion is identical and the device sequence never read the handles.
7261 if t % 2 == 1 {
7262 let rl = cache.recur[il].as_mut().unwrap();
7263 std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
7264 }
7265 if let (Some(checkpoint), Some(states)) = (ckpt, col_states) {
7266 checkpoint.cols[il] = Some(states);
7267 }
7268
7269 // ---- batched gated norm + out-projection at m=T ----
7270 let mixed = if e.uses_q8_1_fast(&la.ssm_out) {
7271 let (gq, gd) = e.gated_rmsnorm_q8_1(
7272 &o_all,
7273 la.ssm_norm.float_data(),
7274 &z,
7275 d_state,
7276 t * num_v,
7277 eps,
7278 )?;
7279 let g0 = e.zeros(0)?;
7280 e.matmul_pre(&la.ssm_out, &gq, &gd, &g0, t)?
7281 } else {
7282 let mut gn = e.uninit(t * value_dim)?;
7283 e.gated_rmsnorm(
7284 &o_all,
7285 la.ssm_norm.float_data(),
7286 &z,
7287 &mut gn,
7288 d_state,
7289 t * num_v,
7290 eps,
7291 )?;
7292 e.matmul(&la.ssm_out, &gn, t)?
7293 };
7294
7295 // ---- residual add + post_attn_norm + FFN at m=T (serving dispatch verbatim) ----
7296 let pnorm = layer.post_attn_norm.float_data();
7297 let mut x1 = e.uninit(t * n_embd)?;
7298 let mut zn = e.uninit(t * n_embd)?;
7299 e.add_rms_norm(x, &mixed, pnorm, &mut x1, &mut zn, n_embd, t, eps)?;
7300 let ffn_out = match &layer.ffn {
7301 crate::hybrid::Ffn::Dense {
7302 ffn_gate,
7303 ffn_up,
7304 ffn_down,
7305 } => {
7306 assert!(
7307 self.cfg.m3.is_none(),
7308 "qwen35 t-parallel verify: M3 swigluoai FFN not yet batched"
7309 );
7310 self.qwen35_tparallel_dense_ffn(e, ffn_gate, ffn_up, ffn_down, &zn, t, n_embd)?
7311 }
7312 crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il_zq8(e, m, &zn, None, t, il as u16)?,
7313 };
7314 let mut x2 = e.uninit(t * n_embd)?;
7315 e.add(&x1, &ffn_out, &mut x2, t * n_embd)?;
7316 // dspark drafter tap (no-op when no sink armed): post-layer residual verify rows
7317 self.dflash_tap(e, cache, il, &x2, t)?;
7318 Ok(x2)
7319 }
7320
7321 /// PP-N STAGE SUBGRAPH of the verify trunk: layers `[lo, hi)` of `decode_step_t_core_stream`'s
7322 /// walk, verbatim. Enters with a MATERIALIZED `[T, n_embd]` residual (no pending fusion pair
7323 /// carried in from outside the range) and exits with the range's final residual materialized
7324 /// (the trailing add executed) — exactly the `decode_layers_eager(lo, hi)` contract, T rows
7325 /// instead of one.
7326 ///
7327 /// EXTRACTED (lane/pp2-spec 2026-08-06) rather than duplicated: `decode_step_t_core_stream` IS
7328 /// the single funnel every verify forward reaches, and its per-layer dispatch MIRRORING (norm
7329 /// fusion per layer, the t>=3/spec_m2 batched-linear window, the fused-q8 FFN chain, the
7330 /// decode-exact projections) is what makes verify bit-identical to eager decode. A second copy
7331 /// for the split arm is how those mirrors drift apart on the next lever. The unsplit body now
7332 /// calls this with `(0, n_layers)`, so the whole-trunk path and every stage range run the SAME
7333 /// code — there is no "split version" of the verify math.
7334 ///
7335 /// Bit-identity of a cut rests on the same kernel-check-pinned identity the eager arm's cut
7336 /// does — `add_rms_norm_q8_1 == add then rms_norm_q8_1` at nrows=T — because the ONLY thing a
7337 /// fence changes is that the cross-layer fusion carry breaks at `hi-1` and is re-materialized
7338 /// as an explicit `add`. `decode-batch-gate --mode ppspec` verifies end-to-end on real weights.
7339 #[allow(clippy::too_many_arguments)]
7340 fn verify_layers(
7341 &self,
7342 e: &Engine,
7343 mut x: CudaSlice<f32>,
7344 lo: usize,
7345 hi: usize,
7346 pos_d: &CudaSlice<i32>,
7347 pos0: usize,
7348 t: usize,
7349 cache: &mut Cache,
7350 mut ckpt: Option<&mut VerifyCkpt>,
7351 stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
7352 graphs: Option<&mut DsparkVerifyGraphs>,
7353 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7354 if self.sliding_gated_moe_batch_program() {
7355 if stream.is_some() {
7356 return Err(
7357 "step35 has no ROUND-STREAM verify arm (the device-counter _dc twins \
7358 cannot express the SWA offset KV view)"
7359 .into(),
7360 );
7361 }
7362 return self.step35_verify_batch_layers(e, x, lo, hi, pos0, t, cache);
7363 }
7364 if self.batched_serving_numeric_class() {
7365 return self.qwen35_verify_batch_layers(
7366 e,
7367 x,
7368 lo,
7369 hi,
7370 pos0,
7371 t,
7372 cache,
7373 ckpt.take(),
7374 stream,
7375 graphs,
7376 );
7377 }
7378 let n_embd = self.cfg.n_embd as usize;
7379 let eps = self.cfg.rms_eps;
7380 // CROSS-LAYER ADD+NORM FUSION (lane/vt-fixes fix 2, mirroring decode_step_h's
7381 // launch-arc form): layer il's post-FFN residual add (x2 = x1 + ffn_out) and layer
7382 // il+1's attn_norm(+quantize) are consecutive row-wise ops — ONE add_rms_norm_q8_1
7383 // launch at nrows=t does all three (bit-identity pinned by the T-row kernel-check
7384 // arms). Carry the un-added (x1, ffn_out) pair; the fused launch materializes x2 (the
7385 // residual the next layer needs) as its `res` output. Falls back to the separate add
7386 // when the next layer is off the fused-q8 path.
7387 let mut pending: Option<(CudaSlice<f32>, CudaSlice<f32>)> = None;
7388 for il in lo..hi {
7389 let layer = &self.layers[il];
7390 // DISPATCH-MIRRORED attn-input RMSNorm (FP-order lesson #8): eager decode fuses the
7391 // 1024-thread rms_norm_q8_1 ONLY when every mixer projection is q8_1-fast; layers with
7392 // Float projections (ssm_beta/ssm_alpha on layers 1/2/4 of the 9B NVFP4 GGUF) take the
7393 // UNFUSED 256-thread rms_norm. The verify norm must mirror that PER-LAYER choice —
7394 // blockDim changes the sum-of-squares reduce order, and the ULP shift amplifies through
7395 // the GDN recurrence into argmax flips (measured: 9B text prompt, 1 ULP at layer 2 ->
7396 // 2.3e-1 logit maxdiff at the head -> K=1..8 divergence at a 0.03-margin token).
7397 let mixer_fast = self.mixer_in_q8_1_fast(e, &layer.mixer);
7398 let norm_fused = std::env::var("MEMRA_NO_FUSE_NORMQ").is_err() && mixer_fast;
7399 // BATCHED EPILOGUE RE-FUSE (lane/vt-fixes fix 2, 2026-08-03): when the norm is
7400 // dispatch-fused AND every consumer of `h` reads only its q8_1 form (Full mixer:
7401 // projections only; Linear mixer: the batched arm — the per-column fallback needs
7402 // f32 h), emit the attn-input norm DIRECTLY as q8_1 via `rms_norm_q8_1` at nrows=t
7403 // (row-indexed kernel — the T-row launch is the per-row m=1 program, kernel-check
7404 // pins bit-identity vs rms_norm_decode -> quantize_q8_1). Kills the standalone
7405 // quantize launch(es) + the f32 h HBM round-trip that decode never pays.
7406 // step35 (Full mixer) is the third case that needs f32 `h`: its verify arm is a
7407 // per-ROW replay of the eager decode mixer, whose `pre_q` contract is a single row —
7408 // a T-row q8_1 pair cannot be handed to it, and re-deriving per-row q8_1 from the f32
7409 // rows is exactly the dispatch being mirrored. Keep step35 on the unfused arm.
7410 let lin_q8_only = match &layer.mixer {
7411 Mixer::Linear(la) => {
7412 (t >= 3 || (t == 2 && spec_m2())) && e.uses_q8_1_fast(&la.ssm_out)
7413 }
7414 Mixer::Full(_) if self.sliding_gated_moe_batch_program() => false,
7415 _ => true,
7416 };
7417 // NOTE decode.rs's take()-first lesson: take the pending pair BEFORE branching so
7418 // a non-fused layer still performs the residual add.
7419 let taken = pending.take();
7420 let (h, h_q8) = if norm_fused && lin_q8_only {
7421 let pair = match taken {
7422 // fused add + attn_norm + q8_1: ONE launch resolves the carried residual
7423 // AND emits this layer's mixer input pre-quantized. res -> x2 (= new x).
7424 Some((x1p, f1p)) => {
7425 let mut x2 = vbuf(e, t * n_embd)?; // fully written (res output)
7426 let p = e.add_rms_norm_q8_1(
7427 &x1p,
7428 &f1p,
7429 layer.attn_norm.float_data(),
7430 &mut x2,
7431 n_embd,
7432 t,
7433 eps,
7434 )?;
7435 x = x2;
7436 p
7437 }
7438 None => e.rms_norm_q8_1(&x, layer.attn_norm.float_data(), n_embd, t, eps)?,
7439 };
7440 (e.zeros(0)?, Some(pair)) // h unused on this path (q8-only consumers)
7441 } else {
7442 if let Some((x1p, f1p)) = taken {
7443 let mut x2 = vbuf(e, t * n_embd)?; // fully written by add
7444 e.add(&x1p, &f1p, &mut x2, t * n_embd)?;
7445 x = x2;
7446 }
7447 let mut h = vbuf(e, t * n_embd)?; // fully written by either rms_norm arm
7448 if norm_fused {
7449 e.rms_norm_decode(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
7450 } else {
7451 e.rms_norm(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
7452 }
7453 (h, None)
7454 };
7455 let h_q8_ref = h_q8.as_ref().map(|(q, d)| (q, d));
7456
7457 let mixed = match &layer.mixer {
7458 Mixer::Full(fa) => self.full_attn_verify(
7459 e,
7460 fa,
7461 &h,
7462 h_q8_ref,
7463 pos_d,
7464 t,
7465 cache,
7466 il,
7467 stream.map(|(_, c)| c),
7468 )?,
7469 Mixer::Mla(_) => crate::hybrid::mla_path_unimplemented("speculative verify"),
7470 Mixer::Kda(_) => crate::hybrid::kda_path_unimplemented("speculative verify"),
7471 Mixer::Linear(la) => {
7472 // BATCHED linear verify (2026-07-03, the MTP-profit lever): one T-token pass —
7473 // batched projections (weight read ONCE, hits the m=2-4 weight-resident matvec),
7474 // carried-state conv (ssm_conv1d_tm_state), GDN prep on the prefill kernels, and
7475 // ONE gdn_scan whose internal sequential t-loop is the SAME recurrence as T
7476 // chained T=1 steps (bit-identical). Falls back to the sequential per-column
7477 // chain when T < d_conv-1 (conv ring update needs T >= pad) — or when ANY
7478 // projection is off the q8_1 fast path: matmul_decode_exact would route a Float
7479 // tensor to cuBLAS at m=t (different FP accumulation than eager's per-token
7480 // GEMV), so mixed-dtype layers stay on the eager-identical per-column chain.
7481 // MEMRA_SPEC_M2 (lane/spec-m2): the t==2 batch rides the same arm — the conv
7482 // wrapper handles t<pad with a pure-copy ring rebuild; see spec_m2() header.
7483 if (t >= 3 || (t == 2 && spec_m2()))
7484 && mixer_fast
7485 && e.uses_q8_1_fast(&la.ssm_out)
7486 {
7487 let want = ckpt.is_some();
7488 let (out, stash) =
7489 self.linear_attn_verify_t(e, la, &h, h_q8_ref, t, cache, il, want)?;
7490 if let (Some(ck), Some(st)) = (ckpt.as_deref_mut(), stash) {
7491 ck.gdn[il] = Some(st);
7492 }
7493 out
7494 } else {
7495 let mut out = vbuf(e, t * n_embd)?; // every col written by copy_into
7496 let mut col_states: Option<Vec<(CudaSlice<f32>, CudaSlice<f32>)>> =
7497 if ckpt.is_some() && t >= 2 {
7498 Some(Vec::with_capacity(t - 1))
7499 } else {
7500 None
7501 };
7502 for col in 0..t {
7503 let mut h_col = vbuf(e, n_embd)?; // fully written by copy_view_into
7504 let src = h.slice(col * n_embd..(col + 1) * n_embd);
7505 e.copy_view_into(&mut h_col, 0, &src, n_embd)?;
7506 let m_col = self.linear_attn_decode(e, la, &h_col, cache, il)?;
7507 e.copy_into(&mut out, col * n_embd, &m_col, n_embd)?;
7508 // REPLAY-FREE ckpt: clone the chain's ACTUAL state after this column
7509 // (pure dtod — cannot change any computed value). Last column skipped:
7510 // rebuild targets are j <= t-1 columns.
7511 if let Some(cs) = col_states.as_mut()
7512 && col + 1 < t
7513 {
7514 let rl = cache.recur[il].as_ref().unwrap();
7515 cs.push((
7516 e.clone_dtod(&rl.conv_state)?,
7517 e.clone_dtod(&rl.ssm_state)?,
7518 ));
7519 }
7520 }
7521 if let (Some(ck), Some(cs)) = (ckpt.as_deref_mut(), col_states) {
7522 // ReplaySSM-assessment instrumentation (2026-07-30): the
7523 // per-column clones are the only true state snapshots left in
7524 // the verify (the batched path stashes INPUTS and replays).
7525 if std::env::var("MEMRA_SPEC_STATS").as_deref() == Ok("1") {
7526 static ONCE: std::sync::Once = std::sync::Once::new();
7527 let bytes: usize =
7528 cs.iter().map(|(c, s)| (c.len() + s.len()) * 4).sum();
7529 ONCE.call_once(|| eprintln!(
7530 "[verify-ckpt] per-column layer il={il}: {} clones, {:.2} MB/layer/round",
7531 cs.len(), bytes as f64 / 1e6));
7532 }
7533 ck.cols[il] = Some(cs);
7534 }
7535 out
7536 }
7537 }
7538 };
7539 if spec_nan_scan_level() >= 2 {
7540 let mixed_width = mixed.len() / t;
7541 nan_scan_rows(
7542 e,
7543 &mixed,
7544 t,
7545 mixed_width,
7546 &format!("verify layer {il} batched ATTN out pos0={pos0}"),
7547 )?;
7548 }
7549
7550 // DISPATCH-MIRRORED post-attn norm: eager residual_norm_ffn fuses add+norm+quant
7551 // (1024-thread add_rms_norm_q8_1) only for Dense FFNs whose gate+up are q8_1-fast;
7552 // otherwise (and for MoE) it runs the 256-thread fused add_rms_norm. Mirror per layer.
7553 let ffn_fuse = match &layer.ffn {
7554 crate::hybrid::Ffn::Dense {
7555 ffn_gate, ffn_up, ..
7556 } => {
7557 std::env::var("MEMRA_NO_FUSE_NORMQ").is_err()
7558 && e.uses_q8_1_fast(ffn_gate)
7559 && e.uses_q8_1_fast(ffn_up)
7560 }
7561 crate::hybrid::Ffn::Moe(_) => false,
7562 };
7563 // BATCHED EPILOGUE RE-FUSE (lane/vt-fixes fix 2): on the ffn_fuse path (Dense,
7564 // gate+up q8_1-fast, non-M3) the FFN input is emitted DIRECTLY as q8_1 by ONE
7565 // add_rms_norm_q8_1 launch at nrows=t (row-indexed kernel: the T-row launch is the
7566 // per-row m=1 program; kernel-check pins bit-identity vs the unfused
7567 // add_f32 -> rms_norm_decode -> quantize_q8_1 chain at T=2/4/5/8) — replacing the
7568 // add + rms_norm_decode launches AND the dual/singles' internal re-quantize.
7569 // M3's swigluoai must keep the f32 chain (the fused SwiGLU epilogue encodes plain
7570 // SiLU), mirroring residual_norm_ffn's m3 guard on the decode path.
7571 // step35: same guard per LAYER. A dense FFN's clamp is the SHEXP array (upstream's
7572 // one build_ffn serves dense + shared expert, llama-graph.cpp:1751), and verify MUST
7573 // mirror decode's dispatch or spec self-consistency fails.
7574 let dense_lim = self.cfg.clamp_shexp_at(il as u32);
7575 let fuse_q8 = ffn_fuse && self.cfg.m3.is_none() && dense_lim.is_none();
7576 let mut x1 = vbuf(e, t * n_embd)?; // fully written by add / add_rms_norm*
7577 let mut z = e.zeros(0)?; // replaced below on the unfused arms
7578 let z_q8 = if fuse_q8 {
7579 Some(e.add_rms_norm_q8_1(
7580 &x,
7581 &mixed,
7582 layer.post_attn_norm.float_data(),
7583 &mut x1,
7584 n_embd,
7585 t,
7586 eps,
7587 )?)
7588 } else {
7589 let mut zf = vbuf(e, t * n_embd)?; // fully written by rms_norm_decode / add_rms_norm
7590 if ffn_fuse {
7591 e.add(&x, &mixed, &mut x1, t * n_embd)?;
7592 e.rms_norm_decode(
7593 &x1,
7594 layer.post_attn_norm.float_data(),
7595 &mut zf,
7596 n_embd,
7597 t,
7598 eps,
7599 )?;
7600 } else {
7601 e.add_rms_norm(
7602 &x,
7603 &mixed,
7604 layer.post_attn_norm.float_data(),
7605 &mut x1,
7606 &mut zf,
7607 n_embd,
7608 t,
7609 eps,
7610 )?;
7611 }
7612 z = zf;
7613 None
7614 };
7615 if spec_nan_scan_level() >= 2 && !z.is_empty() {
7616 nan_scan_rows(
7617 e,
7618 &z,
7619 t,
7620 n_embd,
7621 &format!("verify layer {il} post-attn norm z pos0={pos0}"),
7622 )?;
7623 }
7624 // DECODE-EXACT FFN projections: force MMVQ for gate/up/down at any T to match the
7625 // T=1 decode FP accumulation order. At T>=5 the generic matmul/matmul_pre falls to dp4a
7626 // (128-thread, different FP sum order). At T=2-4 the batched MMVQ is already bit-identical.
7627 let ffn_out = match &layer.ffn {
7628 crate::hybrid::Ffn::Dense {
7629 ffn_gate,
7630 ffn_up,
7631 ffn_down,
7632 } => {
7633 let n_ff = ffn_gate.out_features();
7634 if let Some((zq, zd)) = z_q8.as_ref() {
7635 // FUSED CHAIN (fix 2): pre-quantized z feeds the projections; the SwiGLU
7636 // epilogue emits act pre-quantized for ffn_down (silu_mul_scaled_q8_1,
7637 // bit-identical to silu_mul + quantize — kernel-check-pinned) with the
7638 // NVFP4 macro-scales folded (deferred-scale dual: y*s inline == the
7639 // scale_inplace store, value-exact) — the exact m=1 decode epilogue
7640 // structure at nrows=t.
7641 let pair = e
7642 .matmul_decode_exact_dual_pre(ffn_gate, ffn_up, zq, zd, t)?
7643 .map(|((g, gs), (u, us))| (g, gs, u, us));
7644 let (gate, gs, up, us) = match pair {
7645 Some(x4) => x4,
7646 None => (
7647 e.matmul_decode_exact_pre(ffn_gate, zq, zd, t)?,
7648 1.0, // scale already applied inside _pre
7649 e.matmul_decode_exact_pre(ffn_up, zq, zd, t)?,
7650 1.0,
7651 ),
7652 };
7653 if e.uses_q8_1_fast(ffn_down) {
7654 let (aq, ad) = e.silu_mul_scaled_q8_1(&gate, &up, gs, us, t * n_ff)?;
7655 e.matmul_decode_exact_pre(ffn_down, &aq, &ad, t)?
7656 } else {
7657 let mut act = vbuf(e, t * n_ff)?;
7658 e.silu_mul_scaled(&gate, &up, gs, us, &mut act, t * n_ff)?;
7659 e.matmul_decode_exact(ffn_down, &act, t)?
7660 }
7661 } else {
7662 // UNFUSED (pre-fix) chain — MoE-adjacent/M3/off-fast layers, unchanged.
7663 // DUAL gate+up batched twin (lane/verify-economics, 2026-08-02): one launch
7664 // for the pair at t=2..8 — bit-identical per (tensor,token,row) to the two
7665 // singles (kernel-check pins bitwise; MEMRA_SPEC_DUAL_T=0 reverts). None
7666 // (non-NVFP4 / t outside the tier / seam off) -> the two singles, unchanged.
7667 let (gate, up) =
7668 match e.matmul_decode_exact_dual(ffn_gate, ffn_up, &z, t)? {
7669 Some(pair) => pair,
7670 None => (
7671 e.matmul_decode_exact(ffn_gate, &z, t)?,
7672 e.matmul_decode_exact(ffn_up, &z, t)?,
7673 ),
7674 };
7675 let mut act = vbuf(e, t * n_ff)?; // fully written by ffn_act_lim
7676 Self::ffn_act_lim(
7677 e,
7678 &self.cfg,
7679 &gate,
7680 &up,
7681 1.0,
7682 1.0,
7683 dense_lim,
7684 &mut act,
7685 t * n_ff,
7686 )?;
7687 e.matmul_decode_exact(ffn_down, &act, t)?
7688 }
7689 }
7690 crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il(e, m, &z, t, il as u16)?,
7691 };
7692 if spec_nan_scan_level() >= 2 {
7693 nan_scan_rows(
7694 e,
7695 &ffn_out,
7696 t,
7697 n_embd,
7698 &format!("verify layer {il} batched FFN out pos0={pos0}"),
7699 )?;
7700 }
7701 if spec_nan_scan() {
7702 let mut residual = vbuf(e, t * n_embd)?;
7703 e.add(&x1, &ffn_out, &mut residual, t * n_embd)?;
7704 nan_scan_rows(
7705 e,
7706 &residual,
7707 t,
7708 n_embd,
7709 &format!("verify layer {il} residual pos0={pos0}"),
7710 )?;
7711 }
7712 // CROSS-LAYER fusion: defer this layer's post-FFN residual add — the next layer's
7713 // fused-q8 attn norm folds it in (add_rms_norm_q8_1 == add; rms_norm; quantize,
7714 // kernel-check-pinned at nrows=T). Non-fused next layers add explicitly above.
7715 pending = Some((x1, ffn_out));
7716 }
7717 // RANGE's final add (no next norm INSIDE the range to fuse with; for the
7718 // whole-trunk call that is the last layer, whose next norm is output_norm — f32-out).
7719 if let Some((x1p, f1p)) = pending.take() {
7720 let mut x2 = vbuf(e, t * n_embd)?; // fully written by add
7721 e.add(&x1p, &f1p, &mut x2, t * n_embd)?;
7722 x = x2;
7723 }
7724 Ok(x)
7725 }
7726 /// BATCHED linear-attn verify (T=K+1): the whole layer in ~10 launches instead of T x the
7727 /// T=1 decode chain (T x ~12 launches + T weight reads of the four projections). The GDN
7728 /// recurrence itself is inherently sequential — gdn_scan_s128 runs its internal t-loop with
7729 /// the SAME per-token math as chained T=1 calls (bit-identical state evolution); everything
7730 /// around it (projections, conv, prep, gated norm, out-proj) batches. Advances conv ring +
7731 /// ssm state exactly like T sequential decode steps.
7732 /// `want_stash`: additionally RETAIN the gdn-scan inputs (pure buffer keep-alives, zero extra
7733 /// kernels) so a partial accept can rebuild the state after any column prefix (REPLAY-FREE).
7734 #[allow(clippy::too_many_arguments)]
7735 fn linear_attn_verify_t(
7736 &self,
7737 e: &Engine,
7738 la: &LinearAttnLayer,
7739 h: &CudaSlice<f32>,
7740 h_q8: Option<(&CudaSlice<i8>, &CudaSlice<f32>)>,
7741 t: usize,
7742 cache: &mut Cache,
7743 il: usize,
7744 want_stash: bool,
7745 ) -> Result<(CudaSlice<f32>, Option<GdnStash>), Box<dyn std::error::Error>> {
7746 let cfg = &self.cfg;
7747 let geometry = la.geometry;
7748 let d_state = geometry.key_head_dim as usize;
7749 let num_k = geometry.key_heads as usize;
7750 let num_v = geometry.value_heads as usize;
7751 let d_conv = geometry.conv_kernel as usize;
7752 let key_dim = d_state * num_k;
7753 let conv_dim = key_dim * 2 + geometry.value_head_dim as usize * num_v;
7754 let eps = cfg.rms_eps;
7755 let scale = 1.0 / (d_state as f32).sqrt();
7756
7757 // DECODE-EXACT projections: matmul_decode_exact forces the MMVQ (warp-per-row, 32-thread)
7758 // accumulation order for EVERY m, matching the T=1 decode path bit-for-bit. The generic
7759 // `matmul` at m>=5 falls to dp4a (128-thread, two-level reduce) which has a different FP
7760 // sum order — ULP differences propagate through gdn_scan and flip argmax on the 27B.
7761 // Q8 TRUNK-FUSION at T=1 (35B: wqkv+wqkv_gate both Q8_0): one fused2 launch, bit-identical
7762 // per (tensor,row) to the two m=1 MMVQ dispatches below — decode-exact contract holds.
7763 // VERIFY-TIER TRUNK FUSION (MEMRA_SPEC_FUSED_T, t=2-4): quantize h ONCE for every
7764 // fused-eligible same-input Q8_0 pair of this layer (35B wqkv+wqkv_gate; 9B
7765 // ssm_beta+ssm_alpha) — each fused2 batched launch then replaces two decode-exact
7766 // calls (each of which re-quantizes the same h + runs its own _b2/_b4 launch).
7767 // Bit-identical per (tensor,token,row) — see spec_fused_t().
7768 // BATCHED EPILOGUE RE-FUSE (lane/vt-fixes fix 2): `h_q8` = the attn-input norm emitted
7769 // directly as q8_1 by the caller's fused rms_norm_q8_1 (bit-identical to the unfused
7770 // chain, kernel-check-pinned). When present it REPLACES the standalone quantize below
7771 // and feeds every projection; the caller guaranteed all four input projections are
7772 // q8_1-fast. When absent, the old shared-quantize (fused-t window) stands.
7773 let h_q8_t = if h_q8.is_none()
7774 && spec_fused_t()
7775 && (2..=4).contains(&t)
7776 && ((e.uses_q8_1_fast(&la.wqkv) && e.uses_q8_1_fast(&la.wqkv_gate))
7777 || (e.uses_q8_1_fast(&la.ssm_beta) && e.uses_q8_1_fast(&la.ssm_alpha)))
7778 {
7779 Some(e.quantize_q8_1(h, t, cfg.n_embd as usize)?)
7780 } else {
7781 None
7782 };
7783 // one view: the caller's fused-norm q8 or this fn's own shared quantize.
7784 let hq8_any: Option<(&CudaSlice<i8>, &CudaSlice<f32>)> =
7785 h_q8.or(h_q8_t.as_ref().map(|(q, d)| (q, d)));
7786 let (qkv_mixed, z) = {
7787 let mut fused = None;
7788 if t == 1 && e.uses_q8_1_fast(&la.wqkv) && e.uses_q8_1_fast(&la.wqkv_gate) {
7789 let (hq, hd) = e.quantize_q8_1(h, 1, cfg.n_embd as usize)?;
7790 fused = e.matmul_q8_fused2(&la.wqkv, &la.wqkv_gate, &hq, &hd)?;
7791 } else if let Some((hq, hd)) = hq8_any
7792 && spec_fused_t()
7793 && (2..=4).contains(&t)
7794 {
7795 fused = e.matmul_q8_fused2_t(&la.wqkv, &la.wqkv_gate, hq, hd, t)?;
7796 }
7797 match (fused, hq8_any) {
7798 (Some(pair), _) => pair,
7799 (None, Some((hq, hd))) if h_q8.is_some() => (
7800 e.matmul_decode_exact_pre(&la.wqkv, hq, hd, t)?,
7801 e.matmul_decode_exact_pre(&la.wqkv_gate, hq, hd, t)?,
7802 ),
7803 (None, _) => (
7804 e.matmul_decode_exact(&la.wqkv, h, t)?,
7805 e.matmul_decode_exact(&la.wqkv_gate, h, t)?,
7806 ),
7807 }
7808 };
7809 // beta+alpha DUAL at T=1 (75% of p3 rounds run T=1 verify — p-min chain cuts): the dual
7810 // mr2 kernel is bit-identical per element to the m=1 MMVQ matmul_decode_exact dispatches
7811 // (same warp-per-row body, blockIdx.y picks the weight), so the decode-exact contract
7812 // holds; the run-spec battery is the arbiter. T>1 keeps the per-tensor decode-exact path.
7813 let (beta_raw, alpha) = if t == 1 {
7814 let (hq, hd) = e.quantize_q8_1(h, 1, cfg.n_embd as usize)?;
7815 match e.matmul_pre_dual_noscale(&la.ssm_beta, &la.ssm_alpha, &hq, &hd, 1)? {
7816 Some(((mut b, bs), (mut a, as_))) => {
7817 if bs != 1.0 {
7818 e.scale_inplace(&mut b, bs, la.ssm_beta.out_features())?;
7819 }
7820 if as_ != 1.0 {
7821 e.scale_inplace(&mut a, as_, la.ssm_alpha.out_features())?;
7822 }
7823 (b, a)
7824 }
7825 // Q8_0 fused2 twin (9B stores beta/alpha as Q8_0): DISPATCH-MIRRORS the eager
7826 // decode's beta_alpha closure — the fused body is qmatvec_q8_0_mmvq verbatim,
7827 // bit-identical per row (kernel-check rel=0.00e0 gate), so decode==verify holds.
7828 None => match e.matmul_q8_fused2(&la.ssm_beta, &la.ssm_alpha, &hq, &hd)? {
7829 Some((b, a)) => (b, a),
7830 None => (
7831 e.matmul_decode_exact(&la.ssm_beta, h, 1)?,
7832 e.matmul_decode_exact(&la.ssm_alpha, h, 1)?,
7833 ),
7834 },
7835 }
7836 } else {
7837 // fused-t twin (9B stores beta/alpha as Q8_0): same shared-quantize + one launch
7838 // contract as the wqkv pair above; 35B beta/alpha are Float -> None -> fallback.
7839 let mut nvfp4_fused = None;
7840 let mut q8_fused = None;
7841 if let Some((hq, hd)) = hq8_any {
7842 if t == 3 && std::env::var("MEMRA_NVFP4_AUX_DUAL").as_deref() != Ok("0") {
7843 nvfp4_fused =
7844 e.matmul_decode_exact_dual_pre(&la.ssm_beta, &la.ssm_alpha, hq, hd, t)?;
7845 if nvfp4_fused.is_some() && std::env::var("MEMRA_DEBUG").is_ok() {
7846 static ONCE: std::sync::Once = std::sync::Once::new();
7847 ONCE.call_once(|| {
7848 eprintln!("[memra] NVFP4 beta+alpha batched aux dual ENGAGED (t={t})")
7849 });
7850 }
7851 }
7852 if nvfp4_fused.is_none() && spec_fused_t() && (2..=4).contains(&t) {
7853 q8_fused = e.matmul_q8_fused2_t(&la.ssm_beta, &la.ssm_alpha, hq, hd, t)?;
7854 }
7855 }
7856 if let Some(((mut b, bs), (mut a, as_))) = nvfp4_fused {
7857 if bs != 1.0 {
7858 e.scale_inplace(&mut b, bs, t * la.ssm_beta.out_features())?;
7859 }
7860 if as_ != 1.0 {
7861 e.scale_inplace(&mut a, as_, t * la.ssm_alpha.out_features())?;
7862 }
7863 (b, a)
7864 } else if let Some(pair) = q8_fused {
7865 pair
7866 } else {
7867 match hq8_any {
7868 Some((hq, hd)) if h_q8.is_some() => (
7869 e.matmul_decode_exact_pre(&la.ssm_beta, hq, hd, t)?,
7870 e.matmul_decode_exact_pre(&la.ssm_alpha, hq, hd, t)?,
7871 ),
7872 _ => (
7873 e.matmul_decode_exact(&la.ssm_beta, h, t)?,
7874 e.matmul_decode_exact(&la.ssm_alpha, h, t)?,
7875 ),
7876 }
7877 }
7878 };
7879
7880 // conv with CARRIED state + ring roll (T >= pad rides the input-column update kernel;
7881 // T < pad — the MEMRA_SPEC_M2 t=2 arm — rolls via the pure-copy ring rebuild).
7882 let rl = cache.recur[il].as_mut().unwrap();
7883 let mut conv_out = e.uninit(conv_dim * t)?;
7884 e.ssm_conv1d_tm_state(
7885 &qkv_mixed,
7886 &mut rl.conv_state,
7887 la.ssm_conv1d.float_data(),
7888 &mut conv_out,
7889 conv_dim,
7890 t,
7891 d_conv,
7892 )?;
7893
7894 // GDN prep via the prefill kernels (repack + L2 + sigmoid + glog), T-wide.
7895 let mut q_g = e.uninit(d_state * num_v * t)?;
7896 let mut k_g = e.uninit(d_state * num_v * t)?;
7897 let mut v_g = e.uninit(d_state * num_v * t)?;
7898 e.qkv_to_gdn_repack(
7899 &conv_out, &mut q_g, &mut k_g, &mut v_g, d_state, num_v, num_k, key_dim, t,
7900 )?;
7901 let mut q_l2 = e.uninit(d_state * num_v * t)?;
7902 e.l2_norm_decode(&q_g, &mut q_l2, d_state, num_v * t, eps)?;
7903 let mut k_l2 = e.uninit(d_state * num_v * t)?;
7904 e.l2_norm_decode(&k_g, &mut k_l2, d_state, num_v * t, eps)?;
7905 let mut beta = e.uninit(t * num_v)?;
7906 e.sigmoid(&beta_raw, &mut beta, t * num_v)?;
7907 let mut g_log = e.uninit(t * num_v)?;
7908 e.gdn_glog(
7909 &alpha,
7910 la.ssm_dt.float_data(),
7911 la.ssm_a.float_data(),
7912 &mut g_log,
7913 num_v,
7914 t,
7915 )?;
7916
7917 // ONE gdn_scan over T tokens from the carried state (internal sequential loop ==
7918 // T chained T=1 steps). Ping-pong the resident buffers like eager decode.
7919 let mut o = e.uninit(d_state * num_v * t)?;
7920 {
7921 let crate::cache::RecurLayer {
7922 ssm_state,
7923 ssm_state_alt,
7924 ..
7925 } = rl;
7926 e.gdn_scan_s128(
7927 &q_l2,
7928 &k_l2,
7929 &v_g,
7930 &g_log,
7931 &beta,
7932 ssm_state,
7933 ssm_state_alt,
7934 &mut o,
7935 num_v,
7936 t,
7937 scale,
7938 )?;
7939 }
7940 std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
7941
7942 // gated RMSNorm + out projection, T-wide. FUSED-QUANTIZE ARM (lane/vt-fixes fix 2,
7943 // mirroring the T=1 decode's launch-arc form): when ssm_out rides the q8_1 fast path,
7944 // emit q8_1 straight from the gated norm at nrows=num_v*t (row-indexed kernel, the
7945 // T-wide launch is the per-row program; kernel-check pins bit-identity vs
7946 // gated_rmsnorm -> quantize_q8_1 at T=1 and T=5) and feed the decode-exact dispatch
7947 // pre-quantized — one launch replaces norm + quantize. Fallback = the f32 chain.
7948 let out = if e.uses_q8_1_fast(&la.ssm_out) {
7949 let (gq, gd) =
7950 e.gated_rmsnorm_q8_1(&o, la.ssm_norm.float_data(), &z, d_state, num_v * t, eps)?;
7951 e.matmul_decode_exact_pre(&la.ssm_out, &gq, &gd, t)?
7952 } else {
7953 let mut gn = e.uninit(d_state * num_v * t)?;
7954 e.gated_rmsnorm(
7955 &o,
7956 la.ssm_norm.float_data(),
7957 &z,
7958 &mut gn,
7959 d_state,
7960 num_v * t,
7961 eps,
7962 )?;
7963 // DECODE-EXACT out-projection: same MMVQ path as the T=1 decode (ssm_out at m>=5
7964 // would fall to dp4a with a different FP reduction order — same class of bug as
7965 // the input projs).
7966 e.matmul_decode_exact(&la.ssm_out, &gn, t)?
7967 };
7968 let stash = if want_stash {
7969 Some(GdnStash {
7970 qkv_mixed,
7971 q_l2,
7972 k_l2,
7973 v_g,
7974 g_log,
7975 beta,
7976 })
7977 } else {
7978 None
7979 };
7980 Ok((out, stash))
7981 }
7982
7983 /// REPLAY-FREE partial-accept commit (2026-07-03): make the cache state == "committed through
7984 /// the first `j` verify columns" WITHOUT the legacy rollback + duplicate trunk replay.
7985 /// - Full-attn KV: truncate both the owning-stage shadow and every TP rank to snapshot + j.
7986 /// The verify's appended rows for those columns are bit-identical to what an eager T=1
7987 /// chain writes (the decode-exact contract the verify-probe gates), so keeping them ==
7988 /// replaying them.
7989 /// - Linear layers, batched path: rebuild the conv ring by PURE COPIES (ring holds raw input
7990 /// columns) and the ssm state by a prefix re-run of the SAME gdn_scan kernel (t=j) from the
7991 /// snapshot state over the stash's identical inputs — the kernel's t-loop carries state in
7992 /// registers and writes it once at the end, so iterations 0..j-1 are independent of T:
7993 /// bit-identical to the verify's own state after j tokens == the eager chain state.
7994 /// - Linear layers, per-column path: restore the cloned actual state after column j-1.
7995 /// Caller guarantees 1 <= j <= t-1 (j==0 rounds take the legacy rollback; j==t is full accept).
7996 #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
7997 fn commit_verified_prefix(
7998 &self,
7999 e: &Engine,
8000 cache: &mut Cache,
8001 snap: &crate::cache::CacheSnapshot,
8002 ckpt: &VerifyCkpt,
8003 j: usize,
8004 ) -> Result<(), Box<dyn std::error::Error>> {
8005 // GDN geometry derives lazily inside recurrent-layer arms. Full-attention plans carry no
8006 // recurrent state and must never be forced through a synthetic SSM geometry.
8007 // Engine-bundle slice 1 (DSF-ROUNDCOST-20260820 §1.1): the per-column-arm restores
8008 // are 2 tiny D2D copies per linear layer (~96 dispatches/partial round on the q38
8009 // route). When every cols-arm layer shares uniform state sizes (single ssm cfg —
8010 // always true today), batch them into two `copy_batch_uniform_f32` launches. Bytes,
8011 // buffers and stream order are identical to the per-layer memcpy sequence; the
8012 // kernel-rebuild (gdn-stash) arm below is untouched.
8013 let mut batched_cols = false;
8014 {
8015 use cudarc::driver::DevicePtr;
8016 let s = &e.gpu.stream();
8017 let mut conv_pairs: Vec<(u64, u64)> = Vec::new();
8018 let mut ssm_pairs: Vec<(u64, u64)> = Vec::new();
8019 let (mut conv_words, mut ssm_words) = (0usize, 0usize);
8020 let mut uniform = true;
8021 for il in 0..self.layers.len() {
8022 let Some(rl) = cache.recur[il].as_ref() else {
8023 continue;
8024 };
8025 if ckpt.gdn[il].is_some() {
8026 continue; // kernel-rebuild arm restores below, per layer
8027 }
8028 let Some(cols) = &ckpt.cols[il] else {
8029 continue; // missing-ckpt error surfaces in the main loop
8030 };
8031 let (c, st) = &cols[j - 1];
8032 if conv_pairs.is_empty() {
8033 conv_words = c.len();
8034 ssm_words = st.len();
8035 } else if c.len() != conv_words || st.len() != ssm_words {
8036 uniform = false;
8037 break;
8038 }
8039 let (pc, _g0) = c.device_ptr(s);
8040 let (dc, _g1) = rl.conv_state.device_ptr(s);
8041 let (ps, _g2) = st.device_ptr(s);
8042 let (ds, _g3) = rl.ssm_state.device_ptr(s);
8043 conv_pairs.push((pc, dc));
8044 ssm_pairs.push((ps, ds));
8045 }
8046 if uniform && !conv_pairs.is_empty() {
8047 let n = conv_pairs.len();
8048 let mut t = vec![0u64; 2 * n];
8049 for (k, &(src, dst)) in conv_pairs.iter().enumerate() {
8050 t[k] = src;
8051 t[n + k] = dst;
8052 }
8053 let conv_t = e.htod_u64(&t)?;
8054 for (k, &(src, dst)) in ssm_pairs.iter().enumerate() {
8055 t[k] = src;
8056 t[n + k] = dst;
8057 }
8058 let ssm_t = e.htod_u64(&t)?;
8059 e.copy_batch_uniform_f32(&conv_t, n, conv_words)?;
8060 e.copy_batch_uniform_f32(&ssm_t, n, ssm_words)?;
8061 batched_cols = true;
8062 }
8063 }
8064 for il in 0..self.layers.len() {
8065 if let (Some(kvl), Some(saved)) = (cache.kv[il].as_mut(), snap.kv_len[il]) {
8066 kvl.len = saved + j;
8067 e.set_i32_one(&mut kvl.len_d, kvl.len as i32)?;
8068 }
8069 if let Some(rl) = cache.recur[il].as_mut() {
8070 let Mixer::Linear(linear) = &self.layers[il].mixer else {
8071 return Err(format!("recurrent cache layer {il} has no GDN plan").into());
8072 };
8073 let geometry = linear.geometry;
8074 let d_state = geometry.key_head_dim as usize;
8075 let num_k = geometry.key_heads as usize;
8076 let num_v = geometry.value_heads as usize;
8077 let d_conv = geometry.conv_kernel as usize;
8078 let conv_dim = d_state * num_k * 2 + geometry.value_head_dim as usize * num_v;
8079 let scale = 1.0 / (d_state as f32).sqrt();
8080 if let Some(st) = &ckpt.gdn[il] {
8081 let ring_old = snap.conv[il].as_ref().expect("snapshot missing conv");
8082 let state_in = snap.ssm[il].as_ref().expect("snapshot missing ssm");
8083 e.ssm_conv_ring_rebuild(
8084 &st.qkv_mixed,
8085 ring_old,
8086 &mut rl.conv_state,
8087 conv_dim,
8088 j,
8089 d_conv,
8090 )?;
8091 let mut o = e.uninit(d_state * num_v * j)?; // scan output, discarded
8092 e.gdn_scan_s128(
8093 &st.q_l2,
8094 &st.k_l2,
8095 &st.v_g,
8096 &st.g_log,
8097 &st.beta,
8098 state_in,
8099 &mut rl.ssm_state,
8100 &mut o,
8101 num_v,
8102 j,
8103 scale,
8104 )?;
8105 } else if let Some(cols) = &ckpt.cols[il] {
8106 if !batched_cols {
8107 let (c, s) = &cols[j - 1];
8108 e.copy_into(&mut rl.conv_state, 0, c, c.len())?;
8109 e.copy_into(&mut rl.ssm_state, 0, s, s.len())?;
8110 }
8111 } else {
8112 return Err(
8113 "commit_verified_prefix: verify ckpt missing for linear layer".into(),
8114 );
8115 }
8116 }
8117 }
8118 self.restore_step_tp_kv_verified_prefix(e, cache, snap, j, true)?;
8119 cache.pos = snap.pos + j;
8120 Ok(())
8121 }
8122
8123 /// ROUND-STREAM: recur restore with device-j (the _dc twins; full accept early-exits
8124 /// in-kernel). Requires the batched-linear stash on every linear layer (stream gate).
8125 #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
8126 fn commit_verified_prefix_stream(
8127 &self,
8128 e: &Engine,
8129 cache: &mut Cache,
8130 snap: &crate::cache::CacheSnapshot,
8131 ckpt: &VerifyCkpt,
8132 acc: &CudaSlice<u32>,
8133 base: usize,
8134 t_v: usize,
8135 ) -> Result<(), Box<dyn std::error::Error>> {
8136 for il in 0..self.layers.len() {
8137 if let Some(rl) = cache.recur[il].as_mut() {
8138 let Mixer::Linear(linear) = &self.layers[il].mixer else {
8139 return Err(format!("recurrent cache layer {il} has no GDN plan").into());
8140 };
8141 let geometry = linear.geometry;
8142 let d_state = geometry.key_head_dim as usize;
8143 let num_k = geometry.key_heads as usize;
8144 let num_v = geometry.value_heads as usize;
8145 let d_conv = geometry.conv_kernel as usize;
8146 let conv_dim = d_state * num_k * 2 + geometry.value_head_dim as usize * num_v;
8147 let scale = 1.0 / (d_state as f32).sqrt();
8148 let st = ckpt.gdn[il]
8149 .as_ref()
8150 .ok_or("stream restore: batched-linear stash missing")?;
8151 let ring_old = snap.conv[il].as_ref().expect("snapshot missing conv");
8152 let state_in = snap.ssm[il].as_ref().expect("snapshot missing ssm");
8153 e.ssm_conv_ring_rebuild_dc(
8154 &st.qkv_mixed,
8155 ring_old,
8156 &mut rl.conv_state,
8157 conv_dim,
8158 acc,
8159 base,
8160 t_v,
8161 d_conv,
8162 )?;
8163 let mut o = e.uninit(d_state * num_v * t_v)?;
8164 e.gdn_scan_s128_dc(
8165 &st.q_l2,
8166 &st.k_l2,
8167 &st.v_g,
8168 &st.g_log,
8169 &st.beta,
8170 state_in,
8171 &mut rl.ssm_state,
8172 &mut o,
8173 num_v,
8174 acc,
8175 base,
8176 t_v,
8177 scale,
8178 )?;
8179 }
8180 }
8181 Ok(())
8182 }
8183
8184 /// EAGLE3 aux-capturing verify forward over `tokens` (T) — mirrors `decode_step_t_h` exactly
8185 /// (same KV append, same causal verify, same recur advance) but ALSO clones the aux residual-
8186 /// stream hiddens (blocks in `aux_layers`) for TWO columns: the LAST column (always) and the
8187 /// optional `pred_col` (the EAGLE seed = bonus's predecessor). Returns
8188 /// (all_T_logits host, last_col_aux, pred_col_aux?). Used by the EAGLE3 orchestrator's commit.
8189 #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
8190 pub fn decode_step_t_aux2(
8191 &self,
8192 e: &Engine,
8193 tokens: &[u32],
8194 pos0: usize,
8195 cache: &mut Cache,
8196 aux_layers: &[usize],
8197 pred_col: Option<usize>,
8198 ) -> Result<
8199 (Vec<f32>, Vec<CudaSlice<f32>>, Option<Vec<CudaSlice<f32>>>),
8200 Box<dyn std::error::Error>,
8201 > {
8202 cache.ensure_usable("decode_step_t_aux2")?;
8203 let cfg = &self.cfg;
8204 let n_embd = cfg.n_embd as usize;
8205 let eps = cfg.rms_eps;
8206 let t = tokens.len();
8207 let pos_vec: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
8208 let pos_d = e.htod_i32(&pos_vec)?;
8209 let mut x = e.htod(&self.embd.try_gather(n_embd, tokens)?)?;
8210 let mut aux_last: Vec<CudaSlice<f32>> = Vec::with_capacity(aux_layers.len());
8211 let mut aux_pred: Vec<CudaSlice<f32>> = Vec::new();
8212 let want_pred = pred_col.is_some();
8213
8214 for (il, layer) in self.layers.iter().enumerate() {
8215 // DISPATCH-MIRRORED norms (FP-order lesson #8) — see decode_step_t_h_emb.
8216 let mixer_fast = self.mixer_in_q8_1_fast(e, &layer.mixer);
8217 let norm_fused = std::env::var("MEMRA_NO_FUSE_NORMQ").is_err() && mixer_fast;
8218 let mut h = vbuf(e, t * n_embd)?; // fully written by either rms_norm arm
8219 if norm_fused {
8220 e.rms_norm_decode(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
8221 } else {
8222 e.rms_norm(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
8223 }
8224 let mixed = match &layer.mixer {
8225 Mixer::Full(fa) => {
8226 self.full_attn_verify(e, fa, &h, None, &pos_d, t, cache, il, None)?
8227 }
8228 Mixer::Mla(_) => {
8229 crate::hybrid::mla_path_unimplemented("auxiliary T-parallel decode")
8230 }
8231 Mixer::Kda(_) => crate::hybrid::kda_path_unimplemented("aux decode step"),
8232 Mixer::Linear(la) => {
8233 let mut out = e.zeros(t * n_embd)?;
8234 for col in 0..t {
8235 let mut h_col = e.zeros(n_embd)?;
8236 let src = h.slice(col * n_embd..(col + 1) * n_embd);
8237 e.copy_view_into(&mut h_col, 0, &src, n_embd)?;
8238 let m_col = self.linear_attn_decode(e, la, &h_col, cache, il)?;
8239 e.copy_into(&mut out, col * n_embd, &m_col, n_embd)?;
8240 }
8241 out
8242 }
8243 };
8244 let ffn_fuse = match &layer.ffn {
8245 crate::hybrid::Ffn::Dense {
8246 ffn_gate, ffn_up, ..
8247 } => {
8248 std::env::var("MEMRA_NO_FUSE_NORMQ").is_err()
8249 && e.uses_q8_1_fast(ffn_gate)
8250 && e.uses_q8_1_fast(ffn_up)
8251 }
8252 crate::hybrid::Ffn::Moe(_) => false,
8253 };
8254 let mut x1 = vbuf(e, t * n_embd)?; // fully written by add / add_rms_norm
8255 let mut z = vbuf(e, t * n_embd)?; // fully written by rms_norm_decode / add_rms_norm
8256 if ffn_fuse {
8257 e.add(&x, &mixed, &mut x1, t * n_embd)?;
8258 e.rms_norm_decode(
8259 &x1,
8260 layer.post_attn_norm.float_data(),
8261 &mut z,
8262 n_embd,
8263 t,
8264 eps,
8265 )?;
8266 } else {
8267 e.add_rms_norm(
8268 &x,
8269 &mixed,
8270 layer.post_attn_norm.float_data(),
8271 &mut x1,
8272 &mut z,
8273 n_embd,
8274 t,
8275 eps,
8276 )?;
8277 }
8278 let ffn_out = match &layer.ffn {
8279 crate::hybrid::Ffn::Dense {
8280 ffn_gate,
8281 ffn_up,
8282 ffn_down,
8283 } => {
8284 let n_ff = ffn_gate.out_features();
8285 let gate = e.matmul_decode_exact(ffn_gate, &z, t)?;
8286 let up = e.matmul_decode_exact(ffn_up, &z, t)?;
8287 let mut act = vbuf(e, t * n_ff)?; // fully written by ffn_act_lim
8288 // dense FFN clamp = the SHEXP array (upstream build_ffn serves both).
8289 Self::ffn_act_lim(
8290 e,
8291 &self.cfg,
8292 &gate,
8293 &up,
8294 1.0,
8295 1.0,
8296 self.cfg.clamp_shexp_at(il as u32),
8297 &mut act,
8298 t * n_ff,
8299 )?;
8300 e.matmul_decode_exact(ffn_down, &act, t)?
8301 }
8302 crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il(e, m, &z, t, il as u16)?,
8303 };
8304 let mut x2 = vbuf(e, t * n_embd)?; // fully written by add
8305 e.add(&x1, &ffn_out, &mut x2, t * n_embd)?;
8306 if aux_layers.contains(&il) {
8307 let mut a = e.zeros(n_embd)?;
8308 e.copy_view_into(&mut a, 0, &x2.slice((t - 1) * n_embd..t * n_embd), n_embd)?;
8309 aux_last.push(a);
8310 if let Some(pc) = pred_col {
8311 let mut ap = e.zeros(n_embd)?;
8312 e.copy_view_into(
8313 &mut ap,
8314 0,
8315 &x2.slice(pc * n_embd..(pc + 1) * n_embd),
8316 n_embd,
8317 )?;
8318 aux_pred.push(ap);
8319 }
8320 }
8321 x = x2;
8322 }
8323 let mut hn = vbuf(e, t * n_embd)?; // fully written by rms_norm_decode
8324 e.rms_norm_decode(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
8325 let logits = e.matmul_decode_exact(&self.output, &hn, t)?;
8326 let host = e.dtoh(&logits)?;
8327 cache.pos += t;
8328 Ok((
8329 host,
8330 aux_last,
8331 if want_pred { Some(aux_pred) } else { None },
8332 ))
8333 }
8334
8335 /// step35 SPEC-VERIFY attention over T query tokens — a per-row REPLAY of the eager
8336 /// `step35_decode_attn`.
8337 ///
8338 /// WHY A REPLAY AND NOT A BATCHED TWIN. The verify's whole job is to be bit-identical to what
8339 /// the eager decode would have computed for the same tokens; that is what makes greedy spec
8340 /// decode exact (run-spec asserts token identity for K=1..8). Every other verify arm in this
8341 /// file earns that identity by carefully mirroring dispatch (`matmul_decode_exact` to force
8342 /// MMVQ at any m, per-layer `ffn_fuse` mirroring, per-row `fa_decode` key bounds). step35
8343 /// stacks FOUR more per-layer degrees of freedom on top of that — per-layer `n_head`
8344 /// (64 full / 96 SWA), per-layer rotary width (64 full / 128 SWA), per-layer rope base, and a
8345 /// SEPARATE `attn_gate` tensor whose projection shares the attn-normed input — and its SWA
8346 /// layers attend through a token-OFFSET view whose offset is a function of the ABSOLUTE
8347 /// position of each query row. A batched twin would have to reproduce all of that AND the
8348 /// per-row offset in one launch; the offset alone rules out the existing rows kernels (they
8349 /// take one `base_len`, not a per-row offset).
8350 ///
8351 /// So this arm calls the eager path itself, once per row, on the same cache. Identity is then
8352 /// true BY CONSTRUCTION rather than by mirroring: row r runs exactly the kernel sequence that
8353 /// eager decode step r runs (same projections, same q8_1 fusion decision, same append, same
8354 /// view arithmetic, same `fa_decode_kvmod`, same gate), because it IS that code. Cost: T x the
8355 /// eager decode mixer instead of one batched pass — the same trade the generic arm's `else`
8356 /// per-row loop already accepts when `fa_rows_eligible` says no. Correctness first; a batched
8357 /// step35 twin is a perf lane's job and must be gated against this arm.
8358 ///
8359 /// The `h_q8` pre-quantized pair from the caller's fused norm is NOT forwarded: it is a
8360 /// T-row buffer and `step35_decode_attn`'s `pre_q` contract is one row. Instead each row's
8361 /// f32 `h` slice is handed over and the callee re-derives its own q8_1 exactly as eager decode
8362 /// does (`quantize_q8_1(h, 1, n_embd)`) — which is the dispatch being mirrored. Callers that
8363 /// took the fused arm therefore MUST still pass a live `h`; `step35_verify` asserts that.
8364 #[allow(clippy::too_many_arguments)]
8365 fn step35_verify(
8366 &self,
8367 e: &Engine,
8368 fa: &FullAttnLayer,
8369 h: &CudaSlice<f32>,
8370 h_q8: Option<(&CudaSlice<i8>, &CudaSlice<f32>)>,
8371 t: usize,
8372 cache: &mut Cache,
8373 il: usize,
8374 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8375 let n_embd = self.cfg.n_embd as usize;
8376 // The fused (h-less) attn-norm arm hands `h` as a zero-length placeholder. This arm needs
8377 // the f32 rows, so the caller must not take that lever for step35 — enforced at the call
8378 // site by the sliding-gated-MoE `Mixer::Full(_) => false` arm of
8379 // `lin_q8_only` in `decode_step_t_core_stream`, and asserted here so a future caller
8380 // cannot regress it into silently reading an empty buffer.
8381 assert_eq!(
8382 h.len(),
8383 t * n_embd,
8384 "step35_verify needs the f32 attn-normed rows ([t*n_embd]); the caller took the \
8385 fused q8-only norm arm (h_q8={}) — step35 must stay on the unfused arm",
8386 h_q8.is_some()
8387 );
8388 // ROW WIDTH IS n_embd, NOT n_head*head_dim: `step35_decode_attn` returns the mixer output
8389 // AFTER `wo`, so a row is [n_embd] — the same contract the generic arm's
8390 // `matmul_decode_exact(&fa.wo, &attn_g, t)` return has. Sizing this buffer from the
8391 // per-layer head geometry (8192 on full-attn, 12288 on SWA) instead overran the row on the
8392 // FIRST copy and panicked inside `copy_into`'s `CudaView::slice` unwrap
8393 // (raw/mtp-bt-20260806T212127Z.log frames 12-13).
8394 let mut out = vbuf(e, t * n_embd)?; // each row fully written by the copy below
8395 for r in 0..t {
8396 // Absolute position of this query row. `cache.pos` is the committed length at round
8397 // start and every row before r has already been appended by this loop, so the r-th
8398 // verify token sits at cache.pos + r — the same position eager decode would give it.
8399 let pos_d = e.htod_i32(&[(cache.pos + r) as i32])?;
8400 let mut h_row = vbuf(e, n_embd)?; // fully written by copy_view_into
8401 e.copy_view_into(
8402 &mut h_row,
8403 0,
8404 &h.slice(r * n_embd..(r + 1) * n_embd),
8405 n_embd,
8406 )?;
8407 // THE eager decode mixer: appends this row's K/V at kvl.len, advances it, then
8408 // attends over the (SWA-offset) view. Post-`wo`, same contract as this fn returns.
8409 let o = self.step35_decode_attn(e, fa, il, &h_row, None, &pos_d, cache)?;
8410 debug_assert_eq!(
8411 o.len(),
8412 n_embd,
8413 "step35_decode_attn returns post-wo [n_embd]"
8414 );
8415 e.copy_into(&mut out, r * n_embd, &o, n_embd)?;
8416 }
8417 Ok(out)
8418 }
8419
8420 /// Full-attention mixer over T query tokens with a GROWING resident KV (verify path, §D.3).
8421 /// Appends the T new K/V columns to cache.kv[il] then attends causally over [0..len) via
8422 /// fa_prefill. Token-major [T, kv_dim] projection layout == cache row layout (single copy).
8423 #[allow(clippy::too_many_arguments)]
8424 fn full_attn_verify(
8425 &self,
8426 e: &Engine,
8427 fa: &FullAttnLayer,
8428 h: &CudaSlice<f32>,
8429 h_q8: Option<(&CudaSlice<i8>, &CudaSlice<f32>)>,
8430 pos_d: &CudaSlice<i32>,
8431 t: usize,
8432 cache: &mut Cache,
8433 il: usize,
8434 stream_ctr: Option<&CudaSlice<i32>>,
8435 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8436 // step35: the generic geometry below is wrong for this arch (per-layer n_head, partial
8437 // per-layer rope, the SWA offset view, and a SEPARATE head-wise gate tensor), so it takes
8438 // its own arm. A verify that silently computes different attention than decode defeats the
8439 // whole self-consistency gate, so the arm is a per-row REPLAY of `step35_decode_attn`
8440 // rather than a batched twin — see `step35_verify` for why that is the exactness-correct
8441 // shape and not laziness.
8442 if self.sliding_gated_moe_batch_program() {
8443 if stream_ctr.is_some() {
8444 return Err(
8445 "step35 has no ROUND-STREAM verify arm (the device-counter _dc twins \
8446 cannot express the SWA offset KV view; same root cause as the dc \
8447 decode refusal) — run spec without the stream arm"
8448 .into(),
8449 );
8450 }
8451 return self.step35_verify(e, fa, h, h_q8, t, cache, il);
8452 }
8453 let cfg = &self.cfg;
8454 let geometry = cfg.full_attention_geometry_at(il as u32);
8455 let n_head = geometry.n_head as usize;
8456 let n_head_kv = geometry.n_head_kv as usize;
8457 let head_dim = geometry.head_dim_k as usize;
8458 let eps = cfg.rms_eps;
8459 let scale = geometry.attention_scale();
8460 let n_embd = cfg.n_embd as usize;
8461
8462 // DECODE-EXACT Q/K/V projections: matmul_decode_exact forces the MMVQ (warp-per-row) path
8463 // for every m, matching the T=1 decode's FP accumulation order. matmul_pre at m>=5 would
8464 // fall to dp4a (128-thread, two-level reduce) with a different FP sum order.
8465 // Q8 TRUNK-FUSION at T=1: DISPATCH-MIRRORS the eager decode's fused3 (bit-identical body).
8466 // BATCHED EPILOGUE RE-FUSE (lane/vt-fixes fix 2): `h_q8` = the attn-input norm's q8_1
8467 // form emitted by the fused rms_norm_q8_1 (bit-identical to rms_norm_decode ->
8468 // quantize_q8_1, kernel-check-pinned). When present (caller checked mixer q8_1-fast),
8469 // every projection consumes it — `h` may be a zero-len placeholder and must not be read.
8470 let (qf, mut k, v) = if let Some(mut qkv) = self.full_attn_tp_qkv(e, fa, h, t)? {
8471 let v = qkv.pop().ok_or("full-attention TP verify QKV omitted V")?;
8472 let k = qkv.pop().ok_or("full-attention TP verify QKV omitted K")?;
8473 let q = qkv.pop().ok_or("full-attention TP verify QKV omitted Q")?;
8474 if !qkv.is_empty() {
8475 return Err("full-attention TP verify QKV returned extra projections".into());
8476 }
8477 (q, k, v)
8478 } else {
8479 let mut fused = None;
8480 let qkv_fast =
8481 e.uses_q8_1_fast(&fa.wq) && e.uses_q8_1_fast(&fa.wk) && e.uses_q8_1_fast(&fa.wv);
8482 if t == 1 && qkv_fast {
8483 let (hq_o, hd_o);
8484 let (hq, hd): (&CudaSlice<i8>, &CudaSlice<f32>) = match h_q8 {
8485 Some(p) => p,
8486 None => {
8487 (hq_o, hd_o) = e.quantize_q8_1(h, 1, n_embd)?;
8488 (&hq_o, &hd_o)
8489 }
8490 };
8491 fused = e.matmul_q8_fused3(&fa.wq, &fa.wk, &fa.wv, hq, hd)?;
8492 } else if spec_fused_t() && (2..=4).contains(&t) && qkv_fast {
8493 // VERIFY-TIER TRUNK FUSION (MEMRA_SPEC_FUSED_T): one shared quantize + one
8494 // fused3 batched launch replaces three decode-exact calls (3 re-quantizes of
8495 // the same h + 3 _b2/_b4 launches). Bit-identical per (tensor,token,row).
8496 let (hq_o, hd_o);
8497 let (hq, hd): (&CudaSlice<i8>, &CudaSlice<f32>) = match h_q8 {
8498 Some(p) => p,
8499 None => {
8500 (hq_o, hd_o) = e.quantize_q8_1(h, t, n_embd)?;
8501 (&hq_o, &hd_o)
8502 }
8503 };
8504 fused = e.matmul_q8_fused3_t(&fa.wq, &fa.wk, &fa.wv, hq, hd, t)?;
8505 }
8506 match (fused, h_q8) {
8507 (Some(triple), _) => triple,
8508 // shared pre-quantized activation (q8_1-fast guaranteed by the caller): the
8509 // decode-exact dispatch consumes (hq, hd) instead of re-quantizing 3x.
8510 (None, Some((hq, hd))) if qkv_fast => (
8511 e.matmul_decode_exact_pre(&fa.wq, hq, hd, t)?,
8512 e.matmul_decode_exact_pre(&fa.wk, hq, hd, t)?,
8513 e.matmul_decode_exact_pre(&fa.wv, hq, hd, t)?,
8514 ),
8515 (None, _) => (
8516 e.matmul_decode_exact(&fa.wq, h, t)?,
8517 e.matmul_decode_exact(&fa.wk, h, t)?,
8518 e.matmul_decode_exact(&fa.wv, h, t)?,
8519 ),
8520 }
8521 };
8522 // M3/Hy3 have no attention output gate — wq out is exactly q; skip the split.
8523 let gated = geometry.attention_gate == memra_gguf::config::AttentionGateKind::FusedQ;
8524 let (mut q, gate) = if gated {
8525 let mut q = vbuf(e, t * n_head * head_dim)?; // fully written by q_gate_split
8526 let mut gate = vbuf(e, t * n_head * head_dim)?; // fully written by q_gate_split
8527 e.q_gate_split(&qf, &mut q, &mut gate, head_dim, n_head, t)?;
8528 (q, Some(gate))
8529 } else {
8530 (qf, None)
8531 };
8532
8533 let mut qn = vbuf(e, t * n_head * head_dim)?; // fully written by rms_norm
8534 e.rms_norm(
8535 &q,
8536 fa.q_norm.float_data(),
8537 &mut qn,
8538 head_dim,
8539 n_head * t,
8540 eps,
8541 )?;
8542 q = qn;
8543 let mut kn = vbuf(e, t * n_head_kv * head_dim)?; // fully written by rms_norm
8544 e.rms_norm(
8545 &k,
8546 fa.k_norm.float_data(),
8547 &mut kn,
8548 head_dim,
8549 n_head_kv * t,
8550 eps,
8551 )?;
8552 k = kn;
8553 let rope_dims = geometry.n_rot as usize;
8554 e.rope_neox(
8555 &mut q,
8556 pos_d,
8557 head_dim,
8558 rope_dims,
8559 n_head,
8560 t,
8561 geometry.rope_base,
8562 1.0,
8563 )?;
8564 e.rope_neox(
8565 &mut k,
8566 pos_d,
8567 head_dim,
8568 rope_dims,
8569 n_head_kv,
8570 t,
8571 geometry.rope_base,
8572 1.0,
8573 )?;
8574
8575 // append T new K/V columns to the resident QUANTIZED cache. k/v are token-major [T, kv_dim]
8576 // f32; append-quantize each of the T token rows into the byte cache (q8_0 K / q5_1 V).
8577 let kvl = cache.kv[il].as_mut().unwrap();
8578 let (kv_dim_k, kv_dim_v, ktb, vtb) =
8579 (kvl.kv_dim_k, kvl.kv_dim_v, kvl.k_tok_bytes, kvl.v_tok_bytes);
8580 if let Some(ctr) = stream_ctr {
8581 // stream: ONE batched append at the device counter (rows kernel = the per-view warp
8582 // math on a (block, token) grid, documented byte-identical); host len is a stale
8583 // LOWER BOUND under pre-issue (drain reconciles it).
8584 e.append_kv_quantized_rows_dc(
8585 &k, &v, &mut kvl.k, &mut kvl.v, ctr, t, kv_dim_k, kv_dim_v, ktb, vtb, false,
8586 )?;
8587 } else {
8588 for i in 0..t {
8589 let k_row = k.slice(i * kv_dim_k..(i + 1) * kv_dim_k);
8590 let v_row = v.slice(i * kv_dim_v..(i + 1) * kv_dim_v);
8591 e.append_kv_quantized_view(
8592 &k_row,
8593 &v_row,
8594 &mut kvl.k,
8595 &mut kvl.v,
8596 kvl.len + i,
8597 kv_dim_k,
8598 kv_dim_v,
8599 ktb,
8600 vtb,
8601 false,
8602 )?;
8603 }
8604 kvl.len += t;
8605 }
8606
8607 // BIT-IDENTICAL VERIFY ATTENTION (spec-exactness fix): the FP accumulation order must be
8608 // byte-for-byte identical to the eager decode path. fa_prefill uses a different tile size
8609 // (BLOCK_Q=64, BK=32) and online-softmax structure than fa_decode's split-K + combine,
8610 // which changes FP summation order and can flip argmax at tight logit margins. Query row r
8611 // attends to keys [0..base_len+r+1) — each successive row sees one more key (the causal
8612 // property). This matches eager: decode appends k at len, then fa_decode sees t_kv = len+1
8613 // keys. The verify appends all T tokens first but bounds the key range per row.
8614 //
8615 // MULTI-ROW FUSED PATH (the long-ctx spec fix, 2026-07-03): when every row takes the vec
8616 // kernel (base_len+1 >= FA_VEC_MIN_TKV), ONE fa_decode_rows launch executes the exact
8617 // per-row program for all T rows (grid.z = row, per-row n_splits from the same
8618 // fa_split_keys formula) — replacing T x (2 launches + 2 dtod copies + 5 partial allocs)
8619 // and multiplying resident CTAs by T on a latency-bound kernel. Bit-identical per row by
8620 // construction; kernel-check pins rows-vs-loop byte identity, run-spec is the end gate.
8621 // Short ctx (any row below the vec crossover) and MEMRA_NO_FA_VEC/MEMRA_FA_ROWS_OFF keep the
8622 // per-row loop (whose fa_decode picks scalar/vec per row exactly like eager decode).
8623 let mut attn = vbuf(e, t * n_head * head_dim)?; // fully written by every FA arm below
8624 let base_len = kvl.len - t; // KV len BEFORE this round's T tokens were appended
8625 // T=1 INCLUDED (2026-07-05): p-min cuts the draft to 1 in ~75% of rounds on hard
8626 // (agentic) content — the old t>1 gate sent those rounds to the per-row loop (262us/row
8627 // + q-row copy + per-row allocs vs 93us/row through the fused kernel at grid.z=1, same
8628 // program). nsys accounting: 1088 of 1456 verify FA launches were T=1 escapees.
8629 // LEAN T=1 ARM (MEMRA_SPEC_LEAN, close35): at t==1, q IS one row and fa_decode on it is
8630 // the EXACT eager decode dispatch (vec_q_v2 + combine_f32; the rows pair measured +50us
8631 // at m=1). Byte-identical: kernel-check pins rows-vs-loop identity, and the per-row loop
8632 // at t=1 is fa_decode on the same q with zero-offset copies. Gates arbitrate.
8633 if let Some(ctr) = stream_ctr {
8634 // STREAM ARM: causal base from the device counter; host kvl.len is a stale lower
8635 // bound used only for the split-sizing upper bound (+64 slack covers M pre-issued
8636 // rounds at K<=8). Views span the bound; per-row limits derive in-kernel.
8637 let upper = kvl.len + t + 64;
8638 let k_view = e.view_u8(&kvl.k, (upper.min(cache.max_ctx)) * ktb);
8639 let v_view = e.view_u8(&kvl.v, (upper.min(cache.max_ctx)) * vtb);
8640 e.fa_decode_rows_dc(
8641 &q,
8642 &k_view,
8643 &v_view,
8644 &mut attn,
8645 head_dim,
8646 n_head,
8647 n_head_kv,
8648 ctr,
8649 upper.min(cache.max_ctx),
8650 t,
8651 scale,
8652 ktb,
8653 vtb,
8654 0,
8655 false,
8656 )?;
8657 } else if spec_lean() && t == 1 {
8658 let t_kv = base_len + 1;
8659 let k_view = e.view_u8(&kvl.k, t_kv * ktb);
8660 let v_view = e.view_u8(&kvl.v, t_kv * vtb);
8661 e.fa_decode_kvmod(
8662 &q, &k_view, &v_view, &mut attn, head_dim, n_head, n_head_kv, t_kv, scale, ktb,
8663 vtb, false,
8664 )?;
8665 } else if e.fa_rows_eligible(base_len, head_dim) {
8666 let k_view = e.view_u8(&kvl.k, (base_len + t) * ktb);
8667 let v_view = e.view_u8(&kvl.v, (base_len + t) * vtb);
8668 e.fa_decode_rows(
8669 &q, &k_view, &v_view, &mut attn, head_dim, n_head, n_head_kv, base_len, t, scale,
8670 ktb, vtb, None, false, false, None,
8671 )?;
8672 } else {
8673 for r in 0..t {
8674 let t_kv_r = base_len + r + 1; // this row sees keys [0..t_kv_r)
8675 let k_view_r = e.view_u8(&kvl.k, t_kv_r * ktb);
8676 let v_view_r = e.view_u8(&kvl.v, t_kv_r * vtb);
8677 // copy q row into an owned buffer (fa_decode takes &CudaSlice, not CudaView)
8678 let mut q_row = vbuf(e, n_head * head_dim)?; // fully written by copy_view_into
8679 let q_src = q.slice(r * n_head * head_dim..(r + 1) * n_head * head_dim);
8680 e.copy_view_into(&mut q_row, 0, &q_src, n_head * head_dim)?;
8681 let mut attn_row = vbuf(e, n_head * head_dim)?; // fully written by fa_decode
8682 e.fa_decode_kvmod(
8683 &q_row,
8684 &k_view_r,
8685 &v_view_r,
8686 &mut attn_row,
8687 head_dim,
8688 n_head,
8689 n_head_kv,
8690 t_kv_r,
8691 scale,
8692 ktb,
8693 vtb,
8694 false,
8695 )?;
8696 e.copy_into(
8697 &mut attn,
8698 r * n_head * head_dim,
8699 &attn_row,
8700 n_head * head_dim,
8701 )?;
8702 }
8703 }
8704
8705 let attn_g = match &gate {
8706 Some(gate) => {
8707 let mut gsig = vbuf(e, t * n_head * head_dim)?; // fully written by sigmoid
8708 e.sigmoid(gate, &mut gsig, t * n_head * head_dim)?;
8709 let mut ag = vbuf(e, t * n_head * head_dim)?; // fully written by mul
8710 e.mul(&attn, &gsig, &mut ag, t * n_head * head_dim)?;
8711 ag
8712 }
8713 None => attn,
8714 };
8715 // DECODE-EXACT wo projection: at m>=5 (K=4+ with pending) the generic matmul would use dp4a
8716 // (128-thread, different FP sum order than MMVQ). Force MMVQ for bit-identity with decode.
8717 match self.full_attn_tp_o(e, fa, &attn_g, t)? {
8718 Some(output) => Ok(output),
8719 None => Ok(e.matmul_decode_exact(&fa.wo, &attn_g, t)?),
8720 }
8721 }
8722
8723 /// Context-linear bytes for a plain serving session's trunk cache.
8724 pub fn plain_session_kv_bytes_per_token(&self) -> usize {
8725 crate::cache::cache_bytes_per_token_for_plan(
8726 &self.cfg,
8727 &self.plan,
8728 0,
8729 self.plan.layers.len(),
8730 )
8731 }
8732
8733 /// `(logical bytes/token, ring-capped bytes/token, ring row cap)` for exact admission.
8734 pub fn plain_session_kv_shape(&self) -> (usize, usize, usize) {
8735 (
8736 self.plain_session_kv_bytes_per_token(),
8737 crate::cache::cache_ring_bytes_per_token_for_plan(
8738 &self.cfg,
8739 &self.plan,
8740 0,
8741 self.plan.layers.len(),
8742 ),
8743 crate::cache::cache_ring_row_cap_for_plan(&self.plan),
8744 )
8745 }
8746
8747 /// Context-linear bytes for a speculative serving session: trunk cache plus persistent MTP
8748 /// scratch. With no MTP head this equals the plain coefficient.
8749 pub fn spec_session_kv_bytes_per_token(&self) -> usize {
8750 let scratch = self
8751 .mtp
8752 .iter()
8753 .chain(self.mtp_extra.iter())
8754 .map(|mtp| {
8755 let (_, _, k, v) = mtp_scratch_layout(&self.cfg, mtp.geom.as_ref());
8756 k + v
8757 })
8758 .sum::<usize>();
8759 self.plain_session_kv_bytes_per_token()
8760 .saturating_add(scratch)
8761 }
8762
8763 /// Spec twin of [`HybridModel::plain_session_kv_shape`]; Step35's persistent MTP scratch is
8764 /// capped by the same SWA ring rows as the trunk.
8765 pub fn spec_session_kv_shape(&self) -> (usize, usize, usize) {
8766 let total = self.spec_session_kv_bytes_per_token();
8767 let (_, mut ring, rows) = self.plain_session_kv_shape();
8768 if rows > 0 {
8769 ring = ring.saturating_add(
8770 self.mtp
8771 .iter()
8772 .chain(self.mtp_extra.iter())
8773 .map(|mtp| {
8774 let (_, _, k, v) = mtp_scratch_layout(&self.cfg, mtp.geom.as_ref());
8775 k + v
8776 })
8777 .sum::<usize>(),
8778 );
8779 }
8780 (total, ring, rows)
8781 }
8782
8783 /// Greedy MTP speculative decode (§B). Token-identical to `generate(prompt, max_new)` but uses
8784 /// the NextN head to draft K tokens then verifies them in one batched target forward.
8785 /// Returns (generated tokens, total_drafted, total_accepted) so the caller can report
8786 /// acceptance rate. `k` = draft length per round.
8787 ///
8788 /// GRAPH DRAFT (stage 2 of graph-grade spec): when the model is all-Dense and the MTP head is
8789 /// Dense (no MoE host readbacks), the fixed-shape T=1 MTP forward is CUDA-graph-captured ONCE
8790 /// and replayed per draft step — the ~40 eager launches per drafted token collapse into one
8791 /// graph dispatch; only the 4-byte token id (and 4-byte p-min confidence) round-trip per step.
8792 /// Event tracking is disabled for the whole call (generate_graph pattern) so every buffer the
8793 /// captured graph references is event-free; the spec loop is strictly single-stream.
8794 /// MEMRA_SPEC_NOGRAPH=1 forces the eager draft chain.
8795 /// SAMPLED mode (MEMRA_SPEC_TEMP>0) has its OWN capture (gumbel-perturbed in-graph argmax,
8796 /// device Philox event counter, persistent q retention) — graph-vs-eager sampled streams are
8797 /// bit-identical for the same (seed, prompt, K, temp); see the sampled-graph setup in
8798 /// generate_spec_inner2.
8799 /// Multi-turn session: trunk cache + MTP draft scratch persist across generate calls, so
8800 /// turn N+1 primes ONLY its new suffix (the 124k-conversation daily pattern — re-priming a
8801 /// 32k history costs ~54s; a suffix prime costs seconds). APPEND-ONLY by construction: the
8802 /// hybrid linear-attn states are in-place (no position index), so a session can extend but
8803 /// never rewind — `committed` is the exact token list whose state the caches hold (includes
8804 /// any overshoot tokens past max_new; the caller renders from `committed`, not its own echo).
8805 pub fn new_session(
8806 &self,
8807 e: &Engine,
8808 max_ctx: usize,
8809 ) -> Result<SpecSession, Box<dyn std::error::Error>> {
8810 Ok(SpecSession {
8811 // STAGE-OWNED KV (lane/pp2-spec 2026-08-06): `pp::new_cache`, not `Cache::new`. This
8812 // is the SERVING spec-session path, and with the ppN door open across two cards a
8813 // primary-homed cache makes every remote stage peer-read its OWN KV on every verify
8814 // round — the wrong-card class already fixed on the two batched serving paths
8815 // (worker.rs 2483 / 2837). With the door shut `new_cache` IS `Cache::new` (same
8816 // branch, same allocations), so single-device behavior is byte-unchanged.
8817 cache: crate::pp::new_cache_planned(e, &self.cfg, &self.plan, max_ctx)?,
8818 scratch: self.new_mtp_scratch(e, max_ctx)?,
8819 committed: Vec::new(),
8820 last_h: None,
8821 next_pred: None,
8822 sctr: 0,
8823 uctr: 0,
8824 draft_ctx: None,
8825 pending_tok: None,
8826 turn_ckpt: None,
8827 telem: SpecTelemetryCounters::default(),
8828 capture_at: None,
8829 boundary_captures: Vec::new(),
8830 ckpt_at: None,
8831 capture_disabled: false,
8832 })
8833 }
8834
8835 /// SPEC-ON-CACHE-HIT restore (lane/spec-on-cache-hit, 2026-08-18 — PORT-PLAN item 3,
8836 /// research/cache-spec-design-20260814, scoped to WHOLE-ENTRY restores only): build a
8837 /// SpecSession around a trunk cache the worker already restored from a prefix-cache
8838 /// entry, re-installing the entry's published draft plane as the MTP scratch rows
8839 /// `[0..prefix.len())` and the entry's boundary hidden as `last_h`, then feeding the
8840 /// prompt SUFFIX here — through EXACTLY the plain path's program selection — so the
8841 /// worker always receives a fully-warm continuation session (committed = whole
8842 /// prompt, `next_pred` + `last_h` set; caller sets `next_pred` from the entry's
8843 /// boundary logits on the empty-suffix shape).
8844 ///
8845 /// PROGRAM LAW (the splitiso two-programs class, learned AGAIN in this lane's own
8846 /// gate): the identity target for a converted hit is the PLAIN hit serving the same
8847 /// request, and plain feeds a carried suffix via eager `decode_step` below
8848 /// PRIME_MIN_T and via `prime_cache` at/above it (prefill_tick's arms). The generate
8849 /// path's tokenwise arm routes qwen35-class through the BATCHED T=1 program
8850 /// (`spec_target_step_h`) instead — ULP-different suffix rows, and the gate measured
8851 /// the near-tie flip at generated token ~8 (research/spec-cache-20260818, qwen r3).
8852 /// So the suffix is fed HERE, mirroring prefill_tick arm-for-arm, not handed to the
8853 /// burst prime.
8854 ///
8855 /// SEED RULE (both sampling regimes; lane/sampled-hit-spec 2026-08-19, sampled draw
8856 /// added by lane/sampled-spec-quality 2026-08-19). The boundary token is produced by
8857 /// EXACTLY the rule the cold burst entry applies to its own first token from the same
8858 /// logits row: `argmax` when greedy, and a `sample_boundary_token` draw at Philox
8859 /// counter 0 when sampled. Both shapes are covered — the entry's boundary logits on a
8860 /// full-cover (empty-suffix) hit, this feed's own boundary logits on a suffix hit.
8861 /// That is what keeps a restored session seed-identical to a cold one PER SEED: the
8862 /// cold session draws from the identical row at counter 0 and then runs its rounds from
8863 /// counter 1, so the restored session admits with `sctr = 1` after its own draw.
8864 /// The WORKER owns the one refusal this constructor cannot see — a constrained request.
8865 /// (The penalized-sampled refusal was LIFTED once the burst's penalty window learned to
8866 /// span the session: `committed` here is the WHOLE prompt, so the restored session's
8867 /// window is the cold session's window. It comes back if `MEMRA_SPEC_PEN_SESSION=0`.)
8868 ///
8869 /// NOT the rolled-back partial-restore hazard: the caller restores at exactly the
8870 /// entry's captured endpoint (`e.pos`) through the shipping whole-entry path;
8871 /// mid-entry (`at < e.pos`) trunk restores stay behind MEMRA_PREFIX_PARTIAL_RESTORE
8872 /// and are never routed here.
8873 ///
8874 /// Failure contract: `Err((Some(cache), why))` before any trunk mutation — the
8875 /// worker rebuilds the plain carrier and the hit serves plain, byte-unchanged.
8876 /// `Err((None, why))` after the suffix feed began — the carrier is part-fed and
8877 /// UNUSABLE; the worker serves the request cold-plain (correct, slower) and the
8878 /// entry stays published for the next request.
8879 #[allow(clippy::too_many_arguments)]
8880 #[allow(clippy::result_large_err)] // allow: the fat error type is the diagnostic contract here; boxing it would change the error surface
8881 pub fn spec_session_from_restored(
8882 &self,
8883 e: &Engine,
8884 mut cache: Cache,
8885 prefix: Vec<u32>,
8886 suffix: &[u32],
8887 draft_k: &CudaSlice<u8>,
8888 draft_v: &CudaSlice<u8>,
8889 draft_k_tok_bytes: usize,
8890 draft_v_tok_bytes: usize,
8891 draft_len: usize,
8892 last_h: &[f32],
8893 // The ENTRY's boundary logits row (the full-cover shape's seed source). May be empty
8894 // when a suffix follows — the feed's own logits are the boundary then.
8895 boundary_logits: &[f32],
8896 // The request's sampler, or None for greedy. Owned here so the seed rule lives in
8897 // ONE place instead of being half-applied by the worker.
8898 sampling: Option<SpecSampling>,
8899 require_anchor: bool,
8900 max_ctx: usize,
8901 // STABLE-BOUNDARY REPUBLICATION (lane/frspec-multiturn-cache, 2026-08-21): ABSOLUTE
8902 // prompt position to split the suffix feed at and capture the extended-entry
8903 // publication + this session's `turn_ckpt` — the worker's stable pre-generation
8904 // boundary (`plain_checkpoint_boundary`). None = legacy prompt-end republication.
8905 // WHY: the prompt-end capture below includes the template's live generation header
8906 // (`<|im_start|>assistant\n<think>\n`), which the next turn's re-render replaces, so
8907 // for a hybrid (whole-entry restores only) every extended entry's last ~2 tokens
8908 // diverged from every future prompt and the hit boundary FROZE at the first
8909 // lcp-split entry forever (measured: cached 6811 of 38228 by turn 8, B4).
8910 republish_at: Option<usize>,
8911 ) -> Result<SpecSession, (Option<Cache>, String)> {
8912 let pos = prefix.len();
8913 let fail = |cache: Cache, msg: String| -> Result<SpecSession, (Option<Cache>, String)> {
8914 Err((Some(cache), msg))
8915 };
8916 if let Err(error) = cache.ensure_usable("spec_session_from_restored") {
8917 drop(cache);
8918 return Err((None, error.to_string()));
8919 }
8920 if self.mtp.is_none() {
8921 return fail(cache, "no MTP head attached (nothing to draft with)".into());
8922 }
8923 if pos == 0 {
8924 return fail(cache, "empty committed prefix".into());
8925 }
8926 if cache.pos != pos {
8927 let msg = format!(
8928 "restored cache pos {} != restored prefix len {pos}",
8929 cache.pos
8930 );
8931 return fail(cache, msg);
8932 }
8933 if draft_len != pos {
8934 return fail(
8935 cache,
8936 format!("draft plane len {draft_len} != restored prefix len {pos}"),
8937 );
8938 }
8939 if pos + suffix.len() >= max_ctx {
8940 return fail(
8941 cache,
8942 format!(
8943 "prompt {} + suffix would not leave generation room in ctx {max_ctx}",
8944 pos + suffix.len(),
8945 ),
8946 );
8947 }
8948 let mut scratch = match MtpScratch::new(
8949 e,
8950 &self.cfg,
8951 &self.plan,
8952 max_ctx,
8953 self.mtp.as_ref().and_then(|m| m.geom.as_ref()),
8954 ) {
8955 Ok(s) => s,
8956 Err(err) => return fail(cache, format!("draft scratch alloc failed: {err}")),
8957 };
8958 if scratch.kv.ring.is_some() {
8959 return fail(
8960 cache,
8961 "ring-backed draft scratch (Step35 SWA) cannot take a flat prefix restore".into(),
8962 );
8963 }
8964 if scratch.kv.k_tok_bytes != draft_k_tok_bytes
8965 || scratch.kv.v_tok_bytes != draft_v_tok_bytes
8966 {
8967 return fail(
8968 cache,
8969 format!(
8970 "draft plane layout {draft_k_tok_bytes}/{draft_v_tok_bytes} != scratch \
8971 {}/{} bytes/token (stale entry across a format change)",
8972 scratch.kv.k_tok_bytes, scratch.kv.v_tok_bytes,
8973 ),
8974 );
8975 }
8976 if pos > scratch.cap {
8977 return fail(
8978 cache,
8979 format!(
8980 "draft plane rows {pos} exceed scratch capacity {}",
8981 scratch.cap
8982 ),
8983 );
8984 }
8985 let kb = pos * draft_k_tok_bytes;
8986 let vb = pos * draft_v_tok_bytes;
8987 if draft_k.len() < kb || draft_v.len() < vb {
8988 return fail(
8989 cache,
8990 format!(
8991 "truncated draft plane: K {} < {kb} or V {} < {vb} bytes",
8992 draft_k.len(),
8993 draft_v.len(),
8994 ),
8995 );
8996 }
8997 if kb > 0
8998 && let Err(err) = e.copy_u8_into(&mut scratch.kv.k, 0, draft_k, kb)
8999 {
9000 return fail(cache, format!("draft K restore copy failed: {err}"));
9001 }
9002 if vb > 0
9003 && let Err(err) = e.copy_u8_into(&mut scratch.kv.v, 0, draft_v, vb)
9004 {
9005 return fail(cache, format!("draft V restore copy failed: {err}"));
9006 }
9007 if let Err(err) = scratch.set_len(e, pos) {
9008 return fail(cache, format!("draft scratch len set failed: {err}"));
9009 }
9010 let mut last_h_dev = if last_h.len() == self.cfg.n_embd as usize {
9011 // anchor upload failure is acceptance-only when a suffix feed follows (fill
9012 // row-0 falls back to zeros) but FATAL for an empty-suffix continuation (the
9013 // burst entry asserts committed + last_h + next_pred) — the caller says which.
9014 e.htod(last_h).ok()
9015 } else {
9016 None
9017 };
9018 if require_anchor && last_h_dev.is_none() {
9019 return fail(
9020 cache,
9021 "empty-suffix continuation requires the entry's boundary hidden anchor".into(),
9022 );
9023 }
9024 let mut committed = prefix;
9025 // Set on BOTH shapes below (suffix-fed and full-cover) — never left None, which is
9026 // what the empty-suffix continuation assert in the burst entry requires.
9027 let next_pred;
9028 // Philox: (0,0) at admit exactly like a fresh session; a sampled boundary draw below
9029 // consumes counter 0 and leaves 1, which is the state a cold session reaches after
9030 // drawing its own first token from the same row.
9031 let mut sctr = 0u32;
9032 let sampled = sampling.is_some_and(|s| s.temp > 0.0) && spec_sampled_boundary_on();
9033 // Penalty window for the boundary draw: the last `penalty_last_n` tokens of the WHOLE
9034 // prompt, which is what the cold session's own burst sees (Item 2's window). Built
9035 // after the suffix joins `committed` below.
9036 let mut boundary_captures: Vec<SpecBoundaryCapture> = Vec::new();
9037 let mut restored_turn_ckpt: Option<SpecCheckpoint> = None;
9038 if !suffix.is_empty() {
9039 // ---- SUFFIX FEED, mirroring prefill_tick's program selection exactly ----
9040 // From here on the trunk cache mutates: failures return Err((None, _)) and
9041 // the worker serves the request cold-plain instead of reusing the carrier.
9042 let dirty =
9043 |msg: String| -> Result<SpecSession, (Option<Cache>, String)> { Err((None, msg)) };
9044 let n_embd = self.cfg.n_embd as usize;
9045 let t = suffix.len();
9046 let mut h_rows = match e.uninit(t * n_embd) {
9047 Ok(b) => b,
9048 Err(err) => return fail(cache, format!("suffix hidden buffer alloc: {err}")),
9049 };
9050 // STABLE-BOUNDARY split (see `republish_at`): feed stops at the boundary so the
9051 // in-place GDN conv/ssm state can be snapshotted there — the only moment it
9052 // exists (the cold prime-split law). suffix-relative; None = one-segment legacy.
9053 let b_rel = republish_at
9054 .and_then(|abs| abs.checked_sub(pos))
9055 .filter(|&r| r > 0 && r < t);
9056 let mut feed_logits = Vec::new();
9057 let tokenwise_env = std::env::var("MEMRA_PRIME_TOKENWISE").is_ok()
9058 || e.frozen_cpu_experts_prefer_tokenwise_prime();
9059 let mut fed = 0usize;
9060 for seg_end in [b_rel, Some(t)].into_iter().flatten() {
9061 if seg_end <= fed {
9062 continue;
9063 }
9064 let seg = &suffix[fed..seg_end];
9065 let batched = seg.len() >= crate::hybrid_forward::PRIME_MIN_T && !tokenwise_env;
9066 if batched {
9067 // prefill_tick's prime arm: request-level prime_cache call; tokens still
9068 // queued after this segment ride `queued_after` so Step35 arm selection
9069 // stays keyed to the request's end (tick-seg law).
9070 match self.prime_cache(e, seg, &mut cache, t - seg_end) {
9071 Ok((l, _h_seed, hiddens)) => {
9072 if let Err(err) =
9073 e.copy_into(&mut h_rows, fed * n_embd, &hiddens, seg.len() * n_embd)
9074 {
9075 return dirty(format!("suffix hidden copy: {err}"));
9076 }
9077 feed_logits = l;
9078 }
9079 Err(err) => return dirty(format!("suffix prime failed: {err}")),
9080 }
9081 } else {
9082 // prefill_tick's tokenwise arm: eager decode_step, one token at a time.
9083 for (i, &tok) in seg.iter().enumerate() {
9084 match self.decode_step_h(e, tok, &mut cache) {
9085 Ok((l, h)) => {
9086 if let Err(err) =
9087 e.copy_into(&mut h_rows, (fed + i) * n_embd, &h, n_embd)
9088 {
9089 return dirty(format!("suffix hidden copy: {err}"));
9090 }
9091 feed_logits = l;
9092 }
9093 Err(err) => return dirty(format!("suffix decode_step failed: {err}")),
9094 }
9095 }
9096 }
9097 fed = seg_end;
9098 if Some(seg_end) == b_rel {
9099 // The stable pre-generation boundary: capture the extended-entry
9100 // publication AND this session's own turn checkpoint here instead of at
9101 // prompt-end (both would otherwise carry the volatile live-header tail
9102 // the next re-render replaces). Failure silent, turn_ckpt convention.
9103 debug_assert_eq!(
9104 cache.pos,
9105 pos + seg_end,
9106 "stable-boundary capture off the feed split"
9107 );
9108 if spec_restore_republish_on()
9109 && let Ok(snap) = cache.snapshot(e)
9110 {
9111 boundary_captures.push(SpecBoundaryCapture {
9112 snap,
9113 pos: pos + seg_end,
9114 logits: feed_logits.clone(),
9115 last_h: capture_boundary_hidden(e, &h_rows, seg_end, n_embd),
9116 latent_tails: Vec::new(),
9117 });
9118 }
9119 let anchor: Result<CudaSlice<f32>, Box<dyn std::error::Error>> =
9120 e.uninit(n_embd).and_then(|mut a| {
9121 e.copy_view_into(
9122 &mut a,
9123 0,
9124 &h_rows.slice((seg_end - 1) * n_embd..seg_end * n_embd),
9125 n_embd,
9126 )?;
9127 Ok(a)
9128 });
9129 if let (Ok(snap), Ok(last_h)) = (cache.snapshot(e), anchor) {
9130 restored_turn_ckpt = Some(SpecCheckpoint {
9131 snap,
9132 pos: pos + seg_end,
9133 last_h,
9134 });
9135 }
9136 }
9137 }
9138 // Draft-scratch fill for the suffix rows, predecessor-paired: row `pos` reads
9139 // the entry's boundary anchor (zeros fallback — acceptance-only), row `pos+i`
9140 // reads h_rows[i-1]. Chunked like the generate path's fill (transients scale
9141 // with T). Fill failures are acceptance-only — truncate to the restored rows
9142 // and continue; the burst's own set_len keeps the invariant.
9143 let _mtp = self.mtp.as_ref().expect("mtp checked above"); // invariant check only; the fill below re-reads self.mtp
9144 let (embd_qt, embd_rb) = self.embd.qt_and_row_bytes(n_embd);
9145 let embd_gpu = if spec_host_embd() {
9146 None
9147 } else {
9148 Some(
9149 self.embd_gpu
9150 .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload")),
9151 )
9152 };
9153 let embd_dev = embd_gpu.map(|g| (g, embd_qt, embd_rb));
9154 let fill_chunk = 4096usize;
9155 let mut filled = true;
9156 let mut start = 0usize;
9157 'fill: while start < t {
9158 let end = (start + fill_chunk).min(t);
9159 let tc = end - start;
9160 let Ok(mut phs) = e.zeros(tc * n_embd) else {
9161 filled = false;
9162 break 'fill;
9163 };
9164 let (src_lo, dst_off, n_copy) = if start == 0 {
9165 (0, n_embd, (tc - 1) * n_embd)
9166 } else {
9167 ((start - 1) * n_embd, 0, tc * n_embd)
9168 };
9169 if start == 0
9170 && let Some(lh) = last_h_dev.as_ref()
9171 && e.copy_into(&mut phs, 0, lh, n_embd).is_err()
9172 {
9173 filled = false;
9174 break 'fill;
9175 }
9176 if n_copy > 0
9177 && e.copy_view_into(
9178 &mut phs,
9179 dst_off,
9180 &h_rows.slice(src_lo..src_lo + n_copy),
9181 n_copy,
9182 )
9183 .is_err()
9184 {
9185 filled = false;
9186 break 'fill;
9187 }
9188 if self
9189 .mtp_kv_fill_all(
9190 e,
9191 &suffix[start..end],
9192 &phs,
9193 pos + start,
9194 &mut scratch,
9195 embd_dev,
9196 )
9197 .is_err()
9198 {
9199 filled = false;
9200 break 'fill;
9201 }
9202 start = end;
9203 }
9204 if !filled {
9205 // acceptance-only: drafts over missing suffix rows are cheap and wrong,
9206 // so keep only the restored rows resident and let verify arbitrate.
9207 if let Err(err) = scratch.set_len(e, pos) {
9208 return dirty(format!("scratch truncation after failed fill: {err}"));
9209 }
9210 }
9211 // EXTENDED-ENTRY PUBLICATION (lane/sampled-spec-quality, Item 3 — the fix for
9212 // "a restored spec session never publishes an extended entry", SAMPLED-HIT.md
9213 // finding (d)). Pre-lane, publication was armed only for COLD sessions
9214 // (`spec_resumed == 0` in the worker) and both engine capture sites require a
9215 // non-continuation burst — but a converted hit's first burst IS a continuation,
9216 // so a growing conversation learned exactly ONE boundary and turn 3 could never
9217 // hit a longer prefix than turn 2 did.
9218 //
9219 // WHERE, and why it is safe here: `cache.pos == prefix + suffix` at this exact
9220 // line — the trunk is primed over the whole prompt, nothing is generated, and the
9221 // draft plane rows [0..prompt) are filled just above. That is a complete
9222 // whole-entry boundary (`pos == fed_len`), the same shape the cold seed capture
9223 // publishes; the worker's existing publication sweep picks it up because it is
9224 // keyed on non-empty `boundary_captures` and is sampler- and resume-independent.
9225 // NOT the partial-restore hazard: the boundary is this session's own prompt END,
9226 // never mid-entry, so `entry_pos != fed_len` still refuses on the way back in.
9227 // Failure is SILENT by design (the turn_ckpt / boundary-capture convention):
9228 // publication is an optimization, never a correctness dependency.
9229 //
9230 // SUPERSEDED WHEN `republish_at` FIRED (lane/frspec-multiturn-cache): a prompt-end
9231 // entry's tail is the live generation header the next re-render replaces, so on a
9232 // hybrid (whole-entry restores) it can never serve the conversation's next turn —
9233 // the stable-boundary capture above IS this publication, minus the poisoned tail.
9234 if spec_restore_republish_on() && boundary_captures.is_empty() {
9235 debug_assert_eq!(
9236 cache.pos,
9237 pos + t,
9238 "extended-entry capture must sit at the restored session's prompt end",
9239 );
9240 if let Ok(snap) = cache.snapshot(e) {
9241 boundary_captures.push(SpecBoundaryCapture {
9242 snap,
9243 pos: pos + t,
9244 logits: feed_logits.clone(),
9245 last_h: capture_boundary_hidden(e, &h_rows, t, n_embd),
9246 latent_tails: Vec::new(),
9247 });
9248 }
9249 }
9250 // continuation seed: the feed's boundary logits ARE the plain path's boundary
9251 // logits (same program), so greedy's argmax here is plain's first emitted token,
9252 // and the sampled draw is the cold sampled session's own first token.
9253 next_pred = Some(if sampled {
9254 let sp = sampling.expect("sampled implies a sampler");
9255 // `committed` is still the restored prefix here; the suffix joins it below —
9256 // so this is the last-N window over the WHOLE prompt, exactly the cold
9257 // session's own window at its first token.
9258 let hist = pen_window_seed(&committed, suffix, sp.penalty_last_n);
9259 match sample_boundary_token(
9260 e,
9261 &feed_logits,
9262 &sp,
9263 &hist,
9264 &mut sctr,
9265 "restore-suffix-feed",
9266 ) {
9267 Ok(t) => t,
9268 // the trunk is already fed: hand nothing back, the worker serves the
9269 // request cold-plain. Never fall back to an argmax — that would put a
9270 // greedy token in a sampled stream to save a slow path.
9271 Err(err) => {
9272 return dirty(format!("boundary token draw failed: {err}"));
9273 }
9274 }
9275 } else {
9276 argmax(&feed_logits) as u32
9277 });
9278 let mut lh = match e.uninit(n_embd) {
9279 Ok(b) => b,
9280 Err(err) => return dirty(format!("boundary hidden alloc: {err}")),
9281 };
9282 if let Err(err) = e.copy_view_into(
9283 &mut lh,
9284 0,
9285 &h_rows.slice((t - 1) * n_embd..t * n_embd),
9286 n_embd,
9287 ) {
9288 return dirty(format!("boundary hidden copy: {err}"));
9289 }
9290 last_h_dev = Some(lh);
9291 committed.extend_from_slice(suffix);
9292 } else {
9293 // FULL-COVER shape (empty suffix — the identical-repeat / agent-loop shape): the
9294 // ENTRY's boundary logits are the boundary row, and this is the token the cold
9295 // session emits from that same row. Owned here rather than in the worker so the
9296 // sampled draw cannot be half-applied on one shape (the worker used to argmax it).
9297 if boundary_logits.is_empty() {
9298 return fail(
9299 cache,
9300 "full-cover restore without the entry's boundary logits".into(),
9301 );
9302 }
9303 next_pred = Some(if sampled {
9304 let sp = sampling.expect("sampled implies a sampler");
9305 let hist = pen_window_seed(&committed, &[], sp.penalty_last_n);
9306 match sample_boundary_token(
9307 e,
9308 boundary_logits,
9309 &sp,
9310 &hist,
9311 &mut sctr,
9312 "restore-full-cover",
9313 ) {
9314 Ok(t) => t,
9315 // nothing has been mutated on this shape — hand the carrier back and let
9316 // the hit serve PLAIN (the banked pre-lane path).
9317 Err(err) => {
9318 return fail(cache, format!("boundary token draw failed: {err}"));
9319 }
9320 }
9321 } else {
9322 argmax(boundary_logits) as u32
9323 });
9324 }
9325 Ok(SpecSession {
9326 cache,
9327 scratch,
9328 committed,
9329 last_h: last_h_dev,
9330 next_pred,
9331 sctr,
9332 uctr: 0,
9333 draft_ctx: None,
9334 pending_tok: None,
9335 // Stable-boundary capture from the split feed above (None on the legacy shape):
9336 // a restored session previously parked WITHOUT a checkpoint, so the next turn's
9337 // affinity probe declined ("no turn checkpoint retained") and the conversation
9338 // fell back to the frozen prefix entry forever.
9339 turn_ckpt: restored_turn_ckpt,
9340 telem: SpecTelemetryCounters::default(),
9341 capture_at: None,
9342 boundary_captures,
9343 ckpt_at: None,
9344 capture_disabled: false,
9345 })
9346 }
9347
9348 /// SESSION-AFFINITY REWIND (lane/session-affinity, 2026-08-05): roll `sess` back to its
9349 /// retained prompt-end checkpoint, so a request whose prompt matches
9350 /// `committed[..rewind_pos()]` exactly can resume there and prime only its own delta.
9351 ///
9352 /// EXACTNESS. After this returns, the session is byte-for-byte the state it was in AT that
9353 /// boundary: full-attn KV truncated to it (append-only, position-addressed), GDN conv/ssm
9354 /// restored from the device copy taken there, draft scratch length reset, `committed`
9355 /// truncated, `last_h` = the boundary's predecessor anchor. That is precisely the state a
9356 /// fresh prime of `committed[..pos]` would have produced, so the following suffix prime and
9357 /// every burst after it are identical to a cold run of the same token stream — the
9358 /// committed-tokens-authoritative contract.
9359 ///
9360 /// `next_pred` and `pending_tok` are CLEARED: both describe generation past the boundary,
9361 /// which the rewind discards. The caller therefore must supply a non-empty suffix (a
9362 /// rewound session cannot serve an empty-suffix continuation burst — there is nothing to
9363 /// continue). The persistent draft graph survives: it bakes only session-stable pointers
9364 /// (the scratch KV, the resident embedding), none of which the rewind moves.
9365 ///
9366 /// The checkpoint is CONSUMED (`turn_ckpt` taken): its snapshot buffers are freed here, and
9367 /// this turn's own prime installs a fresh one at the new prompt end. Returns the position
9368 /// rewound to, or `None` when the session holds no checkpoint (caller: full re-prime).
9369 pub fn spec_rewind_to_checkpoint(
9370 &self,
9371 e: &Engine,
9372 sess: &mut SpecSession,
9373 ) -> Result<Option<usize>, Box<dyn std::error::Error>> {
9374 if sess.turn_ckpt.as_ref().is_some_and(|ckpt| {
9375 !sess.cache.can_rollback(&ckpt.snap, 0) || !sess.scratch.can_rewind_to(ckpt.pos)
9376 }) {
9377 return Err(
9378 "SWA ring rewind checkpoint has been lapped; full re-prime required".into(),
9379 );
9380 }
9381 let Some(ckpt) = sess.turn_ckpt.take() else {
9382 return Ok(None);
9383 };
9384 assert!(
9385 ckpt.pos <= sess.committed.len(),
9386 "checkpoint past committed ({} > {})",
9387 ckpt.pos,
9388 sess.committed.len()
9389 );
9390 // Restore through each layer's owning engine. A single primary-engine rollback is not
9391 // sufficient when the serving cache is stage-owned under cross-device PP.
9392 crate::pp::restore_cache_checkpoint(e, self, None, &mut sess.cache, &ckpt.snap)?;
9393 debug_assert_eq!(
9394 sess.cache.pos, ckpt.pos,
9395 "rollback landed off the checkpoint"
9396 );
9397 sess.scratch.set_len(e, ckpt.pos)?;
9398 sess.committed.truncate(ckpt.pos);
9399 sess.last_h = Some(ckpt.last_h);
9400 sess.next_pred = None;
9401 sess.pending_tok = None;
9402 Ok(Some(ckpt.pos))
9403 }
9404
9405 /// Grow a parked speculative session to `target_cap` and rewind it to its retained turn
9406 /// checkpoint without re-priming the checkpoint prefix.
9407 ///
9408 /// The trunk cache is restored exactly like a plain grown cache: append-only full-attention
9409 /// KV rows come from the parked cache, while recurrent state comes from the checkpoint's
9410 /// owned snapshot. The MTP scratch is also context-linear and its rows below the checkpoint
9411 /// remain authoritative, so they are copied into a fresh larger scratch before its length is
9412 /// truncated. Pointer-baking draft graphs are dropped and recaptured on the next burst.
9413 ///
9414 /// All fallible work completes before `sess` is mutated. A failed allocation or copy leaves
9415 /// the parked session intact, allowing the caller one reclaim-and-retry attempt.
9416 pub fn spec_grow_and_rewind_to_checkpoint(
9417 &self,
9418 e: &Engine,
9419 sess: &mut SpecSession,
9420 target_cap: usize,
9421 ) -> Result<Option<usize>, Box<dyn std::error::Error>> {
9422 if target_cap <= sess.cache.max_ctx {
9423 return self.spec_rewind_to_checkpoint(e, sess);
9424 }
9425 let Some(ckpt) = sess.turn_ckpt.as_ref() else {
9426 return Ok(None);
9427 };
9428 if ckpt.pos == 0 || ckpt.pos > sess.committed.len() {
9429 return Err(format!(
9430 "checkpoint pos {} outside committed length {}",
9431 ckpt.pos,
9432 sess.committed.len(),
9433 )
9434 .into());
9435 }
9436 if ckpt.pos > target_cap {
9437 return Err(format!(
9438 "checkpoint pos {} exceeds grown capacity {target_cap}",
9439 ckpt.pos,
9440 )
9441 .into());
9442 }
9443
9444 let mut grown_cache = crate::pp::new_cache_planned(e, &self.cfg, &self.plan, target_cap)?;
9445 let mut grown_scratch = self.new_mtp_scratch(e, target_cap)?;
9446 crate::pp::restore_cache_checkpoint(
9447 e,
9448 self,
9449 Some(&sess.cache),
9450 &mut grown_cache,
9451 &ckpt.snap,
9452 )?;
9453
9454 if sess.scratch.plane_count() != grown_scratch.plane_count() {
9455 return Err("checkpoint draft plane count mismatch".into());
9456 }
9457 for index in 0..sess.scratch.plane_count() {
9458 let (src, _) = sess.scratch.plane(index);
9459 let (dst, _) = grown_scratch.plane_mut(index);
9460 if ckpt.pos > src.len
9461 || src.kv_dim_k != dst.kv_dim_k
9462 || src.kv_dim_v != dst.kv_dim_v
9463 || src.k_tok_bytes != dst.k_tok_bytes
9464 || src.v_tok_bytes != dst.v_tok_bytes
9465 {
9466 return Err(format!(
9467 "checkpoint draft plane {index} layout mismatch (pos {}, source len {})",
9468 ckpt.pos, src.len,
9469 )
9470 .into());
9471 }
9472 match (&src.ring, dst.ring.as_ref()) {
9473 (Some(sring), Some(_)) => {
9474 // Ring-backed draft plane (step35): `ckpt.pos` is absolute and exceeds the
9475 // physical rows once lapped — same class as the trunk-KV restore panic
9476 // (2026-08-29 warm-turn-at-40k). Copy the aligned live window, rebase.
9477 let (new_base, phys) = sring.restore_plan(ckpt.pos).map_err(|err| {
9478 format!("checkpoint draft plane {index} SWA restore refused: {err}")
9479 })?;
9480 let rows = phys.len();
9481 let kb = rows * src.k_tok_bytes;
9482 let vb = rows * src.v_tok_bytes;
9483 if kb > 0 {
9484 e.copy_u8_range_into(
9485 &mut dst.k,
9486 0,
9487 &src.k,
9488 phys.start * src.k_tok_bytes,
9489 kb,
9490 )?;
9491 }
9492 if vb > 0 {
9493 e.copy_u8_range_into(
9494 &mut dst.v,
9495 0,
9496 &src.v,
9497 phys.start * src.v_tok_bytes,
9498 vb,
9499 )?;
9500 }
9501 dst.ring
9502 .as_mut()
9503 .expect("ring presence checked above")
9504 .apply_rebase(new_base);
9505 if let Some(base_d) = dst.base_d.as_mut() {
9506 e.set_i32_one(base_d, new_base as i32)?;
9507 }
9508 }
9509 (None, None) => {
9510 let kb = ckpt.pos * src.k_tok_bytes;
9511 let vb = ckpt.pos * src.v_tok_bytes;
9512 if kb > 0 {
9513 e.copy_u8_into(&mut dst.k, 0, &src.k, kb)?;
9514 }
9515 if vb > 0 {
9516 e.copy_u8_into(&mut dst.v, 0, &src.v, vb)?;
9517 }
9518 }
9519 _ => {
9520 return Err(format!("checkpoint draft plane {index} ring/flat mismatch").into());
9521 }
9522 }
9523 }
9524 grown_scratch.set_len(e, ckpt.pos)?;
9525 // The old scratch is dropped immediately after publication below. Bound its D2D reads
9526 // first; growth happens once per rewritten turn, outside the decode hot loop.
9527 e.stream().synchronize()?;
9528
9529 let ckpt = sess
9530 .turn_ckpt
9531 .take()
9532 .expect("checkpoint remained present through transactional grow");
9533 let pos = ckpt.pos;
9534 sess.cache = grown_cache;
9535 sess.scratch = grown_scratch;
9536 sess.committed.truncate(pos);
9537 sess.last_h = Some(ckpt.last_h);
9538 sess.next_pred = None;
9539 sess.pending_tok = None;
9540 sess.draft_ctx = None;
9541 debug_assert_eq!(sess.cache.pos, pos, "grown rewind landed off checkpoint");
9542 debug_assert!(
9543 (0..sess.scratch.plane_count()).all(|index| sess.scratch.plane(index).0.len == pos),
9544 "grown draft rewind landed off checkpoint"
9545 );
9546 Ok(Some(pos))
9547 }
9548
9549 /// Commit a carried pending bonus (see SpecSession::pending_tok): one T=1 trunk pass
9550 /// (its logits' argmax becomes next_pred) + the draft-KV fill at the carried anchor —
9551 /// byte-identical to the pre-carry session tail. Required before a non-empty-suffix
9552 /// prime, a sampled turn, or parking a session for pool reuse. No-op without a pending.
9553 /// `sampling` is the sampler of the request that will CONSUME the resulting `next_pred`
9554 /// (lane/sampled-spec-quality): this is a boundary site like any other, so a sampled
9555 /// consumer must get a DRAWN token, not an argmax. Pass `None` from the park/demote
9556 /// callers — a pending only ever exists on the GREEDY tail, and the consumer of a
9557 /// park-time flush is a future request whose sampler is not knowable here (residual
9558 /// named at the pool-resume probe in worker.rs and in SAMPLED-QUALITY.md).
9559 pub fn spec_flush_pending(
9560 &self,
9561 e: &Engine,
9562 sess: &mut SpecSession,
9563 sampling: Option<SpecSampling>,
9564 ) -> Result<(), Box<dyn std::error::Error>> {
9565 sess.cache.ensure_usable("spec_flush_pending")?;
9566 let Some(b) = sess.pending_tok.take() else {
9567 return Ok(());
9568 };
9569 if self.mtp.is_none() {
9570 return Err("pending carry requires an MTP head".into());
9571 }
9572 let n_embd = self.cfg.n_embd as usize;
9573 let (embd_qt, embd_rb) = self.embd.qt_and_row_bytes(n_embd);
9574 let embd_gpu = if spec_host_embd() {
9575 None
9576 } else {
9577 Some(
9578 self.embd_gpu
9579 .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload")),
9580 )
9581 };
9582 let embd_dev = embd_gpu.map(|g| (g, embd_qt, embd_rb));
9583 let pos_b = sess.cache.pos;
9584 sess.scratch.set_len(e, pos_b)?;
9585 let (lg_b, hb) = self.spec_target_step_h(e, b, &mut sess.cache)?;
9586 sess.next_pred = Some(match sampling {
9587 Some(sp) if sp.temp > 0.0 && spec_sampled_boundary_on() => {
9588 // window includes `b` itself: it is committed by this pass, and the pre-lane
9589 // code never counted a boundary token in the penalty history at all.
9590 let hist = pen_window_seed(&sess.committed, &[b], sp.penalty_last_n);
9591 sample_boundary_token(e, &lg_b, &sp, &hist, &mut sess.sctr, "flush-pending")?
9592 }
9593 _ => argmax(&lg_b) as u32,
9594 });
9595 let anchor = sess
9596 .last_h
9597 .as_ref()
9598 .expect("pending carry requires last_h (the predecessor-row anchor)");
9599 self.mtp_kv_fill_all(e, &[b], anchor, pos_b, &mut sess.scratch, embd_dev)?;
9600 sess.last_h = Some(hb);
9601 sess.committed.push(b);
9602 Ok(())
9603 }
9604
9605 /// Solo target feed used only at speculative round boundaries. Step35 serving made its
9606 /// staged batched B=1 graph authoritative, so a speculative session must enter and leave
9607 /// rounds through that same graph. Other model families keep their eager T=1 contract.
9608 fn spec_target_step_h(
9609 &self,
9610 e: &Engine,
9611 token: u32,
9612 cache: &mut Cache,
9613 ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9614 cache.ensure_usable("spec_target_step_h")?;
9615 if !self.sliding_gated_moe_batch_program() && !self.batched_serving_numeric_class() {
9616 return self.decode_step_h(e, token, cache);
9617 }
9618 let pos0 = cache.pos;
9619 let (logits, hidden) = self.decode_step_t_core(e, &[token], pos0, cache, None, None)?;
9620 Ok((e.dtoh(&logits)?, hidden))
9621 }
9622
9623 /// The archs whose LIVE B=1 serving runs the generic BATCHED numeric class (decode_step_batch
9624 /// walk + batched head), so their spec verify must run the SAME class. MoE learned this
9625 /// 2026-08-14 AM (4b777ccc5); the dense hybrid reproduced the identical near-tie flip class
9626 /// the same day on Qwen3.8-27B — eager-class verify logits drift from batched-class serving
9627 /// logits ("1 ULP at layer 2 → 2.3e-1 logit maxdiff at the head"), and the GDN recurrence
9628 /// carries the drift until a near-tie flips deep in generation. One predicate so the five
9629 /// dispatch sites cannot drift apart again.
9630 /// Draft-graph head admissibility (lane/draftcost-moe, 2026-08-20): the capture body
9631 /// (`mtp_head_forward_cap`) supports Dense heads and SOFTMAX device-routed resident-MoE
9632 /// heads. Residency alone is insufficient: Hy3/M3/Step sigmoid routing returns selected
9633 /// experts through a host synchronization, which is capture-illegal. Those heads use the
9634 /// exact eager draft chain until a device-only sigmoid expert program lands. Trunk FFN class
9635 /// is irrelevant — the graph body is the HEAD forward only. One predicate for all three
9636 /// eligibility sites so they cannot drift (the serving numeric-class lesson).
9637 fn mtp_graph_capturable(&self) -> bool {
9638 let sigmoid_router = self.cfg.sigmoid_router().is_some();
9639 for head in self.mtp.iter().chain(self.mtp_extra.iter()) {
9640 let reason = match &head.ffn {
9641 crate::hybrid::Ffn::Dense { .. } => None,
9642 crate::hybrid::Ffn::Moe(mo) if mo.dev_exps.is_none() => {
9643 Some("non-resident MoE MTP head")
9644 }
9645 crate::hybrid::Ffn::Moe(_) if sigmoid_router => {
9646 Some("sigmoid-router MoE MTP head requires host-visible routing")
9647 }
9648 crate::hybrid::Ffn::Moe(_) => None,
9649 };
9650 if let Some(reason) = reason {
9651 static NOTICE: std::sync::Once = std::sync::Once::new();
9652 NOTICE.call_once(|| {
9653 eprintln!(
9654 "[spec] draft graph unavailable: {reason}; eager draft chain engaged"
9655 );
9656 });
9657 return false;
9658 }
9659 }
9660 self.mtp.is_some()
9661 }
9662
9663 fn batched_serving_numeric_class(&self) -> bool {
9664 self.plan
9665 .trunk_operations()
9666 .contains(&memra_gguf::model_plan::OperationKind::GatedDeltaNet)
9667 }
9668
9669 /// The family the MTP verify-graph default was measured on: GatedDeltaNet state layers
9670 /// (a `recur` mixer) together with a routed-MoE FFN — Ornith-1.5-35B-A3B and its kin. The
9671 /// server-side twin of this test is `model_forces_spec_replay` (GatedDeltaNet + MoeMlp);
9672 /// keeping the engine's own version structural rather than name-based means a new
9673 /// checkpoint of the same shape inherits the default, and a different shape does not.
9674 /// pub(crate) since lane/graph-launch-guard-sweep-20260831: `dspark_vg_admission_debt`
9675 /// consults it so the MTP-route pool stops escaping the admission charge.
9676 pub(crate) fn vgraph_family_default(&self) -> bool {
9677 let has_linear = self
9678 .layers
9679 .iter()
9680 .any(|l| matches!(l.mixer, Mixer::Linear(_)));
9681 let has_moe = self
9682 .layers
9683 .iter()
9684 .any(|l| matches!(l.ffn, crate::hybrid::Ffn::Moe(_)));
9685 has_linear && has_moe
9686 }
9687
9688 fn sliding_gated_moe_batch_program(&self) -> bool {
9689 self.uses_sliding_gated_moe_program()
9690 }
9691
9692 fn gemma_batch_program(&self) -> bool {
9693 self.uses_gemma_program()
9694 }
9695
9696 /// One spec-decode turn on a live session. `suffix` = the NEW tokens only (turn N+1's user
9697 /// message rendered through the chat template continuation). Returns (new tokens emitted,
9698 /// drafted, accepted); session.committed grows by suffix + emitted.
9699 pub fn generate_spec_session(
9700 &self,
9701 e: &Engine,
9702 sess: &mut SpecSession,
9703 suffix: &[u32],
9704 max_new: usize,
9705 k: usize,
9706 ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
9707 self.generate_spec_session_sampled(e, sess, suffix, max_new, k, None, None)
9708 }
9709
9710 /// Serve-path sampled spec: routes the burst through the rejection-sampling verify with
9711 /// per-SESSION Philox continuity (sess.sctr/uctr). None = env-driven (CLI) or greedy.
9712 /// Filters (top-k/p/min-p) apply SYMMETRICALLY to draft q and verify p — distribution-exact
9713 /// for the filtered target (feat/filtered-spec).
9714 ///
9715 /// `on_commit` (sse-cadence, 2026-08-05): called with each newly-emitted slice of the
9716 /// output — once right after the prime's first token, then once per round commit — so a
9717 /// streaming caller can flush text at round cadence instead of once per burst. The slices
9718 /// are disjoint, in order, and concatenate to exactly the returned token vec. Emission-
9719 /// timing only: token bytes, session state, and exactness are untouched.
9720 ///
9721 /// The returned bool is a CONTINUE-VERDICT (admission yield, 2026-08-06): `false` ends
9722 /// the burst at the current round boundary, exactly as if `max_new` had been reached —
9723 /// the caller's scheduler regains control without waiting the burst out. Burst size is
9724 /// content-neutral (spec-levers battery), so an early exit moves WHEN the burst returns,
9725 /// never what tokens say. The slice may be EMPTY (a poll-only boundary — round-stream
9726 /// drains and the defensive tail flush can land with nothing new committed).
9727 #[allow(clippy::too_many_arguments)]
9728 #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
9729 pub fn generate_spec_session_sampled(
9730 &self,
9731 e: &Engine,
9732 sess: &mut SpecSession,
9733 suffix: &[u32],
9734 max_new: usize,
9735 k: usize,
9736 sampling: Option<SpecSampling>,
9737 on_commit: Option<&mut dyn FnMut(&[u32]) -> bool>,
9738 ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
9739 self.generate_spec_session_sampled_prime_split(
9740 e, sess, suffix, max_new, k, sampling, None, on_commit,
9741 )
9742 }
9743
9744 /// Serve-only cold-prime segmentation twin. `prime_split` is the same stable boundary the
9745 /// plain worker would honor before entering its sub-floor tokenwise tail; warm continuations
9746 /// pass `None` and stay on the existing zero-prime path.
9747 #[allow(clippy::too_many_arguments)]
9748 #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
9749 pub fn generate_spec_session_sampled_prime_split(
9750 &self,
9751 e: &Engine,
9752 sess: &mut SpecSession,
9753 suffix: &[u32],
9754 max_new: usize,
9755 k: usize,
9756 sampling: Option<SpecSampling>,
9757 prime_split: Option<usize>,
9758 on_commit: Option<&mut dyn FnMut(&[u32]) -> bool>,
9759 ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
9760 self.generate_spec_session_constrained_prime_split(
9761 e,
9762 sess,
9763 suffix,
9764 max_new,
9765 k,
9766 sampling,
9767 None,
9768 prime_split,
9769 on_commit,
9770 )
9771 }
9772
9773 /// `generate_spec_session_sampled` + GRAMMAR (constrained decoding, 2026-08-03): the
9774 /// hook truncates acceptance at the first grammar-illegal token AFTER the exactness
9775 /// verify (grammar is an extra rejection rule, ordering like the batched-verify twins)
9776 /// and replaces an illegal bonus with the MASKED argmax of the target's own verify
9777 /// column — token-identical to constrained plain greedy decode. GREEDY only (the
9778 /// worker routes sampled constrained to plain decode). Acceptance under tight grammars
9779 /// may drop (drafter is unconstrained); that is measured, not hidden.
9780 #[allow(clippy::too_many_arguments)]
9781 #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
9782 pub fn generate_spec_session_constrained(
9783 &self,
9784 e: &Engine,
9785 sess: &mut SpecSession,
9786 suffix: &[u32],
9787 max_new: usize,
9788 k: usize,
9789 sampling: Option<SpecSampling>,
9790 constraint: Option<&mut dyn SpecConstraint>,
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_constrained_prime_split(
9794 e, sess, suffix, max_new, k, sampling, constraint, None, on_commit,
9795 )
9796 }
9797
9798 #[allow(clippy::too_many_arguments)]
9799 #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
9800 pub fn generate_spec_session_constrained_prime_split(
9801 &self,
9802 e: &Engine,
9803 sess: &mut SpecSession,
9804 suffix: &[u32],
9805 max_new: usize,
9806 k: usize,
9807 sampling: Option<SpecSampling>,
9808 constraint: Option<&mut dyn SpecConstraint>,
9809 prime_split: Option<usize>,
9810 on_commit: Option<&mut dyn FnMut(&[u32]) -> bool>,
9811 ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
9812 if constraint.is_some() && sampling.is_some_and(|s| s.temp > 0.0) {
9813 return Err(
9814 "constrained spec decode is greedy-only (worker routes sampled \
9815 constrained to plain decode)"
9816 .into(),
9817 );
9818 }
9819 // PENDING-CARRY entry flush: a carried bonus precedes any new suffix in the sequence,
9820 // so it must commit BEFORE the suffix primes; the sampled path doesn't carry (its
9821 // round-0 accept needs the commit pass's logits). Empty-suffix greedy bursts — the
9822 // serve continuation case — consume the carry in-loop with zero solo passes.
9823 if sess.pending_tok.is_some()
9824 && (!suffix.is_empty() || sampling.is_some_and(|s| s.temp > 0.0))
9825 {
9826 self.spec_flush_pending(e, sess, sampling)?;
9827 }
9828
9829 // FULL_PREC forces the EAGER draft: the graph capture would enclose cuBLASLt f32 GEMV
9830 // (the FloatBf16 else-branches) and a bf16_to_f32 dequant alloc — neither is stream-capture
9831 // safe. Eager rides matmul/matmul_decode_exact, which dequant FloatBf16 on use. (§item 2.)
9832 // Multi-head MTP (mtp_extra non-empty) no longer disqualifies: the chain captures
9833 // per-head graphs (lane/step37-draft-graph-serving-20260830, MEMRA_MTP_CHAIN_GRAPH).
9834 let graph_draft = std::env::var("MEMRA_SPEC_NOGRAPH").is_err()
9835 && !spec_host_embd()
9836 && self.mtp_graph_capturable()
9837 && k + 2 < 96
9838 && !crate::model::full_prec_enabled();
9839 let was_tracking = e.ctx().is_event_tracking();
9840 if graph_draft && was_tracking {
9841 unsafe {
9842 e.ctx().disable_event_tracking();
9843 }
9844 }
9845 let r = self.generate_spec_inner2(
9846 e,
9847 suffix,
9848 max_new,
9849 k,
9850 graph_draft,
9851 Some(sess),
9852 sampling,
9853 constraint,
9854 on_commit,
9855 prime_split,
9856 );
9857 if graph_draft && was_tracking {
9858 unsafe {
9859 e.ctx().enable_event_tracking();
9860 }
9861 }
9862 let (out, d, a) = r?;
9863 Ok((out, d, a))
9864 }
9865
9866 pub fn generate_spec(
9867 &self,
9868 e: &Engine,
9869 prompt: &[u32],
9870 max_new: usize,
9871 k: usize,
9872 ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
9873 // glm5 T-parallel verify door (lane/glm5-tparallel-verify): an hc trunk with a
9874 // loaded DRAFT SOURCE — the embedded MTP head OR the DFlash2 drafter
9875 // (lane/glm5-dflash-draft-src) — routes to the glm5 draft->verify->rollback loop —
9876 // MEMRA_GLM5_SPEC=1 only (default OFF; flag row in FLAGS.md). Unset/0 falls
9877 // through to the standing named refusal below, byte-identical to the pre-lane
9878 // binary. Same fail-closed manifest stance as the generic path: an unqualified
9879 // MtpSpec rewrite refuses before any drafting.
9880 if self.hyper.is_some()
9881 && crate::glm_spec::glm5_spec_on()
9882 && (self.mtp.is_some() || self.glm5_dflash.is_some())
9883 {
9884 if !self.rewrite_allowed(memra_gguf::execution_manifest::RewriteSurface::MtpSpec) {
9885 return Err("speculative rewrite is not qualified for this ModelPlan".into());
9886 }
9887 return self.generate_spec_glm5(e, prompt, max_new, k);
9888 }
9889 self.refuse_hyper("generate_spec")?;
9890 if crate::pp::pp_cuts(self.layers.len()).is_some()
9891 && !self.rewrite_allowed(memra_gguf::execution_manifest::RewriteSurface::Pipeline)
9892 {
9893 return Err("pipeline rewrite is not qualified for speculative decode".into());
9894 }
9895 if !self.rewrite_allowed(memra_gguf::execution_manifest::RewriteSurface::MtpSpec) {
9896 return Err("speculative rewrite is not qualified for this ModelPlan".into());
9897 }
9898 // FULL_PREC forces eager (see generate_spec_session note): CUDA graph capture cannot
9899 // enclose cuBLASLt f32 GEMV or the bf16_to_f32 dequant alloc the FloatBf16 path needs.
9900 // Multi-head MTP no longer disqualifies (chain graphs; see generate_spec_session).
9901 let graph_draft = std::env::var("MEMRA_SPEC_NOGRAPH").is_err()
9902 && !spec_host_embd()
9903 && self.mtp_graph_capturable()
9904 && k + 2 < 96
9905 && !crate::model::full_prec_enabled();
9906 if !graph_draft {
9907 return self
9908 .generate_spec_inner2(e, prompt, max_new, k, false, None, None, None, None, None);
9909 }
9910 let was_tracking = e.ctx().is_event_tracking();
9911 if was_tracking {
9912 unsafe {
9913 e.ctx().disable_event_tracking();
9914 }
9915 }
9916 let r =
9917 self.generate_spec_inner2(e, prompt, max_new, k, true, None, None, None, None, None);
9918 if was_tracking {
9919 unsafe {
9920 e.ctx().enable_event_tracking();
9921 }
9922 }
9923 r
9924 }
9925
9926 #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
9927 #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
9928 fn generate_spec_inner2(
9929 &self,
9930 e: &Engine,
9931 prompt: &[u32],
9932 max_new: usize,
9933 k: usize,
9934 graph_draft: bool,
9935 mut sess: Option<&mut SpecSession>,
9936 sampling: Option<SpecSampling>,
9937 mut constraint: Option<&mut dyn SpecConstraint>,
9938 mut on_commit: Option<&mut dyn FnMut(&[u32]) -> bool>,
9939 prime_split: Option<usize>,
9940 ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
9941 assert!(k >= 1, "k must be >= 1");
9942 // sse-cadence flush cursor: everything in out[..flushed] has been handed to on_commit.
9943 let mut flushed = 0usize;
9944 // admission yield (2026-08-06): on_commit's continue-verdict; false = end the burst
9945 // at the next round boundary (same exit as max_new reached — the session tail runs).
9946 // Initialized by the unconditional post-prime flush below.
9947 let mut keep_going;
9948 let mtp = self
9949 .mtp
9950 .as_ref()
9951 .expect("generate_spec requires an MTP head (nextn_predict_layers>0)");
9952 let n_vocab = self.output.out_features();
9953 // FR-Spec: the draft head may be TRIMMED (fewer rows than n_vocab); the draft argmax runs
9954 // over the draft vocab and the winning index maps through d2t to a TARGET token id.
9955 // Everything downstream (verify/accept/commit) sees target ids only — exactness unchanged.
9956 let d_vocab = mtp
9957 .shared_head_head
9958 .as_ref()
9959 .unwrap_or(&self.output)
9960 .out_features();
9961 if !self.mtp_extra.is_empty() {
9962 if self.plan.draft_source != memra_gguf::model_plan::DraftSourcePlan::Embedded
9963 || self.plan.mtp_blocks.len() != self.mtp_head_count()
9964 {
9965 return Err(
9966 "multi-head MTP requires one embedded canonical block per loaded head".into(),
9967 );
9968 }
9969 // TRIMMED chains (2026-08-27): every head must carry the SAME d2t — the ranking is
9970 // token-frequency and head-independent, and every downstream remap (per-step argmax,
9971 // stream pack, sampled d2t_dev) reads head 0's map, so equality is what makes that
9972 // single map correct for the whole chain. Mixed trimmed/untrimmed is refused.
9973 for (offset, head) in self.mtp_extra.iter().enumerate() {
9974 if head.d2t != mtp.d2t
9975 || head
9976 .shared_head_head
9977 .as_ref()
9978 .unwrap_or(&self.output)
9979 .out_features()
9980 != d_vocab
9981 {
9982 return Err(format!(
9983 "embedded MTP head {} has incompatible draft vocabulary",
9984 offset + 1
9985 )
9986 .into());
9987 }
9988 }
9989 eprintln!(
9990 "[mtp-chain] heads={} policy=step-modulo prefix-replay kv=per-head",
9991 self.mtp_head_count()
9992 );
9993 }
9994 let n_embd = self.cfg.n_embd as usize;
9995 // SESSION MODE: reuse the live cache/scratch, prime only the suffix. `base` = tokens
9996 // already committed (their state is in the caches); 0 = fresh single-shot call.
9997 let session_mode = sess.is_some();
9998 let max_ctx = match sess.as_ref() {
9999 Some(s) => s.cache.max_ctx,
10000 None => prompt.len() + max_new + k + 8,
10001 };
10002 let mut own_cache;
10003 let mut own_scratch;
10004 // PREFIX-CACHE capture request threaded out of the session (lane/spec-prefix-cache):
10005 // (requested split, destination list). Single-shot per burst; fresh calls have none.
10006 let mut sess_capture: Option<(Option<usize>, &mut Vec<SpecBoundaryCapture>)> = None;
10007 // STABLE-BOUNDARY turn-checkpoint request (lane/frspec-multiturn-cache): ABSOLUTE
10008 // committed-length position; consumed one-shot like `capture_at`. None = legacy
10009 // prompt-end capture below.
10010 let mut ckpt_req: Option<usize> = None;
10011 // FAIL-SAFE bit threaded out of the session (see `SpecSession::capture_disabled`).
10012 let mut sess_capture_disabled = false;
10013 let (
10014 cache,
10015 scratch,
10016 mut sess_tail,
10017 mut sess_draft_slot,
10018 mut sess_pending_slot,
10019 sess_ckpt_slot,
10020 sess_telem,
10021 ): (
10022 &mut Cache,
10023 &mut MtpScratch,
10024 Option<(
10025 &mut Vec<u32>,
10026 &mut Option<CudaSlice<f32>>,
10027 &mut Option<u32>,
10028 &mut u32,
10029 &mut u32,
10030 )>,
10031 Option<&mut Option<DraftGraphCtx>>,
10032 Option<&mut Option<u32>>,
10033 Option<&mut Option<SpecCheckpoint>>,
10034 Option<&SpecTelemetryCounters>,
10035 ) = match sess.take() {
10036 Some(sr) => {
10037 let SpecSession {
10038 cache,
10039 scratch,
10040 committed,
10041 last_h,
10042 next_pred,
10043 sctr: s_sctr,
10044 uctr: s_uctr,
10045 draft_ctx,
10046 pending_tok,
10047 turn_ckpt,
10048 telem,
10049 capture_at,
10050 boundary_captures,
10051 ckpt_at,
10052 capture_disabled,
10053 } = sr;
10054 sess_capture_disabled = *capture_disabled;
10055 sess_capture = Some((capture_at.take(), boundary_captures));
10056 ckpt_req = ckpt_at.take();
10057 (
10058 cache,
10059 scratch,
10060 Some((committed, last_h, next_pred, s_sctr, s_uctr)),
10061 Some(draft_ctx),
10062 Some(pending_tok),
10063 Some(turn_ckpt),
10064 Some(telem),
10065 )
10066 }
10067 None => {
10068 // STAGE-OWNED KV (lane/pp2-spec 2026-08-06) — see `new_session`. Door shut =
10069 // `Cache::new` verbatim.
10070 own_cache = crate::pp::new_cache_planned(e, &self.cfg, &self.plan, max_ctx)?;
10071 // Persistent scratch = max_ctx rows (~2KB/token quantized).
10072 own_scratch = self.new_mtp_scratch(e, max_ctx)?;
10073 (
10074 &mut own_cache,
10075 &mut own_scratch,
10076 None,
10077 None,
10078 None,
10079 None,
10080 None,
10081 )
10082 }
10083 };
10084 cache.ensure_usable("generate_spec")?;
10085 if scratch.plane_count() != self.mtp_head_count() {
10086 return Err(format!(
10087 "MTP scratch/head count mismatch ({}/{})",
10088 scratch.plane_count(),
10089 self.mtp_head_count()
10090 )
10091 .into());
10092 }
10093 let base = cache.pos;
10094 // PENDING-CARRY consume (2026-08-01): a carried bonus reaches here only on the
10095 // empty-suffix GREEDY continuation path (generate_spec_session_sampled flushed every
10096 // other case). It enters the round loop as round-0's pending — verify col 0 — exactly
10097 // like a mid-burst full-accept boundary: no init feed, no tail commit pass.
10098 let carried_pending: Option<u32> = sess_pending_slot.as_mut().and_then(|s| s.take());
10099 // PERSISTENT DRAFT KV (the only mode since 2026-07-08 — the legacy round-local scratch,
10100 // MEMRA_SPEC_KVLOCAL, measured -35 acceptance pts on the 27B p3 sweep and was removed;
10101 // acceptance-only — exactness is verify's job either way).
10102 // HIDDEN-PAIRING CONVENTION (DEFAULT = predecessor-row, 2026-07-04 — the 27B acceptance
10103 // unlock, +16pts): the MTP head is TRAINED on rows pairing token x_p with the trunk
10104 // hidden of its PREDECESSOR h_{p-1} (the reference engine's mtp_update shifts the target
10105 // hiddens right by one; its draft step 0 feeds (id_last, TRUE hidden of the row id_last
10106 // was sampled from)). memra's historical convention paired SAME-ROW (x_p, h_p) in the fill
10107 // and seeded chain step 0 through an extra MTP pass on a duplicated token (the
10108 // pseudo-seed) — measured 27B p2 K=3 acceptance 0.569 vs 0.731, p3 0.445 vs 0.63+, and
10109 // the chain steps j>=1 were already predecessor-shaped, so ONLY the fill + step-0 seed
10110 // move. The fill shifts by one and the chain seeds from the predecessor's true hidden
10111 // DIRECTLY (vh_seed / vx[j-1]) — the pseudo pass disappears (one MTP-block pass saved
10112 // per round on top of the acceptance win). Draft-quality-only: exactness stays the
10113 // verify's job either way. (The legacy same-row pairing seam, MEMRA_SPEC_HSAME, and its
10114 // pseudo-seed passes were removed 2026-07-08 — predecessor pairing won by +16 acc pts;
10115 // the legacy round-local scratch, MEMRA_SPEC_KVLOCAL, went with it.)
10116 // REPLAY-FREE PARTIAL ACCEPT (default, 2026-07-03): partial rounds keep the verify's own
10117 // bit-identical committed-prefix state (KV truncate + recur rebuild from the VerifyCkpt)
10118 // and leave the bonus PENDING — no duplicate trunk pass (profiled ~0.54 extra full weight
10119 // reads/round at long ctx). MEMRA_SPEC_REPLAY=1 restores the legacy rollback+replay (A/B
10120 // + fallback seam).
10121 // Qwen35-MoE replay pin LIFTED (lane/draftcost-moe, 2026-08-20). The pin's stated
10122 // bar — the retained verify-state commit proven equivalent to sequential serving —
10123 // was waiting on this arch running the serving batched verify class, which the
10124 // t-parallel admission (this lane, increment 1) provided: the VerifyCkpt the
10125 // replay-free commit consumes is now produced by the SAME serving-class verify that
10126 // qualified dense qwen35 on 2026-08-15 (where the per-round duplicate replay
10127 // measured 69 -> 30 tok/s). Qualification receipts (run-spec K=1..8 both arms,
10128 // 8-prompt replay-vs-replay-free canary, long-prompt cell):
10129 // research/draftcost-moe-20260820/RECEIPTS.md. MEMRA_SPEC_REPLAY=1 stays the
10130 // rollback + A/B seam.
10131 let spec_replay = spec_replay_env_enabled();
10132 if constraint.is_some() && spec_replay {
10133 return Err(
10134 "constrained spec decode does not support MEMRA_SPEC_REPLAY=1 \
10135 (legacy replay commits an unmasked bonus)"
10136 .into(),
10137 );
10138 }
10139 // TRUE-HIDDEN REFRESH (default in persistent-draft-KV mode): every round overwrites the
10140 // committed positions' scratch entries from the verify's exact hiddens (mtp_kv_fill batch)
10141 // instead of keeping chain-approximate entries. MEMRA_SPEC_NOREFRESH=1 = legacy (A/B seam).
10142 let refresh = std::env::var("MEMRA_SPEC_NOREFRESH").is_err();
10143 if !refresh && !self.mtp_extra.is_empty() {
10144 return Err("multi-head MTP requires exact accepted-prefix refresh".into());
10145 }
10146
10147 // prime: BATCHED cache prime (prime_cache — the measured #1 e2e gap: tokenwise primed at
10148 // ~102/38 tok/s vs the engine's ~2000-5900 tok/s batched prefill). prime_cache returns the
10149 // full pre-output_norm hidden stack [T, n_embd], which IS prompt_h (the persistent-draft-KV
10150 // mtp_kv_fill input) — no per-token collection needed. Prompts below PRIME_MIN_T, and
10151 // MEMRA_PRIME_TOKENWISE=1, and frozen Hy3 CPU/GPU expert splits take the tokenwise
10152 // decode_step_h loop. The latter avoids transient GPU staging of the spilled expert bank.
10153 // EMPTY-SUFFIX CONTINUATION (serve bursts): a session turn with NO new tokens resumes
10154 // generation exactly where the last turn stopped — no prime at all. The stashed
10155 // `next_pred` plays prime_logits' role: it is the token produced from the logits after
10156 // committed.last() by the same rule this entry applies to a cold prime's last row —
10157 // an argmax when greedy, a `sample_boundary_token` draw when sampled (the burst tail,
10158 // or `spec_session_from_restored` for a converted prefix-cache hit, did the drawing
10159 // where the sampler and the session's Philox counters were live). `last_h` seeds the
10160 // predecessor pairing below. Fresh calls and non-empty suffixes take the normal path.
10161 let continuation = prompt.is_empty();
10162 if continuation {
10163 assert!(session_mode, "empty prompt requires a session");
10164 assert!(
10165 sess_tail
10166 .as_ref()
10167 .is_some_and(|(c, lh, np, _, _)| !c.is_empty()
10168 && lh.is_some()
10169 && (np.is_some() || carried_pending.is_some())),
10170 "empty-suffix continuation needs a primed session (committed + last_h + next_pred|pending)"
10171 );
10172 }
10173 let mut prime_logits;
10174 let mut prompt_h: Option<CudaSlice<f32>> = None;
10175 let t_prime = std::time::Instant::now();
10176 let batched_prime = !continuation
10177 && prompt.len() >= crate::hybrid_forward::PRIME_MIN_T
10178 && std::env::var("MEMRA_PRIME_TOKENWISE").is_err()
10179 && !e.frozen_cpu_experts_prefer_tokenwise_prime();
10180 let prime_split = prime_split.filter(|&split| split > 0 && split < prompt.len());
10181 if prime_split.is_some() && continuation {
10182 return Err("spec prime split requires a non-empty prime".into());
10183 }
10184 // STABLE-BOUNDARY TURN CHECKPOINT stop (lane/frspec-multiturn-cache, 2026-08-21):
10185 // the worker's `ckpt_at` request, ABSOLUTE -> prompt-relative. On WARM bursts
10186 // (base != 0, an affinity-rewound or pool-resumed session priming its own delta)
10187 // this is the only stop; on COLD bursts it usually coincides with `prime_split`
10188 // (both are the plain tier's stable pre-generation boundary). A boundary the prime
10189 // cannot honor (outside this prime's range) silently drops the capture — the
10190 // turn_ckpt convention: the next turn re-primes in full, never a wrong resume.
10191 let ckpt_rel = if continuation {
10192 None
10193 } else {
10194 ckpt_req
10195 .and_then(|abs| abs.checked_sub(base))
10196 .filter(|&r| r > 0 && r < prompt.len())
10197 };
10198 // Prime stops, ordered: each is a boundary the prime halts at so the in-place GDN
10199 // conv/ssm state can be snapshotted there (the only moment it exists). One stop =
10200 // the legacy single-split program, byte-for-byte.
10201 let mut stops: Vec<usize> = Vec::new();
10202 for b in [prime_split, ckpt_rel].into_iter().flatten() {
10203 if !stops.contains(&b) {
10204 stops.push(b);
10205 }
10206 }
10207 stops.sort_unstable();
10208 // Captured at the ckpt stop, installed into the session slot post-prime (replacing
10209 // the legacy prompt-end capture). Some(None) = capture attempted and failed -> the
10210 // slot is cleared (a stale checkpoint would rewind to the WRONG boundary).
10211 let mut ckpt_early: Option<Option<SpecCheckpoint>> = None;
10212 if continuation {
10213 prime_logits = Vec::new();
10214 } else if !stops.is_empty() {
10215 if let Some(&first) = stops.first()
10216 && prime_split == Some(first)
10217 && first < crate::hybrid_forward::PRIME_MIN_T
10218 {
10219 return Err(format!(
10220 "spec prime split {first} is below PRIME_MIN_T {}",
10221 crate::hybrid_forward::PRIME_MIN_T,
10222 )
10223 .into());
10224 }
10225 // Mirror the plain worker's boundary stops exactly. Each segment is a
10226 // request-level prime (`queued_after` keeps Step35 arm selection independent of
10227 // the stops — tick-seg law); a segment below PRIME_MIN_T (and the final tail
10228 // under MEMRA_PRIME_TOKENWISE) takes the same eager tokenwise continuation as
10229 // prefill_tick. Retain every hidden row so the draft scratch fill remains one
10230 // coherent prompt.
10231 let mut h_all = e.uninit(prompt.len() * n_embd)?;
10232 prime_logits = Vec::new();
10233 let mut prev = 0usize;
10234 for seg_end in stops.iter().copied().chain(std::iter::once(prompt.len())) {
10235 if seg_end <= prev {
10236 continue;
10237 }
10238 let seg = &prompt[prev..seg_end];
10239 let is_final = seg_end == prompt.len();
10240 let batched_seg = seg.len() >= crate::hybrid_forward::PRIME_MIN_T
10241 && (!is_final
10242 || (std::env::var("MEMRA_PRIME_TOKENWISE").is_err()
10243 && !e.frozen_cpu_experts_prefer_tokenwise_prime()));
10244 if batched_seg {
10245 let (l, _, h_seg) =
10246 self.prime_cache(e, seg, &mut *cache, prompt.len() - seg_end)?;
10247 e.copy_into(&mut h_all, prev * n_embd, &h_seg, seg.len() * n_embd)?;
10248 prime_logits = l;
10249 } else {
10250 for (i, &tok) in seg.iter().enumerate() {
10251 let (l, h) = self.decode_step_h(e, tok, &mut *cache)?;
10252 e.copy_into(&mut h_all, (prev + i) * n_embd, &h, n_embd)?;
10253 prime_logits = l;
10254 }
10255 }
10256 prev = seg_end;
10257 if is_final {
10258 break;
10259 }
10260 debug_assert_eq!(cache.pos, base + seg_end, "prime stop landed off boundary");
10261 // PREFIX-CACHE BOUNDARY CAPTURE (lane/spec-prefix-cache): the GDN conv/ssm
10262 // states are about to be advanced in place by the next segment, so this is
10263 // the ONLY moment the boundary's recurrent state exists. Capture iff the
10264 // worker requested exactly this stop (cold sessions only — `capture_at` is
10265 // never armed warm). A failed snapshot is silent (turn_ckpt convention) —
10266 // publication is an optimization, never a correctness dependency.
10267 if base == 0
10268 && let Some((requested, slot)) = sess_capture.as_mut()
10269 {
10270 // Publish at the requested miss-LCP stop (the shared-prefix class)
10271 // AND at the stable-boundary stop (the next-turn re-render class,
10272 // lane/frspec-multiturn-cache) — the same boundary set the plain
10273 // prefill tick learns. Without the second entry, the turn after a
10274 // cold re-park could only hit the OLDER lcp entry (the measured
10275 // one-turn transient: t3 restored 607 of 24122 while the plain arm
10276 // rewound to 15222). Dedupe is the worker sweep's has_key.
10277 if (*requested == Some(seg_end) || ckpt_rel == Some(seg_end))
10278 && let Ok(snap) = cache.snapshot(e)
10279 {
10280 slot.push(SpecBoundaryCapture {
10281 snap,
10282 pos: seg_end,
10283 logits: prime_logits.clone(),
10284 // rows [0..seg_end) of h_all are primed — the following
10285 // segments append, never overwrite.
10286 last_h: capture_boundary_hidden(e, &h_all, seg_end, n_embd),
10287 latent_tails: Vec::new(),
10288 });
10289 }
10290 }
10291 // SESSION-AFFINITY TURN CHECKPOINT at the STABLE boundary (see `ckpt_at`):
10292 // same snapshot mechanics, installed post-prime in place of the prompt-end
10293 // capture the re-render class always diverged below.
10294 if ckpt_rel == Some(seg_end) {
10295 let anchor: Result<CudaSlice<f32>, Box<dyn std::error::Error>> =
10296 e.uninit(n_embd).and_then(|mut a| {
10297 e.copy_view_into(
10298 &mut a,
10299 0,
10300 &h_all.slice((seg_end - 1) * n_embd..seg_end * n_embd),
10301 n_embd,
10302 )?;
10303 Ok(a)
10304 });
10305 ckpt_early = Some(match (cache.snapshot(e), anchor) {
10306 (Ok(snap), Ok(last_h)) => Some(SpecCheckpoint {
10307 snap,
10308 pos: base + seg_end,
10309 last_h,
10310 }),
10311 _ => None,
10312 });
10313 }
10314 }
10315 if std::env::var("MEMRA_SPEC_STATS").as_deref() == Ok("1") {
10316 eprintln!(
10317 "[spec-prime] stops={stops:?} tail={}",
10318 prompt.len() - stops.last().copied().unwrap_or(0)
10319 );
10320 }
10321 prompt_h = Some(h_all);
10322 } else if batched_prime {
10323 let (l, _h_seed, hiddens) = self.prime_cache(e, prompt, &mut *cache, 0)?;
10324 prime_logits = l;
10325 prompt_h = Some(hiddens);
10326 } else {
10327 prime_logits = Vec::new();
10328 prompt_h = Some(e.uninit(prompt.len() * n_embd)?);
10329 for (i, &tok) in prompt.iter().enumerate() {
10330 let (l, h) = self.spec_target_step_h(e, tok, &mut *cache)?;
10331 if let Some(ph) = prompt_h.as_mut() {
10332 e.copy_into(ph, i * n_embd, &h, n_embd)?;
10333 }
10334 prime_logits = l;
10335 }
10336 }
10337 e.stream().synchronize()?;
10338 // PREFIX-CACHE SEED CAPTURE (lane/spec-prefix-cache): boundary == prompt end (the seed
10339 // case — no shared-prefix split, publish the whole prompt). The prime just finished, so
10340 // cache.pos == base + prompt.len() and the recurrent state IS the boundary state;
10341 // prime_logits are the boundary logits. Cold sessions only (base == 0) — same law as
10342 // prime_split. The mid-prompt capture above already consumed the request if it matched.
10343 if !continuation
10344 && base == 0
10345 && let Some((requested, slot)) = sess_capture.as_mut()
10346 && *requested == Some(prompt.len())
10347 && slot.is_empty()
10348 {
10349 debug_assert_eq!(cache.pos, prompt.len(), "seed capture off prompt end");
10350 if let Ok(snap) = cache.snapshot(e) {
10351 slot.push(SpecBoundaryCapture {
10352 snap,
10353 pos: prompt.len(),
10354 logits: prime_logits.clone(),
10355 last_h: prompt_h
10356 .as_ref()
10357 .map(|ph| capture_boundary_hidden(e, ph, prompt.len(), n_embd))
10358 .unwrap_or_default(),
10359 latent_tails: Vec::new(),
10360 });
10361 }
10362 }
10363 // Harness timing contract (see crate::PRIME_NANOS): gen-only throughput without the
10364 // prime-subtraction hack.
10365 crate::PRIME_NANOS.store(
10366 t_prime.elapsed().as_nanos() as u64,
10367 std::sync::atomic::Ordering::Relaxed,
10368 );
10369
10370 let (embd_qt, embd_rb) = self.embd.qt_and_row_bytes(n_embd);
10371 // Resident table is fastest when it fits. Large spill deployments can preserve that HBM
10372 // for expert-cache slots and gather only the exact rows needed by MTP/verify from host.
10373 let host_embd = spec_host_embd();
10374 let embd_gpu = if host_embd {
10375 None
10376 } else {
10377 Some(
10378 self.embd_gpu
10379 .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload")),
10380 )
10381 };
10382 let embd_dev = embd_gpu.map(|g| (g, embd_qt, embd_rb));
10383 if host_embd {
10384 eprintln!(
10385 "[spec] host-row embedding: {} bytes kept off HBM",
10386 self.embd.raw.len()
10387 );
10388 }
10389 let mut out: Vec<u32> = Vec::with_capacity(max_new);
10390 let mut total_drafted = 0usize;
10391 let mut total_accepted = 0usize;
10392
10393 // --- SAMPLER FIRST (lane/sampled-spec-quality, 2026-08-19) ---
10394 // The sampler config, the session's Philox counters and the penalty window are parsed
10395 // HERE, above the boundary-token selection, because the boundary token must be drawn
10396 // from the sampler the request asked for. Pre-lane this block sat ~50 lines BELOW the
10397 // selection, which is the whole mechanical reason the boundary token was an argmax:
10398 // the sampler state was not in scope yet. Nothing here depends on the round loop, so
10399 // moving it up is a pure reordering for greedy (`sampled == false` ⇒ every branch
10400 // below takes the argmax path it always took).
10401 // --- SAMPLED SPEC (MEMRA_SPEC_TEMP>0, research/sampled-spec-impl-map.md): rejection-
10402 // sampling verify (Leviathan/Chen) — accept draft x at u < p(x)/q(x), resample from
10403 // norm(max(0,p-q)) on reject, bonus sampled from p on full accept. Counter-based Philox
10404 // everywhere (seed, event) -> reproducible. temp==0/unset = the greedy path, untouched.
10405 let sp = sampling.unwrap_or_else(|| SpecSampling {
10406 temp: std::env::var("MEMRA_SPEC_TEMP")
10407 .ok()
10408 .and_then(|v| v.parse().ok())
10409 .unwrap_or(0.0),
10410 seed: std::env::var("MEMRA_SEED")
10411 .ok()
10412 .and_then(|v| v.parse().ok())
10413 .unwrap_or(42),
10414 top_k: std::env::var("MEMRA_TOP_K")
10415 .ok()
10416 .and_then(|v| v.parse().ok())
10417 .unwrap_or(0),
10418 top_p: std::env::var("MEMRA_TOP_P")
10419 .ok()
10420 .and_then(|v| v.parse().ok())
10421 .unwrap_or(1.0),
10422 min_p: std::env::var("MEMRA_MIN_P")
10423 .ok()
10424 .and_then(|v| v.parse().ok())
10425 .unwrap_or(0.0),
10426 penalty_last_n: std::env::var("MEMRA_PENALTY_LAST_N")
10427 .ok()
10428 .and_then(|v| v.parse().ok())
10429 .unwrap_or(0),
10430 penalty_repeat: std::env::var("MEMRA_PENALTY_REPEAT")
10431 .ok()
10432 .and_then(|v| v.parse().ok())
10433 .unwrap_or(1.0),
10434 penalty_freq: std::env::var("MEMRA_PENALTY_FREQ")
10435 .ok()
10436 .and_then(|v| v.parse().ok())
10437 .unwrap_or(0.0),
10438 penalty_present: std::env::var("MEMRA_PENALTY_PRESENT")
10439 .ok()
10440 .and_then(|v| v.parse().ok())
10441 .unwrap_or(0.0),
10442 });
10443 let (sp_temp, sp_seed) = (sp.temp, sp.seed);
10444 let sampled = sp_temp > 0.0;
10445 // Counters resume from the session (burst continuity: randomness must never repeat
10446 // across generate_spec_session calls); one-shot callers start at (0,0). Read through
10447 // sess_tail — `sess` was take()n into it above, so sess.as_ref() here is always None.
10448 let mut sctr: u32 = sess_tail.as_ref().map(|(_, _, _, s, _)| **s).unwrap_or(0);
10449 let mut uctr: u32 = sess_tail.as_ref().map(|(_, _, _, _, u)| **u).unwrap_or(0);
10450 // Penalties (v2.1): applied to COPIES of q rows and p columns symmetrically (exactness
10451 // for the penalized+filtered target). History = generated tokens, host-tracked window.
10452 let pen_on = sampled
10453 && sp.penalty_last_n > 0
10454 && (sp.penalty_repeat != 1.0 || sp.penalty_freq != 0.0 || sp.penalty_present != 0.0);
10455 // SESSION-SPANNING PENALTY WINDOW (Item 2). Pre-lane this was
10456 // `prompt.iter().rev().take(64).rev()` — the BURST's suffix slice — so a continuation
10457 // burst (the majority of a stream's tokens, and ALL of a converted cache hit's) started
10458 // with an EMPTY penalty history and the client's repetition/frequency/presence penalties
10459 // silently reset at every burst boundary. The window now spans `committed ++ prompt`,
10460 // which is what the API contract says and what the plain sampler's own `history` does.
10461 // Byte-identical to the pre-lane seed for a cold turn-1 burst at the default window.
10462 let mut pen_hist: Vec<u32> = if pen_on {
10463 let sess_hist: &[u32] = if spec_pen_session_on() {
10464 sess_tail
10465 .as_ref()
10466 .map(|(c, ..)| c.as_slice())
10467 .unwrap_or(&[])
10468 } else {
10469 &[] // MEMRA_SPEC_PEN_SESSION=0: pre-lane burst-local window
10470 };
10471 pen_window_seed(sess_hist, prompt, sp.penalty_last_n)
10472 } else {
10473 Vec::new()
10474 };
10475 // First generated token = the BOUNDARY token: greedy takes the argmax of the prompt's
10476 // last logits (== greedy's first token, byte-contract); SAMPLED draws it from the
10477 // request's own filtered/penalized target through the session's Philox stream
10478 // (`sample_boundary_token`, lane/sampled-spec-quality Item 1 — pre-lane this was an
10479 // argmax in both regimes, so ~1 token per burst of a sampled stream was greedy).
10480 // Emit it, then FEED it to establish the loop invariant below.
10481 // PENDING-CARRY: the carried bonus was already emitted by the LAST burst — it becomes
10482 // last_token WITHOUT re-emission, and round 0 consumes it as pending (no init feed).
10483 // CONSTRAINED entry rules: the first emitted token is the MASKED argmax of the
10484 // prompt's last logits (plain constrained-greedy identity); a continuation without
10485 // a carried pending would emit an UNMASKED stashed next_pred — refused loudly (the
10486 // worker never resumes constrained sessions from the pool, so this cannot fire).
10487 if let Some(c) = constraint.as_deref_mut() {
10488 if continuation && carried_pending.is_none() {
10489 return Err("constrained spec continuation requires a carried pending \
10490 (pool resume is unconstrained-only)"
10491 .into());
10492 }
10493 if !continuation {
10494 c.mask_logits(&mut prime_logits)
10495 .map_err(|e2| format!("constraint: {e2}"))?;
10496 }
10497 }
10498 let mut last_token = if let Some(b) = carried_pending {
10499 b
10500 } else if continuation {
10501 // A continuation's boundary token was DRAWN by the burst that stashed it (the
10502 // session tail below), or by `spec_session_from_restored` for a converted
10503 // prefix-cache hit — in both cases from the correct logits row with this same
10504 // session's Philox stream, which is why it can be consumed here as-is.
10505 sess_tail.as_ref().unwrap().2.unwrap()
10506 } else if sampled && constraint.is_none() && spec_sampled_boundary_on() {
10507 sample_boundary_token(e, &prime_logits, &sp, &pen_hist, &mut sctr, "cold-prime")?
10508 } else {
10509 // greedy (byte contract), the rollback door, or constrained (masked-argmax
10510 // identity — the worker routes sampled+constrained to the plain path, and this
10511 // function refuses the combination outright above).
10512 argmax(&prime_logits) as u32
10513 };
10514 if pen_on {
10515 // The boundary token is a GENERATED token: the plain sampler `accept()`s every
10516 // emitted token into its penalty history, and pre-lane the burst's first token
10517 // was invisible to penalties forever (never pushed, and never in `committed`
10518 // until this burst's tail). Covers the carry/continuation seeds too — neither is
10519 // in `committed` yet.
10520 pen_hist.push(last_token);
10521 }
10522 if carried_pending.is_none() {
10523 out.push(last_token);
10524 // grammar advances with every emitted token (carried pendings were consumed
10525 // by the burst that emitted them).
10526 if let Some(c) = constraint.as_deref_mut() {
10527 c.consume(last_token)
10528 .map_err(|e2| format!("constraint: {e2}"))?;
10529 }
10530 }
10531 if continuation {
10532 // draft-KV invariant: entries [0..base) are the session's exact fills; truncate any
10533 // overhang so the chain's first append lands at slot base (== committed.len()).
10534 scratch.set_len(e, base)?;
10535 }
10536 // sse-cadence: hand the caller every not-yet-flushed token (disjoint in-order slices
10537 // concatenating to the full `out`). Called after the prime's first token and after each
10538 // round commit — emission timing only, token bytes untouched. The slice may be EMPTY
10539 // (poll-only boundary: zero-round folds commit nothing new); returns the caller's
10540 // continue-verdict (admission yield, 2026-08-06) — false ends the burst at this round.
10541 #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
10542 fn flush_commit(
10543 cb: &mut Option<&mut dyn FnMut(&[u32]) -> bool>,
10544 out: &[u32],
10545 flushed: &mut usize,
10546 ) -> bool {
10547 if let Some(f) = cb.as_mut() {
10548 let keep = f(&out[*flushed..]);
10549 *flushed = out.len();
10550 keep
10551 } else {
10552 true
10553 }
10554 }
10555 keep_going = flush_commit(&mut on_commit, &out, &mut flushed);
10556 // INVARIANT at loop top: `last_token` is the most-recently-committed/emitted token, its
10557 // KV+recur state IS in `cache` (cache.pos = position right AFTER last_token), `last_pred`
10558 // is the greedy ARGMAX of the logits that predict the token FOLLOWING last_token, and
10559 // `h_seed` = last_token's pre-output_norm hidden. Establish it by feeding last_token once
10560 // (mirrors plain greedy). DEVICE-ARGMAX lever: the accept walk only ever consumes the
10561 // argmax of those logits — never the full vector — so a host u32 replaces the Vec<f32>.
10562 // Trimmed heads: q lives on the trimmed vocab; accept gathers use the TRIMMED index and
10563 // the residual scatters q into target-id space (q=-inf off-trim — the head cannot propose
10564 // those, so their residual mass is p(x), correct by construction).
10565 let d2t_dev: Option<CudaSlice<u32>> = if sampled || crate::spec::spec_stream() {
10566 match &mtp.d2t {
10567 Some(map) => Some(e.htod_u32_v(map)?),
10568 None => None,
10569 }
10570 } else {
10571 None
10572 };
10573 let mut q_full_buf: Option<CudaSlice<f32>> = None;
10574 // host Philox4x32-10 accept-test uniforms: module fn `host_u01` (shared with the
10575 // dspark sampled-admission walk); byte-identical to the closure it replaces.
10576 let mut draft_logits: Vec<CudaSlice<f32>> = Vec::new(); // retained head logits (q), per slot
10577 let mut draft_stats: Vec<(f32, f32, f32)> = Vec::new(); // (row_max, th_e, z_e) per slot
10578 let mut perturb_buf: Option<CudaSlice<f32>> = None; // gumbel scratch (max(n_vocab,d_vocab))
10579 let mut sample_tok = e.alloc_u32_zeroed(1)?; // residual/bonus sample out
10580 let mut col_buf: Option<CudaSlice<f32>> = None; // materialized verify column
10581 let mut pen_hist_d: Option<CudaSlice<u32>> = None;
10582 let mut pcol_buf: Option<CudaSlice<f32>> = None; // penalized p-column scratch
10583 // MEMRA_SPEC_SETUP_TRACE=1 (diagnostics): per-call wall decomposition of the burst
10584 // SETUP + TAIL segments (the round loop's internals are MEMRA_SPEC_PHASE's job) —
10585 // built to pin the serve per-burst fixed cost (research/spec-serving-20260801).
10586 let setup_trace = std::env::var("MEMRA_SPEC_SETUP_TRACE").as_deref() == Ok("1");
10587 let t_ent = std::time::Instant::now();
10588
10589 // SESSION-AFFINITY TURN CHECKPOINT (lane/session-affinity, 2026-08-05): capture the
10590 // PROMPT-END boundary state so a LATER turn can rewind here and re-prime only its own
10591 // delta instead of the whole conversation. See `SpecCheckpoint` for why this boundary is
10592 // the one that matters (a history-rewriting client mutates what the session GENERATED,
10593 // so the next turn's prompt agrees with this one up to exactly here).
10594 //
10595 // WHERE — AND WHY THIS EXACT LINE. Right after the trunk prime, BEFORE the init feed
10596 // (`decode_step_h(last_token)`) and before round 0: the last instant at which the caches
10597 // hold exactly `base + prompt.len()` rows and nothing generated.
10598 //
10599 // This was WRONG in the first cut of this lane: the capture sat after the draft-KV fill,
10600 // which is also after the init feed, so `cache.pos` was `base + prompt.len() + 1` — the
10601 // boundary included the FIRST GENERATED TOKEN. That token is the first thing inside the
10602 // `<think>` block the client strips, so every later turn's diff diverged exactly one
10603 // token below the checkpoint and affinity declined 100% of the time. Measured on the
10604 // owner regime: "history diverged at 12233 of checkpoint 12234". The off-by-one made the
10605 // whole mechanism inert while looking, from the outside, like a working
10606 // correctness-declines-safely path — hence the decline log carries the offsets.
10607 //
10608 // The full-attn planes are `len`-truncatable so the snapshot copies only the GDN conv/ssm
10609 // state (the reason a spec session could not rewind before). The draft scratch needs no
10610 // copy: rows below the boundary are rewritten by the next turn's own fill.
10611 //
10612 // WHEN: non-empty prime only. An empty-suffix continuation burst adds no prompt boundary
10613 // (its "prompt end" IS the previous checkpoint's, already held), so it keeps the existing
10614 // checkpoint rather than replacing it with a strictly worse one.
10615 //
10616 // FAILURE IS SILENT BY DESIGN: on a VRAM-tight rig the snapshot alloc can fail. That
10617 // costs the NEXT turn its rewind (it re-primes fully, today's behavior) and must never
10618 // fail the burst that is already running — so the error is swallowed, loud only under
10619 // MEMRA_DEBUG_SPEC.
10620 //
10621 // STABLE-BOUNDARY OVERRIDE (lane/frspec-multiturn-cache, 2026-08-21): the prompt-end
10622 // posture above was DISPROVED for the think-posture template class — the prompt's own
10623 // tail is the live generation header (`<|im_start|>assistant\n<think>\n`) that the
10624 // next turn's re-render replaces, so the diff diverged a couple tokens BELOW the
10625 // checkpoint and affinity declined 100% of multi-turn agent traffic (the same class
10626 // the plain tier fixed on 2026-08-09 via `plain_checkpoint_boundary`; the port to the
10627 // spec tier is this lane). When the worker armed `ckpt_at`, the capture happened at
10628 // that stop inside the prime above (`ckpt_early`) and is installed here instead;
10629 // capture-attempted-but-failed clears the slot exactly like the legacy arm.
10630 if let Some(slot) = sess_ckpt_slot {
10631 if let Some(early) = ckpt_early {
10632 if early.is_none() && std::env::var("MEMRA_DEBUG_SPEC").is_ok() {
10633 eprintln!(
10634 "[spec] stable-boundary turn checkpoint skipped; \
10635 next turn re-primes in full"
10636 );
10637 }
10638 *slot = early;
10639 } else if !continuation {
10640 let pos = cache.pos;
10641 debug_assert_eq!(
10642 pos,
10643 base + prompt.len(),
10644 "turn checkpoint must sit at the prompt end, before the init feed"
10645 );
10646 let anchor: Result<CudaSlice<f32>, Box<dyn std::error::Error>> =
10647 if let Some(ph) = &prompt_h {
10648 // hidden of the LAST primed row = the predecessor anchor at this
10649 // boundary (exactly what a fresh prime of committed[..pos] leaves in
10650 // last_h, and what the next prime's fill reads for its first row).
10651 let np = prompt.len();
10652 e.uninit(n_embd).and_then(|mut a| {
10653 e.copy_view_into(
10654 &mut a,
10655 0,
10656 &ph.slice((np - 1) * n_embd..np * n_embd),
10657 n_embd,
10658 )?;
10659 Ok(a)
10660 })
10661 } else {
10662 Err("no prompt hiddens".into())
10663 };
10664 match (cache.snapshot(e), anchor) {
10665 (Ok(snap), Ok(last_h)) => {
10666 *slot = Some(SpecCheckpoint { snap, pos, last_h });
10667 }
10668 (s, a) => {
10669 *slot = None; // a stale checkpoint would rewind to the WRONG boundary
10670 if std::env::var("MEMRA_DEBUG_SPEC").is_ok() {
10671 let err = s
10672 .err()
10673 .map(|e| e.to_string())
10674 .or_else(|| a.err().map(|e| e.to_string()))
10675 .unwrap_or_default();
10676 eprintln!(
10677 "[spec] turn checkpoint skipped ({err}); \
10678 next turn re-primes in full"
10679 );
10680 }
10681 }
10682 }
10683 }
10684 }
10685 // INIT FEED — skipped on a pending carry: last_token (the carried bonus) is NOT in the
10686 // caches and must NOT be fed solo; round 0's batched verify commits it as col 0. Its
10687 // seed/anchor hidden is the carried last_h (copied below); last_pred is dead in the
10688 // pending path (t_pred reads verify col 0 — the accept walk overwrites it).
10689 let mut last_pred = 0u32;
10690 let mut last_col_logits: Option<CudaSlice<f32>> = None;
10691 // CONSTRAINED: the init feed's logits back the (n_acc==0, base==0) masked-argmax
10692 // recompute in the grammar-truncation walk — retained host-side, round 0 only.
10693 let mut init_logits_host: Option<Vec<f32>> = None;
10694 let h_seed0: CudaSlice<f32> = if carried_pending.is_none() {
10695 let (init_logits, h) = self.spec_target_step_h(e, last_token, &mut *cache)?;
10696 last_pred = argmax(&init_logits) as u32;
10697 if constraint.is_some() {
10698 init_logits_host = Some(init_logits.clone());
10699 }
10700 // sampled mode: p-distribution after last_token, for the j==0/base==0 accept test.
10701 if sampled {
10702 last_col_logits = Some(e.htod(&init_logits)?);
10703 }
10704 h
10705 } else {
10706 // predecessor-row anchor: hidden of the last COMMITTED row (the carry contract).
10707 let lh = sess_tail
10708 .as_ref()
10709 .unwrap()
10710 .1
10711 .as_ref()
10712 .expect("pending carry requires last_h");
10713 e.clone_dtod(lh)?
10714 };
10715 let t_init = t_ent.elapsed();
10716 let mut last_col_stats: Option<(f32, f32, f32)> = None;
10717 // PERSISTENT h_seed buffer (allocated BEFORE any graph capture so no captured scratch can
10718 // alias it): every path that updates the round seed copies INTO it — no per-round allocs,
10719 // stable pointer for the graph-draft round-start copy.
10720 let mut h_seed_buf = e.clone_dtod(&h_seed0)?;
10721 // Predecessor-pairing trackers: `fill_prev` = trunk hidden AT the last COMMITTED row (the
10722 // predecessor of the next verify's col 0 — the reference's carried pending-h analogue;
10723 // also the predecessor-row hidden for the round-0 legacy-replay seed). At round 0 that
10724 // row is last_token's own (h_seed0). The chain step-0 seed under the pairing default =
10725 // hidden of the row BEFORE last_token = the prompt's last row at round 0 (h_seed_buf
10726 // overwritten below).
10727 let mut fill_prev = e.clone_dtod(&h_seed0)?;
10728 {
10729 if let Some(ph) = &prompt_h {
10730 let np = prompt.len();
10731 e.copy_view_into(
10732 &mut h_seed_buf,
10733 0,
10734 &ph.slice((np - 1) * n_embd..np * n_embd),
10735 n_embd,
10736 )?;
10737 } else if continuation
10738 && let Some((_, lh, _, _, _)) = sess_tail.as_ref()
10739 && let Some(lh) = lh.as_ref()
10740 {
10741 e.copy_into(&mut h_seed_buf, 0, lh, n_embd)?;
10742 }
10743 }
10744 // Persistent device prediction slots for the accept walk (max k+1 verify columns).
10745 let mut preds_d = e.alloc_u32_zeroed(k + 2)?;
10746
10747 let debug_spec = std::env::var("MEMRA_DEBUG_SPEC").is_ok();
10748 // MEMRA_SPEC_STATS=1: per-slot accept histogram + draft-length histogram, printed once at
10749 // the end. Metric normalization vs the reference engine: BOTH engines count
10750 // accepted/drafted where the chain stopped at p-min and the sub-threshold token is
10751 // discarded uncounted — per-slot decay + chain-length mix are the extra dimensions.
10752 let spec_stats = std::env::var("MEMRA_SPEC_STATS").is_ok();
10753 let mut st_drafted = vec![0usize; k];
10754 let mut st_accepted = vec![0usize; k];
10755 let mut st_len_hist = vec![0usize; k + 1];
10756 let mut st_full = 0usize;
10757 // P-MIN CONFIDENCE GATE (MEMRA_SPEC_PMIN, the serve script's --spec-draft-p-min mechanism):
10758 // stop the draft chain early when the head's softmax confidence in its own pick drops
10759 // below p_min. Hoisted above the loop: the graph capture bakes the prob kernels iff on.
10760 static PMIN: std::sync::OnceLock<f32> = std::sync::OnceLock::new();
10761 let p_min = *PMIN.get_or_init(|| {
10762 std::env::var("MEMRA_SPEC_PMIN")
10763 .ok()
10764 .and_then(|v| v.parse().ok())
10765 .unwrap_or(0.0)
10766 });
10767 // ZERO-DRAFT ROUNDS (MEMRA_SPEC_PMIN0=1, vendored from llama.cpp's draft gating): let the
10768 // p-min gate apply at j==0 too, so a low-confidence round drafts NOTHING and the verify
10769 // batch is just the pending bonus (m=1 = a plain decode step). llama's 35B win rides
10770 // exactly this — draft acceptance 76% at mean len 2.5 because unpredictable stretches
10771 // never pay draft+verify overhead. Only legal when a pending bonus exists (an empty
10772 // verify batch is not); the j==0 exemption stays for pending-less rounds.
10773 let pmin0 = std::env::var("MEMRA_SPEC_PMIN0")
10774 .map(|v| v == "1")
10775 .unwrap_or(false);
10776
10777 // --- GRAPH DRAFT setup: persistent I/O buffers + ONE capture (2 warmups inside). The
10778 // warmups mutate scratch len_d / pos / tok / seed — all reset at every round start, so the
10779 // only restore needed is the scratch counter. Capture failure (e.g. a non-capturable
10780 // cuBLAS path in an exotic head) falls back to the eager draft chain.
10781 // PER-SESSION PERSISTENCE (2026-08-01): session calls reuse the DraftGraphCtx parked on
10782 // the SpecSession — the capture (2 warmup head forwards + instantiate) ran ONCE at the
10783 // session's first burst, not per burst (measured ~16ms/burst fixed cost on H100 q27,
10784 // research/spec-serving-20260801). Reuse is pointer-exact: the graph bakes the session's
10785 // own scratch KV (never realloc'd), the model's resident embedding, the OnceLock p_min,
10786 // and the g_* buffers carried in the ctx — replay dispatch is identical to a fresh
10787 // capture, so draft tokens are bit-identical (drafts never decide exactness anyway; the
10788 // verify arbitrates). Single-shot calls (sess=None) build a fresh ctx and drop it.
10789 let mut dctx: DraftGraphCtx = match sess_draft_slot.as_mut().and_then(|s| s.take()) {
10790 Some(c) => c,
10791 None => DraftGraphCtx::new(e, n_embd, if sampled { d_vocab } else { 1 })?,
10792 };
10793 // FAIL-SAFE (step-OOM park replay): pre-mark both fallback flags so no capture arm
10794 // below can fire — LOUD once per replayed session through the standard WARN line.
10795 if sess_capture_disabled {
10796 let reason =
10797 "session replayed after a step-OOM park; draft capture disabled (fail-safe)";
10798 let flip = dctx.failed.mark_greedy(reason);
10799 let flip_s = dctx.failed.mark_sampled(reason);
10800 if let Some(line) = flip.or(flip_s) {
10801 eprintln!("{line}");
10802 }
10803 }
10804 // A session that ran greedy bursts first sized g_q/g_perturb at 1; a sampled resume
10805 // needs d_vocab. Realloc is legal exactly while graph_s is None (nothing baked them).
10806 if sampled && dctx.g_q.len() < d_vocab {
10807 dctx.g_q = e.zeros(d_vocab)?;
10808 dctx.g_perturb = e.zeros(d_vocab)?;
10809 }
10810 // DRAFT-SIDE GRAMMAR MASK (lane/draft-mask, 2026-08-04): the drafter samples the
10811 // grammar's legal set, so proposals are legal BY CONSTRUCTION and the verify-side
10812 // truncation (the correctness backstop) stops cutting every tight-schema round.
10813 // The mask is one node inside the captured draft chain — presence is a CAPTURE-TIME
10814 // shape, so a parked graph of the other shape is dropped and recaptured.
10815 let dmask_on = constraint
10816 .as_deref()
10817 .is_some_and(|c| c.draft_mask_enabled());
10818 let dmask_words = if dmask_on { d_vocab.div_ceil(32) } else { 0 };
10819 if dmask_on && dctx.g_dmask.len() < dmask_words {
10820 dctx.g_dmask = e.alloc_u32_zeroed(dmask_words)?;
10821 dctx.graph = None; // the old capture baked the old (or no) mask pointer
10822 dctx.chain = None; // chain last-row graphs bake the same pointer
10823 dctx.failed.clear_greedy();
10824 dctx.keeper.clear();
10825 }
10826 if (dctx.graph.is_some() || dctx.chain.is_some()) && dctx.graph_masked != dmask_on {
10827 dctx.graph = None;
10828 dctx.chain = None;
10829 dctx.failed.clear_greedy();
10830 dctx.keeper.clear();
10831 }
10832 // MULTI-HEAD CHAIN mode (mtp_extra non-empty — step37's 3-head shipping shape): the
10833 // step-modulo prefix-replay chain captures PER-HEAD single-row graphs
10834 // (`DraftChainGraphs`) instead of the one self-feeding graph below; the single-head
10835 // capture arms are untouched and unreachable in this mode (the launch arms branch the
10836 // same way). This removes the historical `mtp_extra.is_empty()` capture exclusion —
10837 // and with it the silent no-attempt hole: a chain capture that FAILS now trips the
10838 // same LOUD draft-graph WARN as a single-head failure.
10839 let chain_mode = !self.mtp_extra.is_empty();
10840 // ---- PRE-CAPTURE VRAM RESERVE CHECK + PER-SESSION DRAFT-STATE MEASUREMENT ----
10841 // (lane/step37-vram-admission-20260830). `cap_eff0` opens the measurement bracket:
10842 // when any capture succeeds in THIS call, the effective-free delta across the whole
10843 // capture section is recorded as the model's per-session draft-state high-water
10844 // (admission charges it per spec-capable session — this state was charged at ZERO
10845 // before the lane). The reserve check runs BEFORE any capture arm can allocate: a
10846 // refused capture trips the same LOUD once-per-flip WARN class as a failed one, but
10847 // with the card's headroom still intact (the owner's single-session OOM was a capture
10848 // attempt walking the card to the edge and stranding the eager fallback at 5 MiB free).
10849 let cap_eff0 = e
10850 .ctx()
10851 .mem_get_info()
10852 .ok()
10853 .map(|(f, _)| f.saturating_add(e.pool_cached_bytes()));
10854 // Peak instrument for the same bracket: the CAPTURE-TIME peak (warmup transients +
10855 // instantiate scratch, alive together) dwarfs the parked delta — measured on the
10856 // owner shape: a capture whose PARKED state reads ~2.6GB walked a ~7GB-free card to
10857 // OOM mid-capture. Reset the pool watermark here; read it at bracket end.
10858 let _ = e.pool_high_water_reset();
10859 let cap_used0 = e.pool_reserved_used().1;
10860 let mut captured_now = false;
10861 let mut capture_oom_entry_eff: Option<usize> = None;
10862 let capture_need = {
10863 let observed = self.draft_session_admission_bytes();
10864 if observed > 0 {
10865 observed
10866 } else {
10867 draft_capture_bootstrap_estimate(
10868 if chain_mode { self.mtp_head_count() } else { 1 },
10869 k,
10870 d_vocab,
10871 n_embd,
10872 )
10873 }
10874 };
10875 if spec_capture_gate_on()
10876 && graph_draft
10877 && !sampled
10878 && !dctx.failed.greedy_failed()
10879 && ((chain_mode && dctx.chain.is_none() && mtp_chain_graph_on())
10880 || (!chain_mode && dctx.graph.is_none()))
10881 && let Some(reason) = capture_headroom_refusal(e, capture_need)
10882 && let Some(line) = dctx.failed.mark_greedy(&reason)
10883 {
10884 eprintln!("{line}");
10885 }
10886 if graph_draft
10887 && !sampled
10888 && chain_mode
10889 && dctx.chain.is_none()
10890 && !dctx.failed.greedy_failed()
10891 {
10892 if mtp_chain_graph_on() {
10893 let heads_n = self.mtp_head_count();
10894 let DraftGraphCtx {
10895 g_tok,
10896 g_pos,
10897 g_seed,
10898 g_p,
10899 g_dmask,
10900 ..
10901 } = &mut dctx;
10902 if dmask_on {
10903 e.htod_u32_into(g_dmask, &vec![u32::MAX; dmask_words])?;
10904 }
10905 let g_dmask_ro: &CudaSlice<u32> = &*g_dmask;
10906 let with_prob = p_min > 0.0;
10907 // CAPTURE-RETAIN (#68 fix): one keeper for the whole chain — every graph's
10908 // warmup transients stay pinned as long as any of them replays.
10909 let cap_res = (|| -> Result<DraftChainGraphs, Box<dyn std::error::Error>> {
10910 // dcw door: same warmup headroom pre-arm as the single-head capture
10911 // below — every plane, because each head's capture warmups append on
10912 // its OWN plane. INSIDE the fallible closure (vram-admission lane): an
10913 // OOM here used to `?` out of the whole burst as a step error; now it
10914 // is a capture failure — LOUD WARN, eager chain serves.
10915 if step35_draft_dcw_on() {
10916 scratch.ensure_dcw_headroom(e, k + 2)?;
10917 }
10918 let mut interior = Vec::with_capacity(heads_n);
10919 let mut last = Vec::with_capacity(heads_n);
10920 let mut keeper: Vec<Box<dyn std::any::Any + Send>> = Vec::new();
10921 for hi in 0..heads_n {
10922 let head = self.mtp_head_at(hi);
10923 // interior row: KV append + carrier only (`with_head=false` — the
10924 // eager chain discards interior logits too, so this is the same
10925 // consumed-byte program minus the dead full-vocab head matmul).
10926 let (g, keep) = e.capture_graph_retained(|e| {
10927 self.mtp_head_forward_cap(
10928 e,
10929 head,
10930 g_tok,
10931 g_pos,
10932 g_seed,
10933 g_p,
10934 &mut *scratch,
10935 hi,
10936 false,
10937 false,
10938 embd_gpu.expect("graph draft requires resident embedding"),
10939 embd_qt,
10940 embd_rb,
10941 d_vocab,
10942 None,
10943 None,
10944 None,
10945 )
10946 })?;
10947 // the warmups appended rows on plane hi; rewind before the next
10948 // capture so successive warmups never outrun the pre-armed headroom.
10949 scratch.set_plane_len(e, hi, base)?;
10950 interior.push(g);
10951 keeper.extend(keep);
10952 // last row: head matmul + greedy argmax tail (+ p when the policy
10953 // reads it, + the grammar-mask node when constrained).
10954 let (g2, keep2) = e.capture_graph_retained(|e| {
10955 self.mtp_head_forward_cap(
10956 e,
10957 head,
10958 g_tok,
10959 g_pos,
10960 g_seed,
10961 g_p,
10962 &mut *scratch,
10963 hi,
10964 with_prob,
10965 true,
10966 embd_gpu.expect("graph draft requires resident embedding"),
10967 embd_qt,
10968 embd_rb,
10969 d_vocab,
10970 None,
10971 None,
10972 if dmask_on {
10973 Some((g_dmask_ro, dmask_words))
10974 } else {
10975 None
10976 },
10977 )
10978 })?;
10979 scratch.set_plane_len(e, hi, base)?;
10980 last.push(g2);
10981 keeper.extend(keep2);
10982 }
10983 Ok(DraftChainGraphs {
10984 interior,
10985 last,
10986 _keeper: keeper,
10987 })
10988 })();
10989 match cap_res {
10990 Ok(cg) => {
10991 scratch.set_len(e, base)?;
10992 // POSITIVE engagement receipt (the 3a lesson: a WARN-free boot is
10993 // NOT evidence of capture — the captured state must name itself).
10994 eprintln!(
10995 "[mtp-chain-graph] captured mode=greedy heads={heads_n} \
10996 interior={heads_n} last={heads_n} masked={}",
10997 dmask_on as u8
10998 );
10999 dctx.chain = Some(cg);
11000 dctx.graph_masked = dmask_on;
11001 captured_now = true;
11002 }
11003 Err(err) => {
11004 scratch.set_len(e, base)?;
11005 // LOUD flip (audit Q2): a dropped draft graph is a coverage loss,
11006 // never silent — now including the multi-head shipping shape.
11007 // OOM RECOVERY (vram-admission lane): a failed attempt's freed
11008 // transients sit CACHED in the async pool where the driver cannot
11009 // see them; trim them back so the eager fallback (and any driver-
11010 // side allocation) actually has the headroom the free suggests.
11011 let mut reason = err.to_string();
11012 if capture_err_is_oom(&reason) {
11013 capture_oom_entry_eff = capture_oom_entry_eff.max(cap_eff0);
11014 let trimmed = e.pool_trim_to_zero();
11015 if trimmed > 0 {
11016 reason.push_str(&format!(
11017 "; pool trimmed {}MB back to the driver",
11018 trimmed / (1 << 20)
11019 ));
11020 }
11021 }
11022 if let Some(line) = dctx.failed.mark_greedy(&reason) {
11023 eprintln!("{line}");
11024 }
11025 }
11026 }
11027 } else {
11028 // Disarmed by MEMRA_MTP_CHAIN_GRAPH=0: say so once per process — the OFF arm
11029 // must be attributable in a boot log, never inferable from silence.
11030 static NOTE: std::sync::Once = std::sync::Once::new();
11031 NOTE.call_once(|| {
11032 eprintln!(
11033 "[spec] multi-head draft-chain capture disarmed \
11034 (MEMRA_MTP_CHAIN_GRAPH=0); eager chain serves this shape"
11035 );
11036 });
11037 }
11038 }
11039 if graph_draft
11040 && !sampled
11041 && !chain_mode
11042 && dctx.graph.is_none()
11043 && !dctx.failed.greedy_failed()
11044 {
11045 let DraftGraphCtx {
11046 g_tok,
11047 g_pos,
11048 g_seed,
11049 g_p,
11050 g_dmask,
11051 ..
11052 } = &mut dctx;
11053 // capture-time contents: ALL-ONES (ban nothing). A replay only ever runs after the
11054 // host uploads the position's real words, so the warmups stay grammar-free.
11055 if dmask_on {
11056 e.htod_u32_into(g_dmask, &vec![u32::MAX; dmask_words])?;
11057 }
11058 let g_dmask_ro: &CudaSlice<u32> = &*g_dmask;
11059 // CAPTURE-RETAIN (#68 fix): the warmup transients' pool addresses are baked into the
11060 // captured graph; the keeper pins them for the graph's lifetime. capture_graph (non-
11061 // retained) freed them at exit — safe for one-shot generate_spec (nothing else touches
11062 // the pool between replays) but WRONG for sessions: burst-boundary prime/fill/commit
11063 // passes (and, in serve, other sessions) recycle those addresses and the replay then
11064 // clobbers live buffers — the ST serve-spec corruption (research/serve-st-20260803).
11065 let cap_res = (|| {
11066 // dcw door: the capture warmups append device-counter rows the capture body
11067 // cannot rebase for; pre-arm ring headroom host-side (no-op on flat planes /
11068 // room-enough rings, and the door-off path is untouched). INSIDE the fallible
11069 // closure (vram-admission lane): an OOM here is a capture failure, not a
11070 // burst-killing step error.
11071 if step35_draft_dcw_on() {
11072 scratch.ensure_dcw_headroom(e, k + 2)?;
11073 }
11074 e.capture_graph_retained(|e| {
11075 self.mtp_head_forward_cap(
11076 e,
11077 mtp,
11078 g_tok,
11079 g_pos,
11080 g_seed,
11081 g_p,
11082 &mut *scratch,
11083 0,
11084 p_min > 0.0,
11085 true,
11086 embd_gpu.expect("graph draft requires resident embedding"),
11087 embd_qt,
11088 embd_rb,
11089 d_vocab,
11090 None,
11091 None,
11092 if dmask_on {
11093 Some((g_dmask_ro, dmask_words))
11094 } else {
11095 None
11096 },
11097 )
11098 })
11099 })();
11100 match cap_res {
11101 Ok((g, keep)) => {
11102 scratch.set_len(e, base)?;
11103 dctx.graph = Some(g);
11104 dctx.graph_masked = dmask_on;
11105 dctx.keeper = keep;
11106 captured_now = true;
11107 }
11108 Err(err) => {
11109 scratch.set_len(e, base)?;
11110 // LOUD flip (audit Q2): a dropped draft graph is a coverage loss, never
11111 // silent. Once per flip — mark returns None on an already-failed ctx.
11112 let mut reason = err.to_string();
11113 if capture_err_is_oom(&reason) {
11114 capture_oom_entry_eff = capture_oom_entry_eff.max(cap_eff0);
11115 let trimmed = e.pool_trim_to_zero();
11116 if trimmed > 0 {
11117 reason.push_str(&format!(
11118 "; pool trimmed {}MB back to the driver",
11119 trimmed / (1 << 20)
11120 ));
11121 }
11122 }
11123 if let Some(line) = dctx.failed.mark_greedy(&reason) {
11124 eprintln!("{line}");
11125 }
11126 }
11127 }
11128 }
11129 // --- SAMPLED GRAPH DRAFT setup (step 3 of the sampled-spec arc): a SECOND capture, own
11130 // graph object, built only when sampled && graph-eligible — the greedy capture above is
11131 // untouched (and skipped when sampled: its graph would never be launched). Same head
11132 // forward, but the in-graph argmax reads GUMBEL-PERTURBED logits; the Philox event
11133 // counter lives in the persistent device g_ctr (bumped in-graph, host-seeded from sctr
11134 // once per round); the raw head logits land in the persistent g_q for the host's
11135 // per-replay async D2D into the round's q slot (q_slots, K x d_vocab, allocated once).
11136 // seed/temp are capture-time constants — baked into graph_s, so a pool-resumed request
11137 // with a different (seed, temp, k) drops the parked sampled graph and recaptures.
11138 // COST OF THE FRESH-SEED SERVE DEFAULT (dogfood F4, 2026-08-04): omitting `seed` on a
11139 // serve request now draws fresh per-request entropy (it used to default to a pinned 0),
11140 // so a seed-omitting request that RESUMES a parked spec session finds an s_key baked
11141 // with the PREVIOUS request's seed and pays one recapture. Bounded, and it does not
11142 // reopen the ~16ms/burst regression the persistent ctx exists to fix: a session's seed
11143 // is fixed for its whole lifetime (worker.rs reads s.sampler.seed() per burst), so
11144 // this compare misses at most ONCE per resumed request — the first burst recaptures
11145 // and every later burst in that request replays. A client that wants the parked graph
11146 // AND reproducibility supplies an explicit `seed`, honored exactly, which keeps s_key
11147 // stable across its whole conversation.
11148 // COMPOSITION RULE (fspec x gsd merge): the in-graph chain samples from the RAW
11149 // softmax — it can hold neither per-row filter stats nor the varying penalty history.
11150 // The sampled graph therefore engages only in the PURE-TEMP regime; filters/penalties
11151 // force the eager draft (which computes stats/penalties per row).
11152 // KEY THE WHOLE REGIME, not just the baked constants (lane/graph-s-key-exactness-
11153 // 20260819). `s_key` used to be `(seed, temp, k)`; the filters and penalties were left
11154 // out, so a filtered request resuming a session that parked a PURE-TEMP graph kept it —
11155 // and the launch site never re-asked `pure_temp`. See [`SampledGraphKey`] for what that
11156 // costs (an unconditional accept of out-of-head draft tokens, i.e. an exactness bug on
11157 // the request shape the vendor-default flip makes the majority).
11158 let s_key = SampledGraphKey::new(sp_seed, sp_temp, k, sp.top_k, sp.top_p, sp.min_p, pen_on);
11159 let pure_temp = s_key.pure_temp();
11160 // The regime the sampled graph may be captured/launched in: pure-temp always;
11161 // truncation-filtered when the filtered-capture door is on (the filter runs
11162 // IN-GRAPH — lane/step37-draft-graph-serving-20260830); penalties never.
11163 let s_capturable = s_key.graph_capturable();
11164 if sampled && dctx.s_key.is_some_and(|old| old != s_key) {
11165 dctx.graph_s = None;
11166 dctx.chain_s = None;
11167 dctx.failed.clear_sampled();
11168 dctx.s_key = None;
11169 dctx.q_slots.clear();
11170 dctx.keeper_s.clear();
11171 }
11172 // PRE-CAPTURE VRAM RESERVE CHECK, sampled arms (vram-admission lane): same contract
11173 // as the greedy check above — refuse BEFORE allocating, LOUD once, eager serves.
11174 if spec_capture_gate_on()
11175 && graph_draft
11176 && sampled
11177 && s_capturable
11178 && !dctx.failed.sampled_failed()
11179 && ((chain_mode && dctx.chain_s.is_none() && mtp_chain_graph_on())
11180 || (!chain_mode && dctx.graph_s.is_none()))
11181 && let Some(reason) = capture_headroom_refusal(e, capture_need)
11182 && let Some(line) = dctx.failed.mark_sampled(&reason)
11183 {
11184 eprintln!("{line}");
11185 }
11186 // FILTERED capture nodes need q slots sized d_vocab AND the stat slots; the pure-temp
11187 // body leaves g_th/g_z/g_mx untouched (they exist from ctx creation either way).
11188 if graph_draft
11189 && sampled
11190 && s_capturable
11191 && chain_mode
11192 && dctx.chain_s.is_none()
11193 && !dctx.failed.sampled_failed()
11194 {
11195 if mtp_chain_graph_on() {
11196 let heads_n = self.mtp_head_count();
11197 let filtered = s_key.filtered();
11198 let DraftGraphCtx {
11199 g_tok,
11200 g_pos,
11201 g_seed,
11202 g_p,
11203 g_ctr,
11204 g_perturb,
11205 g_q,
11206 g_rows0,
11207 g_th,
11208 g_z,
11209 g_mx,
11210 ..
11211 } = &mut dctx;
11212 let with_prob = p_min > 0.0;
11213 let cap_res = (|| -> Result<DraftChainGraphs, Box<dyn std::error::Error>> {
11214 // dcw pre-arm INSIDE the fallible closure (vram-admission lane): an OOM
11215 // here is a capture failure with the LOUD WARN, never a step error.
11216 if step35_draft_dcw_on() {
11217 scratch.ensure_dcw_headroom(e, k + 2)?;
11218 }
11219 let mut interior = Vec::with_capacity(heads_n);
11220 let mut last = Vec::with_capacity(heads_n);
11221 let mut keeper: Vec<Box<dyn std::any::Any + Send>> = Vec::new();
11222 for hi in 0..heads_n {
11223 let head = self.mtp_head_at(hi);
11224 // interior row: no head, no draw — shared shape with the greedy
11225 // chain's interior, captured per mode for keeper-lifetime hygiene.
11226 let (g, keep) = e.capture_graph_retained(|e| {
11227 self.mtp_head_forward_cap(
11228 e,
11229 head,
11230 g_tok,
11231 g_pos,
11232 g_seed,
11233 g_p,
11234 &mut *scratch,
11235 hi,
11236 false,
11237 false,
11238 embd_gpu.expect("graph draft requires resident embedding"),
11239 embd_qt,
11240 embd_rb,
11241 d_vocab,
11242 None,
11243 None,
11244 None,
11245 )
11246 })?;
11247 scratch.set_plane_len(e, hi, base)?;
11248 interior.push(g);
11249 keeper.extend(keep);
11250 // last row: head matmul + the in-graph categorical draw (filtered
11251 // nodes when the request carries filters).
11252 let (g2, keep2) = e.capture_graph_retained(|e| {
11253 self.mtp_head_forward_cap(
11254 e,
11255 head,
11256 g_tok,
11257 g_pos,
11258 g_seed,
11259 g_p,
11260 &mut *scratch,
11261 hi,
11262 with_prob,
11263 true,
11264 embd_gpu.expect("graph draft requires resident embedding"),
11265 embd_qt,
11266 embd_rb,
11267 d_vocab,
11268 Some(SampledCapArgs {
11269 ctr: &mut *g_ctr,
11270 perturb: &mut *g_perturb,
11271 q_out: &mut *g_q,
11272 seed: sp_seed,
11273 temp: sp_temp,
11274 filt: if filtered {
11275 Some(SampledCapFilter {
11276 rows0: &*g_rows0,
11277 th: &mut *g_th,
11278 z: &mut *g_z,
11279 mx: &mut *g_mx,
11280 top_k: sp.top_k,
11281 top_p: sp.top_p,
11282 min_p: sp.min_p,
11283 })
11284 } else {
11285 None
11286 },
11287 }),
11288 None,
11289 None, // constrained spec is greedy-only
11290 )
11291 })?;
11292 scratch.set_plane_len(e, hi, base)?;
11293 last.push(g2);
11294 keeper.extend(keep2);
11295 }
11296 Ok(DraftChainGraphs {
11297 interior,
11298 last,
11299 _keeper: keeper,
11300 })
11301 })();
11302 match cap_res {
11303 Ok(cg) => {
11304 scratch.set_len(e, base)?;
11305 // NO STRANDED PARTIAL STATE (vram-admission lane): the q-slot allocs
11306 // after a successful capture are themselves fallible on a tight card.
11307 // A mid-loop failure used to `?` out as a step error, leaving orphan
11308 // slots parked on the ctx (wrong count, stale contents) for the next
11309 // capture attempt to stack onto. Allocate all-or-nothing: on failure
11310 // drop the fresh graphs AND the partial slots, mark the LOUD fallback.
11311 dctx.q_slots.clear();
11312 let slots = (0..k)
11313 .map(|_| e.zeros(d_vocab))
11314 .collect::<Result<Vec<_>, _>>();
11315 match slots {
11316 Ok(slots) => {
11317 dctx.q_slots = slots;
11318 eprintln!(
11319 "[mtp-chain-graph] captured mode=sampled heads={heads_n} \
11320 interior={heads_n} last={heads_n} filtered={} key={s_key:?}",
11321 s_key.filtered() as u8
11322 );
11323 dctx.chain_s = Some(cg);
11324 dctx.s_key = Some(s_key);
11325 captured_now = true;
11326 }
11327 Err(err) => {
11328 drop(cg);
11329 dctx.q_slots.clear();
11330 let mut reason = format!("q-slot alloc failed: {err}");
11331 if capture_err_is_oom(&reason) {
11332 capture_oom_entry_eff = capture_oom_entry_eff.max(cap_eff0);
11333 let trimmed = e.pool_trim_to_zero();
11334 if trimmed > 0 {
11335 reason.push_str(&format!(
11336 "; pool trimmed {}MB back to the driver",
11337 trimmed / (1 << 20)
11338 ));
11339 }
11340 }
11341 if let Some(line) = dctx.failed.mark_sampled(&reason) {
11342 eprintln!("{line}");
11343 }
11344 }
11345 }
11346 }
11347 Err(err) => {
11348 scratch.set_len(e, base)?;
11349 let mut reason = err.to_string();
11350 if capture_err_is_oom(&reason) {
11351 capture_oom_entry_eff = capture_oom_entry_eff.max(cap_eff0);
11352 let trimmed = e.pool_trim_to_zero();
11353 if trimmed > 0 {
11354 reason.push_str(&format!(
11355 "; pool trimmed {}MB back to the driver",
11356 trimmed / (1 << 20)
11357 ));
11358 }
11359 }
11360 if let Some(line) = dctx.failed.mark_sampled(&reason) {
11361 eprintln!("{line}");
11362 }
11363 }
11364 }
11365 } else {
11366 static NOTE_S: std::sync::Once = std::sync::Once::new();
11367 NOTE_S.call_once(|| {
11368 eprintln!(
11369 "[spec] multi-head draft-chain capture disarmed \
11370 (MEMRA_MTP_CHAIN_GRAPH=0); eager chain serves this shape"
11371 );
11372 });
11373 }
11374 }
11375 if graph_draft
11376 && sampled
11377 && s_capturable
11378 && !chain_mode
11379 && dctx.graph_s.is_none()
11380 && !dctx.failed.sampled_failed()
11381 {
11382 let filtered = s_key.filtered();
11383 let DraftGraphCtx {
11384 g_tok,
11385 g_pos,
11386 g_seed,
11387 g_p,
11388 g_ctr,
11389 g_perturb,
11390 g_q,
11391 g_rows0,
11392 g_th,
11393 g_z,
11394 g_mx,
11395 ..
11396 } = &mut dctx;
11397 // CAPTURE-RETAIN (#68 fix): same keeper contract as the greedy capture above.
11398 let cap_res = (|| {
11399 // dcw pre-arm INSIDE the fallible closure (vram-admission lane): an OOM
11400 // here is a capture failure with the LOUD WARN, never a step error.
11401 if step35_draft_dcw_on() {
11402 scratch.ensure_dcw_headroom(e, k + 2)?;
11403 }
11404 e.capture_graph_retained(|e| {
11405 self.mtp_head_forward_cap(
11406 e,
11407 mtp,
11408 g_tok,
11409 g_pos,
11410 g_seed,
11411 g_p,
11412 &mut *scratch,
11413 0,
11414 p_min > 0.0,
11415 true,
11416 embd_gpu.expect("graph draft requires resident embedding"),
11417 embd_qt,
11418 embd_rb,
11419 d_vocab,
11420 Some(SampledCapArgs {
11421 ctr: &mut *g_ctr,
11422 perturb: &mut *g_perturb,
11423 q_out: &mut *g_q,
11424 seed: sp_seed,
11425 temp: sp_temp,
11426 filt: if filtered {
11427 Some(SampledCapFilter {
11428 rows0: &*g_rows0,
11429 th: &mut *g_th,
11430 z: &mut *g_z,
11431 mx: &mut *g_mx,
11432 top_k: sp.top_k,
11433 top_p: sp.top_p,
11434 min_p: sp.min_p,
11435 })
11436 } else {
11437 None
11438 },
11439 }),
11440 None,
11441 None, // constrained spec is greedy-only — sampled never carries a hook
11442 )
11443 })
11444 })();
11445 match cap_res {
11446 Ok((g, keep)) => {
11447 scratch.set_len(e, base)?;
11448 // NO STRANDED PARTIAL STATE: all-or-nothing q slots, same contract as
11449 // the chain arm above.
11450 dctx.q_slots.clear();
11451 let slots = (0..k)
11452 .map(|_| e.zeros(d_vocab))
11453 .collect::<Result<Vec<_>, _>>();
11454 match slots {
11455 Ok(slots) => {
11456 dctx.q_slots = slots;
11457 dctx.graph_s = Some(g);
11458 dctx.s_key = Some(s_key);
11459 dctx.keeper_s = keep;
11460 captured_now = true;
11461 }
11462 Err(err) => {
11463 drop(g);
11464 drop(keep);
11465 dctx.q_slots.clear();
11466 let mut reason = format!("q-slot alloc failed: {err}");
11467 if capture_err_is_oom(&reason) {
11468 capture_oom_entry_eff = capture_oom_entry_eff.max(cap_eff0);
11469 let trimmed = e.pool_trim_to_zero();
11470 if trimmed > 0 {
11471 reason.push_str(&format!(
11472 "; pool trimmed {}MB back to the driver",
11473 trimmed / (1 << 20)
11474 ));
11475 }
11476 }
11477 if let Some(line) = dctx.failed.mark_sampled(&reason) {
11478 eprintln!("{line}");
11479 }
11480 }
11481 }
11482 }
11483 Err(err) => {
11484 scratch.set_len(e, base)?;
11485 // LOUD flip (audit Q2): same contract as the greedy capture above.
11486 let mut reason = err.to_string();
11487 if capture_err_is_oom(&reason) {
11488 capture_oom_entry_eff = capture_oom_entry_eff.max(cap_eff0);
11489 let trimmed = e.pool_trim_to_zero();
11490 if trimmed > 0 {
11491 reason.push_str(&format!(
11492 "; pool trimmed {}MB back to the driver",
11493 trimmed / (1 << 20)
11494 ));
11495 }
11496 }
11497 if let Some(line) = dctx.failed.mark_sampled(&reason) {
11498 eprintln!("{line}");
11499 }
11500 }
11501 }
11502 }
11503 // ---- PER-SESSION DRAFT-STATE MEASUREMENT bracket end (vram-admission lane): when a
11504 // capture landed in THIS call, the effective-free delta across the capture section is
11505 // this session's parked draft-graph state (keepers + q slots + instantiated graphs'
11506 // backing). Recorded as a model-owned high-water; admission charges it per
11507 // spec-capable session (see `draft_session_admission_bytes`).
11508 if captured_now
11509 && let Some(eff0) = cap_eff0
11510 && let Ok((f1, _)) = e.ctx().mem_get_info()
11511 {
11512 let eff1 = f1.saturating_add(e.pool_cached_bytes());
11513 let parked_delta = eff0.saturating_sub(eff1);
11514 let (_res_high, used_high) = e.pool_high_water_reset();
11515 let peak_delta = used_high.saturating_sub(cap_used0);
11516 let observed = parked_delta.max(peak_delta);
11517 if observed > 0
11518 && let Some(hw) = self.record_draft_state_bytes(observed)
11519 {
11520 eprintln!(
11521 "[spec] draft-session state high-water: {}MB (max of parked delta {}MB \
11522 and capture-time pool peak {}MB; charged per spec admission and gating \
11523 future captures)",
11524 hw / (1 << 20),
11525 parked_delta / (1 << 20),
11526 peak_delta / (1 << 20),
11527 );
11528 }
11529 }
11530 // FAILURE IS AN OBSERVATION TOO: a capture that OOM'd at entry-effective E proved
11531 // the capture-time peak exceeds E. Feed E into the gauge so every future gate
11532 // refuses at or below the headroom that just failed (self-healing even when the
11533 // boot probe is disarmed and the bootstrap estimate was blind).
11534 if let Some(entry_eff) = capture_oom_entry_eff
11535 && let Some(hw) = self.record_draft_state_bytes(entry_eff)
11536 {
11537 eprintln!(
11538 "[spec] draft-session capture appetite floor raised to {}MB: a capture \
11539 attempt OOM'd with that much effective free (failure-observed bound)",
11540 hw / (1 << 20)
11541 );
11542 }
11543 // ---- EXACTNESS GUARD, the enforceable half (lane/graph-s-key-exactness-20260819,
11544 // widened by lane/step37-draft-graph-serving-20260830) ----
11545 // With the filters and penalties in `s_key`, a graph that SURVIVED the drop above was
11546 // captured under THIS request's exact regime, and capture requires `graph_capturable`
11547 // (pure-temp, or filtered with the in-graph filter nodes; never penalties) — so a
11548 // parked graph implies both. That implication is the whole exactness argument for the
11549 // graph arm, so it is asserted here rather than assumed: a future change that widens
11550 // the capture condition, narrows the key, or copies a `DraftGraphCtx` across regimes
11551 // fails LOUDLY at this line instead of silently drafting from a distribution the
11552 // verify never reconstructs. Release builds refuse the graph (drop it, draft eager)
11553 // rather than launching it; the launch site re-tests the regime independently.
11554 if sampled
11555 && (dctx.graph_s.is_some() || dctx.chain_s.is_some())
11556 && (!s_capturable || dctx.s_key != Some(s_key))
11557 {
11558 debug_assert!(
11559 false,
11560 "sampled draft graph parked under {:?} survived into a request outside its \
11561 capture regime (top_k={} top_p={} min_p={} pen_on={} capturable={}): the \
11562 in-graph draw and the verify's accept test would see different distributions",
11563 dctx.s_key, sp.top_k, sp.top_p, sp.min_p, pen_on, s_capturable,
11564 );
11565 eprintln!(
11566 "[spec] BUG: dropping a parked sampled draft graph that outlived its capture \
11567 regime (s_key={:?}, request top_k={} top_p={} min_p={} pen_on={} \
11568 capturable={}); drafting EAGER — the key must carry every field that shapes q",
11569 dctx.s_key, sp.top_k, sp.top_p, sp.min_p, pen_on, s_capturable,
11570 );
11571 dctx.graph_s = None;
11572 dctx.chain_s = None;
11573 dctx.s_key = None;
11574 dctx.q_slots.clear();
11575 dctx.keeper_s.clear();
11576 }
11577 // SKEY PROBE (MEMRA_SKEY_PROBE=1): the burst-entry facts the reachability question turns
11578 // on — is this request sampled, is it in a regime the sampled graph is legal in, and is
11579 // a graph PARKED from an earlier request of the same session? The launch arms below
11580 // print which chain actually ran, so the probe never restates the condition.
11581 if skey_probe() {
11582 eprintln!(
11583 "[skey] burst sampled={} pure_temp={} capturable={} temp={} top_k={} top_p={} \
11584 min_p={} pen_on={} k={} graph_draft={} graph_s_parked={} chain_s_parked={} \
11585 s_key_parked={:?}",
11586 sampled as u8,
11587 pure_temp as u8,
11588 s_capturable as u8,
11589 sp_temp,
11590 sp.top_k,
11591 sp.top_p,
11592 sp.min_p,
11593 pen_on as u8,
11594 k,
11595 graph_draft as u8,
11596 dctx.graph_s.is_some() as u8,
11597 dctx.chain_s.is_some() as u8,
11598 dctx.s_key,
11599 );
11600 }
11601 let t_cap = t_ent.elapsed();
11602 // PERSISTENT DRAFT KV: fill the MTP block's K/V for every prompt position from the exact
11603 // trunk hiddens collected during prime — ONE batched K/V-only pass (overwrites any
11604 // capture-warmup garbage; capture left len at 0). last_token (the init feed) needs no
11605 // fill: the first chain step processes it and appends its entry at slot prompt.len().
11606 if let Some(ph) = &prompt_h {
11607 // SESSION: rows [0..base) are the previous turns' exact fills (refresh overwrote them
11608 // with true verify hiddens) — truncate any draft overhang, fill ONLY the suffix at
11609 // global positions [base..base+tp). Fresh call: base==0, identical to before.
11610 scratch.set_len(e, base)?;
11611 // CHUNKED FILL (long-ctx OOM fix, 2026-07-05): mtp_kv_fill's transients scale with its
11612 // T (concat = T*2*n_embd*4B — 1.5GB at 40k) and its concat loop is 2*T launches. The
11613 // fill is a pure sequential append, so chunking is exact: each chunk appends its rows
11614 // at pos0=base+start with the identical per-row math. Same knob as the trunk prime.
11615 let tp = prompt.len();
11616 let fill_chunk: usize = if crate::cache::swa_ring_on() {
11617 crate::hybrid_forward::prime_chunk_tokens(tp, self.layers.len())
11618 } else {
11619 // Preserve the flag-OFF schedule byte-for-byte, including the legacy zero value
11620 // meaning one monolithic fill.
11621 std::env::var("MEMRA_PRIME_CHUNK")
11622 .ok()
11623 .and_then(|v| v.parse().ok())
11624 .unwrap_or(4096)
11625 };
11626 let fill_chunk = if fill_chunk == 0 { tp } else { fill_chunk };
11627 // CUDA launch wall (same class as the trunk prime's PRIME_CHUNK_LAUNCH_CAP):
11628 // a fill call's matmuls can land on the grid.y=m dp4a family, and grid.y caps
11629 // at 65,535. This loop has no tail fold, so the raw limit is exact:
11630 // tp <= 65,535 keeps the legacy schedule (monolithic included) byte-for-byte,
11631 // and larger fills — unreachable before the trunk prime's own cap fix — chunk.
11632 let fill_chunk = fill_chunk.min(crate::hybrid_forward::CUDA_GRID_YZ_MAX);
11633 let mut start = 0usize;
11634 while start < tp {
11635 let end = (start + fill_chunk).min(tp);
11636 let tc = end - start;
11637 {
11638 // PREDECESSOR pairing: row i gets h[i-1]; global row 0 a zeros row (the
11639 // reference engine's initial pending-h is zeroed too); a session turn's row 0
11640 // gets the PREVIOUS turn's last committed hidden (sess.last_h). Per chunk:
11641 // rows start..end read h[start-1..end-1] — one dtod into a chunk buffer.
11642 let mut phs = e.zeros(tc * n_embd)?;
11643 let (src_lo, dst_off) = if start == 0 {
11644 (0, n_embd)
11645 } else {
11646 ((start - 1) * n_embd, 0)
11647 };
11648 let n_copy = if start == 0 {
11649 (tc - 1) * n_embd
11650 } else {
11651 tc * n_embd
11652 };
11653 if start == 0
11654 && let Some((_, lh, _, _, _)) = sess_tail.as_ref()
11655 && let Some(lh) = lh.as_ref()
11656 {
11657 e.copy_into(&mut phs, 0, lh, n_embd)?;
11658 }
11659 if n_copy > 0 {
11660 e.copy_view_into(
11661 &mut phs,
11662 dst_off,
11663 &ph.slice(src_lo..src_lo + n_copy),
11664 n_copy,
11665 )?;
11666 }
11667 self.mtp_kv_fill_all(
11668 e,
11669 &prompt[start..end],
11670 &phs,
11671 base + start,
11672 &mut *scratch,
11673 embd_dev,
11674 )?;
11675 }
11676 start = end;
11677 }
11678 }
11679 // MEMRA_PROFILE_SPEC=2: profiler capture starts HERE — after the prime, so an
11680 // `nsys -c cudaProfilerApi` capture contains ONLY the round loop (draft/verify/commit).
11681 // (=1 brackets the whole call in run_spec.rs, prime included.)
11682 if std::env::var("MEMRA_PROFILE_SPEC").as_deref() == Ok("2") {
11683 unsafe extern "C" {
11684 fn cudaProfilerStart() -> i32;
11685 }
11686 unsafe {
11687 cudaProfilerStart();
11688 }
11689 }
11690 // ROUND-STREAM stage (c) 4 (MEMRA_SPEC_STREAM=1, experimental): pre-issued M-round
11691 // bursts with ZERO per-round host readbacks — the accept/seed/rollback/ring kernels
11692 // consume each other's device outputs; the host drains the ring every M rounds. v1
11693 // constraints: greedy, !spec_replay, single-shot, batched-linear layers, no refresh
11694 // fills (acceptance effect A/B-arbitrated), enters from round 1 (pending guaranteed).
11695 // NOTE: not gated on the caller's graph_draft (its trunk_dense conjunct turns the 35B
11696 // MoE off) — the stream capture encloses ONLY the dense MTP head; the head-dense /
11697 // full-prec / k gates are re-derived here and a failed capture degrades to stream-off.
11698 let stream_on = crate::spec::spec_stream()
11699 && !sampled
11700 && !spec_replay
11701 && self.mtp_extra.is_empty()
11702 && constraint.is_none()
11703 && !session_mode
11704 && embd_gpu.is_some()
11705 && !crate::model::full_prec_enabled()
11706 && k + 2 < 96;
11707 let mut stream_graph: Option<cudarc::driver::CudaGraph> = None;
11708 let mut g_tokp2k = e.alloc_u32_zeroed(2 * k.max(1))?;
11709 if stream_on {
11710 let cap = e.capture_graph(|e| {
11711 for j in 0..k.max(1) {
11712 self.mtp_head_forward_cap(
11713 e,
11714 mtp,
11715 &mut dctx.g_tok,
11716 &mut dctx.g_pos,
11717 &mut dctx.g_seed,
11718 &mut dctx.g_p,
11719 &mut *scratch,
11720 0,
11721 true,
11722 true,
11723 embd_gpu.expect("round stream requires resident embedding"),
11724 embd_qt,
11725 embd_rb,
11726 d_vocab,
11727 None,
11728 Some((&mut g_tokp2k, j, d2t_dev.as_ref())),
11729 None, // round-stream requires constraint.is_none() (see stream_on)
11730 )?;
11731 }
11732 Ok(())
11733 });
11734 match cap {
11735 Ok(g) => {
11736 scratch.set_len(e, 0)?;
11737 stream_graph = Some(g);
11738 }
11739 Err(err) => {
11740 scratch.set_len(e, 0)?;
11741 if debug_spec {
11742 eprintln!("[spec] stream-graph capture failed ({err}); stream off");
11743 }
11744 }
11745 }
11746 }
11747 let stream_active = stream_on && stream_graph.is_some();
11748 if debug_spec {
11749 eprintln!(
11750 "[spec] stream_on={stream_on} env={} samp={sampled} dg={} captured={} active={stream_active} session={session_mode} replay={spec_replay}",
11751 crate::spec::spec_stream(),
11752 dctx.graph.is_some(),
11753 stream_graph.is_some()
11754 );
11755 }
11756 let t_v_s = k + 1;
11757 // ROUND-STREAM buffers + ptr tables now live in the model-generic round_stream
11758 // module (extracted 2026-07-12; the gemma burst reuses them).
11759 let sb = crate::round_stream::StreamBufs::new(e, k, crate::spec::spec_stream_m())?;
11760 let crate::round_stream::StreamBufs {
11761 mut vtok_d,
11762 mut brk_d,
11763 mut pend_d,
11764 last_pred_d,
11765 mut pos_ctr,
11766 mut pos_start_d,
11767 mut ring_d,
11768 acc_d: mut stream_acc,
11769 m_rounds,
11770 k: _,
11771 } = sb;
11772 let stream_ptrs: Option<CudaSlice<u64>> = if stream_active {
11773 Some(crate::round_stream::kv_len_ptr_table(
11774 e,
11775 cache,
11776 Some(&pos_ctr),
11777 )?)
11778 } else {
11779 None
11780 };
11781
11782 let t_fill = t_ent.elapsed();
11783 let mut round = 0usize;
11784 // ADAPTIVE DRAFT LENGTH (MEMRA_SPEC_ADAPT=1, opt-in — the gemma_spec accepted-run law,
11785 // ported 2026-08-01): next round's draft depth = last round's accepted run + 1, clamped
11786 // to [floor(pos), k_cap] — a miss shrinks the next draft to the miss point + 1,
11787 // full-accept streaks re-deepen one step per round. NOT the 2026-07-07 acceptance-EMA
11788 // (that arm measured an HONEST LOSS to static per-class optima — 115.0/85.8/73.4 vs
11789 // 121.6/92.7/75.6, EMA lag — and was removed 2026-07-08; rig5090.jsonl has the record).
11790 // The gemma law has no lag class: it reacts within one round, and was worth +7-20% on
11791 // the gemma cells at unchanged exactness (2026-07-10 flip; floor sweep 2026-07-25;
11792 // position key 2026-07-26). Signal = n_acc from the round's EXISTING accept readback —
11793 // zero new syncs; the draft graph is a SINGLE-STEP capture replayed per drafted token,
11794 // so a per-round depth needs no re-capture (unlike gemma's whole-chain graphs). qwen's
11795 // in-round p-min cut already shortens chains mid-round, so gemma's one-round-late p-min
11796 // fold into kc is unnecessary here — the accepted-run law sees the cut via n_acc.
11797 // Exactness is the verify's job at ANY depth (same contract as p-min variable rounds).
11798 // DEFAULT OFF on the qwen path until its cells gate a flip (gemma's is default-on).
11799 // MEASURED 2026-08-01 (H100 GPU-3, interleaved x3, NGEN=256, same-invocation plain
11800 // denominators; research/qwen-adaptive-k-20260801/): REFUTED on the tuned qwen configs.
11801 // q27 K=3+HPOST+PMIN=0.3: short +0.8% (noise; law ~idles, len_hist identical), board
11802 // -1.9%, agentic -0.5%; board PMIN=0 -2.1% (not p-min shadowing — the law itself);
11803 // floor=1 -2.8% (gemma's floor-collapse, reproduced). q35 K=2 board: -6.4% (52/136
11804 // rounds shrink to depth 1; no depth to reclaim at K=2). The gemma direction DOES
11805 // appear at untuned depth-K — q27 K=6 floor=4 +1.5% over fixed K=6 — but stays -3.7%
11806 // below fixed K=3: same verdict class as the retired EMA arm (honest loss to static
11807 // per-class optima). Acceptance-rate rises under the law while tokens/round falls —
11808 // it buys accept-% by adding rounds, and a round's fixed draft+verify cost wins.
11809 // K=1..8 self-consistency PASS both models with the law ON (exactness held).
11810 let adapt = std::env::var("MEMRA_SPEC_ADAPT").as_deref() == Ok("1");
11811 // floor: per-model default keyed on n_embd (gemma's tiering — models with an expensive
11812 // verify keep deep drafts after a miss); MEMRA_SPEC_ADAPT_FLOOR pins it everywhere.
11813 let adapt_floor_env: Option<usize> = std::env::var("MEMRA_SPEC_ADAPT_FLOOR")
11814 .ok()
11815 .and_then(|v| v.parse().ok());
11816 let adapt_floor_default: usize = if self.cfg.n_embd as usize >= 3500 {
11817 4
11818 } else if self.cfg.n_embd as usize >= 2500 {
11819 2
11820 } else {
11821 1
11822 };
11823 let adapt_floor: usize = adapt_floor_env.unwrap_or(adapt_floor_default);
11824 // position key: past floor_ctx a HIGH floor (>=4) relaxes to 1 — forced-deep drafts
11825 // turn net-negative at depth (gemma 31B d1736 evidence); MEMRA_SPEC_FLOOR_CTX moves
11826 // the boundary, an explicit MEMRA_SPEC_ADAPT_FLOOR pins the floor everywhere.
11827 let floor_ctx: usize = std::env::var("MEMRA_SPEC_FLOOR_CTX")
11828 .ok()
11829 .and_then(|v| v.parse().ok())
11830 .unwrap_or(1024);
11831 let floor_at = |pos: usize| -> usize {
11832 if adapt_floor_env.is_some() || pos < floor_ctx {
11833 adapt_floor
11834 } else if adapt_floor >= 4 {
11835 1
11836 } else {
11837 adapt_floor
11838 }
11839 };
11840 // cap: MEMRA_SPEC_CAPMAX (gemma semantics, default 7). Binds only under adapt — the
11841 // fixed-K default path is untouched by this whole block.
11842 let cap_max: usize = std::env::var("MEMRA_SPEC_CAPMAX")
11843 .ok()
11844 .and_then(|v| v.parse().ok())
11845 .unwrap_or(7);
11846 let k_cap = k.min(cap_max).max(1);
11847 let mut kc = k_cap;
11848 // Persistent snapshot buffers are allocated once and refreshed in place.
11849 let mut snap = cache.snapshot(e)?;
11850 // BONUS FOLD (2026-07-04): after a FULL accept the bonus token is NOT committed with a
11851 // separate T=1 trunk pass (a full weight read per round). It stays PENDING and rides as
11852 // column 0 of the NEXT round's verify batch. Under predecessor pairing the next chain
11853 // seeds from the bonus's predecessor's TRUE verify hidden (free — no extra
11854 // pass of any kind). Verify still
11855 // checks every emitted token against the target -> exactness holds by construction; only
11856 // DRAFT QUALITY can shift, which the acceptance numbers arbitrate.
11857 // bonus emitted but not yet committed to cache. A carried pending (see SpecSession::
11858 // pending_tok) enters round 0 directly — the burst boundary becomes a plain round edge.
11859 let mut pending: Option<u32> = carried_pending;
11860 // MEMRA_SPEC_PHASE=1: per-round wall decomposition (draft / verify / accept+commit) —
11861 // no tracing, no extra syncs (each phase is naturally sync-bounded: draft readbacks,
11862 // the verify accept readback). Printed once at loop end via spec-stats.
11863 let anatomy_on = std::env::var("MEMRA_SPEC_PP_ANATOMY").as_deref() == Ok("1");
11864 let phase_on = anatomy_on || std::env::var("MEMRA_SPEC_PHASE").as_deref() == Ok("1");
11865 // MEMRA_SPEC_PHASE_SYNC=1 — reads the phase split correctly, and proves it. `ph_mark` is a
11866 // bare Instant, so `verify-issue` is the host QUEUEING the walk (the GPU is already running
11867 // under it) and `verify-wait` is only the residual drain at the accept readback: one
11868 // overlapped interval cut at the first blocking call, NOT "GPU time" beside "host time".
11869 // Syncing right after the walk is issued moves the whole GPU wall into `verify-issue`. If
11870 // the walk's GPU total is really issue+wait, then with this on verify-issue jumps to that
11871 // sum, verify-wait collapses to the readback alone, and the ROUND WALL DOES NOT MOVE —
11872 // which is what says the queueing time was hidden and is not a target. Diagnostic only.
11873 let phase_sync = std::env::var("MEMRA_SPEC_PHASE_SYNC").as_deref() == Ok("1");
11874 // DRAFT-MASK receipt (lane/draft-mask): speculative-clone wall + rounds, printed with
11875 // spec-stats. The clone is the one cost the design adds per round — measured, not assumed.
11876 let (mut dm_clone_ns, mut dm_rounds) = (0u128, 0usize);
11877 // grammar-truncation counters: how many rounds the verify-side cut fired and how many
11878 // already-verified tokens it threw away. THIS is the quantity draft masking targets.
11879 let (mut dm_cuts, mut dm_cut_tokens) = (0usize, 0usize);
11880 let (mut ph_draft, mut ph_verify, mut ph_rest) = (0f64, 0f64, 0f64);
11881 let mut ph_wait = 0f64;
11882 let mut ph_commit = 0f64;
11883 let mut ph_t = std::time::Instant::now();
11884 let mut ph_mark = |acc: &mut f64, on: bool| {
11885 if on {
11886 let now = std::time::Instant::now();
11887 *acc += (now - ph_t).as_secs_f64();
11888 ph_t = now;
11889 }
11890 };
11891 // MTP-ROUTE VERIFY GRAPHS (`MEMRA_SPEC_VERIFY_GRAPH`, see the flag doc): the
11892 // model-owned capture pool, locked for the whole burst exactly as the dspark serve
11893 // arm holds it — the slab stash is live verify -> commit inside a round, and the
11894 // worker drives rounds from one scheduler thread. PERSISTENT across generations on
11895 // the model (rebuilding per call re-captures the pool per prompt, which is the
11896 // measured way to lose more than the launches cost); the captured bodies are
11897 // cache-independent, every state read going through per-round refreshed pointer
11898 // tables. None = the eager walk, byte-identical.
11899 //
11900 // Never armed together with ROUND-STREAM: the tparallel verify refuses that pair
11901 // loudly, and `stream_active` owns the burst arm above, so the door stays shut
11902 // whenever the stream is live rather than relying on that refusal.
11903 // The lock is taken ONLY when the door is armed: with the flag off this whole block
11904 // is inert, so the default path cannot serialize two spec generations behind a mutex
11905 // it never reads.
11906 let vg_armed =
11907 crate::spec::spec_verify_graph_env().unwrap_or_else(|| self.vgraph_family_default());
11908 let mut vg_guard = if vg_armed && !stream_active {
11909 let mut g = self.dspark_vgraphs.lock().unwrap();
11910 if g.is_none() {
11911 // Size by the WIDEST verify this run can present, which is k+1 and NOT
11912 // k_cap+1: the sampled arm's own window is `t_v_s = k + 1`, so a pool built
11913 // from a smaller adaptive cap gets sliced past its stash rows (a `slice_mut`
11914 // panic in the sampled ON arm, measured before this line said k+1).
11915 let vt_cap = (k.max(k_cap) + 1).max(2);
11916 *g = DsparkVerifyGraphs::new(e, cache, vt_cap, n_embd)?;
11917 if g.is_some() {
11918 // Engagement receipt (the dead-arm lesson): prove the door is LIVE rather
11919 // than trusting that a flag set means a pool built.
11920 eprintln!("[spec-vg] MTP verify-graph pool ENGAGED (vt_cap={vt_cap})");
11921 } else {
11922 eprintln!(
11923 "[spec-vg] MTP verify-graph pool declined (no linear layers, \
11924 non-uniform state, or vt_cap < 2) — eager walk"
11925 );
11926 }
11927 }
11928 Some(g)
11929 } else {
11930 None
11931 };
11932 // Capacity fail-safe: a round wider than the pool was built for must take the eager
11933 // walk, not slice the stash past its rows. The sizing above already covers every
11934 // round this run can present; this keeps a future caller (or a k that grows behind
11935 // the pool's back) on the byte-identical fallback instead of a panic.
11936 let vg_t_cap = vg_guard
11937 .as_ref()
11938 .and_then(|g| g.as_ref())
11939 .map(|g| g.t_capacity())
11940 .unwrap_or(0);
11941 let mut graph_guard_noted = false;
11942 while keep_going && out.len() < max_new {
11943 // GRAPH-LAUNCH HEADROOM GUARD (see GRAPH_LAUNCH_MIN_FREE): below the floor,
11944 // every captured-graph arm in this round yields to its byte-identical eager
11945 // twin instead of feeding cuGraphLaunch a card it segfaults on.
11946 let graph_round_ok = graph_launch_headroom_ok(e);
11947 if !graph_round_ok && !graph_guard_noted {
11948 graph_guard_noted = true;
11949 eprintln!(
11950 "[spec] graph replay suspended: driver free below the {}MB launch floor \
11951 (eager arms serve; cuGraphLaunch segfaults into an exhausted card)",
11952 GRAPH_LAUNCH_MIN_FREE / (1 << 20)
11953 );
11954 }
11955 // MEMRA_SPEC_ROUND_PROF=1: wall of the WHOLE round against the pieces we already
11956 // instrument. Needed because the parts do not add up: the draft step measures 1.27 ms
11957 // ([spec-anatomy] glue 92 / attn 280 / ffn 222 / head 670 us) and the t=2 verify walk
11958 // 25.6 ms ([tcol-prof] attn 10.1 + ffn 15.3), yet a K=1 round takes 177 ms on the
11959 // step37 TP2 stack. This prints where the other ~150 ms lives.
11960 let round_prof = ROUND_PROF
11961 .get_or_init(|| std::env::var("MEMRA_SPEC_ROUND_PROF").as_deref() == Ok("1"));
11962 let round_t0 = round_prof.then(std::time::Instant::now);
11963 // ROUND-STREAM BURST: from round 1 (pending guaranteed by every non-replay arm),
11964 // issue M rounds with zero readbacks, then drain the ring + reconcile mirrors.
11965 if let (true, Some(sg), Some(ptrs)) = (
11966 stream_active && round >= 1 && pending.is_some() && graph_round_ok,
11967 &stream_graph,
11968 &stream_ptrs,
11969 ) {
11970 if debug_spec {
11971 static ONCE: std::sync::Once = std::sync::Once::new();
11972 ONCE.call_once(|| {
11973 eprintln!("[memra] ROUND-STREAM burst engaged (M={m_rounds} k={k})")
11974 });
11975 }
11976 e.set_i32_one(&mut pos_ctr, cache.pos as i32)?;
11977 e.set_u32_one(&mut pend_d, pending.unwrap())?;
11978 e.set_u32_one(&mut ring_d, 0)?; // ring count = 0 (writes element 0)
11979 for _mi in 0..m_rounds {
11980 e.i32_copy_add(&pos_ctr, &mut pos_start_d, 0)?;
11981 cache.snapshot_into(e, &mut snap)?; // device D2Ds, stream-ordered
11982 e.i32_copy_add(&pos_ctr, &mut scratch.kv.len_d, 0)?; // draft-KV rollback
11983 e.i32_copy_add(&pos_ctr, &mut dctx.g_pos, 1)?; // rope pos = pos + base
11984 e.u32_copy(&pend_d, &mut dctx.g_tok)?;
11985 e.copy_into(&mut dctx.g_seed, 0, &h_seed_buf, n_embd)?;
11986 sg.launch()?;
11987 e.spec_assemble_verify(
11988 &g_tokp2k,
11989 &pend_d,
11990 d2t_dev.as_ref(),
11991 &mut vtok_d,
11992 &mut brk_d,
11993 p_min,
11994 k,
11995 pmin0,
11996 )?;
11997 let mut ck = VerifyCkpt::new(self.layers.len());
11998 let dummy = vec![0u32; t_v_s];
11999 let (tl_d, vx) = self.decode_step_t_core_stream(
12000 e,
12001 &dummy,
12002 0,
12003 &mut *cache,
12004 embd_dev,
12005 Some(&mut ck),
12006 Some((&vtok_d, &pos_ctr)),
12007 None,
12008 None,
12009 )?;
12010 for j in 0..t_v_s {
12011 e.argmax_token_device_col(&tl_d, j, n_vocab, &mut preds_d, j)?;
12012 }
12013 e.spec_accept_greedy_dc(
12014 &preds_d,
12015 &vtok_d,
12016 &last_pred_d,
12017 &brk_d,
12018 &mut stream_acc,
12019 )?;
12020 e.spec_seed_gather(&vx, &fill_prev, &stream_acc, &mut h_seed_buf, 1, n_embd)?;
12021 e.copy_into(&mut fill_prev, 0, &h_seed_buf, n_embd)?;
12022 self.commit_verified_prefix_stream(
12023 e,
12024 &mut *cache,
12025 &snap,
12026 &ck,
12027 &stream_acc,
12028 1,
12029 t_v_s,
12030 )?;
12031 e.spec_rollback_stream(
12032 ptrs,
12033 &pos_start_d,
12034 &stream_acc,
12035 1,
12036 self.layers.len() + 1,
12037 )?;
12038 e.spec_ring_commit(&vtok_d, &stream_acc, &brk_d, &mut ring_d, &mut pend_d)?;
12039 }
12040 e.stream().synchronize()?;
12041 let ring_h = e.dtoh_u32(&ring_d)?;
12042 let cnt = ring_h[0] as usize;
12043 for i in 0..cnt {
12044 if out.len() < max_new {
12045 out.push(ring_h[1 + i]);
12046 }
12047 }
12048 let pos_h = e.dtoh_i32(&pos_ctr)?[0] as usize;
12049 for il in 0..self.layers.len() {
12050 if let Some(kvl) = cache.kv[il].as_mut() {
12051 kvl.len = pos_h;
12052 }
12053 }
12054 cache.pos = pos_h;
12055 scratch.kv.len = pos_h;
12056 pending = Some(ring_h[cnt]); // last drained token = the live bonus
12057 last_token = ring_h[cnt];
12058 total_drafted += k * m_rounds; // upper bound (p-min breaks uncounted)
12059 total_accepted += cnt.saturating_sub(m_rounds);
12060 if let Some(t) = sess_telem {
12061 // totals only — the burst's per-round accept counts stayed on device
12062 // (that is the point of the round-stream arm). pos_* untouched.
12063 t.record_totals(m_rounds, k * m_rounds, cnt.saturating_sub(m_rounds));
12064 }
12065 round += m_rounds;
12066 // sse-cadence: the drained ring is committed — flush it at burst-drain cadence.
12067 keep_going = flush_commit(&mut on_commit, &out, &mut flushed);
12068 continue;
12069 }
12070 let pos = cache.pos; // #tokens committed (EXCLUDES a pending bonus)
12071 cache.snapshot_into(e, &mut snap)?; // §C: snapshot BEFORE draft+verify
12072 ph_mark(&mut ph_rest, phase_on);
12073
12074 // --- 1. DRAFT k tokens with the NextN head (autoregressive, T=1 each) ---
12075 // p-min semantics (both paths): stop the chain early when the head's confidence in
12076 // its own pick drops below p_min — the just-drafted token is DISCARDED, but its
12077 // scratch append stands (identical to the eager chain's ordering). j==0 always drafts.
12078 let base0 = if pending.is_some() { 1usize } else { 0usize };
12079 // fixed draft length by default; MEMRA_SPEC_ADAPT=1 drafts at last round's
12080 // accepted run + 1 (the gemma law — see the setup block above the loop).
12081 let k_this = if adapt { kc } else { k };
12082 let mut draft: Vec<u32> = Vec::with_capacity(k);
12083 let mut draft_idx: Vec<u32> = Vec::with_capacity(k); // trimmed-vocab ids (== draft when untrimmed)
12084 {
12085 // Round-start draft-KV sync (BOTH paths). Persistent: truncate/align to the committed
12086 // history — slots 0..P hold entries for the tokens before last_token@P (P = pos +
12087 // base0 - 1); this single set_len IS the draft-side rollback (drops last round's
12088 // rejected drafts and p-min extras via the len mechanism).
12089 scratch.set_len(e, pos + base0 - 1)?;
12090 // dcw door: a captured chain appends k_this device-counter rows (plus the
12091 // pseudo-seed replay) with no host intervention; any ring rebase those appends
12092 // could need happens HERE, host-side, before the replays. The eager arm keeps
12093 // its own per-step prepare, so this is graph-path-only work.
12094 if step35_draft_dcw_on()
12095 && (dctx.graph.is_some()
12096 || dctx.graph_s.is_some()
12097 || dctx.chain.is_some()
12098 || dctx.chain_s.is_some())
12099 {
12100 scratch.ensure_dcw_headroom(e, k_this + 2)?;
12101 }
12102 if pen_on {
12103 // PEN_WINDOW_MAX also bounds the per-round upload and the O(n_hist^2)
12104 // device dedup: the serve window is already PEN_WINDOW_MAX, and this
12105 // defensive min also bounds non-server callers.
12106 let win = sp.penalty_last_n.min(PEN_WINDOW_MAX);
12107 let w0 = pen_hist.len().saturating_sub(win);
12108 pen_hist_d = Some(e.htod_u32_v(&pen_hist[w0..])?);
12109 }
12110 if sampled {
12111 draft_logits.clear();
12112 draft_stats.clear();
12113 }
12114 // DRAFT-SIDE GRAMMAR MASK: clone the committed grammar state ONCE per round; each
12115 // position's mask is computed on that clone and advanced by the PROPOSED token. The
12116 // real state moves only on emission (verify's job), so the emitted stream is
12117 // unchanged — the mask only removes tokens the verify would have truncated anyway.
12118 let mut dmask_live = dmask_on;
12119 if dmask_live {
12120 let t_c = std::time::Instant::now();
12121 constraint
12122 .as_deref_mut()
12123 .unwrap()
12124 .draft_begin()
12125 .map_err(|e2| format!("constraint: {e2}"))?;
12126 dm_clone_ns += t_c.elapsed().as_nanos();
12127 dm_rounds += 1;
12128 }
12129 if let (false, Some(cg)) = (sampled || pen_on || !graph_round_ok, &dctx.chain) {
12130 // GREEDY CHAIN GRAPH (lane/step37-draft-graph-serving-20260830): the
12131 // eager multi-head chain's EXACT launch order — step j rewinds head
12132 // (j % heads)'s plane to the committed length and replays rows 0..=j —
12133 // with each row's whole head-forward as ONE graph launch. The chain
12134 // POLICY (head choice, prefix length, stored-seed feed) is host-side,
12135 // identical to `mtp_chain_forward_dev`, so graph-vs-eager drafts are
12136 // bit-identical by construction (same launcher, same bucket — the dcw
12137 // parity contract). Interior rows launch the head-less graph: their
12138 // logits are dead in the eager chain too, so the consumed bytes match.
12139 let heads_n = self.mtp_head_count();
12140 let committed = pos + base0 - 1;
12141 let mut chain_tokens: Vec<u32> = vec![last_token];
12142 let mut chain_seed_bufs: Vec<CudaSlice<f32>> = vec![e.clone_dtod(&h_seed_buf)?];
12143 for j in 0..k_this {
12144 let index = mtp_chain_head_index(j, heads_n);
12145 if debug_spec {
12146 eprintln!(
12147 "[mtp-chain-step] round={round} j={j} head={index} \
12148 replay_rows={} arm=graph",
12149 chain_tokens.len(),
12150 );
12151 }
12152 scratch.set_plane_len(e, index, committed)?;
12153 e.set_i32_one(&mut dctx.g_pos, (committed + 1) as i32)?;
12154 for row in 0..=j {
12155 e.set_u32_one(&mut dctx.g_tok, chain_tokens[row])?;
12156 e.copy_into(&mut dctx.g_seed, 0, &chain_seed_bufs[row], n_embd)?;
12157 if row < j {
12158 cg.interior[index].launch()?;
12159 } else {
12160 // per-position mask upload before the LAST row only — the
12161 // eager chain applies the mask on is_last exactly the same.
12162 if dmask_live
12163 && !upload_draft_mask(
12164 e,
12165 constraint.as_deref_mut().unwrap(),
12166 &mut dctx.g_dmask,
12167 mtp.d2t.as_ref(),
12168 d_vocab,
12169 dmask_words,
12170 )?
12171 {
12172 e.htod_u32_into(
12173 &mut dctx.g_dmask,
12174 &vec![u32::MAX; dmask_words],
12175 )?;
12176 dmask_live = false;
12177 }
12178 cg.last[index].launch()?;
12179 }
12180 // host mirror (len_d advanced in-graph by the dcw append)
12181 scratch.plane_mut(index).0.len += 1;
12182 }
12183 let idx = e.dtoh_u32_one(&dctx.g_tok)?;
12184 // #87 SENTINEL TRAP (see the single-head graph arm below).
12185 if (idx as usize) >= d_vocab {
12186 let seed_h = e.dtoh(&dctx.g_seed)?;
12187 let seed_nan = seed_h.iter().filter(|v| v.is_nan()).count();
12188 return Err(format!(
12189 "draft(chain-graph) argmax sentinel 0x{idx:08x} >= d_vocab \
12190 {d_vocab} at round {round} j={j} head={index} pos={pos}: \
12191 head-out NaN {seed_nan}/{n_embd} — refusing to dereference \
12192 the embed row (#87 trap)"
12193 )
12194 .into());
12195 }
12196 // multi-head MTP forbids a trimmed head (validated at entry), so the
12197 // draft index IS the target id; keep the map for uniformity.
12198 let d = match &mtp.d2t {
12199 Some(map) => map[idx as usize],
12200 None => idx,
12201 };
12202 let draft_p = if p_min > 0.0 {
12203 Some(e.dtoh(&dctx.g_p)?[0])
12204 } else {
12205 None
12206 };
12207 if let Some(p) = draft_p.filter(|_| p_min > 0.0)
12208 && p < p_min
12209 && (j > 0 || (pmin0 && base0 == 1))
12210 {
12211 break;
12212 }
12213 draft.push(d);
12214 chain_tokens.push(d);
12215 // step j's h_nextn: the last-row graph self-fed it into g_seed —
12216 // snapshot it as the chain history seed for row j+1 (stream-ordered
12217 // after the launch, exactly the eager chain's chain_seeds push).
12218 chain_seed_bufs.push(e.clone_dtod(&dctx.g_seed)?);
12219 // speculative grammar advance (see the single-head graph arm).
12220 if dmask_live
12221 && !constraint
12222 .as_deref_mut()
12223 .unwrap()
12224 .draft_advance(d)
12225 .map_err(|e2| format!("constraint: {e2}"))?
12226 {
12227 e.htod_u32_into(&mut dctx.g_dmask, &vec![u32::MAX; dmask_words])?;
12228 break;
12229 }
12230 }
12231 } else if let (true, Some(cg)) = (
12232 sampled && s_capturable && dctx.s_key == Some(s_key) && graph_round_ok,
12233 &dctx.chain_s,
12234 ) {
12235 if skey_probe() {
12236 eprintln!(
12237 "[skey] chain=graph_chain_s round={round} capturable={} top_k={} \
12238 top_p={} min_p={} s_key_parked={:?}",
12239 s_capturable as u8, sp.top_k, sp.top_p, sp.min_p, dctx.s_key,
12240 );
12241 }
12242 // SAMPLED CHAIN GRAPH: the greedy chain arm's launch order with the
12243 // sampled last-row graphs — in-graph counter bump + (filtered) gumbel
12244 // draw + argmax; q retained per step into q_slots exactly like the
12245 // single-head sampled graph arm. Counter continuity: g_ctr host-seeded
12246 // to sctr-1 once per ROUND; each step's last-row graph bumps it BEFORE
12247 // the perturb, so step j consumes counter sctr+j — the eager Philox
12248 // stream (interior rows never draw, never bump).
12249 let heads_n = self.mtp_head_count();
12250 let committed = pos + base0 - 1;
12251 let filtered_stats_in_graph = s_key.filtered();
12252 let mut chain_tokens: Vec<u32> = vec![last_token];
12253 let mut chain_seed_bufs: Vec<CudaSlice<f32>> = vec![e.clone_dtod(&h_seed_buf)?];
12254 e.set_u32_one(&mut dctx.g_ctr, sctr.wrapping_sub(1))?;
12255 for j in 0..k_this {
12256 let index = mtp_chain_head_index(j, heads_n);
12257 if debug_spec {
12258 eprintln!(
12259 "[mtp-chain-step] round={round} j={j} head={index} \
12260 replay_rows={} arm=graph_s",
12261 chain_tokens.len(),
12262 );
12263 }
12264 scratch.set_plane_len(e, index, committed)?;
12265 e.set_i32_one(&mut dctx.g_pos, (committed + 1) as i32)?;
12266 for row in 0..=j {
12267 e.set_u32_one(&mut dctx.g_tok, chain_tokens[row])?;
12268 e.copy_into(&mut dctx.g_seed, 0, &chain_seed_bufs[row], n_embd)?;
12269 if row < j {
12270 cg.interior[index].launch()?;
12271 } else {
12272 cg.last[index].launch()?;
12273 }
12274 scratch.plane_mut(index).0.len += 1;
12275 }
12276 sctr += 1; // mirrors the in-graph g_ctr bump (eager parity:
12277 // counts the p-min-discarded token too)
12278 // q retention: ONE async D2D of the persistent head-logits buffer
12279 // into this round's slot j (stream-ordered after the replay).
12280 e.copy_into(&mut dctx.q_slots[j], 0, &dctx.g_q, d_vocab)?;
12281 // FILTERED capture: read the in-graph filter_stats scalars back per
12282 // replay instead of a second full-vocab filter_stats per slot post-
12283 // chain — bit-exact (the values the in-graph perturb consumed) and
12284 // measured worth ~5% of vendor-default serving tok/s at K=3. Before
12285 // the p-min break so the discarded slot's stats land too.
12286 if filtered_stats_in_graph {
12287 draft_stats.push((
12288 e.dtoh(&dctx.g_mx)?[0],
12289 e.dtoh(&dctx.g_th)?[0],
12290 e.dtoh(&dctx.g_z)?[0],
12291 ));
12292 }
12293 let idx = e.dtoh_u32_one(&dctx.g_tok)?;
12294 // #87 SENTINEL TRAP (see the single-head graph arms).
12295 if (idx as usize) >= d_vocab {
12296 let seed_h = e.dtoh(&dctx.g_seed)?;
12297 let seed_nan = seed_h.iter().filter(|v| v.is_nan()).count();
12298 return Err(format!(
12299 "draft(chain-graph-sampled) argmax sentinel 0x{idx:08x} >= \
12300 d_vocab {d_vocab} at round {round} j={j} head={index} pos={pos}: \
12301 head-out NaN {seed_nan}/{n_embd} — refusing to dereference the \
12302 embed row (#87 trap)"
12303 )
12304 .into());
12305 }
12306 let d = match &mtp.d2t {
12307 Some(map) => map[idx as usize],
12308 None => idx,
12309 };
12310 draft_idx.push(idx);
12311 if p_min > 0.0 {
12312 let p = e.dtoh(&dctx.g_p)?[0];
12313 if p < p_min && (j > 0 || (pmin0 && base0 == 1)) {
12314 break;
12315 }
12316 }
12317 draft.push(d);
12318 chain_tokens.push(d);
12319 chain_seed_bufs.push(e.clone_dtod(&dctx.g_seed)?);
12320 }
12321 // PURE-TEMP accept path: stats per used slot recomputed from the RETAINED
12322 // q with the SAME filter_stats program the eager arm runs (deployment-
12323 // keyed coop/plain choice, same input bits). The FILTERED graph read its
12324 // stats back per replay above.
12325 if !filtered_stats_in_graph {
12326 for j in 0..draft.len().max(draft_idx.len()) {
12327 let rows0 = e.htod_i32(&[0])?;
12328 let (mut th_d, mut z_d, mut mx_d) =
12329 (e.zeros(1)?, e.zeros(1)?, e.zeros(1)?);
12330 e.filter_stats(
12331 &dctx.q_slots[j],
12332 d_vocab,
12333 &rows0,
12334 &mut th_d,
12335 &mut z_d,
12336 &mut mx_d,
12337 d_vocab,
12338 1,
12339 sp_temp,
12340 sp.top_k,
12341 sp.top_p,
12342 sp.min_p,
12343 )?;
12344 draft_stats.push((
12345 e.dtoh(&mx_d)?[0],
12346 e.dtoh(&th_d)?[0],
12347 e.dtoh(&z_d)?[0],
12348 ));
12349 }
12350 }
12351 } else if let (false, Some(gr)) =
12352 (sampled || pen_on || !graph_round_ok, &dctx.graph)
12353 {
12354 // GRAPH DRAFT: one dispatch per drafted token. The chain feeds itself on-device
12355 // (in-graph argmax -> tok_d -> next replay's embed; h_nextn -> h_seed_d; pos_d
12356 // inc'd in-graph); the host only reads 4B token (+4B p) and decides the break.
12357 e.set_i32_one(&mut dctx.g_pos, (pos + base0) as i32)?;
12358 e.set_u32_one(&mut dctx.g_tok, last_token)?;
12359 e.copy_into(&mut dctx.g_seed, 0, &h_seed_buf, n_embd)?;
12360 for j in 0..k_this {
12361 // per-position mask upload (contents only — the graph's baked pointer is
12362 // dctx.g_dmask). All-ones once masking goes dead mid-chain, so the captured
12363 // mask node degrades to a no-op ban instead of needing a second graph.
12364 if dmask_live
12365 && !upload_draft_mask(
12366 e,
12367 constraint.as_deref_mut().unwrap(),
12368 &mut dctx.g_dmask,
12369 mtp.d2t.as_ref(),
12370 d_vocab,
12371 dmask_words,
12372 )?
12373 {
12374 // no draft-vocab row is grammar-legal here (a trimmed FR-Spec head can
12375 // genuinely miss the legal set): neutralize the captured mask node and
12376 // finish the chain UNMASKED — exactly pre-lane behaviour, never worse.
12377 e.htod_u32_into(&mut dctx.g_dmask, &vec![u32::MAX; dmask_words])?;
12378 dmask_live = false;
12379 }
12380 gr.launch()?;
12381 scratch.kv.len += 1; // host mirror (len_d advanced in-graph)
12382 let idx = e.dtoh_u32_one(&dctx.g_tok)?;
12383 // #87 SENTINEL TRAP: an all-NaN head-logits row leaves the device argmax's
12384 // init sentinel (0x7FFFFFFF) in g_tok — feeding it onward dereferences
12385 // embed_row(sentinel) = table + ~4.6TB (never mapped) inside the NEXT graph
12386 // replay's embed node, and the MMU fault kills the CUDA context for the
12387 // whole process (research/pp2spec-crash-20260807: 3 coredumps, byte-exact
12388 // VA arithmetic). Refuse loudly instead; the diagnostics name the first-NaN
12389 // buffer (g_seed = the verify-side handoff vs head-side compute).
12390 if (idx as usize) >= d_vocab {
12391 // g_seed is SELF-FED (the replay writes h_nextn back into it), so it
12392 // reads as the head's OUTPUT at j; h_seed_buf is the round's INPUT
12393 // seed, untouched since the round-start copy — the pair discriminates
12394 // "seed arrived poisoned" from "head forward produced NaN".
12395 let seed_h = e.dtoh(&dctx.g_seed)?;
12396 let seed_nan = seed_h.iter().filter(|v| v.is_nan()).count();
12397 let in_h = e.dtoh(&h_seed_buf)?;
12398 let in_nan = in_h.iter().filter(|v| v.is_nan()).count();
12399 return Err(format!(
12400 "draft(graph) argmax sentinel 0x{idx:08x} >= d_vocab {d_vocab} at \
12401 round {round} j={j} pos={pos}: head-out NaN {seed_nan}/{n_embd}, \
12402 round-input-seed NaN {in_nan}/{n_embd} — refusing to dereference \
12403 the embed row (#87 trap)"
12404 )
12405 .into());
12406 }
12407 // trimmed draft vocab -> target token id (identity when no d2t map)
12408 let d = match &mtp.d2t {
12409 Some(map) => map[idx as usize],
12410 None => idx,
12411 };
12412 let draft_p = if p_min > 0.0 {
12413 Some(e.dtoh(&dctx.g_p)?[0])
12414 } else {
12415 None
12416 };
12417 if let Some(p) = draft_p.filter(|_| p_min > 0.0)
12418 && p < p_min
12419 && (j > 0 || (pmin0 && base0 == 1))
12420 {
12421 break;
12422 }
12423 draft.push(d);
12424 // with a trimmed head the NEXT embed must read the TARGET id, not the draft
12425 // index the argmax wrote — patch the persistent token buffer (4B htod).
12426 if d != idx {
12427 e.set_u32_one(&mut dctx.g_tok, d)?;
12428 }
12429 // advance the SPECULATIVE state with the proposal; a dead chain drops to
12430 // unmasked drafting for the remaining positions (verify still arbitrates).
12431 // speculative advance; a chain the grammar can no longer follow (EOS
12432 // proposed) ends here. The captured mask node always runs, so a dead chain
12433 // leaves the buffer NEUTRAL (all-ones = ban nothing) before it exits.
12434 if dmask_live
12435 && !constraint
12436 .as_deref_mut()
12437 .unwrap()
12438 .draft_advance(d)
12439 .map_err(|e2| format!("constraint: {e2}"))?
12440 {
12441 e.htod_u32_into(&mut dctx.g_dmask, &vec![u32::MAX; dmask_words])?;
12442 break;
12443 }
12444 }
12445 // REGIME RE-TEST (lane/graph-s-key-exactness-20260819, widened by
12446 // lane/step37-draft-graph-serving-20260830): the sampled graph is legal ONLY
12447 // in the regime it was captured in. The condition used to read
12448 // `(sampled, &dctx.graph_s)` and trusted `s_key` to have dropped anything
12449 // else — which it could not, because the key omitted the filters. Both
12450 // halves are enforced: the key drops a stale graph, and this site refuses to
12451 // launch one whose key differs or whose regime is uncapturable (penalties).
12452 } else if let (true, Some(gr)) = (
12453 sampled && s_capturable && dctx.s_key == Some(s_key) && graph_round_ok,
12454 &dctx.graph_s,
12455 ) {
12456 if skey_probe() {
12457 eprintln!(
12458 "[skey] chain=graph_s round={round} pure_temp={} capturable={} \
12459 top_k={} top_p={} min_p={} s_key_parked={:?}",
12460 pure_temp as u8,
12461 s_capturable as u8,
12462 sp.top_k,
12463 sp.top_p,
12464 sp.min_p,
12465 dctx.s_key,
12466 );
12467 }
12468 // SAMPLED GRAPH DRAFT: one replay per drafted token — head forward + gumbel +
12469 // argmax in ONE dispatch; the host reads 4B token (+4B p), D2Ds q into slot j,
12470 // and decides the break. Event-counter continuity: g_ctr is host-seeded to
12471 // sctr-1 ONCE per round (outside the graph); the in-graph bump runs BEFORE the
12472 // perturb, so replay j consumes counter sctr+j — exactly the eager arm's Philox
12473 // stream. Host sctr advances in lockstep (computed, no readback needed).
12474 e.set_i32_one(&mut dctx.g_pos, (pos + base0) as i32)?;
12475 e.set_u32_one(&mut dctx.g_tok, last_token)?;
12476 e.copy_into(&mut dctx.g_seed, 0, &h_seed_buf, n_embd)?;
12477 e.set_u32_one(&mut dctx.g_ctr, sctr.wrapping_sub(1))?;
12478 let filtered_stats_in_graph = s_key.filtered();
12479 for j in 0..k_this {
12480 gr.launch()?;
12481 scratch.kv.len += 1; // host mirror (len_d advanced in-graph)
12482 sctr += 1; // mirrors the in-graph g_ctr bump (eager parity:
12483 // counts the p-min-discarded token too)
12484 // q retention: ONE async D2D of the persistent head-logits buffer into this
12485 // round's slot j (stream-ordered after the replay, before the next one).
12486 e.copy_into(&mut dctx.q_slots[j], 0, &dctx.g_q, d_vocab)?;
12487 // FILTERED capture: the replay's own filter_stats node already computed
12488 // (th, z, mx) — read the three scalars back instead of paying a SECOND
12489 // full-vocab filter_stats per slot post-chain (measured ~5% of vendor-
12490 // default serving tok/s at K=3). Bit-exact by construction: these are
12491 // the very values the in-graph perturb consumed. Read BEFORE the p-min
12492 // break so the discarded slot's stats land too (accept-path indexing).
12493 if filtered_stats_in_graph {
12494 draft_stats.push((
12495 e.dtoh(&dctx.g_mx)?[0],
12496 e.dtoh(&dctx.g_th)?[0],
12497 e.dtoh(&dctx.g_z)?[0],
12498 ));
12499 }
12500 let idx = e.dtoh_u32_one(&dctx.g_tok)?;
12501 // #87 SENTINEL TRAP (see the greedy graph arm above).
12502 if (idx as usize) >= d_vocab {
12503 let seed_h = e.dtoh(&dctx.g_seed)?;
12504 let seed_nan = seed_h.iter().filter(|v| v.is_nan()).count();
12505 return Err(format!(
12506 "draft(graph-sampled) argmax sentinel 0x{idx:08x} >= d_vocab \
12507 {d_vocab} at round {round} j={j} pos={pos}: round-seed NaN \
12508 {seed_nan}/{n_embd} — refusing to dereference the embed row \
12509 (#87 trap)"
12510 )
12511 .into());
12512 }
12513 let d = match &mtp.d2t {
12514 Some(map) => map[idx as usize],
12515 None => idx,
12516 };
12517 draft_idx.push(idx);
12518 if p_min > 0.0 {
12519 let p = e.dtoh(&dctx.g_p)?[0];
12520 if p < p_min && (j > 0 || (pmin0 && base0 == 1)) {
12521 break;
12522 }
12523 }
12524 draft.push(d);
12525 // trimmed head: the NEXT embed must read the TARGET id (see the greedy arm).
12526 if d != idx {
12527 e.set_u32_one(&mut dctx.g_tok, d)?;
12528 }
12529 }
12530 // PURE-TEMP accept path: fill draft_stats per used slot post-chain (the
12531 // stats degenerate to th=0 / full-Z; one filter_stats launch per slot).
12532 // The FILTERED graph read its stats back per replay above.
12533 if !filtered_stats_in_graph {
12534 for j in 0..draft.len().max(draft_idx.len()) {
12535 let rows0 = e.htod_i32(&[0])?;
12536 let (mut th_d, mut z_d, mut mx_d) =
12537 (e.zeros(1)?, e.zeros(1)?, e.zeros(1)?);
12538 e.filter_stats(
12539 &dctx.q_slots[j],
12540 d_vocab,
12541 &rows0,
12542 &mut th_d,
12543 &mut z_d,
12544 &mut mx_d,
12545 d_vocab,
12546 1,
12547 sp_temp,
12548 sp.top_k,
12549 sp.top_p,
12550 sp.min_p,
12551 )?;
12552 draft_stats.push((
12553 e.dtoh(&mx_d)?[0],
12554 e.dtoh(&th_d)?[0],
12555 e.dtoh(&z_d)?[0],
12556 ));
12557 }
12558 }
12559 } else {
12560 if skey_probe() && sampled {
12561 eprintln!(
12562 "[skey] chain=eager round={round} pure_temp={} top_k={} \
12563 top_p={} min_p={} s_key_parked={:?}",
12564 pure_temp as u8, sp.top_k, sp.top_p, sp.min_p, dctx.s_key,
12565 );
12566 }
12567 // EAGER DRAFT (fallback: MoE head/trunk, huge k, MEMRA_SPEC_NOGRAPH, capture fail).
12568 let chain_heads = !self.mtp_extra.is_empty();
12569 let mut e_tok = last_token;
12570 let mut d_seed = e.clone_dtod(&h_seed_buf)?;
12571 let mut chain_tokens = if chain_heads {
12572 vec![last_token]
12573 } else {
12574 Vec::new()
12575 };
12576 let mut chain_seeds = if chain_heads {
12577 vec![e.clone_dtod(&h_seed_buf)?]
12578 } else {
12579 Vec::new()
12580 };
12581 for j in 0..k_this {
12582 // GPU-ARGMAX DRAFT (2026-07-03): device logits + device argmax + 4-byte token
12583 // read instead of the ~600KB full-vocab dtoh + host argmax per draft token.
12584 let mtp_pos = pos + base0 + j;
12585 // draft-side grammar mask (eager twin of the graph arm's in-graph node).
12586 // A position with no legal draft-vocab row drops to unmasked drafting for
12587 // the rest of the chain (pre-lane behaviour; verify still arbitrates).
12588 if dmask_live {
12589 dmask_live = upload_draft_mask(
12590 e,
12591 constraint.as_deref_mut().unwrap(),
12592 &mut dctx.g_dmask,
12593 mtp.d2t.as_ref(),
12594 d_vocab,
12595 dmask_words,
12596 )?;
12597 }
12598 let mask = if dmask_live {
12599 Some((&dctx.g_dmask, dmask_words))
12600 } else {
12601 None
12602 };
12603 let (dl_d, h_nextn) = if chain_heads {
12604 if debug_spec {
12605 eprintln!(
12606 "[mtp-chain-step] round={round} j={j} head={} replay_rows={}",
12607 mtp_chain_head_index(j, self.mtp_head_count()),
12608 chain_tokens.len(),
12609 );
12610 }
12611 self.mtp_chain_forward_dev(
12612 e,
12613 &chain_tokens,
12614 &chain_seeds,
12615 &mut *scratch,
12616 pos + base0 - 1,
12617 embd_dev,
12618 mask,
12619 )?
12620 } else {
12621 self.mtp_head_forward_dev(
12622 e,
12623 mtp,
12624 e_tok,
12625 &d_seed,
12626 &mut *scratch,
12627 mtp_pos,
12628 embd_dev,
12629 mask,
12630 )?
12631 };
12632 let tok_d = if sampled {
12633 // FILTERED Gumbel-max: stats -> masked perturb -> argmax = one draw from
12634 // the filtered softmax (filters off => th=0, exact v1 semantics).
12635 if perturb_buf.is_none() {
12636 perturb_buf = Some(e.zeros(d_vocab.max(n_vocab))?);
12637 }
12638 let mut q_row = e.clone_dtod(&dl_d)?; // retained q (penalized when on)
12639 if pen_on {
12640 let h = pen_hist_d.as_ref().unwrap();
12641 let nh = h.len();
12642 e.penalize_logits(
12643 &mut q_row,
12644 h,
12645 nh,
12646 sp.penalty_repeat,
12647 sp.penalty_freq,
12648 sp.penalty_present,
12649 d_vocab,
12650 )?;
12651 }
12652 let rows0 = e.htod_i32(&[0])?;
12653 let (mut th_d, mut z_d, mut mx_d) =
12654 (e.zeros(1)?, e.zeros(1)?, e.zeros(1)?);
12655 e.filter_stats(
12656 &q_row, d_vocab, &rows0, &mut th_d, &mut z_d, &mut mx_d, d_vocab,
12657 1, sp_temp, sp.top_k, sp.top_p, sp.min_p,
12658 )?;
12659 let (th, z, mx) =
12660 (e.dtoh(&th_d)?[0], e.dtoh(&z_d)?[0], e.dtoh(&mx_d)?[0]);
12661 let pb = perturb_buf.as_mut().unwrap();
12662 e.gumbel_perturb_filtered(
12663 &q_row, pb, d_vocab, sp_seed, sctr, sp_temp, mx, th,
12664 )?;
12665 sctr += 1;
12666 draft_logits.push(q_row);
12667 draft_stats.push((mx, th, z));
12668 e.argmax_token_device(pb, d_vocab)?
12669 } else {
12670 e.argmax_token_device(&dl_d, d_vocab)?
12671 };
12672 let idx = e.dtoh_u32_one(&tok_d)?;
12673 // #87 SENTINEL TRAP (eager twin — see the graph arm). Extra diagnostics
12674 // here because the eager chain's operands are all readable: dl_d (the head
12675 // logits row) and d_seed (this step's h_seed) name the first-NaN buffer.
12676 if (idx as usize) >= d_vocab {
12677 let dl_h = e.dtoh(&dl_d)?;
12678 let dl_nan = dl_h.iter().filter(|v| v.is_nan()).count();
12679 let seed_h = if chain_heads {
12680 e.dtoh(chain_seeds.last().unwrap())?
12681 } else {
12682 e.dtoh(&d_seed)?
12683 };
12684 let seed_nan = seed_h.iter().filter(|v| v.is_nan()).count();
12685 return Err(format!(
12686 "draft(eager) argmax sentinel 0x{idx:08x} >= d_vocab {d_vocab} at \
12687 round {round} j={j} pos={pos}: head-logits NaN {dl_nan}/{d_vocab}, \
12688 step-seed NaN {seed_nan}/{n_embd} — refusing to dereference the \
12689 embed row (#87 trap)"
12690 )
12691 .into());
12692 }
12693 let d = match &mtp.d2t {
12694 Some(map) => map[idx as usize],
12695 None => idx,
12696 };
12697 if sampled {
12698 draft_idx.push(idx);
12699 }
12700 let draft_p = if p_min > 0.0 {
12701 let p_d = e.prob_of_token_device(&dl_d, &tok_d, d_vocab)?;
12702 Some(e.dtoh(&p_d)?[0])
12703 } else {
12704 None
12705 };
12706 if let Some(p) = draft_p.filter(|_| p_min > 0.0)
12707 && p < p_min
12708 && (j > 0 || (pmin0 && base0 == 1))
12709 {
12710 break;
12711 }
12712 draft.push(d);
12713 if chain_heads {
12714 chain_tokens.push(d);
12715 chain_seeds.push(h_nextn);
12716 } else {
12717 e_tok = d;
12718 d_seed = h_nextn;
12719 }
12720 // speculative advance; a chain the grammar can no longer follow (EOS
12721 // proposed) ends here — the prefix already proposed still rides verify.
12722 if dmask_live
12723 && !constraint
12724 .as_deref_mut()
12725 .unwrap()
12726 .draft_advance(d)
12727 .map_err(|e2| format!("constraint: {e2}"))?
12728 {
12729 break;
12730 }
12731 }
12732 }
12733 }
12734 let k_round = draft.len();
12735
12736 ph_mark(&mut ph_draft, phase_on);
12737 // --- 2. VERIFY: one batched target forward. With a pending bonus, it rides as col 0
12738 // (committing its KV/recur inside the SAME weight read); drafts follow. ---
12739 let verify_tokens: Vec<u32> = match pending {
12740 Some(b) => {
12741 let mut v = Vec::with_capacity(k_round + 1);
12742 v.push(b);
12743 v.extend_from_slice(&draft);
12744 v
12745 }
12746 None => draft.clone(),
12747 };
12748 let base = if pending.is_some() { 1 } else { 0 };
12749 // ckpt (REPLAY-FREE partial accept): retain per-layer state-rebuild inputs alongside
12750 // the verify. Pure buffer keep-alives + dtod clones — kernel work is unchanged.
12751 let mut ckpt = if spec_replay {
12752 None
12753 } else {
12754 Some(VerifyCkpt::new(self.layers.len()))
12755 };
12756 let (tlogits_d, vx) = {
12757 // The serial verify every non-fork round takes — the MTP route's
12758 // verify-graph door. The pool is None unless MEMRA_SPEC_VERIFY_GRAPH armed
12759 // a pool above, and then the walk replays the captured trunk instead of
12760 // re-issuing it launch by launch. `graph_round_ok` is the round's
12761 // headroom snapshot (see GRAPH_LAUNCH_MIN_FREE): below the floor the
12762 // round declines the pool exactly like an over-cap round and rides the
12763 // byte-identical eager walk — the `[spec]` suspension line above
12764 // already named the round.
12765 let vg_round = if verify_tokens.len() <= vg_t_cap && graph_round_ok {
12766 vg_guard.as_mut().and_then(|g| g.as_mut())
12767 } else {
12768 if let Some(g) = vg_guard.as_mut().and_then(|g| g.as_mut()) {
12769 // The commit reads this flag to pick its arm; a round that declines
12770 // the pool must not inherit a stale `true` from the round before it.
12771 g.round_slab = false;
12772 }
12773 None
12774 };
12775 self.decode_step_t_core_vg(
12776 e,
12777 &verify_tokens,
12778 pos,
12779 &mut *cache,
12780 embd_dev,
12781 ckpt.as_mut(),
12782 vg_round,
12783 )?
12784 };
12785
12786 if phase_sync {
12787 e.stream().synchronize()?;
12788 }
12789 ph_mark(&mut ph_verify, phase_on);
12790 // --- 3. GREEDY ACCEPT (walk prefix, stop at first mismatch) ---
12791 // DEVICE-ARGMAX ACCEPT: argmax every verify column ON DEVICE (same 2-pass kernels +
12792 // smallest-index tie-break as host argmax, argmax_gate-validated) and read back ONE
12793 // [T] u32 — replaces the T x n_vocab f32 dtoh + T host argmaxes per round.
12794 // t_pred[j] = target's greedy prediction for the slot after draft[j-1] (j>=1) or after
12795 // last_token (j==0). With a pending bonus, col 0 IS the prediction after last_token
12796 // (== the bonus), so every index shifts by `base` and last_pred is unused.
12797 let t_v = verify_tokens.len();
12798 let mut preds: Vec<u32> = Vec::new();
12799 if !sampled {
12800 for j in 0..t_v {
12801 e.argmax_token_device_col(&tlogits_d, j, n_vocab, &mut preds_d, j)?;
12802 }
12803 preds = e.dtoh_u32(&preds_d)?; // <- the verify-GPU wait lands here
12804 // #87 SENTINEL TRAP, verify side: a sentinel pred becomes the round's bonus =
12805 // next round's last_token = the next chain's embed lookup. Catch it at the
12806 // source with the column named — an all-NaN VERIFY column implicates the
12807 // stage-split trunk (decode_step_t_core_ppn), not the draft head.
12808 if let Some(bad) = preds[..t_v].iter().position(|&p| (p as usize) >= n_vocab) {
12809 let col = &tlogits_d.slice(bad * n_vocab..(bad + 1) * n_vocab);
12810 let mut probe = e.zeros(n_vocab)?;
12811 e.copy_view_into(&mut probe, 0, col, n_vocab)?;
12812 let col_h = e.dtoh(&probe)?;
12813 let col_nan = col_h.iter().filter(|v| v.is_nan()).count();
12814 return Err(format!(
12815 "verify argmax sentinel 0x{:08x} >= n_vocab {n_vocab} at round {round} \
12816 col {bad}/{t_v} pos={pos}: verify-logits col NaN {col_nan}/{n_vocab} \
12817 — the verify TRUNK produced a poisoned column (#87 trap). Run \
12818 MEMRA_SPEC_NAN_SCAN=1 to name the layer that creates it (=2 to split \
12819 that layer into attention and routed MoE). NOT the draft head, and NOT \
12820 the PP stage split this message used to name: pp_cuts() returns None \
12821 without MEMRA_PP_STAGES, so decode_step_t_core_ppn never runs unless \
12822 that variable is set.",
12823 preds[bad]
12824 )
12825 .into());
12826 }
12827 }
12828 ph_mark(&mut ph_wait, phase_on);
12829 let t_pred = |j: usize| -> u32 {
12830 if j == 0 && base == 0 {
12831 last_pred
12832 } else {
12833 // GREEDY-ONLY: `preds` is filled under `if !sampled` above. The debug print
12834 // used to call this from the sampled arm and panicked the worker; it now goes
12835 // through `debug_t_pred0`. Keep the strict index here — in the greedy walk an
12836 // out-of-range pred is a real bug, not something to paper over.
12837 debug_assert!(
12838 !sampled,
12839 "t_pred is greedy-only: `preds` is empty in the sampled arm"
12840 );
12841 preds[base + j - 1]
12842 }
12843 };
12844 let (n_acc, bonus) = if !sampled {
12845 let mut n_acc = 0usize;
12846 #[allow(clippy::needless_range_loop)]
12847 // allow: the explicit index loop keeps the offset arithmetic visible and aligned with the device-side indexing
12848 for j in 0..k_round {
12849 if t_pred(j) == draft[j] {
12850 n_acc += 1;
12851 } else {
12852 break;
12853 }
12854 }
12855 // bonus = target's own token at the first non-accepted slot. n_acc in 0..=k; t_pred
12856 // is defined for j in 0..=k (j==0 -> last_logits, j>=1 -> col j-1, last col = k-1).
12857 (n_acc, t_pred(n_acc))
12858 } else {
12859 // --- SAMPLED ACCEPT (rejection sampling): u_j < p_j(x_j)/q_j(x_j) walk ---
12860 if col_buf.is_none() {
12861 col_buf = Some(e.zeros(n_vocab)?);
12862 }
12863 // FILTERED p_j: per-verify-col stats (one batched filter_stats call), then the
12864 // filtered gather. j==0&&base==0 reads last_col (its own stats row appended).
12865 let mut pj = vec![0f32; k_round.max(1)];
12866 let mut col_stats: Vec<(f32, f32, f32)> = Vec::new(); // (max, th, z) per verify col used
12867 if k_round > 0 {
12868 let mut ids: Vec<u32> = Vec::new();
12869 let mut rows: Vec<i32> = Vec::new();
12870 #[allow(clippy::needless_range_loop)]
12871 // allow: the explicit index loop keeps the offset arithmetic visible and aligned with the device-side indexing
12872 for j in 0..k_round {
12873 if j > 0 || base == 1 {
12874 ids.push(draft[j]);
12875 rows.push((base + j) as i32 - 1);
12876 }
12877 }
12878 if !ids.is_empty() {
12879 let nr = rows.len();
12880 // penalties: materialize the used columns into one contiguous penalized
12881 // buffer (rows remapped 0..nr) so stats+gathers see the penalized p.
12882 // penalties: materialize used columns contiguously, penalize all rows in
12883 // one launch, and point stats+gathers at the penalized buffer (rows 0..nr).
12884 let p_rows: Vec<i32> = if pen_on {
12885 (0..nr as i32).collect()
12886 } else {
12887 rows.clone()
12888 };
12889 if pen_on {
12890 if pcol_buf.as_ref().map(|b| b.len()).unwrap_or(0) < nr * n_vocab {
12891 pcol_buf = Some(e.zeros(nr * n_vocab)?);
12892 }
12893 let pc = pcol_buf.as_mut().unwrap();
12894 for (i2, &r) in rows.iter().enumerate() {
12895 let c = r as usize;
12896 e.copy_view_into(
12897 pc,
12898 i2 * n_vocab,
12899 &tlogits_d.slice(c * n_vocab..(c + 1) * n_vocab),
12900 n_vocab,
12901 )?;
12902 }
12903 let h = pen_hist_d.as_ref().unwrap();
12904 let nh = h.len();
12905 e.penalize_logits_rows(
12906 pc,
12907 h,
12908 nh,
12909 sp.penalty_repeat,
12910 sp.penalty_freq,
12911 sp.penalty_present,
12912 n_vocab,
12913 nr,
12914 )?;
12915 }
12916 let p_src: &CudaSlice<f32> = if pen_on {
12917 pcol_buf.as_ref().unwrap()
12918 } else {
12919 &tlogits_d
12920 };
12921 let rowsd = e.htod_i32(&p_rows)?;
12922 let (mut th_d, mut z_d, mut mx_d) =
12923 (e.zeros(nr)?, e.zeros(nr)?, e.zeros(nr)?);
12924 e.filter_stats(
12925 p_src, n_vocab, &rowsd, &mut th_d, &mut z_d, &mut mx_d, n_vocab, nr,
12926 sp_temp, sp.top_k, sp.top_p, sp.min_p,
12927 )?;
12928 let idsd = e.htod_u32_v(&ids)?;
12929 let mut outd = e.zeros(nr)?;
12930 e.softmax_gather_filtered(
12931 p_src, n_vocab, &idsd, &rowsd, &th_d, &z_d, &mut outd, n_vocab, nr,
12932 sp_temp,
12933 )?;
12934 let outv = e.dtoh(&outd)?;
12935 let (thv, zv, mxv) = (e.dtoh(&th_d)?, e.dtoh(&z_d)?, e.dtoh(&mx_d)?);
12936 let mut oi = 0usize;
12937 #[allow(clippy::needless_range_loop)]
12938 // allow: the explicit index loop keeps the offset arithmetic visible and aligned with the device-side indexing
12939 for j in 0..k_round {
12940 if j > 0 || base == 1 {
12941 pj[j] = outv[oi];
12942 oi += 1;
12943 }
12944 }
12945 col_stats = (0..nr).map(|i| (mxv[i], thv[i], zv[i])).collect();
12946 }
12947 if base == 0 {
12948 let lc: &CudaSlice<f32> = if pen_on {
12949 if col_buf.is_none() {
12950 col_buf = Some(e.zeros(n_vocab)?);
12951 }
12952 let cb = col_buf.as_mut().unwrap();
12953 e.copy_into(
12954 cb,
12955 0,
12956 last_col_logits
12957 .as_ref()
12958 .expect("sampled: last_col_logits unset"),
12959 n_vocab,
12960 )?;
12961 let h = pen_hist_d.as_ref().unwrap();
12962 let nh = h.len();
12963 e.penalize_logits(
12964 cb,
12965 h,
12966 nh,
12967 sp.penalty_repeat,
12968 sp.penalty_freq,
12969 sp.penalty_present,
12970 n_vocab,
12971 )?;
12972 col_buf.as_ref().unwrap()
12973 } else {
12974 last_col_logits
12975 .as_ref()
12976 .expect("sampled: last_col_logits unset")
12977 };
12978 let rows0 = e.htod_i32(&[0])?;
12979 let (mut th_d, mut z_d, mut mx_d) = (e.zeros(1)?, e.zeros(1)?, e.zeros(1)?);
12980 e.filter_stats(
12981 lc, n_vocab, &rows0, &mut th_d, &mut z_d, &mut mx_d, n_vocab, 1,
12982 sp_temp, sp.top_k, sp.top_p, sp.min_p,
12983 )?;
12984 let idsd = e.htod_u32_v(&[draft[0]])?;
12985 let mut outd = e.zeros(1)?;
12986 e.softmax_gather_filtered(
12987 lc, n_vocab, &idsd, &rows0, &th_d, &z_d, &mut outd, n_vocab, 1, sp_temp,
12988 )?;
12989 pj[0] = e.dtoh(&outd)?[0];
12990 last_col_stats =
12991 Some((e.dtoh(&mx_d)?[0], e.dtoh(&th_d)?[0], e.dtoh(&z_d)?[0]));
12992 }
12993 }
12994 // q source: the graph arms (single-head AND chain) retained the head logits
12995 // in the persistent q_slots; the eager arm in per-round draft_logits clones.
12996 // Same raw-logit values either way. FILTERED q_j: stats from draft_stats
12997 // (eager pushes in-chain; the graph arms compute them post-replay from the
12998 // retained q with the same filter_stats program — bit-identical to the
12999 // in-graph stats that shaped the draw, keeping ONE accept path).
13000 let q_bufs: &[CudaSlice<f32>] = if dctx.graph_s.is_some() || dctx.chain_s.is_some()
13001 {
13002 &dctx.q_slots
13003 } else {
13004 &draft_logits
13005 };
13006 let mut n_acc = 0usize;
13007 for j in 0..k_round {
13008 let (qmx, qth, qz) = draft_stats[j];
13009 let idsd = e.htod_u32_v(&[draft_idx[j]])?;
13010 let rowsd = e.htod_i32(&[0])?;
13011 let thd = e.htod(&[qth])?;
13012 let zd = e.htod(&[qz])?;
13013 let _ = qmx;
13014 let mut outd = e.zeros(1)?;
13015 e.softmax_gather_filtered(
13016 &q_bufs[j], d_vocab, &idsd, &rowsd, &thd, &zd, &mut outd, d_vocab, 1,
13017 sp_temp,
13018 )?;
13019 let qj = e.dtoh(&outd)?[0];
13020 let u = host_u01(sp_seed, uctr);
13021 uctr += 1;
13022 let accept = (u as f64) * (qj as f64) < pj[j] as f64;
13023 // SKEY PROBE: q == 0 for the token the draft actually proposed is the
13024 // exactness signature (see `skey_probe`). Impossible when the draft was
13025 // drawn from the same filtered distribution the verify reconstructs here;
13026 // `u * 0 < p` makes it an UNCONDITIONAL accept whenever p > 0.
13027 if skey_probe() && qj == 0.0 {
13028 eprintln!(
13029 "[skey] EXACTNESS q=0 round={round} j={j} draft_tok={} \
13030 draft_idx={} p={:e} u={u} accepted={} th_z={:?}",
13031 draft[j], draft_idx[j], pj[j], accept as u8, draft_stats[j],
13032 );
13033 }
13034 if accept {
13035 n_acc += 1;
13036 } else {
13037 break;
13038 }
13039 }
13040 let bonus = if n_acc == k_round {
13041 // FULL ACCEPT: bonus ~ FILTERED softmax at the last verify column.
13042 let col = base + k_round - 1;
13043 let cb = col_buf.as_mut().unwrap();
13044 e.copy_view_into(
13045 cb,
13046 0,
13047 &tlogits_d.slice(col * n_vocab..(col + 1) * n_vocab),
13048 n_vocab,
13049 )?;
13050 if pen_on {
13051 let h = pen_hist_d.as_ref().unwrap();
13052 let nh = h.len();
13053 e.penalize_logits(
13054 cb,
13055 h,
13056 nh,
13057 sp.penalty_repeat,
13058 sp.penalty_freq,
13059 sp.penalty_present,
13060 n_vocab,
13061 )?;
13062 }
13063 if perturb_buf.is_none() {
13064 perturb_buf = Some(e.zeros(d_vocab.max(n_vocab))?);
13065 }
13066 // STATS MUST COME FROM THIS COLUMN (bug fix 2026-08-05, lane/sampler-
13067 // truncation-fix; receipts research/sampfix-20260805/). The old code reused
13068 // `col_stats.last()` here, which is ALWAYS the wrong row: the gathered set
13069 // covers verify columns 0..=(base+k_round-2) (rows pushed as base+j-1), while
13070 // the full-accept bonus samples column base+k_round-1 — exactly ONE PAST the
13071 // last gathered column, in both base arms. `th` is a threshold in e-units of
13072 // its OWN row's max, so feeding a neighbour's (row_max, th) into
13073 // gumbel_perturb_filtered mis-scales every e0 = exp((x-row_max)/T). When the
13074 // donor column's peak is higher by more than T*ln(1/th), EVERY id fails
13075 // `e0 >= th`, the whole perturbed row becomes -3.4e38, and the 2-pass argmax
13076 // falls through to its smallest-index tie-break => token id 0 ("!") spliced
13077 // mid-word. Fragility is ordered by how large th is: min_p pins th = min_p
13078 // (0.05 => trigger at delta > 2.4 at T=0.8, fires constantly), top_p's
13079 // mass-boundary th is smaller, top_k's k-th-largest th smaller still — which
13080 // is why the head-to-head matrix saw min_p and top_p corrupt while top_k-only
13081 // stayed clean. The pure-temp default regime is immune (th == 0 masks nothing,
13082 // and row_max is unused once nothing is masked), so this fix is a byte-level
13083 // no-op for the untruncated serve default. One extra one-block filter_stats
13084 // per full-accept round is the whole cost.
13085 let (mx, th) = {
13086 let rows0 = e.htod_i32(&[0])?;
13087 let (mut th_d, mut z_d, mut mx_d) = (e.zeros(1)?, e.zeros(1)?, e.zeros(1)?);
13088 let cb0 = col_buf.as_ref().unwrap();
13089 e.filter_stats(
13090 cb0, n_vocab, &rows0, &mut th_d, &mut z_d, &mut mx_d, n_vocab, 1,
13091 sp_temp, sp.top_k, sp.top_p, sp.min_p,
13092 )?;
13093 (e.dtoh(&mx_d)?[0], e.dtoh(&th_d)?[0])
13094 };
13095 let pb = perturb_buf.as_mut().unwrap();
13096 let cb2 = col_buf.as_ref().unwrap();
13097 e.gumbel_perturb_filtered(cb2, pb, n_vocab, sp_seed, sctr, sp_temp, mx, th)?;
13098 sctr += 1;
13099 let td = e.argmax_token_device(pb, n_vocab)?;
13100 e.dtoh_u32_one(&td)?
13101 } else {
13102 // REJECT at n_acc: bonus ~ norm(max(0, softmax_T(p) - softmax_T(q))).
13103 let cb = col_buf.as_mut().unwrap();
13104 if n_acc > 0 || base == 1 {
13105 let col = base + n_acc - 1;
13106 e.copy_view_into(
13107 cb,
13108 0,
13109 &tlogits_d.slice(col * n_vocab..(col + 1) * n_vocab),
13110 n_vocab,
13111 )?;
13112 } else {
13113 let lc = last_col_logits.as_ref().unwrap();
13114 e.copy_into(cb, 0, lc, n_vocab)?;
13115 }
13116 if pen_on {
13117 let h = pen_hist_d.as_ref().unwrap();
13118 let nh = h.len();
13119 e.penalize_logits(
13120 cb,
13121 h,
13122 nh,
13123 sp.penalty_repeat,
13124 sp.penalty_freq,
13125 sp.penalty_present,
13126 n_vocab,
13127 )?;
13128 }
13129 let cb2 = col_buf.as_ref().unwrap();
13130 let sc = sctr;
13131 sctr += 1;
13132 // p-stats for the reject column: from col_stats when the col was gathered,
13133 // else (j==0&&base==0) from last_col_stats.
13134 let p_stats = if n_acc > 0 || base == 1 {
13135 // col index within the gathered set == number of gathered cols before n_acc
13136 let gi = if base == 1 { n_acc } else { n_acc - 1 };
13137 col_stats.get(gi).copied().unwrap_or({
13138 (0.0, 0.0, 1.0) // unreachable: gathered cols always cover the reject slot
13139 })
13140 } else {
13141 last_col_stats.expect("sampled: last_col_stats unset at reject")
13142 };
13143 let q_stats = draft_stats[n_acc];
13144 if let Some(map) = &d2t_dev {
13145 if q_full_buf.is_none() {
13146 q_full_buf = Some(e.zeros(n_vocab)?);
13147 }
13148 let qf = q_full_buf.as_mut().unwrap();
13149 e.scatter_trim_logits(&q_bufs[n_acc], map, qf, d_vocab, n_vocab)?;
13150 let qf2 = q_full_buf.as_ref().unwrap();
13151 e.residual_sample_filtered(
13152 cb2,
13153 Some(qf2),
13154 n_vocab,
13155 sp_temp,
13156 sp_seed,
13157 sc,
13158 p_stats,
13159 q_stats,
13160 &mut sample_tok,
13161 )?;
13162 } else {
13163 e.residual_sample_filtered(
13164 cb2,
13165 Some(&q_bufs[n_acc]),
13166 n_vocab,
13167 sp_temp,
13168 sp_seed,
13169 sc,
13170 p_stats,
13171 q_stats,
13172 &mut sample_tok,
13173 )?;
13174 }
13175 e.dtoh_u32(&sample_tok)?[0]
13176 };
13177 (
13178 n_acc,
13179 guard_vocab_token(
13180 bonus,
13181 n_vocab,
13182 &format!("sampled verify bonus at round {round} pos={pos} n_acc={n_acc}"),
13183 )?,
13184 )
13185 };
13186 // --- 3b. GRAMMAR TRUNCATION (constrained spec, 2026-08-03): the grammar is
13187 // an extra rejection rule AFTER the exactness verify (the batched-verify-twins
13188 // ordering). Walk the accepted drafts through the grammar in commit order; the
13189 // first illegal token truncates acceptance at its slot, and that slot's emission
13190 // is recomputed as the MASKED argmax of the target's own verify column — token-
13191 // identical to constrained plain greedy decode (an unmasked argmax that is
13192 // grammar-legal IS the masked argmax: masking only removes competitors). The
13193 // column D2H (~1MB) is paid only when a cut fires — the tight-grammar cost,
13194 // measured in acceptance numbers, never hidden.
13195 let (n_acc, bonus) = match constraint.as_deref_mut() {
13196 None => (n_acc, bonus),
13197 Some(c) => {
13198 fn ce(e2: String) -> Box<dyn std::error::Error> {
13199 format!("constraint: {e2}").into()
13200 }
13201 let mut na = n_acc;
13202 let mut cut = false;
13203 for (j, &d) in draft.iter().enumerate().take(n_acc) {
13204 if c.is_allowed(d).map_err(ce)? {
13205 c.consume(d).map_err(ce)?;
13206 } else {
13207 na = j;
13208 cut = true;
13209 dm_cut_tokens += n_acc - j;
13210 break;
13211 }
13212 }
13213 if cut {
13214 dm_cuts += 1;
13215 }
13216 let mut bo = bonus;
13217 if cut || !c.is_allowed(bo).map_err(ce)? {
13218 let mut row = if na == 0 && base == 0 {
13219 init_logits_host
13220 .clone()
13221 .ok_or("constraint: init logits missing (round-0 cut)")?
13222 } else {
13223 e.dtoh_view(
13224 &tlogits_d.slice((base + na - 1) * n_vocab..(base + na) * n_vocab),
13225 )?
13226 };
13227 c.mask_logits(&mut row).map_err(ce)?;
13228 bo = argmax(&row) as u32;
13229 }
13230 c.consume(bo).map_err(ce)?;
13231 (na, bo)
13232 }
13233 };
13234 total_drafted += k_round;
13235 total_accepted += n_acc;
13236 if let Some(t) = sess_telem {
13237 // Greedy, rejection-sampling, and grammar truncation all converge here after
13238 // the accept decision is already on host. Fixed-size relaxed atomics only.
13239 t.record_round(k_round, n_acc);
13240 }
13241 if spec_stats {
13242 st_len_hist[k_round] += 1;
13243 #[allow(clippy::needless_range_loop)]
13244 // allow: the explicit index loop keeps the offset arithmetic visible and aligned with the device-side indexing
13245 for j in 0..k_round {
13246 st_drafted[j] += 1;
13247 }
13248 #[allow(clippy::needless_range_loop)]
13249 // allow: the explicit index loop keeps the offset arithmetic visible and aligned with the device-side indexing
13250 for j in 0..n_acc {
13251 st_accepted[j] += 1;
13252 }
13253 if n_acc == k_round {
13254 st_full += 1;
13255 }
13256 }
13257
13258 if debug_spec {
13259 eprintln!(
13260 "[R{round}] pos={pos} out_len={} last_tok={last_token} draft={draft:?} n_acc={n_acc} bonus={bonus} t_pred0={}",
13261 out.len(),
13262 // NOT `t_pred(0)`: `preds` is filled only under `if !sampled` above, so on a
13263 // sampled request round >= 1 (base == 1) indexed an EMPTY vector and PANICKED
13264 // the GPU worker thread — a debug flag that killed the exact regime you would
13265 // set it to investigate. See `debug_t_pred0`.
13266 debug_t_pred0(sampled, base, last_pred, &preds)
13267 );
13268 }
13269
13270 // --- 4. COMMIT: draft[0..n_acc] then bonus (n_acc + 1 tokens) ---
13271 let commit_started = std::time::Instant::now();
13272 // SESSION MODE: every accepted column is already in the CACHE — `out` must carry all
13273 // of them (overshoot past max_new included) or `committed` under-counts the cache rows
13274 // and the next turn's continuation seeds one token off (gate-caught 2026-07-05). The
13275 // single-shot path keeps the cap (its caller truncates + drops the cache anyway).
13276 #[allow(clippy::needless_range_loop)]
13277 // allow: the explicit index loop keeps the offset arithmetic visible and aligned with the device-side indexing
13278 for j in 0..n_acc {
13279 if !session_mode && out.len() >= max_new {
13280 break;
13281 }
13282 out.push(draft[j]);
13283 }
13284 if pen_on {
13285 pen_hist.extend_from_slice(&draft[0..n_acc]);
13286 pen_hist.push(bonus);
13287 }
13288 let bonus_emitted = session_mode || out.len() < max_new;
13289 if bonus_emitted {
13290 out.push(bonus);
13291 }
13292 last_token = bonus;
13293
13294 // --- 5. ROLLBACK + advance (§C) ---
13295 if n_acc == k_round && !spec_replay {
13296 // FULL ACCEPT, BONUS FOLD: all verify columns (pending? + drafts) are committed in
13297 // cache; the NEW bonus stays PENDING for the next round's verify batch — NO extra
13298 // T=1 trunk pass. The next draft chain seeds from the MTP block's h_nextn at the
13299 // bonus position: one MTP-block pass (~1/33 trunk cost) replaces the trunk read.
13300 // last_pred is dead in the pending path (t_pred reads verify col 0).
13301 //
13302 // PERSISTENT DRAFT KV, full-accept fill: the chain covered last_token +
13303 // draft[0..k_round-2] as INPUTS (slots P..P'-2); draft[k_round-1] (slot P'-1) was
13304 // only ever an output, so its entry is MISSING. Fill it from vh_seed — its EXACT
13305 // trunk hidden (the last verify column). set_len first: a p-min break may have
13306 // left one extra chain append at that slot. Partial accepts need NO fill (the
13307 // chain already covered every accepted position; round-start set_len truncates).
13308 self.restore_step_tp_kv_verified_prefix(e, &mut *cache, &snap, t_v, false)?;
13309 let mut vh_seed = e.zeros(n_embd)?;
13310 e.copy_view_into(
13311 &mut vh_seed,
13312 0,
13313 &vx.slice((t_v - 1) * n_embd..t_v * n_embd),
13314 n_embd,
13315 )?;
13316 if refresh {
13317 // TRUE-HIDDEN REFRESH (2026-07-03, the HANDOVER-listed acceptance lever):
13318 // overwrite ALL committed positions' scratch entries with K/V from their EXACT
13319 // verify hiddens — the reference engine's mtp_update fills from true hiddens;
13320 // the full stack (vx) is already resident from the verify. Replaces both the
13321 // chain-approximate entries AND the old last-token-only fill. Acceptance-only
13322 // (draft attention quality); exactness stays the verify's job.
13323 scratch.set_len(e, pos)?;
13324 // PREDECESSOR pairing: row i gets vx[i-1]; row 0 the carried fill_prev
13325 // (hidden of the last committed row before this verify batch).
13326 let mut vxs = e.zeros(t_v * n_embd)?;
13327 e.copy_into(&mut vxs, 0, &fill_prev, n_embd)?;
13328 if t_v > 1 {
13329 e.copy_view_into(
13330 &mut vxs,
13331 n_embd,
13332 &vx.slice(0..(t_v - 1) * n_embd),
13333 (t_v - 1) * n_embd,
13334 )?;
13335 }
13336 self.mtp_kv_fill_all(e, &verify_tokens, &vxs, pos, &mut *scratch, embd_dev)?;
13337 } else {
13338 scratch.set_len(e, pos + base + k_round - 1)?;
13339 // predecessor of the last draft = verify col t_v-2 (or fill_prev at t_v==1)
13340 let mut hp = e.zeros(n_embd)?;
13341 if t_v >= 2 {
13342 e.copy_view_into(
13343 &mut hp,
13344 0,
13345 &vx.slice((t_v - 2) * n_embd..(t_v - 1) * n_embd),
13346 n_embd,
13347 )?;
13348 } else {
13349 e.copy_into(&mut hp, 0, &fill_prev, n_embd)?;
13350 }
13351 self.mtp_kv_fill_all(
13352 e,
13353 &[draft[k_round - 1]],
13354 &hp,
13355 pos + base + k_round - 1,
13356 &mut *scratch,
13357 embd_dev,
13358 )?;
13359 }
13360 // REFERENCE SEEDING: no pseudo pass — the next chain's step 0 IS the
13361 // reference's (id_last, h_prev) draft row; it appends the bonus's scratch
13362 // entry itself. Seed = TRUE hidden of the bonus's predecessor (last verify
13363 // col). Saves one MTP-block pass per round on top of the pairing fix.
13364 e.copy_into(&mut h_seed_buf, 0, &vh_seed, n_embd)?;
13365 e.copy_into(&mut fill_prev, 0, &vh_seed, n_embd)?;
13366 pending = Some(bonus);
13367 if debug_spec {
13368 eprintln!(" -> FULL ACCEPT (bonus pending, prev-h seed)");
13369 }
13370 } else if !spec_replay && base + n_acc >= 1 {
13371 // PARTIAL ACCEPT, REPLAY-FREE (2026-07-03 — the profiled #1 long-ctx spec cost):
13372 // the verify's first j = base+n_acc columns ARE the committed sequence, computed
13373 // bit-identically to eager (decode-exact contract) — so KEEP them: KV truncates to
13374 // pos+j, recurrent state rebuilds from the VerifyCkpt (same-kernel gdn prefix
13375 // re-run / pure state-clone restore), and the bonus stays PENDING exactly like the
13376 // full-accept path — the legacy duplicate trunk replay is gone. The next chain
13377 // seeds from the MTP pseudo-hidden of the bonus, whose seed = the TRUE verify
13378 // hidden of its predecessor (col j-1) — same one-hop pseudo structure as full
13379 // accept (never compounds: the next verify recomputes true hiddens for all
13380 // committed columns).
13381 let j = base + n_acc;
13382 // VERIFY-GRAPH SLAB COMMIT: when the captured trunk ran, the linear layers'
13383 // column stash was written into the graphs ctx's persistent slabs as in-graph
13384 // memcpy nodes, NOT into the per-column VerifyCkpt the cols arm reads — so the
13385 // commit must take the slab twin (same semantics, slab-addressed sources). The
13386 // ctx states which of the two this round produced via `round_slab`; trusting the
13387 // flag rather than the env keeps a round that fell back to the eager walk (a
13388 // capture that declined, a t the pool never captured) on the cols arm.
13389 let slab_commit = vg_guard
13390 .as_ref()
13391 .and_then(|g| g.as_ref())
13392 .map(|g| g.round_slab)
13393 .unwrap_or(false);
13394 if slab_commit {
13395 self.dspark_commit_prefix_slab(
13396 e,
13397 &mut *cache,
13398 &snap,
13399 vg_guard
13400 .as_ref()
13401 .and_then(|g| g.as_ref())
13402 .expect("slab_commit implies a graphs ctx"),
13403 j,
13404 )?;
13405 } else {
13406 self.commit_verified_prefix(e, &mut *cache, &snap, ckpt.as_ref().unwrap(), j)?;
13407 }
13408 let mut seed = e.zeros(n_embd)?;
13409 e.copy_view_into(
13410 &mut seed,
13411 0,
13412 &vx.slice((j - 1) * n_embd..j * n_embd),
13413 n_embd,
13414 )?;
13415 // Draft scratch: TRUE-HIDDEN REFRESH of the committed prefix (see the full-accept
13416 // branch); without it the chain entries stand and only the tail truncates. Either
13417 // way len ends at pos+j so the pseudo append lands at the bonus's slot pos+j
13418 // (persistent mode), rope pos+j+1 (chain convention).
13419 if refresh {
13420 scratch.set_len(e, pos)?;
13421 let mut vxs = e.zeros(j * n_embd)?;
13422 e.copy_into(&mut vxs, 0, &fill_prev, n_embd)?;
13423 if j > 1 {
13424 e.copy_view_into(
13425 &mut vxs,
13426 n_embd,
13427 &vx.slice(0..(j - 1) * n_embd),
13428 (j - 1) * n_embd,
13429 )?;
13430 }
13431 self.mtp_kv_fill_all(
13432 e,
13433 &verify_tokens[0..j],
13434 &vxs,
13435 pos,
13436 &mut *scratch,
13437 embd_dev,
13438 )?;
13439 } else {
13440 scratch.set_len(e, pos + j)?;
13441 }
13442 // REFERENCE SEEDING (see the full-accept branch): seed = TRUE hidden of the
13443 // bonus's predecessor (verify col j-1); no pseudo pass.
13444 e.copy_into(&mut h_seed_buf, 0, &seed, n_embd)?;
13445 e.copy_into(&mut fill_prev, 0, &seed, n_embd)?;
13446 pending = Some(bonus);
13447 if debug_spec {
13448 eprintln!(" -> PARTIAL(replay-free j={j}, bonus pending, prev-h seed)");
13449 }
13450 } else if !spec_replay {
13451 // ZERO ROUND FOLD (2026-07-10, verify-cost target #3): base+n_acc == 0 — a
13452 // pending-less round where nothing was accepted (PMIN0 zero-draft chains after a
13453 // replay/commit, or plain 0-accept rounds at round 0). The old path replayed
13454 // [bonus] through a FULL m=1 trunk+head forward (the 489us full-vocab head pass
13455 // measured at ~0.75/round on PMIN0 configs). Instead: restore the pre-round
13456 // snapshot and let the bonus ride the NEXT round's verify as col 0 — the existing
13457 // base=1 pending machinery, bit-identical by the decode-exact verify contract.
13458 // Seed: the bonus's predecessor is the last COMMITTED token, whose hidden
13459 // fill_prev already carries (same seeding as the 1-token-replay case it replaces).
13460 cache.rollback(e, &snap, 0)?;
13461 scratch.set_len(e, pos)?;
13462 e.copy_into(&mut h_seed_buf, 0, &fill_prev, n_embd)?;
13463 pending = Some(bonus);
13464 if debug_spec {
13465 eprintln!(" -> ZERO-ROUND FOLD (bonus pending, fill_prev seed)");
13466 }
13467 } else {
13468 // PARTIAL ACCEPT, LEGACY REPLAY (seam MEMRA_SPEC_REPLAY=1 — or j==0: nothing of
13469 // this round survives, only possible before the first pending exists, ~round 0):
13470 // restore EVERYTHING to the pre-round snapshot (KV truncate to pos + recur
13471 // restore), then replay the committed prefix pending? ++ draft[0..n_acc] ++
13472 // [bonus] as ONE batched T forward — single weight read, bit-identical to greedy
13473 // (the verify-all-columns path is the same math). Commits the bonus with a TRUE
13474 // trunk hidden.
13475 cache.rollback(e, &snap, 0)?; // accept_len=0: KV len = pos, recur = snapshot
13476 let mut replay: Vec<u32> = Vec::with_capacity(base + n_acc + 1);
13477 if let Some(b) = pending.take() {
13478 replay.push(b);
13479 }
13480 replay.extend_from_slice(&draft[0..n_acc]);
13481 replay.push(bonus);
13482 // Full-stack forward (decode_step_t_core = decode_step_t_h_emb_dev's body):
13483 // Predecessor pairing seeds from the PREDECESSOR row (col len-2) — the same-row path takes the
13484 // last col exactly as before (byte-identical to the old _h_emb_dev call).
13485 let (rl_d, rx) = if self.batched_serving_numeric_class() {
13486 let mut logits = Vec::with_capacity(replay.len() * n_vocab);
13487 let mut hidden = e.uninit(replay.len() * n_embd)?;
13488 for (row, &token) in replay.iter().enumerate() {
13489 let (row_logits, row_hidden) =
13490 self.spec_target_step_h(e, token, &mut *cache)?;
13491 logits.extend_from_slice(&row_logits);
13492 e.dtod_copy_into(&row_hidden, &mut hidden, row * n_embd)?;
13493 }
13494 (e.htod(&logits)?, hidden)
13495 } else {
13496 self.decode_step_t_core(e, &replay, pos, &mut *cache, embd_dev, None)?
13497 };
13498 // last_pred = argmax of the LAST column's logits (predicts the token after `bonus`)
13499 // — device argmax + one 4-byte read instead of the full-vocab column dtoh.
13500 e.argmax_token_device_col(&rl_d, replay.len() - 1, n_vocab, &mut preds_d, 0)?;
13501 last_pred = guard_vocab_token(
13502 e.dtoh_u32(&preds_d)?[0],
13503 n_vocab,
13504 &format!("replay last_pred at round {round} pos={pos}"),
13505 )?;
13506 if sampled {
13507 let lr0 = replay.len();
13508 let lc = last_col_logits
13509 .as_mut()
13510 .expect("sampled: last_col_logits unset");
13511 e.copy_view_into(
13512 lc,
13513 0,
13514 &rl_d.slice((lr0 - 1) * n_vocab..lr0 * n_vocab),
13515 n_vocab,
13516 )?;
13517 }
13518 let lr = replay.len();
13519 if lr >= 2 {
13520 e.copy_view_into(
13521 &mut h_seed_buf,
13522 0,
13523 &rx.slice((lr - 2) * n_embd..(lr - 1) * n_embd),
13524 n_embd,
13525 )?;
13526 } else {
13527 // 1-token replay (round-0 miss): the bonus's predecessor is the OLD
13528 // last_token, whose own-row hidden fill_prev still holds.
13529 e.copy_into(&mut h_seed_buf, 0, &fill_prev, n_embd)?;
13530 }
13531 // the bonus is COMMITTED here — it becomes the last committed row.
13532 let mut rh_last = e.zeros(n_embd)?;
13533 e.copy_view_into(
13534 &mut rh_last,
13535 0,
13536 &rx.slice((lr - 1) * n_embd..lr * n_embd),
13537 n_embd,
13538 )?;
13539 e.copy_into(&mut fill_prev, 0, &rh_last, n_embd)?;
13540 if debug_spec {
13541 eprintln!(" -> PARTIAL(replay={replay:?}), next_pred={last_pred}");
13542 }
13543 }
13544 if anatomy_on {
13545 // Commit/rollback is normally asynchronous on the primary/head stream. Bound it
13546 // only for this diagnostic so it does not disappear into the following draft's
13547 // first token readback.
13548 e.stream().synchronize()?;
13549 ph_commit += commit_started.elapsed().as_secs_f64();
13550 }
13551 // adaptive-K update (host math, zero syncs): next round drafts accepted-run + 1,
13552 // clamped to [floor(pos), k_cap]. cache.pos is post-rollback here (the round's
13553 // final position — the floor's position key reads the committed depth). Burst
13554 // rounds (`continue` above) draft the captured fixed depth and skip this, exactly
13555 // like gemma's burst arm.
13556 if adapt {
13557 let fl_now = floor_at(cache.pos);
13558 kc = (n_acc + 1).clamp(fl_now.min(k_cap), k_cap);
13559 }
13560 ph_mark(&mut ph_rest, phase_on);
13561 if let Some(t0) = round_t0 {
13562 let ms = t0.elapsed().as_secs_f64() * 1e3;
13563 ROUND_MS.fetch_add((ms * 1e3) as u64, std::sync::atomic::Ordering::Relaxed);
13564 let n = ROUND_N.fetch_add(1, std::sync::atomic::Ordering::Relaxed) + 1;
13565 if n.is_multiple_of(32) {
13566 eprintln!(
13567 "[spec-round] rounds={n} avg round wall={:.2} ms (emitted={} drafted so far)",
13568 ROUND_MS.load(std::sync::atomic::Ordering::Relaxed) as f64 / 1e3 / n as f64,
13569 out.len()
13570 );
13571 }
13572 }
13573 round += 1;
13574 // sse-cadence: this round's accepted drafts + bonus are committed (out is
13575 // append-only past step 4) — flush at round cadence.
13576 keep_going = flush_commit(&mut on_commit, &out, &mut flushed);
13577 }
13578 // sse-cadence: nothing below appends to `out`; flush any remainder (defensive).
13579 // (verdict ignored — the burst is over either way; the session tail runs unchanged.)
13580 let _ = flush_commit(&mut on_commit, &out, &mut flushed);
13581
13582 if spec_stats {
13583 let per_slot: Vec<String> = (0..k)
13584 .map(|j| {
13585 if st_drafted[j] > 0 {
13586 format!(
13587 "{}/{}={:.3}",
13588 st_accepted[j],
13589 st_drafted[j],
13590 st_accepted[j] as f64 / st_drafted[j] as f64
13591 )
13592 } else {
13593 "0/0".into()
13594 }
13595 })
13596 .collect();
13597 let acc = if total_drafted > 0 {
13598 total_accepted as f64 / total_drafted as f64
13599 } else {
13600 0.0
13601 };
13602 eprintln!(
13603 "[spec-stats] rounds={round} full_accept={st_full} len_hist={st_len_hist:?} \
13604 per_slot=[{}] total={total_accepted}/{total_drafted}={acc:.3} \
13605 tok_per_round={:.3}",
13606 per_slot.join(" "),
13607 (total_accepted + round) as f64 / round.max(1) as f64
13608 );
13609 }
13610 if constraint.is_some() {
13611 eprintln!(
13612 "[draft-mask] mask_rounds={dm_rounds} clone_total={:.3}ms \
13613 clone_per_round={:.4}ms gram_cuts={dm_cuts}/{round} cut_tokens={dm_cut_tokens}",
13614 dm_clone_ns as f64 / 1e6,
13615 dm_clone_ns as f64 / 1e6 / dm_rounds.max(1) as f64
13616 );
13617 }
13618 if phase_on {
13619 let tot = ph_draft + ph_verify + ph_wait + ph_rest;
13620 eprintln!(
13621 "[spec-phase] draft={:.1}ms ({:.1}%) verify-issue={:.1}ms ({:.1}%) verify-wait={:.1}ms ({:.1}%) commit-host={:.1}ms ({:.1}%) rounds={round}",
13622 ph_draft * 1e3,
13623 ph_draft / tot * 100.0,
13624 ph_verify * 1e3,
13625 ph_verify / tot * 100.0,
13626 ph_wait * 1e3,
13627 ph_wait / tot * 100.0,
13628 ph_rest * 1e3,
13629 ph_rest / tot * 100.0
13630 );
13631 }
13632 if anatomy_on {
13633 let rounds_f = round.max(1) as f64;
13634 let other = (ph_rest - ph_commit).max(0.0);
13635 eprintln!(
13636 "[spec-anatomy] per-round draft={:.3}ms pp-verify={:.3}ms \
13637 verify-accept={:.3}ms commit-rollback={:.3}ms other={:.3}ms rounds={round}",
13638 ph_draft * 1e3 / rounds_f,
13639 ph_verify * 1e3 / rounds_f,
13640 ph_wait * 1e3 / rounds_f,
13641 ph_commit * 1e3 / rounds_f,
13642 other * 1e3 / rounds_f,
13643 );
13644 }
13645 // SESSION TAIL: leave the session in the exact invariant the next turn's suffix prime
13646 // expects — every row in `committed` has trunk KV/recur state AND an exact draft-KV row.
13647 // Park the draft-graph ctx back on the session (the serve-burst fixed-cost fix): the next
13648 // burst replays instead of recapturing. Error paths (`?` above) drop it — recaptured then.
13649 if let Some(slot) = sess_draft_slot.take() {
13650 *slot = Some(dctx);
13651 }
13652 let t_rounds = t_ent.elapsed();
13653 if let Some((committed, last_h, next_pred_slot, sctr_slot, uctr_slot)) = sess_tail.take() {
13654 // NEXT BURST'S BOUNDARY TOKEN (lane/sampled-spec-quality, Item 1). Greedy stashes
13655 // the argmax `last_pred` exactly as before (byte contract). SAMPLED draws the token
13656 // HERE, where the sampler, the session Philox counters and the penalty window are
13657 // all live and the boundary logits row still exists — that is the "make the state
13658 // available" half of the fix; the consuming burst then just emits it. `sctr` is
13659 // written to the session BELOW the draws so the advance is never lost.
13660 *next_pred_slot = Some(last_pred);
13661 let sample_boundary = sampled && constraint.is_none() && spec_sampled_boundary_on();
13662 let mut stashed_pending = false;
13663 if let Some(b) = pending.take() {
13664 if !sampled {
13665 // PENDING-CARRY (2026-08-01): stash the bonus on the session instead of
13666 // committing it with a solo T=1 pass — the next empty-suffix greedy burst
13667 // consumes it as round-0 verify col 0 (a plain round edge; the old tail
13668 // commit + next burst's init feed were 11.6+11.5ms solo trunk passes per
13669 // burst on H100 q27, [spec-setup] trace). b stays in `out` (emitted) but
13670 // OUT of `committed` (cache rows == committed); the consuming call
13671 // prepends it once its verify commits the row. next_pred is unknowable
13672 // without the commit pass — None; callers gate on pending_tok too.
13673 debug_assert_eq!(out.last(), Some(&b), "pending must be the last emitted");
13674 if let Some(slot) = sess_pending_slot.take() {
13675 *slot = Some(b);
13676 }
13677 *next_pred_slot = None;
13678 // fill_prev = hidden of the last COMMITTED row (b's predecessor) — the
13679 // exact chain-seed/fill anchor the consuming burst (or a flush) needs.
13680 *last_h = Some(e.clone_dtod(&fill_prev)?);
13681 stashed_pending = true;
13682 } else {
13683 // SAMPLED tail (unchanged): commit the bonus (one T=1 pass) + draft fill —
13684 // the sampled round-0 accept needs this pass's logits (last_col_logits).
13685 let pos_b = cache.pos;
13686 scratch.set_len(e, pos_b)?;
13687 let (lg_b, hb) = self.spec_target_step_h(e, b, &mut *cache)?;
13688 // after a FULL-accept exit `last_pred` is STALE (it predicted the bonus
13689 // itself — the prediction AFTER the bonus never materialized; it would have
13690 // been the next round's verify col 0). The commit's logits ARE that
13691 // prediction — so they are also the row the next burst's boundary token
13692 // comes off, and (lane/sampled-spec-quality) it is DRAWN from them here.
13693 *next_pred_slot = Some(if sample_boundary {
13694 sample_boundary_token(
13695 e,
13696 &lg_b,
13697 &sp,
13698 &pen_hist,
13699 &mut sctr,
13700 "burst-tail-commit",
13701 )?
13702 } else {
13703 argmax(&lg_b) as u32
13704 });
13705 self.mtp_kv_fill_all(e, &[b], &fill_prev, pos_b, &mut *scratch, embd_dev)?;
13706 *last_h = Some(hb);
13707 }
13708 } else {
13709 // fill_prev tracks the hidden of the last COMMITTED row throughout the loop.
13710 *last_h = Some(e.clone_dtod(&fill_prev)?);
13711 if sample_boundary {
13712 // No pending to commit, so the boundary row is the one `last_pred` was
13713 // argmaxed from and the sampled path keeps it on device: the init feed's
13714 // logits when the burst ran zero rounds, else the legacy-replay path's
13715 // last verify column (both predict the token AFTER the last committed
13716 // row). It is retained precisely because round 0's accept test needs it,
13717 // so the draw costs no extra D2H of the [n_vocab] row.
13718 match last_col_logits.as_ref() {
13719 Some(lc) => {
13720 *next_pred_slot = Some(sample_boundary_token_dev(
13721 e,
13722 lc,
13723 n_vocab,
13724 &sp,
13725 &pen_hist,
13726 &mut sctr,
13727 "burst-tail-nopending",
13728 )?);
13729 }
13730 // NAME THE FALLBACK (house standard): unreachable today — a sampled
13731 // burst always feeds or replays, so the row exists — but if it ever
13732 // is, the stream takes a greedy token and SAYS so rather than
13733 // silently regressing to the pre-lane behaviour.
13734 None => eprintln!(
13735 "[spec-boundary] sampled tail kept the ARGMAX boundary token \
13736 (reason: no retained boundary logits row)"
13737 ),
13738 }
13739 }
13740 }
13741 *sctr_slot = sctr;
13742 *uctr_slot = uctr;
13743 committed.extend_from_slice(prompt);
13744 if let Some(cb) = carried_pending {
13745 // the consumed carry's cache row landed in round 0's verify (every pending
13746 // round commits col 0) — it joins `committed` here, in sequence order.
13747 committed.push(cb);
13748 }
13749 if stashed_pending {
13750 // ZERO-EMIT BURST (2026-08-06 c=8 serve panic, pre-existing since b4aea184):
13751 // `out.len() - 1` underflowed on an EMPTY `out` — "range end index
13752 // 18446744073709551615 out of range for slice of length 0", killing the
13753 // memra-gpu-worker and failing 31 of 32 concurrent requests with "worker closed
13754 // stream". Reachable because `pending` starts as `carried_pending` (a bonus
13755 // stashed by the PREVIOUS burst) while `out` starts empty, and the carry is
13756 // deliberately NOT pushed to `out` (line ~3239: the burst that emitted it already
13757 // did). So a burst that stashes a pending without emitting anything of its own —
13758 // the round loop exits before a push, e.g. the ring drain's `out.len() < max_new`
13759 // guard skipping every token under a tight budget — arrives here with
13760 // out.len() == 0 and stashed_pending == true.
13761 //
13762 // The invariant is unchanged: `committed` gets every emitted token EXCEPT the
13763 // stashed bonus. With nothing emitted, that is nothing — and the carry pushed
13764 // just above is already accounted. Saturating, not a min/assert: an empty `out`
13765 // here is a legitimate burst shape, not a corrupt state.
13766 let emitted = out.len().saturating_sub(1);
13767 committed.extend_from_slice(&out[..emitted]);
13768 } else {
13769 committed.extend_from_slice(&out); // FULL out incl. overshoot — all committed
13770 }
13771 debug_assert_eq!(
13772 cache.pos,
13773 committed.len(),
13774 "session invariant: cache rows == committed tokens"
13775 );
13776 if setup_trace {
13777 e.stream().synchronize()?; // bound the async tail fill in the trace
13778 let t_tail = t_ent.elapsed();
13779 eprintln!(
13780 "[spec-setup] init={:.2}ms cap={:.2}ms fill={:.2}ms rounds={:.2}ms tail={:.2}ms total={:.2}ms out={} cont={}",
13781 t_init.as_secs_f64() * 1e3,
13782 (t_cap - t_init).as_secs_f64() * 1e3,
13783 (t_fill - t_cap).as_secs_f64() * 1e3,
13784 (t_rounds - t_fill).as_secs_f64() * 1e3,
13785 (t_tail - t_rounds).as_secs_f64() * 1e3,
13786 t_tail.as_secs_f64() * 1e3,
13787 out.len(),
13788 continuation
13789 );
13790 }
13791 return Ok((out, total_drafted, total_accepted));
13792 }
13793 out.truncate(max_new);
13794 Ok((out, total_drafted, total_accepted))
13795 }
13796
13797 /// Anchor-bounded DSpark target extraction. The trunk sees the exact generated token tape;
13798 /// only requested hidden rows and target-logit rows cross PCIe. An anchor token at p pairs
13799 /// with the pre-output-norm h[p-1] carrier, exactly as the existing replay/NextN path does.
13800 #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
13801 pub fn extract_dspark_anchors(
13802 &self,
13803 e: &Engine,
13804 tokens: &[u32],
13805 anchor_positions: &[usize],
13806 gamma: usize,
13807 top_k: usize,
13808 chunk: usize,
13809 temperature: f32,
13810 ) -> Result<Vec<DsparkAnchorRecord>, Box<dyn std::error::Error>> {
13811 if tokens.len() < gamma + 2 || gamma == 0 || chunk < 2 {
13812 return Err("DSpark extraction token tape/gamma/chunk is invalid".into());
13813 }
13814 if anchor_positions.windows(2).any(|pair| pair[0] >= pair[1]) {
13815 return Err("DSpark anchor positions must be sorted and unique".into());
13816 }
13817 for &position in anchor_positions {
13818 if position == 0 || position + gamma >= tokens.len() {
13819 return Err(format!(
13820 "DSpark anchor {position} has no predecessor or cannot cover gamma={gamma} in {} tokens",
13821 tokens.len()
13822 )
13823 .into());
13824 }
13825 }
13826
13827 let n_vocab = self.output.out_features();
13828 let n_embd = self.cfg.n_embd as usize;
13829 let mut cache =
13830 crate::pp::new_cache_planned(e, &self.cfg, &self.plan, tokens.len() + gamma + 8)?;
13831 let (embd_qt, embd_rb) = self.embd.qt_and_row_bytes(n_embd);
13832 let embd_gpu = if spec_host_embd() {
13833 None
13834 } else {
13835 Some(
13836 self.embd_gpu
13837 .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload")),
13838 )
13839 };
13840 let embd_dev = embd_gpu.map(|gpu| (gpu, embd_qt, embd_rb));
13841
13842 struct PendingRecord {
13843 position: usize,
13844 hidden: Option<Vec<f32>>,
13845 tokens: Vec<u32>,
13846 target_top_ids: Vec<Option<Vec<u32>>>,
13847 target_top_logits: Vec<Option<Vec<f32>>>,
13848 target_top_probs: Vec<Option<Vec<f32>>>,
13849 target_tail_probs: Vec<Option<f32>>,
13850 }
13851
13852 let mut pending: Vec<PendingRecord> = anchor_positions
13853 .iter()
13854 .map(|&position| PendingRecord {
13855 position,
13856 hidden: None,
13857 tokens: tokens[position..=position + gamma].to_vec(),
13858 target_top_ids: vec![None; gamma],
13859 target_top_logits: vec![None; gamma],
13860 target_top_probs: vec![None; gamma],
13861 target_tail_probs: vec![None; gamma],
13862 })
13863 .collect();
13864
13865 let mut start = 0usize;
13866 while start < tokens.len() {
13867 let end = (start + chunk).min(tokens.len());
13868 let chunk_tokens = &tokens[start..end];
13869 let (target_logits, hidden_rows) =
13870 self.decode_step_t_core(e, chunk_tokens, start, &mut cache, embd_dev, None)?;
13871 for record in &mut pending {
13872 let hidden_position = record.position - 1;
13873 if hidden_position >= start && hidden_position < end {
13874 let local = hidden_position - start;
13875 record.hidden = Some(
13876 e.dtoh_view(&hidden_rows.slice(local * n_embd..(local + 1) * n_embd))?,
13877 );
13878 }
13879 for slot in 0..gamma {
13880 let target_row = record.position + slot;
13881 if target_row < start || target_row >= end {
13882 continue;
13883 }
13884 let local = target_row - start;
13885 let logits =
13886 e.dtoh_view(&target_logits.slice(local * n_vocab..(local + 1) * n_vocab))?;
13887 let (ids, top_logits, probs, tail) =
13888 dspark_sparse_softmax_topk(&logits, top_k, temperature)?;
13889 record.target_top_ids[slot] = Some(ids);
13890 record.target_top_logits[slot] = Some(top_logits);
13891 record.target_top_probs[slot] = Some(probs);
13892 record.target_tail_probs[slot] = Some(tail);
13893 }
13894 }
13895 start = end;
13896 }
13897
13898 pending
13899 .into_iter()
13900 .map(|record| {
13901 let hidden = record
13902 .hidden
13903 .ok_or_else(|| format!("missing DSpark hidden at {}", record.position))?;
13904 let target_top_ids =
13905 flatten_dspark_rows(record.target_top_ids, record.position, "target ids")?;
13906 let target_top_logits = flatten_dspark_rows(
13907 record.target_top_logits,
13908 record.position,
13909 "target logits",
13910 )?;
13911 let target_top_probs =
13912 flatten_dspark_rows(record.target_top_probs, record.position, "target probs")?;
13913 let target_tail_probs = record
13914 .target_tail_probs
13915 .into_iter()
13916 .enumerate()
13917 .map(|(slot, value)| {
13918 value.ok_or_else(|| {
13919 format!("missing DSpark tail at {} slot {slot}", record.position)
13920 })
13921 })
13922 .collect::<Result<Vec<_>, _>>()?;
13923 Ok(DsparkAnchorRecord {
13924 position: record.position,
13925 hidden,
13926 tokens: record.tokens,
13927 target_top_ids,
13928 target_top_logits,
13929 target_top_probs,
13930 target_tail_probs,
13931 })
13932 })
13933 .collect()
13934 }
13935
13936 /// TEACHER-FORCED REPLAY ACCEPTANCE (hqmtp MTP-heal protocol): walk a FIXED token
13937 /// sequence and, at sampled positions, compare the MTP head's K-token draft chain against
13938 /// the trunk's own teacher-forced greedy predictions. Nothing is generated — the context is
13939 /// the corpus text itself, so (a) degenerate self-generated loops cannot inflate acceptance
13940 /// and (b) two arms (bf16 ceiling vs NVFP4) score on IDENTICAL contexts, isolating the
13941 /// quant-induced head/hidden-state mismatch from text drift.
13942 ///
13943 /// Per eval position p (context = tokens[0..=p], predecessor pairing as in spec decode):
13944 /// draft_j = chain token j from (tokens[p], h_{p-1}), then its own drafts — the exact
13945 /// eager spec-decode chain (same mtp_head_forward_dev, same rope positions).
13946 /// target_j = teacher-forced greedy pick for position p+1+j (argmax of the trunk logits
13947 /// at forced context tokens[0..p+j]). For j==0 this equals live spec
13948 /// acceptance; for j>=1 live verify would condition on the drafts, here it
13949 /// conditions on the corpus — deterministic and arm-comparable by design.
13950 ///
13951 /// Returns (rows, bg): one (p, drafts[k], targets[k]) row per eval position (ascending p),
13952 /// plus the full teacher-forced greedy track bg (bg[i] = greedy pick for position i, i>=1)
13953 /// so harnesses can cross-check runs (e.g. different chunk sizes must give identical bg).
13954 ///
13955 /// `hdump`: when Some, every position's pre-output_norm trunk hidden (the exact rows the
13956 /// draft-KV fill pairs from) streams to the file as little-endian f32 [t_total, n_embd] —
13957 /// the head-distillation extraction (hqmtp): the ENGINE is the source of truth for trunk
13958 /// hiddens (HF torch reproductions of the hybrid trunk measured only ~0.5 greedy
13959 /// agreement vs this path — not usable as a training-data source).
13960 #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
13961 pub fn replay_acceptance(
13962 &self,
13963 e: &Engine,
13964 tokens: &[u32],
13965 k: usize,
13966 stride: usize,
13967 chunk: usize,
13968 mut hdump: Option<&mut std::fs::File>,
13969 ) -> Result<(Vec<(usize, Vec<u32>, Vec<u32>)>, Vec<u32>), Box<dyn std::error::Error>> {
13970 assert!(k >= 1 && stride >= 1 && chunk >= 2);
13971 let mtp = self
13972 .mtp
13973 .as_ref()
13974 .expect("replay_acceptance requires an MTP head");
13975 let n_vocab = self.output.out_features();
13976 let d_vocab = mtp
13977 .shared_head_head
13978 .as_ref()
13979 .unwrap_or(&self.output)
13980 .out_features();
13981 let n_embd = self.cfg.n_embd as usize;
13982 let t_total = tokens.len();
13983 assert!(t_total >= 8, "corpus too short ({t_total} tokens)");
13984 // STAGE-OWNED KV (lane/pp2-spec 2026-08-06) — see `new_session`. Door shut = `Cache::new`.
13985 let mut cache = crate::pp::new_cache_planned(e, &self.cfg, &self.plan, t_total + k + 8)?;
13986 let mut scratch = self.new_mtp_scratch(e, t_total + k + 8)?;
13987 let (embd_qt, embd_rb) = self.embd.qt_and_row_bytes(n_embd);
13988 let embd_gpu = if spec_host_embd() {
13989 None
13990 } else {
13991 Some(
13992 self.embd_gpu
13993 .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload")),
13994 )
13995 };
13996 let embd_dev = embd_gpu.map(|g| (g, embd_qt, embd_rb));
13997
13998 // bg[i] = the trunk's greedy pick for position i under the forced context (i >= 1).
13999 let mut bg: Vec<u32> = vec![0; t_total + 1];
14000 let mut rows: Vec<(usize, Vec<u32>, Vec<u32>)> = Vec::new();
14001 let mut prev_last_h = e.zeros(n_embd)?; // predecessor hidden entering the chunk
14002 let mut seed_buf = e.zeros(n_embd)?;
14003 let mut preds_d = e.alloc_u32_zeroed(chunk)?;
14004 let nll_on = std::env::var("MEMRA_REPLAY_NLL").as_deref() == Ok("1");
14005 let (mut nll_sum, mut nll_cnt) = (0f64, 0u64);
14006 let mut s = 0usize;
14007 while s < t_total {
14008 let cend = (s + chunk).min(t_total);
14009 let tc = cend - s;
14010 let ch = &tokens[s..cend];
14011 // 1. forced trunk pass — verify path (decode-exact contract): all-column logits +
14012 // the chunk's true hiddens.
14013 let (tl_d, vx) = self.decode_step_t_core(e, ch, s, &mut cache, embd_dev, None)?;
14014 for j in 0..tc {
14015 e.argmax_token_device_col(&tl_d, j, n_vocab, &mut preds_d, j)?;
14016 }
14017 let preds = e.dtoh_u32(&preds_d)?;
14018 for j in 0..tc {
14019 bg[s + j + 1] = preds[j];
14020 }
14021 // MEMRA_REPLAY_NLL=1: teacher-forced NLL/perplexity over the same forced pass — the
14022 // checkpoint-quality metric (position j's logits score the GOLD next token).
14023 if nll_on {
14024 let jmax = if cend < t_total { tc } else { tc - 1 }; // last pos has no gold next
14025 if jmax > 0 {
14026 let ids: Vec<u32> = (0..jmax).map(|j| tokens[s + j + 1]).collect();
14027 let rows: Vec<i32> = (0..jmax as i32).collect();
14028 let idsd = e.htod_u32_v(&ids)?;
14029 let rowsd = e.htod_i32(&rows)?;
14030 let mut outd = e.zeros(jmax)?;
14031 e.softmax_gather(&tl_d, n_vocab, &idsd, &rowsd, &mut outd, n_vocab, jmax, 1.0)?;
14032 for pr in e.dtoh(&outd)? {
14033 nll_sum += -((pr.max(1e-30)) as f64).ln();
14034 nll_cnt += 1;
14035 }
14036 }
14037 }
14038 if let Some(f) = hdump.as_deref_mut() {
14039 use std::io::Write;
14040 let host: Vec<f32> = e.dtoh(&vx)?;
14041 // bf16 round-to-nearest-even — f32 doubled the disk bill at bulk
14042 // extraction scale (20M tokens x 4096 = 320GB f32 vs 160GB bf16).
14043 let mut bytes = Vec::with_capacity(tc * n_embd * 2);
14044 for v in &host[..tc * n_embd] {
14045 let b = v.to_bits();
14046 let r = b.wrapping_add(0x7FFF + ((b >> 16) & 1));
14047 bytes.extend_from_slice(&((r >> 16) as u16).to_le_bytes());
14048 }
14049 f.write_all(&bytes)?;
14050 }
14051 // CHAINLESS extraction (stride > corpus, the bulk-hdump mode): no chunk ever
14052 // drafts, so the draft-KV fills are pure waste — skip them (2 MTP-block passes
14053 // per token saved; the forced trunk pass + hdump is all the mode needs).
14054 let chainless = stride > t_total;
14055 if chainless {
14056 e.copy_view_into(
14057 &mut prev_last_h,
14058 0,
14059 &vx.slice((tc - 1) * n_embd..tc * n_embd),
14060 n_embd,
14061 )?;
14062 s = cend;
14063 continue;
14064 }
14065 // 2. TRUE predecessor-paired draft-KV fill for the chunk (row i carries h_{i-1};
14066 // row s reads the previous chunk's last true hidden, zeros at corpus start).
14067 let mut vxs = e.zeros(tc * n_embd)?;
14068 e.copy_into(&mut vxs, 0, &prev_last_h, n_embd)?;
14069 if tc > 1 {
14070 e.copy_view_into(
14071 &mut vxs,
14072 n_embd,
14073 &vx.slice(0..(tc - 1) * n_embd),
14074 (tc - 1) * n_embd,
14075 )?;
14076 }
14077 scratch.set_len(e, s)?;
14078 self.mtp_kv_fill_all(e, ch, &vxs, s, &mut scratch, embd_dev)?;
14079 // 3. draft chains at sampled positions, DESCENDING: a chain reads only slots
14080 // [0..p) (true fills) and appends at >= p; the next (smaller-p) chain's set_len
14081 // truncates those approximate appends before they can ever be read.
14082 let ps: Vec<usize> = (s..cend)
14083 .filter(|p| *p >= 1 && *p % stride == 0 && *p + k <= t_total)
14084 .collect();
14085 for &p in ps.iter().rev() {
14086 scratch.set_len(e, p)?;
14087 if p == s {
14088 e.copy_into(&mut seed_buf, 0, &prev_last_h, n_embd)?;
14089 } else {
14090 e.copy_view_into(
14091 &mut seed_buf,
14092 0,
14093 &vx.slice((p - 1 - s) * n_embd..(p - s) * n_embd),
14094 n_embd,
14095 )?;
14096 }
14097 let mut e_tok = tokens[p];
14098 let mut d_seed = e.clone_dtod(&seed_buf)?;
14099 let chain_heads = !self.mtp_extra.is_empty();
14100 let mut chain_tokens = if chain_heads {
14101 vec![tokens[p]]
14102 } else {
14103 Vec::new()
14104 };
14105 let mut chain_seeds = if chain_heads {
14106 vec![e.clone_dtod(&seed_buf)?]
14107 } else {
14108 Vec::new()
14109 };
14110 let mut drafts: Vec<u32> = Vec::with_capacity(k);
14111 for j in 0..k {
14112 let (dl_d, h_nextn) = if chain_heads {
14113 self.mtp_chain_forward_dev(
14114 e,
14115 &chain_tokens,
14116 &chain_seeds,
14117 &mut scratch,
14118 p,
14119 embd_dev,
14120 None,
14121 )?
14122 } else {
14123 self.mtp_head_forward_dev(
14124 e,
14125 mtp,
14126 e_tok,
14127 &d_seed,
14128 &mut scratch,
14129 p + 1 + j,
14130 embd_dev,
14131 None,
14132 )?
14133 };
14134 let tok_d = e.argmax_token_device(&dl_d, d_vocab)?;
14135 let idx = e.dtoh_u32_one(&tok_d)?;
14136 let d = match &mtp.d2t {
14137 Some(map) => map[idx as usize],
14138 None => idx,
14139 };
14140 drafts.push(d);
14141 if chain_heads {
14142 chain_tokens.push(d);
14143 chain_seeds.push(h_nextn);
14144 } else {
14145 e_tok = d;
14146 d_seed = h_nextn;
14147 }
14148 }
14149 // targets may live in a LATER chunk's bg — resolved after the walk.
14150 rows.push((p, drafts, Vec::new()));
14151 }
14152 // 4. restore TRUE entries for the whole chunk (the next chunk's chains and fills
14153 // expect scratch.len == cend with exact rows).
14154 scratch.set_len(e, s)?;
14155 self.mtp_kv_fill_all(e, ch, &vxs, s, &mut scratch, embd_dev)?;
14156 e.copy_view_into(
14157 &mut prev_last_h,
14158 0,
14159 &vx.slice((tc - 1) * n_embd..tc * n_embd),
14160 n_embd,
14161 )?;
14162 s = cend;
14163 }
14164 for (p, drafts, targets) in rows.iter_mut() {
14165 for j in 0..drafts.len() {
14166 targets.push(bg[*p + 1 + j]);
14167 }
14168 }
14169 rows.sort_by_key(|r| r.0);
14170 if nll_cnt > 0 {
14171 let mean = nll_sum / nll_cnt as f64;
14172 println!(
14173 "[replay-nll] tokens={nll_cnt} nll/token={mean:.5} ppl={:.4}",
14174 mean.exp()
14175 );
14176 }
14177 Ok((rows, bg))
14178 }
14179}
14180
14181#[cfg(test)]
14182mod vg_debt_tests {
14183 use super::dspark_vg_debt_projection;
14184
14185 /// TOOTH for the verify-graph admission accounting: the pool's projected remaining
14186 /// growth must be charged (pre-fix, admission charged 0 for a pool measured at
14187 /// 8,852 MiB), the projection must price the MARGINAL cost of one more key rather than
14188 /// extrapolating the pool's one-time shared allocation, and the doors that make growth
14189 /// impossible must zero the debt.
14190 #[test]
14191 fn vg_debt_projects_remaining_growth_and_respects_the_freeze_valves() {
14192 const MIB: usize = 1 << 20;
14193 let d = dspark_vg_debt_projection;
14194 // cold pool: nothing observed, one capture fits inside SPEC_SHRINK_RESERVE.
14195 assert_eq!(d(0, 256, 0, None), 0);
14196 // freeze valve MEMRA_DSPARK_VG_MAX=0: the pool cannot grow.
14197 assert_eq!(d(10, 0, 500 * MIB, None), 0);
14198 // saturated pool: at/past the cap the pool FREEZES, nothing left to reserve.
14199 assert_eq!(d(256, 256, 8852 * MIB, None), 0);
14200 assert_eq!(d(300, 256, 8852 * MIB, None), 0);
14201
14202 // BOOTSTRAP (one observation, growth unmeasurable): at most one more pool's worth.
14203 // The pre-fix mean rule extrapolated 255x here — the measured 8.5 GB phantom.
14204 assert_eq!(d(1, 256, 33 * MIB, None), 33 * MIB);
14205
14206 // MARGINAL, flat pool (the box9 receipt: reserved stayed ~33.6 MiB across captures
14207 // 1..3, so an additional key costs ~nothing and the debt must collapse to ~0 —
14208 // NOT the 8,556/4,261/2,830 MB the mean rule printed).
14209 assert_eq!(d(3, 256, 33 * MIB, Some((1, 33 * MIB))), 0);
14210
14211 // MARGINAL, genuinely growing pool: 40 MiB per new key over 2 keys, 250 slots left.
14212 let debt = d(6, 256, 273 * MIB, Some((4, 193 * MIB)));
14213 assert_eq!(debt, 250 * (40 * MIB));
14214 assert!(
14215 debt > 3 * (1536 * MIB),
14216 "real growth must dwarf SPEC_SHRINK_RESERVE"
14217 );
14218
14219 // a shrinking/recycled reading never becomes a negative charge.
14220 assert_eq!(d(6, 256, 10 * MIB, Some((4, 99 * MIB))), 0);
14221 // a stale observation at the same capture count falls back to bootstrap.
14222 assert_eq!(d(4, 256, 80 * MIB, Some((4, 80 * MIB))), 80 * MIB);
14223 }
14224}
14225
14226#[cfg(test)]
14227mod capture_headroom_tests {
14228 use super::{
14229 CAPTURE_HEADROOM_FLOOR, capture_err_is_oom, capture_headroom_verdict,
14230 draft_capture_bootstrap_estimate,
14231 };
14232
14233 /// TOOTH for the pre-capture reserve check (lane/step37-vram-admission-20260830): a
14234 /// capture attempt must be refused BEFORE it allocates when the device cannot cover its
14235 /// appetite plus the post-capture floor — and pool-cached bytes count as headroom
14236 /// (driver `free` alone under-counts, the wrong direction for a gate that drops
14237 /// coverage).
14238 #[test]
14239 fn capture_reserve_check_refuses_short_devices_and_counts_pool_cache() {
14240 const MIB: usize = 1 << 20;
14241 let need = 900 * MIB;
14242 // Plenty of room: no refusal.
14243 assert_eq!(
14244 capture_headroom_verdict(8_000 * MIB, 0, need, CAPTURE_HEADROOM_FLOOR),
14245 None
14246 );
14247 // The owner's shape: capture appetite would walk the card to the edge — refused,
14248 // with the arithmetic surfaced for the WARN line.
14249 let (required, effective) =
14250 capture_headroom_verdict(1_200 * MIB, 0, need, CAPTURE_HEADROOM_FLOOR)
14251 .expect("short device must refuse");
14252 assert_eq!(required, need + CAPTURE_HEADROOM_FLOOR);
14253 assert_eq!(effective, 1_200 * MIB);
14254 // Pool-cached bytes are real headroom (the trim path makes them driver-visible).
14255 assert_eq!(
14256 capture_headroom_verdict(1_200 * MIB, 7_000 * MIB, need, CAPTURE_HEADROOM_FLOOR),
14257 None
14258 );
14259 // Boundary: exactly enough is enough (>=, never a fencepost refusal).
14260 assert_eq!(
14261 capture_headroom_verdict(
14262 need + CAPTURE_HEADROOM_FLOOR,
14263 0,
14264 need,
14265 CAPTURE_HEADROOM_FLOOR
14266 ),
14267 None
14268 );
14269 // POLICY at the call site (owner-shape receipts, escalated twice on-box): the
14270 // refusal fn is handed 2x the appetite plus TWO floors — a capture may take at
14271 // most half the discretionary headroom, so the card retains a whole capture's
14272 // worth of room after it lands. One floor of slack above one appetite (the shape
14273 // that step-OOM'd on the owner cell) must therefore REFUSE under the call-site
14274 // requirement.
14275 assert!(
14276 capture_headroom_verdict(
14277 need + CAPTURE_HEADROOM_FLOOR + (100 << 20),
14278 0,
14279 2 * need,
14280 CAPTURE_HEADROOM_FLOOR * 2
14281 )
14282 .is_some()
14283 );
14284 }
14285
14286 #[test]
14287 fn bootstrap_estimate_scales_with_heads_and_never_underflows() {
14288 // 3-head chain on a step37-shaped vocab must expect strictly more than one head.
14289 let one = draft_capture_bootstrap_estimate(1, 3, 128_896, 4_096);
14290 let three = draft_capture_bootstrap_estimate(3, 3, 128_896, 4_096);
14291 assert!(three > one);
14292 // Degenerate shapes keep a sane minimum (the estimate feeds a refusal gate; a
14293 // zero-need gate refuses nothing).
14294 assert!(draft_capture_bootstrap_estimate(0, 0, 0, 0) >= 64 << 20);
14295 }
14296
14297 #[test]
14298 fn capture_oom_predicate_matches_the_quoted_driver_text() {
14299 assert!(capture_err_is_oom(
14300 "DriverError(CUDA_ERROR_OUT_OF_MEMORY, \"out of memory\")"
14301 ));
14302 assert!(capture_err_is_oom("allocation failed: out of memory"));
14303 assert!(!capture_err_is_oom("capture produced no graph"));
14304 }
14305}
14306
14307#[cfg(test)]
14308mod mtp_chain_tests {
14309 use super::mtp_chain_head_index;
14310
14311 #[test]
14312 fn embedded_step_heads_cycle_in_declared_order() {
14313 let actual: Vec<usize> = (0..8).map(|step| mtp_chain_head_index(step, 3)).collect();
14314 assert_eq!(actual, [0, 1, 2, 0, 1, 2, 0, 1]);
14315 }
14316
14317 #[test]
14318 fn standalone_draft_remains_single_head() {
14319 assert!((0..8).all(|step| mtp_chain_head_index(step, 1) == 0));
14320 }
14321}
14322
14323#[cfg(test)]
14324mod tp_verified_prefix_tests {
14325 use super::validate_tp_kv_snapshot_shape;
14326 use crate::tp::ResidentTpKvCache;
14327
14328 #[test]
14329 fn snapshot_shape_accepts_matching_tp_presence() {
14330 let layers = vec![
14331 Some(ResidentTpKvCache::new(Vec::new(), 1, 1, 1, 1, 8)),
14332 None,
14333 ];
14334 validate_tp_kv_snapshot_shape(&layers, &[Some(2), None]).unwrap();
14335 }
14336
14337 #[test]
14338 fn snapshot_shape_rejects_changed_tp_presence() {
14339 let layers = vec![Some(ResidentTpKvCache::new(Vec::new(), 1, 1, 1, 1, 8))];
14340 let error = validate_tp_kv_snapshot_shape(&layers, &[None])
14341 .unwrap_err()
14342 .to_string();
14343 assert!(error.contains("changed shape"), "unexpected error: {error}");
14344 }
14345
14346 #[test]
14347 fn step37_dcw_rebase_at_5151_boundary() {
14348 use crate::cache::KvRingAppend;
14349 let mut cache = ResidentTpKvCache::new_swa(Vec::new(), 128, 128, 136, 96, 262_144, 512);
14350 assert_eq!(cache.physical_capacity(), 5151);
14351 cache.publish_hydration(5150, 0).unwrap();
14352 assert_eq!(cache.ring_base(), Some(0));
14353
14354 // Before rebase, attempting to view rows past capacity fails with the exact bug error
14355 let err = cache.physical_range(5150, 5153).unwrap_err();
14356 assert_eq!(
14357 err,
14358 "SWA ring view [5150,5153) is outside resident [0,5151)"
14359 );
14360
14361 let (write_row, would_rebase) = cache.peek_append_ring(3).unwrap();
14362 assert!(would_rebase);
14363 assert_eq!(write_row, 542);
14364
14365 let tx = cache.begin_transaction().unwrap();
14366 let plan = cache.prepare_append(tx, 3).unwrap();
14367 assert_eq!(plan.target(), 5153);
14368 assert_eq!(plan.write_row(), 542);
14369 assert_eq!(
14370 plan.ring_append(),
14371 Some(KvRingAppend::Rebase {
14372 src_row: 4608,
14373 keep_rows: 542,
14374 new_base: 4608,
14375 write_row: 542,
14376 })
14377 );
14378 cache.publish_append_rebase(plan).unwrap();
14379 cache.publish_append_plan(plan).unwrap();
14380 assert_eq!(cache.ring_base(), Some(4608));
14381
14382 // After rebase, physical range is within bounds
14383 let range = cache.physical_range(5150, 5153).unwrap();
14384 assert_eq!(range, 542..545);
14385
14386 let target = cache.commit_target(tx, 3).unwrap();
14387 cache.publish_finalize(tx, target).unwrap();
14388 assert_eq!((cache.committed_len(), cache.staged_len()), (5153, 5153));
14389 }
14390
14391 #[test]
14392 fn step37_dcw_rebase_rollback_preserves_view_and_base() {
14393 let mut cache = ResidentTpKvCache::new_swa(Vec::new(), 128, 128, 136, 96, 262_144, 512);
14394 cache.publish_hydration(5150, 0).unwrap();
14395
14396 let tx = cache.begin_transaction().unwrap();
14397 let plan = cache.prepare_append(tx, 3).unwrap();
14398 cache.publish_append_rebase(plan).unwrap();
14399 cache.publish_append_plan(plan).unwrap();
14400 assert_eq!(cache.ring_base(), Some(4608));
14401
14402 // Rollback 0 rows accepted (pass declined)
14403 let rollback = cache.commit_target(tx, 0).unwrap();
14404 cache.publish_finalize(tx, rollback).unwrap();
14405 assert_eq!((cache.committed_len(), cache.staged_len()), (5150, 5150));
14406 assert_eq!(cache.ring_base(), Some(4608));
14407
14408 // View for base_len (5150) is still valid in resident [4608, 4608 + 5151)
14409 let range = cache.physical_range(4608, 5150).unwrap();
14410 assert_eq!(range, 0..542);
14411 }
14412}
14413
14414#[cfg(test)]
14415mod dspark_sparse_tests {
14416 use super::dspark_sparse_softmax_topk;
14417
14418 #[test]
14419 fn topk_keeps_full_softmax_mass_and_stable_ties() {
14420 let logits = [1.0f32, 3.0, 3.0, -2.0];
14421 let (ids, top_logits, probs, tail) = dspark_sparse_softmax_topk(&logits, 2, 1.0).unwrap();
14422 assert_eq!(ids, vec![1, 2]);
14423 assert_eq!(top_logits, vec![3.0, 3.0]);
14424 let denominator = logits.iter().map(|value| (value - 3.0).exp()).sum::<f32>();
14425 let expected = 1.0 / denominator;
14426 assert!((probs[0] - expected).abs() < 1.0e-6);
14427 assert!((probs[1] - expected).abs() < 1.0e-6);
14428 assert!((tail - (1.0 - 2.0 * expected)).abs() < 1.0e-6);
14429 assert!((probs.iter().sum::<f32>() + tail - 1.0).abs() < 1.0e-6);
14430 }
14431}
14432
14433#[cfg(test)]
14434mod spec_replay_env_tests {
14435 use super::spec_replay_env_on;
14436
14437 #[test]
14438 fn replay_requires_literal_one() {
14439 assert!(!spec_replay_env_on(None));
14440 assert!(!spec_replay_env_on(Some("")));
14441 assert!(!spec_replay_env_on(Some("0")));
14442 assert!(!spec_replay_env_on(Some("true")));
14443 assert!(!spec_replay_env_on(Some("2")));
14444 assert!(spec_replay_env_on(Some("1")));
14445 }
14446}
14447
14448#[cfg(test)]
14449mod telem_tests {
14450 use super::{SPEC_TELEM_POS, SpecTelemetry, SpecTelemetryCounters};
14451
14452 #[test]
14453 fn synthetic_accept_masks_produce_tau_and_position_histogram() {
14454 let counters = SpecTelemetryCounters::default();
14455 for mask in [
14456 [true, true, true],
14457 [true, true, false],
14458 [true, false, false],
14459 [false, false, false],
14460 ] {
14461 let accepted = mask.iter().take_while(|&&value| value).count();
14462 counters.record_round(mask.len(), accepted);
14463 }
14464
14465 let snapshot = counters.snapshot();
14466 assert_eq!(
14467 (snapshot.rounds, snapshot.drafted, snapshot.accepted),
14468 (4, 12, 6)
14469 );
14470 assert_eq!(&snapshot.pos_drafted[..3], &[4, 4, 4]);
14471 assert_eq!(&snapshot.pos_accepted[..3], &[3, 2, 1]);
14472 assert_eq!(snapshot.tau(), 1.5);
14473 assert_eq!(snapshot.pos_drafted[3..], [0; SPEC_TELEM_POS - 3]);
14474 assert_eq!(snapshot.pos_accepted[3..], [0; SPEC_TELEM_POS - 3]);
14475 }
14476
14477 /// The worker's per-burst pattern: stash, accumulate, diff — the delta must isolate
14478 /// exactly the burst's contribution (pool-resumed sessions carry prior requests' counts).
14479 #[test]
14480 fn delta_isolates_burst_contribution() {
14481 let mut t = SpecTelemetry::default();
14482 // "previous request": 2 rounds of k=3, accepts 3 then 1.
14483 for (kr, na) in [(3usize, 3usize), (3, 1)] {
14484 t.rounds += 1;
14485 t.drafted += kr as u64;
14486 t.accepted += na as u64;
14487 for j in 0..kr {
14488 t.pos_drafted[j] += 1;
14489 }
14490 for j in 0..na {
14491 t.pos_accepted[j] += 1;
14492 }
14493 }
14494 let before = t;
14495 // "this burst": 1 round k=3, accepts 2.
14496 t.rounds += 1;
14497 t.drafted += 3;
14498 t.accepted += 2;
14499 for j in 0..3 {
14500 t.pos_drafted[j] += 1;
14501 }
14502 for j in 0..2 {
14503 t.pos_accepted[j] += 1;
14504 }
14505 let d = t.delta_since(&before);
14506 assert_eq!((d.rounds, d.drafted, d.accepted), (1, 3, 2));
14507 assert_eq!(&d.pos_drafted[..3], &[1, 1, 1]);
14508 assert_eq!(&d.pos_accepted[..3], &[1, 1, 0]);
14509 assert_eq!(d.pos_drafted[3..], [0; SPEC_TELEM_POS - 3]);
14510 }
14511
14512 /// merge(delta) then merge(delta2) equals accumulating both — the per-model /metrics
14513 /// aggregation invariant.
14514 #[test]
14515 fn merge_accumulates_fieldwise() {
14516 let mut agg = SpecTelemetry::default();
14517 let mut d1 = SpecTelemetry {
14518 rounds: 2,
14519 drafted: 6,
14520 accepted: 4,
14521 ..Default::default()
14522 };
14523 d1.pos_drafted[0] = 2;
14524 d1.pos_accepted[0] = 2;
14525 let mut d2 = SpecTelemetry {
14526 rounds: 1,
14527 drafted: 3,
14528 accepted: 1,
14529 ..Default::default()
14530 };
14531 d2.pos_drafted[0] = 1;
14532 d2.pos_accepted[0] = 1;
14533 d2.pos_drafted[1] = 1;
14534 agg.merge(&d1);
14535 agg.merge(&d2);
14536 assert_eq!((agg.rounds, agg.drafted, agg.accepted), (3, 9, 5));
14537 assert_eq!(agg.pos_drafted[0], 3);
14538 assert_eq!(agg.pos_accepted[0], 3);
14539 assert_eq!(agg.pos_drafted[1], 1);
14540 assert_eq!(agg.pos_accepted[1], 0);
14541 }
14542
14543 /// Wrong-snapshot diff saturates to zero instead of wrapping — the counters feed a
14544 /// public metrics surface and must never publish a u64-wrapped garbage value.
14545 #[test]
14546 fn delta_saturates_never_wraps() {
14547 let small = SpecTelemetry {
14548 rounds: 1,
14549 drafted: 2,
14550 accepted: 1,
14551 ..Default::default()
14552 };
14553 let big = SpecTelemetry {
14554 rounds: 5,
14555 drafted: 15,
14556 accepted: 9,
14557 ..Default::default()
14558 };
14559 let d = small.delta_since(&big);
14560 assert_eq!((d.rounds, d.drafted, d.accepted), (0, 0, 0));
14561 }
14562}
14563
14564#[cfg(test)]
14565mod draft_graph_fallback_tests {
14566 use super::DraftGraphFallback;
14567
14568 /// The Q2 contract, part (a): a fallback flip is LOUD — exactly once per flip.
14569 #[test]
14570 fn flip_is_loud_once_and_memoized_after() {
14571 let mut f = DraftGraphFallback::default();
14572 let line = f
14573 .mark_greedy("out of memory")
14574 .expect("first flip must return the warn line");
14575 assert!(
14576 line.contains("WARN"),
14577 "flip line must be warn-level: {line}"
14578 );
14579 assert!(
14580 line.contains("out of memory"),
14581 "flip line must carry the reason: {line}"
14582 );
14583 assert!(f.greedy_failed());
14584 // re-marking an already-failed graph is the memoization: quiet, still failed.
14585 assert!(f.mark_greedy("out of memory").is_none());
14586 assert!(f.greedy_failed());
14587 // the two graphs' flags are independent (greedy flip leaves sampled capturable).
14588 assert!(!f.sampled_failed());
14589 let line_s = f
14590 .mark_sampled("capture unsupported")
14591 .expect("sampled flip is its own flip");
14592 assert!(
14593 line_s.contains("sampled"),
14594 "sampled flip names itself: {line_s}"
14595 );
14596 assert!(f.mark_sampled("capture unsupported").is_none());
14597 }
14598
14599 /// The Q2 contract, part (b): resume-from-pool RESETS both flags (fresh capture chance),
14600 /// and says so exactly when there was something to reset.
14601 #[test]
14602 fn reset_on_resume_clears_flags_and_logs_once() {
14603 let mut f = DraftGraphFallback::default();
14604 // clean session: resume is silent, nothing to reset.
14605 assert!(f.reset_on_resume().is_none());
14606 f.mark_greedy("oom").unwrap();
14607 f.mark_sampled("oom").unwrap();
14608 let note = f
14609 .reset_on_resume()
14610 .expect("a set flag must produce the reset note");
14611 assert!(
14612 note.contains("greedy+sampled"),
14613 "note names what was reset: {note}"
14614 );
14615 assert!(
14616 !f.greedy_failed() && !f.sampled_failed(),
14617 "both flags cleared"
14618 );
14619 // and the NEXT failure after a reset is a fresh flip — loud again.
14620 assert!(f.mark_greedy("oom again").is_some());
14621 let note2 = f.reset_on_resume().expect("greedy-only reset");
14622 assert!(note2.contains("(greedy)"), "single-flag note: {note2}");
14623 }
14624
14625 /// Shape-change clears (dmask realloc / mask-shape mismatch / s_key change) stay silent —
14626 /// they precede a fresh capture attempt whose own failure re-flips loudly.
14627 #[test]
14628 fn shape_change_clears_are_silent() {
14629 let mut f = DraftGraphFallback::default();
14630 f.mark_greedy("oom").unwrap();
14631 f.clear_greedy();
14632 assert!(!f.greedy_failed());
14633 f.mark_sampled("oom").unwrap();
14634 f.clear_sampled();
14635 assert!(!f.sampled_failed());
14636 // after a silent clear there is nothing left for resume to report.
14637 assert!(f.reset_on_resume().is_none());
14638 }
14639}
14640
14641/// SAMPLED DRAFT-GRAPH KEY (lane/graph-s-key-exactness-20260819).
14642///
14643/// These are the CPU teeth for an exactness bug whose live reproduction needs a GPU, a trunk, a
14644/// drafter and a two-turn session: the key itself. Every test below fails against the pre-fix key
14645/// `(seed, temp.to_bits(), k)` — `legacy_key` restates it so the collision is explicit rather
14646/// than remembered.
14647#[cfg(test)]
14648mod sampled_graph_key_tests {
14649 use super::{SampledGraphKey, debug_t_pred0};
14650
14651 /// The pre-fix key, verbatim: `let s_key = (sp_seed, sp_temp.to_bits(), k);`
14652 fn legacy_key(k: &SampledGraphKey) -> (u64, u32, usize) {
14653 (k.seed, k.temp_bits, k.k)
14654 }
14655
14656 fn pure_temp_key() -> SampledGraphKey {
14657 // temperature 1.0, filters off — today's serve default, the shape that parks a graph.
14658 SampledGraphKey::new(12345, 1.0, 3, 0, 1.0, 0.0, false)
14659 }
14660
14661 /// THE COLLISION. Two requests that differ ONLY in the truncation filters shared one key, so
14662 /// a parked pure-temp graph survived into a filtered request and the launch site launched it.
14663 #[test]
14664 fn vendor_filters_change_the_key() {
14665 let parked = pure_temp_key();
14666 // qwen3.8 generation_config.json — what the vendor-default flip makes the default shape.
14667 let vendor = SampledGraphKey::new(12345, 1.0, 3, 20, 0.95, 0.0, false);
14668 assert_eq!(
14669 legacy_key(&parked),
14670 legacy_key(&vendor),
14671 "pre-fix key collided: this is the bug, and the reason a test asserts on it",
14672 );
14673 assert_ne!(parked, vendor, "post-fix key must separate the two regimes");
14674 assert!(parked.pure_temp());
14675 assert!(!vendor.pure_temp());
14676 }
14677
14678 /// Each distribution-shaping field alone is enough to drop the parked graph.
14679 #[test]
14680 fn every_filter_field_is_keyed() {
14681 let base = pure_temp_key();
14682 for (what, other) in [
14683 (
14684 "top_k",
14685 SampledGraphKey::new(12345, 1.0, 3, 20, 1.0, 0.0, false),
14686 ),
14687 (
14688 "top_p",
14689 SampledGraphKey::new(12345, 1.0, 3, 0, 0.95, 0.0, false),
14690 ),
14691 (
14692 "min_p",
14693 SampledGraphKey::new(12345, 1.0, 3, 0, 1.0, 0.05, false),
14694 ),
14695 (
14696 "penalties",
14697 SampledGraphKey::new(12345, 1.0, 3, 0, 1.0, 0.0, true),
14698 ),
14699 ] {
14700 assert_ne!(base, other, "{what} must be part of the key");
14701 assert!(!other.pure_temp(), "{what} leaves the pure-temp regime");
14702 assert_eq!(
14703 legacy_key(&base),
14704 legacy_key(&other),
14705 "{what} was invisible to the pre-fix key",
14706 );
14707 }
14708 }
14709
14710 /// The baked constants stay keyed (this half was always right — regression cover for it).
14711 #[test]
14712 fn baked_constants_stay_keyed() {
14713 let base = pure_temp_key();
14714 assert_ne!(
14715 base,
14716 SampledGraphKey::new(999, 1.0, 3, 0, 1.0, 0.0, false),
14717 "seed"
14718 );
14719 assert_ne!(
14720 base,
14721 SampledGraphKey::new(12345, 0.7, 3, 0, 1.0, 0.0, false),
14722 "temp"
14723 );
14724 assert_ne!(
14725 base,
14726 SampledGraphKey::new(12345, 1.0, 4, 0, 1.0, 0.0, false),
14727 "k"
14728 );
14729 // bitwise on temperature: 0.7f32 vs the same value re-derived must NOT differ.
14730 assert_eq!(
14731 SampledGraphKey::new(1, 0.7, 3, 0, 1.0, 0.0, false),
14732 SampledGraphKey::new(1, 7.0 / 10.0, 3, 0, 1.0, 0.0, false),
14733 );
14734 }
14735
14736 /// THE LOAD-BEARING HALF OF THE SEED DECISION (lane/session-resume-sampler-predicate-
14737 /// 20260820). The whole-session resume predicate deliberately does NOT compare `seed`: an
14738 /// omitted serve `seed` draws fresh per-request entropy, so comparing it would refuse every
14739 /// seed-omitting sampled conversation. That is only sound because the one piece of parked state
14740 /// that BAKES the seed — this graph — is re-keyed on it, so a seed change drops and recaptures.
14741 ///
14742 /// This test is the other end of that argument, asserted here rather than remembered in a
14743 /// comment: if a future change dropped `seed` from the key, the resume predicate's exclusion
14744 /// would silently become the unsound thing it is documented not to be.
14745 /// (Paired with `seed_alone_does_not_refuse` in `memra-sampling`.)
14746 #[test]
14747 fn seed_alone_still_rekeys_the_draft_graph() {
14748 let parked = pure_temp_key();
14749 let reseeded = SampledGraphKey::new(999, 1.0, 3, 0, 1.0, 0.0, false);
14750 assert_ne!(
14751 parked, reseeded,
14752 "a seed-only change MUST drop the parked sampled graph — the resume predicate's \
14753 decision not to compare seed rests on exactly this",
14754 );
14755 // Same regime on both sides: the drop is a recapture, not a fall to the eager chain
14756 // because of a filter difference.
14757 assert!(parked.pure_temp() && reseeded.pure_temp());
14758 }
14759
14760 /// `pure_temp()` is the capture guard's predicate, computed from the key so the two cannot
14761 /// drift. The equality below is the invariant the launch-site guard asserts: identical keys
14762 /// agree on the regime, so a graph that survives the drop is legal to launch.
14763 #[test]
14764 fn equal_keys_agree_on_the_regime() {
14765 let a = SampledGraphKey::new(7, 0.8, 3, 20, 0.95, 0.0, false);
14766 let b = SampledGraphKey::new(7, 0.8, 3, 20, 0.95, 0.0, false);
14767 assert_eq!(a, b);
14768 assert_eq!(a.pure_temp(), b.pure_temp());
14769 // top_p slightly above 1.0 (a client sending 1.0 exactly, or an operator default) is
14770 // still the unfiltered regime, matching the original `sp.top_p >= 1.0` test.
14771 assert!(SampledGraphKey::new(7, 0.8, 3, 0, 1.0, 0.0, false).pure_temp());
14772 assert!(SampledGraphKey::new(7, 0.8, 3, 0, 1.5, -1.0, false).pure_temp());
14773 }
14774
14775 /// The WIDENED capture regime (lane/step37-draft-graph-serving-20260830): truncation-
14776 /// filtered shapes are capturable — the filter runs IN-GRAPH (`filter_stats` +
14777 /// `gumbel_perturb_filtered_ctr`), so the draft draws from the same filtered
14778 /// distribution the accept test reconstructs. Penalties never are: the per-round
14779 /// history cannot be baked. The step37 vendor-default shape (temp 0.5 / top_p 0.9) is
14780 /// exactly the previously-excluded regime this lane exists to capture.
14781 #[test]
14782 fn filtered_regimes_are_capturable_penalties_never() {
14783 let vendor = SampledGraphKey::new(12345, 0.5, 3, 0, 0.9, 0.0, false);
14784 assert!(!vendor.pure_temp());
14785 assert!(vendor.filtered());
14786 assert!(
14787 vendor.graph_capturable(),
14788 "the vendor-default filtered shape must be capturable (default door state)",
14789 );
14790 assert!(pure_temp_key().graph_capturable());
14791 assert!(
14792 !pure_temp_key().filtered(),
14793 "pure-temp takes the legacy (filterless) capture body",
14794 );
14795 let pen = SampledGraphKey::new(12345, 0.5, 3, 0, 0.9, 0.0, true);
14796 assert!(
14797 !pen.graph_capturable(),
14798 "penalty history varies per round and can never be baked into a graph",
14799 );
14800 }
14801
14802 /// MEMRA_DEBUG_SPEC on a SAMPLED spec request past round 0: the print must render without
14803 /// indexing the empty greedy `preds` vector (it panicked the GPU worker before this lane).
14804 #[test]
14805 fn debug_print_survives_the_sampled_arm() {
14806 // round >= 1 with a pending bonus == base 1, sampled == `preds` empty.
14807 assert_eq!(debug_t_pred0(true, 1, 4242, &[]), "n/a");
14808 assert_eq!(debug_t_pred0(true, 2, 4242, &[]), "n/a");
14809 // round 0 without a pending bonus still reports last_pred, in both arms.
14810 assert_eq!(debug_t_pred0(true, 0, 4242, &[]), "4242");
14811 assert_eq!(debug_t_pred0(false, 0, 4242, &[7, 8]), "4242");
14812 // greedy keeps the real prediction it always printed.
14813 assert_eq!(debug_t_pred0(false, 1, 4242, &[7, 8]), "7");
14814 assert_eq!(debug_t_pred0(false, 2, 4242, &[7, 8]), "8");
14815 }
14816}