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}
293pub(crate) fn spec_devacc() -> bool {
294 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
295 *ON.get_or_init(|| std::env::var("MEMRA_SPEC_DEVACC").as_deref() == Ok("1"))
296}
297/// Engine-bundle slice 2 (DSF-ROUNDCOST-20260820 §1.1 host/device round trips + §2 rows 2-3),
298/// DEFAULT ON (`MEMRA_DSPARK_DEFER_READBACK=0` reverts): the dspark round's draft-chain DtoH
299/// is DEFERRED past verify dispatch and merged with the verify-argmax readback into ONE host
300/// sync (2 blocking DtoH/round -> 1). Verify embeds DEVICE tokens (`chain_d`) through the
301/// resident embed table — `embed_gather_u32_t`, bit-identical rows to the host gather by its
302/// own pinned contract. The host therefore dispatches snap + the whole verify while the DRAFT
303/// is still executing, instead of blocking ~1.7 ms on the chain and letting the device drain.
304/// Ladder arm only: the confidence policies size vt from a pre-verify head readback (their
305/// chain readback merges into that same sync instead). Exactness unchanged BY CONSTRUCTION —
306/// same tokens, same kernels, same order; E2E + accept-bank gates arbitrate.
307pub(crate) fn dspark_defer_readback_on() -> bool {
308 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
309 *ON.get_or_init(|| {
310 std::env::var("MEMRA_DSPARK_DEFER_READBACK")
311 .map(|v| v != "0")
312 .unwrap_or(true)
313 })
314}
315/// Engine-bundle slice 1 (DSF-ROUNDCOST-20260820 §1.1, lane/dspark-engine-bundle-20260820),
316/// DEFAULT ON (`MEMRA_STATE_COPY_BATCH=0` reverts): batch the dspark round's GDN state
317/// snapshot and partial-accept restore into single `copy_batch_uniform_f32` launches
318/// instead of ~2 memcpy dispatches (+2 alloc_zeros on the snap side) per linear layer per
319/// round — measured 0.67 ms/round snap + 0.25 ms/round commit of pure dispatch on the q38
320/// route. Launch-structure only: bytes, buffers and stream order are unchanged, so
321/// acceptance and streams stay bit-identical (E2E-gated on the B1 packs).
322pub(crate) fn state_copy_batch_on() -> bool {
323 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
324 *ON.get_or_init(|| {
325 std::env::var("MEMRA_STATE_COPY_BATCH")
326 .map(|v| v != "0")
327 .unwrap_or(true)
328 })
329}
330/// Engine-bundle slice 3 + fa-execupdate slice 4c (DSF-ROUNDCOST-20260820 §5 rank 1),
331/// DEFAULT OFF — `MEMRA_DSPARK_VERIFY_GRAPH=1` opts in: per-(segment, vt) CUDA graphs
332/// for the LINEAR-layer runs, plus the full-verify single graph per (vt, rung) when a
333/// round's rows all ride one seqs rung — see [`DsparkVerifyGraphs`]. Requires the
334/// slice-2 deferred path (device tokens); the eager walk is the byte-identical fallback.
335///
336/// MEASURED disposition (box6 card0, agentic pack, 2026-08-20, both slices): exactness
337/// holds everywhere (ALL EXACT, accept lines byte-match the banks, ckpt-gate oracle
338/// green over the graph + slab-commit paths). Slice-3's AUTO_FREE launch-scan limiter
339/// (25.6 us x 16 launches ≈ 0.41 ms/round) is FIXED — the captured bodies' alloc nodes
340/// are balanced by in-graph frees (census 84/84 per segment, 1776/1776 full) so graphs
341/// instantiate USE_NODE_PRIORITY and the scan is gone. What remains at gate scale:
342/// segment graphs +0.1 tok/s over the batched-rows default (114.4 vs 114.3 x5
343/// interleaved — the linear launch overhead was only ~0.1 ms); the FULL-verify graph is
344/// NET NEGATIVE at gate scale (110.6 vs 114.2: ~14-21 (vt, rung) captures/process at
345/// 2 full-walk executions + ~2.9k-node instantiate each eat far more than the ~0.2-0.3
346/// ms/round of remaining launch overhead). The orchestration ceiling of §1.3 is spent —
347/// the fa/append recovery landed DEFAULT-ON as the batched rows arm
348/// (`dspark_fa_rows_on`), not as a graph. The serve-lifetime cell (DSF-ROUNDCOST §9,
349/// nj-ws-solo) measured the amortization: crossover K≈33 requests, steady −0.246
350/// ms/round, −1.25% session wall over 240 requests — and the graphs-serve lane wired
351/// the door into the session arm (`dspark_spec_session_burst`) as a model-owned
352/// capture pool shared across sessions. Stays opt-in pending the owner's default-ON
353/// ratification on the serve-surface battery.
354pub(crate) fn dspark_verify_graph_on() -> bool {
355 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
356 *ON.get_or_init(|| std::env::var("MEMRA_DSPARK_VERIFY_GRAPH").as_deref() == Ok("1"))
357}
358/// MTP-ROUTE verify graphs, DEFAULT ON for the GDN+MoE family since 2026-08-23
359/// (`MEMRA_SPEC_VERIFY_GRAPH=0` is the kill switch, `=1` opts other families in).
360///
361/// The slice-4c capture already lived inside `qwen35_verify_tparallel` and said so in its own
362/// comment — "stream rides the qwen35moe burst, graphs ride the dspark route" — with no caller
363/// on this route. The MTP spec round is that caller.
364///
365/// WHY it is worth a default (receipts: `research/orndecode-20260822/VGRAPH.md`). With
366/// `MEMRA_SPEC_PHASE=1` this route's round reads verify-ISSUE 44-58% and verify-WAIT **0.0%**:
367/// the host is never waiting for the device, it is spending its own time launching the trunk.
368/// Replay collapses that into one graph launch and the phase all but disappears (55-62 ms ->
369/// 8-10 ms per burst).
370///
371/// MEASURED, two host generations, forced ON/OFF, balanced 4+4 boots in both orders:
372/// * current-generation host (9950X, the serving class): OFF 266.0-266.5, ON 318.8-319.5
373/// tok/s — **+19.7%**, no overlap, sub-1% spread per arm; per-round 6.9 -> 5.7 ms.
374/// * Zen 3 host: +3-9% (that rig's own clock drift is wider than the effect, so the ratio
375/// comes from per-round phase totals, which are internal to each boot).
376/// The ON arm lands at ~320 tok/s on BOTH hosts while OFF tracks host speed — the arm moves
377/// the round off the host and onto the device, which is the whole point.
378///
379/// EXACTNESS is structural (same kernels, same order) and gated anyway: a fixed-seed SAMPLED
380/// completion hashes identically ON vs OFF **and across both hosts** (`08941d5bb9762b21`),
381/// greedy seed-pinned likewise, `run-spec` K=1..8 PASS on both arms with identical acceptance
382/// at every K, kernel-check ALL GREEN.
383///
384/// SCOPE, deliberately narrow: default ON only where it was measured — the GatedDeltaNet +
385/// MoE family (`vgraph_family_default`). Qwen3.8-27B is GDN + DENSE mlp and would otherwise
386/// inherit this default unmeasured, which is the family-by-family law this repo keeps; it can
387/// opt in with `=1` once it has its own interleave. Also never armed together with
388/// ROUND-STREAM, and a round wider than the pool declines it for the eager walk.
389pub(crate) fn spec_verify_graph_env() -> Option<bool> {
390 static ON: std::sync::OnceLock<Option<bool>> = std::sync::OnceLock::new();
391 *ON.get_or_init(
392 || match std::env::var("MEMRA_SPEC_VERIFY_GRAPH").as_deref() {
393 Ok("1") => Some(true),
394 Ok("0") => Some(false),
395 _ => None,
396 },
397 )
398}
399/// SERVE-ROUTE twin of [`dspark_verify_graph_on`], DEFAULT ON — owner-ratified
400/// 2026-08-22 on the §10 serve-lifetime battery (DSF-ROUNDCOST-20260820 §10.3:
401/// crossover K=36–43, steady −0.357 ms/round, session wall −1.55..−1.65%, byte-exact
402/// 240/240 ×3 pairs, pool bounded at 8,852 MiB under `MEMRA_DSPARK_VG_MAX`). The env
403/// stays as the kill-switch: `MEMRA_DSPARK_VERIFY_GRAPH=0` restores the eager walk
404/// (byte-identical body); `MEMRA_DSPARK_VG_MAX=0` is the finer freeze valve. The BIN
405/// arm keeps its own opt-in default (`dspark_verify_graph_on`): at gate scale the
406/// capture toll is never repaid (§8 measured disposition — 14–21 captures over a
407/// 256-token run vs the serve session's thousands of rounds), and the two
408/// instruments must keep their own measured dispositions rather than share one flag.
409pub(crate) fn dspark_verify_graph_serve_on() -> bool {
410 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
411 *ON.get_or_init(|| std::env::var("MEMRA_DSPARK_VERIFY_GRAPH").as_deref() != Ok("0"))
412}
413/// Capture-count ceiling for the dspark verify-graph pool (graphs-serve lane) — the
414/// pool's memory policy STATED instead of silently unbounded. The keyspace is
415/// intrinsically finite — segment keys (run_start, vt) ≤ 16 runs x 7 windows, full
416/// keys (vt, rung, hi) ≤ 7 windows x the split-rung ladder (8 rungs at 32k ctx), ~168
417/// on the q38 export — so the default (256) never engages there; the knob is the
418/// safety valve for a future export with a wider ladder. At the ceiling the pool
419/// FREEZES: existing keys keep replaying, rounds needing a new capture run the eager
420/// walk byte-identically (round-atomic — a partial refusal would mix slab- and
421/// cols-stashed layers inside one commit). No eviction by design: destroying a live
422/// exec graph re-opens the stale-address class the indirect tables exist to close,
423/// and the bounded keyspace makes reclaim worthless.
424pub(crate) fn dspark_vg_cap() -> usize {
425 static CAP: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
426 *CAP.get_or_init(|| {
427 std::env::var("MEMRA_DSPARK_VG_MAX")
428 .ok()
429 .and_then(|v| v.parse().ok())
430 .unwrap_or(256)
431 })
432}
433
434/// PROJECTED REMAINING GROWTH of the verify-graph pool, in bytes (lane/hermes-perf-fixes,
435/// 2026-08-23 — the admission accounting the "pool dwarfs spec admission reserve" finding
436/// asks for). The pool was measured at 8,852 MiB at storm-complete on the q38 export while
437/// admission's transient floor (`SPEC_SHRINK_RESERVE`) is 1.5 GiB and never charged for it:
438/// sessions admitted while the pool is cold overcommit VRAM the pool WILL hold, because the
439/// pool grows monotonically (no eviction by design) and is model-owned across sessions.
440///
441/// SELF-MEASURING, no per-model constant (generic-model law — the 8,852 MiB is a q38 number
442/// and proves nothing about another export): the debt is remaining capture slots x the
443/// MARGINAL bytes a capture adds to this device's graph mem pool.
444///
445/// MARGINAL, NOT MEAN — measured correction (box9 on-box receipt, 2026-08-23). The first
446/// version of this used the mean (`reserved / captures`) and the live serve log showed why
447/// that is wrong: with the pool's reservation flat at ~33.6 MiB across captures 1..3, the
448/// mean-based debt printed **8,556 MB, then 4,261, then 2,830** — it extrapolated capture
449/// #1's ONE-TIME shared allocation (staging buffers, stash slabs, pointer tables: sized
450/// once per pool, shared by every key) across all 256 slots. An 8.5 GB phantom reserve at
451/// boot can refuse admissions that would have fit, which is a worse defect than the
452/// under-charge this accounting exists to remove. The marginal reading prices what an
453/// ADDITIONAL key actually costs: two observations `(captures, reserved)` give
454/// `(r1 - r0) / (c1 - c0)`, which is ~0 on an export whose pool does not grow per key and
455/// tracks real growth on one that does.
456///
457/// BOOTSTRAP (only one observation so far, so growth is unmeasurable): reserve one more
458/// pool's worth — `min(remaining x mean, reserved)`. "We have measured `reserved` bytes for
459/// `captures` keys; until growth is measurable, assume at most a doubling" is fail-safe in
460/// the same direction as the old rule without the 255x extrapolation.
461///
462/// Before the FIRST capture the debt is 0 (a single capture lands well inside the existing
463/// 1.5 GiB floor). `cap` is the intrinsic freeze ceiling (`MEMRA_DSPARK_VG_MAX`; =0 freeze
464/// valve => the pool cannot grow => debt 0); at or past the cap the pool FREEZES, so the
465/// debt is 0 there too.
466pub fn dspark_vg_debt_projection(
467 captures: usize,
468 cap: usize,
469 reserved_bytes: usize,
470 prev: Option<(usize, usize)>,
471) -> usize {
472 if captures == 0 || cap == 0 {
473 return 0;
474 }
475 let remaining = cap.saturating_sub(captures);
476 if remaining == 0 {
477 return 0;
478 }
479 match prev {
480 // marginal growth between two observations of the same pool
481 Some((c0, r0)) if captures > c0 => {
482 let marginal = reserved_bytes.saturating_sub(r0) / (captures - c0);
483 remaining.saturating_mul(marginal)
484 }
485 // bootstrap: at most one more pool's worth
486 _ => remaining
487 .saturating_mul(reserved_bytes / captures)
488 .min(reserved_bytes),
489 }
490}
491/// PRE-CAPTURE VRAM RESERVE CHECK door (lane/step37-vram-admission-20260830), DEFAULT ON.
492/// A draft-graph capture attempt on a tight card used to be try-and-fail: the 2 warmup
493/// forwards + instantiate grew the pool to the edge BEFORE the OOM surfaced, and the
494/// "eager fallback" then ran on a card the failed attempt had just exhausted (the owner's
495/// single-session second-prompt OOM: capture WARN followed by 28 step-OOM engine errors,
496/// device at 5 MiB free). With the gate ON, a capture is attempted only when the device's
497/// effective free (driver free + async-pool cached) covers the capture's expected appetite
498/// PLUS a post-capture safety floor — otherwise the session falls back to eager EARLY,
499/// with headroom intact, through the same LOUD once-per-flip WARN. `=0` restores
500/// try-and-fail (diagnostics door; the trim-on-OOM recovery below stays active either way).
501pub fn spec_capture_gate_on() -> bool {
502 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
503 *ON.get_or_init(|| std::env::var("MEMRA_SPEC_CAPTURE_GATE").as_deref() != Ok("0"))
504}
505
506/// Post-capture safety floor the reserve check keeps free ON TOP of the capture's own
507/// appetite: the same measured constant class as the admission transient floor
508/// (capture arenas + verify activations — the admit-oom control fit). A capture that
509/// would leave less than this behind is not worth its eager-coverage risk.
510pub(crate) const CAPTURE_HEADROOM_FLOOR: usize = 1536 << 20;
511
512/// Pure verdict half of the pre-capture reserve check (unit-testable): given the device's
513/// driver-free and pool-cached bytes and the capture's expected `need`, returns
514/// `Some((required, effective))` when the capture must be REFUSED, `None` when it fits.
515pub(crate) fn capture_headroom_verdict(
516 driver_free: usize,
517 pool_cached: usize,
518 need: usize,
519 floor: usize,
520) -> Option<(usize, usize)> {
521 let effective = driver_free.saturating_add(pool_cached);
522 let required = need.saturating_add(floor);
523 (effective < required).then_some((required, effective))
524}
525
526/// Expected device appetite of a draft-graph capture attempt when no measurement exists
527/// yet (bootstrap only — the model-owned high-water gauge takes over after the first
528/// observed capture). Deliberately conservative and shape-derived, never a per-family
529/// constant: per (head, mode) capture the two warmups + capture each walk one head
530/// forward whose dominant transients are a handful of `n_embd` rows and one `d_vocab`
531/// logits row, retained by the keeper; the sampled tail additionally parks
532/// `k` q-slots + perturb/q buffers of `d_vocab` each.
533pub(crate) fn draft_capture_bootstrap_estimate(
534 heads: usize,
535 k: usize,
536 d_vocab: usize,
537 n_embd: usize,
538) -> usize {
539 let per_capture = 3usize // 2 warmups + capture body, each retaining its transients
540 .saturating_mul(d_vocab.saturating_add(8 * n_embd))
541 .saturating_mul(4)
542 .max(32 << 20); // instantiate + driver-side graph backing per capture, floor
543 let captures = heads.max(1).saturating_mul(2); // interior + last per head
544 let sampled_slots = (k.saturating_add(2))
545 .saturating_mul(d_vocab)
546 .saturating_mul(4);
547 captures
548 .saturating_mul(per_capture)
549 .saturating_add(sampled_slots)
550 .max(64 << 20)
551}
552
553/// OOM predicate for capture-failure recovery (engine-side twin of the worker's
554/// `is_cuda_oom` — the same quoted-text contract).
555pub(crate) fn capture_err_is_oom(reason: &str) -> bool {
556 reason.contains("CUDA_ERROR_OUT_OF_MEMORY") || reason.contains("out of memory")
557}
558
559/// Impure half of the pre-capture reserve check: reads the device, trims the async pool
560/// when the driver alone is short but cached blocks would cover it (graph instantiate and
561/// cuBLAS workspaces allocate from the DRIVER, not from our pool — a pool sitting on freed
562/// blocks starves them), and returns the refusal reason line when the capture must not be
563/// attempted. `None` = go ahead.
564pub(crate) fn capture_headroom_refusal(e: &Engine, need: usize) -> Option<String> {
565 let Ok((driver_free, _total)) = e.ctx().mem_get_info() else {
566 return None; // unreadable device: keep the historical try-and-fail behavior
567 };
568 let pool_cached = e.pool_cached_bytes();
569 // A capture may take AT MOST HALF the discretionary headroom: required =
570 // 2x appetite + two floors (owner's contract: "fall back to eager EARLY with headroom
571 // intact"). Measured escalation on the owner-shape cells: one floor of slack let the
572 // capture walk the card to the edge and the burst step-OOM'd immediately; two floors
573 // still allowed a capture whose session then OOM'd on its own admission-charged work,
574 // because the capture had consumed the memory the charge was counting on. Requiring
575 // the appetite TWICE means the card retains a whole capture's worth of room after the
576 // capture lands - enough for the session's charged classes and its peers' bursts. The
577 // capture is an optimization worth ~2-3 ms of TTFT (draft-graph lane receipts); at the
578 // margin it is never worth an OOM incident.
579 let floor = CAPTURE_HEADROOM_FLOOR.saturating_mul(2);
580 let required_need = need.saturating_mul(2);
581 let required = required_need.saturating_add(floor);
582 match capture_headroom_verdict(driver_free, pool_cached, required_need, floor) {
583 Some((required, effective)) => Some(format!(
584 "insufficient VRAM headroom for capture: effective free {}MB (driver {}MB + pool-cached \
585 {}MB) < required {}MB (2x appetite {}MB + floor {}MB); capture skipped pre-attempt",
586 effective / (1 << 20),
587 driver_free / (1 << 20),
588 pool_cached / (1 << 20),
589 required / (1 << 20),
590 need / (1 << 20),
591 floor / (1 << 20),
592 )),
593 None => {
594 if driver_free < required && pool_cached > 0 {
595 let trimmed = e.pool_trim_to_zero();
596 if trimmed > 0 {
597 eprintln!(
598 "[spec] pre-capture pool trim: released {}MB cached back to the driver \
599 (driver free {}MB < required {}MB; instantiate allocates from the driver)",
600 trimmed / (1 << 20),
601 driver_free / (1 << 20),
602 required / (1 << 20),
603 );
604 }
605 }
606 None
607 }
608 }
609}
610
611/// GRAPH-LAUNCH HEADROOM FLOOR (lane/step37-vram-admission-20260830, defect 3 root
612/// cause): `cuGraphLaunch` SEGFAULTS inside libcuda (offset +0x27c87f, a null internal
613/// dereference at address 0x60) when a captured graph is dispatched into a
614/// driver-exhausted card — reproduced on this lane's box with core dumps on BOTH the
615/// pre-lane and lane binaries (multi-active step-OOM squeeze; the crashing thread sits in
616/// `CudaGraph::launch` inside `generate_spec_inner2`). The eager arms fail RECOVERABLY on
617/// the same card (a quoted CUDA OOM the park path handles), so below this driver-free
618/// floor every graph arm yields to eager for the round. A named constant, not a knob: the
619/// winning value is the default and the guard exists to make a driver segfault
620/// unreachable, not to tune anything.
621pub(crate) const GRAPH_LAUNCH_MIN_FREE: usize = 256 << 20;
622
623/// Per-round guard for the floor above. Read failure keeps serving (never a false
624/// refusal from an unreadable device); one `mem_get_info` (~microseconds) per ~25ms round.
625pub(crate) fn graph_launch_headroom_ok(e: &Engine) -> bool {
626 match e.ctx().mem_get_info() {
627 Ok((free, _total)) => free >= GRAPH_LAUNCH_MIN_FREE,
628 Err(_) => true,
629 }
630}
631
632/// One grep-stable suspension line per ROUTE (each call site holds its own
633/// process-lifetime `Once`): every captured-graph launch route below the floor names
634/// itself in the tag while keeping the same `graph replay suspended:` key the step37
635/// admission lane's squeeze cell greps for. The spec-round guard keeps its original
636/// per-generation `[spec]` line; the sweep routes (graph-launch-guard-sweep lane,
637/// 2026-08-31) note once per process — presence is what the gates assert, and a
638/// suspended round is otherwise byte-identical to its eager twin.
639pub(crate) fn graph_replay_suspended_note(route: &str) {
640 eprintln!(
641 "[{route}] graph replay suspended: driver free below the {}MB launch floor \
642 (eager arms serve; cuGraphLaunch segfaults into an exhausted card)",
643 GRAPH_LAUNCH_MIN_FREE / (1 << 20)
644 );
645}
646
647/// Engine-bundle slice 4 (fa-execupdate lane, DSF-ROUNDCOST-20260820 §6 close: "the
648/// residual gap lives in the FULL-ATTENTION per-row section"), DEFAULT ON —
649/// `MEMRA_DSPARK_FA_ROWS=0` reverts to the per-row loop: when every row of a verify
650/// round takes the v4-seqs arm on ONE `fa_split_keys` rung (the straddle law, evaluated
651/// at the round's first and last t_kv — both eligibility gates are intervals in t_kv),
652/// the qwen35 t-parallel verify's per-row KV-append + fa-decode loop collapses into the
653/// z-batched serving twins: ONE `append_quantize_kv_q8_0_q5_1_seqs` + ONE
654/// `fa_decode_vec_q_seqs_v4` + ONE combine per full-attention layer, replacing
655/// T x (4 dtod row copies + append + 3 memsets + main + combine) launches. Bytes are
656/// pinned by the batched-tick increment-2 kernel-check (seqs-vs-per-seq-loop bit
657/// identity: per-row T_kv derives in-kernel from pos_seq[z]; splits >= ns_eff write the
658/// empty partial the combine never reads, so the shared n_splits_max stride changes no
659/// bytes) and re-gated e2e by this lane's battery.
660pub(crate) fn dspark_fa_rows_on() -> bool {
661 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
662 *ON.get_or_init(|| {
663 std::env::var("MEMRA_DSPARK_FA_ROWS")
664 .map(|v| v != "0")
665 .unwrap_or(true)
666 })
667}
668
669/// `t_pred0` for the `MEMRA_DEBUG_SPEC` per-round print, sampled-safe.
670///
671/// `generate_spec_inner2` fills its `preds` vector ONLY on the greedy path (`if !sampled`), and
672/// the per-round debug print was the sole consumer in the sampled arm: `t_pred(0)` survives round
673/// 0 (`base == 0` returns `last_pred`) and from round 1 (`base == 1`, a pending bonus) indexes an
674/// EMPTY vector — `index out of bounds: the len is 0 but the index is 0`, in the GPU worker
675/// thread, which then respawns and reloads weights while the request dies. So any sampled spec
676/// request longer than one round used to kill the worker whenever `MEMRA_DEBUG_SPEC` was set:
677/// the flag crashed precisely the regime it exists to investigate.
678///
679/// Fixed at the print site, not inside the closure, so the greedy accept walk keeps its strict
680/// indexing (an out-of-range pred there is a real bug and must still be loud).
681fn debug_t_pred0(sampled: bool, base: usize, last_pred: u32, preds: &[u32]) -> String {
682 if base == 0 {
683 return last_pred.to_string();
684 }
685 match preds.get(base - 1) {
686 Some(p) => p.to_string(),
687 // sampled: the greedy per-column argmax was never run for this round.
688 None => {
689 debug_assert!(
690 sampled,
691 "greedy spec: preds[{}] missing at base {base}",
692 base - 1
693 );
694 "n/a".to_string()
695 }
696 }
697}
698
699/// `MEMRA_SKEY_PROBE=1` — sampled-draft-graph key probe (lane/graph-s-key-exactness-20260819).
700///
701/// Reports, per burst and per round, which draft chain the sampled arm chose and under which
702/// filter regime, plus the ONE observable that separates a legal filtered draft from a stale
703/// pure-temp graph replayed under filters: an accept test whose gathered `q` is exactly 0.
704/// A draft token sampled from the FILTERED softmax can never gather q=0 (it was drawn from the
705/// kept set), so `q=0` in the verify means the draft came from a distribution the verify does
706/// not believe in — and `u * 0 < p` then accepts it unconditionally.
707///
708/// Its own env var, deliberately NOT `MEMRA_DEBUG_SPEC`: that flag panicked the GPU worker on
709/// any sampled spec request past round 0 until this lane fixed it (§2 of the bank note).
710pub(crate) fn skey_probe() -> bool {
711 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
712 *ON.get_or_init(|| std::env::var("MEMRA_SKEY_PROBE").as_deref() == Ok("1"))
713}
714
715/// GRAMMAR HOOK for constrained spec decode (lane/constrained-full, 2026-08-03). The engine
716/// stays llguidance-agnostic: the server adapts its per-session grammar state behind this
717/// trait. CONTRACT (the verify-side truncation rule — token-identical to constrained plain
718/// greedy decode): the exactness walk runs UNMASKED first; the hook then (a) truncates
719/// acceptance at the first grammar-illegal accepted token, and (b) when the truncation fired
720/// or the bonus is illegal, the engine recomputes that slot as the MASKED argmax of the
721/// target's own verify column (an unmasked argmax that is grammar-legal IS the masked argmax
722/// — masking only removes tokens — so the common case pays nothing). `consume` advances the
723/// state with each EMITTED token in order; EOS handling is the implementor's job (skip).
724pub trait SpecConstraint {
725 /// -inf the current state's banned ids on a HOST logits row (prompt-tail / init-feed
726 /// masked argmax).
727 fn mask_logits(&mut self, logits: &mut [f32]) -> Result<(), String>;
728 /// Packed 32-bit bitset words of the CURRENT state's allowed set (device-mask form).
729 fn mask_words(&mut self) -> Result<Vec<u32>, String>;
730 /// Is `tok` consumable in the CURRENT state?
731 fn is_allowed(&mut self, tok: u32) -> Result<bool, String>;
732 /// Advance the state with an emitted token.
733 fn consume(&mut self, tok: u32) -> Result<(), String>;
734
735 // --- DRAFT-SIDE MASKING (lane/draft-mask, 2026-08-04) ---
736 // The drafter proposed grammar-illegal tokens under tight schemas, so verify-side
737 // truncation cut nearly every round (measured acceptance 0.467-0.513 tight vs 0.62-0.82
738 // loose, research/constrained-full-20260803). These three methods let the engine mask the
739 // DRAFT model's own sampling with the grammar's legal set, so proposals are legal by
740 // construction. The state they walk is a SPECULATIVE CLONE of the session matcher — the
741 // real state is advanced only by `consume` (emitted tokens), so verify-side truncation
742 // stays the correctness backstop and the emitted stream is unchanged by construction
743 // (an accepted draft is the target's unmasked argmax AND grammar-legal, hence the masked
744 // argmax; a cut slot is recomputed as the masked argmax either way).
745 // Default impls = feature OFF (pre-lane behaviour: unmasked drafts).
746
747 /// Is draft-side masking available on this hook? Probed ONCE per burst, before the draft
748 /// graph is captured (the mask is an in-graph node — its presence is a capture-time shape).
749 fn draft_mask_enabled(&self) -> bool {
750 false
751 }
752 /// Start a draft chain: clone the CURRENT (committed) grammar state into the speculative
753 /// slot. Called once per spec round, before the first draft position.
754 fn draft_begin(&mut self) -> Result<(), String> {
755 Ok(())
756 }
757 /// Packed 32-bit bitset words of the SPECULATIVE state's allowed set (target-vocab ids),
758 /// for the draft position about to be sampled. `None` = draft masking off (no-op).
759 fn draft_mask_words(&mut self) -> Result<Option<Vec<u32>>, String> {
760 Ok(None)
761 }
762 /// Advance the SPECULATIVE state with a PROPOSED draft token. `false` = the chain cannot
763 /// continue (EOS proposed, or an unmasked position proposed something illegal) — the
764 /// engine stops drafting; the token already pushed still goes through verify.
765 fn draft_advance(&mut self, _tok: u32) -> Result<bool, String> {
766 Ok(false)
767 }
768}
769
770/// DRAFT-MASK UPLOAD (lane/draft-mask): pull the speculative state's allowed set (TARGET-id
771/// space) from the hook, project it into the DRAFT head's vocab space, and upload it into the
772/// stable device buffer the draft chain reads. Returns false when the chain must stop drafting:
773/// the hook handed out no mask, or NO draft-vocab row is grammar-legal at this position (a
774/// trimmed FR-Spec head genuinely cannot propose a legal token there — masking it would leave
775/// a fully-banned row whose argmax is meaningless, so the round drafts fewer tokens and the
776/// verify emits the masked argmax as usual).
777fn upload_draft_mask(
778 e: &Engine,
779 c: &mut dyn SpecConstraint,
780 dst: &mut CudaSlice<u32>,
781 d2t: Option<&Vec<u32>>,
782 d_vocab: usize,
783 words: usize,
784) -> Result<bool, Box<dyn std::error::Error>> {
785 let Some(tw) = c
786 .draft_mask_words()
787 .map_err(|e2| format!("constraint: {e2}"))?
788 else {
789 return Ok(false);
790 };
791 let bit = |t: usize| -> bool {
792 let w = t >> 5;
793 w < tw.len() && (tw[w] >> (t & 31)) & 1 == 1
794 };
795 let mut buf = vec![0u32; words];
796 match d2t {
797 // TRIMMED draft head: row i proposes target id d2t[i] — permute the mask accordingly.
798 Some(map) => {
799 for (i, &t) in map.iter().enumerate().take(d_vocab) {
800 if bit(t as usize) {
801 buf[i >> 5] |= 1u32 << (i & 31);
802 }
803 }
804 }
805 // UNTRIMMED: draft ids ARE target ids; the packed words transfer verbatim (a short
806 // mask leaves the padded tail zeroed == banned, same rule as constrained::apply_mask).
807 None => {
808 let n = tw.len().min(words);
809 buf[..n].copy_from_slice(&tw[..n]);
810 }
811 }
812 if buf.iter().all(|w| *w == 0) {
813 return Ok(false);
814 }
815 e.htod_u32_into(dst, &buf)?;
816 Ok(true)
817}
818
819/// Keep the full token-embedding table in host memory and upload only the rows needed by each
820/// MTP/verify step. This is an exact memory-capacity seam for very large BF16 vocab tables: host
821/// gather expands the same source bits to f32, and only O(T*n_embd) bytes cross PCIe per step.
822/// CUDA-graph/round-stream draft paths require device token ids and therefore stay disabled.
823pub(crate) fn spec_host_embd() -> bool {
824 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
825 *ON.get_or_init(|| std::env::var("MEMRA_SPEC_HOST_EMBD").as_deref() == Ok("1"))
826}
827
828/// VERIFY-TIER TRUNK LAUNCH-FUSION (default ON since 2026-07-09; MEMRA_SPEC_FUSED_T=0 reverts — lane/close35b): extend
829/// the t=1 fused2/fused3 Q8_0 trunk launches to the batched verify tier (t=2-4, the K=1..3
830/// verify shapes). At t>1 the trunk pairs/triples (35B wqkv+wqkv_gate, wq/wk/wv,
831/// gate_shexp+up_shexp) each run a separate `matmul_decode_exact` — one q8_1 re-quantize of the
832/// SAME activation plus one _b2/_b4 launch per tensor. The fused twins share ONE quantize and
833/// ONE launch per group; per (tensor,token,row) the kernel body is q8_0_mmvq_batched verbatim
834/// with the identical row mapping -> BIT-IDENTICAL by construction (kernel-check pins it,
835/// run-spec K=1..8 + acceptance identity arbitrate e2e).
836pub(crate) fn spec_fused_t() -> bool {
837 static F: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
838 // DEFAULT ON since 2026-07-09 (MEMRA_SPEC_FUSED_T=0 reverts): verify t=2-4 trunk launch-fusion
839 // (fused2/fused3 Q8_0 batched twins, bit-identical by construction — m=1 block-offset split on
840 // the batched body). m=2 marginal token 2117->1762us; 35B daily: p3 +3.7% (crosses llama), p2 +5%.
841 *F.get_or_init(|| {
842 std::env::var("MEMRA_SPEC_FUSED_T")
843 .map(|v| v != "0")
844 .unwrap_or(true)
845 })
846}
847
848/// zeros/uninit switch for verify-path buffers that are FULLY OVERWRITTEN before any read.
849/// Only call this on such buffers — the lean contract is "identical bytes by construction".
850/// TOKEN-ID GUARD for every id that reaches an embed gather (#87 family).
851///
852/// A device argmax seeds its running index with 0x7FFFFFFF and replaces it only through
853/// comparisons, all of which are FALSE against NaN. An all-NaN logits row therefore returns
854/// the sentinel, and the next thing done with a token id is `embed_row(id)` — table +
855/// ~4.6 TB, never mapped, an MMU fault that kills the CUDA context for the whole process
856/// (research/pp2spec-crash-20260807). The draft chain and the GREEDY verify walk already
857/// trap this; the SAMPLED verify bonus, the boundary sampler and the replay arm's last_pred
858/// did not, which is why the recoverable fault on the greedy instrument is a TERMINAL one on
859/// the vendor-default sampled shape we actually serve.
860pub(crate) fn guard_vocab_token(
861 tok: u32,
862 n_vocab: usize,
863 what: &str,
864) -> Result<u32, Box<dyn std::error::Error>> {
865 if (tok as usize) >= n_vocab {
866 return Err(format!(
867 "{what}: token id 0x{tok:08x} >= n_vocab {n_vocab} — an all-NaN logits row left \
868 the device argmax's init sentinel in place; refusing to dereference the embed \
869 row (#87 trap)"
870 )
871 .into());
872 }
873 Ok(tok)
874}
875
876/// SPEC NaN-ORIGIN SCAN (`MEMRA_SPEC_NAN_SCAN=1`, DEFAULT OFF, diagnostic only).
877///
878/// The `#87` trap reports an all-NaN VERIFY logits column, which says the poison reached the
879/// head but not where it entered. With the scan armed the verify walk syncs and reads back
880/// every layer's output, so the FIRST layer whose residual carries a NaN names itself with the
881/// round's row and position. Off by default and never on a serving path: it costs one host
882/// sync + one `t*n_embd` D2H per layer, and the syncs change scheduling (so a run that stops
883/// reproducing under the scan is itself a datum, not an all-clear).
884///
885/// Rollback seam: unset `MEMRA_SPEC_NAN_SCAN` (or set it to 0). Every call site is behind
886/// `spec_nan_scan()`, so the default path keeps the exact launch sequence it had.
887pub(crate) fn spec_nan_scan() -> bool {
888 spec_nan_scan_level() > 0
889}
890
891/// `MEMRA_SPEC_NAN_SCAN` as a LEVEL, not a boolean. `1` scans each layer's residual, which
892/// names the layer. `2` also scans INSIDE the t-column layer body — the per-column attention
893/// output, the deferred-column o-proj/fa2 join, the post-attention norm and the routed-MoE
894/// output — because "layer 20 poisons row 0" does not say whether the attention or the routed
895/// MoE produced it, and those are different bugs with different fixes.
896pub(crate) fn spec_nan_scan_level() -> u8 {
897 static LVL: std::sync::OnceLock<u8> = std::sync::OnceLock::new();
898 *LVL.get_or_init(|| match std::env::var("MEMRA_SPEC_NAN_SCAN").as_deref() {
899 Ok("1") => 1,
900 Ok("2") => 2,
901 _ => 0,
902 })
903}
904
905/// Read back `[rows, cols]` and fail with the first NaN's coordinates. `what` names the
906/// producer (layer index, walk arm) so the error line is the localization.
907/// VERIFY-ARM RECEIPT (rides `MEMRA_SPEC_NAN_SCAN>=1`, bounded to 200 lines).
908///
909/// Names, per trunk layer, WHICH attention arm the t-column walk actually took. This exists
910/// because the level-1 residual scan below sat only on the non-fused tail: the fused
911/// rope+append+fa arm ends in `continue`, so every layer that fused was NEVER SCANNED and
912/// silently read as "clean". A poisoned residual therefore first reported at the next
913/// non-fused layer, which is how "layer 20 creates the poison" could be true of the scan and
914/// false of the engine. Also carries the row-table lookup counter, so "the fused path never
915/// ran" is distinguishable from "it ran and was innocent".
916/// KV-PLANE SCAN (`MEMRA_KV_PLANE_SCAN=1`, DEFAULT OFF, diagnostic only).
917///
918/// Reads back the STAGED rows of a layer's distributed K/V planes and reports the first row
919/// whose quantization scale is not finite. No kernel required: q8_0 blocks are
920/// `[half d][32 x i8]` and q5_1 blocks carry `half d` then `half m`, so the fp16 scale at the
921/// head of each block is host-checkable straight out of the byte plane.
922///
923/// It exists because the level-2 bad-row bitmap says EVERY verify row is non-finite at a
924/// global-attention layer's join, and row r attends a strict superset of row r-1's keys: that
925/// implicates the shared KV history those rows walk, not per-column staging. "The attention
926/// output is NaN" and "the KV history it attends is already NaN" are different bugs with
927/// different owners, and nothing measured so far separates them. A first-corrupt-row index
928/// also dates the corruption against the prime/decode boundary.
929///
930/// Bounded hard: only layers whose geometry has NO window (the global planes), only the first
931/// `MEMRA_KV_PLANE_SCAN_ROUNDS` verify rounds of a process (default 2), and it copies only
932/// `[0, staged_len)`, which is ~1.6 MB at the 1480-token repro rather than the 262144-row
933/// provision. It still syncs per layer, so it is never a serving or a measured-perf arm.
934pub(crate) fn kv_plane_scan_on() -> bool {
935 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
936 *ON.get_or_init(|| std::env::var("MEMRA_KV_PLANE_SCAN").as_deref() == Ok("1"))
937}
938
939fn kv_plane_scan_rounds() -> usize {
940 static R: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
941 *R.get_or_init(|| {
942 std::env::var("MEMRA_KV_PLANE_SCAN_ROUNDS")
943 .ok()
944 .and_then(|v| v.parse().ok())
945 .unwrap_or(2)
946 })
947}
948
949/// First non-finite fp16 block scale in `bytes`, as (block index, raw u16), scanning one
950/// scale every `stride` bytes. Returns None when every block scale is finite.
951fn first_bad_scale(bytes: &[u8], stride: usize) -> Option<(usize, u16)> {
952 if stride == 0 {
953 return None;
954 }
955 for (i, blk) in bytes.chunks_exact(stride).enumerate() {
956 let raw = u16::from_le_bytes([blk[0], blk[1]]);
957 if half_is_non_finite(raw) {
958 return Some((i, raw));
959 }
960 }
961 None
962}
963
964/// IEEE binary16: exponent all ones is Inf or NaN, whatever the mantissa says.
965fn half_is_non_finite(raw: u16) -> bool {
966 (raw & 0x7C00) == 0x7C00
967}
968
969/// Scan one layer's staged K/V planes for a non-finite quantization scale. Returns the
970/// receipt line, or None when the layer is out of scope or every scale is finite.
971pub(crate) fn scan_kv_plane(
972 e: &crate::Engine,
973 distributed: &memra_kv::ResidentTpKvCache,
974 il: usize,
975 pos0: usize,
976) -> Result<(), Box<dyn std::error::Error>> {
977 // One "round" is one pos0, not one layer: the walk visits 45 layers per verify. The
978 // default of 2 rounds is for a fault that shows up immediately; the step37 repro does not
979 // fire until rep 3 or later, i.e. round ~60 of the process, so that arm MUST raise
980 // MEMRA_KV_PLANE_SCAN_ROUNDS or it will scan only the two rounds that were never going to
981 // be poisoned and report a clean history it never looked at.
982 static ROUNDS: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
983 static LAST_POS: std::sync::atomic::AtomicUsize =
984 std::sync::atomic::AtomicUsize::new(usize::MAX);
985 if LAST_POS.swap(pos0, std::sync::atomic::Ordering::Relaxed) != pos0 {
986 ROUNDS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
987 }
988 if ROUNDS.load(std::sync::atomic::Ordering::Relaxed) > kv_plane_scan_rounds() {
989 return Ok(());
990 }
991 let staged = distributed.staged_len();
992 if staged == 0 {
993 return Ok(());
994 }
995 // ENGAGEMENT RECEIPT. This scan prints only on corruption, so `kvbad=0` in a cell would
996 // read the same whether the history was clean or the scan never ran once. Bounded so a
997 // 45-layer walk cannot flood the log.
998 static SEEN: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
999 let seen = SEEN.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
1000 let (ktb, vtb) = (distributed.k_tok_bytes(), distributed.v_tok_bytes());
1001 if seen < 4 {
1002 eprintln!(
1003 "[kv-plane] engaged #{seen} layer {il} pos0={pos0} staged={staged} \
1004 ktok={ktb} vtok={vtb} (scan armed; a corrupt plane prints its own line)"
1005 );
1006 }
1007 for rank in 0..distributed.ranks().len() {
1008 let Some(rc) = distributed.rank(rank) else {
1009 continue;
1010 };
1011 // q8_0 K blocks are [half d][32 x i8] = 34B; q5_1 V blocks lead with half d then half m.
1012 let kbytes = e.dtoh_u8_view(&rc.k().slice(0..staged * ktb))?;
1013 let vbytes = e.dtoh_u8_view(&rc.v().slice(0..staged * vtb))?;
1014 let kbad = first_bad_scale(&kbytes, 34);
1015 let vbad = first_bad_scale(&vbytes, 24);
1016 if kbad.is_some() || vbad.is_some() {
1017 let row = |b: Option<(usize, u16)>, tok: usize| {
1018 b.map(|(i, raw)| format!("blk {i} (row {}) raw={raw:#06x}", i * 34 / tok.max(1)))
1019 .unwrap_or_else(|| "clean".into())
1020 };
1021 eprintln!(
1022 "[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",
1023 row(kbad, ktb),
1024 row(vbad, vtb)
1025 );
1026 return Ok(());
1027 }
1028 }
1029 Ok(())
1030}
1031
1032pub(crate) fn verify_arm_receipt(
1033 arm: &str,
1034 il: usize,
1035 pos0: usize,
1036 t: usize,
1037 staged: Option<usize>,
1038) {
1039 static N: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
1040 if N.fetch_add(1, std::sync::atomic::Ordering::Relaxed) >= 200 {
1041 return;
1042 }
1043 eprintln!(
1044 "[verify-arm] layer {il} arm={arm} pos0={pos0} t={t} staged_len={} rows_tab_lookups={}",
1045 staged.map(|v| v as i64).unwrap_or(-1),
1046 crate::tp::ROWS_TAB_ENGAGED.load(std::sync::atomic::Ordering::Relaxed)
1047 );
1048}
1049
1050pub(crate) fn nan_scan_rows(
1051 e: &Engine,
1052 buf: &CudaSlice<f32>,
1053 rows: usize,
1054 cols: usize,
1055 what: &str,
1056) -> Result<(), Box<dyn std::error::Error>> {
1057 // The readback is also the ATTRIBUTION point for an asynchronous fault: a
1058 // CUDA_ERROR_ILLEGAL_ADDRESS raised by any launch since the previous scan surfaces on this
1059 // sync, and the bare DriverError names nothing. Wrapping it with `what` turns "the process
1060 // died somewhere" into "it died at or before this layer, on this row, at this position".
1061 let host = e.dtoh(buf).map_err(|err| -> Box<dyn std::error::Error> {
1062 format!(
1063 "spec nan-scan: sync at {what} FAILED: {err} — the fault is at or before \
1064 this point in the walk"
1065 )
1066 .into()
1067 })?;
1068 if host.len() < rows * cols {
1069 return Err(format!(
1070 "nan-scan {what}: buffer holds {} < {rows}x{cols}",
1071 host.len()
1072 )
1073 .into());
1074 }
1075 // SCAN EVERY ROW BEFORE REPORTING. A first-hit return says "row 0 is bad" and leaves the
1076 // other rows UNEXAMINED, which is exactly the bit that discriminates the two mechanisms: in
1077 // the t-column verify, row 0 attends keys [0..p+1) and row 1 attends [0..p+2), a strict
1078 // superset, so poison in the SHARED KV history must appear in BOTH rows, while poison in
1079 // per-column staging can appear in one. Report the whole map.
1080 let mut per_row: Vec<usize> = Vec::with_capacity(rows);
1081 let mut first_bad: Option<(usize, usize)> = None;
1082 for r in 0..rows {
1083 let row = &host[r * cols..(r + 1) * cols];
1084 let bad = row.iter().filter(|v| !v.is_finite()).count();
1085 per_row.push(bad);
1086 if bad > 0 && first_bad.is_none() {
1087 first_bad = Some((r, row.iter().position(|v| !v.is_finite()).unwrap_or(0)));
1088 }
1089 }
1090 if let Some((r0, c0)) = first_bad {
1091 let map: String = per_row
1092 .iter()
1093 .map(|&b| if b == 0 { '.' } else { 'X' })
1094 .collect();
1095 return Err(format!(
1096 "spec nan-scan: {what} produced non-finite values — rows[{rows}] map={map} \
1097 counts={per_row:?} of {cols} each; first at row {r0} element {c0}. Both rows bad \
1098 implicates shared state (the KV history this layer reads); one row bad implicates \
1099 per-column staging."
1100 )
1101 .into());
1102 }
1103 Ok(())
1104}
1105
1106fn vbuf(e: &Engine, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
1107 if spec_lean() { e.uninit(n) } else { e.zeros(n) }
1108}
1109
1110/// Scratch KV for the MTP block (one full-attn layer).
1111///
1112/// PERSISTENT MODE (default, 2026-07-03 — the acceptance lever): sized cap = max_ctx and kept in
1113/// sync with the COMMITTED sequence — slot p holds the MTP block's K/V for committed token p
1114/// (roped p+1, the chain's rope convention), so the draft chain's self-attention sees the FULL
1115/// committed history instead of only the current round's 1..K+1 chain tokens (the reference
1116/// engine's "mtp_update" design). Entries come from two sources:
1117/// - chain appends: accepted positions KEEP their chain-computed entries (embedding exact,
1118/// hidden chain-approximate — the reference engine accepts the same);
1119/// - `mtp_kv_fill` batches: prompt positions + the last-draft position on full accept, computed
1120/// from EXACT trunk hiddens (K/V-only MTP-block pass, no attention/FFN/lm_head).
1121/// Rejected drafts / p-min extras / pseudo-seed appends are all discarded by the round-start
1122/// `set_len` truncation (the KvLayer len mechanism — §C rollback for the draft side).
1123/// Multi-turn spec-decode session (2026-07-05): trunk Cache + persistent MTP draft scratch +
1124/// the committed token list, alive across generate_spec_session calls. Turn N+1 primes ONLY its
1125/// suffix (chunked continuation prime over the quantized past) and mtp_kv_fill's its suffix rows,
1126/// then the round loop runs unchanged. `last_h` carries the pre-output_norm hidden of the last
1127/// committed row across turns (the predecessor-pairing seed + fill anchor).
1128/// Per-request sampling config for the sampled-spec serve path.
1129#[derive(Clone, Copy, Debug)]
1130pub struct SpecSampling {
1131 pub temp: f32,
1132 pub seed: u64,
1133 pub top_k: i32, // 0 = off
1134 pub top_p: f32, // 1.0 = off
1135 pub min_p: f32, // 0.0 = off
1136 pub penalty_last_n: usize, // 0 = penalties off
1137 pub penalty_repeat: f32,
1138 pub penalty_freq: f32,
1139 pub penalty_present: f32,
1140}
1141
1142impl SpecSampling {
1143 /// Non-identity penalties requested — THE `pen_on` predicate (one definition; the
1144 /// same group-off rule `SamplerIdentity::of` canonicalizes: a window with neutral
1145 /// coefficients is penalties-absent). Both spec routes and the dspark accept walk
1146 /// key their penalty arms off this.
1147 pub fn pen_on(&self) -> bool {
1148 self.penalty_last_n > 0
1149 && (self.penalty_repeat != 1.0
1150 || self.penalty_freq != 0.0
1151 || self.penalty_present != 0.0)
1152 }
1153}
1154
1155/// Which draft source a spec session is pinned to. The ENGINE-LEVEL half of
1156/// `DraftSourcePlan` (memra-gguf `model_plan.rs`, always general): the plan states what the
1157/// model DECLARES, this states what actually LOADED and therefore what the session runs.
1158/// Pinned at session creation for the session's lifetime.
1159///
1160/// Family-agnostic on purpose (lane/glm5-extract2, the DraftSource seam): glm5 is today's
1161/// consumer with NativeMtp | Dflash2; the hy3/qwen-next spec lanes select through the same
1162/// three-way law instead of re-deriving it. What each family still owns is the per-session
1163/// STATE behind the kind (see `dflash.rs`'s seam note for why that half is not a trait yet).
1164#[derive(Clone, Copy, Debug, PartialEq, Eq)]
1165pub enum DraftSourceKind {
1166 /// The model's own embedded NextN/MTP head.
1167 NativeMtp,
1168 /// A separately loaded DFlash2 block-diffusion drafter
1169 /// ([`crate::dflash::DflashDrafter`]).
1170 Dflash2,
1171}
1172
1173/// The uniform draft-source selection law. Pure — no env, no engine, no family types — so it
1174/// is CPU-gateable and so every spec family answers "which source" the same way.
1175///
1176/// THE LAW, in precedence order:
1177/// 1. A LOADED DFlash2 drafter IS the source. The operator asked for it by name (a set
1178/// drafter flag that cannot load is already a loud boot failure, never a silent
1179/// fallback), and the family's embedded head is deliberately NOT loaded for this source —
1180/// it is a full trunk layer of VRAM.
1181/// 2. Otherwise the embedded head, and only when the PLAN declares an embedded source: a
1182/// loaded head under a plan that does not declare `Embedded` is a load-path bug, not a
1183/// draft source, and it is refused by name rather than drafted from.
1184/// 3. Otherwise there is no draft source and speculative decode must refuse before drafting.
1185pub fn resolve_draft_source_kind(
1186 plan: memra_gguf::model_plan::DraftSourcePlan,
1187 embedded_head_loaded: bool,
1188 dflash_loaded: bool,
1189) -> Result<DraftSourceKind, String> {
1190 use memra_gguf::model_plan::DraftSourcePlan as P;
1191 if dflash_loaded {
1192 return Ok(DraftSourceKind::Dflash2);
1193 }
1194 if embedded_head_loaded {
1195 if plan != P::Embedded {
1196 return Err(format!(
1197 "an embedded draft head is loaded but the ModelPlan declares \
1198 draft_source={plan:?} — refused rather than drafting from a head the plan \
1199 does not claim"
1200 ));
1201 }
1202 return Ok(DraftSourceKind::NativeMtp);
1203 }
1204 Err(format!(
1205 "no draft source loaded (ModelPlan declares draft_source={plan:?}): speculative \
1206 decode has nothing to draft from"
1207 ))
1208}
1209
1210#[cfg(test)]
1211mod draft_source_kind_tests {
1212 use super::{DraftSourceKind, resolve_draft_source_kind};
1213 use memra_gguf::model_plan::DraftSourcePlan as P;
1214
1215 #[test]
1216 fn a_loaded_drafter_wins_over_a_co_loaded_embedded_head() {
1217 // The operator asked for the drafter BY NAME (a set drafter flag that cannot load is
1218 // already a loud boot failure), so it takes precedence under every plan value —
1219 // including ExternalArtifact, which is what a pack declares when the draft weights
1220 // are not in the model file.
1221 for plan in [P::Embedded, P::ExternalArtifact, P::None] {
1222 assert_eq!(
1223 resolve_draft_source_kind(plan, true, true).unwrap(),
1224 DraftSourceKind::Dflash2,
1225 "plan {plan:?}: a loaded drafter must win"
1226 );
1227 assert_eq!(
1228 resolve_draft_source_kind(plan, false, true).unwrap(),
1229 DraftSourceKind::Dflash2
1230 );
1231 }
1232 }
1233
1234 #[test]
1235 fn the_embedded_head_is_the_source_only_under_a_plan_that_claims_it() {
1236 assert_eq!(
1237 resolve_draft_source_kind(P::Embedded, true, false).unwrap(),
1238 DraftSourceKind::NativeMtp
1239 );
1240 // A head loaded under a plan that does not declare Embedded is a LOAD-PATH BUG, not a
1241 // draft source. Unreachable on glm5 today (its pack hardcodes Embedded and the head
1242 // only loads under it) — which is exactly why it is pinned here: an unreachable
1243 // refusal with no arm is an untested refusal, and the next family is the one that
1244 // makes it reachable.
1245 for plan in [P::ExternalArtifact, P::None] {
1246 let err = resolve_draft_source_kind(plan, true, false)
1247 .expect_err("a head under a non-Embedded plan must refuse");
1248 assert!(err.contains("does not claim"), "{err}");
1249 assert!(err.contains(&format!("{plan:?}")), "{err}");
1250 }
1251 }
1252
1253 #[test]
1254 fn nothing_loaded_refuses_before_drafting_and_names_the_plan() {
1255 for plan in [P::Embedded, P::ExternalArtifact, P::None] {
1256 let err =
1257 resolve_draft_source_kind(plan, false, false).expect_err("no source must refuse");
1258 assert!(err.contains("no draft source loaded"), "{err}");
1259 assert!(err.contains(&format!("{plan:?}")), "{err}");
1260 }
1261 }
1262}
1263
1264/// `MEMRA_SPEC_PMIN` break semantics over per-slot draft confidences (the chain break this
1265/// module's drafting loops apply inline: `p < p_min && (j > 0 || pmin0)`): keep the longest
1266/// prefix whose every slot clears `p_min`; slot 0 survives a miss unless PMIN0 arms
1267/// zero-draft rounds. Prefix truncation is forced by the accept rule anyway (a kept slot
1268/// after a dropped one could never commit — the dspark confidence-slot argument). Pure so
1269/// the rule is CPU-gateable; the SHARED K-policy surface every spec family consumes
1270/// (hoisted from the glm5 loop, lane/glm5-extract-general).
1271pub fn spec_conf_keep(q: &[f32], p_min: f32, pmin0: bool) -> usize {
1272 if p_min <= 0.0 {
1273 return q.len();
1274 }
1275 let mut kept = 0usize;
1276 for (j, &qj) in q.iter().enumerate() {
1277 if qj < p_min && (j > 0 || pmin0) {
1278 break;
1279 }
1280 kept += 1;
1281 }
1282 kept
1283}
1284
1285/// Host Philox4x32-10 uniform in (0,1) — mirrors spec_sample.cu's `philox4`/`u01` with the
1286/// ctr_lo tag 0xFFFF_FFFE, so the host accept-test stream never collides with any device
1287/// sampling event (device Gumbel uses (i>>2, stream_pos); device residual uses 0xFFFF_FFFD).
1288/// One value per (seed, ctr) EVENT; callers own the counter discipline. Extracted verbatim
1289/// from generate_spec_inner2's closure for the dspark sampled-admission walk (the two paths
1290/// MUST consume the identical stream construction — two ad-hoc Philox copies drifting apart
1291/// is a distributional bug, not a style problem).
1292pub(crate) fn host_u01(seed: u64, ctr: u32) -> f32 {
1293 let (m0, m1) = (0xD2511F53u32, 0xCD9E8D57u32);
1294 let (mut c0, mut c1, mut c2, mut c3) = (0xFFFF_FFFEu32, ctr, 0u32, 0u32);
1295 let (mut k0, mut k1) = ((seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
1296 for _ in 0..10 {
1297 let (h0, l0) = (((m0 as u64 * c0 as u64) >> 32) as u32, m0.wrapping_mul(c0));
1298 let (h1, l1) = (((m1 as u64 * c2 as u64) >> 32) as u32, m1.wrapping_mul(c2));
1299 let (n0, n1, n2, n3) = (h1 ^ c1 ^ k0, l1, h0 ^ c3 ^ k1, l0);
1300 c0 = n0;
1301 c1 = n1;
1302 c2 = n2;
1303 c3 = n3;
1304 k0 = k0.wrapping_add(0x9E3779B9);
1305 k1 = k1.wrapping_add(0xBB67AE85);
1306 }
1307 (c0 as f32 + 1.0) * (1.0 / 4294967296.0)
1308}
1309
1310/// Tracked draft positions for [`SpecTelemetry`] (serve K defaults to 3; the run-spec gate
1311/// sweeps K=1..8, and MEMRA_SPEC_CAPMAX defaults to 7 — 8 covers every tuned config).
1312pub const SPEC_TELEM_POS: usize = 8;
1313
1314/// Always-on per-draft-position acceptance telemetry (lane/accept-telemetry, 2026-08-05 —
1315/// the llama.cpp #26389 / vLLM spec-decode counter schema, upstream-sweeps 2026-08-05).
1316/// Lives on the [`SpecSession`] and accumulates across bursts; the serve worker diffs a
1317/// stashed copy per burst for its per-model /metrics aggregation and per-request usage.
1318/// Same normalization as the `[spec-stats]` line: p-min-discarded chain tokens are counted
1319/// in NEITHER drafted nor accepted.
1320#[derive(Clone, Copy, Default, Debug)]
1321pub struct SpecTelemetry {
1322 /// verify rounds completed (a round-stream burst counts each of its M rounds).
1323 pub rounds: u64,
1324 /// tokens drafted / accepted across all rounds.
1325 pub drafted: u64,
1326 pub accepted: u64,
1327 /// how often draft position j (0-based within a round's chain) was offered / accepted.
1328 /// Positions >= SPEC_TELEM_POS are untracked (totals still count them). The opt-in
1329 /// round-stream arm (MEMRA_SPEC_STREAM=1) reads back only totals, so under it these
1330 /// arrays cover the standard-path rounds only and their sums may undercount the totals.
1331 pub pos_drafted: [u64; SPEC_TELEM_POS],
1332 pub pos_accepted: [u64; SPEC_TELEM_POS],
1333}
1334
1335impl SpecTelemetry {
1336 /// Fieldwise `self - prev` — the worker's per-burst delta off a copy stashed before the
1337 /// burst call. Saturating: a caller diffing against the wrong snapshot gets zeros, not
1338 /// a wrapped counter.
1339 pub fn delta_since(&self, prev: &SpecTelemetry) -> SpecTelemetry {
1340 let mut d = SpecTelemetry {
1341 rounds: self.rounds.saturating_sub(prev.rounds),
1342 drafted: self.drafted.saturating_sub(prev.drafted),
1343 accepted: self.accepted.saturating_sub(prev.accepted),
1344 ..Default::default()
1345 };
1346 for j in 0..SPEC_TELEM_POS {
1347 d.pos_drafted[j] = self.pos_drafted[j].saturating_sub(prev.pos_drafted[j]);
1348 d.pos_accepted[j] = self.pos_accepted[j].saturating_sub(prev.pos_accepted[j]);
1349 }
1350 d
1351 }
1352 /// Fieldwise `self += d` — the worker's per-model aggregation.
1353 pub fn merge(&mut self, d: &SpecTelemetry) {
1354 self.rounds += d.rounds;
1355 self.drafted += d.drafted;
1356 self.accepted += d.accepted;
1357 for j in 0..SPEC_TELEM_POS {
1358 self.pos_drafted[j] += d.pos_drafted[j];
1359 self.pos_accepted[j] += d.pos_accepted[j];
1360 }
1361 }
1362
1363 /// Mean accepted draft-prefix length per verify round (tau).
1364 pub fn tau(&self) -> f64 {
1365 if self.rounds > 0 {
1366 self.accepted as f64 / self.rounds as f64
1367 } else {
1368 0.0
1369 }
1370 }
1371}
1372
1373/// Session-lifetime atomic acceptance counters. The verifier records only after the greedy or
1374/// rejection-sampling walk has resolved on the host, so these relaxed increments add no GPU
1375/// launch, synchronization, allocation, or ordering dependency to the numeric path.
1376struct SpecTelemetryCounters {
1377 rounds: AtomicU64,
1378 drafted: AtomicU64,
1379 accepted: AtomicU64,
1380 pos_drafted: [AtomicU64; SPEC_TELEM_POS],
1381 pos_accepted: [AtomicU64; SPEC_TELEM_POS],
1382}
1383
1384impl Default for SpecTelemetryCounters {
1385 fn default() -> Self {
1386 Self {
1387 rounds: AtomicU64::new(0),
1388 drafted: AtomicU64::new(0),
1389 accepted: AtomicU64::new(0),
1390 pos_drafted: std::array::from_fn(|_| AtomicU64::new(0)),
1391 pos_accepted: std::array::from_fn(|_| AtomicU64::new(0)),
1392 }
1393 }
1394}
1395
1396impl SpecTelemetryCounters {
1397 fn record_round(&self, drafted: usize, accepted: usize) {
1398 debug_assert!(accepted <= drafted);
1399 self.rounds.fetch_add(1, Ordering::Relaxed);
1400 self.drafted.fetch_add(drafted as u64, Ordering::Relaxed);
1401 self.accepted.fetch_add(accepted as u64, Ordering::Relaxed);
1402 for counter in self.pos_drafted.iter().take(drafted) {
1403 counter.fetch_add(1, Ordering::Relaxed);
1404 }
1405 for counter in self.pos_accepted.iter().take(accepted) {
1406 counter.fetch_add(1, Ordering::Relaxed);
1407 }
1408 }
1409
1410 /// Round-stream keeps each round's accept length on device; retain exact scalar totals while
1411 /// leaving the per-position arrays untouched, matching the pre-existing telemetry contract.
1412 fn record_totals(&self, rounds: usize, drafted: usize, accepted: usize) {
1413 self.rounds.fetch_add(rounds as u64, Ordering::Relaxed);
1414 self.drafted.fetch_add(drafted as u64, Ordering::Relaxed);
1415 self.accepted.fetch_add(accepted as u64, Ordering::Relaxed);
1416 }
1417
1418 fn snapshot(&self) -> SpecTelemetry {
1419 SpecTelemetry {
1420 rounds: self.rounds.load(Ordering::Relaxed),
1421 drafted: self.drafted.load(Ordering::Relaxed),
1422 accepted: self.accepted.load(Ordering::Relaxed),
1423 pos_drafted: std::array::from_fn(|j| self.pos_drafted[j].load(Ordering::Relaxed)),
1424 pos_accepted: std::array::from_fn(|j| self.pos_accepted[j].load(Ordering::Relaxed)),
1425 }
1426 }
1427}
1428
1429pub struct SpecSession {
1430 pub(crate) cache: Cache,
1431 pub(crate) scratch: MtpScratch,
1432 /// Every token whose state the caches hold, in order (prompt turns + generated), INCLUDING
1433 /// overshoot: spec commits accepted drafts past max_new; those rows are in the caches, so the
1434 /// session must count them. Callers render output from this, not from their own echo.
1435 pub committed: Vec<u32>,
1436 /// Pre-output_norm hidden of the LAST committed row (device). None before the first turn.
1437 pub(crate) last_h: Option<CudaSlice<f32>>,
1438 /// Greedy argmax predicting the token AFTER committed.last() (from the last turn's final
1439 /// logits). Fuels empty-suffix continuation bursts (serve): the next turn emits this token
1440 /// first, feeds it, and the round loop resumes without any prime. None before the first turn.
1441 pub next_pred: Option<u32>,
1442 /// SAMPLED-SPEC stream continuity across bursts: Philox event counters persist here so a
1443 /// session's randomness never repeats between generate_spec_session calls. (0,0) at admit.
1444 pub sctr: u32,
1445 pub uctr: u32,
1446 /// PERSISTENT DRAFT-GRAPH CONTEXT (2026-08-01, the serve-burst fixed-cost fix): the captured
1447 /// draft graph(s) + every device I/O buffer they bake, carried ACROSS generate_spec_session
1448 /// calls. Before this, every serve burst re-captured the draft graph (2 warmup forwards +
1449 /// instantiate) — measured ~16ms/burst on H100 q27 (MEMRA_SPEC_BURST sweep,
1450 /// research/spec-serving-20260801). None before the first turn; error paths drop it
1451 /// (next burst recaptures — serve retires errored sessions anyway).
1452 pub(crate) draft_ctx: Option<DraftGraphCtx>,
1453 /// PENDING-CARRY across bursts (2026-08-01, the serve burst-boundary fix): the bonus token
1454 /// emitted by the last round but NOT committed to the caches. The old tail committed it with
1455 /// a solo T=1 trunk pass (+ draft fill), and the next burst's setup fed the stashed next_pred
1456 /// with ANOTHER solo pass — 2x ~11.5ms/burst measured on H100 q27 ([spec-setup] trace).
1457 /// Carrying it lets the next empty-suffix greedy burst consume it as round-0 verify col 0,
1458 /// exactly like a mid-burst full-accept boundary (no solo passes). INVARIANT: when set,
1459 /// `committed` (== cache rows) EXCLUDES this token although it was already emitted in the
1460 /// last burst's output, and `last_h` holds the hidden of the last COMMITTED row (its
1461 /// predecessor — the chain-seed/fill anchor). `next_pred` is None (unknown without the
1462 /// commit pass). Non-empty-suffix or sampled turns must flush first (spec_flush_pending);
1463 /// generate_spec_session_sampled does this at entry, and serve parks only flushed sessions.
1464 pub pending_tok: Option<u32>,
1465 /// SESSION-AFFINITY TURN CHECKPOINT (lane/session-affinity, 2026-08-05): the state at this
1466 /// turn's PROMPT-END boundary, retained so a later turn can REWIND here. See
1467 /// [`SpecCheckpoint`]. Refreshed by every non-empty prime; None until the first one, and on
1468 /// a rig too tight to hold it (a failed capture is silent — resume just isn't available).
1469 pub(crate) turn_ckpt: Option<SpecCheckpoint>,
1470 /// Session-lifetime acceptance telemetry. Relaxed atomics update at the host-side round
1471 /// accounting the loop already does — no syncs, no allocation. NOTE a
1472 /// pool-resumed session carries the PREVIOUS requests' counts; per-request consumers
1473 /// diff with [`SpecTelemetry::delta_since`] around each burst.
1474 telem: SpecTelemetryCounters,
1475 /// PREFIX-CACHE publication request (lane/spec-prefix-cache): worker sets this to the
1476 /// miss-LCP boundary before a cold burst; the prime captures at exactly that split (it must
1477 /// coincide with the burst's `prime_split` or no capture happens). One-shot: consumed by the
1478 /// prime, result lands in `boundary_captures`.
1479 pub capture_at: Option<usize>,
1480 /// The captures the last prime produced (see [`SpecBoundaryCapture`]). Worker drains them
1481 /// post-burst to assemble prefix entries. A failed capture is silent, like `turn_ckpt` —
1482 /// publication just isn't available for that request. Plural since
1483 /// lane/frspec-multiturn-cache (2026-08-21): a cold burst can capture BOTH the miss-LCP
1484 /// split (the shared-prefix class) and the stable pre-generation boundary (the
1485 /// next-turn re-render class) — one entry per stop, exactly the boundary set the plain
1486 /// prefill tick publishes/checkpoints.
1487 pub boundary_captures: Vec<SpecBoundaryCapture>,
1488 /// STABLE-BOUNDARY TURN CHECKPOINT REQUEST (lane/frspec-multiturn-cache, 2026-08-21): the
1489 /// ABSOLUTE committed-length position the next non-empty prime should capture `turn_ckpt`
1490 /// at, instead of prompt-end. The worker sets it to the STABLE PRE-GENERATION boundary
1491 /// (`plain_checkpoint_boundary` — before the live generation header the client rewrites),
1492 /// porting the 2026-08-09 plain-tier fix: a prompt-end spec checkpoint includes the
1493 /// template's live assistant-generation header (`<|im_start|>assistant\n<think>\n`), which
1494 /// the NEXT turn's re-render replaces, so `affinity_match` diverged a couple tokens below
1495 /// the checkpoint and the spec pool declined 100% of multi-turn agent traffic (measured:
1496 /// `spec-affinity: declined (history diverged at 6811 of checkpoint 6813)`,
1497 /// research/multiturn-cache-20260821 B4). One-shot, `capture_at` convention; None = legacy
1498 /// prompt-end capture.
1499 pub ckpt_at: Option<usize>,
1500 /// FAIL-SAFE (lane/step37-vram-admission-20260830, external-review corroboration): set
1501 /// by the worker on a session serving a step-OOM park REPLAY. The burst entry pre-marks
1502 /// the draft-graph fallback so the replay never re-enters the capture path — the capture
1503 /// appetite is part of what drove the card to the OOM, and a replay that recaptures
1504 /// re-runs the incident. If the eager replay still cannot fit, the bounded retry budget
1505 /// exhausts into the honest recoverable Overloaded error instead of looping.
1506 pub capture_disabled: bool,
1507}
1508impl SpecSession {
1509 /// Context capacity of the session's caches (the server's ContextFull guard).
1510 pub fn cache_max_ctx(&self) -> usize {
1511 self.cache.max_ctx
1512 }
1513 /// Read access to the live trunk cache (lane/spec-prefix-cache): the worker slices
1514 /// full-attn KV rows `[0..capture.pos)` out of it when publishing a boundary capture —
1515 /// those rows are append-only for the session's lifetime (rollbacks never truncate below
1516 /// the prime boundary), so no copy was taken at prime time.
1517 pub fn cache_ref(&self) -> &Cache {
1518 &self.cache
1519 }
1520 /// Read access to the persistent draft-scratch plane (lane/spec-on-cache-hit): the
1521 /// worker slices rows `[0..capture.pos)` when publishing a boundary capture, exactly
1522 /// like the trunk KV — draft rows below the prompt end are append-only for the
1523 /// session's lifetime (the prime fill wrote them once; rollbacks reset `len_d` to the
1524 /// committed length, never below the prime boundary, and the true-hidden refresh
1525 /// rewrites generated positions only). Returns `(k, v, k_tok_bytes, v_tok_bytes)`.
1526 /// None when the scratch is ring-backed (Step35 SWA — physical rows are not
1527 /// prefix-addressable; the prefix cache already refuses that class end to end).
1528 pub fn draft_plane_ref(&self) -> Option<(&CudaSlice<u8>, &CudaSlice<u8>, usize, usize)> {
1529 if self.scratch.kv.ring.is_some() {
1530 return None;
1531 }
1532 Some((
1533 &self.scratch.kv.k,
1534 &self.scratch.kv.v,
1535 self.scratch.kv.k_tok_bytes,
1536 self.scratch.kv.v_tok_bytes,
1537 ))
1538 }
1539 /// Snapshot the session's process-local acceptance counters for per-burst diffing.
1540 pub fn telemetry(&self) -> SpecTelemetry {
1541 self.telem.snapshot()
1542 }
1543 /// Committed position this session can REWIND to (its retained prompt-end boundary), if any.
1544 /// A request whose prompt matches `committed[..pos]` exactly can resume from here — see
1545 /// `spec_rewind_to_checkpoint`.
1546 pub fn rewind_pos(&self) -> Option<usize> {
1547 self.turn_ckpt.as_ref().map(|c| c.pos)
1548 }
1549 /// Whether every ring-backed trunk/draft row needed by the retained checkpoint is resident.
1550 pub fn rewind_is_resident(&self) -> bool {
1551 self.turn_ckpt.as_ref().is_some_and(|ckpt| {
1552 self.cache.can_rollback(&ckpt.snap, 0) && self.scratch.can_rewind_to(ckpt.pos)
1553 })
1554 }
1555 /// Is this session in the DEMOTION-READY shape (see [`SpecSession::into_demoted`])?
1556 /// `false` means a carried pending must be flushed first (`spec_flush_pending`), or the
1557 /// session has never run a turn and has no prediction to hand over.
1558 pub fn demote_ready(&self) -> bool {
1559 self.pending_tok.is_none() && self.next_pred.is_some()
1560 }
1561 /// Does this session hold a carried pending bonus (flush required before a handoff/park)?
1562 pub fn has_pending(&self) -> bool {
1563 self.pending_tok.is_some()
1564 }
1565 /// Committed row count == cache rows (the session invariant), for the caller's own
1566 /// `fed`-length cross-check at a handoff boundary.
1567 pub fn committed_len(&self) -> usize {
1568 self.committed.len()
1569 }
1570 /// DEMOTION HANDOFF (lane/spec-gate, 2026-08-07): consume this session and hand its trunk
1571 /// cache + next-token prediction to the plain batched-decode path.
1572 ///
1573 /// WHY THIS IS EXACT (greedy). The invariant at a burst boundary is `cache.pos ==
1574 /// committed.len()`: every committed row has trunk KV + recurrent state, exactly as a plain
1575 /// tokenwise prime of the same `committed` sequence would have left it (that is the
1576 /// session-tail contract, and the same property `spec_rewind_to_checkpoint` and the reuse
1577 /// pool already rely on). `next_pred` is the argmax of the verify's logits for the LAST
1578 /// committed row — and verify-column logits are bit-identical to plain decode's logits at
1579 /// that position, because `matmul_decode_exact` bit-identity IS the basis of the greedy
1580 /// accept walk. So handing (cache, next_pred) to the batched path continues the stream from
1581 /// a state indistinguishable from one the batched path produced itself: the batched tick
1582 /// emits `next_pred`, feeds it into this same cache, and decodes on.
1583 ///
1584 /// `None` when the session is not in the handoff shape — a carried pending (its bonus row is
1585 /// NOT in the cache, so `spec_flush_pending` must commit it first) or no `next_pred` yet
1586 /// (never bursted). Callers must not force it: a half-committed cache handed to the batched
1587 /// path would silently skip a token.
1588 ///
1589 /// The MTP draft scratch, the persistent draft-graph context and the turn checkpoint are
1590 /// DROPPED here (freeing their VRAM): the batched path never drafts, and this handoff is
1591 /// one-way by design — there is no cheap symmetric re-promotion (rebuilding the draft KV
1592 /// would mean an `mtp_kv_fill` over the whole committed history).
1593 pub fn into_demoted(self) -> Option<(Cache, u32)> {
1594 if self.pending_tok.is_some() || self.cache.tainted {
1595 return None;
1596 }
1597 let np = self.next_pred?;
1598 debug_assert_eq!(
1599 self.cache.pos,
1600 self.committed.len(),
1601 "demotion handoff: cache rows != committed tokens"
1602 );
1603 Some((self.cache, np))
1604 }
1605 /// Pool-resume hook (audit Q2): clear the parked draft-graph failure memoization so a
1606 /// NEW request resuming this session gets one fresh capture chance — a transient-pressure
1607 /// capture failure must not persist for the pool's whole lifetime (the TRT #16072 class).
1608 /// Logs once iff a flag was actually set; a no-fallback resume is silent and free.
1609 pub fn reset_graph_fallback_on_resume(&mut self) {
1610 if let Some(line) = self
1611 .draft_ctx
1612 .as_mut()
1613 .and_then(|c| c.failed.reset_on_resume())
1614 {
1615 eprintln!("{line}");
1616 }
1617 }
1618}
1619
1620/// A session's PROMPT-END boundary state, the rewind target for session-affinity resume.
1621///
1622/// WHY THIS BOUNDARY, AND WHY IT IS THE ONLY ONE WORTH KEEPING. The rewrite class this lane
1623/// exists for (a client that strips `<think>` blocks out of prior assistant turns) mutates the
1624/// text the session GENERATED, never the prompt it was given. So turn N's prompt agrees with
1625/// turn N-1's committed tokens up to almost exactly where turn N-1's generation began — the
1626/// prompt-end boundary. Keeping a checkpoint there means the next turn re-primes only its own
1627/// delta (the rewritten answer + the new user turn) instead of the whole conversation.
1628///
1629/// WHAT IT MUST HOLD. Full-attn KV is append-only and position-addressed, so rewinding it is a
1630/// `len` truncation (no data). Linear-attn (GDN) conv/ssm state is mutated IN PLACE with no
1631/// position index, so it must be a real device COPY — that copy is the entire reason a spec
1632/// session could not previously rewind. The MTP draft scratch needs no copy either: its rows
1633/// below the boundary were written by this turn's fill and are never revisited (the per-round
1634/// true-hidden refresh only rewrites the CURRENT burst's committed positions), so rewinding it
1635/// is also just a `len` reset. `last_h` is the hidden of the last row below the boundary — the
1636/// predecessor-pairing anchor the next prime's fill reads for its first row.
1637///
1638/// COST: one `Cache::snapshot` per TURN, on a code path that already takes one per ROUND.
1639pub(crate) struct SpecCheckpoint {
1640 snap: crate::cache::CacheSnapshot,
1641 /// Committed length at the boundary (== cache.pos there, the session invariant).
1642 pos: usize,
1643 /// Pre-output_norm hidden of row `pos - 1`.
1644 last_h: CudaSlice<f32>,
1645}
1646
1647/// PREFIX-CACHE BOUNDARY CAPTURE (lane/spec-prefix-cache, 2026-08-14): the state a spec session
1648/// records at its cold-prime split so the WORKER can publish a cross-request prefix entry —
1649/// the commit-gated-publication port (research/cache-spec-design-20260814/PORT-PLAN.md item 1).
1650/// Only the pieces that are DESTROYED by continuing the prime need copies here: the in-place
1651/// GDN conv/ssm states (via `Cache::snapshot`, same mechanism as [`SpecCheckpoint`]) and the
1652/// boundary logits. Full-attn KV rows `[0..pos)` and draft-scratch rows `[0..pos)` are
1653/// append-only for the session's lifetime (rollbacks never truncate below the prime boundary),
1654/// so the worker slices those from the live caches post-burst instead of copying at prime time.
1655pub struct SpecBoundaryCapture {
1656 pub snap: crate::cache::CacheSnapshot,
1657 /// Token boundary (== cache.pos at capture; == the worker's miss-LCP split).
1658 pub pos: usize,
1659 /// Full-vocab logits after the prefix prime — the entry's boundary logits.
1660 pub logits: Vec<f32>,
1661 /// Pre-output_norm trunk hidden of row `pos - 1` (lane/spec-on-cache-hit): the
1662 /// predecessor-pairing anchor a RESTORED spec session's first suffix-fill row reads
1663 /// (the `SpecSession::last_h` convention). Empty = unavailable (capture stays valid;
1664 /// the fill's zeros row-0 fallback covers it at a bounded acceptance cost).
1665 pub last_h: Vec<f32>,
1666 /// Per-layer latent boundary tails (lane/glm5-prefix-latent2, 2026-09-01): the
1667 /// generation-destroyed slice of each MLA/DSA layer's boundary state, captured eagerly
1668 /// so the worker's DEFERRED publication can slice the append-only planes from the live
1669 /// cache (`LatentKvLayer::snapshot_plane_at`). EMPTY on every two-plane model — the
1670 /// pre-field captures are byte-identical; a latent-bearing cache with an EMPTY vec here
1671 /// keeps the publisher's loud refusal (the fail-closed door stays shut).
1672 pub latent_tails: Vec<Option<crate::cache::LatentTailCapture>>,
1673}
1674
1675/// D2H one hidden row out of a `[T, n_embd]` prime hidden stack — the boundary anchor a
1676/// spec boundary capture carries for later restored-session fills. Failure is silent
1677/// (`turn_ckpt` convention): the capture publishes without an anchor.
1678pub(crate) fn capture_boundary_hidden(
1679 e: &Engine,
1680 h_rows: &CudaSlice<f32>,
1681 pos: usize,
1682 n_embd: usize,
1683) -> Vec<f32> {
1684 if pos == 0 || h_rows.len() < pos * n_embd {
1685 return Vec::new();
1686 }
1687 let Ok(mut row) = e.uninit(n_embd) else {
1688 return Vec::new();
1689 };
1690 if e.copy_view_into(
1691 &mut row,
1692 0,
1693 &h_rows.slice((pos - 1) * n_embd..pos * n_embd),
1694 n_embd,
1695 )
1696 .is_err()
1697 {
1698 return Vec::new();
1699 }
1700 e.dtoh(&row).unwrap_or_default()
1701}
1702
1703/// ROLLBACK DOOR for sampled BOUNDARY tokens (lane/sampled-spec-quality, 2026-08-19).
1704/// Default ON: the token a burst emits at its own boundary is drawn from the request's
1705/// sampler. `MEMRA_SPEC_SAMPLED_BOUNDARY=0` restores the pre-lane posture (an ARGMAX at
1706/// every boundary) without touching greedy, which is byte-unaffected either way.
1707pub fn spec_sampled_boundary_on() -> bool {
1708 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1709 *ON.get_or_init(|| std::env::var("MEMRA_SPEC_SAMPLED_BOUNDARY").as_deref() != Ok("0"))
1710}
1711
1712/// ROLLBACK DOOR for SESSION-SPANNING penalty history (lane/sampled-spec-quality).
1713/// Default ON: `pen_hist` is seeded from the session's committed tail, so repetition /
1714/// frequency / presence penalties see the whole stream. `MEMRA_SPEC_PEN_SESSION=0`
1715/// restores the pre-lane posture (each burst restarts the window from its own prompt
1716/// slice, i.e. from NOTHING on a continuation burst) — and with the door shut the worker
1717/// must keep refusing penalized sampled prefix-cache restores, because the restored
1718/// session's continuation burst is handed no prompt slice at all.
1719pub fn spec_pen_session_on() -> bool {
1720 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1721 *ON.get_or_init(|| std::env::var("MEMRA_SPEC_PEN_SESSION").as_deref() != Ok("0"))
1722}
1723
1724/// ROLLBACK DOOR for extended-entry publication from a RESTORED session
1725/// (lane/sampled-spec-quality, Item 3). Default ON: a converted prefix-cache hit that fed a
1726/// suffix captures its own prompt-end boundary so the NEXT turn can hit a longer prefix.
1727/// `MEMRA_SPEC_RESTORE_REPUBLISH=0` restores the pre-lane posture (a namespace learns exactly
1728/// one boundary and never advances it). Whole-entry semantics only — the boundary is the
1729/// restored session's own prompt end, so `entry_pos != fed_len` still refuses on the way in.
1730pub fn spec_restore_republish_on() -> bool {
1731 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1732 *ON.get_or_init(|| std::env::var("MEMRA_SPEC_RESTORE_REPUBLISH").as_deref() != Ok("0"))
1733}
1734
1735/// Diagnostics: name every boundary token on stderr (`MEMRA_SPEC_BOUNDARY_TRACE=1`), with
1736/// the argmax the pre-lane code would have emitted from the same row. This is how the
1737/// lane MEASURES the boundary rate and the deviation rate instead of estimating them.
1738fn spec_boundary_trace() -> bool {
1739 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1740 *ON.get_or_init(|| std::env::var("MEMRA_SPEC_BOUNDARY_TRACE").as_deref() == Ok("1"))
1741}
1742
1743/// llama-parity floor for the penalty window when the request does not ask for a bigger
1744/// one (`repeat_last_n` default). The serve API arms `penalty_last_n = PEN_WINDOW_MAX` for any
1745/// non-identity penalty, so this floor only matters to explicit small windows and to the
1746/// CLI env path.
1747const PEN_WINDOW_FLOOR: usize = 64;
1748
1749/// CEILING on the penalty window, and it is a COST bound, not a semantic preference.
1750/// `penalize_logits_f32` (cu/spec_sample.cu) dedups on device by having thread `i` scan
1751/// `hist[0..i]`, so a pass is O(n_hist²) and it runs ~3x per verify round (the q rows, the
1752/// p column, the bonus column). The serve API uses this same bound for every non-identity
1753/// penalty so host/plain, sparse-device, and speculative sampling cannot change logits on
1754/// admission demotion. An uncapped 128k-token history would put ~1.7e10
1755/// comparisons per pass, tens of ms per round, i.e. penalties would silently destroy decode
1756/// throughput on exactly the long-context requests that most want them. 8192 keeps a pass
1757/// at ~7e7 comparisons (tens of microseconds) while still being **128x wider than the
1758/// pre-lane effective window** (64 prompt-tail tokens + whatever the current burst had
1759/// generated). A request that genuinely needs a window beyond this wants host-side dedup +
1760/// counts through a new kernel signature — a follow-up lane, named here rather than hidden.
1761/// `pub` since lane/dspark-penalized-sampled-20260821: the dspark route's accept walk and
1762/// the dspark_sample_gate binary trim their uploads with the SAME cap — a second constant
1763/// is a second thing to drift.
1764pub const PEN_WINDOW_MAX: usize = 8192;
1765
1766/// Seed a penalty window over the SESSION, not the burst (lane/sampled-spec-quality,
1767/// Item 2). The window is the last `max(penalty_last_n, 64)` tokens of
1768/// `session_committed ++ burst_prompt` — for a cold turn-1 burst (`session_committed`
1769/// empty, default `penalty_last_n`) that is byte-identically the pre-lane
1770/// `prompt.iter().rev().take(64).rev()`; for a continuation burst it is the stream the
1771/// client actually asked us to penalize, where the pre-lane code had NOTHING.
1772/// `pub` since lane/dspark-penalized-sampled-20260821: the dspark route seeds its session
1773/// window through the SAME function (one definition of "the window" across both spec
1774/// routes and the gate binary's trunk-only reference arm).
1775pub fn pen_window_seed(
1776 session_committed: &[u32],
1777 burst_prompt: &[u32],
1778 penalty_last_n: usize,
1779) -> Vec<u32> {
1780 let win = penalty_last_n.clamp(PEN_WINDOW_FLOOR, PEN_WINDOW_MAX);
1781 let take_prompt = burst_prompt.len().min(win);
1782 let take_sess = (win - take_prompt).min(session_committed.len());
1783 let mut hist = Vec::with_capacity(take_sess + take_prompt);
1784 hist.extend_from_slice(&session_committed[session_committed.len() - take_sess..]);
1785 hist.extend_from_slice(&burst_prompt[burst_prompt.len() - take_prompt..]);
1786 hist
1787}
1788
1789/// Draw a BOUNDARY token from the target distribution the request asked for
1790/// (lane/sampled-spec-quality, Item 1) — the fix for "sampled spec emits an ARGMAX token at
1791/// every burst boundary".
1792///
1793/// WHY THIS EXISTS. A spec burst's first emitted token is not produced by the accept walk:
1794/// it comes off a logits row that already exists (the prime's last row on a cold burst; the
1795/// row after the last committed token on a continuation burst; the prefix-cache entry's
1796/// boundary row on a restored one). Pre-lane that token was `argmax` in BOTH sampling
1797/// regimes, so a sampled stream took a greedy token once per burst — measured, not
1798/// estimated, in research/spec-cache-20260818/SAMPLED-QUALITY.md. At temperature > 0 the
1799/// customer asked for a sampled token, so this draws one.
1800///
1801/// THE PROGRAM IS THE FULL-ACCEPT BONUS'S PROGRAM, deliberately: penalize the row (over the
1802/// session's window), take this row's OWN filter stats (the sampfix-20260805 law — stats
1803/// from a neighbour row mis-scale every `e0` and can wipe the row to token 0), gumbel-perturb
1804/// with the session's Philox stream at `*sctr`, argmax the perturbed row. Reusing the bonus's
1805/// composition means `sample_check`'s distributional oracle covers this draw too, and the
1806/// boundary token is drawn from the same filtered/penalized `p` the accept walk targets.
1807///
1808/// THE STREAM IS THE SESSION'S, NOT A FRESH ONE. `sctr` is the caller's live counter and is
1809/// advanced by exactly one, so a boundary draw consumes the next value in the same Philox
1810/// stream the accept walk uses — never a second, independently seeded stream (which would be
1811/// a new distributional bug: two streams from one seed correlate wherever their counters
1812/// collide). That also makes a restored session's boundary draw at `sctr == 0` bit-identical
1813/// to the cold session's own first draw from the same logits row, which is what preserves the
1814/// sampled-hit lane's per-seed hit==cold byte identity.
1815#[allow(clippy::too_many_arguments)]
1816pub fn sample_boundary_token_dev(
1817 e: &Engine,
1818 logits: &CudaSlice<f32>,
1819 n_vocab: usize,
1820 sp: &SpecSampling,
1821 pen_hist: &[u32],
1822 sctr: &mut u32,
1823 site: &str,
1824) -> Result<u32, Box<dyn std::error::Error>> {
1825 debug_assert!(
1826 sp.temp > 0.0,
1827 "boundary sampling is the sampled regime only"
1828 );
1829 // Own copy: penalize_logits mutates in place and the caller's row is live state
1830 // (prime_logits back the constrained recompute; last_col_logits backs round 0's accept).
1831 let mut col = e.zeros(n_vocab)?;
1832 e.copy_into(&mut col, 0, logits, n_vocab)?;
1833 let pen_on = sp.penalty_last_n > 0
1834 && (sp.penalty_repeat != 1.0 || sp.penalty_freq != 0.0 || sp.penalty_present != 0.0);
1835 if pen_on && !pen_hist.is_empty() {
1836 // window trim mirrors the round loop's own upload (`pen_hist[w0..]`), cap included.
1837 let w0 = pen_hist
1838 .len()
1839 .saturating_sub(sp.penalty_last_n.min(PEN_WINDOW_MAX));
1840 let hist = &pen_hist[w0..];
1841 let hd = e.htod_u32_v(hist)?;
1842 e.penalize_logits(
1843 &mut col,
1844 &hd,
1845 hist.len(),
1846 sp.penalty_repeat,
1847 sp.penalty_freq,
1848 sp.penalty_present,
1849 n_vocab,
1850 )?;
1851 }
1852 let rows0 = e.htod_i32(&[0])?;
1853 let (mut th_d, mut z_d, mut mx_d) = (e.zeros(1)?, e.zeros(1)?, e.zeros(1)?);
1854 e.filter_stats(
1855 &col, n_vocab, &rows0, &mut th_d, &mut z_d, &mut mx_d, n_vocab, 1, sp.temp, sp.top_k,
1856 sp.top_p, sp.min_p,
1857 )?;
1858 let (th, mx) = (e.dtoh(&th_d)?[0], e.dtoh(&mx_d)?[0]);
1859 let mut perturb = e.zeros(n_vocab)?;
1860 e.gumbel_perturb_filtered(&col, &mut perturb, n_vocab, sp.seed, *sctr, sp.temp, mx, th)?;
1861 *sctr = sctr.wrapping_add(1);
1862 let td = e.argmax_token_device(&perturb, n_vocab)?;
1863 let tok = guard_vocab_token(
1864 e.dtoh_u32_one(&td)?,
1865 n_vocab,
1866 &format!("sampled boundary token (site={site})"),
1867 )?;
1868 if spec_boundary_trace() {
1869 // the pre-lane token, from the SAME row, so the deviation rate is measurable.
1870 let raw = e.argmax_token_device(logits, n_vocab)?;
1871 let greedy = e.dtoh_u32_one(&raw)?;
1872 eprintln!(
1873 "[spec-boundary] site={site} sampled={tok} argmax={greedy} \
1874 deviates={} temp={} sctr={}",
1875 (tok != greedy) as u8,
1876 sp.temp,
1877 sctr.wrapping_sub(1),
1878 );
1879 }
1880 Ok(tok)
1881}
1882
1883/// Host-row twin of [`sample_boundary_token_dev`] (the prime / feed / entry rows arrive as
1884/// host `Vec<f32>`).
1885#[allow(clippy::too_many_arguments)]
1886pub fn sample_boundary_token(
1887 e: &Engine,
1888 logits: &[f32],
1889 sp: &SpecSampling,
1890 pen_hist: &[u32],
1891 sctr: &mut u32,
1892 site: &str,
1893) -> Result<u32, Box<dyn std::error::Error>> {
1894 let n_vocab = logits.len();
1895 let d = e.htod(logits)?;
1896 sample_boundary_token_dev(e, &d, n_vocab, sp, pen_hist, sctr, site)
1897}
1898
1899struct SpecPipeTraceClock {
1900 pair: usize,
1901 started: std::time::Instant,
1902}
1903
1904#[derive(Clone)]
1905struct SpecPipeTraceCtx {
1906 clock: std::sync::Arc<SpecPipeTraceClock>,
1907 round: usize,
1908 lane: usize,
1909}
1910
1911struct SpecPipeTraceMarker {
1912 trace: SpecPipeTraceCtx,
1913 phase: &'static str,
1914 edge: &'static str,
1915 slot: Option<usize>,
1916}
1917
1918unsafe extern "C" fn spec_pipe_trace_marker(raw: *mut std::ffi::c_void) {
1919 let marker = unsafe { Box::from_raw(raw.cast::<SpecPipeTraceMarker>()) };
1920 let lane = if marker.trace.lane == 0 { "A" } else { "B" };
1921 let slot = marker
1922 .slot
1923 .map(|v| v.to_string())
1924 .unwrap_or_else(|| "-".into());
1925 let t_ms = marker.trace.clock.started.elapsed().as_secs_f64() * 1e3;
1926 use std::io::Write as _;
1927 let stderr = std::io::stderr();
1928 let mut stderr = stderr.lock();
1929 let _ = writeln!(
1930 stderr,
1931 "[spec-pipe-timeline] pair={} round={} lane={lane} phase={} edge={} \
1932 slot={slot} t_ms={t_ms:.3}",
1933 marker.trace.clock.pair, marker.trace.round, marker.phase, marker.edge,
1934 );
1935}
1936
1937fn enqueue_spec_pipe_trace_marker(
1938 stream: &cudarc::driver::CudaStream,
1939 trace: Option<&SpecPipeTraceCtx>,
1940 phase: &'static str,
1941 edge: &'static str,
1942 slot: Option<usize>,
1943) -> Result<(), Box<dyn std::error::Error>> {
1944 let Some(trace) = trace else {
1945 return Ok(());
1946 };
1947 let marker = Box::new(SpecPipeTraceMarker {
1948 trace: trace.clone(),
1949 phase,
1950 edge,
1951 slot,
1952 });
1953 let raw = Box::into_raw(marker);
1954 let result = unsafe {
1955 cudarc::driver::result::stream::launch_host_function(
1956 stream.cu_stream(),
1957 spec_pipe_trace_marker,
1958 raw.cast(),
1959 )
1960 };
1961 if let Err(err) = result {
1962 unsafe {
1963 drop(Box::from_raw(raw));
1964 }
1965 return Err(err.into());
1966 }
1967 Ok(())
1968}
1969
1970#[derive(Default)]
1971struct SpecPipeProgress {
1972 setup_done: [bool; 2],
1973 draft_done: [usize; 2],
1974 stage0_done: [usize; 2],
1975 verify_done: [usize; 2],
1976 accept_done: [usize; 2],
1977 finished: [bool; 2],
1978 aborted: bool,
1979}
1980
1981/// Host-side issue coordinator for the reduced two-session speculative pipeline. Each session
1982/// keeps its existing call stack and round locals; this object only orders phase entry. The
1983/// primary mutex spans whole draft/accept/tail issue regions so Engine's single-stream scratch
1984/// cannot be interleaved by the two host threads.
1985struct SpecPipeSync {
1986 progress: std::sync::Mutex<SpecPipeProgress>,
1987 changed: std::sync::Condvar,
1988 primary: std::sync::Mutex<()>,
1989 trace: Option<std::sync::Arc<SpecPipeTraceClock>>,
1990}
1991
1992impl SpecPipeSync {
1993 fn new() -> Self {
1994 static TRACE_PAIR: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
1995 let trace = (std::env::var("MEMRA_SPEC_PIPE_TRACE").as_deref() == Ok("1")).then(|| {
1996 std::sync::Arc::new(SpecPipeTraceClock {
1997 pair: TRACE_PAIR.fetch_add(1, std::sync::atomic::Ordering::Relaxed) + 1,
1998 started: std::time::Instant::now(),
1999 })
2000 });
2001 Self {
2002 progress: std::sync::Mutex::new(SpecPipeProgress::default()),
2003 changed: std::sync::Condvar::new(),
2004 primary: std::sync::Mutex::new(()),
2005 trace,
2006 }
2007 }
2008}
2009
2010#[derive(Clone)]
2011struct SpecPipeLane {
2012 sync: std::sync::Arc<SpecPipeSync>,
2013 lane: usize,
2014 rt: &'static crate::pp::PpNRt,
2015 walk_permit: crate::pp::PpWalkPermit,
2016}
2017
2018struct SpecPipePrimaryGuard<'a> {
2019 _primary: std::sync::MutexGuard<'a, ()>,
2020 _walk: crate::pp::PpWalkBorrowGuard,
2021}
2022
2023impl SpecPipeLane {
2024 fn peer(&self) -> usize {
2025 1 - self.lane
2026 }
2027
2028 fn aborted() -> Box<dyn std::error::Error> {
2029 "paired speculative peer aborted".into()
2030 }
2031
2032 fn trace(&self, round: usize) -> Option<SpecPipeTraceCtx> {
2033 self.sync.trace.as_ref().map(|clock| SpecPipeTraceCtx {
2034 clock: clock.clone(),
2035 round,
2036 lane: self.lane,
2037 })
2038 }
2039
2040 fn setup_begin(&self) -> Result<crate::pp::PpWalkBorrowGuard, Box<dyn std::error::Error>> {
2041 let mut p = self.sync.progress.lock().unwrap();
2042 while !p.aborted && self.lane == 1 && !p.setup_done[0] && !p.finished[0] {
2043 p = self.sync.changed.wait(p).unwrap();
2044 }
2045 if p.aborted {
2046 Err(Self::aborted())
2047 } else {
2048 drop(p);
2049 self.rt.borrow_walk(&self.walk_permit, "spec_pipe/setup")
2050 }
2051 }
2052
2053 fn setup_end(&self) {
2054 let mut p = self.sync.progress.lock().unwrap();
2055 p.setup_done[self.lane] = true;
2056 self.sync.changed.notify_all();
2057 }
2058
2059 fn draft_begin(
2060 &self,
2061 round: usize,
2062 ) -> Result<SpecPipePrimaryGuard<'_>, Box<dyn std::error::Error>> {
2063 let peer = self.peer();
2064 let mut p = self.sync.progress.lock().unwrap();
2065 loop {
2066 if p.aborted {
2067 return Err(Self::aborted());
2068 }
2069 let setup_ready =
2070 (p.setup_done[0] || p.finished[0]) && (p.setup_done[1] || p.finished[1]);
2071 let prior_ready = p.accept_done[self.lane] >= round
2072 && (p.accept_done[peer] >= round || p.finished[peer]);
2073 let turn_ready = if self.lane == 0 {
2074 true
2075 } else {
2076 p.draft_done[0] > round || p.finished[0]
2077 };
2078 if setup_ready && prior_ready && turn_ready {
2079 break;
2080 }
2081 p = self.sync.changed.wait(p).unwrap();
2082 }
2083 drop(p);
2084 let primary = self.sync.primary.lock().unwrap();
2085 let walk = self.rt.borrow_walk(&self.walk_permit, "spec_pipe/draft")?;
2086 Ok(SpecPipePrimaryGuard {
2087 _primary: primary,
2088 _walk: walk,
2089 })
2090 }
2091
2092 fn draft_end(&self, round: usize) {
2093 let mut p = self.sync.progress.lock().unwrap();
2094 p.draft_done[self.lane] = round + 1;
2095 self.sync.changed.notify_all();
2096 }
2097
2098 /// Admit stage 0 and return whether this lane owns the interval's one reverse fence.
2099 /// Lane B releases as soon as lane A has issued its boundary TX, not after A's full body.
2100 fn stage0_begin(&self, round: usize) -> Result<bool, Box<dyn std::error::Error>> {
2101 let peer = self.peer();
2102 let mut p = self.sync.progress.lock().unwrap();
2103 loop {
2104 if p.aborted {
2105 return Err(Self::aborted());
2106 }
2107 let ready = if self.lane == 0 {
2108 p.draft_done[0] > round && (p.draft_done[1] > round || p.finished[1])
2109 } else {
2110 p.draft_done[1] > round && (p.stage0_done[0] > round || p.finished[0])
2111 };
2112 if ready {
2113 return Ok(self.lane == 0 || p.finished[peer]);
2114 }
2115 p = self.sync.changed.wait(p).unwrap();
2116 }
2117 }
2118
2119 fn stage0_end(&self, round: usize) {
2120 let mut p = self.sync.progress.lock().unwrap();
2121 p.stage0_done[self.lane] = round + 1;
2122 self.sync.changed.notify_all();
2123 }
2124
2125 /// Stage 1 is single-owner per engine. A proceeds immediately after its own ticket; B waits
2126 /// for A's full stage1/head issue so only A.S1 and B.S0 can overlap.
2127 fn stage1_begin(&self, round: usize) -> Result<(), Box<dyn std::error::Error>> {
2128 let mut p = self.sync.progress.lock().unwrap();
2129 while !p.aborted
2130 && !(p.stage0_done[self.lane] > round
2131 && (self.lane == 0 || p.verify_done[0] > round || p.finished[0]))
2132 {
2133 p = self.sync.changed.wait(p).unwrap();
2134 }
2135 if p.aborted {
2136 Err(Self::aborted())
2137 } else {
2138 Ok(())
2139 }
2140 }
2141
2142 fn verify_end(&self, round: usize) {
2143 let mut p = self.sync.progress.lock().unwrap();
2144 p.verify_done[self.lane] = round + 1;
2145 self.sync.changed.notify_all();
2146 }
2147
2148 fn accept_begin(
2149 &self,
2150 round: usize,
2151 ) -> Result<SpecPipePrimaryGuard<'_>, Box<dyn std::error::Error>> {
2152 let mut p = self.sync.progress.lock().unwrap();
2153 loop {
2154 if p.aborted {
2155 return Err(Self::aborted());
2156 }
2157 let ready = if self.lane == 0 {
2158 p.verify_done[0] > round && (p.verify_done[1] > round || p.finished[1])
2159 } else {
2160 p.verify_done[1] > round && (p.accept_done[0] > round || p.finished[0])
2161 };
2162 if ready {
2163 break;
2164 }
2165 p = self.sync.changed.wait(p).unwrap();
2166 }
2167 drop(p);
2168 let primary = self.sync.primary.lock().unwrap();
2169 let walk = self.rt.borrow_walk(&self.walk_permit, "spec_pipe/accept")?;
2170 Ok(SpecPipePrimaryGuard {
2171 _primary: primary,
2172 _walk: walk,
2173 })
2174 }
2175
2176 fn accept_end(&self, round: usize) {
2177 let mut p = self.sync.progress.lock().unwrap();
2178 p.accept_done[self.lane] = round + 1;
2179 self.sync.changed.notify_all();
2180 }
2181
2182 fn primary(&self) -> Result<SpecPipePrimaryGuard<'_>, Box<dyn std::error::Error>> {
2183 let primary = self.sync.primary.lock().unwrap();
2184 let walk = self.rt.borrow_walk(&self.walk_permit, "spec_pipe/tail")?;
2185 Ok(SpecPipePrimaryGuard {
2186 _primary: primary,
2187 _walk: walk,
2188 })
2189 }
2190
2191 fn coordinated_walk(&self) -> Result<crate::pp::PpWalkBorrowGuard, Box<dyn std::error::Error>> {
2192 self.rt
2193 .borrow_walk(&self.walk_permit, "spec_pipe/coordinated_verify")
2194 }
2195
2196 fn finish(&self, failed: bool) {
2197 let mut p = self.sync.progress.lock().unwrap();
2198 p.finished[self.lane] = true;
2199 p.aborted |= failed;
2200 self.sync.changed.notify_all();
2201 }
2202}
2203
2204struct SpecPipeFinish<'a> {
2205 lane: &'a SpecPipeLane,
2206 closed: bool,
2207}
2208
2209impl<'a> SpecPipeFinish<'a> {
2210 fn new(lane: &'a SpecPipeLane) -> Self {
2211 Self {
2212 lane,
2213 closed: false,
2214 }
2215 }
2216
2217 fn close(&mut self, failed: bool) {
2218 self.lane.finish(failed);
2219 self.closed = true;
2220 }
2221}
2222
2223impl Drop for SpecPipeFinish<'_> {
2224 fn drop(&mut self) {
2225 if !self.closed {
2226 self.lane.finish(true);
2227 }
2228 }
2229}
2230
2231/// Scoped transfer of one exclusively-borrowed session to the second host issue thread.
2232/// `CudaGraph` is not marked Send by cudarc because its raw driver handles carry no automatic
2233/// trait. CUDA driver graph handles are context-scoped rather than OS-thread-affine; the caller
2234/// binds that context before touching the session, joins before returning, and never aliases the
2235/// pointer. Keep this exception local to the experimental pair call instead of marking the public
2236/// session type Send.
2237struct SpecPipeSessionPtr(*mut SpecSession);
2238
2239unsafe impl Send for SpecPipeSessionPtr {}
2240
2241impl SpecPipeSessionPtr {
2242 unsafe fn get_mut(&mut self) -> &mut SpecSession {
2243 unsafe { &mut *self.0 }
2244 }
2245}
2246
2247/// Per-session persistent draft-graph context: the captured CUDA graph(s) plus the device
2248/// buffers whose POINTERS the capture bakes. Reuse legality: the greedy capture bakes only
2249/// session-stable pointers (the session's own MtpScratch KV — allocated once, never realloc'd;
2250/// the model's resident embedding; the process-wide OnceLock p_min) and the g_* buffers held
2251/// HERE — so one capture serves the session's whole lifetime. The sampled capture additionally
2252/// bakes (seed, temp) as capture-time constants and needs k q-slots — keyed by `s_key`, dropped
2253/// and recaptured when a pool-resumed request changes them. `*_failed` memoizes a failed capture
2254/// so the eager fallback doesn't pay a doomed capture attempt every burst.
2255/// Capture identity of the parked SAMPLED draft graph (`DraftGraphCtx::graph_s`).
2256///
2257/// EXACTNESS, not perf (lane/graph-s-key-exactness-20260819; receipts
2258/// `research/spec-cache-20260818/GRAPH-S-KEY.md`). Two classes of field live here, both
2259/// load-bearing:
2260///
2261/// - **Baked constants.** `seed` and `temp` are capture-time constants INSIDE the graph and `k`
2262/// sizes the q slots its replays write. A resumed request changing any of them must recapture.
2263/// This is all the key used to carry.
2264/// - **Regime fields.** `top_k`/`top_p`/`min_p`/`pen_on` are not baked, but they decide whether
2265/// the captured graph is a legal draft chain AT ALL. The in-graph draw is one gumbel-max over
2266/// the RAW softmax (`gumbel_perturb_ctr`, unfiltered by construction), while the verify builds
2267/// the accept test's `q` from `filter_stats(q_slots, top_k, top_p, min_p)`. If those disagree
2268/// the accept test evaluates a distribution the draft was never sampled from: a draft token
2269/// below the filter threshold gathers `q = 0` (`softmax_gather_filtered_f32`,
2270/// `cu/spec_sample.cu`) and `u * 0 < p` accepts it UNCONDITIONALLY.
2271///
2272/// Omitting the regime fields was reachable — not through the prefix-cache spec restore (that
2273/// path is greedy-only, `memra-server` `spec_restore_convertible`), but through WHOLE-SESSION
2274/// spec reuse: a parked `SpecSession` carries this `DraftGraphCtx`, and the pool-resume probe
2275/// applies no sampler predicate at all. Turn 1 pure-temp parks a graph; turn 2 of the same
2276/// conversation, same explicit seed and temperature, adds `top_p`/`top_k` and inherits it.
2277#[derive(Clone, Copy, PartialEq, Eq, Debug)]
2278pub(crate) struct SampledGraphKey {
2279 seed: u64,
2280 temp_bits: u32,
2281 k: usize,
2282 top_k: i32,
2283 top_p_bits: u32,
2284 min_p_bits: u32,
2285 pen_on: bool,
2286}
2287
2288impl SampledGraphKey {
2289 pub(crate) fn new(
2290 seed: u64,
2291 temp: f32,
2292 k: usize,
2293 top_k: i32,
2294 top_p: f32,
2295 min_p: f32,
2296 pen_on: bool,
2297 ) -> Self {
2298 SampledGraphKey {
2299 seed,
2300 temp_bits: temp.to_bits(),
2301 k,
2302 top_k,
2303 top_p_bits: top_p.to_bits(),
2304 min_p_bits: min_p.to_bits(),
2305 pen_on,
2306 }
2307 }
2308
2309 /// The one regime the PURE-TEMP in-graph sampled chain may stand in for the eager one:
2310 /// nothing but temperature shapes `q`. Computed FROM THE KEY so the capture guard, the
2311 /// launch guard and the key can never drift apart (they were three separate expressions
2312 /// before this lane, and the launch site simply forgot to ask).
2313 pub(crate) fn pure_temp(&self) -> bool {
2314 self.top_k == 0
2315 && f32::from_bits(self.top_p_bits) >= 1.0
2316 && f32::from_bits(self.min_p_bits) <= 0.0
2317 && !self.pen_on
2318 }
2319
2320 /// Truncation filters active — the capture body needs the IN-GRAPH filter nodes
2321 /// (`filter_stats` + `gumbel_perturb_filtered_ctr`) so the draft draws from the same
2322 /// filtered distribution the accept test reconstructs. Meaningful only when
2323 /// `graph_capturable`; penalties never reach a capture body.
2324 pub(crate) fn filtered(&self) -> bool {
2325 !self.pure_temp()
2326 }
2327
2328 /// May the sampled draft graph be CAPTURED (and a parked one LAUNCHED) for this regime?
2329 /// Pure-temp always; filtered regimes when the filtered-capture door is on
2330 /// (lane/step37-draft-graph-serving-20260830); penalties never — the per-round history
2331 /// cannot be baked into a graph, and composing a raw-softmax (or stale-history) draw
2332 /// with a penalized accept test is the unconditional-accept exactness bug. Computed FROM
2333 /// THE KEY for the same no-drift reason as `pure_temp`.
2334 pub(crate) fn graph_capturable(&self) -> bool {
2335 !self.pen_on && (self.pure_temp() || spec_graph_filtered_on())
2336 }
2337}
2338
2339/// Per-head captured graphs for the MULTI-HEAD MTP draft chain (step-modulo prefix-replay,
2340/// lane/step37-draft-graph-serving-20260830). The chain POLICY — which head serves step j,
2341/// how long the replayed prefix is, which stored seed feeds row r — stays HOST-SIDE in the
2342/// launch loop, exactly `mtp_chain_forward_dev`'s order; the graphs capture ONE head-row
2343/// forward each, on the head's OWN scratch plane:
2344/// - `interior[i]`: head i, `with_head=false` — KV append + carrier only. Interior rows'
2345/// logits are dead in the eager chain too (`mtp_chain_forward_dev` keeps only the last
2346/// row), so skipping the head matmul changes no consumed byte and removes the eager
2347/// chain's per-replay-row full-vocab matmul.
2348/// - `last[i]`: head i, `with_head=true` + the mode's tail (greedy argmax, or the sampled
2349/// gumbel draw — filtered in-graph when the request carries filters).
2350///
2351/// One `DraftChainGraphs` per MODE (greedy vs sampled), owning its keeper: dropping the
2352/// sampled chain on an s_key change never invalidates the greedy one.
2353struct DraftChainGraphs {
2354 interior: Vec<cudarc::driver::CudaGraph>,
2355 last: Vec<cudarc::driver::CudaGraph>,
2356 /// Never read: exists to OWN the captured graphs' backing buffers for as long as the
2357 /// graphs replay (the capture-retain law; same class as `DsparkSegGraph::_keeper`).
2358 _keeper: Vec<Box<dyn std::any::Any + Send>>,
2359}
2360
2361/// Sampled-tail capture pack for `mtp_head_forward_cap`: the persistent buffers and baked
2362/// constants of the in-graph categorical draw. `filt: None` = the PURE-TEMP body (gumbel
2363/// over the raw softmax), byte-identical to the pre-lane capture; `Some` adds the in-graph
2364/// truncation filter (`filter_stats` + `gumbel_perturb_filtered_ctr`) so the draft draws
2365/// from the same filtered distribution the accept test reconstructs
2366/// (lane/step37-draft-graph-serving-20260830).
2367struct SampledCapArgs<'a> {
2368 ctr: &'a mut CudaSlice<u32>,
2369 perturb: &'a mut CudaSlice<f32>,
2370 q_out: &'a mut CudaSlice<f32>,
2371 seed: u64,
2372 temp: f32,
2373 filt: Option<SampledCapFilter<'a>>,
2374}
2375
2376/// In-graph truncation-filter nodes: the stat slots `filter_stats` fills and the perturb
2377/// reads, plus the filter constants baked into the capture (they live in `s_key`, so a
2378/// request whose filters differ drops the parked graph before this ever goes stale).
2379struct SampledCapFilter<'a> {
2380 rows0: &'a CudaSlice<i32>,
2381 th: &'a mut CudaSlice<f32>,
2382 z: &'a mut CudaSlice<f32>,
2383 mx: &'a mut CudaSlice<f32>,
2384 top_k: i32,
2385 top_p: f32,
2386 min_p: f32,
2387}
2388
2389pub(crate) struct DraftGraphCtx {
2390 g_tok: CudaSlice<u32>,
2391 g_pos: CudaSlice<i32>,
2392 g_seed: CudaSlice<f32>,
2393 g_p: CudaSlice<f32>,
2394 g_ctr: CudaSlice<u32>,
2395 g_q: CudaSlice<f32>,
2396 g_perturb: CudaSlice<f32>,
2397 /// IN-GRAPH filter-stat slots (filtered sampled capture): `filter_stats` writes
2398 /// (th, z, mx) here inside the graph; `gumbel_perturb_filtered_ctr` reads (mx, th) from
2399 /// the same slots. Persistent so the baked pointers survive replays. `g_rows0` is the
2400 /// constant row-index-0 the single-row `filter_stats` launch reads (a captured memcpy
2401 /// source must not be a host temporary).
2402 g_rows0: CudaSlice<i32>,
2403 g_th: CudaSlice<f32>,
2404 g_z: CudaSlice<f32>,
2405 g_mx: CudaSlice<f32>,
2406 q_slots: Vec<CudaSlice<f32>>,
2407 /// DRAFT-SIDE GRAMMAR MASK (lane/draft-mask): packed allowed-set words over the DRAFT
2408 /// head's vocab, at a STABLE address so the captured draft graph's mask node reads the
2409 /// per-position contents the host re-uploads before each replay (the graph-promote
2410 /// pattern from decode.rs). Empty unless the session drafts under a grammar.
2411 g_dmask: CudaSlice<u32>,
2412 /// was `graph` captured WITH the mask node? A parked graph of the wrong shape is dropped.
2413 /// Covers the multi-head `chain` too (single-head and chain are mutually exclusive for a
2414 /// given model, so one flag serves whichever is active).
2415 graph_masked: bool,
2416 graph: Option<cudarc::driver::CudaGraph>,
2417 graph_s: Option<cudarc::driver::CudaGraph>,
2418 /// Multi-head chain graphs (see [`DraftChainGraphs`]): greedy and sampled chains, the
2419 /// chain twins of `graph` / `graph_s`. `chain_s`'s capture identity is `s_key` (shared
2420 /// with `graph_s` — a session is either single-head or chain, never both), and it obeys
2421 /// the same drop rules (key mismatch, penalty regime, mask-shape change).
2422 chain: Option<DraftChainGraphs>,
2423 chain_s: Option<DraftChainGraphs>,
2424 /// Failed-capture memoization for both graphs — LOUD on flip, cleared on pool resume
2425 /// (audit Q2, the TRT #16072 silent-permanent-coverage-loss class).
2426 failed: DraftGraphFallback,
2427 /// Capture identity of `graph_s` — see [`SampledGraphKey`]. `None` iff no sampled graph is
2428 /// parked; a request whose key differs drops the parked graph (and its q slots/keeper).
2429 s_key: Option<SampledGraphKey>,
2430 /// CAPTURE-RETAIN keepers (#68 root cause, 2026-08-04): the warmup-run transients whose
2431 /// pool addresses the captured graph(s) bake. Without these, the transients return to the
2432 /// pool at capture-body exit and later work (burst-boundary prime/fill/commit passes, or a
2433 /// co-served session in the worker) reuses those addresses — the persisted graph's replay
2434 /// then reads/writes live unrelated buffers (exactness corruption, first seen as the ST
2435 /// serve-spec 4B graph-arm corruption; one-shot CLI calls never re-shuffled the pool, which
2436 /// is why run-spec K=1..8 passed on the same checkpoint). Same fix class as
2437 /// capture_graph_retained's gemma/decode.rs sites — hold as long as the graph replays.
2438 keeper: Vec<Box<dyn std::any::Any + Send>>,
2439 keeper_s: Vec<Box<dyn std::any::Any + Send>>,
2440}
2441
2442/// Failed-capture memoization for the two draft graphs (audit Q2, 2026-08-05 — the
2443/// TRT #16072 trap class: pressure-triggered, silent, long-lived coverage loss).
2444///
2445/// Three contracts:
2446/// - LOUD FLIP: `mark_*` returns the warn line exactly on the false→true transition
2447/// (returned, not printed, so the once-per-flip contract is unit-testable); the caller
2448/// `eprintln!`s it UNCONDITIONALLY — a dropped draft graph is never silent. Re-marking
2449/// an already-failed graph returns None (the per-burst memoization that keeps the eager
2450/// fallback from paying a doomed capture attempt every burst).
2451/// - RESET ON RESUME: `reset_on_resume` clears both flags — a parked session resumed by a
2452/// NEW request gets one fresh capture chance instead of carrying a transient-pressure
2453/// failure for the pool's whole lifetime. Returns the note line only when a flag was
2454/// actually set (quiet on the common clean-resume path).
2455/// - Shape-change clears (`clear_*`) stay silent, exactly as before: they precede a fresh
2456/// capture attempt whose own failure would re-flip loudly.
2457#[derive(Default)]
2458pub(crate) struct DraftGraphFallback {
2459 greedy: bool,
2460 sampled: bool,
2461}
2462impl DraftGraphFallback {
2463 fn mark_greedy(&mut self, reason: &str) -> Option<String> {
2464 if self.greedy {
2465 return None;
2466 }
2467 self.greedy = true;
2468 Some(format!(
2469 "[spec] WARN: draft-graph capture failed ({reason}); eager fallback until session resume"
2470 ))
2471 }
2472 fn mark_sampled(&mut self, reason: &str) -> Option<String> {
2473 if self.sampled {
2474 return None;
2475 }
2476 self.sampled = true;
2477 Some(format!(
2478 "[spec] WARN: sampled draft-graph capture failed ({reason}); eager fallback until session resume"
2479 ))
2480 }
2481 fn greedy_failed(&self) -> bool {
2482 self.greedy
2483 }
2484 fn sampled_failed(&self) -> bool {
2485 self.sampled
2486 }
2487 fn clear_greedy(&mut self) {
2488 self.greedy = false;
2489 }
2490 fn clear_sampled(&mut self) {
2491 self.sampled = false;
2492 }
2493 /// Pool-resume reset: both graphs get a fresh capture chance. Some(note) iff any flag
2494 /// was set (so clean resumes stay quiet).
2495 pub(crate) fn reset_on_resume(&mut self) -> Option<String> {
2496 if !self.greedy && !self.sampled {
2497 return None;
2498 }
2499 let which = match (self.greedy, self.sampled) {
2500 (true, true) => "greedy+sampled",
2501 (true, false) => "greedy",
2502 _ => "sampled",
2503 };
2504 self.greedy = false;
2505 self.sampled = false;
2506 Some(format!(
2507 "[spec] draft-graph fallback reset on session resume ({which}); recapture eligible"
2508 ))
2509 }
2510}
2511
2512impl DraftGraphCtx {
2513 fn new(e: &Engine, n_embd: usize, qlen: usize) -> Result<Self, Box<dyn std::error::Error>> {
2514 Ok(DraftGraphCtx {
2515 g_tok: e.alloc_u32_zeroed(1)?,
2516 g_pos: e.htod_i32(&[0])?,
2517 g_seed: e.zeros(n_embd)?,
2518 g_p: e.zeros(1)?,
2519 g_ctr: e.alloc_u32_zeroed(1)?,
2520 g_q: e.zeros(qlen)?,
2521 g_perturb: e.zeros(qlen)?,
2522 g_rows0: e.htod_i32(&[0])?,
2523 g_th: e.zeros(1)?,
2524 g_z: e.zeros(1)?,
2525 g_mx: e.zeros(1)?,
2526 q_slots: Vec::new(),
2527 g_dmask: e.alloc_u32_zeroed(1)?,
2528 graph_masked: false,
2529 graph: None,
2530 graph_s: None,
2531 chain: None,
2532 chain_s: None,
2533 failed: DraftGraphFallback::default(),
2534 s_key: None,
2535 keeper: Vec::new(),
2536 keeper_s: Vec::new(),
2537 })
2538 }
2539}
2540
2541pub(crate) struct MtpScratch {
2542 kv: KvLayer,
2543 /// Logical row capacity. On the graph/DC draft path it also doubles as fa_decode_dc's
2544 /// bucket_max: n_splits is sized from it ONCE, so the graph captured at round 0 stays valid
2545 /// for every later t_kv. Step35 refuses that path and may back this logical extent with the
2546 /// smaller host-indexed SWA ring instead.
2547 cap: usize,
2548 extra: Vec<MtpScratchPlane>,
2549}
2550
2551struct MtpScratchPlane {
2552 kv: KvLayer,
2553 cap: usize,
2554}
2555
2556fn mtp_scratch_layout(
2557 cfg: &memra_gguf::config::ModelConfig,
2558 geom: Option<&crate::hybrid::DraftGeom>,
2559) -> (usize, usize, usize, usize) {
2560 // Student draft heads carry fewer KV heads (head_dim unchanged) -> smaller scratch rows.
2561 let n_head_kv = geom.map(|g| g.n_head_kv).unwrap_or(cfg.n_head_kv as usize);
2562 let head_dim_k = cfg.head_dim_k as usize;
2563 let head_dim_v = cfg.head_dim_v as usize;
2564 assert!(
2565 head_dim_k.is_multiple_of(32) && head_dim_v.is_multiple_of(32),
2566 "KVQUANT requires head_dim%32==0 (MTP scratch)"
2567 );
2568 let kv_dim_k = head_dim_k * n_head_kv;
2569 let kv_dim_v = head_dim_v * n_head_kv;
2570 // The fp8-KV arm deliberately does not reach the draft scratch; keep the exact format
2571 // policy shared with `MtpScratch::new` so admission scales the same allocation.
2572 let (kbb, vbb) = crate::kv_blk_bytes();
2573 let k_tok_bytes = (kv_dim_k / 32) * kbb;
2574 let v_tok_bytes = (kv_dim_v / 32) * vbb;
2575 (kv_dim_k, kv_dim_v, k_tok_bytes, v_tok_bytes)
2576}
2577
2578fn mtp_chain_head_index(step: usize, head_count: usize) -> usize {
2579 assert!(head_count > 0, "MTP chain requires at least one head");
2580 step % head_count
2581}
2582
2583impl MtpScratch {
2584 fn alloc_plane(
2585 e: &Engine,
2586 cfg: &memra_gguf::config::ModelConfig,
2587 plan: &memra_gguf::model_plan::ModelPlan,
2588 cap: usize,
2589 geom: Option<&crate::hybrid::DraftGeom>,
2590 ) -> Result<MtpScratchPlane, Box<dyn std::error::Error>> {
2591 let (kv_dim_k, kv_dim_v, k_tok_bytes, v_tok_bytes) = mtp_scratch_layout(cfg, geom);
2592 let ring = if crate::cache::swa_ring_on()
2593 && crate::plan_backend::decode_batch_program(plan)
2594 == crate::plan_backend::DecodeBatchProgram::SlidingGatedMoe
2595 {
2596 let window = plan
2597 .layers
2598 .iter()
2599 .find_map(|layer| match layer.attention {
2600 memra_gguf::model_plan::AttentionPlan::SlidingWindow { window, .. } => {
2601 Some(window as usize)
2602 }
2603 _ => None,
2604 })
2605 .ok_or("sliding-gated-MoE draft scratch has no sliding-window layer")?;
2606 Some(crate::cache::KvRing::new(
2607 crate::cache::swa_ring_rows(window, cap),
2608 window,
2609 ))
2610 } else {
2611 None
2612 };
2613 let alloc_rows = ring.as_ref().map(crate::cache::KvRing::rows).unwrap_or(cap);
2614 // Ring-backed planes arm the device base mirror for the dcw draft arm (see
2615 // KvLayer::base_d): the captured chain derives its physical rows from
2616 // (len_d, base_d, window) with zero per-token node updates.
2617 let base_d = match ring.as_ref() {
2618 Some(_) => Some(e.htod_i32(&[0])?),
2619 None => None,
2620 };
2621 Ok(MtpScratchPlane {
2622 kv: KvLayer {
2623 k: e.alloc_u8(alloc_rows * k_tok_bytes)?,
2624 v: e.alloc_u8(alloc_rows * v_tok_bytes)?,
2625 kv_dim_k,
2626 kv_dim_v,
2627 k_tok_bytes,
2628 v_tok_bytes,
2629 len: 0,
2630 ring,
2631 len_d: e.htod_i32(&[0])?,
2632 base_d,
2633 },
2634 cap,
2635 })
2636 }
2637
2638 fn new(
2639 e: &Engine,
2640 cfg: &memra_gguf::config::ModelConfig,
2641 plan: &memra_gguf::model_plan::ModelPlan,
2642 cap: usize,
2643 geom: Option<&crate::hybrid::DraftGeom>,
2644 ) -> Result<Self, Box<dyn std::error::Error>> {
2645 // env-selected KV formats (default 34/24). The fp8-KV arm (MEMRA_KV_FP8) deliberately
2646 // does NOT reach the draft scratch: fp8 drafts drifted acceptance 69-88% -> 46%
2647 // (2026-07-12 A/B); the scratch is tiny, so it keeps baseline q8_0/q5_1 numerics
2648 // while the TRUNK cache carries the fp8 depth win. Scratch append/fa pass g=false.
2649 let primary = Self::alloc_plane(e, cfg, plan, cap, geom)?;
2650 Ok(MtpScratch {
2651 kv: primary.kv,
2652 cap: primary.cap,
2653 extra: Vec::new(),
2654 })
2655 }
2656
2657 fn push_plane(
2658 &mut self,
2659 e: &Engine,
2660 cfg: &memra_gguf::config::ModelConfig,
2661 plan: &memra_gguf::model_plan::ModelPlan,
2662 geom: Option<&crate::hybrid::DraftGeom>,
2663 ) -> Result<(), Box<dyn std::error::Error>> {
2664 self.extra
2665 .push(Self::alloc_plane(e, cfg, plan, self.cap, geom)?);
2666 Ok(())
2667 }
2668
2669 fn plane_count(&self) -> usize {
2670 1 + self.extra.len()
2671 }
2672
2673 fn plane(&self, index: usize) -> (&KvLayer, usize) {
2674 if index == 0 {
2675 (&self.kv, self.cap)
2676 } else {
2677 let plane = &self.extra[index - 1];
2678 (&plane.kv, plane.cap)
2679 }
2680 }
2681
2682 fn plane_mut(&mut self, index: usize) -> (&mut KvLayer, usize) {
2683 if index == 0 {
2684 (&mut self.kv, self.cap)
2685 } else {
2686 let plane = &mut self.extra[index - 1];
2687 (&mut plane.kv, plane.cap)
2688 }
2689 }
2690
2691 // #[track_caller]: set_len/set_plane_len have eight call sites (checkpoint restore, spec
2692 // rollback, session grow, seed replay ...) and the lap failure needs to say WHICH one, not
2693 // just that a rewind was refused.
2694 #[track_caller]
2695 fn set_plane_len(
2696 &mut self,
2697 e: &Engine,
2698 index: usize,
2699 n: usize,
2700 ) -> Result<(), Box<dyn std::error::Error>> {
2701 let caller = std::panic::Location::caller();
2702 let (kv, cap) = self.plane_mut(index);
2703 if let Some(ring) = kv.ring.as_ref()
2704 && !ring.can_rewind_to(n)
2705 {
2706 // NAME THE NUMBERS (2026-08-28). This error is a step37 serving blocker on the
2707 // vendor-default shape and it fires from more than one call path with more than
2708 // one trigger: a long generation walks the checkpoint out of the ring, but a
2709 // ~4.5k-token prompt also fails within 5 s of prime, which accumulation cannot
2710 // explain. A bare message forced two rounds of guessing; the operands make each
2711 // trigger name itself.
2712 let raw = n.saturating_sub(ring.window().saturating_sub(1));
2713 return Err(format!(
2714 "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})",
2715 ring.window(),
2716 ring.base(),
2717 ring.rows(),
2718 raw & !31usize,
2719 )
2720 .into());
2721 }
2722 kv.len = n;
2723 e.set_i32_one(&mut kv.len_d, n as i32)
2724 }
2725
2726 /// Set BOTH length counters: the host mirror AND the device len_d the captured append/fa read
2727 /// (a 4-byte in-place htod — the counter pointer is baked into the graph, never realloc'd).
2728 /// This is the ONLY truncation/rollback mechanism the persistent draft KV needs.
2729 #[track_caller]
2730 fn set_len(&mut self, e: &Engine, n: usize) -> Result<(), Box<dyn std::error::Error>> {
2731 let caller = std::panic::Location::caller();
2732 if !self.can_rewind_to(n) {
2733 // set_plane_len re-checks and reports the operands; call it so the failure carries
2734 // which plane refused and why, instead of this bare aggregate.
2735 for index in 0..self.plane_count() {
2736 self.set_plane_len(e, index, n)?;
2737 }
2738 return Err(format!(
2739 "SWA ring MTP checkpoint has been lapped; full re-prime required (aggregate rewind_to={n}, no single plane reported, called from {caller})"
2740 )
2741 .into());
2742 }
2743 for index in 0..self.plane_count() {
2744 self.set_plane_len(e, index, n)?;
2745 }
2746 Ok(())
2747 }
2748
2749 fn can_rewind_to(&self, n: usize) -> bool {
2750 (0..self.plane_count()).all(|index| {
2751 self.plane(index)
2752 .0
2753 .ring
2754 .as_ref()
2755 .is_none_or(|ring| ring.can_rewind_to(n))
2756 })
2757 }
2758
2759 /// Pre-arm ring headroom for `rows` upcoming DEVICE-COUNTER appends (the dcw draft arm):
2760 /// a captured chain cannot rebase mid-replay, so any rebase the coming appends could need
2761 /// happens HERE, host-side, before the capture warmups or the round's replays (the rebase
2762 /// arm of `prepare_kv_append` also refreshes the plane's `base_d` device mirror). No-op on
2763 /// flat planes and when the ring already has room; `len` is untouched either way.
2764 fn ensure_dcw_headroom(
2765 &mut self,
2766 e: &Engine,
2767 rows: usize,
2768 ) -> Result<(), Box<dyn std::error::Error>> {
2769 for index in 0..self.plane_count() {
2770 let (kv, _) = self.plane_mut(index);
2771 let Some(ring) = kv.ring.as_ref() else {
2772 continue;
2773 };
2774 let retain = memra_kv::swa_retain_from(kv.len, ring.window(), ring.base());
2775 e.prepare_kv_append(kv, retain, rows)?;
2776 }
2777 Ok(())
2778 }
2779}
2780
2781/// Retained verify intermediates for the REPLAY-FREE partial accept (2026-07-03, the profiled
2782/// #1 spec cost at long ctx: the partial-accept replay was a DUPLICATE trunk pass — ~0.54 extra
2783/// full weight reads per round — recomputing columns the verify had already produced
2784/// bit-identically). Holds, per linear layer, everything needed to rebuild its recurrent state
2785/// to "after the first j verify columns" WITHOUT re-running the trunk:
2786/// - BATCHED-path layers (`gdn`): the exact token-major inputs the round's ONE gdn_scan
2787/// consumed. A prefix re-run of the SAME kernel (t=j) from the snapshot state is bit-identical
2788/// to the first j iterations of the verify's scan — the kernel's t-loop carries state in
2789/// registers and iteration t never depends on T. `qkv_mixed` (the conv input) feeds the
2790/// pure-copy ring rebuild.
2791/// - PER-COLUMN-path layers (`cols`): dtod clones of (conv_state, ssm_state) taken after each
2792/// column 0..t-2 — pure copies of the actual chain states (the last column is never a rebuild
2793/// target: j <= t-1).
2794/// Full-attn layers need nothing: their verify KV rows are bit-identical to eager's (the
2795/// decode-exact contract; verify-probe pins it), so rollback = len truncation.
2796struct GdnStash {
2797 qkv_mixed: CudaSlice<f32>, // [t, conv_dim] token-major (conv input)
2798 q_l2: CudaSlice<f32>,
2799 k_l2: CudaSlice<f32>,
2800 v_g: CudaSlice<f32>, // [t, num_v, d_state]
2801 g_log: CudaSlice<f32>,
2802 beta: CudaSlice<f32>, // [t, num_v]
2803}
2804pub(crate) struct VerifyCkpt {
2805 gdn: Vec<Option<GdnStash>>, // [n_layer], Some iff batched linear path ran
2806 #[allow(clippy::type_complexity)]
2807 // allow: one-shot composite type; naming it would hide the shape that matters at the call site
2808 cols: Vec<Option<Vec<(CudaSlice<f32>, CudaSlice<f32>)>>>, // [n_layer][col] = (conv, ssm) after col
2809}
2810/// Opaque handle for the dspark round (dflash.rs) — VerifyCkpt stays spec-private.
2811pub(crate) struct DsparkVerifyCkpt(VerifyCkpt);
2812
2813/// Engine-bundle slice 3 (DSF-ROUNDCOST-20260820 §2 row 4 / §5 rank 1): bucketed CUDA
2814/// graphs for the dspark verify's LINEAR-layer segments. The measured verify is ~2,800
2815/// eager launches whose residual cost is DEVICE-side per-launch overhead (slice 2 proved
2816/// host dispatch is not the binder: fully-deferred dispatch bought ~0 wall). The 48 GDN
2817/// layers between full-attention layers are shape-static given vt — no positions, no
2818/// t_kv, state addressed through pointer tables — so runs of them capture per
2819/// (segment, vt) and replay as ONE graph launch each. Full-attention layers stay eager
2820/// (their per-row append/fa arm picks are t_kv-driven — the exec-update extension).
2821///
2822/// Per round out-of-graph: one pointer-table refresh (gdn ping-pong moves the canonical
2823/// handles), one input-staging copy per segment, host parity bookkeeping. Captured via
2824/// `capture_graph_retained` (2 warmups + capture, keeper retains warmup transients so
2825/// pool addresses stay stable); the warmups EXECUTE, so segment conv/ssm state is saved
2826/// before and restored after — the graph's first real launch starts from the exact
2827/// pre-round state. The ckpt column stash rides persistent slabs (written inside the
2828/// graph as memcpy nodes); commit reads them via `dspark_commit_prefix_slab`.
2829/// `MEMRA_DSPARK_VERIFY_GRAPH=0` reverts to the eager walk (byte-identical body).
2830pub(crate) struct DsparkVerifyGraphs {
2831 /// Linear-attention layer indices ascending; `lin_pos[il]` = index into the vecs.
2832 lin: Vec<usize>,
2833 lin_pos: std::collections::HashMap<usize, usize>,
2834 /// [n_lin x 6] pointer table (conv, s0, s1, conv, s1, s0 per layer), refreshed per
2835 /// verify from the live handles; layer il's slice starts at lin_pos[il]*6.
2836 table_all: CudaSlice<u64>,
2837 host_table: Vec<u64>,
2838 /// Persistent per-layer ckpt stash slabs: row r of the verify at slab offset
2839 /// r*words. Shared by every (segment, vt) bucket — one verify runs at a time.
2840 stash_conv: Vec<CudaSlice<f32>>,
2841 stash_ssm: Vec<CudaSlice<f32>>,
2842 conv_words: usize,
2843 ssm_words: usize,
2844 /// Per-vt input/output staging (stable addresses the graphs bake).
2845 stage: std::collections::HashMap<usize, (CudaSlice<f32>, CudaSlice<f32>)>,
2846 /// Per-vt dflash tap-sink buffers — the captured segments bake the tap dst address,
2847 /// so the sink buffer must live (and persist) with the graphs, not with the round.
2848 pub(crate) tap_bufs: std::collections::HashMap<usize, CudaSlice<f32>>,
2849 graphs: std::collections::HashMap<(usize, usize), DsparkSegGraph>,
2850 /// Warmup-corruption guard scratch: pre-capture conv/ssm of every linear layer
2851 /// (sized n_lin — the slice-4c full-verify warmups execute the whole walk).
2852 save_conv: CudaSlice<f32>,
2853 save_ssm: CudaSlice<f32>,
2854 max_run: usize,
2855 n_embd: usize,
2856 /// Set by the verify walk: this round's linear ckpt lives in the slabs (the caller
2857 /// commits through `dspark_commit_prefix_slab` instead of the cols arm).
2858 pub(crate) round_slab: bool,
2859 // ---- slice 4c: full-verify single graph per (vt, rung) ----
2860 /// Full-attention layer indices ascending; `fa_pos[il]` = index into the vec.
2861 fa: Vec<usize>,
2862 fa_pos: std::collections::HashMap<usize, usize>,
2863 /// [n_fa x 2 x t_cap] interleaved (k,v) base-pointer pairs, refreshed per verify;
2864 /// layer il's slice starts at `fa_pos[il] * 2 * t_cap` (the seqs twins read pairs
2865 /// [2z], z < t <= t_cap, so one t_cap-sized table serves every vt).
2866 fa_table: CudaSlice<u64>,
2867 fa_host_table: Vec<u64>,
2868 t_cap: usize,
2869 /// Per-vt position staging for the captured bodies — contents refreshed per round
2870 /// (rope reads row r; the seqs twins derive append slot and T_kv per z from it).
2871 pos_stage: std::collections::HashMap<usize, CudaSlice<i32>>,
2872 /// Full-verify graphs keyed (vt, rung_end, hi).
2873 full: std::collections::HashMap<(usize, usize, usize), DsparkSegGraph>,
2874 /// Largest n with every layer in [0, n) linear or full-attention (walk coverage).
2875 covered: usize,
2876 /// Every layer in [0, n) is linear or full-attention (no MLA/unknown mixers) — the
2877 /// full-verify capture walks all of them.
2878 walk_uniform: bool,
2879 /// Last `(captures, device graph-mem reserved bytes)` reading taken by
2880 /// `HybridModel::dspark_vg_admission_debt` — the two-point base of the MARGINAL debt
2881 /// projection (see `dspark_vg_debt_projection`; a mean-based reading extrapolated the
2882 /// pool's one-time shared allocation and reserved 8.5 GB of phantom VRAM).
2883 debt_obs: Option<(usize, usize)>,
2884}
2885
2886struct DsparkSegGraph {
2887 graph: cudarc::driver::CudaGraph,
2888 _keeper: Vec<Box<dyn std::any::Any + Send>>,
2889}
2890
2891/// Per-call arguments of [`HybridModel::qwen35_tparallel_fa_layer`] — one struct so the
2892/// eager walk and the slice-4c captured full-verify graphs hand the SAME body its two
2893/// modes without a second copy of the math.
2894pub(crate) struct FaLayerArgs<'a> {
2895 /// [T] per-row positions (device): rope reads them row-indexed; the seqs twins read
2896 /// them per-z (append slot = pos, T_kv = pos + 1).
2897 pub pos_d: &'a CudaSlice<i32>,
2898 /// Verify-level lazy per-row 1-element position buffers — only the per-row fallback
2899 /// arm builds/uses them (graph mode refuses that arm).
2900 pub pos_rows: &'a mut Option<Vec<CudaSlice<i32>>>,
2901 pub pos0: usize,
2902 pub seqs_append: bool,
2903 pub batch_fa_on: bool,
2904 /// Some((kv pointer table, offset-in-u64s, rung_end)) = captured-graph mode.
2905 pub graph_cap: Option<(&'a CudaSlice<u64>, usize, usize)>,
2906 /// ROUND-STREAM (lane/draftcost-moe, v0.100 train merge): Some((token stream, device
2907 /// round counter)) routes the FA attend through the dc rows kernels and the Linear
2908 /// mixer through `linear_attn_verify_t` (the stream arms the old inline body carried).
2909 /// Never armed together with `graph_cap` (the verify-level merge guard refuses).
2910 pub stream: Option<(&'a CudaSlice<u32>, &'a CudaSlice<i32>)>,
2911 /// VerifyCkpt for the stream-Linear arm's GdnStash install; None in graph mode and
2912 /// for FA layers that never touch it.
2913 pub ckpt: Option<&'a mut VerifyCkpt>,
2914}
2915
2916// SAFETY: `CudaGraph` is not marked Send by cudarc because its raw driver handles carry
2917// no automatic trait; CUDA driver graph handles are context-scoped rather than
2918// OS-thread-affine (the SpecPipeSessionPtr precedent above). The ctx lives in
2919// `HybridModel::dspark_vgraphs` behind a Mutex and every touch happens on the engine's
2920// single decode-stream thread.
2921unsafe impl Send for DsparkVerifyGraphs {}
2922
2923impl DsparkVerifyGraphs {
2924 /// Live capture count (segment + full graphs) — the denominator of
2925 /// [`dspark_vg_debt_projection`]'s observed bytes/capture mean.
2926 pub(crate) fn captures(&self) -> usize {
2927 self.graphs.len() + self.full.len()
2928 }
2929
2930 /// Take the marginal-growth debt reading and record this observation for the next one.
2931 /// Called under the pool mutex by `HybridModel::dspark_vg_admission_debt`.
2932 pub(crate) fn admission_debt(&mut self, reserved_bytes: usize) -> usize {
2933 let captures = self.captures();
2934 let debt =
2935 dspark_vg_debt_projection(captures, dspark_vg_cap(), reserved_bytes, self.debt_obs);
2936 if captures > 0 {
2937 match self.debt_obs {
2938 Some((c0, _)) if captures <= c0 => {}
2939 _ => self.debt_obs = Some((captures, reserved_bytes)),
2940 }
2941 }
2942 debt
2943 }
2944
2945 /// Build for this cache's shape. None when there are no linear layers, sizes are
2946 /// non-uniform, or the trunk keeps a gemma4 config (never on the qwen35 family).
2947 pub(crate) fn new(
2948 e: &Engine,
2949 cache: &Cache,
2950 t_max: usize,
2951 n_embd: usize,
2952 ) -> Result<Option<Self>, Box<dyn std::error::Error>> {
2953 let lin: Vec<usize> = (0..cache.recur.len())
2954 .filter(|&il| cache.recur[il].is_some())
2955 .collect();
2956 if lin.is_empty() || t_max < 2 {
2957 return Ok(None);
2958 }
2959 let first = cache.recur[lin[0]].as_ref().unwrap();
2960 let (conv_words, ssm_words) = (first.conv_state.len(), first.ssm_state.len());
2961 for &il in &lin {
2962 let rl = cache.recur[il].as_ref().unwrap();
2963 if rl.conv_state.len() != conv_words || rl.ssm_state.len() != ssm_words {
2964 return Ok(None);
2965 }
2966 }
2967 let n = lin.len();
2968 let mut lin_pos = std::collections::HashMap::with_capacity(n);
2969 for (k, &il) in lin.iter().enumerate() {
2970 lin_pos.insert(il, k);
2971 }
2972 // longest run of consecutive linear layers (save-scratch sizing)
2973 let mut max_run = 1usize;
2974 let mut run = 1usize;
2975 for w in lin.windows(2) {
2976 if w[1] == w[0] + 1 {
2977 run += 1;
2978 max_run = max_run.max(run);
2979 } else {
2980 run = 1;
2981 }
2982 }
2983 let rows = t_max - 1;
2984 let mut stash_conv = Vec::with_capacity(n);
2985 let mut stash_ssm = Vec::with_capacity(n);
2986 for _ in 0..n {
2987 stash_conv.push(e.uninit(rows * conv_words)?);
2988 stash_ssm.push(e.uninit(rows * ssm_words)?);
2989 }
2990 let host_table = vec![0u64; n * 6];
2991 let table_all = e.htod_u64(&host_table)?;
2992 // slice 4c: full-attention census for the full-verify graphs.
2993 let fa: Vec<usize> = (0..cache.kv.len())
2994 .filter(|&il| cache.kv[il].is_some())
2995 .collect();
2996 let mut fa_pos = std::collections::HashMap::with_capacity(fa.len());
2997 for (k, &il) in fa.iter().enumerate() {
2998 fa_pos.insert(il, k);
2999 }
3000 let n_layers = cache.kv.len().max(cache.recur.len());
3001 // exactly one of (linear state, kv cache) per layer — no MLA/unknown mixers.
3002 let walk_uniform = (0..n_layers).all(|il| {
3003 cache.recur.get(il).is_some_and(|r| r.is_some())
3004 != cache.kv.get(il).is_some_and(|k| k.is_some())
3005 });
3006 // Contiguous covered prefix: the largest n such that every layer in [0, n) is
3007 // linear or full-attention. The TRUNK walk is [0, layers.len()) and the cache
3008 // vecs can carry EXTRA state slots past it (the q38 export keeps the MTP head
3009 // layer's kv at the tail — hi == lin+fa never held, the s4c battery's zero
3010 // 'full' captures). The full-graph guard is walk coverage, not slot arithmetic.
3011 let covered = (0..n_layers)
3012 .take_while(|il| lin_pos.contains_key(il) || fa_pos.contains_key(il))
3013 .count();
3014 let t_cap = t_max;
3015 let fa_host_table = vec![0u64; fa.len() * 2 * t_cap];
3016 let fa_table = e.htod_u64(&fa_host_table)?;
3017 Ok(Some(Self {
3018 lin,
3019 lin_pos,
3020 table_all,
3021 host_table,
3022 stash_conv,
3023 stash_ssm,
3024 conv_words,
3025 ssm_words,
3026 stage: std::collections::HashMap::new(),
3027 tap_bufs: std::collections::HashMap::new(),
3028 graphs: std::collections::HashMap::new(),
3029 save_conv: e.uninit(n * conv_words)?,
3030 save_ssm: e.uninit(n * ssm_words)?,
3031 max_run,
3032 n_embd,
3033 round_slab: false,
3034 fa,
3035 fa_pos,
3036 fa_table,
3037 fa_host_table,
3038 t_cap,
3039 pos_stage: std::collections::HashMap::new(),
3040 full: std::collections::HashMap::new(),
3041 covered,
3042 walk_uniform,
3043 debt_obs: None,
3044 }))
3045 }
3046
3047 /// Rebuild the pointer tables from the live handles (once per verify — the gdn
3048 /// ping-pong swaps the canonical/alt handles between rounds; a fresh generation's
3049 /// cache buffers land at new addresses; a stale table would read the wrong state).
3050 pub(crate) fn refresh_tables(
3051 &mut self,
3052 e: &Engine,
3053 cache: &Cache,
3054 ) -> Result<(), Box<dyn std::error::Error>> {
3055 use cudarc::driver::DevicePtr;
3056 {
3057 let s = &e.gpu.stream();
3058 for (k, &il) in self.lin.iter().enumerate() {
3059 let rl = cache.recur[il].as_ref().unwrap();
3060 let (pc, _g0) = rl.conv_state.device_ptr(s);
3061 let (p0, _g1) = rl.ssm_state.device_ptr(s);
3062 let (p1, _g2) = rl.ssm_state_alt.device_ptr(s);
3063 let o = k * 6;
3064 self.host_table[o] = pc;
3065 self.host_table[o + 1] = p0;
3066 self.host_table[o + 2] = p1;
3067 self.host_table[o + 3] = pc;
3068 self.host_table[o + 4] = p1;
3069 self.host_table[o + 5] = p0;
3070 }
3071 for (k, &il) in self.fa.iter().enumerate() {
3072 let kvl = cache.kv[il].as_ref().unwrap();
3073 let (pk, _g0) = kvl.k.device_ptr(s);
3074 let (pv, _g1) = kvl.v.device_ptr(s);
3075 let o = k * 2 * self.t_cap;
3076 for z in 0..self.t_cap {
3077 self.fa_host_table[o + 2 * z] = pk;
3078 self.fa_host_table[o + 2 * z + 1] = pv;
3079 }
3080 }
3081 }
3082 e.htod_u64_into(&self.host_table, &mut self.table_all)?;
3083 if !self.fa_host_table.is_empty() {
3084 e.htod_u64_into(&self.fa_host_table, &mut self.fa_table)?;
3085 }
3086 Ok(())
3087 }
3088
3089 /// Slice 4c eligibility: Some(rung_end) when this round can replay (or capture) a
3090 /// full-verify graph — the whole walk [lo, hi) is covered, every layer is linear or
3091 /// full-attention, and ALL of the round's per-row t_kv values take the v4-seqs arm
3092 /// on ONE `fa_split_keys` ladder step that the rung also sits on (the straddle law;
3093 /// both gates are t_kv intervals, so ends-inside means all-inside). The rung is the
3094 /// round's next power of two — grid/partial sizing only (`n_splits_max` is pure
3095 /// stride; splits >= ns_eff write the empty partial the combine never reads), so one
3096 /// captured graph is bit-identical for every round the rung covers.
3097 #[allow(clippy::too_many_arguments)]
3098 pub(crate) fn full_rung(
3099 &self,
3100 model: &crate::hybrid::HybridModel,
3101 cache: &Cache,
3102 lo: usize,
3103 hi: usize,
3104 t: usize,
3105 seqs_arms_on: bool,
3106 ) -> Option<usize> {
3107 if std::env::var("MEMRA_DSPARK_FULLG_DEBUG").as_deref() == Ok("1") {
3108 static ONCE: std::sync::Once = std::sync::Once::new();
3109 let len0 = self
3110 .fa
3111 .first()
3112 .and_then(|&il| cache.kv[il].as_ref())
3113 .map(|k| k.len);
3114 ONCE.call_once(|| {
3115 eprintln!(
3116 "[fullg-debug] walk_uniform={} covered={} seqs_arms_on={} fa_rows_on={} t={} lo={} hi={} lin={} fa={} t_cap={} len0={:?}",
3117 self.walk_uniform, self.covered, seqs_arms_on, dspark_fa_rows_on(), t, lo, hi,
3118 self.lin.len(), self.fa.len(), self.t_cap, len0
3119 );
3120 });
3121 }
3122 if !self.walk_uniform
3123 || !seqs_arms_on
3124 || !dspark_fa_rows_on()
3125 || t < 2
3126 || lo != 0
3127 || hi > self.covered
3128 || t > self.t_cap
3129 || self.fa.is_empty()
3130 {
3131 return None;
3132 }
3133 let cfg = &model.cfg;
3134 let head_dim_global = cfg.head_dim_k as usize;
3135 let nkv = cfg.n_head_kv as usize;
3136 let kvl0 = cache.kv[self.fa[0]].as_ref().unwrap();
3137 // the z-batched twins read stacked rows at the cache's kv dims — must equal the
3138 // projection stride (the body's guard, hoisted so ineligible models fall back
3139 // instead of refusing mid-capture).
3140 let geom = cfg.full_attention_geometry_at(self.fa[0] as u32);
3141 let kv_dim = geom.n_head_kv as usize * geom.head_dim_k as usize;
3142 if kvl0.kv_dim_k != kv_dim || kvl0.kv_dim_v != kv_dim {
3143 return None;
3144 }
3145 let len0 = kvl0.len;
3146 let (t_kv_first, t_kv_last) = (len0 + 1, len0 + t);
3147 if !crate::fa_seqs_eligible(t_kv_first, head_dim_global)
3148 || !crate::fa_seqs_eligible(t_kv_last, head_dim_global)
3149 || crate::fa_split_keys(t_kv_first, nkv) != crate::fa_split_keys(t_kv_last, nkv)
3150 {
3151 return None;
3152 }
3153 let rung = t_kv_last.next_power_of_two().max(256);
3154 if crate::fa_split_keys(rung, nkv) != crate::fa_split_keys(t_kv_last, nkv) {
3155 return None;
3156 }
3157 Some(rung)
3158 }
3159
3160 /// Run the WHOLE verify walk [lo, hi) as one captured graph at (vt=t, rung): stage
3161 /// the residual + refresh the per-vt position staging, capture on first encounter
3162 /// (2 executing warmups bracketed by a full linear-state save/restore; KV warmup
3163 /// appends write the exact slots the replay writes — idempotent), launch, then apply
3164 /// the host bookkeeping the captured body skipped (per-linear-layer parity swap for
3165 /// odd t, per-fa-layer len bump). Returns the fresh residual.
3166 #[allow(clippy::too_many_arguments)]
3167 #[allow(clippy::map_entry)] // allow: the init body is fallible (`?`); Entry::or_insert_with cannot propagate errors
3168 pub(crate) fn run_full(
3169 &mut self,
3170 model: &crate::hybrid::HybridModel,
3171 e: &Engine,
3172 lo: usize,
3173 hi: usize,
3174 x: &CudaSlice<f32>,
3175 t: usize,
3176 pos0: usize,
3177 rung: usize,
3178 cache: &mut Cache,
3179 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3180 let n_embd = self.n_embd;
3181 if !self.stage.contains_key(&t) {
3182 let xin = e.uninit(t * n_embd)?;
3183 let xout = e.uninit(t * n_embd)?;
3184 self.stage.insert(t, (xin, xout));
3185 }
3186 if !self.pos_stage.contains_key(&t) {
3187 self.pos_stage.insert(t, e.htod_i32(&vec![0i32; t])?);
3188 }
3189 // Per-round refresh: position contents + input staging (both addresses are baked
3190 // by the captured bodies; only their CONTENTS change round to round).
3191 {
3192 let pos_host: Vec<i32> = (0..t).map(|r| (pos0 + r) as i32).collect();
3193 let pb = self.pos_stage.get_mut(&t).unwrap();
3194 e.htod_i32_into(pb, &pos_host)?;
3195 let (xin, _) = self.stage.get_mut(&t).unwrap();
3196 e.copy_into(xin, 0, x, t * n_embd)?;
3197 }
3198 let key = (t, rung, hi);
3199 if !self.full.contains_key(&key) {
3200 // The warmups EXECUTE the whole walk on live state — save every linear
3201 // layer's conv + canonical ssm first, restore after (KV needs no restore:
3202 // graph mode never bumps host lens and the appends write this round's own
3203 // slots).
3204 for (k, &il) in self.lin.iter().enumerate() {
3205 let rl = cache.recur[il].as_ref().unwrap();
3206 e.copy_into(
3207 &mut self.save_conv,
3208 k * self.conv_words,
3209 &rl.conv_state,
3210 self.conv_words,
3211 )?;
3212 e.copy_into(
3213 &mut self.save_ssm,
3214 k * self.ssm_words,
3215 &rl.ssm_state,
3216 self.ssm_words,
3217 )?;
3218 }
3219 let (graph, keeper) = {
3220 let table_all = &self.table_all;
3221 let lin_pos = &self.lin_pos;
3222 let fa_pos = &self.fa_pos;
3223 let fa_table = &self.fa_table;
3224 let t_cap = self.t_cap;
3225 let stash_conv = &mut self.stash_conv;
3226 let stash_ssm = &mut self.stash_ssm;
3227 let pos_d: &CudaSlice<i32> = &self.pos_stage[&t];
3228 let (xin, xout) = self
3229 .stage
3230 .get_mut(&t)
3231 .map(|(a, b)| (&*a, b))
3232 .expect("stage bucket created above");
3233 let cache_ref: &mut Cache = cache;
3234 let iflag = if std::env::var("MEMRA_DSPARK_VG_AUTOFREE").as_deref() == Ok("1") {
3235 cudarc::driver::sys::CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH
3236 } else {
3237 cudarc::driver::sys::CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_USE_NODE_PRIORITY
3238 };
3239 e.capture_graph_retained_flags(iflag, move |e| {
3240 let mut xc: Option<CudaSlice<f32>> = None;
3241 for il in lo..hi {
3242 let xr: &CudaSlice<f32> = xc.as_ref().unwrap_or(xin);
3243 let nx = if let Some(&k) = lin_pos.get(&il) {
3244 model.qwen35_tparallel_linear_layer(
3245 e,
3246 il,
3247 xr,
3248 t,
3249 cache_ref,
3250 None,
3251 Some((&mut stash_conv[k], &mut stash_ssm[k])),
3252 Some((table_all, k * 6)),
3253 )?
3254 } else if let Some(&kf) = fa_pos.get(&il) {
3255 let mut no_rows: Option<Vec<CudaSlice<i32>>> = None;
3256 model.qwen35_tparallel_fa_layer(
3257 e,
3258 il,
3259 xr,
3260 t,
3261 cache_ref,
3262 FaLayerArgs {
3263 pos_d,
3264 pos_rows: &mut no_rows,
3265 pos0,
3266 seqs_append: true,
3267 batch_fa_on: true,
3268 graph_cap: Some((fa_table, kf * 2 * t_cap, rung)),
3269 stream: None,
3270 ckpt: None,
3271 },
3272 )?
3273 } else {
3274 return Err(format!(
3275 "run_full: layer {il} is neither linear nor full-attention"
3276 )
3277 .into());
3278 };
3279 xc = Some(nx);
3280 }
3281 e.copy_into(xout, 0, xc.as_ref().unwrap(), t * n_embd)?;
3282 Ok(())
3283 })?
3284 };
3285 // Undo the net host parity motion of the 3 body runs (each run swaps iff t
3286 // is odd -> 3 runs = net one swap), then restore the device state the
3287 // warmups consumed (walk scope only — layers past hi never executed). The
3288 // launch below then behaves exactly like one run.
3289 if t % 2 == 1 {
3290 for &il in &self.lin {
3291 if il < lo || il >= hi {
3292 continue;
3293 }
3294 let rl = cache.recur[il].as_mut().unwrap();
3295 std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
3296 }
3297 }
3298 for (k, &il) in self.lin.iter().enumerate() {
3299 if il < lo || il >= hi {
3300 continue;
3301 }
3302 let rl = cache.recur[il].as_mut().unwrap();
3303 let (cw, sw) = (self.conv_words, self.ssm_words);
3304 {
3305 let sv = e.view(&self.save_conv, self.lin.len() * cw);
3306 let win = sv.slice(k * cw..(k + 1) * cw);
3307 e.copy_view_into(&mut rl.conv_state, 0, &win, cw)?;
3308 }
3309 {
3310 let sv = e.view(&self.save_ssm, self.lin.len() * sw);
3311 let win = sv.slice(k * sw..(k + 1) * sw);
3312 e.copy_view_into(&mut rl.ssm_state, 0, &win, sw)?;
3313 }
3314 }
3315 if std::env::var("MEMRA_GRAPH_CENSUS").as_deref() == Ok("1")
3316 && let Ok(c) = crate::graph_update::node_census(&graph)
3317 {
3318 eprintln!("[dspark-vg-census] full vt={t} rung={rung} {c:?}");
3319 }
3320 self.full.insert(
3321 key,
3322 DsparkSegGraph {
3323 graph,
3324 _keeper: keeper,
3325 },
3326 );
3327 }
3328 self.full[&key].graph.launch()?;
3329 // Host bookkeeping for the replayed body (captured host code does not re-run):
3330 // gdn parity swap per linear layer (t odd), kv len bump per fa layer — scoped
3331 // to the WALK [lo, hi): the cache can carry extra state slots past it (the MTP
3332 // head layer's kv) that the walk never touches.
3333 if t % 2 == 1 {
3334 for &il in &self.lin {
3335 if il < lo || il >= hi {
3336 continue;
3337 }
3338 let rl = cache.recur[il].as_mut().unwrap();
3339 std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
3340 }
3341 }
3342 for &il in &self.fa {
3343 if il < lo || il >= hi {
3344 continue;
3345 }
3346 cache.kv[il].as_mut().unwrap().len += t;
3347 }
3348 let (_, xout) = self.stage.get(&t).unwrap();
3349 let mut out = e.uninit(t * n_embd)?;
3350 e.copy_into(&mut out, 0, xout, t * n_embd)?;
3351 Ok(out)
3352 }
3353
3354 /// Run layers [start, end) (all linear) as one captured graph at this vt: stage the
3355 /// residual into the bucket's x_in, capture on first encounter (2 executing warmups
3356 /// bracketed by a segment state save/restore), launch, then apply the host parity
3357 /// bookkeeping the captured body would have done. Returns the fresh residual.
3358 #[allow(clippy::too_many_arguments)]
3359 #[allow(clippy::map_entry)] // allow: the init body is fallible (`?`); Entry::or_insert_with cannot propagate errors
3360 fn run_segment(
3361 &mut self,
3362 model: &crate::hybrid::HybridModel,
3363 e: &Engine,
3364 start: usize,
3365 end: usize,
3366 x: &CudaSlice<f32>,
3367 t: usize,
3368 cache: &mut Cache,
3369 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3370 let n_embd = self.n_embd;
3371 debug_assert!(end - start <= self.max_run);
3372 if !self.stage.contains_key(&t) {
3373 let xin = e.uninit(t * n_embd)?;
3374 let xout = e.uninit(t * n_embd)?;
3375 self.stage.insert(t, (xin, xout));
3376 }
3377 // Stage the residual at the bucket's baked input address.
3378 {
3379 let (xin, _) = self.stage.get_mut(&t).unwrap();
3380 e.copy_into(xin, 0, x, t * n_embd)?;
3381 }
3382 let key = (start, t);
3383 if !self.graphs.contains_key(&key) {
3384 // The 2 warmups EXECUTE the segment on live state — save conv + the canonical
3385 // ssm of every segment layer first, restore after, so the graph's first real
3386 // launch starts from the exact pre-round state (bytes gated e2e).
3387 for (k, il) in (start..end).enumerate() {
3388 let rl = cache.recur[il].as_ref().unwrap();
3389 e.copy_into(
3390 &mut self.save_conv,
3391 k * self.conv_words,
3392 &rl.conv_state,
3393 self.conv_words,
3394 )?;
3395 e.copy_into(
3396 &mut self.save_ssm,
3397 k * self.ssm_words,
3398 &rl.ssm_state,
3399 self.ssm_words,
3400 )?;
3401 }
3402 let (graph, keeper) = {
3403 let table_all = &self.table_all;
3404 let lin_pos = &self.lin_pos;
3405 let stash_conv = &mut self.stash_conv;
3406 let stash_ssm = &mut self.stash_ssm;
3407 let (xin, xout) = self
3408 .stage
3409 .get_mut(&t)
3410 .map(|(a, b)| (&*a, b))
3411 .expect("stage bucket created above");
3412 let cache_ref: &mut Cache = cache;
3413 // Slice 4 (fa-execupdate lane): USE_NODE_PRIORITY instead of
3414 // AUTO_FREE_ON_LAUNCH. The slice-3 measured limiter was AUTO_FREE's
3415 // launch-time mem-pool scan — 25.6 us per cuGraphLaunch x 16 segments
3416 // = ~0.41 ms/round, most of the eager-launch savings. The captured
3417 // body's cuMemAllocAsync transients are BALANCED by in-graph frees
3418 // (every transient drops inside the capture region — the generic
3419 // capture path's census precedent, 1589/1589), so AUTO_FREE has
3420 // nothing to reclaim and the graph is legal to instantiate without
3421 // it; PRIORITY is the flag the gemma slotted door ships for exactly
3422 // this reason (both alternatives drop the scan; UPLOAD via
3423 // cuGraphInstantiateWithFlags is WithParams-only and refused).
3424 // MEMRA_DSPARK_VG_AUTOFREE=1 reverts; MEMRA_GRAPH_CENSUS=1 prints
3425 // the node census at capture (the ALLOC==FREE receipt).
3426 let iflag = if std::env::var("MEMRA_DSPARK_VG_AUTOFREE").as_deref() == Ok("1") {
3427 cudarc::driver::sys::CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH
3428 } else {
3429 cudarc::driver::sys::CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_USE_NODE_PRIORITY
3430 };
3431 e.capture_graph_retained_flags(iflag, move |e| {
3432 let mut xc: Option<CudaSlice<f32>> = None;
3433 for il in start..end {
3434 let k = lin_pos[&il];
3435 let xr: &CudaSlice<f32> = xc.as_ref().unwrap_or(xin);
3436 let nx = model.qwen35_tparallel_linear_layer(
3437 e,
3438 il,
3439 xr,
3440 t,
3441 cache_ref,
3442 None,
3443 Some((&mut stash_conv[k], &mut stash_ssm[k])),
3444 Some((table_all, k * 6)),
3445 )?;
3446 xc = Some(nx);
3447 }
3448 e.copy_into(xout, 0, xc.as_ref().unwrap(), t * n_embd)?;
3449 Ok(())
3450 })?
3451 };
3452 // Undo the net host parity motion of the 3 body runs (each run swaps iff t
3453 // is odd -> 3 runs = net one swap), then restore the device state the
3454 // warmups consumed. The launch below then behaves exactly like one run.
3455 if t % 2 == 1 {
3456 for il in start..end {
3457 let rl = cache.recur[il].as_mut().unwrap();
3458 std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
3459 }
3460 }
3461 for (k, il) in (start..end).enumerate() {
3462 let rl = cache.recur[il].as_mut().unwrap();
3463 let (cw, sw) = (self.conv_words, self.ssm_words);
3464 {
3465 let sv = e.view(&self.save_conv, self.lin.len() * cw);
3466 let win = sv.slice(k * cw..(k + 1) * cw);
3467 e.copy_view_into(&mut rl.conv_state, 0, &win, cw)?;
3468 }
3469 {
3470 let sv = e.view(&self.save_ssm, self.lin.len() * sw);
3471 let win = sv.slice(k * sw..(k + 1) * sw);
3472 e.copy_view_into(&mut rl.ssm_state, 0, &win, sw)?;
3473 }
3474 }
3475 if std::env::var("MEMRA_GRAPH_CENSUS").as_deref() == Ok("1")
3476 && let Ok(c) = crate::graph_update::node_census(&graph)
3477 {
3478 eprintln!("[dspark-vg-census] seg={start}..{end} vt={t} {c:?}");
3479 }
3480 self.graphs.insert(
3481 key,
3482 DsparkSegGraph {
3483 graph,
3484 _keeper: keeper,
3485 },
3486 );
3487 }
3488 self.graphs[&key].graph.launch()?;
3489 // Host parity bookkeeping for the replayed body (the captured host swaps do not
3490 // re-run at replay).
3491 if t % 2 == 1 {
3492 for il in start..end {
3493 let rl = cache.recur[il].as_mut().unwrap();
3494 std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
3495 }
3496 }
3497 let (_, xout) = self.stage.get(&t).unwrap();
3498 let mut out = e.uninit(t * n_embd)?;
3499 e.copy_into(&mut out, 0, xout, t * n_embd)?;
3500 Ok(out)
3501 }
3502
3503 /// Pool freeze check (`dspark_vg_cap`): below the ceiling new keys may capture.
3504 fn can_capture(&self) -> bool {
3505 self.graphs.len() + self.full.len() < dspark_vg_cap()
3506 }
3507
3508 /// Round-atomic segment-door readiness: TRUE when this round's walk can ride the
3509 /// per-(segment, vt) graphs without a NEW capture past the pool ceiling — every
3510 /// linear run in [lo, hi) already has its (run_start, t) key, or capture is still
3511 /// allowed. FALSE sends the WHOLE round down the eager cols-ckpt walk: a partial
3512 /// refusal would stash some layers in the ctx slabs and others in the round's cols
3513 /// while one commit reads only one of them.
3514 pub(crate) fn segments_ready(
3515 &self,
3516 model: &crate::hybrid::HybridModel,
3517 lo: usize,
3518 hi: usize,
3519 t: usize,
3520 ) -> bool {
3521 if self.can_capture() {
3522 return true;
3523 }
3524 let mut il = lo;
3525 while il < hi {
3526 if matches!(model.layers[il].mixer, Mixer::Linear(_)) {
3527 let start = il;
3528 while il < hi && matches!(model.layers[il].mixer, Mixer::Linear(_)) {
3529 il += 1;
3530 }
3531 if !self.graphs.contains_key(&(start, t)) {
3532 return false;
3533 }
3534 } else {
3535 il += 1;
3536 }
3537 }
3538 true
3539 }
3540
3541 /// Widest verify window this pool was built for. A caller whose round exceeds it must
3542 /// take the eager walk: the stash slabs hold `t_capacity() - 1` column rows, and slicing
3543 /// past them is a panic rather than a refusal.
3544 pub(crate) fn t_capacity(&self) -> usize {
3545 self.t_cap
3546 }
3547
3548 /// Slab row (conv, ssm) device pointers + lengths for the commit restore of column
3549 /// `row` (0-based) of layer `il`. None for non-linear layers.
3550 pub(crate) fn slab_row(
3551 &self,
3552 e: &Engine,
3553 il: usize,
3554 row: usize,
3555 ) -> Option<(u64, u64, usize, usize)> {
3556 use cudarc::driver::DevicePtr;
3557 let k = *self.lin_pos.get(&il)?;
3558 let s = &e.gpu.stream();
3559 let (pc, _g0) = self.stash_conv[k].device_ptr(s);
3560 let (ps, _g1) = self.stash_ssm[k].device_ptr(s);
3561 Some((
3562 pc + (row * self.conv_words * 4) as u64,
3563 ps + (row * self.ssm_words * 4) as u64,
3564 self.conv_words,
3565 self.ssm_words,
3566 ))
3567 }
3568}
3569
3570impl VerifyCkpt {
3571 fn new(n_layer: usize) -> Self {
3572 VerifyCkpt {
3573 gdn: (0..n_layer).map(|_| None).collect(),
3574 cols: (0..n_layer).map(|_| None).collect(),
3575 }
3576 }
3577}
3578
3579/// The stage-0/TX half of one PP verify. The boundary slot is the ownership token: stage 1
3580/// consumes exactly the slot selected by `tx()` / `tx_pipelined()`, never a slot inferred from
3581/// a logical round number.
3582struct VerifyBoundaryTicket {
3583 rt: &'static crate::pp::PpNRt,
3584 caller_stream: std::sync::Arc<cudarc::driver::CudaStream>,
3585 slot: usize,
3586 pos0: usize,
3587 t: usize,
3588 payload: usize,
3589 n_st: usize,
3590 pipelined: bool,
3591 pp_anatomy: bool,
3592 pp_started: std::time::Instant,
3593 reverse_ms: f64,
3594 stage0_ms: f64,
3595 tx_ms: f64,
3596 trace: Option<SpecPipeTraceCtx>,
3597 _walk_owner: crate::pp::PpWalkLease,
3598}
3599
3600/// Explicit OPTIPIPE diagnostic control. Forced modes are set only by `optipipe-gate`; the
3601/// increment-2 controller can also be armed by the server's fresh-process research door.
3602#[derive(Clone, Copy, Debug, PartialEq, Eq)]
3603pub enum OptiForkGateMode {
3604 Disabled,
3605 Hit,
3606 Miss,
3607 Alternate,
3608 Abort,
3609 Controller,
3610}
3611
3612static OPTI_FORK_GATE_MODE: std::sync::atomic::AtomicU8 = std::sync::atomic::AtomicU8::new(0);
3613static OPTI_CONTROLLER_THRESHOLD: std::sync::atomic::AtomicU32 =
3614 std::sync::atomic::AtomicU32::new(0);
3615static OPTI_FORK_ATTEMPTS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
3616static OPTI_FORK_HITS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
3617static OPTI_FORK_MISSES: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
3618static OPTI_FORK_ABORT_DRAINS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
3619static OPTI_FORK_REFUSALS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
3620static OPTI_GATE_CHECKS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
3621static OPTI_GATE_ADMITS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
3622static OPTI_GATE_REJECTS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
3623static OPTI_RECONCILES: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
3624static OPTI_WASTED_DRAFT_TOKENS: std::sync::atomic::AtomicU64 =
3625 std::sync::atomic::AtomicU64::new(0);
3626static OPTI_SHADOW_DRAFT_TOKENS: std::sync::atomic::AtomicU64 =
3627 std::sync::atomic::AtomicU64::new(0);
3628static OPTI_BREAKER_TRIPS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
3629
3630impl OptiForkGateMode {
3631 fn code(self) -> u8 {
3632 match self {
3633 Self::Disabled => 0,
3634 Self::Hit => 1,
3635 Self::Miss => 2,
3636 Self::Alternate => 3,
3637 Self::Abort => 4,
3638 Self::Controller => 5,
3639 }
3640 }
3641
3642 fn configured() -> Self {
3643 match OPTI_FORK_GATE_MODE.load(std::sync::atomic::Ordering::Relaxed) {
3644 1 => Self::Hit,
3645 2 => Self::Miss,
3646 3 => Self::Alternate,
3647 4 => Self::Abort,
3648 5 => Self::Controller,
3649 _ => Self::Disabled,
3650 }
3651 }
3652
3653 fn action(self, generation: u64) -> OptiForkAction {
3654 match self {
3655 Self::Hit => OptiForkAction::Hit,
3656 Self::Miss => OptiForkAction::Miss,
3657 Self::Alternate if generation & 1 == 0 => OptiForkAction::Hit,
3658 Self::Alternate => OptiForkAction::Miss,
3659 Self::Abort => OptiForkAction::Abort,
3660 Self::Disabled | Self::Controller => {
3661 unreachable!("non-forced mode cannot choose a forced fork action")
3662 }
3663 }
3664 }
3665
3666 fn is_forced(self) -> bool {
3667 matches!(self, Self::Hit | Self::Miss | Self::Alternate | Self::Abort)
3668 }
3669}
3670
3671/// Arm or disarm the forced harness. Serving uses only `set_optipipe_controller_threshold`.
3672pub fn set_optipipe_gate_mode(mode: OptiForkGateMode) {
3673 OPTI_FORK_GATE_MODE.store(mode.code(), std::sync::atomic::Ordering::Relaxed);
3674}
3675
3676/// Arm the increment-2 diagnostic controller. The threshold applies to the uncalibrated
3677/// two-token draft-probability product. Serving can call this only through its explicit
3678/// fresh-process research door; the absent-door default remains byte-for-byte disabled.
3679pub fn set_optipipe_controller_threshold(threshold: f32) {
3680 assert!(threshold.is_finite() && (0.0..=1.0).contains(&threshold));
3681 OPTI_CONTROLLER_THRESHOLD.store(threshold.to_bits(), std::sync::atomic::Ordering::Relaxed);
3682 set_optipipe_gate_mode(OptiForkGateMode::Controller);
3683}
3684
3685#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
3686pub struct OptiForkGateStats {
3687 pub attempts: u64,
3688 pub hits: u64,
3689 pub misses: u64,
3690 pub abort_drains: u64,
3691 pub refusals: u64,
3692 pub gate_checks: u64,
3693 pub gate_admits: u64,
3694 pub gate_rejects: u64,
3695 pub reconciles: u64,
3696 pub wasted_draft_tokens: u64,
3697 pub shadow_draft_tokens: u64,
3698 pub breaker_trips: u64,
3699}
3700
3701#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
3702pub struct OptiForkStateIdentity {
3703 pub trunk_kv_bytes: usize,
3704 pub recurrent_bytes: usize,
3705 pub scratch_kv_bytes: usize,
3706 pub hidden_bytes: usize,
3707}
3708
3709pub fn reset_optipipe_gate_stats() {
3710 for counter in [
3711 &OPTI_FORK_ATTEMPTS,
3712 &OPTI_FORK_HITS,
3713 &OPTI_FORK_MISSES,
3714 &OPTI_FORK_ABORT_DRAINS,
3715 &OPTI_FORK_REFUSALS,
3716 &OPTI_GATE_CHECKS,
3717 &OPTI_GATE_ADMITS,
3718 &OPTI_GATE_REJECTS,
3719 &OPTI_RECONCILES,
3720 &OPTI_WASTED_DRAFT_TOKENS,
3721 &OPTI_SHADOW_DRAFT_TOKENS,
3722 &OPTI_BREAKER_TRIPS,
3723 ] {
3724 counter.store(0, std::sync::atomic::Ordering::Relaxed);
3725 }
3726}
3727
3728pub fn optipipe_gate_stats() -> OptiForkGateStats {
3729 let load = |v: &std::sync::atomic::AtomicU64| v.load(std::sync::atomic::Ordering::Relaxed);
3730 OptiForkGateStats {
3731 attempts: load(&OPTI_FORK_ATTEMPTS),
3732 hits: load(&OPTI_FORK_HITS),
3733 misses: load(&OPTI_FORK_MISSES),
3734 abort_drains: load(&OPTI_FORK_ABORT_DRAINS),
3735 refusals: load(&OPTI_FORK_REFUSALS),
3736 gate_checks: load(&OPTI_GATE_CHECKS),
3737 gate_admits: load(&OPTI_GATE_ADMITS),
3738 gate_rejects: load(&OPTI_GATE_REJECTS),
3739 reconciles: load(&OPTI_RECONCILES),
3740 wasted_draft_tokens: load(&OPTI_WASTED_DRAFT_TOKENS),
3741 shadow_draft_tokens: load(&OPTI_SHADOW_DRAFT_TOKENS),
3742 breaker_trips: load(&OPTI_BREAKER_TRIPS),
3743 }
3744}
3745
3746#[derive(Clone, Copy, Debug)]
3747struct OptiControllerPolicy {
3748 threshold: f32,
3749 consecutive_misses: u8,
3750 breaker_tripped: bool,
3751}
3752
3753impl OptiControllerPolicy {
3754 fn configured() -> Self {
3755 Self {
3756 threshold: f32::from_bits(
3757 OPTI_CONTROLLER_THRESHOLD.load(std::sync::atomic::Ordering::Relaxed),
3758 ),
3759 consecutive_misses: 0,
3760 breaker_tripped: false,
3761 }
3762 }
3763
3764 fn admit(&self, q_proxy: f32) -> bool {
3765 q_proxy.is_finite()
3766 && (0.0..=1.0).contains(&q_proxy)
3767 && (self.threshold == 0.0 || (!self.breaker_tripped && q_proxy >= self.threshold))
3768 }
3769
3770 /// Returns true exactly when this resolution newly trips the three-miss breaker.
3771 fn resolve(&mut self, hit: bool) -> bool {
3772 // q*=0 is the lane's explicit unconditional measurement arm. Its purpose is to price
3773 // every optimistic opportunity, so the safety breaker is measured separately and must
3774 // not silently turn this arm into "three attempts then serial".
3775 if self.threshold == 0.0 {
3776 self.consecutive_misses = 0;
3777 return false;
3778 }
3779 if hit {
3780 self.consecutive_misses = 0;
3781 return false;
3782 }
3783 self.consecutive_misses = self.consecutive_misses.saturating_add(1);
3784 if !self.breaker_tripped && self.consecutive_misses >= 3 {
3785 self.breaker_tripped = true;
3786 return true;
3787 }
3788 false
3789 }
3790}
3791
3792#[derive(Clone, Copy, Debug, PartialEq, Eq)]
3793enum OptiForkAction {
3794 Hit,
3795 Miss,
3796 Abort,
3797}
3798
3799#[derive(Clone, Copy, Debug, PartialEq, Eq)]
3800struct OptiForkGeneration {
3801 id: u64,
3802 slot: usize,
3803}
3804
3805#[derive(Default)]
3806struct OptiForkGenerationTracker {
3807 next: u64,
3808 live: [Option<u64>; 2],
3809}
3810
3811impl OptiForkGenerationTracker {
3812 fn reserve(&mut self) -> Result<OptiForkGeneration, Box<dyn std::error::Error>> {
3813 let generation = OptiForkGeneration {
3814 id: self.next,
3815 slot: (self.next & 1) as usize,
3816 };
3817 if let Some(live) = self.live[generation.slot] {
3818 return Err(format!(
3819 "optipipe snapshot slot {} still owns generation {live}; refusing to overwrite it",
3820 generation.slot,
3821 )
3822 .into());
3823 }
3824 self.next += 1;
3825 self.live[generation.slot] = Some(generation.id);
3826 Ok(generation)
3827 }
3828
3829 fn retire(&mut self, generation: OptiForkGeneration) -> Result<(), Box<dyn std::error::Error>> {
3830 match self.live[generation.slot] {
3831 Some(id) if id == generation.id => {
3832 self.live[generation.slot] = None;
3833 Ok(())
3834 }
3835 other => Err(format!(
3836 "optipipe generation teardown mismatch: ticket={} slot={} live={other:?}",
3837 generation.id, generation.slot,
3838 )
3839 .into()),
3840 }
3841 }
3842}
3843
3844struct OptiForkSeedGeneration {
3845 h_seed: CudaSlice<f32>,
3846 fill_prev: CudaSlice<f32>,
3847 scratch_len: usize,
3848}
3849
3850/// Allocate or refresh one full checkpoint through the engine that owns each PP stage. The
3851/// generic cache helper accepts one device and therefore cannot copy GDN state split across
3852/// devices. KV lengths and position stay host metadata; only recurrent buffers need stage-local
3853/// device ownership.
3854fn opti_snapshot_stage_owned(
3855 e: &Engine,
3856 cache: &Cache,
3857 rt: &'static crate::pp::PpNRt,
3858 fence: &[usize],
3859) -> Result<crate::cache::CacheSnapshot, Box<dyn std::error::Error>> {
3860 let n = cache.kv.len();
3861 let mut snapshot = crate::cache::CacheSnapshot {
3862 kv_len: vec![None; n],
3863 tp_kv_len: vec![None; n],
3864 conv: (0..n).map(|_| None).collect(),
3865 ssm: (0..n).map(|_| None).collect(),
3866 pos: cache.pos,
3867 };
3868 opti_snapshot_stage_owned_into(e, cache, rt, fence, &mut snapshot)?;
3869 Ok(snapshot)
3870}
3871
3872fn opti_snapshot_stage_owned_into(
3873 e: &Engine,
3874 cache: &Cache,
3875 rt: &'static crate::pp::PpNRt,
3876 fence: &[usize],
3877 snapshot: &mut crate::cache::CacheSnapshot,
3878) -> Result<(), Box<dyn std::error::Error>> {
3879 if fence.len() != rt.n_stages() + 1
3880 || snapshot.kv_len.len() != cache.kv.len()
3881 || snapshot.tp_kv_len.len() != cache.tp_kv.len()
3882 {
3883 return Err("optipipe stage-owned snapshot shape mismatch".into());
3884 }
3885 for stage in 0..rt.n_stages() {
3886 opti_snapshot_one_stage_owned_into(e, cache, rt, fence, stage, snapshot)?;
3887 }
3888 snapshot.pos = cache.pos;
3889 Ok(())
3890}
3891
3892/// Refresh one PP stage of a checkpoint. Increment 2 uses this split form so stage 0's
3893/// optimistic post-N state is captured before N+1 stage 0 is queued, while stage 1's matching
3894/// post-N state is captured only after N stage 1 is enqueued. Calling the all-stage helper at
3895/// either point would capture one side of the fork at the wrong generation.
3896fn opti_snapshot_one_stage_owned_into(
3897 e: &Engine,
3898 cache: &Cache,
3899 rt: &'static crate::pp::PpNRt,
3900 fence: &[usize],
3901 stage: usize,
3902 snapshot: &mut crate::cache::CacheSnapshot,
3903) -> Result<(), Box<dyn std::error::Error>> {
3904 if fence.len() != rt.n_stages() + 1
3905 || snapshot.kv_len.len() != cache.kv.len()
3906 || snapshot.tp_kv_len.len() != cache.tp_kv.len()
3907 || stage >= rt.n_stages()
3908 {
3909 return Err("optipipe single-stage snapshot shape mismatch".into());
3910 }
3911 let _scope = rt.enter(stage);
3912 let owner = rt.engine(stage, e);
3913 for il in fence[stage]..fence[stage + 1] {
3914 snapshot.kv_len[il] = cache.kv[il].as_ref().map(|kv| kv.len);
3915 snapshot.tp_kv_len[il] = cache.tp_kv[il]
3916 .as_ref()
3917 .map(crate::tp::ResidentTpKvCache::committed_len);
3918 match &cache.recur[il] {
3919 Some(recur) => {
3920 match snapshot.conv[il].as_mut() {
3921 Some(dst) => {
3922 owner.copy_into(dst, 0, &recur.conv_state, recur.conv_state.len())?
3923 }
3924 None => snapshot.conv[il] = Some(owner.clone_dtod(&recur.conv_state)?),
3925 }
3926 match snapshot.ssm[il].as_mut() {
3927 Some(dst) => {
3928 owner.copy_into(dst, 0, &recur.ssm_state, recur.ssm_state.len())?
3929 }
3930 None => snapshot.ssm[il] = Some(owner.clone_dtod(&recur.ssm_state)?),
3931 }
3932 }
3933 None if snapshot.conv[il].is_some() || snapshot.ssm[il].is_some() => {
3934 return Err(
3935 format!("optipipe stage-owned snapshot layer {il} changed shape").into(),
3936 );
3937 }
3938 None => {}
3939 }
3940 }
3941 snapshot.pos = cache.pos;
3942 Ok(())
3943}
3944
3945/// Increment-1 persistent fork state. Exactly two snapshot/seed slots alternate; a live ticket
3946/// names its generation and keeps teardown fail-closed. Only stage 0 is allowed to mutate before
3947/// resolve, so the reconcile tables and conditional restores are stage-local.
3948struct OptiForkState {
3949 mode: OptiForkGateMode,
3950 controller: Option<OptiControllerPolicy>,
3951 generations: OptiForkGenerationTracker,
3952 active_snapshot_slot: usize,
3953 alternate_snapshot: crate::cache::CacheSnapshot,
3954 seeds: [OptiForkSeedGeneration; 2],
3955 rt: &'static crate::pp::PpNRt,
3956 fence: [usize; 3],
3957 split: usize,
3958 len_ptrs: CudaSlice<u64>,
3959 saved_lens: CudaSlice<i32>,
3960 forced_acc: CudaSlice<u32>,
3961 valid: CudaSlice<u32>,
3962 stage0_stream: std::sync::Arc<cudarc::driver::CudaStream>,
3963 logical_payload_bytes: [usize; 2],
3964}
3965
3966struct OptiForkTicket {
3967 generation: OptiForkGeneration,
3968 boundary: Option<VerifyBoundaryTicket>,
3969 drain: std::sync::Arc<cudarc::driver::CudaStream>,
3970 settled: bool,
3971}
3972
3973struct OptiControllerTicket {
3974 generation: OptiForkGeneration,
3975 boundary: Option<VerifyBoundaryTicket>,
3976 ckpt: Option<VerifyCkpt>,
3977 verify_tokens: [u32; 2],
3978 draft_prob: f32,
3979 eager_seed: Option<CudaSlice<f32>>,
3980 q_proxy: f32,
3981 scratch_len: usize,
3982 issued_at: std::time::Instant,
3983 drain: std::sync::Arc<cudarc::driver::CudaStream>,
3984 settled: bool,
3985}
3986
3987struct OptiControllerPrepared {
3988 verify_tokens: [u32; 2],
3989 draft_prob: f32,
3990 eager_seed: Option<CudaSlice<f32>>,
3991 q_proxy: f32,
3992 scratch_len: usize,
3993}
3994
3995impl OptiControllerTicket {
3996 fn take_boundary(&mut self) -> VerifyBoundaryTicket {
3997 self.boundary
3998 .take()
3999 .expect("controller boundary ticket already consumed")
4000 }
4001
4002 fn take_ckpt(&mut self) -> VerifyCkpt {
4003 self.ckpt
4004 .take()
4005 .expect("controller verify checkpoint already consumed")
4006 }
4007
4008 fn take_eager_seed(&mut self) -> Option<CudaSlice<f32>> {
4009 self.eager_seed.take()
4010 }
4011
4012 fn settle(&mut self) {
4013 self.settled = true;
4014 }
4015}
4016
4017impl Drop for OptiControllerTicket {
4018 fn drop(&mut self) {
4019 if !self.settled {
4020 let _ = self.drain.synchronize();
4021 OPTI_FORK_ABORT_DRAINS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
4022 }
4023 }
4024}
4025
4026impl OptiForkTicket {
4027 fn take_boundary(&mut self) -> VerifyBoundaryTicket {
4028 self.boundary
4029 .take()
4030 .expect("fork ticket boundary already consumed")
4031 }
4032
4033 fn settle(&mut self) {
4034 self.settled = true;
4035 }
4036}
4037
4038impl Drop for OptiForkTicket {
4039 fn drop(&mut self) {
4040 if !self.settled {
4041 let _ = self.drain.synchronize();
4042 OPTI_FORK_ABORT_DRAINS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
4043 }
4044 }
4045}
4046
4047impl OptiForkState {
4048 #[allow(clippy::too_many_arguments)]
4049 fn new(
4050 e: &Engine,
4051 cache: &Cache,
4052 mode: OptiForkGateMode,
4053 alternate_snapshot: crate::cache::CacheSnapshot,
4054 h_seed: &CudaSlice<f32>,
4055 fill_prev: &CudaSlice<f32>,
4056 rt: &'static crate::pp::PpNRt,
4057 split: usize,
4058 n_layer: usize,
4059 ) -> Result<Self, Box<dyn std::error::Error>> {
4060 let fence = [0, split, n_layer];
4061 let mut logical_payload_bytes = [0usize; 2];
4062 for stage in 0..2 {
4063 for il in fence[stage]..fence[stage + 1] {
4064 logical_payload_bytes[stage] += alternate_snapshot.conv[il]
4065 .as_ref()
4066 .map_or(0, |v| v.len() * std::mem::size_of::<f32>());
4067 logical_payload_bytes[stage] += alternate_snapshot.ssm[il]
4068 .as_ref()
4069 .map_or(0, |v| v.len() * std::mem::size_of::<f32>());
4070 }
4071 }
4072 let seeds = [
4073 OptiForkSeedGeneration {
4074 h_seed: e.clone_dtod(h_seed)?,
4075 fill_prev: e.clone_dtod(fill_prev)?,
4076 scratch_len: 0,
4077 },
4078 OptiForkSeedGeneration {
4079 h_seed: e.clone_dtod(h_seed)?,
4080 fill_prev: e.clone_dtod(fill_prev)?,
4081 scratch_len: 0,
4082 },
4083 ];
4084 let (len_ptrs, saved_lens, forced_acc, valid, stage0_stream) = {
4085 let _stage = rt.enter(0);
4086 let e0 = rt.engine(0, e);
4087 (
4088 crate::round_stream::kv_len_ptr_table_range(e0, cache, 0..split, None)?,
4089 e0.htod_i32(&vec![0; split])?,
4090 e0.alloc_u32_zeroed(2)?,
4091 e0.alloc_u32_zeroed(1)?,
4092 e0.stream(),
4093 )
4094 };
4095 logical_payload_bytes[0] += seeds
4096 .iter()
4097 .map(|seed| (seed.h_seed.len() + seed.fill_prev.len()) * std::mem::size_of::<f32>())
4098 .sum::<usize>();
4099 logical_payload_bytes[0] += len_ptrs.len() * std::mem::size_of::<u64>()
4100 + saved_lens.len() * std::mem::size_of::<i32>()
4101 + forced_acc.len() * std::mem::size_of::<u32>()
4102 + valid.len() * std::mem::size_of::<u32>();
4103 Ok(Self {
4104 mode,
4105 controller: (mode == OptiForkGateMode::Controller)
4106 .then(OptiControllerPolicy::configured),
4107 generations: OptiForkGenerationTracker::default(),
4108 active_snapshot_slot: 0,
4109 alternate_snapshot,
4110 seeds,
4111 rt,
4112 fence,
4113 split,
4114 len_ptrs,
4115 saved_lens,
4116 forced_acc,
4117 valid,
4118 stage0_stream,
4119 logical_payload_bytes,
4120 })
4121 }
4122
4123 fn reserve(
4124 &mut self,
4125 current_snapshot: &mut crate::cache::CacheSnapshot,
4126 ) -> Result<OptiForkGeneration, Box<dyn std::error::Error>> {
4127 let generation = self.generations.reserve()?;
4128 if generation.slot != self.active_snapshot_slot {
4129 std::mem::swap(current_snapshot, &mut self.alternate_snapshot);
4130 self.active_snapshot_slot = generation.slot;
4131 }
4132 Ok(generation)
4133 }
4134
4135 fn capture_seed(
4136 &mut self,
4137 e: &Engine,
4138 generation: OptiForkGeneration,
4139 h_seed: &CudaSlice<f32>,
4140 fill_prev: &CudaSlice<f32>,
4141 scratch_len: usize,
4142 ) -> Result<(), Box<dyn std::error::Error>> {
4143 let seed = &mut self.seeds[generation.slot];
4144 e.copy_into(&mut seed.h_seed, 0, h_seed, h_seed.len())?;
4145 e.copy_into(&mut seed.fill_prev, 0, fill_prev, fill_prev.len())?;
4146 seed.scratch_len = scratch_len;
4147 Ok(())
4148 }
4149
4150 fn ticket(
4151 &self,
4152 generation: OptiForkGeneration,
4153 boundary: VerifyBoundaryTicket,
4154 ) -> OptiForkTicket {
4155 OptiForkTicket {
4156 generation,
4157 boundary: Some(boundary),
4158 drain: self.stage0_stream.clone(),
4159 settled: false,
4160 }
4161 }
4162
4163 #[allow(clippy::too_many_arguments)]
4164 fn controller_ticket(
4165 &self,
4166 generation: OptiForkGeneration,
4167 boundary: VerifyBoundaryTicket,
4168 ckpt: VerifyCkpt,
4169 verify_tokens: [u32; 2],
4170 draft_prob: f32,
4171 eager_seed: Option<CudaSlice<f32>>,
4172 q_proxy: f32,
4173 scratch_len: usize,
4174 ) -> OptiControllerTicket {
4175 OptiControllerTicket {
4176 generation,
4177 boundary: Some(boundary),
4178 ckpt: Some(ckpt),
4179 verify_tokens,
4180 draft_prob,
4181 eager_seed,
4182 q_proxy,
4183 scratch_len,
4184 issued_at: std::time::Instant::now(),
4185 drain: self.stage0_stream.clone(),
4186 settled: false,
4187 }
4188 }
4189
4190 fn reserve_successor(&mut self) -> Result<OptiForkGeneration, Box<dyn std::error::Error>> {
4191 self.generations.reserve()
4192 }
4193
4194 fn successor_snapshot_mut(&mut self) -> &mut crate::cache::CacheSnapshot {
4195 &mut self.alternate_snapshot
4196 }
4197
4198 fn promote_successor_snapshot(
4199 &mut self,
4200 current_snapshot: &mut crate::cache::CacheSnapshot,
4201 generation: OptiForkGeneration,
4202 ) {
4203 std::mem::swap(current_snapshot, &mut self.alternate_snapshot);
4204 self.active_snapshot_slot = generation.slot;
4205 }
4206
4207 fn queue_actual_reconcile(
4208 &mut self,
4209 e: &Engine,
4210 snapshot: &crate::cache::CacheSnapshot,
4211 acc: &CudaSlice<u32>,
4212 optimistic_pending: u32,
4213 base: usize,
4214 ) -> Result<(), Box<dyn std::error::Error>> {
4215 let saved: Vec<i32> = (0..self.split)
4216 .map(|il| snapshot.kv_len[il].map(|v| v as i32).unwrap_or(0))
4217 .collect();
4218 // Serving keeps the caller/accept walk on the head (stage-1) device. Record the accept
4219 // decision point there and append a wait to stage 0 after its optimistic successor/TX;
4220 // the validity/reconcile kernels must never peer-read acc before it is written. The
4221 // increment-1 harness uses primary stage 0, where stream order already provides this.
4222 if self.rt.engine(0, e).ctx().ordinal() != e.ctx().ordinal() {
4223 self.rt.fence_stages_behind(&e.stream())?;
4224 }
4225 let _stage = self.rt.enter(0);
4226 let e0 = self.rt.engine(0, e);
4227 e0.htod_i32_into(&mut self.saved_lens, &saved)?;
4228 e0.spec_fork_valid(acc, optimistic_pending, &mut self.valid)?;
4229 e0.spec_fork_reconcile_kv(
4230 &self.len_ptrs,
4231 &self.saved_lens,
4232 acc,
4233 &self.valid,
4234 base,
4235 self.split,
4236 )
4237 }
4238
4239 fn finish_actual_reconcile(
4240 &mut self,
4241 e: &Engine,
4242 cache: &mut Cache,
4243 snapshot: &crate::cache::CacheSnapshot,
4244 n_acc: usize,
4245 base: usize,
4246 hit: bool,
4247 ) -> Result<(), Box<dyn std::error::Error>> {
4248 if hit {
4249 return Ok(());
4250 }
4251 let len_delta = base + n_acc;
4252 for il in 0..self.split {
4253 if let (Some(kv), Some(saved)) = (cache.kv[il].as_mut(), snapshot.kv_len[il]) {
4254 kv.len = saved + len_delta;
4255 }
4256 }
4257 {
4258 let _stage = self.rt.enter(1);
4259 let e1 = self.rt.engine(1, e);
4260 for il in self.split..self.fence[2] {
4261 if let (Some(kv), Some(saved)) = (cache.kv[il].as_mut(), snapshot.kv_len[il]) {
4262 kv.len = saved + len_delta;
4263 e1.set_i32_one(&mut kv.len_d, kv.len as i32)?;
4264 }
4265 }
4266 }
4267 self.rt.publish_to(0, &e.stream())?;
4268 Ok(())
4269 }
4270
4271 fn cancel_controller_ticket(
4272 &mut self,
4273 e: &Engine,
4274 cache: &mut Cache,
4275 scratch: &mut MtpScratch,
4276 snapshot: &crate::cache::CacheSnapshot,
4277 ticket: &mut OptiControllerTicket,
4278 ) -> Result<(), Box<dyn std::error::Error>> {
4279 {
4280 let _stage = self.rt.enter(0);
4281 let e0 = self.rt.engine(0, e);
4282 for il in 0..self.split {
4283 if let (Some(kv), Some(saved)) = (cache.kv[il].as_mut(), snapshot.kv_len[il]) {
4284 kv.len = saved;
4285 e0.set_i32_one(&mut kv.len_d, saved as i32)?;
4286 }
4287 }
4288 }
4289 scratch.set_len(e, snapshot.pos)?;
4290 ticket.settle();
4291 self.generations.retire(ticket.generation)?;
4292 OPTI_FORK_ABORT_DRAINS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
4293 OPTI_WASTED_DRAFT_TOKENS.fetch_add(2, std::sync::atomic::Ordering::Relaxed);
4294 eprintln!(
4295 "[opti-controller] tail-drain generation={} slot={}",
4296 ticket.generation.id, ticket.generation.slot,
4297 );
4298 Ok(())
4299 }
4300
4301 #[allow(clippy::too_many_arguments)]
4302 fn reconcile(
4303 &mut self,
4304 e: &Engine,
4305 cache: &mut Cache,
4306 scratch: &mut MtpScratch,
4307 snapshot: &crate::cache::CacheSnapshot,
4308 h_seed: &mut CudaSlice<f32>,
4309 fill_prev: &mut CudaSlice<f32>,
4310 generation: OptiForkGeneration,
4311 action: OptiForkAction,
4312 optimistic_pending: u32,
4313 ) -> Result<(), Box<dyn std::error::Error>> {
4314 debug_assert!(action != OptiForkAction::Abort);
4315 let miss_started = std::time::Instant::now();
4316 let keep = action == OptiForkAction::Hit;
4317 let saved: Vec<i32> = (0..self.split)
4318 .map(|il| snapshot.kv_len[il].map(|v| v as i32).unwrap_or(0))
4319 .collect();
4320 let seed = &self.seeds[generation.slot];
4321 {
4322 let _stage = self.rt.enter(0);
4323 let e0 = self.rt.engine(0, e);
4324 e0.htod_i32_into(&mut self.saved_lens, &saved)?;
4325 let forced = if keep {
4326 [1u32, optimistic_pending]
4327 } else {
4328 [0u32, optimistic_pending]
4329 };
4330 e0.htod_u32_into(&mut self.forced_acc, &forced)?;
4331 e0.spec_fork_valid(&self.forced_acc, optimistic_pending, &mut self.valid)?;
4332 e0.spec_fork_reconcile_kv(
4333 &self.len_ptrs,
4334 &self.saved_lens,
4335 &self.forced_acc,
4336 &self.valid,
4337 0,
4338 self.split,
4339 )?;
4340 for il in 0..self.split {
4341 if let Some(recur) = cache.recur[il].as_mut() {
4342 let conv = snapshot.conv[il]
4343 .as_ref()
4344 .ok_or("optipipe stage0 snapshot missing conv state")?;
4345 let ssm = snapshot.ssm[il]
4346 .as_ref()
4347 .ok_or("optipipe stage0 snapshot missing ssm state")?;
4348 e0.spec_fork_restore_f32(conv, &mut recur.conv_state, &self.valid)?;
4349 e0.spec_fork_restore_f32(ssm, &mut recur.ssm_state, &self.valid)?;
4350 }
4351 }
4352 e0.spec_fork_restore_f32(&seed.h_seed, h_seed, &self.valid)?;
4353 e0.spec_fork_restore_f32(&seed.fill_prev, fill_prev, &self.valid)?;
4354 }
4355
4356 if keep {
4357 OPTI_FORK_HITS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
4358 return Ok(());
4359 }
4360
4361 for il in 0..self.split {
4362 if let (Some(kv), Some(saved)) = (cache.kv[il].as_mut(), snapshot.kv_len[il]) {
4363 kv.len = saved;
4364 }
4365 }
4366 scratch.set_len(e, seed.scratch_len)?;
4367 // Targeted E_restart: publish only stage 0's reconcile to the caller, then bound the
4368 // forced diagnostic so the retained number is the actual miss cost, not enqueue time.
4369 let caller = e.stream();
4370 self.rt.publish_to(0, &caller)?;
4371 caller.synchronize()?;
4372 let miss_ms = miss_started.elapsed().as_secs_f64() * 1e3;
4373 eprintln!(
4374 "[opti-fork-reconcile] generation={} slot={} miss_ms={miss_ms:.3}",
4375 generation.id, generation.slot,
4376 );
4377 OPTI_FORK_MISSES.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
4378 Ok(())
4379 }
4380
4381 fn retire(&mut self, generation: OptiForkGeneration) -> Result<(), Box<dyn std::error::Error>> {
4382 self.generations.retire(generation)
4383 }
4384}
4385
4386/// MEMRA_SPEC_ROUND_PROF counters: whole-round wall, so the round can be weighed against the
4387/// draft-step ([spec-anatomy]) and verify-walk ([tcol-prof]) splits we already print.
4388static ROUND_PROF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
4389static ROUND_MS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
4390static ROUND_N: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
4391
4392fn validate_tp_kv_snapshot_shape(
4393 tp_kv: &[Option<crate::tp::ResidentTpKvCache>],
4394 saved_lens: &[Option<usize>],
4395) -> Result<(), Box<dyn std::error::Error>> {
4396 if tp_kv.len() != saved_lens.len() {
4397 return Err("spec TP KV snapshot shape mismatch".into());
4398 }
4399 for (layer, (cache, saved)) in tp_kv.iter().zip(saved_lens).enumerate() {
4400 if cache.is_some() != saved.is_some() {
4401 return Err(
4402 format!("spec TP KV layer {layer} changed shape since its snapshot").into(),
4403 );
4404 }
4405 }
4406 Ok(())
4407}
4408
4409impl HybridModel {
4410 fn restore_step_tp_kv_verified_prefix(
4411 &self,
4412 e: &Engine,
4413 cache: &mut Cache,
4414 snap: &crate::cache::CacheSnapshot,
4415 accepted: usize,
4416 ) -> Result<(), Box<dyn std::error::Error>> {
4417 validate_tp_kv_snapshot_shape(&cache.tp_kv, &snap.tp_kv_len)?;
4418 e.stream().synchronize()?;
4419 {
4420 let (local_layers, distributed_layers) = (&cache.kv, &mut cache.tp_kv);
4421 let stream = e.gpu.stream();
4422 let mut runtime: Option<std::sync::Arc<crate::tp::TpE4m3HostBounce>> = None;
4423 let mut uniform_runtime = true;
4424 let mut batch = Vec::new();
4425 for (il, (distributed_slot, local_slot)) in distributed_layers
4426 .iter_mut()
4427 .zip(local_layers.iter())
4428 .enumerate()
4429 {
4430 let (Some(distributed), Some(saved)) =
4431 (distributed_slot.as_mut(), snap.tp_kv_len[il])
4432 else {
4433 continue;
4434 };
4435 let target = saved
4436 .checked_add(accepted)
4437 .ok_or("spec TP KV batch restore length overflow")?;
4438 let local = local_slot
4439 .as_ref()
4440 .ok_or_else(|| format!("spec TP KV layer {il} lost its owning cache"))?;
4441 if local.len < target {
4442 return Err(format!(
4443 "spec TP KV layer {il} local length {} precedes restore target {target}",
4444 local.len
4445 )
4446 .into());
4447 }
4448 let physical = local.physical_rows(saved, target)?;
4449 if physical.len() != accepted {
4450 return Err(format!(
4451 "spec TP KV layer {il} restore [{saved},{target}) is not contiguous"
4452 )
4453 .into());
4454 }
4455 let Mixer::Full(fa) = &self.layers[il].mixer else {
4456 return Err(format!("spec TP KV layer {il} is not full attention").into());
4457 };
4458 let tp = fa
4459 .step_tp_qkv
4460 .as_ref()
4461 .ok_or_else(|| format!("spec TP KV layer {il} lost its TP runtime"))?;
4462 if let Some(first) = runtime.as_ref() {
4463 if !std::sync::Arc::ptr_eq(first, &tp.runtime) {
4464 uniform_runtime = false;
4465 break;
4466 }
4467 } else {
4468 runtime = Some(tp.runtime.clone());
4469 }
4470 use cudarc::driver::DevicePtr;
4471 let (k_base, _k_guard) = local.k.device_ptr(&stream);
4472 let (v_base, _v_guard) = local.v.device_ptr(&stream);
4473 batch.push(crate::tp::TpKvVerifiedLayer {
4474 cache: distributed,
4475 start: saved,
4476 logical_len: target,
4477 source_k_raw: k_base + (physical.start * local.k_tok_bytes) as u64,
4478 source_v_raw: v_base + (physical.start * local.v_tok_bytes) as u64,
4479 source_k_tok_bytes: local.k_tok_bytes,
4480 source_v_tok_bytes: local.v_tok_bytes,
4481 });
4482 }
4483 if uniform_runtime
4484 && let Some(runtime) = runtime
4485 && runtime.restore_tp_kv_layers_from_device(&mut batch)?
4486 {
4487 return Ok(());
4488 }
4489 }
4490 for il in 0..self.layers.len() {
4491 let (Some(distributed), Some(saved)) = (cache.tp_kv[il].as_mut(), snap.tp_kv_len[il])
4492 else {
4493 continue;
4494 };
4495 let target = saved
4496 .checked_add(accepted)
4497 .ok_or("spec TP KV restore length overflow")?;
4498 let local = cache.kv[il]
4499 .as_ref()
4500 .ok_or_else(|| format!("spec TP KV layer {il} lost its owning cache"))?;
4501 if local.len < target {
4502 return Err(format!(
4503 "spec TP KV layer {il} local length {} precedes restore target {target}",
4504 local.len
4505 )
4506 .into());
4507 }
4508 let physical = local.physical_rows(saved, target)?;
4509 if physical.len() != accepted {
4510 return Err(format!(
4511 "spec TP KV layer {il} restore [{saved},{target}) is not contiguous"
4512 )
4513 .into());
4514 }
4515 use cudarc::driver::DevicePtr;
4516 let stream = e.gpu.stream();
4517 let (k_base, _k_guard) = local.k.device_ptr(&stream);
4518 let (v_base, _v_guard) = local.v.device_ptr(&stream);
4519 let k_raw = k_base + (physical.start * local.k_tok_bytes) as u64;
4520 let v_raw = v_base + (physical.start * local.v_tok_bytes) as u64;
4521 let Mixer::Full(fa) = &self.layers[il].mixer else {
4522 return Err(format!("spec TP KV layer {il} is not full attention").into());
4523 };
4524 let tp = fa
4525 .step_tp_qkv
4526 .as_ref()
4527 .ok_or_else(|| format!("spec TP KV layer {il} lost its TP runtime"))?;
4528 tp.runtime.restore_tp_kv_rows_from_device(
4529 distributed,
4530 saved,
4531 target,
4532 k_raw,
4533 v_raw,
4534 local.k_tok_bytes,
4535 local.v_tok_bytes,
4536 )?;
4537 }
4538 Ok(())
4539 }
4540
4541 fn mtp_head_count(&self) -> usize {
4542 usize::from(self.mtp.is_some()) + self.mtp_extra.len()
4543 }
4544
4545 fn mtp_head_at(&self, index: usize) -> &MtpHead {
4546 if index == 0 {
4547 self.mtp.as_ref().expect("MTP head 0 is unavailable")
4548 } else {
4549 &self.mtp_extra[index - 1]
4550 }
4551 }
4552
4553 fn new_mtp_scratch(
4554 &self,
4555 e: &Engine,
4556 cap: usize,
4557 ) -> Result<MtpScratch, Box<dyn std::error::Error>> {
4558 let mut scratch = MtpScratch::new(
4559 e,
4560 &self.cfg,
4561 &self.plan,
4562 cap,
4563 self.mtp.as_ref().and_then(|head| head.geom.as_ref()),
4564 )?;
4565 for head in &self.mtp_extra {
4566 scratch.push_plane(e, &self.cfg, &self.plan, head.geom.as_ref())?;
4567 }
4568 Ok(scratch)
4569 }
4570
4571 fn opti_graph_draft_step(
4572 &self,
4573 e: &Engine,
4574 mtp: &MtpHead,
4575 dctx: &mut DraftGraphCtx,
4576 scratch: &mut MtpScratch,
4577 d_vocab: usize,
4578 ) -> Result<(u32, f32), Box<dyn std::error::Error>> {
4579 // dcw door: one replay appends one device-counter row; pre-arm ring headroom
4580 // host-side before launching (no-op on flat planes).
4581 if step35_draft_dcw_on() {
4582 scratch.ensure_dcw_headroom(e, 2)?;
4583 }
4584 dctx.graph
4585 .as_ref()
4586 .ok_or("optipipe controller requires the greedy draft graph")?
4587 .launch()?;
4588 scratch.kv.len += 1;
4589 let idx = e.dtoh_u32_one(&dctx.g_tok)?;
4590 if (idx as usize) >= d_vocab {
4591 return Err(
4592 format!("optipipe draft argmax sentinel 0x{idx:08x} >= d_vocab {d_vocab}").into(),
4593 );
4594 }
4595 let probability = e.dtoh(&dctx.g_p)?[0];
4596 if !probability.is_finite() || !(0.0..=1.0).contains(&probability) {
4597 return Err(format!("optipipe draft probability is invalid: {probability}").into());
4598 }
4599 let token = match &mtp.d2t {
4600 Some(map) => map[idx as usize],
4601 None => idx,
4602 };
4603 if token != idx {
4604 e.set_u32_one(&mut dctx.g_tok, token)?;
4605 }
4606 Ok((token, probability))
4607 }
4608
4609 #[allow(clippy::too_many_arguments)]
4610 fn opti_controller_draft_step(
4611 &self,
4612 e: &Engine,
4613 mtp: &MtpHead,
4614 dctx: &mut DraftGraphCtx,
4615 scratch: &mut MtpScratch,
4616 d_vocab: usize,
4617 eager_state: &mut Option<(u32, CudaSlice<f32>)>,
4618 eager_pos: usize,
4619 embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
4620 round_graph_ok: bool,
4621 ) -> Result<(u32, f32), Box<dyn std::error::Error>> {
4622 // GRAPH-LAUNCH HEADROOM GUARD (see GRAPH_LAUNCH_MIN_FREE): `round_graph_ok` is
4623 // the round's headroom snapshot. Below the floor the main draft arm already ran
4624 // eager (13651-class gate), which seeded `eager_state`, so the controller probe
4625 // rides its eager twin below instead of replaying the draft graph into an
4626 // exhausted card. The seed-unavailable Err beneath stays the recoverable
4627 // fail-closed for the shapes that never seed it.
4628 if dctx.graph.is_some() && round_graph_ok {
4629 return self.opti_graph_draft_step(e, mtp, dctx, scratch, d_vocab);
4630 }
4631 let (input_token, input_seed) = eager_state
4632 .take()
4633 .ok_or("optipipe eager continuation seed is unavailable")?;
4634 let (logits, next_seed) = self.mtp_head_forward_dev(
4635 e,
4636 mtp,
4637 input_token,
4638 &input_seed,
4639 scratch,
4640 eager_pos,
4641 embd_dev,
4642 None,
4643 )?;
4644 let token_d = e.argmax_token_device(&logits, d_vocab)?;
4645 let idx = e.dtoh_u32_one(&token_d)?;
4646 if (idx as usize) >= d_vocab {
4647 return Err(format!(
4648 "optipipe eager draft argmax sentinel 0x{idx:08x} >= d_vocab {d_vocab}"
4649 )
4650 .into());
4651 }
4652 let probability_d = e.prob_of_token_device(&logits, &token_d, d_vocab)?;
4653 let probability = e.dtoh(&probability_d)?[0];
4654 if !probability.is_finite() || !(0.0..=1.0).contains(&probability) {
4655 return Err(
4656 format!("optipipe eager draft probability is invalid: {probability}").into(),
4657 );
4658 }
4659 let token = match &mtp.d2t {
4660 Some(map) => map[idx as usize],
4661 None => idx,
4662 };
4663 *eager_state = Some((token, next_seed));
4664 Ok((token, probability))
4665 }
4666
4667 /// NextN head forward for ONE draft token (§A ops 1-13, T=1).
4668 /// Inputs: `e_tok` = the token to predict FROM (last committed / previous draft); `h_seed` =
4669 /// the trunk's pre-output_norm hidden of that token (§A op 2 input). `mtp_pos` = absolute
4670 /// position of the token being predicted from. Returns (draft_logits[n_vocab] host, h_nextn dev).
4671 /// `h_nextn` (§A op 10) becomes `h_seed` for the next autoregressive draft step.
4672 /// Device-resident: returns draft logits ON DEVICE (no [n_vocab] dtoh). The greedy draft
4673 /// loop only needs argmax — paired with `argmax_token_device` this cuts the ~600KB logits
4674 /// transfer + host argmax per draft token from the K-token draft chain.
4675 #[allow(clippy::too_many_arguments)]
4676 fn mtp_head_forward_dev(
4677 &self,
4678 e: &Engine,
4679 mtp: &MtpHead,
4680 e_tok: u32,
4681 h_seed: &CudaSlice<f32>,
4682 scratch: &mut MtpScratch,
4683 mtp_pos: usize,
4684 embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
4685 mask: Option<(&CudaSlice<u32>, usize)>,
4686 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4687 self.mtp_head_forward_dev_at(e, mtp, e_tok, h_seed, scratch, 0, mtp_pos, embd_dev, mask)
4688 }
4689
4690 #[allow(clippy::too_many_arguments)]
4691 fn mtp_head_forward_dev_at(
4692 &self,
4693 e: &Engine,
4694 mtp: &MtpHead,
4695 e_tok: u32,
4696 h_seed: &CudaSlice<f32>,
4697 scratch: &mut MtpScratch,
4698 scratch_index: usize,
4699 mtp_pos: usize,
4700 embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
4701 // DRAFT-SIDE GRAMMAR MASK (lane/draft-mask): (packed draft-vocab allowed set, words).
4702 // Applied to the head logits BEFORE they are returned, so every consumer (argmax,
4703 // gumbel draw, p-min prob) sees the grammar-legal row. None = unmasked (pre-lane).
4704 mask: Option<(&CudaSlice<u32>, usize)>,
4705 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4706 // MEMRA_SPEC_ANATOMY=1 — eager-step phase timers (diagnostic only). Phase boundaries
4707 // sync the stream, so absolute time inflates; the BREAKDOWN is the signal. Cumulative
4708 // summary on stderr every 128 steps: glue (embed..attn_norm), attn, ffn, head.
4709 use std::sync::atomic::{AtomicU64, Ordering::Relaxed};
4710 static ANAT_NS: [AtomicU64; 5] = [
4711 AtomicU64::new(0),
4712 AtomicU64::new(0),
4713 AtomicU64::new(0),
4714 AtomicU64::new(0),
4715 AtomicU64::new(0),
4716 ];
4717 static ANAT_STEPS: AtomicU64 = AtomicU64::new(0);
4718 let anat = {
4719 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
4720 *ON.get_or_init(|| std::env::var("MEMRA_SPEC_ANATOMY").as_deref() == Ok("1"))
4721 };
4722 if anat {
4723 e.stream().synchronize()?; // drain prior queue so phase 0 starts clean
4724 }
4725 let t_all = std::time::Instant::now();
4726 let mut t_ph = std::time::Instant::now();
4727 let anat_mark = |i: usize,
4728 e: &Engine,
4729 t: &mut std::time::Instant|
4730 -> Result<(), Box<dyn std::error::Error>> {
4731 if anat {
4732 e.stream().synchronize()?;
4733 ANAT_NS[i].fetch_add(t.elapsed().as_nanos() as u64, Relaxed);
4734 *t = std::time::Instant::now();
4735 }
4736 Ok(())
4737 };
4738 let cfg = &self.cfg;
4739 let n_embd = cfg.n_embd as usize;
4740 // Distilled-student geometry: the block runs at the INNER width `di` (eh_proj out /
4741 // attn / ffn); the n_embd interface (embed, norms in, carrier out, head in) is unchanged.
4742 let di = mtp.geom.as_ref().map(|g| g.d_inner).unwrap_or(n_embd);
4743 let eps = cfg.rms_eps;
4744 let pos_d = e.htod_i32(&[mtp_pos as i32])?;
4745
4746 // op A: a resident table transfers one 4B token id. The exact host-row capacity path
4747 // expands this one row on CPU and transfers n_embd f32 values instead.
4748 let e_emb = match embd_dev {
4749 Some((g, qt, rb)) => e.embed_gather_device_t(g, &[e_tok], n_embd, qt, rb)?,
4750 None => e.htod(&self.embd.gather(n_embd, &[e_tok]))?,
4751 };
4752
4753 // op 1/2: e_norm = RMSNorm(e, enorm); h_norm = RMSNorm(h_seed, hnorm)
4754 let mut e_norm = e.zeros(n_embd)?;
4755 e.rms_norm(&e_emb, mtp.enorm.float_data(), &mut e_norm, n_embd, 1, eps)?;
4756 let mut h_norm = e.zeros(n_embd)?;
4757 e.rms_norm(h_seed, mtp.hnorm.float_data(), &mut h_norm, n_embd, 1, eps)?;
4758
4759 // op 3: concat = [e_norm ; h_norm] -> [2*n_embd], e_norm in [0,n_embd), h_norm in [n_embd,2n_embd)
4760 let mut concat = e.zeros(2 * n_embd)?;
4761 e.copy_into(&mut concat, 0, &e_norm, n_embd)?;
4762 e.copy_into(&mut concat, n_embd, &h_norm, n_embd)?;
4763
4764 // op 4: inpSA = eh_proj @ concat (eh_proj [2*n_embd, n_embd]) -> [n_embd]
4765 let inp_sa = e.matmul(&mtp.eh_proj, &concat, 1)?;
4766
4767 // op 5: a_norm = RMSNorm(inpSA, attn_norm)
4768 let mut a_norm = e.zeros(di)?;
4769 e.rms_norm(&inp_sa, mtp.attn_norm.float_data(), &mut a_norm, di, 1, eps)?;
4770 anat_mark(0, e, &mut t_ph)?;
4771
4772 // op 6: attention on the scratch KV. SAME dc launcher as the graph path (bucket_max =
4773 // scratch.cap, length from the device len_d) so eager drafts match graph drafts
4774 // bit-for-bit at any t_kv (the parity gate). Host len mirrored here (the dc append
4775 // advances only the device counter).
4776 let attn_out = match (&mtp.mixer, mtp.step35.as_ref()) {
4777 // step35 MTP block, dcw door armed: the SAME windowed device-counter launcher as
4778 // the captured chain (draft parity by construction). Per-step ring headroom runs
4779 // HERE (eager is host-len work, a rebase is legal); host len mirrored like the
4780 // plain dc arm below.
4781 (Mixer::Full(fa), Some(g))
4782 if self.step35_dcw_eligible(g, scratch.plane(scratch_index).1) =>
4783 {
4784 {
4785 let (kv, _) = scratch.plane_mut(scratch_index);
4786 let retain = match kv.ring.as_ref() {
4787 Some(ring) => memra_kv::swa_retain_from(kv.len, ring.window(), ring.base()),
4788 None => 0,
4789 };
4790 e.prepare_kv_append(kv, retain, 1)?;
4791 }
4792 let out =
4793 self.mtp_step35_attn_dcw(e, fa, g, &a_norm, &pos_d, scratch, scratch_index)?;
4794 scratch.plane_mut(scratch_index).0.len += 1;
4795 out
4796 }
4797 // step35 MTP block, door off (MEMRA_STEP35_DRAFT_DCW=0 rollback) or class-
4798 // ineligible: PER-LAYER geometry + a separate head-wise gate + an SWA window,
4799 // none of which the plain dc launcher can express (see `mtp_step35_attn`).
4800 // Host-len arm. Advances BOTH the
4801 // host len and the device counter itself (unlike the dc arm, whose host-side
4802 // mirror the caller does).
4803 (Mixer::Full(fa), Some(g)) => {
4804 self.mtp_step35_attn(e, fa, g, &a_norm, &pos_d, scratch, scratch_index)?
4805 }
4806 (Mixer::Full(fa), None) => {
4807 let out = self.mtp_full_attn_dc(
4808 e,
4809 fa,
4810 &a_norm,
4811 &pos_d,
4812 scratch,
4813 scratch_index,
4814 mtp.geom.as_ref(),
4815 )?;
4816 scratch.plane_mut(scratch_index).0.len += 1;
4817 out
4818 }
4819 (Mixer::Linear(_), _) => {
4820 panic!("MTP block is full-attn in qwen35; linear MTP not supported")
4821 }
4822 (Mixer::Mla(_), _) => crate::hybrid::mla_path_unimplemented("MTP head forward"),
4823 (Mixer::Kda(_), _) => crate::hybrid::kda_path_unimplemented("MTP head forward"),
4824 };
4825 anat_mark(1, e, &mut t_ph)?;
4826
4827 // op 7: x1 = inpSA + attn_out
4828 let mut x1 = e.zeros(di)?;
4829 e.add(&inp_sa, &attn_out, &mut x1, di)?;
4830
4831 // op 8: z = RMSNorm(x1, post_attn_norm) (pre-FFN norm)
4832 let mut z = e.zeros(di)?;
4833 e.rms_norm(&x1, mtp.post_attn_norm.float_data(), &mut z, di, 1, eps)?;
4834
4835 // op 9: FFN (Dense or MoE) — same as the trunk decode FFN
4836 let ffn_out = match &mtp.ffn {
4837 crate::hybrid::Ffn::Dense {
4838 ffn_gate,
4839 ffn_up,
4840 ffn_down,
4841 } => {
4842 let n_ff = ffn_gate.out_features();
4843 let (gate, up) = if e.uses_q8_1_fast(ffn_gate) && e.uses_q8_1_fast(ffn_up) {
4844 let (zq, zd) = e.quantize_q8_1(&z, 1, di)?;
4845 (
4846 e.matmul_pre(ffn_gate, &zq, &zd, &z, 1)?,
4847 e.matmul_pre(ffn_up, &zq, &zd, &z, 1)?,
4848 )
4849 } else {
4850 (e.matmul(ffn_gate, &z, 1)?, e.matmul(ffn_up, &z, 1)?)
4851 };
4852 let mut act = e.zeros(n_ff)?;
4853 // step35: a DENSE FFN reads the per-layer SHEXP clamp (upstream's one `build_ffn`
4854 // serves the dense MLP and the shared expert off `swiglu_clamp_shexp` —
4855 // llama-graph.cpp:1751), resolved for the MTP block's OWN index. Every other arch
4856 // passes None, which is `ffn_act`'s dispatch verbatim.
4857 Self::ffn_act_lim(
4858 e,
4859 &self.cfg,
4860 &gate,
4861 &up,
4862 1.0,
4863 1.0,
4864 mtp.step35
4865 .as_ref()
4866 .and_then(|s| s.clamp_shexp)
4867 .map(SwigluClamp::Post),
4868 &mut act,
4869 n_ff,
4870 )?;
4871 e.matmul(ffn_down, &act, 1)?
4872 }
4873 // MTP head is a distinct block — key its experts under a separate layer index (u16::MAX)
4874 // so they never alias trunk layer 0's cache keys.
4875 crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il(e, m, &z, 1, u16::MAX)?,
4876 };
4877 anat_mark(2, e, &mut t_ph)?;
4878
4879 // op 10: h_nextn = x1 + ffn_out (at di)
4880 let mut h_inner = e.zeros(di)?;
4881 e.add(&x1, &ffn_out, &mut h_inner, di)?;
4882
4883 // op 10.5 (student): up-project the inner hidden back to n_embd — training semantics:
4884 // the chain carrier AND the head input are out_up(h_inner) (pre-final-norm).
4885 let h_nextn = match mtp.geom.as_ref() {
4886 Some(g) => e.matmul(&g.out_up, &h_inner, 1)?,
4887 None => h_inner,
4888 };
4889
4890 // op 11: final = RMSNorm(h_nextn, shared_head_norm OR output_norm)
4891 let final_norm = mtp.shared_head_norm.as_ref().unwrap_or(&self.output_norm);
4892 let mut final_h = e.zeros(n_embd)?;
4893 e.rms_norm(
4894 &h_nextn,
4895 final_norm.float_data(),
4896 &mut final_h,
4897 n_embd,
4898 1,
4899 eps,
4900 )?;
4901
4902 // op 12: draft_logits = (shared_head_head OR output) @ final — stays ON DEVICE.
4903 let head = mtp.shared_head_head.as_ref().unwrap_or(&self.output);
4904 let mut logits = e.matmul(head, &final_h, 1)?;
4905 // op 12b (lane/draft-mask): grammar mask over the DRAFT vocab, applied here so the
4906 // caller's argmax / gumbel draw / p-min prob all read the grammar-legal row.
4907 if let Some((mask_d, mw)) = mask {
4908 let d_vocab = head.out_features();
4909 e.mask_logits_col(&mut logits, mask_d, 0, d_vocab, mw)?;
4910 }
4911 anat_mark(3, e, &mut t_ph)?;
4912 if anat {
4913 ANAT_NS[4].fetch_add(t_all.elapsed().as_nanos() as u64, Relaxed);
4914 let n = ANAT_STEPS.fetch_add(1, Relaxed) + 1;
4915 if n.is_multiple_of(128) {
4916 let us = |i: usize| ANAT_NS[i].load(Relaxed) / n / 1000;
4917 eprintln!(
4918 "[spec-anatomy] steps={n} avg us/step: glue={} attn={} ffn={} head={} total={}",
4919 us(0),
4920 us(1),
4921 us(2),
4922 us(3),
4923 us(4)
4924 );
4925 }
4926 }
4927 // Chain recurrence hand-over: pre-norm h_nextn (default) or post-norm final_h
4928 // (MEMRA_SPEC_HPOST — llama.cpp #24025's t_h_nextn is taken AFTER the head norm).
4929 Ok((logits, if spec_hpost() { final_h } else { h_nextn }))
4930 }
4931
4932 /// One NextN/MTP draft step for an **MLA-mixer** MTP block (glm5_next class: MLA + own
4933 /// k-pool indexer + MoE, serial residual — the NextN layer carries no hc_* tensors), on
4934 /// the model `Cache`'s own MTP latent plane rather than the full-attn `MtpScratch` the
4935 /// qwen35/step35 chain uses. Gate: `glm5_mtp_head_gpu` (engine vs `memra_reference`
4936 /// `execute_mtp`, teacher-forced walk, eh_proj-transpose and h_seed-off-by-one red arms).
4937 ///
4938 /// The interface, stated precisely for the verify arc:
4939 /// - `h_seed`: `[n_embd]` f32 device — the trunk's COLLAPSED PRE-output_norm hidden of
4940 /// the position whose next token is being drafted (MTP-PLAN §A; exactly what
4941 /// `prime_cache`/`decode_step` return for hc models). `MEMRA_SPEC_HPOST` flips both
4942 /// this producer and the returned carrier to the post-norm variant, same as the dev path.
4943 /// - `e_tok`: the token at the seeded position's SUCCESSOR — the token the trunk just
4944 /// sampled/accepted (reference oracle pairing: `fused[i] = eh_proj([enorm(embed(ids[i]));
4945 /// hnorm(trunk_hidden[i])])`, i.e. this call with `e_tok = ids[i]`, `h_seed = h[i]`,
4946 /// `mtp_pos = i` reproduces the reference's row `i`).
4947 /// - `mtp_pos`: the absolute position this step appends to the MTP block's latent plane;
4948 /// must equal that plane's current length (the plane advances by ONE row per call inside
4949 /// `mla_attn_cached`; rollback on rejection = the verify arc's latent-plane len reset).
4950 /// - returns `(draft_logits [n_vocab], carrier [n_embd])` on device. glm5_next ships no
4951 /// private MTP head, so the logits ride the trunk `lm_head` (full vocab, no d2t).
4952 pub fn mtp_head_forward_mla_cached(
4953 &self,
4954 e: &Engine,
4955 depth: usize,
4956 e_tok: u32,
4957 h_seed: &CudaSlice<f32>,
4958 cache: &mut Cache,
4959 mtp_pos: usize,
4960 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4961 if depth >= self.mtp_head_count() {
4962 return Err(format!(
4963 "MTP depth {depth} out of range: {} embedded head(s) loaded \
4964 (is MEMRA_GLM5_MTP=1 set for a glm5_next model?)",
4965 self.mtp_head_count()
4966 )
4967 .into());
4968 }
4969 let mtp = self.mtp_head_at(depth);
4970 let block = self
4971 .plan
4972 .mtp_blocks
4973 .get(depth)
4974 .ok_or_else(|| format!("ModelPlan declares no MTP block at depth {depth}"))?;
4975 let il = block.layer.index as usize;
4976 let Mixer::Mla(mla) = &mtp.mixer else {
4977 return Err(
4978 "mtp_head_forward_mla_cached serves MLA-mixer MTP blocks only; full-attn \
4979 blocks take mtp_head_forward_dev's scratch path"
4980 .into(),
4981 );
4982 };
4983 if matches!(mtp.ffn, crate::hybrid::Ffn::Dense { .. }) {
4984 return Err(
4985 "MLA-mixer MTP block with a Dense FFN has no gated arm yet (glm5_next and \
4986 glm-dsa NextN blocks are MoE); refusing rather than running ungated math"
4987 .into(),
4988 );
4989 }
4990 let plane_len = cache
4991 .latent
4992 .get(il)
4993 .and_then(|plane| plane.as_ref())
4994 .map(|plane| plane.len)
4995 .ok_or_else(|| {
4996 format!(
4997 "MTP block layer {il} has no latent cache plane — the Cache must be \
4998 built from a plan whose mtp_blocks declare StatePlan::LatentKvCache"
4999 )
5000 })?;
5001 if mtp_pos != plane_len {
5002 return Err(format!(
5003 "MTP draft position {mtp_pos} != the MTP latent plane's length {plane_len} — \
5004 the plane advances one row per draft step and rolls back by len reset; a \
5005 skipped or repeated position would attend the wrong horizon"
5006 )
5007 .into());
5008 }
5009
5010 let cfg = &self.cfg;
5011 let n_embd = cfg.n_embd as usize;
5012 let eps = cfg.rms_eps;
5013 let pos_d = e.htod_i32(&[mtp_pos as i32])?;
5014
5015 // Same op chain as `mtp_head_forward_dev_at` (ops 1-12), same kernels — only the
5016 // attention arm differs: `mla_attn_cached` on the plan's own MTP plane instead of
5017 // `mtp_full_attn_dc` on the MtpScratch.
5018 let e_emb = e.htod(&self.embd.gather(n_embd, &[e_tok]))?;
5019 let mut e_norm = e.zeros(n_embd)?;
5020 e.rms_norm(&e_emb, mtp.enorm.float_data(), &mut e_norm, n_embd, 1, eps)?;
5021 let mut h_norm = e.zeros(n_embd)?;
5022 e.rms_norm(h_seed, mtp.hnorm.float_data(), &mut h_norm, n_embd, 1, eps)?;
5023
5024 let mut concat = e.zeros(2 * n_embd)?;
5025 e.copy_into(&mut concat, 0, &e_norm, n_embd)?;
5026 e.copy_into(&mut concat, n_embd, &h_norm, n_embd)?;
5027 let inp_sa = e.matmul(&mtp.eh_proj, &concat, 1)?;
5028
5029 let mut a_norm = e.zeros(n_embd)?;
5030 e.rms_norm(
5031 &inp_sa,
5032 mtp.attn_norm.float_data(),
5033 &mut a_norm,
5034 n_embd,
5035 1,
5036 eps,
5037 )?;
5038 let attn_out = self.mla_attn_cached(e, mla, &a_norm, &pos_d, 1, il, cache)?;
5039
5040 let mut x1 = e.zeros(n_embd)?;
5041 e.add(&inp_sa, &attn_out, &mut x1, n_embd)?;
5042 let mut z = e.zeros(n_embd)?;
5043 e.rms_norm(&x1, mtp.post_attn_norm.float_data(), &mut z, n_embd, 1, eps)?;
5044 let ffn_out = match &mtp.ffn {
5045 // Distinct block — key its experts off the trunk layers' cache keys (dev-path rule).
5046 crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il(e, m, &z, 1, u16::MAX)?,
5047 crate::hybrid::Ffn::Dense { .. } => unreachable!("refused above"),
5048 };
5049 let mut h_nextn = e.zeros(n_embd)?;
5050 e.add(&x1, &ffn_out, &mut h_nextn, n_embd)?;
5051
5052 let final_norm = mtp.shared_head_norm.as_ref().unwrap_or(&self.output_norm);
5053 let mut final_h = e.zeros(n_embd)?;
5054 e.rms_norm(
5055 &h_nextn,
5056 final_norm.float_data(),
5057 &mut final_h,
5058 n_embd,
5059 1,
5060 eps,
5061 )?;
5062 let head = mtp.shared_head_head.as_ref().unwrap_or(&self.output);
5063 let logits = e.matmul(head, &final_h, 1)?;
5064 Ok((logits, if spec_hpost() { final_h } else { h_nextn }))
5065 }
5066
5067 #[allow(clippy::too_many_arguments)]
5068 fn mtp_chain_forward_dev(
5069 &self,
5070 e: &Engine,
5071 tokens: &[u32],
5072 seeds: &[CudaSlice<f32>],
5073 scratch: &mut MtpScratch,
5074 committed_scratch_len: usize,
5075 embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
5076 mask: Option<(&CudaSlice<u32>, usize)>,
5077 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
5078 if tokens.is_empty() || tokens.len() != seeds.len() {
5079 return Err("multi-head MTP prefix tokens/seeds are malformed".into());
5080 }
5081 let index = mtp_chain_head_index(tokens.len() - 1, self.mtp_head_count());
5082 let head = self.mtp_head_at(index);
5083 scratch.set_plane_len(e, index, committed_scratch_len)?;
5084
5085 let mut last = None;
5086 for row in 0..tokens.len() {
5087 let is_last = row + 1 == tokens.len();
5088 last = Some(self.mtp_head_forward_dev_at(
5089 e,
5090 head,
5091 tokens[row],
5092 &seeds[row],
5093 scratch,
5094 index,
5095 committed_scratch_len + row + 1,
5096 embd_dev,
5097 if is_last { mask } else { None },
5098 )?);
5099 }
5100 Ok(last.expect("non-empty MTP prefix produced no row"))
5101 }
5102
5103 /// step35 MTP-block attention, T=1, on the scratch KV — the EAGER-ONLY twin of
5104 /// `mtp_full_attn_dc`. Three things force a separate arm rather than a geometry parameter on
5105 /// the dc path, and all three are properties of this arch's MTP block:
5106 ///
5107 /// 1. **The SWA window.** Block 45 is an SWA-type block (`sliding_window_pattern[45]=true`,
5108 /// window 512). Windowed decode in memra is a token-aligned VIEW OFFSET into the quantized
5109 /// cache (the gemma4 R6 / `step35_decode_attn` pattern: keys carry absolute rope and the
5110 /// mask is purely positional, so one query at `len-1` attending the last `win` rows IS the
5111 /// windowed result). `fa_decode_dc` takes the key count from a DEVICE counter and always
5112 /// starts at row 0 — it cannot express a nonzero offset. The windowed dc arm is
5113 /// `mtp_step35_attn_dcw` (`fa_decode_dcw`, doored via MEMRA_STEP35_DRAFT_DCW —
5114 /// default ON since lane/step37-draft-graph-serving-20260830); this host-len arm is
5115 /// the =0 rollback and the class-ineligibility fallback.
5116 /// 2. **Per-layer head count.** 96 q heads over 8 KV (GQA 12) at this block, vs the trunk's 64
5117 /// on its full-attn layers. The trunk cfg's `n_head` scalar is the MAX over layers, and the
5118 /// trunk ARTIFACT's per-layer arrays stop at index 44 — so the count must come from the
5119 /// resolved `Step35MtpGeom`, never from `cfg`.
5120 /// 3. **The separate head-wise gate.** `blk.45.attn_gate.weight [n_embd, 96]` produces one
5121 /// sigmoid scalar per head (broadcast over head_dim) — `attn_head_gate`, not the qwen35
5122 /// fused-into-wq `q_gate_split` form the dc arm handles.
5123 ///
5124 /// DOOR STATE: with MEMRA_STEP35_DRAFT_DCW=0 (or a sub-eligible kernel class),
5125 /// `mtp_head_forward_cap` refuses step35 heads explicitly (rather than silently capturing
5126 /// a window-less, wrong-past-`win` graph) and this eager chain IS the served path. With
5127 /// the door armed (the default), BOTH draft modes run the `mtp_step35_attn_dcw` twin
5128 /// instead of this arm.
5129 ///
5130 /// Unlike the dc arm this advances BOTH the host `kv.len` and the device counter, so the
5131 /// caller must not mirror.
5132 #[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
5133 fn mtp_step35_attn(
5134 &self,
5135 e: &Engine,
5136 fa: &FullAttnLayer,
5137 g: &crate::hybrid::Step35MtpGeom,
5138 h: &CudaSlice<f32>,
5139 pos_d: &CudaSlice<i32>,
5140 scratch: &mut MtpScratch,
5141 scratch_index: usize,
5142 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5143 let (nh, nkv, hd) = (g.n_head, g.n_head_kv, self.cfg.head_dim_k as usize);
5144 // MTP-GEOM RECEIPT, once per process, on the SERVED draft path. Slot-0 acceptance is
5145 // 0.725 here against 0.994 for vLLM MTP3 on the same checkpoint family and card class, and
5146 // the first three explanations for that gap were all wrong: head assignment (step-modulo
5147 // is index 0 at K=1, correct), MEMRA_SPEC_HPOST (identical 84/116 both arms), and this
5148 // block's geometry. Geometry was the one that could have failed SILENTLY — a wrong window
5149 // makes the draft attend the whole context instead of Step-3.7's 512, stays fluent, and
5150 // shows up only as acceptance — so it gets a standing receipt rather than another reading
5151 // of the source. Prints the resolved Step35MtpGeom the served path actually runs on;
5152 // `full_attention_geometry_at`'s missing-row fallback (window: None) does NOT reach here.
5153 {
5154 static ONCE: std::sync::OnceLock<()> = std::sync::OnceLock::new();
5155 ONCE.get_or_init(|| {
5156 eprintln!(
5157 "[mtp-geom] arm=eager block={} swa={} window={} n_head={nh} n_head_kv={nkv} \
5158 head_dim_k={hd} n_rot={} rope_base={} clamp_shexp={:?}",
5159 g.il, g.swa, g.window, g.n_rot, g.rope_base, g.clamp_shexp,
5160 );
5161 });
5162 }
5163 let eps = self.cfg.rms_eps;
5164 let scale = 1.0 / (hd as f32).sqrt(); // step35.cpp:255 kq_scale
5165 let n_embd = self.cfg.n_embd as usize;
5166 let gw = fa
5167 .attn_gate
5168 .as_ref()
5169 .ok_or("step35 MTP block is missing attn_gate.weight (head-wise attention gate)")?;
5170
5171 let (q0, k0, v0, gt) = if e.uses_q8_1_fast(&fa.wq)
5172 && e.uses_q8_1_fast(&fa.wk)
5173 && e.uses_q8_1_fast(&fa.wv)
5174 && e.uses_q8_1_fast(gw)
5175 {
5176 let (hq, hdq) = e.quantize_q8_1(h, 1, n_embd)?;
5177 let (a, b, c) = match e.matmul_q8_fused3(&fa.wq, &fa.wk, &fa.wv, &hq, &hdq)? {
5178 Some(t3) => t3,
5179 None => (
5180 e.matmul_pre(&fa.wq, &hq, &hdq, h, 1)?,
5181 e.matmul_pre(&fa.wk, &hq, &hdq, h, 1)?,
5182 e.matmul_pre(&fa.wv, &hq, &hdq, h, 1)?,
5183 ),
5184 };
5185 (a, b, c, e.matmul_pre(gw, &hq, &hdq, h, 1)?)
5186 } else {
5187 (
5188 e.matmul(&fa.wq, h, 1)?,
5189 e.matmul(&fa.wk, h, 1)?,
5190 e.matmul(&fa.wv, h, 1)?,
5191 e.matmul(gw, h, 1)?,
5192 )
5193 };
5194
5195 let mut q = e.uninit(nh * hd)?;
5196 e.rms_norm(&q0, fa.q_norm.float_data(), &mut q, hd, nh, eps)?;
5197 let mut k = e.uninit(nkv * hd)?;
5198 e.rms_norm(&k0, fa.k_norm.float_data(), &mut k, hd, nkv, eps)?;
5199 // `rope_freqs.weight` (llama3 factors) applies to the FULL-attn layers ONLY; SWA passes
5200 // null (llama-hparams / step35.cpp). Block 45 is SWA, so `ff` is None there — but read
5201 // the resolved flag, not the constant, so an all-full sibling stays correct.
5202 let ff = if g.swa {
5203 None
5204 } else {
5205 self.step35_aux.as_ref().and_then(|a| a.rope_freqs(e))
5206 };
5207 #[cfg(debug_assertions)]
5208 if let Some(ff) = ff {
5209 crate::debug_assert_tensor_stream_device(ff, &e.stream(), "mtp_step35_attn.rope_freqs");
5210 }
5211 e.rope_neox2(
5212 &mut q,
5213 &mut k,
5214 pos_d,
5215 hd,
5216 g.n_rot,
5217 nh,
5218 nkv,
5219 1,
5220 g.rope_base,
5221 1.0,
5222 ff,
5223 )?;
5224
5225 // Append at the HOST slot, then re-stamp the device counter: the eager chain has the
5226 // length on the host anyway, and the windowed view below needs it there to compute the
5227 // offset. The device counter is kept in lockstep so `mtp_kv_fill`'s `set_i32_one` and any
5228 // dc-family consumer of this scratch still agree.
5229 let (kv, scratch_cap) = scratch.plane_mut(scratch_index);
5230 assert!(
5231 kv.len < scratch_cap,
5232 "step35 MTP scratch overflow ({} >= {})",
5233 kv.len,
5234 scratch_cap
5235 );
5236 let next_len = kv.len + 1;
5237 let (off, t_kv) = if g.swa && next_len > g.window {
5238 (next_len - g.window, g.window)
5239 } else {
5240 (0, next_len)
5241 };
5242 // `off`/`t_kv` stay the ATTENTION view; the retain is a separate, lower bound so the
5243 // rewind that follows this append is still resident. THIS is the only site that rebases
5244 // this plane (MEMRA_KV_REBASE_TRACE, one run: 1 rebase, all from here), so it is the site
5245 // that decides `base` for everyone.
5246 let retain_from = match kv.ring.as_ref() {
5247 Some(ring) => memra_kv::swa_retain_from(kv.len, ring.window(), ring.base()),
5248 None => off & !31usize,
5249 };
5250 let write_row = e.prepare_kv_append(kv, retain_from, 1)?;
5251 e.append_kv_quantized(
5252 &k,
5253 &v0,
5254 &mut kv.k,
5255 &mut kv.v,
5256 write_row,
5257 kv.kv_dim_k,
5258 kv.kv_dim_v,
5259 kv.k_tok_bytes,
5260 kv.v_tok_bytes,
5261 false,
5262 )?;
5263 kv.len = next_len;
5264 e.set_i32_one(&mut kv.len_d, kv.len as i32)?;
5265 // SWA view offset (see note 1). The draft chain is short (k+2 rows), but the scratch is
5266 // PERSISTENT across rounds — `mtp_kv_fill` leaves one row per committed token behind, so
5267 // `kv.len` tracks absolute position and crosses 512 in any real generation. The window is
5268 // therefore live, not theoretical.
5269 let physical = kv.physical_rows(off, off + t_kv)?;
5270 let k_view = e.view_u8_range(
5271 &kv.k,
5272 physical.start * kv.k_tok_bytes,
5273 physical.end * kv.k_tok_bytes,
5274 );
5275 let v_view = e.view_u8_range(
5276 &kv.v,
5277 physical.start * kv.v_tok_bytes,
5278 physical.end * kv.v_tok_bytes,
5279 );
5280 let mut attn = e.uninit(nh * hd)?;
5281 e.fa_decode_kvmod(
5282 &q,
5283 &k_view,
5284 &v_view,
5285 &mut attn,
5286 hd,
5287 nh,
5288 nkv,
5289 t_kv,
5290 scale,
5291 kv.k_tok_bytes,
5292 kv.v_tok_bytes,
5293 false,
5294 )?;
5295
5296 let mut ag = e.uninit(nh * hd)?;
5297 e.attn_head_gate(&attn, >, &mut ag, None, hd, nh, 1)?;
5298 e.matmul(&fa.wo, &ag, 1)
5299 }
5300
5301 /// The dcw draft arm's kernel-class precondition, mirrored from `fa_decode_dcw`'s own
5302 /// refusal plus the v3 walk's format contract (`fa_v3_active`), so the DEV dispatch can
5303 /// never pick an arm the launcher would refuse mid-chain (the eager chain has no graceful
5304 /// fallback point) and the CAP site refuses with the named reason instead.
5305 ///
5306 /// `cap` = the SESSION's scratch-plane row capacity: the launcher's vec gate reads
5307 /// `bucket_max = min(window, cap)`, so a SMALL session (tiny prompt + tiny max_tokens,
5308 /// e.g. a max_tokens=8 probe: cap ~62 < the 96 vec floor) is OUTSIDE the dcw domain even
5309 /// though the WINDOW clears the floor. Mirroring the window alone shipped exactly that
5310 /// hole when the door default flipped ON (2026-08-30, vision-cell receipt: sampled
5311 /// capture WARN + `[engine-error] fa_decode_dcw supports the default v3-vec class only`
5312 /// hard-failing the burst — the eager dcw arm has no graceful fallback point). Sub-floor
5313 /// sessions now take the host-len kvmod arm, byte-for-byte the door-off serving.
5314 fn step35_dcw_eligible(&self, g: &crate::hybrid::Step35MtpGeom, cap: usize) -> bool {
5315 let hd = self.cfg.head_dim_k as usize;
5316 step35_draft_dcw_on()
5317 && g.swa
5318 && g.window.min(cap) >= crate::fa_vec_min_tkv()
5319 && std::env::var("MEMRA_NO_FA_VEC").is_err()
5320 && crate::fa_v3_active(hd)
5321 && hd <= 256
5322 && hd.is_multiple_of(32)
5323 }
5324
5325 /// step35 MTP-block attention, T=1, on the scratch KV: the WINDOWED DEVICE-COUNTER twin
5326 /// of `mtp_step35_attn`, serving BOTH draft paths when `step35_draft_dcw_on`. Write slot,
5327 /// key bound and SWA view offset all derive from device state (`len_d`, `base_d` written
5328 /// only at host-side rebases, and the block's `window`), so ONE captured graph serves the
5329 /// whole chain and replays see KV growth through the counter: the `mtp_full_attn_dc`
5330 /// contract plus the view offset the plain `_dc` kernel could not express (the old
5331 /// capture-refusal root cause). The three step35 properties stay per-geom exactly as in
5332 /// the eager twin: nh/nkv from `Step35MtpGeom`, the separate head-wise gate
5333 /// (`attn_head_gate`), per-layer rope width/base with SWA passing null freqs.
5334 ///
5335 /// bucket_max = min(cap, window): the windowed view never exceeds `window` rows, so the
5336 /// capture-time grid stays valid for every replayed len, and the kernel derives ns_eff
5337 /// from the LIVE T_kv at the fixed split_keys (one-partition law). Both arms call THIS
5338 /// launcher at THIS bucket, so eager and captured drafts are bit-identical by
5339 /// construction; vs the retired-by-flag `mtp_step35_attn` the only numeric-class deltas
5340 /// are the sub-vec-floor region (t_kv < 96: kvmod ran scalar, dcw stays vec) and any
5341 /// live-len split-ladder rung below the bucket's, both draft-side only (the verify
5342 /// arbitrates emitted bytes; acceptance is gated by the battery).
5343 ///
5344 /// Host len is NOT advanced here (graph contract); callers mirror. The EAGER caller runs
5345 /// `prepare_kv_append` per step (ring headroom, rebase legal there); the CAPTURED path
5346 /// pre-arms headroom at capture time and round start (`MtpScratch::ensure_dcw_headroom`)
5347 /// because a rebase is host work no captured chain may contain.
5348 #[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
5349 fn mtp_step35_attn_dcw(
5350 &self,
5351 e: &Engine,
5352 fa: &FullAttnLayer,
5353 g: &crate::hybrid::Step35MtpGeom,
5354 h: &CudaSlice<f32>,
5355 pos_d: &CudaSlice<i32>,
5356 scratch: &mut MtpScratch,
5357 scratch_index: usize,
5358 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5359 let (nh, nkv, hd) = (g.n_head, g.n_head_kv, self.cfg.head_dim_k as usize);
5360 // MTP-GEOM RECEIPT (dcw twin of the `mtp_step35_attn` receipt): once per process,
5361 // naming the arm, so a serving log proves WHICH draft attention program ran (the
5362 // engagement receipt for the flag door, both directions).
5363 {
5364 static ONCE: std::sync::OnceLock<()> = std::sync::OnceLock::new();
5365 ONCE.get_or_init(|| {
5366 eprintln!(
5367 "[mtp-geom] arm=dcw block={} swa={} window={} n_head={nh} n_head_kv={nkv} \
5368 head_dim_k={hd} n_rot={} rope_base={} clamp_shexp={:?}",
5369 g.il, g.swa, g.window, g.n_rot, g.rope_base, g.clamp_shexp,
5370 );
5371 });
5372 }
5373 let eps = self.cfg.rms_eps;
5374 let scale = 1.0 / (hd as f32).sqrt(); // step35.cpp:255 kq_scale
5375 let n_embd = self.cfg.n_embd as usize;
5376 let gw = fa
5377 .attn_gate
5378 .as_ref()
5379 .ok_or("step35 MTP block is missing attn_gate.weight (head-wise attention gate)")?;
5380
5381 let (q0, k0, v0, gt) = if e.uses_q8_1_fast(&fa.wq)
5382 && e.uses_q8_1_fast(&fa.wk)
5383 && e.uses_q8_1_fast(&fa.wv)
5384 && e.uses_q8_1_fast(gw)
5385 {
5386 let (hq, hdq) = e.quantize_q8_1(h, 1, n_embd)?;
5387 let (a, b, c) = match e.matmul_q8_fused3(&fa.wq, &fa.wk, &fa.wv, &hq, &hdq)? {
5388 Some(t3) => t3,
5389 None => (
5390 e.matmul_pre(&fa.wq, &hq, &hdq, h, 1)?,
5391 e.matmul_pre(&fa.wk, &hq, &hdq, h, 1)?,
5392 e.matmul_pre(&fa.wv, &hq, &hdq, h, 1)?,
5393 ),
5394 };
5395 (a, b, c, e.matmul_pre(gw, &hq, &hdq, h, 1)?)
5396 } else {
5397 (
5398 e.matmul(&fa.wq, h, 1)?,
5399 e.matmul(&fa.wk, h, 1)?,
5400 e.matmul(&fa.wv, h, 1)?,
5401 e.matmul(gw, h, 1)?,
5402 )
5403 };
5404
5405 let mut q = e.zeros(nh * hd)?;
5406 e.rms_norm(&q0, fa.q_norm.float_data(), &mut q, hd, nh, eps)?;
5407 let mut k = e.zeros(nkv * hd)?;
5408 e.rms_norm(&k0, fa.k_norm.float_data(), &mut k, hd, nkv, eps)?;
5409 // rope_freqs (llama3 factors) apply to the FULL-attn layers ONLY; SWA passes null
5410 // (the eager twin's rule, resolved from the flag, not the constant).
5411 let ff = if g.swa {
5412 None
5413 } else {
5414 self.step35_aux.as_ref().and_then(|a| a.rope_freqs(e))
5415 };
5416 #[cfg(debug_assertions)]
5417 if let Some(ff) = ff {
5418 crate::debug_assert_tensor_stream_device(
5419 ff,
5420 &e.stream(),
5421 "mtp_step35_attn_dcw.rope_freqs",
5422 );
5423 }
5424 e.rope_neox2(
5425 &mut q,
5426 &mut k,
5427 pos_d,
5428 hd,
5429 g.n_rot,
5430 nh,
5431 nkv,
5432 1,
5433 g.rope_base,
5434 1.0,
5435 ff,
5436 )?;
5437
5438 let (kv, cap) = scratch.plane_mut(scratch_index);
5439 // Append at the DEVICE slot's PHYSICAL row (len_d - base_d), then advance the counter
5440 // in-graph. Physical room is the callers' headroom contract (see the fn doc).
5441 e.append_kv_quantized_dcw(
5442 &k,
5443 &v0,
5444 &mut kv.k,
5445 &mut kv.v,
5446 &kv.len_d,
5447 kv.base_d.as_ref(),
5448 kv.kv_dim_k,
5449 kv.kv_dim_v,
5450 kv.k_tok_bytes,
5451 kv.v_tok_bytes,
5452 )?;
5453 e.inc_seqlen(&mut kv.len_d)?;
5454 // Full-buffer views (any in-round physical row stays in range under the headroom
5455 // contract); the kernel bounds and offsets the key range from (len_d, base_d, window).
5456 let k_view = e.view_u8(&kv.k, kv.k.len());
5457 let v_view = e.view_u8(&kv.v, kv.v.len());
5458 let bucket = g.window.min(cap);
5459 let mut attn = e.zeros(nh * hd)?;
5460 e.fa_decode_dcw(
5461 &q,
5462 &k_view,
5463 &v_view,
5464 &mut attn,
5465 hd,
5466 nh,
5467 nkv,
5468 &kv.len_d,
5469 kv.base_d.as_ref(),
5470 if g.swa { g.window } else { 0 },
5471 bucket,
5472 scale,
5473 kv.k_tok_bytes,
5474 kv.v_tok_bytes,
5475 None,
5476 )?;
5477
5478 let mut ag = e.zeros(nh * hd)?;
5479 e.attn_head_gate(&attn, >, &mut ag, None, hd, nh, 1)?;
5480 e.matmul(&fa.wo, &ag, 1)
5481 }
5482
5483 /// MTP-block full attention, T=1, on the scratch KV (BOTH draft paths — eager and graph):
5484 /// the scratch write slot and the attention bound come from `scratch.kv.len_d` (device i32[1])
5485 /// so the launch args are FIXED across draft steps — ONE captured graph serves the whole
5486 /// chain, and replays keep seeing KV growth through the device counter (no recapture).
5487 /// Geometry contract: n_splits is sized from `scratch.cap` (the persistent capacity); splits
5488 /// whose key range lies beyond the device t_kv exit empty and the shared combine skips them
5489 /// (fa_decode_dc bit-correct-for-any-t_kv<=bucket_max contract). The eager path uses the SAME
5490 /// launcher with the SAME bucket_max -> identical dispatch -> bit-identical draft tokens (the
5491 /// graph-vs-eager parity gate). Host len is NOT advanced here (graph contract); callers mirror.
5492 #[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
5493 fn mtp_full_attn_dc(
5494 &self,
5495 e: &Engine,
5496 fa: &FullAttnLayer,
5497 h: &CudaSlice<f32>,
5498 pos_d: &CudaSlice<i32>,
5499 scratch: &mut MtpScratch,
5500 scratch_index: usize,
5501 geom: Option<&crate::hybrid::DraftGeom>,
5502 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5503 let cfg = &self.cfg;
5504 let mtp_il = cfg.n_layer.saturating_sub(cfg.nextn_predict_layers);
5505 let geometry = cfg.full_attention_geometry_at(mtp_il);
5506 let n_head = geom.map(|g| g.n_head).unwrap_or(geometry.n_head as usize);
5507 let n_head_kv = geom
5508 .map(|g| g.n_head_kv)
5509 .unwrap_or(geometry.n_head_kv as usize);
5510 let head_dim = geometry.head_dim_k as usize;
5511 let eps = cfg.rms_eps;
5512 let scale = geometry.attention_scale();
5513 let n_embd = geom.map(|g| g.d_inner).unwrap_or(cfg.n_embd as usize);
5514 let bucket_max = scratch.plane(scratch_index).1;
5515
5516 let (qf, mut k, v) =
5517 if e.uses_q8_1_fast(&fa.wq) && e.uses_q8_1_fast(&fa.wk) && e.uses_q8_1_fast(&fa.wv) {
5518 let (hq, hd) = e.quantize_q8_1(h, 1, n_embd)?;
5519 (
5520 e.matmul_pre(&fa.wq, &hq, &hd, h, 1)?,
5521 e.matmul_pre(&fa.wk, &hq, &hd, h, 1)?,
5522 e.matmul_pre(&fa.wv, &hq, &hd, h, 1)?,
5523 )
5524 } else {
5525 (
5526 e.matmul(&fa.wq, h, 1)?,
5527 e.matmul(&fa.wk, h, 1)?,
5528 e.matmul(&fa.wv, h, 1)?,
5529 )
5530 };
5531 // M3/Hy3 have no attention output gate — wq out is exactly q; skip the split.
5532 let gated = geometry.attention_gate == memra_gguf::config::AttentionGateKind::FusedQ;
5533 let (mut q, gate) = if gated {
5534 let mut q = e.zeros(n_head * head_dim)?;
5535 let mut gate = e.zeros(n_head * head_dim)?;
5536 e.q_gate_split(&qf, &mut q, &mut gate, head_dim, n_head, 1)?;
5537 (q, Some(gate))
5538 } else {
5539 (qf, None)
5540 };
5541
5542 let mut qn = e.zeros(n_head * head_dim)?;
5543 e.rms_norm(&q, fa.q_norm.float_data(), &mut qn, head_dim, n_head, eps)?;
5544 q = qn;
5545 let mut kn = e.zeros(n_head_kv * head_dim)?;
5546 e.rms_norm(
5547 &k,
5548 fa.k_norm.float_data(),
5549 &mut kn,
5550 head_dim,
5551 n_head_kv,
5552 eps,
5553 )?;
5554 k = kn;
5555 let rope_dims = geometry.n_rot as usize;
5556 e.rope_neox(
5557 &mut q,
5558 pos_d,
5559 head_dim,
5560 rope_dims,
5561 n_head,
5562 1,
5563 geometry.rope_base,
5564 1.0,
5565 )?;
5566 e.rope_neox(
5567 &mut k,
5568 pos_d,
5569 head_dim,
5570 rope_dims,
5571 n_head_kv,
5572 1,
5573 geometry.rope_base,
5574 1.0,
5575 )?;
5576
5577 let kv = scratch.plane_mut(scratch_index).0;
5578 // append at the DEVICE slot (kv.len_d == old len), then advance the counter in-graph.
5579 e.append_kv_quantized_dc(
5580 &k,
5581 &v,
5582 &mut kv.k,
5583 &mut kv.v,
5584 &kv.len_d,
5585 kv.kv_dim_k,
5586 kv.kv_dim_v,
5587 kv.k_tok_bytes,
5588 kv.v_tok_bytes,
5589 false,
5590 )?;
5591 e.inc_seqlen(&mut kv.len_d)?;
5592 // full-buffer views (any in-round t_kv stays in range on replay); the kernel bounds the
5593 // key range from the device counter.
5594 let k_view = e.view_u8(&kv.k, kv.k.len());
5595 let v_view = e.view_u8(&kv.v, kv.v.len());
5596 let (ktb, vtb) = (kv.k_tok_bytes, kv.v_tok_bytes);
5597 let mut attn = e.zeros(n_head * head_dim)?;
5598 e.fa_decode_dc(
5599 &q, &k_view, &v_view, &mut attn, head_dim, n_head, n_head_kv, &kv.len_d, bucket_max,
5600 scale, ktb, vtb, false,
5601 )?;
5602
5603 let attn_g = match &gate {
5604 Some(gate) => {
5605 let mut gsig = e.zeros(n_head * head_dim)?;
5606 e.sigmoid(gate, &mut gsig, n_head * head_dim)?;
5607 let mut ag = e.zeros(n_head * head_dim)?;
5608 e.mul(&attn, &gsig, &mut ag, n_head * head_dim)?;
5609 ag
5610 }
5611 None => attn,
5612 };
5613 e.matmul(&fa.wo, &attn_g, 1)
5614 }
5615
5616 /// PERSISTENT-DRAFT-KV fill (the reference engine's "mtp_update" analogue): compute the MTP
5617 /// block's K/V for `tokens` (committed tokens at positions pos0..pos0+T) from their EXACT
5618 /// trunk hiddens `h` ([T, n_embd] token-major, pre-output_norm) and append at slots pos0.. of
5619 /// the scratch KV. K/V-ONLY — ops A/1-5 plus the K-side of op 6 (wk/wv + k_norm + rope +
5620 /// quantized append); no wq/attention/FFN/lm_head, so per-token cost ~= eh_proj + wk/wv (a
5621 /// small fraction of one trunk layer), T-batched. Rope follows the chain convention
5622 /// rope(token@p) = p+1. Runs at round boundaries OUTSIDE the captured graph in BOTH draft
5623 /// modes -> draft parity by construction. Caller must have scratch.kv.len == pos0.
5624 #[allow(clippy::too_many_arguments)]
5625 fn mtp_kv_fill_at(
5626 &self,
5627 e: &Engine,
5628 mtp: &MtpHead,
5629 tokens: &[u32],
5630 h: &CudaSlice<f32>,
5631 pos0: usize,
5632 scratch: &mut MtpScratch,
5633 scratch_index: usize,
5634 embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
5635 ) -> Result<(), Box<dyn std::error::Error>> {
5636 let cfg = &self.cfg;
5637 let n_embd = cfg.n_embd as usize;
5638 let eps = cfg.rms_eps;
5639 let t = tokens.len();
5640 let (scratch_kv, scratch_cap) = scratch.plane(scratch_index);
5641 assert_eq!(scratch_kv.len, pos0, "mtp_kv_fill: append slot mismatch");
5642 assert!(pos0 + t <= scratch_cap, "mtp_kv_fill: scratch overflow");
5643 let Mixer::Full(fa) = &mtp.mixer else {
5644 panic!("MTP block is full-attn in qwen35; linear MTP not supported")
5645 };
5646 let pos_vec: Vec<i32> = (0..t).map(|i| (pos0 + i + 1) as i32).collect();
5647 let pos_d = e.htod_i32(&pos_vec)?;
5648
5649 // ops A/1/2: embed + the two input norms, T-wide.
5650 let e_emb = match embd_dev {
5651 Some((g, qt, rb)) => e.embed_gather_device_t(g, tokens, n_embd, qt, rb)?,
5652 None => e.htod(&self.embd.gather(n_embd, tokens))?,
5653 };
5654 let mut e_norm = e.zeros(t * n_embd)?;
5655 e.rms_norm(&e_emb, mtp.enorm.float_data(), &mut e_norm, n_embd, t, eps)?;
5656 let mut h_norm = e.zeros(t * n_embd)?;
5657 e.rms_norm(h, mtp.hnorm.float_data(), &mut h_norm, n_embd, t, eps)?;
5658
5659 // op 3: per-row [e_norm ; h_norm] concat, token-major [T, 2*n_embd].
5660 let mut concat = e.zeros(t * 2 * n_embd)?;
5661 for i in 0..t {
5662 e.copy_view_into(
5663 &mut concat,
5664 i * 2 * n_embd,
5665 &e_norm.slice(i * n_embd..(i + 1) * n_embd),
5666 n_embd,
5667 )?;
5668 e.copy_view_into(
5669 &mut concat,
5670 i * 2 * n_embd + n_embd,
5671 &h_norm.slice(i * n_embd..(i + 1) * n_embd),
5672 n_embd,
5673 )?;
5674 }
5675
5676 // ops 4/5: eh_proj + attn_norm, T-wide (at the student inner width when geom is set).
5677 let di = mtp.geom.as_ref().map(|g| g.d_inner).unwrap_or(n_embd);
5678 let inp_sa = e.matmul(&mtp.eh_proj, &concat, t)?;
5679 let mut a_norm = e.zeros(t * di)?;
5680 e.rms_norm(&inp_sa, mtp.attn_norm.float_data(), &mut a_norm, di, t, eps)?;
5681
5682 // op 6 (K/V half): wk/wv + k_norm + rope + per-row quantized append. No wq/attention —
5683 // the fill only has to leave correct K/V rows behind for later chains to attend over.
5684 let n_head_kv = mtp
5685 .geom
5686 .as_ref()
5687 .map(|g| g.n_head_kv)
5688 .or(mtp.step35.as_ref().map(|s| s.n_head_kv))
5689 .unwrap_or_else(|| {
5690 let mtp_il = cfg.n_layer.saturating_sub(cfg.nextn_predict_layers);
5691 cfg.full_attention_geometry_at(mtp_il).n_head_kv as usize
5692 });
5693 let mtp_il = cfg.n_layer.saturating_sub(cfg.nextn_predict_layers);
5694 let geometry = cfg.full_attention_geometry_at(mtp_il);
5695 let head_dim = geometry.head_dim_k as usize;
5696 let mut k = e.matmul(&fa.wk, &a_norm, t)?;
5697 let v = e.matmul(&fa.wv, &a_norm, t)?;
5698 let mut kn = e.zeros(t * n_head_kv * head_dim)?;
5699 e.rms_norm(
5700 &k,
5701 fa.k_norm.float_data(),
5702 &mut kn,
5703 head_dim,
5704 n_head_kv * t,
5705 eps,
5706 )?;
5707 k = kn;
5708 // step35: rotary width AND base are per-layer, and the MTP block's values come from the
5709 // resolved `Step35MtpGeom` — NOT from `cfg.rope_dim_count`/`cfg.rope_freq_base`, which
5710 // carry the arch defaults (128 / 5e6, i.e. the FULL-attn layers' base). Getting this wrong
5711 // writes K rows the attention arm then re-derives at a different theta: correct-looking
5712 // output with dead acceptance, invisible to the exactness gates.
5713 let (rope_dims, rope_base, ff) = match mtp.step35.as_ref() {
5714 Some(s) => (
5715 s.n_rot,
5716 s.rope_base,
5717 if s.swa {
5718 None
5719 } else {
5720 self.step35_aux.as_ref().and_then(|a| a.rope_freqs(e))
5721 },
5722 ),
5723 None => (geometry.n_rot as usize, geometry.rope_base, None),
5724 };
5725 #[cfg(debug_assertions)]
5726 if let Some(ff) = ff {
5727 crate::debug_assert_tensor_stream_device(ff, &e.stream(), "mtp_kv_fill.rope_freqs");
5728 }
5729 match ff {
5730 Some(f) => e.rope_neox_ff(
5731 &mut k, &pos_d, head_dim, rope_dims, n_head_kv, t, rope_base, 1.0, f,
5732 )?,
5733 None => e.rope_neox(
5734 &mut k, &pos_d, head_dim, rope_dims, n_head_kv, t, rope_base, 1.0,
5735 )?,
5736 }
5737
5738 let kv = scratch.plane_mut(scratch_index).0;
5739 // Match the trunk prime contract: a chunk may need the aligned window immediately before
5740 // its first row, so preserve that prefix when the physical tail rebases at wrap.
5741 let retain_from = kv
5742 .ring
5743 .as_ref()
5744 .map(|ring| memra_kv::swa_retain_from(pos0, ring.window(), ring.base()))
5745 .unwrap_or(0);
5746 let write_row = e.prepare_kv_append(kv, retain_from, t)?;
5747 for i in 0..t {
5748 let k_row = k.slice(i * kv.kv_dim_k..(i + 1) * kv.kv_dim_k);
5749 let v_row = v.slice(i * kv.kv_dim_v..(i + 1) * kv.kv_dim_v);
5750 e.append_kv_quantized_view(
5751 &k_row,
5752 &v_row,
5753 &mut kv.k,
5754 &mut kv.v,
5755 write_row + i,
5756 kv.kv_dim_k,
5757 kv.kv_dim_v,
5758 kv.k_tok_bytes,
5759 kv.v_tok_bytes,
5760 false,
5761 )?;
5762 }
5763 kv.len = pos0 + t;
5764 e.set_i32_one(&mut kv.len_d, kv.len as i32)?;
5765 Ok(())
5766 }
5767
5768 #[allow(clippy::too_many_arguments)]
5769 fn mtp_kv_fill_all(
5770 &self,
5771 e: &Engine,
5772 tokens: &[u32],
5773 h: &CudaSlice<f32>,
5774 pos0: usize,
5775 scratch: &mut MtpScratch,
5776 embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
5777 ) -> Result<(), Box<dyn std::error::Error>> {
5778 debug_assert_eq!(self.mtp_head_count(), scratch.plane_count());
5779 for index in 0..self.mtp_head_count() {
5780 self.mtp_kv_fill_at(
5781 e,
5782 self.mtp_head_at(index),
5783 tokens,
5784 h,
5785 pos0,
5786 scratch,
5787 index,
5788 embd_dev,
5789 )?;
5790 }
5791 Ok(())
5792 }
5793
5794 /// CAPTURE body for the GRAPH DRAFT (stage 2 of graph-grade spec): ONE MTP head forward with
5795 /// every varying input device-resident —
5796 /// - token id from the persistent `tok_d` (the previous replay's in-graph argmax wrote it,
5797 /// so the chain feeds itself; the host reads the same 4 bytes for the draft list),
5798 /// - h_seed from the persistent `h_seed_d` (h_nextn is copied BACK into it at the end),
5799 /// - rope pos from the persistent `pos_d` counter (inc'd in-graph),
5800 /// - scratch KV slot/bound from `scratch.kv.len_d` (see mtp_full_attn_dc).
5801 /// The p-min confidence lands in the persistent `p_d` iff `with_prob` (env is fixed per run).
5802 /// Same kernels, same dispatch as the eager mtp_head_forward_dev chain -> same draft tokens
5803 /// (exactness never depends on drafts — the verify arbitrates — but acceptance parity does).
5804 /// `with_head=false` captures the HEAD-LESS twin for the pseudo-seed replay (2026-07-03):
5805 /// the pseudo pass only needs h_nextn (op 10) + the scratch append — the lm_head read
5806 /// (~1.06ms q6_K on the 9B), argmax and prob are dead weight there. h_nextn's inputs are
5807 /// untouched, so the seed value is identical; round-start resets overwrite tok_d/p_d anyway.
5808 /// `sampled_cap` = Some((ctr_d, perturb_d, q_out_d, seed, temp)) captures the SAMPLED twin
5809 /// (step 3 of the sampled-spec arc): head logits are retained in the persistent `q_out_d`
5810 /// (host D2Ds them to the round's q slot after each replay), the DEVICE event counter is
5811 /// bumped in-graph, and the argmax reads GUMBEL-PERTURBED logits — one categorical draw per
5812 /// replay, bit-identical to the eager arm's gumbel_perturb at the same (seed, sctr, temp).
5813 /// seed/temp are capture-time constants (fixed per generate call, like p_min).
5814 #[allow(clippy::too_many_arguments)]
5815 #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
5816 fn mtp_head_forward_cap(
5817 &self,
5818 e: &Engine,
5819 mtp: &MtpHead,
5820 tok_d: &mut CudaSlice<u32>,
5821 pos_d: &mut CudaSlice<i32>,
5822 h_seed_d: &mut CudaSlice<f32>,
5823 p_d: &mut CudaSlice<f32>,
5824 scratch: &mut MtpScratch,
5825 // Which scratch plane this head appends to / attends over: 0 for the single-head
5826 // chain (every pre-lane caller), the head's own plane index for the multi-head
5827 // chain graphs (each head owns one plane — `mtp_chain_forward_dev`'s contract).
5828 scratch_index: usize,
5829 with_prob: bool,
5830 with_head: bool,
5831 embd_gpu: &CudaSlice<u8>,
5832 embd_qt: i32,
5833 embd_rb: usize,
5834 d_vocab: usize,
5835 sampled_cap: Option<SampledCapArgs<'_>>,
5836 stream_pack: Option<(&mut CudaSlice<u32>, usize, Option<&CudaSlice<u32>>)>,
5837 // DRAFT-SIDE GRAMMAR MASK (lane/draft-mask): (packed draft-vocab allowed-set buffer,
5838 // word count). Captured as ONE mask_logits_f32 node between the head matmul and the
5839 // in-graph argmax; the buffer address is baked, its CONTENTS are re-uploaded by the
5840 // host before every replay (the decode.rs graph-mask pattern). All-ones contents = a
5841 // no-op ban, so a position the grammar cannot constrain costs one pass over the row.
5842 mask_cap: Option<(&CudaSlice<u32>, usize)>,
5843 ) -> Result<(), Box<dyn std::error::Error>> {
5844 let cfg = &self.cfg;
5845 let n_embd = cfg.n_embd as usize;
5846 // step35: capturable through the WINDOWED device-counter arm (`mtp_step35_attn_dcw`)
5847 // once the dcw door is armed and the v3-vec class is live. Without the door this stays
5848 // the deliberate, named refusal: the plain `_dc` attention's key bound always starts at
5849 // row 0, cannot express this block's SWA view offset, and a captured chain would
5850 // silently attend OUTSIDE the window once the persistent scratch passes 512 rows.
5851 // Returning Err (not a panic) is what the capture sites already handle by degrading to
5852 // the eager chain (`mtp_head_forward_dev` -> `mtp_step35_attn`).
5853 // ROUND-STREAM stays refused EITHER WAY: the stream VERIFY has no step35 twin (see the
5854 // step35_verify refusal), so a stream capture that succeeded here would only move the
5855 // failure from capture time (graceful stream-off) to serve time (a failed round).
5856 if let Some(g) = mtp.step35.as_ref() {
5857 if stream_pack.is_some() {
5858 return Err(
5859 "step35 has no ROUND-STREAM draft arm (the stream verify has no step35 \
5860 twin); stream off"
5861 .into(),
5862 );
5863 }
5864 if !self.step35_dcw_eligible(g, scratch.plane(scratch_index).1) {
5865 return Err(format!(
5866 "step35 has no captured draft chain (fa_decode_dc cannot express the MTP \
5867 block's SWA view offset; the windowed dcw capture needs \
5868 MEMRA_STEP35_DRAFT_DCW armed [default ON, =0 disarms] and the v3-vec \
5869 class live at bucket=min(window {}, scratch cap {})) - the eager draft \
5870 chain serves this shape",
5871 g.window,
5872 scratch.plane(scratch_index).1,
5873 )
5874 .into());
5875 }
5876 }
5877 // student inner width (see mtp_head_forward_dev) — interface dims stay n_embd.
5878 let di = mtp.geom.as_ref().map(|g| g.d_inner).unwrap_or(n_embd);
5879 let eps = cfg.rms_eps;
5880 let e_emb = e.embed_gather_device(embd_gpu, tok_d, n_embd, embd_qt, embd_rb)?;
5881 let mut e_norm = e.zeros(n_embd)?;
5882 e.rms_norm(&e_emb, mtp.enorm.float_data(), &mut e_norm, n_embd, 1, eps)?;
5883 let mut h_norm = e.zeros(n_embd)?;
5884 e.rms_norm(
5885 &*h_seed_d,
5886 mtp.hnorm.float_data(),
5887 &mut h_norm,
5888 n_embd,
5889 1,
5890 eps,
5891 )?;
5892 let mut concat = e.zeros(2 * n_embd)?;
5893 e.copy_into(&mut concat, 0, &e_norm, n_embd)?;
5894 e.copy_into(&mut concat, n_embd, &h_norm, n_embd)?;
5895 let inp_sa = e.matmul(&mtp.eh_proj, &concat, 1)?;
5896 let mut a_norm = e.zeros(di)?;
5897 e.rms_norm(&inp_sa, mtp.attn_norm.float_data(), &mut a_norm, di, 1, eps)?;
5898 let attn_out = match (&mtp.mixer, mtp.step35.as_ref()) {
5899 // step35 (eligibility already enforced by the refusal above): the windowed dcw
5900 // arm, the SAME launcher the eager dev arm runs when the door is armed. No host
5901 // work here (this is the capture body); headroom is the callers' pre-arm.
5902 (Mixer::Full(fa), Some(g)) => {
5903 self.mtp_step35_attn_dcw(e, fa, g, &a_norm, pos_d, scratch, scratch_index)?
5904 }
5905 (Mixer::Full(fa), None) => self.mtp_full_attn_dc(
5906 e,
5907 fa,
5908 &a_norm,
5909 pos_d,
5910 scratch,
5911 scratch_index,
5912 mtp.geom.as_ref(),
5913 )?,
5914 (Mixer::Linear(_), _) => {
5915 panic!("MTP block is full-attn in qwen35; linear MTP not supported")
5916 }
5917 (Mixer::Mla(_), _) => {
5918 crate::hybrid::mla_path_unimplemented("captured MTP head forward")
5919 }
5920 (Mixer::Kda(_), _) => {
5921 crate::hybrid::kda_path_unimplemented("captured MTP head forward")
5922 }
5923 };
5924 let mut x1 = e.zeros(di)?;
5925 e.add(&inp_sa, &attn_out, &mut x1, di)?;
5926 let mut z = e.zeros(di)?;
5927 e.rms_norm(&x1, mtp.post_attn_norm.float_data(), &mut z, di, 1, eps)?;
5928 let ffn_out = match &mtp.ffn {
5929 crate::hybrid::Ffn::Dense {
5930 ffn_gate,
5931 ffn_up,
5932 ffn_down,
5933 } => {
5934 let n_ff = ffn_gate.out_features();
5935 let (gate, up) = if e.uses_q8_1_fast(ffn_gate) && e.uses_q8_1_fast(ffn_up) {
5936 let (zq, zd) = e.quantize_q8_1(&z, 1, di)?;
5937 (
5938 e.matmul_pre(ffn_gate, &zq, &zd, &z, 1)?,
5939 e.matmul_pre(ffn_up, &zq, &zd, &z, 1)?,
5940 )
5941 } else {
5942 (e.matmul(ffn_gate, &z, 1)?, e.matmul(ffn_up, &z, 1)?)
5943 };
5944 let mut act = e.zeros(n_ff)?;
5945 // step35: the dense FFN reads the per-layer SHEXP clamp, resolved for the MTP
5946 // block's own index (the mtp_head_forward_dev rule; None for every other arch,
5947 // which is `ffn_act`'s dispatch verbatim). The eager and captured chains must
5948 // run the ONE activation program.
5949 Self::ffn_act_lim(
5950 e,
5951 &self.cfg,
5952 &gate,
5953 &up,
5954 1.0,
5955 1.0,
5956 mtp.step35
5957 .as_ref()
5958 .and_then(|s| s.clamp_shexp)
5959 .map(SwigluClamp::Post),
5960 &mut act,
5961 n_ff,
5962 )?;
5963 e.matmul(ffn_down, &act, 1)?
5964 }
5965 // ROUND-STREAM: a softmax-routed resident MoE takes the zero-D2H device router +
5966 // expert program and is capture-legal. Sigmoid-routed MoE (Hy3/M3/Step) still
5967 // selects through the host-visible sigmoid router; capturing that stream sync
5968 // invalidates CUDA capture, so it stays on the eager draft chain even when every
5969 // expert is resident. Non-resident (SLRU-lock) is likewise rejected.
5970 crate::hybrid::Ffn::Moe(m)
5971 if m.dev_exps.is_some() && self.cfg.sigmoid_router().is_none() =>
5972 {
5973 self.moe_ffn_il(e, m, &z, 1, u16::MAX)?
5974 }
5975 crate::hybrid::Ffn::Moe(_) => {
5976 return Err(
5977 "graph draft requires a Dense or device-routed resident-MoE MTP FFN".into(),
5978 );
5979 }
5980 };
5981 let mut h_inner = e.zeros(di)?;
5982 e.add(&x1, &ffn_out, &mut h_inner, di)?;
5983 // student: up-project back to n_embd (carrier + head input; see mtp_head_forward_dev).
5984 let h_nextn = match mtp.geom.as_ref() {
5985 Some(g) => e.matmul(&g.out_up, &h_inner, 1)?,
5986 None => h_inner,
5987 };
5988 // MEMRA_SPEC_HPOST needs final_h even head-less (it IS the next seed under that convention).
5989 let final_h = if with_head || spec_hpost() {
5990 let final_norm = mtp.shared_head_norm.as_ref().unwrap_or(&self.output_norm);
5991 let mut fh = e.zeros(n_embd)?;
5992 e.rms_norm(&h_nextn, final_norm.float_data(), &mut fh, n_embd, 1, eps)?;
5993 Some(fh)
5994 } else {
5995 None
5996 };
5997 if with_head {
5998 let head = mtp.shared_head_head.as_ref().unwrap_or(&self.output);
5999 let mut logits = e.matmul(head, final_h.as_ref().unwrap(), 1)?;
6000 // DRAFT-SIDE GRAMMAR MASK: ban the grammar-illegal draft ids IN the captured chain,
6001 // before the argmax — proposals become legal by construction. Contents-only
6002 // per-replay upload keeps the capture valid.
6003 if let Some((mask_d, mw)) = mask_cap {
6004 e.mask_logits_col(&mut logits, mask_d, 0, d_vocab, mw)?;
6005 }
6006 if let Some(SampledCapArgs {
6007 ctr: ctr_d,
6008 perturb: perturb_d,
6009 q_out: q_out_d,
6010 seed,
6011 temp,
6012 filt,
6013 }) = sampled_cap
6014 {
6015 // SAMPLED chain: retain q (raw head logits -> persistent q_out_d; the matmul's
6016 // own buffer is pool-recycled after the capture body returns, so it can't be the
6017 // retention target), bump the device event counter, gumbel-perturb reading it,
6018 // and argmax the PERTURBED logits into tok_d — the in-graph categorical draw.
6019 e.copy_into(q_out_d, 0, &logits, d_vocab)?;
6020 e.sctr_inc(ctr_d)?;
6021 match filt {
6022 // PURE-TEMP: gumbel over the raw softmax — byte-identical to the
6023 // pre-lane capture body.
6024 None => e.gumbel_perturb_ctr(&logits, perturb_d, d_vocab, seed, ctr_d, temp)?,
6025 // FILTERED (lane/step37-draft-graph-serving-20260830): the SAME
6026 // filter_stats program the eager arm and the accept path run (the
6027 // wrapper's coop/plain choice is deployment-keyed, never per-call), then
6028 // the device-stat/device-counter perturb twin — the draft draws from the
6029 // exact filtered distribution the verify gathers `q` from. q was
6030 // retained ABOVE, pre-perturb, so the accept path's post-replay stats
6031 // recompute (same kernel, same bits) reconstructs these th/z exactly.
6032 Some(f) => {
6033 e.filter_stats(
6034 &logits, d_vocab, f.rows0, f.th, f.z, f.mx, d_vocab, 1, temp, f.top_k,
6035 f.top_p, f.min_p,
6036 )?;
6037 e.gumbel_perturb_filtered_ctr(
6038 &logits, perturb_d, d_vocab, seed, ctr_d, temp, f.mx, f.th,
6039 )?;
6040 }
6041 }
6042 e.argmax_token_device_into(perturb_d, tok_d, d_vocab)?;
6043 // p-min prob = the head's RAW softmax confidence in the SAMPLED pick — same
6044 // semantics as the eager sampled arm's prob_of_token_device(dl_d, tok_d).
6045 if with_prob {
6046 e.prob_of_token_device_into(&logits, tok_d, p_d, d_vocab)?;
6047 }
6048 } else {
6049 // draft token -> persistent tok_d (next replay's embed reads it; host reads the 4 bytes).
6050 e.argmax_token_device_into(&logits, tok_d, d_vocab)?;
6051 // p-min under a draft mask reads the MASKED row: confidence relative to the
6052 // grammar-LEGAL alternatives (illegal ids leave the softmax denominator), which
6053 // is the right semantics for "does the drafter know what comes next here" and
6054 // the same row the pick came from. Draft-quality only — verify arbitrates.
6055 if with_prob {
6056 e.prob_of_token_device_into(&logits, tok_d, p_d, d_vocab)?;
6057 }
6058 }
6059 }
6060 // ROUND-STREAM K-chain: pack (tok, p) into slot j, then remap tok through d2t so the
6061 // NEXT chained body's embed reads the TARGET id — zero host involvement per step.
6062 if let Some((out, slot, d2t)) = stream_pack {
6063 e.pack_tok_p(tok_d, p_d, out, slot)?;
6064 if let Some(map) = d2t {
6065 e.tok_map_u32(tok_d, map)?;
6066 }
6067 }
6068 // Next draft step's h_seed: pre-norm h_nextn (default) or post-norm final_h (HPOST).
6069 if spec_hpost() {
6070 e.copy_into(h_seed_d, 0, final_h.as_ref().unwrap(), n_embd)?;
6071 } else {
6072 e.copy_into(h_seed_d, 0, &h_nextn, n_embd)?;
6073 }
6074 // advance the draft rope position in-graph.
6075 e.inc_seqlen(pos_d)?;
6076 Ok(())
6077 }
6078
6079 /// Batched target verify forward over `tokens` at positions `pos0..pos0+T` (§D.3, T=K+1).
6080 /// Returns ALL T logit columns (host f32, [T*n_vocab]); appends T cols to every full-attn KV
6081 /// and advances every linear-attn recur state by T steps (the recur steps are SEQUENTIAL T=1).
6082 /// Advances `cache.pos` by T.
6083 pub fn decode_step_t(
6084 &self,
6085 e: &Engine,
6086 tokens: &[u32],
6087 pos0: usize,
6088 cache: &mut Cache,
6089 ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
6090 if self.is_gemma4_e4b() {
6091 return Ok(self.gemma4_e4b_decode_step_t_h(e, tokens, pos0, cache)?.0);
6092 }
6093 if self.gemma_batch_program() {
6094 return self.gemma4_decode_step_t(e, tokens, pos0, cache);
6095 }
6096 Ok(self.decode_step_t_h(e, tokens, pos0, cache)?.0)
6097 }
6098
6099 /// Like `decode_step_t` but ALSO returns the LAST column's pre-output_norm hidden (h_seed for
6100 /// the next draft round). This lets partial-accept replay run as ONE batched T=(n_acc+1) forward
6101 /// (single weight read) instead of n_acc+1 separate T=1 decode_steps (n_acc+1 weight reads).
6102 /// At batch=1 decode is bandwidth-bound, so batching the replay is THE MTP profitability lever.
6103 pub fn decode_step_t_h(
6104 &self,
6105 e: &Engine,
6106 tokens: &[u32],
6107 pos0: usize,
6108 cache: &mut Cache,
6109 ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
6110 self.decode_step_t_h_emb(e, tokens, pos0, cache, None)
6111 }
6112
6113 /// Like `decode_step_t_h` with an optional RESIDENT embed table (spec hot loop): device
6114 /// gather instead of host dequant + [T, n_embd] f32 htod. Bit-identical rows.
6115 pub fn decode_step_t_h_emb(
6116 &self,
6117 e: &Engine,
6118 tokens: &[u32],
6119 pos0: usize,
6120 cache: &mut Cache,
6121 embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
6122 ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
6123 let (logits_d, h_seed) = self.decode_step_t_h_emb_dev(e, tokens, pos0, cache, embd_dev)?;
6124 Ok((e.dtoh(&logits_d)?, h_seed))
6125 }
6126
6127 /// DEVICE-LOGITS verify forward (spec device-argmax lever): identical kernel chain to
6128 /// `decode_step_t_h_emb` but returns the [T, n_vocab] logits ON DEVICE — the accept walk
6129 /// argmaxes each column on-device and reads back ONE [T] u32 instead of dtoh'ing the full
6130 /// T x n_vocab f32 block (~1-4 MB + T host argmaxes, every round). Kernel dispatch is
6131 /// UNCHANGED (same decode-exact kernels); only the post-logits transfer moves.
6132 pub fn decode_step_t_h_emb_dev(
6133 &self,
6134 e: &Engine,
6135 tokens: &[u32],
6136 pos0: usize,
6137 cache: &mut Cache,
6138 embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
6139 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
6140 cache.ensure_usable("decode_step_t")?;
6141 let n_embd = self.cfg.n_embd as usize;
6142 let t = tokens.len();
6143 let (logits, x) = self.decode_step_t_core(e, tokens, pos0, cache, embd_dev, None)?;
6144 // h_seed for the next round = LAST column's pre-output_norm hidden ([n_embd]).
6145 let mut hs = vbuf(e, n_embd)?; // fully written by copy_view_into below
6146 e.copy_view_into(&mut hs, 0, &x.slice((t - 1) * n_embd..t * n_embd), n_embd)?;
6147 Ok((logits, hs))
6148 }
6149
6150 /// CORE verify forward: the `decode_step_t_h_emb_dev` kernel chain, returning the FULL
6151 /// pre-output_norm hidden stack x ([T, n_embd], any column extractable) and optionally
6152 /// filling a `VerifyCkpt` (retained per-layer state-rebuild inputs) for the REPLAY-FREE
6153 /// partial accept. `ckpt: None` => byte-for-byte the old behavior (the ckpt writes are pure
6154 /// retains/copies — they never change what any kernel computes).
6155 fn decode_step_t_core(
6156 &self,
6157 e: &Engine,
6158 tokens: &[u32],
6159 pos0: usize,
6160 cache: &mut Cache,
6161 embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
6162 mut ckpt: Option<&mut VerifyCkpt>,
6163 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
6164 self.decode_step_t_core_stream(
6165 e,
6166 tokens,
6167 pos0,
6168 cache,
6169 embd_dev,
6170 ckpt.take(),
6171 None,
6172 None,
6173 None,
6174 None,
6175 )
6176 }
6177
6178 /// [`Self::decode_step_t_core`] with the MTP route's verify-graph pool armed
6179 /// (`MEMRA_SPEC_VERIFY_GRAPH`). `graphs: None` reproduces `decode_step_t_core`
6180 /// argument-for-argument, so the eager walk stays the byte-identical fallback.
6181 #[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
6182 fn decode_step_t_core_vg(
6183 &self,
6184 e: &Engine,
6185 tokens: &[u32],
6186 pos0: usize,
6187 cache: &mut Cache,
6188 embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
6189 mut ckpt: Option<&mut VerifyCkpt>,
6190 graphs: Option<&mut DsparkVerifyGraphs>,
6191 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
6192 self.decode_step_t_core_stream(
6193 e,
6194 tokens,
6195 pos0,
6196 cache,
6197 embd_dev,
6198 ckpt.take(),
6199 None,
6200 None,
6201 None,
6202 graphs,
6203 )
6204 }
6205
6206 /// Increment-0 two-session PP seam: release the peer after this lane's stage-0 boundary TX.
6207 /// The two independent sessions keep their own cache/checkpoint state; only issue order moves.
6208 #[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
6209 fn decode_step_t_core_pipelined(
6210 &self,
6211 e: &Engine,
6212 tokens: &[u32],
6213 pos0: usize,
6214 cache: &mut Cache,
6215 embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
6216 mut ckpt: Option<&mut VerifyCkpt>,
6217 pipe: &SpecPipeLane,
6218 round: usize,
6219 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
6220 let fence = crate::pp::pp_cuts(self.layers.len())
6221 .ok_or("two-session speculative pipeline requires a PP stage cut")?;
6222 if crate::pp::pp2_streams_off() || !crate::pp::spec_pp_on() {
6223 return Err("two-session speculative pipeline requires the PP verify split".into());
6224 }
6225 let interval_fence = pipe.stage0_begin(round)?;
6226 let _walk = pipe.coordinated_walk()?;
6227 let ticket = self.verify_stage0_issue(
6228 e,
6229 tokens,
6230 pos0,
6231 cache,
6232 embd_dev,
6233 ckpt.as_deref_mut(),
6234 None,
6235 &fence,
6236 Some(interval_fence),
6237 pipe.trace(round),
6238 )?;
6239 pipe.stage0_end(round);
6240 pipe.stage1_begin(round)?;
6241 let result = self.verify_stage1_finish(e, ticket, cache, ckpt, None, &fence, true)?;
6242 pipe.verify_end(round);
6243 Ok(result)
6244 }
6245
6246 /// ROUND-STREAM stage (c) 4: `stream` = (device verify tokens [t], device pos counter) —
6247 /// when Some, rope positions come from pos_iota over the counter, the embed gathers the
6248 /// device tokens, and full_attn_verify routes appends/FA through the _dc twins reading the
6249 /// SAME counter (every layer's kvl.len == cache.pos, one counter drives all three). The
6250 /// host `tokens`/`pos0` args still size buffers (t is FIXED K+1 in stream mode).
6251 /// `vtok_dev` (engine-bundle slice 2): device verify tokens for the EMBED only —
6252 /// unlike `stream` mode it changes nothing else (host pos iota, host-len KV appends).
6253 /// `tokens` then only sizes buffers (the dummy-slice pattern the round-stream arm uses).
6254 #[allow(clippy::too_many_arguments)]
6255 fn decode_step_t_core_stream(
6256 &self,
6257 e: &Engine,
6258 tokens: &[u32],
6259 pos0: usize,
6260 cache: &mut Cache,
6261 embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
6262 mut ckpt: Option<&mut VerifyCkpt>,
6263 stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
6264 pp_pipe: Option<bool>,
6265 vtok_dev: Option<&CudaSlice<u32>>,
6266 graphs: Option<&mut DsparkVerifyGraphs>,
6267 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
6268 // PP DOOR (lane/pp2-spec 2026-08-06): the verify trunk now takes its OWN stage split,
6269 // exactly as the eager and batched steps do. This is the single funnel every verify
6270 // forward reaches (decode_step_t / _h / _h_emb / _h_emb_dev / _core all land here), so
6271 // wiring it here wires the whole spec surface — the draft/accept/commit machinery above
6272 // is untouched.
6273 //
6274 // History: pp2-hardening (2026-08-06) made this funnel FAIL CLOSED, because its trunk
6275 // walk was unsplit on one stream and a sharded cross-device placement peer-read every
6276 // remote layer's weights on every spec round (measured 13.9-28x on the batched twin).
6277 // The refusal below survives to cover the residue — MEMRA_SPEC_PP=0, MEMRA_PP_STREAMS=0,
6278 // or a placement whose PpNRt fails to build — so a config that would still walk the
6279 // whole trunk on one stream refuses instead of regressing 28x.
6280 if let Some(fence) = crate::pp::pp_cuts(self.layers.len())
6281 && !crate::pp::pp2_streams_off()
6282 && crate::pp::spec_pp_on()
6283 {
6284 if vtok_dev.is_some() {
6285 return Err(
6286 "device-token dspark verify (slice-2 deferred readback) has no PP \
6287 stage-split arm; set MEMRA_DSPARK_DEFER_READBACK=0 or run the dspark \
6288 route on one device"
6289 .into(),
6290 );
6291 }
6292 return self.decode_step_t_core_ppn(
6293 e,
6294 tokens,
6295 pos0,
6296 cache,
6297 embd_dev,
6298 ckpt.take(),
6299 stream,
6300 &fence,
6301 pp_pipe,
6302 );
6303 }
6304 crate::pp::refuse_unsplit_if_remote(
6305 "decode_step_t (spec verify)",
6306 "drop MEMRA_SPEC_PP=0 / MEMRA_PP_STREAMS=0 so the verify trunk takes its OWN stage \
6307 split (decode_step_t_core_ppn); or run spec on one device",
6308 )?;
6309 let cfg = &self.cfg;
6310 let n_embd = cfg.n_embd as usize;
6311 let eps = cfg.rms_eps;
6312 let t = tokens.len();
6313 let pos_d = match stream {
6314 Some((_, ctr)) => {
6315 let mut p = e.alloc_uninit::<i32>(t)?;
6316 e.pos_iota(ctr, &mut p, t)?;
6317 p
6318 }
6319 None => {
6320 let pos_vec: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
6321 e.htod_i32(&pos_vec)?
6322 }
6323 };
6324
6325 // embed T tokens -> [T, n_embd] token-major (device gather on the spec hot loop)
6326 let x = match (stream, embd_dev) {
6327 (Some((vtok, _)), Some((g, qt, rb))) => {
6328 e.embed_gather_device_td(g, vtok, t, n_embd, qt, rb)?
6329 }
6330 (None, Some((g, qt, rb))) => match vtok_dev {
6331 // slice 2: device verify tokens, same embed_gather_u32_t kernel —
6332 // bit-identical rows to the host-token arm (same per-dtype deq).
6333 Some(vt_d) => e.embed_gather_device_td(g, vt_d, t, n_embd, qt, rb)?,
6334 None => e.embed_gather_device_t(g, tokens, n_embd, qt, rb)?,
6335 },
6336 _ => {
6337 assert!(
6338 vtok_dev.is_none(),
6339 "device-token verify requires the resident embed table (embd_dev)"
6340 );
6341 e.htod(&self.embd.gather(n_embd, tokens))?
6342 }
6343 };
6344
6345 // TRUNK WALK: layers [0, n_layers) through the SAME range-scoped subgraph the PP-N
6346 // stage split calls per stage (`verify_layers`) — one code path, so the split cannot
6347 // drift from the unsplit dispatch mirroring. lane/pp2-spec 2026-08-06.
6348 let x = self.verify_layers(
6349 e,
6350 x,
6351 0,
6352 self.layers.len(),
6353 &pos_d,
6354 pos0,
6355 t,
6356 cache,
6357 ckpt.take(),
6358 stream,
6359 graphs,
6360 )?;
6361 if spec_nan_scan() {
6362 nan_scan_rows(e, &x, t, n_embd, &format!("verify trunk exit pos0={pos0}"))?;
6363 }
6364
6365 let mut hn = vbuf(e, t * n_embd)?;
6366 // Stage-A door: with the serving-class row-outer verify walk, the TAIL must be the
6367 // t=1 decode program per row too (rms_norm t=1 + the single-row bf16 head — the
6368 // split head's concat is receipted bit-identical to it). The batched cuBLASLt head
6369 // is a different ULP class and flips near-tie argmaxes off the greedy tape.
6370 let eager_tail = self.sliding_gated_moe_batch_program() && spec_verify_eager_on();
6371 if eager_tail {
6372 let n_vocab = self.cfg.n_vocab as usize;
6373 // MEMRA_SPEC_HEAD_ROWS=1 — THE VERIFY TAIL'S REDUNDANT HEAD READ.
6374 //
6375 // The loop below runs the head at m=1 once PER COLUMN, so the LM head's weights are
6376 // streamed t times per verify pass. On step37 that head is ~0.49 GiB per card after the
6377 // rank split, ~1.07 ms of pure re-read at t=2 and worse at every wider t — which is a
6378 // large part of why the fixed K ladder LOSES (K=1 81.2 > K=2 73.1 > K=3 62.7 tok/s).
6379 //
6380 // The loop's justification is the comment above: the batched cuBLASLt head is a
6381 // different ULP class and flips near-tie argmaxes off the greedy tape. That is true of
6382 // cuBLASLt and it does NOT apply here, because a FloatBf16 head at 1..=32 rows never
6383 // reaches cuBLASLt: `matmul` routes it to `matvec_bf16_rows_into` (lib.rs:12248), whose
6384 // own doc says `matvec_bf16_f32acc_x4_rows` "runs the t=1 decode head program PER ROW
6385 // (identical dot + reduce), so decode/verify tiers keep the t=1 numeric class". Under
6386 // the W8 doors both widths route to the q8 mirror instead, and the t-column mirror is
6387 // documented "bit-identical to t single-row calls". So the batched form is the SAME
6388 // arithmetic per row on both paths, with one weight read instead of t.
6389 //
6390 // rms_norm is row-wise, so norm(t) is per-row identical to t x norm(1) by construction.
6391 //
6392 // DEFAULT OFF for exactly one turn of the crank: "bit-identical by two documented
6393 // claims" is still an argument. The greedy byte tape decides, and the door flips only
6394 // once the tape is a receipt.
6395 if head_rows_on() {
6396 e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
6397 let logits = e.matmul(&self.output, &hn, t)?;
6398 if stream.is_none() {
6399 cache.pos += t;
6400 }
6401 return Ok((logits, if spec_hpost() { hn } else { x }));
6402 }
6403 let mut logits = vbuf(e, t * n_vocab)?;
6404 for r in 0..t {
6405 let mut row = e.uninit(n_embd)?;
6406 e.dtod_copy_view(&x.slice(r * n_embd..(r + 1) * n_embd), &mut row)?;
6407 let mut hr = e.uninit(n_embd)?;
6408 e.rms_norm(&row, self.output_norm.float_data(), &mut hr, n_embd, 1, eps)?;
6409 let lr = e.matmul(&self.output, &hr, 1)?;
6410 e.dtod_copy_into(&lr, &mut logits, r * n_vocab)?;
6411 e.dtod_copy_into(&hr, &mut hn, r * n_embd)?;
6412 }
6413 if stream.is_none() {
6414 cache.pos += t;
6415 }
6416 return Ok((logits, if spec_hpost() { hn } else { x }));
6417 }
6418 let serving_head =
6419 self.sliding_gated_moe_batch_program() || self.batched_serving_numeric_class();
6420 let logits = if serving_head {
6421 // Step35 and the qwen35 family (MoE 2026-08-14 AM, dense-hybrid same day PM — the
6422 // Q3.8 bring-up reproduced the identical near-tie class on dense: eager-class verify
6423 // vs batched-class live serving, ULP drift amplified through the GDN recurrence)
6424 // serve one batched numeric class at every live width, including B=1. Keep the
6425 // verify head in that same class; other generic families retain the decode-exact
6426 // head that their run-spec contract pins.
6427 e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
6428 e.matmul(&self.output, &hn, t)?
6429 } else {
6430 e.rms_norm_decode(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
6431 e.matmul_decode_exact(&self.output, &hn, t)?
6432 };
6433 // stream: the device pos counter owns position; host mirror reconciles at drain.
6434 if stream.is_none() {
6435 cache.pos += t;
6436 }
6437 // Hidden stack for seeds/refresh-fills: pre-norm x (default) or post-norm hn (HPOST).
6438 Ok((logits, if spec_hpost() { hn } else { x }))
6439 }
6440
6441 /// THE VERIFY TRUNK OVER PP-N (lane/pp2-spec 2026-08-06): `decode_step_t_core_stream`'s walk
6442 /// as N stage subgraphs, each on its own engine/stream (and, under `MEMRA_PP_DEVICES`, its own
6443 /// device), with a `[T, n_embd]` boundary transfer between them. T = K+1 (the verify batch),
6444 /// so this is the batched-boundary shape the pp2-batch lane's grow-only slots already handle
6445 /// (`tx(b, x, t*n_embd)`; the slot grows to the high-water T and the transport moves exactly
6446 /// the payload).
6447 ///
6448 /// Structure is `decode_step_batch_ppn`'s, which is `decode_step_h_ppn`'s. FOUR THINGS ARE
6449 /// PER-STAGE and each for a measured reason (see `decode_step_batch_ppn`'s header for the
6450 /// receipts):
6451 ///
6452 /// 1. THE ENGINE (`rt.engine(s, e)`) — `Engine` owns lazily-grown stable-pointer scratch
6453 /// (`fa_part_pool`, `fa_vf16_scratch`, `argmax_partials`) that is single-stream-safe BY
6454 /// DESIGN. Two stage streams through one Engine is the 2026-08-02 shared-scratch race
6455 /// (35% flake, nondeterministic all-logits divergence). `PpNRt::build` gives every stage
6456 /// s>0 its own Engine even on the primary device; honouring it here is what scopes the
6457 /// pools. The verify path allocates MORE of that scratch than eager decode does (FA at
6458 /// m=T, and the per-layer `GdnStash` retains), so this is load-bearing, not inherited.
6459 ///
6460 /// 2. `pos_d` — each stage uploads its OWN copy of the T rope positions on ITS stream, so the
6461 /// buffer is allocated, consumed and freed on one stream. In `stream` mode that means each
6462 /// stage runs its own `pos_iota` over the SHARED device counter (`pos_ctr`): the counter is
6463 /// read-only during the forward (the round's `inc`/`copy_add` happen outside it), so every
6464 /// stage derives the identical iota, and each stage's own output buffer is stream-local.
6465 ///
6466 /// 3. THE EMBED lives with stage 0 (`self.embd` / `embd_gpu` are host/primary-side; the
6467 /// sharded loader leaves the table with stage 0 by construction).
6468 ///
6469 /// 4. THE HEAD (`output_norm` + `output`) runs on the LAST stage — the sharded loader uploaded
6470 /// both through that stage's engine (`hybrid.rs`: `e_head = layer_engine(e, n_trunk,
6471 /// n_trunk-1)`), so reading them anywhere else is a peer read of the biggest tensor in the
6472 /// model, every round.
6473 ///
6474 /// WHAT STAYS ON THE PRIMARY, deliberately: the returned logits and hidden stack `x`. Both are
6475 /// last-stage-allocated device buffers, and every consumer (the device argmax walk, the accept
6476 /// kernels, `spec_seed_gather`, the ckpt rebuild in `commit_verified_prefix`) reads them
6477 /// through the primary context by UVA — the same read the batched serving epilogue's
6478 /// `last_logits_dev` park does. Those consumers are per-round O(T x n_vocab) and O(n_embd),
6479 /// not per-layer, so they are not the 28x class; splitting them is a separate lane.
6480 ///
6481 /// The MTP HEAD (draft side) is NOT split: it is one block, it lives wherever the loader put
6482 /// it (`load_mtp` uses the primary engine), and it is ~1-2 GB against the trunk's tens. Draft
6483 /// placement is measured, not assumed — see `research/pp2-spec-20260806`.
6484 ///
6485 /// EXACTNESS: PP-N adds ZERO deviation. Each stage runs the SAME kernels on the SAME bytes in
6486 /// the same order via the SAME `verify_layers` the unsplit body calls; the only change is
6487 /// where the residual is materialized, and the boundary is a straight f32 copy (dtod
6488 /// same-device / `cudaMemcpyPeerAsync` cross-device, no conversion). So the split MUST be
6489 /// BIT-IDENTICAL to the unsplit verify at the same T, in both placement orders. Gate:
6490 /// `decode-batch-gate --mode ppspec`. Acceptance counts are a DERIVED consequence — greedy
6491 /// accept argmaxes these logits, so bit-identical logits force identical accept walks; the
6492 /// `run-spec` K=1..8 arm checks that end-to-end rather than trusting the implication.
6493 #[allow(clippy::too_many_arguments)]
6494 fn decode_step_t_core_ppn(
6495 &self,
6496 e: &Engine,
6497 tokens: &[u32],
6498 pos0: usize,
6499 cache: &mut Cache,
6500 embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
6501 mut ckpt: Option<&mut VerifyCkpt>,
6502 stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
6503 fence: &[usize],
6504 pp_pipe: Option<bool>,
6505 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
6506 let ticket = self.verify_stage0_issue(
6507 e,
6508 tokens,
6509 pos0,
6510 cache,
6511 embd_dev,
6512 ckpt.as_deref_mut(),
6513 stream,
6514 fence,
6515 pp_pipe,
6516 None,
6517 )?;
6518 self.verify_stage1_finish(e, ticket, cache, ckpt, stream, fence, true)
6519 }
6520
6521 /// Enqueue embed, stage 0, and the first boundary TX, then return the actual boundary slot.
6522 /// The ordinary PP verify wrapper calls `verify_stage1_finish` immediately after this return.
6523 #[allow(clippy::too_many_arguments)]
6524 fn verify_stage0_issue(
6525 &self,
6526 e: &Engine,
6527 tokens: &[u32],
6528 pos0: usize,
6529 cache: &mut Cache,
6530 embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
6531 ckpt: Option<&mut VerifyCkpt>,
6532 stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
6533 fence: &[usize],
6534 pp_pipe: Option<bool>,
6535 trace: Option<SpecPipeTraceCtx>,
6536 ) -> Result<VerifyBoundaryTicket, Box<dyn std::error::Error>> {
6537 assert!(
6538 !self.is_gemma4_e4b() && !self.gemma_batch_program(),
6539 "decode_step_t_core_ppn covers the hybrid non-gemma4 verify trunk only \
6540 (the gemma4 arms have their own decode_step_t twins)"
6541 );
6542 if crate::pp::pp_host_bounce_active() && (stream.is_some() || embd_dev.is_some()) {
6543 return Err(
6544 "decode_step_t_core_ppn: refused with MEMRA_PP_HOST_BOUNCE=1 — the trunk \
6545 boundary itself is host-staged, but device-resident verify still peer-reads \
6546 primary-device token/position/embedding buffers from stage 0. Run plain PP \
6547 serving on this host class; spec requires local per-stage inputs first."
6548 .into(),
6549 );
6550 }
6551 let rt = crate::pp::PpNRt::get(e)?;
6552 // Pipelined callers do not bypass ownership: their explicit coordinator borrow makes
6553 // this acquire clone the same active generation. Ordinary callers acquire a fresh lease.
6554 let walk_owner = rt.acquire_walk("verify_stage0_issue")?;
6555 let n_st = fence.len() - 1;
6556 assert_eq!(
6557 rt.n_stages(),
6558 n_st,
6559 "PpNRt stage count {} != fence stages {n_st}",
6560 rt.n_stages()
6561 );
6562 let n_embd = self.cfg.n_embd as usize;
6563 let t = tokens.len();
6564 let payload = t * n_embd;
6565 if pp_pipe.is_some() {
6566 assert_eq!(n_st, 2, "spec pipeline requires exactly two PP stages");
6567 }
6568 // One-shot lane diagnostic: force natural PP-2 boundaries to completion so the server
6569 // log can price stage 0, the peer hop, the RX copy, and stage 1 + head separately without
6570 // nsys. The ordinary path keeps every enqueue asynchronous. N>2 is deliberately excluded:
6571 // the report below names exactly two stages and must never imply it measured middle ones.
6572 let pp_anatomy = n_st == 2 && std::env::var("MEMRA_SPEC_PP_ANATOMY").as_deref() == Ok("1");
6573 let pp_started = std::time::Instant::now();
6574 let (mut reverse_ms, mut stage0_ms, mut tx_ms) = (0.0f64, 0.0f64, 0.0f64);
6575 // The CALLER's ambient stream, captured BEFORE any `rt.enter()` pushes a stage stream:
6576 // this body returns DEVICE-RESIDENT buffers (the device-argmax accept walk's contract),
6577 // so the exit needs the same publication the boundaries get — see `PpNRt::publish_to`.
6578 // Taken here, not at the end, because inside the last-stage scope `e.stream()` IS the
6579 // stage stream and the wait would self-order into a no-op.
6580 let caller_stream = e.stream();
6581 // #87 ROOT-CAUSE FENCE (lane/pp2spec-crash): the PREVIOUS round's stage-allocated
6582 // outputs (logits/hidden/ckpt stashes) freed stream-ordered on the STAGE streams while
6583 // the primary stream still holds queued reads of them — with event tracking elided,
6584 // nothing stops the pool from reusing those blocks for THIS round's stage allocations,
6585 // whose writes then race the queued reads (measured: 13/4096-NaN random-bits garbage in
6586 // the spec round seed; the full anatomy is on `PpNRt::fence_stages_behind`). Order every
6587 // stage stream behind the caller before enqueueing new stage work.
6588 let reverse_started = std::time::Instant::now();
6589 if pp_pipe != Some(false) {
6590 rt.fence_stages_behind(&caller_stream)?;
6591 }
6592 if pp_pipe == Some(true) {
6593 // Both session verifies must alternate boundary slots even when the ordinary
6594 // decode overlap experiment is off. Prewarm before A's stage 0 so B cannot grow
6595 // slot 1 by synchronizing the RX stream while A's stage 1 is in flight.
6596 rt.prepare_overlap_slots(0, payload)?;
6597 }
6598 if pp_anatomy {
6599 // Drain the reverse-publication dependency before timing stage 0 itself. At c=1 this
6600 // prices any primary-stream rollback/refresh tail inherited from the prior round.
6601 for s in 0..n_st {
6602 let _st = rt.enter(s);
6603 rt.engine(s, e).stream().synchronize()?;
6604 }
6605 reverse_ms = reverse_started.elapsed().as_secs_f64() * 1e3;
6606 }
6607
6608 // Per-stage rope positions: in host mode the same [T] iota each stage uploads itself; in
6609 // stream mode each stage's own `pos_iota` over the shared read-only device counter.
6610 let stage_pos = |es: &Engine| -> Result<CudaSlice<i32>, Box<dyn std::error::Error>> {
6611 match stream {
6612 Some((_, ctr)) => {
6613 let mut p = es.alloc_uninit::<i32>(t)?;
6614 es.pos_iota(ctr, &mut p, t)?;
6615 Ok(p)
6616 }
6617 None => {
6618 let pos_vec: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
6619 es.htod_i32(&pos_vec)
6620 }
6621 }
6622 };
6623
6624 // ---- STAGE 0: embed (the table lives with stage 0) + layers [0, fence[1]) + TX ----
6625 let slot = {
6626 let _st0 = rt.enter(0);
6627 let e0 = rt.engine(0, e);
6628 enqueue_spec_pipe_trace_marker(&e0.stream(), trace.as_ref(), "S0", "start", None)?;
6629 let stage0_started = std::time::Instant::now();
6630 let pos_d = stage_pos(e0)?;
6631 let x = match (stream, embd_dev) {
6632 (Some((vtok, _)), Some((g, qt, rb))) => {
6633 e0.embed_gather_device_td(g, vtok, t, n_embd, qt, rb)?
6634 }
6635 (None, Some((g, qt, rb))) => e0.embed_gather_device_t(g, tokens, n_embd, qt, rb)?,
6636 _ => e0.htod(&self.embd.gather(n_embd, tokens))?,
6637 };
6638 let x = self.verify_layers(
6639 e0, x, fence[0], fence[1], &pos_d, pos0, t, cache, ckpt, stream, None,
6640 )?;
6641 if pp_anatomy {
6642 e0.stream().synchronize()?;
6643 stage0_ms = stage0_started.elapsed().as_secs_f64() * 1e3;
6644 }
6645 let tx_started = std::time::Instant::now();
6646 let slot = if pp_pipe.is_some() {
6647 rt.tx_pipelined(0, &x, payload)?
6648 } else {
6649 rt.tx(0, &x, payload)?
6650 };
6651 enqueue_spec_pipe_trace_marker(&e0.stream(), trace.as_ref(), "S0", "end", Some(slot))?;
6652 if pp_anatomy {
6653 e0.stream().synchronize()?;
6654 tx_ms = tx_started.elapsed().as_secs_f64() * 1e3;
6655 }
6656 slot
6657 // x + pos_d drop here: freed stream-ordered on stage-0's stream after use.
6658 };
6659
6660 Ok(VerifyBoundaryTicket {
6661 rt,
6662 caller_stream,
6663 slot,
6664 pos0,
6665 t,
6666 payload,
6667 n_st,
6668 pipelined: pp_pipe.is_some(),
6669 pp_anatomy,
6670 pp_started,
6671 reverse_ms,
6672 stage0_ms,
6673 tx_ms,
6674 trace,
6675 _walk_owner: walk_owner,
6676 })
6677 }
6678
6679 /// Consume a stage-0 boundary ticket and enqueue the remaining PP stages plus the head.
6680 /// On PP-2 this is exactly stage 1; PP-N keeps its pre-existing middle-stage walk here.
6681 #[allow(clippy::too_many_arguments)]
6682 fn verify_stage1_finish(
6683 &self,
6684 e: &Engine,
6685 ticket: VerifyBoundaryTicket,
6686 cache: &mut Cache,
6687 mut ckpt: Option<&mut VerifyCkpt>,
6688 stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
6689 fence: &[usize],
6690 publish_to_caller: bool,
6691 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
6692 let VerifyBoundaryTicket {
6693 rt,
6694 caller_stream,
6695 slot,
6696 pos0,
6697 t,
6698 payload,
6699 n_st,
6700 pipelined,
6701 pp_anatomy,
6702 pp_started,
6703 reverse_ms,
6704 stage0_ms,
6705 tx_ms,
6706 trace,
6707 _walk_owner,
6708 } = ticket;
6709 let n_embd = self.cfg.n_embd as usize;
6710 let eps = self.cfg.rms_eps;
6711 let mut slot = slot;
6712 let (mut rx_ms, mut stage1_ms) = (0.0f64, 0.0f64);
6713 let stage_pos = |es: &Engine| -> Result<CudaSlice<i32>, Box<dyn std::error::Error>> {
6714 match stream {
6715 Some((_, ctr)) => {
6716 let mut p = es.alloc_uninit::<i32>(t)?;
6717 es.pos_iota(ctr, &mut p, t)?;
6718 Ok(p)
6719 }
6720 None => {
6721 let pos_vec: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
6722 es.htod_i32(&pos_vec)
6723 }
6724 }
6725 };
6726
6727 // ---- MIDDLE STAGES: RX boundary s-1 -> range -> TX boundary s ----
6728 for s in 1..n_st - 1 {
6729 let _st = rt.enter(s);
6730 let es = rt.engine(s, e);
6731 let pos_d = stage_pos(es)?;
6732 let x = rt.rx(s - 1, slot, payload)?;
6733 let x = self.verify_layers(
6734 es,
6735 x,
6736 fence[s],
6737 fence[s + 1],
6738 &pos_d,
6739 pos0,
6740 t,
6741 cache,
6742 ckpt.as_deref_mut(),
6743 stream,
6744 None,
6745 )?;
6746 slot = if pipelined {
6747 rt.tx_pipelined(s, &x, payload)?
6748 } else {
6749 rt.tx(s, &x, payload)?
6750 };
6751 }
6752
6753 // ---- LAST STAGE: RX + final range + output_norm + lm head ----
6754 let _stl = rt.enter(n_st - 1);
6755 let el = rt.engine(n_st - 1, e);
6756 let pos_d = stage_pos(el)?;
6757 let rx_started = std::time::Instant::now();
6758 let x = rt.rx(n_st - 2, slot, payload)?;
6759 if pp_anatomy {
6760 el.stream().synchronize()?;
6761 rx_ms = rx_started.elapsed().as_secs_f64() * 1e3;
6762 }
6763 enqueue_spec_pipe_trace_marker(&el.stream(), trace.as_ref(), "S1", "start", Some(slot))?;
6764 let stage1_started = std::time::Instant::now();
6765 let x = self.verify_layers(
6766 el,
6767 x,
6768 fence[n_st - 1],
6769 fence[n_st],
6770 &pos_d,
6771 pos0,
6772 t,
6773 cache,
6774 ckpt,
6775 stream,
6776 None,
6777 )?;
6778
6779 let mut hn = vbuf(el, payload)?;
6780 let logits = if self.sliding_gated_moe_batch_program() {
6781 // The PP Step35 serving path uses rms_norm + matmul for B=1 as well as B>1.
6782 // Verify must not switch numeric class merely because the same session speculates.
6783 el.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
6784 el.matmul(&self.output, &hn, t)?
6785 } else {
6786 el.rms_norm_decode(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
6787 el.matmul_decode_exact(&self.output, &hn, t)?
6788 };
6789 enqueue_spec_pipe_trace_marker(&el.stream(), trace.as_ref(), "S1", "end", Some(slot))?;
6790 if pp_anatomy {
6791 el.stream().synchronize()?;
6792 stage1_ms = stage1_started.elapsed().as_secs_f64() * 1e3;
6793 }
6794 // EXIT PUBLICATION: both returned buffers are still being produced on the last stage's
6795 // stream. Order the caller's stream behind that work before the buffers escape this
6796 // scope (the 2026-08-06 same-device ppspec find: without it the caller's primary-stream
6797 // consumer read unwritten logits — nondeterministic, one-device-only, and it poisoned
6798 // the following arm's KV in the same process).
6799 if publish_to_caller {
6800 rt.publish_to(n_st - 1, &caller_stream)?;
6801 }
6802 if pp_anatomy {
6803 if publish_to_caller {
6804 caller_stream.synchronize()?;
6805 }
6806 eprintln!(
6807 "[spec-pp-anatomy] t={t} reverse={reverse_ms:.3}ms stage0={stage0_ms:.3}ms \
6808 tx={tx_ms:.3}ms rx={rx_ms:.3}ms stage1-head={stage1_ms:.3}ms total={:.3}ms",
6809 pp_started.elapsed().as_secs_f64() * 1e3,
6810 );
6811 }
6812 // stream: the device pos counter owns position; host mirror reconciles at drain.
6813 if stream.is_none() {
6814 cache.pos += t;
6815 }
6816 Ok((logits, if spec_hpost() { hn } else { x }))
6817 }
6818
6819 /// Step3.5/Step3.7 verify trunk in the serving batched numeric class.
6820 ///
6821 /// `step35_decode_batch_layers` is now authoritative at every live serving width, including
6822 /// B=1 (lane/cx-b1fix). The older verify walk deliberately mirrored the eager T=1 class:
6823 /// it replayed `step35_decode_attn` per row and used the eager/decode-exact FFN dispatch.
6824 /// Those classes are individually stable, but a near-tie prompt can choose different greedy
6825 /// bytes when a request moves from batched plain serving into speculative verify. Run the
6826 /// same authoritative B=1 stage subgraph for each verify row here. Rows still advance
6827 /// layer-by-layer, so every layer sees the preceding verify rows in its attention cache while
6828 /// every norm/projection/FFN uses exactly the live serving dispatch.
6829 #[allow(clippy::too_many_arguments)]
6830 /// PRIME-BY-T-ROWS (MEMRA_PRIME_TROWS=1): prefill the prompt through the same-session
6831 /// t-row walk in 32-row chunks — every row runs the t=1 decode program bit-for-bit
6832 /// (the TOKENWISE-prime ORACLE class), so this door is exact against the exactness
6833 /// reference while replacing the host-canonical per-token prime. Requires the walk
6834 /// doors (MEMRA_SPEC_VERIFY_EAGER/TCOL); returns the prime contract trio.
6835 #[allow(clippy::type_complexity)]
6836 pub(crate) fn step35_prime_trows(
6837 &self,
6838 e: &Engine,
6839 tokens: &[u32],
6840 cache: &mut Cache,
6841 ) -> Result<Option<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
6842 {
6843 let dbg = std::env::var("MEMRA_SPEC_FA2_DEBUG").as_deref() == Ok("1");
6844 if !prime_trows_on() {
6845 return Ok(None);
6846 }
6847 if !self.uses_sliding_gated_moe_program()
6848 || cache.pos != 0
6849 || cache.dflash_taps.is_some()
6850 || !spec_verify_eager_on()
6851 || !spec_verify_tcol_on()
6852 {
6853 if dbg {
6854 eprintln!(
6855 "[prime-trows] refuse: program={} pos={} taps={} eager={:?} tcol={:?}",
6856 self.uses_sliding_gated_moe_program(),
6857 cache.pos,
6858 cache.dflash_taps.is_some(),
6859 std::env::var("MEMRA_SPEC_VERIFY_EAGER").ok(),
6860 std::env::var("MEMRA_SPEC_VERIFY_TCOL").ok()
6861 );
6862 }
6863 return Ok(None);
6864 }
6865 let n_embd = self.cfg.n_embd as usize;
6866 let n_layers = self.layers.len();
6867 let t_total = tokens.len();
6868 let Some(embd_gpu) = self.embd_gpu_try(e) else {
6869 if dbg {
6870 eprintln!("[prime-trows] refuse: no device embed table");
6871 }
6872 return Ok(None);
6873 };
6874 let embd_qtype = match self.embd.ggml_type {
6875 memra_gguf::GgmlType::BF16 => crate::QT_BF16,
6876 memra_gguf::GgmlType::Q8_0 => crate::QT_Q8_0,
6877 other => {
6878 if dbg {
6879 eprintln!("[prime-trows] refuse: embed dtype {other:?}");
6880 }
6881 return Ok(None);
6882 }
6883 };
6884 let embd_row_bytes = self.embd.raw.len() / self.cfg.n_vocab as usize;
6885 // Chunk plan: 32-row chunks; a 1-token tail folds into the previous chunk
6886 // (the walk floor is t >= 2).
6887 let mut bounds = Vec::new();
6888 let mut start = 0usize;
6889 while start < t_total {
6890 let mut end = (start + 32).min(t_total);
6891 if t_total - end == 1 {
6892 end -= 1;
6893 }
6894 bounds.push((start, end));
6895 start = end;
6896 }
6897 if bounds.iter().any(|(a, b)| b - a < 2) {
6898 return Ok(None); // degenerate short prompt keeps the ordinary prime
6899 }
6900 let mut hiddens = e.uninit(t_total * n_embd)?;
6901 let mut last: Option<CudaSlice<f32>> = None;
6902 for &(a, b) in &bounds {
6903 let tc = b - a;
6904 let tok_d = e.stream().clone_htod(&tokens[a..b])?;
6905 let x =
6906 e.embed_gather_device_td(embd_gpu, &tok_d, tc, n_embd, embd_qtype, embd_row_bytes)?;
6907 let out = self.step35_verify_batch_layers(e, x, 0, n_layers, a, tc, cache)?;
6908 e.copy_into(&mut hiddens, a * n_embd, &out, tc * n_embd)?;
6909 if b == t_total {
6910 let mut h = e.uninit(n_embd)?;
6911 e.dtod_copy_view(&out.slice((tc - 1) * n_embd..tc * n_embd), &mut h)?;
6912 last = Some(h);
6913 }
6914 }
6915 let h_seed = last.expect("last chunk produced the seed row");
6916 let mut hn = e.uninit(n_embd)?;
6917 e.rms_norm_decode(
6918 &h_seed,
6919 self.output_norm.float_data(),
6920 &mut hn,
6921 n_embd,
6922 1,
6923 self.cfg.rms_eps,
6924 )?;
6925 let logits_d = e.matmul_decode_exact(&self.output, &hn, 1)?;
6926 let logits = e.dtoh(&logits_d)?;
6927 cache.pos = t_total;
6928 Ok(Some((logits, h_seed, hiddens)))
6929 }
6930
6931 #[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
6932 fn step35_verify_batch_layers(
6933 &self,
6934 e: &Engine,
6935 mut x: CudaSlice<f32>,
6936 lo: usize,
6937 hi: usize,
6938 pos0: usize,
6939 t: usize,
6940 cache: &mut Cache,
6941 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6942 let n_embd = self.cfg.n_embd as usize;
6943 if !self.uses_sliding_gated_moe_program() {
6944 return Err(
6945 "serving-class verify requires sliding-gated-MoE canonical operations".into(),
6946 );
6947 }
6948 // SERVING-CLASS VERIFY (MEMRA_SPEC_VERIFY_EAGER=1, step37 MTP bring-up): each verify
6949 // column rides decode_layers_eager — the EXACT t=1 program live serving runs (all TP2
6950 // doors) — row-outer, so row r's appends land before row r+1 attends: bit-equal to
6951 // plain greedy by construction. Only the unsplit full-range walk qualifies; PP splits
6952 // and the tap path keep the batch-layer class.
6953 static VE: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
6954 let eager_verify =
6955 *VE.get_or_init(spec_verify_eager_on) && lo == 0 && hi == self.layers.len();
6956 if eager_verify {
6957 // T-COLUMN LAYER-OUTER WALK (MEMRA_SPEC_VERIFY_TCOL=1): per layer, one t-grid
6958 // attn norm + ONE weight-amortized QKV(+gate) over all T columns, then each
6959 // column runs the UNMODIFIED t=1 attention program via the col-select door and
6960 // the ordinary residual/FFN body. Values per column are bit-equal to the
6961 // row-outer walk: rms over the materialized residual == the fused add+norm
6962 // (kernel_check identity), the tcol kernel's per-column FP order == the t=1
6963 // kernel, and every downstream op IS the t=1 program.
6964 static TCOL: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
6965 let tcol = *TCOL.get_or_init(spec_verify_tcol_on);
6966 // T > 32 (prefill-class): run the SAME walk in 32-row chunks — each chunk's
6967 // rows are the t=1 program bit-for-bit and the rope pass advances the cache,
6968 // so a chunked call is value-identical to the row-outer loop it replaces.
6969 static TROWS_PREFILL: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
6970 // MEMRA_STEP_GEMM_PRIME outranks the walk: with the grouped GEMM prime armed, the
6971 // t-row walk defers so the batch path (GEMM trunk + grouped MoE) takes the prompt —
6972 // flag precedence between two existing doors, not a new flag. Without this, both
6973 // doors ON meant the walk still won and the GEMM prime needed PRIME_TROWS=0 by hand.
6974 let trows_prefill =
6975 *TROWS_PREFILL.get_or_init(|| prime_trows_on() && !crate::step_gemm_prime_on());
6976 // MEMRA_PRIME_TROWS_T=<w>: chunk width, default 8 = the REAL cap of this walk.
6977 // The workspace slabs go to 32 rows, but `matvec_bf16_qkvg_tcol_into` refuses
6978 // t > 8 (compile-time-T twins exist for 2/4/8 only; the runtime-t kernel spills
6979 // its accumulators to local memory), so a wider chunk fails the request with
6980 // "matvec_bf16_qkvg_tcol geometry" — which is exactly how the first server-path
6981 // TROWS arm died. Measured at 193 tokens: w=8 2.459 s, w=4 2.574 s.
6982 static TROWS_W: std::sync::OnceLock<Result<usize, String>> = std::sync::OnceLock::new();
6983 let trows_w = match TROWS_W.get_or_init(|| {
6984 let value = std::env::var("MEMRA_PRIME_TROWS_T").ok();
6985 parse_prime_trows_width(value.as_deref())
6986 }) {
6987 Ok(width) => *width,
6988 Err(err) => return Err(err.clone().into()),
6989 };
6990 if tcol && trows_prefill && t > trows_w {
6991 // One-time engagement receipt: without it a prefill gate cannot tell a
6992 // chunked walk from the row-outer fallback it is supposed to replace
6993 // (the first PRIME_TROWS gate passed vacuously on exactly that).
6994 static SEEN: std::sync::atomic::AtomicBool =
6995 std::sync::atomic::AtomicBool::new(false);
6996 if !SEEN.swap(true, std::sync::atomic::Ordering::Relaxed) {
6997 eprintln!(
6998 "[prime-trows] ENGAGED t={t} width={trows_w} chunks={} layers={}..{}",
6999 t.div_ceil(trows_w),
7000 lo,
7001 hi
7002 );
7003 }
7004 let mut out = e.uninit(t * n_embd)?;
7005 let mut start = 0usize;
7006 while start < t {
7007 let mut end = (start + trows_w).min(t);
7008 if t - end == 1 {
7009 end -= 1;
7010 }
7011 let tc = end - start;
7012 let mut xc = e.uninit(tc * n_embd)?;
7013 e.dtod_copy_view(&x.slice(start * n_embd..end * n_embd), &mut xc)?;
7014 let oc =
7015 self.step35_verify_batch_layers(e, xc, lo, hi, pos0 + start, tc, cache)?;
7016 e.copy_into(&mut out, start * n_embd, &oc, tc * n_embd)?;
7017 start = end;
7018 }
7019 return Ok(out);
7020 }
7021 if tcol && (2..=32).contains(&t) {
7022 // MEMRA_TCOL_PROF=1: synchronized per-segment wall profile of the walk
7023 // (norm+QKV precompute / per-col attention / per-col residual+FFN). The
7024 // syncs serialize the stream, so the split is for TARGETING amortization
7025 // work only — never a perf claim.
7026 static PROF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
7027 let prof =
7028 *PROF.get_or_init(|| std::env::var("MEMRA_TCOL_PROF").as_deref() == Ok("1"));
7029 let mut prof_ms = [0f64; 3];
7030 let eps = self.cfg.rms_eps;
7031 let mut x_t = x;
7032 let mut h_t = e.uninit(t * n_embd)?;
7033 let mut h_row = e.uninit(n_embd)?; // real row: the non-dcw fallback reads it
7034 // Per-column pos buffers hoisted out of the layer loop (a per-col-per-layer
7035 // pageable htod was an in-stream engine turnaround x t x 45).
7036 let mut pos_rows = Vec::with_capacity(t);
7037 for r in 0..t {
7038 pos_rows.push(e.htod_i32(&[(pos0 + r) as i32])?);
7039 }
7040 let mut ok = true;
7041 // MEMRA_TCOL_OPROJ=1: defer each column's o_proj — the finish seam
7042 // stashes `gated` instead of joining per column; one b4_tcol per rank +
7043 // one slab join produce every column's `mixed` after the attention pass.
7044 // Bit-exact per column (t=1 b4 program per column; elementwise join).
7045 // MEMRA_TCOL_FFN=1: today this only IMPLIES the o_proj defer above. Its
7046 // named feature, the two-column device-routed FFN sweep, rode the
7047 // slot-major v2 TP banks and was REMOVED with the MEMRA_NVFP4_BANK_V2 door
7048 // (2026-08-29, research/step37-bankv2-removal-20260829): the v2 layout
7049 // changed generated text in serving. The flag itself stays because it is
7050 // family-armed in the step37 serving defaults and killing it here would
7051 // silently drop the o_proj defer from the qualified serving shape.
7052 static FFN2: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
7053 let ffn_batch = *FFN2.get_or_init(tcol_ffn_on);
7054 let oproj_batch = crate::tp::tcol_oproj_on() || ffn_batch;
7055 // MEMRA_SPEC_FA2=1 (T=2 only): eligible layers defer BOTH columns' fa —
7056 // the per-column pass norms/ropes/appends and stashes q+gate, then one
7057 // shared-KV fa_decode_dcw2 per rank + the o_proj join produce the
7058 // [2, o_out] mixed slab. The precheck runs before arming (stashing is
7059 // unrecoverable); ineligible/boundary layers run the ordinary program.
7060 let fa2 = crate::tp::spec_fa2_on() && t <= 32;
7061 let mut mixed_row = e.uninit(n_embd)?;
7062 let mut pos_staged = false;
7063 for il in lo..hi {
7064 let layer = &self.layers[il];
7065 // BEFORE this layer touches its planes: is the history it is about to
7066 // attend already poisoned? Global (non-ring) layers only, which are the
7067 // ones the level-2 bitmap implicates.
7068 if kv_plane_scan_on()
7069 && self.step35_geom(il).window.is_none()
7070 && let Some(distributed) = cache.tp_kv[il].as_ref()
7071 {
7072 scan_kv_plane(e, distributed, il, pos0)?;
7073 }
7074 let fa2_layer = fa2 && self.step35_fa_rows_precheck(cache, il, pos0, t)?;
7075 let mut seg = std::time::Instant::now();
7076 e.rms_norm(&x_t, layer.attn_norm.float_data(), &mut h_t, n_embd, t, eps)?;
7077 if !self.step35_verify_qkv_precompute(e, il, &h_t, t)? {
7078 ok = false;
7079 break;
7080 }
7081 // FULL t-row attention pass (rope/append + fa + combine + o_proj in
7082 // 3 launches/rank): same-session rows, slot = len-base+r, one len
7083 // advance by t. Host cache bookkeeping mirrors the per-column tail.
7084 if fa2_layer
7085 && let Some(mixed_t) =
7086 self.step35_verify_rope_fa_pass(e, il, cache, pos0, t, !pos_staged)?
7087 {
7088 pos_staged = true;
7089 {
7090 let tp_kv = cache.tp_kv[il]
7091 .as_mut()
7092 .expect("precheck verified the distributed cache");
7093 let transaction = tp_kv.begin_transaction()?;
7094 let crate::hybrid::Mixer::Full(fa) = &layer.mixer else {
7095 return Err("verify rope pass expects full attention".into());
7096 };
7097 let tp = fa
7098 .step_tp_qkv
7099 .as_ref()
7100 .ok_or("verify rope pass lost its TP state")?;
7101 let empty: [CudaSlice<f32>; 0] = [];
7102 tp.runtime.append_tp_kv_transaction_inner(
7103 tp_kv,
7104 transaction,
7105 &empty,
7106 &empty,
7107 t,
7108 true,
7109 )?;
7110 tp.runtime
7111 .commit_tp_kv_transaction_external(tp_kv, transaction, t)?;
7112 if let Some(local) = cache.kv[il].as_mut() {
7113 local.len = pos0 + t;
7114 if !crate::tp::len_mirror_lazy_on() {
7115 e.set_i32_one(&mut local.len_d, local.len as i32)?;
7116 }
7117 }
7118 }
7119 if prof {
7120 e.stream().synchronize()?;
7121 prof_ms[1] += seg.elapsed().as_secs_f64() * 1e3;
7122 seg = std::time::Instant::now();
7123 }
7124 let o_out = mixed_t.len() / t;
7125 let mut next = e.uninit(t * n_embd)?;
7126 {
7127 for r in 0..t {
7128 e.dtod_copy_view(
7129 &mixed_t.slice(r * o_out..(r + 1) * o_out),
7130 &mut mixed_row,
7131 )?;
7132 let mut x_row = e.uninit(n_embd)?;
7133 e.dtod_copy_view(
7134 &x_t.slice(r * n_embd..(r + 1) * n_embd),
7135 &mut x_row,
7136 )?;
7137 let (x1, ffn_out) = self.residual_norm_ffn(
7138 e, layer, &x_row, &mixed_row, n_embd, il, eps,
7139 )?;
7140 let mut x2 = e.uninit(n_embd)?;
7141 e.add(&x1, &ffn_out, &mut x2, n_embd)?;
7142 e.dtod_copy_into(&x2, &mut next, r * n_embd)?;
7143 }
7144 }
7145 if prof {
7146 e.stream().synchronize()?;
7147 prof_ms[2] += seg.elapsed().as_secs_f64() * 1e3;
7148 }
7149 x_t = next;
7150 if spec_nan_scan() {
7151 // The scan MUST sit on this arm too. It used to live only on
7152 // the non-fused tail, so a fused layer's poison was first
7153 // reported by the next non-fused layer.
7154 verify_arm_receipt(
7155 "fused",
7156 il,
7157 pos0,
7158 t,
7159 cache.tp_kv[il].as_ref().map(|d| d.staged_len()),
7160 );
7161 nan_scan_rows(
7162 e,
7163 &x_t,
7164 t,
7165 n_embd,
7166 &format!("tcol layer {il} pos0={pos0} arm=fused"),
7167 )?;
7168 }
7169 continue;
7170 }
7171 if prof {
7172 e.stream().synchronize()?;
7173 prof_ms[0] += seg.elapsed().as_secs_f64() * 1e3;
7174 seg = std::time::Instant::now();
7175 }
7176 let mut next = e.uninit(t * n_embd)?;
7177 // Columns whose o_proj was deferred (their FFN runs after the join).
7178 // A NON-deferred column's FFN must run INSIDE the column loop: the
7179 // oproj-tail handoff is a single cell that the same column's
7180 // residual_norm_ffn consumes before the next column's finish.
7181 let mut deferred: Vec<usize> = Vec::new();
7182 let mut fa2_deferred: Vec<usize> = Vec::new();
7183 let ffn_col = |r: usize,
7184 mixed: &CudaSlice<f32>,
7185 next: &mut CudaSlice<f32>|
7186 -> Result<(), Box<dyn std::error::Error>> {
7187 let mut x_row = e.uninit(n_embd)?;
7188 e.dtod_copy_view(&x_t.slice(r * n_embd..(r + 1) * n_embd), &mut x_row)?;
7189 let (x1, ffn_out) =
7190 self.residual_norm_ffn(e, layer, &x_row, mixed, n_embd, il, eps)?;
7191 if spec_nan_scan_level() >= 2 {
7192 nan_scan_rows(
7193 e,
7194 &ffn_out,
7195 1,
7196 n_embd,
7197 &format!("tcol layer {il} col {r} per-column FFN out"),
7198 )?;
7199 }
7200 let mut x2 = e.uninit(n_embd)?;
7201 e.add(&x1, &ffn_out, &mut x2, n_embd)?;
7202 e.dtod_copy_into(&x2, next, r * n_embd)?;
7203 Ok(())
7204 };
7205 #[allow(clippy::needless_range_loop)]
7206 // allow: the explicit index loop keeps the offset arithmetic visible and aligned with the device-side indexing
7207 for r in 0..t {
7208 e.dtod_copy_view(&h_t.slice(r * n_embd..(r + 1) * n_embd), &mut h_row)?;
7209 let row_pos = &pos_rows[r];
7210 crate::tp::set_verify_tcol(Some(r));
7211 if fa2_layer {
7212 crate::tp::set_spec_fa2_defer(Some(r));
7213 } else if oproj_batch {
7214 crate::tp::set_tcol_oproj_defer(Some(r));
7215 }
7216 let mixed = match &layer.mixer {
7217 crate::hybrid::Mixer::Full(fa) => {
7218 self.full_attn_decode(e, fa, &h_row, row_pos, pos0 + r, cache, il)
7219 }
7220 _ => Err("step35 verify expects full attention".into()),
7221 };
7222 crate::tp::set_verify_tcol(None);
7223 crate::tp::set_spec_fa2_defer(None);
7224 crate::tp::set_tcol_oproj_defer(None);
7225 let mixed = mixed?;
7226 if fa2_layer && crate::tp::take_spec_fa2_stashed() {
7227 fa2_deferred.push(r);
7228 } else if oproj_batch && crate::tp::take_tcol_oproj_stashed() {
7229 deferred.push(r);
7230 } else {
7231 if spec_nan_scan_level() >= 2 {
7232 let cols = mixed.len();
7233 nan_scan_rows(
7234 e,
7235 &mixed,
7236 1,
7237 cols,
7238 &format!("tcol layer {il} col {r} per-column ATTN out"),
7239 )?;
7240 }
7241 ffn_col(r, &mixed, &mut next)?;
7242 }
7243 }
7244 if !fa2_deferred.is_empty() && fa2_deferred.len() != t {
7245 // The precheck guarantees both columns stash or neither; a strict
7246 // subset means a column's output was never produced anywhere.
7247 return Err("spec fa2 stash engaged for a subset of columns".into());
7248 }
7249 if prof {
7250 e.stream().synchronize()?;
7251 prof_ms[1] += seg.elapsed().as_secs_f64() * 1e3;
7252 seg = std::time::Instant::now();
7253 }
7254 if !fa2_deferred.is_empty() {
7255 deferred = fa2_deferred;
7256 }
7257 if !deferred.is_empty() {
7258 let mixed_t = if fa2_layer {
7259 self.step35_verify_fa_rows_join(e, il, cache, pos0, t)?
7260 } else {
7261 self.step35_verify_oproj_tcol(e, il, t)?
7262 };
7263 let o_out = mixed_t.len() / t;
7264 if spec_nan_scan_level() >= 2 {
7265 nan_scan_rows(
7266 e,
7267 &mixed_t,
7268 t,
7269 o_out,
7270 &format!("tcol layer {il} JOINED attn over deferred cols"),
7271 )?;
7272 }
7273 // Batched t=2 residual+MoE: one t-grid add_rms_norm (per-row
7274 // program == t=1; bit-identical to the oproj-tail join per the
7275 // M2 verbatim-program contract) feeding the two-column routed
7276 // sweep. Ineligible layers (dense FFN, non-nvfp4) fall through
7277 // to the per-column body.
7278 {
7279 for &r in &deferred {
7280 e.dtod_copy_view(
7281 &mixed_t.slice(r * o_out..(r + 1) * o_out),
7282 &mut mixed_row,
7283 )?;
7284 ffn_col(r, &mixed_row, &mut next)?;
7285 }
7286 }
7287 }
7288 if prof {
7289 e.stream().synchronize()?;
7290 prof_ms[2] += seg.elapsed().as_secs_f64() * 1e3;
7291 }
7292 x_t = next;
7293 if spec_nan_scan() {
7294 verify_arm_receipt(
7295 if fa2_layer { "join" } else { "percol" },
7296 il,
7297 pos0,
7298 t,
7299 cache.tp_kv[il].as_ref().map(|d| d.staged_len()),
7300 );
7301 nan_scan_rows(
7302 e,
7303 &x_t,
7304 t,
7305 n_embd,
7306 &format!(
7307 "tcol layer {il} pos0={pos0} arm={}",
7308 if fa2_layer { "join" } else { "percol" }
7309 ),
7310 )?;
7311 }
7312 }
7313 if prof {
7314 eprintln!(
7315 "[tcol-prof] t={t} norm+qkv={:.3}ms attn={:.3}ms ffn={:.3}ms",
7316 prof_ms[0], prof_ms[1], prof_ms[2]
7317 );
7318 }
7319 if ok {
7320 return Ok(x_t);
7321 }
7322 // fall through to the row-outer walk on ineligible layers
7323 x = x_t;
7324 }
7325 let mut next = e.uninit(t * n_embd)?;
7326 let scan = spec_nan_scan();
7327 for r in 0..t {
7328 let mut row = e.uninit(n_embd)?;
7329 e.dtod_copy_view(&x.slice(r * n_embd..(r + 1) * n_embd), &mut row)?;
7330 let row_pos = e.htod_i32(&[(pos0 + r) as i32])?;
7331 let out = if scan {
7332 // Diagnostic arm: the same range walked one layer at a time so the first
7333 // poisoned layer names itself. `decode_layers_eager(lo, hi)` is range-scoped
7334 // and executes its trailing residual add, so a per-layer chain is the same
7335 // program with the cross-layer add+norm fusion unrolled.
7336 nan_scan_rows(
7337 e,
7338 &row,
7339 1,
7340 n_embd,
7341 &format!("embed row r={r} pos={}", pos0 + r),
7342 )?;
7343 let mut acc = row;
7344 for il in lo..hi {
7345 acc = self.decode_layers_eager(
7346 e,
7347 acc,
7348 il,
7349 il + 1,
7350 &row_pos,
7351 pos0 + r,
7352 cache,
7353 )?;
7354 nan_scan_rows(
7355 e,
7356 &acc,
7357 1,
7358 n_embd,
7359 &format!("row-outer layer {il} r={r} pos={}", pos0 + r),
7360 )?;
7361 }
7362 acc
7363 } else {
7364 self.decode_layers_eager(e, row, lo, hi, &row_pos, pos0 + r, cache)?
7365 };
7366 e.dtod_copy_into(&out, &mut next, r * n_embd)?;
7367 }
7368 // dflash taps are NOT produced on this arm (they need per-layer hiddens the
7369 // row-outer walk does not materialize); the door is a step37 MTP bring-up
7370 // surface where taps are unused.
7371 return Ok(next);
7372 }
7373 let mut ph_last = std::time::Instant::now();
7374 for il in lo..hi {
7375 let mut next = e.uninit(t * n_embd)?;
7376 for r in 0..t {
7377 let mut row = e.uninit(n_embd)?;
7378 e.dtod_copy_view(&x.slice(r * n_embd..(r + 1) * n_embd), &mut row)?;
7379 // The caller owns this verify's position. During controller overlap, cache.pos
7380 // still describes generation N while this stage-0 walk belongs to N+1.
7381 let row_pos = e.htod_i32(&[(pos0 + r) as i32])?;
7382 let mut one = [&mut *cache];
7383 let out = self.step35_decode_batch_layers(
7384 e,
7385 row,
7386 &mut one,
7387 &[(pos0 + r) as i32],
7388 &row_pos,
7389 il,
7390 il + 1,
7391 &mut ph_last,
7392 )?;
7393 e.dtod_copy_into(&out, &mut next, r * n_embd)?;
7394 }
7395 self.dflash_tap(e, cache, il, &next, t)?;
7396 x = next;
7397 if spec_nan_scan() {
7398 nan_scan_rows(e, &x, t, n_embd, &format!("batch-layer {il} pos0={pos0}"))?;
7399 }
7400 }
7401 Ok(x)
7402 }
7403
7404 /// DSpark drafter verify (lane/dspark-q38-recover): one t-row forward through the
7405 /// SERVING-CLASS verify funnel (`decode_step_t_core_stream` — the same numeric class
7406 /// MTP verify rides, GDN state advanced in place), returning per-row argmax tokens.
7407 /// Advances `cache.pos += t`; the caller owns snapshot/rollback (block acceptance is
7408 /// prefix-keep, not all-or-nothing).
7409 pub(crate) fn dspark_verify_t_am(
7410 &self,
7411 e: &Engine,
7412 tokens: &[u32],
7413 pos0: usize,
7414 cache: &mut Cache,
7415 ) -> Result<Vec<u32>, Box<dyn std::error::Error>> {
7416 let (logits, _hn) = self.decode_step_t_core_stream(
7417 e, tokens, pos0, cache, None, None, None, None, None, None,
7418 )?;
7419 let t = tokens.len();
7420 let v = self.output.out_features();
7421 let mut am_d = e.stream().alloc_zeros::<u32>(t)?;
7422 for r in 0..t {
7423 e.argmax_token_device_col(&logits, r, v, &mut am_d, r)?;
7424 }
7425 e.dtoh_u32(&am_d)
7426 }
7427
7428 /// DSpark verify returning the RAW verify logits [t, n_vocab] (device-resident) instead
7429 /// of per-row argmaxes — the sampled-admission arm's input (rejection-sampling accept
7430 /// gathers filtered p from these columns; lane/dspark-sampled-admission-20260820). Same
7431 /// forward as `dspark_verify_t_am`; the greedy arm keeps its argmax wrapper untouched.
7432 pub(crate) fn dspark_verify_t_logits(
7433 &self,
7434 e: &Engine,
7435 tokens: &[u32],
7436 pos0: usize,
7437 cache: &mut Cache,
7438 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7439 let (logits, _hn) = self.decode_step_t_core_stream(
7440 e, tokens, pos0, cache, None, None, None, None, None, None,
7441 )?;
7442 Ok(logits)
7443 }
7444
7445 /// DSpark verify with the MTP column-stash armed: identical forward to
7446 /// `dspark_verify_t_am`, but fills a `VerifyCkpt` so a partial accept can restore
7447 /// column state directly (`dspark_commit_prefix`) instead of snapshot-replay.
7448 /// The ckpt type is opaque outside spec.rs (newtype) — dflash.rs threads it through.
7449 pub(crate) fn dspark_verify_t_am_ckpt(
7450 &self,
7451 e: &Engine,
7452 tokens: &[u32],
7453 pos0: usize,
7454 cache: &mut Cache,
7455 ) -> Result<(Vec<u32>, DsparkVerifyCkpt), Box<dyn std::error::Error>> {
7456 let mut ck = VerifyCkpt::new(self.layers.len());
7457 let (logits, _hn) = self.decode_step_t_core_stream(
7458 e,
7459 tokens,
7460 pos0,
7461 cache,
7462 None,
7463 Some(&mut ck),
7464 None,
7465 None,
7466 None,
7467 None,
7468 )?;
7469 let t = tokens.len();
7470 let v = self.output.out_features();
7471 let mut am_d = e.stream().alloc_zeros::<u32>(t)?;
7472 for r in 0..t {
7473 e.argmax_token_device_col(&logits, r, v, &mut am_d, r)?;
7474 }
7475 Ok((e.dtoh_u32(&am_d)?, DsparkVerifyCkpt(ck)))
7476 }
7477
7478 /// Engine-bundle slice 2: `dspark_verify_t_am_ckpt` with DEVICE tokens and NO readback.
7479 /// The verify tokens are the round's `chain_d` (cand layout: [anchor, drafts...]); the
7480 /// embed gathers its first `t` entries on-device (`embed_gather_u32_t` — bit-identical
7481 /// rows to the host arm), so the host never blocks on the draft chain before dispatching
7482 /// verify. Returns the device per-row argmax buffer; the caller merges its readback with
7483 /// the chain's into ONE sync. Forward, ckpt fill and argmax walk are `_ckpt` verbatim.
7484 #[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
7485 pub(crate) fn dspark_verify_t_am_ckpt_dev(
7486 &self,
7487 e: &Engine,
7488 vtok: &CudaSlice<u32>,
7489 t: usize,
7490 pos0: usize,
7491 cache: &mut Cache,
7492 embd_dev: (&CudaSlice<u8>, i32, usize),
7493 graphs: Option<&mut DsparkVerifyGraphs>,
7494 ) -> Result<(CudaSlice<u32>, DsparkVerifyCkpt), Box<dyn std::error::Error>> {
7495 debug_assert!(
7496 vtok.len() >= t,
7497 "verify window exceeds the device token buffer"
7498 );
7499 // The slab flag is a per-round statement: clear it here so a verify that never
7500 // reaches the graphs door (rowwise env, a non-tparallel arm) cannot leave a
7501 // stale `true` steering the commit at slabs the round never wrote.
7502 let mut graphs = graphs;
7503 if let Some(g) = graphs.as_deref_mut() {
7504 g.round_slab = false;
7505 }
7506 let mut ck = VerifyCkpt::new(self.layers.len());
7507 // Dummy host tokens size the funnel; the embed reads `vtok` (the round-stream
7508 // arm's established pattern — spec.rs stream-mode verify does the same).
7509 let dummy = vec![0u32; t];
7510 let (logits, _hn) = self.decode_step_t_core_stream(
7511 e,
7512 &dummy,
7513 pos0,
7514 cache,
7515 Some(embd_dev),
7516 Some(&mut ck),
7517 None,
7518 None,
7519 Some(vtok),
7520 graphs,
7521 )?;
7522 let v = self.output.out_features();
7523 let mut am_d = e.stream().alloc_zeros::<u32>(t)?;
7524 for r in 0..t {
7525 e.argmax_token_device_col(&logits, r, v, &mut am_d, r)?;
7526 }
7527 Ok((am_d, DsparkVerifyCkpt(ck)))
7528 }
7529
7530 /// Ckpt-armed twin of [`Self::dspark_verify_t_logits`] (sampled-admission arm).
7531 pub(crate) fn dspark_verify_t_logits_ckpt(
7532 &self,
7533 e: &Engine,
7534 tokens: &[u32],
7535 pos0: usize,
7536 cache: &mut Cache,
7537 ) -> Result<(CudaSlice<f32>, DsparkVerifyCkpt), Box<dyn std::error::Error>> {
7538 let mut ck = VerifyCkpt::new(self.layers.len());
7539 let (logits, _hn) = self.decode_step_t_core_stream(
7540 e,
7541 tokens,
7542 pos0,
7543 cache,
7544 None,
7545 Some(&mut ck),
7546 None,
7547 None,
7548 None,
7549 None,
7550 )?;
7551 Ok((logits, DsparkVerifyCkpt(ck)))
7552 }
7553
7554 /// Restore the round to `keep` accepted columns from the verify stash: KV lens and
7555 /// pos from the pre-verify snapshot + keep, GDN conv/ssm from the stashed column
7556 /// state — no replay forward. The exact `commit_verified_prefix` the MTP path ships.
7557 pub(crate) fn dspark_commit_prefix(
7558 &self,
7559 e: &Engine,
7560 cache: &mut Cache,
7561 snap: &crate::cache::CacheSnapshot,
7562 ckpt: &DsparkVerifyCkpt,
7563 keep: usize,
7564 ) -> Result<(), Box<dyn std::error::Error>> {
7565 self.commit_verified_prefix(e, cache, snap, &ckpt.0, keep, false, None)
7566 }
7567
7568 /// Slice-3 commit twin: restore to `keep` accepted columns when the round's linear
7569 /// column stash lives in the graphs ctx's persistent slabs (`DsparkVerifyGraphs`) —
7570 /// the cols arm's exact semantics (KV lens + pos from the snapshot, GDN conv/ssm
7571 /// from the stash of column keep-1), slab-addressed and batched into two copy
7572 /// launches. `MEMRA_STATE_COPY_BATCH=0` falls back to per-layer view copies.
7573 pub(crate) fn dspark_commit_prefix_slab(
7574 &self,
7575 e: &Engine,
7576 cache: &mut Cache,
7577 snap: &crate::cache::CacheSnapshot,
7578 ctx: &DsparkVerifyGraphs,
7579 keep: usize,
7580 ) -> Result<(), Box<dyn std::error::Error>> {
7581 use cudarc::driver::DevicePtr;
7582 debug_assert!(keep >= 1, "keep==0 rounds take the legacy rollback");
7583 let mut conv_src: Vec<u64> = Vec::new();
7584 let mut ssm_src: Vec<u64> = Vec::new();
7585 let mut conv_dst: Vec<u64> = Vec::new();
7586 let mut ssm_dst: Vec<u64> = Vec::new();
7587 for il in 0..self.layers.len() {
7588 if let (Some(kvl), Some(saved)) = (cache.kv[il].as_mut(), snap.kv_len[il]) {
7589 kvl.len = saved + keep;
7590 e.set_i32_one(&mut kvl.len_d, kvl.len as i32)?;
7591 }
7592 if let Some(rl) = cache.recur[il].as_ref() {
7593 let (pc, ps, _cw, _sw) = ctx
7594 .slab_row(e, il, keep - 1)
7595 .ok_or("slab commit: linear layer missing from the graphs ctx")?;
7596 conv_src.push(pc);
7597 ssm_src.push(ps);
7598 let st = &e.gpu.stream();
7599 let (dc, _g0) = rl.conv_state.device_ptr(st);
7600 let (ds, _g1) = rl.ssm_state.device_ptr(st);
7601 conv_dst.push(dc);
7602 ssm_dst.push(ds);
7603 }
7604 }
7605 let n = conv_src.len();
7606 if n > 0 {
7607 if state_copy_batch_on() {
7608 let mut tt = vec![0u64; 2 * n];
7609 tt[..n].copy_from_slice(&conv_src);
7610 tt[n..].copy_from_slice(&conv_dst);
7611 let ct = e.htod_u64(&tt)?;
7612 tt[..n].copy_from_slice(&ssm_src);
7613 tt[n..].copy_from_slice(&ssm_dst);
7614 let st = e.htod_u64(&tt)?;
7615 e.copy_batch_uniform_f32(&ct, n, ctx.conv_words)?;
7616 e.copy_batch_uniform_f32(&st, n, ctx.ssm_words)?;
7617 } else {
7618 let (cw, sw) = (ctx.conv_words, ctx.ssm_words);
7619 let row = keep - 1;
7620 for il in 0..self.layers.len() {
7621 let Some(rl) = cache.recur[il].as_mut() else {
7622 continue;
7623 };
7624 let k = ctx.lin_pos[&il];
7625 {
7626 let sv = e.view(&ctx.stash_conv[k], (row + 1) * cw);
7627 let win = sv.slice(row * cw..(row + 1) * cw);
7628 e.copy_view_into(&mut rl.conv_state, 0, &win, cw)?;
7629 }
7630 {
7631 let sv = e.view(&ctx.stash_ssm[k], (row + 1) * sw);
7632 let win = sv.slice(row * sw..(row + 1) * sw);
7633 e.copy_view_into(&mut rl.ssm_state, 0, &win, sw)?;
7634 }
7635 }
7636 }
7637 }
7638 cache.pos = snap.pos + keep;
7639 Ok(())
7640 }
7641
7642 /// Qwen35-family verify trunk in the live serving numeric class.
7643 ///
7644 /// Serving intentionally keeps this architecture in the generic batched program even at
7645 /// B=1. The older verify walk used its own mirrored dispatch and can flip near-tie argmaxes.
7646 ///
7647 /// Two arms, one numeric class:
7648 /// - DENSE GDN (`DenseMlp`, t<=16): `qwen35_verify_tparallel` — the weight ops (norms,
7649 /// projections, FFN) hoist to m=T through the exact-tier batched kernels whose per-row
7650 /// program IS the m=1 program (`matmul_pre == fused2 per (tensor,row); _bN mmvq per-row
7651 /// == m=1` — decode_batch.rs v2 note), while the state ops (conv ring, gdn scan, KV
7652 /// append, fa decode) stay a per-row loop running the b_n=1 serving kernels with each
7653 /// row's own t_kv-driven arm pick (the straddle law: every row executes the exact
7654 /// program its isolated serving step would). One weight read per layer per round
7655 /// instead of T — this is what makes MTP profitable in the exact class (the per-row
7656 /// walk measured verify(K+1) ~= (K+1) plain steps: 69 -> 44 tok/s served, 2026-08-15).
7657 /// - MoE / t>16 / `MEMRA_SPEC_VERIFY_ROWWISE=1`: the per-row replay of the authoritative
7658 /// serving layer body, preserving single-session autoregressive cache order (the
7659 /// correctness reference; also the rollback seam for the t-parallel arm).
7660 ///
7661 /// Bit-identity of the t-parallel arm vs the rowwise arm is gated by spec-serve-gate
7662 /// (zero differing logits at T=1..4, K arms) + the 8-prompt ON/OFF canary before ship.
7663 #[allow(clippy::too_many_arguments)]
7664 fn qwen35_verify_batch_layers(
7665 &self,
7666 e: &Engine,
7667 x: CudaSlice<f32>,
7668 lo: usize,
7669 hi: usize,
7670 pos0: usize,
7671 t: usize,
7672 cache: &mut Cache,
7673 ckpt: Option<&mut VerifyCkpt>,
7674 stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
7675 graphs: Option<&mut DsparkVerifyGraphs>,
7676 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7677 // Qwen35Moe admitted 2026-08-20 (lane/draftcost-moe): the t-parallel arm already
7678 // carries the MoE FFN (`moe_ffn_il_zq8` at m=T) and the GDN per-row state loop; the
7679 // arch fence was a qualification gate, not a mechanism gap. Measured disease on the
7680 // 35B-A3B class: rowwise verify ~= 5.6 ms per drafted token (one full trunk step
7681 // each) — the same (K+1)-plain-steps wall the dense admission fixed on 2026-08-15.
7682 // Rollback seam unchanged: MEMRA_SPEC_VERIFY_ROWWISE=1.
7683 let rowwise = std::env::var("MEMRA_SPEC_VERIFY_ROWWISE").as_deref() == Ok("1")
7684 || !self.batched_serving_numeric_class()
7685 || t > 16;
7686 if rowwise {
7687 if stream.is_some() {
7688 // rowwise replays per row with host cache.pos — irreconcilable with a
7689 // device position counter. Burst callers must keep t <= 16 and the
7690 // ROWWISE env unset; refusing beats silently mispositioned rows.
7691 return Err("qwen35 rowwise verify has no ROUND-STREAM arm \
7692 (t > 16 or MEMRA_SPEC_VERIFY_ROWWISE=1)"
7693 .into());
7694 }
7695 self.qwen35_verify_rowwise(e, x, lo, hi, pos0, t, cache, ckpt)
7696 } else {
7697 self.qwen35_verify_tparallel(e, x, lo, hi, pos0, t, cache, ckpt, stream, graphs)
7698 }
7699 }
7700
7701 /// The per-row correctness reference: replay each verify row through the authoritative
7702 /// serving layer body (`decode_batch_layers` at b_n=1). T full weight reads per layer.
7703 #[allow(clippy::too_many_arguments)]
7704 fn qwen35_verify_rowwise(
7705 &self,
7706 e: &Engine,
7707 mut x: CudaSlice<f32>,
7708 lo: usize,
7709 hi: usize,
7710 pos0: usize,
7711 t: usize,
7712 cache: &mut Cache,
7713 mut ckpt: Option<&mut VerifyCkpt>,
7714 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7715 let n_embd = self.cfg.n_embd as usize;
7716 let saved_pos = cache.pos;
7717 let mut ph_last = std::time::Instant::now();
7718 for il in lo..hi {
7719 let mut next = e.uninit(t * n_embd)?;
7720 let mut col_states: Option<Vec<(CudaSlice<f32>, CudaSlice<f32>)>> =
7721 if ckpt.is_some() && t >= 2 && matches!(self.layers[il].mixer, Mixer::Linear(_)) {
7722 Some(Vec::with_capacity(t - 1))
7723 } else {
7724 None
7725 };
7726 for r in 0..t {
7727 cache.pos = pos0 + r;
7728 let mut row = e.uninit(n_embd)?;
7729 e.dtod_copy_view(&x.slice(r * n_embd..(r + 1) * n_embd), &mut row)?;
7730 let row_pos = e.htod_i32(&[(pos0 + r) as i32])?;
7731 let mut one = [&mut *cache];
7732 let ctx = self.batch_layer_ctx(e, &one, il, il + 1)?;
7733 let out = match self.decode_batch_layers(
7734 e,
7735 row,
7736 &mut one,
7737 &ctx,
7738 &row_pos,
7739 &mut ph_last,
7740 ) {
7741 Ok(out) => out,
7742 Err(error) => {
7743 cache.pos = saved_pos;
7744 return Err(error);
7745 }
7746 };
7747 e.dtod_copy_into(&out, &mut next, r * n_embd)?;
7748 if r + 1 < t
7749 && let Some(states) = col_states.as_mut()
7750 {
7751 let recur = cache.recur[il]
7752 .as_ref()
7753 .ok_or("Qwen35-MoE linear verify layer has no recurrent state")?;
7754 states.push((
7755 e.clone_dtod(&recur.conv_state)?,
7756 e.clone_dtod(&recur.ssm_state)?,
7757 ));
7758 }
7759 }
7760 if let (Some(checkpoint), Some(states)) = (ckpt.as_deref_mut(), col_states) {
7761 checkpoint.cols[il] = Some(states);
7762 }
7763 x = next;
7764 }
7765 cache.pos = saved_pos;
7766 Ok(x)
7767 }
7768
7769 /// T-PARALLEL VERIFY IN THE SERVING NUMERIC CLASS (lane/tparallel-verify, 2026-08-15).
7770 ///
7771 /// The weight ops run ONCE per layer at m=T; the state ops run per row through the same
7772 /// b_n=1 serving kernels the rowwise replay uses. Per-row bit-identity rests on the two
7773 /// pins the serving batch tier already carries:
7774 /// * `matmul_pre` / `_bN` mmvq: per-row program == m=1 program (decode_batch.rs v2 note,
7775 /// kernel-check pinned) — so a [T, n_embd] projection row equals the row projected
7776 /// alone;
7777 /// * row-indexed norms/elementwise (`rms_norm`, `quantize_q8_1`, `add_rms_norm`,
7778 /// `gated_rmsnorm[_q8_1]`, `silu_mul`, `rope_neox` with per-row positions): the T-row
7779 /// launch is the per-row program (same pin the generic verify's fused norms rely on).
7780 /// The sequential dependencies keep their exact serving order: the conv ring / gdn scan
7781 /// chain state row -> row through the `_b` kernels at b_n=1 (ping-pong via a 6-entry
7782 /// alternating pointer table, host handles swapped per row so VerifyCkpt clones the
7783 /// canonical state exactly as the rowwise arm does), and each row's KV append + fa decode
7784 /// picks its arm from ITS OWN t_kv (append: format-only; fa: `fa_seqs_eligible` + its own
7785 /// `fa_split_keys` rung at b_n=1) — the straddle law per row, so every row executes the
7786 /// program its isolated B=1 serving step would.
7787 ///
7788 /// Cost: 1 weight read per layer per round + T state micro-launches, vs the rowwise arm's
7789 /// T weight reads. Gated bit-identical vs the rowwise arm by spec-serve-gate + canary.
7790 #[allow(clippy::too_many_arguments)]
7791 fn qwen35_verify_tparallel(
7792 &self,
7793 e: &Engine,
7794 mut x: CudaSlice<f32>,
7795 lo: usize,
7796 hi: usize,
7797 pos0: usize,
7798 t: usize,
7799 cache: &mut Cache,
7800 mut ckpt: Option<&mut VerifyCkpt>,
7801 stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
7802 mut graphs: Option<&mut DsparkVerifyGraphs>,
7803 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7804 let seqs_append =
7805 std::env::var("MEMRA_BATCH_APPEND").as_deref() != Ok("0") && !Engine::kv_fp8_on();
7806 let batch_fa_on = std::env::var("MEMRA_BATCH_FA").as_deref() != Ok("0");
7807
7808 // Merge guard (v0.98 train, re-affirmed on the v0.100 train over slice 4c): the
7809 // ROUND-STREAM arm (lane/draftcost-moe, device position counter) and the dspark
7810 // verify graphs (engine-bundle slice 3 / trunk slice 4c) have no common caller —
7811 // stream rides the qwen35moe burst, graphs ride the dspark route. If a future
7812 // caller arms both, refuse loudly instead of silently dropping the graphs ctx
7813 // (the stream linear arm takes linear_attn_verify_t, not the graphed segment or
7814 // full-verify bodies).
7815 if stream.is_some() && graphs.is_some() {
7816 return Err(
7817 "qwen35 tparallel verify: ROUND-STREAM and dspark verify graphs \
7818 cannot arm together"
7819 .into(),
7820 );
7821 }
7822 // Engine-bundle slice 3 + slice 4c: with a graphs ctx armed, pointer tables are
7823 // refreshed once per verify (the gdn ping-pong moves handles; a fresh generation
7824 // moves the kv caches). Then:
7825 // - slice 4c: when the WHOLE round rides one seqs rung (every row batchable, one
7826 // split-ladder step, rung covers the round), the ENTIRE walk replays as ONE
7827 // full-verify graph per (vt, rung) — linear layers through the shared
7828 // `qwen35_tparallel_linear_layer` body, full-attention layers through the
7829 // shared `qwen35_tparallel_fa_layer` body in graph mode.
7830 // - fallback (straddle rounds, below the vec floor, partial walks): runs of
7831 // consecutive LINEAR layers replay the slice-3 per-(segment, vt) graphs and
7832 // the full-attention layers run eager (batched rows when eligible).
7833 //
7834 // GRAPH-LAUNCH HEADROOM GUARD (see GRAPH_LAUNCH_MIN_FREE): the dspark verify
7835 // graphs replay through this walk from THREE callers — the MTP spec round's vg
7836 // door (already dropped per round by `graph_round_ok` before it gets here), the
7837 // dspark one-shot, and the dspark SERVE round (default ON since v0.108). Below
7838 // the driver-free floor the WHOLE round takes the byte-identical eager
7839 // cols-ckpt walk — the same drop-the-ctx fallback the pool ceiling already
7840 // takes — instead of feeding cuGraphLaunch a card it segfaults on.
7841 if let Some(g) = graphs.as_deref_mut()
7842 && !graph_launch_headroom_ok(e)
7843 {
7844 g.round_slab = false;
7845 graphs = None;
7846 static NOTED: std::sync::Once = std::sync::Once::new();
7847 NOTED.call_once(|| graph_replay_suspended_note("dspark-vg"));
7848 }
7849 if let Some(g) = graphs.as_deref_mut() {
7850 g.refresh_tables(e, cache)?;
7851 g.round_slab = false;
7852 if let Some(rung) = g.full_rung(self, cache, lo, hi, t, seqs_append && batch_fa_on) {
7853 // Pool ceiling (dspark_vg_cap): an existing key always replays; a NEW
7854 // full capture past the ceiling falls through to the segment/eager arms.
7855 if g.full.contains_key(&(t, rung, hi)) || g.can_capture() {
7856 let out = g.run_full(self, e, lo, hi, &x, t, pos0, rung, cache)?;
7857 g.round_slab = true;
7858 return Ok(out);
7859 }
7860 }
7861 // Round-atomic ceiling check for the segment door: if any linear run in this
7862 // walk would need a NEW capture past the ceiling, the whole round runs the
7863 // eager cols-ckpt walk (mixing slab- and cols-stashed layers in one round
7864 // would corrupt the commit).
7865 if !g.segments_ready(self, lo, hi, t) {
7866 graphs = None;
7867 }
7868 }
7869 // STREAM (2b, lane/draftcost-moe): positions come from the device round counter
7870 // (pos_iota / i32_copy_add) so a burst round needs no host position knowledge.
7871 let pos_d = match stream {
7872 Some((_, ctr)) => {
7873 let mut p = e.alloc_uninit::<i32>(t)?;
7874 e.pos_iota(ctr, &mut p, t)?;
7875 p
7876 }
7877 None => {
7878 let pos_host: Vec<i32> = (0..t).map(|r| (pos0 + r) as i32).collect();
7879 e.htod_i32(&pos_host)?
7880 }
7881 };
7882 // Per-row 1-element position buffers, built ONCE per verify (the append/fa wrappers
7883 // take owned pos slices; building these inside the layer x row loops cost 16xT H2Ds).
7884 // LAZY since slice 4: the batched fa/append arm never touches them — they are built
7885 // on the first per-row fallback layer only (stream-aware there; the stream FA arm
7886 // rides the dc rows kernels and never reaches the fallback).
7887 let mut pos_rows: Option<Vec<CudaSlice<i32>>> = None;
7888 let mut il = lo;
7889 while il < hi {
7890 if graphs.is_some() && matches!(self.layers[il].mixer, Mixer::Linear(_)) {
7891 let mut end = il;
7892 while end < hi && matches!(self.layers[end].mixer, Mixer::Linear(_)) {
7893 end += 1;
7894 }
7895 let g = graphs.as_deref_mut().expect("checked above");
7896 x = g.run_segment(self, e, il, end, &x, t, cache)?;
7897 g.round_slab = true;
7898 il = end;
7899 continue;
7900 }
7901 let layer = &self.layers[il];
7902 if stream.is_none() && matches!(layer.mixer, Mixer::Linear(_)) {
7903 // Eager linear layer (no graphs ctx): the shared body, legacy cols-ckpt arm.
7904 // Under ROUND-STREAM the linear layers ride the fa-body match's stream arm
7905 // below (linear_attn_verify_t — the stream COMMIT needs its GdnStash).
7906 x = self.qwen35_tparallel_linear_layer(
7907 e,
7908 il,
7909 &x,
7910 t,
7911 cache,
7912 ckpt.as_deref_mut(),
7913 None,
7914 None,
7915 )?;
7916 il += 1;
7917 continue;
7918 }
7919 // Full-attention (or stream-Linear, or MLA-refusing) layer: the extracted
7920 // shared body — eager arm (fresh per-verify pos/table, exact t_kv sizing,
7921 // in-body len bump). The slice-4c captured full-verify graphs run the SAME
7922 // body in graph mode; under ROUND-STREAM the body's dc-rows / GDN stream arms
7923 // run (lane/draftcost-moe).
7924 x = self.qwen35_tparallel_fa_layer(
7925 e,
7926 il,
7927 &x,
7928 t,
7929 cache,
7930 FaLayerArgs {
7931 pos_d: &pos_d,
7932 pos_rows: &mut pos_rows,
7933 pos0,
7934 seqs_append,
7935 batch_fa_on,
7936 graph_cap: None,
7937 stream,
7938 ckpt: ckpt.as_deref_mut(),
7939 },
7940 )?;
7941 il += 1;
7942 }
7943 Ok(x)
7944 }
7945
7946 /// SHARED dense-FFN body for the qwen35 t-parallel layers (trunk-kernels slice B) —
7947 /// ONE copy for the fa and linear layer bodies (the verify_layers extraction lesson).
7948 /// Dual arm (MEMRA_TK_FFN_DUAL, default on): gate+up in ONE dual launch from the
7949 /// pre-quantized activation with macro-scales DEFERRED into the fused SwiGLU+q8_1
7950 /// epilogue, then ffn_down from the fused (aq, ad) — the q27 verify chain verbatim.
7951 /// Every door is the bit-identical proven one: `matmul_decode_exact_dual_pre` (per
7952 /// (tensor,token,row) == the two singles), `silu_mul_scaled_q8_1` (y*s inline == the
7953 /// scale_inplace store, value-exact; fused quantize == quantize_q8_1 bytes),
7954 /// `matmul_decode_exact_pre` (dispatch mirror of the singles' q8_1-fast tail).
7955 /// Dual-refused (t outside 2..=7, non-NVFP4, layout mismatch) or seam off -> the
7956 /// original singles chain, byte-for-byte.
7957 #[allow(clippy::too_many_arguments)]
7958 fn qwen35_tparallel_dense_ffn(
7959 &self,
7960 e: &Engine,
7961 ffn_gate: &crate::model::GpuTensor,
7962 ffn_up: &crate::model::GpuTensor,
7963 ffn_down: &crate::model::GpuTensor,
7964 zn: &CudaSlice<f32>,
7965 t: usize,
7966 n_embd: usize,
7967 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7968 let n_ff = ffn_gate.out_features();
7969 let (zq, zd) = e.quantize_q8_1(zn, t, n_embd)?;
7970 if Engine::tk_ffn_dual_on()
7971 && let Some(((g, gs), (u, us))) =
7972 e.matmul_decode_exact_dual_pre(ffn_gate, ffn_up, &zq, &zd, t)?
7973 {
7974 if e.uses_q8_1_fast(ffn_down) {
7975 let (aq, ad) = e.silu_mul_scaled_q8_1(&g, &u, gs, us, t * n_ff)?;
7976 return e.matmul_decode_exact_pre(ffn_down, &aq, &ad, t);
7977 }
7978 let mut act = e.uninit(t * n_ff)?;
7979 e.silu_mul_scaled(&g, &u, gs, us, &mut act, t * n_ff)?;
7980 let (aq, ad) = e.quantize_q8_1(&act, t, n_ff)?;
7981 return e.matmul_pre(ffn_down, &aq, &ad, &act, t);
7982 }
7983 // v1 singles chain (seam off or dual-refused) — the pre-slice-B body verbatim.
7984 let g = e.matmul_pre(ffn_gate, &zq, &zd, zn, t)?;
7985 let u = e.matmul_pre(ffn_up, &zq, &zd, zn, t)?;
7986 let mut act = e.uninit(t * n_ff)?;
7987 e.silu_mul(&g, &u, &mut act, t * n_ff)?;
7988 let (aq, ad) = e.quantize_q8_1(&act, t, n_ff)?;
7989 e.matmul_pre(ffn_down, &aq, &ad, &act, t)
7990 }
7991
7992 /// ONE t-parallel FULL-ATTENTION layer (attn_norm + fa mixer + post_attn_norm + FFN +
7993 /// tap) — extracted from the walk exactly like `qwen35_tparallel_linear_layer` so the
7994 /// eager walk and the slice-4c captured full-verify graphs execute the SAME body (a
7995 /// second copy is how dispatch mirrors drift — the verify_layers extraction lesson).
7996 ///
7997 /// `args.graph_cap = Some((table, off, rung_end))` is the captured-graph mode:
7998 /// - kv base-pointer pairs come from the ctx-owned persistent table at `off` (a fresh
7999 /// generation's cache lands at new addresses that only the per-verify table refresh
8000 /// knows — the slice-3 baked-address lesson);
8001 /// - the seqs twins size partials/grid at `rung_end` and pin `split_keys` to the
8002 /// rung's ladder value: `n_splits_max` is pure stride, splits >= ns_eff write the
8003 /// EMPTY partial the combine never reads, and every per-row T_kv derives in-kernel
8004 /// from `pos_seq[z]` — so one captured launch replays bit-identically for every
8005 /// round whose rows all sit inside the rung;
8006 /// - the host len bump moves to the replay caller (captured host code does not
8007 /// re-run at replay).
8008 /// Graph mode REFUSES any round the batched arm cannot take: the per-row fallback
8009 /// host-branches on t_kv and must never be captured.
8010 #[allow(clippy::too_many_arguments)]
8011 fn qwen35_tparallel_fa_layer(
8012 &self,
8013 e: &Engine,
8014 il: usize,
8015 x: &CudaSlice<f32>,
8016 t: usize,
8017 cache: &mut Cache,
8018 args: FaLayerArgs<'_>,
8019 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8020 use cudarc::driver::DevicePtr;
8021 let cfg = &self.cfg;
8022 let n_embd = cfg.n_embd as usize;
8023 let eps = cfg.rms_eps;
8024 let head_dim_global = cfg.head_dim_k as usize;
8025 let layer = &self.layers[il];
8026 let FaLayerArgs {
8027 pos_d,
8028 pos_rows,
8029 pos0,
8030 seqs_append,
8031 batch_fa_on,
8032 graph_cap,
8033 stream,
8034 ckpt,
8035 } = args;
8036
8037 // ---- attn_norm + q8_1 quantize at m=T (row-indexed == per-row) ----
8038 let anorm = layer.attn_norm.float_data();
8039 let mut xn = e.uninit(t * n_embd)?;
8040 e.rms_norm(x, anorm, &mut xn, n_embd, t, eps)?;
8041 let (hq, hd) = e.quantize_q8_1(&xn, t, n_embd)?;
8042
8043 let mixed: CudaSlice<f32> = match &layer.mixer {
8044 Mixer::Mla(_) => crate::hybrid::mla_path_unimplemented("tensor-parallel attention"),
8045 Mixer::Kda(_) => crate::hybrid::kda_path_unimplemented("T-parallel attention"),
8046 // STREAM ARM (2b, lane/draftcost-moe): under a device position counter the
8047 // per-row serving-kernel chain cannot run (host state swaps keyed on host
8048 // row index are fine, but the stream COMMIT needs the GdnStash for its _dc
8049 // rebuild — the per-row chain only produces per-column clones). GDN rides
8050 // `linear_attn_verify_t`: batched q8_1-class projections, stash-producing,
8051 // and its one-scan recurrence is pinned bit-identical to T chained T=1
8052 // steps (its header + kernel-check). Position-independent, so no counter
8053 // plumbing is needed. Guards mirror the generic call site exactly.
8054 Mixer::Linear(la) if stream.is_some() => {
8055 if !(t >= 3 || (t == 2 && spec_m2()))
8056 || !self.mixer_in_q8_1_fast(e, &layer.mixer)
8057 || !e.uses_q8_1_fast(&la.ssm_out)
8058 {
8059 return Err("qwen35 stream verify: GDN batched arm requires t>=3 \
8060 (or MEMRA_SPEC_M2 at t=2) and q8_1-fast projections"
8061 .into());
8062 }
8063 let want = ckpt.is_some();
8064 let (out, stash) =
8065 self.linear_attn_verify_t(e, la, &xn, Some((&hq, &hd)), t, cache, il, want)?;
8066 if let (Some(ck), Some(st)) = (ckpt, stash) {
8067 ck.gdn[il] = Some(st);
8068 }
8069 out
8070 }
8071 Mixer::Linear(_) => {
8072 unreachable!("linear layers ride qwen35_tparallel_linear_layer")
8073 }
8074 Mixer::Full(fa) => {
8075 let geometry = cfg.full_attention_geometry_at(il as u32);
8076 let n_head = geometry.n_head as usize;
8077 let n_head_kv = geometry.n_head_kv as usize;
8078 let head_dim = geometry.head_dim_k as usize;
8079 let rope_dims = geometry.n_rot as usize;
8080 let rope_base = geometry.rope_base;
8081 let scale = geometry.attention_scale();
8082 // Batched projections: one weight read serves all T rows.
8083 // GROUP-3 twin (trunk-kernels slice D): q/k/v in ONE launch — the group4
8084 // kernel with n3=0, bit-identical per (tensor, token, row) to the three
8085 // singles; refused or MEMRA_TK_FA_GROUP=0 -> singles byte-for-byte.
8086 let (qf, mut k, v) = match e.matmul_decode_exact_group3_pre(
8087 [&fa.wq, &fa.wk, &fa.wv],
8088 &hq,
8089 &hd,
8090 t,
8091 )? {
8092 Some(mut g3) => {
8093 let v = g3.pop().unwrap();
8094 let k = g3.pop().unwrap();
8095 let qf = g3.pop().unwrap();
8096 (qf, k, v)
8097 }
8098 None => (
8099 e.matmul_pre(&fa.wq, &hq, &hd, &xn, t)?,
8100 e.matmul_pre(&fa.wk, &hq, &hd, &xn, t)?,
8101 e.matmul_pre(&fa.wv, &hq, &hd, &xn, t)?,
8102 ),
8103 };
8104 let gated =
8105 geometry.attention_gate == memra_gguf::config::AttentionGateKind::FusedQ;
8106 let (mut q, gate) = if gated {
8107 let mut qs = e.uninit(t * n_head * head_dim)?;
8108 let mut gs = e.uninit(t * n_head * head_dim)?;
8109 e.q_gate_split(&qf, &mut qs, &mut gs, head_dim, n_head, t)?;
8110 (qs, Some(gs))
8111 } else {
8112 (qf, None)
8113 };
8114 let mut qn = e.uninit(t * n_head * head_dim)?;
8115 e.rms_norm(
8116 &q,
8117 fa.q_norm.float_data(),
8118 &mut qn,
8119 head_dim,
8120 t * n_head,
8121 eps,
8122 )?;
8123 q = qn;
8124 let mut kn = e.uninit(t * n_head_kv * head_dim)?;
8125 e.rms_norm(
8126 &k,
8127 fa.k_norm.float_data(),
8128 &mut kn,
8129 head_dim,
8130 t * n_head_kv,
8131 eps,
8132 )?;
8133 k = kn;
8134 e.rope_neox(
8135 &mut q, pos_d, head_dim, rope_dims, n_head, t, rope_base, 1.0,
8136 )?;
8137 e.rope_neox(
8138 &mut k, pos_d, head_dim, rope_dims, n_head_kv, t, rope_base, 1.0,
8139 )?;
8140
8141 // Per-row append + attend: row r sees rows 0..r in KV (causal within the
8142 // draft), each through the b_n=1 serving kernels at its own t_kv.
8143 let q_dim = n_head * head_dim;
8144 let kv_dim = n_head_kv * head_dim;
8145 let mut attn = e.uninit(t * q_dim)?;
8146 let (kdk, kdv, ktb, vtb, len0, kv_local) = {
8147 let kvl = cache.kv[il].as_ref().unwrap();
8148 // [2T] interleaved k,v base pointers: entry pair z serves row z of
8149 // the batched twins; the per-row fallback reads pair 0 (same cache
8150 // for every row of one layer). Graph mode reads the ctx table.
8151 let local: Option<CudaSlice<u64>> = match graph_cap {
8152 Some(_) => None,
8153 None => {
8154 let s = &e.gpu.stream();
8155 let (pk, _g) = kvl.k.device_ptr(s);
8156 let (pv, _g2) = kvl.v.device_ptr(s);
8157 let mut tbl = Vec::with_capacity(2 * t);
8158 for _ in 0..t {
8159 tbl.push(pk);
8160 tbl.push(pv);
8161 }
8162 Some(e.htod_u64(&tbl)?)
8163 }
8164 };
8165 (
8166 kvl.kv_dim_k,
8167 kvl.kv_dim_v,
8168 kvl.k_tok_bytes,
8169 kvl.v_tok_bytes,
8170 kvl.len,
8171 local,
8172 )
8173 };
8174 let (kv_tbl, kv_off): (&CudaSlice<u64>, usize) = match graph_cap {
8175 Some((tb, off, _)) => (tb, off),
8176 None => (kv_local.as_ref().expect("built above"), 0),
8177 };
8178 // Slice 4 (fa/append rows — see dspark_fa_rows_on): the whole per-row
8179 // section batches into the z-batched serving twins when every row of
8180 // this round takes the v4-seqs arm on ONE fa_split_keys rung. Both
8181 // guards are evaluated at the round's FIRST and LAST t_kv — the
8182 // eligibility window (vec floor .. v4 max) and each split-ladder rung
8183 // are intervals in t_kv, so ends-inside means all-inside (the straddle
8184 // law). Appending all T rows before any attend is read-equivalent to
8185 // the interleaved order: row r's walk reads keys 0..len0+r only, and
8186 // rows > r land at slots it never touches; every written cache row is
8187 // the per-token appender's exact warp program (kernel-check pinned).
8188 let t_kv_first = len0 + 1;
8189 let t_kv_last = len0 + t;
8190 let rows_batched = t >= 2
8191 && seqs_append
8192 && batch_fa_on
8193 && dspark_fa_rows_on()
8194 // the z-batched twins read stacked rows at the CACHE's kv dims;
8195 // the projection stack is [T, n_head_kv*head_dim] — they must be
8196 // the same stride or row z misaligns (true for this family; the
8197 // guard keeps any asymmetric-kv model on the per-row loop).
8198 && kdk == kv_dim
8199 && kdv == kv_dim
8200 && crate::fa_seqs_eligible(t_kv_first, head_dim_global)
8201 && crate::fa_seqs_eligible(t_kv_last, head_dim_global)
8202 && crate::fa_split_keys(t_kv_first, cfg.n_head_kv as usize)
8203 == crate::fa_split_keys(t_kv_last, cfg.n_head_kv as usize);
8204 // Sizing: eager = exact round bound; graph mode = the rung end (stride +
8205 // grid only — bytes proven equal above). Capture-time invariants refuse
8206 // loudly rather than bake a divergent body.
8207 let (size_kv_max, sp) = match graph_cap {
8208 Some((_, _, rung)) => {
8209 if !rows_batched {
8210 return Err(format!(
8211 "fa graph capture: layer {il} round is not batchable \
8212 (t_kv {t_kv_first}..{t_kv_last}) — the per-row fallback \
8213 must never be captured"
8214 )
8215 .into());
8216 }
8217 let sp_r = crate::fa_split_keys(rung, cfg.n_head_kv as usize);
8218 if t_kv_last > rung
8219 || sp_r != crate::fa_split_keys(t_kv_last, cfg.n_head_kv as usize)
8220 {
8221 return Err(format!(
8222 "fa graph capture: rung {rung} does not cover round \
8223 t_kv {t_kv_first}..{t_kv_last} on one split ladder step"
8224 )
8225 .into());
8226 }
8227 (rung, sp_r)
8228 }
8229 None => (
8230 t_kv_last,
8231 crate::fa_split_keys(t_kv_last, cfg.n_head_kv as usize),
8232 ),
8233 };
8234 if let Some((_, ctr)) = stream {
8235 // STREAM ARM (2b): one batched dc append + the multi-row dc attention
8236 // — the generic stream arm's exact shape (rows kernels are pinned
8237 // byte-identical to the per-row programs by kernel-check). Host len
8238 // stays a stale lower bound; the burst drain reconciles it.
8239 let kvl = cache.kv[il].as_mut().unwrap();
8240 e.append_kv_quantized_rows_dc(
8241 &k,
8242 &v,
8243 &mut kvl.k,
8244 &mut kvl.v,
8245 ctr,
8246 t,
8247 kdk,
8248 kdv,
8249 ktb,
8250 vtb,
8251 Engine::kv_fp8_on(),
8252 )?;
8253 let upper = (kvl.len + t + 64).min(cache.max_ctx);
8254 let k_view = e.view_u8(&kvl.k, upper * ktb);
8255 let v_view = e.view_u8(&kvl.v, upper * vtb);
8256 e.fa_decode_rows_dc(
8257 &q, &k_view, &v_view, &mut attn, head_dim, n_head, n_head_kv, ctr, upper,
8258 t, scale, ktb, vtb, 0, false,
8259 )?;
8260 } else if rows_batched {
8261 e.append_kv_quantized_seqs(
8262 &k,
8263 &v,
8264 &kv_tbl.slice(kv_off..kv_off + 2 * t),
8265 pos_d,
8266 t,
8267 kdk,
8268 kdv,
8269 ktb,
8270 vtb,
8271 )?;
8272 if graph_cap.is_none() {
8273 cache.kv[il].as_mut().unwrap().len += t;
8274 }
8275 e.fa_decode_batch_seqs_v4(
8276 &q,
8277 &kv_tbl.slice(kv_off..kv_off + 2 * t),
8278 pos_d,
8279 &mut attn,
8280 head_dim,
8281 n_head,
8282 n_head_kv,
8283 t,
8284 size_kv_max,
8285 scale,
8286 sp,
8287 ktb,
8288 vtb,
8289 )?;
8290 } else {
8291 if pos_rows.is_none() {
8292 // Stream-aware for symmetry with pos_d (the stream FA arm rides
8293 // the dc rows kernels above and never reaches this fallback).
8294 *pos_rows = Some(match stream {
8295 Some((_, ctr)) => (0..t)
8296 .map(|r| {
8297 let mut b = e.alloc_uninit::<i32>(1)?;
8298 e.i32_copy_add(ctr, &mut b, r as i32)?;
8299 Ok(b)
8300 })
8301 .collect::<Result<_, Box<dyn std::error::Error>>>()?,
8302 None => (0..t)
8303 .map(|r| e.htod_i32(&[(pos0 + r) as i32]))
8304 .collect::<Result<_, _>>()?,
8305 });
8306 }
8307 let pos_rows = pos_rows.as_ref().unwrap();
8308 #[allow(clippy::needless_range_loop)]
8309 // allow: the explicit index loop keeps the offset arithmetic visible and aligned with the device-side indexing
8310 for r in 0..t {
8311 // Owned per-row scratch: the b_n=1 kernels take packed batch buffers
8312 // whose row 0 is this row (arithmetic-free materialization copies,
8313 // same as decode's per-seq fallback arm).
8314 let mut k_row = e.uninit(kv_dim)?;
8315 e.dtod_copy_view(&k.slice(r * kv_dim..(r + 1) * kv_dim), &mut k_row)?;
8316 let mut v_row = e.uninit(kv_dim)?;
8317 e.dtod_copy_view(&v.slice(r * kv_dim..(r + 1) * kv_dim), &mut v_row)?;
8318 let pos_row = &pos_rows[r];
8319 let kvl = cache.kv[il].as_mut().unwrap();
8320 if seqs_append {
8321 e.append_kv_quantized_seqs(
8322 &k_row,
8323 &v_row,
8324 &kv_tbl.slice(kv_off..kv_off + 2),
8325 pos_row,
8326 1,
8327 kdk,
8328 kdv,
8329 ktb,
8330 vtb,
8331 )?;
8332 kvl.len += 1;
8333 } else {
8334 e.append_kv_quantized_view(
8335 &k_row.slice(0..kv_dim),
8336 &v_row.slice(0..kv_dim),
8337 &mut kvl.k,
8338 &mut kvl.v,
8339 kvl.len,
8340 kvl.kv_dim_k,
8341 kvl.kv_dim_v,
8342 kvl.k_tok_bytes,
8343 kvl.v_tok_bytes,
8344 Engine::kv_fp8_on(),
8345 )?;
8346 kvl.len += 1;
8347 }
8348 let t_kv = kvl.len;
8349 let mut q_row = e.uninit(q_dim)?;
8350 e.dtod_copy_view(&q.slice(r * q_dim..(r + 1) * q_dim), &mut q_row)?;
8351 let mut a_row = e.uninit(q_dim)?;
8352 if batch_fa_on && crate::fa_seqs_eligible(t_kv, head_dim_global) {
8353 let sp0_r = crate::fa_split_keys(t_kv, cfg.n_head_kv as usize);
8354 e.fa_decode_batch_seqs_v4(
8355 &q_row,
8356 &kv_tbl.slice(kv_off..kv_off + 2),
8357 pos_row,
8358 &mut a_row,
8359 head_dim,
8360 n_head,
8361 n_head_kv,
8362 1,
8363 t_kv,
8364 scale,
8365 sp0_r,
8366 ktb,
8367 vtb,
8368 )?;
8369 } else {
8370 let k_view = e.view_u8(&kvl.k, t_kv * kvl.k_tok_bytes);
8371 let v_view = e.view_u8(&kvl.v, t_kv * kvl.v_tok_bytes);
8372 let mut a_view = a_row.slice_mut(0..q_dim);
8373 e.fa_decode_kvmod_view(
8374 &q_row.slice(0..q_dim),
8375 &k_view,
8376 &v_view,
8377 &mut a_view,
8378 head_dim,
8379 n_head,
8380 n_head_kv,
8381 t_kv,
8382 scale,
8383 kvl.k_tok_bytes,
8384 kvl.v_tok_bytes,
8385 Engine::kv_fp8_on(),
8386 )?;
8387 }
8388 e.dtod_copy_into(&a_row, &mut attn, r * q_dim)?;
8389 }
8390 }
8391
8392 // Output gate (element-wise) + o-proj at m=T.
8393 let attn_g = match &gate {
8394 Some(g) => {
8395 let n = t * q_dim;
8396 let mut gsig = e.uninit(n)?;
8397 e.sigmoid(g, &mut gsig, n)?;
8398 let mut ag = e.uninit(n)?;
8399 e.mul(&attn, &gsig, &mut ag, n)?;
8400 ag
8401 }
8402 None => attn,
8403 };
8404 e.matmul(&fa.wo, &attn_g, t)?
8405 }
8406 };
8407
8408 // ---- residual add + post_attn_norm + FFN at m=T (serving dispatch verbatim) ----
8409 let pnorm = layer.post_attn_norm.float_data();
8410 let mut x1 = e.uninit(t * n_embd)?;
8411 let mut zn = e.uninit(t * n_embd)?;
8412 e.add_rms_norm(x, &mixed, pnorm, &mut x1, &mut zn, n_embd, t, eps)?;
8413 let ffn_out = match &layer.ffn {
8414 crate::hybrid::Ffn::Dense {
8415 ffn_gate,
8416 ffn_up,
8417 ffn_down,
8418 } => {
8419 assert!(
8420 self.cfg.m3.is_none(),
8421 "qwen35 t-parallel verify: M3 swigluoai FFN not yet batched"
8422 );
8423 self.qwen35_tparallel_dense_ffn(e, ffn_gate, ffn_up, ffn_down, &zn, t, n_embd)?
8424 }
8425 crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il_zq8(e, m, &zn, None, t, il as u16)?,
8426 };
8427 let mut x2 = e.uninit(t * n_embd)?;
8428 e.add(&x1, &ffn_out, &mut x2, t * n_embd)?;
8429 // dspark drafter tap (no-op when no sink armed): post-layer residual verify rows
8430 self.dflash_tap(e, cache, il, &x2, t)?;
8431 Ok(x2)
8432 }
8433
8434 /// ONE t-parallel LINEAR layer (attn_norm + gdn mixer + post_attn_norm + FFN + tap) —
8435 /// the exact body the old in-loop Linear arm ran, extracted so the eager walk and the
8436 /// slice-3 captured segments execute the SAME code (a second copy is how dispatch
8437 /// mirrors drift — the verify_layers extraction lesson). Two deliberate changes, both
8438 /// bit-identical by construction:
8439 /// - the gdn ping-pong host swap moves from per-row to ONE end-of-body swap (t odd):
8440 /// the device sequence is driven entirely by the 6-entry pointer table, which
8441 /// already encodes both parities; the ckpt stash reads name row r's out buffer
8442 /// directly (r even -> alt handle, odd -> canonical) — the same physical bytes the
8443 /// legacy post-swap clone read.
8444 /// - `stash` (slice-3 ctx): persistent per-layer slabs written by copy_into instead of
8445 /// per-row clone_dtod allocs — same bytes, capture-legal (no per-round host objects).
8446 /// `table_src` = (persistent pointer table, offset) when the ctx owns the tables;
8447 /// None builds the per-verify table exactly as before.
8448 #[allow(clippy::too_many_arguments)]
8449 fn qwen35_tparallel_linear_layer(
8450 &self,
8451 e: &Engine,
8452 il: usize,
8453 x: &CudaSlice<f32>,
8454 t: usize,
8455 cache: &mut Cache,
8456 ckpt: Option<&mut VerifyCkpt>,
8457 stash: Option<(&mut CudaSlice<f32>, &mut CudaSlice<f32>)>,
8458 table_src: Option<(&CudaSlice<u64>, usize)>,
8459 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8460 use cudarc::driver::DevicePtr;
8461 let cfg = &self.cfg;
8462 let n_embd = cfg.n_embd as usize;
8463 let eps = cfg.rms_eps;
8464 let layer = &self.layers[il];
8465 let Mixer::Linear(la) = &layer.mixer else {
8466 return Err("qwen35_tparallel_linear_layer on a non-linear layer".into());
8467 };
8468 // ---- attn_norm + q8_1 quantize at m=T (row-indexed == per-row) ----
8469 let anorm = layer.attn_norm.float_data();
8470 let mut xn = e.uninit(t * n_embd)?;
8471 e.rms_norm(x, anorm, &mut xn, n_embd, t, eps)?;
8472 let (hq, hd) = e.quantize_q8_1(&xn, t, n_embd)?;
8473
8474 let geometry = la.geometry;
8475 let d_state = geometry.key_head_dim as usize;
8476 let num_k = geometry.key_heads as usize;
8477 let num_v = geometry.value_heads as usize;
8478 let d_conv = geometry.conv_kernel as usize;
8479 let key_dim = d_state * num_k;
8480 let value_dim = geometry.value_head_dim as usize * num_v;
8481 let conv_dim = key_dim * 2 + value_dim;
8482 let gdn_scale = 1.0 / (d_state as f32).sqrt();
8483
8484 // ---- batched projections: one weight read for all T rows ----
8485 // GROUP-4 twin (trunk-kernels slice C): the whole 4-tuple in ONE launch, bit-identical
8486 // per (tensor, token, row) to the four singles; refused (layout/tier) or
8487 // MEMRA_TK_GDN_GROUP=0 -> the singles chain byte-for-byte.
8488 let (qkv_mixed, z, beta_raw, alpha) = match e.matmul_decode_exact_group4_pre(
8489 [&la.wqkv, &la.wqkv_gate, &la.ssm_beta, &la.ssm_alpha],
8490 &hq,
8491 &hd,
8492 t,
8493 )? {
8494 Some(mut g4) => {
8495 let alpha = g4.pop().unwrap();
8496 let beta_raw = g4.pop().unwrap();
8497 let z = g4.pop().unwrap();
8498 let qkv_mixed = g4.pop().unwrap();
8499 (qkv_mixed, z, beta_raw, alpha)
8500 }
8501 None => (
8502 e.matmul_pre(&la.wqkv, &hq, &hd, &xn, t)?,
8503 e.matmul_pre(&la.wqkv_gate, &hq, &hd, &xn, t)?,
8504 e.matmul_pre(&la.ssm_beta, &hq, &hd, &xn, t)?,
8505 e.matmul_pre(&la.ssm_alpha, &hq, &hd, &xn, t)?,
8506 ),
8507 };
8508 let beta_w = la.ssm_beta.out_features();
8509 let alpha_w = la.ssm_alpha.out_features();
8510 let qkv_w = la.wqkv.out_features();
8511
8512 // ---- per-row state chain through the b_n=1 serving kernels ----
8513 // 6-entry alternating pointer table expresses the ping-pong without a rebuild per
8514 // row: even rows scan s0 -> s1, odd rows s1 -> s0.
8515 let table_local: Option<CudaSlice<u64>> = match table_src {
8516 Some(_) => None,
8517 None => {
8518 let rl = cache.recur[il].as_ref().unwrap();
8519 let s = &e.gpu.stream();
8520 let (pc, _g0) = rl.conv_state.device_ptr(s);
8521 let (p0, _g1) = rl.ssm_state.device_ptr(s);
8522 let (p1, _g2) = rl.ssm_state_alt.device_ptr(s);
8523 Some(e.htod_u64(&[pc, p0, p1, pc, p1, p0])?)
8524 }
8525 };
8526 let (table, toff): (&CudaSlice<u64>, usize) = match table_src {
8527 Some((tb, off)) => (tb, off),
8528 None => (table_local.as_ref().unwrap(), 0),
8529 };
8530 let mut o_all = e.uninit(t * value_dim)?;
8531 let mut col_states: Option<Vec<(CudaSlice<f32>, CudaSlice<f32>)>> =
8532 if ckpt.is_some() && stash.is_none() && t >= 2 {
8533 Some(Vec::with_capacity(t - 1))
8534 } else {
8535 None
8536 };
8537 let mut stash = stash;
8538 // Per-row scratch reused across rows (uninit is cheap but not free at
8539 // 48 layers x T rows); row inputs/outputs pass as VIEWS into the packed
8540 // [T, ...] buffers — zero arithmetic-free copies in this loop.
8541 let mut conv_out = e.uninit(conv_dim)?;
8542 let mut q_l2 = e.uninit(value_dim)?;
8543 let mut k_l2 = e.uninit(value_dim)?;
8544 let mut v_gd = e.uninit(value_dim)?;
8545 let mut beta_b = e.uninit(num_v)?;
8546 let mut g_log = e.uninit(num_v)?;
8547 for r in 0..t {
8548 let base = toff + if r % 2 == 0 { 0 } else { 3 };
8549 let conv_view = table.slice(base..base + 1);
8550 let in_view = table.slice(base + 1..base + 2);
8551 let out_view = table.slice(base + 2..base + 3);
8552 e.ssm_conv1d_fused_decode_b_view(
8553 &qkv_mixed.slice(r * qkv_w..(r + 1) * qkv_w),
8554 &conv_view,
8555 la.ssm_conv1d.float_data(),
8556 &mut conv_out,
8557 conv_dim,
8558 d_conv,
8559 1,
8560 )?;
8561 e.gdn_prep_decode_b_view(
8562 &conv_out,
8563 &beta_raw.slice(r * beta_w..(r + 1) * beta_w),
8564 &alpha.slice(r * alpha_w..(r + 1) * alpha_w),
8565 la.ssm_dt.float_data(),
8566 la.ssm_a.float_data(),
8567 &mut q_l2,
8568 &mut k_l2,
8569 &mut v_gd,
8570 &mut beta_b,
8571 &mut g_log,
8572 d_state,
8573 num_v,
8574 num_k,
8575 key_dim,
8576 eps,
8577 conv_dim,
8578 1,
8579 )?;
8580 let mut o_row = o_all.slice_mut(r * value_dim..(r + 1) * value_dim);
8581 e.gdn_scan_s128_batched_view(
8582 &q_l2, &k_l2, &v_gd, &g_log, &beta_b, &in_view, &out_view, &mut o_row, num_v, 1,
8583 gdn_scale,
8584 )?;
8585 if r + 1 < t {
8586 // Row r's out buffer: even rows write s1 (the alt handle — no swaps ran),
8587 // odd rows write s0 — the same physical state the legacy post-swap
8588 // canonical clone read.
8589 let rl = cache.recur[il]
8590 .as_ref()
8591 .ok_or("qwen35 linear verify layer has no recurrent state")?;
8592 let ssm_src = if r % 2 == 0 {
8593 &rl.ssm_state_alt
8594 } else {
8595 &rl.ssm_state
8596 };
8597 match stash.as_mut() {
8598 Some((conv_slab, ssm_slab)) => {
8599 // BOTH stash reads go through the pointer table at run time: the
8600 // ssm handles ping-pong between rounds, and the ctx (with its
8601 // captured graphs) outlives the Cache — a fresh generation's
8602 // conv/ssm buffers land at new addresses that only the per-round
8603 // table refresh knows. A baked direct copy would read freed
8604 // memory (parity was the slice-3 smoke divergence; cache
8605 // lifetime is the cross-generation twin).
8606 e.copy_indirect_src_f32(
8607 &conv_view,
8608 conv_slab,
8609 r * conv_dim * (d_conv - 1),
8610 conv_dim * (d_conv - 1),
8611 )?;
8612 // The ssm handles PING-PONG between rounds: a captured direct
8613 // copy would bake the capture-time physical buffer and read the
8614 // wrong parity after any odd-vt round (the slice-3 smoke
8615 // divergence). Read the src address from row r's OUT table
8616 // entry at run time — the same entry the scan just wrote.
8617 e.copy_indirect_src_f32(
8618 &out_view,
8619 ssm_slab,
8620 r * d_state * d_state * num_v,
8621 d_state * d_state * num_v,
8622 )?;
8623 }
8624 None => {
8625 if let Some(states) = col_states.as_mut() {
8626 states.push((e.clone_dtod(&rl.conv_state)?, e.clone_dtod(ssm_src)?));
8627 }
8628 }
8629 }
8630 }
8631 }
8632 // ONE end-of-body parity swap (t odd) — the legacy loop swapped per row; the net
8633 // handle motion is identical and the device sequence never read the handles.
8634 if t % 2 == 1 {
8635 let rl = cache.recur[il].as_mut().unwrap();
8636 std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
8637 }
8638 if let (Some(checkpoint), Some(states)) = (ckpt, col_states) {
8639 checkpoint.cols[il] = Some(states);
8640 }
8641
8642 // ---- batched gated norm + out-projection at m=T ----
8643 let mixed = if e.uses_q8_1_fast(&la.ssm_out) {
8644 let (gq, gd) = e.gated_rmsnorm_q8_1(
8645 &o_all,
8646 la.ssm_norm.float_data(),
8647 &z,
8648 d_state,
8649 t * num_v,
8650 eps,
8651 )?;
8652 let g0 = e.zeros(0)?;
8653 e.matmul_pre(&la.ssm_out, &gq, &gd, &g0, t)?
8654 } else {
8655 let mut gn = e.uninit(t * value_dim)?;
8656 e.gated_rmsnorm(
8657 &o_all,
8658 la.ssm_norm.float_data(),
8659 &z,
8660 &mut gn,
8661 d_state,
8662 t * num_v,
8663 eps,
8664 )?;
8665 e.matmul(&la.ssm_out, &gn, t)?
8666 };
8667
8668 // ---- residual add + post_attn_norm + FFN at m=T (serving dispatch verbatim) ----
8669 let pnorm = layer.post_attn_norm.float_data();
8670 let mut x1 = e.uninit(t * n_embd)?;
8671 let mut zn = e.uninit(t * n_embd)?;
8672 e.add_rms_norm(x, &mixed, pnorm, &mut x1, &mut zn, n_embd, t, eps)?;
8673 let ffn_out = match &layer.ffn {
8674 crate::hybrid::Ffn::Dense {
8675 ffn_gate,
8676 ffn_up,
8677 ffn_down,
8678 } => {
8679 assert!(
8680 self.cfg.m3.is_none(),
8681 "qwen35 t-parallel verify: M3 swigluoai FFN not yet batched"
8682 );
8683 self.qwen35_tparallel_dense_ffn(e, ffn_gate, ffn_up, ffn_down, &zn, t, n_embd)?
8684 }
8685 crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il_zq8(e, m, &zn, None, t, il as u16)?,
8686 };
8687 let mut x2 = e.uninit(t * n_embd)?;
8688 e.add(&x1, &ffn_out, &mut x2, t * n_embd)?;
8689 // dspark drafter tap (no-op when no sink armed): post-layer residual verify rows
8690 self.dflash_tap(e, cache, il, &x2, t)?;
8691 Ok(x2)
8692 }
8693
8694 /// PP-N STAGE SUBGRAPH of the verify trunk: layers `[lo, hi)` of `decode_step_t_core_stream`'s
8695 /// walk, verbatim. Enters with a MATERIALIZED `[T, n_embd]` residual (no pending fusion pair
8696 /// carried in from outside the range) and exits with the range's final residual materialized
8697 /// (the trailing add executed) — exactly the `decode_layers_eager(lo, hi)` contract, T rows
8698 /// instead of one.
8699 ///
8700 /// EXTRACTED (lane/pp2-spec 2026-08-06) rather than duplicated: `decode_step_t_core_stream` IS
8701 /// the single funnel every verify forward reaches, and its per-layer dispatch MIRRORING (norm
8702 /// fusion per layer, the t>=3/spec_m2 batched-linear window, the fused-q8 FFN chain, the
8703 /// decode-exact projections) is what makes verify bit-identical to eager decode. A second copy
8704 /// for the split arm is how those mirrors drift apart on the next lever. The unsplit body now
8705 /// calls this with `(0, n_layers)`, so the whole-trunk path and every stage range run the SAME
8706 /// code — there is no "split version" of the verify math.
8707 ///
8708 /// Bit-identity of a cut rests on the same kernel-check-pinned identity the eager arm's cut
8709 /// does — `add_rms_norm_q8_1 == add then rms_norm_q8_1` at nrows=T — because the ONLY thing a
8710 /// fence changes is that the cross-layer fusion carry breaks at `hi-1` and is re-materialized
8711 /// as an explicit `add`. `decode-batch-gate --mode ppspec` verifies end-to-end on real weights.
8712 #[allow(clippy::too_many_arguments)]
8713 fn verify_layers(
8714 &self,
8715 e: &Engine,
8716 mut x: CudaSlice<f32>,
8717 lo: usize,
8718 hi: usize,
8719 pos_d: &CudaSlice<i32>,
8720 pos0: usize,
8721 t: usize,
8722 cache: &mut Cache,
8723 mut ckpt: Option<&mut VerifyCkpt>,
8724 stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
8725 graphs: Option<&mut DsparkVerifyGraphs>,
8726 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8727 if self.sliding_gated_moe_batch_program() {
8728 if stream.is_some() {
8729 return Err(
8730 "step35 has no ROUND-STREAM verify arm (the device-counter _dc twins \
8731 cannot express the SWA offset KV view)"
8732 .into(),
8733 );
8734 }
8735 return self.step35_verify_batch_layers(e, x, lo, hi, pos0, t, cache);
8736 }
8737 if self.batched_serving_numeric_class() {
8738 return self.qwen35_verify_batch_layers(
8739 e,
8740 x,
8741 lo,
8742 hi,
8743 pos0,
8744 t,
8745 cache,
8746 ckpt.take(),
8747 stream,
8748 graphs,
8749 );
8750 }
8751 let n_embd = self.cfg.n_embd as usize;
8752 let eps = self.cfg.rms_eps;
8753 // CROSS-LAYER ADD+NORM FUSION (lane/vt-fixes fix 2, mirroring decode_step_h's
8754 // launch-arc form): layer il's post-FFN residual add (x2 = x1 + ffn_out) and layer
8755 // il+1's attn_norm(+quantize) are consecutive row-wise ops — ONE add_rms_norm_q8_1
8756 // launch at nrows=t does all three (bit-identity pinned by the T-row kernel-check
8757 // arms). Carry the un-added (x1, ffn_out) pair; the fused launch materializes x2 (the
8758 // residual the next layer needs) as its `res` output. Falls back to the separate add
8759 // when the next layer is off the fused-q8 path.
8760 let mut pending: Option<(CudaSlice<f32>, CudaSlice<f32>)> = None;
8761 for il in lo..hi {
8762 let layer = &self.layers[il];
8763 // DISPATCH-MIRRORED attn-input RMSNorm (FP-order lesson #8): eager decode fuses the
8764 // 1024-thread rms_norm_q8_1 ONLY when every mixer projection is q8_1-fast; layers with
8765 // Float projections (ssm_beta/ssm_alpha on layers 1/2/4 of the 9B NVFP4 GGUF) take the
8766 // UNFUSED 256-thread rms_norm. The verify norm must mirror that PER-LAYER choice —
8767 // blockDim changes the sum-of-squares reduce order, and the ULP shift amplifies through
8768 // the GDN recurrence into argmax flips (measured: 9B text prompt, 1 ULP at layer 2 ->
8769 // 2.3e-1 logit maxdiff at the head -> K=1..8 divergence at a 0.03-margin token).
8770 let mixer_fast = self.mixer_in_q8_1_fast(e, &layer.mixer);
8771 let norm_fused = std::env::var("MEMRA_NO_FUSE_NORMQ").is_err() && mixer_fast;
8772 // BATCHED EPILOGUE RE-FUSE (lane/vt-fixes fix 2, 2026-08-03): when the norm is
8773 // dispatch-fused AND every consumer of `h` reads only its q8_1 form (Full mixer:
8774 // projections only; Linear mixer: the batched arm — the per-column fallback needs
8775 // f32 h), emit the attn-input norm DIRECTLY as q8_1 via `rms_norm_q8_1` at nrows=t
8776 // (row-indexed kernel — the T-row launch is the per-row m=1 program, kernel-check
8777 // pins bit-identity vs rms_norm_decode -> quantize_q8_1). Kills the standalone
8778 // quantize launch(es) + the f32 h HBM round-trip that decode never pays.
8779 // step35 (Full mixer) is the third case that needs f32 `h`: its verify arm is a
8780 // per-ROW replay of the eager decode mixer, whose `pre_q` contract is a single row —
8781 // a T-row q8_1 pair cannot be handed to it, and re-deriving per-row q8_1 from the f32
8782 // rows is exactly the dispatch being mirrored. Keep step35 on the unfused arm.
8783 let lin_q8_only = match &layer.mixer {
8784 Mixer::Linear(la) => {
8785 (t >= 3 || (t == 2 && spec_m2())) && e.uses_q8_1_fast(&la.ssm_out)
8786 }
8787 Mixer::Full(_) if self.sliding_gated_moe_batch_program() => false,
8788 _ => true,
8789 };
8790 // NOTE decode.rs's take()-first lesson: take the pending pair BEFORE branching so
8791 // a non-fused layer still performs the residual add.
8792 let taken = pending.take();
8793 let (h, h_q8) = if norm_fused && lin_q8_only {
8794 let pair = match taken {
8795 // fused add + attn_norm + q8_1: ONE launch resolves the carried residual
8796 // AND emits this layer's mixer input pre-quantized. res -> x2 (= new x).
8797 Some((x1p, f1p)) => {
8798 let mut x2 = vbuf(e, t * n_embd)?; // fully written (res output)
8799 let p = e.add_rms_norm_q8_1(
8800 &x1p,
8801 &f1p,
8802 layer.attn_norm.float_data(),
8803 &mut x2,
8804 n_embd,
8805 t,
8806 eps,
8807 )?;
8808 x = x2;
8809 p
8810 }
8811 None => e.rms_norm_q8_1(&x, layer.attn_norm.float_data(), n_embd, t, eps)?,
8812 };
8813 (e.zeros(0)?, Some(pair)) // h unused on this path (q8-only consumers)
8814 } else {
8815 if let Some((x1p, f1p)) = taken {
8816 let mut x2 = vbuf(e, t * n_embd)?; // fully written by add
8817 e.add(&x1p, &f1p, &mut x2, t * n_embd)?;
8818 x = x2;
8819 }
8820 let mut h = vbuf(e, t * n_embd)?; // fully written by either rms_norm arm
8821 if norm_fused {
8822 e.rms_norm_decode(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
8823 } else {
8824 e.rms_norm(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
8825 }
8826 (h, None)
8827 };
8828 let h_q8_ref = h_q8.as_ref().map(|(q, d)| (q, d));
8829
8830 let mixed = match &layer.mixer {
8831 Mixer::Full(fa) => self.full_attn_verify(
8832 e,
8833 fa,
8834 &h,
8835 h_q8_ref,
8836 pos_d,
8837 t,
8838 cache,
8839 il,
8840 stream.map(|(_, c)| c),
8841 )?,
8842 Mixer::Mla(_) => crate::hybrid::mla_path_unimplemented("speculative verify"),
8843 Mixer::Kda(_) => crate::hybrid::kda_path_unimplemented("speculative verify"),
8844 Mixer::Linear(la) => {
8845 // BATCHED linear verify (2026-07-03, the MTP-profit lever): one T-token pass —
8846 // batched projections (weight read ONCE, hits the m=2-4 weight-resident matvec),
8847 // carried-state conv (ssm_conv1d_tm_state), GDN prep on the prefill kernels, and
8848 // ONE gdn_scan whose internal sequential t-loop is the SAME recurrence as T
8849 // chained T=1 steps (bit-identical). Falls back to the sequential per-column
8850 // chain when T < d_conv-1 (conv ring update needs T >= pad) — or when ANY
8851 // projection is off the q8_1 fast path: matmul_decode_exact would route a Float
8852 // tensor to cuBLAS at m=t (different FP accumulation than eager's per-token
8853 // GEMV), so mixed-dtype layers stay on the eager-identical per-column chain.
8854 // MEMRA_SPEC_M2 (lane/spec-m2): the t==2 batch rides the same arm — the conv
8855 // wrapper handles t<pad with a pure-copy ring rebuild; see spec_m2() header.
8856 if (t >= 3 || (t == 2 && spec_m2()))
8857 && mixer_fast
8858 && e.uses_q8_1_fast(&la.ssm_out)
8859 {
8860 let want = ckpt.is_some();
8861 let (out, stash) =
8862 self.linear_attn_verify_t(e, la, &h, h_q8_ref, t, cache, il, want)?;
8863 if let (Some(ck), Some(st)) = (ckpt.as_deref_mut(), stash) {
8864 ck.gdn[il] = Some(st);
8865 }
8866 out
8867 } else {
8868 let mut out = vbuf(e, t * n_embd)?; // every col written by copy_into
8869 let mut col_states: Option<Vec<(CudaSlice<f32>, CudaSlice<f32>)>> =
8870 if ckpt.is_some() && t >= 2 {
8871 Some(Vec::with_capacity(t - 1))
8872 } else {
8873 None
8874 };
8875 for col in 0..t {
8876 let mut h_col = vbuf(e, n_embd)?; // fully written by copy_view_into
8877 let src = h.slice(col * n_embd..(col + 1) * n_embd);
8878 e.copy_view_into(&mut h_col, 0, &src, n_embd)?;
8879 let m_col = self.linear_attn_decode(e, la, &h_col, cache, il)?;
8880 e.copy_into(&mut out, col * n_embd, &m_col, n_embd)?;
8881 // REPLAY-FREE ckpt: clone the chain's ACTUAL state after this column
8882 // (pure dtod — cannot change any computed value). Last column skipped:
8883 // rebuild targets are j <= t-1 columns.
8884 if let Some(cs) = col_states.as_mut()
8885 && col + 1 < t
8886 {
8887 let rl = cache.recur[il].as_ref().unwrap();
8888 cs.push((
8889 e.clone_dtod(&rl.conv_state)?,
8890 e.clone_dtod(&rl.ssm_state)?,
8891 ));
8892 }
8893 }
8894 if let (Some(ck), Some(cs)) = (ckpt.as_deref_mut(), col_states) {
8895 // ReplaySSM-assessment instrumentation (2026-07-30): the
8896 // per-column clones are the only true state snapshots left in
8897 // the verify (the batched path stashes INPUTS and replays).
8898 if std::env::var("MEMRA_SPEC_STATS").as_deref() == Ok("1") {
8899 static ONCE: std::sync::Once = std::sync::Once::new();
8900 let bytes: usize =
8901 cs.iter().map(|(c, s)| (c.len() + s.len()) * 4).sum();
8902 ONCE.call_once(|| eprintln!(
8903 "[verify-ckpt] per-column layer il={il}: {} clones, {:.2} MB/layer/round",
8904 cs.len(), bytes as f64 / 1e6));
8905 }
8906 ck.cols[il] = Some(cs);
8907 }
8908 out
8909 }
8910 }
8911 };
8912 if spec_nan_scan_level() >= 2 {
8913 let mixed_width = mixed.len() / t;
8914 nan_scan_rows(
8915 e,
8916 &mixed,
8917 t,
8918 mixed_width,
8919 &format!("verify layer {il} batched ATTN out pos0={pos0}"),
8920 )?;
8921 }
8922
8923 // DISPATCH-MIRRORED post-attn norm: eager residual_norm_ffn fuses add+norm+quant
8924 // (1024-thread add_rms_norm_q8_1) only for Dense FFNs whose gate+up are q8_1-fast;
8925 // otherwise (and for MoE) it runs the 256-thread fused add_rms_norm. Mirror per layer.
8926 let ffn_fuse = match &layer.ffn {
8927 crate::hybrid::Ffn::Dense {
8928 ffn_gate, ffn_up, ..
8929 } => {
8930 std::env::var("MEMRA_NO_FUSE_NORMQ").is_err()
8931 && e.uses_q8_1_fast(ffn_gate)
8932 && e.uses_q8_1_fast(ffn_up)
8933 }
8934 crate::hybrid::Ffn::Moe(_) => false,
8935 };
8936 // BATCHED EPILOGUE RE-FUSE (lane/vt-fixes fix 2): on the ffn_fuse path (Dense,
8937 // gate+up q8_1-fast, non-M3) the FFN input is emitted DIRECTLY as q8_1 by ONE
8938 // add_rms_norm_q8_1 launch at nrows=t (row-indexed kernel: the T-row launch is the
8939 // per-row m=1 program; kernel-check pins bit-identity vs the unfused
8940 // add_f32 -> rms_norm_decode -> quantize_q8_1 chain at T=2/4/5/8) — replacing the
8941 // add + rms_norm_decode launches AND the dual/singles' internal re-quantize.
8942 // M3's swigluoai must keep the f32 chain (the fused SwiGLU epilogue encodes plain
8943 // SiLU), mirroring residual_norm_ffn's m3 guard on the decode path.
8944 // step35: same guard per LAYER. A dense FFN's clamp is the SHEXP array (upstream's
8945 // one build_ffn serves dense + shared expert, llama-graph.cpp:1751), and verify MUST
8946 // mirror decode's dispatch or spec self-consistency fails.
8947 let dense_lim = self.cfg.clamp_shexp_at(il as u32);
8948 let fuse_q8 = ffn_fuse && self.cfg.m3.is_none() && dense_lim.is_none();
8949 let mut x1 = vbuf(e, t * n_embd)?; // fully written by add / add_rms_norm*
8950 let mut z = e.zeros(0)?; // replaced below on the unfused arms
8951 let z_q8 = if fuse_q8 {
8952 Some(e.add_rms_norm_q8_1(
8953 &x,
8954 &mixed,
8955 layer.post_attn_norm.float_data(),
8956 &mut x1,
8957 n_embd,
8958 t,
8959 eps,
8960 )?)
8961 } else {
8962 let mut zf = vbuf(e, t * n_embd)?; // fully written by rms_norm_decode / add_rms_norm
8963 if ffn_fuse {
8964 e.add(&x, &mixed, &mut x1, t * n_embd)?;
8965 e.rms_norm_decode(
8966 &x1,
8967 layer.post_attn_norm.float_data(),
8968 &mut zf,
8969 n_embd,
8970 t,
8971 eps,
8972 )?;
8973 } else {
8974 e.add_rms_norm(
8975 &x,
8976 &mixed,
8977 layer.post_attn_norm.float_data(),
8978 &mut x1,
8979 &mut zf,
8980 n_embd,
8981 t,
8982 eps,
8983 )?;
8984 }
8985 z = zf;
8986 None
8987 };
8988 if spec_nan_scan_level() >= 2 && !z.is_empty() {
8989 nan_scan_rows(
8990 e,
8991 &z,
8992 t,
8993 n_embd,
8994 &format!("verify layer {il} post-attn norm z pos0={pos0}"),
8995 )?;
8996 }
8997 // DECODE-EXACT FFN projections: force MMVQ for gate/up/down at any T to match the
8998 // T=1 decode FP accumulation order. At T>=5 the generic matmul/matmul_pre falls to dp4a
8999 // (128-thread, different FP sum order). At T=2-4 the batched MMVQ is already bit-identical.
9000 let ffn_out = match &layer.ffn {
9001 crate::hybrid::Ffn::Dense {
9002 ffn_gate,
9003 ffn_up,
9004 ffn_down,
9005 } => {
9006 let n_ff = ffn_gate.out_features();
9007 if let Some((zq, zd)) = z_q8.as_ref() {
9008 // FUSED CHAIN (fix 2): pre-quantized z feeds the projections; the SwiGLU
9009 // epilogue emits act pre-quantized for ffn_down (silu_mul_scaled_q8_1,
9010 // bit-identical to silu_mul + quantize — kernel-check-pinned) with the
9011 // NVFP4 macro-scales folded (deferred-scale dual: y*s inline == the
9012 // scale_inplace store, value-exact) — the exact m=1 decode epilogue
9013 // structure at nrows=t.
9014 let pair = e
9015 .matmul_decode_exact_dual_pre(ffn_gate, ffn_up, zq, zd, t)?
9016 .map(|((g, gs), (u, us))| (g, gs, u, us));
9017 let (gate, gs, up, us) = match pair {
9018 Some(x4) => x4,
9019 None => (
9020 e.matmul_decode_exact_pre(ffn_gate, zq, zd, t)?,
9021 1.0, // scale already applied inside _pre
9022 e.matmul_decode_exact_pre(ffn_up, zq, zd, t)?,
9023 1.0,
9024 ),
9025 };
9026 if e.uses_q8_1_fast(ffn_down) {
9027 let (aq, ad) = e.silu_mul_scaled_q8_1(&gate, &up, gs, us, t * n_ff)?;
9028 e.matmul_decode_exact_pre(ffn_down, &aq, &ad, t)?
9029 } else {
9030 let mut act = vbuf(e, t * n_ff)?;
9031 e.silu_mul_scaled(&gate, &up, gs, us, &mut act, t * n_ff)?;
9032 e.matmul_decode_exact(ffn_down, &act, t)?
9033 }
9034 } else {
9035 // UNFUSED (pre-fix) chain — MoE-adjacent/M3/off-fast layers, unchanged.
9036 // DUAL gate+up batched twin (lane/verify-economics, 2026-08-02): one launch
9037 // for the pair at t=2..8 — bit-identical per (tensor,token,row) to the two
9038 // singles (kernel-check pins bitwise; MEMRA_SPEC_DUAL_T=0 reverts). None
9039 // (non-NVFP4 / t outside the tier / seam off) -> the two singles, unchanged.
9040 let (gate, up) =
9041 match e.matmul_decode_exact_dual(ffn_gate, ffn_up, &z, t)? {
9042 Some(pair) => pair,
9043 None => (
9044 e.matmul_decode_exact(ffn_gate, &z, t)?,
9045 e.matmul_decode_exact(ffn_up, &z, t)?,
9046 ),
9047 };
9048 let mut act = vbuf(e, t * n_ff)?; // fully written by ffn_act_lim
9049 Self::ffn_act_lim(
9050 e,
9051 &self.cfg,
9052 &gate,
9053 &up,
9054 1.0,
9055 1.0,
9056 dense_lim,
9057 &mut act,
9058 t * n_ff,
9059 )?;
9060 e.matmul_decode_exact(ffn_down, &act, t)?
9061 }
9062 }
9063 crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il(e, m, &z, t, il as u16)?,
9064 };
9065 if spec_nan_scan_level() >= 2 {
9066 nan_scan_rows(
9067 e,
9068 &ffn_out,
9069 t,
9070 n_embd,
9071 &format!("verify layer {il} batched FFN out pos0={pos0}"),
9072 )?;
9073 }
9074 if spec_nan_scan() {
9075 let mut residual = vbuf(e, t * n_embd)?;
9076 e.add(&x1, &ffn_out, &mut residual, t * n_embd)?;
9077 nan_scan_rows(
9078 e,
9079 &residual,
9080 t,
9081 n_embd,
9082 &format!("verify layer {il} residual pos0={pos0}"),
9083 )?;
9084 }
9085 // CROSS-LAYER fusion: defer this layer's post-FFN residual add — the next layer's
9086 // fused-q8 attn norm folds it in (add_rms_norm_q8_1 == add; rms_norm; quantize,
9087 // kernel-check-pinned at nrows=T). Non-fused next layers add explicitly above.
9088 pending = Some((x1, ffn_out));
9089 }
9090 // RANGE's final add (no next norm INSIDE the range to fuse with; for the
9091 // whole-trunk call that is the last layer, whose next norm is output_norm — f32-out).
9092 if let Some((x1p, f1p)) = pending.take() {
9093 let mut x2 = vbuf(e, t * n_embd)?; // fully written by add
9094 e.add(&x1p, &f1p, &mut x2, t * n_embd)?;
9095 x = x2;
9096 }
9097 Ok(x)
9098 }
9099 /// BATCHED linear-attn verify (T=K+1): the whole layer in ~10 launches instead of T x the
9100 /// T=1 decode chain (T x ~12 launches + T weight reads of the four projections). The GDN
9101 /// recurrence itself is inherently sequential — gdn_scan_s128 runs its internal t-loop with
9102 /// the SAME per-token math as chained T=1 calls (bit-identical state evolution); everything
9103 /// around it (projections, conv, prep, gated norm, out-proj) batches. Advances conv ring +
9104 /// ssm state exactly like T sequential decode steps.
9105 /// `want_stash`: additionally RETAIN the gdn-scan inputs (pure buffer keep-alives, zero extra
9106 /// kernels) so a partial accept can rebuild the state after any column prefix (REPLAY-FREE).
9107 #[allow(clippy::too_many_arguments)]
9108 fn linear_attn_verify_t(
9109 &self,
9110 e: &Engine,
9111 la: &LinearAttnLayer,
9112 h: &CudaSlice<f32>,
9113 h_q8: Option<(&CudaSlice<i8>, &CudaSlice<f32>)>,
9114 t: usize,
9115 cache: &mut Cache,
9116 il: usize,
9117 want_stash: bool,
9118 ) -> Result<(CudaSlice<f32>, Option<GdnStash>), Box<dyn std::error::Error>> {
9119 let cfg = &self.cfg;
9120 let geometry = la.geometry;
9121 let d_state = geometry.key_head_dim as usize;
9122 let num_k = geometry.key_heads as usize;
9123 let num_v = geometry.value_heads as usize;
9124 let d_conv = geometry.conv_kernel as usize;
9125 let key_dim = d_state * num_k;
9126 let conv_dim = key_dim * 2 + geometry.value_head_dim as usize * num_v;
9127 let eps = cfg.rms_eps;
9128 let scale = 1.0 / (d_state as f32).sqrt();
9129
9130 // DECODE-EXACT projections: matmul_decode_exact forces the MMVQ (warp-per-row, 32-thread)
9131 // accumulation order for EVERY m, matching the T=1 decode path bit-for-bit. The generic
9132 // `matmul` at m>=5 falls to dp4a (128-thread, two-level reduce) which has a different FP
9133 // sum order — ULP differences propagate through gdn_scan and flip argmax on the 27B.
9134 // Q8 TRUNK-FUSION at T=1 (35B: wqkv+wqkv_gate both Q8_0): one fused2 launch, bit-identical
9135 // per (tensor,row) to the two m=1 MMVQ dispatches below — decode-exact contract holds.
9136 // VERIFY-TIER TRUNK FUSION (MEMRA_SPEC_FUSED_T, t=2-4): quantize h ONCE for every
9137 // fused-eligible same-input Q8_0 pair of this layer (35B wqkv+wqkv_gate; 9B
9138 // ssm_beta+ssm_alpha) — each fused2 batched launch then replaces two decode-exact
9139 // calls (each of which re-quantizes the same h + runs its own _b2/_b4 launch).
9140 // Bit-identical per (tensor,token,row) — see spec_fused_t().
9141 // BATCHED EPILOGUE RE-FUSE (lane/vt-fixes fix 2): `h_q8` = the attn-input norm emitted
9142 // directly as q8_1 by the caller's fused rms_norm_q8_1 (bit-identical to the unfused
9143 // chain, kernel-check-pinned). When present it REPLACES the standalone quantize below
9144 // and feeds every projection; the caller guaranteed all four input projections are
9145 // q8_1-fast. When absent, the old shared-quantize (fused-t window) stands.
9146 let h_q8_t = if h_q8.is_none()
9147 && spec_fused_t()
9148 && (2..=4).contains(&t)
9149 && ((e.uses_q8_1_fast(&la.wqkv) && e.uses_q8_1_fast(&la.wqkv_gate))
9150 || (e.uses_q8_1_fast(&la.ssm_beta) && e.uses_q8_1_fast(&la.ssm_alpha)))
9151 {
9152 Some(e.quantize_q8_1(h, t, cfg.n_embd as usize)?)
9153 } else {
9154 None
9155 };
9156 // one view: the caller's fused-norm q8 or this fn's own shared quantize.
9157 let hq8_any: Option<(&CudaSlice<i8>, &CudaSlice<f32>)> =
9158 h_q8.or(h_q8_t.as_ref().map(|(q, d)| (q, d)));
9159 let (qkv_mixed, z) = {
9160 let mut fused = None;
9161 if t == 1 && e.uses_q8_1_fast(&la.wqkv) && e.uses_q8_1_fast(&la.wqkv_gate) {
9162 let (hq, hd) = e.quantize_q8_1(h, 1, cfg.n_embd as usize)?;
9163 fused = e.matmul_q8_fused2(&la.wqkv, &la.wqkv_gate, &hq, &hd)?;
9164 } else if let Some((hq, hd)) = hq8_any
9165 && spec_fused_t()
9166 && (2..=4).contains(&t)
9167 {
9168 fused = e.matmul_q8_fused2_t(&la.wqkv, &la.wqkv_gate, hq, hd, t)?;
9169 }
9170 match (fused, hq8_any) {
9171 (Some(pair), _) => pair,
9172 (None, Some((hq, hd))) if h_q8.is_some() => (
9173 e.matmul_decode_exact_pre(&la.wqkv, hq, hd, t)?,
9174 e.matmul_decode_exact_pre(&la.wqkv_gate, hq, hd, t)?,
9175 ),
9176 (None, _) => (
9177 e.matmul_decode_exact(&la.wqkv, h, t)?,
9178 e.matmul_decode_exact(&la.wqkv_gate, h, t)?,
9179 ),
9180 }
9181 };
9182 // beta+alpha DUAL at T=1 (75% of p3 rounds run T=1 verify — p-min chain cuts): the dual
9183 // mr2 kernel is bit-identical per element to the m=1 MMVQ matmul_decode_exact dispatches
9184 // (same warp-per-row body, blockIdx.y picks the weight), so the decode-exact contract
9185 // holds; the run-spec battery is the arbiter. T>1 keeps the per-tensor decode-exact path.
9186 let (beta_raw, alpha) = if t == 1 {
9187 let (hq, hd) = e.quantize_q8_1(h, 1, cfg.n_embd as usize)?;
9188 match e.matmul_pre_dual_noscale(&la.ssm_beta, &la.ssm_alpha, &hq, &hd, 1)? {
9189 Some(((mut b, bs), (mut a, as_))) => {
9190 if bs != 1.0 {
9191 e.scale_inplace(&mut b, bs, la.ssm_beta.out_features())?;
9192 }
9193 if as_ != 1.0 {
9194 e.scale_inplace(&mut a, as_, la.ssm_alpha.out_features())?;
9195 }
9196 (b, a)
9197 }
9198 // Q8_0 fused2 twin (9B stores beta/alpha as Q8_0): DISPATCH-MIRRORS the eager
9199 // decode's beta_alpha closure — the fused body is qmatvec_q8_0_mmvq verbatim,
9200 // bit-identical per row (kernel-check rel=0.00e0 gate), so decode==verify holds.
9201 None => match e.matmul_q8_fused2(&la.ssm_beta, &la.ssm_alpha, &hq, &hd)? {
9202 Some((b, a)) => (b, a),
9203 None => (
9204 e.matmul_decode_exact(&la.ssm_beta, h, 1)?,
9205 e.matmul_decode_exact(&la.ssm_alpha, h, 1)?,
9206 ),
9207 },
9208 }
9209 } else {
9210 // fused-t twin (9B stores beta/alpha as Q8_0): same shared-quantize + one launch
9211 // contract as the wqkv pair above; 35B beta/alpha are Float -> None -> fallback.
9212 let mut nvfp4_fused = None;
9213 let mut q8_fused = None;
9214 if let Some((hq, hd)) = hq8_any {
9215 if t == 3 && std::env::var("MEMRA_NVFP4_AUX_DUAL").as_deref() != Ok("0") {
9216 nvfp4_fused =
9217 e.matmul_decode_exact_dual_pre(&la.ssm_beta, &la.ssm_alpha, hq, hd, t)?;
9218 if nvfp4_fused.is_some() && std::env::var("MEMRA_DEBUG").is_ok() {
9219 static ONCE: std::sync::Once = std::sync::Once::new();
9220 ONCE.call_once(|| {
9221 eprintln!("[memra] NVFP4 beta+alpha batched aux dual ENGAGED (t={t})")
9222 });
9223 }
9224 }
9225 if nvfp4_fused.is_none() && spec_fused_t() && (2..=4).contains(&t) {
9226 q8_fused = e.matmul_q8_fused2_t(&la.ssm_beta, &la.ssm_alpha, hq, hd, t)?;
9227 }
9228 }
9229 if let Some(((mut b, bs), (mut a, as_))) = nvfp4_fused {
9230 if bs != 1.0 {
9231 e.scale_inplace(&mut b, bs, t * la.ssm_beta.out_features())?;
9232 }
9233 if as_ != 1.0 {
9234 e.scale_inplace(&mut a, as_, t * la.ssm_alpha.out_features())?;
9235 }
9236 (b, a)
9237 } else if let Some(pair) = q8_fused {
9238 pair
9239 } else {
9240 match hq8_any {
9241 Some((hq, hd)) if h_q8.is_some() => (
9242 e.matmul_decode_exact_pre(&la.ssm_beta, hq, hd, t)?,
9243 e.matmul_decode_exact_pre(&la.ssm_alpha, hq, hd, t)?,
9244 ),
9245 _ => (
9246 e.matmul_decode_exact(&la.ssm_beta, h, t)?,
9247 e.matmul_decode_exact(&la.ssm_alpha, h, t)?,
9248 ),
9249 }
9250 }
9251 };
9252
9253 // conv with CARRIED state + ring roll (T >= pad rides the input-column update kernel;
9254 // T < pad — the MEMRA_SPEC_M2 t=2 arm — rolls via the pure-copy ring rebuild).
9255 let rl = cache.recur[il].as_mut().unwrap();
9256 let mut conv_out = e.uninit(conv_dim * t)?;
9257 e.ssm_conv1d_tm_state(
9258 &qkv_mixed,
9259 &mut rl.conv_state,
9260 la.ssm_conv1d.float_data(),
9261 &mut conv_out,
9262 conv_dim,
9263 t,
9264 d_conv,
9265 )?;
9266
9267 // GDN prep via the prefill kernels (repack + L2 + sigmoid + glog), T-wide.
9268 let mut q_g = e.uninit(d_state * num_v * t)?;
9269 let mut k_g = e.uninit(d_state * num_v * t)?;
9270 let mut v_g = e.uninit(d_state * num_v * t)?;
9271 e.qkv_to_gdn_repack(
9272 &conv_out, &mut q_g, &mut k_g, &mut v_g, d_state, num_v, num_k, key_dim, t,
9273 )?;
9274 let mut q_l2 = e.uninit(d_state * num_v * t)?;
9275 e.l2_norm_decode(&q_g, &mut q_l2, d_state, num_v * t, eps)?;
9276 let mut k_l2 = e.uninit(d_state * num_v * t)?;
9277 e.l2_norm_decode(&k_g, &mut k_l2, d_state, num_v * t, eps)?;
9278 let mut beta = e.uninit(t * num_v)?;
9279 e.sigmoid(&beta_raw, &mut beta, t * num_v)?;
9280 let mut g_log = e.uninit(t * num_v)?;
9281 e.gdn_glog(
9282 &alpha,
9283 la.ssm_dt.float_data(),
9284 la.ssm_a.float_data(),
9285 &mut g_log,
9286 num_v,
9287 t,
9288 )?;
9289
9290 // ONE gdn_scan over T tokens from the carried state (internal sequential loop ==
9291 // T chained T=1 steps). Ping-pong the resident buffers like eager decode.
9292 let mut o = e.uninit(d_state * num_v * t)?;
9293 {
9294 let crate::cache::RecurLayer {
9295 ssm_state,
9296 ssm_state_alt,
9297 ..
9298 } = rl;
9299 e.gdn_scan_s128(
9300 &q_l2,
9301 &k_l2,
9302 &v_g,
9303 &g_log,
9304 &beta,
9305 ssm_state,
9306 ssm_state_alt,
9307 &mut o,
9308 num_v,
9309 t,
9310 scale,
9311 )?;
9312 }
9313 std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
9314
9315 // gated RMSNorm + out projection, T-wide. FUSED-QUANTIZE ARM (lane/vt-fixes fix 2,
9316 // mirroring the T=1 decode's launch-arc form): when ssm_out rides the q8_1 fast path,
9317 // emit q8_1 straight from the gated norm at nrows=num_v*t (row-indexed kernel, the
9318 // T-wide launch is the per-row program; kernel-check pins bit-identity vs
9319 // gated_rmsnorm -> quantize_q8_1 at T=1 and T=5) and feed the decode-exact dispatch
9320 // pre-quantized — one launch replaces norm + quantize. Fallback = the f32 chain.
9321 let out = if e.uses_q8_1_fast(&la.ssm_out) {
9322 let (gq, gd) =
9323 e.gated_rmsnorm_q8_1(&o, la.ssm_norm.float_data(), &z, d_state, num_v * t, eps)?;
9324 e.matmul_decode_exact_pre(&la.ssm_out, &gq, &gd, t)?
9325 } else {
9326 let mut gn = e.uninit(d_state * num_v * t)?;
9327 e.gated_rmsnorm(
9328 &o,
9329 la.ssm_norm.float_data(),
9330 &z,
9331 &mut gn,
9332 d_state,
9333 num_v * t,
9334 eps,
9335 )?;
9336 // DECODE-EXACT out-projection: same MMVQ path as the T=1 decode (ssm_out at m>=5
9337 // would fall to dp4a with a different FP reduction order — same class of bug as
9338 // the input projs).
9339 e.matmul_decode_exact(&la.ssm_out, &gn, t)?
9340 };
9341 let stash = if want_stash {
9342 Some(GdnStash {
9343 qkv_mixed,
9344 q_l2,
9345 k_l2,
9346 v_g,
9347 g_log,
9348 beta,
9349 })
9350 } else {
9351 None
9352 };
9353 Ok((out, stash))
9354 }
9355
9356 /// REPLAY-FREE partial-accept commit (2026-07-03): make the cache state == "committed through
9357 /// the first `j` verify columns" WITHOUT the legacy rollback + duplicate trunk replay.
9358 /// - Full-attn KV: truncate both the owning-stage shadow and every TP rank to snapshot + j.
9359 /// The verify's appended rows for those columns are bit-identical to what an eager T=1
9360 /// chain writes (the decode-exact contract the verify-probe gates), so keeping them ==
9361 /// replaying them.
9362 /// - Linear layers, batched path: rebuild the conv ring by PURE COPIES (ring holds raw input
9363 /// columns) and the ssm state by a prefix re-run of the SAME gdn_scan kernel (t=j) from the
9364 /// snapshot state over the stash's identical inputs — the kernel's t-loop carries state in
9365 /// registers and writes it once at the end, so iterations 0..j-1 are independent of T:
9366 /// bit-identical to the verify's own state after j tokens == the eager chain state.
9367 /// - Linear layers, per-column path: restore the cloned actual state after column j-1.
9368 /// Caller guarantees 1 <= j <= t-1 (j==0 rounds take the legacy rollback; j==t is full accept).
9369 #[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
9370 fn commit_verified_prefix(
9371 &self,
9372 e: &Engine,
9373 cache: &mut Cache,
9374 snap: &crate::cache::CacheSnapshot,
9375 ckpt: &VerifyCkpt,
9376 j: usize,
9377 kv_lens_done: bool,
9378 dev_j: Option<(&CudaSlice<u32>, usize, usize)>,
9379 ) -> Result<(), Box<dyn std::error::Error>> {
9380 // GDN geometry derives lazily inside recurrent-layer arms. Full-attention plans carry no
9381 // recurrent state and must never be forced through a synthetic SSM geometry.
9382 // Engine-bundle slice 1 (DSF-ROUNDCOST-20260820 §1.1): the per-column-arm restores
9383 // are 2 tiny D2D copies per linear layer (~96 dispatches/partial round on the q38
9384 // route). When every cols-arm layer shares uniform state sizes (single ssm cfg —
9385 // always true today), batch them into two `copy_batch_uniform_f32` launches. Bytes,
9386 // buffers and stream order are identical to the per-layer memcpy sequence; the
9387 // kernel-rebuild (gdn-stash) arm below is untouched. MEMRA_STATE_COPY_BATCH=0 reverts.
9388 let mut batched_cols = false;
9389 if state_copy_batch_on() && dev_j.is_none() {
9390 use cudarc::driver::DevicePtr;
9391 let s = &e.gpu.stream();
9392 let mut conv_pairs: Vec<(u64, u64)> = Vec::new();
9393 let mut ssm_pairs: Vec<(u64, u64)> = Vec::new();
9394 let (mut conv_words, mut ssm_words) = (0usize, 0usize);
9395 let mut uniform = true;
9396 for il in 0..self.layers.len() {
9397 let Some(rl) = cache.recur[il].as_ref() else {
9398 continue;
9399 };
9400 if ckpt.gdn[il].is_some() {
9401 continue; // kernel-rebuild arm restores below, per layer
9402 }
9403 let Some(cols) = &ckpt.cols[il] else {
9404 continue; // missing-ckpt error surfaces in the main loop
9405 };
9406 let (c, st) = &cols[j - 1];
9407 if conv_pairs.is_empty() {
9408 conv_words = c.len();
9409 ssm_words = st.len();
9410 } else if c.len() != conv_words || st.len() != ssm_words {
9411 uniform = false;
9412 break;
9413 }
9414 let (pc, _g0) = c.device_ptr(s);
9415 let (dc, _g1) = rl.conv_state.device_ptr(s);
9416 let (ps, _g2) = st.device_ptr(s);
9417 let (ds, _g3) = rl.ssm_state.device_ptr(s);
9418 conv_pairs.push((pc, dc));
9419 ssm_pairs.push((ps, ds));
9420 }
9421 if uniform && !conv_pairs.is_empty() {
9422 let n = conv_pairs.len();
9423 let mut t = vec![0u64; 2 * n];
9424 for (k, &(src, dst)) in conv_pairs.iter().enumerate() {
9425 t[k] = src;
9426 t[n + k] = dst;
9427 }
9428 let conv_t = e.htod_u64(&t)?;
9429 for (k, &(src, dst)) in ssm_pairs.iter().enumerate() {
9430 t[k] = src;
9431 t[n + k] = dst;
9432 }
9433 let ssm_t = e.htod_u64(&t)?;
9434 e.copy_batch_uniform_f32(&conv_t, n, conv_words)?;
9435 e.copy_batch_uniform_f32(&ssm_t, n, ssm_words)?;
9436 batched_cols = true;
9437 }
9438 }
9439 for il in 0..self.layers.len() {
9440 if let (Some(kvl), Some(saved)) = (cache.kv[il].as_mut(), snap.kv_len[il]) {
9441 kvl.len = saved + j;
9442 // devacc 3a: spec_rollback_kv already wrote len_d on-device (same value).
9443 if !kv_lens_done {
9444 e.set_i32_one(&mut kvl.len_d, kvl.len as i32)?;
9445 }
9446 }
9447 if let Some(rl) = cache.recur[il].as_mut() {
9448 let Mixer::Linear(linear) = &self.layers[il].mixer else {
9449 return Err(format!("recurrent cache layer {il} has no GDN plan").into());
9450 };
9451 let geometry = linear.geometry;
9452 let d_state = geometry.key_head_dim as usize;
9453 let num_k = geometry.key_heads as usize;
9454 let num_v = geometry.value_heads as usize;
9455 let d_conv = geometry.conv_kernel as usize;
9456 let conv_dim = d_state * num_k * 2 + geometry.value_head_dim as usize * num_v;
9457 let scale = 1.0 / (d_state as f32).sqrt();
9458 if let Some(st) = &ckpt.gdn[il] {
9459 let ring_old = snap.conv[il].as_ref().expect("snapshot missing conv");
9460 let state_in = snap.ssm[il].as_ref().expect("snapshot missing ssm");
9461 if let Some((acc, base, t_v)) = dev_j {
9462 // 3b: j read on-device (_dc twins, same bodies; full accept early-exits).
9463 e.ssm_conv_ring_rebuild_dc(
9464 &st.qkv_mixed,
9465 ring_old,
9466 &mut rl.conv_state,
9467 conv_dim,
9468 acc,
9469 base,
9470 t_v,
9471 d_conv,
9472 )?;
9473 let mut o = e.uninit(d_state * num_v * j.max(1))?;
9474 e.gdn_scan_s128_dc(
9475 &st.q_l2,
9476 &st.k_l2,
9477 &st.v_g,
9478 &st.g_log,
9479 &st.beta,
9480 state_in,
9481 &mut rl.ssm_state,
9482 &mut o,
9483 num_v,
9484 acc,
9485 base,
9486 t_v,
9487 scale,
9488 )?;
9489 } else {
9490 e.ssm_conv_ring_rebuild(
9491 &st.qkv_mixed,
9492 ring_old,
9493 &mut rl.conv_state,
9494 conv_dim,
9495 j,
9496 d_conv,
9497 )?;
9498 let mut o = e.uninit(d_state * num_v * j)?; // scan output, discarded
9499 e.gdn_scan_s128(
9500 &st.q_l2,
9501 &st.k_l2,
9502 &st.v_g,
9503 &st.g_log,
9504 &st.beta,
9505 state_in,
9506 &mut rl.ssm_state,
9507 &mut o,
9508 num_v,
9509 j,
9510 scale,
9511 )?;
9512 }
9513 } else if let Some(cols) = &ckpt.cols[il] {
9514 if !batched_cols {
9515 let (c, s) = &cols[j - 1];
9516 e.copy_into(&mut rl.conv_state, 0, c, c.len())?;
9517 e.copy_into(&mut rl.ssm_state, 0, s, s.len())?;
9518 }
9519 } else {
9520 return Err(
9521 "commit_verified_prefix: verify ckpt missing for linear layer".into(),
9522 );
9523 }
9524 }
9525 }
9526 self.restore_step_tp_kv_verified_prefix(e, cache, snap, j)?;
9527 cache.pos = snap.pos + j;
9528 Ok(())
9529 }
9530
9531 /// ROUND-STREAM: recur restore with device-j (the _dc twins; full accept early-exits
9532 /// in-kernel). Requires the batched-linear stash on every linear layer (stream gate).
9533 #[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
9534 fn commit_verified_prefix_stream(
9535 &self,
9536 e: &Engine,
9537 cache: &mut Cache,
9538 snap: &crate::cache::CacheSnapshot,
9539 ckpt: &VerifyCkpt,
9540 acc: &CudaSlice<u32>,
9541 base: usize,
9542 t_v: usize,
9543 ) -> Result<(), Box<dyn std::error::Error>> {
9544 for il in 0..self.layers.len() {
9545 if let Some(rl) = cache.recur[il].as_mut() {
9546 let Mixer::Linear(linear) = &self.layers[il].mixer else {
9547 return Err(format!("recurrent cache layer {il} has no GDN plan").into());
9548 };
9549 let geometry = linear.geometry;
9550 let d_state = geometry.key_head_dim as usize;
9551 let num_k = geometry.key_heads as usize;
9552 let num_v = geometry.value_heads as usize;
9553 let d_conv = geometry.conv_kernel as usize;
9554 let conv_dim = d_state * num_k * 2 + geometry.value_head_dim as usize * num_v;
9555 let scale = 1.0 / (d_state as f32).sqrt();
9556 let st = ckpt.gdn[il]
9557 .as_ref()
9558 .ok_or("stream restore: batched-linear stash missing")?;
9559 let ring_old = snap.conv[il].as_ref().expect("snapshot missing conv");
9560 let state_in = snap.ssm[il].as_ref().expect("snapshot missing ssm");
9561 e.ssm_conv_ring_rebuild_dc(
9562 &st.qkv_mixed,
9563 ring_old,
9564 &mut rl.conv_state,
9565 conv_dim,
9566 acc,
9567 base,
9568 t_v,
9569 d_conv,
9570 )?;
9571 let mut o = e.uninit(d_state * num_v * t_v)?;
9572 e.gdn_scan_s128_dc(
9573 &st.q_l2,
9574 &st.k_l2,
9575 &st.v_g,
9576 &st.g_log,
9577 &st.beta,
9578 state_in,
9579 &mut rl.ssm_state,
9580 &mut o,
9581 num_v,
9582 acc,
9583 base,
9584 t_v,
9585 scale,
9586 )?;
9587 }
9588 }
9589 Ok(())
9590 }
9591
9592 /// EAGLE3 aux-capturing verify forward over `tokens` (T) — mirrors `decode_step_t_h` exactly
9593 /// (same KV append, same causal verify, same recur advance) but ALSO clones the aux residual-
9594 /// stream hiddens (blocks in `aux_layers`) for TWO columns: the LAST column (always) and the
9595 /// optional `pred_col` (the EAGLE seed = bonus's predecessor). Returns
9596 /// (all_T_logits host, last_col_aux, pred_col_aux?). Used by the EAGLE3 orchestrator's commit.
9597 #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
9598 pub fn decode_step_t_aux2(
9599 &self,
9600 e: &Engine,
9601 tokens: &[u32],
9602 pos0: usize,
9603 cache: &mut Cache,
9604 aux_layers: &[usize],
9605 pred_col: Option<usize>,
9606 ) -> Result<
9607 (Vec<f32>, Vec<CudaSlice<f32>>, Option<Vec<CudaSlice<f32>>>),
9608 Box<dyn std::error::Error>,
9609 > {
9610 cache.ensure_usable("decode_step_t_aux2")?;
9611 let cfg = &self.cfg;
9612 let n_embd = cfg.n_embd as usize;
9613 let eps = cfg.rms_eps;
9614 let t = tokens.len();
9615 let pos_vec: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
9616 let pos_d = e.htod_i32(&pos_vec)?;
9617 let mut x = e.htod(&self.embd.gather(n_embd, tokens))?;
9618 let mut aux_last: Vec<CudaSlice<f32>> = Vec::with_capacity(aux_layers.len());
9619 let mut aux_pred: Vec<CudaSlice<f32>> = Vec::new();
9620 let want_pred = pred_col.is_some();
9621
9622 for (il, layer) in self.layers.iter().enumerate() {
9623 // DISPATCH-MIRRORED norms (FP-order lesson #8) — see decode_step_t_h_emb.
9624 let mixer_fast = self.mixer_in_q8_1_fast(e, &layer.mixer);
9625 let norm_fused = std::env::var("MEMRA_NO_FUSE_NORMQ").is_err() && mixer_fast;
9626 let mut h = vbuf(e, t * n_embd)?; // fully written by either rms_norm arm
9627 if norm_fused {
9628 e.rms_norm_decode(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
9629 } else {
9630 e.rms_norm(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
9631 }
9632 let mixed = match &layer.mixer {
9633 Mixer::Full(fa) => {
9634 self.full_attn_verify(e, fa, &h, None, &pos_d, t, cache, il, None)?
9635 }
9636 Mixer::Mla(_) => {
9637 crate::hybrid::mla_path_unimplemented("auxiliary T-parallel decode")
9638 }
9639 Mixer::Kda(_) => crate::hybrid::kda_path_unimplemented("aux decode step"),
9640 Mixer::Linear(la) => {
9641 let mut out = e.zeros(t * n_embd)?;
9642 for col in 0..t {
9643 let mut h_col = e.zeros(n_embd)?;
9644 let src = h.slice(col * n_embd..(col + 1) * n_embd);
9645 e.copy_view_into(&mut h_col, 0, &src, n_embd)?;
9646 let m_col = self.linear_attn_decode(e, la, &h_col, cache, il)?;
9647 e.copy_into(&mut out, col * n_embd, &m_col, n_embd)?;
9648 }
9649 out
9650 }
9651 };
9652 let ffn_fuse = match &layer.ffn {
9653 crate::hybrid::Ffn::Dense {
9654 ffn_gate, ffn_up, ..
9655 } => {
9656 std::env::var("MEMRA_NO_FUSE_NORMQ").is_err()
9657 && e.uses_q8_1_fast(ffn_gate)
9658 && e.uses_q8_1_fast(ffn_up)
9659 }
9660 crate::hybrid::Ffn::Moe(_) => false,
9661 };
9662 let mut x1 = vbuf(e, t * n_embd)?; // fully written by add / add_rms_norm
9663 let mut z = vbuf(e, t * n_embd)?; // fully written by rms_norm_decode / add_rms_norm
9664 if ffn_fuse {
9665 e.add(&x, &mixed, &mut x1, t * n_embd)?;
9666 e.rms_norm_decode(
9667 &x1,
9668 layer.post_attn_norm.float_data(),
9669 &mut z,
9670 n_embd,
9671 t,
9672 eps,
9673 )?;
9674 } else {
9675 e.add_rms_norm(
9676 &x,
9677 &mixed,
9678 layer.post_attn_norm.float_data(),
9679 &mut x1,
9680 &mut z,
9681 n_embd,
9682 t,
9683 eps,
9684 )?;
9685 }
9686 let ffn_out = match &layer.ffn {
9687 crate::hybrid::Ffn::Dense {
9688 ffn_gate,
9689 ffn_up,
9690 ffn_down,
9691 } => {
9692 let n_ff = ffn_gate.out_features();
9693 let gate = e.matmul_decode_exact(ffn_gate, &z, t)?;
9694 let up = e.matmul_decode_exact(ffn_up, &z, t)?;
9695 let mut act = vbuf(e, t * n_ff)?; // fully written by ffn_act_lim
9696 // dense FFN clamp = the SHEXP array (upstream build_ffn serves both).
9697 Self::ffn_act_lim(
9698 e,
9699 &self.cfg,
9700 &gate,
9701 &up,
9702 1.0,
9703 1.0,
9704 self.cfg.clamp_shexp_at(il as u32),
9705 &mut act,
9706 t * n_ff,
9707 )?;
9708 e.matmul_decode_exact(ffn_down, &act, t)?
9709 }
9710 crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il(e, m, &z, t, il as u16)?,
9711 };
9712 let mut x2 = vbuf(e, t * n_embd)?; // fully written by add
9713 e.add(&x1, &ffn_out, &mut x2, t * n_embd)?;
9714 if aux_layers.contains(&il) {
9715 let mut a = e.zeros(n_embd)?;
9716 e.copy_view_into(&mut a, 0, &x2.slice((t - 1) * n_embd..t * n_embd), n_embd)?;
9717 aux_last.push(a);
9718 if let Some(pc) = pred_col {
9719 let mut ap = e.zeros(n_embd)?;
9720 e.copy_view_into(
9721 &mut ap,
9722 0,
9723 &x2.slice(pc * n_embd..(pc + 1) * n_embd),
9724 n_embd,
9725 )?;
9726 aux_pred.push(ap);
9727 }
9728 }
9729 x = x2;
9730 }
9731 let mut hn = vbuf(e, t * n_embd)?; // fully written by rms_norm_decode
9732 e.rms_norm_decode(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
9733 let logits = e.matmul_decode_exact(&self.output, &hn, t)?;
9734 let host = e.dtoh(&logits)?;
9735 cache.pos += t;
9736 Ok((
9737 host,
9738 aux_last,
9739 if want_pred { Some(aux_pred) } else { None },
9740 ))
9741 }
9742
9743 /// step35 SPEC-VERIFY attention over T query tokens — a per-row REPLAY of the eager
9744 /// `step35_decode_attn`.
9745 ///
9746 /// WHY A REPLAY AND NOT A BATCHED TWIN. The verify's whole job is to be bit-identical to what
9747 /// the eager decode would have computed for the same tokens; that is what makes greedy spec
9748 /// decode exact (run-spec asserts token identity for K=1..8). Every other verify arm in this
9749 /// file earns that identity by carefully mirroring dispatch (`matmul_decode_exact` to force
9750 /// MMVQ at any m, per-layer `ffn_fuse` mirroring, per-row `fa_decode` key bounds). step35
9751 /// stacks FOUR more per-layer degrees of freedom on top of that — per-layer `n_head`
9752 /// (64 full / 96 SWA), per-layer rotary width (64 full / 128 SWA), per-layer rope base, and a
9753 /// SEPARATE `attn_gate` tensor whose projection shares the attn-normed input — and its SWA
9754 /// layers attend through a token-OFFSET view whose offset is a function of the ABSOLUTE
9755 /// position of each query row. A batched twin would have to reproduce all of that AND the
9756 /// per-row offset in one launch; the offset alone rules out the existing rows kernels (they
9757 /// take one `base_len`, not a per-row offset).
9758 ///
9759 /// So this arm calls the eager path itself, once per row, on the same cache. Identity is then
9760 /// true BY CONSTRUCTION rather than by mirroring: row r runs exactly the kernel sequence that
9761 /// eager decode step r runs (same projections, same q8_1 fusion decision, same append, same
9762 /// view arithmetic, same `fa_decode_kvmod`, same gate), because it IS that code. Cost: T x the
9763 /// eager decode mixer instead of one batched pass — the same trade the generic arm's `else`
9764 /// per-row loop already accepts when `fa_rows_eligible` says no. Correctness first; a batched
9765 /// step35 twin is a perf lane's job and must be gated against this arm.
9766 ///
9767 /// The `h_q8` pre-quantized pair from the caller's fused norm is NOT forwarded: it is a
9768 /// T-row buffer and `step35_decode_attn`'s `pre_q` contract is one row. Instead each row's
9769 /// f32 `h` slice is handed over and the callee re-derives its own q8_1 exactly as eager decode
9770 /// does (`quantize_q8_1(h, 1, n_embd)`) — which is the dispatch being mirrored. Callers that
9771 /// took the fused arm therefore MUST still pass a live `h`; `step35_verify` asserts that.
9772 #[allow(clippy::too_many_arguments)]
9773 fn step35_verify(
9774 &self,
9775 e: &Engine,
9776 fa: &FullAttnLayer,
9777 h: &CudaSlice<f32>,
9778 h_q8: Option<(&CudaSlice<i8>, &CudaSlice<f32>)>,
9779 t: usize,
9780 cache: &mut Cache,
9781 il: usize,
9782 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9783 let n_embd = self.cfg.n_embd as usize;
9784 // The fused (h-less) attn-norm arm hands `h` as a zero-length placeholder. This arm needs
9785 // the f32 rows, so the caller must not take that lever for step35 — enforced at the call
9786 // site by the sliding-gated-MoE `Mixer::Full(_) => false` arm of
9787 // `lin_q8_only` in `decode_step_t_core_stream`, and asserted here so a future caller
9788 // cannot regress it into silently reading an empty buffer.
9789 assert_eq!(
9790 h.len(),
9791 t * n_embd,
9792 "step35_verify needs the f32 attn-normed rows ([t*n_embd]); the caller took the \
9793 fused q8-only norm arm (h_q8={}) — step35 must stay on the unfused arm",
9794 h_q8.is_some()
9795 );
9796 // ROW WIDTH IS n_embd, NOT n_head*head_dim: `step35_decode_attn` returns the mixer output
9797 // AFTER `wo`, so a row is [n_embd] — the same contract the generic arm's
9798 // `matmul_decode_exact(&fa.wo, &attn_g, t)` return has. Sizing this buffer from the
9799 // per-layer head geometry (8192 on full-attn, 12288 on SWA) instead overran the row on the
9800 // FIRST copy and panicked inside `copy_into`'s `CudaView::slice` unwrap
9801 // (raw/mtp-bt-20260806T212127Z.log frames 12-13).
9802 let mut out = vbuf(e, t * n_embd)?; // each row fully written by the copy below
9803 for r in 0..t {
9804 // Absolute position of this query row. `cache.pos` is the committed length at round
9805 // start and every row before r has already been appended by this loop, so the r-th
9806 // verify token sits at cache.pos + r — the same position eager decode would give it.
9807 let pos_d = e.htod_i32(&[(cache.pos + r) as i32])?;
9808 let mut h_row = vbuf(e, n_embd)?; // fully written by copy_view_into
9809 e.copy_view_into(
9810 &mut h_row,
9811 0,
9812 &h.slice(r * n_embd..(r + 1) * n_embd),
9813 n_embd,
9814 )?;
9815 // THE eager decode mixer: appends this row's K/V at kvl.len, advances it, then
9816 // attends over the (SWA-offset) view. Post-`wo`, same contract as this fn returns.
9817 let o = self.step35_decode_attn(e, fa, il, &h_row, None, &pos_d, cache)?;
9818 debug_assert_eq!(
9819 o.len(),
9820 n_embd,
9821 "step35_decode_attn returns post-wo [n_embd]"
9822 );
9823 e.copy_into(&mut out, r * n_embd, &o, n_embd)?;
9824 }
9825 Ok(out)
9826 }
9827
9828 /// Full-attention mixer over T query tokens with a GROWING resident KV (verify path, §D.3).
9829 /// Appends the T new K/V columns to cache.kv[il] then attends causally over [0..len) via
9830 /// fa_prefill. Token-major [T, kv_dim] projection layout == cache row layout (single copy).
9831 #[allow(clippy::too_many_arguments)]
9832 fn full_attn_verify(
9833 &self,
9834 e: &Engine,
9835 fa: &FullAttnLayer,
9836 h: &CudaSlice<f32>,
9837 h_q8: Option<(&CudaSlice<i8>, &CudaSlice<f32>)>,
9838 pos_d: &CudaSlice<i32>,
9839 t: usize,
9840 cache: &mut Cache,
9841 il: usize,
9842 stream_ctr: Option<&CudaSlice<i32>>,
9843 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9844 // step35: the generic geometry below is wrong for this arch (per-layer n_head, partial
9845 // per-layer rope, the SWA offset view, and a SEPARATE head-wise gate tensor), so it takes
9846 // its own arm. A verify that silently computes different attention than decode defeats the
9847 // whole self-consistency gate, so the arm is a per-row REPLAY of `step35_decode_attn`
9848 // rather than a batched twin — see `step35_verify` for why that is the exactness-correct
9849 // shape and not laziness.
9850 if self.sliding_gated_moe_batch_program() {
9851 if stream_ctr.is_some() {
9852 return Err(
9853 "step35 has no ROUND-STREAM verify arm (the device-counter _dc twins \
9854 cannot express the SWA offset KV view; same root cause as the dc \
9855 decode refusal) — run spec without the stream arm"
9856 .into(),
9857 );
9858 }
9859 return self.step35_verify(e, fa, h, h_q8, t, cache, il);
9860 }
9861 let cfg = &self.cfg;
9862 let geometry = cfg.full_attention_geometry_at(il as u32);
9863 let n_head = geometry.n_head as usize;
9864 let n_head_kv = geometry.n_head_kv as usize;
9865 let head_dim = geometry.head_dim_k as usize;
9866 let eps = cfg.rms_eps;
9867 let scale = geometry.attention_scale();
9868 let n_embd = cfg.n_embd as usize;
9869
9870 // DECODE-EXACT Q/K/V projections: matmul_decode_exact forces the MMVQ (warp-per-row) path
9871 // for every m, matching the T=1 decode's FP accumulation order. matmul_pre at m>=5 would
9872 // fall to dp4a (128-thread, two-level reduce) with a different FP sum order.
9873 // Q8 TRUNK-FUSION at T=1: DISPATCH-MIRRORS the eager decode's fused3 (bit-identical body).
9874 // BATCHED EPILOGUE RE-FUSE (lane/vt-fixes fix 2): `h_q8` = the attn-input norm's q8_1
9875 // form emitted by the fused rms_norm_q8_1 (bit-identical to rms_norm_decode ->
9876 // quantize_q8_1, kernel-check-pinned). When present (caller checked mixer q8_1-fast),
9877 // every projection consumes it — `h` may be a zero-len placeholder and must not be read.
9878 let (qf, mut k, v) = if let Some(mut qkv) = self.full_attn_tp_qkv(e, fa, h, t)? {
9879 let v = qkv.pop().ok_or("full-attention TP verify QKV omitted V")?;
9880 let k = qkv.pop().ok_or("full-attention TP verify QKV omitted K")?;
9881 let q = qkv.pop().ok_or("full-attention TP verify QKV omitted Q")?;
9882 if !qkv.is_empty() {
9883 return Err("full-attention TP verify QKV returned extra projections".into());
9884 }
9885 (q, k, v)
9886 } else {
9887 let mut fused = None;
9888 let qkv_fast =
9889 e.uses_q8_1_fast(&fa.wq) && e.uses_q8_1_fast(&fa.wk) && e.uses_q8_1_fast(&fa.wv);
9890 if t == 1 && qkv_fast {
9891 let (hq_o, hd_o);
9892 let (hq, hd): (&CudaSlice<i8>, &CudaSlice<f32>) = match h_q8 {
9893 Some(p) => p,
9894 None => {
9895 (hq_o, hd_o) = e.quantize_q8_1(h, 1, n_embd)?;
9896 (&hq_o, &hd_o)
9897 }
9898 };
9899 fused = e.matmul_q8_fused3(&fa.wq, &fa.wk, &fa.wv, hq, hd)?;
9900 } else if spec_fused_t() && (2..=4).contains(&t) && qkv_fast {
9901 // VERIFY-TIER TRUNK FUSION (MEMRA_SPEC_FUSED_T): one shared quantize + one
9902 // fused3 batched launch replaces three decode-exact calls (3 re-quantizes of
9903 // the same h + 3 _b2/_b4 launches). Bit-identical per (tensor,token,row).
9904 let (hq_o, hd_o);
9905 let (hq, hd): (&CudaSlice<i8>, &CudaSlice<f32>) = match h_q8 {
9906 Some(p) => p,
9907 None => {
9908 (hq_o, hd_o) = e.quantize_q8_1(h, t, n_embd)?;
9909 (&hq_o, &hd_o)
9910 }
9911 };
9912 fused = e.matmul_q8_fused3_t(&fa.wq, &fa.wk, &fa.wv, hq, hd, t)?;
9913 }
9914 match (fused, h_q8) {
9915 (Some(triple), _) => triple,
9916 // shared pre-quantized activation (q8_1-fast guaranteed by the caller): the
9917 // decode-exact dispatch consumes (hq, hd) instead of re-quantizing 3x.
9918 (None, Some((hq, hd))) if qkv_fast => (
9919 e.matmul_decode_exact_pre(&fa.wq, hq, hd, t)?,
9920 e.matmul_decode_exact_pre(&fa.wk, hq, hd, t)?,
9921 e.matmul_decode_exact_pre(&fa.wv, hq, hd, t)?,
9922 ),
9923 (None, _) => (
9924 e.matmul_decode_exact(&fa.wq, h, t)?,
9925 e.matmul_decode_exact(&fa.wk, h, t)?,
9926 e.matmul_decode_exact(&fa.wv, h, t)?,
9927 ),
9928 }
9929 };
9930 // M3/Hy3 have no attention output gate — wq out is exactly q; skip the split.
9931 let gated = geometry.attention_gate == memra_gguf::config::AttentionGateKind::FusedQ;
9932 let (mut q, gate) = if gated {
9933 let mut q = vbuf(e, t * n_head * head_dim)?; // fully written by q_gate_split
9934 let mut gate = vbuf(e, t * n_head * head_dim)?; // fully written by q_gate_split
9935 e.q_gate_split(&qf, &mut q, &mut gate, head_dim, n_head, t)?;
9936 (q, Some(gate))
9937 } else {
9938 (qf, None)
9939 };
9940
9941 let mut qn = vbuf(e, t * n_head * head_dim)?; // fully written by rms_norm
9942 e.rms_norm(
9943 &q,
9944 fa.q_norm.float_data(),
9945 &mut qn,
9946 head_dim,
9947 n_head * t,
9948 eps,
9949 )?;
9950 q = qn;
9951 let mut kn = vbuf(e, t * n_head_kv * head_dim)?; // fully written by rms_norm
9952 e.rms_norm(
9953 &k,
9954 fa.k_norm.float_data(),
9955 &mut kn,
9956 head_dim,
9957 n_head_kv * t,
9958 eps,
9959 )?;
9960 k = kn;
9961 let rope_dims = geometry.n_rot as usize;
9962 e.rope_neox(
9963 &mut q,
9964 pos_d,
9965 head_dim,
9966 rope_dims,
9967 n_head,
9968 t,
9969 geometry.rope_base,
9970 1.0,
9971 )?;
9972 e.rope_neox(
9973 &mut k,
9974 pos_d,
9975 head_dim,
9976 rope_dims,
9977 n_head_kv,
9978 t,
9979 geometry.rope_base,
9980 1.0,
9981 )?;
9982
9983 // append T new K/V columns to the resident QUANTIZED cache. k/v are token-major [T, kv_dim]
9984 // f32; append-quantize each of the T token rows into the byte cache (q8_0 K / q5_1 V).
9985 let kvl = cache.kv[il].as_mut().unwrap();
9986 let (kv_dim_k, kv_dim_v, ktb, vtb) =
9987 (kvl.kv_dim_k, kvl.kv_dim_v, kvl.k_tok_bytes, kvl.v_tok_bytes);
9988 if let Some(ctr) = stream_ctr {
9989 // stream: ONE batched append at the device counter (rows kernel = the per-view warp
9990 // math on a (block, token) grid, documented byte-identical); host len is a stale
9991 // LOWER BOUND under pre-issue (drain reconciles it).
9992 e.append_kv_quantized_rows_dc(
9993 &k,
9994 &v,
9995 &mut kvl.k,
9996 &mut kvl.v,
9997 ctr,
9998 t,
9999 kv_dim_k,
10000 kv_dim_v,
10001 ktb,
10002 vtb,
10003 crate::Engine::kv_fp8_on(),
10004 )?;
10005 } else {
10006 for i in 0..t {
10007 let k_row = k.slice(i * kv_dim_k..(i + 1) * kv_dim_k);
10008 let v_row = v.slice(i * kv_dim_v..(i + 1) * kv_dim_v);
10009 e.append_kv_quantized_view(
10010 &k_row,
10011 &v_row,
10012 &mut kvl.k,
10013 &mut kvl.v,
10014 kvl.len + i,
10015 kv_dim_k,
10016 kv_dim_v,
10017 ktb,
10018 vtb,
10019 crate::Engine::kv_fp8_on(),
10020 )?;
10021 }
10022 kvl.len += t;
10023 }
10024
10025 // BIT-IDENTICAL VERIFY ATTENTION (spec-exactness fix): the FP accumulation order must be
10026 // byte-for-byte identical to the eager decode path. fa_prefill uses a different tile size
10027 // (BLOCK_Q=64, BK=32) and online-softmax structure than fa_decode's split-K + combine,
10028 // which changes FP summation order and can flip argmax at tight logit margins. Query row r
10029 // attends to keys [0..base_len+r+1) — each successive row sees one more key (the causal
10030 // property). This matches eager: decode appends k at len, then fa_decode sees t_kv = len+1
10031 // keys. The verify appends all T tokens first but bounds the key range per row.
10032 //
10033 // MULTI-ROW FUSED PATH (the long-ctx spec fix, 2026-07-03): when every row takes the vec
10034 // kernel (base_len+1 >= FA_VEC_MIN_TKV), ONE fa_decode_rows launch executes the exact
10035 // per-row program for all T rows (grid.z = row, per-row n_splits from the same
10036 // fa_split_keys formula) — replacing T x (2 launches + 2 dtod copies + 5 partial allocs)
10037 // and multiplying resident CTAs by T on a latency-bound kernel. Bit-identical per row by
10038 // construction; kernel-check pins rows-vs-loop byte identity, run-spec is the end gate.
10039 // Short ctx (any row below the vec crossover) and MEMRA_NO_FA_VEC/MEMRA_FA_ROWS_OFF keep the
10040 // per-row loop (whose fa_decode picks scalar/vec per row exactly like eager decode).
10041 let mut attn = vbuf(e, t * n_head * head_dim)?; // fully written by every FA arm below
10042 let base_len = kvl.len - t; // KV len BEFORE this round's T tokens were appended
10043 // T=1 INCLUDED (2026-07-05): p-min cuts the draft to 1 in ~75% of rounds on hard
10044 // (agentic) content — the old t>1 gate sent those rounds to the per-row loop (262us/row
10045 // + q-row copy + per-row allocs vs 93us/row through the fused kernel at grid.z=1, same
10046 // program). nsys accounting: 1088 of 1456 verify FA launches were T=1 escapees.
10047 // LEAN T=1 ARM (MEMRA_SPEC_LEAN, close35): at t==1, q IS one row and fa_decode on it is
10048 // the EXACT eager decode dispatch (vec_q_v2 + combine_f32; the rows pair measured +50us
10049 // at m=1). Byte-identical: kernel-check pins rows-vs-loop identity, and the per-row loop
10050 // at t=1 is fa_decode on the same q with zero-offset copies. Gates arbitrate.
10051 if let Some(ctr) = stream_ctr {
10052 // STREAM ARM: causal base from the device counter; host kvl.len is a stale lower
10053 // bound used only for the split-sizing upper bound (+64 slack covers M pre-issued
10054 // rounds at K<=8). Views span the bound; per-row limits derive in-kernel.
10055 let upper = kvl.len + t + 64;
10056 let k_view = e.view_u8(&kvl.k, (upper.min(cache.max_ctx)) * ktb);
10057 let v_view = e.view_u8(&kvl.v, (upper.min(cache.max_ctx)) * vtb);
10058 e.fa_decode_rows_dc(
10059 &q,
10060 &k_view,
10061 &v_view,
10062 &mut attn,
10063 head_dim,
10064 n_head,
10065 n_head_kv,
10066 ctr,
10067 upper.min(cache.max_ctx),
10068 t,
10069 scale,
10070 ktb,
10071 vtb,
10072 0,
10073 false,
10074 )?;
10075 } else if spec_lean() && t == 1 {
10076 let t_kv = base_len + 1;
10077 let k_view = e.view_u8(&kvl.k, t_kv * ktb);
10078 let v_view = e.view_u8(&kvl.v, t_kv * vtb);
10079 e.fa_decode_kvmod(
10080 &q,
10081 &k_view,
10082 &v_view,
10083 &mut attn,
10084 head_dim,
10085 n_head,
10086 n_head_kv,
10087 t_kv,
10088 scale,
10089 ktb,
10090 vtb,
10091 crate::Engine::kv_fp8_on(),
10092 )?;
10093 } else if e.fa_rows_eligible(base_len, head_dim) {
10094 let k_view = e.view_u8(&kvl.k, (base_len + t) * ktb);
10095 let v_view = e.view_u8(&kvl.v, (base_len + t) * vtb);
10096 e.fa_decode_rows(
10097 &q,
10098 &k_view,
10099 &v_view,
10100 &mut attn,
10101 head_dim,
10102 n_head,
10103 n_head_kv,
10104 base_len,
10105 t,
10106 scale,
10107 ktb,
10108 vtb,
10109 None,
10110 false,
10111 crate::Engine::kv_fp8_on(),
10112 None,
10113 )?;
10114 } else {
10115 for r in 0..t {
10116 let t_kv_r = base_len + r + 1; // this row sees keys [0..t_kv_r)
10117 let k_view_r = e.view_u8(&kvl.k, t_kv_r * ktb);
10118 let v_view_r = e.view_u8(&kvl.v, t_kv_r * vtb);
10119 // copy q row into an owned buffer (fa_decode takes &CudaSlice, not CudaView)
10120 let mut q_row = vbuf(e, n_head * head_dim)?; // fully written by copy_view_into
10121 let q_src = q.slice(r * n_head * head_dim..(r + 1) * n_head * head_dim);
10122 e.copy_view_into(&mut q_row, 0, &q_src, n_head * head_dim)?;
10123 let mut attn_row = vbuf(e, n_head * head_dim)?; // fully written by fa_decode
10124 e.fa_decode_kvmod(
10125 &q_row,
10126 &k_view_r,
10127 &v_view_r,
10128 &mut attn_row,
10129 head_dim,
10130 n_head,
10131 n_head_kv,
10132 t_kv_r,
10133 scale,
10134 ktb,
10135 vtb,
10136 crate::Engine::kv_fp8_on(),
10137 )?;
10138 e.copy_into(
10139 &mut attn,
10140 r * n_head * head_dim,
10141 &attn_row,
10142 n_head * head_dim,
10143 )?;
10144 }
10145 }
10146
10147 let attn_g = match &gate {
10148 Some(gate) => {
10149 let mut gsig = vbuf(e, t * n_head * head_dim)?; // fully written by sigmoid
10150 e.sigmoid(gate, &mut gsig, t * n_head * head_dim)?;
10151 let mut ag = vbuf(e, t * n_head * head_dim)?; // fully written by mul
10152 e.mul(&attn, &gsig, &mut ag, t * n_head * head_dim)?;
10153 ag
10154 }
10155 None => attn,
10156 };
10157 // DECODE-EXACT wo projection: at m>=5 (K=4+ with pending) the generic matmul would use dp4a
10158 // (128-thread, different FP sum order than MMVQ). Force MMVQ for bit-identity with decode.
10159 match self.full_attn_tp_o(e, fa, &attn_g, t)? {
10160 Some(output) => Ok(output),
10161 None => Ok(e.matmul_decode_exact(&fa.wo, &attn_g, t)?),
10162 }
10163 }
10164
10165 /// Context-linear bytes for a plain serving session's trunk cache.
10166 pub fn plain_session_kv_bytes_per_token(&self) -> usize {
10167 crate::cache::cache_bytes_per_token_for_plan(
10168 &self.cfg,
10169 &self.plan,
10170 0,
10171 self.plan.layers.len(),
10172 )
10173 }
10174
10175 /// `(logical bytes/token, ring-capped bytes/token, ring row cap)` for exact admission.
10176 pub fn plain_session_kv_shape(&self) -> (usize, usize, usize) {
10177 (
10178 self.plain_session_kv_bytes_per_token(),
10179 crate::cache::cache_ring_bytes_per_token_for_plan(
10180 &self.cfg,
10181 &self.plan,
10182 0,
10183 self.plan.layers.len(),
10184 ),
10185 crate::cache::cache_ring_row_cap_for_plan(&self.plan),
10186 )
10187 }
10188
10189 /// Context-linear bytes for a speculative serving session: trunk cache plus persistent MTP
10190 /// scratch. With no MTP head this equals the plain coefficient.
10191 pub fn spec_session_kv_bytes_per_token(&self) -> usize {
10192 let scratch = self
10193 .mtp
10194 .iter()
10195 .chain(self.mtp_extra.iter())
10196 .map(|mtp| {
10197 let (_, _, k, v) = mtp_scratch_layout(&self.cfg, mtp.geom.as_ref());
10198 k + v
10199 })
10200 .sum::<usize>();
10201 self.plain_session_kv_bytes_per_token()
10202 .saturating_add(scratch)
10203 }
10204
10205 /// Spec twin of [`HybridModel::plain_session_kv_shape`]; Step35's persistent MTP scratch is
10206 /// capped by the same SWA ring rows as the trunk.
10207 pub fn spec_session_kv_shape(&self) -> (usize, usize, usize) {
10208 let total = self.spec_session_kv_bytes_per_token();
10209 let (_, mut ring, rows) = self.plain_session_kv_shape();
10210 if rows > 0 {
10211 ring = ring.saturating_add(
10212 self.mtp
10213 .iter()
10214 .chain(self.mtp_extra.iter())
10215 .map(|mtp| {
10216 let (_, _, k, v) = mtp_scratch_layout(&self.cfg, mtp.geom.as_ref());
10217 k + v
10218 })
10219 .sum::<usize>(),
10220 );
10221 }
10222 (total, ring, rows)
10223 }
10224
10225 /// Greedy MTP speculative decode (§B). Token-identical to `generate(prompt, max_new)` but uses
10226 /// the NextN head to draft K tokens then verifies them in one batched target forward.
10227 /// Returns (generated tokens, total_drafted, total_accepted) so the caller can report
10228 /// acceptance rate. `k` = draft length per round.
10229 ///
10230 /// GRAPH DRAFT (stage 2 of graph-grade spec): when the model is all-Dense and the MTP head is
10231 /// Dense (no MoE host readbacks), the fixed-shape T=1 MTP forward is CUDA-graph-captured ONCE
10232 /// and replayed per draft step — the ~40 eager launches per drafted token collapse into one
10233 /// graph dispatch; only the 4-byte token id (and 4-byte p-min confidence) round-trip per step.
10234 /// Event tracking is disabled for the whole call (generate_graph pattern) so every buffer the
10235 /// captured graph references is event-free; the spec loop is strictly single-stream.
10236 /// MEMRA_SPEC_NOGRAPH=1 forces the eager draft chain.
10237 /// SAMPLED mode (MEMRA_SPEC_TEMP>0) has its OWN capture (gumbel-perturbed in-graph argmax,
10238 /// device Philox event counter, persistent q retention) — graph-vs-eager sampled streams are
10239 /// bit-identical for the same (seed, prompt, K, temp); see the sampled-graph setup in
10240 /// generate_spec_inner2.
10241 /// Multi-turn session: trunk cache + MTP draft scratch persist across generate calls, so
10242 /// turn N+1 primes ONLY its new suffix (the 124k-conversation daily pattern — re-priming a
10243 /// 32k history costs ~54s; a suffix prime costs seconds). APPEND-ONLY by construction: the
10244 /// hybrid linear-attn states are in-place (no position index), so a session can extend but
10245 /// never rewind — `committed` is the exact token list whose state the caches hold (includes
10246 /// any overshoot tokens past max_new; the caller renders from `committed`, not its own echo).
10247 pub fn new_session(
10248 &self,
10249 e: &Engine,
10250 max_ctx: usize,
10251 ) -> Result<SpecSession, Box<dyn std::error::Error>> {
10252 Ok(SpecSession {
10253 // STAGE-OWNED KV (lane/pp2-spec 2026-08-06): `pp::new_cache`, not `Cache::new`. This
10254 // is the SERVING spec-session path, and with the ppN door open across two cards a
10255 // primary-homed cache makes every remote stage peer-read its OWN KV on every verify
10256 // round — the wrong-card class already fixed on the two batched serving paths
10257 // (worker.rs 2483 / 2837). With the door shut `new_cache` IS `Cache::new` (same
10258 // branch, same allocations), so single-device behavior is byte-unchanged.
10259 cache: crate::pp::new_cache_planned(e, &self.cfg, &self.plan, max_ctx)?,
10260 scratch: self.new_mtp_scratch(e, max_ctx)?,
10261 committed: Vec::new(),
10262 last_h: None,
10263 next_pred: None,
10264 sctr: 0,
10265 uctr: 0,
10266 draft_ctx: None,
10267 pending_tok: None,
10268 turn_ckpt: None,
10269 telem: SpecTelemetryCounters::default(),
10270 capture_at: None,
10271 boundary_captures: Vec::new(),
10272 ckpt_at: None,
10273 capture_disabled: false,
10274 })
10275 }
10276
10277 /// SPEC-ON-CACHE-HIT restore (lane/spec-on-cache-hit, 2026-08-18 — PORT-PLAN item 3,
10278 /// research/cache-spec-design-20260814, scoped to WHOLE-ENTRY restores only): build a
10279 /// SpecSession around a trunk cache the worker already restored from a prefix-cache
10280 /// entry, re-installing the entry's published draft plane as the MTP scratch rows
10281 /// `[0..prefix.len())` and the entry's boundary hidden as `last_h`, then feeding the
10282 /// prompt SUFFIX here — through EXACTLY the plain path's program selection — so the
10283 /// worker always receives a fully-warm continuation session (committed = whole
10284 /// prompt, `next_pred` + `last_h` set; caller sets `next_pred` from the entry's
10285 /// boundary logits on the empty-suffix shape).
10286 ///
10287 /// PROGRAM LAW (the splitiso two-programs class, learned AGAIN in this lane's own
10288 /// gate): the identity target for a converted hit is the PLAIN hit serving the same
10289 /// request, and plain feeds a carried suffix via eager `decode_step` below
10290 /// PRIME_MIN_T and via `prime_cache` at/above it (prefill_tick's arms). The generate
10291 /// path's tokenwise arm routes qwen35-class through the BATCHED T=1 program
10292 /// (`spec_target_step_h`) instead — ULP-different suffix rows, and the gate measured
10293 /// the near-tie flip at generated token ~8 (research/spec-cache-20260818, qwen r3).
10294 /// So the suffix is fed HERE, mirroring prefill_tick arm-for-arm, not handed to the
10295 /// burst prime.
10296 ///
10297 /// SEED RULE (both sampling regimes; lane/sampled-hit-spec 2026-08-19, sampled draw
10298 /// added by lane/sampled-spec-quality 2026-08-19). The boundary token is produced by
10299 /// EXACTLY the rule the cold burst entry applies to its own first token from the same
10300 /// logits row: `argmax` when greedy, and a `sample_boundary_token` draw at Philox
10301 /// counter 0 when sampled. Both shapes are covered — the entry's boundary logits on a
10302 /// full-cover (empty-suffix) hit, this feed's own boundary logits on a suffix hit.
10303 /// That is what keeps a restored session seed-identical to a cold one PER SEED: the
10304 /// cold session draws from the identical row at counter 0 and then runs its rounds from
10305 /// counter 1, so the restored session admits with `sctr = 1` after its own draw.
10306 /// The WORKER owns the one refusal this constructor cannot see — a constrained request.
10307 /// (The penalized-sampled refusal was LIFTED once the burst's penalty window learned to
10308 /// span the session: `committed` here is the WHOLE prompt, so the restored session's
10309 /// window is the cold session's window. It comes back if `MEMRA_SPEC_PEN_SESSION=0`.)
10310 ///
10311 /// NOT the rolled-back partial-restore hazard: the caller restores at exactly the
10312 /// entry's captured endpoint (`e.pos`) through the shipping whole-entry path;
10313 /// mid-entry (`at < e.pos`) trunk restores stay behind MEMRA_PREFIX_PARTIAL_RESTORE
10314 /// and are never routed here.
10315 ///
10316 /// Failure contract: `Err((Some(cache), why))` before any trunk mutation — the
10317 /// worker rebuilds the plain carrier and the hit serves plain, byte-unchanged.
10318 /// `Err((None, why))` after the suffix feed began — the carrier is part-fed and
10319 /// UNUSABLE; the worker serves the request cold-plain (correct, slower) and the
10320 /// entry stays published for the next request.
10321 #[allow(clippy::too_many_arguments)]
10322 #[allow(clippy::result_large_err)] // allow: the fat error type is the diagnostic contract here; boxing it would change the error surface
10323 pub fn spec_session_from_restored(
10324 &self,
10325 e: &Engine,
10326 mut cache: Cache,
10327 prefix: Vec<u32>,
10328 suffix: &[u32],
10329 draft_k: &CudaSlice<u8>,
10330 draft_v: &CudaSlice<u8>,
10331 draft_k_tok_bytes: usize,
10332 draft_v_tok_bytes: usize,
10333 draft_len: usize,
10334 last_h: &[f32],
10335 // The ENTRY's boundary logits row (the full-cover shape's seed source). May be empty
10336 // when a suffix follows — the feed's own logits are the boundary then.
10337 boundary_logits: &[f32],
10338 // The request's sampler, or None for greedy. Owned here so the seed rule lives in
10339 // ONE place instead of being half-applied by the worker.
10340 sampling: Option<SpecSampling>,
10341 require_anchor: bool,
10342 max_ctx: usize,
10343 // STABLE-BOUNDARY REPUBLICATION (lane/frspec-multiturn-cache, 2026-08-21): ABSOLUTE
10344 // prompt position to split the suffix feed at and capture the extended-entry
10345 // publication + this session's `turn_ckpt` — the worker's stable pre-generation
10346 // boundary (`plain_checkpoint_boundary`). None = legacy prompt-end republication.
10347 // WHY: the prompt-end capture below includes the template's live generation header
10348 // (`<|im_start|>assistant\n<think>\n`), which the next turn's re-render replaces, so
10349 // for a hybrid (whole-entry restores only) every extended entry's last ~2 tokens
10350 // diverged from every future prompt and the hit boundary FROZE at the first
10351 // lcp-split entry forever (measured: cached 6811 of 38228 by turn 8, B4).
10352 republish_at: Option<usize>,
10353 ) -> Result<SpecSession, (Option<Cache>, String)> {
10354 let pos = prefix.len();
10355 let fail = |cache: Cache, msg: String| -> Result<SpecSession, (Option<Cache>, String)> {
10356 Err((Some(cache), msg))
10357 };
10358 if let Err(error) = cache.ensure_usable("spec_session_from_restored") {
10359 drop(cache);
10360 return Err((None, error.to_string()));
10361 }
10362 if self.mtp.is_none() {
10363 return fail(cache, "no MTP head attached (nothing to draft with)".into());
10364 }
10365 if pos == 0 {
10366 return fail(cache, "empty committed prefix".into());
10367 }
10368 if cache.pos != pos {
10369 let msg = format!(
10370 "restored cache pos {} != restored prefix len {pos}",
10371 cache.pos
10372 );
10373 return fail(cache, msg);
10374 }
10375 if draft_len != pos {
10376 return fail(
10377 cache,
10378 format!("draft plane len {draft_len} != restored prefix len {pos}"),
10379 );
10380 }
10381 if pos + suffix.len() >= max_ctx {
10382 return fail(
10383 cache,
10384 format!(
10385 "prompt {} + suffix would not leave generation room in ctx {max_ctx}",
10386 pos + suffix.len(),
10387 ),
10388 );
10389 }
10390 let mut scratch = match MtpScratch::new(
10391 e,
10392 &self.cfg,
10393 &self.plan,
10394 max_ctx,
10395 self.mtp.as_ref().and_then(|m| m.geom.as_ref()),
10396 ) {
10397 Ok(s) => s,
10398 Err(err) => return fail(cache, format!("draft scratch alloc failed: {err}")),
10399 };
10400 if scratch.kv.ring.is_some() {
10401 return fail(
10402 cache,
10403 "ring-backed draft scratch (Step35 SWA) cannot take a flat prefix restore".into(),
10404 );
10405 }
10406 if scratch.kv.k_tok_bytes != draft_k_tok_bytes
10407 || scratch.kv.v_tok_bytes != draft_v_tok_bytes
10408 {
10409 return fail(
10410 cache,
10411 format!(
10412 "draft plane layout {draft_k_tok_bytes}/{draft_v_tok_bytes} != scratch \
10413 {}/{} bytes/token (stale entry across a format change)",
10414 scratch.kv.k_tok_bytes, scratch.kv.v_tok_bytes,
10415 ),
10416 );
10417 }
10418 if pos > scratch.cap {
10419 return fail(
10420 cache,
10421 format!(
10422 "draft plane rows {pos} exceed scratch capacity {}",
10423 scratch.cap
10424 ),
10425 );
10426 }
10427 let kb = pos * draft_k_tok_bytes;
10428 let vb = pos * draft_v_tok_bytes;
10429 if draft_k.len() < kb || draft_v.len() < vb {
10430 return fail(
10431 cache,
10432 format!(
10433 "truncated draft plane: K {} < {kb} or V {} < {vb} bytes",
10434 draft_k.len(),
10435 draft_v.len(),
10436 ),
10437 );
10438 }
10439 if kb > 0
10440 && let Err(err) = e.copy_u8_into(&mut scratch.kv.k, 0, draft_k, kb)
10441 {
10442 return fail(cache, format!("draft K restore copy failed: {err}"));
10443 }
10444 if vb > 0
10445 && let Err(err) = e.copy_u8_into(&mut scratch.kv.v, 0, draft_v, vb)
10446 {
10447 return fail(cache, format!("draft V restore copy failed: {err}"));
10448 }
10449 if let Err(err) = scratch.set_len(e, pos) {
10450 return fail(cache, format!("draft scratch len set failed: {err}"));
10451 }
10452 let mut last_h_dev = if last_h.len() == self.cfg.n_embd as usize {
10453 // anchor upload failure is acceptance-only when a suffix feed follows (fill
10454 // row-0 falls back to zeros) but FATAL for an empty-suffix continuation (the
10455 // burst entry asserts committed + last_h + next_pred) — the caller says which.
10456 e.htod(last_h).ok()
10457 } else {
10458 None
10459 };
10460 if require_anchor && last_h_dev.is_none() {
10461 return fail(
10462 cache,
10463 "empty-suffix continuation requires the entry's boundary hidden anchor".into(),
10464 );
10465 }
10466 let mut committed = prefix;
10467 // Set on BOTH shapes below (suffix-fed and full-cover) — never left None, which is
10468 // what the empty-suffix continuation assert in the burst entry requires.
10469 let next_pred;
10470 // Philox: (0,0) at admit exactly like a fresh session; a sampled boundary draw below
10471 // consumes counter 0 and leaves 1, which is the state a cold session reaches after
10472 // drawing its own first token from the same row.
10473 let mut sctr = 0u32;
10474 let sampled = sampling.is_some_and(|s| s.temp > 0.0) && spec_sampled_boundary_on();
10475 // Penalty window for the boundary draw: the last `penalty_last_n` tokens of the WHOLE
10476 // prompt, which is what the cold session's own burst sees (Item 2's window). Built
10477 // after the suffix joins `committed` below.
10478 let mut boundary_captures: Vec<SpecBoundaryCapture> = Vec::new();
10479 let mut restored_turn_ckpt: Option<SpecCheckpoint> = None;
10480 if !suffix.is_empty() {
10481 // ---- SUFFIX FEED, mirroring prefill_tick's program selection exactly ----
10482 // From here on the trunk cache mutates: failures return Err((None, _)) and
10483 // the worker serves the request cold-plain instead of reusing the carrier.
10484 let dirty =
10485 |msg: String| -> Result<SpecSession, (Option<Cache>, String)> { Err((None, msg)) };
10486 let n_embd = self.cfg.n_embd as usize;
10487 let t = suffix.len();
10488 let mut h_rows = match e.uninit(t * n_embd) {
10489 Ok(b) => b,
10490 Err(err) => return fail(cache, format!("suffix hidden buffer alloc: {err}")),
10491 };
10492 // STABLE-BOUNDARY split (see `republish_at`): feed stops at the boundary so the
10493 // in-place GDN conv/ssm state can be snapshotted there — the only moment it
10494 // exists (the cold prime-split law). suffix-relative; None = one-segment legacy.
10495 let b_rel = republish_at
10496 .and_then(|abs| abs.checked_sub(pos))
10497 .filter(|&r| r > 0 && r < t);
10498 let mut feed_logits = Vec::new();
10499 let tokenwise_env = std::env::var("MEMRA_PRIME_TOKENWISE").is_ok()
10500 || e.frozen_cpu_experts_prefer_tokenwise_prime();
10501 let mut fed = 0usize;
10502 for seg_end in [b_rel, Some(t)].into_iter().flatten() {
10503 if seg_end <= fed {
10504 continue;
10505 }
10506 let seg = &suffix[fed..seg_end];
10507 let batched = seg.len() >= crate::hybrid_forward::PRIME_MIN_T && !tokenwise_env;
10508 if batched {
10509 // prefill_tick's prime arm: request-level prime_cache call; tokens still
10510 // queued after this segment ride `queued_after` so Step35 arm selection
10511 // stays keyed to the request's end (tick-seg law).
10512 match self.prime_cache(e, seg, &mut cache, t - seg_end) {
10513 Ok((l, _h_seed, hiddens)) => {
10514 if let Err(err) =
10515 e.copy_into(&mut h_rows, fed * n_embd, &hiddens, seg.len() * n_embd)
10516 {
10517 return dirty(format!("suffix hidden copy: {err}"));
10518 }
10519 feed_logits = l;
10520 }
10521 Err(err) => return dirty(format!("suffix prime failed: {err}")),
10522 }
10523 } else {
10524 // prefill_tick's tokenwise arm: eager decode_step, one token at a time.
10525 for (i, &tok) in seg.iter().enumerate() {
10526 match self.decode_step_h(e, tok, &mut cache) {
10527 Ok((l, h)) => {
10528 if let Err(err) =
10529 e.copy_into(&mut h_rows, (fed + i) * n_embd, &h, n_embd)
10530 {
10531 return dirty(format!("suffix hidden copy: {err}"));
10532 }
10533 feed_logits = l;
10534 }
10535 Err(err) => return dirty(format!("suffix decode_step failed: {err}")),
10536 }
10537 }
10538 }
10539 fed = seg_end;
10540 if Some(seg_end) == b_rel {
10541 // The stable pre-generation boundary: capture the extended-entry
10542 // publication AND this session's own turn checkpoint here instead of at
10543 // prompt-end (both would otherwise carry the volatile live-header tail
10544 // the next re-render replaces). Failure silent, turn_ckpt convention.
10545 debug_assert_eq!(
10546 cache.pos,
10547 pos + seg_end,
10548 "stable-boundary capture off the feed split"
10549 );
10550 if spec_restore_republish_on()
10551 && let Ok(snap) = cache.snapshot(e)
10552 {
10553 boundary_captures.push(SpecBoundaryCapture {
10554 snap,
10555 pos: pos + seg_end,
10556 logits: feed_logits.clone(),
10557 last_h: capture_boundary_hidden(e, &h_rows, seg_end, n_embd),
10558 latent_tails: Vec::new(),
10559 });
10560 }
10561 let anchor: Result<CudaSlice<f32>, Box<dyn std::error::Error>> =
10562 e.uninit(n_embd).and_then(|mut a| {
10563 e.copy_view_into(
10564 &mut a,
10565 0,
10566 &h_rows.slice((seg_end - 1) * n_embd..seg_end * n_embd),
10567 n_embd,
10568 )?;
10569 Ok(a)
10570 });
10571 if let (Ok(snap), Ok(last_h)) = (cache.snapshot(e), anchor) {
10572 restored_turn_ckpt = Some(SpecCheckpoint {
10573 snap,
10574 pos: pos + seg_end,
10575 last_h,
10576 });
10577 }
10578 }
10579 }
10580 // Draft-scratch fill for the suffix rows, predecessor-paired: row `pos` reads
10581 // the entry's boundary anchor (zeros fallback — acceptance-only), row `pos+i`
10582 // reads h_rows[i-1]. Chunked like the generate path's fill (transients scale
10583 // with T). Fill failures are acceptance-only — truncate to the restored rows
10584 // and continue; the burst's own set_len keeps the invariant.
10585 let _mtp = self.mtp.as_ref().expect("mtp checked above"); // invariant check only; the fill below re-reads self.mtp
10586 let (embd_qt, embd_rb) = self.embd.qt_and_row_bytes(n_embd);
10587 let embd_gpu = if spec_host_embd() {
10588 None
10589 } else {
10590 Some(
10591 self.embd_gpu
10592 .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload")),
10593 )
10594 };
10595 let embd_dev = embd_gpu.map(|g| (g, embd_qt, embd_rb));
10596 let fill_chunk = 4096usize;
10597 let mut filled = true;
10598 let mut start = 0usize;
10599 'fill: while start < t {
10600 let end = (start + fill_chunk).min(t);
10601 let tc = end - start;
10602 let Ok(mut phs) = e.zeros(tc * n_embd) else {
10603 filled = false;
10604 break 'fill;
10605 };
10606 let (src_lo, dst_off, n_copy) = if start == 0 {
10607 (0, n_embd, (tc - 1) * n_embd)
10608 } else {
10609 ((start - 1) * n_embd, 0, tc * n_embd)
10610 };
10611 if start == 0
10612 && let Some(lh) = last_h_dev.as_ref()
10613 && e.copy_into(&mut phs, 0, lh, n_embd).is_err()
10614 {
10615 filled = false;
10616 break 'fill;
10617 }
10618 if n_copy > 0
10619 && e.copy_view_into(
10620 &mut phs,
10621 dst_off,
10622 &h_rows.slice(src_lo..src_lo + n_copy),
10623 n_copy,
10624 )
10625 .is_err()
10626 {
10627 filled = false;
10628 break 'fill;
10629 }
10630 if self
10631 .mtp_kv_fill_all(
10632 e,
10633 &suffix[start..end],
10634 &phs,
10635 pos + start,
10636 &mut scratch,
10637 embd_dev,
10638 )
10639 .is_err()
10640 {
10641 filled = false;
10642 break 'fill;
10643 }
10644 start = end;
10645 }
10646 if !filled {
10647 // acceptance-only: drafts over missing suffix rows are cheap and wrong,
10648 // so keep only the restored rows resident and let verify arbitrate.
10649 if let Err(err) = scratch.set_len(e, pos) {
10650 return dirty(format!("scratch truncation after failed fill: {err}"));
10651 }
10652 }
10653 // EXTENDED-ENTRY PUBLICATION (lane/sampled-spec-quality, Item 3 — the fix for
10654 // "a restored spec session never publishes an extended entry", SAMPLED-HIT.md
10655 // finding (d)). Pre-lane, publication was armed only for COLD sessions
10656 // (`spec_resumed == 0` in the worker) and both engine capture sites require a
10657 // non-continuation burst — but a converted hit's first burst IS a continuation,
10658 // so a growing conversation learned exactly ONE boundary and turn 3 could never
10659 // hit a longer prefix than turn 2 did.
10660 //
10661 // WHERE, and why it is safe here: `cache.pos == prefix + suffix` at this exact
10662 // line — the trunk is primed over the whole prompt, nothing is generated, and the
10663 // draft plane rows [0..prompt) are filled just above. That is a complete
10664 // whole-entry boundary (`pos == fed_len`), the same shape the cold seed capture
10665 // publishes; the worker's existing publication sweep picks it up because it is
10666 // keyed on non-empty `boundary_captures` and is sampler- and resume-independent.
10667 // NOT the partial-restore hazard: the boundary is this session's own prompt END,
10668 // never mid-entry, so `entry_pos != fed_len` still refuses on the way back in.
10669 // Failure is SILENT by design (the turn_ckpt / boundary-capture convention):
10670 // publication is an optimization, never a correctness dependency.
10671 //
10672 // SUPERSEDED WHEN `republish_at` FIRED (lane/frspec-multiturn-cache): a prompt-end
10673 // entry's tail is the live generation header the next re-render replaces, so on a
10674 // hybrid (whole-entry restores) it can never serve the conversation's next turn —
10675 // the stable-boundary capture above IS this publication, minus the poisoned tail.
10676 if spec_restore_republish_on() && boundary_captures.is_empty() {
10677 debug_assert_eq!(
10678 cache.pos,
10679 pos + t,
10680 "extended-entry capture must sit at the restored session's prompt end",
10681 );
10682 if let Ok(snap) = cache.snapshot(e) {
10683 boundary_captures.push(SpecBoundaryCapture {
10684 snap,
10685 pos: pos + t,
10686 logits: feed_logits.clone(),
10687 last_h: capture_boundary_hidden(e, &h_rows, t, n_embd),
10688 latent_tails: Vec::new(),
10689 });
10690 }
10691 }
10692 // continuation seed: the feed's boundary logits ARE the plain path's boundary
10693 // logits (same program), so greedy's argmax here is plain's first emitted token,
10694 // and the sampled draw is the cold sampled session's own first token.
10695 next_pred = Some(if sampled {
10696 let sp = sampling.expect("sampled implies a sampler");
10697 // `committed` is still the restored prefix here; the suffix joins it below —
10698 // so this is the last-N window over the WHOLE prompt, exactly the cold
10699 // session's own window at its first token.
10700 let hist = pen_window_seed(&committed, suffix, sp.penalty_last_n);
10701 match sample_boundary_token(
10702 e,
10703 &feed_logits,
10704 &sp,
10705 &hist,
10706 &mut sctr,
10707 "restore-suffix-feed",
10708 ) {
10709 Ok(t) => t,
10710 // the trunk is already fed: hand nothing back, the worker serves the
10711 // request cold-plain. Never fall back to an argmax — that would put a
10712 // greedy token in a sampled stream to save a slow path.
10713 Err(err) => {
10714 return dirty(format!("boundary token draw failed: {err}"));
10715 }
10716 }
10717 } else {
10718 argmax(&feed_logits) as u32
10719 });
10720 let mut lh = match e.uninit(n_embd) {
10721 Ok(b) => b,
10722 Err(err) => return dirty(format!("boundary hidden alloc: {err}")),
10723 };
10724 if let Err(err) = e.copy_view_into(
10725 &mut lh,
10726 0,
10727 &h_rows.slice((t - 1) * n_embd..t * n_embd),
10728 n_embd,
10729 ) {
10730 return dirty(format!("boundary hidden copy: {err}"));
10731 }
10732 last_h_dev = Some(lh);
10733 committed.extend_from_slice(suffix);
10734 } else {
10735 // FULL-COVER shape (empty suffix — the identical-repeat / agent-loop shape): the
10736 // ENTRY's boundary logits are the boundary row, and this is the token the cold
10737 // session emits from that same row. Owned here rather than in the worker so the
10738 // sampled draw cannot be half-applied on one shape (the worker used to argmax it).
10739 if boundary_logits.is_empty() {
10740 return fail(
10741 cache,
10742 "full-cover restore without the entry's boundary logits".into(),
10743 );
10744 }
10745 next_pred = Some(if sampled {
10746 let sp = sampling.expect("sampled implies a sampler");
10747 let hist = pen_window_seed(&committed, &[], sp.penalty_last_n);
10748 match sample_boundary_token(
10749 e,
10750 boundary_logits,
10751 &sp,
10752 &hist,
10753 &mut sctr,
10754 "restore-full-cover",
10755 ) {
10756 Ok(t) => t,
10757 // nothing has been mutated on this shape — hand the carrier back and let
10758 // the hit serve PLAIN (the banked pre-lane path).
10759 Err(err) => {
10760 return fail(cache, format!("boundary token draw failed: {err}"));
10761 }
10762 }
10763 } else {
10764 argmax(boundary_logits) as u32
10765 });
10766 }
10767 Ok(SpecSession {
10768 cache,
10769 scratch,
10770 committed,
10771 last_h: last_h_dev,
10772 next_pred,
10773 sctr,
10774 uctr: 0,
10775 draft_ctx: None,
10776 pending_tok: None,
10777 // Stable-boundary capture from the split feed above (None on the legacy shape):
10778 // a restored session previously parked WITHOUT a checkpoint, so the next turn's
10779 // affinity probe declined ("no turn checkpoint retained") and the conversation
10780 // fell back to the frozen prefix entry forever.
10781 turn_ckpt: restored_turn_ckpt,
10782 telem: SpecTelemetryCounters::default(),
10783 capture_at: None,
10784 boundary_captures,
10785 ckpt_at: None,
10786 capture_disabled: false,
10787 })
10788 }
10789
10790 /// Forced-gate exact state comparison. This intentionally reads the real live prefixes from
10791 /// their owning PP devices: matching emitted ids alone would miss a stale `len_d`, recurrent
10792 /// snapshot, or draft-KV row that only corrupts the following round.
10793 pub fn optipipe_compare_session_state(
10794 &self,
10795 e: &Engine,
10796 reference: &SpecSession,
10797 candidate: &SpecSession,
10798 ) -> Result<OptiForkStateIdentity, Box<dyn std::error::Error>> {
10799 fn fail(what: &str) -> Box<dyn std::error::Error> {
10800 format!("optipipe state mismatch: {what}").into()
10801 }
10802 fn same_f32(a: &[f32], b: &[f32]) -> bool {
10803 a.len() == b.len() && a.iter().zip(b).all(|(x, y)| x.to_bits() == y.to_bits())
10804 }
10805 fn compare_layers(
10806 es: &Engine,
10807 range: std::ops::Range<usize>,
10808 reference: &SpecSession,
10809 candidate: &SpecSession,
10810 report: &mut OptiForkStateIdentity,
10811 ) -> Result<(), Box<dyn std::error::Error>> {
10812 for il in range {
10813 match (&reference.cache.kv[il], &candidate.cache.kv[il]) {
10814 (Some(a), Some(b)) => {
10815 if a.len != b.len {
10816 return Err(fail(&format!(
10817 "layer {il} host KV len {} != {}",
10818 a.len, b.len
10819 )));
10820 }
10821 let ad = es.dtoh_i32(&a.len_d)?;
10822 let bd = es.dtoh_i32(&b.len_d)?;
10823 if ad != bd || ad.first().copied() != Some(a.len as i32) {
10824 return Err(fail(&format!(
10825 "layer {il} device KV len {ad:?} != {bd:?} (host={})",
10826 a.len,
10827 )));
10828 }
10829 let kb = a.len * a.k_tok_bytes;
10830 let vb = a.len * a.v_tok_bytes;
10831 if kb > 0 {
10832 let ak = es.dtoh_u8_view(&a.k.slice(0..kb))?;
10833 let bk = es.dtoh_u8_view(&b.k.slice(0..kb))?;
10834 if ak != bk {
10835 let at = ak.iter().zip(&bk).position(|(x, y)| x != y).unwrap();
10836 return Err(fail(&format!(
10837 "layer {il} K bytes at byte {at} row {} offset {}: {} != {}",
10838 at / a.k_tok_bytes,
10839 at % a.k_tok_bytes,
10840 ak[at],
10841 bk[at],
10842 )));
10843 }
10844 }
10845 if vb > 0 {
10846 let av = es.dtoh_u8_view(&a.v.slice(0..vb))?;
10847 let bv = es.dtoh_u8_view(&b.v.slice(0..vb))?;
10848 if av != bv {
10849 let at = av.iter().zip(&bv).position(|(x, y)| x != y).unwrap();
10850 return Err(fail(&format!(
10851 "layer {il} V bytes at byte {at} row {} offset {}: {} != {}",
10852 at / a.v_tok_bytes,
10853 at % a.v_tok_bytes,
10854 av[at],
10855 bv[at],
10856 )));
10857 }
10858 }
10859 report.trunk_kv_bytes += kb + vb;
10860 }
10861 (None, None) => {}
10862 _ => return Err(fail(&format!("layer {il} KV presence"))),
10863 }
10864 match (&reference.cache.recur[il], &candidate.cache.recur[il]) {
10865 (Some(a), Some(b)) => {
10866 let ac = es.dtoh(&a.conv_state)?;
10867 let bc = es.dtoh(&b.conv_state)?;
10868 if !same_f32(&ac, &bc) {
10869 return Err(fail(&format!("layer {il} conv state")));
10870 }
10871 let as_ = es.dtoh(&a.ssm_state)?;
10872 let bs = es.dtoh(&b.ssm_state)?;
10873 if !same_f32(&as_, &bs) {
10874 return Err(fail(&format!("layer {il} SSM state")));
10875 }
10876 report.recurrent_bytes += (ac.len() + as_.len()) * 4;
10877 }
10878 (None, None) => {}
10879 _ => return Err(fail(&format!("layer {il} recurrent presence"))),
10880 }
10881 }
10882 Ok(())
10883 }
10884
10885 if reference.committed != candidate.committed {
10886 return Err(fail("committed token ids"));
10887 }
10888 if reference.cache.pos != candidate.cache.pos
10889 || reference.cache.max_ctx != candidate.cache.max_ctx
10890 {
10891 return Err(fail("cache pos/capacity"));
10892 }
10893 if reference.pending_tok != candidate.pending_tok
10894 || reference.next_pred != candidate.next_pred
10895 || reference.sctr != candidate.sctr
10896 || reference.uctr != candidate.uctr
10897 {
10898 return Err(fail("pending/prediction/counter tail"));
10899 }
10900
10901 let mut report = OptiForkStateIdentity::default();
10902 if let Some(fence) = crate::pp::pp_cuts(self.layers.len()) {
10903 let rt = crate::pp::PpNRt::get(e)?;
10904 for stage in 0..rt.n_stages() {
10905 let _scope = rt.enter(stage);
10906 compare_layers(
10907 rt.engine(stage, e),
10908 fence[stage]..fence[stage + 1],
10909 reference,
10910 candidate,
10911 &mut report,
10912 )?;
10913 }
10914 } else {
10915 compare_layers(e, 0..self.layers.len(), reference, candidate, &mut report)?;
10916 }
10917
10918 if reference.scratch.plane_count() != candidate.scratch.plane_count() {
10919 return Err(fail("draft scratch plane count"));
10920 }
10921 for index in 0..reference.scratch.plane_count() {
10922 let (a, _) = reference.scratch.plane(index);
10923 let (b, _) = candidate.scratch.plane(index);
10924 if a.len != b.len
10925 || a.kv_dim_k != b.kv_dim_k
10926 || a.kv_dim_v != b.kv_dim_v
10927 || a.k_tok_bytes != b.k_tok_bytes
10928 || a.v_tok_bytes != b.v_tok_bytes
10929 || e.dtoh_i32(&a.len_d)? != e.dtoh_i32(&b.len_d)?
10930 {
10931 return Err(fail(&format!("draft scratch plane {index} length/layout")));
10932 }
10933 let kb = a.len * a.k_tok_bytes;
10934 let vb = a.len * a.v_tok_bytes;
10935 if kb > 0 && e.dtoh_u8_view(&a.k.slice(0..kb))? != e.dtoh_u8_view(&b.k.slice(0..kb))? {
10936 return Err(fail(&format!("draft scratch plane {index} K bytes")));
10937 }
10938 if vb > 0 && e.dtoh_u8_view(&a.v.slice(0..vb))? != e.dtoh_u8_view(&b.v.slice(0..vb))? {
10939 return Err(fail(&format!("draft scratch plane {index} V bytes")));
10940 }
10941 report.scratch_kv_bytes += kb + vb;
10942 }
10943
10944 match (&reference.last_h, &candidate.last_h) {
10945 (Some(a), Some(b)) => {
10946 let ah = e.dtoh(a)?;
10947 let bh = e.dtoh(b)?;
10948 if !same_f32(&ah, &bh) {
10949 return Err(fail("last hidden/seed bytes"));
10950 }
10951 report.hidden_bytes = ah.len() * 4;
10952 }
10953 (None, None) => {}
10954 _ => return Err(fail("last hidden/seed presence")),
10955 }
10956 Ok(report)
10957 }
10958
10959 /// SESSION-AFFINITY REWIND (lane/session-affinity, 2026-08-05): roll `sess` back to its
10960 /// retained prompt-end checkpoint, so a request whose prompt matches
10961 /// `committed[..rewind_pos()]` exactly can resume there and prime only its own delta.
10962 ///
10963 /// EXACTNESS. After this returns, the session is byte-for-byte the state it was in AT that
10964 /// boundary: full-attn KV truncated to it (append-only, position-addressed), GDN conv/ssm
10965 /// restored from the device copy taken there, draft scratch length reset, `committed`
10966 /// truncated, `last_h` = the boundary's predecessor anchor. That is precisely the state a
10967 /// fresh prime of `committed[..pos]` would have produced, so the following suffix prime and
10968 /// every burst after it are identical to a cold run of the same token stream — the
10969 /// committed-tokens-authoritative contract.
10970 ///
10971 /// `next_pred` and `pending_tok` are CLEARED: both describe generation past the boundary,
10972 /// which the rewind discards. The caller therefore must supply a non-empty suffix (a
10973 /// rewound session cannot serve an empty-suffix continuation burst — there is nothing to
10974 /// continue). The persistent draft graph survives: it bakes only session-stable pointers
10975 /// (the scratch KV, the resident embedding), none of which the rewind moves.
10976 ///
10977 /// The checkpoint is CONSUMED (`turn_ckpt` taken): its snapshot buffers are freed here, and
10978 /// this turn's own prime installs a fresh one at the new prompt end. Returns the position
10979 /// rewound to, or `None` when the session holds no checkpoint (caller: full re-prime).
10980 pub fn spec_rewind_to_checkpoint(
10981 &self,
10982 e: &Engine,
10983 sess: &mut SpecSession,
10984 ) -> Result<Option<usize>, Box<dyn std::error::Error>> {
10985 if sess.turn_ckpt.as_ref().is_some_and(|ckpt| {
10986 !sess.cache.can_rollback(&ckpt.snap, 0) || !sess.scratch.can_rewind_to(ckpt.pos)
10987 }) {
10988 return Err(
10989 "SWA ring rewind checkpoint has been lapped; full re-prime required".into(),
10990 );
10991 }
10992 let Some(ckpt) = sess.turn_ckpt.take() else {
10993 return Ok(None);
10994 };
10995 assert!(
10996 ckpt.pos <= sess.committed.len(),
10997 "checkpoint past committed ({} > {})",
10998 ckpt.pos,
10999 sess.committed.len()
11000 );
11001 // Restore through each layer's owning engine. A single primary-engine rollback is not
11002 // sufficient when the serving cache is stage-owned under cross-device PP.
11003 crate::pp::restore_cache_checkpoint(e, self, None, &mut sess.cache, &ckpt.snap)?;
11004 debug_assert_eq!(
11005 sess.cache.pos, ckpt.pos,
11006 "rollback landed off the checkpoint"
11007 );
11008 sess.scratch.set_len(e, ckpt.pos)?;
11009 sess.committed.truncate(ckpt.pos);
11010 sess.last_h = Some(ckpt.last_h);
11011 sess.next_pred = None;
11012 sess.pending_tok = None;
11013 Ok(Some(ckpt.pos))
11014 }
11015
11016 /// Grow a parked speculative session to `target_cap` and rewind it to its retained turn
11017 /// checkpoint without re-priming the checkpoint prefix.
11018 ///
11019 /// The trunk cache is restored exactly like a plain grown cache: append-only full-attention
11020 /// KV rows come from the parked cache, while recurrent state comes from the checkpoint's
11021 /// owned snapshot. The MTP scratch is also context-linear and its rows below the checkpoint
11022 /// remain authoritative, so they are copied into a fresh larger scratch before its length is
11023 /// truncated. Pointer-baking draft graphs are dropped and recaptured on the next burst.
11024 ///
11025 /// All fallible work completes before `sess` is mutated. A failed allocation or copy leaves
11026 /// the parked session intact, allowing the caller one reclaim-and-retry attempt.
11027 pub fn spec_grow_and_rewind_to_checkpoint(
11028 &self,
11029 e: &Engine,
11030 sess: &mut SpecSession,
11031 target_cap: usize,
11032 ) -> Result<Option<usize>, Box<dyn std::error::Error>> {
11033 if target_cap <= sess.cache.max_ctx {
11034 return self.spec_rewind_to_checkpoint(e, sess);
11035 }
11036 let Some(ckpt) = sess.turn_ckpt.as_ref() else {
11037 return Ok(None);
11038 };
11039 if ckpt.pos == 0 || ckpt.pos > sess.committed.len() {
11040 return Err(format!(
11041 "checkpoint pos {} outside committed length {}",
11042 ckpt.pos,
11043 sess.committed.len(),
11044 )
11045 .into());
11046 }
11047 if ckpt.pos > target_cap {
11048 return Err(format!(
11049 "checkpoint pos {} exceeds grown capacity {target_cap}",
11050 ckpt.pos,
11051 )
11052 .into());
11053 }
11054
11055 let mut grown_cache = crate::pp::new_cache_planned(e, &self.cfg, &self.plan, target_cap)?;
11056 let mut grown_scratch = self.new_mtp_scratch(e, target_cap)?;
11057 crate::pp::restore_cache_checkpoint(
11058 e,
11059 self,
11060 Some(&sess.cache),
11061 &mut grown_cache,
11062 &ckpt.snap,
11063 )?;
11064
11065 if sess.scratch.plane_count() != grown_scratch.plane_count() {
11066 return Err("checkpoint draft plane count mismatch".into());
11067 }
11068 for index in 0..sess.scratch.plane_count() {
11069 let (src, _) = sess.scratch.plane(index);
11070 let (dst, _) = grown_scratch.plane_mut(index);
11071 if ckpt.pos > src.len
11072 || src.kv_dim_k != dst.kv_dim_k
11073 || src.kv_dim_v != dst.kv_dim_v
11074 || src.k_tok_bytes != dst.k_tok_bytes
11075 || src.v_tok_bytes != dst.v_tok_bytes
11076 {
11077 return Err(format!(
11078 "checkpoint draft plane {index} layout mismatch (pos {}, source len {})",
11079 ckpt.pos, src.len,
11080 )
11081 .into());
11082 }
11083 match (&src.ring, dst.ring.as_ref()) {
11084 (Some(sring), Some(_)) => {
11085 // Ring-backed draft plane (step35): `ckpt.pos` is absolute and exceeds the
11086 // physical rows once lapped — same class as the trunk-KV restore panic
11087 // (2026-08-29 warm-turn-at-40k). Copy the aligned live window, rebase.
11088 let (new_base, phys) = sring.restore_plan(ckpt.pos).map_err(|err| {
11089 format!("checkpoint draft plane {index} SWA restore refused: {err}")
11090 })?;
11091 let rows = phys.len();
11092 let kb = rows * src.k_tok_bytes;
11093 let vb = rows * src.v_tok_bytes;
11094 if kb > 0 {
11095 e.copy_u8_range_into(
11096 &mut dst.k,
11097 0,
11098 &src.k,
11099 phys.start * src.k_tok_bytes,
11100 kb,
11101 )?;
11102 }
11103 if vb > 0 {
11104 e.copy_u8_range_into(
11105 &mut dst.v,
11106 0,
11107 &src.v,
11108 phys.start * src.v_tok_bytes,
11109 vb,
11110 )?;
11111 }
11112 dst.ring
11113 .as_mut()
11114 .expect("ring presence checked above")
11115 .apply_rebase(new_base);
11116 if let Some(base_d) = dst.base_d.as_mut() {
11117 e.set_i32_one(base_d, new_base as i32)?;
11118 }
11119 }
11120 (None, None) => {
11121 let kb = ckpt.pos * src.k_tok_bytes;
11122 let vb = ckpt.pos * src.v_tok_bytes;
11123 if kb > 0 {
11124 e.copy_u8_into(&mut dst.k, 0, &src.k, kb)?;
11125 }
11126 if vb > 0 {
11127 e.copy_u8_into(&mut dst.v, 0, &src.v, vb)?;
11128 }
11129 }
11130 _ => {
11131 return Err(format!("checkpoint draft plane {index} ring/flat mismatch").into());
11132 }
11133 }
11134 }
11135 grown_scratch.set_len(e, ckpt.pos)?;
11136 // The old scratch is dropped immediately after publication below. Bound its D2D reads
11137 // first; growth happens once per rewritten turn, outside the decode hot loop.
11138 e.stream().synchronize()?;
11139
11140 let ckpt = sess
11141 .turn_ckpt
11142 .take()
11143 .expect("checkpoint remained present through transactional grow");
11144 let pos = ckpt.pos;
11145 sess.cache = grown_cache;
11146 sess.scratch = grown_scratch;
11147 sess.committed.truncate(pos);
11148 sess.last_h = Some(ckpt.last_h);
11149 sess.next_pred = None;
11150 sess.pending_tok = None;
11151 sess.draft_ctx = None;
11152 debug_assert_eq!(sess.cache.pos, pos, "grown rewind landed off checkpoint");
11153 debug_assert!(
11154 (0..sess.scratch.plane_count()).all(|index| sess.scratch.plane(index).0.len == pos),
11155 "grown draft rewind landed off checkpoint"
11156 );
11157 Ok(Some(pos))
11158 }
11159
11160 /// Commit a carried pending bonus (see SpecSession::pending_tok): one T=1 trunk pass
11161 /// (its logits' argmax becomes next_pred) + the draft-KV fill at the carried anchor —
11162 /// byte-identical to the pre-carry session tail. Required before a non-empty-suffix
11163 /// prime, a sampled turn, or parking a session for pool reuse. No-op without a pending.
11164 /// `sampling` is the sampler of the request that will CONSUME the resulting `next_pred`
11165 /// (lane/sampled-spec-quality): this is a boundary site like any other, so a sampled
11166 /// consumer must get a DRAWN token, not an argmax. Pass `None` from the park/demote
11167 /// callers — a pending only ever exists on the GREEDY tail, and the consumer of a
11168 /// park-time flush is a future request whose sampler is not knowable here (residual
11169 /// named at the pool-resume probe in worker.rs and in SAMPLED-QUALITY.md).
11170 pub fn spec_flush_pending(
11171 &self,
11172 e: &Engine,
11173 sess: &mut SpecSession,
11174 sampling: Option<SpecSampling>,
11175 ) -> Result<(), Box<dyn std::error::Error>> {
11176 sess.cache.ensure_usable("spec_flush_pending")?;
11177 let Some(b) = sess.pending_tok.take() else {
11178 return Ok(());
11179 };
11180 if self.mtp.is_none() {
11181 return Err("pending carry requires an MTP head".into());
11182 }
11183 let n_embd = self.cfg.n_embd as usize;
11184 let (embd_qt, embd_rb) = self.embd.qt_and_row_bytes(n_embd);
11185 let embd_gpu = if spec_host_embd() {
11186 None
11187 } else {
11188 Some(
11189 self.embd_gpu
11190 .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload")),
11191 )
11192 };
11193 let embd_dev = embd_gpu.map(|g| (g, embd_qt, embd_rb));
11194 let pos_b = sess.cache.pos;
11195 sess.scratch.set_len(e, pos_b)?;
11196 let (lg_b, hb) = self.spec_target_step_h(e, b, &mut sess.cache)?;
11197 sess.next_pred = Some(match sampling {
11198 Some(sp) if sp.temp > 0.0 && spec_sampled_boundary_on() => {
11199 // window includes `b` itself: it is committed by this pass, and the pre-lane
11200 // code never counted a boundary token in the penalty history at all.
11201 let hist = pen_window_seed(&sess.committed, &[b], sp.penalty_last_n);
11202 sample_boundary_token(e, &lg_b, &sp, &hist, &mut sess.sctr, "flush-pending")?
11203 }
11204 _ => argmax(&lg_b) as u32,
11205 });
11206 let anchor = sess
11207 .last_h
11208 .as_ref()
11209 .expect("pending carry requires last_h (the predecessor-row anchor)");
11210 self.mtp_kv_fill_all(e, &[b], anchor, pos_b, &mut sess.scratch, embd_dev)?;
11211 sess.last_h = Some(hb);
11212 sess.committed.push(b);
11213 Ok(())
11214 }
11215
11216 /// Solo target feed used only at speculative round boundaries. Step35 serving made its
11217 /// staged batched B=1 graph authoritative, so a speculative session must enter and leave
11218 /// rounds through that same graph. Other model families keep their eager T=1 contract.
11219 fn spec_target_step_h(
11220 &self,
11221 e: &Engine,
11222 token: u32,
11223 cache: &mut Cache,
11224 ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
11225 cache.ensure_usable("spec_target_step_h")?;
11226 if !self.sliding_gated_moe_batch_program() && !self.batched_serving_numeric_class() {
11227 return self.decode_step_h(e, token, cache);
11228 }
11229 let pos0 = cache.pos;
11230 let (logits, hidden) = self.decode_step_t_core(e, &[token], pos0, cache, None, None)?;
11231 Ok((e.dtoh(&logits)?, hidden))
11232 }
11233
11234 /// The archs whose LIVE B=1 serving runs the generic BATCHED numeric class (decode_step_batch
11235 /// walk + batched head), so their spec verify must run the SAME class. MoE learned this
11236 /// 2026-08-14 AM (4b777ccc5); the dense hybrid reproduced the identical near-tie flip class
11237 /// the same day on Qwen3.8-27B — eager-class verify logits drift from batched-class serving
11238 /// logits ("1 ULP at layer 2 → 2.3e-1 logit maxdiff at the head"), and the GDN recurrence
11239 /// carries the drift until a near-tie flips deep in generation. One predicate so the five
11240 /// dispatch sites cannot drift apart again.
11241 /// Draft-graph head admissibility (lane/draftcost-moe, 2026-08-20): the capture body
11242 /// (`mtp_head_forward_cap`) supports Dense heads and SOFTMAX device-routed resident-MoE
11243 /// heads. Residency alone is insufficient: Hy3/M3/Step sigmoid routing returns selected
11244 /// experts through a host synchronization, which is capture-illegal. Those heads use the
11245 /// exact eager draft chain until a device-only sigmoid expert program lands. Trunk FFN class
11246 /// is irrelevant — the graph body is the HEAD forward only. One predicate for all three
11247 /// eligibility sites so they cannot drift (the serving numeric-class lesson).
11248 fn mtp_graph_capturable(&self) -> bool {
11249 let sigmoid_router = self.cfg.sigmoid_router().is_some();
11250 for head in self.mtp.iter().chain(self.mtp_extra.iter()) {
11251 let reason = match &head.ffn {
11252 crate::hybrid::Ffn::Dense { .. } => None,
11253 crate::hybrid::Ffn::Moe(mo) if mo.dev_exps.is_none() => {
11254 Some("non-resident MoE MTP head")
11255 }
11256 crate::hybrid::Ffn::Moe(_) if sigmoid_router => {
11257 Some("sigmoid-router MoE MTP head requires host-visible routing")
11258 }
11259 crate::hybrid::Ffn::Moe(_) => None,
11260 };
11261 if let Some(reason) = reason {
11262 static NOTICE: std::sync::Once = std::sync::Once::new();
11263 NOTICE.call_once(|| {
11264 eprintln!(
11265 "[spec] draft graph unavailable: {reason}; eager draft chain engaged"
11266 );
11267 });
11268 return false;
11269 }
11270 }
11271 self.mtp.is_some()
11272 }
11273
11274 fn batched_serving_numeric_class(&self) -> bool {
11275 self.plan
11276 .trunk_operations()
11277 .contains(&memra_gguf::model_plan::OperationKind::GatedDeltaNet)
11278 }
11279
11280 /// The family the MTP verify-graph default was measured on: GatedDeltaNet state layers
11281 /// (a `recur` mixer) together with a routed-MoE FFN — Ornith-1.5-35B-A3B and its kin. The
11282 /// server-side twin of this test is `model_forces_spec_replay` (GatedDeltaNet + MoeMlp);
11283 /// keeping the engine's own version structural rather than name-based means a new
11284 /// checkpoint of the same shape inherits the default, and a different shape does not.
11285 /// pub(crate) since lane/graph-launch-guard-sweep-20260831: `dspark_vg_admission_debt`
11286 /// consults it so the MTP-route pool stops escaping the admission charge.
11287 pub(crate) fn vgraph_family_default(&self) -> bool {
11288 let has_linear = self
11289 .layers
11290 .iter()
11291 .any(|l| matches!(l.mixer, Mixer::Linear(_)));
11292 let has_moe = self
11293 .layers
11294 .iter()
11295 .any(|l| matches!(l.ffn, crate::hybrid::Ffn::Moe(_)));
11296 has_linear && has_moe
11297 }
11298
11299 fn sliding_gated_moe_batch_program(&self) -> bool {
11300 self.uses_sliding_gated_moe_program()
11301 }
11302
11303 fn gemma_batch_program(&self) -> bool {
11304 self.uses_gemma_program()
11305 }
11306
11307 /// Reduced-matrix admission for increment 1. This deliberately does not change the PP-2
11308 /// serving policy: the worker calls it only after `MEMRA_SPEC_PIPE=1` and an explicit spec
11309 /// session already exist.
11310 pub fn spec_pipe_available(&self, e: &Engine) -> bool {
11311 if std::env::var("MEMRA_SPEC_PIPE").as_deref() != Ok("1")
11312 || !spec_devacc()
11313 || spec_replay_env_enabled()
11314 || spec_stream()
11315 || std::env::var("MEMRA_SPEC_ADAPT").as_deref() == Ok("1")
11316 || std::env::var("MEMRA_SPEC_PMIN0").as_deref() == Ok("1")
11317 || std::env::var("MEMRA_SPEC_PP_ANATOMY").as_deref() == Ok("1")
11318 || std::env::var("MEMRA_SPEC_PMIN")
11319 .ok()
11320 .and_then(|v| v.parse::<f32>().ok())
11321 .unwrap_or(0.0)
11322 > 0.0
11323 || self.is_gemma4_e4b()
11324 || self.gemma_batch_program()
11325 || self.mtp.is_none()
11326 || !self.mtp_extra.is_empty()
11327 // Both paired lanes would otherwise hold the model-global verify-graph mutex across
11328 // setup and wait for each other. Independent graph pools are future work; the pair
11329 // requires the explicit eager-verify arm today.
11330 || crate::spec::spec_verify_graph_env()
11331 .unwrap_or_else(|| self.vgraph_family_default())
11332 {
11333 return false;
11334 }
11335 let Some(cuts) = crate::pp::pp_cuts(self.layers.len()) else {
11336 return false;
11337 };
11338 if cuts.len() != 3 || crate::pp::pp2_streams_off() || !crate::pp::spec_pp_on() {
11339 return false;
11340 }
11341 crate::pp::PpNRt::get(e)
11342 .map(|rt| rt.n_stages() == 2 && rt.cross_device())
11343 .unwrap_or(false)
11344 }
11345
11346 /// Two warm greedy continuation bursts over one PP-2 interval coordinator. The two existing
11347 /// `generate_spec_inner2` call stacks own all per-session round locals; only phase issue order
11348 /// changes. No callback is accepted in increment 1 — the worker publishes each completed burst.
11349 #[allow(clippy::too_many_arguments)]
11350 #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
11351 pub fn generate_spec_session_pair(
11352 &self,
11353 e: &Engine,
11354 sess_a: &mut SpecSession,
11355 max_new_a: usize,
11356 k_a: usize,
11357 sess_b: &mut SpecSession,
11358 max_new_b: usize,
11359 k_b: usize,
11360 ) -> Result<((Vec<u32>, usize, usize), (Vec<u32>, usize, usize)), Box<dyn std::error::Error>>
11361 {
11362 self.refuse_hyper("generate_spec_session_pair")?;
11363 if !self.spec_pipe_available(e) {
11364 return Err("two-session speculative pipeline is outside its reduced matrix".into());
11365 }
11366 let rt = crate::pp::PpNRt::get(e)?;
11367 let pp_walk = rt.acquire_walk("generate_spec_session_pair")?;
11368 let pp_permit = rt.walk_permit(&pp_walk, "generate_spec_session_pair")?;
11369 if max_new_a == 0 || max_new_b == 0 || k_a == 0 || k_b == 0 {
11370 return Err(
11371 "two-session speculative pipeline requires non-empty positive-K bursts".into(),
11372 );
11373 }
11374 for sess in [&*sess_a, &*sess_b] {
11375 if sess.committed.is_empty()
11376 || sess.last_h.is_none()
11377 || (sess.next_pred.is_none() && sess.pending_tok.is_none())
11378 {
11379 return Err("two-session speculative pipeline requires warm continuations".into());
11380 }
11381 }
11382
11383 let graph_ok = std::env::var("MEMRA_SPEC_NOGRAPH").is_err()
11384 && !spec_host_embd()
11385 && self.mtp_graph_capturable()
11386 && self.mtp_extra.is_empty()
11387 && !crate::model::full_prec_enabled();
11388 let graph_a = graph_ok && k_a + 2 < 96;
11389 let graph_b = graph_ok && k_b + 2 < 96;
11390 let was_tracking = e.ctx().is_event_tracking();
11391 if (graph_a || graph_b) && was_tracking {
11392 unsafe {
11393 e.ctx().disable_event_tracking();
11394 }
11395 }
11396
11397 static LOGGED: std::sync::Once = std::sync::Once::new();
11398 LOGGED.call_once(|| {
11399 eprintln!("[spec-pipe] two-session PP-2 continuation pipeline engaged");
11400 });
11401 let sync = std::sync::Arc::new(SpecPipeSync::new());
11402 let lane_a = SpecPipeLane {
11403 sync: sync.clone(),
11404 lane: 0,
11405 rt,
11406 walk_permit: pp_permit.clone(),
11407 };
11408 let lane_b = SpecPipeLane {
11409 sync,
11410 lane: 1,
11411 rt,
11412 walk_permit: pp_permit,
11413 };
11414 let mut sess_b_ptr = SpecPipeSessionPtr(sess_b as *mut SpecSession);
11415 let (result_a, result_b) = std::thread::scope(|scope| {
11416 let b = scope.spawn(move || {
11417 let mut finish = SpecPipeFinish::new(&lane_b);
11418 let sess_b = unsafe { sess_b_ptr.get_mut() };
11419 let result = (|| -> Result<_, String> {
11420 e.ctx().bind_to_thread().map_err(|err| err.to_string())?;
11421 self.generate_spec_inner2(
11422 e,
11423 &[],
11424 max_new_b,
11425 k_b,
11426 graph_b,
11427 Some(sess_b),
11428 None,
11429 None,
11430 None,
11431 None,
11432 Some(&lane_b),
11433 )
11434 .map_err(|err| err.to_string())
11435 })();
11436 finish.close(result.is_err());
11437 result
11438 });
11439 let mut finish = SpecPipeFinish::new(&lane_a);
11440 let result_a = self.generate_spec_inner2(
11441 e,
11442 &[],
11443 max_new_a,
11444 k_a,
11445 graph_a,
11446 Some(sess_a),
11447 None,
11448 None,
11449 None,
11450 None,
11451 Some(&lane_a),
11452 );
11453 finish.close(result_a.is_err());
11454 let result_b = b
11455 .join()
11456 .map_err(|_| "paired speculative session B panicked".to_string())
11457 .and_then(|r| r);
11458 (result_a, result_b)
11459 });
11460
11461 if (graph_a || graph_b) && was_tracking {
11462 unsafe {
11463 e.ctx().enable_event_tracking();
11464 }
11465 }
11466 let result_a = result_a?;
11467 let result_b = result_b.map_err(|err| -> Box<dyn std::error::Error> { err.into() })?;
11468 Ok((result_a, result_b))
11469 }
11470
11471 /// One spec-decode turn on a live session. `suffix` = the NEW tokens only (turn N+1's user
11472 /// message rendered through the chat template continuation). Returns (new tokens emitted,
11473 /// drafted, accepted); session.committed grows by suffix + emitted.
11474 pub fn generate_spec_session(
11475 &self,
11476 e: &Engine,
11477 sess: &mut SpecSession,
11478 suffix: &[u32],
11479 max_new: usize,
11480 k: usize,
11481 ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
11482 self.generate_spec_session_sampled(e, sess, suffix, max_new, k, None, None)
11483 }
11484
11485 /// Serve-path sampled spec: routes the burst through the rejection-sampling verify with
11486 /// per-SESSION Philox continuity (sess.sctr/uctr). None = env-driven (CLI) or greedy.
11487 /// Filters (top-k/p/min-p) apply SYMMETRICALLY to draft q and verify p — distribution-exact
11488 /// for the filtered target (feat/filtered-spec).
11489 ///
11490 /// `on_commit` (sse-cadence, 2026-08-05): called with each newly-emitted slice of the
11491 /// output — once right after the prime's first token, then once per round commit — so a
11492 /// streaming caller can flush text at round cadence instead of once per burst. The slices
11493 /// are disjoint, in order, and concatenate to exactly the returned token vec. Emission-
11494 /// timing only: token bytes, session state, and exactness are untouched.
11495 ///
11496 /// The returned bool is a CONTINUE-VERDICT (admission yield, 2026-08-06): `false` ends
11497 /// the burst at the current round boundary, exactly as if `max_new` had been reached —
11498 /// the caller's scheduler regains control without waiting the burst out. Burst size is
11499 /// content-neutral (spec-levers battery), so an early exit moves WHEN the burst returns,
11500 /// never what tokens say. The slice may be EMPTY (a poll-only boundary — round-stream
11501 /// drains and the defensive tail flush can land with nothing new committed).
11502 #[allow(clippy::too_many_arguments)]
11503 #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
11504 pub fn generate_spec_session_sampled(
11505 &self,
11506 e: &Engine,
11507 sess: &mut SpecSession,
11508 suffix: &[u32],
11509 max_new: usize,
11510 k: usize,
11511 sampling: Option<SpecSampling>,
11512 on_commit: Option<&mut dyn FnMut(&[u32]) -> bool>,
11513 ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
11514 self.generate_spec_session_sampled_prime_split(
11515 e, sess, suffix, max_new, k, sampling, None, on_commit,
11516 )
11517 }
11518
11519 /// Serve-only cold-prime segmentation twin. `prime_split` is the same stable boundary the
11520 /// plain worker would honor before entering its sub-floor tokenwise tail; warm continuations
11521 /// pass `None` and stay on the existing zero-prime path.
11522 #[allow(clippy::too_many_arguments)]
11523 #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
11524 pub fn generate_spec_session_sampled_prime_split(
11525 &self,
11526 e: &Engine,
11527 sess: &mut SpecSession,
11528 suffix: &[u32],
11529 max_new: usize,
11530 k: usize,
11531 sampling: Option<SpecSampling>,
11532 prime_split: Option<usize>,
11533 on_commit: Option<&mut dyn FnMut(&[u32]) -> bool>,
11534 ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
11535 self.generate_spec_session_constrained_prime_split(
11536 e,
11537 sess,
11538 suffix,
11539 max_new,
11540 k,
11541 sampling,
11542 None,
11543 prime_split,
11544 on_commit,
11545 )
11546 }
11547
11548 /// `generate_spec_session_sampled` + GRAMMAR (constrained decoding, 2026-08-03): the
11549 /// hook truncates acceptance at the first grammar-illegal token AFTER the exactness
11550 /// verify (grammar is an extra rejection rule, ordering like the batched-verify twins)
11551 /// and replaces an illegal bonus with the MASKED argmax of the target's own verify
11552 /// column — token-identical to constrained plain greedy decode. GREEDY only (the
11553 /// worker routes sampled constrained to plain decode). Acceptance under tight grammars
11554 /// may drop (drafter is unconstrained); that is measured, not hidden.
11555 #[allow(clippy::too_many_arguments)]
11556 #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
11557 pub fn generate_spec_session_constrained(
11558 &self,
11559 e: &Engine,
11560 sess: &mut SpecSession,
11561 suffix: &[u32],
11562 max_new: usize,
11563 k: usize,
11564 sampling: Option<SpecSampling>,
11565 constraint: Option<&mut dyn SpecConstraint>,
11566 on_commit: Option<&mut dyn FnMut(&[u32]) -> bool>,
11567 ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
11568 self.generate_spec_session_constrained_prime_split(
11569 e, sess, suffix, max_new, k, sampling, constraint, None, on_commit,
11570 )
11571 }
11572
11573 #[allow(clippy::too_many_arguments)]
11574 #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
11575 pub fn generate_spec_session_constrained_prime_split(
11576 &self,
11577 e: &Engine,
11578 sess: &mut SpecSession,
11579 suffix: &[u32],
11580 max_new: usize,
11581 k: usize,
11582 sampling: Option<SpecSampling>,
11583 constraint: Option<&mut dyn SpecConstraint>,
11584 prime_split: Option<usize>,
11585 on_commit: Option<&mut dyn FnMut(&[u32]) -> bool>,
11586 ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
11587 if constraint.is_some() && sampling.is_some_and(|s| s.temp > 0.0) {
11588 return Err(
11589 "constrained spec decode is greedy-only (worker routes sampled \
11590 constrained to plain decode)"
11591 .into(),
11592 );
11593 }
11594 // PENDING-CARRY entry flush: a carried bonus precedes any new suffix in the sequence,
11595 // so it must commit BEFORE the suffix primes; the sampled path doesn't carry (its
11596 // round-0 accept needs the commit pass's logits). Empty-suffix greedy bursts — the
11597 // serve continuation case — consume the carry in-loop with zero solo passes.
11598 if sess.pending_tok.is_some()
11599 && (!suffix.is_empty() || sampling.is_some_and(|s| s.temp > 0.0))
11600 {
11601 self.spec_flush_pending(e, sess, sampling)?;
11602 }
11603
11604 // FULL_PREC forces the EAGER draft: the graph capture would enclose cuBLASLt f32 GEMV
11605 // (the FloatBf16 else-branches) and a bf16_to_f32 dequant alloc — neither is stream-capture
11606 // safe. Eager rides matmul/matmul_decode_exact, which dequant FloatBf16 on use. (§item 2.)
11607 // Multi-head MTP (mtp_extra non-empty) no longer disqualifies: the chain captures
11608 // per-head graphs (lane/step37-draft-graph-serving-20260830, MEMRA_MTP_CHAIN_GRAPH).
11609 let graph_draft = std::env::var("MEMRA_SPEC_NOGRAPH").is_err()
11610 && !spec_host_embd()
11611 && self.mtp_graph_capturable()
11612 && k + 2 < 96
11613 && !crate::model::full_prec_enabled();
11614 let was_tracking = e.ctx().is_event_tracking();
11615 if graph_draft && was_tracking {
11616 unsafe {
11617 e.ctx().disable_event_tracking();
11618 }
11619 }
11620 let r = self.generate_spec_inner2(
11621 e,
11622 suffix,
11623 max_new,
11624 k,
11625 graph_draft,
11626 Some(sess),
11627 sampling,
11628 constraint,
11629 on_commit,
11630 prime_split,
11631 None,
11632 );
11633 if graph_draft && was_tracking {
11634 unsafe {
11635 e.ctx().enable_event_tracking();
11636 }
11637 }
11638 let (out, d, a) = r?;
11639 Ok((out, d, a))
11640 }
11641
11642 pub fn generate_spec(
11643 &self,
11644 e: &Engine,
11645 prompt: &[u32],
11646 max_new: usize,
11647 k: usize,
11648 ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
11649 // glm5 T-parallel verify door (lane/glm5-tparallel-verify): an hc trunk with a
11650 // loaded DRAFT SOURCE — the embedded MTP head OR the DFlash2 drafter
11651 // (lane/glm5-dflash-draft-src) — routes to the glm5 draft->verify->rollback loop —
11652 // MEMRA_GLM5_SPEC=1 only (default OFF; flag row in FLAGS.md). Unset/0 falls
11653 // through to the standing named refusal below, byte-identical to the pre-lane
11654 // binary. Same fail-closed manifest stance as the generic path: an unqualified
11655 // MtpSpec rewrite refuses before any drafting.
11656 if self.hyper.is_some()
11657 && crate::glm_spec::glm5_spec_on()
11658 && (self.mtp.is_some() || self.glm5_dflash.is_some())
11659 {
11660 if !self.rewrite_allowed(memra_gguf::execution_manifest::RewriteSurface::MtpSpec) {
11661 return Err("speculative rewrite is not qualified for this ModelPlan".into());
11662 }
11663 return self.generate_spec_glm5(e, prompt, max_new, k);
11664 }
11665 self.refuse_hyper("generate_spec")?;
11666 if crate::pp::pp_cuts(self.layers.len()).is_some()
11667 && !self.rewrite_allowed(memra_gguf::execution_manifest::RewriteSurface::Pipeline)
11668 {
11669 return Err("pipeline rewrite is not qualified for speculative decode".into());
11670 }
11671 if !self.rewrite_allowed(memra_gguf::execution_manifest::RewriteSurface::MtpSpec) {
11672 return Err("speculative rewrite is not qualified for this ModelPlan".into());
11673 }
11674 // FULL_PREC forces eager (see generate_spec_session note): CUDA graph capture cannot
11675 // enclose cuBLASLt f32 GEMV or the bf16_to_f32 dequant alloc the FloatBf16 path needs.
11676 // Multi-head MTP no longer disqualifies (chain graphs; see generate_spec_session).
11677 let graph_draft = std::env::var("MEMRA_SPEC_NOGRAPH").is_err()
11678 && !spec_host_embd()
11679 && self.mtp_graph_capturable()
11680 && k + 2 < 96
11681 && !crate::model::full_prec_enabled();
11682 if !graph_draft {
11683 return self.generate_spec_inner2(
11684 e, prompt, max_new, k, false, None, None, None, None, None, None,
11685 );
11686 }
11687 let was_tracking = e.ctx().is_event_tracking();
11688 if was_tracking {
11689 unsafe {
11690 e.ctx().disable_event_tracking();
11691 }
11692 }
11693 let r = self.generate_spec_inner2(
11694 e, prompt, max_new, k, true, None, None, None, None, None, None,
11695 );
11696 if was_tracking {
11697 unsafe {
11698 e.ctx().enable_event_tracking();
11699 }
11700 }
11701 r
11702 }
11703
11704 #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
11705 #[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
11706 fn generate_spec_inner2(
11707 &self,
11708 e: &Engine,
11709 prompt: &[u32],
11710 max_new: usize,
11711 k: usize,
11712 graph_draft: bool,
11713 mut sess: Option<&mut SpecSession>,
11714 sampling: Option<SpecSampling>,
11715 mut constraint: Option<&mut dyn SpecConstraint>,
11716 mut on_commit: Option<&mut dyn FnMut(&[u32]) -> bool>,
11717 prime_split: Option<usize>,
11718 pipe: Option<&SpecPipeLane>,
11719 ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
11720 assert!(k >= 1, "k must be >= 1");
11721 let pipe_setup_walk = match pipe {
11722 Some(p) => Some(p.setup_begin()?),
11723 None => None,
11724 };
11725 // sse-cadence flush cursor: everything in out[..flushed] has been handed to on_commit.
11726 let mut flushed = 0usize;
11727 // admission yield (2026-08-06): on_commit's continue-verdict; false = end the burst
11728 // at the next round boundary (same exit as max_new reached — the session tail runs).
11729 // Initialized by the unconditional post-prime flush below.
11730 let mut keep_going;
11731 let mtp = self
11732 .mtp
11733 .as_ref()
11734 .expect("generate_spec requires an MTP head (nextn_predict_layers>0)");
11735 let n_vocab = self.output.out_features();
11736 // FR-Spec: the draft head may be TRIMMED (fewer rows than n_vocab); the draft argmax runs
11737 // over the draft vocab and the winning index maps through d2t to a TARGET token id.
11738 // Everything downstream (verify/accept/commit) sees target ids only — exactness unchanged.
11739 let d_vocab = mtp
11740 .shared_head_head
11741 .as_ref()
11742 .unwrap_or(&self.output)
11743 .out_features();
11744 if !self.mtp_extra.is_empty() {
11745 if self.plan.draft_source != memra_gguf::model_plan::DraftSourcePlan::Embedded
11746 || self.plan.mtp_blocks.len() != self.mtp_head_count()
11747 {
11748 return Err(
11749 "multi-head MTP requires one embedded canonical block per loaded head".into(),
11750 );
11751 }
11752 // TRIMMED chains (2026-08-27): every head must carry the SAME d2t — the ranking is
11753 // token-frequency and head-independent, and every downstream remap (per-step argmax,
11754 // stream pack, sampled d2t_dev) reads head 0's map, so equality is what makes that
11755 // single map correct for the whole chain. Mixed trimmed/untrimmed is refused.
11756 for (offset, head) in self.mtp_extra.iter().enumerate() {
11757 if head.d2t != mtp.d2t
11758 || head
11759 .shared_head_head
11760 .as_ref()
11761 .unwrap_or(&self.output)
11762 .out_features()
11763 != d_vocab
11764 {
11765 return Err(format!(
11766 "embedded MTP head {} has incompatible draft vocabulary",
11767 offset + 1
11768 )
11769 .into());
11770 }
11771 }
11772 eprintln!(
11773 "[mtp-chain] heads={} policy=step-modulo prefix-replay kv=per-head",
11774 self.mtp_head_count()
11775 );
11776 }
11777 let n_embd = self.cfg.n_embd as usize;
11778 // SESSION MODE: reuse the live cache/scratch, prime only the suffix. `base` = tokens
11779 // already committed (their state is in the caches); 0 = fresh single-shot call.
11780 let session_mode = sess.is_some();
11781 let max_ctx = match sess.as_ref() {
11782 Some(s) => s.cache.max_ctx,
11783 None => prompt.len() + max_new + k + 8,
11784 };
11785 let mut own_cache;
11786 let mut own_scratch;
11787 // PREFIX-CACHE capture request threaded out of the session (lane/spec-prefix-cache):
11788 // (requested split, destination list). Single-shot per burst; fresh calls have none.
11789 let mut sess_capture: Option<(Option<usize>, &mut Vec<SpecBoundaryCapture>)> = None;
11790 // STABLE-BOUNDARY turn-checkpoint request (lane/frspec-multiturn-cache): ABSOLUTE
11791 // committed-length position; consumed one-shot like `capture_at`. None = legacy
11792 // prompt-end capture below.
11793 let mut ckpt_req: Option<usize> = None;
11794 // FAIL-SAFE bit threaded out of the session (see `SpecSession::capture_disabled`).
11795 let mut sess_capture_disabled = false;
11796 let (
11797 cache,
11798 scratch,
11799 mut sess_tail,
11800 mut sess_draft_slot,
11801 mut sess_pending_slot,
11802 sess_ckpt_slot,
11803 sess_telem,
11804 ): (
11805 &mut Cache,
11806 &mut MtpScratch,
11807 Option<(
11808 &mut Vec<u32>,
11809 &mut Option<CudaSlice<f32>>,
11810 &mut Option<u32>,
11811 &mut u32,
11812 &mut u32,
11813 )>,
11814 Option<&mut Option<DraftGraphCtx>>,
11815 Option<&mut Option<u32>>,
11816 Option<&mut Option<SpecCheckpoint>>,
11817 Option<&SpecTelemetryCounters>,
11818 ) = match sess.take() {
11819 Some(sr) => {
11820 let SpecSession {
11821 cache,
11822 scratch,
11823 committed,
11824 last_h,
11825 next_pred,
11826 sctr: s_sctr,
11827 uctr: s_uctr,
11828 draft_ctx,
11829 pending_tok,
11830 turn_ckpt,
11831 telem,
11832 capture_at,
11833 boundary_captures,
11834 ckpt_at,
11835 capture_disabled,
11836 } = sr;
11837 sess_capture_disabled = *capture_disabled;
11838 sess_capture = Some((capture_at.take(), boundary_captures));
11839 ckpt_req = ckpt_at.take();
11840 (
11841 cache,
11842 scratch,
11843 Some((committed, last_h, next_pred, s_sctr, s_uctr)),
11844 Some(draft_ctx),
11845 Some(pending_tok),
11846 Some(turn_ckpt),
11847 Some(telem),
11848 )
11849 }
11850 None => {
11851 // STAGE-OWNED KV (lane/pp2-spec 2026-08-06) — see `new_session`. Door shut =
11852 // `Cache::new` verbatim.
11853 own_cache = crate::pp::new_cache_planned(e, &self.cfg, &self.plan, max_ctx)?;
11854 // Persistent scratch = max_ctx rows (~2KB/token quantized).
11855 own_scratch = self.new_mtp_scratch(e, max_ctx)?;
11856 (
11857 &mut own_cache,
11858 &mut own_scratch,
11859 None,
11860 None,
11861 None,
11862 None,
11863 None,
11864 )
11865 }
11866 };
11867 cache.ensure_usable("generate_spec")?;
11868 if scratch.plane_count() != self.mtp_head_count() {
11869 return Err(format!(
11870 "MTP scratch/head count mismatch ({}/{})",
11871 scratch.plane_count(),
11872 self.mtp_head_count()
11873 )
11874 .into());
11875 }
11876 let base = cache.pos;
11877 // PENDING-CARRY consume (2026-08-01): a carried bonus reaches here only on the
11878 // empty-suffix GREEDY continuation path (generate_spec_session_sampled flushed every
11879 // other case). It enters the round loop as round-0's pending — verify col 0 — exactly
11880 // like a mid-burst full-accept boundary: no init feed, no tail commit pass.
11881 let carried_pending: Option<u32> = sess_pending_slot.as_mut().and_then(|s| s.take());
11882 // PERSISTENT DRAFT KV (the only mode since 2026-07-08 — the legacy round-local scratch,
11883 // MEMRA_SPEC_KVLOCAL, measured -35 acceptance pts on the 27B p3 sweep and was removed;
11884 // acceptance-only — exactness is verify's job either way).
11885 // HIDDEN-PAIRING CONVENTION (DEFAULT = predecessor-row, 2026-07-04 — the 27B acceptance
11886 // unlock, +16pts): the MTP head is TRAINED on rows pairing token x_p with the trunk
11887 // hidden of its PREDECESSOR h_{p-1} (the reference engine's mtp_update shifts the target
11888 // hiddens right by one; its draft step 0 feeds (id_last, TRUE hidden of the row id_last
11889 // was sampled from)). memra's historical convention paired SAME-ROW (x_p, h_p) in the fill
11890 // and seeded chain step 0 through an extra MTP pass on a duplicated token (the
11891 // pseudo-seed) — measured 27B p2 K=3 acceptance 0.569 vs 0.731, p3 0.445 vs 0.63+, and
11892 // the chain steps j>=1 were already predecessor-shaped, so ONLY the fill + step-0 seed
11893 // move. The fill shifts by one and the chain seeds from the predecessor's true hidden
11894 // DIRECTLY (vh_seed / vx[j-1]) — the pseudo pass disappears (one MTP-block pass saved
11895 // per round on top of the acceptance win). Draft-quality-only: exactness stays the
11896 // verify's job either way. (The legacy same-row pairing seam, MEMRA_SPEC_HSAME, and its
11897 // pseudo-seed passes were removed 2026-07-08 — predecessor pairing won by +16 acc pts;
11898 // the legacy round-local scratch, MEMRA_SPEC_KVLOCAL, went with it.)
11899 // REPLAY-FREE PARTIAL ACCEPT (default, 2026-07-03): partial rounds keep the verify's own
11900 // bit-identical committed-prefix state (KV truncate + recur rebuild from the VerifyCkpt)
11901 // and leave the bonus PENDING — no duplicate trunk pass (profiled ~0.54 extra full weight
11902 // reads/round at long ctx). MEMRA_SPEC_REPLAY=1 restores the legacy rollback+replay (A/B
11903 // + fallback seam).
11904 // Qwen35-MoE replay pin LIFTED (lane/draftcost-moe, 2026-08-20). The pin's stated
11905 // bar — the retained verify-state commit proven equivalent to sequential serving —
11906 // was waiting on this arch running the serving batched verify class, which the
11907 // t-parallel admission (this lane, increment 1) provided: the VerifyCkpt the
11908 // replay-free commit consumes is now produced by the SAME serving-class verify that
11909 // qualified dense qwen35 on 2026-08-15 (where the per-round duplicate replay
11910 // measured 69 -> 30 tok/s). Qualification receipts (run-spec K=1..8 both arms,
11911 // 8-prompt replay-vs-replay-free canary, long-prompt cell):
11912 // research/draftcost-moe-20260820/RECEIPTS.md. MEMRA_SPEC_REPLAY=1 stays the
11913 // rollback + A/B seam.
11914 let spec_replay = spec_replay_env_enabled();
11915 if constraint.is_some() && spec_replay {
11916 return Err(
11917 "constrained spec decode does not support MEMRA_SPEC_REPLAY=1 \
11918 (legacy replay commits an unmasked bonus)"
11919 .into(),
11920 );
11921 }
11922 // TRUE-HIDDEN REFRESH (default in persistent-draft-KV mode): every round overwrites the
11923 // committed positions' scratch entries from the verify's exact hiddens (mtp_kv_fill batch)
11924 // instead of keeping chain-approximate entries. MEMRA_SPEC_NOREFRESH=1 = legacy (A/B seam).
11925 let refresh = std::env::var("MEMRA_SPEC_NOREFRESH").is_err();
11926 if !refresh && !self.mtp_extra.is_empty() {
11927 return Err("multi-head MTP requires exact accepted-prefix refresh".into());
11928 }
11929
11930 // prime: BATCHED cache prime (prime_cache — the measured #1 e2e gap: tokenwise primed at
11931 // ~102/38 tok/s vs the engine's ~2000-5900 tok/s batched prefill). prime_cache returns the
11932 // full pre-output_norm hidden stack [T, n_embd], which IS prompt_h (the persistent-draft-KV
11933 // mtp_kv_fill input) — no per-token collection needed. Prompts below PRIME_MIN_T, and
11934 // MEMRA_PRIME_TOKENWISE=1, and frozen Hy3 CPU/GPU expert splits take the tokenwise
11935 // decode_step_h loop. The latter avoids transient GPU staging of the spilled expert bank.
11936 // EMPTY-SUFFIX CONTINUATION (serve bursts): a session turn with NO new tokens resumes
11937 // generation exactly where the last turn stopped — no prime at all. The stashed
11938 // `next_pred` plays prime_logits' role: it is the token produced from the logits after
11939 // committed.last() by the same rule this entry applies to a cold prime's last row —
11940 // an argmax when greedy, a `sample_boundary_token` draw when sampled (the burst tail,
11941 // or `spec_session_from_restored` for a converted prefix-cache hit, did the drawing
11942 // where the sampler and the session's Philox counters were live). `last_h` seeds the
11943 // predecessor pairing below. Fresh calls and non-empty suffixes take the normal path.
11944 let continuation = prompt.is_empty();
11945 if continuation {
11946 assert!(session_mode, "empty prompt requires a session");
11947 assert!(
11948 sess_tail
11949 .as_ref()
11950 .is_some_and(|(c, lh, np, _, _)| !c.is_empty()
11951 && lh.is_some()
11952 && (np.is_some() || carried_pending.is_some())),
11953 "empty-suffix continuation needs a primed session (committed + last_h + next_pred|pending)"
11954 );
11955 }
11956 let mut prime_logits;
11957 let mut prompt_h: Option<CudaSlice<f32>> = None;
11958 let t_prime = std::time::Instant::now();
11959 let batched_prime = !continuation
11960 && prompt.len() >= crate::hybrid_forward::PRIME_MIN_T
11961 && std::env::var("MEMRA_PRIME_TOKENWISE").is_err()
11962 && !e.frozen_cpu_experts_prefer_tokenwise_prime();
11963 let prime_split = prime_split.filter(|&split| split > 0 && split < prompt.len());
11964 if prime_split.is_some() && continuation {
11965 return Err("spec prime split requires a non-empty prime".into());
11966 }
11967 // STABLE-BOUNDARY TURN CHECKPOINT stop (lane/frspec-multiturn-cache, 2026-08-21):
11968 // the worker's `ckpt_at` request, ABSOLUTE -> prompt-relative. On WARM bursts
11969 // (base != 0, an affinity-rewound or pool-resumed session priming its own delta)
11970 // this is the only stop; on COLD bursts it usually coincides with `prime_split`
11971 // (both are the plain tier's stable pre-generation boundary). A boundary the prime
11972 // cannot honor (outside this prime's range) silently drops the capture — the
11973 // turn_ckpt convention: the next turn re-primes in full, never a wrong resume.
11974 let ckpt_rel = if continuation {
11975 None
11976 } else {
11977 ckpt_req
11978 .and_then(|abs| abs.checked_sub(base))
11979 .filter(|&r| r > 0 && r < prompt.len())
11980 };
11981 // Prime stops, ordered: each is a boundary the prime halts at so the in-place GDN
11982 // conv/ssm state can be snapshotted there (the only moment it exists). One stop =
11983 // the legacy single-split program, byte-for-byte.
11984 let mut stops: Vec<usize> = Vec::new();
11985 for b in [prime_split, ckpt_rel].into_iter().flatten() {
11986 if !stops.contains(&b) {
11987 stops.push(b);
11988 }
11989 }
11990 stops.sort_unstable();
11991 // Captured at the ckpt stop, installed into the session slot post-prime (replacing
11992 // the legacy prompt-end capture). Some(None) = capture attempted and failed -> the
11993 // slot is cleared (a stale checkpoint would rewind to the WRONG boundary).
11994 let mut ckpt_early: Option<Option<SpecCheckpoint>> = None;
11995 if continuation {
11996 prime_logits = Vec::new();
11997 } else if !stops.is_empty() {
11998 if let Some(&first) = stops.first()
11999 && prime_split == Some(first)
12000 && first < crate::hybrid_forward::PRIME_MIN_T
12001 {
12002 return Err(format!(
12003 "spec prime split {first} is below PRIME_MIN_T {}",
12004 crate::hybrid_forward::PRIME_MIN_T,
12005 )
12006 .into());
12007 }
12008 // Mirror the plain worker's boundary stops exactly. Each segment is a
12009 // request-level prime (`queued_after` keeps Step35 arm selection independent of
12010 // the stops — tick-seg law); a segment below PRIME_MIN_T (and the final tail
12011 // under MEMRA_PRIME_TOKENWISE) takes the same eager tokenwise continuation as
12012 // prefill_tick. Retain every hidden row so the draft scratch fill remains one
12013 // coherent prompt.
12014 let mut h_all = e.uninit(prompt.len() * n_embd)?;
12015 prime_logits = Vec::new();
12016 let mut prev = 0usize;
12017 for seg_end in stops.iter().copied().chain(std::iter::once(prompt.len())) {
12018 if seg_end <= prev {
12019 continue;
12020 }
12021 let seg = &prompt[prev..seg_end];
12022 let is_final = seg_end == prompt.len();
12023 let batched_seg = seg.len() >= crate::hybrid_forward::PRIME_MIN_T
12024 && (!is_final
12025 || (std::env::var("MEMRA_PRIME_TOKENWISE").is_err()
12026 && !e.frozen_cpu_experts_prefer_tokenwise_prime()));
12027 if batched_seg {
12028 let (l, _, h_seg) =
12029 self.prime_cache(e, seg, &mut *cache, prompt.len() - seg_end)?;
12030 e.copy_into(&mut h_all, prev * n_embd, &h_seg, seg.len() * n_embd)?;
12031 prime_logits = l;
12032 } else {
12033 for (i, &tok) in seg.iter().enumerate() {
12034 let (l, h) = self.decode_step_h(e, tok, &mut *cache)?;
12035 e.copy_into(&mut h_all, (prev + i) * n_embd, &h, n_embd)?;
12036 prime_logits = l;
12037 }
12038 }
12039 prev = seg_end;
12040 if is_final {
12041 break;
12042 }
12043 debug_assert_eq!(cache.pos, base + seg_end, "prime stop landed off boundary");
12044 // PREFIX-CACHE BOUNDARY CAPTURE (lane/spec-prefix-cache): the GDN conv/ssm
12045 // states are about to be advanced in place by the next segment, so this is
12046 // the ONLY moment the boundary's recurrent state exists. Capture iff the
12047 // worker requested exactly this stop (cold sessions only — `capture_at` is
12048 // never armed warm). A failed snapshot is silent (turn_ckpt convention) —
12049 // publication is an optimization, never a correctness dependency.
12050 if base == 0
12051 && let Some((requested, slot)) = sess_capture.as_mut()
12052 {
12053 // Publish at the requested miss-LCP stop (the shared-prefix class)
12054 // AND at the stable-boundary stop (the next-turn re-render class,
12055 // lane/frspec-multiturn-cache) — the same boundary set the plain
12056 // prefill tick learns. Without the second entry, the turn after a
12057 // cold re-park could only hit the OLDER lcp entry (the measured
12058 // one-turn transient: t3 restored 607 of 24122 while the plain arm
12059 // rewound to 15222). Dedupe is the worker sweep's has_key.
12060 if (*requested == Some(seg_end) || ckpt_rel == Some(seg_end))
12061 && let Ok(snap) = cache.snapshot(e)
12062 {
12063 slot.push(SpecBoundaryCapture {
12064 snap,
12065 pos: seg_end,
12066 logits: prime_logits.clone(),
12067 // rows [0..seg_end) of h_all are primed — the following
12068 // segments append, never overwrite.
12069 last_h: capture_boundary_hidden(e, &h_all, seg_end, n_embd),
12070 latent_tails: Vec::new(),
12071 });
12072 }
12073 }
12074 // SESSION-AFFINITY TURN CHECKPOINT at the STABLE boundary (see `ckpt_at`):
12075 // same snapshot mechanics, installed post-prime in place of the prompt-end
12076 // capture the re-render class always diverged below.
12077 if ckpt_rel == Some(seg_end) {
12078 let anchor: Result<CudaSlice<f32>, Box<dyn std::error::Error>> =
12079 e.uninit(n_embd).and_then(|mut a| {
12080 e.copy_view_into(
12081 &mut a,
12082 0,
12083 &h_all.slice((seg_end - 1) * n_embd..seg_end * n_embd),
12084 n_embd,
12085 )?;
12086 Ok(a)
12087 });
12088 ckpt_early = Some(match (cache.snapshot(e), anchor) {
12089 (Ok(snap), Ok(last_h)) => Some(SpecCheckpoint {
12090 snap,
12091 pos: base + seg_end,
12092 last_h,
12093 }),
12094 _ => None,
12095 });
12096 }
12097 }
12098 if std::env::var("MEMRA_SPEC_STATS").as_deref() == Ok("1") {
12099 eprintln!(
12100 "[spec-prime] stops={stops:?} tail={}",
12101 prompt.len() - stops.last().copied().unwrap_or(0)
12102 );
12103 }
12104 prompt_h = Some(h_all);
12105 } else if batched_prime {
12106 let (l, _h_seed, hiddens) = self.prime_cache(e, prompt, &mut *cache, 0)?;
12107 prime_logits = l;
12108 prompt_h = Some(hiddens);
12109 } else {
12110 prime_logits = Vec::new();
12111 prompt_h = Some(e.uninit(prompt.len() * n_embd)?);
12112 for (i, &tok) in prompt.iter().enumerate() {
12113 let (l, h) = self.spec_target_step_h(e, tok, &mut *cache)?;
12114 if let Some(ph) = prompt_h.as_mut() {
12115 e.copy_into(ph, i * n_embd, &h, n_embd)?;
12116 }
12117 prime_logits = l;
12118 }
12119 }
12120 e.stream().synchronize()?;
12121 // PREFIX-CACHE SEED CAPTURE (lane/spec-prefix-cache): boundary == prompt end (the seed
12122 // case — no shared-prefix split, publish the whole prompt). The prime just finished, so
12123 // cache.pos == base + prompt.len() and the recurrent state IS the boundary state;
12124 // prime_logits are the boundary logits. Cold sessions only (base == 0) — same law as
12125 // prime_split. The mid-prompt capture above already consumed the request if it matched.
12126 if !continuation
12127 && base == 0
12128 && let Some((requested, slot)) = sess_capture.as_mut()
12129 && *requested == Some(prompt.len())
12130 && slot.is_empty()
12131 {
12132 debug_assert_eq!(cache.pos, prompt.len(), "seed capture off prompt end");
12133 if let Ok(snap) = cache.snapshot(e) {
12134 slot.push(SpecBoundaryCapture {
12135 snap,
12136 pos: prompt.len(),
12137 logits: prime_logits.clone(),
12138 last_h: prompt_h
12139 .as_ref()
12140 .map(|ph| capture_boundary_hidden(e, ph, prompt.len(), n_embd))
12141 .unwrap_or_default(),
12142 latent_tails: Vec::new(),
12143 });
12144 }
12145 }
12146 // Harness timing contract (see crate::PRIME_NANOS): gen-only throughput without the
12147 // prime-subtraction hack.
12148 crate::PRIME_NANOS.store(
12149 t_prime.elapsed().as_nanos() as u64,
12150 std::sync::atomic::Ordering::Relaxed,
12151 );
12152
12153 let (embd_qt, embd_rb) = self.embd.qt_and_row_bytes(n_embd);
12154 // Resident table is fastest when it fits. Large spill deployments can preserve that HBM
12155 // for expert-cache slots and gather only the exact rows needed by MTP/verify from host.
12156 let host_embd = spec_host_embd();
12157 let embd_gpu = if host_embd {
12158 None
12159 } else {
12160 Some(
12161 self.embd_gpu
12162 .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload")),
12163 )
12164 };
12165 let embd_dev = embd_gpu.map(|g| (g, embd_qt, embd_rb));
12166 if host_embd {
12167 eprintln!(
12168 "[spec] host-row embedding: {} bytes kept off HBM",
12169 self.embd.raw.len()
12170 );
12171 }
12172 let mut out: Vec<u32> = Vec::with_capacity(max_new);
12173 let mut total_drafted = 0usize;
12174 let mut total_accepted = 0usize;
12175
12176 // --- SAMPLER FIRST (lane/sampled-spec-quality, 2026-08-19) ---
12177 // The sampler config, the session's Philox counters and the penalty window are parsed
12178 // HERE, above the boundary-token selection, because the boundary token must be drawn
12179 // from the sampler the request asked for. Pre-lane this block sat ~50 lines BELOW the
12180 // selection, which is the whole mechanical reason the boundary token was an argmax:
12181 // the sampler state was not in scope yet. Nothing here depends on the round loop, so
12182 // moving it up is a pure reordering for greedy (`sampled == false` ⇒ every branch
12183 // below takes the argmax path it always took).
12184 // --- SAMPLED SPEC (MEMRA_SPEC_TEMP>0, research/sampled-spec-impl-map.md): rejection-
12185 // sampling verify (Leviathan/Chen) — accept draft x at u < p(x)/q(x), resample from
12186 // norm(max(0,p-q)) on reject, bonus sampled from p on full accept. Counter-based Philox
12187 // everywhere (seed, event) -> reproducible. temp==0/unset = the greedy path, untouched.
12188 let sp = sampling.unwrap_or_else(|| SpecSampling {
12189 temp: std::env::var("MEMRA_SPEC_TEMP")
12190 .ok()
12191 .and_then(|v| v.parse().ok())
12192 .unwrap_or(0.0),
12193 seed: std::env::var("MEMRA_SEED")
12194 .ok()
12195 .and_then(|v| v.parse().ok())
12196 .unwrap_or(42),
12197 top_k: std::env::var("MEMRA_TOP_K")
12198 .ok()
12199 .and_then(|v| v.parse().ok())
12200 .unwrap_or(0),
12201 top_p: std::env::var("MEMRA_TOP_P")
12202 .ok()
12203 .and_then(|v| v.parse().ok())
12204 .unwrap_or(1.0),
12205 min_p: std::env::var("MEMRA_MIN_P")
12206 .ok()
12207 .and_then(|v| v.parse().ok())
12208 .unwrap_or(0.0),
12209 penalty_last_n: std::env::var("MEMRA_PENALTY_LAST_N")
12210 .ok()
12211 .and_then(|v| v.parse().ok())
12212 .unwrap_or(0),
12213 penalty_repeat: std::env::var("MEMRA_PENALTY_REPEAT")
12214 .ok()
12215 .and_then(|v| v.parse().ok())
12216 .unwrap_or(1.0),
12217 penalty_freq: std::env::var("MEMRA_PENALTY_FREQ")
12218 .ok()
12219 .and_then(|v| v.parse().ok())
12220 .unwrap_or(0.0),
12221 penalty_present: std::env::var("MEMRA_PENALTY_PRESENT")
12222 .ok()
12223 .and_then(|v| v.parse().ok())
12224 .unwrap_or(0.0),
12225 });
12226 let (sp_temp, sp_seed) = (sp.temp, sp.seed);
12227 let sampled = sp_temp > 0.0;
12228 // Counters resume from the session (burst continuity: randomness must never repeat
12229 // across generate_spec_session calls); one-shot callers start at (0,0). Read through
12230 // sess_tail — `sess` was take()n into it above, so sess.as_ref() here is always None.
12231 let mut sctr: u32 = sess_tail.as_ref().map(|(_, _, _, s, _)| **s).unwrap_or(0);
12232 let mut uctr: u32 = sess_tail.as_ref().map(|(_, _, _, _, u)| **u).unwrap_or(0);
12233 // Penalties (v2.1): applied to COPIES of q rows and p columns symmetrically (exactness
12234 // for the penalized+filtered target). History = generated tokens, host-tracked window.
12235 let pen_on = sampled
12236 && sp.penalty_last_n > 0
12237 && (sp.penalty_repeat != 1.0 || sp.penalty_freq != 0.0 || sp.penalty_present != 0.0);
12238 // SESSION-SPANNING PENALTY WINDOW (Item 2). Pre-lane this was
12239 // `prompt.iter().rev().take(64).rev()` — the BURST's suffix slice — so a continuation
12240 // burst (the majority of a stream's tokens, and ALL of a converted cache hit's) started
12241 // with an EMPTY penalty history and the client's repetition/frequency/presence penalties
12242 // silently reset at every burst boundary. The window now spans `committed ++ prompt`,
12243 // which is what the API contract says and what the plain sampler's own `history` does.
12244 // Byte-identical to the pre-lane seed for a cold turn-1 burst at the default window.
12245 let mut pen_hist: Vec<u32> = if pen_on {
12246 let sess_hist: &[u32] = if spec_pen_session_on() {
12247 sess_tail
12248 .as_ref()
12249 .map(|(c, ..)| c.as_slice())
12250 .unwrap_or(&[])
12251 } else {
12252 &[] // MEMRA_SPEC_PEN_SESSION=0: pre-lane burst-local window
12253 };
12254 pen_window_seed(sess_hist, prompt, sp.penalty_last_n)
12255 } else {
12256 Vec::new()
12257 };
12258 // First generated token = the BOUNDARY token: greedy takes the argmax of the prompt's
12259 // last logits (== greedy's first token, byte-contract); SAMPLED draws it from the
12260 // request's own filtered/penalized target through the session's Philox stream
12261 // (`sample_boundary_token`, lane/sampled-spec-quality Item 1 — pre-lane this was an
12262 // argmax in both regimes, so ~1 token per burst of a sampled stream was greedy).
12263 // Emit it, then FEED it to establish the loop invariant below.
12264 // PENDING-CARRY: the carried bonus was already emitted by the LAST burst — it becomes
12265 // last_token WITHOUT re-emission, and round 0 consumes it as pending (no init feed).
12266 // CONSTRAINED entry rules: the first emitted token is the MASKED argmax of the
12267 // prompt's last logits (plain constrained-greedy identity); a continuation without
12268 // a carried pending would emit an UNMASKED stashed next_pred — refused loudly (the
12269 // worker never resumes constrained sessions from the pool, so this cannot fire).
12270 if let Some(c) = constraint.as_deref_mut() {
12271 if continuation && carried_pending.is_none() {
12272 return Err("constrained spec continuation requires a carried pending \
12273 (pool resume is unconstrained-only)"
12274 .into());
12275 }
12276 if !continuation {
12277 c.mask_logits(&mut prime_logits)
12278 .map_err(|e2| format!("constraint: {e2}"))?;
12279 }
12280 }
12281 let mut last_token = if let Some(b) = carried_pending {
12282 b
12283 } else if continuation {
12284 // A continuation's boundary token was DRAWN by the burst that stashed it (the
12285 // session tail below), or by `spec_session_from_restored` for a converted
12286 // prefix-cache hit — in both cases from the correct logits row with this same
12287 // session's Philox stream, which is why it can be consumed here as-is.
12288 sess_tail.as_ref().unwrap().2.unwrap()
12289 } else if sampled && constraint.is_none() && spec_sampled_boundary_on() {
12290 sample_boundary_token(e, &prime_logits, &sp, &pen_hist, &mut sctr, "cold-prime")?
12291 } else {
12292 // greedy (byte contract), the rollback door, or constrained (masked-argmax
12293 // identity — the worker routes sampled+constrained to the plain path, and this
12294 // function refuses the combination outright above).
12295 argmax(&prime_logits) as u32
12296 };
12297 if pen_on {
12298 // The boundary token is a GENERATED token: the plain sampler `accept()`s every
12299 // emitted token into its penalty history, and pre-lane the burst's first token
12300 // was invisible to penalties forever (never pushed, and never in `committed`
12301 // until this burst's tail). Covers the carry/continuation seeds too — neither is
12302 // in `committed` yet.
12303 pen_hist.push(last_token);
12304 }
12305 if carried_pending.is_none() {
12306 out.push(last_token);
12307 // grammar advances with every emitted token (carried pendings were consumed
12308 // by the burst that emitted them).
12309 if let Some(c) = constraint.as_deref_mut() {
12310 c.consume(last_token)
12311 .map_err(|e2| format!("constraint: {e2}"))?;
12312 }
12313 }
12314 if continuation {
12315 // draft-KV invariant: entries [0..base) are the session's exact fills; truncate any
12316 // overhang so the chain's first append lands at slot base (== committed.len()).
12317 scratch.set_len(e, base)?;
12318 }
12319 // sse-cadence: hand the caller every not-yet-flushed token (disjoint in-order slices
12320 // concatenating to the full `out`). Called after the prime's first token and after each
12321 // round commit — emission timing only, token bytes untouched. The slice may be EMPTY
12322 // (poll-only boundary: zero-round folds commit nothing new); returns the caller's
12323 // continue-verdict (admission yield, 2026-08-06) — false ends the burst at this round.
12324 #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
12325 fn flush_commit(
12326 cb: &mut Option<&mut dyn FnMut(&[u32]) -> bool>,
12327 out: &[u32],
12328 flushed: &mut usize,
12329 ) -> bool {
12330 if let Some(f) = cb.as_mut() {
12331 let keep = f(&out[*flushed..]);
12332 *flushed = out.len();
12333 keep
12334 } else {
12335 true
12336 }
12337 }
12338 keep_going = flush_commit(&mut on_commit, &out, &mut flushed);
12339 // INVARIANT at loop top: `last_token` is the most-recently-committed/emitted token, its
12340 // KV+recur state IS in `cache` (cache.pos = position right AFTER last_token), `last_pred`
12341 // is the greedy ARGMAX of the logits that predict the token FOLLOWING last_token, and
12342 // `h_seed` = last_token's pre-output_norm hidden. Establish it by feeding last_token once
12343 // (mirrors plain greedy). DEVICE-ARGMAX lever: the accept walk only ever consumes the
12344 // argmax of those logits — never the full vector — so a host u32 replaces the Vec<f32>.
12345 // Trimmed heads: q lives on the trimmed vocab; accept gathers use the TRIMMED index and
12346 // the residual scatters q into target-id space (q=-inf off-trim — the head cannot propose
12347 // those, so their residual mass is p(x), correct by construction).
12348 let d2t_dev: Option<CudaSlice<u32>> = if sampled || crate::spec::spec_stream() {
12349 match &mtp.d2t {
12350 Some(map) => Some(e.htod_u32_v(map)?),
12351 None => None,
12352 }
12353 } else {
12354 None
12355 };
12356 let mut q_full_buf: Option<CudaSlice<f32>> = None;
12357 // host Philox4x32-10 accept-test uniforms: module fn `host_u01` (shared with the
12358 // dspark sampled-admission walk); byte-identical to the closure it replaces.
12359 let mut draft_logits: Vec<CudaSlice<f32>> = Vec::new(); // retained head logits (q), per slot
12360 let mut draft_stats: Vec<(f32, f32, f32)> = Vec::new(); // (row_max, th_e, z_e) per slot
12361 let mut perturb_buf: Option<CudaSlice<f32>> = None; // gumbel scratch (max(n_vocab,d_vocab))
12362 let mut sample_tok = e.alloc_u32_zeroed(1)?; // residual/bonus sample out
12363 let mut col_buf: Option<CudaSlice<f32>> = None; // materialized verify column
12364 let mut pen_hist_d: Option<CudaSlice<u32>> = None;
12365 let mut pcol_buf: Option<CudaSlice<f32>> = None; // penalized p-column scratch
12366 // MEMRA_SPEC_SETUP_TRACE=1 (diagnostics): per-call wall decomposition of the burst
12367 // SETUP + TAIL segments (the round loop's internals are MEMRA_SPEC_PHASE's job) —
12368 // built to pin the serve per-burst fixed cost (research/spec-serving-20260801).
12369 let setup_trace = std::env::var("MEMRA_SPEC_SETUP_TRACE").as_deref() == Ok("1");
12370 let t_ent = std::time::Instant::now();
12371
12372 // SESSION-AFFINITY TURN CHECKPOINT (lane/session-affinity, 2026-08-05): capture the
12373 // PROMPT-END boundary state so a LATER turn can rewind here and re-prime only its own
12374 // delta instead of the whole conversation. See `SpecCheckpoint` for why this boundary is
12375 // the one that matters (a history-rewriting client mutates what the session GENERATED,
12376 // so the next turn's prompt agrees with this one up to exactly here).
12377 //
12378 // WHERE — AND WHY THIS EXACT LINE. Right after the trunk prime, BEFORE the init feed
12379 // (`decode_step_h(last_token)`) and before round 0: the last instant at which the caches
12380 // hold exactly `base + prompt.len()` rows and nothing generated.
12381 //
12382 // This was WRONG in the first cut of this lane: the capture sat after the draft-KV fill,
12383 // which is also after the init feed, so `cache.pos` was `base + prompt.len() + 1` — the
12384 // boundary included the FIRST GENERATED TOKEN. That token is the first thing inside the
12385 // `<think>` block the client strips, so every later turn's diff diverged exactly one
12386 // token below the checkpoint and affinity declined 100% of the time. Measured on the
12387 // owner regime: "history diverged at 12233 of checkpoint 12234". The off-by-one made the
12388 // whole mechanism inert while looking, from the outside, like a working
12389 // correctness-declines-safely path — hence the decline log carries the offsets.
12390 //
12391 // The full-attn planes are `len`-truncatable so the snapshot copies only the GDN conv/ssm
12392 // state (the reason a spec session could not rewind before). The draft scratch needs no
12393 // copy: rows below the boundary are rewritten by the next turn's own fill.
12394 //
12395 // WHEN: non-empty prime only. An empty-suffix continuation burst adds no prompt boundary
12396 // (its "prompt end" IS the previous checkpoint's, already held), so it keeps the existing
12397 // checkpoint rather than replacing it with a strictly worse one.
12398 //
12399 // FAILURE IS SILENT BY DESIGN: on a VRAM-tight rig the snapshot alloc can fail. That
12400 // costs the NEXT turn its rewind (it re-primes fully, today's behavior) and must never
12401 // fail the burst that is already running — so the error is swallowed, loud only under
12402 // MEMRA_DEBUG_SPEC.
12403 //
12404 // STABLE-BOUNDARY OVERRIDE (lane/frspec-multiturn-cache, 2026-08-21): the prompt-end
12405 // posture above was DISPROVED for the think-posture template class — the prompt's own
12406 // tail is the live generation header (`<|im_start|>assistant\n<think>\n`) that the
12407 // next turn's re-render replaces, so the diff diverged a couple tokens BELOW the
12408 // checkpoint and affinity declined 100% of multi-turn agent traffic (the same class
12409 // the plain tier fixed on 2026-08-09 via `plain_checkpoint_boundary`; the port to the
12410 // spec tier is this lane). When the worker armed `ckpt_at`, the capture happened at
12411 // that stop inside the prime above (`ckpt_early`) and is installed here instead;
12412 // capture-attempted-but-failed clears the slot exactly like the legacy arm.
12413 if let Some(slot) = sess_ckpt_slot {
12414 if let Some(early) = ckpt_early {
12415 if early.is_none() && std::env::var("MEMRA_DEBUG_SPEC").is_ok() {
12416 eprintln!(
12417 "[spec] stable-boundary turn checkpoint skipped; \
12418 next turn re-primes in full"
12419 );
12420 }
12421 *slot = early;
12422 } else if !continuation {
12423 let pos = cache.pos;
12424 debug_assert_eq!(
12425 pos,
12426 base + prompt.len(),
12427 "turn checkpoint must sit at the prompt end, before the init feed"
12428 );
12429 let anchor: Result<CudaSlice<f32>, Box<dyn std::error::Error>> =
12430 if let Some(ph) = &prompt_h {
12431 // hidden of the LAST primed row = the predecessor anchor at this
12432 // boundary (exactly what a fresh prime of committed[..pos] leaves in
12433 // last_h, and what the next prime's fill reads for its first row).
12434 let np = prompt.len();
12435 e.uninit(n_embd).and_then(|mut a| {
12436 e.copy_view_into(
12437 &mut a,
12438 0,
12439 &ph.slice((np - 1) * n_embd..np * n_embd),
12440 n_embd,
12441 )?;
12442 Ok(a)
12443 })
12444 } else {
12445 Err("no prompt hiddens".into())
12446 };
12447 match (cache.snapshot(e), anchor) {
12448 (Ok(snap), Ok(last_h)) => {
12449 *slot = Some(SpecCheckpoint { snap, pos, last_h });
12450 }
12451 (s, a) => {
12452 *slot = None; // a stale checkpoint would rewind to the WRONG boundary
12453 if std::env::var("MEMRA_DEBUG_SPEC").is_ok() {
12454 let err = s
12455 .err()
12456 .map(|e| e.to_string())
12457 .or_else(|| a.err().map(|e| e.to_string()))
12458 .unwrap_or_default();
12459 eprintln!(
12460 "[spec] turn checkpoint skipped ({err}); \
12461 next turn re-primes in full"
12462 );
12463 }
12464 }
12465 }
12466 }
12467 }
12468 // INIT FEED — skipped on a pending carry: last_token (the carried bonus) is NOT in the
12469 // caches and must NOT be fed solo; round 0's batched verify commits it as col 0. Its
12470 // seed/anchor hidden is the carried last_h (copied below); last_pred is dead in the
12471 // pending path (t_pred reads verify col 0 — the accept walk overwrites it).
12472 let mut last_pred = 0u32;
12473 let mut last_col_logits: Option<CudaSlice<f32>> = None;
12474 // CONSTRAINED: the init feed's logits back the (n_acc==0, base==0) masked-argmax
12475 // recompute in the grammar-truncation walk — retained host-side, round 0 only.
12476 let mut init_logits_host: Option<Vec<f32>> = None;
12477 let h_seed0: CudaSlice<f32> = if carried_pending.is_none() {
12478 let (init_logits, h) = self.spec_target_step_h(e, last_token, &mut *cache)?;
12479 last_pred = argmax(&init_logits) as u32;
12480 if constraint.is_some() {
12481 init_logits_host = Some(init_logits.clone());
12482 }
12483 // sampled mode: p-distribution after last_token, for the j==0/base==0 accept test.
12484 if sampled {
12485 last_col_logits = Some(e.htod(&init_logits)?);
12486 }
12487 h
12488 } else {
12489 // predecessor-row anchor: hidden of the last COMMITTED row (the carry contract).
12490 let lh = sess_tail
12491 .as_ref()
12492 .unwrap()
12493 .1
12494 .as_ref()
12495 .expect("pending carry requires last_h");
12496 e.clone_dtod(lh)?
12497 };
12498 let t_init = t_ent.elapsed();
12499 let mut last_col_stats: Option<(f32, f32, f32)> = None;
12500 // PERSISTENT h_seed buffer (allocated BEFORE any graph capture so no captured scratch can
12501 // alias it): every path that updates the round seed copies INTO it — no per-round allocs,
12502 // stable pointer for the graph-draft round-start copy.
12503 let mut h_seed_buf = e.clone_dtod(&h_seed0)?;
12504 // Predecessor-pairing trackers: `fill_prev` = trunk hidden AT the last COMMITTED row (the
12505 // predecessor of the next verify's col 0 — the reference's carried pending-h analogue;
12506 // also the predecessor-row hidden for the round-0 legacy-replay seed). At round 0 that
12507 // row is last_token's own (h_seed0). The chain step-0 seed under the pairing default =
12508 // hidden of the row BEFORE last_token = the prompt's last row at round 0 (h_seed_buf
12509 // overwritten below).
12510 let mut fill_prev = e.clone_dtod(&h_seed0)?;
12511 {
12512 if let Some(ph) = &prompt_h {
12513 let np = prompt.len();
12514 e.copy_view_into(
12515 &mut h_seed_buf,
12516 0,
12517 &ph.slice((np - 1) * n_embd..np * n_embd),
12518 n_embd,
12519 )?;
12520 } else if continuation
12521 && let Some((_, lh, _, _, _)) = sess_tail.as_ref()
12522 && let Some(lh) = lh.as_ref()
12523 {
12524 e.copy_into(&mut h_seed_buf, 0, lh, n_embd)?;
12525 }
12526 }
12527 // Persistent device prediction slots for the accept walk (max k+1 verify columns).
12528 let mut preds_d = e.alloc_u32_zeroed(k + 2)?;
12529
12530 let debug_spec = std::env::var("MEMRA_DEBUG_SPEC").is_ok();
12531 let fork_mode = OptiForkGateMode::configured();
12532 // MEMRA_SPEC_STATS=1: per-slot accept histogram + draft-length histogram, printed once at
12533 // the end. Metric normalization vs the reference engine: BOTH engines count
12534 // accepted/drafted where the chain stopped at p-min and the sub-threshold token is
12535 // discarded uncounted — per-slot decay + chain-length mix are the extra dimensions.
12536 let spec_stats = std::env::var("MEMRA_SPEC_STATS").is_ok();
12537 let mut st_drafted = vec![0usize; k];
12538 let mut st_accepted = vec![0usize; k];
12539 let mut st_len_hist = vec![0usize; k + 1];
12540 let mut st_full = 0usize;
12541 // P-MIN CONFIDENCE GATE (MEMRA_SPEC_PMIN, the serve script's --spec-draft-p-min mechanism):
12542 // stop the draft chain early when the head's softmax confidence in its own pick drops
12543 // below p_min. Hoisted above the loop: the graph capture bakes the prob kernels iff on.
12544 static PMIN: std::sync::OnceLock<f32> = std::sync::OnceLock::new();
12545 let p_min = *PMIN.get_or_init(|| {
12546 std::env::var("MEMRA_SPEC_PMIN")
12547 .ok()
12548 .and_then(|v| v.parse().ok())
12549 .unwrap_or(0.0)
12550 });
12551 // ZERO-DRAFT ROUNDS (MEMRA_SPEC_PMIN0=1, vendored from llama.cpp's draft gating): let the
12552 // p-min gate apply at j==0 too, so a low-confidence round drafts NOTHING and the verify
12553 // batch is just the pending bonus (m=1 = a plain decode step). llama's 35B win rides
12554 // exactly this — draft acceptance 76% at mean len 2.5 because unpredictable stretches
12555 // never pay draft+verify overhead. Only legal when a pending bonus exists (an empty
12556 // verify batch is not); the j==0 exemption stays for pending-less rounds.
12557 let pmin0 = std::env::var("MEMRA_SPEC_PMIN0")
12558 .map(|v| v == "1")
12559 .unwrap_or(false);
12560
12561 // --- GRAPH DRAFT setup: persistent I/O buffers + ONE capture (2 warmups inside). The
12562 // warmups mutate scratch len_d / pos / tok / seed — all reset at every round start, so the
12563 // only restore needed is the scratch counter. Capture failure (e.g. a non-capturable
12564 // cuBLAS path in an exotic head) falls back to the eager draft chain.
12565 // PER-SESSION PERSISTENCE (2026-08-01): session calls reuse the DraftGraphCtx parked on
12566 // the SpecSession — the capture (2 warmup head forwards + instantiate) ran ONCE at the
12567 // session's first burst, not per burst (measured ~16ms/burst fixed cost on H100 q27,
12568 // research/spec-serving-20260801). Reuse is pointer-exact: the graph bakes the session's
12569 // own scratch KV (never realloc'd), the model's resident embedding, the OnceLock p_min,
12570 // and the g_* buffers carried in the ctx — replay dispatch is identical to a fresh
12571 // capture, so draft tokens are bit-identical (drafts never decide exactness anyway; the
12572 // verify arbitrates). Single-shot calls (sess=None) build a fresh ctx and drop it.
12573 let mut dctx: DraftGraphCtx = match sess_draft_slot.as_mut().and_then(|s| s.take()) {
12574 Some(c) => c,
12575 None => DraftGraphCtx::new(e, n_embd, if sampled { d_vocab } else { 1 })?,
12576 };
12577 // FAIL-SAFE (step-OOM park replay): pre-mark both fallback flags so no capture arm
12578 // below can fire — LOUD once per replayed session through the standard WARN line.
12579 if sess_capture_disabled {
12580 let reason =
12581 "session replayed after a step-OOM park; draft capture disabled (fail-safe)";
12582 let flip = dctx.failed.mark_greedy(reason);
12583 let flip_s = dctx.failed.mark_sampled(reason);
12584 if let Some(line) = flip.or(flip_s) {
12585 eprintln!("{line}");
12586 }
12587 }
12588 // A session that ran greedy bursts first sized g_q/g_perturb at 1; a sampled resume
12589 // needs d_vocab. Realloc is legal exactly while graph_s is None (nothing baked them).
12590 if sampled && dctx.g_q.len() < d_vocab {
12591 dctx.g_q = e.zeros(d_vocab)?;
12592 dctx.g_perturb = e.zeros(d_vocab)?;
12593 }
12594 // DRAFT-SIDE GRAMMAR MASK (lane/draft-mask, 2026-08-04): the drafter samples the
12595 // grammar's legal set, so proposals are legal BY CONSTRUCTION and the verify-side
12596 // truncation (the correctness backstop) stops cutting every tight-schema round.
12597 // The mask is one node inside the captured draft chain — presence is a CAPTURE-TIME
12598 // shape, so a parked graph of the other shape is dropped and recaptured.
12599 let dmask_on = constraint
12600 .as_deref()
12601 .is_some_and(|c| c.draft_mask_enabled());
12602 let dmask_words = if dmask_on { d_vocab.div_ceil(32) } else { 0 };
12603 if dmask_on && dctx.g_dmask.len() < dmask_words {
12604 dctx.g_dmask = e.alloc_u32_zeroed(dmask_words)?;
12605 dctx.graph = None; // the old capture baked the old (or no) mask pointer
12606 dctx.chain = None; // chain last-row graphs bake the same pointer
12607 dctx.failed.clear_greedy();
12608 dctx.keeper.clear();
12609 }
12610 if (dctx.graph.is_some() || dctx.chain.is_some()) && dctx.graph_masked != dmask_on {
12611 dctx.graph = None;
12612 dctx.chain = None;
12613 dctx.failed.clear_greedy();
12614 dctx.keeper.clear();
12615 }
12616 // MULTI-HEAD CHAIN mode (mtp_extra non-empty — step37's 3-head shipping shape): the
12617 // step-modulo prefix-replay chain captures PER-HEAD single-row graphs
12618 // (`DraftChainGraphs`) instead of the one self-feeding graph below; the single-head
12619 // capture arms are untouched and unreachable in this mode (the launch arms branch the
12620 // same way). This removes the historical `mtp_extra.is_empty()` capture exclusion —
12621 // and with it the silent no-attempt hole: a chain capture that FAILS now trips the
12622 // same LOUD draft-graph WARN as a single-head failure.
12623 let chain_mode = !self.mtp_extra.is_empty();
12624 // ---- PRE-CAPTURE VRAM RESERVE CHECK + PER-SESSION DRAFT-STATE MEASUREMENT ----
12625 // (lane/step37-vram-admission-20260830). `cap_eff0` opens the measurement bracket:
12626 // when any capture succeeds in THIS call, the effective-free delta across the whole
12627 // capture section is recorded as the model's per-session draft-state high-water
12628 // (admission charges it per spec-capable session — this state was charged at ZERO
12629 // before the lane). The reserve check runs BEFORE any capture arm can allocate: a
12630 // refused capture trips the same LOUD once-per-flip WARN class as a failed one, but
12631 // with the card's headroom still intact (the owner's single-session OOM was a capture
12632 // attempt walking the card to the edge and stranding the eager fallback at 5 MiB free).
12633 let cap_eff0 = e
12634 .ctx()
12635 .mem_get_info()
12636 .ok()
12637 .map(|(f, _)| f.saturating_add(e.pool_cached_bytes()));
12638 // Peak instrument for the same bracket: the CAPTURE-TIME peak (warmup transients +
12639 // instantiate scratch, alive together) dwarfs the parked delta — measured on the
12640 // owner shape: a capture whose PARKED state reads ~2.6GB walked a ~7GB-free card to
12641 // OOM mid-capture. Reset the pool watermark here; read it at bracket end.
12642 let _ = e.pool_high_water_reset();
12643 let cap_used0 = e.pool_reserved_used().1;
12644 let mut captured_now = false;
12645 let mut capture_oom_entry_eff: Option<usize> = None;
12646 let capture_need = {
12647 let observed = self.draft_session_admission_bytes();
12648 if observed > 0 {
12649 observed
12650 } else {
12651 draft_capture_bootstrap_estimate(
12652 if chain_mode { self.mtp_head_count() } else { 1 },
12653 k,
12654 d_vocab,
12655 n_embd,
12656 )
12657 }
12658 };
12659 if spec_capture_gate_on()
12660 && graph_draft
12661 && !sampled
12662 && !dctx.failed.greedy_failed()
12663 && ((chain_mode && dctx.chain.is_none() && mtp_chain_graph_on())
12664 || (!chain_mode && dctx.graph.is_none()))
12665 && let Some(reason) = capture_headroom_refusal(e, capture_need)
12666 && let Some(line) = dctx.failed.mark_greedy(&reason)
12667 {
12668 eprintln!("{line}");
12669 }
12670 if graph_draft
12671 && !sampled
12672 && chain_mode
12673 && dctx.chain.is_none()
12674 && !dctx.failed.greedy_failed()
12675 {
12676 if mtp_chain_graph_on() {
12677 let heads_n = self.mtp_head_count();
12678 let DraftGraphCtx {
12679 g_tok,
12680 g_pos,
12681 g_seed,
12682 g_p,
12683 g_dmask,
12684 ..
12685 } = &mut dctx;
12686 if dmask_on {
12687 e.htod_u32_into(g_dmask, &vec![u32::MAX; dmask_words])?;
12688 }
12689 let g_dmask_ro: &CudaSlice<u32> = &*g_dmask;
12690 let with_prob = p_min > 0.0;
12691 // CAPTURE-RETAIN (#68 fix): one keeper for the whole chain — every graph's
12692 // warmup transients stay pinned as long as any of them replays.
12693 let cap_res = (|| -> Result<DraftChainGraphs, Box<dyn std::error::Error>> {
12694 // dcw door: same warmup headroom pre-arm as the single-head capture
12695 // below — every plane, because each head's capture warmups append on
12696 // its OWN plane. INSIDE the fallible closure (vram-admission lane): an
12697 // OOM here used to `?` out of the whole burst as a step error; now it
12698 // is a capture failure — LOUD WARN, eager chain serves.
12699 if step35_draft_dcw_on() {
12700 scratch.ensure_dcw_headroom(e, k + 2)?;
12701 }
12702 let mut interior = Vec::with_capacity(heads_n);
12703 let mut last = Vec::with_capacity(heads_n);
12704 let mut keeper: Vec<Box<dyn std::any::Any + Send>> = Vec::new();
12705 for hi in 0..heads_n {
12706 let head = self.mtp_head_at(hi);
12707 // interior row: KV append + carrier only (`with_head=false` — the
12708 // eager chain discards interior logits too, so this is the same
12709 // consumed-byte program minus the dead full-vocab head matmul).
12710 let (g, keep) = e.capture_graph_retained(|e| {
12711 self.mtp_head_forward_cap(
12712 e,
12713 head,
12714 g_tok,
12715 g_pos,
12716 g_seed,
12717 g_p,
12718 &mut *scratch,
12719 hi,
12720 false,
12721 false,
12722 embd_gpu.expect("graph draft requires resident embedding"),
12723 embd_qt,
12724 embd_rb,
12725 d_vocab,
12726 None,
12727 None,
12728 None,
12729 )
12730 })?;
12731 // the warmups appended rows on plane hi; rewind before the next
12732 // capture so successive warmups never outrun the pre-armed headroom.
12733 scratch.set_plane_len(e, hi, base)?;
12734 interior.push(g);
12735 keeper.extend(keep);
12736 // last row: head matmul + greedy argmax tail (+ p when the policy
12737 // reads it, + the grammar-mask node when constrained).
12738 let (g2, keep2) = e.capture_graph_retained(|e| {
12739 self.mtp_head_forward_cap(
12740 e,
12741 head,
12742 g_tok,
12743 g_pos,
12744 g_seed,
12745 g_p,
12746 &mut *scratch,
12747 hi,
12748 with_prob,
12749 true,
12750 embd_gpu.expect("graph draft requires resident embedding"),
12751 embd_qt,
12752 embd_rb,
12753 d_vocab,
12754 None,
12755 None,
12756 if dmask_on {
12757 Some((g_dmask_ro, dmask_words))
12758 } else {
12759 None
12760 },
12761 )
12762 })?;
12763 scratch.set_plane_len(e, hi, base)?;
12764 last.push(g2);
12765 keeper.extend(keep2);
12766 }
12767 Ok(DraftChainGraphs {
12768 interior,
12769 last,
12770 _keeper: keeper,
12771 })
12772 })();
12773 match cap_res {
12774 Ok(cg) => {
12775 scratch.set_len(e, base)?;
12776 // POSITIVE engagement receipt (the 3a lesson: a WARN-free boot is
12777 // NOT evidence of capture — the captured state must name itself).
12778 eprintln!(
12779 "[mtp-chain-graph] captured mode=greedy heads={heads_n} \
12780 interior={heads_n} last={heads_n} masked={}",
12781 dmask_on as u8
12782 );
12783 dctx.chain = Some(cg);
12784 dctx.graph_masked = dmask_on;
12785 captured_now = true;
12786 }
12787 Err(err) => {
12788 scratch.set_len(e, base)?;
12789 // LOUD flip (audit Q2): a dropped draft graph is a coverage loss,
12790 // never silent — now including the multi-head shipping shape.
12791 // OOM RECOVERY (vram-admission lane): a failed attempt's freed
12792 // transients sit CACHED in the async pool where the driver cannot
12793 // see them; trim them back so the eager fallback (and any driver-
12794 // side allocation) actually has the headroom the free suggests.
12795 let mut reason = err.to_string();
12796 if capture_err_is_oom(&reason) {
12797 capture_oom_entry_eff = capture_oom_entry_eff.max(cap_eff0);
12798 let trimmed = e.pool_trim_to_zero();
12799 if trimmed > 0 {
12800 reason.push_str(&format!(
12801 "; pool trimmed {}MB back to the driver",
12802 trimmed / (1 << 20)
12803 ));
12804 }
12805 }
12806 if let Some(line) = dctx.failed.mark_greedy(&reason) {
12807 eprintln!("{line}");
12808 }
12809 }
12810 }
12811 } else {
12812 // Disarmed by MEMRA_MTP_CHAIN_GRAPH=0: say so once per process — the OFF arm
12813 // must be attributable in a boot log, never inferable from silence.
12814 static NOTE: std::sync::Once = std::sync::Once::new();
12815 NOTE.call_once(|| {
12816 eprintln!(
12817 "[spec] multi-head draft-chain capture disarmed \
12818 (MEMRA_MTP_CHAIN_GRAPH=0); eager chain serves this shape"
12819 );
12820 });
12821 }
12822 }
12823 if graph_draft
12824 && !sampled
12825 && !chain_mode
12826 && dctx.graph.is_none()
12827 && !dctx.failed.greedy_failed()
12828 {
12829 let DraftGraphCtx {
12830 g_tok,
12831 g_pos,
12832 g_seed,
12833 g_p,
12834 g_dmask,
12835 ..
12836 } = &mut dctx;
12837 // capture-time contents: ALL-ONES (ban nothing). A replay only ever runs after the
12838 // host uploads the position's real words, so the warmups stay grammar-free.
12839 if dmask_on {
12840 e.htod_u32_into(g_dmask, &vec![u32::MAX; dmask_words])?;
12841 }
12842 let g_dmask_ro: &CudaSlice<u32> = &*g_dmask;
12843 // CAPTURE-RETAIN (#68 fix): the warmup transients' pool addresses are baked into the
12844 // captured graph; the keeper pins them for the graph's lifetime. capture_graph (non-
12845 // retained) freed them at exit — safe for one-shot generate_spec (nothing else touches
12846 // the pool between replays) but WRONG for sessions: burst-boundary prime/fill/commit
12847 // passes (and, in serve, other sessions) recycle those addresses and the replay then
12848 // clobbers live buffers — the ST serve-spec corruption (research/serve-st-20260803).
12849 let cap_res = (|| {
12850 // dcw door: the capture warmups append device-counter rows the capture body
12851 // cannot rebase for; pre-arm ring headroom host-side (no-op on flat planes /
12852 // room-enough rings, and the door-off path is untouched). INSIDE the fallible
12853 // closure (vram-admission lane): an OOM here is a capture failure, not a
12854 // burst-killing step error.
12855 if step35_draft_dcw_on() {
12856 scratch.ensure_dcw_headroom(e, k + 2)?;
12857 }
12858 e.capture_graph_retained(|e| {
12859 self.mtp_head_forward_cap(
12860 e,
12861 mtp,
12862 g_tok,
12863 g_pos,
12864 g_seed,
12865 g_p,
12866 &mut *scratch,
12867 0,
12868 p_min > 0.0 || fork_mode == OptiForkGateMode::Controller,
12869 true,
12870 embd_gpu.expect("graph draft requires resident embedding"),
12871 embd_qt,
12872 embd_rb,
12873 d_vocab,
12874 None,
12875 None,
12876 if dmask_on {
12877 Some((g_dmask_ro, dmask_words))
12878 } else {
12879 None
12880 },
12881 )
12882 })
12883 })();
12884 match cap_res {
12885 Ok((g, keep)) => {
12886 scratch.set_len(e, base)?;
12887 dctx.graph = Some(g);
12888 dctx.graph_masked = dmask_on;
12889 dctx.keeper = keep;
12890 captured_now = true;
12891 }
12892 Err(err) => {
12893 scratch.set_len(e, base)?;
12894 // LOUD flip (audit Q2): a dropped draft graph is a coverage loss, never
12895 // silent. Once per flip — mark returns None on an already-failed ctx.
12896 let mut reason = err.to_string();
12897 if capture_err_is_oom(&reason) {
12898 capture_oom_entry_eff = capture_oom_entry_eff.max(cap_eff0);
12899 let trimmed = e.pool_trim_to_zero();
12900 if trimmed > 0 {
12901 reason.push_str(&format!(
12902 "; pool trimmed {}MB back to the driver",
12903 trimmed / (1 << 20)
12904 ));
12905 }
12906 }
12907 if let Some(line) = dctx.failed.mark_greedy(&reason) {
12908 eprintln!("{line}");
12909 }
12910 }
12911 }
12912 }
12913 // --- SAMPLED GRAPH DRAFT setup (step 3 of the sampled-spec arc): a SECOND capture, own
12914 // graph object, built only when sampled && graph-eligible — the greedy capture above is
12915 // untouched (and skipped when sampled: its graph would never be launched). Same head
12916 // forward, but the in-graph argmax reads GUMBEL-PERTURBED logits; the Philox event
12917 // counter lives in the persistent device g_ctr (bumped in-graph, host-seeded from sctr
12918 // once per round); the raw head logits land in the persistent g_q for the host's
12919 // per-replay async D2D into the round's q slot (q_slots, K x d_vocab, allocated once).
12920 // seed/temp are capture-time constants — baked into graph_s, so a pool-resumed request
12921 // with a different (seed, temp, k) drops the parked sampled graph and recaptures.
12922 // COST OF THE FRESH-SEED SERVE DEFAULT (dogfood F4, 2026-08-04): omitting `seed` on a
12923 // serve request now draws fresh per-request entropy (it used to default to a pinned 0),
12924 // so a seed-omitting request that RESUMES a parked spec session finds an s_key baked
12925 // with the PREVIOUS request's seed and pays one recapture. Bounded, and it does not
12926 // reopen the ~16ms/burst regression the persistent ctx exists to fix: a session's seed
12927 // is fixed for its whole lifetime (worker.rs reads s.sampler.seed() per burst), so
12928 // this compare misses at most ONCE per resumed request — the first burst recaptures
12929 // and every later burst in that request replays. A client that wants the parked graph
12930 // AND reproducibility supplies an explicit `seed`, honored exactly, which keeps s_key
12931 // stable across its whole conversation.
12932 // COMPOSITION RULE (fspec x gsd merge): the in-graph chain samples from the RAW
12933 // softmax — it can hold neither per-row filter stats nor the varying penalty history.
12934 // The sampled graph therefore engages only in the PURE-TEMP regime; filters/penalties
12935 // force the eager draft (which computes stats/penalties per row).
12936 // KEY THE WHOLE REGIME, not just the baked constants (lane/graph-s-key-exactness-
12937 // 20260819). `s_key` used to be `(seed, temp, k)`; the filters and penalties were left
12938 // out, so a filtered request resuming a session that parked a PURE-TEMP graph kept it —
12939 // and the launch site never re-asked `pure_temp`. See [`SampledGraphKey`] for what that
12940 // costs (an unconditional accept of out-of-head draft tokens, i.e. an exactness bug on
12941 // the request shape the vendor-default flip makes the majority).
12942 let s_key = SampledGraphKey::new(sp_seed, sp_temp, k, sp.top_k, sp.top_p, sp.min_p, pen_on);
12943 let pure_temp = s_key.pure_temp();
12944 // The regime the sampled graph may be captured/launched in: pure-temp always;
12945 // truncation-filtered when the filtered-capture door is on (the filter runs
12946 // IN-GRAPH — lane/step37-draft-graph-serving-20260830); penalties never.
12947 let s_capturable = s_key.graph_capturable();
12948 if sampled && dctx.s_key.is_some_and(|old| old != s_key) {
12949 dctx.graph_s = None;
12950 dctx.chain_s = None;
12951 dctx.failed.clear_sampled();
12952 dctx.s_key = None;
12953 dctx.q_slots.clear();
12954 dctx.keeper_s.clear();
12955 }
12956 // PRE-CAPTURE VRAM RESERVE CHECK, sampled arms (vram-admission lane): same contract
12957 // as the greedy check above — refuse BEFORE allocating, LOUD once, eager serves.
12958 if spec_capture_gate_on()
12959 && graph_draft
12960 && sampled
12961 && s_capturable
12962 && !dctx.failed.sampled_failed()
12963 && ((chain_mode && dctx.chain_s.is_none() && mtp_chain_graph_on())
12964 || (!chain_mode && dctx.graph_s.is_none()))
12965 && let Some(reason) = capture_headroom_refusal(e, capture_need)
12966 && let Some(line) = dctx.failed.mark_sampled(&reason)
12967 {
12968 eprintln!("{line}");
12969 }
12970 // FILTERED capture nodes need q slots sized d_vocab AND the stat slots; the pure-temp
12971 // body leaves g_th/g_z/g_mx untouched (they exist from ctx creation either way).
12972 if graph_draft
12973 && sampled
12974 && s_capturable
12975 && chain_mode
12976 && dctx.chain_s.is_none()
12977 && !dctx.failed.sampled_failed()
12978 {
12979 if mtp_chain_graph_on() {
12980 let heads_n = self.mtp_head_count();
12981 let filtered = s_key.filtered();
12982 let DraftGraphCtx {
12983 g_tok,
12984 g_pos,
12985 g_seed,
12986 g_p,
12987 g_ctr,
12988 g_perturb,
12989 g_q,
12990 g_rows0,
12991 g_th,
12992 g_z,
12993 g_mx,
12994 ..
12995 } = &mut dctx;
12996 let with_prob = p_min > 0.0;
12997 let cap_res = (|| -> Result<DraftChainGraphs, Box<dyn std::error::Error>> {
12998 // dcw pre-arm INSIDE the fallible closure (vram-admission lane): an OOM
12999 // here is a capture failure with the LOUD WARN, never a step error.
13000 if step35_draft_dcw_on() {
13001 scratch.ensure_dcw_headroom(e, k + 2)?;
13002 }
13003 let mut interior = Vec::with_capacity(heads_n);
13004 let mut last = Vec::with_capacity(heads_n);
13005 let mut keeper: Vec<Box<dyn std::any::Any + Send>> = Vec::new();
13006 for hi in 0..heads_n {
13007 let head = self.mtp_head_at(hi);
13008 // interior row: no head, no draw — shared shape with the greedy
13009 // chain's interior, captured per mode for keeper-lifetime hygiene.
13010 let (g, keep) = e.capture_graph_retained(|e| {
13011 self.mtp_head_forward_cap(
13012 e,
13013 head,
13014 g_tok,
13015 g_pos,
13016 g_seed,
13017 g_p,
13018 &mut *scratch,
13019 hi,
13020 false,
13021 false,
13022 embd_gpu.expect("graph draft requires resident embedding"),
13023 embd_qt,
13024 embd_rb,
13025 d_vocab,
13026 None,
13027 None,
13028 None,
13029 )
13030 })?;
13031 scratch.set_plane_len(e, hi, base)?;
13032 interior.push(g);
13033 keeper.extend(keep);
13034 // last row: head matmul + the in-graph categorical draw (filtered
13035 // nodes when the request carries filters).
13036 let (g2, keep2) = e.capture_graph_retained(|e| {
13037 self.mtp_head_forward_cap(
13038 e,
13039 head,
13040 g_tok,
13041 g_pos,
13042 g_seed,
13043 g_p,
13044 &mut *scratch,
13045 hi,
13046 with_prob,
13047 true,
13048 embd_gpu.expect("graph draft requires resident embedding"),
13049 embd_qt,
13050 embd_rb,
13051 d_vocab,
13052 Some(SampledCapArgs {
13053 ctr: &mut *g_ctr,
13054 perturb: &mut *g_perturb,
13055 q_out: &mut *g_q,
13056 seed: sp_seed,
13057 temp: sp_temp,
13058 filt: if filtered {
13059 Some(SampledCapFilter {
13060 rows0: &*g_rows0,
13061 th: &mut *g_th,
13062 z: &mut *g_z,
13063 mx: &mut *g_mx,
13064 top_k: sp.top_k,
13065 top_p: sp.top_p,
13066 min_p: sp.min_p,
13067 })
13068 } else {
13069 None
13070 },
13071 }),
13072 None,
13073 None, // constrained spec is greedy-only
13074 )
13075 })?;
13076 scratch.set_plane_len(e, hi, base)?;
13077 last.push(g2);
13078 keeper.extend(keep2);
13079 }
13080 Ok(DraftChainGraphs {
13081 interior,
13082 last,
13083 _keeper: keeper,
13084 })
13085 })();
13086 match cap_res {
13087 Ok(cg) => {
13088 scratch.set_len(e, base)?;
13089 // NO STRANDED PARTIAL STATE (vram-admission lane): the q-slot allocs
13090 // after a successful capture are themselves fallible on a tight card.
13091 // A mid-loop failure used to `?` out as a step error, leaving orphan
13092 // slots parked on the ctx (wrong count, stale contents) for the next
13093 // capture attempt to stack onto. Allocate all-or-nothing: on failure
13094 // drop the fresh graphs AND the partial slots, mark the LOUD fallback.
13095 dctx.q_slots.clear();
13096 let slots = (0..k)
13097 .map(|_| e.zeros(d_vocab))
13098 .collect::<Result<Vec<_>, _>>();
13099 match slots {
13100 Ok(slots) => {
13101 dctx.q_slots = slots;
13102 eprintln!(
13103 "[mtp-chain-graph] captured mode=sampled heads={heads_n} \
13104 interior={heads_n} last={heads_n} filtered={} key={s_key:?}",
13105 s_key.filtered() as u8
13106 );
13107 dctx.chain_s = Some(cg);
13108 dctx.s_key = Some(s_key);
13109 captured_now = true;
13110 }
13111 Err(err) => {
13112 drop(cg);
13113 dctx.q_slots.clear();
13114 let mut reason = format!("q-slot alloc failed: {err}");
13115 if capture_err_is_oom(&reason) {
13116 capture_oom_entry_eff = capture_oom_entry_eff.max(cap_eff0);
13117 let trimmed = e.pool_trim_to_zero();
13118 if trimmed > 0 {
13119 reason.push_str(&format!(
13120 "; pool trimmed {}MB back to the driver",
13121 trimmed / (1 << 20)
13122 ));
13123 }
13124 }
13125 if let Some(line) = dctx.failed.mark_sampled(&reason) {
13126 eprintln!("{line}");
13127 }
13128 }
13129 }
13130 }
13131 Err(err) => {
13132 scratch.set_len(e, base)?;
13133 let mut reason = err.to_string();
13134 if capture_err_is_oom(&reason) {
13135 capture_oom_entry_eff = capture_oom_entry_eff.max(cap_eff0);
13136 let trimmed = e.pool_trim_to_zero();
13137 if trimmed > 0 {
13138 reason.push_str(&format!(
13139 "; pool trimmed {}MB back to the driver",
13140 trimmed / (1 << 20)
13141 ));
13142 }
13143 }
13144 if let Some(line) = dctx.failed.mark_sampled(&reason) {
13145 eprintln!("{line}");
13146 }
13147 }
13148 }
13149 } else {
13150 static NOTE_S: std::sync::Once = std::sync::Once::new();
13151 NOTE_S.call_once(|| {
13152 eprintln!(
13153 "[spec] multi-head draft-chain capture disarmed \
13154 (MEMRA_MTP_CHAIN_GRAPH=0); eager chain serves this shape"
13155 );
13156 });
13157 }
13158 }
13159 if graph_draft
13160 && sampled
13161 && s_capturable
13162 && !chain_mode
13163 && dctx.graph_s.is_none()
13164 && !dctx.failed.sampled_failed()
13165 {
13166 let filtered = s_key.filtered();
13167 let DraftGraphCtx {
13168 g_tok,
13169 g_pos,
13170 g_seed,
13171 g_p,
13172 g_ctr,
13173 g_perturb,
13174 g_q,
13175 g_rows0,
13176 g_th,
13177 g_z,
13178 g_mx,
13179 ..
13180 } = &mut dctx;
13181 // CAPTURE-RETAIN (#68 fix): same keeper contract as the greedy capture above.
13182 let cap_res = (|| {
13183 // dcw pre-arm INSIDE the fallible closure (vram-admission lane): an OOM
13184 // here is a capture failure with the LOUD WARN, never a step error.
13185 if step35_draft_dcw_on() {
13186 scratch.ensure_dcw_headroom(e, k + 2)?;
13187 }
13188 e.capture_graph_retained(|e| {
13189 self.mtp_head_forward_cap(
13190 e,
13191 mtp,
13192 g_tok,
13193 g_pos,
13194 g_seed,
13195 g_p,
13196 &mut *scratch,
13197 0,
13198 p_min > 0.0,
13199 true,
13200 embd_gpu.expect("graph draft requires resident embedding"),
13201 embd_qt,
13202 embd_rb,
13203 d_vocab,
13204 Some(SampledCapArgs {
13205 ctr: &mut *g_ctr,
13206 perturb: &mut *g_perturb,
13207 q_out: &mut *g_q,
13208 seed: sp_seed,
13209 temp: sp_temp,
13210 filt: if filtered {
13211 Some(SampledCapFilter {
13212 rows0: &*g_rows0,
13213 th: &mut *g_th,
13214 z: &mut *g_z,
13215 mx: &mut *g_mx,
13216 top_k: sp.top_k,
13217 top_p: sp.top_p,
13218 min_p: sp.min_p,
13219 })
13220 } else {
13221 None
13222 },
13223 }),
13224 None,
13225 None, // constrained spec is greedy-only — sampled never carries a hook
13226 )
13227 })
13228 })();
13229 match cap_res {
13230 Ok((g, keep)) => {
13231 scratch.set_len(e, base)?;
13232 // NO STRANDED PARTIAL STATE: all-or-nothing q slots, same contract as
13233 // the chain arm above.
13234 dctx.q_slots.clear();
13235 let slots = (0..k)
13236 .map(|_| e.zeros(d_vocab))
13237 .collect::<Result<Vec<_>, _>>();
13238 match slots {
13239 Ok(slots) => {
13240 dctx.q_slots = slots;
13241 dctx.graph_s = Some(g);
13242 dctx.s_key = Some(s_key);
13243 dctx.keeper_s = keep;
13244 captured_now = true;
13245 }
13246 Err(err) => {
13247 drop(g);
13248 drop(keep);
13249 dctx.q_slots.clear();
13250 let mut reason = format!("q-slot alloc failed: {err}");
13251 if capture_err_is_oom(&reason) {
13252 capture_oom_entry_eff = capture_oom_entry_eff.max(cap_eff0);
13253 let trimmed = e.pool_trim_to_zero();
13254 if trimmed > 0 {
13255 reason.push_str(&format!(
13256 "; pool trimmed {}MB back to the driver",
13257 trimmed / (1 << 20)
13258 ));
13259 }
13260 }
13261 if let Some(line) = dctx.failed.mark_sampled(&reason) {
13262 eprintln!("{line}");
13263 }
13264 }
13265 }
13266 }
13267 Err(err) => {
13268 scratch.set_len(e, base)?;
13269 // LOUD flip (audit Q2): same contract as the greedy capture above.
13270 let mut reason = err.to_string();
13271 if capture_err_is_oom(&reason) {
13272 capture_oom_entry_eff = capture_oom_entry_eff.max(cap_eff0);
13273 let trimmed = e.pool_trim_to_zero();
13274 if trimmed > 0 {
13275 reason.push_str(&format!(
13276 "; pool trimmed {}MB back to the driver",
13277 trimmed / (1 << 20)
13278 ));
13279 }
13280 }
13281 if let Some(line) = dctx.failed.mark_sampled(&reason) {
13282 eprintln!("{line}");
13283 }
13284 }
13285 }
13286 }
13287 // ---- PER-SESSION DRAFT-STATE MEASUREMENT bracket end (vram-admission lane): when a
13288 // capture landed in THIS call, the effective-free delta across the capture section is
13289 // this session's parked draft-graph state (keepers + q slots + instantiated graphs'
13290 // backing). Recorded as a model-owned high-water; admission charges it per
13291 // spec-capable session (see `draft_session_admission_bytes`).
13292 if captured_now
13293 && let Some(eff0) = cap_eff0
13294 && let Ok((f1, _)) = e.ctx().mem_get_info()
13295 {
13296 let eff1 = f1.saturating_add(e.pool_cached_bytes());
13297 let parked_delta = eff0.saturating_sub(eff1);
13298 let (_res_high, used_high) = e.pool_high_water_reset();
13299 let peak_delta = used_high.saturating_sub(cap_used0);
13300 let observed = parked_delta.max(peak_delta);
13301 if observed > 0
13302 && let Some(hw) = self.record_draft_state_bytes(observed)
13303 {
13304 eprintln!(
13305 "[spec] draft-session state high-water: {}MB (max of parked delta {}MB \
13306 and capture-time pool peak {}MB; charged per spec admission and gating \
13307 future captures)",
13308 hw / (1 << 20),
13309 parked_delta / (1 << 20),
13310 peak_delta / (1 << 20),
13311 );
13312 }
13313 }
13314 // FAILURE IS AN OBSERVATION TOO: a capture that OOM'd at entry-effective E proved
13315 // the capture-time peak exceeds E. Feed E into the gauge so every future gate
13316 // refuses at or below the headroom that just failed (self-healing even when the
13317 // boot probe is disarmed and the bootstrap estimate was blind).
13318 if let Some(entry_eff) = capture_oom_entry_eff
13319 && let Some(hw) = self.record_draft_state_bytes(entry_eff)
13320 {
13321 eprintln!(
13322 "[spec] draft-session capture appetite floor raised to {}MB: a capture \
13323 attempt OOM'd with that much effective free (failure-observed bound)",
13324 hw / (1 << 20)
13325 );
13326 }
13327 // ---- EXACTNESS GUARD, the enforceable half (lane/graph-s-key-exactness-20260819,
13328 // widened by lane/step37-draft-graph-serving-20260830) ----
13329 // With the filters and penalties in `s_key`, a graph that SURVIVED the drop above was
13330 // captured under THIS request's exact regime, and capture requires `graph_capturable`
13331 // (pure-temp, or filtered with the in-graph filter nodes; never penalties) — so a
13332 // parked graph implies both. That implication is the whole exactness argument for the
13333 // graph arm, so it is asserted here rather than assumed: a future change that widens
13334 // the capture condition, narrows the key, or copies a `DraftGraphCtx` across regimes
13335 // fails LOUDLY at this line instead of silently drafting from a distribution the
13336 // verify never reconstructs. Release builds refuse the graph (drop it, draft eager)
13337 // rather than launching it; the launch site re-tests the regime independently.
13338 if sampled
13339 && (dctx.graph_s.is_some() || dctx.chain_s.is_some())
13340 && (!s_capturable || dctx.s_key != Some(s_key))
13341 {
13342 debug_assert!(
13343 false,
13344 "sampled draft graph parked under {:?} survived into a request outside its \
13345 capture regime (top_k={} top_p={} min_p={} pen_on={} capturable={}): the \
13346 in-graph draw and the verify's accept test would see different distributions",
13347 dctx.s_key, sp.top_k, sp.top_p, sp.min_p, pen_on, s_capturable,
13348 );
13349 eprintln!(
13350 "[spec] BUG: dropping a parked sampled draft graph that outlived its capture \
13351 regime (s_key={:?}, request top_k={} top_p={} min_p={} pen_on={} \
13352 capturable={}); drafting EAGER — the key must carry every field that shapes q",
13353 dctx.s_key, sp.top_k, sp.top_p, sp.min_p, pen_on, s_capturable,
13354 );
13355 dctx.graph_s = None;
13356 dctx.chain_s = None;
13357 dctx.s_key = None;
13358 dctx.q_slots.clear();
13359 dctx.keeper_s.clear();
13360 }
13361 // SKEY PROBE (MEMRA_SKEY_PROBE=1): the burst-entry facts the reachability question turns
13362 // on — is this request sampled, is it in a regime the sampled graph is legal in, and is
13363 // a graph PARKED from an earlier request of the same session? The launch arms below
13364 // print which chain actually ran, so the probe never restates the condition.
13365 if skey_probe() {
13366 eprintln!(
13367 "[skey] burst sampled={} pure_temp={} capturable={} temp={} top_k={} top_p={} \
13368 min_p={} pen_on={} k={} graph_draft={} graph_s_parked={} chain_s_parked={} \
13369 s_key_parked={:?}",
13370 sampled as u8,
13371 pure_temp as u8,
13372 s_capturable as u8,
13373 sp_temp,
13374 sp.top_k,
13375 sp.top_p,
13376 sp.min_p,
13377 pen_on as u8,
13378 k,
13379 graph_draft as u8,
13380 dctx.graph_s.is_some() as u8,
13381 dctx.chain_s.is_some() as u8,
13382 dctx.s_key,
13383 );
13384 }
13385 let t_cap = t_ent.elapsed();
13386 // PERSISTENT DRAFT KV: fill the MTP block's K/V for every prompt position from the exact
13387 // trunk hiddens collected during prime — ONE batched K/V-only pass (overwrites any
13388 // capture-warmup garbage; capture left len at 0). last_token (the init feed) needs no
13389 // fill: the first chain step processes it and appends its entry at slot prompt.len().
13390 if let Some(ph) = &prompt_h {
13391 // SESSION: rows [0..base) are the previous turns' exact fills (refresh overwrote them
13392 // with true verify hiddens) — truncate any draft overhang, fill ONLY the suffix at
13393 // global positions [base..base+tp). Fresh call: base==0, identical to before.
13394 scratch.set_len(e, base)?;
13395 // CHUNKED FILL (long-ctx OOM fix, 2026-07-05): mtp_kv_fill's transients scale with its
13396 // T (concat = T*2*n_embd*4B — 1.5GB at 40k) and its concat loop is 2*T launches. The
13397 // fill is a pure sequential append, so chunking is exact: each chunk appends its rows
13398 // at pos0=base+start with the identical per-row math. Same knob as the trunk prime.
13399 let tp = prompt.len();
13400 let fill_chunk: usize = if crate::cache::swa_ring_on() {
13401 crate::hybrid_forward::prime_chunk_tokens(tp, self.layers.len())
13402 } else {
13403 // Preserve the flag-OFF schedule byte-for-byte, including the legacy zero value
13404 // meaning one monolithic fill.
13405 std::env::var("MEMRA_PRIME_CHUNK")
13406 .ok()
13407 .and_then(|v| v.parse().ok())
13408 .unwrap_or(4096)
13409 };
13410 let fill_chunk = if fill_chunk == 0 { tp } else { fill_chunk };
13411 // CUDA launch wall (same class as the trunk prime's PRIME_CHUNK_LAUNCH_CAP):
13412 // a fill call's matmuls can land on the grid.y=m dp4a family, and grid.y caps
13413 // at 65,535. This loop has no tail fold, so the raw limit is exact:
13414 // tp <= 65,535 keeps the legacy schedule (monolithic included) byte-for-byte,
13415 // and larger fills — unreachable before the trunk prime's own cap fix — chunk.
13416 let fill_chunk = fill_chunk.min(crate::hybrid_forward::CUDA_GRID_YZ_MAX);
13417 let mut start = 0usize;
13418 while start < tp {
13419 let end = (start + fill_chunk).min(tp);
13420 let tc = end - start;
13421 {
13422 // PREDECESSOR pairing: row i gets h[i-1]; global row 0 a zeros row (the
13423 // reference engine's initial pending-h is zeroed too); a session turn's row 0
13424 // gets the PREVIOUS turn's last committed hidden (sess.last_h). Per chunk:
13425 // rows start..end read h[start-1..end-1] — one dtod into a chunk buffer.
13426 let mut phs = e.zeros(tc * n_embd)?;
13427 let (src_lo, dst_off) = if start == 0 {
13428 (0, n_embd)
13429 } else {
13430 ((start - 1) * n_embd, 0)
13431 };
13432 let n_copy = if start == 0 {
13433 (tc - 1) * n_embd
13434 } else {
13435 tc * n_embd
13436 };
13437 if start == 0
13438 && let Some((_, lh, _, _, _)) = sess_tail.as_ref()
13439 && let Some(lh) = lh.as_ref()
13440 {
13441 e.copy_into(&mut phs, 0, lh, n_embd)?;
13442 }
13443 if n_copy > 0 {
13444 e.copy_view_into(
13445 &mut phs,
13446 dst_off,
13447 &ph.slice(src_lo..src_lo + n_copy),
13448 n_copy,
13449 )?;
13450 }
13451 self.mtp_kv_fill_all(
13452 e,
13453 &prompt[start..end],
13454 &phs,
13455 base + start,
13456 &mut *scratch,
13457 embd_dev,
13458 )?;
13459 }
13460 start = end;
13461 }
13462 }
13463 // MEMRA_PROFILE_SPEC=2: profiler capture starts HERE — after the prime, so an
13464 // `nsys -c cudaProfilerApi` capture contains ONLY the round loop (draft/verify/commit).
13465 // (=1 brackets the whole call in run_spec.rs, prime included.)
13466 if std::env::var("MEMRA_PROFILE_SPEC").as_deref() == Ok("2") {
13467 unsafe extern "C" {
13468 fn cudaProfilerStart() -> i32;
13469 }
13470 unsafe {
13471 cudaProfilerStart();
13472 }
13473 }
13474 // ROUND-STREAM stage (c) 4 (MEMRA_SPEC_STREAM=1, experimental): pre-issued M-round
13475 // bursts with ZERO per-round host readbacks — the accept/seed/rollback/ring kernels
13476 // consume each other's device outputs; the host drains the ring every M rounds. v1
13477 // constraints: greedy, !spec_replay, single-shot, batched-linear layers, no refresh
13478 // fills (acceptance effect A/B-arbitrated), enters from round 1 (pending guaranteed).
13479 // NOTE: not gated on the caller's graph_draft (its trunk_dense conjunct turns the 35B
13480 // MoE off) — the stream capture encloses ONLY the dense MTP head; the head-dense /
13481 // full-prec / k gates are re-derived here and a failed capture degrades to stream-off.
13482 let stream_on = crate::spec::spec_stream()
13483 && !sampled
13484 && !spec_replay
13485 && self.mtp_extra.is_empty()
13486 && constraint.is_none()
13487 && !session_mode
13488 && embd_gpu.is_some()
13489 && !crate::model::full_prec_enabled()
13490 && k + 2 < 96;
13491 let mut stream_graph: Option<cudarc::driver::CudaGraph> = None;
13492 let mut g_tokp2k = e.alloc_u32_zeroed(2 * k.max(1))?;
13493 if stream_on {
13494 let cap = e.capture_graph(|e| {
13495 for j in 0..k.max(1) {
13496 self.mtp_head_forward_cap(
13497 e,
13498 mtp,
13499 &mut dctx.g_tok,
13500 &mut dctx.g_pos,
13501 &mut dctx.g_seed,
13502 &mut dctx.g_p,
13503 &mut *scratch,
13504 0,
13505 true,
13506 true,
13507 embd_gpu.expect("round stream requires resident embedding"),
13508 embd_qt,
13509 embd_rb,
13510 d_vocab,
13511 None,
13512 Some((&mut g_tokp2k, j, d2t_dev.as_ref())),
13513 None, // round-stream requires constraint.is_none() (see stream_on)
13514 )?;
13515 }
13516 Ok(())
13517 });
13518 match cap {
13519 Ok(g) => {
13520 scratch.set_len(e, 0)?;
13521 stream_graph = Some(g);
13522 }
13523 Err(err) => {
13524 scratch.set_len(e, 0)?;
13525 if debug_spec {
13526 eprintln!("[spec] stream-graph capture failed ({err}); stream off");
13527 }
13528 }
13529 }
13530 }
13531 let stream_active = stream_on && stream_graph.is_some();
13532 if debug_spec {
13533 eprintln!(
13534 "[spec] stream_on={stream_on} env={} samp={sampled} dg={} captured={} active={stream_active} session={session_mode} replay={spec_replay}",
13535 crate::spec::spec_stream(),
13536 dctx.graph.is_some(),
13537 stream_graph.is_some()
13538 );
13539 }
13540 let t_v_s = k + 1;
13541 // ROUND-STREAM buffers + ptr tables now live in the model-generic round_stream
13542 // module (extracted 2026-07-12; the gemma burst reuses them).
13543 let sb = crate::round_stream::StreamBufs::new(e, k, crate::spec::spec_stream_m())?;
13544 let crate::round_stream::StreamBufs {
13545 mut vtok_d,
13546 mut brk_d,
13547 mut pend_d,
13548 last_pred_d,
13549 mut pos_ctr,
13550 mut pos_start_d,
13551 mut ring_d,
13552 acc_d: mut stream_acc,
13553 m_rounds,
13554 k: _,
13555 } = sb;
13556 let stream_ptrs: Option<CudaSlice<u64>> = if stream_active {
13557 Some(crate::round_stream::kv_len_ptr_table(
13558 e,
13559 cache,
13560 Some(&pos_ctr),
13561 )?)
13562 } else {
13563 None
13564 };
13565
13566 let t_fill = t_ent.elapsed();
13567 let mut round = 0usize;
13568 // ADAPTIVE DRAFT LENGTH (MEMRA_SPEC_ADAPT=1, opt-in — the gemma_spec accepted-run law,
13569 // ported 2026-08-01): next round's draft depth = last round's accepted run + 1, clamped
13570 // to [floor(pos), k_cap] — a miss shrinks the next draft to the miss point + 1,
13571 // full-accept streaks re-deepen one step per round. NOT the 2026-07-07 acceptance-EMA
13572 // (that arm measured an HONEST LOSS to static per-class optima — 115.0/85.8/73.4 vs
13573 // 121.6/92.7/75.6, EMA lag — and was removed 2026-07-08; rig5090.jsonl has the record).
13574 // The gemma law has no lag class: it reacts within one round, and was worth +7-20% on
13575 // the gemma cells at unchanged exactness (2026-07-10 flip; floor sweep 2026-07-25;
13576 // position key 2026-07-26). Signal = n_acc from the round's EXISTING accept readback —
13577 // zero new syncs; the draft graph is a SINGLE-STEP capture replayed per drafted token,
13578 // so a per-round depth needs no re-capture (unlike gemma's whole-chain graphs). qwen's
13579 // in-round p-min cut already shortens chains mid-round, so gemma's one-round-late p-min
13580 // fold into kc is unnecessary here — the accepted-run law sees the cut via n_acc.
13581 // Exactness is the verify's job at ANY depth (same contract as p-min variable rounds).
13582 // DEFAULT OFF on the qwen path until its cells gate a flip (gemma's is default-on).
13583 // MEASURED 2026-08-01 (H100 GPU-3, interleaved x3, NGEN=256, same-invocation plain
13584 // denominators; research/qwen-adaptive-k-20260801/): REFUTED on the tuned qwen configs.
13585 // q27 K=3+HPOST+PMIN=0.3: short +0.8% (noise; law ~idles, len_hist identical), board
13586 // -1.9%, agentic -0.5%; board PMIN=0 -2.1% (not p-min shadowing — the law itself);
13587 // floor=1 -2.8% (gemma's floor-collapse, reproduced). q35 K=2 board: -6.4% (52/136
13588 // rounds shrink to depth 1; no depth to reclaim at K=2). The gemma direction DOES
13589 // appear at untuned depth-K — q27 K=6 floor=4 +1.5% over fixed K=6 — but stays -3.7%
13590 // below fixed K=3: same verdict class as the retired EMA arm (honest loss to static
13591 // per-class optima). Acceptance-rate rises under the law while tokens/round falls —
13592 // it buys accept-% by adding rounds, and a round's fixed draft+verify cost wins.
13593 // K=1..8 self-consistency PASS both models with the law ON (exactness held).
13594 let adapt = std::env::var("MEMRA_SPEC_ADAPT").as_deref() == Ok("1");
13595 // floor: per-model default keyed on n_embd (gemma's tiering — models with an expensive
13596 // verify keep deep drafts after a miss); MEMRA_SPEC_ADAPT_FLOOR pins it everywhere.
13597 let adapt_floor_env: Option<usize> = std::env::var("MEMRA_SPEC_ADAPT_FLOOR")
13598 .ok()
13599 .and_then(|v| v.parse().ok());
13600 let adapt_floor_default: usize = if self.cfg.n_embd as usize >= 3500 {
13601 4
13602 } else if self.cfg.n_embd as usize >= 2500 {
13603 2
13604 } else {
13605 1
13606 };
13607 let adapt_floor: usize = adapt_floor_env.unwrap_or(adapt_floor_default);
13608 // position key: past floor_ctx a HIGH floor (>=4) relaxes to 1 — forced-deep drafts
13609 // turn net-negative at depth (gemma 31B d1736 evidence); MEMRA_SPEC_FLOOR_CTX moves
13610 // the boundary, an explicit MEMRA_SPEC_ADAPT_FLOOR pins the floor everywhere.
13611 let floor_ctx: usize = std::env::var("MEMRA_SPEC_FLOOR_CTX")
13612 .ok()
13613 .and_then(|v| v.parse().ok())
13614 .unwrap_or(1024);
13615 let floor_at = |pos: usize| -> usize {
13616 if adapt_floor_env.is_some() || pos < floor_ctx {
13617 adapt_floor
13618 } else if adapt_floor >= 4 {
13619 1
13620 } else {
13621 adapt_floor
13622 }
13623 };
13624 // cap: MEMRA_SPEC_CAPMAX (gemma semantics, default 7). Binds only under adapt — the
13625 // fixed-K default path is untouched by this whole block.
13626 let cap_max: usize = std::env::var("MEMRA_SPEC_CAPMAX")
13627 .ok()
13628 .and_then(|v| v.parse().ok())
13629 .unwrap_or(7);
13630 let k_cap = k.min(cap_max).max(1);
13631 let mut kc = k_cap;
13632 let mut opti_fork: Option<OptiForkState> = None;
13633 let mut _opti_walk: Option<crate::pp::PpWalkLease> = None;
13634 let mut _opti_walk_borrow: Option<crate::pp::PpWalkBorrowGuard> = None;
13635 let mut fork_snapshot: Option<crate::cache::CacheSnapshot> = None;
13636 if fork_mode != OptiForkGateMode::Disabled {
13637 let fence = crate::pp::pp_cuts(self.layers.len());
13638 let refusal = if !session_mode {
13639 Some("not-session")
13640 } else if k != 1 || adapt {
13641 Some("requires-fixed-k1")
13642 } else if sampled || constraint.is_some() || spec_replay {
13643 Some("sampled-constrained-or-replay")
13644 } else if pipe.is_some() {
13645 Some("two-session-pipeline")
13646 } else if !spec_devacc() {
13647 Some("requires-device-accept")
13648 } else if stream_active || crate::spec::spec_stream() {
13649 Some("round-stream")
13650 } else if !self.mtp_extra.is_empty() {
13651 Some("multi-head-mtp")
13652 } else if crate::cache::swa_ring_on() || cache.has_swa_ring() {
13653 Some("swa-ring")
13654 } else if crate::pp::pp_host_bounce_active() {
13655 Some("host-bounce")
13656 } else if fork_mode == OptiForkGateMode::Controller
13657 && cache.recur.iter().any(Option::is_some)
13658 {
13659 Some("controller-requires-zero-recurrent-state")
13660 } else if fence.as_ref().is_none_or(|f| f.len() != 3) {
13661 Some("requires-pp2")
13662 } else {
13663 None
13664 };
13665 if let Some(reason) = refusal {
13666 OPTI_FORK_REFUSALS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
13667 eprintln!("[opti-fork] refused reason={reason}");
13668 } else {
13669 let fence = fence.expect("validated PP-2 fence");
13670 let rt = crate::pp::PpNRt::get(e)?;
13671 let primary_stage0 = rt.engine(0, e).ctx().ordinal() == e.ctx().ordinal();
13672 let primary_stage1 = rt.engine(1, e).ctx().ordinal() == e.ctx().ordinal();
13673 let primary_supported =
13674 primary_stage0 || (fork_mode == OptiForkGateMode::Controller && primary_stage1);
13675 if !rt.cross_device() || !primary_supported {
13676 OPTI_FORK_REFUSALS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
13677 eprintln!("[opti-fork] refused reason=requires-supported-primary-cross-device");
13678 } else {
13679 // The optimistic controller can keep two boundary tickets in flight. Give
13680 // every nested verify an explicit borrow of one whole-walk generation; no
13681 // `pp_pipe` boolean is allowed to bypass ownership on its own.
13682 let walk = rt.acquire_walk("opti_fork_coordinator")?;
13683 let permit = rt.walk_permit(&walk, "opti_fork_coordinator")?;
13684 let borrow = rt.borrow_walk(&permit, "opti_fork_coordinator")?;
13685 // Both recurrent snapshots and both seed generations are allocated before
13686 // the first fork, each through its owning PP stage. Allocation failure
13687 // therefore happens before any optimistic state mutation can occur.
13688 let current_snapshot = opti_snapshot_stage_owned(e, cache, rt, &fence)?;
13689 let alternate_snapshot = opti_snapshot_stage_owned(e, cache, rt, &fence)?;
13690 let fork = OptiForkState::new(
13691 e,
13692 cache,
13693 fork_mode,
13694 alternate_snapshot,
13695 &h_seed_buf,
13696 &fill_prev,
13697 rt,
13698 fence[1],
13699 self.layers.len(),
13700 )?;
13701 eprintln!(
13702 "[opti-fork] armed mode={fork_mode:?} snapshots=2 seeds=2 split={} \
13703 payload_dev0={} payload_dev1={} q_threshold={:.3}",
13704 fence[1],
13705 fork.logical_payload_bytes[0],
13706 fork.logical_payload_bytes[1],
13707 fork.controller.map_or(0.0, |policy| policy.threshold),
13708 );
13709 fork_snapshot = Some(current_snapshot);
13710 opti_fork = Some(fork);
13711 _opti_walk = Some(walk);
13712 _opti_walk_borrow = Some(borrow);
13713 }
13714 }
13715 }
13716 // Persistent snapshot buffers are allocated once and refreshed in place. The fork arm
13717 // uses stage-owned snapshots; refused/disabled arms retain the existing generic helper.
13718 let mut snap = match fork_snapshot {
13719 Some(snapshot) => snapshot,
13720 None => cache.snapshot(e)?,
13721 };
13722 let mut carried_opti: Option<OptiControllerTicket> = None;
13723 // ROUND-STREAM stage (b) 3a: device table of per-layer kvl.len_d pointers (stable — the
13724 // cache never reallocates len_d; see cache.rs "stable pointer" note). 0 = no KV layer.
13725 let kv_len_ptrs: Option<CudaSlice<u64>> = if spec_devacc() && !spec_replay {
13726 Some(crate::round_stream::kv_len_ptr_table(e, cache, None)?)
13727 } else {
13728 None
13729 };
13730 // BONUS FOLD (2026-07-04): after a FULL accept the bonus token is NOT committed with a
13731 // separate T=1 trunk pass (a full weight read per round). It stays PENDING and rides as
13732 // column 0 of the NEXT round's verify batch. Under predecessor pairing the next chain
13733 // seeds from the bonus's predecessor's TRUE verify hidden (free — no extra
13734 // pass of any kind). Verify still
13735 // checks every emitted token against the target -> exactness holds by construction; only
13736 // DRAFT QUALITY can shift, which the acceptance numbers arbitrate.
13737 // bonus emitted but not yet committed to cache. A carried pending (see SpecSession::
13738 // pending_tok) enters round 0 directly — the burst boundary becomes a plain round edge.
13739 let mut pending: Option<u32> = carried_pending;
13740 // MEMRA_SPEC_PHASE=1: per-round wall decomposition (draft / verify / accept+commit) —
13741 // no tracing, no extra syncs (each phase is naturally sync-bounded: draft readbacks,
13742 // the verify accept readback). Printed once at loop end via spec-stats.
13743 let anatomy_on = std::env::var("MEMRA_SPEC_PP_ANATOMY").as_deref() == Ok("1");
13744 let phase_on = anatomy_on || std::env::var("MEMRA_SPEC_PHASE").as_deref() == Ok("1");
13745 // MEMRA_SPEC_PHASE_SYNC=1 — reads the phase split correctly, and proves it. `ph_mark` is a
13746 // bare Instant, so `verify-issue` is the host QUEUEING the walk (the GPU is already running
13747 // under it) and `verify-wait` is only the residual drain at the accept readback: one
13748 // overlapped interval cut at the first blocking call, NOT "GPU time" beside "host time".
13749 // Syncing right after the walk is issued moves the whole GPU wall into `verify-issue`. If
13750 // the walk's GPU total is really issue+wait, then with this on verify-issue jumps to that
13751 // sum, verify-wait collapses to the readback alone, and the ROUND WALL DOES NOT MOVE —
13752 // which is what says the queueing time was hidden and is not a target. Diagnostic only.
13753 let phase_sync = std::env::var("MEMRA_SPEC_PHASE_SYNC").as_deref() == Ok("1");
13754 // DRAFT-MASK receipt (lane/draft-mask): speculative-clone wall + rounds, printed with
13755 // spec-stats. The clone is the one cost the design adds per round — measured, not assumed.
13756 let (mut dm_clone_ns, mut dm_rounds) = (0u128, 0usize);
13757 // grammar-truncation counters: how many rounds the verify-side cut fired and how many
13758 // already-verified tokens it threw away. THIS is the quantity draft masking targets.
13759 let (mut dm_cuts, mut dm_cut_tokens) = (0usize, 0usize);
13760 let (mut ph_draft, mut ph_verify, mut ph_rest) = (0f64, 0f64, 0f64);
13761 let mut ph_wait = 0f64;
13762 let mut ph_commit = 0f64;
13763 let mut ph_t = std::time::Instant::now();
13764 let mut ph_mark = |acc: &mut f64, on: bool| {
13765 if on {
13766 let now = std::time::Instant::now();
13767 *acc += (now - ph_t).as_secs_f64();
13768 ph_t = now;
13769 }
13770 };
13771 // MTP-ROUTE VERIFY GRAPHS (`MEMRA_SPEC_VERIFY_GRAPH`, see the flag doc): the
13772 // model-owned capture pool, locked for the whole burst exactly as the dspark serve
13773 // arm holds it — the slab stash is live verify -> commit inside a round, and the
13774 // worker drives rounds from one scheduler thread. PERSISTENT across generations on
13775 // the model (rebuilding per call re-captures the pool per prompt, which is the
13776 // measured way to lose more than the launches cost); the captured bodies are
13777 // cache-independent, every state read going through per-round refreshed pointer
13778 // tables. None = the eager walk, byte-identical.
13779 //
13780 // Never armed together with ROUND-STREAM: the tparallel verify refuses that pair
13781 // loudly, and `stream_active` owns the burst arm above, so the door stays shut
13782 // whenever the stream is live rather than relying on that refusal.
13783 // The lock is taken ONLY when the door is armed: with the flag off this whole block
13784 // is inert, so the default path cannot serialize two spec generations behind a mutex
13785 // it never reads.
13786 let vg_armed =
13787 crate::spec::spec_verify_graph_env().unwrap_or_else(|| self.vgraph_family_default());
13788 let mut vg_guard = if vg_armed && !stream_active {
13789 let mut g = self.dspark_vgraphs.lock().unwrap();
13790 if g.is_none() {
13791 // Size by the WIDEST verify this run can present, which is k+1 and NOT
13792 // k_cap+1: the sampled arm's own window is `t_v_s = k + 1`, so a pool built
13793 // from a smaller adaptive cap gets sliced past its stash rows (a `slice_mut`
13794 // panic in the sampled ON arm, measured before this line said k+1).
13795 let vt_cap = (k.max(k_cap) + 1).max(2);
13796 *g = DsparkVerifyGraphs::new(e, cache, vt_cap, n_embd)?;
13797 if g.is_some() {
13798 // Engagement receipt (the dead-arm lesson): prove the door is LIVE rather
13799 // than trusting that a flag set means a pool built.
13800 eprintln!("[spec-vg] MTP verify-graph pool ENGAGED (vt_cap={vt_cap})");
13801 } else {
13802 eprintln!(
13803 "[spec-vg] MTP verify-graph pool declined (no linear layers, \
13804 non-uniform state, or vt_cap < 2) — eager walk"
13805 );
13806 }
13807 }
13808 Some(g)
13809 } else {
13810 None
13811 };
13812 // Capacity fail-safe: a round wider than the pool was built for must take the eager
13813 // walk, not slice the stash past its rows. The sizing above already covers every
13814 // round this run can present; this keeps a future caller (or a k that grows behind
13815 // the pool's back) on the byte-identical fallback instead of a panic.
13816 let vg_t_cap = vg_guard
13817 .as_ref()
13818 .and_then(|g| g.as_ref())
13819 .map(|g| g.t_capacity())
13820 .unwrap_or(0);
13821 if let Some(p) = pipe {
13822 p.setup_end();
13823 }
13824 drop(pipe_setup_walk);
13825 let mut graph_guard_noted = false;
13826 while keep_going && out.len() < max_new {
13827 // GRAPH-LAUNCH HEADROOM GUARD (see GRAPH_LAUNCH_MIN_FREE): below the floor,
13828 // every captured-graph arm in this round yields to its byte-identical eager
13829 // twin instead of feeding cuGraphLaunch a card it segfaults on.
13830 let graph_round_ok = graph_launch_headroom_ok(e);
13831 if !graph_round_ok && !graph_guard_noted {
13832 graph_guard_noted = true;
13833 eprintln!(
13834 "[spec] graph replay suspended: driver free below the {}MB launch floor \
13835 (eager arms serve; cuGraphLaunch segfaults into an exhausted card)",
13836 GRAPH_LAUNCH_MIN_FREE / (1 << 20)
13837 );
13838 }
13839 // MEMRA_SPEC_ROUND_PROF=1: wall of the WHOLE round against the pieces we already
13840 // instrument. Needed because the parts do not add up: the draft step measures 1.27 ms
13841 // ([spec-anatomy] glue 92 / attn 280 / ffn 222 / head 670 us) and the t=2 verify walk
13842 // 25.6 ms ([tcol-prof] attn 10.1 + ffn 15.3), yet a K=1 round takes 177 ms on the
13843 // step37 TP2 stack. This prints where the other ~150 ms lives.
13844 let round_prof = ROUND_PROF
13845 .get_or_init(|| std::env::var("MEMRA_SPEC_ROUND_PROF").as_deref() == Ok("1"));
13846 let round_t0 = round_prof.then(std::time::Instant::now);
13847 // ROUND-STREAM BURST: from round 1 (pending guaranteed by every non-replay arm),
13848 // issue M rounds with zero readbacks, then drain the ring + reconcile mirrors.
13849 if let (true, Some(sg), Some(ptrs)) = (
13850 stream_active && round >= 1 && pending.is_some() && graph_round_ok,
13851 &stream_graph,
13852 &stream_ptrs,
13853 ) {
13854 if debug_spec {
13855 static ONCE: std::sync::Once = std::sync::Once::new();
13856 ONCE.call_once(|| {
13857 eprintln!("[memra] ROUND-STREAM burst engaged (M={m_rounds} k={k})")
13858 });
13859 }
13860 e.set_i32_one(&mut pos_ctr, cache.pos as i32)?;
13861 e.set_u32_one(&mut pend_d, pending.unwrap())?;
13862 e.set_u32_one(&mut ring_d, 0)?; // ring count = 0 (writes element 0)
13863 for _mi in 0..m_rounds {
13864 e.i32_copy_add(&pos_ctr, &mut pos_start_d, 0)?;
13865 cache.snapshot_into(e, &mut snap)?; // device D2Ds, stream-ordered
13866 e.i32_copy_add(&pos_ctr, &mut scratch.kv.len_d, 0)?; // draft-KV rollback
13867 e.i32_copy_add(&pos_ctr, &mut dctx.g_pos, 1)?; // rope pos = pos + base
13868 e.u32_copy(&pend_d, &mut dctx.g_tok)?;
13869 e.copy_into(&mut dctx.g_seed, 0, &h_seed_buf, n_embd)?;
13870 sg.launch()?;
13871 e.spec_assemble_verify(
13872 &g_tokp2k,
13873 &pend_d,
13874 d2t_dev.as_ref(),
13875 &mut vtok_d,
13876 &mut brk_d,
13877 p_min,
13878 k,
13879 pmin0,
13880 )?;
13881 let mut ck = VerifyCkpt::new(self.layers.len());
13882 let dummy = vec![0u32; t_v_s];
13883 let (tl_d, vx) = self.decode_step_t_core_stream(
13884 e,
13885 &dummy,
13886 0,
13887 &mut *cache,
13888 embd_dev,
13889 Some(&mut ck),
13890 Some((&vtok_d, &pos_ctr)),
13891 None,
13892 None,
13893 None,
13894 )?;
13895 for j in 0..t_v_s {
13896 e.argmax_token_device_col(&tl_d, j, n_vocab, &mut preds_d, j)?;
13897 }
13898 e.spec_accept_greedy_dc(
13899 &preds_d,
13900 &vtok_d,
13901 &last_pred_d,
13902 &brk_d,
13903 &mut stream_acc,
13904 )?;
13905 e.spec_seed_gather(&vx, &fill_prev, &stream_acc, &mut h_seed_buf, 1, n_embd)?;
13906 e.copy_into(&mut fill_prev, 0, &h_seed_buf, n_embd)?;
13907 self.commit_verified_prefix_stream(
13908 e,
13909 &mut *cache,
13910 &snap,
13911 &ck,
13912 &stream_acc,
13913 1,
13914 t_v_s,
13915 )?;
13916 e.spec_rollback_stream(
13917 ptrs,
13918 &pos_start_d,
13919 &stream_acc,
13920 1,
13921 self.layers.len() + 1,
13922 )?;
13923 e.spec_ring_commit(&vtok_d, &stream_acc, &brk_d, &mut ring_d, &mut pend_d)?;
13924 }
13925 e.stream().synchronize()?;
13926 let ring_h = e.dtoh_u32(&ring_d)?;
13927 let cnt = ring_h[0] as usize;
13928 for i in 0..cnt {
13929 if out.len() < max_new {
13930 out.push(ring_h[1 + i]);
13931 }
13932 }
13933 let pos_h = e.dtoh_i32(&pos_ctr)?[0] as usize;
13934 for il in 0..self.layers.len() {
13935 if let Some(kvl) = cache.kv[il].as_mut() {
13936 kvl.len = pos_h;
13937 }
13938 }
13939 cache.pos = pos_h;
13940 scratch.kv.len = pos_h;
13941 pending = Some(ring_h[cnt]); // last drained token = the live bonus
13942 last_token = ring_h[cnt];
13943 total_drafted += k * m_rounds; // upper bound (p-min breaks uncounted)
13944 total_accepted += cnt.saturating_sub(m_rounds);
13945 if let Some(t) = sess_telem {
13946 // totals only — the burst's per-round accept counts stayed on device
13947 // (that is the point of the round-stream arm). pos_* untouched.
13948 t.record_totals(m_rounds, k * m_rounds, cnt.saturating_sub(m_rounds));
13949 }
13950 round += m_rounds;
13951 // sse-cadence: the drained ring is committed — flush it at burst-drain cadence.
13952 keep_going = flush_commit(&mut on_commit, &out, &mut flushed);
13953 continue;
13954 }
13955 let pipe_draft = match pipe {
13956 Some(p) => Some(p.draft_begin(round)?),
13957 None => None,
13958 };
13959 let pos = cache.pos; // #tokens committed (EXCLUDES a pending bonus)
13960 let mut current_opti = carried_opti.take();
13961 let mut fork_generation = if current_opti.is_none() && pending.is_some() {
13962 match opti_fork.as_mut() {
13963 Some(fork) if fork.mode.is_forced() => Some(fork.reserve(&mut snap)?),
13964 None => None,
13965 Some(_) => None,
13966 }
13967 } else {
13968 None
13969 };
13970 if current_opti.is_none() {
13971 if let Some(fork) = opti_fork.as_ref() {
13972 opti_snapshot_stage_owned_into(e, cache, fork.rt, &fork.fence, &mut snap)?;
13973 } else {
13974 cache.snapshot_into(e, &mut snap)?;
13975 }
13976 } else if snap.pos != pos {
13977 return Err(format!(
13978 "optipipe carried snapshot pos {} != current pos {pos}",
13979 snap.pos
13980 )
13981 .into());
13982 } // §C: snapshot BEFORE draft+verify (already retained for a carried successor)
13983 ph_mark(&mut ph_rest, phase_on);
13984
13985 // --- 1. DRAFT k tokens with the NextN head (autoregressive, T=1 each) ---
13986 // p-min semantics (both paths): stop the chain early when the head's confidence in
13987 // its own pick drops below p_min — the just-drafted token is DISCARDED, but its
13988 // scratch append stands (identical to the eager chain's ordering). j==0 always drafts.
13989 let base0 = if pending.is_some() { 1usize } else { 0usize };
13990 // fixed draft length by default; MEMRA_SPEC_ADAPT=1 drafts at last round's
13991 // accepted run + 1 (the gemma law — see the setup block above the loop).
13992 let k_this = if adapt { kc } else { k };
13993 let mut draft: Vec<u32> = Vec::with_capacity(k);
13994 let mut draft_idx: Vec<u32> = Vec::with_capacity(k); // trimmed-vocab ids (== draft when untrimmed)
13995 let mut controller_draft_prob: Option<f32> = None;
13996 let mut controller_eager_state: Option<(u32, CudaSlice<f32>)> = None;
13997 if let Some(ticket) = current_opti.as_mut() {
13998 let carried_pending = pending.ok_or("optipipe carried successor lost pending")?;
13999 if ticket.verify_tokens[0] != carried_pending {
14000 return Err(format!(
14001 "optipipe carried pending mismatch: ticket={} live={carried_pending}",
14002 ticket.verify_tokens[0],
14003 )
14004 .into());
14005 }
14006 draft.push(ticket.verify_tokens[1]);
14007 controller_draft_prob = Some(ticket.draft_prob);
14008 controller_eager_state = ticket
14009 .take_eager_seed()
14010 .map(|seed| (ticket.verify_tokens[1], seed));
14011 } else {
14012 // Round-start draft-KV sync (BOTH paths). Persistent: truncate/align to the committed
14013 // history — slots 0..P hold entries for the tokens before last_token@P (P = pos +
14014 // base0 - 1); this single set_len IS the draft-side rollback (drops last round's
14015 // rejected drafts and p-min extras via the len mechanism).
14016 scratch.set_len(e, pos + base0 - 1)?;
14017 // dcw door: a captured chain appends k_this device-counter rows (plus the
14018 // pseudo-seed replay) with no host intervention; any ring rebase those appends
14019 // could need happens HERE, host-side, before the replays. The eager arm keeps
14020 // its own per-step prepare, so this is graph-path-only work.
14021 if step35_draft_dcw_on()
14022 && (dctx.graph.is_some()
14023 || dctx.graph_s.is_some()
14024 || dctx.chain.is_some()
14025 || dctx.chain_s.is_some())
14026 {
14027 scratch.ensure_dcw_headroom(e, k_this + 2)?;
14028 }
14029 if pen_on {
14030 // PEN_WINDOW_MAX also bounds the per-round upload and the O(n_hist^2)
14031 // device dedup: the serve window is already PEN_WINDOW_MAX, and this
14032 // defensive min also bounds non-server callers.
14033 let win = sp.penalty_last_n.min(PEN_WINDOW_MAX);
14034 let w0 = pen_hist.len().saturating_sub(win);
14035 pen_hist_d = Some(e.htod_u32_v(&pen_hist[w0..])?);
14036 }
14037 if sampled {
14038 draft_logits.clear();
14039 draft_stats.clear();
14040 }
14041 // DRAFT-SIDE GRAMMAR MASK: clone the committed grammar state ONCE per round; each
14042 // position's mask is computed on that clone and advanced by the PROPOSED token. The
14043 // real state moves only on emission (verify's job), so the emitted stream is
14044 // unchanged — the mask only removes tokens the verify would have truncated anyway.
14045 let mut dmask_live = dmask_on;
14046 if dmask_live {
14047 let t_c = std::time::Instant::now();
14048 constraint
14049 .as_deref_mut()
14050 .unwrap()
14051 .draft_begin()
14052 .map_err(|e2| format!("constraint: {e2}"))?;
14053 dm_clone_ns += t_c.elapsed().as_nanos();
14054 dm_rounds += 1;
14055 }
14056 if let (false, Some(cg)) = (sampled || pen_on || !graph_round_ok, &dctx.chain) {
14057 // GREEDY CHAIN GRAPH (lane/step37-draft-graph-serving-20260830): the
14058 // eager multi-head chain's EXACT launch order — step j rewinds head
14059 // (j % heads)'s plane to the committed length and replays rows 0..=j —
14060 // with each row's whole head-forward as ONE graph launch. The chain
14061 // POLICY (head choice, prefix length, stored-seed feed) is host-side,
14062 // identical to `mtp_chain_forward_dev`, so graph-vs-eager drafts are
14063 // bit-identical by construction (same launcher, same bucket — the dcw
14064 // parity contract). Interior rows launch the head-less graph: their
14065 // logits are dead in the eager chain too, so the consumed bytes match.
14066 let heads_n = self.mtp_head_count();
14067 let committed = pos + base0 - 1;
14068 let mut chain_tokens: Vec<u32> = vec![last_token];
14069 let mut chain_seed_bufs: Vec<CudaSlice<f32>> = vec![e.clone_dtod(&h_seed_buf)?];
14070 for j in 0..k_this {
14071 let index = mtp_chain_head_index(j, heads_n);
14072 if debug_spec {
14073 eprintln!(
14074 "[mtp-chain-step] round={round} j={j} head={index} \
14075 replay_rows={} arm=graph",
14076 chain_tokens.len(),
14077 );
14078 }
14079 scratch.set_plane_len(e, index, committed)?;
14080 e.set_i32_one(&mut dctx.g_pos, (committed + 1) as i32)?;
14081 for row in 0..=j {
14082 e.set_u32_one(&mut dctx.g_tok, chain_tokens[row])?;
14083 e.copy_into(&mut dctx.g_seed, 0, &chain_seed_bufs[row], n_embd)?;
14084 if row < j {
14085 cg.interior[index].launch()?;
14086 } else {
14087 // per-position mask upload before the LAST row only — the
14088 // eager chain applies the mask on is_last exactly the same.
14089 if dmask_live
14090 && !upload_draft_mask(
14091 e,
14092 constraint.as_deref_mut().unwrap(),
14093 &mut dctx.g_dmask,
14094 mtp.d2t.as_ref(),
14095 d_vocab,
14096 dmask_words,
14097 )?
14098 {
14099 e.htod_u32_into(
14100 &mut dctx.g_dmask,
14101 &vec![u32::MAX; dmask_words],
14102 )?;
14103 dmask_live = false;
14104 }
14105 cg.last[index].launch()?;
14106 }
14107 // host mirror (len_d advanced in-graph by the dcw append)
14108 scratch.plane_mut(index).0.len += 1;
14109 }
14110 let idx = e.dtoh_u32_one(&dctx.g_tok)?;
14111 // #87 SENTINEL TRAP (see the single-head graph arm below).
14112 if (idx as usize) >= d_vocab {
14113 let seed_h = e.dtoh(&dctx.g_seed)?;
14114 let seed_nan = seed_h.iter().filter(|v| v.is_nan()).count();
14115 return Err(format!(
14116 "draft(chain-graph) argmax sentinel 0x{idx:08x} >= d_vocab \
14117 {d_vocab} at round {round} j={j} head={index} pos={pos}: \
14118 head-out NaN {seed_nan}/{n_embd} — refusing to dereference \
14119 the embed row (#87 trap)"
14120 )
14121 .into());
14122 }
14123 // multi-head MTP forbids a trimmed head (validated at entry), so the
14124 // draft index IS the target id; keep the map for uniformity.
14125 let d = match &mtp.d2t {
14126 Some(map) => map[idx as usize],
14127 None => idx,
14128 };
14129 let draft_p = if p_min > 0.0 {
14130 Some(e.dtoh(&dctx.g_p)?[0])
14131 } else {
14132 None
14133 };
14134 if j == 0 {
14135 controller_draft_prob = draft_p;
14136 }
14137 if let Some(p) = draft_p.filter(|_| p_min > 0.0)
14138 && p < p_min
14139 && (j > 0 || (pmin0 && base0 == 1))
14140 {
14141 break;
14142 }
14143 draft.push(d);
14144 chain_tokens.push(d);
14145 // step j's h_nextn: the last-row graph self-fed it into g_seed —
14146 // snapshot it as the chain history seed for row j+1 (stream-ordered
14147 // after the launch, exactly the eager chain's chain_seeds push).
14148 chain_seed_bufs.push(e.clone_dtod(&dctx.g_seed)?);
14149 // speculative grammar advance (see the single-head graph arm).
14150 if dmask_live
14151 && !constraint
14152 .as_deref_mut()
14153 .unwrap()
14154 .draft_advance(d)
14155 .map_err(|e2| format!("constraint: {e2}"))?
14156 {
14157 e.htod_u32_into(&mut dctx.g_dmask, &vec![u32::MAX; dmask_words])?;
14158 break;
14159 }
14160 }
14161 } else if let (true, Some(cg)) = (
14162 sampled && s_capturable && dctx.s_key == Some(s_key) && graph_round_ok,
14163 &dctx.chain_s,
14164 ) {
14165 if skey_probe() {
14166 eprintln!(
14167 "[skey] chain=graph_chain_s round={round} capturable={} top_k={} \
14168 top_p={} min_p={} s_key_parked={:?}",
14169 s_capturable as u8, sp.top_k, sp.top_p, sp.min_p, dctx.s_key,
14170 );
14171 }
14172 // SAMPLED CHAIN GRAPH: the greedy chain arm's launch order with the
14173 // sampled last-row graphs — in-graph counter bump + (filtered) gumbel
14174 // draw + argmax; q retained per step into q_slots exactly like the
14175 // single-head sampled graph arm. Counter continuity: g_ctr host-seeded
14176 // to sctr-1 once per ROUND; each step's last-row graph bumps it BEFORE
14177 // the perturb, so step j consumes counter sctr+j — the eager Philox
14178 // stream (interior rows never draw, never bump).
14179 let heads_n = self.mtp_head_count();
14180 let committed = pos + base0 - 1;
14181 let filtered_stats_in_graph = s_key.filtered();
14182 let mut chain_tokens: Vec<u32> = vec![last_token];
14183 let mut chain_seed_bufs: Vec<CudaSlice<f32>> = vec![e.clone_dtod(&h_seed_buf)?];
14184 e.set_u32_one(&mut dctx.g_ctr, sctr.wrapping_sub(1))?;
14185 for j in 0..k_this {
14186 let index = mtp_chain_head_index(j, heads_n);
14187 if debug_spec {
14188 eprintln!(
14189 "[mtp-chain-step] round={round} j={j} head={index} \
14190 replay_rows={} arm=graph_s",
14191 chain_tokens.len(),
14192 );
14193 }
14194 scratch.set_plane_len(e, index, committed)?;
14195 e.set_i32_one(&mut dctx.g_pos, (committed + 1) as i32)?;
14196 for row in 0..=j {
14197 e.set_u32_one(&mut dctx.g_tok, chain_tokens[row])?;
14198 e.copy_into(&mut dctx.g_seed, 0, &chain_seed_bufs[row], n_embd)?;
14199 if row < j {
14200 cg.interior[index].launch()?;
14201 } else {
14202 cg.last[index].launch()?;
14203 }
14204 scratch.plane_mut(index).0.len += 1;
14205 }
14206 sctr += 1; // mirrors the in-graph g_ctr bump (eager parity:
14207 // counts the p-min-discarded token too)
14208 // q retention: ONE async D2D of the persistent head-logits buffer
14209 // into this round's slot j (stream-ordered after the replay).
14210 e.copy_into(&mut dctx.q_slots[j], 0, &dctx.g_q, d_vocab)?;
14211 // FILTERED capture: read the in-graph filter_stats scalars back per
14212 // replay instead of a second full-vocab filter_stats per slot post-
14213 // chain — bit-exact (the values the in-graph perturb consumed) and
14214 // measured worth ~5% of vendor-default serving tok/s at K=3. Before
14215 // the p-min break so the discarded slot's stats land too.
14216 if filtered_stats_in_graph {
14217 draft_stats.push((
14218 e.dtoh(&dctx.g_mx)?[0],
14219 e.dtoh(&dctx.g_th)?[0],
14220 e.dtoh(&dctx.g_z)?[0],
14221 ));
14222 }
14223 let idx = e.dtoh_u32_one(&dctx.g_tok)?;
14224 // #87 SENTINEL TRAP (see the single-head graph arms).
14225 if (idx as usize) >= d_vocab {
14226 let seed_h = e.dtoh(&dctx.g_seed)?;
14227 let seed_nan = seed_h.iter().filter(|v| v.is_nan()).count();
14228 return Err(format!(
14229 "draft(chain-graph-sampled) argmax sentinel 0x{idx:08x} >= \
14230 d_vocab {d_vocab} at round {round} j={j} head={index} pos={pos}: \
14231 head-out NaN {seed_nan}/{n_embd} — refusing to dereference the \
14232 embed row (#87 trap)"
14233 )
14234 .into());
14235 }
14236 let d = match &mtp.d2t {
14237 Some(map) => map[idx as usize],
14238 None => idx,
14239 };
14240 draft_idx.push(idx);
14241 if p_min > 0.0 {
14242 let p = e.dtoh(&dctx.g_p)?[0];
14243 if p < p_min && (j > 0 || (pmin0 && base0 == 1)) {
14244 break;
14245 }
14246 }
14247 draft.push(d);
14248 chain_tokens.push(d);
14249 chain_seed_bufs.push(e.clone_dtod(&dctx.g_seed)?);
14250 }
14251 // PURE-TEMP accept path: stats per used slot recomputed from the RETAINED
14252 // q with the SAME filter_stats program the eager arm runs (deployment-
14253 // keyed coop/plain choice, same input bits). The FILTERED graph read its
14254 // stats back per replay above.
14255 if !filtered_stats_in_graph {
14256 for j in 0..draft.len().max(draft_idx.len()) {
14257 let rows0 = e.htod_i32(&[0])?;
14258 let (mut th_d, mut z_d, mut mx_d) =
14259 (e.zeros(1)?, e.zeros(1)?, e.zeros(1)?);
14260 e.filter_stats(
14261 &dctx.q_slots[j],
14262 d_vocab,
14263 &rows0,
14264 &mut th_d,
14265 &mut z_d,
14266 &mut mx_d,
14267 d_vocab,
14268 1,
14269 sp_temp,
14270 sp.top_k,
14271 sp.top_p,
14272 sp.min_p,
14273 )?;
14274 draft_stats.push((
14275 e.dtoh(&mx_d)?[0],
14276 e.dtoh(&th_d)?[0],
14277 e.dtoh(&z_d)?[0],
14278 ));
14279 }
14280 }
14281 } else if let (false, Some(gr)) =
14282 (sampled || pen_on || !graph_round_ok, &dctx.graph)
14283 {
14284 // GRAPH DRAFT: one dispatch per drafted token. The chain feeds itself on-device
14285 // (in-graph argmax -> tok_d -> next replay's embed; h_nextn -> h_seed_d; pos_d
14286 // inc'd in-graph); the host only reads 4B token (+4B p) and decides the break.
14287 e.set_i32_one(&mut dctx.g_pos, (pos + base0) as i32)?;
14288 e.set_u32_one(&mut dctx.g_tok, last_token)?;
14289 e.copy_into(&mut dctx.g_seed, 0, &h_seed_buf, n_embd)?;
14290 for j in 0..k_this {
14291 // per-position mask upload (contents only — the graph's baked pointer is
14292 // dctx.g_dmask). All-ones once masking goes dead mid-chain, so the captured
14293 // mask node degrades to a no-op ban instead of needing a second graph.
14294 if dmask_live
14295 && !upload_draft_mask(
14296 e,
14297 constraint.as_deref_mut().unwrap(),
14298 &mut dctx.g_dmask,
14299 mtp.d2t.as_ref(),
14300 d_vocab,
14301 dmask_words,
14302 )?
14303 {
14304 // no draft-vocab row is grammar-legal here (a trimmed FR-Spec head can
14305 // genuinely miss the legal set): neutralize the captured mask node and
14306 // finish the chain UNMASKED — exactly pre-lane behaviour, never worse.
14307 e.htod_u32_into(&mut dctx.g_dmask, &vec![u32::MAX; dmask_words])?;
14308 dmask_live = false;
14309 }
14310 gr.launch()?;
14311 scratch.kv.len += 1; // host mirror (len_d advanced in-graph)
14312 let idx = e.dtoh_u32_one(&dctx.g_tok)?;
14313 // #87 SENTINEL TRAP: an all-NaN head-logits row leaves the device argmax's
14314 // init sentinel (0x7FFFFFFF) in g_tok — feeding it onward dereferences
14315 // embed_row(sentinel) = table + ~4.6TB (never mapped) inside the NEXT graph
14316 // replay's embed node, and the MMU fault kills the CUDA context for the
14317 // whole process (research/pp2spec-crash-20260807: 3 coredumps, byte-exact
14318 // VA arithmetic). Refuse loudly instead; the diagnostics name the first-NaN
14319 // buffer (g_seed = the verify-side handoff vs head-side compute).
14320 if (idx as usize) >= d_vocab {
14321 // g_seed is SELF-FED (the replay writes h_nextn back into it), so it
14322 // reads as the head's OUTPUT at j; h_seed_buf is the round's INPUT
14323 // seed, untouched since the round-start copy — the pair discriminates
14324 // "seed arrived poisoned" from "head forward produced NaN".
14325 let seed_h = e.dtoh(&dctx.g_seed)?;
14326 let seed_nan = seed_h.iter().filter(|v| v.is_nan()).count();
14327 let in_h = e.dtoh(&h_seed_buf)?;
14328 let in_nan = in_h.iter().filter(|v| v.is_nan()).count();
14329 return Err(format!(
14330 "draft(graph) argmax sentinel 0x{idx:08x} >= d_vocab {d_vocab} at \
14331 round {round} j={j} pos={pos}: head-out NaN {seed_nan}/{n_embd}, \
14332 round-input-seed NaN {in_nan}/{n_embd} — refusing to dereference \
14333 the embed row (#87 trap)"
14334 )
14335 .into());
14336 }
14337 // trimmed draft vocab -> target token id (identity when no d2t map)
14338 let d = match &mtp.d2t {
14339 Some(map) => map[idx as usize],
14340 None => idx,
14341 };
14342 let draft_p = if p_min > 0.0
14343 || opti_fork
14344 .as_ref()
14345 .is_some_and(|fork| fork.controller.is_some())
14346 {
14347 Some(e.dtoh(&dctx.g_p)?[0])
14348 } else {
14349 None
14350 };
14351 if j == 0 {
14352 controller_draft_prob = draft_p;
14353 }
14354 if let Some(p) = draft_p.filter(|_| p_min > 0.0)
14355 && p < p_min
14356 && (j > 0 || (pmin0 && base0 == 1))
14357 {
14358 break;
14359 }
14360 draft.push(d);
14361 // with a trimmed head the NEXT embed must read the TARGET id, not the draft
14362 // index the argmax wrote — patch the persistent token buffer (4B htod).
14363 if d != idx {
14364 e.set_u32_one(&mut dctx.g_tok, d)?;
14365 }
14366 // advance the SPECULATIVE state with the proposal; a dead chain drops to
14367 // unmasked drafting for the remaining positions (verify still arbitrates).
14368 // speculative advance; a chain the grammar can no longer follow (EOS
14369 // proposed) ends here. The captured mask node always runs, so a dead chain
14370 // leaves the buffer NEUTRAL (all-ones = ban nothing) before it exits.
14371 if dmask_live
14372 && !constraint
14373 .as_deref_mut()
14374 .unwrap()
14375 .draft_advance(d)
14376 .map_err(|e2| format!("constraint: {e2}"))?
14377 {
14378 e.htod_u32_into(&mut dctx.g_dmask, &vec![u32::MAX; dmask_words])?;
14379 break;
14380 }
14381 }
14382 // REGIME RE-TEST (lane/graph-s-key-exactness-20260819, widened by
14383 // lane/step37-draft-graph-serving-20260830): the sampled graph is legal ONLY
14384 // in the regime it was captured in. The condition used to read
14385 // `(sampled, &dctx.graph_s)` and trusted `s_key` to have dropped anything
14386 // else — which it could not, because the key omitted the filters. Both
14387 // halves are enforced: the key drops a stale graph, and this site refuses to
14388 // launch one whose key differs or whose regime is uncapturable (penalties).
14389 } else if let (true, Some(gr)) = (
14390 sampled && s_capturable && dctx.s_key == Some(s_key) && graph_round_ok,
14391 &dctx.graph_s,
14392 ) {
14393 if skey_probe() {
14394 eprintln!(
14395 "[skey] chain=graph_s round={round} pure_temp={} capturable={} \
14396 top_k={} top_p={} min_p={} s_key_parked={:?}",
14397 pure_temp as u8,
14398 s_capturable as u8,
14399 sp.top_k,
14400 sp.top_p,
14401 sp.min_p,
14402 dctx.s_key,
14403 );
14404 }
14405 // SAMPLED GRAPH DRAFT: one replay per drafted token — head forward + gumbel +
14406 // argmax in ONE dispatch; the host reads 4B token (+4B p), D2Ds q into slot j,
14407 // and decides the break. Event-counter continuity: g_ctr is host-seeded to
14408 // sctr-1 ONCE per round (outside the graph); the in-graph bump runs BEFORE the
14409 // perturb, so replay j consumes counter sctr+j — exactly the eager arm's Philox
14410 // stream. Host sctr advances in lockstep (computed, no readback needed).
14411 e.set_i32_one(&mut dctx.g_pos, (pos + base0) as i32)?;
14412 e.set_u32_one(&mut dctx.g_tok, last_token)?;
14413 e.copy_into(&mut dctx.g_seed, 0, &h_seed_buf, n_embd)?;
14414 e.set_u32_one(&mut dctx.g_ctr, sctr.wrapping_sub(1))?;
14415 let filtered_stats_in_graph = s_key.filtered();
14416 for j in 0..k_this {
14417 gr.launch()?;
14418 scratch.kv.len += 1; // host mirror (len_d advanced in-graph)
14419 sctr += 1; // mirrors the in-graph g_ctr bump (eager parity:
14420 // counts the p-min-discarded token too)
14421 // q retention: ONE async D2D of the persistent head-logits buffer into this
14422 // round's slot j (stream-ordered after the replay, before the next one).
14423 e.copy_into(&mut dctx.q_slots[j], 0, &dctx.g_q, d_vocab)?;
14424 // FILTERED capture: the replay's own filter_stats node already computed
14425 // (th, z, mx) — read the three scalars back instead of paying a SECOND
14426 // full-vocab filter_stats per slot post-chain (measured ~5% of vendor-
14427 // default serving tok/s at K=3). Bit-exact by construction: these are
14428 // the very values the in-graph perturb consumed. Read BEFORE the p-min
14429 // break so the discarded slot's stats land too (accept-path indexing).
14430 if filtered_stats_in_graph {
14431 draft_stats.push((
14432 e.dtoh(&dctx.g_mx)?[0],
14433 e.dtoh(&dctx.g_th)?[0],
14434 e.dtoh(&dctx.g_z)?[0],
14435 ));
14436 }
14437 let idx = e.dtoh_u32_one(&dctx.g_tok)?;
14438 // #87 SENTINEL TRAP (see the greedy graph arm above).
14439 if (idx as usize) >= d_vocab {
14440 let seed_h = e.dtoh(&dctx.g_seed)?;
14441 let seed_nan = seed_h.iter().filter(|v| v.is_nan()).count();
14442 return Err(format!(
14443 "draft(graph-sampled) argmax sentinel 0x{idx:08x} >= d_vocab \
14444 {d_vocab} at round {round} j={j} pos={pos}: round-seed NaN \
14445 {seed_nan}/{n_embd} — refusing to dereference the embed row \
14446 (#87 trap)"
14447 )
14448 .into());
14449 }
14450 let d = match &mtp.d2t {
14451 Some(map) => map[idx as usize],
14452 None => idx,
14453 };
14454 draft_idx.push(idx);
14455 if p_min > 0.0 {
14456 let p = e.dtoh(&dctx.g_p)?[0];
14457 if p < p_min && (j > 0 || (pmin0 && base0 == 1)) {
14458 break;
14459 }
14460 }
14461 draft.push(d);
14462 // trimmed head: the NEXT embed must read the TARGET id (see the greedy arm).
14463 if d != idx {
14464 e.set_u32_one(&mut dctx.g_tok, d)?;
14465 }
14466 }
14467 // PURE-TEMP accept path: fill draft_stats per used slot post-chain (the
14468 // stats degenerate to th=0 / full-Z; one filter_stats launch per slot).
14469 // The FILTERED graph read its stats back per replay above.
14470 if !filtered_stats_in_graph {
14471 for j in 0..draft.len().max(draft_idx.len()) {
14472 let rows0 = e.htod_i32(&[0])?;
14473 let (mut th_d, mut z_d, mut mx_d) =
14474 (e.zeros(1)?, e.zeros(1)?, e.zeros(1)?);
14475 e.filter_stats(
14476 &dctx.q_slots[j],
14477 d_vocab,
14478 &rows0,
14479 &mut th_d,
14480 &mut z_d,
14481 &mut mx_d,
14482 d_vocab,
14483 1,
14484 sp_temp,
14485 sp.top_k,
14486 sp.top_p,
14487 sp.min_p,
14488 )?;
14489 draft_stats.push((
14490 e.dtoh(&mx_d)?[0],
14491 e.dtoh(&th_d)?[0],
14492 e.dtoh(&z_d)?[0],
14493 ));
14494 }
14495 }
14496 } else {
14497 if skey_probe() && sampled {
14498 eprintln!(
14499 "[skey] chain=eager round={round} pure_temp={} top_k={} \
14500 top_p={} min_p={} s_key_parked={:?}",
14501 pure_temp as u8, sp.top_k, sp.top_p, sp.min_p, dctx.s_key,
14502 );
14503 }
14504 // EAGER DRAFT (fallback: MoE head/trunk, huge k, MEMRA_SPEC_NOGRAPH, capture fail).
14505 let chain_heads = !self.mtp_extra.is_empty();
14506 let mut e_tok = last_token;
14507 let mut d_seed = e.clone_dtod(&h_seed_buf)?;
14508 let mut chain_tokens = if chain_heads {
14509 vec![last_token]
14510 } else {
14511 Vec::new()
14512 };
14513 let mut chain_seeds = if chain_heads {
14514 vec![e.clone_dtod(&h_seed_buf)?]
14515 } else {
14516 Vec::new()
14517 };
14518 for j in 0..k_this {
14519 // GPU-ARGMAX DRAFT (2026-07-03): device logits + device argmax + 4-byte token
14520 // read instead of the ~600KB full-vocab dtoh + host argmax per draft token.
14521 let mtp_pos = pos + base0 + j;
14522 // draft-side grammar mask (eager twin of the graph arm's in-graph node).
14523 // A position with no legal draft-vocab row drops to unmasked drafting for
14524 // the rest of the chain (pre-lane behaviour; verify still arbitrates).
14525 if dmask_live {
14526 dmask_live = upload_draft_mask(
14527 e,
14528 constraint.as_deref_mut().unwrap(),
14529 &mut dctx.g_dmask,
14530 mtp.d2t.as_ref(),
14531 d_vocab,
14532 dmask_words,
14533 )?;
14534 }
14535 let mask = if dmask_live {
14536 Some((&dctx.g_dmask, dmask_words))
14537 } else {
14538 None
14539 };
14540 let (dl_d, h_nextn) = if chain_heads {
14541 if debug_spec {
14542 eprintln!(
14543 "[mtp-chain-step] round={round} j={j} head={} replay_rows={}",
14544 mtp_chain_head_index(j, self.mtp_head_count()),
14545 chain_tokens.len(),
14546 );
14547 }
14548 self.mtp_chain_forward_dev(
14549 e,
14550 &chain_tokens,
14551 &chain_seeds,
14552 &mut *scratch,
14553 pos + base0 - 1,
14554 embd_dev,
14555 mask,
14556 )?
14557 } else {
14558 self.mtp_head_forward_dev(
14559 e,
14560 mtp,
14561 e_tok,
14562 &d_seed,
14563 &mut *scratch,
14564 mtp_pos,
14565 embd_dev,
14566 mask,
14567 )?
14568 };
14569 let tok_d = if sampled {
14570 // FILTERED Gumbel-max: stats -> masked perturb -> argmax = one draw from
14571 // the filtered softmax (filters off => th=0, exact v1 semantics).
14572 if perturb_buf.is_none() {
14573 perturb_buf = Some(e.zeros(d_vocab.max(n_vocab))?);
14574 }
14575 let mut q_row = e.clone_dtod(&dl_d)?; // retained q (penalized when on)
14576 if pen_on {
14577 let h = pen_hist_d.as_ref().unwrap();
14578 let nh = h.len();
14579 e.penalize_logits(
14580 &mut q_row,
14581 h,
14582 nh,
14583 sp.penalty_repeat,
14584 sp.penalty_freq,
14585 sp.penalty_present,
14586 d_vocab,
14587 )?;
14588 }
14589 let rows0 = e.htod_i32(&[0])?;
14590 let (mut th_d, mut z_d, mut mx_d) =
14591 (e.zeros(1)?, e.zeros(1)?, e.zeros(1)?);
14592 e.filter_stats(
14593 &q_row, d_vocab, &rows0, &mut th_d, &mut z_d, &mut mx_d, d_vocab,
14594 1, sp_temp, sp.top_k, sp.top_p, sp.min_p,
14595 )?;
14596 let (th, z, mx) =
14597 (e.dtoh(&th_d)?[0], e.dtoh(&z_d)?[0], e.dtoh(&mx_d)?[0]);
14598 let pb = perturb_buf.as_mut().unwrap();
14599 e.gumbel_perturb_filtered(
14600 &q_row, pb, d_vocab, sp_seed, sctr, sp_temp, mx, th,
14601 )?;
14602 sctr += 1;
14603 draft_logits.push(q_row);
14604 draft_stats.push((mx, th, z));
14605 e.argmax_token_device(pb, d_vocab)?
14606 } else {
14607 e.argmax_token_device(&dl_d, d_vocab)?
14608 };
14609 let idx = e.dtoh_u32_one(&tok_d)?;
14610 // #87 SENTINEL TRAP (eager twin — see the graph arm). Extra diagnostics
14611 // here because the eager chain's operands are all readable: dl_d (the head
14612 // logits row) and d_seed (this step's h_seed) name the first-NaN buffer.
14613 if (idx as usize) >= d_vocab {
14614 let dl_h = e.dtoh(&dl_d)?;
14615 let dl_nan = dl_h.iter().filter(|v| v.is_nan()).count();
14616 let seed_h = if chain_heads {
14617 e.dtoh(chain_seeds.last().unwrap())?
14618 } else {
14619 e.dtoh(&d_seed)?
14620 };
14621 let seed_nan = seed_h.iter().filter(|v| v.is_nan()).count();
14622 return Err(format!(
14623 "draft(eager) argmax sentinel 0x{idx:08x} >= d_vocab {d_vocab} at \
14624 round {round} j={j} pos={pos}: head-logits NaN {dl_nan}/{d_vocab}, \
14625 step-seed NaN {seed_nan}/{n_embd} — refusing to dereference the \
14626 embed row (#87 trap)"
14627 )
14628 .into());
14629 }
14630 let d = match &mtp.d2t {
14631 Some(map) => map[idx as usize],
14632 None => idx,
14633 };
14634 if sampled {
14635 draft_idx.push(idx);
14636 }
14637 let draft_p = if p_min > 0.0
14638 || opti_fork
14639 .as_ref()
14640 .is_some_and(|fork| fork.controller.is_some())
14641 {
14642 let p_d = e.prob_of_token_device(&dl_d, &tok_d, d_vocab)?;
14643 Some(e.dtoh(&p_d)?[0])
14644 } else {
14645 None
14646 };
14647 if j == 0 {
14648 controller_draft_prob = draft_p;
14649 }
14650 if let Some(p) = draft_p.filter(|_| p_min > 0.0)
14651 && p < p_min
14652 && (j > 0 || (pmin0 && base0 == 1))
14653 {
14654 break;
14655 }
14656 draft.push(d);
14657 if chain_heads {
14658 chain_tokens.push(d);
14659 chain_seeds.push(h_nextn);
14660 } else {
14661 e_tok = d;
14662 d_seed = h_nextn;
14663 }
14664 // speculative advance; a chain the grammar can no longer follow (EOS
14665 // proposed) ends here — the prefix already proposed still rides verify.
14666 if dmask_live
14667 && !constraint
14668 .as_deref_mut()
14669 .unwrap()
14670 .draft_advance(d)
14671 .map_err(|e2| format!("constraint: {e2}"))?
14672 {
14673 break;
14674 }
14675 }
14676 if !chain_heads
14677 && opti_fork
14678 .as_ref()
14679 .is_some_and(|fork| fork.controller.is_some())
14680 {
14681 controller_eager_state = Some((e_tok, d_seed));
14682 }
14683 }
14684 }
14685 let k_round = draft.len();
14686 if let Some(p) = pipe {
14687 p.draft_end(round);
14688 }
14689 drop(pipe_draft);
14690
14691 ph_mark(&mut ph_draft, phase_on);
14692 // --- 2. VERIFY: one batched target forward. With a pending bonus, it rides as col 0
14693 // (committing its KV/recur inside the SAME weight read); drafts follow. ---
14694 let verify_tokens: Vec<u32> = match pending {
14695 Some(b) => {
14696 let mut v = Vec::with_capacity(k_round + 1);
14697 v.push(b);
14698 v.extend_from_slice(&draft);
14699 v
14700 }
14701 None => draft.clone(),
14702 };
14703 let base = if pending.is_some() { 1 } else { 0 };
14704 // ckpt (REPLAY-FREE partial accept): retain per-layer state-rebuild inputs alongside
14705 // the verify. Pure buffer keep-alives + dtod clones — kernel work is unchanged.
14706 let mut ckpt = if let Some(ticket) = current_opti.as_mut() {
14707 Some(ticket.take_ckpt())
14708 } else if spec_replay {
14709 None
14710 } else {
14711 Some(VerifyCkpt::new(self.layers.len()))
14712 };
14713 let controller_can_probe = base == 1
14714 && k_round == 1
14715 && out.len().saturating_add(2) < max_new
14716 && controller_draft_prob.is_some()
14717 && opti_fork
14718 .as_ref()
14719 .and_then(|fork| fork.controller.as_ref())
14720 .is_some_and(|policy| !policy.breaker_tripped);
14721 let mut successor_attempt: Option<OptiControllerTicket> = None;
14722 let mut rejected_probe: Option<(f32, u32)> = None;
14723 let mut controller_prepared: Option<OptiControllerPrepared> = None;
14724 if controller_can_probe {
14725 // Prepare d2/q and, on admission, d3 before either current verify half is
14726 // issued. N stage 0 can then be followed immediately by N+1 stage 0; once N's
14727 // boundary fires, those dev0 launches overlap N stage 1 on dev1. Preparing on
14728 // the primary stream after N stage 1 would serialize the supposed pipeline.
14729 let eager_pos = scratch.kv.len + 1;
14730 let (optimistic_pending, pending_probability) = self.opti_controller_draft_step(
14731 e,
14732 mtp,
14733 &mut dctx,
14734 &mut *scratch,
14735 d_vocab,
14736 &mut controller_eager_state,
14737 eager_pos,
14738 embd_dev,
14739 graph_round_ok,
14740 )?;
14741 let first_probability = controller_draft_prob
14742 .ok_or("optipipe controller probe lost first-token probability")?;
14743 let q_proxy = first_probability * pending_probability;
14744 OPTI_GATE_CHECKS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
14745 OPTI_SHADOW_DRAFT_TOKENS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
14746 let admitted = opti_fork
14747 .as_ref()
14748 .and_then(|fork| fork.controller.as_ref())
14749 .ok_or("optipipe controller policy disappeared")?
14750 .admit(q_proxy);
14751 if admitted {
14752 OPTI_GATE_ADMITS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
14753 let eager_pos = scratch.kv.len + 1;
14754 let (optimistic_draft, optimistic_draft_probability) = self
14755 .opti_controller_draft_step(
14756 e,
14757 mtp,
14758 &mut dctx,
14759 &mut *scratch,
14760 d_vocab,
14761 &mut controller_eager_state,
14762 eager_pos,
14763 embd_dev,
14764 graph_round_ok,
14765 )?;
14766 OPTI_SHADOW_DRAFT_TOKENS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
14767 let eager_seed = controller_eager_state.take().map(|(token, seed)| {
14768 debug_assert_eq!(token, optimistic_draft);
14769 seed
14770 });
14771 controller_prepared = Some(OptiControllerPrepared {
14772 verify_tokens: [optimistic_pending, optimistic_draft],
14773 draft_prob: optimistic_draft_probability,
14774 eager_seed,
14775 q_proxy,
14776 scratch_len: scratch.kv.len,
14777 });
14778 } else {
14779 OPTI_GATE_REJECTS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
14780 OPTI_WASTED_DRAFT_TOKENS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
14781 rejected_probe = Some((q_proxy, optimistic_pending));
14782 eprintln!(
14783 "[opti-controller] reject q={q_proxy:.6} threshold={:.3}",
14784 opti_fork
14785 .as_ref()
14786 .and_then(|fork| fork.controller.as_ref())
14787 .expect("controller policy")
14788 .threshold,
14789 );
14790 }
14791 }
14792 let fork_attempt = match fork_generation.take() {
14793 Some(generation) if base == 1 && k_round == 1 => Some(generation),
14794 Some(generation) => {
14795 opti_fork
14796 .as_mut()
14797 .expect("fork generation without fork state")
14798 .retire(generation)?;
14799 None
14800 }
14801 None => None,
14802 };
14803 let (tlogits_d, vx) = if let Some(p) = pipe {
14804 self.decode_step_t_core_pipelined(
14805 e,
14806 &verify_tokens,
14807 pos,
14808 &mut *cache,
14809 embd_dev,
14810 ckpt.as_mut(),
14811 p,
14812 round,
14813 )?
14814 } else if controller_can_probe {
14815 let fence = opti_fork
14816 .as_ref()
14817 .ok_or("optipipe controller probe lost fork state")?
14818 .fence;
14819 let boundary = match current_opti.as_mut() {
14820 Some(ticket) => ticket.take_boundary(),
14821 None => self.verify_stage0_issue(
14822 e,
14823 &verify_tokens,
14824 pos,
14825 &mut *cache,
14826 embd_dev,
14827 ckpt.as_mut(),
14828 None,
14829 &fence,
14830 Some(true),
14831 None,
14832 )?,
14833 };
14834 if let Some(prepared) = controller_prepared.take() {
14835 let generation = {
14836 let fork = opti_fork
14837 .as_mut()
14838 .ok_or("optipipe controller admission lost fork state")?;
14839 let generation = fork.reserve_successor()?;
14840 let rt = fork.rt;
14841 let snapshot_fence = fork.fence;
14842 opti_snapshot_one_stage_owned_into(
14843 e,
14844 cache,
14845 rt,
14846 &snapshot_fence,
14847 0,
14848 fork.successor_snapshot_mut(),
14849 )?;
14850 generation
14851 };
14852 let mut successor_ckpt = VerifyCkpt::new(self.layers.len());
14853 let successor_boundary = self.verify_stage0_issue(
14854 e,
14855 &prepared.verify_tokens,
14856 pos + verify_tokens.len(),
14857 &mut *cache,
14858 embd_dev,
14859 Some(&mut successor_ckpt),
14860 None,
14861 &fence,
14862 Some(false),
14863 None,
14864 )?;
14865 OPTI_FORK_ATTEMPTS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
14866 let fork = opti_fork
14867 .as_ref()
14868 .ok_or("optipipe controller ticket lost fork state")?;
14869 successor_attempt = Some(fork.controller_ticket(
14870 generation,
14871 successor_boundary,
14872 successor_ckpt,
14873 prepared.verify_tokens,
14874 prepared.draft_prob,
14875 prepared.eager_seed,
14876 prepared.q_proxy,
14877 prepared.scratch_len,
14878 ));
14879 eprintln!(
14880 "[opti-controller] issue generation={} q={:.6} threshold={:.3} \
14881 verify={:?}",
14882 generation.id,
14883 prepared.q_proxy,
14884 fork.controller.expect("controller policy").threshold,
14885 prepared.verify_tokens,
14886 );
14887 }
14888 let result = self.verify_stage1_finish(
14889 e,
14890 boundary,
14891 &mut *cache,
14892 ckpt.as_mut(),
14893 None,
14894 &fence,
14895 successor_attempt.is_none(),
14896 )?;
14897 if let Some(ticket) = current_opti.as_mut() {
14898 ticket.settle();
14899 }
14900 if successor_attempt.is_some() {
14901 let fork = opti_fork
14902 .as_mut()
14903 .ok_or("optipipe successor snapshot lost fork state")?;
14904 let rt = fork.rt;
14905 let snapshot_fence = fork.fence;
14906 opti_snapshot_one_stage_owned_into(
14907 e,
14908 cache,
14909 rt,
14910 &snapshot_fence,
14911 1,
14912 fork.successor_snapshot_mut(),
14913 )?;
14914 // Publish N only after both independent successor-state queues are complete.
14915 fork.rt.publish_to(1, &e.stream())?;
14916 }
14917 result
14918 } else if let Some(ticket) = current_opti.as_mut() {
14919 let fork = opti_fork
14920 .as_mut()
14921 .ok_or("optipipe carried controller ticket lost fork state")?;
14922 let boundary = ticket.take_boundary();
14923 let result = self.verify_stage1_finish(
14924 e,
14925 boundary,
14926 &mut *cache,
14927 ckpt.as_mut(),
14928 None,
14929 &fork.fence,
14930 true,
14931 )?;
14932 ticket.settle();
14933 result
14934 } else if let Some(generation) = fork_attempt {
14935 let fork = opti_fork
14936 .as_mut()
14937 .expect("fork generation without fork state");
14938 fork.capture_seed(e, generation, &h_seed_buf, &fill_prev, scratch.kv.len)?;
14939 let action = fork.mode.action(generation.id);
14940 let boundary = self.verify_stage0_issue(
14941 e,
14942 &verify_tokens,
14943 pos,
14944 &mut *cache,
14945 embd_dev,
14946 ckpt.as_mut(),
14947 None,
14948 &fork.fence,
14949 Some(true),
14950 None,
14951 )?;
14952 OPTI_FORK_ATTEMPTS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
14953 let mut ticket = fork.ticket(generation, boundary);
14954 if action == OptiForkAction::Abort {
14955 return Err(format!(
14956 "optipipe forced abort with generation {} stage0 in flight",
14957 generation.id,
14958 )
14959 .into());
14960 }
14961 fork.reconcile(
14962 e,
14963 &mut *cache,
14964 &mut *scratch,
14965 &snap,
14966 &mut h_seed_buf,
14967 &mut fill_prev,
14968 generation,
14969 action,
14970 verify_tokens[0],
14971 )?;
14972 let result = if action == OptiForkAction::Hit {
14973 let boundary = ticket.take_boundary();
14974 self.verify_stage1_finish(
14975 e,
14976 boundary,
14977 &mut *cache,
14978 ckpt.as_mut(),
14979 None,
14980 &fork.fence,
14981 true,
14982 )?
14983 } else {
14984 // The optimistic boundary slot has no reader. Re-run the unchanged serial
14985 // verify only after E_restart published the restored stage-0 state.
14986 self.decode_step_t_core(
14987 e,
14988 &verify_tokens,
14989 pos,
14990 &mut *cache,
14991 embd_dev,
14992 ckpt.as_mut(),
14993 )?
14994 };
14995 ticket.settle();
14996 debug_assert_eq!(ticket.generation, generation);
14997 fork.retire(generation)?;
14998 result
14999 } else {
15000 // The serial verify every non-fork round takes — the MTP route's
15001 // verify-graph door. The pool is None unless MEMRA_SPEC_VERIFY_GRAPH armed
15002 // a pool above, and then the walk replays the captured trunk instead of
15003 // re-issuing it launch by launch. `graph_round_ok` is the round's
15004 // headroom snapshot (see GRAPH_LAUNCH_MIN_FREE): below the floor the
15005 // round declines the pool exactly like an over-cap round and rides the
15006 // byte-identical eager walk — the `[spec]` suspension line above
15007 // already named the round.
15008 let vg_round = if verify_tokens.len() <= vg_t_cap && graph_round_ok {
15009 vg_guard.as_mut().and_then(|g| g.as_mut())
15010 } else {
15011 if let Some(g) = vg_guard.as_mut().and_then(|g| g.as_mut()) {
15012 // The commit reads this flag to pick its arm; a round that declines
15013 // the pool must not inherit a stale `true` from the round before it.
15014 g.round_slab = false;
15015 }
15016 None
15017 };
15018 self.decode_step_t_core_vg(
15019 e,
15020 &verify_tokens,
15021 pos,
15022 &mut *cache,
15023 embd_dev,
15024 ckpt.as_mut(),
15025 vg_round,
15026 )?
15027 };
15028 let pipe_accept = match pipe {
15029 Some(p) => Some(p.accept_begin(round)?),
15030 None => None,
15031 };
15032
15033 if phase_sync {
15034 e.stream().synchronize()?;
15035 }
15036 ph_mark(&mut ph_verify, phase_on);
15037 // --- 3. GREEDY ACCEPT (walk prefix, stop at first mismatch) ---
15038 // DEVICE-ARGMAX ACCEPT: argmax every verify column ON DEVICE (same 2-pass kernels +
15039 // smallest-index tie-break as host argmax, argmax_gate-validated) and read back ONE
15040 // [T] u32 — replaces the T x n_vocab f32 dtoh + T host argmaxes per round.
15041 // t_pred[j] = target's greedy prediction for the slot after draft[j-1] (j>=1) or after
15042 // last_token (j==0). With a pending bonus, col 0 IS the prediction after last_token
15043 // (== the bonus), so every index shifts by `base` and last_pred is unused.
15044 let t_v = verify_tokens.len();
15045 let mut preds: Vec<u32> = Vec::new();
15046 if !sampled {
15047 for j in 0..t_v {
15048 e.argmax_token_device_col(&tlogits_d, j, n_vocab, &mut preds_d, j)?;
15049 }
15050 preds = e.dtoh_u32(&preds_d)?; // <- the verify-GPU wait lands here
15051 // #87 SENTINEL TRAP, verify side: a sentinel pred becomes the round's bonus =
15052 // next round's last_token = the next chain's embed lookup. Catch it at the
15053 // source with the column named — an all-NaN VERIFY column implicates the
15054 // stage-split trunk (decode_step_t_core_ppn), not the draft head.
15055 if let Some(bad) = preds[..t_v].iter().position(|&p| (p as usize) >= n_vocab) {
15056 let col = &tlogits_d.slice(bad * n_vocab..(bad + 1) * n_vocab);
15057 let mut probe = e.zeros(n_vocab)?;
15058 e.copy_view_into(&mut probe, 0, col, n_vocab)?;
15059 let col_h = e.dtoh(&probe)?;
15060 let col_nan = col_h.iter().filter(|v| v.is_nan()).count();
15061 return Err(format!(
15062 "verify argmax sentinel 0x{:08x} >= n_vocab {n_vocab} at round {round} \
15063 col {bad}/{t_v} pos={pos}: verify-logits col NaN {col_nan}/{n_vocab} \
15064 — the verify TRUNK produced a poisoned column (#87 trap). Run \
15065 MEMRA_SPEC_NAN_SCAN=1 to name the layer that creates it (=2 to split \
15066 that layer into attention and routed MoE). NOT the draft head, and NOT \
15067 the PP stage split this message used to name: pp_cuts() returns None \
15068 without MEMRA_PP_STAGES, so decode_step_t_core_ppn never runs unless \
15069 that variable is set.",
15070 preds[bad]
15071 )
15072 .into());
15073 }
15074 }
15075 ph_mark(&mut ph_wait, phase_on);
15076 let t_pred = |j: usize| -> u32 {
15077 if j == 0 && base == 0 {
15078 last_pred
15079 } else {
15080 // GREEDY-ONLY: `preds` is filled under `if !sampled` above. The debug print
15081 // used to call this from the sampled arm and panicked the worker; it now goes
15082 // through `debug_t_pred0`. Keep the strict index here — in the greedy walk an
15083 // out-of-range pred is a real bug, not something to paper over.
15084 debug_assert!(
15085 !sampled,
15086 "t_pred is greedy-only: `preds` is empty in the sampled arm"
15087 );
15088 preds[base + j - 1]
15089 }
15090 };
15091 let mut devacc_seeded = false;
15092 let mut devacc_acc: Option<CudaSlice<u32>> = None;
15093 let (n_acc, bonus) = if !sampled {
15094 // ROUND-STREAM stage (a) (MEMRA_SPEC_DEVACC=1 opt-in): the walk runs ON DEVICE
15095 // (spec_accept_greedy, verbatim rule) and the host reads back 8B (n_acc, bonus)
15096 // instead of the [T] preds. Same sync count — machinery for stages (b)/(c),
15097 // gated on token identity vs the host walk (the arms below are bit-equal rules).
15098 if crate::spec::spec_devacc() && k_round > 0 && !spec_replay && constraint.is_none()
15099 {
15100 let draft_d = e.htod_u32_v(&draft)?;
15101 let mut acc_out = e.alloc_u32_zeroed(2)?;
15102 e.spec_accept_greedy(
15103 &preds_d,
15104 &draft_d,
15105 last_pred,
15106 base,
15107 k_round,
15108 &mut acc_out,
15109 )?;
15110 devacc_acc = Some(acc_out.clone());
15111 // stage (b): next-round seed gathered ON DEVICE from acc_out before the host
15112 // ever reads n_acc (j=base+n_acc -> vx col j-1; j==0 -> fill_prev). The three
15113 // non-replay commit arms skip their host-offset seed copies (guarded below);
15114 // the legacy spec_replay arm keeps its own rx-based seeding (excluded here).
15115 // NOTE: fill_prev is NOT updated here — the commit arms' TRUE-HIDDEN
15116 // REFRESH reads the OLD fill_prev (predecessor of this round's verify batch);
15117 // the update lands after the arms (devacc_seeded guard below).
15118 e.spec_seed_gather(&vx, &fill_prev, &acc_out, &mut h_seed_buf, base, n_embd)?;
15119 // 3a: KV lens roll back on device (len = saved + base + n_acc, all arms'
15120 // unified rule; full accept rewrites the verify-left value). Host mirrors
15121 // update after the readback; commit_verified_prefix skips its len_d writes.
15122 if let Some(successor) = successor_attempt.as_ref() {
15123 opti_fork
15124 .as_mut()
15125 .ok_or("optipipe successor reconcile lost fork state")?
15126 .queue_actual_reconcile(
15127 e,
15128 &snap,
15129 &acc_out,
15130 successor.verify_tokens[0],
15131 base,
15132 )?;
15133 } else if let Some(ptrs) = &kv_len_ptrs {
15134 let saved: Vec<i32> = (0..self.layers.len())
15135 .map(|il| snap.kv_len[il].map(|v| v as i32).unwrap_or(0))
15136 .collect();
15137 let saved_d = e.htod_i32(&saved)?;
15138 e.spec_rollback_kv(ptrs, &saved_d, &acc_out, base, self.layers.len())?;
15139 }
15140 devacc_seeded = true;
15141 let ab = e.dtoh_u32(&acc_out)?;
15142 (ab[0] as usize, ab[1])
15143 } else {
15144 let mut n_acc = 0usize;
15145 #[allow(clippy::needless_range_loop)]
15146 // allow: the explicit index loop keeps the offset arithmetic visible and aligned with the device-side indexing
15147 for j in 0..k_round {
15148 if t_pred(j) == draft[j] {
15149 n_acc += 1;
15150 } else {
15151 break;
15152 }
15153 }
15154 // bonus = target's own token at the first non-accepted slot. n_acc in 0..=k; t_pred
15155 // is defined for j in 0..=k (j==0 -> last_logits, j>=1 -> col j-1, last col = k-1).
15156 (n_acc, t_pred(n_acc))
15157 }
15158 } else {
15159 // --- SAMPLED ACCEPT (rejection sampling): u_j < p_j(x_j)/q_j(x_j) walk ---
15160 if col_buf.is_none() {
15161 col_buf = Some(e.zeros(n_vocab)?);
15162 }
15163 // FILTERED p_j: per-verify-col stats (one batched filter_stats call), then the
15164 // filtered gather. j==0&&base==0 reads last_col (its own stats row appended).
15165 let mut pj = vec![0f32; k_round.max(1)];
15166 let mut col_stats: Vec<(f32, f32, f32)> = Vec::new(); // (max, th, z) per verify col used
15167 if k_round > 0 {
15168 let mut ids: Vec<u32> = Vec::new();
15169 let mut rows: Vec<i32> = Vec::new();
15170 #[allow(clippy::needless_range_loop)]
15171 // allow: the explicit index loop keeps the offset arithmetic visible and aligned with the device-side indexing
15172 for j in 0..k_round {
15173 if j > 0 || base == 1 {
15174 ids.push(draft[j]);
15175 rows.push((base + j) as i32 - 1);
15176 }
15177 }
15178 if !ids.is_empty() {
15179 let nr = rows.len();
15180 // penalties: materialize the used columns into one contiguous penalized
15181 // buffer (rows remapped 0..nr) so stats+gathers see the penalized p.
15182 // penalties: materialize used columns contiguously, penalize all rows in
15183 // one launch, and point stats+gathers at the penalized buffer (rows 0..nr).
15184 let p_rows: Vec<i32> = if pen_on {
15185 (0..nr as i32).collect()
15186 } else {
15187 rows.clone()
15188 };
15189 if pen_on {
15190 if pcol_buf.as_ref().map(|b| b.len()).unwrap_or(0) < nr * n_vocab {
15191 pcol_buf = Some(e.zeros(nr * n_vocab)?);
15192 }
15193 let pc = pcol_buf.as_mut().unwrap();
15194 for (i2, &r) in rows.iter().enumerate() {
15195 let c = r as usize;
15196 e.copy_view_into(
15197 pc,
15198 i2 * n_vocab,
15199 &tlogits_d.slice(c * n_vocab..(c + 1) * n_vocab),
15200 n_vocab,
15201 )?;
15202 }
15203 let h = pen_hist_d.as_ref().unwrap();
15204 let nh = h.len();
15205 e.penalize_logits_rows(
15206 pc,
15207 h,
15208 nh,
15209 sp.penalty_repeat,
15210 sp.penalty_freq,
15211 sp.penalty_present,
15212 n_vocab,
15213 nr,
15214 )?;
15215 }
15216 let p_src: &CudaSlice<f32> = if pen_on {
15217 pcol_buf.as_ref().unwrap()
15218 } else {
15219 &tlogits_d
15220 };
15221 let rowsd = e.htod_i32(&p_rows)?;
15222 let (mut th_d, mut z_d, mut mx_d) =
15223 (e.zeros(nr)?, e.zeros(nr)?, e.zeros(nr)?);
15224 e.filter_stats(
15225 p_src, n_vocab, &rowsd, &mut th_d, &mut z_d, &mut mx_d, n_vocab, nr,
15226 sp_temp, sp.top_k, sp.top_p, sp.min_p,
15227 )?;
15228 let idsd = e.htod_u32_v(&ids)?;
15229 let mut outd = e.zeros(nr)?;
15230 e.softmax_gather_filtered(
15231 p_src, n_vocab, &idsd, &rowsd, &th_d, &z_d, &mut outd, n_vocab, nr,
15232 sp_temp,
15233 )?;
15234 let outv = e.dtoh(&outd)?;
15235 let (thv, zv, mxv) = (e.dtoh(&th_d)?, e.dtoh(&z_d)?, e.dtoh(&mx_d)?);
15236 let mut oi = 0usize;
15237 #[allow(clippy::needless_range_loop)]
15238 // allow: the explicit index loop keeps the offset arithmetic visible and aligned with the device-side indexing
15239 for j in 0..k_round {
15240 if j > 0 || base == 1 {
15241 pj[j] = outv[oi];
15242 oi += 1;
15243 }
15244 }
15245 col_stats = (0..nr).map(|i| (mxv[i], thv[i], zv[i])).collect();
15246 }
15247 if base == 0 {
15248 let lc: &CudaSlice<f32> = if pen_on {
15249 if col_buf.is_none() {
15250 col_buf = Some(e.zeros(n_vocab)?);
15251 }
15252 let cb = col_buf.as_mut().unwrap();
15253 e.copy_into(
15254 cb,
15255 0,
15256 last_col_logits
15257 .as_ref()
15258 .expect("sampled: last_col_logits unset"),
15259 n_vocab,
15260 )?;
15261 let h = pen_hist_d.as_ref().unwrap();
15262 let nh = h.len();
15263 e.penalize_logits(
15264 cb,
15265 h,
15266 nh,
15267 sp.penalty_repeat,
15268 sp.penalty_freq,
15269 sp.penalty_present,
15270 n_vocab,
15271 )?;
15272 col_buf.as_ref().unwrap()
15273 } else {
15274 last_col_logits
15275 .as_ref()
15276 .expect("sampled: last_col_logits unset")
15277 };
15278 let rows0 = e.htod_i32(&[0])?;
15279 let (mut th_d, mut z_d, mut mx_d) = (e.zeros(1)?, e.zeros(1)?, e.zeros(1)?);
15280 e.filter_stats(
15281 lc, n_vocab, &rows0, &mut th_d, &mut z_d, &mut mx_d, n_vocab, 1,
15282 sp_temp, sp.top_k, sp.top_p, sp.min_p,
15283 )?;
15284 let idsd = e.htod_u32_v(&[draft[0]])?;
15285 let mut outd = e.zeros(1)?;
15286 e.softmax_gather_filtered(
15287 lc, n_vocab, &idsd, &rows0, &th_d, &z_d, &mut outd, n_vocab, 1, sp_temp,
15288 )?;
15289 pj[0] = e.dtoh(&outd)?[0];
15290 last_col_stats =
15291 Some((e.dtoh(&mx_d)?[0], e.dtoh(&th_d)?[0], e.dtoh(&z_d)?[0]));
15292 }
15293 }
15294 // q source: the graph arms (single-head AND chain) retained the head logits
15295 // in the persistent q_slots; the eager arm in per-round draft_logits clones.
15296 // Same raw-logit values either way. FILTERED q_j: stats from draft_stats
15297 // (eager pushes in-chain; the graph arms compute them post-replay from the
15298 // retained q with the same filter_stats program — bit-identical to the
15299 // in-graph stats that shaped the draw, keeping ONE accept path).
15300 let q_bufs: &[CudaSlice<f32>] = if dctx.graph_s.is_some() || dctx.chain_s.is_some()
15301 {
15302 &dctx.q_slots
15303 } else {
15304 &draft_logits
15305 };
15306 let mut n_acc = 0usize;
15307 for j in 0..k_round {
15308 let (qmx, qth, qz) = draft_stats[j];
15309 let idsd = e.htod_u32_v(&[draft_idx[j]])?;
15310 let rowsd = e.htod_i32(&[0])?;
15311 let thd = e.htod(&[qth])?;
15312 let zd = e.htod(&[qz])?;
15313 let _ = qmx;
15314 let mut outd = e.zeros(1)?;
15315 e.softmax_gather_filtered(
15316 &q_bufs[j], d_vocab, &idsd, &rowsd, &thd, &zd, &mut outd, d_vocab, 1,
15317 sp_temp,
15318 )?;
15319 let qj = e.dtoh(&outd)?[0];
15320 let u = host_u01(sp_seed, uctr);
15321 uctr += 1;
15322 let accept = (u as f64) * (qj as f64) < pj[j] as f64;
15323 // SKEY PROBE: q == 0 for the token the draft actually proposed is the
15324 // exactness signature (see `skey_probe`). Impossible when the draft was
15325 // drawn from the same filtered distribution the verify reconstructs here;
15326 // `u * 0 < p` makes it an UNCONDITIONAL accept whenever p > 0.
15327 if skey_probe() && qj == 0.0 {
15328 eprintln!(
15329 "[skey] EXACTNESS q=0 round={round} j={j} draft_tok={} \
15330 draft_idx={} p={:e} u={u} accepted={} th_z={:?}",
15331 draft[j], draft_idx[j], pj[j], accept as u8, draft_stats[j],
15332 );
15333 }
15334 if accept {
15335 n_acc += 1;
15336 } else {
15337 break;
15338 }
15339 }
15340 let bonus = if n_acc == k_round {
15341 // FULL ACCEPT: bonus ~ FILTERED softmax at the last verify column.
15342 let col = base + k_round - 1;
15343 let cb = col_buf.as_mut().unwrap();
15344 e.copy_view_into(
15345 cb,
15346 0,
15347 &tlogits_d.slice(col * n_vocab..(col + 1) * n_vocab),
15348 n_vocab,
15349 )?;
15350 if pen_on {
15351 let h = pen_hist_d.as_ref().unwrap();
15352 let nh = h.len();
15353 e.penalize_logits(
15354 cb,
15355 h,
15356 nh,
15357 sp.penalty_repeat,
15358 sp.penalty_freq,
15359 sp.penalty_present,
15360 n_vocab,
15361 )?;
15362 }
15363 if perturb_buf.is_none() {
15364 perturb_buf = Some(e.zeros(d_vocab.max(n_vocab))?);
15365 }
15366 // STATS MUST COME FROM THIS COLUMN (bug fix 2026-08-05, lane/sampler-
15367 // truncation-fix; receipts research/sampfix-20260805/). The old code reused
15368 // `col_stats.last()` here, which is ALWAYS the wrong row: the gathered set
15369 // covers verify columns 0..=(base+k_round-2) (rows pushed as base+j-1), while
15370 // the full-accept bonus samples column base+k_round-1 — exactly ONE PAST the
15371 // last gathered column, in both base arms. `th` is a threshold in e-units of
15372 // its OWN row's max, so feeding a neighbour's (row_max, th) into
15373 // gumbel_perturb_filtered mis-scales every e0 = exp((x-row_max)/T). When the
15374 // donor column's peak is higher by more than T*ln(1/th), EVERY id fails
15375 // `e0 >= th`, the whole perturbed row becomes -3.4e38, and the 2-pass argmax
15376 // falls through to its smallest-index tie-break => token id 0 ("!") spliced
15377 // mid-word. Fragility is ordered by how large th is: min_p pins th = min_p
15378 // (0.05 => trigger at delta > 2.4 at T=0.8, fires constantly), top_p's
15379 // mass-boundary th is smaller, top_k's k-th-largest th smaller still — which
15380 // is why the head-to-head matrix saw min_p and top_p corrupt while top_k-only
15381 // stayed clean. The pure-temp default regime is immune (th == 0 masks nothing,
15382 // and row_max is unused once nothing is masked), so this fix is a byte-level
15383 // no-op for the untruncated serve default. One extra one-block filter_stats
15384 // per full-accept round is the whole cost.
15385 let (mx, th) = {
15386 let rows0 = e.htod_i32(&[0])?;
15387 let (mut th_d, mut z_d, mut mx_d) = (e.zeros(1)?, e.zeros(1)?, e.zeros(1)?);
15388 let cb0 = col_buf.as_ref().unwrap();
15389 e.filter_stats(
15390 cb0, n_vocab, &rows0, &mut th_d, &mut z_d, &mut mx_d, n_vocab, 1,
15391 sp_temp, sp.top_k, sp.top_p, sp.min_p,
15392 )?;
15393 (e.dtoh(&mx_d)?[0], e.dtoh(&th_d)?[0])
15394 };
15395 let pb = perturb_buf.as_mut().unwrap();
15396 let cb2 = col_buf.as_ref().unwrap();
15397 e.gumbel_perturb_filtered(cb2, pb, n_vocab, sp_seed, sctr, sp_temp, mx, th)?;
15398 sctr += 1;
15399 let td = e.argmax_token_device(pb, n_vocab)?;
15400 e.dtoh_u32_one(&td)?
15401 } else {
15402 // REJECT at n_acc: bonus ~ norm(max(0, softmax_T(p) - softmax_T(q))).
15403 let cb = col_buf.as_mut().unwrap();
15404 if n_acc > 0 || base == 1 {
15405 let col = base + n_acc - 1;
15406 e.copy_view_into(
15407 cb,
15408 0,
15409 &tlogits_d.slice(col * n_vocab..(col + 1) * n_vocab),
15410 n_vocab,
15411 )?;
15412 } else {
15413 let lc = last_col_logits.as_ref().unwrap();
15414 e.copy_into(cb, 0, lc, n_vocab)?;
15415 }
15416 if pen_on {
15417 let h = pen_hist_d.as_ref().unwrap();
15418 let nh = h.len();
15419 e.penalize_logits(
15420 cb,
15421 h,
15422 nh,
15423 sp.penalty_repeat,
15424 sp.penalty_freq,
15425 sp.penalty_present,
15426 n_vocab,
15427 )?;
15428 }
15429 let cb2 = col_buf.as_ref().unwrap();
15430 let sc = sctr;
15431 sctr += 1;
15432 // p-stats for the reject column: from col_stats when the col was gathered,
15433 // else (j==0&&base==0) from last_col_stats.
15434 let p_stats = if n_acc > 0 || base == 1 {
15435 // col index within the gathered set == number of gathered cols before n_acc
15436 let gi = if base == 1 { n_acc } else { n_acc - 1 };
15437 col_stats.get(gi).copied().unwrap_or({
15438 (0.0, 0.0, 1.0) // unreachable: gathered cols always cover the reject slot
15439 })
15440 } else {
15441 last_col_stats.expect("sampled: last_col_stats unset at reject")
15442 };
15443 let q_stats = draft_stats[n_acc];
15444 if let Some(map) = &d2t_dev {
15445 if q_full_buf.is_none() {
15446 q_full_buf = Some(e.zeros(n_vocab)?);
15447 }
15448 let qf = q_full_buf.as_mut().unwrap();
15449 e.scatter_trim_logits(&q_bufs[n_acc], map, qf, d_vocab, n_vocab)?;
15450 let qf2 = q_full_buf.as_ref().unwrap();
15451 e.residual_sample_filtered(
15452 cb2,
15453 Some(qf2),
15454 n_vocab,
15455 sp_temp,
15456 sp_seed,
15457 sc,
15458 p_stats,
15459 q_stats,
15460 &mut sample_tok,
15461 )?;
15462 } else {
15463 e.residual_sample_filtered(
15464 cb2,
15465 Some(&q_bufs[n_acc]),
15466 n_vocab,
15467 sp_temp,
15468 sp_seed,
15469 sc,
15470 p_stats,
15471 q_stats,
15472 &mut sample_tok,
15473 )?;
15474 }
15475 e.dtoh_u32(&sample_tok)?[0]
15476 };
15477 (
15478 n_acc,
15479 guard_vocab_token(
15480 bonus,
15481 n_vocab,
15482 &format!("sampled verify bonus at round {round} pos={pos} n_acc={n_acc}"),
15483 )?,
15484 )
15485 };
15486 // --- 3b. GRAMMAR TRUNCATION (constrained spec, 2026-08-03): the grammar is
15487 // an extra rejection rule AFTER the exactness verify (the batched-verify-twins
15488 // ordering). Walk the accepted drafts through the grammar in commit order; the
15489 // first illegal token truncates acceptance at its slot, and that slot's emission
15490 // is recomputed as the MASKED argmax of the target's own verify column — token-
15491 // identical to constrained plain greedy decode (an unmasked argmax that is
15492 // grammar-legal IS the masked argmax: masking only removes competitors). The
15493 // column D2H (~1MB) is paid only when a cut fires — the tight-grammar cost,
15494 // measured in acceptance numbers, never hidden.
15495 let (n_acc, bonus) = match constraint.as_deref_mut() {
15496 None => (n_acc, bonus),
15497 Some(c) => {
15498 fn ce(e2: String) -> Box<dyn std::error::Error> {
15499 format!("constraint: {e2}").into()
15500 }
15501 let mut na = n_acc;
15502 let mut cut = false;
15503 for (j, &d) in draft.iter().enumerate().take(n_acc) {
15504 if c.is_allowed(d).map_err(ce)? {
15505 c.consume(d).map_err(ce)?;
15506 } else {
15507 na = j;
15508 cut = true;
15509 dm_cut_tokens += n_acc - j;
15510 break;
15511 }
15512 }
15513 if cut {
15514 dm_cuts += 1;
15515 }
15516 let mut bo = bonus;
15517 if cut || !c.is_allowed(bo).map_err(ce)? {
15518 let mut row = if na == 0 && base == 0 {
15519 init_logits_host
15520 .clone()
15521 .ok_or("constraint: init logits missing (round-0 cut)")?
15522 } else {
15523 e.dtoh_view(
15524 &tlogits_d.slice((base + na - 1) * n_vocab..(base + na) * n_vocab),
15525 )?
15526 };
15527 c.mask_logits(&mut row).map_err(ce)?;
15528 bo = argmax(&row) as u32;
15529 }
15530 c.consume(bo).map_err(ce)?;
15531 (na, bo)
15532 }
15533 };
15534 let mut successor_valid = false;
15535 if let Some((q_proxy, expected_d2)) = rejected_probe {
15536 let v_n = n_acc == 1 && bonus == expected_d2;
15537 eprintln!(
15538 "[opti-controller] shadow q={q_proxy:.6} admitted=false v_n={v_n} \
15539 expected_d2={expected_d2} n_acc={n_acc} bonus={bonus}",
15540 );
15541 }
15542 if let Some(successor) = successor_attempt.as_ref() {
15543 successor_valid = n_acc == 1 && bonus == successor.verify_tokens[0];
15544 let generation = successor.generation;
15545 let q_proxy = successor.q_proxy;
15546 let expected_pending = successor.verify_tokens[0];
15547 let resolution_ms = successor.issued_at.elapsed().as_secs_f64() * 1e3;
15548 let fork = opti_fork
15549 .as_mut()
15550 .ok_or("optipipe successor resolution lost fork state")?;
15551 fork.finish_actual_reconcile(e, &mut *cache, &snap, n_acc, base, successor_valid)?;
15552 if successor_valid {
15553 OPTI_FORK_HITS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
15554 } else {
15555 OPTI_FORK_MISSES.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
15556 OPTI_RECONCILES.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
15557 OPTI_WASTED_DRAFT_TOKENS.fetch_add(2, std::sync::atomic::Ordering::Relaxed);
15558 }
15559 let breaker_tripped = fork
15560 .controller
15561 .as_mut()
15562 .expect("controller policy")
15563 .resolve(successor_valid);
15564 if breaker_tripped {
15565 OPTI_BREAKER_TRIPS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
15566 }
15567 eprintln!(
15568 "[opti-controller] resolve generation={} hit={} q={q_proxy:.6} \
15569 expected_pending={expected_pending} n_acc={n_acc} bonus={bonus} \
15570 resolution_ms={resolution_ms:.3} reconcile={} breaker={}",
15571 generation.id, successor_valid, !successor_valid, breaker_tripped,
15572 );
15573 if !successor_valid {
15574 let mut successor = successor_attempt
15575 .take()
15576 .expect("controller successor disappeared on miss");
15577 successor.settle();
15578 fork.retire(generation)?;
15579 }
15580 }
15581 total_drafted += k_round;
15582 total_accepted += n_acc;
15583 if let Some(t) = sess_telem {
15584 // Greedy, rejection-sampling, and grammar truncation all converge here after
15585 // the accept decision is already on host. Fixed-size relaxed atomics only.
15586 t.record_round(k_round, n_acc);
15587 }
15588 if spec_stats {
15589 st_len_hist[k_round] += 1;
15590 #[allow(clippy::needless_range_loop)]
15591 // allow: the explicit index loop keeps the offset arithmetic visible and aligned with the device-side indexing
15592 for j in 0..k_round {
15593 st_drafted[j] += 1;
15594 }
15595 #[allow(clippy::needless_range_loop)]
15596 // allow: the explicit index loop keeps the offset arithmetic visible and aligned with the device-side indexing
15597 for j in 0..n_acc {
15598 st_accepted[j] += 1;
15599 }
15600 if n_acc == k_round {
15601 st_full += 1;
15602 }
15603 }
15604
15605 if debug_spec {
15606 eprintln!(
15607 "[R{round}] pos={pos} out_len={} last_tok={last_token} draft={draft:?} n_acc={n_acc} bonus={bonus} t_pred0={}",
15608 out.len(),
15609 // NOT `t_pred(0)`: `preds` is filled only under `if !sampled` above, so on a
15610 // sampled request round >= 1 (base == 1) indexed an EMPTY vector and PANICKED
15611 // the GPU worker thread — a debug flag that killed the exact regime you would
15612 // set it to investigate. See `debug_t_pred0`.
15613 debug_t_pred0(sampled, base, last_pred, &preds)
15614 );
15615 }
15616
15617 // --- 4. COMMIT: draft[0..n_acc] then bonus (n_acc + 1 tokens) ---
15618 let commit_started = std::time::Instant::now();
15619 // SESSION MODE: every accepted column is already in the CACHE — `out` must carry all
15620 // of them (overshoot past max_new included) or `committed` under-counts the cache rows
15621 // and the next turn's continuation seeds one token off (gate-caught 2026-07-05). The
15622 // single-shot path keeps the cap (its caller truncates + drops the cache anyway).
15623 #[allow(clippy::needless_range_loop)]
15624 // allow: the explicit index loop keeps the offset arithmetic visible and aligned with the device-side indexing
15625 for j in 0..n_acc {
15626 if !session_mode && out.len() >= max_new {
15627 break;
15628 }
15629 out.push(draft[j]);
15630 }
15631 if pen_on {
15632 pen_hist.extend_from_slice(&draft[0..n_acc]);
15633 pen_hist.push(bonus);
15634 }
15635 let bonus_emitted = session_mode || out.len() < max_new;
15636 if bonus_emitted {
15637 out.push(bonus);
15638 }
15639 last_token = bonus;
15640
15641 // --- 5. ROLLBACK + advance (§C) ---
15642 if n_acc == k_round && !spec_replay {
15643 // FULL ACCEPT, BONUS FOLD: all verify columns (pending? + drafts) are committed in
15644 // cache; the NEW bonus stays PENDING for the next round's verify batch — NO extra
15645 // T=1 trunk pass. The next draft chain seeds from the MTP block's h_nextn at the
15646 // bonus position: one MTP-block pass (~1/33 trunk cost) replaces the trunk read.
15647 // last_pred is dead in the pending path (t_pred reads verify col 0).
15648 //
15649 // PERSISTENT DRAFT KV, full-accept fill: the chain covered last_token +
15650 // draft[0..k_round-2] as INPUTS (slots P..P'-2); draft[k_round-1] (slot P'-1) was
15651 // only ever an output, so its entry is MISSING. Fill it from vh_seed — its EXACT
15652 // trunk hidden (the last verify column). set_len first: a p-min break may have
15653 // left one extra chain append at that slot. Partial accepts need NO fill (the
15654 // chain already covered every accepted position; round-start set_len truncates).
15655 self.restore_step_tp_kv_verified_prefix(e, &mut *cache, &snap, t_v)?;
15656 let mut vh_seed = e.zeros(n_embd)?;
15657 e.copy_view_into(
15658 &mut vh_seed,
15659 0,
15660 &vx.slice((t_v - 1) * n_embd..t_v * n_embd),
15661 n_embd,
15662 )?;
15663 if refresh {
15664 // TRUE-HIDDEN REFRESH (2026-07-03, the HANDOVER-listed acceptance lever):
15665 // overwrite ALL committed positions' scratch entries with K/V from their EXACT
15666 // verify hiddens — the reference engine's mtp_update fills from true hiddens;
15667 // the full stack (vx) is already resident from the verify. Replaces both the
15668 // chain-approximate entries AND the old last-token-only fill. Acceptance-only
15669 // (draft attention quality); exactness stays the verify's job.
15670 scratch.set_len(e, pos)?;
15671 // PREDECESSOR pairing: row i gets vx[i-1]; row 0 the carried fill_prev
15672 // (hidden of the last committed row before this verify batch).
15673 let mut vxs = e.zeros(t_v * n_embd)?;
15674 e.copy_into(&mut vxs, 0, &fill_prev, n_embd)?;
15675 if t_v > 1 {
15676 e.copy_view_into(
15677 &mut vxs,
15678 n_embd,
15679 &vx.slice(0..(t_v - 1) * n_embd),
15680 (t_v - 1) * n_embd,
15681 )?;
15682 }
15683 self.mtp_kv_fill_all(e, &verify_tokens, &vxs, pos, &mut *scratch, embd_dev)?;
15684 } else {
15685 scratch.set_len(e, pos + base + k_round - 1)?;
15686 // predecessor of the last draft = verify col t_v-2 (or fill_prev at t_v==1)
15687 let mut hp = e.zeros(n_embd)?;
15688 if t_v >= 2 {
15689 e.copy_view_into(
15690 &mut hp,
15691 0,
15692 &vx.slice((t_v - 2) * n_embd..(t_v - 1) * n_embd),
15693 n_embd,
15694 )?;
15695 } else {
15696 e.copy_into(&mut hp, 0, &fill_prev, n_embd)?;
15697 }
15698 self.mtp_kv_fill_all(
15699 e,
15700 &[draft[k_round - 1]],
15701 &hp,
15702 pos + base + k_round - 1,
15703 &mut *scratch,
15704 embd_dev,
15705 )?;
15706 }
15707 // REFERENCE SEEDING: no pseudo pass — the next chain's step 0 IS the
15708 // reference's (id_last, h_prev) draft row; it appends the bonus's scratch
15709 // entry itself. Seed = TRUE hidden of the bonus's predecessor (last verify
15710 // col). Saves one MTP-block pass per round on top of the pairing fix.
15711 if !devacc_seeded {
15712 e.copy_into(&mut h_seed_buf, 0, &vh_seed, n_embd)?;
15713 e.copy_into(&mut fill_prev, 0, &vh_seed, n_embd)?;
15714 }
15715 pending = Some(bonus);
15716 if debug_spec {
15717 eprintln!(" -> FULL ACCEPT (bonus pending, prev-h seed)");
15718 }
15719 } else if !spec_replay && base + n_acc >= 1 {
15720 // PARTIAL ACCEPT, REPLAY-FREE (2026-07-03 — the profiled #1 long-ctx spec cost):
15721 // the verify's first j = base+n_acc columns ARE the committed sequence, computed
15722 // bit-identically to eager (decode-exact contract) — so KEEP them: KV truncates to
15723 // pos+j, recurrent state rebuilds from the VerifyCkpt (same-kernel gdn prefix
15724 // re-run / pure state-clone restore), and the bonus stays PENDING exactly like the
15725 // full-accept path — the legacy duplicate trunk replay is gone. The next chain
15726 // seeds from the MTP pseudo-hidden of the bonus, whose seed = the TRUE verify
15727 // hidden of its predecessor (col j-1) — same one-hop pseudo structure as full
15728 // accept (never compounds: the next verify recomputes true hiddens for all
15729 // committed columns).
15730 let j = base + n_acc;
15731 // VERIFY-GRAPH SLAB COMMIT: when the captured trunk ran, the linear layers'
15732 // column stash was written into the graphs ctx's persistent slabs as in-graph
15733 // memcpy nodes, NOT into the per-column VerifyCkpt the cols arm reads — so the
15734 // commit must take the slab twin (same semantics, slab-addressed sources). The
15735 // ctx states which of the two this round produced via `round_slab`; trusting the
15736 // flag rather than the env keeps a round that fell back to the eager walk (a
15737 // capture that declined, a t the pool never captured) on the cols arm.
15738 let slab_commit = vg_guard
15739 .as_ref()
15740 .and_then(|g| g.as_ref())
15741 .map(|g| g.round_slab)
15742 .unwrap_or(false);
15743 if slab_commit {
15744 self.dspark_commit_prefix_slab(
15745 e,
15746 &mut *cache,
15747 &snap,
15748 vg_guard
15749 .as_ref()
15750 .and_then(|g| g.as_ref())
15751 .expect("slab_commit implies a graphs ctx"),
15752 j,
15753 )?;
15754 } else {
15755 self.commit_verified_prefix(
15756 e,
15757 &mut *cache,
15758 &snap,
15759 ckpt.as_ref().unwrap(),
15760 j,
15761 devacc_seeded,
15762 if devacc_seeded {
15763 devacc_acc.as_ref().map(|a| (a, base, t_v))
15764 } else {
15765 None
15766 },
15767 )?;
15768 }
15769 let mut seed = e.zeros(n_embd)?;
15770 e.copy_view_into(
15771 &mut seed,
15772 0,
15773 &vx.slice((j - 1) * n_embd..j * n_embd),
15774 n_embd,
15775 )?;
15776 // Draft scratch: TRUE-HIDDEN REFRESH of the committed prefix (see the full-accept
15777 // branch); without it the chain entries stand and only the tail truncates. Either
15778 // way len ends at pos+j so the pseudo append lands at the bonus's slot pos+j
15779 // (persistent mode), rope pos+j+1 (chain convention).
15780 if refresh {
15781 scratch.set_len(e, pos)?;
15782 let mut vxs = e.zeros(j * n_embd)?;
15783 e.copy_into(&mut vxs, 0, &fill_prev, n_embd)?;
15784 if j > 1 {
15785 e.copy_view_into(
15786 &mut vxs,
15787 n_embd,
15788 &vx.slice(0..(j - 1) * n_embd),
15789 (j - 1) * n_embd,
15790 )?;
15791 }
15792 self.mtp_kv_fill_all(
15793 e,
15794 &verify_tokens[0..j],
15795 &vxs,
15796 pos,
15797 &mut *scratch,
15798 embd_dev,
15799 )?;
15800 } else {
15801 scratch.set_len(e, pos + j)?;
15802 }
15803 // REFERENCE SEEDING (see the full-accept branch): seed = TRUE hidden of the
15804 // bonus's predecessor (verify col j-1); no pseudo pass.
15805 if !devacc_seeded {
15806 e.copy_into(&mut h_seed_buf, 0, &seed, n_embd)?;
15807 e.copy_into(&mut fill_prev, 0, &seed, n_embd)?;
15808 }
15809 pending = Some(bonus);
15810 if debug_spec {
15811 eprintln!(" -> PARTIAL(replay-free j={j}, bonus pending, prev-h seed)");
15812 }
15813 } else if !spec_replay {
15814 // ZERO ROUND FOLD (2026-07-10, verify-cost target #3): base+n_acc == 0 — a
15815 // pending-less round where nothing was accepted (PMIN0 zero-draft chains after a
15816 // replay/commit, or plain 0-accept rounds at round 0). The old path replayed
15817 // [bonus] through a FULL m=1 trunk+head forward (the 489us full-vocab head pass
15818 // measured at ~0.75/round on PMIN0 configs). Instead: restore the pre-round
15819 // snapshot and let the bonus ride the NEXT round's verify as col 0 — the existing
15820 // base=1 pending machinery, bit-identical by the decode-exact verify contract.
15821 // Seed: the bonus's predecessor is the last COMMITTED token, whose hidden
15822 // fill_prev already carries (same seeding as the 1-token-replay case it replaces).
15823 cache.rollback(e, &snap, 0)?;
15824 scratch.set_len(e, pos)?;
15825 e.copy_into(&mut h_seed_buf, 0, &fill_prev, n_embd)?;
15826 pending = Some(bonus);
15827 if debug_spec {
15828 eprintln!(" -> ZERO-ROUND FOLD (bonus pending, fill_prev seed)");
15829 }
15830 } else {
15831 // PARTIAL ACCEPT, LEGACY REPLAY (seam MEMRA_SPEC_REPLAY=1 — or j==0: nothing of
15832 // this round survives, only possible before the first pending exists, ~round 0):
15833 // restore EVERYTHING to the pre-round snapshot (KV truncate to pos + recur
15834 // restore), then replay the committed prefix pending? ++ draft[0..n_acc] ++
15835 // [bonus] as ONE batched T forward — single weight read, bit-identical to greedy
15836 // (the verify-all-columns path is the same math). Commits the bonus with a TRUE
15837 // trunk hidden.
15838 cache.rollback(e, &snap, 0)?; // accept_len=0: KV len = pos, recur = snapshot
15839 let mut replay: Vec<u32> = Vec::with_capacity(base + n_acc + 1);
15840 if let Some(b) = pending.take() {
15841 replay.push(b);
15842 }
15843 replay.extend_from_slice(&draft[0..n_acc]);
15844 replay.push(bonus);
15845 // Full-stack forward (decode_step_t_core = decode_step_t_h_emb_dev's body):
15846 // Predecessor pairing seeds from the PREDECESSOR row (col len-2) — the same-row path takes the
15847 // last col exactly as before (byte-identical to the old _h_emb_dev call).
15848 let (rl_d, rx) = if self.batched_serving_numeric_class() {
15849 let mut logits = Vec::with_capacity(replay.len() * n_vocab);
15850 let mut hidden = e.uninit(replay.len() * n_embd)?;
15851 for (row, &token) in replay.iter().enumerate() {
15852 let (row_logits, row_hidden) =
15853 self.spec_target_step_h(e, token, &mut *cache)?;
15854 logits.extend_from_slice(&row_logits);
15855 e.dtod_copy_into(&row_hidden, &mut hidden, row * n_embd)?;
15856 }
15857 (e.htod(&logits)?, hidden)
15858 } else {
15859 self.decode_step_t_core(e, &replay, pos, &mut *cache, embd_dev, None)?
15860 };
15861 // last_pred = argmax of the LAST column's logits (predicts the token after `bonus`)
15862 // — device argmax + one 4-byte read instead of the full-vocab column dtoh.
15863 e.argmax_token_device_col(&rl_d, replay.len() - 1, n_vocab, &mut preds_d, 0)?;
15864 last_pred = guard_vocab_token(
15865 e.dtoh_u32(&preds_d)?[0],
15866 n_vocab,
15867 &format!("replay last_pred at round {round} pos={pos}"),
15868 )?;
15869 if sampled {
15870 let lr0 = replay.len();
15871 let lc = last_col_logits
15872 .as_mut()
15873 .expect("sampled: last_col_logits unset");
15874 e.copy_view_into(
15875 lc,
15876 0,
15877 &rl_d.slice((lr0 - 1) * n_vocab..lr0 * n_vocab),
15878 n_vocab,
15879 )?;
15880 }
15881 let lr = replay.len();
15882 if lr >= 2 {
15883 e.copy_view_into(
15884 &mut h_seed_buf,
15885 0,
15886 &rx.slice((lr - 2) * n_embd..(lr - 1) * n_embd),
15887 n_embd,
15888 )?;
15889 } else {
15890 // 1-token replay (round-0 miss): the bonus's predecessor is the OLD
15891 // last_token, whose own-row hidden fill_prev still holds.
15892 e.copy_into(&mut h_seed_buf, 0, &fill_prev, n_embd)?;
15893 }
15894 // the bonus is COMMITTED here — it becomes the last committed row.
15895 let mut rh_last = e.zeros(n_embd)?;
15896 e.copy_view_into(
15897 &mut rh_last,
15898 0,
15899 &rx.slice((lr - 1) * n_embd..lr * n_embd),
15900 n_embd,
15901 )?;
15902 e.copy_into(&mut fill_prev, 0, &rh_last, n_embd)?;
15903 if debug_spec {
15904 eprintln!(" -> PARTIAL(replay={replay:?}), next_pred={last_pred}");
15905 }
15906 }
15907 if devacc_seeded {
15908 // stage (b) epilogue: fill_prev takes the gathered seed AFTER the refresh fills
15909 // consumed the old value (both slots carry the same value in every non-replay arm).
15910 e.copy_into(&mut fill_prev, 0, &h_seed_buf, n_embd)?;
15911 }
15912 if successor_valid {
15913 let optimistic_scratch_len = successor_attempt
15914 .as_ref()
15915 .expect("valid controller successor disappeared")
15916 .scratch_len;
15917 // The normal current-round commit refreshed/truncated the logical scratch tail.
15918 // Its optimistic successor row was already written physically, so restoring only
15919 // the retained logical length makes that row live for the carried round.
15920 scratch.set_len(e, optimistic_scratch_len)?;
15921 }
15922 if let Some(current) = current_opti.take() {
15923 opti_fork
15924 .as_mut()
15925 .ok_or("optipipe current retirement lost fork state")?
15926 .retire(current.generation)?;
15927 }
15928 if successor_valid {
15929 let successor = successor_attempt
15930 .take()
15931 .expect("valid controller successor disappeared before promotion");
15932 let generation = successor.generation;
15933 opti_fork
15934 .as_mut()
15935 .ok_or("optipipe successor promotion lost fork state")?
15936 .promote_successor_snapshot(&mut snap, generation);
15937 carried_opti = Some(successor);
15938 }
15939 if anatomy_on {
15940 // Commit/rollback is normally asynchronous on the primary/head stream. Bound it
15941 // only for this diagnostic so it does not disappear into the following draft's
15942 // first token readback.
15943 e.stream().synchronize()?;
15944 ph_commit += commit_started.elapsed().as_secs_f64();
15945 }
15946 // adaptive-K update (host math, zero syncs): next round drafts accepted-run + 1,
15947 // clamped to [floor(pos), k_cap]. cache.pos is post-rollback here (the round's
15948 // final position — the floor's position key reads the committed depth). Burst
15949 // rounds (`continue` above) draft the captured fixed depth and skip this, exactly
15950 // like gemma's burst arm.
15951 if adapt {
15952 let fl_now = floor_at(cache.pos);
15953 kc = (n_acc + 1).clamp(fl_now.min(k_cap), k_cap);
15954 }
15955 ph_mark(&mut ph_rest, phase_on);
15956 if let Some(p) = pipe {
15957 p.accept_end(round);
15958 }
15959 drop(pipe_accept);
15960 if let Some(t0) = round_t0 {
15961 let ms = t0.elapsed().as_secs_f64() * 1e3;
15962 ROUND_MS.fetch_add((ms * 1e3) as u64, std::sync::atomic::Ordering::Relaxed);
15963 let n = ROUND_N.fetch_add(1, std::sync::atomic::Ordering::Relaxed) + 1;
15964 if n.is_multiple_of(32) {
15965 eprintln!(
15966 "[spec-round] rounds={n} avg round wall={:.2} ms (emitted={} drafted so far)",
15967 ROUND_MS.load(std::sync::atomic::Ordering::Relaxed) as f64 / 1e3 / n as f64,
15968 out.len()
15969 );
15970 }
15971 }
15972 round += 1;
15973 // sse-cadence: this round's accepted drafts + bonus are committed (out is
15974 // append-only past step 4) — flush at round cadence.
15975 keep_going = flush_commit(&mut on_commit, &out, &mut flushed);
15976 }
15977 if let Some(mut ticket) = carried_opti.take() {
15978 opti_fork
15979 .as_mut()
15980 .ok_or("optipipe tail drain lost fork state")?
15981 .cancel_controller_ticket(e, &mut *cache, &mut *scratch, &snap, &mut ticket)?;
15982 }
15983 // sse-cadence: nothing below appends to `out`; flush any remainder (defensive).
15984 // (verdict ignored — the burst is over either way; the session tail runs unchanged.)
15985 let _ = flush_commit(&mut on_commit, &out, &mut flushed);
15986
15987 if spec_stats {
15988 let per_slot: Vec<String> = (0..k)
15989 .map(|j| {
15990 if st_drafted[j] > 0 {
15991 format!(
15992 "{}/{}={:.3}",
15993 st_accepted[j],
15994 st_drafted[j],
15995 st_accepted[j] as f64 / st_drafted[j] as f64
15996 )
15997 } else {
15998 "0/0".into()
15999 }
16000 })
16001 .collect();
16002 let acc = if total_drafted > 0 {
16003 total_accepted as f64 / total_drafted as f64
16004 } else {
16005 0.0
16006 };
16007 eprintln!(
16008 "[spec-stats] rounds={round} full_accept={st_full} len_hist={st_len_hist:?} \
16009 per_slot=[{}] total={total_accepted}/{total_drafted}={acc:.3} \
16010 tok_per_round={:.3}",
16011 per_slot.join(" "),
16012 (total_accepted + round) as f64 / round.max(1) as f64
16013 );
16014 }
16015 if constraint.is_some() {
16016 eprintln!(
16017 "[draft-mask] mask_rounds={dm_rounds} clone_total={:.3}ms \
16018 clone_per_round={:.4}ms gram_cuts={dm_cuts}/{round} cut_tokens={dm_cut_tokens}",
16019 dm_clone_ns as f64 / 1e6,
16020 dm_clone_ns as f64 / 1e6 / dm_rounds.max(1) as f64
16021 );
16022 }
16023 if phase_on {
16024 let tot = ph_draft + ph_verify + ph_wait + ph_rest;
16025 eprintln!(
16026 "[spec-phase] draft={:.1}ms ({:.1}%) verify-issue={:.1}ms ({:.1}%) verify-wait={:.1}ms ({:.1}%) commit-host={:.1}ms ({:.1}%) rounds={round}",
16027 ph_draft * 1e3,
16028 ph_draft / tot * 100.0,
16029 ph_verify * 1e3,
16030 ph_verify / tot * 100.0,
16031 ph_wait * 1e3,
16032 ph_wait / tot * 100.0,
16033 ph_rest * 1e3,
16034 ph_rest / tot * 100.0
16035 );
16036 }
16037 if anatomy_on {
16038 let rounds_f = round.max(1) as f64;
16039 let other = (ph_rest - ph_commit).max(0.0);
16040 eprintln!(
16041 "[spec-anatomy] per-round draft={:.3}ms pp-verify={:.3}ms \
16042 verify-accept={:.3}ms commit-rollback={:.3}ms other={:.3}ms rounds={round}",
16043 ph_draft * 1e3 / rounds_f,
16044 ph_verify * 1e3 / rounds_f,
16045 ph_wait * 1e3 / rounds_f,
16046 ph_commit * 1e3 / rounds_f,
16047 other * 1e3 / rounds_f,
16048 );
16049 }
16050 let _pipe_tail = pipe.map(|p| p.primary()).transpose()?;
16051 // SESSION TAIL: leave the session in the exact invariant the next turn's suffix prime
16052 // expects — every row in `committed` has trunk KV/recur state AND an exact draft-KV row.
16053 // Park the draft-graph ctx back on the session (the serve-burst fixed-cost fix): the next
16054 // burst replays instead of recapturing. Error paths (`?` above) drop it — recaptured then.
16055 if let Some(slot) = sess_draft_slot.take() {
16056 *slot = Some(dctx);
16057 }
16058 let t_rounds = t_ent.elapsed();
16059 if let Some((committed, last_h, next_pred_slot, sctr_slot, uctr_slot)) = sess_tail.take() {
16060 // NEXT BURST'S BOUNDARY TOKEN (lane/sampled-spec-quality, Item 1). Greedy stashes
16061 // the argmax `last_pred` exactly as before (byte contract). SAMPLED draws the token
16062 // HERE, where the sampler, the session Philox counters and the penalty window are
16063 // all live and the boundary logits row still exists — that is the "make the state
16064 // available" half of the fix; the consuming burst then just emits it. `sctr` is
16065 // written to the session BELOW the draws so the advance is never lost.
16066 *next_pred_slot = Some(last_pred);
16067 let sample_boundary = sampled && constraint.is_none() && spec_sampled_boundary_on();
16068 let mut stashed_pending = false;
16069 if let Some(b) = pending.take() {
16070 if !sampled {
16071 // PENDING-CARRY (2026-08-01): stash the bonus on the session instead of
16072 // committing it with a solo T=1 pass — the next empty-suffix greedy burst
16073 // consumes it as round-0 verify col 0 (a plain round edge; the old tail
16074 // commit + next burst's init feed were 11.6+11.5ms solo trunk passes per
16075 // burst on H100 q27, [spec-setup] trace). b stays in `out` (emitted) but
16076 // OUT of `committed` (cache rows == committed); the consuming call
16077 // prepends it once its verify commits the row. next_pred is unknowable
16078 // without the commit pass — None; callers gate on pending_tok too.
16079 debug_assert_eq!(out.last(), Some(&b), "pending must be the last emitted");
16080 if let Some(slot) = sess_pending_slot.take() {
16081 *slot = Some(b);
16082 }
16083 *next_pred_slot = None;
16084 // fill_prev = hidden of the last COMMITTED row (b's predecessor) — the
16085 // exact chain-seed/fill anchor the consuming burst (or a flush) needs.
16086 *last_h = Some(e.clone_dtod(&fill_prev)?);
16087 stashed_pending = true;
16088 } else {
16089 // SAMPLED tail (unchanged): commit the bonus (one T=1 pass) + draft fill —
16090 // the sampled round-0 accept needs this pass's logits (last_col_logits).
16091 let pos_b = cache.pos;
16092 scratch.set_len(e, pos_b)?;
16093 let (lg_b, hb) = self.spec_target_step_h(e, b, &mut *cache)?;
16094 // after a FULL-accept exit `last_pred` is STALE (it predicted the bonus
16095 // itself — the prediction AFTER the bonus never materialized; it would have
16096 // been the next round's verify col 0). The commit's logits ARE that
16097 // prediction — so they are also the row the next burst's boundary token
16098 // comes off, and (lane/sampled-spec-quality) it is DRAWN from them here.
16099 *next_pred_slot = Some(if sample_boundary {
16100 sample_boundary_token(
16101 e,
16102 &lg_b,
16103 &sp,
16104 &pen_hist,
16105 &mut sctr,
16106 "burst-tail-commit",
16107 )?
16108 } else {
16109 argmax(&lg_b) as u32
16110 });
16111 self.mtp_kv_fill_all(e, &[b], &fill_prev, pos_b, &mut *scratch, embd_dev)?;
16112 *last_h = Some(hb);
16113 }
16114 } else {
16115 // fill_prev tracks the hidden of the last COMMITTED row throughout the loop.
16116 *last_h = Some(e.clone_dtod(&fill_prev)?);
16117 if sample_boundary {
16118 // No pending to commit, so the boundary row is the one `last_pred` was
16119 // argmaxed from and the sampled path keeps it on device: the init feed's
16120 // logits when the burst ran zero rounds, else the legacy-replay path's
16121 // last verify column (both predict the token AFTER the last committed
16122 // row). It is retained precisely because round 0's accept test needs it,
16123 // so the draw costs no extra D2H of the [n_vocab] row.
16124 match last_col_logits.as_ref() {
16125 Some(lc) => {
16126 *next_pred_slot = Some(sample_boundary_token_dev(
16127 e,
16128 lc,
16129 n_vocab,
16130 &sp,
16131 &pen_hist,
16132 &mut sctr,
16133 "burst-tail-nopending",
16134 )?);
16135 }
16136 // NAME THE FALLBACK (house standard): unreachable today — a sampled
16137 // burst always feeds or replays, so the row exists — but if it ever
16138 // is, the stream takes a greedy token and SAYS so rather than
16139 // silently regressing to the pre-lane behaviour.
16140 None => eprintln!(
16141 "[spec-boundary] sampled tail kept the ARGMAX boundary token \
16142 (reason: no retained boundary logits row)"
16143 ),
16144 }
16145 }
16146 }
16147 *sctr_slot = sctr;
16148 *uctr_slot = uctr;
16149 committed.extend_from_slice(prompt);
16150 if let Some(cb) = carried_pending {
16151 // the consumed carry's cache row landed in round 0's verify (every pending
16152 // round commits col 0) — it joins `committed` here, in sequence order.
16153 committed.push(cb);
16154 }
16155 if stashed_pending {
16156 // ZERO-EMIT BURST (2026-08-06 c=8 serve panic, pre-existing since b4aea184):
16157 // `out.len() - 1` underflowed on an EMPTY `out` — "range end index
16158 // 18446744073709551615 out of range for slice of length 0", killing the
16159 // memra-gpu-worker and failing 31 of 32 concurrent requests with "worker closed
16160 // stream". Reachable because `pending` starts as `carried_pending` (a bonus
16161 // stashed by the PREVIOUS burst) while `out` starts empty, and the carry is
16162 // deliberately NOT pushed to `out` (line ~3239: the burst that emitted it already
16163 // did). So a burst that stashes a pending without emitting anything of its own —
16164 // the round loop exits before a push, e.g. the ring drain's `out.len() < max_new`
16165 // guard skipping every token under a tight budget — arrives here with
16166 // out.len() == 0 and stashed_pending == true.
16167 //
16168 // The invariant is unchanged: `committed` gets every emitted token EXCEPT the
16169 // stashed bonus. With nothing emitted, that is nothing — and the carry pushed
16170 // just above is already accounted. Saturating, not a min/assert: an empty `out`
16171 // here is a legitimate burst shape, not a corrupt state.
16172 let emitted = out.len().saturating_sub(1);
16173 committed.extend_from_slice(&out[..emitted]);
16174 } else {
16175 committed.extend_from_slice(&out); // FULL out incl. overshoot — all committed
16176 }
16177 debug_assert_eq!(
16178 cache.pos,
16179 committed.len(),
16180 "session invariant: cache rows == committed tokens"
16181 );
16182 if setup_trace {
16183 e.stream().synchronize()?; // bound the async tail fill in the trace
16184 let t_tail = t_ent.elapsed();
16185 eprintln!(
16186 "[spec-setup] init={:.2}ms cap={:.2}ms fill={:.2}ms rounds={:.2}ms tail={:.2}ms total={:.2}ms out={} cont={}",
16187 t_init.as_secs_f64() * 1e3,
16188 (t_cap - t_init).as_secs_f64() * 1e3,
16189 (t_fill - t_cap).as_secs_f64() * 1e3,
16190 (t_rounds - t_fill).as_secs_f64() * 1e3,
16191 (t_tail - t_rounds).as_secs_f64() * 1e3,
16192 t_tail.as_secs_f64() * 1e3,
16193 out.len(),
16194 continuation
16195 );
16196 }
16197 return Ok((out, total_drafted, total_accepted));
16198 }
16199 out.truncate(max_new);
16200 Ok((out, total_drafted, total_accepted))
16201 }
16202
16203 /// Anchor-bounded DSpark target extraction. The trunk sees the exact generated token tape;
16204 /// only requested hidden rows and target-logit rows cross PCIe. An anchor token at p pairs
16205 /// with the pre-output-norm h[p-1] carrier, exactly as the existing replay/NextN path does.
16206 #[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
16207 pub fn extract_dspark_anchors(
16208 &self,
16209 e: &Engine,
16210 tokens: &[u32],
16211 anchor_positions: &[usize],
16212 gamma: usize,
16213 top_k: usize,
16214 chunk: usize,
16215 temperature: f32,
16216 ) -> Result<Vec<DsparkAnchorRecord>, Box<dyn std::error::Error>> {
16217 if tokens.len() < gamma + 2 || gamma == 0 || chunk < 2 {
16218 return Err("DSpark extraction token tape/gamma/chunk is invalid".into());
16219 }
16220 if anchor_positions.windows(2).any(|pair| pair[0] >= pair[1]) {
16221 return Err("DSpark anchor positions must be sorted and unique".into());
16222 }
16223 for &position in anchor_positions {
16224 if position == 0 || position + gamma >= tokens.len() {
16225 return Err(format!(
16226 "DSpark anchor {position} has no predecessor or cannot cover gamma={gamma} in {} tokens",
16227 tokens.len()
16228 )
16229 .into());
16230 }
16231 }
16232
16233 let n_vocab = self.output.out_features();
16234 let n_embd = self.cfg.n_embd as usize;
16235 let mut cache =
16236 crate::pp::new_cache_planned(e, &self.cfg, &self.plan, tokens.len() + gamma + 8)?;
16237 let (embd_qt, embd_rb) = self.embd.qt_and_row_bytes(n_embd);
16238 let embd_gpu = if spec_host_embd() {
16239 None
16240 } else {
16241 Some(
16242 self.embd_gpu
16243 .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload")),
16244 )
16245 };
16246 let embd_dev = embd_gpu.map(|gpu| (gpu, embd_qt, embd_rb));
16247
16248 struct PendingRecord {
16249 position: usize,
16250 hidden: Option<Vec<f32>>,
16251 tokens: Vec<u32>,
16252 target_top_ids: Vec<Option<Vec<u32>>>,
16253 target_top_logits: Vec<Option<Vec<f32>>>,
16254 target_top_probs: Vec<Option<Vec<f32>>>,
16255 target_tail_probs: Vec<Option<f32>>,
16256 }
16257
16258 let mut pending: Vec<PendingRecord> = anchor_positions
16259 .iter()
16260 .map(|&position| PendingRecord {
16261 position,
16262 hidden: None,
16263 tokens: tokens[position..=position + gamma].to_vec(),
16264 target_top_ids: vec![None; gamma],
16265 target_top_logits: vec![None; gamma],
16266 target_top_probs: vec![None; gamma],
16267 target_tail_probs: vec![None; gamma],
16268 })
16269 .collect();
16270
16271 let mut start = 0usize;
16272 while start < tokens.len() {
16273 let end = (start + chunk).min(tokens.len());
16274 let chunk_tokens = &tokens[start..end];
16275 let (target_logits, hidden_rows) =
16276 self.decode_step_t_core(e, chunk_tokens, start, &mut cache, embd_dev, None)?;
16277 for record in &mut pending {
16278 let hidden_position = record.position - 1;
16279 if hidden_position >= start && hidden_position < end {
16280 let local = hidden_position - start;
16281 record.hidden = Some(
16282 e.dtoh_view(&hidden_rows.slice(local * n_embd..(local + 1) * n_embd))?,
16283 );
16284 }
16285 for slot in 0..gamma {
16286 let target_row = record.position + slot;
16287 if target_row < start || target_row >= end {
16288 continue;
16289 }
16290 let local = target_row - start;
16291 let logits =
16292 e.dtoh_view(&target_logits.slice(local * n_vocab..(local + 1) * n_vocab))?;
16293 let (ids, top_logits, probs, tail) =
16294 dspark_sparse_softmax_topk(&logits, top_k, temperature)?;
16295 record.target_top_ids[slot] = Some(ids);
16296 record.target_top_logits[slot] = Some(top_logits);
16297 record.target_top_probs[slot] = Some(probs);
16298 record.target_tail_probs[slot] = Some(tail);
16299 }
16300 }
16301 start = end;
16302 }
16303
16304 pending
16305 .into_iter()
16306 .map(|record| {
16307 let hidden = record
16308 .hidden
16309 .ok_or_else(|| format!("missing DSpark hidden at {}", record.position))?;
16310 let target_top_ids =
16311 flatten_dspark_rows(record.target_top_ids, record.position, "target ids")?;
16312 let target_top_logits = flatten_dspark_rows(
16313 record.target_top_logits,
16314 record.position,
16315 "target logits",
16316 )?;
16317 let target_top_probs =
16318 flatten_dspark_rows(record.target_top_probs, record.position, "target probs")?;
16319 let target_tail_probs = record
16320 .target_tail_probs
16321 .into_iter()
16322 .enumerate()
16323 .map(|(slot, value)| {
16324 value.ok_or_else(|| {
16325 format!("missing DSpark tail at {} slot {slot}", record.position)
16326 })
16327 })
16328 .collect::<Result<Vec<_>, _>>()?;
16329 Ok(DsparkAnchorRecord {
16330 position: record.position,
16331 hidden,
16332 tokens: record.tokens,
16333 target_top_ids,
16334 target_top_logits,
16335 target_top_probs,
16336 target_tail_probs,
16337 })
16338 })
16339 .collect()
16340 }
16341
16342 /// TEACHER-FORCED REPLAY ACCEPTANCE (hqmtp MTP-heal protocol): walk a FIXED token
16343 /// sequence and, at sampled positions, compare the MTP head's K-token draft chain against
16344 /// the trunk's own teacher-forced greedy predictions. Nothing is generated — the context is
16345 /// the corpus text itself, so (a) degenerate self-generated loops cannot inflate acceptance
16346 /// and (b) two arms (bf16 ceiling vs NVFP4) score on IDENTICAL contexts, isolating the
16347 /// quant-induced head/hidden-state mismatch from text drift.
16348 ///
16349 /// Per eval position p (context = tokens[0..=p], predecessor pairing as in spec decode):
16350 /// draft_j = chain token j from (tokens[p], h_{p-1}), then its own drafts — the exact
16351 /// eager spec-decode chain (same mtp_head_forward_dev, same rope positions).
16352 /// target_j = teacher-forced greedy pick for position p+1+j (argmax of the trunk logits
16353 /// at forced context tokens[0..p+j]). For j==0 this equals live spec
16354 /// acceptance; for j>=1 live verify would condition on the drafts, here it
16355 /// conditions on the corpus — deterministic and arm-comparable by design.
16356 ///
16357 /// Returns (rows, bg): one (p, drafts[k], targets[k]) row per eval position (ascending p),
16358 /// plus the full teacher-forced greedy track bg (bg[i] = greedy pick for position i, i>=1)
16359 /// so harnesses can cross-check runs (e.g. different chunk sizes must give identical bg).
16360 ///
16361 /// `hdump`: when Some, every position's pre-output_norm trunk hidden (the exact rows the
16362 /// draft-KV fill pairs from) streams to the file as little-endian f32 [t_total, n_embd] —
16363 /// the head-distillation extraction (hqmtp): the ENGINE is the source of truth for trunk
16364 /// hiddens (HF torch reproductions of the hybrid trunk measured only ~0.5 greedy
16365 /// agreement vs this path — not usable as a training-data source).
16366 #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
16367 pub fn replay_acceptance(
16368 &self,
16369 e: &Engine,
16370 tokens: &[u32],
16371 k: usize,
16372 stride: usize,
16373 chunk: usize,
16374 mut hdump: Option<&mut std::fs::File>,
16375 ) -> Result<(Vec<(usize, Vec<u32>, Vec<u32>)>, Vec<u32>), Box<dyn std::error::Error>> {
16376 assert!(k >= 1 && stride >= 1 && chunk >= 2);
16377 let mtp = self
16378 .mtp
16379 .as_ref()
16380 .expect("replay_acceptance requires an MTP head");
16381 let n_vocab = self.output.out_features();
16382 let d_vocab = mtp
16383 .shared_head_head
16384 .as_ref()
16385 .unwrap_or(&self.output)
16386 .out_features();
16387 let n_embd = self.cfg.n_embd as usize;
16388 let t_total = tokens.len();
16389 assert!(t_total >= 8, "corpus too short ({t_total} tokens)");
16390 // STAGE-OWNED KV (lane/pp2-spec 2026-08-06) — see `new_session`. Door shut = `Cache::new`.
16391 let mut cache = crate::pp::new_cache_planned(e, &self.cfg, &self.plan, t_total + k + 8)?;
16392 let mut scratch = self.new_mtp_scratch(e, t_total + k + 8)?;
16393 let (embd_qt, embd_rb) = self.embd.qt_and_row_bytes(n_embd);
16394 let embd_gpu = if spec_host_embd() {
16395 None
16396 } else {
16397 Some(
16398 self.embd_gpu
16399 .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload")),
16400 )
16401 };
16402 let embd_dev = embd_gpu.map(|g| (g, embd_qt, embd_rb));
16403
16404 // bg[i] = the trunk's greedy pick for position i under the forced context (i >= 1).
16405 let mut bg: Vec<u32> = vec![0; t_total + 1];
16406 let mut rows: Vec<(usize, Vec<u32>, Vec<u32>)> = Vec::new();
16407 let mut prev_last_h = e.zeros(n_embd)?; // predecessor hidden entering the chunk
16408 let mut seed_buf = e.zeros(n_embd)?;
16409 let mut preds_d = e.alloc_u32_zeroed(chunk)?;
16410 let nll_on = std::env::var("MEMRA_REPLAY_NLL").as_deref() == Ok("1");
16411 let (mut nll_sum, mut nll_cnt) = (0f64, 0u64);
16412 let mut s = 0usize;
16413 while s < t_total {
16414 let cend = (s + chunk).min(t_total);
16415 let tc = cend - s;
16416 let ch = &tokens[s..cend];
16417 // 1. forced trunk pass — verify path (decode-exact contract): all-column logits +
16418 // the chunk's true hiddens.
16419 let (tl_d, vx) = self.decode_step_t_core(e, ch, s, &mut cache, embd_dev, None)?;
16420 for j in 0..tc {
16421 e.argmax_token_device_col(&tl_d, j, n_vocab, &mut preds_d, j)?;
16422 }
16423 let preds = e.dtoh_u32(&preds_d)?;
16424 for j in 0..tc {
16425 bg[s + j + 1] = preds[j];
16426 }
16427 // MEMRA_REPLAY_NLL=1: teacher-forced NLL/perplexity over the same forced pass — the
16428 // checkpoint-quality metric (position j's logits score the GOLD next token).
16429 if nll_on {
16430 let jmax = if cend < t_total { tc } else { tc - 1 }; // last pos has no gold next
16431 if jmax > 0 {
16432 let ids: Vec<u32> = (0..jmax).map(|j| tokens[s + j + 1]).collect();
16433 let rows: Vec<i32> = (0..jmax as i32).collect();
16434 let idsd = e.htod_u32_v(&ids)?;
16435 let rowsd = e.htod_i32(&rows)?;
16436 let mut outd = e.zeros(jmax)?;
16437 e.softmax_gather(&tl_d, n_vocab, &idsd, &rowsd, &mut outd, n_vocab, jmax, 1.0)?;
16438 for pr in e.dtoh(&outd)? {
16439 nll_sum += -((pr.max(1e-30)) as f64).ln();
16440 nll_cnt += 1;
16441 }
16442 }
16443 }
16444 if let Some(f) = hdump.as_deref_mut() {
16445 use std::io::Write;
16446 let host: Vec<f32> = e.dtoh(&vx)?;
16447 // bf16 round-to-nearest-even — f32 doubled the disk bill at bulk
16448 // extraction scale (20M tokens x 4096 = 320GB f32 vs 160GB bf16).
16449 let mut bytes = Vec::with_capacity(tc * n_embd * 2);
16450 for v in &host[..tc * n_embd] {
16451 let b = v.to_bits();
16452 let r = b.wrapping_add(0x7FFF + ((b >> 16) & 1));
16453 bytes.extend_from_slice(&((r >> 16) as u16).to_le_bytes());
16454 }
16455 f.write_all(&bytes)?;
16456 }
16457 // CHAINLESS extraction (stride > corpus, the bulk-hdump mode): no chunk ever
16458 // drafts, so the draft-KV fills are pure waste — skip them (2 MTP-block passes
16459 // per token saved; the forced trunk pass + hdump is all the mode needs).
16460 let chainless = stride > t_total;
16461 if chainless {
16462 e.copy_view_into(
16463 &mut prev_last_h,
16464 0,
16465 &vx.slice((tc - 1) * n_embd..tc * n_embd),
16466 n_embd,
16467 )?;
16468 s = cend;
16469 continue;
16470 }
16471 // 2. TRUE predecessor-paired draft-KV fill for the chunk (row i carries h_{i-1};
16472 // row s reads the previous chunk's last true hidden, zeros at corpus start).
16473 let mut vxs = e.zeros(tc * n_embd)?;
16474 e.copy_into(&mut vxs, 0, &prev_last_h, n_embd)?;
16475 if tc > 1 {
16476 e.copy_view_into(
16477 &mut vxs,
16478 n_embd,
16479 &vx.slice(0..(tc - 1) * n_embd),
16480 (tc - 1) * n_embd,
16481 )?;
16482 }
16483 scratch.set_len(e, s)?;
16484 self.mtp_kv_fill_all(e, ch, &vxs, s, &mut scratch, embd_dev)?;
16485 // 3. draft chains at sampled positions, DESCENDING: a chain reads only slots
16486 // [0..p) (true fills) and appends at >= p; the next (smaller-p) chain's set_len
16487 // truncates those approximate appends before they can ever be read.
16488 let ps: Vec<usize> = (s..cend)
16489 .filter(|p| *p >= 1 && *p % stride == 0 && *p + k <= t_total)
16490 .collect();
16491 for &p in ps.iter().rev() {
16492 scratch.set_len(e, p)?;
16493 if p == s {
16494 e.copy_into(&mut seed_buf, 0, &prev_last_h, n_embd)?;
16495 } else {
16496 e.copy_view_into(
16497 &mut seed_buf,
16498 0,
16499 &vx.slice((p - 1 - s) * n_embd..(p - s) * n_embd),
16500 n_embd,
16501 )?;
16502 }
16503 let mut e_tok = tokens[p];
16504 let mut d_seed = e.clone_dtod(&seed_buf)?;
16505 let chain_heads = !self.mtp_extra.is_empty();
16506 let mut chain_tokens = if chain_heads {
16507 vec![tokens[p]]
16508 } else {
16509 Vec::new()
16510 };
16511 let mut chain_seeds = if chain_heads {
16512 vec![e.clone_dtod(&seed_buf)?]
16513 } else {
16514 Vec::new()
16515 };
16516 let mut drafts: Vec<u32> = Vec::with_capacity(k);
16517 for j in 0..k {
16518 let (dl_d, h_nextn) = if chain_heads {
16519 self.mtp_chain_forward_dev(
16520 e,
16521 &chain_tokens,
16522 &chain_seeds,
16523 &mut scratch,
16524 p,
16525 embd_dev,
16526 None,
16527 )?
16528 } else {
16529 self.mtp_head_forward_dev(
16530 e,
16531 mtp,
16532 e_tok,
16533 &d_seed,
16534 &mut scratch,
16535 p + 1 + j,
16536 embd_dev,
16537 None,
16538 )?
16539 };
16540 let tok_d = e.argmax_token_device(&dl_d, d_vocab)?;
16541 let idx = e.dtoh_u32_one(&tok_d)?;
16542 let d = match &mtp.d2t {
16543 Some(map) => map[idx as usize],
16544 None => idx,
16545 };
16546 drafts.push(d);
16547 if chain_heads {
16548 chain_tokens.push(d);
16549 chain_seeds.push(h_nextn);
16550 } else {
16551 e_tok = d;
16552 d_seed = h_nextn;
16553 }
16554 }
16555 // targets may live in a LATER chunk's bg — resolved after the walk.
16556 rows.push((p, drafts, Vec::new()));
16557 }
16558 // 4. restore TRUE entries for the whole chunk (the next chunk's chains and fills
16559 // expect scratch.len == cend with exact rows).
16560 scratch.set_len(e, s)?;
16561 self.mtp_kv_fill_all(e, ch, &vxs, s, &mut scratch, embd_dev)?;
16562 e.copy_view_into(
16563 &mut prev_last_h,
16564 0,
16565 &vx.slice((tc - 1) * n_embd..tc * n_embd),
16566 n_embd,
16567 )?;
16568 s = cend;
16569 }
16570 for (p, drafts, targets) in rows.iter_mut() {
16571 for j in 0..drafts.len() {
16572 targets.push(bg[*p + 1 + j]);
16573 }
16574 }
16575 rows.sort_by_key(|r| r.0);
16576 if nll_cnt > 0 {
16577 let mean = nll_sum / nll_cnt as f64;
16578 println!(
16579 "[replay-nll] tokens={nll_cnt} nll/token={mean:.5} ppl={:.4}",
16580 mean.exp()
16581 );
16582 }
16583 Ok((rows, bg))
16584 }
16585}
16586
16587#[cfg(test)]
16588mod vg_debt_tests {
16589 use super::dspark_vg_debt_projection;
16590
16591 /// TOOTH for the verify-graph admission accounting: the pool's projected remaining
16592 /// growth must be charged (pre-fix, admission charged 0 for a pool measured at
16593 /// 8,852 MiB), the projection must price the MARGINAL cost of one more key rather than
16594 /// extrapolating the pool's one-time shared allocation, and the doors that make growth
16595 /// impossible must zero the debt.
16596 #[test]
16597 fn vg_debt_projects_remaining_growth_and_respects_the_freeze_valves() {
16598 const MIB: usize = 1 << 20;
16599 let d = dspark_vg_debt_projection;
16600 // cold pool: nothing observed, one capture fits inside SPEC_SHRINK_RESERVE.
16601 assert_eq!(d(0, 256, 0, None), 0);
16602 // freeze valve MEMRA_DSPARK_VG_MAX=0: the pool cannot grow.
16603 assert_eq!(d(10, 0, 500 * MIB, None), 0);
16604 // saturated pool: at/past the cap the pool FREEZES, nothing left to reserve.
16605 assert_eq!(d(256, 256, 8852 * MIB, None), 0);
16606 assert_eq!(d(300, 256, 8852 * MIB, None), 0);
16607
16608 // BOOTSTRAP (one observation, growth unmeasurable): at most one more pool's worth.
16609 // The pre-fix mean rule extrapolated 255x here — the measured 8.5 GB phantom.
16610 assert_eq!(d(1, 256, 33 * MIB, None), 33 * MIB);
16611
16612 // MARGINAL, flat pool (the box9 receipt: reserved stayed ~33.6 MiB across captures
16613 // 1..3, so an additional key costs ~nothing and the debt must collapse to ~0 —
16614 // NOT the 8,556/4,261/2,830 MB the mean rule printed).
16615 assert_eq!(d(3, 256, 33 * MIB, Some((1, 33 * MIB))), 0);
16616
16617 // MARGINAL, genuinely growing pool: 40 MiB per new key over 2 keys, 250 slots left.
16618 let debt = d(6, 256, 273 * MIB, Some((4, 193 * MIB)));
16619 assert_eq!(debt, 250 * (40 * MIB));
16620 assert!(
16621 debt > 3 * (1536 * MIB),
16622 "real growth must dwarf SPEC_SHRINK_RESERVE"
16623 );
16624
16625 // a shrinking/recycled reading never becomes a negative charge.
16626 assert_eq!(d(6, 256, 10 * MIB, Some((4, 99 * MIB))), 0);
16627 // a stale observation at the same capture count falls back to bootstrap.
16628 assert_eq!(d(4, 256, 80 * MIB, Some((4, 80 * MIB))), 80 * MIB);
16629 }
16630}
16631
16632#[cfg(test)]
16633mod capture_headroom_tests {
16634 use super::{
16635 CAPTURE_HEADROOM_FLOOR, capture_err_is_oom, capture_headroom_verdict,
16636 draft_capture_bootstrap_estimate,
16637 };
16638
16639 /// TOOTH for the pre-capture reserve check (lane/step37-vram-admission-20260830): a
16640 /// capture attempt must be refused BEFORE it allocates when the device cannot cover its
16641 /// appetite plus the post-capture floor — and pool-cached bytes count as headroom
16642 /// (driver `free` alone under-counts, the wrong direction for a gate that drops
16643 /// coverage).
16644 #[test]
16645 fn capture_reserve_check_refuses_short_devices_and_counts_pool_cache() {
16646 const MIB: usize = 1 << 20;
16647 let need = 900 * MIB;
16648 // Plenty of room: no refusal.
16649 assert_eq!(
16650 capture_headroom_verdict(8_000 * MIB, 0, need, CAPTURE_HEADROOM_FLOOR),
16651 None
16652 );
16653 // The owner's shape: capture appetite would walk the card to the edge — refused,
16654 // with the arithmetic surfaced for the WARN line.
16655 let (required, effective) =
16656 capture_headroom_verdict(1_200 * MIB, 0, need, CAPTURE_HEADROOM_FLOOR)
16657 .expect("short device must refuse");
16658 assert_eq!(required, need + CAPTURE_HEADROOM_FLOOR);
16659 assert_eq!(effective, 1_200 * MIB);
16660 // Pool-cached bytes are real headroom (the trim path makes them driver-visible).
16661 assert_eq!(
16662 capture_headroom_verdict(1_200 * MIB, 7_000 * MIB, need, CAPTURE_HEADROOM_FLOOR),
16663 None
16664 );
16665 // Boundary: exactly enough is enough (>=, never a fencepost refusal).
16666 assert_eq!(
16667 capture_headroom_verdict(
16668 need + CAPTURE_HEADROOM_FLOOR,
16669 0,
16670 need,
16671 CAPTURE_HEADROOM_FLOOR
16672 ),
16673 None
16674 );
16675 // POLICY at the call site (owner-shape receipts, escalated twice on-box): the
16676 // refusal fn is handed 2x the appetite plus TWO floors — a capture may take at
16677 // most half the discretionary headroom, so the card retains a whole capture's
16678 // worth of room after it lands. One floor of slack above one appetite (the shape
16679 // that step-OOM'd on the owner cell) must therefore REFUSE under the call-site
16680 // requirement.
16681 assert!(
16682 capture_headroom_verdict(
16683 need + CAPTURE_HEADROOM_FLOOR + (100 << 20),
16684 0,
16685 2 * need,
16686 CAPTURE_HEADROOM_FLOOR * 2
16687 )
16688 .is_some()
16689 );
16690 }
16691
16692 #[test]
16693 fn bootstrap_estimate_scales_with_heads_and_never_underflows() {
16694 // 3-head chain on a step37-shaped vocab must expect strictly more than one head.
16695 let one = draft_capture_bootstrap_estimate(1, 3, 128_896, 4_096);
16696 let three = draft_capture_bootstrap_estimate(3, 3, 128_896, 4_096);
16697 assert!(three > one);
16698 // Degenerate shapes keep a sane minimum (the estimate feeds a refusal gate; a
16699 // zero-need gate refuses nothing).
16700 assert!(draft_capture_bootstrap_estimate(0, 0, 0, 0) >= 64 << 20);
16701 }
16702
16703 #[test]
16704 fn capture_oom_predicate_matches_the_quoted_driver_text() {
16705 assert!(capture_err_is_oom(
16706 "DriverError(CUDA_ERROR_OUT_OF_MEMORY, \"out of memory\")"
16707 ));
16708 assert!(capture_err_is_oom("allocation failed: out of memory"));
16709 assert!(!capture_err_is_oom("capture produced no graph"));
16710 }
16711}
16712
16713#[cfg(test)]
16714mod mtp_chain_tests {
16715 use super::mtp_chain_head_index;
16716
16717 #[test]
16718 fn embedded_step_heads_cycle_in_declared_order() {
16719 let actual: Vec<usize> = (0..8).map(|step| mtp_chain_head_index(step, 3)).collect();
16720 assert_eq!(actual, [0, 1, 2, 0, 1, 2, 0, 1]);
16721 }
16722
16723 #[test]
16724 fn standalone_draft_remains_single_head() {
16725 assert!((0..8).all(|step| mtp_chain_head_index(step, 1) == 0));
16726 }
16727}
16728
16729#[cfg(test)]
16730mod tp_verified_prefix_tests {
16731 use super::validate_tp_kv_snapshot_shape;
16732 use crate::tp::ResidentTpKvCache;
16733
16734 #[test]
16735 fn snapshot_shape_accepts_matching_tp_presence() {
16736 let layers = vec![
16737 Some(ResidentTpKvCache::new(Vec::new(), 1, 1, 1, 1, 8)),
16738 None,
16739 ];
16740 validate_tp_kv_snapshot_shape(&layers, &[Some(2), None]).unwrap();
16741 }
16742
16743 #[test]
16744 fn snapshot_shape_rejects_changed_tp_presence() {
16745 let layers = vec![Some(ResidentTpKvCache::new(Vec::new(), 1, 1, 1, 1, 8))];
16746 let error = validate_tp_kv_snapshot_shape(&layers, &[None])
16747 .unwrap_err()
16748 .to_string();
16749 assert!(error.contains("changed shape"), "unexpected error: {error}");
16750 }
16751}
16752
16753#[cfg(test)]
16754mod dspark_sparse_tests {
16755 use super::dspark_sparse_softmax_topk;
16756
16757 #[test]
16758 fn topk_keeps_full_softmax_mass_and_stable_ties() {
16759 let logits = [1.0f32, 3.0, 3.0, -2.0];
16760 let (ids, top_logits, probs, tail) = dspark_sparse_softmax_topk(&logits, 2, 1.0).unwrap();
16761 assert_eq!(ids, vec![1, 2]);
16762 assert_eq!(top_logits, vec![3.0, 3.0]);
16763 let denominator = logits.iter().map(|value| (value - 3.0).exp()).sum::<f32>();
16764 let expected = 1.0 / denominator;
16765 assert!((probs[0] - expected).abs() < 1.0e-6);
16766 assert!((probs[1] - expected).abs() < 1.0e-6);
16767 assert!((tail - (1.0 - 2.0 * expected)).abs() < 1.0e-6);
16768 assert!((probs.iter().sum::<f32>() + tail - 1.0).abs() < 1.0e-6);
16769 }
16770}
16771
16772#[cfg(test)]
16773mod spec_replay_env_tests {
16774 use super::spec_replay_env_on;
16775
16776 #[test]
16777 fn replay_requires_literal_one() {
16778 assert!(!spec_replay_env_on(None));
16779 assert!(!spec_replay_env_on(Some("")));
16780 assert!(!spec_replay_env_on(Some("0")));
16781 assert!(!spec_replay_env_on(Some("true")));
16782 assert!(!spec_replay_env_on(Some("2")));
16783 assert!(spec_replay_env_on(Some("1")));
16784 }
16785}
16786
16787#[cfg(test)]
16788mod telem_tests {
16789 use super::{SPEC_TELEM_POS, SpecTelemetry, SpecTelemetryCounters};
16790
16791 #[test]
16792 fn synthetic_accept_masks_produce_tau_and_position_histogram() {
16793 let counters = SpecTelemetryCounters::default();
16794 for mask in [
16795 [true, true, true],
16796 [true, true, false],
16797 [true, false, false],
16798 [false, false, false],
16799 ] {
16800 let accepted = mask.iter().take_while(|&&value| value).count();
16801 counters.record_round(mask.len(), accepted);
16802 }
16803
16804 let snapshot = counters.snapshot();
16805 assert_eq!(
16806 (snapshot.rounds, snapshot.drafted, snapshot.accepted),
16807 (4, 12, 6)
16808 );
16809 assert_eq!(&snapshot.pos_drafted[..3], &[4, 4, 4]);
16810 assert_eq!(&snapshot.pos_accepted[..3], &[3, 2, 1]);
16811 assert_eq!(snapshot.tau(), 1.5);
16812 assert_eq!(snapshot.pos_drafted[3..], [0; SPEC_TELEM_POS - 3]);
16813 assert_eq!(snapshot.pos_accepted[3..], [0; SPEC_TELEM_POS - 3]);
16814 }
16815
16816 /// The worker's per-burst pattern: stash, accumulate, diff — the delta must isolate
16817 /// exactly the burst's contribution (pool-resumed sessions carry prior requests' counts).
16818 #[test]
16819 fn delta_isolates_burst_contribution() {
16820 let mut t = SpecTelemetry::default();
16821 // "previous request": 2 rounds of k=3, accepts 3 then 1.
16822 for (kr, na) in [(3usize, 3usize), (3, 1)] {
16823 t.rounds += 1;
16824 t.drafted += kr as u64;
16825 t.accepted += na as u64;
16826 for j in 0..kr {
16827 t.pos_drafted[j] += 1;
16828 }
16829 for j in 0..na {
16830 t.pos_accepted[j] += 1;
16831 }
16832 }
16833 let before = t;
16834 // "this burst": 1 round k=3, accepts 2.
16835 t.rounds += 1;
16836 t.drafted += 3;
16837 t.accepted += 2;
16838 for j in 0..3 {
16839 t.pos_drafted[j] += 1;
16840 }
16841 for j in 0..2 {
16842 t.pos_accepted[j] += 1;
16843 }
16844 let d = t.delta_since(&before);
16845 assert_eq!((d.rounds, d.drafted, d.accepted), (1, 3, 2));
16846 assert_eq!(&d.pos_drafted[..3], &[1, 1, 1]);
16847 assert_eq!(&d.pos_accepted[..3], &[1, 1, 0]);
16848 assert_eq!(d.pos_drafted[3..], [0; SPEC_TELEM_POS - 3]);
16849 }
16850
16851 /// merge(delta) then merge(delta2) equals accumulating both — the per-model /metrics
16852 /// aggregation invariant.
16853 #[test]
16854 fn merge_accumulates_fieldwise() {
16855 let mut agg = SpecTelemetry::default();
16856 let mut d1 = SpecTelemetry {
16857 rounds: 2,
16858 drafted: 6,
16859 accepted: 4,
16860 ..Default::default()
16861 };
16862 d1.pos_drafted[0] = 2;
16863 d1.pos_accepted[0] = 2;
16864 let mut d2 = SpecTelemetry {
16865 rounds: 1,
16866 drafted: 3,
16867 accepted: 1,
16868 ..Default::default()
16869 };
16870 d2.pos_drafted[0] = 1;
16871 d2.pos_accepted[0] = 1;
16872 d2.pos_drafted[1] = 1;
16873 agg.merge(&d1);
16874 agg.merge(&d2);
16875 assert_eq!((agg.rounds, agg.drafted, agg.accepted), (3, 9, 5));
16876 assert_eq!(agg.pos_drafted[0], 3);
16877 assert_eq!(agg.pos_accepted[0], 3);
16878 assert_eq!(agg.pos_drafted[1], 1);
16879 assert_eq!(agg.pos_accepted[1], 0);
16880 }
16881
16882 /// Wrong-snapshot diff saturates to zero instead of wrapping — the counters feed a
16883 /// public metrics surface and must never publish a u64-wrapped garbage value.
16884 #[test]
16885 fn delta_saturates_never_wraps() {
16886 let small = SpecTelemetry {
16887 rounds: 1,
16888 drafted: 2,
16889 accepted: 1,
16890 ..Default::default()
16891 };
16892 let big = SpecTelemetry {
16893 rounds: 5,
16894 drafted: 15,
16895 accepted: 9,
16896 ..Default::default()
16897 };
16898 let d = small.delta_since(&big);
16899 assert_eq!((d.rounds, d.drafted, d.accepted), (0, 0, 0));
16900 }
16901}
16902
16903#[cfg(test)]
16904mod opti_fork_tests {
16905 use super::{
16906 OptiControllerPolicy, OptiForkAction, OptiForkGateMode, OptiForkGenerationTracker,
16907 };
16908
16909 #[test]
16910 fn controller_threshold_and_three_miss_breaker_are_exact() {
16911 let mut policy = OptiControllerPolicy {
16912 threshold: 0.7,
16913 consecutive_misses: 0,
16914 breaker_tripped: false,
16915 };
16916 assert!(!policy.admit(0.699_999));
16917 assert!(policy.admit(0.7));
16918 assert!(!policy.resolve(false));
16919 assert!(!policy.resolve(false));
16920 assert!(policy.resolve(false));
16921 assert!(policy.breaker_tripped);
16922 assert!(!policy.admit(1.0));
16923 assert!(
16924 !policy.resolve(true),
16925 "a resolved hit cannot re-arm a tripped request"
16926 );
16927 assert!(policy.breaker_tripped);
16928 }
16929
16930 #[test]
16931 fn zero_threshold_is_the_true_unconditional_measurement_arm() {
16932 let mut policy = OptiControllerPolicy {
16933 threshold: 0.0,
16934 consecutive_misses: 0,
16935 breaker_tripped: false,
16936 };
16937 for _ in 0..16 {
16938 assert!(policy.admit(0.0));
16939 assert!(!policy.resolve(false));
16940 }
16941 for invalid in [f32::NAN, f32::INFINITY, -0.01, 1.01] {
16942 assert!(
16943 !policy.admit(invalid),
16944 "invalid q proxy must fail closed: {invalid}"
16945 );
16946 }
16947 assert!(!policy.breaker_tripped);
16948 assert_eq!(policy.consecutive_misses, 0);
16949 }
16950
16951 #[test]
16952 fn alternating_mode_flips_by_generation_not_round_parity() {
16953 assert_eq!(OptiForkGateMode::Alternate.action(0), OptiForkAction::Hit);
16954 assert_eq!(OptiForkGateMode::Alternate.action(1), OptiForkAction::Miss);
16955 assert_eq!(OptiForkGateMode::Alternate.action(8), OptiForkAction::Hit);
16956 assert_eq!(OptiForkGateMode::Alternate.action(9), OptiForkAction::Miss);
16957 }
16958
16959 #[test]
16960 fn live_generation_cannot_be_overwritten() {
16961 let mut tracker = OptiForkGenerationTracker::default();
16962 let g0 = tracker.reserve().unwrap();
16963 let g1 = tracker.reserve().unwrap();
16964 let err = tracker.reserve().unwrap_err().to_string();
16965 assert!(
16966 err.contains("still owns generation 0"),
16967 "unexpected error: {err}"
16968 );
16969 tracker.retire(g0).unwrap();
16970 let g2 = tracker.reserve().unwrap();
16971 assert_eq!((g2.id, g2.slot), (2, 0));
16972 tracker.retire(g1).unwrap();
16973 tracker.retire(g2).unwrap();
16974 }
16975
16976 #[test]
16977 fn teardown_rejects_a_stale_generation_tag() {
16978 let mut tracker = OptiForkGenerationTracker::default();
16979 let g0 = tracker.reserve().unwrap();
16980 tracker.retire(g0).unwrap();
16981 let err = tracker.retire(g0).unwrap_err().to_string();
16982 assert!(err.contains("teardown mismatch"), "unexpected error: {err}");
16983 }
16984}
16985
16986#[cfg(test)]
16987mod draft_graph_fallback_tests {
16988 use super::DraftGraphFallback;
16989
16990 /// The Q2 contract, part (a): a fallback flip is LOUD — exactly once per flip.
16991 #[test]
16992 fn flip_is_loud_once_and_memoized_after() {
16993 let mut f = DraftGraphFallback::default();
16994 let line = f
16995 .mark_greedy("out of memory")
16996 .expect("first flip must return the warn line");
16997 assert!(
16998 line.contains("WARN"),
16999 "flip line must be warn-level: {line}"
17000 );
17001 assert!(
17002 line.contains("out of memory"),
17003 "flip line must carry the reason: {line}"
17004 );
17005 assert!(f.greedy_failed());
17006 // re-marking an already-failed graph is the memoization: quiet, still failed.
17007 assert!(f.mark_greedy("out of memory").is_none());
17008 assert!(f.greedy_failed());
17009 // the two graphs' flags are independent (greedy flip leaves sampled capturable).
17010 assert!(!f.sampled_failed());
17011 let line_s = f
17012 .mark_sampled("capture unsupported")
17013 .expect("sampled flip is its own flip");
17014 assert!(
17015 line_s.contains("sampled"),
17016 "sampled flip names itself: {line_s}"
17017 );
17018 assert!(f.mark_sampled("capture unsupported").is_none());
17019 }
17020
17021 /// The Q2 contract, part (b): resume-from-pool RESETS both flags (fresh capture chance),
17022 /// and says so exactly when there was something to reset.
17023 #[test]
17024 fn reset_on_resume_clears_flags_and_logs_once() {
17025 let mut f = DraftGraphFallback::default();
17026 // clean session: resume is silent, nothing to reset.
17027 assert!(f.reset_on_resume().is_none());
17028 f.mark_greedy("oom").unwrap();
17029 f.mark_sampled("oom").unwrap();
17030 let note = f
17031 .reset_on_resume()
17032 .expect("a set flag must produce the reset note");
17033 assert!(
17034 note.contains("greedy+sampled"),
17035 "note names what was reset: {note}"
17036 );
17037 assert!(
17038 !f.greedy_failed() && !f.sampled_failed(),
17039 "both flags cleared"
17040 );
17041 // and the NEXT failure after a reset is a fresh flip — loud again.
17042 assert!(f.mark_greedy("oom again").is_some());
17043 let note2 = f.reset_on_resume().expect("greedy-only reset");
17044 assert!(note2.contains("(greedy)"), "single-flag note: {note2}");
17045 }
17046
17047 /// Shape-change clears (dmask realloc / mask-shape mismatch / s_key change) stay silent —
17048 /// they precede a fresh capture attempt whose own failure re-flips loudly.
17049 #[test]
17050 fn shape_change_clears_are_silent() {
17051 let mut f = DraftGraphFallback::default();
17052 f.mark_greedy("oom").unwrap();
17053 f.clear_greedy();
17054 assert!(!f.greedy_failed());
17055 f.mark_sampled("oom").unwrap();
17056 f.clear_sampled();
17057 assert!(!f.sampled_failed());
17058 // after a silent clear there is nothing left for resume to report.
17059 assert!(f.reset_on_resume().is_none());
17060 }
17061}
17062
17063/// SAMPLED DRAFT-GRAPH KEY (lane/graph-s-key-exactness-20260819).
17064///
17065/// These are the CPU teeth for an exactness bug whose live reproduction needs a GPU, a trunk, a
17066/// drafter and a two-turn session: the key itself. Every test below fails against the pre-fix key
17067/// `(seed, temp.to_bits(), k)` — `legacy_key` restates it so the collision is explicit rather
17068/// than remembered.
17069#[cfg(test)]
17070mod sampled_graph_key_tests {
17071 use super::{SampledGraphKey, debug_t_pred0};
17072
17073 /// The pre-fix key, verbatim: `let s_key = (sp_seed, sp_temp.to_bits(), k);`
17074 fn legacy_key(k: &SampledGraphKey) -> (u64, u32, usize) {
17075 (k.seed, k.temp_bits, k.k)
17076 }
17077
17078 fn pure_temp_key() -> SampledGraphKey {
17079 // temperature 1.0, filters off — today's serve default, the shape that parks a graph.
17080 SampledGraphKey::new(12345, 1.0, 3, 0, 1.0, 0.0, false)
17081 }
17082
17083 /// THE COLLISION. Two requests that differ ONLY in the truncation filters shared one key, so
17084 /// a parked pure-temp graph survived into a filtered request and the launch site launched it.
17085 #[test]
17086 fn vendor_filters_change_the_key() {
17087 let parked = pure_temp_key();
17088 // qwen3.8 generation_config.json — what the vendor-default flip makes the default shape.
17089 let vendor = SampledGraphKey::new(12345, 1.0, 3, 20, 0.95, 0.0, false);
17090 assert_eq!(
17091 legacy_key(&parked),
17092 legacy_key(&vendor),
17093 "pre-fix key collided: this is the bug, and the reason a test asserts on it",
17094 );
17095 assert_ne!(parked, vendor, "post-fix key must separate the two regimes");
17096 assert!(parked.pure_temp());
17097 assert!(!vendor.pure_temp());
17098 }
17099
17100 /// Each distribution-shaping field alone is enough to drop the parked graph.
17101 #[test]
17102 fn every_filter_field_is_keyed() {
17103 let base = pure_temp_key();
17104 for (what, other) in [
17105 (
17106 "top_k",
17107 SampledGraphKey::new(12345, 1.0, 3, 20, 1.0, 0.0, false),
17108 ),
17109 (
17110 "top_p",
17111 SampledGraphKey::new(12345, 1.0, 3, 0, 0.95, 0.0, false),
17112 ),
17113 (
17114 "min_p",
17115 SampledGraphKey::new(12345, 1.0, 3, 0, 1.0, 0.05, false),
17116 ),
17117 (
17118 "penalties",
17119 SampledGraphKey::new(12345, 1.0, 3, 0, 1.0, 0.0, true),
17120 ),
17121 ] {
17122 assert_ne!(base, other, "{what} must be part of the key");
17123 assert!(!other.pure_temp(), "{what} leaves the pure-temp regime");
17124 assert_eq!(
17125 legacy_key(&base),
17126 legacy_key(&other),
17127 "{what} was invisible to the pre-fix key",
17128 );
17129 }
17130 }
17131
17132 /// The baked constants stay keyed (this half was always right — regression cover for it).
17133 #[test]
17134 fn baked_constants_stay_keyed() {
17135 let base = pure_temp_key();
17136 assert_ne!(
17137 base,
17138 SampledGraphKey::new(999, 1.0, 3, 0, 1.0, 0.0, false),
17139 "seed"
17140 );
17141 assert_ne!(
17142 base,
17143 SampledGraphKey::new(12345, 0.7, 3, 0, 1.0, 0.0, false),
17144 "temp"
17145 );
17146 assert_ne!(
17147 base,
17148 SampledGraphKey::new(12345, 1.0, 4, 0, 1.0, 0.0, false),
17149 "k"
17150 );
17151 // bitwise on temperature: 0.7f32 vs the same value re-derived must NOT differ.
17152 assert_eq!(
17153 SampledGraphKey::new(1, 0.7, 3, 0, 1.0, 0.0, false),
17154 SampledGraphKey::new(1, 7.0 / 10.0, 3, 0, 1.0, 0.0, false),
17155 );
17156 }
17157
17158 /// THE LOAD-BEARING HALF OF THE SEED DECISION (lane/session-resume-sampler-predicate-
17159 /// 20260820). The whole-session resume predicate deliberately does NOT compare `seed`: an
17160 /// omitted serve `seed` draws fresh per-request entropy, so comparing it would refuse every
17161 /// seed-omitting sampled conversation. That is only sound because the one piece of parked state
17162 /// that BAKES the seed — this graph — is re-keyed on it, so a seed change drops and recaptures.
17163 ///
17164 /// This test is the other end of that argument, asserted here rather than remembered in a
17165 /// comment: if a future change dropped `seed` from the key, the resume predicate's exclusion
17166 /// would silently become the unsound thing it is documented not to be.
17167 /// (Paired with `seed_alone_does_not_refuse` in `memra-sampling`.)
17168 #[test]
17169 fn seed_alone_still_rekeys_the_draft_graph() {
17170 let parked = pure_temp_key();
17171 let reseeded = SampledGraphKey::new(999, 1.0, 3, 0, 1.0, 0.0, false);
17172 assert_ne!(
17173 parked, reseeded,
17174 "a seed-only change MUST drop the parked sampled graph — the resume predicate's \
17175 decision not to compare seed rests on exactly this",
17176 );
17177 // Same regime on both sides: the drop is a recapture, not a fall to the eager chain
17178 // because of a filter difference.
17179 assert!(parked.pure_temp() && reseeded.pure_temp());
17180 }
17181
17182 /// `pure_temp()` is the capture guard's predicate, computed from the key so the two cannot
17183 /// drift. The equality below is the invariant the launch-site guard asserts: identical keys
17184 /// agree on the regime, so a graph that survives the drop is legal to launch.
17185 #[test]
17186 fn equal_keys_agree_on_the_regime() {
17187 let a = SampledGraphKey::new(7, 0.8, 3, 20, 0.95, 0.0, false);
17188 let b = SampledGraphKey::new(7, 0.8, 3, 20, 0.95, 0.0, false);
17189 assert_eq!(a, b);
17190 assert_eq!(a.pure_temp(), b.pure_temp());
17191 // top_p slightly above 1.0 (a client sending 1.0 exactly, or an operator default) is
17192 // still the unfiltered regime, matching the original `sp.top_p >= 1.0` test.
17193 assert!(SampledGraphKey::new(7, 0.8, 3, 0, 1.0, 0.0, false).pure_temp());
17194 assert!(SampledGraphKey::new(7, 0.8, 3, 0, 1.5, -1.0, false).pure_temp());
17195 }
17196
17197 /// The WIDENED capture regime (lane/step37-draft-graph-serving-20260830): truncation-
17198 /// filtered shapes are capturable — the filter runs IN-GRAPH (`filter_stats` +
17199 /// `gumbel_perturb_filtered_ctr`), so the draft draws from the same filtered
17200 /// distribution the accept test reconstructs. Penalties never are: the per-round
17201 /// history cannot be baked. The step37 vendor-default shape (temp 0.5 / top_p 0.9) is
17202 /// exactly the previously-excluded regime this lane exists to capture.
17203 #[test]
17204 fn filtered_regimes_are_capturable_penalties_never() {
17205 let vendor = SampledGraphKey::new(12345, 0.5, 3, 0, 0.9, 0.0, false);
17206 assert!(!vendor.pure_temp());
17207 assert!(vendor.filtered());
17208 assert!(
17209 vendor.graph_capturable(),
17210 "the vendor-default filtered shape must be capturable (default door state)",
17211 );
17212 assert!(pure_temp_key().graph_capturable());
17213 assert!(
17214 !pure_temp_key().filtered(),
17215 "pure-temp takes the legacy (filterless) capture body",
17216 );
17217 let pen = SampledGraphKey::new(12345, 0.5, 3, 0, 0.9, 0.0, true);
17218 assert!(
17219 !pen.graph_capturable(),
17220 "penalty history varies per round and can never be baked into a graph",
17221 );
17222 }
17223
17224 /// MEMRA_DEBUG_SPEC on a SAMPLED spec request past round 0: the print must render without
17225 /// indexing the empty greedy `preds` vector (it panicked the GPU worker before this lane).
17226 #[test]
17227 fn debug_print_survives_the_sampled_arm() {
17228 // round >= 1 with a pending bonus == base 1, sampled == `preds` empty.
17229 assert_eq!(debug_t_pred0(true, 1, 4242, &[]), "n/a");
17230 assert_eq!(debug_t_pred0(true, 2, 4242, &[]), "n/a");
17231 // round 0 without a pending bonus still reports last_pred, in both arms.
17232 assert_eq!(debug_t_pred0(true, 0, 4242, &[]), "4242");
17233 assert_eq!(debug_t_pred0(false, 0, 4242, &[7, 8]), "4242");
17234 // greedy keeps the real prediction it always printed.
17235 assert_eq!(debug_t_pred0(false, 1, 4242, &[7, 8]), "7");
17236 assert_eq!(debug_t_pred0(false, 2, 4242, &[7, 8]), "8");
17237 }
17238}