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; default OFF until the gate
32/// battery banks receipts, per the new-flags law). ON routes the step35 MTP block's draft
33/// attention through the WINDOWED device-counter family (`append_kv_quantized_dcw` +
34/// `fa_decode_dcw`, the step TP graph arc's kernels), which derives the SWA view entirely
35/// from device state (len_d, base_d, window): exactly the view offset the old capture
36/// refusal said `fa_decode_dc` could not express. BOTH draft modes switch together: eager
37/// and captured run the ONE launcher at the ONE bucket (min(cap, window)), so graph-vs-eager
38/// draft parity holds by construction (the `mtp_full_attn_dc` precedent). OFF keeps today's
39/// serving byte-for-byte: the host-len eager arm (`mtp_step35_attn`) plus the named capture
40/// refusal. Rollback seam: unset the env (or set 0) and restart; no state survives.
41fn step35_draft_dcw_on() -> bool {
42 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
43 *ON.get_or_init(|| std::env::var("MEMRA_STEP35_DRAFT_DCW").as_deref() == Ok("1"))
44}
45
46fn parse_prime_trows_width(value: Option<&str>) -> Result<usize, String> {
47 let Some(raw) = value else {
48 return Ok(8);
49 };
50 let width = raw
51 .parse::<usize>()
52 .map_err(|_| format!("MEMRA_PRIME_TROWS_T must be an integer in 2..=8, got {raw:?}"))?;
53 if !(2..=8).contains(&width) {
54 return Err(format!("MEMRA_PRIME_TROWS_T must be in 2..=8, got {width}"));
55 }
56 Ok(width)
57}
58
59#[cfg(test)]
60mod prime_trows_width_tests {
61 #[test]
62 fn width_defaults_to_eight_and_refuses_invalid_operator_values() {
63 assert_eq!(super::parse_prime_trows_width(None), Ok(8));
64 assert_eq!(super::parse_prime_trows_width(Some("2")), Ok(2));
65 assert_eq!(super::parse_prime_trows_width(Some("8")), Ok(8));
66 for invalid in ["", "1", "9", "32", "wide"] {
67 let err = super::parse_prime_trows_width(Some(invalid)).unwrap_err();
68 assert!(err.contains("MEMRA_PRIME_TROWS_T"), "{err}");
69 assert!(err.contains("2..=8"), "{err}");
70 }
71 }
72}
73
74/// One compact, anchor-bounded DSpark supervision record. `tokens[0]` is the anchor at p and
75/// `hidden` is its predecessor carrier h[p-1], matching the live NextN/DSpark pairing. Target
76/// rows p..p+gamma-1 score tokens p+1..p+gamma. They are the full-target softmax's top-k
77/// entries; `target_tail_probs[j]` is the probability mass outside those rows. All flattened
78/// target arrays are `[gamma, top_k]` in row-major order.
79pub struct DsparkAnchorRecord {
80 pub position: usize,
81 pub hidden: Vec<f32>,
82 pub tokens: Vec<u32>,
83 pub target_top_ids: Vec<u32>,
84 pub target_top_logits: Vec<f32>,
85 pub target_top_probs: Vec<f32>,
86 pub target_tail_probs: Vec<f32>,
87}
88
89fn dspark_sparse_softmax_topk(
90 logits: &[f32],
91 top_k: usize,
92 temperature: f32,
93) -> Result<(Vec<u32>, Vec<f32>, Vec<f32>, f32), Box<dyn std::error::Error>> {
94 if logits.is_empty() || top_k == 0 || top_k > logits.len() || temperature <= 0.0 {
95 return Err("invalid DSpark sparse-softmax shape or temperature".into());
96 }
97 if logits.iter().any(|value| !value.is_finite()) {
98 return Err("DSpark target logits contain a non-finite value".into());
99 }
100 let mut ranked: Vec<(u32, f32)> = logits
101 .iter()
102 .copied()
103 .enumerate()
104 .map(|(index, value)| (index as u32, value))
105 .collect();
106 let compare = |left: &(u32, f32), right: &(u32, f32)| {
107 right.1.total_cmp(&left.1).then(left.0.cmp(&right.0))
108 };
109 ranked.select_nth_unstable_by(top_k - 1, compare);
110 ranked[..top_k].sort_unstable_by(compare);
111
112 let max_logit = logits.iter().copied().fold(f32::NEG_INFINITY, f32::max);
113 let inv_temperature = 1.0f64 / temperature as f64;
114 let denominator: f64 = logits
115 .iter()
116 .map(|value| (((*value - max_logit) as f64) * inv_temperature).exp())
117 .sum();
118 let ids: Vec<u32> = ranked[..top_k].iter().map(|(index, _)| *index).collect();
119 let top_logits: Vec<f32> = ranked[..top_k].iter().map(|(_, value)| *value).collect();
120 let top_probs: Vec<f32> = top_logits
121 .iter()
122 .map(|value| ((((value - max_logit) as f64) * inv_temperature).exp() / denominator) as f32)
123 .collect();
124 let top_mass: f64 = top_probs.iter().map(|value| *value as f64).sum();
125 let tail = (1.0f64 - top_mass).clamp(0.0, 1.0) as f32;
126 Ok((ids, top_logits, top_probs, tail))
127}
128
129fn flatten_dspark_rows<T>(
130 rows: Vec<Option<Vec<T>>>,
131 position: usize,
132 label: &str,
133) -> Result<Vec<T>, Box<dyn std::error::Error>> {
134 let mut flattened = Vec::new();
135 for (slot, row) in rows.into_iter().enumerate() {
136 flattened.extend(
137 row.ok_or_else(|| format!("missing DSpark {label} at {position} slot {slot}"))?,
138 );
139 }
140 Ok(flattened)
141}
142
143/// H-SEED CONVENTION (MEMRA_SPEC_HPOST=1): feed the MTP head the POST-norm hidden — trunk rows
144/// hand over `output_norm(x)` and the draft chain recurrence hands over `shared_head_norm(h_nextn)`
145/// (= final_h) — matching the reference engines: llama.cpp #24025 ("qwen35: use post-norm hidden
146/// state for MTP", t_h_nextn is taken AFTER the final norm in both trunk and MTP graphs) and
147/// SGLang's qwen3_5_mtp (spec_info.hidden_states = the target model's post-norm output). memra's
148/// historical convention (default, MTP-PLAN §A) is PRE-norm x. Draft-quality-only: exactness is
149/// the verify's job either way; acceptance arbitrates. OnceLock: read once, hot-loop safe.
150/// `MEMRA_SPEC_HEAD_ROWS=1` — batch the verify tail's LM head over its t columns instead of running
151/// it at m=1 once per column. See the call site in `decode_step_t_core_stream` for why the batched
152/// form is the same per-row arithmetic (the bf16/q8 rows twins, not cuBLASLt) and what it costs
153/// today: the head is re-streamed t times per verify pass. Default off until the byte tape says so.
154pub(crate) fn head_rows_on() -> bool {
155 static ENV: std::sync::OnceLock<Option<bool>> = std::sync::OnceLock::new();
156 crate::step37_door(&ENV, "MEMRA_SPEC_HEAD_ROWS")
157}
158
159/// The serving walk's own doors, tri-stated the same way (owner flip 2026-08-27): env forces,
160/// unset takes the step37 family default. Call sites are the t-row verify walk itself.
161pub(crate) fn spec_verify_eager_on() -> bool {
162 static ENV: std::sync::OnceLock<Option<bool>> = std::sync::OnceLock::new();
163 crate::step37_door(&ENV, "MEMRA_SPEC_VERIFY_EAGER")
164}
165
166pub(crate) fn spec_verify_tcol_on() -> bool {
167 static ENV: std::sync::OnceLock<Option<bool>> = std::sync::OnceLock::new();
168 crate::step37_door(&ENV, "MEMRA_SPEC_VERIFY_TCOL")
169}
170
171/// NOT family-armed (2026-08-27): the walk's prime leaves its sub-32 TAIL chunk out of the
172/// DISTRIBUTED kv, so the server refuses before decode with "cache lengths diverged
173/// local=N distributed=floor(N/32)*32" for every prompt whose token count is not a multiple of
174/// 32 — i.e. nearly all real traffic. Isolated on the server route: defaults ERR (local=445
175/// distributed=416), MEMRA_PRIME_TROWS=0 OK. It was default-OFF before the 2026-08-27 flip and
176/// goes back to opt-in until the tail append is fixed and gated ON THE SERVER ROUTE, not just
177/// run-gen (run-gen calls decode_step_t on the whole prompt and never exercises this path — the
178/// reason a run-gen-only receipt could not see it). The GEMM prime supersedes it on this route.
179pub(crate) fn prime_trows_on() -> bool {
180 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
181 *ON.get_or_init(|| std::env::var("MEMRA_PRIME_TROWS").as_deref() == Ok("1"))
182}
183
184pub(crate) fn tcol_ffn_on() -> bool {
185 static ENV: std::sync::OnceLock<Option<bool>> = std::sync::OnceLock::new();
186 crate::step37_door(&ENV, "MEMRA_TCOL_FFN")
187}
188
189pub(crate) fn spec_hpost() -> bool {
190 static H: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
191 *H.get_or_init(|| {
192 std::env::var("MEMRA_SPEC_HPOST")
193 .map(|v| v != "0")
194 .unwrap_or(false)
195 })
196}
197
198/// LEAN VERIFY (default ON since 2026-07-08; MEMRA_SPEC_LEAN=0 reverts — close35 lane): the verify m-scaling
199/// probe + nsys diff showed the verify t-path pays ~1.0ms/call at m=1 over eager decode on the
200/// 35B, and the kernels are NOT the cause (dev-MoE identical, kernel-time delta only +179us).
201/// The overhead is (a) ~250 extra cuMemsetD8Async/call from `e.zeros()` on buffers every kernel
202/// fully overwrites (~0.9ms host issue + ~0.35ms GPU) and (b) the t=1 FA rows dispatch (rows_v2 +
203/// combine_rows, +50us vs the eager fa_decode pair). This flag switches (a) fully-overwritten
204/// verify buffers to `e.uninit` (identical bytes: every element is written before read) and
205/// (b) t==1 verify FA to the eager `fa_decode` entry (byte-identical: kernel-check pins the
206/// rows-vs-loop identity and the per-row loop at t=1 IS fa_decode on the same q). Gates arbitrate.
207pub(crate) fn spec_lean() -> bool {
208 static L: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
209 // DEFAULT ON since 2026-07-08 (MEMRA_SPEC_LEAN=0 reverts): bit-identical (buffers fully
210 // overwritten; gates green incl maxdiff-identical run-gen) and measured +2.4% e2e p3 /
211 // +1.5% p2 at the daily 35B config. m=1 verify now costs eager-decode parity.
212 *L.get_or_init(|| {
213 std::env::var("MEMRA_SPEC_LEAN")
214 .map(|v| v != "0")
215 .unwrap_or(true)
216 })
217}
218
219/// SMALL-M BATCHED VERIFY (default ON since 2026-07-09; MEMRA_SPEC_M2=0 reverts — lane/spec-m2): extend the
220/// batched linear-attn verify arm down to t=2 and batch the MoE dev token loop over a
221/// grid.z=token axis at every verify t. The close35 m-scaling probe put the m=2 verify tier at
222/// x1.54 of m=1 (llama x1.14); the per-column linear chain (t<3) and the serial MoE dev token
223/// loop are the two launch-structure causes. Both changes are LAUNCH-STRUCTURE ONLY:
224/// (a) the batched conv's t<pad ring update is pure copies (ssm_conv_ring_rebuild from a cloned
225/// ring — the ring stores raw input columns); every arithmetic kernel is the same one the
226/// t>=3 arm already runs (matmul_decode_exact bit-identical at m=2-4, gdn_scan's internal
227/// t-loop == chained T=1 steps);
228/// (b) the MoE dev-rows twins run the serial loop's per-token warp program with tok-offset
229/// pointers (same sel/w/aq/ad bytes, same dot order, same slot-ordered FMA chain).
230/// Gates arbitrate: run-spec K=1..8 self-consistency (35B+9B), kernel-check, run-gen argmax.
231pub(crate) fn spec_m2() -> bool {
232 static M: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
233 // DEFAULT ON since 2026-07-09 (MEMRA_SPEC_M2=0 reverts): launch-structure only — t=2
234 // batched linear arm (ring-roll copies, zero new FP order) + MoE dev-rows kernels
235 // (grid.z=token, 4 launches/layer at any verify t). Acceptance bit-identical at every K;
236 // 35B p2 +3.4% / p3 +3.6%; the profitable-K plateau widens (new optimum K=3 at 223).
237 *M.get_or_init(|| {
238 std::env::var("MEMRA_SPEC_M2")
239 .map(|v| v != "0")
240 .unwrap_or(true)
241 })
242}
243pub(crate) fn spec_stream() -> bool {
244 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
245 *ON.get_or_init(|| std::env::var("MEMRA_SPEC_STREAM").as_deref() == Ok("1"))
246}
247pub(crate) fn spec_stream_m() -> usize {
248 static M: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
249 *M.get_or_init(|| {
250 std::env::var("MEMRA_SPEC_STREAM_M")
251 .ok()
252 .and_then(|v| v.parse().ok())
253 .unwrap_or(4)
254 })
255}
256pub(crate) fn spec_devacc() -> bool {
257 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
258 *ON.get_or_init(|| std::env::var("MEMRA_SPEC_DEVACC").as_deref() == Ok("1"))
259}
260/// Engine-bundle slice 2 (DSF-ROUNDCOST-20260820 §1.1 host/device round trips + §2 rows 2-3),
261/// DEFAULT ON (`MEMRA_DSPARK_DEFER_READBACK=0` reverts): the dspark round's draft-chain DtoH
262/// is DEFERRED past verify dispatch and merged with the verify-argmax readback into ONE host
263/// sync (2 blocking DtoH/round -> 1). Verify embeds DEVICE tokens (`chain_d`) through the
264/// resident embed table — `embed_gather_u32_t`, bit-identical rows to the host gather by its
265/// own pinned contract. The host therefore dispatches snap + the whole verify while the DRAFT
266/// is still executing, instead of blocking ~1.7 ms on the chain and letting the device drain.
267/// Ladder arm only: the confidence policies size vt from a pre-verify head readback (their
268/// chain readback merges into that same sync instead). Exactness unchanged BY CONSTRUCTION —
269/// same tokens, same kernels, same order; E2E + accept-bank gates arbitrate.
270pub(crate) fn dspark_defer_readback_on() -> bool {
271 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
272 *ON.get_or_init(|| {
273 std::env::var("MEMRA_DSPARK_DEFER_READBACK")
274 .map(|v| v != "0")
275 .unwrap_or(true)
276 })
277}
278/// Engine-bundle slice 1 (DSF-ROUNDCOST-20260820 §1.1, lane/dspark-engine-bundle-20260820),
279/// DEFAULT ON (`MEMRA_STATE_COPY_BATCH=0` reverts): batch the dspark round's GDN state
280/// snapshot and partial-accept restore into single `copy_batch_uniform_f32` launches
281/// instead of ~2 memcpy dispatches (+2 alloc_zeros on the snap side) per linear layer per
282/// round — measured 0.67 ms/round snap + 0.25 ms/round commit of pure dispatch on the q38
283/// route. Launch-structure only: bytes, buffers and stream order are unchanged, so
284/// acceptance and streams stay bit-identical (E2E-gated on the B1 packs).
285pub(crate) fn state_copy_batch_on() -> bool {
286 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
287 *ON.get_or_init(|| {
288 std::env::var("MEMRA_STATE_COPY_BATCH")
289 .map(|v| v != "0")
290 .unwrap_or(true)
291 })
292}
293/// Engine-bundle slice 3 + fa-execupdate slice 4c (DSF-ROUNDCOST-20260820 §5 rank 1),
294/// DEFAULT OFF — `MEMRA_DSPARK_VERIFY_GRAPH=1` opts in: per-(segment, vt) CUDA graphs
295/// for the LINEAR-layer runs, plus the full-verify single graph per (vt, rung) when a
296/// round's rows all ride one seqs rung — see [`DsparkVerifyGraphs`]. Requires the
297/// slice-2 deferred path (device tokens); the eager walk is the byte-identical fallback.
298///
299/// MEASURED disposition (box6 card0, agentic pack, 2026-08-20, both slices): exactness
300/// holds everywhere (ALL EXACT, accept lines byte-match the banks, ckpt-gate oracle
301/// green over the graph + slab-commit paths). Slice-3's AUTO_FREE launch-scan limiter
302/// (25.6 us x 16 launches ≈ 0.41 ms/round) is FIXED — the captured bodies' alloc nodes
303/// are balanced by in-graph frees (census 84/84 per segment, 1776/1776 full) so graphs
304/// instantiate USE_NODE_PRIORITY and the scan is gone. What remains at gate scale:
305/// segment graphs +0.1 tok/s over the batched-rows default (114.4 vs 114.3 x5
306/// interleaved — the linear launch overhead was only ~0.1 ms); the FULL-verify graph is
307/// NET NEGATIVE at gate scale (110.6 vs 114.2: ~14-21 (vt, rung) captures/process at
308/// 2 full-walk executions + ~2.9k-node instantiate each eat far more than the ~0.2-0.3
309/// ms/round of remaining launch overhead). The orchestration ceiling of §1.3 is spent —
310/// the fa/append recovery landed DEFAULT-ON as the batched rows arm
311/// (`dspark_fa_rows_on`), not as a graph. The serve-lifetime cell (DSF-ROUNDCOST §9,
312/// nj-ws-solo) measured the amortization: crossover K≈33 requests, steady −0.246
313/// ms/round, −1.25% session wall over 240 requests — and the graphs-serve lane wired
314/// the door into the session arm (`dspark_spec_session_burst`) as a model-owned
315/// capture pool shared across sessions. Stays opt-in pending the owner's default-ON
316/// ratification on the serve-surface battery.
317pub(crate) fn dspark_verify_graph_on() -> bool {
318 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
319 *ON.get_or_init(|| std::env::var("MEMRA_DSPARK_VERIFY_GRAPH").as_deref() == Ok("1"))
320}
321/// MTP-ROUTE verify graphs, DEFAULT ON for the GDN+MoE family since 2026-08-23
322/// (`MEMRA_SPEC_VERIFY_GRAPH=0` is the kill switch, `=1` opts other families in).
323///
324/// The slice-4c capture already lived inside `qwen35_verify_tparallel` and said so in its own
325/// comment — "stream rides the qwen35moe burst, graphs ride the dspark route" — with no caller
326/// on this route. The MTP spec round is that caller.
327///
328/// WHY it is worth a default (receipts: `research/orndecode-20260822/VGRAPH.md`). With
329/// `MEMRA_SPEC_PHASE=1` this route's round reads verify-ISSUE 44-58% and verify-WAIT **0.0%**:
330/// the host is never waiting for the device, it is spending its own time launching the trunk.
331/// Replay collapses that into one graph launch and the phase all but disappears (55-62 ms ->
332/// 8-10 ms per burst).
333///
334/// MEASURED, two host generations, forced ON/OFF, balanced 4+4 boots in both orders:
335/// * current-generation host (9950X, the serving class): OFF 266.0-266.5, ON 318.8-319.5
336/// tok/s — **+19.7%**, no overlap, sub-1% spread per arm; per-round 6.9 -> 5.7 ms.
337/// * Zen 3 host: +3-9% (that rig's own clock drift is wider than the effect, so the ratio
338/// comes from per-round phase totals, which are internal to each boot).
339/// The ON arm lands at ~320 tok/s on BOTH hosts while OFF tracks host speed — the arm moves
340/// the round off the host and onto the device, which is the whole point.
341///
342/// EXACTNESS is structural (same kernels, same order) and gated anyway: a fixed-seed SAMPLED
343/// completion hashes identically ON vs OFF **and across both hosts** (`08941d5bb9762b21`),
344/// greedy seed-pinned likewise, `run-spec` K=1..8 PASS on both arms with identical acceptance
345/// at every K, kernel-check ALL GREEN.
346///
347/// SCOPE, deliberately narrow: default ON only where it was measured — the GatedDeltaNet +
348/// MoE family (`vgraph_family_default`). Qwen3.8-27B is GDN + DENSE mlp and would otherwise
349/// inherit this default unmeasured, which is the family-by-family law this repo keeps; it can
350/// opt in with `=1` once it has its own interleave. Also never armed together with
351/// ROUND-STREAM, and a round wider than the pool declines it for the eager walk.
352pub(crate) fn spec_verify_graph_env() -> Option<bool> {
353 static ON: std::sync::OnceLock<Option<bool>> = std::sync::OnceLock::new();
354 *ON.get_or_init(
355 || match std::env::var("MEMRA_SPEC_VERIFY_GRAPH").as_deref() {
356 Ok("1") => Some(true),
357 Ok("0") => Some(false),
358 _ => None,
359 },
360 )
361}
362/// SERVE-ROUTE twin of [`dspark_verify_graph_on`], DEFAULT ON — owner-ratified
363/// 2026-08-22 on the §10 serve-lifetime battery (DSF-ROUNDCOST-20260820 §10.3:
364/// crossover K=36–43, steady −0.357 ms/round, session wall −1.55..−1.65%, byte-exact
365/// 240/240 ×3 pairs, pool bounded at 8,852 MiB under `MEMRA_DSPARK_VG_MAX`). The env
366/// stays as the kill-switch: `MEMRA_DSPARK_VERIFY_GRAPH=0` restores the eager walk
367/// (byte-identical body); `MEMRA_DSPARK_VG_MAX=0` is the finer freeze valve. The BIN
368/// arm keeps its own opt-in default (`dspark_verify_graph_on`): at gate scale the
369/// capture toll is never repaid (§8 measured disposition — 14–21 captures over a
370/// 256-token run vs the serve session's thousands of rounds), and the two
371/// instruments must keep their own measured dispositions rather than share one flag.
372pub(crate) fn dspark_verify_graph_serve_on() -> bool {
373 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
374 *ON.get_or_init(|| std::env::var("MEMRA_DSPARK_VERIFY_GRAPH").as_deref() != Ok("0"))
375}
376/// Capture-count ceiling for the dspark verify-graph pool (graphs-serve lane) — the
377/// pool's memory policy STATED instead of silently unbounded. The keyspace is
378/// intrinsically finite — segment keys (run_start, vt) ≤ 16 runs x 7 windows, full
379/// keys (vt, rung, hi) ≤ 7 windows x the split-rung ladder (8 rungs at 32k ctx), ~168
380/// on the q38 export — so the default (256) never engages there; the knob is the
381/// safety valve for a future export with a wider ladder. At the ceiling the pool
382/// FREEZES: existing keys keep replaying, rounds needing a new capture run the eager
383/// walk byte-identically (round-atomic — a partial refusal would mix slab- and
384/// cols-stashed layers inside one commit). No eviction by design: destroying a live
385/// exec graph re-opens the stale-address class the indirect tables exist to close,
386/// and the bounded keyspace makes reclaim worthless.
387pub(crate) fn dspark_vg_cap() -> usize {
388 static CAP: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
389 *CAP.get_or_init(|| {
390 std::env::var("MEMRA_DSPARK_VG_MAX")
391 .ok()
392 .and_then(|v| v.parse().ok())
393 .unwrap_or(256)
394 })
395}
396
397/// PROJECTED REMAINING GROWTH of the verify-graph pool, in bytes (lane/hermes-perf-fixes,
398/// 2026-08-23 — the admission accounting the "pool dwarfs spec admission reserve" finding
399/// asks for). The pool was measured at 8,852 MiB at storm-complete on the q38 export while
400/// admission's transient floor (`SPEC_SHRINK_RESERVE`) is 1.5 GiB and never charged for it:
401/// sessions admitted while the pool is cold overcommit VRAM the pool WILL hold, because the
402/// pool grows monotonically (no eviction by design) and is model-owned across sessions.
403///
404/// SELF-MEASURING, no per-model constant (generic-model law — the 8,852 MiB is a q38 number
405/// and proves nothing about another export): the debt is remaining capture slots x the
406/// MARGINAL bytes a capture adds to this device's graph mem pool.
407///
408/// MARGINAL, NOT MEAN — measured correction (box9 on-box receipt, 2026-08-23). The first
409/// version of this used the mean (`reserved / captures`) and the live serve log showed why
410/// that is wrong: with the pool's reservation flat at ~33.6 MiB across captures 1..3, the
411/// mean-based debt printed **8,556 MB, then 4,261, then 2,830** — it extrapolated capture
412/// #1's ONE-TIME shared allocation (staging buffers, stash slabs, pointer tables: sized
413/// once per pool, shared by every key) across all 256 slots. An 8.5 GB phantom reserve at
414/// boot can refuse admissions that would have fit, which is a worse defect than the
415/// under-charge this accounting exists to remove. The marginal reading prices what an
416/// ADDITIONAL key actually costs: two observations `(captures, reserved)` give
417/// `(r1 - r0) / (c1 - c0)`, which is ~0 on an export whose pool does not grow per key and
418/// tracks real growth on one that does.
419///
420/// BOOTSTRAP (only one observation so far, so growth is unmeasurable): reserve one more
421/// pool's worth — `min(remaining x mean, reserved)`. "We have measured `reserved` bytes for
422/// `captures` keys; until growth is measurable, assume at most a doubling" is fail-safe in
423/// the same direction as the old rule without the 255x extrapolation.
424///
425/// Before the FIRST capture the debt is 0 (a single capture lands well inside the existing
426/// 1.5 GiB floor). `cap` is the intrinsic freeze ceiling (`MEMRA_DSPARK_VG_MAX`; =0 freeze
427/// valve => the pool cannot grow => debt 0); at or past the cap the pool FREEZES, so the
428/// debt is 0 there too.
429pub fn dspark_vg_debt_projection(
430 captures: usize,
431 cap: usize,
432 reserved_bytes: usize,
433 prev: Option<(usize, usize)>,
434) -> usize {
435 if captures == 0 || cap == 0 {
436 return 0;
437 }
438 let remaining = cap.saturating_sub(captures);
439 if remaining == 0 {
440 return 0;
441 }
442 match prev {
443 // marginal growth between two observations of the same pool
444 Some((c0, r0)) if captures > c0 => {
445 let marginal = reserved_bytes.saturating_sub(r0) / (captures - c0);
446 remaining.saturating_mul(marginal)
447 }
448 // bootstrap: at most one more pool's worth
449 _ => remaining
450 .saturating_mul(reserved_bytes / captures)
451 .min(reserved_bytes),
452 }
453}
454/// Engine-bundle slice 4 (fa-execupdate lane, DSF-ROUNDCOST-20260820 §6 close: "the
455/// residual gap lives in the FULL-ATTENTION per-row section"), DEFAULT ON —
456/// `MEMRA_DSPARK_FA_ROWS=0` reverts to the per-row loop: when every row of a verify
457/// round takes the v4-seqs arm on ONE `fa_split_keys` rung (the straddle law, evaluated
458/// at the round's first and last t_kv — both eligibility gates are intervals in t_kv),
459/// the qwen35 t-parallel verify's per-row KV-append + fa-decode loop collapses into the
460/// z-batched serving twins: ONE `append_quantize_kv_q8_0_q5_1_seqs` + ONE
461/// `fa_decode_vec_q_seqs_v4` + ONE combine per full-attention layer, replacing
462/// T x (4 dtod row copies + append + 3 memsets + main + combine) launches. Bytes are
463/// pinned by the batched-tick increment-2 kernel-check (seqs-vs-per-seq-loop bit
464/// identity: per-row T_kv derives in-kernel from pos_seq[z]; splits >= ns_eff write the
465/// empty partial the combine never reads, so the shared n_splits_max stride changes no
466/// bytes) and re-gated e2e by this lane's battery.
467pub(crate) fn dspark_fa_rows_on() -> bool {
468 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
469 *ON.get_or_init(|| {
470 std::env::var("MEMRA_DSPARK_FA_ROWS")
471 .map(|v| v != "0")
472 .unwrap_or(true)
473 })
474}
475
476/// `t_pred0` for the `MEMRA_DEBUG_SPEC` per-round print, sampled-safe.
477///
478/// `generate_spec_inner2` fills its `preds` vector ONLY on the greedy path (`if !sampled`), and
479/// the per-round debug print was the sole consumer in the sampled arm: `t_pred(0)` survives round
480/// 0 (`base == 0` returns `last_pred`) and from round 1 (`base == 1`, a pending bonus) indexes an
481/// EMPTY vector — `index out of bounds: the len is 0 but the index is 0`, in the GPU worker
482/// thread, which then respawns and reloads weights while the request dies. So any sampled spec
483/// request longer than one round used to kill the worker whenever `MEMRA_DEBUG_SPEC` was set:
484/// the flag crashed precisely the regime it exists to investigate.
485///
486/// Fixed at the print site, not inside the closure, so the greedy accept walk keeps its strict
487/// indexing (an out-of-range pred there is a real bug and must still be loud).
488fn debug_t_pred0(sampled: bool, base: usize, last_pred: u32, preds: &[u32]) -> String {
489 if base == 0 {
490 return last_pred.to_string();
491 }
492 match preds.get(base - 1) {
493 Some(p) => p.to_string(),
494 // sampled: the greedy per-column argmax was never run for this round.
495 None => {
496 debug_assert!(
497 sampled,
498 "greedy spec: preds[{}] missing at base {base}",
499 base - 1
500 );
501 "n/a".to_string()
502 }
503 }
504}
505
506/// `MEMRA_SKEY_PROBE=1` — sampled-draft-graph key probe (lane/graph-s-key-exactness-20260819).
507///
508/// Reports, per burst and per round, which draft chain the sampled arm chose and under which
509/// filter regime, plus the ONE observable that separates a legal filtered draft from a stale
510/// pure-temp graph replayed under filters: an accept test whose gathered `q` is exactly 0.
511/// A draft token sampled from the FILTERED softmax can never gather q=0 (it was drawn from the
512/// kept set), so `q=0` in the verify means the draft came from a distribution the verify does
513/// not believe in — and `u * 0 < p` then accepts it unconditionally.
514///
515/// Its own env var, deliberately NOT `MEMRA_DEBUG_SPEC`: that flag panicked the GPU worker on
516/// any sampled spec request past round 0 until this lane fixed it (§2 of the bank note).
517pub(crate) fn skey_probe() -> bool {
518 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
519 *ON.get_or_init(|| std::env::var("MEMRA_SKEY_PROBE").as_deref() == Ok("1"))
520}
521
522/// GRAMMAR HOOK for constrained spec decode (lane/constrained-full, 2026-08-03). The engine
523/// stays llguidance-agnostic: the server adapts its per-session grammar state behind this
524/// trait. CONTRACT (the verify-side truncation rule — token-identical to constrained plain
525/// greedy decode): the exactness walk runs UNMASKED first; the hook then (a) truncates
526/// acceptance at the first grammar-illegal accepted token, and (b) when the truncation fired
527/// or the bonus is illegal, the engine recomputes that slot as the MASKED argmax of the
528/// target's own verify column (an unmasked argmax that is grammar-legal IS the masked argmax
529/// — masking only removes tokens — so the common case pays nothing). `consume` advances the
530/// state with each EMITTED token in order; EOS handling is the implementor's job (skip).
531pub trait SpecConstraint {
532 /// -inf the current state's banned ids on a HOST logits row (prompt-tail / init-feed
533 /// masked argmax).
534 fn mask_logits(&mut self, logits: &mut [f32]) -> Result<(), String>;
535 /// Packed 32-bit bitset words of the CURRENT state's allowed set (device-mask form).
536 fn mask_words(&mut self) -> Result<Vec<u32>, String>;
537 /// Is `tok` consumable in the CURRENT state?
538 fn is_allowed(&mut self, tok: u32) -> Result<bool, String>;
539 /// Advance the state with an emitted token.
540 fn consume(&mut self, tok: u32) -> Result<(), String>;
541
542 // --- DRAFT-SIDE MASKING (lane/draft-mask, 2026-08-04) ---
543 // The drafter proposed grammar-illegal tokens under tight schemas, so verify-side
544 // truncation cut nearly every round (measured acceptance 0.467-0.513 tight vs 0.62-0.82
545 // loose, research/constrained-full-20260803). These three methods let the engine mask the
546 // DRAFT model's own sampling with the grammar's legal set, so proposals are legal by
547 // construction. The state they walk is a SPECULATIVE CLONE of the session matcher — the
548 // real state is advanced only by `consume` (emitted tokens), so verify-side truncation
549 // stays the correctness backstop and the emitted stream is unchanged by construction
550 // (an accepted draft is the target's unmasked argmax AND grammar-legal, hence the masked
551 // argmax; a cut slot is recomputed as the masked argmax either way).
552 // Default impls = feature OFF (pre-lane behaviour: unmasked drafts).
553
554 /// Is draft-side masking available on this hook? Probed ONCE per burst, before the draft
555 /// graph is captured (the mask is an in-graph node — its presence is a capture-time shape).
556 fn draft_mask_enabled(&self) -> bool {
557 false
558 }
559 /// Start a draft chain: clone the CURRENT (committed) grammar state into the speculative
560 /// slot. Called once per spec round, before the first draft position.
561 fn draft_begin(&mut self) -> Result<(), String> {
562 Ok(())
563 }
564 /// Packed 32-bit bitset words of the SPECULATIVE state's allowed set (target-vocab ids),
565 /// for the draft position about to be sampled. `None` = draft masking off (no-op).
566 fn draft_mask_words(&mut self) -> Result<Option<Vec<u32>>, String> {
567 Ok(None)
568 }
569 /// Advance the SPECULATIVE state with a PROPOSED draft token. `false` = the chain cannot
570 /// continue (EOS proposed, or an unmasked position proposed something illegal) — the
571 /// engine stops drafting; the token already pushed still goes through verify.
572 fn draft_advance(&mut self, _tok: u32) -> Result<bool, String> {
573 Ok(false)
574 }
575}
576
577/// DRAFT-MASK UPLOAD (lane/draft-mask): pull the speculative state's allowed set (TARGET-id
578/// space) from the hook, project it into the DRAFT head's vocab space, and upload it into the
579/// stable device buffer the draft chain reads. Returns false when the chain must stop drafting:
580/// the hook handed out no mask, or NO draft-vocab row is grammar-legal at this position (a
581/// trimmed FR-Spec head genuinely cannot propose a legal token there — masking it would leave
582/// a fully-banned row whose argmax is meaningless, so the round drafts fewer tokens and the
583/// verify emits the masked argmax as usual).
584fn upload_draft_mask(
585 e: &Engine,
586 c: &mut dyn SpecConstraint,
587 dst: &mut CudaSlice<u32>,
588 d2t: Option<&Vec<u32>>,
589 d_vocab: usize,
590 words: usize,
591) -> Result<bool, Box<dyn std::error::Error>> {
592 let Some(tw) = c
593 .draft_mask_words()
594 .map_err(|e2| format!("constraint: {e2}"))?
595 else {
596 return Ok(false);
597 };
598 let bit = |t: usize| -> bool {
599 let w = t >> 5;
600 w < tw.len() && (tw[w] >> (t & 31)) & 1 == 1
601 };
602 let mut buf = vec![0u32; words];
603 match d2t {
604 // TRIMMED draft head: row i proposes target id d2t[i] — permute the mask accordingly.
605 Some(map) => {
606 for (i, &t) in map.iter().enumerate().take(d_vocab) {
607 if bit(t as usize) {
608 buf[i >> 5] |= 1u32 << (i & 31);
609 }
610 }
611 }
612 // UNTRIMMED: draft ids ARE target ids; the packed words transfer verbatim (a short
613 // mask leaves the padded tail zeroed == banned, same rule as constrained::apply_mask).
614 None => {
615 let n = tw.len().min(words);
616 buf[..n].copy_from_slice(&tw[..n]);
617 }
618 }
619 if buf.iter().all(|w| *w == 0) {
620 return Ok(false);
621 }
622 e.htod_u32_into(dst, &buf)?;
623 Ok(true)
624}
625
626/// Keep the full token-embedding table in host memory and upload only the rows needed by each
627/// MTP/verify step. This is an exact memory-capacity seam for very large BF16 vocab tables: host
628/// gather expands the same source bits to f32, and only O(T*n_embd) bytes cross PCIe per step.
629/// CUDA-graph/round-stream draft paths require device token ids and therefore stay disabled.
630pub(crate) fn spec_host_embd() -> bool {
631 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
632 *ON.get_or_init(|| std::env::var("MEMRA_SPEC_HOST_EMBD").as_deref() == Ok("1"))
633}
634
635/// VERIFY-TIER TRUNK LAUNCH-FUSION (default ON since 2026-07-09; MEMRA_SPEC_FUSED_T=0 reverts — lane/close35b): extend
636/// the t=1 fused2/fused3 Q8_0 trunk launches to the batched verify tier (t=2-4, the K=1..3
637/// verify shapes). At t>1 the trunk pairs/triples (35B wqkv+wqkv_gate, wq/wk/wv,
638/// gate_shexp+up_shexp) each run a separate `matmul_decode_exact` — one q8_1 re-quantize of the
639/// SAME activation plus one _b2/_b4 launch per tensor. The fused twins share ONE quantize and
640/// ONE launch per group; per (tensor,token,row) the kernel body is q8_0_mmvq_batched verbatim
641/// with the identical row mapping -> BIT-IDENTICAL by construction (kernel-check pins it,
642/// run-spec K=1..8 + acceptance identity arbitrate e2e).
643pub(crate) fn spec_fused_t() -> bool {
644 static F: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
645 // DEFAULT ON since 2026-07-09 (MEMRA_SPEC_FUSED_T=0 reverts): verify t=2-4 trunk launch-fusion
646 // (fused2/fused3 Q8_0 batched twins, bit-identical by construction — m=1 block-offset split on
647 // the batched body). m=2 marginal token 2117->1762us; 35B daily: p3 +3.7% (crosses llama), p2 +5%.
648 *F.get_or_init(|| {
649 std::env::var("MEMRA_SPEC_FUSED_T")
650 .map(|v| v != "0")
651 .unwrap_or(true)
652 })
653}
654
655/// zeros/uninit switch for verify-path buffers that are FULLY OVERWRITTEN before any read.
656/// Only call this on such buffers — the lean contract is "identical bytes by construction".
657/// TOKEN-ID GUARD for every id that reaches an embed gather (#87 family).
658///
659/// A device argmax seeds its running index with 0x7FFFFFFF and replaces it only through
660/// comparisons, all of which are FALSE against NaN. An all-NaN logits row therefore returns
661/// the sentinel, and the next thing done with a token id is `embed_row(id)` — table +
662/// ~4.6 TB, never mapped, an MMU fault that kills the CUDA context for the whole process
663/// (research/pp2spec-crash-20260807). The draft chain and the GREEDY verify walk already
664/// trap this; the SAMPLED verify bonus, the boundary sampler and the replay arm's last_pred
665/// did not, which is why the recoverable fault on the greedy instrument is a TERMINAL one on
666/// the vendor-default sampled shape we actually serve.
667pub(crate) fn guard_vocab_token(
668 tok: u32,
669 n_vocab: usize,
670 what: &str,
671) -> Result<u32, Box<dyn std::error::Error>> {
672 if (tok as usize) >= n_vocab {
673 return Err(format!(
674 "{what}: token id 0x{tok:08x} >= n_vocab {n_vocab} — an all-NaN logits row left \
675 the device argmax's init sentinel in place; refusing to dereference the embed \
676 row (#87 trap)"
677 )
678 .into());
679 }
680 Ok(tok)
681}
682
683/// SPEC NaN-ORIGIN SCAN (`MEMRA_SPEC_NAN_SCAN=1`, DEFAULT OFF, diagnostic only).
684///
685/// The `#87` trap reports an all-NaN VERIFY logits column, which says the poison reached the
686/// head but not where it entered. With the scan armed the verify walk syncs and reads back
687/// every layer's output, so the FIRST layer whose residual carries a NaN names itself with the
688/// round's row and position. Off by default and never on a serving path: it costs one host
689/// sync + one `t*n_embd` D2H per layer, and the syncs change scheduling (so a run that stops
690/// reproducing under the scan is itself a datum, not an all-clear).
691///
692/// Rollback seam: unset `MEMRA_SPEC_NAN_SCAN` (or set it to 0). Every call site is behind
693/// `spec_nan_scan()`, so the default path keeps the exact launch sequence it had.
694pub(crate) fn spec_nan_scan() -> bool {
695 spec_nan_scan_level() > 0
696}
697
698/// `MEMRA_SPEC_NAN_SCAN` as a LEVEL, not a boolean. `1` scans each layer's residual, which
699/// names the layer. `2` also scans INSIDE the t-column layer body — the per-column attention
700/// output, the deferred-column o-proj/fa2 join, the post-attention norm and the routed-MoE
701/// output — because "layer 20 poisons row 0" does not say whether the attention or the routed
702/// MoE produced it, and those are different bugs with different fixes.
703pub(crate) fn spec_nan_scan_level() -> u8 {
704 static LVL: std::sync::OnceLock<u8> = std::sync::OnceLock::new();
705 *LVL.get_or_init(|| match std::env::var("MEMRA_SPEC_NAN_SCAN").as_deref() {
706 Ok("1") => 1,
707 Ok("2") => 2,
708 _ => 0,
709 })
710}
711
712/// Read back `[rows, cols]` and fail with the first NaN's coordinates. `what` names the
713/// producer (layer index, walk arm) so the error line is the localization.
714/// VERIFY-ARM RECEIPT (rides `MEMRA_SPEC_NAN_SCAN>=1`, bounded to 200 lines).
715///
716/// Names, per trunk layer, WHICH attention arm the t-column walk actually took. This exists
717/// because the level-1 residual scan below sat only on the non-fused tail: the fused
718/// rope+append+fa arm ends in `continue`, so every layer that fused was NEVER SCANNED and
719/// silently read as "clean". A poisoned residual therefore first reported at the next
720/// non-fused layer, which is how "layer 20 creates the poison" could be true of the scan and
721/// false of the engine. Also carries the row-table lookup counter, so "the fused path never
722/// ran" is distinguishable from "it ran and was innocent".
723/// KV-PLANE SCAN (`MEMRA_KV_PLANE_SCAN=1`, DEFAULT OFF, diagnostic only).
724///
725/// Reads back the STAGED rows of a layer's distributed K/V planes and reports the first row
726/// whose quantization scale is not finite. No kernel required: q8_0 blocks are
727/// `[half d][32 x i8]` and q5_1 blocks carry `half d` then `half m`, so the fp16 scale at the
728/// head of each block is host-checkable straight out of the byte plane.
729///
730/// It exists because the level-2 bad-row bitmap says EVERY verify row is non-finite at a
731/// global-attention layer's join, and row r attends a strict superset of row r-1's keys: that
732/// implicates the shared KV history those rows walk, not per-column staging. "The attention
733/// output is NaN" and "the KV history it attends is already NaN" are different bugs with
734/// different owners, and nothing measured so far separates them. A first-corrupt-row index
735/// also dates the corruption against the prime/decode boundary.
736///
737/// Bounded hard: only layers whose geometry has NO window (the global planes), only the first
738/// `MEMRA_KV_PLANE_SCAN_ROUNDS` verify rounds of a process (default 2), and it copies only
739/// `[0, staged_len)`, which is ~1.6 MB at the 1480-token repro rather than the 262144-row
740/// provision. It still syncs per layer, so it is never a serving or a measured-perf arm.
741pub(crate) fn kv_plane_scan_on() -> bool {
742 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
743 *ON.get_or_init(|| std::env::var("MEMRA_KV_PLANE_SCAN").as_deref() == Ok("1"))
744}
745
746fn kv_plane_scan_rounds() -> usize {
747 static R: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
748 *R.get_or_init(|| {
749 std::env::var("MEMRA_KV_PLANE_SCAN_ROUNDS")
750 .ok()
751 .and_then(|v| v.parse().ok())
752 .unwrap_or(2)
753 })
754}
755
756/// First non-finite fp16 block scale in `bytes`, as (block index, raw u16), scanning one
757/// scale every `stride` bytes. Returns None when every block scale is finite.
758fn first_bad_scale(bytes: &[u8], stride: usize) -> Option<(usize, u16)> {
759 if stride == 0 {
760 return None;
761 }
762 for (i, blk) in bytes.chunks_exact(stride).enumerate() {
763 let raw = u16::from_le_bytes([blk[0], blk[1]]);
764 if half_is_non_finite(raw) {
765 return Some((i, raw));
766 }
767 }
768 None
769}
770
771/// IEEE binary16: exponent all ones is Inf or NaN, whatever the mantissa says.
772fn half_is_non_finite(raw: u16) -> bool {
773 (raw & 0x7C00) == 0x7C00
774}
775
776/// Scan one layer's staged K/V planes for a non-finite quantization scale. Returns the
777/// receipt line, or None when the layer is out of scope or every scale is finite.
778pub(crate) fn scan_kv_plane(
779 e: &crate::Engine,
780 distributed: &memra_kv::ResidentTpKvCache,
781 il: usize,
782 pos0: usize,
783) -> Result<(), Box<dyn std::error::Error>> {
784 // One "round" is one pos0, not one layer: the walk visits 45 layers per verify. The
785 // default of 2 rounds is for a fault that shows up immediately; the step37 repro does not
786 // fire until rep 3 or later, i.e. round ~60 of the process, so that arm MUST raise
787 // MEMRA_KV_PLANE_SCAN_ROUNDS or it will scan only the two rounds that were never going to
788 // be poisoned and report a clean history it never looked at.
789 static ROUNDS: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
790 static LAST_POS: std::sync::atomic::AtomicUsize =
791 std::sync::atomic::AtomicUsize::new(usize::MAX);
792 if LAST_POS.swap(pos0, std::sync::atomic::Ordering::Relaxed) != pos0 {
793 ROUNDS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
794 }
795 if ROUNDS.load(std::sync::atomic::Ordering::Relaxed) > kv_plane_scan_rounds() {
796 return Ok(());
797 }
798 let staged = distributed.staged_len();
799 if staged == 0 {
800 return Ok(());
801 }
802 // ENGAGEMENT RECEIPT. This scan prints only on corruption, so `kvbad=0` in a cell would
803 // read the same whether the history was clean or the scan never ran once. Bounded so a
804 // 45-layer walk cannot flood the log.
805 static SEEN: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
806 let seen = SEEN.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
807 let (ktb, vtb) = (distributed.k_tok_bytes(), distributed.v_tok_bytes());
808 if seen < 4 {
809 eprintln!(
810 "[kv-plane] engaged #{seen} layer {il} pos0={pos0} staged={staged} \
811 ktok={ktb} vtok={vtb} (scan armed; a corrupt plane prints its own line)"
812 );
813 }
814 for rank in 0..distributed.ranks().len() {
815 let Some(rc) = distributed.rank(rank) else {
816 continue;
817 };
818 // q8_0 K blocks are [half d][32 x i8] = 34B; q5_1 V blocks lead with half d then half m.
819 let kbytes = e.dtoh_u8_view(&rc.k().slice(0..staged * ktb))?;
820 let vbytes = e.dtoh_u8_view(&rc.v().slice(0..staged * vtb))?;
821 let kbad = first_bad_scale(&kbytes, 34);
822 let vbad = first_bad_scale(&vbytes, 24);
823 if kbad.is_some() || vbad.is_some() {
824 let row = |b: Option<(usize, u16)>, tok: usize| {
825 b.map(|(i, raw)| format!("blk {i} (row {}) raw={raw:#06x}", i * 34 / tok.max(1)))
826 .unwrap_or_else(|| "clean".into())
827 };
828 eprintln!(
829 "[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",
830 row(kbad, ktb),
831 row(vbad, vtb)
832 );
833 return Ok(());
834 }
835 }
836 Ok(())
837}
838
839pub(crate) fn verify_arm_receipt(
840 arm: &str,
841 il: usize,
842 pos0: usize,
843 t: usize,
844 staged: Option<usize>,
845) {
846 static N: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
847 if N.fetch_add(1, std::sync::atomic::Ordering::Relaxed) >= 200 {
848 return;
849 }
850 eprintln!(
851 "[verify-arm] layer {il} arm={arm} pos0={pos0} t={t} staged_len={} rows_tab_lookups={}",
852 staged.map(|v| v as i64).unwrap_or(-1),
853 crate::tp::ROWS_TAB_ENGAGED.load(std::sync::atomic::Ordering::Relaxed)
854 );
855}
856
857pub(crate) fn nan_scan_rows(
858 e: &Engine,
859 buf: &CudaSlice<f32>,
860 rows: usize,
861 cols: usize,
862 what: &str,
863) -> Result<(), Box<dyn std::error::Error>> {
864 // The readback is also the ATTRIBUTION point for an asynchronous fault: a
865 // CUDA_ERROR_ILLEGAL_ADDRESS raised by any launch since the previous scan surfaces on this
866 // sync, and the bare DriverError names nothing. Wrapping it with `what` turns "the process
867 // died somewhere" into "it died at or before this layer, on this row, at this position".
868 let host = e.dtoh(buf).map_err(|err| -> Box<dyn std::error::Error> {
869 format!(
870 "spec nan-scan: sync at {what} FAILED: {err} — the fault is at or before \
871 this point in the walk"
872 )
873 .into()
874 })?;
875 if host.len() < rows * cols {
876 return Err(format!(
877 "nan-scan {what}: buffer holds {} < {rows}x{cols}",
878 host.len()
879 )
880 .into());
881 }
882 // SCAN EVERY ROW BEFORE REPORTING. A first-hit return says "row 0 is bad" and leaves the
883 // other rows UNEXAMINED, which is exactly the bit that discriminates the two mechanisms: in
884 // the t-column verify, row 0 attends keys [0..p+1) and row 1 attends [0..p+2), a strict
885 // superset, so poison in the SHARED KV history must appear in BOTH rows, while poison in
886 // per-column staging can appear in one. Report the whole map.
887 let mut per_row: Vec<usize> = Vec::with_capacity(rows);
888 let mut first_bad: Option<(usize, usize)> = None;
889 for r in 0..rows {
890 let row = &host[r * cols..(r + 1) * cols];
891 let bad = row.iter().filter(|v| !v.is_finite()).count();
892 per_row.push(bad);
893 if bad > 0 && first_bad.is_none() {
894 first_bad = Some((r, row.iter().position(|v| !v.is_finite()).unwrap_or(0)));
895 }
896 }
897 if let Some((r0, c0)) = first_bad {
898 let map: String = per_row
899 .iter()
900 .map(|&b| if b == 0 { '.' } else { 'X' })
901 .collect();
902 return Err(format!(
903 "spec nan-scan: {what} produced non-finite values — rows[{rows}] map={map} \
904 counts={per_row:?} of {cols} each; first at row {r0} element {c0}. Both rows bad \
905 implicates shared state (the KV history this layer reads); one row bad implicates \
906 per-column staging."
907 )
908 .into());
909 }
910 Ok(())
911}
912
913fn vbuf(e: &Engine, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
914 if spec_lean() { e.uninit(n) } else { e.zeros(n) }
915}
916
917/// Scratch KV for the MTP block (one full-attn layer).
918///
919/// PERSISTENT MODE (default, 2026-07-03 — the acceptance lever): sized cap = max_ctx and kept in
920/// sync with the COMMITTED sequence — slot p holds the MTP block's K/V for committed token p
921/// (roped p+1, the chain's rope convention), so the draft chain's self-attention sees the FULL
922/// committed history instead of only the current round's 1..K+1 chain tokens (the reference
923/// engine's "mtp_update" design). Entries come from two sources:
924/// - chain appends: accepted positions KEEP their chain-computed entries (embedding exact,
925/// hidden chain-approximate — the reference engine accepts the same);
926/// - `mtp_kv_fill` batches: prompt positions + the last-draft position on full accept, computed
927/// from EXACT trunk hiddens (K/V-only MTP-block pass, no attention/FFN/lm_head).
928/// Rejected drafts / p-min extras / pseudo-seed appends are all discarded by the round-start
929/// `set_len` truncation (the KvLayer len mechanism — §C rollback for the draft side).
930/// Multi-turn spec-decode session (2026-07-05): trunk Cache + persistent MTP draft scratch +
931/// the committed token list, alive across generate_spec_session calls. Turn N+1 primes ONLY its
932/// suffix (chunked continuation prime over the quantized past) and mtp_kv_fill's its suffix rows,
933/// then the round loop runs unchanged. `last_h` carries the pre-output_norm hidden of the last
934/// committed row across turns (the predecessor-pairing seed + fill anchor).
935/// Per-request sampling config for the sampled-spec serve path.
936#[derive(Clone, Copy, Debug)]
937pub struct SpecSampling {
938 pub temp: f32,
939 pub seed: u64,
940 pub top_k: i32, // 0 = off
941 pub top_p: f32, // 1.0 = off
942 pub min_p: f32, // 0.0 = off
943 pub penalty_last_n: usize, // 0 = penalties off
944 pub penalty_repeat: f32,
945 pub penalty_freq: f32,
946 pub penalty_present: f32,
947}
948
949impl SpecSampling {
950 /// Non-identity penalties requested — THE `pen_on` predicate (one definition; the
951 /// same group-off rule `SamplerIdentity::of` canonicalizes: a window with neutral
952 /// coefficients is penalties-absent). Both spec routes and the dspark accept walk
953 /// key their penalty arms off this.
954 pub fn pen_on(&self) -> bool {
955 self.penalty_last_n > 0
956 && (self.penalty_repeat != 1.0
957 || self.penalty_freq != 0.0
958 || self.penalty_present != 0.0)
959 }
960}
961
962/// Host Philox4x32-10 uniform in (0,1) — mirrors spec_sample.cu's `philox4`/`u01` with the
963/// ctr_lo tag 0xFFFF_FFFE, so the host accept-test stream never collides with any device
964/// sampling event (device Gumbel uses (i>>2, stream_pos); device residual uses 0xFFFF_FFFD).
965/// One value per (seed, ctr) EVENT; callers own the counter discipline. Extracted verbatim
966/// from generate_spec_inner2's closure for the dspark sampled-admission walk (the two paths
967/// MUST consume the identical stream construction — two ad-hoc Philox copies drifting apart
968/// is a distributional bug, not a style problem).
969pub(crate) fn host_u01(seed: u64, ctr: u32) -> f32 {
970 let (m0, m1) = (0xD2511F53u32, 0xCD9E8D57u32);
971 let (mut c0, mut c1, mut c2, mut c3) = (0xFFFF_FFFEu32, ctr, 0u32, 0u32);
972 let (mut k0, mut k1) = ((seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
973 for _ in 0..10 {
974 let (h0, l0) = (((m0 as u64 * c0 as u64) >> 32) as u32, m0.wrapping_mul(c0));
975 let (h1, l1) = (((m1 as u64 * c2 as u64) >> 32) as u32, m1.wrapping_mul(c2));
976 let (n0, n1, n2, n3) = (h1 ^ c1 ^ k0, l1, h0 ^ c3 ^ k1, l0);
977 c0 = n0;
978 c1 = n1;
979 c2 = n2;
980 c3 = n3;
981 k0 = k0.wrapping_add(0x9E3779B9);
982 k1 = k1.wrapping_add(0xBB67AE85);
983 }
984 (c0 as f32 + 1.0) * (1.0 / 4294967296.0)
985}
986
987/// Tracked draft positions for [`SpecTelemetry`] (serve K defaults to 3; the run-spec gate
988/// sweeps K=1..8, and MEMRA_SPEC_CAPMAX defaults to 7 — 8 covers every tuned config).
989pub const SPEC_TELEM_POS: usize = 8;
990
991/// Always-on per-draft-position acceptance telemetry (lane/accept-telemetry, 2026-08-05 —
992/// the llama.cpp #26389 / vLLM spec-decode counter schema, upstream-sweeps 2026-08-05).
993/// Lives on the [`SpecSession`] and accumulates across bursts; the serve worker diffs a
994/// stashed copy per burst for its per-model /metrics aggregation and per-request usage.
995/// Same normalization as the `[spec-stats]` line: p-min-discarded chain tokens are counted
996/// in NEITHER drafted nor accepted.
997#[derive(Clone, Copy, Default, Debug)]
998pub struct SpecTelemetry {
999 /// verify rounds completed (a round-stream burst counts each of its M rounds).
1000 pub rounds: u64,
1001 /// tokens drafted / accepted across all rounds.
1002 pub drafted: u64,
1003 pub accepted: u64,
1004 /// how often draft position j (0-based within a round's chain) was offered / accepted.
1005 /// Positions >= SPEC_TELEM_POS are untracked (totals still count them). The opt-in
1006 /// round-stream arm (MEMRA_SPEC_STREAM=1) reads back only totals, so under it these
1007 /// arrays cover the standard-path rounds only and their sums may undercount the totals.
1008 pub pos_drafted: [u64; SPEC_TELEM_POS],
1009 pub pos_accepted: [u64; SPEC_TELEM_POS],
1010}
1011
1012impl SpecTelemetry {
1013 /// Fieldwise `self - prev` — the worker's per-burst delta off a copy stashed before the
1014 /// burst call. Saturating: a caller diffing against the wrong snapshot gets zeros, not
1015 /// a wrapped counter.
1016 pub fn delta_since(&self, prev: &SpecTelemetry) -> SpecTelemetry {
1017 let mut d = SpecTelemetry {
1018 rounds: self.rounds.saturating_sub(prev.rounds),
1019 drafted: self.drafted.saturating_sub(prev.drafted),
1020 accepted: self.accepted.saturating_sub(prev.accepted),
1021 ..Default::default()
1022 };
1023 for j in 0..SPEC_TELEM_POS {
1024 d.pos_drafted[j] = self.pos_drafted[j].saturating_sub(prev.pos_drafted[j]);
1025 d.pos_accepted[j] = self.pos_accepted[j].saturating_sub(prev.pos_accepted[j]);
1026 }
1027 d
1028 }
1029 /// Fieldwise `self += d` — the worker's per-model aggregation.
1030 pub fn merge(&mut self, d: &SpecTelemetry) {
1031 self.rounds += d.rounds;
1032 self.drafted += d.drafted;
1033 self.accepted += d.accepted;
1034 for j in 0..SPEC_TELEM_POS {
1035 self.pos_drafted[j] += d.pos_drafted[j];
1036 self.pos_accepted[j] += d.pos_accepted[j];
1037 }
1038 }
1039
1040 /// Mean accepted draft-prefix length per verify round (tau).
1041 pub fn tau(&self) -> f64 {
1042 if self.rounds > 0 {
1043 self.accepted as f64 / self.rounds as f64
1044 } else {
1045 0.0
1046 }
1047 }
1048}
1049
1050/// Session-lifetime atomic acceptance counters. The verifier records only after the greedy or
1051/// rejection-sampling walk has resolved on the host, so these relaxed increments add no GPU
1052/// launch, synchronization, allocation, or ordering dependency to the numeric path.
1053struct SpecTelemetryCounters {
1054 rounds: AtomicU64,
1055 drafted: AtomicU64,
1056 accepted: AtomicU64,
1057 pos_drafted: [AtomicU64; SPEC_TELEM_POS],
1058 pos_accepted: [AtomicU64; SPEC_TELEM_POS],
1059}
1060
1061impl Default for SpecTelemetryCounters {
1062 fn default() -> Self {
1063 Self {
1064 rounds: AtomicU64::new(0),
1065 drafted: AtomicU64::new(0),
1066 accepted: AtomicU64::new(0),
1067 pos_drafted: std::array::from_fn(|_| AtomicU64::new(0)),
1068 pos_accepted: std::array::from_fn(|_| AtomicU64::new(0)),
1069 }
1070 }
1071}
1072
1073impl SpecTelemetryCounters {
1074 fn record_round(&self, drafted: usize, accepted: usize) {
1075 debug_assert!(accepted <= drafted);
1076 self.rounds.fetch_add(1, Ordering::Relaxed);
1077 self.drafted.fetch_add(drafted as u64, Ordering::Relaxed);
1078 self.accepted.fetch_add(accepted as u64, Ordering::Relaxed);
1079 for counter in self.pos_drafted.iter().take(drafted) {
1080 counter.fetch_add(1, Ordering::Relaxed);
1081 }
1082 for counter in self.pos_accepted.iter().take(accepted) {
1083 counter.fetch_add(1, Ordering::Relaxed);
1084 }
1085 }
1086
1087 /// Round-stream keeps each round's accept length on device; retain exact scalar totals while
1088 /// leaving the per-position arrays untouched, matching the pre-existing telemetry contract.
1089 fn record_totals(&self, rounds: usize, drafted: usize, accepted: usize) {
1090 self.rounds.fetch_add(rounds as u64, Ordering::Relaxed);
1091 self.drafted.fetch_add(drafted as u64, Ordering::Relaxed);
1092 self.accepted.fetch_add(accepted as u64, Ordering::Relaxed);
1093 }
1094
1095 fn snapshot(&self) -> SpecTelemetry {
1096 SpecTelemetry {
1097 rounds: self.rounds.load(Ordering::Relaxed),
1098 drafted: self.drafted.load(Ordering::Relaxed),
1099 accepted: self.accepted.load(Ordering::Relaxed),
1100 pos_drafted: std::array::from_fn(|j| self.pos_drafted[j].load(Ordering::Relaxed)),
1101 pos_accepted: std::array::from_fn(|j| self.pos_accepted[j].load(Ordering::Relaxed)),
1102 }
1103 }
1104}
1105
1106pub struct SpecSession {
1107 pub(crate) cache: Cache,
1108 pub(crate) scratch: MtpScratch,
1109 /// Every token whose state the caches hold, in order (prompt turns + generated), INCLUDING
1110 /// overshoot: spec commits accepted drafts past max_new; those rows are in the caches, so the
1111 /// session must count them. Callers render output from this, not from their own echo.
1112 pub committed: Vec<u32>,
1113 /// Pre-output_norm hidden of the LAST committed row (device). None before the first turn.
1114 pub(crate) last_h: Option<CudaSlice<f32>>,
1115 /// Greedy argmax predicting the token AFTER committed.last() (from the last turn's final
1116 /// logits). Fuels empty-suffix continuation bursts (serve): the next turn emits this token
1117 /// first, feeds it, and the round loop resumes without any prime. None before the first turn.
1118 pub next_pred: Option<u32>,
1119 /// SAMPLED-SPEC stream continuity across bursts: Philox event counters persist here so a
1120 /// session's randomness never repeats between generate_spec_session calls. (0,0) at admit.
1121 pub sctr: u32,
1122 pub uctr: u32,
1123 /// PERSISTENT DRAFT-GRAPH CONTEXT (2026-08-01, the serve-burst fixed-cost fix): the captured
1124 /// draft graph(s) + every device I/O buffer they bake, carried ACROSS generate_spec_session
1125 /// calls. Before this, every serve burst re-captured the draft graph (2 warmup forwards +
1126 /// instantiate) — measured ~16ms/burst on H100 q27 (MEMRA_SPEC_BURST sweep,
1127 /// research/spec-serving-20260801). None before the first turn; error paths drop it
1128 /// (next burst recaptures — serve retires errored sessions anyway).
1129 pub(crate) draft_ctx: Option<DraftGraphCtx>,
1130 /// PENDING-CARRY across bursts (2026-08-01, the serve burst-boundary fix): the bonus token
1131 /// emitted by the last round but NOT committed to the caches. The old tail committed it with
1132 /// a solo T=1 trunk pass (+ draft fill), and the next burst's setup fed the stashed next_pred
1133 /// with ANOTHER solo pass — 2x ~11.5ms/burst measured on H100 q27 ([spec-setup] trace).
1134 /// Carrying it lets the next empty-suffix greedy burst consume it as round-0 verify col 0,
1135 /// exactly like a mid-burst full-accept boundary (no solo passes). INVARIANT: when set,
1136 /// `committed` (== cache rows) EXCLUDES this token although it was already emitted in the
1137 /// last burst's output, and `last_h` holds the hidden of the last COMMITTED row (its
1138 /// predecessor — the chain-seed/fill anchor). `next_pred` is None (unknown without the
1139 /// commit pass). Non-empty-suffix or sampled turns must flush first (spec_flush_pending);
1140 /// generate_spec_session_sampled does this at entry, and serve parks only flushed sessions.
1141 pub pending_tok: Option<u32>,
1142 /// SESSION-AFFINITY TURN CHECKPOINT (lane/session-affinity, 2026-08-05): the state at this
1143 /// turn's PROMPT-END boundary, retained so a later turn can REWIND here. See
1144 /// [`SpecCheckpoint`]. Refreshed by every non-empty prime; None until the first one, and on
1145 /// a rig too tight to hold it (a failed capture is silent — resume just isn't available).
1146 pub(crate) turn_ckpt: Option<SpecCheckpoint>,
1147 /// Session-lifetime acceptance telemetry. Relaxed atomics update at the host-side round
1148 /// accounting the loop already does — no syncs, no allocation. NOTE a
1149 /// pool-resumed session carries the PREVIOUS requests' counts; per-request consumers
1150 /// diff with [`SpecTelemetry::delta_since`] around each burst.
1151 telem: SpecTelemetryCounters,
1152 /// PREFIX-CACHE publication request (lane/spec-prefix-cache): worker sets this to the
1153 /// miss-LCP boundary before a cold burst; the prime captures at exactly that split (it must
1154 /// coincide with the burst's `prime_split` or no capture happens). One-shot: consumed by the
1155 /// prime, result lands in `boundary_captures`.
1156 pub capture_at: Option<usize>,
1157 /// The captures the last prime produced (see [`SpecBoundaryCapture`]). Worker drains them
1158 /// post-burst to assemble prefix entries. A failed capture is silent, like `turn_ckpt` —
1159 /// publication just isn't available for that request. Plural since
1160 /// lane/frspec-multiturn-cache (2026-08-21): a cold burst can capture BOTH the miss-LCP
1161 /// split (the shared-prefix class) and the stable pre-generation boundary (the
1162 /// next-turn re-render class) — one entry per stop, exactly the boundary set the plain
1163 /// prefill tick publishes/checkpoints.
1164 pub boundary_captures: Vec<SpecBoundaryCapture>,
1165 /// STABLE-BOUNDARY TURN CHECKPOINT REQUEST (lane/frspec-multiturn-cache, 2026-08-21): the
1166 /// ABSOLUTE committed-length position the next non-empty prime should capture `turn_ckpt`
1167 /// at, instead of prompt-end. The worker sets it to the STABLE PRE-GENERATION boundary
1168 /// (`plain_checkpoint_boundary` — before the live generation header the client rewrites),
1169 /// porting the 2026-08-09 plain-tier fix: a prompt-end spec checkpoint includes the
1170 /// template's live assistant-generation header (`<|im_start|>assistant\n<think>\n`), which
1171 /// the NEXT turn's re-render replaces, so `affinity_match` diverged a couple tokens below
1172 /// the checkpoint and the spec pool declined 100% of multi-turn agent traffic (measured:
1173 /// `spec-affinity: declined (history diverged at 6811 of checkpoint 6813)`,
1174 /// research/multiturn-cache-20260821 B4). One-shot, `capture_at` convention; None = legacy
1175 /// prompt-end capture.
1176 pub ckpt_at: Option<usize>,
1177}
1178impl SpecSession {
1179 /// Context capacity of the session's caches (the server's ContextFull guard).
1180 pub fn cache_max_ctx(&self) -> usize {
1181 self.cache.max_ctx
1182 }
1183 /// Read access to the live trunk cache (lane/spec-prefix-cache): the worker slices
1184 /// full-attn KV rows `[0..capture.pos)` out of it when publishing a boundary capture —
1185 /// those rows are append-only for the session's lifetime (rollbacks never truncate below
1186 /// the prime boundary), so no copy was taken at prime time.
1187 pub fn cache_ref(&self) -> &Cache {
1188 &self.cache
1189 }
1190 /// Read access to the persistent draft-scratch plane (lane/spec-on-cache-hit): the
1191 /// worker slices rows `[0..capture.pos)` when publishing a boundary capture, exactly
1192 /// like the trunk KV — draft rows below the prompt end are append-only for the
1193 /// session's lifetime (the prime fill wrote them once; rollbacks reset `len_d` to the
1194 /// committed length, never below the prime boundary, and the true-hidden refresh
1195 /// rewrites generated positions only). Returns `(k, v, k_tok_bytes, v_tok_bytes)`.
1196 /// None when the scratch is ring-backed (Step35 SWA — physical rows are not
1197 /// prefix-addressable; the prefix cache already refuses that class end to end).
1198 pub fn draft_plane_ref(&self) -> Option<(&CudaSlice<u8>, &CudaSlice<u8>, usize, usize)> {
1199 if self.scratch.kv.ring.is_some() {
1200 return None;
1201 }
1202 Some((
1203 &self.scratch.kv.k,
1204 &self.scratch.kv.v,
1205 self.scratch.kv.k_tok_bytes,
1206 self.scratch.kv.v_tok_bytes,
1207 ))
1208 }
1209 /// Snapshot the session's process-local acceptance counters for per-burst diffing.
1210 pub fn telemetry(&self) -> SpecTelemetry {
1211 self.telem.snapshot()
1212 }
1213 /// Committed position this session can REWIND to (its retained prompt-end boundary), if any.
1214 /// A request whose prompt matches `committed[..pos]` exactly can resume from here — see
1215 /// `spec_rewind_to_checkpoint`.
1216 pub fn rewind_pos(&self) -> Option<usize> {
1217 self.turn_ckpt.as_ref().map(|c| c.pos)
1218 }
1219 /// Whether every ring-backed trunk/draft row needed by the retained checkpoint is resident.
1220 pub fn rewind_is_resident(&self) -> bool {
1221 self.turn_ckpt.as_ref().is_some_and(|ckpt| {
1222 self.cache.can_rollback(&ckpt.snap, 0) && self.scratch.can_rewind_to(ckpt.pos)
1223 })
1224 }
1225 /// Is this session in the DEMOTION-READY shape (see [`SpecSession::into_demoted`])?
1226 /// `false` means a carried pending must be flushed first (`spec_flush_pending`), or the
1227 /// session has never run a turn and has no prediction to hand over.
1228 pub fn demote_ready(&self) -> bool {
1229 self.pending_tok.is_none() && self.next_pred.is_some()
1230 }
1231 /// Does this session hold a carried pending bonus (flush required before a handoff/park)?
1232 pub fn has_pending(&self) -> bool {
1233 self.pending_tok.is_some()
1234 }
1235 /// Committed row count == cache rows (the session invariant), for the caller's own
1236 /// `fed`-length cross-check at a handoff boundary.
1237 pub fn committed_len(&self) -> usize {
1238 self.committed.len()
1239 }
1240 /// DEMOTION HANDOFF (lane/spec-gate, 2026-08-07): consume this session and hand its trunk
1241 /// cache + next-token prediction to the plain batched-decode path.
1242 ///
1243 /// WHY THIS IS EXACT (greedy). The invariant at a burst boundary is `cache.pos ==
1244 /// committed.len()`: every committed row has trunk KV + recurrent state, exactly as a plain
1245 /// tokenwise prime of the same `committed` sequence would have left it (that is the
1246 /// session-tail contract, and the same property `spec_rewind_to_checkpoint` and the reuse
1247 /// pool already rely on). `next_pred` is the argmax of the verify's logits for the LAST
1248 /// committed row — and verify-column logits are bit-identical to plain decode's logits at
1249 /// that position, because `matmul_decode_exact` bit-identity IS the basis of the greedy
1250 /// accept walk. So handing (cache, next_pred) to the batched path continues the stream from
1251 /// a state indistinguishable from one the batched path produced itself: the batched tick
1252 /// emits `next_pred`, feeds it into this same cache, and decodes on.
1253 ///
1254 /// `None` when the session is not in the handoff shape — a carried pending (its bonus row is
1255 /// NOT in the cache, so `spec_flush_pending` must commit it first) or no `next_pred` yet
1256 /// (never bursted). Callers must not force it: a half-committed cache handed to the batched
1257 /// path would silently skip a token.
1258 ///
1259 /// The MTP draft scratch, the persistent draft-graph context and the turn checkpoint are
1260 /// DROPPED here (freeing their VRAM): the batched path never drafts, and this handoff is
1261 /// one-way by design — there is no cheap symmetric re-promotion (rebuilding the draft KV
1262 /// would mean an `mtp_kv_fill` over the whole committed history).
1263 pub fn into_demoted(self) -> Option<(Cache, u32)> {
1264 if self.pending_tok.is_some() {
1265 return None;
1266 }
1267 let np = self.next_pred?;
1268 debug_assert_eq!(
1269 self.cache.pos,
1270 self.committed.len(),
1271 "demotion handoff: cache rows != committed tokens"
1272 );
1273 Some((self.cache, np))
1274 }
1275 /// Pool-resume hook (audit Q2): clear the parked draft-graph failure memoization so a
1276 /// NEW request resuming this session gets one fresh capture chance — a transient-pressure
1277 /// capture failure must not persist for the pool's whole lifetime (the TRT #16072 class).
1278 /// Logs once iff a flag was actually set; a no-fallback resume is silent and free.
1279 pub fn reset_graph_fallback_on_resume(&mut self) {
1280 if let Some(line) = self
1281 .draft_ctx
1282 .as_mut()
1283 .and_then(|c| c.failed.reset_on_resume())
1284 {
1285 eprintln!("{line}");
1286 }
1287 }
1288}
1289
1290/// A session's PROMPT-END boundary state, the rewind target for session-affinity resume.
1291///
1292/// WHY THIS BOUNDARY, AND WHY IT IS THE ONLY ONE WORTH KEEPING. The rewrite class this lane
1293/// exists for (a client that strips `<think>` blocks out of prior assistant turns) mutates the
1294/// text the session GENERATED, never the prompt it was given. So turn N's prompt agrees with
1295/// turn N-1's committed tokens up to almost exactly where turn N-1's generation began — the
1296/// prompt-end boundary. Keeping a checkpoint there means the next turn re-primes only its own
1297/// delta (the rewritten answer + the new user turn) instead of the whole conversation.
1298///
1299/// WHAT IT MUST HOLD. Full-attn KV is append-only and position-addressed, so rewinding it is a
1300/// `len` truncation (no data). Linear-attn (GDN) conv/ssm state is mutated IN PLACE with no
1301/// position index, so it must be a real device COPY — that copy is the entire reason a spec
1302/// session could not previously rewind. The MTP draft scratch needs no copy either: its rows
1303/// below the boundary were written by this turn's fill and are never revisited (the per-round
1304/// true-hidden refresh only rewrites the CURRENT burst's committed positions), so rewinding it
1305/// is also just a `len` reset. `last_h` is the hidden of the last row below the boundary — the
1306/// predecessor-pairing anchor the next prime's fill reads for its first row.
1307///
1308/// COST: one `Cache::snapshot` per TURN, on a code path that already takes one per ROUND.
1309pub(crate) struct SpecCheckpoint {
1310 snap: crate::cache::CacheSnapshot,
1311 /// Committed length at the boundary (== cache.pos there, the session invariant).
1312 pos: usize,
1313 /// Pre-output_norm hidden of row `pos - 1`.
1314 last_h: CudaSlice<f32>,
1315}
1316
1317/// PREFIX-CACHE BOUNDARY CAPTURE (lane/spec-prefix-cache, 2026-08-14): the state a spec session
1318/// records at its cold-prime split so the WORKER can publish a cross-request prefix entry —
1319/// the commit-gated-publication port (research/cache-spec-design-20260814/PORT-PLAN.md item 1).
1320/// Only the pieces that are DESTROYED by continuing the prime need copies here: the in-place
1321/// GDN conv/ssm states (via `Cache::snapshot`, same mechanism as [`SpecCheckpoint`]) and the
1322/// boundary logits. Full-attn KV rows `[0..pos)` and draft-scratch rows `[0..pos)` are
1323/// append-only for the session's lifetime (rollbacks never truncate below the prime boundary),
1324/// so the worker slices those from the live caches post-burst instead of copying at prime time.
1325pub struct SpecBoundaryCapture {
1326 pub snap: crate::cache::CacheSnapshot,
1327 /// Token boundary (== cache.pos at capture; == the worker's miss-LCP split).
1328 pub pos: usize,
1329 /// Full-vocab logits after the prefix prime — the entry's boundary logits.
1330 pub logits: Vec<f32>,
1331 /// Pre-output_norm trunk hidden of row `pos - 1` (lane/spec-on-cache-hit): the
1332 /// predecessor-pairing anchor a RESTORED spec session's first suffix-fill row reads
1333 /// (the `SpecSession::last_h` convention). Empty = unavailable (capture stays valid;
1334 /// the fill's zeros row-0 fallback covers it at a bounded acceptance cost).
1335 pub last_h: Vec<f32>,
1336}
1337
1338/// D2H one hidden row out of a `[T, n_embd]` prime hidden stack — the boundary anchor a
1339/// spec boundary capture carries for later restored-session fills. Failure is silent
1340/// (`turn_ckpt` convention): the capture publishes without an anchor.
1341fn capture_boundary_hidden(
1342 e: &Engine,
1343 h_rows: &CudaSlice<f32>,
1344 pos: usize,
1345 n_embd: usize,
1346) -> Vec<f32> {
1347 if pos == 0 || h_rows.len() < pos * n_embd {
1348 return Vec::new();
1349 }
1350 let Ok(mut row) = e.uninit(n_embd) else {
1351 return Vec::new();
1352 };
1353 if e.copy_view_into(
1354 &mut row,
1355 0,
1356 &h_rows.slice((pos - 1) * n_embd..pos * n_embd),
1357 n_embd,
1358 )
1359 .is_err()
1360 {
1361 return Vec::new();
1362 }
1363 e.dtoh(&row).unwrap_or_default()
1364}
1365
1366/// ROLLBACK DOOR for sampled BOUNDARY tokens (lane/sampled-spec-quality, 2026-08-19).
1367/// Default ON: the token a burst emits at its own boundary is drawn from the request's
1368/// sampler. `MEMRA_SPEC_SAMPLED_BOUNDARY=0` restores the pre-lane posture (an ARGMAX at
1369/// every boundary) without touching greedy, which is byte-unaffected either way.
1370pub fn spec_sampled_boundary_on() -> bool {
1371 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1372 *ON.get_or_init(|| std::env::var("MEMRA_SPEC_SAMPLED_BOUNDARY").as_deref() != Ok("0"))
1373}
1374
1375/// ROLLBACK DOOR for SESSION-SPANNING penalty history (lane/sampled-spec-quality).
1376/// Default ON: `pen_hist` is seeded from the session's committed tail, so repetition /
1377/// frequency / presence penalties see the whole stream. `MEMRA_SPEC_PEN_SESSION=0`
1378/// restores the pre-lane posture (each burst restarts the window from its own prompt
1379/// slice, i.e. from NOTHING on a continuation burst) — and with the door shut the worker
1380/// must keep refusing penalized sampled prefix-cache restores, because the restored
1381/// session's continuation burst is handed no prompt slice at all.
1382pub fn spec_pen_session_on() -> bool {
1383 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1384 *ON.get_or_init(|| std::env::var("MEMRA_SPEC_PEN_SESSION").as_deref() != Ok("0"))
1385}
1386
1387/// ROLLBACK DOOR for extended-entry publication from a RESTORED session
1388/// (lane/sampled-spec-quality, Item 3). Default ON: a converted prefix-cache hit that fed a
1389/// suffix captures its own prompt-end boundary so the NEXT turn can hit a longer prefix.
1390/// `MEMRA_SPEC_RESTORE_REPUBLISH=0` restores the pre-lane posture (a namespace learns exactly
1391/// one boundary and never advances it). Whole-entry semantics only — the boundary is the
1392/// restored session's own prompt end, so `entry_pos != fed_len` still refuses on the way in.
1393pub fn spec_restore_republish_on() -> bool {
1394 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1395 *ON.get_or_init(|| std::env::var("MEMRA_SPEC_RESTORE_REPUBLISH").as_deref() != Ok("0"))
1396}
1397
1398/// Diagnostics: name every boundary token on stderr (`MEMRA_SPEC_BOUNDARY_TRACE=1`), with
1399/// the argmax the pre-lane code would have emitted from the same row. This is how the
1400/// lane MEASURES the boundary rate and the deviation rate instead of estimating them.
1401fn spec_boundary_trace() -> bool {
1402 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1403 *ON.get_or_init(|| std::env::var("MEMRA_SPEC_BOUNDARY_TRACE").as_deref() == Ok("1"))
1404}
1405
1406/// llama-parity floor for the penalty window when the request does not ask for a bigger
1407/// one (`repeat_last_n` default). The serve API arms `penalty_last_n = PEN_WINDOW_MAX` for any
1408/// non-identity penalty, so this floor only matters to explicit small windows and to the
1409/// CLI env path.
1410const PEN_WINDOW_FLOOR: usize = 64;
1411
1412/// CEILING on the penalty window, and it is a COST bound, not a semantic preference.
1413/// `penalize_logits_f32` (cu/spec_sample.cu) dedups on device by having thread `i` scan
1414/// `hist[0..i]`, so a pass is O(n_hist²) and it runs ~3x per verify round (the q rows, the
1415/// p column, the bonus column). The serve API uses this same bound for every non-identity
1416/// penalty so host/plain, sparse-device, and speculative sampling cannot change logits on
1417/// admission demotion. An uncapped 128k-token history would put ~1.7e10
1418/// comparisons per pass, tens of ms per round, i.e. penalties would silently destroy decode
1419/// throughput on exactly the long-context requests that most want them. 8192 keeps a pass
1420/// at ~7e7 comparisons (tens of microseconds) while still being **128x wider than the
1421/// pre-lane effective window** (64 prompt-tail tokens + whatever the current burst had
1422/// generated). A request that genuinely needs a window beyond this wants host-side dedup +
1423/// counts through a new kernel signature — a follow-up lane, named here rather than hidden.
1424/// `pub` since lane/dspark-penalized-sampled-20260821: the dspark route's accept walk and
1425/// the dspark_sample_gate binary trim their uploads with the SAME cap — a second constant
1426/// is a second thing to drift.
1427pub const PEN_WINDOW_MAX: usize = 8192;
1428
1429/// Seed a penalty window over the SESSION, not the burst (lane/sampled-spec-quality,
1430/// Item 2). The window is the last `max(penalty_last_n, 64)` tokens of
1431/// `session_committed ++ burst_prompt` — for a cold turn-1 burst (`session_committed`
1432/// empty, default `penalty_last_n`) that is byte-identically the pre-lane
1433/// `prompt.iter().rev().take(64).rev()`; for a continuation burst it is the stream the
1434/// client actually asked us to penalize, where the pre-lane code had NOTHING.
1435/// `pub` since lane/dspark-penalized-sampled-20260821: the dspark route seeds its session
1436/// window through the SAME function (one definition of "the window" across both spec
1437/// routes and the gate binary's trunk-only reference arm).
1438pub fn pen_window_seed(
1439 session_committed: &[u32],
1440 burst_prompt: &[u32],
1441 penalty_last_n: usize,
1442) -> Vec<u32> {
1443 let win = penalty_last_n.clamp(PEN_WINDOW_FLOOR, PEN_WINDOW_MAX);
1444 let take_prompt = burst_prompt.len().min(win);
1445 let take_sess = (win - take_prompt).min(session_committed.len());
1446 let mut hist = Vec::with_capacity(take_sess + take_prompt);
1447 hist.extend_from_slice(&session_committed[session_committed.len() - take_sess..]);
1448 hist.extend_from_slice(&burst_prompt[burst_prompt.len() - take_prompt..]);
1449 hist
1450}
1451
1452/// Draw a BOUNDARY token from the target distribution the request asked for
1453/// (lane/sampled-spec-quality, Item 1) — the fix for "sampled spec emits an ARGMAX token at
1454/// every burst boundary".
1455///
1456/// WHY THIS EXISTS. A spec burst's first emitted token is not produced by the accept walk:
1457/// it comes off a logits row that already exists (the prime's last row on a cold burst; the
1458/// row after the last committed token on a continuation burst; the prefix-cache entry's
1459/// boundary row on a restored one). Pre-lane that token was `argmax` in BOTH sampling
1460/// regimes, so a sampled stream took a greedy token once per burst — measured, not
1461/// estimated, in research/spec-cache-20260818/SAMPLED-QUALITY.md. At temperature > 0 the
1462/// customer asked for a sampled token, so this draws one.
1463///
1464/// THE PROGRAM IS THE FULL-ACCEPT BONUS'S PROGRAM, deliberately: penalize the row (over the
1465/// session's window), take this row's OWN filter stats (the sampfix-20260805 law — stats
1466/// from a neighbour row mis-scale every `e0` and can wipe the row to token 0), gumbel-perturb
1467/// with the session's Philox stream at `*sctr`, argmax the perturbed row. Reusing the bonus's
1468/// composition means `sample_check`'s distributional oracle covers this draw too, and the
1469/// boundary token is drawn from the same filtered/penalized `p` the accept walk targets.
1470///
1471/// THE STREAM IS THE SESSION'S, NOT A FRESH ONE. `sctr` is the caller's live counter and is
1472/// advanced by exactly one, so a boundary draw consumes the next value in the same Philox
1473/// stream the accept walk uses — never a second, independently seeded stream (which would be
1474/// a new distributional bug: two streams from one seed correlate wherever their counters
1475/// collide). That also makes a restored session's boundary draw at `sctr == 0` bit-identical
1476/// to the cold session's own first draw from the same logits row, which is what preserves the
1477/// sampled-hit lane's per-seed hit==cold byte identity.
1478#[allow(clippy::too_many_arguments)]
1479pub fn sample_boundary_token_dev(
1480 e: &Engine,
1481 logits: &CudaSlice<f32>,
1482 n_vocab: usize,
1483 sp: &SpecSampling,
1484 pen_hist: &[u32],
1485 sctr: &mut u32,
1486 site: &str,
1487) -> Result<u32, Box<dyn std::error::Error>> {
1488 debug_assert!(
1489 sp.temp > 0.0,
1490 "boundary sampling is the sampled regime only"
1491 );
1492 // Own copy: penalize_logits mutates in place and the caller's row is live state
1493 // (prime_logits back the constrained recompute; last_col_logits backs round 0's accept).
1494 let mut col = e.zeros(n_vocab)?;
1495 e.copy_into(&mut col, 0, logits, n_vocab)?;
1496 let pen_on = sp.penalty_last_n > 0
1497 && (sp.penalty_repeat != 1.0 || sp.penalty_freq != 0.0 || sp.penalty_present != 0.0);
1498 if pen_on && !pen_hist.is_empty() {
1499 // window trim mirrors the round loop's own upload (`pen_hist[w0..]`), cap included.
1500 let w0 = pen_hist
1501 .len()
1502 .saturating_sub(sp.penalty_last_n.min(PEN_WINDOW_MAX));
1503 let hist = &pen_hist[w0..];
1504 let hd = e.htod_u32_v(hist)?;
1505 e.penalize_logits(
1506 &mut col,
1507 &hd,
1508 hist.len(),
1509 sp.penalty_repeat,
1510 sp.penalty_freq,
1511 sp.penalty_present,
1512 n_vocab,
1513 )?;
1514 }
1515 let rows0 = e.htod_i32(&[0])?;
1516 let (mut th_d, mut z_d, mut mx_d) = (e.zeros(1)?, e.zeros(1)?, e.zeros(1)?);
1517 e.filter_stats(
1518 &col, n_vocab, &rows0, &mut th_d, &mut z_d, &mut mx_d, n_vocab, 1, sp.temp, sp.top_k,
1519 sp.top_p, sp.min_p,
1520 )?;
1521 let (th, mx) = (e.dtoh(&th_d)?[0], e.dtoh(&mx_d)?[0]);
1522 let mut perturb = e.zeros(n_vocab)?;
1523 e.gumbel_perturb_filtered(&col, &mut perturb, n_vocab, sp.seed, *sctr, sp.temp, mx, th)?;
1524 *sctr = sctr.wrapping_add(1);
1525 let td = e.argmax_token_device(&perturb, n_vocab)?;
1526 let tok = guard_vocab_token(
1527 e.dtoh_u32_one(&td)?,
1528 n_vocab,
1529 &format!("sampled boundary token (site={site})"),
1530 )?;
1531 if spec_boundary_trace() {
1532 // the pre-lane token, from the SAME row, so the deviation rate is measurable.
1533 let raw = e.argmax_token_device(logits, n_vocab)?;
1534 let greedy = e.dtoh_u32_one(&raw)?;
1535 eprintln!(
1536 "[spec-boundary] site={site} sampled={tok} argmax={greedy} \
1537 deviates={} temp={} sctr={}",
1538 (tok != greedy) as u8,
1539 sp.temp,
1540 sctr.wrapping_sub(1),
1541 );
1542 }
1543 Ok(tok)
1544}
1545
1546/// Host-row twin of [`sample_boundary_token_dev`] (the prime / feed / entry rows arrive as
1547/// host `Vec<f32>`).
1548#[allow(clippy::too_many_arguments)]
1549pub fn sample_boundary_token(
1550 e: &Engine,
1551 logits: &[f32],
1552 sp: &SpecSampling,
1553 pen_hist: &[u32],
1554 sctr: &mut u32,
1555 site: &str,
1556) -> Result<u32, Box<dyn std::error::Error>> {
1557 let n_vocab = logits.len();
1558 let d = e.htod(logits)?;
1559 sample_boundary_token_dev(e, &d, n_vocab, sp, pen_hist, sctr, site)
1560}
1561
1562struct SpecPipeTraceClock {
1563 pair: usize,
1564 started: std::time::Instant,
1565}
1566
1567#[derive(Clone)]
1568struct SpecPipeTraceCtx {
1569 clock: std::sync::Arc<SpecPipeTraceClock>,
1570 round: usize,
1571 lane: usize,
1572}
1573
1574struct SpecPipeTraceMarker {
1575 trace: SpecPipeTraceCtx,
1576 phase: &'static str,
1577 edge: &'static str,
1578 slot: Option<usize>,
1579}
1580
1581unsafe extern "C" fn spec_pipe_trace_marker(raw: *mut std::ffi::c_void) {
1582 let marker = unsafe { Box::from_raw(raw.cast::<SpecPipeTraceMarker>()) };
1583 let lane = if marker.trace.lane == 0 { "A" } else { "B" };
1584 let slot = marker
1585 .slot
1586 .map(|v| v.to_string())
1587 .unwrap_or_else(|| "-".into());
1588 let t_ms = marker.trace.clock.started.elapsed().as_secs_f64() * 1e3;
1589 use std::io::Write as _;
1590 let stderr = std::io::stderr();
1591 let mut stderr = stderr.lock();
1592 let _ = writeln!(
1593 stderr,
1594 "[spec-pipe-timeline] pair={} round={} lane={lane} phase={} edge={} \
1595 slot={slot} t_ms={t_ms:.3}",
1596 marker.trace.clock.pair, marker.trace.round, marker.phase, marker.edge,
1597 );
1598}
1599
1600fn enqueue_spec_pipe_trace_marker(
1601 stream: &cudarc::driver::CudaStream,
1602 trace: Option<&SpecPipeTraceCtx>,
1603 phase: &'static str,
1604 edge: &'static str,
1605 slot: Option<usize>,
1606) -> Result<(), Box<dyn std::error::Error>> {
1607 let Some(trace) = trace else {
1608 return Ok(());
1609 };
1610 let marker = Box::new(SpecPipeTraceMarker {
1611 trace: trace.clone(),
1612 phase,
1613 edge,
1614 slot,
1615 });
1616 let raw = Box::into_raw(marker);
1617 let result = unsafe {
1618 cudarc::driver::result::stream::launch_host_function(
1619 stream.cu_stream(),
1620 spec_pipe_trace_marker,
1621 raw.cast(),
1622 )
1623 };
1624 if let Err(err) = result {
1625 unsafe {
1626 drop(Box::from_raw(raw));
1627 }
1628 return Err(err.into());
1629 }
1630 Ok(())
1631}
1632
1633#[derive(Default)]
1634struct SpecPipeProgress {
1635 setup_done: [bool; 2],
1636 draft_done: [usize; 2],
1637 stage0_done: [usize; 2],
1638 verify_done: [usize; 2],
1639 accept_done: [usize; 2],
1640 finished: [bool; 2],
1641 aborted: bool,
1642}
1643
1644/// Host-side issue coordinator for the reduced two-session speculative pipeline. Each session
1645/// keeps its existing call stack and round locals; this object only orders phase entry. The
1646/// primary mutex spans whole draft/accept/tail issue regions so Engine's single-stream scratch
1647/// cannot be interleaved by the two host threads.
1648struct SpecPipeSync {
1649 progress: std::sync::Mutex<SpecPipeProgress>,
1650 changed: std::sync::Condvar,
1651 primary: std::sync::Mutex<()>,
1652 trace: Option<std::sync::Arc<SpecPipeTraceClock>>,
1653}
1654
1655impl SpecPipeSync {
1656 fn new() -> Self {
1657 static TRACE_PAIR: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
1658 let trace = (std::env::var("MEMRA_SPEC_PIPE_TRACE").as_deref() == Ok("1")).then(|| {
1659 std::sync::Arc::new(SpecPipeTraceClock {
1660 pair: TRACE_PAIR.fetch_add(1, std::sync::atomic::Ordering::Relaxed) + 1,
1661 started: std::time::Instant::now(),
1662 })
1663 });
1664 Self {
1665 progress: std::sync::Mutex::new(SpecPipeProgress::default()),
1666 changed: std::sync::Condvar::new(),
1667 primary: std::sync::Mutex::new(()),
1668 trace,
1669 }
1670 }
1671}
1672
1673#[derive(Clone)]
1674struct SpecPipeLane {
1675 sync: std::sync::Arc<SpecPipeSync>,
1676 lane: usize,
1677}
1678
1679impl SpecPipeLane {
1680 fn peer(&self) -> usize {
1681 1 - self.lane
1682 }
1683
1684 fn aborted() -> Box<dyn std::error::Error> {
1685 "paired speculative peer aborted".into()
1686 }
1687
1688 fn trace(&self, round: usize) -> Option<SpecPipeTraceCtx> {
1689 self.sync.trace.as_ref().map(|clock| SpecPipeTraceCtx {
1690 clock: clock.clone(),
1691 round,
1692 lane: self.lane,
1693 })
1694 }
1695
1696 fn setup_begin(&self) -> Result<(), Box<dyn std::error::Error>> {
1697 let mut p = self.sync.progress.lock().unwrap();
1698 while !p.aborted && self.lane == 1 && !p.setup_done[0] && !p.finished[0] {
1699 p = self.sync.changed.wait(p).unwrap();
1700 }
1701 if p.aborted {
1702 Err(Self::aborted())
1703 } else {
1704 Ok(())
1705 }
1706 }
1707
1708 fn setup_end(&self) {
1709 let mut p = self.sync.progress.lock().unwrap();
1710 p.setup_done[self.lane] = true;
1711 self.sync.changed.notify_all();
1712 }
1713
1714 fn draft_begin(
1715 &self,
1716 round: usize,
1717 ) -> Result<std::sync::MutexGuard<'_, ()>, Box<dyn std::error::Error>> {
1718 let peer = self.peer();
1719 let mut p = self.sync.progress.lock().unwrap();
1720 loop {
1721 if p.aborted {
1722 return Err(Self::aborted());
1723 }
1724 let setup_ready =
1725 (p.setup_done[0] || p.finished[0]) && (p.setup_done[1] || p.finished[1]);
1726 let prior_ready = p.accept_done[self.lane] >= round
1727 && (p.accept_done[peer] >= round || p.finished[peer]);
1728 let turn_ready = if self.lane == 0 {
1729 true
1730 } else {
1731 p.draft_done[0] > round || p.finished[0]
1732 };
1733 if setup_ready && prior_ready && turn_ready {
1734 break;
1735 }
1736 p = self.sync.changed.wait(p).unwrap();
1737 }
1738 drop(p);
1739 Ok(self.sync.primary.lock().unwrap())
1740 }
1741
1742 fn draft_end(&self, round: usize) {
1743 let mut p = self.sync.progress.lock().unwrap();
1744 p.draft_done[self.lane] = round + 1;
1745 self.sync.changed.notify_all();
1746 }
1747
1748 /// Admit stage 0 and return whether this lane owns the interval's one reverse fence.
1749 /// Lane B releases as soon as lane A has issued its boundary TX, not after A's full body.
1750 fn stage0_begin(&self, round: usize) -> Result<bool, Box<dyn std::error::Error>> {
1751 let peer = self.peer();
1752 let mut p = self.sync.progress.lock().unwrap();
1753 loop {
1754 if p.aborted {
1755 return Err(Self::aborted());
1756 }
1757 let ready = if self.lane == 0 {
1758 p.draft_done[0] > round && (p.draft_done[1] > round || p.finished[1])
1759 } else {
1760 p.draft_done[1] > round && (p.stage0_done[0] > round || p.finished[0])
1761 };
1762 if ready {
1763 return Ok(self.lane == 0 || p.finished[peer]);
1764 }
1765 p = self.sync.changed.wait(p).unwrap();
1766 }
1767 }
1768
1769 fn stage0_end(&self, round: usize) {
1770 let mut p = self.sync.progress.lock().unwrap();
1771 p.stage0_done[self.lane] = round + 1;
1772 self.sync.changed.notify_all();
1773 }
1774
1775 /// Stage 1 is single-owner per engine. A proceeds immediately after its own ticket; B waits
1776 /// for A's full stage1/head issue so only A.S1 and B.S0 can overlap.
1777 fn stage1_begin(&self, round: usize) -> Result<(), Box<dyn std::error::Error>> {
1778 let mut p = self.sync.progress.lock().unwrap();
1779 while !p.aborted
1780 && !(p.stage0_done[self.lane] > round
1781 && (self.lane == 0 || p.verify_done[0] > round || p.finished[0]))
1782 {
1783 p = self.sync.changed.wait(p).unwrap();
1784 }
1785 if p.aborted {
1786 Err(Self::aborted())
1787 } else {
1788 Ok(())
1789 }
1790 }
1791
1792 fn verify_end(&self, round: usize) {
1793 let mut p = self.sync.progress.lock().unwrap();
1794 p.verify_done[self.lane] = round + 1;
1795 self.sync.changed.notify_all();
1796 }
1797
1798 fn accept_begin(
1799 &self,
1800 round: usize,
1801 ) -> Result<std::sync::MutexGuard<'_, ()>, Box<dyn std::error::Error>> {
1802 let mut p = self.sync.progress.lock().unwrap();
1803 loop {
1804 if p.aborted {
1805 return Err(Self::aborted());
1806 }
1807 let ready = if self.lane == 0 {
1808 p.verify_done[0] > round && (p.verify_done[1] > round || p.finished[1])
1809 } else {
1810 p.verify_done[1] > round && (p.accept_done[0] > round || p.finished[0])
1811 };
1812 if ready {
1813 break;
1814 }
1815 p = self.sync.changed.wait(p).unwrap();
1816 }
1817 drop(p);
1818 Ok(self.sync.primary.lock().unwrap())
1819 }
1820
1821 fn accept_end(&self, round: usize) {
1822 let mut p = self.sync.progress.lock().unwrap();
1823 p.accept_done[self.lane] = round + 1;
1824 self.sync.changed.notify_all();
1825 }
1826
1827 fn primary(&self) -> std::sync::MutexGuard<'_, ()> {
1828 self.sync.primary.lock().unwrap()
1829 }
1830
1831 fn finish(&self, failed: bool) {
1832 let mut p = self.sync.progress.lock().unwrap();
1833 p.finished[self.lane] = true;
1834 p.aborted |= failed;
1835 self.sync.changed.notify_all();
1836 }
1837}
1838
1839struct SpecPipeFinish<'a> {
1840 lane: &'a SpecPipeLane,
1841 closed: bool,
1842}
1843
1844impl<'a> SpecPipeFinish<'a> {
1845 fn new(lane: &'a SpecPipeLane) -> Self {
1846 Self {
1847 lane,
1848 closed: false,
1849 }
1850 }
1851
1852 fn close(&mut self, failed: bool) {
1853 self.lane.finish(failed);
1854 self.closed = true;
1855 }
1856}
1857
1858impl Drop for SpecPipeFinish<'_> {
1859 fn drop(&mut self) {
1860 if !self.closed {
1861 self.lane.finish(true);
1862 }
1863 }
1864}
1865
1866/// Scoped transfer of one exclusively-borrowed session to the second host issue thread.
1867/// `CudaGraph` is not marked Send by cudarc because its raw driver handles carry no automatic
1868/// trait. CUDA driver graph handles are context-scoped rather than OS-thread-affine; the caller
1869/// binds that context before touching the session, joins before returning, and never aliases the
1870/// pointer. Keep this exception local to the experimental pair call instead of marking the public
1871/// session type Send.
1872struct SpecPipeSessionPtr(*mut SpecSession);
1873
1874unsafe impl Send for SpecPipeSessionPtr {}
1875
1876impl SpecPipeSessionPtr {
1877 unsafe fn get_mut(&mut self) -> &mut SpecSession {
1878 unsafe { &mut *self.0 }
1879 }
1880}
1881
1882/// Per-session persistent draft-graph context: the captured CUDA graph(s) plus the device
1883/// buffers whose POINTERS the capture bakes. Reuse legality: the greedy capture bakes only
1884/// session-stable pointers (the session's own MtpScratch KV — allocated once, never realloc'd;
1885/// the model's resident embedding; the process-wide OnceLock p_min) and the g_* buffers held
1886/// HERE — so one capture serves the session's whole lifetime. The sampled capture additionally
1887/// bakes (seed, temp) as capture-time constants and needs k q-slots — keyed by `s_key`, dropped
1888/// and recaptured when a pool-resumed request changes them. `*_failed` memoizes a failed capture
1889/// so the eager fallback doesn't pay a doomed capture attempt every burst.
1890/// Capture identity of the parked SAMPLED draft graph (`DraftGraphCtx::graph_s`).
1891///
1892/// EXACTNESS, not perf (lane/graph-s-key-exactness-20260819; receipts
1893/// `research/spec-cache-20260818/GRAPH-S-KEY.md`). Two classes of field live here, both
1894/// load-bearing:
1895///
1896/// - **Baked constants.** `seed` and `temp` are capture-time constants INSIDE the graph and `k`
1897/// sizes the q slots its replays write. A resumed request changing any of them must recapture.
1898/// This is all the key used to carry.
1899/// - **Regime fields.** `top_k`/`top_p`/`min_p`/`pen_on` are not baked, but they decide whether
1900/// the captured graph is a legal draft chain AT ALL. The in-graph draw is one gumbel-max over
1901/// the RAW softmax (`gumbel_perturb_ctr`, unfiltered by construction), while the verify builds
1902/// the accept test's `q` from `filter_stats(q_slots, top_k, top_p, min_p)`. If those disagree
1903/// the accept test evaluates a distribution the draft was never sampled from: a draft token
1904/// below the filter threshold gathers `q = 0` (`softmax_gather_filtered_f32`,
1905/// `cu/spec_sample.cu`) and `u * 0 < p` accepts it UNCONDITIONALLY.
1906///
1907/// Omitting the regime fields was reachable — not through the prefix-cache spec restore (that
1908/// path is greedy-only, `memra-server` `spec_restore_convertible`), but through WHOLE-SESSION
1909/// spec reuse: a parked `SpecSession` carries this `DraftGraphCtx`, and the pool-resume probe
1910/// applies no sampler predicate at all. Turn 1 pure-temp parks a graph; turn 2 of the same
1911/// conversation, same explicit seed and temperature, adds `top_p`/`top_k` and inherits it.
1912#[derive(Clone, Copy, PartialEq, Eq, Debug)]
1913pub(crate) struct SampledGraphKey {
1914 seed: u64,
1915 temp_bits: u32,
1916 k: usize,
1917 top_k: i32,
1918 top_p_bits: u32,
1919 min_p_bits: u32,
1920 pen_on: bool,
1921}
1922
1923impl SampledGraphKey {
1924 pub(crate) fn new(
1925 seed: u64,
1926 temp: f32,
1927 k: usize,
1928 top_k: i32,
1929 top_p: f32,
1930 min_p: f32,
1931 pen_on: bool,
1932 ) -> Self {
1933 SampledGraphKey {
1934 seed,
1935 temp_bits: temp.to_bits(),
1936 k,
1937 top_k,
1938 top_p_bits: top_p.to_bits(),
1939 min_p_bits: min_p.to_bits(),
1940 pen_on,
1941 }
1942 }
1943
1944 /// The one regime the in-graph sampled chain may stand in for the eager one: nothing but
1945 /// temperature shapes `q`. Computed FROM THE KEY so the capture guard, the launch guard and
1946 /// the key can never drift apart (they were three separate expressions before this lane, and
1947 /// the launch site simply forgot to ask).
1948 pub(crate) fn pure_temp(&self) -> bool {
1949 self.top_k == 0
1950 && f32::from_bits(self.top_p_bits) >= 1.0
1951 && f32::from_bits(self.min_p_bits) <= 0.0
1952 && !self.pen_on
1953 }
1954}
1955
1956pub(crate) struct DraftGraphCtx {
1957 g_tok: CudaSlice<u32>,
1958 g_pos: CudaSlice<i32>,
1959 g_seed: CudaSlice<f32>,
1960 g_p: CudaSlice<f32>,
1961 g_ctr: CudaSlice<u32>,
1962 g_q: CudaSlice<f32>,
1963 g_perturb: CudaSlice<f32>,
1964 q_slots: Vec<CudaSlice<f32>>,
1965 /// DRAFT-SIDE GRAMMAR MASK (lane/draft-mask): packed allowed-set words over the DRAFT
1966 /// head's vocab, at a STABLE address so the captured draft graph's mask node reads the
1967 /// per-position contents the host re-uploads before each replay (the graph-promote
1968 /// pattern from decode.rs). Empty unless the session drafts under a grammar.
1969 g_dmask: CudaSlice<u32>,
1970 /// was `graph` captured WITH the mask node? A parked graph of the wrong shape is dropped.
1971 graph_masked: bool,
1972 graph: Option<cudarc::driver::CudaGraph>,
1973 graph_s: Option<cudarc::driver::CudaGraph>,
1974 /// Failed-capture memoization for both graphs — LOUD on flip, cleared on pool resume
1975 /// (audit Q2, the TRT #16072 silent-permanent-coverage-loss class).
1976 failed: DraftGraphFallback,
1977 /// Capture identity of `graph_s` — see [`SampledGraphKey`]. `None` iff no sampled graph is
1978 /// parked; a request whose key differs drops the parked graph (and its q slots/keeper).
1979 s_key: Option<SampledGraphKey>,
1980 /// CAPTURE-RETAIN keepers (#68 root cause, 2026-08-04): the warmup-run transients whose
1981 /// pool addresses the captured graph(s) bake. Without these, the transients return to the
1982 /// pool at capture-body exit and later work (burst-boundary prime/fill/commit passes, or a
1983 /// co-served session in the worker) reuses those addresses — the persisted graph's replay
1984 /// then reads/writes live unrelated buffers (exactness corruption, first seen as the ST
1985 /// serve-spec 4B graph-arm corruption; one-shot CLI calls never re-shuffled the pool, which
1986 /// is why run-spec K=1..8 passed on the same checkpoint). Same fix class as
1987 /// capture_graph_retained's gemma/decode.rs sites — hold as long as the graph replays.
1988 keeper: Vec<Box<dyn std::any::Any + Send>>,
1989 keeper_s: Vec<Box<dyn std::any::Any + Send>>,
1990}
1991
1992/// Failed-capture memoization for the two draft graphs (audit Q2, 2026-08-05 — the
1993/// TRT #16072 trap class: pressure-triggered, silent, long-lived coverage loss).
1994///
1995/// Three contracts:
1996/// - LOUD FLIP: `mark_*` returns the warn line exactly on the false→true transition
1997/// (returned, not printed, so the once-per-flip contract is unit-testable); the caller
1998/// `eprintln!`s it UNCONDITIONALLY — a dropped draft graph is never silent. Re-marking
1999/// an already-failed graph returns None (the per-burst memoization that keeps the eager
2000/// fallback from paying a doomed capture attempt every burst).
2001/// - RESET ON RESUME: `reset_on_resume` clears both flags — a parked session resumed by a
2002/// NEW request gets one fresh capture chance instead of carrying a transient-pressure
2003/// failure for the pool's whole lifetime. Returns the note line only when a flag was
2004/// actually set (quiet on the common clean-resume path).
2005/// - Shape-change clears (`clear_*`) stay silent, exactly as before: they precede a fresh
2006/// capture attempt whose own failure would re-flip loudly.
2007#[derive(Default)]
2008pub(crate) struct DraftGraphFallback {
2009 greedy: bool,
2010 sampled: bool,
2011}
2012impl DraftGraphFallback {
2013 fn mark_greedy(&mut self, reason: &str) -> Option<String> {
2014 if self.greedy {
2015 return None;
2016 }
2017 self.greedy = true;
2018 Some(format!(
2019 "[spec] WARN: draft-graph capture failed ({reason}); eager fallback until session resume"
2020 ))
2021 }
2022 fn mark_sampled(&mut self, reason: &str) -> Option<String> {
2023 if self.sampled {
2024 return None;
2025 }
2026 self.sampled = true;
2027 Some(format!(
2028 "[spec] WARN: sampled draft-graph capture failed ({reason}); eager fallback until session resume"
2029 ))
2030 }
2031 fn greedy_failed(&self) -> bool {
2032 self.greedy
2033 }
2034 fn sampled_failed(&self) -> bool {
2035 self.sampled
2036 }
2037 fn clear_greedy(&mut self) {
2038 self.greedy = false;
2039 }
2040 fn clear_sampled(&mut self) {
2041 self.sampled = false;
2042 }
2043 /// Pool-resume reset: both graphs get a fresh capture chance. Some(note) iff any flag
2044 /// was set (so clean resumes stay quiet).
2045 pub(crate) fn reset_on_resume(&mut self) -> Option<String> {
2046 if !self.greedy && !self.sampled {
2047 return None;
2048 }
2049 let which = match (self.greedy, self.sampled) {
2050 (true, true) => "greedy+sampled",
2051 (true, false) => "greedy",
2052 _ => "sampled",
2053 };
2054 self.greedy = false;
2055 self.sampled = false;
2056 Some(format!(
2057 "[spec] draft-graph fallback reset on session resume ({which}); recapture eligible"
2058 ))
2059 }
2060}
2061
2062impl DraftGraphCtx {
2063 fn new(e: &Engine, n_embd: usize, qlen: usize) -> Result<Self, Box<dyn std::error::Error>> {
2064 Ok(DraftGraphCtx {
2065 g_tok: e.alloc_u32_zeroed(1)?,
2066 g_pos: e.htod_i32(&[0])?,
2067 g_seed: e.zeros(n_embd)?,
2068 g_p: e.zeros(1)?,
2069 g_ctr: e.alloc_u32_zeroed(1)?,
2070 g_q: e.zeros(qlen)?,
2071 g_perturb: e.zeros(qlen)?,
2072 q_slots: Vec::new(),
2073 g_dmask: e.alloc_u32_zeroed(1)?,
2074 graph_masked: false,
2075 graph: None,
2076 graph_s: None,
2077 failed: DraftGraphFallback::default(),
2078 s_key: None,
2079 keeper: Vec::new(),
2080 keeper_s: Vec::new(),
2081 })
2082 }
2083}
2084
2085pub(crate) struct MtpScratch {
2086 kv: KvLayer,
2087 /// Logical row capacity. On the graph/DC draft path it also doubles as fa_decode_dc's
2088 /// bucket_max: n_splits is sized from it ONCE, so the graph captured at round 0 stays valid
2089 /// for every later t_kv. Step35 refuses that path and may back this logical extent with the
2090 /// smaller host-indexed SWA ring instead.
2091 cap: usize,
2092 extra: Vec<MtpScratchPlane>,
2093}
2094
2095struct MtpScratchPlane {
2096 kv: KvLayer,
2097 cap: usize,
2098}
2099
2100fn mtp_scratch_layout(
2101 cfg: &memra_gguf::config::ModelConfig,
2102 geom: Option<&crate::hybrid::DraftGeom>,
2103) -> (usize, usize, usize, usize) {
2104 // Student draft heads carry fewer KV heads (head_dim unchanged) -> smaller scratch rows.
2105 let n_head_kv = geom.map(|g| g.n_head_kv).unwrap_or(cfg.n_head_kv as usize);
2106 let head_dim_k = cfg.head_dim_k as usize;
2107 let head_dim_v = cfg.head_dim_v as usize;
2108 assert!(
2109 head_dim_k % 32 == 0 && head_dim_v % 32 == 0,
2110 "KVQUANT requires head_dim%32==0 (MTP scratch)"
2111 );
2112 let kv_dim_k = head_dim_k * n_head_kv;
2113 let kv_dim_v = head_dim_v * n_head_kv;
2114 // The fp8-KV arm deliberately does not reach the draft scratch; keep the exact format
2115 // policy shared with `MtpScratch::new` so admission scales the same allocation.
2116 let (kbb, vbb) = crate::kv_blk_bytes();
2117 let k_tok_bytes = (kv_dim_k / 32) * kbb;
2118 let v_tok_bytes = (kv_dim_v / 32) * vbb;
2119 (kv_dim_k, kv_dim_v, k_tok_bytes, v_tok_bytes)
2120}
2121
2122fn mtp_chain_head_index(step: usize, head_count: usize) -> usize {
2123 assert!(head_count > 0, "MTP chain requires at least one head");
2124 step % head_count
2125}
2126
2127impl MtpScratch {
2128 fn alloc_plane(
2129 e: &Engine,
2130 cfg: &memra_gguf::config::ModelConfig,
2131 plan: &memra_gguf::model_plan::ModelPlan,
2132 cap: usize,
2133 geom: Option<&crate::hybrid::DraftGeom>,
2134 ) -> Result<MtpScratchPlane, Box<dyn std::error::Error>> {
2135 let (kv_dim_k, kv_dim_v, k_tok_bytes, v_tok_bytes) = mtp_scratch_layout(cfg, geom);
2136 let ring = if crate::cache::swa_ring_on()
2137 && crate::plan_backend::decode_batch_program(plan)
2138 == crate::plan_backend::DecodeBatchProgram::SlidingGatedMoe
2139 {
2140 let window = plan
2141 .layers
2142 .iter()
2143 .find_map(|layer| match layer.attention {
2144 memra_gguf::model_plan::AttentionPlan::SlidingWindow { window, .. } => {
2145 Some(window as usize)
2146 }
2147 _ => None,
2148 })
2149 .ok_or("sliding-gated-MoE draft scratch has no sliding-window layer")?;
2150 Some(crate::cache::KvRing::new(
2151 crate::cache::swa_ring_rows(window, cap),
2152 window,
2153 ))
2154 } else {
2155 None
2156 };
2157 let alloc_rows = ring.as_ref().map(crate::cache::KvRing::rows).unwrap_or(cap);
2158 // Ring-backed planes arm the device base mirror for the dcw draft arm (see
2159 // KvLayer::base_d): the captured chain derives its physical rows from
2160 // (len_d, base_d, window) with zero per-token node updates.
2161 let base_d = match ring.as_ref() {
2162 Some(_) => Some(e.htod_i32(&[0])?),
2163 None => None,
2164 };
2165 Ok(MtpScratchPlane {
2166 kv: KvLayer {
2167 k: e.alloc_u8(alloc_rows * k_tok_bytes)?,
2168 v: e.alloc_u8(alloc_rows * v_tok_bytes)?,
2169 kv_dim_k,
2170 kv_dim_v,
2171 k_tok_bytes,
2172 v_tok_bytes,
2173 len: 0,
2174 ring,
2175 len_d: e.htod_i32(&[0])?,
2176 base_d,
2177 },
2178 cap,
2179 })
2180 }
2181
2182 fn new(
2183 e: &Engine,
2184 cfg: &memra_gguf::config::ModelConfig,
2185 plan: &memra_gguf::model_plan::ModelPlan,
2186 cap: usize,
2187 geom: Option<&crate::hybrid::DraftGeom>,
2188 ) -> Result<Self, Box<dyn std::error::Error>> {
2189 // env-selected KV formats (default 34/24). The fp8-KV arm (MEMRA_KV_FP8) deliberately
2190 // does NOT reach the draft scratch: fp8 drafts drifted acceptance 69-88% -> 46%
2191 // (2026-07-12 A/B); the scratch is tiny, so it keeps baseline q8_0/q5_1 numerics
2192 // while the TRUNK cache carries the fp8 depth win. Scratch append/fa pass g=false.
2193 let primary = Self::alloc_plane(e, cfg, plan, cap, geom)?;
2194 Ok(MtpScratch {
2195 kv: primary.kv,
2196 cap: primary.cap,
2197 extra: Vec::new(),
2198 })
2199 }
2200
2201 fn push_plane(
2202 &mut self,
2203 e: &Engine,
2204 cfg: &memra_gguf::config::ModelConfig,
2205 plan: &memra_gguf::model_plan::ModelPlan,
2206 geom: Option<&crate::hybrid::DraftGeom>,
2207 ) -> Result<(), Box<dyn std::error::Error>> {
2208 self.extra
2209 .push(Self::alloc_plane(e, cfg, plan, self.cap, geom)?);
2210 Ok(())
2211 }
2212
2213 fn plane_count(&self) -> usize {
2214 1 + self.extra.len()
2215 }
2216
2217 fn plane(&self, index: usize) -> (&KvLayer, usize) {
2218 if index == 0 {
2219 (&self.kv, self.cap)
2220 } else {
2221 let plane = &self.extra[index - 1];
2222 (&plane.kv, plane.cap)
2223 }
2224 }
2225
2226 fn plane_mut(&mut self, index: usize) -> (&mut KvLayer, usize) {
2227 if index == 0 {
2228 (&mut self.kv, self.cap)
2229 } else {
2230 let plane = &mut self.extra[index - 1];
2231 (&mut plane.kv, plane.cap)
2232 }
2233 }
2234
2235 // #[track_caller]: set_len/set_plane_len have eight call sites (checkpoint restore, spec
2236 // rollback, session grow, seed replay ...) and the lap failure needs to say WHICH one, not
2237 // just that a rewind was refused.
2238 #[track_caller]
2239 fn set_plane_len(
2240 &mut self,
2241 e: &Engine,
2242 index: usize,
2243 n: usize,
2244 ) -> Result<(), Box<dyn std::error::Error>> {
2245 let caller = std::panic::Location::caller();
2246 let (kv, cap) = self.plane_mut(index);
2247 if let Some(ring) = kv.ring.as_ref() {
2248 if !ring.can_rewind_to(n) {
2249 // NAME THE NUMBERS (2026-08-28). This error is a step37 serving blocker on the
2250 // vendor-default shape and it fires from more than one call path with more than
2251 // one trigger: a long generation walks the checkpoint out of the ring, but a
2252 // ~4.5k-token prompt also fails within 5 s of prime, which accumulation cannot
2253 // explain. A bare message forced two rounds of guessing; the operands make each
2254 // trigger name itself.
2255 let raw = n.saturating_sub(ring.window().saturating_sub(1));
2256 return Err(format!(
2257 "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})",
2258 ring.window(),
2259 ring.base(),
2260 ring.rows(),
2261 raw & !31usize,
2262 )
2263 .into());
2264 }
2265 }
2266 kv.len = n;
2267 e.set_i32_one(&mut kv.len_d, n as i32)
2268 }
2269
2270 /// Set BOTH length counters: the host mirror AND the device len_d the captured append/fa read
2271 /// (a 4-byte in-place htod — the counter pointer is baked into the graph, never realloc'd).
2272 /// This is the ONLY truncation/rollback mechanism the persistent draft KV needs.
2273 #[track_caller]
2274 fn set_len(&mut self, e: &Engine, n: usize) -> Result<(), Box<dyn std::error::Error>> {
2275 let caller = std::panic::Location::caller();
2276 if !self.can_rewind_to(n) {
2277 // set_plane_len re-checks and reports the operands; call it so the failure carries
2278 // which plane refused and why, instead of this bare aggregate.
2279 for index in 0..self.plane_count() {
2280 self.set_plane_len(e, index, n)?;
2281 }
2282 return Err(format!(
2283 "SWA ring MTP checkpoint has been lapped; full re-prime required (aggregate rewind_to={n}, no single plane reported, called from {caller})"
2284 )
2285 .into());
2286 }
2287 for index in 0..self.plane_count() {
2288 self.set_plane_len(e, index, n)?;
2289 }
2290 Ok(())
2291 }
2292
2293 fn can_rewind_to(&self, n: usize) -> bool {
2294 (0..self.plane_count()).all(|index| {
2295 self.plane(index)
2296 .0
2297 .ring
2298 .as_ref()
2299 .is_none_or(|ring| ring.can_rewind_to(n))
2300 })
2301 }
2302
2303 /// Pre-arm ring headroom for `rows` upcoming DEVICE-COUNTER appends (the dcw draft arm):
2304 /// a captured chain cannot rebase mid-replay, so any rebase the coming appends could need
2305 /// happens HERE, host-side, before the capture warmups or the round's replays (the rebase
2306 /// arm of `prepare_kv_append` also refreshes the plane's `base_d` device mirror). No-op on
2307 /// flat planes and when the ring already has room; `len` is untouched either way.
2308 fn ensure_dcw_headroom(
2309 &mut self,
2310 e: &Engine,
2311 rows: usize,
2312 ) -> Result<(), Box<dyn std::error::Error>> {
2313 for index in 0..self.plane_count() {
2314 let (kv, _) = self.plane_mut(index);
2315 let Some(ring) = kv.ring.as_ref() else {
2316 continue;
2317 };
2318 let retain = memra_kv::swa_retain_from(kv.len, ring.window(), ring.base());
2319 e.prepare_kv_append(kv, retain, rows)?;
2320 }
2321 Ok(())
2322 }
2323}
2324
2325/// Retained verify intermediates for the REPLAY-FREE partial accept (2026-07-03, the profiled
2326/// #1 spec cost at long ctx: the partial-accept replay was a DUPLICATE trunk pass — ~0.54 extra
2327/// full weight reads per round — recomputing columns the verify had already produced
2328/// bit-identically). Holds, per linear layer, everything needed to rebuild its recurrent state
2329/// to "after the first j verify columns" WITHOUT re-running the trunk:
2330/// - BATCHED-path layers (`gdn`): the exact token-major inputs the round's ONE gdn_scan
2331/// consumed. A prefix re-run of the SAME kernel (t=j) from the snapshot state is bit-identical
2332/// to the first j iterations of the verify's scan — the kernel's t-loop carries state in
2333/// registers and iteration t never depends on T. `qkv_mixed` (the conv input) feeds the
2334/// pure-copy ring rebuild.
2335/// - PER-COLUMN-path layers (`cols`): dtod clones of (conv_state, ssm_state) taken after each
2336/// column 0..t-2 — pure copies of the actual chain states (the last column is never a rebuild
2337/// target: j <= t-1).
2338/// Full-attn layers need nothing: their verify KV rows are bit-identical to eager's (the
2339/// decode-exact contract; verify-probe pins it), so rollback = len truncation.
2340struct GdnStash {
2341 qkv_mixed: CudaSlice<f32>, // [t, conv_dim] token-major (conv input)
2342 q_l2: CudaSlice<f32>,
2343 k_l2: CudaSlice<f32>,
2344 v_g: CudaSlice<f32>, // [t, num_v, d_state]
2345 g_log: CudaSlice<f32>,
2346 beta: CudaSlice<f32>, // [t, num_v]
2347}
2348pub(crate) struct VerifyCkpt {
2349 gdn: Vec<Option<GdnStash>>, // [n_layer], Some iff batched linear path ran
2350 cols: Vec<Option<Vec<(CudaSlice<f32>, CudaSlice<f32>)>>>, // [n_layer][col] = (conv, ssm) after col
2351}
2352/// Opaque handle for the dspark round (dflash.rs) — VerifyCkpt stays spec-private.
2353pub(crate) struct DsparkVerifyCkpt(VerifyCkpt);
2354
2355/// Engine-bundle slice 3 (DSF-ROUNDCOST-20260820 §2 row 4 / §5 rank 1): bucketed CUDA
2356/// graphs for the dspark verify's LINEAR-layer segments. The measured verify is ~2,800
2357/// eager launches whose residual cost is DEVICE-side per-launch overhead (slice 2 proved
2358/// host dispatch is not the binder: fully-deferred dispatch bought ~0 wall). The 48 GDN
2359/// layers between full-attention layers are shape-static given vt — no positions, no
2360/// t_kv, state addressed through pointer tables — so runs of them capture per
2361/// (segment, vt) and replay as ONE graph launch each. Full-attention layers stay eager
2362/// (their per-row append/fa arm picks are t_kv-driven — the exec-update extension).
2363///
2364/// Per round out-of-graph: one pointer-table refresh (gdn ping-pong moves the canonical
2365/// handles), one input-staging copy per segment, host parity bookkeeping. Captured via
2366/// `capture_graph_retained` (2 warmups + capture, keeper retains warmup transients so
2367/// pool addresses stay stable); the warmups EXECUTE, so segment conv/ssm state is saved
2368/// before and restored after — the graph's first real launch starts from the exact
2369/// pre-round state. The ckpt column stash rides persistent slabs (written inside the
2370/// graph as memcpy nodes); commit reads them via `dspark_commit_prefix_slab`.
2371/// `MEMRA_DSPARK_VERIFY_GRAPH=0` reverts to the eager walk (byte-identical body).
2372pub(crate) struct DsparkVerifyGraphs {
2373 /// Linear-attention layer indices ascending; `lin_pos[il]` = index into the vecs.
2374 lin: Vec<usize>,
2375 lin_pos: std::collections::HashMap<usize, usize>,
2376 /// [n_lin x 6] pointer table (conv, s0, s1, conv, s1, s0 per layer), refreshed per
2377 /// verify from the live handles; layer il's slice starts at lin_pos[il]*6.
2378 table_all: CudaSlice<u64>,
2379 host_table: Vec<u64>,
2380 /// Persistent per-layer ckpt stash slabs: row r of the verify at slab offset
2381 /// r*words. Shared by every (segment, vt) bucket — one verify runs at a time.
2382 stash_conv: Vec<CudaSlice<f32>>,
2383 stash_ssm: Vec<CudaSlice<f32>>,
2384 conv_words: usize,
2385 ssm_words: usize,
2386 /// Per-vt input/output staging (stable addresses the graphs bake).
2387 stage: std::collections::HashMap<usize, (CudaSlice<f32>, CudaSlice<f32>)>,
2388 /// Per-vt dflash tap-sink buffers — the captured segments bake the tap dst address,
2389 /// so the sink buffer must live (and persist) with the graphs, not with the round.
2390 pub(crate) tap_bufs: std::collections::HashMap<usize, CudaSlice<f32>>,
2391 graphs: std::collections::HashMap<(usize, usize), DsparkSegGraph>,
2392 /// Warmup-corruption guard scratch: pre-capture conv/ssm of every linear layer
2393 /// (sized n_lin — the slice-4c full-verify warmups execute the whole walk).
2394 save_conv: CudaSlice<f32>,
2395 save_ssm: CudaSlice<f32>,
2396 max_run: usize,
2397 n_embd: usize,
2398 /// Set by the verify walk: this round's linear ckpt lives in the slabs (the caller
2399 /// commits through `dspark_commit_prefix_slab` instead of the cols arm).
2400 pub(crate) round_slab: bool,
2401 // ---- slice 4c: full-verify single graph per (vt, rung) ----
2402 /// Full-attention layer indices ascending; `fa_pos[il]` = index into the vec.
2403 fa: Vec<usize>,
2404 fa_pos: std::collections::HashMap<usize, usize>,
2405 /// [n_fa x 2 x t_cap] interleaved (k,v) base-pointer pairs, refreshed per verify;
2406 /// layer il's slice starts at `fa_pos[il] * 2 * t_cap` (the seqs twins read pairs
2407 /// [2z], z < t <= t_cap, so one t_cap-sized table serves every vt).
2408 fa_table: CudaSlice<u64>,
2409 fa_host_table: Vec<u64>,
2410 t_cap: usize,
2411 /// Per-vt position staging for the captured bodies — contents refreshed per round
2412 /// (rope reads row r; the seqs twins derive append slot and T_kv per z from it).
2413 pos_stage: std::collections::HashMap<usize, CudaSlice<i32>>,
2414 /// Full-verify graphs keyed (vt, rung_end, hi).
2415 full: std::collections::HashMap<(usize, usize, usize), DsparkSegGraph>,
2416 /// Largest n with every layer in [0, n) linear or full-attention (walk coverage).
2417 covered: usize,
2418 /// Every layer in [0, n) is linear or full-attention (no MLA/unknown mixers) — the
2419 /// full-verify capture walks all of them.
2420 walk_uniform: bool,
2421 /// Last `(captures, device graph-mem reserved bytes)` reading taken by
2422 /// `HybridModel::dspark_vg_admission_debt` — the two-point base of the MARGINAL debt
2423 /// projection (see `dspark_vg_debt_projection`; a mean-based reading extrapolated the
2424 /// pool's one-time shared allocation and reserved 8.5 GB of phantom VRAM).
2425 debt_obs: Option<(usize, usize)>,
2426}
2427
2428struct DsparkSegGraph {
2429 graph: cudarc::driver::CudaGraph,
2430 _keeper: Vec<Box<dyn std::any::Any + Send>>,
2431}
2432
2433/// Per-call arguments of [`HybridModel::qwen35_tparallel_fa_layer`] — one struct so the
2434/// eager walk and the slice-4c captured full-verify graphs hand the SAME body its two
2435/// modes without a second copy of the math.
2436pub(crate) struct FaLayerArgs<'a> {
2437 /// [T] per-row positions (device): rope reads them row-indexed; the seqs twins read
2438 /// them per-z (append slot = pos, T_kv = pos + 1).
2439 pub pos_d: &'a CudaSlice<i32>,
2440 /// Verify-level lazy per-row 1-element position buffers — only the per-row fallback
2441 /// arm builds/uses them (graph mode refuses that arm).
2442 pub pos_rows: &'a mut Option<Vec<CudaSlice<i32>>>,
2443 pub pos0: usize,
2444 pub seqs_append: bool,
2445 pub batch_fa_on: bool,
2446 /// Some((kv pointer table, offset-in-u64s, rung_end)) = captured-graph mode.
2447 pub graph_cap: Option<(&'a CudaSlice<u64>, usize, usize)>,
2448 /// ROUND-STREAM (lane/draftcost-moe, v0.100 train merge): Some((token stream, device
2449 /// round counter)) routes the FA attend through the dc rows kernels and the Linear
2450 /// mixer through `linear_attn_verify_t` (the stream arms the old inline body carried).
2451 /// Never armed together with `graph_cap` (the verify-level merge guard refuses).
2452 pub stream: Option<(&'a CudaSlice<u32>, &'a CudaSlice<i32>)>,
2453 /// VerifyCkpt for the stream-Linear arm's GdnStash install; None in graph mode and
2454 /// for FA layers that never touch it.
2455 pub ckpt: Option<&'a mut VerifyCkpt>,
2456}
2457
2458// SAFETY: `CudaGraph` is not marked Send by cudarc because its raw driver handles carry
2459// no automatic trait; CUDA driver graph handles are context-scoped rather than
2460// OS-thread-affine (the SpecPipeSessionPtr precedent above). The ctx lives in
2461// `HybridModel::dspark_vgraphs` behind a Mutex and every touch happens on the engine's
2462// single decode-stream thread.
2463unsafe impl Send for DsparkVerifyGraphs {}
2464
2465impl DsparkVerifyGraphs {
2466 /// Live capture count (segment + full graphs) — the denominator of
2467 /// [`dspark_vg_debt_projection`]'s observed bytes/capture mean.
2468 pub(crate) fn captures(&self) -> usize {
2469 self.graphs.len() + self.full.len()
2470 }
2471
2472 /// Take the marginal-growth debt reading and record this observation for the next one.
2473 /// Called under the pool mutex by `HybridModel::dspark_vg_admission_debt`.
2474 pub(crate) fn admission_debt(&mut self, reserved_bytes: usize) -> usize {
2475 let captures = self.captures();
2476 let debt =
2477 dspark_vg_debt_projection(captures, dspark_vg_cap(), reserved_bytes, self.debt_obs);
2478 if captures > 0 {
2479 match self.debt_obs {
2480 Some((c0, _)) if captures <= c0 => {}
2481 _ => self.debt_obs = Some((captures, reserved_bytes)),
2482 }
2483 }
2484 debt
2485 }
2486
2487 /// Build for this cache's shape. None when there are no linear layers, sizes are
2488 /// non-uniform, or the trunk keeps a gemma4 config (never on the qwen35 family).
2489 pub(crate) fn new(
2490 e: &Engine,
2491 cache: &Cache,
2492 t_max: usize,
2493 n_embd: usize,
2494 ) -> Result<Option<Self>, Box<dyn std::error::Error>> {
2495 let lin: Vec<usize> = (0..cache.recur.len())
2496 .filter(|&il| cache.recur[il].is_some())
2497 .collect();
2498 if lin.is_empty() || t_max < 2 {
2499 return Ok(None);
2500 }
2501 let first = cache.recur[lin[0]].as_ref().unwrap();
2502 let (conv_words, ssm_words) = (first.conv_state.len(), first.ssm_state.len());
2503 for &il in &lin {
2504 let rl = cache.recur[il].as_ref().unwrap();
2505 if rl.conv_state.len() != conv_words || rl.ssm_state.len() != ssm_words {
2506 return Ok(None);
2507 }
2508 }
2509 let n = lin.len();
2510 let mut lin_pos = std::collections::HashMap::with_capacity(n);
2511 for (k, &il) in lin.iter().enumerate() {
2512 lin_pos.insert(il, k);
2513 }
2514 // longest run of consecutive linear layers (save-scratch sizing)
2515 let mut max_run = 1usize;
2516 let mut run = 1usize;
2517 for w in lin.windows(2) {
2518 if w[1] == w[0] + 1 {
2519 run += 1;
2520 max_run = max_run.max(run);
2521 } else {
2522 run = 1;
2523 }
2524 }
2525 let rows = t_max - 1;
2526 let mut stash_conv = Vec::with_capacity(n);
2527 let mut stash_ssm = Vec::with_capacity(n);
2528 for _ in 0..n {
2529 stash_conv.push(e.uninit(rows * conv_words)?);
2530 stash_ssm.push(e.uninit(rows * ssm_words)?);
2531 }
2532 let host_table = vec![0u64; n * 6];
2533 let table_all = e.htod_u64(&host_table)?;
2534 // slice 4c: full-attention census for the full-verify graphs.
2535 let fa: Vec<usize> = (0..cache.kv.len())
2536 .filter(|&il| cache.kv[il].is_some())
2537 .collect();
2538 let mut fa_pos = std::collections::HashMap::with_capacity(fa.len());
2539 for (k, &il) in fa.iter().enumerate() {
2540 fa_pos.insert(il, k);
2541 }
2542 let n_layers = cache.kv.len().max(cache.recur.len());
2543 // exactly one of (linear state, kv cache) per layer — no MLA/unknown mixers.
2544 let walk_uniform = (0..n_layers).all(|il| {
2545 cache.recur.get(il).is_some_and(|r| r.is_some())
2546 != cache.kv.get(il).is_some_and(|k| k.is_some())
2547 });
2548 // Contiguous covered prefix: the largest n such that every layer in [0, n) is
2549 // linear or full-attention. The TRUNK walk is [0, layers.len()) and the cache
2550 // vecs can carry EXTRA state slots past it (the q38 export keeps the MTP head
2551 // layer's kv at the tail — hi == lin+fa never held, the s4c battery's zero
2552 // 'full' captures). The full-graph guard is walk coverage, not slot arithmetic.
2553 let covered = (0..n_layers)
2554 .take_while(|il| lin_pos.contains_key(il) || fa_pos.contains_key(il))
2555 .count();
2556 let t_cap = t_max;
2557 let fa_host_table = vec![0u64; fa.len() * 2 * t_cap];
2558 let fa_table = e.htod_u64(&fa_host_table)?;
2559 Ok(Some(Self {
2560 lin,
2561 lin_pos,
2562 table_all,
2563 host_table,
2564 stash_conv,
2565 stash_ssm,
2566 conv_words,
2567 ssm_words,
2568 stage: std::collections::HashMap::new(),
2569 tap_bufs: std::collections::HashMap::new(),
2570 graphs: std::collections::HashMap::new(),
2571 save_conv: e.uninit(n * conv_words)?,
2572 save_ssm: e.uninit(n * ssm_words)?,
2573 max_run,
2574 n_embd,
2575 round_slab: false,
2576 fa,
2577 fa_pos,
2578 fa_table,
2579 fa_host_table,
2580 t_cap,
2581 pos_stage: std::collections::HashMap::new(),
2582 full: std::collections::HashMap::new(),
2583 covered,
2584 walk_uniform,
2585 debt_obs: None,
2586 }))
2587 }
2588
2589 /// Rebuild the pointer tables from the live handles (once per verify — the gdn
2590 /// ping-pong swaps the canonical/alt handles between rounds; a fresh generation's
2591 /// cache buffers land at new addresses; a stale table would read the wrong state).
2592 pub(crate) fn refresh_tables(
2593 &mut self,
2594 e: &Engine,
2595 cache: &Cache,
2596 ) -> Result<(), Box<dyn std::error::Error>> {
2597 use cudarc::driver::DevicePtr;
2598 {
2599 let s = &e.gpu.stream();
2600 for (k, &il) in self.lin.iter().enumerate() {
2601 let rl = cache.recur[il].as_ref().unwrap();
2602 let (pc, _g0) = rl.conv_state.device_ptr(s);
2603 let (p0, _g1) = rl.ssm_state.device_ptr(s);
2604 let (p1, _g2) = rl.ssm_state_alt.device_ptr(s);
2605 let o = k * 6;
2606 self.host_table[o] = pc as u64;
2607 self.host_table[o + 1] = p0 as u64;
2608 self.host_table[o + 2] = p1 as u64;
2609 self.host_table[o + 3] = pc as u64;
2610 self.host_table[o + 4] = p1 as u64;
2611 self.host_table[o + 5] = p0 as u64;
2612 }
2613 for (k, &il) in self.fa.iter().enumerate() {
2614 let kvl = cache.kv[il].as_ref().unwrap();
2615 let (pk, _g0) = kvl.k.device_ptr(s);
2616 let (pv, _g1) = kvl.v.device_ptr(s);
2617 let o = k * 2 * self.t_cap;
2618 for z in 0..self.t_cap {
2619 self.fa_host_table[o + 2 * z] = pk as u64;
2620 self.fa_host_table[o + 2 * z + 1] = pv as u64;
2621 }
2622 }
2623 }
2624 e.htod_u64_into(&self.host_table, &mut self.table_all)?;
2625 if !self.fa_host_table.is_empty() {
2626 e.htod_u64_into(&self.fa_host_table, &mut self.fa_table)?;
2627 }
2628 Ok(())
2629 }
2630
2631 /// Slice 4c eligibility: Some(rung_end) when this round can replay (or capture) a
2632 /// full-verify graph — the whole walk [lo, hi) is covered, every layer is linear or
2633 /// full-attention, and ALL of the round's per-row t_kv values take the v4-seqs arm
2634 /// on ONE `fa_split_keys` ladder step that the rung also sits on (the straddle law;
2635 /// both gates are t_kv intervals, so ends-inside means all-inside). The rung is the
2636 /// round's next power of two — grid/partial sizing only (`n_splits_max` is pure
2637 /// stride; splits >= ns_eff write the empty partial the combine never reads), so one
2638 /// captured graph is bit-identical for every round the rung covers.
2639 #[allow(clippy::too_many_arguments)]
2640 pub(crate) fn full_rung(
2641 &self,
2642 model: &crate::hybrid::HybridModel,
2643 cache: &Cache,
2644 lo: usize,
2645 hi: usize,
2646 t: usize,
2647 seqs_arms_on: bool,
2648 ) -> Option<usize> {
2649 if std::env::var("MEMRA_DSPARK_FULLG_DEBUG").as_deref() == Ok("1") {
2650 static ONCE: std::sync::Once = std::sync::Once::new();
2651 let len0 = self
2652 .fa
2653 .first()
2654 .and_then(|&il| cache.kv[il].as_ref())
2655 .map(|k| k.len);
2656 ONCE.call_once(|| {
2657 eprintln!(
2658 "[fullg-debug] walk_uniform={} covered={} seqs_arms_on={} fa_rows_on={} t={} lo={} hi={} lin={} fa={} t_cap={} len0={:?}",
2659 self.walk_uniform, self.covered, seqs_arms_on, dspark_fa_rows_on(), t, lo, hi,
2660 self.lin.len(), self.fa.len(), self.t_cap, len0
2661 );
2662 });
2663 }
2664 if !self.walk_uniform
2665 || !seqs_arms_on
2666 || !dspark_fa_rows_on()
2667 || t < 2
2668 || lo != 0
2669 || hi > self.covered
2670 || t > self.t_cap
2671 || self.fa.is_empty()
2672 {
2673 return None;
2674 }
2675 let cfg = &model.cfg;
2676 let head_dim_global = cfg.head_dim_k as usize;
2677 let nkv = cfg.n_head_kv as usize;
2678 let kvl0 = cache.kv[self.fa[0]].as_ref().unwrap();
2679 // the z-batched twins read stacked rows at the cache's kv dims — must equal the
2680 // projection stride (the body's guard, hoisted so ineligible models fall back
2681 // instead of refusing mid-capture).
2682 let geom = cfg.full_attention_geometry_at(self.fa[0] as u32);
2683 let kv_dim = geom.n_head_kv as usize * geom.head_dim_k as usize;
2684 if kvl0.kv_dim_k != kv_dim || kvl0.kv_dim_v != kv_dim {
2685 return None;
2686 }
2687 let len0 = kvl0.len;
2688 let (t_kv_first, t_kv_last) = (len0 + 1, len0 + t);
2689 if !crate::fa_seqs_eligible(t_kv_first, head_dim_global)
2690 || !crate::fa_seqs_eligible(t_kv_last, head_dim_global)
2691 || crate::fa_split_keys(t_kv_first, nkv) != crate::fa_split_keys(t_kv_last, nkv)
2692 {
2693 return None;
2694 }
2695 let rung = t_kv_last.next_power_of_two().max(256);
2696 if crate::fa_split_keys(rung, nkv) != crate::fa_split_keys(t_kv_last, nkv) {
2697 return None;
2698 }
2699 Some(rung)
2700 }
2701
2702 /// Run the WHOLE verify walk [lo, hi) as one captured graph at (vt=t, rung): stage
2703 /// the residual + refresh the per-vt position staging, capture on first encounter
2704 /// (2 executing warmups bracketed by a full linear-state save/restore; KV warmup
2705 /// appends write the exact slots the replay writes — idempotent), launch, then apply
2706 /// the host bookkeeping the captured body skipped (per-linear-layer parity swap for
2707 /// odd t, per-fa-layer len bump). Returns the fresh residual.
2708 #[allow(clippy::too_many_arguments)]
2709 pub(crate) fn run_full(
2710 &mut self,
2711 model: &crate::hybrid::HybridModel,
2712 e: &Engine,
2713 lo: usize,
2714 hi: usize,
2715 x: &CudaSlice<f32>,
2716 t: usize,
2717 pos0: usize,
2718 rung: usize,
2719 cache: &mut Cache,
2720 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2721 let n_embd = self.n_embd;
2722 if !self.stage.contains_key(&t) {
2723 let xin = e.uninit(t * n_embd)?;
2724 let xout = e.uninit(t * n_embd)?;
2725 self.stage.insert(t, (xin, xout));
2726 }
2727 if !self.pos_stage.contains_key(&t) {
2728 self.pos_stage.insert(t, e.htod_i32(&vec![0i32; t])?);
2729 }
2730 // Per-round refresh: position contents + input staging (both addresses are baked
2731 // by the captured bodies; only their CONTENTS change round to round).
2732 {
2733 let pos_host: Vec<i32> = (0..t).map(|r| (pos0 + r) as i32).collect();
2734 let pb = self.pos_stage.get_mut(&t).unwrap();
2735 e.htod_i32_into(pb, &pos_host)?;
2736 let (xin, _) = self.stage.get_mut(&t).unwrap();
2737 e.copy_into(xin, 0, x, t * n_embd)?;
2738 }
2739 let key = (t, rung, hi);
2740 if !self.full.contains_key(&key) {
2741 // The warmups EXECUTE the whole walk on live state — save every linear
2742 // layer's conv + canonical ssm first, restore after (KV needs no restore:
2743 // graph mode never bumps host lens and the appends write this round's own
2744 // slots).
2745 for (k, &il) in self.lin.iter().enumerate() {
2746 let rl = cache.recur[il].as_ref().unwrap();
2747 e.copy_into(
2748 &mut self.save_conv,
2749 k * self.conv_words,
2750 &rl.conv_state,
2751 self.conv_words,
2752 )?;
2753 e.copy_into(
2754 &mut self.save_ssm,
2755 k * self.ssm_words,
2756 &rl.ssm_state,
2757 self.ssm_words,
2758 )?;
2759 }
2760 let (graph, keeper) = {
2761 let table_all = &self.table_all;
2762 let lin_pos = &self.lin_pos;
2763 let fa_pos = &self.fa_pos;
2764 let fa_table = &self.fa_table;
2765 let t_cap = self.t_cap;
2766 let stash_conv = &mut self.stash_conv;
2767 let stash_ssm = &mut self.stash_ssm;
2768 let pos_d: &CudaSlice<i32> = &self.pos_stage[&t];
2769 let (xin, xout) = self
2770 .stage
2771 .get_mut(&t)
2772 .map(|(a, b)| (&*a, b))
2773 .expect("stage bucket created above");
2774 let cache_ref: &mut Cache = cache;
2775 let iflag = if std::env::var("MEMRA_DSPARK_VG_AUTOFREE").as_deref() == Ok("1") {
2776 cudarc::driver::sys::CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH
2777 } else {
2778 cudarc::driver::sys::CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_USE_NODE_PRIORITY
2779 };
2780 e.capture_graph_retained_flags(iflag, move |e| {
2781 let mut xc: Option<CudaSlice<f32>> = None;
2782 for il in lo..hi {
2783 let xr: &CudaSlice<f32> = xc.as_ref().unwrap_or(xin);
2784 let nx = if let Some(&k) = lin_pos.get(&il) {
2785 model.qwen35_tparallel_linear_layer(
2786 e,
2787 il,
2788 xr,
2789 t,
2790 cache_ref,
2791 None,
2792 Some((&mut stash_conv[k], &mut stash_ssm[k])),
2793 Some((table_all, k * 6)),
2794 )?
2795 } else if let Some(&kf) = fa_pos.get(&il) {
2796 let mut no_rows: Option<Vec<CudaSlice<i32>>> = None;
2797 model.qwen35_tparallel_fa_layer(
2798 e,
2799 il,
2800 xr,
2801 t,
2802 cache_ref,
2803 FaLayerArgs {
2804 pos_d,
2805 pos_rows: &mut no_rows,
2806 pos0,
2807 seqs_append: true,
2808 batch_fa_on: true,
2809 graph_cap: Some((fa_table, kf * 2 * t_cap, rung)),
2810 stream: None,
2811 ckpt: None,
2812 },
2813 )?
2814 } else {
2815 return Err(format!(
2816 "run_full: layer {il} is neither linear nor full-attention"
2817 )
2818 .into());
2819 };
2820 xc = Some(nx);
2821 }
2822 e.copy_into(xout, 0, xc.as_ref().unwrap(), t * n_embd)?;
2823 Ok(())
2824 })?
2825 };
2826 // Undo the net host parity motion of the 3 body runs (each run swaps iff t
2827 // is odd -> 3 runs = net one swap), then restore the device state the
2828 // warmups consumed (walk scope only — layers past hi never executed). The
2829 // launch below then behaves exactly like one run.
2830 if t % 2 == 1 {
2831 for &il in &self.lin {
2832 if il < lo || il >= hi {
2833 continue;
2834 }
2835 let rl = cache.recur[il].as_mut().unwrap();
2836 std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
2837 }
2838 }
2839 for (k, &il) in self.lin.iter().enumerate() {
2840 if il < lo || il >= hi {
2841 continue;
2842 }
2843 let rl = cache.recur[il].as_mut().unwrap();
2844 let (cw, sw) = (self.conv_words, self.ssm_words);
2845 {
2846 let sv = e.view(&self.save_conv, self.lin.len() * cw);
2847 let win = sv.slice(k * cw..(k + 1) * cw);
2848 e.copy_view_into(&mut rl.conv_state, 0, &win, cw)?;
2849 }
2850 {
2851 let sv = e.view(&self.save_ssm, self.lin.len() * sw);
2852 let win = sv.slice(k * sw..(k + 1) * sw);
2853 e.copy_view_into(&mut rl.ssm_state, 0, &win, sw)?;
2854 }
2855 }
2856 if std::env::var("MEMRA_GRAPH_CENSUS").as_deref() == Ok("1") {
2857 if let Ok(c) = crate::graph_update::node_census(&graph) {
2858 eprintln!("[dspark-vg-census] full vt={t} rung={rung} {c:?}");
2859 }
2860 }
2861 self.full.insert(
2862 key,
2863 DsparkSegGraph {
2864 graph,
2865 _keeper: keeper,
2866 },
2867 );
2868 }
2869 self.full[&key].graph.launch()?;
2870 // Host bookkeeping for the replayed body (captured host code does not re-run):
2871 // gdn parity swap per linear layer (t odd), kv len bump per fa layer — scoped
2872 // to the WALK [lo, hi): the cache can carry extra state slots past it (the MTP
2873 // head layer's kv) that the walk never touches.
2874 if t % 2 == 1 {
2875 for &il in &self.lin {
2876 if il < lo || il >= hi {
2877 continue;
2878 }
2879 let rl = cache.recur[il].as_mut().unwrap();
2880 std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
2881 }
2882 }
2883 for &il in &self.fa {
2884 if il < lo || il >= hi {
2885 continue;
2886 }
2887 cache.kv[il].as_mut().unwrap().len += t;
2888 }
2889 let (_, xout) = self.stage.get(&t).unwrap();
2890 let mut out = e.uninit(t * n_embd)?;
2891 e.copy_into(&mut out, 0, xout, t * n_embd)?;
2892 Ok(out)
2893 }
2894
2895 /// Run layers [start, end) (all linear) as one captured graph at this vt: stage the
2896 /// residual into the bucket's x_in, capture on first encounter (2 executing warmups
2897 /// bracketed by a segment state save/restore), launch, then apply the host parity
2898 /// bookkeeping the captured body would have done. Returns the fresh residual.
2899 #[allow(clippy::too_many_arguments)]
2900 fn run_segment(
2901 &mut self,
2902 model: &crate::hybrid::HybridModel,
2903 e: &Engine,
2904 start: usize,
2905 end: usize,
2906 x: &CudaSlice<f32>,
2907 t: usize,
2908 cache: &mut Cache,
2909 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2910 let n_embd = self.n_embd;
2911 debug_assert!(end - start <= self.max_run);
2912 if !self.stage.contains_key(&t) {
2913 let xin = e.uninit(t * n_embd)?;
2914 let xout = e.uninit(t * n_embd)?;
2915 self.stage.insert(t, (xin, xout));
2916 }
2917 // Stage the residual at the bucket's baked input address.
2918 {
2919 let (xin, _) = self.stage.get_mut(&t).unwrap();
2920 e.copy_into(xin, 0, x, t * n_embd)?;
2921 }
2922 let key = (start, t);
2923 if !self.graphs.contains_key(&key) {
2924 // The 2 warmups EXECUTE the segment on live state — save conv + the canonical
2925 // ssm of every segment layer first, restore after, so the graph's first real
2926 // launch starts from the exact pre-round state (bytes gated e2e).
2927 for (k, il) in (start..end).enumerate() {
2928 let rl = cache.recur[il].as_ref().unwrap();
2929 e.copy_into(
2930 &mut self.save_conv,
2931 k * self.conv_words,
2932 &rl.conv_state,
2933 self.conv_words,
2934 )?;
2935 e.copy_into(
2936 &mut self.save_ssm,
2937 k * self.ssm_words,
2938 &rl.ssm_state,
2939 self.ssm_words,
2940 )?;
2941 }
2942 let (graph, keeper) = {
2943 let table_all = &self.table_all;
2944 let lin_pos = &self.lin_pos;
2945 let stash_conv = &mut self.stash_conv;
2946 let stash_ssm = &mut self.stash_ssm;
2947 let (xin, xout) = self
2948 .stage
2949 .get_mut(&t)
2950 .map(|(a, b)| (&*a, b))
2951 .expect("stage bucket created above");
2952 let cache_ref: &mut Cache = cache;
2953 // Slice 4 (fa-execupdate lane): USE_NODE_PRIORITY instead of
2954 // AUTO_FREE_ON_LAUNCH. The slice-3 measured limiter was AUTO_FREE's
2955 // launch-time mem-pool scan — 25.6 us per cuGraphLaunch x 16 segments
2956 // = ~0.41 ms/round, most of the eager-launch savings. The captured
2957 // body's cuMemAllocAsync transients are BALANCED by in-graph frees
2958 // (every transient drops inside the capture region — the generic
2959 // capture path's census precedent, 1589/1589), so AUTO_FREE has
2960 // nothing to reclaim and the graph is legal to instantiate without
2961 // it; PRIORITY is the flag the gemma slotted door ships for exactly
2962 // this reason (both alternatives drop the scan; UPLOAD via
2963 // cuGraphInstantiateWithFlags is WithParams-only and refused).
2964 // MEMRA_DSPARK_VG_AUTOFREE=1 reverts; MEMRA_GRAPH_CENSUS=1 prints
2965 // the node census at capture (the ALLOC==FREE receipt).
2966 let iflag = if std::env::var("MEMRA_DSPARK_VG_AUTOFREE").as_deref() == Ok("1") {
2967 cudarc::driver::sys::CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH
2968 } else {
2969 cudarc::driver::sys::CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_USE_NODE_PRIORITY
2970 };
2971 e.capture_graph_retained_flags(iflag, move |e| {
2972 let mut xc: Option<CudaSlice<f32>> = None;
2973 for il in start..end {
2974 let k = lin_pos[&il];
2975 let xr: &CudaSlice<f32> = xc.as_ref().unwrap_or(xin);
2976 let nx = model.qwen35_tparallel_linear_layer(
2977 e,
2978 il,
2979 xr,
2980 t,
2981 cache_ref,
2982 None,
2983 Some((&mut stash_conv[k], &mut stash_ssm[k])),
2984 Some((table_all, k * 6)),
2985 )?;
2986 xc = Some(nx);
2987 }
2988 e.copy_into(xout, 0, xc.as_ref().unwrap(), t * n_embd)?;
2989 Ok(())
2990 })?
2991 };
2992 // Undo the net host parity motion of the 3 body runs (each run swaps iff t
2993 // is odd -> 3 runs = net one swap), then restore the device state the
2994 // warmups consumed. The launch below then behaves exactly like one run.
2995 if t % 2 == 1 {
2996 for il in start..end {
2997 let rl = cache.recur[il].as_mut().unwrap();
2998 std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
2999 }
3000 }
3001 for (k, il) in (start..end).enumerate() {
3002 let rl = cache.recur[il].as_mut().unwrap();
3003 let (cw, sw) = (self.conv_words, self.ssm_words);
3004 {
3005 let sv = e.view(&self.save_conv, self.lin.len() * cw);
3006 let win = sv.slice(k * cw..(k + 1) * cw);
3007 e.copy_view_into(&mut rl.conv_state, 0, &win, cw)?;
3008 }
3009 {
3010 let sv = e.view(&self.save_ssm, self.lin.len() * sw);
3011 let win = sv.slice(k * sw..(k + 1) * sw);
3012 e.copy_view_into(&mut rl.ssm_state, 0, &win, sw)?;
3013 }
3014 }
3015 if std::env::var("MEMRA_GRAPH_CENSUS").as_deref() == Ok("1") {
3016 if let Ok(c) = crate::graph_update::node_census(&graph) {
3017 eprintln!("[dspark-vg-census] seg={start}..{end} vt={t} {c:?}");
3018 }
3019 }
3020 self.graphs.insert(
3021 key,
3022 DsparkSegGraph {
3023 graph,
3024 _keeper: keeper,
3025 },
3026 );
3027 }
3028 self.graphs[&key].graph.launch()?;
3029 // Host parity bookkeeping for the replayed body (the captured host swaps do not
3030 // re-run at replay).
3031 if t % 2 == 1 {
3032 for il in start..end {
3033 let rl = cache.recur[il].as_mut().unwrap();
3034 std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
3035 }
3036 }
3037 let (_, xout) = self.stage.get(&t).unwrap();
3038 let mut out = e.uninit(t * n_embd)?;
3039 e.copy_into(&mut out, 0, xout, t * n_embd)?;
3040 Ok(out)
3041 }
3042
3043 /// Pool freeze check (`dspark_vg_cap`): below the ceiling new keys may capture.
3044 fn can_capture(&self) -> bool {
3045 self.graphs.len() + self.full.len() < dspark_vg_cap()
3046 }
3047
3048 /// Round-atomic segment-door readiness: TRUE when this round's walk can ride the
3049 /// per-(segment, vt) graphs without a NEW capture past the pool ceiling — every
3050 /// linear run in [lo, hi) already has its (run_start, t) key, or capture is still
3051 /// allowed. FALSE sends the WHOLE round down the eager cols-ckpt walk: a partial
3052 /// refusal would stash some layers in the ctx slabs and others in the round's cols
3053 /// while one commit reads only one of them.
3054 pub(crate) fn segments_ready(
3055 &self,
3056 model: &crate::hybrid::HybridModel,
3057 lo: usize,
3058 hi: usize,
3059 t: usize,
3060 ) -> bool {
3061 if self.can_capture() {
3062 return true;
3063 }
3064 let mut il = lo;
3065 while il < hi {
3066 if matches!(model.layers[il].mixer, Mixer::Linear(_)) {
3067 let start = il;
3068 while il < hi && matches!(model.layers[il].mixer, Mixer::Linear(_)) {
3069 il += 1;
3070 }
3071 if !self.graphs.contains_key(&(start, t)) {
3072 return false;
3073 }
3074 } else {
3075 il += 1;
3076 }
3077 }
3078 true
3079 }
3080
3081 /// Widest verify window this pool was built for. A caller whose round exceeds it must
3082 /// take the eager walk: the stash slabs hold `t_capacity() - 1` column rows, and slicing
3083 /// past them is a panic rather than a refusal.
3084 pub(crate) fn t_capacity(&self) -> usize {
3085 self.t_cap
3086 }
3087
3088 /// Slab row (conv, ssm) device pointers + lengths for the commit restore of column
3089 /// `row` (0-based) of layer `il`. None for non-linear layers.
3090 pub(crate) fn slab_row(
3091 &self,
3092 e: &Engine,
3093 il: usize,
3094 row: usize,
3095 ) -> Option<(u64, u64, usize, usize)> {
3096 use cudarc::driver::DevicePtr;
3097 let k = *self.lin_pos.get(&il)?;
3098 let s = &e.gpu.stream();
3099 let (pc, _g0) = self.stash_conv[k].device_ptr(s);
3100 let (ps, _g1) = self.stash_ssm[k].device_ptr(s);
3101 Some((
3102 pc as u64 + (row * self.conv_words * 4) as u64,
3103 ps as u64 + (row * self.ssm_words * 4) as u64,
3104 self.conv_words,
3105 self.ssm_words,
3106 ))
3107 }
3108}
3109
3110impl VerifyCkpt {
3111 fn new(n_layer: usize) -> Self {
3112 VerifyCkpt {
3113 gdn: (0..n_layer).map(|_| None).collect(),
3114 cols: (0..n_layer).map(|_| None).collect(),
3115 }
3116 }
3117}
3118
3119/// The stage-0/TX half of one PP verify. The boundary slot is the ownership token: stage 1
3120/// consumes exactly the slot selected by `tx()` / `tx_pipelined()`, never a slot inferred from
3121/// a logical round number.
3122struct VerifyBoundaryTicket {
3123 rt: &'static crate::pp::PpNRt,
3124 caller_stream: std::sync::Arc<cudarc::driver::CudaStream>,
3125 slot: usize,
3126 pos0: usize,
3127 t: usize,
3128 payload: usize,
3129 n_st: usize,
3130 pipelined: bool,
3131 pp_anatomy: bool,
3132 pp_started: std::time::Instant,
3133 reverse_ms: f64,
3134 stage0_ms: f64,
3135 tx_ms: f64,
3136 trace: Option<SpecPipeTraceCtx>,
3137}
3138
3139/// Explicit OPTIPIPE diagnostic control. Forced modes are set only by `optipipe-gate`; the
3140/// increment-2 controller can also be armed by the server's fresh-process research door.
3141#[derive(Clone, Copy, Debug, PartialEq, Eq)]
3142pub enum OptiForkGateMode {
3143 Disabled,
3144 Hit,
3145 Miss,
3146 Alternate,
3147 Abort,
3148 Controller,
3149}
3150
3151static OPTI_FORK_GATE_MODE: std::sync::atomic::AtomicU8 = std::sync::atomic::AtomicU8::new(0);
3152static OPTI_CONTROLLER_THRESHOLD: std::sync::atomic::AtomicU32 =
3153 std::sync::atomic::AtomicU32::new(0);
3154static OPTI_FORK_ATTEMPTS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
3155static OPTI_FORK_HITS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
3156static OPTI_FORK_MISSES: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
3157static OPTI_FORK_ABORT_DRAINS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
3158static OPTI_FORK_REFUSALS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
3159static OPTI_GATE_CHECKS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
3160static OPTI_GATE_ADMITS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
3161static OPTI_GATE_REJECTS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
3162static OPTI_RECONCILES: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
3163static OPTI_WASTED_DRAFT_TOKENS: std::sync::atomic::AtomicU64 =
3164 std::sync::atomic::AtomicU64::new(0);
3165static OPTI_SHADOW_DRAFT_TOKENS: std::sync::atomic::AtomicU64 =
3166 std::sync::atomic::AtomicU64::new(0);
3167static OPTI_BREAKER_TRIPS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
3168
3169impl OptiForkGateMode {
3170 fn code(self) -> u8 {
3171 match self {
3172 Self::Disabled => 0,
3173 Self::Hit => 1,
3174 Self::Miss => 2,
3175 Self::Alternate => 3,
3176 Self::Abort => 4,
3177 Self::Controller => 5,
3178 }
3179 }
3180
3181 fn configured() -> Self {
3182 match OPTI_FORK_GATE_MODE.load(std::sync::atomic::Ordering::Relaxed) {
3183 1 => Self::Hit,
3184 2 => Self::Miss,
3185 3 => Self::Alternate,
3186 4 => Self::Abort,
3187 5 => Self::Controller,
3188 _ => Self::Disabled,
3189 }
3190 }
3191
3192 fn action(self, generation: u64) -> OptiForkAction {
3193 match self {
3194 Self::Hit => OptiForkAction::Hit,
3195 Self::Miss => OptiForkAction::Miss,
3196 Self::Alternate if generation & 1 == 0 => OptiForkAction::Hit,
3197 Self::Alternate => OptiForkAction::Miss,
3198 Self::Abort => OptiForkAction::Abort,
3199 Self::Disabled | Self::Controller => {
3200 unreachable!("non-forced mode cannot choose a forced fork action")
3201 }
3202 }
3203 }
3204
3205 fn is_forced(self) -> bool {
3206 matches!(self, Self::Hit | Self::Miss | Self::Alternate | Self::Abort)
3207 }
3208}
3209
3210/// Arm or disarm the forced harness. Serving uses only `set_optipipe_controller_threshold`.
3211pub fn set_optipipe_gate_mode(mode: OptiForkGateMode) {
3212 OPTI_FORK_GATE_MODE.store(mode.code(), std::sync::atomic::Ordering::Relaxed);
3213}
3214
3215/// Arm the increment-2 diagnostic controller. The threshold applies to the uncalibrated
3216/// two-token draft-probability product. Serving can call this only through its explicit
3217/// fresh-process research door; the absent-door default remains byte-for-byte disabled.
3218pub fn set_optipipe_controller_threshold(threshold: f32) {
3219 assert!(threshold.is_finite() && (0.0..=1.0).contains(&threshold));
3220 OPTI_CONTROLLER_THRESHOLD.store(threshold.to_bits(), std::sync::atomic::Ordering::Relaxed);
3221 set_optipipe_gate_mode(OptiForkGateMode::Controller);
3222}
3223
3224#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
3225pub struct OptiForkGateStats {
3226 pub attempts: u64,
3227 pub hits: u64,
3228 pub misses: u64,
3229 pub abort_drains: u64,
3230 pub refusals: u64,
3231 pub gate_checks: u64,
3232 pub gate_admits: u64,
3233 pub gate_rejects: u64,
3234 pub reconciles: u64,
3235 pub wasted_draft_tokens: u64,
3236 pub shadow_draft_tokens: u64,
3237 pub breaker_trips: u64,
3238}
3239
3240#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
3241pub struct OptiForkStateIdentity {
3242 pub trunk_kv_bytes: usize,
3243 pub recurrent_bytes: usize,
3244 pub scratch_kv_bytes: usize,
3245 pub hidden_bytes: usize,
3246}
3247
3248pub fn reset_optipipe_gate_stats() {
3249 for counter in [
3250 &OPTI_FORK_ATTEMPTS,
3251 &OPTI_FORK_HITS,
3252 &OPTI_FORK_MISSES,
3253 &OPTI_FORK_ABORT_DRAINS,
3254 &OPTI_FORK_REFUSALS,
3255 &OPTI_GATE_CHECKS,
3256 &OPTI_GATE_ADMITS,
3257 &OPTI_GATE_REJECTS,
3258 &OPTI_RECONCILES,
3259 &OPTI_WASTED_DRAFT_TOKENS,
3260 &OPTI_SHADOW_DRAFT_TOKENS,
3261 &OPTI_BREAKER_TRIPS,
3262 ] {
3263 counter.store(0, std::sync::atomic::Ordering::Relaxed);
3264 }
3265}
3266
3267pub fn optipipe_gate_stats() -> OptiForkGateStats {
3268 let load = |v: &std::sync::atomic::AtomicU64| v.load(std::sync::atomic::Ordering::Relaxed);
3269 OptiForkGateStats {
3270 attempts: load(&OPTI_FORK_ATTEMPTS),
3271 hits: load(&OPTI_FORK_HITS),
3272 misses: load(&OPTI_FORK_MISSES),
3273 abort_drains: load(&OPTI_FORK_ABORT_DRAINS),
3274 refusals: load(&OPTI_FORK_REFUSALS),
3275 gate_checks: load(&OPTI_GATE_CHECKS),
3276 gate_admits: load(&OPTI_GATE_ADMITS),
3277 gate_rejects: load(&OPTI_GATE_REJECTS),
3278 reconciles: load(&OPTI_RECONCILES),
3279 wasted_draft_tokens: load(&OPTI_WASTED_DRAFT_TOKENS),
3280 shadow_draft_tokens: load(&OPTI_SHADOW_DRAFT_TOKENS),
3281 breaker_trips: load(&OPTI_BREAKER_TRIPS),
3282 }
3283}
3284
3285#[derive(Clone, Copy, Debug)]
3286struct OptiControllerPolicy {
3287 threshold: f32,
3288 consecutive_misses: u8,
3289 breaker_tripped: bool,
3290}
3291
3292impl OptiControllerPolicy {
3293 fn configured() -> Self {
3294 Self {
3295 threshold: f32::from_bits(
3296 OPTI_CONTROLLER_THRESHOLD.load(std::sync::atomic::Ordering::Relaxed),
3297 ),
3298 consecutive_misses: 0,
3299 breaker_tripped: false,
3300 }
3301 }
3302
3303 fn admit(&self, q_proxy: f32) -> bool {
3304 q_proxy.is_finite()
3305 && (0.0..=1.0).contains(&q_proxy)
3306 && (self.threshold == 0.0 || (!self.breaker_tripped && q_proxy >= self.threshold))
3307 }
3308
3309 /// Returns true exactly when this resolution newly trips the three-miss breaker.
3310 fn resolve(&mut self, hit: bool) -> bool {
3311 // q*=0 is the lane's explicit unconditional measurement arm. Its purpose is to price
3312 // every optimistic opportunity, so the safety breaker is measured separately and must
3313 // not silently turn this arm into "three attempts then serial".
3314 if self.threshold == 0.0 {
3315 self.consecutive_misses = 0;
3316 return false;
3317 }
3318 if hit {
3319 self.consecutive_misses = 0;
3320 return false;
3321 }
3322 self.consecutive_misses = self.consecutive_misses.saturating_add(1);
3323 if !self.breaker_tripped && self.consecutive_misses >= 3 {
3324 self.breaker_tripped = true;
3325 return true;
3326 }
3327 false
3328 }
3329}
3330
3331#[derive(Clone, Copy, Debug, PartialEq, Eq)]
3332enum OptiForkAction {
3333 Hit,
3334 Miss,
3335 Abort,
3336}
3337
3338#[derive(Clone, Copy, Debug, PartialEq, Eq)]
3339struct OptiForkGeneration {
3340 id: u64,
3341 slot: usize,
3342}
3343
3344#[derive(Default)]
3345struct OptiForkGenerationTracker {
3346 next: u64,
3347 live: [Option<u64>; 2],
3348}
3349
3350impl OptiForkGenerationTracker {
3351 fn reserve(&mut self) -> Result<OptiForkGeneration, Box<dyn std::error::Error>> {
3352 let generation = OptiForkGeneration {
3353 id: self.next,
3354 slot: (self.next & 1) as usize,
3355 };
3356 if let Some(live) = self.live[generation.slot] {
3357 return Err(format!(
3358 "optipipe snapshot slot {} still owns generation {live}; refusing to overwrite it",
3359 generation.slot,
3360 )
3361 .into());
3362 }
3363 self.next += 1;
3364 self.live[generation.slot] = Some(generation.id);
3365 Ok(generation)
3366 }
3367
3368 fn retire(&mut self, generation: OptiForkGeneration) -> Result<(), Box<dyn std::error::Error>> {
3369 match self.live[generation.slot] {
3370 Some(id) if id == generation.id => {
3371 self.live[generation.slot] = None;
3372 Ok(())
3373 }
3374 other => Err(format!(
3375 "optipipe generation teardown mismatch: ticket={} slot={} live={other:?}",
3376 generation.id, generation.slot,
3377 )
3378 .into()),
3379 }
3380 }
3381}
3382
3383struct OptiForkSeedGeneration {
3384 h_seed: CudaSlice<f32>,
3385 fill_prev: CudaSlice<f32>,
3386 scratch_len: usize,
3387}
3388
3389/// Allocate or refresh one full checkpoint through the engine that owns each PP stage. The
3390/// generic cache helper accepts one device and therefore cannot copy GDN state split across
3391/// devices. KV lengths and position stay host metadata; only recurrent buffers need stage-local
3392/// device ownership.
3393fn opti_snapshot_stage_owned(
3394 e: &Engine,
3395 cache: &Cache,
3396 rt: &'static crate::pp::PpNRt,
3397 fence: &[usize],
3398) -> Result<crate::cache::CacheSnapshot, Box<dyn std::error::Error>> {
3399 let n = cache.kv.len();
3400 let mut snapshot = crate::cache::CacheSnapshot {
3401 kv_len: vec![None; n],
3402 tp_kv_len: vec![None; n],
3403 conv: (0..n).map(|_| None).collect(),
3404 ssm: (0..n).map(|_| None).collect(),
3405 pos: cache.pos,
3406 };
3407 opti_snapshot_stage_owned_into(e, cache, rt, fence, &mut snapshot)?;
3408 Ok(snapshot)
3409}
3410
3411fn opti_snapshot_stage_owned_into(
3412 e: &Engine,
3413 cache: &Cache,
3414 rt: &'static crate::pp::PpNRt,
3415 fence: &[usize],
3416 snapshot: &mut crate::cache::CacheSnapshot,
3417) -> Result<(), Box<dyn std::error::Error>> {
3418 if fence.len() != rt.n_stages() + 1
3419 || snapshot.kv_len.len() != cache.kv.len()
3420 || snapshot.tp_kv_len.len() != cache.tp_kv.len()
3421 {
3422 return Err("optipipe stage-owned snapshot shape mismatch".into());
3423 }
3424 for stage in 0..rt.n_stages() {
3425 opti_snapshot_one_stage_owned_into(e, cache, rt, fence, stage, snapshot)?;
3426 }
3427 snapshot.pos = cache.pos;
3428 Ok(())
3429}
3430
3431/// Refresh one PP stage of a checkpoint. Increment 2 uses this split form so stage 0's
3432/// optimistic post-N state is captured before N+1 stage 0 is queued, while stage 1's matching
3433/// post-N state is captured only after N stage 1 is enqueued. Calling the all-stage helper at
3434/// either point would capture one side of the fork at the wrong generation.
3435fn opti_snapshot_one_stage_owned_into(
3436 e: &Engine,
3437 cache: &Cache,
3438 rt: &'static crate::pp::PpNRt,
3439 fence: &[usize],
3440 stage: usize,
3441 snapshot: &mut crate::cache::CacheSnapshot,
3442) -> Result<(), Box<dyn std::error::Error>> {
3443 if fence.len() != rt.n_stages() + 1
3444 || snapshot.kv_len.len() != cache.kv.len()
3445 || snapshot.tp_kv_len.len() != cache.tp_kv.len()
3446 || stage >= rt.n_stages()
3447 {
3448 return Err("optipipe single-stage snapshot shape mismatch".into());
3449 }
3450 let _scope = rt.enter(stage);
3451 let owner = rt.engine(stage, e);
3452 for il in fence[stage]..fence[stage + 1] {
3453 snapshot.kv_len[il] = cache.kv[il].as_ref().map(|kv| kv.len);
3454 snapshot.tp_kv_len[il] = cache.tp_kv[il]
3455 .as_ref()
3456 .map(crate::tp::ResidentTpKvCache::committed_len);
3457 match &cache.recur[il] {
3458 Some(recur) => {
3459 match snapshot.conv[il].as_mut() {
3460 Some(dst) => {
3461 owner.copy_into(dst, 0, &recur.conv_state, recur.conv_state.len())?
3462 }
3463 None => snapshot.conv[il] = Some(owner.clone_dtod(&recur.conv_state)?),
3464 }
3465 match snapshot.ssm[il].as_mut() {
3466 Some(dst) => {
3467 owner.copy_into(dst, 0, &recur.ssm_state, recur.ssm_state.len())?
3468 }
3469 None => snapshot.ssm[il] = Some(owner.clone_dtod(&recur.ssm_state)?),
3470 }
3471 }
3472 None if snapshot.conv[il].is_some() || snapshot.ssm[il].is_some() => {
3473 return Err(
3474 format!("optipipe stage-owned snapshot layer {il} changed shape").into(),
3475 );
3476 }
3477 None => {}
3478 }
3479 }
3480 snapshot.pos = cache.pos;
3481 Ok(())
3482}
3483
3484/// Increment-1 persistent fork state. Exactly two snapshot/seed slots alternate; a live ticket
3485/// names its generation and keeps teardown fail-closed. Only stage 0 is allowed to mutate before
3486/// resolve, so the reconcile tables and conditional restores are stage-local.
3487struct OptiForkState {
3488 mode: OptiForkGateMode,
3489 controller: Option<OptiControllerPolicy>,
3490 generations: OptiForkGenerationTracker,
3491 active_snapshot_slot: usize,
3492 alternate_snapshot: crate::cache::CacheSnapshot,
3493 seeds: [OptiForkSeedGeneration; 2],
3494 rt: &'static crate::pp::PpNRt,
3495 fence: [usize; 3],
3496 split: usize,
3497 len_ptrs: CudaSlice<u64>,
3498 saved_lens: CudaSlice<i32>,
3499 forced_acc: CudaSlice<u32>,
3500 valid: CudaSlice<u32>,
3501 stage0_stream: std::sync::Arc<cudarc::driver::CudaStream>,
3502 logical_payload_bytes: [usize; 2],
3503}
3504
3505struct OptiForkTicket {
3506 generation: OptiForkGeneration,
3507 boundary: Option<VerifyBoundaryTicket>,
3508 drain: std::sync::Arc<cudarc::driver::CudaStream>,
3509 settled: bool,
3510}
3511
3512struct OptiControllerTicket {
3513 generation: OptiForkGeneration,
3514 boundary: Option<VerifyBoundaryTicket>,
3515 ckpt: Option<VerifyCkpt>,
3516 verify_tokens: [u32; 2],
3517 draft_prob: f32,
3518 eager_seed: Option<CudaSlice<f32>>,
3519 q_proxy: f32,
3520 scratch_len: usize,
3521 issued_at: std::time::Instant,
3522 drain: std::sync::Arc<cudarc::driver::CudaStream>,
3523 settled: bool,
3524}
3525
3526struct OptiControllerPrepared {
3527 verify_tokens: [u32; 2],
3528 draft_prob: f32,
3529 eager_seed: Option<CudaSlice<f32>>,
3530 q_proxy: f32,
3531 scratch_len: usize,
3532}
3533
3534impl OptiControllerTicket {
3535 fn take_boundary(&mut self) -> VerifyBoundaryTicket {
3536 self.boundary
3537 .take()
3538 .expect("controller boundary ticket already consumed")
3539 }
3540
3541 fn take_ckpt(&mut self) -> VerifyCkpt {
3542 self.ckpt
3543 .take()
3544 .expect("controller verify checkpoint already consumed")
3545 }
3546
3547 fn take_eager_seed(&mut self) -> Option<CudaSlice<f32>> {
3548 self.eager_seed.take()
3549 }
3550
3551 fn settle(&mut self) {
3552 self.settled = true;
3553 }
3554}
3555
3556impl Drop for OptiControllerTicket {
3557 fn drop(&mut self) {
3558 if !self.settled {
3559 let _ = self.drain.synchronize();
3560 OPTI_FORK_ABORT_DRAINS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
3561 }
3562 }
3563}
3564
3565impl OptiForkTicket {
3566 fn take_boundary(&mut self) -> VerifyBoundaryTicket {
3567 self.boundary
3568 .take()
3569 .expect("fork ticket boundary already consumed")
3570 }
3571
3572 fn settle(&mut self) {
3573 self.settled = true;
3574 }
3575}
3576
3577impl Drop for OptiForkTicket {
3578 fn drop(&mut self) {
3579 if !self.settled {
3580 let _ = self.drain.synchronize();
3581 OPTI_FORK_ABORT_DRAINS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
3582 }
3583 }
3584}
3585
3586impl OptiForkState {
3587 #[allow(clippy::too_many_arguments)]
3588 fn new(
3589 e: &Engine,
3590 cache: &Cache,
3591 mode: OptiForkGateMode,
3592 alternate_snapshot: crate::cache::CacheSnapshot,
3593 h_seed: &CudaSlice<f32>,
3594 fill_prev: &CudaSlice<f32>,
3595 rt: &'static crate::pp::PpNRt,
3596 split: usize,
3597 n_layer: usize,
3598 ) -> Result<Self, Box<dyn std::error::Error>> {
3599 let fence = [0, split, n_layer];
3600 let mut logical_payload_bytes = [0usize; 2];
3601 for stage in 0..2 {
3602 for il in fence[stage]..fence[stage + 1] {
3603 logical_payload_bytes[stage] += alternate_snapshot.conv[il]
3604 .as_ref()
3605 .map_or(0, |v| v.len() * std::mem::size_of::<f32>());
3606 logical_payload_bytes[stage] += alternate_snapshot.ssm[il]
3607 .as_ref()
3608 .map_or(0, |v| v.len() * std::mem::size_of::<f32>());
3609 }
3610 }
3611 let seeds = [
3612 OptiForkSeedGeneration {
3613 h_seed: e.clone_dtod(h_seed)?,
3614 fill_prev: e.clone_dtod(fill_prev)?,
3615 scratch_len: 0,
3616 },
3617 OptiForkSeedGeneration {
3618 h_seed: e.clone_dtod(h_seed)?,
3619 fill_prev: e.clone_dtod(fill_prev)?,
3620 scratch_len: 0,
3621 },
3622 ];
3623 let (len_ptrs, saved_lens, forced_acc, valid, stage0_stream) = {
3624 let _stage = rt.enter(0);
3625 let e0 = rt.engine(0, e);
3626 (
3627 crate::round_stream::kv_len_ptr_table_range(e0, cache, 0..split, None)?,
3628 e0.htod_i32(&vec![0; split])?,
3629 e0.alloc_u32_zeroed(2)?,
3630 e0.alloc_u32_zeroed(1)?,
3631 e0.stream(),
3632 )
3633 };
3634 logical_payload_bytes[0] += seeds
3635 .iter()
3636 .map(|seed| (seed.h_seed.len() + seed.fill_prev.len()) * std::mem::size_of::<f32>())
3637 .sum::<usize>();
3638 logical_payload_bytes[0] += len_ptrs.len() * std::mem::size_of::<u64>()
3639 + saved_lens.len() * std::mem::size_of::<i32>()
3640 + forced_acc.len() * std::mem::size_of::<u32>()
3641 + valid.len() * std::mem::size_of::<u32>();
3642 Ok(Self {
3643 mode,
3644 controller: (mode == OptiForkGateMode::Controller)
3645 .then(OptiControllerPolicy::configured),
3646 generations: OptiForkGenerationTracker::default(),
3647 active_snapshot_slot: 0,
3648 alternate_snapshot,
3649 seeds,
3650 rt,
3651 fence,
3652 split,
3653 len_ptrs,
3654 saved_lens,
3655 forced_acc,
3656 valid,
3657 stage0_stream,
3658 logical_payload_bytes,
3659 })
3660 }
3661
3662 fn reserve(
3663 &mut self,
3664 current_snapshot: &mut crate::cache::CacheSnapshot,
3665 ) -> Result<OptiForkGeneration, Box<dyn std::error::Error>> {
3666 let generation = self.generations.reserve()?;
3667 if generation.slot != self.active_snapshot_slot {
3668 std::mem::swap(current_snapshot, &mut self.alternate_snapshot);
3669 self.active_snapshot_slot = generation.slot;
3670 }
3671 Ok(generation)
3672 }
3673
3674 fn capture_seed(
3675 &mut self,
3676 e: &Engine,
3677 generation: OptiForkGeneration,
3678 h_seed: &CudaSlice<f32>,
3679 fill_prev: &CudaSlice<f32>,
3680 scratch_len: usize,
3681 ) -> Result<(), Box<dyn std::error::Error>> {
3682 let seed = &mut self.seeds[generation.slot];
3683 e.copy_into(&mut seed.h_seed, 0, h_seed, h_seed.len())?;
3684 e.copy_into(&mut seed.fill_prev, 0, fill_prev, fill_prev.len())?;
3685 seed.scratch_len = scratch_len;
3686 Ok(())
3687 }
3688
3689 fn ticket(
3690 &self,
3691 generation: OptiForkGeneration,
3692 boundary: VerifyBoundaryTicket,
3693 ) -> OptiForkTicket {
3694 OptiForkTicket {
3695 generation,
3696 boundary: Some(boundary),
3697 drain: self.stage0_stream.clone(),
3698 settled: false,
3699 }
3700 }
3701
3702 #[allow(clippy::too_many_arguments)]
3703 fn controller_ticket(
3704 &self,
3705 generation: OptiForkGeneration,
3706 boundary: VerifyBoundaryTicket,
3707 ckpt: VerifyCkpt,
3708 verify_tokens: [u32; 2],
3709 draft_prob: f32,
3710 eager_seed: Option<CudaSlice<f32>>,
3711 q_proxy: f32,
3712 scratch_len: usize,
3713 ) -> OptiControllerTicket {
3714 OptiControllerTicket {
3715 generation,
3716 boundary: Some(boundary),
3717 ckpt: Some(ckpt),
3718 verify_tokens,
3719 draft_prob,
3720 eager_seed,
3721 q_proxy,
3722 scratch_len,
3723 issued_at: std::time::Instant::now(),
3724 drain: self.stage0_stream.clone(),
3725 settled: false,
3726 }
3727 }
3728
3729 fn reserve_successor(&mut self) -> Result<OptiForkGeneration, Box<dyn std::error::Error>> {
3730 self.generations.reserve()
3731 }
3732
3733 fn successor_snapshot_mut(&mut self) -> &mut crate::cache::CacheSnapshot {
3734 &mut self.alternate_snapshot
3735 }
3736
3737 fn promote_successor_snapshot(
3738 &mut self,
3739 current_snapshot: &mut crate::cache::CacheSnapshot,
3740 generation: OptiForkGeneration,
3741 ) {
3742 std::mem::swap(current_snapshot, &mut self.alternate_snapshot);
3743 self.active_snapshot_slot = generation.slot;
3744 }
3745
3746 fn queue_actual_reconcile(
3747 &mut self,
3748 e: &Engine,
3749 snapshot: &crate::cache::CacheSnapshot,
3750 acc: &CudaSlice<u32>,
3751 optimistic_pending: u32,
3752 base: usize,
3753 ) -> Result<(), Box<dyn std::error::Error>> {
3754 let saved: Vec<i32> = (0..self.split)
3755 .map(|il| snapshot.kv_len[il].map(|v| v as i32).unwrap_or(0))
3756 .collect();
3757 // Serving keeps the caller/accept walk on the head (stage-1) device. Record the accept
3758 // decision point there and append a wait to stage 0 after its optimistic successor/TX;
3759 // the validity/reconcile kernels must never peer-read acc before it is written. The
3760 // increment-1 harness uses primary stage 0, where stream order already provides this.
3761 if self.rt.engine(0, e).ctx().ordinal() != e.ctx().ordinal() {
3762 self.rt.fence_stages_behind(&e.stream())?;
3763 }
3764 let _stage = self.rt.enter(0);
3765 let e0 = self.rt.engine(0, e);
3766 e0.htod_i32_into(&mut self.saved_lens, &saved)?;
3767 e0.spec_fork_valid(acc, optimistic_pending, &mut self.valid)?;
3768 e0.spec_fork_reconcile_kv(
3769 &self.len_ptrs,
3770 &self.saved_lens,
3771 acc,
3772 &self.valid,
3773 base,
3774 self.split,
3775 )
3776 }
3777
3778 fn finish_actual_reconcile(
3779 &mut self,
3780 e: &Engine,
3781 cache: &mut Cache,
3782 snapshot: &crate::cache::CacheSnapshot,
3783 n_acc: usize,
3784 base: usize,
3785 hit: bool,
3786 ) -> Result<(), Box<dyn std::error::Error>> {
3787 if hit {
3788 return Ok(());
3789 }
3790 let len_delta = base + n_acc;
3791 for il in 0..self.split {
3792 if let (Some(kv), Some(saved)) = (cache.kv[il].as_mut(), snapshot.kv_len[il]) {
3793 kv.len = saved + len_delta;
3794 }
3795 }
3796 {
3797 let _stage = self.rt.enter(1);
3798 let e1 = self.rt.engine(1, e);
3799 for il in self.split..self.fence[2] {
3800 if let (Some(kv), Some(saved)) = (cache.kv[il].as_mut(), snapshot.kv_len[il]) {
3801 kv.len = saved + len_delta;
3802 e1.set_i32_one(&mut kv.len_d, kv.len as i32)?;
3803 }
3804 }
3805 }
3806 self.rt.publish_to(0, &e.stream())?;
3807 Ok(())
3808 }
3809
3810 fn cancel_controller_ticket(
3811 &mut self,
3812 e: &Engine,
3813 cache: &mut Cache,
3814 scratch: &mut MtpScratch,
3815 snapshot: &crate::cache::CacheSnapshot,
3816 ticket: &mut OptiControllerTicket,
3817 ) -> Result<(), Box<dyn std::error::Error>> {
3818 {
3819 let _stage = self.rt.enter(0);
3820 let e0 = self.rt.engine(0, e);
3821 for il in 0..self.split {
3822 if let (Some(kv), Some(saved)) = (cache.kv[il].as_mut(), snapshot.kv_len[il]) {
3823 kv.len = saved;
3824 e0.set_i32_one(&mut kv.len_d, saved as i32)?;
3825 }
3826 }
3827 }
3828 scratch.set_len(e, snapshot.pos)?;
3829 ticket.settle();
3830 self.generations.retire(ticket.generation)?;
3831 OPTI_FORK_ABORT_DRAINS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
3832 OPTI_WASTED_DRAFT_TOKENS.fetch_add(2, std::sync::atomic::Ordering::Relaxed);
3833 eprintln!(
3834 "[opti-controller] tail-drain generation={} slot={}",
3835 ticket.generation.id, ticket.generation.slot,
3836 );
3837 Ok(())
3838 }
3839
3840 #[allow(clippy::too_many_arguments)]
3841 fn reconcile(
3842 &mut self,
3843 e: &Engine,
3844 cache: &mut Cache,
3845 scratch: &mut MtpScratch,
3846 snapshot: &crate::cache::CacheSnapshot,
3847 h_seed: &mut CudaSlice<f32>,
3848 fill_prev: &mut CudaSlice<f32>,
3849 generation: OptiForkGeneration,
3850 action: OptiForkAction,
3851 optimistic_pending: u32,
3852 ) -> Result<(), Box<dyn std::error::Error>> {
3853 debug_assert!(action != OptiForkAction::Abort);
3854 let miss_started = std::time::Instant::now();
3855 let keep = action == OptiForkAction::Hit;
3856 let saved: Vec<i32> = (0..self.split)
3857 .map(|il| snapshot.kv_len[il].map(|v| v as i32).unwrap_or(0))
3858 .collect();
3859 let seed = &self.seeds[generation.slot];
3860 {
3861 let _stage = self.rt.enter(0);
3862 let e0 = self.rt.engine(0, e);
3863 e0.htod_i32_into(&mut self.saved_lens, &saved)?;
3864 let forced = if keep {
3865 [1u32, optimistic_pending]
3866 } else {
3867 [0u32, optimistic_pending]
3868 };
3869 e0.htod_u32_into(&mut self.forced_acc, &forced)?;
3870 e0.spec_fork_valid(&self.forced_acc, optimistic_pending, &mut self.valid)?;
3871 e0.spec_fork_reconcile_kv(
3872 &self.len_ptrs,
3873 &self.saved_lens,
3874 &self.forced_acc,
3875 &self.valid,
3876 0,
3877 self.split,
3878 )?;
3879 for il in 0..self.split {
3880 if let Some(recur) = cache.recur[il].as_mut() {
3881 let conv = snapshot.conv[il]
3882 .as_ref()
3883 .ok_or("optipipe stage0 snapshot missing conv state")?;
3884 let ssm = snapshot.ssm[il]
3885 .as_ref()
3886 .ok_or("optipipe stage0 snapshot missing ssm state")?;
3887 e0.spec_fork_restore_f32(conv, &mut recur.conv_state, &self.valid)?;
3888 e0.spec_fork_restore_f32(ssm, &mut recur.ssm_state, &self.valid)?;
3889 }
3890 }
3891 e0.spec_fork_restore_f32(&seed.h_seed, h_seed, &self.valid)?;
3892 e0.spec_fork_restore_f32(&seed.fill_prev, fill_prev, &self.valid)?;
3893 }
3894
3895 if keep {
3896 OPTI_FORK_HITS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
3897 return Ok(());
3898 }
3899
3900 for il in 0..self.split {
3901 if let (Some(kv), Some(saved)) = (cache.kv[il].as_mut(), snapshot.kv_len[il]) {
3902 kv.len = saved;
3903 }
3904 }
3905 scratch.set_len(e, seed.scratch_len)?;
3906 // Targeted E_restart: publish only stage 0's reconcile to the caller, then bound the
3907 // forced diagnostic so the retained number is the actual miss cost, not enqueue time.
3908 let caller = e.stream();
3909 self.rt.publish_to(0, &caller)?;
3910 caller.synchronize()?;
3911 let miss_ms = miss_started.elapsed().as_secs_f64() * 1e3;
3912 eprintln!(
3913 "[opti-fork-reconcile] generation={} slot={} miss_ms={miss_ms:.3}",
3914 generation.id, generation.slot,
3915 );
3916 OPTI_FORK_MISSES.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
3917 Ok(())
3918 }
3919
3920 fn retire(&mut self, generation: OptiForkGeneration) -> Result<(), Box<dyn std::error::Error>> {
3921 self.generations.retire(generation)
3922 }
3923}
3924
3925fn rewind_tp_kv_verified_prefix(
3926 tp_kv: &mut [Option<crate::tp::ResidentTpKvCache>],
3927 saved_lens: &[Option<usize>],
3928 accepted: usize,
3929) -> Result<(), Box<dyn std::error::Error>> {
3930 if tp_kv.len() != saved_lens.len() {
3931 return Err("spec TP KV snapshot shape mismatch".into());
3932 }
3933 for (layer, (cache, saved)) in tp_kv.iter_mut().zip(saved_lens).enumerate() {
3934 match (cache.as_mut(), *saved) {
3935 (Some(cache), Some(saved)) => {
3936 let target = saved
3937 .checked_add(accepted)
3938 .ok_or("spec TP KV committed length overflow")?;
3939 cache.rewind_to(target)?;
3940 }
3941 (None, None) => {}
3942 _ => {
3943 return Err(
3944 format!("spec TP KV layer {layer} changed shape since its snapshot").into(),
3945 );
3946 }
3947 }
3948 }
3949 Ok(())
3950}
3951
3952/// MEMRA_WALK_SCRATCH=1: persistent per-layer temporaries for the verify walk. The walk allocated
3953/// `next`, `x1_t`, `z_t` and `x2_t` fresh EVERY LAYER — ~4 allocations x 45 layers x 2 ranks per
3954/// round. This session measured a device allocation at ~20 us (the sampled split head arrived
3955/// slower than the unsplit one purely on five allocations per token), so ~200 allocations is
3956/// ~4 ms of the 8.6 ms/round the [spec-phase] host-issue term reports. Stable addresses are also
3957/// the precondition for ever capturing this walk in a CUDA graph.
3958struct WalkScratch {
3959 dev: usize,
3960 cap: usize,
3961 x1: CudaSlice<f32>,
3962 z: CudaSlice<f32>,
3963}
3964static WALK_SCRATCH: std::sync::Mutex<Option<WalkScratch>> = std::sync::Mutex::new(None);
3965
3966fn walk_scratch_on() -> bool {
3967 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
3968 *ON.get_or_init(|| std::env::var("MEMRA_WALK_SCRATCH").as_deref() == Ok("1"))
3969}
3970
3971/// MEMRA_SPEC_ROUND_PROF counters: whole-round wall, so the round can be weighed against the
3972/// draft-step ([spec-anatomy]) and verify-walk ([tcol-prof]) splits we already print.
3973static ROUND_PROF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
3974static ROUND_MS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
3975static ROUND_N: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
3976
3977impl HybridModel {
3978 fn mtp_head_count(&self) -> usize {
3979 usize::from(self.mtp.is_some()) + self.mtp_extra.len()
3980 }
3981
3982 fn mtp_head_at(&self, index: usize) -> &MtpHead {
3983 if index == 0 {
3984 self.mtp.as_ref().expect("MTP head 0 is unavailable")
3985 } else {
3986 &self.mtp_extra[index - 1]
3987 }
3988 }
3989
3990 fn new_mtp_scratch(
3991 &self,
3992 e: &Engine,
3993 cap: usize,
3994 ) -> Result<MtpScratch, Box<dyn std::error::Error>> {
3995 let mut scratch = MtpScratch::new(
3996 e,
3997 &self.cfg,
3998 &self.plan,
3999 cap,
4000 self.mtp.as_ref().and_then(|head| head.geom.as_ref()),
4001 )?;
4002 for head in &self.mtp_extra {
4003 scratch.push_plane(e, &self.cfg, &self.plan, head.geom.as_ref())?;
4004 }
4005 Ok(scratch)
4006 }
4007
4008 fn opti_graph_draft_step(
4009 &self,
4010 e: &Engine,
4011 mtp: &MtpHead,
4012 dctx: &mut DraftGraphCtx,
4013 scratch: &mut MtpScratch,
4014 d_vocab: usize,
4015 ) -> Result<(u32, f32), Box<dyn std::error::Error>> {
4016 // dcw door: one replay appends one device-counter row; pre-arm ring headroom
4017 // host-side before launching (no-op on flat planes).
4018 if step35_draft_dcw_on() {
4019 scratch.ensure_dcw_headroom(e, 2)?;
4020 }
4021 dctx.graph
4022 .as_ref()
4023 .ok_or("optipipe controller requires the greedy draft graph")?
4024 .launch()?;
4025 scratch.kv.len += 1;
4026 let idx = e.dtoh_u32_one(&dctx.g_tok)?;
4027 if (idx as usize) >= d_vocab {
4028 return Err(
4029 format!("optipipe draft argmax sentinel 0x{idx:08x} >= d_vocab {d_vocab}").into(),
4030 );
4031 }
4032 let probability = e.dtoh(&dctx.g_p)?[0];
4033 if !probability.is_finite() || !(0.0..=1.0).contains(&probability) {
4034 return Err(format!("optipipe draft probability is invalid: {probability}").into());
4035 }
4036 let token = match &mtp.d2t {
4037 Some(map) => map[idx as usize],
4038 None => idx,
4039 };
4040 if token != idx {
4041 e.set_u32_one(&mut dctx.g_tok, token)?;
4042 }
4043 Ok((token, probability))
4044 }
4045
4046 #[allow(clippy::too_many_arguments)]
4047 fn opti_controller_draft_step(
4048 &self,
4049 e: &Engine,
4050 mtp: &MtpHead,
4051 dctx: &mut DraftGraphCtx,
4052 scratch: &mut MtpScratch,
4053 d_vocab: usize,
4054 eager_state: &mut Option<(u32, CudaSlice<f32>)>,
4055 eager_pos: usize,
4056 embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
4057 ) -> Result<(u32, f32), Box<dyn std::error::Error>> {
4058 if dctx.graph.is_some() {
4059 return self.opti_graph_draft_step(e, mtp, dctx, scratch, d_vocab);
4060 }
4061 let (input_token, input_seed) = eager_state
4062 .take()
4063 .ok_or("optipipe eager continuation seed is unavailable")?;
4064 let (logits, next_seed) = self.mtp_head_forward_dev(
4065 e,
4066 mtp,
4067 input_token,
4068 &input_seed,
4069 scratch,
4070 eager_pos,
4071 embd_dev,
4072 None,
4073 )?;
4074 let token_d = e.argmax_token_device(&logits, d_vocab)?;
4075 let idx = e.dtoh_u32_one(&token_d)?;
4076 if (idx as usize) >= d_vocab {
4077 return Err(format!(
4078 "optipipe eager draft argmax sentinel 0x{idx:08x} >= d_vocab {d_vocab}"
4079 )
4080 .into());
4081 }
4082 let probability_d = e.prob_of_token_device(&logits, &token_d, d_vocab)?;
4083 let probability = e.dtoh(&probability_d)?[0];
4084 if !probability.is_finite() || !(0.0..=1.0).contains(&probability) {
4085 return Err(
4086 format!("optipipe eager draft probability is invalid: {probability}").into(),
4087 );
4088 }
4089 let token = match &mtp.d2t {
4090 Some(map) => map[idx as usize],
4091 None => idx,
4092 };
4093 *eager_state = Some((token, next_seed));
4094 Ok((token, probability))
4095 }
4096
4097 /// NextN head forward for ONE draft token (§A ops 1-13, T=1).
4098 /// Inputs: `e_tok` = the token to predict FROM (last committed / previous draft); `h_seed` =
4099 /// the trunk's pre-output_norm hidden of that token (§A op 2 input). `mtp_pos` = absolute
4100 /// position of the token being predicted from. Returns (draft_logits[n_vocab] host, h_nextn dev).
4101 /// `h_nextn` (§A op 10) becomes `h_seed` for the next autoregressive draft step.
4102 /// Device-resident: returns draft logits ON DEVICE (no [n_vocab] dtoh). The greedy draft
4103 /// loop only needs argmax — paired with `argmax_token_device` this cuts the ~600KB logits
4104 /// transfer + host argmax per draft token from the K-token draft chain.
4105 #[allow(clippy::too_many_arguments)]
4106 fn mtp_head_forward_dev(
4107 &self,
4108 e: &Engine,
4109 mtp: &MtpHead,
4110 e_tok: u32,
4111 h_seed: &CudaSlice<f32>,
4112 scratch: &mut MtpScratch,
4113 mtp_pos: usize,
4114 embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
4115 mask: Option<(&CudaSlice<u32>, usize)>,
4116 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4117 self.mtp_head_forward_dev_at(e, mtp, e_tok, h_seed, scratch, 0, mtp_pos, embd_dev, mask)
4118 }
4119
4120 #[allow(clippy::too_many_arguments)]
4121 fn mtp_head_forward_dev_at(
4122 &self,
4123 e: &Engine,
4124 mtp: &MtpHead,
4125 e_tok: u32,
4126 h_seed: &CudaSlice<f32>,
4127 scratch: &mut MtpScratch,
4128 scratch_index: usize,
4129 mtp_pos: usize,
4130 embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
4131 // DRAFT-SIDE GRAMMAR MASK (lane/draft-mask): (packed draft-vocab allowed set, words).
4132 // Applied to the head logits BEFORE they are returned, so every consumer (argmax,
4133 // gumbel draw, p-min prob) sees the grammar-legal row. None = unmasked (pre-lane).
4134 mask: Option<(&CudaSlice<u32>, usize)>,
4135 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4136 // MEMRA_SPEC_ANATOMY=1 — eager-step phase timers (diagnostic only). Phase boundaries
4137 // sync the stream, so absolute time inflates; the BREAKDOWN is the signal. Cumulative
4138 // summary on stderr every 128 steps: glue (embed..attn_norm), attn, ffn, head.
4139 use std::sync::atomic::{AtomicU64, Ordering::Relaxed};
4140 static ANAT_NS: [AtomicU64; 5] = [
4141 AtomicU64::new(0),
4142 AtomicU64::new(0),
4143 AtomicU64::new(0),
4144 AtomicU64::new(0),
4145 AtomicU64::new(0),
4146 ];
4147 static ANAT_STEPS: AtomicU64 = AtomicU64::new(0);
4148 let anat = {
4149 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
4150 *ON.get_or_init(|| std::env::var("MEMRA_SPEC_ANATOMY").as_deref() == Ok("1"))
4151 };
4152 if anat {
4153 e.stream().synchronize()?; // drain prior queue so phase 0 starts clean
4154 }
4155 let t_all = std::time::Instant::now();
4156 let mut t_ph = std::time::Instant::now();
4157 let mut anat_mark = |i: usize,
4158 e: &Engine,
4159 t: &mut std::time::Instant|
4160 -> Result<(), Box<dyn std::error::Error>> {
4161 if anat {
4162 e.stream().synchronize()?;
4163 ANAT_NS[i].fetch_add(t.elapsed().as_nanos() as u64, Relaxed);
4164 *t = std::time::Instant::now();
4165 }
4166 Ok(())
4167 };
4168 let cfg = &self.cfg;
4169 let n_embd = cfg.n_embd as usize;
4170 // Distilled-student geometry: the block runs at the INNER width `di` (eh_proj out /
4171 // attn / ffn); the n_embd interface (embed, norms in, carrier out, head in) is unchanged.
4172 let di = mtp.geom.as_ref().map(|g| g.d_inner).unwrap_or(n_embd);
4173 let eps = cfg.rms_eps;
4174 let pos_d = e.htod_i32(&[mtp_pos as i32])?;
4175
4176 // op A: a resident table transfers one 4B token id. The exact host-row capacity path
4177 // expands this one row on CPU and transfers n_embd f32 values instead.
4178 let e_emb = match embd_dev {
4179 Some((g, qt, rb)) => e.embed_gather_device_t(g, &[e_tok], n_embd, qt, rb)?,
4180 None => e.htod(&self.embd.gather(n_embd, &[e_tok]))?,
4181 };
4182
4183 // op 1/2: e_norm = RMSNorm(e, enorm); h_norm = RMSNorm(h_seed, hnorm)
4184 let mut e_norm = e.zeros(n_embd)?;
4185 e.rms_norm(&e_emb, mtp.enorm.float_data(), &mut e_norm, n_embd, 1, eps)?;
4186 let mut h_norm = e.zeros(n_embd)?;
4187 e.rms_norm(h_seed, mtp.hnorm.float_data(), &mut h_norm, n_embd, 1, eps)?;
4188
4189 // op 3: concat = [e_norm ; h_norm] -> [2*n_embd], e_norm in [0,n_embd), h_norm in [n_embd,2n_embd)
4190 let mut concat = e.zeros(2 * n_embd)?;
4191 e.copy_into(&mut concat, 0, &e_norm, n_embd)?;
4192 e.copy_into(&mut concat, n_embd, &h_norm, n_embd)?;
4193
4194 // op 4: inpSA = eh_proj @ concat (eh_proj [2*n_embd, n_embd]) -> [n_embd]
4195 let inp_sa = e.matmul(&mtp.eh_proj, &concat, 1)?;
4196
4197 // op 5: a_norm = RMSNorm(inpSA, attn_norm)
4198 let mut a_norm = e.zeros(di)?;
4199 e.rms_norm(&inp_sa, mtp.attn_norm.float_data(), &mut a_norm, di, 1, eps)?;
4200 anat_mark(0, e, &mut t_ph)?;
4201
4202 // op 6: attention on the scratch KV. SAME dc launcher as the graph path (bucket_max =
4203 // scratch.cap, length from the device len_d) so eager drafts match graph drafts
4204 // bit-for-bit at any t_kv (the parity gate). Host len mirrored here (the dc append
4205 // advances only the device counter).
4206 let attn_out = match (&mtp.mixer, mtp.step35.as_ref()) {
4207 // step35 MTP block, dcw door armed: the SAME windowed device-counter launcher as
4208 // the captured chain (draft parity by construction). Per-step ring headroom runs
4209 // HERE (eager is host-len work, a rebase is legal); host len mirrored like the
4210 // plain dc arm below.
4211 (Mixer::Full(fa), Some(g)) if self.step35_dcw_eligible(g) => {
4212 {
4213 let (kv, _) = scratch.plane_mut(scratch_index);
4214 let retain = match kv.ring.as_ref() {
4215 Some(ring) => memra_kv::swa_retain_from(kv.len, ring.window(), ring.base()),
4216 None => 0,
4217 };
4218 e.prepare_kv_append(kv, retain, 1)?;
4219 }
4220 let out =
4221 self.mtp_step35_attn_dcw(e, fa, g, &a_norm, &pos_d, scratch, scratch_index)?;
4222 scratch.plane_mut(scratch_index).0.len += 1;
4223 out
4224 }
4225 // step35 MTP block, door off (the shipping default until receipts): PER-LAYER
4226 // geometry + a separate head-wise gate + an SWA window, none of which the plain dc
4227 // launcher can express (see `mtp_step35_attn`). Host-len arm. Advances BOTH the
4228 // host len and the device counter itself (unlike the dc arm, whose host-side
4229 // mirror the caller does).
4230 (Mixer::Full(fa), Some(g)) => {
4231 self.mtp_step35_attn(e, fa, g, &a_norm, &pos_d, scratch, scratch_index)?
4232 }
4233 (Mixer::Full(fa), None) => {
4234 let out = self.mtp_full_attn_dc(
4235 e,
4236 fa,
4237 &a_norm,
4238 &pos_d,
4239 scratch,
4240 scratch_index,
4241 mtp.geom.as_ref(),
4242 )?;
4243 scratch.plane_mut(scratch_index).0.len += 1;
4244 out
4245 }
4246 (Mixer::Linear(_), _) => {
4247 panic!("MTP block is full-attn in qwen35; linear MTP not supported")
4248 }
4249 (Mixer::Mla(_), _) => crate::hybrid::mla_forward_unimplemented(),
4250 };
4251 anat_mark(1, e, &mut t_ph)?;
4252
4253 // op 7: x1 = inpSA + attn_out
4254 let mut x1 = e.zeros(di)?;
4255 e.add(&inp_sa, &attn_out, &mut x1, di)?;
4256
4257 // op 8: z = RMSNorm(x1, post_attn_norm) (pre-FFN norm)
4258 let mut z = e.zeros(di)?;
4259 e.rms_norm(&x1, mtp.post_attn_norm.float_data(), &mut z, di, 1, eps)?;
4260
4261 // op 9: FFN (Dense or MoE) — same as the trunk decode FFN
4262 let ffn_out = match &mtp.ffn {
4263 crate::hybrid::Ffn::Dense {
4264 ffn_gate,
4265 ffn_up,
4266 ffn_down,
4267 } => {
4268 let n_ff = ffn_gate.out_features();
4269 let (gate, up) = if e.uses_q8_1_fast(ffn_gate) && e.uses_q8_1_fast(ffn_up) {
4270 let (zq, zd) = e.quantize_q8_1(&z, 1, di)?;
4271 (
4272 e.matmul_pre(ffn_gate, &zq, &zd, &z, 1)?,
4273 e.matmul_pre(ffn_up, &zq, &zd, &z, 1)?,
4274 )
4275 } else {
4276 (e.matmul(ffn_gate, &z, 1)?, e.matmul(ffn_up, &z, 1)?)
4277 };
4278 let mut act = e.zeros(n_ff)?;
4279 // step35: a DENSE FFN reads the per-layer SHEXP clamp (upstream's one `build_ffn`
4280 // serves the dense MLP and the shared expert off `swiglu_clamp_shexp` —
4281 // llama-graph.cpp:1751), resolved for the MTP block's OWN index. Every other arch
4282 // passes None, which is `ffn_act`'s dispatch verbatim.
4283 Self::ffn_act_lim(
4284 e,
4285 &self.cfg,
4286 &gate,
4287 &up,
4288 1.0,
4289 1.0,
4290 mtp.step35.as_ref().and_then(|s| s.clamp_shexp),
4291 &mut act,
4292 n_ff,
4293 )?;
4294 e.matmul(ffn_down, &act, 1)?
4295 }
4296 // MTP head is a distinct block — key its experts under a separate layer index (u16::MAX)
4297 // so they never alias trunk layer 0's cache keys.
4298 crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il(e, m, &z, 1, u16::MAX)?,
4299 };
4300 anat_mark(2, e, &mut t_ph)?;
4301
4302 // op 10: h_nextn = x1 + ffn_out (at di)
4303 let mut h_inner = e.zeros(di)?;
4304 e.add(&x1, &ffn_out, &mut h_inner, di)?;
4305
4306 // op 10.5 (student): up-project the inner hidden back to n_embd — training semantics:
4307 // the chain carrier AND the head input are out_up(h_inner) (pre-final-norm).
4308 let h_nextn = match mtp.geom.as_ref() {
4309 Some(g) => e.matmul(&g.out_up, &h_inner, 1)?,
4310 None => h_inner,
4311 };
4312
4313 // op 11: final = RMSNorm(h_nextn, shared_head_norm OR output_norm)
4314 let final_norm = mtp.shared_head_norm.as_ref().unwrap_or(&self.output_norm);
4315 let mut final_h = e.zeros(n_embd)?;
4316 e.rms_norm(
4317 &h_nextn,
4318 final_norm.float_data(),
4319 &mut final_h,
4320 n_embd,
4321 1,
4322 eps,
4323 )?;
4324
4325 // op 12: draft_logits = (shared_head_head OR output) @ final — stays ON DEVICE.
4326 let head = mtp.shared_head_head.as_ref().unwrap_or(&self.output);
4327 let mut logits = e.matmul(head, &final_h, 1)?;
4328 // op 12b (lane/draft-mask): grammar mask over the DRAFT vocab, applied here so the
4329 // caller's argmax / gumbel draw / p-min prob all read the grammar-legal row.
4330 if let Some((mask_d, mw)) = mask {
4331 let d_vocab = head.out_features();
4332 e.mask_logits_col(&mut logits, mask_d, 0, d_vocab, mw)?;
4333 }
4334 anat_mark(3, e, &mut t_ph)?;
4335 if anat {
4336 ANAT_NS[4].fetch_add(t_all.elapsed().as_nanos() as u64, Relaxed);
4337 let n = ANAT_STEPS.fetch_add(1, Relaxed) + 1;
4338 if n % 128 == 0 {
4339 let us = |i: usize| ANAT_NS[i].load(Relaxed) / n / 1000;
4340 eprintln!(
4341 "[spec-anatomy] steps={n} avg us/step: glue={} attn={} ffn={} head={} total={}",
4342 us(0),
4343 us(1),
4344 us(2),
4345 us(3),
4346 us(4)
4347 );
4348 }
4349 }
4350 // Chain recurrence hand-over: pre-norm h_nextn (default) or post-norm final_h
4351 // (MEMRA_SPEC_HPOST — llama.cpp #24025's t_h_nextn is taken AFTER the head norm).
4352 Ok((logits, if spec_hpost() { final_h } else { h_nextn }))
4353 }
4354
4355 #[allow(clippy::too_many_arguments)]
4356 fn mtp_chain_forward_dev(
4357 &self,
4358 e: &Engine,
4359 tokens: &[u32],
4360 seeds: &[CudaSlice<f32>],
4361 scratch: &mut MtpScratch,
4362 committed_scratch_len: usize,
4363 embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
4364 mask: Option<(&CudaSlice<u32>, usize)>,
4365 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4366 if tokens.is_empty() || tokens.len() != seeds.len() {
4367 return Err("multi-head MTP prefix tokens/seeds are malformed".into());
4368 }
4369 let index = mtp_chain_head_index(tokens.len() - 1, self.mtp_head_count());
4370 let head = self.mtp_head_at(index);
4371 scratch.set_plane_len(e, index, committed_scratch_len)?;
4372
4373 let mut last = None;
4374 for row in 0..tokens.len() {
4375 let is_last = row + 1 == tokens.len();
4376 last = Some(self.mtp_head_forward_dev_at(
4377 e,
4378 head,
4379 tokens[row],
4380 &seeds[row],
4381 scratch,
4382 index,
4383 committed_scratch_len + row + 1,
4384 embd_dev,
4385 if is_last { mask } else { None },
4386 )?);
4387 }
4388 Ok(last.expect("non-empty MTP prefix produced no row"))
4389 }
4390
4391 /// step35 MTP-block attention, T=1, on the scratch KV — the EAGER-ONLY twin of
4392 /// `mtp_full_attn_dc`. Three things force a separate arm rather than a geometry parameter on
4393 /// the dc path, and all three are properties of this arch's MTP block:
4394 ///
4395 /// 1. **The SWA window.** Block 45 is an SWA-type block (`sliding_window_pattern[45]=true`,
4396 /// window 512). Windowed decode in memra is a token-aligned VIEW OFFSET into the quantized
4397 /// cache (the gemma4 R6 / `step35_decode_attn` pattern: keys carry absolute rope and the
4398 /// mask is purely positional, so one query at `len-1` attending the last `win` rows IS the
4399 /// windowed result). `fa_decode_dc` takes the key count from a DEVICE counter and always
4400 /// starts at row 0 — it cannot express a nonzero offset. The windowed dc arm is
4401 /// `mtp_step35_attn_dcw` (`fa_decode_dcw`, flag-doored via MEMRA_STEP35_DRAFT_DCW);
4402 /// this host-len arm remains the default until the door's receipts land.
4403 /// 2. **Per-layer head count.** 96 q heads over 8 KV (GQA 12) at this block, vs the trunk's 64
4404 /// on its full-attn layers. The trunk cfg's `n_head` scalar is the MAX over layers, and the
4405 /// trunk ARTIFACT's per-layer arrays stop at index 44 — so the count must come from the
4406 /// resolved `Step35MtpGeom`, never from `cfg`.
4407 /// 3. **The separate head-wise gate.** `blk.45.attn_gate.weight [n_embd, 96]` produces one
4408 /// sigmoid scalar per head (broadcast over head_dim) — `attn_head_gate`, not the qwen35
4409 /// fused-into-wq `q_gate_split` form the dc arm handles.
4410 ///
4411 /// DOOR STATE: with MEMRA_STEP35_DRAFT_DCW unset, `mtp_head_forward_cap` refuses step35
4412 /// heads explicitly (rather than silently capturing a window-less, wrong-past-`win` graph)
4413 /// and this eager chain IS the served path. With the door armed, BOTH draft modes run the
4414 /// `mtp_step35_attn_dcw` twin instead of this arm.
4415 ///
4416 /// Unlike the dc arm this advances BOTH the host `kv.len` and the device counter, so the
4417 /// caller must not mirror.
4418 fn mtp_step35_attn(
4419 &self,
4420 e: &Engine,
4421 fa: &FullAttnLayer,
4422 g: &crate::hybrid::Step35MtpGeom,
4423 h: &CudaSlice<f32>,
4424 pos_d: &CudaSlice<i32>,
4425 scratch: &mut MtpScratch,
4426 scratch_index: usize,
4427 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4428 let (nh, nkv, hd) = (g.n_head, g.n_head_kv, self.cfg.head_dim_k as usize);
4429 // MTP-GEOM RECEIPT, once per process, on the SERVED draft path. Slot-0 acceptance is
4430 // 0.725 here against 0.994 for vLLM MTP3 on the same checkpoint family and card class, and
4431 // the first three explanations for that gap were all wrong: head assignment (step-modulo
4432 // is index 0 at K=1, correct), MEMRA_SPEC_HPOST (identical 84/116 both arms), and this
4433 // block's geometry. Geometry was the one that could have failed SILENTLY — a wrong window
4434 // makes the draft attend the whole context instead of Step-3.7's 512, stays fluent, and
4435 // shows up only as acceptance — so it gets a standing receipt rather than another reading
4436 // of the source. Prints the resolved Step35MtpGeom the served path actually runs on;
4437 // `full_attention_geometry_at`'s missing-row fallback (window: None) does NOT reach here.
4438 {
4439 static ONCE: std::sync::OnceLock<()> = std::sync::OnceLock::new();
4440 ONCE.get_or_init(|| {
4441 eprintln!(
4442 "[mtp-geom] arm=eager block={} swa={} window={} n_head={nh} n_head_kv={nkv} \
4443 head_dim_k={hd} n_rot={} rope_base={} clamp_shexp={:?}",
4444 g.il, g.swa, g.window, g.n_rot, g.rope_base, g.clamp_shexp,
4445 );
4446 });
4447 }
4448 let eps = self.cfg.rms_eps;
4449 let scale = 1.0 / (hd as f32).sqrt(); // step35.cpp:255 kq_scale
4450 let n_embd = self.cfg.n_embd as usize;
4451 let gw = fa
4452 .attn_gate
4453 .as_ref()
4454 .ok_or("step35 MTP block is missing attn_gate.weight (head-wise attention gate)")?;
4455
4456 let (q0, k0, v0, gt) = if e.uses_q8_1_fast(&fa.wq)
4457 && e.uses_q8_1_fast(&fa.wk)
4458 && e.uses_q8_1_fast(&fa.wv)
4459 && e.uses_q8_1_fast(gw)
4460 {
4461 let (hq, hdq) = e.quantize_q8_1(h, 1, n_embd)?;
4462 let (a, b, c) = match e.matmul_q8_fused3(&fa.wq, &fa.wk, &fa.wv, &hq, &hdq)? {
4463 Some(t3) => t3,
4464 None => (
4465 e.matmul_pre(&fa.wq, &hq, &hdq, h, 1)?,
4466 e.matmul_pre(&fa.wk, &hq, &hdq, h, 1)?,
4467 e.matmul_pre(&fa.wv, &hq, &hdq, h, 1)?,
4468 ),
4469 };
4470 (a, b, c, e.matmul_pre(gw, &hq, &hdq, h, 1)?)
4471 } else {
4472 (
4473 e.matmul(&fa.wq, h, 1)?,
4474 e.matmul(&fa.wk, h, 1)?,
4475 e.matmul(&fa.wv, h, 1)?,
4476 e.matmul(gw, h, 1)?,
4477 )
4478 };
4479
4480 let mut q = e.uninit(nh * hd)?;
4481 e.rms_norm(&q0, fa.q_norm.float_data(), &mut q, hd, nh, eps)?;
4482 let mut k = e.uninit(nkv * hd)?;
4483 e.rms_norm(&k0, fa.k_norm.float_data(), &mut k, hd, nkv, eps)?;
4484 // `rope_freqs.weight` (llama3 factors) applies to the FULL-attn layers ONLY; SWA passes
4485 // null (llama-hparams / step35.cpp). Block 45 is SWA, so `ff` is None there — but read
4486 // the resolved flag, not the constant, so an all-full sibling stays correct.
4487 let ff = if g.swa {
4488 None
4489 } else {
4490 self.step35_aux.as_ref().and_then(|a| a.rope_freqs(e))
4491 };
4492 #[cfg(debug_assertions)]
4493 if let Some(ff) = ff {
4494 crate::debug_assert_tensor_stream_device(ff, &e.stream(), "mtp_step35_attn.rope_freqs");
4495 }
4496 e.rope_neox2(
4497 &mut q,
4498 &mut k,
4499 pos_d,
4500 hd,
4501 g.n_rot,
4502 nh,
4503 nkv,
4504 1,
4505 g.rope_base,
4506 1.0,
4507 ff,
4508 )?;
4509
4510 // Append at the HOST slot, then re-stamp the device counter: the eager chain has the
4511 // length on the host anyway, and the windowed view below needs it there to compute the
4512 // offset. The device counter is kept in lockstep so `mtp_kv_fill`'s `set_i32_one` and any
4513 // dc-family consumer of this scratch still agree.
4514 let (kv, scratch_cap) = scratch.plane_mut(scratch_index);
4515 assert!(
4516 kv.len < scratch_cap,
4517 "step35 MTP scratch overflow ({} >= {})",
4518 kv.len,
4519 scratch_cap
4520 );
4521 let next_len = kv.len + 1;
4522 let (off, t_kv) = if g.swa && next_len > g.window {
4523 (next_len - g.window, g.window)
4524 } else {
4525 (0, next_len)
4526 };
4527 // `off`/`t_kv` stay the ATTENTION view; the retain is a separate, lower bound so the
4528 // rewind that follows this append is still resident. THIS is the only site that rebases
4529 // this plane (MEMRA_KV_REBASE_TRACE, one run: 1 rebase, all from here), so it is the site
4530 // that decides `base` for everyone.
4531 let retain_from = match kv.ring.as_ref() {
4532 Some(ring) => memra_kv::swa_retain_from(kv.len, ring.window(), ring.base()),
4533 None => off & !31usize,
4534 };
4535 let write_row = e.prepare_kv_append(kv, retain_from, 1)?;
4536 e.append_kv_quantized(
4537 &k,
4538 &v0,
4539 &mut kv.k,
4540 &mut kv.v,
4541 write_row,
4542 kv.kv_dim_k,
4543 kv.kv_dim_v,
4544 kv.k_tok_bytes,
4545 kv.v_tok_bytes,
4546 false,
4547 )?;
4548 kv.len = next_len;
4549 e.set_i32_one(&mut kv.len_d, kv.len as i32)?;
4550 // SWA view offset (see note 1). The draft chain is short (k+2 rows), but the scratch is
4551 // PERSISTENT across rounds — `mtp_kv_fill` leaves one row per committed token behind, so
4552 // `kv.len` tracks absolute position and crosses 512 in any real generation. The window is
4553 // therefore live, not theoretical.
4554 let physical = kv.physical_rows(off, off + t_kv)?;
4555 let k_view = e.view_u8_range(
4556 &kv.k,
4557 physical.start * kv.k_tok_bytes,
4558 physical.end * kv.k_tok_bytes,
4559 );
4560 let v_view = e.view_u8_range(
4561 &kv.v,
4562 physical.start * kv.v_tok_bytes,
4563 physical.end * kv.v_tok_bytes,
4564 );
4565 let mut attn = e.uninit(nh * hd)?;
4566 e.fa_decode_kvmod(
4567 &q,
4568 &k_view,
4569 &v_view,
4570 &mut attn,
4571 hd,
4572 nh,
4573 nkv,
4574 t_kv,
4575 scale,
4576 kv.k_tok_bytes,
4577 kv.v_tok_bytes,
4578 false,
4579 )?;
4580
4581 let mut ag = e.uninit(nh * hd)?;
4582 e.attn_head_gate(&attn, >, &mut ag, None, hd, nh, 1)?;
4583 Ok(e.matmul(&fa.wo, &ag, 1)?)
4584 }
4585
4586 /// The dcw draft arm's kernel-class precondition, mirrored from `fa_decode_dcw`'s own
4587 /// refusal plus the v3 walk's format contract (`fa_v3_active`), so the DEV dispatch can
4588 /// never pick an arm the launcher would refuse mid-chain (the eager chain has no graceful
4589 /// fallback point) and the CAP site refuses with the named reason instead.
4590 fn step35_dcw_eligible(&self, g: &crate::hybrid::Step35MtpGeom) -> bool {
4591 let hd = self.cfg.head_dim_k as usize;
4592 step35_draft_dcw_on()
4593 && g.swa
4594 && g.window >= crate::fa_vec_min_tkv()
4595 && std::env::var("MEMRA_NO_FA_VEC").is_err()
4596 && crate::fa_v3_active(hd)
4597 && hd <= 256
4598 }
4599
4600 /// step35 MTP-block attention, T=1, on the scratch KV: the WINDOWED DEVICE-COUNTER twin
4601 /// of `mtp_step35_attn`, serving BOTH draft paths when `step35_draft_dcw_on`. Write slot,
4602 /// key bound and SWA view offset all derive from device state (`len_d`, `base_d` written
4603 /// only at host-side rebases, and the block's `window`), so ONE captured graph serves the
4604 /// whole chain and replays see KV growth through the counter: the `mtp_full_attn_dc`
4605 /// contract plus the view offset the plain `_dc` kernel could not express (the old
4606 /// capture-refusal root cause). The three step35 properties stay per-geom exactly as in
4607 /// the eager twin: nh/nkv from `Step35MtpGeom`, the separate head-wise gate
4608 /// (`attn_head_gate`), per-layer rope width/base with SWA passing null freqs.
4609 ///
4610 /// bucket_max = min(cap, window): the windowed view never exceeds `window` rows, so the
4611 /// capture-time grid stays valid for every replayed len, and the kernel derives ns_eff
4612 /// from the LIVE T_kv at the fixed split_keys (one-partition law). Both arms call THIS
4613 /// launcher at THIS bucket, so eager and captured drafts are bit-identical by
4614 /// construction; vs the retired-by-flag `mtp_step35_attn` the only numeric-class deltas
4615 /// are the sub-vec-floor region (t_kv < 96: kvmod ran scalar, dcw stays vec) and any
4616 /// live-len split-ladder rung below the bucket's, both draft-side only (the verify
4617 /// arbitrates emitted bytes; acceptance is gated by the battery).
4618 ///
4619 /// Host len is NOT advanced here (graph contract); callers mirror. The EAGER caller runs
4620 /// `prepare_kv_append` per step (ring headroom, rebase legal there); the CAPTURED path
4621 /// pre-arms headroom at capture time and round start (`MtpScratch::ensure_dcw_headroom`)
4622 /// because a rebase is host work no captured chain may contain.
4623 fn mtp_step35_attn_dcw(
4624 &self,
4625 e: &Engine,
4626 fa: &FullAttnLayer,
4627 g: &crate::hybrid::Step35MtpGeom,
4628 h: &CudaSlice<f32>,
4629 pos_d: &CudaSlice<i32>,
4630 scratch: &mut MtpScratch,
4631 scratch_index: usize,
4632 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4633 let (nh, nkv, hd) = (g.n_head, g.n_head_kv, self.cfg.head_dim_k as usize);
4634 // MTP-GEOM RECEIPT (dcw twin of the `mtp_step35_attn` receipt): once per process,
4635 // naming the arm, so a serving log proves WHICH draft attention program ran (the
4636 // engagement receipt for the flag door, both directions).
4637 {
4638 static ONCE: std::sync::OnceLock<()> = std::sync::OnceLock::new();
4639 ONCE.get_or_init(|| {
4640 eprintln!(
4641 "[mtp-geom] arm=dcw block={} swa={} window={} n_head={nh} n_head_kv={nkv} \
4642 head_dim_k={hd} n_rot={} rope_base={} clamp_shexp={:?}",
4643 g.il, g.swa, g.window, g.n_rot, g.rope_base, g.clamp_shexp,
4644 );
4645 });
4646 }
4647 let eps = self.cfg.rms_eps;
4648 let scale = 1.0 / (hd as f32).sqrt(); // step35.cpp:255 kq_scale
4649 let n_embd = self.cfg.n_embd as usize;
4650 let gw = fa
4651 .attn_gate
4652 .as_ref()
4653 .ok_or("step35 MTP block is missing attn_gate.weight (head-wise attention gate)")?;
4654
4655 let (q0, k0, v0, gt) = if e.uses_q8_1_fast(&fa.wq)
4656 && e.uses_q8_1_fast(&fa.wk)
4657 && e.uses_q8_1_fast(&fa.wv)
4658 && e.uses_q8_1_fast(gw)
4659 {
4660 let (hq, hdq) = e.quantize_q8_1(h, 1, n_embd)?;
4661 let (a, b, c) = match e.matmul_q8_fused3(&fa.wq, &fa.wk, &fa.wv, &hq, &hdq)? {
4662 Some(t3) => t3,
4663 None => (
4664 e.matmul_pre(&fa.wq, &hq, &hdq, h, 1)?,
4665 e.matmul_pre(&fa.wk, &hq, &hdq, h, 1)?,
4666 e.matmul_pre(&fa.wv, &hq, &hdq, h, 1)?,
4667 ),
4668 };
4669 (a, b, c, e.matmul_pre(gw, &hq, &hdq, h, 1)?)
4670 } else {
4671 (
4672 e.matmul(&fa.wq, h, 1)?,
4673 e.matmul(&fa.wk, h, 1)?,
4674 e.matmul(&fa.wv, h, 1)?,
4675 e.matmul(gw, h, 1)?,
4676 )
4677 };
4678
4679 let mut q = e.zeros(nh * hd)?;
4680 e.rms_norm(&q0, fa.q_norm.float_data(), &mut q, hd, nh, eps)?;
4681 let mut k = e.zeros(nkv * hd)?;
4682 e.rms_norm(&k0, fa.k_norm.float_data(), &mut k, hd, nkv, eps)?;
4683 // rope_freqs (llama3 factors) apply to the FULL-attn layers ONLY; SWA passes null
4684 // (the eager twin's rule, resolved from the flag, not the constant).
4685 let ff = if g.swa {
4686 None
4687 } else {
4688 self.step35_aux.as_ref().and_then(|a| a.rope_freqs(e))
4689 };
4690 #[cfg(debug_assertions)]
4691 if let Some(ff) = ff {
4692 crate::debug_assert_tensor_stream_device(
4693 ff,
4694 &e.stream(),
4695 "mtp_step35_attn_dcw.rope_freqs",
4696 );
4697 }
4698 e.rope_neox2(
4699 &mut q,
4700 &mut k,
4701 pos_d,
4702 hd,
4703 g.n_rot,
4704 nh,
4705 nkv,
4706 1,
4707 g.rope_base,
4708 1.0,
4709 ff,
4710 )?;
4711
4712 let (kv, cap) = scratch.plane_mut(scratch_index);
4713 // Append at the DEVICE slot's PHYSICAL row (len_d - base_d), then advance the counter
4714 // in-graph. Physical room is the callers' headroom contract (see the fn doc).
4715 e.append_kv_quantized_dcw(
4716 &k,
4717 &v0,
4718 &mut kv.k,
4719 &mut kv.v,
4720 &kv.len_d,
4721 kv.base_d.as_ref(),
4722 kv.kv_dim_k,
4723 kv.kv_dim_v,
4724 kv.k_tok_bytes,
4725 kv.v_tok_bytes,
4726 )?;
4727 e.inc_seqlen(&mut kv.len_d)?;
4728 // Full-buffer views (any in-round physical row stays in range under the headroom
4729 // contract); the kernel bounds and offsets the key range from (len_d, base_d, window).
4730 let k_view = e.view_u8(&kv.k, kv.k.len());
4731 let v_view = e.view_u8(&kv.v, kv.v.len());
4732 let bucket = g.window.min(cap);
4733 let mut attn = e.zeros(nh * hd)?;
4734 e.fa_decode_dcw(
4735 &q,
4736 &k_view,
4737 &v_view,
4738 &mut attn,
4739 hd,
4740 nh,
4741 nkv,
4742 &kv.len_d,
4743 kv.base_d.as_ref(),
4744 if g.swa { g.window } else { 0 },
4745 bucket,
4746 scale,
4747 kv.k_tok_bytes,
4748 kv.v_tok_bytes,
4749 None,
4750 )?;
4751
4752 let mut ag = e.zeros(nh * hd)?;
4753 e.attn_head_gate(&attn, >, &mut ag, None, hd, nh, 1)?;
4754 Ok(e.matmul(&fa.wo, &ag, 1)?)
4755 }
4756
4757 /// MTP-block full attention, T=1, on the scratch KV (BOTH draft paths — eager and graph):
4758 /// the scratch write slot and the attention bound come from `scratch.kv.len_d` (device i32[1])
4759 /// so the launch args are FIXED across draft steps — ONE captured graph serves the whole
4760 /// chain, and replays keep seeing KV growth through the device counter (no recapture).
4761 /// Geometry contract: n_splits is sized from `scratch.cap` (the persistent capacity); splits
4762 /// whose key range lies beyond the device t_kv exit empty and the shared combine skips them
4763 /// (fa_decode_dc bit-correct-for-any-t_kv<=bucket_max contract). The eager path uses the SAME
4764 /// launcher with the SAME bucket_max -> identical dispatch -> bit-identical draft tokens (the
4765 /// graph-vs-eager parity gate). Host len is NOT advanced here (graph contract); callers mirror.
4766 fn mtp_full_attn_dc(
4767 &self,
4768 e: &Engine,
4769 fa: &FullAttnLayer,
4770 h: &CudaSlice<f32>,
4771 pos_d: &CudaSlice<i32>,
4772 scratch: &mut MtpScratch,
4773 scratch_index: usize,
4774 geom: Option<&crate::hybrid::DraftGeom>,
4775 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4776 let cfg = &self.cfg;
4777 let mtp_il = cfg.n_layer.saturating_sub(cfg.nextn_predict_layers);
4778 let geometry = cfg.full_attention_geometry_at(mtp_il);
4779 let n_head = geom.map(|g| g.n_head).unwrap_or(geometry.n_head as usize);
4780 let n_head_kv = geom
4781 .map(|g| g.n_head_kv)
4782 .unwrap_or(geometry.n_head_kv as usize);
4783 let head_dim = geometry.head_dim_k as usize;
4784 let eps = cfg.rms_eps;
4785 let scale = geometry.attention_scale();
4786 let n_embd = geom.map(|g| g.d_inner).unwrap_or(cfg.n_embd as usize);
4787 let bucket_max = scratch.plane(scratch_index).1;
4788
4789 let (qf, mut k, v) =
4790 if e.uses_q8_1_fast(&fa.wq) && e.uses_q8_1_fast(&fa.wk) && e.uses_q8_1_fast(&fa.wv) {
4791 let (hq, hd) = e.quantize_q8_1(h, 1, n_embd)?;
4792 (
4793 e.matmul_pre(&fa.wq, &hq, &hd, h, 1)?,
4794 e.matmul_pre(&fa.wk, &hq, &hd, h, 1)?,
4795 e.matmul_pre(&fa.wv, &hq, &hd, h, 1)?,
4796 )
4797 } else {
4798 (
4799 e.matmul(&fa.wq, h, 1)?,
4800 e.matmul(&fa.wk, h, 1)?,
4801 e.matmul(&fa.wv, h, 1)?,
4802 )
4803 };
4804 // M3/Hy3 have no attention output gate — wq out is exactly q; skip the split.
4805 let gated = geometry.attention_gate == memra_gguf::config::AttentionGateKind::FusedQ;
4806 let (mut q, gate) = if gated {
4807 let mut q = e.zeros(n_head * head_dim)?;
4808 let mut gate = e.zeros(n_head * head_dim)?;
4809 e.q_gate_split(&qf, &mut q, &mut gate, head_dim, n_head, 1)?;
4810 (q, Some(gate))
4811 } else {
4812 (qf, None)
4813 };
4814
4815 let mut qn = e.zeros(n_head * head_dim)?;
4816 e.rms_norm(&q, fa.q_norm.float_data(), &mut qn, head_dim, n_head, eps)?;
4817 q = qn;
4818 let mut kn = e.zeros(n_head_kv * head_dim)?;
4819 e.rms_norm(
4820 &k,
4821 fa.k_norm.float_data(),
4822 &mut kn,
4823 head_dim,
4824 n_head_kv,
4825 eps,
4826 )?;
4827 k = kn;
4828 let rope_dims = geometry.n_rot as usize;
4829 e.rope_neox(
4830 &mut q,
4831 pos_d,
4832 head_dim,
4833 rope_dims,
4834 n_head,
4835 1,
4836 geometry.rope_base,
4837 1.0,
4838 )?;
4839 e.rope_neox(
4840 &mut k,
4841 pos_d,
4842 head_dim,
4843 rope_dims,
4844 n_head_kv,
4845 1,
4846 geometry.rope_base,
4847 1.0,
4848 )?;
4849
4850 let kv = scratch.plane_mut(scratch_index).0;
4851 // append at the DEVICE slot (kv.len_d == old len), then advance the counter in-graph.
4852 e.append_kv_quantized_dc(
4853 &k,
4854 &v,
4855 &mut kv.k,
4856 &mut kv.v,
4857 &kv.len_d,
4858 kv.kv_dim_k,
4859 kv.kv_dim_v,
4860 kv.k_tok_bytes,
4861 kv.v_tok_bytes,
4862 false,
4863 )?;
4864 e.inc_seqlen(&mut kv.len_d)?;
4865 // full-buffer views (any in-round t_kv stays in range on replay); the kernel bounds the
4866 // key range from the device counter.
4867 let k_view = e.view_u8(&kv.k, kv.k.len());
4868 let v_view = e.view_u8(&kv.v, kv.v.len());
4869 let (ktb, vtb) = (kv.k_tok_bytes, kv.v_tok_bytes);
4870 let mut attn = e.zeros(n_head * head_dim)?;
4871 e.fa_decode_dc(
4872 &q, &k_view, &v_view, &mut attn, head_dim, n_head, n_head_kv, &kv.len_d, bucket_max,
4873 scale, ktb, vtb, false,
4874 )?;
4875
4876 let attn_g = match &gate {
4877 Some(gate) => {
4878 let mut gsig = e.zeros(n_head * head_dim)?;
4879 e.sigmoid(gate, &mut gsig, n_head * head_dim)?;
4880 let mut ag = e.zeros(n_head * head_dim)?;
4881 e.mul(&attn, &gsig, &mut ag, n_head * head_dim)?;
4882 ag
4883 }
4884 None => attn,
4885 };
4886 Ok(e.matmul(&fa.wo, &attn_g, 1)?)
4887 }
4888
4889 /// PERSISTENT-DRAFT-KV fill (the reference engine's "mtp_update" analogue): compute the MTP
4890 /// block's K/V for `tokens` (committed tokens at positions pos0..pos0+T) from their EXACT
4891 /// trunk hiddens `h` ([T, n_embd] token-major, pre-output_norm) and append at slots pos0.. of
4892 /// the scratch KV. K/V-ONLY — ops A/1-5 plus the K-side of op 6 (wk/wv + k_norm + rope +
4893 /// quantized append); no wq/attention/FFN/lm_head, so per-token cost ~= eh_proj + wk/wv (a
4894 /// small fraction of one trunk layer), T-batched. Rope follows the chain convention
4895 /// rope(token@p) = p+1. Runs at round boundaries OUTSIDE the captured graph in BOTH draft
4896 /// modes -> draft parity by construction. Caller must have scratch.kv.len == pos0.
4897 #[allow(clippy::too_many_arguments)]
4898 fn mtp_kv_fill_at(
4899 &self,
4900 e: &Engine,
4901 mtp: &MtpHead,
4902 tokens: &[u32],
4903 h: &CudaSlice<f32>,
4904 pos0: usize,
4905 scratch: &mut MtpScratch,
4906 scratch_index: usize,
4907 embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
4908 ) -> Result<(), Box<dyn std::error::Error>> {
4909 let cfg = &self.cfg;
4910 let n_embd = cfg.n_embd as usize;
4911 let eps = cfg.rms_eps;
4912 let t = tokens.len();
4913 let (scratch_kv, scratch_cap) = scratch.plane(scratch_index);
4914 assert_eq!(scratch_kv.len, pos0, "mtp_kv_fill: append slot mismatch");
4915 assert!(pos0 + t <= scratch_cap, "mtp_kv_fill: scratch overflow");
4916 let Mixer::Full(fa) = &mtp.mixer else {
4917 panic!("MTP block is full-attn in qwen35; linear MTP not supported")
4918 };
4919 let pos_vec: Vec<i32> = (0..t).map(|i| (pos0 + i + 1) as i32).collect();
4920 let pos_d = e.htod_i32(&pos_vec)?;
4921
4922 // ops A/1/2: embed + the two input norms, T-wide.
4923 let e_emb = match embd_dev {
4924 Some((g, qt, rb)) => e.embed_gather_device_t(g, tokens, n_embd, qt, rb)?,
4925 None => e.htod(&self.embd.gather(n_embd, tokens))?,
4926 };
4927 let mut e_norm = e.zeros(t * n_embd)?;
4928 e.rms_norm(&e_emb, mtp.enorm.float_data(), &mut e_norm, n_embd, t, eps)?;
4929 let mut h_norm = e.zeros(t * n_embd)?;
4930 e.rms_norm(h, mtp.hnorm.float_data(), &mut h_norm, n_embd, t, eps)?;
4931
4932 // op 3: per-row [e_norm ; h_norm] concat, token-major [T, 2*n_embd].
4933 let mut concat = e.zeros(t * 2 * n_embd)?;
4934 for i in 0..t {
4935 e.copy_view_into(
4936 &mut concat,
4937 i * 2 * n_embd,
4938 &e_norm.slice(i * n_embd..(i + 1) * n_embd),
4939 n_embd,
4940 )?;
4941 e.copy_view_into(
4942 &mut concat,
4943 i * 2 * n_embd + n_embd,
4944 &h_norm.slice(i * n_embd..(i + 1) * n_embd),
4945 n_embd,
4946 )?;
4947 }
4948
4949 // ops 4/5: eh_proj + attn_norm, T-wide (at the student inner width when geom is set).
4950 let di = mtp.geom.as_ref().map(|g| g.d_inner).unwrap_or(n_embd);
4951 let inp_sa = e.matmul(&mtp.eh_proj, &concat, t)?;
4952 let mut a_norm = e.zeros(t * di)?;
4953 e.rms_norm(&inp_sa, mtp.attn_norm.float_data(), &mut a_norm, di, t, eps)?;
4954
4955 // op 6 (K/V half): wk/wv + k_norm + rope + per-row quantized append. No wq/attention —
4956 // the fill only has to leave correct K/V rows behind for later chains to attend over.
4957 let n_head_kv = mtp
4958 .geom
4959 .as_ref()
4960 .map(|g| g.n_head_kv)
4961 .or(mtp.step35.as_ref().map(|s| s.n_head_kv))
4962 .unwrap_or_else(|| {
4963 let mtp_il = cfg.n_layer.saturating_sub(cfg.nextn_predict_layers);
4964 cfg.full_attention_geometry_at(mtp_il).n_head_kv as usize
4965 });
4966 let mtp_il = cfg.n_layer.saturating_sub(cfg.nextn_predict_layers);
4967 let geometry = cfg.full_attention_geometry_at(mtp_il);
4968 let head_dim = geometry.head_dim_k as usize;
4969 let mut k = e.matmul(&fa.wk, &a_norm, t)?;
4970 let v = e.matmul(&fa.wv, &a_norm, t)?;
4971 let mut kn = e.zeros(t * n_head_kv * head_dim)?;
4972 e.rms_norm(
4973 &k,
4974 fa.k_norm.float_data(),
4975 &mut kn,
4976 head_dim,
4977 n_head_kv * t,
4978 eps,
4979 )?;
4980 k = kn;
4981 // step35: rotary width AND base are per-layer, and the MTP block's values come from the
4982 // resolved `Step35MtpGeom` — NOT from `cfg.rope_dim_count`/`cfg.rope_freq_base`, which
4983 // carry the arch defaults (128 / 5e6, i.e. the FULL-attn layers' base). Getting this wrong
4984 // writes K rows the attention arm then re-derives at a different theta: correct-looking
4985 // output with dead acceptance, invisible to the exactness gates.
4986 let (rope_dims, rope_base, ff) = match mtp.step35.as_ref() {
4987 Some(s) => (
4988 s.n_rot,
4989 s.rope_base,
4990 if s.swa {
4991 None
4992 } else {
4993 self.step35_aux.as_ref().and_then(|a| a.rope_freqs(e))
4994 },
4995 ),
4996 None => (geometry.n_rot as usize, geometry.rope_base, None),
4997 };
4998 #[cfg(debug_assertions)]
4999 if let Some(ff) = ff {
5000 crate::debug_assert_tensor_stream_device(ff, &e.stream(), "mtp_kv_fill.rope_freqs");
5001 }
5002 match ff {
5003 Some(f) => e.rope_neox_ff(
5004 &mut k, &pos_d, head_dim, rope_dims, n_head_kv, t, rope_base, 1.0, f,
5005 )?,
5006 None => e.rope_neox(
5007 &mut k, &pos_d, head_dim, rope_dims, n_head_kv, t, rope_base, 1.0,
5008 )?,
5009 }
5010
5011 let kv = scratch.plane_mut(scratch_index).0;
5012 // Match the trunk prime contract: a chunk may need the aligned window immediately before
5013 // its first row, so preserve that prefix when the physical tail rebases at wrap.
5014 let retain_from = kv
5015 .ring
5016 .as_ref()
5017 .map(|ring| memra_kv::swa_retain_from(pos0, ring.window(), ring.base()))
5018 .unwrap_or(0);
5019 let write_row = e.prepare_kv_append(kv, retain_from, t)?;
5020 for i in 0..t {
5021 let k_row = k.slice(i * kv.kv_dim_k..(i + 1) * kv.kv_dim_k);
5022 let v_row = v.slice(i * kv.kv_dim_v..(i + 1) * kv.kv_dim_v);
5023 e.append_kv_quantized_view(
5024 &k_row,
5025 &v_row,
5026 &mut kv.k,
5027 &mut kv.v,
5028 write_row + i,
5029 kv.kv_dim_k,
5030 kv.kv_dim_v,
5031 kv.k_tok_bytes,
5032 kv.v_tok_bytes,
5033 false,
5034 )?;
5035 }
5036 kv.len = pos0 + t;
5037 e.set_i32_one(&mut kv.len_d, kv.len as i32)?;
5038 Ok(())
5039 }
5040
5041 #[allow(clippy::too_many_arguments)]
5042 fn mtp_kv_fill_all(
5043 &self,
5044 e: &Engine,
5045 tokens: &[u32],
5046 h: &CudaSlice<f32>,
5047 pos0: usize,
5048 scratch: &mut MtpScratch,
5049 embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
5050 ) -> Result<(), Box<dyn std::error::Error>> {
5051 debug_assert_eq!(self.mtp_head_count(), scratch.plane_count());
5052 for index in 0..self.mtp_head_count() {
5053 self.mtp_kv_fill_at(
5054 e,
5055 self.mtp_head_at(index),
5056 tokens,
5057 h,
5058 pos0,
5059 scratch,
5060 index,
5061 embd_dev,
5062 )?;
5063 }
5064 Ok(())
5065 }
5066
5067 /// CAPTURE body for the GRAPH DRAFT (stage 2 of graph-grade spec): ONE MTP head forward with
5068 /// every varying input device-resident —
5069 /// - token id from the persistent `tok_d` (the previous replay's in-graph argmax wrote it,
5070 /// so the chain feeds itself; the host reads the same 4 bytes for the draft list),
5071 /// - h_seed from the persistent `h_seed_d` (h_nextn is copied BACK into it at the end),
5072 /// - rope pos from the persistent `pos_d` counter (inc'd in-graph),
5073 /// - scratch KV slot/bound from `scratch.kv.len_d` (see mtp_full_attn_dc).
5074 /// The p-min confidence lands in the persistent `p_d` iff `with_prob` (env is fixed per run).
5075 /// Same kernels, same dispatch as the eager mtp_head_forward_dev chain -> same draft tokens
5076 /// (exactness never depends on drafts — the verify arbitrates — but acceptance parity does).
5077 /// `with_head=false` captures the HEAD-LESS twin for the pseudo-seed replay (2026-07-03):
5078 /// the pseudo pass only needs h_nextn (op 10) + the scratch append — the lm_head read
5079 /// (~1.06ms q6_K on the 9B), argmax and prob are dead weight there. h_nextn's inputs are
5080 /// untouched, so the seed value is identical; round-start resets overwrite tok_d/p_d anyway.
5081 /// `sampled_cap` = Some((ctr_d, perturb_d, q_out_d, seed, temp)) captures the SAMPLED twin
5082 /// (step 3 of the sampled-spec arc): head logits are retained in the persistent `q_out_d`
5083 /// (host D2Ds them to the round's q slot after each replay), the DEVICE event counter is
5084 /// bumped in-graph, and the argmax reads GUMBEL-PERTURBED logits — one categorical draw per
5085 /// replay, bit-identical to the eager arm's gumbel_perturb at the same (seed, sctr, temp).
5086 /// seed/temp are capture-time constants (fixed per generate call, like p_min).
5087 #[allow(clippy::too_many_arguments)]
5088 fn mtp_head_forward_cap(
5089 &self,
5090 e: &Engine,
5091 mtp: &MtpHead,
5092 tok_d: &mut CudaSlice<u32>,
5093 pos_d: &mut CudaSlice<i32>,
5094 h_seed_d: &mut CudaSlice<f32>,
5095 p_d: &mut CudaSlice<f32>,
5096 scratch: &mut MtpScratch,
5097 with_prob: bool,
5098 with_head: bool,
5099 embd_gpu: &CudaSlice<u8>,
5100 embd_qt: i32,
5101 embd_rb: usize,
5102 d_vocab: usize,
5103 sampled_cap: Option<(
5104 &mut CudaSlice<u32>,
5105 &mut CudaSlice<f32>,
5106 &mut CudaSlice<f32>,
5107 u64,
5108 f32,
5109 )>,
5110 stream_pack: Option<(&mut CudaSlice<u32>, usize, Option<&CudaSlice<u32>>)>,
5111 // DRAFT-SIDE GRAMMAR MASK (lane/draft-mask): (packed draft-vocab allowed-set buffer,
5112 // word count). Captured as ONE mask_logits_f32 node between the head matmul and the
5113 // in-graph argmax; the buffer address is baked, its CONTENTS are re-uploaded by the
5114 // host before every replay (the decode.rs graph-mask pattern). All-ones contents = a
5115 // no-op ban, so a position the grammar cannot constrain costs one pass over the row.
5116 mask_cap: Option<(&CudaSlice<u32>, usize)>,
5117 ) -> Result<(), Box<dyn std::error::Error>> {
5118 let cfg = &self.cfg;
5119 let n_embd = cfg.n_embd as usize;
5120 // step35: capturable through the WINDOWED device-counter arm (`mtp_step35_attn_dcw`)
5121 // once the dcw door is armed and the v3-vec class is live. Without the door this stays
5122 // the deliberate, named refusal: the plain `_dc` attention's key bound always starts at
5123 // row 0, cannot express this block's SWA view offset, and a captured chain would
5124 // silently attend OUTSIDE the window once the persistent scratch passes 512 rows.
5125 // Returning Err (not a panic) is what the capture sites already handle by degrading to
5126 // the eager chain (`mtp_head_forward_dev` -> `mtp_step35_attn`).
5127 // ROUND-STREAM stays refused EITHER WAY: the stream VERIFY has no step35 twin (see the
5128 // step35_verify refusal), so a stream capture that succeeded here would only move the
5129 // failure from capture time (graceful stream-off) to serve time (a failed round).
5130 if let Some(g) = mtp.step35.as_ref() {
5131 if stream_pack.is_some() {
5132 return Err(
5133 "step35 has no ROUND-STREAM draft arm (the stream verify has no step35 \
5134 twin); stream off"
5135 .into(),
5136 );
5137 }
5138 if !self.step35_dcw_eligible(g) {
5139 return Err(
5140 "step35 has no captured draft chain (fa_decode_dc cannot express the MTP \
5141 block's SWA view offset; MEMRA_STEP35_DRAFT_DCW=1 arms the windowed dcw \
5142 capture when the v3-vec class is live) - the eager draft chain serves \
5143 this arch"
5144 .into(),
5145 );
5146 }
5147 }
5148 // student inner width (see mtp_head_forward_dev) — interface dims stay n_embd.
5149 let di = mtp.geom.as_ref().map(|g| g.d_inner).unwrap_or(n_embd);
5150 let eps = cfg.rms_eps;
5151 let e_emb = e.embed_gather_device(embd_gpu, tok_d, n_embd, embd_qt, embd_rb)?;
5152 let mut e_norm = e.zeros(n_embd)?;
5153 e.rms_norm(&e_emb, mtp.enorm.float_data(), &mut e_norm, n_embd, 1, eps)?;
5154 let mut h_norm = e.zeros(n_embd)?;
5155 e.rms_norm(
5156 &*h_seed_d,
5157 mtp.hnorm.float_data(),
5158 &mut h_norm,
5159 n_embd,
5160 1,
5161 eps,
5162 )?;
5163 let mut concat = e.zeros(2 * n_embd)?;
5164 e.copy_into(&mut concat, 0, &e_norm, n_embd)?;
5165 e.copy_into(&mut concat, n_embd, &h_norm, n_embd)?;
5166 let inp_sa = e.matmul(&mtp.eh_proj, &concat, 1)?;
5167 let mut a_norm = e.zeros(di)?;
5168 e.rms_norm(&inp_sa, mtp.attn_norm.float_data(), &mut a_norm, di, 1, eps)?;
5169 let attn_out = match (&mtp.mixer, mtp.step35.as_ref()) {
5170 // step35 (eligibility already enforced by the refusal above): the windowed dcw
5171 // arm, the SAME launcher the eager dev arm runs when the door is armed. No host
5172 // work here (this is the capture body); headroom is the callers' pre-arm.
5173 (Mixer::Full(fa), Some(g)) => {
5174 self.mtp_step35_attn_dcw(e, fa, g, &a_norm, pos_d, scratch, 0)?
5175 }
5176 (Mixer::Full(fa), None) => {
5177 self.mtp_full_attn_dc(e, fa, &a_norm, pos_d, scratch, 0, mtp.geom.as_ref())?
5178 }
5179 (Mixer::Linear(_), _) => {
5180 panic!("MTP block is full-attn in qwen35; linear MTP not supported")
5181 }
5182 (Mixer::Mla(_), _) => crate::hybrid::mla_forward_unimplemented(),
5183 };
5184 let mut x1 = e.zeros(di)?;
5185 e.add(&inp_sa, &attn_out, &mut x1, di)?;
5186 let mut z = e.zeros(di)?;
5187 e.rms_norm(&x1, mtp.post_attn_norm.float_data(), &mut z, di, 1, eps)?;
5188 let ffn_out = match &mtp.ffn {
5189 crate::hybrid::Ffn::Dense {
5190 ffn_gate,
5191 ffn_up,
5192 ffn_down,
5193 } => {
5194 let n_ff = ffn_gate.out_features();
5195 let (gate, up) = if e.uses_q8_1_fast(ffn_gate) && e.uses_q8_1_fast(ffn_up) {
5196 let (zq, zd) = e.quantize_q8_1(&z, 1, di)?;
5197 (
5198 e.matmul_pre(ffn_gate, &zq, &zd, &z, 1)?,
5199 e.matmul_pre(ffn_up, &zq, &zd, &z, 1)?,
5200 )
5201 } else {
5202 (e.matmul(ffn_gate, &z, 1)?, e.matmul(ffn_up, &z, 1)?)
5203 };
5204 let mut act = e.zeros(n_ff)?;
5205 // step35: the dense FFN reads the per-layer SHEXP clamp, resolved for the MTP
5206 // block's own index (the mtp_head_forward_dev rule; None for every other arch,
5207 // which is `ffn_act`'s dispatch verbatim). The eager and captured chains must
5208 // run the ONE activation program.
5209 Self::ffn_act_lim(
5210 e,
5211 &self.cfg,
5212 &gate,
5213 &up,
5214 1.0,
5215 1.0,
5216 mtp.step35.as_ref().and_then(|s| s.clamp_shexp),
5217 &mut act,
5218 n_ff,
5219 )?;
5220 e.matmul(ffn_down, &act, 1)?
5221 }
5222 // ROUND-STREAM: the 35B NextN block carries a MoE FFN. With RESIDENT experts the
5223 // dev path is pure device launches (device top-k + rows kernels, ZERO-DtoH by
5224 // design) — capture-legal. Non-resident (SLRU-lock) stays rejected: the capture
5225 // error arm degrades the caller to eager/stream-off.
5226 crate::hybrid::Ffn::Moe(m) if m.dev_exps.is_some() => {
5227 self.moe_ffn_il(e, m, &z, 1, u16::MAX)?
5228 }
5229 crate::hybrid::Ffn::Moe(_) => {
5230 return Err("graph draft requires a Dense (or resident-MoE) MTP FFN".into());
5231 }
5232 };
5233 let mut h_inner = e.zeros(di)?;
5234 e.add(&x1, &ffn_out, &mut h_inner, di)?;
5235 // student: up-project back to n_embd (carrier + head input; see mtp_head_forward_dev).
5236 let h_nextn = match mtp.geom.as_ref() {
5237 Some(g) => e.matmul(&g.out_up, &h_inner, 1)?,
5238 None => h_inner,
5239 };
5240 // MEMRA_SPEC_HPOST needs final_h even head-less (it IS the next seed under that convention).
5241 let final_h = if with_head || spec_hpost() {
5242 let final_norm = mtp.shared_head_norm.as_ref().unwrap_or(&self.output_norm);
5243 let mut fh = e.zeros(n_embd)?;
5244 e.rms_norm(&h_nextn, final_norm.float_data(), &mut fh, n_embd, 1, eps)?;
5245 Some(fh)
5246 } else {
5247 None
5248 };
5249 if with_head {
5250 let head = mtp.shared_head_head.as_ref().unwrap_or(&self.output);
5251 let mut logits = e.matmul(head, final_h.as_ref().unwrap(), 1)?;
5252 // DRAFT-SIDE GRAMMAR MASK: ban the grammar-illegal draft ids IN the captured chain,
5253 // before the argmax — proposals become legal by construction. Contents-only
5254 // per-replay upload keeps the capture valid.
5255 if let Some((mask_d, mw)) = mask_cap {
5256 e.mask_logits_col(&mut logits, mask_d, 0, d_vocab, mw)?;
5257 }
5258 if let Some((ctr_d, perturb_d, q_out_d, seed, temp)) = sampled_cap {
5259 // SAMPLED chain: retain q (raw head logits -> persistent q_out_d; the matmul's
5260 // own buffer is pool-recycled after the capture body returns, so it can't be the
5261 // retention target), bump the device event counter, gumbel-perturb reading it,
5262 // and argmax the PERTURBED logits into tok_d — the in-graph categorical draw.
5263 e.copy_into(q_out_d, 0, &logits, d_vocab)?;
5264 e.sctr_inc(ctr_d)?;
5265 e.gumbel_perturb_ctr(&logits, perturb_d, d_vocab, seed, ctr_d, temp)?;
5266 e.argmax_token_device_into(perturb_d, tok_d, d_vocab)?;
5267 // p-min prob = the head's RAW softmax confidence in the SAMPLED pick — same
5268 // semantics as the eager sampled arm's prob_of_token_device(dl_d, tok_d).
5269 if with_prob {
5270 e.prob_of_token_device_into(&logits, tok_d, p_d, d_vocab)?;
5271 }
5272 } else {
5273 // draft token -> persistent tok_d (next replay's embed reads it; host reads the 4 bytes).
5274 e.argmax_token_device_into(&logits, tok_d, d_vocab)?;
5275 // p-min under a draft mask reads the MASKED row: confidence relative to the
5276 // grammar-LEGAL alternatives (illegal ids leave the softmax denominator), which
5277 // is the right semantics for "does the drafter know what comes next here" and
5278 // the same row the pick came from. Draft-quality only — verify arbitrates.
5279 if with_prob {
5280 e.prob_of_token_device_into(&logits, tok_d, p_d, d_vocab)?;
5281 }
5282 }
5283 }
5284 // ROUND-STREAM K-chain: pack (tok, p) into slot j, then remap tok through d2t so the
5285 // NEXT chained body's embed reads the TARGET id — zero host involvement per step.
5286 if let Some((out, slot, d2t)) = stream_pack {
5287 e.pack_tok_p(tok_d, p_d, out, slot)?;
5288 if let Some(map) = d2t {
5289 e.tok_map_u32(tok_d, map)?;
5290 }
5291 }
5292 // Next draft step's h_seed: pre-norm h_nextn (default) or post-norm final_h (HPOST).
5293 if spec_hpost() {
5294 e.copy_into(h_seed_d, 0, final_h.as_ref().unwrap(), n_embd)?;
5295 } else {
5296 e.copy_into(h_seed_d, 0, &h_nextn, n_embd)?;
5297 }
5298 // advance the draft rope position in-graph.
5299 e.inc_seqlen(pos_d)?;
5300 Ok(())
5301 }
5302
5303 /// Batched target verify forward over `tokens` at positions `pos0..pos0+T` (§D.3, T=K+1).
5304 /// Returns ALL T logit columns (host f32, [T*n_vocab]); appends T cols to every full-attn KV
5305 /// and advances every linear-attn recur state by T steps (the recur steps are SEQUENTIAL T=1).
5306 /// Advances `cache.pos` by T.
5307 pub fn decode_step_t(
5308 &self,
5309 e: &Engine,
5310 tokens: &[u32],
5311 pos0: usize,
5312 cache: &mut Cache,
5313 ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
5314 if self.is_gemma4_e4b() {
5315 return Ok(self.gemma4_e4b_decode_step_t_h(e, tokens, pos0, cache)?.0);
5316 }
5317 if self.gemma_batch_program() {
5318 return self.gemma4_decode_step_t(e, tokens, pos0, cache);
5319 }
5320 Ok(self.decode_step_t_h(e, tokens, pos0, cache)?.0)
5321 }
5322
5323 /// Like `decode_step_t` but ALSO returns the LAST column's pre-output_norm hidden (h_seed for
5324 /// the next draft round). This lets partial-accept replay run as ONE batched T=(n_acc+1) forward
5325 /// (single weight read) instead of n_acc+1 separate T=1 decode_steps (n_acc+1 weight reads).
5326 /// At batch=1 decode is bandwidth-bound, so batching the replay is THE MTP profitability lever.
5327 pub fn decode_step_t_h(
5328 &self,
5329 e: &Engine,
5330 tokens: &[u32],
5331 pos0: usize,
5332 cache: &mut Cache,
5333 ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
5334 self.decode_step_t_h_emb(e, tokens, pos0, cache, None)
5335 }
5336
5337 /// Like `decode_step_t_h` with an optional RESIDENT embed table (spec hot loop): device
5338 /// gather instead of host dequant + [T, n_embd] f32 htod. Bit-identical rows.
5339 pub fn decode_step_t_h_emb(
5340 &self,
5341 e: &Engine,
5342 tokens: &[u32],
5343 pos0: usize,
5344 cache: &mut Cache,
5345 embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
5346 ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
5347 let (logits_d, h_seed) = self.decode_step_t_h_emb_dev(e, tokens, pos0, cache, embd_dev)?;
5348 Ok((e.dtoh(&logits_d)?, h_seed))
5349 }
5350
5351 /// DEVICE-LOGITS verify forward (spec device-argmax lever): identical kernel chain to
5352 /// `decode_step_t_h_emb` but returns the [T, n_vocab] logits ON DEVICE — the accept walk
5353 /// argmaxes each column on-device and reads back ONE [T] u32 instead of dtoh'ing the full
5354 /// T x n_vocab f32 block (~1-4 MB + T host argmaxes, every round). Kernel dispatch is
5355 /// UNCHANGED (same decode-exact kernels); only the post-logits transfer moves.
5356 pub fn decode_step_t_h_emb_dev(
5357 &self,
5358 e: &Engine,
5359 tokens: &[u32],
5360 pos0: usize,
5361 cache: &mut Cache,
5362 embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
5363 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
5364 let n_embd = self.cfg.n_embd as usize;
5365 let t = tokens.len();
5366 let (logits, x) = self.decode_step_t_core(e, tokens, pos0, cache, embd_dev, None)?;
5367 // h_seed for the next round = LAST column's pre-output_norm hidden ([n_embd]).
5368 let mut hs = vbuf(e, n_embd)?; // fully written by copy_view_into below
5369 e.copy_view_into(&mut hs, 0, &x.slice((t - 1) * n_embd..t * n_embd), n_embd)?;
5370 Ok((logits, hs))
5371 }
5372
5373 /// CORE verify forward: the `decode_step_t_h_emb_dev` kernel chain, returning the FULL
5374 /// pre-output_norm hidden stack x ([T, n_embd], any column extractable) and optionally
5375 /// filling a `VerifyCkpt` (retained per-layer state-rebuild inputs) for the REPLAY-FREE
5376 /// partial accept. `ckpt: None` => byte-for-byte the old behavior (the ckpt writes are pure
5377 /// retains/copies — they never change what any kernel computes).
5378 fn decode_step_t_core(
5379 &self,
5380 e: &Engine,
5381 tokens: &[u32],
5382 pos0: usize,
5383 cache: &mut Cache,
5384 embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
5385 mut ckpt: Option<&mut VerifyCkpt>,
5386 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
5387 self.decode_step_t_core_stream(
5388 e,
5389 tokens,
5390 pos0,
5391 cache,
5392 embd_dev,
5393 ckpt.take(),
5394 None,
5395 None,
5396 None,
5397 None,
5398 )
5399 }
5400
5401 /// [`Self::decode_step_t_core`] with the MTP route's verify-graph pool armed
5402 /// (`MEMRA_SPEC_VERIFY_GRAPH`). `graphs: None` reproduces `decode_step_t_core`
5403 /// argument-for-argument, so the eager walk stays the byte-identical fallback.
5404 fn decode_step_t_core_vg(
5405 &self,
5406 e: &Engine,
5407 tokens: &[u32],
5408 pos0: usize,
5409 cache: &mut Cache,
5410 embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
5411 mut ckpt: Option<&mut VerifyCkpt>,
5412 graphs: Option<&mut DsparkVerifyGraphs>,
5413 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
5414 self.decode_step_t_core_stream(
5415 e,
5416 tokens,
5417 pos0,
5418 cache,
5419 embd_dev,
5420 ckpt.take(),
5421 None,
5422 None,
5423 None,
5424 graphs,
5425 )
5426 }
5427
5428 /// Increment-0 two-session PP seam: release the peer after this lane's stage-0 boundary TX.
5429 /// The two independent sessions keep their own cache/checkpoint state; only issue order moves.
5430 fn decode_step_t_core_pipelined(
5431 &self,
5432 e: &Engine,
5433 tokens: &[u32],
5434 pos0: usize,
5435 cache: &mut Cache,
5436 embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
5437 mut ckpt: Option<&mut VerifyCkpt>,
5438 pipe: &SpecPipeLane,
5439 round: usize,
5440 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
5441 let fence = crate::pp::pp_cuts(self.layers.len())
5442 .ok_or("two-session speculative pipeline requires a PP stage cut")?;
5443 if crate::pp::pp2_streams_off() || !crate::pp::spec_pp_on() {
5444 return Err("two-session speculative pipeline requires the PP verify split".into());
5445 }
5446 let interval_fence = pipe.stage0_begin(round)?;
5447 let ticket = self.verify_stage0_issue(
5448 e,
5449 tokens,
5450 pos0,
5451 cache,
5452 embd_dev,
5453 ckpt.as_deref_mut(),
5454 None,
5455 &fence,
5456 Some(interval_fence),
5457 pipe.trace(round),
5458 )?;
5459 pipe.stage0_end(round);
5460 pipe.stage1_begin(round)?;
5461 let result = self.verify_stage1_finish(e, ticket, cache, ckpt, None, &fence, true)?;
5462 pipe.verify_end(round);
5463 Ok(result)
5464 }
5465
5466 /// ROUND-STREAM stage (c) 4: `stream` = (device verify tokens [t], device pos counter) —
5467 /// when Some, rope positions come from pos_iota over the counter, the embed gathers the
5468 /// device tokens, and full_attn_verify routes appends/FA through the _dc twins reading the
5469 /// SAME counter (every layer's kvl.len == cache.pos, one counter drives all three). The
5470 /// host `tokens`/`pos0` args still size buffers (t is FIXED K+1 in stream mode).
5471 /// `vtok_dev` (engine-bundle slice 2): device verify tokens for the EMBED only —
5472 /// unlike `stream` mode it changes nothing else (host pos iota, host-len KV appends).
5473 /// `tokens` then only sizes buffers (the dummy-slice pattern the round-stream arm uses).
5474 #[allow(clippy::too_many_arguments)]
5475 fn decode_step_t_core_stream(
5476 &self,
5477 e: &Engine,
5478 tokens: &[u32],
5479 pos0: usize,
5480 cache: &mut Cache,
5481 embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
5482 mut ckpt: Option<&mut VerifyCkpt>,
5483 stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
5484 pp_pipe: Option<bool>,
5485 vtok_dev: Option<&CudaSlice<u32>>,
5486 graphs: Option<&mut DsparkVerifyGraphs>,
5487 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
5488 // PP DOOR (lane/pp2-spec 2026-08-06): the verify trunk now takes its OWN stage split,
5489 // exactly as the eager and batched steps do. This is the single funnel every verify
5490 // forward reaches (decode_step_t / _h / _h_emb / _h_emb_dev / _core all land here), so
5491 // wiring it here wires the whole spec surface — the draft/accept/commit machinery above
5492 // is untouched.
5493 //
5494 // History: pp2-hardening (2026-08-06) made this funnel FAIL CLOSED, because its trunk
5495 // walk was unsplit on one stream and a sharded cross-device placement peer-read every
5496 // remote layer's weights on every spec round (measured 13.9-28x on the batched twin).
5497 // The refusal below survives to cover the residue — MEMRA_SPEC_PP=0, MEMRA_PP_STREAMS=0,
5498 // or a placement whose PpNRt fails to build — so a config that would still walk the
5499 // whole trunk on one stream refuses instead of regressing 28x.
5500 if let Some(fence) = crate::pp::pp_cuts(self.layers.len()) {
5501 if !crate::pp::pp2_streams_off() && crate::pp::spec_pp_on() {
5502 if vtok_dev.is_some() {
5503 return Err(
5504 "device-token dspark verify (slice-2 deferred readback) has no PP \
5505 stage-split arm; set MEMRA_DSPARK_DEFER_READBACK=0 or run the dspark \
5506 route on one device"
5507 .into(),
5508 );
5509 }
5510 return self.decode_step_t_core_ppn(
5511 e,
5512 tokens,
5513 pos0,
5514 cache,
5515 embd_dev,
5516 ckpt.take(),
5517 stream,
5518 &fence,
5519 pp_pipe,
5520 );
5521 }
5522 }
5523 crate::pp::refuse_unsplit_if_remote(
5524 "decode_step_t (spec verify)",
5525 "drop MEMRA_SPEC_PP=0 / MEMRA_PP_STREAMS=0 so the verify trunk takes its OWN stage \
5526 split (decode_step_t_core_ppn); or run spec on one device",
5527 )?;
5528 let cfg = &self.cfg;
5529 let n_embd = cfg.n_embd as usize;
5530 let eps = cfg.rms_eps;
5531 let t = tokens.len();
5532 let pos_d = match stream {
5533 Some((_, ctr)) => {
5534 let mut p = e.alloc_uninit::<i32>(t)?;
5535 e.pos_iota(ctr, &mut p, t)?;
5536 p
5537 }
5538 None => {
5539 let pos_vec: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
5540 e.htod_i32(&pos_vec)?
5541 }
5542 };
5543
5544 // embed T tokens -> [T, n_embd] token-major (device gather on the spec hot loop)
5545 let x = match (stream, embd_dev) {
5546 (Some((vtok, _)), Some((g, qt, rb))) => {
5547 e.embed_gather_device_td(g, vtok, t, n_embd, qt, rb)?
5548 }
5549 (None, Some((g, qt, rb))) => match vtok_dev {
5550 // slice 2: device verify tokens, same embed_gather_u32_t kernel —
5551 // bit-identical rows to the host-token arm (same per-dtype deq).
5552 Some(vt_d) => e.embed_gather_device_td(g, vt_d, t, n_embd, qt, rb)?,
5553 None => e.embed_gather_device_t(g, tokens, n_embd, qt, rb)?,
5554 },
5555 _ => {
5556 assert!(
5557 vtok_dev.is_none(),
5558 "device-token verify requires the resident embed table (embd_dev)"
5559 );
5560 e.htod(&self.embd.gather(n_embd, tokens))?
5561 }
5562 };
5563
5564 // TRUNK WALK: layers [0, n_layers) through the SAME range-scoped subgraph the PP-N
5565 // stage split calls per stage (`verify_layers`) — one code path, so the split cannot
5566 // drift from the unsplit dispatch mirroring. lane/pp2-spec 2026-08-06.
5567 let x = self.verify_layers(
5568 e,
5569 x,
5570 0,
5571 self.layers.len(),
5572 &pos_d,
5573 pos0,
5574 t,
5575 cache,
5576 ckpt.take(),
5577 stream,
5578 graphs,
5579 )?;
5580 if spec_nan_scan() {
5581 nan_scan_rows(e, &x, t, n_embd, &format!("verify trunk exit pos0={pos0}"))?;
5582 }
5583
5584 let mut hn = vbuf(e, t * n_embd)?;
5585 // Stage-A door: with the serving-class row-outer verify walk, the TAIL must be the
5586 // t=1 decode program per row too (rms_norm t=1 + the single-row bf16 head — the
5587 // split head's concat is receipted bit-identical to it). The batched cuBLASLt head
5588 // is a different ULP class and flips near-tie argmaxes off the greedy tape.
5589 let eager_tail = self.sliding_gated_moe_batch_program() && spec_verify_eager_on();
5590 if eager_tail {
5591 let n_vocab = self.cfg.n_vocab as usize;
5592 // MEMRA_SPEC_HEAD_ROWS=1 — THE VERIFY TAIL'S REDUNDANT HEAD READ.
5593 //
5594 // The loop below runs the head at m=1 once PER COLUMN, so the LM head's weights are
5595 // streamed t times per verify pass. On step37 that head is ~0.49 GiB per card after the
5596 // rank split, ~1.07 ms of pure re-read at t=2 and worse at every wider t — which is a
5597 // large part of why the fixed K ladder LOSES (K=1 81.2 > K=2 73.1 > K=3 62.7 tok/s).
5598 //
5599 // The loop's justification is the comment above: the batched cuBLASLt head is a
5600 // different ULP class and flips near-tie argmaxes off the greedy tape. That is true of
5601 // cuBLASLt and it does NOT apply here, because a FloatBf16 head at 1..=32 rows never
5602 // reaches cuBLASLt: `matmul` routes it to `matvec_bf16_rows_into` (lib.rs:12248), whose
5603 // own doc says `matvec_bf16_f32acc_x4_rows` "runs the t=1 decode head program PER ROW
5604 // (identical dot + reduce), so decode/verify tiers keep the t=1 numeric class". Under
5605 // the W8 doors both widths route to the q8 mirror instead, and the t-column mirror is
5606 // documented "bit-identical to t single-row calls". So the batched form is the SAME
5607 // arithmetic per row on both paths, with one weight read instead of t.
5608 //
5609 // rms_norm is row-wise, so norm(t) is per-row identical to t x norm(1) by construction.
5610 //
5611 // DEFAULT OFF for exactly one turn of the crank: "bit-identical by two documented
5612 // claims" is still an argument. The greedy byte tape decides, and the door flips only
5613 // once the tape is a receipt.
5614 if head_rows_on() {
5615 e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
5616 let logits = e.matmul(&self.output, &hn, t)?;
5617 if stream.is_none() {
5618 cache.pos += t;
5619 }
5620 return Ok((logits, if spec_hpost() { hn } else { x }));
5621 }
5622 let mut logits = vbuf(e, t * n_vocab)?;
5623 for r in 0..t {
5624 let mut row = e.uninit(n_embd)?;
5625 e.dtod_copy_view(&x.slice(r * n_embd..(r + 1) * n_embd), &mut row)?;
5626 let mut hr = e.uninit(n_embd)?;
5627 e.rms_norm(&row, self.output_norm.float_data(), &mut hr, n_embd, 1, eps)?;
5628 let lr = e.matmul(&self.output, &hr, 1)?;
5629 e.dtod_copy_into(&lr, &mut logits, r * n_vocab)?;
5630 e.dtod_copy_into(&hr, &mut hn, r * n_embd)?;
5631 }
5632 if stream.is_none() {
5633 cache.pos += t;
5634 }
5635 return Ok((logits, if spec_hpost() { hn } else { x }));
5636 }
5637 let serving_head =
5638 self.sliding_gated_moe_batch_program() || self.batched_serving_numeric_class();
5639 let logits = if serving_head {
5640 // Step35 and the qwen35 family (MoE 2026-08-14 AM, dense-hybrid same day PM — the
5641 // Q3.8 bring-up reproduced the identical near-tie class on dense: eager-class verify
5642 // vs batched-class live serving, ULP drift amplified through the GDN recurrence)
5643 // serve one batched numeric class at every live width, including B=1. Keep the
5644 // verify head in that same class; other generic families retain the decode-exact
5645 // head that their run-spec contract pins.
5646 e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
5647 e.matmul(&self.output, &hn, t)?
5648 } else {
5649 e.rms_norm_decode(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
5650 e.matmul_decode_exact(&self.output, &hn, t)?
5651 };
5652 // stream: the device pos counter owns position; host mirror reconciles at drain.
5653 if stream.is_none() {
5654 cache.pos += t;
5655 }
5656 // Hidden stack for seeds/refresh-fills: pre-norm x (default) or post-norm hn (HPOST).
5657 Ok((logits, if spec_hpost() { hn } else { x }))
5658 }
5659
5660 /// THE VERIFY TRUNK OVER PP-N (lane/pp2-spec 2026-08-06): `decode_step_t_core_stream`'s walk
5661 /// as N stage subgraphs, each on its own engine/stream (and, under `MEMRA_PP_DEVICES`, its own
5662 /// device), with a `[T, n_embd]` boundary transfer between them. T = K+1 (the verify batch),
5663 /// so this is the batched-boundary shape the pp2-batch lane's grow-only slots already handle
5664 /// (`tx(b, x, t*n_embd)`; the slot grows to the high-water T and the transport moves exactly
5665 /// the payload).
5666 ///
5667 /// Structure is `decode_step_batch_ppn`'s, which is `decode_step_h_ppn`'s. FOUR THINGS ARE
5668 /// PER-STAGE and each for a measured reason (see `decode_step_batch_ppn`'s header for the
5669 /// receipts):
5670 ///
5671 /// 1. THE ENGINE (`rt.engine(s, e)`) — `Engine` owns lazily-grown stable-pointer scratch
5672 /// (`fa_part_pool`, `fa_vf16_scratch`, `argmax_partials`) that is single-stream-safe BY
5673 /// DESIGN. Two stage streams through one Engine is the 2026-08-02 shared-scratch race
5674 /// (35% flake, nondeterministic all-logits divergence). `PpNRt::build` gives every stage
5675 /// s>0 its own Engine even on the primary device; honouring it here is what scopes the
5676 /// pools. The verify path allocates MORE of that scratch than eager decode does (FA at
5677 /// m=T, and the per-layer `GdnStash` retains), so this is load-bearing, not inherited.
5678 ///
5679 /// 2. `pos_d` — each stage uploads its OWN copy of the T rope positions on ITS stream, so the
5680 /// buffer is allocated, consumed and freed on one stream. In `stream` mode that means each
5681 /// stage runs its own `pos_iota` over the SHARED device counter (`pos_ctr`): the counter is
5682 /// read-only during the forward (the round's `inc`/`copy_add` happen outside it), so every
5683 /// stage derives the identical iota, and each stage's own output buffer is stream-local.
5684 ///
5685 /// 3. THE EMBED lives with stage 0 (`self.embd` / `embd_gpu` are host/primary-side; the
5686 /// sharded loader leaves the table with stage 0 by construction).
5687 ///
5688 /// 4. THE HEAD (`output_norm` + `output`) runs on the LAST stage — the sharded loader uploaded
5689 /// both through that stage's engine (`hybrid.rs`: `e_head = layer_engine(e, n_trunk,
5690 /// n_trunk-1)`), so reading them anywhere else is a peer read of the biggest tensor in the
5691 /// model, every round.
5692 ///
5693 /// WHAT STAYS ON THE PRIMARY, deliberately: the returned logits and hidden stack `x`. Both are
5694 /// last-stage-allocated device buffers, and every consumer (the device argmax walk, the accept
5695 /// kernels, `spec_seed_gather`, the ckpt rebuild in `commit_verified_prefix`) reads them
5696 /// through the primary context by UVA — the same read the batched serving epilogue's
5697 /// `last_logits_dev` park does. Those consumers are per-round O(T x n_vocab) and O(n_embd),
5698 /// not per-layer, so they are not the 28x class; splitting them is a separate lane.
5699 ///
5700 /// The MTP HEAD (draft side) is NOT split: it is one block, it lives wherever the loader put
5701 /// it (`load_mtp` uses the primary engine), and it is ~1-2 GB against the trunk's tens. Draft
5702 /// placement is measured, not assumed — see `research/pp2-spec-20260806`.
5703 ///
5704 /// EXACTNESS: PP-N adds ZERO deviation. Each stage runs the SAME kernels on the SAME bytes in
5705 /// the same order via the SAME `verify_layers` the unsplit body calls; the only change is
5706 /// where the residual is materialized, and the boundary is a straight f32 copy (dtod
5707 /// same-device / `cudaMemcpyPeerAsync` cross-device, no conversion). So the split MUST be
5708 /// BIT-IDENTICAL to the unsplit verify at the same T, in both placement orders. Gate:
5709 /// `decode-batch-gate --mode ppspec`. Acceptance counts are a DERIVED consequence — greedy
5710 /// accept argmaxes these logits, so bit-identical logits force identical accept walks; the
5711 /// `run-spec` K=1..8 arm checks that end-to-end rather than trusting the implication.
5712 #[allow(clippy::too_many_arguments)]
5713 fn decode_step_t_core_ppn(
5714 &self,
5715 e: &Engine,
5716 tokens: &[u32],
5717 pos0: usize,
5718 cache: &mut Cache,
5719 embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
5720 mut ckpt: Option<&mut VerifyCkpt>,
5721 stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
5722 fence: &[usize],
5723 pp_pipe: Option<bool>,
5724 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
5725 let ticket = self.verify_stage0_issue(
5726 e,
5727 tokens,
5728 pos0,
5729 cache,
5730 embd_dev,
5731 ckpt.as_deref_mut(),
5732 stream,
5733 fence,
5734 pp_pipe,
5735 None,
5736 )?;
5737 self.verify_stage1_finish(e, ticket, cache, ckpt, stream, fence, true)
5738 }
5739
5740 /// Enqueue embed, stage 0, and the first boundary TX, then return the actual boundary slot.
5741 /// The ordinary PP verify wrapper calls `verify_stage1_finish` immediately after this return.
5742 #[allow(clippy::too_many_arguments)]
5743 fn verify_stage0_issue(
5744 &self,
5745 e: &Engine,
5746 tokens: &[u32],
5747 pos0: usize,
5748 cache: &mut Cache,
5749 embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
5750 mut ckpt: Option<&mut VerifyCkpt>,
5751 stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
5752 fence: &[usize],
5753 pp_pipe: Option<bool>,
5754 trace: Option<SpecPipeTraceCtx>,
5755 ) -> Result<VerifyBoundaryTicket, Box<dyn std::error::Error>> {
5756 assert!(
5757 !self.is_gemma4_e4b() && !self.gemma_batch_program(),
5758 "decode_step_t_core_ppn covers the hybrid non-gemma4 verify trunk only \
5759 (the gemma4 arms have their own decode_step_t twins)"
5760 );
5761 if crate::pp::pp_host_bounce_active() && (stream.is_some() || embd_dev.is_some()) {
5762 return Err(
5763 "decode_step_t_core_ppn: refused with MEMRA_PP_HOST_BOUNCE=1 — the trunk \
5764 boundary itself is host-staged, but device-resident verify still peer-reads \
5765 primary-device token/position/embedding buffers from stage 0. Run plain PP \
5766 serving on this host class; spec requires local per-stage inputs first."
5767 .into(),
5768 );
5769 }
5770 let rt = crate::pp::PpNRt::get(e)?;
5771 let n_st = fence.len() - 1;
5772 assert_eq!(
5773 rt.n_stages(),
5774 n_st,
5775 "PpNRt stage count {} != fence stages {n_st}",
5776 rt.n_stages()
5777 );
5778 let n_embd = self.cfg.n_embd as usize;
5779 let t = tokens.len();
5780 let payload = t * n_embd;
5781 if pp_pipe.is_some() {
5782 assert_eq!(n_st, 2, "spec pipeline requires exactly two PP stages");
5783 }
5784 // One-shot lane diagnostic: force natural PP-2 boundaries to completion so the server
5785 // log can price stage 0, the peer hop, the RX copy, and stage 1 + head separately without
5786 // nsys. The ordinary path keeps every enqueue asynchronous. N>2 is deliberately excluded:
5787 // the report below names exactly two stages and must never imply it measured middle ones.
5788 let pp_anatomy = n_st == 2 && std::env::var("MEMRA_SPEC_PP_ANATOMY").as_deref() == Ok("1");
5789 let pp_started = std::time::Instant::now();
5790 let (mut reverse_ms, mut stage0_ms, mut tx_ms) = (0.0f64, 0.0f64, 0.0f64);
5791 // The CALLER's ambient stream, captured BEFORE any `rt.enter()` pushes a stage stream:
5792 // this body returns DEVICE-RESIDENT buffers (the device-argmax accept walk's contract),
5793 // so the exit needs the same publication the boundaries get — see `PpNRt::publish_to`.
5794 // Taken here, not at the end, because inside the last-stage scope `e.stream()` IS the
5795 // stage stream and the wait would self-order into a no-op.
5796 let caller_stream = e.stream();
5797 // #87 ROOT-CAUSE FENCE (lane/pp2spec-crash): the PREVIOUS round's stage-allocated
5798 // outputs (logits/hidden/ckpt stashes) freed stream-ordered on the STAGE streams while
5799 // the primary stream still holds queued reads of them — with event tracking elided,
5800 // nothing stops the pool from reusing those blocks for THIS round's stage allocations,
5801 // whose writes then race the queued reads (measured: 13/4096-NaN random-bits garbage in
5802 // the spec round seed; the full anatomy is on `PpNRt::fence_stages_behind`). Order every
5803 // stage stream behind the caller before enqueueing new stage work.
5804 let reverse_started = std::time::Instant::now();
5805 if pp_pipe != Some(false) {
5806 rt.fence_stages_behind(&caller_stream)?;
5807 }
5808 if pp_pipe == Some(true) {
5809 // Both session verifies must alternate boundary slots even when the ordinary
5810 // decode overlap experiment is off. Prewarm before A's stage 0 so B cannot grow
5811 // slot 1 by synchronizing the RX stream while A's stage 1 is in flight.
5812 rt.prepare_overlap_slots(0, payload)?;
5813 }
5814 if pp_anatomy {
5815 // Drain the reverse-publication dependency before timing stage 0 itself. At c=1 this
5816 // prices any primary-stream rollback/refresh tail inherited from the prior round.
5817 for s in 0..n_st {
5818 let _st = rt.enter(s);
5819 rt.engine(s, e).stream().synchronize()?;
5820 }
5821 reverse_ms = reverse_started.elapsed().as_secs_f64() * 1e3;
5822 }
5823
5824 // Per-stage rope positions: in host mode the same [T] iota each stage uploads itself; in
5825 // stream mode each stage's own `pos_iota` over the shared read-only device counter.
5826 let stage_pos = |es: &Engine| -> Result<CudaSlice<i32>, Box<dyn std::error::Error>> {
5827 match stream {
5828 Some((_, ctr)) => {
5829 let mut p = es.alloc_uninit::<i32>(t)?;
5830 es.pos_iota(ctr, &mut p, t)?;
5831 Ok(p)
5832 }
5833 None => {
5834 let pos_vec: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
5835 es.htod_i32(&pos_vec)
5836 }
5837 }
5838 };
5839
5840 // ---- STAGE 0: embed (the table lives with stage 0) + layers [0, fence[1]) + TX ----
5841 let slot = {
5842 let _st0 = rt.enter(0);
5843 let e0 = rt.engine(0, e);
5844 enqueue_spec_pipe_trace_marker(&e0.stream(), trace.as_ref(), "S0", "start", None)?;
5845 let stage0_started = std::time::Instant::now();
5846 let pos_d = stage_pos(e0)?;
5847 let x = match (stream, embd_dev) {
5848 (Some((vtok, _)), Some((g, qt, rb))) => {
5849 e0.embed_gather_device_td(g, vtok, t, n_embd, qt, rb)?
5850 }
5851 (None, Some((g, qt, rb))) => e0.embed_gather_device_t(g, tokens, n_embd, qt, rb)?,
5852 _ => e0.htod(&self.embd.gather(n_embd, tokens))?,
5853 };
5854 let x = self.verify_layers(
5855 e0,
5856 x,
5857 fence[0],
5858 fence[1],
5859 &pos_d,
5860 pos0,
5861 t,
5862 cache,
5863 ckpt.as_deref_mut(),
5864 stream,
5865 None,
5866 )?;
5867 if pp_anatomy {
5868 e0.stream().synchronize()?;
5869 stage0_ms = stage0_started.elapsed().as_secs_f64() * 1e3;
5870 }
5871 let tx_started = std::time::Instant::now();
5872 let slot = if pp_pipe.is_some() {
5873 rt.tx_pipelined(0, &x, payload)?
5874 } else {
5875 rt.tx(0, &x, payload)?
5876 };
5877 enqueue_spec_pipe_trace_marker(&e0.stream(), trace.as_ref(), "S0", "end", Some(slot))?;
5878 if pp_anatomy {
5879 e0.stream().synchronize()?;
5880 tx_ms = tx_started.elapsed().as_secs_f64() * 1e3;
5881 }
5882 slot
5883 // x + pos_d drop here: freed stream-ordered on stage-0's stream after use.
5884 };
5885
5886 Ok(VerifyBoundaryTicket {
5887 rt,
5888 caller_stream,
5889 slot,
5890 pos0,
5891 t,
5892 payload,
5893 n_st,
5894 pipelined: pp_pipe.is_some(),
5895 pp_anatomy,
5896 pp_started,
5897 reverse_ms,
5898 stage0_ms,
5899 tx_ms,
5900 trace,
5901 })
5902 }
5903
5904 /// Consume a stage-0 boundary ticket and enqueue the remaining PP stages plus the head.
5905 /// On PP-2 this is exactly stage 1; PP-N keeps its pre-existing middle-stage walk here.
5906 #[allow(clippy::too_many_arguments)]
5907 fn verify_stage1_finish(
5908 &self,
5909 e: &Engine,
5910 ticket: VerifyBoundaryTicket,
5911 cache: &mut Cache,
5912 mut ckpt: Option<&mut VerifyCkpt>,
5913 stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
5914 fence: &[usize],
5915 publish_to_caller: bool,
5916 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
5917 let VerifyBoundaryTicket {
5918 rt,
5919 caller_stream,
5920 slot,
5921 pos0,
5922 t,
5923 payload,
5924 n_st,
5925 pipelined,
5926 pp_anatomy,
5927 pp_started,
5928 reverse_ms,
5929 stage0_ms,
5930 tx_ms,
5931 trace,
5932 } = ticket;
5933 let n_embd = self.cfg.n_embd as usize;
5934 let eps = self.cfg.rms_eps;
5935 let mut slot = slot;
5936 let (mut rx_ms, mut stage1_ms) = (0.0f64, 0.0f64);
5937 let stage_pos = |es: &Engine| -> Result<CudaSlice<i32>, Box<dyn std::error::Error>> {
5938 match stream {
5939 Some((_, ctr)) => {
5940 let mut p = es.alloc_uninit::<i32>(t)?;
5941 es.pos_iota(ctr, &mut p, t)?;
5942 Ok(p)
5943 }
5944 None => {
5945 let pos_vec: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
5946 es.htod_i32(&pos_vec)
5947 }
5948 }
5949 };
5950
5951 // ---- MIDDLE STAGES: RX boundary s-1 -> range -> TX boundary s ----
5952 for s in 1..n_st - 1 {
5953 let _st = rt.enter(s);
5954 let es = rt.engine(s, e);
5955 let pos_d = stage_pos(es)?;
5956 let x = rt.rx(s - 1, slot, payload)?;
5957 let x = self.verify_layers(
5958 es,
5959 x,
5960 fence[s],
5961 fence[s + 1],
5962 &pos_d,
5963 pos0,
5964 t,
5965 cache,
5966 ckpt.as_deref_mut(),
5967 stream,
5968 None,
5969 )?;
5970 slot = if pipelined {
5971 rt.tx_pipelined(s, &x, payload)?
5972 } else {
5973 rt.tx(s, &x, payload)?
5974 };
5975 }
5976
5977 // ---- LAST STAGE: RX + final range + output_norm + lm head ----
5978 let _stl = rt.enter(n_st - 1);
5979 let el = rt.engine(n_st - 1, e);
5980 let pos_d = stage_pos(el)?;
5981 let rx_started = std::time::Instant::now();
5982 let x = rt.rx(n_st - 2, slot, payload)?;
5983 if pp_anatomy {
5984 el.stream().synchronize()?;
5985 rx_ms = rx_started.elapsed().as_secs_f64() * 1e3;
5986 }
5987 enqueue_spec_pipe_trace_marker(&el.stream(), trace.as_ref(), "S1", "start", Some(slot))?;
5988 let stage1_started = std::time::Instant::now();
5989 let x = self.verify_layers(
5990 el,
5991 x,
5992 fence[n_st - 1],
5993 fence[n_st],
5994 &pos_d,
5995 pos0,
5996 t,
5997 cache,
5998 ckpt.as_deref_mut(),
5999 stream,
6000 None,
6001 )?;
6002
6003 let mut hn = vbuf(el, payload)?;
6004 let logits = if self.sliding_gated_moe_batch_program() {
6005 // The PP Step35 serving path uses rms_norm + matmul for B=1 as well as B>1.
6006 // Verify must not switch numeric class merely because the same session speculates.
6007 el.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
6008 el.matmul(&self.output, &hn, t)?
6009 } else {
6010 el.rms_norm_decode(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
6011 el.matmul_decode_exact(&self.output, &hn, t)?
6012 };
6013 enqueue_spec_pipe_trace_marker(&el.stream(), trace.as_ref(), "S1", "end", Some(slot))?;
6014 if pp_anatomy {
6015 el.stream().synchronize()?;
6016 stage1_ms = stage1_started.elapsed().as_secs_f64() * 1e3;
6017 }
6018 // EXIT PUBLICATION: both returned buffers are still being produced on the last stage's
6019 // stream. Order the caller's stream behind that work before the buffers escape this
6020 // scope (the 2026-08-06 same-device ppspec find: without it the caller's primary-stream
6021 // consumer read unwritten logits — nondeterministic, one-device-only, and it poisoned
6022 // the following arm's KV in the same process).
6023 if publish_to_caller {
6024 rt.publish_to(n_st - 1, &caller_stream)?;
6025 }
6026 if pp_anatomy {
6027 if publish_to_caller {
6028 caller_stream.synchronize()?;
6029 }
6030 eprintln!(
6031 "[spec-pp-anatomy] t={t} reverse={reverse_ms:.3}ms stage0={stage0_ms:.3}ms \
6032 tx={tx_ms:.3}ms rx={rx_ms:.3}ms stage1-head={stage1_ms:.3}ms total={:.3}ms",
6033 pp_started.elapsed().as_secs_f64() * 1e3,
6034 );
6035 }
6036 // stream: the device pos counter owns position; host mirror reconciles at drain.
6037 if stream.is_none() {
6038 cache.pos += t;
6039 }
6040 Ok((logits, if spec_hpost() { hn } else { x }))
6041 }
6042
6043 /// Step3.5/Step3.7 verify trunk in the serving batched numeric class.
6044 ///
6045 /// `step35_decode_batch_layers` is now authoritative at every live serving width, including
6046 /// B=1 (lane/cx-b1fix). The older verify walk deliberately mirrored the eager T=1 class:
6047 /// it replayed `step35_decode_attn` per row and used the eager/decode-exact FFN dispatch.
6048 /// Those classes are individually stable, but a near-tie prompt can choose different greedy
6049 /// bytes when a request moves from batched plain serving into speculative verify. Run the
6050 /// same authoritative B=1 stage subgraph for each verify row here. Rows still advance
6051 /// layer-by-layer, so every layer sees the preceding verify rows in its attention cache while
6052 /// every norm/projection/FFN uses exactly the live serving dispatch.
6053 #[allow(clippy::too_many_arguments)]
6054 /// PRIME-BY-T-ROWS (MEMRA_PRIME_TROWS=1): prefill the prompt through the same-session
6055 /// t-row walk in 32-row chunks — every row runs the t=1 decode program bit-for-bit
6056 /// (the TOKENWISE-prime ORACLE class), so this door is exact against the exactness
6057 /// reference while replacing the host-canonical per-token prime. Requires the walk
6058 /// doors (MEMRA_SPEC_VERIFY_EAGER/TCOL); returns the prime contract trio.
6059 #[allow(clippy::type_complexity)]
6060 pub(crate) fn step35_prime_trows(
6061 &self,
6062 e: &Engine,
6063 tokens: &[u32],
6064 cache: &mut Cache,
6065 ) -> Result<Option<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
6066 {
6067 let dbg = std::env::var("MEMRA_SPEC_FA2_DEBUG").as_deref() == Ok("1");
6068 if !prime_trows_on() {
6069 return Ok(None);
6070 }
6071 if !self.uses_sliding_gated_moe_program()
6072 || cache.pos != 0
6073 || cache.dflash_taps.is_some()
6074 || !spec_verify_eager_on()
6075 || !spec_verify_tcol_on()
6076 {
6077 if dbg {
6078 eprintln!(
6079 "[prime-trows] refuse: program={} pos={} taps={} eager={:?} tcol={:?}",
6080 self.uses_sliding_gated_moe_program(),
6081 cache.pos,
6082 cache.dflash_taps.is_some(),
6083 std::env::var("MEMRA_SPEC_VERIFY_EAGER").ok(),
6084 std::env::var("MEMRA_SPEC_VERIFY_TCOL").ok()
6085 );
6086 }
6087 return Ok(None);
6088 }
6089 let n_embd = self.cfg.n_embd as usize;
6090 let n_layers = self.layers.len();
6091 let t_total = tokens.len();
6092 let Some(embd_gpu) = self.embd_gpu_try(e) else {
6093 if dbg {
6094 eprintln!("[prime-trows] refuse: no device embed table");
6095 }
6096 return Ok(None);
6097 };
6098 let embd_qtype = match self.embd.ggml_type {
6099 memra_gguf::GgmlType::BF16 => crate::QT_BF16,
6100 memra_gguf::GgmlType::Q8_0 => crate::QT_Q8_0,
6101 other => {
6102 if dbg {
6103 eprintln!("[prime-trows] refuse: embed dtype {other:?}");
6104 }
6105 return Ok(None);
6106 }
6107 };
6108 let embd_row_bytes = self.embd.raw.len() / self.cfg.n_vocab as usize;
6109 // Chunk plan: 32-row chunks; a 1-token tail folds into the previous chunk
6110 // (the walk floor is t >= 2).
6111 let mut bounds = Vec::new();
6112 let mut start = 0usize;
6113 while start < t_total {
6114 let mut end = (start + 32).min(t_total);
6115 if t_total - end == 1 {
6116 end -= 1;
6117 }
6118 bounds.push((start, end));
6119 start = end;
6120 }
6121 if bounds.iter().any(|(a, b)| b - a < 2) {
6122 return Ok(None); // degenerate short prompt keeps the ordinary prime
6123 }
6124 let mut hiddens = e.uninit(t_total * n_embd)?;
6125 let mut last: Option<CudaSlice<f32>> = None;
6126 for &(a, b) in &bounds {
6127 let tc = b - a;
6128 let tok_d = e.stream().clone_htod(&tokens[a..b])?;
6129 let x =
6130 e.embed_gather_device_td(embd_gpu, &tok_d, tc, n_embd, embd_qtype, embd_row_bytes)?;
6131 let out = self.step35_verify_batch_layers(e, x, 0, n_layers, a, tc, cache)?;
6132 e.copy_into(&mut hiddens, a * n_embd, &out, tc * n_embd)?;
6133 if b == t_total {
6134 let mut h = e.uninit(n_embd)?;
6135 e.dtod_copy_view(&out.slice((tc - 1) * n_embd..tc * n_embd), &mut h)?;
6136 last = Some(h);
6137 }
6138 }
6139 let h_seed = last.expect("last chunk produced the seed row");
6140 let mut hn = e.uninit(n_embd)?;
6141 e.rms_norm_decode(
6142 &h_seed,
6143 self.output_norm.float_data(),
6144 &mut hn,
6145 n_embd,
6146 1,
6147 self.cfg.rms_eps,
6148 )?;
6149 let logits_d = e.matmul_decode_exact(&self.output, &hn, 1)?;
6150 let logits = e.dtoh(&logits_d)?;
6151 cache.pos = t_total;
6152 Ok(Some((logits, h_seed, hiddens)))
6153 }
6154
6155 fn step35_verify_batch_layers(
6156 &self,
6157 e: &Engine,
6158 mut x: CudaSlice<f32>,
6159 lo: usize,
6160 hi: usize,
6161 pos0: usize,
6162 t: usize,
6163 cache: &mut Cache,
6164 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6165 let n_embd = self.cfg.n_embd as usize;
6166 if !self.uses_sliding_gated_moe_program() {
6167 return Err(
6168 "serving-class verify requires sliding-gated-MoE canonical operations".into(),
6169 );
6170 }
6171 // SERVING-CLASS VERIFY (MEMRA_SPEC_VERIFY_EAGER=1, step37 MTP bring-up): each verify
6172 // column rides decode_layers_eager — the EXACT t=1 program live serving runs (all TP2
6173 // doors) — row-outer, so row r's appends land before row r+1 attends: bit-equal to
6174 // plain greedy by construction. Only the unsplit full-range walk qualifies; PP splits
6175 // and the tap path keep the batch-layer class.
6176 static VE: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
6177 let eager_verify =
6178 *VE.get_or_init(spec_verify_eager_on) && lo == 0 && hi == self.layers.len();
6179 if eager_verify {
6180 // T-COLUMN LAYER-OUTER WALK (MEMRA_SPEC_VERIFY_TCOL=1): per layer, one t-grid
6181 // attn norm + ONE weight-amortized QKV(+gate) over all T columns, then each
6182 // column runs the UNMODIFIED t=1 attention program via the col-select door and
6183 // the ordinary residual/FFN body. Values per column are bit-equal to the
6184 // row-outer walk: rms over the materialized residual == the fused add+norm
6185 // (kernel_check identity), the tcol kernel's per-column FP order == the t=1
6186 // kernel, and every downstream op IS the t=1 program.
6187 static TCOL: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
6188 let tcol = *TCOL.get_or_init(spec_verify_tcol_on);
6189 // T > 32 (prefill-class): run the SAME walk in 32-row chunks — each chunk's
6190 // rows are the t=1 program bit-for-bit and the rope pass advances the cache,
6191 // so a chunked call is value-identical to the row-outer loop it replaces.
6192 static TROWS_PREFILL: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
6193 // MEMRA_STEP_GEMM_PRIME outranks the walk: with the grouped GEMM prime armed, the
6194 // t-row walk defers so the batch path (GEMM trunk + grouped MoE) takes the prompt —
6195 // flag precedence between two existing doors, not a new flag. Without this, both
6196 // doors ON meant the walk still won and the GEMM prime needed PRIME_TROWS=0 by hand.
6197 let trows_prefill =
6198 *TROWS_PREFILL.get_or_init(|| prime_trows_on() && !crate::step_gemm_prime_on());
6199 // MEMRA_PRIME_TROWS_T=<w>: chunk width, default 8 = the REAL cap of this walk.
6200 // The workspace slabs go to 32 rows, but `matvec_bf16_qkvg_tcol_into` refuses
6201 // t > 8 (compile-time-T twins exist for 2/4/8 only; the runtime-t kernel spills
6202 // its accumulators to local memory), so a wider chunk fails the request with
6203 // "matvec_bf16_qkvg_tcol geometry" — which is exactly how the first server-path
6204 // TROWS arm died. Measured at 193 tokens: w=8 2.459 s, w=4 2.574 s.
6205 static TROWS_W: std::sync::OnceLock<Result<usize, String>> = std::sync::OnceLock::new();
6206 let trows_w = match TROWS_W.get_or_init(|| {
6207 let value = std::env::var("MEMRA_PRIME_TROWS_T").ok();
6208 parse_prime_trows_width(value.as_deref())
6209 }) {
6210 Ok(width) => *width,
6211 Err(err) => return Err(err.clone().into()),
6212 };
6213 if tcol && trows_prefill && t > trows_w {
6214 // One-time engagement receipt: without it a prefill gate cannot tell a
6215 // chunked walk from the row-outer fallback it is supposed to replace
6216 // (the first PRIME_TROWS gate passed vacuously on exactly that).
6217 static SEEN: std::sync::atomic::AtomicBool =
6218 std::sync::atomic::AtomicBool::new(false);
6219 if !SEEN.swap(true, std::sync::atomic::Ordering::Relaxed) {
6220 eprintln!(
6221 "[prime-trows] ENGAGED t={t} width={trows_w} chunks={} layers={}..{}",
6222 t.div_ceil(trows_w),
6223 lo,
6224 hi
6225 );
6226 }
6227 let mut out = e.uninit(t * n_embd)?;
6228 let mut start = 0usize;
6229 while start < t {
6230 let mut end = (start + trows_w).min(t);
6231 if t - end == 1 {
6232 end -= 1;
6233 }
6234 let tc = end - start;
6235 let mut xc = e.uninit(tc * n_embd)?;
6236 e.dtod_copy_view(&x.slice(start * n_embd..end * n_embd), &mut xc)?;
6237 let oc =
6238 self.step35_verify_batch_layers(e, xc, lo, hi, pos0 + start, tc, cache)?;
6239 e.copy_into(&mut out, start * n_embd, &oc, tc * n_embd)?;
6240 start = end;
6241 }
6242 return Ok(out);
6243 }
6244 if tcol && t >= 2 && t <= 32 {
6245 // MEMRA_TCOL_PROF=1: synchronized per-segment wall profile of the walk
6246 // (norm+QKV precompute / per-col attention / per-col residual+FFN). The
6247 // syncs serialize the stream, so the split is for TARGETING amortization
6248 // work only — never a perf claim.
6249 static PROF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
6250 let prof =
6251 *PROF.get_or_init(|| std::env::var("MEMRA_TCOL_PROF").as_deref() == Ok("1"));
6252 let mut prof_ms = [0f64; 3];
6253 let eps = self.cfg.rms_eps;
6254 let mut x_t = x;
6255 let mut h_t = e.uninit(t * n_embd)?;
6256 let mut h_row = e.uninit(n_embd)?; // real row: the non-dcw fallback reads it
6257 // Per-column pos buffers hoisted out of the layer loop (a per-col-per-layer
6258 // pageable htod was an in-stream engine turnaround x t x 45).
6259 let mut pos_rows = Vec::with_capacity(t);
6260 for r in 0..t {
6261 pos_rows.push(e.htod_i32(&[(pos0 + r) as i32])?);
6262 }
6263 let mut ok = true;
6264 // MEMRA_TCOL_OPROJ=1: defer each column's o_proj — the finish seam
6265 // stashes `gated` instead of joining per column; one b4_tcol per rank +
6266 // one slab join produce every column's `mixed` after the attention pass.
6267 // Bit-exact per column (t=1 b4 program per column; elementwise join).
6268 // MEMRA_TCOL_FFN=1 (implies the o_proj defer): when every column of a
6269 // MoE layer deferred, the residual norm runs as one t-grid launch
6270 // (per-row program == t=1) and the FFN as ONE two-column device-routed
6271 // sweep + per-column shexp — the two columns' expert weights dedup
6272 // through L2 instead of reading HBM twice.
6273 static FFN2: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
6274 let ffn_batch = *FFN2.get_or_init(tcol_ffn_on);
6275 let oproj_batch = crate::tp::tcol_oproj_on() || ffn_batch;
6276 // MEMRA_SPEC_FA2=1 (T=2 only): eligible layers defer BOTH columns' fa —
6277 // the per-column pass norms/ropes/appends and stashes q+gate, then one
6278 // shared-KV fa_decode_dcw2 per rank + the o_proj join produce the
6279 // [2, o_out] mixed slab. The precheck runs before arming (stashing is
6280 // unrecoverable); ineligible/boundary layers run the ordinary program.
6281 let fa2 = crate::tp::spec_fa2_on() && t <= 32;
6282 let mut mixed_row = e.uninit(n_embd)?;
6283 let mut pos_staged = false;
6284 for il in lo..hi {
6285 let layer = &self.layers[il];
6286 // BEFORE this layer touches its planes: is the history it is about to
6287 // attend already poisoned? Global (non-ring) layers only, which are the
6288 // ones the level-2 bitmap implicates.
6289 if kv_plane_scan_on() && self.step35_geom(il).window.is_none() {
6290 if let Some(distributed) = cache.tp_kv[il].as_ref() {
6291 scan_kv_plane(e, distributed, il, pos0)?;
6292 }
6293 }
6294 let fa2_layer = fa2 && self.step35_fa_rows_precheck(cache, il, pos0, t)?;
6295 let mut seg = std::time::Instant::now();
6296 e.rms_norm(&x_t, layer.attn_norm.float_data(), &mut h_t, n_embd, t, eps)?;
6297 if !self.step35_verify_qkv_precompute(e, il, &h_t, t)? {
6298 ok = false;
6299 break;
6300 }
6301 // FULL t-row attention pass (rope/append + fa + combine + o_proj in
6302 // 3 launches/rank): same-session rows, slot = len-base+r, one len
6303 // advance by t. Host cache bookkeeping mirrors the per-column tail.
6304 if fa2_layer {
6305 if let Some(mixed_t) =
6306 self.step35_verify_rope_fa_pass(e, il, cache, pos0, t, !pos_staged)?
6307 {
6308 pos_staged = true;
6309 {
6310 let tp_kv = cache.tp_kv[il]
6311 .as_mut()
6312 .expect("precheck verified the distributed cache");
6313 let transaction = tp_kv.begin_transaction()?;
6314 let crate::hybrid::Mixer::Full(fa) = &layer.mixer else {
6315 return Err("verify rope pass expects full attention".into());
6316 };
6317 let tp = fa
6318 .step_tp_qkv
6319 .as_ref()
6320 .ok_or("verify rope pass lost its TP state")?;
6321 let empty: [CudaSlice<f32>; 0] = [];
6322 tp.runtime.append_tp_kv_transaction_inner(
6323 tp_kv,
6324 transaction,
6325 &empty,
6326 &empty,
6327 t,
6328 true,
6329 )?;
6330 tp.runtime.commit_tp_kv_transaction_external(
6331 tp_kv,
6332 transaction,
6333 t,
6334 )?;
6335 if let Some(local) = cache.kv[il].as_mut() {
6336 local.len = pos0 + t;
6337 if !crate::tp::len_mirror_lazy_on() {
6338 e.set_i32_one(&mut local.len_d, local.len as i32)?;
6339 }
6340 }
6341 }
6342 if prof {
6343 e.stream().synchronize()?;
6344 prof_ms[1] += seg.elapsed().as_secs_f64() * 1e3;
6345 seg = std::time::Instant::now();
6346 }
6347 let o_out = mixed_t.len() / t;
6348 let mut next = e.uninit(t * n_embd)?;
6349 let mut batched = false;
6350 if ffn_batch && o_out == n_embd {
6351 // MEMRA_WALK_SCRATCH=1: reuse persistent slabs instead of four
6352 // fresh allocations per layer (see WalkScratch). Same kernels,
6353 // same order, same values — only the buffers' provenance changes.
6354 if walk_scratch_on() {
6355 let mut guard = WALK_SCRATCH
6356 .lock()
6357 .map_err(|_| "walk scratch lock poisoned")?;
6358 let dev = e.ctx().ordinal();
6359 if guard
6360 .as_ref()
6361 .is_none_or(|w| w.dev != dev || w.cap < t * n_embd)
6362 {
6363 *guard = Some(WalkScratch {
6364 dev,
6365 cap: 32 * n_embd,
6366 x1: e.uninit(32 * n_embd)?,
6367 z: e.uninit(32 * n_embd)?,
6368 });
6369 }
6370 let w = guard.as_mut().expect("armed above");
6371 // `add_rms_norm` and `step35_verify_moe_tn` both take an
6372 // explicit element count and tolerate a longer slab, so the
6373 // persistent buffers drop straight in: same kernels, same
6374 // order, same values, two fewer allocations per layer.
6375 e.add_rms_norm(
6376 &x_t,
6377 &mixed_t,
6378 layer.post_attn_norm.float_data(),
6379 &mut w.x1,
6380 &mut w.z,
6381 n_embd,
6382 t,
6383 eps,
6384 )?;
6385 if spec_nan_scan_level() >= 2 {
6386 nan_scan_rows(
6387 e,
6388 &w.z,
6389 t,
6390 n_embd,
6391 &format!("tcol layer {il} post-attn norm z"),
6392 )?;
6393 }
6394 if let Some(ffn_t) =
6395 self.step35_verify_moe_tn(e, il, &w.z, t)?
6396 {
6397 if spec_nan_scan_level() >= 2 {
6398 nan_scan_rows(
6399 e,
6400 &ffn_t,
6401 t,
6402 n_embd,
6403 &format!("tcol layer {il} batched routed-MoE out"),
6404 )?;
6405 }
6406 let mut x2_t = e.uninit(t * n_embd)?;
6407 e.add(&w.x1, &ffn_t, &mut x2_t, t * n_embd)?;
6408 next = x2_t;
6409 batched = true;
6410 }
6411 } else {
6412 let mut x1_t = e.uninit(t * n_embd)?;
6413 let mut z_t = e.uninit(t * n_embd)?;
6414 e.add_rms_norm(
6415 &x_t,
6416 &mixed_t,
6417 layer.post_attn_norm.float_data(),
6418 &mut x1_t,
6419 &mut z_t,
6420 n_embd,
6421 t,
6422 eps,
6423 )?;
6424 if spec_nan_scan_level() >= 2 {
6425 nan_scan_rows(
6426 e,
6427 &z_t,
6428 t,
6429 n_embd,
6430 &format!("tcol layer {il} post-attn norm z"),
6431 )?;
6432 }
6433 if let Some(ffn_t) =
6434 self.step35_verify_moe_tn(e, il, &z_t, t)?
6435 {
6436 if spec_nan_scan_level() >= 2 {
6437 nan_scan_rows(
6438 e,
6439 &ffn_t,
6440 t,
6441 n_embd,
6442 &format!("tcol layer {il} batched routed-MoE out"),
6443 )?;
6444 }
6445 let mut x2_t = e.uninit(t * n_embd)?;
6446 e.add(&x1_t, &ffn_t, &mut x2_t, t * n_embd)?;
6447 next = x2_t;
6448 batched = true;
6449 }
6450 }
6451 }
6452 if !batched {
6453 for r in 0..t {
6454 e.dtod_copy_view(
6455 &mixed_t.slice(r * o_out..(r + 1) * o_out),
6456 &mut mixed_row,
6457 )?;
6458 let mut x_row = e.uninit(n_embd)?;
6459 e.dtod_copy_view(
6460 &x_t.slice(r * n_embd..(r + 1) * n_embd),
6461 &mut x_row,
6462 )?;
6463 let (x1, ffn_out) = self.residual_norm_ffn(
6464 e, layer, &x_row, &mixed_row, n_embd, il, eps,
6465 )?;
6466 let mut x2 = e.uninit(n_embd)?;
6467 e.add(&x1, &ffn_out, &mut x2, n_embd)?;
6468 e.dtod_copy_into(&x2, &mut next, r * n_embd)?;
6469 }
6470 }
6471 if prof {
6472 e.stream().synchronize()?;
6473 prof_ms[2] += seg.elapsed().as_secs_f64() * 1e3;
6474 }
6475 x_t = next;
6476 if spec_nan_scan() {
6477 // The scan MUST sit on this arm too. It used to live only on
6478 // the non-fused tail, so a fused layer's poison was first
6479 // reported by the next non-fused layer.
6480 verify_arm_receipt(
6481 "fused",
6482 il,
6483 pos0,
6484 t,
6485 cache.tp_kv[il].as_ref().map(|d| d.staged_len()),
6486 );
6487 nan_scan_rows(
6488 e,
6489 &x_t,
6490 t,
6491 n_embd,
6492 &format!("tcol layer {il} pos0={pos0} arm=fused"),
6493 )?;
6494 }
6495 continue;
6496 }
6497 }
6498 if prof {
6499 e.stream().synchronize()?;
6500 prof_ms[0] += seg.elapsed().as_secs_f64() * 1e3;
6501 seg = std::time::Instant::now();
6502 }
6503 let mut next = e.uninit(t * n_embd)?;
6504 // Columns whose o_proj was deferred (their FFN runs after the join).
6505 // A NON-deferred column's FFN must run INSIDE the column loop: the
6506 // oproj-tail handoff is a single cell that the same column's
6507 // residual_norm_ffn consumes before the next column's finish.
6508 let mut deferred: Vec<usize> = Vec::new();
6509 let mut fa2_deferred: Vec<usize> = Vec::new();
6510 let mut ffn_col =
6511 |r: usize,
6512 mixed: &CudaSlice<f32>,
6513 next: &mut CudaSlice<f32>|
6514 -> Result<(), Box<dyn std::error::Error>> {
6515 let mut x_row = e.uninit(n_embd)?;
6516 e.dtod_copy_view(&x_t.slice(r * n_embd..(r + 1) * n_embd), &mut x_row)?;
6517 let (x1, ffn_out) =
6518 self.residual_norm_ffn(e, layer, &x_row, mixed, n_embd, il, eps)?;
6519 if spec_nan_scan_level() >= 2 {
6520 nan_scan_rows(
6521 e,
6522 &ffn_out,
6523 1,
6524 n_embd,
6525 &format!("tcol layer {il} col {r} per-column FFN out"),
6526 )?;
6527 }
6528 let mut x2 = e.uninit(n_embd)?;
6529 e.add(&x1, &ffn_out, &mut x2, n_embd)?;
6530 e.dtod_copy_into(&x2, next, r * n_embd)?;
6531 Ok(())
6532 };
6533 for r in 0..t {
6534 e.dtod_copy_view(&h_t.slice(r * n_embd..(r + 1) * n_embd), &mut h_row)?;
6535 let row_pos = &pos_rows[r];
6536 crate::tp::set_verify_tcol(Some(r));
6537 if fa2_layer {
6538 crate::tp::set_spec_fa2_defer(Some(r));
6539 } else if oproj_batch {
6540 crate::tp::set_tcol_oproj_defer(Some(r));
6541 }
6542 let mixed = match &layer.mixer {
6543 crate::hybrid::Mixer::Full(fa) => {
6544 self.full_attn_decode(e, fa, &h_row, row_pos, pos0 + r, cache, il)
6545 }
6546 _ => Err("step35 verify expects full attention".into()),
6547 };
6548 crate::tp::set_verify_tcol(None);
6549 crate::tp::set_spec_fa2_defer(None);
6550 crate::tp::set_tcol_oproj_defer(None);
6551 let mixed = mixed?;
6552 if fa2_layer && crate::tp::take_spec_fa2_stashed() {
6553 fa2_deferred.push(r);
6554 } else if oproj_batch && crate::tp::take_tcol_oproj_stashed() {
6555 deferred.push(r);
6556 } else {
6557 if spec_nan_scan_level() >= 2 {
6558 let cols = mixed.len();
6559 nan_scan_rows(
6560 e,
6561 &mixed,
6562 1,
6563 cols,
6564 &format!("tcol layer {il} col {r} per-column ATTN out"),
6565 )?;
6566 }
6567 ffn_col(r, &mixed, &mut next)?;
6568 }
6569 }
6570 if !fa2_deferred.is_empty() && fa2_deferred.len() != t {
6571 // The precheck guarantees both columns stash or neither; a strict
6572 // subset means a column's output was never produced anywhere.
6573 return Err("spec fa2 stash engaged for a subset of columns".into());
6574 }
6575 if prof {
6576 e.stream().synchronize()?;
6577 prof_ms[1] += seg.elapsed().as_secs_f64() * 1e3;
6578 seg = std::time::Instant::now();
6579 }
6580 if !fa2_deferred.is_empty() {
6581 deferred = fa2_deferred;
6582 }
6583 if !deferred.is_empty() {
6584 let mixed_t = if fa2_layer {
6585 self.step35_verify_fa_rows_join(e, il, cache, pos0, t)?
6586 } else {
6587 self.step35_verify_oproj_tcol(e, il, t)?
6588 };
6589 let o_out = mixed_t.len() / t;
6590 if spec_nan_scan_level() >= 2 {
6591 nan_scan_rows(
6592 e,
6593 &mixed_t,
6594 t,
6595 o_out,
6596 &format!("tcol layer {il} JOINED attn over deferred cols"),
6597 )?;
6598 }
6599 // Batched t=2 residual+MoE: one t-grid add_rms_norm (per-row
6600 // program == t=1; bit-identical to the oproj-tail join per the
6601 // M2 verbatim-program contract) feeding the two-column routed
6602 // sweep. Ineligible layers (dense FFN, non-nvfp4) fall through
6603 // to the per-column body.
6604 let mut batched = false;
6605 if ffn_batch && deferred.len() == t && o_out == n_embd {
6606 let mut x1_t = e.uninit(t * n_embd)?;
6607 let mut z_t = e.uninit(t * n_embd)?;
6608 e.add_rms_norm(
6609 &x_t,
6610 &mixed_t,
6611 layer.post_attn_norm.float_data(),
6612 &mut x1_t,
6613 &mut z_t,
6614 n_embd,
6615 t,
6616 eps,
6617 )?;
6618 if let Some(ffn_t) = self.step35_verify_moe_tn(e, il, &z_t, t)? {
6619 let mut x2_t = e.uninit(t * n_embd)?;
6620 e.add(&x1_t, &ffn_t, &mut x2_t, t * n_embd)?;
6621 next = x2_t;
6622 batched = true;
6623 }
6624 }
6625 if !batched {
6626 for &r in &deferred {
6627 e.dtod_copy_view(
6628 &mixed_t.slice(r * o_out..(r + 1) * o_out),
6629 &mut mixed_row,
6630 )?;
6631 ffn_col(r, &mixed_row, &mut next)?;
6632 }
6633 }
6634 }
6635 if prof {
6636 e.stream().synchronize()?;
6637 prof_ms[2] += seg.elapsed().as_secs_f64() * 1e3;
6638 }
6639 drop(ffn_col);
6640 x_t = next;
6641 if spec_nan_scan() {
6642 verify_arm_receipt(
6643 if fa2_layer { "join" } else { "percol" },
6644 il,
6645 pos0,
6646 t,
6647 cache.tp_kv[il].as_ref().map(|d| d.staged_len()),
6648 );
6649 nan_scan_rows(
6650 e,
6651 &x_t,
6652 t,
6653 n_embd,
6654 &format!(
6655 "tcol layer {il} pos0={pos0} arm={}",
6656 if fa2_layer { "join" } else { "percol" }
6657 ),
6658 )?;
6659 }
6660 }
6661 if prof {
6662 eprintln!(
6663 "[tcol-prof] t={t} norm+qkv={:.3}ms attn={:.3}ms ffn={:.3}ms",
6664 prof_ms[0], prof_ms[1], prof_ms[2]
6665 );
6666 }
6667 if ok {
6668 return Ok(x_t);
6669 }
6670 // fall through to the row-outer walk on ineligible layers
6671 x = x_t;
6672 }
6673 let mut next = e.uninit(t * n_embd)?;
6674 let scan = spec_nan_scan();
6675 for r in 0..t {
6676 let mut row = e.uninit(n_embd)?;
6677 e.dtod_copy_view(&x.slice(r * n_embd..(r + 1) * n_embd), &mut row)?;
6678 let row_pos = e.htod_i32(&[(pos0 + r) as i32])?;
6679 let out = if scan {
6680 // Diagnostic arm: the same range walked one layer at a time so the first
6681 // poisoned layer names itself. `decode_layers_eager(lo, hi)` is range-scoped
6682 // and executes its trailing residual add, so a per-layer chain is the same
6683 // program with the cross-layer add+norm fusion unrolled.
6684 nan_scan_rows(
6685 e,
6686 &row,
6687 1,
6688 n_embd,
6689 &format!("embed row r={r} pos={}", pos0 + r),
6690 )?;
6691 let mut acc = row;
6692 for il in lo..hi {
6693 acc = self.decode_layers_eager(
6694 e,
6695 acc,
6696 il,
6697 il + 1,
6698 &row_pos,
6699 pos0 + r,
6700 cache,
6701 )?;
6702 nan_scan_rows(
6703 e,
6704 &acc,
6705 1,
6706 n_embd,
6707 &format!("row-outer layer {il} r={r} pos={}", pos0 + r),
6708 )?;
6709 }
6710 acc
6711 } else {
6712 self.decode_layers_eager(e, row, lo, hi, &row_pos, pos0 + r, cache)?
6713 };
6714 e.dtod_copy_into(&out, &mut next, r * n_embd)?;
6715 }
6716 // dflash taps are NOT produced on this arm (they need per-layer hiddens the
6717 // row-outer walk does not materialize); the door is a step37 MTP bring-up
6718 // surface where taps are unused.
6719 return Ok(next);
6720 }
6721 let mut ph_last = std::time::Instant::now();
6722 for il in lo..hi {
6723 let mut next = e.uninit(t * n_embd)?;
6724 for r in 0..t {
6725 let mut row = e.uninit(n_embd)?;
6726 e.dtod_copy_view(&x.slice(r * n_embd..(r + 1) * n_embd), &mut row)?;
6727 // The caller owns this verify's position. During controller overlap, cache.pos
6728 // still describes generation N while this stage-0 walk belongs to N+1.
6729 let row_pos = e.htod_i32(&[(pos0 + r) as i32])?;
6730 let mut one = [&mut *cache];
6731 let out = self.step35_decode_batch_layers(
6732 e,
6733 row,
6734 &mut one,
6735 &[(pos0 + r) as i32],
6736 &row_pos,
6737 il,
6738 il + 1,
6739 &mut ph_last,
6740 )?;
6741 e.dtod_copy_into(&out, &mut next, r * n_embd)?;
6742 }
6743 self.dflash_tap(e, cache, il, &next, t)?;
6744 x = next;
6745 if spec_nan_scan() {
6746 nan_scan_rows(e, &x, t, n_embd, &format!("batch-layer {il} pos0={pos0}"))?;
6747 }
6748 }
6749 Ok(x)
6750 }
6751
6752 /// DSpark drafter verify (lane/dspark-q38-recover): one t-row forward through the
6753 /// SERVING-CLASS verify funnel (`decode_step_t_core_stream` — the same numeric class
6754 /// MTP verify rides, GDN state advanced in place), returning per-row argmax tokens.
6755 /// Advances `cache.pos += t`; the caller owns snapshot/rollback (block acceptance is
6756 /// prefix-keep, not all-or-nothing).
6757 pub(crate) fn dspark_verify_t_am(
6758 &self,
6759 e: &Engine,
6760 tokens: &[u32],
6761 pos0: usize,
6762 cache: &mut Cache,
6763 ) -> Result<Vec<u32>, Box<dyn std::error::Error>> {
6764 let (logits, _hn) = self.decode_step_t_core_stream(
6765 e, tokens, pos0, cache, None, None, None, None, None, None,
6766 )?;
6767 let t = tokens.len();
6768 let v = self.output.out_features();
6769 let mut am_d = e.stream().alloc_zeros::<u32>(t)?;
6770 for r in 0..t {
6771 e.argmax_token_device_col(&logits, r, v, &mut am_d, r)?;
6772 }
6773 Ok(e.dtoh_u32(&am_d)?)
6774 }
6775
6776 /// DSpark verify returning the RAW verify logits [t, n_vocab] (device-resident) instead
6777 /// of per-row argmaxes — the sampled-admission arm's input (rejection-sampling accept
6778 /// gathers filtered p from these columns; lane/dspark-sampled-admission-20260820). Same
6779 /// forward as `dspark_verify_t_am`; the greedy arm keeps its argmax wrapper untouched.
6780 pub(crate) fn dspark_verify_t_logits(
6781 &self,
6782 e: &Engine,
6783 tokens: &[u32],
6784 pos0: usize,
6785 cache: &mut Cache,
6786 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6787 let (logits, _hn) = self.decode_step_t_core_stream(
6788 e, tokens, pos0, cache, None, None, None, None, None, None,
6789 )?;
6790 Ok(logits)
6791 }
6792
6793 /// DSpark verify with the MTP column-stash armed: identical forward to
6794 /// `dspark_verify_t_am`, but fills a `VerifyCkpt` so a partial accept can restore
6795 /// column state directly (`dspark_commit_prefix`) instead of snapshot-replay.
6796 /// The ckpt type is opaque outside spec.rs (newtype) — dflash.rs threads it through.
6797 pub(crate) fn dspark_verify_t_am_ckpt(
6798 &self,
6799 e: &Engine,
6800 tokens: &[u32],
6801 pos0: usize,
6802 cache: &mut Cache,
6803 ) -> Result<(Vec<u32>, DsparkVerifyCkpt), Box<dyn std::error::Error>> {
6804 let mut ck = VerifyCkpt::new(self.layers.len());
6805 let (logits, _hn) = self.decode_step_t_core_stream(
6806 e,
6807 tokens,
6808 pos0,
6809 cache,
6810 None,
6811 Some(&mut ck),
6812 None,
6813 None,
6814 None,
6815 None,
6816 )?;
6817 let t = tokens.len();
6818 let v = self.output.out_features();
6819 let mut am_d = e.stream().alloc_zeros::<u32>(t)?;
6820 for r in 0..t {
6821 e.argmax_token_device_col(&logits, r, v, &mut am_d, r)?;
6822 }
6823 Ok((e.dtoh_u32(&am_d)?, DsparkVerifyCkpt(ck)))
6824 }
6825
6826 /// Engine-bundle slice 2: `dspark_verify_t_am_ckpt` with DEVICE tokens and NO readback.
6827 /// The verify tokens are the round's `chain_d` (cand layout: [anchor, drafts...]); the
6828 /// embed gathers its first `t` entries on-device (`embed_gather_u32_t` — bit-identical
6829 /// rows to the host arm), so the host never blocks on the draft chain before dispatching
6830 /// verify. Returns the device per-row argmax buffer; the caller merges its readback with
6831 /// the chain's into ONE sync. Forward, ckpt fill and argmax walk are `_ckpt` verbatim.
6832 pub(crate) fn dspark_verify_t_am_ckpt_dev(
6833 &self,
6834 e: &Engine,
6835 vtok: &CudaSlice<u32>,
6836 t: usize,
6837 pos0: usize,
6838 cache: &mut Cache,
6839 embd_dev: (&CudaSlice<u8>, i32, usize),
6840 graphs: Option<&mut DsparkVerifyGraphs>,
6841 ) -> Result<(CudaSlice<u32>, DsparkVerifyCkpt), Box<dyn std::error::Error>> {
6842 debug_assert!(
6843 vtok.len() >= t,
6844 "verify window exceeds the device token buffer"
6845 );
6846 // The slab flag is a per-round statement: clear it here so a verify that never
6847 // reaches the graphs door (rowwise env, a non-tparallel arm) cannot leave a
6848 // stale `true` steering the commit at slabs the round never wrote.
6849 let mut graphs = graphs;
6850 if let Some(g) = graphs.as_deref_mut() {
6851 g.round_slab = false;
6852 }
6853 let mut ck = VerifyCkpt::new(self.layers.len());
6854 // Dummy host tokens size the funnel; the embed reads `vtok` (the round-stream
6855 // arm's established pattern — spec.rs stream-mode verify does the same).
6856 let dummy = vec![0u32; t];
6857 let (logits, _hn) = self.decode_step_t_core_stream(
6858 e,
6859 &dummy,
6860 pos0,
6861 cache,
6862 Some(embd_dev),
6863 Some(&mut ck),
6864 None,
6865 None,
6866 Some(vtok),
6867 graphs,
6868 )?;
6869 let v = self.output.out_features();
6870 let mut am_d = e.stream().alloc_zeros::<u32>(t)?;
6871 for r in 0..t {
6872 e.argmax_token_device_col(&logits, r, v, &mut am_d, r)?;
6873 }
6874 Ok((am_d, DsparkVerifyCkpt(ck)))
6875 }
6876
6877 /// Ckpt-armed twin of [`Self::dspark_verify_t_logits`] (sampled-admission arm).
6878 pub(crate) fn dspark_verify_t_logits_ckpt(
6879 &self,
6880 e: &Engine,
6881 tokens: &[u32],
6882 pos0: usize,
6883 cache: &mut Cache,
6884 ) -> Result<(CudaSlice<f32>, DsparkVerifyCkpt), Box<dyn std::error::Error>> {
6885 let mut ck = VerifyCkpt::new(self.layers.len());
6886 let (logits, _hn) = self.decode_step_t_core_stream(
6887 e,
6888 tokens,
6889 pos0,
6890 cache,
6891 None,
6892 Some(&mut ck),
6893 None,
6894 None,
6895 None,
6896 None,
6897 )?;
6898 Ok((logits, DsparkVerifyCkpt(ck)))
6899 }
6900
6901 /// Restore the round to `keep` accepted columns from the verify stash: KV lens and
6902 /// pos from the pre-verify snapshot + keep, GDN conv/ssm from the stashed column
6903 /// state — no replay forward. The exact `commit_verified_prefix` the MTP path ships.
6904 pub(crate) fn dspark_commit_prefix(
6905 &self,
6906 e: &Engine,
6907 cache: &mut Cache,
6908 snap: &crate::cache::CacheSnapshot,
6909 ckpt: &DsparkVerifyCkpt,
6910 keep: usize,
6911 ) -> Result<(), Box<dyn std::error::Error>> {
6912 self.commit_verified_prefix(e, cache, snap, &ckpt.0, keep, false, None)
6913 }
6914
6915 /// Slice-3 commit twin: restore to `keep` accepted columns when the round's linear
6916 /// column stash lives in the graphs ctx's persistent slabs (`DsparkVerifyGraphs`) —
6917 /// the cols arm's exact semantics (KV lens + pos from the snapshot, GDN conv/ssm
6918 /// from the stash of column keep-1), slab-addressed and batched into two copy
6919 /// launches. `MEMRA_STATE_COPY_BATCH=0` falls back to per-layer view copies.
6920 pub(crate) fn dspark_commit_prefix_slab(
6921 &self,
6922 e: &Engine,
6923 cache: &mut Cache,
6924 snap: &crate::cache::CacheSnapshot,
6925 ctx: &DsparkVerifyGraphs,
6926 keep: usize,
6927 ) -> Result<(), Box<dyn std::error::Error>> {
6928 use cudarc::driver::DevicePtr;
6929 debug_assert!(keep >= 1, "keep==0 rounds take the legacy rollback");
6930 let mut conv_src: Vec<u64> = Vec::new();
6931 let mut ssm_src: Vec<u64> = Vec::new();
6932 let mut conv_dst: Vec<u64> = Vec::new();
6933 let mut ssm_dst: Vec<u64> = Vec::new();
6934 for il in 0..self.layers.len() {
6935 if let (Some(kvl), Some(saved)) = (cache.kv[il].as_mut(), snap.kv_len[il]) {
6936 kvl.len = saved + keep;
6937 e.set_i32_one(&mut kvl.len_d, kvl.len as i32)?;
6938 }
6939 if let Some(rl) = cache.recur[il].as_ref() {
6940 let (pc, ps, _cw, _sw) = ctx
6941 .slab_row(e, il, keep - 1)
6942 .ok_or("slab commit: linear layer missing from the graphs ctx")?;
6943 conv_src.push(pc);
6944 ssm_src.push(ps);
6945 let st = &e.gpu.stream();
6946 let (dc, _g0) = rl.conv_state.device_ptr(st);
6947 let (ds, _g1) = rl.ssm_state.device_ptr(st);
6948 conv_dst.push(dc as u64);
6949 ssm_dst.push(ds as u64);
6950 }
6951 }
6952 let n = conv_src.len();
6953 if n > 0 {
6954 if state_copy_batch_on() {
6955 let mut tt = vec![0u64; 2 * n];
6956 tt[..n].copy_from_slice(&conv_src);
6957 tt[n..].copy_from_slice(&conv_dst);
6958 let ct = e.htod_u64(&tt)?;
6959 tt[..n].copy_from_slice(&ssm_src);
6960 tt[n..].copy_from_slice(&ssm_dst);
6961 let st = e.htod_u64(&tt)?;
6962 e.copy_batch_uniform_f32(&ct, n, ctx.conv_words)?;
6963 e.copy_batch_uniform_f32(&st, n, ctx.ssm_words)?;
6964 } else {
6965 let (cw, sw) = (ctx.conv_words, ctx.ssm_words);
6966 let row = keep - 1;
6967 for il in 0..self.layers.len() {
6968 let Some(rl) = cache.recur[il].as_mut() else {
6969 continue;
6970 };
6971 let k = ctx.lin_pos[&il];
6972 {
6973 let sv = e.view(&ctx.stash_conv[k], (row + 1) * cw);
6974 let win = sv.slice(row * cw..(row + 1) * cw);
6975 e.copy_view_into(&mut rl.conv_state, 0, &win, cw)?;
6976 }
6977 {
6978 let sv = e.view(&ctx.stash_ssm[k], (row + 1) * sw);
6979 let win = sv.slice(row * sw..(row + 1) * sw);
6980 e.copy_view_into(&mut rl.ssm_state, 0, &win, sw)?;
6981 }
6982 }
6983 }
6984 }
6985 cache.pos = snap.pos + keep;
6986 Ok(())
6987 }
6988
6989 /// Qwen35-family verify trunk in the live serving numeric class.
6990 ///
6991 /// Serving intentionally keeps this architecture in the generic batched program even at
6992 /// B=1. The older verify walk used its own mirrored dispatch and can flip near-tie argmaxes.
6993 ///
6994 /// Two arms, one numeric class:
6995 /// - DENSE GDN (`DenseMlp`, t<=16): `qwen35_verify_tparallel` — the weight ops (norms,
6996 /// projections, FFN) hoist to m=T through the exact-tier batched kernels whose per-row
6997 /// program IS the m=1 program (`matmul_pre == fused2 per (tensor,row); _bN mmvq per-row
6998 /// == m=1` — decode_batch.rs v2 note), while the state ops (conv ring, gdn scan, KV
6999 /// append, fa decode) stay a per-row loop running the b_n=1 serving kernels with each
7000 /// row's own t_kv-driven arm pick (the straddle law: every row executes the exact
7001 /// program its isolated serving step would). One weight read per layer per round
7002 /// instead of T — this is what makes MTP profitable in the exact class (the per-row
7003 /// walk measured verify(K+1) ~= (K+1) plain steps: 69 -> 44 tok/s served, 2026-08-15).
7004 /// - MoE / t>16 / `MEMRA_SPEC_VERIFY_ROWWISE=1`: the per-row replay of the authoritative
7005 /// serving layer body, preserving single-session autoregressive cache order (the
7006 /// correctness reference; also the rollback seam for the t-parallel arm).
7007 ///
7008 /// Bit-identity of the t-parallel arm vs the rowwise arm is gated by spec-serve-gate
7009 /// (zero differing logits at T=1..4, K arms) + the 8-prompt ON/OFF canary before ship.
7010 #[allow(clippy::too_many_arguments)]
7011 fn qwen35_verify_batch_layers(
7012 &self,
7013 e: &Engine,
7014 x: CudaSlice<f32>,
7015 lo: usize,
7016 hi: usize,
7017 pos0: usize,
7018 t: usize,
7019 cache: &mut Cache,
7020 ckpt: Option<&mut VerifyCkpt>,
7021 stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
7022 graphs: Option<&mut DsparkVerifyGraphs>,
7023 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7024 // Qwen35Moe admitted 2026-08-20 (lane/draftcost-moe): the t-parallel arm already
7025 // carries the MoE FFN (`moe_ffn_il_zq8` at m=T) and the GDN per-row state loop; the
7026 // arch fence was a qualification gate, not a mechanism gap. Measured disease on the
7027 // 35B-A3B class: rowwise verify ~= 5.6 ms per drafted token (one full trunk step
7028 // each) — the same (K+1)-plain-steps wall the dense admission fixed on 2026-08-15.
7029 // Rollback seam unchanged: MEMRA_SPEC_VERIFY_ROWWISE=1.
7030 let rowwise = std::env::var("MEMRA_SPEC_VERIFY_ROWWISE").as_deref() == Ok("1")
7031 || !self.batched_serving_numeric_class()
7032 || t > 16;
7033 if rowwise {
7034 if stream.is_some() {
7035 // rowwise replays per row with host cache.pos — irreconcilable with a
7036 // device position counter. Burst callers must keep t <= 16 and the
7037 // ROWWISE env unset; refusing beats silently mispositioned rows.
7038 return Err("qwen35 rowwise verify has no ROUND-STREAM arm \
7039 (t > 16 or MEMRA_SPEC_VERIFY_ROWWISE=1)"
7040 .into());
7041 }
7042 self.qwen35_verify_rowwise(e, x, lo, hi, pos0, t, cache, ckpt)
7043 } else {
7044 self.qwen35_verify_tparallel(e, x, lo, hi, pos0, t, cache, ckpt, stream, graphs)
7045 }
7046 }
7047
7048 /// The per-row correctness reference: replay each verify row through the authoritative
7049 /// serving layer body (`decode_batch_layers` at b_n=1). T full weight reads per layer.
7050 #[allow(clippy::too_many_arguments)]
7051 fn qwen35_verify_rowwise(
7052 &self,
7053 e: &Engine,
7054 mut x: CudaSlice<f32>,
7055 lo: usize,
7056 hi: usize,
7057 pos0: usize,
7058 t: usize,
7059 cache: &mut Cache,
7060 mut ckpt: Option<&mut VerifyCkpt>,
7061 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7062 let n_embd = self.cfg.n_embd as usize;
7063 let saved_pos = cache.pos;
7064 let mut ph_last = std::time::Instant::now();
7065 for il in lo..hi {
7066 let mut next = e.uninit(t * n_embd)?;
7067 let mut col_states: Option<Vec<(CudaSlice<f32>, CudaSlice<f32>)>> =
7068 if ckpt.is_some() && t >= 2 && matches!(self.layers[il].mixer, Mixer::Linear(_)) {
7069 Some(Vec::with_capacity(t - 1))
7070 } else {
7071 None
7072 };
7073 for r in 0..t {
7074 cache.pos = pos0 + r;
7075 let mut row = e.uninit(n_embd)?;
7076 e.dtod_copy_view(&x.slice(r * n_embd..(r + 1) * n_embd), &mut row)?;
7077 let row_pos = e.htod_i32(&[(pos0 + r) as i32])?;
7078 let mut one = [&mut *cache];
7079 let ctx = self.batch_layer_ctx(e, &one, il, il + 1)?;
7080 let out = match self.decode_batch_layers(
7081 e,
7082 row,
7083 &mut one,
7084 &ctx,
7085 &row_pos,
7086 &mut ph_last,
7087 ) {
7088 Ok(out) => out,
7089 Err(error) => {
7090 cache.pos = saved_pos;
7091 return Err(error);
7092 }
7093 };
7094 e.dtod_copy_into(&out, &mut next, r * n_embd)?;
7095 if r + 1 < t {
7096 if let Some(states) = col_states.as_mut() {
7097 let recur = cache.recur[il]
7098 .as_ref()
7099 .ok_or("Qwen35-MoE linear verify layer has no recurrent state")?;
7100 states.push((
7101 e.clone_dtod(&recur.conv_state)?,
7102 e.clone_dtod(&recur.ssm_state)?,
7103 ));
7104 }
7105 }
7106 }
7107 if let (Some(checkpoint), Some(states)) = (ckpt.as_deref_mut(), col_states) {
7108 checkpoint.cols[il] = Some(states);
7109 }
7110 x = next;
7111 }
7112 cache.pos = saved_pos;
7113 Ok(x)
7114 }
7115
7116 /// T-PARALLEL VERIFY IN THE SERVING NUMERIC CLASS (lane/tparallel-verify, 2026-08-15).
7117 ///
7118 /// The weight ops run ONCE per layer at m=T; the state ops run per row through the same
7119 /// b_n=1 serving kernels the rowwise replay uses. Per-row bit-identity rests on the two
7120 /// pins the serving batch tier already carries:
7121 /// * `matmul_pre` / `_bN` mmvq: per-row program == m=1 program (decode_batch.rs v2 note,
7122 /// kernel-check pinned) — so a [T, n_embd] projection row equals the row projected
7123 /// alone;
7124 /// * row-indexed norms/elementwise (`rms_norm`, `quantize_q8_1`, `add_rms_norm`,
7125 /// `gated_rmsnorm[_q8_1]`, `silu_mul`, `rope_neox` with per-row positions): the T-row
7126 /// launch is the per-row program (same pin the generic verify's fused norms rely on).
7127 /// The sequential dependencies keep their exact serving order: the conv ring / gdn scan
7128 /// chain state row -> row through the `_b` kernels at b_n=1 (ping-pong via a 6-entry
7129 /// alternating pointer table, host handles swapped per row so VerifyCkpt clones the
7130 /// canonical state exactly as the rowwise arm does), and each row's KV append + fa decode
7131 /// picks its arm from ITS OWN t_kv (append: format-only; fa: `fa_seqs_eligible` + its own
7132 /// `fa_split_keys` rung at b_n=1) — the straddle law per row, so every row executes the
7133 /// program its isolated B=1 serving step would.
7134 ///
7135 /// Cost: 1 weight read per layer per round + T state micro-launches, vs the rowwise arm's
7136 /// T weight reads. Gated bit-identical vs the rowwise arm by spec-serve-gate + canary.
7137 #[allow(clippy::too_many_arguments)]
7138 fn qwen35_verify_tparallel(
7139 &self,
7140 e: &Engine,
7141 mut x: CudaSlice<f32>,
7142 lo: usize,
7143 hi: usize,
7144 pos0: usize,
7145 t: usize,
7146 cache: &mut Cache,
7147 mut ckpt: Option<&mut VerifyCkpt>,
7148 stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
7149 mut graphs: Option<&mut DsparkVerifyGraphs>,
7150 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7151 let seqs_append =
7152 std::env::var("MEMRA_BATCH_APPEND").as_deref() != Ok("0") && !Engine::kv_fp8_on();
7153 let batch_fa_on = std::env::var("MEMRA_BATCH_FA").as_deref() != Ok("0");
7154
7155 // Merge guard (v0.98 train, re-affirmed on the v0.100 train over slice 4c): the
7156 // ROUND-STREAM arm (lane/draftcost-moe, device position counter) and the dspark
7157 // verify graphs (engine-bundle slice 3 / trunk slice 4c) have no common caller —
7158 // stream rides the qwen35moe burst, graphs ride the dspark route. If a future
7159 // caller arms both, refuse loudly instead of silently dropping the graphs ctx
7160 // (the stream linear arm takes linear_attn_verify_t, not the graphed segment or
7161 // full-verify bodies).
7162 if stream.is_some() && graphs.is_some() {
7163 return Err(
7164 "qwen35 tparallel verify: ROUND-STREAM and dspark verify graphs \
7165 cannot arm together"
7166 .into(),
7167 );
7168 }
7169 // Engine-bundle slice 3 + slice 4c: with a graphs ctx armed, pointer tables are
7170 // refreshed once per verify (the gdn ping-pong moves handles; a fresh generation
7171 // moves the kv caches). Then:
7172 // - slice 4c: when the WHOLE round rides one seqs rung (every row batchable, one
7173 // split-ladder step, rung covers the round), the ENTIRE walk replays as ONE
7174 // full-verify graph per (vt, rung) — linear layers through the shared
7175 // `qwen35_tparallel_linear_layer` body, full-attention layers through the
7176 // shared `qwen35_tparallel_fa_layer` body in graph mode.
7177 // - fallback (straddle rounds, below the vec floor, partial walks): runs of
7178 // consecutive LINEAR layers replay the slice-3 per-(segment, vt) graphs and
7179 // the full-attention layers run eager (batched rows when eligible).
7180 if let Some(g) = graphs.as_deref_mut() {
7181 g.refresh_tables(e, cache)?;
7182 g.round_slab = false;
7183 if let Some(rung) = g.full_rung(self, cache, lo, hi, t, seqs_append && batch_fa_on) {
7184 // Pool ceiling (dspark_vg_cap): an existing key always replays; a NEW
7185 // full capture past the ceiling falls through to the segment/eager arms.
7186 if g.full.contains_key(&(t, rung, hi)) || g.can_capture() {
7187 let out = g.run_full(self, e, lo, hi, &x, t, pos0, rung, cache)?;
7188 g.round_slab = true;
7189 return Ok(out);
7190 }
7191 }
7192 // Round-atomic ceiling check for the segment door: if any linear run in this
7193 // walk would need a NEW capture past the ceiling, the whole round runs the
7194 // eager cols-ckpt walk (mixing slab- and cols-stashed layers in one round
7195 // would corrupt the commit).
7196 if !g.segments_ready(self, lo, hi, t) {
7197 graphs = None;
7198 }
7199 }
7200 // STREAM (2b, lane/draftcost-moe): positions come from the device round counter
7201 // (pos_iota / i32_copy_add) so a burst round needs no host position knowledge.
7202 let pos_d = match stream {
7203 Some((_, ctr)) => {
7204 let mut p = e.alloc_uninit::<i32>(t)?;
7205 e.pos_iota(ctr, &mut p, t)?;
7206 p
7207 }
7208 None => {
7209 let pos_host: Vec<i32> = (0..t).map(|r| (pos0 + r) as i32).collect();
7210 e.htod_i32(&pos_host)?
7211 }
7212 };
7213 // Per-row 1-element position buffers, built ONCE per verify (the append/fa wrappers
7214 // take owned pos slices; building these inside the layer x row loops cost 16xT H2Ds).
7215 // LAZY since slice 4: the batched fa/append arm never touches them — they are built
7216 // on the first per-row fallback layer only (stream-aware there; the stream FA arm
7217 // rides the dc rows kernels and never reaches the fallback).
7218 let mut pos_rows: Option<Vec<CudaSlice<i32>>> = None;
7219 let mut il = lo;
7220 while il < hi {
7221 if graphs.is_some() && matches!(self.layers[il].mixer, Mixer::Linear(_)) {
7222 let mut end = il;
7223 while end < hi && matches!(self.layers[end].mixer, Mixer::Linear(_)) {
7224 end += 1;
7225 }
7226 let g = graphs.as_deref_mut().expect("checked above");
7227 x = g.run_segment(self, e, il, end, &x, t, cache)?;
7228 g.round_slab = true;
7229 il = end;
7230 continue;
7231 }
7232 let layer = &self.layers[il];
7233 if stream.is_none() && matches!(layer.mixer, Mixer::Linear(_)) {
7234 // Eager linear layer (no graphs ctx): the shared body, legacy cols-ckpt arm.
7235 // Under ROUND-STREAM the linear layers ride the fa-body match's stream arm
7236 // below (linear_attn_verify_t — the stream COMMIT needs its GdnStash).
7237 x = self.qwen35_tparallel_linear_layer(
7238 e,
7239 il,
7240 &x,
7241 t,
7242 cache,
7243 ckpt.as_deref_mut(),
7244 None,
7245 None,
7246 )?;
7247 il += 1;
7248 continue;
7249 }
7250 // Full-attention (or stream-Linear, or MLA-refusing) layer: the extracted
7251 // shared body — eager arm (fresh per-verify pos/table, exact t_kv sizing,
7252 // in-body len bump). The slice-4c captured full-verify graphs run the SAME
7253 // body in graph mode; under ROUND-STREAM the body's dc-rows / GDN stream arms
7254 // run (lane/draftcost-moe).
7255 x = self.qwen35_tparallel_fa_layer(
7256 e,
7257 il,
7258 &x,
7259 t,
7260 cache,
7261 FaLayerArgs {
7262 pos_d: &pos_d,
7263 pos_rows: &mut pos_rows,
7264 pos0,
7265 seqs_append,
7266 batch_fa_on,
7267 graph_cap: None,
7268 stream,
7269 ckpt: ckpt.as_deref_mut(),
7270 },
7271 )?;
7272 il += 1;
7273 }
7274 Ok(x)
7275 }
7276
7277 /// SHARED dense-FFN body for the qwen35 t-parallel layers (trunk-kernels slice B) —
7278 /// ONE copy for the fa and linear layer bodies (the verify_layers extraction lesson).
7279 /// Dual arm (MEMRA_TK_FFN_DUAL, default on): gate+up in ONE dual launch from the
7280 /// pre-quantized activation with macro-scales DEFERRED into the fused SwiGLU+q8_1
7281 /// epilogue, then ffn_down from the fused (aq, ad) — the q27 verify chain verbatim.
7282 /// Every door is the bit-identical proven one: `matmul_decode_exact_dual_pre` (per
7283 /// (tensor,token,row) == the two singles), `silu_mul_scaled_q8_1` (y*s inline == the
7284 /// scale_inplace store, value-exact; fused quantize == quantize_q8_1 bytes),
7285 /// `matmul_decode_exact_pre` (dispatch mirror of the singles' q8_1-fast tail).
7286 /// Dual-refused (t outside 2..=7, non-NVFP4, layout mismatch) or seam off -> the
7287 /// original singles chain, byte-for-byte.
7288 #[allow(clippy::too_many_arguments)]
7289 fn qwen35_tparallel_dense_ffn(
7290 &self,
7291 e: &Engine,
7292 ffn_gate: &crate::model::GpuTensor,
7293 ffn_up: &crate::model::GpuTensor,
7294 ffn_down: &crate::model::GpuTensor,
7295 zn: &CudaSlice<f32>,
7296 t: usize,
7297 n_embd: usize,
7298 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7299 let n_ff = ffn_gate.out_features();
7300 let (zq, zd) = e.quantize_q8_1(zn, t, n_embd)?;
7301 if Engine::tk_ffn_dual_on() {
7302 if let Some(((g, gs), (u, us))) =
7303 e.matmul_decode_exact_dual_pre(ffn_gate, ffn_up, &zq, &zd, t)?
7304 {
7305 if e.uses_q8_1_fast(ffn_down) {
7306 let (aq, ad) = e.silu_mul_scaled_q8_1(&g, &u, gs, us, t * n_ff)?;
7307 return e.matmul_decode_exact_pre(ffn_down, &aq, &ad, t);
7308 }
7309 let mut act = e.uninit(t * n_ff)?;
7310 e.silu_mul_scaled(&g, &u, gs, us, &mut act, t * n_ff)?;
7311 let (aq, ad) = e.quantize_q8_1(&act, t, n_ff)?;
7312 return e.matmul_pre(ffn_down, &aq, &ad, &act, t);
7313 }
7314 }
7315 // v1 singles chain (seam off or dual-refused) — the pre-slice-B body verbatim.
7316 let g = e.matmul_pre(ffn_gate, &zq, &zd, zn, t)?;
7317 let u = e.matmul_pre(ffn_up, &zq, &zd, zn, t)?;
7318 let mut act = e.uninit(t * n_ff)?;
7319 e.silu_mul(&g, &u, &mut act, t * n_ff)?;
7320 let (aq, ad) = e.quantize_q8_1(&act, t, n_ff)?;
7321 e.matmul_pre(ffn_down, &aq, &ad, &act, t)
7322 }
7323
7324 /// ONE t-parallel FULL-ATTENTION layer (attn_norm + fa mixer + post_attn_norm + FFN +
7325 /// tap) — extracted from the walk exactly like `qwen35_tparallel_linear_layer` so the
7326 /// eager walk and the slice-4c captured full-verify graphs execute the SAME body (a
7327 /// second copy is how dispatch mirrors drift — the verify_layers extraction lesson).
7328 ///
7329 /// `args.graph_cap = Some((table, off, rung_end))` is the captured-graph mode:
7330 /// - kv base-pointer pairs come from the ctx-owned persistent table at `off` (a fresh
7331 /// generation's cache lands at new addresses that only the per-verify table refresh
7332 /// knows — the slice-3 baked-address lesson);
7333 /// - the seqs twins size partials/grid at `rung_end` and pin `split_keys` to the
7334 /// rung's ladder value: `n_splits_max` is pure stride, splits >= ns_eff write the
7335 /// EMPTY partial the combine never reads, and every per-row T_kv derives in-kernel
7336 /// from `pos_seq[z]` — so one captured launch replays bit-identically for every
7337 /// round whose rows all sit inside the rung;
7338 /// - the host len bump moves to the replay caller (captured host code does not
7339 /// re-run at replay).
7340 /// Graph mode REFUSES any round the batched arm cannot take: the per-row fallback
7341 /// host-branches on t_kv and must never be captured.
7342 #[allow(clippy::too_many_arguments)]
7343 fn qwen35_tparallel_fa_layer(
7344 &self,
7345 e: &Engine,
7346 il: usize,
7347 x: &CudaSlice<f32>,
7348 t: usize,
7349 cache: &mut Cache,
7350 args: FaLayerArgs<'_>,
7351 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7352 use cudarc::driver::DevicePtr;
7353 let cfg = &self.cfg;
7354 let n_embd = cfg.n_embd as usize;
7355 let eps = cfg.rms_eps;
7356 let head_dim_global = cfg.head_dim_k as usize;
7357 let layer = &self.layers[il];
7358 let FaLayerArgs {
7359 pos_d,
7360 pos_rows,
7361 pos0,
7362 seqs_append,
7363 batch_fa_on,
7364 graph_cap,
7365 stream,
7366 mut ckpt,
7367 } = args;
7368
7369 // ---- attn_norm + q8_1 quantize at m=T (row-indexed == per-row) ----
7370 let anorm = layer.attn_norm.float_data();
7371 let mut xn = e.uninit(t * n_embd)?;
7372 e.rms_norm(x, anorm, &mut xn, n_embd, t, eps)?;
7373 let (hq, hd) = e.quantize_q8_1(&xn, t, n_embd)?;
7374
7375 let mixed: CudaSlice<f32> = match &layer.mixer {
7376 Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
7377 // STREAM ARM (2b, lane/draftcost-moe): under a device position counter the
7378 // per-row serving-kernel chain cannot run (host state swaps keyed on host
7379 // row index are fine, but the stream COMMIT needs the GdnStash for its _dc
7380 // rebuild — the per-row chain only produces per-column clones). GDN rides
7381 // `linear_attn_verify_t`: batched q8_1-class projections, stash-producing,
7382 // and its one-scan recurrence is pinned bit-identical to T chained T=1
7383 // steps (its header + kernel-check). Position-independent, so no counter
7384 // plumbing is needed. Guards mirror the generic call site exactly.
7385 Mixer::Linear(la) if stream.is_some() => {
7386 if !(t >= 3 || (t == 2 && spec_m2()))
7387 || !self.mixer_in_q8_1_fast(e, &layer.mixer)
7388 || !e.uses_q8_1_fast(&la.ssm_out)
7389 {
7390 return Err("qwen35 stream verify: GDN batched arm requires t>=3 \
7391 (or MEMRA_SPEC_M2 at t=2) and q8_1-fast projections"
7392 .into());
7393 }
7394 let want = ckpt.is_some();
7395 let (out, stash) =
7396 self.linear_attn_verify_t(e, la, &xn, Some((&hq, &hd)), t, cache, il, want)?;
7397 if let (Some(ck), Some(st)) = (ckpt.as_deref_mut(), stash) {
7398 ck.gdn[il] = Some(st);
7399 }
7400 out
7401 }
7402 Mixer::Linear(_) => {
7403 unreachable!("linear layers ride qwen35_tparallel_linear_layer")
7404 }
7405 Mixer::Full(fa) => {
7406 let geometry = cfg.full_attention_geometry_at(il as u32);
7407 let n_head = geometry.n_head as usize;
7408 let n_head_kv = geometry.n_head_kv as usize;
7409 let head_dim = geometry.head_dim_k as usize;
7410 let rope_dims = geometry.n_rot as usize;
7411 let rope_base = geometry.rope_base;
7412 let scale = geometry.attention_scale();
7413 // Batched projections: one weight read serves all T rows.
7414 // GROUP-3 twin (trunk-kernels slice D): q/k/v in ONE launch — the group4
7415 // kernel with n3=0, bit-identical per (tensor, token, row) to the three
7416 // singles; refused or MEMRA_TK_FA_GROUP=0 -> singles byte-for-byte.
7417 let (qf, mut k, v) = match e.matmul_decode_exact_group3_pre(
7418 [&fa.wq, &fa.wk, &fa.wv],
7419 &hq,
7420 &hd,
7421 t,
7422 )? {
7423 Some(mut g3) => {
7424 let v = g3.pop().unwrap();
7425 let k = g3.pop().unwrap();
7426 let qf = g3.pop().unwrap();
7427 (qf, k, v)
7428 }
7429 None => (
7430 e.matmul_pre(&fa.wq, &hq, &hd, &xn, t)?,
7431 e.matmul_pre(&fa.wk, &hq, &hd, &xn, t)?,
7432 e.matmul_pre(&fa.wv, &hq, &hd, &xn, t)?,
7433 ),
7434 };
7435 let gated =
7436 geometry.attention_gate == memra_gguf::config::AttentionGateKind::FusedQ;
7437 let (mut q, gate) = if gated {
7438 let mut qs = e.uninit(t * n_head * head_dim)?;
7439 let mut gs = e.uninit(t * n_head * head_dim)?;
7440 e.q_gate_split(&qf, &mut qs, &mut gs, head_dim, n_head, t)?;
7441 (qs, Some(gs))
7442 } else {
7443 (qf, None)
7444 };
7445 let mut qn = e.uninit(t * n_head * head_dim)?;
7446 e.rms_norm(
7447 &q,
7448 fa.q_norm.float_data(),
7449 &mut qn,
7450 head_dim,
7451 t * n_head,
7452 eps,
7453 )?;
7454 q = qn;
7455 let mut kn = e.uninit(t * n_head_kv * head_dim)?;
7456 e.rms_norm(
7457 &k,
7458 fa.k_norm.float_data(),
7459 &mut kn,
7460 head_dim,
7461 t * n_head_kv,
7462 eps,
7463 )?;
7464 k = kn;
7465 e.rope_neox(
7466 &mut q, pos_d, head_dim, rope_dims, n_head, t, rope_base, 1.0,
7467 )?;
7468 e.rope_neox(
7469 &mut k, pos_d, head_dim, rope_dims, n_head_kv, t, rope_base, 1.0,
7470 )?;
7471
7472 // Per-row append + attend: row r sees rows 0..r in KV (causal within the
7473 // draft), each through the b_n=1 serving kernels at its own t_kv.
7474 let q_dim = n_head * head_dim;
7475 let kv_dim = n_head_kv * head_dim;
7476 let mut attn = e.uninit(t * q_dim)?;
7477 let (kdk, kdv, ktb, vtb, len0, kv_local) = {
7478 let kvl = cache.kv[il].as_ref().unwrap();
7479 // [2T] interleaved k,v base pointers: entry pair z serves row z of
7480 // the batched twins; the per-row fallback reads pair 0 (same cache
7481 // for every row of one layer). Graph mode reads the ctx table.
7482 let local: Option<CudaSlice<u64>> = match graph_cap {
7483 Some(_) => None,
7484 None => {
7485 let s = &e.gpu.stream();
7486 let (pk, _g) = kvl.k.device_ptr(s);
7487 let (pv, _g2) = kvl.v.device_ptr(s);
7488 let mut tbl = Vec::with_capacity(2 * t);
7489 for _ in 0..t {
7490 tbl.push(pk as u64);
7491 tbl.push(pv as u64);
7492 }
7493 Some(e.htod_u64(&tbl)?)
7494 }
7495 };
7496 (
7497 kvl.kv_dim_k,
7498 kvl.kv_dim_v,
7499 kvl.k_tok_bytes,
7500 kvl.v_tok_bytes,
7501 kvl.len,
7502 local,
7503 )
7504 };
7505 let (kv_tbl, kv_off): (&CudaSlice<u64>, usize) = match graph_cap {
7506 Some((tb, off, _)) => (tb, off),
7507 None => (kv_local.as_ref().expect("built above"), 0),
7508 };
7509 // Slice 4 (fa/append rows — see dspark_fa_rows_on): the whole per-row
7510 // section batches into the z-batched serving twins when every row of
7511 // this round takes the v4-seqs arm on ONE fa_split_keys rung. Both
7512 // guards are evaluated at the round's FIRST and LAST t_kv — the
7513 // eligibility window (vec floor .. v4 max) and each split-ladder rung
7514 // are intervals in t_kv, so ends-inside means all-inside (the straddle
7515 // law). Appending all T rows before any attend is read-equivalent to
7516 // the interleaved order: row r's walk reads keys 0..len0+r only, and
7517 // rows > r land at slots it never touches; every written cache row is
7518 // the per-token appender's exact warp program (kernel-check pinned).
7519 let t_kv_first = len0 + 1;
7520 let t_kv_last = len0 + t;
7521 let rows_batched = t >= 2
7522 && seqs_append
7523 && batch_fa_on
7524 && dspark_fa_rows_on()
7525 // the z-batched twins read stacked rows at the CACHE's kv dims;
7526 // the projection stack is [T, n_head_kv*head_dim] — they must be
7527 // the same stride or row z misaligns (true for this family; the
7528 // guard keeps any asymmetric-kv model on the per-row loop).
7529 && kdk == kv_dim
7530 && kdv == kv_dim
7531 && crate::fa_seqs_eligible(t_kv_first, head_dim_global)
7532 && crate::fa_seqs_eligible(t_kv_last, head_dim_global)
7533 && crate::fa_split_keys(t_kv_first, cfg.n_head_kv as usize)
7534 == crate::fa_split_keys(t_kv_last, cfg.n_head_kv as usize);
7535 // Sizing: eager = exact round bound; graph mode = the rung end (stride +
7536 // grid only — bytes proven equal above). Capture-time invariants refuse
7537 // loudly rather than bake a divergent body.
7538 let (size_kv_max, sp) = match graph_cap {
7539 Some((_, _, rung)) => {
7540 if !rows_batched {
7541 return Err(format!(
7542 "fa graph capture: layer {il} round is not batchable \
7543 (t_kv {t_kv_first}..{t_kv_last}) — the per-row fallback \
7544 must never be captured"
7545 )
7546 .into());
7547 }
7548 let sp_r = crate::fa_split_keys(rung, cfg.n_head_kv as usize);
7549 if t_kv_last > rung
7550 || sp_r != crate::fa_split_keys(t_kv_last, cfg.n_head_kv as usize)
7551 {
7552 return Err(format!(
7553 "fa graph capture: rung {rung} does not cover round \
7554 t_kv {t_kv_first}..{t_kv_last} on one split ladder step"
7555 )
7556 .into());
7557 }
7558 (rung, sp_r)
7559 }
7560 None => (
7561 t_kv_last,
7562 crate::fa_split_keys(t_kv_last, cfg.n_head_kv as usize),
7563 ),
7564 };
7565 if let Some((_, ctr)) = stream {
7566 // STREAM ARM (2b): one batched dc append + the multi-row dc attention
7567 // — the generic stream arm's exact shape (rows kernels are pinned
7568 // byte-identical to the per-row programs by kernel-check). Host len
7569 // stays a stale lower bound; the burst drain reconciles it.
7570 let kvl = cache.kv[il].as_mut().unwrap();
7571 e.append_kv_quantized_rows_dc(
7572 &k,
7573 &v,
7574 &mut kvl.k,
7575 &mut kvl.v,
7576 ctr,
7577 t,
7578 kdk,
7579 kdv,
7580 ktb,
7581 vtb,
7582 Engine::kv_fp8_on(),
7583 )?;
7584 let upper = (kvl.len + t + 64).min(cache.max_ctx);
7585 let k_view = e.view_u8(&kvl.k, upper * ktb);
7586 let v_view = e.view_u8(&kvl.v, upper * vtb);
7587 e.fa_decode_rows_dc(
7588 &q, &k_view, &v_view, &mut attn, head_dim, n_head, n_head_kv, ctr, upper,
7589 t, scale, ktb, vtb, 0, false,
7590 )?;
7591 } else if rows_batched {
7592 e.append_kv_quantized_seqs(
7593 &k,
7594 &v,
7595 &kv_tbl.slice(kv_off..kv_off + 2 * t),
7596 pos_d,
7597 t,
7598 kdk,
7599 kdv,
7600 ktb,
7601 vtb,
7602 )?;
7603 if graph_cap.is_none() {
7604 cache.kv[il].as_mut().unwrap().len += t;
7605 }
7606 e.fa_decode_batch_seqs_v4(
7607 &q,
7608 &kv_tbl.slice(kv_off..kv_off + 2 * t),
7609 pos_d,
7610 &mut attn,
7611 head_dim,
7612 n_head,
7613 n_head_kv,
7614 t,
7615 size_kv_max,
7616 scale,
7617 sp,
7618 ktb,
7619 vtb,
7620 )?;
7621 } else {
7622 if pos_rows.is_none() {
7623 // Stream-aware for symmetry with pos_d (the stream FA arm rides
7624 // the dc rows kernels above and never reaches this fallback).
7625 *pos_rows = Some(match stream {
7626 Some((_, ctr)) => (0..t)
7627 .map(|r| {
7628 let mut b = e.alloc_uninit::<i32>(1)?;
7629 e.i32_copy_add(ctr, &mut b, r as i32)?;
7630 Ok(b)
7631 })
7632 .collect::<Result<_, Box<dyn std::error::Error>>>()?,
7633 None => (0..t)
7634 .map(|r| e.htod_i32(&[(pos0 + r) as i32]))
7635 .collect::<Result<_, _>>()?,
7636 });
7637 }
7638 let pos_rows = pos_rows.as_ref().unwrap();
7639 for r in 0..t {
7640 // Owned per-row scratch: the b_n=1 kernels take packed batch buffers
7641 // whose row 0 is this row (arithmetic-free materialization copies,
7642 // same as decode's per-seq fallback arm).
7643 let mut k_row = e.uninit(kv_dim)?;
7644 e.dtod_copy_view(&k.slice(r * kv_dim..(r + 1) * kv_dim), &mut k_row)?;
7645 let mut v_row = e.uninit(kv_dim)?;
7646 e.dtod_copy_view(&v.slice(r * kv_dim..(r + 1) * kv_dim), &mut v_row)?;
7647 let pos_row = &pos_rows[r];
7648 let kvl = cache.kv[il].as_mut().unwrap();
7649 if seqs_append {
7650 e.append_kv_quantized_seqs(
7651 &k_row,
7652 &v_row,
7653 &kv_tbl.slice(kv_off..kv_off + 2),
7654 pos_row,
7655 1,
7656 kdk,
7657 kdv,
7658 ktb,
7659 vtb,
7660 )?;
7661 kvl.len += 1;
7662 } else {
7663 e.append_kv_quantized_view(
7664 &k_row.slice(0..kv_dim),
7665 &v_row.slice(0..kv_dim),
7666 &mut kvl.k,
7667 &mut kvl.v,
7668 kvl.len,
7669 kvl.kv_dim_k,
7670 kvl.kv_dim_v,
7671 kvl.k_tok_bytes,
7672 kvl.v_tok_bytes,
7673 Engine::kv_fp8_on(),
7674 )?;
7675 kvl.len += 1;
7676 }
7677 let t_kv = kvl.len;
7678 let mut q_row = e.uninit(q_dim)?;
7679 e.dtod_copy_view(&q.slice(r * q_dim..(r + 1) * q_dim), &mut q_row)?;
7680 let mut a_row = e.uninit(q_dim)?;
7681 if batch_fa_on && crate::fa_seqs_eligible(t_kv, head_dim_global) {
7682 let sp0_r = crate::fa_split_keys(t_kv, cfg.n_head_kv as usize);
7683 e.fa_decode_batch_seqs_v4(
7684 &q_row,
7685 &kv_tbl.slice(kv_off..kv_off + 2),
7686 pos_row,
7687 &mut a_row,
7688 head_dim,
7689 n_head,
7690 n_head_kv,
7691 1,
7692 t_kv,
7693 scale,
7694 sp0_r,
7695 ktb,
7696 vtb,
7697 )?;
7698 } else {
7699 let k_view = e.view_u8(&kvl.k, t_kv * kvl.k_tok_bytes);
7700 let v_view = e.view_u8(&kvl.v, t_kv * kvl.v_tok_bytes);
7701 let mut a_view = a_row.slice_mut(0..q_dim);
7702 e.fa_decode_kvmod_view(
7703 &q_row.slice(0..q_dim),
7704 &k_view,
7705 &v_view,
7706 &mut a_view,
7707 head_dim,
7708 n_head,
7709 n_head_kv,
7710 t_kv,
7711 scale,
7712 kvl.k_tok_bytes,
7713 kvl.v_tok_bytes,
7714 Engine::kv_fp8_on(),
7715 )?;
7716 }
7717 e.dtod_copy_into(&a_row, &mut attn, r * q_dim)?;
7718 }
7719 }
7720
7721 // Output gate (element-wise) + o-proj at m=T.
7722 let attn_g = match &gate {
7723 Some(g) => {
7724 let n = t * q_dim;
7725 let mut gsig = e.uninit(n)?;
7726 e.sigmoid(g, &mut gsig, n)?;
7727 let mut ag = e.uninit(n)?;
7728 e.mul(&attn, &gsig, &mut ag, n)?;
7729 ag
7730 }
7731 None => attn,
7732 };
7733 e.matmul(&fa.wo, &attn_g, t)?
7734 }
7735 };
7736
7737 // ---- residual add + post_attn_norm + FFN at m=T (serving dispatch verbatim) ----
7738 let pnorm = layer.post_attn_norm.float_data();
7739 let mut x1 = e.uninit(t * n_embd)?;
7740 let mut zn = e.uninit(t * n_embd)?;
7741 e.add_rms_norm(x, &mixed, pnorm, &mut x1, &mut zn, n_embd, t, eps)?;
7742 let ffn_out = match &layer.ffn {
7743 crate::hybrid::Ffn::Dense {
7744 ffn_gate,
7745 ffn_up,
7746 ffn_down,
7747 } => {
7748 assert!(
7749 self.cfg.m3.is_none(),
7750 "qwen35 t-parallel verify: M3 swigluoai FFN not yet batched"
7751 );
7752 self.qwen35_tparallel_dense_ffn(e, ffn_gate, ffn_up, ffn_down, &zn, t, n_embd)?
7753 }
7754 crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il_zq8(e, m, &zn, None, t, il as u16)?,
7755 };
7756 let mut x2 = e.uninit(t * n_embd)?;
7757 e.add(&x1, &ffn_out, &mut x2, t * n_embd)?;
7758 // dspark drafter tap (no-op when no sink armed): post-layer residual verify rows
7759 self.dflash_tap(e, cache, il, &x2, t)?;
7760 Ok(x2)
7761 }
7762
7763 /// ONE t-parallel LINEAR layer (attn_norm + gdn mixer + post_attn_norm + FFN + tap) —
7764 /// the exact body the old in-loop Linear arm ran, extracted so the eager walk and the
7765 /// slice-3 captured segments execute the SAME code (a second copy is how dispatch
7766 /// mirrors drift — the verify_layers extraction lesson). Two deliberate changes, both
7767 /// bit-identical by construction:
7768 /// - the gdn ping-pong host swap moves from per-row to ONE end-of-body swap (t odd):
7769 /// the device sequence is driven entirely by the 6-entry pointer table, which
7770 /// already encodes both parities; the ckpt stash reads name row r's out buffer
7771 /// directly (r even -> alt handle, odd -> canonical) — the same physical bytes the
7772 /// legacy post-swap clone read.
7773 /// - `stash` (slice-3 ctx): persistent per-layer slabs written by copy_into instead of
7774 /// per-row clone_dtod allocs — same bytes, capture-legal (no per-round host objects).
7775 /// `table_src` = (persistent pointer table, offset) when the ctx owns the tables;
7776 /// None builds the per-verify table exactly as before.
7777 #[allow(clippy::too_many_arguments)]
7778 fn qwen35_tparallel_linear_layer(
7779 &self,
7780 e: &Engine,
7781 il: usize,
7782 x: &CudaSlice<f32>,
7783 t: usize,
7784 cache: &mut Cache,
7785 mut ckpt: Option<&mut VerifyCkpt>,
7786 stash: Option<(&mut CudaSlice<f32>, &mut CudaSlice<f32>)>,
7787 table_src: Option<(&CudaSlice<u64>, usize)>,
7788 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7789 use cudarc::driver::DevicePtr;
7790 let cfg = &self.cfg;
7791 let n_embd = cfg.n_embd as usize;
7792 let eps = cfg.rms_eps;
7793 let layer = &self.layers[il];
7794 let Mixer::Linear(la) = &layer.mixer else {
7795 return Err("qwen35_tparallel_linear_layer on a non-linear layer".into());
7796 };
7797 // ---- attn_norm + q8_1 quantize at m=T (row-indexed == per-row) ----
7798 let anorm = layer.attn_norm.float_data();
7799 let mut xn = e.uninit(t * n_embd)?;
7800 e.rms_norm(x, anorm, &mut xn, n_embd, t, eps)?;
7801 let (hq, hd) = e.quantize_q8_1(&xn, t, n_embd)?;
7802
7803 let geometry = la.geometry;
7804 let d_state = geometry.key_head_dim as usize;
7805 let num_k = geometry.key_heads as usize;
7806 let num_v = geometry.value_heads as usize;
7807 let d_conv = geometry.conv_kernel as usize;
7808 let key_dim = d_state * num_k;
7809 let value_dim = geometry.value_head_dim as usize * num_v;
7810 let conv_dim = key_dim * 2 + value_dim;
7811 let gdn_scale = 1.0 / (d_state as f32).sqrt();
7812
7813 // ---- batched projections: one weight read for all T rows ----
7814 // GROUP-4 twin (trunk-kernels slice C): the whole 4-tuple in ONE launch, bit-identical
7815 // per (tensor, token, row) to the four singles; refused (layout/tier) or
7816 // MEMRA_TK_GDN_GROUP=0 -> the singles chain byte-for-byte.
7817 let (qkv_mixed, z, beta_raw, alpha) = match e.matmul_decode_exact_group4_pre(
7818 [&la.wqkv, &la.wqkv_gate, &la.ssm_beta, &la.ssm_alpha],
7819 &hq,
7820 &hd,
7821 t,
7822 )? {
7823 Some(mut g4) => {
7824 let alpha = g4.pop().unwrap();
7825 let beta_raw = g4.pop().unwrap();
7826 let z = g4.pop().unwrap();
7827 let qkv_mixed = g4.pop().unwrap();
7828 (qkv_mixed, z, beta_raw, alpha)
7829 }
7830 None => (
7831 e.matmul_pre(&la.wqkv, &hq, &hd, &xn, t)?,
7832 e.matmul_pre(&la.wqkv_gate, &hq, &hd, &xn, t)?,
7833 e.matmul_pre(&la.ssm_beta, &hq, &hd, &xn, t)?,
7834 e.matmul_pre(&la.ssm_alpha, &hq, &hd, &xn, t)?,
7835 ),
7836 };
7837 let beta_w = la.ssm_beta.out_features();
7838 let alpha_w = la.ssm_alpha.out_features();
7839 let qkv_w = la.wqkv.out_features();
7840
7841 // ---- per-row state chain through the b_n=1 serving kernels ----
7842 // 6-entry alternating pointer table expresses the ping-pong without a rebuild per
7843 // row: even rows scan s0 -> s1, odd rows s1 -> s0.
7844 let table_local: Option<CudaSlice<u64>> = match table_src {
7845 Some(_) => None,
7846 None => {
7847 let rl = cache.recur[il].as_ref().unwrap();
7848 let s = &e.gpu.stream();
7849 let (pc, _g0) = rl.conv_state.device_ptr(s);
7850 let (p0, _g1) = rl.ssm_state.device_ptr(s);
7851 let (p1, _g2) = rl.ssm_state_alt.device_ptr(s);
7852 Some(e.htod_u64(&[
7853 pc as u64, p0 as u64, p1 as u64, pc as u64, p1 as u64, p0 as u64,
7854 ])?)
7855 }
7856 };
7857 let (table, toff): (&CudaSlice<u64>, usize) = match table_src {
7858 Some((tb, off)) => (tb, off),
7859 None => (table_local.as_ref().unwrap(), 0),
7860 };
7861 let mut o_all = e.uninit(t * value_dim)?;
7862 let mut col_states: Option<Vec<(CudaSlice<f32>, CudaSlice<f32>)>> =
7863 if ckpt.is_some() && stash.is_none() && t >= 2 {
7864 Some(Vec::with_capacity(t - 1))
7865 } else {
7866 None
7867 };
7868 let mut stash = stash;
7869 // Per-row scratch reused across rows (uninit is cheap but not free at
7870 // 48 layers x T rows); row inputs/outputs pass as VIEWS into the packed
7871 // [T, ...] buffers — zero arithmetic-free copies in this loop.
7872 let mut conv_out = e.uninit(conv_dim)?;
7873 let mut q_l2 = e.uninit(value_dim)?;
7874 let mut k_l2 = e.uninit(value_dim)?;
7875 let mut v_gd = e.uninit(value_dim)?;
7876 let mut beta_b = e.uninit(num_v)?;
7877 let mut g_log = e.uninit(num_v)?;
7878 for r in 0..t {
7879 let base = toff + if r % 2 == 0 { 0 } else { 3 };
7880 let conv_view = table.slice(base..base + 1);
7881 let in_view = table.slice(base + 1..base + 2);
7882 let out_view = table.slice(base + 2..base + 3);
7883 e.ssm_conv1d_fused_decode_b_view(
7884 &qkv_mixed.slice(r * qkv_w..(r + 1) * qkv_w),
7885 &conv_view,
7886 la.ssm_conv1d.float_data(),
7887 &mut conv_out,
7888 conv_dim,
7889 d_conv,
7890 1,
7891 )?;
7892 e.gdn_prep_decode_b_view(
7893 &conv_out,
7894 &beta_raw.slice(r * beta_w..(r + 1) * beta_w),
7895 &alpha.slice(r * alpha_w..(r + 1) * alpha_w),
7896 la.ssm_dt.float_data(),
7897 la.ssm_a.float_data(),
7898 &mut q_l2,
7899 &mut k_l2,
7900 &mut v_gd,
7901 &mut beta_b,
7902 &mut g_log,
7903 d_state,
7904 num_v,
7905 num_k,
7906 key_dim,
7907 eps,
7908 conv_dim,
7909 1,
7910 )?;
7911 let mut o_row = o_all.slice_mut(r * value_dim..(r + 1) * value_dim);
7912 e.gdn_scan_s128_batched_view(
7913 &q_l2, &k_l2, &v_gd, &g_log, &beta_b, &in_view, &out_view, &mut o_row, num_v, 1,
7914 gdn_scale,
7915 )?;
7916 if r + 1 < t {
7917 // Row r's out buffer: even rows write s1 (the alt handle — no swaps ran),
7918 // odd rows write s0 — the same physical state the legacy post-swap
7919 // canonical clone read.
7920 let rl = cache.recur[il]
7921 .as_ref()
7922 .ok_or("qwen35 linear verify layer has no recurrent state")?;
7923 let ssm_src = if r % 2 == 0 {
7924 &rl.ssm_state_alt
7925 } else {
7926 &rl.ssm_state
7927 };
7928 match stash.as_mut() {
7929 Some((conv_slab, ssm_slab)) => {
7930 // BOTH stash reads go through the pointer table at run time: the
7931 // ssm handles ping-pong between rounds, and the ctx (with its
7932 // captured graphs) outlives the Cache — a fresh generation's
7933 // conv/ssm buffers land at new addresses that only the per-round
7934 // table refresh knows. A baked direct copy would read freed
7935 // memory (parity was the slice-3 smoke divergence; cache
7936 // lifetime is the cross-generation twin).
7937 e.copy_indirect_src_f32(
7938 &conv_view,
7939 conv_slab,
7940 r * conv_dim * (d_conv - 1),
7941 conv_dim * (d_conv - 1),
7942 )?;
7943 // The ssm handles PING-PONG between rounds: a captured direct
7944 // copy would bake the capture-time physical buffer and read the
7945 // wrong parity after any odd-vt round (the slice-3 smoke
7946 // divergence). Read the src address from row r's OUT table
7947 // entry at run time — the same entry the scan just wrote.
7948 e.copy_indirect_src_f32(
7949 &out_view,
7950 ssm_slab,
7951 r * d_state * d_state * num_v,
7952 d_state * d_state * num_v,
7953 )?;
7954 }
7955 None => {
7956 if let Some(states) = col_states.as_mut() {
7957 states.push((e.clone_dtod(&rl.conv_state)?, e.clone_dtod(ssm_src)?));
7958 }
7959 }
7960 }
7961 }
7962 }
7963 // ONE end-of-body parity swap (t odd) — the legacy loop swapped per row; the net
7964 // handle motion is identical and the device sequence never read the handles.
7965 if t % 2 == 1 {
7966 let rl = cache.recur[il].as_mut().unwrap();
7967 std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
7968 }
7969 if let (Some(checkpoint), Some(states)) = (ckpt.as_deref_mut(), col_states) {
7970 checkpoint.cols[il] = Some(states);
7971 }
7972
7973 // ---- batched gated norm + out-projection at m=T ----
7974 let mixed = if e.uses_q8_1_fast(&la.ssm_out) {
7975 let (gq, gd) = e.gated_rmsnorm_q8_1(
7976 &o_all,
7977 la.ssm_norm.float_data(),
7978 &z,
7979 d_state,
7980 t * num_v,
7981 eps,
7982 )?;
7983 let g0 = e.zeros(0)?;
7984 e.matmul_pre(&la.ssm_out, &gq, &gd, &g0, t)?
7985 } else {
7986 let mut gn = e.uninit(t * value_dim)?;
7987 e.gated_rmsnorm(
7988 &o_all,
7989 la.ssm_norm.float_data(),
7990 &z,
7991 &mut gn,
7992 d_state,
7993 t * num_v,
7994 eps,
7995 )?;
7996 e.matmul(&la.ssm_out, &gn, t)?
7997 };
7998
7999 // ---- residual add + post_attn_norm + FFN at m=T (serving dispatch verbatim) ----
8000 let pnorm = layer.post_attn_norm.float_data();
8001 let mut x1 = e.uninit(t * n_embd)?;
8002 let mut zn = e.uninit(t * n_embd)?;
8003 e.add_rms_norm(x, &mixed, pnorm, &mut x1, &mut zn, n_embd, t, eps)?;
8004 let ffn_out = match &layer.ffn {
8005 crate::hybrid::Ffn::Dense {
8006 ffn_gate,
8007 ffn_up,
8008 ffn_down,
8009 } => {
8010 assert!(
8011 self.cfg.m3.is_none(),
8012 "qwen35 t-parallel verify: M3 swigluoai FFN not yet batched"
8013 );
8014 self.qwen35_tparallel_dense_ffn(e, ffn_gate, ffn_up, ffn_down, &zn, t, n_embd)?
8015 }
8016 crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il_zq8(e, m, &zn, None, t, il as u16)?,
8017 };
8018 let mut x2 = e.uninit(t * n_embd)?;
8019 e.add(&x1, &ffn_out, &mut x2, t * n_embd)?;
8020 // dspark drafter tap (no-op when no sink armed): post-layer residual verify rows
8021 self.dflash_tap(e, cache, il, &x2, t)?;
8022 Ok(x2)
8023 }
8024
8025 /// PP-N STAGE SUBGRAPH of the verify trunk: layers `[lo, hi)` of `decode_step_t_core_stream`'s
8026 /// walk, verbatim. Enters with a MATERIALIZED `[T, n_embd]` residual (no pending fusion pair
8027 /// carried in from outside the range) and exits with the range's final residual materialized
8028 /// (the trailing add executed) — exactly the `decode_layers_eager(lo, hi)` contract, T rows
8029 /// instead of one.
8030 ///
8031 /// EXTRACTED (lane/pp2-spec 2026-08-06) rather than duplicated: `decode_step_t_core_stream` IS
8032 /// the single funnel every verify forward reaches, and its per-layer dispatch MIRRORING (norm
8033 /// fusion per layer, the t>=3/spec_m2 batched-linear window, the fused-q8 FFN chain, the
8034 /// decode-exact projections) is what makes verify bit-identical to eager decode. A second copy
8035 /// for the split arm is how those mirrors drift apart on the next lever. The unsplit body now
8036 /// calls this with `(0, n_layers)`, so the whole-trunk path and every stage range run the SAME
8037 /// code — there is no "split version" of the verify math.
8038 ///
8039 /// Bit-identity of a cut rests on the same kernel-check-pinned identity the eager arm's cut
8040 /// does — `add_rms_norm_q8_1 == add then rms_norm_q8_1` at nrows=T — because the ONLY thing a
8041 /// fence changes is that the cross-layer fusion carry breaks at `hi-1` and is re-materialized
8042 /// as an explicit `add`. `decode-batch-gate --mode ppspec` verifies end-to-end on real weights.
8043 #[allow(clippy::too_many_arguments)]
8044 fn verify_layers(
8045 &self,
8046 e: &Engine,
8047 mut x: CudaSlice<f32>,
8048 lo: usize,
8049 hi: usize,
8050 pos_d: &CudaSlice<i32>,
8051 pos0: usize,
8052 t: usize,
8053 cache: &mut Cache,
8054 mut ckpt: Option<&mut VerifyCkpt>,
8055 stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
8056 graphs: Option<&mut DsparkVerifyGraphs>,
8057 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8058 if self.sliding_gated_moe_batch_program() {
8059 if stream.is_some() {
8060 return Err(
8061 "step35 has no ROUND-STREAM verify arm (the device-counter _dc twins \
8062 cannot express the SWA offset KV view)"
8063 .into(),
8064 );
8065 }
8066 return self.step35_verify_batch_layers(e, x, lo, hi, pos0, t, cache);
8067 }
8068 if self.batched_serving_numeric_class() {
8069 return self.qwen35_verify_batch_layers(
8070 e,
8071 x,
8072 lo,
8073 hi,
8074 pos0,
8075 t,
8076 cache,
8077 ckpt.take(),
8078 stream,
8079 graphs,
8080 );
8081 }
8082 let n_embd = self.cfg.n_embd as usize;
8083 let eps = self.cfg.rms_eps;
8084 // CROSS-LAYER ADD+NORM FUSION (lane/vt-fixes fix 2, mirroring decode_step_h's
8085 // launch-arc form): layer il's post-FFN residual add (x2 = x1 + ffn_out) and layer
8086 // il+1's attn_norm(+quantize) are consecutive row-wise ops — ONE add_rms_norm_q8_1
8087 // launch at nrows=t does all three (bit-identity pinned by the T-row kernel-check
8088 // arms). Carry the un-added (x1, ffn_out) pair; the fused launch materializes x2 (the
8089 // residual the next layer needs) as its `res` output. Falls back to the separate add
8090 // when the next layer is off the fused-q8 path.
8091 let mut pending: Option<(CudaSlice<f32>, CudaSlice<f32>)> = None;
8092 for il in lo..hi {
8093 let layer = &self.layers[il];
8094 // DISPATCH-MIRRORED attn-input RMSNorm (FP-order lesson #8): eager decode fuses the
8095 // 1024-thread rms_norm_q8_1 ONLY when every mixer projection is q8_1-fast; layers with
8096 // Float projections (ssm_beta/ssm_alpha on layers 1/2/4 of the 9B NVFP4 GGUF) take the
8097 // UNFUSED 256-thread rms_norm. The verify norm must mirror that PER-LAYER choice —
8098 // blockDim changes the sum-of-squares reduce order, and the ULP shift amplifies through
8099 // the GDN recurrence into argmax flips (measured: 9B text prompt, 1 ULP at layer 2 ->
8100 // 2.3e-1 logit maxdiff at the head -> K=1..8 divergence at a 0.03-margin token).
8101 let mixer_fast = self.mixer_in_q8_1_fast(e, &layer.mixer);
8102 let norm_fused = std::env::var("MEMRA_NO_FUSE_NORMQ").is_err() && mixer_fast;
8103 // BATCHED EPILOGUE RE-FUSE (lane/vt-fixes fix 2, 2026-08-03): when the norm is
8104 // dispatch-fused AND every consumer of `h` reads only its q8_1 form (Full mixer:
8105 // projections only; Linear mixer: the batched arm — the per-column fallback needs
8106 // f32 h), emit the attn-input norm DIRECTLY as q8_1 via `rms_norm_q8_1` at nrows=t
8107 // (row-indexed kernel — the T-row launch is the per-row m=1 program, kernel-check
8108 // pins bit-identity vs rms_norm_decode -> quantize_q8_1). Kills the standalone
8109 // quantize launch(es) + the f32 h HBM round-trip that decode never pays.
8110 // step35 (Full mixer) is the third case that needs f32 `h`: its verify arm is a
8111 // per-ROW replay of the eager decode mixer, whose `pre_q` contract is a single row —
8112 // a T-row q8_1 pair cannot be handed to it, and re-deriving per-row q8_1 from the f32
8113 // rows is exactly the dispatch being mirrored. Keep step35 on the unfused arm.
8114 let lin_q8_only = match &layer.mixer {
8115 Mixer::Linear(la) => {
8116 (t >= 3 || (t == 2 && spec_m2())) && e.uses_q8_1_fast(&la.ssm_out)
8117 }
8118 Mixer::Full(_) if self.sliding_gated_moe_batch_program() => false,
8119 _ => true,
8120 };
8121 // NOTE decode.rs's take()-first lesson: take the pending pair BEFORE branching so
8122 // a non-fused layer still performs the residual add.
8123 let taken = pending.take();
8124 let (h, h_q8) = if norm_fused && lin_q8_only {
8125 let pair = match taken {
8126 // fused add + attn_norm + q8_1: ONE launch resolves the carried residual
8127 // AND emits this layer's mixer input pre-quantized. res -> x2 (= new x).
8128 Some((x1p, f1p)) => {
8129 let mut x2 = vbuf(e, t * n_embd)?; // fully written (res output)
8130 let p = e.add_rms_norm_q8_1(
8131 &x1p,
8132 &f1p,
8133 layer.attn_norm.float_data(),
8134 &mut x2,
8135 n_embd,
8136 t,
8137 eps,
8138 )?;
8139 x = x2;
8140 p
8141 }
8142 None => e.rms_norm_q8_1(&x, layer.attn_norm.float_data(), n_embd, t, eps)?,
8143 };
8144 (e.zeros(0)?, Some(pair)) // h unused on this path (q8-only consumers)
8145 } else {
8146 if let Some((x1p, f1p)) = taken {
8147 let mut x2 = vbuf(e, t * n_embd)?; // fully written by add
8148 e.add(&x1p, &f1p, &mut x2, t * n_embd)?;
8149 x = x2;
8150 }
8151 let mut h = vbuf(e, t * n_embd)?; // fully written by either rms_norm arm
8152 if norm_fused {
8153 e.rms_norm_decode(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
8154 } else {
8155 e.rms_norm(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
8156 }
8157 (h, None)
8158 };
8159 let h_q8_ref = h_q8.as_ref().map(|(q, d)| (q, d));
8160
8161 let mixed = match &layer.mixer {
8162 Mixer::Full(fa) => self.full_attn_verify(
8163 e,
8164 fa,
8165 &h,
8166 h_q8_ref,
8167 pos_d,
8168 t,
8169 cache,
8170 il,
8171 stream.map(|(_, c)| c),
8172 )?,
8173 Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
8174 Mixer::Linear(la) => {
8175 // BATCHED linear verify (2026-07-03, the MTP-profit lever): one T-token pass —
8176 // batched projections (weight read ONCE, hits the m=2-4 weight-resident matvec),
8177 // carried-state conv (ssm_conv1d_tm_state), GDN prep on the prefill kernels, and
8178 // ONE gdn_scan whose internal sequential t-loop is the SAME recurrence as T
8179 // chained T=1 steps (bit-identical). Falls back to the sequential per-column
8180 // chain when T < d_conv-1 (conv ring update needs T >= pad) — or when ANY
8181 // projection is off the q8_1 fast path: matmul_decode_exact would route a Float
8182 // tensor to cuBLAS at m=t (different FP accumulation than eager's per-token
8183 // GEMV), so mixed-dtype layers stay on the eager-identical per-column chain.
8184 // MEMRA_SPEC_M2 (lane/spec-m2): the t==2 batch rides the same arm — the conv
8185 // wrapper handles t<pad with a pure-copy ring rebuild; see spec_m2() header.
8186 if (t >= 3 || (t == 2 && spec_m2()))
8187 && mixer_fast
8188 && e.uses_q8_1_fast(&la.ssm_out)
8189 {
8190 let want = ckpt.is_some();
8191 let (out, stash) =
8192 self.linear_attn_verify_t(e, la, &h, h_q8_ref, t, cache, il, want)?;
8193 if let (Some(ck), Some(st)) = (ckpt.as_deref_mut(), stash) {
8194 ck.gdn[il] = Some(st);
8195 }
8196 out
8197 } else {
8198 let mut out = vbuf(e, t * n_embd)?; // every col written by copy_into
8199 let mut col_states: Option<Vec<(CudaSlice<f32>, CudaSlice<f32>)>> =
8200 if ckpt.is_some() && t >= 2 {
8201 Some(Vec::with_capacity(t - 1))
8202 } else {
8203 None
8204 };
8205 for col in 0..t {
8206 let mut h_col = vbuf(e, n_embd)?; // fully written by copy_view_into
8207 let src = h.slice(col * n_embd..(col + 1) * n_embd);
8208 e.copy_view_into(&mut h_col, 0, &src, n_embd)?;
8209 let m_col = self.linear_attn_decode(e, la, &h_col, cache, il)?;
8210 e.copy_into(&mut out, col * n_embd, &m_col, n_embd)?;
8211 // REPLAY-FREE ckpt: clone the chain's ACTUAL state after this column
8212 // (pure dtod — cannot change any computed value). Last column skipped:
8213 // rebuild targets are j <= t-1 columns.
8214 if let Some(cs) = col_states.as_mut() {
8215 if col + 1 < t {
8216 let rl = cache.recur[il].as_ref().unwrap();
8217 cs.push((
8218 e.clone_dtod(&rl.conv_state)?,
8219 e.clone_dtod(&rl.ssm_state)?,
8220 ));
8221 }
8222 }
8223 }
8224 if let (Some(ck), Some(cs)) = (ckpt.as_deref_mut(), col_states) {
8225 // ReplaySSM-assessment instrumentation (2026-07-30): the
8226 // per-column clones are the only true state snapshots left in
8227 // the verify (the batched path stashes INPUTS and replays).
8228 if std::env::var("MEMRA_SPEC_STATS").as_deref() == Ok("1") {
8229 static ONCE: std::sync::Once = std::sync::Once::new();
8230 let bytes: usize =
8231 cs.iter().map(|(c, s)| (c.len() + s.len()) * 4).sum();
8232 ONCE.call_once(|| eprintln!(
8233 "[verify-ckpt] per-column layer il={il}: {} clones, {:.2} MB/layer/round",
8234 cs.len(), bytes as f64 / 1e6));
8235 }
8236 ck.cols[il] = Some(cs);
8237 }
8238 out
8239 }
8240 }
8241 };
8242
8243 // DISPATCH-MIRRORED post-attn norm: eager residual_norm_ffn fuses add+norm+quant
8244 // (1024-thread add_rms_norm_q8_1) only for Dense FFNs whose gate+up are q8_1-fast;
8245 // otherwise (and for MoE) it runs the 256-thread fused add_rms_norm. Mirror per layer.
8246 let ffn_fuse = match &layer.ffn {
8247 crate::hybrid::Ffn::Dense {
8248 ffn_gate, ffn_up, ..
8249 } => {
8250 std::env::var("MEMRA_NO_FUSE_NORMQ").is_err()
8251 && e.uses_q8_1_fast(ffn_gate)
8252 && e.uses_q8_1_fast(ffn_up)
8253 }
8254 crate::hybrid::Ffn::Moe(_) => false,
8255 };
8256 // BATCHED EPILOGUE RE-FUSE (lane/vt-fixes fix 2): on the ffn_fuse path (Dense,
8257 // gate+up q8_1-fast, non-M3) the FFN input is emitted DIRECTLY as q8_1 by ONE
8258 // add_rms_norm_q8_1 launch at nrows=t (row-indexed kernel: the T-row launch is the
8259 // per-row m=1 program; kernel-check pins bit-identity vs the unfused
8260 // add_f32 -> rms_norm_decode -> quantize_q8_1 chain at T=2/4/5/8) — replacing the
8261 // add + rms_norm_decode launches AND the dual/singles' internal re-quantize.
8262 // M3's swigluoai must keep the f32 chain (the fused SwiGLU epilogue encodes plain
8263 // SiLU), mirroring residual_norm_ffn's m3 guard on the decode path.
8264 // step35: same guard per LAYER. A dense FFN's clamp is the SHEXP array (upstream's
8265 // one build_ffn serves dense + shared expert, llama-graph.cpp:1751), and verify MUST
8266 // mirror decode's dispatch or spec self-consistency fails.
8267 let dense_lim = self.cfg.clamp_shexp_at(il as u32);
8268 let fuse_q8 = ffn_fuse && self.cfg.m3.is_none() && dense_lim.is_none();
8269 let mut x1 = vbuf(e, t * n_embd)?; // fully written by add / add_rms_norm*
8270 let mut z = e.zeros(0)?; // replaced below on the unfused arms
8271 let z_q8 = if fuse_q8 {
8272 Some(e.add_rms_norm_q8_1(
8273 &x,
8274 &mixed,
8275 layer.post_attn_norm.float_data(),
8276 &mut x1,
8277 n_embd,
8278 t,
8279 eps,
8280 )?)
8281 } else {
8282 let mut zf = vbuf(e, t * n_embd)?; // fully written by rms_norm_decode / add_rms_norm
8283 if ffn_fuse {
8284 e.add(&x, &mixed, &mut x1, t * n_embd)?;
8285 e.rms_norm_decode(
8286 &x1,
8287 layer.post_attn_norm.float_data(),
8288 &mut zf,
8289 n_embd,
8290 t,
8291 eps,
8292 )?;
8293 } else {
8294 e.add_rms_norm(
8295 &x,
8296 &mixed,
8297 layer.post_attn_norm.float_data(),
8298 &mut x1,
8299 &mut zf,
8300 n_embd,
8301 t,
8302 eps,
8303 )?;
8304 }
8305 z = zf;
8306 None
8307 };
8308 // DECODE-EXACT FFN projections: force MMVQ for gate/up/down at any T to match the
8309 // T=1 decode FP accumulation order. At T>=5 the generic matmul/matmul_pre falls to dp4a
8310 // (128-thread, different FP sum order). At T=2-4 the batched MMVQ is already bit-identical.
8311 let ffn_out = match &layer.ffn {
8312 crate::hybrid::Ffn::Dense {
8313 ffn_gate,
8314 ffn_up,
8315 ffn_down,
8316 } => {
8317 let n_ff = ffn_gate.out_features();
8318 if let Some((zq, zd)) = z_q8.as_ref() {
8319 // FUSED CHAIN (fix 2): pre-quantized z feeds the projections; the SwiGLU
8320 // epilogue emits act pre-quantized for ffn_down (silu_mul_scaled_q8_1,
8321 // bit-identical to silu_mul + quantize — kernel-check-pinned) with the
8322 // NVFP4 macro-scales folded (deferred-scale dual: y*s inline == the
8323 // scale_inplace store, value-exact) — the exact m=1 decode epilogue
8324 // structure at nrows=t.
8325 let pair =
8326 match e.matmul_decode_exact_dual_pre(ffn_gate, ffn_up, zq, zd, t)? {
8327 Some(((g, gs), (u, us))) => Some((g, gs, u, us)),
8328 None => None,
8329 };
8330 let (gate, gs, up, us) = match pair {
8331 Some(x4) => x4,
8332 None => (
8333 e.matmul_decode_exact_pre(ffn_gate, zq, zd, t)?,
8334 1.0, // scale already applied inside _pre
8335 e.matmul_decode_exact_pre(ffn_up, zq, zd, t)?,
8336 1.0,
8337 ),
8338 };
8339 if e.uses_q8_1_fast(ffn_down) {
8340 let (aq, ad) = e.silu_mul_scaled_q8_1(&gate, &up, gs, us, t * n_ff)?;
8341 e.matmul_decode_exact_pre(ffn_down, &aq, &ad, t)?
8342 } else {
8343 let mut act = vbuf(e, t * n_ff)?;
8344 e.silu_mul_scaled(&gate, &up, gs, us, &mut act, t * n_ff)?;
8345 e.matmul_decode_exact(ffn_down, &act, t)?
8346 }
8347 } else {
8348 // UNFUSED (pre-fix) chain — MoE-adjacent/M3/off-fast layers, unchanged.
8349 // DUAL gate+up batched twin (lane/verify-economics, 2026-08-02): one launch
8350 // for the pair at t=2..8 — bit-identical per (tensor,token,row) to the two
8351 // singles (kernel-check pins bitwise; MEMRA_SPEC_DUAL_T=0 reverts). None
8352 // (non-NVFP4 / t outside the tier / seam off) -> the two singles, unchanged.
8353 let (gate, up) =
8354 match e.matmul_decode_exact_dual(ffn_gate, ffn_up, &z, t)? {
8355 Some(pair) => pair,
8356 None => (
8357 e.matmul_decode_exact(ffn_gate, &z, t)?,
8358 e.matmul_decode_exact(ffn_up, &z, t)?,
8359 ),
8360 };
8361 let mut act = vbuf(e, t * n_ff)?; // fully written by ffn_act_lim
8362 Self::ffn_act_lim(
8363 e,
8364 &self.cfg,
8365 &gate,
8366 &up,
8367 1.0,
8368 1.0,
8369 dense_lim,
8370 &mut act,
8371 t * n_ff,
8372 )?;
8373 e.matmul_decode_exact(ffn_down, &act, t)?
8374 }
8375 }
8376 crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il(e, m, &z, t, il as u16)?,
8377 };
8378 // CROSS-LAYER fusion: defer this layer's post-FFN residual add — the next layer's
8379 // fused-q8 attn norm folds it in (add_rms_norm_q8_1 == add; rms_norm; quantize,
8380 // kernel-check-pinned at nrows=T). Non-fused next layers add explicitly above.
8381 pending = Some((x1, ffn_out));
8382 }
8383 // RANGE's final add (no next norm INSIDE the range to fuse with; for the
8384 // whole-trunk call that is the last layer, whose next norm is output_norm — f32-out).
8385 if let Some((x1p, f1p)) = pending.take() {
8386 let mut x2 = vbuf(e, t * n_embd)?; // fully written by add
8387 e.add(&x1p, &f1p, &mut x2, t * n_embd)?;
8388 x = x2;
8389 }
8390 Ok(x)
8391 }
8392 /// BATCHED linear-attn verify (T=K+1): the whole layer in ~10 launches instead of T x the
8393 /// T=1 decode chain (T x ~12 launches + T weight reads of the four projections). The GDN
8394 /// recurrence itself is inherently sequential — gdn_scan_s128 runs its internal t-loop with
8395 /// the SAME per-token math as chained T=1 calls (bit-identical state evolution); everything
8396 /// around it (projections, conv, prep, gated norm, out-proj) batches. Advances conv ring +
8397 /// ssm state exactly like T sequential decode steps.
8398 /// `want_stash`: additionally RETAIN the gdn-scan inputs (pure buffer keep-alives, zero extra
8399 /// kernels) so a partial accept can rebuild the state after any column prefix (REPLAY-FREE).
8400 #[allow(clippy::too_many_arguments)]
8401 fn linear_attn_verify_t(
8402 &self,
8403 e: &Engine,
8404 la: &LinearAttnLayer,
8405 h: &CudaSlice<f32>,
8406 h_q8: Option<(&CudaSlice<i8>, &CudaSlice<f32>)>,
8407 t: usize,
8408 cache: &mut Cache,
8409 il: usize,
8410 want_stash: bool,
8411 ) -> Result<(CudaSlice<f32>, Option<GdnStash>), Box<dyn std::error::Error>> {
8412 let cfg = &self.cfg;
8413 let geometry = la.geometry;
8414 let d_state = geometry.key_head_dim as usize;
8415 let num_k = geometry.key_heads as usize;
8416 let num_v = geometry.value_heads as usize;
8417 let d_conv = geometry.conv_kernel as usize;
8418 let key_dim = d_state * num_k;
8419 let conv_dim = key_dim * 2 + geometry.value_head_dim as usize * num_v;
8420 let eps = cfg.rms_eps;
8421 let scale = 1.0 / (d_state as f32).sqrt();
8422
8423 // DECODE-EXACT projections: matmul_decode_exact forces the MMVQ (warp-per-row, 32-thread)
8424 // accumulation order for EVERY m, matching the T=1 decode path bit-for-bit. The generic
8425 // `matmul` at m>=5 falls to dp4a (128-thread, two-level reduce) which has a different FP
8426 // sum order — ULP differences propagate through gdn_scan and flip argmax on the 27B.
8427 // Q8 TRUNK-FUSION at T=1 (35B: wqkv+wqkv_gate both Q8_0): one fused2 launch, bit-identical
8428 // per (tensor,row) to the two m=1 MMVQ dispatches below — decode-exact contract holds.
8429 // VERIFY-TIER TRUNK FUSION (MEMRA_SPEC_FUSED_T, t=2-4): quantize h ONCE for every
8430 // fused-eligible same-input Q8_0 pair of this layer (35B wqkv+wqkv_gate; 9B
8431 // ssm_beta+ssm_alpha) — each fused2 batched launch then replaces two decode-exact
8432 // calls (each of which re-quantizes the same h + runs its own _b2/_b4 launch).
8433 // Bit-identical per (tensor,token,row) — see spec_fused_t().
8434 // BATCHED EPILOGUE RE-FUSE (lane/vt-fixes fix 2): `h_q8` = the attn-input norm emitted
8435 // directly as q8_1 by the caller's fused rms_norm_q8_1 (bit-identical to the unfused
8436 // chain, kernel-check-pinned). When present it REPLACES the standalone quantize below
8437 // and feeds every projection; the caller guaranteed all four input projections are
8438 // q8_1-fast. When absent, the old shared-quantize (fused-t window) stands.
8439 let h_q8_t = if h_q8.is_none()
8440 && spec_fused_t()
8441 && (2..=4).contains(&t)
8442 && ((e.uses_q8_1_fast(&la.wqkv) && e.uses_q8_1_fast(&la.wqkv_gate))
8443 || (e.uses_q8_1_fast(&la.ssm_beta) && e.uses_q8_1_fast(&la.ssm_alpha)))
8444 {
8445 Some(e.quantize_q8_1(h, t, cfg.n_embd as usize)?)
8446 } else {
8447 None
8448 };
8449 // one view: the caller's fused-norm q8 or this fn's own shared quantize.
8450 let hq8_any: Option<(&CudaSlice<i8>, &CudaSlice<f32>)> =
8451 h_q8.or(h_q8_t.as_ref().map(|(q, d)| (q, d)));
8452 let (qkv_mixed, z) = {
8453 let mut fused = None;
8454 if t == 1 && e.uses_q8_1_fast(&la.wqkv) && e.uses_q8_1_fast(&la.wqkv_gate) {
8455 let (hq, hd) = e.quantize_q8_1(h, 1, cfg.n_embd as usize)?;
8456 fused = e.matmul_q8_fused2(&la.wqkv, &la.wqkv_gate, &hq, &hd)?;
8457 } else if let Some((hq, hd)) = hq8_any {
8458 if spec_fused_t() && (2..=4).contains(&t) {
8459 fused = e.matmul_q8_fused2_t(&la.wqkv, &la.wqkv_gate, hq, hd, t)?;
8460 }
8461 }
8462 match (fused, hq8_any) {
8463 (Some(pair), _) => pair,
8464 (None, Some((hq, hd))) if h_q8.is_some() => (
8465 e.matmul_decode_exact_pre(&la.wqkv, hq, hd, t)?,
8466 e.matmul_decode_exact_pre(&la.wqkv_gate, hq, hd, t)?,
8467 ),
8468 (None, _) => (
8469 e.matmul_decode_exact(&la.wqkv, h, t)?,
8470 e.matmul_decode_exact(&la.wqkv_gate, h, t)?,
8471 ),
8472 }
8473 };
8474 // beta+alpha DUAL at T=1 (75% of p3 rounds run T=1 verify — p-min chain cuts): the dual
8475 // mr2 kernel is bit-identical per element to the m=1 MMVQ matmul_decode_exact dispatches
8476 // (same warp-per-row body, blockIdx.y picks the weight), so the decode-exact contract
8477 // holds; the run-spec battery is the arbiter. T>1 keeps the per-tensor decode-exact path.
8478 let (beta_raw, alpha) = if t == 1 {
8479 let (hq, hd) = e.quantize_q8_1(h, 1, cfg.n_embd as usize)?;
8480 match e.matmul_pre_dual_noscale(&la.ssm_beta, &la.ssm_alpha, &hq, &hd, 1)? {
8481 Some(((mut b, bs), (mut a, as_))) => {
8482 if bs != 1.0 {
8483 e.scale_inplace(&mut b, bs, la.ssm_beta.out_features())?;
8484 }
8485 if as_ != 1.0 {
8486 e.scale_inplace(&mut a, as_, la.ssm_alpha.out_features())?;
8487 }
8488 (b, a)
8489 }
8490 // Q8_0 fused2 twin (9B stores beta/alpha as Q8_0): DISPATCH-MIRRORS the eager
8491 // decode's beta_alpha closure — the fused body is qmatvec_q8_0_mmvq verbatim,
8492 // bit-identical per row (kernel-check rel=0.00e0 gate), so decode==verify holds.
8493 None => match e.matmul_q8_fused2(&la.ssm_beta, &la.ssm_alpha, &hq, &hd)? {
8494 Some((b, a)) => (b, a),
8495 None => (
8496 e.matmul_decode_exact(&la.ssm_beta, h, 1)?,
8497 e.matmul_decode_exact(&la.ssm_alpha, h, 1)?,
8498 ),
8499 },
8500 }
8501 } else {
8502 // fused-t twin (9B stores beta/alpha as Q8_0): same shared-quantize + one launch
8503 // contract as the wqkv pair above; 35B beta/alpha are Float -> None -> fallback.
8504 let mut nvfp4_fused = None;
8505 let mut q8_fused = None;
8506 if let Some((hq, hd)) = hq8_any {
8507 if t == 3 && std::env::var("MEMRA_NVFP4_AUX_DUAL").as_deref() != Ok("0") {
8508 nvfp4_fused =
8509 e.matmul_decode_exact_dual_pre(&la.ssm_beta, &la.ssm_alpha, hq, hd, t)?;
8510 if nvfp4_fused.is_some() && std::env::var("MEMRA_DEBUG").is_ok() {
8511 static ONCE: std::sync::Once = std::sync::Once::new();
8512 ONCE.call_once(|| {
8513 eprintln!("[memra] NVFP4 beta+alpha batched aux dual ENGAGED (t={t})")
8514 });
8515 }
8516 }
8517 if nvfp4_fused.is_none() && spec_fused_t() && (2..=4).contains(&t) {
8518 q8_fused = e.matmul_q8_fused2_t(&la.ssm_beta, &la.ssm_alpha, hq, hd, t)?;
8519 }
8520 }
8521 if let Some(((mut b, bs), (mut a, as_))) = nvfp4_fused {
8522 if bs != 1.0 {
8523 e.scale_inplace(&mut b, bs, t * la.ssm_beta.out_features())?;
8524 }
8525 if as_ != 1.0 {
8526 e.scale_inplace(&mut a, as_, t * la.ssm_alpha.out_features())?;
8527 }
8528 (b, a)
8529 } else if let Some(pair) = q8_fused {
8530 pair
8531 } else {
8532 match hq8_any {
8533 Some((hq, hd)) if h_q8.is_some() => (
8534 e.matmul_decode_exact_pre(&la.ssm_beta, hq, hd, t)?,
8535 e.matmul_decode_exact_pre(&la.ssm_alpha, hq, hd, t)?,
8536 ),
8537 _ => (
8538 e.matmul_decode_exact(&la.ssm_beta, h, t)?,
8539 e.matmul_decode_exact(&la.ssm_alpha, h, t)?,
8540 ),
8541 }
8542 }
8543 };
8544
8545 // conv with CARRIED state + ring roll (T >= pad rides the input-column update kernel;
8546 // T < pad — the MEMRA_SPEC_M2 t=2 arm — rolls via the pure-copy ring rebuild).
8547 let rl = cache.recur[il].as_mut().unwrap();
8548 let mut conv_out = e.uninit(conv_dim * t)?;
8549 e.ssm_conv1d_tm_state(
8550 &qkv_mixed,
8551 &mut rl.conv_state,
8552 la.ssm_conv1d.float_data(),
8553 &mut conv_out,
8554 conv_dim,
8555 t,
8556 d_conv,
8557 )?;
8558
8559 // GDN prep via the prefill kernels (repack + L2 + sigmoid + glog), T-wide.
8560 let mut q_g = e.uninit(d_state * num_v * t)?;
8561 let mut k_g = e.uninit(d_state * num_v * t)?;
8562 let mut v_g = e.uninit(d_state * num_v * t)?;
8563 e.qkv_to_gdn_repack(
8564 &conv_out, &mut q_g, &mut k_g, &mut v_g, d_state, num_v, num_k, key_dim, t,
8565 )?;
8566 let mut q_l2 = e.uninit(d_state * num_v * t)?;
8567 e.l2_norm_decode(&q_g, &mut q_l2, d_state, num_v * t, eps)?;
8568 let mut k_l2 = e.uninit(d_state * num_v * t)?;
8569 e.l2_norm_decode(&k_g, &mut k_l2, d_state, num_v * t, eps)?;
8570 let mut beta = e.uninit(t * num_v)?;
8571 e.sigmoid(&beta_raw, &mut beta, t * num_v)?;
8572 let mut g_log = e.uninit(t * num_v)?;
8573 e.gdn_glog(
8574 &alpha,
8575 la.ssm_dt.float_data(),
8576 la.ssm_a.float_data(),
8577 &mut g_log,
8578 num_v,
8579 t,
8580 )?;
8581
8582 // ONE gdn_scan over T tokens from the carried state (internal sequential loop ==
8583 // T chained T=1 steps). Ping-pong the resident buffers like eager decode.
8584 let mut o = e.uninit(d_state * num_v * t)?;
8585 {
8586 let crate::cache::RecurLayer {
8587 ssm_state,
8588 ssm_state_alt,
8589 ..
8590 } = rl;
8591 e.gdn_scan_s128(
8592 &q_l2,
8593 &k_l2,
8594 &v_g,
8595 &g_log,
8596 &beta,
8597 ssm_state,
8598 ssm_state_alt,
8599 &mut o,
8600 num_v,
8601 t,
8602 scale,
8603 )?;
8604 }
8605 std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
8606
8607 // gated RMSNorm + out projection, T-wide. FUSED-QUANTIZE ARM (lane/vt-fixes fix 2,
8608 // mirroring the T=1 decode's launch-arc form): when ssm_out rides the q8_1 fast path,
8609 // emit q8_1 straight from the gated norm at nrows=num_v*t (row-indexed kernel, the
8610 // T-wide launch is the per-row program; kernel-check pins bit-identity vs
8611 // gated_rmsnorm -> quantize_q8_1 at T=1 and T=5) and feed the decode-exact dispatch
8612 // pre-quantized — one launch replaces norm + quantize. Fallback = the f32 chain.
8613 let out = if e.uses_q8_1_fast(&la.ssm_out) {
8614 let (gq, gd) =
8615 e.gated_rmsnorm_q8_1(&o, la.ssm_norm.float_data(), &z, d_state, num_v * t, eps)?;
8616 e.matmul_decode_exact_pre(&la.ssm_out, &gq, &gd, t)?
8617 } else {
8618 let mut gn = e.uninit(d_state * num_v * t)?;
8619 e.gated_rmsnorm(
8620 &o,
8621 la.ssm_norm.float_data(),
8622 &z,
8623 &mut gn,
8624 d_state,
8625 num_v * t,
8626 eps,
8627 )?;
8628 // DECODE-EXACT out-projection: same MMVQ path as the T=1 decode (ssm_out at m>=5
8629 // would fall to dp4a with a different FP reduction order — same class of bug as
8630 // the input projs).
8631 e.matmul_decode_exact(&la.ssm_out, &gn, t)?
8632 };
8633 let stash = if want_stash {
8634 Some(GdnStash {
8635 qkv_mixed,
8636 q_l2,
8637 k_l2,
8638 v_g,
8639 g_log,
8640 beta,
8641 })
8642 } else {
8643 None
8644 };
8645 Ok((out, stash))
8646 }
8647
8648 /// REPLAY-FREE partial-accept commit (2026-07-03): make the cache state == "committed through
8649 /// the first `j` verify columns" WITHOUT the legacy rollback + duplicate trunk replay.
8650 /// - Full-attn KV: truncate both the owning-stage shadow and every TP rank to snapshot + j.
8651 /// The verify's appended rows for those columns are bit-identical to what an eager T=1
8652 /// chain writes (the decode-exact contract the verify-probe gates), so keeping them ==
8653 /// replaying them.
8654 /// - Linear layers, batched path: rebuild the conv ring by PURE COPIES (ring holds raw input
8655 /// columns) and the ssm state by a prefix re-run of the SAME gdn_scan kernel (t=j) from the
8656 /// snapshot state over the stash's identical inputs — the kernel's t-loop carries state in
8657 /// registers and writes it once at the end, so iterations 0..j-1 are independent of T:
8658 /// bit-identical to the verify's own state after j tokens == the eager chain state.
8659 /// - Linear layers, per-column path: restore the cloned actual state after column j-1.
8660 /// Caller guarantees 1 <= j <= t-1 (j==0 rounds take the legacy rollback; j==t is full accept).
8661 fn commit_verified_prefix(
8662 &self,
8663 e: &Engine,
8664 cache: &mut Cache,
8665 snap: &crate::cache::CacheSnapshot,
8666 ckpt: &VerifyCkpt,
8667 j: usize,
8668 kv_lens_done: bool,
8669 dev_j: Option<(&CudaSlice<u32>, usize, usize)>,
8670 ) -> Result<(), Box<dyn std::error::Error>> {
8671 // GDN geometry derives lazily inside recurrent-layer arms. Full-attention plans carry no
8672 // recurrent state and must never be forced through a synthetic SSM geometry.
8673 // Engine-bundle slice 1 (DSF-ROUNDCOST-20260820 §1.1): the per-column-arm restores
8674 // are 2 tiny D2D copies per linear layer (~96 dispatches/partial round on the q38
8675 // route). When every cols-arm layer shares uniform state sizes (single ssm cfg —
8676 // always true today), batch them into two `copy_batch_uniform_f32` launches. Bytes,
8677 // buffers and stream order are identical to the per-layer memcpy sequence; the
8678 // kernel-rebuild (gdn-stash) arm below is untouched. MEMRA_STATE_COPY_BATCH=0 reverts.
8679 let mut batched_cols = false;
8680 if state_copy_batch_on() && dev_j.is_none() {
8681 use cudarc::driver::DevicePtr;
8682 let s = &e.gpu.stream();
8683 let mut conv_pairs: Vec<(u64, u64)> = Vec::new();
8684 let mut ssm_pairs: Vec<(u64, u64)> = Vec::new();
8685 let (mut conv_words, mut ssm_words) = (0usize, 0usize);
8686 let mut uniform = true;
8687 for il in 0..self.layers.len() {
8688 let Some(rl) = cache.recur[il].as_ref() else {
8689 continue;
8690 };
8691 if ckpt.gdn[il].is_some() {
8692 continue; // kernel-rebuild arm restores below, per layer
8693 }
8694 let Some(cols) = &ckpt.cols[il] else {
8695 continue; // missing-ckpt error surfaces in the main loop
8696 };
8697 let (c, st) = &cols[j - 1];
8698 if conv_pairs.is_empty() {
8699 conv_words = c.len();
8700 ssm_words = st.len();
8701 } else if c.len() != conv_words || st.len() != ssm_words {
8702 uniform = false;
8703 break;
8704 }
8705 let (pc, _g0) = c.device_ptr(s);
8706 let (dc, _g1) = rl.conv_state.device_ptr(s);
8707 let (ps, _g2) = st.device_ptr(s);
8708 let (ds, _g3) = rl.ssm_state.device_ptr(s);
8709 conv_pairs.push((pc as u64, dc as u64));
8710 ssm_pairs.push((ps as u64, ds as u64));
8711 }
8712 if uniform && !conv_pairs.is_empty() {
8713 let n = conv_pairs.len();
8714 let mut t = vec![0u64; 2 * n];
8715 for (k, &(src, dst)) in conv_pairs.iter().enumerate() {
8716 t[k] = src;
8717 t[n + k] = dst;
8718 }
8719 let conv_t = e.htod_u64(&t)?;
8720 for (k, &(src, dst)) in ssm_pairs.iter().enumerate() {
8721 t[k] = src;
8722 t[n + k] = dst;
8723 }
8724 let ssm_t = e.htod_u64(&t)?;
8725 e.copy_batch_uniform_f32(&conv_t, n, conv_words)?;
8726 e.copy_batch_uniform_f32(&ssm_t, n, ssm_words)?;
8727 batched_cols = true;
8728 }
8729 }
8730 rewind_tp_kv_verified_prefix(&mut cache.tp_kv, &snap.tp_kv_len, j)?;
8731 for il in 0..self.layers.len() {
8732 if let (Some(kvl), Some(saved)) = (cache.kv[il].as_mut(), snap.kv_len[il]) {
8733 kvl.len = saved + j;
8734 // devacc 3a: spec_rollback_kv already wrote len_d on-device (same value).
8735 if !kv_lens_done {
8736 e.set_i32_one(&mut kvl.len_d, kvl.len as i32)?;
8737 }
8738 }
8739 if let Some(rl) = cache.recur[il].as_mut() {
8740 let Mixer::Linear(linear) = &self.layers[il].mixer else {
8741 return Err(format!("recurrent cache layer {il} has no GDN plan").into());
8742 };
8743 let geometry = linear.geometry;
8744 let d_state = geometry.key_head_dim as usize;
8745 let num_k = geometry.key_heads as usize;
8746 let num_v = geometry.value_heads as usize;
8747 let d_conv = geometry.conv_kernel as usize;
8748 let conv_dim = d_state * num_k * 2 + geometry.value_head_dim as usize * num_v;
8749 let scale = 1.0 / (d_state as f32).sqrt();
8750 if let Some(st) = &ckpt.gdn[il] {
8751 let ring_old = snap.conv[il].as_ref().expect("snapshot missing conv");
8752 let state_in = snap.ssm[il].as_ref().expect("snapshot missing ssm");
8753 if let Some((acc, base, t_v)) = dev_j {
8754 // 3b: j read on-device (_dc twins, same bodies; full accept early-exits).
8755 e.ssm_conv_ring_rebuild_dc(
8756 &st.qkv_mixed,
8757 ring_old,
8758 &mut rl.conv_state,
8759 conv_dim,
8760 acc,
8761 base,
8762 t_v,
8763 d_conv,
8764 )?;
8765 let mut o = e.uninit(d_state * num_v * j.max(1))?;
8766 e.gdn_scan_s128_dc(
8767 &st.q_l2,
8768 &st.k_l2,
8769 &st.v_g,
8770 &st.g_log,
8771 &st.beta,
8772 state_in,
8773 &mut rl.ssm_state,
8774 &mut o,
8775 num_v,
8776 acc,
8777 base,
8778 t_v,
8779 scale,
8780 )?;
8781 } else {
8782 e.ssm_conv_ring_rebuild(
8783 &st.qkv_mixed,
8784 ring_old,
8785 &mut rl.conv_state,
8786 conv_dim,
8787 j,
8788 d_conv,
8789 )?;
8790 let mut o = e.uninit(d_state * num_v * j)?; // scan output, discarded
8791 e.gdn_scan_s128(
8792 &st.q_l2,
8793 &st.k_l2,
8794 &st.v_g,
8795 &st.g_log,
8796 &st.beta,
8797 state_in,
8798 &mut rl.ssm_state,
8799 &mut o,
8800 num_v,
8801 j,
8802 scale,
8803 )?;
8804 }
8805 } else if let Some(cols) = &ckpt.cols[il] {
8806 if !batched_cols {
8807 let (c, s) = &cols[j - 1];
8808 e.copy_into(&mut rl.conv_state, 0, c, c.len())?;
8809 e.copy_into(&mut rl.ssm_state, 0, s, s.len())?;
8810 }
8811 } else {
8812 return Err(
8813 "commit_verified_prefix: verify ckpt missing for linear layer".into(),
8814 );
8815 }
8816 }
8817 }
8818 cache.pos = snap.pos + j;
8819 Ok(())
8820 }
8821
8822 /// ROUND-STREAM: recur restore with device-j (the _dc twins; full accept early-exits
8823 /// in-kernel). Requires the batched-linear stash on every linear layer (stream gate).
8824 fn commit_verified_prefix_stream(
8825 &self,
8826 e: &Engine,
8827 cache: &mut Cache,
8828 snap: &crate::cache::CacheSnapshot,
8829 ckpt: &VerifyCkpt,
8830 acc: &CudaSlice<u32>,
8831 base: usize,
8832 t_v: usize,
8833 ) -> Result<(), Box<dyn std::error::Error>> {
8834 for il in 0..self.layers.len() {
8835 if let Some(rl) = cache.recur[il].as_mut() {
8836 let Mixer::Linear(linear) = &self.layers[il].mixer else {
8837 return Err(format!("recurrent cache layer {il} has no GDN plan").into());
8838 };
8839 let geometry = linear.geometry;
8840 let d_state = geometry.key_head_dim as usize;
8841 let num_k = geometry.key_heads as usize;
8842 let num_v = geometry.value_heads as usize;
8843 let d_conv = geometry.conv_kernel as usize;
8844 let conv_dim = d_state * num_k * 2 + geometry.value_head_dim as usize * num_v;
8845 let scale = 1.0 / (d_state as f32).sqrt();
8846 let st = ckpt.gdn[il]
8847 .as_ref()
8848 .ok_or("stream restore: batched-linear stash missing")?;
8849 let ring_old = snap.conv[il].as_ref().expect("snapshot missing conv");
8850 let state_in = snap.ssm[il].as_ref().expect("snapshot missing ssm");
8851 e.ssm_conv_ring_rebuild_dc(
8852 &st.qkv_mixed,
8853 ring_old,
8854 &mut rl.conv_state,
8855 conv_dim,
8856 acc,
8857 base,
8858 t_v,
8859 d_conv,
8860 )?;
8861 let mut o = e.uninit(d_state * num_v * t_v)?;
8862 e.gdn_scan_s128_dc(
8863 &st.q_l2,
8864 &st.k_l2,
8865 &st.v_g,
8866 &st.g_log,
8867 &st.beta,
8868 state_in,
8869 &mut rl.ssm_state,
8870 &mut o,
8871 num_v,
8872 acc,
8873 base,
8874 t_v,
8875 scale,
8876 )?;
8877 }
8878 }
8879 Ok(())
8880 }
8881
8882 /// EAGLE3 aux-capturing verify forward over `tokens` (T) — mirrors `decode_step_t_h` exactly
8883 /// (same KV append, same causal verify, same recur advance) but ALSO clones the aux residual-
8884 /// stream hiddens (blocks in `aux_layers`) for TWO columns: the LAST column (always) and the
8885 /// optional `pred_col` (the EAGLE seed = bonus's predecessor). Returns
8886 /// (all_T_logits host, last_col_aux, pred_col_aux?). Used by the EAGLE3 orchestrator's commit.
8887 pub fn decode_step_t_aux2(
8888 &self,
8889 e: &Engine,
8890 tokens: &[u32],
8891 pos0: usize,
8892 cache: &mut Cache,
8893 aux_layers: &[usize],
8894 pred_col: Option<usize>,
8895 ) -> Result<
8896 (Vec<f32>, Vec<CudaSlice<f32>>, Option<Vec<CudaSlice<f32>>>),
8897 Box<dyn std::error::Error>,
8898 > {
8899 let cfg = &self.cfg;
8900 let n_embd = cfg.n_embd as usize;
8901 let eps = cfg.rms_eps;
8902 let t = tokens.len();
8903 let pos_vec: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
8904 let pos_d = e.htod_i32(&pos_vec)?;
8905 let mut x = e.htod(&self.embd.gather(n_embd, tokens))?;
8906 let mut aux_last: Vec<CudaSlice<f32>> = Vec::with_capacity(aux_layers.len());
8907 let mut aux_pred: Vec<CudaSlice<f32>> = Vec::new();
8908 let want_pred = pred_col.is_some();
8909
8910 for (il, layer) in self.layers.iter().enumerate() {
8911 // DISPATCH-MIRRORED norms (FP-order lesson #8) — see decode_step_t_h_emb.
8912 let mixer_fast = self.mixer_in_q8_1_fast(e, &layer.mixer);
8913 let norm_fused = std::env::var("MEMRA_NO_FUSE_NORMQ").is_err() && mixer_fast;
8914 let mut h = vbuf(e, t * n_embd)?; // fully written by either rms_norm arm
8915 if norm_fused {
8916 e.rms_norm_decode(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
8917 } else {
8918 e.rms_norm(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
8919 }
8920 let mixed = match &layer.mixer {
8921 Mixer::Full(fa) => {
8922 self.full_attn_verify(e, fa, &h, None, &pos_d, t, cache, il, None)?
8923 }
8924 Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
8925 Mixer::Linear(la) => {
8926 let mut out = e.zeros(t * n_embd)?;
8927 for col in 0..t {
8928 let mut h_col = e.zeros(n_embd)?;
8929 let src = h.slice(col * n_embd..(col + 1) * n_embd);
8930 e.copy_view_into(&mut h_col, 0, &src, n_embd)?;
8931 let m_col = self.linear_attn_decode(e, la, &h_col, cache, il)?;
8932 e.copy_into(&mut out, col * n_embd, &m_col, n_embd)?;
8933 }
8934 out
8935 }
8936 };
8937 let ffn_fuse = match &layer.ffn {
8938 crate::hybrid::Ffn::Dense {
8939 ffn_gate, ffn_up, ..
8940 } => {
8941 std::env::var("MEMRA_NO_FUSE_NORMQ").is_err()
8942 && e.uses_q8_1_fast(ffn_gate)
8943 && e.uses_q8_1_fast(ffn_up)
8944 }
8945 crate::hybrid::Ffn::Moe(_) => false,
8946 };
8947 let mut x1 = vbuf(e, t * n_embd)?; // fully written by add / add_rms_norm
8948 let mut z = vbuf(e, t * n_embd)?; // fully written by rms_norm_decode / add_rms_norm
8949 if ffn_fuse {
8950 e.add(&x, &mixed, &mut x1, t * n_embd)?;
8951 e.rms_norm_decode(
8952 &x1,
8953 layer.post_attn_norm.float_data(),
8954 &mut z,
8955 n_embd,
8956 t,
8957 eps,
8958 )?;
8959 } else {
8960 e.add_rms_norm(
8961 &x,
8962 &mixed,
8963 layer.post_attn_norm.float_data(),
8964 &mut x1,
8965 &mut z,
8966 n_embd,
8967 t,
8968 eps,
8969 )?;
8970 }
8971 let ffn_out = match &layer.ffn {
8972 crate::hybrid::Ffn::Dense {
8973 ffn_gate,
8974 ffn_up,
8975 ffn_down,
8976 } => {
8977 let n_ff = ffn_gate.out_features();
8978 let gate = e.matmul_decode_exact(ffn_gate, &z, t)?;
8979 let up = e.matmul_decode_exact(ffn_up, &z, t)?;
8980 let mut act = vbuf(e, t * n_ff)?; // fully written by ffn_act_lim
8981 // dense FFN clamp = the SHEXP array (upstream build_ffn serves both).
8982 Self::ffn_act_lim(
8983 e,
8984 &self.cfg,
8985 &gate,
8986 &up,
8987 1.0,
8988 1.0,
8989 self.cfg.clamp_shexp_at(il as u32),
8990 &mut act,
8991 t * n_ff,
8992 )?;
8993 e.matmul_decode_exact(ffn_down, &act, t)?
8994 }
8995 crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il(e, m, &z, t, il as u16)?,
8996 };
8997 let mut x2 = vbuf(e, t * n_embd)?; // fully written by add
8998 e.add(&x1, &ffn_out, &mut x2, t * n_embd)?;
8999 if aux_layers.contains(&il) {
9000 let mut a = e.zeros(n_embd)?;
9001 e.copy_view_into(&mut a, 0, &x2.slice((t - 1) * n_embd..t * n_embd), n_embd)?;
9002 aux_last.push(a);
9003 if let Some(pc) = pred_col {
9004 let mut ap = e.zeros(n_embd)?;
9005 e.copy_view_into(
9006 &mut ap,
9007 0,
9008 &x2.slice(pc * n_embd..(pc + 1) * n_embd),
9009 n_embd,
9010 )?;
9011 aux_pred.push(ap);
9012 }
9013 }
9014 x = x2;
9015 }
9016 let mut hn = vbuf(e, t * n_embd)?; // fully written by rms_norm_decode
9017 e.rms_norm_decode(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
9018 let logits = e.matmul_decode_exact(&self.output, &hn, t)?;
9019 let host = e.dtoh(&logits)?;
9020 cache.pos += t;
9021 Ok((
9022 host,
9023 aux_last,
9024 if want_pred { Some(aux_pred) } else { None },
9025 ))
9026 }
9027
9028 /// step35 SPEC-VERIFY attention over T query tokens — a per-row REPLAY of the eager
9029 /// `step35_decode_attn`.
9030 ///
9031 /// WHY A REPLAY AND NOT A BATCHED TWIN. The verify's whole job is to be bit-identical to what
9032 /// the eager decode would have computed for the same tokens; that is what makes greedy spec
9033 /// decode exact (run-spec asserts token identity for K=1..8). Every other verify arm in this
9034 /// file earns that identity by carefully mirroring dispatch (`matmul_decode_exact` to force
9035 /// MMVQ at any m, per-layer `ffn_fuse` mirroring, per-row `fa_decode` key bounds). step35
9036 /// stacks FOUR more per-layer degrees of freedom on top of that — per-layer `n_head`
9037 /// (64 full / 96 SWA), per-layer rotary width (64 full / 128 SWA), per-layer rope base, and a
9038 /// SEPARATE `attn_gate` tensor whose projection shares the attn-normed input — and its SWA
9039 /// layers attend through a token-OFFSET view whose offset is a function of the ABSOLUTE
9040 /// position of each query row. A batched twin would have to reproduce all of that AND the
9041 /// per-row offset in one launch; the offset alone rules out the existing rows kernels (they
9042 /// take one `base_len`, not a per-row offset).
9043 ///
9044 /// So this arm calls the eager path itself, once per row, on the same cache. Identity is then
9045 /// true BY CONSTRUCTION rather than by mirroring: row r runs exactly the kernel sequence that
9046 /// eager decode step r runs (same projections, same q8_1 fusion decision, same append, same
9047 /// view arithmetic, same `fa_decode_kvmod`, same gate), because it IS that code. Cost: T x the
9048 /// eager decode mixer instead of one batched pass — the same trade the generic arm's `else`
9049 /// per-row loop already accepts when `fa_rows_eligible` says no. Correctness first; a batched
9050 /// step35 twin is a perf lane's job and must be gated against this arm.
9051 ///
9052 /// The `h_q8` pre-quantized pair from the caller's fused norm is NOT forwarded: it is a
9053 /// T-row buffer and `step35_decode_attn`'s `pre_q` contract is one row. Instead each row's
9054 /// f32 `h` slice is handed over and the callee re-derives its own q8_1 exactly as eager decode
9055 /// does (`quantize_q8_1(h, 1, n_embd)`) — which is the dispatch being mirrored. Callers that
9056 /// took the fused arm therefore MUST still pass a live `h`; `step35_verify` asserts that.
9057 #[allow(clippy::too_many_arguments)]
9058 fn step35_verify(
9059 &self,
9060 e: &Engine,
9061 fa: &FullAttnLayer,
9062 h: &CudaSlice<f32>,
9063 h_q8: Option<(&CudaSlice<i8>, &CudaSlice<f32>)>,
9064 t: usize,
9065 cache: &mut Cache,
9066 il: usize,
9067 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9068 let n_embd = self.cfg.n_embd as usize;
9069 // The fused (h-less) attn-norm arm hands `h` as a zero-length placeholder. This arm needs
9070 // the f32 rows, so the caller must not take that lever for step35 — enforced at the call
9071 // site by the sliding-gated-MoE `Mixer::Full(_) => false` arm of
9072 // `lin_q8_only` in `decode_step_t_core_stream`, and asserted here so a future caller
9073 // cannot regress it into silently reading an empty buffer.
9074 assert_eq!(
9075 h.len(),
9076 t * n_embd,
9077 "step35_verify needs the f32 attn-normed rows ([t*n_embd]); the caller took the \
9078 fused q8-only norm arm (h_q8={}) — step35 must stay on the unfused arm",
9079 h_q8.is_some()
9080 );
9081 // ROW WIDTH IS n_embd, NOT n_head*head_dim: `step35_decode_attn` returns the mixer output
9082 // AFTER `wo`, so a row is [n_embd] — the same contract the generic arm's
9083 // `matmul_decode_exact(&fa.wo, &attn_g, t)` return has. Sizing this buffer from the
9084 // per-layer head geometry (8192 on full-attn, 12288 on SWA) instead overran the row on the
9085 // FIRST copy and panicked inside `copy_into`'s `CudaView::slice` unwrap
9086 // (raw/mtp-bt-20260806T212127Z.log frames 12-13).
9087 let mut out = vbuf(e, t * n_embd)?; // each row fully written by the copy below
9088 for r in 0..t {
9089 // Absolute position of this query row. `cache.pos` is the committed length at round
9090 // start and every row before r has already been appended by this loop, so the r-th
9091 // verify token sits at cache.pos + r — the same position eager decode would give it.
9092 let pos_d = e.htod_i32(&[(cache.pos + r) as i32])?;
9093 let mut h_row = vbuf(e, n_embd)?; // fully written by copy_view_into
9094 e.copy_view_into(
9095 &mut h_row,
9096 0,
9097 &h.slice(r * n_embd..(r + 1) * n_embd),
9098 n_embd,
9099 )?;
9100 // THE eager decode mixer: appends this row's K/V at kvl.len, advances it, then
9101 // attends over the (SWA-offset) view. Post-`wo`, same contract as this fn returns.
9102 let o = self.step35_decode_attn(e, fa, il, &h_row, None, &pos_d, cache)?;
9103 debug_assert_eq!(
9104 o.len(),
9105 n_embd,
9106 "step35_decode_attn returns post-wo [n_embd]"
9107 );
9108 e.copy_into(&mut out, r * n_embd, &o, n_embd)?;
9109 }
9110 Ok(out)
9111 }
9112
9113 /// Full-attention mixer over T query tokens with a GROWING resident KV (verify path, §D.3).
9114 /// Appends the T new K/V columns to cache.kv[il] then attends causally over [0..len) via
9115 /// fa_prefill. Token-major [T, kv_dim] projection layout == cache row layout (single copy).
9116 #[allow(clippy::too_many_arguments)]
9117 fn full_attn_verify(
9118 &self,
9119 e: &Engine,
9120 fa: &FullAttnLayer,
9121 h: &CudaSlice<f32>,
9122 h_q8: Option<(&CudaSlice<i8>, &CudaSlice<f32>)>,
9123 pos_d: &CudaSlice<i32>,
9124 t: usize,
9125 cache: &mut Cache,
9126 il: usize,
9127 stream_ctr: Option<&CudaSlice<i32>>,
9128 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9129 // step35: the generic geometry below is wrong for this arch (per-layer n_head, partial
9130 // per-layer rope, the SWA offset view, and a SEPARATE head-wise gate tensor), so it takes
9131 // its own arm. A verify that silently computes different attention than decode defeats the
9132 // whole self-consistency gate, so the arm is a per-row REPLAY of `step35_decode_attn`
9133 // rather than a batched twin — see `step35_verify` for why that is the exactness-correct
9134 // shape and not laziness.
9135 if self.sliding_gated_moe_batch_program() {
9136 if stream_ctr.is_some() {
9137 return Err(
9138 "step35 has no ROUND-STREAM verify arm (the device-counter _dc twins \
9139 cannot express the SWA offset KV view; same root cause as the dc \
9140 decode refusal) — run spec without the stream arm"
9141 .into(),
9142 );
9143 }
9144 return self.step35_verify(e, fa, h, h_q8, t, cache, il);
9145 }
9146 let cfg = &self.cfg;
9147 let geometry = cfg.full_attention_geometry_at(il as u32);
9148 let n_head = geometry.n_head as usize;
9149 let n_head_kv = geometry.n_head_kv as usize;
9150 let head_dim = geometry.head_dim_k as usize;
9151 let eps = cfg.rms_eps;
9152 let scale = geometry.attention_scale();
9153 let n_embd = cfg.n_embd as usize;
9154
9155 // DECODE-EXACT Q/K/V projections: matmul_decode_exact forces the MMVQ (warp-per-row) path
9156 // for every m, matching the T=1 decode's FP accumulation order. matmul_pre at m>=5 would
9157 // fall to dp4a (128-thread, two-level reduce) with a different FP sum order.
9158 // Q8 TRUNK-FUSION at T=1: DISPATCH-MIRRORS the eager decode's fused3 (bit-identical body).
9159 // BATCHED EPILOGUE RE-FUSE (lane/vt-fixes fix 2): `h_q8` = the attn-input norm's q8_1
9160 // form emitted by the fused rms_norm_q8_1 (bit-identical to rms_norm_decode ->
9161 // quantize_q8_1, kernel-check-pinned). When present (caller checked mixer q8_1-fast),
9162 // every projection consumes it — `h` may be a zero-len placeholder and must not be read.
9163 let (qf, mut k, v) = {
9164 let mut fused = None;
9165 let qkv_fast =
9166 e.uses_q8_1_fast(&fa.wq) && e.uses_q8_1_fast(&fa.wk) && e.uses_q8_1_fast(&fa.wv);
9167 if t == 1 && qkv_fast {
9168 let (hq_o, hd_o);
9169 let (hq, hd): (&CudaSlice<i8>, &CudaSlice<f32>) = match h_q8 {
9170 Some(p) => p,
9171 None => {
9172 (hq_o, hd_o) = e.quantize_q8_1(h, 1, n_embd)?;
9173 (&hq_o, &hd_o)
9174 }
9175 };
9176 fused = e.matmul_q8_fused3(&fa.wq, &fa.wk, &fa.wv, hq, hd)?;
9177 } else if spec_fused_t() && (2..=4).contains(&t) && qkv_fast {
9178 // VERIFY-TIER TRUNK FUSION (MEMRA_SPEC_FUSED_T): one shared quantize + one
9179 // fused3 batched launch replaces three decode-exact calls (3 re-quantizes of
9180 // the same h + 3 _b2/_b4 launches). Bit-identical per (tensor,token,row).
9181 let (hq_o, hd_o);
9182 let (hq, hd): (&CudaSlice<i8>, &CudaSlice<f32>) = match h_q8 {
9183 Some(p) => p,
9184 None => {
9185 (hq_o, hd_o) = e.quantize_q8_1(h, t, n_embd)?;
9186 (&hq_o, &hd_o)
9187 }
9188 };
9189 fused = e.matmul_q8_fused3_t(&fa.wq, &fa.wk, &fa.wv, hq, hd, t)?;
9190 }
9191 match (fused, h_q8) {
9192 (Some(triple), _) => triple,
9193 // shared pre-quantized activation (q8_1-fast guaranteed by the caller): the
9194 // decode-exact dispatch consumes (hq, hd) instead of re-quantizing 3x.
9195 (None, Some((hq, hd))) if qkv_fast => (
9196 e.matmul_decode_exact_pre(&fa.wq, hq, hd, t)?,
9197 e.matmul_decode_exact_pre(&fa.wk, hq, hd, t)?,
9198 e.matmul_decode_exact_pre(&fa.wv, hq, hd, t)?,
9199 ),
9200 (None, _) => (
9201 e.matmul_decode_exact(&fa.wq, h, t)?,
9202 e.matmul_decode_exact(&fa.wk, h, t)?,
9203 e.matmul_decode_exact(&fa.wv, h, t)?,
9204 ),
9205 }
9206 };
9207 // M3/Hy3 have no attention output gate — wq out is exactly q; skip the split.
9208 let gated = geometry.attention_gate == memra_gguf::config::AttentionGateKind::FusedQ;
9209 let (mut q, gate) = if gated {
9210 let mut q = vbuf(e, t * n_head * head_dim)?; // fully written by q_gate_split
9211 let mut gate = vbuf(e, t * n_head * head_dim)?; // fully written by q_gate_split
9212 e.q_gate_split(&qf, &mut q, &mut gate, head_dim, n_head, t)?;
9213 (q, Some(gate))
9214 } else {
9215 (qf, None)
9216 };
9217
9218 let mut qn = vbuf(e, t * n_head * head_dim)?; // fully written by rms_norm
9219 e.rms_norm(
9220 &q,
9221 fa.q_norm.float_data(),
9222 &mut qn,
9223 head_dim,
9224 n_head * t,
9225 eps,
9226 )?;
9227 q = qn;
9228 let mut kn = vbuf(e, t * n_head_kv * head_dim)?; // fully written by rms_norm
9229 e.rms_norm(
9230 &k,
9231 fa.k_norm.float_data(),
9232 &mut kn,
9233 head_dim,
9234 n_head_kv * t,
9235 eps,
9236 )?;
9237 k = kn;
9238 let rope_dims = geometry.n_rot as usize;
9239 e.rope_neox(
9240 &mut q,
9241 pos_d,
9242 head_dim,
9243 rope_dims,
9244 n_head,
9245 t,
9246 geometry.rope_base,
9247 1.0,
9248 )?;
9249 e.rope_neox(
9250 &mut k,
9251 pos_d,
9252 head_dim,
9253 rope_dims,
9254 n_head_kv,
9255 t,
9256 geometry.rope_base,
9257 1.0,
9258 )?;
9259
9260 // append T new K/V columns to the resident QUANTIZED cache. k/v are token-major [T, kv_dim]
9261 // f32; append-quantize each of the T token rows into the byte cache (q8_0 K / q5_1 V).
9262 let kvl = cache.kv[il].as_mut().unwrap();
9263 let (kv_dim_k, kv_dim_v, ktb, vtb) =
9264 (kvl.kv_dim_k, kvl.kv_dim_v, kvl.k_tok_bytes, kvl.v_tok_bytes);
9265 if let Some(ctr) = stream_ctr {
9266 // stream: ONE batched append at the device counter (rows kernel = the per-view warp
9267 // math on a (block, token) grid, documented byte-identical); host len is a stale
9268 // LOWER BOUND under pre-issue (drain reconciles it).
9269 e.append_kv_quantized_rows_dc(
9270 &k,
9271 &v,
9272 &mut kvl.k,
9273 &mut kvl.v,
9274 ctr,
9275 t,
9276 kv_dim_k,
9277 kv_dim_v,
9278 ktb,
9279 vtb,
9280 crate::Engine::kv_fp8_on(),
9281 )?;
9282 } else {
9283 for i in 0..t {
9284 let k_row = k.slice(i * kv_dim_k..(i + 1) * kv_dim_k);
9285 let v_row = v.slice(i * kv_dim_v..(i + 1) * kv_dim_v);
9286 e.append_kv_quantized_view(
9287 &k_row,
9288 &v_row,
9289 &mut kvl.k,
9290 &mut kvl.v,
9291 kvl.len + i,
9292 kv_dim_k,
9293 kv_dim_v,
9294 ktb,
9295 vtb,
9296 crate::Engine::kv_fp8_on(),
9297 )?;
9298 }
9299 kvl.len += t;
9300 }
9301
9302 // BIT-IDENTICAL VERIFY ATTENTION (spec-exactness fix): the FP accumulation order must be
9303 // byte-for-byte identical to the eager decode path. fa_prefill uses a different tile size
9304 // (BLOCK_Q=64, BK=32) and online-softmax structure than fa_decode's split-K + combine,
9305 // which changes FP summation order and can flip argmax at tight logit margins. Query row r
9306 // attends to keys [0..base_len+r+1) — each successive row sees one more key (the causal
9307 // property). This matches eager: decode appends k at len, then fa_decode sees t_kv = len+1
9308 // keys. The verify appends all T tokens first but bounds the key range per row.
9309 //
9310 // MULTI-ROW FUSED PATH (the long-ctx spec fix, 2026-07-03): when every row takes the vec
9311 // kernel (base_len+1 >= FA_VEC_MIN_TKV), ONE fa_decode_rows launch executes the exact
9312 // per-row program for all T rows (grid.z = row, per-row n_splits from the same
9313 // fa_split_keys formula) — replacing T x (2 launches + 2 dtod copies + 5 partial allocs)
9314 // and multiplying resident CTAs by T on a latency-bound kernel. Bit-identical per row by
9315 // construction; kernel-check pins rows-vs-loop byte identity, run-spec is the end gate.
9316 // Short ctx (any row below the vec crossover) and MEMRA_NO_FA_VEC/MEMRA_FA_ROWS_OFF keep the
9317 // per-row loop (whose fa_decode picks scalar/vec per row exactly like eager decode).
9318 let mut attn = vbuf(e, t * n_head * head_dim)?; // fully written by every FA arm below
9319 let base_len = kvl.len - t; // KV len BEFORE this round's T tokens were appended
9320 // T=1 INCLUDED (2026-07-05): p-min cuts the draft to 1 in ~75% of rounds on hard
9321 // (agentic) content — the old t>1 gate sent those rounds to the per-row loop (262us/row
9322 // + q-row copy + per-row allocs vs 93us/row through the fused kernel at grid.z=1, same
9323 // program). nsys accounting: 1088 of 1456 verify FA launches were T=1 escapees.
9324 // LEAN T=1 ARM (MEMRA_SPEC_LEAN, close35): at t==1, q IS one row and fa_decode on it is
9325 // the EXACT eager decode dispatch (vec_q_v2 + combine_f32; the rows pair measured +50us
9326 // at m=1). Byte-identical: kernel-check pins rows-vs-loop identity, and the per-row loop
9327 // at t=1 is fa_decode on the same q with zero-offset copies. Gates arbitrate.
9328 if let Some(ctr) = stream_ctr {
9329 // STREAM ARM: causal base from the device counter; host kvl.len is a stale lower
9330 // bound used only for the split-sizing upper bound (+64 slack covers M pre-issued
9331 // rounds at K<=8). Views span the bound; per-row limits derive in-kernel.
9332 let upper = kvl.len + t + 64;
9333 let k_view = e.view_u8(&kvl.k, (upper.min(cache.max_ctx)) * ktb);
9334 let v_view = e.view_u8(&kvl.v, (upper.min(cache.max_ctx)) * vtb);
9335 e.fa_decode_rows_dc(
9336 &q,
9337 &k_view,
9338 &v_view,
9339 &mut attn,
9340 head_dim,
9341 n_head,
9342 n_head_kv,
9343 ctr,
9344 upper.min(cache.max_ctx),
9345 t,
9346 scale,
9347 ktb,
9348 vtb,
9349 0,
9350 false,
9351 )?;
9352 } else if spec_lean() && t == 1 {
9353 let t_kv = base_len + 1;
9354 let k_view = e.view_u8(&kvl.k, t_kv * ktb);
9355 let v_view = e.view_u8(&kvl.v, t_kv * vtb);
9356 e.fa_decode_kvmod(
9357 &q,
9358 &k_view,
9359 &v_view,
9360 &mut attn,
9361 head_dim,
9362 n_head,
9363 n_head_kv,
9364 t_kv,
9365 scale,
9366 ktb,
9367 vtb,
9368 crate::Engine::kv_fp8_on(),
9369 )?;
9370 } else if e.fa_rows_eligible(base_len, head_dim) {
9371 let k_view = e.view_u8(&kvl.k, (base_len + t) * ktb);
9372 let v_view = e.view_u8(&kvl.v, (base_len + t) * vtb);
9373 e.fa_decode_rows(
9374 &q,
9375 &k_view,
9376 &v_view,
9377 &mut attn,
9378 head_dim,
9379 n_head,
9380 n_head_kv,
9381 base_len,
9382 t,
9383 scale,
9384 ktb,
9385 vtb,
9386 None,
9387 false,
9388 crate::Engine::kv_fp8_on(),
9389 None,
9390 )?;
9391 } else {
9392 for r in 0..t {
9393 let t_kv_r = base_len + r + 1; // this row sees keys [0..t_kv_r)
9394 let k_view_r = e.view_u8(&kvl.k, t_kv_r * ktb);
9395 let v_view_r = e.view_u8(&kvl.v, t_kv_r * vtb);
9396 // copy q row into an owned buffer (fa_decode takes &CudaSlice, not CudaView)
9397 let mut q_row = vbuf(e, n_head * head_dim)?; // fully written by copy_view_into
9398 let q_src = q.slice(r * n_head * head_dim..(r + 1) * n_head * head_dim);
9399 e.copy_view_into(&mut q_row, 0, &q_src, n_head * head_dim)?;
9400 let mut attn_row = vbuf(e, n_head * head_dim)?; // fully written by fa_decode
9401 e.fa_decode_kvmod(
9402 &q_row,
9403 &k_view_r,
9404 &v_view_r,
9405 &mut attn_row,
9406 head_dim,
9407 n_head,
9408 n_head_kv,
9409 t_kv_r,
9410 scale,
9411 ktb,
9412 vtb,
9413 crate::Engine::kv_fp8_on(),
9414 )?;
9415 e.copy_into(
9416 &mut attn,
9417 r * n_head * head_dim,
9418 &attn_row,
9419 n_head * head_dim,
9420 )?;
9421 }
9422 }
9423
9424 let attn_g = match &gate {
9425 Some(gate) => {
9426 let mut gsig = vbuf(e, t * n_head * head_dim)?; // fully written by sigmoid
9427 e.sigmoid(gate, &mut gsig, t * n_head * head_dim)?;
9428 let mut ag = vbuf(e, t * n_head * head_dim)?; // fully written by mul
9429 e.mul(&attn, &gsig, &mut ag, t * n_head * head_dim)?;
9430 ag
9431 }
9432 None => attn,
9433 };
9434 // DECODE-EXACT wo projection: at m>=5 (K=4+ with pending) the generic matmul would use dp4a
9435 // (128-thread, different FP sum order than MMVQ). Force MMVQ for bit-identity with decode.
9436 Ok(e.matmul_decode_exact(&fa.wo, &attn_g, t)?)
9437 }
9438
9439 /// Context-linear bytes for a plain serving session's trunk cache.
9440 pub fn plain_session_kv_bytes_per_token(&self) -> usize {
9441 crate::cache::cache_bytes_per_token_for_plan(
9442 &self.cfg,
9443 &self.plan,
9444 0,
9445 self.plan.layers.len(),
9446 )
9447 }
9448
9449 /// `(logical bytes/token, ring-capped bytes/token, ring row cap)` for exact admission.
9450 pub fn plain_session_kv_shape(&self) -> (usize, usize, usize) {
9451 (
9452 self.plain_session_kv_bytes_per_token(),
9453 crate::cache::cache_ring_bytes_per_token_for_plan(
9454 &self.cfg,
9455 &self.plan,
9456 0,
9457 self.plan.layers.len(),
9458 ),
9459 crate::cache::cache_ring_row_cap_for_plan(&self.plan),
9460 )
9461 }
9462
9463 /// Context-linear bytes for a speculative serving session: trunk cache plus persistent MTP
9464 /// scratch. With no MTP head this equals the plain coefficient.
9465 pub fn spec_session_kv_bytes_per_token(&self) -> usize {
9466 let scratch = self
9467 .mtp
9468 .iter()
9469 .chain(self.mtp_extra.iter())
9470 .map(|mtp| {
9471 let (_, _, k, v) = mtp_scratch_layout(&self.cfg, mtp.geom.as_ref());
9472 k + v
9473 })
9474 .sum::<usize>();
9475 self.plain_session_kv_bytes_per_token()
9476 .saturating_add(scratch)
9477 }
9478
9479 /// Spec twin of [`HybridModel::plain_session_kv_shape`]; Step35's persistent MTP scratch is
9480 /// capped by the same SWA ring rows as the trunk.
9481 pub fn spec_session_kv_shape(&self) -> (usize, usize, usize) {
9482 let total = self.spec_session_kv_bytes_per_token();
9483 let (_, mut ring, rows) = self.plain_session_kv_shape();
9484 if rows > 0 {
9485 ring = ring.saturating_add(
9486 self.mtp
9487 .iter()
9488 .chain(self.mtp_extra.iter())
9489 .map(|mtp| {
9490 let (_, _, k, v) = mtp_scratch_layout(&self.cfg, mtp.geom.as_ref());
9491 k + v
9492 })
9493 .sum::<usize>(),
9494 );
9495 }
9496 (total, ring, rows)
9497 }
9498
9499 /// Greedy MTP speculative decode (§B). Token-identical to `generate(prompt, max_new)` but uses
9500 /// the NextN head to draft K tokens then verifies them in one batched target forward.
9501 /// Returns (generated tokens, total_drafted, total_accepted) so the caller can report
9502 /// acceptance rate. `k` = draft length per round.
9503 ///
9504 /// GRAPH DRAFT (stage 2 of graph-grade spec): when the model is all-Dense and the MTP head is
9505 /// Dense (no MoE host readbacks), the fixed-shape T=1 MTP forward is CUDA-graph-captured ONCE
9506 /// and replayed per draft step — the ~40 eager launches per drafted token collapse into one
9507 /// graph dispatch; only the 4-byte token id (and 4-byte p-min confidence) round-trip per step.
9508 /// Event tracking is disabled for the whole call (generate_graph pattern) so every buffer the
9509 /// captured graph references is event-free; the spec loop is strictly single-stream.
9510 /// MEMRA_SPEC_NOGRAPH=1 forces the eager draft chain.
9511 /// SAMPLED mode (MEMRA_SPEC_TEMP>0) has its OWN capture (gumbel-perturbed in-graph argmax,
9512 /// device Philox event counter, persistent q retention) — graph-vs-eager sampled streams are
9513 /// bit-identical for the same (seed, prompt, K, temp); see the sampled-graph setup in
9514 /// generate_spec_inner2.
9515 /// Multi-turn session: trunk cache + MTP draft scratch persist across generate calls, so
9516 /// turn N+1 primes ONLY its new suffix (the 124k-conversation daily pattern — re-priming a
9517 /// 32k history costs ~54s; a suffix prime costs seconds). APPEND-ONLY by construction: the
9518 /// hybrid linear-attn states are in-place (no position index), so a session can extend but
9519 /// never rewind — `committed` is the exact token list whose state the caches hold (includes
9520 /// any overshoot tokens past max_new; the caller renders from `committed`, not its own echo).
9521 pub fn new_session(
9522 &self,
9523 e: &Engine,
9524 max_ctx: usize,
9525 ) -> Result<SpecSession, Box<dyn std::error::Error>> {
9526 Ok(SpecSession {
9527 // STAGE-OWNED KV (lane/pp2-spec 2026-08-06): `pp::new_cache`, not `Cache::new`. This
9528 // is the SERVING spec-session path, and with the ppN door open across two cards a
9529 // primary-homed cache makes every remote stage peer-read its OWN KV on every verify
9530 // round — the wrong-card class already fixed on the two batched serving paths
9531 // (worker.rs 2483 / 2837). With the door shut `new_cache` IS `Cache::new` (same
9532 // branch, same allocations), so single-device behavior is byte-unchanged.
9533 cache: crate::pp::new_cache_planned(e, &self.cfg, &self.plan, max_ctx)?,
9534 scratch: self.new_mtp_scratch(e, max_ctx)?,
9535 committed: Vec::new(),
9536 last_h: None,
9537 next_pred: None,
9538 sctr: 0,
9539 uctr: 0,
9540 draft_ctx: None,
9541 pending_tok: None,
9542 turn_ckpt: None,
9543 telem: SpecTelemetryCounters::default(),
9544 capture_at: None,
9545 boundary_captures: Vec::new(),
9546 ckpt_at: None,
9547 })
9548 }
9549
9550 /// SPEC-ON-CACHE-HIT restore (lane/spec-on-cache-hit, 2026-08-18 — PORT-PLAN item 3,
9551 /// research/cache-spec-design-20260814, scoped to WHOLE-ENTRY restores only): build a
9552 /// SpecSession around a trunk cache the worker already restored from a prefix-cache
9553 /// entry, re-installing the entry's published draft plane as the MTP scratch rows
9554 /// `[0..prefix.len())` and the entry's boundary hidden as `last_h`, then feeding the
9555 /// prompt SUFFIX here — through EXACTLY the plain path's program selection — so the
9556 /// worker always receives a fully-warm continuation session (committed = whole
9557 /// prompt, `next_pred` + `last_h` set; caller sets `next_pred` from the entry's
9558 /// boundary logits on the empty-suffix shape).
9559 ///
9560 /// PROGRAM LAW (the splitiso two-programs class, learned AGAIN in this lane's own
9561 /// gate): the identity target for a converted hit is the PLAIN hit serving the same
9562 /// request, and plain feeds a carried suffix via eager `decode_step` below
9563 /// PRIME_MIN_T and via `prime_cache` at/above it (prefill_tick's arms). The generate
9564 /// path's tokenwise arm routes qwen35-class through the BATCHED T=1 program
9565 /// (`spec_target_step_h`) instead — ULP-different suffix rows, and the gate measured
9566 /// the near-tie flip at generated token ~8 (research/spec-cache-20260818, qwen r3).
9567 /// So the suffix is fed HERE, mirroring prefill_tick arm-for-arm, not handed to the
9568 /// burst prime.
9569 ///
9570 /// SEED RULE (both sampling regimes; lane/sampled-hit-spec 2026-08-19, sampled draw
9571 /// added by lane/sampled-spec-quality 2026-08-19). The boundary token is produced by
9572 /// EXACTLY the rule the cold burst entry applies to its own first token from the same
9573 /// logits row: `argmax` when greedy, and a `sample_boundary_token` draw at Philox
9574 /// counter 0 when sampled. Both shapes are covered — the entry's boundary logits on a
9575 /// full-cover (empty-suffix) hit, this feed's own boundary logits on a suffix hit.
9576 /// That is what keeps a restored session seed-identical to a cold one PER SEED: the
9577 /// cold session draws from the identical row at counter 0 and then runs its rounds from
9578 /// counter 1, so the restored session admits with `sctr = 1` after its own draw.
9579 /// The WORKER owns the one refusal this constructor cannot see — a constrained request.
9580 /// (The penalized-sampled refusal was LIFTED once the burst's penalty window learned to
9581 /// span the session: `committed` here is the WHOLE prompt, so the restored session's
9582 /// window is the cold session's window. It comes back if `MEMRA_SPEC_PEN_SESSION=0`.)
9583 ///
9584 /// NOT the rolled-back partial-restore hazard: the caller restores at exactly the
9585 /// entry's captured endpoint (`e.pos`) through the shipping whole-entry path;
9586 /// mid-entry (`at < e.pos`) trunk restores stay behind MEMRA_PREFIX_PARTIAL_RESTORE
9587 /// and are never routed here.
9588 ///
9589 /// Failure contract: `Err((Some(cache), why))` before any trunk mutation — the
9590 /// worker rebuilds the plain carrier and the hit serves plain, byte-unchanged.
9591 /// `Err((None, why))` after the suffix feed began — the carrier is part-fed and
9592 /// UNUSABLE; the worker serves the request cold-plain (correct, slower) and the
9593 /// entry stays published for the next request.
9594 #[allow(clippy::too_many_arguments)]
9595 pub fn spec_session_from_restored(
9596 &self,
9597 e: &Engine,
9598 mut cache: Cache,
9599 prefix: Vec<u32>,
9600 suffix: &[u32],
9601 draft_k: &CudaSlice<u8>,
9602 draft_v: &CudaSlice<u8>,
9603 draft_k_tok_bytes: usize,
9604 draft_v_tok_bytes: usize,
9605 draft_len: usize,
9606 last_h: &[f32],
9607 // The ENTRY's boundary logits row (the full-cover shape's seed source). May be empty
9608 // when a suffix follows — the feed's own logits are the boundary then.
9609 boundary_logits: &[f32],
9610 // The request's sampler, or None for greedy. Owned here so the seed rule lives in
9611 // ONE place instead of being half-applied by the worker.
9612 sampling: Option<SpecSampling>,
9613 require_anchor: bool,
9614 max_ctx: usize,
9615 // STABLE-BOUNDARY REPUBLICATION (lane/frspec-multiturn-cache, 2026-08-21): ABSOLUTE
9616 // prompt position to split the suffix feed at and capture the extended-entry
9617 // publication + this session's `turn_ckpt` — the worker's stable pre-generation
9618 // boundary (`plain_checkpoint_boundary`). None = legacy prompt-end republication.
9619 // WHY: the prompt-end capture below includes the template's live generation header
9620 // (`<|im_start|>assistant\n<think>\n`), which the next turn's re-render replaces, so
9621 // for a hybrid (whole-entry restores only) every extended entry's last ~2 tokens
9622 // diverged from every future prompt and the hit boundary FROZE at the first
9623 // lcp-split entry forever (measured: cached 6811 of 38228 by turn 8, B4).
9624 republish_at: Option<usize>,
9625 ) -> Result<SpecSession, (Option<Cache>, String)> {
9626 let pos = prefix.len();
9627 let fail = |cache: Cache, msg: String| -> Result<SpecSession, (Option<Cache>, String)> {
9628 Err((Some(cache), msg))
9629 };
9630 if self.mtp.is_none() {
9631 return fail(cache, "no MTP head attached (nothing to draft with)".into());
9632 }
9633 if pos == 0 {
9634 return fail(cache, "empty committed prefix".into());
9635 }
9636 if cache.pos != pos {
9637 let msg = format!(
9638 "restored cache pos {} != restored prefix len {pos}",
9639 cache.pos
9640 );
9641 return fail(cache, msg);
9642 }
9643 if draft_len != pos {
9644 return fail(
9645 cache,
9646 format!("draft plane len {draft_len} != restored prefix len {pos}"),
9647 );
9648 }
9649 if pos + suffix.len() >= max_ctx {
9650 return fail(
9651 cache,
9652 format!(
9653 "prompt {} + suffix would not leave generation room in ctx {max_ctx}",
9654 pos + suffix.len(),
9655 ),
9656 );
9657 }
9658 let mut scratch = match MtpScratch::new(
9659 e,
9660 &self.cfg,
9661 &self.plan,
9662 max_ctx,
9663 self.mtp.as_ref().and_then(|m| m.geom.as_ref()),
9664 ) {
9665 Ok(s) => s,
9666 Err(err) => return fail(cache, format!("draft scratch alloc failed: {err}")),
9667 };
9668 if scratch.kv.ring.is_some() {
9669 return fail(
9670 cache,
9671 "ring-backed draft scratch (Step35 SWA) cannot take a flat prefix restore".into(),
9672 );
9673 }
9674 if scratch.kv.k_tok_bytes != draft_k_tok_bytes
9675 || scratch.kv.v_tok_bytes != draft_v_tok_bytes
9676 {
9677 return fail(
9678 cache,
9679 format!(
9680 "draft plane layout {draft_k_tok_bytes}/{draft_v_tok_bytes} != scratch \
9681 {}/{} bytes/token (stale entry across a format change)",
9682 scratch.kv.k_tok_bytes, scratch.kv.v_tok_bytes,
9683 ),
9684 );
9685 }
9686 if pos > scratch.cap {
9687 return fail(
9688 cache,
9689 format!(
9690 "draft plane rows {pos} exceed scratch capacity {}",
9691 scratch.cap
9692 ),
9693 );
9694 }
9695 let kb = pos * draft_k_tok_bytes;
9696 let vb = pos * draft_v_tok_bytes;
9697 if draft_k.len() < kb || draft_v.len() < vb {
9698 return fail(
9699 cache,
9700 format!(
9701 "truncated draft plane: K {} < {kb} or V {} < {vb} bytes",
9702 draft_k.len(),
9703 draft_v.len(),
9704 ),
9705 );
9706 }
9707 if kb > 0 {
9708 if let Err(err) = e.copy_u8_into(&mut scratch.kv.k, 0, draft_k, kb) {
9709 return fail(cache, format!("draft K restore copy failed: {err}"));
9710 }
9711 }
9712 if vb > 0 {
9713 if let Err(err) = e.copy_u8_into(&mut scratch.kv.v, 0, draft_v, vb) {
9714 return fail(cache, format!("draft V restore copy failed: {err}"));
9715 }
9716 }
9717 if let Err(err) = scratch.set_len(e, pos) {
9718 return fail(cache, format!("draft scratch len set failed: {err}"));
9719 }
9720 let mut last_h_dev = if last_h.len() == self.cfg.n_embd as usize {
9721 // anchor upload failure is acceptance-only when a suffix feed follows (fill
9722 // row-0 falls back to zeros) but FATAL for an empty-suffix continuation (the
9723 // burst entry asserts committed + last_h + next_pred) — the caller says which.
9724 e.htod(last_h).ok()
9725 } else {
9726 None
9727 };
9728 if require_anchor && last_h_dev.is_none() {
9729 return fail(
9730 cache,
9731 "empty-suffix continuation requires the entry's boundary hidden anchor".into(),
9732 );
9733 }
9734 let mut committed = prefix;
9735 // Set on BOTH shapes below (suffix-fed and full-cover) — never left None, which is
9736 // what the empty-suffix continuation assert in the burst entry requires.
9737 let next_pred;
9738 // Philox: (0,0) at admit exactly like a fresh session; a sampled boundary draw below
9739 // consumes counter 0 and leaves 1, which is the state a cold session reaches after
9740 // drawing its own first token from the same row.
9741 let mut sctr = 0u32;
9742 let sampled = sampling.is_some_and(|s| s.temp > 0.0) && spec_sampled_boundary_on();
9743 // Penalty window for the boundary draw: the last `penalty_last_n` tokens of the WHOLE
9744 // prompt, which is what the cold session's own burst sees (Item 2's window). Built
9745 // after the suffix joins `committed` below.
9746 let mut boundary_captures: Vec<SpecBoundaryCapture> = Vec::new();
9747 let mut restored_turn_ckpt: Option<SpecCheckpoint> = None;
9748 if !suffix.is_empty() {
9749 // ---- SUFFIX FEED, mirroring prefill_tick's program selection exactly ----
9750 // From here on the trunk cache mutates: failures return Err((None, _)) and
9751 // the worker serves the request cold-plain instead of reusing the carrier.
9752 let dirty =
9753 |msg: String| -> Result<SpecSession, (Option<Cache>, String)> { Err((None, msg)) };
9754 let n_embd = self.cfg.n_embd as usize;
9755 let t = suffix.len();
9756 let mut h_rows = match e.uninit(t * n_embd) {
9757 Ok(b) => b,
9758 Err(err) => return fail(cache, format!("suffix hidden buffer alloc: {err}")),
9759 };
9760 // STABLE-BOUNDARY split (see `republish_at`): feed stops at the boundary so the
9761 // in-place GDN conv/ssm state can be snapshotted there — the only moment it
9762 // exists (the cold prime-split law). suffix-relative; None = one-segment legacy.
9763 let b_rel = republish_at
9764 .and_then(|abs| abs.checked_sub(pos))
9765 .filter(|&r| r > 0 && r < t);
9766 let mut feed_logits = Vec::new();
9767 let tokenwise_env = std::env::var("MEMRA_PRIME_TOKENWISE").is_ok()
9768 || e.frozen_cpu_experts_prefer_tokenwise_prime();
9769 let mut fed = 0usize;
9770 for seg_end in [b_rel, Some(t)].into_iter().flatten() {
9771 if seg_end <= fed {
9772 continue;
9773 }
9774 let seg = &suffix[fed..seg_end];
9775 let batched = seg.len() >= crate::hybrid_forward::PRIME_MIN_T && !tokenwise_env;
9776 if batched {
9777 // prefill_tick's prime arm: request-level prime_cache call; tokens still
9778 // queued after this segment ride `queued_after` so Step35 arm selection
9779 // stays keyed to the request's end (tick-seg law).
9780 match self.prime_cache(e, seg, &mut cache, t - seg_end) {
9781 Ok((l, _h_seed, hiddens)) => {
9782 if let Err(err) =
9783 e.copy_into(&mut h_rows, fed * n_embd, &hiddens, seg.len() * n_embd)
9784 {
9785 return dirty(format!("suffix hidden copy: {err}"));
9786 }
9787 feed_logits = l;
9788 }
9789 Err(err) => return dirty(format!("suffix prime failed: {err}")),
9790 }
9791 } else {
9792 // prefill_tick's tokenwise arm: eager decode_step, one token at a time.
9793 for (i, &tok) in seg.iter().enumerate() {
9794 match self.decode_step_h(e, tok, &mut cache) {
9795 Ok((l, h)) => {
9796 if let Err(err) =
9797 e.copy_into(&mut h_rows, (fed + i) * n_embd, &h, n_embd)
9798 {
9799 return dirty(format!("suffix hidden copy: {err}"));
9800 }
9801 feed_logits = l;
9802 }
9803 Err(err) => return dirty(format!("suffix decode_step failed: {err}")),
9804 }
9805 }
9806 }
9807 fed = seg_end;
9808 if Some(seg_end) == b_rel {
9809 // The stable pre-generation boundary: capture the extended-entry
9810 // publication AND this session's own turn checkpoint here instead of at
9811 // prompt-end (both would otherwise carry the volatile live-header tail
9812 // the next re-render replaces). Failure silent, turn_ckpt convention.
9813 debug_assert_eq!(
9814 cache.pos,
9815 pos + seg_end,
9816 "stable-boundary capture off the feed split"
9817 );
9818 if spec_restore_republish_on() {
9819 if let Ok(snap) = cache.snapshot(e) {
9820 boundary_captures.push(SpecBoundaryCapture {
9821 snap,
9822 pos: pos + seg_end,
9823 logits: feed_logits.clone(),
9824 last_h: capture_boundary_hidden(e, &h_rows, seg_end, n_embd),
9825 });
9826 }
9827 }
9828 let anchor: Result<CudaSlice<f32>, Box<dyn std::error::Error>> =
9829 e.uninit(n_embd).and_then(|mut a| {
9830 e.copy_view_into(
9831 &mut a,
9832 0,
9833 &h_rows.slice((seg_end - 1) * n_embd..seg_end * n_embd),
9834 n_embd,
9835 )?;
9836 Ok(a)
9837 });
9838 if let (Ok(snap), Ok(last_h)) = (cache.snapshot(e), anchor) {
9839 restored_turn_ckpt = Some(SpecCheckpoint {
9840 snap,
9841 pos: pos + seg_end,
9842 last_h,
9843 });
9844 }
9845 }
9846 }
9847 // Draft-scratch fill for the suffix rows, predecessor-paired: row `pos` reads
9848 // the entry's boundary anchor (zeros fallback — acceptance-only), row `pos+i`
9849 // reads h_rows[i-1]. Chunked like the generate path's fill (transients scale
9850 // with T). Fill failures are acceptance-only — truncate to the restored rows
9851 // and continue; the burst's own set_len keeps the invariant.
9852 let mtp = self.mtp.as_ref().expect("mtp checked above");
9853 let (embd_qt, embd_rb) = self.embd.qt_and_row_bytes(n_embd);
9854 let embd_gpu = if spec_host_embd() {
9855 None
9856 } else {
9857 Some(
9858 self.embd_gpu
9859 .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload")),
9860 )
9861 };
9862 let embd_dev = embd_gpu.map(|g| (g, embd_qt, embd_rb));
9863 let fill_chunk = 4096usize;
9864 let mut filled = true;
9865 let mut start = 0usize;
9866 'fill: while start < t {
9867 let end = (start + fill_chunk).min(t);
9868 let tc = end - start;
9869 let Ok(mut phs) = e.zeros(tc * n_embd) else {
9870 filled = false;
9871 break 'fill;
9872 };
9873 let (src_lo, dst_off, n_copy) = if start == 0 {
9874 (0, n_embd, (tc - 1) * n_embd)
9875 } else {
9876 ((start - 1) * n_embd, 0, tc * n_embd)
9877 };
9878 if start == 0 {
9879 if let Some(lh) = last_h_dev.as_ref() {
9880 if e.copy_into(&mut phs, 0, lh, n_embd).is_err() {
9881 filled = false;
9882 break 'fill;
9883 }
9884 }
9885 }
9886 if n_copy > 0
9887 && e.copy_view_into(
9888 &mut phs,
9889 dst_off,
9890 &h_rows.slice(src_lo..src_lo + n_copy),
9891 n_copy,
9892 )
9893 .is_err()
9894 {
9895 filled = false;
9896 break 'fill;
9897 }
9898 if self
9899 .mtp_kv_fill_all(
9900 e,
9901 &suffix[start..end],
9902 &phs,
9903 pos + start,
9904 &mut scratch,
9905 embd_dev,
9906 )
9907 .is_err()
9908 {
9909 filled = false;
9910 break 'fill;
9911 }
9912 start = end;
9913 }
9914 if !filled {
9915 // acceptance-only: drafts over missing suffix rows are cheap and wrong,
9916 // so keep only the restored rows resident and let verify arbitrate.
9917 if let Err(err) = scratch.set_len(e, pos) {
9918 return dirty(format!("scratch truncation after failed fill: {err}"));
9919 }
9920 }
9921 // EXTENDED-ENTRY PUBLICATION (lane/sampled-spec-quality, Item 3 — the fix for
9922 // "a restored spec session never publishes an extended entry", SAMPLED-HIT.md
9923 // finding (d)). Pre-lane, publication was armed only for COLD sessions
9924 // (`spec_resumed == 0` in the worker) and both engine capture sites require a
9925 // non-continuation burst — but a converted hit's first burst IS a continuation,
9926 // so a growing conversation learned exactly ONE boundary and turn 3 could never
9927 // hit a longer prefix than turn 2 did.
9928 //
9929 // WHERE, and why it is safe here: `cache.pos == prefix + suffix` at this exact
9930 // line — the trunk is primed over the whole prompt, nothing is generated, and the
9931 // draft plane rows [0..prompt) are filled just above. That is a complete
9932 // whole-entry boundary (`pos == fed_len`), the same shape the cold seed capture
9933 // publishes; the worker's existing publication sweep picks it up because it is
9934 // keyed on non-empty `boundary_captures` and is sampler- and resume-independent.
9935 // NOT the partial-restore hazard: the boundary is this session's own prompt END,
9936 // never mid-entry, so `entry_pos != fed_len` still refuses on the way back in.
9937 // Failure is SILENT by design (the turn_ckpt / boundary-capture convention):
9938 // publication is an optimization, never a correctness dependency.
9939 //
9940 // SUPERSEDED WHEN `republish_at` FIRED (lane/frspec-multiturn-cache): a prompt-end
9941 // entry's tail is the live generation header the next re-render replaces, so on a
9942 // hybrid (whole-entry restores) it can never serve the conversation's next turn —
9943 // the stable-boundary capture above IS this publication, minus the poisoned tail.
9944 if spec_restore_republish_on() && boundary_captures.is_empty() {
9945 debug_assert_eq!(
9946 cache.pos,
9947 pos + t,
9948 "extended-entry capture must sit at the restored session's prompt end",
9949 );
9950 if let Ok(snap) = cache.snapshot(e) {
9951 boundary_captures.push(SpecBoundaryCapture {
9952 snap,
9953 pos: pos + t,
9954 logits: feed_logits.clone(),
9955 last_h: capture_boundary_hidden(e, &h_rows, t, n_embd),
9956 });
9957 }
9958 }
9959 // continuation seed: the feed's boundary logits ARE the plain path's boundary
9960 // logits (same program), so greedy's argmax here is plain's first emitted token,
9961 // and the sampled draw is the cold sampled session's own first token.
9962 next_pred = Some(if sampled {
9963 let sp = sampling.expect("sampled implies a sampler");
9964 // `committed` is still the restored prefix here; the suffix joins it below —
9965 // so this is the last-N window over the WHOLE prompt, exactly the cold
9966 // session's own window at its first token.
9967 let hist = pen_window_seed(&committed, suffix, sp.penalty_last_n);
9968 match sample_boundary_token(
9969 e,
9970 &feed_logits,
9971 &sp,
9972 &hist,
9973 &mut sctr,
9974 "restore-suffix-feed",
9975 ) {
9976 Ok(t) => t,
9977 // the trunk is already fed: hand nothing back, the worker serves the
9978 // request cold-plain. Never fall back to an argmax — that would put a
9979 // greedy token in a sampled stream to save a slow path.
9980 Err(err) => {
9981 return dirty(format!("boundary token draw failed: {err}"));
9982 }
9983 }
9984 } else {
9985 argmax(&feed_logits) as u32
9986 });
9987 let mut lh = match e.uninit(n_embd) {
9988 Ok(b) => b,
9989 Err(err) => return dirty(format!("boundary hidden alloc: {err}")),
9990 };
9991 if let Err(err) = e.copy_view_into(
9992 &mut lh,
9993 0,
9994 &h_rows.slice((t - 1) * n_embd..t * n_embd),
9995 n_embd,
9996 ) {
9997 return dirty(format!("boundary hidden copy: {err}"));
9998 }
9999 last_h_dev = Some(lh);
10000 committed.extend_from_slice(suffix);
10001 } else {
10002 // FULL-COVER shape (empty suffix — the identical-repeat / agent-loop shape): the
10003 // ENTRY's boundary logits are the boundary row, and this is the token the cold
10004 // session emits from that same row. Owned here rather than in the worker so the
10005 // sampled draw cannot be half-applied on one shape (the worker used to argmax it).
10006 if boundary_logits.is_empty() {
10007 return fail(
10008 cache,
10009 "full-cover restore without the entry's boundary logits".into(),
10010 );
10011 }
10012 next_pred = Some(if sampled {
10013 let sp = sampling.expect("sampled implies a sampler");
10014 let hist = pen_window_seed(&committed, &[], sp.penalty_last_n);
10015 match sample_boundary_token(
10016 e,
10017 boundary_logits,
10018 &sp,
10019 &hist,
10020 &mut sctr,
10021 "restore-full-cover",
10022 ) {
10023 Ok(t) => t,
10024 // nothing has been mutated on this shape — hand the carrier back and let
10025 // the hit serve PLAIN (the banked pre-lane path).
10026 Err(err) => {
10027 return fail(cache, format!("boundary token draw failed: {err}"));
10028 }
10029 }
10030 } else {
10031 argmax(boundary_logits) as u32
10032 });
10033 }
10034 Ok(SpecSession {
10035 cache,
10036 scratch,
10037 committed,
10038 last_h: last_h_dev,
10039 next_pred,
10040 sctr,
10041 uctr: 0,
10042 draft_ctx: None,
10043 pending_tok: None,
10044 // Stable-boundary capture from the split feed above (None on the legacy shape):
10045 // a restored session previously parked WITHOUT a checkpoint, so the next turn's
10046 // affinity probe declined ("no turn checkpoint retained") and the conversation
10047 // fell back to the frozen prefix entry forever.
10048 turn_ckpt: restored_turn_ckpt,
10049 telem: SpecTelemetryCounters::default(),
10050 capture_at: None,
10051 boundary_captures,
10052 ckpt_at: None,
10053 })
10054 }
10055
10056 /// Forced-gate exact state comparison. This intentionally reads the real live prefixes from
10057 /// their owning PP devices: matching emitted ids alone would miss a stale `len_d`, recurrent
10058 /// snapshot, or draft-KV row that only corrupts the following round.
10059 pub fn optipipe_compare_session_state(
10060 &self,
10061 e: &Engine,
10062 reference: &SpecSession,
10063 candidate: &SpecSession,
10064 ) -> Result<OptiForkStateIdentity, Box<dyn std::error::Error>> {
10065 fn fail(what: &str) -> Box<dyn std::error::Error> {
10066 format!("optipipe state mismatch: {what}").into()
10067 }
10068 fn same_f32(a: &[f32], b: &[f32]) -> bool {
10069 a.len() == b.len() && a.iter().zip(b).all(|(x, y)| x.to_bits() == y.to_bits())
10070 }
10071 fn compare_layers(
10072 es: &Engine,
10073 range: std::ops::Range<usize>,
10074 reference: &SpecSession,
10075 candidate: &SpecSession,
10076 report: &mut OptiForkStateIdentity,
10077 ) -> Result<(), Box<dyn std::error::Error>> {
10078 for il in range {
10079 match (&reference.cache.kv[il], &candidate.cache.kv[il]) {
10080 (Some(a), Some(b)) => {
10081 if a.len != b.len {
10082 return Err(fail(&format!(
10083 "layer {il} host KV len {} != {}",
10084 a.len, b.len
10085 )));
10086 }
10087 let ad = es.dtoh_i32(&a.len_d)?;
10088 let bd = es.dtoh_i32(&b.len_d)?;
10089 if ad != bd || ad.first().copied() != Some(a.len as i32) {
10090 return Err(fail(&format!(
10091 "layer {il} device KV len {ad:?} != {bd:?} (host={})",
10092 a.len,
10093 )));
10094 }
10095 let kb = a.len * a.k_tok_bytes;
10096 let vb = a.len * a.v_tok_bytes;
10097 if kb > 0 {
10098 let ak = es.dtoh_u8_view(&a.k.slice(0..kb))?;
10099 let bk = es.dtoh_u8_view(&b.k.slice(0..kb))?;
10100 if ak != bk {
10101 let at = ak.iter().zip(&bk).position(|(x, y)| x != y).unwrap();
10102 return Err(fail(&format!(
10103 "layer {il} K bytes at byte {at} row {} offset {}: {} != {}",
10104 at / a.k_tok_bytes,
10105 at % a.k_tok_bytes,
10106 ak[at],
10107 bk[at],
10108 )));
10109 }
10110 }
10111 if vb > 0 {
10112 let av = es.dtoh_u8_view(&a.v.slice(0..vb))?;
10113 let bv = es.dtoh_u8_view(&b.v.slice(0..vb))?;
10114 if av != bv {
10115 let at = av.iter().zip(&bv).position(|(x, y)| x != y).unwrap();
10116 return Err(fail(&format!(
10117 "layer {il} V bytes at byte {at} row {} offset {}: {} != {}",
10118 at / a.v_tok_bytes,
10119 at % a.v_tok_bytes,
10120 av[at],
10121 bv[at],
10122 )));
10123 }
10124 }
10125 report.trunk_kv_bytes += kb + vb;
10126 }
10127 (None, None) => {}
10128 _ => return Err(fail(&format!("layer {il} KV presence"))),
10129 }
10130 match (&reference.cache.recur[il], &candidate.cache.recur[il]) {
10131 (Some(a), Some(b)) => {
10132 let ac = es.dtoh(&a.conv_state)?;
10133 let bc = es.dtoh(&b.conv_state)?;
10134 if !same_f32(&ac, &bc) {
10135 return Err(fail(&format!("layer {il} conv state")));
10136 }
10137 let as_ = es.dtoh(&a.ssm_state)?;
10138 let bs = es.dtoh(&b.ssm_state)?;
10139 if !same_f32(&as_, &bs) {
10140 return Err(fail(&format!("layer {il} SSM state")));
10141 }
10142 report.recurrent_bytes += (ac.len() + as_.len()) * 4;
10143 }
10144 (None, None) => {}
10145 _ => return Err(fail(&format!("layer {il} recurrent presence"))),
10146 }
10147 }
10148 Ok(())
10149 }
10150
10151 if reference.committed != candidate.committed {
10152 return Err(fail("committed token ids"));
10153 }
10154 if reference.cache.pos != candidate.cache.pos
10155 || reference.cache.max_ctx != candidate.cache.max_ctx
10156 {
10157 return Err(fail("cache pos/capacity"));
10158 }
10159 if reference.pending_tok != candidate.pending_tok
10160 || reference.next_pred != candidate.next_pred
10161 || reference.sctr != candidate.sctr
10162 || reference.uctr != candidate.uctr
10163 {
10164 return Err(fail("pending/prediction/counter tail"));
10165 }
10166
10167 let mut report = OptiForkStateIdentity::default();
10168 if let Some(fence) = crate::pp::pp_cuts(self.layers.len()) {
10169 let rt = crate::pp::PpNRt::get(e)?;
10170 for stage in 0..rt.n_stages() {
10171 let _scope = rt.enter(stage);
10172 compare_layers(
10173 rt.engine(stage, e),
10174 fence[stage]..fence[stage + 1],
10175 reference,
10176 candidate,
10177 &mut report,
10178 )?;
10179 }
10180 } else {
10181 compare_layers(e, 0..self.layers.len(), reference, candidate, &mut report)?;
10182 }
10183
10184 if reference.scratch.plane_count() != candidate.scratch.plane_count() {
10185 return Err(fail("draft scratch plane count"));
10186 }
10187 for index in 0..reference.scratch.plane_count() {
10188 let (a, _) = reference.scratch.plane(index);
10189 let (b, _) = candidate.scratch.plane(index);
10190 if a.len != b.len
10191 || a.kv_dim_k != b.kv_dim_k
10192 || a.kv_dim_v != b.kv_dim_v
10193 || a.k_tok_bytes != b.k_tok_bytes
10194 || a.v_tok_bytes != b.v_tok_bytes
10195 || e.dtoh_i32(&a.len_d)? != e.dtoh_i32(&b.len_d)?
10196 {
10197 return Err(fail(&format!("draft scratch plane {index} length/layout")));
10198 }
10199 let kb = a.len * a.k_tok_bytes;
10200 let vb = a.len * a.v_tok_bytes;
10201 if kb > 0 && e.dtoh_u8_view(&a.k.slice(0..kb))? != e.dtoh_u8_view(&b.k.slice(0..kb))? {
10202 return Err(fail(&format!("draft scratch plane {index} K bytes")));
10203 }
10204 if vb > 0 && e.dtoh_u8_view(&a.v.slice(0..vb))? != e.dtoh_u8_view(&b.v.slice(0..vb))? {
10205 return Err(fail(&format!("draft scratch plane {index} V bytes")));
10206 }
10207 report.scratch_kv_bytes += kb + vb;
10208 }
10209
10210 match (&reference.last_h, &candidate.last_h) {
10211 (Some(a), Some(b)) => {
10212 let ah = e.dtoh(a)?;
10213 let bh = e.dtoh(b)?;
10214 if !same_f32(&ah, &bh) {
10215 return Err(fail("last hidden/seed bytes"));
10216 }
10217 report.hidden_bytes = ah.len() * 4;
10218 }
10219 (None, None) => {}
10220 _ => return Err(fail("last hidden/seed presence")),
10221 }
10222 Ok(report)
10223 }
10224
10225 /// SESSION-AFFINITY REWIND (lane/session-affinity, 2026-08-05): roll `sess` back to its
10226 /// retained prompt-end checkpoint, so a request whose prompt matches
10227 /// `committed[..rewind_pos()]` exactly can resume there and prime only its own delta.
10228 ///
10229 /// EXACTNESS. After this returns, the session is byte-for-byte the state it was in AT that
10230 /// boundary: full-attn KV truncated to it (append-only, position-addressed), GDN conv/ssm
10231 /// restored from the device copy taken there, draft scratch length reset, `committed`
10232 /// truncated, `last_h` = the boundary's predecessor anchor. That is precisely the state a
10233 /// fresh prime of `committed[..pos]` would have produced, so the following suffix prime and
10234 /// every burst after it are identical to a cold run of the same token stream — the
10235 /// committed-tokens-authoritative contract.
10236 ///
10237 /// `next_pred` and `pending_tok` are CLEARED: both describe generation past the boundary,
10238 /// which the rewind discards. The caller therefore must supply a non-empty suffix (a
10239 /// rewound session cannot serve an empty-suffix continuation burst — there is nothing to
10240 /// continue). The persistent draft graph survives: it bakes only session-stable pointers
10241 /// (the scratch KV, the resident embedding), none of which the rewind moves.
10242 ///
10243 /// The checkpoint is CONSUMED (`turn_ckpt` taken): its snapshot buffers are freed here, and
10244 /// this turn's own prime installs a fresh one at the new prompt end. Returns the position
10245 /// rewound to, or `None` when the session holds no checkpoint (caller: full re-prime).
10246 pub fn spec_rewind_to_checkpoint(
10247 &self,
10248 e: &Engine,
10249 sess: &mut SpecSession,
10250 ) -> Result<Option<usize>, Box<dyn std::error::Error>> {
10251 if sess.turn_ckpt.as_ref().is_some_and(|ckpt| {
10252 !sess.cache.can_rollback(&ckpt.snap, 0) || !sess.scratch.can_rewind_to(ckpt.pos)
10253 }) {
10254 return Err(
10255 "SWA ring rewind checkpoint has been lapped; full re-prime required".into(),
10256 );
10257 }
10258 let Some(ckpt) = sess.turn_ckpt.take() else {
10259 return Ok(None);
10260 };
10261 assert!(
10262 ckpt.pos <= sess.committed.len(),
10263 "checkpoint past committed ({} > {})",
10264 ckpt.pos,
10265 sess.committed.len()
10266 );
10267 // Restore through each layer's owning engine. A single primary-engine rollback is not
10268 // sufficient when the serving cache is stage-owned under cross-device PP.
10269 crate::pp::restore_cache_checkpoint(e, self, None, &mut sess.cache, &ckpt.snap)?;
10270 debug_assert_eq!(
10271 sess.cache.pos, ckpt.pos,
10272 "rollback landed off the checkpoint"
10273 );
10274 sess.scratch.set_len(e, ckpt.pos)?;
10275 sess.committed.truncate(ckpt.pos);
10276 sess.last_h = Some(ckpt.last_h);
10277 sess.next_pred = None;
10278 sess.pending_tok = None;
10279 Ok(Some(ckpt.pos))
10280 }
10281
10282 /// Grow a parked speculative session to `target_cap` and rewind it to its retained turn
10283 /// checkpoint without re-priming the checkpoint prefix.
10284 ///
10285 /// The trunk cache is restored exactly like a plain grown cache: append-only full-attention
10286 /// KV rows come from the parked cache, while recurrent state comes from the checkpoint's
10287 /// owned snapshot. The MTP scratch is also context-linear and its rows below the checkpoint
10288 /// remain authoritative, so they are copied into a fresh larger scratch before its length is
10289 /// truncated. Pointer-baking draft graphs are dropped and recaptured on the next burst.
10290 ///
10291 /// All fallible work completes before `sess` is mutated. A failed allocation or copy leaves
10292 /// the parked session intact, allowing the caller one reclaim-and-retry attempt.
10293 pub fn spec_grow_and_rewind_to_checkpoint(
10294 &self,
10295 e: &Engine,
10296 sess: &mut SpecSession,
10297 target_cap: usize,
10298 ) -> Result<Option<usize>, Box<dyn std::error::Error>> {
10299 if target_cap <= sess.cache.max_ctx {
10300 return self.spec_rewind_to_checkpoint(e, sess);
10301 }
10302 let Some(ckpt) = sess.turn_ckpt.as_ref() else {
10303 return Ok(None);
10304 };
10305 if ckpt.pos == 0 || ckpt.pos > sess.committed.len() {
10306 return Err(format!(
10307 "checkpoint pos {} outside committed length {}",
10308 ckpt.pos,
10309 sess.committed.len(),
10310 )
10311 .into());
10312 }
10313 if ckpt.pos > target_cap {
10314 return Err(format!(
10315 "checkpoint pos {} exceeds grown capacity {target_cap}",
10316 ckpt.pos,
10317 )
10318 .into());
10319 }
10320
10321 let mut grown_cache = crate::pp::new_cache_planned(e, &self.cfg, &self.plan, target_cap)?;
10322 let mut grown_scratch = self.new_mtp_scratch(e, target_cap)?;
10323 crate::pp::restore_cache_checkpoint(
10324 e,
10325 self,
10326 Some(&sess.cache),
10327 &mut grown_cache,
10328 &ckpt.snap,
10329 )?;
10330
10331 if sess.scratch.plane_count() != grown_scratch.plane_count() {
10332 return Err("checkpoint draft plane count mismatch".into());
10333 }
10334 for index in 0..sess.scratch.plane_count() {
10335 let (src, _) = sess.scratch.plane(index);
10336 let (dst, _) = grown_scratch.plane_mut(index);
10337 if ckpt.pos > src.len
10338 || src.kv_dim_k != dst.kv_dim_k
10339 || src.kv_dim_v != dst.kv_dim_v
10340 || src.k_tok_bytes != dst.k_tok_bytes
10341 || src.v_tok_bytes != dst.v_tok_bytes
10342 {
10343 return Err(format!(
10344 "checkpoint draft plane {index} layout mismatch (pos {}, source len {})",
10345 ckpt.pos, src.len,
10346 )
10347 .into());
10348 }
10349 match (&src.ring, dst.ring.as_ref()) {
10350 (Some(sring), Some(_)) => {
10351 // Ring-backed draft plane (step35): `ckpt.pos` is absolute and exceeds the
10352 // physical rows once lapped — same class as the trunk-KV restore panic
10353 // (2026-08-29 warm-turn-at-40k). Copy the aligned live window, rebase.
10354 let (new_base, phys) = sring.restore_plan(ckpt.pos).map_err(|err| {
10355 format!("checkpoint draft plane {index} SWA restore refused: {err}")
10356 })?;
10357 let rows = phys.len();
10358 let kb = rows * src.k_tok_bytes;
10359 let vb = rows * src.v_tok_bytes;
10360 if kb > 0 {
10361 e.copy_u8_range_into(
10362 &mut dst.k,
10363 0,
10364 &src.k,
10365 phys.start * src.k_tok_bytes,
10366 kb,
10367 )?;
10368 }
10369 if vb > 0 {
10370 e.copy_u8_range_into(
10371 &mut dst.v,
10372 0,
10373 &src.v,
10374 phys.start * src.v_tok_bytes,
10375 vb,
10376 )?;
10377 }
10378 dst.ring
10379 .as_mut()
10380 .expect("ring presence checked above")
10381 .apply_rebase(new_base);
10382 if let Some(base_d) = dst.base_d.as_mut() {
10383 e.set_i32_one(base_d, new_base as i32)?;
10384 }
10385 }
10386 (None, None) => {
10387 let kb = ckpt.pos * src.k_tok_bytes;
10388 let vb = ckpt.pos * src.v_tok_bytes;
10389 if kb > 0 {
10390 e.copy_u8_into(&mut dst.k, 0, &src.k, kb)?;
10391 }
10392 if vb > 0 {
10393 e.copy_u8_into(&mut dst.v, 0, &src.v, vb)?;
10394 }
10395 }
10396 _ => {
10397 return Err(format!("checkpoint draft plane {index} ring/flat mismatch").into());
10398 }
10399 }
10400 }
10401 grown_scratch.set_len(e, ckpt.pos)?;
10402 // The old scratch is dropped immediately after publication below. Bound its D2D reads
10403 // first; growth happens once per rewritten turn, outside the decode hot loop.
10404 e.stream().synchronize()?;
10405
10406 let ckpt = sess
10407 .turn_ckpt
10408 .take()
10409 .expect("checkpoint remained present through transactional grow");
10410 let pos = ckpt.pos;
10411 sess.cache = grown_cache;
10412 sess.scratch = grown_scratch;
10413 sess.committed.truncate(pos);
10414 sess.last_h = Some(ckpt.last_h);
10415 sess.next_pred = None;
10416 sess.pending_tok = None;
10417 sess.draft_ctx = None;
10418 debug_assert_eq!(sess.cache.pos, pos, "grown rewind landed off checkpoint");
10419 debug_assert!(
10420 (0..sess.scratch.plane_count()).all(|index| sess.scratch.plane(index).0.len == pos),
10421 "grown draft rewind landed off checkpoint"
10422 );
10423 Ok(Some(pos))
10424 }
10425
10426 /// Commit a carried pending bonus (see SpecSession::pending_tok): one T=1 trunk pass
10427 /// (its logits' argmax becomes next_pred) + the draft-KV fill at the carried anchor —
10428 /// byte-identical to the pre-carry session tail. Required before a non-empty-suffix
10429 /// prime, a sampled turn, or parking a session for pool reuse. No-op without a pending.
10430 /// `sampling` is the sampler of the request that will CONSUME the resulting `next_pred`
10431 /// (lane/sampled-spec-quality): this is a boundary site like any other, so a sampled
10432 /// consumer must get a DRAWN token, not an argmax. Pass `None` from the park/demote
10433 /// callers — a pending only ever exists on the GREEDY tail, and the consumer of a
10434 /// park-time flush is a future request whose sampler is not knowable here (residual
10435 /// named at the pool-resume probe in worker.rs and in SAMPLED-QUALITY.md).
10436 pub fn spec_flush_pending(
10437 &self,
10438 e: &Engine,
10439 sess: &mut SpecSession,
10440 sampling: Option<SpecSampling>,
10441 ) -> Result<(), Box<dyn std::error::Error>> {
10442 let Some(b) = sess.pending_tok.take() else {
10443 return Ok(());
10444 };
10445 if self.mtp.is_none() {
10446 return Err("pending carry requires an MTP head".into());
10447 }
10448 let n_embd = self.cfg.n_embd as usize;
10449 let (embd_qt, embd_rb) = self.embd.qt_and_row_bytes(n_embd);
10450 let embd_gpu = if spec_host_embd() {
10451 None
10452 } else {
10453 Some(
10454 self.embd_gpu
10455 .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload")),
10456 )
10457 };
10458 let embd_dev = embd_gpu.map(|g| (g, embd_qt, embd_rb));
10459 let pos_b = sess.cache.pos;
10460 sess.scratch.set_len(e, pos_b)?;
10461 let (lg_b, hb) = self.spec_target_step_h(e, b, &mut sess.cache)?;
10462 sess.next_pred = Some(match sampling {
10463 Some(sp) if sp.temp > 0.0 && spec_sampled_boundary_on() => {
10464 // window includes `b` itself: it is committed by this pass, and the pre-lane
10465 // code never counted a boundary token in the penalty history at all.
10466 let hist = pen_window_seed(&sess.committed, &[b], sp.penalty_last_n);
10467 sample_boundary_token(e, &lg_b, &sp, &hist, &mut sess.sctr, "flush-pending")?
10468 }
10469 _ => argmax(&lg_b) as u32,
10470 });
10471 let anchor = sess
10472 .last_h
10473 .as_ref()
10474 .expect("pending carry requires last_h (the predecessor-row anchor)");
10475 self.mtp_kv_fill_all(e, &[b], anchor, pos_b, &mut sess.scratch, embd_dev)?;
10476 sess.last_h = Some(hb);
10477 sess.committed.push(b);
10478 Ok(())
10479 }
10480
10481 /// Solo target feed used only at speculative round boundaries. Step35 serving made its
10482 /// staged batched B=1 graph authoritative, so a speculative session must enter and leave
10483 /// rounds through that same graph. Other model families keep their eager T=1 contract.
10484 fn spec_target_step_h(
10485 &self,
10486 e: &Engine,
10487 token: u32,
10488 cache: &mut Cache,
10489 ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
10490 if !self.sliding_gated_moe_batch_program() && !self.batched_serving_numeric_class() {
10491 return self.decode_step_h(e, token, cache);
10492 }
10493 let pos0 = cache.pos;
10494 let (logits, hidden) = self.decode_step_t_core(e, &[token], pos0, cache, None, None)?;
10495 Ok((e.dtoh(&logits)?, hidden))
10496 }
10497
10498 /// The archs whose LIVE B=1 serving runs the generic BATCHED numeric class (decode_step_batch
10499 /// walk + batched head), so their spec verify must run the SAME class. MoE learned this
10500 /// 2026-08-14 AM (4b777ccc5); the dense hybrid reproduced the identical near-tie flip class
10501 /// the same day on Qwen3.8-27B — eager-class verify logits drift from batched-class serving
10502 /// logits ("1 ULP at layer 2 → 2.3e-1 logit maxdiff at the head"), and the GDN recurrence
10503 /// carries the drift until a near-tie flips deep in generation. One predicate so the five
10504 /// dispatch sites cannot drift apart again.
10505 /// Draft-graph head admissibility (lane/draftcost-moe, 2026-08-20): the capture body
10506 /// (`mtp_head_forward_cap`) supports Dense heads AND resident-MoE heads
10507 /// (`Ffn::Moe(m) if m.dev_exps.is_some()`); non-resident MoE still refuses inside the
10508 /// capture and the caller falls back to the eager chain by design. Trunk FFN class is
10509 /// irrelevant — the graph body is the HEAD forward only. One predicate for all three
10510 /// eligibility sites so they cannot drift (the serving numeric-class lesson).
10511 fn mtp_graph_capturable(&self) -> bool {
10512 self.mtp
10513 .as_ref()
10514 .map(|m| match &m.ffn {
10515 crate::hybrid::Ffn::Dense { .. } => true,
10516 crate::hybrid::Ffn::Moe(mo) => mo.dev_exps.is_some(),
10517 })
10518 .unwrap_or(false)
10519 }
10520
10521 fn batched_serving_numeric_class(&self) -> bool {
10522 self.plan
10523 .trunk_operations()
10524 .contains(&memra_gguf::model_plan::OperationKind::GatedDeltaNet)
10525 }
10526
10527 /// The family the MTP verify-graph default was measured on: GatedDeltaNet state layers
10528 /// (a `recur` mixer) together with a routed-MoE FFN — Ornith-1.5-35B-A3B and its kin. The
10529 /// server-side twin of this test is `model_forces_spec_replay` (GatedDeltaNet + MoeMlp);
10530 /// keeping the engine's own version structural rather than name-based means a new
10531 /// checkpoint of the same shape inherits the default, and a different shape does not.
10532 fn vgraph_family_default(&self) -> bool {
10533 let has_linear = self
10534 .layers
10535 .iter()
10536 .any(|l| matches!(l.mixer, Mixer::Linear(_)));
10537 let has_moe = self
10538 .layers
10539 .iter()
10540 .any(|l| matches!(l.ffn, crate::hybrid::Ffn::Moe(_)));
10541 has_linear && has_moe
10542 }
10543
10544 fn sliding_gated_moe_batch_program(&self) -> bool {
10545 self.uses_sliding_gated_moe_program()
10546 }
10547
10548 fn gemma_batch_program(&self) -> bool {
10549 self.uses_gemma_program()
10550 }
10551
10552 /// Reduced-matrix admission for increment 1. This deliberately does not change the PP-2
10553 /// serving policy: the worker calls it only after `MEMRA_SPEC_PIPE=1` and an explicit spec
10554 /// session already exist.
10555 pub fn spec_pipe_available(&self, e: &Engine) -> bool {
10556 if std::env::var("MEMRA_SPEC_PIPE").as_deref() != Ok("1")
10557 || !spec_devacc()
10558 || spec_replay_env_enabled()
10559 || spec_stream()
10560 || std::env::var("MEMRA_SPEC_ADAPT").as_deref() == Ok("1")
10561 || std::env::var("MEMRA_SPEC_PMIN0").as_deref() == Ok("1")
10562 || std::env::var("MEMRA_SPEC_PP_ANATOMY").as_deref() == Ok("1")
10563 || std::env::var("MEMRA_SPEC_PMIN")
10564 .ok()
10565 .and_then(|v| v.parse::<f32>().ok())
10566 .unwrap_or(0.0)
10567 > 0.0
10568 || self.is_gemma4_e4b()
10569 || self.gemma_batch_program()
10570 || self.mtp.is_none()
10571 || !self.mtp_extra.is_empty()
10572 {
10573 return false;
10574 }
10575 let Some(cuts) = crate::pp::pp_cuts(self.layers.len()) else {
10576 return false;
10577 };
10578 if cuts.len() != 3 || crate::pp::pp2_streams_off() || !crate::pp::spec_pp_on() {
10579 return false;
10580 }
10581 crate::pp::PpNRt::get(e)
10582 .map(|rt| rt.n_stages() == 2 && rt.cross_device())
10583 .unwrap_or(false)
10584 }
10585
10586 /// Two warm greedy continuation bursts over one PP-2 interval coordinator. The two existing
10587 /// `generate_spec_inner2` call stacks own all per-session round locals; only phase issue order
10588 /// changes. No callback is accepted in increment 1 — the worker publishes each completed burst.
10589 #[allow(clippy::too_many_arguments)]
10590 pub fn generate_spec_session_pair(
10591 &self,
10592 e: &Engine,
10593 sess_a: &mut SpecSession,
10594 max_new_a: usize,
10595 k_a: usize,
10596 sess_b: &mut SpecSession,
10597 max_new_b: usize,
10598 k_b: usize,
10599 ) -> Result<((Vec<u32>, usize, usize), (Vec<u32>, usize, usize)), Box<dyn std::error::Error>>
10600 {
10601 if !self.spec_pipe_available(e) {
10602 return Err("two-session speculative pipeline is outside its reduced matrix".into());
10603 }
10604 if max_new_a == 0 || max_new_b == 0 || k_a == 0 || k_b == 0 {
10605 return Err(
10606 "two-session speculative pipeline requires non-empty positive-K bursts".into(),
10607 );
10608 }
10609 for sess in [&*sess_a, &*sess_b] {
10610 if sess.committed.is_empty()
10611 || sess.last_h.is_none()
10612 || (sess.next_pred.is_none() && sess.pending_tok.is_none())
10613 {
10614 return Err("two-session speculative pipeline requires warm continuations".into());
10615 }
10616 }
10617
10618 let graph_ok = std::env::var("MEMRA_SPEC_NOGRAPH").is_err()
10619 && !spec_host_embd()
10620 && self.mtp_graph_capturable()
10621 && self.mtp_extra.is_empty()
10622 && !crate::model::full_prec_enabled();
10623 let graph_a = graph_ok && k_a + 2 < 96;
10624 let graph_b = graph_ok && k_b + 2 < 96;
10625 let was_tracking = e.ctx().is_event_tracking();
10626 if (graph_a || graph_b) && was_tracking {
10627 unsafe {
10628 e.ctx().disable_event_tracking();
10629 }
10630 }
10631
10632 static LOGGED: std::sync::Once = std::sync::Once::new();
10633 LOGGED.call_once(|| {
10634 eprintln!("[spec-pipe] two-session PP-2 continuation pipeline engaged");
10635 });
10636 let sync = std::sync::Arc::new(SpecPipeSync::new());
10637 let lane_a = SpecPipeLane {
10638 sync: sync.clone(),
10639 lane: 0,
10640 };
10641 let lane_b = SpecPipeLane { sync, lane: 1 };
10642 let mut sess_b_ptr = SpecPipeSessionPtr(sess_b as *mut SpecSession);
10643 let (result_a, result_b) = std::thread::scope(|scope| {
10644 let b = scope.spawn(move || {
10645 let mut finish = SpecPipeFinish::new(&lane_b);
10646 let sess_b = unsafe { sess_b_ptr.get_mut() };
10647 let result = e
10648 .ctx()
10649 .bind_to_thread()
10650 .map_err(|err| err.to_string())
10651 .and_then(|_| {
10652 self.generate_spec_inner2(
10653 e,
10654 &[],
10655 max_new_b,
10656 k_b,
10657 graph_b,
10658 Some(sess_b),
10659 None,
10660 None,
10661 None,
10662 None,
10663 Some(&lane_b),
10664 )
10665 .map_err(|err| err.to_string())
10666 });
10667 finish.close(result.is_err());
10668 result
10669 });
10670 let mut finish = SpecPipeFinish::new(&lane_a);
10671 let result_a = self.generate_spec_inner2(
10672 e,
10673 &[],
10674 max_new_a,
10675 k_a,
10676 graph_a,
10677 Some(sess_a),
10678 None,
10679 None,
10680 None,
10681 None,
10682 Some(&lane_a),
10683 );
10684 finish.close(result_a.is_err());
10685 let result_b = b
10686 .join()
10687 .map_err(|_| "paired speculative session B panicked".to_string())
10688 .and_then(|r| r);
10689 (result_a, result_b)
10690 });
10691
10692 if (graph_a || graph_b) && was_tracking {
10693 unsafe {
10694 e.ctx().enable_event_tracking();
10695 }
10696 }
10697 let result_a = result_a?;
10698 let result_b = result_b.map_err(|err| -> Box<dyn std::error::Error> { err.into() })?;
10699 Ok((result_a, result_b))
10700 }
10701
10702 /// One spec-decode turn on a live session. `suffix` = the NEW tokens only (turn N+1's user
10703 /// message rendered through the chat template continuation). Returns (new tokens emitted,
10704 /// drafted, accepted); session.committed grows by suffix + emitted.
10705 pub fn generate_spec_session(
10706 &self,
10707 e: &Engine,
10708 sess: &mut SpecSession,
10709 suffix: &[u32],
10710 max_new: usize,
10711 k: usize,
10712 ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
10713 self.generate_spec_session_sampled(e, sess, suffix, max_new, k, None, None)
10714 }
10715
10716 /// Serve-path sampled spec: routes the burst through the rejection-sampling verify with
10717 /// per-SESSION Philox continuity (sess.sctr/uctr). None = env-driven (CLI) or greedy.
10718 /// Filters (top-k/p/min-p) apply SYMMETRICALLY to draft q and verify p — distribution-exact
10719 /// for the filtered target (feat/filtered-spec).
10720 ///
10721 /// `on_commit` (sse-cadence, 2026-08-05): called with each newly-emitted slice of the
10722 /// output — once right after the prime's first token, then once per round commit — so a
10723 /// streaming caller can flush text at round cadence instead of once per burst. The slices
10724 /// are disjoint, in order, and concatenate to exactly the returned token vec. Emission-
10725 /// timing only: token bytes, session state, and exactness are untouched.
10726 ///
10727 /// The returned bool is a CONTINUE-VERDICT (admission yield, 2026-08-06): `false` ends
10728 /// the burst at the current round boundary, exactly as if `max_new` had been reached —
10729 /// the caller's scheduler regains control without waiting the burst out. Burst size is
10730 /// content-neutral (spec-levers battery), so an early exit moves WHEN the burst returns,
10731 /// never what tokens say. The slice may be EMPTY (a poll-only boundary — round-stream
10732 /// drains and the defensive tail flush can land with nothing new committed).
10733 #[allow(clippy::too_many_arguments)]
10734 pub fn generate_spec_session_sampled(
10735 &self,
10736 e: &Engine,
10737 sess: &mut SpecSession,
10738 suffix: &[u32],
10739 max_new: usize,
10740 k: usize,
10741 sampling: Option<SpecSampling>,
10742 on_commit: Option<&mut dyn FnMut(&[u32]) -> bool>,
10743 ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
10744 self.generate_spec_session_sampled_prime_split(
10745 e, sess, suffix, max_new, k, sampling, None, on_commit,
10746 )
10747 }
10748
10749 /// Serve-only cold-prime segmentation twin. `prime_split` is the same stable boundary the
10750 /// plain worker would honor before entering its sub-floor tokenwise tail; warm continuations
10751 /// pass `None` and stay on the existing zero-prime path.
10752 #[allow(clippy::too_many_arguments)]
10753 pub fn generate_spec_session_sampled_prime_split(
10754 &self,
10755 e: &Engine,
10756 sess: &mut SpecSession,
10757 suffix: &[u32],
10758 max_new: usize,
10759 k: usize,
10760 sampling: Option<SpecSampling>,
10761 prime_split: Option<usize>,
10762 on_commit: Option<&mut dyn FnMut(&[u32]) -> bool>,
10763 ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
10764 self.generate_spec_session_constrained_prime_split(
10765 e,
10766 sess,
10767 suffix,
10768 max_new,
10769 k,
10770 sampling,
10771 None,
10772 prime_split,
10773 on_commit,
10774 )
10775 }
10776
10777 /// `generate_spec_session_sampled` + GRAMMAR (constrained decoding, 2026-08-03): the
10778 /// hook truncates acceptance at the first grammar-illegal token AFTER the exactness
10779 /// verify (grammar is an extra rejection rule, ordering like the batched-verify twins)
10780 /// and replaces an illegal bonus with the MASKED argmax of the target's own verify
10781 /// column — token-identical to constrained plain greedy decode. GREEDY only (the
10782 /// worker routes sampled constrained to plain decode). Acceptance under tight grammars
10783 /// may drop (drafter is unconstrained); that is measured, not hidden.
10784 #[allow(clippy::too_many_arguments)]
10785 pub fn generate_spec_session_constrained(
10786 &self,
10787 e: &Engine,
10788 sess: &mut SpecSession,
10789 suffix: &[u32],
10790 max_new: usize,
10791 k: usize,
10792 sampling: Option<SpecSampling>,
10793 constraint: Option<&mut dyn SpecConstraint>,
10794 on_commit: Option<&mut dyn FnMut(&[u32]) -> bool>,
10795 ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
10796 self.generate_spec_session_constrained_prime_split(
10797 e, sess, suffix, max_new, k, sampling, constraint, None, on_commit,
10798 )
10799 }
10800
10801 #[allow(clippy::too_many_arguments)]
10802 pub fn generate_spec_session_constrained_prime_split(
10803 &self,
10804 e: &Engine,
10805 sess: &mut SpecSession,
10806 suffix: &[u32],
10807 max_new: usize,
10808 k: usize,
10809 sampling: Option<SpecSampling>,
10810 constraint: Option<&mut dyn SpecConstraint>,
10811 prime_split: Option<usize>,
10812 on_commit: Option<&mut dyn FnMut(&[u32]) -> bool>,
10813 ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
10814 if constraint.is_some() && sampling.is_some_and(|s| s.temp > 0.0) {
10815 return Err(
10816 "constrained spec decode is greedy-only (worker routes sampled \
10817 constrained to plain decode)"
10818 .into(),
10819 );
10820 }
10821 // PENDING-CARRY entry flush: a carried bonus precedes any new suffix in the sequence,
10822 // so it must commit BEFORE the suffix primes; the sampled path doesn't carry (its
10823 // round-0 accept needs the commit pass's logits). Empty-suffix greedy bursts — the
10824 // serve continuation case — consume the carry in-loop with zero solo passes.
10825 if sess.pending_tok.is_some()
10826 && (!suffix.is_empty() || sampling.map_or(false, |s| s.temp > 0.0))
10827 {
10828 self.spec_flush_pending(e, sess, sampling)?;
10829 }
10830
10831 // FULL_PREC forces the EAGER draft: the graph capture would enclose cuBLASLt f32 GEMV
10832 // (the FloatBf16 else-branches) and a bf16_to_f32 dequant alloc — neither is stream-capture
10833 // safe. Eager rides matmul/matmul_decode_exact, which dequant FloatBf16 on use. (§item 2.)
10834 let graph_draft = std::env::var("MEMRA_SPEC_NOGRAPH").is_err()
10835 && !spec_host_embd()
10836 && self.mtp_graph_capturable()
10837 && self.mtp_extra.is_empty()
10838 && k + 2 < 96
10839 && !crate::model::full_prec_enabled();
10840 let was_tracking = e.ctx().is_event_tracking();
10841 if graph_draft && was_tracking {
10842 unsafe {
10843 e.ctx().disable_event_tracking();
10844 }
10845 }
10846 let r = self.generate_spec_inner2(
10847 e,
10848 suffix,
10849 max_new,
10850 k,
10851 graph_draft,
10852 Some(sess),
10853 sampling,
10854 constraint,
10855 on_commit,
10856 prime_split,
10857 None,
10858 );
10859 if graph_draft && was_tracking {
10860 unsafe {
10861 e.ctx().enable_event_tracking();
10862 }
10863 }
10864 let (out, d, a) = r?;
10865 Ok((out, d, a))
10866 }
10867
10868 pub fn generate_spec(
10869 &self,
10870 e: &Engine,
10871 prompt: &[u32],
10872 max_new: usize,
10873 k: usize,
10874 ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
10875 if crate::pp::pp_cuts(self.layers.len()).is_some()
10876 && !self.rewrite_allowed(memra_gguf::execution_manifest::RewriteSurface::Pipeline)
10877 {
10878 return Err("pipeline rewrite is not qualified for speculative decode".into());
10879 }
10880 if !self.rewrite_allowed(memra_gguf::execution_manifest::RewriteSurface::MtpSpec) {
10881 return Err("speculative rewrite is not qualified for this ModelPlan".into());
10882 }
10883 // FULL_PREC forces eager (see generate_spec_session note): CUDA graph capture cannot
10884 // enclose cuBLASLt f32 GEMV or the bf16_to_f32 dequant alloc the FloatBf16 path needs.
10885 let graph_draft = std::env::var("MEMRA_SPEC_NOGRAPH").is_err()
10886 && !spec_host_embd()
10887 && self.mtp_graph_capturable()
10888 && self.mtp_extra.is_empty()
10889 && k + 2 < 96
10890 && !crate::model::full_prec_enabled();
10891 if !graph_draft {
10892 return self.generate_spec_inner2(
10893 e, prompt, max_new, k, false, None, None, None, None, None, None,
10894 );
10895 }
10896 let was_tracking = e.ctx().is_event_tracking();
10897 if was_tracking {
10898 unsafe {
10899 e.ctx().disable_event_tracking();
10900 }
10901 }
10902 let r = self.generate_spec_inner2(
10903 e, prompt, max_new, k, true, None, None, None, None, None, None,
10904 );
10905 if was_tracking {
10906 unsafe {
10907 e.ctx().enable_event_tracking();
10908 }
10909 }
10910 r
10911 }
10912
10913 fn generate_spec_inner2(
10914 &self,
10915 e: &Engine,
10916 prompt: &[u32],
10917 max_new: usize,
10918 k: usize,
10919 graph_draft: bool,
10920 mut sess: Option<&mut SpecSession>,
10921 sampling: Option<SpecSampling>,
10922 mut constraint: Option<&mut dyn SpecConstraint>,
10923 mut on_commit: Option<&mut dyn FnMut(&[u32]) -> bool>,
10924 prime_split: Option<usize>,
10925 pipe: Option<&SpecPipeLane>,
10926 ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
10927 assert!(k >= 1, "k must be >= 1");
10928 if let Some(p) = pipe {
10929 p.setup_begin()?;
10930 }
10931 // sse-cadence flush cursor: everything in out[..flushed] has been handed to on_commit.
10932 let mut flushed = 0usize;
10933 // admission yield (2026-08-06): on_commit's continue-verdict; false = end the burst
10934 // at the next round boundary (same exit as max_new reached — the session tail runs).
10935 // Initialized by the unconditional post-prime flush below.
10936 let mut keep_going;
10937 let mtp = self
10938 .mtp
10939 .as_ref()
10940 .expect("generate_spec requires an MTP head (nextn_predict_layers>0)");
10941 let n_vocab = self.output.out_features();
10942 // FR-Spec: the draft head may be TRIMMED (fewer rows than n_vocab); the draft argmax runs
10943 // over the draft vocab and the winning index maps through d2t to a TARGET token id.
10944 // Everything downstream (verify/accept/commit) sees target ids only — exactness unchanged.
10945 let d_vocab = mtp
10946 .shared_head_head
10947 .as_ref()
10948 .unwrap_or(&self.output)
10949 .out_features();
10950 if !self.mtp_extra.is_empty() {
10951 if self.plan.draft_source != memra_gguf::model_plan::DraftSourcePlan::Embedded
10952 || self.plan.mtp_blocks.len() != self.mtp_head_count()
10953 {
10954 return Err(
10955 "multi-head MTP requires one embedded canonical block per loaded head".into(),
10956 );
10957 }
10958 // TRIMMED chains (2026-08-27): every head must carry the SAME d2t — the ranking is
10959 // token-frequency and head-independent, and every downstream remap (per-step argmax,
10960 // stream pack, sampled d2t_dev) reads head 0's map, so equality is what makes that
10961 // single map correct for the whole chain. Mixed trimmed/untrimmed is refused.
10962 for (offset, head) in self.mtp_extra.iter().enumerate() {
10963 if head.d2t != mtp.d2t
10964 || head
10965 .shared_head_head
10966 .as_ref()
10967 .unwrap_or(&self.output)
10968 .out_features()
10969 != d_vocab
10970 {
10971 return Err(format!(
10972 "embedded MTP head {} has incompatible draft vocabulary",
10973 offset + 1
10974 )
10975 .into());
10976 }
10977 }
10978 eprintln!(
10979 "[mtp-chain] heads={} policy=step-modulo prefix-replay kv=per-head",
10980 self.mtp_head_count()
10981 );
10982 }
10983 let n_embd = self.cfg.n_embd as usize;
10984 // SESSION MODE: reuse the live cache/scratch, prime only the suffix. `base` = tokens
10985 // already committed (their state is in the caches); 0 = fresh single-shot call.
10986 let session_mode = sess.is_some();
10987 let max_ctx = match sess.as_ref() {
10988 Some(s) => s.cache.max_ctx,
10989 None => prompt.len() + max_new + k + 8,
10990 };
10991 let mut own_cache;
10992 let mut own_scratch;
10993 // PREFIX-CACHE capture request threaded out of the session (lane/spec-prefix-cache):
10994 // (requested split, destination list). Single-shot per burst; fresh calls have none.
10995 let mut sess_capture: Option<(Option<usize>, &mut Vec<SpecBoundaryCapture>)> = None;
10996 // STABLE-BOUNDARY turn-checkpoint request (lane/frspec-multiturn-cache): ABSOLUTE
10997 // committed-length position; consumed one-shot like `capture_at`. None = legacy
10998 // prompt-end capture below.
10999 let mut ckpt_req: Option<usize> = None;
11000 let (
11001 cache,
11002 scratch,
11003 mut sess_tail,
11004 mut sess_draft_slot,
11005 mut sess_pending_slot,
11006 sess_ckpt_slot,
11007 sess_telem,
11008 ): (
11009 &mut Cache,
11010 &mut MtpScratch,
11011 Option<(
11012 &mut Vec<u32>,
11013 &mut Option<CudaSlice<f32>>,
11014 &mut Option<u32>,
11015 &mut u32,
11016 &mut u32,
11017 )>,
11018 Option<&mut Option<DraftGraphCtx>>,
11019 Option<&mut Option<u32>>,
11020 Option<&mut Option<SpecCheckpoint>>,
11021 Option<&SpecTelemetryCounters>,
11022 ) = match sess.take() {
11023 Some(sr) => {
11024 let SpecSession {
11025 cache,
11026 scratch,
11027 committed,
11028 last_h,
11029 next_pred,
11030 sctr: s_sctr,
11031 uctr: s_uctr,
11032 draft_ctx,
11033 pending_tok,
11034 turn_ckpt,
11035 telem,
11036 capture_at,
11037 boundary_captures,
11038 ckpt_at,
11039 } = sr;
11040 sess_capture = Some((capture_at.take(), boundary_captures));
11041 ckpt_req = ckpt_at.take();
11042 (
11043 cache,
11044 scratch,
11045 Some((committed, last_h, next_pred, s_sctr, s_uctr)),
11046 Some(draft_ctx),
11047 Some(pending_tok),
11048 Some(turn_ckpt),
11049 Some(telem),
11050 )
11051 }
11052 None => {
11053 // STAGE-OWNED KV (lane/pp2-spec 2026-08-06) — see `new_session`. Door shut =
11054 // `Cache::new` verbatim.
11055 own_cache = crate::pp::new_cache_planned(e, &self.cfg, &self.plan, max_ctx)?;
11056 // Persistent scratch = max_ctx rows (~2KB/token quantized).
11057 own_scratch = self.new_mtp_scratch(e, max_ctx)?;
11058 (
11059 &mut own_cache,
11060 &mut own_scratch,
11061 None,
11062 None,
11063 None,
11064 None,
11065 None,
11066 )
11067 }
11068 };
11069 if scratch.plane_count() != self.mtp_head_count() {
11070 return Err(format!(
11071 "MTP scratch/head count mismatch ({}/{})",
11072 scratch.plane_count(),
11073 self.mtp_head_count()
11074 )
11075 .into());
11076 }
11077 let base = cache.pos;
11078 // PENDING-CARRY consume (2026-08-01): a carried bonus reaches here only on the
11079 // empty-suffix GREEDY continuation path (generate_spec_session_sampled flushed every
11080 // other case). It enters the round loop as round-0's pending — verify col 0 — exactly
11081 // like a mid-burst full-accept boundary: no init feed, no tail commit pass.
11082 let carried_pending: Option<u32> = sess_pending_slot.as_mut().and_then(|s| s.take());
11083 // PERSISTENT DRAFT KV (the only mode since 2026-07-08 — the legacy round-local scratch,
11084 // MEMRA_SPEC_KVLOCAL, measured -35 acceptance pts on the 27B p3 sweep and was removed;
11085 // acceptance-only — exactness is verify's job either way).
11086 // HIDDEN-PAIRING CONVENTION (DEFAULT = predecessor-row, 2026-07-04 — the 27B acceptance
11087 // unlock, +16pts): the MTP head is TRAINED on rows pairing token x_p with the trunk
11088 // hidden of its PREDECESSOR h_{p-1} (the reference engine's mtp_update shifts the target
11089 // hiddens right by one; its draft step 0 feeds (id_last, TRUE hidden of the row id_last
11090 // was sampled from)). memra's historical convention paired SAME-ROW (x_p, h_p) in the fill
11091 // and seeded chain step 0 through an extra MTP pass on a duplicated token (the
11092 // pseudo-seed) — measured 27B p2 K=3 acceptance 0.569 vs 0.731, p3 0.445 vs 0.63+, and
11093 // the chain steps j>=1 were already predecessor-shaped, so ONLY the fill + step-0 seed
11094 // move. The fill shifts by one and the chain seeds from the predecessor's true hidden
11095 // DIRECTLY (vh_seed / vx[j-1]) — the pseudo pass disappears (one MTP-block pass saved
11096 // per round on top of the acceptance win). Draft-quality-only: exactness stays the
11097 // verify's job either way. (The legacy same-row pairing seam, MEMRA_SPEC_HSAME, and its
11098 // pseudo-seed passes were removed 2026-07-08 — predecessor pairing won by +16 acc pts;
11099 // the legacy round-local scratch, MEMRA_SPEC_KVLOCAL, went with it.)
11100 // REPLAY-FREE PARTIAL ACCEPT (default, 2026-07-03): partial rounds keep the verify's own
11101 // bit-identical committed-prefix state (KV truncate + recur rebuild from the VerifyCkpt)
11102 // and leave the bonus PENDING — no duplicate trunk pass (profiled ~0.54 extra full weight
11103 // reads/round at long ctx). MEMRA_SPEC_REPLAY=1 restores the legacy rollback+replay (A/B
11104 // + fallback seam).
11105 // Qwen35-MoE replay pin LIFTED (lane/draftcost-moe, 2026-08-20). The pin's stated
11106 // bar — the retained verify-state commit proven equivalent to sequential serving —
11107 // was waiting on this arch running the serving batched verify class, which the
11108 // t-parallel admission (this lane, increment 1) provided: the VerifyCkpt the
11109 // replay-free commit consumes is now produced by the SAME serving-class verify that
11110 // qualified dense qwen35 on 2026-08-15 (where the per-round duplicate replay
11111 // measured 69 -> 30 tok/s). Qualification receipts (run-spec K=1..8 both arms,
11112 // 8-prompt replay-vs-replay-free canary, long-prompt cell):
11113 // research/draftcost-moe-20260820/RECEIPTS.md. MEMRA_SPEC_REPLAY=1 stays the
11114 // rollback + A/B seam.
11115 let spec_replay = spec_replay_env_enabled();
11116 if constraint.is_some() && spec_replay {
11117 return Err(
11118 "constrained spec decode does not support MEMRA_SPEC_REPLAY=1 \
11119 (legacy replay commits an unmasked bonus)"
11120 .into(),
11121 );
11122 }
11123 // TRUE-HIDDEN REFRESH (default in persistent-draft-KV mode): every round overwrites the
11124 // committed positions' scratch entries from the verify's exact hiddens (mtp_kv_fill batch)
11125 // instead of keeping chain-approximate entries. MEMRA_SPEC_NOREFRESH=1 = legacy (A/B seam).
11126 let refresh = std::env::var("MEMRA_SPEC_NOREFRESH").is_err();
11127 if !refresh && !self.mtp_extra.is_empty() {
11128 return Err("multi-head MTP requires exact accepted-prefix refresh".into());
11129 }
11130
11131 // prime: BATCHED cache prime (prime_cache — the measured #1 e2e gap: tokenwise primed at
11132 // ~102/38 tok/s vs the engine's ~2000-5900 tok/s batched prefill). prime_cache returns the
11133 // full pre-output_norm hidden stack [T, n_embd], which IS prompt_h (the persistent-draft-KV
11134 // mtp_kv_fill input) — no per-token collection needed. Prompts below PRIME_MIN_T, and
11135 // MEMRA_PRIME_TOKENWISE=1, and frozen Hy3 CPU/GPU expert splits take the tokenwise
11136 // decode_step_h loop. The latter avoids transient GPU staging of the spilled expert bank.
11137 // EMPTY-SUFFIX CONTINUATION (serve bursts): a session turn with NO new tokens resumes
11138 // generation exactly where the last turn stopped — no prime at all. The stashed
11139 // `next_pred` plays prime_logits' role: it is the token produced from the logits after
11140 // committed.last() by the same rule this entry applies to a cold prime's last row —
11141 // an argmax when greedy, a `sample_boundary_token` draw when sampled (the burst tail,
11142 // or `spec_session_from_restored` for a converted prefix-cache hit, did the drawing
11143 // where the sampler and the session's Philox counters were live). `last_h` seeds the
11144 // predecessor pairing below. Fresh calls and non-empty suffixes take the normal path.
11145 let continuation = prompt.is_empty();
11146 if continuation {
11147 assert!(session_mode, "empty prompt requires a session");
11148 assert!(
11149 sess_tail
11150 .as_ref()
11151 .map_or(false, |(c, lh, np, _, _)| !c.is_empty()
11152 && lh.is_some()
11153 && (np.is_some() || carried_pending.is_some())),
11154 "empty-suffix continuation needs a primed session (committed + last_h + next_pred|pending)"
11155 );
11156 }
11157 let mut prime_logits;
11158 let mut prompt_h: Option<CudaSlice<f32>> = None;
11159 let t_prime = std::time::Instant::now();
11160 let batched_prime = !continuation
11161 && prompt.len() >= crate::hybrid_forward::PRIME_MIN_T
11162 && std::env::var("MEMRA_PRIME_TOKENWISE").is_err()
11163 && !e.frozen_cpu_experts_prefer_tokenwise_prime();
11164 let prime_split = prime_split.filter(|&split| split > 0 && split < prompt.len());
11165 if prime_split.is_some() && continuation {
11166 return Err("spec prime split requires a non-empty prime".into());
11167 }
11168 // STABLE-BOUNDARY TURN CHECKPOINT stop (lane/frspec-multiturn-cache, 2026-08-21):
11169 // the worker's `ckpt_at` request, ABSOLUTE -> prompt-relative. On WARM bursts
11170 // (base != 0, an affinity-rewound or pool-resumed session priming its own delta)
11171 // this is the only stop; on COLD bursts it usually coincides with `prime_split`
11172 // (both are the plain tier's stable pre-generation boundary). A boundary the prime
11173 // cannot honor (outside this prime's range) silently drops the capture — the
11174 // turn_ckpt convention: the next turn re-primes in full, never a wrong resume.
11175 let ckpt_rel = if continuation {
11176 None
11177 } else {
11178 ckpt_req
11179 .and_then(|abs| abs.checked_sub(base))
11180 .filter(|&r| r > 0 && r < prompt.len())
11181 };
11182 // Prime stops, ordered: each is a boundary the prime halts at so the in-place GDN
11183 // conv/ssm state can be snapshotted there (the only moment it exists). One stop =
11184 // the legacy single-split program, byte-for-byte.
11185 let mut stops: Vec<usize> = Vec::new();
11186 for b in [prime_split, ckpt_rel].into_iter().flatten() {
11187 if !stops.contains(&b) {
11188 stops.push(b);
11189 }
11190 }
11191 stops.sort_unstable();
11192 // Captured at the ckpt stop, installed into the session slot post-prime (replacing
11193 // the legacy prompt-end capture). Some(None) = capture attempted and failed -> the
11194 // slot is cleared (a stale checkpoint would rewind to the WRONG boundary).
11195 let mut ckpt_early: Option<Option<SpecCheckpoint>> = None;
11196 if continuation {
11197 prime_logits = Vec::new();
11198 } else if !stops.is_empty() {
11199 if let Some(&first) = stops.first() {
11200 if prime_split == Some(first) && first < crate::hybrid_forward::PRIME_MIN_T {
11201 return Err(format!(
11202 "spec prime split {first} is below PRIME_MIN_T {}",
11203 crate::hybrid_forward::PRIME_MIN_T,
11204 )
11205 .into());
11206 }
11207 }
11208 // Mirror the plain worker's boundary stops exactly. Each segment is a
11209 // request-level prime (`queued_after` keeps Step35 arm selection independent of
11210 // the stops — tick-seg law); a segment below PRIME_MIN_T (and the final tail
11211 // under MEMRA_PRIME_TOKENWISE) takes the same eager tokenwise continuation as
11212 // prefill_tick. Retain every hidden row so the draft scratch fill remains one
11213 // coherent prompt.
11214 let mut h_all = e.uninit(prompt.len() * n_embd)?;
11215 prime_logits = Vec::new();
11216 let mut prev = 0usize;
11217 for seg_end in stops.iter().copied().chain(std::iter::once(prompt.len())) {
11218 if seg_end <= prev {
11219 continue;
11220 }
11221 let seg = &prompt[prev..seg_end];
11222 let is_final = seg_end == prompt.len();
11223 let batched_seg = seg.len() >= crate::hybrid_forward::PRIME_MIN_T
11224 && (!is_final
11225 || (std::env::var("MEMRA_PRIME_TOKENWISE").is_err()
11226 && !e.frozen_cpu_experts_prefer_tokenwise_prime()));
11227 if batched_seg {
11228 let (l, _, h_seg) =
11229 self.prime_cache(e, seg, &mut *cache, prompt.len() - seg_end)?;
11230 e.copy_into(&mut h_all, prev * n_embd, &h_seg, seg.len() * n_embd)?;
11231 prime_logits = l;
11232 } else {
11233 for (i, &tok) in seg.iter().enumerate() {
11234 let (l, h) = self.decode_step_h(e, tok, &mut *cache)?;
11235 e.copy_into(&mut h_all, (prev + i) * n_embd, &h, n_embd)?;
11236 prime_logits = l;
11237 }
11238 }
11239 prev = seg_end;
11240 if is_final {
11241 break;
11242 }
11243 debug_assert_eq!(cache.pos, base + seg_end, "prime stop landed off boundary");
11244 // PREFIX-CACHE BOUNDARY CAPTURE (lane/spec-prefix-cache): the GDN conv/ssm
11245 // states are about to be advanced in place by the next segment, so this is
11246 // the ONLY moment the boundary's recurrent state exists. Capture iff the
11247 // worker requested exactly this stop (cold sessions only — `capture_at` is
11248 // never armed warm). A failed snapshot is silent (turn_ckpt convention) —
11249 // publication is an optimization, never a correctness dependency.
11250 if base == 0 {
11251 if let Some((requested, slot)) = sess_capture.as_mut() {
11252 // Publish at the requested miss-LCP stop (the shared-prefix class)
11253 // AND at the stable-boundary stop (the next-turn re-render class,
11254 // lane/frspec-multiturn-cache) — the same boundary set the plain
11255 // prefill tick learns. Without the second entry, the turn after a
11256 // cold re-park could only hit the OLDER lcp entry (the measured
11257 // one-turn transient: t3 restored 607 of 24122 while the plain arm
11258 // rewound to 15222). Dedupe is the worker sweep's has_key.
11259 if *requested == Some(seg_end) || ckpt_rel == Some(seg_end) {
11260 if let Ok(snap) = cache.snapshot(e) {
11261 slot.push(SpecBoundaryCapture {
11262 snap,
11263 pos: seg_end,
11264 logits: prime_logits.clone(),
11265 // rows [0..seg_end) of h_all are primed — the following
11266 // segments append, never overwrite.
11267 last_h: capture_boundary_hidden(e, &h_all, seg_end, n_embd),
11268 });
11269 }
11270 }
11271 }
11272 }
11273 // SESSION-AFFINITY TURN CHECKPOINT at the STABLE boundary (see `ckpt_at`):
11274 // same snapshot mechanics, installed post-prime in place of the prompt-end
11275 // capture the re-render class always diverged below.
11276 if ckpt_rel == Some(seg_end) {
11277 let anchor: Result<CudaSlice<f32>, Box<dyn std::error::Error>> =
11278 e.uninit(n_embd).and_then(|mut a| {
11279 e.copy_view_into(
11280 &mut a,
11281 0,
11282 &h_all.slice((seg_end - 1) * n_embd..seg_end * n_embd),
11283 n_embd,
11284 )?;
11285 Ok(a)
11286 });
11287 ckpt_early = Some(match (cache.snapshot(e), anchor) {
11288 (Ok(snap), Ok(last_h)) => Some(SpecCheckpoint {
11289 snap,
11290 pos: base + seg_end,
11291 last_h,
11292 }),
11293 _ => None,
11294 });
11295 }
11296 }
11297 if std::env::var("MEMRA_SPEC_STATS").as_deref() == Ok("1") {
11298 eprintln!(
11299 "[spec-prime] stops={stops:?} tail={}",
11300 prompt.len() - stops.last().copied().unwrap_or(0)
11301 );
11302 }
11303 prompt_h = Some(h_all);
11304 } else if batched_prime {
11305 let (l, _h_seed, hiddens) = self.prime_cache(e, prompt, &mut *cache, 0)?;
11306 prime_logits = l;
11307 prompt_h = Some(hiddens);
11308 } else {
11309 prime_logits = Vec::new();
11310 prompt_h = Some(e.uninit(prompt.len() * n_embd)?);
11311 for (i, &tok) in prompt.iter().enumerate() {
11312 let (l, h) = self.spec_target_step_h(e, tok, &mut *cache)?;
11313 if let Some(ph) = prompt_h.as_mut() {
11314 e.copy_into(ph, i * n_embd, &h, n_embd)?;
11315 }
11316 prime_logits = l;
11317 }
11318 }
11319 e.stream().synchronize()?;
11320 // PREFIX-CACHE SEED CAPTURE (lane/spec-prefix-cache): boundary == prompt end (the seed
11321 // case — no shared-prefix split, publish the whole prompt). The prime just finished, so
11322 // cache.pos == base + prompt.len() and the recurrent state IS the boundary state;
11323 // prime_logits are the boundary logits. Cold sessions only (base == 0) — same law as
11324 // prime_split. The mid-prompt capture above already consumed the request if it matched.
11325 if !continuation && base == 0 {
11326 if let Some((requested, slot)) = sess_capture.as_mut() {
11327 if *requested == Some(prompt.len()) && slot.is_empty() {
11328 debug_assert_eq!(cache.pos, prompt.len(), "seed capture off prompt end");
11329 if let Ok(snap) = cache.snapshot(e) {
11330 slot.push(SpecBoundaryCapture {
11331 snap,
11332 pos: prompt.len(),
11333 logits: prime_logits.clone(),
11334 last_h: prompt_h
11335 .as_ref()
11336 .map(|ph| capture_boundary_hidden(e, ph, prompt.len(), n_embd))
11337 .unwrap_or_default(),
11338 });
11339 }
11340 }
11341 }
11342 }
11343 // Harness timing contract (see crate::PRIME_NANOS): gen-only throughput without the
11344 // prime-subtraction hack.
11345 crate::PRIME_NANOS.store(
11346 t_prime.elapsed().as_nanos() as u64,
11347 std::sync::atomic::Ordering::Relaxed,
11348 );
11349
11350 let (embd_qt, embd_rb) = self.embd.qt_and_row_bytes(n_embd);
11351 // Resident table is fastest when it fits. Large spill deployments can preserve that HBM
11352 // for expert-cache slots and gather only the exact rows needed by MTP/verify from host.
11353 let host_embd = spec_host_embd();
11354 let embd_gpu = if host_embd {
11355 None
11356 } else {
11357 Some(
11358 self.embd_gpu
11359 .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload")),
11360 )
11361 };
11362 let embd_dev = embd_gpu.map(|g| (g, embd_qt, embd_rb));
11363 if host_embd {
11364 eprintln!(
11365 "[spec] host-row embedding: {} bytes kept off HBM",
11366 self.embd.raw.len()
11367 );
11368 }
11369 let mut out: Vec<u32> = Vec::with_capacity(max_new);
11370 let mut total_drafted = 0usize;
11371 let mut total_accepted = 0usize;
11372
11373 // --- SAMPLER FIRST (lane/sampled-spec-quality, 2026-08-19) ---
11374 // The sampler config, the session's Philox counters and the penalty window are parsed
11375 // HERE, above the boundary-token selection, because the boundary token must be drawn
11376 // from the sampler the request asked for. Pre-lane this block sat ~50 lines BELOW the
11377 // selection, which is the whole mechanical reason the boundary token was an argmax:
11378 // the sampler state was not in scope yet. Nothing here depends on the round loop, so
11379 // moving it up is a pure reordering for greedy (`sampled == false` ⇒ every branch
11380 // below takes the argmax path it always took).
11381 // --- SAMPLED SPEC (MEMRA_SPEC_TEMP>0, research/sampled-spec-impl-map.md): rejection-
11382 // sampling verify (Leviathan/Chen) — accept draft x at u < p(x)/q(x), resample from
11383 // norm(max(0,p-q)) on reject, bonus sampled from p on full accept. Counter-based Philox
11384 // everywhere (seed, event) -> reproducible. temp==0/unset = the greedy path, untouched.
11385 let sp = sampling.unwrap_or_else(|| SpecSampling {
11386 temp: std::env::var("MEMRA_SPEC_TEMP")
11387 .ok()
11388 .and_then(|v| v.parse().ok())
11389 .unwrap_or(0.0),
11390 seed: std::env::var("MEMRA_SEED")
11391 .ok()
11392 .and_then(|v| v.parse().ok())
11393 .unwrap_or(42),
11394 top_k: std::env::var("MEMRA_TOP_K")
11395 .ok()
11396 .and_then(|v| v.parse().ok())
11397 .unwrap_or(0),
11398 top_p: std::env::var("MEMRA_TOP_P")
11399 .ok()
11400 .and_then(|v| v.parse().ok())
11401 .unwrap_or(1.0),
11402 min_p: std::env::var("MEMRA_MIN_P")
11403 .ok()
11404 .and_then(|v| v.parse().ok())
11405 .unwrap_or(0.0),
11406 penalty_last_n: std::env::var("MEMRA_PENALTY_LAST_N")
11407 .ok()
11408 .and_then(|v| v.parse().ok())
11409 .unwrap_or(0),
11410 penalty_repeat: std::env::var("MEMRA_PENALTY_REPEAT")
11411 .ok()
11412 .and_then(|v| v.parse().ok())
11413 .unwrap_or(1.0),
11414 penalty_freq: std::env::var("MEMRA_PENALTY_FREQ")
11415 .ok()
11416 .and_then(|v| v.parse().ok())
11417 .unwrap_or(0.0),
11418 penalty_present: std::env::var("MEMRA_PENALTY_PRESENT")
11419 .ok()
11420 .and_then(|v| v.parse().ok())
11421 .unwrap_or(0.0),
11422 });
11423 let (sp_temp, sp_seed) = (sp.temp, sp.seed);
11424 let sampled = sp_temp > 0.0;
11425 // Counters resume from the session (burst continuity: randomness must never repeat
11426 // across generate_spec_session calls); one-shot callers start at (0,0). Read through
11427 // sess_tail — `sess` was take()n into it above, so sess.as_ref() here is always None.
11428 let mut sctr: u32 = sess_tail.as_ref().map(|(_, _, _, s, _)| **s).unwrap_or(0);
11429 let mut uctr: u32 = sess_tail.as_ref().map(|(_, _, _, _, u)| **u).unwrap_or(0);
11430 // Penalties (v2.1): applied to COPIES of q rows and p columns symmetrically (exactness
11431 // for the penalized+filtered target). History = generated tokens, host-tracked window.
11432 let pen_on = sampled
11433 && sp.penalty_last_n > 0
11434 && (sp.penalty_repeat != 1.0 || sp.penalty_freq != 0.0 || sp.penalty_present != 0.0);
11435 // SESSION-SPANNING PENALTY WINDOW (Item 2). Pre-lane this was
11436 // `prompt.iter().rev().take(64).rev()` — the BURST's suffix slice — so a continuation
11437 // burst (the majority of a stream's tokens, and ALL of a converted cache hit's) started
11438 // with an EMPTY penalty history and the client's repetition/frequency/presence penalties
11439 // silently reset at every burst boundary. The window now spans `committed ++ prompt`,
11440 // which is what the API contract says and what the plain sampler's own `history` does.
11441 // Byte-identical to the pre-lane seed for a cold turn-1 burst at the default window.
11442 let mut pen_hist: Vec<u32> = if pen_on {
11443 let sess_hist: &[u32] = if spec_pen_session_on() {
11444 sess_tail
11445 .as_ref()
11446 .map(|(c, ..)| c.as_slice())
11447 .unwrap_or(&[])
11448 } else {
11449 &[] // MEMRA_SPEC_PEN_SESSION=0: pre-lane burst-local window
11450 };
11451 pen_window_seed(sess_hist, prompt, sp.penalty_last_n)
11452 } else {
11453 Vec::new()
11454 };
11455 // First generated token = the BOUNDARY token: greedy takes the argmax of the prompt's
11456 // last logits (== greedy's first token, byte-contract); SAMPLED draws it from the
11457 // request's own filtered/penalized target through the session's Philox stream
11458 // (`sample_boundary_token`, lane/sampled-spec-quality Item 1 — pre-lane this was an
11459 // argmax in both regimes, so ~1 token per burst of a sampled stream was greedy).
11460 // Emit it, then FEED it to establish the loop invariant below.
11461 // PENDING-CARRY: the carried bonus was already emitted by the LAST burst — it becomes
11462 // last_token WITHOUT re-emission, and round 0 consumes it as pending (no init feed).
11463 // CONSTRAINED entry rules: the first emitted token is the MASKED argmax of the
11464 // prompt's last logits (plain constrained-greedy identity); a continuation without
11465 // a carried pending would emit an UNMASKED stashed next_pred — refused loudly (the
11466 // worker never resumes constrained sessions from the pool, so this cannot fire).
11467 if let Some(c) = constraint.as_deref_mut() {
11468 if continuation && carried_pending.is_none() {
11469 return Err("constrained spec continuation requires a carried pending \
11470 (pool resume is unconstrained-only)"
11471 .into());
11472 }
11473 if !continuation {
11474 c.mask_logits(&mut prime_logits)
11475 .map_err(|e2| format!("constraint: {e2}"))?;
11476 }
11477 }
11478 let mut last_token = if let Some(b) = carried_pending {
11479 b
11480 } else if continuation {
11481 // A continuation's boundary token was DRAWN by the burst that stashed it (the
11482 // session tail below), or by `spec_session_from_restored` for a converted
11483 // prefix-cache hit — in both cases from the correct logits row with this same
11484 // session's Philox stream, which is why it can be consumed here as-is.
11485 sess_tail.as_ref().unwrap().2.unwrap()
11486 } else if sampled && constraint.is_none() && spec_sampled_boundary_on() {
11487 sample_boundary_token(e, &prime_logits, &sp, &pen_hist, &mut sctr, "cold-prime")?
11488 } else {
11489 // greedy (byte contract), the rollback door, or constrained (masked-argmax
11490 // identity — the worker routes sampled+constrained to the plain path, and this
11491 // function refuses the combination outright above).
11492 argmax(&prime_logits) as u32
11493 };
11494 if pen_on {
11495 // The boundary token is a GENERATED token: the plain sampler `accept()`s every
11496 // emitted token into its penalty history, and pre-lane the burst's first token
11497 // was invisible to penalties forever (never pushed, and never in `committed`
11498 // until this burst's tail). Covers the carry/continuation seeds too — neither is
11499 // in `committed` yet.
11500 pen_hist.push(last_token);
11501 }
11502 if carried_pending.is_none() {
11503 out.push(last_token);
11504 // grammar advances with every emitted token (carried pendings were consumed
11505 // by the burst that emitted them).
11506 if let Some(c) = constraint.as_deref_mut() {
11507 c.consume(last_token)
11508 .map_err(|e2| format!("constraint: {e2}"))?;
11509 }
11510 }
11511 if continuation {
11512 // draft-KV invariant: entries [0..base) are the session's exact fills; truncate any
11513 // overhang so the chain's first append lands at slot base (== committed.len()).
11514 scratch.set_len(e, base)?;
11515 }
11516 // sse-cadence: hand the caller every not-yet-flushed token (disjoint in-order slices
11517 // concatenating to the full `out`). Called after the prime's first token and after each
11518 // round commit — emission timing only, token bytes untouched. The slice may be EMPTY
11519 // (poll-only boundary: zero-round folds commit nothing new); returns the caller's
11520 // continue-verdict (admission yield, 2026-08-06) — false ends the burst at this round.
11521 fn flush_commit(
11522 cb: &mut Option<&mut dyn FnMut(&[u32]) -> bool>,
11523 out: &[u32],
11524 flushed: &mut usize,
11525 ) -> bool {
11526 if let Some(f) = cb.as_mut() {
11527 let keep = f(&out[*flushed..]);
11528 *flushed = out.len();
11529 keep
11530 } else {
11531 true
11532 }
11533 }
11534 keep_going = flush_commit(&mut on_commit, &out, &mut flushed);
11535 // INVARIANT at loop top: `last_token` is the most-recently-committed/emitted token, its
11536 // KV+recur state IS in `cache` (cache.pos = position right AFTER last_token), `last_pred`
11537 // is the greedy ARGMAX of the logits that predict the token FOLLOWING last_token, and
11538 // `h_seed` = last_token's pre-output_norm hidden. Establish it by feeding last_token once
11539 // (mirrors plain greedy). DEVICE-ARGMAX lever: the accept walk only ever consumes the
11540 // argmax of those logits — never the full vector — so a host u32 replaces the Vec<f32>.
11541 // Trimmed heads: q lives on the trimmed vocab; accept gathers use the TRIMMED index and
11542 // the residual scatters q into target-id space (q=-inf off-trim — the head cannot propose
11543 // those, so their residual mass is p(x), correct by construction).
11544 let d2t_dev: Option<CudaSlice<u32>> = if sampled || crate::spec::spec_stream() {
11545 match &mtp.d2t {
11546 Some(map) => Some(e.htod_u32_v(map)?),
11547 None => None,
11548 }
11549 } else {
11550 None
11551 };
11552 let mut q_full_buf: Option<CudaSlice<f32>> = None;
11553 // host Philox4x32-10 accept-test uniforms: module fn `host_u01` (shared with the
11554 // dspark sampled-admission walk); byte-identical to the closure it replaces.
11555 let mut draft_logits: Vec<CudaSlice<f32>> = Vec::new(); // retained head logits (q), per slot
11556 let mut draft_stats: Vec<(f32, f32, f32)> = Vec::new(); // (row_max, th_e, z_e) per slot
11557 let mut perturb_buf: Option<CudaSlice<f32>> = None; // gumbel scratch (max(n_vocab,d_vocab))
11558 let mut sample_tok = e.alloc_u32_zeroed(1)?; // residual/bonus sample out
11559 let mut col_buf: Option<CudaSlice<f32>> = None; // materialized verify column
11560 let mut pen_hist_d: Option<CudaSlice<u32>> = None;
11561 let mut pcol_buf: Option<CudaSlice<f32>> = None; // penalized p-column scratch
11562 // MEMRA_SPEC_SETUP_TRACE=1 (diagnostics): per-call wall decomposition of the burst
11563 // SETUP + TAIL segments (the round loop's internals are MEMRA_SPEC_PHASE's job) —
11564 // built to pin the serve per-burst fixed cost (research/spec-serving-20260801).
11565 let setup_trace = std::env::var("MEMRA_SPEC_SETUP_TRACE").as_deref() == Ok("1");
11566 let t_ent = std::time::Instant::now();
11567
11568 // SESSION-AFFINITY TURN CHECKPOINT (lane/session-affinity, 2026-08-05): capture the
11569 // PROMPT-END boundary state so a LATER turn can rewind here and re-prime only its own
11570 // delta instead of the whole conversation. See `SpecCheckpoint` for why this boundary is
11571 // the one that matters (a history-rewriting client mutates what the session GENERATED,
11572 // so the next turn's prompt agrees with this one up to exactly here).
11573 //
11574 // WHERE — AND WHY THIS EXACT LINE. Right after the trunk prime, BEFORE the init feed
11575 // (`decode_step_h(last_token)`) and before round 0: the last instant at which the caches
11576 // hold exactly `base + prompt.len()` rows and nothing generated.
11577 //
11578 // This was WRONG in the first cut of this lane: the capture sat after the draft-KV fill,
11579 // which is also after the init feed, so `cache.pos` was `base + prompt.len() + 1` — the
11580 // boundary included the FIRST GENERATED TOKEN. That token is the first thing inside the
11581 // `<think>` block the client strips, so every later turn's diff diverged exactly one
11582 // token below the checkpoint and affinity declined 100% of the time. Measured on the
11583 // owner regime: "history diverged at 12233 of checkpoint 12234". The off-by-one made the
11584 // whole mechanism inert while looking, from the outside, like a working
11585 // correctness-declines-safely path — hence the decline log carries the offsets.
11586 //
11587 // The full-attn planes are `len`-truncatable so the snapshot copies only the GDN conv/ssm
11588 // state (the reason a spec session could not rewind before). The draft scratch needs no
11589 // copy: rows below the boundary are rewritten by the next turn's own fill.
11590 //
11591 // WHEN: non-empty prime only. An empty-suffix continuation burst adds no prompt boundary
11592 // (its "prompt end" IS the previous checkpoint's, already held), so it keeps the existing
11593 // checkpoint rather than replacing it with a strictly worse one.
11594 //
11595 // FAILURE IS SILENT BY DESIGN: on a VRAM-tight rig the snapshot alloc can fail. That
11596 // costs the NEXT turn its rewind (it re-primes fully, today's behavior) and must never
11597 // fail the burst that is already running — so the error is swallowed, loud only under
11598 // MEMRA_DEBUG_SPEC.
11599 //
11600 // STABLE-BOUNDARY OVERRIDE (lane/frspec-multiturn-cache, 2026-08-21): the prompt-end
11601 // posture above was DISPROVED for the think-posture template class — the prompt's own
11602 // tail is the live generation header (`<|im_start|>assistant\n<think>\n`) that the
11603 // next turn's re-render replaces, so the diff diverged a couple tokens BELOW the
11604 // checkpoint and affinity declined 100% of multi-turn agent traffic (the same class
11605 // the plain tier fixed on 2026-08-09 via `plain_checkpoint_boundary`; the port to the
11606 // spec tier is this lane). When the worker armed `ckpt_at`, the capture happened at
11607 // that stop inside the prime above (`ckpt_early`) and is installed here instead;
11608 // capture-attempted-but-failed clears the slot exactly like the legacy arm.
11609 if let Some(slot) = sess_ckpt_slot {
11610 if let Some(early) = ckpt_early {
11611 if early.is_none() && std::env::var("MEMRA_DEBUG_SPEC").is_ok() {
11612 eprintln!(
11613 "[spec] stable-boundary turn checkpoint skipped; \
11614 next turn re-primes in full"
11615 );
11616 }
11617 *slot = early;
11618 } else if !continuation {
11619 let pos = cache.pos;
11620 debug_assert_eq!(
11621 pos,
11622 base + prompt.len(),
11623 "turn checkpoint must sit at the prompt end, before the init feed"
11624 );
11625 let anchor: Result<CudaSlice<f32>, Box<dyn std::error::Error>> =
11626 if let Some(ph) = &prompt_h {
11627 // hidden of the LAST primed row = the predecessor anchor at this
11628 // boundary (exactly what a fresh prime of committed[..pos] leaves in
11629 // last_h, and what the next prime's fill reads for its first row).
11630 let np = prompt.len();
11631 e.uninit(n_embd).and_then(|mut a| {
11632 e.copy_view_into(
11633 &mut a,
11634 0,
11635 &ph.slice((np - 1) * n_embd..np * n_embd),
11636 n_embd,
11637 )?;
11638 Ok(a)
11639 })
11640 } else {
11641 Err("no prompt hiddens".into())
11642 };
11643 match (cache.snapshot(e), anchor) {
11644 (Ok(snap), Ok(last_h)) => {
11645 *slot = Some(SpecCheckpoint { snap, pos, last_h });
11646 }
11647 (s, a) => {
11648 *slot = None; // a stale checkpoint would rewind to the WRONG boundary
11649 if std::env::var("MEMRA_DEBUG_SPEC").is_ok() {
11650 let err = s
11651 .err()
11652 .map(|e| e.to_string())
11653 .or_else(|| a.err().map(|e| e.to_string()))
11654 .unwrap_or_default();
11655 eprintln!(
11656 "[spec] turn checkpoint skipped ({err}); \
11657 next turn re-primes in full"
11658 );
11659 }
11660 }
11661 }
11662 }
11663 }
11664 // INIT FEED — skipped on a pending carry: last_token (the carried bonus) is NOT in the
11665 // caches and must NOT be fed solo; round 0's batched verify commits it as col 0. Its
11666 // seed/anchor hidden is the carried last_h (copied below); last_pred is dead in the
11667 // pending path (t_pred reads verify col 0 — the accept walk overwrites it).
11668 let mut last_pred = 0u32;
11669 let mut last_col_logits: Option<CudaSlice<f32>> = None;
11670 // CONSTRAINED: the init feed's logits back the (n_acc==0, base==0) masked-argmax
11671 // recompute in the grammar-truncation walk — retained host-side, round 0 only.
11672 let mut init_logits_host: Option<Vec<f32>> = None;
11673 let h_seed0: CudaSlice<f32> = if carried_pending.is_none() {
11674 let (init_logits, h) = self.spec_target_step_h(e, last_token, &mut *cache)?;
11675 last_pred = argmax(&init_logits) as u32;
11676 if constraint.is_some() {
11677 init_logits_host = Some(init_logits.clone());
11678 }
11679 // sampled mode: p-distribution after last_token, for the j==0/base==0 accept test.
11680 if sampled {
11681 last_col_logits = Some(e.htod(&init_logits)?);
11682 }
11683 h
11684 } else {
11685 // predecessor-row anchor: hidden of the last COMMITTED row (the carry contract).
11686 let lh = sess_tail
11687 .as_ref()
11688 .unwrap()
11689 .1
11690 .as_ref()
11691 .expect("pending carry requires last_h");
11692 e.clone_dtod(lh)?
11693 };
11694 let t_init = t_ent.elapsed();
11695 let mut last_col_stats: Option<(f32, f32, f32)> = None;
11696 // PERSISTENT h_seed buffer (allocated BEFORE any graph capture so no captured scratch can
11697 // alias it): every path that updates the round seed copies INTO it — no per-round allocs,
11698 // stable pointer for the graph-draft round-start copy.
11699 let mut h_seed_buf = e.clone_dtod(&h_seed0)?;
11700 // Predecessor-pairing trackers: `fill_prev` = trunk hidden AT the last COMMITTED row (the
11701 // predecessor of the next verify's col 0 — the reference's carried pending-h analogue;
11702 // also the predecessor-row hidden for the round-0 legacy-replay seed). At round 0 that
11703 // row is last_token's own (h_seed0). The chain step-0 seed under the pairing default =
11704 // hidden of the row BEFORE last_token = the prompt's last row at round 0 (h_seed_buf
11705 // overwritten below).
11706 let mut fill_prev = e.clone_dtod(&h_seed0)?;
11707 {
11708 if let Some(ph) = &prompt_h {
11709 let np = prompt.len();
11710 e.copy_view_into(
11711 &mut h_seed_buf,
11712 0,
11713 &ph.slice((np - 1) * n_embd..np * n_embd),
11714 n_embd,
11715 )?;
11716 } else if continuation {
11717 if let Some((_, lh, _, _, _)) = sess_tail.as_ref() {
11718 if let Some(lh) = lh.as_ref() {
11719 e.copy_into(&mut h_seed_buf, 0, lh, n_embd)?;
11720 }
11721 }
11722 }
11723 }
11724 // Persistent device prediction slots for the accept walk (max k+1 verify columns).
11725 let mut preds_d = e.alloc_u32_zeroed(k + 2)?;
11726
11727 let debug_spec = std::env::var("MEMRA_DEBUG_SPEC").is_ok();
11728 let fork_mode = OptiForkGateMode::configured();
11729 // MEMRA_SPEC_STATS=1: per-slot accept histogram + draft-length histogram, printed once at
11730 // the end. Metric normalization vs the reference engine: BOTH engines count
11731 // accepted/drafted where the chain stopped at p-min and the sub-threshold token is
11732 // discarded uncounted — per-slot decay + chain-length mix are the extra dimensions.
11733 let spec_stats = std::env::var("MEMRA_SPEC_STATS").is_ok();
11734 let mut st_drafted = vec![0usize; k];
11735 let mut st_accepted = vec![0usize; k];
11736 let mut st_len_hist = vec![0usize; k + 1];
11737 let mut st_full = 0usize;
11738 // P-MIN CONFIDENCE GATE (MEMRA_SPEC_PMIN, the serve script's --spec-draft-p-min mechanism):
11739 // stop the draft chain early when the head's softmax confidence in its own pick drops
11740 // below p_min. Hoisted above the loop: the graph capture bakes the prob kernels iff on.
11741 static PMIN: std::sync::OnceLock<f32> = std::sync::OnceLock::new();
11742 let p_min = *PMIN.get_or_init(|| {
11743 std::env::var("MEMRA_SPEC_PMIN")
11744 .ok()
11745 .and_then(|v| v.parse().ok())
11746 .unwrap_or(0.0)
11747 });
11748 // ZERO-DRAFT ROUNDS (MEMRA_SPEC_PMIN0=1, vendored from llama.cpp's draft gating): let the
11749 // p-min gate apply at j==0 too, so a low-confidence round drafts NOTHING and the verify
11750 // batch is just the pending bonus (m=1 = a plain decode step). llama's 35B win rides
11751 // exactly this — draft acceptance 76% at mean len 2.5 because unpredictable stretches
11752 // never pay draft+verify overhead. Only legal when a pending bonus exists (an empty
11753 // verify batch is not); the j==0 exemption stays for pending-less rounds.
11754 let pmin0 = std::env::var("MEMRA_SPEC_PMIN0")
11755 .map(|v| v == "1")
11756 .unwrap_or(false);
11757
11758 // --- GRAPH DRAFT setup: persistent I/O buffers + ONE capture (2 warmups inside). The
11759 // warmups mutate scratch len_d / pos / tok / seed — all reset at every round start, so the
11760 // only restore needed is the scratch counter. Capture failure (e.g. a non-capturable
11761 // cuBLAS path in an exotic head) falls back to the eager draft chain.
11762 // PER-SESSION PERSISTENCE (2026-08-01): session calls reuse the DraftGraphCtx parked on
11763 // the SpecSession — the capture (2 warmup head forwards + instantiate) ran ONCE at the
11764 // session's first burst, not per burst (measured ~16ms/burst fixed cost on H100 q27,
11765 // research/spec-serving-20260801). Reuse is pointer-exact: the graph bakes the session's
11766 // own scratch KV (never realloc'd), the model's resident embedding, the OnceLock p_min,
11767 // and the g_* buffers carried in the ctx — replay dispatch is identical to a fresh
11768 // capture, so draft tokens are bit-identical (drafts never decide exactness anyway; the
11769 // verify arbitrates). Single-shot calls (sess=None) build a fresh ctx and drop it.
11770 let mut dctx: DraftGraphCtx = match sess_draft_slot.as_mut().and_then(|s| s.take()) {
11771 Some(c) => c,
11772 None => DraftGraphCtx::new(e, n_embd, if sampled { d_vocab } else { 1 })?,
11773 };
11774 // A session that ran greedy bursts first sized g_q/g_perturb at 1; a sampled resume
11775 // needs d_vocab. Realloc is legal exactly while graph_s is None (nothing baked them).
11776 if sampled && dctx.g_q.len() < d_vocab {
11777 dctx.g_q = e.zeros(d_vocab)?;
11778 dctx.g_perturb = e.zeros(d_vocab)?;
11779 }
11780 // DRAFT-SIDE GRAMMAR MASK (lane/draft-mask, 2026-08-04): the drafter samples the
11781 // grammar's legal set, so proposals are legal BY CONSTRUCTION and the verify-side
11782 // truncation (the correctness backstop) stops cutting every tight-schema round.
11783 // The mask is one node inside the captured draft chain — presence is a CAPTURE-TIME
11784 // shape, so a parked graph of the other shape is dropped and recaptured.
11785 let dmask_on = constraint
11786 .as_deref()
11787 .is_some_and(|c| c.draft_mask_enabled());
11788 let dmask_words = if dmask_on { d_vocab.div_ceil(32) } else { 0 };
11789 if dmask_on && dctx.g_dmask.len() < dmask_words {
11790 dctx.g_dmask = e.alloc_u32_zeroed(dmask_words)?;
11791 dctx.graph = None; // the old capture baked the old (or no) mask pointer
11792 dctx.failed.clear_greedy();
11793 dctx.keeper.clear();
11794 }
11795 if dctx.graph.is_some() && dctx.graph_masked != dmask_on {
11796 dctx.graph = None;
11797 dctx.failed.clear_greedy();
11798 dctx.keeper.clear();
11799 }
11800 if graph_draft && !sampled && dctx.graph.is_none() && !dctx.failed.greedy_failed() {
11801 // dcw door: the capture warmups append device-counter rows the capture body cannot
11802 // rebase for; pre-arm ring headroom host-side (no-op on flat planes / room-enough
11803 // rings, and the door-off path is untouched).
11804 if step35_draft_dcw_on() {
11805 scratch.ensure_dcw_headroom(e, k + 2)?;
11806 }
11807 let DraftGraphCtx {
11808 g_tok,
11809 g_pos,
11810 g_seed,
11811 g_p,
11812 g_dmask,
11813 ..
11814 } = &mut dctx;
11815 // capture-time contents: ALL-ONES (ban nothing). A replay only ever runs after the
11816 // host uploads the position's real words, so the warmups stay grammar-free.
11817 if dmask_on {
11818 e.htod_u32_into(g_dmask, &vec![u32::MAX; dmask_words])?;
11819 }
11820 let g_dmask_ro: &CudaSlice<u32> = &*g_dmask;
11821 // CAPTURE-RETAIN (#68 fix): the warmup transients' pool addresses are baked into the
11822 // captured graph; the keeper pins them for the graph's lifetime. capture_graph (non-
11823 // retained) freed them at exit — safe for one-shot generate_spec (nothing else touches
11824 // the pool between replays) but WRONG for sessions: burst-boundary prime/fill/commit
11825 // passes (and, in serve, other sessions) recycle those addresses and the replay then
11826 // clobbers live buffers — the ST serve-spec corruption (research/serve-st-20260803).
11827 let cap_res = e.capture_graph_retained(|e| {
11828 self.mtp_head_forward_cap(
11829 e,
11830 mtp,
11831 g_tok,
11832 g_pos,
11833 g_seed,
11834 g_p,
11835 &mut *scratch,
11836 p_min > 0.0 || fork_mode == OptiForkGateMode::Controller,
11837 true,
11838 embd_gpu.expect("graph draft requires resident embedding"),
11839 embd_qt,
11840 embd_rb,
11841 d_vocab,
11842 None,
11843 None,
11844 if dmask_on {
11845 Some((g_dmask_ro, dmask_words))
11846 } else {
11847 None
11848 },
11849 )
11850 });
11851 match cap_res {
11852 Ok((g, keep)) => {
11853 scratch.set_len(e, base)?;
11854 dctx.graph = Some(g);
11855 dctx.graph_masked = dmask_on;
11856 dctx.keeper = keep;
11857 }
11858 Err(err) => {
11859 scratch.set_len(e, base)?;
11860 // LOUD flip (audit Q2): a dropped draft graph is a coverage loss, never
11861 // silent. Once per flip — mark returns None on an already-failed ctx.
11862 if let Some(line) = dctx.failed.mark_greedy(&err.to_string()) {
11863 eprintln!("{line}");
11864 }
11865 }
11866 }
11867 }
11868 // --- SAMPLED GRAPH DRAFT setup (step 3 of the sampled-spec arc): a SECOND capture, own
11869 // graph object, built only when sampled && graph-eligible — the greedy capture above is
11870 // untouched (and skipped when sampled: its graph would never be launched). Same head
11871 // forward, but the in-graph argmax reads GUMBEL-PERTURBED logits; the Philox event
11872 // counter lives in the persistent device g_ctr (bumped in-graph, host-seeded from sctr
11873 // once per round); the raw head logits land in the persistent g_q for the host's
11874 // per-replay async D2D into the round's q slot (q_slots, K x d_vocab, allocated once).
11875 // seed/temp are capture-time constants — baked into graph_s, so a pool-resumed request
11876 // with a different (seed, temp, k) drops the parked sampled graph and recaptures.
11877 // COST OF THE FRESH-SEED SERVE DEFAULT (dogfood F4, 2026-08-04): omitting `seed` on a
11878 // serve request now draws fresh per-request entropy (it used to default to a pinned 0),
11879 // so a seed-omitting request that RESUMES a parked spec session finds an s_key baked
11880 // with the PREVIOUS request's seed and pays one recapture. Bounded, and it does not
11881 // reopen the ~16ms/burst regression the persistent ctx exists to fix: a session's seed
11882 // is fixed for its whole lifetime (worker.rs reads s.sampler.seed() per burst), so
11883 // this compare misses at most ONCE per resumed request — the first burst recaptures
11884 // and every later burst in that request replays. A client that wants the parked graph
11885 // AND reproducibility supplies an explicit `seed`, honored exactly, which keeps s_key
11886 // stable across its whole conversation.
11887 // COMPOSITION RULE (fspec x gsd merge): the in-graph chain samples from the RAW
11888 // softmax — it can hold neither per-row filter stats nor the varying penalty history.
11889 // The sampled graph therefore engages only in the PURE-TEMP regime; filters/penalties
11890 // force the eager draft (which computes stats/penalties per row).
11891 // KEY THE WHOLE REGIME, not just the baked constants (lane/graph-s-key-exactness-
11892 // 20260819). `s_key` used to be `(seed, temp, k)`; the filters and penalties were left
11893 // out, so a filtered request resuming a session that parked a PURE-TEMP graph kept it —
11894 // and the launch site never re-asked `pure_temp`. See [`SampledGraphKey`] for what that
11895 // costs (an unconditional accept of out-of-head draft tokens, i.e. an exactness bug on
11896 // the request shape the vendor-default flip makes the majority).
11897 let s_key = SampledGraphKey::new(sp_seed, sp_temp, k, sp.top_k, sp.top_p, sp.min_p, pen_on);
11898 let pure_temp = s_key.pure_temp();
11899 if sampled && dctx.s_key.is_some_and(|old| old != s_key) {
11900 dctx.graph_s = None;
11901 dctx.failed.clear_sampled();
11902 dctx.s_key = None;
11903 dctx.q_slots.clear();
11904 dctx.keeper_s.clear();
11905 }
11906 if graph_draft
11907 && sampled
11908 && pure_temp
11909 && dctx.graph_s.is_none()
11910 && !dctx.failed.sampled_failed()
11911 {
11912 // dcw door: same warmup headroom pre-arm as the greedy capture above.
11913 if step35_draft_dcw_on() {
11914 scratch.ensure_dcw_headroom(e, k + 2)?;
11915 }
11916 let DraftGraphCtx {
11917 g_tok,
11918 g_pos,
11919 g_seed,
11920 g_p,
11921 g_ctr,
11922 g_perturb,
11923 g_q,
11924 ..
11925 } = &mut dctx;
11926 // CAPTURE-RETAIN (#68 fix): same keeper contract as the greedy capture above.
11927 let cap_res = e.capture_graph_retained(|e| {
11928 self.mtp_head_forward_cap(
11929 e,
11930 mtp,
11931 g_tok,
11932 g_pos,
11933 g_seed,
11934 g_p,
11935 &mut *scratch,
11936 p_min > 0.0,
11937 true,
11938 embd_gpu.expect("graph draft requires resident embedding"),
11939 embd_qt,
11940 embd_rb,
11941 d_vocab,
11942 Some((g_ctr, g_perturb, g_q, sp_seed, sp_temp)),
11943 None,
11944 None, // constrained spec is greedy-only — sampled never carries a hook
11945 )
11946 });
11947 match cap_res {
11948 Ok((g, keep)) => {
11949 scratch.set_len(e, base)?;
11950 for _ in 0..k {
11951 dctx.q_slots.push(e.zeros(d_vocab)?);
11952 }
11953 dctx.graph_s = Some(g);
11954 dctx.s_key = Some(s_key);
11955 dctx.keeper_s = keep;
11956 }
11957 Err(err) => {
11958 scratch.set_len(e, base)?;
11959 // LOUD flip (audit Q2): same contract as the greedy capture above.
11960 if let Some(line) = dctx.failed.mark_sampled(&err.to_string()) {
11961 eprintln!("{line}");
11962 }
11963 }
11964 }
11965 }
11966 // ---- EXACTNESS GUARD, the enforceable half (lane/graph-s-key-exactness-20260819) ----
11967 // With the filters and penalties in `s_key`, a graph that SURVIVED the drop above was
11968 // captured under this request's exact regime, and capture requires `pure_temp` — so a
11969 // parked graph implies `pure_temp`. That implication is the whole exactness argument for
11970 // the graph arm, so it is asserted here rather than assumed: a future change that widens
11971 // the capture condition, narrows the key, or copies a `DraftGraphCtx` across regimes
11972 // fails LOUDLY at this line instead of silently drafting from the raw softmax while the
11973 // verify applies filter stats. Release builds refuse the graph (drop it, draft eager)
11974 // rather than launching it; the launch site re-tests `pure_temp` independently.
11975 if sampled && !pure_temp && dctx.graph_s.is_some() {
11976 debug_assert!(
11977 false,
11978 "sampled draft graph parked under {:?} survived into a FILTERED request \
11979 (top_k={} top_p={} min_p={} pen_on={}): the in-graph chain draws from the RAW \
11980 softmax, so the verify's filtered q would test a distribution the draft was \
11981 never sampled from",
11982 dctx.s_key, sp.top_k, sp.top_p, sp.min_p, pen_on,
11983 );
11984 eprintln!(
11985 "[spec] BUG: dropping a parked sampled draft graph that outlived its capture \
11986 regime (s_key={:?}, request top_k={} top_p={} min_p={} pen_on={}); drafting \
11987 EAGER — the key must carry every field that shapes q",
11988 dctx.s_key, sp.top_k, sp.top_p, sp.min_p, pen_on,
11989 );
11990 dctx.graph_s = None;
11991 dctx.s_key = None;
11992 dctx.q_slots.clear();
11993 dctx.keeper_s.clear();
11994 }
11995 // SKEY PROBE (MEMRA_SKEY_PROBE=1): the burst-entry facts the reachability question turns
11996 // on — is this request sampled, is it in the pure-temp regime the sampled graph is only
11997 // legal in, and is a graph PARKED from an earlier request of the same session? The launch
11998 // arms below print which chain actually ran, so the probe never restates the condition.
11999 if skey_probe() {
12000 eprintln!(
12001 "[skey] burst sampled={} pure_temp={} temp={} top_k={} top_p={} min_p={} \
12002 pen_on={} k={} graph_draft={} graph_s_parked={} s_key_parked={:?}",
12003 sampled as u8,
12004 pure_temp as u8,
12005 sp_temp,
12006 sp.top_k,
12007 sp.top_p,
12008 sp.min_p,
12009 pen_on as u8,
12010 k,
12011 graph_draft as u8,
12012 dctx.graph_s.is_some() as u8,
12013 dctx.s_key,
12014 );
12015 }
12016 let t_cap = t_ent.elapsed();
12017 // PERSISTENT DRAFT KV: fill the MTP block's K/V for every prompt position from the exact
12018 // trunk hiddens collected during prime — ONE batched K/V-only pass (overwrites any
12019 // capture-warmup garbage; capture left len at 0). last_token (the init feed) needs no
12020 // fill: the first chain step processes it and appends its entry at slot prompt.len().
12021 if let Some(ph) = &prompt_h {
12022 // SESSION: rows [0..base) are the previous turns' exact fills (refresh overwrote them
12023 // with true verify hiddens) — truncate any draft overhang, fill ONLY the suffix at
12024 // global positions [base..base+tp). Fresh call: base==0, identical to before.
12025 scratch.set_len(e, base)?;
12026 // CHUNKED FILL (long-ctx OOM fix, 2026-07-05): mtp_kv_fill's transients scale with its
12027 // T (concat = T*2*n_embd*4B — 1.5GB at 40k) and its concat loop is 2*T launches. The
12028 // fill is a pure sequential append, so chunking is exact: each chunk appends its rows
12029 // at pos0=base+start with the identical per-row math. Same knob as the trunk prime.
12030 let tp = prompt.len();
12031 let fill_chunk: usize = if crate::cache::swa_ring_on() {
12032 crate::hybrid_forward::prime_chunk_tokens(tp, self.layers.len())
12033 } else {
12034 // Preserve the flag-OFF schedule byte-for-byte, including the legacy zero value
12035 // meaning one monolithic fill.
12036 std::env::var("MEMRA_PRIME_CHUNK")
12037 .ok()
12038 .and_then(|v| v.parse().ok())
12039 .unwrap_or(4096)
12040 };
12041 let fill_chunk = if fill_chunk == 0 { tp } else { fill_chunk };
12042 let mut start = 0usize;
12043 while start < tp {
12044 let end = (start + fill_chunk).min(tp);
12045 let tc = end - start;
12046 {
12047 // PREDECESSOR pairing: row i gets h[i-1]; global row 0 a zeros row (the
12048 // reference engine's initial pending-h is zeroed too); a session turn's row 0
12049 // gets the PREVIOUS turn's last committed hidden (sess.last_h). Per chunk:
12050 // rows start..end read h[start-1..end-1] — one dtod into a chunk buffer.
12051 let mut phs = e.zeros(tc * n_embd)?;
12052 let (src_lo, dst_off) = if start == 0 {
12053 (0, n_embd)
12054 } else {
12055 ((start - 1) * n_embd, 0)
12056 };
12057 let n_copy = if start == 0 {
12058 (tc - 1) * n_embd
12059 } else {
12060 tc * n_embd
12061 };
12062 if start == 0 {
12063 if let Some((_, lh, _, _, _)) = sess_tail.as_ref() {
12064 if let Some(lh) = lh.as_ref() {
12065 e.copy_into(&mut phs, 0, lh, n_embd)?;
12066 }
12067 }
12068 }
12069 if n_copy > 0 {
12070 e.copy_view_into(
12071 &mut phs,
12072 dst_off,
12073 &ph.slice(src_lo..src_lo + n_copy),
12074 n_copy,
12075 )?;
12076 }
12077 self.mtp_kv_fill_all(
12078 e,
12079 &prompt[start..end],
12080 &phs,
12081 base + start,
12082 &mut *scratch,
12083 embd_dev,
12084 )?;
12085 }
12086 start = end;
12087 }
12088 }
12089 // MEMRA_PROFILE_SPEC=2: profiler capture starts HERE — after the prime, so an
12090 // `nsys -c cudaProfilerApi` capture contains ONLY the round loop (draft/verify/commit).
12091 // (=1 brackets the whole call in run_spec.rs, prime included.)
12092 if std::env::var("MEMRA_PROFILE_SPEC").as_deref() == Ok("2") {
12093 unsafe extern "C" {
12094 fn cudaProfilerStart() -> i32;
12095 }
12096 unsafe {
12097 cudaProfilerStart();
12098 }
12099 }
12100 // ROUND-STREAM stage (c) 4 (MEMRA_SPEC_STREAM=1, experimental): pre-issued M-round
12101 // bursts with ZERO per-round host readbacks — the accept/seed/rollback/ring kernels
12102 // consume each other's device outputs; the host drains the ring every M rounds. v1
12103 // constraints: greedy, !spec_replay, single-shot, batched-linear layers, no refresh
12104 // fills (acceptance effect A/B-arbitrated), enters from round 1 (pending guaranteed).
12105 // NOTE: not gated on the caller's graph_draft (its trunk_dense conjunct turns the 35B
12106 // MoE off) — the stream capture encloses ONLY the dense MTP head; the head-dense /
12107 // full-prec / k gates are re-derived here and a failed capture degrades to stream-off.
12108 let stream_on = crate::spec::spec_stream()
12109 && !sampled
12110 && !spec_replay
12111 && self.mtp_extra.is_empty()
12112 && constraint.is_none()
12113 && !session_mode
12114 && embd_gpu.is_some()
12115 && !crate::model::full_prec_enabled()
12116 && k + 2 < 96;
12117 let mut stream_graph: Option<cudarc::driver::CudaGraph> = None;
12118 let mut g_tokp2k = e.alloc_u32_zeroed(2 * k.max(1))?;
12119 if stream_on {
12120 let cap = e.capture_graph(|e| {
12121 for j in 0..k.max(1) {
12122 self.mtp_head_forward_cap(
12123 e,
12124 mtp,
12125 &mut dctx.g_tok,
12126 &mut dctx.g_pos,
12127 &mut dctx.g_seed,
12128 &mut dctx.g_p,
12129 &mut *scratch,
12130 true,
12131 true,
12132 embd_gpu.expect("round stream requires resident embedding"),
12133 embd_qt,
12134 embd_rb,
12135 d_vocab,
12136 None,
12137 Some((&mut g_tokp2k, j, d2t_dev.as_ref())),
12138 None, // round-stream requires constraint.is_none() (see stream_on)
12139 )?;
12140 }
12141 Ok(())
12142 });
12143 match cap {
12144 Ok(g) => {
12145 scratch.set_len(e, 0)?;
12146 stream_graph = Some(g);
12147 }
12148 Err(err) => {
12149 scratch.set_len(e, 0)?;
12150 if debug_spec {
12151 eprintln!("[spec] stream-graph capture failed ({err}); stream off");
12152 }
12153 }
12154 }
12155 }
12156 let stream_active = stream_on && stream_graph.is_some();
12157 if debug_spec {
12158 eprintln!(
12159 "[spec] stream_on={stream_on} env={} samp={sampled} dg={} captured={} active={stream_active} session={session_mode} replay={spec_replay}",
12160 crate::spec::spec_stream(),
12161 dctx.graph.is_some(),
12162 stream_graph.is_some()
12163 );
12164 }
12165 let t_v_s = k + 1;
12166 // ROUND-STREAM buffers + ptr tables now live in the model-generic round_stream
12167 // module (extracted 2026-07-12; the gemma burst reuses them).
12168 let sb = crate::round_stream::StreamBufs::new(e, k, crate::spec::spec_stream_m())?;
12169 let crate::round_stream::StreamBufs {
12170 mut vtok_d,
12171 mut brk_d,
12172 mut pend_d,
12173 last_pred_d,
12174 mut pos_ctr,
12175 mut pos_start_d,
12176 mut ring_d,
12177 acc_d: mut stream_acc,
12178 m_rounds,
12179 k: _,
12180 } = sb;
12181 let stream_ptrs: Option<CudaSlice<u64>> = if stream_active {
12182 Some(crate::round_stream::kv_len_ptr_table(
12183 e,
12184 cache,
12185 Some(&pos_ctr),
12186 )?)
12187 } else {
12188 None
12189 };
12190
12191 let t_fill = t_ent.elapsed();
12192 let mut round = 0usize;
12193 // ADAPTIVE DRAFT LENGTH (MEMRA_SPEC_ADAPT=1, opt-in — the gemma_spec accepted-run law,
12194 // ported 2026-08-01): next round's draft depth = last round's accepted run + 1, clamped
12195 // to [floor(pos), k_cap] — a miss shrinks the next draft to the miss point + 1,
12196 // full-accept streaks re-deepen one step per round. NOT the 2026-07-07 acceptance-EMA
12197 // (that arm measured an HONEST LOSS to static per-class optima — 115.0/85.8/73.4 vs
12198 // 121.6/92.7/75.6, EMA lag — and was removed 2026-07-08; rig5090.jsonl has the record).
12199 // The gemma law has no lag class: it reacts within one round, and was worth +7-20% on
12200 // the gemma cells at unchanged exactness (2026-07-10 flip; floor sweep 2026-07-25;
12201 // position key 2026-07-26). Signal = n_acc from the round's EXISTING accept readback —
12202 // zero new syncs; the draft graph is a SINGLE-STEP capture replayed per drafted token,
12203 // so a per-round depth needs no re-capture (unlike gemma's whole-chain graphs). qwen's
12204 // in-round p-min cut already shortens chains mid-round, so gemma's one-round-late p-min
12205 // fold into kc is unnecessary here — the accepted-run law sees the cut via n_acc.
12206 // Exactness is the verify's job at ANY depth (same contract as p-min variable rounds).
12207 // DEFAULT OFF on the qwen path until its cells gate a flip (gemma's is default-on).
12208 // MEASURED 2026-08-01 (H100 GPU-3, interleaved x3, NGEN=256, same-invocation plain
12209 // denominators; research/qwen-adaptive-k-20260801/): REFUTED on the tuned qwen configs.
12210 // q27 K=3+HPOST+PMIN=0.3: short +0.8% (noise; law ~idles, len_hist identical), board
12211 // -1.9%, agentic -0.5%; board PMIN=0 -2.1% (not p-min shadowing — the law itself);
12212 // floor=1 -2.8% (gemma's floor-collapse, reproduced). q35 K=2 board: -6.4% (52/136
12213 // rounds shrink to depth 1; no depth to reclaim at K=2). The gemma direction DOES
12214 // appear at untuned depth-K — q27 K=6 floor=4 +1.5% over fixed K=6 — but stays -3.7%
12215 // below fixed K=3: same verdict class as the retired EMA arm (honest loss to static
12216 // per-class optima). Acceptance-rate rises under the law while tokens/round falls —
12217 // it buys accept-% by adding rounds, and a round's fixed draft+verify cost wins.
12218 // K=1..8 self-consistency PASS both models with the law ON (exactness held).
12219 let adapt = std::env::var("MEMRA_SPEC_ADAPT").as_deref() == Ok("1");
12220 // floor: per-model default keyed on n_embd (gemma's tiering — models with an expensive
12221 // verify keep deep drafts after a miss); MEMRA_SPEC_ADAPT_FLOOR pins it everywhere.
12222 let adapt_floor_env: Option<usize> = std::env::var("MEMRA_SPEC_ADAPT_FLOOR")
12223 .ok()
12224 .and_then(|v| v.parse().ok());
12225 let adapt_floor_default: usize = if self.cfg.n_embd as usize >= 3500 {
12226 4
12227 } else if self.cfg.n_embd as usize >= 2500 {
12228 2
12229 } else {
12230 1
12231 };
12232 let adapt_floor: usize = adapt_floor_env.unwrap_or(adapt_floor_default);
12233 // position key: past floor_ctx a HIGH floor (>=4) relaxes to 1 — forced-deep drafts
12234 // turn net-negative at depth (gemma 31B d1736 evidence); MEMRA_SPEC_FLOOR_CTX moves
12235 // the boundary, an explicit MEMRA_SPEC_ADAPT_FLOOR pins the floor everywhere.
12236 let floor_ctx: usize = std::env::var("MEMRA_SPEC_FLOOR_CTX")
12237 .ok()
12238 .and_then(|v| v.parse().ok())
12239 .unwrap_or(1024);
12240 let floor_at = |pos: usize| -> usize {
12241 if adapt_floor_env.is_some() || pos < floor_ctx {
12242 adapt_floor
12243 } else if adapt_floor >= 4 {
12244 1
12245 } else {
12246 adapt_floor
12247 }
12248 };
12249 // cap: MEMRA_SPEC_CAPMAX (gemma semantics, default 7). Binds only under adapt — the
12250 // fixed-K default path is untouched by this whole block.
12251 let cap_max: usize = std::env::var("MEMRA_SPEC_CAPMAX")
12252 .ok()
12253 .and_then(|v| v.parse().ok())
12254 .unwrap_or(7);
12255 let k_cap = k.min(cap_max).max(1);
12256 let mut kc = k_cap;
12257 let mut opti_fork: Option<OptiForkState> = None;
12258 let mut fork_snapshot: Option<crate::cache::CacheSnapshot> = None;
12259 if fork_mode != OptiForkGateMode::Disabled {
12260 let fence = crate::pp::pp_cuts(self.layers.len());
12261 let refusal = if !session_mode {
12262 Some("not-session")
12263 } else if k != 1 || adapt {
12264 Some("requires-fixed-k1")
12265 } else if sampled || constraint.is_some() || spec_replay {
12266 Some("sampled-constrained-or-replay")
12267 } else if pipe.is_some() {
12268 Some("two-session-pipeline")
12269 } else if !spec_devacc() {
12270 Some("requires-device-accept")
12271 } else if stream_active || crate::spec::spec_stream() {
12272 Some("round-stream")
12273 } else if !self.mtp_extra.is_empty() {
12274 Some("multi-head-mtp")
12275 } else if crate::cache::swa_ring_on() || cache.has_swa_ring() {
12276 Some("swa-ring")
12277 } else if crate::pp::pp_host_bounce_active() {
12278 Some("host-bounce")
12279 } else if fork_mode == OptiForkGateMode::Controller
12280 && cache.recur.iter().any(Option::is_some)
12281 {
12282 Some("controller-requires-zero-recurrent-state")
12283 } else if fence.as_ref().is_none_or(|f| f.len() != 3) {
12284 Some("requires-pp2")
12285 } else {
12286 None
12287 };
12288 if let Some(reason) = refusal {
12289 OPTI_FORK_REFUSALS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
12290 eprintln!("[opti-fork] refused reason={reason}");
12291 } else {
12292 let fence = fence.expect("validated PP-2 fence");
12293 let rt = crate::pp::PpNRt::get(e)?;
12294 let primary_stage0 = rt.engine(0, e).ctx().ordinal() == e.ctx().ordinal();
12295 let primary_stage1 = rt.engine(1, e).ctx().ordinal() == e.ctx().ordinal();
12296 let primary_supported =
12297 primary_stage0 || (fork_mode == OptiForkGateMode::Controller && primary_stage1);
12298 if !rt.cross_device() || !primary_supported {
12299 OPTI_FORK_REFUSALS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
12300 eprintln!("[opti-fork] refused reason=requires-supported-primary-cross-device");
12301 } else {
12302 // Both recurrent snapshots and both seed generations are allocated before
12303 // the first fork, each through its owning PP stage. Allocation failure
12304 // therefore happens before any optimistic state mutation can occur.
12305 let current_snapshot = opti_snapshot_stage_owned(e, cache, rt, &fence)?;
12306 let alternate_snapshot = opti_snapshot_stage_owned(e, cache, rt, &fence)?;
12307 let fork = OptiForkState::new(
12308 e,
12309 cache,
12310 fork_mode,
12311 alternate_snapshot,
12312 &h_seed_buf,
12313 &fill_prev,
12314 rt,
12315 fence[1],
12316 self.layers.len(),
12317 )?;
12318 eprintln!(
12319 "[opti-fork] armed mode={fork_mode:?} snapshots=2 seeds=2 split={} \
12320 payload_dev0={} payload_dev1={} q_threshold={:.3}",
12321 fence[1],
12322 fork.logical_payload_bytes[0],
12323 fork.logical_payload_bytes[1],
12324 fork.controller.map_or(0.0, |policy| policy.threshold),
12325 );
12326 fork_snapshot = Some(current_snapshot);
12327 opti_fork = Some(fork);
12328 }
12329 }
12330 }
12331 // Persistent snapshot buffers are allocated once and refreshed in place. The fork arm
12332 // uses stage-owned snapshots; refused/disabled arms retain the existing generic helper.
12333 let mut snap = match fork_snapshot {
12334 Some(snapshot) => snapshot,
12335 None => cache.snapshot(e)?,
12336 };
12337 let mut carried_opti: Option<OptiControllerTicket> = None;
12338 // ROUND-STREAM stage (b) 3a: device table of per-layer kvl.len_d pointers (stable — the
12339 // cache never reallocates len_d; see cache.rs "stable pointer" note). 0 = no KV layer.
12340 let kv_len_ptrs: Option<CudaSlice<u64>> = if spec_devacc() && !spec_replay {
12341 Some(crate::round_stream::kv_len_ptr_table(e, cache, None)?)
12342 } else {
12343 None
12344 };
12345 // BONUS FOLD (2026-07-04): after a FULL accept the bonus token is NOT committed with a
12346 // separate T=1 trunk pass (a full weight read per round). It stays PENDING and rides as
12347 // column 0 of the NEXT round's verify batch. Under predecessor pairing the next chain
12348 // seeds from the bonus's predecessor's TRUE verify hidden (free — no extra
12349 // pass of any kind). Verify still
12350 // checks every emitted token against the target -> exactness holds by construction; only
12351 // DRAFT QUALITY can shift, which the acceptance numbers arbitrate.
12352 // bonus emitted but not yet committed to cache. A carried pending (see SpecSession::
12353 // pending_tok) enters round 0 directly — the burst boundary becomes a plain round edge.
12354 let mut pending: Option<u32> = carried_pending;
12355 // MEMRA_SPEC_PHASE=1: per-round wall decomposition (draft / verify / accept+commit) —
12356 // no tracing, no extra syncs (each phase is naturally sync-bounded: draft readbacks,
12357 // the verify accept readback). Printed once at loop end via spec-stats.
12358 let anatomy_on = std::env::var("MEMRA_SPEC_PP_ANATOMY").as_deref() == Ok("1");
12359 let phase_on = anatomy_on || std::env::var("MEMRA_SPEC_PHASE").as_deref() == Ok("1");
12360 // MEMRA_SPEC_PHASE_SYNC=1 — reads the phase split correctly, and proves it. `ph_mark` is a
12361 // bare Instant, so `verify-issue` is the host QUEUEING the walk (the GPU is already running
12362 // under it) and `verify-wait` is only the residual drain at the accept readback: one
12363 // overlapped interval cut at the first blocking call, NOT "GPU time" beside "host time".
12364 // Syncing right after the walk is issued moves the whole GPU wall into `verify-issue`. If
12365 // the walk's GPU total is really issue+wait, then with this on verify-issue jumps to that
12366 // sum, verify-wait collapses to the readback alone, and the ROUND WALL DOES NOT MOVE —
12367 // which is what says the queueing time was hidden and is not a target. Diagnostic only.
12368 let phase_sync = std::env::var("MEMRA_SPEC_PHASE_SYNC").as_deref() == Ok("1");
12369 // DRAFT-MASK receipt (lane/draft-mask): speculative-clone wall + rounds, printed with
12370 // spec-stats. The clone is the one cost the design adds per round — measured, not assumed.
12371 let (mut dm_clone_ns, mut dm_rounds) = (0u128, 0usize);
12372 // grammar-truncation counters: how many rounds the verify-side cut fired and how many
12373 // already-verified tokens it threw away. THIS is the quantity draft masking targets.
12374 let (mut dm_cuts, mut dm_cut_tokens) = (0usize, 0usize);
12375 let (mut ph_draft, mut ph_verify, mut ph_rest) = (0f64, 0f64, 0f64);
12376 let mut ph_wait = 0f64;
12377 let mut ph_commit = 0f64;
12378 let mut ph_t = std::time::Instant::now();
12379 let mut ph_mark = |acc: &mut f64, on: bool| {
12380 if on {
12381 let now = std::time::Instant::now();
12382 *acc += (now - ph_t).as_secs_f64();
12383 ph_t = now;
12384 }
12385 };
12386 // MTP-ROUTE VERIFY GRAPHS (`MEMRA_SPEC_VERIFY_GRAPH`, see the flag doc): the
12387 // model-owned capture pool, locked for the whole burst exactly as the dspark serve
12388 // arm holds it — the slab stash is live verify -> commit inside a round, and the
12389 // worker drives rounds from one scheduler thread. PERSISTENT across generations on
12390 // the model (rebuilding per call re-captures the pool per prompt, which is the
12391 // measured way to lose more than the launches cost); the captured bodies are
12392 // cache-independent, every state read going through per-round refreshed pointer
12393 // tables. None = the eager walk, byte-identical.
12394 //
12395 // Never armed together with ROUND-STREAM: the tparallel verify refuses that pair
12396 // loudly, and `stream_active` owns the burst arm above, so the door stays shut
12397 // whenever the stream is live rather than relying on that refusal.
12398 // The lock is taken ONLY when the door is armed: with the flag off this whole block
12399 // is inert, so the default path cannot serialize two spec generations behind a mutex
12400 // it never reads.
12401 let vg_armed =
12402 crate::spec::spec_verify_graph_env().unwrap_or_else(|| self.vgraph_family_default());
12403 let mut vg_guard = if vg_armed && !stream_active {
12404 let mut g = self.dspark_vgraphs.lock().unwrap();
12405 if g.is_none() {
12406 // Size by the WIDEST verify this run can present, which is k+1 and NOT
12407 // k_cap+1: the sampled arm's own window is `t_v_s = k + 1`, so a pool built
12408 // from a smaller adaptive cap gets sliced past its stash rows (a `slice_mut`
12409 // panic in the sampled ON arm, measured before this line said k+1).
12410 let vt_cap = (k.max(k_cap) + 1).max(2);
12411 *g = DsparkVerifyGraphs::new(e, cache, vt_cap, n_embd)?;
12412 if g.is_some() {
12413 // Engagement receipt (the dead-arm lesson): prove the door is LIVE rather
12414 // than trusting that a flag set means a pool built.
12415 eprintln!("[spec-vg] MTP verify-graph pool ENGAGED (vt_cap={vt_cap})");
12416 } else {
12417 eprintln!(
12418 "[spec-vg] MTP verify-graph pool declined (no linear layers, \
12419 non-uniform state, or vt_cap < 2) — eager walk"
12420 );
12421 }
12422 }
12423 Some(g)
12424 } else {
12425 None
12426 };
12427 // Capacity fail-safe: a round wider than the pool was built for must take the eager
12428 // walk, not slice the stash past its rows. The sizing above already covers every
12429 // round this run can present; this keeps a future caller (or a k that grows behind
12430 // the pool's back) on the byte-identical fallback instead of a panic.
12431 let vg_t_cap = vg_guard
12432 .as_ref()
12433 .and_then(|g| g.as_ref())
12434 .map(|g| g.t_capacity())
12435 .unwrap_or(0);
12436 if let Some(p) = pipe {
12437 p.setup_end();
12438 }
12439 while keep_going && out.len() < max_new {
12440 // MEMRA_SPEC_ROUND_PROF=1: wall of the WHOLE round against the pieces we already
12441 // instrument. Needed because the parts do not add up: the draft step measures 1.27 ms
12442 // ([spec-anatomy] glue 92 / attn 280 / ffn 222 / head 670 us) and the t=2 verify walk
12443 // 25.6 ms ([tcol-prof] attn 10.1 + ffn 15.3), yet a K=1 round takes 177 ms on the
12444 // step37 TP2 stack. This prints where the other ~150 ms lives.
12445 let round_prof = ROUND_PROF
12446 .get_or_init(|| std::env::var("MEMRA_SPEC_ROUND_PROF").as_deref() == Ok("1"));
12447 let round_t0 = round_prof.then(std::time::Instant::now);
12448 // ROUND-STREAM BURST: from round 1 (pending guaranteed by every non-replay arm),
12449 // issue M rounds with zero readbacks, then drain the ring + reconcile mirrors.
12450 if let (true, Some(sg), Some(ptrs)) = (
12451 stream_active && round >= 1 && pending.is_some(),
12452 &stream_graph,
12453 &stream_ptrs,
12454 ) {
12455 if debug_spec {
12456 static ONCE: std::sync::Once = std::sync::Once::new();
12457 ONCE.call_once(|| {
12458 eprintln!("[memra] ROUND-STREAM burst engaged (M={m_rounds} k={k})")
12459 });
12460 }
12461 e.set_i32_one(&mut pos_ctr, cache.pos as i32)?;
12462 e.set_u32_one(&mut pend_d, pending.unwrap())?;
12463 e.set_u32_one(&mut ring_d, 0)?; // ring count = 0 (writes element 0)
12464 for _mi in 0..m_rounds {
12465 e.i32_copy_add(&pos_ctr, &mut pos_start_d, 0)?;
12466 cache.snapshot_into(e, &mut snap)?; // device D2Ds, stream-ordered
12467 e.i32_copy_add(&pos_ctr, &mut scratch.kv.len_d, 0)?; // draft-KV rollback
12468 e.i32_copy_add(&pos_ctr, &mut dctx.g_pos, 1)?; // rope pos = pos + base
12469 e.u32_copy(&pend_d, &mut dctx.g_tok)?;
12470 e.copy_into(&mut dctx.g_seed, 0, &h_seed_buf, n_embd)?;
12471 sg.launch()?;
12472 e.spec_assemble_verify(
12473 &g_tokp2k,
12474 &pend_d,
12475 d2t_dev.as_ref(),
12476 &mut vtok_d,
12477 &mut brk_d,
12478 p_min,
12479 k,
12480 pmin0,
12481 )?;
12482 let mut ck = VerifyCkpt::new(self.layers.len());
12483 let dummy = vec![0u32; t_v_s];
12484 let (tl_d, vx) = self.decode_step_t_core_stream(
12485 e,
12486 &dummy,
12487 0,
12488 &mut *cache,
12489 embd_dev,
12490 Some(&mut ck),
12491 Some((&vtok_d, &pos_ctr)),
12492 None,
12493 None,
12494 None,
12495 )?;
12496 for j in 0..t_v_s {
12497 e.argmax_token_device_col(&tl_d, j, n_vocab, &mut preds_d, j)?;
12498 }
12499 e.spec_accept_greedy_dc(
12500 &preds_d,
12501 &vtok_d,
12502 &last_pred_d,
12503 &brk_d,
12504 &mut stream_acc,
12505 )?;
12506 e.spec_seed_gather(&vx, &fill_prev, &stream_acc, &mut h_seed_buf, 1, n_embd)?;
12507 e.copy_into(&mut fill_prev, 0, &h_seed_buf, n_embd)?;
12508 self.commit_verified_prefix_stream(
12509 e,
12510 &mut *cache,
12511 &snap,
12512 &ck,
12513 &stream_acc,
12514 1,
12515 t_v_s,
12516 )?;
12517 e.spec_rollback_stream(
12518 ptrs,
12519 &pos_start_d,
12520 &stream_acc,
12521 1,
12522 self.layers.len() + 1,
12523 )?;
12524 e.spec_ring_commit(&vtok_d, &stream_acc, &brk_d, &mut ring_d, &mut pend_d)?;
12525 }
12526 e.stream().synchronize()?;
12527 let ring_h = e.dtoh_u32(&ring_d)?;
12528 let cnt = ring_h[0] as usize;
12529 for i in 0..cnt {
12530 if out.len() < max_new {
12531 out.push(ring_h[1 + i]);
12532 }
12533 }
12534 let pos_h = e.dtoh_i32(&pos_ctr)?[0] as usize;
12535 for il in 0..self.layers.len() {
12536 if let Some(kvl) = cache.kv[il].as_mut() {
12537 kvl.len = pos_h;
12538 }
12539 }
12540 cache.pos = pos_h;
12541 scratch.kv.len = pos_h;
12542 pending = Some(ring_h[cnt]); // last drained token = the live bonus
12543 last_token = ring_h[cnt];
12544 total_drafted += k * m_rounds; // upper bound (p-min breaks uncounted)
12545 total_accepted += cnt.saturating_sub(m_rounds);
12546 if let Some(t) = sess_telem {
12547 // totals only — the burst's per-round accept counts stayed on device
12548 // (that is the point of the round-stream arm). pos_* untouched.
12549 t.record_totals(m_rounds, k * m_rounds, cnt.saturating_sub(m_rounds));
12550 }
12551 round += m_rounds;
12552 // sse-cadence: the drained ring is committed — flush it at burst-drain cadence.
12553 keep_going = flush_commit(&mut on_commit, &out, &mut flushed);
12554 continue;
12555 }
12556 let pipe_draft = match pipe {
12557 Some(p) => Some(p.draft_begin(round)?),
12558 None => None,
12559 };
12560 let pos = cache.pos; // #tokens committed (EXCLUDES a pending bonus)
12561 let mut current_opti = carried_opti.take();
12562 let mut fork_generation = if current_opti.is_none() && pending.is_some() {
12563 match opti_fork.as_mut() {
12564 Some(fork) if fork.mode.is_forced() => Some(fork.reserve(&mut snap)?),
12565 None => None,
12566 Some(_) => None,
12567 }
12568 } else {
12569 None
12570 };
12571 if current_opti.is_none() {
12572 if let Some(fork) = opti_fork.as_ref() {
12573 opti_snapshot_stage_owned_into(e, cache, fork.rt, &fork.fence, &mut snap)?;
12574 } else {
12575 cache.snapshot_into(e, &mut snap)?;
12576 }
12577 } else if snap.pos != pos {
12578 return Err(format!(
12579 "optipipe carried snapshot pos {} != current pos {pos}",
12580 snap.pos
12581 )
12582 .into());
12583 } // §C: snapshot BEFORE draft+verify (already retained for a carried successor)
12584 ph_mark(&mut ph_rest, phase_on);
12585
12586 // --- 1. DRAFT k tokens with the NextN head (autoregressive, T=1 each) ---
12587 // p-min semantics (both paths): stop the chain early when the head's confidence in
12588 // its own pick drops below p_min — the just-drafted token is DISCARDED, but its
12589 // scratch append stands (identical to the eager chain's ordering). j==0 always drafts.
12590 let base0 = if pending.is_some() { 1usize } else { 0usize };
12591 // fixed draft length by default; MEMRA_SPEC_ADAPT=1 drafts at last round's
12592 // accepted run + 1 (the gemma law — see the setup block above the loop).
12593 let k_this = if adapt { kc } else { k };
12594 let mut draft: Vec<u32> = Vec::with_capacity(k);
12595 let mut draft_idx: Vec<u32> = Vec::with_capacity(k); // trimmed-vocab ids (== draft when untrimmed)
12596 let mut controller_draft_prob: Option<f32> = None;
12597 let mut controller_eager_state: Option<(u32, CudaSlice<f32>)> = None;
12598 if let Some(ticket) = current_opti.as_mut() {
12599 let carried_pending = pending.ok_or("optipipe carried successor lost pending")?;
12600 if ticket.verify_tokens[0] != carried_pending {
12601 return Err(format!(
12602 "optipipe carried pending mismatch: ticket={} live={carried_pending}",
12603 ticket.verify_tokens[0],
12604 )
12605 .into());
12606 }
12607 draft.push(ticket.verify_tokens[1]);
12608 controller_draft_prob = Some(ticket.draft_prob);
12609 controller_eager_state = ticket
12610 .take_eager_seed()
12611 .map(|seed| (ticket.verify_tokens[1], seed));
12612 } else {
12613 // Round-start draft-KV sync (BOTH paths). Persistent: truncate/align to the committed
12614 // history — slots 0..P hold entries for the tokens before last_token@P (P = pos +
12615 // base0 - 1); this single set_len IS the draft-side rollback (drops last round's
12616 // rejected drafts and p-min extras via the len mechanism).
12617 scratch.set_len(e, pos + base0 - 1)?;
12618 // dcw door: a captured chain appends k_this device-counter rows (plus the
12619 // pseudo-seed replay) with no host intervention; any ring rebase those appends
12620 // could need happens HERE, host-side, before the replays. The eager arm keeps
12621 // its own per-step prepare, so this is graph-path-only work.
12622 if step35_draft_dcw_on() && (dctx.graph.is_some() || dctx.graph_s.is_some()) {
12623 scratch.ensure_dcw_headroom(e, k_this + 2)?;
12624 }
12625 if pen_on {
12626 // PEN_WINDOW_MAX also bounds the per-round upload and the O(n_hist^2)
12627 // device dedup: the serve window is already PEN_WINDOW_MAX, and this
12628 // defensive min also bounds non-server callers.
12629 let win = sp.penalty_last_n.min(PEN_WINDOW_MAX);
12630 let w0 = pen_hist.len().saturating_sub(win);
12631 pen_hist_d = Some(e.htod_u32_v(&pen_hist[w0..])?);
12632 }
12633 if sampled {
12634 draft_logits.clear();
12635 draft_stats.clear();
12636 }
12637 // DRAFT-SIDE GRAMMAR MASK: clone the committed grammar state ONCE per round; each
12638 // position's mask is computed on that clone and advanced by the PROPOSED token. The
12639 // real state moves only on emission (verify's job), so the emitted stream is
12640 // unchanged — the mask only removes tokens the verify would have truncated anyway.
12641 let mut dmask_live = dmask_on;
12642 if dmask_live {
12643 let t_c = std::time::Instant::now();
12644 constraint
12645 .as_deref_mut()
12646 .unwrap()
12647 .draft_begin()
12648 .map_err(|e2| format!("constraint: {e2}"))?;
12649 dm_clone_ns += t_c.elapsed().as_nanos();
12650 dm_rounds += 1;
12651 }
12652 if let (false, Some(gr)) = (sampled || pen_on, &dctx.graph) {
12653 // GRAPH DRAFT: one dispatch per drafted token. The chain feeds itself on-device
12654 // (in-graph argmax -> tok_d -> next replay's embed; h_nextn -> h_seed_d; pos_d
12655 // inc'd in-graph); the host only reads 4B token (+4B p) and decides the break.
12656 e.set_i32_one(&mut dctx.g_pos, (pos + base0) as i32)?;
12657 e.set_u32_one(&mut dctx.g_tok, last_token)?;
12658 e.copy_into(&mut dctx.g_seed, 0, &h_seed_buf, n_embd)?;
12659 for j in 0..k_this {
12660 // per-position mask upload (contents only — the graph's baked pointer is
12661 // dctx.g_dmask). All-ones once masking goes dead mid-chain, so the captured
12662 // mask node degrades to a no-op ban instead of needing a second graph.
12663 if dmask_live
12664 && !upload_draft_mask(
12665 e,
12666 constraint.as_deref_mut().unwrap(),
12667 &mut dctx.g_dmask,
12668 mtp.d2t.as_ref(),
12669 d_vocab,
12670 dmask_words,
12671 )?
12672 {
12673 // no draft-vocab row is grammar-legal here (a trimmed FR-Spec head can
12674 // genuinely miss the legal set): neutralize the captured mask node and
12675 // finish the chain UNMASKED — exactly pre-lane behaviour, never worse.
12676 e.htod_u32_into(&mut dctx.g_dmask, &vec![u32::MAX; dmask_words])?;
12677 dmask_live = false;
12678 }
12679 gr.launch()?;
12680 scratch.kv.len += 1; // host mirror (len_d advanced in-graph)
12681 let idx = e.dtoh_u32_one(&dctx.g_tok)?;
12682 // #87 SENTINEL TRAP: an all-NaN head-logits row leaves the device argmax's
12683 // init sentinel (0x7FFFFFFF) in g_tok — feeding it onward dereferences
12684 // embed_row(sentinel) = table + ~4.6TB (never mapped) inside the NEXT graph
12685 // replay's embed node, and the MMU fault kills the CUDA context for the
12686 // whole process (research/pp2spec-crash-20260807: 3 coredumps, byte-exact
12687 // VA arithmetic). Refuse loudly instead; the diagnostics name the first-NaN
12688 // buffer (g_seed = the verify-side handoff vs head-side compute).
12689 if (idx as usize) >= d_vocab {
12690 // g_seed is SELF-FED (the replay writes h_nextn back into it), so it
12691 // reads as the head's OUTPUT at j; h_seed_buf is the round's INPUT
12692 // seed, untouched since the round-start copy — the pair discriminates
12693 // "seed arrived poisoned" from "head forward produced NaN".
12694 let seed_h = e.dtoh(&dctx.g_seed)?;
12695 let seed_nan = seed_h.iter().filter(|v| v.is_nan()).count();
12696 let in_h = e.dtoh(&h_seed_buf)?;
12697 let in_nan = in_h.iter().filter(|v| v.is_nan()).count();
12698 return Err(format!(
12699 "draft(graph) argmax sentinel 0x{idx:08x} >= d_vocab {d_vocab} at \
12700 round {round} j={j} pos={pos}: head-out NaN {seed_nan}/{n_embd}, \
12701 round-input-seed NaN {in_nan}/{n_embd} — refusing to dereference \
12702 the embed row (#87 trap)"
12703 )
12704 .into());
12705 }
12706 // trimmed draft vocab -> target token id (identity when no d2t map)
12707 let d = match &mtp.d2t {
12708 Some(map) => map[idx as usize],
12709 None => idx,
12710 };
12711 let draft_p = if p_min > 0.0
12712 || opti_fork
12713 .as_ref()
12714 .is_some_and(|fork| fork.controller.is_some())
12715 {
12716 Some(e.dtoh(&dctx.g_p)?[0])
12717 } else {
12718 None
12719 };
12720 if j == 0 {
12721 controller_draft_prob = draft_p;
12722 }
12723 if let Some(p) = draft_p.filter(|_| p_min > 0.0) {
12724 if p < p_min && (j > 0 || (pmin0 && base0 == 1)) {
12725 break;
12726 }
12727 }
12728 draft.push(d);
12729 // with a trimmed head the NEXT embed must read the TARGET id, not the draft
12730 // index the argmax wrote — patch the persistent token buffer (4B htod).
12731 if d != idx {
12732 e.set_u32_one(&mut dctx.g_tok, d)?;
12733 }
12734 // advance the SPECULATIVE state with the proposal; a dead chain drops to
12735 // unmasked drafting for the remaining positions (verify still arbitrates).
12736 // speculative advance; a chain the grammar can no longer follow (EOS
12737 // proposed) ends here. The captured mask node always runs, so a dead chain
12738 // leaves the buffer NEUTRAL (all-ones = ban nothing) before it exits.
12739 if dmask_live
12740 && !constraint
12741 .as_deref_mut()
12742 .unwrap()
12743 .draft_advance(d)
12744 .map_err(|e2| format!("constraint: {e2}"))?
12745 {
12746 e.htod_u32_into(&mut dctx.g_dmask, &vec![u32::MAX; dmask_words])?;
12747 break;
12748 }
12749 }
12750 // PURE-TEMP RE-TEST (lane/graph-s-key-exactness-20260819): the sampled graph is
12751 // legal ONLY in the regime it was captured in. The condition used to read
12752 // `(sampled, &dctx.graph_s)` and trusted `s_key` to have dropped anything else —
12753 // which it could not, because the key omitted the filters. Both halves are now
12754 // enforced: the key drops a stale graph, and this site refuses to launch one.
12755 } else if let (true, Some(gr)) = (sampled && pure_temp, &dctx.graph_s) {
12756 if skey_probe() {
12757 eprintln!(
12758 "[skey] chain=graph_s round={round} pure_temp={} top_k={} \
12759 top_p={} min_p={} s_key_parked={:?}",
12760 pure_temp as u8, sp.top_k, sp.top_p, sp.min_p, dctx.s_key,
12761 );
12762 }
12763 // SAMPLED GRAPH DRAFT: one replay per drafted token — head forward + gumbel +
12764 // argmax in ONE dispatch; the host reads 4B token (+4B p), D2Ds q into slot j,
12765 // and decides the break. Event-counter continuity: g_ctr is host-seeded to
12766 // sctr-1 ONCE per round (outside the graph); the in-graph bump runs BEFORE the
12767 // perturb, so replay j consumes counter sctr+j — exactly the eager arm's Philox
12768 // stream. Host sctr advances in lockstep (computed, no readback needed).
12769 e.set_i32_one(&mut dctx.g_pos, (pos + base0) as i32)?;
12770 e.set_u32_one(&mut dctx.g_tok, last_token)?;
12771 e.copy_into(&mut dctx.g_seed, 0, &h_seed_buf, n_embd)?;
12772 e.set_u32_one(&mut dctx.g_ctr, sctr.wrapping_sub(1))?;
12773 for j in 0..k_this {
12774 gr.launch()?;
12775 scratch.kv.len += 1; // host mirror (len_d advanced in-graph)
12776 sctr += 1; // mirrors the in-graph g_ctr bump (eager parity:
12777 // counts the p-min-discarded token too)
12778 // q retention: ONE async D2D of the persistent head-logits buffer into this
12779 // round's slot j (stream-ordered after the replay, before the next one).
12780 e.copy_into(&mut dctx.q_slots[j], 0, &dctx.g_q, d_vocab)?;
12781 let idx = e.dtoh_u32_one(&dctx.g_tok)?;
12782 // #87 SENTINEL TRAP (see the greedy graph arm above).
12783 if (idx as usize) >= d_vocab {
12784 let seed_h = e.dtoh(&dctx.g_seed)?;
12785 let seed_nan = seed_h.iter().filter(|v| v.is_nan()).count();
12786 return Err(format!(
12787 "draft(graph-sampled) argmax sentinel 0x{idx:08x} >= d_vocab \
12788 {d_vocab} at round {round} j={j} pos={pos}: round-seed NaN \
12789 {seed_nan}/{n_embd} — refusing to dereference the embed row \
12790 (#87 trap)"
12791 )
12792 .into());
12793 }
12794 let d = match &mtp.d2t {
12795 Some(map) => map[idx as usize],
12796 None => idx,
12797 };
12798 draft_idx.push(idx);
12799 if p_min > 0.0 {
12800 let p = e.dtoh(&dctx.g_p)?[0];
12801 if p < p_min && (j > 0 || (pmin0 && base0 == 1)) {
12802 break;
12803 }
12804 }
12805 draft.push(d);
12806 // trimmed head: the NEXT embed must read the TARGET id (see the greedy arm).
12807 if d != idx {
12808 e.set_u32_one(&mut dctx.g_tok, d)?;
12809 }
12810 }
12811 // uniform accept path: fill draft_stats per used slot (pure-temp regime — the
12812 // stats degenerate to th=0 / full-Z; one filter_stats launch per slot, tiny).
12813 for j in 0..draft.len().max(draft_idx.len()) {
12814 let rows0 = e.htod_i32(&[0])?;
12815 let (mut th_d, mut z_d, mut mx_d) = (e.zeros(1)?, e.zeros(1)?, e.zeros(1)?);
12816 e.filter_stats(
12817 &dctx.q_slots[j],
12818 d_vocab,
12819 &rows0,
12820 &mut th_d,
12821 &mut z_d,
12822 &mut mx_d,
12823 d_vocab,
12824 1,
12825 sp_temp,
12826 sp.top_k,
12827 sp.top_p,
12828 sp.min_p,
12829 )?;
12830 draft_stats.push((e.dtoh(&mx_d)?[0], e.dtoh(&th_d)?[0], e.dtoh(&z_d)?[0]));
12831 }
12832 } else {
12833 if skey_probe() && sampled {
12834 eprintln!(
12835 "[skey] chain=eager round={round} pure_temp={} top_k={} \
12836 top_p={} min_p={} s_key_parked={:?}",
12837 pure_temp as u8, sp.top_k, sp.top_p, sp.min_p, dctx.s_key,
12838 );
12839 }
12840 // EAGER DRAFT (fallback: MoE head/trunk, huge k, MEMRA_SPEC_NOGRAPH, capture fail).
12841 let chain_heads = !self.mtp_extra.is_empty();
12842 let mut e_tok = last_token;
12843 let mut d_seed = e.clone_dtod(&h_seed_buf)?;
12844 let mut chain_tokens = if chain_heads {
12845 vec![last_token]
12846 } else {
12847 Vec::new()
12848 };
12849 let mut chain_seeds = if chain_heads {
12850 vec![e.clone_dtod(&h_seed_buf)?]
12851 } else {
12852 Vec::new()
12853 };
12854 for j in 0..k_this {
12855 // GPU-ARGMAX DRAFT (2026-07-03): device logits + device argmax + 4-byte token
12856 // read instead of the ~600KB full-vocab dtoh + host argmax per draft token.
12857 let mtp_pos = pos + base0 + j;
12858 // draft-side grammar mask (eager twin of the graph arm's in-graph node).
12859 // A position with no legal draft-vocab row drops to unmasked drafting for
12860 // the rest of the chain (pre-lane behaviour; verify still arbitrates).
12861 if dmask_live {
12862 dmask_live = upload_draft_mask(
12863 e,
12864 constraint.as_deref_mut().unwrap(),
12865 &mut dctx.g_dmask,
12866 mtp.d2t.as_ref(),
12867 d_vocab,
12868 dmask_words,
12869 )?;
12870 }
12871 let mask = if dmask_live {
12872 Some((&dctx.g_dmask, dmask_words))
12873 } else {
12874 None
12875 };
12876 let (dl_d, h_nextn) = if chain_heads {
12877 if debug_spec {
12878 eprintln!(
12879 "[mtp-chain-step] round={round} j={j} head={} replay_rows={}",
12880 mtp_chain_head_index(j, self.mtp_head_count()),
12881 chain_tokens.len(),
12882 );
12883 }
12884 self.mtp_chain_forward_dev(
12885 e,
12886 &chain_tokens,
12887 &chain_seeds,
12888 &mut *scratch,
12889 pos + base0 - 1,
12890 embd_dev,
12891 mask,
12892 )?
12893 } else {
12894 self.mtp_head_forward_dev(
12895 e,
12896 mtp,
12897 e_tok,
12898 &d_seed,
12899 &mut *scratch,
12900 mtp_pos,
12901 embd_dev,
12902 mask,
12903 )?
12904 };
12905 let tok_d = if sampled {
12906 // FILTERED Gumbel-max: stats -> masked perturb -> argmax = one draw from
12907 // the filtered softmax (filters off => th=0, exact v1 semantics).
12908 if perturb_buf.is_none() {
12909 perturb_buf = Some(e.zeros(d_vocab.max(n_vocab))?);
12910 }
12911 let mut q_row = e.clone_dtod(&dl_d)?; // retained q (penalized when on)
12912 if pen_on {
12913 let h = pen_hist_d.as_ref().unwrap();
12914 let nh = h.len();
12915 e.penalize_logits(
12916 &mut q_row,
12917 h,
12918 nh,
12919 sp.penalty_repeat,
12920 sp.penalty_freq,
12921 sp.penalty_present,
12922 d_vocab,
12923 )?;
12924 }
12925 let rows0 = e.htod_i32(&[0])?;
12926 let (mut th_d, mut z_d, mut mx_d) =
12927 (e.zeros(1)?, e.zeros(1)?, e.zeros(1)?);
12928 e.filter_stats(
12929 &q_row, d_vocab, &rows0, &mut th_d, &mut z_d, &mut mx_d, d_vocab,
12930 1, sp_temp, sp.top_k, sp.top_p, sp.min_p,
12931 )?;
12932 let (th, z, mx) =
12933 (e.dtoh(&th_d)?[0], e.dtoh(&z_d)?[0], e.dtoh(&mx_d)?[0]);
12934 let pb = perturb_buf.as_mut().unwrap();
12935 e.gumbel_perturb_filtered(
12936 &q_row, pb, d_vocab, sp_seed, sctr, sp_temp, mx, th,
12937 )?;
12938 sctr += 1;
12939 draft_logits.push(q_row);
12940 draft_stats.push((mx, th, z));
12941 e.argmax_token_device(pb, d_vocab)?
12942 } else {
12943 e.argmax_token_device(&dl_d, d_vocab)?
12944 };
12945 let idx = e.dtoh_u32_one(&tok_d)?;
12946 // #87 SENTINEL TRAP (eager twin — see the graph arm). Extra diagnostics
12947 // here because the eager chain's operands are all readable: dl_d (the head
12948 // logits row) and d_seed (this step's h_seed) name the first-NaN buffer.
12949 if (idx as usize) >= d_vocab {
12950 let dl_h = e.dtoh(&dl_d)?;
12951 let dl_nan = dl_h.iter().filter(|v| v.is_nan()).count();
12952 let seed_h = if chain_heads {
12953 e.dtoh(chain_seeds.last().unwrap())?
12954 } else {
12955 e.dtoh(&d_seed)?
12956 };
12957 let seed_nan = seed_h.iter().filter(|v| v.is_nan()).count();
12958 return Err(format!(
12959 "draft(eager) argmax sentinel 0x{idx:08x} >= d_vocab {d_vocab} at \
12960 round {round} j={j} pos={pos}: head-logits NaN {dl_nan}/{d_vocab}, \
12961 step-seed NaN {seed_nan}/{n_embd} — refusing to dereference the \
12962 embed row (#87 trap)"
12963 )
12964 .into());
12965 }
12966 let d = match &mtp.d2t {
12967 Some(map) => map[idx as usize],
12968 None => idx,
12969 };
12970 if sampled {
12971 draft_idx.push(idx);
12972 }
12973 let draft_p = if p_min > 0.0
12974 || opti_fork
12975 .as_ref()
12976 .is_some_and(|fork| fork.controller.is_some())
12977 {
12978 let p_d = e.prob_of_token_device(&dl_d, &tok_d, d_vocab)?;
12979 Some(e.dtoh(&p_d)?[0])
12980 } else {
12981 None
12982 };
12983 if j == 0 {
12984 controller_draft_prob = draft_p;
12985 }
12986 if let Some(p) = draft_p.filter(|_| p_min > 0.0) {
12987 if p < p_min && (j > 0 || (pmin0 && base0 == 1)) {
12988 break;
12989 }
12990 }
12991 draft.push(d);
12992 if chain_heads {
12993 chain_tokens.push(d);
12994 chain_seeds.push(h_nextn);
12995 } else {
12996 e_tok = d;
12997 d_seed = h_nextn;
12998 }
12999 // speculative advance; a chain the grammar can no longer follow (EOS
13000 // proposed) ends here — the prefix already proposed still rides verify.
13001 if dmask_live
13002 && !constraint
13003 .as_deref_mut()
13004 .unwrap()
13005 .draft_advance(d)
13006 .map_err(|e2| format!("constraint: {e2}"))?
13007 {
13008 break;
13009 }
13010 }
13011 if !chain_heads
13012 && opti_fork
13013 .as_ref()
13014 .is_some_and(|fork| fork.controller.is_some())
13015 {
13016 controller_eager_state = Some((e_tok, d_seed));
13017 }
13018 }
13019 }
13020 let k_round = draft.len();
13021 if let Some(p) = pipe {
13022 p.draft_end(round);
13023 }
13024 drop(pipe_draft);
13025
13026 ph_mark(&mut ph_draft, phase_on);
13027 // --- 2. VERIFY: one batched target forward. With a pending bonus, it rides as col 0
13028 // (committing its KV/recur inside the SAME weight read); drafts follow. ---
13029 let verify_tokens: Vec<u32> = match pending {
13030 Some(b) => {
13031 let mut v = Vec::with_capacity(k_round + 1);
13032 v.push(b);
13033 v.extend_from_slice(&draft);
13034 v
13035 }
13036 None => draft.clone(),
13037 };
13038 let base = if pending.is_some() { 1 } else { 0 };
13039 // ckpt (REPLAY-FREE partial accept): retain per-layer state-rebuild inputs alongside
13040 // the verify. Pure buffer keep-alives + dtod clones — kernel work is unchanged.
13041 let mut ckpt = if let Some(ticket) = current_opti.as_mut() {
13042 Some(ticket.take_ckpt())
13043 } else if spec_replay {
13044 None
13045 } else {
13046 Some(VerifyCkpt::new(self.layers.len()))
13047 };
13048 let controller_can_probe = base == 1
13049 && k_round == 1
13050 && out.len().saturating_add(2) < max_new
13051 && controller_draft_prob.is_some()
13052 && opti_fork
13053 .as_ref()
13054 .and_then(|fork| fork.controller.as_ref())
13055 .is_some_and(|policy| !policy.breaker_tripped);
13056 let mut successor_attempt: Option<OptiControllerTicket> = None;
13057 let mut rejected_probe: Option<(f32, u32)> = None;
13058 let mut controller_prepared: Option<OptiControllerPrepared> = None;
13059 if controller_can_probe {
13060 // Prepare d2/q and, on admission, d3 before either current verify half is
13061 // issued. N stage 0 can then be followed immediately by N+1 stage 0; once N's
13062 // boundary fires, those dev0 launches overlap N stage 1 on dev1. Preparing on
13063 // the primary stream after N stage 1 would serialize the supposed pipeline.
13064 let eager_pos = scratch.kv.len + 1;
13065 let (optimistic_pending, pending_probability) = self.opti_controller_draft_step(
13066 e,
13067 mtp,
13068 &mut dctx,
13069 &mut *scratch,
13070 d_vocab,
13071 &mut controller_eager_state,
13072 eager_pos,
13073 embd_dev,
13074 )?;
13075 let first_probability = controller_draft_prob
13076 .ok_or("optipipe controller probe lost first-token probability")?;
13077 let q_proxy = first_probability * pending_probability;
13078 OPTI_GATE_CHECKS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
13079 OPTI_SHADOW_DRAFT_TOKENS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
13080 let admitted = opti_fork
13081 .as_ref()
13082 .and_then(|fork| fork.controller.as_ref())
13083 .ok_or("optipipe controller policy disappeared")?
13084 .admit(q_proxy);
13085 if admitted {
13086 OPTI_GATE_ADMITS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
13087 let eager_pos = scratch.kv.len + 1;
13088 let (optimistic_draft, optimistic_draft_probability) = self
13089 .opti_controller_draft_step(
13090 e,
13091 mtp,
13092 &mut dctx,
13093 &mut *scratch,
13094 d_vocab,
13095 &mut controller_eager_state,
13096 eager_pos,
13097 embd_dev,
13098 )?;
13099 OPTI_SHADOW_DRAFT_TOKENS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
13100 let eager_seed = controller_eager_state.take().map(|(token, seed)| {
13101 debug_assert_eq!(token, optimistic_draft);
13102 seed
13103 });
13104 controller_prepared = Some(OptiControllerPrepared {
13105 verify_tokens: [optimistic_pending, optimistic_draft],
13106 draft_prob: optimistic_draft_probability,
13107 eager_seed,
13108 q_proxy,
13109 scratch_len: scratch.kv.len,
13110 });
13111 } else {
13112 OPTI_GATE_REJECTS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
13113 OPTI_WASTED_DRAFT_TOKENS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
13114 rejected_probe = Some((q_proxy, optimistic_pending));
13115 eprintln!(
13116 "[opti-controller] reject q={q_proxy:.6} threshold={:.3}",
13117 opti_fork
13118 .as_ref()
13119 .and_then(|fork| fork.controller.as_ref())
13120 .expect("controller policy")
13121 .threshold,
13122 );
13123 }
13124 }
13125 let fork_attempt = match fork_generation.take() {
13126 Some(generation) if base == 1 && k_round == 1 => Some(generation),
13127 Some(generation) => {
13128 opti_fork
13129 .as_mut()
13130 .expect("fork generation without fork state")
13131 .retire(generation)?;
13132 None
13133 }
13134 None => None,
13135 };
13136 let (tlogits_d, vx) = if let Some(p) = pipe {
13137 self.decode_step_t_core_pipelined(
13138 e,
13139 &verify_tokens,
13140 pos,
13141 &mut *cache,
13142 embd_dev,
13143 ckpt.as_mut(),
13144 p,
13145 round,
13146 )?
13147 } else if controller_can_probe {
13148 let fence = opti_fork
13149 .as_ref()
13150 .ok_or("optipipe controller probe lost fork state")?
13151 .fence;
13152 let boundary = match current_opti.as_mut() {
13153 Some(ticket) => ticket.take_boundary(),
13154 None => self.verify_stage0_issue(
13155 e,
13156 &verify_tokens,
13157 pos,
13158 &mut *cache,
13159 embd_dev,
13160 ckpt.as_mut(),
13161 None,
13162 &fence,
13163 Some(true),
13164 None,
13165 )?,
13166 };
13167 if let Some(prepared) = controller_prepared.take() {
13168 let generation = {
13169 let fork = opti_fork
13170 .as_mut()
13171 .ok_or("optipipe controller admission lost fork state")?;
13172 let generation = fork.reserve_successor()?;
13173 let rt = fork.rt;
13174 let snapshot_fence = fork.fence;
13175 opti_snapshot_one_stage_owned_into(
13176 e,
13177 cache,
13178 rt,
13179 &snapshot_fence,
13180 0,
13181 fork.successor_snapshot_mut(),
13182 )?;
13183 generation
13184 };
13185 let mut successor_ckpt = VerifyCkpt::new(self.layers.len());
13186 let successor_boundary = self.verify_stage0_issue(
13187 e,
13188 &prepared.verify_tokens,
13189 pos + verify_tokens.len(),
13190 &mut *cache,
13191 embd_dev,
13192 Some(&mut successor_ckpt),
13193 None,
13194 &fence,
13195 Some(false),
13196 None,
13197 )?;
13198 OPTI_FORK_ATTEMPTS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
13199 let fork = opti_fork
13200 .as_ref()
13201 .ok_or("optipipe controller ticket lost fork state")?;
13202 successor_attempt = Some(fork.controller_ticket(
13203 generation,
13204 successor_boundary,
13205 successor_ckpt,
13206 prepared.verify_tokens,
13207 prepared.draft_prob,
13208 prepared.eager_seed,
13209 prepared.q_proxy,
13210 prepared.scratch_len,
13211 ));
13212 eprintln!(
13213 "[opti-controller] issue generation={} q={:.6} threshold={:.3} \
13214 verify={:?}",
13215 generation.id,
13216 prepared.q_proxy,
13217 fork.controller.expect("controller policy").threshold,
13218 prepared.verify_tokens,
13219 );
13220 }
13221 let result = self.verify_stage1_finish(
13222 e,
13223 boundary,
13224 &mut *cache,
13225 ckpt.as_mut(),
13226 None,
13227 &fence,
13228 successor_attempt.is_none(),
13229 )?;
13230 if let Some(ticket) = current_opti.as_mut() {
13231 ticket.settle();
13232 }
13233 if successor_attempt.is_some() {
13234 let fork = opti_fork
13235 .as_mut()
13236 .ok_or("optipipe successor snapshot lost fork state")?;
13237 let rt = fork.rt;
13238 let snapshot_fence = fork.fence;
13239 opti_snapshot_one_stage_owned_into(
13240 e,
13241 cache,
13242 rt,
13243 &snapshot_fence,
13244 1,
13245 fork.successor_snapshot_mut(),
13246 )?;
13247 // Publish N only after both independent successor-state queues are complete.
13248 fork.rt.publish_to(1, &e.stream())?;
13249 }
13250 result
13251 } else if let Some(ticket) = current_opti.as_mut() {
13252 let fork = opti_fork
13253 .as_mut()
13254 .ok_or("optipipe carried controller ticket lost fork state")?;
13255 let boundary = ticket.take_boundary();
13256 let result = self.verify_stage1_finish(
13257 e,
13258 boundary,
13259 &mut *cache,
13260 ckpt.as_mut(),
13261 None,
13262 &fork.fence,
13263 true,
13264 )?;
13265 ticket.settle();
13266 result
13267 } else if let Some(generation) = fork_attempt {
13268 let fork = opti_fork
13269 .as_mut()
13270 .expect("fork generation without fork state");
13271 fork.capture_seed(e, generation, &h_seed_buf, &fill_prev, scratch.kv.len)?;
13272 let action = fork.mode.action(generation.id);
13273 let boundary = self.verify_stage0_issue(
13274 e,
13275 &verify_tokens,
13276 pos,
13277 &mut *cache,
13278 embd_dev,
13279 ckpt.as_mut(),
13280 None,
13281 &fork.fence,
13282 Some(true),
13283 None,
13284 )?;
13285 OPTI_FORK_ATTEMPTS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
13286 let mut ticket = fork.ticket(generation, boundary);
13287 if action == OptiForkAction::Abort {
13288 return Err(format!(
13289 "optipipe forced abort with generation {} stage0 in flight",
13290 generation.id,
13291 )
13292 .into());
13293 }
13294 fork.reconcile(
13295 e,
13296 &mut *cache,
13297 &mut *scratch,
13298 &snap,
13299 &mut h_seed_buf,
13300 &mut fill_prev,
13301 generation,
13302 action,
13303 verify_tokens[0],
13304 )?;
13305 let result = if action == OptiForkAction::Hit {
13306 let boundary = ticket.take_boundary();
13307 self.verify_stage1_finish(
13308 e,
13309 boundary,
13310 &mut *cache,
13311 ckpt.as_mut(),
13312 None,
13313 &fork.fence,
13314 true,
13315 )?
13316 } else {
13317 // The optimistic boundary slot has no reader. Re-run the unchanged serial
13318 // verify only after E_restart published the restored stage-0 state.
13319 self.decode_step_t_core(
13320 e,
13321 &verify_tokens,
13322 pos,
13323 &mut *cache,
13324 embd_dev,
13325 ckpt.as_mut(),
13326 )?
13327 };
13328 ticket.settle();
13329 debug_assert_eq!(ticket.generation, generation);
13330 fork.retire(generation)?;
13331 result
13332 } else {
13333 // The serial verify every non-fork round takes — the MTP route's
13334 // verify-graph door. The pool is None unless MEMRA_SPEC_VERIFY_GRAPH armed
13335 // a pool above, and then the walk replays the captured trunk instead of
13336 // re-issuing it launch by launch.
13337 let vg_round = if verify_tokens.len() <= vg_t_cap {
13338 vg_guard.as_mut().and_then(|g| g.as_mut())
13339 } else {
13340 if let Some(g) = vg_guard.as_mut().and_then(|g| g.as_mut()) {
13341 // The commit reads this flag to pick its arm; a round that declines
13342 // the pool must not inherit a stale `true` from the round before it.
13343 g.round_slab = false;
13344 }
13345 None
13346 };
13347 self.decode_step_t_core_vg(
13348 e,
13349 &verify_tokens,
13350 pos,
13351 &mut *cache,
13352 embd_dev,
13353 ckpt.as_mut(),
13354 vg_round,
13355 )?
13356 };
13357 let pipe_accept = match pipe {
13358 Some(p) => Some(p.accept_begin(round)?),
13359 None => None,
13360 };
13361
13362 if phase_sync {
13363 e.stream().synchronize()?;
13364 }
13365 ph_mark(&mut ph_verify, phase_on);
13366 // --- 3. GREEDY ACCEPT (walk prefix, stop at first mismatch) ---
13367 // DEVICE-ARGMAX ACCEPT: argmax every verify column ON DEVICE (same 2-pass kernels +
13368 // smallest-index tie-break as host argmax, argmax_gate-validated) and read back ONE
13369 // [T] u32 — replaces the T x n_vocab f32 dtoh + T host argmaxes per round.
13370 // t_pred[j] = target's greedy prediction for the slot after draft[j-1] (j>=1) or after
13371 // last_token (j==0). With a pending bonus, col 0 IS the prediction after last_token
13372 // (== the bonus), so every index shifts by `base` and last_pred is unused.
13373 let t_v = verify_tokens.len();
13374 let mut preds: Vec<u32> = Vec::new();
13375 if !sampled {
13376 for j in 0..t_v {
13377 e.argmax_token_device_col(&tlogits_d, j, n_vocab, &mut preds_d, j)?;
13378 }
13379 preds = e.dtoh_u32(&preds_d)?; // <- the verify-GPU wait lands here
13380 // #87 SENTINEL TRAP, verify side: a sentinel pred becomes the round's bonus =
13381 // next round's last_token = the next chain's embed lookup. Catch it at the
13382 // source with the column named — an all-NaN VERIFY column implicates the
13383 // stage-split trunk (decode_step_t_core_ppn), not the draft head.
13384 if let Some(bad) = preds[..t_v].iter().position(|&p| (p as usize) >= n_vocab) {
13385 let col = &tlogits_d.slice(bad * n_vocab..(bad + 1) * n_vocab);
13386 let mut probe = e.zeros(n_vocab)?;
13387 e.copy_view_into(&mut probe, 0, col, n_vocab)?;
13388 let col_h = e.dtoh(&probe)?;
13389 let col_nan = col_h.iter().filter(|v| v.is_nan()).count();
13390 return Err(format!(
13391 "verify argmax sentinel 0x{:08x} >= n_vocab {n_vocab} at round {round} \
13392 col {bad}/{t_v} pos={pos}: verify-logits col NaN {col_nan}/{n_vocab} \
13393 — the verify TRUNK produced a poisoned column (#87 trap). Run \
13394 MEMRA_SPEC_NAN_SCAN=1 to name the layer that creates it (=2 to split \
13395 that layer into attention and routed MoE). NOT the draft head, and NOT \
13396 the PP stage split this message used to name: pp_cuts() returns None \
13397 without MEMRA_PP_STAGES, so decode_step_t_core_ppn never runs unless \
13398 that variable is set.",
13399 preds[bad]
13400 )
13401 .into());
13402 }
13403 }
13404 ph_mark(&mut ph_wait, phase_on);
13405 let t_pred = |j: usize| -> u32 {
13406 if j == 0 && base == 0 {
13407 last_pred
13408 } else {
13409 // GREEDY-ONLY: `preds` is filled under `if !sampled` above. The debug print
13410 // used to call this from the sampled arm and panicked the worker; it now goes
13411 // through `debug_t_pred0`. Keep the strict index here — in the greedy walk an
13412 // out-of-range pred is a real bug, not something to paper over.
13413 debug_assert!(
13414 !sampled,
13415 "t_pred is greedy-only: `preds` is empty in the sampled arm"
13416 );
13417 preds[base + j - 1]
13418 }
13419 };
13420 let mut devacc_seeded = false;
13421 let mut devacc_acc: Option<CudaSlice<u32>> = None;
13422 let (n_acc, bonus) = if !sampled {
13423 // ROUND-STREAM stage (a) (MEMRA_SPEC_DEVACC=1 opt-in): the walk runs ON DEVICE
13424 // (spec_accept_greedy, verbatim rule) and the host reads back 8B (n_acc, bonus)
13425 // instead of the [T] preds. Same sync count — machinery for stages (b)/(c),
13426 // gated on token identity vs the host walk (the arms below are bit-equal rules).
13427 if crate::spec::spec_devacc() && k_round > 0 && !spec_replay && constraint.is_none()
13428 {
13429 let draft_d = e.htod_u32_v(&draft)?;
13430 let mut acc_out = e.alloc_u32_zeroed(2)?;
13431 e.spec_accept_greedy(
13432 &preds_d,
13433 &draft_d,
13434 last_pred,
13435 base,
13436 k_round,
13437 &mut acc_out,
13438 )?;
13439 devacc_acc = Some(acc_out.clone());
13440 // stage (b): next-round seed gathered ON DEVICE from acc_out before the host
13441 // ever reads n_acc (j=base+n_acc -> vx col j-1; j==0 -> fill_prev). The three
13442 // non-replay commit arms skip their host-offset seed copies (guarded below);
13443 // the legacy spec_replay arm keeps its own rx-based seeding (excluded here).
13444 // NOTE: fill_prev is NOT updated here — the commit arms' TRUE-HIDDEN
13445 // REFRESH reads the OLD fill_prev (predecessor of this round's verify batch);
13446 // the update lands after the arms (devacc_seeded guard below).
13447 e.spec_seed_gather(&vx, &fill_prev, &acc_out, &mut h_seed_buf, base, n_embd)?;
13448 // 3a: KV lens roll back on device (len = saved + base + n_acc, all arms'
13449 // unified rule; full accept rewrites the verify-left value). Host mirrors
13450 // update after the readback; commit_verified_prefix skips its len_d writes.
13451 if let Some(successor) = successor_attempt.as_ref() {
13452 opti_fork
13453 .as_mut()
13454 .ok_or("optipipe successor reconcile lost fork state")?
13455 .queue_actual_reconcile(
13456 e,
13457 &snap,
13458 &acc_out,
13459 successor.verify_tokens[0],
13460 base,
13461 )?;
13462 } else if let Some(ptrs) = &kv_len_ptrs {
13463 let saved: Vec<i32> = (0..self.layers.len())
13464 .map(|il| snap.kv_len[il].map(|v| v as i32).unwrap_or(0))
13465 .collect();
13466 let saved_d = e.htod_i32(&saved)?;
13467 e.spec_rollback_kv(ptrs, &saved_d, &acc_out, base, self.layers.len())?;
13468 }
13469 devacc_seeded = true;
13470 let ab = e.dtoh_u32(&acc_out)?;
13471 (ab[0] as usize, ab[1])
13472 } else {
13473 let mut n_acc = 0usize;
13474 for j in 0..k_round {
13475 if t_pred(j) == draft[j] {
13476 n_acc += 1;
13477 } else {
13478 break;
13479 }
13480 }
13481 // bonus = target's own token at the first non-accepted slot. n_acc in 0..=k; t_pred
13482 // is defined for j in 0..=k (j==0 -> last_logits, j>=1 -> col j-1, last col = k-1).
13483 (n_acc, t_pred(n_acc))
13484 }
13485 } else {
13486 // --- SAMPLED ACCEPT (rejection sampling): u_j < p_j(x_j)/q_j(x_j) walk ---
13487 if col_buf.is_none() {
13488 col_buf = Some(e.zeros(n_vocab)?);
13489 }
13490 // FILTERED p_j: per-verify-col stats (one batched filter_stats call), then the
13491 // filtered gather. j==0&&base==0 reads last_col (its own stats row appended).
13492 let mut pj = vec![0f32; k_round.max(1)];
13493 let mut col_stats: Vec<(f32, f32, f32)> = Vec::new(); // (max, th, z) per verify col used
13494 if k_round > 0 {
13495 let mut ids: Vec<u32> = Vec::new();
13496 let mut rows: Vec<i32> = Vec::new();
13497 for j in 0..k_round {
13498 if j > 0 || base == 1 {
13499 ids.push(draft[j]);
13500 rows.push((base + j) as i32 - 1);
13501 }
13502 }
13503 if !ids.is_empty() {
13504 let nr = rows.len();
13505 // penalties: materialize the used columns into one contiguous penalized
13506 // buffer (rows remapped 0..nr) so stats+gathers see the penalized p.
13507 // penalties: materialize used columns contiguously, penalize all rows in
13508 // one launch, and point stats+gathers at the penalized buffer (rows 0..nr).
13509 let p_rows: Vec<i32> = if pen_on {
13510 (0..nr as i32).collect()
13511 } else {
13512 rows.clone()
13513 };
13514 if pen_on {
13515 if pcol_buf.as_ref().map(|b| b.len()).unwrap_or(0) < nr * n_vocab {
13516 pcol_buf = Some(e.zeros(nr * n_vocab)?);
13517 }
13518 let pc = pcol_buf.as_mut().unwrap();
13519 for (i2, &r) in rows.iter().enumerate() {
13520 let c = r as usize;
13521 e.copy_view_into(
13522 pc,
13523 i2 * n_vocab,
13524 &tlogits_d.slice(c * n_vocab..(c + 1) * n_vocab),
13525 n_vocab,
13526 )?;
13527 }
13528 let h = pen_hist_d.as_ref().unwrap();
13529 let nh = h.len();
13530 e.penalize_logits_rows(
13531 pc,
13532 h,
13533 nh,
13534 sp.penalty_repeat,
13535 sp.penalty_freq,
13536 sp.penalty_present,
13537 n_vocab,
13538 nr,
13539 )?;
13540 }
13541 let p_src: &CudaSlice<f32> = if pen_on {
13542 pcol_buf.as_ref().unwrap()
13543 } else {
13544 &tlogits_d
13545 };
13546 let rowsd = e.htod_i32(&p_rows)?;
13547 let (mut th_d, mut z_d, mut mx_d) =
13548 (e.zeros(nr)?, e.zeros(nr)?, e.zeros(nr)?);
13549 e.filter_stats(
13550 p_src, n_vocab, &rowsd, &mut th_d, &mut z_d, &mut mx_d, n_vocab, nr,
13551 sp_temp, sp.top_k, sp.top_p, sp.min_p,
13552 )?;
13553 let idsd = e.htod_u32_v(&ids)?;
13554 let mut outd = e.zeros(nr)?;
13555 e.softmax_gather_filtered(
13556 p_src, n_vocab, &idsd, &rowsd, &th_d, &z_d, &mut outd, n_vocab, nr,
13557 sp_temp,
13558 )?;
13559 let outv = e.dtoh(&outd)?;
13560 let (thv, zv, mxv) = (e.dtoh(&th_d)?, e.dtoh(&z_d)?, e.dtoh(&mx_d)?);
13561 let mut oi = 0usize;
13562 for j in 0..k_round {
13563 if j > 0 || base == 1 {
13564 pj[j] = outv[oi];
13565 oi += 1;
13566 }
13567 }
13568 col_stats = (0..nr).map(|i| (mxv[i], thv[i], zv[i])).collect();
13569 }
13570 if base == 0 {
13571 let lc: &CudaSlice<f32> = if pen_on {
13572 if col_buf.is_none() {
13573 col_buf = Some(e.zeros(n_vocab)?);
13574 }
13575 let cb = col_buf.as_mut().unwrap();
13576 e.copy_into(
13577 cb,
13578 0,
13579 last_col_logits
13580 .as_ref()
13581 .expect("sampled: last_col_logits unset"),
13582 n_vocab,
13583 )?;
13584 let h = pen_hist_d.as_ref().unwrap();
13585 let nh = h.len();
13586 e.penalize_logits(
13587 cb,
13588 h,
13589 nh,
13590 sp.penalty_repeat,
13591 sp.penalty_freq,
13592 sp.penalty_present,
13593 n_vocab,
13594 )?;
13595 col_buf.as_ref().unwrap()
13596 } else {
13597 last_col_logits
13598 .as_ref()
13599 .expect("sampled: last_col_logits unset")
13600 };
13601 let rows0 = e.htod_i32(&[0])?;
13602 let (mut th_d, mut z_d, mut mx_d) = (e.zeros(1)?, e.zeros(1)?, e.zeros(1)?);
13603 e.filter_stats(
13604 lc, n_vocab, &rows0, &mut th_d, &mut z_d, &mut mx_d, n_vocab, 1,
13605 sp_temp, sp.top_k, sp.top_p, sp.min_p,
13606 )?;
13607 let idsd = e.htod_u32_v(&[draft[0]])?;
13608 let mut outd = e.zeros(1)?;
13609 e.softmax_gather_filtered(
13610 lc, n_vocab, &idsd, &rows0, &th_d, &z_d, &mut outd, n_vocab, 1, sp_temp,
13611 )?;
13612 pj[0] = e.dtoh(&outd)?[0];
13613 last_col_stats =
13614 Some((e.dtoh(&mx_d)?[0], e.dtoh(&th_d)?[0], e.dtoh(&z_d)?[0]));
13615 }
13616 }
13617 // q source: the graph arm retained the head logits in the persistent q_slots;
13618 // the eager arm in per-round draft_logits clones. Same raw-logit values either way.
13619 // FILTERED q_j: stats from draft_stats (eager pushes in-chain; the graph arm
13620 // computes them post-replay — graph engages only filter/penalty-free, so the
13621 // stats degenerate to th=0/full-Z there, keeping ONE accept path).
13622 let q_bufs: &[CudaSlice<f32>] = if dctx.graph_s.is_some() {
13623 &dctx.q_slots
13624 } else {
13625 &draft_logits
13626 };
13627 let mut n_acc = 0usize;
13628 for j in 0..k_round {
13629 let (qmx, qth, qz) = draft_stats[j];
13630 let idsd = e.htod_u32_v(&[draft_idx[j]])?;
13631 let rowsd = e.htod_i32(&[0])?;
13632 let thd = e.htod(&[qth])?;
13633 let zd = e.htod(&[qz])?;
13634 let _ = qmx;
13635 let mut outd = e.zeros(1)?;
13636 e.softmax_gather_filtered(
13637 &q_bufs[j], d_vocab, &idsd, &rowsd, &thd, &zd, &mut outd, d_vocab, 1,
13638 sp_temp,
13639 )?;
13640 let qj = e.dtoh(&outd)?[0];
13641 let u = host_u01(sp_seed, uctr);
13642 uctr += 1;
13643 let accept = (u as f64) * (qj as f64) < pj[j] as f64;
13644 // SKEY PROBE: q == 0 for the token the draft actually proposed is the
13645 // exactness signature (see `skey_probe`). Impossible when the draft was
13646 // drawn from the same filtered distribution the verify reconstructs here;
13647 // `u * 0 < p` makes it an UNCONDITIONAL accept whenever p > 0.
13648 if skey_probe() && qj == 0.0 {
13649 eprintln!(
13650 "[skey] EXACTNESS q=0 round={round} j={j} draft_tok={} \
13651 draft_idx={} p={:e} u={u} accepted={} th_z={:?}",
13652 draft[j], draft_idx[j], pj[j], accept as u8, draft_stats[j],
13653 );
13654 }
13655 if accept {
13656 n_acc += 1;
13657 } else {
13658 break;
13659 }
13660 }
13661 let bonus = if n_acc == k_round {
13662 // FULL ACCEPT: bonus ~ FILTERED softmax at the last verify column.
13663 let col = base + k_round - 1;
13664 let cb = col_buf.as_mut().unwrap();
13665 e.copy_view_into(
13666 cb,
13667 0,
13668 &tlogits_d.slice(col * n_vocab..(col + 1) * n_vocab),
13669 n_vocab,
13670 )?;
13671 if pen_on {
13672 let h = pen_hist_d.as_ref().unwrap();
13673 let nh = h.len();
13674 e.penalize_logits(
13675 cb,
13676 h,
13677 nh,
13678 sp.penalty_repeat,
13679 sp.penalty_freq,
13680 sp.penalty_present,
13681 n_vocab,
13682 )?;
13683 }
13684 if perturb_buf.is_none() {
13685 perturb_buf = Some(e.zeros(d_vocab.max(n_vocab))?);
13686 }
13687 // STATS MUST COME FROM THIS COLUMN (bug fix 2026-08-05, lane/sampler-
13688 // truncation-fix; receipts research/sampfix-20260805/). The old code reused
13689 // `col_stats.last()` here, which is ALWAYS the wrong row: the gathered set
13690 // covers verify columns 0..=(base+k_round-2) (rows pushed as base+j-1), while
13691 // the full-accept bonus samples column base+k_round-1 — exactly ONE PAST the
13692 // last gathered column, in both base arms. `th` is a threshold in e-units of
13693 // its OWN row's max, so feeding a neighbour's (row_max, th) into
13694 // gumbel_perturb_filtered mis-scales every e0 = exp((x-row_max)/T). When the
13695 // donor column's peak is higher by more than T*ln(1/th), EVERY id fails
13696 // `e0 >= th`, the whole perturbed row becomes -3.4e38, and the 2-pass argmax
13697 // falls through to its smallest-index tie-break => token id 0 ("!") spliced
13698 // mid-word. Fragility is ordered by how large th is: min_p pins th = min_p
13699 // (0.05 => trigger at delta > 2.4 at T=0.8, fires constantly), top_p's
13700 // mass-boundary th is smaller, top_k's k-th-largest th smaller still — which
13701 // is why the head-to-head matrix saw min_p and top_p corrupt while top_k-only
13702 // stayed clean. The pure-temp default regime is immune (th == 0 masks nothing,
13703 // and row_max is unused once nothing is masked), so this fix is a byte-level
13704 // no-op for the untruncated serve default. One extra one-block filter_stats
13705 // per full-accept round is the whole cost.
13706 let (mx, th) = {
13707 let rows0 = e.htod_i32(&[0])?;
13708 let (mut th_d, mut z_d, mut mx_d) = (e.zeros(1)?, e.zeros(1)?, e.zeros(1)?);
13709 let cb0 = col_buf.as_ref().unwrap();
13710 e.filter_stats(
13711 cb0, n_vocab, &rows0, &mut th_d, &mut z_d, &mut mx_d, n_vocab, 1,
13712 sp_temp, sp.top_k, sp.top_p, sp.min_p,
13713 )?;
13714 (e.dtoh(&mx_d)?[0], e.dtoh(&th_d)?[0])
13715 };
13716 let pb = perturb_buf.as_mut().unwrap();
13717 let cb2 = col_buf.as_ref().unwrap();
13718 e.gumbel_perturb_filtered(cb2, pb, n_vocab, sp_seed, sctr, sp_temp, mx, th)?;
13719 sctr += 1;
13720 let td = e.argmax_token_device(pb, n_vocab)?;
13721 e.dtoh_u32_one(&td)?
13722 } else {
13723 // REJECT at n_acc: bonus ~ norm(max(0, softmax_T(p) - softmax_T(q))).
13724 let cb = col_buf.as_mut().unwrap();
13725 if n_acc > 0 || base == 1 {
13726 let col = base + n_acc - 1;
13727 e.copy_view_into(
13728 cb,
13729 0,
13730 &tlogits_d.slice(col * n_vocab..(col + 1) * n_vocab),
13731 n_vocab,
13732 )?;
13733 } else {
13734 let lc = last_col_logits.as_ref().unwrap();
13735 e.copy_into(cb, 0, lc, n_vocab)?;
13736 }
13737 if pen_on {
13738 let h = pen_hist_d.as_ref().unwrap();
13739 let nh = h.len();
13740 e.penalize_logits(
13741 cb,
13742 h,
13743 nh,
13744 sp.penalty_repeat,
13745 sp.penalty_freq,
13746 sp.penalty_present,
13747 n_vocab,
13748 )?;
13749 }
13750 let cb2 = col_buf.as_ref().unwrap();
13751 let sc = sctr;
13752 sctr += 1;
13753 // p-stats for the reject column: from col_stats when the col was gathered,
13754 // else (j==0&&base==0) from last_col_stats.
13755 let p_stats = if n_acc > 0 || base == 1 {
13756 // col index within the gathered set == number of gathered cols before n_acc
13757 let gi = if base == 1 { n_acc } else { n_acc - 1 };
13758 col_stats.get(gi).copied().unwrap_or_else(|| {
13759 (0.0, 0.0, 1.0) // unreachable: gathered cols always cover the reject slot
13760 })
13761 } else {
13762 last_col_stats.expect("sampled: last_col_stats unset at reject")
13763 };
13764 let q_stats = draft_stats[n_acc];
13765 if let Some(map) = &d2t_dev {
13766 if q_full_buf.is_none() {
13767 q_full_buf = Some(e.zeros(n_vocab)?);
13768 }
13769 let qf = q_full_buf.as_mut().unwrap();
13770 e.scatter_trim_logits(&q_bufs[n_acc], map, qf, d_vocab, n_vocab)?;
13771 let qf2 = q_full_buf.as_ref().unwrap();
13772 e.residual_sample_filtered(
13773 cb2,
13774 Some(qf2),
13775 n_vocab,
13776 sp_temp,
13777 sp_seed,
13778 sc,
13779 p_stats,
13780 q_stats,
13781 &mut sample_tok,
13782 )?;
13783 } else {
13784 e.residual_sample_filtered(
13785 cb2,
13786 Some(&q_bufs[n_acc]),
13787 n_vocab,
13788 sp_temp,
13789 sp_seed,
13790 sc,
13791 p_stats,
13792 q_stats,
13793 &mut sample_tok,
13794 )?;
13795 }
13796 e.dtoh_u32(&sample_tok)?[0]
13797 };
13798 (
13799 n_acc,
13800 guard_vocab_token(
13801 bonus,
13802 n_vocab,
13803 &format!("sampled verify bonus at round {round} pos={pos} n_acc={n_acc}"),
13804 )?,
13805 )
13806 };
13807 // --- 3b. GRAMMAR TRUNCATION (constrained spec, 2026-08-03): the grammar is
13808 // an extra rejection rule AFTER the exactness verify (the batched-verify-twins
13809 // ordering). Walk the accepted drafts through the grammar in commit order; the
13810 // first illegal token truncates acceptance at its slot, and that slot's emission
13811 // is recomputed as the MASKED argmax of the target's own verify column — token-
13812 // identical to constrained plain greedy decode (an unmasked argmax that is
13813 // grammar-legal IS the masked argmax: masking only removes competitors). The
13814 // column D2H (~1MB) is paid only when a cut fires — the tight-grammar cost,
13815 // measured in acceptance numbers, never hidden.
13816 let (n_acc, bonus) = match constraint.as_deref_mut() {
13817 None => (n_acc, bonus),
13818 Some(c) => {
13819 fn ce(e2: String) -> Box<dyn std::error::Error> {
13820 format!("constraint: {e2}").into()
13821 }
13822 let mut na = n_acc;
13823 let mut cut = false;
13824 for (j, &d) in draft.iter().enumerate().take(n_acc) {
13825 if c.is_allowed(d).map_err(ce)? {
13826 c.consume(d).map_err(ce)?;
13827 } else {
13828 na = j;
13829 cut = true;
13830 dm_cut_tokens += n_acc - j;
13831 break;
13832 }
13833 }
13834 if cut {
13835 dm_cuts += 1;
13836 }
13837 let mut bo = bonus;
13838 if cut || !c.is_allowed(bo).map_err(ce)? {
13839 let mut row = if na == 0 && base == 0 {
13840 init_logits_host
13841 .clone()
13842 .ok_or("constraint: init logits missing (round-0 cut)")?
13843 } else {
13844 e.dtoh_view(
13845 &tlogits_d.slice((base + na - 1) * n_vocab..(base + na) * n_vocab),
13846 )?
13847 };
13848 c.mask_logits(&mut row).map_err(ce)?;
13849 bo = argmax(&row) as u32;
13850 }
13851 c.consume(bo).map_err(ce)?;
13852 (na, bo)
13853 }
13854 };
13855 let mut successor_valid = false;
13856 if let Some((q_proxy, expected_d2)) = rejected_probe {
13857 let v_n = n_acc == 1 && bonus == expected_d2;
13858 eprintln!(
13859 "[opti-controller] shadow q={q_proxy:.6} admitted=false v_n={v_n} \
13860 expected_d2={expected_d2} n_acc={n_acc} bonus={bonus}",
13861 );
13862 }
13863 if let Some(successor) = successor_attempt.as_ref() {
13864 successor_valid = n_acc == 1 && bonus == successor.verify_tokens[0];
13865 let generation = successor.generation;
13866 let q_proxy = successor.q_proxy;
13867 let expected_pending = successor.verify_tokens[0];
13868 let resolution_ms = successor.issued_at.elapsed().as_secs_f64() * 1e3;
13869 let fork = opti_fork
13870 .as_mut()
13871 .ok_or("optipipe successor resolution lost fork state")?;
13872 fork.finish_actual_reconcile(e, &mut *cache, &snap, n_acc, base, successor_valid)?;
13873 if successor_valid {
13874 OPTI_FORK_HITS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
13875 } else {
13876 OPTI_FORK_MISSES.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
13877 OPTI_RECONCILES.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
13878 OPTI_WASTED_DRAFT_TOKENS.fetch_add(2, std::sync::atomic::Ordering::Relaxed);
13879 }
13880 let breaker_tripped = fork
13881 .controller
13882 .as_mut()
13883 .expect("controller policy")
13884 .resolve(successor_valid);
13885 if breaker_tripped {
13886 OPTI_BREAKER_TRIPS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
13887 }
13888 eprintln!(
13889 "[opti-controller] resolve generation={} hit={} q={q_proxy:.6} \
13890 expected_pending={expected_pending} n_acc={n_acc} bonus={bonus} \
13891 resolution_ms={resolution_ms:.3} reconcile={} breaker={}",
13892 generation.id, successor_valid, !successor_valid, breaker_tripped,
13893 );
13894 if !successor_valid {
13895 let mut successor = successor_attempt
13896 .take()
13897 .expect("controller successor disappeared on miss");
13898 successor.settle();
13899 fork.retire(generation)?;
13900 }
13901 }
13902 total_drafted += k_round;
13903 total_accepted += n_acc;
13904 if let Some(t) = sess_telem {
13905 // Greedy, rejection-sampling, and grammar truncation all converge here after
13906 // the accept decision is already on host. Fixed-size relaxed atomics only.
13907 t.record_round(k_round, n_acc);
13908 }
13909 if spec_stats {
13910 st_len_hist[k_round] += 1;
13911 for j in 0..k_round {
13912 st_drafted[j] += 1;
13913 }
13914 for j in 0..n_acc {
13915 st_accepted[j] += 1;
13916 }
13917 if n_acc == k_round {
13918 st_full += 1;
13919 }
13920 }
13921
13922 if debug_spec {
13923 eprintln!(
13924 "[R{round}] pos={pos} out_len={} last_tok={last_token} draft={draft:?} n_acc={n_acc} bonus={bonus} t_pred0={}",
13925 out.len(),
13926 // NOT `t_pred(0)`: `preds` is filled only under `if !sampled` above, so on a
13927 // sampled request round >= 1 (base == 1) indexed an EMPTY vector and PANICKED
13928 // the GPU worker thread — a debug flag that killed the exact regime you would
13929 // set it to investigate. See `debug_t_pred0`.
13930 debug_t_pred0(sampled, base, last_pred, &preds)
13931 );
13932 }
13933
13934 // --- 4. COMMIT: draft[0..n_acc] then bonus (n_acc + 1 tokens) ---
13935 let commit_started = std::time::Instant::now();
13936 // SESSION MODE: every accepted column is already in the CACHE — `out` must carry all
13937 // of them (overshoot past max_new included) or `committed` under-counts the cache rows
13938 // and the next turn's continuation seeds one token off (gate-caught 2026-07-05). The
13939 // single-shot path keeps the cap (its caller truncates + drops the cache anyway).
13940 for j in 0..n_acc {
13941 if !session_mode && out.len() >= max_new {
13942 break;
13943 }
13944 out.push(draft[j]);
13945 }
13946 if pen_on {
13947 pen_hist.extend_from_slice(&draft[0..n_acc]);
13948 pen_hist.push(bonus);
13949 }
13950 let bonus_emitted = session_mode || out.len() < max_new;
13951 if bonus_emitted {
13952 out.push(bonus);
13953 }
13954 last_token = bonus;
13955
13956 // --- 5. ROLLBACK + advance (§C) ---
13957 if n_acc == k_round && !spec_replay {
13958 // FULL ACCEPT, BONUS FOLD: all verify columns (pending? + drafts) are committed in
13959 // cache; the NEW bonus stays PENDING for the next round's verify batch — NO extra
13960 // T=1 trunk pass. The next draft chain seeds from the MTP block's h_nextn at the
13961 // bonus position: one MTP-block pass (~1/33 trunk cost) replaces the trunk read.
13962 // last_pred is dead in the pending path (t_pred reads verify col 0).
13963 //
13964 // PERSISTENT DRAFT KV, full-accept fill: the chain covered last_token +
13965 // draft[0..k_round-2] as INPUTS (slots P..P'-2); draft[k_round-1] (slot P'-1) was
13966 // only ever an output, so its entry is MISSING. Fill it from vh_seed — its EXACT
13967 // trunk hidden (the last verify column). set_len first: a p-min break may have
13968 // left one extra chain append at that slot. Partial accepts need NO fill (the
13969 // chain already covered every accepted position; round-start set_len truncates).
13970 let mut vh_seed = e.zeros(n_embd)?;
13971 e.copy_view_into(
13972 &mut vh_seed,
13973 0,
13974 &vx.slice((t_v - 1) * n_embd..t_v * n_embd),
13975 n_embd,
13976 )?;
13977 if refresh {
13978 // TRUE-HIDDEN REFRESH (2026-07-03, the HANDOVER-listed acceptance lever):
13979 // overwrite ALL committed positions' scratch entries with K/V from their EXACT
13980 // verify hiddens — the reference engine's mtp_update fills from true hiddens;
13981 // the full stack (vx) is already resident from the verify. Replaces both the
13982 // chain-approximate entries AND the old last-token-only fill. Acceptance-only
13983 // (draft attention quality); exactness stays the verify's job.
13984 scratch.set_len(e, pos)?;
13985 // PREDECESSOR pairing: row i gets vx[i-1]; row 0 the carried fill_prev
13986 // (hidden of the last committed row before this verify batch).
13987 let mut vxs = e.zeros(t_v * n_embd)?;
13988 e.copy_into(&mut vxs, 0, &fill_prev, n_embd)?;
13989 if t_v > 1 {
13990 e.copy_view_into(
13991 &mut vxs,
13992 n_embd,
13993 &vx.slice(0..(t_v - 1) * n_embd),
13994 (t_v - 1) * n_embd,
13995 )?;
13996 }
13997 self.mtp_kv_fill_all(e, &verify_tokens, &vxs, pos, &mut *scratch, embd_dev)?;
13998 } else {
13999 scratch.set_len(e, pos + base + k_round - 1)?;
14000 // predecessor of the last draft = verify col t_v-2 (or fill_prev at t_v==1)
14001 let mut hp = e.zeros(n_embd)?;
14002 if t_v >= 2 {
14003 e.copy_view_into(
14004 &mut hp,
14005 0,
14006 &vx.slice((t_v - 2) * n_embd..(t_v - 1) * n_embd),
14007 n_embd,
14008 )?;
14009 } else {
14010 e.copy_into(&mut hp, 0, &fill_prev, n_embd)?;
14011 }
14012 self.mtp_kv_fill_all(
14013 e,
14014 &[draft[k_round - 1]],
14015 &hp,
14016 pos + base + k_round - 1,
14017 &mut *scratch,
14018 embd_dev,
14019 )?;
14020 }
14021 // REFERENCE SEEDING: no pseudo pass — the next chain's step 0 IS the
14022 // reference's (id_last, h_prev) draft row; it appends the bonus's scratch
14023 // entry itself. Seed = TRUE hidden of the bonus's predecessor (last verify
14024 // col). Saves one MTP-block pass per round on top of the pairing fix.
14025 if !devacc_seeded {
14026 e.copy_into(&mut h_seed_buf, 0, &vh_seed, n_embd)?;
14027 e.copy_into(&mut fill_prev, 0, &vh_seed, n_embd)?;
14028 }
14029 pending = Some(bonus);
14030 if debug_spec {
14031 eprintln!(" -> FULL ACCEPT (bonus pending, prev-h seed)");
14032 }
14033 } else if !spec_replay && base + n_acc >= 1 {
14034 // PARTIAL ACCEPT, REPLAY-FREE (2026-07-03 — the profiled #1 long-ctx spec cost):
14035 // the verify's first j = base+n_acc columns ARE the committed sequence, computed
14036 // bit-identically to eager (decode-exact contract) — so KEEP them: KV truncates to
14037 // pos+j, recurrent state rebuilds from the VerifyCkpt (same-kernel gdn prefix
14038 // re-run / pure state-clone restore), and the bonus stays PENDING exactly like the
14039 // full-accept path — the legacy duplicate trunk replay is gone. The next chain
14040 // seeds from the MTP pseudo-hidden of the bonus, whose seed = the TRUE verify
14041 // hidden of its predecessor (col j-1) — same one-hop pseudo structure as full
14042 // accept (never compounds: the next verify recomputes true hiddens for all
14043 // committed columns).
14044 let j = base + n_acc;
14045 // VERIFY-GRAPH SLAB COMMIT: when the captured trunk ran, the linear layers'
14046 // column stash was written into the graphs ctx's persistent slabs as in-graph
14047 // memcpy nodes, NOT into the per-column VerifyCkpt the cols arm reads — so the
14048 // commit must take the slab twin (same semantics, slab-addressed sources). The
14049 // ctx states which of the two this round produced via `round_slab`; trusting the
14050 // flag rather than the env keeps a round that fell back to the eager walk (a
14051 // capture that declined, a t the pool never captured) on the cols arm.
14052 let slab_commit = vg_guard
14053 .as_ref()
14054 .and_then(|g| g.as_ref())
14055 .map(|g| g.round_slab)
14056 .unwrap_or(false);
14057 if slab_commit {
14058 self.dspark_commit_prefix_slab(
14059 e,
14060 &mut *cache,
14061 &snap,
14062 vg_guard
14063 .as_ref()
14064 .and_then(|g| g.as_ref())
14065 .expect("slab_commit implies a graphs ctx"),
14066 j,
14067 )?;
14068 } else {
14069 self.commit_verified_prefix(
14070 e,
14071 &mut *cache,
14072 &snap,
14073 ckpt.as_ref().unwrap(),
14074 j,
14075 devacc_seeded,
14076 if devacc_seeded {
14077 devacc_acc.as_ref().map(|a| (a, base, t_v))
14078 } else {
14079 None
14080 },
14081 )?;
14082 }
14083 let mut seed = e.zeros(n_embd)?;
14084 e.copy_view_into(
14085 &mut seed,
14086 0,
14087 &vx.slice((j - 1) * n_embd..j * n_embd),
14088 n_embd,
14089 )?;
14090 // Draft scratch: TRUE-HIDDEN REFRESH of the committed prefix (see the full-accept
14091 // branch); without it the chain entries stand and only the tail truncates. Either
14092 // way len ends at pos+j so the pseudo append lands at the bonus's slot pos+j
14093 // (persistent mode), rope pos+j+1 (chain convention).
14094 if refresh {
14095 scratch.set_len(e, pos)?;
14096 let mut vxs = e.zeros(j * n_embd)?;
14097 e.copy_into(&mut vxs, 0, &fill_prev, n_embd)?;
14098 if j > 1 {
14099 e.copy_view_into(
14100 &mut vxs,
14101 n_embd,
14102 &vx.slice(0..(j - 1) * n_embd),
14103 (j - 1) * n_embd,
14104 )?;
14105 }
14106 self.mtp_kv_fill_all(
14107 e,
14108 &verify_tokens[0..j],
14109 &vxs,
14110 pos,
14111 &mut *scratch,
14112 embd_dev,
14113 )?;
14114 } else {
14115 scratch.set_len(e, pos + j)?;
14116 }
14117 // REFERENCE SEEDING (see the full-accept branch): seed = TRUE hidden of the
14118 // bonus's predecessor (verify col j-1); no pseudo pass.
14119 if !devacc_seeded {
14120 e.copy_into(&mut h_seed_buf, 0, &seed, n_embd)?;
14121 e.copy_into(&mut fill_prev, 0, &seed, n_embd)?;
14122 }
14123 pending = Some(bonus);
14124 if debug_spec {
14125 eprintln!(" -> PARTIAL(replay-free j={j}, bonus pending, prev-h seed)");
14126 }
14127 } else if !spec_replay {
14128 // ZERO ROUND FOLD (2026-07-10, verify-cost target #3): base+n_acc == 0 — a
14129 // pending-less round where nothing was accepted (PMIN0 zero-draft chains after a
14130 // replay/commit, or plain 0-accept rounds at round 0). The old path replayed
14131 // [bonus] through a FULL m=1 trunk+head forward (the 489us full-vocab head pass
14132 // measured at ~0.75/round on PMIN0 configs). Instead: restore the pre-round
14133 // snapshot and let the bonus ride the NEXT round's verify as col 0 — the existing
14134 // base=1 pending machinery, bit-identical by the decode-exact verify contract.
14135 // Seed: the bonus's predecessor is the last COMMITTED token, whose hidden
14136 // fill_prev already carries (same seeding as the 1-token-replay case it replaces).
14137 cache.rollback(e, &snap, 0)?;
14138 scratch.set_len(e, pos)?;
14139 e.copy_into(&mut h_seed_buf, 0, &fill_prev, n_embd)?;
14140 pending = Some(bonus);
14141 if debug_spec {
14142 eprintln!(" -> ZERO-ROUND FOLD (bonus pending, fill_prev seed)");
14143 }
14144 } else {
14145 // PARTIAL ACCEPT, LEGACY REPLAY (seam MEMRA_SPEC_REPLAY=1 — or j==0: nothing of
14146 // this round survives, only possible before the first pending exists, ~round 0):
14147 // restore EVERYTHING to the pre-round snapshot (KV truncate to pos + recur
14148 // restore), then replay the committed prefix pending? ++ draft[0..n_acc] ++
14149 // [bonus] as ONE batched T forward — single weight read, bit-identical to greedy
14150 // (the verify-all-columns path is the same math). Commits the bonus with a TRUE
14151 // trunk hidden.
14152 cache.rollback(e, &snap, 0)?; // accept_len=0: KV len = pos, recur = snapshot
14153 let mut replay: Vec<u32> = Vec::with_capacity(base + n_acc + 1);
14154 if let Some(b) = pending.take() {
14155 replay.push(b);
14156 }
14157 replay.extend_from_slice(&draft[0..n_acc]);
14158 replay.push(bonus);
14159 // Full-stack forward (decode_step_t_core = decode_step_t_h_emb_dev's body):
14160 // Predecessor pairing seeds from the PREDECESSOR row (col len-2) — the same-row path takes the
14161 // last col exactly as before (byte-identical to the old _h_emb_dev call).
14162 let (rl_d, rx) = if self.batched_serving_numeric_class() {
14163 let mut logits = Vec::with_capacity(replay.len() * n_vocab);
14164 let mut hidden = e.uninit(replay.len() * n_embd)?;
14165 for (row, &token) in replay.iter().enumerate() {
14166 let (row_logits, row_hidden) =
14167 self.spec_target_step_h(e, token, &mut *cache)?;
14168 logits.extend_from_slice(&row_logits);
14169 e.dtod_copy_into(&row_hidden, &mut hidden, row * n_embd)?;
14170 }
14171 (e.htod(&logits)?, hidden)
14172 } else {
14173 self.decode_step_t_core(e, &replay, pos, &mut *cache, embd_dev, None)?
14174 };
14175 // last_pred = argmax of the LAST column's logits (predicts the token after `bonus`)
14176 // — device argmax + one 4-byte read instead of the full-vocab column dtoh.
14177 e.argmax_token_device_col(&rl_d, replay.len() - 1, n_vocab, &mut preds_d, 0)?;
14178 last_pred = guard_vocab_token(
14179 e.dtoh_u32(&preds_d)?[0],
14180 n_vocab,
14181 &format!("replay last_pred at round {round} pos={pos}"),
14182 )?;
14183 if sampled {
14184 let lr0 = replay.len();
14185 let lc = last_col_logits
14186 .as_mut()
14187 .expect("sampled: last_col_logits unset");
14188 e.copy_view_into(
14189 lc,
14190 0,
14191 &rl_d.slice((lr0 - 1) * n_vocab..lr0 * n_vocab),
14192 n_vocab,
14193 )?;
14194 }
14195 let lr = replay.len();
14196 if lr >= 2 {
14197 e.copy_view_into(
14198 &mut h_seed_buf,
14199 0,
14200 &rx.slice((lr - 2) * n_embd..(lr - 1) * n_embd),
14201 n_embd,
14202 )?;
14203 } else {
14204 // 1-token replay (round-0 miss): the bonus's predecessor is the OLD
14205 // last_token, whose own-row hidden fill_prev still holds.
14206 e.copy_into(&mut h_seed_buf, 0, &fill_prev, n_embd)?;
14207 }
14208 // the bonus is COMMITTED here — it becomes the last committed row.
14209 let mut rh_last = e.zeros(n_embd)?;
14210 e.copy_view_into(
14211 &mut rh_last,
14212 0,
14213 &rx.slice((lr - 1) * n_embd..lr * n_embd),
14214 n_embd,
14215 )?;
14216 e.copy_into(&mut fill_prev, 0, &rh_last, n_embd)?;
14217 if debug_spec {
14218 eprintln!(" -> PARTIAL(replay={replay:?}), next_pred={last_pred}");
14219 }
14220 }
14221 if devacc_seeded {
14222 // stage (b) epilogue: fill_prev takes the gathered seed AFTER the refresh fills
14223 // consumed the old value (both slots carry the same value in every non-replay arm).
14224 e.copy_into(&mut fill_prev, 0, &h_seed_buf, n_embd)?;
14225 }
14226 if successor_valid {
14227 let optimistic_scratch_len = successor_attempt
14228 .as_ref()
14229 .expect("valid controller successor disappeared")
14230 .scratch_len;
14231 // The normal current-round commit refreshed/truncated the logical scratch tail.
14232 // Its optimistic successor row was already written physically, so restoring only
14233 // the retained logical length makes that row live for the carried round.
14234 scratch.set_len(e, optimistic_scratch_len)?;
14235 }
14236 if let Some(current) = current_opti.take() {
14237 opti_fork
14238 .as_mut()
14239 .ok_or("optipipe current retirement lost fork state")?
14240 .retire(current.generation)?;
14241 }
14242 if successor_valid {
14243 let successor = successor_attempt
14244 .take()
14245 .expect("valid controller successor disappeared before promotion");
14246 let generation = successor.generation;
14247 opti_fork
14248 .as_mut()
14249 .ok_or("optipipe successor promotion lost fork state")?
14250 .promote_successor_snapshot(&mut snap, generation);
14251 carried_opti = Some(successor);
14252 }
14253 if anatomy_on {
14254 // Commit/rollback is normally asynchronous on the primary/head stream. Bound it
14255 // only for this diagnostic so it does not disappear into the following draft's
14256 // first token readback.
14257 e.stream().synchronize()?;
14258 ph_commit += commit_started.elapsed().as_secs_f64();
14259 }
14260 // adaptive-K update (host math, zero syncs): next round drafts accepted-run + 1,
14261 // clamped to [floor(pos), k_cap]. cache.pos is post-rollback here (the round's
14262 // final position — the floor's position key reads the committed depth). Burst
14263 // rounds (`continue` above) draft the captured fixed depth and skip this, exactly
14264 // like gemma's burst arm.
14265 if adapt {
14266 let fl_now = floor_at(cache.pos);
14267 kc = (n_acc + 1).clamp(fl_now.min(k_cap), k_cap);
14268 }
14269 ph_mark(&mut ph_rest, phase_on);
14270 if let Some(p) = pipe {
14271 p.accept_end(round);
14272 }
14273 drop(pipe_accept);
14274 if let Some(t0) = round_t0 {
14275 let ms = t0.elapsed().as_secs_f64() * 1e3;
14276 ROUND_MS.fetch_add((ms * 1e3) as u64, std::sync::atomic::Ordering::Relaxed);
14277 let n = ROUND_N.fetch_add(1, std::sync::atomic::Ordering::Relaxed) + 1;
14278 if n % 32 == 0 {
14279 eprintln!(
14280 "[spec-round] rounds={n} avg round wall={:.2} ms (emitted={} drafted so far)",
14281 ROUND_MS.load(std::sync::atomic::Ordering::Relaxed) as f64 / 1e3 / n as f64,
14282 out.len()
14283 );
14284 }
14285 }
14286 round += 1;
14287 // sse-cadence: this round's accepted drafts + bonus are committed (out is
14288 // append-only past step 4) — flush at round cadence.
14289 keep_going = flush_commit(&mut on_commit, &out, &mut flushed);
14290 }
14291 if let Some(mut ticket) = carried_opti.take() {
14292 opti_fork
14293 .as_mut()
14294 .ok_or("optipipe tail drain lost fork state")?
14295 .cancel_controller_ticket(e, &mut *cache, &mut *scratch, &snap, &mut ticket)?;
14296 }
14297 // sse-cadence: nothing below appends to `out`; flush any remainder (defensive).
14298 // (verdict ignored — the burst is over either way; the session tail runs unchanged.)
14299 let _ = flush_commit(&mut on_commit, &out, &mut flushed);
14300
14301 if spec_stats {
14302 let per_slot: Vec<String> = (0..k)
14303 .map(|j| {
14304 if st_drafted[j] > 0 {
14305 format!(
14306 "{}/{}={:.3}",
14307 st_accepted[j],
14308 st_drafted[j],
14309 st_accepted[j] as f64 / st_drafted[j] as f64
14310 )
14311 } else {
14312 "0/0".into()
14313 }
14314 })
14315 .collect();
14316 let acc = if total_drafted > 0 {
14317 total_accepted as f64 / total_drafted as f64
14318 } else {
14319 0.0
14320 };
14321 eprintln!(
14322 "[spec-stats] rounds={round} full_accept={st_full} len_hist={st_len_hist:?} \
14323 per_slot=[{}] total={total_accepted}/{total_drafted}={acc:.3} \
14324 tok_per_round={:.3}",
14325 per_slot.join(" "),
14326 (total_accepted + round) as f64 / round.max(1) as f64
14327 );
14328 }
14329 if constraint.is_some() {
14330 eprintln!(
14331 "[draft-mask] mask_rounds={dm_rounds} clone_total={:.3}ms \
14332 clone_per_round={:.4}ms gram_cuts={dm_cuts}/{round} cut_tokens={dm_cut_tokens}",
14333 dm_clone_ns as f64 / 1e6,
14334 dm_clone_ns as f64 / 1e6 / dm_rounds.max(1) as f64
14335 );
14336 }
14337 if phase_on {
14338 let tot = ph_draft + ph_verify + ph_wait + ph_rest;
14339 eprintln!(
14340 "[spec-phase] draft={:.1}ms ({:.1}%) verify-issue={:.1}ms ({:.1}%) verify-wait={:.1}ms ({:.1}%) commit-host={:.1}ms ({:.1}%) rounds={round}",
14341 ph_draft * 1e3,
14342 ph_draft / tot * 100.0,
14343 ph_verify * 1e3,
14344 ph_verify / tot * 100.0,
14345 ph_wait * 1e3,
14346 ph_wait / tot * 100.0,
14347 ph_rest * 1e3,
14348 ph_rest / tot * 100.0
14349 );
14350 }
14351 if anatomy_on {
14352 let rounds_f = round.max(1) as f64;
14353 let other = (ph_rest - ph_commit).max(0.0);
14354 eprintln!(
14355 "[spec-anatomy] per-round draft={:.3}ms pp-verify={:.3}ms \
14356 verify-accept={:.3}ms commit-rollback={:.3}ms other={:.3}ms rounds={round}",
14357 ph_draft * 1e3 / rounds_f,
14358 ph_verify * 1e3 / rounds_f,
14359 ph_wait * 1e3 / rounds_f,
14360 ph_commit * 1e3 / rounds_f,
14361 other * 1e3 / rounds_f,
14362 );
14363 }
14364 let _pipe_tail = pipe.map(|p| p.primary());
14365 // SESSION TAIL: leave the session in the exact invariant the next turn's suffix prime
14366 // expects — every row in `committed` has trunk KV/recur state AND an exact draft-KV row.
14367 // Park the draft-graph ctx back on the session (the serve-burst fixed-cost fix): the next
14368 // burst replays instead of recapturing. Error paths (`?` above) drop it — recaptured then.
14369 if let Some(slot) = sess_draft_slot.take() {
14370 *slot = Some(dctx);
14371 }
14372 let t_rounds = t_ent.elapsed();
14373 if let Some((committed, last_h, next_pred_slot, sctr_slot, uctr_slot)) = sess_tail.take() {
14374 // NEXT BURST'S BOUNDARY TOKEN (lane/sampled-spec-quality, Item 1). Greedy stashes
14375 // the argmax `last_pred` exactly as before (byte contract). SAMPLED draws the token
14376 // HERE, where the sampler, the session Philox counters and the penalty window are
14377 // all live and the boundary logits row still exists — that is the "make the state
14378 // available" half of the fix; the consuming burst then just emits it. `sctr` is
14379 // written to the session BELOW the draws so the advance is never lost.
14380 *next_pred_slot = Some(last_pred);
14381 let sample_boundary = sampled && constraint.is_none() && spec_sampled_boundary_on();
14382 let mut stashed_pending = false;
14383 if let Some(b) = pending.take() {
14384 if !sampled {
14385 // PENDING-CARRY (2026-08-01): stash the bonus on the session instead of
14386 // committing it with a solo T=1 pass — the next empty-suffix greedy burst
14387 // consumes it as round-0 verify col 0 (a plain round edge; the old tail
14388 // commit + next burst's init feed were 11.6+11.5ms solo trunk passes per
14389 // burst on H100 q27, [spec-setup] trace). b stays in `out` (emitted) but
14390 // OUT of `committed` (cache rows == committed); the consuming call
14391 // prepends it once its verify commits the row. next_pred is unknowable
14392 // without the commit pass — None; callers gate on pending_tok too.
14393 debug_assert_eq!(out.last(), Some(&b), "pending must be the last emitted");
14394 if let Some(slot) = sess_pending_slot.take() {
14395 *slot = Some(b);
14396 }
14397 *next_pred_slot = None;
14398 // fill_prev = hidden of the last COMMITTED row (b's predecessor) — the
14399 // exact chain-seed/fill anchor the consuming burst (or a flush) needs.
14400 *last_h = Some(e.clone_dtod(&fill_prev)?);
14401 stashed_pending = true;
14402 } else {
14403 // SAMPLED tail (unchanged): commit the bonus (one T=1 pass) + draft fill —
14404 // the sampled round-0 accept needs this pass's logits (last_col_logits).
14405 let pos_b = cache.pos;
14406 scratch.set_len(e, pos_b)?;
14407 let (lg_b, hb) = self.spec_target_step_h(e, b, &mut *cache)?;
14408 // after a FULL-accept exit `last_pred` is STALE (it predicted the bonus
14409 // itself — the prediction AFTER the bonus never materialized; it would have
14410 // been the next round's verify col 0). The commit's logits ARE that
14411 // prediction — so they are also the row the next burst's boundary token
14412 // comes off, and (lane/sampled-spec-quality) it is DRAWN from them here.
14413 *next_pred_slot = Some(if sample_boundary {
14414 sample_boundary_token(
14415 e,
14416 &lg_b,
14417 &sp,
14418 &pen_hist,
14419 &mut sctr,
14420 "burst-tail-commit",
14421 )?
14422 } else {
14423 argmax(&lg_b) as u32
14424 });
14425 self.mtp_kv_fill_all(e, &[b], &fill_prev, pos_b, &mut *scratch, embd_dev)?;
14426 *last_h = Some(hb);
14427 }
14428 } else {
14429 // fill_prev tracks the hidden of the last COMMITTED row throughout the loop.
14430 *last_h = Some(e.clone_dtod(&fill_prev)?);
14431 if sample_boundary {
14432 // No pending to commit, so the boundary row is the one `last_pred` was
14433 // argmaxed from and the sampled path keeps it on device: the init feed's
14434 // logits when the burst ran zero rounds, else the legacy-replay path's
14435 // last verify column (both predict the token AFTER the last committed
14436 // row). It is retained precisely because round 0's accept test needs it,
14437 // so the draw costs no extra D2H of the [n_vocab] row.
14438 match last_col_logits.as_ref() {
14439 Some(lc) => {
14440 *next_pred_slot = Some(sample_boundary_token_dev(
14441 e,
14442 lc,
14443 n_vocab,
14444 &sp,
14445 &pen_hist,
14446 &mut sctr,
14447 "burst-tail-nopending",
14448 )?);
14449 }
14450 // NAME THE FALLBACK (house standard): unreachable today — a sampled
14451 // burst always feeds or replays, so the row exists — but if it ever
14452 // is, the stream takes a greedy token and SAYS so rather than
14453 // silently regressing to the pre-lane behaviour.
14454 None => eprintln!(
14455 "[spec-boundary] sampled tail kept the ARGMAX boundary token \
14456 (reason: no retained boundary logits row)"
14457 ),
14458 }
14459 }
14460 }
14461 *sctr_slot = sctr;
14462 *uctr_slot = uctr;
14463 committed.extend_from_slice(prompt);
14464 if let Some(cb) = carried_pending {
14465 // the consumed carry's cache row landed in round 0's verify (every pending
14466 // round commits col 0) — it joins `committed` here, in sequence order.
14467 committed.push(cb);
14468 }
14469 if stashed_pending {
14470 // ZERO-EMIT BURST (2026-08-06 c=8 serve panic, pre-existing since b4aea184):
14471 // `out.len() - 1` underflowed on an EMPTY `out` — "range end index
14472 // 18446744073709551615 out of range for slice of length 0", killing the
14473 // memra-gpu-worker and failing 31 of 32 concurrent requests with "worker closed
14474 // stream". Reachable because `pending` starts as `carried_pending` (a bonus
14475 // stashed by the PREVIOUS burst) while `out` starts empty, and the carry is
14476 // deliberately NOT pushed to `out` (line ~3239: the burst that emitted it already
14477 // did). So a burst that stashes a pending without emitting anything of its own —
14478 // the round loop exits before a push, e.g. the ring drain's `out.len() < max_new`
14479 // guard skipping every token under a tight budget — arrives here with
14480 // out.len() == 0 and stashed_pending == true.
14481 //
14482 // The invariant is unchanged: `committed` gets every emitted token EXCEPT the
14483 // stashed bonus. With nothing emitted, that is nothing — and the carry pushed
14484 // just above is already accounted. Saturating, not a min/assert: an empty `out`
14485 // here is a legitimate burst shape, not a corrupt state.
14486 let emitted = out.len().saturating_sub(1);
14487 committed.extend_from_slice(&out[..emitted]);
14488 } else {
14489 committed.extend_from_slice(&out); // FULL out incl. overshoot — all committed
14490 }
14491 debug_assert_eq!(
14492 cache.pos,
14493 committed.len(),
14494 "session invariant: cache rows == committed tokens"
14495 );
14496 if setup_trace {
14497 e.stream().synchronize()?; // bound the async tail fill in the trace
14498 let t_tail = t_ent.elapsed();
14499 eprintln!(
14500 "[spec-setup] init={:.2}ms cap={:.2}ms fill={:.2}ms rounds={:.2}ms tail={:.2}ms total={:.2}ms out={} cont={}",
14501 t_init.as_secs_f64() * 1e3,
14502 (t_cap - t_init).as_secs_f64() * 1e3,
14503 (t_fill - t_cap).as_secs_f64() * 1e3,
14504 (t_rounds - t_fill).as_secs_f64() * 1e3,
14505 (t_tail - t_rounds).as_secs_f64() * 1e3,
14506 t_tail.as_secs_f64() * 1e3,
14507 out.len(),
14508 continuation
14509 );
14510 }
14511 return Ok((out, total_drafted, total_accepted));
14512 }
14513 out.truncate(max_new);
14514 Ok((out, total_drafted, total_accepted))
14515 }
14516
14517 /// Anchor-bounded DSpark target extraction. The trunk sees the exact generated token tape;
14518 /// only requested hidden rows and target-logit rows cross PCIe. An anchor token at p pairs
14519 /// with the pre-output-norm h[p-1] carrier, exactly as the existing replay/NextN path does.
14520 pub fn extract_dspark_anchors(
14521 &self,
14522 e: &Engine,
14523 tokens: &[u32],
14524 anchor_positions: &[usize],
14525 gamma: usize,
14526 top_k: usize,
14527 chunk: usize,
14528 temperature: f32,
14529 ) -> Result<Vec<DsparkAnchorRecord>, Box<dyn std::error::Error>> {
14530 if tokens.len() < gamma + 2 || gamma == 0 || chunk < 2 {
14531 return Err("DSpark extraction token tape/gamma/chunk is invalid".into());
14532 }
14533 if anchor_positions.windows(2).any(|pair| pair[0] >= pair[1]) {
14534 return Err("DSpark anchor positions must be sorted and unique".into());
14535 }
14536 for &position in anchor_positions {
14537 if position == 0 || position + gamma >= tokens.len() {
14538 return Err(format!(
14539 "DSpark anchor {position} has no predecessor or cannot cover gamma={gamma} in {} tokens",
14540 tokens.len()
14541 )
14542 .into());
14543 }
14544 }
14545
14546 let n_vocab = self.output.out_features();
14547 let n_embd = self.cfg.n_embd as usize;
14548 let mut cache =
14549 crate::pp::new_cache_planned(e, &self.cfg, &self.plan, tokens.len() + gamma + 8)?;
14550 let (embd_qt, embd_rb) = self.embd.qt_and_row_bytes(n_embd);
14551 let embd_gpu = if spec_host_embd() {
14552 None
14553 } else {
14554 Some(
14555 self.embd_gpu
14556 .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload")),
14557 )
14558 };
14559 let embd_dev = embd_gpu.map(|gpu| (gpu, embd_qt, embd_rb));
14560
14561 struct PendingRecord {
14562 position: usize,
14563 hidden: Option<Vec<f32>>,
14564 tokens: Vec<u32>,
14565 target_top_ids: Vec<Option<Vec<u32>>>,
14566 target_top_logits: Vec<Option<Vec<f32>>>,
14567 target_top_probs: Vec<Option<Vec<f32>>>,
14568 target_tail_probs: Vec<Option<f32>>,
14569 }
14570
14571 let mut pending: Vec<PendingRecord> = anchor_positions
14572 .iter()
14573 .map(|&position| PendingRecord {
14574 position,
14575 hidden: None,
14576 tokens: tokens[position..=position + gamma].to_vec(),
14577 target_top_ids: vec![None; gamma],
14578 target_top_logits: vec![None; gamma],
14579 target_top_probs: vec![None; gamma],
14580 target_tail_probs: vec![None; gamma],
14581 })
14582 .collect();
14583
14584 let mut start = 0usize;
14585 while start < tokens.len() {
14586 let end = (start + chunk).min(tokens.len());
14587 let chunk_tokens = &tokens[start..end];
14588 let (target_logits, hidden_rows) =
14589 self.decode_step_t_core(e, chunk_tokens, start, &mut cache, embd_dev, None)?;
14590 for record in &mut pending {
14591 let hidden_position = record.position - 1;
14592 if hidden_position >= start && hidden_position < end {
14593 let local = hidden_position - start;
14594 record.hidden = Some(
14595 e.dtoh_view(&hidden_rows.slice(local * n_embd..(local + 1) * n_embd))?,
14596 );
14597 }
14598 for slot in 0..gamma {
14599 let target_row = record.position + slot;
14600 if target_row < start || target_row >= end {
14601 continue;
14602 }
14603 let local = target_row - start;
14604 let logits =
14605 e.dtoh_view(&target_logits.slice(local * n_vocab..(local + 1) * n_vocab))?;
14606 let (ids, top_logits, probs, tail) =
14607 dspark_sparse_softmax_topk(&logits, top_k, temperature)?;
14608 record.target_top_ids[slot] = Some(ids);
14609 record.target_top_logits[slot] = Some(top_logits);
14610 record.target_top_probs[slot] = Some(probs);
14611 record.target_tail_probs[slot] = Some(tail);
14612 }
14613 }
14614 start = end;
14615 }
14616
14617 pending
14618 .into_iter()
14619 .map(|record| {
14620 let hidden = record
14621 .hidden
14622 .ok_or_else(|| format!("missing DSpark hidden at {}", record.position))?;
14623 let target_top_ids =
14624 flatten_dspark_rows(record.target_top_ids, record.position, "target ids")?;
14625 let target_top_logits = flatten_dspark_rows(
14626 record.target_top_logits,
14627 record.position,
14628 "target logits",
14629 )?;
14630 let target_top_probs =
14631 flatten_dspark_rows(record.target_top_probs, record.position, "target probs")?;
14632 let target_tail_probs = record
14633 .target_tail_probs
14634 .into_iter()
14635 .enumerate()
14636 .map(|(slot, value)| {
14637 value.ok_or_else(|| {
14638 format!("missing DSpark tail at {} slot {slot}", record.position)
14639 })
14640 })
14641 .collect::<Result<Vec<_>, _>>()?;
14642 Ok(DsparkAnchorRecord {
14643 position: record.position,
14644 hidden,
14645 tokens: record.tokens,
14646 target_top_ids,
14647 target_top_logits,
14648 target_top_probs,
14649 target_tail_probs,
14650 })
14651 })
14652 .collect()
14653 }
14654
14655 /// TEACHER-FORCED REPLAY ACCEPTANCE (hqmtp MTP-heal protocol): walk a FIXED token
14656 /// sequence and, at sampled positions, compare the MTP head's K-token draft chain against
14657 /// the trunk's own teacher-forced greedy predictions. Nothing is generated — the context is
14658 /// the corpus text itself, so (a) degenerate self-generated loops cannot inflate acceptance
14659 /// and (b) two arms (bf16 ceiling vs NVFP4) score on IDENTICAL contexts, isolating the
14660 /// quant-induced head/hidden-state mismatch from text drift.
14661 ///
14662 /// Per eval position p (context = tokens[0..=p], predecessor pairing as in spec decode):
14663 /// draft_j = chain token j from (tokens[p], h_{p-1}), then its own drafts — the exact
14664 /// eager spec-decode chain (same mtp_head_forward_dev, same rope positions).
14665 /// target_j = teacher-forced greedy pick for position p+1+j (argmax of the trunk logits
14666 /// at forced context tokens[0..p+j]). For j==0 this equals live spec
14667 /// acceptance; for j>=1 live verify would condition on the drafts, here it
14668 /// conditions on the corpus — deterministic and arm-comparable by design.
14669 ///
14670 /// Returns (rows, bg): one (p, drafts[k], targets[k]) row per eval position (ascending p),
14671 /// plus the full teacher-forced greedy track bg (bg[i] = greedy pick for position i, i>=1)
14672 /// so harnesses can cross-check runs (e.g. different chunk sizes must give identical bg).
14673 ///
14674 /// `hdump`: when Some, every position's pre-output_norm trunk hidden (the exact rows the
14675 /// draft-KV fill pairs from) streams to the file as little-endian f32 [t_total, n_embd] —
14676 /// the head-distillation extraction (hqmtp): the ENGINE is the source of truth for trunk
14677 /// hiddens (HF torch reproductions of the hybrid trunk measured only ~0.5 greedy
14678 /// agreement vs this path — not usable as a training-data source).
14679 pub fn replay_acceptance(
14680 &self,
14681 e: &Engine,
14682 tokens: &[u32],
14683 k: usize,
14684 stride: usize,
14685 chunk: usize,
14686 mut hdump: Option<&mut std::fs::File>,
14687 ) -> Result<(Vec<(usize, Vec<u32>, Vec<u32>)>, Vec<u32>), Box<dyn std::error::Error>> {
14688 assert!(k >= 1 && stride >= 1 && chunk >= 2);
14689 let mtp = self
14690 .mtp
14691 .as_ref()
14692 .expect("replay_acceptance requires an MTP head");
14693 let n_vocab = self.output.out_features();
14694 let d_vocab = mtp
14695 .shared_head_head
14696 .as_ref()
14697 .unwrap_or(&self.output)
14698 .out_features();
14699 let n_embd = self.cfg.n_embd as usize;
14700 let t_total = tokens.len();
14701 assert!(t_total >= 8, "corpus too short ({t_total} tokens)");
14702 // STAGE-OWNED KV (lane/pp2-spec 2026-08-06) — see `new_session`. Door shut = `Cache::new`.
14703 let mut cache = crate::pp::new_cache_planned(e, &self.cfg, &self.plan, t_total + k + 8)?;
14704 let mut scratch = self.new_mtp_scratch(e, t_total + k + 8)?;
14705 let (embd_qt, embd_rb) = self.embd.qt_and_row_bytes(n_embd);
14706 let embd_gpu = if spec_host_embd() {
14707 None
14708 } else {
14709 Some(
14710 self.embd_gpu
14711 .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload")),
14712 )
14713 };
14714 let embd_dev = embd_gpu.map(|g| (g, embd_qt, embd_rb));
14715
14716 // bg[i] = the trunk's greedy pick for position i under the forced context (i >= 1).
14717 let mut bg: Vec<u32> = vec![0; t_total + 1];
14718 let mut rows: Vec<(usize, Vec<u32>, Vec<u32>)> = Vec::new();
14719 let mut prev_last_h = e.zeros(n_embd)?; // predecessor hidden entering the chunk
14720 let mut seed_buf = e.zeros(n_embd)?;
14721 let mut preds_d = e.alloc_u32_zeroed(chunk)?;
14722 let nll_on = std::env::var("MEMRA_REPLAY_NLL").as_deref() == Ok("1");
14723 let (mut nll_sum, mut nll_cnt) = (0f64, 0u64);
14724 let mut s = 0usize;
14725 while s < t_total {
14726 let cend = (s + chunk).min(t_total);
14727 let tc = cend - s;
14728 let ch = &tokens[s..cend];
14729 // 1. forced trunk pass — verify path (decode-exact contract): all-column logits +
14730 // the chunk's true hiddens.
14731 let (tl_d, vx) = self.decode_step_t_core(e, ch, s, &mut cache, embd_dev, None)?;
14732 for j in 0..tc {
14733 e.argmax_token_device_col(&tl_d, j, n_vocab, &mut preds_d, j)?;
14734 }
14735 let preds = e.dtoh_u32(&preds_d)?;
14736 for j in 0..tc {
14737 bg[s + j + 1] = preds[j];
14738 }
14739 // MEMRA_REPLAY_NLL=1: teacher-forced NLL/perplexity over the same forced pass — the
14740 // checkpoint-quality metric (position j's logits score the GOLD next token).
14741 if nll_on {
14742 let jmax = if cend < t_total { tc } else { tc - 1 }; // last pos has no gold next
14743 if jmax > 0 {
14744 let ids: Vec<u32> = (0..jmax).map(|j| tokens[s + j + 1]).collect();
14745 let rows: Vec<i32> = (0..jmax as i32).collect();
14746 let idsd = e.htod_u32_v(&ids)?;
14747 let rowsd = e.htod_i32(&rows)?;
14748 let mut outd = e.zeros(jmax)?;
14749 e.softmax_gather(&tl_d, n_vocab, &idsd, &rowsd, &mut outd, n_vocab, jmax, 1.0)?;
14750 for pr in e.dtoh(&outd)? {
14751 nll_sum += -((pr.max(1e-30)) as f64).ln();
14752 nll_cnt += 1;
14753 }
14754 }
14755 }
14756 if let Some(f) = hdump.as_deref_mut() {
14757 use std::io::Write;
14758 let host: Vec<f32> = e.dtoh(&vx)?;
14759 // bf16 round-to-nearest-even — f32 doubled the disk bill at bulk
14760 // extraction scale (20M tokens x 4096 = 320GB f32 vs 160GB bf16).
14761 let mut bytes = Vec::with_capacity(tc * n_embd * 2);
14762 for v in &host[..tc * n_embd] {
14763 let b = v.to_bits();
14764 let r = b.wrapping_add(0x7FFF + ((b >> 16) & 1));
14765 bytes.extend_from_slice(&((r >> 16) as u16).to_le_bytes());
14766 }
14767 f.write_all(&bytes)?;
14768 }
14769 // CHAINLESS extraction (stride > corpus, the bulk-hdump mode): no chunk ever
14770 // drafts, so the draft-KV fills are pure waste — skip them (2 MTP-block passes
14771 // per token saved; the forced trunk pass + hdump is all the mode needs).
14772 let chainless = stride > t_total;
14773 if chainless {
14774 e.copy_view_into(
14775 &mut prev_last_h,
14776 0,
14777 &vx.slice((tc - 1) * n_embd..tc * n_embd),
14778 n_embd,
14779 )?;
14780 s = cend;
14781 continue;
14782 }
14783 // 2. TRUE predecessor-paired draft-KV fill for the chunk (row i carries h_{i-1};
14784 // row s reads the previous chunk's last true hidden, zeros at corpus start).
14785 let mut vxs = e.zeros(tc * n_embd)?;
14786 e.copy_into(&mut vxs, 0, &prev_last_h, n_embd)?;
14787 if tc > 1 {
14788 e.copy_view_into(
14789 &mut vxs,
14790 n_embd,
14791 &vx.slice(0..(tc - 1) * n_embd),
14792 (tc - 1) * n_embd,
14793 )?;
14794 }
14795 scratch.set_len(e, s)?;
14796 self.mtp_kv_fill_all(e, ch, &vxs, s, &mut scratch, embd_dev)?;
14797 // 3. draft chains at sampled positions, DESCENDING: a chain reads only slots
14798 // [0..p) (true fills) and appends at >= p; the next (smaller-p) chain's set_len
14799 // truncates those approximate appends before they can ever be read.
14800 let ps: Vec<usize> = (s..cend)
14801 .filter(|p| *p >= 1 && *p % stride == 0 && *p + k <= t_total)
14802 .collect();
14803 for &p in ps.iter().rev() {
14804 scratch.set_len(e, p)?;
14805 if p == s {
14806 e.copy_into(&mut seed_buf, 0, &prev_last_h, n_embd)?;
14807 } else {
14808 e.copy_view_into(
14809 &mut seed_buf,
14810 0,
14811 &vx.slice((p - 1 - s) * n_embd..(p - s) * n_embd),
14812 n_embd,
14813 )?;
14814 }
14815 let mut e_tok = tokens[p];
14816 let mut d_seed = e.clone_dtod(&seed_buf)?;
14817 let chain_heads = !self.mtp_extra.is_empty();
14818 let mut chain_tokens = if chain_heads {
14819 vec![tokens[p]]
14820 } else {
14821 Vec::new()
14822 };
14823 let mut chain_seeds = if chain_heads {
14824 vec![e.clone_dtod(&seed_buf)?]
14825 } else {
14826 Vec::new()
14827 };
14828 let mut drafts: Vec<u32> = Vec::with_capacity(k);
14829 for j in 0..k {
14830 let (dl_d, h_nextn) = if chain_heads {
14831 self.mtp_chain_forward_dev(
14832 e,
14833 &chain_tokens,
14834 &chain_seeds,
14835 &mut scratch,
14836 p,
14837 embd_dev,
14838 None,
14839 )?
14840 } else {
14841 self.mtp_head_forward_dev(
14842 e,
14843 mtp,
14844 e_tok,
14845 &d_seed,
14846 &mut scratch,
14847 p + 1 + j,
14848 embd_dev,
14849 None,
14850 )?
14851 };
14852 let tok_d = e.argmax_token_device(&dl_d, d_vocab)?;
14853 let idx = e.dtoh_u32_one(&tok_d)?;
14854 let d = match &mtp.d2t {
14855 Some(map) => map[idx as usize],
14856 None => idx,
14857 };
14858 drafts.push(d);
14859 if chain_heads {
14860 chain_tokens.push(d);
14861 chain_seeds.push(h_nextn);
14862 } else {
14863 e_tok = d;
14864 d_seed = h_nextn;
14865 }
14866 }
14867 // targets may live in a LATER chunk's bg — resolved after the walk.
14868 rows.push((p, drafts, Vec::new()));
14869 }
14870 // 4. restore TRUE entries for the whole chunk (the next chunk's chains and fills
14871 // expect scratch.len == cend with exact rows).
14872 scratch.set_len(e, s)?;
14873 self.mtp_kv_fill_all(e, ch, &vxs, s, &mut scratch, embd_dev)?;
14874 e.copy_view_into(
14875 &mut prev_last_h,
14876 0,
14877 &vx.slice((tc - 1) * n_embd..tc * n_embd),
14878 n_embd,
14879 )?;
14880 s = cend;
14881 }
14882 for (p, drafts, targets) in rows.iter_mut() {
14883 for j in 0..drafts.len() {
14884 targets.push(bg[*p + 1 + j]);
14885 }
14886 }
14887 rows.sort_by_key(|r| r.0);
14888 if nll_cnt > 0 {
14889 let mean = nll_sum / nll_cnt as f64;
14890 println!(
14891 "[replay-nll] tokens={nll_cnt} nll/token={mean:.5} ppl={:.4}",
14892 mean.exp()
14893 );
14894 }
14895 Ok((rows, bg))
14896 }
14897}
14898
14899#[cfg(test)]
14900mod vg_debt_tests {
14901 use super::dspark_vg_debt_projection;
14902
14903 /// TOOTH for the verify-graph admission accounting: the pool's projected remaining
14904 /// growth must be charged (pre-fix, admission charged 0 for a pool measured at
14905 /// 8,852 MiB), the projection must price the MARGINAL cost of one more key rather than
14906 /// extrapolating the pool's one-time shared allocation, and the doors that make growth
14907 /// impossible must zero the debt.
14908 #[test]
14909 fn vg_debt_projects_remaining_growth_and_respects_the_freeze_valves() {
14910 const MIB: usize = 1 << 20;
14911 let d = dspark_vg_debt_projection;
14912 // cold pool: nothing observed, one capture fits inside SPEC_SHRINK_RESERVE.
14913 assert_eq!(d(0, 256, 0, None), 0);
14914 // freeze valve MEMRA_DSPARK_VG_MAX=0: the pool cannot grow.
14915 assert_eq!(d(10, 0, 500 * MIB, None), 0);
14916 // saturated pool: at/past the cap the pool FREEZES, nothing left to reserve.
14917 assert_eq!(d(256, 256, 8852 * MIB, None), 0);
14918 assert_eq!(d(300, 256, 8852 * MIB, None), 0);
14919
14920 // BOOTSTRAP (one observation, growth unmeasurable): at most one more pool's worth.
14921 // The pre-fix mean rule extrapolated 255x here — the measured 8.5 GB phantom.
14922 assert_eq!(d(1, 256, 33 * MIB, None), 33 * MIB);
14923
14924 // MARGINAL, flat pool (the box9 receipt: reserved stayed ~33.6 MiB across captures
14925 // 1..3, so an additional key costs ~nothing and the debt must collapse to ~0 —
14926 // NOT the 8,556/4,261/2,830 MB the mean rule printed).
14927 assert_eq!(d(3, 256, 33 * MIB, Some((1, 33 * MIB))), 0);
14928
14929 // MARGINAL, genuinely growing pool: 40 MiB per new key over 2 keys, 250 slots left.
14930 let debt = d(6, 256, 273 * MIB, Some((4, 193 * MIB)));
14931 assert_eq!(debt, 250 * (40 * MIB));
14932 assert!(
14933 debt > 3 * (1536 * MIB),
14934 "real growth must dwarf SPEC_SHRINK_RESERVE"
14935 );
14936
14937 // a shrinking/recycled reading never becomes a negative charge.
14938 assert_eq!(d(6, 256, 10 * MIB, Some((4, 99 * MIB))), 0);
14939 // a stale observation at the same capture count falls back to bootstrap.
14940 assert_eq!(d(4, 256, 80 * MIB, Some((4, 80 * MIB))), 80 * MIB);
14941 }
14942}
14943
14944#[cfg(test)]
14945mod mtp_chain_tests {
14946 use super::mtp_chain_head_index;
14947
14948 #[test]
14949 fn embedded_step_heads_cycle_in_declared_order() {
14950 let actual: Vec<usize> = (0..8).map(|step| mtp_chain_head_index(step, 3)).collect();
14951 assert_eq!(actual, [0, 1, 2, 0, 1, 2, 0, 1]);
14952 }
14953
14954 #[test]
14955 fn standalone_draft_remains_single_head() {
14956 assert!((0..8).all(|step| mtp_chain_head_index(step, 1) == 0));
14957 }
14958}
14959
14960#[cfg(test)]
14961mod tp_verified_prefix_tests {
14962 use super::rewind_tp_kv_verified_prefix;
14963 use crate::tp::ResidentTpKvCache;
14964
14965 fn cache_with_committed_len(committed: usize) -> ResidentTpKvCache {
14966 let mut cache = ResidentTpKvCache::new(Vec::new(), 1, 1, 1, 1, 8);
14967 let transaction = cache.begin_transaction().unwrap();
14968 let target = cache.append_target(transaction, committed).unwrap();
14969 cache.publish_append(transaction, target).unwrap();
14970 let target = cache.commit_target(transaction, committed).unwrap();
14971 cache.publish_finalize(transaction, target).unwrap();
14972 cache
14973 }
14974
14975 #[test]
14976 fn replay_free_prefix_rewinds_tp_visibility_to_snapshot_plus_accepts() {
14977 let mut layers = vec![Some(cache_with_committed_len(5)), None];
14978 rewind_tp_kv_verified_prefix(&mut layers, &[Some(2), None], 1).unwrap();
14979 let cache = layers[0].as_ref().unwrap();
14980 assert_eq!(cache.committed_len(), 3);
14981 assert_eq!(cache.staged_len(), 3);
14982 }
14983
14984 #[test]
14985 fn replay_free_prefix_rejects_a_changed_tp_cache_shape() {
14986 let mut layers = vec![Some(cache_with_committed_len(1))];
14987 let error = rewind_tp_kv_verified_prefix(&mut layers, &[None], 1)
14988 .unwrap_err()
14989 .to_string();
14990 assert!(error.contains("changed shape"), "unexpected error: {error}");
14991 }
14992}
14993
14994#[cfg(test)]
14995mod dspark_sparse_tests {
14996 use super::dspark_sparse_softmax_topk;
14997
14998 #[test]
14999 fn topk_keeps_full_softmax_mass_and_stable_ties() {
15000 let logits = [1.0f32, 3.0, 3.0, -2.0];
15001 let (ids, top_logits, probs, tail) = dspark_sparse_softmax_topk(&logits, 2, 1.0).unwrap();
15002 assert_eq!(ids, vec![1, 2]);
15003 assert_eq!(top_logits, vec![3.0, 3.0]);
15004 let denominator = logits.iter().map(|value| (value - 3.0).exp()).sum::<f32>();
15005 let expected = 1.0 / denominator;
15006 assert!((probs[0] - expected).abs() < 1.0e-6);
15007 assert!((probs[1] - expected).abs() < 1.0e-6);
15008 assert!((tail - (1.0 - 2.0 * expected)).abs() < 1.0e-6);
15009 assert!((probs.iter().sum::<f32>() + tail - 1.0).abs() < 1.0e-6);
15010 }
15011}
15012
15013#[cfg(test)]
15014mod spec_replay_env_tests {
15015 use super::spec_replay_env_on;
15016
15017 #[test]
15018 fn replay_requires_literal_one() {
15019 assert!(!spec_replay_env_on(None));
15020 assert!(!spec_replay_env_on(Some("")));
15021 assert!(!spec_replay_env_on(Some("0")));
15022 assert!(!spec_replay_env_on(Some("true")));
15023 assert!(!spec_replay_env_on(Some("2")));
15024 assert!(spec_replay_env_on(Some("1")));
15025 }
15026}
15027
15028#[cfg(test)]
15029mod telem_tests {
15030 use super::{SPEC_TELEM_POS, SpecTelemetry, SpecTelemetryCounters};
15031
15032 #[test]
15033 fn synthetic_accept_masks_produce_tau_and_position_histogram() {
15034 let counters = SpecTelemetryCounters::default();
15035 for mask in [
15036 [true, true, true],
15037 [true, true, false],
15038 [true, false, false],
15039 [false, false, false],
15040 ] {
15041 let accepted = mask.iter().take_while(|&&value| value).count();
15042 counters.record_round(mask.len(), accepted);
15043 }
15044
15045 let snapshot = counters.snapshot();
15046 assert_eq!(
15047 (snapshot.rounds, snapshot.drafted, snapshot.accepted),
15048 (4, 12, 6)
15049 );
15050 assert_eq!(&snapshot.pos_drafted[..3], &[4, 4, 4]);
15051 assert_eq!(&snapshot.pos_accepted[..3], &[3, 2, 1]);
15052 assert_eq!(snapshot.tau(), 1.5);
15053 assert_eq!(snapshot.pos_drafted[3..], [0; SPEC_TELEM_POS - 3]);
15054 assert_eq!(snapshot.pos_accepted[3..], [0; SPEC_TELEM_POS - 3]);
15055 }
15056
15057 /// The worker's per-burst pattern: stash, accumulate, diff — the delta must isolate
15058 /// exactly the burst's contribution (pool-resumed sessions carry prior requests' counts).
15059 #[test]
15060 fn delta_isolates_burst_contribution() {
15061 let mut t = SpecTelemetry::default();
15062 // "previous request": 2 rounds of k=3, accepts 3 then 1.
15063 for (kr, na) in [(3usize, 3usize), (3, 1)] {
15064 t.rounds += 1;
15065 t.drafted += kr as u64;
15066 t.accepted += na as u64;
15067 for j in 0..kr {
15068 t.pos_drafted[j] += 1;
15069 }
15070 for j in 0..na {
15071 t.pos_accepted[j] += 1;
15072 }
15073 }
15074 let before = t;
15075 // "this burst": 1 round k=3, accepts 2.
15076 t.rounds += 1;
15077 t.drafted += 3;
15078 t.accepted += 2;
15079 for j in 0..3 {
15080 t.pos_drafted[j] += 1;
15081 }
15082 for j in 0..2 {
15083 t.pos_accepted[j] += 1;
15084 }
15085 let d = t.delta_since(&before);
15086 assert_eq!((d.rounds, d.drafted, d.accepted), (1, 3, 2));
15087 assert_eq!(&d.pos_drafted[..3], &[1, 1, 1]);
15088 assert_eq!(&d.pos_accepted[..3], &[1, 1, 0]);
15089 assert_eq!(d.pos_drafted[3..], [0; SPEC_TELEM_POS - 3]);
15090 }
15091
15092 /// merge(delta) then merge(delta2) equals accumulating both — the per-model /metrics
15093 /// aggregation invariant.
15094 #[test]
15095 fn merge_accumulates_fieldwise() {
15096 let mut agg = SpecTelemetry::default();
15097 let mut d1 = SpecTelemetry {
15098 rounds: 2,
15099 drafted: 6,
15100 accepted: 4,
15101 ..Default::default()
15102 };
15103 d1.pos_drafted[0] = 2;
15104 d1.pos_accepted[0] = 2;
15105 let mut d2 = SpecTelemetry {
15106 rounds: 1,
15107 drafted: 3,
15108 accepted: 1,
15109 ..Default::default()
15110 };
15111 d2.pos_drafted[0] = 1;
15112 d2.pos_accepted[0] = 1;
15113 d2.pos_drafted[1] = 1;
15114 agg.merge(&d1);
15115 agg.merge(&d2);
15116 assert_eq!((agg.rounds, agg.drafted, agg.accepted), (3, 9, 5));
15117 assert_eq!(agg.pos_drafted[0], 3);
15118 assert_eq!(agg.pos_accepted[0], 3);
15119 assert_eq!(agg.pos_drafted[1], 1);
15120 assert_eq!(agg.pos_accepted[1], 0);
15121 }
15122
15123 /// Wrong-snapshot diff saturates to zero instead of wrapping — the counters feed a
15124 /// public metrics surface and must never publish a u64-wrapped garbage value.
15125 #[test]
15126 fn delta_saturates_never_wraps() {
15127 let small = SpecTelemetry {
15128 rounds: 1,
15129 drafted: 2,
15130 accepted: 1,
15131 ..Default::default()
15132 };
15133 let big = SpecTelemetry {
15134 rounds: 5,
15135 drafted: 15,
15136 accepted: 9,
15137 ..Default::default()
15138 };
15139 let d = small.delta_since(&big);
15140 assert_eq!((d.rounds, d.drafted, d.accepted), (0, 0, 0));
15141 }
15142}
15143
15144#[cfg(test)]
15145mod opti_fork_tests {
15146 use super::{
15147 OptiControllerPolicy, OptiForkAction, OptiForkGateMode, OptiForkGenerationTracker,
15148 };
15149
15150 #[test]
15151 fn controller_threshold_and_three_miss_breaker_are_exact() {
15152 let mut policy = OptiControllerPolicy {
15153 threshold: 0.7,
15154 consecutive_misses: 0,
15155 breaker_tripped: false,
15156 };
15157 assert!(!policy.admit(0.699_999));
15158 assert!(policy.admit(0.7));
15159 assert!(!policy.resolve(false));
15160 assert!(!policy.resolve(false));
15161 assert!(policy.resolve(false));
15162 assert!(policy.breaker_tripped);
15163 assert!(!policy.admit(1.0));
15164 assert!(
15165 !policy.resolve(true),
15166 "a resolved hit cannot re-arm a tripped request"
15167 );
15168 assert!(policy.breaker_tripped);
15169 }
15170
15171 #[test]
15172 fn zero_threshold_is_the_true_unconditional_measurement_arm() {
15173 let mut policy = OptiControllerPolicy {
15174 threshold: 0.0,
15175 consecutive_misses: 0,
15176 breaker_tripped: false,
15177 };
15178 for _ in 0..16 {
15179 assert!(policy.admit(0.0));
15180 assert!(!policy.resolve(false));
15181 }
15182 for invalid in [f32::NAN, f32::INFINITY, -0.01, 1.01] {
15183 assert!(
15184 !policy.admit(invalid),
15185 "invalid q proxy must fail closed: {invalid}"
15186 );
15187 }
15188 assert!(!policy.breaker_tripped);
15189 assert_eq!(policy.consecutive_misses, 0);
15190 }
15191
15192 #[test]
15193 fn alternating_mode_flips_by_generation_not_round_parity() {
15194 assert_eq!(OptiForkGateMode::Alternate.action(0), OptiForkAction::Hit);
15195 assert_eq!(OptiForkGateMode::Alternate.action(1), OptiForkAction::Miss);
15196 assert_eq!(OptiForkGateMode::Alternate.action(8), OptiForkAction::Hit);
15197 assert_eq!(OptiForkGateMode::Alternate.action(9), OptiForkAction::Miss);
15198 }
15199
15200 #[test]
15201 fn live_generation_cannot_be_overwritten() {
15202 let mut tracker = OptiForkGenerationTracker::default();
15203 let g0 = tracker.reserve().unwrap();
15204 let g1 = tracker.reserve().unwrap();
15205 let err = tracker.reserve().unwrap_err().to_string();
15206 assert!(
15207 err.contains("still owns generation 0"),
15208 "unexpected error: {err}"
15209 );
15210 tracker.retire(g0).unwrap();
15211 let g2 = tracker.reserve().unwrap();
15212 assert_eq!((g2.id, g2.slot), (2, 0));
15213 tracker.retire(g1).unwrap();
15214 tracker.retire(g2).unwrap();
15215 }
15216
15217 #[test]
15218 fn teardown_rejects_a_stale_generation_tag() {
15219 let mut tracker = OptiForkGenerationTracker::default();
15220 let g0 = tracker.reserve().unwrap();
15221 tracker.retire(g0).unwrap();
15222 let err = tracker.retire(g0).unwrap_err().to_string();
15223 assert!(err.contains("teardown mismatch"), "unexpected error: {err}");
15224 }
15225}
15226
15227#[cfg(test)]
15228mod draft_graph_fallback_tests {
15229 use super::DraftGraphFallback;
15230
15231 /// The Q2 contract, part (a): a fallback flip is LOUD — exactly once per flip.
15232 #[test]
15233 fn flip_is_loud_once_and_memoized_after() {
15234 let mut f = DraftGraphFallback::default();
15235 let line = f
15236 .mark_greedy("out of memory")
15237 .expect("first flip must return the warn line");
15238 assert!(
15239 line.contains("WARN"),
15240 "flip line must be warn-level: {line}"
15241 );
15242 assert!(
15243 line.contains("out of memory"),
15244 "flip line must carry the reason: {line}"
15245 );
15246 assert!(f.greedy_failed());
15247 // re-marking an already-failed graph is the memoization: quiet, still failed.
15248 assert!(f.mark_greedy("out of memory").is_none());
15249 assert!(f.greedy_failed());
15250 // the two graphs' flags are independent (greedy flip leaves sampled capturable).
15251 assert!(!f.sampled_failed());
15252 let line_s = f
15253 .mark_sampled("capture unsupported")
15254 .expect("sampled flip is its own flip");
15255 assert!(
15256 line_s.contains("sampled"),
15257 "sampled flip names itself: {line_s}"
15258 );
15259 assert!(f.mark_sampled("capture unsupported").is_none());
15260 }
15261
15262 /// The Q2 contract, part (b): resume-from-pool RESETS both flags (fresh capture chance),
15263 /// and says so exactly when there was something to reset.
15264 #[test]
15265 fn reset_on_resume_clears_flags_and_logs_once() {
15266 let mut f = DraftGraphFallback::default();
15267 // clean session: resume is silent, nothing to reset.
15268 assert!(f.reset_on_resume().is_none());
15269 f.mark_greedy("oom").unwrap();
15270 f.mark_sampled("oom").unwrap();
15271 let note = f
15272 .reset_on_resume()
15273 .expect("a set flag must produce the reset note");
15274 assert!(
15275 note.contains("greedy+sampled"),
15276 "note names what was reset: {note}"
15277 );
15278 assert!(
15279 !f.greedy_failed() && !f.sampled_failed(),
15280 "both flags cleared"
15281 );
15282 // and the NEXT failure after a reset is a fresh flip — loud again.
15283 assert!(f.mark_greedy("oom again").is_some());
15284 let note2 = f.reset_on_resume().expect("greedy-only reset");
15285 assert!(note2.contains("(greedy)"), "single-flag note: {note2}");
15286 }
15287
15288 /// Shape-change clears (dmask realloc / mask-shape mismatch / s_key change) stay silent —
15289 /// they precede a fresh capture attempt whose own failure re-flips loudly.
15290 #[test]
15291 fn shape_change_clears_are_silent() {
15292 let mut f = DraftGraphFallback::default();
15293 f.mark_greedy("oom").unwrap();
15294 f.clear_greedy();
15295 assert!(!f.greedy_failed());
15296 f.mark_sampled("oom").unwrap();
15297 f.clear_sampled();
15298 assert!(!f.sampled_failed());
15299 // after a silent clear there is nothing left for resume to report.
15300 assert!(f.reset_on_resume().is_none());
15301 }
15302}
15303
15304/// SAMPLED DRAFT-GRAPH KEY (lane/graph-s-key-exactness-20260819).
15305///
15306/// These are the CPU teeth for an exactness bug whose live reproduction needs a GPU, a trunk, a
15307/// drafter and a two-turn session: the key itself. Every test below fails against the pre-fix key
15308/// `(seed, temp.to_bits(), k)` — `legacy_key` restates it so the collision is explicit rather
15309/// than remembered.
15310#[cfg(test)]
15311mod sampled_graph_key_tests {
15312 use super::{SampledGraphKey, debug_t_pred0};
15313
15314 /// The pre-fix key, verbatim: `let s_key = (sp_seed, sp_temp.to_bits(), k);`
15315 fn legacy_key(k: &SampledGraphKey) -> (u64, u32, usize) {
15316 (k.seed, k.temp_bits, k.k)
15317 }
15318
15319 fn pure_temp_key() -> SampledGraphKey {
15320 // temperature 1.0, filters off — today's serve default, the shape that parks a graph.
15321 SampledGraphKey::new(12345, 1.0, 3, 0, 1.0, 0.0, false)
15322 }
15323
15324 /// THE COLLISION. Two requests that differ ONLY in the truncation filters shared one key, so
15325 /// a parked pure-temp graph survived into a filtered request and the launch site launched it.
15326 #[test]
15327 fn vendor_filters_change_the_key() {
15328 let parked = pure_temp_key();
15329 // qwen3.8 generation_config.json — what the vendor-default flip makes the default shape.
15330 let vendor = SampledGraphKey::new(12345, 1.0, 3, 20, 0.95, 0.0, false);
15331 assert_eq!(
15332 legacy_key(&parked),
15333 legacy_key(&vendor),
15334 "pre-fix key collided: this is the bug, and the reason a test asserts on it",
15335 );
15336 assert_ne!(parked, vendor, "post-fix key must separate the two regimes");
15337 assert!(parked.pure_temp());
15338 assert!(!vendor.pure_temp());
15339 }
15340
15341 /// Each distribution-shaping field alone is enough to drop the parked graph.
15342 #[test]
15343 fn every_filter_field_is_keyed() {
15344 let base = pure_temp_key();
15345 for (what, other) in [
15346 (
15347 "top_k",
15348 SampledGraphKey::new(12345, 1.0, 3, 20, 1.0, 0.0, false),
15349 ),
15350 (
15351 "top_p",
15352 SampledGraphKey::new(12345, 1.0, 3, 0, 0.95, 0.0, false),
15353 ),
15354 (
15355 "min_p",
15356 SampledGraphKey::new(12345, 1.0, 3, 0, 1.0, 0.05, false),
15357 ),
15358 (
15359 "penalties",
15360 SampledGraphKey::new(12345, 1.0, 3, 0, 1.0, 0.0, true),
15361 ),
15362 ] {
15363 assert_ne!(base, other, "{what} must be part of the key");
15364 assert!(!other.pure_temp(), "{what} leaves the pure-temp regime");
15365 assert_eq!(
15366 legacy_key(&base),
15367 legacy_key(&other),
15368 "{what} was invisible to the pre-fix key",
15369 );
15370 }
15371 }
15372
15373 /// The baked constants stay keyed (this half was always right — regression cover for it).
15374 #[test]
15375 fn baked_constants_stay_keyed() {
15376 let base = pure_temp_key();
15377 assert_ne!(
15378 base,
15379 SampledGraphKey::new(999, 1.0, 3, 0, 1.0, 0.0, false),
15380 "seed"
15381 );
15382 assert_ne!(
15383 base,
15384 SampledGraphKey::new(12345, 0.7, 3, 0, 1.0, 0.0, false),
15385 "temp"
15386 );
15387 assert_ne!(
15388 base,
15389 SampledGraphKey::new(12345, 1.0, 4, 0, 1.0, 0.0, false),
15390 "k"
15391 );
15392 // bitwise on temperature: 0.7f32 vs the same value re-derived must NOT differ.
15393 assert_eq!(
15394 SampledGraphKey::new(1, 0.7, 3, 0, 1.0, 0.0, false),
15395 SampledGraphKey::new(1, 7.0 / 10.0, 3, 0, 1.0, 0.0, false),
15396 );
15397 }
15398
15399 /// THE LOAD-BEARING HALF OF THE SEED DECISION (lane/session-resume-sampler-predicate-
15400 /// 20260820). The whole-session resume predicate deliberately does NOT compare `seed`: an
15401 /// omitted serve `seed` draws fresh per-request entropy, so comparing it would refuse every
15402 /// seed-omitting sampled conversation. That is only sound because the one piece of parked state
15403 /// that BAKES the seed — this graph — is re-keyed on it, so a seed change drops and recaptures.
15404 ///
15405 /// This test is the other end of that argument, asserted here rather than remembered in a
15406 /// comment: if a future change dropped `seed` from the key, the resume predicate's exclusion
15407 /// would silently become the unsound thing it is documented not to be.
15408 /// (Paired with `seed_alone_does_not_refuse` in `memra-sampling`.)
15409 #[test]
15410 fn seed_alone_still_rekeys_the_draft_graph() {
15411 let parked = pure_temp_key();
15412 let reseeded = SampledGraphKey::new(999, 1.0, 3, 0, 1.0, 0.0, false);
15413 assert_ne!(
15414 parked, reseeded,
15415 "a seed-only change MUST drop the parked sampled graph — the resume predicate's \
15416 decision not to compare seed rests on exactly this",
15417 );
15418 // Same regime on both sides: the drop is a recapture, not a fall to the eager chain
15419 // because of a filter difference.
15420 assert!(parked.pure_temp() && reseeded.pure_temp());
15421 }
15422
15423 /// `pure_temp()` is the capture guard's predicate, computed from the key so the two cannot
15424 /// drift. The equality below is the invariant the launch-site guard asserts: identical keys
15425 /// agree on the regime, so a graph that survives the drop is legal to launch.
15426 #[test]
15427 fn equal_keys_agree_on_the_regime() {
15428 let a = SampledGraphKey::new(7, 0.8, 3, 20, 0.95, 0.0, false);
15429 let b = SampledGraphKey::new(7, 0.8, 3, 20, 0.95, 0.0, false);
15430 assert_eq!(a, b);
15431 assert_eq!(a.pure_temp(), b.pure_temp());
15432 // top_p slightly above 1.0 (a client sending 1.0 exactly, or an operator default) is
15433 // still the unfiltered regime, matching the original `sp.top_p >= 1.0` test.
15434 assert!(SampledGraphKey::new(7, 0.8, 3, 0, 1.0, 0.0, false).pure_temp());
15435 assert!(SampledGraphKey::new(7, 0.8, 3, 0, 1.5, -1.0, false).pure_temp());
15436 }
15437
15438 /// MEMRA_DEBUG_SPEC on a SAMPLED spec request past round 0: the print must render without
15439 /// indexing the empty greedy `preds` vector (it panicked the GPU worker before this lane).
15440 #[test]
15441 fn debug_print_survives_the_sampled_arm() {
15442 // round >= 1 with a pending bonus == base 1, sampled == `preds` empty.
15443 assert_eq!(debug_t_pred0(true, 1, 4242, &[]), "n/a");
15444 assert_eq!(debug_t_pred0(true, 2, 4242, &[]), "n/a");
15445 // round 0 without a pending bonus still reports last_pred, in both arms.
15446 assert_eq!(debug_t_pred0(true, 0, 4242, &[]), "4242");
15447 assert_eq!(debug_t_pred0(false, 0, 4242, &[7, 8]), "4242");
15448 // greedy keeps the real prediction it always printed.
15449 assert_eq!(debug_t_pred0(false, 1, 4242, &[7, 8]), "7");
15450 assert_eq!(debug_t_pred0(false, 2, 4242, &[7, 8]), "8");
15451 }
15452}