memra_engine/spec.rs
1//! Qwen3.5 MTP (NextN) greedy speculative decode (research/mtp/MTP-PLAN.md §A/§B/§C/§D).
2//!
3//! Greedy spec decode is MATHEMATICALLY EXACT: the accepted+bonus token stream is token-for-token
4//! identical to plain greedy `generate`. This module provides:
5//! - `mtp_head_forward` (§A, T=1): one NextN draft-token forward.
6//! - `decode_step_t` (§D.3, T=K+1): batched target verify forward, all-column logits.
7//! - `generate_spec` (§B): the draft/verify/accept/rollback orchestrator.
8//! Cache snapshot/rollback lives in cache.rs (§D.4). The MTP head uses its OWN scratch KV (§D.6),
9//! PERSISTENT over the committed sequence (see `MtpScratch`).
10
11use crate::Engine;
12use crate::cache::{Cache, KvLayer};
13use crate::forward::argmax;
14use crate::hybrid::{FullAttnLayer, HybridModel, LinearAttnLayer, Mixer, MtpHead};
15use cudarc::driver::CudaSlice;
16use std::sync::atomic::{AtomicU64, Ordering};
17
18/// Parse the documented `MEMRA_SPEC_REPLAY=1` rollback seam.
19///
20/// Keep this shared with serving admission so `=0` cannot select replay in one
21/// layer while another layer treats it as disabled.
22pub fn spec_replay_env_on(value: Option<&str>) -> bool {
23 value == Some("1")
24}
25
26pub fn spec_replay_env_enabled() -> bool {
27 let value = std::env::var("MEMRA_SPEC_REPLAY").ok();
28 spec_replay_env_on(value.as_deref())
29}
30
31/// step35 dcw draft-chain door (lane/step37-draft-graph-20260829). ON routes the step35 MTP
32/// block's draft attention through the WINDOWED device-counter family
33/// (`append_kv_quantized_dcw` + `fa_decode_dcw`, the step TP graph arc's kernels), which
34/// derives the SWA view entirely from device state (len_d, base_d, window): exactly the view
35/// offset the old capture refusal said `fa_decode_dc` could not express. BOTH draft modes
36/// switch together: eager and captured run the ONE launcher at the ONE bucket
37/// (min(cap, window)), so graph-vs-eager draft parity holds by construction (the
38/// `mtp_full_attn_dc` precedent).
39///
40/// DEFAULT ON since lane/step37-draft-graph-serving-20260830: the 20260829 lane shipped it
41/// OFF because it enabled nothing at the shipping head count (capture was structurally
42/// unreachable at heads=3); with the multi-head chain capture and the in-graph filtered
43/// sampler landed, this door is the kernel prerequisite for the captured chain on the
44/// QUALIFIED serving shape, and the exactness battery (greedy K=1..8 identity, per-K
45/// acceptance identity, seeded sampled twins) banks on the ON arm. Rollback seam:
46/// MEMRA_STEP35_DRAFT_DCW=0 restores the host-len eager arm (`mtp_step35_attn`) plus the
47/// named capture refusal, byte-for-byte the pre-lane serving; no state survives restart.
48fn step35_draft_dcw_on() -> bool {
49 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
50 *ON.get_or_init(|| std::env::var("MEMRA_STEP35_DRAFT_DCW").as_deref() != Ok("0"))
51}
52
53/// Multi-head MTP draft-chain capture door (lane/step37-draft-graph-serving-20260830,
54/// default ON — receipts in the lane RESULTS). ON lets the step-modulo prefix-replay chain
55/// (`mtp_extra` non-empty, the step37 3-head shipping shape) capture per-head single-row
56/// CUDA graphs and replay them in the exact eager launch order; the chain POLICY (head
57/// selection, prefix length, seed history) stays host-side, so graph-vs-eager drafts are
58/// bit-identical by construction. A failed capture degrades LOUDLY to the eager chain (the
59/// draft-graph WARN contract). OFF (=0) keeps the eager chain as the only multi-head path —
60/// the pre-lane serving byte-for-byte. Single-head capture is untouched by this door.
61fn mtp_chain_graph_on() -> bool {
62 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
63 *ON.get_or_init(|| std::env::var("MEMRA_MTP_CHAIN_GRAPH").as_deref() != Ok("0"))
64}
65
66/// In-graph FILTERED sampled draft door (lane/step37-draft-graph-serving-20260830, default
67/// ON — receipts in the lane RESULTS). ON widens the sampled draft-graph capture from the
68/// pure-temp regime to every truncation-filtered regime (top_k / top_p / min_p): the capture
69/// body runs `filter_stats` + `gumbel_perturb_filtered_ctr` IN-GRAPH, so the draft draws
70/// from the SAME filtered distribution the verify's accept test reconstructs (the
71/// graph-s-key exactness law, now satisfied inside the graph instead of by refusing it).
72/// Penalties stay eager either way (the history varies per round and cannot be baked).
73/// The pure-temp capture body is UNTOUCHED by this door (byte-identical to the pre-lane
74/// graph). OFF (=0) restores the pure-temp-only capture guard: filtered requests draft
75/// eager, byte-for-byte the pre-lane behavior.
76fn spec_graph_filtered_on() -> bool {
77 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
78 *ON.get_or_init(|| std::env::var("MEMRA_SPEC_GRAPH_FILTERED").as_deref() != Ok("0"))
79}
80
81fn parse_prime_trows_width(value: Option<&str>) -> Result<usize, String> {
82 let Some(raw) = value else {
83 return Ok(8);
84 };
85 let width = raw
86 .parse::<usize>()
87 .map_err(|_| format!("MEMRA_PRIME_TROWS_T must be an integer in 2..=8, got {raw:?}"))?;
88 if !(2..=8).contains(&width) {
89 return Err(format!("MEMRA_PRIME_TROWS_T must be in 2..=8, got {width}"));
90 }
91 Ok(width)
92}
93
94#[cfg(test)]
95mod prime_trows_width_tests {
96 #[test]
97 fn width_defaults_to_eight_and_refuses_invalid_operator_values() {
98 assert_eq!(super::parse_prime_trows_width(None), Ok(8));
99 assert_eq!(super::parse_prime_trows_width(Some("2")), Ok(2));
100 assert_eq!(super::parse_prime_trows_width(Some("8")), Ok(8));
101 for invalid in ["", "1", "9", "32", "wide"] {
102 let err = super::parse_prime_trows_width(Some(invalid)).unwrap_err();
103 assert!(err.contains("MEMRA_PRIME_TROWS_T"), "{err}");
104 assert!(err.contains("2..=8"), "{err}");
105 }
106 }
107}
108
109/// One compact, anchor-bounded DSpark supervision record. `tokens[0]` is the anchor at p and
110/// `hidden` is its predecessor carrier h[p-1], matching the live NextN/DSpark pairing. Target
111/// rows p..p+gamma-1 score tokens p+1..p+gamma. They are the full-target softmax's top-k
112/// entries; `target_tail_probs[j]` is the probability mass outside those rows. All flattened
113/// target arrays are `[gamma, top_k]` in row-major order.
114pub struct DsparkAnchorRecord {
115 pub position: usize,
116 pub hidden: Vec<f32>,
117 pub tokens: Vec<u32>,
118 pub target_top_ids: Vec<u32>,
119 pub target_top_logits: Vec<f32>,
120 pub target_top_probs: Vec<f32>,
121 pub target_tail_probs: Vec<f32>,
122}
123
124fn dspark_sparse_softmax_topk(
125 logits: &[f32],
126 top_k: usize,
127 temperature: f32,
128) -> Result<(Vec<u32>, Vec<f32>, Vec<f32>, f32), Box<dyn std::error::Error>> {
129 if logits.is_empty() || top_k == 0 || top_k > logits.len() || temperature <= 0.0 {
130 return Err("invalid DSpark sparse-softmax shape or temperature".into());
131 }
132 if logits.iter().any(|value| !value.is_finite()) {
133 return Err("DSpark target logits contain a non-finite value".into());
134 }
135 let mut ranked: Vec<(u32, f32)> = logits
136 .iter()
137 .copied()
138 .enumerate()
139 .map(|(index, value)| (index as u32, value))
140 .collect();
141 let compare = |left: &(u32, f32), right: &(u32, f32)| {
142 right.1.total_cmp(&left.1).then(left.0.cmp(&right.0))
143 };
144 ranked.select_nth_unstable_by(top_k - 1, compare);
145 ranked[..top_k].sort_unstable_by(compare);
146
147 let max_logit = logits.iter().copied().fold(f32::NEG_INFINITY, f32::max);
148 let inv_temperature = 1.0f64 / temperature as f64;
149 let denominator: f64 = logits
150 .iter()
151 .map(|value| (((*value - max_logit) as f64) * inv_temperature).exp())
152 .sum();
153 let ids: Vec<u32> = ranked[..top_k].iter().map(|(index, _)| *index).collect();
154 let top_logits: Vec<f32> = ranked[..top_k].iter().map(|(_, value)| *value).collect();
155 let top_probs: Vec<f32> = top_logits
156 .iter()
157 .map(|value| ((((value - max_logit) as f64) * inv_temperature).exp() / denominator) as f32)
158 .collect();
159 let top_mass: f64 = top_probs.iter().map(|value| *value as f64).sum();
160 let tail = (1.0f64 - top_mass).clamp(0.0, 1.0) as f32;
161 Ok((ids, top_logits, top_probs, tail))
162}
163
164fn flatten_dspark_rows<T>(
165 rows: Vec<Option<Vec<T>>>,
166 position: usize,
167 label: &str,
168) -> Result<Vec<T>, Box<dyn std::error::Error>> {
169 let mut flattened = Vec::new();
170 for (slot, row) in rows.into_iter().enumerate() {
171 flattened.extend(
172 row.ok_or_else(|| format!("missing DSpark {label} at {position} slot {slot}"))?,
173 );
174 }
175 Ok(flattened)
176}
177
178/// H-SEED CONVENTION (MEMRA_SPEC_HPOST=1): feed the MTP head the POST-norm hidden — trunk rows
179/// hand over `output_norm(x)` and the draft chain recurrence hands over `shared_head_norm(h_nextn)`
180/// (= final_h) — matching the reference engines: llama.cpp #24025 ("qwen35: use post-norm hidden
181/// state for MTP", t_h_nextn is taken AFTER the final norm in both trunk and MTP graphs) and
182/// SGLang's qwen3_5_mtp (spec_info.hidden_states = the target model's post-norm output). memra's
183/// historical convention (default, MTP-PLAN §A) is PRE-norm x. Draft-quality-only: exactness is
184/// the verify's job either way; acceptance arbitrates. OnceLock: read once, hot-loop safe.
185/// `MEMRA_SPEC_HEAD_ROWS=1` — batch the verify tail's LM head over its t columns instead of running
186/// it at m=1 once per column. See the call site in `decode_step_t_core_stream` for why the batched
187/// form is the same per-row arithmetic (the bf16/q8 rows twins, not cuBLASLt) and what it costs
188/// today: the head is re-streamed t times per verify pass. Default off until the byte tape says so.
189pub(crate) fn head_rows_on() -> bool {
190 static ENV: std::sync::OnceLock<Option<bool>> = std::sync::OnceLock::new();
191 crate::step37_door(&ENV, "MEMRA_SPEC_HEAD_ROWS")
192}
193
194/// The serving walk's own doors, tri-stated the same way (owner flip 2026-08-27): env forces,
195/// unset takes the step37 family default. Call sites are the t-row verify walk itself.
196pub(crate) fn spec_verify_eager_on() -> bool {
197 static ENV: std::sync::OnceLock<Option<bool>> = std::sync::OnceLock::new();
198 crate::step37_door(&ENV, "MEMRA_SPEC_VERIFY_EAGER")
199}
200
201pub(crate) fn spec_verify_tcol_on() -> bool {
202 static ENV: std::sync::OnceLock<Option<bool>> = std::sync::OnceLock::new();
203 crate::step37_door(&ENV, "MEMRA_SPEC_VERIFY_TCOL")
204}
205
206/// NOT family-armed (2026-08-27): the walk's prime leaves its sub-32 TAIL chunk out of the
207/// DISTRIBUTED kv, so the server refuses before decode with "cache lengths diverged
208/// local=N distributed=floor(N/32)*32" for every prompt whose token count is not a multiple of
209/// 32 — i.e. nearly all real traffic. Isolated on the server route: defaults ERR (local=445
210/// distributed=416), MEMRA_PRIME_TROWS=0 OK. It was default-OFF before the 2026-08-27 flip and
211/// goes back to opt-in until the tail append is fixed and gated ON THE SERVER ROUTE, not just
212/// run-gen (run-gen calls decode_step_t on the whole prompt and never exercises this path — the
213/// reason a run-gen-only receipt could not see it). The GEMM prime supersedes it on this route.
214pub(crate) fn prime_trows_on() -> bool {
215 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
216 *ON.get_or_init(|| std::env::var("MEMRA_PRIME_TROWS").as_deref() == Ok("1"))
217}
218
219pub(crate) fn tcol_ffn_on() -> bool {
220 static ENV: std::sync::OnceLock<Option<bool>> = std::sync::OnceLock::new();
221 crate::step37_door(&ENV, "MEMRA_TCOL_FFN")
222}
223
224pub(crate) fn spec_hpost() -> bool {
225 static H: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
226 *H.get_or_init(|| {
227 std::env::var("MEMRA_SPEC_HPOST")
228 .map(|v| v != "0")
229 .unwrap_or(false)
230 })
231}
232
233/// LEAN VERIFY (default ON since 2026-07-08; MEMRA_SPEC_LEAN=0 reverts — close35 lane): the verify m-scaling
234/// probe + nsys diff showed the verify t-path pays ~1.0ms/call at m=1 over eager decode on the
235/// 35B, and the kernels are NOT the cause (dev-MoE identical, kernel-time delta only +179us).
236/// The overhead is (a) ~250 extra cuMemsetD8Async/call from `e.zeros()` on buffers every kernel
237/// fully overwrites (~0.9ms host issue + ~0.35ms GPU) and (b) the t=1 FA rows dispatch (rows_v2 +
238/// combine_rows, +50us vs the eager fa_decode pair). This flag switches (a) fully-overwritten
239/// verify buffers to `e.uninit` (identical bytes: every element is written before read) and
240/// (b) t==1 verify FA to the eager `fa_decode` entry (byte-identical: kernel-check pins the
241/// rows-vs-loop identity and the per-row loop at t=1 IS fa_decode on the same q). Gates arbitrate.
242pub(crate) fn spec_lean() -> bool {
243 static L: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
244 // DEFAULT ON since 2026-07-08 (MEMRA_SPEC_LEAN=0 reverts): bit-identical (buffers fully
245 // overwritten; gates green incl maxdiff-identical run-gen) and measured +2.4% e2e p3 /
246 // +1.5% p2 at the daily 35B config. m=1 verify now costs eager-decode parity.
247 *L.get_or_init(|| {
248 std::env::var("MEMRA_SPEC_LEAN")
249 .map(|v| v != "0")
250 .unwrap_or(true)
251 })
252}
253
254/// SMALL-M BATCHED VERIFY (default ON since 2026-07-09; MEMRA_SPEC_M2=0 reverts — lane/spec-m2): extend the
255/// batched linear-attn verify arm down to t=2 and batch the MoE dev token loop over a
256/// grid.z=token axis at every verify t. The close35 m-scaling probe put the m=2 verify tier at
257/// x1.54 of m=1 (llama x1.14); the per-column linear chain (t<3) and the serial MoE dev token
258/// loop are the two launch-structure causes. Both changes are LAUNCH-STRUCTURE ONLY:
259/// (a) the batched conv's t<pad ring update is pure copies (ssm_conv_ring_rebuild from a cloned
260/// ring — the ring stores raw input columns); every arithmetic kernel is the same one the
261/// t>=3 arm already runs (matmul_decode_exact bit-identical at m=2-4, gdn_scan's internal
262/// t-loop == chained T=1 steps);
263/// (b) the MoE dev-rows twins run the serial loop's per-token warp program with tok-offset
264/// pointers (same sel/w/aq/ad bytes, same dot order, same slot-ordered FMA chain).
265/// Gates arbitrate: run-spec K=1..8 self-consistency (35B+9B), kernel-check, run-gen argmax.
266pub(crate) fn spec_m2() -> bool {
267 static M: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
268 // DEFAULT ON since 2026-07-09 (MEMRA_SPEC_M2=0 reverts): launch-structure only — t=2
269 // batched linear arm (ring-roll copies, zero new FP order) + MoE dev-rows kernels
270 // (grid.z=token, 4 launches/layer at any verify t). Acceptance bit-identical at every K;
271 // 35B p2 +3.4% / p3 +3.6%; the profitable-K plateau widens (new optimum K=3 at 223).
272 *M.get_or_init(|| {
273 std::env::var("MEMRA_SPEC_M2")
274 .map(|v| v != "0")
275 .unwrap_or(true)
276 })
277}
278pub(crate) fn spec_stream() -> bool {
279 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
280 *ON.get_or_init(|| std::env::var("MEMRA_SPEC_STREAM").as_deref() == Ok("1"))
281}
282pub(crate) fn spec_stream_m() -> usize {
283 static M: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
284 *M.get_or_init(|| {
285 std::env::var("MEMRA_SPEC_STREAM_M")
286 .ok()
287 .and_then(|v| v.parse().ok())
288 .unwrap_or(4)
289 })
290}
291pub(crate) fn spec_devacc() -> bool {
292 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
293 *ON.get_or_init(|| std::env::var("MEMRA_SPEC_DEVACC").as_deref() == Ok("1"))
294}
295/// Engine-bundle slice 2 (DSF-ROUNDCOST-20260820 §1.1 host/device round trips + §2 rows 2-3),
296/// DEFAULT ON (`MEMRA_DSPARK_DEFER_READBACK=0` reverts): the dspark round's draft-chain DtoH
297/// is DEFERRED past verify dispatch and merged with the verify-argmax readback into ONE host
298/// sync (2 blocking DtoH/round -> 1). Verify embeds DEVICE tokens (`chain_d`) through the
299/// resident embed table — `embed_gather_u32_t`, bit-identical rows to the host gather by its
300/// own pinned contract. The host therefore dispatches snap + the whole verify while the DRAFT
301/// is still executing, instead of blocking ~1.7 ms on the chain and letting the device drain.
302/// Ladder arm only: the confidence policies size vt from a pre-verify head readback (their
303/// chain readback merges into that same sync instead). Exactness unchanged BY CONSTRUCTION —
304/// same tokens, same kernels, same order; E2E + accept-bank gates arbitrate.
305pub(crate) fn dspark_defer_readback_on() -> bool {
306 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
307 *ON.get_or_init(|| {
308 std::env::var("MEMRA_DSPARK_DEFER_READBACK")
309 .map(|v| v != "0")
310 .unwrap_or(true)
311 })
312}
313/// Engine-bundle slice 1 (DSF-ROUNDCOST-20260820 §1.1, lane/dspark-engine-bundle-20260820),
314/// DEFAULT ON (`MEMRA_STATE_COPY_BATCH=0` reverts): batch the dspark round's GDN state
315/// snapshot and partial-accept restore into single `copy_batch_uniform_f32` launches
316/// instead of ~2 memcpy dispatches (+2 alloc_zeros on the snap side) per linear layer per
317/// round — measured 0.67 ms/round snap + 0.25 ms/round commit of pure dispatch on the q38
318/// route. Launch-structure only: bytes, buffers and stream order are unchanged, so
319/// acceptance and streams stay bit-identical (E2E-gated on the B1 packs).
320pub(crate) fn state_copy_batch_on() -> bool {
321 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
322 *ON.get_or_init(|| {
323 std::env::var("MEMRA_STATE_COPY_BATCH")
324 .map(|v| v != "0")
325 .unwrap_or(true)
326 })
327}
328/// Engine-bundle slice 3 + fa-execupdate slice 4c (DSF-ROUNDCOST-20260820 §5 rank 1),
329/// DEFAULT OFF — `MEMRA_DSPARK_VERIFY_GRAPH=1` opts in: per-(segment, vt) CUDA graphs
330/// for the LINEAR-layer runs, plus the full-verify single graph per (vt, rung) when a
331/// round's rows all ride one seqs rung — see [`DsparkVerifyGraphs`]. Requires the
332/// slice-2 deferred path (device tokens); the eager walk is the byte-identical fallback.
333///
334/// MEASURED disposition (box6 card0, agentic pack, 2026-08-20, both slices): exactness
335/// holds everywhere (ALL EXACT, accept lines byte-match the banks, ckpt-gate oracle
336/// green over the graph + slab-commit paths). Slice-3's AUTO_FREE launch-scan limiter
337/// (25.6 us x 16 launches ≈ 0.41 ms/round) is FIXED — the captured bodies' alloc nodes
338/// are balanced by in-graph frees (census 84/84 per segment, 1776/1776 full) so graphs
339/// instantiate USE_NODE_PRIORITY and the scan is gone. What remains at gate scale:
340/// segment graphs +0.1 tok/s over the batched-rows default (114.4 vs 114.3 x5
341/// interleaved — the linear launch overhead was only ~0.1 ms); the FULL-verify graph is
342/// NET NEGATIVE at gate scale (110.6 vs 114.2: ~14-21 (vt, rung) captures/process at
343/// 2 full-walk executions + ~2.9k-node instantiate each eat far more than the ~0.2-0.3
344/// ms/round of remaining launch overhead). The orchestration ceiling of §1.3 is spent —
345/// the fa/append recovery landed DEFAULT-ON as the batched rows arm
346/// (`dspark_fa_rows_on`), not as a graph. The serve-lifetime cell (DSF-ROUNDCOST §9,
347/// nj-ws-solo) measured the amortization: crossover K≈33 requests, steady −0.246
348/// ms/round, −1.25% session wall over 240 requests — and the graphs-serve lane wired
349/// the door into the session arm (`dspark_spec_session_burst`) as a model-owned
350/// capture pool shared across sessions. Stays opt-in pending the owner's default-ON
351/// ratification on the serve-surface battery.
352pub(crate) fn dspark_verify_graph_on() -> bool {
353 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
354 *ON.get_or_init(|| std::env::var("MEMRA_DSPARK_VERIFY_GRAPH").as_deref() == Ok("1"))
355}
356/// MTP-ROUTE verify graphs, DEFAULT ON for the GDN+MoE family since 2026-08-23
357/// (`MEMRA_SPEC_VERIFY_GRAPH=0` is the kill switch, `=1` opts other families in).
358///
359/// The slice-4c capture already lived inside `qwen35_verify_tparallel` and said so in its own
360/// comment — "stream rides the qwen35moe burst, graphs ride the dspark route" — with no caller
361/// on this route. The MTP spec round is that caller.
362///
363/// WHY it is worth a default (receipts: `research/orndecode-20260822/VGRAPH.md`). With
364/// `MEMRA_SPEC_PHASE=1` this route's round reads verify-ISSUE 44-58% and verify-WAIT **0.0%**:
365/// the host is never waiting for the device, it is spending its own time launching the trunk.
366/// Replay collapses that into one graph launch and the phase all but disappears (55-62 ms ->
367/// 8-10 ms per burst).
368///
369/// MEASURED, two host generations, forced ON/OFF, balanced 4+4 boots in both orders:
370/// * current-generation host (9950X, the serving class): OFF 266.0-266.5, ON 318.8-319.5
371/// tok/s — **+19.7%**, no overlap, sub-1% spread per arm; per-round 6.9 -> 5.7 ms.
372/// * Zen 3 host: +3-9% (that rig's own clock drift is wider than the effect, so the ratio
373/// comes from per-round phase totals, which are internal to each boot).
374/// The ON arm lands at ~320 tok/s on BOTH hosts while OFF tracks host speed — the arm moves
375/// the round off the host and onto the device, which is the whole point.
376///
377/// EXACTNESS is structural (same kernels, same order) and gated anyway: a fixed-seed SAMPLED
378/// completion hashes identically ON vs OFF **and across both hosts** (`08941d5bb9762b21`),
379/// greedy seed-pinned likewise, `run-spec` K=1..8 PASS on both arms with identical acceptance
380/// at every K, kernel-check ALL GREEN.
381///
382/// SCOPE, deliberately narrow: default ON only where it was measured — the GatedDeltaNet +
383/// MoE family (`vgraph_family_default`). Qwen3.8-27B is GDN + DENSE mlp and would otherwise
384/// inherit this default unmeasured, which is the family-by-family law this repo keeps; it can
385/// opt in with `=1` once it has its own interleave. Also never armed together with
386/// ROUND-STREAM, and a round wider than the pool declines it for the eager walk.
387pub(crate) fn spec_verify_graph_env() -> Option<bool> {
388 static ON: std::sync::OnceLock<Option<bool>> = std::sync::OnceLock::new();
389 *ON.get_or_init(
390 || match std::env::var("MEMRA_SPEC_VERIFY_GRAPH").as_deref() {
391 Ok("1") => Some(true),
392 Ok("0") => Some(false),
393 _ => None,
394 },
395 )
396}
397/// SERVE-ROUTE twin of [`dspark_verify_graph_on`], DEFAULT ON — owner-ratified
398/// 2026-08-22 on the §10 serve-lifetime battery (DSF-ROUNDCOST-20260820 §10.3:
399/// crossover K=36–43, steady −0.357 ms/round, session wall −1.55..−1.65%, byte-exact
400/// 240/240 ×3 pairs, pool bounded at 8,852 MiB under `MEMRA_DSPARK_VG_MAX`). The env
401/// stays as the kill-switch: `MEMRA_DSPARK_VERIFY_GRAPH=0` restores the eager walk
402/// (byte-identical body); `MEMRA_DSPARK_VG_MAX=0` is the finer freeze valve. The BIN
403/// arm keeps its own opt-in default (`dspark_verify_graph_on`): at gate scale the
404/// capture toll is never repaid (§8 measured disposition — 14–21 captures over a
405/// 256-token run vs the serve session's thousands of rounds), and the two
406/// instruments must keep their own measured dispositions rather than share one flag.
407pub(crate) fn dspark_verify_graph_serve_on() -> bool {
408 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
409 *ON.get_or_init(|| std::env::var("MEMRA_DSPARK_VERIFY_GRAPH").as_deref() != Ok("0"))
410}
411/// Capture-count ceiling for the dspark verify-graph pool (graphs-serve lane) — the
412/// pool's memory policy STATED instead of silently unbounded. The keyspace is
413/// intrinsically finite — segment keys (run_start, vt) ≤ 16 runs x 7 windows, full
414/// keys (vt, rung, hi) ≤ 7 windows x the split-rung ladder (8 rungs at 32k ctx), ~168
415/// on the q38 export — so the default (256) never engages there; the knob is the
416/// safety valve for a future export with a wider ladder. At the ceiling the pool
417/// FREEZES: existing keys keep replaying, rounds needing a new capture run the eager
418/// walk byte-identically (round-atomic — a partial refusal would mix slab- and
419/// cols-stashed layers inside one commit). No eviction by design: destroying a live
420/// exec graph re-opens the stale-address class the indirect tables exist to close,
421/// and the bounded keyspace makes reclaim worthless.
422pub(crate) fn dspark_vg_cap() -> usize {
423 static CAP: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
424 *CAP.get_or_init(|| {
425 std::env::var("MEMRA_DSPARK_VG_MAX")
426 .ok()
427 .and_then(|v| v.parse().ok())
428 .unwrap_or(256)
429 })
430}
431
432/// PROJECTED REMAINING GROWTH of the verify-graph pool, in bytes (lane/hermes-perf-fixes,
433/// 2026-08-23 — the admission accounting the "pool dwarfs spec admission reserve" finding
434/// asks for). The pool was measured at 8,852 MiB at storm-complete on the q38 export while
435/// admission's transient floor (`SPEC_SHRINK_RESERVE`) is 1.5 GiB and never charged for it:
436/// sessions admitted while the pool is cold overcommit VRAM the pool WILL hold, because the
437/// pool grows monotonically (no eviction by design) and is model-owned across sessions.
438///
439/// SELF-MEASURING, no per-model constant (generic-model law — the 8,852 MiB is a q38 number
440/// and proves nothing about another export): the debt is remaining capture slots x the
441/// MARGINAL bytes a capture adds to this device's graph mem pool.
442///
443/// MARGINAL, NOT MEAN — measured correction (box9 on-box receipt, 2026-08-23). The first
444/// version of this used the mean (`reserved / captures`) and the live serve log showed why
445/// that is wrong: with the pool's reservation flat at ~33.6 MiB across captures 1..3, the
446/// mean-based debt printed **8,556 MB, then 4,261, then 2,830** — it extrapolated capture
447/// #1's ONE-TIME shared allocation (staging buffers, stash slabs, pointer tables: sized
448/// once per pool, shared by every key) across all 256 slots. An 8.5 GB phantom reserve at
449/// boot can refuse admissions that would have fit, which is a worse defect than the
450/// under-charge this accounting exists to remove. The marginal reading prices what an
451/// ADDITIONAL key actually costs: two observations `(captures, reserved)` give
452/// `(r1 - r0) / (c1 - c0)`, which is ~0 on an export whose pool does not grow per key and
453/// tracks real growth on one that does.
454///
455/// BOOTSTRAP (only one observation so far, so growth is unmeasurable): reserve one more
456/// pool's worth — `min(remaining x mean, reserved)`. "We have measured `reserved` bytes for
457/// `captures` keys; until growth is measurable, assume at most a doubling" is fail-safe in
458/// the same direction as the old rule without the 255x extrapolation.
459///
460/// Before the FIRST capture the debt is 0 (a single capture lands well inside the existing
461/// 1.5 GiB floor). `cap` is the intrinsic freeze ceiling (`MEMRA_DSPARK_VG_MAX`; =0 freeze
462/// valve => the pool cannot grow => debt 0); at or past the cap the pool FREEZES, so the
463/// debt is 0 there too.
464pub fn dspark_vg_debt_projection(
465 captures: usize,
466 cap: usize,
467 reserved_bytes: usize,
468 prev: Option<(usize, usize)>,
469) -> usize {
470 if captures == 0 || cap == 0 {
471 return 0;
472 }
473 let remaining = cap.saturating_sub(captures);
474 if remaining == 0 {
475 return 0;
476 }
477 match prev {
478 // marginal growth between two observations of the same pool
479 Some((c0, r0)) if captures > c0 => {
480 let marginal = reserved_bytes.saturating_sub(r0) / (captures - c0);
481 remaining.saturating_mul(marginal)
482 }
483 // bootstrap: at most one more pool's worth
484 _ => remaining
485 .saturating_mul(reserved_bytes / captures)
486 .min(reserved_bytes),
487 }
488}
489/// Engine-bundle slice 4 (fa-execupdate lane, DSF-ROUNDCOST-20260820 §6 close: "the
490/// residual gap lives in the FULL-ATTENTION per-row section"), DEFAULT ON —
491/// `MEMRA_DSPARK_FA_ROWS=0` reverts to the per-row loop: when every row of a verify
492/// round takes the v4-seqs arm on ONE `fa_split_keys` rung (the straddle law, evaluated
493/// at the round's first and last t_kv — both eligibility gates are intervals in t_kv),
494/// the qwen35 t-parallel verify's per-row KV-append + fa-decode loop collapses into the
495/// z-batched serving twins: ONE `append_quantize_kv_q8_0_q5_1_seqs` + ONE
496/// `fa_decode_vec_q_seqs_v4` + ONE combine per full-attention layer, replacing
497/// T x (4 dtod row copies + append + 3 memsets + main + combine) launches. Bytes are
498/// pinned by the batched-tick increment-2 kernel-check (seqs-vs-per-seq-loop bit
499/// identity: per-row T_kv derives in-kernel from pos_seq[z]; splits >= ns_eff write the
500/// empty partial the combine never reads, so the shared n_splits_max stride changes no
501/// bytes) and re-gated e2e by this lane's battery.
502pub(crate) fn dspark_fa_rows_on() -> bool {
503 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
504 *ON.get_or_init(|| {
505 std::env::var("MEMRA_DSPARK_FA_ROWS")
506 .map(|v| v != "0")
507 .unwrap_or(true)
508 })
509}
510
511/// `t_pred0` for the `MEMRA_DEBUG_SPEC` per-round print, sampled-safe.
512///
513/// `generate_spec_inner2` fills its `preds` vector ONLY on the greedy path (`if !sampled`), and
514/// the per-round debug print was the sole consumer in the sampled arm: `t_pred(0)` survives round
515/// 0 (`base == 0` returns `last_pred`) and from round 1 (`base == 1`, a pending bonus) indexes an
516/// EMPTY vector — `index out of bounds: the len is 0 but the index is 0`, in the GPU worker
517/// thread, which then respawns and reloads weights while the request dies. So any sampled spec
518/// request longer than one round used to kill the worker whenever `MEMRA_DEBUG_SPEC` was set:
519/// the flag crashed precisely the regime it exists to investigate.
520///
521/// Fixed at the print site, not inside the closure, so the greedy accept walk keeps its strict
522/// indexing (an out-of-range pred there is a real bug and must still be loud).
523fn debug_t_pred0(sampled: bool, base: usize, last_pred: u32, preds: &[u32]) -> String {
524 if base == 0 {
525 return last_pred.to_string();
526 }
527 match preds.get(base - 1) {
528 Some(p) => p.to_string(),
529 // sampled: the greedy per-column argmax was never run for this round.
530 None => {
531 debug_assert!(
532 sampled,
533 "greedy spec: preds[{}] missing at base {base}",
534 base - 1
535 );
536 "n/a".to_string()
537 }
538 }
539}
540
541/// `MEMRA_SKEY_PROBE=1` — sampled-draft-graph key probe (lane/graph-s-key-exactness-20260819).
542///
543/// Reports, per burst and per round, which draft chain the sampled arm chose and under which
544/// filter regime, plus the ONE observable that separates a legal filtered draft from a stale
545/// pure-temp graph replayed under filters: an accept test whose gathered `q` is exactly 0.
546/// A draft token sampled from the FILTERED softmax can never gather q=0 (it was drawn from the
547/// kept set), so `q=0` in the verify means the draft came from a distribution the verify does
548/// not believe in — and `u * 0 < p` then accepts it unconditionally.
549///
550/// Its own env var, deliberately NOT `MEMRA_DEBUG_SPEC`: that flag panicked the GPU worker on
551/// any sampled spec request past round 0 until this lane fixed it (§2 of the bank note).
552pub(crate) fn skey_probe() -> bool {
553 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
554 *ON.get_or_init(|| std::env::var("MEMRA_SKEY_PROBE").as_deref() == Ok("1"))
555}
556
557/// GRAMMAR HOOK for constrained spec decode (lane/constrained-full, 2026-08-03). The engine
558/// stays llguidance-agnostic: the server adapts its per-session grammar state behind this
559/// trait. CONTRACT (the verify-side truncation rule — token-identical to constrained plain
560/// greedy decode): the exactness walk runs UNMASKED first; the hook then (a) truncates
561/// acceptance at the first grammar-illegal accepted token, and (b) when the truncation fired
562/// or the bonus is illegal, the engine recomputes that slot as the MASKED argmax of the
563/// target's own verify column (an unmasked argmax that is grammar-legal IS the masked argmax
564/// — masking only removes tokens — so the common case pays nothing). `consume` advances the
565/// state with each EMITTED token in order; EOS handling is the implementor's job (skip).
566pub trait SpecConstraint {
567 /// -inf the current state's banned ids on a HOST logits row (prompt-tail / init-feed
568 /// masked argmax).
569 fn mask_logits(&mut self, logits: &mut [f32]) -> Result<(), String>;
570 /// Packed 32-bit bitset words of the CURRENT state's allowed set (device-mask form).
571 fn mask_words(&mut self) -> Result<Vec<u32>, String>;
572 /// Is `tok` consumable in the CURRENT state?
573 fn is_allowed(&mut self, tok: u32) -> Result<bool, String>;
574 /// Advance the state with an emitted token.
575 fn consume(&mut self, tok: u32) -> Result<(), String>;
576
577 // --- DRAFT-SIDE MASKING (lane/draft-mask, 2026-08-04) ---
578 // The drafter proposed grammar-illegal tokens under tight schemas, so verify-side
579 // truncation cut nearly every round (measured acceptance 0.467-0.513 tight vs 0.62-0.82
580 // loose, research/constrained-full-20260803). These three methods let the engine mask the
581 // DRAFT model's own sampling with the grammar's legal set, so proposals are legal by
582 // construction. The state they walk is a SPECULATIVE CLONE of the session matcher — the
583 // real state is advanced only by `consume` (emitted tokens), so verify-side truncation
584 // stays the correctness backstop and the emitted stream is unchanged by construction
585 // (an accepted draft is the target's unmasked argmax AND grammar-legal, hence the masked
586 // argmax; a cut slot is recomputed as the masked argmax either way).
587 // Default impls = feature OFF (pre-lane behaviour: unmasked drafts).
588
589 /// Is draft-side masking available on this hook? Probed ONCE per burst, before the draft
590 /// graph is captured (the mask is an in-graph node — its presence is a capture-time shape).
591 fn draft_mask_enabled(&self) -> bool {
592 false
593 }
594 /// Start a draft chain: clone the CURRENT (committed) grammar state into the speculative
595 /// slot. Called once per spec round, before the first draft position.
596 fn draft_begin(&mut self) -> Result<(), String> {
597 Ok(())
598 }
599 /// Packed 32-bit bitset words of the SPECULATIVE state's allowed set (target-vocab ids),
600 /// for the draft position about to be sampled. `None` = draft masking off (no-op).
601 fn draft_mask_words(&mut self) -> Result<Option<Vec<u32>>, String> {
602 Ok(None)
603 }
604 /// Advance the SPECULATIVE state with a PROPOSED draft token. `false` = the chain cannot
605 /// continue (EOS proposed, or an unmasked position proposed something illegal) — the
606 /// engine stops drafting; the token already pushed still goes through verify.
607 fn draft_advance(&mut self, _tok: u32) -> Result<bool, String> {
608 Ok(false)
609 }
610}
611
612/// DRAFT-MASK UPLOAD (lane/draft-mask): pull the speculative state's allowed set (TARGET-id
613/// space) from the hook, project it into the DRAFT head's vocab space, and upload it into the
614/// stable device buffer the draft chain reads. Returns false when the chain must stop drafting:
615/// the hook handed out no mask, or NO draft-vocab row is grammar-legal at this position (a
616/// trimmed FR-Spec head genuinely cannot propose a legal token there — masking it would leave
617/// a fully-banned row whose argmax is meaningless, so the round drafts fewer tokens and the
618/// verify emits the masked argmax as usual).
619fn upload_draft_mask(
620 e: &Engine,
621 c: &mut dyn SpecConstraint,
622 dst: &mut CudaSlice<u32>,
623 d2t: Option<&Vec<u32>>,
624 d_vocab: usize,
625 words: usize,
626) -> Result<bool, Box<dyn std::error::Error>> {
627 let Some(tw) = c
628 .draft_mask_words()
629 .map_err(|e2| format!("constraint: {e2}"))?
630 else {
631 return Ok(false);
632 };
633 let bit = |t: usize| -> bool {
634 let w = t >> 5;
635 w < tw.len() && (tw[w] >> (t & 31)) & 1 == 1
636 };
637 let mut buf = vec![0u32; words];
638 match d2t {
639 // TRIMMED draft head: row i proposes target id d2t[i] — permute the mask accordingly.
640 Some(map) => {
641 for (i, &t) in map.iter().enumerate().take(d_vocab) {
642 if bit(t as usize) {
643 buf[i >> 5] |= 1u32 << (i & 31);
644 }
645 }
646 }
647 // UNTRIMMED: draft ids ARE target ids; the packed words transfer verbatim (a short
648 // mask leaves the padded tail zeroed == banned, same rule as constrained::apply_mask).
649 None => {
650 let n = tw.len().min(words);
651 buf[..n].copy_from_slice(&tw[..n]);
652 }
653 }
654 if buf.iter().all(|w| *w == 0) {
655 return Ok(false);
656 }
657 e.htod_u32_into(dst, &buf)?;
658 Ok(true)
659}
660
661/// Keep the full token-embedding table in host memory and upload only the rows needed by each
662/// MTP/verify step. This is an exact memory-capacity seam for very large BF16 vocab tables: host
663/// gather expands the same source bits to f32, and only O(T*n_embd) bytes cross PCIe per step.
664/// CUDA-graph/round-stream draft paths require device token ids and therefore stay disabled.
665pub(crate) fn spec_host_embd() -> bool {
666 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
667 *ON.get_or_init(|| std::env::var("MEMRA_SPEC_HOST_EMBD").as_deref() == Ok("1"))
668}
669
670/// VERIFY-TIER TRUNK LAUNCH-FUSION (default ON since 2026-07-09; MEMRA_SPEC_FUSED_T=0 reverts — lane/close35b): extend
671/// the t=1 fused2/fused3 Q8_0 trunk launches to the batched verify tier (t=2-4, the K=1..3
672/// verify shapes). At t>1 the trunk pairs/triples (35B wqkv+wqkv_gate, wq/wk/wv,
673/// gate_shexp+up_shexp) each run a separate `matmul_decode_exact` — one q8_1 re-quantize of the
674/// SAME activation plus one _b2/_b4 launch per tensor. The fused twins share ONE quantize and
675/// ONE launch per group; per (tensor,token,row) the kernel body is q8_0_mmvq_batched verbatim
676/// with the identical row mapping -> BIT-IDENTICAL by construction (kernel-check pins it,
677/// run-spec K=1..8 + acceptance identity arbitrate e2e).
678pub(crate) fn spec_fused_t() -> bool {
679 static F: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
680 // DEFAULT ON since 2026-07-09 (MEMRA_SPEC_FUSED_T=0 reverts): verify t=2-4 trunk launch-fusion
681 // (fused2/fused3 Q8_0 batched twins, bit-identical by construction — m=1 block-offset split on
682 // the batched body). m=2 marginal token 2117->1762us; 35B daily: p3 +3.7% (crosses llama), p2 +5%.
683 *F.get_or_init(|| {
684 std::env::var("MEMRA_SPEC_FUSED_T")
685 .map(|v| v != "0")
686 .unwrap_or(true)
687 })
688}
689
690/// zeros/uninit switch for verify-path buffers that are FULLY OVERWRITTEN before any read.
691/// Only call this on such buffers — the lean contract is "identical bytes by construction".
692/// TOKEN-ID GUARD for every id that reaches an embed gather (#87 family).
693///
694/// A device argmax seeds its running index with 0x7FFFFFFF and replaces it only through
695/// comparisons, all of which are FALSE against NaN. An all-NaN logits row therefore returns
696/// the sentinel, and the next thing done with a token id is `embed_row(id)` — table +
697/// ~4.6 TB, never mapped, an MMU fault that kills the CUDA context for the whole process
698/// (research/pp2spec-crash-20260807). The draft chain and the GREEDY verify walk already
699/// trap this; the SAMPLED verify bonus, the boundary sampler and the replay arm's last_pred
700/// did not, which is why the recoverable fault on the greedy instrument is a TERMINAL one on
701/// the vendor-default sampled shape we actually serve.
702pub(crate) fn guard_vocab_token(
703 tok: u32,
704 n_vocab: usize,
705 what: &str,
706) -> Result<u32, Box<dyn std::error::Error>> {
707 if (tok as usize) >= n_vocab {
708 return Err(format!(
709 "{what}: token id 0x{tok:08x} >= n_vocab {n_vocab} — an all-NaN logits row left \
710 the device argmax's init sentinel in place; refusing to dereference the embed \
711 row (#87 trap)"
712 )
713 .into());
714 }
715 Ok(tok)
716}
717
718/// SPEC NaN-ORIGIN SCAN (`MEMRA_SPEC_NAN_SCAN=1`, DEFAULT OFF, diagnostic only).
719///
720/// The `#87` trap reports an all-NaN VERIFY logits column, which says the poison reached the
721/// head but not where it entered. With the scan armed the verify walk syncs and reads back
722/// every layer's output, so the FIRST layer whose residual carries a NaN names itself with the
723/// round's row and position. Off by default and never on a serving path: it costs one host
724/// sync + one `t*n_embd` D2H per layer, and the syncs change scheduling (so a run that stops
725/// reproducing under the scan is itself a datum, not an all-clear).
726///
727/// Rollback seam: unset `MEMRA_SPEC_NAN_SCAN` (or set it to 0). Every call site is behind
728/// `spec_nan_scan()`, so the default path keeps the exact launch sequence it had.
729pub(crate) fn spec_nan_scan() -> bool {
730 spec_nan_scan_level() > 0
731}
732
733/// `MEMRA_SPEC_NAN_SCAN` as a LEVEL, not a boolean. `1` scans each layer's residual, which
734/// names the layer. `2` also scans INSIDE the t-column layer body — the per-column attention
735/// output, the deferred-column o-proj/fa2 join, the post-attention norm and the routed-MoE
736/// output — because "layer 20 poisons row 0" does not say whether the attention or the routed
737/// MoE produced it, and those are different bugs with different fixes.
738pub(crate) fn spec_nan_scan_level() -> u8 {
739 static LVL: std::sync::OnceLock<u8> = std::sync::OnceLock::new();
740 *LVL.get_or_init(|| match std::env::var("MEMRA_SPEC_NAN_SCAN").as_deref() {
741 Ok("1") => 1,
742 Ok("2") => 2,
743 _ => 0,
744 })
745}
746
747/// Read back `[rows, cols]` and fail with the first NaN's coordinates. `what` names the
748/// producer (layer index, walk arm) so the error line is the localization.
749/// VERIFY-ARM RECEIPT (rides `MEMRA_SPEC_NAN_SCAN>=1`, bounded to 200 lines).
750///
751/// Names, per trunk layer, WHICH attention arm the t-column walk actually took. This exists
752/// because the level-1 residual scan below sat only on the non-fused tail: the fused
753/// rope+append+fa arm ends in `continue`, so every layer that fused was NEVER SCANNED and
754/// silently read as "clean". A poisoned residual therefore first reported at the next
755/// non-fused layer, which is how "layer 20 creates the poison" could be true of the scan and
756/// false of the engine. Also carries the row-table lookup counter, so "the fused path never
757/// ran" is distinguishable from "it ran and was innocent".
758/// KV-PLANE SCAN (`MEMRA_KV_PLANE_SCAN=1`, DEFAULT OFF, diagnostic only).
759///
760/// Reads back the STAGED rows of a layer's distributed K/V planes and reports the first row
761/// whose quantization scale is not finite. No kernel required: q8_0 blocks are
762/// `[half d][32 x i8]` and q5_1 blocks carry `half d` then `half m`, so the fp16 scale at the
763/// head of each block is host-checkable straight out of the byte plane.
764///
765/// It exists because the level-2 bad-row bitmap says EVERY verify row is non-finite at a
766/// global-attention layer's join, and row r attends a strict superset of row r-1's keys: that
767/// implicates the shared KV history those rows walk, not per-column staging. "The attention
768/// output is NaN" and "the KV history it attends is already NaN" are different bugs with
769/// different owners, and nothing measured so far separates them. A first-corrupt-row index
770/// also dates the corruption against the prime/decode boundary.
771///
772/// Bounded hard: only layers whose geometry has NO window (the global planes), only the first
773/// `MEMRA_KV_PLANE_SCAN_ROUNDS` verify rounds of a process (default 2), and it copies only
774/// `[0, staged_len)`, which is ~1.6 MB at the 1480-token repro rather than the 262144-row
775/// provision. It still syncs per layer, so it is never a serving or a measured-perf arm.
776pub(crate) fn kv_plane_scan_on() -> bool {
777 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
778 *ON.get_or_init(|| std::env::var("MEMRA_KV_PLANE_SCAN").as_deref() == Ok("1"))
779}
780
781fn kv_plane_scan_rounds() -> usize {
782 static R: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
783 *R.get_or_init(|| {
784 std::env::var("MEMRA_KV_PLANE_SCAN_ROUNDS")
785 .ok()
786 .and_then(|v| v.parse().ok())
787 .unwrap_or(2)
788 })
789}
790
791/// First non-finite fp16 block scale in `bytes`, as (block index, raw u16), scanning one
792/// scale every `stride` bytes. Returns None when every block scale is finite.
793fn first_bad_scale(bytes: &[u8], stride: usize) -> Option<(usize, u16)> {
794 if stride == 0 {
795 return None;
796 }
797 for (i, blk) in bytes.chunks_exact(stride).enumerate() {
798 let raw = u16::from_le_bytes([blk[0], blk[1]]);
799 if half_is_non_finite(raw) {
800 return Some((i, raw));
801 }
802 }
803 None
804}
805
806/// IEEE binary16: exponent all ones is Inf or NaN, whatever the mantissa says.
807fn half_is_non_finite(raw: u16) -> bool {
808 (raw & 0x7C00) == 0x7C00
809}
810
811/// Scan one layer's staged K/V planes for a non-finite quantization scale. Returns the
812/// receipt line, or None when the layer is out of scope or every scale is finite.
813pub(crate) fn scan_kv_plane(
814 e: &crate::Engine,
815 distributed: &memra_kv::ResidentTpKvCache,
816 il: usize,
817 pos0: usize,
818) -> Result<(), Box<dyn std::error::Error>> {
819 // One "round" is one pos0, not one layer: the walk visits 45 layers per verify. The
820 // default of 2 rounds is for a fault that shows up immediately; the step37 repro does not
821 // fire until rep 3 or later, i.e. round ~60 of the process, so that arm MUST raise
822 // MEMRA_KV_PLANE_SCAN_ROUNDS or it will scan only the two rounds that were never going to
823 // be poisoned and report a clean history it never looked at.
824 static ROUNDS: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
825 static LAST_POS: std::sync::atomic::AtomicUsize =
826 std::sync::atomic::AtomicUsize::new(usize::MAX);
827 if LAST_POS.swap(pos0, std::sync::atomic::Ordering::Relaxed) != pos0 {
828 ROUNDS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
829 }
830 if ROUNDS.load(std::sync::atomic::Ordering::Relaxed) > kv_plane_scan_rounds() {
831 return Ok(());
832 }
833 let staged = distributed.staged_len();
834 if staged == 0 {
835 return Ok(());
836 }
837 // ENGAGEMENT RECEIPT. This scan prints only on corruption, so `kvbad=0` in a cell would
838 // read the same whether the history was clean or the scan never ran once. Bounded so a
839 // 45-layer walk cannot flood the log.
840 static SEEN: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
841 let seen = SEEN.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
842 let (ktb, vtb) = (distributed.k_tok_bytes(), distributed.v_tok_bytes());
843 if seen < 4 {
844 eprintln!(
845 "[kv-plane] engaged #{seen} layer {il} pos0={pos0} staged={staged} \
846 ktok={ktb} vtok={vtb} (scan armed; a corrupt plane prints its own line)"
847 );
848 }
849 for rank in 0..distributed.ranks().len() {
850 let Some(rc) = distributed.rank(rank) else {
851 continue;
852 };
853 // q8_0 K blocks are [half d][32 x i8] = 34B; q5_1 V blocks lead with half d then half m.
854 let kbytes = e.dtoh_u8_view(&rc.k().slice(0..staged * ktb))?;
855 let vbytes = e.dtoh_u8_view(&rc.v().slice(0..staged * vtb))?;
856 let kbad = first_bad_scale(&kbytes, 34);
857 let vbad = first_bad_scale(&vbytes, 24);
858 if kbad.is_some() || vbad.is_some() {
859 let row = |b: Option<(usize, u16)>, tok: usize| {
860 b.map(|(i, raw)| format!("blk {i} (row {}) raw={raw:#06x}", i * 34 / tok.max(1)))
861 .unwrap_or_else(|| "clean".into())
862 };
863 eprintln!(
864 "[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",
865 row(kbad, ktb),
866 row(vbad, vtb)
867 );
868 return Ok(());
869 }
870 }
871 Ok(())
872}
873
874pub(crate) fn verify_arm_receipt(
875 arm: &str,
876 il: usize,
877 pos0: usize,
878 t: usize,
879 staged: Option<usize>,
880) {
881 static N: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
882 if N.fetch_add(1, std::sync::atomic::Ordering::Relaxed) >= 200 {
883 return;
884 }
885 eprintln!(
886 "[verify-arm] layer {il} arm={arm} pos0={pos0} t={t} staged_len={} rows_tab_lookups={}",
887 staged.map(|v| v as i64).unwrap_or(-1),
888 crate::tp::ROWS_TAB_ENGAGED.load(std::sync::atomic::Ordering::Relaxed)
889 );
890}
891
892pub(crate) fn nan_scan_rows(
893 e: &Engine,
894 buf: &CudaSlice<f32>,
895 rows: usize,
896 cols: usize,
897 what: &str,
898) -> Result<(), Box<dyn std::error::Error>> {
899 // The readback is also the ATTRIBUTION point for an asynchronous fault: a
900 // CUDA_ERROR_ILLEGAL_ADDRESS raised by any launch since the previous scan surfaces on this
901 // sync, and the bare DriverError names nothing. Wrapping it with `what` turns "the process
902 // died somewhere" into "it died at or before this layer, on this row, at this position".
903 let host = e.dtoh(buf).map_err(|err| -> Box<dyn std::error::Error> {
904 format!(
905 "spec nan-scan: sync at {what} FAILED: {err} — the fault is at or before \
906 this point in the walk"
907 )
908 .into()
909 })?;
910 if host.len() < rows * cols {
911 return Err(format!(
912 "nan-scan {what}: buffer holds {} < {rows}x{cols}",
913 host.len()
914 )
915 .into());
916 }
917 // SCAN EVERY ROW BEFORE REPORTING. A first-hit return says "row 0 is bad" and leaves the
918 // other rows UNEXAMINED, which is exactly the bit that discriminates the two mechanisms: in
919 // the t-column verify, row 0 attends keys [0..p+1) and row 1 attends [0..p+2), a strict
920 // superset, so poison in the SHARED KV history must appear in BOTH rows, while poison in
921 // per-column staging can appear in one. Report the whole map.
922 let mut per_row: Vec<usize> = Vec::with_capacity(rows);
923 let mut first_bad: Option<(usize, usize)> = None;
924 for r in 0..rows {
925 let row = &host[r * cols..(r + 1) * cols];
926 let bad = row.iter().filter(|v| !v.is_finite()).count();
927 per_row.push(bad);
928 if bad > 0 && first_bad.is_none() {
929 first_bad = Some((r, row.iter().position(|v| !v.is_finite()).unwrap_or(0)));
930 }
931 }
932 if let Some((r0, c0)) = first_bad {
933 let map: String = per_row
934 .iter()
935 .map(|&b| if b == 0 { '.' } else { 'X' })
936 .collect();
937 return Err(format!(
938 "spec nan-scan: {what} produced non-finite values — rows[{rows}] map={map} \
939 counts={per_row:?} of {cols} each; first at row {r0} element {c0}. Both rows bad \
940 implicates shared state (the KV history this layer reads); one row bad implicates \
941 per-column staging."
942 )
943 .into());
944 }
945 Ok(())
946}
947
948fn vbuf(e: &Engine, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
949 if spec_lean() { e.uninit(n) } else { e.zeros(n) }
950}
951
952/// Scratch KV for the MTP block (one full-attn layer).
953///
954/// PERSISTENT MODE (default, 2026-07-03 — the acceptance lever): sized cap = max_ctx and kept in
955/// sync with the COMMITTED sequence — slot p holds the MTP block's K/V for committed token p
956/// (roped p+1, the chain's rope convention), so the draft chain's self-attention sees the FULL
957/// committed history instead of only the current round's 1..K+1 chain tokens (the reference
958/// engine's "mtp_update" design). Entries come from two sources:
959/// - chain appends: accepted positions KEEP their chain-computed entries (embedding exact,
960/// hidden chain-approximate — the reference engine accepts the same);
961/// - `mtp_kv_fill` batches: prompt positions + the last-draft position on full accept, computed
962/// from EXACT trunk hiddens (K/V-only MTP-block pass, no attention/FFN/lm_head).
963/// Rejected drafts / p-min extras / pseudo-seed appends are all discarded by the round-start
964/// `set_len` truncation (the KvLayer len mechanism — §C rollback for the draft side).
965/// Multi-turn spec-decode session (2026-07-05): trunk Cache + persistent MTP draft scratch +
966/// the committed token list, alive across generate_spec_session calls. Turn N+1 primes ONLY its
967/// suffix (chunked continuation prime over the quantized past) and mtp_kv_fill's its suffix rows,
968/// then the round loop runs unchanged. `last_h` carries the pre-output_norm hidden of the last
969/// committed row across turns (the predecessor-pairing seed + fill anchor).
970/// Per-request sampling config for the sampled-spec serve path.
971#[derive(Clone, Copy, Debug)]
972pub struct SpecSampling {
973 pub temp: f32,
974 pub seed: u64,
975 pub top_k: i32, // 0 = off
976 pub top_p: f32, // 1.0 = off
977 pub min_p: f32, // 0.0 = off
978 pub penalty_last_n: usize, // 0 = penalties off
979 pub penalty_repeat: f32,
980 pub penalty_freq: f32,
981 pub penalty_present: f32,
982}
983
984impl SpecSampling {
985 /// Non-identity penalties requested — THE `pen_on` predicate (one definition; the
986 /// same group-off rule `SamplerIdentity::of` canonicalizes: a window with neutral
987 /// coefficients is penalties-absent). Both spec routes and the dspark accept walk
988 /// key their penalty arms off this.
989 pub fn pen_on(&self) -> bool {
990 self.penalty_last_n > 0
991 && (self.penalty_repeat != 1.0
992 || self.penalty_freq != 0.0
993 || self.penalty_present != 0.0)
994 }
995}
996
997/// Host Philox4x32-10 uniform in (0,1) — mirrors spec_sample.cu's `philox4`/`u01` with the
998/// ctr_lo tag 0xFFFF_FFFE, so the host accept-test stream never collides with any device
999/// sampling event (device Gumbel uses (i>>2, stream_pos); device residual uses 0xFFFF_FFFD).
1000/// One value per (seed, ctr) EVENT; callers own the counter discipline. Extracted verbatim
1001/// from generate_spec_inner2's closure for the dspark sampled-admission walk (the two paths
1002/// MUST consume the identical stream construction — two ad-hoc Philox copies drifting apart
1003/// is a distributional bug, not a style problem).
1004pub(crate) fn host_u01(seed: u64, ctr: u32) -> f32 {
1005 let (m0, m1) = (0xD2511F53u32, 0xCD9E8D57u32);
1006 let (mut c0, mut c1, mut c2, mut c3) = (0xFFFF_FFFEu32, ctr, 0u32, 0u32);
1007 let (mut k0, mut k1) = ((seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
1008 for _ in 0..10 {
1009 let (h0, l0) = (((m0 as u64 * c0 as u64) >> 32) as u32, m0.wrapping_mul(c0));
1010 let (h1, l1) = (((m1 as u64 * c2 as u64) >> 32) as u32, m1.wrapping_mul(c2));
1011 let (n0, n1, n2, n3) = (h1 ^ c1 ^ k0, l1, h0 ^ c3 ^ k1, l0);
1012 c0 = n0;
1013 c1 = n1;
1014 c2 = n2;
1015 c3 = n3;
1016 k0 = k0.wrapping_add(0x9E3779B9);
1017 k1 = k1.wrapping_add(0xBB67AE85);
1018 }
1019 (c0 as f32 + 1.0) * (1.0 / 4294967296.0)
1020}
1021
1022/// Tracked draft positions for [`SpecTelemetry`] (serve K defaults to 3; the run-spec gate
1023/// sweeps K=1..8, and MEMRA_SPEC_CAPMAX defaults to 7 — 8 covers every tuned config).
1024pub const SPEC_TELEM_POS: usize = 8;
1025
1026/// Always-on per-draft-position acceptance telemetry (lane/accept-telemetry, 2026-08-05 —
1027/// the llama.cpp #26389 / vLLM spec-decode counter schema, upstream-sweeps 2026-08-05).
1028/// Lives on the [`SpecSession`] and accumulates across bursts; the serve worker diffs a
1029/// stashed copy per burst for its per-model /metrics aggregation and per-request usage.
1030/// Same normalization as the `[spec-stats]` line: p-min-discarded chain tokens are counted
1031/// in NEITHER drafted nor accepted.
1032#[derive(Clone, Copy, Default, Debug)]
1033pub struct SpecTelemetry {
1034 /// verify rounds completed (a round-stream burst counts each of its M rounds).
1035 pub rounds: u64,
1036 /// tokens drafted / accepted across all rounds.
1037 pub drafted: u64,
1038 pub accepted: u64,
1039 /// how often draft position j (0-based within a round's chain) was offered / accepted.
1040 /// Positions >= SPEC_TELEM_POS are untracked (totals still count them). The opt-in
1041 /// round-stream arm (MEMRA_SPEC_STREAM=1) reads back only totals, so under it these
1042 /// arrays cover the standard-path rounds only and their sums may undercount the totals.
1043 pub pos_drafted: [u64; SPEC_TELEM_POS],
1044 pub pos_accepted: [u64; SPEC_TELEM_POS],
1045}
1046
1047impl SpecTelemetry {
1048 /// Fieldwise `self - prev` — the worker's per-burst delta off a copy stashed before the
1049 /// burst call. Saturating: a caller diffing against the wrong snapshot gets zeros, not
1050 /// a wrapped counter.
1051 pub fn delta_since(&self, prev: &SpecTelemetry) -> SpecTelemetry {
1052 let mut d = SpecTelemetry {
1053 rounds: self.rounds.saturating_sub(prev.rounds),
1054 drafted: self.drafted.saturating_sub(prev.drafted),
1055 accepted: self.accepted.saturating_sub(prev.accepted),
1056 ..Default::default()
1057 };
1058 for j in 0..SPEC_TELEM_POS {
1059 d.pos_drafted[j] = self.pos_drafted[j].saturating_sub(prev.pos_drafted[j]);
1060 d.pos_accepted[j] = self.pos_accepted[j].saturating_sub(prev.pos_accepted[j]);
1061 }
1062 d
1063 }
1064 /// Fieldwise `self += d` — the worker's per-model aggregation.
1065 pub fn merge(&mut self, d: &SpecTelemetry) {
1066 self.rounds += d.rounds;
1067 self.drafted += d.drafted;
1068 self.accepted += d.accepted;
1069 for j in 0..SPEC_TELEM_POS {
1070 self.pos_drafted[j] += d.pos_drafted[j];
1071 self.pos_accepted[j] += d.pos_accepted[j];
1072 }
1073 }
1074
1075 /// Mean accepted draft-prefix length per verify round (tau).
1076 pub fn tau(&self) -> f64 {
1077 if self.rounds > 0 {
1078 self.accepted as f64 / self.rounds as f64
1079 } else {
1080 0.0
1081 }
1082 }
1083}
1084
1085/// Session-lifetime atomic acceptance counters. The verifier records only after the greedy or
1086/// rejection-sampling walk has resolved on the host, so these relaxed increments add no GPU
1087/// launch, synchronization, allocation, or ordering dependency to the numeric path.
1088struct SpecTelemetryCounters {
1089 rounds: AtomicU64,
1090 drafted: AtomicU64,
1091 accepted: AtomicU64,
1092 pos_drafted: [AtomicU64; SPEC_TELEM_POS],
1093 pos_accepted: [AtomicU64; SPEC_TELEM_POS],
1094}
1095
1096impl Default for SpecTelemetryCounters {
1097 fn default() -> Self {
1098 Self {
1099 rounds: AtomicU64::new(0),
1100 drafted: AtomicU64::new(0),
1101 accepted: AtomicU64::new(0),
1102 pos_drafted: std::array::from_fn(|_| AtomicU64::new(0)),
1103 pos_accepted: std::array::from_fn(|_| AtomicU64::new(0)),
1104 }
1105 }
1106}
1107
1108impl SpecTelemetryCounters {
1109 fn record_round(&self, drafted: usize, accepted: usize) {
1110 debug_assert!(accepted <= drafted);
1111 self.rounds.fetch_add(1, Ordering::Relaxed);
1112 self.drafted.fetch_add(drafted as u64, Ordering::Relaxed);
1113 self.accepted.fetch_add(accepted as u64, Ordering::Relaxed);
1114 for counter in self.pos_drafted.iter().take(drafted) {
1115 counter.fetch_add(1, Ordering::Relaxed);
1116 }
1117 for counter in self.pos_accepted.iter().take(accepted) {
1118 counter.fetch_add(1, Ordering::Relaxed);
1119 }
1120 }
1121
1122 /// Round-stream keeps each round's accept length on device; retain exact scalar totals while
1123 /// leaving the per-position arrays untouched, matching the pre-existing telemetry contract.
1124 fn record_totals(&self, rounds: usize, drafted: usize, accepted: usize) {
1125 self.rounds.fetch_add(rounds as u64, Ordering::Relaxed);
1126 self.drafted.fetch_add(drafted as u64, Ordering::Relaxed);
1127 self.accepted.fetch_add(accepted as u64, Ordering::Relaxed);
1128 }
1129
1130 fn snapshot(&self) -> SpecTelemetry {
1131 SpecTelemetry {
1132 rounds: self.rounds.load(Ordering::Relaxed),
1133 drafted: self.drafted.load(Ordering::Relaxed),
1134 accepted: self.accepted.load(Ordering::Relaxed),
1135 pos_drafted: std::array::from_fn(|j| self.pos_drafted[j].load(Ordering::Relaxed)),
1136 pos_accepted: std::array::from_fn(|j| self.pos_accepted[j].load(Ordering::Relaxed)),
1137 }
1138 }
1139}
1140
1141pub struct SpecSession {
1142 pub(crate) cache: Cache,
1143 pub(crate) scratch: MtpScratch,
1144 /// Every token whose state the caches hold, in order (prompt turns + generated), INCLUDING
1145 /// overshoot: spec commits accepted drafts past max_new; those rows are in the caches, so the
1146 /// session must count them. Callers render output from this, not from their own echo.
1147 pub committed: Vec<u32>,
1148 /// Pre-output_norm hidden of the LAST committed row (device). None before the first turn.
1149 pub(crate) last_h: Option<CudaSlice<f32>>,
1150 /// Greedy argmax predicting the token AFTER committed.last() (from the last turn's final
1151 /// logits). Fuels empty-suffix continuation bursts (serve): the next turn emits this token
1152 /// first, feeds it, and the round loop resumes without any prime. None before the first turn.
1153 pub next_pred: Option<u32>,
1154 /// SAMPLED-SPEC stream continuity across bursts: Philox event counters persist here so a
1155 /// session's randomness never repeats between generate_spec_session calls. (0,0) at admit.
1156 pub sctr: u32,
1157 pub uctr: u32,
1158 /// PERSISTENT DRAFT-GRAPH CONTEXT (2026-08-01, the serve-burst fixed-cost fix): the captured
1159 /// draft graph(s) + every device I/O buffer they bake, carried ACROSS generate_spec_session
1160 /// calls. Before this, every serve burst re-captured the draft graph (2 warmup forwards +
1161 /// instantiate) — measured ~16ms/burst on H100 q27 (MEMRA_SPEC_BURST sweep,
1162 /// research/spec-serving-20260801). None before the first turn; error paths drop it
1163 /// (next burst recaptures — serve retires errored sessions anyway).
1164 pub(crate) draft_ctx: Option<DraftGraphCtx>,
1165 /// PENDING-CARRY across bursts (2026-08-01, the serve burst-boundary fix): the bonus token
1166 /// emitted by the last round but NOT committed to the caches. The old tail committed it with
1167 /// a solo T=1 trunk pass (+ draft fill), and the next burst's setup fed the stashed next_pred
1168 /// with ANOTHER solo pass — 2x ~11.5ms/burst measured on H100 q27 ([spec-setup] trace).
1169 /// Carrying it lets the next empty-suffix greedy burst consume it as round-0 verify col 0,
1170 /// exactly like a mid-burst full-accept boundary (no solo passes). INVARIANT: when set,
1171 /// `committed` (== cache rows) EXCLUDES this token although it was already emitted in the
1172 /// last burst's output, and `last_h` holds the hidden of the last COMMITTED row (its
1173 /// predecessor — the chain-seed/fill anchor). `next_pred` is None (unknown without the
1174 /// commit pass). Non-empty-suffix or sampled turns must flush first (spec_flush_pending);
1175 /// generate_spec_session_sampled does this at entry, and serve parks only flushed sessions.
1176 pub pending_tok: Option<u32>,
1177 /// SESSION-AFFINITY TURN CHECKPOINT (lane/session-affinity, 2026-08-05): the state at this
1178 /// turn's PROMPT-END boundary, retained so a later turn can REWIND here. See
1179 /// [`SpecCheckpoint`]. Refreshed by every non-empty prime; None until the first one, and on
1180 /// a rig too tight to hold it (a failed capture is silent — resume just isn't available).
1181 pub(crate) turn_ckpt: Option<SpecCheckpoint>,
1182 /// Session-lifetime acceptance telemetry. Relaxed atomics update at the host-side round
1183 /// accounting the loop already does — no syncs, no allocation. NOTE a
1184 /// pool-resumed session carries the PREVIOUS requests' counts; per-request consumers
1185 /// diff with [`SpecTelemetry::delta_since`] around each burst.
1186 telem: SpecTelemetryCounters,
1187 /// PREFIX-CACHE publication request (lane/spec-prefix-cache): worker sets this to the
1188 /// miss-LCP boundary before a cold burst; the prime captures at exactly that split (it must
1189 /// coincide with the burst's `prime_split` or no capture happens). One-shot: consumed by the
1190 /// prime, result lands in `boundary_captures`.
1191 pub capture_at: Option<usize>,
1192 /// The captures the last prime produced (see [`SpecBoundaryCapture`]). Worker drains them
1193 /// post-burst to assemble prefix entries. A failed capture is silent, like `turn_ckpt` —
1194 /// publication just isn't available for that request. Plural since
1195 /// lane/frspec-multiturn-cache (2026-08-21): a cold burst can capture BOTH the miss-LCP
1196 /// split (the shared-prefix class) and the stable pre-generation boundary (the
1197 /// next-turn re-render class) — one entry per stop, exactly the boundary set the plain
1198 /// prefill tick publishes/checkpoints.
1199 pub boundary_captures: Vec<SpecBoundaryCapture>,
1200 /// STABLE-BOUNDARY TURN CHECKPOINT REQUEST (lane/frspec-multiturn-cache, 2026-08-21): the
1201 /// ABSOLUTE committed-length position the next non-empty prime should capture `turn_ckpt`
1202 /// at, instead of prompt-end. The worker sets it to the STABLE PRE-GENERATION boundary
1203 /// (`plain_checkpoint_boundary` — before the live generation header the client rewrites),
1204 /// porting the 2026-08-09 plain-tier fix: a prompt-end spec checkpoint includes the
1205 /// template's live assistant-generation header (`<|im_start|>assistant\n<think>\n`), which
1206 /// the NEXT turn's re-render replaces, so `affinity_match` diverged a couple tokens below
1207 /// the checkpoint and the spec pool declined 100% of multi-turn agent traffic (measured:
1208 /// `spec-affinity: declined (history diverged at 6811 of checkpoint 6813)`,
1209 /// research/multiturn-cache-20260821 B4). One-shot, `capture_at` convention; None = legacy
1210 /// prompt-end capture.
1211 pub ckpt_at: Option<usize>,
1212}
1213impl SpecSession {
1214 /// Context capacity of the session's caches (the server's ContextFull guard).
1215 pub fn cache_max_ctx(&self) -> usize {
1216 self.cache.max_ctx
1217 }
1218 /// Read access to the live trunk cache (lane/spec-prefix-cache): the worker slices
1219 /// full-attn KV rows `[0..capture.pos)` out of it when publishing a boundary capture —
1220 /// those rows are append-only for the session's lifetime (rollbacks never truncate below
1221 /// the prime boundary), so no copy was taken at prime time.
1222 pub fn cache_ref(&self) -> &Cache {
1223 &self.cache
1224 }
1225 /// Read access to the persistent draft-scratch plane (lane/spec-on-cache-hit): the
1226 /// worker slices rows `[0..capture.pos)` when publishing a boundary capture, exactly
1227 /// like the trunk KV — draft rows below the prompt end are append-only for the
1228 /// session's lifetime (the prime fill wrote them once; rollbacks reset `len_d` to the
1229 /// committed length, never below the prime boundary, and the true-hidden refresh
1230 /// rewrites generated positions only). Returns `(k, v, k_tok_bytes, v_tok_bytes)`.
1231 /// None when the scratch is ring-backed (Step35 SWA — physical rows are not
1232 /// prefix-addressable; the prefix cache already refuses that class end to end).
1233 pub fn draft_plane_ref(&self) -> Option<(&CudaSlice<u8>, &CudaSlice<u8>, usize, usize)> {
1234 if self.scratch.kv.ring.is_some() {
1235 return None;
1236 }
1237 Some((
1238 &self.scratch.kv.k,
1239 &self.scratch.kv.v,
1240 self.scratch.kv.k_tok_bytes,
1241 self.scratch.kv.v_tok_bytes,
1242 ))
1243 }
1244 /// Snapshot the session's process-local acceptance counters for per-burst diffing.
1245 pub fn telemetry(&self) -> SpecTelemetry {
1246 self.telem.snapshot()
1247 }
1248 /// Committed position this session can REWIND to (its retained prompt-end boundary), if any.
1249 /// A request whose prompt matches `committed[..pos]` exactly can resume from here — see
1250 /// `spec_rewind_to_checkpoint`.
1251 pub fn rewind_pos(&self) -> Option<usize> {
1252 self.turn_ckpt.as_ref().map(|c| c.pos)
1253 }
1254 /// Whether every ring-backed trunk/draft row needed by the retained checkpoint is resident.
1255 pub fn rewind_is_resident(&self) -> bool {
1256 self.turn_ckpt.as_ref().is_some_and(|ckpt| {
1257 self.cache.can_rollback(&ckpt.snap, 0) && self.scratch.can_rewind_to(ckpt.pos)
1258 })
1259 }
1260 /// Is this session in the DEMOTION-READY shape (see [`SpecSession::into_demoted`])?
1261 /// `false` means a carried pending must be flushed first (`spec_flush_pending`), or the
1262 /// session has never run a turn and has no prediction to hand over.
1263 pub fn demote_ready(&self) -> bool {
1264 self.pending_tok.is_none() && self.next_pred.is_some()
1265 }
1266 /// Does this session hold a carried pending bonus (flush required before a handoff/park)?
1267 pub fn has_pending(&self) -> bool {
1268 self.pending_tok.is_some()
1269 }
1270 /// Committed row count == cache rows (the session invariant), for the caller's own
1271 /// `fed`-length cross-check at a handoff boundary.
1272 pub fn committed_len(&self) -> usize {
1273 self.committed.len()
1274 }
1275 /// DEMOTION HANDOFF (lane/spec-gate, 2026-08-07): consume this session and hand its trunk
1276 /// cache + next-token prediction to the plain batched-decode path.
1277 ///
1278 /// WHY THIS IS EXACT (greedy). The invariant at a burst boundary is `cache.pos ==
1279 /// committed.len()`: every committed row has trunk KV + recurrent state, exactly as a plain
1280 /// tokenwise prime of the same `committed` sequence would have left it (that is the
1281 /// session-tail contract, and the same property `spec_rewind_to_checkpoint` and the reuse
1282 /// pool already rely on). `next_pred` is the argmax of the verify's logits for the LAST
1283 /// committed row — and verify-column logits are bit-identical to plain decode's logits at
1284 /// that position, because `matmul_decode_exact` bit-identity IS the basis of the greedy
1285 /// accept walk. So handing (cache, next_pred) to the batched path continues the stream from
1286 /// a state indistinguishable from one the batched path produced itself: the batched tick
1287 /// emits `next_pred`, feeds it into this same cache, and decodes on.
1288 ///
1289 /// `None` when the session is not in the handoff shape — a carried pending (its bonus row is
1290 /// NOT in the cache, so `spec_flush_pending` must commit it first) or no `next_pred` yet
1291 /// (never bursted). Callers must not force it: a half-committed cache handed to the batched
1292 /// path would silently skip a token.
1293 ///
1294 /// The MTP draft scratch, the persistent draft-graph context and the turn checkpoint are
1295 /// DROPPED here (freeing their VRAM): the batched path never drafts, and this handoff is
1296 /// one-way by design — there is no cheap symmetric re-promotion (rebuilding the draft KV
1297 /// would mean an `mtp_kv_fill` over the whole committed history).
1298 pub fn into_demoted(self) -> Option<(Cache, u32)> {
1299 if self.pending_tok.is_some() {
1300 return None;
1301 }
1302 let np = self.next_pred?;
1303 debug_assert_eq!(
1304 self.cache.pos,
1305 self.committed.len(),
1306 "demotion handoff: cache rows != committed tokens"
1307 );
1308 Some((self.cache, np))
1309 }
1310 /// Pool-resume hook (audit Q2): clear the parked draft-graph failure memoization so a
1311 /// NEW request resuming this session gets one fresh capture chance — a transient-pressure
1312 /// capture failure must not persist for the pool's whole lifetime (the TRT #16072 class).
1313 /// Logs once iff a flag was actually set; a no-fallback resume is silent and free.
1314 pub fn reset_graph_fallback_on_resume(&mut self) {
1315 if let Some(line) = self
1316 .draft_ctx
1317 .as_mut()
1318 .and_then(|c| c.failed.reset_on_resume())
1319 {
1320 eprintln!("{line}");
1321 }
1322 }
1323}
1324
1325/// A session's PROMPT-END boundary state, the rewind target for session-affinity resume.
1326///
1327/// WHY THIS BOUNDARY, AND WHY IT IS THE ONLY ONE WORTH KEEPING. The rewrite class this lane
1328/// exists for (a client that strips `<think>` blocks out of prior assistant turns) mutates the
1329/// text the session GENERATED, never the prompt it was given. So turn N's prompt agrees with
1330/// turn N-1's committed tokens up to almost exactly where turn N-1's generation began — the
1331/// prompt-end boundary. Keeping a checkpoint there means the next turn re-primes only its own
1332/// delta (the rewritten answer + the new user turn) instead of the whole conversation.
1333///
1334/// WHAT IT MUST HOLD. Full-attn KV is append-only and position-addressed, so rewinding it is a
1335/// `len` truncation (no data). Linear-attn (GDN) conv/ssm state is mutated IN PLACE with no
1336/// position index, so it must be a real device COPY — that copy is the entire reason a spec
1337/// session could not previously rewind. The MTP draft scratch needs no copy either: its rows
1338/// below the boundary were written by this turn's fill and are never revisited (the per-round
1339/// true-hidden refresh only rewrites the CURRENT burst's committed positions), so rewinding it
1340/// is also just a `len` reset. `last_h` is the hidden of the last row below the boundary — the
1341/// predecessor-pairing anchor the next prime's fill reads for its first row.
1342///
1343/// COST: one `Cache::snapshot` per TURN, on a code path that already takes one per ROUND.
1344pub(crate) struct SpecCheckpoint {
1345 snap: crate::cache::CacheSnapshot,
1346 /// Committed length at the boundary (== cache.pos there, the session invariant).
1347 pos: usize,
1348 /// Pre-output_norm hidden of row `pos - 1`.
1349 last_h: CudaSlice<f32>,
1350}
1351
1352/// PREFIX-CACHE BOUNDARY CAPTURE (lane/spec-prefix-cache, 2026-08-14): the state a spec session
1353/// records at its cold-prime split so the WORKER can publish a cross-request prefix entry —
1354/// the commit-gated-publication port (research/cache-spec-design-20260814/PORT-PLAN.md item 1).
1355/// Only the pieces that are DESTROYED by continuing the prime need copies here: the in-place
1356/// GDN conv/ssm states (via `Cache::snapshot`, same mechanism as [`SpecCheckpoint`]) and the
1357/// boundary logits. Full-attn KV rows `[0..pos)` and draft-scratch rows `[0..pos)` are
1358/// append-only for the session's lifetime (rollbacks never truncate below the prime boundary),
1359/// so the worker slices those from the live caches post-burst instead of copying at prime time.
1360pub struct SpecBoundaryCapture {
1361 pub snap: crate::cache::CacheSnapshot,
1362 /// Token boundary (== cache.pos at capture; == the worker's miss-LCP split).
1363 pub pos: usize,
1364 /// Full-vocab logits after the prefix prime — the entry's boundary logits.
1365 pub logits: Vec<f32>,
1366 /// Pre-output_norm trunk hidden of row `pos - 1` (lane/spec-on-cache-hit): the
1367 /// predecessor-pairing anchor a RESTORED spec session's first suffix-fill row reads
1368 /// (the `SpecSession::last_h` convention). Empty = unavailable (capture stays valid;
1369 /// the fill's zeros row-0 fallback covers it at a bounded acceptance cost).
1370 pub last_h: Vec<f32>,
1371}
1372
1373/// D2H one hidden row out of a `[T, n_embd]` prime hidden stack — the boundary anchor a
1374/// spec boundary capture carries for later restored-session fills. Failure is silent
1375/// (`turn_ckpt` convention): the capture publishes without an anchor.
1376fn capture_boundary_hidden(
1377 e: &Engine,
1378 h_rows: &CudaSlice<f32>,
1379 pos: usize,
1380 n_embd: usize,
1381) -> Vec<f32> {
1382 if pos == 0 || h_rows.len() < pos * n_embd {
1383 return Vec::new();
1384 }
1385 let Ok(mut row) = e.uninit(n_embd) else {
1386 return Vec::new();
1387 };
1388 if e.copy_view_into(
1389 &mut row,
1390 0,
1391 &h_rows.slice((pos - 1) * n_embd..pos * n_embd),
1392 n_embd,
1393 )
1394 .is_err()
1395 {
1396 return Vec::new();
1397 }
1398 e.dtoh(&row).unwrap_or_default()
1399}
1400
1401/// ROLLBACK DOOR for sampled BOUNDARY tokens (lane/sampled-spec-quality, 2026-08-19).
1402/// Default ON: the token a burst emits at its own boundary is drawn from the request's
1403/// sampler. `MEMRA_SPEC_SAMPLED_BOUNDARY=0` restores the pre-lane posture (an ARGMAX at
1404/// every boundary) without touching greedy, which is byte-unaffected either way.
1405pub fn spec_sampled_boundary_on() -> bool {
1406 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1407 *ON.get_or_init(|| std::env::var("MEMRA_SPEC_SAMPLED_BOUNDARY").as_deref() != Ok("0"))
1408}
1409
1410/// ROLLBACK DOOR for SESSION-SPANNING penalty history (lane/sampled-spec-quality).
1411/// Default ON: `pen_hist` is seeded from the session's committed tail, so repetition /
1412/// frequency / presence penalties see the whole stream. `MEMRA_SPEC_PEN_SESSION=0`
1413/// restores the pre-lane posture (each burst restarts the window from its own prompt
1414/// slice, i.e. from NOTHING on a continuation burst) — and with the door shut the worker
1415/// must keep refusing penalized sampled prefix-cache restores, because the restored
1416/// session's continuation burst is handed no prompt slice at all.
1417pub fn spec_pen_session_on() -> bool {
1418 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1419 *ON.get_or_init(|| std::env::var("MEMRA_SPEC_PEN_SESSION").as_deref() != Ok("0"))
1420}
1421
1422/// ROLLBACK DOOR for extended-entry publication from a RESTORED session
1423/// (lane/sampled-spec-quality, Item 3). Default ON: a converted prefix-cache hit that fed a
1424/// suffix captures its own prompt-end boundary so the NEXT turn can hit a longer prefix.
1425/// `MEMRA_SPEC_RESTORE_REPUBLISH=0` restores the pre-lane posture (a namespace learns exactly
1426/// one boundary and never advances it). Whole-entry semantics only — the boundary is the
1427/// restored session's own prompt end, so `entry_pos != fed_len` still refuses on the way in.
1428pub fn spec_restore_republish_on() -> bool {
1429 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1430 *ON.get_or_init(|| std::env::var("MEMRA_SPEC_RESTORE_REPUBLISH").as_deref() != Ok("0"))
1431}
1432
1433/// Diagnostics: name every boundary token on stderr (`MEMRA_SPEC_BOUNDARY_TRACE=1`), with
1434/// the argmax the pre-lane code would have emitted from the same row. This is how the
1435/// lane MEASURES the boundary rate and the deviation rate instead of estimating them.
1436fn spec_boundary_trace() -> bool {
1437 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1438 *ON.get_or_init(|| std::env::var("MEMRA_SPEC_BOUNDARY_TRACE").as_deref() == Ok("1"))
1439}
1440
1441/// llama-parity floor for the penalty window when the request does not ask for a bigger
1442/// one (`repeat_last_n` default). The serve API arms `penalty_last_n = PEN_WINDOW_MAX` for any
1443/// non-identity penalty, so this floor only matters to explicit small windows and to the
1444/// CLI env path.
1445const PEN_WINDOW_FLOOR: usize = 64;
1446
1447/// CEILING on the penalty window, and it is a COST bound, not a semantic preference.
1448/// `penalize_logits_f32` (cu/spec_sample.cu) dedups on device by having thread `i` scan
1449/// `hist[0..i]`, so a pass is O(n_hist²) and it runs ~3x per verify round (the q rows, the
1450/// p column, the bonus column). The serve API uses this same bound for every non-identity
1451/// penalty so host/plain, sparse-device, and speculative sampling cannot change logits on
1452/// admission demotion. An uncapped 128k-token history would put ~1.7e10
1453/// comparisons per pass, tens of ms per round, i.e. penalties would silently destroy decode
1454/// throughput on exactly the long-context requests that most want them. 8192 keeps a pass
1455/// at ~7e7 comparisons (tens of microseconds) while still being **128x wider than the
1456/// pre-lane effective window** (64 prompt-tail tokens + whatever the current burst had
1457/// generated). A request that genuinely needs a window beyond this wants host-side dedup +
1458/// counts through a new kernel signature — a follow-up lane, named here rather than hidden.
1459/// `pub` since lane/dspark-penalized-sampled-20260821: the dspark route's accept walk and
1460/// the dspark_sample_gate binary trim their uploads with the SAME cap — a second constant
1461/// is a second thing to drift.
1462pub const PEN_WINDOW_MAX: usize = 8192;
1463
1464/// Seed a penalty window over the SESSION, not the burst (lane/sampled-spec-quality,
1465/// Item 2). The window is the last `max(penalty_last_n, 64)` tokens of
1466/// `session_committed ++ burst_prompt` — for a cold turn-1 burst (`session_committed`
1467/// empty, default `penalty_last_n`) that is byte-identically the pre-lane
1468/// `prompt.iter().rev().take(64).rev()`; for a continuation burst it is the stream the
1469/// client actually asked us to penalize, where the pre-lane code had NOTHING.
1470/// `pub` since lane/dspark-penalized-sampled-20260821: the dspark route seeds its session
1471/// window through the SAME function (one definition of "the window" across both spec
1472/// routes and the gate binary's trunk-only reference arm).
1473pub fn pen_window_seed(
1474 session_committed: &[u32],
1475 burst_prompt: &[u32],
1476 penalty_last_n: usize,
1477) -> Vec<u32> {
1478 let win = penalty_last_n.clamp(PEN_WINDOW_FLOOR, PEN_WINDOW_MAX);
1479 let take_prompt = burst_prompt.len().min(win);
1480 let take_sess = (win - take_prompt).min(session_committed.len());
1481 let mut hist = Vec::with_capacity(take_sess + take_prompt);
1482 hist.extend_from_slice(&session_committed[session_committed.len() - take_sess..]);
1483 hist.extend_from_slice(&burst_prompt[burst_prompt.len() - take_prompt..]);
1484 hist
1485}
1486
1487/// Draw a BOUNDARY token from the target distribution the request asked for
1488/// (lane/sampled-spec-quality, Item 1) — the fix for "sampled spec emits an ARGMAX token at
1489/// every burst boundary".
1490///
1491/// WHY THIS EXISTS. A spec burst's first emitted token is not produced by the accept walk:
1492/// it comes off a logits row that already exists (the prime's last row on a cold burst; the
1493/// row after the last committed token on a continuation burst; the prefix-cache entry's
1494/// boundary row on a restored one). Pre-lane that token was `argmax` in BOTH sampling
1495/// regimes, so a sampled stream took a greedy token once per burst — measured, not
1496/// estimated, in research/spec-cache-20260818/SAMPLED-QUALITY.md. At temperature > 0 the
1497/// customer asked for a sampled token, so this draws one.
1498///
1499/// THE PROGRAM IS THE FULL-ACCEPT BONUS'S PROGRAM, deliberately: penalize the row (over the
1500/// session's window), take this row's OWN filter stats (the sampfix-20260805 law — stats
1501/// from a neighbour row mis-scale every `e0` and can wipe the row to token 0), gumbel-perturb
1502/// with the session's Philox stream at `*sctr`, argmax the perturbed row. Reusing the bonus's
1503/// composition means `sample_check`'s distributional oracle covers this draw too, and the
1504/// boundary token is drawn from the same filtered/penalized `p` the accept walk targets.
1505///
1506/// THE STREAM IS THE SESSION'S, NOT A FRESH ONE. `sctr` is the caller's live counter and is
1507/// advanced by exactly one, so a boundary draw consumes the next value in the same Philox
1508/// stream the accept walk uses — never a second, independently seeded stream (which would be
1509/// a new distributional bug: two streams from one seed correlate wherever their counters
1510/// collide). That also makes a restored session's boundary draw at `sctr == 0` bit-identical
1511/// to the cold session's own first draw from the same logits row, which is what preserves the
1512/// sampled-hit lane's per-seed hit==cold byte identity.
1513#[allow(clippy::too_many_arguments)]
1514pub fn sample_boundary_token_dev(
1515 e: &Engine,
1516 logits: &CudaSlice<f32>,
1517 n_vocab: usize,
1518 sp: &SpecSampling,
1519 pen_hist: &[u32],
1520 sctr: &mut u32,
1521 site: &str,
1522) -> Result<u32, Box<dyn std::error::Error>> {
1523 debug_assert!(
1524 sp.temp > 0.0,
1525 "boundary sampling is the sampled regime only"
1526 );
1527 // Own copy: penalize_logits mutates in place and the caller's row is live state
1528 // (prime_logits back the constrained recompute; last_col_logits backs round 0's accept).
1529 let mut col = e.zeros(n_vocab)?;
1530 e.copy_into(&mut col, 0, logits, n_vocab)?;
1531 let pen_on = sp.penalty_last_n > 0
1532 && (sp.penalty_repeat != 1.0 || sp.penalty_freq != 0.0 || sp.penalty_present != 0.0);
1533 if pen_on && !pen_hist.is_empty() {
1534 // window trim mirrors the round loop's own upload (`pen_hist[w0..]`), cap included.
1535 let w0 = pen_hist
1536 .len()
1537 .saturating_sub(sp.penalty_last_n.min(PEN_WINDOW_MAX));
1538 let hist = &pen_hist[w0..];
1539 let hd = e.htod_u32_v(hist)?;
1540 e.penalize_logits(
1541 &mut col,
1542 &hd,
1543 hist.len(),
1544 sp.penalty_repeat,
1545 sp.penalty_freq,
1546 sp.penalty_present,
1547 n_vocab,
1548 )?;
1549 }
1550 let rows0 = e.htod_i32(&[0])?;
1551 let (mut th_d, mut z_d, mut mx_d) = (e.zeros(1)?, e.zeros(1)?, e.zeros(1)?);
1552 e.filter_stats(
1553 &col, n_vocab, &rows0, &mut th_d, &mut z_d, &mut mx_d, n_vocab, 1, sp.temp, sp.top_k,
1554 sp.top_p, sp.min_p,
1555 )?;
1556 let (th, mx) = (e.dtoh(&th_d)?[0], e.dtoh(&mx_d)?[0]);
1557 let mut perturb = e.zeros(n_vocab)?;
1558 e.gumbel_perturb_filtered(&col, &mut perturb, n_vocab, sp.seed, *sctr, sp.temp, mx, th)?;
1559 *sctr = sctr.wrapping_add(1);
1560 let td = e.argmax_token_device(&perturb, n_vocab)?;
1561 let tok = guard_vocab_token(
1562 e.dtoh_u32_one(&td)?,
1563 n_vocab,
1564 &format!("sampled boundary token (site={site})"),
1565 )?;
1566 if spec_boundary_trace() {
1567 // the pre-lane token, from the SAME row, so the deviation rate is measurable.
1568 let raw = e.argmax_token_device(logits, n_vocab)?;
1569 let greedy = e.dtoh_u32_one(&raw)?;
1570 eprintln!(
1571 "[spec-boundary] site={site} sampled={tok} argmax={greedy} \
1572 deviates={} temp={} sctr={}",
1573 (tok != greedy) as u8,
1574 sp.temp,
1575 sctr.wrapping_sub(1),
1576 );
1577 }
1578 Ok(tok)
1579}
1580
1581/// Host-row twin of [`sample_boundary_token_dev`] (the prime / feed / entry rows arrive as
1582/// host `Vec<f32>`).
1583#[allow(clippy::too_many_arguments)]
1584pub fn sample_boundary_token(
1585 e: &Engine,
1586 logits: &[f32],
1587 sp: &SpecSampling,
1588 pen_hist: &[u32],
1589 sctr: &mut u32,
1590 site: &str,
1591) -> Result<u32, Box<dyn std::error::Error>> {
1592 let n_vocab = logits.len();
1593 let d = e.htod(logits)?;
1594 sample_boundary_token_dev(e, &d, n_vocab, sp, pen_hist, sctr, site)
1595}
1596
1597struct SpecPipeTraceClock {
1598 pair: usize,
1599 started: std::time::Instant,
1600}
1601
1602#[derive(Clone)]
1603struct SpecPipeTraceCtx {
1604 clock: std::sync::Arc<SpecPipeTraceClock>,
1605 round: usize,
1606 lane: usize,
1607}
1608
1609struct SpecPipeTraceMarker {
1610 trace: SpecPipeTraceCtx,
1611 phase: &'static str,
1612 edge: &'static str,
1613 slot: Option<usize>,
1614}
1615
1616unsafe extern "C" fn spec_pipe_trace_marker(raw: *mut std::ffi::c_void) {
1617 let marker = unsafe { Box::from_raw(raw.cast::<SpecPipeTraceMarker>()) };
1618 let lane = if marker.trace.lane == 0 { "A" } else { "B" };
1619 let slot = marker
1620 .slot
1621 .map(|v| v.to_string())
1622 .unwrap_or_else(|| "-".into());
1623 let t_ms = marker.trace.clock.started.elapsed().as_secs_f64() * 1e3;
1624 use std::io::Write as _;
1625 let stderr = std::io::stderr();
1626 let mut stderr = stderr.lock();
1627 let _ = writeln!(
1628 stderr,
1629 "[spec-pipe-timeline] pair={} round={} lane={lane} phase={} edge={} \
1630 slot={slot} t_ms={t_ms:.3}",
1631 marker.trace.clock.pair, marker.trace.round, marker.phase, marker.edge,
1632 );
1633}
1634
1635fn enqueue_spec_pipe_trace_marker(
1636 stream: &cudarc::driver::CudaStream,
1637 trace: Option<&SpecPipeTraceCtx>,
1638 phase: &'static str,
1639 edge: &'static str,
1640 slot: Option<usize>,
1641) -> Result<(), Box<dyn std::error::Error>> {
1642 let Some(trace) = trace else {
1643 return Ok(());
1644 };
1645 let marker = Box::new(SpecPipeTraceMarker {
1646 trace: trace.clone(),
1647 phase,
1648 edge,
1649 slot,
1650 });
1651 let raw = Box::into_raw(marker);
1652 let result = unsafe {
1653 cudarc::driver::result::stream::launch_host_function(
1654 stream.cu_stream(),
1655 spec_pipe_trace_marker,
1656 raw.cast(),
1657 )
1658 };
1659 if let Err(err) = result {
1660 unsafe {
1661 drop(Box::from_raw(raw));
1662 }
1663 return Err(err.into());
1664 }
1665 Ok(())
1666}
1667
1668#[derive(Default)]
1669struct SpecPipeProgress {
1670 setup_done: [bool; 2],
1671 draft_done: [usize; 2],
1672 stage0_done: [usize; 2],
1673 verify_done: [usize; 2],
1674 accept_done: [usize; 2],
1675 finished: [bool; 2],
1676 aborted: bool,
1677}
1678
1679/// Host-side issue coordinator for the reduced two-session speculative pipeline. Each session
1680/// keeps its existing call stack and round locals; this object only orders phase entry. The
1681/// primary mutex spans whole draft/accept/tail issue regions so Engine's single-stream scratch
1682/// cannot be interleaved by the two host threads.
1683struct SpecPipeSync {
1684 progress: std::sync::Mutex<SpecPipeProgress>,
1685 changed: std::sync::Condvar,
1686 primary: std::sync::Mutex<()>,
1687 trace: Option<std::sync::Arc<SpecPipeTraceClock>>,
1688}
1689
1690impl SpecPipeSync {
1691 fn new() -> Self {
1692 static TRACE_PAIR: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
1693 let trace = (std::env::var("MEMRA_SPEC_PIPE_TRACE").as_deref() == Ok("1")).then(|| {
1694 std::sync::Arc::new(SpecPipeTraceClock {
1695 pair: TRACE_PAIR.fetch_add(1, std::sync::atomic::Ordering::Relaxed) + 1,
1696 started: std::time::Instant::now(),
1697 })
1698 });
1699 Self {
1700 progress: std::sync::Mutex::new(SpecPipeProgress::default()),
1701 changed: std::sync::Condvar::new(),
1702 primary: std::sync::Mutex::new(()),
1703 trace,
1704 }
1705 }
1706}
1707
1708#[derive(Clone)]
1709struct SpecPipeLane {
1710 sync: std::sync::Arc<SpecPipeSync>,
1711 lane: usize,
1712}
1713
1714impl SpecPipeLane {
1715 fn peer(&self) -> usize {
1716 1 - self.lane
1717 }
1718
1719 fn aborted() -> Box<dyn std::error::Error> {
1720 "paired speculative peer aborted".into()
1721 }
1722
1723 fn trace(&self, round: usize) -> Option<SpecPipeTraceCtx> {
1724 self.sync.trace.as_ref().map(|clock| SpecPipeTraceCtx {
1725 clock: clock.clone(),
1726 round,
1727 lane: self.lane,
1728 })
1729 }
1730
1731 fn setup_begin(&self) -> Result<(), Box<dyn std::error::Error>> {
1732 let mut p = self.sync.progress.lock().unwrap();
1733 while !p.aborted && self.lane == 1 && !p.setup_done[0] && !p.finished[0] {
1734 p = self.sync.changed.wait(p).unwrap();
1735 }
1736 if p.aborted {
1737 Err(Self::aborted())
1738 } else {
1739 Ok(())
1740 }
1741 }
1742
1743 fn setup_end(&self) {
1744 let mut p = self.sync.progress.lock().unwrap();
1745 p.setup_done[self.lane] = true;
1746 self.sync.changed.notify_all();
1747 }
1748
1749 fn draft_begin(
1750 &self,
1751 round: usize,
1752 ) -> Result<std::sync::MutexGuard<'_, ()>, Box<dyn std::error::Error>> {
1753 let peer = self.peer();
1754 let mut p = self.sync.progress.lock().unwrap();
1755 loop {
1756 if p.aborted {
1757 return Err(Self::aborted());
1758 }
1759 let setup_ready =
1760 (p.setup_done[0] || p.finished[0]) && (p.setup_done[1] || p.finished[1]);
1761 let prior_ready = p.accept_done[self.lane] >= round
1762 && (p.accept_done[peer] >= round || p.finished[peer]);
1763 let turn_ready = if self.lane == 0 {
1764 true
1765 } else {
1766 p.draft_done[0] > round || p.finished[0]
1767 };
1768 if setup_ready && prior_ready && turn_ready {
1769 break;
1770 }
1771 p = self.sync.changed.wait(p).unwrap();
1772 }
1773 drop(p);
1774 Ok(self.sync.primary.lock().unwrap())
1775 }
1776
1777 fn draft_end(&self, round: usize) {
1778 let mut p = self.sync.progress.lock().unwrap();
1779 p.draft_done[self.lane] = round + 1;
1780 self.sync.changed.notify_all();
1781 }
1782
1783 /// Admit stage 0 and return whether this lane owns the interval's one reverse fence.
1784 /// Lane B releases as soon as lane A has issued its boundary TX, not after A's full body.
1785 fn stage0_begin(&self, round: usize) -> Result<bool, Box<dyn std::error::Error>> {
1786 let peer = self.peer();
1787 let mut p = self.sync.progress.lock().unwrap();
1788 loop {
1789 if p.aborted {
1790 return Err(Self::aborted());
1791 }
1792 let ready = if self.lane == 0 {
1793 p.draft_done[0] > round && (p.draft_done[1] > round || p.finished[1])
1794 } else {
1795 p.draft_done[1] > round && (p.stage0_done[0] > round || p.finished[0])
1796 };
1797 if ready {
1798 return Ok(self.lane == 0 || p.finished[peer]);
1799 }
1800 p = self.sync.changed.wait(p).unwrap();
1801 }
1802 }
1803
1804 fn stage0_end(&self, round: usize) {
1805 let mut p = self.sync.progress.lock().unwrap();
1806 p.stage0_done[self.lane] = round + 1;
1807 self.sync.changed.notify_all();
1808 }
1809
1810 /// Stage 1 is single-owner per engine. A proceeds immediately after its own ticket; B waits
1811 /// for A's full stage1/head issue so only A.S1 and B.S0 can overlap.
1812 fn stage1_begin(&self, round: usize) -> Result<(), Box<dyn std::error::Error>> {
1813 let mut p = self.sync.progress.lock().unwrap();
1814 while !p.aborted
1815 && !(p.stage0_done[self.lane] > round
1816 && (self.lane == 0 || p.verify_done[0] > round || p.finished[0]))
1817 {
1818 p = self.sync.changed.wait(p).unwrap();
1819 }
1820 if p.aborted {
1821 Err(Self::aborted())
1822 } else {
1823 Ok(())
1824 }
1825 }
1826
1827 fn verify_end(&self, round: usize) {
1828 let mut p = self.sync.progress.lock().unwrap();
1829 p.verify_done[self.lane] = round + 1;
1830 self.sync.changed.notify_all();
1831 }
1832
1833 fn accept_begin(
1834 &self,
1835 round: usize,
1836 ) -> Result<std::sync::MutexGuard<'_, ()>, Box<dyn std::error::Error>> {
1837 let mut p = self.sync.progress.lock().unwrap();
1838 loop {
1839 if p.aborted {
1840 return Err(Self::aborted());
1841 }
1842 let ready = if self.lane == 0 {
1843 p.verify_done[0] > round && (p.verify_done[1] > round || p.finished[1])
1844 } else {
1845 p.verify_done[1] > round && (p.accept_done[0] > round || p.finished[0])
1846 };
1847 if ready {
1848 break;
1849 }
1850 p = self.sync.changed.wait(p).unwrap();
1851 }
1852 drop(p);
1853 Ok(self.sync.primary.lock().unwrap())
1854 }
1855
1856 fn accept_end(&self, round: usize) {
1857 let mut p = self.sync.progress.lock().unwrap();
1858 p.accept_done[self.lane] = round + 1;
1859 self.sync.changed.notify_all();
1860 }
1861
1862 fn primary(&self) -> std::sync::MutexGuard<'_, ()> {
1863 self.sync.primary.lock().unwrap()
1864 }
1865
1866 fn finish(&self, failed: bool) {
1867 let mut p = self.sync.progress.lock().unwrap();
1868 p.finished[self.lane] = true;
1869 p.aborted |= failed;
1870 self.sync.changed.notify_all();
1871 }
1872}
1873
1874struct SpecPipeFinish<'a> {
1875 lane: &'a SpecPipeLane,
1876 closed: bool,
1877}
1878
1879impl<'a> SpecPipeFinish<'a> {
1880 fn new(lane: &'a SpecPipeLane) -> Self {
1881 Self {
1882 lane,
1883 closed: false,
1884 }
1885 }
1886
1887 fn close(&mut self, failed: bool) {
1888 self.lane.finish(failed);
1889 self.closed = true;
1890 }
1891}
1892
1893impl Drop for SpecPipeFinish<'_> {
1894 fn drop(&mut self) {
1895 if !self.closed {
1896 self.lane.finish(true);
1897 }
1898 }
1899}
1900
1901/// Scoped transfer of one exclusively-borrowed session to the second host issue thread.
1902/// `CudaGraph` is not marked Send by cudarc because its raw driver handles carry no automatic
1903/// trait. CUDA driver graph handles are context-scoped rather than OS-thread-affine; the caller
1904/// binds that context before touching the session, joins before returning, and never aliases the
1905/// pointer. Keep this exception local to the experimental pair call instead of marking the public
1906/// session type Send.
1907struct SpecPipeSessionPtr(*mut SpecSession);
1908
1909unsafe impl Send for SpecPipeSessionPtr {}
1910
1911impl SpecPipeSessionPtr {
1912 unsafe fn get_mut(&mut self) -> &mut SpecSession {
1913 unsafe { &mut *self.0 }
1914 }
1915}
1916
1917/// Per-session persistent draft-graph context: the captured CUDA graph(s) plus the device
1918/// buffers whose POINTERS the capture bakes. Reuse legality: the greedy capture bakes only
1919/// session-stable pointers (the session's own MtpScratch KV — allocated once, never realloc'd;
1920/// the model's resident embedding; the process-wide OnceLock p_min) and the g_* buffers held
1921/// HERE — so one capture serves the session's whole lifetime. The sampled capture additionally
1922/// bakes (seed, temp) as capture-time constants and needs k q-slots — keyed by `s_key`, dropped
1923/// and recaptured when a pool-resumed request changes them. `*_failed` memoizes a failed capture
1924/// so the eager fallback doesn't pay a doomed capture attempt every burst.
1925/// Capture identity of the parked SAMPLED draft graph (`DraftGraphCtx::graph_s`).
1926///
1927/// EXACTNESS, not perf (lane/graph-s-key-exactness-20260819; receipts
1928/// `research/spec-cache-20260818/GRAPH-S-KEY.md`). Two classes of field live here, both
1929/// load-bearing:
1930///
1931/// - **Baked constants.** `seed` and `temp` are capture-time constants INSIDE the graph and `k`
1932/// sizes the q slots its replays write. A resumed request changing any of them must recapture.
1933/// This is all the key used to carry.
1934/// - **Regime fields.** `top_k`/`top_p`/`min_p`/`pen_on` are not baked, but they decide whether
1935/// the captured graph is a legal draft chain AT ALL. The in-graph draw is one gumbel-max over
1936/// the RAW softmax (`gumbel_perturb_ctr`, unfiltered by construction), while the verify builds
1937/// the accept test's `q` from `filter_stats(q_slots, top_k, top_p, min_p)`. If those disagree
1938/// the accept test evaluates a distribution the draft was never sampled from: a draft token
1939/// below the filter threshold gathers `q = 0` (`softmax_gather_filtered_f32`,
1940/// `cu/spec_sample.cu`) and `u * 0 < p` accepts it UNCONDITIONALLY.
1941///
1942/// Omitting the regime fields was reachable — not through the prefix-cache spec restore (that
1943/// path is greedy-only, `memra-server` `spec_restore_convertible`), but through WHOLE-SESSION
1944/// spec reuse: a parked `SpecSession` carries this `DraftGraphCtx`, and the pool-resume probe
1945/// applies no sampler predicate at all. Turn 1 pure-temp parks a graph; turn 2 of the same
1946/// conversation, same explicit seed and temperature, adds `top_p`/`top_k` and inherits it.
1947#[derive(Clone, Copy, PartialEq, Eq, Debug)]
1948pub(crate) struct SampledGraphKey {
1949 seed: u64,
1950 temp_bits: u32,
1951 k: usize,
1952 top_k: i32,
1953 top_p_bits: u32,
1954 min_p_bits: u32,
1955 pen_on: bool,
1956}
1957
1958impl SampledGraphKey {
1959 pub(crate) fn new(
1960 seed: u64,
1961 temp: f32,
1962 k: usize,
1963 top_k: i32,
1964 top_p: f32,
1965 min_p: f32,
1966 pen_on: bool,
1967 ) -> Self {
1968 SampledGraphKey {
1969 seed,
1970 temp_bits: temp.to_bits(),
1971 k,
1972 top_k,
1973 top_p_bits: top_p.to_bits(),
1974 min_p_bits: min_p.to_bits(),
1975 pen_on,
1976 }
1977 }
1978
1979 /// The one regime the PURE-TEMP in-graph sampled chain may stand in for the eager one:
1980 /// nothing but temperature shapes `q`. Computed FROM THE KEY so the capture guard, the
1981 /// launch guard and the key can never drift apart (they were three separate expressions
1982 /// before this lane, and the launch site simply forgot to ask).
1983 pub(crate) fn pure_temp(&self) -> bool {
1984 self.top_k == 0
1985 && f32::from_bits(self.top_p_bits) >= 1.0
1986 && f32::from_bits(self.min_p_bits) <= 0.0
1987 && !self.pen_on
1988 }
1989
1990 /// Truncation filters active — the capture body needs the IN-GRAPH filter nodes
1991 /// (`filter_stats` + `gumbel_perturb_filtered_ctr`) so the draft draws from the same
1992 /// filtered distribution the accept test reconstructs. Meaningful only when
1993 /// `graph_capturable`; penalties never reach a capture body.
1994 pub(crate) fn filtered(&self) -> bool {
1995 !self.pure_temp()
1996 }
1997
1998 /// May the sampled draft graph be CAPTURED (and a parked one LAUNCHED) for this regime?
1999 /// Pure-temp always; filtered regimes when the filtered-capture door is on
2000 /// (lane/step37-draft-graph-serving-20260830); penalties never — the per-round history
2001 /// cannot be baked into a graph, and composing a raw-softmax (or stale-history) draw
2002 /// with a penalized accept test is the unconditional-accept exactness bug. Computed FROM
2003 /// THE KEY for the same no-drift reason as `pure_temp`.
2004 pub(crate) fn graph_capturable(&self) -> bool {
2005 !self.pen_on && (self.pure_temp() || spec_graph_filtered_on())
2006 }
2007}
2008
2009/// Per-head captured graphs for the MULTI-HEAD MTP draft chain (step-modulo prefix-replay,
2010/// lane/step37-draft-graph-serving-20260830). The chain POLICY — which head serves step j,
2011/// how long the replayed prefix is, which stored seed feeds row r — stays HOST-SIDE in the
2012/// launch loop, exactly `mtp_chain_forward_dev`'s order; the graphs capture ONE head-row
2013/// forward each, on the head's OWN scratch plane:
2014/// - `interior[i]`: head i, `with_head=false` — KV append + carrier only. Interior rows'
2015/// logits are dead in the eager chain too (`mtp_chain_forward_dev` keeps only the last
2016/// row), so skipping the head matmul changes no consumed byte and removes the eager
2017/// chain's per-replay-row full-vocab matmul.
2018/// - `last[i]`: head i, `with_head=true` + the mode's tail (greedy argmax, or the sampled
2019/// gumbel draw — filtered in-graph when the request carries filters).
2020/// One `DraftChainGraphs` per MODE (greedy vs sampled), owning its keeper: dropping the
2021/// sampled chain on an s_key change never invalidates the greedy one.
2022struct DraftChainGraphs {
2023 interior: Vec<cudarc::driver::CudaGraph>,
2024 last: Vec<cudarc::driver::CudaGraph>,
2025 keeper: Vec<Box<dyn std::any::Any + Send>>,
2026}
2027
2028/// Sampled-tail capture pack for `mtp_head_forward_cap`: the persistent buffers and baked
2029/// constants of the in-graph categorical draw. `filt: None` = the PURE-TEMP body (gumbel
2030/// over the raw softmax), byte-identical to the pre-lane capture; `Some` adds the in-graph
2031/// truncation filter (`filter_stats` + `gumbel_perturb_filtered_ctr`) so the draft draws
2032/// from the same filtered distribution the accept test reconstructs
2033/// (lane/step37-draft-graph-serving-20260830).
2034struct SampledCapArgs<'a> {
2035 ctr: &'a mut CudaSlice<u32>,
2036 perturb: &'a mut CudaSlice<f32>,
2037 q_out: &'a mut CudaSlice<f32>,
2038 seed: u64,
2039 temp: f32,
2040 filt: Option<SampledCapFilter<'a>>,
2041}
2042
2043/// In-graph truncation-filter nodes: the stat slots `filter_stats` fills and the perturb
2044/// reads, plus the filter constants baked into the capture (they live in `s_key`, so a
2045/// request whose filters differ drops the parked graph before this ever goes stale).
2046struct SampledCapFilter<'a> {
2047 rows0: &'a CudaSlice<i32>,
2048 th: &'a mut CudaSlice<f32>,
2049 z: &'a mut CudaSlice<f32>,
2050 mx: &'a mut CudaSlice<f32>,
2051 top_k: i32,
2052 top_p: f32,
2053 min_p: f32,
2054}
2055
2056pub(crate) struct DraftGraphCtx {
2057 g_tok: CudaSlice<u32>,
2058 g_pos: CudaSlice<i32>,
2059 g_seed: CudaSlice<f32>,
2060 g_p: CudaSlice<f32>,
2061 g_ctr: CudaSlice<u32>,
2062 g_q: CudaSlice<f32>,
2063 g_perturb: CudaSlice<f32>,
2064 /// IN-GRAPH filter-stat slots (filtered sampled capture): `filter_stats` writes
2065 /// (th, z, mx) here inside the graph; `gumbel_perturb_filtered_ctr` reads (mx, th) from
2066 /// the same slots. Persistent so the baked pointers survive replays. `g_rows0` is the
2067 /// constant row-index-0 the single-row `filter_stats` launch reads (a captured memcpy
2068 /// source must not be a host temporary).
2069 g_rows0: CudaSlice<i32>,
2070 g_th: CudaSlice<f32>,
2071 g_z: CudaSlice<f32>,
2072 g_mx: CudaSlice<f32>,
2073 q_slots: Vec<CudaSlice<f32>>,
2074 /// DRAFT-SIDE GRAMMAR MASK (lane/draft-mask): packed allowed-set words over the DRAFT
2075 /// head's vocab, at a STABLE address so the captured draft graph's mask node reads the
2076 /// per-position contents the host re-uploads before each replay (the graph-promote
2077 /// pattern from decode.rs). Empty unless the session drafts under a grammar.
2078 g_dmask: CudaSlice<u32>,
2079 /// was `graph` captured WITH the mask node? A parked graph of the wrong shape is dropped.
2080 /// Covers the multi-head `chain` too (single-head and chain are mutually exclusive for a
2081 /// given model, so one flag serves whichever is active).
2082 graph_masked: bool,
2083 graph: Option<cudarc::driver::CudaGraph>,
2084 graph_s: Option<cudarc::driver::CudaGraph>,
2085 /// Multi-head chain graphs (see [`DraftChainGraphs`]): greedy and sampled chains, the
2086 /// chain twins of `graph` / `graph_s`. `chain_s`'s capture identity is `s_key` (shared
2087 /// with `graph_s` — a session is either single-head or chain, never both), and it obeys
2088 /// the same drop rules (key mismatch, penalty regime, mask-shape change).
2089 chain: Option<DraftChainGraphs>,
2090 chain_s: Option<DraftChainGraphs>,
2091 /// Failed-capture memoization for both graphs — LOUD on flip, cleared on pool resume
2092 /// (audit Q2, the TRT #16072 silent-permanent-coverage-loss class).
2093 failed: DraftGraphFallback,
2094 /// Capture identity of `graph_s` — see [`SampledGraphKey`]. `None` iff no sampled graph is
2095 /// parked; a request whose key differs drops the parked graph (and its q slots/keeper).
2096 s_key: Option<SampledGraphKey>,
2097 /// CAPTURE-RETAIN keepers (#68 root cause, 2026-08-04): the warmup-run transients whose
2098 /// pool addresses the captured graph(s) bake. Without these, the transients return to the
2099 /// pool at capture-body exit and later work (burst-boundary prime/fill/commit passes, or a
2100 /// co-served session in the worker) reuses those addresses — the persisted graph's replay
2101 /// then reads/writes live unrelated buffers (exactness corruption, first seen as the ST
2102 /// serve-spec 4B graph-arm corruption; one-shot CLI calls never re-shuffled the pool, which
2103 /// is why run-spec K=1..8 passed on the same checkpoint). Same fix class as
2104 /// capture_graph_retained's gemma/decode.rs sites — hold as long as the graph replays.
2105 keeper: Vec<Box<dyn std::any::Any + Send>>,
2106 keeper_s: Vec<Box<dyn std::any::Any + Send>>,
2107}
2108
2109/// Failed-capture memoization for the two draft graphs (audit Q2, 2026-08-05 — the
2110/// TRT #16072 trap class: pressure-triggered, silent, long-lived coverage loss).
2111///
2112/// Three contracts:
2113/// - LOUD FLIP: `mark_*` returns the warn line exactly on the false→true transition
2114/// (returned, not printed, so the once-per-flip contract is unit-testable); the caller
2115/// `eprintln!`s it UNCONDITIONALLY — a dropped draft graph is never silent. Re-marking
2116/// an already-failed graph returns None (the per-burst memoization that keeps the eager
2117/// fallback from paying a doomed capture attempt every burst).
2118/// - RESET ON RESUME: `reset_on_resume` clears both flags — a parked session resumed by a
2119/// NEW request gets one fresh capture chance instead of carrying a transient-pressure
2120/// failure for the pool's whole lifetime. Returns the note line only when a flag was
2121/// actually set (quiet on the common clean-resume path).
2122/// - Shape-change clears (`clear_*`) stay silent, exactly as before: they precede a fresh
2123/// capture attempt whose own failure would re-flip loudly.
2124#[derive(Default)]
2125pub(crate) struct DraftGraphFallback {
2126 greedy: bool,
2127 sampled: bool,
2128}
2129impl DraftGraphFallback {
2130 fn mark_greedy(&mut self, reason: &str) -> Option<String> {
2131 if self.greedy {
2132 return None;
2133 }
2134 self.greedy = true;
2135 Some(format!(
2136 "[spec] WARN: draft-graph capture failed ({reason}); eager fallback until session resume"
2137 ))
2138 }
2139 fn mark_sampled(&mut self, reason: &str) -> Option<String> {
2140 if self.sampled {
2141 return None;
2142 }
2143 self.sampled = true;
2144 Some(format!(
2145 "[spec] WARN: sampled draft-graph capture failed ({reason}); eager fallback until session resume"
2146 ))
2147 }
2148 fn greedy_failed(&self) -> bool {
2149 self.greedy
2150 }
2151 fn sampled_failed(&self) -> bool {
2152 self.sampled
2153 }
2154 fn clear_greedy(&mut self) {
2155 self.greedy = false;
2156 }
2157 fn clear_sampled(&mut self) {
2158 self.sampled = false;
2159 }
2160 /// Pool-resume reset: both graphs get a fresh capture chance. Some(note) iff any flag
2161 /// was set (so clean resumes stay quiet).
2162 pub(crate) fn reset_on_resume(&mut self) -> Option<String> {
2163 if !self.greedy && !self.sampled {
2164 return None;
2165 }
2166 let which = match (self.greedy, self.sampled) {
2167 (true, true) => "greedy+sampled",
2168 (true, false) => "greedy",
2169 _ => "sampled",
2170 };
2171 self.greedy = false;
2172 self.sampled = false;
2173 Some(format!(
2174 "[spec] draft-graph fallback reset on session resume ({which}); recapture eligible"
2175 ))
2176 }
2177}
2178
2179impl DraftGraphCtx {
2180 fn new(e: &Engine, n_embd: usize, qlen: usize) -> Result<Self, Box<dyn std::error::Error>> {
2181 Ok(DraftGraphCtx {
2182 g_tok: e.alloc_u32_zeroed(1)?,
2183 g_pos: e.htod_i32(&[0])?,
2184 g_seed: e.zeros(n_embd)?,
2185 g_p: e.zeros(1)?,
2186 g_ctr: e.alloc_u32_zeroed(1)?,
2187 g_q: e.zeros(qlen)?,
2188 g_perturb: e.zeros(qlen)?,
2189 g_rows0: e.htod_i32(&[0])?,
2190 g_th: e.zeros(1)?,
2191 g_z: e.zeros(1)?,
2192 g_mx: e.zeros(1)?,
2193 q_slots: Vec::new(),
2194 g_dmask: e.alloc_u32_zeroed(1)?,
2195 graph_masked: false,
2196 graph: None,
2197 graph_s: None,
2198 chain: None,
2199 chain_s: None,
2200 failed: DraftGraphFallback::default(),
2201 s_key: None,
2202 keeper: Vec::new(),
2203 keeper_s: Vec::new(),
2204 })
2205 }
2206}
2207
2208pub(crate) struct MtpScratch {
2209 kv: KvLayer,
2210 /// Logical row capacity. On the graph/DC draft path it also doubles as fa_decode_dc's
2211 /// bucket_max: n_splits is sized from it ONCE, so the graph captured at round 0 stays valid
2212 /// for every later t_kv. Step35 refuses that path and may back this logical extent with the
2213 /// smaller host-indexed SWA ring instead.
2214 cap: usize,
2215 extra: Vec<MtpScratchPlane>,
2216}
2217
2218struct MtpScratchPlane {
2219 kv: KvLayer,
2220 cap: usize,
2221}
2222
2223fn mtp_scratch_layout(
2224 cfg: &memra_gguf::config::ModelConfig,
2225 geom: Option<&crate::hybrid::DraftGeom>,
2226) -> (usize, usize, usize, usize) {
2227 // Student draft heads carry fewer KV heads (head_dim unchanged) -> smaller scratch rows.
2228 let n_head_kv = geom.map(|g| g.n_head_kv).unwrap_or(cfg.n_head_kv as usize);
2229 let head_dim_k = cfg.head_dim_k as usize;
2230 let head_dim_v = cfg.head_dim_v as usize;
2231 assert!(
2232 head_dim_k % 32 == 0 && head_dim_v % 32 == 0,
2233 "KVQUANT requires head_dim%32==0 (MTP scratch)"
2234 );
2235 let kv_dim_k = head_dim_k * n_head_kv;
2236 let kv_dim_v = head_dim_v * n_head_kv;
2237 // The fp8-KV arm deliberately does not reach the draft scratch; keep the exact format
2238 // policy shared with `MtpScratch::new` so admission scales the same allocation.
2239 let (kbb, vbb) = crate::kv_blk_bytes();
2240 let k_tok_bytes = (kv_dim_k / 32) * kbb;
2241 let v_tok_bytes = (kv_dim_v / 32) * vbb;
2242 (kv_dim_k, kv_dim_v, k_tok_bytes, v_tok_bytes)
2243}
2244
2245fn mtp_chain_head_index(step: usize, head_count: usize) -> usize {
2246 assert!(head_count > 0, "MTP chain requires at least one head");
2247 step % head_count
2248}
2249
2250impl MtpScratch {
2251 fn alloc_plane(
2252 e: &Engine,
2253 cfg: &memra_gguf::config::ModelConfig,
2254 plan: &memra_gguf::model_plan::ModelPlan,
2255 cap: usize,
2256 geom: Option<&crate::hybrid::DraftGeom>,
2257 ) -> Result<MtpScratchPlane, Box<dyn std::error::Error>> {
2258 let (kv_dim_k, kv_dim_v, k_tok_bytes, v_tok_bytes) = mtp_scratch_layout(cfg, geom);
2259 let ring = if crate::cache::swa_ring_on()
2260 && crate::plan_backend::decode_batch_program(plan)
2261 == crate::plan_backend::DecodeBatchProgram::SlidingGatedMoe
2262 {
2263 let window = plan
2264 .layers
2265 .iter()
2266 .find_map(|layer| match layer.attention {
2267 memra_gguf::model_plan::AttentionPlan::SlidingWindow { window, .. } => {
2268 Some(window as usize)
2269 }
2270 _ => None,
2271 })
2272 .ok_or("sliding-gated-MoE draft scratch has no sliding-window layer")?;
2273 Some(crate::cache::KvRing::new(
2274 crate::cache::swa_ring_rows(window, cap),
2275 window,
2276 ))
2277 } else {
2278 None
2279 };
2280 let alloc_rows = ring.as_ref().map(crate::cache::KvRing::rows).unwrap_or(cap);
2281 // Ring-backed planes arm the device base mirror for the dcw draft arm (see
2282 // KvLayer::base_d): the captured chain derives its physical rows from
2283 // (len_d, base_d, window) with zero per-token node updates.
2284 let base_d = match ring.as_ref() {
2285 Some(_) => Some(e.htod_i32(&[0])?),
2286 None => None,
2287 };
2288 Ok(MtpScratchPlane {
2289 kv: KvLayer {
2290 k: e.alloc_u8(alloc_rows * k_tok_bytes)?,
2291 v: e.alloc_u8(alloc_rows * v_tok_bytes)?,
2292 kv_dim_k,
2293 kv_dim_v,
2294 k_tok_bytes,
2295 v_tok_bytes,
2296 len: 0,
2297 ring,
2298 len_d: e.htod_i32(&[0])?,
2299 base_d,
2300 },
2301 cap,
2302 })
2303 }
2304
2305 fn new(
2306 e: &Engine,
2307 cfg: &memra_gguf::config::ModelConfig,
2308 plan: &memra_gguf::model_plan::ModelPlan,
2309 cap: usize,
2310 geom: Option<&crate::hybrid::DraftGeom>,
2311 ) -> Result<Self, Box<dyn std::error::Error>> {
2312 // env-selected KV formats (default 34/24). The fp8-KV arm (MEMRA_KV_FP8) deliberately
2313 // does NOT reach the draft scratch: fp8 drafts drifted acceptance 69-88% -> 46%
2314 // (2026-07-12 A/B); the scratch is tiny, so it keeps baseline q8_0/q5_1 numerics
2315 // while the TRUNK cache carries the fp8 depth win. Scratch append/fa pass g=false.
2316 let primary = Self::alloc_plane(e, cfg, plan, cap, geom)?;
2317 Ok(MtpScratch {
2318 kv: primary.kv,
2319 cap: primary.cap,
2320 extra: Vec::new(),
2321 })
2322 }
2323
2324 fn push_plane(
2325 &mut self,
2326 e: &Engine,
2327 cfg: &memra_gguf::config::ModelConfig,
2328 plan: &memra_gguf::model_plan::ModelPlan,
2329 geom: Option<&crate::hybrid::DraftGeom>,
2330 ) -> Result<(), Box<dyn std::error::Error>> {
2331 self.extra
2332 .push(Self::alloc_plane(e, cfg, plan, self.cap, geom)?);
2333 Ok(())
2334 }
2335
2336 fn plane_count(&self) -> usize {
2337 1 + self.extra.len()
2338 }
2339
2340 fn plane(&self, index: usize) -> (&KvLayer, usize) {
2341 if index == 0 {
2342 (&self.kv, self.cap)
2343 } else {
2344 let plane = &self.extra[index - 1];
2345 (&plane.kv, plane.cap)
2346 }
2347 }
2348
2349 fn plane_mut(&mut self, index: usize) -> (&mut KvLayer, usize) {
2350 if index == 0 {
2351 (&mut self.kv, self.cap)
2352 } else {
2353 let plane = &mut self.extra[index - 1];
2354 (&mut plane.kv, plane.cap)
2355 }
2356 }
2357
2358 // #[track_caller]: set_len/set_plane_len have eight call sites (checkpoint restore, spec
2359 // rollback, session grow, seed replay ...) and the lap failure needs to say WHICH one, not
2360 // just that a rewind was refused.
2361 #[track_caller]
2362 fn set_plane_len(
2363 &mut self,
2364 e: &Engine,
2365 index: usize,
2366 n: usize,
2367 ) -> Result<(), Box<dyn std::error::Error>> {
2368 let caller = std::panic::Location::caller();
2369 let (kv, cap) = self.plane_mut(index);
2370 if let Some(ring) = kv.ring.as_ref() {
2371 if !ring.can_rewind_to(n) {
2372 // NAME THE NUMBERS (2026-08-28). This error is a step37 serving blocker on the
2373 // vendor-default shape and it fires from more than one call path with more than
2374 // one trigger: a long generation walks the checkpoint out of the ring, but a
2375 // ~4.5k-token prompt also fails within 5 s of prime, which accumulation cannot
2376 // explain. A bare message forced two rounds of guessing; the operands make each
2377 // trigger name itself.
2378 let raw = n.saturating_sub(ring.window().saturating_sub(1));
2379 return Err(format!(
2380 "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})",
2381 ring.window(),
2382 ring.base(),
2383 ring.rows(),
2384 raw & !31usize,
2385 )
2386 .into());
2387 }
2388 }
2389 kv.len = n;
2390 e.set_i32_one(&mut kv.len_d, n as i32)
2391 }
2392
2393 /// Set BOTH length counters: the host mirror AND the device len_d the captured append/fa read
2394 /// (a 4-byte in-place htod — the counter pointer is baked into the graph, never realloc'd).
2395 /// This is the ONLY truncation/rollback mechanism the persistent draft KV needs.
2396 #[track_caller]
2397 fn set_len(&mut self, e: &Engine, n: usize) -> Result<(), Box<dyn std::error::Error>> {
2398 let caller = std::panic::Location::caller();
2399 if !self.can_rewind_to(n) {
2400 // set_plane_len re-checks and reports the operands; call it so the failure carries
2401 // which plane refused and why, instead of this bare aggregate.
2402 for index in 0..self.plane_count() {
2403 self.set_plane_len(e, index, n)?;
2404 }
2405 return Err(format!(
2406 "SWA ring MTP checkpoint has been lapped; full re-prime required (aggregate rewind_to={n}, no single plane reported, called from {caller})"
2407 )
2408 .into());
2409 }
2410 for index in 0..self.plane_count() {
2411 self.set_plane_len(e, index, n)?;
2412 }
2413 Ok(())
2414 }
2415
2416 fn can_rewind_to(&self, n: usize) -> bool {
2417 (0..self.plane_count()).all(|index| {
2418 self.plane(index)
2419 .0
2420 .ring
2421 .as_ref()
2422 .is_none_or(|ring| ring.can_rewind_to(n))
2423 })
2424 }
2425
2426 /// Pre-arm ring headroom for `rows` upcoming DEVICE-COUNTER appends (the dcw draft arm):
2427 /// a captured chain cannot rebase mid-replay, so any rebase the coming appends could need
2428 /// happens HERE, host-side, before the capture warmups or the round's replays (the rebase
2429 /// arm of `prepare_kv_append` also refreshes the plane's `base_d` device mirror). No-op on
2430 /// flat planes and when the ring already has room; `len` is untouched either way.
2431 fn ensure_dcw_headroom(
2432 &mut self,
2433 e: &Engine,
2434 rows: usize,
2435 ) -> Result<(), Box<dyn std::error::Error>> {
2436 for index in 0..self.plane_count() {
2437 let (kv, _) = self.plane_mut(index);
2438 let Some(ring) = kv.ring.as_ref() else {
2439 continue;
2440 };
2441 let retain = memra_kv::swa_retain_from(kv.len, ring.window(), ring.base());
2442 e.prepare_kv_append(kv, retain, rows)?;
2443 }
2444 Ok(())
2445 }
2446}
2447
2448/// Retained verify intermediates for the REPLAY-FREE partial accept (2026-07-03, the profiled
2449/// #1 spec cost at long ctx: the partial-accept replay was a DUPLICATE trunk pass — ~0.54 extra
2450/// full weight reads per round — recomputing columns the verify had already produced
2451/// bit-identically). Holds, per linear layer, everything needed to rebuild its recurrent state
2452/// to "after the first j verify columns" WITHOUT re-running the trunk:
2453/// - BATCHED-path layers (`gdn`): the exact token-major inputs the round's ONE gdn_scan
2454/// consumed. A prefix re-run of the SAME kernel (t=j) from the snapshot state is bit-identical
2455/// to the first j iterations of the verify's scan — the kernel's t-loop carries state in
2456/// registers and iteration t never depends on T. `qkv_mixed` (the conv input) feeds the
2457/// pure-copy ring rebuild.
2458/// - PER-COLUMN-path layers (`cols`): dtod clones of (conv_state, ssm_state) taken after each
2459/// column 0..t-2 — pure copies of the actual chain states (the last column is never a rebuild
2460/// target: j <= t-1).
2461/// Full-attn layers need nothing: their verify KV rows are bit-identical to eager's (the
2462/// decode-exact contract; verify-probe pins it), so rollback = len truncation.
2463struct GdnStash {
2464 qkv_mixed: CudaSlice<f32>, // [t, conv_dim] token-major (conv input)
2465 q_l2: CudaSlice<f32>,
2466 k_l2: CudaSlice<f32>,
2467 v_g: CudaSlice<f32>, // [t, num_v, d_state]
2468 g_log: CudaSlice<f32>,
2469 beta: CudaSlice<f32>, // [t, num_v]
2470}
2471pub(crate) struct VerifyCkpt {
2472 gdn: Vec<Option<GdnStash>>, // [n_layer], Some iff batched linear path ran
2473 cols: Vec<Option<Vec<(CudaSlice<f32>, CudaSlice<f32>)>>>, // [n_layer][col] = (conv, ssm) after col
2474}
2475/// Opaque handle for the dspark round (dflash.rs) — VerifyCkpt stays spec-private.
2476pub(crate) struct DsparkVerifyCkpt(VerifyCkpt);
2477
2478/// Engine-bundle slice 3 (DSF-ROUNDCOST-20260820 §2 row 4 / §5 rank 1): bucketed CUDA
2479/// graphs for the dspark verify's LINEAR-layer segments. The measured verify is ~2,800
2480/// eager launches whose residual cost is DEVICE-side per-launch overhead (slice 2 proved
2481/// host dispatch is not the binder: fully-deferred dispatch bought ~0 wall). The 48 GDN
2482/// layers between full-attention layers are shape-static given vt — no positions, no
2483/// t_kv, state addressed through pointer tables — so runs of them capture per
2484/// (segment, vt) and replay as ONE graph launch each. Full-attention layers stay eager
2485/// (their per-row append/fa arm picks are t_kv-driven — the exec-update extension).
2486///
2487/// Per round out-of-graph: one pointer-table refresh (gdn ping-pong moves the canonical
2488/// handles), one input-staging copy per segment, host parity bookkeeping. Captured via
2489/// `capture_graph_retained` (2 warmups + capture, keeper retains warmup transients so
2490/// pool addresses stay stable); the warmups EXECUTE, so segment conv/ssm state is saved
2491/// before and restored after — the graph's first real launch starts from the exact
2492/// pre-round state. The ckpt column stash rides persistent slabs (written inside the
2493/// graph as memcpy nodes); commit reads them via `dspark_commit_prefix_slab`.
2494/// `MEMRA_DSPARK_VERIFY_GRAPH=0` reverts to the eager walk (byte-identical body).
2495pub(crate) struct DsparkVerifyGraphs {
2496 /// Linear-attention layer indices ascending; `lin_pos[il]` = index into the vecs.
2497 lin: Vec<usize>,
2498 lin_pos: std::collections::HashMap<usize, usize>,
2499 /// [n_lin x 6] pointer table (conv, s0, s1, conv, s1, s0 per layer), refreshed per
2500 /// verify from the live handles; layer il's slice starts at lin_pos[il]*6.
2501 table_all: CudaSlice<u64>,
2502 host_table: Vec<u64>,
2503 /// Persistent per-layer ckpt stash slabs: row r of the verify at slab offset
2504 /// r*words. Shared by every (segment, vt) bucket — one verify runs at a time.
2505 stash_conv: Vec<CudaSlice<f32>>,
2506 stash_ssm: Vec<CudaSlice<f32>>,
2507 conv_words: usize,
2508 ssm_words: usize,
2509 /// Per-vt input/output staging (stable addresses the graphs bake).
2510 stage: std::collections::HashMap<usize, (CudaSlice<f32>, CudaSlice<f32>)>,
2511 /// Per-vt dflash tap-sink buffers — the captured segments bake the tap dst address,
2512 /// so the sink buffer must live (and persist) with the graphs, not with the round.
2513 pub(crate) tap_bufs: std::collections::HashMap<usize, CudaSlice<f32>>,
2514 graphs: std::collections::HashMap<(usize, usize), DsparkSegGraph>,
2515 /// Warmup-corruption guard scratch: pre-capture conv/ssm of every linear layer
2516 /// (sized n_lin — the slice-4c full-verify warmups execute the whole walk).
2517 save_conv: CudaSlice<f32>,
2518 save_ssm: CudaSlice<f32>,
2519 max_run: usize,
2520 n_embd: usize,
2521 /// Set by the verify walk: this round's linear ckpt lives in the slabs (the caller
2522 /// commits through `dspark_commit_prefix_slab` instead of the cols arm).
2523 pub(crate) round_slab: bool,
2524 // ---- slice 4c: full-verify single graph per (vt, rung) ----
2525 /// Full-attention layer indices ascending; `fa_pos[il]` = index into the vec.
2526 fa: Vec<usize>,
2527 fa_pos: std::collections::HashMap<usize, usize>,
2528 /// [n_fa x 2 x t_cap] interleaved (k,v) base-pointer pairs, refreshed per verify;
2529 /// layer il's slice starts at `fa_pos[il] * 2 * t_cap` (the seqs twins read pairs
2530 /// [2z], z < t <= t_cap, so one t_cap-sized table serves every vt).
2531 fa_table: CudaSlice<u64>,
2532 fa_host_table: Vec<u64>,
2533 t_cap: usize,
2534 /// Per-vt position staging for the captured bodies — contents refreshed per round
2535 /// (rope reads row r; the seqs twins derive append slot and T_kv per z from it).
2536 pos_stage: std::collections::HashMap<usize, CudaSlice<i32>>,
2537 /// Full-verify graphs keyed (vt, rung_end, hi).
2538 full: std::collections::HashMap<(usize, usize, usize), DsparkSegGraph>,
2539 /// Largest n with every layer in [0, n) linear or full-attention (walk coverage).
2540 covered: usize,
2541 /// Every layer in [0, n) is linear or full-attention (no MLA/unknown mixers) — the
2542 /// full-verify capture walks all of them.
2543 walk_uniform: bool,
2544 /// Last `(captures, device graph-mem reserved bytes)` reading taken by
2545 /// `HybridModel::dspark_vg_admission_debt` — the two-point base of the MARGINAL debt
2546 /// projection (see `dspark_vg_debt_projection`; a mean-based reading extrapolated the
2547 /// pool's one-time shared allocation and reserved 8.5 GB of phantom VRAM).
2548 debt_obs: Option<(usize, usize)>,
2549}
2550
2551struct DsparkSegGraph {
2552 graph: cudarc::driver::CudaGraph,
2553 _keeper: Vec<Box<dyn std::any::Any + Send>>,
2554}
2555
2556/// Per-call arguments of [`HybridModel::qwen35_tparallel_fa_layer`] — one struct so the
2557/// eager walk and the slice-4c captured full-verify graphs hand the SAME body its two
2558/// modes without a second copy of the math.
2559pub(crate) struct FaLayerArgs<'a> {
2560 /// [T] per-row positions (device): rope reads them row-indexed; the seqs twins read
2561 /// them per-z (append slot = pos, T_kv = pos + 1).
2562 pub pos_d: &'a CudaSlice<i32>,
2563 /// Verify-level lazy per-row 1-element position buffers — only the per-row fallback
2564 /// arm builds/uses them (graph mode refuses that arm).
2565 pub pos_rows: &'a mut Option<Vec<CudaSlice<i32>>>,
2566 pub pos0: usize,
2567 pub seqs_append: bool,
2568 pub batch_fa_on: bool,
2569 /// Some((kv pointer table, offset-in-u64s, rung_end)) = captured-graph mode.
2570 pub graph_cap: Option<(&'a CudaSlice<u64>, usize, usize)>,
2571 /// ROUND-STREAM (lane/draftcost-moe, v0.100 train merge): Some((token stream, device
2572 /// round counter)) routes the FA attend through the dc rows kernels and the Linear
2573 /// mixer through `linear_attn_verify_t` (the stream arms the old inline body carried).
2574 /// Never armed together with `graph_cap` (the verify-level merge guard refuses).
2575 pub stream: Option<(&'a CudaSlice<u32>, &'a CudaSlice<i32>)>,
2576 /// VerifyCkpt for the stream-Linear arm's GdnStash install; None in graph mode and
2577 /// for FA layers that never touch it.
2578 pub ckpt: Option<&'a mut VerifyCkpt>,
2579}
2580
2581// SAFETY: `CudaGraph` is not marked Send by cudarc because its raw driver handles carry
2582// no automatic trait; CUDA driver graph handles are context-scoped rather than
2583// OS-thread-affine (the SpecPipeSessionPtr precedent above). The ctx lives in
2584// `HybridModel::dspark_vgraphs` behind a Mutex and every touch happens on the engine's
2585// single decode-stream thread.
2586unsafe impl Send for DsparkVerifyGraphs {}
2587
2588impl DsparkVerifyGraphs {
2589 /// Live capture count (segment + full graphs) — the denominator of
2590 /// [`dspark_vg_debt_projection`]'s observed bytes/capture mean.
2591 pub(crate) fn captures(&self) -> usize {
2592 self.graphs.len() + self.full.len()
2593 }
2594
2595 /// Take the marginal-growth debt reading and record this observation for the next one.
2596 /// Called under the pool mutex by `HybridModel::dspark_vg_admission_debt`.
2597 pub(crate) fn admission_debt(&mut self, reserved_bytes: usize) -> usize {
2598 let captures = self.captures();
2599 let debt =
2600 dspark_vg_debt_projection(captures, dspark_vg_cap(), reserved_bytes, self.debt_obs);
2601 if captures > 0 {
2602 match self.debt_obs {
2603 Some((c0, _)) if captures <= c0 => {}
2604 _ => self.debt_obs = Some((captures, reserved_bytes)),
2605 }
2606 }
2607 debt
2608 }
2609
2610 /// Build for this cache's shape. None when there are no linear layers, sizes are
2611 /// non-uniform, or the trunk keeps a gemma4 config (never on the qwen35 family).
2612 pub(crate) fn new(
2613 e: &Engine,
2614 cache: &Cache,
2615 t_max: usize,
2616 n_embd: usize,
2617 ) -> Result<Option<Self>, Box<dyn std::error::Error>> {
2618 let lin: Vec<usize> = (0..cache.recur.len())
2619 .filter(|&il| cache.recur[il].is_some())
2620 .collect();
2621 if lin.is_empty() || t_max < 2 {
2622 return Ok(None);
2623 }
2624 let first = cache.recur[lin[0]].as_ref().unwrap();
2625 let (conv_words, ssm_words) = (first.conv_state.len(), first.ssm_state.len());
2626 for &il in &lin {
2627 let rl = cache.recur[il].as_ref().unwrap();
2628 if rl.conv_state.len() != conv_words || rl.ssm_state.len() != ssm_words {
2629 return Ok(None);
2630 }
2631 }
2632 let n = lin.len();
2633 let mut lin_pos = std::collections::HashMap::with_capacity(n);
2634 for (k, &il) in lin.iter().enumerate() {
2635 lin_pos.insert(il, k);
2636 }
2637 // longest run of consecutive linear layers (save-scratch sizing)
2638 let mut max_run = 1usize;
2639 let mut run = 1usize;
2640 for w in lin.windows(2) {
2641 if w[1] == w[0] + 1 {
2642 run += 1;
2643 max_run = max_run.max(run);
2644 } else {
2645 run = 1;
2646 }
2647 }
2648 let rows = t_max - 1;
2649 let mut stash_conv = Vec::with_capacity(n);
2650 let mut stash_ssm = Vec::with_capacity(n);
2651 for _ in 0..n {
2652 stash_conv.push(e.uninit(rows * conv_words)?);
2653 stash_ssm.push(e.uninit(rows * ssm_words)?);
2654 }
2655 let host_table = vec![0u64; n * 6];
2656 let table_all = e.htod_u64(&host_table)?;
2657 // slice 4c: full-attention census for the full-verify graphs.
2658 let fa: Vec<usize> = (0..cache.kv.len())
2659 .filter(|&il| cache.kv[il].is_some())
2660 .collect();
2661 let mut fa_pos = std::collections::HashMap::with_capacity(fa.len());
2662 for (k, &il) in fa.iter().enumerate() {
2663 fa_pos.insert(il, k);
2664 }
2665 let n_layers = cache.kv.len().max(cache.recur.len());
2666 // exactly one of (linear state, kv cache) per layer — no MLA/unknown mixers.
2667 let walk_uniform = (0..n_layers).all(|il| {
2668 cache.recur.get(il).is_some_and(|r| r.is_some())
2669 != cache.kv.get(il).is_some_and(|k| k.is_some())
2670 });
2671 // Contiguous covered prefix: the largest n such that every layer in [0, n) is
2672 // linear or full-attention. The TRUNK walk is [0, layers.len()) and the cache
2673 // vecs can carry EXTRA state slots past it (the q38 export keeps the MTP head
2674 // layer's kv at the tail — hi == lin+fa never held, the s4c battery's zero
2675 // 'full' captures). The full-graph guard is walk coverage, not slot arithmetic.
2676 let covered = (0..n_layers)
2677 .take_while(|il| lin_pos.contains_key(il) || fa_pos.contains_key(il))
2678 .count();
2679 let t_cap = t_max;
2680 let fa_host_table = vec![0u64; fa.len() * 2 * t_cap];
2681 let fa_table = e.htod_u64(&fa_host_table)?;
2682 Ok(Some(Self {
2683 lin,
2684 lin_pos,
2685 table_all,
2686 host_table,
2687 stash_conv,
2688 stash_ssm,
2689 conv_words,
2690 ssm_words,
2691 stage: std::collections::HashMap::new(),
2692 tap_bufs: std::collections::HashMap::new(),
2693 graphs: std::collections::HashMap::new(),
2694 save_conv: e.uninit(n * conv_words)?,
2695 save_ssm: e.uninit(n * ssm_words)?,
2696 max_run,
2697 n_embd,
2698 round_slab: false,
2699 fa,
2700 fa_pos,
2701 fa_table,
2702 fa_host_table,
2703 t_cap,
2704 pos_stage: std::collections::HashMap::new(),
2705 full: std::collections::HashMap::new(),
2706 covered,
2707 walk_uniform,
2708 debt_obs: None,
2709 }))
2710 }
2711
2712 /// Rebuild the pointer tables from the live handles (once per verify — the gdn
2713 /// ping-pong swaps the canonical/alt handles between rounds; a fresh generation's
2714 /// cache buffers land at new addresses; a stale table would read the wrong state).
2715 pub(crate) fn refresh_tables(
2716 &mut self,
2717 e: &Engine,
2718 cache: &Cache,
2719 ) -> Result<(), Box<dyn std::error::Error>> {
2720 use cudarc::driver::DevicePtr;
2721 {
2722 let s = &e.gpu.stream();
2723 for (k, &il) in self.lin.iter().enumerate() {
2724 let rl = cache.recur[il].as_ref().unwrap();
2725 let (pc, _g0) = rl.conv_state.device_ptr(s);
2726 let (p0, _g1) = rl.ssm_state.device_ptr(s);
2727 let (p1, _g2) = rl.ssm_state_alt.device_ptr(s);
2728 let o = k * 6;
2729 self.host_table[o] = pc as u64;
2730 self.host_table[o + 1] = p0 as u64;
2731 self.host_table[o + 2] = p1 as u64;
2732 self.host_table[o + 3] = pc as u64;
2733 self.host_table[o + 4] = p1 as u64;
2734 self.host_table[o + 5] = p0 as u64;
2735 }
2736 for (k, &il) in self.fa.iter().enumerate() {
2737 let kvl = cache.kv[il].as_ref().unwrap();
2738 let (pk, _g0) = kvl.k.device_ptr(s);
2739 let (pv, _g1) = kvl.v.device_ptr(s);
2740 let o = k * 2 * self.t_cap;
2741 for z in 0..self.t_cap {
2742 self.fa_host_table[o + 2 * z] = pk as u64;
2743 self.fa_host_table[o + 2 * z + 1] = pv as u64;
2744 }
2745 }
2746 }
2747 e.htod_u64_into(&self.host_table, &mut self.table_all)?;
2748 if !self.fa_host_table.is_empty() {
2749 e.htod_u64_into(&self.fa_host_table, &mut self.fa_table)?;
2750 }
2751 Ok(())
2752 }
2753
2754 /// Slice 4c eligibility: Some(rung_end) when this round can replay (or capture) a
2755 /// full-verify graph — the whole walk [lo, hi) is covered, every layer is linear or
2756 /// full-attention, and ALL of the round's per-row t_kv values take the v4-seqs arm
2757 /// on ONE `fa_split_keys` ladder step that the rung also sits on (the straddle law;
2758 /// both gates are t_kv intervals, so ends-inside means all-inside). The rung is the
2759 /// round's next power of two — grid/partial sizing only (`n_splits_max` is pure
2760 /// stride; splits >= ns_eff write the empty partial the combine never reads), so one
2761 /// captured graph is bit-identical for every round the rung covers.
2762 #[allow(clippy::too_many_arguments)]
2763 pub(crate) fn full_rung(
2764 &self,
2765 model: &crate::hybrid::HybridModel,
2766 cache: &Cache,
2767 lo: usize,
2768 hi: usize,
2769 t: usize,
2770 seqs_arms_on: bool,
2771 ) -> Option<usize> {
2772 if std::env::var("MEMRA_DSPARK_FULLG_DEBUG").as_deref() == Ok("1") {
2773 static ONCE: std::sync::Once = std::sync::Once::new();
2774 let len0 = self
2775 .fa
2776 .first()
2777 .and_then(|&il| cache.kv[il].as_ref())
2778 .map(|k| k.len);
2779 ONCE.call_once(|| {
2780 eprintln!(
2781 "[fullg-debug] walk_uniform={} covered={} seqs_arms_on={} fa_rows_on={} t={} lo={} hi={} lin={} fa={} t_cap={} len0={:?}",
2782 self.walk_uniform, self.covered, seqs_arms_on, dspark_fa_rows_on(), t, lo, hi,
2783 self.lin.len(), self.fa.len(), self.t_cap, len0
2784 );
2785 });
2786 }
2787 if !self.walk_uniform
2788 || !seqs_arms_on
2789 || !dspark_fa_rows_on()
2790 || t < 2
2791 || lo != 0
2792 || hi > self.covered
2793 || t > self.t_cap
2794 || self.fa.is_empty()
2795 {
2796 return None;
2797 }
2798 let cfg = &model.cfg;
2799 let head_dim_global = cfg.head_dim_k as usize;
2800 let nkv = cfg.n_head_kv as usize;
2801 let kvl0 = cache.kv[self.fa[0]].as_ref().unwrap();
2802 // the z-batched twins read stacked rows at the cache's kv dims — must equal the
2803 // projection stride (the body's guard, hoisted so ineligible models fall back
2804 // instead of refusing mid-capture).
2805 let geom = cfg.full_attention_geometry_at(self.fa[0] as u32);
2806 let kv_dim = geom.n_head_kv as usize * geom.head_dim_k as usize;
2807 if kvl0.kv_dim_k != kv_dim || kvl0.kv_dim_v != kv_dim {
2808 return None;
2809 }
2810 let len0 = kvl0.len;
2811 let (t_kv_first, t_kv_last) = (len0 + 1, len0 + t);
2812 if !crate::fa_seqs_eligible(t_kv_first, head_dim_global)
2813 || !crate::fa_seqs_eligible(t_kv_last, head_dim_global)
2814 || crate::fa_split_keys(t_kv_first, nkv) != crate::fa_split_keys(t_kv_last, nkv)
2815 {
2816 return None;
2817 }
2818 let rung = t_kv_last.next_power_of_two().max(256);
2819 if crate::fa_split_keys(rung, nkv) != crate::fa_split_keys(t_kv_last, nkv) {
2820 return None;
2821 }
2822 Some(rung)
2823 }
2824
2825 /// Run the WHOLE verify walk [lo, hi) as one captured graph at (vt=t, rung): stage
2826 /// the residual + refresh the per-vt position staging, capture on first encounter
2827 /// (2 executing warmups bracketed by a full linear-state save/restore; KV warmup
2828 /// appends write the exact slots the replay writes — idempotent), launch, then apply
2829 /// the host bookkeeping the captured body skipped (per-linear-layer parity swap for
2830 /// odd t, per-fa-layer len bump). Returns the fresh residual.
2831 #[allow(clippy::too_many_arguments)]
2832 pub(crate) fn run_full(
2833 &mut self,
2834 model: &crate::hybrid::HybridModel,
2835 e: &Engine,
2836 lo: usize,
2837 hi: usize,
2838 x: &CudaSlice<f32>,
2839 t: usize,
2840 pos0: usize,
2841 rung: usize,
2842 cache: &mut Cache,
2843 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2844 let n_embd = self.n_embd;
2845 if !self.stage.contains_key(&t) {
2846 let xin = e.uninit(t * n_embd)?;
2847 let xout = e.uninit(t * n_embd)?;
2848 self.stage.insert(t, (xin, xout));
2849 }
2850 if !self.pos_stage.contains_key(&t) {
2851 self.pos_stage.insert(t, e.htod_i32(&vec![0i32; t])?);
2852 }
2853 // Per-round refresh: position contents + input staging (both addresses are baked
2854 // by the captured bodies; only their CONTENTS change round to round).
2855 {
2856 let pos_host: Vec<i32> = (0..t).map(|r| (pos0 + r) as i32).collect();
2857 let pb = self.pos_stage.get_mut(&t).unwrap();
2858 e.htod_i32_into(pb, &pos_host)?;
2859 let (xin, _) = self.stage.get_mut(&t).unwrap();
2860 e.copy_into(xin, 0, x, t * n_embd)?;
2861 }
2862 let key = (t, rung, hi);
2863 if !self.full.contains_key(&key) {
2864 // The warmups EXECUTE the whole walk on live state — save every linear
2865 // layer's conv + canonical ssm first, restore after (KV needs no restore:
2866 // graph mode never bumps host lens and the appends write this round's own
2867 // slots).
2868 for (k, &il) in self.lin.iter().enumerate() {
2869 let rl = cache.recur[il].as_ref().unwrap();
2870 e.copy_into(
2871 &mut self.save_conv,
2872 k * self.conv_words,
2873 &rl.conv_state,
2874 self.conv_words,
2875 )?;
2876 e.copy_into(
2877 &mut self.save_ssm,
2878 k * self.ssm_words,
2879 &rl.ssm_state,
2880 self.ssm_words,
2881 )?;
2882 }
2883 let (graph, keeper) = {
2884 let table_all = &self.table_all;
2885 let lin_pos = &self.lin_pos;
2886 let fa_pos = &self.fa_pos;
2887 let fa_table = &self.fa_table;
2888 let t_cap = self.t_cap;
2889 let stash_conv = &mut self.stash_conv;
2890 let stash_ssm = &mut self.stash_ssm;
2891 let pos_d: &CudaSlice<i32> = &self.pos_stage[&t];
2892 let (xin, xout) = self
2893 .stage
2894 .get_mut(&t)
2895 .map(|(a, b)| (&*a, b))
2896 .expect("stage bucket created above");
2897 let cache_ref: &mut Cache = cache;
2898 let iflag = if std::env::var("MEMRA_DSPARK_VG_AUTOFREE").as_deref() == Ok("1") {
2899 cudarc::driver::sys::CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH
2900 } else {
2901 cudarc::driver::sys::CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_USE_NODE_PRIORITY
2902 };
2903 e.capture_graph_retained_flags(iflag, move |e| {
2904 let mut xc: Option<CudaSlice<f32>> = None;
2905 for il in lo..hi {
2906 let xr: &CudaSlice<f32> = xc.as_ref().unwrap_or(xin);
2907 let nx = if let Some(&k) = lin_pos.get(&il) {
2908 model.qwen35_tparallel_linear_layer(
2909 e,
2910 il,
2911 xr,
2912 t,
2913 cache_ref,
2914 None,
2915 Some((&mut stash_conv[k], &mut stash_ssm[k])),
2916 Some((table_all, k * 6)),
2917 )?
2918 } else if let Some(&kf) = fa_pos.get(&il) {
2919 let mut no_rows: Option<Vec<CudaSlice<i32>>> = None;
2920 model.qwen35_tparallel_fa_layer(
2921 e,
2922 il,
2923 xr,
2924 t,
2925 cache_ref,
2926 FaLayerArgs {
2927 pos_d,
2928 pos_rows: &mut no_rows,
2929 pos0,
2930 seqs_append: true,
2931 batch_fa_on: true,
2932 graph_cap: Some((fa_table, kf * 2 * t_cap, rung)),
2933 stream: None,
2934 ckpt: None,
2935 },
2936 )?
2937 } else {
2938 return Err(format!(
2939 "run_full: layer {il} is neither linear nor full-attention"
2940 )
2941 .into());
2942 };
2943 xc = Some(nx);
2944 }
2945 e.copy_into(xout, 0, xc.as_ref().unwrap(), t * n_embd)?;
2946 Ok(())
2947 })?
2948 };
2949 // Undo the net host parity motion of the 3 body runs (each run swaps iff t
2950 // is odd -> 3 runs = net one swap), then restore the device state the
2951 // warmups consumed (walk scope only — layers past hi never executed). The
2952 // launch below then behaves exactly like one run.
2953 if t % 2 == 1 {
2954 for &il in &self.lin {
2955 if il < lo || il >= hi {
2956 continue;
2957 }
2958 let rl = cache.recur[il].as_mut().unwrap();
2959 std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
2960 }
2961 }
2962 for (k, &il) in self.lin.iter().enumerate() {
2963 if il < lo || il >= hi {
2964 continue;
2965 }
2966 let rl = cache.recur[il].as_mut().unwrap();
2967 let (cw, sw) = (self.conv_words, self.ssm_words);
2968 {
2969 let sv = e.view(&self.save_conv, self.lin.len() * cw);
2970 let win = sv.slice(k * cw..(k + 1) * cw);
2971 e.copy_view_into(&mut rl.conv_state, 0, &win, cw)?;
2972 }
2973 {
2974 let sv = e.view(&self.save_ssm, self.lin.len() * sw);
2975 let win = sv.slice(k * sw..(k + 1) * sw);
2976 e.copy_view_into(&mut rl.ssm_state, 0, &win, sw)?;
2977 }
2978 }
2979 if std::env::var("MEMRA_GRAPH_CENSUS").as_deref() == Ok("1") {
2980 if let Ok(c) = crate::graph_update::node_census(&graph) {
2981 eprintln!("[dspark-vg-census] full vt={t} rung={rung} {c:?}");
2982 }
2983 }
2984 self.full.insert(
2985 key,
2986 DsparkSegGraph {
2987 graph,
2988 _keeper: keeper,
2989 },
2990 );
2991 }
2992 self.full[&key].graph.launch()?;
2993 // Host bookkeeping for the replayed body (captured host code does not re-run):
2994 // gdn parity swap per linear layer (t odd), kv len bump per fa layer — scoped
2995 // to the WALK [lo, hi): the cache can carry extra state slots past it (the MTP
2996 // head layer's kv) that the walk never touches.
2997 if t % 2 == 1 {
2998 for &il in &self.lin {
2999 if il < lo || il >= hi {
3000 continue;
3001 }
3002 let rl = cache.recur[il].as_mut().unwrap();
3003 std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
3004 }
3005 }
3006 for &il in &self.fa {
3007 if il < lo || il >= hi {
3008 continue;
3009 }
3010 cache.kv[il].as_mut().unwrap().len += t;
3011 }
3012 let (_, xout) = self.stage.get(&t).unwrap();
3013 let mut out = e.uninit(t * n_embd)?;
3014 e.copy_into(&mut out, 0, xout, t * n_embd)?;
3015 Ok(out)
3016 }
3017
3018 /// Run layers [start, end) (all linear) as one captured graph at this vt: stage the
3019 /// residual into the bucket's x_in, capture on first encounter (2 executing warmups
3020 /// bracketed by a segment state save/restore), launch, then apply the host parity
3021 /// bookkeeping the captured body would have done. Returns the fresh residual.
3022 #[allow(clippy::too_many_arguments)]
3023 fn run_segment(
3024 &mut self,
3025 model: &crate::hybrid::HybridModel,
3026 e: &Engine,
3027 start: usize,
3028 end: usize,
3029 x: &CudaSlice<f32>,
3030 t: usize,
3031 cache: &mut Cache,
3032 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3033 let n_embd = self.n_embd;
3034 debug_assert!(end - start <= self.max_run);
3035 if !self.stage.contains_key(&t) {
3036 let xin = e.uninit(t * n_embd)?;
3037 let xout = e.uninit(t * n_embd)?;
3038 self.stage.insert(t, (xin, xout));
3039 }
3040 // Stage the residual at the bucket's baked input address.
3041 {
3042 let (xin, _) = self.stage.get_mut(&t).unwrap();
3043 e.copy_into(xin, 0, x, t * n_embd)?;
3044 }
3045 let key = (start, t);
3046 if !self.graphs.contains_key(&key) {
3047 // The 2 warmups EXECUTE the segment on live state — save conv + the canonical
3048 // ssm of every segment layer first, restore after, so the graph's first real
3049 // launch starts from the exact pre-round state (bytes gated e2e).
3050 for (k, il) in (start..end).enumerate() {
3051 let rl = cache.recur[il].as_ref().unwrap();
3052 e.copy_into(
3053 &mut self.save_conv,
3054 k * self.conv_words,
3055 &rl.conv_state,
3056 self.conv_words,
3057 )?;
3058 e.copy_into(
3059 &mut self.save_ssm,
3060 k * self.ssm_words,
3061 &rl.ssm_state,
3062 self.ssm_words,
3063 )?;
3064 }
3065 let (graph, keeper) = {
3066 let table_all = &self.table_all;
3067 let lin_pos = &self.lin_pos;
3068 let stash_conv = &mut self.stash_conv;
3069 let stash_ssm = &mut self.stash_ssm;
3070 let (xin, xout) = self
3071 .stage
3072 .get_mut(&t)
3073 .map(|(a, b)| (&*a, b))
3074 .expect("stage bucket created above");
3075 let cache_ref: &mut Cache = cache;
3076 // Slice 4 (fa-execupdate lane): USE_NODE_PRIORITY instead of
3077 // AUTO_FREE_ON_LAUNCH. The slice-3 measured limiter was AUTO_FREE's
3078 // launch-time mem-pool scan — 25.6 us per cuGraphLaunch x 16 segments
3079 // = ~0.41 ms/round, most of the eager-launch savings. The captured
3080 // body's cuMemAllocAsync transients are BALANCED by in-graph frees
3081 // (every transient drops inside the capture region — the generic
3082 // capture path's census precedent, 1589/1589), so AUTO_FREE has
3083 // nothing to reclaim and the graph is legal to instantiate without
3084 // it; PRIORITY is the flag the gemma slotted door ships for exactly
3085 // this reason (both alternatives drop the scan; UPLOAD via
3086 // cuGraphInstantiateWithFlags is WithParams-only and refused).
3087 // MEMRA_DSPARK_VG_AUTOFREE=1 reverts; MEMRA_GRAPH_CENSUS=1 prints
3088 // the node census at capture (the ALLOC==FREE receipt).
3089 let iflag = if std::env::var("MEMRA_DSPARK_VG_AUTOFREE").as_deref() == Ok("1") {
3090 cudarc::driver::sys::CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH
3091 } else {
3092 cudarc::driver::sys::CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_USE_NODE_PRIORITY
3093 };
3094 e.capture_graph_retained_flags(iflag, move |e| {
3095 let mut xc: Option<CudaSlice<f32>> = None;
3096 for il in start..end {
3097 let k = lin_pos[&il];
3098 let xr: &CudaSlice<f32> = xc.as_ref().unwrap_or(xin);
3099 let nx = model.qwen35_tparallel_linear_layer(
3100 e,
3101 il,
3102 xr,
3103 t,
3104 cache_ref,
3105 None,
3106 Some((&mut stash_conv[k], &mut stash_ssm[k])),
3107 Some((table_all, k * 6)),
3108 )?;
3109 xc = Some(nx);
3110 }
3111 e.copy_into(xout, 0, xc.as_ref().unwrap(), t * n_embd)?;
3112 Ok(())
3113 })?
3114 };
3115 // Undo the net host parity motion of the 3 body runs (each run swaps iff t
3116 // is odd -> 3 runs = net one swap), then restore the device state the
3117 // warmups consumed. The launch below then behaves exactly like one run.
3118 if t % 2 == 1 {
3119 for il in start..end {
3120 let rl = cache.recur[il].as_mut().unwrap();
3121 std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
3122 }
3123 }
3124 for (k, il) in (start..end).enumerate() {
3125 let rl = cache.recur[il].as_mut().unwrap();
3126 let (cw, sw) = (self.conv_words, self.ssm_words);
3127 {
3128 let sv = e.view(&self.save_conv, self.lin.len() * cw);
3129 let win = sv.slice(k * cw..(k + 1) * cw);
3130 e.copy_view_into(&mut rl.conv_state, 0, &win, cw)?;
3131 }
3132 {
3133 let sv = e.view(&self.save_ssm, self.lin.len() * sw);
3134 let win = sv.slice(k * sw..(k + 1) * sw);
3135 e.copy_view_into(&mut rl.ssm_state, 0, &win, sw)?;
3136 }
3137 }
3138 if std::env::var("MEMRA_GRAPH_CENSUS").as_deref() == Ok("1") {
3139 if let Ok(c) = crate::graph_update::node_census(&graph) {
3140 eprintln!("[dspark-vg-census] seg={start}..{end} vt={t} {c:?}");
3141 }
3142 }
3143 self.graphs.insert(
3144 key,
3145 DsparkSegGraph {
3146 graph,
3147 _keeper: keeper,
3148 },
3149 );
3150 }
3151 self.graphs[&key].graph.launch()?;
3152 // Host parity bookkeeping for the replayed body (the captured host swaps do not
3153 // re-run at replay).
3154 if t % 2 == 1 {
3155 for il in start..end {
3156 let rl = cache.recur[il].as_mut().unwrap();
3157 std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
3158 }
3159 }
3160 let (_, xout) = self.stage.get(&t).unwrap();
3161 let mut out = e.uninit(t * n_embd)?;
3162 e.copy_into(&mut out, 0, xout, t * n_embd)?;
3163 Ok(out)
3164 }
3165
3166 /// Pool freeze check (`dspark_vg_cap`): below the ceiling new keys may capture.
3167 fn can_capture(&self) -> bool {
3168 self.graphs.len() + self.full.len() < dspark_vg_cap()
3169 }
3170
3171 /// Round-atomic segment-door readiness: TRUE when this round's walk can ride the
3172 /// per-(segment, vt) graphs without a NEW capture past the pool ceiling — every
3173 /// linear run in [lo, hi) already has its (run_start, t) key, or capture is still
3174 /// allowed. FALSE sends the WHOLE round down the eager cols-ckpt walk: a partial
3175 /// refusal would stash some layers in the ctx slabs and others in the round's cols
3176 /// while one commit reads only one of them.
3177 pub(crate) fn segments_ready(
3178 &self,
3179 model: &crate::hybrid::HybridModel,
3180 lo: usize,
3181 hi: usize,
3182 t: usize,
3183 ) -> bool {
3184 if self.can_capture() {
3185 return true;
3186 }
3187 let mut il = lo;
3188 while il < hi {
3189 if matches!(model.layers[il].mixer, Mixer::Linear(_)) {
3190 let start = il;
3191 while il < hi && matches!(model.layers[il].mixer, Mixer::Linear(_)) {
3192 il += 1;
3193 }
3194 if !self.graphs.contains_key(&(start, t)) {
3195 return false;
3196 }
3197 } else {
3198 il += 1;
3199 }
3200 }
3201 true
3202 }
3203
3204 /// Widest verify window this pool was built for. A caller whose round exceeds it must
3205 /// take the eager walk: the stash slabs hold `t_capacity() - 1` column rows, and slicing
3206 /// past them is a panic rather than a refusal.
3207 pub(crate) fn t_capacity(&self) -> usize {
3208 self.t_cap
3209 }
3210
3211 /// Slab row (conv, ssm) device pointers + lengths for the commit restore of column
3212 /// `row` (0-based) of layer `il`. None for non-linear layers.
3213 pub(crate) fn slab_row(
3214 &self,
3215 e: &Engine,
3216 il: usize,
3217 row: usize,
3218 ) -> Option<(u64, u64, usize, usize)> {
3219 use cudarc::driver::DevicePtr;
3220 let k = *self.lin_pos.get(&il)?;
3221 let s = &e.gpu.stream();
3222 let (pc, _g0) = self.stash_conv[k].device_ptr(s);
3223 let (ps, _g1) = self.stash_ssm[k].device_ptr(s);
3224 Some((
3225 pc as u64 + (row * self.conv_words * 4) as u64,
3226 ps as u64 + (row * self.ssm_words * 4) as u64,
3227 self.conv_words,
3228 self.ssm_words,
3229 ))
3230 }
3231}
3232
3233impl VerifyCkpt {
3234 fn new(n_layer: usize) -> Self {
3235 VerifyCkpt {
3236 gdn: (0..n_layer).map(|_| None).collect(),
3237 cols: (0..n_layer).map(|_| None).collect(),
3238 }
3239 }
3240}
3241
3242/// The stage-0/TX half of one PP verify. The boundary slot is the ownership token: stage 1
3243/// consumes exactly the slot selected by `tx()` / `tx_pipelined()`, never a slot inferred from
3244/// a logical round number.
3245struct VerifyBoundaryTicket {
3246 rt: &'static crate::pp::PpNRt,
3247 caller_stream: std::sync::Arc<cudarc::driver::CudaStream>,
3248 slot: usize,
3249 pos0: usize,
3250 t: usize,
3251 payload: usize,
3252 n_st: usize,
3253 pipelined: bool,
3254 pp_anatomy: bool,
3255 pp_started: std::time::Instant,
3256 reverse_ms: f64,
3257 stage0_ms: f64,
3258 tx_ms: f64,
3259 trace: Option<SpecPipeTraceCtx>,
3260}
3261
3262/// Explicit OPTIPIPE diagnostic control. Forced modes are set only by `optipipe-gate`; the
3263/// increment-2 controller can also be armed by the server's fresh-process research door.
3264#[derive(Clone, Copy, Debug, PartialEq, Eq)]
3265pub enum OptiForkGateMode {
3266 Disabled,
3267 Hit,
3268 Miss,
3269 Alternate,
3270 Abort,
3271 Controller,
3272}
3273
3274static OPTI_FORK_GATE_MODE: std::sync::atomic::AtomicU8 = std::sync::atomic::AtomicU8::new(0);
3275static OPTI_CONTROLLER_THRESHOLD: std::sync::atomic::AtomicU32 =
3276 std::sync::atomic::AtomicU32::new(0);
3277static OPTI_FORK_ATTEMPTS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
3278static OPTI_FORK_HITS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
3279static OPTI_FORK_MISSES: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
3280static OPTI_FORK_ABORT_DRAINS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
3281static OPTI_FORK_REFUSALS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
3282static OPTI_GATE_CHECKS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
3283static OPTI_GATE_ADMITS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
3284static OPTI_GATE_REJECTS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
3285static OPTI_RECONCILES: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
3286static OPTI_WASTED_DRAFT_TOKENS: std::sync::atomic::AtomicU64 =
3287 std::sync::atomic::AtomicU64::new(0);
3288static OPTI_SHADOW_DRAFT_TOKENS: std::sync::atomic::AtomicU64 =
3289 std::sync::atomic::AtomicU64::new(0);
3290static OPTI_BREAKER_TRIPS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
3291
3292impl OptiForkGateMode {
3293 fn code(self) -> u8 {
3294 match self {
3295 Self::Disabled => 0,
3296 Self::Hit => 1,
3297 Self::Miss => 2,
3298 Self::Alternate => 3,
3299 Self::Abort => 4,
3300 Self::Controller => 5,
3301 }
3302 }
3303
3304 fn configured() -> Self {
3305 match OPTI_FORK_GATE_MODE.load(std::sync::atomic::Ordering::Relaxed) {
3306 1 => Self::Hit,
3307 2 => Self::Miss,
3308 3 => Self::Alternate,
3309 4 => Self::Abort,
3310 5 => Self::Controller,
3311 _ => Self::Disabled,
3312 }
3313 }
3314
3315 fn action(self, generation: u64) -> OptiForkAction {
3316 match self {
3317 Self::Hit => OptiForkAction::Hit,
3318 Self::Miss => OptiForkAction::Miss,
3319 Self::Alternate if generation & 1 == 0 => OptiForkAction::Hit,
3320 Self::Alternate => OptiForkAction::Miss,
3321 Self::Abort => OptiForkAction::Abort,
3322 Self::Disabled | Self::Controller => {
3323 unreachable!("non-forced mode cannot choose a forced fork action")
3324 }
3325 }
3326 }
3327
3328 fn is_forced(self) -> bool {
3329 matches!(self, Self::Hit | Self::Miss | Self::Alternate | Self::Abort)
3330 }
3331}
3332
3333/// Arm or disarm the forced harness. Serving uses only `set_optipipe_controller_threshold`.
3334pub fn set_optipipe_gate_mode(mode: OptiForkGateMode) {
3335 OPTI_FORK_GATE_MODE.store(mode.code(), std::sync::atomic::Ordering::Relaxed);
3336}
3337
3338/// Arm the increment-2 diagnostic controller. The threshold applies to the uncalibrated
3339/// two-token draft-probability product. Serving can call this only through its explicit
3340/// fresh-process research door; the absent-door default remains byte-for-byte disabled.
3341pub fn set_optipipe_controller_threshold(threshold: f32) {
3342 assert!(threshold.is_finite() && (0.0..=1.0).contains(&threshold));
3343 OPTI_CONTROLLER_THRESHOLD.store(threshold.to_bits(), std::sync::atomic::Ordering::Relaxed);
3344 set_optipipe_gate_mode(OptiForkGateMode::Controller);
3345}
3346
3347#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
3348pub struct OptiForkGateStats {
3349 pub attempts: u64,
3350 pub hits: u64,
3351 pub misses: u64,
3352 pub abort_drains: u64,
3353 pub refusals: u64,
3354 pub gate_checks: u64,
3355 pub gate_admits: u64,
3356 pub gate_rejects: u64,
3357 pub reconciles: u64,
3358 pub wasted_draft_tokens: u64,
3359 pub shadow_draft_tokens: u64,
3360 pub breaker_trips: u64,
3361}
3362
3363#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
3364pub struct OptiForkStateIdentity {
3365 pub trunk_kv_bytes: usize,
3366 pub recurrent_bytes: usize,
3367 pub scratch_kv_bytes: usize,
3368 pub hidden_bytes: usize,
3369}
3370
3371pub fn reset_optipipe_gate_stats() {
3372 for counter in [
3373 &OPTI_FORK_ATTEMPTS,
3374 &OPTI_FORK_HITS,
3375 &OPTI_FORK_MISSES,
3376 &OPTI_FORK_ABORT_DRAINS,
3377 &OPTI_FORK_REFUSALS,
3378 &OPTI_GATE_CHECKS,
3379 &OPTI_GATE_ADMITS,
3380 &OPTI_GATE_REJECTS,
3381 &OPTI_RECONCILES,
3382 &OPTI_WASTED_DRAFT_TOKENS,
3383 &OPTI_SHADOW_DRAFT_TOKENS,
3384 &OPTI_BREAKER_TRIPS,
3385 ] {
3386 counter.store(0, std::sync::atomic::Ordering::Relaxed);
3387 }
3388}
3389
3390pub fn optipipe_gate_stats() -> OptiForkGateStats {
3391 let load = |v: &std::sync::atomic::AtomicU64| v.load(std::sync::atomic::Ordering::Relaxed);
3392 OptiForkGateStats {
3393 attempts: load(&OPTI_FORK_ATTEMPTS),
3394 hits: load(&OPTI_FORK_HITS),
3395 misses: load(&OPTI_FORK_MISSES),
3396 abort_drains: load(&OPTI_FORK_ABORT_DRAINS),
3397 refusals: load(&OPTI_FORK_REFUSALS),
3398 gate_checks: load(&OPTI_GATE_CHECKS),
3399 gate_admits: load(&OPTI_GATE_ADMITS),
3400 gate_rejects: load(&OPTI_GATE_REJECTS),
3401 reconciles: load(&OPTI_RECONCILES),
3402 wasted_draft_tokens: load(&OPTI_WASTED_DRAFT_TOKENS),
3403 shadow_draft_tokens: load(&OPTI_SHADOW_DRAFT_TOKENS),
3404 breaker_trips: load(&OPTI_BREAKER_TRIPS),
3405 }
3406}
3407
3408#[derive(Clone, Copy, Debug)]
3409struct OptiControllerPolicy {
3410 threshold: f32,
3411 consecutive_misses: u8,
3412 breaker_tripped: bool,
3413}
3414
3415impl OptiControllerPolicy {
3416 fn configured() -> Self {
3417 Self {
3418 threshold: f32::from_bits(
3419 OPTI_CONTROLLER_THRESHOLD.load(std::sync::atomic::Ordering::Relaxed),
3420 ),
3421 consecutive_misses: 0,
3422 breaker_tripped: false,
3423 }
3424 }
3425
3426 fn admit(&self, q_proxy: f32) -> bool {
3427 q_proxy.is_finite()
3428 && (0.0..=1.0).contains(&q_proxy)
3429 && (self.threshold == 0.0 || (!self.breaker_tripped && q_proxy >= self.threshold))
3430 }
3431
3432 /// Returns true exactly when this resolution newly trips the three-miss breaker.
3433 fn resolve(&mut self, hit: bool) -> bool {
3434 // q*=0 is the lane's explicit unconditional measurement arm. Its purpose is to price
3435 // every optimistic opportunity, so the safety breaker is measured separately and must
3436 // not silently turn this arm into "three attempts then serial".
3437 if self.threshold == 0.0 {
3438 self.consecutive_misses = 0;
3439 return false;
3440 }
3441 if hit {
3442 self.consecutive_misses = 0;
3443 return false;
3444 }
3445 self.consecutive_misses = self.consecutive_misses.saturating_add(1);
3446 if !self.breaker_tripped && self.consecutive_misses >= 3 {
3447 self.breaker_tripped = true;
3448 return true;
3449 }
3450 false
3451 }
3452}
3453
3454#[derive(Clone, Copy, Debug, PartialEq, Eq)]
3455enum OptiForkAction {
3456 Hit,
3457 Miss,
3458 Abort,
3459}
3460
3461#[derive(Clone, Copy, Debug, PartialEq, Eq)]
3462struct OptiForkGeneration {
3463 id: u64,
3464 slot: usize,
3465}
3466
3467#[derive(Default)]
3468struct OptiForkGenerationTracker {
3469 next: u64,
3470 live: [Option<u64>; 2],
3471}
3472
3473impl OptiForkGenerationTracker {
3474 fn reserve(&mut self) -> Result<OptiForkGeneration, Box<dyn std::error::Error>> {
3475 let generation = OptiForkGeneration {
3476 id: self.next,
3477 slot: (self.next & 1) as usize,
3478 };
3479 if let Some(live) = self.live[generation.slot] {
3480 return Err(format!(
3481 "optipipe snapshot slot {} still owns generation {live}; refusing to overwrite it",
3482 generation.slot,
3483 )
3484 .into());
3485 }
3486 self.next += 1;
3487 self.live[generation.slot] = Some(generation.id);
3488 Ok(generation)
3489 }
3490
3491 fn retire(&mut self, generation: OptiForkGeneration) -> Result<(), Box<dyn std::error::Error>> {
3492 match self.live[generation.slot] {
3493 Some(id) if id == generation.id => {
3494 self.live[generation.slot] = None;
3495 Ok(())
3496 }
3497 other => Err(format!(
3498 "optipipe generation teardown mismatch: ticket={} slot={} live={other:?}",
3499 generation.id, generation.slot,
3500 )
3501 .into()),
3502 }
3503 }
3504}
3505
3506struct OptiForkSeedGeneration {
3507 h_seed: CudaSlice<f32>,
3508 fill_prev: CudaSlice<f32>,
3509 scratch_len: usize,
3510}
3511
3512/// Allocate or refresh one full checkpoint through the engine that owns each PP stage. The
3513/// generic cache helper accepts one device and therefore cannot copy GDN state split across
3514/// devices. KV lengths and position stay host metadata; only recurrent buffers need stage-local
3515/// device ownership.
3516fn opti_snapshot_stage_owned(
3517 e: &Engine,
3518 cache: &Cache,
3519 rt: &'static crate::pp::PpNRt,
3520 fence: &[usize],
3521) -> Result<crate::cache::CacheSnapshot, Box<dyn std::error::Error>> {
3522 let n = cache.kv.len();
3523 let mut snapshot = crate::cache::CacheSnapshot {
3524 kv_len: vec![None; n],
3525 tp_kv_len: vec![None; n],
3526 conv: (0..n).map(|_| None).collect(),
3527 ssm: (0..n).map(|_| None).collect(),
3528 pos: cache.pos,
3529 };
3530 opti_snapshot_stage_owned_into(e, cache, rt, fence, &mut snapshot)?;
3531 Ok(snapshot)
3532}
3533
3534fn opti_snapshot_stage_owned_into(
3535 e: &Engine,
3536 cache: &Cache,
3537 rt: &'static crate::pp::PpNRt,
3538 fence: &[usize],
3539 snapshot: &mut crate::cache::CacheSnapshot,
3540) -> Result<(), Box<dyn std::error::Error>> {
3541 if fence.len() != rt.n_stages() + 1
3542 || snapshot.kv_len.len() != cache.kv.len()
3543 || snapshot.tp_kv_len.len() != cache.tp_kv.len()
3544 {
3545 return Err("optipipe stage-owned snapshot shape mismatch".into());
3546 }
3547 for stage in 0..rt.n_stages() {
3548 opti_snapshot_one_stage_owned_into(e, cache, rt, fence, stage, snapshot)?;
3549 }
3550 snapshot.pos = cache.pos;
3551 Ok(())
3552}
3553
3554/// Refresh one PP stage of a checkpoint. Increment 2 uses this split form so stage 0's
3555/// optimistic post-N state is captured before N+1 stage 0 is queued, while stage 1's matching
3556/// post-N state is captured only after N stage 1 is enqueued. Calling the all-stage helper at
3557/// either point would capture one side of the fork at the wrong generation.
3558fn opti_snapshot_one_stage_owned_into(
3559 e: &Engine,
3560 cache: &Cache,
3561 rt: &'static crate::pp::PpNRt,
3562 fence: &[usize],
3563 stage: usize,
3564 snapshot: &mut crate::cache::CacheSnapshot,
3565) -> Result<(), Box<dyn std::error::Error>> {
3566 if fence.len() != rt.n_stages() + 1
3567 || snapshot.kv_len.len() != cache.kv.len()
3568 || snapshot.tp_kv_len.len() != cache.tp_kv.len()
3569 || stage >= rt.n_stages()
3570 {
3571 return Err("optipipe single-stage snapshot shape mismatch".into());
3572 }
3573 let _scope = rt.enter(stage);
3574 let owner = rt.engine(stage, e);
3575 for il in fence[stage]..fence[stage + 1] {
3576 snapshot.kv_len[il] = cache.kv[il].as_ref().map(|kv| kv.len);
3577 snapshot.tp_kv_len[il] = cache.tp_kv[il]
3578 .as_ref()
3579 .map(crate::tp::ResidentTpKvCache::committed_len);
3580 match &cache.recur[il] {
3581 Some(recur) => {
3582 match snapshot.conv[il].as_mut() {
3583 Some(dst) => {
3584 owner.copy_into(dst, 0, &recur.conv_state, recur.conv_state.len())?
3585 }
3586 None => snapshot.conv[il] = Some(owner.clone_dtod(&recur.conv_state)?),
3587 }
3588 match snapshot.ssm[il].as_mut() {
3589 Some(dst) => {
3590 owner.copy_into(dst, 0, &recur.ssm_state, recur.ssm_state.len())?
3591 }
3592 None => snapshot.ssm[il] = Some(owner.clone_dtod(&recur.ssm_state)?),
3593 }
3594 }
3595 None if snapshot.conv[il].is_some() || snapshot.ssm[il].is_some() => {
3596 return Err(
3597 format!("optipipe stage-owned snapshot layer {il} changed shape").into(),
3598 );
3599 }
3600 None => {}
3601 }
3602 }
3603 snapshot.pos = cache.pos;
3604 Ok(())
3605}
3606
3607/// Increment-1 persistent fork state. Exactly two snapshot/seed slots alternate; a live ticket
3608/// names its generation and keeps teardown fail-closed. Only stage 0 is allowed to mutate before
3609/// resolve, so the reconcile tables and conditional restores are stage-local.
3610struct OptiForkState {
3611 mode: OptiForkGateMode,
3612 controller: Option<OptiControllerPolicy>,
3613 generations: OptiForkGenerationTracker,
3614 active_snapshot_slot: usize,
3615 alternate_snapshot: crate::cache::CacheSnapshot,
3616 seeds: [OptiForkSeedGeneration; 2],
3617 rt: &'static crate::pp::PpNRt,
3618 fence: [usize; 3],
3619 split: usize,
3620 len_ptrs: CudaSlice<u64>,
3621 saved_lens: CudaSlice<i32>,
3622 forced_acc: CudaSlice<u32>,
3623 valid: CudaSlice<u32>,
3624 stage0_stream: std::sync::Arc<cudarc::driver::CudaStream>,
3625 logical_payload_bytes: [usize; 2],
3626}
3627
3628struct OptiForkTicket {
3629 generation: OptiForkGeneration,
3630 boundary: Option<VerifyBoundaryTicket>,
3631 drain: std::sync::Arc<cudarc::driver::CudaStream>,
3632 settled: bool,
3633}
3634
3635struct OptiControllerTicket {
3636 generation: OptiForkGeneration,
3637 boundary: Option<VerifyBoundaryTicket>,
3638 ckpt: Option<VerifyCkpt>,
3639 verify_tokens: [u32; 2],
3640 draft_prob: f32,
3641 eager_seed: Option<CudaSlice<f32>>,
3642 q_proxy: f32,
3643 scratch_len: usize,
3644 issued_at: std::time::Instant,
3645 drain: std::sync::Arc<cudarc::driver::CudaStream>,
3646 settled: bool,
3647}
3648
3649struct OptiControllerPrepared {
3650 verify_tokens: [u32; 2],
3651 draft_prob: f32,
3652 eager_seed: Option<CudaSlice<f32>>,
3653 q_proxy: f32,
3654 scratch_len: usize,
3655}
3656
3657impl OptiControllerTicket {
3658 fn take_boundary(&mut self) -> VerifyBoundaryTicket {
3659 self.boundary
3660 .take()
3661 .expect("controller boundary ticket already consumed")
3662 }
3663
3664 fn take_ckpt(&mut self) -> VerifyCkpt {
3665 self.ckpt
3666 .take()
3667 .expect("controller verify checkpoint already consumed")
3668 }
3669
3670 fn take_eager_seed(&mut self) -> Option<CudaSlice<f32>> {
3671 self.eager_seed.take()
3672 }
3673
3674 fn settle(&mut self) {
3675 self.settled = true;
3676 }
3677}
3678
3679impl Drop for OptiControllerTicket {
3680 fn drop(&mut self) {
3681 if !self.settled {
3682 let _ = self.drain.synchronize();
3683 OPTI_FORK_ABORT_DRAINS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
3684 }
3685 }
3686}
3687
3688impl OptiForkTicket {
3689 fn take_boundary(&mut self) -> VerifyBoundaryTicket {
3690 self.boundary
3691 .take()
3692 .expect("fork ticket boundary already consumed")
3693 }
3694
3695 fn settle(&mut self) {
3696 self.settled = true;
3697 }
3698}
3699
3700impl Drop for OptiForkTicket {
3701 fn drop(&mut self) {
3702 if !self.settled {
3703 let _ = self.drain.synchronize();
3704 OPTI_FORK_ABORT_DRAINS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
3705 }
3706 }
3707}
3708
3709impl OptiForkState {
3710 #[allow(clippy::too_many_arguments)]
3711 fn new(
3712 e: &Engine,
3713 cache: &Cache,
3714 mode: OptiForkGateMode,
3715 alternate_snapshot: crate::cache::CacheSnapshot,
3716 h_seed: &CudaSlice<f32>,
3717 fill_prev: &CudaSlice<f32>,
3718 rt: &'static crate::pp::PpNRt,
3719 split: usize,
3720 n_layer: usize,
3721 ) -> Result<Self, Box<dyn std::error::Error>> {
3722 let fence = [0, split, n_layer];
3723 let mut logical_payload_bytes = [0usize; 2];
3724 for stage in 0..2 {
3725 for il in fence[stage]..fence[stage + 1] {
3726 logical_payload_bytes[stage] += alternate_snapshot.conv[il]
3727 .as_ref()
3728 .map_or(0, |v| v.len() * std::mem::size_of::<f32>());
3729 logical_payload_bytes[stage] += alternate_snapshot.ssm[il]
3730 .as_ref()
3731 .map_or(0, |v| v.len() * std::mem::size_of::<f32>());
3732 }
3733 }
3734 let seeds = [
3735 OptiForkSeedGeneration {
3736 h_seed: e.clone_dtod(h_seed)?,
3737 fill_prev: e.clone_dtod(fill_prev)?,
3738 scratch_len: 0,
3739 },
3740 OptiForkSeedGeneration {
3741 h_seed: e.clone_dtod(h_seed)?,
3742 fill_prev: e.clone_dtod(fill_prev)?,
3743 scratch_len: 0,
3744 },
3745 ];
3746 let (len_ptrs, saved_lens, forced_acc, valid, stage0_stream) = {
3747 let _stage = rt.enter(0);
3748 let e0 = rt.engine(0, e);
3749 (
3750 crate::round_stream::kv_len_ptr_table_range(e0, cache, 0..split, None)?,
3751 e0.htod_i32(&vec![0; split])?,
3752 e0.alloc_u32_zeroed(2)?,
3753 e0.alloc_u32_zeroed(1)?,
3754 e0.stream(),
3755 )
3756 };
3757 logical_payload_bytes[0] += seeds
3758 .iter()
3759 .map(|seed| (seed.h_seed.len() + seed.fill_prev.len()) * std::mem::size_of::<f32>())
3760 .sum::<usize>();
3761 logical_payload_bytes[0] += len_ptrs.len() * std::mem::size_of::<u64>()
3762 + saved_lens.len() * std::mem::size_of::<i32>()
3763 + forced_acc.len() * std::mem::size_of::<u32>()
3764 + valid.len() * std::mem::size_of::<u32>();
3765 Ok(Self {
3766 mode,
3767 controller: (mode == OptiForkGateMode::Controller)
3768 .then(OptiControllerPolicy::configured),
3769 generations: OptiForkGenerationTracker::default(),
3770 active_snapshot_slot: 0,
3771 alternate_snapshot,
3772 seeds,
3773 rt,
3774 fence,
3775 split,
3776 len_ptrs,
3777 saved_lens,
3778 forced_acc,
3779 valid,
3780 stage0_stream,
3781 logical_payload_bytes,
3782 })
3783 }
3784
3785 fn reserve(
3786 &mut self,
3787 current_snapshot: &mut crate::cache::CacheSnapshot,
3788 ) -> Result<OptiForkGeneration, Box<dyn std::error::Error>> {
3789 let generation = self.generations.reserve()?;
3790 if generation.slot != self.active_snapshot_slot {
3791 std::mem::swap(current_snapshot, &mut self.alternate_snapshot);
3792 self.active_snapshot_slot = generation.slot;
3793 }
3794 Ok(generation)
3795 }
3796
3797 fn capture_seed(
3798 &mut self,
3799 e: &Engine,
3800 generation: OptiForkGeneration,
3801 h_seed: &CudaSlice<f32>,
3802 fill_prev: &CudaSlice<f32>,
3803 scratch_len: usize,
3804 ) -> Result<(), Box<dyn std::error::Error>> {
3805 let seed = &mut self.seeds[generation.slot];
3806 e.copy_into(&mut seed.h_seed, 0, h_seed, h_seed.len())?;
3807 e.copy_into(&mut seed.fill_prev, 0, fill_prev, fill_prev.len())?;
3808 seed.scratch_len = scratch_len;
3809 Ok(())
3810 }
3811
3812 fn ticket(
3813 &self,
3814 generation: OptiForkGeneration,
3815 boundary: VerifyBoundaryTicket,
3816 ) -> OptiForkTicket {
3817 OptiForkTicket {
3818 generation,
3819 boundary: Some(boundary),
3820 drain: self.stage0_stream.clone(),
3821 settled: false,
3822 }
3823 }
3824
3825 #[allow(clippy::too_many_arguments)]
3826 fn controller_ticket(
3827 &self,
3828 generation: OptiForkGeneration,
3829 boundary: VerifyBoundaryTicket,
3830 ckpt: VerifyCkpt,
3831 verify_tokens: [u32; 2],
3832 draft_prob: f32,
3833 eager_seed: Option<CudaSlice<f32>>,
3834 q_proxy: f32,
3835 scratch_len: usize,
3836 ) -> OptiControllerTicket {
3837 OptiControllerTicket {
3838 generation,
3839 boundary: Some(boundary),
3840 ckpt: Some(ckpt),
3841 verify_tokens,
3842 draft_prob,
3843 eager_seed,
3844 q_proxy,
3845 scratch_len,
3846 issued_at: std::time::Instant::now(),
3847 drain: self.stage0_stream.clone(),
3848 settled: false,
3849 }
3850 }
3851
3852 fn reserve_successor(&mut self) -> Result<OptiForkGeneration, Box<dyn std::error::Error>> {
3853 self.generations.reserve()
3854 }
3855
3856 fn successor_snapshot_mut(&mut self) -> &mut crate::cache::CacheSnapshot {
3857 &mut self.alternate_snapshot
3858 }
3859
3860 fn promote_successor_snapshot(
3861 &mut self,
3862 current_snapshot: &mut crate::cache::CacheSnapshot,
3863 generation: OptiForkGeneration,
3864 ) {
3865 std::mem::swap(current_snapshot, &mut self.alternate_snapshot);
3866 self.active_snapshot_slot = generation.slot;
3867 }
3868
3869 fn queue_actual_reconcile(
3870 &mut self,
3871 e: &Engine,
3872 snapshot: &crate::cache::CacheSnapshot,
3873 acc: &CudaSlice<u32>,
3874 optimistic_pending: u32,
3875 base: usize,
3876 ) -> Result<(), Box<dyn std::error::Error>> {
3877 let saved: Vec<i32> = (0..self.split)
3878 .map(|il| snapshot.kv_len[il].map(|v| v as i32).unwrap_or(0))
3879 .collect();
3880 // Serving keeps the caller/accept walk on the head (stage-1) device. Record the accept
3881 // decision point there and append a wait to stage 0 after its optimistic successor/TX;
3882 // the validity/reconcile kernels must never peer-read acc before it is written. The
3883 // increment-1 harness uses primary stage 0, where stream order already provides this.
3884 if self.rt.engine(0, e).ctx().ordinal() != e.ctx().ordinal() {
3885 self.rt.fence_stages_behind(&e.stream())?;
3886 }
3887 let _stage = self.rt.enter(0);
3888 let e0 = self.rt.engine(0, e);
3889 e0.htod_i32_into(&mut self.saved_lens, &saved)?;
3890 e0.spec_fork_valid(acc, optimistic_pending, &mut self.valid)?;
3891 e0.spec_fork_reconcile_kv(
3892 &self.len_ptrs,
3893 &self.saved_lens,
3894 acc,
3895 &self.valid,
3896 base,
3897 self.split,
3898 )
3899 }
3900
3901 fn finish_actual_reconcile(
3902 &mut self,
3903 e: &Engine,
3904 cache: &mut Cache,
3905 snapshot: &crate::cache::CacheSnapshot,
3906 n_acc: usize,
3907 base: usize,
3908 hit: bool,
3909 ) -> Result<(), Box<dyn std::error::Error>> {
3910 if hit {
3911 return Ok(());
3912 }
3913 let len_delta = base + n_acc;
3914 for il in 0..self.split {
3915 if let (Some(kv), Some(saved)) = (cache.kv[il].as_mut(), snapshot.kv_len[il]) {
3916 kv.len = saved + len_delta;
3917 }
3918 }
3919 {
3920 let _stage = self.rt.enter(1);
3921 let e1 = self.rt.engine(1, e);
3922 for il in self.split..self.fence[2] {
3923 if let (Some(kv), Some(saved)) = (cache.kv[il].as_mut(), snapshot.kv_len[il]) {
3924 kv.len = saved + len_delta;
3925 e1.set_i32_one(&mut kv.len_d, kv.len as i32)?;
3926 }
3927 }
3928 }
3929 self.rt.publish_to(0, &e.stream())?;
3930 Ok(())
3931 }
3932
3933 fn cancel_controller_ticket(
3934 &mut self,
3935 e: &Engine,
3936 cache: &mut Cache,
3937 scratch: &mut MtpScratch,
3938 snapshot: &crate::cache::CacheSnapshot,
3939 ticket: &mut OptiControllerTicket,
3940 ) -> Result<(), Box<dyn std::error::Error>> {
3941 {
3942 let _stage = self.rt.enter(0);
3943 let e0 = self.rt.engine(0, e);
3944 for il in 0..self.split {
3945 if let (Some(kv), Some(saved)) = (cache.kv[il].as_mut(), snapshot.kv_len[il]) {
3946 kv.len = saved;
3947 e0.set_i32_one(&mut kv.len_d, saved as i32)?;
3948 }
3949 }
3950 }
3951 scratch.set_len(e, snapshot.pos)?;
3952 ticket.settle();
3953 self.generations.retire(ticket.generation)?;
3954 OPTI_FORK_ABORT_DRAINS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
3955 OPTI_WASTED_DRAFT_TOKENS.fetch_add(2, std::sync::atomic::Ordering::Relaxed);
3956 eprintln!(
3957 "[opti-controller] tail-drain generation={} slot={}",
3958 ticket.generation.id, ticket.generation.slot,
3959 );
3960 Ok(())
3961 }
3962
3963 #[allow(clippy::too_many_arguments)]
3964 fn reconcile(
3965 &mut self,
3966 e: &Engine,
3967 cache: &mut Cache,
3968 scratch: &mut MtpScratch,
3969 snapshot: &crate::cache::CacheSnapshot,
3970 h_seed: &mut CudaSlice<f32>,
3971 fill_prev: &mut CudaSlice<f32>,
3972 generation: OptiForkGeneration,
3973 action: OptiForkAction,
3974 optimistic_pending: u32,
3975 ) -> Result<(), Box<dyn std::error::Error>> {
3976 debug_assert!(action != OptiForkAction::Abort);
3977 let miss_started = std::time::Instant::now();
3978 let keep = action == OptiForkAction::Hit;
3979 let saved: Vec<i32> = (0..self.split)
3980 .map(|il| snapshot.kv_len[il].map(|v| v as i32).unwrap_or(0))
3981 .collect();
3982 let seed = &self.seeds[generation.slot];
3983 {
3984 let _stage = self.rt.enter(0);
3985 let e0 = self.rt.engine(0, e);
3986 e0.htod_i32_into(&mut self.saved_lens, &saved)?;
3987 let forced = if keep {
3988 [1u32, optimistic_pending]
3989 } else {
3990 [0u32, optimistic_pending]
3991 };
3992 e0.htod_u32_into(&mut self.forced_acc, &forced)?;
3993 e0.spec_fork_valid(&self.forced_acc, optimistic_pending, &mut self.valid)?;
3994 e0.spec_fork_reconcile_kv(
3995 &self.len_ptrs,
3996 &self.saved_lens,
3997 &self.forced_acc,
3998 &self.valid,
3999 0,
4000 self.split,
4001 )?;
4002 for il in 0..self.split {
4003 if let Some(recur) = cache.recur[il].as_mut() {
4004 let conv = snapshot.conv[il]
4005 .as_ref()
4006 .ok_or("optipipe stage0 snapshot missing conv state")?;
4007 let ssm = snapshot.ssm[il]
4008 .as_ref()
4009 .ok_or("optipipe stage0 snapshot missing ssm state")?;
4010 e0.spec_fork_restore_f32(conv, &mut recur.conv_state, &self.valid)?;
4011 e0.spec_fork_restore_f32(ssm, &mut recur.ssm_state, &self.valid)?;
4012 }
4013 }
4014 e0.spec_fork_restore_f32(&seed.h_seed, h_seed, &self.valid)?;
4015 e0.spec_fork_restore_f32(&seed.fill_prev, fill_prev, &self.valid)?;
4016 }
4017
4018 if keep {
4019 OPTI_FORK_HITS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
4020 return Ok(());
4021 }
4022
4023 for il in 0..self.split {
4024 if let (Some(kv), Some(saved)) = (cache.kv[il].as_mut(), snapshot.kv_len[il]) {
4025 kv.len = saved;
4026 }
4027 }
4028 scratch.set_len(e, seed.scratch_len)?;
4029 // Targeted E_restart: publish only stage 0's reconcile to the caller, then bound the
4030 // forced diagnostic so the retained number is the actual miss cost, not enqueue time.
4031 let caller = e.stream();
4032 self.rt.publish_to(0, &caller)?;
4033 caller.synchronize()?;
4034 let miss_ms = miss_started.elapsed().as_secs_f64() * 1e3;
4035 eprintln!(
4036 "[opti-fork-reconcile] generation={} slot={} miss_ms={miss_ms:.3}",
4037 generation.id, generation.slot,
4038 );
4039 OPTI_FORK_MISSES.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
4040 Ok(())
4041 }
4042
4043 fn retire(&mut self, generation: OptiForkGeneration) -> Result<(), Box<dyn std::error::Error>> {
4044 self.generations.retire(generation)
4045 }
4046}
4047
4048fn rewind_tp_kv_verified_prefix(
4049 tp_kv: &mut [Option<crate::tp::ResidentTpKvCache>],
4050 saved_lens: &[Option<usize>],
4051 accepted: usize,
4052) -> Result<(), Box<dyn std::error::Error>> {
4053 if tp_kv.len() != saved_lens.len() {
4054 return Err("spec TP KV snapshot shape mismatch".into());
4055 }
4056 for (layer, (cache, saved)) in tp_kv.iter_mut().zip(saved_lens).enumerate() {
4057 match (cache.as_mut(), *saved) {
4058 (Some(cache), Some(saved)) => {
4059 let target = saved
4060 .checked_add(accepted)
4061 .ok_or("spec TP KV committed length overflow")?;
4062 cache.rewind_to(target)?;
4063 }
4064 (None, None) => {}
4065 _ => {
4066 return Err(
4067 format!("spec TP KV layer {layer} changed shape since its snapshot").into(),
4068 );
4069 }
4070 }
4071 }
4072 Ok(())
4073}
4074
4075/// MEMRA_SPEC_ROUND_PROF counters: whole-round wall, so the round can be weighed against the
4076/// draft-step ([spec-anatomy]) and verify-walk ([tcol-prof]) splits we already print.
4077static ROUND_PROF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
4078static ROUND_MS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
4079static ROUND_N: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
4080
4081impl HybridModel {
4082 fn mtp_head_count(&self) -> usize {
4083 usize::from(self.mtp.is_some()) + self.mtp_extra.len()
4084 }
4085
4086 fn mtp_head_at(&self, index: usize) -> &MtpHead {
4087 if index == 0 {
4088 self.mtp.as_ref().expect("MTP head 0 is unavailable")
4089 } else {
4090 &self.mtp_extra[index - 1]
4091 }
4092 }
4093
4094 fn new_mtp_scratch(
4095 &self,
4096 e: &Engine,
4097 cap: usize,
4098 ) -> Result<MtpScratch, Box<dyn std::error::Error>> {
4099 let mut scratch = MtpScratch::new(
4100 e,
4101 &self.cfg,
4102 &self.plan,
4103 cap,
4104 self.mtp.as_ref().and_then(|head| head.geom.as_ref()),
4105 )?;
4106 for head in &self.mtp_extra {
4107 scratch.push_plane(e, &self.cfg, &self.plan, head.geom.as_ref())?;
4108 }
4109 Ok(scratch)
4110 }
4111
4112 fn opti_graph_draft_step(
4113 &self,
4114 e: &Engine,
4115 mtp: &MtpHead,
4116 dctx: &mut DraftGraphCtx,
4117 scratch: &mut MtpScratch,
4118 d_vocab: usize,
4119 ) -> Result<(u32, f32), Box<dyn std::error::Error>> {
4120 // dcw door: one replay appends one device-counter row; pre-arm ring headroom
4121 // host-side before launching (no-op on flat planes).
4122 if step35_draft_dcw_on() {
4123 scratch.ensure_dcw_headroom(e, 2)?;
4124 }
4125 dctx.graph
4126 .as_ref()
4127 .ok_or("optipipe controller requires the greedy draft graph")?
4128 .launch()?;
4129 scratch.kv.len += 1;
4130 let idx = e.dtoh_u32_one(&dctx.g_tok)?;
4131 if (idx as usize) >= d_vocab {
4132 return Err(
4133 format!("optipipe draft argmax sentinel 0x{idx:08x} >= d_vocab {d_vocab}").into(),
4134 );
4135 }
4136 let probability = e.dtoh(&dctx.g_p)?[0];
4137 if !probability.is_finite() || !(0.0..=1.0).contains(&probability) {
4138 return Err(format!("optipipe draft probability is invalid: {probability}").into());
4139 }
4140 let token = match &mtp.d2t {
4141 Some(map) => map[idx as usize],
4142 None => idx,
4143 };
4144 if token != idx {
4145 e.set_u32_one(&mut dctx.g_tok, token)?;
4146 }
4147 Ok((token, probability))
4148 }
4149
4150 #[allow(clippy::too_many_arguments)]
4151 fn opti_controller_draft_step(
4152 &self,
4153 e: &Engine,
4154 mtp: &MtpHead,
4155 dctx: &mut DraftGraphCtx,
4156 scratch: &mut MtpScratch,
4157 d_vocab: usize,
4158 eager_state: &mut Option<(u32, CudaSlice<f32>)>,
4159 eager_pos: usize,
4160 embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
4161 ) -> Result<(u32, f32), Box<dyn std::error::Error>> {
4162 if dctx.graph.is_some() {
4163 return self.opti_graph_draft_step(e, mtp, dctx, scratch, d_vocab);
4164 }
4165 let (input_token, input_seed) = eager_state
4166 .take()
4167 .ok_or("optipipe eager continuation seed is unavailable")?;
4168 let (logits, next_seed) = self.mtp_head_forward_dev(
4169 e,
4170 mtp,
4171 input_token,
4172 &input_seed,
4173 scratch,
4174 eager_pos,
4175 embd_dev,
4176 None,
4177 )?;
4178 let token_d = e.argmax_token_device(&logits, d_vocab)?;
4179 let idx = e.dtoh_u32_one(&token_d)?;
4180 if (idx as usize) >= d_vocab {
4181 return Err(format!(
4182 "optipipe eager draft argmax sentinel 0x{idx:08x} >= d_vocab {d_vocab}"
4183 )
4184 .into());
4185 }
4186 let probability_d = e.prob_of_token_device(&logits, &token_d, d_vocab)?;
4187 let probability = e.dtoh(&probability_d)?[0];
4188 if !probability.is_finite() || !(0.0..=1.0).contains(&probability) {
4189 return Err(
4190 format!("optipipe eager draft probability is invalid: {probability}").into(),
4191 );
4192 }
4193 let token = match &mtp.d2t {
4194 Some(map) => map[idx as usize],
4195 None => idx,
4196 };
4197 *eager_state = Some((token, next_seed));
4198 Ok((token, probability))
4199 }
4200
4201 /// NextN head forward for ONE draft token (§A ops 1-13, T=1).
4202 /// Inputs: `e_tok` = the token to predict FROM (last committed / previous draft); `h_seed` =
4203 /// the trunk's pre-output_norm hidden of that token (§A op 2 input). `mtp_pos` = absolute
4204 /// position of the token being predicted from. Returns (draft_logits[n_vocab] host, h_nextn dev).
4205 /// `h_nextn` (§A op 10) becomes `h_seed` for the next autoregressive draft step.
4206 /// Device-resident: returns draft logits ON DEVICE (no [n_vocab] dtoh). The greedy draft
4207 /// loop only needs argmax — paired with `argmax_token_device` this cuts the ~600KB logits
4208 /// transfer + host argmax per draft token from the K-token draft chain.
4209 #[allow(clippy::too_many_arguments)]
4210 fn mtp_head_forward_dev(
4211 &self,
4212 e: &Engine,
4213 mtp: &MtpHead,
4214 e_tok: u32,
4215 h_seed: &CudaSlice<f32>,
4216 scratch: &mut MtpScratch,
4217 mtp_pos: usize,
4218 embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
4219 mask: Option<(&CudaSlice<u32>, usize)>,
4220 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4221 self.mtp_head_forward_dev_at(e, mtp, e_tok, h_seed, scratch, 0, mtp_pos, embd_dev, mask)
4222 }
4223
4224 #[allow(clippy::too_many_arguments)]
4225 fn mtp_head_forward_dev_at(
4226 &self,
4227 e: &Engine,
4228 mtp: &MtpHead,
4229 e_tok: u32,
4230 h_seed: &CudaSlice<f32>,
4231 scratch: &mut MtpScratch,
4232 scratch_index: usize,
4233 mtp_pos: usize,
4234 embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
4235 // DRAFT-SIDE GRAMMAR MASK (lane/draft-mask): (packed draft-vocab allowed set, words).
4236 // Applied to the head logits BEFORE they are returned, so every consumer (argmax,
4237 // gumbel draw, p-min prob) sees the grammar-legal row. None = unmasked (pre-lane).
4238 mask: Option<(&CudaSlice<u32>, usize)>,
4239 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4240 // MEMRA_SPEC_ANATOMY=1 — eager-step phase timers (diagnostic only). Phase boundaries
4241 // sync the stream, so absolute time inflates; the BREAKDOWN is the signal. Cumulative
4242 // summary on stderr every 128 steps: glue (embed..attn_norm), attn, ffn, head.
4243 use std::sync::atomic::{AtomicU64, Ordering::Relaxed};
4244 static ANAT_NS: [AtomicU64; 5] = [
4245 AtomicU64::new(0),
4246 AtomicU64::new(0),
4247 AtomicU64::new(0),
4248 AtomicU64::new(0),
4249 AtomicU64::new(0),
4250 ];
4251 static ANAT_STEPS: AtomicU64 = AtomicU64::new(0);
4252 let anat = {
4253 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
4254 *ON.get_or_init(|| std::env::var("MEMRA_SPEC_ANATOMY").as_deref() == Ok("1"))
4255 };
4256 if anat {
4257 e.stream().synchronize()?; // drain prior queue so phase 0 starts clean
4258 }
4259 let t_all = std::time::Instant::now();
4260 let mut t_ph = std::time::Instant::now();
4261 let mut anat_mark = |i: usize,
4262 e: &Engine,
4263 t: &mut std::time::Instant|
4264 -> Result<(), Box<dyn std::error::Error>> {
4265 if anat {
4266 e.stream().synchronize()?;
4267 ANAT_NS[i].fetch_add(t.elapsed().as_nanos() as u64, Relaxed);
4268 *t = std::time::Instant::now();
4269 }
4270 Ok(())
4271 };
4272 let cfg = &self.cfg;
4273 let n_embd = cfg.n_embd as usize;
4274 // Distilled-student geometry: the block runs at the INNER width `di` (eh_proj out /
4275 // attn / ffn); the n_embd interface (embed, norms in, carrier out, head in) is unchanged.
4276 let di = mtp.geom.as_ref().map(|g| g.d_inner).unwrap_or(n_embd);
4277 let eps = cfg.rms_eps;
4278 let pos_d = e.htod_i32(&[mtp_pos as i32])?;
4279
4280 // op A: a resident table transfers one 4B token id. The exact host-row capacity path
4281 // expands this one row on CPU and transfers n_embd f32 values instead.
4282 let e_emb = match embd_dev {
4283 Some((g, qt, rb)) => e.embed_gather_device_t(g, &[e_tok], n_embd, qt, rb)?,
4284 None => e.htod(&self.embd.gather(n_embd, &[e_tok]))?,
4285 };
4286
4287 // op 1/2: e_norm = RMSNorm(e, enorm); h_norm = RMSNorm(h_seed, hnorm)
4288 let mut e_norm = e.zeros(n_embd)?;
4289 e.rms_norm(&e_emb, mtp.enorm.float_data(), &mut e_norm, n_embd, 1, eps)?;
4290 let mut h_norm = e.zeros(n_embd)?;
4291 e.rms_norm(h_seed, mtp.hnorm.float_data(), &mut h_norm, n_embd, 1, eps)?;
4292
4293 // op 3: concat = [e_norm ; h_norm] -> [2*n_embd], e_norm in [0,n_embd), h_norm in [n_embd,2n_embd)
4294 let mut concat = e.zeros(2 * n_embd)?;
4295 e.copy_into(&mut concat, 0, &e_norm, n_embd)?;
4296 e.copy_into(&mut concat, n_embd, &h_norm, n_embd)?;
4297
4298 // op 4: inpSA = eh_proj @ concat (eh_proj [2*n_embd, n_embd]) -> [n_embd]
4299 let inp_sa = e.matmul(&mtp.eh_proj, &concat, 1)?;
4300
4301 // op 5: a_norm = RMSNorm(inpSA, attn_norm)
4302 let mut a_norm = e.zeros(di)?;
4303 e.rms_norm(&inp_sa, mtp.attn_norm.float_data(), &mut a_norm, di, 1, eps)?;
4304 anat_mark(0, e, &mut t_ph)?;
4305
4306 // op 6: attention on the scratch KV. SAME dc launcher as the graph path (bucket_max =
4307 // scratch.cap, length from the device len_d) so eager drafts match graph drafts
4308 // bit-for-bit at any t_kv (the parity gate). Host len mirrored here (the dc append
4309 // advances only the device counter).
4310 let attn_out = match (&mtp.mixer, mtp.step35.as_ref()) {
4311 // step35 MTP block, dcw door armed: the SAME windowed device-counter launcher as
4312 // the captured chain (draft parity by construction). Per-step ring headroom runs
4313 // HERE (eager is host-len work, a rebase is legal); host len mirrored like the
4314 // plain dc arm below.
4315 (Mixer::Full(fa), Some(g))
4316 if self.step35_dcw_eligible(g, scratch.plane(scratch_index).1) =>
4317 {
4318 {
4319 let (kv, _) = scratch.plane_mut(scratch_index);
4320 let retain = match kv.ring.as_ref() {
4321 Some(ring) => memra_kv::swa_retain_from(kv.len, ring.window(), ring.base()),
4322 None => 0,
4323 };
4324 e.prepare_kv_append(kv, retain, 1)?;
4325 }
4326 let out =
4327 self.mtp_step35_attn_dcw(e, fa, g, &a_norm, &pos_d, scratch, scratch_index)?;
4328 scratch.plane_mut(scratch_index).0.len += 1;
4329 out
4330 }
4331 // step35 MTP block, door off (MEMRA_STEP35_DRAFT_DCW=0 rollback) or class-
4332 // ineligible: PER-LAYER geometry + a separate head-wise gate + an SWA window,
4333 // none of which the plain dc launcher can express (see `mtp_step35_attn`).
4334 // Host-len arm. Advances BOTH the
4335 // host len and the device counter itself (unlike the dc arm, whose host-side
4336 // mirror the caller does).
4337 (Mixer::Full(fa), Some(g)) => {
4338 self.mtp_step35_attn(e, fa, g, &a_norm, &pos_d, scratch, scratch_index)?
4339 }
4340 (Mixer::Full(fa), None) => {
4341 let out = self.mtp_full_attn_dc(
4342 e,
4343 fa,
4344 &a_norm,
4345 &pos_d,
4346 scratch,
4347 scratch_index,
4348 mtp.geom.as_ref(),
4349 )?;
4350 scratch.plane_mut(scratch_index).0.len += 1;
4351 out
4352 }
4353 (Mixer::Linear(_), _) => {
4354 panic!("MTP block is full-attn in qwen35; linear MTP not supported")
4355 }
4356 (Mixer::Mla(_), _) => crate::hybrid::mla_forward_unimplemented(),
4357 };
4358 anat_mark(1, e, &mut t_ph)?;
4359
4360 // op 7: x1 = inpSA + attn_out
4361 let mut x1 = e.zeros(di)?;
4362 e.add(&inp_sa, &attn_out, &mut x1, di)?;
4363
4364 // op 8: z = RMSNorm(x1, post_attn_norm) (pre-FFN norm)
4365 let mut z = e.zeros(di)?;
4366 e.rms_norm(&x1, mtp.post_attn_norm.float_data(), &mut z, di, 1, eps)?;
4367
4368 // op 9: FFN (Dense or MoE) — same as the trunk decode FFN
4369 let ffn_out = match &mtp.ffn {
4370 crate::hybrid::Ffn::Dense {
4371 ffn_gate,
4372 ffn_up,
4373 ffn_down,
4374 } => {
4375 let n_ff = ffn_gate.out_features();
4376 let (gate, up) = if e.uses_q8_1_fast(ffn_gate) && e.uses_q8_1_fast(ffn_up) {
4377 let (zq, zd) = e.quantize_q8_1(&z, 1, di)?;
4378 (
4379 e.matmul_pre(ffn_gate, &zq, &zd, &z, 1)?,
4380 e.matmul_pre(ffn_up, &zq, &zd, &z, 1)?,
4381 )
4382 } else {
4383 (e.matmul(ffn_gate, &z, 1)?, e.matmul(ffn_up, &z, 1)?)
4384 };
4385 let mut act = e.zeros(n_ff)?;
4386 // step35: a DENSE FFN reads the per-layer SHEXP clamp (upstream's one `build_ffn`
4387 // serves the dense MLP and the shared expert off `swiglu_clamp_shexp` —
4388 // llama-graph.cpp:1751), resolved for the MTP block's OWN index. Every other arch
4389 // passes None, which is `ffn_act`'s dispatch verbatim.
4390 Self::ffn_act_lim(
4391 e,
4392 &self.cfg,
4393 &gate,
4394 &up,
4395 1.0,
4396 1.0,
4397 mtp.step35.as_ref().and_then(|s| s.clamp_shexp),
4398 &mut act,
4399 n_ff,
4400 )?;
4401 e.matmul(ffn_down, &act, 1)?
4402 }
4403 // MTP head is a distinct block — key its experts under a separate layer index (u16::MAX)
4404 // so they never alias trunk layer 0's cache keys.
4405 crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il(e, m, &z, 1, u16::MAX)?,
4406 };
4407 anat_mark(2, e, &mut t_ph)?;
4408
4409 // op 10: h_nextn = x1 + ffn_out (at di)
4410 let mut h_inner = e.zeros(di)?;
4411 e.add(&x1, &ffn_out, &mut h_inner, di)?;
4412
4413 // op 10.5 (student): up-project the inner hidden back to n_embd — training semantics:
4414 // the chain carrier AND the head input are out_up(h_inner) (pre-final-norm).
4415 let h_nextn = match mtp.geom.as_ref() {
4416 Some(g) => e.matmul(&g.out_up, &h_inner, 1)?,
4417 None => h_inner,
4418 };
4419
4420 // op 11: final = RMSNorm(h_nextn, shared_head_norm OR output_norm)
4421 let final_norm = mtp.shared_head_norm.as_ref().unwrap_or(&self.output_norm);
4422 let mut final_h = e.zeros(n_embd)?;
4423 e.rms_norm(
4424 &h_nextn,
4425 final_norm.float_data(),
4426 &mut final_h,
4427 n_embd,
4428 1,
4429 eps,
4430 )?;
4431
4432 // op 12: draft_logits = (shared_head_head OR output) @ final — stays ON DEVICE.
4433 let head = mtp.shared_head_head.as_ref().unwrap_or(&self.output);
4434 let mut logits = e.matmul(head, &final_h, 1)?;
4435 // op 12b (lane/draft-mask): grammar mask over the DRAFT vocab, applied here so the
4436 // caller's argmax / gumbel draw / p-min prob all read the grammar-legal row.
4437 if let Some((mask_d, mw)) = mask {
4438 let d_vocab = head.out_features();
4439 e.mask_logits_col(&mut logits, mask_d, 0, d_vocab, mw)?;
4440 }
4441 anat_mark(3, e, &mut t_ph)?;
4442 if anat {
4443 ANAT_NS[4].fetch_add(t_all.elapsed().as_nanos() as u64, Relaxed);
4444 let n = ANAT_STEPS.fetch_add(1, Relaxed) + 1;
4445 if n % 128 == 0 {
4446 let us = |i: usize| ANAT_NS[i].load(Relaxed) / n / 1000;
4447 eprintln!(
4448 "[spec-anatomy] steps={n} avg us/step: glue={} attn={} ffn={} head={} total={}",
4449 us(0),
4450 us(1),
4451 us(2),
4452 us(3),
4453 us(4)
4454 );
4455 }
4456 }
4457 // Chain recurrence hand-over: pre-norm h_nextn (default) or post-norm final_h
4458 // (MEMRA_SPEC_HPOST — llama.cpp #24025's t_h_nextn is taken AFTER the head norm).
4459 Ok((logits, if spec_hpost() { final_h } else { h_nextn }))
4460 }
4461
4462 #[allow(clippy::too_many_arguments)]
4463 fn mtp_chain_forward_dev(
4464 &self,
4465 e: &Engine,
4466 tokens: &[u32],
4467 seeds: &[CudaSlice<f32>],
4468 scratch: &mut MtpScratch,
4469 committed_scratch_len: usize,
4470 embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
4471 mask: Option<(&CudaSlice<u32>, usize)>,
4472 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4473 if tokens.is_empty() || tokens.len() != seeds.len() {
4474 return Err("multi-head MTP prefix tokens/seeds are malformed".into());
4475 }
4476 let index = mtp_chain_head_index(tokens.len() - 1, self.mtp_head_count());
4477 let head = self.mtp_head_at(index);
4478 scratch.set_plane_len(e, index, committed_scratch_len)?;
4479
4480 let mut last = None;
4481 for row in 0..tokens.len() {
4482 let is_last = row + 1 == tokens.len();
4483 last = Some(self.mtp_head_forward_dev_at(
4484 e,
4485 head,
4486 tokens[row],
4487 &seeds[row],
4488 scratch,
4489 index,
4490 committed_scratch_len + row + 1,
4491 embd_dev,
4492 if is_last { mask } else { None },
4493 )?);
4494 }
4495 Ok(last.expect("non-empty MTP prefix produced no row"))
4496 }
4497
4498 /// step35 MTP-block attention, T=1, on the scratch KV — the EAGER-ONLY twin of
4499 /// `mtp_full_attn_dc`. Three things force a separate arm rather than a geometry parameter on
4500 /// the dc path, and all three are properties of this arch's MTP block:
4501 ///
4502 /// 1. **The SWA window.** Block 45 is an SWA-type block (`sliding_window_pattern[45]=true`,
4503 /// window 512). Windowed decode in memra is a token-aligned VIEW OFFSET into the quantized
4504 /// cache (the gemma4 R6 / `step35_decode_attn` pattern: keys carry absolute rope and the
4505 /// mask is purely positional, so one query at `len-1` attending the last `win` rows IS the
4506 /// windowed result). `fa_decode_dc` takes the key count from a DEVICE counter and always
4507 /// starts at row 0 — it cannot express a nonzero offset. The windowed dc arm is
4508 /// `mtp_step35_attn_dcw` (`fa_decode_dcw`, doored via MEMRA_STEP35_DRAFT_DCW —
4509 /// default ON since lane/step37-draft-graph-serving-20260830); this host-len arm is
4510 /// the =0 rollback and the class-ineligibility fallback.
4511 /// 2. **Per-layer head count.** 96 q heads over 8 KV (GQA 12) at this block, vs the trunk's 64
4512 /// on its full-attn layers. The trunk cfg's `n_head` scalar is the MAX over layers, and the
4513 /// trunk ARTIFACT's per-layer arrays stop at index 44 — so the count must come from the
4514 /// resolved `Step35MtpGeom`, never from `cfg`.
4515 /// 3. **The separate head-wise gate.** `blk.45.attn_gate.weight [n_embd, 96]` produces one
4516 /// sigmoid scalar per head (broadcast over head_dim) — `attn_head_gate`, not the qwen35
4517 /// fused-into-wq `q_gate_split` form the dc arm handles.
4518 ///
4519 /// DOOR STATE: with MEMRA_STEP35_DRAFT_DCW=0 (or a sub-eligible kernel class),
4520 /// `mtp_head_forward_cap` refuses step35 heads explicitly (rather than silently capturing
4521 /// a window-less, wrong-past-`win` graph) and this eager chain IS the served path. With
4522 /// the door armed (the default), BOTH draft modes run the `mtp_step35_attn_dcw` twin
4523 /// instead of this arm.
4524 ///
4525 /// Unlike the dc arm this advances BOTH the host `kv.len` and the device counter, so the
4526 /// caller must not mirror.
4527 fn mtp_step35_attn(
4528 &self,
4529 e: &Engine,
4530 fa: &FullAttnLayer,
4531 g: &crate::hybrid::Step35MtpGeom,
4532 h: &CudaSlice<f32>,
4533 pos_d: &CudaSlice<i32>,
4534 scratch: &mut MtpScratch,
4535 scratch_index: usize,
4536 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4537 let (nh, nkv, hd) = (g.n_head, g.n_head_kv, self.cfg.head_dim_k as usize);
4538 // MTP-GEOM RECEIPT, once per process, on the SERVED draft path. Slot-0 acceptance is
4539 // 0.725 here against 0.994 for vLLM MTP3 on the same checkpoint family and card class, and
4540 // the first three explanations for that gap were all wrong: head assignment (step-modulo
4541 // is index 0 at K=1, correct), MEMRA_SPEC_HPOST (identical 84/116 both arms), and this
4542 // block's geometry. Geometry was the one that could have failed SILENTLY — a wrong window
4543 // makes the draft attend the whole context instead of Step-3.7's 512, stays fluent, and
4544 // shows up only as acceptance — so it gets a standing receipt rather than another reading
4545 // of the source. Prints the resolved Step35MtpGeom the served path actually runs on;
4546 // `full_attention_geometry_at`'s missing-row fallback (window: None) does NOT reach here.
4547 {
4548 static ONCE: std::sync::OnceLock<()> = std::sync::OnceLock::new();
4549 ONCE.get_or_init(|| {
4550 eprintln!(
4551 "[mtp-geom] arm=eager block={} swa={} window={} n_head={nh} n_head_kv={nkv} \
4552 head_dim_k={hd} n_rot={} rope_base={} clamp_shexp={:?}",
4553 g.il, g.swa, g.window, g.n_rot, g.rope_base, g.clamp_shexp,
4554 );
4555 });
4556 }
4557 let eps = self.cfg.rms_eps;
4558 let scale = 1.0 / (hd as f32).sqrt(); // step35.cpp:255 kq_scale
4559 let n_embd = self.cfg.n_embd as usize;
4560 let gw = fa
4561 .attn_gate
4562 .as_ref()
4563 .ok_or("step35 MTP block is missing attn_gate.weight (head-wise attention gate)")?;
4564
4565 let (q0, k0, v0, gt) = if e.uses_q8_1_fast(&fa.wq)
4566 && e.uses_q8_1_fast(&fa.wk)
4567 && e.uses_q8_1_fast(&fa.wv)
4568 && e.uses_q8_1_fast(gw)
4569 {
4570 let (hq, hdq) = e.quantize_q8_1(h, 1, n_embd)?;
4571 let (a, b, c) = match e.matmul_q8_fused3(&fa.wq, &fa.wk, &fa.wv, &hq, &hdq)? {
4572 Some(t3) => t3,
4573 None => (
4574 e.matmul_pre(&fa.wq, &hq, &hdq, h, 1)?,
4575 e.matmul_pre(&fa.wk, &hq, &hdq, h, 1)?,
4576 e.matmul_pre(&fa.wv, &hq, &hdq, h, 1)?,
4577 ),
4578 };
4579 (a, b, c, e.matmul_pre(gw, &hq, &hdq, h, 1)?)
4580 } else {
4581 (
4582 e.matmul(&fa.wq, h, 1)?,
4583 e.matmul(&fa.wk, h, 1)?,
4584 e.matmul(&fa.wv, h, 1)?,
4585 e.matmul(gw, h, 1)?,
4586 )
4587 };
4588
4589 let mut q = e.uninit(nh * hd)?;
4590 e.rms_norm(&q0, fa.q_norm.float_data(), &mut q, hd, nh, eps)?;
4591 let mut k = e.uninit(nkv * hd)?;
4592 e.rms_norm(&k0, fa.k_norm.float_data(), &mut k, hd, nkv, eps)?;
4593 // `rope_freqs.weight` (llama3 factors) applies to the FULL-attn layers ONLY; SWA passes
4594 // null (llama-hparams / step35.cpp). Block 45 is SWA, so `ff` is None there — but read
4595 // the resolved flag, not the constant, so an all-full sibling stays correct.
4596 let ff = if g.swa {
4597 None
4598 } else {
4599 self.step35_aux.as_ref().and_then(|a| a.rope_freqs(e))
4600 };
4601 #[cfg(debug_assertions)]
4602 if let Some(ff) = ff {
4603 crate::debug_assert_tensor_stream_device(ff, &e.stream(), "mtp_step35_attn.rope_freqs");
4604 }
4605 e.rope_neox2(
4606 &mut q,
4607 &mut k,
4608 pos_d,
4609 hd,
4610 g.n_rot,
4611 nh,
4612 nkv,
4613 1,
4614 g.rope_base,
4615 1.0,
4616 ff,
4617 )?;
4618
4619 // Append at the HOST slot, then re-stamp the device counter: the eager chain has the
4620 // length on the host anyway, and the windowed view below needs it there to compute the
4621 // offset. The device counter is kept in lockstep so `mtp_kv_fill`'s `set_i32_one` and any
4622 // dc-family consumer of this scratch still agree.
4623 let (kv, scratch_cap) = scratch.plane_mut(scratch_index);
4624 assert!(
4625 kv.len < scratch_cap,
4626 "step35 MTP scratch overflow ({} >= {})",
4627 kv.len,
4628 scratch_cap
4629 );
4630 let next_len = kv.len + 1;
4631 let (off, t_kv) = if g.swa && next_len > g.window {
4632 (next_len - g.window, g.window)
4633 } else {
4634 (0, next_len)
4635 };
4636 // `off`/`t_kv` stay the ATTENTION view; the retain is a separate, lower bound so the
4637 // rewind that follows this append is still resident. THIS is the only site that rebases
4638 // this plane (MEMRA_KV_REBASE_TRACE, one run: 1 rebase, all from here), so it is the site
4639 // that decides `base` for everyone.
4640 let retain_from = match kv.ring.as_ref() {
4641 Some(ring) => memra_kv::swa_retain_from(kv.len, ring.window(), ring.base()),
4642 None => off & !31usize,
4643 };
4644 let write_row = e.prepare_kv_append(kv, retain_from, 1)?;
4645 e.append_kv_quantized(
4646 &k,
4647 &v0,
4648 &mut kv.k,
4649 &mut kv.v,
4650 write_row,
4651 kv.kv_dim_k,
4652 kv.kv_dim_v,
4653 kv.k_tok_bytes,
4654 kv.v_tok_bytes,
4655 false,
4656 )?;
4657 kv.len = next_len;
4658 e.set_i32_one(&mut kv.len_d, kv.len as i32)?;
4659 // SWA view offset (see note 1). The draft chain is short (k+2 rows), but the scratch is
4660 // PERSISTENT across rounds — `mtp_kv_fill` leaves one row per committed token behind, so
4661 // `kv.len` tracks absolute position and crosses 512 in any real generation. The window is
4662 // therefore live, not theoretical.
4663 let physical = kv.physical_rows(off, off + t_kv)?;
4664 let k_view = e.view_u8_range(
4665 &kv.k,
4666 physical.start * kv.k_tok_bytes,
4667 physical.end * kv.k_tok_bytes,
4668 );
4669 let v_view = e.view_u8_range(
4670 &kv.v,
4671 physical.start * kv.v_tok_bytes,
4672 physical.end * kv.v_tok_bytes,
4673 );
4674 let mut attn = e.uninit(nh * hd)?;
4675 e.fa_decode_kvmod(
4676 &q,
4677 &k_view,
4678 &v_view,
4679 &mut attn,
4680 hd,
4681 nh,
4682 nkv,
4683 t_kv,
4684 scale,
4685 kv.k_tok_bytes,
4686 kv.v_tok_bytes,
4687 false,
4688 )?;
4689
4690 let mut ag = e.uninit(nh * hd)?;
4691 e.attn_head_gate(&attn, >, &mut ag, None, hd, nh, 1)?;
4692 Ok(e.matmul(&fa.wo, &ag, 1)?)
4693 }
4694
4695 /// The dcw draft arm's kernel-class precondition, mirrored from `fa_decode_dcw`'s own
4696 /// refusal plus the v3 walk's format contract (`fa_v3_active`), so the DEV dispatch can
4697 /// never pick an arm the launcher would refuse mid-chain (the eager chain has no graceful
4698 /// fallback point) and the CAP site refuses with the named reason instead.
4699 ///
4700 /// `cap` = the SESSION's scratch-plane row capacity: the launcher's vec gate reads
4701 /// `bucket_max = min(window, cap)`, so a SMALL session (tiny prompt + tiny max_tokens,
4702 /// e.g. a max_tokens=8 probe: cap ~62 < the 96 vec floor) is OUTSIDE the dcw domain even
4703 /// though the WINDOW clears the floor. Mirroring the window alone shipped exactly that
4704 /// hole when the door default flipped ON (2026-08-30, vision-cell receipt: sampled
4705 /// capture WARN + `[engine-error] fa_decode_dcw supports the default v3-vec class only`
4706 /// hard-failing the burst — the eager dcw arm has no graceful fallback point). Sub-floor
4707 /// sessions now take the host-len kvmod arm, byte-for-byte the door-off serving.
4708 fn step35_dcw_eligible(&self, g: &crate::hybrid::Step35MtpGeom, cap: usize) -> bool {
4709 let hd = self.cfg.head_dim_k as usize;
4710 step35_draft_dcw_on()
4711 && g.swa
4712 && g.window.min(cap) >= crate::fa_vec_min_tkv()
4713 && std::env::var("MEMRA_NO_FA_VEC").is_err()
4714 && crate::fa_v3_active(hd)
4715 && hd <= 256
4716 && hd % 32 == 0
4717 }
4718
4719 /// step35 MTP-block attention, T=1, on the scratch KV: the WINDOWED DEVICE-COUNTER twin
4720 /// of `mtp_step35_attn`, serving BOTH draft paths when `step35_draft_dcw_on`. Write slot,
4721 /// key bound and SWA view offset all derive from device state (`len_d`, `base_d` written
4722 /// only at host-side rebases, and the block's `window`), so ONE captured graph serves the
4723 /// whole chain and replays see KV growth through the counter: the `mtp_full_attn_dc`
4724 /// contract plus the view offset the plain `_dc` kernel could not express (the old
4725 /// capture-refusal root cause). The three step35 properties stay per-geom exactly as in
4726 /// the eager twin: nh/nkv from `Step35MtpGeom`, the separate head-wise gate
4727 /// (`attn_head_gate`), per-layer rope width/base with SWA passing null freqs.
4728 ///
4729 /// bucket_max = min(cap, window): the windowed view never exceeds `window` rows, so the
4730 /// capture-time grid stays valid for every replayed len, and the kernel derives ns_eff
4731 /// from the LIVE T_kv at the fixed split_keys (one-partition law). Both arms call THIS
4732 /// launcher at THIS bucket, so eager and captured drafts are bit-identical by
4733 /// construction; vs the retired-by-flag `mtp_step35_attn` the only numeric-class deltas
4734 /// are the sub-vec-floor region (t_kv < 96: kvmod ran scalar, dcw stays vec) and any
4735 /// live-len split-ladder rung below the bucket's, both draft-side only (the verify
4736 /// arbitrates emitted bytes; acceptance is gated by the battery).
4737 ///
4738 /// Host len is NOT advanced here (graph contract); callers mirror. The EAGER caller runs
4739 /// `prepare_kv_append` per step (ring headroom, rebase legal there); the CAPTURED path
4740 /// pre-arms headroom at capture time and round start (`MtpScratch::ensure_dcw_headroom`)
4741 /// because a rebase is host work no captured chain may contain.
4742 fn mtp_step35_attn_dcw(
4743 &self,
4744 e: &Engine,
4745 fa: &FullAttnLayer,
4746 g: &crate::hybrid::Step35MtpGeom,
4747 h: &CudaSlice<f32>,
4748 pos_d: &CudaSlice<i32>,
4749 scratch: &mut MtpScratch,
4750 scratch_index: usize,
4751 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4752 let (nh, nkv, hd) = (g.n_head, g.n_head_kv, self.cfg.head_dim_k as usize);
4753 // MTP-GEOM RECEIPT (dcw twin of the `mtp_step35_attn` receipt): once per process,
4754 // naming the arm, so a serving log proves WHICH draft attention program ran (the
4755 // engagement receipt for the flag door, both directions).
4756 {
4757 static ONCE: std::sync::OnceLock<()> = std::sync::OnceLock::new();
4758 ONCE.get_or_init(|| {
4759 eprintln!(
4760 "[mtp-geom] arm=dcw block={} swa={} window={} n_head={nh} n_head_kv={nkv} \
4761 head_dim_k={hd} n_rot={} rope_base={} clamp_shexp={:?}",
4762 g.il, g.swa, g.window, g.n_rot, g.rope_base, g.clamp_shexp,
4763 );
4764 });
4765 }
4766 let eps = self.cfg.rms_eps;
4767 let scale = 1.0 / (hd as f32).sqrt(); // step35.cpp:255 kq_scale
4768 let n_embd = self.cfg.n_embd as usize;
4769 let gw = fa
4770 .attn_gate
4771 .as_ref()
4772 .ok_or("step35 MTP block is missing attn_gate.weight (head-wise attention gate)")?;
4773
4774 let (q0, k0, v0, gt) = if e.uses_q8_1_fast(&fa.wq)
4775 && e.uses_q8_1_fast(&fa.wk)
4776 && e.uses_q8_1_fast(&fa.wv)
4777 && e.uses_q8_1_fast(gw)
4778 {
4779 let (hq, hdq) = e.quantize_q8_1(h, 1, n_embd)?;
4780 let (a, b, c) = match e.matmul_q8_fused3(&fa.wq, &fa.wk, &fa.wv, &hq, &hdq)? {
4781 Some(t3) => t3,
4782 None => (
4783 e.matmul_pre(&fa.wq, &hq, &hdq, h, 1)?,
4784 e.matmul_pre(&fa.wk, &hq, &hdq, h, 1)?,
4785 e.matmul_pre(&fa.wv, &hq, &hdq, h, 1)?,
4786 ),
4787 };
4788 (a, b, c, e.matmul_pre(gw, &hq, &hdq, h, 1)?)
4789 } else {
4790 (
4791 e.matmul(&fa.wq, h, 1)?,
4792 e.matmul(&fa.wk, h, 1)?,
4793 e.matmul(&fa.wv, h, 1)?,
4794 e.matmul(gw, h, 1)?,
4795 )
4796 };
4797
4798 let mut q = e.zeros(nh * hd)?;
4799 e.rms_norm(&q0, fa.q_norm.float_data(), &mut q, hd, nh, eps)?;
4800 let mut k = e.zeros(nkv * hd)?;
4801 e.rms_norm(&k0, fa.k_norm.float_data(), &mut k, hd, nkv, eps)?;
4802 // rope_freqs (llama3 factors) apply to the FULL-attn layers ONLY; SWA passes null
4803 // (the eager twin's rule, resolved from the flag, not the constant).
4804 let ff = if g.swa {
4805 None
4806 } else {
4807 self.step35_aux.as_ref().and_then(|a| a.rope_freqs(e))
4808 };
4809 #[cfg(debug_assertions)]
4810 if let Some(ff) = ff {
4811 crate::debug_assert_tensor_stream_device(
4812 ff,
4813 &e.stream(),
4814 "mtp_step35_attn_dcw.rope_freqs",
4815 );
4816 }
4817 e.rope_neox2(
4818 &mut q,
4819 &mut k,
4820 pos_d,
4821 hd,
4822 g.n_rot,
4823 nh,
4824 nkv,
4825 1,
4826 g.rope_base,
4827 1.0,
4828 ff,
4829 )?;
4830
4831 let (kv, cap) = scratch.plane_mut(scratch_index);
4832 // Append at the DEVICE slot's PHYSICAL row (len_d - base_d), then advance the counter
4833 // in-graph. Physical room is the callers' headroom contract (see the fn doc).
4834 e.append_kv_quantized_dcw(
4835 &k,
4836 &v0,
4837 &mut kv.k,
4838 &mut kv.v,
4839 &kv.len_d,
4840 kv.base_d.as_ref(),
4841 kv.kv_dim_k,
4842 kv.kv_dim_v,
4843 kv.k_tok_bytes,
4844 kv.v_tok_bytes,
4845 )?;
4846 e.inc_seqlen(&mut kv.len_d)?;
4847 // Full-buffer views (any in-round physical row stays in range under the headroom
4848 // contract); the kernel bounds and offsets the key range from (len_d, base_d, window).
4849 let k_view = e.view_u8(&kv.k, kv.k.len());
4850 let v_view = e.view_u8(&kv.v, kv.v.len());
4851 let bucket = g.window.min(cap);
4852 let mut attn = e.zeros(nh * hd)?;
4853 e.fa_decode_dcw(
4854 &q,
4855 &k_view,
4856 &v_view,
4857 &mut attn,
4858 hd,
4859 nh,
4860 nkv,
4861 &kv.len_d,
4862 kv.base_d.as_ref(),
4863 if g.swa { g.window } else { 0 },
4864 bucket,
4865 scale,
4866 kv.k_tok_bytes,
4867 kv.v_tok_bytes,
4868 None,
4869 )?;
4870
4871 let mut ag = e.zeros(nh * hd)?;
4872 e.attn_head_gate(&attn, >, &mut ag, None, hd, nh, 1)?;
4873 Ok(e.matmul(&fa.wo, &ag, 1)?)
4874 }
4875
4876 /// MTP-block full attention, T=1, on the scratch KV (BOTH draft paths — eager and graph):
4877 /// the scratch write slot and the attention bound come from `scratch.kv.len_d` (device i32[1])
4878 /// so the launch args are FIXED across draft steps — ONE captured graph serves the whole
4879 /// chain, and replays keep seeing KV growth through the device counter (no recapture).
4880 /// Geometry contract: n_splits is sized from `scratch.cap` (the persistent capacity); splits
4881 /// whose key range lies beyond the device t_kv exit empty and the shared combine skips them
4882 /// (fa_decode_dc bit-correct-for-any-t_kv<=bucket_max contract). The eager path uses the SAME
4883 /// launcher with the SAME bucket_max -> identical dispatch -> bit-identical draft tokens (the
4884 /// graph-vs-eager parity gate). Host len is NOT advanced here (graph contract); callers mirror.
4885 fn mtp_full_attn_dc(
4886 &self,
4887 e: &Engine,
4888 fa: &FullAttnLayer,
4889 h: &CudaSlice<f32>,
4890 pos_d: &CudaSlice<i32>,
4891 scratch: &mut MtpScratch,
4892 scratch_index: usize,
4893 geom: Option<&crate::hybrid::DraftGeom>,
4894 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4895 let cfg = &self.cfg;
4896 let mtp_il = cfg.n_layer.saturating_sub(cfg.nextn_predict_layers);
4897 let geometry = cfg.full_attention_geometry_at(mtp_il);
4898 let n_head = geom.map(|g| g.n_head).unwrap_or(geometry.n_head as usize);
4899 let n_head_kv = geom
4900 .map(|g| g.n_head_kv)
4901 .unwrap_or(geometry.n_head_kv as usize);
4902 let head_dim = geometry.head_dim_k as usize;
4903 let eps = cfg.rms_eps;
4904 let scale = geometry.attention_scale();
4905 let n_embd = geom.map(|g| g.d_inner).unwrap_or(cfg.n_embd as usize);
4906 let bucket_max = scratch.plane(scratch_index).1;
4907
4908 let (qf, mut k, v) =
4909 if e.uses_q8_1_fast(&fa.wq) && e.uses_q8_1_fast(&fa.wk) && e.uses_q8_1_fast(&fa.wv) {
4910 let (hq, hd) = e.quantize_q8_1(h, 1, n_embd)?;
4911 (
4912 e.matmul_pre(&fa.wq, &hq, &hd, h, 1)?,
4913 e.matmul_pre(&fa.wk, &hq, &hd, h, 1)?,
4914 e.matmul_pre(&fa.wv, &hq, &hd, h, 1)?,
4915 )
4916 } else {
4917 (
4918 e.matmul(&fa.wq, h, 1)?,
4919 e.matmul(&fa.wk, h, 1)?,
4920 e.matmul(&fa.wv, h, 1)?,
4921 )
4922 };
4923 // M3/Hy3 have no attention output gate — wq out is exactly q; skip the split.
4924 let gated = geometry.attention_gate == memra_gguf::config::AttentionGateKind::FusedQ;
4925 let (mut q, gate) = if gated {
4926 let mut q = e.zeros(n_head * head_dim)?;
4927 let mut gate = e.zeros(n_head * head_dim)?;
4928 e.q_gate_split(&qf, &mut q, &mut gate, head_dim, n_head, 1)?;
4929 (q, Some(gate))
4930 } else {
4931 (qf, None)
4932 };
4933
4934 let mut qn = e.zeros(n_head * head_dim)?;
4935 e.rms_norm(&q, fa.q_norm.float_data(), &mut qn, head_dim, n_head, eps)?;
4936 q = qn;
4937 let mut kn = e.zeros(n_head_kv * head_dim)?;
4938 e.rms_norm(
4939 &k,
4940 fa.k_norm.float_data(),
4941 &mut kn,
4942 head_dim,
4943 n_head_kv,
4944 eps,
4945 )?;
4946 k = kn;
4947 let rope_dims = geometry.n_rot as usize;
4948 e.rope_neox(
4949 &mut q,
4950 pos_d,
4951 head_dim,
4952 rope_dims,
4953 n_head,
4954 1,
4955 geometry.rope_base,
4956 1.0,
4957 )?;
4958 e.rope_neox(
4959 &mut k,
4960 pos_d,
4961 head_dim,
4962 rope_dims,
4963 n_head_kv,
4964 1,
4965 geometry.rope_base,
4966 1.0,
4967 )?;
4968
4969 let kv = scratch.plane_mut(scratch_index).0;
4970 // append at the DEVICE slot (kv.len_d == old len), then advance the counter in-graph.
4971 e.append_kv_quantized_dc(
4972 &k,
4973 &v,
4974 &mut kv.k,
4975 &mut kv.v,
4976 &kv.len_d,
4977 kv.kv_dim_k,
4978 kv.kv_dim_v,
4979 kv.k_tok_bytes,
4980 kv.v_tok_bytes,
4981 false,
4982 )?;
4983 e.inc_seqlen(&mut kv.len_d)?;
4984 // full-buffer views (any in-round t_kv stays in range on replay); the kernel bounds the
4985 // key range from the device counter.
4986 let k_view = e.view_u8(&kv.k, kv.k.len());
4987 let v_view = e.view_u8(&kv.v, kv.v.len());
4988 let (ktb, vtb) = (kv.k_tok_bytes, kv.v_tok_bytes);
4989 let mut attn = e.zeros(n_head * head_dim)?;
4990 e.fa_decode_dc(
4991 &q, &k_view, &v_view, &mut attn, head_dim, n_head, n_head_kv, &kv.len_d, bucket_max,
4992 scale, ktb, vtb, false,
4993 )?;
4994
4995 let attn_g = match &gate {
4996 Some(gate) => {
4997 let mut gsig = e.zeros(n_head * head_dim)?;
4998 e.sigmoid(gate, &mut gsig, n_head * head_dim)?;
4999 let mut ag = e.zeros(n_head * head_dim)?;
5000 e.mul(&attn, &gsig, &mut ag, n_head * head_dim)?;
5001 ag
5002 }
5003 None => attn,
5004 };
5005 Ok(e.matmul(&fa.wo, &attn_g, 1)?)
5006 }
5007
5008 /// PERSISTENT-DRAFT-KV fill (the reference engine's "mtp_update" analogue): compute the MTP
5009 /// block's K/V for `tokens` (committed tokens at positions pos0..pos0+T) from their EXACT
5010 /// trunk hiddens `h` ([T, n_embd] token-major, pre-output_norm) and append at slots pos0.. of
5011 /// the scratch KV. K/V-ONLY — ops A/1-5 plus the K-side of op 6 (wk/wv + k_norm + rope +
5012 /// quantized append); no wq/attention/FFN/lm_head, so per-token cost ~= eh_proj + wk/wv (a
5013 /// small fraction of one trunk layer), T-batched. Rope follows the chain convention
5014 /// rope(token@p) = p+1. Runs at round boundaries OUTSIDE the captured graph in BOTH draft
5015 /// modes -> draft parity by construction. Caller must have scratch.kv.len == pos0.
5016 #[allow(clippy::too_many_arguments)]
5017 fn mtp_kv_fill_at(
5018 &self,
5019 e: &Engine,
5020 mtp: &MtpHead,
5021 tokens: &[u32],
5022 h: &CudaSlice<f32>,
5023 pos0: usize,
5024 scratch: &mut MtpScratch,
5025 scratch_index: usize,
5026 embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
5027 ) -> Result<(), Box<dyn std::error::Error>> {
5028 let cfg = &self.cfg;
5029 let n_embd = cfg.n_embd as usize;
5030 let eps = cfg.rms_eps;
5031 let t = tokens.len();
5032 let (scratch_kv, scratch_cap) = scratch.plane(scratch_index);
5033 assert_eq!(scratch_kv.len, pos0, "mtp_kv_fill: append slot mismatch");
5034 assert!(pos0 + t <= scratch_cap, "mtp_kv_fill: scratch overflow");
5035 let Mixer::Full(fa) = &mtp.mixer else {
5036 panic!("MTP block is full-attn in qwen35; linear MTP not supported")
5037 };
5038 let pos_vec: Vec<i32> = (0..t).map(|i| (pos0 + i + 1) as i32).collect();
5039 let pos_d = e.htod_i32(&pos_vec)?;
5040
5041 // ops A/1/2: embed + the two input norms, T-wide.
5042 let e_emb = match embd_dev {
5043 Some((g, qt, rb)) => e.embed_gather_device_t(g, tokens, n_embd, qt, rb)?,
5044 None => e.htod(&self.embd.gather(n_embd, tokens))?,
5045 };
5046 let mut e_norm = e.zeros(t * n_embd)?;
5047 e.rms_norm(&e_emb, mtp.enorm.float_data(), &mut e_norm, n_embd, t, eps)?;
5048 let mut h_norm = e.zeros(t * n_embd)?;
5049 e.rms_norm(h, mtp.hnorm.float_data(), &mut h_norm, n_embd, t, eps)?;
5050
5051 // op 3: per-row [e_norm ; h_norm] concat, token-major [T, 2*n_embd].
5052 let mut concat = e.zeros(t * 2 * n_embd)?;
5053 for i in 0..t {
5054 e.copy_view_into(
5055 &mut concat,
5056 i * 2 * n_embd,
5057 &e_norm.slice(i * n_embd..(i + 1) * n_embd),
5058 n_embd,
5059 )?;
5060 e.copy_view_into(
5061 &mut concat,
5062 i * 2 * n_embd + n_embd,
5063 &h_norm.slice(i * n_embd..(i + 1) * n_embd),
5064 n_embd,
5065 )?;
5066 }
5067
5068 // ops 4/5: eh_proj + attn_norm, T-wide (at the student inner width when geom is set).
5069 let di = mtp.geom.as_ref().map(|g| g.d_inner).unwrap_or(n_embd);
5070 let inp_sa = e.matmul(&mtp.eh_proj, &concat, t)?;
5071 let mut a_norm = e.zeros(t * di)?;
5072 e.rms_norm(&inp_sa, mtp.attn_norm.float_data(), &mut a_norm, di, t, eps)?;
5073
5074 // op 6 (K/V half): wk/wv + k_norm + rope + per-row quantized append. No wq/attention —
5075 // the fill only has to leave correct K/V rows behind for later chains to attend over.
5076 let n_head_kv = mtp
5077 .geom
5078 .as_ref()
5079 .map(|g| g.n_head_kv)
5080 .or(mtp.step35.as_ref().map(|s| s.n_head_kv))
5081 .unwrap_or_else(|| {
5082 let mtp_il = cfg.n_layer.saturating_sub(cfg.nextn_predict_layers);
5083 cfg.full_attention_geometry_at(mtp_il).n_head_kv as usize
5084 });
5085 let mtp_il = cfg.n_layer.saturating_sub(cfg.nextn_predict_layers);
5086 let geometry = cfg.full_attention_geometry_at(mtp_il);
5087 let head_dim = geometry.head_dim_k as usize;
5088 let mut k = e.matmul(&fa.wk, &a_norm, t)?;
5089 let v = e.matmul(&fa.wv, &a_norm, t)?;
5090 let mut kn = e.zeros(t * n_head_kv * head_dim)?;
5091 e.rms_norm(
5092 &k,
5093 fa.k_norm.float_data(),
5094 &mut kn,
5095 head_dim,
5096 n_head_kv * t,
5097 eps,
5098 )?;
5099 k = kn;
5100 // step35: rotary width AND base are per-layer, and the MTP block's values come from the
5101 // resolved `Step35MtpGeom` — NOT from `cfg.rope_dim_count`/`cfg.rope_freq_base`, which
5102 // carry the arch defaults (128 / 5e6, i.e. the FULL-attn layers' base). Getting this wrong
5103 // writes K rows the attention arm then re-derives at a different theta: correct-looking
5104 // output with dead acceptance, invisible to the exactness gates.
5105 let (rope_dims, rope_base, ff) = match mtp.step35.as_ref() {
5106 Some(s) => (
5107 s.n_rot,
5108 s.rope_base,
5109 if s.swa {
5110 None
5111 } else {
5112 self.step35_aux.as_ref().and_then(|a| a.rope_freqs(e))
5113 },
5114 ),
5115 None => (geometry.n_rot as usize, geometry.rope_base, None),
5116 };
5117 #[cfg(debug_assertions)]
5118 if let Some(ff) = ff {
5119 crate::debug_assert_tensor_stream_device(ff, &e.stream(), "mtp_kv_fill.rope_freqs");
5120 }
5121 match ff {
5122 Some(f) => e.rope_neox_ff(
5123 &mut k, &pos_d, head_dim, rope_dims, n_head_kv, t, rope_base, 1.0, f,
5124 )?,
5125 None => e.rope_neox(
5126 &mut k, &pos_d, head_dim, rope_dims, n_head_kv, t, rope_base, 1.0,
5127 )?,
5128 }
5129
5130 let kv = scratch.plane_mut(scratch_index).0;
5131 // Match the trunk prime contract: a chunk may need the aligned window immediately before
5132 // its first row, so preserve that prefix when the physical tail rebases at wrap.
5133 let retain_from = kv
5134 .ring
5135 .as_ref()
5136 .map(|ring| memra_kv::swa_retain_from(pos0, ring.window(), ring.base()))
5137 .unwrap_or(0);
5138 let write_row = e.prepare_kv_append(kv, retain_from, t)?;
5139 for i in 0..t {
5140 let k_row = k.slice(i * kv.kv_dim_k..(i + 1) * kv.kv_dim_k);
5141 let v_row = v.slice(i * kv.kv_dim_v..(i + 1) * kv.kv_dim_v);
5142 e.append_kv_quantized_view(
5143 &k_row,
5144 &v_row,
5145 &mut kv.k,
5146 &mut kv.v,
5147 write_row + i,
5148 kv.kv_dim_k,
5149 kv.kv_dim_v,
5150 kv.k_tok_bytes,
5151 kv.v_tok_bytes,
5152 false,
5153 )?;
5154 }
5155 kv.len = pos0 + t;
5156 e.set_i32_one(&mut kv.len_d, kv.len as i32)?;
5157 Ok(())
5158 }
5159
5160 #[allow(clippy::too_many_arguments)]
5161 fn mtp_kv_fill_all(
5162 &self,
5163 e: &Engine,
5164 tokens: &[u32],
5165 h: &CudaSlice<f32>,
5166 pos0: usize,
5167 scratch: &mut MtpScratch,
5168 embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
5169 ) -> Result<(), Box<dyn std::error::Error>> {
5170 debug_assert_eq!(self.mtp_head_count(), scratch.plane_count());
5171 for index in 0..self.mtp_head_count() {
5172 self.mtp_kv_fill_at(
5173 e,
5174 self.mtp_head_at(index),
5175 tokens,
5176 h,
5177 pos0,
5178 scratch,
5179 index,
5180 embd_dev,
5181 )?;
5182 }
5183 Ok(())
5184 }
5185
5186 /// CAPTURE body for the GRAPH DRAFT (stage 2 of graph-grade spec): ONE MTP head forward with
5187 /// every varying input device-resident —
5188 /// - token id from the persistent `tok_d` (the previous replay's in-graph argmax wrote it,
5189 /// so the chain feeds itself; the host reads the same 4 bytes for the draft list),
5190 /// - h_seed from the persistent `h_seed_d` (h_nextn is copied BACK into it at the end),
5191 /// - rope pos from the persistent `pos_d` counter (inc'd in-graph),
5192 /// - scratch KV slot/bound from `scratch.kv.len_d` (see mtp_full_attn_dc).
5193 /// The p-min confidence lands in the persistent `p_d` iff `with_prob` (env is fixed per run).
5194 /// Same kernels, same dispatch as the eager mtp_head_forward_dev chain -> same draft tokens
5195 /// (exactness never depends on drafts — the verify arbitrates — but acceptance parity does).
5196 /// `with_head=false` captures the HEAD-LESS twin for the pseudo-seed replay (2026-07-03):
5197 /// the pseudo pass only needs h_nextn (op 10) + the scratch append — the lm_head read
5198 /// (~1.06ms q6_K on the 9B), argmax and prob are dead weight there. h_nextn's inputs are
5199 /// untouched, so the seed value is identical; round-start resets overwrite tok_d/p_d anyway.
5200 /// `sampled_cap` = Some((ctr_d, perturb_d, q_out_d, seed, temp)) captures the SAMPLED twin
5201 /// (step 3 of the sampled-spec arc): head logits are retained in the persistent `q_out_d`
5202 /// (host D2Ds them to the round's q slot after each replay), the DEVICE event counter is
5203 /// bumped in-graph, and the argmax reads GUMBEL-PERTURBED logits — one categorical draw per
5204 /// replay, bit-identical to the eager arm's gumbel_perturb at the same (seed, sctr, temp).
5205 /// seed/temp are capture-time constants (fixed per generate call, like p_min).
5206 #[allow(clippy::too_many_arguments)]
5207 fn mtp_head_forward_cap(
5208 &self,
5209 e: &Engine,
5210 mtp: &MtpHead,
5211 tok_d: &mut CudaSlice<u32>,
5212 pos_d: &mut CudaSlice<i32>,
5213 h_seed_d: &mut CudaSlice<f32>,
5214 p_d: &mut CudaSlice<f32>,
5215 scratch: &mut MtpScratch,
5216 // Which scratch plane this head appends to / attends over: 0 for the single-head
5217 // chain (every pre-lane caller), the head's own plane index for the multi-head
5218 // chain graphs (each head owns one plane — `mtp_chain_forward_dev`'s contract).
5219 scratch_index: usize,
5220 with_prob: bool,
5221 with_head: bool,
5222 embd_gpu: &CudaSlice<u8>,
5223 embd_qt: i32,
5224 embd_rb: usize,
5225 d_vocab: usize,
5226 sampled_cap: Option<SampledCapArgs<'_>>,
5227 stream_pack: Option<(&mut CudaSlice<u32>, usize, Option<&CudaSlice<u32>>)>,
5228 // DRAFT-SIDE GRAMMAR MASK (lane/draft-mask): (packed draft-vocab allowed-set buffer,
5229 // word count). Captured as ONE mask_logits_f32 node between the head matmul and the
5230 // in-graph argmax; the buffer address is baked, its CONTENTS are re-uploaded by the
5231 // host before every replay (the decode.rs graph-mask pattern). All-ones contents = a
5232 // no-op ban, so a position the grammar cannot constrain costs one pass over the row.
5233 mask_cap: Option<(&CudaSlice<u32>, usize)>,
5234 ) -> Result<(), Box<dyn std::error::Error>> {
5235 let cfg = &self.cfg;
5236 let n_embd = cfg.n_embd as usize;
5237 // step35: capturable through the WINDOWED device-counter arm (`mtp_step35_attn_dcw`)
5238 // once the dcw door is armed and the v3-vec class is live. Without the door this stays
5239 // the deliberate, named refusal: the plain `_dc` attention's key bound always starts at
5240 // row 0, cannot express this block's SWA view offset, and a captured chain would
5241 // silently attend OUTSIDE the window once the persistent scratch passes 512 rows.
5242 // Returning Err (not a panic) is what the capture sites already handle by degrading to
5243 // the eager chain (`mtp_head_forward_dev` -> `mtp_step35_attn`).
5244 // ROUND-STREAM stays refused EITHER WAY: the stream VERIFY has no step35 twin (see the
5245 // step35_verify refusal), so a stream capture that succeeded here would only move the
5246 // failure from capture time (graceful stream-off) to serve time (a failed round).
5247 if let Some(g) = mtp.step35.as_ref() {
5248 if stream_pack.is_some() {
5249 return Err(
5250 "step35 has no ROUND-STREAM draft arm (the stream verify has no step35 \
5251 twin); stream off"
5252 .into(),
5253 );
5254 }
5255 if !self.step35_dcw_eligible(g, scratch.plane(scratch_index).1) {
5256 return Err(format!(
5257 "step35 has no captured draft chain (fa_decode_dc cannot express the MTP \
5258 block's SWA view offset; the windowed dcw capture needs \
5259 MEMRA_STEP35_DRAFT_DCW armed [default ON, =0 disarms] and the v3-vec \
5260 class live at bucket=min(window {}, scratch cap {})) - the eager draft \
5261 chain serves this shape",
5262 g.window,
5263 scratch.plane(scratch_index).1,
5264 )
5265 .into());
5266 }
5267 }
5268 // student inner width (see mtp_head_forward_dev) — interface dims stay n_embd.
5269 let di = mtp.geom.as_ref().map(|g| g.d_inner).unwrap_or(n_embd);
5270 let eps = cfg.rms_eps;
5271 let e_emb = e.embed_gather_device(embd_gpu, tok_d, n_embd, embd_qt, embd_rb)?;
5272 let mut e_norm = e.zeros(n_embd)?;
5273 e.rms_norm(&e_emb, mtp.enorm.float_data(), &mut e_norm, n_embd, 1, eps)?;
5274 let mut h_norm = e.zeros(n_embd)?;
5275 e.rms_norm(
5276 &*h_seed_d,
5277 mtp.hnorm.float_data(),
5278 &mut h_norm,
5279 n_embd,
5280 1,
5281 eps,
5282 )?;
5283 let mut concat = e.zeros(2 * n_embd)?;
5284 e.copy_into(&mut concat, 0, &e_norm, n_embd)?;
5285 e.copy_into(&mut concat, n_embd, &h_norm, n_embd)?;
5286 let inp_sa = e.matmul(&mtp.eh_proj, &concat, 1)?;
5287 let mut a_norm = e.zeros(di)?;
5288 e.rms_norm(&inp_sa, mtp.attn_norm.float_data(), &mut a_norm, di, 1, eps)?;
5289 let attn_out = match (&mtp.mixer, mtp.step35.as_ref()) {
5290 // step35 (eligibility already enforced by the refusal above): the windowed dcw
5291 // arm, the SAME launcher the eager dev arm runs when the door is armed. No host
5292 // work here (this is the capture body); headroom is the callers' pre-arm.
5293 (Mixer::Full(fa), Some(g)) => {
5294 self.mtp_step35_attn_dcw(e, fa, g, &a_norm, pos_d, scratch, scratch_index)?
5295 }
5296 (Mixer::Full(fa), None) => self.mtp_full_attn_dc(
5297 e,
5298 fa,
5299 &a_norm,
5300 pos_d,
5301 scratch,
5302 scratch_index,
5303 mtp.geom.as_ref(),
5304 )?,
5305 (Mixer::Linear(_), _) => {
5306 panic!("MTP block is full-attn in qwen35; linear MTP not supported")
5307 }
5308 (Mixer::Mla(_), _) => crate::hybrid::mla_forward_unimplemented(),
5309 };
5310 let mut x1 = e.zeros(di)?;
5311 e.add(&inp_sa, &attn_out, &mut x1, di)?;
5312 let mut z = e.zeros(di)?;
5313 e.rms_norm(&x1, mtp.post_attn_norm.float_data(), &mut z, di, 1, eps)?;
5314 let ffn_out = match &mtp.ffn {
5315 crate::hybrid::Ffn::Dense {
5316 ffn_gate,
5317 ffn_up,
5318 ffn_down,
5319 } => {
5320 let n_ff = ffn_gate.out_features();
5321 let (gate, up) = if e.uses_q8_1_fast(ffn_gate) && e.uses_q8_1_fast(ffn_up) {
5322 let (zq, zd) = e.quantize_q8_1(&z, 1, di)?;
5323 (
5324 e.matmul_pre(ffn_gate, &zq, &zd, &z, 1)?,
5325 e.matmul_pre(ffn_up, &zq, &zd, &z, 1)?,
5326 )
5327 } else {
5328 (e.matmul(ffn_gate, &z, 1)?, e.matmul(ffn_up, &z, 1)?)
5329 };
5330 let mut act = e.zeros(n_ff)?;
5331 // step35: the dense FFN reads the per-layer SHEXP clamp, resolved for the MTP
5332 // block's own index (the mtp_head_forward_dev rule; None for every other arch,
5333 // which is `ffn_act`'s dispatch verbatim). The eager and captured chains must
5334 // run the ONE activation program.
5335 Self::ffn_act_lim(
5336 e,
5337 &self.cfg,
5338 &gate,
5339 &up,
5340 1.0,
5341 1.0,
5342 mtp.step35.as_ref().and_then(|s| s.clamp_shexp),
5343 &mut act,
5344 n_ff,
5345 )?;
5346 e.matmul(ffn_down, &act, 1)?
5347 }
5348 // ROUND-STREAM: the 35B NextN block carries a MoE FFN. With RESIDENT experts the
5349 // dev path is pure device launches (device top-k + rows kernels, ZERO-DtoH by
5350 // design) — capture-legal. Non-resident (SLRU-lock) stays rejected: the capture
5351 // error arm degrades the caller to eager/stream-off.
5352 crate::hybrid::Ffn::Moe(m) if m.dev_exps.is_some() => {
5353 self.moe_ffn_il(e, m, &z, 1, u16::MAX)?
5354 }
5355 crate::hybrid::Ffn::Moe(_) => {
5356 return Err("graph draft requires a Dense (or resident-MoE) MTP FFN".into());
5357 }
5358 };
5359 let mut h_inner = e.zeros(di)?;
5360 e.add(&x1, &ffn_out, &mut h_inner, di)?;
5361 // student: up-project back to n_embd (carrier + head input; see mtp_head_forward_dev).
5362 let h_nextn = match mtp.geom.as_ref() {
5363 Some(g) => e.matmul(&g.out_up, &h_inner, 1)?,
5364 None => h_inner,
5365 };
5366 // MEMRA_SPEC_HPOST needs final_h even head-less (it IS the next seed under that convention).
5367 let final_h = if with_head || spec_hpost() {
5368 let final_norm = mtp.shared_head_norm.as_ref().unwrap_or(&self.output_norm);
5369 let mut fh = e.zeros(n_embd)?;
5370 e.rms_norm(&h_nextn, final_norm.float_data(), &mut fh, n_embd, 1, eps)?;
5371 Some(fh)
5372 } else {
5373 None
5374 };
5375 if with_head {
5376 let head = mtp.shared_head_head.as_ref().unwrap_or(&self.output);
5377 let mut logits = e.matmul(head, final_h.as_ref().unwrap(), 1)?;
5378 // DRAFT-SIDE GRAMMAR MASK: ban the grammar-illegal draft ids IN the captured chain,
5379 // before the argmax — proposals become legal by construction. Contents-only
5380 // per-replay upload keeps the capture valid.
5381 if let Some((mask_d, mw)) = mask_cap {
5382 e.mask_logits_col(&mut logits, mask_d, 0, d_vocab, mw)?;
5383 }
5384 if let Some(SampledCapArgs {
5385 ctr: ctr_d,
5386 perturb: perturb_d,
5387 q_out: q_out_d,
5388 seed,
5389 temp,
5390 filt,
5391 }) = sampled_cap
5392 {
5393 // SAMPLED chain: retain q (raw head logits -> persistent q_out_d; the matmul's
5394 // own buffer is pool-recycled after the capture body returns, so it can't be the
5395 // retention target), bump the device event counter, gumbel-perturb reading it,
5396 // and argmax the PERTURBED logits into tok_d — the in-graph categorical draw.
5397 e.copy_into(q_out_d, 0, &logits, d_vocab)?;
5398 e.sctr_inc(ctr_d)?;
5399 match filt {
5400 // PURE-TEMP: gumbel over the raw softmax — byte-identical to the
5401 // pre-lane capture body.
5402 None => e.gumbel_perturb_ctr(&logits, perturb_d, d_vocab, seed, ctr_d, temp)?,
5403 // FILTERED (lane/step37-draft-graph-serving-20260830): the SAME
5404 // filter_stats program the eager arm and the accept path run (the
5405 // wrapper's coop/plain choice is deployment-keyed, never per-call), then
5406 // the device-stat/device-counter perturb twin — the draft draws from the
5407 // exact filtered distribution the verify gathers `q` from. q was
5408 // retained ABOVE, pre-perturb, so the accept path's post-replay stats
5409 // recompute (same kernel, same bits) reconstructs these th/z exactly.
5410 Some(f) => {
5411 e.filter_stats(
5412 &logits, d_vocab, f.rows0, f.th, f.z, f.mx, d_vocab, 1, temp, f.top_k,
5413 f.top_p, f.min_p,
5414 )?;
5415 e.gumbel_perturb_filtered_ctr(
5416 &logits, perturb_d, d_vocab, seed, ctr_d, temp, f.mx, f.th,
5417 )?;
5418 }
5419 }
5420 e.argmax_token_device_into(perturb_d, tok_d, d_vocab)?;
5421 // p-min prob = the head's RAW softmax confidence in the SAMPLED pick — same
5422 // semantics as the eager sampled arm's prob_of_token_device(dl_d, tok_d).
5423 if with_prob {
5424 e.prob_of_token_device_into(&logits, tok_d, p_d, d_vocab)?;
5425 }
5426 } else {
5427 // draft token -> persistent tok_d (next replay's embed reads it; host reads the 4 bytes).
5428 e.argmax_token_device_into(&logits, tok_d, d_vocab)?;
5429 // p-min under a draft mask reads the MASKED row: confidence relative to the
5430 // grammar-LEGAL alternatives (illegal ids leave the softmax denominator), which
5431 // is the right semantics for "does the drafter know what comes next here" and
5432 // the same row the pick came from. Draft-quality only — verify arbitrates.
5433 if with_prob {
5434 e.prob_of_token_device_into(&logits, tok_d, p_d, d_vocab)?;
5435 }
5436 }
5437 }
5438 // ROUND-STREAM K-chain: pack (tok, p) into slot j, then remap tok through d2t so the
5439 // NEXT chained body's embed reads the TARGET id — zero host involvement per step.
5440 if let Some((out, slot, d2t)) = stream_pack {
5441 e.pack_tok_p(tok_d, p_d, out, slot)?;
5442 if let Some(map) = d2t {
5443 e.tok_map_u32(tok_d, map)?;
5444 }
5445 }
5446 // Next draft step's h_seed: pre-norm h_nextn (default) or post-norm final_h (HPOST).
5447 if spec_hpost() {
5448 e.copy_into(h_seed_d, 0, final_h.as_ref().unwrap(), n_embd)?;
5449 } else {
5450 e.copy_into(h_seed_d, 0, &h_nextn, n_embd)?;
5451 }
5452 // advance the draft rope position in-graph.
5453 e.inc_seqlen(pos_d)?;
5454 Ok(())
5455 }
5456
5457 /// Batched target verify forward over `tokens` at positions `pos0..pos0+T` (§D.3, T=K+1).
5458 /// Returns ALL T logit columns (host f32, [T*n_vocab]); appends T cols to every full-attn KV
5459 /// and advances every linear-attn recur state by T steps (the recur steps are SEQUENTIAL T=1).
5460 /// Advances `cache.pos` by T.
5461 pub fn decode_step_t(
5462 &self,
5463 e: &Engine,
5464 tokens: &[u32],
5465 pos0: usize,
5466 cache: &mut Cache,
5467 ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
5468 if self.is_gemma4_e4b() {
5469 return Ok(self.gemma4_e4b_decode_step_t_h(e, tokens, pos0, cache)?.0);
5470 }
5471 if self.gemma_batch_program() {
5472 return self.gemma4_decode_step_t(e, tokens, pos0, cache);
5473 }
5474 Ok(self.decode_step_t_h(e, tokens, pos0, cache)?.0)
5475 }
5476
5477 /// Like `decode_step_t` but ALSO returns the LAST column's pre-output_norm hidden (h_seed for
5478 /// the next draft round). This lets partial-accept replay run as ONE batched T=(n_acc+1) forward
5479 /// (single weight read) instead of n_acc+1 separate T=1 decode_steps (n_acc+1 weight reads).
5480 /// At batch=1 decode is bandwidth-bound, so batching the replay is THE MTP profitability lever.
5481 pub fn decode_step_t_h(
5482 &self,
5483 e: &Engine,
5484 tokens: &[u32],
5485 pos0: usize,
5486 cache: &mut Cache,
5487 ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
5488 self.decode_step_t_h_emb(e, tokens, pos0, cache, None)
5489 }
5490
5491 /// Like `decode_step_t_h` with an optional RESIDENT embed table (spec hot loop): device
5492 /// gather instead of host dequant + [T, n_embd] f32 htod. Bit-identical rows.
5493 pub fn decode_step_t_h_emb(
5494 &self,
5495 e: &Engine,
5496 tokens: &[u32],
5497 pos0: usize,
5498 cache: &mut Cache,
5499 embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
5500 ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
5501 let (logits_d, h_seed) = self.decode_step_t_h_emb_dev(e, tokens, pos0, cache, embd_dev)?;
5502 Ok((e.dtoh(&logits_d)?, h_seed))
5503 }
5504
5505 /// DEVICE-LOGITS verify forward (spec device-argmax lever): identical kernel chain to
5506 /// `decode_step_t_h_emb` but returns the [T, n_vocab] logits ON DEVICE — the accept walk
5507 /// argmaxes each column on-device and reads back ONE [T] u32 instead of dtoh'ing the full
5508 /// T x n_vocab f32 block (~1-4 MB + T host argmaxes, every round). Kernel dispatch is
5509 /// UNCHANGED (same decode-exact kernels); only the post-logits transfer moves.
5510 pub fn decode_step_t_h_emb_dev(
5511 &self,
5512 e: &Engine,
5513 tokens: &[u32],
5514 pos0: usize,
5515 cache: &mut Cache,
5516 embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
5517 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
5518 let n_embd = self.cfg.n_embd as usize;
5519 let t = tokens.len();
5520 let (logits, x) = self.decode_step_t_core(e, tokens, pos0, cache, embd_dev, None)?;
5521 // h_seed for the next round = LAST column's pre-output_norm hidden ([n_embd]).
5522 let mut hs = vbuf(e, n_embd)?; // fully written by copy_view_into below
5523 e.copy_view_into(&mut hs, 0, &x.slice((t - 1) * n_embd..t * n_embd), n_embd)?;
5524 Ok((logits, hs))
5525 }
5526
5527 /// CORE verify forward: the `decode_step_t_h_emb_dev` kernel chain, returning the FULL
5528 /// pre-output_norm hidden stack x ([T, n_embd], any column extractable) and optionally
5529 /// filling a `VerifyCkpt` (retained per-layer state-rebuild inputs) for the REPLAY-FREE
5530 /// partial accept. `ckpt: None` => byte-for-byte the old behavior (the ckpt writes are pure
5531 /// retains/copies — they never change what any kernel computes).
5532 fn decode_step_t_core(
5533 &self,
5534 e: &Engine,
5535 tokens: &[u32],
5536 pos0: usize,
5537 cache: &mut Cache,
5538 embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
5539 mut ckpt: Option<&mut VerifyCkpt>,
5540 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
5541 self.decode_step_t_core_stream(
5542 e,
5543 tokens,
5544 pos0,
5545 cache,
5546 embd_dev,
5547 ckpt.take(),
5548 None,
5549 None,
5550 None,
5551 None,
5552 )
5553 }
5554
5555 /// [`Self::decode_step_t_core`] with the MTP route's verify-graph pool armed
5556 /// (`MEMRA_SPEC_VERIFY_GRAPH`). `graphs: None` reproduces `decode_step_t_core`
5557 /// argument-for-argument, so the eager walk stays the byte-identical fallback.
5558 fn decode_step_t_core_vg(
5559 &self,
5560 e: &Engine,
5561 tokens: &[u32],
5562 pos0: usize,
5563 cache: &mut Cache,
5564 embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
5565 mut ckpt: Option<&mut VerifyCkpt>,
5566 graphs: Option<&mut DsparkVerifyGraphs>,
5567 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
5568 self.decode_step_t_core_stream(
5569 e,
5570 tokens,
5571 pos0,
5572 cache,
5573 embd_dev,
5574 ckpt.take(),
5575 None,
5576 None,
5577 None,
5578 graphs,
5579 )
5580 }
5581
5582 /// Increment-0 two-session PP seam: release the peer after this lane's stage-0 boundary TX.
5583 /// The two independent sessions keep their own cache/checkpoint state; only issue order moves.
5584 fn decode_step_t_core_pipelined(
5585 &self,
5586 e: &Engine,
5587 tokens: &[u32],
5588 pos0: usize,
5589 cache: &mut Cache,
5590 embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
5591 mut ckpt: Option<&mut VerifyCkpt>,
5592 pipe: &SpecPipeLane,
5593 round: usize,
5594 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
5595 let fence = crate::pp::pp_cuts(self.layers.len())
5596 .ok_or("two-session speculative pipeline requires a PP stage cut")?;
5597 if crate::pp::pp2_streams_off() || !crate::pp::spec_pp_on() {
5598 return Err("two-session speculative pipeline requires the PP verify split".into());
5599 }
5600 let interval_fence = pipe.stage0_begin(round)?;
5601 let ticket = self.verify_stage0_issue(
5602 e,
5603 tokens,
5604 pos0,
5605 cache,
5606 embd_dev,
5607 ckpt.as_deref_mut(),
5608 None,
5609 &fence,
5610 Some(interval_fence),
5611 pipe.trace(round),
5612 )?;
5613 pipe.stage0_end(round);
5614 pipe.stage1_begin(round)?;
5615 let result = self.verify_stage1_finish(e, ticket, cache, ckpt, None, &fence, true)?;
5616 pipe.verify_end(round);
5617 Ok(result)
5618 }
5619
5620 /// ROUND-STREAM stage (c) 4: `stream` = (device verify tokens [t], device pos counter) —
5621 /// when Some, rope positions come from pos_iota over the counter, the embed gathers the
5622 /// device tokens, and full_attn_verify routes appends/FA through the _dc twins reading the
5623 /// SAME counter (every layer's kvl.len == cache.pos, one counter drives all three). The
5624 /// host `tokens`/`pos0` args still size buffers (t is FIXED K+1 in stream mode).
5625 /// `vtok_dev` (engine-bundle slice 2): device verify tokens for the EMBED only —
5626 /// unlike `stream` mode it changes nothing else (host pos iota, host-len KV appends).
5627 /// `tokens` then only sizes buffers (the dummy-slice pattern the round-stream arm uses).
5628 #[allow(clippy::too_many_arguments)]
5629 fn decode_step_t_core_stream(
5630 &self,
5631 e: &Engine,
5632 tokens: &[u32],
5633 pos0: usize,
5634 cache: &mut Cache,
5635 embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
5636 mut ckpt: Option<&mut VerifyCkpt>,
5637 stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
5638 pp_pipe: Option<bool>,
5639 vtok_dev: Option<&CudaSlice<u32>>,
5640 graphs: Option<&mut DsparkVerifyGraphs>,
5641 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
5642 // PP DOOR (lane/pp2-spec 2026-08-06): the verify trunk now takes its OWN stage split,
5643 // exactly as the eager and batched steps do. This is the single funnel every verify
5644 // forward reaches (decode_step_t / _h / _h_emb / _h_emb_dev / _core all land here), so
5645 // wiring it here wires the whole spec surface — the draft/accept/commit machinery above
5646 // is untouched.
5647 //
5648 // History: pp2-hardening (2026-08-06) made this funnel FAIL CLOSED, because its trunk
5649 // walk was unsplit on one stream and a sharded cross-device placement peer-read every
5650 // remote layer's weights on every spec round (measured 13.9-28x on the batched twin).
5651 // The refusal below survives to cover the residue — MEMRA_SPEC_PP=0, MEMRA_PP_STREAMS=0,
5652 // or a placement whose PpNRt fails to build — so a config that would still walk the
5653 // whole trunk on one stream refuses instead of regressing 28x.
5654 if let Some(fence) = crate::pp::pp_cuts(self.layers.len()) {
5655 if !crate::pp::pp2_streams_off() && crate::pp::spec_pp_on() {
5656 if vtok_dev.is_some() {
5657 return Err(
5658 "device-token dspark verify (slice-2 deferred readback) has no PP \
5659 stage-split arm; set MEMRA_DSPARK_DEFER_READBACK=0 or run the dspark \
5660 route on one device"
5661 .into(),
5662 );
5663 }
5664 return self.decode_step_t_core_ppn(
5665 e,
5666 tokens,
5667 pos0,
5668 cache,
5669 embd_dev,
5670 ckpt.take(),
5671 stream,
5672 &fence,
5673 pp_pipe,
5674 );
5675 }
5676 }
5677 crate::pp::refuse_unsplit_if_remote(
5678 "decode_step_t (spec verify)",
5679 "drop MEMRA_SPEC_PP=0 / MEMRA_PP_STREAMS=0 so the verify trunk takes its OWN stage \
5680 split (decode_step_t_core_ppn); or run spec on one device",
5681 )?;
5682 let cfg = &self.cfg;
5683 let n_embd = cfg.n_embd as usize;
5684 let eps = cfg.rms_eps;
5685 let t = tokens.len();
5686 let pos_d = match stream {
5687 Some((_, ctr)) => {
5688 let mut p = e.alloc_uninit::<i32>(t)?;
5689 e.pos_iota(ctr, &mut p, t)?;
5690 p
5691 }
5692 None => {
5693 let pos_vec: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
5694 e.htod_i32(&pos_vec)?
5695 }
5696 };
5697
5698 // embed T tokens -> [T, n_embd] token-major (device gather on the spec hot loop)
5699 let x = match (stream, embd_dev) {
5700 (Some((vtok, _)), Some((g, qt, rb))) => {
5701 e.embed_gather_device_td(g, vtok, t, n_embd, qt, rb)?
5702 }
5703 (None, Some((g, qt, rb))) => match vtok_dev {
5704 // slice 2: device verify tokens, same embed_gather_u32_t kernel —
5705 // bit-identical rows to the host-token arm (same per-dtype deq).
5706 Some(vt_d) => e.embed_gather_device_td(g, vt_d, t, n_embd, qt, rb)?,
5707 None => e.embed_gather_device_t(g, tokens, n_embd, qt, rb)?,
5708 },
5709 _ => {
5710 assert!(
5711 vtok_dev.is_none(),
5712 "device-token verify requires the resident embed table (embd_dev)"
5713 );
5714 e.htod(&self.embd.gather(n_embd, tokens))?
5715 }
5716 };
5717
5718 // TRUNK WALK: layers [0, n_layers) through the SAME range-scoped subgraph the PP-N
5719 // stage split calls per stage (`verify_layers`) — one code path, so the split cannot
5720 // drift from the unsplit dispatch mirroring. lane/pp2-spec 2026-08-06.
5721 let x = self.verify_layers(
5722 e,
5723 x,
5724 0,
5725 self.layers.len(),
5726 &pos_d,
5727 pos0,
5728 t,
5729 cache,
5730 ckpt.take(),
5731 stream,
5732 graphs,
5733 )?;
5734 if spec_nan_scan() {
5735 nan_scan_rows(e, &x, t, n_embd, &format!("verify trunk exit pos0={pos0}"))?;
5736 }
5737
5738 let mut hn = vbuf(e, t * n_embd)?;
5739 // Stage-A door: with the serving-class row-outer verify walk, the TAIL must be the
5740 // t=1 decode program per row too (rms_norm t=1 + the single-row bf16 head — the
5741 // split head's concat is receipted bit-identical to it). The batched cuBLASLt head
5742 // is a different ULP class and flips near-tie argmaxes off the greedy tape.
5743 let eager_tail = self.sliding_gated_moe_batch_program() && spec_verify_eager_on();
5744 if eager_tail {
5745 let n_vocab = self.cfg.n_vocab as usize;
5746 // MEMRA_SPEC_HEAD_ROWS=1 — THE VERIFY TAIL'S REDUNDANT HEAD READ.
5747 //
5748 // The loop below runs the head at m=1 once PER COLUMN, so the LM head's weights are
5749 // streamed t times per verify pass. On step37 that head is ~0.49 GiB per card after the
5750 // rank split, ~1.07 ms of pure re-read at t=2 and worse at every wider t — which is a
5751 // large part of why the fixed K ladder LOSES (K=1 81.2 > K=2 73.1 > K=3 62.7 tok/s).
5752 //
5753 // The loop's justification is the comment above: the batched cuBLASLt head is a
5754 // different ULP class and flips near-tie argmaxes off the greedy tape. That is true of
5755 // cuBLASLt and it does NOT apply here, because a FloatBf16 head at 1..=32 rows never
5756 // reaches cuBLASLt: `matmul` routes it to `matvec_bf16_rows_into` (lib.rs:12248), whose
5757 // own doc says `matvec_bf16_f32acc_x4_rows` "runs the t=1 decode head program PER ROW
5758 // (identical dot + reduce), so decode/verify tiers keep the t=1 numeric class". Under
5759 // the W8 doors both widths route to the q8 mirror instead, and the t-column mirror is
5760 // documented "bit-identical to t single-row calls". So the batched form is the SAME
5761 // arithmetic per row on both paths, with one weight read instead of t.
5762 //
5763 // rms_norm is row-wise, so norm(t) is per-row identical to t x norm(1) by construction.
5764 //
5765 // DEFAULT OFF for exactly one turn of the crank: "bit-identical by two documented
5766 // claims" is still an argument. The greedy byte tape decides, and the door flips only
5767 // once the tape is a receipt.
5768 if head_rows_on() {
5769 e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
5770 let logits = e.matmul(&self.output, &hn, t)?;
5771 if stream.is_none() {
5772 cache.pos += t;
5773 }
5774 return Ok((logits, if spec_hpost() { hn } else { x }));
5775 }
5776 let mut logits = vbuf(e, t * n_vocab)?;
5777 for r in 0..t {
5778 let mut row = e.uninit(n_embd)?;
5779 e.dtod_copy_view(&x.slice(r * n_embd..(r + 1) * n_embd), &mut row)?;
5780 let mut hr = e.uninit(n_embd)?;
5781 e.rms_norm(&row, self.output_norm.float_data(), &mut hr, n_embd, 1, eps)?;
5782 let lr = e.matmul(&self.output, &hr, 1)?;
5783 e.dtod_copy_into(&lr, &mut logits, r * n_vocab)?;
5784 e.dtod_copy_into(&hr, &mut hn, r * n_embd)?;
5785 }
5786 if stream.is_none() {
5787 cache.pos += t;
5788 }
5789 return Ok((logits, if spec_hpost() { hn } else { x }));
5790 }
5791 let serving_head =
5792 self.sliding_gated_moe_batch_program() || self.batched_serving_numeric_class();
5793 let logits = if serving_head {
5794 // Step35 and the qwen35 family (MoE 2026-08-14 AM, dense-hybrid same day PM — the
5795 // Q3.8 bring-up reproduced the identical near-tie class on dense: eager-class verify
5796 // vs batched-class live serving, ULP drift amplified through the GDN recurrence)
5797 // serve one batched numeric class at every live width, including B=1. Keep the
5798 // verify head in that same class; other generic families retain the decode-exact
5799 // head that their run-spec contract pins.
5800 e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
5801 e.matmul(&self.output, &hn, t)?
5802 } else {
5803 e.rms_norm_decode(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
5804 e.matmul_decode_exact(&self.output, &hn, t)?
5805 };
5806 // stream: the device pos counter owns position; host mirror reconciles at drain.
5807 if stream.is_none() {
5808 cache.pos += t;
5809 }
5810 // Hidden stack for seeds/refresh-fills: pre-norm x (default) or post-norm hn (HPOST).
5811 Ok((logits, if spec_hpost() { hn } else { x }))
5812 }
5813
5814 /// THE VERIFY TRUNK OVER PP-N (lane/pp2-spec 2026-08-06): `decode_step_t_core_stream`'s walk
5815 /// as N stage subgraphs, each on its own engine/stream (and, under `MEMRA_PP_DEVICES`, its own
5816 /// device), with a `[T, n_embd]` boundary transfer between them. T = K+1 (the verify batch),
5817 /// so this is the batched-boundary shape the pp2-batch lane's grow-only slots already handle
5818 /// (`tx(b, x, t*n_embd)`; the slot grows to the high-water T and the transport moves exactly
5819 /// the payload).
5820 ///
5821 /// Structure is `decode_step_batch_ppn`'s, which is `decode_step_h_ppn`'s. FOUR THINGS ARE
5822 /// PER-STAGE and each for a measured reason (see `decode_step_batch_ppn`'s header for the
5823 /// receipts):
5824 ///
5825 /// 1. THE ENGINE (`rt.engine(s, e)`) — `Engine` owns lazily-grown stable-pointer scratch
5826 /// (`fa_part_pool`, `fa_vf16_scratch`, `argmax_partials`) that is single-stream-safe BY
5827 /// DESIGN. Two stage streams through one Engine is the 2026-08-02 shared-scratch race
5828 /// (35% flake, nondeterministic all-logits divergence). `PpNRt::build` gives every stage
5829 /// s>0 its own Engine even on the primary device; honouring it here is what scopes the
5830 /// pools. The verify path allocates MORE of that scratch than eager decode does (FA at
5831 /// m=T, and the per-layer `GdnStash` retains), so this is load-bearing, not inherited.
5832 ///
5833 /// 2. `pos_d` — each stage uploads its OWN copy of the T rope positions on ITS stream, so the
5834 /// buffer is allocated, consumed and freed on one stream. In `stream` mode that means each
5835 /// stage runs its own `pos_iota` over the SHARED device counter (`pos_ctr`): the counter is
5836 /// read-only during the forward (the round's `inc`/`copy_add` happen outside it), so every
5837 /// stage derives the identical iota, and each stage's own output buffer is stream-local.
5838 ///
5839 /// 3. THE EMBED lives with stage 0 (`self.embd` / `embd_gpu` are host/primary-side; the
5840 /// sharded loader leaves the table with stage 0 by construction).
5841 ///
5842 /// 4. THE HEAD (`output_norm` + `output`) runs on the LAST stage — the sharded loader uploaded
5843 /// both through that stage's engine (`hybrid.rs`: `e_head = layer_engine(e, n_trunk,
5844 /// n_trunk-1)`), so reading them anywhere else is a peer read of the biggest tensor in the
5845 /// model, every round.
5846 ///
5847 /// WHAT STAYS ON THE PRIMARY, deliberately: the returned logits and hidden stack `x`. Both are
5848 /// last-stage-allocated device buffers, and every consumer (the device argmax walk, the accept
5849 /// kernels, `spec_seed_gather`, the ckpt rebuild in `commit_verified_prefix`) reads them
5850 /// through the primary context by UVA — the same read the batched serving epilogue's
5851 /// `last_logits_dev` park does. Those consumers are per-round O(T x n_vocab) and O(n_embd),
5852 /// not per-layer, so they are not the 28x class; splitting them is a separate lane.
5853 ///
5854 /// The MTP HEAD (draft side) is NOT split: it is one block, it lives wherever the loader put
5855 /// it (`load_mtp` uses the primary engine), and it is ~1-2 GB against the trunk's tens. Draft
5856 /// placement is measured, not assumed — see `research/pp2-spec-20260806`.
5857 ///
5858 /// EXACTNESS: PP-N adds ZERO deviation. Each stage runs the SAME kernels on the SAME bytes in
5859 /// the same order via the SAME `verify_layers` the unsplit body calls; the only change is
5860 /// where the residual is materialized, and the boundary is a straight f32 copy (dtod
5861 /// same-device / `cudaMemcpyPeerAsync` cross-device, no conversion). So the split MUST be
5862 /// BIT-IDENTICAL to the unsplit verify at the same T, in both placement orders. Gate:
5863 /// `decode-batch-gate --mode ppspec`. Acceptance counts are a DERIVED consequence — greedy
5864 /// accept argmaxes these logits, so bit-identical logits force identical accept walks; the
5865 /// `run-spec` K=1..8 arm checks that end-to-end rather than trusting the implication.
5866 #[allow(clippy::too_many_arguments)]
5867 fn decode_step_t_core_ppn(
5868 &self,
5869 e: &Engine,
5870 tokens: &[u32],
5871 pos0: usize,
5872 cache: &mut Cache,
5873 embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
5874 mut ckpt: Option<&mut VerifyCkpt>,
5875 stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
5876 fence: &[usize],
5877 pp_pipe: Option<bool>,
5878 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
5879 let ticket = self.verify_stage0_issue(
5880 e,
5881 tokens,
5882 pos0,
5883 cache,
5884 embd_dev,
5885 ckpt.as_deref_mut(),
5886 stream,
5887 fence,
5888 pp_pipe,
5889 None,
5890 )?;
5891 self.verify_stage1_finish(e, ticket, cache, ckpt, stream, fence, true)
5892 }
5893
5894 /// Enqueue embed, stage 0, and the first boundary TX, then return the actual boundary slot.
5895 /// The ordinary PP verify wrapper calls `verify_stage1_finish` immediately after this return.
5896 #[allow(clippy::too_many_arguments)]
5897 fn verify_stage0_issue(
5898 &self,
5899 e: &Engine,
5900 tokens: &[u32],
5901 pos0: usize,
5902 cache: &mut Cache,
5903 embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
5904 mut ckpt: Option<&mut VerifyCkpt>,
5905 stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
5906 fence: &[usize],
5907 pp_pipe: Option<bool>,
5908 trace: Option<SpecPipeTraceCtx>,
5909 ) -> Result<VerifyBoundaryTicket, Box<dyn std::error::Error>> {
5910 assert!(
5911 !self.is_gemma4_e4b() && !self.gemma_batch_program(),
5912 "decode_step_t_core_ppn covers the hybrid non-gemma4 verify trunk only \
5913 (the gemma4 arms have their own decode_step_t twins)"
5914 );
5915 if crate::pp::pp_host_bounce_active() && (stream.is_some() || embd_dev.is_some()) {
5916 return Err(
5917 "decode_step_t_core_ppn: refused with MEMRA_PP_HOST_BOUNCE=1 — the trunk \
5918 boundary itself is host-staged, but device-resident verify still peer-reads \
5919 primary-device token/position/embedding buffers from stage 0. Run plain PP \
5920 serving on this host class; spec requires local per-stage inputs first."
5921 .into(),
5922 );
5923 }
5924 let rt = crate::pp::PpNRt::get(e)?;
5925 let n_st = fence.len() - 1;
5926 assert_eq!(
5927 rt.n_stages(),
5928 n_st,
5929 "PpNRt stage count {} != fence stages {n_st}",
5930 rt.n_stages()
5931 );
5932 let n_embd = self.cfg.n_embd as usize;
5933 let t = tokens.len();
5934 let payload = t * n_embd;
5935 if pp_pipe.is_some() {
5936 assert_eq!(n_st, 2, "spec pipeline requires exactly two PP stages");
5937 }
5938 // One-shot lane diagnostic: force natural PP-2 boundaries to completion so the server
5939 // log can price stage 0, the peer hop, the RX copy, and stage 1 + head separately without
5940 // nsys. The ordinary path keeps every enqueue asynchronous. N>2 is deliberately excluded:
5941 // the report below names exactly two stages and must never imply it measured middle ones.
5942 let pp_anatomy = n_st == 2 && std::env::var("MEMRA_SPEC_PP_ANATOMY").as_deref() == Ok("1");
5943 let pp_started = std::time::Instant::now();
5944 let (mut reverse_ms, mut stage0_ms, mut tx_ms) = (0.0f64, 0.0f64, 0.0f64);
5945 // The CALLER's ambient stream, captured BEFORE any `rt.enter()` pushes a stage stream:
5946 // this body returns DEVICE-RESIDENT buffers (the device-argmax accept walk's contract),
5947 // so the exit needs the same publication the boundaries get — see `PpNRt::publish_to`.
5948 // Taken here, not at the end, because inside the last-stage scope `e.stream()` IS the
5949 // stage stream and the wait would self-order into a no-op.
5950 let caller_stream = e.stream();
5951 // #87 ROOT-CAUSE FENCE (lane/pp2spec-crash): the PREVIOUS round's stage-allocated
5952 // outputs (logits/hidden/ckpt stashes) freed stream-ordered on the STAGE streams while
5953 // the primary stream still holds queued reads of them — with event tracking elided,
5954 // nothing stops the pool from reusing those blocks for THIS round's stage allocations,
5955 // whose writes then race the queued reads (measured: 13/4096-NaN random-bits garbage in
5956 // the spec round seed; the full anatomy is on `PpNRt::fence_stages_behind`). Order every
5957 // stage stream behind the caller before enqueueing new stage work.
5958 let reverse_started = std::time::Instant::now();
5959 if pp_pipe != Some(false) {
5960 rt.fence_stages_behind(&caller_stream)?;
5961 }
5962 if pp_pipe == Some(true) {
5963 // Both session verifies must alternate boundary slots even when the ordinary
5964 // decode overlap experiment is off. Prewarm before A's stage 0 so B cannot grow
5965 // slot 1 by synchronizing the RX stream while A's stage 1 is in flight.
5966 rt.prepare_overlap_slots(0, payload)?;
5967 }
5968 if pp_anatomy {
5969 // Drain the reverse-publication dependency before timing stage 0 itself. At c=1 this
5970 // prices any primary-stream rollback/refresh tail inherited from the prior round.
5971 for s in 0..n_st {
5972 let _st = rt.enter(s);
5973 rt.engine(s, e).stream().synchronize()?;
5974 }
5975 reverse_ms = reverse_started.elapsed().as_secs_f64() * 1e3;
5976 }
5977
5978 // Per-stage rope positions: in host mode the same [T] iota each stage uploads itself; in
5979 // stream mode each stage's own `pos_iota` over the shared read-only device counter.
5980 let stage_pos = |es: &Engine| -> Result<CudaSlice<i32>, Box<dyn std::error::Error>> {
5981 match stream {
5982 Some((_, ctr)) => {
5983 let mut p = es.alloc_uninit::<i32>(t)?;
5984 es.pos_iota(ctr, &mut p, t)?;
5985 Ok(p)
5986 }
5987 None => {
5988 let pos_vec: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
5989 es.htod_i32(&pos_vec)
5990 }
5991 }
5992 };
5993
5994 // ---- STAGE 0: embed (the table lives with stage 0) + layers [0, fence[1]) + TX ----
5995 let slot = {
5996 let _st0 = rt.enter(0);
5997 let e0 = rt.engine(0, e);
5998 enqueue_spec_pipe_trace_marker(&e0.stream(), trace.as_ref(), "S0", "start", None)?;
5999 let stage0_started = std::time::Instant::now();
6000 let pos_d = stage_pos(e0)?;
6001 let x = match (stream, embd_dev) {
6002 (Some((vtok, _)), Some((g, qt, rb))) => {
6003 e0.embed_gather_device_td(g, vtok, t, n_embd, qt, rb)?
6004 }
6005 (None, Some((g, qt, rb))) => e0.embed_gather_device_t(g, tokens, n_embd, qt, rb)?,
6006 _ => e0.htod(&self.embd.gather(n_embd, tokens))?,
6007 };
6008 let x = self.verify_layers(
6009 e0,
6010 x,
6011 fence[0],
6012 fence[1],
6013 &pos_d,
6014 pos0,
6015 t,
6016 cache,
6017 ckpt.as_deref_mut(),
6018 stream,
6019 None,
6020 )?;
6021 if pp_anatomy {
6022 e0.stream().synchronize()?;
6023 stage0_ms = stage0_started.elapsed().as_secs_f64() * 1e3;
6024 }
6025 let tx_started = std::time::Instant::now();
6026 let slot = if pp_pipe.is_some() {
6027 rt.tx_pipelined(0, &x, payload)?
6028 } else {
6029 rt.tx(0, &x, payload)?
6030 };
6031 enqueue_spec_pipe_trace_marker(&e0.stream(), trace.as_ref(), "S0", "end", Some(slot))?;
6032 if pp_anatomy {
6033 e0.stream().synchronize()?;
6034 tx_ms = tx_started.elapsed().as_secs_f64() * 1e3;
6035 }
6036 slot
6037 // x + pos_d drop here: freed stream-ordered on stage-0's stream after use.
6038 };
6039
6040 Ok(VerifyBoundaryTicket {
6041 rt,
6042 caller_stream,
6043 slot,
6044 pos0,
6045 t,
6046 payload,
6047 n_st,
6048 pipelined: pp_pipe.is_some(),
6049 pp_anatomy,
6050 pp_started,
6051 reverse_ms,
6052 stage0_ms,
6053 tx_ms,
6054 trace,
6055 })
6056 }
6057
6058 /// Consume a stage-0 boundary ticket and enqueue the remaining PP stages plus the head.
6059 /// On PP-2 this is exactly stage 1; PP-N keeps its pre-existing middle-stage walk here.
6060 #[allow(clippy::too_many_arguments)]
6061 fn verify_stage1_finish(
6062 &self,
6063 e: &Engine,
6064 ticket: VerifyBoundaryTicket,
6065 cache: &mut Cache,
6066 mut ckpt: Option<&mut VerifyCkpt>,
6067 stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
6068 fence: &[usize],
6069 publish_to_caller: bool,
6070 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
6071 let VerifyBoundaryTicket {
6072 rt,
6073 caller_stream,
6074 slot,
6075 pos0,
6076 t,
6077 payload,
6078 n_st,
6079 pipelined,
6080 pp_anatomy,
6081 pp_started,
6082 reverse_ms,
6083 stage0_ms,
6084 tx_ms,
6085 trace,
6086 } = ticket;
6087 let n_embd = self.cfg.n_embd as usize;
6088 let eps = self.cfg.rms_eps;
6089 let mut slot = slot;
6090 let (mut rx_ms, mut stage1_ms) = (0.0f64, 0.0f64);
6091 let stage_pos = |es: &Engine| -> Result<CudaSlice<i32>, Box<dyn std::error::Error>> {
6092 match stream {
6093 Some((_, ctr)) => {
6094 let mut p = es.alloc_uninit::<i32>(t)?;
6095 es.pos_iota(ctr, &mut p, t)?;
6096 Ok(p)
6097 }
6098 None => {
6099 let pos_vec: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
6100 es.htod_i32(&pos_vec)
6101 }
6102 }
6103 };
6104
6105 // ---- MIDDLE STAGES: RX boundary s-1 -> range -> TX boundary s ----
6106 for s in 1..n_st - 1 {
6107 let _st = rt.enter(s);
6108 let es = rt.engine(s, e);
6109 let pos_d = stage_pos(es)?;
6110 let x = rt.rx(s - 1, slot, payload)?;
6111 let x = self.verify_layers(
6112 es,
6113 x,
6114 fence[s],
6115 fence[s + 1],
6116 &pos_d,
6117 pos0,
6118 t,
6119 cache,
6120 ckpt.as_deref_mut(),
6121 stream,
6122 None,
6123 )?;
6124 slot = if pipelined {
6125 rt.tx_pipelined(s, &x, payload)?
6126 } else {
6127 rt.tx(s, &x, payload)?
6128 };
6129 }
6130
6131 // ---- LAST STAGE: RX + final range + output_norm + lm head ----
6132 let _stl = rt.enter(n_st - 1);
6133 let el = rt.engine(n_st - 1, e);
6134 let pos_d = stage_pos(el)?;
6135 let rx_started = std::time::Instant::now();
6136 let x = rt.rx(n_st - 2, slot, payload)?;
6137 if pp_anatomy {
6138 el.stream().synchronize()?;
6139 rx_ms = rx_started.elapsed().as_secs_f64() * 1e3;
6140 }
6141 enqueue_spec_pipe_trace_marker(&el.stream(), trace.as_ref(), "S1", "start", Some(slot))?;
6142 let stage1_started = std::time::Instant::now();
6143 let x = self.verify_layers(
6144 el,
6145 x,
6146 fence[n_st - 1],
6147 fence[n_st],
6148 &pos_d,
6149 pos0,
6150 t,
6151 cache,
6152 ckpt.as_deref_mut(),
6153 stream,
6154 None,
6155 )?;
6156
6157 let mut hn = vbuf(el, payload)?;
6158 let logits = if self.sliding_gated_moe_batch_program() {
6159 // The PP Step35 serving path uses rms_norm + matmul for B=1 as well as B>1.
6160 // Verify must not switch numeric class merely because the same session speculates.
6161 el.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
6162 el.matmul(&self.output, &hn, t)?
6163 } else {
6164 el.rms_norm_decode(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
6165 el.matmul_decode_exact(&self.output, &hn, t)?
6166 };
6167 enqueue_spec_pipe_trace_marker(&el.stream(), trace.as_ref(), "S1", "end", Some(slot))?;
6168 if pp_anatomy {
6169 el.stream().synchronize()?;
6170 stage1_ms = stage1_started.elapsed().as_secs_f64() * 1e3;
6171 }
6172 // EXIT PUBLICATION: both returned buffers are still being produced on the last stage's
6173 // stream. Order the caller's stream behind that work before the buffers escape this
6174 // scope (the 2026-08-06 same-device ppspec find: without it the caller's primary-stream
6175 // consumer read unwritten logits — nondeterministic, one-device-only, and it poisoned
6176 // the following arm's KV in the same process).
6177 if publish_to_caller {
6178 rt.publish_to(n_st - 1, &caller_stream)?;
6179 }
6180 if pp_anatomy {
6181 if publish_to_caller {
6182 caller_stream.synchronize()?;
6183 }
6184 eprintln!(
6185 "[spec-pp-anatomy] t={t} reverse={reverse_ms:.3}ms stage0={stage0_ms:.3}ms \
6186 tx={tx_ms:.3}ms rx={rx_ms:.3}ms stage1-head={stage1_ms:.3}ms total={:.3}ms",
6187 pp_started.elapsed().as_secs_f64() * 1e3,
6188 );
6189 }
6190 // stream: the device pos counter owns position; host mirror reconciles at drain.
6191 if stream.is_none() {
6192 cache.pos += t;
6193 }
6194 Ok((logits, if spec_hpost() { hn } else { x }))
6195 }
6196
6197 /// Step3.5/Step3.7 verify trunk in the serving batched numeric class.
6198 ///
6199 /// `step35_decode_batch_layers` is now authoritative at every live serving width, including
6200 /// B=1 (lane/cx-b1fix). The older verify walk deliberately mirrored the eager T=1 class:
6201 /// it replayed `step35_decode_attn` per row and used the eager/decode-exact FFN dispatch.
6202 /// Those classes are individually stable, but a near-tie prompt can choose different greedy
6203 /// bytes when a request moves from batched plain serving into speculative verify. Run the
6204 /// same authoritative B=1 stage subgraph for each verify row here. Rows still advance
6205 /// layer-by-layer, so every layer sees the preceding verify rows in its attention cache while
6206 /// every norm/projection/FFN uses exactly the live serving dispatch.
6207 #[allow(clippy::too_many_arguments)]
6208 /// PRIME-BY-T-ROWS (MEMRA_PRIME_TROWS=1): prefill the prompt through the same-session
6209 /// t-row walk in 32-row chunks — every row runs the t=1 decode program bit-for-bit
6210 /// (the TOKENWISE-prime ORACLE class), so this door is exact against the exactness
6211 /// reference while replacing the host-canonical per-token prime. Requires the walk
6212 /// doors (MEMRA_SPEC_VERIFY_EAGER/TCOL); returns the prime contract trio.
6213 #[allow(clippy::type_complexity)]
6214 pub(crate) fn step35_prime_trows(
6215 &self,
6216 e: &Engine,
6217 tokens: &[u32],
6218 cache: &mut Cache,
6219 ) -> Result<Option<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
6220 {
6221 let dbg = std::env::var("MEMRA_SPEC_FA2_DEBUG").as_deref() == Ok("1");
6222 if !prime_trows_on() {
6223 return Ok(None);
6224 }
6225 if !self.uses_sliding_gated_moe_program()
6226 || cache.pos != 0
6227 || cache.dflash_taps.is_some()
6228 || !spec_verify_eager_on()
6229 || !spec_verify_tcol_on()
6230 {
6231 if dbg {
6232 eprintln!(
6233 "[prime-trows] refuse: program={} pos={} taps={} eager={:?} tcol={:?}",
6234 self.uses_sliding_gated_moe_program(),
6235 cache.pos,
6236 cache.dflash_taps.is_some(),
6237 std::env::var("MEMRA_SPEC_VERIFY_EAGER").ok(),
6238 std::env::var("MEMRA_SPEC_VERIFY_TCOL").ok()
6239 );
6240 }
6241 return Ok(None);
6242 }
6243 let n_embd = self.cfg.n_embd as usize;
6244 let n_layers = self.layers.len();
6245 let t_total = tokens.len();
6246 let Some(embd_gpu) = self.embd_gpu_try(e) else {
6247 if dbg {
6248 eprintln!("[prime-trows] refuse: no device embed table");
6249 }
6250 return Ok(None);
6251 };
6252 let embd_qtype = match self.embd.ggml_type {
6253 memra_gguf::GgmlType::BF16 => crate::QT_BF16,
6254 memra_gguf::GgmlType::Q8_0 => crate::QT_Q8_0,
6255 other => {
6256 if dbg {
6257 eprintln!("[prime-trows] refuse: embed dtype {other:?}");
6258 }
6259 return Ok(None);
6260 }
6261 };
6262 let embd_row_bytes = self.embd.raw.len() / self.cfg.n_vocab as usize;
6263 // Chunk plan: 32-row chunks; a 1-token tail folds into the previous chunk
6264 // (the walk floor is t >= 2).
6265 let mut bounds = Vec::new();
6266 let mut start = 0usize;
6267 while start < t_total {
6268 let mut end = (start + 32).min(t_total);
6269 if t_total - end == 1 {
6270 end -= 1;
6271 }
6272 bounds.push((start, end));
6273 start = end;
6274 }
6275 if bounds.iter().any(|(a, b)| b - a < 2) {
6276 return Ok(None); // degenerate short prompt keeps the ordinary prime
6277 }
6278 let mut hiddens = e.uninit(t_total * n_embd)?;
6279 let mut last: Option<CudaSlice<f32>> = None;
6280 for &(a, b) in &bounds {
6281 let tc = b - a;
6282 let tok_d = e.stream().clone_htod(&tokens[a..b])?;
6283 let x =
6284 e.embed_gather_device_td(embd_gpu, &tok_d, tc, n_embd, embd_qtype, embd_row_bytes)?;
6285 let out = self.step35_verify_batch_layers(e, x, 0, n_layers, a, tc, cache)?;
6286 e.copy_into(&mut hiddens, a * n_embd, &out, tc * n_embd)?;
6287 if b == t_total {
6288 let mut h = e.uninit(n_embd)?;
6289 e.dtod_copy_view(&out.slice((tc - 1) * n_embd..tc * n_embd), &mut h)?;
6290 last = Some(h);
6291 }
6292 }
6293 let h_seed = last.expect("last chunk produced the seed row");
6294 let mut hn = e.uninit(n_embd)?;
6295 e.rms_norm_decode(
6296 &h_seed,
6297 self.output_norm.float_data(),
6298 &mut hn,
6299 n_embd,
6300 1,
6301 self.cfg.rms_eps,
6302 )?;
6303 let logits_d = e.matmul_decode_exact(&self.output, &hn, 1)?;
6304 let logits = e.dtoh(&logits_d)?;
6305 cache.pos = t_total;
6306 Ok(Some((logits, h_seed, hiddens)))
6307 }
6308
6309 fn step35_verify_batch_layers(
6310 &self,
6311 e: &Engine,
6312 mut x: CudaSlice<f32>,
6313 lo: usize,
6314 hi: usize,
6315 pos0: usize,
6316 t: usize,
6317 cache: &mut Cache,
6318 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6319 let n_embd = self.cfg.n_embd as usize;
6320 if !self.uses_sliding_gated_moe_program() {
6321 return Err(
6322 "serving-class verify requires sliding-gated-MoE canonical operations".into(),
6323 );
6324 }
6325 // SERVING-CLASS VERIFY (MEMRA_SPEC_VERIFY_EAGER=1, step37 MTP bring-up): each verify
6326 // column rides decode_layers_eager — the EXACT t=1 program live serving runs (all TP2
6327 // doors) — row-outer, so row r's appends land before row r+1 attends: bit-equal to
6328 // plain greedy by construction. Only the unsplit full-range walk qualifies; PP splits
6329 // and the tap path keep the batch-layer class.
6330 static VE: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
6331 let eager_verify =
6332 *VE.get_or_init(spec_verify_eager_on) && lo == 0 && hi == self.layers.len();
6333 if eager_verify {
6334 // T-COLUMN LAYER-OUTER WALK (MEMRA_SPEC_VERIFY_TCOL=1): per layer, one t-grid
6335 // attn norm + ONE weight-amortized QKV(+gate) over all T columns, then each
6336 // column runs the UNMODIFIED t=1 attention program via the col-select door and
6337 // the ordinary residual/FFN body. Values per column are bit-equal to the
6338 // row-outer walk: rms over the materialized residual == the fused add+norm
6339 // (kernel_check identity), the tcol kernel's per-column FP order == the t=1
6340 // kernel, and every downstream op IS the t=1 program.
6341 static TCOL: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
6342 let tcol = *TCOL.get_or_init(spec_verify_tcol_on);
6343 // T > 32 (prefill-class): run the SAME walk in 32-row chunks — each chunk's
6344 // rows are the t=1 program bit-for-bit and the rope pass advances the cache,
6345 // so a chunked call is value-identical to the row-outer loop it replaces.
6346 static TROWS_PREFILL: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
6347 // MEMRA_STEP_GEMM_PRIME outranks the walk: with the grouped GEMM prime armed, the
6348 // t-row walk defers so the batch path (GEMM trunk + grouped MoE) takes the prompt —
6349 // flag precedence between two existing doors, not a new flag. Without this, both
6350 // doors ON meant the walk still won and the GEMM prime needed PRIME_TROWS=0 by hand.
6351 let trows_prefill =
6352 *TROWS_PREFILL.get_or_init(|| prime_trows_on() && !crate::step_gemm_prime_on());
6353 // MEMRA_PRIME_TROWS_T=<w>: chunk width, default 8 = the REAL cap of this walk.
6354 // The workspace slabs go to 32 rows, but `matvec_bf16_qkvg_tcol_into` refuses
6355 // t > 8 (compile-time-T twins exist for 2/4/8 only; the runtime-t kernel spills
6356 // its accumulators to local memory), so a wider chunk fails the request with
6357 // "matvec_bf16_qkvg_tcol geometry" — which is exactly how the first server-path
6358 // TROWS arm died. Measured at 193 tokens: w=8 2.459 s, w=4 2.574 s.
6359 static TROWS_W: std::sync::OnceLock<Result<usize, String>> = std::sync::OnceLock::new();
6360 let trows_w = match TROWS_W.get_or_init(|| {
6361 let value = std::env::var("MEMRA_PRIME_TROWS_T").ok();
6362 parse_prime_trows_width(value.as_deref())
6363 }) {
6364 Ok(width) => *width,
6365 Err(err) => return Err(err.clone().into()),
6366 };
6367 if tcol && trows_prefill && t > trows_w {
6368 // One-time engagement receipt: without it a prefill gate cannot tell a
6369 // chunked walk from the row-outer fallback it is supposed to replace
6370 // (the first PRIME_TROWS gate passed vacuously on exactly that).
6371 static SEEN: std::sync::atomic::AtomicBool =
6372 std::sync::atomic::AtomicBool::new(false);
6373 if !SEEN.swap(true, std::sync::atomic::Ordering::Relaxed) {
6374 eprintln!(
6375 "[prime-trows] ENGAGED t={t} width={trows_w} chunks={} layers={}..{}",
6376 t.div_ceil(trows_w),
6377 lo,
6378 hi
6379 );
6380 }
6381 let mut out = e.uninit(t * n_embd)?;
6382 let mut start = 0usize;
6383 while start < t {
6384 let mut end = (start + trows_w).min(t);
6385 if t - end == 1 {
6386 end -= 1;
6387 }
6388 let tc = end - start;
6389 let mut xc = e.uninit(tc * n_embd)?;
6390 e.dtod_copy_view(&x.slice(start * n_embd..end * n_embd), &mut xc)?;
6391 let oc =
6392 self.step35_verify_batch_layers(e, xc, lo, hi, pos0 + start, tc, cache)?;
6393 e.copy_into(&mut out, start * n_embd, &oc, tc * n_embd)?;
6394 start = end;
6395 }
6396 return Ok(out);
6397 }
6398 if tcol && t >= 2 && t <= 32 {
6399 // MEMRA_TCOL_PROF=1: synchronized per-segment wall profile of the walk
6400 // (norm+QKV precompute / per-col attention / per-col residual+FFN). The
6401 // syncs serialize the stream, so the split is for TARGETING amortization
6402 // work only — never a perf claim.
6403 static PROF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
6404 let prof =
6405 *PROF.get_or_init(|| std::env::var("MEMRA_TCOL_PROF").as_deref() == Ok("1"));
6406 let mut prof_ms = [0f64; 3];
6407 let eps = self.cfg.rms_eps;
6408 let mut x_t = x;
6409 let mut h_t = e.uninit(t * n_embd)?;
6410 let mut h_row = e.uninit(n_embd)?; // real row: the non-dcw fallback reads it
6411 // Per-column pos buffers hoisted out of the layer loop (a per-col-per-layer
6412 // pageable htod was an in-stream engine turnaround x t x 45).
6413 let mut pos_rows = Vec::with_capacity(t);
6414 for r in 0..t {
6415 pos_rows.push(e.htod_i32(&[(pos0 + r) as i32])?);
6416 }
6417 let mut ok = true;
6418 // MEMRA_TCOL_OPROJ=1: defer each column's o_proj — the finish seam
6419 // stashes `gated` instead of joining per column; one b4_tcol per rank +
6420 // one slab join produce every column's `mixed` after the attention pass.
6421 // Bit-exact per column (t=1 b4 program per column; elementwise join).
6422 // MEMRA_TCOL_FFN=1: today this only IMPLIES the o_proj defer above. Its
6423 // named feature, the two-column device-routed FFN sweep, rode the
6424 // slot-major v2 TP banks and was REMOVED with the MEMRA_NVFP4_BANK_V2 door
6425 // (2026-08-29, research/step37-bankv2-removal-20260829): the v2 layout
6426 // changed generated text in serving. The flag itself stays because it is
6427 // family-armed in the step37 serving defaults and killing it here would
6428 // silently drop the o_proj defer from the qualified serving shape.
6429 static FFN2: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
6430 let ffn_batch = *FFN2.get_or_init(tcol_ffn_on);
6431 let oproj_batch = crate::tp::tcol_oproj_on() || ffn_batch;
6432 // MEMRA_SPEC_FA2=1 (T=2 only): eligible layers defer BOTH columns' fa —
6433 // the per-column pass norms/ropes/appends and stashes q+gate, then one
6434 // shared-KV fa_decode_dcw2 per rank + the o_proj join produce the
6435 // [2, o_out] mixed slab. The precheck runs before arming (stashing is
6436 // unrecoverable); ineligible/boundary layers run the ordinary program.
6437 let fa2 = crate::tp::spec_fa2_on() && t <= 32;
6438 let mut mixed_row = e.uninit(n_embd)?;
6439 let mut pos_staged = false;
6440 for il in lo..hi {
6441 let layer = &self.layers[il];
6442 // BEFORE this layer touches its planes: is the history it is about to
6443 // attend already poisoned? Global (non-ring) layers only, which are the
6444 // ones the level-2 bitmap implicates.
6445 if kv_plane_scan_on() && self.step35_geom(il).window.is_none() {
6446 if let Some(distributed) = cache.tp_kv[il].as_ref() {
6447 scan_kv_plane(e, distributed, il, pos0)?;
6448 }
6449 }
6450 let fa2_layer = fa2 && self.step35_fa_rows_precheck(cache, il, pos0, t)?;
6451 let mut seg = std::time::Instant::now();
6452 e.rms_norm(&x_t, layer.attn_norm.float_data(), &mut h_t, n_embd, t, eps)?;
6453 if !self.step35_verify_qkv_precompute(e, il, &h_t, t)? {
6454 ok = false;
6455 break;
6456 }
6457 // FULL t-row attention pass (rope/append + fa + combine + o_proj in
6458 // 3 launches/rank): same-session rows, slot = len-base+r, one len
6459 // advance by t. Host cache bookkeeping mirrors the per-column tail.
6460 if fa2_layer {
6461 if let Some(mixed_t) =
6462 self.step35_verify_rope_fa_pass(e, il, cache, pos0, t, !pos_staged)?
6463 {
6464 pos_staged = true;
6465 {
6466 let tp_kv = cache.tp_kv[il]
6467 .as_mut()
6468 .expect("precheck verified the distributed cache");
6469 let transaction = tp_kv.begin_transaction()?;
6470 let crate::hybrid::Mixer::Full(fa) = &layer.mixer else {
6471 return Err("verify rope pass expects full attention".into());
6472 };
6473 let tp = fa
6474 .step_tp_qkv
6475 .as_ref()
6476 .ok_or("verify rope pass lost its TP state")?;
6477 let empty: [CudaSlice<f32>; 0] = [];
6478 tp.runtime.append_tp_kv_transaction_inner(
6479 tp_kv,
6480 transaction,
6481 &empty,
6482 &empty,
6483 t,
6484 true,
6485 )?;
6486 tp.runtime.commit_tp_kv_transaction_external(
6487 tp_kv,
6488 transaction,
6489 t,
6490 )?;
6491 if let Some(local) = cache.kv[il].as_mut() {
6492 local.len = pos0 + t;
6493 if !crate::tp::len_mirror_lazy_on() {
6494 e.set_i32_one(&mut local.len_d, local.len as i32)?;
6495 }
6496 }
6497 }
6498 if prof {
6499 e.stream().synchronize()?;
6500 prof_ms[1] += seg.elapsed().as_secs_f64() * 1e3;
6501 seg = std::time::Instant::now();
6502 }
6503 let o_out = mixed_t.len() / t;
6504 let mut next = e.uninit(t * n_embd)?;
6505 {
6506 for r in 0..t {
6507 e.dtod_copy_view(
6508 &mixed_t.slice(r * o_out..(r + 1) * o_out),
6509 &mut mixed_row,
6510 )?;
6511 let mut x_row = e.uninit(n_embd)?;
6512 e.dtod_copy_view(
6513 &x_t.slice(r * n_embd..(r + 1) * n_embd),
6514 &mut x_row,
6515 )?;
6516 let (x1, ffn_out) = self.residual_norm_ffn(
6517 e, layer, &x_row, &mixed_row, n_embd, il, eps,
6518 )?;
6519 let mut x2 = e.uninit(n_embd)?;
6520 e.add(&x1, &ffn_out, &mut x2, n_embd)?;
6521 e.dtod_copy_into(&x2, &mut next, r * n_embd)?;
6522 }
6523 }
6524 if prof {
6525 e.stream().synchronize()?;
6526 prof_ms[2] += seg.elapsed().as_secs_f64() * 1e3;
6527 }
6528 x_t = next;
6529 if spec_nan_scan() {
6530 // The scan MUST sit on this arm too. It used to live only on
6531 // the non-fused tail, so a fused layer's poison was first
6532 // reported by the next non-fused layer.
6533 verify_arm_receipt(
6534 "fused",
6535 il,
6536 pos0,
6537 t,
6538 cache.tp_kv[il].as_ref().map(|d| d.staged_len()),
6539 );
6540 nan_scan_rows(
6541 e,
6542 &x_t,
6543 t,
6544 n_embd,
6545 &format!("tcol layer {il} pos0={pos0} arm=fused"),
6546 )?;
6547 }
6548 continue;
6549 }
6550 }
6551 if prof {
6552 e.stream().synchronize()?;
6553 prof_ms[0] += seg.elapsed().as_secs_f64() * 1e3;
6554 seg = std::time::Instant::now();
6555 }
6556 let mut next = e.uninit(t * n_embd)?;
6557 // Columns whose o_proj was deferred (their FFN runs after the join).
6558 // A NON-deferred column's FFN must run INSIDE the column loop: the
6559 // oproj-tail handoff is a single cell that the same column's
6560 // residual_norm_ffn consumes before the next column's finish.
6561 let mut deferred: Vec<usize> = Vec::new();
6562 let mut fa2_deferred: Vec<usize> = Vec::new();
6563 let mut ffn_col =
6564 |r: usize,
6565 mixed: &CudaSlice<f32>,
6566 next: &mut CudaSlice<f32>|
6567 -> Result<(), Box<dyn std::error::Error>> {
6568 let mut x_row = e.uninit(n_embd)?;
6569 e.dtod_copy_view(&x_t.slice(r * n_embd..(r + 1) * n_embd), &mut x_row)?;
6570 let (x1, ffn_out) =
6571 self.residual_norm_ffn(e, layer, &x_row, mixed, n_embd, il, eps)?;
6572 if spec_nan_scan_level() >= 2 {
6573 nan_scan_rows(
6574 e,
6575 &ffn_out,
6576 1,
6577 n_embd,
6578 &format!("tcol layer {il} col {r} per-column FFN out"),
6579 )?;
6580 }
6581 let mut x2 = e.uninit(n_embd)?;
6582 e.add(&x1, &ffn_out, &mut x2, n_embd)?;
6583 e.dtod_copy_into(&x2, next, r * n_embd)?;
6584 Ok(())
6585 };
6586 for r in 0..t {
6587 e.dtod_copy_view(&h_t.slice(r * n_embd..(r + 1) * n_embd), &mut h_row)?;
6588 let row_pos = &pos_rows[r];
6589 crate::tp::set_verify_tcol(Some(r));
6590 if fa2_layer {
6591 crate::tp::set_spec_fa2_defer(Some(r));
6592 } else if oproj_batch {
6593 crate::tp::set_tcol_oproj_defer(Some(r));
6594 }
6595 let mixed = match &layer.mixer {
6596 crate::hybrid::Mixer::Full(fa) => {
6597 self.full_attn_decode(e, fa, &h_row, row_pos, pos0 + r, cache, il)
6598 }
6599 _ => Err("step35 verify expects full attention".into()),
6600 };
6601 crate::tp::set_verify_tcol(None);
6602 crate::tp::set_spec_fa2_defer(None);
6603 crate::tp::set_tcol_oproj_defer(None);
6604 let mixed = mixed?;
6605 if fa2_layer && crate::tp::take_spec_fa2_stashed() {
6606 fa2_deferred.push(r);
6607 } else if oproj_batch && crate::tp::take_tcol_oproj_stashed() {
6608 deferred.push(r);
6609 } else {
6610 if spec_nan_scan_level() >= 2 {
6611 let cols = mixed.len();
6612 nan_scan_rows(
6613 e,
6614 &mixed,
6615 1,
6616 cols,
6617 &format!("tcol layer {il} col {r} per-column ATTN out"),
6618 )?;
6619 }
6620 ffn_col(r, &mixed, &mut next)?;
6621 }
6622 }
6623 if !fa2_deferred.is_empty() && fa2_deferred.len() != t {
6624 // The precheck guarantees both columns stash or neither; a strict
6625 // subset means a column's output was never produced anywhere.
6626 return Err("spec fa2 stash engaged for a subset of columns".into());
6627 }
6628 if prof {
6629 e.stream().synchronize()?;
6630 prof_ms[1] += seg.elapsed().as_secs_f64() * 1e3;
6631 seg = std::time::Instant::now();
6632 }
6633 if !fa2_deferred.is_empty() {
6634 deferred = fa2_deferred;
6635 }
6636 if !deferred.is_empty() {
6637 let mixed_t = if fa2_layer {
6638 self.step35_verify_fa_rows_join(e, il, cache, pos0, t)?
6639 } else {
6640 self.step35_verify_oproj_tcol(e, il, t)?
6641 };
6642 let o_out = mixed_t.len() / t;
6643 if spec_nan_scan_level() >= 2 {
6644 nan_scan_rows(
6645 e,
6646 &mixed_t,
6647 t,
6648 o_out,
6649 &format!("tcol layer {il} JOINED attn over deferred cols"),
6650 )?;
6651 }
6652 // Batched t=2 residual+MoE: one t-grid add_rms_norm (per-row
6653 // program == t=1; bit-identical to the oproj-tail join per the
6654 // M2 verbatim-program contract) feeding the two-column routed
6655 // sweep. Ineligible layers (dense FFN, non-nvfp4) fall through
6656 // to the per-column body.
6657 {
6658 for &r in &deferred {
6659 e.dtod_copy_view(
6660 &mixed_t.slice(r * o_out..(r + 1) * o_out),
6661 &mut mixed_row,
6662 )?;
6663 ffn_col(r, &mixed_row, &mut next)?;
6664 }
6665 }
6666 }
6667 if prof {
6668 e.stream().synchronize()?;
6669 prof_ms[2] += seg.elapsed().as_secs_f64() * 1e3;
6670 }
6671 drop(ffn_col);
6672 x_t = next;
6673 if spec_nan_scan() {
6674 verify_arm_receipt(
6675 if fa2_layer { "join" } else { "percol" },
6676 il,
6677 pos0,
6678 t,
6679 cache.tp_kv[il].as_ref().map(|d| d.staged_len()),
6680 );
6681 nan_scan_rows(
6682 e,
6683 &x_t,
6684 t,
6685 n_embd,
6686 &format!(
6687 "tcol layer {il} pos0={pos0} arm={}",
6688 if fa2_layer { "join" } else { "percol" }
6689 ),
6690 )?;
6691 }
6692 }
6693 if prof {
6694 eprintln!(
6695 "[tcol-prof] t={t} norm+qkv={:.3}ms attn={:.3}ms ffn={:.3}ms",
6696 prof_ms[0], prof_ms[1], prof_ms[2]
6697 );
6698 }
6699 if ok {
6700 return Ok(x_t);
6701 }
6702 // fall through to the row-outer walk on ineligible layers
6703 x = x_t;
6704 }
6705 let mut next = e.uninit(t * n_embd)?;
6706 let scan = spec_nan_scan();
6707 for r in 0..t {
6708 let mut row = e.uninit(n_embd)?;
6709 e.dtod_copy_view(&x.slice(r * n_embd..(r + 1) * n_embd), &mut row)?;
6710 let row_pos = e.htod_i32(&[(pos0 + r) as i32])?;
6711 let out = if scan {
6712 // Diagnostic arm: the same range walked one layer at a time so the first
6713 // poisoned layer names itself. `decode_layers_eager(lo, hi)` is range-scoped
6714 // and executes its trailing residual add, so a per-layer chain is the same
6715 // program with the cross-layer add+norm fusion unrolled.
6716 nan_scan_rows(
6717 e,
6718 &row,
6719 1,
6720 n_embd,
6721 &format!("embed row r={r} pos={}", pos0 + r),
6722 )?;
6723 let mut acc = row;
6724 for il in lo..hi {
6725 acc = self.decode_layers_eager(
6726 e,
6727 acc,
6728 il,
6729 il + 1,
6730 &row_pos,
6731 pos0 + r,
6732 cache,
6733 )?;
6734 nan_scan_rows(
6735 e,
6736 &acc,
6737 1,
6738 n_embd,
6739 &format!("row-outer layer {il} r={r} pos={}", pos0 + r),
6740 )?;
6741 }
6742 acc
6743 } else {
6744 self.decode_layers_eager(e, row, lo, hi, &row_pos, pos0 + r, cache)?
6745 };
6746 e.dtod_copy_into(&out, &mut next, r * n_embd)?;
6747 }
6748 // dflash taps are NOT produced on this arm (they need per-layer hiddens the
6749 // row-outer walk does not materialize); the door is a step37 MTP bring-up
6750 // surface where taps are unused.
6751 return Ok(next);
6752 }
6753 let mut ph_last = std::time::Instant::now();
6754 for il in lo..hi {
6755 let mut next = e.uninit(t * n_embd)?;
6756 for r in 0..t {
6757 let mut row = e.uninit(n_embd)?;
6758 e.dtod_copy_view(&x.slice(r * n_embd..(r + 1) * n_embd), &mut row)?;
6759 // The caller owns this verify's position. During controller overlap, cache.pos
6760 // still describes generation N while this stage-0 walk belongs to N+1.
6761 let row_pos = e.htod_i32(&[(pos0 + r) as i32])?;
6762 let mut one = [&mut *cache];
6763 let out = self.step35_decode_batch_layers(
6764 e,
6765 row,
6766 &mut one,
6767 &[(pos0 + r) as i32],
6768 &row_pos,
6769 il,
6770 il + 1,
6771 &mut ph_last,
6772 )?;
6773 e.dtod_copy_into(&out, &mut next, r * n_embd)?;
6774 }
6775 self.dflash_tap(e, cache, il, &next, t)?;
6776 x = next;
6777 if spec_nan_scan() {
6778 nan_scan_rows(e, &x, t, n_embd, &format!("batch-layer {il} pos0={pos0}"))?;
6779 }
6780 }
6781 Ok(x)
6782 }
6783
6784 /// DSpark drafter verify (lane/dspark-q38-recover): one t-row forward through the
6785 /// SERVING-CLASS verify funnel (`decode_step_t_core_stream` — the same numeric class
6786 /// MTP verify rides, GDN state advanced in place), returning per-row argmax tokens.
6787 /// Advances `cache.pos += t`; the caller owns snapshot/rollback (block acceptance is
6788 /// prefix-keep, not all-or-nothing).
6789 pub(crate) fn dspark_verify_t_am(
6790 &self,
6791 e: &Engine,
6792 tokens: &[u32],
6793 pos0: usize,
6794 cache: &mut Cache,
6795 ) -> Result<Vec<u32>, Box<dyn std::error::Error>> {
6796 let (logits, _hn) = self.decode_step_t_core_stream(
6797 e, tokens, pos0, cache, None, None, None, None, None, None,
6798 )?;
6799 let t = tokens.len();
6800 let v = self.output.out_features();
6801 let mut am_d = e.stream().alloc_zeros::<u32>(t)?;
6802 for r in 0..t {
6803 e.argmax_token_device_col(&logits, r, v, &mut am_d, r)?;
6804 }
6805 Ok(e.dtoh_u32(&am_d)?)
6806 }
6807
6808 /// DSpark verify returning the RAW verify logits [t, n_vocab] (device-resident) instead
6809 /// of per-row argmaxes — the sampled-admission arm's input (rejection-sampling accept
6810 /// gathers filtered p from these columns; lane/dspark-sampled-admission-20260820). Same
6811 /// forward as `dspark_verify_t_am`; the greedy arm keeps its argmax wrapper untouched.
6812 pub(crate) fn dspark_verify_t_logits(
6813 &self,
6814 e: &Engine,
6815 tokens: &[u32],
6816 pos0: usize,
6817 cache: &mut Cache,
6818 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6819 let (logits, _hn) = self.decode_step_t_core_stream(
6820 e, tokens, pos0, cache, None, None, None, None, None, None,
6821 )?;
6822 Ok(logits)
6823 }
6824
6825 /// DSpark verify with the MTP column-stash armed: identical forward to
6826 /// `dspark_verify_t_am`, but fills a `VerifyCkpt` so a partial accept can restore
6827 /// column state directly (`dspark_commit_prefix`) instead of snapshot-replay.
6828 /// The ckpt type is opaque outside spec.rs (newtype) — dflash.rs threads it through.
6829 pub(crate) fn dspark_verify_t_am_ckpt(
6830 &self,
6831 e: &Engine,
6832 tokens: &[u32],
6833 pos0: usize,
6834 cache: &mut Cache,
6835 ) -> Result<(Vec<u32>, DsparkVerifyCkpt), Box<dyn std::error::Error>> {
6836 let mut ck = VerifyCkpt::new(self.layers.len());
6837 let (logits, _hn) = self.decode_step_t_core_stream(
6838 e,
6839 tokens,
6840 pos0,
6841 cache,
6842 None,
6843 Some(&mut ck),
6844 None,
6845 None,
6846 None,
6847 None,
6848 )?;
6849 let t = tokens.len();
6850 let v = self.output.out_features();
6851 let mut am_d = e.stream().alloc_zeros::<u32>(t)?;
6852 for r in 0..t {
6853 e.argmax_token_device_col(&logits, r, v, &mut am_d, r)?;
6854 }
6855 Ok((e.dtoh_u32(&am_d)?, DsparkVerifyCkpt(ck)))
6856 }
6857
6858 /// Engine-bundle slice 2: `dspark_verify_t_am_ckpt` with DEVICE tokens and NO readback.
6859 /// The verify tokens are the round's `chain_d` (cand layout: [anchor, drafts...]); the
6860 /// embed gathers its first `t` entries on-device (`embed_gather_u32_t` — bit-identical
6861 /// rows to the host arm), so the host never blocks on the draft chain before dispatching
6862 /// verify. Returns the device per-row argmax buffer; the caller merges its readback with
6863 /// the chain's into ONE sync. Forward, ckpt fill and argmax walk are `_ckpt` verbatim.
6864 pub(crate) fn dspark_verify_t_am_ckpt_dev(
6865 &self,
6866 e: &Engine,
6867 vtok: &CudaSlice<u32>,
6868 t: usize,
6869 pos0: usize,
6870 cache: &mut Cache,
6871 embd_dev: (&CudaSlice<u8>, i32, usize),
6872 graphs: Option<&mut DsparkVerifyGraphs>,
6873 ) -> Result<(CudaSlice<u32>, DsparkVerifyCkpt), Box<dyn std::error::Error>> {
6874 debug_assert!(
6875 vtok.len() >= t,
6876 "verify window exceeds the device token buffer"
6877 );
6878 // The slab flag is a per-round statement: clear it here so a verify that never
6879 // reaches the graphs door (rowwise env, a non-tparallel arm) cannot leave a
6880 // stale `true` steering the commit at slabs the round never wrote.
6881 let mut graphs = graphs;
6882 if let Some(g) = graphs.as_deref_mut() {
6883 g.round_slab = false;
6884 }
6885 let mut ck = VerifyCkpt::new(self.layers.len());
6886 // Dummy host tokens size the funnel; the embed reads `vtok` (the round-stream
6887 // arm's established pattern — spec.rs stream-mode verify does the same).
6888 let dummy = vec![0u32; t];
6889 let (logits, _hn) = self.decode_step_t_core_stream(
6890 e,
6891 &dummy,
6892 pos0,
6893 cache,
6894 Some(embd_dev),
6895 Some(&mut ck),
6896 None,
6897 None,
6898 Some(vtok),
6899 graphs,
6900 )?;
6901 let v = self.output.out_features();
6902 let mut am_d = e.stream().alloc_zeros::<u32>(t)?;
6903 for r in 0..t {
6904 e.argmax_token_device_col(&logits, r, v, &mut am_d, r)?;
6905 }
6906 Ok((am_d, DsparkVerifyCkpt(ck)))
6907 }
6908
6909 /// Ckpt-armed twin of [`Self::dspark_verify_t_logits`] (sampled-admission arm).
6910 pub(crate) fn dspark_verify_t_logits_ckpt(
6911 &self,
6912 e: &Engine,
6913 tokens: &[u32],
6914 pos0: usize,
6915 cache: &mut Cache,
6916 ) -> Result<(CudaSlice<f32>, DsparkVerifyCkpt), Box<dyn std::error::Error>> {
6917 let mut ck = VerifyCkpt::new(self.layers.len());
6918 let (logits, _hn) = self.decode_step_t_core_stream(
6919 e,
6920 tokens,
6921 pos0,
6922 cache,
6923 None,
6924 Some(&mut ck),
6925 None,
6926 None,
6927 None,
6928 None,
6929 )?;
6930 Ok((logits, DsparkVerifyCkpt(ck)))
6931 }
6932
6933 /// Restore the round to `keep` accepted columns from the verify stash: KV lens and
6934 /// pos from the pre-verify snapshot + keep, GDN conv/ssm from the stashed column
6935 /// state — no replay forward. The exact `commit_verified_prefix` the MTP path ships.
6936 pub(crate) fn dspark_commit_prefix(
6937 &self,
6938 e: &Engine,
6939 cache: &mut Cache,
6940 snap: &crate::cache::CacheSnapshot,
6941 ckpt: &DsparkVerifyCkpt,
6942 keep: usize,
6943 ) -> Result<(), Box<dyn std::error::Error>> {
6944 self.commit_verified_prefix(e, cache, snap, &ckpt.0, keep, false, None)
6945 }
6946
6947 /// Slice-3 commit twin: restore to `keep` accepted columns when the round's linear
6948 /// column stash lives in the graphs ctx's persistent slabs (`DsparkVerifyGraphs`) —
6949 /// the cols arm's exact semantics (KV lens + pos from the snapshot, GDN conv/ssm
6950 /// from the stash of column keep-1), slab-addressed and batched into two copy
6951 /// launches. `MEMRA_STATE_COPY_BATCH=0` falls back to per-layer view copies.
6952 pub(crate) fn dspark_commit_prefix_slab(
6953 &self,
6954 e: &Engine,
6955 cache: &mut Cache,
6956 snap: &crate::cache::CacheSnapshot,
6957 ctx: &DsparkVerifyGraphs,
6958 keep: usize,
6959 ) -> Result<(), Box<dyn std::error::Error>> {
6960 use cudarc::driver::DevicePtr;
6961 debug_assert!(keep >= 1, "keep==0 rounds take the legacy rollback");
6962 let mut conv_src: Vec<u64> = Vec::new();
6963 let mut ssm_src: Vec<u64> = Vec::new();
6964 let mut conv_dst: Vec<u64> = Vec::new();
6965 let mut ssm_dst: Vec<u64> = Vec::new();
6966 for il in 0..self.layers.len() {
6967 if let (Some(kvl), Some(saved)) = (cache.kv[il].as_mut(), snap.kv_len[il]) {
6968 kvl.len = saved + keep;
6969 e.set_i32_one(&mut kvl.len_d, kvl.len as i32)?;
6970 }
6971 if let Some(rl) = cache.recur[il].as_ref() {
6972 let (pc, ps, _cw, _sw) = ctx
6973 .slab_row(e, il, keep - 1)
6974 .ok_or("slab commit: linear layer missing from the graphs ctx")?;
6975 conv_src.push(pc);
6976 ssm_src.push(ps);
6977 let st = &e.gpu.stream();
6978 let (dc, _g0) = rl.conv_state.device_ptr(st);
6979 let (ds, _g1) = rl.ssm_state.device_ptr(st);
6980 conv_dst.push(dc as u64);
6981 ssm_dst.push(ds as u64);
6982 }
6983 }
6984 let n = conv_src.len();
6985 if n > 0 {
6986 if state_copy_batch_on() {
6987 let mut tt = vec![0u64; 2 * n];
6988 tt[..n].copy_from_slice(&conv_src);
6989 tt[n..].copy_from_slice(&conv_dst);
6990 let ct = e.htod_u64(&tt)?;
6991 tt[..n].copy_from_slice(&ssm_src);
6992 tt[n..].copy_from_slice(&ssm_dst);
6993 let st = e.htod_u64(&tt)?;
6994 e.copy_batch_uniform_f32(&ct, n, ctx.conv_words)?;
6995 e.copy_batch_uniform_f32(&st, n, ctx.ssm_words)?;
6996 } else {
6997 let (cw, sw) = (ctx.conv_words, ctx.ssm_words);
6998 let row = keep - 1;
6999 for il in 0..self.layers.len() {
7000 let Some(rl) = cache.recur[il].as_mut() else {
7001 continue;
7002 };
7003 let k = ctx.lin_pos[&il];
7004 {
7005 let sv = e.view(&ctx.stash_conv[k], (row + 1) * cw);
7006 let win = sv.slice(row * cw..(row + 1) * cw);
7007 e.copy_view_into(&mut rl.conv_state, 0, &win, cw)?;
7008 }
7009 {
7010 let sv = e.view(&ctx.stash_ssm[k], (row + 1) * sw);
7011 let win = sv.slice(row * sw..(row + 1) * sw);
7012 e.copy_view_into(&mut rl.ssm_state, 0, &win, sw)?;
7013 }
7014 }
7015 }
7016 }
7017 cache.pos = snap.pos + keep;
7018 Ok(())
7019 }
7020
7021 /// Qwen35-family verify trunk in the live serving numeric class.
7022 ///
7023 /// Serving intentionally keeps this architecture in the generic batched program even at
7024 /// B=1. The older verify walk used its own mirrored dispatch and can flip near-tie argmaxes.
7025 ///
7026 /// Two arms, one numeric class:
7027 /// - DENSE GDN (`DenseMlp`, t<=16): `qwen35_verify_tparallel` — the weight ops (norms,
7028 /// projections, FFN) hoist to m=T through the exact-tier batched kernels whose per-row
7029 /// program IS the m=1 program (`matmul_pre == fused2 per (tensor,row); _bN mmvq per-row
7030 /// == m=1` — decode_batch.rs v2 note), while the state ops (conv ring, gdn scan, KV
7031 /// append, fa decode) stay a per-row loop running the b_n=1 serving kernels with each
7032 /// row's own t_kv-driven arm pick (the straddle law: every row executes the exact
7033 /// program its isolated serving step would). One weight read per layer per round
7034 /// instead of T — this is what makes MTP profitable in the exact class (the per-row
7035 /// walk measured verify(K+1) ~= (K+1) plain steps: 69 -> 44 tok/s served, 2026-08-15).
7036 /// - MoE / t>16 / `MEMRA_SPEC_VERIFY_ROWWISE=1`: the per-row replay of the authoritative
7037 /// serving layer body, preserving single-session autoregressive cache order (the
7038 /// correctness reference; also the rollback seam for the t-parallel arm).
7039 ///
7040 /// Bit-identity of the t-parallel arm vs the rowwise arm is gated by spec-serve-gate
7041 /// (zero differing logits at T=1..4, K arms) + the 8-prompt ON/OFF canary before ship.
7042 #[allow(clippy::too_many_arguments)]
7043 fn qwen35_verify_batch_layers(
7044 &self,
7045 e: &Engine,
7046 x: CudaSlice<f32>,
7047 lo: usize,
7048 hi: usize,
7049 pos0: usize,
7050 t: usize,
7051 cache: &mut Cache,
7052 ckpt: Option<&mut VerifyCkpt>,
7053 stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
7054 graphs: Option<&mut DsparkVerifyGraphs>,
7055 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7056 // Qwen35Moe admitted 2026-08-20 (lane/draftcost-moe): the t-parallel arm already
7057 // carries the MoE FFN (`moe_ffn_il_zq8` at m=T) and the GDN per-row state loop; the
7058 // arch fence was a qualification gate, not a mechanism gap. Measured disease on the
7059 // 35B-A3B class: rowwise verify ~= 5.6 ms per drafted token (one full trunk step
7060 // each) — the same (K+1)-plain-steps wall the dense admission fixed on 2026-08-15.
7061 // Rollback seam unchanged: MEMRA_SPEC_VERIFY_ROWWISE=1.
7062 let rowwise = std::env::var("MEMRA_SPEC_VERIFY_ROWWISE").as_deref() == Ok("1")
7063 || !self.batched_serving_numeric_class()
7064 || t > 16;
7065 if rowwise {
7066 if stream.is_some() {
7067 // rowwise replays per row with host cache.pos — irreconcilable with a
7068 // device position counter. Burst callers must keep t <= 16 and the
7069 // ROWWISE env unset; refusing beats silently mispositioned rows.
7070 return Err("qwen35 rowwise verify has no ROUND-STREAM arm \
7071 (t > 16 or MEMRA_SPEC_VERIFY_ROWWISE=1)"
7072 .into());
7073 }
7074 self.qwen35_verify_rowwise(e, x, lo, hi, pos0, t, cache, ckpt)
7075 } else {
7076 self.qwen35_verify_tparallel(e, x, lo, hi, pos0, t, cache, ckpt, stream, graphs)
7077 }
7078 }
7079
7080 /// The per-row correctness reference: replay each verify row through the authoritative
7081 /// serving layer body (`decode_batch_layers` at b_n=1). T full weight reads per layer.
7082 #[allow(clippy::too_many_arguments)]
7083 fn qwen35_verify_rowwise(
7084 &self,
7085 e: &Engine,
7086 mut x: CudaSlice<f32>,
7087 lo: usize,
7088 hi: usize,
7089 pos0: usize,
7090 t: usize,
7091 cache: &mut Cache,
7092 mut ckpt: Option<&mut VerifyCkpt>,
7093 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7094 let n_embd = self.cfg.n_embd as usize;
7095 let saved_pos = cache.pos;
7096 let mut ph_last = std::time::Instant::now();
7097 for il in lo..hi {
7098 let mut next = e.uninit(t * n_embd)?;
7099 let mut col_states: Option<Vec<(CudaSlice<f32>, CudaSlice<f32>)>> =
7100 if ckpt.is_some() && t >= 2 && matches!(self.layers[il].mixer, Mixer::Linear(_)) {
7101 Some(Vec::with_capacity(t - 1))
7102 } else {
7103 None
7104 };
7105 for r in 0..t {
7106 cache.pos = pos0 + r;
7107 let mut row = e.uninit(n_embd)?;
7108 e.dtod_copy_view(&x.slice(r * n_embd..(r + 1) * n_embd), &mut row)?;
7109 let row_pos = e.htod_i32(&[(pos0 + r) as i32])?;
7110 let mut one = [&mut *cache];
7111 let ctx = self.batch_layer_ctx(e, &one, il, il + 1)?;
7112 let out = match self.decode_batch_layers(
7113 e,
7114 row,
7115 &mut one,
7116 &ctx,
7117 &row_pos,
7118 &mut ph_last,
7119 ) {
7120 Ok(out) => out,
7121 Err(error) => {
7122 cache.pos = saved_pos;
7123 return Err(error);
7124 }
7125 };
7126 e.dtod_copy_into(&out, &mut next, r * n_embd)?;
7127 if r + 1 < t {
7128 if let Some(states) = col_states.as_mut() {
7129 let recur = cache.recur[il]
7130 .as_ref()
7131 .ok_or("Qwen35-MoE linear verify layer has no recurrent state")?;
7132 states.push((
7133 e.clone_dtod(&recur.conv_state)?,
7134 e.clone_dtod(&recur.ssm_state)?,
7135 ));
7136 }
7137 }
7138 }
7139 if let (Some(checkpoint), Some(states)) = (ckpt.as_deref_mut(), col_states) {
7140 checkpoint.cols[il] = Some(states);
7141 }
7142 x = next;
7143 }
7144 cache.pos = saved_pos;
7145 Ok(x)
7146 }
7147
7148 /// T-PARALLEL VERIFY IN THE SERVING NUMERIC CLASS (lane/tparallel-verify, 2026-08-15).
7149 ///
7150 /// The weight ops run ONCE per layer at m=T; the state ops run per row through the same
7151 /// b_n=1 serving kernels the rowwise replay uses. Per-row bit-identity rests on the two
7152 /// pins the serving batch tier already carries:
7153 /// * `matmul_pre` / `_bN` mmvq: per-row program == m=1 program (decode_batch.rs v2 note,
7154 /// kernel-check pinned) — so a [T, n_embd] projection row equals the row projected
7155 /// alone;
7156 /// * row-indexed norms/elementwise (`rms_norm`, `quantize_q8_1`, `add_rms_norm`,
7157 /// `gated_rmsnorm[_q8_1]`, `silu_mul`, `rope_neox` with per-row positions): the T-row
7158 /// launch is the per-row program (same pin the generic verify's fused norms rely on).
7159 /// The sequential dependencies keep their exact serving order: the conv ring / gdn scan
7160 /// chain state row -> row through the `_b` kernels at b_n=1 (ping-pong via a 6-entry
7161 /// alternating pointer table, host handles swapped per row so VerifyCkpt clones the
7162 /// canonical state exactly as the rowwise arm does), and each row's KV append + fa decode
7163 /// picks its arm from ITS OWN t_kv (append: format-only; fa: `fa_seqs_eligible` + its own
7164 /// `fa_split_keys` rung at b_n=1) — the straddle law per row, so every row executes the
7165 /// program its isolated B=1 serving step would.
7166 ///
7167 /// Cost: 1 weight read per layer per round + T state micro-launches, vs the rowwise arm's
7168 /// T weight reads. Gated bit-identical vs the rowwise arm by spec-serve-gate + canary.
7169 #[allow(clippy::too_many_arguments)]
7170 fn qwen35_verify_tparallel(
7171 &self,
7172 e: &Engine,
7173 mut x: CudaSlice<f32>,
7174 lo: usize,
7175 hi: usize,
7176 pos0: usize,
7177 t: usize,
7178 cache: &mut Cache,
7179 mut ckpt: Option<&mut VerifyCkpt>,
7180 stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
7181 mut graphs: Option<&mut DsparkVerifyGraphs>,
7182 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7183 let seqs_append =
7184 std::env::var("MEMRA_BATCH_APPEND").as_deref() != Ok("0") && !Engine::kv_fp8_on();
7185 let batch_fa_on = std::env::var("MEMRA_BATCH_FA").as_deref() != Ok("0");
7186
7187 // Merge guard (v0.98 train, re-affirmed on the v0.100 train over slice 4c): the
7188 // ROUND-STREAM arm (lane/draftcost-moe, device position counter) and the dspark
7189 // verify graphs (engine-bundle slice 3 / trunk slice 4c) have no common caller —
7190 // stream rides the qwen35moe burst, graphs ride the dspark route. If a future
7191 // caller arms both, refuse loudly instead of silently dropping the graphs ctx
7192 // (the stream linear arm takes linear_attn_verify_t, not the graphed segment or
7193 // full-verify bodies).
7194 if stream.is_some() && graphs.is_some() {
7195 return Err(
7196 "qwen35 tparallel verify: ROUND-STREAM and dspark verify graphs \
7197 cannot arm together"
7198 .into(),
7199 );
7200 }
7201 // Engine-bundle slice 3 + slice 4c: with a graphs ctx armed, pointer tables are
7202 // refreshed once per verify (the gdn ping-pong moves handles; a fresh generation
7203 // moves the kv caches). Then:
7204 // - slice 4c: when the WHOLE round rides one seqs rung (every row batchable, one
7205 // split-ladder step, rung covers the round), the ENTIRE walk replays as ONE
7206 // full-verify graph per (vt, rung) — linear layers through the shared
7207 // `qwen35_tparallel_linear_layer` body, full-attention layers through the
7208 // shared `qwen35_tparallel_fa_layer` body in graph mode.
7209 // - fallback (straddle rounds, below the vec floor, partial walks): runs of
7210 // consecutive LINEAR layers replay the slice-3 per-(segment, vt) graphs and
7211 // the full-attention layers run eager (batched rows when eligible).
7212 if let Some(g) = graphs.as_deref_mut() {
7213 g.refresh_tables(e, cache)?;
7214 g.round_slab = false;
7215 if let Some(rung) = g.full_rung(self, cache, lo, hi, t, seqs_append && batch_fa_on) {
7216 // Pool ceiling (dspark_vg_cap): an existing key always replays; a NEW
7217 // full capture past the ceiling falls through to the segment/eager arms.
7218 if g.full.contains_key(&(t, rung, hi)) || g.can_capture() {
7219 let out = g.run_full(self, e, lo, hi, &x, t, pos0, rung, cache)?;
7220 g.round_slab = true;
7221 return Ok(out);
7222 }
7223 }
7224 // Round-atomic ceiling check for the segment door: if any linear run in this
7225 // walk would need a NEW capture past the ceiling, the whole round runs the
7226 // eager cols-ckpt walk (mixing slab- and cols-stashed layers in one round
7227 // would corrupt the commit).
7228 if !g.segments_ready(self, lo, hi, t) {
7229 graphs = None;
7230 }
7231 }
7232 // STREAM (2b, lane/draftcost-moe): positions come from the device round counter
7233 // (pos_iota / i32_copy_add) so a burst round needs no host position knowledge.
7234 let pos_d = match stream {
7235 Some((_, ctr)) => {
7236 let mut p = e.alloc_uninit::<i32>(t)?;
7237 e.pos_iota(ctr, &mut p, t)?;
7238 p
7239 }
7240 None => {
7241 let pos_host: Vec<i32> = (0..t).map(|r| (pos0 + r) as i32).collect();
7242 e.htod_i32(&pos_host)?
7243 }
7244 };
7245 // Per-row 1-element position buffers, built ONCE per verify (the append/fa wrappers
7246 // take owned pos slices; building these inside the layer x row loops cost 16xT H2Ds).
7247 // LAZY since slice 4: the batched fa/append arm never touches them — they are built
7248 // on the first per-row fallback layer only (stream-aware there; the stream FA arm
7249 // rides the dc rows kernels and never reaches the fallback).
7250 let mut pos_rows: Option<Vec<CudaSlice<i32>>> = None;
7251 let mut il = lo;
7252 while il < hi {
7253 if graphs.is_some() && matches!(self.layers[il].mixer, Mixer::Linear(_)) {
7254 let mut end = il;
7255 while end < hi && matches!(self.layers[end].mixer, Mixer::Linear(_)) {
7256 end += 1;
7257 }
7258 let g = graphs.as_deref_mut().expect("checked above");
7259 x = g.run_segment(self, e, il, end, &x, t, cache)?;
7260 g.round_slab = true;
7261 il = end;
7262 continue;
7263 }
7264 let layer = &self.layers[il];
7265 if stream.is_none() && matches!(layer.mixer, Mixer::Linear(_)) {
7266 // Eager linear layer (no graphs ctx): the shared body, legacy cols-ckpt arm.
7267 // Under ROUND-STREAM the linear layers ride the fa-body match's stream arm
7268 // below (linear_attn_verify_t — the stream COMMIT needs its GdnStash).
7269 x = self.qwen35_tparallel_linear_layer(
7270 e,
7271 il,
7272 &x,
7273 t,
7274 cache,
7275 ckpt.as_deref_mut(),
7276 None,
7277 None,
7278 )?;
7279 il += 1;
7280 continue;
7281 }
7282 // Full-attention (or stream-Linear, or MLA-refusing) layer: the extracted
7283 // shared body — eager arm (fresh per-verify pos/table, exact t_kv sizing,
7284 // in-body len bump). The slice-4c captured full-verify graphs run the SAME
7285 // body in graph mode; under ROUND-STREAM the body's dc-rows / GDN stream arms
7286 // run (lane/draftcost-moe).
7287 x = self.qwen35_tparallel_fa_layer(
7288 e,
7289 il,
7290 &x,
7291 t,
7292 cache,
7293 FaLayerArgs {
7294 pos_d: &pos_d,
7295 pos_rows: &mut pos_rows,
7296 pos0,
7297 seqs_append,
7298 batch_fa_on,
7299 graph_cap: None,
7300 stream,
7301 ckpt: ckpt.as_deref_mut(),
7302 },
7303 )?;
7304 il += 1;
7305 }
7306 Ok(x)
7307 }
7308
7309 /// SHARED dense-FFN body for the qwen35 t-parallel layers (trunk-kernels slice B) —
7310 /// ONE copy for the fa and linear layer bodies (the verify_layers extraction lesson).
7311 /// Dual arm (MEMRA_TK_FFN_DUAL, default on): gate+up in ONE dual launch from the
7312 /// pre-quantized activation with macro-scales DEFERRED into the fused SwiGLU+q8_1
7313 /// epilogue, then ffn_down from the fused (aq, ad) — the q27 verify chain verbatim.
7314 /// Every door is the bit-identical proven one: `matmul_decode_exact_dual_pre` (per
7315 /// (tensor,token,row) == the two singles), `silu_mul_scaled_q8_1` (y*s inline == the
7316 /// scale_inplace store, value-exact; fused quantize == quantize_q8_1 bytes),
7317 /// `matmul_decode_exact_pre` (dispatch mirror of the singles' q8_1-fast tail).
7318 /// Dual-refused (t outside 2..=7, non-NVFP4, layout mismatch) or seam off -> the
7319 /// original singles chain, byte-for-byte.
7320 #[allow(clippy::too_many_arguments)]
7321 fn qwen35_tparallel_dense_ffn(
7322 &self,
7323 e: &Engine,
7324 ffn_gate: &crate::model::GpuTensor,
7325 ffn_up: &crate::model::GpuTensor,
7326 ffn_down: &crate::model::GpuTensor,
7327 zn: &CudaSlice<f32>,
7328 t: usize,
7329 n_embd: usize,
7330 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7331 let n_ff = ffn_gate.out_features();
7332 let (zq, zd) = e.quantize_q8_1(zn, t, n_embd)?;
7333 if Engine::tk_ffn_dual_on() {
7334 if let Some(((g, gs), (u, us))) =
7335 e.matmul_decode_exact_dual_pre(ffn_gate, ffn_up, &zq, &zd, t)?
7336 {
7337 if e.uses_q8_1_fast(ffn_down) {
7338 let (aq, ad) = e.silu_mul_scaled_q8_1(&g, &u, gs, us, t * n_ff)?;
7339 return e.matmul_decode_exact_pre(ffn_down, &aq, &ad, t);
7340 }
7341 let mut act = e.uninit(t * n_ff)?;
7342 e.silu_mul_scaled(&g, &u, gs, us, &mut act, t * n_ff)?;
7343 let (aq, ad) = e.quantize_q8_1(&act, t, n_ff)?;
7344 return e.matmul_pre(ffn_down, &aq, &ad, &act, t);
7345 }
7346 }
7347 // v1 singles chain (seam off or dual-refused) — the pre-slice-B body verbatim.
7348 let g = e.matmul_pre(ffn_gate, &zq, &zd, zn, t)?;
7349 let u = e.matmul_pre(ffn_up, &zq, &zd, zn, t)?;
7350 let mut act = e.uninit(t * n_ff)?;
7351 e.silu_mul(&g, &u, &mut act, t * n_ff)?;
7352 let (aq, ad) = e.quantize_q8_1(&act, t, n_ff)?;
7353 e.matmul_pre(ffn_down, &aq, &ad, &act, t)
7354 }
7355
7356 /// ONE t-parallel FULL-ATTENTION layer (attn_norm + fa mixer + post_attn_norm + FFN +
7357 /// tap) — extracted from the walk exactly like `qwen35_tparallel_linear_layer` so the
7358 /// eager walk and the slice-4c captured full-verify graphs execute the SAME body (a
7359 /// second copy is how dispatch mirrors drift — the verify_layers extraction lesson).
7360 ///
7361 /// `args.graph_cap = Some((table, off, rung_end))` is the captured-graph mode:
7362 /// - kv base-pointer pairs come from the ctx-owned persistent table at `off` (a fresh
7363 /// generation's cache lands at new addresses that only the per-verify table refresh
7364 /// knows — the slice-3 baked-address lesson);
7365 /// - the seqs twins size partials/grid at `rung_end` and pin `split_keys` to the
7366 /// rung's ladder value: `n_splits_max` is pure stride, splits >= ns_eff write the
7367 /// EMPTY partial the combine never reads, and every per-row T_kv derives in-kernel
7368 /// from `pos_seq[z]` — so one captured launch replays bit-identically for every
7369 /// round whose rows all sit inside the rung;
7370 /// - the host len bump moves to the replay caller (captured host code does not
7371 /// re-run at replay).
7372 /// Graph mode REFUSES any round the batched arm cannot take: the per-row fallback
7373 /// host-branches on t_kv and must never be captured.
7374 #[allow(clippy::too_many_arguments)]
7375 fn qwen35_tparallel_fa_layer(
7376 &self,
7377 e: &Engine,
7378 il: usize,
7379 x: &CudaSlice<f32>,
7380 t: usize,
7381 cache: &mut Cache,
7382 args: FaLayerArgs<'_>,
7383 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7384 use cudarc::driver::DevicePtr;
7385 let cfg = &self.cfg;
7386 let n_embd = cfg.n_embd as usize;
7387 let eps = cfg.rms_eps;
7388 let head_dim_global = cfg.head_dim_k as usize;
7389 let layer = &self.layers[il];
7390 let FaLayerArgs {
7391 pos_d,
7392 pos_rows,
7393 pos0,
7394 seqs_append,
7395 batch_fa_on,
7396 graph_cap,
7397 stream,
7398 mut ckpt,
7399 } = args;
7400
7401 // ---- attn_norm + q8_1 quantize at m=T (row-indexed == per-row) ----
7402 let anorm = layer.attn_norm.float_data();
7403 let mut xn = e.uninit(t * n_embd)?;
7404 e.rms_norm(x, anorm, &mut xn, n_embd, t, eps)?;
7405 let (hq, hd) = e.quantize_q8_1(&xn, t, n_embd)?;
7406
7407 let mixed: CudaSlice<f32> = match &layer.mixer {
7408 Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
7409 // STREAM ARM (2b, lane/draftcost-moe): under a device position counter the
7410 // per-row serving-kernel chain cannot run (host state swaps keyed on host
7411 // row index are fine, but the stream COMMIT needs the GdnStash for its _dc
7412 // rebuild — the per-row chain only produces per-column clones). GDN rides
7413 // `linear_attn_verify_t`: batched q8_1-class projections, stash-producing,
7414 // and its one-scan recurrence is pinned bit-identical to T chained T=1
7415 // steps (its header + kernel-check). Position-independent, so no counter
7416 // plumbing is needed. Guards mirror the generic call site exactly.
7417 Mixer::Linear(la) if stream.is_some() => {
7418 if !(t >= 3 || (t == 2 && spec_m2()))
7419 || !self.mixer_in_q8_1_fast(e, &layer.mixer)
7420 || !e.uses_q8_1_fast(&la.ssm_out)
7421 {
7422 return Err("qwen35 stream verify: GDN batched arm requires t>=3 \
7423 (or MEMRA_SPEC_M2 at t=2) and q8_1-fast projections"
7424 .into());
7425 }
7426 let want = ckpt.is_some();
7427 let (out, stash) =
7428 self.linear_attn_verify_t(e, la, &xn, Some((&hq, &hd)), t, cache, il, want)?;
7429 if let (Some(ck), Some(st)) = (ckpt.as_deref_mut(), stash) {
7430 ck.gdn[il] = Some(st);
7431 }
7432 out
7433 }
7434 Mixer::Linear(_) => {
7435 unreachable!("linear layers ride qwen35_tparallel_linear_layer")
7436 }
7437 Mixer::Full(fa) => {
7438 let geometry = cfg.full_attention_geometry_at(il as u32);
7439 let n_head = geometry.n_head as usize;
7440 let n_head_kv = geometry.n_head_kv as usize;
7441 let head_dim = geometry.head_dim_k as usize;
7442 let rope_dims = geometry.n_rot as usize;
7443 let rope_base = geometry.rope_base;
7444 let scale = geometry.attention_scale();
7445 // Batched projections: one weight read serves all T rows.
7446 // GROUP-3 twin (trunk-kernels slice D): q/k/v in ONE launch — the group4
7447 // kernel with n3=0, bit-identical per (tensor, token, row) to the three
7448 // singles; refused or MEMRA_TK_FA_GROUP=0 -> singles byte-for-byte.
7449 let (qf, mut k, v) = match e.matmul_decode_exact_group3_pre(
7450 [&fa.wq, &fa.wk, &fa.wv],
7451 &hq,
7452 &hd,
7453 t,
7454 )? {
7455 Some(mut g3) => {
7456 let v = g3.pop().unwrap();
7457 let k = g3.pop().unwrap();
7458 let qf = g3.pop().unwrap();
7459 (qf, k, v)
7460 }
7461 None => (
7462 e.matmul_pre(&fa.wq, &hq, &hd, &xn, t)?,
7463 e.matmul_pre(&fa.wk, &hq, &hd, &xn, t)?,
7464 e.matmul_pre(&fa.wv, &hq, &hd, &xn, t)?,
7465 ),
7466 };
7467 let gated =
7468 geometry.attention_gate == memra_gguf::config::AttentionGateKind::FusedQ;
7469 let (mut q, gate) = if gated {
7470 let mut qs = e.uninit(t * n_head * head_dim)?;
7471 let mut gs = e.uninit(t * n_head * head_dim)?;
7472 e.q_gate_split(&qf, &mut qs, &mut gs, head_dim, n_head, t)?;
7473 (qs, Some(gs))
7474 } else {
7475 (qf, None)
7476 };
7477 let mut qn = e.uninit(t * n_head * head_dim)?;
7478 e.rms_norm(
7479 &q,
7480 fa.q_norm.float_data(),
7481 &mut qn,
7482 head_dim,
7483 t * n_head,
7484 eps,
7485 )?;
7486 q = qn;
7487 let mut kn = e.uninit(t * n_head_kv * head_dim)?;
7488 e.rms_norm(
7489 &k,
7490 fa.k_norm.float_data(),
7491 &mut kn,
7492 head_dim,
7493 t * n_head_kv,
7494 eps,
7495 )?;
7496 k = kn;
7497 e.rope_neox(
7498 &mut q, pos_d, head_dim, rope_dims, n_head, t, rope_base, 1.0,
7499 )?;
7500 e.rope_neox(
7501 &mut k, pos_d, head_dim, rope_dims, n_head_kv, t, rope_base, 1.0,
7502 )?;
7503
7504 // Per-row append + attend: row r sees rows 0..r in KV (causal within the
7505 // draft), each through the b_n=1 serving kernels at its own t_kv.
7506 let q_dim = n_head * head_dim;
7507 let kv_dim = n_head_kv * head_dim;
7508 let mut attn = e.uninit(t * q_dim)?;
7509 let (kdk, kdv, ktb, vtb, len0, kv_local) = {
7510 let kvl = cache.kv[il].as_ref().unwrap();
7511 // [2T] interleaved k,v base pointers: entry pair z serves row z of
7512 // the batched twins; the per-row fallback reads pair 0 (same cache
7513 // for every row of one layer). Graph mode reads the ctx table.
7514 let local: Option<CudaSlice<u64>> = match graph_cap {
7515 Some(_) => None,
7516 None => {
7517 let s = &e.gpu.stream();
7518 let (pk, _g) = kvl.k.device_ptr(s);
7519 let (pv, _g2) = kvl.v.device_ptr(s);
7520 let mut tbl = Vec::with_capacity(2 * t);
7521 for _ in 0..t {
7522 tbl.push(pk as u64);
7523 tbl.push(pv as u64);
7524 }
7525 Some(e.htod_u64(&tbl)?)
7526 }
7527 };
7528 (
7529 kvl.kv_dim_k,
7530 kvl.kv_dim_v,
7531 kvl.k_tok_bytes,
7532 kvl.v_tok_bytes,
7533 kvl.len,
7534 local,
7535 )
7536 };
7537 let (kv_tbl, kv_off): (&CudaSlice<u64>, usize) = match graph_cap {
7538 Some((tb, off, _)) => (tb, off),
7539 None => (kv_local.as_ref().expect("built above"), 0),
7540 };
7541 // Slice 4 (fa/append rows — see dspark_fa_rows_on): the whole per-row
7542 // section batches into the z-batched serving twins when every row of
7543 // this round takes the v4-seqs arm on ONE fa_split_keys rung. Both
7544 // guards are evaluated at the round's FIRST and LAST t_kv — the
7545 // eligibility window (vec floor .. v4 max) and each split-ladder rung
7546 // are intervals in t_kv, so ends-inside means all-inside (the straddle
7547 // law). Appending all T rows before any attend is read-equivalent to
7548 // the interleaved order: row r's walk reads keys 0..len0+r only, and
7549 // rows > r land at slots it never touches; every written cache row is
7550 // the per-token appender's exact warp program (kernel-check pinned).
7551 let t_kv_first = len0 + 1;
7552 let t_kv_last = len0 + t;
7553 let rows_batched = t >= 2
7554 && seqs_append
7555 && batch_fa_on
7556 && dspark_fa_rows_on()
7557 // the z-batched twins read stacked rows at the CACHE's kv dims;
7558 // the projection stack is [T, n_head_kv*head_dim] — they must be
7559 // the same stride or row z misaligns (true for this family; the
7560 // guard keeps any asymmetric-kv model on the per-row loop).
7561 && kdk == kv_dim
7562 && kdv == kv_dim
7563 && crate::fa_seqs_eligible(t_kv_first, head_dim_global)
7564 && crate::fa_seqs_eligible(t_kv_last, head_dim_global)
7565 && crate::fa_split_keys(t_kv_first, cfg.n_head_kv as usize)
7566 == crate::fa_split_keys(t_kv_last, cfg.n_head_kv as usize);
7567 // Sizing: eager = exact round bound; graph mode = the rung end (stride +
7568 // grid only — bytes proven equal above). Capture-time invariants refuse
7569 // loudly rather than bake a divergent body.
7570 let (size_kv_max, sp) = match graph_cap {
7571 Some((_, _, rung)) => {
7572 if !rows_batched {
7573 return Err(format!(
7574 "fa graph capture: layer {il} round is not batchable \
7575 (t_kv {t_kv_first}..{t_kv_last}) — the per-row fallback \
7576 must never be captured"
7577 )
7578 .into());
7579 }
7580 let sp_r = crate::fa_split_keys(rung, cfg.n_head_kv as usize);
7581 if t_kv_last > rung
7582 || sp_r != crate::fa_split_keys(t_kv_last, cfg.n_head_kv as usize)
7583 {
7584 return Err(format!(
7585 "fa graph capture: rung {rung} does not cover round \
7586 t_kv {t_kv_first}..{t_kv_last} on one split ladder step"
7587 )
7588 .into());
7589 }
7590 (rung, sp_r)
7591 }
7592 None => (
7593 t_kv_last,
7594 crate::fa_split_keys(t_kv_last, cfg.n_head_kv as usize),
7595 ),
7596 };
7597 if let Some((_, ctr)) = stream {
7598 // STREAM ARM (2b): one batched dc append + the multi-row dc attention
7599 // — the generic stream arm's exact shape (rows kernels are pinned
7600 // byte-identical to the per-row programs by kernel-check). Host len
7601 // stays a stale lower bound; the burst drain reconciles it.
7602 let kvl = cache.kv[il].as_mut().unwrap();
7603 e.append_kv_quantized_rows_dc(
7604 &k,
7605 &v,
7606 &mut kvl.k,
7607 &mut kvl.v,
7608 ctr,
7609 t,
7610 kdk,
7611 kdv,
7612 ktb,
7613 vtb,
7614 Engine::kv_fp8_on(),
7615 )?;
7616 let upper = (kvl.len + t + 64).min(cache.max_ctx);
7617 let k_view = e.view_u8(&kvl.k, upper * ktb);
7618 let v_view = e.view_u8(&kvl.v, upper * vtb);
7619 e.fa_decode_rows_dc(
7620 &q, &k_view, &v_view, &mut attn, head_dim, n_head, n_head_kv, ctr, upper,
7621 t, scale, ktb, vtb, 0, false,
7622 )?;
7623 } else if rows_batched {
7624 e.append_kv_quantized_seqs(
7625 &k,
7626 &v,
7627 &kv_tbl.slice(kv_off..kv_off + 2 * t),
7628 pos_d,
7629 t,
7630 kdk,
7631 kdv,
7632 ktb,
7633 vtb,
7634 )?;
7635 if graph_cap.is_none() {
7636 cache.kv[il].as_mut().unwrap().len += t;
7637 }
7638 e.fa_decode_batch_seqs_v4(
7639 &q,
7640 &kv_tbl.slice(kv_off..kv_off + 2 * t),
7641 pos_d,
7642 &mut attn,
7643 head_dim,
7644 n_head,
7645 n_head_kv,
7646 t,
7647 size_kv_max,
7648 scale,
7649 sp,
7650 ktb,
7651 vtb,
7652 )?;
7653 } else {
7654 if pos_rows.is_none() {
7655 // Stream-aware for symmetry with pos_d (the stream FA arm rides
7656 // the dc rows kernels above and never reaches this fallback).
7657 *pos_rows = Some(match stream {
7658 Some((_, ctr)) => (0..t)
7659 .map(|r| {
7660 let mut b = e.alloc_uninit::<i32>(1)?;
7661 e.i32_copy_add(ctr, &mut b, r as i32)?;
7662 Ok(b)
7663 })
7664 .collect::<Result<_, Box<dyn std::error::Error>>>()?,
7665 None => (0..t)
7666 .map(|r| e.htod_i32(&[(pos0 + r) as i32]))
7667 .collect::<Result<_, _>>()?,
7668 });
7669 }
7670 let pos_rows = pos_rows.as_ref().unwrap();
7671 for r in 0..t {
7672 // Owned per-row scratch: the b_n=1 kernels take packed batch buffers
7673 // whose row 0 is this row (arithmetic-free materialization copies,
7674 // same as decode's per-seq fallback arm).
7675 let mut k_row = e.uninit(kv_dim)?;
7676 e.dtod_copy_view(&k.slice(r * kv_dim..(r + 1) * kv_dim), &mut k_row)?;
7677 let mut v_row = e.uninit(kv_dim)?;
7678 e.dtod_copy_view(&v.slice(r * kv_dim..(r + 1) * kv_dim), &mut v_row)?;
7679 let pos_row = &pos_rows[r];
7680 let kvl = cache.kv[il].as_mut().unwrap();
7681 if seqs_append {
7682 e.append_kv_quantized_seqs(
7683 &k_row,
7684 &v_row,
7685 &kv_tbl.slice(kv_off..kv_off + 2),
7686 pos_row,
7687 1,
7688 kdk,
7689 kdv,
7690 ktb,
7691 vtb,
7692 )?;
7693 kvl.len += 1;
7694 } else {
7695 e.append_kv_quantized_view(
7696 &k_row.slice(0..kv_dim),
7697 &v_row.slice(0..kv_dim),
7698 &mut kvl.k,
7699 &mut kvl.v,
7700 kvl.len,
7701 kvl.kv_dim_k,
7702 kvl.kv_dim_v,
7703 kvl.k_tok_bytes,
7704 kvl.v_tok_bytes,
7705 Engine::kv_fp8_on(),
7706 )?;
7707 kvl.len += 1;
7708 }
7709 let t_kv = kvl.len;
7710 let mut q_row = e.uninit(q_dim)?;
7711 e.dtod_copy_view(&q.slice(r * q_dim..(r + 1) * q_dim), &mut q_row)?;
7712 let mut a_row = e.uninit(q_dim)?;
7713 if batch_fa_on && crate::fa_seqs_eligible(t_kv, head_dim_global) {
7714 let sp0_r = crate::fa_split_keys(t_kv, cfg.n_head_kv as usize);
7715 e.fa_decode_batch_seqs_v4(
7716 &q_row,
7717 &kv_tbl.slice(kv_off..kv_off + 2),
7718 pos_row,
7719 &mut a_row,
7720 head_dim,
7721 n_head,
7722 n_head_kv,
7723 1,
7724 t_kv,
7725 scale,
7726 sp0_r,
7727 ktb,
7728 vtb,
7729 )?;
7730 } else {
7731 let k_view = e.view_u8(&kvl.k, t_kv * kvl.k_tok_bytes);
7732 let v_view = e.view_u8(&kvl.v, t_kv * kvl.v_tok_bytes);
7733 let mut a_view = a_row.slice_mut(0..q_dim);
7734 e.fa_decode_kvmod_view(
7735 &q_row.slice(0..q_dim),
7736 &k_view,
7737 &v_view,
7738 &mut a_view,
7739 head_dim,
7740 n_head,
7741 n_head_kv,
7742 t_kv,
7743 scale,
7744 kvl.k_tok_bytes,
7745 kvl.v_tok_bytes,
7746 Engine::kv_fp8_on(),
7747 )?;
7748 }
7749 e.dtod_copy_into(&a_row, &mut attn, r * q_dim)?;
7750 }
7751 }
7752
7753 // Output gate (element-wise) + o-proj at m=T.
7754 let attn_g = match &gate {
7755 Some(g) => {
7756 let n = t * q_dim;
7757 let mut gsig = e.uninit(n)?;
7758 e.sigmoid(g, &mut gsig, n)?;
7759 let mut ag = e.uninit(n)?;
7760 e.mul(&attn, &gsig, &mut ag, n)?;
7761 ag
7762 }
7763 None => attn,
7764 };
7765 e.matmul(&fa.wo, &attn_g, t)?
7766 }
7767 };
7768
7769 // ---- residual add + post_attn_norm + FFN at m=T (serving dispatch verbatim) ----
7770 let pnorm = layer.post_attn_norm.float_data();
7771 let mut x1 = e.uninit(t * n_embd)?;
7772 let mut zn = e.uninit(t * n_embd)?;
7773 e.add_rms_norm(x, &mixed, pnorm, &mut x1, &mut zn, n_embd, t, eps)?;
7774 let ffn_out = match &layer.ffn {
7775 crate::hybrid::Ffn::Dense {
7776 ffn_gate,
7777 ffn_up,
7778 ffn_down,
7779 } => {
7780 assert!(
7781 self.cfg.m3.is_none(),
7782 "qwen35 t-parallel verify: M3 swigluoai FFN not yet batched"
7783 );
7784 self.qwen35_tparallel_dense_ffn(e, ffn_gate, ffn_up, ffn_down, &zn, t, n_embd)?
7785 }
7786 crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il_zq8(e, m, &zn, None, t, il as u16)?,
7787 };
7788 let mut x2 = e.uninit(t * n_embd)?;
7789 e.add(&x1, &ffn_out, &mut x2, t * n_embd)?;
7790 // dspark drafter tap (no-op when no sink armed): post-layer residual verify rows
7791 self.dflash_tap(e, cache, il, &x2, t)?;
7792 Ok(x2)
7793 }
7794
7795 /// ONE t-parallel LINEAR layer (attn_norm + gdn mixer + post_attn_norm + FFN + tap) —
7796 /// the exact body the old in-loop Linear arm ran, extracted so the eager walk and the
7797 /// slice-3 captured segments execute the SAME code (a second copy is how dispatch
7798 /// mirrors drift — the verify_layers extraction lesson). Two deliberate changes, both
7799 /// bit-identical by construction:
7800 /// - the gdn ping-pong host swap moves from per-row to ONE end-of-body swap (t odd):
7801 /// the device sequence is driven entirely by the 6-entry pointer table, which
7802 /// already encodes both parities; the ckpt stash reads name row r's out buffer
7803 /// directly (r even -> alt handle, odd -> canonical) — the same physical bytes the
7804 /// legacy post-swap clone read.
7805 /// - `stash` (slice-3 ctx): persistent per-layer slabs written by copy_into instead of
7806 /// per-row clone_dtod allocs — same bytes, capture-legal (no per-round host objects).
7807 /// `table_src` = (persistent pointer table, offset) when the ctx owns the tables;
7808 /// None builds the per-verify table exactly as before.
7809 #[allow(clippy::too_many_arguments)]
7810 fn qwen35_tparallel_linear_layer(
7811 &self,
7812 e: &Engine,
7813 il: usize,
7814 x: &CudaSlice<f32>,
7815 t: usize,
7816 cache: &mut Cache,
7817 mut ckpt: Option<&mut VerifyCkpt>,
7818 stash: Option<(&mut CudaSlice<f32>, &mut CudaSlice<f32>)>,
7819 table_src: Option<(&CudaSlice<u64>, usize)>,
7820 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7821 use cudarc::driver::DevicePtr;
7822 let cfg = &self.cfg;
7823 let n_embd = cfg.n_embd as usize;
7824 let eps = cfg.rms_eps;
7825 let layer = &self.layers[il];
7826 let Mixer::Linear(la) = &layer.mixer else {
7827 return Err("qwen35_tparallel_linear_layer on a non-linear layer".into());
7828 };
7829 // ---- attn_norm + q8_1 quantize at m=T (row-indexed == per-row) ----
7830 let anorm = layer.attn_norm.float_data();
7831 let mut xn = e.uninit(t * n_embd)?;
7832 e.rms_norm(x, anorm, &mut xn, n_embd, t, eps)?;
7833 let (hq, hd) = e.quantize_q8_1(&xn, t, n_embd)?;
7834
7835 let geometry = la.geometry;
7836 let d_state = geometry.key_head_dim as usize;
7837 let num_k = geometry.key_heads as usize;
7838 let num_v = geometry.value_heads as usize;
7839 let d_conv = geometry.conv_kernel as usize;
7840 let key_dim = d_state * num_k;
7841 let value_dim = geometry.value_head_dim as usize * num_v;
7842 let conv_dim = key_dim * 2 + value_dim;
7843 let gdn_scale = 1.0 / (d_state as f32).sqrt();
7844
7845 // ---- batched projections: one weight read for all T rows ----
7846 // GROUP-4 twin (trunk-kernels slice C): the whole 4-tuple in ONE launch, bit-identical
7847 // per (tensor, token, row) to the four singles; refused (layout/tier) or
7848 // MEMRA_TK_GDN_GROUP=0 -> the singles chain byte-for-byte.
7849 let (qkv_mixed, z, beta_raw, alpha) = match e.matmul_decode_exact_group4_pre(
7850 [&la.wqkv, &la.wqkv_gate, &la.ssm_beta, &la.ssm_alpha],
7851 &hq,
7852 &hd,
7853 t,
7854 )? {
7855 Some(mut g4) => {
7856 let alpha = g4.pop().unwrap();
7857 let beta_raw = g4.pop().unwrap();
7858 let z = g4.pop().unwrap();
7859 let qkv_mixed = g4.pop().unwrap();
7860 (qkv_mixed, z, beta_raw, alpha)
7861 }
7862 None => (
7863 e.matmul_pre(&la.wqkv, &hq, &hd, &xn, t)?,
7864 e.matmul_pre(&la.wqkv_gate, &hq, &hd, &xn, t)?,
7865 e.matmul_pre(&la.ssm_beta, &hq, &hd, &xn, t)?,
7866 e.matmul_pre(&la.ssm_alpha, &hq, &hd, &xn, t)?,
7867 ),
7868 };
7869 let beta_w = la.ssm_beta.out_features();
7870 let alpha_w = la.ssm_alpha.out_features();
7871 let qkv_w = la.wqkv.out_features();
7872
7873 // ---- per-row state chain through the b_n=1 serving kernels ----
7874 // 6-entry alternating pointer table expresses the ping-pong without a rebuild per
7875 // row: even rows scan s0 -> s1, odd rows s1 -> s0.
7876 let table_local: Option<CudaSlice<u64>> = match table_src {
7877 Some(_) => None,
7878 None => {
7879 let rl = cache.recur[il].as_ref().unwrap();
7880 let s = &e.gpu.stream();
7881 let (pc, _g0) = rl.conv_state.device_ptr(s);
7882 let (p0, _g1) = rl.ssm_state.device_ptr(s);
7883 let (p1, _g2) = rl.ssm_state_alt.device_ptr(s);
7884 Some(e.htod_u64(&[
7885 pc as u64, p0 as u64, p1 as u64, pc as u64, p1 as u64, p0 as u64,
7886 ])?)
7887 }
7888 };
7889 let (table, toff): (&CudaSlice<u64>, usize) = match table_src {
7890 Some((tb, off)) => (tb, off),
7891 None => (table_local.as_ref().unwrap(), 0),
7892 };
7893 let mut o_all = e.uninit(t * value_dim)?;
7894 let mut col_states: Option<Vec<(CudaSlice<f32>, CudaSlice<f32>)>> =
7895 if ckpt.is_some() && stash.is_none() && t >= 2 {
7896 Some(Vec::with_capacity(t - 1))
7897 } else {
7898 None
7899 };
7900 let mut stash = stash;
7901 // Per-row scratch reused across rows (uninit is cheap but not free at
7902 // 48 layers x T rows); row inputs/outputs pass as VIEWS into the packed
7903 // [T, ...] buffers — zero arithmetic-free copies in this loop.
7904 let mut conv_out = e.uninit(conv_dim)?;
7905 let mut q_l2 = e.uninit(value_dim)?;
7906 let mut k_l2 = e.uninit(value_dim)?;
7907 let mut v_gd = e.uninit(value_dim)?;
7908 let mut beta_b = e.uninit(num_v)?;
7909 let mut g_log = e.uninit(num_v)?;
7910 for r in 0..t {
7911 let base = toff + if r % 2 == 0 { 0 } else { 3 };
7912 let conv_view = table.slice(base..base + 1);
7913 let in_view = table.slice(base + 1..base + 2);
7914 let out_view = table.slice(base + 2..base + 3);
7915 e.ssm_conv1d_fused_decode_b_view(
7916 &qkv_mixed.slice(r * qkv_w..(r + 1) * qkv_w),
7917 &conv_view,
7918 la.ssm_conv1d.float_data(),
7919 &mut conv_out,
7920 conv_dim,
7921 d_conv,
7922 1,
7923 )?;
7924 e.gdn_prep_decode_b_view(
7925 &conv_out,
7926 &beta_raw.slice(r * beta_w..(r + 1) * beta_w),
7927 &alpha.slice(r * alpha_w..(r + 1) * alpha_w),
7928 la.ssm_dt.float_data(),
7929 la.ssm_a.float_data(),
7930 &mut q_l2,
7931 &mut k_l2,
7932 &mut v_gd,
7933 &mut beta_b,
7934 &mut g_log,
7935 d_state,
7936 num_v,
7937 num_k,
7938 key_dim,
7939 eps,
7940 conv_dim,
7941 1,
7942 )?;
7943 let mut o_row = o_all.slice_mut(r * value_dim..(r + 1) * value_dim);
7944 e.gdn_scan_s128_batched_view(
7945 &q_l2, &k_l2, &v_gd, &g_log, &beta_b, &in_view, &out_view, &mut o_row, num_v, 1,
7946 gdn_scale,
7947 )?;
7948 if r + 1 < t {
7949 // Row r's out buffer: even rows write s1 (the alt handle — no swaps ran),
7950 // odd rows write s0 — the same physical state the legacy post-swap
7951 // canonical clone read.
7952 let rl = cache.recur[il]
7953 .as_ref()
7954 .ok_or("qwen35 linear verify layer has no recurrent state")?;
7955 let ssm_src = if r % 2 == 0 {
7956 &rl.ssm_state_alt
7957 } else {
7958 &rl.ssm_state
7959 };
7960 match stash.as_mut() {
7961 Some((conv_slab, ssm_slab)) => {
7962 // BOTH stash reads go through the pointer table at run time: the
7963 // ssm handles ping-pong between rounds, and the ctx (with its
7964 // captured graphs) outlives the Cache — a fresh generation's
7965 // conv/ssm buffers land at new addresses that only the per-round
7966 // table refresh knows. A baked direct copy would read freed
7967 // memory (parity was the slice-3 smoke divergence; cache
7968 // lifetime is the cross-generation twin).
7969 e.copy_indirect_src_f32(
7970 &conv_view,
7971 conv_slab,
7972 r * conv_dim * (d_conv - 1),
7973 conv_dim * (d_conv - 1),
7974 )?;
7975 // The ssm handles PING-PONG between rounds: a captured direct
7976 // copy would bake the capture-time physical buffer and read the
7977 // wrong parity after any odd-vt round (the slice-3 smoke
7978 // divergence). Read the src address from row r's OUT table
7979 // entry at run time — the same entry the scan just wrote.
7980 e.copy_indirect_src_f32(
7981 &out_view,
7982 ssm_slab,
7983 r * d_state * d_state * num_v,
7984 d_state * d_state * num_v,
7985 )?;
7986 }
7987 None => {
7988 if let Some(states) = col_states.as_mut() {
7989 states.push((e.clone_dtod(&rl.conv_state)?, e.clone_dtod(ssm_src)?));
7990 }
7991 }
7992 }
7993 }
7994 }
7995 // ONE end-of-body parity swap (t odd) — the legacy loop swapped per row; the net
7996 // handle motion is identical and the device sequence never read the handles.
7997 if t % 2 == 1 {
7998 let rl = cache.recur[il].as_mut().unwrap();
7999 std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
8000 }
8001 if let (Some(checkpoint), Some(states)) = (ckpt.as_deref_mut(), col_states) {
8002 checkpoint.cols[il] = Some(states);
8003 }
8004
8005 // ---- batched gated norm + out-projection at m=T ----
8006 let mixed = if e.uses_q8_1_fast(&la.ssm_out) {
8007 let (gq, gd) = e.gated_rmsnorm_q8_1(
8008 &o_all,
8009 la.ssm_norm.float_data(),
8010 &z,
8011 d_state,
8012 t * num_v,
8013 eps,
8014 )?;
8015 let g0 = e.zeros(0)?;
8016 e.matmul_pre(&la.ssm_out, &gq, &gd, &g0, t)?
8017 } else {
8018 let mut gn = e.uninit(t * value_dim)?;
8019 e.gated_rmsnorm(
8020 &o_all,
8021 la.ssm_norm.float_data(),
8022 &z,
8023 &mut gn,
8024 d_state,
8025 t * num_v,
8026 eps,
8027 )?;
8028 e.matmul(&la.ssm_out, &gn, t)?
8029 };
8030
8031 // ---- residual add + post_attn_norm + FFN at m=T (serving dispatch verbatim) ----
8032 let pnorm = layer.post_attn_norm.float_data();
8033 let mut x1 = e.uninit(t * n_embd)?;
8034 let mut zn = e.uninit(t * n_embd)?;
8035 e.add_rms_norm(x, &mixed, pnorm, &mut x1, &mut zn, n_embd, t, eps)?;
8036 let ffn_out = match &layer.ffn {
8037 crate::hybrid::Ffn::Dense {
8038 ffn_gate,
8039 ffn_up,
8040 ffn_down,
8041 } => {
8042 assert!(
8043 self.cfg.m3.is_none(),
8044 "qwen35 t-parallel verify: M3 swigluoai FFN not yet batched"
8045 );
8046 self.qwen35_tparallel_dense_ffn(e, ffn_gate, ffn_up, ffn_down, &zn, t, n_embd)?
8047 }
8048 crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il_zq8(e, m, &zn, None, t, il as u16)?,
8049 };
8050 let mut x2 = e.uninit(t * n_embd)?;
8051 e.add(&x1, &ffn_out, &mut x2, t * n_embd)?;
8052 // dspark drafter tap (no-op when no sink armed): post-layer residual verify rows
8053 self.dflash_tap(e, cache, il, &x2, t)?;
8054 Ok(x2)
8055 }
8056
8057 /// PP-N STAGE SUBGRAPH of the verify trunk: layers `[lo, hi)` of `decode_step_t_core_stream`'s
8058 /// walk, verbatim. Enters with a MATERIALIZED `[T, n_embd]` residual (no pending fusion pair
8059 /// carried in from outside the range) and exits with the range's final residual materialized
8060 /// (the trailing add executed) — exactly the `decode_layers_eager(lo, hi)` contract, T rows
8061 /// instead of one.
8062 ///
8063 /// EXTRACTED (lane/pp2-spec 2026-08-06) rather than duplicated: `decode_step_t_core_stream` IS
8064 /// the single funnel every verify forward reaches, and its per-layer dispatch MIRRORING (norm
8065 /// fusion per layer, the t>=3/spec_m2 batched-linear window, the fused-q8 FFN chain, the
8066 /// decode-exact projections) is what makes verify bit-identical to eager decode. A second copy
8067 /// for the split arm is how those mirrors drift apart on the next lever. The unsplit body now
8068 /// calls this with `(0, n_layers)`, so the whole-trunk path and every stage range run the SAME
8069 /// code — there is no "split version" of the verify math.
8070 ///
8071 /// Bit-identity of a cut rests on the same kernel-check-pinned identity the eager arm's cut
8072 /// does — `add_rms_norm_q8_1 == add then rms_norm_q8_1` at nrows=T — because the ONLY thing a
8073 /// fence changes is that the cross-layer fusion carry breaks at `hi-1` and is re-materialized
8074 /// as an explicit `add`. `decode-batch-gate --mode ppspec` verifies end-to-end on real weights.
8075 #[allow(clippy::too_many_arguments)]
8076 fn verify_layers(
8077 &self,
8078 e: &Engine,
8079 mut x: CudaSlice<f32>,
8080 lo: usize,
8081 hi: usize,
8082 pos_d: &CudaSlice<i32>,
8083 pos0: usize,
8084 t: usize,
8085 cache: &mut Cache,
8086 mut ckpt: Option<&mut VerifyCkpt>,
8087 stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
8088 graphs: Option<&mut DsparkVerifyGraphs>,
8089 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8090 if self.sliding_gated_moe_batch_program() {
8091 if stream.is_some() {
8092 return Err(
8093 "step35 has no ROUND-STREAM verify arm (the device-counter _dc twins \
8094 cannot express the SWA offset KV view)"
8095 .into(),
8096 );
8097 }
8098 return self.step35_verify_batch_layers(e, x, lo, hi, pos0, t, cache);
8099 }
8100 if self.batched_serving_numeric_class() {
8101 return self.qwen35_verify_batch_layers(
8102 e,
8103 x,
8104 lo,
8105 hi,
8106 pos0,
8107 t,
8108 cache,
8109 ckpt.take(),
8110 stream,
8111 graphs,
8112 );
8113 }
8114 let n_embd = self.cfg.n_embd as usize;
8115 let eps = self.cfg.rms_eps;
8116 // CROSS-LAYER ADD+NORM FUSION (lane/vt-fixes fix 2, mirroring decode_step_h's
8117 // launch-arc form): layer il's post-FFN residual add (x2 = x1 + ffn_out) and layer
8118 // il+1's attn_norm(+quantize) are consecutive row-wise ops — ONE add_rms_norm_q8_1
8119 // launch at nrows=t does all three (bit-identity pinned by the T-row kernel-check
8120 // arms). Carry the un-added (x1, ffn_out) pair; the fused launch materializes x2 (the
8121 // residual the next layer needs) as its `res` output. Falls back to the separate add
8122 // when the next layer is off the fused-q8 path.
8123 let mut pending: Option<(CudaSlice<f32>, CudaSlice<f32>)> = None;
8124 for il in lo..hi {
8125 let layer = &self.layers[il];
8126 // DISPATCH-MIRRORED attn-input RMSNorm (FP-order lesson #8): eager decode fuses the
8127 // 1024-thread rms_norm_q8_1 ONLY when every mixer projection is q8_1-fast; layers with
8128 // Float projections (ssm_beta/ssm_alpha on layers 1/2/4 of the 9B NVFP4 GGUF) take the
8129 // UNFUSED 256-thread rms_norm. The verify norm must mirror that PER-LAYER choice —
8130 // blockDim changes the sum-of-squares reduce order, and the ULP shift amplifies through
8131 // the GDN recurrence into argmax flips (measured: 9B text prompt, 1 ULP at layer 2 ->
8132 // 2.3e-1 logit maxdiff at the head -> K=1..8 divergence at a 0.03-margin token).
8133 let mixer_fast = self.mixer_in_q8_1_fast(e, &layer.mixer);
8134 let norm_fused = std::env::var("MEMRA_NO_FUSE_NORMQ").is_err() && mixer_fast;
8135 // BATCHED EPILOGUE RE-FUSE (lane/vt-fixes fix 2, 2026-08-03): when the norm is
8136 // dispatch-fused AND every consumer of `h` reads only its q8_1 form (Full mixer:
8137 // projections only; Linear mixer: the batched arm — the per-column fallback needs
8138 // f32 h), emit the attn-input norm DIRECTLY as q8_1 via `rms_norm_q8_1` at nrows=t
8139 // (row-indexed kernel — the T-row launch is the per-row m=1 program, kernel-check
8140 // pins bit-identity vs rms_norm_decode -> quantize_q8_1). Kills the standalone
8141 // quantize launch(es) + the f32 h HBM round-trip that decode never pays.
8142 // step35 (Full mixer) is the third case that needs f32 `h`: its verify arm is a
8143 // per-ROW replay of the eager decode mixer, whose `pre_q` contract is a single row —
8144 // a T-row q8_1 pair cannot be handed to it, and re-deriving per-row q8_1 from the f32
8145 // rows is exactly the dispatch being mirrored. Keep step35 on the unfused arm.
8146 let lin_q8_only = match &layer.mixer {
8147 Mixer::Linear(la) => {
8148 (t >= 3 || (t == 2 && spec_m2())) && e.uses_q8_1_fast(&la.ssm_out)
8149 }
8150 Mixer::Full(_) if self.sliding_gated_moe_batch_program() => false,
8151 _ => true,
8152 };
8153 // NOTE decode.rs's take()-first lesson: take the pending pair BEFORE branching so
8154 // a non-fused layer still performs the residual add.
8155 let taken = pending.take();
8156 let (h, h_q8) = if norm_fused && lin_q8_only {
8157 let pair = match taken {
8158 // fused add + attn_norm + q8_1: ONE launch resolves the carried residual
8159 // AND emits this layer's mixer input pre-quantized. res -> x2 (= new x).
8160 Some((x1p, f1p)) => {
8161 let mut x2 = vbuf(e, t * n_embd)?; // fully written (res output)
8162 let p = e.add_rms_norm_q8_1(
8163 &x1p,
8164 &f1p,
8165 layer.attn_norm.float_data(),
8166 &mut x2,
8167 n_embd,
8168 t,
8169 eps,
8170 )?;
8171 x = x2;
8172 p
8173 }
8174 None => e.rms_norm_q8_1(&x, layer.attn_norm.float_data(), n_embd, t, eps)?,
8175 };
8176 (e.zeros(0)?, Some(pair)) // h unused on this path (q8-only consumers)
8177 } else {
8178 if let Some((x1p, f1p)) = taken {
8179 let mut x2 = vbuf(e, t * n_embd)?; // fully written by add
8180 e.add(&x1p, &f1p, &mut x2, t * n_embd)?;
8181 x = x2;
8182 }
8183 let mut h = vbuf(e, t * n_embd)?; // fully written by either rms_norm arm
8184 if norm_fused {
8185 e.rms_norm_decode(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
8186 } else {
8187 e.rms_norm(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
8188 }
8189 (h, None)
8190 };
8191 let h_q8_ref = h_q8.as_ref().map(|(q, d)| (q, d));
8192
8193 let mixed = match &layer.mixer {
8194 Mixer::Full(fa) => self.full_attn_verify(
8195 e,
8196 fa,
8197 &h,
8198 h_q8_ref,
8199 pos_d,
8200 t,
8201 cache,
8202 il,
8203 stream.map(|(_, c)| c),
8204 )?,
8205 Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
8206 Mixer::Linear(la) => {
8207 // BATCHED linear verify (2026-07-03, the MTP-profit lever): one T-token pass —
8208 // batched projections (weight read ONCE, hits the m=2-4 weight-resident matvec),
8209 // carried-state conv (ssm_conv1d_tm_state), GDN prep on the prefill kernels, and
8210 // ONE gdn_scan whose internal sequential t-loop is the SAME recurrence as T
8211 // chained T=1 steps (bit-identical). Falls back to the sequential per-column
8212 // chain when T < d_conv-1 (conv ring update needs T >= pad) — or when ANY
8213 // projection is off the q8_1 fast path: matmul_decode_exact would route a Float
8214 // tensor to cuBLAS at m=t (different FP accumulation than eager's per-token
8215 // GEMV), so mixed-dtype layers stay on the eager-identical per-column chain.
8216 // MEMRA_SPEC_M2 (lane/spec-m2): the t==2 batch rides the same arm — the conv
8217 // wrapper handles t<pad with a pure-copy ring rebuild; see spec_m2() header.
8218 if (t >= 3 || (t == 2 && spec_m2()))
8219 && mixer_fast
8220 && e.uses_q8_1_fast(&la.ssm_out)
8221 {
8222 let want = ckpt.is_some();
8223 let (out, stash) =
8224 self.linear_attn_verify_t(e, la, &h, h_q8_ref, t, cache, il, want)?;
8225 if let (Some(ck), Some(st)) = (ckpt.as_deref_mut(), stash) {
8226 ck.gdn[il] = Some(st);
8227 }
8228 out
8229 } else {
8230 let mut out = vbuf(e, t * n_embd)?; // every col written by copy_into
8231 let mut col_states: Option<Vec<(CudaSlice<f32>, CudaSlice<f32>)>> =
8232 if ckpt.is_some() && t >= 2 {
8233 Some(Vec::with_capacity(t - 1))
8234 } else {
8235 None
8236 };
8237 for col in 0..t {
8238 let mut h_col = vbuf(e, n_embd)?; // fully written by copy_view_into
8239 let src = h.slice(col * n_embd..(col + 1) * n_embd);
8240 e.copy_view_into(&mut h_col, 0, &src, n_embd)?;
8241 let m_col = self.linear_attn_decode(e, la, &h_col, cache, il)?;
8242 e.copy_into(&mut out, col * n_embd, &m_col, n_embd)?;
8243 // REPLAY-FREE ckpt: clone the chain's ACTUAL state after this column
8244 // (pure dtod — cannot change any computed value). Last column skipped:
8245 // rebuild targets are j <= t-1 columns.
8246 if let Some(cs) = col_states.as_mut() {
8247 if col + 1 < t {
8248 let rl = cache.recur[il].as_ref().unwrap();
8249 cs.push((
8250 e.clone_dtod(&rl.conv_state)?,
8251 e.clone_dtod(&rl.ssm_state)?,
8252 ));
8253 }
8254 }
8255 }
8256 if let (Some(ck), Some(cs)) = (ckpt.as_deref_mut(), col_states) {
8257 // ReplaySSM-assessment instrumentation (2026-07-30): the
8258 // per-column clones are the only true state snapshots left in
8259 // the verify (the batched path stashes INPUTS and replays).
8260 if std::env::var("MEMRA_SPEC_STATS").as_deref() == Ok("1") {
8261 static ONCE: std::sync::Once = std::sync::Once::new();
8262 let bytes: usize =
8263 cs.iter().map(|(c, s)| (c.len() + s.len()) * 4).sum();
8264 ONCE.call_once(|| eprintln!(
8265 "[verify-ckpt] per-column layer il={il}: {} clones, {:.2} MB/layer/round",
8266 cs.len(), bytes as f64 / 1e6));
8267 }
8268 ck.cols[il] = Some(cs);
8269 }
8270 out
8271 }
8272 }
8273 };
8274
8275 // DISPATCH-MIRRORED post-attn norm: eager residual_norm_ffn fuses add+norm+quant
8276 // (1024-thread add_rms_norm_q8_1) only for Dense FFNs whose gate+up are q8_1-fast;
8277 // otherwise (and for MoE) it runs the 256-thread fused add_rms_norm. Mirror per layer.
8278 let ffn_fuse = match &layer.ffn {
8279 crate::hybrid::Ffn::Dense {
8280 ffn_gate, ffn_up, ..
8281 } => {
8282 std::env::var("MEMRA_NO_FUSE_NORMQ").is_err()
8283 && e.uses_q8_1_fast(ffn_gate)
8284 && e.uses_q8_1_fast(ffn_up)
8285 }
8286 crate::hybrid::Ffn::Moe(_) => false,
8287 };
8288 // BATCHED EPILOGUE RE-FUSE (lane/vt-fixes fix 2): on the ffn_fuse path (Dense,
8289 // gate+up q8_1-fast, non-M3) the FFN input is emitted DIRECTLY as q8_1 by ONE
8290 // add_rms_norm_q8_1 launch at nrows=t (row-indexed kernel: the T-row launch is the
8291 // per-row m=1 program; kernel-check pins bit-identity vs the unfused
8292 // add_f32 -> rms_norm_decode -> quantize_q8_1 chain at T=2/4/5/8) — replacing the
8293 // add + rms_norm_decode launches AND the dual/singles' internal re-quantize.
8294 // M3's swigluoai must keep the f32 chain (the fused SwiGLU epilogue encodes plain
8295 // SiLU), mirroring residual_norm_ffn's m3 guard on the decode path.
8296 // step35: same guard per LAYER. A dense FFN's clamp is the SHEXP array (upstream's
8297 // one build_ffn serves dense + shared expert, llama-graph.cpp:1751), and verify MUST
8298 // mirror decode's dispatch or spec self-consistency fails.
8299 let dense_lim = self.cfg.clamp_shexp_at(il as u32);
8300 let fuse_q8 = ffn_fuse && self.cfg.m3.is_none() && dense_lim.is_none();
8301 let mut x1 = vbuf(e, t * n_embd)?; // fully written by add / add_rms_norm*
8302 let mut z = e.zeros(0)?; // replaced below on the unfused arms
8303 let z_q8 = if fuse_q8 {
8304 Some(e.add_rms_norm_q8_1(
8305 &x,
8306 &mixed,
8307 layer.post_attn_norm.float_data(),
8308 &mut x1,
8309 n_embd,
8310 t,
8311 eps,
8312 )?)
8313 } else {
8314 let mut zf = vbuf(e, t * n_embd)?; // fully written by rms_norm_decode / add_rms_norm
8315 if ffn_fuse {
8316 e.add(&x, &mixed, &mut x1, t * n_embd)?;
8317 e.rms_norm_decode(
8318 &x1,
8319 layer.post_attn_norm.float_data(),
8320 &mut zf,
8321 n_embd,
8322 t,
8323 eps,
8324 )?;
8325 } else {
8326 e.add_rms_norm(
8327 &x,
8328 &mixed,
8329 layer.post_attn_norm.float_data(),
8330 &mut x1,
8331 &mut zf,
8332 n_embd,
8333 t,
8334 eps,
8335 )?;
8336 }
8337 z = zf;
8338 None
8339 };
8340 // DECODE-EXACT FFN projections: force MMVQ for gate/up/down at any T to match the
8341 // T=1 decode FP accumulation order. At T>=5 the generic matmul/matmul_pre falls to dp4a
8342 // (128-thread, different FP sum order). At T=2-4 the batched MMVQ is already bit-identical.
8343 let ffn_out = match &layer.ffn {
8344 crate::hybrid::Ffn::Dense {
8345 ffn_gate,
8346 ffn_up,
8347 ffn_down,
8348 } => {
8349 let n_ff = ffn_gate.out_features();
8350 if let Some((zq, zd)) = z_q8.as_ref() {
8351 // FUSED CHAIN (fix 2): pre-quantized z feeds the projections; the SwiGLU
8352 // epilogue emits act pre-quantized for ffn_down (silu_mul_scaled_q8_1,
8353 // bit-identical to silu_mul + quantize — kernel-check-pinned) with the
8354 // NVFP4 macro-scales folded (deferred-scale dual: y*s inline == the
8355 // scale_inplace store, value-exact) — the exact m=1 decode epilogue
8356 // structure at nrows=t.
8357 let pair =
8358 match e.matmul_decode_exact_dual_pre(ffn_gate, ffn_up, zq, zd, t)? {
8359 Some(((g, gs), (u, us))) => Some((g, gs, u, us)),
8360 None => None,
8361 };
8362 let (gate, gs, up, us) = match pair {
8363 Some(x4) => x4,
8364 None => (
8365 e.matmul_decode_exact_pre(ffn_gate, zq, zd, t)?,
8366 1.0, // scale already applied inside _pre
8367 e.matmul_decode_exact_pre(ffn_up, zq, zd, t)?,
8368 1.0,
8369 ),
8370 };
8371 if e.uses_q8_1_fast(ffn_down) {
8372 let (aq, ad) = e.silu_mul_scaled_q8_1(&gate, &up, gs, us, t * n_ff)?;
8373 e.matmul_decode_exact_pre(ffn_down, &aq, &ad, t)?
8374 } else {
8375 let mut act = vbuf(e, t * n_ff)?;
8376 e.silu_mul_scaled(&gate, &up, gs, us, &mut act, t * n_ff)?;
8377 e.matmul_decode_exact(ffn_down, &act, t)?
8378 }
8379 } else {
8380 // UNFUSED (pre-fix) chain — MoE-adjacent/M3/off-fast layers, unchanged.
8381 // DUAL gate+up batched twin (lane/verify-economics, 2026-08-02): one launch
8382 // for the pair at t=2..8 — bit-identical per (tensor,token,row) to the two
8383 // singles (kernel-check pins bitwise; MEMRA_SPEC_DUAL_T=0 reverts). None
8384 // (non-NVFP4 / t outside the tier / seam off) -> the two singles, unchanged.
8385 let (gate, up) =
8386 match e.matmul_decode_exact_dual(ffn_gate, ffn_up, &z, t)? {
8387 Some(pair) => pair,
8388 None => (
8389 e.matmul_decode_exact(ffn_gate, &z, t)?,
8390 e.matmul_decode_exact(ffn_up, &z, t)?,
8391 ),
8392 };
8393 let mut act = vbuf(e, t * n_ff)?; // fully written by ffn_act_lim
8394 Self::ffn_act_lim(
8395 e,
8396 &self.cfg,
8397 &gate,
8398 &up,
8399 1.0,
8400 1.0,
8401 dense_lim,
8402 &mut act,
8403 t * n_ff,
8404 )?;
8405 e.matmul_decode_exact(ffn_down, &act, t)?
8406 }
8407 }
8408 crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il(e, m, &z, t, il as u16)?,
8409 };
8410 // CROSS-LAYER fusion: defer this layer's post-FFN residual add — the next layer's
8411 // fused-q8 attn norm folds it in (add_rms_norm_q8_1 == add; rms_norm; quantize,
8412 // kernel-check-pinned at nrows=T). Non-fused next layers add explicitly above.
8413 pending = Some((x1, ffn_out));
8414 }
8415 // RANGE's final add (no next norm INSIDE the range to fuse with; for the
8416 // whole-trunk call that is the last layer, whose next norm is output_norm — f32-out).
8417 if let Some((x1p, f1p)) = pending.take() {
8418 let mut x2 = vbuf(e, t * n_embd)?; // fully written by add
8419 e.add(&x1p, &f1p, &mut x2, t * n_embd)?;
8420 x = x2;
8421 }
8422 Ok(x)
8423 }
8424 /// BATCHED linear-attn verify (T=K+1): the whole layer in ~10 launches instead of T x the
8425 /// T=1 decode chain (T x ~12 launches + T weight reads of the four projections). The GDN
8426 /// recurrence itself is inherently sequential — gdn_scan_s128 runs its internal t-loop with
8427 /// the SAME per-token math as chained T=1 calls (bit-identical state evolution); everything
8428 /// around it (projections, conv, prep, gated norm, out-proj) batches. Advances conv ring +
8429 /// ssm state exactly like T sequential decode steps.
8430 /// `want_stash`: additionally RETAIN the gdn-scan inputs (pure buffer keep-alives, zero extra
8431 /// kernels) so a partial accept can rebuild the state after any column prefix (REPLAY-FREE).
8432 #[allow(clippy::too_many_arguments)]
8433 fn linear_attn_verify_t(
8434 &self,
8435 e: &Engine,
8436 la: &LinearAttnLayer,
8437 h: &CudaSlice<f32>,
8438 h_q8: Option<(&CudaSlice<i8>, &CudaSlice<f32>)>,
8439 t: usize,
8440 cache: &mut Cache,
8441 il: usize,
8442 want_stash: bool,
8443 ) -> Result<(CudaSlice<f32>, Option<GdnStash>), Box<dyn std::error::Error>> {
8444 let cfg = &self.cfg;
8445 let geometry = la.geometry;
8446 let d_state = geometry.key_head_dim as usize;
8447 let num_k = geometry.key_heads as usize;
8448 let num_v = geometry.value_heads as usize;
8449 let d_conv = geometry.conv_kernel as usize;
8450 let key_dim = d_state * num_k;
8451 let conv_dim = key_dim * 2 + geometry.value_head_dim as usize * num_v;
8452 let eps = cfg.rms_eps;
8453 let scale = 1.0 / (d_state as f32).sqrt();
8454
8455 // DECODE-EXACT projections: matmul_decode_exact forces the MMVQ (warp-per-row, 32-thread)
8456 // accumulation order for EVERY m, matching the T=1 decode path bit-for-bit. The generic
8457 // `matmul` at m>=5 falls to dp4a (128-thread, two-level reduce) which has a different FP
8458 // sum order — ULP differences propagate through gdn_scan and flip argmax on the 27B.
8459 // Q8 TRUNK-FUSION at T=1 (35B: wqkv+wqkv_gate both Q8_0): one fused2 launch, bit-identical
8460 // per (tensor,row) to the two m=1 MMVQ dispatches below — decode-exact contract holds.
8461 // VERIFY-TIER TRUNK FUSION (MEMRA_SPEC_FUSED_T, t=2-4): quantize h ONCE for every
8462 // fused-eligible same-input Q8_0 pair of this layer (35B wqkv+wqkv_gate; 9B
8463 // ssm_beta+ssm_alpha) — each fused2 batched launch then replaces two decode-exact
8464 // calls (each of which re-quantizes the same h + runs its own _b2/_b4 launch).
8465 // Bit-identical per (tensor,token,row) — see spec_fused_t().
8466 // BATCHED EPILOGUE RE-FUSE (lane/vt-fixes fix 2): `h_q8` = the attn-input norm emitted
8467 // directly as q8_1 by the caller's fused rms_norm_q8_1 (bit-identical to the unfused
8468 // chain, kernel-check-pinned). When present it REPLACES the standalone quantize below
8469 // and feeds every projection; the caller guaranteed all four input projections are
8470 // q8_1-fast. When absent, the old shared-quantize (fused-t window) stands.
8471 let h_q8_t = if h_q8.is_none()
8472 && spec_fused_t()
8473 && (2..=4).contains(&t)
8474 && ((e.uses_q8_1_fast(&la.wqkv) && e.uses_q8_1_fast(&la.wqkv_gate))
8475 || (e.uses_q8_1_fast(&la.ssm_beta) && e.uses_q8_1_fast(&la.ssm_alpha)))
8476 {
8477 Some(e.quantize_q8_1(h, t, cfg.n_embd as usize)?)
8478 } else {
8479 None
8480 };
8481 // one view: the caller's fused-norm q8 or this fn's own shared quantize.
8482 let hq8_any: Option<(&CudaSlice<i8>, &CudaSlice<f32>)> =
8483 h_q8.or(h_q8_t.as_ref().map(|(q, d)| (q, d)));
8484 let (qkv_mixed, z) = {
8485 let mut fused = None;
8486 if t == 1 && e.uses_q8_1_fast(&la.wqkv) && e.uses_q8_1_fast(&la.wqkv_gate) {
8487 let (hq, hd) = e.quantize_q8_1(h, 1, cfg.n_embd as usize)?;
8488 fused = e.matmul_q8_fused2(&la.wqkv, &la.wqkv_gate, &hq, &hd)?;
8489 } else if let Some((hq, hd)) = hq8_any {
8490 if spec_fused_t() && (2..=4).contains(&t) {
8491 fused = e.matmul_q8_fused2_t(&la.wqkv, &la.wqkv_gate, hq, hd, t)?;
8492 }
8493 }
8494 match (fused, hq8_any) {
8495 (Some(pair), _) => pair,
8496 (None, Some((hq, hd))) if h_q8.is_some() => (
8497 e.matmul_decode_exact_pre(&la.wqkv, hq, hd, t)?,
8498 e.matmul_decode_exact_pre(&la.wqkv_gate, hq, hd, t)?,
8499 ),
8500 (None, _) => (
8501 e.matmul_decode_exact(&la.wqkv, h, t)?,
8502 e.matmul_decode_exact(&la.wqkv_gate, h, t)?,
8503 ),
8504 }
8505 };
8506 // beta+alpha DUAL at T=1 (75% of p3 rounds run T=1 verify — p-min chain cuts): the dual
8507 // mr2 kernel is bit-identical per element to the m=1 MMVQ matmul_decode_exact dispatches
8508 // (same warp-per-row body, blockIdx.y picks the weight), so the decode-exact contract
8509 // holds; the run-spec battery is the arbiter. T>1 keeps the per-tensor decode-exact path.
8510 let (beta_raw, alpha) = if t == 1 {
8511 let (hq, hd) = e.quantize_q8_1(h, 1, cfg.n_embd as usize)?;
8512 match e.matmul_pre_dual_noscale(&la.ssm_beta, &la.ssm_alpha, &hq, &hd, 1)? {
8513 Some(((mut b, bs), (mut a, as_))) => {
8514 if bs != 1.0 {
8515 e.scale_inplace(&mut b, bs, la.ssm_beta.out_features())?;
8516 }
8517 if as_ != 1.0 {
8518 e.scale_inplace(&mut a, as_, la.ssm_alpha.out_features())?;
8519 }
8520 (b, a)
8521 }
8522 // Q8_0 fused2 twin (9B stores beta/alpha as Q8_0): DISPATCH-MIRRORS the eager
8523 // decode's beta_alpha closure — the fused body is qmatvec_q8_0_mmvq verbatim,
8524 // bit-identical per row (kernel-check rel=0.00e0 gate), so decode==verify holds.
8525 None => match e.matmul_q8_fused2(&la.ssm_beta, &la.ssm_alpha, &hq, &hd)? {
8526 Some((b, a)) => (b, a),
8527 None => (
8528 e.matmul_decode_exact(&la.ssm_beta, h, 1)?,
8529 e.matmul_decode_exact(&la.ssm_alpha, h, 1)?,
8530 ),
8531 },
8532 }
8533 } else {
8534 // fused-t twin (9B stores beta/alpha as Q8_0): same shared-quantize + one launch
8535 // contract as the wqkv pair above; 35B beta/alpha are Float -> None -> fallback.
8536 let mut nvfp4_fused = None;
8537 let mut q8_fused = None;
8538 if let Some((hq, hd)) = hq8_any {
8539 if t == 3 && std::env::var("MEMRA_NVFP4_AUX_DUAL").as_deref() != Ok("0") {
8540 nvfp4_fused =
8541 e.matmul_decode_exact_dual_pre(&la.ssm_beta, &la.ssm_alpha, hq, hd, t)?;
8542 if nvfp4_fused.is_some() && std::env::var("MEMRA_DEBUG").is_ok() {
8543 static ONCE: std::sync::Once = std::sync::Once::new();
8544 ONCE.call_once(|| {
8545 eprintln!("[memra] NVFP4 beta+alpha batched aux dual ENGAGED (t={t})")
8546 });
8547 }
8548 }
8549 if nvfp4_fused.is_none() && spec_fused_t() && (2..=4).contains(&t) {
8550 q8_fused = e.matmul_q8_fused2_t(&la.ssm_beta, &la.ssm_alpha, hq, hd, t)?;
8551 }
8552 }
8553 if let Some(((mut b, bs), (mut a, as_))) = nvfp4_fused {
8554 if bs != 1.0 {
8555 e.scale_inplace(&mut b, bs, t * la.ssm_beta.out_features())?;
8556 }
8557 if as_ != 1.0 {
8558 e.scale_inplace(&mut a, as_, t * la.ssm_alpha.out_features())?;
8559 }
8560 (b, a)
8561 } else if let Some(pair) = q8_fused {
8562 pair
8563 } else {
8564 match hq8_any {
8565 Some((hq, hd)) if h_q8.is_some() => (
8566 e.matmul_decode_exact_pre(&la.ssm_beta, hq, hd, t)?,
8567 e.matmul_decode_exact_pre(&la.ssm_alpha, hq, hd, t)?,
8568 ),
8569 _ => (
8570 e.matmul_decode_exact(&la.ssm_beta, h, t)?,
8571 e.matmul_decode_exact(&la.ssm_alpha, h, t)?,
8572 ),
8573 }
8574 }
8575 };
8576
8577 // conv with CARRIED state + ring roll (T >= pad rides the input-column update kernel;
8578 // T < pad — the MEMRA_SPEC_M2 t=2 arm — rolls via the pure-copy ring rebuild).
8579 let rl = cache.recur[il].as_mut().unwrap();
8580 let mut conv_out = e.uninit(conv_dim * t)?;
8581 e.ssm_conv1d_tm_state(
8582 &qkv_mixed,
8583 &mut rl.conv_state,
8584 la.ssm_conv1d.float_data(),
8585 &mut conv_out,
8586 conv_dim,
8587 t,
8588 d_conv,
8589 )?;
8590
8591 // GDN prep via the prefill kernels (repack + L2 + sigmoid + glog), T-wide.
8592 let mut q_g = e.uninit(d_state * num_v * t)?;
8593 let mut k_g = e.uninit(d_state * num_v * t)?;
8594 let mut v_g = e.uninit(d_state * num_v * t)?;
8595 e.qkv_to_gdn_repack(
8596 &conv_out, &mut q_g, &mut k_g, &mut v_g, d_state, num_v, num_k, key_dim, t,
8597 )?;
8598 let mut q_l2 = e.uninit(d_state * num_v * t)?;
8599 e.l2_norm_decode(&q_g, &mut q_l2, d_state, num_v * t, eps)?;
8600 let mut k_l2 = e.uninit(d_state * num_v * t)?;
8601 e.l2_norm_decode(&k_g, &mut k_l2, d_state, num_v * t, eps)?;
8602 let mut beta = e.uninit(t * num_v)?;
8603 e.sigmoid(&beta_raw, &mut beta, t * num_v)?;
8604 let mut g_log = e.uninit(t * num_v)?;
8605 e.gdn_glog(
8606 &alpha,
8607 la.ssm_dt.float_data(),
8608 la.ssm_a.float_data(),
8609 &mut g_log,
8610 num_v,
8611 t,
8612 )?;
8613
8614 // ONE gdn_scan over T tokens from the carried state (internal sequential loop ==
8615 // T chained T=1 steps). Ping-pong the resident buffers like eager decode.
8616 let mut o = e.uninit(d_state * num_v * t)?;
8617 {
8618 let crate::cache::RecurLayer {
8619 ssm_state,
8620 ssm_state_alt,
8621 ..
8622 } = rl;
8623 e.gdn_scan_s128(
8624 &q_l2,
8625 &k_l2,
8626 &v_g,
8627 &g_log,
8628 &beta,
8629 ssm_state,
8630 ssm_state_alt,
8631 &mut o,
8632 num_v,
8633 t,
8634 scale,
8635 )?;
8636 }
8637 std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
8638
8639 // gated RMSNorm + out projection, T-wide. FUSED-QUANTIZE ARM (lane/vt-fixes fix 2,
8640 // mirroring the T=1 decode's launch-arc form): when ssm_out rides the q8_1 fast path,
8641 // emit q8_1 straight from the gated norm at nrows=num_v*t (row-indexed kernel, the
8642 // T-wide launch is the per-row program; kernel-check pins bit-identity vs
8643 // gated_rmsnorm -> quantize_q8_1 at T=1 and T=5) and feed the decode-exact dispatch
8644 // pre-quantized — one launch replaces norm + quantize. Fallback = the f32 chain.
8645 let out = if e.uses_q8_1_fast(&la.ssm_out) {
8646 let (gq, gd) =
8647 e.gated_rmsnorm_q8_1(&o, la.ssm_norm.float_data(), &z, d_state, num_v * t, eps)?;
8648 e.matmul_decode_exact_pre(&la.ssm_out, &gq, &gd, t)?
8649 } else {
8650 let mut gn = e.uninit(d_state * num_v * t)?;
8651 e.gated_rmsnorm(
8652 &o,
8653 la.ssm_norm.float_data(),
8654 &z,
8655 &mut gn,
8656 d_state,
8657 num_v * t,
8658 eps,
8659 )?;
8660 // DECODE-EXACT out-projection: same MMVQ path as the T=1 decode (ssm_out at m>=5
8661 // would fall to dp4a with a different FP reduction order — same class of bug as
8662 // the input projs).
8663 e.matmul_decode_exact(&la.ssm_out, &gn, t)?
8664 };
8665 let stash = if want_stash {
8666 Some(GdnStash {
8667 qkv_mixed,
8668 q_l2,
8669 k_l2,
8670 v_g,
8671 g_log,
8672 beta,
8673 })
8674 } else {
8675 None
8676 };
8677 Ok((out, stash))
8678 }
8679
8680 /// REPLAY-FREE partial-accept commit (2026-07-03): make the cache state == "committed through
8681 /// the first `j` verify columns" WITHOUT the legacy rollback + duplicate trunk replay.
8682 /// - Full-attn KV: truncate both the owning-stage shadow and every TP rank to snapshot + j.
8683 /// The verify's appended rows for those columns are bit-identical to what an eager T=1
8684 /// chain writes (the decode-exact contract the verify-probe gates), so keeping them ==
8685 /// replaying them.
8686 /// - Linear layers, batched path: rebuild the conv ring by PURE COPIES (ring holds raw input
8687 /// columns) and the ssm state by a prefix re-run of the SAME gdn_scan kernel (t=j) from the
8688 /// snapshot state over the stash's identical inputs — the kernel's t-loop carries state in
8689 /// registers and writes it once at the end, so iterations 0..j-1 are independent of T:
8690 /// bit-identical to the verify's own state after j tokens == the eager chain state.
8691 /// - Linear layers, per-column path: restore the cloned actual state after column j-1.
8692 /// Caller guarantees 1 <= j <= t-1 (j==0 rounds take the legacy rollback; j==t is full accept).
8693 fn commit_verified_prefix(
8694 &self,
8695 e: &Engine,
8696 cache: &mut Cache,
8697 snap: &crate::cache::CacheSnapshot,
8698 ckpt: &VerifyCkpt,
8699 j: usize,
8700 kv_lens_done: bool,
8701 dev_j: Option<(&CudaSlice<u32>, usize, usize)>,
8702 ) -> Result<(), Box<dyn std::error::Error>> {
8703 // GDN geometry derives lazily inside recurrent-layer arms. Full-attention plans carry no
8704 // recurrent state and must never be forced through a synthetic SSM geometry.
8705 // Engine-bundle slice 1 (DSF-ROUNDCOST-20260820 §1.1): the per-column-arm restores
8706 // are 2 tiny D2D copies per linear layer (~96 dispatches/partial round on the q38
8707 // route). When every cols-arm layer shares uniform state sizes (single ssm cfg —
8708 // always true today), batch them into two `copy_batch_uniform_f32` launches. Bytes,
8709 // buffers and stream order are identical to the per-layer memcpy sequence; the
8710 // kernel-rebuild (gdn-stash) arm below is untouched. MEMRA_STATE_COPY_BATCH=0 reverts.
8711 let mut batched_cols = false;
8712 if state_copy_batch_on() && dev_j.is_none() {
8713 use cudarc::driver::DevicePtr;
8714 let s = &e.gpu.stream();
8715 let mut conv_pairs: Vec<(u64, u64)> = Vec::new();
8716 let mut ssm_pairs: Vec<(u64, u64)> = Vec::new();
8717 let (mut conv_words, mut ssm_words) = (0usize, 0usize);
8718 let mut uniform = true;
8719 for il in 0..self.layers.len() {
8720 let Some(rl) = cache.recur[il].as_ref() else {
8721 continue;
8722 };
8723 if ckpt.gdn[il].is_some() {
8724 continue; // kernel-rebuild arm restores below, per layer
8725 }
8726 let Some(cols) = &ckpt.cols[il] else {
8727 continue; // missing-ckpt error surfaces in the main loop
8728 };
8729 let (c, st) = &cols[j - 1];
8730 if conv_pairs.is_empty() {
8731 conv_words = c.len();
8732 ssm_words = st.len();
8733 } else if c.len() != conv_words || st.len() != ssm_words {
8734 uniform = false;
8735 break;
8736 }
8737 let (pc, _g0) = c.device_ptr(s);
8738 let (dc, _g1) = rl.conv_state.device_ptr(s);
8739 let (ps, _g2) = st.device_ptr(s);
8740 let (ds, _g3) = rl.ssm_state.device_ptr(s);
8741 conv_pairs.push((pc as u64, dc as u64));
8742 ssm_pairs.push((ps as u64, ds as u64));
8743 }
8744 if uniform && !conv_pairs.is_empty() {
8745 let n = conv_pairs.len();
8746 let mut t = vec![0u64; 2 * n];
8747 for (k, &(src, dst)) in conv_pairs.iter().enumerate() {
8748 t[k] = src;
8749 t[n + k] = dst;
8750 }
8751 let conv_t = e.htod_u64(&t)?;
8752 for (k, &(src, dst)) in ssm_pairs.iter().enumerate() {
8753 t[k] = src;
8754 t[n + k] = dst;
8755 }
8756 let ssm_t = e.htod_u64(&t)?;
8757 e.copy_batch_uniform_f32(&conv_t, n, conv_words)?;
8758 e.copy_batch_uniform_f32(&ssm_t, n, ssm_words)?;
8759 batched_cols = true;
8760 }
8761 }
8762 rewind_tp_kv_verified_prefix(&mut cache.tp_kv, &snap.tp_kv_len, j)?;
8763 for il in 0..self.layers.len() {
8764 if let (Some(kvl), Some(saved)) = (cache.kv[il].as_mut(), snap.kv_len[il]) {
8765 kvl.len = saved + j;
8766 // devacc 3a: spec_rollback_kv already wrote len_d on-device (same value).
8767 if !kv_lens_done {
8768 e.set_i32_one(&mut kvl.len_d, kvl.len as i32)?;
8769 }
8770 }
8771 if let Some(rl) = cache.recur[il].as_mut() {
8772 let Mixer::Linear(linear) = &self.layers[il].mixer else {
8773 return Err(format!("recurrent cache layer {il} has no GDN plan").into());
8774 };
8775 let geometry = linear.geometry;
8776 let d_state = geometry.key_head_dim as usize;
8777 let num_k = geometry.key_heads as usize;
8778 let num_v = geometry.value_heads as usize;
8779 let d_conv = geometry.conv_kernel as usize;
8780 let conv_dim = d_state * num_k * 2 + geometry.value_head_dim as usize * num_v;
8781 let scale = 1.0 / (d_state as f32).sqrt();
8782 if let Some(st) = &ckpt.gdn[il] {
8783 let ring_old = snap.conv[il].as_ref().expect("snapshot missing conv");
8784 let state_in = snap.ssm[il].as_ref().expect("snapshot missing ssm");
8785 if let Some((acc, base, t_v)) = dev_j {
8786 // 3b: j read on-device (_dc twins, same bodies; full accept early-exits).
8787 e.ssm_conv_ring_rebuild_dc(
8788 &st.qkv_mixed,
8789 ring_old,
8790 &mut rl.conv_state,
8791 conv_dim,
8792 acc,
8793 base,
8794 t_v,
8795 d_conv,
8796 )?;
8797 let mut o = e.uninit(d_state * num_v * j.max(1))?;
8798 e.gdn_scan_s128_dc(
8799 &st.q_l2,
8800 &st.k_l2,
8801 &st.v_g,
8802 &st.g_log,
8803 &st.beta,
8804 state_in,
8805 &mut rl.ssm_state,
8806 &mut o,
8807 num_v,
8808 acc,
8809 base,
8810 t_v,
8811 scale,
8812 )?;
8813 } else {
8814 e.ssm_conv_ring_rebuild(
8815 &st.qkv_mixed,
8816 ring_old,
8817 &mut rl.conv_state,
8818 conv_dim,
8819 j,
8820 d_conv,
8821 )?;
8822 let mut o = e.uninit(d_state * num_v * j)?; // scan output, discarded
8823 e.gdn_scan_s128(
8824 &st.q_l2,
8825 &st.k_l2,
8826 &st.v_g,
8827 &st.g_log,
8828 &st.beta,
8829 state_in,
8830 &mut rl.ssm_state,
8831 &mut o,
8832 num_v,
8833 j,
8834 scale,
8835 )?;
8836 }
8837 } else if let Some(cols) = &ckpt.cols[il] {
8838 if !batched_cols {
8839 let (c, s) = &cols[j - 1];
8840 e.copy_into(&mut rl.conv_state, 0, c, c.len())?;
8841 e.copy_into(&mut rl.ssm_state, 0, s, s.len())?;
8842 }
8843 } else {
8844 return Err(
8845 "commit_verified_prefix: verify ckpt missing for linear layer".into(),
8846 );
8847 }
8848 }
8849 }
8850 cache.pos = snap.pos + j;
8851 Ok(())
8852 }
8853
8854 /// ROUND-STREAM: recur restore with device-j (the _dc twins; full accept early-exits
8855 /// in-kernel). Requires the batched-linear stash on every linear layer (stream gate).
8856 fn commit_verified_prefix_stream(
8857 &self,
8858 e: &Engine,
8859 cache: &mut Cache,
8860 snap: &crate::cache::CacheSnapshot,
8861 ckpt: &VerifyCkpt,
8862 acc: &CudaSlice<u32>,
8863 base: usize,
8864 t_v: usize,
8865 ) -> Result<(), Box<dyn std::error::Error>> {
8866 for il in 0..self.layers.len() {
8867 if let Some(rl) = cache.recur[il].as_mut() {
8868 let Mixer::Linear(linear) = &self.layers[il].mixer else {
8869 return Err(format!("recurrent cache layer {il} has no GDN plan").into());
8870 };
8871 let geometry = linear.geometry;
8872 let d_state = geometry.key_head_dim as usize;
8873 let num_k = geometry.key_heads as usize;
8874 let num_v = geometry.value_heads as usize;
8875 let d_conv = geometry.conv_kernel as usize;
8876 let conv_dim = d_state * num_k * 2 + geometry.value_head_dim as usize * num_v;
8877 let scale = 1.0 / (d_state as f32).sqrt();
8878 let st = ckpt.gdn[il]
8879 .as_ref()
8880 .ok_or("stream restore: batched-linear stash missing")?;
8881 let ring_old = snap.conv[il].as_ref().expect("snapshot missing conv");
8882 let state_in = snap.ssm[il].as_ref().expect("snapshot missing ssm");
8883 e.ssm_conv_ring_rebuild_dc(
8884 &st.qkv_mixed,
8885 ring_old,
8886 &mut rl.conv_state,
8887 conv_dim,
8888 acc,
8889 base,
8890 t_v,
8891 d_conv,
8892 )?;
8893 let mut o = e.uninit(d_state * num_v * t_v)?;
8894 e.gdn_scan_s128_dc(
8895 &st.q_l2,
8896 &st.k_l2,
8897 &st.v_g,
8898 &st.g_log,
8899 &st.beta,
8900 state_in,
8901 &mut rl.ssm_state,
8902 &mut o,
8903 num_v,
8904 acc,
8905 base,
8906 t_v,
8907 scale,
8908 )?;
8909 }
8910 }
8911 Ok(())
8912 }
8913
8914 /// EAGLE3 aux-capturing verify forward over `tokens` (T) — mirrors `decode_step_t_h` exactly
8915 /// (same KV append, same causal verify, same recur advance) but ALSO clones the aux residual-
8916 /// stream hiddens (blocks in `aux_layers`) for TWO columns: the LAST column (always) and the
8917 /// optional `pred_col` (the EAGLE seed = bonus's predecessor). Returns
8918 /// (all_T_logits host, last_col_aux, pred_col_aux?). Used by the EAGLE3 orchestrator's commit.
8919 pub fn decode_step_t_aux2(
8920 &self,
8921 e: &Engine,
8922 tokens: &[u32],
8923 pos0: usize,
8924 cache: &mut Cache,
8925 aux_layers: &[usize],
8926 pred_col: Option<usize>,
8927 ) -> Result<
8928 (Vec<f32>, Vec<CudaSlice<f32>>, Option<Vec<CudaSlice<f32>>>),
8929 Box<dyn std::error::Error>,
8930 > {
8931 let cfg = &self.cfg;
8932 let n_embd = cfg.n_embd as usize;
8933 let eps = cfg.rms_eps;
8934 let t = tokens.len();
8935 let pos_vec: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
8936 let pos_d = e.htod_i32(&pos_vec)?;
8937 let mut x = e.htod(&self.embd.gather(n_embd, tokens))?;
8938 let mut aux_last: Vec<CudaSlice<f32>> = Vec::with_capacity(aux_layers.len());
8939 let mut aux_pred: Vec<CudaSlice<f32>> = Vec::new();
8940 let want_pred = pred_col.is_some();
8941
8942 for (il, layer) in self.layers.iter().enumerate() {
8943 // DISPATCH-MIRRORED norms (FP-order lesson #8) — see decode_step_t_h_emb.
8944 let mixer_fast = self.mixer_in_q8_1_fast(e, &layer.mixer);
8945 let norm_fused = std::env::var("MEMRA_NO_FUSE_NORMQ").is_err() && mixer_fast;
8946 let mut h = vbuf(e, t * n_embd)?; // fully written by either rms_norm arm
8947 if norm_fused {
8948 e.rms_norm_decode(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
8949 } else {
8950 e.rms_norm(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
8951 }
8952 let mixed = match &layer.mixer {
8953 Mixer::Full(fa) => {
8954 self.full_attn_verify(e, fa, &h, None, &pos_d, t, cache, il, None)?
8955 }
8956 Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
8957 Mixer::Linear(la) => {
8958 let mut out = e.zeros(t * n_embd)?;
8959 for col in 0..t {
8960 let mut h_col = e.zeros(n_embd)?;
8961 let src = h.slice(col * n_embd..(col + 1) * n_embd);
8962 e.copy_view_into(&mut h_col, 0, &src, n_embd)?;
8963 let m_col = self.linear_attn_decode(e, la, &h_col, cache, il)?;
8964 e.copy_into(&mut out, col * n_embd, &m_col, n_embd)?;
8965 }
8966 out
8967 }
8968 };
8969 let ffn_fuse = match &layer.ffn {
8970 crate::hybrid::Ffn::Dense {
8971 ffn_gate, ffn_up, ..
8972 } => {
8973 std::env::var("MEMRA_NO_FUSE_NORMQ").is_err()
8974 && e.uses_q8_1_fast(ffn_gate)
8975 && e.uses_q8_1_fast(ffn_up)
8976 }
8977 crate::hybrid::Ffn::Moe(_) => false,
8978 };
8979 let mut x1 = vbuf(e, t * n_embd)?; // fully written by add / add_rms_norm
8980 let mut z = vbuf(e, t * n_embd)?; // fully written by rms_norm_decode / add_rms_norm
8981 if ffn_fuse {
8982 e.add(&x, &mixed, &mut x1, t * n_embd)?;
8983 e.rms_norm_decode(
8984 &x1,
8985 layer.post_attn_norm.float_data(),
8986 &mut z,
8987 n_embd,
8988 t,
8989 eps,
8990 )?;
8991 } else {
8992 e.add_rms_norm(
8993 &x,
8994 &mixed,
8995 layer.post_attn_norm.float_data(),
8996 &mut x1,
8997 &mut z,
8998 n_embd,
8999 t,
9000 eps,
9001 )?;
9002 }
9003 let ffn_out = match &layer.ffn {
9004 crate::hybrid::Ffn::Dense {
9005 ffn_gate,
9006 ffn_up,
9007 ffn_down,
9008 } => {
9009 let n_ff = ffn_gate.out_features();
9010 let gate = e.matmul_decode_exact(ffn_gate, &z, t)?;
9011 let up = e.matmul_decode_exact(ffn_up, &z, t)?;
9012 let mut act = vbuf(e, t * n_ff)?; // fully written by ffn_act_lim
9013 // dense FFN clamp = the SHEXP array (upstream build_ffn serves both).
9014 Self::ffn_act_lim(
9015 e,
9016 &self.cfg,
9017 &gate,
9018 &up,
9019 1.0,
9020 1.0,
9021 self.cfg.clamp_shexp_at(il as u32),
9022 &mut act,
9023 t * n_ff,
9024 )?;
9025 e.matmul_decode_exact(ffn_down, &act, t)?
9026 }
9027 crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il(e, m, &z, t, il as u16)?,
9028 };
9029 let mut x2 = vbuf(e, t * n_embd)?; // fully written by add
9030 e.add(&x1, &ffn_out, &mut x2, t * n_embd)?;
9031 if aux_layers.contains(&il) {
9032 let mut a = e.zeros(n_embd)?;
9033 e.copy_view_into(&mut a, 0, &x2.slice((t - 1) * n_embd..t * n_embd), n_embd)?;
9034 aux_last.push(a);
9035 if let Some(pc) = pred_col {
9036 let mut ap = e.zeros(n_embd)?;
9037 e.copy_view_into(
9038 &mut ap,
9039 0,
9040 &x2.slice(pc * n_embd..(pc + 1) * n_embd),
9041 n_embd,
9042 )?;
9043 aux_pred.push(ap);
9044 }
9045 }
9046 x = x2;
9047 }
9048 let mut hn = vbuf(e, t * n_embd)?; // fully written by rms_norm_decode
9049 e.rms_norm_decode(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
9050 let logits = e.matmul_decode_exact(&self.output, &hn, t)?;
9051 let host = e.dtoh(&logits)?;
9052 cache.pos += t;
9053 Ok((
9054 host,
9055 aux_last,
9056 if want_pred { Some(aux_pred) } else { None },
9057 ))
9058 }
9059
9060 /// step35 SPEC-VERIFY attention over T query tokens — a per-row REPLAY of the eager
9061 /// `step35_decode_attn`.
9062 ///
9063 /// WHY A REPLAY AND NOT A BATCHED TWIN. The verify's whole job is to be bit-identical to what
9064 /// the eager decode would have computed for the same tokens; that is what makes greedy spec
9065 /// decode exact (run-spec asserts token identity for K=1..8). Every other verify arm in this
9066 /// file earns that identity by carefully mirroring dispatch (`matmul_decode_exact` to force
9067 /// MMVQ at any m, per-layer `ffn_fuse` mirroring, per-row `fa_decode` key bounds). step35
9068 /// stacks FOUR more per-layer degrees of freedom on top of that — per-layer `n_head`
9069 /// (64 full / 96 SWA), per-layer rotary width (64 full / 128 SWA), per-layer rope base, and a
9070 /// SEPARATE `attn_gate` tensor whose projection shares the attn-normed input — and its SWA
9071 /// layers attend through a token-OFFSET view whose offset is a function of the ABSOLUTE
9072 /// position of each query row. A batched twin would have to reproduce all of that AND the
9073 /// per-row offset in one launch; the offset alone rules out the existing rows kernels (they
9074 /// take one `base_len`, not a per-row offset).
9075 ///
9076 /// So this arm calls the eager path itself, once per row, on the same cache. Identity is then
9077 /// true BY CONSTRUCTION rather than by mirroring: row r runs exactly the kernel sequence that
9078 /// eager decode step r runs (same projections, same q8_1 fusion decision, same append, same
9079 /// view arithmetic, same `fa_decode_kvmod`, same gate), because it IS that code. Cost: T x the
9080 /// eager decode mixer instead of one batched pass — the same trade the generic arm's `else`
9081 /// per-row loop already accepts when `fa_rows_eligible` says no. Correctness first; a batched
9082 /// step35 twin is a perf lane's job and must be gated against this arm.
9083 ///
9084 /// The `h_q8` pre-quantized pair from the caller's fused norm is NOT forwarded: it is a
9085 /// T-row buffer and `step35_decode_attn`'s `pre_q` contract is one row. Instead each row's
9086 /// f32 `h` slice is handed over and the callee re-derives its own q8_1 exactly as eager decode
9087 /// does (`quantize_q8_1(h, 1, n_embd)`) — which is the dispatch being mirrored. Callers that
9088 /// took the fused arm therefore MUST still pass a live `h`; `step35_verify` asserts that.
9089 #[allow(clippy::too_many_arguments)]
9090 fn step35_verify(
9091 &self,
9092 e: &Engine,
9093 fa: &FullAttnLayer,
9094 h: &CudaSlice<f32>,
9095 h_q8: Option<(&CudaSlice<i8>, &CudaSlice<f32>)>,
9096 t: usize,
9097 cache: &mut Cache,
9098 il: usize,
9099 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9100 let n_embd = self.cfg.n_embd as usize;
9101 // The fused (h-less) attn-norm arm hands `h` as a zero-length placeholder. This arm needs
9102 // the f32 rows, so the caller must not take that lever for step35 — enforced at the call
9103 // site by the sliding-gated-MoE `Mixer::Full(_) => false` arm of
9104 // `lin_q8_only` in `decode_step_t_core_stream`, and asserted here so a future caller
9105 // cannot regress it into silently reading an empty buffer.
9106 assert_eq!(
9107 h.len(),
9108 t * n_embd,
9109 "step35_verify needs the f32 attn-normed rows ([t*n_embd]); the caller took the \
9110 fused q8-only norm arm (h_q8={}) — step35 must stay on the unfused arm",
9111 h_q8.is_some()
9112 );
9113 // ROW WIDTH IS n_embd, NOT n_head*head_dim: `step35_decode_attn` returns the mixer output
9114 // AFTER `wo`, so a row is [n_embd] — the same contract the generic arm's
9115 // `matmul_decode_exact(&fa.wo, &attn_g, t)` return has. Sizing this buffer from the
9116 // per-layer head geometry (8192 on full-attn, 12288 on SWA) instead overran the row on the
9117 // FIRST copy and panicked inside `copy_into`'s `CudaView::slice` unwrap
9118 // (raw/mtp-bt-20260806T212127Z.log frames 12-13).
9119 let mut out = vbuf(e, t * n_embd)?; // each row fully written by the copy below
9120 for r in 0..t {
9121 // Absolute position of this query row. `cache.pos` is the committed length at round
9122 // start and every row before r has already been appended by this loop, so the r-th
9123 // verify token sits at cache.pos + r — the same position eager decode would give it.
9124 let pos_d = e.htod_i32(&[(cache.pos + r) as i32])?;
9125 let mut h_row = vbuf(e, n_embd)?; // fully written by copy_view_into
9126 e.copy_view_into(
9127 &mut h_row,
9128 0,
9129 &h.slice(r * n_embd..(r + 1) * n_embd),
9130 n_embd,
9131 )?;
9132 // THE eager decode mixer: appends this row's K/V at kvl.len, advances it, then
9133 // attends over the (SWA-offset) view. Post-`wo`, same contract as this fn returns.
9134 let o = self.step35_decode_attn(e, fa, il, &h_row, None, &pos_d, cache)?;
9135 debug_assert_eq!(
9136 o.len(),
9137 n_embd,
9138 "step35_decode_attn returns post-wo [n_embd]"
9139 );
9140 e.copy_into(&mut out, r * n_embd, &o, n_embd)?;
9141 }
9142 Ok(out)
9143 }
9144
9145 /// Full-attention mixer over T query tokens with a GROWING resident KV (verify path, §D.3).
9146 /// Appends the T new K/V columns to cache.kv[il] then attends causally over [0..len) via
9147 /// fa_prefill. Token-major [T, kv_dim] projection layout == cache row layout (single copy).
9148 #[allow(clippy::too_many_arguments)]
9149 fn full_attn_verify(
9150 &self,
9151 e: &Engine,
9152 fa: &FullAttnLayer,
9153 h: &CudaSlice<f32>,
9154 h_q8: Option<(&CudaSlice<i8>, &CudaSlice<f32>)>,
9155 pos_d: &CudaSlice<i32>,
9156 t: usize,
9157 cache: &mut Cache,
9158 il: usize,
9159 stream_ctr: Option<&CudaSlice<i32>>,
9160 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9161 // step35: the generic geometry below is wrong for this arch (per-layer n_head, partial
9162 // per-layer rope, the SWA offset view, and a SEPARATE head-wise gate tensor), so it takes
9163 // its own arm. A verify that silently computes different attention than decode defeats the
9164 // whole self-consistency gate, so the arm is a per-row REPLAY of `step35_decode_attn`
9165 // rather than a batched twin — see `step35_verify` for why that is the exactness-correct
9166 // shape and not laziness.
9167 if self.sliding_gated_moe_batch_program() {
9168 if stream_ctr.is_some() {
9169 return Err(
9170 "step35 has no ROUND-STREAM verify arm (the device-counter _dc twins \
9171 cannot express the SWA offset KV view; same root cause as the dc \
9172 decode refusal) — run spec without the stream arm"
9173 .into(),
9174 );
9175 }
9176 return self.step35_verify(e, fa, h, h_q8, t, cache, il);
9177 }
9178 let cfg = &self.cfg;
9179 let geometry = cfg.full_attention_geometry_at(il as u32);
9180 let n_head = geometry.n_head as usize;
9181 let n_head_kv = geometry.n_head_kv as usize;
9182 let head_dim = geometry.head_dim_k as usize;
9183 let eps = cfg.rms_eps;
9184 let scale = geometry.attention_scale();
9185 let n_embd = cfg.n_embd as usize;
9186
9187 // DECODE-EXACT Q/K/V projections: matmul_decode_exact forces the MMVQ (warp-per-row) path
9188 // for every m, matching the T=1 decode's FP accumulation order. matmul_pre at m>=5 would
9189 // fall to dp4a (128-thread, two-level reduce) with a different FP sum order.
9190 // Q8 TRUNK-FUSION at T=1: DISPATCH-MIRRORS the eager decode's fused3 (bit-identical body).
9191 // BATCHED EPILOGUE RE-FUSE (lane/vt-fixes fix 2): `h_q8` = the attn-input norm's q8_1
9192 // form emitted by the fused rms_norm_q8_1 (bit-identical to rms_norm_decode ->
9193 // quantize_q8_1, kernel-check-pinned). When present (caller checked mixer q8_1-fast),
9194 // every projection consumes it — `h` may be a zero-len placeholder and must not be read.
9195 let (qf, mut k, v) = {
9196 let mut fused = None;
9197 let qkv_fast =
9198 e.uses_q8_1_fast(&fa.wq) && e.uses_q8_1_fast(&fa.wk) && e.uses_q8_1_fast(&fa.wv);
9199 if t == 1 && qkv_fast {
9200 let (hq_o, hd_o);
9201 let (hq, hd): (&CudaSlice<i8>, &CudaSlice<f32>) = match h_q8 {
9202 Some(p) => p,
9203 None => {
9204 (hq_o, hd_o) = e.quantize_q8_1(h, 1, n_embd)?;
9205 (&hq_o, &hd_o)
9206 }
9207 };
9208 fused = e.matmul_q8_fused3(&fa.wq, &fa.wk, &fa.wv, hq, hd)?;
9209 } else if spec_fused_t() && (2..=4).contains(&t) && qkv_fast {
9210 // VERIFY-TIER TRUNK FUSION (MEMRA_SPEC_FUSED_T): one shared quantize + one
9211 // fused3 batched launch replaces three decode-exact calls (3 re-quantizes of
9212 // the same h + 3 _b2/_b4 launches). Bit-identical per (tensor,token,row).
9213 let (hq_o, hd_o);
9214 let (hq, hd): (&CudaSlice<i8>, &CudaSlice<f32>) = match h_q8 {
9215 Some(p) => p,
9216 None => {
9217 (hq_o, hd_o) = e.quantize_q8_1(h, t, n_embd)?;
9218 (&hq_o, &hd_o)
9219 }
9220 };
9221 fused = e.matmul_q8_fused3_t(&fa.wq, &fa.wk, &fa.wv, hq, hd, t)?;
9222 }
9223 match (fused, h_q8) {
9224 (Some(triple), _) => triple,
9225 // shared pre-quantized activation (q8_1-fast guaranteed by the caller): the
9226 // decode-exact dispatch consumes (hq, hd) instead of re-quantizing 3x.
9227 (None, Some((hq, hd))) if qkv_fast => (
9228 e.matmul_decode_exact_pre(&fa.wq, hq, hd, t)?,
9229 e.matmul_decode_exact_pre(&fa.wk, hq, hd, t)?,
9230 e.matmul_decode_exact_pre(&fa.wv, hq, hd, t)?,
9231 ),
9232 (None, _) => (
9233 e.matmul_decode_exact(&fa.wq, h, t)?,
9234 e.matmul_decode_exact(&fa.wk, h, t)?,
9235 e.matmul_decode_exact(&fa.wv, h, t)?,
9236 ),
9237 }
9238 };
9239 // M3/Hy3 have no attention output gate — wq out is exactly q; skip the split.
9240 let gated = geometry.attention_gate == memra_gguf::config::AttentionGateKind::FusedQ;
9241 let (mut q, gate) = if gated {
9242 let mut q = vbuf(e, t * n_head * head_dim)?; // fully written by q_gate_split
9243 let mut gate = vbuf(e, t * n_head * head_dim)?; // fully written by q_gate_split
9244 e.q_gate_split(&qf, &mut q, &mut gate, head_dim, n_head, t)?;
9245 (q, Some(gate))
9246 } else {
9247 (qf, None)
9248 };
9249
9250 let mut qn = vbuf(e, t * n_head * head_dim)?; // fully written by rms_norm
9251 e.rms_norm(
9252 &q,
9253 fa.q_norm.float_data(),
9254 &mut qn,
9255 head_dim,
9256 n_head * t,
9257 eps,
9258 )?;
9259 q = qn;
9260 let mut kn = vbuf(e, t * n_head_kv * head_dim)?; // fully written by rms_norm
9261 e.rms_norm(
9262 &k,
9263 fa.k_norm.float_data(),
9264 &mut kn,
9265 head_dim,
9266 n_head_kv * t,
9267 eps,
9268 )?;
9269 k = kn;
9270 let rope_dims = geometry.n_rot as usize;
9271 e.rope_neox(
9272 &mut q,
9273 pos_d,
9274 head_dim,
9275 rope_dims,
9276 n_head,
9277 t,
9278 geometry.rope_base,
9279 1.0,
9280 )?;
9281 e.rope_neox(
9282 &mut k,
9283 pos_d,
9284 head_dim,
9285 rope_dims,
9286 n_head_kv,
9287 t,
9288 geometry.rope_base,
9289 1.0,
9290 )?;
9291
9292 // append T new K/V columns to the resident QUANTIZED cache. k/v are token-major [T, kv_dim]
9293 // f32; append-quantize each of the T token rows into the byte cache (q8_0 K / q5_1 V).
9294 let kvl = cache.kv[il].as_mut().unwrap();
9295 let (kv_dim_k, kv_dim_v, ktb, vtb) =
9296 (kvl.kv_dim_k, kvl.kv_dim_v, kvl.k_tok_bytes, kvl.v_tok_bytes);
9297 if let Some(ctr) = stream_ctr {
9298 // stream: ONE batched append at the device counter (rows kernel = the per-view warp
9299 // math on a (block, token) grid, documented byte-identical); host len is a stale
9300 // LOWER BOUND under pre-issue (drain reconciles it).
9301 e.append_kv_quantized_rows_dc(
9302 &k,
9303 &v,
9304 &mut kvl.k,
9305 &mut kvl.v,
9306 ctr,
9307 t,
9308 kv_dim_k,
9309 kv_dim_v,
9310 ktb,
9311 vtb,
9312 crate::Engine::kv_fp8_on(),
9313 )?;
9314 } else {
9315 for i in 0..t {
9316 let k_row = k.slice(i * kv_dim_k..(i + 1) * kv_dim_k);
9317 let v_row = v.slice(i * kv_dim_v..(i + 1) * kv_dim_v);
9318 e.append_kv_quantized_view(
9319 &k_row,
9320 &v_row,
9321 &mut kvl.k,
9322 &mut kvl.v,
9323 kvl.len + i,
9324 kv_dim_k,
9325 kv_dim_v,
9326 ktb,
9327 vtb,
9328 crate::Engine::kv_fp8_on(),
9329 )?;
9330 }
9331 kvl.len += t;
9332 }
9333
9334 // BIT-IDENTICAL VERIFY ATTENTION (spec-exactness fix): the FP accumulation order must be
9335 // byte-for-byte identical to the eager decode path. fa_prefill uses a different tile size
9336 // (BLOCK_Q=64, BK=32) and online-softmax structure than fa_decode's split-K + combine,
9337 // which changes FP summation order and can flip argmax at tight logit margins. Query row r
9338 // attends to keys [0..base_len+r+1) — each successive row sees one more key (the causal
9339 // property). This matches eager: decode appends k at len, then fa_decode sees t_kv = len+1
9340 // keys. The verify appends all T tokens first but bounds the key range per row.
9341 //
9342 // MULTI-ROW FUSED PATH (the long-ctx spec fix, 2026-07-03): when every row takes the vec
9343 // kernel (base_len+1 >= FA_VEC_MIN_TKV), ONE fa_decode_rows launch executes the exact
9344 // per-row program for all T rows (grid.z = row, per-row n_splits from the same
9345 // fa_split_keys formula) — replacing T x (2 launches + 2 dtod copies + 5 partial allocs)
9346 // and multiplying resident CTAs by T on a latency-bound kernel. Bit-identical per row by
9347 // construction; kernel-check pins rows-vs-loop byte identity, run-spec is the end gate.
9348 // Short ctx (any row below the vec crossover) and MEMRA_NO_FA_VEC/MEMRA_FA_ROWS_OFF keep the
9349 // per-row loop (whose fa_decode picks scalar/vec per row exactly like eager decode).
9350 let mut attn = vbuf(e, t * n_head * head_dim)?; // fully written by every FA arm below
9351 let base_len = kvl.len - t; // KV len BEFORE this round's T tokens were appended
9352 // T=1 INCLUDED (2026-07-05): p-min cuts the draft to 1 in ~75% of rounds on hard
9353 // (agentic) content — the old t>1 gate sent those rounds to the per-row loop (262us/row
9354 // + q-row copy + per-row allocs vs 93us/row through the fused kernel at grid.z=1, same
9355 // program). nsys accounting: 1088 of 1456 verify FA launches were T=1 escapees.
9356 // LEAN T=1 ARM (MEMRA_SPEC_LEAN, close35): at t==1, q IS one row and fa_decode on it is
9357 // the EXACT eager decode dispatch (vec_q_v2 + combine_f32; the rows pair measured +50us
9358 // at m=1). Byte-identical: kernel-check pins rows-vs-loop identity, and the per-row loop
9359 // at t=1 is fa_decode on the same q with zero-offset copies. Gates arbitrate.
9360 if let Some(ctr) = stream_ctr {
9361 // STREAM ARM: causal base from the device counter; host kvl.len is a stale lower
9362 // bound used only for the split-sizing upper bound (+64 slack covers M pre-issued
9363 // rounds at K<=8). Views span the bound; per-row limits derive in-kernel.
9364 let upper = kvl.len + t + 64;
9365 let k_view = e.view_u8(&kvl.k, (upper.min(cache.max_ctx)) * ktb);
9366 let v_view = e.view_u8(&kvl.v, (upper.min(cache.max_ctx)) * vtb);
9367 e.fa_decode_rows_dc(
9368 &q,
9369 &k_view,
9370 &v_view,
9371 &mut attn,
9372 head_dim,
9373 n_head,
9374 n_head_kv,
9375 ctr,
9376 upper.min(cache.max_ctx),
9377 t,
9378 scale,
9379 ktb,
9380 vtb,
9381 0,
9382 false,
9383 )?;
9384 } else if spec_lean() && t == 1 {
9385 let t_kv = base_len + 1;
9386 let k_view = e.view_u8(&kvl.k, t_kv * ktb);
9387 let v_view = e.view_u8(&kvl.v, t_kv * vtb);
9388 e.fa_decode_kvmod(
9389 &q,
9390 &k_view,
9391 &v_view,
9392 &mut attn,
9393 head_dim,
9394 n_head,
9395 n_head_kv,
9396 t_kv,
9397 scale,
9398 ktb,
9399 vtb,
9400 crate::Engine::kv_fp8_on(),
9401 )?;
9402 } else if e.fa_rows_eligible(base_len, head_dim) {
9403 let k_view = e.view_u8(&kvl.k, (base_len + t) * ktb);
9404 let v_view = e.view_u8(&kvl.v, (base_len + t) * vtb);
9405 e.fa_decode_rows(
9406 &q,
9407 &k_view,
9408 &v_view,
9409 &mut attn,
9410 head_dim,
9411 n_head,
9412 n_head_kv,
9413 base_len,
9414 t,
9415 scale,
9416 ktb,
9417 vtb,
9418 None,
9419 false,
9420 crate::Engine::kv_fp8_on(),
9421 None,
9422 )?;
9423 } else {
9424 for r in 0..t {
9425 let t_kv_r = base_len + r + 1; // this row sees keys [0..t_kv_r)
9426 let k_view_r = e.view_u8(&kvl.k, t_kv_r * ktb);
9427 let v_view_r = e.view_u8(&kvl.v, t_kv_r * vtb);
9428 // copy q row into an owned buffer (fa_decode takes &CudaSlice, not CudaView)
9429 let mut q_row = vbuf(e, n_head * head_dim)?; // fully written by copy_view_into
9430 let q_src = q.slice(r * n_head * head_dim..(r + 1) * n_head * head_dim);
9431 e.copy_view_into(&mut q_row, 0, &q_src, n_head * head_dim)?;
9432 let mut attn_row = vbuf(e, n_head * head_dim)?; // fully written by fa_decode
9433 e.fa_decode_kvmod(
9434 &q_row,
9435 &k_view_r,
9436 &v_view_r,
9437 &mut attn_row,
9438 head_dim,
9439 n_head,
9440 n_head_kv,
9441 t_kv_r,
9442 scale,
9443 ktb,
9444 vtb,
9445 crate::Engine::kv_fp8_on(),
9446 )?;
9447 e.copy_into(
9448 &mut attn,
9449 r * n_head * head_dim,
9450 &attn_row,
9451 n_head * head_dim,
9452 )?;
9453 }
9454 }
9455
9456 let attn_g = match &gate {
9457 Some(gate) => {
9458 let mut gsig = vbuf(e, t * n_head * head_dim)?; // fully written by sigmoid
9459 e.sigmoid(gate, &mut gsig, t * n_head * head_dim)?;
9460 let mut ag = vbuf(e, t * n_head * head_dim)?; // fully written by mul
9461 e.mul(&attn, &gsig, &mut ag, t * n_head * head_dim)?;
9462 ag
9463 }
9464 None => attn,
9465 };
9466 // DECODE-EXACT wo projection: at m>=5 (K=4+ with pending) the generic matmul would use dp4a
9467 // (128-thread, different FP sum order than MMVQ). Force MMVQ for bit-identity with decode.
9468 Ok(e.matmul_decode_exact(&fa.wo, &attn_g, t)?)
9469 }
9470
9471 /// Context-linear bytes for a plain serving session's trunk cache.
9472 pub fn plain_session_kv_bytes_per_token(&self) -> usize {
9473 crate::cache::cache_bytes_per_token_for_plan(
9474 &self.cfg,
9475 &self.plan,
9476 0,
9477 self.plan.layers.len(),
9478 )
9479 }
9480
9481 /// `(logical bytes/token, ring-capped bytes/token, ring row cap)` for exact admission.
9482 pub fn plain_session_kv_shape(&self) -> (usize, usize, usize) {
9483 (
9484 self.plain_session_kv_bytes_per_token(),
9485 crate::cache::cache_ring_bytes_per_token_for_plan(
9486 &self.cfg,
9487 &self.plan,
9488 0,
9489 self.plan.layers.len(),
9490 ),
9491 crate::cache::cache_ring_row_cap_for_plan(&self.plan),
9492 )
9493 }
9494
9495 /// Context-linear bytes for a speculative serving session: trunk cache plus persistent MTP
9496 /// scratch. With no MTP head this equals the plain coefficient.
9497 pub fn spec_session_kv_bytes_per_token(&self) -> usize {
9498 let scratch = self
9499 .mtp
9500 .iter()
9501 .chain(self.mtp_extra.iter())
9502 .map(|mtp| {
9503 let (_, _, k, v) = mtp_scratch_layout(&self.cfg, mtp.geom.as_ref());
9504 k + v
9505 })
9506 .sum::<usize>();
9507 self.plain_session_kv_bytes_per_token()
9508 .saturating_add(scratch)
9509 }
9510
9511 /// Spec twin of [`HybridModel::plain_session_kv_shape`]; Step35's persistent MTP scratch is
9512 /// capped by the same SWA ring rows as the trunk.
9513 pub fn spec_session_kv_shape(&self) -> (usize, usize, usize) {
9514 let total = self.spec_session_kv_bytes_per_token();
9515 let (_, mut ring, rows) = self.plain_session_kv_shape();
9516 if rows > 0 {
9517 ring = ring.saturating_add(
9518 self.mtp
9519 .iter()
9520 .chain(self.mtp_extra.iter())
9521 .map(|mtp| {
9522 let (_, _, k, v) = mtp_scratch_layout(&self.cfg, mtp.geom.as_ref());
9523 k + v
9524 })
9525 .sum::<usize>(),
9526 );
9527 }
9528 (total, ring, rows)
9529 }
9530
9531 /// Greedy MTP speculative decode (§B). Token-identical to `generate(prompt, max_new)` but uses
9532 /// the NextN head to draft K tokens then verifies them in one batched target forward.
9533 /// Returns (generated tokens, total_drafted, total_accepted) so the caller can report
9534 /// acceptance rate. `k` = draft length per round.
9535 ///
9536 /// GRAPH DRAFT (stage 2 of graph-grade spec): when the model is all-Dense and the MTP head is
9537 /// Dense (no MoE host readbacks), the fixed-shape T=1 MTP forward is CUDA-graph-captured ONCE
9538 /// and replayed per draft step — the ~40 eager launches per drafted token collapse into one
9539 /// graph dispatch; only the 4-byte token id (and 4-byte p-min confidence) round-trip per step.
9540 /// Event tracking is disabled for the whole call (generate_graph pattern) so every buffer the
9541 /// captured graph references is event-free; the spec loop is strictly single-stream.
9542 /// MEMRA_SPEC_NOGRAPH=1 forces the eager draft chain.
9543 /// SAMPLED mode (MEMRA_SPEC_TEMP>0) has its OWN capture (gumbel-perturbed in-graph argmax,
9544 /// device Philox event counter, persistent q retention) — graph-vs-eager sampled streams are
9545 /// bit-identical for the same (seed, prompt, K, temp); see the sampled-graph setup in
9546 /// generate_spec_inner2.
9547 /// Multi-turn session: trunk cache + MTP draft scratch persist across generate calls, so
9548 /// turn N+1 primes ONLY its new suffix (the 124k-conversation daily pattern — re-priming a
9549 /// 32k history costs ~54s; a suffix prime costs seconds). APPEND-ONLY by construction: the
9550 /// hybrid linear-attn states are in-place (no position index), so a session can extend but
9551 /// never rewind — `committed` is the exact token list whose state the caches hold (includes
9552 /// any overshoot tokens past max_new; the caller renders from `committed`, not its own echo).
9553 pub fn new_session(
9554 &self,
9555 e: &Engine,
9556 max_ctx: usize,
9557 ) -> Result<SpecSession, Box<dyn std::error::Error>> {
9558 Ok(SpecSession {
9559 // STAGE-OWNED KV (lane/pp2-spec 2026-08-06): `pp::new_cache`, not `Cache::new`. This
9560 // is the SERVING spec-session path, and with the ppN door open across two cards a
9561 // primary-homed cache makes every remote stage peer-read its OWN KV on every verify
9562 // round — the wrong-card class already fixed on the two batched serving paths
9563 // (worker.rs 2483 / 2837). With the door shut `new_cache` IS `Cache::new` (same
9564 // branch, same allocations), so single-device behavior is byte-unchanged.
9565 cache: crate::pp::new_cache_planned(e, &self.cfg, &self.plan, max_ctx)?,
9566 scratch: self.new_mtp_scratch(e, max_ctx)?,
9567 committed: Vec::new(),
9568 last_h: None,
9569 next_pred: None,
9570 sctr: 0,
9571 uctr: 0,
9572 draft_ctx: None,
9573 pending_tok: None,
9574 turn_ckpt: None,
9575 telem: SpecTelemetryCounters::default(),
9576 capture_at: None,
9577 boundary_captures: Vec::new(),
9578 ckpt_at: None,
9579 })
9580 }
9581
9582 /// SPEC-ON-CACHE-HIT restore (lane/spec-on-cache-hit, 2026-08-18 — PORT-PLAN item 3,
9583 /// research/cache-spec-design-20260814, scoped to WHOLE-ENTRY restores only): build a
9584 /// SpecSession around a trunk cache the worker already restored from a prefix-cache
9585 /// entry, re-installing the entry's published draft plane as the MTP scratch rows
9586 /// `[0..prefix.len())` and the entry's boundary hidden as `last_h`, then feeding the
9587 /// prompt SUFFIX here — through EXACTLY the plain path's program selection — so the
9588 /// worker always receives a fully-warm continuation session (committed = whole
9589 /// prompt, `next_pred` + `last_h` set; caller sets `next_pred` from the entry's
9590 /// boundary logits on the empty-suffix shape).
9591 ///
9592 /// PROGRAM LAW (the splitiso two-programs class, learned AGAIN in this lane's own
9593 /// gate): the identity target for a converted hit is the PLAIN hit serving the same
9594 /// request, and plain feeds a carried suffix via eager `decode_step` below
9595 /// PRIME_MIN_T and via `prime_cache` at/above it (prefill_tick's arms). The generate
9596 /// path's tokenwise arm routes qwen35-class through the BATCHED T=1 program
9597 /// (`spec_target_step_h`) instead — ULP-different suffix rows, and the gate measured
9598 /// the near-tie flip at generated token ~8 (research/spec-cache-20260818, qwen r3).
9599 /// So the suffix is fed HERE, mirroring prefill_tick arm-for-arm, not handed to the
9600 /// burst prime.
9601 ///
9602 /// SEED RULE (both sampling regimes; lane/sampled-hit-spec 2026-08-19, sampled draw
9603 /// added by lane/sampled-spec-quality 2026-08-19). The boundary token is produced by
9604 /// EXACTLY the rule the cold burst entry applies to its own first token from the same
9605 /// logits row: `argmax` when greedy, and a `sample_boundary_token` draw at Philox
9606 /// counter 0 when sampled. Both shapes are covered — the entry's boundary logits on a
9607 /// full-cover (empty-suffix) hit, this feed's own boundary logits on a suffix hit.
9608 /// That is what keeps a restored session seed-identical to a cold one PER SEED: the
9609 /// cold session draws from the identical row at counter 0 and then runs its rounds from
9610 /// counter 1, so the restored session admits with `sctr = 1` after its own draw.
9611 /// The WORKER owns the one refusal this constructor cannot see — a constrained request.
9612 /// (The penalized-sampled refusal was LIFTED once the burst's penalty window learned to
9613 /// span the session: `committed` here is the WHOLE prompt, so the restored session's
9614 /// window is the cold session's window. It comes back if `MEMRA_SPEC_PEN_SESSION=0`.)
9615 ///
9616 /// NOT the rolled-back partial-restore hazard: the caller restores at exactly the
9617 /// entry's captured endpoint (`e.pos`) through the shipping whole-entry path;
9618 /// mid-entry (`at < e.pos`) trunk restores stay behind MEMRA_PREFIX_PARTIAL_RESTORE
9619 /// and are never routed here.
9620 ///
9621 /// Failure contract: `Err((Some(cache), why))` before any trunk mutation — the
9622 /// worker rebuilds the plain carrier and the hit serves plain, byte-unchanged.
9623 /// `Err((None, why))` after the suffix feed began — the carrier is part-fed and
9624 /// UNUSABLE; the worker serves the request cold-plain (correct, slower) and the
9625 /// entry stays published for the next request.
9626 #[allow(clippy::too_many_arguments)]
9627 pub fn spec_session_from_restored(
9628 &self,
9629 e: &Engine,
9630 mut cache: Cache,
9631 prefix: Vec<u32>,
9632 suffix: &[u32],
9633 draft_k: &CudaSlice<u8>,
9634 draft_v: &CudaSlice<u8>,
9635 draft_k_tok_bytes: usize,
9636 draft_v_tok_bytes: usize,
9637 draft_len: usize,
9638 last_h: &[f32],
9639 // The ENTRY's boundary logits row (the full-cover shape's seed source). May be empty
9640 // when a suffix follows — the feed's own logits are the boundary then.
9641 boundary_logits: &[f32],
9642 // The request's sampler, or None for greedy. Owned here so the seed rule lives in
9643 // ONE place instead of being half-applied by the worker.
9644 sampling: Option<SpecSampling>,
9645 require_anchor: bool,
9646 max_ctx: usize,
9647 // STABLE-BOUNDARY REPUBLICATION (lane/frspec-multiturn-cache, 2026-08-21): ABSOLUTE
9648 // prompt position to split the suffix feed at and capture the extended-entry
9649 // publication + this session's `turn_ckpt` — the worker's stable pre-generation
9650 // boundary (`plain_checkpoint_boundary`). None = legacy prompt-end republication.
9651 // WHY: the prompt-end capture below includes the template's live generation header
9652 // (`<|im_start|>assistant\n<think>\n`), which the next turn's re-render replaces, so
9653 // for a hybrid (whole-entry restores only) every extended entry's last ~2 tokens
9654 // diverged from every future prompt and the hit boundary FROZE at the first
9655 // lcp-split entry forever (measured: cached 6811 of 38228 by turn 8, B4).
9656 republish_at: Option<usize>,
9657 ) -> Result<SpecSession, (Option<Cache>, String)> {
9658 let pos = prefix.len();
9659 let fail = |cache: Cache, msg: String| -> Result<SpecSession, (Option<Cache>, String)> {
9660 Err((Some(cache), msg))
9661 };
9662 if self.mtp.is_none() {
9663 return fail(cache, "no MTP head attached (nothing to draft with)".into());
9664 }
9665 if pos == 0 {
9666 return fail(cache, "empty committed prefix".into());
9667 }
9668 if cache.pos != pos {
9669 let msg = format!(
9670 "restored cache pos {} != restored prefix len {pos}",
9671 cache.pos
9672 );
9673 return fail(cache, msg);
9674 }
9675 if draft_len != pos {
9676 return fail(
9677 cache,
9678 format!("draft plane len {draft_len} != restored prefix len {pos}"),
9679 );
9680 }
9681 if pos + suffix.len() >= max_ctx {
9682 return fail(
9683 cache,
9684 format!(
9685 "prompt {} + suffix would not leave generation room in ctx {max_ctx}",
9686 pos + suffix.len(),
9687 ),
9688 );
9689 }
9690 let mut scratch = match MtpScratch::new(
9691 e,
9692 &self.cfg,
9693 &self.plan,
9694 max_ctx,
9695 self.mtp.as_ref().and_then(|m| m.geom.as_ref()),
9696 ) {
9697 Ok(s) => s,
9698 Err(err) => return fail(cache, format!("draft scratch alloc failed: {err}")),
9699 };
9700 if scratch.kv.ring.is_some() {
9701 return fail(
9702 cache,
9703 "ring-backed draft scratch (Step35 SWA) cannot take a flat prefix restore".into(),
9704 );
9705 }
9706 if scratch.kv.k_tok_bytes != draft_k_tok_bytes
9707 || scratch.kv.v_tok_bytes != draft_v_tok_bytes
9708 {
9709 return fail(
9710 cache,
9711 format!(
9712 "draft plane layout {draft_k_tok_bytes}/{draft_v_tok_bytes} != scratch \
9713 {}/{} bytes/token (stale entry across a format change)",
9714 scratch.kv.k_tok_bytes, scratch.kv.v_tok_bytes,
9715 ),
9716 );
9717 }
9718 if pos > scratch.cap {
9719 return fail(
9720 cache,
9721 format!(
9722 "draft plane rows {pos} exceed scratch capacity {}",
9723 scratch.cap
9724 ),
9725 );
9726 }
9727 let kb = pos * draft_k_tok_bytes;
9728 let vb = pos * draft_v_tok_bytes;
9729 if draft_k.len() < kb || draft_v.len() < vb {
9730 return fail(
9731 cache,
9732 format!(
9733 "truncated draft plane: K {} < {kb} or V {} < {vb} bytes",
9734 draft_k.len(),
9735 draft_v.len(),
9736 ),
9737 );
9738 }
9739 if kb > 0 {
9740 if let Err(err) = e.copy_u8_into(&mut scratch.kv.k, 0, draft_k, kb) {
9741 return fail(cache, format!("draft K restore copy failed: {err}"));
9742 }
9743 }
9744 if vb > 0 {
9745 if let Err(err) = e.copy_u8_into(&mut scratch.kv.v, 0, draft_v, vb) {
9746 return fail(cache, format!("draft V restore copy failed: {err}"));
9747 }
9748 }
9749 if let Err(err) = scratch.set_len(e, pos) {
9750 return fail(cache, format!("draft scratch len set failed: {err}"));
9751 }
9752 let mut last_h_dev = if last_h.len() == self.cfg.n_embd as usize {
9753 // anchor upload failure is acceptance-only when a suffix feed follows (fill
9754 // row-0 falls back to zeros) but FATAL for an empty-suffix continuation (the
9755 // burst entry asserts committed + last_h + next_pred) — the caller says which.
9756 e.htod(last_h).ok()
9757 } else {
9758 None
9759 };
9760 if require_anchor && last_h_dev.is_none() {
9761 return fail(
9762 cache,
9763 "empty-suffix continuation requires the entry's boundary hidden anchor".into(),
9764 );
9765 }
9766 let mut committed = prefix;
9767 // Set on BOTH shapes below (suffix-fed and full-cover) — never left None, which is
9768 // what the empty-suffix continuation assert in the burst entry requires.
9769 let next_pred;
9770 // Philox: (0,0) at admit exactly like a fresh session; a sampled boundary draw below
9771 // consumes counter 0 and leaves 1, which is the state a cold session reaches after
9772 // drawing its own first token from the same row.
9773 let mut sctr = 0u32;
9774 let sampled = sampling.is_some_and(|s| s.temp > 0.0) && spec_sampled_boundary_on();
9775 // Penalty window for the boundary draw: the last `penalty_last_n` tokens of the WHOLE
9776 // prompt, which is what the cold session's own burst sees (Item 2's window). Built
9777 // after the suffix joins `committed` below.
9778 let mut boundary_captures: Vec<SpecBoundaryCapture> = Vec::new();
9779 let mut restored_turn_ckpt: Option<SpecCheckpoint> = None;
9780 if !suffix.is_empty() {
9781 // ---- SUFFIX FEED, mirroring prefill_tick's program selection exactly ----
9782 // From here on the trunk cache mutates: failures return Err((None, _)) and
9783 // the worker serves the request cold-plain instead of reusing the carrier.
9784 let dirty =
9785 |msg: String| -> Result<SpecSession, (Option<Cache>, String)> { Err((None, msg)) };
9786 let n_embd = self.cfg.n_embd as usize;
9787 let t = suffix.len();
9788 let mut h_rows = match e.uninit(t * n_embd) {
9789 Ok(b) => b,
9790 Err(err) => return fail(cache, format!("suffix hidden buffer alloc: {err}")),
9791 };
9792 // STABLE-BOUNDARY split (see `republish_at`): feed stops at the boundary so the
9793 // in-place GDN conv/ssm state can be snapshotted there — the only moment it
9794 // exists (the cold prime-split law). suffix-relative; None = one-segment legacy.
9795 let b_rel = republish_at
9796 .and_then(|abs| abs.checked_sub(pos))
9797 .filter(|&r| r > 0 && r < t);
9798 let mut feed_logits = Vec::new();
9799 let tokenwise_env = std::env::var("MEMRA_PRIME_TOKENWISE").is_ok()
9800 || e.frozen_cpu_experts_prefer_tokenwise_prime();
9801 let mut fed = 0usize;
9802 for seg_end in [b_rel, Some(t)].into_iter().flatten() {
9803 if seg_end <= fed {
9804 continue;
9805 }
9806 let seg = &suffix[fed..seg_end];
9807 let batched = seg.len() >= crate::hybrid_forward::PRIME_MIN_T && !tokenwise_env;
9808 if batched {
9809 // prefill_tick's prime arm: request-level prime_cache call; tokens still
9810 // queued after this segment ride `queued_after` so Step35 arm selection
9811 // stays keyed to the request's end (tick-seg law).
9812 match self.prime_cache(e, seg, &mut cache, t - seg_end) {
9813 Ok((l, _h_seed, hiddens)) => {
9814 if let Err(err) =
9815 e.copy_into(&mut h_rows, fed * n_embd, &hiddens, seg.len() * n_embd)
9816 {
9817 return dirty(format!("suffix hidden copy: {err}"));
9818 }
9819 feed_logits = l;
9820 }
9821 Err(err) => return dirty(format!("suffix prime failed: {err}")),
9822 }
9823 } else {
9824 // prefill_tick's tokenwise arm: eager decode_step, one token at a time.
9825 for (i, &tok) in seg.iter().enumerate() {
9826 match self.decode_step_h(e, tok, &mut cache) {
9827 Ok((l, h)) => {
9828 if let Err(err) =
9829 e.copy_into(&mut h_rows, (fed + i) * n_embd, &h, n_embd)
9830 {
9831 return dirty(format!("suffix hidden copy: {err}"));
9832 }
9833 feed_logits = l;
9834 }
9835 Err(err) => return dirty(format!("suffix decode_step failed: {err}")),
9836 }
9837 }
9838 }
9839 fed = seg_end;
9840 if Some(seg_end) == b_rel {
9841 // The stable pre-generation boundary: capture the extended-entry
9842 // publication AND this session's own turn checkpoint here instead of at
9843 // prompt-end (both would otherwise carry the volatile live-header tail
9844 // the next re-render replaces). Failure silent, turn_ckpt convention.
9845 debug_assert_eq!(
9846 cache.pos,
9847 pos + seg_end,
9848 "stable-boundary capture off the feed split"
9849 );
9850 if spec_restore_republish_on() {
9851 if let Ok(snap) = cache.snapshot(e) {
9852 boundary_captures.push(SpecBoundaryCapture {
9853 snap,
9854 pos: pos + seg_end,
9855 logits: feed_logits.clone(),
9856 last_h: capture_boundary_hidden(e, &h_rows, seg_end, n_embd),
9857 });
9858 }
9859 }
9860 let anchor: Result<CudaSlice<f32>, Box<dyn std::error::Error>> =
9861 e.uninit(n_embd).and_then(|mut a| {
9862 e.copy_view_into(
9863 &mut a,
9864 0,
9865 &h_rows.slice((seg_end - 1) * n_embd..seg_end * n_embd),
9866 n_embd,
9867 )?;
9868 Ok(a)
9869 });
9870 if let (Ok(snap), Ok(last_h)) = (cache.snapshot(e), anchor) {
9871 restored_turn_ckpt = Some(SpecCheckpoint {
9872 snap,
9873 pos: pos + seg_end,
9874 last_h,
9875 });
9876 }
9877 }
9878 }
9879 // Draft-scratch fill for the suffix rows, predecessor-paired: row `pos` reads
9880 // the entry's boundary anchor (zeros fallback — acceptance-only), row `pos+i`
9881 // reads h_rows[i-1]. Chunked like the generate path's fill (transients scale
9882 // with T). Fill failures are acceptance-only — truncate to the restored rows
9883 // and continue; the burst's own set_len keeps the invariant.
9884 let mtp = self.mtp.as_ref().expect("mtp checked above");
9885 let (embd_qt, embd_rb) = self.embd.qt_and_row_bytes(n_embd);
9886 let embd_gpu = if spec_host_embd() {
9887 None
9888 } else {
9889 Some(
9890 self.embd_gpu
9891 .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload")),
9892 )
9893 };
9894 let embd_dev = embd_gpu.map(|g| (g, embd_qt, embd_rb));
9895 let fill_chunk = 4096usize;
9896 let mut filled = true;
9897 let mut start = 0usize;
9898 'fill: while start < t {
9899 let end = (start + fill_chunk).min(t);
9900 let tc = end - start;
9901 let Ok(mut phs) = e.zeros(tc * n_embd) else {
9902 filled = false;
9903 break 'fill;
9904 };
9905 let (src_lo, dst_off, n_copy) = if start == 0 {
9906 (0, n_embd, (tc - 1) * n_embd)
9907 } else {
9908 ((start - 1) * n_embd, 0, tc * n_embd)
9909 };
9910 if start == 0 {
9911 if let Some(lh) = last_h_dev.as_ref() {
9912 if e.copy_into(&mut phs, 0, lh, n_embd).is_err() {
9913 filled = false;
9914 break 'fill;
9915 }
9916 }
9917 }
9918 if n_copy > 0
9919 && e.copy_view_into(
9920 &mut phs,
9921 dst_off,
9922 &h_rows.slice(src_lo..src_lo + n_copy),
9923 n_copy,
9924 )
9925 .is_err()
9926 {
9927 filled = false;
9928 break 'fill;
9929 }
9930 if self
9931 .mtp_kv_fill_all(
9932 e,
9933 &suffix[start..end],
9934 &phs,
9935 pos + start,
9936 &mut scratch,
9937 embd_dev,
9938 )
9939 .is_err()
9940 {
9941 filled = false;
9942 break 'fill;
9943 }
9944 start = end;
9945 }
9946 if !filled {
9947 // acceptance-only: drafts over missing suffix rows are cheap and wrong,
9948 // so keep only the restored rows resident and let verify arbitrate.
9949 if let Err(err) = scratch.set_len(e, pos) {
9950 return dirty(format!("scratch truncation after failed fill: {err}"));
9951 }
9952 }
9953 // EXTENDED-ENTRY PUBLICATION (lane/sampled-spec-quality, Item 3 — the fix for
9954 // "a restored spec session never publishes an extended entry", SAMPLED-HIT.md
9955 // finding (d)). Pre-lane, publication was armed only for COLD sessions
9956 // (`spec_resumed == 0` in the worker) and both engine capture sites require a
9957 // non-continuation burst — but a converted hit's first burst IS a continuation,
9958 // so a growing conversation learned exactly ONE boundary and turn 3 could never
9959 // hit a longer prefix than turn 2 did.
9960 //
9961 // WHERE, and why it is safe here: `cache.pos == prefix + suffix` at this exact
9962 // line — the trunk is primed over the whole prompt, nothing is generated, and the
9963 // draft plane rows [0..prompt) are filled just above. That is a complete
9964 // whole-entry boundary (`pos == fed_len`), the same shape the cold seed capture
9965 // publishes; the worker's existing publication sweep picks it up because it is
9966 // keyed on non-empty `boundary_captures` and is sampler- and resume-independent.
9967 // NOT the partial-restore hazard: the boundary is this session's own prompt END,
9968 // never mid-entry, so `entry_pos != fed_len` still refuses on the way back in.
9969 // Failure is SILENT by design (the turn_ckpt / boundary-capture convention):
9970 // publication is an optimization, never a correctness dependency.
9971 //
9972 // SUPERSEDED WHEN `republish_at` FIRED (lane/frspec-multiturn-cache): a prompt-end
9973 // entry's tail is the live generation header the next re-render replaces, so on a
9974 // hybrid (whole-entry restores) it can never serve the conversation's next turn —
9975 // the stable-boundary capture above IS this publication, minus the poisoned tail.
9976 if spec_restore_republish_on() && boundary_captures.is_empty() {
9977 debug_assert_eq!(
9978 cache.pos,
9979 pos + t,
9980 "extended-entry capture must sit at the restored session's prompt end",
9981 );
9982 if let Ok(snap) = cache.snapshot(e) {
9983 boundary_captures.push(SpecBoundaryCapture {
9984 snap,
9985 pos: pos + t,
9986 logits: feed_logits.clone(),
9987 last_h: capture_boundary_hidden(e, &h_rows, t, n_embd),
9988 });
9989 }
9990 }
9991 // continuation seed: the feed's boundary logits ARE the plain path's boundary
9992 // logits (same program), so greedy's argmax here is plain's first emitted token,
9993 // and the sampled draw is the cold sampled session's own first token.
9994 next_pred = Some(if sampled {
9995 let sp = sampling.expect("sampled implies a sampler");
9996 // `committed` is still the restored prefix here; the suffix joins it below —
9997 // so this is the last-N window over the WHOLE prompt, exactly the cold
9998 // session's own window at its first token.
9999 let hist = pen_window_seed(&committed, suffix, sp.penalty_last_n);
10000 match sample_boundary_token(
10001 e,
10002 &feed_logits,
10003 &sp,
10004 &hist,
10005 &mut sctr,
10006 "restore-suffix-feed",
10007 ) {
10008 Ok(t) => t,
10009 // the trunk is already fed: hand nothing back, the worker serves the
10010 // request cold-plain. Never fall back to an argmax — that would put a
10011 // greedy token in a sampled stream to save a slow path.
10012 Err(err) => {
10013 return dirty(format!("boundary token draw failed: {err}"));
10014 }
10015 }
10016 } else {
10017 argmax(&feed_logits) as u32
10018 });
10019 let mut lh = match e.uninit(n_embd) {
10020 Ok(b) => b,
10021 Err(err) => return dirty(format!("boundary hidden alloc: {err}")),
10022 };
10023 if let Err(err) = e.copy_view_into(
10024 &mut lh,
10025 0,
10026 &h_rows.slice((t - 1) * n_embd..t * n_embd),
10027 n_embd,
10028 ) {
10029 return dirty(format!("boundary hidden copy: {err}"));
10030 }
10031 last_h_dev = Some(lh);
10032 committed.extend_from_slice(suffix);
10033 } else {
10034 // FULL-COVER shape (empty suffix — the identical-repeat / agent-loop shape): the
10035 // ENTRY's boundary logits are the boundary row, and this is the token the cold
10036 // session emits from that same row. Owned here rather than in the worker so the
10037 // sampled draw cannot be half-applied on one shape (the worker used to argmax it).
10038 if boundary_logits.is_empty() {
10039 return fail(
10040 cache,
10041 "full-cover restore without the entry's boundary logits".into(),
10042 );
10043 }
10044 next_pred = Some(if sampled {
10045 let sp = sampling.expect("sampled implies a sampler");
10046 let hist = pen_window_seed(&committed, &[], sp.penalty_last_n);
10047 match sample_boundary_token(
10048 e,
10049 boundary_logits,
10050 &sp,
10051 &hist,
10052 &mut sctr,
10053 "restore-full-cover",
10054 ) {
10055 Ok(t) => t,
10056 // nothing has been mutated on this shape — hand the carrier back and let
10057 // the hit serve PLAIN (the banked pre-lane path).
10058 Err(err) => {
10059 return fail(cache, format!("boundary token draw failed: {err}"));
10060 }
10061 }
10062 } else {
10063 argmax(boundary_logits) as u32
10064 });
10065 }
10066 Ok(SpecSession {
10067 cache,
10068 scratch,
10069 committed,
10070 last_h: last_h_dev,
10071 next_pred,
10072 sctr,
10073 uctr: 0,
10074 draft_ctx: None,
10075 pending_tok: None,
10076 // Stable-boundary capture from the split feed above (None on the legacy shape):
10077 // a restored session previously parked WITHOUT a checkpoint, so the next turn's
10078 // affinity probe declined ("no turn checkpoint retained") and the conversation
10079 // fell back to the frozen prefix entry forever.
10080 turn_ckpt: restored_turn_ckpt,
10081 telem: SpecTelemetryCounters::default(),
10082 capture_at: None,
10083 boundary_captures,
10084 ckpt_at: None,
10085 })
10086 }
10087
10088 /// Forced-gate exact state comparison. This intentionally reads the real live prefixes from
10089 /// their owning PP devices: matching emitted ids alone would miss a stale `len_d`, recurrent
10090 /// snapshot, or draft-KV row that only corrupts the following round.
10091 pub fn optipipe_compare_session_state(
10092 &self,
10093 e: &Engine,
10094 reference: &SpecSession,
10095 candidate: &SpecSession,
10096 ) -> Result<OptiForkStateIdentity, Box<dyn std::error::Error>> {
10097 fn fail(what: &str) -> Box<dyn std::error::Error> {
10098 format!("optipipe state mismatch: {what}").into()
10099 }
10100 fn same_f32(a: &[f32], b: &[f32]) -> bool {
10101 a.len() == b.len() && a.iter().zip(b).all(|(x, y)| x.to_bits() == y.to_bits())
10102 }
10103 fn compare_layers(
10104 es: &Engine,
10105 range: std::ops::Range<usize>,
10106 reference: &SpecSession,
10107 candidate: &SpecSession,
10108 report: &mut OptiForkStateIdentity,
10109 ) -> Result<(), Box<dyn std::error::Error>> {
10110 for il in range {
10111 match (&reference.cache.kv[il], &candidate.cache.kv[il]) {
10112 (Some(a), Some(b)) => {
10113 if a.len != b.len {
10114 return Err(fail(&format!(
10115 "layer {il} host KV len {} != {}",
10116 a.len, b.len
10117 )));
10118 }
10119 let ad = es.dtoh_i32(&a.len_d)?;
10120 let bd = es.dtoh_i32(&b.len_d)?;
10121 if ad != bd || ad.first().copied() != Some(a.len as i32) {
10122 return Err(fail(&format!(
10123 "layer {il} device KV len {ad:?} != {bd:?} (host={})",
10124 a.len,
10125 )));
10126 }
10127 let kb = a.len * a.k_tok_bytes;
10128 let vb = a.len * a.v_tok_bytes;
10129 if kb > 0 {
10130 let ak = es.dtoh_u8_view(&a.k.slice(0..kb))?;
10131 let bk = es.dtoh_u8_view(&b.k.slice(0..kb))?;
10132 if ak != bk {
10133 let at = ak.iter().zip(&bk).position(|(x, y)| x != y).unwrap();
10134 return Err(fail(&format!(
10135 "layer {il} K bytes at byte {at} row {} offset {}: {} != {}",
10136 at / a.k_tok_bytes,
10137 at % a.k_tok_bytes,
10138 ak[at],
10139 bk[at],
10140 )));
10141 }
10142 }
10143 if vb > 0 {
10144 let av = es.dtoh_u8_view(&a.v.slice(0..vb))?;
10145 let bv = es.dtoh_u8_view(&b.v.slice(0..vb))?;
10146 if av != bv {
10147 let at = av.iter().zip(&bv).position(|(x, y)| x != y).unwrap();
10148 return Err(fail(&format!(
10149 "layer {il} V bytes at byte {at} row {} offset {}: {} != {}",
10150 at / a.v_tok_bytes,
10151 at % a.v_tok_bytes,
10152 av[at],
10153 bv[at],
10154 )));
10155 }
10156 }
10157 report.trunk_kv_bytes += kb + vb;
10158 }
10159 (None, None) => {}
10160 _ => return Err(fail(&format!("layer {il} KV presence"))),
10161 }
10162 match (&reference.cache.recur[il], &candidate.cache.recur[il]) {
10163 (Some(a), Some(b)) => {
10164 let ac = es.dtoh(&a.conv_state)?;
10165 let bc = es.dtoh(&b.conv_state)?;
10166 if !same_f32(&ac, &bc) {
10167 return Err(fail(&format!("layer {il} conv state")));
10168 }
10169 let as_ = es.dtoh(&a.ssm_state)?;
10170 let bs = es.dtoh(&b.ssm_state)?;
10171 if !same_f32(&as_, &bs) {
10172 return Err(fail(&format!("layer {il} SSM state")));
10173 }
10174 report.recurrent_bytes += (ac.len() + as_.len()) * 4;
10175 }
10176 (None, None) => {}
10177 _ => return Err(fail(&format!("layer {il} recurrent presence"))),
10178 }
10179 }
10180 Ok(())
10181 }
10182
10183 if reference.committed != candidate.committed {
10184 return Err(fail("committed token ids"));
10185 }
10186 if reference.cache.pos != candidate.cache.pos
10187 || reference.cache.max_ctx != candidate.cache.max_ctx
10188 {
10189 return Err(fail("cache pos/capacity"));
10190 }
10191 if reference.pending_tok != candidate.pending_tok
10192 || reference.next_pred != candidate.next_pred
10193 || reference.sctr != candidate.sctr
10194 || reference.uctr != candidate.uctr
10195 {
10196 return Err(fail("pending/prediction/counter tail"));
10197 }
10198
10199 let mut report = OptiForkStateIdentity::default();
10200 if let Some(fence) = crate::pp::pp_cuts(self.layers.len()) {
10201 let rt = crate::pp::PpNRt::get(e)?;
10202 for stage in 0..rt.n_stages() {
10203 let _scope = rt.enter(stage);
10204 compare_layers(
10205 rt.engine(stage, e),
10206 fence[stage]..fence[stage + 1],
10207 reference,
10208 candidate,
10209 &mut report,
10210 )?;
10211 }
10212 } else {
10213 compare_layers(e, 0..self.layers.len(), reference, candidate, &mut report)?;
10214 }
10215
10216 if reference.scratch.plane_count() != candidate.scratch.plane_count() {
10217 return Err(fail("draft scratch plane count"));
10218 }
10219 for index in 0..reference.scratch.plane_count() {
10220 let (a, _) = reference.scratch.plane(index);
10221 let (b, _) = candidate.scratch.plane(index);
10222 if a.len != b.len
10223 || a.kv_dim_k != b.kv_dim_k
10224 || a.kv_dim_v != b.kv_dim_v
10225 || a.k_tok_bytes != b.k_tok_bytes
10226 || a.v_tok_bytes != b.v_tok_bytes
10227 || e.dtoh_i32(&a.len_d)? != e.dtoh_i32(&b.len_d)?
10228 {
10229 return Err(fail(&format!("draft scratch plane {index} length/layout")));
10230 }
10231 let kb = a.len * a.k_tok_bytes;
10232 let vb = a.len * a.v_tok_bytes;
10233 if kb > 0 && e.dtoh_u8_view(&a.k.slice(0..kb))? != e.dtoh_u8_view(&b.k.slice(0..kb))? {
10234 return Err(fail(&format!("draft scratch plane {index} K bytes")));
10235 }
10236 if vb > 0 && e.dtoh_u8_view(&a.v.slice(0..vb))? != e.dtoh_u8_view(&b.v.slice(0..vb))? {
10237 return Err(fail(&format!("draft scratch plane {index} V bytes")));
10238 }
10239 report.scratch_kv_bytes += kb + vb;
10240 }
10241
10242 match (&reference.last_h, &candidate.last_h) {
10243 (Some(a), Some(b)) => {
10244 let ah = e.dtoh(a)?;
10245 let bh = e.dtoh(b)?;
10246 if !same_f32(&ah, &bh) {
10247 return Err(fail("last hidden/seed bytes"));
10248 }
10249 report.hidden_bytes = ah.len() * 4;
10250 }
10251 (None, None) => {}
10252 _ => return Err(fail("last hidden/seed presence")),
10253 }
10254 Ok(report)
10255 }
10256
10257 /// SESSION-AFFINITY REWIND (lane/session-affinity, 2026-08-05): roll `sess` back to its
10258 /// retained prompt-end checkpoint, so a request whose prompt matches
10259 /// `committed[..rewind_pos()]` exactly can resume there and prime only its own delta.
10260 ///
10261 /// EXACTNESS. After this returns, the session is byte-for-byte the state it was in AT that
10262 /// boundary: full-attn KV truncated to it (append-only, position-addressed), GDN conv/ssm
10263 /// restored from the device copy taken there, draft scratch length reset, `committed`
10264 /// truncated, `last_h` = the boundary's predecessor anchor. That is precisely the state a
10265 /// fresh prime of `committed[..pos]` would have produced, so the following suffix prime and
10266 /// every burst after it are identical to a cold run of the same token stream — the
10267 /// committed-tokens-authoritative contract.
10268 ///
10269 /// `next_pred` and `pending_tok` are CLEARED: both describe generation past the boundary,
10270 /// which the rewind discards. The caller therefore must supply a non-empty suffix (a
10271 /// rewound session cannot serve an empty-suffix continuation burst — there is nothing to
10272 /// continue). The persistent draft graph survives: it bakes only session-stable pointers
10273 /// (the scratch KV, the resident embedding), none of which the rewind moves.
10274 ///
10275 /// The checkpoint is CONSUMED (`turn_ckpt` taken): its snapshot buffers are freed here, and
10276 /// this turn's own prime installs a fresh one at the new prompt end. Returns the position
10277 /// rewound to, or `None` when the session holds no checkpoint (caller: full re-prime).
10278 pub fn spec_rewind_to_checkpoint(
10279 &self,
10280 e: &Engine,
10281 sess: &mut SpecSession,
10282 ) -> Result<Option<usize>, Box<dyn std::error::Error>> {
10283 if sess.turn_ckpt.as_ref().is_some_and(|ckpt| {
10284 !sess.cache.can_rollback(&ckpt.snap, 0) || !sess.scratch.can_rewind_to(ckpt.pos)
10285 }) {
10286 return Err(
10287 "SWA ring rewind checkpoint has been lapped; full re-prime required".into(),
10288 );
10289 }
10290 let Some(ckpt) = sess.turn_ckpt.take() else {
10291 return Ok(None);
10292 };
10293 assert!(
10294 ckpt.pos <= sess.committed.len(),
10295 "checkpoint past committed ({} > {})",
10296 ckpt.pos,
10297 sess.committed.len()
10298 );
10299 // Restore through each layer's owning engine. A single primary-engine rollback is not
10300 // sufficient when the serving cache is stage-owned under cross-device PP.
10301 crate::pp::restore_cache_checkpoint(e, self, None, &mut sess.cache, &ckpt.snap)?;
10302 debug_assert_eq!(
10303 sess.cache.pos, ckpt.pos,
10304 "rollback landed off the checkpoint"
10305 );
10306 sess.scratch.set_len(e, ckpt.pos)?;
10307 sess.committed.truncate(ckpt.pos);
10308 sess.last_h = Some(ckpt.last_h);
10309 sess.next_pred = None;
10310 sess.pending_tok = None;
10311 Ok(Some(ckpt.pos))
10312 }
10313
10314 /// Grow a parked speculative session to `target_cap` and rewind it to its retained turn
10315 /// checkpoint without re-priming the checkpoint prefix.
10316 ///
10317 /// The trunk cache is restored exactly like a plain grown cache: append-only full-attention
10318 /// KV rows come from the parked cache, while recurrent state comes from the checkpoint's
10319 /// owned snapshot. The MTP scratch is also context-linear and its rows below the checkpoint
10320 /// remain authoritative, so they are copied into a fresh larger scratch before its length is
10321 /// truncated. Pointer-baking draft graphs are dropped and recaptured on the next burst.
10322 ///
10323 /// All fallible work completes before `sess` is mutated. A failed allocation or copy leaves
10324 /// the parked session intact, allowing the caller one reclaim-and-retry attempt.
10325 pub fn spec_grow_and_rewind_to_checkpoint(
10326 &self,
10327 e: &Engine,
10328 sess: &mut SpecSession,
10329 target_cap: usize,
10330 ) -> Result<Option<usize>, Box<dyn std::error::Error>> {
10331 if target_cap <= sess.cache.max_ctx {
10332 return self.spec_rewind_to_checkpoint(e, sess);
10333 }
10334 let Some(ckpt) = sess.turn_ckpt.as_ref() else {
10335 return Ok(None);
10336 };
10337 if ckpt.pos == 0 || ckpt.pos > sess.committed.len() {
10338 return Err(format!(
10339 "checkpoint pos {} outside committed length {}",
10340 ckpt.pos,
10341 sess.committed.len(),
10342 )
10343 .into());
10344 }
10345 if ckpt.pos > target_cap {
10346 return Err(format!(
10347 "checkpoint pos {} exceeds grown capacity {target_cap}",
10348 ckpt.pos,
10349 )
10350 .into());
10351 }
10352
10353 let mut grown_cache = crate::pp::new_cache_planned(e, &self.cfg, &self.plan, target_cap)?;
10354 let mut grown_scratch = self.new_mtp_scratch(e, target_cap)?;
10355 crate::pp::restore_cache_checkpoint(
10356 e,
10357 self,
10358 Some(&sess.cache),
10359 &mut grown_cache,
10360 &ckpt.snap,
10361 )?;
10362
10363 if sess.scratch.plane_count() != grown_scratch.plane_count() {
10364 return Err("checkpoint draft plane count mismatch".into());
10365 }
10366 for index in 0..sess.scratch.plane_count() {
10367 let (src, _) = sess.scratch.plane(index);
10368 let (dst, _) = grown_scratch.plane_mut(index);
10369 if ckpt.pos > src.len
10370 || src.kv_dim_k != dst.kv_dim_k
10371 || src.kv_dim_v != dst.kv_dim_v
10372 || src.k_tok_bytes != dst.k_tok_bytes
10373 || src.v_tok_bytes != dst.v_tok_bytes
10374 {
10375 return Err(format!(
10376 "checkpoint draft plane {index} layout mismatch (pos {}, source len {})",
10377 ckpt.pos, src.len,
10378 )
10379 .into());
10380 }
10381 match (&src.ring, dst.ring.as_ref()) {
10382 (Some(sring), Some(_)) => {
10383 // Ring-backed draft plane (step35): `ckpt.pos` is absolute and exceeds the
10384 // physical rows once lapped — same class as the trunk-KV restore panic
10385 // (2026-08-29 warm-turn-at-40k). Copy the aligned live window, rebase.
10386 let (new_base, phys) = sring.restore_plan(ckpt.pos).map_err(|err| {
10387 format!("checkpoint draft plane {index} SWA restore refused: {err}")
10388 })?;
10389 let rows = phys.len();
10390 let kb = rows * src.k_tok_bytes;
10391 let vb = rows * src.v_tok_bytes;
10392 if kb > 0 {
10393 e.copy_u8_range_into(
10394 &mut dst.k,
10395 0,
10396 &src.k,
10397 phys.start * src.k_tok_bytes,
10398 kb,
10399 )?;
10400 }
10401 if vb > 0 {
10402 e.copy_u8_range_into(
10403 &mut dst.v,
10404 0,
10405 &src.v,
10406 phys.start * src.v_tok_bytes,
10407 vb,
10408 )?;
10409 }
10410 dst.ring
10411 .as_mut()
10412 .expect("ring presence checked above")
10413 .apply_rebase(new_base);
10414 if let Some(base_d) = dst.base_d.as_mut() {
10415 e.set_i32_one(base_d, new_base as i32)?;
10416 }
10417 }
10418 (None, None) => {
10419 let kb = ckpt.pos * src.k_tok_bytes;
10420 let vb = ckpt.pos * src.v_tok_bytes;
10421 if kb > 0 {
10422 e.copy_u8_into(&mut dst.k, 0, &src.k, kb)?;
10423 }
10424 if vb > 0 {
10425 e.copy_u8_into(&mut dst.v, 0, &src.v, vb)?;
10426 }
10427 }
10428 _ => {
10429 return Err(format!("checkpoint draft plane {index} ring/flat mismatch").into());
10430 }
10431 }
10432 }
10433 grown_scratch.set_len(e, ckpt.pos)?;
10434 // The old scratch is dropped immediately after publication below. Bound its D2D reads
10435 // first; growth happens once per rewritten turn, outside the decode hot loop.
10436 e.stream().synchronize()?;
10437
10438 let ckpt = sess
10439 .turn_ckpt
10440 .take()
10441 .expect("checkpoint remained present through transactional grow");
10442 let pos = ckpt.pos;
10443 sess.cache = grown_cache;
10444 sess.scratch = grown_scratch;
10445 sess.committed.truncate(pos);
10446 sess.last_h = Some(ckpt.last_h);
10447 sess.next_pred = None;
10448 sess.pending_tok = None;
10449 sess.draft_ctx = None;
10450 debug_assert_eq!(sess.cache.pos, pos, "grown rewind landed off checkpoint");
10451 debug_assert!(
10452 (0..sess.scratch.plane_count()).all(|index| sess.scratch.plane(index).0.len == pos),
10453 "grown draft rewind landed off checkpoint"
10454 );
10455 Ok(Some(pos))
10456 }
10457
10458 /// Commit a carried pending bonus (see SpecSession::pending_tok): one T=1 trunk pass
10459 /// (its logits' argmax becomes next_pred) + the draft-KV fill at the carried anchor —
10460 /// byte-identical to the pre-carry session tail. Required before a non-empty-suffix
10461 /// prime, a sampled turn, or parking a session for pool reuse. No-op without a pending.
10462 /// `sampling` is the sampler of the request that will CONSUME the resulting `next_pred`
10463 /// (lane/sampled-spec-quality): this is a boundary site like any other, so a sampled
10464 /// consumer must get a DRAWN token, not an argmax. Pass `None` from the park/demote
10465 /// callers — a pending only ever exists on the GREEDY tail, and the consumer of a
10466 /// park-time flush is a future request whose sampler is not knowable here (residual
10467 /// named at the pool-resume probe in worker.rs and in SAMPLED-QUALITY.md).
10468 pub fn spec_flush_pending(
10469 &self,
10470 e: &Engine,
10471 sess: &mut SpecSession,
10472 sampling: Option<SpecSampling>,
10473 ) -> Result<(), Box<dyn std::error::Error>> {
10474 let Some(b) = sess.pending_tok.take() else {
10475 return Ok(());
10476 };
10477 if self.mtp.is_none() {
10478 return Err("pending carry requires an MTP head".into());
10479 }
10480 let n_embd = self.cfg.n_embd as usize;
10481 let (embd_qt, embd_rb) = self.embd.qt_and_row_bytes(n_embd);
10482 let embd_gpu = if spec_host_embd() {
10483 None
10484 } else {
10485 Some(
10486 self.embd_gpu
10487 .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload")),
10488 )
10489 };
10490 let embd_dev = embd_gpu.map(|g| (g, embd_qt, embd_rb));
10491 let pos_b = sess.cache.pos;
10492 sess.scratch.set_len(e, pos_b)?;
10493 let (lg_b, hb) = self.spec_target_step_h(e, b, &mut sess.cache)?;
10494 sess.next_pred = Some(match sampling {
10495 Some(sp) if sp.temp > 0.0 && spec_sampled_boundary_on() => {
10496 // window includes `b` itself: it is committed by this pass, and the pre-lane
10497 // code never counted a boundary token in the penalty history at all.
10498 let hist = pen_window_seed(&sess.committed, &[b], sp.penalty_last_n);
10499 sample_boundary_token(e, &lg_b, &sp, &hist, &mut sess.sctr, "flush-pending")?
10500 }
10501 _ => argmax(&lg_b) as u32,
10502 });
10503 let anchor = sess
10504 .last_h
10505 .as_ref()
10506 .expect("pending carry requires last_h (the predecessor-row anchor)");
10507 self.mtp_kv_fill_all(e, &[b], anchor, pos_b, &mut sess.scratch, embd_dev)?;
10508 sess.last_h = Some(hb);
10509 sess.committed.push(b);
10510 Ok(())
10511 }
10512
10513 /// Solo target feed used only at speculative round boundaries. Step35 serving made its
10514 /// staged batched B=1 graph authoritative, so a speculative session must enter and leave
10515 /// rounds through that same graph. Other model families keep their eager T=1 contract.
10516 fn spec_target_step_h(
10517 &self,
10518 e: &Engine,
10519 token: u32,
10520 cache: &mut Cache,
10521 ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
10522 if !self.sliding_gated_moe_batch_program() && !self.batched_serving_numeric_class() {
10523 return self.decode_step_h(e, token, cache);
10524 }
10525 let pos0 = cache.pos;
10526 let (logits, hidden) = self.decode_step_t_core(e, &[token], pos0, cache, None, None)?;
10527 Ok((e.dtoh(&logits)?, hidden))
10528 }
10529
10530 /// The archs whose LIVE B=1 serving runs the generic BATCHED numeric class (decode_step_batch
10531 /// walk + batched head), so their spec verify must run the SAME class. MoE learned this
10532 /// 2026-08-14 AM (4b777ccc5); the dense hybrid reproduced the identical near-tie flip class
10533 /// the same day on Qwen3.8-27B — eager-class verify logits drift from batched-class serving
10534 /// logits ("1 ULP at layer 2 → 2.3e-1 logit maxdiff at the head"), and the GDN recurrence
10535 /// carries the drift until a near-tie flips deep in generation. One predicate so the five
10536 /// dispatch sites cannot drift apart again.
10537 /// Draft-graph head admissibility (lane/draftcost-moe, 2026-08-20): the capture body
10538 /// (`mtp_head_forward_cap`) supports Dense heads AND resident-MoE heads
10539 /// (`Ffn::Moe(m) if m.dev_exps.is_some()`); non-resident MoE still refuses inside the
10540 /// capture and the caller falls back to the eager chain by design. Trunk FFN class is
10541 /// irrelevant — the graph body is the HEAD forward only. One predicate for all three
10542 /// eligibility sites so they cannot drift (the serving numeric-class lesson).
10543 fn mtp_graph_capturable(&self) -> bool {
10544 // EVERY loaded head must be capturable: the multi-head chain graphs capture each
10545 // head's forward (lane/step37-draft-graph-serving-20260830), so one SLRU-locked MoE
10546 // head anywhere in the chain refuses capture for the whole chain (loudly, via the
10547 // capture-site WARN) rather than capturing a subset the launch order cannot honor.
10548 let head_ok = |m: &MtpHead| match &m.ffn {
10549 crate::hybrid::Ffn::Dense { .. } => true,
10550 crate::hybrid::Ffn::Moe(mo) => mo.dev_exps.is_some(),
10551 };
10552 self.mtp.as_ref().map(&head_ok).unwrap_or(false) && self.mtp_extra.iter().all(head_ok)
10553 }
10554
10555 fn batched_serving_numeric_class(&self) -> bool {
10556 self.plan
10557 .trunk_operations()
10558 .contains(&memra_gguf::model_plan::OperationKind::GatedDeltaNet)
10559 }
10560
10561 /// The family the MTP verify-graph default was measured on: GatedDeltaNet state layers
10562 /// (a `recur` mixer) together with a routed-MoE FFN — Ornith-1.5-35B-A3B and its kin. The
10563 /// server-side twin of this test is `model_forces_spec_replay` (GatedDeltaNet + MoeMlp);
10564 /// keeping the engine's own version structural rather than name-based means a new
10565 /// checkpoint of the same shape inherits the default, and a different shape does not.
10566 fn vgraph_family_default(&self) -> bool {
10567 let has_linear = self
10568 .layers
10569 .iter()
10570 .any(|l| matches!(l.mixer, Mixer::Linear(_)));
10571 let has_moe = self
10572 .layers
10573 .iter()
10574 .any(|l| matches!(l.ffn, crate::hybrid::Ffn::Moe(_)));
10575 has_linear && has_moe
10576 }
10577
10578 fn sliding_gated_moe_batch_program(&self) -> bool {
10579 self.uses_sliding_gated_moe_program()
10580 }
10581
10582 fn gemma_batch_program(&self) -> bool {
10583 self.uses_gemma_program()
10584 }
10585
10586 /// Reduced-matrix admission for increment 1. This deliberately does not change the PP-2
10587 /// serving policy: the worker calls it only after `MEMRA_SPEC_PIPE=1` and an explicit spec
10588 /// session already exist.
10589 pub fn spec_pipe_available(&self, e: &Engine) -> bool {
10590 if std::env::var("MEMRA_SPEC_PIPE").as_deref() != Ok("1")
10591 || !spec_devacc()
10592 || spec_replay_env_enabled()
10593 || spec_stream()
10594 || std::env::var("MEMRA_SPEC_ADAPT").as_deref() == Ok("1")
10595 || std::env::var("MEMRA_SPEC_PMIN0").as_deref() == Ok("1")
10596 || std::env::var("MEMRA_SPEC_PP_ANATOMY").as_deref() == Ok("1")
10597 || std::env::var("MEMRA_SPEC_PMIN")
10598 .ok()
10599 .and_then(|v| v.parse::<f32>().ok())
10600 .unwrap_or(0.0)
10601 > 0.0
10602 || self.is_gemma4_e4b()
10603 || self.gemma_batch_program()
10604 || self.mtp.is_none()
10605 || !self.mtp_extra.is_empty()
10606 {
10607 return false;
10608 }
10609 let Some(cuts) = crate::pp::pp_cuts(self.layers.len()) else {
10610 return false;
10611 };
10612 if cuts.len() != 3 || crate::pp::pp2_streams_off() || !crate::pp::spec_pp_on() {
10613 return false;
10614 }
10615 crate::pp::PpNRt::get(e)
10616 .map(|rt| rt.n_stages() == 2 && rt.cross_device())
10617 .unwrap_or(false)
10618 }
10619
10620 /// Two warm greedy continuation bursts over one PP-2 interval coordinator. The two existing
10621 /// `generate_spec_inner2` call stacks own all per-session round locals; only phase issue order
10622 /// changes. No callback is accepted in increment 1 — the worker publishes each completed burst.
10623 #[allow(clippy::too_many_arguments)]
10624 pub fn generate_spec_session_pair(
10625 &self,
10626 e: &Engine,
10627 sess_a: &mut SpecSession,
10628 max_new_a: usize,
10629 k_a: usize,
10630 sess_b: &mut SpecSession,
10631 max_new_b: usize,
10632 k_b: usize,
10633 ) -> Result<((Vec<u32>, usize, usize), (Vec<u32>, usize, usize)), Box<dyn std::error::Error>>
10634 {
10635 if !self.spec_pipe_available(e) {
10636 return Err("two-session speculative pipeline is outside its reduced matrix".into());
10637 }
10638 if max_new_a == 0 || max_new_b == 0 || k_a == 0 || k_b == 0 {
10639 return Err(
10640 "two-session speculative pipeline requires non-empty positive-K bursts".into(),
10641 );
10642 }
10643 for sess in [&*sess_a, &*sess_b] {
10644 if sess.committed.is_empty()
10645 || sess.last_h.is_none()
10646 || (sess.next_pred.is_none() && sess.pending_tok.is_none())
10647 {
10648 return Err("two-session speculative pipeline requires warm continuations".into());
10649 }
10650 }
10651
10652 let graph_ok = std::env::var("MEMRA_SPEC_NOGRAPH").is_err()
10653 && !spec_host_embd()
10654 && self.mtp_graph_capturable()
10655 && self.mtp_extra.is_empty()
10656 && !crate::model::full_prec_enabled();
10657 let graph_a = graph_ok && k_a + 2 < 96;
10658 let graph_b = graph_ok && k_b + 2 < 96;
10659 let was_tracking = e.ctx().is_event_tracking();
10660 if (graph_a || graph_b) && was_tracking {
10661 unsafe {
10662 e.ctx().disable_event_tracking();
10663 }
10664 }
10665
10666 static LOGGED: std::sync::Once = std::sync::Once::new();
10667 LOGGED.call_once(|| {
10668 eprintln!("[spec-pipe] two-session PP-2 continuation pipeline engaged");
10669 });
10670 let sync = std::sync::Arc::new(SpecPipeSync::new());
10671 let lane_a = SpecPipeLane {
10672 sync: sync.clone(),
10673 lane: 0,
10674 };
10675 let lane_b = SpecPipeLane { sync, lane: 1 };
10676 let mut sess_b_ptr = SpecPipeSessionPtr(sess_b as *mut SpecSession);
10677 let (result_a, result_b) = std::thread::scope(|scope| {
10678 let b = scope.spawn(move || {
10679 let mut finish = SpecPipeFinish::new(&lane_b);
10680 let sess_b = unsafe { sess_b_ptr.get_mut() };
10681 let result = e
10682 .ctx()
10683 .bind_to_thread()
10684 .map_err(|err| err.to_string())
10685 .and_then(|_| {
10686 self.generate_spec_inner2(
10687 e,
10688 &[],
10689 max_new_b,
10690 k_b,
10691 graph_b,
10692 Some(sess_b),
10693 None,
10694 None,
10695 None,
10696 None,
10697 Some(&lane_b),
10698 )
10699 .map_err(|err| err.to_string())
10700 });
10701 finish.close(result.is_err());
10702 result
10703 });
10704 let mut finish = SpecPipeFinish::new(&lane_a);
10705 let result_a = self.generate_spec_inner2(
10706 e,
10707 &[],
10708 max_new_a,
10709 k_a,
10710 graph_a,
10711 Some(sess_a),
10712 None,
10713 None,
10714 None,
10715 None,
10716 Some(&lane_a),
10717 );
10718 finish.close(result_a.is_err());
10719 let result_b = b
10720 .join()
10721 .map_err(|_| "paired speculative session B panicked".to_string())
10722 .and_then(|r| r);
10723 (result_a, result_b)
10724 });
10725
10726 if (graph_a || graph_b) && was_tracking {
10727 unsafe {
10728 e.ctx().enable_event_tracking();
10729 }
10730 }
10731 let result_a = result_a?;
10732 let result_b = result_b.map_err(|err| -> Box<dyn std::error::Error> { err.into() })?;
10733 Ok((result_a, result_b))
10734 }
10735
10736 /// One spec-decode turn on a live session. `suffix` = the NEW tokens only (turn N+1's user
10737 /// message rendered through the chat template continuation). Returns (new tokens emitted,
10738 /// drafted, accepted); session.committed grows by suffix + emitted.
10739 pub fn generate_spec_session(
10740 &self,
10741 e: &Engine,
10742 sess: &mut SpecSession,
10743 suffix: &[u32],
10744 max_new: usize,
10745 k: usize,
10746 ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
10747 self.generate_spec_session_sampled(e, sess, suffix, max_new, k, None, None)
10748 }
10749
10750 /// Serve-path sampled spec: routes the burst through the rejection-sampling verify with
10751 /// per-SESSION Philox continuity (sess.sctr/uctr). None = env-driven (CLI) or greedy.
10752 /// Filters (top-k/p/min-p) apply SYMMETRICALLY to draft q and verify p — distribution-exact
10753 /// for the filtered target (feat/filtered-spec).
10754 ///
10755 /// `on_commit` (sse-cadence, 2026-08-05): called with each newly-emitted slice of the
10756 /// output — once right after the prime's first token, then once per round commit — so a
10757 /// streaming caller can flush text at round cadence instead of once per burst. The slices
10758 /// are disjoint, in order, and concatenate to exactly the returned token vec. Emission-
10759 /// timing only: token bytes, session state, and exactness are untouched.
10760 ///
10761 /// The returned bool is a CONTINUE-VERDICT (admission yield, 2026-08-06): `false` ends
10762 /// the burst at the current round boundary, exactly as if `max_new` had been reached —
10763 /// the caller's scheduler regains control without waiting the burst out. Burst size is
10764 /// content-neutral (spec-levers battery), so an early exit moves WHEN the burst returns,
10765 /// never what tokens say. The slice may be EMPTY (a poll-only boundary — round-stream
10766 /// drains and the defensive tail flush can land with nothing new committed).
10767 #[allow(clippy::too_many_arguments)]
10768 pub fn generate_spec_session_sampled(
10769 &self,
10770 e: &Engine,
10771 sess: &mut SpecSession,
10772 suffix: &[u32],
10773 max_new: usize,
10774 k: usize,
10775 sampling: Option<SpecSampling>,
10776 on_commit: Option<&mut dyn FnMut(&[u32]) -> bool>,
10777 ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
10778 self.generate_spec_session_sampled_prime_split(
10779 e, sess, suffix, max_new, k, sampling, None, on_commit,
10780 )
10781 }
10782
10783 /// Serve-only cold-prime segmentation twin. `prime_split` is the same stable boundary the
10784 /// plain worker would honor before entering its sub-floor tokenwise tail; warm continuations
10785 /// pass `None` and stay on the existing zero-prime path.
10786 #[allow(clippy::too_many_arguments)]
10787 pub fn generate_spec_session_sampled_prime_split(
10788 &self,
10789 e: &Engine,
10790 sess: &mut SpecSession,
10791 suffix: &[u32],
10792 max_new: usize,
10793 k: usize,
10794 sampling: Option<SpecSampling>,
10795 prime_split: Option<usize>,
10796 on_commit: Option<&mut dyn FnMut(&[u32]) -> bool>,
10797 ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
10798 self.generate_spec_session_constrained_prime_split(
10799 e,
10800 sess,
10801 suffix,
10802 max_new,
10803 k,
10804 sampling,
10805 None,
10806 prime_split,
10807 on_commit,
10808 )
10809 }
10810
10811 /// `generate_spec_session_sampled` + GRAMMAR (constrained decoding, 2026-08-03): the
10812 /// hook truncates acceptance at the first grammar-illegal token AFTER the exactness
10813 /// verify (grammar is an extra rejection rule, ordering like the batched-verify twins)
10814 /// and replaces an illegal bonus with the MASKED argmax of the target's own verify
10815 /// column — token-identical to constrained plain greedy decode. GREEDY only (the
10816 /// worker routes sampled constrained to plain decode). Acceptance under tight grammars
10817 /// may drop (drafter is unconstrained); that is measured, not hidden.
10818 #[allow(clippy::too_many_arguments)]
10819 pub fn generate_spec_session_constrained(
10820 &self,
10821 e: &Engine,
10822 sess: &mut SpecSession,
10823 suffix: &[u32],
10824 max_new: usize,
10825 k: usize,
10826 sampling: Option<SpecSampling>,
10827 constraint: Option<&mut dyn SpecConstraint>,
10828 on_commit: Option<&mut dyn FnMut(&[u32]) -> bool>,
10829 ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
10830 self.generate_spec_session_constrained_prime_split(
10831 e, sess, suffix, max_new, k, sampling, constraint, None, on_commit,
10832 )
10833 }
10834
10835 #[allow(clippy::too_many_arguments)]
10836 pub fn generate_spec_session_constrained_prime_split(
10837 &self,
10838 e: &Engine,
10839 sess: &mut SpecSession,
10840 suffix: &[u32],
10841 max_new: usize,
10842 k: usize,
10843 sampling: Option<SpecSampling>,
10844 constraint: Option<&mut dyn SpecConstraint>,
10845 prime_split: Option<usize>,
10846 on_commit: Option<&mut dyn FnMut(&[u32]) -> bool>,
10847 ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
10848 if constraint.is_some() && sampling.is_some_and(|s| s.temp > 0.0) {
10849 return Err(
10850 "constrained spec decode is greedy-only (worker routes sampled \
10851 constrained to plain decode)"
10852 .into(),
10853 );
10854 }
10855 // PENDING-CARRY entry flush: a carried bonus precedes any new suffix in the sequence,
10856 // so it must commit BEFORE the suffix primes; the sampled path doesn't carry (its
10857 // round-0 accept needs the commit pass's logits). Empty-suffix greedy bursts — the
10858 // serve continuation case — consume the carry in-loop with zero solo passes.
10859 if sess.pending_tok.is_some()
10860 && (!suffix.is_empty() || sampling.map_or(false, |s| s.temp > 0.0))
10861 {
10862 self.spec_flush_pending(e, sess, sampling)?;
10863 }
10864
10865 // FULL_PREC forces the EAGER draft: the graph capture would enclose cuBLASLt f32 GEMV
10866 // (the FloatBf16 else-branches) and a bf16_to_f32 dequant alloc — neither is stream-capture
10867 // safe. Eager rides matmul/matmul_decode_exact, which dequant FloatBf16 on use. (§item 2.)
10868 // Multi-head MTP (mtp_extra non-empty) no longer disqualifies: the chain captures
10869 // per-head graphs (lane/step37-draft-graph-serving-20260830, MEMRA_MTP_CHAIN_GRAPH).
10870 let graph_draft = std::env::var("MEMRA_SPEC_NOGRAPH").is_err()
10871 && !spec_host_embd()
10872 && self.mtp_graph_capturable()
10873 && k + 2 < 96
10874 && !crate::model::full_prec_enabled();
10875 let was_tracking = e.ctx().is_event_tracking();
10876 if graph_draft && was_tracking {
10877 unsafe {
10878 e.ctx().disable_event_tracking();
10879 }
10880 }
10881 let r = self.generate_spec_inner2(
10882 e,
10883 suffix,
10884 max_new,
10885 k,
10886 graph_draft,
10887 Some(sess),
10888 sampling,
10889 constraint,
10890 on_commit,
10891 prime_split,
10892 None,
10893 );
10894 if graph_draft && was_tracking {
10895 unsafe {
10896 e.ctx().enable_event_tracking();
10897 }
10898 }
10899 let (out, d, a) = r?;
10900 Ok((out, d, a))
10901 }
10902
10903 pub fn generate_spec(
10904 &self,
10905 e: &Engine,
10906 prompt: &[u32],
10907 max_new: usize,
10908 k: usize,
10909 ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
10910 if crate::pp::pp_cuts(self.layers.len()).is_some()
10911 && !self.rewrite_allowed(memra_gguf::execution_manifest::RewriteSurface::Pipeline)
10912 {
10913 return Err("pipeline rewrite is not qualified for speculative decode".into());
10914 }
10915 if !self.rewrite_allowed(memra_gguf::execution_manifest::RewriteSurface::MtpSpec) {
10916 return Err("speculative rewrite is not qualified for this ModelPlan".into());
10917 }
10918 // FULL_PREC forces eager (see generate_spec_session note): CUDA graph capture cannot
10919 // enclose cuBLASLt f32 GEMV or the bf16_to_f32 dequant alloc the FloatBf16 path needs.
10920 // Multi-head MTP no longer disqualifies (chain graphs; see generate_spec_session).
10921 let graph_draft = std::env::var("MEMRA_SPEC_NOGRAPH").is_err()
10922 && !spec_host_embd()
10923 && self.mtp_graph_capturable()
10924 && k + 2 < 96
10925 && !crate::model::full_prec_enabled();
10926 if !graph_draft {
10927 return self.generate_spec_inner2(
10928 e, prompt, max_new, k, false, None, None, None, None, None, None,
10929 );
10930 }
10931 let was_tracking = e.ctx().is_event_tracking();
10932 if was_tracking {
10933 unsafe {
10934 e.ctx().disable_event_tracking();
10935 }
10936 }
10937 let r = self.generate_spec_inner2(
10938 e, prompt, max_new, k, true, None, None, None, None, None, None,
10939 );
10940 if was_tracking {
10941 unsafe {
10942 e.ctx().enable_event_tracking();
10943 }
10944 }
10945 r
10946 }
10947
10948 fn generate_spec_inner2(
10949 &self,
10950 e: &Engine,
10951 prompt: &[u32],
10952 max_new: usize,
10953 k: usize,
10954 graph_draft: bool,
10955 mut sess: Option<&mut SpecSession>,
10956 sampling: Option<SpecSampling>,
10957 mut constraint: Option<&mut dyn SpecConstraint>,
10958 mut on_commit: Option<&mut dyn FnMut(&[u32]) -> bool>,
10959 prime_split: Option<usize>,
10960 pipe: Option<&SpecPipeLane>,
10961 ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
10962 assert!(k >= 1, "k must be >= 1");
10963 if let Some(p) = pipe {
10964 p.setup_begin()?;
10965 }
10966 // sse-cadence flush cursor: everything in out[..flushed] has been handed to on_commit.
10967 let mut flushed = 0usize;
10968 // admission yield (2026-08-06): on_commit's continue-verdict; false = end the burst
10969 // at the next round boundary (same exit as max_new reached — the session tail runs).
10970 // Initialized by the unconditional post-prime flush below.
10971 let mut keep_going;
10972 let mtp = self
10973 .mtp
10974 .as_ref()
10975 .expect("generate_spec requires an MTP head (nextn_predict_layers>0)");
10976 let n_vocab = self.output.out_features();
10977 // FR-Spec: the draft head may be TRIMMED (fewer rows than n_vocab); the draft argmax runs
10978 // over the draft vocab and the winning index maps through d2t to a TARGET token id.
10979 // Everything downstream (verify/accept/commit) sees target ids only — exactness unchanged.
10980 let d_vocab = mtp
10981 .shared_head_head
10982 .as_ref()
10983 .unwrap_or(&self.output)
10984 .out_features();
10985 if !self.mtp_extra.is_empty() {
10986 if self.plan.draft_source != memra_gguf::model_plan::DraftSourcePlan::Embedded
10987 || self.plan.mtp_blocks.len() != self.mtp_head_count()
10988 {
10989 return Err(
10990 "multi-head MTP requires one embedded canonical block per loaded head".into(),
10991 );
10992 }
10993 // TRIMMED chains (2026-08-27): every head must carry the SAME d2t — the ranking is
10994 // token-frequency and head-independent, and every downstream remap (per-step argmax,
10995 // stream pack, sampled d2t_dev) reads head 0's map, so equality is what makes that
10996 // single map correct for the whole chain. Mixed trimmed/untrimmed is refused.
10997 for (offset, head) in self.mtp_extra.iter().enumerate() {
10998 if head.d2t != mtp.d2t
10999 || head
11000 .shared_head_head
11001 .as_ref()
11002 .unwrap_or(&self.output)
11003 .out_features()
11004 != d_vocab
11005 {
11006 return Err(format!(
11007 "embedded MTP head {} has incompatible draft vocabulary",
11008 offset + 1
11009 )
11010 .into());
11011 }
11012 }
11013 eprintln!(
11014 "[mtp-chain] heads={} policy=step-modulo prefix-replay kv=per-head",
11015 self.mtp_head_count()
11016 );
11017 }
11018 let n_embd = self.cfg.n_embd as usize;
11019 // SESSION MODE: reuse the live cache/scratch, prime only the suffix. `base` = tokens
11020 // already committed (their state is in the caches); 0 = fresh single-shot call.
11021 let session_mode = sess.is_some();
11022 let max_ctx = match sess.as_ref() {
11023 Some(s) => s.cache.max_ctx,
11024 None => prompt.len() + max_new + k + 8,
11025 };
11026 let mut own_cache;
11027 let mut own_scratch;
11028 // PREFIX-CACHE capture request threaded out of the session (lane/spec-prefix-cache):
11029 // (requested split, destination list). Single-shot per burst; fresh calls have none.
11030 let mut sess_capture: Option<(Option<usize>, &mut Vec<SpecBoundaryCapture>)> = None;
11031 // STABLE-BOUNDARY turn-checkpoint request (lane/frspec-multiturn-cache): ABSOLUTE
11032 // committed-length position; consumed one-shot like `capture_at`. None = legacy
11033 // prompt-end capture below.
11034 let mut ckpt_req: Option<usize> = None;
11035 let (
11036 cache,
11037 scratch,
11038 mut sess_tail,
11039 mut sess_draft_slot,
11040 mut sess_pending_slot,
11041 sess_ckpt_slot,
11042 sess_telem,
11043 ): (
11044 &mut Cache,
11045 &mut MtpScratch,
11046 Option<(
11047 &mut Vec<u32>,
11048 &mut Option<CudaSlice<f32>>,
11049 &mut Option<u32>,
11050 &mut u32,
11051 &mut u32,
11052 )>,
11053 Option<&mut Option<DraftGraphCtx>>,
11054 Option<&mut Option<u32>>,
11055 Option<&mut Option<SpecCheckpoint>>,
11056 Option<&SpecTelemetryCounters>,
11057 ) = match sess.take() {
11058 Some(sr) => {
11059 let SpecSession {
11060 cache,
11061 scratch,
11062 committed,
11063 last_h,
11064 next_pred,
11065 sctr: s_sctr,
11066 uctr: s_uctr,
11067 draft_ctx,
11068 pending_tok,
11069 turn_ckpt,
11070 telem,
11071 capture_at,
11072 boundary_captures,
11073 ckpt_at,
11074 } = sr;
11075 sess_capture = Some((capture_at.take(), boundary_captures));
11076 ckpt_req = ckpt_at.take();
11077 (
11078 cache,
11079 scratch,
11080 Some((committed, last_h, next_pred, s_sctr, s_uctr)),
11081 Some(draft_ctx),
11082 Some(pending_tok),
11083 Some(turn_ckpt),
11084 Some(telem),
11085 )
11086 }
11087 None => {
11088 // STAGE-OWNED KV (lane/pp2-spec 2026-08-06) — see `new_session`. Door shut =
11089 // `Cache::new` verbatim.
11090 own_cache = crate::pp::new_cache_planned(e, &self.cfg, &self.plan, max_ctx)?;
11091 // Persistent scratch = max_ctx rows (~2KB/token quantized).
11092 own_scratch = self.new_mtp_scratch(e, max_ctx)?;
11093 (
11094 &mut own_cache,
11095 &mut own_scratch,
11096 None,
11097 None,
11098 None,
11099 None,
11100 None,
11101 )
11102 }
11103 };
11104 if scratch.plane_count() != self.mtp_head_count() {
11105 return Err(format!(
11106 "MTP scratch/head count mismatch ({}/{})",
11107 scratch.plane_count(),
11108 self.mtp_head_count()
11109 )
11110 .into());
11111 }
11112 let base = cache.pos;
11113 // PENDING-CARRY consume (2026-08-01): a carried bonus reaches here only on the
11114 // empty-suffix GREEDY continuation path (generate_spec_session_sampled flushed every
11115 // other case). It enters the round loop as round-0's pending — verify col 0 — exactly
11116 // like a mid-burst full-accept boundary: no init feed, no tail commit pass.
11117 let carried_pending: Option<u32> = sess_pending_slot.as_mut().and_then(|s| s.take());
11118 // PERSISTENT DRAFT KV (the only mode since 2026-07-08 — the legacy round-local scratch,
11119 // MEMRA_SPEC_KVLOCAL, measured -35 acceptance pts on the 27B p3 sweep and was removed;
11120 // acceptance-only — exactness is verify's job either way).
11121 // HIDDEN-PAIRING CONVENTION (DEFAULT = predecessor-row, 2026-07-04 — the 27B acceptance
11122 // unlock, +16pts): the MTP head is TRAINED on rows pairing token x_p with the trunk
11123 // hidden of its PREDECESSOR h_{p-1} (the reference engine's mtp_update shifts the target
11124 // hiddens right by one; its draft step 0 feeds (id_last, TRUE hidden of the row id_last
11125 // was sampled from)). memra's historical convention paired SAME-ROW (x_p, h_p) in the fill
11126 // and seeded chain step 0 through an extra MTP pass on a duplicated token (the
11127 // pseudo-seed) — measured 27B p2 K=3 acceptance 0.569 vs 0.731, p3 0.445 vs 0.63+, and
11128 // the chain steps j>=1 were already predecessor-shaped, so ONLY the fill + step-0 seed
11129 // move. The fill shifts by one and the chain seeds from the predecessor's true hidden
11130 // DIRECTLY (vh_seed / vx[j-1]) — the pseudo pass disappears (one MTP-block pass saved
11131 // per round on top of the acceptance win). Draft-quality-only: exactness stays the
11132 // verify's job either way. (The legacy same-row pairing seam, MEMRA_SPEC_HSAME, and its
11133 // pseudo-seed passes were removed 2026-07-08 — predecessor pairing won by +16 acc pts;
11134 // the legacy round-local scratch, MEMRA_SPEC_KVLOCAL, went with it.)
11135 // REPLAY-FREE PARTIAL ACCEPT (default, 2026-07-03): partial rounds keep the verify's own
11136 // bit-identical committed-prefix state (KV truncate + recur rebuild from the VerifyCkpt)
11137 // and leave the bonus PENDING — no duplicate trunk pass (profiled ~0.54 extra full weight
11138 // reads/round at long ctx). MEMRA_SPEC_REPLAY=1 restores the legacy rollback+replay (A/B
11139 // + fallback seam).
11140 // Qwen35-MoE replay pin LIFTED (lane/draftcost-moe, 2026-08-20). The pin's stated
11141 // bar — the retained verify-state commit proven equivalent to sequential serving —
11142 // was waiting on this arch running the serving batched verify class, which the
11143 // t-parallel admission (this lane, increment 1) provided: the VerifyCkpt the
11144 // replay-free commit consumes is now produced by the SAME serving-class verify that
11145 // qualified dense qwen35 on 2026-08-15 (where the per-round duplicate replay
11146 // measured 69 -> 30 tok/s). Qualification receipts (run-spec K=1..8 both arms,
11147 // 8-prompt replay-vs-replay-free canary, long-prompt cell):
11148 // research/draftcost-moe-20260820/RECEIPTS.md. MEMRA_SPEC_REPLAY=1 stays the
11149 // rollback + A/B seam.
11150 let spec_replay = spec_replay_env_enabled();
11151 if constraint.is_some() && spec_replay {
11152 return Err(
11153 "constrained spec decode does not support MEMRA_SPEC_REPLAY=1 \
11154 (legacy replay commits an unmasked bonus)"
11155 .into(),
11156 );
11157 }
11158 // TRUE-HIDDEN REFRESH (default in persistent-draft-KV mode): every round overwrites the
11159 // committed positions' scratch entries from the verify's exact hiddens (mtp_kv_fill batch)
11160 // instead of keeping chain-approximate entries. MEMRA_SPEC_NOREFRESH=1 = legacy (A/B seam).
11161 let refresh = std::env::var("MEMRA_SPEC_NOREFRESH").is_err();
11162 if !refresh && !self.mtp_extra.is_empty() {
11163 return Err("multi-head MTP requires exact accepted-prefix refresh".into());
11164 }
11165
11166 // prime: BATCHED cache prime (prime_cache — the measured #1 e2e gap: tokenwise primed at
11167 // ~102/38 tok/s vs the engine's ~2000-5900 tok/s batched prefill). prime_cache returns the
11168 // full pre-output_norm hidden stack [T, n_embd], which IS prompt_h (the persistent-draft-KV
11169 // mtp_kv_fill input) — no per-token collection needed. Prompts below PRIME_MIN_T, and
11170 // MEMRA_PRIME_TOKENWISE=1, and frozen Hy3 CPU/GPU expert splits take the tokenwise
11171 // decode_step_h loop. The latter avoids transient GPU staging of the spilled expert bank.
11172 // EMPTY-SUFFIX CONTINUATION (serve bursts): a session turn with NO new tokens resumes
11173 // generation exactly where the last turn stopped — no prime at all. The stashed
11174 // `next_pred` plays prime_logits' role: it is the token produced from the logits after
11175 // committed.last() by the same rule this entry applies to a cold prime's last row —
11176 // an argmax when greedy, a `sample_boundary_token` draw when sampled (the burst tail,
11177 // or `spec_session_from_restored` for a converted prefix-cache hit, did the drawing
11178 // where the sampler and the session's Philox counters were live). `last_h` seeds the
11179 // predecessor pairing below. Fresh calls and non-empty suffixes take the normal path.
11180 let continuation = prompt.is_empty();
11181 if continuation {
11182 assert!(session_mode, "empty prompt requires a session");
11183 assert!(
11184 sess_tail
11185 .as_ref()
11186 .map_or(false, |(c, lh, np, _, _)| !c.is_empty()
11187 && lh.is_some()
11188 && (np.is_some() || carried_pending.is_some())),
11189 "empty-suffix continuation needs a primed session (committed + last_h + next_pred|pending)"
11190 );
11191 }
11192 let mut prime_logits;
11193 let mut prompt_h: Option<CudaSlice<f32>> = None;
11194 let t_prime = std::time::Instant::now();
11195 let batched_prime = !continuation
11196 && prompt.len() >= crate::hybrid_forward::PRIME_MIN_T
11197 && std::env::var("MEMRA_PRIME_TOKENWISE").is_err()
11198 && !e.frozen_cpu_experts_prefer_tokenwise_prime();
11199 let prime_split = prime_split.filter(|&split| split > 0 && split < prompt.len());
11200 if prime_split.is_some() && continuation {
11201 return Err("spec prime split requires a non-empty prime".into());
11202 }
11203 // STABLE-BOUNDARY TURN CHECKPOINT stop (lane/frspec-multiturn-cache, 2026-08-21):
11204 // the worker's `ckpt_at` request, ABSOLUTE -> prompt-relative. On WARM bursts
11205 // (base != 0, an affinity-rewound or pool-resumed session priming its own delta)
11206 // this is the only stop; on COLD bursts it usually coincides with `prime_split`
11207 // (both are the plain tier's stable pre-generation boundary). A boundary the prime
11208 // cannot honor (outside this prime's range) silently drops the capture — the
11209 // turn_ckpt convention: the next turn re-primes in full, never a wrong resume.
11210 let ckpt_rel = if continuation {
11211 None
11212 } else {
11213 ckpt_req
11214 .and_then(|abs| abs.checked_sub(base))
11215 .filter(|&r| r > 0 && r < prompt.len())
11216 };
11217 // Prime stops, ordered: each is a boundary the prime halts at so the in-place GDN
11218 // conv/ssm state can be snapshotted there (the only moment it exists). One stop =
11219 // the legacy single-split program, byte-for-byte.
11220 let mut stops: Vec<usize> = Vec::new();
11221 for b in [prime_split, ckpt_rel].into_iter().flatten() {
11222 if !stops.contains(&b) {
11223 stops.push(b);
11224 }
11225 }
11226 stops.sort_unstable();
11227 // Captured at the ckpt stop, installed into the session slot post-prime (replacing
11228 // the legacy prompt-end capture). Some(None) = capture attempted and failed -> the
11229 // slot is cleared (a stale checkpoint would rewind to the WRONG boundary).
11230 let mut ckpt_early: Option<Option<SpecCheckpoint>> = None;
11231 if continuation {
11232 prime_logits = Vec::new();
11233 } else if !stops.is_empty() {
11234 if let Some(&first) = stops.first() {
11235 if prime_split == Some(first) && first < crate::hybrid_forward::PRIME_MIN_T {
11236 return Err(format!(
11237 "spec prime split {first} is below PRIME_MIN_T {}",
11238 crate::hybrid_forward::PRIME_MIN_T,
11239 )
11240 .into());
11241 }
11242 }
11243 // Mirror the plain worker's boundary stops exactly. Each segment is a
11244 // request-level prime (`queued_after` keeps Step35 arm selection independent of
11245 // the stops — tick-seg law); a segment below PRIME_MIN_T (and the final tail
11246 // under MEMRA_PRIME_TOKENWISE) takes the same eager tokenwise continuation as
11247 // prefill_tick. Retain every hidden row so the draft scratch fill remains one
11248 // coherent prompt.
11249 let mut h_all = e.uninit(prompt.len() * n_embd)?;
11250 prime_logits = Vec::new();
11251 let mut prev = 0usize;
11252 for seg_end in stops.iter().copied().chain(std::iter::once(prompt.len())) {
11253 if seg_end <= prev {
11254 continue;
11255 }
11256 let seg = &prompt[prev..seg_end];
11257 let is_final = seg_end == prompt.len();
11258 let batched_seg = seg.len() >= crate::hybrid_forward::PRIME_MIN_T
11259 && (!is_final
11260 || (std::env::var("MEMRA_PRIME_TOKENWISE").is_err()
11261 && !e.frozen_cpu_experts_prefer_tokenwise_prime()));
11262 if batched_seg {
11263 let (l, _, h_seg) =
11264 self.prime_cache(e, seg, &mut *cache, prompt.len() - seg_end)?;
11265 e.copy_into(&mut h_all, prev * n_embd, &h_seg, seg.len() * n_embd)?;
11266 prime_logits = l;
11267 } else {
11268 for (i, &tok) in seg.iter().enumerate() {
11269 let (l, h) = self.decode_step_h(e, tok, &mut *cache)?;
11270 e.copy_into(&mut h_all, (prev + i) * n_embd, &h, n_embd)?;
11271 prime_logits = l;
11272 }
11273 }
11274 prev = seg_end;
11275 if is_final {
11276 break;
11277 }
11278 debug_assert_eq!(cache.pos, base + seg_end, "prime stop landed off boundary");
11279 // PREFIX-CACHE BOUNDARY CAPTURE (lane/spec-prefix-cache): the GDN conv/ssm
11280 // states are about to be advanced in place by the next segment, so this is
11281 // the ONLY moment the boundary's recurrent state exists. Capture iff the
11282 // worker requested exactly this stop (cold sessions only — `capture_at` is
11283 // never armed warm). A failed snapshot is silent (turn_ckpt convention) —
11284 // publication is an optimization, never a correctness dependency.
11285 if base == 0 {
11286 if let Some((requested, slot)) = sess_capture.as_mut() {
11287 // Publish at the requested miss-LCP stop (the shared-prefix class)
11288 // AND at the stable-boundary stop (the next-turn re-render class,
11289 // lane/frspec-multiturn-cache) — the same boundary set the plain
11290 // prefill tick learns. Without the second entry, the turn after a
11291 // cold re-park could only hit the OLDER lcp entry (the measured
11292 // one-turn transient: t3 restored 607 of 24122 while the plain arm
11293 // rewound to 15222). Dedupe is the worker sweep's has_key.
11294 if *requested == Some(seg_end) || ckpt_rel == Some(seg_end) {
11295 if let Ok(snap) = cache.snapshot(e) {
11296 slot.push(SpecBoundaryCapture {
11297 snap,
11298 pos: seg_end,
11299 logits: prime_logits.clone(),
11300 // rows [0..seg_end) of h_all are primed — the following
11301 // segments append, never overwrite.
11302 last_h: capture_boundary_hidden(e, &h_all, seg_end, n_embd),
11303 });
11304 }
11305 }
11306 }
11307 }
11308 // SESSION-AFFINITY TURN CHECKPOINT at the STABLE boundary (see `ckpt_at`):
11309 // same snapshot mechanics, installed post-prime in place of the prompt-end
11310 // capture the re-render class always diverged below.
11311 if ckpt_rel == Some(seg_end) {
11312 let anchor: Result<CudaSlice<f32>, Box<dyn std::error::Error>> =
11313 e.uninit(n_embd).and_then(|mut a| {
11314 e.copy_view_into(
11315 &mut a,
11316 0,
11317 &h_all.slice((seg_end - 1) * n_embd..seg_end * n_embd),
11318 n_embd,
11319 )?;
11320 Ok(a)
11321 });
11322 ckpt_early = Some(match (cache.snapshot(e), anchor) {
11323 (Ok(snap), Ok(last_h)) => Some(SpecCheckpoint {
11324 snap,
11325 pos: base + seg_end,
11326 last_h,
11327 }),
11328 _ => None,
11329 });
11330 }
11331 }
11332 if std::env::var("MEMRA_SPEC_STATS").as_deref() == Ok("1") {
11333 eprintln!(
11334 "[spec-prime] stops={stops:?} tail={}",
11335 prompt.len() - stops.last().copied().unwrap_or(0)
11336 );
11337 }
11338 prompt_h = Some(h_all);
11339 } else if batched_prime {
11340 let (l, _h_seed, hiddens) = self.prime_cache(e, prompt, &mut *cache, 0)?;
11341 prime_logits = l;
11342 prompt_h = Some(hiddens);
11343 } else {
11344 prime_logits = Vec::new();
11345 prompt_h = Some(e.uninit(prompt.len() * n_embd)?);
11346 for (i, &tok) in prompt.iter().enumerate() {
11347 let (l, h) = self.spec_target_step_h(e, tok, &mut *cache)?;
11348 if let Some(ph) = prompt_h.as_mut() {
11349 e.copy_into(ph, i * n_embd, &h, n_embd)?;
11350 }
11351 prime_logits = l;
11352 }
11353 }
11354 e.stream().synchronize()?;
11355 // PREFIX-CACHE SEED CAPTURE (lane/spec-prefix-cache): boundary == prompt end (the seed
11356 // case — no shared-prefix split, publish the whole prompt). The prime just finished, so
11357 // cache.pos == base + prompt.len() and the recurrent state IS the boundary state;
11358 // prime_logits are the boundary logits. Cold sessions only (base == 0) — same law as
11359 // prime_split. The mid-prompt capture above already consumed the request if it matched.
11360 if !continuation && base == 0 {
11361 if let Some((requested, slot)) = sess_capture.as_mut() {
11362 if *requested == Some(prompt.len()) && slot.is_empty() {
11363 debug_assert_eq!(cache.pos, prompt.len(), "seed capture off prompt end");
11364 if let Ok(snap) = cache.snapshot(e) {
11365 slot.push(SpecBoundaryCapture {
11366 snap,
11367 pos: prompt.len(),
11368 logits: prime_logits.clone(),
11369 last_h: prompt_h
11370 .as_ref()
11371 .map(|ph| capture_boundary_hidden(e, ph, prompt.len(), n_embd))
11372 .unwrap_or_default(),
11373 });
11374 }
11375 }
11376 }
11377 }
11378 // Harness timing contract (see crate::PRIME_NANOS): gen-only throughput without the
11379 // prime-subtraction hack.
11380 crate::PRIME_NANOS.store(
11381 t_prime.elapsed().as_nanos() as u64,
11382 std::sync::atomic::Ordering::Relaxed,
11383 );
11384
11385 let (embd_qt, embd_rb) = self.embd.qt_and_row_bytes(n_embd);
11386 // Resident table is fastest when it fits. Large spill deployments can preserve that HBM
11387 // for expert-cache slots and gather only the exact rows needed by MTP/verify from host.
11388 let host_embd = spec_host_embd();
11389 let embd_gpu = if host_embd {
11390 None
11391 } else {
11392 Some(
11393 self.embd_gpu
11394 .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload")),
11395 )
11396 };
11397 let embd_dev = embd_gpu.map(|g| (g, embd_qt, embd_rb));
11398 if host_embd {
11399 eprintln!(
11400 "[spec] host-row embedding: {} bytes kept off HBM",
11401 self.embd.raw.len()
11402 );
11403 }
11404 let mut out: Vec<u32> = Vec::with_capacity(max_new);
11405 let mut total_drafted = 0usize;
11406 let mut total_accepted = 0usize;
11407
11408 // --- SAMPLER FIRST (lane/sampled-spec-quality, 2026-08-19) ---
11409 // The sampler config, the session's Philox counters and the penalty window are parsed
11410 // HERE, above the boundary-token selection, because the boundary token must be drawn
11411 // from the sampler the request asked for. Pre-lane this block sat ~50 lines BELOW the
11412 // selection, which is the whole mechanical reason the boundary token was an argmax:
11413 // the sampler state was not in scope yet. Nothing here depends on the round loop, so
11414 // moving it up is a pure reordering for greedy (`sampled == false` ⇒ every branch
11415 // below takes the argmax path it always took).
11416 // --- SAMPLED SPEC (MEMRA_SPEC_TEMP>0, research/sampled-spec-impl-map.md): rejection-
11417 // sampling verify (Leviathan/Chen) — accept draft x at u < p(x)/q(x), resample from
11418 // norm(max(0,p-q)) on reject, bonus sampled from p on full accept. Counter-based Philox
11419 // everywhere (seed, event) -> reproducible. temp==0/unset = the greedy path, untouched.
11420 let sp = sampling.unwrap_or_else(|| SpecSampling {
11421 temp: std::env::var("MEMRA_SPEC_TEMP")
11422 .ok()
11423 .and_then(|v| v.parse().ok())
11424 .unwrap_or(0.0),
11425 seed: std::env::var("MEMRA_SEED")
11426 .ok()
11427 .and_then(|v| v.parse().ok())
11428 .unwrap_or(42),
11429 top_k: std::env::var("MEMRA_TOP_K")
11430 .ok()
11431 .and_then(|v| v.parse().ok())
11432 .unwrap_or(0),
11433 top_p: std::env::var("MEMRA_TOP_P")
11434 .ok()
11435 .and_then(|v| v.parse().ok())
11436 .unwrap_or(1.0),
11437 min_p: std::env::var("MEMRA_MIN_P")
11438 .ok()
11439 .and_then(|v| v.parse().ok())
11440 .unwrap_or(0.0),
11441 penalty_last_n: std::env::var("MEMRA_PENALTY_LAST_N")
11442 .ok()
11443 .and_then(|v| v.parse().ok())
11444 .unwrap_or(0),
11445 penalty_repeat: std::env::var("MEMRA_PENALTY_REPEAT")
11446 .ok()
11447 .and_then(|v| v.parse().ok())
11448 .unwrap_or(1.0),
11449 penalty_freq: std::env::var("MEMRA_PENALTY_FREQ")
11450 .ok()
11451 .and_then(|v| v.parse().ok())
11452 .unwrap_or(0.0),
11453 penalty_present: std::env::var("MEMRA_PENALTY_PRESENT")
11454 .ok()
11455 .and_then(|v| v.parse().ok())
11456 .unwrap_or(0.0),
11457 });
11458 let (sp_temp, sp_seed) = (sp.temp, sp.seed);
11459 let sampled = sp_temp > 0.0;
11460 // Counters resume from the session (burst continuity: randomness must never repeat
11461 // across generate_spec_session calls); one-shot callers start at (0,0). Read through
11462 // sess_tail — `sess` was take()n into it above, so sess.as_ref() here is always None.
11463 let mut sctr: u32 = sess_tail.as_ref().map(|(_, _, _, s, _)| **s).unwrap_or(0);
11464 let mut uctr: u32 = sess_tail.as_ref().map(|(_, _, _, _, u)| **u).unwrap_or(0);
11465 // Penalties (v2.1): applied to COPIES of q rows and p columns symmetrically (exactness
11466 // for the penalized+filtered target). History = generated tokens, host-tracked window.
11467 let pen_on = sampled
11468 && sp.penalty_last_n > 0
11469 && (sp.penalty_repeat != 1.0 || sp.penalty_freq != 0.0 || sp.penalty_present != 0.0);
11470 // SESSION-SPANNING PENALTY WINDOW (Item 2). Pre-lane this was
11471 // `prompt.iter().rev().take(64).rev()` — the BURST's suffix slice — so a continuation
11472 // burst (the majority of a stream's tokens, and ALL of a converted cache hit's) started
11473 // with an EMPTY penalty history and the client's repetition/frequency/presence penalties
11474 // silently reset at every burst boundary. The window now spans `committed ++ prompt`,
11475 // which is what the API contract says and what the plain sampler's own `history` does.
11476 // Byte-identical to the pre-lane seed for a cold turn-1 burst at the default window.
11477 let mut pen_hist: Vec<u32> = if pen_on {
11478 let sess_hist: &[u32] = if spec_pen_session_on() {
11479 sess_tail
11480 .as_ref()
11481 .map(|(c, ..)| c.as_slice())
11482 .unwrap_or(&[])
11483 } else {
11484 &[] // MEMRA_SPEC_PEN_SESSION=0: pre-lane burst-local window
11485 };
11486 pen_window_seed(sess_hist, prompt, sp.penalty_last_n)
11487 } else {
11488 Vec::new()
11489 };
11490 // First generated token = the BOUNDARY token: greedy takes the argmax of the prompt's
11491 // last logits (== greedy's first token, byte-contract); SAMPLED draws it from the
11492 // request's own filtered/penalized target through the session's Philox stream
11493 // (`sample_boundary_token`, lane/sampled-spec-quality Item 1 — pre-lane this was an
11494 // argmax in both regimes, so ~1 token per burst of a sampled stream was greedy).
11495 // Emit it, then FEED it to establish the loop invariant below.
11496 // PENDING-CARRY: the carried bonus was already emitted by the LAST burst — it becomes
11497 // last_token WITHOUT re-emission, and round 0 consumes it as pending (no init feed).
11498 // CONSTRAINED entry rules: the first emitted token is the MASKED argmax of the
11499 // prompt's last logits (plain constrained-greedy identity); a continuation without
11500 // a carried pending would emit an UNMASKED stashed next_pred — refused loudly (the
11501 // worker never resumes constrained sessions from the pool, so this cannot fire).
11502 if let Some(c) = constraint.as_deref_mut() {
11503 if continuation && carried_pending.is_none() {
11504 return Err("constrained spec continuation requires a carried pending \
11505 (pool resume is unconstrained-only)"
11506 .into());
11507 }
11508 if !continuation {
11509 c.mask_logits(&mut prime_logits)
11510 .map_err(|e2| format!("constraint: {e2}"))?;
11511 }
11512 }
11513 let mut last_token = if let Some(b) = carried_pending {
11514 b
11515 } else if continuation {
11516 // A continuation's boundary token was DRAWN by the burst that stashed it (the
11517 // session tail below), or by `spec_session_from_restored` for a converted
11518 // prefix-cache hit — in both cases from the correct logits row with this same
11519 // session's Philox stream, which is why it can be consumed here as-is.
11520 sess_tail.as_ref().unwrap().2.unwrap()
11521 } else if sampled && constraint.is_none() && spec_sampled_boundary_on() {
11522 sample_boundary_token(e, &prime_logits, &sp, &pen_hist, &mut sctr, "cold-prime")?
11523 } else {
11524 // greedy (byte contract), the rollback door, or constrained (masked-argmax
11525 // identity — the worker routes sampled+constrained to the plain path, and this
11526 // function refuses the combination outright above).
11527 argmax(&prime_logits) as u32
11528 };
11529 if pen_on {
11530 // The boundary token is a GENERATED token: the plain sampler `accept()`s every
11531 // emitted token into its penalty history, and pre-lane the burst's first token
11532 // was invisible to penalties forever (never pushed, and never in `committed`
11533 // until this burst's tail). Covers the carry/continuation seeds too — neither is
11534 // in `committed` yet.
11535 pen_hist.push(last_token);
11536 }
11537 if carried_pending.is_none() {
11538 out.push(last_token);
11539 // grammar advances with every emitted token (carried pendings were consumed
11540 // by the burst that emitted them).
11541 if let Some(c) = constraint.as_deref_mut() {
11542 c.consume(last_token)
11543 .map_err(|e2| format!("constraint: {e2}"))?;
11544 }
11545 }
11546 if continuation {
11547 // draft-KV invariant: entries [0..base) are the session's exact fills; truncate any
11548 // overhang so the chain's first append lands at slot base (== committed.len()).
11549 scratch.set_len(e, base)?;
11550 }
11551 // sse-cadence: hand the caller every not-yet-flushed token (disjoint in-order slices
11552 // concatenating to the full `out`). Called after the prime's first token and after each
11553 // round commit — emission timing only, token bytes untouched. The slice may be EMPTY
11554 // (poll-only boundary: zero-round folds commit nothing new); returns the caller's
11555 // continue-verdict (admission yield, 2026-08-06) — false ends the burst at this round.
11556 fn flush_commit(
11557 cb: &mut Option<&mut dyn FnMut(&[u32]) -> bool>,
11558 out: &[u32],
11559 flushed: &mut usize,
11560 ) -> bool {
11561 if let Some(f) = cb.as_mut() {
11562 let keep = f(&out[*flushed..]);
11563 *flushed = out.len();
11564 keep
11565 } else {
11566 true
11567 }
11568 }
11569 keep_going = flush_commit(&mut on_commit, &out, &mut flushed);
11570 // INVARIANT at loop top: `last_token` is the most-recently-committed/emitted token, its
11571 // KV+recur state IS in `cache` (cache.pos = position right AFTER last_token), `last_pred`
11572 // is the greedy ARGMAX of the logits that predict the token FOLLOWING last_token, and
11573 // `h_seed` = last_token's pre-output_norm hidden. Establish it by feeding last_token once
11574 // (mirrors plain greedy). DEVICE-ARGMAX lever: the accept walk only ever consumes the
11575 // argmax of those logits — never the full vector — so a host u32 replaces the Vec<f32>.
11576 // Trimmed heads: q lives on the trimmed vocab; accept gathers use the TRIMMED index and
11577 // the residual scatters q into target-id space (q=-inf off-trim — the head cannot propose
11578 // those, so their residual mass is p(x), correct by construction).
11579 let d2t_dev: Option<CudaSlice<u32>> = if sampled || crate::spec::spec_stream() {
11580 match &mtp.d2t {
11581 Some(map) => Some(e.htod_u32_v(map)?),
11582 None => None,
11583 }
11584 } else {
11585 None
11586 };
11587 let mut q_full_buf: Option<CudaSlice<f32>> = None;
11588 // host Philox4x32-10 accept-test uniforms: module fn `host_u01` (shared with the
11589 // dspark sampled-admission walk); byte-identical to the closure it replaces.
11590 let mut draft_logits: Vec<CudaSlice<f32>> = Vec::new(); // retained head logits (q), per slot
11591 let mut draft_stats: Vec<(f32, f32, f32)> = Vec::new(); // (row_max, th_e, z_e) per slot
11592 let mut perturb_buf: Option<CudaSlice<f32>> = None; // gumbel scratch (max(n_vocab,d_vocab))
11593 let mut sample_tok = e.alloc_u32_zeroed(1)?; // residual/bonus sample out
11594 let mut col_buf: Option<CudaSlice<f32>> = None; // materialized verify column
11595 let mut pen_hist_d: Option<CudaSlice<u32>> = None;
11596 let mut pcol_buf: Option<CudaSlice<f32>> = None; // penalized p-column scratch
11597 // MEMRA_SPEC_SETUP_TRACE=1 (diagnostics): per-call wall decomposition of the burst
11598 // SETUP + TAIL segments (the round loop's internals are MEMRA_SPEC_PHASE's job) —
11599 // built to pin the serve per-burst fixed cost (research/spec-serving-20260801).
11600 let setup_trace = std::env::var("MEMRA_SPEC_SETUP_TRACE").as_deref() == Ok("1");
11601 let t_ent = std::time::Instant::now();
11602
11603 // SESSION-AFFINITY TURN CHECKPOINT (lane/session-affinity, 2026-08-05): capture the
11604 // PROMPT-END boundary state so a LATER turn can rewind here and re-prime only its own
11605 // delta instead of the whole conversation. See `SpecCheckpoint` for why this boundary is
11606 // the one that matters (a history-rewriting client mutates what the session GENERATED,
11607 // so the next turn's prompt agrees with this one up to exactly here).
11608 //
11609 // WHERE — AND WHY THIS EXACT LINE. Right after the trunk prime, BEFORE the init feed
11610 // (`decode_step_h(last_token)`) and before round 0: the last instant at which the caches
11611 // hold exactly `base + prompt.len()` rows and nothing generated.
11612 //
11613 // This was WRONG in the first cut of this lane: the capture sat after the draft-KV fill,
11614 // which is also after the init feed, so `cache.pos` was `base + prompt.len() + 1` — the
11615 // boundary included the FIRST GENERATED TOKEN. That token is the first thing inside the
11616 // `<think>` block the client strips, so every later turn's diff diverged exactly one
11617 // token below the checkpoint and affinity declined 100% of the time. Measured on the
11618 // owner regime: "history diverged at 12233 of checkpoint 12234". The off-by-one made the
11619 // whole mechanism inert while looking, from the outside, like a working
11620 // correctness-declines-safely path — hence the decline log carries the offsets.
11621 //
11622 // The full-attn planes are `len`-truncatable so the snapshot copies only the GDN conv/ssm
11623 // state (the reason a spec session could not rewind before). The draft scratch needs no
11624 // copy: rows below the boundary are rewritten by the next turn's own fill.
11625 //
11626 // WHEN: non-empty prime only. An empty-suffix continuation burst adds no prompt boundary
11627 // (its "prompt end" IS the previous checkpoint's, already held), so it keeps the existing
11628 // checkpoint rather than replacing it with a strictly worse one.
11629 //
11630 // FAILURE IS SILENT BY DESIGN: on a VRAM-tight rig the snapshot alloc can fail. That
11631 // costs the NEXT turn its rewind (it re-primes fully, today's behavior) and must never
11632 // fail the burst that is already running — so the error is swallowed, loud only under
11633 // MEMRA_DEBUG_SPEC.
11634 //
11635 // STABLE-BOUNDARY OVERRIDE (lane/frspec-multiturn-cache, 2026-08-21): the prompt-end
11636 // posture above was DISPROVED for the think-posture template class — the prompt's own
11637 // tail is the live generation header (`<|im_start|>assistant\n<think>\n`) that the
11638 // next turn's re-render replaces, so the diff diverged a couple tokens BELOW the
11639 // checkpoint and affinity declined 100% of multi-turn agent traffic (the same class
11640 // the plain tier fixed on 2026-08-09 via `plain_checkpoint_boundary`; the port to the
11641 // spec tier is this lane). When the worker armed `ckpt_at`, the capture happened at
11642 // that stop inside the prime above (`ckpt_early`) and is installed here instead;
11643 // capture-attempted-but-failed clears the slot exactly like the legacy arm.
11644 if let Some(slot) = sess_ckpt_slot {
11645 if let Some(early) = ckpt_early {
11646 if early.is_none() && std::env::var("MEMRA_DEBUG_SPEC").is_ok() {
11647 eprintln!(
11648 "[spec] stable-boundary turn checkpoint skipped; \
11649 next turn re-primes in full"
11650 );
11651 }
11652 *slot = early;
11653 } else if !continuation {
11654 let pos = cache.pos;
11655 debug_assert_eq!(
11656 pos,
11657 base + prompt.len(),
11658 "turn checkpoint must sit at the prompt end, before the init feed"
11659 );
11660 let anchor: Result<CudaSlice<f32>, Box<dyn std::error::Error>> =
11661 if let Some(ph) = &prompt_h {
11662 // hidden of the LAST primed row = the predecessor anchor at this
11663 // boundary (exactly what a fresh prime of committed[..pos] leaves in
11664 // last_h, and what the next prime's fill reads for its first row).
11665 let np = prompt.len();
11666 e.uninit(n_embd).and_then(|mut a| {
11667 e.copy_view_into(
11668 &mut a,
11669 0,
11670 &ph.slice((np - 1) * n_embd..np * n_embd),
11671 n_embd,
11672 )?;
11673 Ok(a)
11674 })
11675 } else {
11676 Err("no prompt hiddens".into())
11677 };
11678 match (cache.snapshot(e), anchor) {
11679 (Ok(snap), Ok(last_h)) => {
11680 *slot = Some(SpecCheckpoint { snap, pos, last_h });
11681 }
11682 (s, a) => {
11683 *slot = None; // a stale checkpoint would rewind to the WRONG boundary
11684 if std::env::var("MEMRA_DEBUG_SPEC").is_ok() {
11685 let err = s
11686 .err()
11687 .map(|e| e.to_string())
11688 .or_else(|| a.err().map(|e| e.to_string()))
11689 .unwrap_or_default();
11690 eprintln!(
11691 "[spec] turn checkpoint skipped ({err}); \
11692 next turn re-primes in full"
11693 );
11694 }
11695 }
11696 }
11697 }
11698 }
11699 // INIT FEED — skipped on a pending carry: last_token (the carried bonus) is NOT in the
11700 // caches and must NOT be fed solo; round 0's batched verify commits it as col 0. Its
11701 // seed/anchor hidden is the carried last_h (copied below); last_pred is dead in the
11702 // pending path (t_pred reads verify col 0 — the accept walk overwrites it).
11703 let mut last_pred = 0u32;
11704 let mut last_col_logits: Option<CudaSlice<f32>> = None;
11705 // CONSTRAINED: the init feed's logits back the (n_acc==0, base==0) masked-argmax
11706 // recompute in the grammar-truncation walk — retained host-side, round 0 only.
11707 let mut init_logits_host: Option<Vec<f32>> = None;
11708 let h_seed0: CudaSlice<f32> = if carried_pending.is_none() {
11709 let (init_logits, h) = self.spec_target_step_h(e, last_token, &mut *cache)?;
11710 last_pred = argmax(&init_logits) as u32;
11711 if constraint.is_some() {
11712 init_logits_host = Some(init_logits.clone());
11713 }
11714 // sampled mode: p-distribution after last_token, for the j==0/base==0 accept test.
11715 if sampled {
11716 last_col_logits = Some(e.htod(&init_logits)?);
11717 }
11718 h
11719 } else {
11720 // predecessor-row anchor: hidden of the last COMMITTED row (the carry contract).
11721 let lh = sess_tail
11722 .as_ref()
11723 .unwrap()
11724 .1
11725 .as_ref()
11726 .expect("pending carry requires last_h");
11727 e.clone_dtod(lh)?
11728 };
11729 let t_init = t_ent.elapsed();
11730 let mut last_col_stats: Option<(f32, f32, f32)> = None;
11731 // PERSISTENT h_seed buffer (allocated BEFORE any graph capture so no captured scratch can
11732 // alias it): every path that updates the round seed copies INTO it — no per-round allocs,
11733 // stable pointer for the graph-draft round-start copy.
11734 let mut h_seed_buf = e.clone_dtod(&h_seed0)?;
11735 // Predecessor-pairing trackers: `fill_prev` = trunk hidden AT the last COMMITTED row (the
11736 // predecessor of the next verify's col 0 — the reference's carried pending-h analogue;
11737 // also the predecessor-row hidden for the round-0 legacy-replay seed). At round 0 that
11738 // row is last_token's own (h_seed0). The chain step-0 seed under the pairing default =
11739 // hidden of the row BEFORE last_token = the prompt's last row at round 0 (h_seed_buf
11740 // overwritten below).
11741 let mut fill_prev = e.clone_dtod(&h_seed0)?;
11742 {
11743 if let Some(ph) = &prompt_h {
11744 let np = prompt.len();
11745 e.copy_view_into(
11746 &mut h_seed_buf,
11747 0,
11748 &ph.slice((np - 1) * n_embd..np * n_embd),
11749 n_embd,
11750 )?;
11751 } else if continuation {
11752 if let Some((_, lh, _, _, _)) = sess_tail.as_ref() {
11753 if let Some(lh) = lh.as_ref() {
11754 e.copy_into(&mut h_seed_buf, 0, lh, n_embd)?;
11755 }
11756 }
11757 }
11758 }
11759 // Persistent device prediction slots for the accept walk (max k+1 verify columns).
11760 let mut preds_d = e.alloc_u32_zeroed(k + 2)?;
11761
11762 let debug_spec = std::env::var("MEMRA_DEBUG_SPEC").is_ok();
11763 let fork_mode = OptiForkGateMode::configured();
11764 // MEMRA_SPEC_STATS=1: per-slot accept histogram + draft-length histogram, printed once at
11765 // the end. Metric normalization vs the reference engine: BOTH engines count
11766 // accepted/drafted where the chain stopped at p-min and the sub-threshold token is
11767 // discarded uncounted — per-slot decay + chain-length mix are the extra dimensions.
11768 let spec_stats = std::env::var("MEMRA_SPEC_STATS").is_ok();
11769 let mut st_drafted = vec![0usize; k];
11770 let mut st_accepted = vec![0usize; k];
11771 let mut st_len_hist = vec![0usize; k + 1];
11772 let mut st_full = 0usize;
11773 // P-MIN CONFIDENCE GATE (MEMRA_SPEC_PMIN, the serve script's --spec-draft-p-min mechanism):
11774 // stop the draft chain early when the head's softmax confidence in its own pick drops
11775 // below p_min. Hoisted above the loop: the graph capture bakes the prob kernels iff on.
11776 static PMIN: std::sync::OnceLock<f32> = std::sync::OnceLock::new();
11777 let p_min = *PMIN.get_or_init(|| {
11778 std::env::var("MEMRA_SPEC_PMIN")
11779 .ok()
11780 .and_then(|v| v.parse().ok())
11781 .unwrap_or(0.0)
11782 });
11783 // ZERO-DRAFT ROUNDS (MEMRA_SPEC_PMIN0=1, vendored from llama.cpp's draft gating): let the
11784 // p-min gate apply at j==0 too, so a low-confidence round drafts NOTHING and the verify
11785 // batch is just the pending bonus (m=1 = a plain decode step). llama's 35B win rides
11786 // exactly this — draft acceptance 76% at mean len 2.5 because unpredictable stretches
11787 // never pay draft+verify overhead. Only legal when a pending bonus exists (an empty
11788 // verify batch is not); the j==0 exemption stays for pending-less rounds.
11789 let pmin0 = std::env::var("MEMRA_SPEC_PMIN0")
11790 .map(|v| v == "1")
11791 .unwrap_or(false);
11792
11793 // --- GRAPH DRAFT setup: persistent I/O buffers + ONE capture (2 warmups inside). The
11794 // warmups mutate scratch len_d / pos / tok / seed — all reset at every round start, so the
11795 // only restore needed is the scratch counter. Capture failure (e.g. a non-capturable
11796 // cuBLAS path in an exotic head) falls back to the eager draft chain.
11797 // PER-SESSION PERSISTENCE (2026-08-01): session calls reuse the DraftGraphCtx parked on
11798 // the SpecSession — the capture (2 warmup head forwards + instantiate) ran ONCE at the
11799 // session's first burst, not per burst (measured ~16ms/burst fixed cost on H100 q27,
11800 // research/spec-serving-20260801). Reuse is pointer-exact: the graph bakes the session's
11801 // own scratch KV (never realloc'd), the model's resident embedding, the OnceLock p_min,
11802 // and the g_* buffers carried in the ctx — replay dispatch is identical to a fresh
11803 // capture, so draft tokens are bit-identical (drafts never decide exactness anyway; the
11804 // verify arbitrates). Single-shot calls (sess=None) build a fresh ctx and drop it.
11805 let mut dctx: DraftGraphCtx = match sess_draft_slot.as_mut().and_then(|s| s.take()) {
11806 Some(c) => c,
11807 None => DraftGraphCtx::new(e, n_embd, if sampled { d_vocab } else { 1 })?,
11808 };
11809 // A session that ran greedy bursts first sized g_q/g_perturb at 1; a sampled resume
11810 // needs d_vocab. Realloc is legal exactly while graph_s is None (nothing baked them).
11811 if sampled && dctx.g_q.len() < d_vocab {
11812 dctx.g_q = e.zeros(d_vocab)?;
11813 dctx.g_perturb = e.zeros(d_vocab)?;
11814 }
11815 // DRAFT-SIDE GRAMMAR MASK (lane/draft-mask, 2026-08-04): the drafter samples the
11816 // grammar's legal set, so proposals are legal BY CONSTRUCTION and the verify-side
11817 // truncation (the correctness backstop) stops cutting every tight-schema round.
11818 // The mask is one node inside the captured draft chain — presence is a CAPTURE-TIME
11819 // shape, so a parked graph of the other shape is dropped and recaptured.
11820 let dmask_on = constraint
11821 .as_deref()
11822 .is_some_and(|c| c.draft_mask_enabled());
11823 let dmask_words = if dmask_on { d_vocab.div_ceil(32) } else { 0 };
11824 if dmask_on && dctx.g_dmask.len() < dmask_words {
11825 dctx.g_dmask = e.alloc_u32_zeroed(dmask_words)?;
11826 dctx.graph = None; // the old capture baked the old (or no) mask pointer
11827 dctx.chain = None; // chain last-row graphs bake the same pointer
11828 dctx.failed.clear_greedy();
11829 dctx.keeper.clear();
11830 }
11831 if (dctx.graph.is_some() || dctx.chain.is_some()) && dctx.graph_masked != dmask_on {
11832 dctx.graph = None;
11833 dctx.chain = None;
11834 dctx.failed.clear_greedy();
11835 dctx.keeper.clear();
11836 }
11837 // MULTI-HEAD CHAIN mode (mtp_extra non-empty — step37's 3-head shipping shape): the
11838 // step-modulo prefix-replay chain captures PER-HEAD single-row graphs
11839 // (`DraftChainGraphs`) instead of the one self-feeding graph below; the single-head
11840 // capture arms are untouched and unreachable in this mode (the launch arms branch the
11841 // same way). This removes the historical `mtp_extra.is_empty()` capture exclusion —
11842 // and with it the silent no-attempt hole: a chain capture that FAILS now trips the
11843 // same LOUD draft-graph WARN as a single-head failure.
11844 let chain_mode = !self.mtp_extra.is_empty();
11845 if graph_draft
11846 && !sampled
11847 && chain_mode
11848 && dctx.chain.is_none()
11849 && !dctx.failed.greedy_failed()
11850 {
11851 if mtp_chain_graph_on() {
11852 // dcw door: same warmup headroom pre-arm as the single-head capture below —
11853 // every plane, because each head's capture warmups append on its OWN plane.
11854 if step35_draft_dcw_on() {
11855 scratch.ensure_dcw_headroom(e, k + 2)?;
11856 }
11857 let heads_n = self.mtp_head_count();
11858 let DraftGraphCtx {
11859 g_tok,
11860 g_pos,
11861 g_seed,
11862 g_p,
11863 g_dmask,
11864 ..
11865 } = &mut dctx;
11866 if dmask_on {
11867 e.htod_u32_into(g_dmask, &vec![u32::MAX; dmask_words])?;
11868 }
11869 let g_dmask_ro: &CudaSlice<u32> = &*g_dmask;
11870 let with_prob = p_min > 0.0;
11871 // CAPTURE-RETAIN (#68 fix): one keeper for the whole chain — every graph's
11872 // warmup transients stay pinned as long as any of them replays.
11873 let cap_res = (|| -> Result<DraftChainGraphs, Box<dyn std::error::Error>> {
11874 let mut interior = Vec::with_capacity(heads_n);
11875 let mut last = Vec::with_capacity(heads_n);
11876 let mut keeper: Vec<Box<dyn std::any::Any + Send>> = Vec::new();
11877 for hi in 0..heads_n {
11878 let head = self.mtp_head_at(hi);
11879 // interior row: KV append + carrier only (`with_head=false` — the
11880 // eager chain discards interior logits too, so this is the same
11881 // consumed-byte program minus the dead full-vocab head matmul).
11882 let (g, keep) = e.capture_graph_retained(|e| {
11883 self.mtp_head_forward_cap(
11884 e,
11885 head,
11886 g_tok,
11887 g_pos,
11888 g_seed,
11889 g_p,
11890 &mut *scratch,
11891 hi,
11892 false,
11893 false,
11894 embd_gpu.expect("graph draft requires resident embedding"),
11895 embd_qt,
11896 embd_rb,
11897 d_vocab,
11898 None,
11899 None,
11900 None,
11901 )
11902 })?;
11903 // the warmups appended rows on plane hi; rewind before the next
11904 // capture so successive warmups never outrun the pre-armed headroom.
11905 scratch.set_plane_len(e, hi, base)?;
11906 interior.push(g);
11907 keeper.extend(keep);
11908 // last row: head matmul + greedy argmax tail (+ p when the policy
11909 // reads it, + the grammar-mask node when constrained).
11910 let (g2, keep2) = e.capture_graph_retained(|e| {
11911 self.mtp_head_forward_cap(
11912 e,
11913 head,
11914 g_tok,
11915 g_pos,
11916 g_seed,
11917 g_p,
11918 &mut *scratch,
11919 hi,
11920 with_prob,
11921 true,
11922 embd_gpu.expect("graph draft requires resident embedding"),
11923 embd_qt,
11924 embd_rb,
11925 d_vocab,
11926 None,
11927 None,
11928 if dmask_on {
11929 Some((g_dmask_ro, dmask_words))
11930 } else {
11931 None
11932 },
11933 )
11934 })?;
11935 scratch.set_plane_len(e, hi, base)?;
11936 last.push(g2);
11937 keeper.extend(keep2);
11938 }
11939 Ok(DraftChainGraphs {
11940 interior,
11941 last,
11942 keeper,
11943 })
11944 })();
11945 match cap_res {
11946 Ok(cg) => {
11947 scratch.set_len(e, base)?;
11948 // POSITIVE engagement receipt (the 3a lesson: a WARN-free boot is
11949 // NOT evidence of capture — the captured state must name itself).
11950 eprintln!(
11951 "[mtp-chain-graph] captured mode=greedy heads={heads_n} \
11952 interior={heads_n} last={heads_n} masked={}",
11953 dmask_on as u8
11954 );
11955 dctx.chain = Some(cg);
11956 dctx.graph_masked = dmask_on;
11957 }
11958 Err(err) => {
11959 scratch.set_len(e, base)?;
11960 // LOUD flip (audit Q2): a dropped draft graph is a coverage loss,
11961 // never silent — now including the multi-head shipping shape.
11962 if let Some(line) = dctx.failed.mark_greedy(&err.to_string()) {
11963 eprintln!("{line}");
11964 }
11965 }
11966 }
11967 } else {
11968 // Disarmed by MEMRA_MTP_CHAIN_GRAPH=0: say so once per process — the OFF arm
11969 // must be attributable in a boot log, never inferable from silence.
11970 static NOTE: std::sync::Once = std::sync::Once::new();
11971 NOTE.call_once(|| {
11972 eprintln!(
11973 "[spec] multi-head draft-chain capture disarmed \
11974 (MEMRA_MTP_CHAIN_GRAPH=0); eager chain serves this shape"
11975 );
11976 });
11977 }
11978 }
11979 if graph_draft
11980 && !sampled
11981 && !chain_mode
11982 && dctx.graph.is_none()
11983 && !dctx.failed.greedy_failed()
11984 {
11985 // dcw door: the capture warmups append device-counter rows the capture body cannot
11986 // rebase for; pre-arm ring headroom host-side (no-op on flat planes / room-enough
11987 // rings, and the door-off path is untouched).
11988 if step35_draft_dcw_on() {
11989 scratch.ensure_dcw_headroom(e, k + 2)?;
11990 }
11991 let DraftGraphCtx {
11992 g_tok,
11993 g_pos,
11994 g_seed,
11995 g_p,
11996 g_dmask,
11997 ..
11998 } = &mut dctx;
11999 // capture-time contents: ALL-ONES (ban nothing). A replay only ever runs after the
12000 // host uploads the position's real words, so the warmups stay grammar-free.
12001 if dmask_on {
12002 e.htod_u32_into(g_dmask, &vec![u32::MAX; dmask_words])?;
12003 }
12004 let g_dmask_ro: &CudaSlice<u32> = &*g_dmask;
12005 // CAPTURE-RETAIN (#68 fix): the warmup transients' pool addresses are baked into the
12006 // captured graph; the keeper pins them for the graph's lifetime. capture_graph (non-
12007 // retained) freed them at exit — safe for one-shot generate_spec (nothing else touches
12008 // the pool between replays) but WRONG for sessions: burst-boundary prime/fill/commit
12009 // passes (and, in serve, other sessions) recycle those addresses and the replay then
12010 // clobbers live buffers — the ST serve-spec corruption (research/serve-st-20260803).
12011 let cap_res = e.capture_graph_retained(|e| {
12012 self.mtp_head_forward_cap(
12013 e,
12014 mtp,
12015 g_tok,
12016 g_pos,
12017 g_seed,
12018 g_p,
12019 &mut *scratch,
12020 0,
12021 p_min > 0.0 || fork_mode == OptiForkGateMode::Controller,
12022 true,
12023 embd_gpu.expect("graph draft requires resident embedding"),
12024 embd_qt,
12025 embd_rb,
12026 d_vocab,
12027 None,
12028 None,
12029 if dmask_on {
12030 Some((g_dmask_ro, dmask_words))
12031 } else {
12032 None
12033 },
12034 )
12035 });
12036 match cap_res {
12037 Ok((g, keep)) => {
12038 scratch.set_len(e, base)?;
12039 dctx.graph = Some(g);
12040 dctx.graph_masked = dmask_on;
12041 dctx.keeper = keep;
12042 }
12043 Err(err) => {
12044 scratch.set_len(e, base)?;
12045 // LOUD flip (audit Q2): a dropped draft graph is a coverage loss, never
12046 // silent. Once per flip — mark returns None on an already-failed ctx.
12047 if let Some(line) = dctx.failed.mark_greedy(&err.to_string()) {
12048 eprintln!("{line}");
12049 }
12050 }
12051 }
12052 }
12053 // --- SAMPLED GRAPH DRAFT setup (step 3 of the sampled-spec arc): a SECOND capture, own
12054 // graph object, built only when sampled && graph-eligible — the greedy capture above is
12055 // untouched (and skipped when sampled: its graph would never be launched). Same head
12056 // forward, but the in-graph argmax reads GUMBEL-PERTURBED logits; the Philox event
12057 // counter lives in the persistent device g_ctr (bumped in-graph, host-seeded from sctr
12058 // once per round); the raw head logits land in the persistent g_q for the host's
12059 // per-replay async D2D into the round's q slot (q_slots, K x d_vocab, allocated once).
12060 // seed/temp are capture-time constants — baked into graph_s, so a pool-resumed request
12061 // with a different (seed, temp, k) drops the parked sampled graph and recaptures.
12062 // COST OF THE FRESH-SEED SERVE DEFAULT (dogfood F4, 2026-08-04): omitting `seed` on a
12063 // serve request now draws fresh per-request entropy (it used to default to a pinned 0),
12064 // so a seed-omitting request that RESUMES a parked spec session finds an s_key baked
12065 // with the PREVIOUS request's seed and pays one recapture. Bounded, and it does not
12066 // reopen the ~16ms/burst regression the persistent ctx exists to fix: a session's seed
12067 // is fixed for its whole lifetime (worker.rs reads s.sampler.seed() per burst), so
12068 // this compare misses at most ONCE per resumed request — the first burst recaptures
12069 // and every later burst in that request replays. A client that wants the parked graph
12070 // AND reproducibility supplies an explicit `seed`, honored exactly, which keeps s_key
12071 // stable across its whole conversation.
12072 // COMPOSITION RULE (fspec x gsd merge): the in-graph chain samples from the RAW
12073 // softmax — it can hold neither per-row filter stats nor the varying penalty history.
12074 // The sampled graph therefore engages only in the PURE-TEMP regime; filters/penalties
12075 // force the eager draft (which computes stats/penalties per row).
12076 // KEY THE WHOLE REGIME, not just the baked constants (lane/graph-s-key-exactness-
12077 // 20260819). `s_key` used to be `(seed, temp, k)`; the filters and penalties were left
12078 // out, so a filtered request resuming a session that parked a PURE-TEMP graph kept it —
12079 // and the launch site never re-asked `pure_temp`. See [`SampledGraphKey`] for what that
12080 // costs (an unconditional accept of out-of-head draft tokens, i.e. an exactness bug on
12081 // the request shape the vendor-default flip makes the majority).
12082 let s_key = SampledGraphKey::new(sp_seed, sp_temp, k, sp.top_k, sp.top_p, sp.min_p, pen_on);
12083 let pure_temp = s_key.pure_temp();
12084 // The regime the sampled graph may be captured/launched in: pure-temp always;
12085 // truncation-filtered when the filtered-capture door is on (the filter runs
12086 // IN-GRAPH — lane/step37-draft-graph-serving-20260830); penalties never.
12087 let s_capturable = s_key.graph_capturable();
12088 if sampled && dctx.s_key.is_some_and(|old| old != s_key) {
12089 dctx.graph_s = None;
12090 dctx.chain_s = None;
12091 dctx.failed.clear_sampled();
12092 dctx.s_key = None;
12093 dctx.q_slots.clear();
12094 dctx.keeper_s.clear();
12095 }
12096 // FILTERED capture nodes need q slots sized d_vocab AND the stat slots; the pure-temp
12097 // body leaves g_th/g_z/g_mx untouched (they exist from ctx creation either way).
12098 if graph_draft
12099 && sampled
12100 && s_capturable
12101 && chain_mode
12102 && dctx.chain_s.is_none()
12103 && !dctx.failed.sampled_failed()
12104 {
12105 if mtp_chain_graph_on() {
12106 if step35_draft_dcw_on() {
12107 scratch.ensure_dcw_headroom(e, k + 2)?;
12108 }
12109 let heads_n = self.mtp_head_count();
12110 let filtered = s_key.filtered();
12111 let DraftGraphCtx {
12112 g_tok,
12113 g_pos,
12114 g_seed,
12115 g_p,
12116 g_ctr,
12117 g_perturb,
12118 g_q,
12119 g_rows0,
12120 g_th,
12121 g_z,
12122 g_mx,
12123 ..
12124 } = &mut dctx;
12125 let with_prob = p_min > 0.0;
12126 let cap_res = (|| -> Result<DraftChainGraphs, Box<dyn std::error::Error>> {
12127 let mut interior = Vec::with_capacity(heads_n);
12128 let mut last = Vec::with_capacity(heads_n);
12129 let mut keeper: Vec<Box<dyn std::any::Any + Send>> = Vec::new();
12130 for hi in 0..heads_n {
12131 let head = self.mtp_head_at(hi);
12132 // interior row: no head, no draw — shared shape with the greedy
12133 // chain's interior, captured per mode for keeper-lifetime hygiene.
12134 let (g, keep) = e.capture_graph_retained(|e| {
12135 self.mtp_head_forward_cap(
12136 e,
12137 head,
12138 g_tok,
12139 g_pos,
12140 g_seed,
12141 g_p,
12142 &mut *scratch,
12143 hi,
12144 false,
12145 false,
12146 embd_gpu.expect("graph draft requires resident embedding"),
12147 embd_qt,
12148 embd_rb,
12149 d_vocab,
12150 None,
12151 None,
12152 None,
12153 )
12154 })?;
12155 scratch.set_plane_len(e, hi, base)?;
12156 interior.push(g);
12157 keeper.extend(keep);
12158 // last row: head matmul + the in-graph categorical draw (filtered
12159 // nodes when the request carries filters).
12160 let (g2, keep2) = e.capture_graph_retained(|e| {
12161 self.mtp_head_forward_cap(
12162 e,
12163 head,
12164 g_tok,
12165 g_pos,
12166 g_seed,
12167 g_p,
12168 &mut *scratch,
12169 hi,
12170 with_prob,
12171 true,
12172 embd_gpu.expect("graph draft requires resident embedding"),
12173 embd_qt,
12174 embd_rb,
12175 d_vocab,
12176 Some(SampledCapArgs {
12177 ctr: &mut *g_ctr,
12178 perturb: &mut *g_perturb,
12179 q_out: &mut *g_q,
12180 seed: sp_seed,
12181 temp: sp_temp,
12182 filt: if filtered {
12183 Some(SampledCapFilter {
12184 rows0: &*g_rows0,
12185 th: &mut *g_th,
12186 z: &mut *g_z,
12187 mx: &mut *g_mx,
12188 top_k: sp.top_k,
12189 top_p: sp.top_p,
12190 min_p: sp.min_p,
12191 })
12192 } else {
12193 None
12194 },
12195 }),
12196 None,
12197 None, // constrained spec is greedy-only
12198 )
12199 })?;
12200 scratch.set_plane_len(e, hi, base)?;
12201 last.push(g2);
12202 keeper.extend(keep2);
12203 }
12204 Ok(DraftChainGraphs {
12205 interior,
12206 last,
12207 keeper,
12208 })
12209 })();
12210 match cap_res {
12211 Ok(cg) => {
12212 scratch.set_len(e, base)?;
12213 for _ in 0..k {
12214 dctx.q_slots.push(e.zeros(d_vocab)?);
12215 }
12216 eprintln!(
12217 "[mtp-chain-graph] captured mode=sampled heads={heads_n} \
12218 interior={heads_n} last={heads_n} filtered={} key={s_key:?}",
12219 s_key.filtered() as u8
12220 );
12221 dctx.chain_s = Some(cg);
12222 dctx.s_key = Some(s_key);
12223 }
12224 Err(err) => {
12225 scratch.set_len(e, base)?;
12226 if let Some(line) = dctx.failed.mark_sampled(&err.to_string()) {
12227 eprintln!("{line}");
12228 }
12229 }
12230 }
12231 } else {
12232 static NOTE_S: std::sync::Once = std::sync::Once::new();
12233 NOTE_S.call_once(|| {
12234 eprintln!(
12235 "[spec] multi-head draft-chain capture disarmed \
12236 (MEMRA_MTP_CHAIN_GRAPH=0); eager chain serves this shape"
12237 );
12238 });
12239 }
12240 }
12241 if graph_draft
12242 && sampled
12243 && s_capturable
12244 && !chain_mode
12245 && dctx.graph_s.is_none()
12246 && !dctx.failed.sampled_failed()
12247 {
12248 // dcw door: same warmup headroom pre-arm as the greedy capture above.
12249 if step35_draft_dcw_on() {
12250 scratch.ensure_dcw_headroom(e, k + 2)?;
12251 }
12252 let filtered = s_key.filtered();
12253 let DraftGraphCtx {
12254 g_tok,
12255 g_pos,
12256 g_seed,
12257 g_p,
12258 g_ctr,
12259 g_perturb,
12260 g_q,
12261 g_rows0,
12262 g_th,
12263 g_z,
12264 g_mx,
12265 ..
12266 } = &mut dctx;
12267 // CAPTURE-RETAIN (#68 fix): same keeper contract as the greedy capture above.
12268 let cap_res = e.capture_graph_retained(|e| {
12269 self.mtp_head_forward_cap(
12270 e,
12271 mtp,
12272 g_tok,
12273 g_pos,
12274 g_seed,
12275 g_p,
12276 &mut *scratch,
12277 0,
12278 p_min > 0.0,
12279 true,
12280 embd_gpu.expect("graph draft requires resident embedding"),
12281 embd_qt,
12282 embd_rb,
12283 d_vocab,
12284 Some(SampledCapArgs {
12285 ctr: &mut *g_ctr,
12286 perturb: &mut *g_perturb,
12287 q_out: &mut *g_q,
12288 seed: sp_seed,
12289 temp: sp_temp,
12290 filt: if filtered {
12291 Some(SampledCapFilter {
12292 rows0: &*g_rows0,
12293 th: &mut *g_th,
12294 z: &mut *g_z,
12295 mx: &mut *g_mx,
12296 top_k: sp.top_k,
12297 top_p: sp.top_p,
12298 min_p: sp.min_p,
12299 })
12300 } else {
12301 None
12302 },
12303 }),
12304 None,
12305 None, // constrained spec is greedy-only — sampled never carries a hook
12306 )
12307 });
12308 match cap_res {
12309 Ok((g, keep)) => {
12310 scratch.set_len(e, base)?;
12311 for _ in 0..k {
12312 dctx.q_slots.push(e.zeros(d_vocab)?);
12313 }
12314 dctx.graph_s = Some(g);
12315 dctx.s_key = Some(s_key);
12316 dctx.keeper_s = keep;
12317 }
12318 Err(err) => {
12319 scratch.set_len(e, base)?;
12320 // LOUD flip (audit Q2): same contract as the greedy capture above.
12321 if let Some(line) = dctx.failed.mark_sampled(&err.to_string()) {
12322 eprintln!("{line}");
12323 }
12324 }
12325 }
12326 }
12327 // ---- EXACTNESS GUARD, the enforceable half (lane/graph-s-key-exactness-20260819,
12328 // widened by lane/step37-draft-graph-serving-20260830) ----
12329 // With the filters and penalties in `s_key`, a graph that SURVIVED the drop above was
12330 // captured under THIS request's exact regime, and capture requires `graph_capturable`
12331 // (pure-temp, or filtered with the in-graph filter nodes; never penalties) — so a
12332 // parked graph implies both. That implication is the whole exactness argument for the
12333 // graph arm, so it is asserted here rather than assumed: a future change that widens
12334 // the capture condition, narrows the key, or copies a `DraftGraphCtx` across regimes
12335 // fails LOUDLY at this line instead of silently drafting from a distribution the
12336 // verify never reconstructs. Release builds refuse the graph (drop it, draft eager)
12337 // rather than launching it; the launch site re-tests the regime independently.
12338 if sampled
12339 && (dctx.graph_s.is_some() || dctx.chain_s.is_some())
12340 && (!s_capturable || dctx.s_key != Some(s_key))
12341 {
12342 debug_assert!(
12343 false,
12344 "sampled draft graph parked under {:?} survived into a request outside its \
12345 capture regime (top_k={} top_p={} min_p={} pen_on={} capturable={}): the \
12346 in-graph draw and the verify's accept test would see different distributions",
12347 dctx.s_key, sp.top_k, sp.top_p, sp.min_p, pen_on, s_capturable,
12348 );
12349 eprintln!(
12350 "[spec] BUG: dropping a parked sampled draft graph that outlived its capture \
12351 regime (s_key={:?}, request top_k={} top_p={} min_p={} pen_on={} \
12352 capturable={}); drafting EAGER — the key must carry every field that shapes q",
12353 dctx.s_key, sp.top_k, sp.top_p, sp.min_p, pen_on, s_capturable,
12354 );
12355 dctx.graph_s = None;
12356 dctx.chain_s = None;
12357 dctx.s_key = None;
12358 dctx.q_slots.clear();
12359 dctx.keeper_s.clear();
12360 }
12361 // SKEY PROBE (MEMRA_SKEY_PROBE=1): the burst-entry facts the reachability question turns
12362 // on — is this request sampled, is it in a regime the sampled graph is legal in, and is
12363 // a graph PARKED from an earlier request of the same session? The launch arms below
12364 // print which chain actually ran, so the probe never restates the condition.
12365 if skey_probe() {
12366 eprintln!(
12367 "[skey] burst sampled={} pure_temp={} capturable={} temp={} top_k={} top_p={} \
12368 min_p={} pen_on={} k={} graph_draft={} graph_s_parked={} chain_s_parked={} \
12369 s_key_parked={:?}",
12370 sampled as u8,
12371 pure_temp as u8,
12372 s_capturable as u8,
12373 sp_temp,
12374 sp.top_k,
12375 sp.top_p,
12376 sp.min_p,
12377 pen_on as u8,
12378 k,
12379 graph_draft as u8,
12380 dctx.graph_s.is_some() as u8,
12381 dctx.chain_s.is_some() as u8,
12382 dctx.s_key,
12383 );
12384 }
12385 let t_cap = t_ent.elapsed();
12386 // PERSISTENT DRAFT KV: fill the MTP block's K/V for every prompt position from the exact
12387 // trunk hiddens collected during prime — ONE batched K/V-only pass (overwrites any
12388 // capture-warmup garbage; capture left len at 0). last_token (the init feed) needs no
12389 // fill: the first chain step processes it and appends its entry at slot prompt.len().
12390 if let Some(ph) = &prompt_h {
12391 // SESSION: rows [0..base) are the previous turns' exact fills (refresh overwrote them
12392 // with true verify hiddens) — truncate any draft overhang, fill ONLY the suffix at
12393 // global positions [base..base+tp). Fresh call: base==0, identical to before.
12394 scratch.set_len(e, base)?;
12395 // CHUNKED FILL (long-ctx OOM fix, 2026-07-05): mtp_kv_fill's transients scale with its
12396 // T (concat = T*2*n_embd*4B — 1.5GB at 40k) and its concat loop is 2*T launches. The
12397 // fill is a pure sequential append, so chunking is exact: each chunk appends its rows
12398 // at pos0=base+start with the identical per-row math. Same knob as the trunk prime.
12399 let tp = prompt.len();
12400 let fill_chunk: usize = if crate::cache::swa_ring_on() {
12401 crate::hybrid_forward::prime_chunk_tokens(tp, self.layers.len())
12402 } else {
12403 // Preserve the flag-OFF schedule byte-for-byte, including the legacy zero value
12404 // meaning one monolithic fill.
12405 std::env::var("MEMRA_PRIME_CHUNK")
12406 .ok()
12407 .and_then(|v| v.parse().ok())
12408 .unwrap_or(4096)
12409 };
12410 let fill_chunk = if fill_chunk == 0 { tp } else { fill_chunk };
12411 let mut start = 0usize;
12412 while start < tp {
12413 let end = (start + fill_chunk).min(tp);
12414 let tc = end - start;
12415 {
12416 // PREDECESSOR pairing: row i gets h[i-1]; global row 0 a zeros row (the
12417 // reference engine's initial pending-h is zeroed too); a session turn's row 0
12418 // gets the PREVIOUS turn's last committed hidden (sess.last_h). Per chunk:
12419 // rows start..end read h[start-1..end-1] — one dtod into a chunk buffer.
12420 let mut phs = e.zeros(tc * n_embd)?;
12421 let (src_lo, dst_off) = if start == 0 {
12422 (0, n_embd)
12423 } else {
12424 ((start - 1) * n_embd, 0)
12425 };
12426 let n_copy = if start == 0 {
12427 (tc - 1) * n_embd
12428 } else {
12429 tc * n_embd
12430 };
12431 if start == 0 {
12432 if let Some((_, lh, _, _, _)) = sess_tail.as_ref() {
12433 if let Some(lh) = lh.as_ref() {
12434 e.copy_into(&mut phs, 0, lh, n_embd)?;
12435 }
12436 }
12437 }
12438 if n_copy > 0 {
12439 e.copy_view_into(
12440 &mut phs,
12441 dst_off,
12442 &ph.slice(src_lo..src_lo + n_copy),
12443 n_copy,
12444 )?;
12445 }
12446 self.mtp_kv_fill_all(
12447 e,
12448 &prompt[start..end],
12449 &phs,
12450 base + start,
12451 &mut *scratch,
12452 embd_dev,
12453 )?;
12454 }
12455 start = end;
12456 }
12457 }
12458 // MEMRA_PROFILE_SPEC=2: profiler capture starts HERE — after the prime, so an
12459 // `nsys -c cudaProfilerApi` capture contains ONLY the round loop (draft/verify/commit).
12460 // (=1 brackets the whole call in run_spec.rs, prime included.)
12461 if std::env::var("MEMRA_PROFILE_SPEC").as_deref() == Ok("2") {
12462 unsafe extern "C" {
12463 fn cudaProfilerStart() -> i32;
12464 }
12465 unsafe {
12466 cudaProfilerStart();
12467 }
12468 }
12469 // ROUND-STREAM stage (c) 4 (MEMRA_SPEC_STREAM=1, experimental): pre-issued M-round
12470 // bursts with ZERO per-round host readbacks — the accept/seed/rollback/ring kernels
12471 // consume each other's device outputs; the host drains the ring every M rounds. v1
12472 // constraints: greedy, !spec_replay, single-shot, batched-linear layers, no refresh
12473 // fills (acceptance effect A/B-arbitrated), enters from round 1 (pending guaranteed).
12474 // NOTE: not gated on the caller's graph_draft (its trunk_dense conjunct turns the 35B
12475 // MoE off) — the stream capture encloses ONLY the dense MTP head; the head-dense /
12476 // full-prec / k gates are re-derived here and a failed capture degrades to stream-off.
12477 let stream_on = crate::spec::spec_stream()
12478 && !sampled
12479 && !spec_replay
12480 && self.mtp_extra.is_empty()
12481 && constraint.is_none()
12482 && !session_mode
12483 && embd_gpu.is_some()
12484 && !crate::model::full_prec_enabled()
12485 && k + 2 < 96;
12486 let mut stream_graph: Option<cudarc::driver::CudaGraph> = None;
12487 let mut g_tokp2k = e.alloc_u32_zeroed(2 * k.max(1))?;
12488 if stream_on {
12489 let cap = e.capture_graph(|e| {
12490 for j in 0..k.max(1) {
12491 self.mtp_head_forward_cap(
12492 e,
12493 mtp,
12494 &mut dctx.g_tok,
12495 &mut dctx.g_pos,
12496 &mut dctx.g_seed,
12497 &mut dctx.g_p,
12498 &mut *scratch,
12499 0,
12500 true,
12501 true,
12502 embd_gpu.expect("round stream requires resident embedding"),
12503 embd_qt,
12504 embd_rb,
12505 d_vocab,
12506 None,
12507 Some((&mut g_tokp2k, j, d2t_dev.as_ref())),
12508 None, // round-stream requires constraint.is_none() (see stream_on)
12509 )?;
12510 }
12511 Ok(())
12512 });
12513 match cap {
12514 Ok(g) => {
12515 scratch.set_len(e, 0)?;
12516 stream_graph = Some(g);
12517 }
12518 Err(err) => {
12519 scratch.set_len(e, 0)?;
12520 if debug_spec {
12521 eprintln!("[spec] stream-graph capture failed ({err}); stream off");
12522 }
12523 }
12524 }
12525 }
12526 let stream_active = stream_on && stream_graph.is_some();
12527 if debug_spec {
12528 eprintln!(
12529 "[spec] stream_on={stream_on} env={} samp={sampled} dg={} captured={} active={stream_active} session={session_mode} replay={spec_replay}",
12530 crate::spec::spec_stream(),
12531 dctx.graph.is_some(),
12532 stream_graph.is_some()
12533 );
12534 }
12535 let t_v_s = k + 1;
12536 // ROUND-STREAM buffers + ptr tables now live in the model-generic round_stream
12537 // module (extracted 2026-07-12; the gemma burst reuses them).
12538 let sb = crate::round_stream::StreamBufs::new(e, k, crate::spec::spec_stream_m())?;
12539 let crate::round_stream::StreamBufs {
12540 mut vtok_d,
12541 mut brk_d,
12542 mut pend_d,
12543 last_pred_d,
12544 mut pos_ctr,
12545 mut pos_start_d,
12546 mut ring_d,
12547 acc_d: mut stream_acc,
12548 m_rounds,
12549 k: _,
12550 } = sb;
12551 let stream_ptrs: Option<CudaSlice<u64>> = if stream_active {
12552 Some(crate::round_stream::kv_len_ptr_table(
12553 e,
12554 cache,
12555 Some(&pos_ctr),
12556 )?)
12557 } else {
12558 None
12559 };
12560
12561 let t_fill = t_ent.elapsed();
12562 let mut round = 0usize;
12563 // ADAPTIVE DRAFT LENGTH (MEMRA_SPEC_ADAPT=1, opt-in — the gemma_spec accepted-run law,
12564 // ported 2026-08-01): next round's draft depth = last round's accepted run + 1, clamped
12565 // to [floor(pos), k_cap] — a miss shrinks the next draft to the miss point + 1,
12566 // full-accept streaks re-deepen one step per round. NOT the 2026-07-07 acceptance-EMA
12567 // (that arm measured an HONEST LOSS to static per-class optima — 115.0/85.8/73.4 vs
12568 // 121.6/92.7/75.6, EMA lag — and was removed 2026-07-08; rig5090.jsonl has the record).
12569 // The gemma law has no lag class: it reacts within one round, and was worth +7-20% on
12570 // the gemma cells at unchanged exactness (2026-07-10 flip; floor sweep 2026-07-25;
12571 // position key 2026-07-26). Signal = n_acc from the round's EXISTING accept readback —
12572 // zero new syncs; the draft graph is a SINGLE-STEP capture replayed per drafted token,
12573 // so a per-round depth needs no re-capture (unlike gemma's whole-chain graphs). qwen's
12574 // in-round p-min cut already shortens chains mid-round, so gemma's one-round-late p-min
12575 // fold into kc is unnecessary here — the accepted-run law sees the cut via n_acc.
12576 // Exactness is the verify's job at ANY depth (same contract as p-min variable rounds).
12577 // DEFAULT OFF on the qwen path until its cells gate a flip (gemma's is default-on).
12578 // MEASURED 2026-08-01 (H100 GPU-3, interleaved x3, NGEN=256, same-invocation plain
12579 // denominators; research/qwen-adaptive-k-20260801/): REFUTED on the tuned qwen configs.
12580 // q27 K=3+HPOST+PMIN=0.3: short +0.8% (noise; law ~idles, len_hist identical), board
12581 // -1.9%, agentic -0.5%; board PMIN=0 -2.1% (not p-min shadowing — the law itself);
12582 // floor=1 -2.8% (gemma's floor-collapse, reproduced). q35 K=2 board: -6.4% (52/136
12583 // rounds shrink to depth 1; no depth to reclaim at K=2). The gemma direction DOES
12584 // appear at untuned depth-K — q27 K=6 floor=4 +1.5% over fixed K=6 — but stays -3.7%
12585 // below fixed K=3: same verdict class as the retired EMA arm (honest loss to static
12586 // per-class optima). Acceptance-rate rises under the law while tokens/round falls —
12587 // it buys accept-% by adding rounds, and a round's fixed draft+verify cost wins.
12588 // K=1..8 self-consistency PASS both models with the law ON (exactness held).
12589 let adapt = std::env::var("MEMRA_SPEC_ADAPT").as_deref() == Ok("1");
12590 // floor: per-model default keyed on n_embd (gemma's tiering — models with an expensive
12591 // verify keep deep drafts after a miss); MEMRA_SPEC_ADAPT_FLOOR pins it everywhere.
12592 let adapt_floor_env: Option<usize> = std::env::var("MEMRA_SPEC_ADAPT_FLOOR")
12593 .ok()
12594 .and_then(|v| v.parse().ok());
12595 let adapt_floor_default: usize = if self.cfg.n_embd as usize >= 3500 {
12596 4
12597 } else if self.cfg.n_embd as usize >= 2500 {
12598 2
12599 } else {
12600 1
12601 };
12602 let adapt_floor: usize = adapt_floor_env.unwrap_or(adapt_floor_default);
12603 // position key: past floor_ctx a HIGH floor (>=4) relaxes to 1 — forced-deep drafts
12604 // turn net-negative at depth (gemma 31B d1736 evidence); MEMRA_SPEC_FLOOR_CTX moves
12605 // the boundary, an explicit MEMRA_SPEC_ADAPT_FLOOR pins the floor everywhere.
12606 let floor_ctx: usize = std::env::var("MEMRA_SPEC_FLOOR_CTX")
12607 .ok()
12608 .and_then(|v| v.parse().ok())
12609 .unwrap_or(1024);
12610 let floor_at = |pos: usize| -> usize {
12611 if adapt_floor_env.is_some() || pos < floor_ctx {
12612 adapt_floor
12613 } else if adapt_floor >= 4 {
12614 1
12615 } else {
12616 adapt_floor
12617 }
12618 };
12619 // cap: MEMRA_SPEC_CAPMAX (gemma semantics, default 7). Binds only under adapt — the
12620 // fixed-K default path is untouched by this whole block.
12621 let cap_max: usize = std::env::var("MEMRA_SPEC_CAPMAX")
12622 .ok()
12623 .and_then(|v| v.parse().ok())
12624 .unwrap_or(7);
12625 let k_cap = k.min(cap_max).max(1);
12626 let mut kc = k_cap;
12627 let mut opti_fork: Option<OptiForkState> = None;
12628 let mut fork_snapshot: Option<crate::cache::CacheSnapshot> = None;
12629 if fork_mode != OptiForkGateMode::Disabled {
12630 let fence = crate::pp::pp_cuts(self.layers.len());
12631 let refusal = if !session_mode {
12632 Some("not-session")
12633 } else if k != 1 || adapt {
12634 Some("requires-fixed-k1")
12635 } else if sampled || constraint.is_some() || spec_replay {
12636 Some("sampled-constrained-or-replay")
12637 } else if pipe.is_some() {
12638 Some("two-session-pipeline")
12639 } else if !spec_devacc() {
12640 Some("requires-device-accept")
12641 } else if stream_active || crate::spec::spec_stream() {
12642 Some("round-stream")
12643 } else if !self.mtp_extra.is_empty() {
12644 Some("multi-head-mtp")
12645 } else if crate::cache::swa_ring_on() || cache.has_swa_ring() {
12646 Some("swa-ring")
12647 } else if crate::pp::pp_host_bounce_active() {
12648 Some("host-bounce")
12649 } else if fork_mode == OptiForkGateMode::Controller
12650 && cache.recur.iter().any(Option::is_some)
12651 {
12652 Some("controller-requires-zero-recurrent-state")
12653 } else if fence.as_ref().is_none_or(|f| f.len() != 3) {
12654 Some("requires-pp2")
12655 } else {
12656 None
12657 };
12658 if let Some(reason) = refusal {
12659 OPTI_FORK_REFUSALS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
12660 eprintln!("[opti-fork] refused reason={reason}");
12661 } else {
12662 let fence = fence.expect("validated PP-2 fence");
12663 let rt = crate::pp::PpNRt::get(e)?;
12664 let primary_stage0 = rt.engine(0, e).ctx().ordinal() == e.ctx().ordinal();
12665 let primary_stage1 = rt.engine(1, e).ctx().ordinal() == e.ctx().ordinal();
12666 let primary_supported =
12667 primary_stage0 || (fork_mode == OptiForkGateMode::Controller && primary_stage1);
12668 if !rt.cross_device() || !primary_supported {
12669 OPTI_FORK_REFUSALS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
12670 eprintln!("[opti-fork] refused reason=requires-supported-primary-cross-device");
12671 } else {
12672 // Both recurrent snapshots and both seed generations are allocated before
12673 // the first fork, each through its owning PP stage. Allocation failure
12674 // therefore happens before any optimistic state mutation can occur.
12675 let current_snapshot = opti_snapshot_stage_owned(e, cache, rt, &fence)?;
12676 let alternate_snapshot = opti_snapshot_stage_owned(e, cache, rt, &fence)?;
12677 let fork = OptiForkState::new(
12678 e,
12679 cache,
12680 fork_mode,
12681 alternate_snapshot,
12682 &h_seed_buf,
12683 &fill_prev,
12684 rt,
12685 fence[1],
12686 self.layers.len(),
12687 )?;
12688 eprintln!(
12689 "[opti-fork] armed mode={fork_mode:?} snapshots=2 seeds=2 split={} \
12690 payload_dev0={} payload_dev1={} q_threshold={:.3}",
12691 fence[1],
12692 fork.logical_payload_bytes[0],
12693 fork.logical_payload_bytes[1],
12694 fork.controller.map_or(0.0, |policy| policy.threshold),
12695 );
12696 fork_snapshot = Some(current_snapshot);
12697 opti_fork = Some(fork);
12698 }
12699 }
12700 }
12701 // Persistent snapshot buffers are allocated once and refreshed in place. The fork arm
12702 // uses stage-owned snapshots; refused/disabled arms retain the existing generic helper.
12703 let mut snap = match fork_snapshot {
12704 Some(snapshot) => snapshot,
12705 None => cache.snapshot(e)?,
12706 };
12707 let mut carried_opti: Option<OptiControllerTicket> = None;
12708 // ROUND-STREAM stage (b) 3a: device table of per-layer kvl.len_d pointers (stable — the
12709 // cache never reallocates len_d; see cache.rs "stable pointer" note). 0 = no KV layer.
12710 let kv_len_ptrs: Option<CudaSlice<u64>> = if spec_devacc() && !spec_replay {
12711 Some(crate::round_stream::kv_len_ptr_table(e, cache, None)?)
12712 } else {
12713 None
12714 };
12715 // BONUS FOLD (2026-07-04): after a FULL accept the bonus token is NOT committed with a
12716 // separate T=1 trunk pass (a full weight read per round). It stays PENDING and rides as
12717 // column 0 of the NEXT round's verify batch. Under predecessor pairing the next chain
12718 // seeds from the bonus's predecessor's TRUE verify hidden (free — no extra
12719 // pass of any kind). Verify still
12720 // checks every emitted token against the target -> exactness holds by construction; only
12721 // DRAFT QUALITY can shift, which the acceptance numbers arbitrate.
12722 // bonus emitted but not yet committed to cache. A carried pending (see SpecSession::
12723 // pending_tok) enters round 0 directly — the burst boundary becomes a plain round edge.
12724 let mut pending: Option<u32> = carried_pending;
12725 // MEMRA_SPEC_PHASE=1: per-round wall decomposition (draft / verify / accept+commit) —
12726 // no tracing, no extra syncs (each phase is naturally sync-bounded: draft readbacks,
12727 // the verify accept readback). Printed once at loop end via spec-stats.
12728 let anatomy_on = std::env::var("MEMRA_SPEC_PP_ANATOMY").as_deref() == Ok("1");
12729 let phase_on = anatomy_on || std::env::var("MEMRA_SPEC_PHASE").as_deref() == Ok("1");
12730 // MEMRA_SPEC_PHASE_SYNC=1 — reads the phase split correctly, and proves it. `ph_mark` is a
12731 // bare Instant, so `verify-issue` is the host QUEUEING the walk (the GPU is already running
12732 // under it) and `verify-wait` is only the residual drain at the accept readback: one
12733 // overlapped interval cut at the first blocking call, NOT "GPU time" beside "host time".
12734 // Syncing right after the walk is issued moves the whole GPU wall into `verify-issue`. If
12735 // the walk's GPU total is really issue+wait, then with this on verify-issue jumps to that
12736 // sum, verify-wait collapses to the readback alone, and the ROUND WALL DOES NOT MOVE —
12737 // which is what says the queueing time was hidden and is not a target. Diagnostic only.
12738 let phase_sync = std::env::var("MEMRA_SPEC_PHASE_SYNC").as_deref() == Ok("1");
12739 // DRAFT-MASK receipt (lane/draft-mask): speculative-clone wall + rounds, printed with
12740 // spec-stats. The clone is the one cost the design adds per round — measured, not assumed.
12741 let (mut dm_clone_ns, mut dm_rounds) = (0u128, 0usize);
12742 // grammar-truncation counters: how many rounds the verify-side cut fired and how many
12743 // already-verified tokens it threw away. THIS is the quantity draft masking targets.
12744 let (mut dm_cuts, mut dm_cut_tokens) = (0usize, 0usize);
12745 let (mut ph_draft, mut ph_verify, mut ph_rest) = (0f64, 0f64, 0f64);
12746 let mut ph_wait = 0f64;
12747 let mut ph_commit = 0f64;
12748 let mut ph_t = std::time::Instant::now();
12749 let mut ph_mark = |acc: &mut f64, on: bool| {
12750 if on {
12751 let now = std::time::Instant::now();
12752 *acc += (now - ph_t).as_secs_f64();
12753 ph_t = now;
12754 }
12755 };
12756 // MTP-ROUTE VERIFY GRAPHS (`MEMRA_SPEC_VERIFY_GRAPH`, see the flag doc): the
12757 // model-owned capture pool, locked for the whole burst exactly as the dspark serve
12758 // arm holds it — the slab stash is live verify -> commit inside a round, and the
12759 // worker drives rounds from one scheduler thread. PERSISTENT across generations on
12760 // the model (rebuilding per call re-captures the pool per prompt, which is the
12761 // measured way to lose more than the launches cost); the captured bodies are
12762 // cache-independent, every state read going through per-round refreshed pointer
12763 // tables. None = the eager walk, byte-identical.
12764 //
12765 // Never armed together with ROUND-STREAM: the tparallel verify refuses that pair
12766 // loudly, and `stream_active` owns the burst arm above, so the door stays shut
12767 // whenever the stream is live rather than relying on that refusal.
12768 // The lock is taken ONLY when the door is armed: with the flag off this whole block
12769 // is inert, so the default path cannot serialize two spec generations behind a mutex
12770 // it never reads.
12771 let vg_armed =
12772 crate::spec::spec_verify_graph_env().unwrap_or_else(|| self.vgraph_family_default());
12773 let mut vg_guard = if vg_armed && !stream_active {
12774 let mut g = self.dspark_vgraphs.lock().unwrap();
12775 if g.is_none() {
12776 // Size by the WIDEST verify this run can present, which is k+1 and NOT
12777 // k_cap+1: the sampled arm's own window is `t_v_s = k + 1`, so a pool built
12778 // from a smaller adaptive cap gets sliced past its stash rows (a `slice_mut`
12779 // panic in the sampled ON arm, measured before this line said k+1).
12780 let vt_cap = (k.max(k_cap) + 1).max(2);
12781 *g = DsparkVerifyGraphs::new(e, cache, vt_cap, n_embd)?;
12782 if g.is_some() {
12783 // Engagement receipt (the dead-arm lesson): prove the door is LIVE rather
12784 // than trusting that a flag set means a pool built.
12785 eprintln!("[spec-vg] MTP verify-graph pool ENGAGED (vt_cap={vt_cap})");
12786 } else {
12787 eprintln!(
12788 "[spec-vg] MTP verify-graph pool declined (no linear layers, \
12789 non-uniform state, or vt_cap < 2) — eager walk"
12790 );
12791 }
12792 }
12793 Some(g)
12794 } else {
12795 None
12796 };
12797 // Capacity fail-safe: a round wider than the pool was built for must take the eager
12798 // walk, not slice the stash past its rows. The sizing above already covers every
12799 // round this run can present; this keeps a future caller (or a k that grows behind
12800 // the pool's back) on the byte-identical fallback instead of a panic.
12801 let vg_t_cap = vg_guard
12802 .as_ref()
12803 .and_then(|g| g.as_ref())
12804 .map(|g| g.t_capacity())
12805 .unwrap_or(0);
12806 if let Some(p) = pipe {
12807 p.setup_end();
12808 }
12809 while keep_going && out.len() < max_new {
12810 // MEMRA_SPEC_ROUND_PROF=1: wall of the WHOLE round against the pieces we already
12811 // instrument. Needed because the parts do not add up: the draft step measures 1.27 ms
12812 // ([spec-anatomy] glue 92 / attn 280 / ffn 222 / head 670 us) and the t=2 verify walk
12813 // 25.6 ms ([tcol-prof] attn 10.1 + ffn 15.3), yet a K=1 round takes 177 ms on the
12814 // step37 TP2 stack. This prints where the other ~150 ms lives.
12815 let round_prof = ROUND_PROF
12816 .get_or_init(|| std::env::var("MEMRA_SPEC_ROUND_PROF").as_deref() == Ok("1"));
12817 let round_t0 = round_prof.then(std::time::Instant::now);
12818 // ROUND-STREAM BURST: from round 1 (pending guaranteed by every non-replay arm),
12819 // issue M rounds with zero readbacks, then drain the ring + reconcile mirrors.
12820 if let (true, Some(sg), Some(ptrs)) = (
12821 stream_active && round >= 1 && pending.is_some(),
12822 &stream_graph,
12823 &stream_ptrs,
12824 ) {
12825 if debug_spec {
12826 static ONCE: std::sync::Once = std::sync::Once::new();
12827 ONCE.call_once(|| {
12828 eprintln!("[memra] ROUND-STREAM burst engaged (M={m_rounds} k={k})")
12829 });
12830 }
12831 e.set_i32_one(&mut pos_ctr, cache.pos as i32)?;
12832 e.set_u32_one(&mut pend_d, pending.unwrap())?;
12833 e.set_u32_one(&mut ring_d, 0)?; // ring count = 0 (writes element 0)
12834 for _mi in 0..m_rounds {
12835 e.i32_copy_add(&pos_ctr, &mut pos_start_d, 0)?;
12836 cache.snapshot_into(e, &mut snap)?; // device D2Ds, stream-ordered
12837 e.i32_copy_add(&pos_ctr, &mut scratch.kv.len_d, 0)?; // draft-KV rollback
12838 e.i32_copy_add(&pos_ctr, &mut dctx.g_pos, 1)?; // rope pos = pos + base
12839 e.u32_copy(&pend_d, &mut dctx.g_tok)?;
12840 e.copy_into(&mut dctx.g_seed, 0, &h_seed_buf, n_embd)?;
12841 sg.launch()?;
12842 e.spec_assemble_verify(
12843 &g_tokp2k,
12844 &pend_d,
12845 d2t_dev.as_ref(),
12846 &mut vtok_d,
12847 &mut brk_d,
12848 p_min,
12849 k,
12850 pmin0,
12851 )?;
12852 let mut ck = VerifyCkpt::new(self.layers.len());
12853 let dummy = vec![0u32; t_v_s];
12854 let (tl_d, vx) = self.decode_step_t_core_stream(
12855 e,
12856 &dummy,
12857 0,
12858 &mut *cache,
12859 embd_dev,
12860 Some(&mut ck),
12861 Some((&vtok_d, &pos_ctr)),
12862 None,
12863 None,
12864 None,
12865 )?;
12866 for j in 0..t_v_s {
12867 e.argmax_token_device_col(&tl_d, j, n_vocab, &mut preds_d, j)?;
12868 }
12869 e.spec_accept_greedy_dc(
12870 &preds_d,
12871 &vtok_d,
12872 &last_pred_d,
12873 &brk_d,
12874 &mut stream_acc,
12875 )?;
12876 e.spec_seed_gather(&vx, &fill_prev, &stream_acc, &mut h_seed_buf, 1, n_embd)?;
12877 e.copy_into(&mut fill_prev, 0, &h_seed_buf, n_embd)?;
12878 self.commit_verified_prefix_stream(
12879 e,
12880 &mut *cache,
12881 &snap,
12882 &ck,
12883 &stream_acc,
12884 1,
12885 t_v_s,
12886 )?;
12887 e.spec_rollback_stream(
12888 ptrs,
12889 &pos_start_d,
12890 &stream_acc,
12891 1,
12892 self.layers.len() + 1,
12893 )?;
12894 e.spec_ring_commit(&vtok_d, &stream_acc, &brk_d, &mut ring_d, &mut pend_d)?;
12895 }
12896 e.stream().synchronize()?;
12897 let ring_h = e.dtoh_u32(&ring_d)?;
12898 let cnt = ring_h[0] as usize;
12899 for i in 0..cnt {
12900 if out.len() < max_new {
12901 out.push(ring_h[1 + i]);
12902 }
12903 }
12904 let pos_h = e.dtoh_i32(&pos_ctr)?[0] as usize;
12905 for il in 0..self.layers.len() {
12906 if let Some(kvl) = cache.kv[il].as_mut() {
12907 kvl.len = pos_h;
12908 }
12909 }
12910 cache.pos = pos_h;
12911 scratch.kv.len = pos_h;
12912 pending = Some(ring_h[cnt]); // last drained token = the live bonus
12913 last_token = ring_h[cnt];
12914 total_drafted += k * m_rounds; // upper bound (p-min breaks uncounted)
12915 total_accepted += cnt.saturating_sub(m_rounds);
12916 if let Some(t) = sess_telem {
12917 // totals only — the burst's per-round accept counts stayed on device
12918 // (that is the point of the round-stream arm). pos_* untouched.
12919 t.record_totals(m_rounds, k * m_rounds, cnt.saturating_sub(m_rounds));
12920 }
12921 round += m_rounds;
12922 // sse-cadence: the drained ring is committed — flush it at burst-drain cadence.
12923 keep_going = flush_commit(&mut on_commit, &out, &mut flushed);
12924 continue;
12925 }
12926 let pipe_draft = match pipe {
12927 Some(p) => Some(p.draft_begin(round)?),
12928 None => None,
12929 };
12930 let pos = cache.pos; // #tokens committed (EXCLUDES a pending bonus)
12931 let mut current_opti = carried_opti.take();
12932 let mut fork_generation = if current_opti.is_none() && pending.is_some() {
12933 match opti_fork.as_mut() {
12934 Some(fork) if fork.mode.is_forced() => Some(fork.reserve(&mut snap)?),
12935 None => None,
12936 Some(_) => None,
12937 }
12938 } else {
12939 None
12940 };
12941 if current_opti.is_none() {
12942 if let Some(fork) = opti_fork.as_ref() {
12943 opti_snapshot_stage_owned_into(e, cache, fork.rt, &fork.fence, &mut snap)?;
12944 } else {
12945 cache.snapshot_into(e, &mut snap)?;
12946 }
12947 } else if snap.pos != pos {
12948 return Err(format!(
12949 "optipipe carried snapshot pos {} != current pos {pos}",
12950 snap.pos
12951 )
12952 .into());
12953 } // §C: snapshot BEFORE draft+verify (already retained for a carried successor)
12954 ph_mark(&mut ph_rest, phase_on);
12955
12956 // --- 1. DRAFT k tokens with the NextN head (autoregressive, T=1 each) ---
12957 // p-min semantics (both paths): stop the chain early when the head's confidence in
12958 // its own pick drops below p_min — the just-drafted token is DISCARDED, but its
12959 // scratch append stands (identical to the eager chain's ordering). j==0 always drafts.
12960 let base0 = if pending.is_some() { 1usize } else { 0usize };
12961 // fixed draft length by default; MEMRA_SPEC_ADAPT=1 drafts at last round's
12962 // accepted run + 1 (the gemma law — see the setup block above the loop).
12963 let k_this = if adapt { kc } else { k };
12964 let mut draft: Vec<u32> = Vec::with_capacity(k);
12965 let mut draft_idx: Vec<u32> = Vec::with_capacity(k); // trimmed-vocab ids (== draft when untrimmed)
12966 let mut controller_draft_prob: Option<f32> = None;
12967 let mut controller_eager_state: Option<(u32, CudaSlice<f32>)> = None;
12968 if let Some(ticket) = current_opti.as_mut() {
12969 let carried_pending = pending.ok_or("optipipe carried successor lost pending")?;
12970 if ticket.verify_tokens[0] != carried_pending {
12971 return Err(format!(
12972 "optipipe carried pending mismatch: ticket={} live={carried_pending}",
12973 ticket.verify_tokens[0],
12974 )
12975 .into());
12976 }
12977 draft.push(ticket.verify_tokens[1]);
12978 controller_draft_prob = Some(ticket.draft_prob);
12979 controller_eager_state = ticket
12980 .take_eager_seed()
12981 .map(|seed| (ticket.verify_tokens[1], seed));
12982 } else {
12983 // Round-start draft-KV sync (BOTH paths). Persistent: truncate/align to the committed
12984 // history — slots 0..P hold entries for the tokens before last_token@P (P = pos +
12985 // base0 - 1); this single set_len IS the draft-side rollback (drops last round's
12986 // rejected drafts and p-min extras via the len mechanism).
12987 scratch.set_len(e, pos + base0 - 1)?;
12988 // dcw door: a captured chain appends k_this device-counter rows (plus the
12989 // pseudo-seed replay) with no host intervention; any ring rebase those appends
12990 // could need happens HERE, host-side, before the replays. The eager arm keeps
12991 // its own per-step prepare, so this is graph-path-only work.
12992 if step35_draft_dcw_on()
12993 && (dctx.graph.is_some()
12994 || dctx.graph_s.is_some()
12995 || dctx.chain.is_some()
12996 || dctx.chain_s.is_some())
12997 {
12998 scratch.ensure_dcw_headroom(e, k_this + 2)?;
12999 }
13000 if pen_on {
13001 // PEN_WINDOW_MAX also bounds the per-round upload and the O(n_hist^2)
13002 // device dedup: the serve window is already PEN_WINDOW_MAX, and this
13003 // defensive min also bounds non-server callers.
13004 let win = sp.penalty_last_n.min(PEN_WINDOW_MAX);
13005 let w0 = pen_hist.len().saturating_sub(win);
13006 pen_hist_d = Some(e.htod_u32_v(&pen_hist[w0..])?);
13007 }
13008 if sampled {
13009 draft_logits.clear();
13010 draft_stats.clear();
13011 }
13012 // DRAFT-SIDE GRAMMAR MASK: clone the committed grammar state ONCE per round; each
13013 // position's mask is computed on that clone and advanced by the PROPOSED token. The
13014 // real state moves only on emission (verify's job), so the emitted stream is
13015 // unchanged — the mask only removes tokens the verify would have truncated anyway.
13016 let mut dmask_live = dmask_on;
13017 if dmask_live {
13018 let t_c = std::time::Instant::now();
13019 constraint
13020 .as_deref_mut()
13021 .unwrap()
13022 .draft_begin()
13023 .map_err(|e2| format!("constraint: {e2}"))?;
13024 dm_clone_ns += t_c.elapsed().as_nanos();
13025 dm_rounds += 1;
13026 }
13027 if let (false, Some(cg)) = (sampled || pen_on, &dctx.chain) {
13028 // GREEDY CHAIN GRAPH (lane/step37-draft-graph-serving-20260830): the
13029 // eager multi-head chain's EXACT launch order — step j rewinds head
13030 // (j % heads)'s plane to the committed length and replays rows 0..=j —
13031 // with each row's whole head-forward as ONE graph launch. The chain
13032 // POLICY (head choice, prefix length, stored-seed feed) is host-side,
13033 // identical to `mtp_chain_forward_dev`, so graph-vs-eager drafts are
13034 // bit-identical by construction (same launcher, same bucket — the dcw
13035 // parity contract). Interior rows launch the head-less graph: their
13036 // logits are dead in the eager chain too, so the consumed bytes match.
13037 let heads_n = self.mtp_head_count();
13038 let committed = pos + base0 - 1;
13039 let mut chain_tokens: Vec<u32> = vec![last_token];
13040 let mut chain_seed_bufs: Vec<CudaSlice<f32>> = vec![e.clone_dtod(&h_seed_buf)?];
13041 for j in 0..k_this {
13042 let index = mtp_chain_head_index(j, heads_n);
13043 if debug_spec {
13044 eprintln!(
13045 "[mtp-chain-step] round={round} j={j} head={index} \
13046 replay_rows={} arm=graph",
13047 chain_tokens.len(),
13048 );
13049 }
13050 scratch.set_plane_len(e, index, committed)?;
13051 e.set_i32_one(&mut dctx.g_pos, (committed + 1) as i32)?;
13052 for row in 0..=j {
13053 e.set_u32_one(&mut dctx.g_tok, chain_tokens[row])?;
13054 e.copy_into(&mut dctx.g_seed, 0, &chain_seed_bufs[row], n_embd)?;
13055 if row < j {
13056 cg.interior[index].launch()?;
13057 } else {
13058 // per-position mask upload before the LAST row only — the
13059 // eager chain applies the mask on is_last exactly the same.
13060 if dmask_live
13061 && !upload_draft_mask(
13062 e,
13063 constraint.as_deref_mut().unwrap(),
13064 &mut dctx.g_dmask,
13065 mtp.d2t.as_ref(),
13066 d_vocab,
13067 dmask_words,
13068 )?
13069 {
13070 e.htod_u32_into(
13071 &mut dctx.g_dmask,
13072 &vec![u32::MAX; dmask_words],
13073 )?;
13074 dmask_live = false;
13075 }
13076 cg.last[index].launch()?;
13077 }
13078 // host mirror (len_d advanced in-graph by the dcw append)
13079 scratch.plane_mut(index).0.len += 1;
13080 }
13081 let idx = e.dtoh_u32_one(&dctx.g_tok)?;
13082 // #87 SENTINEL TRAP (see the single-head graph arm below).
13083 if (idx as usize) >= d_vocab {
13084 let seed_h = e.dtoh(&dctx.g_seed)?;
13085 let seed_nan = seed_h.iter().filter(|v| v.is_nan()).count();
13086 return Err(format!(
13087 "draft(chain-graph) argmax sentinel 0x{idx:08x} >= d_vocab \
13088 {d_vocab} at round {round} j={j} head={index} pos={pos}: \
13089 head-out NaN {seed_nan}/{n_embd} — refusing to dereference \
13090 the embed row (#87 trap)"
13091 )
13092 .into());
13093 }
13094 // multi-head MTP forbids a trimmed head (validated at entry), so the
13095 // draft index IS the target id; keep the map for uniformity.
13096 let d = match &mtp.d2t {
13097 Some(map) => map[idx as usize],
13098 None => idx,
13099 };
13100 let draft_p = if p_min > 0.0 {
13101 Some(e.dtoh(&dctx.g_p)?[0])
13102 } else {
13103 None
13104 };
13105 if j == 0 {
13106 controller_draft_prob = draft_p;
13107 }
13108 if let Some(p) = draft_p.filter(|_| p_min > 0.0) {
13109 if p < p_min && (j > 0 || (pmin0 && base0 == 1)) {
13110 break;
13111 }
13112 }
13113 draft.push(d);
13114 chain_tokens.push(d);
13115 // step j's h_nextn: the last-row graph self-fed it into g_seed —
13116 // snapshot it as the chain history seed for row j+1 (stream-ordered
13117 // after the launch, exactly the eager chain's chain_seeds push).
13118 chain_seed_bufs.push(e.clone_dtod(&dctx.g_seed)?);
13119 // speculative grammar advance (see the single-head graph arm).
13120 if dmask_live
13121 && !constraint
13122 .as_deref_mut()
13123 .unwrap()
13124 .draft_advance(d)
13125 .map_err(|e2| format!("constraint: {e2}"))?
13126 {
13127 e.htod_u32_into(&mut dctx.g_dmask, &vec![u32::MAX; dmask_words])?;
13128 break;
13129 }
13130 }
13131 } else if let (true, Some(cg)) = (
13132 sampled && s_capturable && dctx.s_key == Some(s_key),
13133 &dctx.chain_s,
13134 ) {
13135 if skey_probe() {
13136 eprintln!(
13137 "[skey] chain=graph_chain_s round={round} capturable={} top_k={} \
13138 top_p={} min_p={} s_key_parked={:?}",
13139 s_capturable as u8, sp.top_k, sp.top_p, sp.min_p, dctx.s_key,
13140 );
13141 }
13142 // SAMPLED CHAIN GRAPH: the greedy chain arm's launch order with the
13143 // sampled last-row graphs — in-graph counter bump + (filtered) gumbel
13144 // draw + argmax; q retained per step into q_slots exactly like the
13145 // single-head sampled graph arm. Counter continuity: g_ctr host-seeded
13146 // to sctr-1 once per ROUND; each step's last-row graph bumps it BEFORE
13147 // the perturb, so step j consumes counter sctr+j — the eager Philox
13148 // stream (interior rows never draw, never bump).
13149 let heads_n = self.mtp_head_count();
13150 let committed = pos + base0 - 1;
13151 let filtered_stats_in_graph = s_key.filtered();
13152 let mut chain_tokens: Vec<u32> = vec![last_token];
13153 let mut chain_seed_bufs: Vec<CudaSlice<f32>> = vec![e.clone_dtod(&h_seed_buf)?];
13154 e.set_u32_one(&mut dctx.g_ctr, sctr.wrapping_sub(1))?;
13155 for j in 0..k_this {
13156 let index = mtp_chain_head_index(j, heads_n);
13157 if debug_spec {
13158 eprintln!(
13159 "[mtp-chain-step] round={round} j={j} head={index} \
13160 replay_rows={} arm=graph_s",
13161 chain_tokens.len(),
13162 );
13163 }
13164 scratch.set_plane_len(e, index, committed)?;
13165 e.set_i32_one(&mut dctx.g_pos, (committed + 1) as i32)?;
13166 for row in 0..=j {
13167 e.set_u32_one(&mut dctx.g_tok, chain_tokens[row])?;
13168 e.copy_into(&mut dctx.g_seed, 0, &chain_seed_bufs[row], n_embd)?;
13169 if row < j {
13170 cg.interior[index].launch()?;
13171 } else {
13172 cg.last[index].launch()?;
13173 }
13174 scratch.plane_mut(index).0.len += 1;
13175 }
13176 sctr += 1; // mirrors the in-graph g_ctr bump (eager parity:
13177 // counts the p-min-discarded token too)
13178 // q retention: ONE async D2D of the persistent head-logits buffer
13179 // into this round's slot j (stream-ordered after the replay).
13180 e.copy_into(&mut dctx.q_slots[j], 0, &dctx.g_q, d_vocab)?;
13181 // FILTERED capture: read the in-graph filter_stats scalars back per
13182 // replay instead of a second full-vocab filter_stats per slot post-
13183 // chain — bit-exact (the values the in-graph perturb consumed) and
13184 // measured worth ~5% of vendor-default serving tok/s at K=3. Before
13185 // the p-min break so the discarded slot's stats land too.
13186 if filtered_stats_in_graph {
13187 draft_stats.push((
13188 e.dtoh(&dctx.g_mx)?[0],
13189 e.dtoh(&dctx.g_th)?[0],
13190 e.dtoh(&dctx.g_z)?[0],
13191 ));
13192 }
13193 let idx = e.dtoh_u32_one(&dctx.g_tok)?;
13194 // #87 SENTINEL TRAP (see the single-head graph arms).
13195 if (idx as usize) >= d_vocab {
13196 let seed_h = e.dtoh(&dctx.g_seed)?;
13197 let seed_nan = seed_h.iter().filter(|v| v.is_nan()).count();
13198 return Err(format!(
13199 "draft(chain-graph-sampled) argmax sentinel 0x{idx:08x} >= \
13200 d_vocab {d_vocab} at round {round} j={j} head={index} pos={pos}: \
13201 head-out NaN {seed_nan}/{n_embd} — refusing to dereference the \
13202 embed row (#87 trap)"
13203 )
13204 .into());
13205 }
13206 let d = match &mtp.d2t {
13207 Some(map) => map[idx as usize],
13208 None => idx,
13209 };
13210 draft_idx.push(idx);
13211 if p_min > 0.0 {
13212 let p = e.dtoh(&dctx.g_p)?[0];
13213 if p < p_min && (j > 0 || (pmin0 && base0 == 1)) {
13214 break;
13215 }
13216 }
13217 draft.push(d);
13218 chain_tokens.push(d);
13219 chain_seed_bufs.push(e.clone_dtod(&dctx.g_seed)?);
13220 }
13221 // PURE-TEMP accept path: stats per used slot recomputed from the RETAINED
13222 // q with the SAME filter_stats program the eager arm runs (deployment-
13223 // keyed coop/plain choice, same input bits). The FILTERED graph read its
13224 // stats back per replay above.
13225 if !filtered_stats_in_graph {
13226 for j in 0..draft.len().max(draft_idx.len()) {
13227 let rows0 = e.htod_i32(&[0])?;
13228 let (mut th_d, mut z_d, mut mx_d) =
13229 (e.zeros(1)?, e.zeros(1)?, e.zeros(1)?);
13230 e.filter_stats(
13231 &dctx.q_slots[j],
13232 d_vocab,
13233 &rows0,
13234 &mut th_d,
13235 &mut z_d,
13236 &mut mx_d,
13237 d_vocab,
13238 1,
13239 sp_temp,
13240 sp.top_k,
13241 sp.top_p,
13242 sp.min_p,
13243 )?;
13244 draft_stats.push((
13245 e.dtoh(&mx_d)?[0],
13246 e.dtoh(&th_d)?[0],
13247 e.dtoh(&z_d)?[0],
13248 ));
13249 }
13250 }
13251 } else if let (false, Some(gr)) = (sampled || pen_on, &dctx.graph) {
13252 // GRAPH DRAFT: one dispatch per drafted token. The chain feeds itself on-device
13253 // (in-graph argmax -> tok_d -> next replay's embed; h_nextn -> h_seed_d; pos_d
13254 // inc'd in-graph); the host only reads 4B token (+4B p) and decides the break.
13255 e.set_i32_one(&mut dctx.g_pos, (pos + base0) as i32)?;
13256 e.set_u32_one(&mut dctx.g_tok, last_token)?;
13257 e.copy_into(&mut dctx.g_seed, 0, &h_seed_buf, n_embd)?;
13258 for j in 0..k_this {
13259 // per-position mask upload (contents only — the graph's baked pointer is
13260 // dctx.g_dmask). All-ones once masking goes dead mid-chain, so the captured
13261 // mask node degrades to a no-op ban instead of needing a second graph.
13262 if dmask_live
13263 && !upload_draft_mask(
13264 e,
13265 constraint.as_deref_mut().unwrap(),
13266 &mut dctx.g_dmask,
13267 mtp.d2t.as_ref(),
13268 d_vocab,
13269 dmask_words,
13270 )?
13271 {
13272 // no draft-vocab row is grammar-legal here (a trimmed FR-Spec head can
13273 // genuinely miss the legal set): neutralize the captured mask node and
13274 // finish the chain UNMASKED — exactly pre-lane behaviour, never worse.
13275 e.htod_u32_into(&mut dctx.g_dmask, &vec![u32::MAX; dmask_words])?;
13276 dmask_live = false;
13277 }
13278 gr.launch()?;
13279 scratch.kv.len += 1; // host mirror (len_d advanced in-graph)
13280 let idx = e.dtoh_u32_one(&dctx.g_tok)?;
13281 // #87 SENTINEL TRAP: an all-NaN head-logits row leaves the device argmax's
13282 // init sentinel (0x7FFFFFFF) in g_tok — feeding it onward dereferences
13283 // embed_row(sentinel) = table + ~4.6TB (never mapped) inside the NEXT graph
13284 // replay's embed node, and the MMU fault kills the CUDA context for the
13285 // whole process (research/pp2spec-crash-20260807: 3 coredumps, byte-exact
13286 // VA arithmetic). Refuse loudly instead; the diagnostics name the first-NaN
13287 // buffer (g_seed = the verify-side handoff vs head-side compute).
13288 if (idx as usize) >= d_vocab {
13289 // g_seed is SELF-FED (the replay writes h_nextn back into it), so it
13290 // reads as the head's OUTPUT at j; h_seed_buf is the round's INPUT
13291 // seed, untouched since the round-start copy — the pair discriminates
13292 // "seed arrived poisoned" from "head forward produced NaN".
13293 let seed_h = e.dtoh(&dctx.g_seed)?;
13294 let seed_nan = seed_h.iter().filter(|v| v.is_nan()).count();
13295 let in_h = e.dtoh(&h_seed_buf)?;
13296 let in_nan = in_h.iter().filter(|v| v.is_nan()).count();
13297 return Err(format!(
13298 "draft(graph) argmax sentinel 0x{idx:08x} >= d_vocab {d_vocab} at \
13299 round {round} j={j} pos={pos}: head-out NaN {seed_nan}/{n_embd}, \
13300 round-input-seed NaN {in_nan}/{n_embd} — refusing to dereference \
13301 the embed row (#87 trap)"
13302 )
13303 .into());
13304 }
13305 // trimmed draft vocab -> target token id (identity when no d2t map)
13306 let d = match &mtp.d2t {
13307 Some(map) => map[idx as usize],
13308 None => idx,
13309 };
13310 let draft_p = if p_min > 0.0
13311 || opti_fork
13312 .as_ref()
13313 .is_some_and(|fork| fork.controller.is_some())
13314 {
13315 Some(e.dtoh(&dctx.g_p)?[0])
13316 } else {
13317 None
13318 };
13319 if j == 0 {
13320 controller_draft_prob = draft_p;
13321 }
13322 if let Some(p) = draft_p.filter(|_| p_min > 0.0) {
13323 if p < p_min && (j > 0 || (pmin0 && base0 == 1)) {
13324 break;
13325 }
13326 }
13327 draft.push(d);
13328 // with a trimmed head the NEXT embed must read the TARGET id, not the draft
13329 // index the argmax wrote — patch the persistent token buffer (4B htod).
13330 if d != idx {
13331 e.set_u32_one(&mut dctx.g_tok, d)?;
13332 }
13333 // advance the SPECULATIVE state with the proposal; a dead chain drops to
13334 // unmasked drafting for the remaining positions (verify still arbitrates).
13335 // speculative advance; a chain the grammar can no longer follow (EOS
13336 // proposed) ends here. The captured mask node always runs, so a dead chain
13337 // leaves the buffer NEUTRAL (all-ones = ban nothing) before it exits.
13338 if dmask_live
13339 && !constraint
13340 .as_deref_mut()
13341 .unwrap()
13342 .draft_advance(d)
13343 .map_err(|e2| format!("constraint: {e2}"))?
13344 {
13345 e.htod_u32_into(&mut dctx.g_dmask, &vec![u32::MAX; dmask_words])?;
13346 break;
13347 }
13348 }
13349 // REGIME RE-TEST (lane/graph-s-key-exactness-20260819, widened by
13350 // lane/step37-draft-graph-serving-20260830): the sampled graph is legal ONLY
13351 // in the regime it was captured in. The condition used to read
13352 // `(sampled, &dctx.graph_s)` and trusted `s_key` to have dropped anything
13353 // else — which it could not, because the key omitted the filters. Both
13354 // halves are enforced: the key drops a stale graph, and this site refuses to
13355 // launch one whose key differs or whose regime is uncapturable (penalties).
13356 } else if let (true, Some(gr)) = (
13357 sampled && s_capturable && dctx.s_key == Some(s_key),
13358 &dctx.graph_s,
13359 ) {
13360 if skey_probe() {
13361 eprintln!(
13362 "[skey] chain=graph_s round={round} pure_temp={} capturable={} \
13363 top_k={} top_p={} min_p={} s_key_parked={:?}",
13364 pure_temp as u8,
13365 s_capturable as u8,
13366 sp.top_k,
13367 sp.top_p,
13368 sp.min_p,
13369 dctx.s_key,
13370 );
13371 }
13372 // SAMPLED GRAPH DRAFT: one replay per drafted token — head forward + gumbel +
13373 // argmax in ONE dispatch; the host reads 4B token (+4B p), D2Ds q into slot j,
13374 // and decides the break. Event-counter continuity: g_ctr is host-seeded to
13375 // sctr-1 ONCE per round (outside the graph); the in-graph bump runs BEFORE the
13376 // perturb, so replay j consumes counter sctr+j — exactly the eager arm's Philox
13377 // stream. Host sctr advances in lockstep (computed, no readback needed).
13378 e.set_i32_one(&mut dctx.g_pos, (pos + base0) as i32)?;
13379 e.set_u32_one(&mut dctx.g_tok, last_token)?;
13380 e.copy_into(&mut dctx.g_seed, 0, &h_seed_buf, n_embd)?;
13381 e.set_u32_one(&mut dctx.g_ctr, sctr.wrapping_sub(1))?;
13382 let filtered_stats_in_graph = s_key.filtered();
13383 for j in 0..k_this {
13384 gr.launch()?;
13385 scratch.kv.len += 1; // host mirror (len_d advanced in-graph)
13386 sctr += 1; // mirrors the in-graph g_ctr bump (eager parity:
13387 // counts the p-min-discarded token too)
13388 // q retention: ONE async D2D of the persistent head-logits buffer into this
13389 // round's slot j (stream-ordered after the replay, before the next one).
13390 e.copy_into(&mut dctx.q_slots[j], 0, &dctx.g_q, d_vocab)?;
13391 // FILTERED capture: the replay's own filter_stats node already computed
13392 // (th, z, mx) — read the three scalars back instead of paying a SECOND
13393 // full-vocab filter_stats per slot post-chain (measured ~5% of vendor-
13394 // default serving tok/s at K=3). Bit-exact by construction: these are
13395 // the very values the in-graph perturb consumed. Read BEFORE the p-min
13396 // break so the discarded slot's stats land too (accept-path indexing).
13397 if filtered_stats_in_graph {
13398 draft_stats.push((
13399 e.dtoh(&dctx.g_mx)?[0],
13400 e.dtoh(&dctx.g_th)?[0],
13401 e.dtoh(&dctx.g_z)?[0],
13402 ));
13403 }
13404 let idx = e.dtoh_u32_one(&dctx.g_tok)?;
13405 // #87 SENTINEL TRAP (see the greedy graph arm above).
13406 if (idx as usize) >= d_vocab {
13407 let seed_h = e.dtoh(&dctx.g_seed)?;
13408 let seed_nan = seed_h.iter().filter(|v| v.is_nan()).count();
13409 return Err(format!(
13410 "draft(graph-sampled) argmax sentinel 0x{idx:08x} >= d_vocab \
13411 {d_vocab} at round {round} j={j} pos={pos}: round-seed NaN \
13412 {seed_nan}/{n_embd} — refusing to dereference the embed row \
13413 (#87 trap)"
13414 )
13415 .into());
13416 }
13417 let d = match &mtp.d2t {
13418 Some(map) => map[idx as usize],
13419 None => idx,
13420 };
13421 draft_idx.push(idx);
13422 if p_min > 0.0 {
13423 let p = e.dtoh(&dctx.g_p)?[0];
13424 if p < p_min && (j > 0 || (pmin0 && base0 == 1)) {
13425 break;
13426 }
13427 }
13428 draft.push(d);
13429 // trimmed head: the NEXT embed must read the TARGET id (see the greedy arm).
13430 if d != idx {
13431 e.set_u32_one(&mut dctx.g_tok, d)?;
13432 }
13433 }
13434 // PURE-TEMP accept path: fill draft_stats per used slot post-chain (the
13435 // stats degenerate to th=0 / full-Z; one filter_stats launch per slot).
13436 // The FILTERED graph read its stats back per replay above.
13437 if !filtered_stats_in_graph {
13438 for j in 0..draft.len().max(draft_idx.len()) {
13439 let rows0 = e.htod_i32(&[0])?;
13440 let (mut th_d, mut z_d, mut mx_d) =
13441 (e.zeros(1)?, e.zeros(1)?, e.zeros(1)?);
13442 e.filter_stats(
13443 &dctx.q_slots[j],
13444 d_vocab,
13445 &rows0,
13446 &mut th_d,
13447 &mut z_d,
13448 &mut mx_d,
13449 d_vocab,
13450 1,
13451 sp_temp,
13452 sp.top_k,
13453 sp.top_p,
13454 sp.min_p,
13455 )?;
13456 draft_stats.push((
13457 e.dtoh(&mx_d)?[0],
13458 e.dtoh(&th_d)?[0],
13459 e.dtoh(&z_d)?[0],
13460 ));
13461 }
13462 }
13463 } else {
13464 if skey_probe() && sampled {
13465 eprintln!(
13466 "[skey] chain=eager round={round} pure_temp={} top_k={} \
13467 top_p={} min_p={} s_key_parked={:?}",
13468 pure_temp as u8, sp.top_k, sp.top_p, sp.min_p, dctx.s_key,
13469 );
13470 }
13471 // EAGER DRAFT (fallback: MoE head/trunk, huge k, MEMRA_SPEC_NOGRAPH, capture fail).
13472 let chain_heads = !self.mtp_extra.is_empty();
13473 let mut e_tok = last_token;
13474 let mut d_seed = e.clone_dtod(&h_seed_buf)?;
13475 let mut chain_tokens = if chain_heads {
13476 vec![last_token]
13477 } else {
13478 Vec::new()
13479 };
13480 let mut chain_seeds = if chain_heads {
13481 vec![e.clone_dtod(&h_seed_buf)?]
13482 } else {
13483 Vec::new()
13484 };
13485 for j in 0..k_this {
13486 // GPU-ARGMAX DRAFT (2026-07-03): device logits + device argmax + 4-byte token
13487 // read instead of the ~600KB full-vocab dtoh + host argmax per draft token.
13488 let mtp_pos = pos + base0 + j;
13489 // draft-side grammar mask (eager twin of the graph arm's in-graph node).
13490 // A position with no legal draft-vocab row drops to unmasked drafting for
13491 // the rest of the chain (pre-lane behaviour; verify still arbitrates).
13492 if dmask_live {
13493 dmask_live = upload_draft_mask(
13494 e,
13495 constraint.as_deref_mut().unwrap(),
13496 &mut dctx.g_dmask,
13497 mtp.d2t.as_ref(),
13498 d_vocab,
13499 dmask_words,
13500 )?;
13501 }
13502 let mask = if dmask_live {
13503 Some((&dctx.g_dmask, dmask_words))
13504 } else {
13505 None
13506 };
13507 let (dl_d, h_nextn) = if chain_heads {
13508 if debug_spec {
13509 eprintln!(
13510 "[mtp-chain-step] round={round} j={j} head={} replay_rows={}",
13511 mtp_chain_head_index(j, self.mtp_head_count()),
13512 chain_tokens.len(),
13513 );
13514 }
13515 self.mtp_chain_forward_dev(
13516 e,
13517 &chain_tokens,
13518 &chain_seeds,
13519 &mut *scratch,
13520 pos + base0 - 1,
13521 embd_dev,
13522 mask,
13523 )?
13524 } else {
13525 self.mtp_head_forward_dev(
13526 e,
13527 mtp,
13528 e_tok,
13529 &d_seed,
13530 &mut *scratch,
13531 mtp_pos,
13532 embd_dev,
13533 mask,
13534 )?
13535 };
13536 let tok_d = if sampled {
13537 // FILTERED Gumbel-max: stats -> masked perturb -> argmax = one draw from
13538 // the filtered softmax (filters off => th=0, exact v1 semantics).
13539 if perturb_buf.is_none() {
13540 perturb_buf = Some(e.zeros(d_vocab.max(n_vocab))?);
13541 }
13542 let mut q_row = e.clone_dtod(&dl_d)?; // retained q (penalized when on)
13543 if pen_on {
13544 let h = pen_hist_d.as_ref().unwrap();
13545 let nh = h.len();
13546 e.penalize_logits(
13547 &mut q_row,
13548 h,
13549 nh,
13550 sp.penalty_repeat,
13551 sp.penalty_freq,
13552 sp.penalty_present,
13553 d_vocab,
13554 )?;
13555 }
13556 let rows0 = e.htod_i32(&[0])?;
13557 let (mut th_d, mut z_d, mut mx_d) =
13558 (e.zeros(1)?, e.zeros(1)?, e.zeros(1)?);
13559 e.filter_stats(
13560 &q_row, d_vocab, &rows0, &mut th_d, &mut z_d, &mut mx_d, d_vocab,
13561 1, sp_temp, sp.top_k, sp.top_p, sp.min_p,
13562 )?;
13563 let (th, z, mx) =
13564 (e.dtoh(&th_d)?[0], e.dtoh(&z_d)?[0], e.dtoh(&mx_d)?[0]);
13565 let pb = perturb_buf.as_mut().unwrap();
13566 e.gumbel_perturb_filtered(
13567 &q_row, pb, d_vocab, sp_seed, sctr, sp_temp, mx, th,
13568 )?;
13569 sctr += 1;
13570 draft_logits.push(q_row);
13571 draft_stats.push((mx, th, z));
13572 e.argmax_token_device(pb, d_vocab)?
13573 } else {
13574 e.argmax_token_device(&dl_d, d_vocab)?
13575 };
13576 let idx = e.dtoh_u32_one(&tok_d)?;
13577 // #87 SENTINEL TRAP (eager twin — see the graph arm). Extra diagnostics
13578 // here because the eager chain's operands are all readable: dl_d (the head
13579 // logits row) and d_seed (this step's h_seed) name the first-NaN buffer.
13580 if (idx as usize) >= d_vocab {
13581 let dl_h = e.dtoh(&dl_d)?;
13582 let dl_nan = dl_h.iter().filter(|v| v.is_nan()).count();
13583 let seed_h = if chain_heads {
13584 e.dtoh(chain_seeds.last().unwrap())?
13585 } else {
13586 e.dtoh(&d_seed)?
13587 };
13588 let seed_nan = seed_h.iter().filter(|v| v.is_nan()).count();
13589 return Err(format!(
13590 "draft(eager) argmax sentinel 0x{idx:08x} >= d_vocab {d_vocab} at \
13591 round {round} j={j} pos={pos}: head-logits NaN {dl_nan}/{d_vocab}, \
13592 step-seed NaN {seed_nan}/{n_embd} — refusing to dereference the \
13593 embed row (#87 trap)"
13594 )
13595 .into());
13596 }
13597 let d = match &mtp.d2t {
13598 Some(map) => map[idx as usize],
13599 None => idx,
13600 };
13601 if sampled {
13602 draft_idx.push(idx);
13603 }
13604 let draft_p = if p_min > 0.0
13605 || opti_fork
13606 .as_ref()
13607 .is_some_and(|fork| fork.controller.is_some())
13608 {
13609 let p_d = e.prob_of_token_device(&dl_d, &tok_d, d_vocab)?;
13610 Some(e.dtoh(&p_d)?[0])
13611 } else {
13612 None
13613 };
13614 if j == 0 {
13615 controller_draft_prob = draft_p;
13616 }
13617 if let Some(p) = draft_p.filter(|_| p_min > 0.0) {
13618 if p < p_min && (j > 0 || (pmin0 && base0 == 1)) {
13619 break;
13620 }
13621 }
13622 draft.push(d);
13623 if chain_heads {
13624 chain_tokens.push(d);
13625 chain_seeds.push(h_nextn);
13626 } else {
13627 e_tok = d;
13628 d_seed = h_nextn;
13629 }
13630 // speculative advance; a chain the grammar can no longer follow (EOS
13631 // proposed) ends here — the prefix already proposed still rides verify.
13632 if dmask_live
13633 && !constraint
13634 .as_deref_mut()
13635 .unwrap()
13636 .draft_advance(d)
13637 .map_err(|e2| format!("constraint: {e2}"))?
13638 {
13639 break;
13640 }
13641 }
13642 if !chain_heads
13643 && opti_fork
13644 .as_ref()
13645 .is_some_and(|fork| fork.controller.is_some())
13646 {
13647 controller_eager_state = Some((e_tok, d_seed));
13648 }
13649 }
13650 }
13651 let k_round = draft.len();
13652 if let Some(p) = pipe {
13653 p.draft_end(round);
13654 }
13655 drop(pipe_draft);
13656
13657 ph_mark(&mut ph_draft, phase_on);
13658 // --- 2. VERIFY: one batched target forward. With a pending bonus, it rides as col 0
13659 // (committing its KV/recur inside the SAME weight read); drafts follow. ---
13660 let verify_tokens: Vec<u32> = match pending {
13661 Some(b) => {
13662 let mut v = Vec::with_capacity(k_round + 1);
13663 v.push(b);
13664 v.extend_from_slice(&draft);
13665 v
13666 }
13667 None => draft.clone(),
13668 };
13669 let base = if pending.is_some() { 1 } else { 0 };
13670 // ckpt (REPLAY-FREE partial accept): retain per-layer state-rebuild inputs alongside
13671 // the verify. Pure buffer keep-alives + dtod clones — kernel work is unchanged.
13672 let mut ckpt = if let Some(ticket) = current_opti.as_mut() {
13673 Some(ticket.take_ckpt())
13674 } else if spec_replay {
13675 None
13676 } else {
13677 Some(VerifyCkpt::new(self.layers.len()))
13678 };
13679 let controller_can_probe = base == 1
13680 && k_round == 1
13681 && out.len().saturating_add(2) < max_new
13682 && controller_draft_prob.is_some()
13683 && opti_fork
13684 .as_ref()
13685 .and_then(|fork| fork.controller.as_ref())
13686 .is_some_and(|policy| !policy.breaker_tripped);
13687 let mut successor_attempt: Option<OptiControllerTicket> = None;
13688 let mut rejected_probe: Option<(f32, u32)> = None;
13689 let mut controller_prepared: Option<OptiControllerPrepared> = None;
13690 if controller_can_probe {
13691 // Prepare d2/q and, on admission, d3 before either current verify half is
13692 // issued. N stage 0 can then be followed immediately by N+1 stage 0; once N's
13693 // boundary fires, those dev0 launches overlap N stage 1 on dev1. Preparing on
13694 // the primary stream after N stage 1 would serialize the supposed pipeline.
13695 let eager_pos = scratch.kv.len + 1;
13696 let (optimistic_pending, pending_probability) = self.opti_controller_draft_step(
13697 e,
13698 mtp,
13699 &mut dctx,
13700 &mut *scratch,
13701 d_vocab,
13702 &mut controller_eager_state,
13703 eager_pos,
13704 embd_dev,
13705 )?;
13706 let first_probability = controller_draft_prob
13707 .ok_or("optipipe controller probe lost first-token probability")?;
13708 let q_proxy = first_probability * pending_probability;
13709 OPTI_GATE_CHECKS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
13710 OPTI_SHADOW_DRAFT_TOKENS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
13711 let admitted = opti_fork
13712 .as_ref()
13713 .and_then(|fork| fork.controller.as_ref())
13714 .ok_or("optipipe controller policy disappeared")?
13715 .admit(q_proxy);
13716 if admitted {
13717 OPTI_GATE_ADMITS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
13718 let eager_pos = scratch.kv.len + 1;
13719 let (optimistic_draft, optimistic_draft_probability) = self
13720 .opti_controller_draft_step(
13721 e,
13722 mtp,
13723 &mut dctx,
13724 &mut *scratch,
13725 d_vocab,
13726 &mut controller_eager_state,
13727 eager_pos,
13728 embd_dev,
13729 )?;
13730 OPTI_SHADOW_DRAFT_TOKENS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
13731 let eager_seed = controller_eager_state.take().map(|(token, seed)| {
13732 debug_assert_eq!(token, optimistic_draft);
13733 seed
13734 });
13735 controller_prepared = Some(OptiControllerPrepared {
13736 verify_tokens: [optimistic_pending, optimistic_draft],
13737 draft_prob: optimistic_draft_probability,
13738 eager_seed,
13739 q_proxy,
13740 scratch_len: scratch.kv.len,
13741 });
13742 } else {
13743 OPTI_GATE_REJECTS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
13744 OPTI_WASTED_DRAFT_TOKENS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
13745 rejected_probe = Some((q_proxy, optimistic_pending));
13746 eprintln!(
13747 "[opti-controller] reject q={q_proxy:.6} threshold={:.3}",
13748 opti_fork
13749 .as_ref()
13750 .and_then(|fork| fork.controller.as_ref())
13751 .expect("controller policy")
13752 .threshold,
13753 );
13754 }
13755 }
13756 let fork_attempt = match fork_generation.take() {
13757 Some(generation) if base == 1 && k_round == 1 => Some(generation),
13758 Some(generation) => {
13759 opti_fork
13760 .as_mut()
13761 .expect("fork generation without fork state")
13762 .retire(generation)?;
13763 None
13764 }
13765 None => None,
13766 };
13767 let (tlogits_d, vx) = if let Some(p) = pipe {
13768 self.decode_step_t_core_pipelined(
13769 e,
13770 &verify_tokens,
13771 pos,
13772 &mut *cache,
13773 embd_dev,
13774 ckpt.as_mut(),
13775 p,
13776 round,
13777 )?
13778 } else if controller_can_probe {
13779 let fence = opti_fork
13780 .as_ref()
13781 .ok_or("optipipe controller probe lost fork state")?
13782 .fence;
13783 let boundary = match current_opti.as_mut() {
13784 Some(ticket) => ticket.take_boundary(),
13785 None => self.verify_stage0_issue(
13786 e,
13787 &verify_tokens,
13788 pos,
13789 &mut *cache,
13790 embd_dev,
13791 ckpt.as_mut(),
13792 None,
13793 &fence,
13794 Some(true),
13795 None,
13796 )?,
13797 };
13798 if let Some(prepared) = controller_prepared.take() {
13799 let generation = {
13800 let fork = opti_fork
13801 .as_mut()
13802 .ok_or("optipipe controller admission lost fork state")?;
13803 let generation = fork.reserve_successor()?;
13804 let rt = fork.rt;
13805 let snapshot_fence = fork.fence;
13806 opti_snapshot_one_stage_owned_into(
13807 e,
13808 cache,
13809 rt,
13810 &snapshot_fence,
13811 0,
13812 fork.successor_snapshot_mut(),
13813 )?;
13814 generation
13815 };
13816 let mut successor_ckpt = VerifyCkpt::new(self.layers.len());
13817 let successor_boundary = self.verify_stage0_issue(
13818 e,
13819 &prepared.verify_tokens,
13820 pos + verify_tokens.len(),
13821 &mut *cache,
13822 embd_dev,
13823 Some(&mut successor_ckpt),
13824 None,
13825 &fence,
13826 Some(false),
13827 None,
13828 )?;
13829 OPTI_FORK_ATTEMPTS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
13830 let fork = opti_fork
13831 .as_ref()
13832 .ok_or("optipipe controller ticket lost fork state")?;
13833 successor_attempt = Some(fork.controller_ticket(
13834 generation,
13835 successor_boundary,
13836 successor_ckpt,
13837 prepared.verify_tokens,
13838 prepared.draft_prob,
13839 prepared.eager_seed,
13840 prepared.q_proxy,
13841 prepared.scratch_len,
13842 ));
13843 eprintln!(
13844 "[opti-controller] issue generation={} q={:.6} threshold={:.3} \
13845 verify={:?}",
13846 generation.id,
13847 prepared.q_proxy,
13848 fork.controller.expect("controller policy").threshold,
13849 prepared.verify_tokens,
13850 );
13851 }
13852 let result = self.verify_stage1_finish(
13853 e,
13854 boundary,
13855 &mut *cache,
13856 ckpt.as_mut(),
13857 None,
13858 &fence,
13859 successor_attempt.is_none(),
13860 )?;
13861 if let Some(ticket) = current_opti.as_mut() {
13862 ticket.settle();
13863 }
13864 if successor_attempt.is_some() {
13865 let fork = opti_fork
13866 .as_mut()
13867 .ok_or("optipipe successor snapshot lost fork state")?;
13868 let rt = fork.rt;
13869 let snapshot_fence = fork.fence;
13870 opti_snapshot_one_stage_owned_into(
13871 e,
13872 cache,
13873 rt,
13874 &snapshot_fence,
13875 1,
13876 fork.successor_snapshot_mut(),
13877 )?;
13878 // Publish N only after both independent successor-state queues are complete.
13879 fork.rt.publish_to(1, &e.stream())?;
13880 }
13881 result
13882 } else if let Some(ticket) = current_opti.as_mut() {
13883 let fork = opti_fork
13884 .as_mut()
13885 .ok_or("optipipe carried controller ticket lost fork state")?;
13886 let boundary = ticket.take_boundary();
13887 let result = self.verify_stage1_finish(
13888 e,
13889 boundary,
13890 &mut *cache,
13891 ckpt.as_mut(),
13892 None,
13893 &fork.fence,
13894 true,
13895 )?;
13896 ticket.settle();
13897 result
13898 } else if let Some(generation) = fork_attempt {
13899 let fork = opti_fork
13900 .as_mut()
13901 .expect("fork generation without fork state");
13902 fork.capture_seed(e, generation, &h_seed_buf, &fill_prev, scratch.kv.len)?;
13903 let action = fork.mode.action(generation.id);
13904 let boundary = self.verify_stage0_issue(
13905 e,
13906 &verify_tokens,
13907 pos,
13908 &mut *cache,
13909 embd_dev,
13910 ckpt.as_mut(),
13911 None,
13912 &fork.fence,
13913 Some(true),
13914 None,
13915 )?;
13916 OPTI_FORK_ATTEMPTS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
13917 let mut ticket = fork.ticket(generation, boundary);
13918 if action == OptiForkAction::Abort {
13919 return Err(format!(
13920 "optipipe forced abort with generation {} stage0 in flight",
13921 generation.id,
13922 )
13923 .into());
13924 }
13925 fork.reconcile(
13926 e,
13927 &mut *cache,
13928 &mut *scratch,
13929 &snap,
13930 &mut h_seed_buf,
13931 &mut fill_prev,
13932 generation,
13933 action,
13934 verify_tokens[0],
13935 )?;
13936 let result = if action == OptiForkAction::Hit {
13937 let boundary = ticket.take_boundary();
13938 self.verify_stage1_finish(
13939 e,
13940 boundary,
13941 &mut *cache,
13942 ckpt.as_mut(),
13943 None,
13944 &fork.fence,
13945 true,
13946 )?
13947 } else {
13948 // The optimistic boundary slot has no reader. Re-run the unchanged serial
13949 // verify only after E_restart published the restored stage-0 state.
13950 self.decode_step_t_core(
13951 e,
13952 &verify_tokens,
13953 pos,
13954 &mut *cache,
13955 embd_dev,
13956 ckpt.as_mut(),
13957 )?
13958 };
13959 ticket.settle();
13960 debug_assert_eq!(ticket.generation, generation);
13961 fork.retire(generation)?;
13962 result
13963 } else {
13964 // The serial verify every non-fork round takes — the MTP route's
13965 // verify-graph door. The pool is None unless MEMRA_SPEC_VERIFY_GRAPH armed
13966 // a pool above, and then the walk replays the captured trunk instead of
13967 // re-issuing it launch by launch.
13968 let vg_round = if verify_tokens.len() <= vg_t_cap {
13969 vg_guard.as_mut().and_then(|g| g.as_mut())
13970 } else {
13971 if let Some(g) = vg_guard.as_mut().and_then(|g| g.as_mut()) {
13972 // The commit reads this flag to pick its arm; a round that declines
13973 // the pool must not inherit a stale `true` from the round before it.
13974 g.round_slab = false;
13975 }
13976 None
13977 };
13978 self.decode_step_t_core_vg(
13979 e,
13980 &verify_tokens,
13981 pos,
13982 &mut *cache,
13983 embd_dev,
13984 ckpt.as_mut(),
13985 vg_round,
13986 )?
13987 };
13988 let pipe_accept = match pipe {
13989 Some(p) => Some(p.accept_begin(round)?),
13990 None => None,
13991 };
13992
13993 if phase_sync {
13994 e.stream().synchronize()?;
13995 }
13996 ph_mark(&mut ph_verify, phase_on);
13997 // --- 3. GREEDY ACCEPT (walk prefix, stop at first mismatch) ---
13998 // DEVICE-ARGMAX ACCEPT: argmax every verify column ON DEVICE (same 2-pass kernels +
13999 // smallest-index tie-break as host argmax, argmax_gate-validated) and read back ONE
14000 // [T] u32 — replaces the T x n_vocab f32 dtoh + T host argmaxes per round.
14001 // t_pred[j] = target's greedy prediction for the slot after draft[j-1] (j>=1) or after
14002 // last_token (j==0). With a pending bonus, col 0 IS the prediction after last_token
14003 // (== the bonus), so every index shifts by `base` and last_pred is unused.
14004 let t_v = verify_tokens.len();
14005 let mut preds: Vec<u32> = Vec::new();
14006 if !sampled {
14007 for j in 0..t_v {
14008 e.argmax_token_device_col(&tlogits_d, j, n_vocab, &mut preds_d, j)?;
14009 }
14010 preds = e.dtoh_u32(&preds_d)?; // <- the verify-GPU wait lands here
14011 // #87 SENTINEL TRAP, verify side: a sentinel pred becomes the round's bonus =
14012 // next round's last_token = the next chain's embed lookup. Catch it at the
14013 // source with the column named — an all-NaN VERIFY column implicates the
14014 // stage-split trunk (decode_step_t_core_ppn), not the draft head.
14015 if let Some(bad) = preds[..t_v].iter().position(|&p| (p as usize) >= n_vocab) {
14016 let col = &tlogits_d.slice(bad * n_vocab..(bad + 1) * n_vocab);
14017 let mut probe = e.zeros(n_vocab)?;
14018 e.copy_view_into(&mut probe, 0, col, n_vocab)?;
14019 let col_h = e.dtoh(&probe)?;
14020 let col_nan = col_h.iter().filter(|v| v.is_nan()).count();
14021 return Err(format!(
14022 "verify argmax sentinel 0x{:08x} >= n_vocab {n_vocab} at round {round} \
14023 col {bad}/{t_v} pos={pos}: verify-logits col NaN {col_nan}/{n_vocab} \
14024 — the verify TRUNK produced a poisoned column (#87 trap). Run \
14025 MEMRA_SPEC_NAN_SCAN=1 to name the layer that creates it (=2 to split \
14026 that layer into attention and routed MoE). NOT the draft head, and NOT \
14027 the PP stage split this message used to name: pp_cuts() returns None \
14028 without MEMRA_PP_STAGES, so decode_step_t_core_ppn never runs unless \
14029 that variable is set.",
14030 preds[bad]
14031 )
14032 .into());
14033 }
14034 }
14035 ph_mark(&mut ph_wait, phase_on);
14036 let t_pred = |j: usize| -> u32 {
14037 if j == 0 && base == 0 {
14038 last_pred
14039 } else {
14040 // GREEDY-ONLY: `preds` is filled under `if !sampled` above. The debug print
14041 // used to call this from the sampled arm and panicked the worker; it now goes
14042 // through `debug_t_pred0`. Keep the strict index here — in the greedy walk an
14043 // out-of-range pred is a real bug, not something to paper over.
14044 debug_assert!(
14045 !sampled,
14046 "t_pred is greedy-only: `preds` is empty in the sampled arm"
14047 );
14048 preds[base + j - 1]
14049 }
14050 };
14051 let mut devacc_seeded = false;
14052 let mut devacc_acc: Option<CudaSlice<u32>> = None;
14053 let (n_acc, bonus) = if !sampled {
14054 // ROUND-STREAM stage (a) (MEMRA_SPEC_DEVACC=1 opt-in): the walk runs ON DEVICE
14055 // (spec_accept_greedy, verbatim rule) and the host reads back 8B (n_acc, bonus)
14056 // instead of the [T] preds. Same sync count — machinery for stages (b)/(c),
14057 // gated on token identity vs the host walk (the arms below are bit-equal rules).
14058 if crate::spec::spec_devacc() && k_round > 0 && !spec_replay && constraint.is_none()
14059 {
14060 let draft_d = e.htod_u32_v(&draft)?;
14061 let mut acc_out = e.alloc_u32_zeroed(2)?;
14062 e.spec_accept_greedy(
14063 &preds_d,
14064 &draft_d,
14065 last_pred,
14066 base,
14067 k_round,
14068 &mut acc_out,
14069 )?;
14070 devacc_acc = Some(acc_out.clone());
14071 // stage (b): next-round seed gathered ON DEVICE from acc_out before the host
14072 // ever reads n_acc (j=base+n_acc -> vx col j-1; j==0 -> fill_prev). The three
14073 // non-replay commit arms skip their host-offset seed copies (guarded below);
14074 // the legacy spec_replay arm keeps its own rx-based seeding (excluded here).
14075 // NOTE: fill_prev is NOT updated here — the commit arms' TRUE-HIDDEN
14076 // REFRESH reads the OLD fill_prev (predecessor of this round's verify batch);
14077 // the update lands after the arms (devacc_seeded guard below).
14078 e.spec_seed_gather(&vx, &fill_prev, &acc_out, &mut h_seed_buf, base, n_embd)?;
14079 // 3a: KV lens roll back on device (len = saved + base + n_acc, all arms'
14080 // unified rule; full accept rewrites the verify-left value). Host mirrors
14081 // update after the readback; commit_verified_prefix skips its len_d writes.
14082 if let Some(successor) = successor_attempt.as_ref() {
14083 opti_fork
14084 .as_mut()
14085 .ok_or("optipipe successor reconcile lost fork state")?
14086 .queue_actual_reconcile(
14087 e,
14088 &snap,
14089 &acc_out,
14090 successor.verify_tokens[0],
14091 base,
14092 )?;
14093 } else if let Some(ptrs) = &kv_len_ptrs {
14094 let saved: Vec<i32> = (0..self.layers.len())
14095 .map(|il| snap.kv_len[il].map(|v| v as i32).unwrap_or(0))
14096 .collect();
14097 let saved_d = e.htod_i32(&saved)?;
14098 e.spec_rollback_kv(ptrs, &saved_d, &acc_out, base, self.layers.len())?;
14099 }
14100 devacc_seeded = true;
14101 let ab = e.dtoh_u32(&acc_out)?;
14102 (ab[0] as usize, ab[1])
14103 } else {
14104 let mut n_acc = 0usize;
14105 for j in 0..k_round {
14106 if t_pred(j) == draft[j] {
14107 n_acc += 1;
14108 } else {
14109 break;
14110 }
14111 }
14112 // bonus = target's own token at the first non-accepted slot. n_acc in 0..=k; t_pred
14113 // is defined for j in 0..=k (j==0 -> last_logits, j>=1 -> col j-1, last col = k-1).
14114 (n_acc, t_pred(n_acc))
14115 }
14116 } else {
14117 // --- SAMPLED ACCEPT (rejection sampling): u_j < p_j(x_j)/q_j(x_j) walk ---
14118 if col_buf.is_none() {
14119 col_buf = Some(e.zeros(n_vocab)?);
14120 }
14121 // FILTERED p_j: per-verify-col stats (one batched filter_stats call), then the
14122 // filtered gather. j==0&&base==0 reads last_col (its own stats row appended).
14123 let mut pj = vec![0f32; k_round.max(1)];
14124 let mut col_stats: Vec<(f32, f32, f32)> = Vec::new(); // (max, th, z) per verify col used
14125 if k_round > 0 {
14126 let mut ids: Vec<u32> = Vec::new();
14127 let mut rows: Vec<i32> = Vec::new();
14128 for j in 0..k_round {
14129 if j > 0 || base == 1 {
14130 ids.push(draft[j]);
14131 rows.push((base + j) as i32 - 1);
14132 }
14133 }
14134 if !ids.is_empty() {
14135 let nr = rows.len();
14136 // penalties: materialize the used columns into one contiguous penalized
14137 // buffer (rows remapped 0..nr) so stats+gathers see the penalized p.
14138 // penalties: materialize used columns contiguously, penalize all rows in
14139 // one launch, and point stats+gathers at the penalized buffer (rows 0..nr).
14140 let p_rows: Vec<i32> = if pen_on {
14141 (0..nr as i32).collect()
14142 } else {
14143 rows.clone()
14144 };
14145 if pen_on {
14146 if pcol_buf.as_ref().map(|b| b.len()).unwrap_or(0) < nr * n_vocab {
14147 pcol_buf = Some(e.zeros(nr * n_vocab)?);
14148 }
14149 let pc = pcol_buf.as_mut().unwrap();
14150 for (i2, &r) in rows.iter().enumerate() {
14151 let c = r as usize;
14152 e.copy_view_into(
14153 pc,
14154 i2 * n_vocab,
14155 &tlogits_d.slice(c * n_vocab..(c + 1) * n_vocab),
14156 n_vocab,
14157 )?;
14158 }
14159 let h = pen_hist_d.as_ref().unwrap();
14160 let nh = h.len();
14161 e.penalize_logits_rows(
14162 pc,
14163 h,
14164 nh,
14165 sp.penalty_repeat,
14166 sp.penalty_freq,
14167 sp.penalty_present,
14168 n_vocab,
14169 nr,
14170 )?;
14171 }
14172 let p_src: &CudaSlice<f32> = if pen_on {
14173 pcol_buf.as_ref().unwrap()
14174 } else {
14175 &tlogits_d
14176 };
14177 let rowsd = e.htod_i32(&p_rows)?;
14178 let (mut th_d, mut z_d, mut mx_d) =
14179 (e.zeros(nr)?, e.zeros(nr)?, e.zeros(nr)?);
14180 e.filter_stats(
14181 p_src, n_vocab, &rowsd, &mut th_d, &mut z_d, &mut mx_d, n_vocab, nr,
14182 sp_temp, sp.top_k, sp.top_p, sp.min_p,
14183 )?;
14184 let idsd = e.htod_u32_v(&ids)?;
14185 let mut outd = e.zeros(nr)?;
14186 e.softmax_gather_filtered(
14187 p_src, n_vocab, &idsd, &rowsd, &th_d, &z_d, &mut outd, n_vocab, nr,
14188 sp_temp,
14189 )?;
14190 let outv = e.dtoh(&outd)?;
14191 let (thv, zv, mxv) = (e.dtoh(&th_d)?, e.dtoh(&z_d)?, e.dtoh(&mx_d)?);
14192 let mut oi = 0usize;
14193 for j in 0..k_round {
14194 if j > 0 || base == 1 {
14195 pj[j] = outv[oi];
14196 oi += 1;
14197 }
14198 }
14199 col_stats = (0..nr).map(|i| (mxv[i], thv[i], zv[i])).collect();
14200 }
14201 if base == 0 {
14202 let lc: &CudaSlice<f32> = if pen_on {
14203 if col_buf.is_none() {
14204 col_buf = Some(e.zeros(n_vocab)?);
14205 }
14206 let cb = col_buf.as_mut().unwrap();
14207 e.copy_into(
14208 cb,
14209 0,
14210 last_col_logits
14211 .as_ref()
14212 .expect("sampled: last_col_logits unset"),
14213 n_vocab,
14214 )?;
14215 let h = pen_hist_d.as_ref().unwrap();
14216 let nh = h.len();
14217 e.penalize_logits(
14218 cb,
14219 h,
14220 nh,
14221 sp.penalty_repeat,
14222 sp.penalty_freq,
14223 sp.penalty_present,
14224 n_vocab,
14225 )?;
14226 col_buf.as_ref().unwrap()
14227 } else {
14228 last_col_logits
14229 .as_ref()
14230 .expect("sampled: last_col_logits unset")
14231 };
14232 let rows0 = e.htod_i32(&[0])?;
14233 let (mut th_d, mut z_d, mut mx_d) = (e.zeros(1)?, e.zeros(1)?, e.zeros(1)?);
14234 e.filter_stats(
14235 lc, n_vocab, &rows0, &mut th_d, &mut z_d, &mut mx_d, n_vocab, 1,
14236 sp_temp, sp.top_k, sp.top_p, sp.min_p,
14237 )?;
14238 let idsd = e.htod_u32_v(&[draft[0]])?;
14239 let mut outd = e.zeros(1)?;
14240 e.softmax_gather_filtered(
14241 lc, n_vocab, &idsd, &rows0, &th_d, &z_d, &mut outd, n_vocab, 1, sp_temp,
14242 )?;
14243 pj[0] = e.dtoh(&outd)?[0];
14244 last_col_stats =
14245 Some((e.dtoh(&mx_d)?[0], e.dtoh(&th_d)?[0], e.dtoh(&z_d)?[0]));
14246 }
14247 }
14248 // q source: the graph arms (single-head AND chain) retained the head logits
14249 // in the persistent q_slots; the eager arm in per-round draft_logits clones.
14250 // Same raw-logit values either way. FILTERED q_j: stats from draft_stats
14251 // (eager pushes in-chain; the graph arms compute them post-replay from the
14252 // retained q with the same filter_stats program — bit-identical to the
14253 // in-graph stats that shaped the draw, keeping ONE accept path).
14254 let q_bufs: &[CudaSlice<f32>] = if dctx.graph_s.is_some() || dctx.chain_s.is_some()
14255 {
14256 &dctx.q_slots
14257 } else {
14258 &draft_logits
14259 };
14260 let mut n_acc = 0usize;
14261 for j in 0..k_round {
14262 let (qmx, qth, qz) = draft_stats[j];
14263 let idsd = e.htod_u32_v(&[draft_idx[j]])?;
14264 let rowsd = e.htod_i32(&[0])?;
14265 let thd = e.htod(&[qth])?;
14266 let zd = e.htod(&[qz])?;
14267 let _ = qmx;
14268 let mut outd = e.zeros(1)?;
14269 e.softmax_gather_filtered(
14270 &q_bufs[j], d_vocab, &idsd, &rowsd, &thd, &zd, &mut outd, d_vocab, 1,
14271 sp_temp,
14272 )?;
14273 let qj = e.dtoh(&outd)?[0];
14274 let u = host_u01(sp_seed, uctr);
14275 uctr += 1;
14276 let accept = (u as f64) * (qj as f64) < pj[j] as f64;
14277 // SKEY PROBE: q == 0 for the token the draft actually proposed is the
14278 // exactness signature (see `skey_probe`). Impossible when the draft was
14279 // drawn from the same filtered distribution the verify reconstructs here;
14280 // `u * 0 < p` makes it an UNCONDITIONAL accept whenever p > 0.
14281 if skey_probe() && qj == 0.0 {
14282 eprintln!(
14283 "[skey] EXACTNESS q=0 round={round} j={j} draft_tok={} \
14284 draft_idx={} p={:e} u={u} accepted={} th_z={:?}",
14285 draft[j], draft_idx[j], pj[j], accept as u8, draft_stats[j],
14286 );
14287 }
14288 if accept {
14289 n_acc += 1;
14290 } else {
14291 break;
14292 }
14293 }
14294 let bonus = if n_acc == k_round {
14295 // FULL ACCEPT: bonus ~ FILTERED softmax at the last verify column.
14296 let col = base + k_round - 1;
14297 let cb = col_buf.as_mut().unwrap();
14298 e.copy_view_into(
14299 cb,
14300 0,
14301 &tlogits_d.slice(col * n_vocab..(col + 1) * n_vocab),
14302 n_vocab,
14303 )?;
14304 if pen_on {
14305 let h = pen_hist_d.as_ref().unwrap();
14306 let nh = h.len();
14307 e.penalize_logits(
14308 cb,
14309 h,
14310 nh,
14311 sp.penalty_repeat,
14312 sp.penalty_freq,
14313 sp.penalty_present,
14314 n_vocab,
14315 )?;
14316 }
14317 if perturb_buf.is_none() {
14318 perturb_buf = Some(e.zeros(d_vocab.max(n_vocab))?);
14319 }
14320 // STATS MUST COME FROM THIS COLUMN (bug fix 2026-08-05, lane/sampler-
14321 // truncation-fix; receipts research/sampfix-20260805/). The old code reused
14322 // `col_stats.last()` here, which is ALWAYS the wrong row: the gathered set
14323 // covers verify columns 0..=(base+k_round-2) (rows pushed as base+j-1), while
14324 // the full-accept bonus samples column base+k_round-1 — exactly ONE PAST the
14325 // last gathered column, in both base arms. `th` is a threshold in e-units of
14326 // its OWN row's max, so feeding a neighbour's (row_max, th) into
14327 // gumbel_perturb_filtered mis-scales every e0 = exp((x-row_max)/T). When the
14328 // donor column's peak is higher by more than T*ln(1/th), EVERY id fails
14329 // `e0 >= th`, the whole perturbed row becomes -3.4e38, and the 2-pass argmax
14330 // falls through to its smallest-index tie-break => token id 0 ("!") spliced
14331 // mid-word. Fragility is ordered by how large th is: min_p pins th = min_p
14332 // (0.05 => trigger at delta > 2.4 at T=0.8, fires constantly), top_p's
14333 // mass-boundary th is smaller, top_k's k-th-largest th smaller still — which
14334 // is why the head-to-head matrix saw min_p and top_p corrupt while top_k-only
14335 // stayed clean. The pure-temp default regime is immune (th == 0 masks nothing,
14336 // and row_max is unused once nothing is masked), so this fix is a byte-level
14337 // no-op for the untruncated serve default. One extra one-block filter_stats
14338 // per full-accept round is the whole cost.
14339 let (mx, th) = {
14340 let rows0 = e.htod_i32(&[0])?;
14341 let (mut th_d, mut z_d, mut mx_d) = (e.zeros(1)?, e.zeros(1)?, e.zeros(1)?);
14342 let cb0 = col_buf.as_ref().unwrap();
14343 e.filter_stats(
14344 cb0, n_vocab, &rows0, &mut th_d, &mut z_d, &mut mx_d, n_vocab, 1,
14345 sp_temp, sp.top_k, sp.top_p, sp.min_p,
14346 )?;
14347 (e.dtoh(&mx_d)?[0], e.dtoh(&th_d)?[0])
14348 };
14349 let pb = perturb_buf.as_mut().unwrap();
14350 let cb2 = col_buf.as_ref().unwrap();
14351 e.gumbel_perturb_filtered(cb2, pb, n_vocab, sp_seed, sctr, sp_temp, mx, th)?;
14352 sctr += 1;
14353 let td = e.argmax_token_device(pb, n_vocab)?;
14354 e.dtoh_u32_one(&td)?
14355 } else {
14356 // REJECT at n_acc: bonus ~ norm(max(0, softmax_T(p) - softmax_T(q))).
14357 let cb = col_buf.as_mut().unwrap();
14358 if n_acc > 0 || base == 1 {
14359 let col = base + n_acc - 1;
14360 e.copy_view_into(
14361 cb,
14362 0,
14363 &tlogits_d.slice(col * n_vocab..(col + 1) * n_vocab),
14364 n_vocab,
14365 )?;
14366 } else {
14367 let lc = last_col_logits.as_ref().unwrap();
14368 e.copy_into(cb, 0, lc, n_vocab)?;
14369 }
14370 if pen_on {
14371 let h = pen_hist_d.as_ref().unwrap();
14372 let nh = h.len();
14373 e.penalize_logits(
14374 cb,
14375 h,
14376 nh,
14377 sp.penalty_repeat,
14378 sp.penalty_freq,
14379 sp.penalty_present,
14380 n_vocab,
14381 )?;
14382 }
14383 let cb2 = col_buf.as_ref().unwrap();
14384 let sc = sctr;
14385 sctr += 1;
14386 // p-stats for the reject column: from col_stats when the col was gathered,
14387 // else (j==0&&base==0) from last_col_stats.
14388 let p_stats = if n_acc > 0 || base == 1 {
14389 // col index within the gathered set == number of gathered cols before n_acc
14390 let gi = if base == 1 { n_acc } else { n_acc - 1 };
14391 col_stats.get(gi).copied().unwrap_or_else(|| {
14392 (0.0, 0.0, 1.0) // unreachable: gathered cols always cover the reject slot
14393 })
14394 } else {
14395 last_col_stats.expect("sampled: last_col_stats unset at reject")
14396 };
14397 let q_stats = draft_stats[n_acc];
14398 if let Some(map) = &d2t_dev {
14399 if q_full_buf.is_none() {
14400 q_full_buf = Some(e.zeros(n_vocab)?);
14401 }
14402 let qf = q_full_buf.as_mut().unwrap();
14403 e.scatter_trim_logits(&q_bufs[n_acc], map, qf, d_vocab, n_vocab)?;
14404 let qf2 = q_full_buf.as_ref().unwrap();
14405 e.residual_sample_filtered(
14406 cb2,
14407 Some(qf2),
14408 n_vocab,
14409 sp_temp,
14410 sp_seed,
14411 sc,
14412 p_stats,
14413 q_stats,
14414 &mut sample_tok,
14415 )?;
14416 } else {
14417 e.residual_sample_filtered(
14418 cb2,
14419 Some(&q_bufs[n_acc]),
14420 n_vocab,
14421 sp_temp,
14422 sp_seed,
14423 sc,
14424 p_stats,
14425 q_stats,
14426 &mut sample_tok,
14427 )?;
14428 }
14429 e.dtoh_u32(&sample_tok)?[0]
14430 };
14431 (
14432 n_acc,
14433 guard_vocab_token(
14434 bonus,
14435 n_vocab,
14436 &format!("sampled verify bonus at round {round} pos={pos} n_acc={n_acc}"),
14437 )?,
14438 )
14439 };
14440 // --- 3b. GRAMMAR TRUNCATION (constrained spec, 2026-08-03): the grammar is
14441 // an extra rejection rule AFTER the exactness verify (the batched-verify-twins
14442 // ordering). Walk the accepted drafts through the grammar in commit order; the
14443 // first illegal token truncates acceptance at its slot, and that slot's emission
14444 // is recomputed as the MASKED argmax of the target's own verify column — token-
14445 // identical to constrained plain greedy decode (an unmasked argmax that is
14446 // grammar-legal IS the masked argmax: masking only removes competitors). The
14447 // column D2H (~1MB) is paid only when a cut fires — the tight-grammar cost,
14448 // measured in acceptance numbers, never hidden.
14449 let (n_acc, bonus) = match constraint.as_deref_mut() {
14450 None => (n_acc, bonus),
14451 Some(c) => {
14452 fn ce(e2: String) -> Box<dyn std::error::Error> {
14453 format!("constraint: {e2}").into()
14454 }
14455 let mut na = n_acc;
14456 let mut cut = false;
14457 for (j, &d) in draft.iter().enumerate().take(n_acc) {
14458 if c.is_allowed(d).map_err(ce)? {
14459 c.consume(d).map_err(ce)?;
14460 } else {
14461 na = j;
14462 cut = true;
14463 dm_cut_tokens += n_acc - j;
14464 break;
14465 }
14466 }
14467 if cut {
14468 dm_cuts += 1;
14469 }
14470 let mut bo = bonus;
14471 if cut || !c.is_allowed(bo).map_err(ce)? {
14472 let mut row = if na == 0 && base == 0 {
14473 init_logits_host
14474 .clone()
14475 .ok_or("constraint: init logits missing (round-0 cut)")?
14476 } else {
14477 e.dtoh_view(
14478 &tlogits_d.slice((base + na - 1) * n_vocab..(base + na) * n_vocab),
14479 )?
14480 };
14481 c.mask_logits(&mut row).map_err(ce)?;
14482 bo = argmax(&row) as u32;
14483 }
14484 c.consume(bo).map_err(ce)?;
14485 (na, bo)
14486 }
14487 };
14488 let mut successor_valid = false;
14489 if let Some((q_proxy, expected_d2)) = rejected_probe {
14490 let v_n = n_acc == 1 && bonus == expected_d2;
14491 eprintln!(
14492 "[opti-controller] shadow q={q_proxy:.6} admitted=false v_n={v_n} \
14493 expected_d2={expected_d2} n_acc={n_acc} bonus={bonus}",
14494 );
14495 }
14496 if let Some(successor) = successor_attempt.as_ref() {
14497 successor_valid = n_acc == 1 && bonus == successor.verify_tokens[0];
14498 let generation = successor.generation;
14499 let q_proxy = successor.q_proxy;
14500 let expected_pending = successor.verify_tokens[0];
14501 let resolution_ms = successor.issued_at.elapsed().as_secs_f64() * 1e3;
14502 let fork = opti_fork
14503 .as_mut()
14504 .ok_or("optipipe successor resolution lost fork state")?;
14505 fork.finish_actual_reconcile(e, &mut *cache, &snap, n_acc, base, successor_valid)?;
14506 if successor_valid {
14507 OPTI_FORK_HITS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
14508 } else {
14509 OPTI_FORK_MISSES.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
14510 OPTI_RECONCILES.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
14511 OPTI_WASTED_DRAFT_TOKENS.fetch_add(2, std::sync::atomic::Ordering::Relaxed);
14512 }
14513 let breaker_tripped = fork
14514 .controller
14515 .as_mut()
14516 .expect("controller policy")
14517 .resolve(successor_valid);
14518 if breaker_tripped {
14519 OPTI_BREAKER_TRIPS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
14520 }
14521 eprintln!(
14522 "[opti-controller] resolve generation={} hit={} q={q_proxy:.6} \
14523 expected_pending={expected_pending} n_acc={n_acc} bonus={bonus} \
14524 resolution_ms={resolution_ms:.3} reconcile={} breaker={}",
14525 generation.id, successor_valid, !successor_valid, breaker_tripped,
14526 );
14527 if !successor_valid {
14528 let mut successor = successor_attempt
14529 .take()
14530 .expect("controller successor disappeared on miss");
14531 successor.settle();
14532 fork.retire(generation)?;
14533 }
14534 }
14535 total_drafted += k_round;
14536 total_accepted += n_acc;
14537 if let Some(t) = sess_telem {
14538 // Greedy, rejection-sampling, and grammar truncation all converge here after
14539 // the accept decision is already on host. Fixed-size relaxed atomics only.
14540 t.record_round(k_round, n_acc);
14541 }
14542 if spec_stats {
14543 st_len_hist[k_round] += 1;
14544 for j in 0..k_round {
14545 st_drafted[j] += 1;
14546 }
14547 for j in 0..n_acc {
14548 st_accepted[j] += 1;
14549 }
14550 if n_acc == k_round {
14551 st_full += 1;
14552 }
14553 }
14554
14555 if debug_spec {
14556 eprintln!(
14557 "[R{round}] pos={pos} out_len={} last_tok={last_token} draft={draft:?} n_acc={n_acc} bonus={bonus} t_pred0={}",
14558 out.len(),
14559 // NOT `t_pred(0)`: `preds` is filled only under `if !sampled` above, so on a
14560 // sampled request round >= 1 (base == 1) indexed an EMPTY vector and PANICKED
14561 // the GPU worker thread — a debug flag that killed the exact regime you would
14562 // set it to investigate. See `debug_t_pred0`.
14563 debug_t_pred0(sampled, base, last_pred, &preds)
14564 );
14565 }
14566
14567 // --- 4. COMMIT: draft[0..n_acc] then bonus (n_acc + 1 tokens) ---
14568 let commit_started = std::time::Instant::now();
14569 // SESSION MODE: every accepted column is already in the CACHE — `out` must carry all
14570 // of them (overshoot past max_new included) or `committed` under-counts the cache rows
14571 // and the next turn's continuation seeds one token off (gate-caught 2026-07-05). The
14572 // single-shot path keeps the cap (its caller truncates + drops the cache anyway).
14573 for j in 0..n_acc {
14574 if !session_mode && out.len() >= max_new {
14575 break;
14576 }
14577 out.push(draft[j]);
14578 }
14579 if pen_on {
14580 pen_hist.extend_from_slice(&draft[0..n_acc]);
14581 pen_hist.push(bonus);
14582 }
14583 let bonus_emitted = session_mode || out.len() < max_new;
14584 if bonus_emitted {
14585 out.push(bonus);
14586 }
14587 last_token = bonus;
14588
14589 // --- 5. ROLLBACK + advance (§C) ---
14590 if n_acc == k_round && !spec_replay {
14591 // FULL ACCEPT, BONUS FOLD: all verify columns (pending? + drafts) are committed in
14592 // cache; the NEW bonus stays PENDING for the next round's verify batch — NO extra
14593 // T=1 trunk pass. The next draft chain seeds from the MTP block's h_nextn at the
14594 // bonus position: one MTP-block pass (~1/33 trunk cost) replaces the trunk read.
14595 // last_pred is dead in the pending path (t_pred reads verify col 0).
14596 //
14597 // PERSISTENT DRAFT KV, full-accept fill: the chain covered last_token +
14598 // draft[0..k_round-2] as INPUTS (slots P..P'-2); draft[k_round-1] (slot P'-1) was
14599 // only ever an output, so its entry is MISSING. Fill it from vh_seed — its EXACT
14600 // trunk hidden (the last verify column). set_len first: a p-min break may have
14601 // left one extra chain append at that slot. Partial accepts need NO fill (the
14602 // chain already covered every accepted position; round-start set_len truncates).
14603 let mut vh_seed = e.zeros(n_embd)?;
14604 e.copy_view_into(
14605 &mut vh_seed,
14606 0,
14607 &vx.slice((t_v - 1) * n_embd..t_v * n_embd),
14608 n_embd,
14609 )?;
14610 if refresh {
14611 // TRUE-HIDDEN REFRESH (2026-07-03, the HANDOVER-listed acceptance lever):
14612 // overwrite ALL committed positions' scratch entries with K/V from their EXACT
14613 // verify hiddens — the reference engine's mtp_update fills from true hiddens;
14614 // the full stack (vx) is already resident from the verify. Replaces both the
14615 // chain-approximate entries AND the old last-token-only fill. Acceptance-only
14616 // (draft attention quality); exactness stays the verify's job.
14617 scratch.set_len(e, pos)?;
14618 // PREDECESSOR pairing: row i gets vx[i-1]; row 0 the carried fill_prev
14619 // (hidden of the last committed row before this verify batch).
14620 let mut vxs = e.zeros(t_v * n_embd)?;
14621 e.copy_into(&mut vxs, 0, &fill_prev, n_embd)?;
14622 if t_v > 1 {
14623 e.copy_view_into(
14624 &mut vxs,
14625 n_embd,
14626 &vx.slice(0..(t_v - 1) * n_embd),
14627 (t_v - 1) * n_embd,
14628 )?;
14629 }
14630 self.mtp_kv_fill_all(e, &verify_tokens, &vxs, pos, &mut *scratch, embd_dev)?;
14631 } else {
14632 scratch.set_len(e, pos + base + k_round - 1)?;
14633 // predecessor of the last draft = verify col t_v-2 (or fill_prev at t_v==1)
14634 let mut hp = e.zeros(n_embd)?;
14635 if t_v >= 2 {
14636 e.copy_view_into(
14637 &mut hp,
14638 0,
14639 &vx.slice((t_v - 2) * n_embd..(t_v - 1) * n_embd),
14640 n_embd,
14641 )?;
14642 } else {
14643 e.copy_into(&mut hp, 0, &fill_prev, n_embd)?;
14644 }
14645 self.mtp_kv_fill_all(
14646 e,
14647 &[draft[k_round - 1]],
14648 &hp,
14649 pos + base + k_round - 1,
14650 &mut *scratch,
14651 embd_dev,
14652 )?;
14653 }
14654 // REFERENCE SEEDING: no pseudo pass — the next chain's step 0 IS the
14655 // reference's (id_last, h_prev) draft row; it appends the bonus's scratch
14656 // entry itself. Seed = TRUE hidden of the bonus's predecessor (last verify
14657 // col). Saves one MTP-block pass per round on top of the pairing fix.
14658 if !devacc_seeded {
14659 e.copy_into(&mut h_seed_buf, 0, &vh_seed, n_embd)?;
14660 e.copy_into(&mut fill_prev, 0, &vh_seed, n_embd)?;
14661 }
14662 pending = Some(bonus);
14663 if debug_spec {
14664 eprintln!(" -> FULL ACCEPT (bonus pending, prev-h seed)");
14665 }
14666 } else if !spec_replay && base + n_acc >= 1 {
14667 // PARTIAL ACCEPT, REPLAY-FREE (2026-07-03 — the profiled #1 long-ctx spec cost):
14668 // the verify's first j = base+n_acc columns ARE the committed sequence, computed
14669 // bit-identically to eager (decode-exact contract) — so KEEP them: KV truncates to
14670 // pos+j, recurrent state rebuilds from the VerifyCkpt (same-kernel gdn prefix
14671 // re-run / pure state-clone restore), and the bonus stays PENDING exactly like the
14672 // full-accept path — the legacy duplicate trunk replay is gone. The next chain
14673 // seeds from the MTP pseudo-hidden of the bonus, whose seed = the TRUE verify
14674 // hidden of its predecessor (col j-1) — same one-hop pseudo structure as full
14675 // accept (never compounds: the next verify recomputes true hiddens for all
14676 // committed columns).
14677 let j = base + n_acc;
14678 // VERIFY-GRAPH SLAB COMMIT: when the captured trunk ran, the linear layers'
14679 // column stash was written into the graphs ctx's persistent slabs as in-graph
14680 // memcpy nodes, NOT into the per-column VerifyCkpt the cols arm reads — so the
14681 // commit must take the slab twin (same semantics, slab-addressed sources). The
14682 // ctx states which of the two this round produced via `round_slab`; trusting the
14683 // flag rather than the env keeps a round that fell back to the eager walk (a
14684 // capture that declined, a t the pool never captured) on the cols arm.
14685 let slab_commit = vg_guard
14686 .as_ref()
14687 .and_then(|g| g.as_ref())
14688 .map(|g| g.round_slab)
14689 .unwrap_or(false);
14690 if slab_commit {
14691 self.dspark_commit_prefix_slab(
14692 e,
14693 &mut *cache,
14694 &snap,
14695 vg_guard
14696 .as_ref()
14697 .and_then(|g| g.as_ref())
14698 .expect("slab_commit implies a graphs ctx"),
14699 j,
14700 )?;
14701 } else {
14702 self.commit_verified_prefix(
14703 e,
14704 &mut *cache,
14705 &snap,
14706 ckpt.as_ref().unwrap(),
14707 j,
14708 devacc_seeded,
14709 if devacc_seeded {
14710 devacc_acc.as_ref().map(|a| (a, base, t_v))
14711 } else {
14712 None
14713 },
14714 )?;
14715 }
14716 let mut seed = e.zeros(n_embd)?;
14717 e.copy_view_into(
14718 &mut seed,
14719 0,
14720 &vx.slice((j - 1) * n_embd..j * n_embd),
14721 n_embd,
14722 )?;
14723 // Draft scratch: TRUE-HIDDEN REFRESH of the committed prefix (see the full-accept
14724 // branch); without it the chain entries stand and only the tail truncates. Either
14725 // way len ends at pos+j so the pseudo append lands at the bonus's slot pos+j
14726 // (persistent mode), rope pos+j+1 (chain convention).
14727 if refresh {
14728 scratch.set_len(e, pos)?;
14729 let mut vxs = e.zeros(j * n_embd)?;
14730 e.copy_into(&mut vxs, 0, &fill_prev, n_embd)?;
14731 if j > 1 {
14732 e.copy_view_into(
14733 &mut vxs,
14734 n_embd,
14735 &vx.slice(0..(j - 1) * n_embd),
14736 (j - 1) * n_embd,
14737 )?;
14738 }
14739 self.mtp_kv_fill_all(
14740 e,
14741 &verify_tokens[0..j],
14742 &vxs,
14743 pos,
14744 &mut *scratch,
14745 embd_dev,
14746 )?;
14747 } else {
14748 scratch.set_len(e, pos + j)?;
14749 }
14750 // REFERENCE SEEDING (see the full-accept branch): seed = TRUE hidden of the
14751 // bonus's predecessor (verify col j-1); no pseudo pass.
14752 if !devacc_seeded {
14753 e.copy_into(&mut h_seed_buf, 0, &seed, n_embd)?;
14754 e.copy_into(&mut fill_prev, 0, &seed, n_embd)?;
14755 }
14756 pending = Some(bonus);
14757 if debug_spec {
14758 eprintln!(" -> PARTIAL(replay-free j={j}, bonus pending, prev-h seed)");
14759 }
14760 } else if !spec_replay {
14761 // ZERO ROUND FOLD (2026-07-10, verify-cost target #3): base+n_acc == 0 — a
14762 // pending-less round where nothing was accepted (PMIN0 zero-draft chains after a
14763 // replay/commit, or plain 0-accept rounds at round 0). The old path replayed
14764 // [bonus] through a FULL m=1 trunk+head forward (the 489us full-vocab head pass
14765 // measured at ~0.75/round on PMIN0 configs). Instead: restore the pre-round
14766 // snapshot and let the bonus ride the NEXT round's verify as col 0 — the existing
14767 // base=1 pending machinery, bit-identical by the decode-exact verify contract.
14768 // Seed: the bonus's predecessor is the last COMMITTED token, whose hidden
14769 // fill_prev already carries (same seeding as the 1-token-replay case it replaces).
14770 cache.rollback(e, &snap, 0)?;
14771 scratch.set_len(e, pos)?;
14772 e.copy_into(&mut h_seed_buf, 0, &fill_prev, n_embd)?;
14773 pending = Some(bonus);
14774 if debug_spec {
14775 eprintln!(" -> ZERO-ROUND FOLD (bonus pending, fill_prev seed)");
14776 }
14777 } else {
14778 // PARTIAL ACCEPT, LEGACY REPLAY (seam MEMRA_SPEC_REPLAY=1 — or j==0: nothing of
14779 // this round survives, only possible before the first pending exists, ~round 0):
14780 // restore EVERYTHING to the pre-round snapshot (KV truncate to pos + recur
14781 // restore), then replay the committed prefix pending? ++ draft[0..n_acc] ++
14782 // [bonus] as ONE batched T forward — single weight read, bit-identical to greedy
14783 // (the verify-all-columns path is the same math). Commits the bonus with a TRUE
14784 // trunk hidden.
14785 cache.rollback(e, &snap, 0)?; // accept_len=0: KV len = pos, recur = snapshot
14786 let mut replay: Vec<u32> = Vec::with_capacity(base + n_acc + 1);
14787 if let Some(b) = pending.take() {
14788 replay.push(b);
14789 }
14790 replay.extend_from_slice(&draft[0..n_acc]);
14791 replay.push(bonus);
14792 // Full-stack forward (decode_step_t_core = decode_step_t_h_emb_dev's body):
14793 // Predecessor pairing seeds from the PREDECESSOR row (col len-2) — the same-row path takes the
14794 // last col exactly as before (byte-identical to the old _h_emb_dev call).
14795 let (rl_d, rx) = if self.batched_serving_numeric_class() {
14796 let mut logits = Vec::with_capacity(replay.len() * n_vocab);
14797 let mut hidden = e.uninit(replay.len() * n_embd)?;
14798 for (row, &token) in replay.iter().enumerate() {
14799 let (row_logits, row_hidden) =
14800 self.spec_target_step_h(e, token, &mut *cache)?;
14801 logits.extend_from_slice(&row_logits);
14802 e.dtod_copy_into(&row_hidden, &mut hidden, row * n_embd)?;
14803 }
14804 (e.htod(&logits)?, hidden)
14805 } else {
14806 self.decode_step_t_core(e, &replay, pos, &mut *cache, embd_dev, None)?
14807 };
14808 // last_pred = argmax of the LAST column's logits (predicts the token after `bonus`)
14809 // — device argmax + one 4-byte read instead of the full-vocab column dtoh.
14810 e.argmax_token_device_col(&rl_d, replay.len() - 1, n_vocab, &mut preds_d, 0)?;
14811 last_pred = guard_vocab_token(
14812 e.dtoh_u32(&preds_d)?[0],
14813 n_vocab,
14814 &format!("replay last_pred at round {round} pos={pos}"),
14815 )?;
14816 if sampled {
14817 let lr0 = replay.len();
14818 let lc = last_col_logits
14819 .as_mut()
14820 .expect("sampled: last_col_logits unset");
14821 e.copy_view_into(
14822 lc,
14823 0,
14824 &rl_d.slice((lr0 - 1) * n_vocab..lr0 * n_vocab),
14825 n_vocab,
14826 )?;
14827 }
14828 let lr = replay.len();
14829 if lr >= 2 {
14830 e.copy_view_into(
14831 &mut h_seed_buf,
14832 0,
14833 &rx.slice((lr - 2) * n_embd..(lr - 1) * n_embd),
14834 n_embd,
14835 )?;
14836 } else {
14837 // 1-token replay (round-0 miss): the bonus's predecessor is the OLD
14838 // last_token, whose own-row hidden fill_prev still holds.
14839 e.copy_into(&mut h_seed_buf, 0, &fill_prev, n_embd)?;
14840 }
14841 // the bonus is COMMITTED here — it becomes the last committed row.
14842 let mut rh_last = e.zeros(n_embd)?;
14843 e.copy_view_into(
14844 &mut rh_last,
14845 0,
14846 &rx.slice((lr - 1) * n_embd..lr * n_embd),
14847 n_embd,
14848 )?;
14849 e.copy_into(&mut fill_prev, 0, &rh_last, n_embd)?;
14850 if debug_spec {
14851 eprintln!(" -> PARTIAL(replay={replay:?}), next_pred={last_pred}");
14852 }
14853 }
14854 if devacc_seeded {
14855 // stage (b) epilogue: fill_prev takes the gathered seed AFTER the refresh fills
14856 // consumed the old value (both slots carry the same value in every non-replay arm).
14857 e.copy_into(&mut fill_prev, 0, &h_seed_buf, n_embd)?;
14858 }
14859 if successor_valid {
14860 let optimistic_scratch_len = successor_attempt
14861 .as_ref()
14862 .expect("valid controller successor disappeared")
14863 .scratch_len;
14864 // The normal current-round commit refreshed/truncated the logical scratch tail.
14865 // Its optimistic successor row was already written physically, so restoring only
14866 // the retained logical length makes that row live for the carried round.
14867 scratch.set_len(e, optimistic_scratch_len)?;
14868 }
14869 if let Some(current) = current_opti.take() {
14870 opti_fork
14871 .as_mut()
14872 .ok_or("optipipe current retirement lost fork state")?
14873 .retire(current.generation)?;
14874 }
14875 if successor_valid {
14876 let successor = successor_attempt
14877 .take()
14878 .expect("valid controller successor disappeared before promotion");
14879 let generation = successor.generation;
14880 opti_fork
14881 .as_mut()
14882 .ok_or("optipipe successor promotion lost fork state")?
14883 .promote_successor_snapshot(&mut snap, generation);
14884 carried_opti = Some(successor);
14885 }
14886 if anatomy_on {
14887 // Commit/rollback is normally asynchronous on the primary/head stream. Bound it
14888 // only for this diagnostic so it does not disappear into the following draft's
14889 // first token readback.
14890 e.stream().synchronize()?;
14891 ph_commit += commit_started.elapsed().as_secs_f64();
14892 }
14893 // adaptive-K update (host math, zero syncs): next round drafts accepted-run + 1,
14894 // clamped to [floor(pos), k_cap]. cache.pos is post-rollback here (the round's
14895 // final position — the floor's position key reads the committed depth). Burst
14896 // rounds (`continue` above) draft the captured fixed depth and skip this, exactly
14897 // like gemma's burst arm.
14898 if adapt {
14899 let fl_now = floor_at(cache.pos);
14900 kc = (n_acc + 1).clamp(fl_now.min(k_cap), k_cap);
14901 }
14902 ph_mark(&mut ph_rest, phase_on);
14903 if let Some(p) = pipe {
14904 p.accept_end(round);
14905 }
14906 drop(pipe_accept);
14907 if let Some(t0) = round_t0 {
14908 let ms = t0.elapsed().as_secs_f64() * 1e3;
14909 ROUND_MS.fetch_add((ms * 1e3) as u64, std::sync::atomic::Ordering::Relaxed);
14910 let n = ROUND_N.fetch_add(1, std::sync::atomic::Ordering::Relaxed) + 1;
14911 if n % 32 == 0 {
14912 eprintln!(
14913 "[spec-round] rounds={n} avg round wall={:.2} ms (emitted={} drafted so far)",
14914 ROUND_MS.load(std::sync::atomic::Ordering::Relaxed) as f64 / 1e3 / n as f64,
14915 out.len()
14916 );
14917 }
14918 }
14919 round += 1;
14920 // sse-cadence: this round's accepted drafts + bonus are committed (out is
14921 // append-only past step 4) — flush at round cadence.
14922 keep_going = flush_commit(&mut on_commit, &out, &mut flushed);
14923 }
14924 if let Some(mut ticket) = carried_opti.take() {
14925 opti_fork
14926 .as_mut()
14927 .ok_or("optipipe tail drain lost fork state")?
14928 .cancel_controller_ticket(e, &mut *cache, &mut *scratch, &snap, &mut ticket)?;
14929 }
14930 // sse-cadence: nothing below appends to `out`; flush any remainder (defensive).
14931 // (verdict ignored — the burst is over either way; the session tail runs unchanged.)
14932 let _ = flush_commit(&mut on_commit, &out, &mut flushed);
14933
14934 if spec_stats {
14935 let per_slot: Vec<String> = (0..k)
14936 .map(|j| {
14937 if st_drafted[j] > 0 {
14938 format!(
14939 "{}/{}={:.3}",
14940 st_accepted[j],
14941 st_drafted[j],
14942 st_accepted[j] as f64 / st_drafted[j] as f64
14943 )
14944 } else {
14945 "0/0".into()
14946 }
14947 })
14948 .collect();
14949 let acc = if total_drafted > 0 {
14950 total_accepted as f64 / total_drafted as f64
14951 } else {
14952 0.0
14953 };
14954 eprintln!(
14955 "[spec-stats] rounds={round} full_accept={st_full} len_hist={st_len_hist:?} \
14956 per_slot=[{}] total={total_accepted}/{total_drafted}={acc:.3} \
14957 tok_per_round={:.3}",
14958 per_slot.join(" "),
14959 (total_accepted + round) as f64 / round.max(1) as f64
14960 );
14961 }
14962 if constraint.is_some() {
14963 eprintln!(
14964 "[draft-mask] mask_rounds={dm_rounds} clone_total={:.3}ms \
14965 clone_per_round={:.4}ms gram_cuts={dm_cuts}/{round} cut_tokens={dm_cut_tokens}",
14966 dm_clone_ns as f64 / 1e6,
14967 dm_clone_ns as f64 / 1e6 / dm_rounds.max(1) as f64
14968 );
14969 }
14970 if phase_on {
14971 let tot = ph_draft + ph_verify + ph_wait + ph_rest;
14972 eprintln!(
14973 "[spec-phase] draft={:.1}ms ({:.1}%) verify-issue={:.1}ms ({:.1}%) verify-wait={:.1}ms ({:.1}%) commit-host={:.1}ms ({:.1}%) rounds={round}",
14974 ph_draft * 1e3,
14975 ph_draft / tot * 100.0,
14976 ph_verify * 1e3,
14977 ph_verify / tot * 100.0,
14978 ph_wait * 1e3,
14979 ph_wait / tot * 100.0,
14980 ph_rest * 1e3,
14981 ph_rest / tot * 100.0
14982 );
14983 }
14984 if anatomy_on {
14985 let rounds_f = round.max(1) as f64;
14986 let other = (ph_rest - ph_commit).max(0.0);
14987 eprintln!(
14988 "[spec-anatomy] per-round draft={:.3}ms pp-verify={:.3}ms \
14989 verify-accept={:.3}ms commit-rollback={:.3}ms other={:.3}ms rounds={round}",
14990 ph_draft * 1e3 / rounds_f,
14991 ph_verify * 1e3 / rounds_f,
14992 ph_wait * 1e3 / rounds_f,
14993 ph_commit * 1e3 / rounds_f,
14994 other * 1e3 / rounds_f,
14995 );
14996 }
14997 let _pipe_tail = pipe.map(|p| p.primary());
14998 // SESSION TAIL: leave the session in the exact invariant the next turn's suffix prime
14999 // expects — every row in `committed` has trunk KV/recur state AND an exact draft-KV row.
15000 // Park the draft-graph ctx back on the session (the serve-burst fixed-cost fix): the next
15001 // burst replays instead of recapturing. Error paths (`?` above) drop it — recaptured then.
15002 if let Some(slot) = sess_draft_slot.take() {
15003 *slot = Some(dctx);
15004 }
15005 let t_rounds = t_ent.elapsed();
15006 if let Some((committed, last_h, next_pred_slot, sctr_slot, uctr_slot)) = sess_tail.take() {
15007 // NEXT BURST'S BOUNDARY TOKEN (lane/sampled-spec-quality, Item 1). Greedy stashes
15008 // the argmax `last_pred` exactly as before (byte contract). SAMPLED draws the token
15009 // HERE, where the sampler, the session Philox counters and the penalty window are
15010 // all live and the boundary logits row still exists — that is the "make the state
15011 // available" half of the fix; the consuming burst then just emits it. `sctr` is
15012 // written to the session BELOW the draws so the advance is never lost.
15013 *next_pred_slot = Some(last_pred);
15014 let sample_boundary = sampled && constraint.is_none() && spec_sampled_boundary_on();
15015 let mut stashed_pending = false;
15016 if let Some(b) = pending.take() {
15017 if !sampled {
15018 // PENDING-CARRY (2026-08-01): stash the bonus on the session instead of
15019 // committing it with a solo T=1 pass — the next empty-suffix greedy burst
15020 // consumes it as round-0 verify col 0 (a plain round edge; the old tail
15021 // commit + next burst's init feed were 11.6+11.5ms solo trunk passes per
15022 // burst on H100 q27, [spec-setup] trace). b stays in `out` (emitted) but
15023 // OUT of `committed` (cache rows == committed); the consuming call
15024 // prepends it once its verify commits the row. next_pred is unknowable
15025 // without the commit pass — None; callers gate on pending_tok too.
15026 debug_assert_eq!(out.last(), Some(&b), "pending must be the last emitted");
15027 if let Some(slot) = sess_pending_slot.take() {
15028 *slot = Some(b);
15029 }
15030 *next_pred_slot = None;
15031 // fill_prev = hidden of the last COMMITTED row (b's predecessor) — the
15032 // exact chain-seed/fill anchor the consuming burst (or a flush) needs.
15033 *last_h = Some(e.clone_dtod(&fill_prev)?);
15034 stashed_pending = true;
15035 } else {
15036 // SAMPLED tail (unchanged): commit the bonus (one T=1 pass) + draft fill —
15037 // the sampled round-0 accept needs this pass's logits (last_col_logits).
15038 let pos_b = cache.pos;
15039 scratch.set_len(e, pos_b)?;
15040 let (lg_b, hb) = self.spec_target_step_h(e, b, &mut *cache)?;
15041 // after a FULL-accept exit `last_pred` is STALE (it predicted the bonus
15042 // itself — the prediction AFTER the bonus never materialized; it would have
15043 // been the next round's verify col 0). The commit's logits ARE that
15044 // prediction — so they are also the row the next burst's boundary token
15045 // comes off, and (lane/sampled-spec-quality) it is DRAWN from them here.
15046 *next_pred_slot = Some(if sample_boundary {
15047 sample_boundary_token(
15048 e,
15049 &lg_b,
15050 &sp,
15051 &pen_hist,
15052 &mut sctr,
15053 "burst-tail-commit",
15054 )?
15055 } else {
15056 argmax(&lg_b) as u32
15057 });
15058 self.mtp_kv_fill_all(e, &[b], &fill_prev, pos_b, &mut *scratch, embd_dev)?;
15059 *last_h = Some(hb);
15060 }
15061 } else {
15062 // fill_prev tracks the hidden of the last COMMITTED row throughout the loop.
15063 *last_h = Some(e.clone_dtod(&fill_prev)?);
15064 if sample_boundary {
15065 // No pending to commit, so the boundary row is the one `last_pred` was
15066 // argmaxed from and the sampled path keeps it on device: the init feed's
15067 // logits when the burst ran zero rounds, else the legacy-replay path's
15068 // last verify column (both predict the token AFTER the last committed
15069 // row). It is retained precisely because round 0's accept test needs it,
15070 // so the draw costs no extra D2H of the [n_vocab] row.
15071 match last_col_logits.as_ref() {
15072 Some(lc) => {
15073 *next_pred_slot = Some(sample_boundary_token_dev(
15074 e,
15075 lc,
15076 n_vocab,
15077 &sp,
15078 &pen_hist,
15079 &mut sctr,
15080 "burst-tail-nopending",
15081 )?);
15082 }
15083 // NAME THE FALLBACK (house standard): unreachable today — a sampled
15084 // burst always feeds or replays, so the row exists — but if it ever
15085 // is, the stream takes a greedy token and SAYS so rather than
15086 // silently regressing to the pre-lane behaviour.
15087 None => eprintln!(
15088 "[spec-boundary] sampled tail kept the ARGMAX boundary token \
15089 (reason: no retained boundary logits row)"
15090 ),
15091 }
15092 }
15093 }
15094 *sctr_slot = sctr;
15095 *uctr_slot = uctr;
15096 committed.extend_from_slice(prompt);
15097 if let Some(cb) = carried_pending {
15098 // the consumed carry's cache row landed in round 0's verify (every pending
15099 // round commits col 0) — it joins `committed` here, in sequence order.
15100 committed.push(cb);
15101 }
15102 if stashed_pending {
15103 // ZERO-EMIT BURST (2026-08-06 c=8 serve panic, pre-existing since b4aea184):
15104 // `out.len() - 1` underflowed on an EMPTY `out` — "range end index
15105 // 18446744073709551615 out of range for slice of length 0", killing the
15106 // memra-gpu-worker and failing 31 of 32 concurrent requests with "worker closed
15107 // stream". Reachable because `pending` starts as `carried_pending` (a bonus
15108 // stashed by the PREVIOUS burst) while `out` starts empty, and the carry is
15109 // deliberately NOT pushed to `out` (line ~3239: the burst that emitted it already
15110 // did). So a burst that stashes a pending without emitting anything of its own —
15111 // the round loop exits before a push, e.g. the ring drain's `out.len() < max_new`
15112 // guard skipping every token under a tight budget — arrives here with
15113 // out.len() == 0 and stashed_pending == true.
15114 //
15115 // The invariant is unchanged: `committed` gets every emitted token EXCEPT the
15116 // stashed bonus. With nothing emitted, that is nothing — and the carry pushed
15117 // just above is already accounted. Saturating, not a min/assert: an empty `out`
15118 // here is a legitimate burst shape, not a corrupt state.
15119 let emitted = out.len().saturating_sub(1);
15120 committed.extend_from_slice(&out[..emitted]);
15121 } else {
15122 committed.extend_from_slice(&out); // FULL out incl. overshoot — all committed
15123 }
15124 debug_assert_eq!(
15125 cache.pos,
15126 committed.len(),
15127 "session invariant: cache rows == committed tokens"
15128 );
15129 if setup_trace {
15130 e.stream().synchronize()?; // bound the async tail fill in the trace
15131 let t_tail = t_ent.elapsed();
15132 eprintln!(
15133 "[spec-setup] init={:.2}ms cap={:.2}ms fill={:.2}ms rounds={:.2}ms tail={:.2}ms total={:.2}ms out={} cont={}",
15134 t_init.as_secs_f64() * 1e3,
15135 (t_cap - t_init).as_secs_f64() * 1e3,
15136 (t_fill - t_cap).as_secs_f64() * 1e3,
15137 (t_rounds - t_fill).as_secs_f64() * 1e3,
15138 (t_tail - t_rounds).as_secs_f64() * 1e3,
15139 t_tail.as_secs_f64() * 1e3,
15140 out.len(),
15141 continuation
15142 );
15143 }
15144 return Ok((out, total_drafted, total_accepted));
15145 }
15146 out.truncate(max_new);
15147 Ok((out, total_drafted, total_accepted))
15148 }
15149
15150 /// Anchor-bounded DSpark target extraction. The trunk sees the exact generated token tape;
15151 /// only requested hidden rows and target-logit rows cross PCIe. An anchor token at p pairs
15152 /// with the pre-output-norm h[p-1] carrier, exactly as the existing replay/NextN path does.
15153 pub fn extract_dspark_anchors(
15154 &self,
15155 e: &Engine,
15156 tokens: &[u32],
15157 anchor_positions: &[usize],
15158 gamma: usize,
15159 top_k: usize,
15160 chunk: usize,
15161 temperature: f32,
15162 ) -> Result<Vec<DsparkAnchorRecord>, Box<dyn std::error::Error>> {
15163 if tokens.len() < gamma + 2 || gamma == 0 || chunk < 2 {
15164 return Err("DSpark extraction token tape/gamma/chunk is invalid".into());
15165 }
15166 if anchor_positions.windows(2).any(|pair| pair[0] >= pair[1]) {
15167 return Err("DSpark anchor positions must be sorted and unique".into());
15168 }
15169 for &position in anchor_positions {
15170 if position == 0 || position + gamma >= tokens.len() {
15171 return Err(format!(
15172 "DSpark anchor {position} has no predecessor or cannot cover gamma={gamma} in {} tokens",
15173 tokens.len()
15174 )
15175 .into());
15176 }
15177 }
15178
15179 let n_vocab = self.output.out_features();
15180 let n_embd = self.cfg.n_embd as usize;
15181 let mut cache =
15182 crate::pp::new_cache_planned(e, &self.cfg, &self.plan, tokens.len() + gamma + 8)?;
15183 let (embd_qt, embd_rb) = self.embd.qt_and_row_bytes(n_embd);
15184 let embd_gpu = if spec_host_embd() {
15185 None
15186 } else {
15187 Some(
15188 self.embd_gpu
15189 .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload")),
15190 )
15191 };
15192 let embd_dev = embd_gpu.map(|gpu| (gpu, embd_qt, embd_rb));
15193
15194 struct PendingRecord {
15195 position: usize,
15196 hidden: Option<Vec<f32>>,
15197 tokens: Vec<u32>,
15198 target_top_ids: Vec<Option<Vec<u32>>>,
15199 target_top_logits: Vec<Option<Vec<f32>>>,
15200 target_top_probs: Vec<Option<Vec<f32>>>,
15201 target_tail_probs: Vec<Option<f32>>,
15202 }
15203
15204 let mut pending: Vec<PendingRecord> = anchor_positions
15205 .iter()
15206 .map(|&position| PendingRecord {
15207 position,
15208 hidden: None,
15209 tokens: tokens[position..=position + gamma].to_vec(),
15210 target_top_ids: vec![None; gamma],
15211 target_top_logits: vec![None; gamma],
15212 target_top_probs: vec![None; gamma],
15213 target_tail_probs: vec![None; gamma],
15214 })
15215 .collect();
15216
15217 let mut start = 0usize;
15218 while start < tokens.len() {
15219 let end = (start + chunk).min(tokens.len());
15220 let chunk_tokens = &tokens[start..end];
15221 let (target_logits, hidden_rows) =
15222 self.decode_step_t_core(e, chunk_tokens, start, &mut cache, embd_dev, None)?;
15223 for record in &mut pending {
15224 let hidden_position = record.position - 1;
15225 if hidden_position >= start && hidden_position < end {
15226 let local = hidden_position - start;
15227 record.hidden = Some(
15228 e.dtoh_view(&hidden_rows.slice(local * n_embd..(local + 1) * n_embd))?,
15229 );
15230 }
15231 for slot in 0..gamma {
15232 let target_row = record.position + slot;
15233 if target_row < start || target_row >= end {
15234 continue;
15235 }
15236 let local = target_row - start;
15237 let logits =
15238 e.dtoh_view(&target_logits.slice(local * n_vocab..(local + 1) * n_vocab))?;
15239 let (ids, top_logits, probs, tail) =
15240 dspark_sparse_softmax_topk(&logits, top_k, temperature)?;
15241 record.target_top_ids[slot] = Some(ids);
15242 record.target_top_logits[slot] = Some(top_logits);
15243 record.target_top_probs[slot] = Some(probs);
15244 record.target_tail_probs[slot] = Some(tail);
15245 }
15246 }
15247 start = end;
15248 }
15249
15250 pending
15251 .into_iter()
15252 .map(|record| {
15253 let hidden = record
15254 .hidden
15255 .ok_or_else(|| format!("missing DSpark hidden at {}", record.position))?;
15256 let target_top_ids =
15257 flatten_dspark_rows(record.target_top_ids, record.position, "target ids")?;
15258 let target_top_logits = flatten_dspark_rows(
15259 record.target_top_logits,
15260 record.position,
15261 "target logits",
15262 )?;
15263 let target_top_probs =
15264 flatten_dspark_rows(record.target_top_probs, record.position, "target probs")?;
15265 let target_tail_probs = record
15266 .target_tail_probs
15267 .into_iter()
15268 .enumerate()
15269 .map(|(slot, value)| {
15270 value.ok_or_else(|| {
15271 format!("missing DSpark tail at {} slot {slot}", record.position)
15272 })
15273 })
15274 .collect::<Result<Vec<_>, _>>()?;
15275 Ok(DsparkAnchorRecord {
15276 position: record.position,
15277 hidden,
15278 tokens: record.tokens,
15279 target_top_ids,
15280 target_top_logits,
15281 target_top_probs,
15282 target_tail_probs,
15283 })
15284 })
15285 .collect()
15286 }
15287
15288 /// TEACHER-FORCED REPLAY ACCEPTANCE (hqmtp MTP-heal protocol): walk a FIXED token
15289 /// sequence and, at sampled positions, compare the MTP head's K-token draft chain against
15290 /// the trunk's own teacher-forced greedy predictions. Nothing is generated — the context is
15291 /// the corpus text itself, so (a) degenerate self-generated loops cannot inflate acceptance
15292 /// and (b) two arms (bf16 ceiling vs NVFP4) score on IDENTICAL contexts, isolating the
15293 /// quant-induced head/hidden-state mismatch from text drift.
15294 ///
15295 /// Per eval position p (context = tokens[0..=p], predecessor pairing as in spec decode):
15296 /// draft_j = chain token j from (tokens[p], h_{p-1}), then its own drafts — the exact
15297 /// eager spec-decode chain (same mtp_head_forward_dev, same rope positions).
15298 /// target_j = teacher-forced greedy pick for position p+1+j (argmax of the trunk logits
15299 /// at forced context tokens[0..p+j]). For j==0 this equals live spec
15300 /// acceptance; for j>=1 live verify would condition on the drafts, here it
15301 /// conditions on the corpus — deterministic and arm-comparable by design.
15302 ///
15303 /// Returns (rows, bg): one (p, drafts[k], targets[k]) row per eval position (ascending p),
15304 /// plus the full teacher-forced greedy track bg (bg[i] = greedy pick for position i, i>=1)
15305 /// so harnesses can cross-check runs (e.g. different chunk sizes must give identical bg).
15306 ///
15307 /// `hdump`: when Some, every position's pre-output_norm trunk hidden (the exact rows the
15308 /// draft-KV fill pairs from) streams to the file as little-endian f32 [t_total, n_embd] —
15309 /// the head-distillation extraction (hqmtp): the ENGINE is the source of truth for trunk
15310 /// hiddens (HF torch reproductions of the hybrid trunk measured only ~0.5 greedy
15311 /// agreement vs this path — not usable as a training-data source).
15312 pub fn replay_acceptance(
15313 &self,
15314 e: &Engine,
15315 tokens: &[u32],
15316 k: usize,
15317 stride: usize,
15318 chunk: usize,
15319 mut hdump: Option<&mut std::fs::File>,
15320 ) -> Result<(Vec<(usize, Vec<u32>, Vec<u32>)>, Vec<u32>), Box<dyn std::error::Error>> {
15321 assert!(k >= 1 && stride >= 1 && chunk >= 2);
15322 let mtp = self
15323 .mtp
15324 .as_ref()
15325 .expect("replay_acceptance requires an MTP head");
15326 let n_vocab = self.output.out_features();
15327 let d_vocab = mtp
15328 .shared_head_head
15329 .as_ref()
15330 .unwrap_or(&self.output)
15331 .out_features();
15332 let n_embd = self.cfg.n_embd as usize;
15333 let t_total = tokens.len();
15334 assert!(t_total >= 8, "corpus too short ({t_total} tokens)");
15335 // STAGE-OWNED KV (lane/pp2-spec 2026-08-06) — see `new_session`. Door shut = `Cache::new`.
15336 let mut cache = crate::pp::new_cache_planned(e, &self.cfg, &self.plan, t_total + k + 8)?;
15337 let mut scratch = self.new_mtp_scratch(e, t_total + k + 8)?;
15338 let (embd_qt, embd_rb) = self.embd.qt_and_row_bytes(n_embd);
15339 let embd_gpu = if spec_host_embd() {
15340 None
15341 } else {
15342 Some(
15343 self.embd_gpu
15344 .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload")),
15345 )
15346 };
15347 let embd_dev = embd_gpu.map(|g| (g, embd_qt, embd_rb));
15348
15349 // bg[i] = the trunk's greedy pick for position i under the forced context (i >= 1).
15350 let mut bg: Vec<u32> = vec![0; t_total + 1];
15351 let mut rows: Vec<(usize, Vec<u32>, Vec<u32>)> = Vec::new();
15352 let mut prev_last_h = e.zeros(n_embd)?; // predecessor hidden entering the chunk
15353 let mut seed_buf = e.zeros(n_embd)?;
15354 let mut preds_d = e.alloc_u32_zeroed(chunk)?;
15355 let nll_on = std::env::var("MEMRA_REPLAY_NLL").as_deref() == Ok("1");
15356 let (mut nll_sum, mut nll_cnt) = (0f64, 0u64);
15357 let mut s = 0usize;
15358 while s < t_total {
15359 let cend = (s + chunk).min(t_total);
15360 let tc = cend - s;
15361 let ch = &tokens[s..cend];
15362 // 1. forced trunk pass — verify path (decode-exact contract): all-column logits +
15363 // the chunk's true hiddens.
15364 let (tl_d, vx) = self.decode_step_t_core(e, ch, s, &mut cache, embd_dev, None)?;
15365 for j in 0..tc {
15366 e.argmax_token_device_col(&tl_d, j, n_vocab, &mut preds_d, j)?;
15367 }
15368 let preds = e.dtoh_u32(&preds_d)?;
15369 for j in 0..tc {
15370 bg[s + j + 1] = preds[j];
15371 }
15372 // MEMRA_REPLAY_NLL=1: teacher-forced NLL/perplexity over the same forced pass — the
15373 // checkpoint-quality metric (position j's logits score the GOLD next token).
15374 if nll_on {
15375 let jmax = if cend < t_total { tc } else { tc - 1 }; // last pos has no gold next
15376 if jmax > 0 {
15377 let ids: Vec<u32> = (0..jmax).map(|j| tokens[s + j + 1]).collect();
15378 let rows: Vec<i32> = (0..jmax as i32).collect();
15379 let idsd = e.htod_u32_v(&ids)?;
15380 let rowsd = e.htod_i32(&rows)?;
15381 let mut outd = e.zeros(jmax)?;
15382 e.softmax_gather(&tl_d, n_vocab, &idsd, &rowsd, &mut outd, n_vocab, jmax, 1.0)?;
15383 for pr in e.dtoh(&outd)? {
15384 nll_sum += -((pr.max(1e-30)) as f64).ln();
15385 nll_cnt += 1;
15386 }
15387 }
15388 }
15389 if let Some(f) = hdump.as_deref_mut() {
15390 use std::io::Write;
15391 let host: Vec<f32> = e.dtoh(&vx)?;
15392 // bf16 round-to-nearest-even — f32 doubled the disk bill at bulk
15393 // extraction scale (20M tokens x 4096 = 320GB f32 vs 160GB bf16).
15394 let mut bytes = Vec::with_capacity(tc * n_embd * 2);
15395 for v in &host[..tc * n_embd] {
15396 let b = v.to_bits();
15397 let r = b.wrapping_add(0x7FFF + ((b >> 16) & 1));
15398 bytes.extend_from_slice(&((r >> 16) as u16).to_le_bytes());
15399 }
15400 f.write_all(&bytes)?;
15401 }
15402 // CHAINLESS extraction (stride > corpus, the bulk-hdump mode): no chunk ever
15403 // drafts, so the draft-KV fills are pure waste — skip them (2 MTP-block passes
15404 // per token saved; the forced trunk pass + hdump is all the mode needs).
15405 let chainless = stride > t_total;
15406 if chainless {
15407 e.copy_view_into(
15408 &mut prev_last_h,
15409 0,
15410 &vx.slice((tc - 1) * n_embd..tc * n_embd),
15411 n_embd,
15412 )?;
15413 s = cend;
15414 continue;
15415 }
15416 // 2. TRUE predecessor-paired draft-KV fill for the chunk (row i carries h_{i-1};
15417 // row s reads the previous chunk's last true hidden, zeros at corpus start).
15418 let mut vxs = e.zeros(tc * n_embd)?;
15419 e.copy_into(&mut vxs, 0, &prev_last_h, n_embd)?;
15420 if tc > 1 {
15421 e.copy_view_into(
15422 &mut vxs,
15423 n_embd,
15424 &vx.slice(0..(tc - 1) * n_embd),
15425 (tc - 1) * n_embd,
15426 )?;
15427 }
15428 scratch.set_len(e, s)?;
15429 self.mtp_kv_fill_all(e, ch, &vxs, s, &mut scratch, embd_dev)?;
15430 // 3. draft chains at sampled positions, DESCENDING: a chain reads only slots
15431 // [0..p) (true fills) and appends at >= p; the next (smaller-p) chain's set_len
15432 // truncates those approximate appends before they can ever be read.
15433 let ps: Vec<usize> = (s..cend)
15434 .filter(|p| *p >= 1 && *p % stride == 0 && *p + k <= t_total)
15435 .collect();
15436 for &p in ps.iter().rev() {
15437 scratch.set_len(e, p)?;
15438 if p == s {
15439 e.copy_into(&mut seed_buf, 0, &prev_last_h, n_embd)?;
15440 } else {
15441 e.copy_view_into(
15442 &mut seed_buf,
15443 0,
15444 &vx.slice((p - 1 - s) * n_embd..(p - s) * n_embd),
15445 n_embd,
15446 )?;
15447 }
15448 let mut e_tok = tokens[p];
15449 let mut d_seed = e.clone_dtod(&seed_buf)?;
15450 let chain_heads = !self.mtp_extra.is_empty();
15451 let mut chain_tokens = if chain_heads {
15452 vec![tokens[p]]
15453 } else {
15454 Vec::new()
15455 };
15456 let mut chain_seeds = if chain_heads {
15457 vec![e.clone_dtod(&seed_buf)?]
15458 } else {
15459 Vec::new()
15460 };
15461 let mut drafts: Vec<u32> = Vec::with_capacity(k);
15462 for j in 0..k {
15463 let (dl_d, h_nextn) = if chain_heads {
15464 self.mtp_chain_forward_dev(
15465 e,
15466 &chain_tokens,
15467 &chain_seeds,
15468 &mut scratch,
15469 p,
15470 embd_dev,
15471 None,
15472 )?
15473 } else {
15474 self.mtp_head_forward_dev(
15475 e,
15476 mtp,
15477 e_tok,
15478 &d_seed,
15479 &mut scratch,
15480 p + 1 + j,
15481 embd_dev,
15482 None,
15483 )?
15484 };
15485 let tok_d = e.argmax_token_device(&dl_d, d_vocab)?;
15486 let idx = e.dtoh_u32_one(&tok_d)?;
15487 let d = match &mtp.d2t {
15488 Some(map) => map[idx as usize],
15489 None => idx,
15490 };
15491 drafts.push(d);
15492 if chain_heads {
15493 chain_tokens.push(d);
15494 chain_seeds.push(h_nextn);
15495 } else {
15496 e_tok = d;
15497 d_seed = h_nextn;
15498 }
15499 }
15500 // targets may live in a LATER chunk's bg — resolved after the walk.
15501 rows.push((p, drafts, Vec::new()));
15502 }
15503 // 4. restore TRUE entries for the whole chunk (the next chunk's chains and fills
15504 // expect scratch.len == cend with exact rows).
15505 scratch.set_len(e, s)?;
15506 self.mtp_kv_fill_all(e, ch, &vxs, s, &mut scratch, embd_dev)?;
15507 e.copy_view_into(
15508 &mut prev_last_h,
15509 0,
15510 &vx.slice((tc - 1) * n_embd..tc * n_embd),
15511 n_embd,
15512 )?;
15513 s = cend;
15514 }
15515 for (p, drafts, targets) in rows.iter_mut() {
15516 for j in 0..drafts.len() {
15517 targets.push(bg[*p + 1 + j]);
15518 }
15519 }
15520 rows.sort_by_key(|r| r.0);
15521 if nll_cnt > 0 {
15522 let mean = nll_sum / nll_cnt as f64;
15523 println!(
15524 "[replay-nll] tokens={nll_cnt} nll/token={mean:.5} ppl={:.4}",
15525 mean.exp()
15526 );
15527 }
15528 Ok((rows, bg))
15529 }
15530}
15531
15532#[cfg(test)]
15533mod vg_debt_tests {
15534 use super::dspark_vg_debt_projection;
15535
15536 /// TOOTH for the verify-graph admission accounting: the pool's projected remaining
15537 /// growth must be charged (pre-fix, admission charged 0 for a pool measured at
15538 /// 8,852 MiB), the projection must price the MARGINAL cost of one more key rather than
15539 /// extrapolating the pool's one-time shared allocation, and the doors that make growth
15540 /// impossible must zero the debt.
15541 #[test]
15542 fn vg_debt_projects_remaining_growth_and_respects_the_freeze_valves() {
15543 const MIB: usize = 1 << 20;
15544 let d = dspark_vg_debt_projection;
15545 // cold pool: nothing observed, one capture fits inside SPEC_SHRINK_RESERVE.
15546 assert_eq!(d(0, 256, 0, None), 0);
15547 // freeze valve MEMRA_DSPARK_VG_MAX=0: the pool cannot grow.
15548 assert_eq!(d(10, 0, 500 * MIB, None), 0);
15549 // saturated pool: at/past the cap the pool FREEZES, nothing left to reserve.
15550 assert_eq!(d(256, 256, 8852 * MIB, None), 0);
15551 assert_eq!(d(300, 256, 8852 * MIB, None), 0);
15552
15553 // BOOTSTRAP (one observation, growth unmeasurable): at most one more pool's worth.
15554 // The pre-fix mean rule extrapolated 255x here — the measured 8.5 GB phantom.
15555 assert_eq!(d(1, 256, 33 * MIB, None), 33 * MIB);
15556
15557 // MARGINAL, flat pool (the box9 receipt: reserved stayed ~33.6 MiB across captures
15558 // 1..3, so an additional key costs ~nothing and the debt must collapse to ~0 —
15559 // NOT the 8,556/4,261/2,830 MB the mean rule printed).
15560 assert_eq!(d(3, 256, 33 * MIB, Some((1, 33 * MIB))), 0);
15561
15562 // MARGINAL, genuinely growing pool: 40 MiB per new key over 2 keys, 250 slots left.
15563 let debt = d(6, 256, 273 * MIB, Some((4, 193 * MIB)));
15564 assert_eq!(debt, 250 * (40 * MIB));
15565 assert!(
15566 debt > 3 * (1536 * MIB),
15567 "real growth must dwarf SPEC_SHRINK_RESERVE"
15568 );
15569
15570 // a shrinking/recycled reading never becomes a negative charge.
15571 assert_eq!(d(6, 256, 10 * MIB, Some((4, 99 * MIB))), 0);
15572 // a stale observation at the same capture count falls back to bootstrap.
15573 assert_eq!(d(4, 256, 80 * MIB, Some((4, 80 * MIB))), 80 * MIB);
15574 }
15575}
15576
15577#[cfg(test)]
15578mod mtp_chain_tests {
15579 use super::mtp_chain_head_index;
15580
15581 #[test]
15582 fn embedded_step_heads_cycle_in_declared_order() {
15583 let actual: Vec<usize> = (0..8).map(|step| mtp_chain_head_index(step, 3)).collect();
15584 assert_eq!(actual, [0, 1, 2, 0, 1, 2, 0, 1]);
15585 }
15586
15587 #[test]
15588 fn standalone_draft_remains_single_head() {
15589 assert!((0..8).all(|step| mtp_chain_head_index(step, 1) == 0));
15590 }
15591}
15592
15593#[cfg(test)]
15594mod tp_verified_prefix_tests {
15595 use super::rewind_tp_kv_verified_prefix;
15596 use crate::tp::ResidentTpKvCache;
15597
15598 fn cache_with_committed_len(committed: usize) -> ResidentTpKvCache {
15599 let mut cache = ResidentTpKvCache::new(Vec::new(), 1, 1, 1, 1, 8);
15600 let transaction = cache.begin_transaction().unwrap();
15601 let target = cache.append_target(transaction, committed).unwrap();
15602 cache.publish_append(transaction, target).unwrap();
15603 let target = cache.commit_target(transaction, committed).unwrap();
15604 cache.publish_finalize(transaction, target).unwrap();
15605 cache
15606 }
15607
15608 #[test]
15609 fn replay_free_prefix_rewinds_tp_visibility_to_snapshot_plus_accepts() {
15610 let mut layers = vec![Some(cache_with_committed_len(5)), None];
15611 rewind_tp_kv_verified_prefix(&mut layers, &[Some(2), None], 1).unwrap();
15612 let cache = layers[0].as_ref().unwrap();
15613 assert_eq!(cache.committed_len(), 3);
15614 assert_eq!(cache.staged_len(), 3);
15615 }
15616
15617 #[test]
15618 fn replay_free_prefix_rejects_a_changed_tp_cache_shape() {
15619 let mut layers = vec![Some(cache_with_committed_len(1))];
15620 let error = rewind_tp_kv_verified_prefix(&mut layers, &[None], 1)
15621 .unwrap_err()
15622 .to_string();
15623 assert!(error.contains("changed shape"), "unexpected error: {error}");
15624 }
15625}
15626
15627#[cfg(test)]
15628mod dspark_sparse_tests {
15629 use super::dspark_sparse_softmax_topk;
15630
15631 #[test]
15632 fn topk_keeps_full_softmax_mass_and_stable_ties() {
15633 let logits = [1.0f32, 3.0, 3.0, -2.0];
15634 let (ids, top_logits, probs, tail) = dspark_sparse_softmax_topk(&logits, 2, 1.0).unwrap();
15635 assert_eq!(ids, vec![1, 2]);
15636 assert_eq!(top_logits, vec![3.0, 3.0]);
15637 let denominator = logits.iter().map(|value| (value - 3.0).exp()).sum::<f32>();
15638 let expected = 1.0 / denominator;
15639 assert!((probs[0] - expected).abs() < 1.0e-6);
15640 assert!((probs[1] - expected).abs() < 1.0e-6);
15641 assert!((tail - (1.0 - 2.0 * expected)).abs() < 1.0e-6);
15642 assert!((probs.iter().sum::<f32>() + tail - 1.0).abs() < 1.0e-6);
15643 }
15644}
15645
15646#[cfg(test)]
15647mod spec_replay_env_tests {
15648 use super::spec_replay_env_on;
15649
15650 #[test]
15651 fn replay_requires_literal_one() {
15652 assert!(!spec_replay_env_on(None));
15653 assert!(!spec_replay_env_on(Some("")));
15654 assert!(!spec_replay_env_on(Some("0")));
15655 assert!(!spec_replay_env_on(Some("true")));
15656 assert!(!spec_replay_env_on(Some("2")));
15657 assert!(spec_replay_env_on(Some("1")));
15658 }
15659}
15660
15661#[cfg(test)]
15662mod telem_tests {
15663 use super::{SPEC_TELEM_POS, SpecTelemetry, SpecTelemetryCounters};
15664
15665 #[test]
15666 fn synthetic_accept_masks_produce_tau_and_position_histogram() {
15667 let counters = SpecTelemetryCounters::default();
15668 for mask in [
15669 [true, true, true],
15670 [true, true, false],
15671 [true, false, false],
15672 [false, false, false],
15673 ] {
15674 let accepted = mask.iter().take_while(|&&value| value).count();
15675 counters.record_round(mask.len(), accepted);
15676 }
15677
15678 let snapshot = counters.snapshot();
15679 assert_eq!(
15680 (snapshot.rounds, snapshot.drafted, snapshot.accepted),
15681 (4, 12, 6)
15682 );
15683 assert_eq!(&snapshot.pos_drafted[..3], &[4, 4, 4]);
15684 assert_eq!(&snapshot.pos_accepted[..3], &[3, 2, 1]);
15685 assert_eq!(snapshot.tau(), 1.5);
15686 assert_eq!(snapshot.pos_drafted[3..], [0; SPEC_TELEM_POS - 3]);
15687 assert_eq!(snapshot.pos_accepted[3..], [0; SPEC_TELEM_POS - 3]);
15688 }
15689
15690 /// The worker's per-burst pattern: stash, accumulate, diff — the delta must isolate
15691 /// exactly the burst's contribution (pool-resumed sessions carry prior requests' counts).
15692 #[test]
15693 fn delta_isolates_burst_contribution() {
15694 let mut t = SpecTelemetry::default();
15695 // "previous request": 2 rounds of k=3, accepts 3 then 1.
15696 for (kr, na) in [(3usize, 3usize), (3, 1)] {
15697 t.rounds += 1;
15698 t.drafted += kr as u64;
15699 t.accepted += na as u64;
15700 for j in 0..kr {
15701 t.pos_drafted[j] += 1;
15702 }
15703 for j in 0..na {
15704 t.pos_accepted[j] += 1;
15705 }
15706 }
15707 let before = t;
15708 // "this burst": 1 round k=3, accepts 2.
15709 t.rounds += 1;
15710 t.drafted += 3;
15711 t.accepted += 2;
15712 for j in 0..3 {
15713 t.pos_drafted[j] += 1;
15714 }
15715 for j in 0..2 {
15716 t.pos_accepted[j] += 1;
15717 }
15718 let d = t.delta_since(&before);
15719 assert_eq!((d.rounds, d.drafted, d.accepted), (1, 3, 2));
15720 assert_eq!(&d.pos_drafted[..3], &[1, 1, 1]);
15721 assert_eq!(&d.pos_accepted[..3], &[1, 1, 0]);
15722 assert_eq!(d.pos_drafted[3..], [0; SPEC_TELEM_POS - 3]);
15723 }
15724
15725 /// merge(delta) then merge(delta2) equals accumulating both — the per-model /metrics
15726 /// aggregation invariant.
15727 #[test]
15728 fn merge_accumulates_fieldwise() {
15729 let mut agg = SpecTelemetry::default();
15730 let mut d1 = SpecTelemetry {
15731 rounds: 2,
15732 drafted: 6,
15733 accepted: 4,
15734 ..Default::default()
15735 };
15736 d1.pos_drafted[0] = 2;
15737 d1.pos_accepted[0] = 2;
15738 let mut d2 = SpecTelemetry {
15739 rounds: 1,
15740 drafted: 3,
15741 accepted: 1,
15742 ..Default::default()
15743 };
15744 d2.pos_drafted[0] = 1;
15745 d2.pos_accepted[0] = 1;
15746 d2.pos_drafted[1] = 1;
15747 agg.merge(&d1);
15748 agg.merge(&d2);
15749 assert_eq!((agg.rounds, agg.drafted, agg.accepted), (3, 9, 5));
15750 assert_eq!(agg.pos_drafted[0], 3);
15751 assert_eq!(agg.pos_accepted[0], 3);
15752 assert_eq!(agg.pos_drafted[1], 1);
15753 assert_eq!(agg.pos_accepted[1], 0);
15754 }
15755
15756 /// Wrong-snapshot diff saturates to zero instead of wrapping — the counters feed a
15757 /// public metrics surface and must never publish a u64-wrapped garbage value.
15758 #[test]
15759 fn delta_saturates_never_wraps() {
15760 let small = SpecTelemetry {
15761 rounds: 1,
15762 drafted: 2,
15763 accepted: 1,
15764 ..Default::default()
15765 };
15766 let big = SpecTelemetry {
15767 rounds: 5,
15768 drafted: 15,
15769 accepted: 9,
15770 ..Default::default()
15771 };
15772 let d = small.delta_since(&big);
15773 assert_eq!((d.rounds, d.drafted, d.accepted), (0, 0, 0));
15774 }
15775}
15776
15777#[cfg(test)]
15778mod opti_fork_tests {
15779 use super::{
15780 OptiControllerPolicy, OptiForkAction, OptiForkGateMode, OptiForkGenerationTracker,
15781 };
15782
15783 #[test]
15784 fn controller_threshold_and_three_miss_breaker_are_exact() {
15785 let mut policy = OptiControllerPolicy {
15786 threshold: 0.7,
15787 consecutive_misses: 0,
15788 breaker_tripped: false,
15789 };
15790 assert!(!policy.admit(0.699_999));
15791 assert!(policy.admit(0.7));
15792 assert!(!policy.resolve(false));
15793 assert!(!policy.resolve(false));
15794 assert!(policy.resolve(false));
15795 assert!(policy.breaker_tripped);
15796 assert!(!policy.admit(1.0));
15797 assert!(
15798 !policy.resolve(true),
15799 "a resolved hit cannot re-arm a tripped request"
15800 );
15801 assert!(policy.breaker_tripped);
15802 }
15803
15804 #[test]
15805 fn zero_threshold_is_the_true_unconditional_measurement_arm() {
15806 let mut policy = OptiControllerPolicy {
15807 threshold: 0.0,
15808 consecutive_misses: 0,
15809 breaker_tripped: false,
15810 };
15811 for _ in 0..16 {
15812 assert!(policy.admit(0.0));
15813 assert!(!policy.resolve(false));
15814 }
15815 for invalid in [f32::NAN, f32::INFINITY, -0.01, 1.01] {
15816 assert!(
15817 !policy.admit(invalid),
15818 "invalid q proxy must fail closed: {invalid}"
15819 );
15820 }
15821 assert!(!policy.breaker_tripped);
15822 assert_eq!(policy.consecutive_misses, 0);
15823 }
15824
15825 #[test]
15826 fn alternating_mode_flips_by_generation_not_round_parity() {
15827 assert_eq!(OptiForkGateMode::Alternate.action(0), OptiForkAction::Hit);
15828 assert_eq!(OptiForkGateMode::Alternate.action(1), OptiForkAction::Miss);
15829 assert_eq!(OptiForkGateMode::Alternate.action(8), OptiForkAction::Hit);
15830 assert_eq!(OptiForkGateMode::Alternate.action(9), OptiForkAction::Miss);
15831 }
15832
15833 #[test]
15834 fn live_generation_cannot_be_overwritten() {
15835 let mut tracker = OptiForkGenerationTracker::default();
15836 let g0 = tracker.reserve().unwrap();
15837 let g1 = tracker.reserve().unwrap();
15838 let err = tracker.reserve().unwrap_err().to_string();
15839 assert!(
15840 err.contains("still owns generation 0"),
15841 "unexpected error: {err}"
15842 );
15843 tracker.retire(g0).unwrap();
15844 let g2 = tracker.reserve().unwrap();
15845 assert_eq!((g2.id, g2.slot), (2, 0));
15846 tracker.retire(g1).unwrap();
15847 tracker.retire(g2).unwrap();
15848 }
15849
15850 #[test]
15851 fn teardown_rejects_a_stale_generation_tag() {
15852 let mut tracker = OptiForkGenerationTracker::default();
15853 let g0 = tracker.reserve().unwrap();
15854 tracker.retire(g0).unwrap();
15855 let err = tracker.retire(g0).unwrap_err().to_string();
15856 assert!(err.contains("teardown mismatch"), "unexpected error: {err}");
15857 }
15858}
15859
15860#[cfg(test)]
15861mod draft_graph_fallback_tests {
15862 use super::DraftGraphFallback;
15863
15864 /// The Q2 contract, part (a): a fallback flip is LOUD — exactly once per flip.
15865 #[test]
15866 fn flip_is_loud_once_and_memoized_after() {
15867 let mut f = DraftGraphFallback::default();
15868 let line = f
15869 .mark_greedy("out of memory")
15870 .expect("first flip must return the warn line");
15871 assert!(
15872 line.contains("WARN"),
15873 "flip line must be warn-level: {line}"
15874 );
15875 assert!(
15876 line.contains("out of memory"),
15877 "flip line must carry the reason: {line}"
15878 );
15879 assert!(f.greedy_failed());
15880 // re-marking an already-failed graph is the memoization: quiet, still failed.
15881 assert!(f.mark_greedy("out of memory").is_none());
15882 assert!(f.greedy_failed());
15883 // the two graphs' flags are independent (greedy flip leaves sampled capturable).
15884 assert!(!f.sampled_failed());
15885 let line_s = f
15886 .mark_sampled("capture unsupported")
15887 .expect("sampled flip is its own flip");
15888 assert!(
15889 line_s.contains("sampled"),
15890 "sampled flip names itself: {line_s}"
15891 );
15892 assert!(f.mark_sampled("capture unsupported").is_none());
15893 }
15894
15895 /// The Q2 contract, part (b): resume-from-pool RESETS both flags (fresh capture chance),
15896 /// and says so exactly when there was something to reset.
15897 #[test]
15898 fn reset_on_resume_clears_flags_and_logs_once() {
15899 let mut f = DraftGraphFallback::default();
15900 // clean session: resume is silent, nothing to reset.
15901 assert!(f.reset_on_resume().is_none());
15902 f.mark_greedy("oom").unwrap();
15903 f.mark_sampled("oom").unwrap();
15904 let note = f
15905 .reset_on_resume()
15906 .expect("a set flag must produce the reset note");
15907 assert!(
15908 note.contains("greedy+sampled"),
15909 "note names what was reset: {note}"
15910 );
15911 assert!(
15912 !f.greedy_failed() && !f.sampled_failed(),
15913 "both flags cleared"
15914 );
15915 // and the NEXT failure after a reset is a fresh flip — loud again.
15916 assert!(f.mark_greedy("oom again").is_some());
15917 let note2 = f.reset_on_resume().expect("greedy-only reset");
15918 assert!(note2.contains("(greedy)"), "single-flag note: {note2}");
15919 }
15920
15921 /// Shape-change clears (dmask realloc / mask-shape mismatch / s_key change) stay silent —
15922 /// they precede a fresh capture attempt whose own failure re-flips loudly.
15923 #[test]
15924 fn shape_change_clears_are_silent() {
15925 let mut f = DraftGraphFallback::default();
15926 f.mark_greedy("oom").unwrap();
15927 f.clear_greedy();
15928 assert!(!f.greedy_failed());
15929 f.mark_sampled("oom").unwrap();
15930 f.clear_sampled();
15931 assert!(!f.sampled_failed());
15932 // after a silent clear there is nothing left for resume to report.
15933 assert!(f.reset_on_resume().is_none());
15934 }
15935}
15936
15937/// SAMPLED DRAFT-GRAPH KEY (lane/graph-s-key-exactness-20260819).
15938///
15939/// These are the CPU teeth for an exactness bug whose live reproduction needs a GPU, a trunk, a
15940/// drafter and a two-turn session: the key itself. Every test below fails against the pre-fix key
15941/// `(seed, temp.to_bits(), k)` — `legacy_key` restates it so the collision is explicit rather
15942/// than remembered.
15943#[cfg(test)]
15944mod sampled_graph_key_tests {
15945 use super::{SampledGraphKey, debug_t_pred0};
15946
15947 /// The pre-fix key, verbatim: `let s_key = (sp_seed, sp_temp.to_bits(), k);`
15948 fn legacy_key(k: &SampledGraphKey) -> (u64, u32, usize) {
15949 (k.seed, k.temp_bits, k.k)
15950 }
15951
15952 fn pure_temp_key() -> SampledGraphKey {
15953 // temperature 1.0, filters off — today's serve default, the shape that parks a graph.
15954 SampledGraphKey::new(12345, 1.0, 3, 0, 1.0, 0.0, false)
15955 }
15956
15957 /// THE COLLISION. Two requests that differ ONLY in the truncation filters shared one key, so
15958 /// a parked pure-temp graph survived into a filtered request and the launch site launched it.
15959 #[test]
15960 fn vendor_filters_change_the_key() {
15961 let parked = pure_temp_key();
15962 // qwen3.8 generation_config.json — what the vendor-default flip makes the default shape.
15963 let vendor = SampledGraphKey::new(12345, 1.0, 3, 20, 0.95, 0.0, false);
15964 assert_eq!(
15965 legacy_key(&parked),
15966 legacy_key(&vendor),
15967 "pre-fix key collided: this is the bug, and the reason a test asserts on it",
15968 );
15969 assert_ne!(parked, vendor, "post-fix key must separate the two regimes");
15970 assert!(parked.pure_temp());
15971 assert!(!vendor.pure_temp());
15972 }
15973
15974 /// Each distribution-shaping field alone is enough to drop the parked graph.
15975 #[test]
15976 fn every_filter_field_is_keyed() {
15977 let base = pure_temp_key();
15978 for (what, other) in [
15979 (
15980 "top_k",
15981 SampledGraphKey::new(12345, 1.0, 3, 20, 1.0, 0.0, false),
15982 ),
15983 (
15984 "top_p",
15985 SampledGraphKey::new(12345, 1.0, 3, 0, 0.95, 0.0, false),
15986 ),
15987 (
15988 "min_p",
15989 SampledGraphKey::new(12345, 1.0, 3, 0, 1.0, 0.05, false),
15990 ),
15991 (
15992 "penalties",
15993 SampledGraphKey::new(12345, 1.0, 3, 0, 1.0, 0.0, true),
15994 ),
15995 ] {
15996 assert_ne!(base, other, "{what} must be part of the key");
15997 assert!(!other.pure_temp(), "{what} leaves the pure-temp regime");
15998 assert_eq!(
15999 legacy_key(&base),
16000 legacy_key(&other),
16001 "{what} was invisible to the pre-fix key",
16002 );
16003 }
16004 }
16005
16006 /// The baked constants stay keyed (this half was always right — regression cover for it).
16007 #[test]
16008 fn baked_constants_stay_keyed() {
16009 let base = pure_temp_key();
16010 assert_ne!(
16011 base,
16012 SampledGraphKey::new(999, 1.0, 3, 0, 1.0, 0.0, false),
16013 "seed"
16014 );
16015 assert_ne!(
16016 base,
16017 SampledGraphKey::new(12345, 0.7, 3, 0, 1.0, 0.0, false),
16018 "temp"
16019 );
16020 assert_ne!(
16021 base,
16022 SampledGraphKey::new(12345, 1.0, 4, 0, 1.0, 0.0, false),
16023 "k"
16024 );
16025 // bitwise on temperature: 0.7f32 vs the same value re-derived must NOT differ.
16026 assert_eq!(
16027 SampledGraphKey::new(1, 0.7, 3, 0, 1.0, 0.0, false),
16028 SampledGraphKey::new(1, 7.0 / 10.0, 3, 0, 1.0, 0.0, false),
16029 );
16030 }
16031
16032 /// THE LOAD-BEARING HALF OF THE SEED DECISION (lane/session-resume-sampler-predicate-
16033 /// 20260820). The whole-session resume predicate deliberately does NOT compare `seed`: an
16034 /// omitted serve `seed` draws fresh per-request entropy, so comparing it would refuse every
16035 /// seed-omitting sampled conversation. That is only sound because the one piece of parked state
16036 /// that BAKES the seed — this graph — is re-keyed on it, so a seed change drops and recaptures.
16037 ///
16038 /// This test is the other end of that argument, asserted here rather than remembered in a
16039 /// comment: if a future change dropped `seed` from the key, the resume predicate's exclusion
16040 /// would silently become the unsound thing it is documented not to be.
16041 /// (Paired with `seed_alone_does_not_refuse` in `memra-sampling`.)
16042 #[test]
16043 fn seed_alone_still_rekeys_the_draft_graph() {
16044 let parked = pure_temp_key();
16045 let reseeded = SampledGraphKey::new(999, 1.0, 3, 0, 1.0, 0.0, false);
16046 assert_ne!(
16047 parked, reseeded,
16048 "a seed-only change MUST drop the parked sampled graph — the resume predicate's \
16049 decision not to compare seed rests on exactly this",
16050 );
16051 // Same regime on both sides: the drop is a recapture, not a fall to the eager chain
16052 // because of a filter difference.
16053 assert!(parked.pure_temp() && reseeded.pure_temp());
16054 }
16055
16056 /// `pure_temp()` is the capture guard's predicate, computed from the key so the two cannot
16057 /// drift. The equality below is the invariant the launch-site guard asserts: identical keys
16058 /// agree on the regime, so a graph that survives the drop is legal to launch.
16059 #[test]
16060 fn equal_keys_agree_on_the_regime() {
16061 let a = SampledGraphKey::new(7, 0.8, 3, 20, 0.95, 0.0, false);
16062 let b = SampledGraphKey::new(7, 0.8, 3, 20, 0.95, 0.0, false);
16063 assert_eq!(a, b);
16064 assert_eq!(a.pure_temp(), b.pure_temp());
16065 // top_p slightly above 1.0 (a client sending 1.0 exactly, or an operator default) is
16066 // still the unfiltered regime, matching the original `sp.top_p >= 1.0` test.
16067 assert!(SampledGraphKey::new(7, 0.8, 3, 0, 1.0, 0.0, false).pure_temp());
16068 assert!(SampledGraphKey::new(7, 0.8, 3, 0, 1.5, -1.0, false).pure_temp());
16069 }
16070
16071 /// The WIDENED capture regime (lane/step37-draft-graph-serving-20260830): truncation-
16072 /// filtered shapes are capturable — the filter runs IN-GRAPH (`filter_stats` +
16073 /// `gumbel_perturb_filtered_ctr`), so the draft draws from the same filtered
16074 /// distribution the accept test reconstructs. Penalties never are: the per-round
16075 /// history cannot be baked. The step37 vendor-default shape (temp 0.5 / top_p 0.9) is
16076 /// exactly the previously-excluded regime this lane exists to capture.
16077 #[test]
16078 fn filtered_regimes_are_capturable_penalties_never() {
16079 let vendor = SampledGraphKey::new(12345, 0.5, 3, 0, 0.9, 0.0, false);
16080 assert!(!vendor.pure_temp());
16081 assert!(vendor.filtered());
16082 assert!(
16083 vendor.graph_capturable(),
16084 "the vendor-default filtered shape must be capturable (default door state)",
16085 );
16086 assert!(pure_temp_key().graph_capturable());
16087 assert!(
16088 !pure_temp_key().filtered(),
16089 "pure-temp takes the legacy (filterless) capture body",
16090 );
16091 let pen = SampledGraphKey::new(12345, 0.5, 3, 0, 0.9, 0.0, true);
16092 assert!(
16093 !pen.graph_capturable(),
16094 "penalty history varies per round and can never be baked into a graph",
16095 );
16096 }
16097
16098 /// MEMRA_DEBUG_SPEC on a SAMPLED spec request past round 0: the print must render without
16099 /// indexing the empty greedy `preds` vector (it panicked the GPU worker before this lane).
16100 #[test]
16101 fn debug_print_survives_the_sampled_arm() {
16102 // round >= 1 with a pending bonus == base 1, sampled == `preds` empty.
16103 assert_eq!(debug_t_pred0(true, 1, 4242, &[]), "n/a");
16104 assert_eq!(debug_t_pred0(true, 2, 4242, &[]), "n/a");
16105 // round 0 without a pending bonus still reports last_pred, in both arms.
16106 assert_eq!(debug_t_pred0(true, 0, 4242, &[]), "4242");
16107 assert_eq!(debug_t_pred0(false, 0, 4242, &[7, 8]), "4242");
16108 // greedy keeps the real prediction it always printed.
16109 assert_eq!(debug_t_pred0(false, 1, 4242, &[7, 8]), "7");
16110 assert_eq!(debug_t_pred0(false, 2, 4242, &[7, 8]), "8");
16111 }
16112}