memra_engine/glm_spec.rs
1//! glm5_next T-parallel speculative verify: the acceptance/rollback machinery that turns the
2//! native MTP draft head (`mtp_head_forward_mla_cached`, lane/glm5-mtp-remint) into served
3//! speculative decoding on the HyperConnections trunk.
4//!
5//! THE WALK-REUSE DECISION, stated once and load-bearing everywhere below: the verify of K
6//! drafted tokens is a t=K+1 walk over the SAME trunk, and it rides the BATCHED-DECODE walk's
7//! kernel classes (`hyper_batch_range_decode`, lane/glm53-batched-decode), NOT the prefill /
8//! prime walk. Why:
9//!
10//! * The batched walk's row-parallel ops are per-row BIT-EXACT vs the isolated t=1 decode
11//! step at every width in `1..=hyper_batch_cap()` (= `PRIME_MIN_T - 1` = 15, the shexp
12//! decode-exact knee, measured in `batched-decode-gate/31-KNEE-b16-forced.log`): the hc
13//! glue is block-per-token, the hc-mix GEMM runs per-row m=1 (`hyper::pre_exact`, the
14//! lt_ndep law), the MoE FFN's router/experts/shexp are per-row programs below
15//! `PRIME_MIN_T`, and the lm_head is `matmul_decode_exact`. That per-row exactness is
16//! what makes spec-vs-plain BYTE IDENTITY achievable at all.
17//! * The prime walk is a different numeric class: `hyper::pre` batches the mix GEMM
18//! (cuBLASLt n-dependent reduction), the FFN takes the prefill dispatch
19//! (`moe_ffn_il_prefill` / grouped GEMMs), and KDA's prefill conv arm reads the chunk
20//! window instead of the decode ring program. None of those are per-row bit-identical
21//! to the t=1 decode chain, so a prime-walk verify could never pass a byte-identity
22//! gate against plain decode.
23//! * The knee bounds K: K+1 <= 15, i.e. K <= 14 — the DFlash2 probe's K<=7 drafter and
24//! upstream's 5/7-draft MTP configs fit with margin.
25//!
26//! The ONE difference from `hyper_batch_range_decode`: the batched walk runs B independent
27//! sessions (each row -> its OWN cache at its own single position); the verify walk runs
28//! K+1 SEQUENTIAL positions of ONE session, so the mixers chain state row -> row through the
29//! one cache (KDA: the t=1 recurrent step per row, exactly the serving decode program; MLA:
30//! the t=1 `mla_attn_cached` append+attend per row, so row r attends rows `0..pos0+r+1` —
31//! causal within the drafted block by construction). This is the "hyper rows-walk with
32//! causal verify appends" the batched-decode lane's standing refusals named as missing.
33//!
34//! VERIFY-ROW BATCHING (lane/glm5-verify-batch, 2026-08-30 — the flip re-battery's named
35//! flip condition: "the verify walk must stop paying one plain-step per row", ~24-26 ms/row
36//! measured): `MEMRA_GLM5_VERIFY_BATCH` (default ON) restructures the mixer walk PER LAYER
37//! while keeping the sequential contract exactly where the math demands it. Each KDA layer
38//! runs ONE t=K+1 `kda_core` call — projections/gates/conv batched through the decode-exact
39//! matmul classes (`matmul_rows_exact`; the bf16 tcols twin reads each weight ONCE for all
40//! t rows) with the recurrence SEQUENTIAL INSIDE one `memra_kda_scan_s128` launch (the
41//! in-kernel T-loop over register state IS the chained t=1 program). Each MLA layer runs
42//! ONE t=K+1 `mla_attn_cached_rows_exact` call (per-query causal kpool selection + gathered
43//! attention by construction). Rollback on the batched arm: conv ring = pre-round snapshot
44//! + re-roll(T=keep) over stolen raw rows; ssm = ONE scan replay at T=keep from the
45//! pre-round snapshot over the stolen batched inputs (`kda::KdaRowsStash`). `0` = the
46//! per-row walk below, byte-for-byte — the rollback seam. Per-row byte identity vs the
47//! plain tape is held by this file's standing batteries running the batched arm, plus the
48//! kernel bit-gates in `tests/glm5_verify_batch_gpu.rs`.
49//!
50//! ROLLBACK (the hard part, per the engine survey's upstream reading — vLLM keeps
51//! num_spec+1 KDA state columns and commits the last accepted one; SGLang's ReplaySSM keeps
52//! an input ring and replays the accepted prefix):
53//!
54//! * KDA recurrent state: SNAPSHOT + SCAN-INPUT REPLAY (lane/glm5-loop-port port 3 —
55//! the GdnStash/ReplaySSM diet this doc used to name as the follow-up, landed). The
56//! walk clones the resident ssm state ONCE per layer per round (before row 0) and
57//! STEALS each row's scan-input buffers (`kda::KdaScanInputs`, ~160 KB/row/layer,
58//! zero copies — the step allocated them either way); accept-j REBUILDS the state by
59//! re-issuing rows 0..=j's original t=1 `memra_kda_scan_s128` launches from the
60//! snapshot (`kda::kda_scan_replay`) — byte-identical to the retired per-row clone
61//! by construction, since each replay is the very launch that produced it. Full
62//! accept keeps the resident state, no work. Memory, stated: one glm5_next KDA state
63//! is 64 heads x 128 x 128 f32 = 4 MiB, x34 KDA layers = 136 MiB per round (was 136
64//! MiB PER COLUMN — ~0.95 GiB of transient clones at K=7, retired to ~136 MiB + K x
65//! ~160 KB x 34 of stolen stash). The conv ring stays per-row cloned (288 KiB/row/
66//! layer, 1.4% of the ssm plane — not worth a replay arm).
67//! * MLA latent rows: TRUNCATE (append-only, position-addressed): `len = snap + keep`,
68//! device mirror in lock-step.
69//! * kpool index planes: the tail ring drains IN-CALL (`mla_kpool_indices`), so pool keys
70//! over drafted rows may FINALIZE during the verify walk. Rollback clamps
71//! `index_pools_ready` via `truncate_index_pool_keys` (the clamp the field's own doc
72//! demands of "the same code that shortens len"); the next call rebuilds keys for the
73//! re-appended rows from `[ready, len/pool)`. The residency tripwire in
74//! `mla_kpool_indices` fails LOUDLY if any rewind forgets this — that tripwire is the
75//! red arm of the rollback gate.
76//! * MTP draft plane (il = n_trunk): len reset by the loop (LANE.md contract: "rollback =
77//! plane len reset"), same pool-key clamp.
78//!
79//! Gate: `tests/glm5_tparallel_verify_gpu.rs` — accept-j-then-continue byte identity vs the
80//! never-drafted sequential path for every j in 0..=K, red-proven with a stale-KDA-state
81//! mutation and a pool-key-finalized-past-j mutation; plus end-to-end spec-vs-plain greedy
82//! tape identity at K=1..7 with forced-rejection positions, red-proven by disabling
83//! rollback. The SERVED shape (worker-sized bursts over one [`Glm5SpecSession`], state
84//! carried across burst boundaries, sampled twin, EOS, receipt log red/green) is gated by
85//! `tests/glm5_spec_session_gpu.rs`.
86//!
87//! SAMPLED ACCEPTANCE (landed, lane/glm5-spec-routing 2026-08-30): greedy
88//! longest-matching-prefix below is the byte-deterministic instrument
89//! (greedy-is-the-instrument law). The sampled arm applies memra's existing spec sampled
90//! contract — `spec::SpecSampling` with the rejection-sampling accept walk
91//! (`u_j * q_j(x_j) < p_j(x_j)`, host Philox4x32-10 stream tag 0xFFFF_FFFE via
92//! `spec::host_u01`, residual resampling on the first rejection) — the same contract
93//! `generate_spec_inner2`'s MEMRA_SPEC_TEMP>0 route and the dspark sampled-admission walk
94//! consume. It plugs in at exactly one seam (`glm5_sampled_accept`, the accept rule over
95//! the verify logit rows, with the draft chain drawn from the SAME filtered distribution
96//! the q gather reads); nothing in the walk or the rollback changes. Philox counters live
97//! ON the [`Glm5SpecSession`] so randomness never repeats across serve bursts (the
98//! session-continuity law; burst-split invariance is pinned by `glm5_spec_session_gpu`).
99//! PENALIZED sampled requests are refused — no penalty arm yet; the worker keeps them on
100//! the plain path.
101//!
102//! FR-SPEC VOCAB MASKING (owner addition, 2026-08-30 — the house spec recipe, the q38 way):
103//! the loop consumes the existing `MEMRA_FRSPEC_TRIM` contract, no new flag. The loader
104//! already reaches this head: the trim match in `hybrid.rs` consumes the embedded head
105//! regardless of arch, `frspec_trim_own_head_name(n_trunk)` misses (glm5_next ships no
106//! private MTP lm_head) and the gather falls back to the trunk `output.weight` /
107//! `token_embd.weight` — which for glm5_next is EXACT BY CONTRACT, not merely by tying:
108//! the MTP block projects through the trunk lm_head (LANE.md; the draft gate pins
109//! `shared_head_head.is_none()` untrimmed). With a ranks artifact loaded,
110//! `mtp_head_forward_mla_cached` already projects through the gathered rows (its head is
111//! `shared_head_head.unwrap_or(trunk)`), so the DRAFT logits arrive `[n_ranks]`; this
112//! loop's seam is the remap: every draft argmax is a RANK id and maps through `d2t` back
113//! to the true vocab BEFORE it is drafted, chained (e_tok drives the embedding gather),
114//! or verified. THE VERIFY WALK STAYS FULL-VOCAB AND UNTOUCHED — a trimmed draft can only
115//! change WHICH tokens get drafted, never how they verify; that invariant is the whole
116//! design (q38's measured skipped-remap defect was 0/248 acceptance with every exactness
117//! gate green — silent, which is why the gate below makes it loud). Rank artifacts for
118//! glm5 are an INPUT DEPENDENCY: the corpus mint (SXC pools through GLM's tokenizer,
119//! per traffic class, the q38 plain-text format) is the owner's CPU-only lane; the
120//! self-trim d2t arm needs no external artifact and lands here first.
121//!
122//! PPN (lane/glm5-ppn-verify, 2026-08-30): the verify walk owns its stage split exactly as
123//! the batched decode walk does — `glm5_verify_rows_ppn` mirrors
124//! `decode_step_batch_hyper_ppn` (per-stage engine, per-stage pos_rows, ONE
125//! `[t, streams, n_embd]` boundary payload per cut; row chaining is per-LAYER through the
126//! one cache, so a straight layer-range split preserves it exactly). Rollback restores each
127//! stage's layers through that stage's engine on its stream; the MTP block, its latent
128//! plane and every draft-chain/accept-side op ride the LAST stage's engine
129//! (`glm5_head_engine` — where the loader puts the lm head and `pp::new_cache*` puts the
130//! trailing MTP plane), so the h_seed carrier never bounces devices. Gate:
131//! `glm5-spec-ppn-gate` (the tparallel battery under the split, stages=2 and 3, red-proven;
132//! the cross-device twin is the box arm). Worker admission bounds sharded placements to the
133//! GATED stage set (`glm5_sharded_placement_admits`, worker.rs) — everything else stays
134//! fail-closed by name.
135//!
136//! DRAFT SOURCE SEAM (lane/glm5-dflash-draft-src, 2026-08-30): the session's drafts come
137//! from ONE of two sources, selected at load and pinned for the session — everything from
138//! the verify walk on (accept, rollback, commit, receipts, K policy) is SHARED and
139//! source-blind, which is where the exactness invariant lives (a draft source can only
140//! move acceptance, never output):
141//!
142//! * `NativeMtp` (existing): the embedded NextN head chains K drafts through
143//! `mtp_head_forward_mla_cached`; requires `MEMRA_GLM5_MTP=1`.
144//! * `Dflash2` (`MEMRA_GLM5_DFLASH=<dir-or-hf-spec>`): the pinned
145//! incoai/GLM-5.3-Flash-DFlash2 block-diffusion drafter (owner holds WRITTEN APPROVAL
146//! from the DFlash2 owners, 2026-08-30, for use beyond probe/eval). It REUSES the
147//! shipped q38 DFlash2 machinery verbatim (`DflashDraft`: `ctx_features` ->
148//! `ingest_ctx` -> `forward_round` -> `dflash2_propose_*`, mask-fill harvest, selector
149//! walk); the ONE glm5-specific input is the drafter's measured feature contract — the
150//! STREAM-MEAN (`hc_contract`) of the COMPLETED trunk layer output at the drafter
151//! config's `target_layer_ids` (plan layers 5,14,24,33,42 on the real artifact), the
152//! exact definition the probe's `MEMRA_TRACE_LAYER_ROWS` capture seam banked 0.73
153//! acc@1 / 3.06 tokens-per-cycle against
154//! (research/glm53-flash-bringup-20260827/dflash2-probe-20260829/RECEIPTS.md). The
155//! features flow through [`crate::cache::HcTapSink`], a HOST sink filled by the hc
156//! prime walks and this file's verify walk (host-staged so a ppN split needs no
157//! cross-device tap plumbing; the drafter itself runs on the HEAD engine, where the
158//! trunk lm_head it projects through lives). THE NATIVE MTP HEAD IS NOT LOADED for
159//! this source (the q38 pattern — a full MoE trunk layer of VRAM back); the plan's
160//! trailing MTP cache plane still allocates (plan-structural, ~`ctx * latent_width`
161//! f32 per declared block — named cost, not forked). Sampled route: the drafter's
162//! selector proposal records its true q (`DsparkDraftSample::Selector`) and the accept
163//! rides the SAME `dspark_accept_sampled` rejection walk the q38 serve route ships,
164//! with this session's Philox counters (`uctr` selector/accept draws, `sctr`
165//! bonus/residual) so randomness never repeats across bursts. K is bounded by the
166//! drafter block (K <= block_size-1 = 7): the worker clamps, the burst refuses loudly.
167//! Selection receipt: boot logs `[glm5-spec] draft source = native-mtp` or
168//! `[glm5-spec] draft source = dflash2 @ <sha8>`; both flags off = plain serving
169//! (fail-closed warn). `MEMRA_GLM5_DFLASH_GATE_RED=tap-shift` is a GATE INSTRUMENT
170//! (never a serving flag): it shifts every tap layer +1 to red-prove that a wrong
171//! feature input collapses acceptance while the tape stays byte-identical.
172//!
173//! SERVING EXPOSURE (lane/glm5-spec-routing, 2026-08-30): `MEMRA_GLM5_SPEC` (default OFF,
174//! FLAGS.md row) is the ONE master flag — it routes `generate_spec` here for hc trunks
175//! with a loaded MTP head AND arms the worker route (`glm5_spec_capable` +
176//! `step_glm5_spec` driving [`Glm5SpecSession`] bursts). OFF = the named `refuse_hyper`
177//! refusal and zero `[glm5-spec]` log lines, byte-identical serving. The MTP_SPEC
178//! capability manifest remains deliberately UNEXTENDED — worker `mtp_spec_capable` stays
179//! false for glm5_next plans; the serving capability lives in its OWN manifest
180//! (`GLM5_SPEC`, execution_manifest.rs) whose table names exactly the glm5_next class, and
181//! a SEALED production bundle still fails closed until it banks a `glm5-spec.v1` rewrite
182//! receipt (the real-artifact qualification lane's job).
183
184use crate::Engine;
185use crate::cache::{Cache, HcTapSink};
186use crate::dflash::{DflashDraft, DflashKv, DsparkDraftSample};
187use crate::forward::argmax;
188use crate::hybrid::{HybridModel, Mixer};
189use crate::spec::SpecSampling;
190use crate::spec_phase::SpecPhaseNs;
191use cudarc::driver::CudaSlice;
192
193type Res<T> = Result<T, Box<dyn std::error::Error>>;
194
195/// `MEMRA_GLM5_SPEC=1` routes `generate_spec` to the glm5 T-parallel loop. Default OFF:
196/// unset/0 keeps the standing `refuse_hyper` refusal, so serving is byte-identical to the
197/// pre-lane binary. Read once (worker chunk policies read their flags the same way).
198pub fn glm5_spec_on() -> bool {
199 use std::sync::OnceLock;
200 static ON: OnceLock<bool> = OnceLock::new();
201 *ON.get_or_init(|| std::env::var("MEMRA_GLM5_SPEC").as_deref() == Ok("1"))
202}
203
204// Per-burst phase attribution moved to `crate::spec_phase` (lane/glm5-extract-general):
205// the draft/verify/accept/roll/maint split is spec-family-generic. This loop consumes
206// `MEMRA_SPEC_TRACE` (glm5 alias `MEMRA_GLM5_SPEC_TRACE` stays honored) and passes its
207// own `[glm5-phase]` / `[glm5-phase-v]` tags so every banked receipt keeps its shape.
208
209/// `MEMRA_GLM5_VERIFY_BATCH` (default ON, lane/glm5-verify-batch): the per-LAYER batched
210/// mixer walk — one t=K+1 KDA call per layer (projections/conv/gates batched through the
211/// decode-exact classes, the recurrence sequential INSIDE one scan launch) and one
212/// t=K+1 rows-exact MLA call per layer, replacing the per-row mixer loop. `0` restores
213/// the per-row walk byte-for-byte — the rollback seam. Deliberate default (new-flags
214/// law): the walk only exists behind `MEMRA_GLM5_SPEC` (default OFF in prod), per-row
215/// byte identity is bit-gated on the rig (`glm5_tparallel_verify_gpu` +
216/// `glm5_verify_batch_gpu`), and the box re-battery A/Bs this seam in one build. Read
217/// PER CALL (the `MEMRA_KDA_FUSED_PROJ` per-call precedent) so gates drive both arms in
218/// one process; one env read per verify walk.
219pub fn glm5_verify_batch_on() -> bool {
220 std::env::var("MEMRA_GLM5_VERIFY_BATCH").as_deref() != Ok("0")
221}
222
223/// `MEMRA_GLM5_SPEC_PREFIX` (default OFF, lane/glm5-prefix-latent2 2026-09-01): the glm5
224/// spec x prefix-cache interplay — BOTH sides, one flag (the `MEMRA_DSPARK_PREFIX_RESTORE`
225/// precedent): (capture) DFlash2-source sessions take a prompt-boundary capture at creation
226/// for the worker's deferred prefix publication, and (restore) the worker may convert a
227/// prefix hit into `glm5_spec_session_from_restored` instead of demoting it to the plain
228/// route. Requires `MEMRA_PREFIX_LATENT=1` too — a capture the worker's publisher would
229/// refuse (latent entries need the plane flag) is pure waste, so this predicate ANDs both
230/// env reads. DEFAULT OFF BY DESIGN (new-flags law): the restored-session program is
231/// unmeasured until the box battery banks restored-vs-cold byte identity on the
232/// continuation; unset restores the pre-lane posture exactly (spec sessions never capture,
233/// hits demote to plain). Read once per process.
234pub fn glm5_spec_prefix_on() -> bool {
235 use std::sync::OnceLock;
236 static ON: OnceLock<bool> = OnceLock::new();
237 *ON.get_or_init(|| {
238 let sp = std::env::var("MEMRA_GLM5_SPEC_PREFIX").as_deref() == Ok("1");
239 let pl = std::env::var("MEMRA_PREFIX_LATENT").as_deref() == Ok("1");
240 if sp && !pl {
241 // A mis-built recipe would otherwise only surface through the battery's
242 // receipt gates (PR #96 review round 2, minor) — say it at first read.
243 eprintln!(
244 "[glm5-spec] MEMRA_GLM5_SPEC_PREFIX=1 is INERT: it requires \
245 MEMRA_PREFIX_LATENT=1 (latent entries could never publish without it)"
246 );
247 }
248 sp && pl
249 })
250}
251
252/// `MEMRA_GLM5_SPEC_TP` (default OFF, lane/glm5-composition 2026-09-01): admit glm5 spec
253/// SESSIONS on a `MEMRA_GLM5_TP`-armed model. DEFAULT OFF BY DESIGN (new-flags law): the
254/// composition's verify/rollback wiring is rig-gated for correctness (per-rank KDA
255/// snapshot/replay, per-replica MLA latent truncation — `glm5-tp-gate` arms S*), but it has
256/// ZERO real-artifact receipts and the TP serving wiring is still the named box increment;
257/// an unmeasured composition does not default ON. `=1` lifts ONLY the session co-refusal —
258/// every other admission law holds (draft source required, batched verify walk required:
259/// the per-row rollback seam carries no TP arm and refuses by name). Read per session
260/// creation. Rollback seam: unset (the co-refusal is restored verbatim).
261pub fn glm5_spec_tp_on() -> bool {
262 std::env::var("MEMRA_GLM5_SPEC_TP").as_deref() == Ok("1")
263}
264
265/// `MEMRA_SPEC_PMIN`, honored by the glm5 loop (loop-port 2 — the step37 shipping family,
266/// `MEMRA_SPEC_PMIN=0.5 MEMRA_SPEC_PMIN0=1` is what step37 serves; NO new flag): stop the
267/// draft chain early when the drafter's confidence in its own pick drops below p_min.
268/// Native chain: p = the head's softmax confidence in its pick (the spec.rs `g_p`
269/// statistic, `prob_of_token_device`). DFlash2: q = the selector's recorded per-slot
270/// candidate-set confidence (`q_chosen`; T=1 twin on the greedy walk) — the owner's
271/// "take only high confidence offers" tau-slot form, truncated PRE-verify.
272/// Unset/0 = OFF (today's rounds, byte-identical). The VALUE is a per-model measurement
273/// (spec.rs bank: q27 PMIN=0.3 was -1.9% on one pack; step37 ships 0.5) — the box-B tau
274/// ladder prices glm5's.
275pub(crate) fn glm5_pmin() -> f32 {
276 use std::sync::OnceLock;
277 static P: OnceLock<f32> = OnceLock::new();
278 *P.get_or_init(|| {
279 std::env::var("MEMRA_SPEC_PMIN")
280 .ok()
281 .and_then(|v| v.parse().ok())
282 .unwrap_or(0.0)
283 })
284}
285
286/// `MEMRA_SPEC_PMIN0=1` (llama.cpp's draft gating, vendored via spec.rs): the p-min gate
287/// applies at slot 0 too, so a low-confidence round drafts NOTHING and the verify batch is
288/// just the anchor row — m=1 = a plain decode step. Always legal for glm5 (the anchor row
289/// exists every round). "llama's 35B win rides exactly this — draft acceptance 76% at mean
290/// len 2.5 because unpredictable stretches never pay draft+verify overhead" (spec.rs).
291pub(crate) fn glm5_pmin0() -> bool {
292 use std::sync::OnceLock;
293 static P: OnceLock<bool> = OnceLock::new();
294 *P.get_or_init(|| std::env::var("MEMRA_SPEC_PMIN0").as_deref() == Ok("1"))
295}
296
297/// MEMRA_SPEC_PMIN break semantics — hoisted to the shared K-policy surface
298/// ([`crate::spec::spec_conf_keep`], lane/glm5-extract-general); re-exported here so the
299/// glm5 gates and call sites keep their name.
300pub use crate::spec::spec_conf_keep as glm5_conf_keep;
301
302/// The loaded DFlash2 drafter (module doc, DRAFT SOURCE SEAM): the model-level half of the
303/// `Dflash2` draft source — weights loaded ONCE per model on the head engine (`hybrid.rs`,
304/// `MEMRA_GLM5_DFLASH`); per-session state lives in [`Glm5DraftState`].
305///
306/// HOISTED to the general seam ([`crate::dflash::DflashDrafter`], lane/glm5-extract2): the
307/// holder is `{ drafter weights, byte-identity pin }` with nothing glm5 in it. Re-exported
308/// here under its old name so glm5's call sites and gates keep the name they were written
309/// against, exactly as `glm5_conf_keep` does above.
310pub use crate::dflash::DflashDrafter as Glm5DflashDrafter;
311
312/// Per-session draft-source state (module doc, DRAFT SOURCE SEAM). Selected at session
313/// creation from the model's loaded sources and pinned for the session's lifetime.
314pub(crate) enum Glm5DraftState {
315 /// Embedded NextN head: state = the MTP latent plane + `Glm5SpecSession::pending`
316 /// (token, h_seed) pairs — the pre-seam program, byte-identical.
317 NativeMtp,
318 /// DFlash2 block-diffusion drafter: state = the drafter's own ctx-feature KV cache
319 /// plus host feature rows not yet ingested. Invariant at every round boundary:
320 /// `kv.len + pending.len()/(taps.len()*n_embd) == committed.len()` — the drafter's
321 /// context is exactly the committed tokens (the probe's `F_feat[new_lo:start]` walk).
322 Dflash2 {
323 kv: DflashKv,
324 /// Committed-position feature rows awaiting ingest, `[n, n_taps*n_embd]` host
325 /// (the prompt's prime taps at session start; each round's kept verify taps after).
326 pending: Vec<f32>,
327 /// Resolved tap layers (drafter config `target_layer_ids`, red-arm shift applied).
328 taps: Vec<usize>,
329 },
330}
331
332/// The retained q side of one round's draft chain — what the sampled accept walk consumes.
333/// Greedy rounds carry `None` (the accept is the byte-deterministic prefix walk).
334enum Glm5DraftQ {
335 None,
336 /// Native MTP chain: per-slot retained draft logits (rank space under a trim) + the
337 /// filtered stats of the distribution each draft was drawn from.
338 Mtp {
339 draft_idx: Vec<u32>,
340 draft_logits: Vec<CudaSlice<f32>>,
341 draft_stats: Vec<(f32, f32, f32)>,
342 },
343 /// DFlash2 selector proposal (the recorded candidate-set q) + the retained draft-logit
344 /// rows `dl` (`dspark_accept_sampled`'s buffer contract; unread on the Selector q path).
345 Selector {
346 prop: DsparkDraftSample,
347 dl: CudaSlice<f32>,
348 },
349}
350
351/// Resolve the drafter's tap layers against the trunk: the drafter config's
352/// `target_layer_ids` are memra PLAN layer indices whose COMPLETED output feeds the fc
353/// (the probe's capture convention: `MEMRA_TRACE_LAYER_ROWS_LAYERS=5,14,24,33,42` == the
354/// drafter's own `target_layer_ids`, asserted 1:1 in `score_dflash2.py`).
355/// `MEMRA_GLM5_DFLASH_GATE_RED=tap-shift` is the RED-ARM INSTRUMENT: every tap moves +1
356/// layer — deliberately wrong features whose acceptance collapse the gate asserts while
357/// the output tape stays byte-identical. Unknown values refuse loudly.
358///
359/// The resolution itself is the general seam ([`crate::dflash::resolve_tap_layers`],
360/// lane/glm5-extract2); what stays here is glm5's OWN red arm — the gate instrument reads its
361/// env, prints its `[glm5-spec]` tag, and hands the shift in as a parameter. Error bytes are
362/// unchanged ("glm5 DFlash2" is the `what` label).
363fn glm5_dflash_tap_layers(draft: &DflashDraft, n_trunk: usize) -> Res<Vec<usize>> {
364 let shift = match std::env::var("MEMRA_GLM5_DFLASH_GATE_RED").ok().as_deref() {
365 Some("tap-shift") => {
366 eprintln!(
367 "[glm5-spec] RED-ARM tap-shift: drafter tap layers shifted +1 (gate \
368 instrument, never a serving flag)"
369 );
370 1
371 }
372 Some("") | None => 0,
373 Some(other) => {
374 return Err(format!(
375 "MEMRA_GLM5_DFLASH_GATE_RED={other:?}: unknown red arm (want tap-shift)"
376 )
377 .into());
378 }
379 };
380 Ok(crate::dflash::resolve_tap_layers(
381 &draft.cfg.target_layer_ids,
382 n_trunk,
383 shift,
384 "glm5 DFlash2",
385 )?)
386}
387
388/// Pre-round state checkpoint for one glm5 verify round. Captured by `glm5_verify_rows`
389/// BEFORE any row runs; consumed by `glm5_verify_rollback`.
390///
391/// Covers exactly the state planes a glm5_next trunk mutates in a verify round:
392/// - `latent_len`: per-layer MLA latent length at round start (rollback = truncate).
393/// - KDA state (loop-port 3, the module doc's GdnStash/ReplaySSM diet LANDED): the old
394/// per-row (conv, ssm) column clones — 4 MiB x 34 layers per COLUMN, ~0.95 GiB of
395/// transient at K=7 — are replaced by
396/// * `kda_ssm_snap[il]`: ONE recurrent-state clone per layer per round (the state
397/// BEFORE row 0),
398/// * `kda_scan_stash[il][r]`: row r's scan-input buffers, STOLEN from the step (zero
399/// copies, ~160 KB/row/layer — `kda::KdaScanInputs`), rows `0..t-1` except the last
400/// (`keep == t` needs no restore, so row `t-1` is never a replay target),
401/// * `kda_conv_cols[il][r]`: the conv ring stays PER-ROW CLONED (288 KiB, 1.4% of the
402/// ssm plane it rode beside — not worth a replay arm).
403///
404/// Partial-accept rollback REPLAYS rows `0..keep` from the snapshot
405/// (`kda::kda_scan_replay`): each replay is the original t=1 scan launch re-issued over
406/// the very buffers that row consumed, so the rebuilt state is byte-identical to the
407/// clone it replaces by construction. Under a ppN split every clone/stash lives on its
408/// layer's OWNING stage engine; rollback restores through the same per-stage seam.
409/// - `pos`: `cache.pos` at round start.
410///
411/// glm5_next has no Full/Linear trunk mixers (the walk refuses them by name), so `kv`,
412/// `tp_kv` and GDN stashes have no arm here — growing one is a deliberate extension with
413/// its own gate, not a silent default.
414pub struct Glm5VerifyCkpt {
415 pos: usize,
416 latent_len: Vec<Option<usize>>,
417 /// Per-row conv-ring clones, rows `0..t-1` except the last (doc above). PER-ROW walk
418 /// only (`MEMRA_GLM5_VERIFY_BATCH=0`); the batched walk fills `kda_rows` instead.
419 kda_conv_cols: Vec<Option<Vec<CudaSlice<f32>>>>,
420 /// The recurrent state BEFORE row 0, one clone per KDA layer per round (doc above).
421 /// BOTH walks fill this — it is the batched replay's scan base too.
422 kda_ssm_snap: Vec<Option<CudaSlice<f32>>>,
423 /// Stolen per-row scan inputs, rows `0..t-1` except the last (doc above). PER-ROW
424 /// walk only.
425 kda_scan_stash: Vec<Option<Vec<crate::kda::KdaScanInputs>>>,
426 /// BATCHED walk (lane/glm5-verify-batch): one [`crate::kda::KdaRowsStash`] per KDA
427 /// layer per round — ring snapshot + stolen raw conv rows + stolen batched scan
428 /// inputs; rollback re-rolls the ring and replays the scan ONCE at T=keep.
429 kda_rows: Vec<Option<crate::kda::KdaRowsStash>>,
430 /// glm5 TP composition (lane/glm5-composition): per-rank rollback material of each
431 /// SHARDED KDA layer's batched verify call — the rank-indexed twin of
432 /// (`kda_ssm_snap`, `kda_rows`), restored through each rank's own engine. `None` on
433 /// every unsharded layer.
434 kda_tp: Vec<Option<crate::glm5_tp::Glm5TpKdaVerifyStash>>,
435 /// Row count of the walk that filled this ckpt; rollback validates `keep` against it.
436 rows: usize,
437}
438
439impl Glm5VerifyCkpt {
440 /// GATE RECEIPT (wiring anchor, not a serving surface): how many KDA layers filled
441 /// the BATCHED rows stash vs the PER-ROW column stash — the flag A/B gate asserts
442 /// the arm it set actually ran (wiring-assertions-match-prose law: anchor on the
443 /// invocation's artifact, never the log prose).
444 pub fn kda_stash_kinds(&self) -> (usize, usize) {
445 (
446 self.kda_rows.iter().filter(|s| s.is_some()).count(),
447 self.kda_conv_cols.iter().filter(|s| s.is_some()).count(),
448 )
449 }
450}
451
452/// Position buffers for one verify walk range (per stage engine under a split — the
453/// per-stage pos_d law): `all` = the `[t]` vector the BATCHED per-layer mixer calls
454/// consume; `rows` = the per-row single-position buffers of the per-row arm, built only
455/// when that arm can run (flag off) — the batched arm never reads them.
456struct Glm5VerifyPos {
457 pos0: usize,
458 t: usize,
459 all: CudaSlice<i32>,
460 rows: Vec<CudaSlice<i32>>,
461}
462
463impl Glm5VerifyPos {
464 fn new(e: &Engine, pos0: usize, t: usize) -> Res<Self> {
465 let v: Vec<i32> = (0..t as i32).map(|r| pos0 as i32 + r).collect();
466 let all = e.htod_i32(&v)?;
467 let rows = if glm5_verify_batch_on() && t > 1 {
468 Vec::new()
469 } else {
470 (0..t)
471 .map(|r| e.htod_i32(&[(pos0 + r) as i32]))
472 .collect::<Result<_, _>>()?
473 };
474 Ok(Self { pos0, t, all, rows })
475 }
476}
477
478impl HybridModel {
479 /// THE T-PARALLEL VERIFY WALK: score `tokens` (row 0 = the last committed token, rows
480 /// 1..t = the K drafted tokens) in ONE forward over the hc trunk at positions
481 /// `cache.pos .. cache.pos + t`, in the batched-decode kernel classes (module doc).
482 ///
483 /// Returns `(logits [t, n_vocab] device, collapsed [t, n_embd] device, ckpt)`:
484 /// `logits` row r is bit-identical to the plain `decode_step_hyper` logits after
485 /// consuming `tokens[r]` at that position (the gate's bar); `collapsed` row r is the
486 /// pre-output_norm hidden — the MTP `h_seed` for position `cache.pos + r`.
487 ///
488 /// State effects: every trunk MLA plane appends `t` rows; every trunk KDA state
489 /// advances `t` steps (per-step columns stashed in the ckpt); `cache.pos` is NOT moved
490 /// (rollback owns it). The MTP block's plane (il = n_trunk) is untouched.
491 pub fn glm5_verify_rows(
492 &self,
493 e: &Engine,
494 tokens: &[u32],
495 cache: &mut Cache,
496 ) -> Res<(CudaSlice<f32>, CudaSlice<f32>, Glm5VerifyCkpt)> {
497 let topology = *self
498 .hyper
499 .as_ref()
500 .ok_or("glm5_verify_rows on a model with no HyperConnections topology")?;
501 let t = tokens.len();
502 let cap = Self::hyper_batch_cap();
503 if t == 0 {
504 return Err("glm5_verify_rows: empty verify row set".into());
505 }
506 if t > cap {
507 return Err(format!(
508 "glm5_verify_rows: t={t} > cap {cap} — at t >= PRIME_MIN_T (16) the MoE \
509 shared-expert trio crosses off the decode-exact class (the batched-decode \
510 gate's measured B=16 knee), so per-row bit-identity vs plain decode breaks. \
511 K <= cap-1 drafts per round"
512 )
513 .into());
514 }
515 let mut any_sharded = false;
516 for (il, layer) in self.layers.iter().enumerate() {
517 match &layer.mixer {
518 Mixer::Kda(la) => any_sharded |= la.tp.is_some(),
519 Mixer::Mla(mla) => any_sharded |= mla.tp.is_some(),
520 _ => {
521 return Err(format!(
522 "glm5_verify_rows: trunk layer {il} is not a KDA or MLA mixer — the \
523 rollback contract below is built and gated for glm5_next's two state \
524 classes only; a Full/Linear arm needs its own ckpt plane and gate"
525 )
526 .into());
527 }
528 }
529 }
530 // spec x TP composition (lane/glm5-composition): the per-row walk carries no TP
531 // rollback arm — a sharded trunk demands the BATCHED walk at t > 1 (t = 1 rounds
532 // ride the TP decode walk below; full accept is the only legal outcome there).
533 if any_sharded && t > 1 && !glm5_verify_batch_on() {
534 return Err(
535 "glm5_verify_rows: the trunk is glm5-TP-SHARDED and MEMRA_GLM5_VERIFY_BATCH=0 — \
536 the per-row rollback seam carries no TP arm; the spec x TP composition \
537 requires the batched verify walk (unset MEMRA_GLM5_VERIFY_BATCH or run \
538 without the TP door)"
539 .into(),
540 );
541 }
542
543 let n_embd = self.cfg.n_embd as usize;
544 let pos0 = cache.pos;
545
546 // Ckpt BEFORE any state moves.
547 let mut ckpt = Glm5VerifyCkpt {
548 pos: pos0,
549 latent_len: cache
550 .latent
551 .iter()
552 .take(self.layers.len())
553 .map(|plane| plane.as_ref().map(|plane| plane.len))
554 .collect(),
555 kda_conv_cols: (0..self.layers.len()).map(|_| None).collect(),
556 kda_ssm_snap: (0..self.layers.len()).map(|_| None).collect(),
557 kda_scan_stash: (0..self.layers.len()).map(|_| None).collect(),
558 kda_rows: (0..self.layers.len()).map(|_| None).collect(),
559 kda_tp: (0..self.layers.len()).map(|_| None).collect(),
560 rows: t,
561 };
562
563 // ppN door — the verify walk owns its stage split exactly as the batched decode
564 // walk does (`decode_step_batch_hyper_ppn`, decode_batch.rs). Loud refusal on an
565 // unqualified pipeline rewrite, never a single-engine walk over stage-sharded
566 // weights.
567 if let Some(fence) = crate::pp::pp_cuts(self.layers.len()) {
568 if !self.rewrite_allowed(memra_gguf::execution_manifest::RewriteSurface::Pipeline) {
569 return Err("pipeline rewrite is not qualified for this ModelPlan".into());
570 }
571 return self.glm5_verify_rows_ppn(e, tokens, cache, ckpt, &topology, &fence);
572 }
573
574 let pos = Glm5VerifyPos::new(e, pos0, t)?;
575 let embedded = e.htod(&self.embd.gather(n_embd, tokens))?;
576 let x = crate::hyper::expand(e, &topology, &embedded, t, n_embd)?;
577 let x = self.glm5_verify_range(
578 e,
579 &topology,
580 x,
581 0,
582 self.layers.len(),
583 &pos,
584 cache,
585 &mut ckpt,
586 )?;
587 let (logits, collapsed) = self.glm5_verify_head(e, &topology, &x, t)?;
588 Ok((logits, collapsed, ckpt))
589 }
590
591 /// One hc layer RANGE `[lo, hi)` of the verify walk — the body `glm5_verify_rows` ran
592 /// inline before the ppN twin landed, extracted so the unsplit walk and every pipeline
593 /// stage run the SAME code over their own range (the `hyper_range_decode` /
594 /// `decode_batch_layers` precedent: bit-identity between the arms is then structural,
595 /// not a coincidence of two maintained copies). At `lo=0, hi=n_layers` the launch
596 /// sequence is identical to the pre-extraction walk.
597 ///
598 /// KDA ckpt columns are cloned THROUGH `e` — under a split that is the owning stage's
599 /// engine, so each column lives on the device (and is ordered on the stream) that owns
600 /// its layer's state; `glm5_verify_rollback` restores through the same per-stage seam.
601 #[allow(clippy::too_many_arguments)]
602 // allow: the parameter list mirrors the range-walk call contract its siblings share
603 fn glm5_verify_range(
604 &self,
605 e: &Engine,
606 topology: &crate::hyper::HyperTopology,
607 mut x: CudaSlice<f32>,
608 lo: usize,
609 hi: usize,
610 pos: &Glm5VerifyPos,
611 cache: &mut Cache,
612 ckpt: &mut Glm5VerifyCkpt,
613 ) -> Res<CudaSlice<f32>> {
614 let t = pos.t;
615 let n_embd = self.cfg.n_embd as usize;
616 let eps = self.cfg.rms_eps;
617 // THE BATCHED MIXER ARM (lane/glm5-verify-batch, default ON): one t=K+1 call per
618 // layer per class instead of the per-row loop — KDA batches projections/conv/
619 // gates through the decode-exact classes with the recurrence sequential INSIDE
620 // one scan launch; MLA runs the SAME cached core at t rows on the rows-exact
621 // matmul classes (per-query causal selection + attention by construction).
622 // `0` = the per-row walk below, byte-for-byte (the rollback seam). Engagement is
623 // a receipt, announced once per process.
624 let batch = glm5_verify_batch_on() && t > 1;
625 {
626 static SAID: std::sync::Once = std::sync::Once::new();
627 SAID.call_once(|| {
628 if batch {
629 eprintln!(
630 "[glm5-spec] verify walk BATCHED per layer: kda=one t-call (scan \
631 sequential in-kernel), mla=rows-exact t-call, head=rows-exact, \
632 moe=pairs rows-call where qualified \
633 (MEMRA_GLM5_VERIFY_BATCH default ON)"
634 );
635 } else {
636 eprintln!("[glm5-spec] verify walk PER-ROW (MEMRA_GLM5_VERIFY_BATCH=0 or t=1)");
637 }
638 });
639 }
640 let trace_v = crate::spec_phase::spec_trace_level() >= 2;
641 // Sub-phase clock (trace level 2 only): drain the walking stream so the elapsed
642 // ns lands in the mixer-class bucket — shares, never walls.
643 let vclock = |on: bool| -> Option<std::time::Instant> {
644 on.then(|| {
645 let _ = e.stream().synchronize();
646 std::time::Instant::now()
647 })
648 };
649 for il in lo..hi {
650 let layer = &self.layers[il];
651 let hyper = layer.hyper.as_ref().ok_or_else(|| {
652 format!("layer {il} carries no hyper-connection weights under an hc plan")
653 })?;
654
655 let (y, mix) = crate::hyper::pre_exact(e, topology, &hyper.attn, &x, t, n_embd)?;
656 let mut h = e.uninit(t * n_embd)?;
657 e.rms_norm(&y, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
658 // The batched arm refuses per-layer only for an MLA layer WITHOUT the DSA
659 // indexer: the absorbed t>1 attention arm has no per-row bit-identity claim
660 // at this seam, so it stays on the per-row loop by name (glm5_next always
661 // carries the indexer, so this is a foreign-geometry guard, not a live path).
662 let layer_batched = batch
663 && match &layer.mixer {
664 Mixer::Kda(_) => true,
665 Mixer::Mla(mla) => mla.index.is_some(),
666 Mixer::Full(_) | Mixer::Linear(_) => unreachable!("refused at walk entry"),
667 };
668 let mixed = if layer_batched {
669 match &layer.mixer {
670 // spec x TP composition: sharded mixers ride the TP verify walks —
671 // per-rank batched rows calls, column-parallel-over-gather joins on
672 // the rows-exact classes, per-rank rollback stash into the ckpt.
673 Mixer::Kda(la) if la.tp.is_some() => {
674 let t0 = vclock(trace_v);
675 let mut scan_ns = 0u64;
676 let (out, stash) = crate::glm5_tp::kda_tp_verify_rows(
677 e,
678 la,
679 &h,
680 t,
681 eps,
682 cache,
683 il,
684 trace_v.then_some(&mut scan_ns),
685 )?;
686 ckpt.kda_tp[il] = Some(stash);
687 if let Some(t0) = t0 {
688 let _ = e.stream().synchronize();
689 use std::sync::atomic::Ordering;
690 crate::spec_phase::V_KDA_NS
691 .fetch_add(t0.elapsed().as_nanos() as u64, Ordering::Relaxed);
692 crate::spec_phase::V_KDA_SCAN_NS.fetch_add(scan_ns, Ordering::Relaxed);
693 }
694 out
695 }
696 Mixer::Mla(mla) if mla.tp.is_some() => {
697 let t0 = vclock(trace_v);
698 let out =
699 self.mla_tp_attn_cached(e, mla, &h, &pos.all, t, il, cache, true)?;
700 if let Some(t0) = t0 {
701 let _ = e.stream().synchronize();
702 use std::sync::atomic::Ordering;
703 crate::spec_phase::V_MLA_NS
704 .fetch_add(t0.elapsed().as_nanos() as u64, Ordering::Relaxed);
705 }
706 out
707 }
708 Mixer::Kda(la) => {
709 // Pre-round snapshot: ONE ssm clone per layer per round, BEFORE
710 // the batched call advances the resident state (ckpt doc; also
711 // the batched rollback's scan-replay base).
712 {
713 let rl = cache.recur[il]
714 .as_ref()
715 .ok_or("glm5 verify KDA layer has no recurrent state")?;
716 ckpt.kda_ssm_snap[il] = Some(e.clone_dtod(&rl.ssm_state)?);
717 }
718 let t0 = vclock(trace_v);
719 let mut scan_ns = 0u64;
720 let (out, stash) = crate::kda::kda_verify_rows_cached(
721 e,
722 la,
723 &h,
724 t,
725 eps,
726 cache,
727 il,
728 trace_v.then_some(&mut scan_ns),
729 )?;
730 ckpt.kda_rows[il] = Some(stash);
731 if let Some(t0) = t0 {
732 let _ = e.stream().synchronize();
733 use std::sync::atomic::Ordering;
734 crate::spec_phase::V_KDA_NS
735 .fetch_add(t0.elapsed().as_nanos() as u64, Ordering::Relaxed);
736 crate::spec_phase::V_KDA_SCAN_NS.fetch_add(scan_ns, Ordering::Relaxed);
737 }
738 out
739 }
740 Mixer::Mla(mla) => {
741 let t0 = vclock(trace_v);
742 let out =
743 self.mla_attn_cached_rows_exact(e, mla, &h, &pos.all, t, il, cache)?;
744 if let Some(t0) = t0 {
745 let _ = e.stream().synchronize();
746 use std::sync::atomic::Ordering;
747 crate::spec_phase::V_MLA_NS
748 .fetch_add(t0.elapsed().as_nanos() as u64, Ordering::Relaxed);
749 }
750 out
751 }
752 Mixer::Full(_) | Mixer::Linear(_) => unreachable!("refused at walk entry"),
753 }
754 } else {
755 // ---- PER-ROW mixer walk (the rollback seam; also t=1 rounds and the
756 // no-indexer MLA guard): row r's state input is row r-1's state output
757 // (KDA) / rows 0..pos0+r (MLA latent) — each row the SAME t=1 call its
758 // plain decode step makes.
759 // h_row is hoisted out of the row loop (loop-port 3): the mixer consumes
760 // it in stream order before the next row's overwrite, so ONE buffer per
761 // layer replaces t allocations (stream-ordered pool churn is the dsv4
762 // lesson).
763 let mut mixed = e.uninit(t * n_embd)?;
764 let mut h_row = e.uninit(n_embd)?;
765 #[allow(clippy::needless_range_loop)]
766 // allow: r is the sequential row cursor (slices h, offsets pos); iterating pos buffers would hide the row-chaining contract
767 for r in 0..t {
768 e.dtod_copy_view(&h.slice(r * n_embd..(r + 1) * n_embd), &mut h_row)?;
769 // The per-row position buffer: prebuilt when the per-row arm owns the
770 // walk; built on demand for the rare per-layer refusal under batch.
771 let pos_row: CudaSlice<i32>;
772 let pos_r = if let Some(p) = pos.rows.get(r) {
773 p
774 } else {
775 pos_row = e.htod_i32(&[(pos.pos0 + r) as i32])?;
776 &pos_row
777 };
778 let out_row = match &layer.mixer {
779 // spec x TP composition on the per-row arm. KDA shards reach
780 // here at t == 1 ONLY, asserted locally: the walk-entry guard is
781 // a FLAG check two frames up (MEMRA_GLM5_VERIFY_BATCH=0 at t>1
782 // refuses), and under the batched flag layer_batched is
783 // unconditionally true for KDA — but neither is a structural
784 // invariant of THIS arm (#80 review's latent-trap finding). A
785 // sharded NO-INDEXER MLA layer legally lands here at any t
786 // (append + truncate rollback covers every keep; foreign
787 // geometry, never a live glm5_next path).
788 Mixer::Kda(la) if la.tp.is_some() => {
789 if t > 1 {
790 return Err(format!(
791 "glm5 verify per-row arm reached a sharded KDA \
792 layer {il} at t={t}: no per-rank rollback stash \
793 exists on this arm (walk-entry guard bypassed?)"
794 )
795 .into());
796 }
797 crate::glm5_tp::kda_tp_cached(
798 e,
799 la,
800 &h_row,
801 1,
802 eps,
803 cache,
804 il,
805 crate::kda::ConvArm::Decode,
806 )?
807 }
808 Mixer::Mla(mla) if mla.tp.is_some() => {
809 self.mla_tp_attn_cached(e, mla, &h_row, pos_r, 1, il, cache, false)?
810 }
811 Mixer::Kda(la) => {
812 // Pre-round snapshot: ONE ssm clone per layer per round taken
813 // before row 0 mutates the resident state (loop-port 3).
814 if r == 0 && t > 1 {
815 let rl = cache.recur[il]
816 .as_ref()
817 .ok_or("glm5 verify KDA layer has no recurrent state")?;
818 ckpt.kda_ssm_snap[il] = Some(e.clone_dtod(&rl.ssm_state)?);
819 }
820 if r + 1 < t {
821 // Steal the row's scan inputs for the replay stash (zero
822 // copies); clone only the small conv ring per row.
823 let (out, inputs) = crate::kda::kda_decode_cached_stash(
824 e, la, &h_row, eps, cache, il,
825 )?;
826 let rl = cache.recur[il]
827 .as_ref()
828 .ok_or("glm5 verify KDA layer has no recurrent state")?;
829 ckpt.kda_conv_cols[il]
830 .get_or_insert_with(Vec::new)
831 .push(e.clone_dtod(&rl.conv_state)?);
832 ckpt.kda_scan_stash[il]
833 .get_or_insert_with(Vec::new)
834 .push(inputs);
835 out
836 } else {
837 crate::kda::kda_decode_cached(e, la, &h_row, eps, cache, il)?
838 }
839 }
840 Mixer::Mla(mla) => {
841 self.mla_attn_cached(e, mla, &h_row, pos_r, 1, il, cache)?
842 }
843 // Refused at entry; unreachable keeps the match total without a silent arm.
844 Mixer::Full(_) | Mixer::Linear(_) => unreachable!("refused at walk entry"),
845 };
846 e.copy_into(&mut mixed, r * n_embd, &out_row, n_embd)?;
847 }
848 mixed
849 };
850 x = crate::hyper::post(e, topology, &mixed, &x, &mix, t, n_embd)?;
851
852 let (y, mix) = crate::hyper::pre_exact(e, topology, &hyper.mlp, &x, t, n_embd)?;
853 let mut z = e.uninit(t * n_embd)?;
854 e.rms_norm(
855 &y,
856 layer.post_attn_norm.float_data(),
857 &mut z,
858 n_embd,
859 t,
860 eps,
861 )?;
862 // FFN branch: `batch` arms the pairs-shaped batched MoE across the t rows
863 // (lane/glm5-vrest — fail-closed inside to the byte-identical sequential
864 // loop); the =0 arm keeps the pre-lane per-(token,expert) class. Clocked
865 // into the vffn sub-bucket at trace level 2 (batched arm only, like vkda).
866 let t0 = vclock(trace_v && batch);
867 let ffn_out = self.hyper_ffn_branch_batch(e, layer, &z, t, il, batch)?;
868 if let Some(t0) = t0 {
869 let _ = e.stream().synchronize();
870 use std::sync::atomic::Ordering;
871 crate::spec_phase::V_FFN_NS
872 .fetch_add(t0.elapsed().as_nanos() as u64, Ordering::Relaxed);
873 }
874 x = crate::hyper::post(e, topology, &ffn_out, &x, &mix, t, n_embd)?;
875 // glm5 DFlash2 feature tap (module doc, DRAFT SOURCE SEAM): the verify rows'
876 // contracted completed-layer outputs are next round's drafter context.
877 self.glm5_hc_tap(e, cache, topology, il, &x, t)?;
878 }
879 Ok(x)
880 }
881
882 /// Write one tapped layer's CONTRACTED completed output into the armed
883 /// [`HcTapSink`] — the glm5 DFlash2 drafter's measured feature contract (stream-mean
884 /// over the hyper streams, the probe's `hc_contract` capture definition). Staged
885 /// through the WALKING engine `e` (the owning stage engine under a ppN split), so the
886 /// sink is placement-invariant. One Option check when unarmed; nothing else pays.
887 ///
888 /// Two staging arms (loop-port 1):
889 /// * `device_stage` (the verify-round sink): ONE async D2D into the slot's device
890 /// buffer — the walk never blocks; the round drains all slots post-walk in its
891 /// single sync point (`glm5_tap_drain`). Kills the five in-walk DtoHs the 3way
892 /// window priced into the 31.6 ms fixed round cost (map row #17).
893 /// * host-staged (prime sinks): the pre-port behavior — per-chunk DtoH, amortized
894 /// over the prime's >= 256-row chunks.
895 pub(crate) fn glm5_hc_tap(
896 &self,
897 e: &Engine,
898 cache: &mut Cache,
899 topology: &crate::hyper::HyperTopology,
900 il: usize,
901 x: &CudaSlice<f32>,
902 t: usize,
903 ) -> Res<()> {
904 let Some(sink) = cache.hc_taps.as_mut() else {
905 return Ok(());
906 };
907 let Some(slot) = sink.layer_ids.iter().position(|&l| l == il) else {
908 return Ok(());
909 };
910 let h = sink.hidden;
911 let n_taps = sink.layer_ids.len();
912 // Sink-relative row of this walk's row 0 (doc on `HcTapSink::origin`): fresh-prompt
913 // sinks have origin 0 and this is exactly the pre-field arithmetic; a suffix-prime
914 // sink is anchored at the restored boundary. A base below the origin is a caller
915 // bug (a walk over rows the sink does not cover) — refuse, never wrap.
916 let base = sink.base.checked_sub(sink.origin).ok_or_else(|| {
917 format!(
918 "hc tap base {} below sink origin {} (walk outside the sink's window)",
919 sink.base, sink.origin,
920 )
921 })?;
922 debug_assert!(
923 base + t <= sink.t,
924 "hc tap window {base}+{t} exceeds sink {}",
925 sink.t
926 );
927 let contracted = crate::hyper::contract_mean(e, topology, x, t, h)?;
928 if sink.device_stage {
929 // Lazy slot buffer on the WRITING engine (this layer always walks on one
930 // stage, so the buffer's device is stable for the sink's lifetime). Every
931 // walk row writes every tapped layer, so the buffer is fully covered by the
932 // walk that armed the sink.
933 if sink.dev[slot].is_none() {
934 sink.dev[slot] = Some(e.uninit(sink.t * h)?);
935 }
936 let buf = sink.dev[slot].as_mut().expect("just filled");
937 e.copy_into(buf, base * h, &contracted, t * h)?;
938 return Ok(());
939 }
940 let host = e.dtoh(&contracted)?;
941 for r in 0..t {
942 let dst = (base + r) * n_taps * h + slot * h;
943 sink.rows[dst..dst + h].copy_from_slice(&host[r * h..(r + 1) * h]);
944 }
945 Ok(())
946 }
947
948 /// Drain a device-staged tap sink into its host `rows` — the round's ONE post-walk
949 /// sync point for tap features (loop-port 1). Each slot reads back through its
950 /// layer's OWNING engine (the stage engine under a live split, the caller's engine
951 /// otherwise); the verify walk's terminal drain has already retired every stage's
952 /// writes (stream program order: the slot copy precedes its stage's TX, and the
953 /// TX-wait chain covers it transitively — the pp.rs multi-stream law).
954 fn glm5_tap_drain(&self, e: &Engine, sink: &mut HcTapSink) -> Res<()> {
955 if !sink.device_stage {
956 return Ok(());
957 }
958 let h = sink.hidden;
959 let n_taps = sink.layer_ids.len();
960 let split = match crate::pp::pp_cuts(self.layers.len()) {
961 Some(fence) if !crate::pp::pp2_streams_off() => {
962 Some((crate::pp::PpNRt::get(e)?, fence))
963 }
964 _ => None,
965 };
966 for slot in 0..n_taps {
967 let Some(buf) = sink.dev[slot].take() else {
968 continue;
969 };
970 let il = sink.layer_ids[slot];
971 let es = match split.as_ref() {
972 Some((rt, fence)) => {
973 let stage = fence
974 .windows(2)
975 .position(|w| il >= w[0] && il < w[1])
976 .ok_or_else(|| format!("tap layer {il} outside every stage range"))?;
977 rt.engine(stage, e)
978 }
979 None => e,
980 };
981 let host = es.dtoh(&buf)?;
982 for r in 0..sink.t {
983 let dst = r * n_taps * h + slot * h;
984 sink.rows[dst..dst + h].copy_from_slice(&host[r * h..(r + 1) * h]);
985 }
986 }
987 Ok(())
988 }
989
990 /// Trunk exit of the verify walk, the batched head's decode-exact form
991 /// (`hyper_batch_head_logits`), with the collapsed pre-output_norm rows kept — they are
992 /// the h_seeds the MTP head re-seeds from (LANE.md §A). Under a split this runs on the
993 /// LAST stage's engine, where the loader put `output_norm` + the lm head.
994 fn glm5_verify_head(
995 &self,
996 e: &Engine,
997 topology: &crate::hyper::HyperTopology,
998 x: &CudaSlice<f32>,
999 t: usize,
1000 ) -> Res<(CudaSlice<f32>, CudaSlice<f32>)> {
1001 let n_embd = self.cfg.n_embd as usize;
1002 let eps = self.cfg.rms_eps;
1003 let collapsed =
1004 crate::hyper::collapse(e, topology, self.hyper_head.as_ref(), x, t, n_embd)?;
1005 let mut hn = e.uninit(t * n_embd)?;
1006 e.rms_norm(
1007 &collapsed,
1008 self.output_norm.float_data(),
1009 &mut hn,
1010 n_embd,
1011 t,
1012 eps,
1013 )?;
1014 // Under the batched walk the lm head rides the rows-exact classes too (the bf16
1015 // tcols twin reads the 1.27 GB head ONCE per round instead of once per row);
1016 // per-row bits unchanged by contract, the tcols bit-gate holds it.
1017 let logits = if glm5_verify_batch_on() && t > 1 {
1018 e.matmul_rows_exact(&self.output, &hn, t)?
1019 } else {
1020 e.matmul_decode_exact(&self.output, &hn, t)?
1021 };
1022 Ok((logits, collapsed))
1023 }
1024
1025 /// ppN twin of the verify walk (lane/glm5-ppn-verify, 2026-08-30), mirroring
1026 /// `decode_step_batch_hyper_ppn` (decode_batch.rs): the t=K+1 rows walk as N stage
1027 /// subgraphs — per-stage engine, per-stage pos_rows uploads, ONE `[t, streams, n_embd]`
1028 /// boundary payload per fence cut. Row chaining is per-LAYER through the one cache
1029 /// (row r+1 at layer il depends only on row r at layer il), so a straight layer-range
1030 /// split preserves it exactly; no row ever crosses a boundary individually. Head +
1031 /// collapsed rows land on the LAST stage's engine — where the loader put the lm head
1032 /// and where `pp::new_cache*` places the MTP plane the re-seed feeds.
1033 ///
1034 /// DRAIN CONTRACT: this walk returns DEVICE buffers with no terminal dtoh (unlike its
1035 /// decode twins, whose epilogue reads back on the last stage's stream), so it owns the
1036 /// settle — the per-stage arm synchronizes the LAST stage's stream before returning.
1037 /// The TX-wait chain transitively covers every earlier stage (pp.rs multi-stream law),
1038 /// so the logits, the collapsed rows AND the ckpt's per-stage KDA columns are all safe
1039 /// for consumption from the caller's streams after this returns.
1040 #[allow(clippy::too_many_arguments)]
1041 // allow: the parameter list mirrors its decode twin's stage-walk contract
1042 fn glm5_verify_rows_ppn(
1043 &self,
1044 e: &Engine,
1045 tokens: &[u32],
1046 cache: &mut Cache,
1047 mut ckpt: Glm5VerifyCkpt,
1048 topology: &crate::hyper::HyperTopology,
1049 fence: &[usize],
1050 ) -> Res<(CudaSlice<f32>, CudaSlice<f32>, Glm5VerifyCkpt)> {
1051 let t = tokens.len();
1052 let n_embd = self.cfg.n_embd as usize;
1053 let pos0 = ckpt.pos;
1054 let payload = t * topology.streams * n_embd;
1055 // Position buffers through THIS stage's engine (the per-stage pos_d law:
1056 // allocated, consumed and freed on one stage's stream).
1057 let pos_on = |eng: &Engine| -> Res<Glm5VerifyPos> { Glm5VerifyPos::new(eng, pos0, t) };
1058
1059 // Same-stream seam (MEMRA_PP_STREAMS=0): one engine, boundary copies between
1060 // ranges — the shape every hc ppN walk uses for this knob.
1061 if crate::pp::pp2_streams_off() {
1062 let pos = pos_on(e)?;
1063 let embedded = e.htod(&self.embd.gather(n_embd, tokens))?;
1064 let mut x = crate::hyper::expand(e, topology, &embedded, t, n_embd)?;
1065 x =
1066 self.glm5_verify_range(e, topology, x, fence[0], fence[1], &pos, cache, &mut ckpt)?;
1067 for s in 1..fence.len() - 1 {
1068 let boundary_tx = e.clone_dtod(&x)?;
1069 let boundary_rx = e.clone_dtod(&boundary_tx)?;
1070 x = self.glm5_verify_range(
1071 e,
1072 topology,
1073 boundary_rx,
1074 fence[s],
1075 fence[s + 1],
1076 &pos,
1077 cache,
1078 &mut ckpt,
1079 )?;
1080 }
1081 let (logits, collapsed) = self.glm5_verify_head(e, topology, &x, t)?;
1082 return Ok((logits, collapsed, ckpt));
1083 }
1084
1085 let rt = crate::pp::PpNRt::get(e)?;
1086 let n_st = fence.len() - 1;
1087 assert_eq!(
1088 rt.n_stages(),
1089 n_st,
1090 "PpNRt stage count {} != fence stages {n_st}",
1091 rt.n_stages()
1092 );
1093 // #87 reverse publication (see decode_step_batch_ppn): order every stage stream
1094 // behind the caller before this body's first stage allocation.
1095 rt.fence_stages_behind(&e.stream())?;
1096
1097 // ---- STAGE 0: embed + expand (no weights) + layers [0, fence[1]) + TX ----
1098 let mut slot = {
1099 let _st0 = rt.enter(0);
1100 let e0 = rt.engine(0, e);
1101 let pos = pos_on(e0)?;
1102 let embedded = e0.htod(&self.embd.gather(n_embd, tokens))?;
1103 let x = crate::hyper::expand(e0, topology, &embedded, t, n_embd)?;
1104 let x = self
1105 .glm5_verify_range(e0, topology, x, fence[0], fence[1], &pos, cache, &mut ckpt)?;
1106 rt.tx(0, &x, payload)?
1107 };
1108
1109 // ---- MIDDLE STAGES: RX -> range -> TX ----
1110 for s in 1..n_st - 1 {
1111 let _st = rt.enter(s);
1112 let es = rt.engine(s, e);
1113 let pos = pos_on(es)?;
1114 let x = rt.rx(s - 1, slot, payload)?;
1115 let x = self.glm5_verify_range(
1116 es,
1117 topology,
1118 x,
1119 fence[s],
1120 fence[s + 1],
1121 &pos,
1122 cache,
1123 &mut ckpt,
1124 )?;
1125 slot = rt.tx(s, &x, payload)?;
1126 }
1127
1128 // ---- LAST STAGE: RX + final range + collapse/head + the drain (doc above) ----
1129 let _stl = rt.enter(n_st - 1);
1130 let el = rt.engine(n_st - 1, e);
1131 let pos = pos_on(el)?;
1132 let x = rt.rx(n_st - 2, slot, payload)?;
1133 let x = self.glm5_verify_range(
1134 el,
1135 topology,
1136 x,
1137 fence[n_st - 1],
1138 fence[n_st],
1139 &pos,
1140 cache,
1141 &mut ckpt,
1142 )?;
1143 let (logits, collapsed) = self.glm5_verify_head(el, topology, &x, t)?;
1144 // el.stream() under the enter-guard IS the stage stream (memra_runtime ambient
1145 // override) — this drain settles the whole walk transitively.
1146 el.stream().synchronize()?;
1147 drop(_stl);
1148 // EXIT PUBLICATION (lane/glm5-accrace): the drain above settles the LAST stage, and
1149 // the TX-wait chain covers every earlier stage's work only UP TO its `ev_tx`. Each
1150 // earlier stage's stream still holds the stage-scope tail its locals enqueue when
1151 // they drop under the override (`pos`, the boundary residual, the per-layer
1152 // transients, this round's ckpt clones). The caller resumes and allocates for the
1153 // accept walk and the MTP re-seed, so it must be ordered behind ALL stages.
1154 self.glm5_publish_stages(e)?;
1155 Ok((logits, collapsed, ckpt))
1156 }
1157
1158 /// The engine that owns the trunk exit (collapse + output_norm + lm head), the MTP
1159 /// block's weights AND its latent plane under a ppN split: the LAST stage's engine —
1160 /// `hybrid.rs` uploads the head there (`pp::layer_engine(e, n_trunk, n_trunk - 1)`)
1161 /// and `pp::new_cache*` maps trailing MTP/NextN planes to the last stage. Door shut or
1162 /// the same-stream seam: the caller's engine, unchanged (single-device callers pay
1163 /// nothing — `e` is returned by identity).
1164 fn glm5_head_engine<'e>(&self, e: &'e Engine) -> Res<&'e Engine> {
1165 match crate::pp::pp_cuts(self.layers.len()) {
1166 Some(fence) if !crate::pp::pp2_streams_off() => {
1167 let rt = crate::pp::PpNRt::get(e)?;
1168 Ok(rt.engine(fence.len() - 2, e))
1169 }
1170 _ => Ok(e),
1171 }
1172 }
1173
1174 /// Roll the trunk back to exactly `keep` accepted verify rows (1 <= keep <= t; keep =
1175 /// j+1: the always-committed anchor row plus j accepted drafts).
1176 ///
1177 /// - MLA latent planes: `len = snapshot + keep` (truncate; rows are position-addressed
1178 /// and append-only, so the kept rows ARE what a plain decode chain would have
1179 /// written — the decode-exact contract), device `len_d` in lock-step, and
1180 /// `truncate_index_pool_keys(pool)` clamps pool-key finality to what the shortened
1181 /// `len` still justifies (the tail-ring residency tripwire fires on the next call if
1182 /// this clamp is ever skipped).
1183 /// - KDA state: restore column keep-1 (state after the last kept row); full accept
1184 /// (keep == t) keeps the resident state — the columns are clones OF it.
1185 /// - `cache.pos = snapshot + keep`.
1186 ///
1187 /// Under a live ppN split each stage's layers restore THROUGH that stage's engine ON
1188 /// its stream: the state planes and the ckpt columns live on the owning stage's device
1189 /// (per-stage `KvDev` allocation; per-stage clones in the walk), and enqueuing the
1190 /// restores on the same stage streams the walk writes on orders them relative to the
1191 /// walk without any extra fence.
1192 ///
1193 /// THE EXIT PUBLICATION IS NOT OPTIONAL (lane/glm5-accrace 2026-09-01). This body used
1194 /// to return with the restores merely ENQUEUED on the stage streams, on the reasoning
1195 /// that "the next walk's own entry fence covers the primary-stream seam". It does not:
1196 /// `fence_stages_behind` orders the STAGE streams behind the CALLER, and everything the
1197 /// round does after this point — the MTP plane reset, the h_seed rows, the next round's
1198 /// whole draft chain, the next SESSION's cache allocation and prime — runs on the
1199 /// CALLER's stream and ALLOCATES. cudarc's drops carry no read guard, so the pool could
1200 /// hand the caller a block whose stage-stream lifetime had not retired and the caller's
1201 /// writes landed under queued rollback work.
1202 ///
1203 /// MEASURED CONSEQUENCE, and why a "rollback ordering" bug showed up as a PRIME bug:
1204 /// with per-stage streams on one device the hc ppN prime over a fixed 24-token prompt
1205 /// returned three distinct logit fingerprints inside one process (a third of all primes
1206 /// non-canonical); downstream, one glm5 spec round lost an acceptance silently
1207 /// (14/42 -> 13/42) and the e2e tape diverged. Publishing here took non-canonical primes
1208 /// from 20/110 to 2/110 in an interleaved A/B, and the walk's own exit publication
1209 /// closed the remainder. Receipts:
1210 /// `research/glm53-flash-bringup-20260827/accrace-20260901/LANE.md`.
1211 pub fn glm5_verify_rollback(
1212 &self,
1213 e: &Engine,
1214 cache: &mut Cache,
1215 ckpt: &Glm5VerifyCkpt,
1216 keep: usize,
1217 ) -> Res<()> {
1218 if keep == 0 || keep > ckpt.rows {
1219 return Err(format!(
1220 "glm5_verify_rollback: keep={keep} outside 1..={} (the anchor row is always \
1221 committed; keep = accepted drafts + 1)",
1222 ckpt.rows
1223 )
1224 .into());
1225 }
1226 match crate::pp::pp_cuts(self.layers.len()) {
1227 Some(fence) if !crate::pp::pp2_streams_off() => {
1228 let rt = crate::pp::PpNRt::get(e)?;
1229 for s in 0..fence.len() - 1 {
1230 let _st = rt.enter(s);
1231 let es = rt.engine(s, e);
1232 for il in fence[s]..fence[s + 1] {
1233 self.glm5_rollback_layer(es, cache, ckpt, keep, il)?;
1234 }
1235 }
1236 // EXIT PUBLICATION (doc above): every stage stream, to the caller's.
1237 self.glm5_publish_stages(e)?;
1238 }
1239 _ => {
1240 for il in 0..self.layers.len() {
1241 self.glm5_rollback_layer(e, cache, ckpt, keep, il)?;
1242 }
1243 }
1244 }
1245 cache.pos = ckpt.pos + keep;
1246 Ok(())
1247 }
1248
1249 /// Restore ONE trunk layer to the ckpt's `keep`-row state (the per-plane contract in
1250 /// [`Self::glm5_verify_rollback`]'s doc). `e` is the layer's OWNING engine — the stage
1251 /// engine under a split, the caller's engine otherwise.
1252 fn glm5_rollback_layer(
1253 &self,
1254 e: &Engine,
1255 cache: &mut Cache,
1256 ckpt: &Glm5VerifyCkpt,
1257 keep: usize,
1258 il: usize,
1259 ) -> Res<()> {
1260 match &self.layers[il].mixer {
1261 Mixer::Mla(mla) => {
1262 if keep == ckpt.rows {
1263 // Full accept: the walk already advanced len AND the len_d device
1264 // mirror to saved + rows on the canonical plane and every replica
1265 // (append-time stores), so the restore below would rewrite unchanged
1266 // values — ~11 synchronizing pageable 4-byte copies per round on the
1267 // HOT outcome (the KDA arms' early-out twin; #82 review).
1268 return Ok(());
1269 }
1270 let saved = ckpt.latent_len[il].ok_or_else(|| {
1271 format!("glm5_verify_rollback: MLA layer {il} missing from the ckpt")
1272 })?;
1273 let plane = cache.latent[il].as_mut().ok_or_else(|| {
1274 format!("glm5_verify_rollback: MLA layer {il} has no latent plane")
1275 })?;
1276 plane.len = saved + keep;
1277 let len_i32 =
1278 i32::try_from(plane.len).map_err(|_| "latent length exceeds i32 mirror")?;
1279 // Door H (`MEMRA_GLM5_HTOD_DIET`): async `i32_set_k` instead of the synchronizing
1280 // pageable 4-byte copy — 11 of these per round, and unconditional (unlike the
1281 // KDA arm, which short-circuits when `keep == rows`).
1282 e.i32_mirror_store(&mut plane.len_d, len_i32)?;
1283 if let Some(indexer) = mla.index.as_ref() {
1284 plane.truncate_index_pool_keys(indexer.geom.pool);
1285 }
1286 // spec x TP composition: the PEER latent replicas append in lock-step with
1287 // the canonical plane (the TP walk's construction), so the same truncation
1288 // restores each of them — through its own rank's engine for the device
1289 // `len_d` mirror. Full accept skips the loop (lens already read
1290 // saved + rows — the KDA arm's early-out twin; each skipped store is a
1291 // synchronizing pageable copy per rank per layer on the HOT outcome).
1292 // Missing replicas after a verify walk are a wiring bug and refuse by
1293 // name, never a silent canonical-only restore (#80 review hardening).
1294 if let Some(tp) = mla.tp.as_ref() {
1295 let replicas = cache.glm5_tp_latent_peer[il].as_mut().ok_or_else(|| {
1296 format!(
1297 "glm5_verify_rollback: sharded MLA layer {il} has no peer \
1298 latent replicas (the TP verify walk hydrates them; a \
1299 rollback without them would silently restore the canonical \
1300 plane only)"
1301 )
1302 })?;
1303 for (i, replica) in replicas.iter_mut().enumerate() {
1304 replica.len = saved + keep;
1305 tp.rt.peers[i].i32_mirror_store(&mut replica.len_d, len_i32)?;
1306 if let Some(indexer) = mla.index.as_ref() {
1307 replica.truncate_index_pool_keys(indexer.geom.pool);
1308 }
1309 }
1310 }
1311 }
1312 Mixer::Kda(la) if la.tp.is_some() => {
1313 if keep == ckpt.rows {
1314 return Ok(()); // resident per-rank states ARE the post-keep states
1315 }
1316 let stash = ckpt.kda_tp[il].as_ref().ok_or_else(|| {
1317 format!(
1318 "glm5_verify_rollback: sharded KDA layer {il} has no per-rank stash \
1319 (the batched TP verify walk fills it; the per-row arm is refused \
1320 at walk entry)"
1321 )
1322 })?;
1323 crate::glm5_tp::kda_tp_verify_rollback(e, la, stash, keep, cache, il)?;
1324 }
1325 Mixer::Kda(la) => {
1326 if keep == ckpt.rows {
1327 return Ok(()); // resident state IS the state after the last kept row
1328 }
1329 // BATCHED-walk stash (lane/glm5-verify-batch): ring restore + re-roll,
1330 // then ONE scan replay at T=keep from the pre-round snapshot.
1331 if let Some(stash) = ckpt.kda_rows[il].as_ref() {
1332 let snap = ckpt.kda_ssm_snap[il].as_ref().ok_or_else(|| {
1333 format!("glm5_verify_rollback: KDA layer {il} has no ssm snapshot")
1334 })?;
1335 return crate::kda::kda_verify_rollback_rows(
1336 e, la, snap, stash, keep, cache, il,
1337 );
1338 }
1339 // Conv ring: restore the cloned column (unchanged — 288 KiB).
1340 let conv_cols = ckpt.kda_conv_cols[il].as_ref().ok_or_else(|| {
1341 format!("glm5_verify_rollback: KDA layer {il} has no conv columns")
1342 })?;
1343 let conv = &conv_cols[keep - 1];
1344 {
1345 let rl = cache.recur[il].as_mut().ok_or_else(|| {
1346 format!("glm5_verify_rollback: KDA layer {il} has no recurrent state")
1347 })?;
1348 e.copy_into(&mut rl.conv_state, 0, conv, conv.len())?;
1349 }
1350 // Recurrent state: REPLAY rows 0..keep from the pre-round snapshot
1351 // (loop-port 3; ckpt doc) — each replay re-issues that row's original
1352 // t=1 scan over its stolen inputs, so the rebuilt state is byte-identical
1353 // to the per-row clone this retires.
1354 let snap = ckpt.kda_ssm_snap[il].as_ref().ok_or_else(|| {
1355 format!("glm5_verify_rollback: KDA layer {il} has no ssm snapshot")
1356 })?;
1357 let stash = ckpt.kda_scan_stash[il].as_ref().ok_or_else(|| {
1358 format!("glm5_verify_rollback: KDA layer {il} has no scan stash")
1359 })?;
1360 crate::kda::kda_scan_replay(e, la, snap, &stash[..keep], cache, il)?;
1361 }
1362 Mixer::Full(_) | Mixer::Linear(_) => {
1363 return Err(format!(
1364 "glm5_verify_rollback: layer {il} mixer class was refused at walk \
1365 entry and cannot appear in a ckpt"
1366 )
1367 .into());
1368 }
1369 }
1370 Ok(())
1371 }
1372
1373 /// Reset the MTP draft plane (il = n_trunk) to `len` rows — the LANE.md rollback
1374 /// contract ("one row per step, rollback = plane len reset") plus the same pool-key
1375 /// clamp every len-shortening path owes the tail ring. The plane lives on the LAST
1376 /// stage under a split (`pp::new_cache*` maps trailing MTP planes there), so the
1377 /// device mirror writes through the head engine.
1378 fn glm5_mtp_plane_reset(&self, e: &Engine, cache: &mut Cache, len: usize) -> Res<()> {
1379 let e = self.glm5_head_engine(e)?;
1380 let mtp = self
1381 .mtp
1382 .as_ref()
1383 .ok_or("glm5_mtp_plane_reset with no MTP head loaded")?;
1384 let il = self
1385 .plan
1386 .mtp_blocks
1387 .first()
1388 .ok_or("ModelPlan declares no MTP block")?
1389 .layer
1390 .index as usize;
1391 let plane = cache
1392 .latent
1393 .get_mut(il)
1394 .and_then(|plane| plane.as_mut())
1395 .ok_or_else(|| format!("MTP block layer {il} has no latent cache plane"))?;
1396 if len > plane.len {
1397 return Err(format!(
1398 "glm5_mtp_plane_reset: target {len} is past the plane's {} rows — a reset \
1399 only ever shortens",
1400 plane.len
1401 )
1402 .into());
1403 }
1404 plane.len = len;
1405 let len_i32 = i32::try_from(len).map_err(|_| "latent length exceeds i32 mirror")?;
1406 e.stream().memcpy_htod(&[len_i32], &mut plane.len_d)?;
1407 if let Mixer::Mla(mla) = &mtp.mixer
1408 && let Some(indexer) = mla.index.as_ref()
1409 {
1410 plane.truncate_index_pool_keys(indexer.geom.pool);
1411 }
1412 Ok(())
1413 }
1414
1415 /// Single-shot glm5 speculative generation: draft (MTP head, K steps) -> verify (one
1416 /// t=K+1 walk) -> accept j (greedy longest matching prefix) -> rollback -> re-seed.
1417 /// Returns `(tokens, drafted, accepted)` — `generate_spec`'s contract. One-shot form:
1418 /// builds a [`Glm5SpecSession`] over a fresh cache and drives it to `max_new` — the
1419 /// SAME round machinery the serve worker bursts, so the tparallel gate's byte-identity
1420 /// pins cover the served path's rounds too.
1421 pub fn generate_spec_glm5(
1422 &self,
1423 e: &Engine,
1424 prompt: &[u32],
1425 max_new: usize,
1426 k: usize,
1427 ) -> Res<(Vec<u32>, usize, usize)> {
1428 self.generate_spec_glm5_gated(e, prompt, max_new, k, Glm5SpecKnobs::default())
1429 }
1430
1431 /// `generate_spec_glm5` with GATE INSTRUMENTS (never a serving surface): a draft
1432 /// override for deterministic forced-accept / forced-reject rounds, and a
1433 /// rollback-disable arm that red-proves the end-to-end byte-identity gate.
1434 pub fn generate_spec_glm5_gated(
1435 &self,
1436 e: &Engine,
1437 prompt: &[u32],
1438 max_new: usize,
1439 k: usize,
1440 mut knobs: Glm5SpecKnobs<'_>,
1441 ) -> Res<(Vec<u32>, usize, usize)> {
1442 let cap = Self::hyper_batch_cap();
1443 if k == 0 || k + 1 > cap {
1444 return Err(format!(
1445 "generate_spec_glm5: k={k} outside 1..={} (verify rows = k+1 must stay \
1446 inside the decode-exact knee, cap {cap})",
1447 cap - 1
1448 )
1449 .into());
1450 }
1451 if max_new == 0 {
1452 return Ok((Vec::new(), 0, 0));
1453 }
1454 let max_ctx = prompt.len() + max_new + k + 8;
1455 let mut sess = self.glm5_spec_session_new(e, prompt, max_ctx, None)?;
1456 let mut out: Vec<u32> = Vec::with_capacity(max_new + k);
1457 let mut drafted = 0usize;
1458 let mut accepted = 0usize;
1459 while out.len() < max_new && !sess.finished() {
1460 let (burst, d, a) = self.glm5_spec_session_burst_gated(
1461 e,
1462 &mut sess,
1463 max_new - out.len(),
1464 k,
1465 &[],
1466 &mut knobs,
1467 )?;
1468 if burst.is_empty() {
1469 break; // ctx guard tripped with nothing new — never spin
1470 }
1471 out.extend(burst);
1472 drafted += d;
1473 accepted += a;
1474 }
1475 out.truncate(max_new);
1476 Ok((out, drafted, accepted))
1477 }
1478
1479 /// BATCHED MTP-PLANE WARM (loop-port fold-in; doc at the call site in
1480 /// `glm5_spec_session_new`): fill the NextN block's latent plane with rows for pairs
1481 /// `(tokens_next[i], hiddens row i)`, i in `0..t`, in chunked t-parallel passes —
1482 /// ops 1-7 of `mtp_head_forward_mla_cached` batched over the chunk (embed gather,
1483 /// enorm/hnorm, the eh_proj concat via `place_rows_strided`, attn_norm), then ONE
1484 /// `mla_attn_cached` append per chunk (the prime-class t>1 arm the trunk's own MLA
1485 /// layers warm through; its attention output is discarded — the plane rows are the
1486 /// product). The MoE FFN, final norm and lm-head of the per-token chain are never
1487 /// run: they fed nothing but the (discarded) draft logits of prompt positions.
1488 fn glm5_mtp_plane_fill(
1489 &self,
1490 e: &Engine,
1491 tokens_next: &[u32],
1492 hiddens: &CudaSlice<f32>,
1493 t: usize,
1494 cache: &mut Cache,
1495 ) -> Res<()> {
1496 let mtp = self
1497 .mtp
1498 .as_ref()
1499 .ok_or("glm5_mtp_plane_fill with no MTP head loaded")?;
1500 let il = self
1501 .plan
1502 .mtp_blocks
1503 .first()
1504 .ok_or("ModelPlan declares no MTP block")?
1505 .layer
1506 .index as usize;
1507 let Mixer::Mla(mla) = &mtp.mixer else {
1508 return Err("glm5_mtp_plane_fill serves MLA-mixer MTP blocks only".into());
1509 };
1510 if tokens_next.len() < t {
1511 return Err(format!(
1512 "glm5_mtp_plane_fill: {t} rows requested over {} successor tokens",
1513 tokens_next.len()
1514 )
1515 .into());
1516 }
1517 let n_embd = self.cfg.n_embd as usize;
1518 let eps = self.cfg.rms_eps;
1519 // Chunk bound: the trunk prime's workspace discipline — bounds the t>1 attention
1520 // workspace and the transient buffers below without changing the append semantics
1521 // (`mla_attn_cached` appends at the plane's running length either way).
1522 const CHUNK: usize = 512;
1523 let mut done = 0usize;
1524 while done < t {
1525 let tc = (t - done).min(CHUNK);
1526 let e_emb = e.htod(&self.embd.gather(n_embd, &tokens_next[done..done + tc]))?;
1527 let mut e_norm = e.uninit(tc * n_embd)?;
1528 e.rms_norm(&e_emb, mtp.enorm.float_data(), &mut e_norm, n_embd, tc, eps)?;
1529 // hnorm over the chunk's hidden rows (one contiguous view copy — rms_norm
1530 // takes an owned-slice operand).
1531 let hv = e.view(hiddens, (done + tc) * n_embd);
1532 let mut h_rows = e.uninit(tc * n_embd)?;
1533 e.copy_view_into(
1534 &mut h_rows,
1535 0,
1536 &hv.slice(done * n_embd..(done + tc) * n_embd),
1537 tc * n_embd,
1538 )?;
1539 let mut h_norm = e.uninit(tc * n_embd)?;
1540 e.rms_norm(
1541 &h_rows,
1542 mtp.hnorm.float_data(),
1543 &mut h_norm,
1544 n_embd,
1545 tc,
1546 eps,
1547 )?;
1548 // concat rows [tc, 2*n_embd] = [enorm ; hnorm] — two strided placements.
1549 let mut concat = e.uninit(tc * 2 * n_embd)?;
1550 e.place_rows_strided(&e_norm, &mut concat, n_embd, tc, 2 * n_embd, 0)?;
1551 e.place_rows_strided(&h_norm, &mut concat, n_embd, tc, 2 * n_embd, n_embd)?;
1552 let inp_sa = e.matmul(&mtp.eh_proj, &concat, tc)?;
1553 let mut a_norm = e.uninit(tc * n_embd)?;
1554 e.rms_norm(
1555 &inp_sa,
1556 mtp.attn_norm.float_data(),
1557 &mut a_norm,
1558 n_embd,
1559 tc,
1560 eps,
1561 )?;
1562 let pos: Vec<i32> = (done as i32..(done + tc) as i32).collect();
1563 let pos_d = e.htod_i32(&pos)?;
1564 let _ = self.mla_attn_cached(e, mla, &a_norm, &pos_d, tc, il, cache)?;
1565 done += tc;
1566 }
1567 Ok(())
1568 }
1569
1570 /// Row `row` of a `[rows, n_embd]` device stack, copied into its own `[n_embd]` buffer
1571 /// (the MTP `h_seed` handoff shape).
1572 fn glm5_seed_row(
1573 &self,
1574 e: &Engine,
1575 src: &CudaSlice<f32>,
1576 rows: usize,
1577 row: usize,
1578 ) -> Res<CudaSlice<f32>> {
1579 let n_embd = self.cfg.n_embd as usize;
1580 let stack = e.view(src, rows * n_embd);
1581 let view = stack.slice(row * n_embd..(row + 1) * n_embd);
1582 let mut seed = e.uninit(n_embd)?;
1583 e.copy_view_into(&mut seed, 0, &view, n_embd)?;
1584 Ok(seed)
1585 }
1586
1587 /// SERVED-SESSION ENTRY (lane/glm5-spec-routing): prime the prompt, warm the MTP plane,
1588 /// draw the boundary token, and hand back a [`Glm5SpecSession`] the worker bursts.
1589 ///
1590 /// `sampling`: `None` / `temp <= 0` = the greedy byte-contract route (the instrument);
1591 /// `Some` with `temp > 0` = the sampled route — the boundary token, the draft chain and
1592 /// the accept walk all draw through the session's own Philox counters (`sctr` device
1593 /// events, `uctr` host accept-test uniforms via `spec::host_u01`, tag 0xFFFF_FFFE), so a
1594 /// session's randomness never repeats across bursts (the session-continuity law).
1595 /// PENALIZED sampled requests are refused loudly — the glm5 accept walk has no penalty
1596 /// arm yet; worker admission keeps them on the plain path (same split as dspark's
1597 /// penalized-greedy exclusion).
1598 pub fn glm5_spec_session_new(
1599 &self,
1600 e: &Engine,
1601 prompt: &[u32],
1602 ctx_cap: usize,
1603 sampling: Option<SpecSampling>,
1604 ) -> Res<Glm5SpecSession> {
1605 if self.hyper.is_none() {
1606 return Err("generate_spec_glm5 requires a HyperConnections trunk".into());
1607 }
1608 // Two parallel/spec programs on one model never silently coexist unless the
1609 // composition is EXPLICITLY armed: the spec x TP verify/rollback wiring
1610 // (lane/glm5-composition) exists and is rig-gated, but it has zero real-artifact
1611 // receipts, so sessions on a SHARDED model stay co-refused unless
1612 // MEMRA_GLM5_SPEC_TP=1 lifts the refusal (default OFF by design — the FLAGS row).
1613 // The predicate is the MODEL's own sharding (the same per-layer truth the verify
1614 // walk keys on), never the MEMRA_GLM5_TP env: sharding is a load-time property,
1615 // and an env read here is bypassable after load (set/load/unset) and spuriously
1616 // refuses an UNSHARDED model in a process that still carries the env (the #80
1617 // review's confirmed finding).
1618 let tp_sharded = self.layers.iter().any(|l| match &l.mixer {
1619 Mixer::Kda(la) => la.tp.is_some(),
1620 Mixer::Mla(mla) => mla.tp.is_some(),
1621 _ => false,
1622 });
1623 if tp_sharded {
1624 if !glm5_spec_tp_on() {
1625 return Err(
1626 "glm5 spec is co-refused on a MEMRA_GLM5_TP-sharded model: set \
1627 MEMRA_GLM5_SPEC_TP=1 to run the gated spec x TP composition \
1628 (default OFF — zero real-artifact receipts; every other admission \
1629 law still holds)"
1630 .into(),
1631 );
1632 }
1633 if !glm5_verify_batch_on() {
1634 return Err("MEMRA_GLM5_SPEC_TP=1 requires the BATCHED verify walk \
1635 (MEMRA_GLM5_VERIFY_BATCH must not be 0): the per-row rollback seam \
1636 carries no TP arm"
1637 .into());
1638 }
1639 // The ARMED announce prints AFTER the last admission law below — a session
1640 // refused later (draft source, ppN qualification, penalties, ctx) must never
1641 // log the engagement receipt (the fleet's prints-ARMED-then-serves-plain
1642 // trap class; rig-gates/03 caught SF3 logging it on a refused session).
1643 }
1644 // DRAFT-SOURCE SELECTION — through the GENERAL law
1645 // ([`crate::spec::resolve_draft_source_kind`], lane/glm5-extract2): DFlash2 when the
1646 // drafter is loaded (MEMRA_GLM5_DFLASH — it wins over a co-loaded MTP head, the boot
1647 // receipt states the selection), native MTP otherwise; neither = the loud refusal
1648 // below, whose bytes name the glm5 flags because the FAMILY owns the how-to-arm text.
1649 // The native MTP head is NOT required for the DFlash2 source (the q38 pattern).
1650 let dflash_src = self.glm5_dflash.as_ref();
1651 let source_kind = crate::spec::resolve_draft_source_kind(
1652 self.plan.draft_source,
1653 self.mtp.is_some(),
1654 dflash_src.is_some(),
1655 )
1656 .map_err(|why| {
1657 // The general law says WHY there is no usable source; the family owns the
1658 // how-to-arm text. Composed so the sentence is true on BOTH of the law's refusal
1659 // branches (nothing loaded / a head loaded under a plan that does not claim it),
1660 // rather than asserting "requires a draft source" at an operator whom the bracket
1661 // then tells a head IS loaded.
1662 format!(
1663 "generate_spec_glm5 cannot select a draft source ({why}). Arm one: the \
1664 embedded MTP head (MEMRA_GLM5_MTP=1; a full MoE layer, unloaded by default) \
1665 or the DFlash2 drafter (MEMRA_GLM5_DFLASH=<dir-or-hf-spec>)"
1666 )
1667 })?;
1668 if prompt.len() < 2 {
1669 return Err(
1670 "generate_spec_glm5 needs a prompt of >= 2 tokens (the MTP plane warms on \
1671 (token[i+1], hidden[i]) pairs)"
1672 .into(),
1673 );
1674 }
1675 if dflash_src.is_none() && crate::spec::spec_hpost() {
1676 // MTP-carrier-specific refusal: the DFlash2 source consumes tapped trunk
1677 // features, not the h_seed carrier, so the flag has nothing to flip there.
1678 return Err(
1679 "generate_spec_glm5 has no MEMRA_SPEC_HPOST arm: the flag flips the MTP \
1680 carrier to the post-norm hidden, but this loop seeds every committed pair \
1681 from the trunk's PRE-output_norm collapsed rows (LANE.md §A). Mixing the \
1682 two silently degrades drafts; the HPOST twin needs its own gate before it \
1683 may run"
1684 .into(),
1685 );
1686 }
1687 // ppN split (lane/glm5-ppn-verify): the verify walk, the rollback and the MTP
1688 // chain all run under the split now — but an UNQUALIFIED pipeline rewrite still
1689 // refuses loudly at the session seam, before any cache is allocated over
1690 // stage-sharded weights (worker admission additionally bounds the stage count to
1691 // the gated set; see glm5_spec_capable).
1692 if crate::pp::pp_cuts(self.layers.len()).is_some()
1693 && !self.rewrite_allowed(memra_gguf::execution_manifest::RewriteSurface::Pipeline)
1694 {
1695 return Err("pipeline rewrite is not qualified for this ModelPlan".into());
1696 }
1697 let sampling = sampling.filter(|sp| sp.temp > 0.0);
1698 if let Some(sp) = sampling.as_ref()
1699 && sp.pen_on()
1700 {
1701 return Err(
1702 "glm5 spec has no penalty arm yet: penalized sampled requests serve on the \
1703 plain path (worker admission owns the exclusion; silently dropping the \
1704 request's penalties is the failure class this refusal prevents)"
1705 .into(),
1706 );
1707 }
1708 // Room for the prompt, the anchor row and at least one verify round.
1709 if prompt.len() + 4 > ctx_cap {
1710 return Err(format!(
1711 "glm5 spec session needs ctx for prompt {} + anchor + one verify round, \
1712 cap {ctx_cap}",
1713 prompt.len()
1714 )
1715 .into());
1716 }
1717 let n_vocab = self.output.out_features();
1718 // FR-SPEC TRIM (module doc): a loaded `MEMRA_FRSPEC_TRIM` artifact means the draft
1719 // head projects over gathered top-N rows and every draft pick is a RANK id that
1720 // must remap through d2t to the true vocabulary BEFORE it is chained or verified.
1721 // The verify walk stays full-vocab regardless — the invariant under gate.
1722 if let Some(map) = self.glm5_d2t() {
1723 if map.iter().any(|&t| t as usize >= n_vocab) {
1724 return Err(format!(
1725 "glm5 FR-Spec d2t carries a token id >= n_vocab {n_vocab} — the ranks \
1726 artifact was minted for a different vocabulary"
1727 )
1728 .into());
1729 }
1730 // Engagement receipt (the dspark trim receipt's shape): the server-log line
1731 // the trim arm's per-session engagement is verified by.
1732 eprintln!(
1733 "[glm5-spec] draft head TRIMMED to {} rows (FR-Spec d2t engaged)",
1734 map.len(),
1735 );
1736 }
1737 // Confidence-gate engagement receipt (loop-port 2; the deploy-gate greps this —
1738 // never-serve-greedy law's receipt discipline): armed iff MEMRA_SPEC_PMIN > 0.
1739 if glm5_pmin() > 0.0 {
1740 eprintln!(
1741 "[glm5-spec] draft confidence gate armed: PMIN={:.3} PMIN0={} (native \
1742 chain p-of-pick; DFlash2 selector-q tau-slot truncation)",
1743 glm5_pmin(),
1744 glm5_pmin0() as u8,
1745 );
1746 }
1747 // The MTP plane index is a NATIVE-arm need; the DFlash2 source never touches the
1748 // plane (it still allocates below — plan-structural, the named cost in the module
1749 // doc — but nothing reads or resets it).
1750 let mtp_il = match dflash_src {
1751 Some(_) => None,
1752 None => Some(
1753 self.plan
1754 .mtp_blocks
1755 .first()
1756 .ok_or("ModelPlan declares no MTP block")?
1757 .layer
1758 .index as usize,
1759 ),
1760 };
1761 // Stage-owned allocation under a split (each layer's planes on its stage's device,
1762 // trailing MTP plane on the last stage); door shut = plain `Cache::new_planned`.
1763 let mut cache = crate::pp::new_cache_planned(e, &self.cfg, &self.plan, ctx_cap)?;
1764
1765 // ---- prime, boundary token, draft-source warm over the prompt ----
1766 // `prime_cache` routes to its own ppN twin under the split; `hiddens` is owned by
1767 // the LAST stage's engine (its published contract) — exactly where the MTP chain
1768 // below runs, so the warm consumes it with no device bounce. DFlash2 source: the
1769 // prime walk fills the armed HcTapSink with every prompt row's contracted tap
1770 // features (the drafter's context; round 1 ingests them into its own KV).
1771 let plen = prompt.len();
1772 let n_embd = self.cfg.n_embd as usize;
1773 let tap_layers = match dflash_src {
1774 Some(dr) => {
1775 let taps = glm5_dflash_tap_layers(&dr.draft, self.layers.len())?;
1776 cache.hc_taps = Some(HcTapSink::new(taps.clone(), n_embd, plen));
1777 Some(taps)
1778 }
1779 None => None,
1780 };
1781 let (logits0, _seed, hiddens) = self.prime_cache(e, prompt, &mut cache, 0)?;
1782 let eh = self.glm5_head_engine(e)?;
1783 // Prompt-boundary capture (lane/glm5-prefix-latent2): taken NOW — after the prime
1784 // filled every plane to the boundary, before the anchor/draft machinery below and
1785 // before any burst mutates the conv/ssm state or laps the tail ring. DFlash2-only:
1786 // the native arm's plane fill moves the MTP latent layer past the boundary before a
1787 // capture could be taken, and restore refuses the native source anyway. A refusal
1788 // drops the capture loudly and the session serves regardless.
1789 let prefix_capture = if glm5_spec_prefix_on() && dflash_src.is_some() {
1790 self.glm5_prefix_boundary_capture(e, eh, &cache, &logits0, &hiddens, plen)
1791 } else {
1792 None
1793 };
1794 let mut sctr = 0u32;
1795 let anchor = match sampling.as_ref() {
1796 Some(sp) => {
1797 crate::spec::sample_boundary_token(eh, &logits0, sp, &[], &mut sctr, "glm5-prime")?
1798 }
1799 None => argmax(&logits0) as u32,
1800 };
1801
1802 // Keyed on the KIND the general law returned, not on a second local re-derivation:
1803 // a seam whose answer is recomputed by its consumer is decoration. (`tap_layers` is
1804 // `Some` exactly when `dflash_src` is, and the law returns `Dflash2` exactly then, so
1805 // this is the same program the pre-seam code ran — the `_` arm's refusal is the
1806 // never-taken proof of that rather than a silent fallback.)
1807 let (draft, pending) = match (source_kind, dflash_src, tap_layers) {
1808 (crate::spec::DraftSourceKind::Dflash2, Some(dr), Some(taps)) => {
1809 let sink = cache
1810 .hc_taps
1811 .take()
1812 .ok_or("glm5 dflash prime tap sink vanished")?;
1813 // Drafter ctx KV on the HEAD engine (where the drafter weights loaded and
1814 // every round's chain runs); prompt feature rows ride `pending` so round 1
1815 // ingests them through the one chunked path.
1816 let kv = DflashKv::new(eh, &dr.draft.cfg, ctx_cap)?;
1817 (
1818 Glm5DraftState::Dflash2 {
1819 kv,
1820 pending: sink.rows,
1821 taps,
1822 },
1823 Vec::new(),
1824 )
1825 }
1826 (crate::spec::DraftSourceKind::Dflash2, _, _) => {
1827 return Err(
1828 "the draft-source law selected DFlash2 but this session resolved no tap \
1829 layers — a load-path bug, refused instead of silently drafting from the \
1830 MTP plane (a VANISHED tap sink is a different failure, caught by name in \
1831 the Dflash2 arm itself)"
1832 .into(),
1833 );
1834 }
1835 (crate::spec::DraftSourceKind::NativeMtp, _, _) => {
1836 // BATCHED PLANE WARM (loop-port fold-in — the map's #4, the spec.rs
1837 // `mtp_kv_fill_all` pattern re-aimed at the MLA plane): pairs
1838 // (prompt[i+1], h_i) at plane pos i, i in 0..P-1, filled in CHUNKED
1839 // t-parallel passes instead of P-1 sequential full-block forwards. The
1840 // sequential warm ran ~400 tok/s — the measured +2.5 s TTFT per 1k
1841 // prompt tokens, spec-battery flip condition 1 by name. MTP rows are
1842 // INDEPENDENT given the trunk hiddens (no row-to-row recurrence — the
1843 // plane is the only carrier), so the fill is exact in structure; the
1844 // t>1 attention takes the prime-class program, which can only move
1845 // DRAFTS, never output (verify arbitrates; the byte-identity batteries
1846 // stay the proof).
1847 self.glm5_mtp_plane_fill(eh, &prompt[1..], &hiddens, plen - 1, &mut cache)?;
1848 // pending = committed (token, h_seed) pairs not yet fed to the MTP plane.
1849 // The LAST pair's logits are the next round's first draft — the re-warm
1850 // doubles as draft 1.
1851 let pending = vec![(anchor, self.glm5_seed_row(eh, &hiddens, plen, plen - 1)?)];
1852 (Glm5DraftState::NativeMtp, pending)
1853 }
1854 };
1855 // Composition engagement receipt — printed immediately before the session is
1856 // RETURNED (after every admission law, the d2t vocabulary check, the cache
1857 // allocation and the prompt prime), so a grep for this line counts sessions that
1858 // actually opened; a refusal or a failure anywhere above never logs it (the #82
1859 // review moved it here after finding four fallible steps below its first home).
1860 if tp_sharded {
1861 eprintln!(
1862 "[glm5-spec] spec x TP composition ARMED (MEMRA_GLM5_SPEC_TP=1): verify \
1863 rows ride the TP shards; rollback restores per-rank planes \
1864 performance_claim=false"
1865 );
1866 }
1867 Ok(Glm5SpecSession {
1868 cache,
1869 committed: prompt.to_vec(),
1870 anchor,
1871 anchor_emitted: false,
1872 pending,
1873 draft,
1874 sampling,
1875 sctr,
1876 uctr: 0,
1877 rounds: 0,
1878 done: false,
1879 max_ctx: ctx_cap,
1880 mtp_il,
1881 prefix_capture,
1882 })
1883 }
1884
1885 /// Boundary capture for the DEFERRED prefix publication (lane/glm5-prefix-latent2,
1886 /// 2026-09-01): the generation-destroyed state at `pos == plen` — conv/ssm via
1887 /// `Cache::snapshot` (D2D copies), per-layer latent tails via `snapshot_tail`, the
1888 /// prime's boundary logits, the pre-output_norm boundary hidden. `None` (loud) on any
1889 /// refusal — a capture is an optimization the session must never fail on. The
1890 /// append-only planes (latent rows, final pool keys, full-attn KV) are NOT copied here:
1891 /// the worker slices them from the live cache at publish (`snapshot_plane_at`), which is
1892 /// legal because the glm5 verify rollback never truncates below the prime boundary.
1893 fn glm5_prefix_boundary_capture(
1894 &self,
1895 e: &Engine,
1896 eh: &Engine,
1897 cache: &Cache,
1898 logits0: &[f32],
1899 hiddens: &CudaSlice<f32>,
1900 plen: usize,
1901 ) -> Option<crate::spec::SpecBoundaryCapture> {
1902 debug_assert_eq!(
1903 cache.pos, plen,
1904 "boundary capture must sit at the prime boundary"
1905 );
1906 let snap = match cache.snapshot(e) {
1907 Ok(s) => s,
1908 Err(err) => {
1909 eprintln!("[glm5-spec] prefix boundary capture SKIPPED (cache snapshot: {err})");
1910 return None;
1911 }
1912 };
1913 // The ONLY latent layer that may legitimately sit empty at the boundary is the
1914 // MTP/NextN plane (allocated, never executed by the trunk on the DFlash2 arm).
1915 // Identity-keyed, not length-keyed (PR #96 review round 2, finding 1): a TRUNK
1916 // plane empty at the boundary is a regression that must refuse the capture, or
1917 // three length-keyed layers downstream would each wave it through and reproduce
1918 // the parent lane's fabrication shape per-layer.
1919 let mtp_plane_il = self.plan.mtp_blocks.first().map(|b| b.layer.index as usize);
1920 let mut latent_tails = Vec::with_capacity(cache.latent.len());
1921 for (il, l) in cache.latent.iter().enumerate() {
1922 match l {
1923 Some(l) if l.len == 0 && cache.pos > 0 => {
1924 if Some(il) != mtp_plane_il {
1925 eprintln!(
1926 "[glm5-spec] prefix boundary capture SKIPPED (trunk latent \
1927 layer {il} is EMPTY at the boundary — not the MTP plane; a \
1928 capture would publish an absent history for a live layer)"
1929 );
1930 return None;
1931 }
1932 latent_tails.push(None)
1933 }
1934 Some(l) => {
1935 if l.len != plen {
1936 eprintln!(
1937 "[glm5-spec] prefix boundary capture SKIPPED (latent layer {il} \
1938 len {} != boundary {plen})",
1939 l.len,
1940 );
1941 return None;
1942 }
1943 match l.snapshot_tail(e) {
1944 Ok(t) => latent_tails.push(Some(t)),
1945 Err(err) => {
1946 eprintln!(
1947 "[glm5-spec] prefix boundary capture SKIPPED (latent layer \
1948 {il}: {err})"
1949 );
1950 return None;
1951 }
1952 }
1953 }
1954 None => latent_tails.push(None),
1955 }
1956 }
1957 let last_h =
1958 crate::spec::capture_boundary_hidden(eh, hiddens, plen, self.cfg.n_embd as usize);
1959 Some(crate::spec::SpecBoundaryCapture {
1960 snap,
1961 pos: plen,
1962 logits: logits0.to_vec(),
1963 last_h,
1964 latent_tails,
1965 })
1966 }
1967
1968 /// Re-arm a glm5 spec session from a RESTORED trunk cache plus a published DFlash2
1969 /// drafter tail — the glm5 twin of `dspark_spec_session_from_restored`, EXTENDED with
1970 /// the suffix prime the multi-turn shape needs (lane/glm5-prefix-latent2, 2026-09-01).
1971 ///
1972 /// WHY IT IS EQUIVALENT TO A COLD PRIME, field by field:
1973 /// * `cache` — the caller's whole-entry restored trunk cache at `fed.len()` (KDA
1974 /// conv/ssm + MLA latent rows + kpool keys + tail ring, the parent lane's restore),
1975 /// and the SUFFIX primes onto it through `prime_cache` — the same continuation
1976 /// program every chunk after the first of a cold prime runs.
1977 /// * drafter — `dkv` is rebuilt from the entry's tail into the SAME absolute rows
1978 /// (`DflashKv::from_tail`, caller-side while the prefix cache is borrowable), and the
1979 /// suffix's tap rows ride `pending` exactly as a cold session's prompt rows do — the
1980 /// drafter's context is the committed tokens either way (a truncated tail below the
1981 /// window can only move ACCEPTANCE, never output: verify arbitrates).
1982 /// * anchor — drawn from the SUFFIX prime's boundary logits with the request's own
1983 /// sampler, exactly the cold composition; Philox counters fresh (a restore is a NEW
1984 /// session — the frspec continuity law, the dspark restore's own convention).
1985 /// * republish — with `glm5_spec_prefix_on()` the session takes a NEW boundary capture
1986 /// at `fed + suffix`, so the next turn hits a DEEPER prefix (the
1987 /// MEMRA_SPEC_RESTORE_REPUBLISH posture; the worker's has_key dedupe drops equals).
1988 ///
1989 /// Refuses (never asserts) whenever the restored halves disagree — a caller that gets
1990 /// `Err` serves the plain hit (correct, slower).
1991 #[allow(clippy::too_many_arguments)]
1992 pub fn glm5_spec_session_from_restored(
1993 &self,
1994 e: &Engine,
1995 mut cache: Cache,
1996 fed: &[u32],
1997 suffix: &[u32],
1998 dkv: DflashKv,
1999 ctx_cap: usize,
2000 sampling: Option<SpecSampling>,
2001 ) -> Res<Glm5SpecSession> {
2002 if self.hyper.is_none() {
2003 return Err("glm5_spec_session_from_restored requires a HyperConnections trunk".into());
2004 }
2005 // The composition laws a cold session enforces hold here too, fail-closed and by
2006 // name — a restored session must never be the door around an admission law.
2007 let tp_sharded = self.layers.iter().any(|l| match &l.mixer {
2008 Mixer::Kda(la) => la.tp.is_some(),
2009 Mixer::Mla(mla) => mla.tp.is_some(),
2010 _ => false,
2011 });
2012 if tp_sharded {
2013 return Err(
2014 "restored glm5 spec sessions carry no TP arm (the spec x TP composition is \
2015 cold-session gated only); the plain hit serves"
2016 .into(),
2017 );
2018 }
2019 if crate::pp::pp_cuts(self.layers.len()).is_some()
2020 && !self.rewrite_allowed(memra_gguf::execution_manifest::RewriteSurface::Pipeline)
2021 {
2022 return Err("pipeline rewrite is not qualified for this ModelPlan".into());
2023 }
2024 // DFlash2 source ONLY: the native MTP plane fill consumes trunk hiddens the restored
2025 // range does not have (re-running the trunk over it would be a second prime — the
2026 // whole cost this restore exists to avoid).
2027 let dr = self.glm5_dflash.as_ref().ok_or(
2028 "restored glm5 spec sessions require the DFlash2 drafter (MEMRA_GLM5_DFLASH): \
2029 the native MTP plane cannot be re-warmed from restored KV",
2030 )?;
2031 let source = crate::spec::resolve_draft_source_kind(
2032 self.plan.draft_source,
2033 self.mtp.is_some(),
2034 true,
2035 )
2036 .map_err(|why| format!("restored glm5 spec session has no draft source ({why})"))?;
2037 if !matches!(source, crate::spec::DraftSourceKind::Dflash2) {
2038 return Err(
2039 "restored glm5 spec sessions require the DFlash2 draft source; the plan \
2040 selected another"
2041 .into(),
2042 );
2043 }
2044 let sampling = sampling.filter(|sp| sp.temp > 0.0);
2045 if let Some(sp) = sampling.as_ref()
2046 && sp.pen_on()
2047 {
2048 return Err(
2049 "glm5 spec has no penalty arm: penalized requests serve on the plain path".into(),
2050 );
2051 }
2052 if fed.is_empty() || suffix.is_empty() {
2053 return Err(
2054 "restored glm5 spec session needs a non-empty restored prefix AND a \
2055 non-empty suffix (empty-suffix full-cover hits keep the plain \
2056 boundary-logits resume)"
2057 .into(),
2058 );
2059 }
2060 if cache.pos != fed.len() {
2061 return Err(format!(
2062 "restored glm5 spec session needs a whole-entry trunk cache: cache.pos {} \
2063 != restored prefix {}",
2064 cache.pos,
2065 fed.len(),
2066 )
2067 .into());
2068 }
2069 if dkv.len != fed.len() {
2070 return Err(format!(
2071 "restored draft KV len {} != restored prefix {}",
2072 dkv.len,
2073 fed.len(),
2074 )
2075 .into());
2076 }
2077 if dkv.cap != ctx_cap {
2078 return Err(
2079 format!("restored draft KV cap {} != session ctx {ctx_cap}", dkv.cap,).into(),
2080 );
2081 }
2082 // Room for prefix + suffix, the anchor row and at least one verify round.
2083 if fed.len() + suffix.len() + 4 > ctx_cap {
2084 return Err(format!(
2085 "restored glm5 spec session needs ctx for prefix {} + suffix {} + anchor + \
2086 one verify round, cap {ctx_cap}",
2087 fed.len(),
2088 suffix.len(),
2089 )
2090 .into());
2091 }
2092 let n_vocab = self.output.out_features();
2093 if let Some(map) = self.glm5_d2t() {
2094 if map.iter().any(|&t| t as usize >= n_vocab) {
2095 return Err(format!(
2096 "glm5 FR-Spec d2t carries a token id >= n_vocab {n_vocab} — the ranks \
2097 artifact was minted for a different vocabulary"
2098 )
2099 .into());
2100 }
2101 eprintln!(
2102 "[glm5-spec] draft head TRIMMED to {} rows (FR-Spec d2t engaged)",
2103 map.len(),
2104 );
2105 }
2106 if glm5_pmin() > 0.0 {
2107 eprintln!(
2108 "[glm5-spec] draft confidence gate armed: PMIN={:.3} PMIN0={} (native \
2109 chain p-of-pick; DFlash2 selector-q tau-slot truncation)",
2110 glm5_pmin(),
2111 glm5_pmin0() as u8,
2112 );
2113 }
2114 // ---- suffix prime over the restored planes (the continuation program), taps armed
2115 // for exactly the suffix rows (`HcTapSink::origin` anchors the sink at the restored
2116 // boundary; the chunked walk's absolute bases rebase through it).
2117 let n_embd = self.cfg.n_embd as usize;
2118 let taps = glm5_dflash_tap_layers(&dr.draft, self.layers.len())?;
2119 cache.hc_taps = Some(HcTapSink::new_at(
2120 taps.clone(),
2121 n_embd,
2122 suffix.len(),
2123 fed.len(),
2124 ));
2125 let (logits_s, _seed, hiddens) = self.prime_cache(e, suffix, &mut cache, 0)?;
2126 let eh = self.glm5_head_engine(e)?;
2127 // Republish capture at the NEW (deeper) boundary — pos == fed + suffix here.
2128 let prefix_capture = if glm5_spec_prefix_on() {
2129 self.glm5_prefix_boundary_capture(e, eh, &cache, &logits_s, &hiddens, cache.pos)
2130 } else {
2131 None
2132 };
2133 let mut sctr = 0u32;
2134 let anchor = match sampling.as_ref() {
2135 Some(sp) => crate::spec::sample_boundary_token(
2136 eh,
2137 &logits_s,
2138 sp,
2139 &[],
2140 &mut sctr,
2141 "glm5-restore",
2142 )?,
2143 None => argmax(&logits_s) as u32,
2144 };
2145 let sink = cache
2146 .hc_taps
2147 .take()
2148 .ok_or("glm5 restored-session suffix tap sink vanished")?;
2149 let mut committed = Vec::with_capacity(fed.len() + suffix.len());
2150 committed.extend_from_slice(fed);
2151 committed.extend_from_slice(suffix);
2152 // Engagement receipt (the dspark restore's shape — the deploy gate greps this; a
2153 // cached_tokens number alone cannot distinguish a spec restore from a plain hit).
2154 eprintln!(
2155 "[glm5-spec] RESTORED session: {} prefix tokens + {} suffix from cache — no \
2156 cold prime (drafter tail rows {})",
2157 fed.len(),
2158 suffix.len(),
2159 dkv.len,
2160 );
2161 Ok(Glm5SpecSession {
2162 cache,
2163 committed,
2164 anchor,
2165 anchor_emitted: false,
2166 pending: Vec::new(),
2167 draft: Glm5DraftState::Dflash2 {
2168 kv: dkv,
2169 pending: sink.rows,
2170 taps,
2171 },
2172 sampling,
2173 sctr,
2174 uctr: 0,
2175 rounds: 0,
2176 done: false,
2177 max_ctx: ctx_cap,
2178 mtp_il: None,
2179 prefix_capture,
2180 })
2181 }
2182
2183 /// The loaded FR-Spec draft->target map, when a trim artifact actually landed on the
2184 /// embedded head (None = full-vocab head, rank id == token id).
2185 fn glm5_d2t(&self) -> Option<&[u32]> {
2186 self.mtp
2187 .as_ref()
2188 .and_then(|head| head.d2t.as_deref())
2189 .filter(|map| !map.is_empty())
2190 }
2191
2192 /// ONE serve burst (the worker's per-tick call, `step_glm5_spec`): rounds of
2193 /// draft(K) -> `glm5_verify_rows` -> accept -> rollback/commit until `target` new
2194 /// tokens are out, EOS commits, or the context guard trips. Returns
2195 /// `(burst, drafted, accepted)`; the burst may overshoot `target` by up to K (a
2196 /// round commits j+1 tokens atomically — the engine surplus stays committed in the
2197 /// session cache and the WORKER clamps public emission to the request budget, the
2198 /// SpecSession overshoot contract).
2199 pub fn glm5_spec_session_burst(
2200 &self,
2201 e: &Engine,
2202 sess: &mut Glm5SpecSession,
2203 target: usize,
2204 k: usize,
2205 eos: &[u32],
2206 ) -> Res<(Vec<u32>, usize, usize)> {
2207 self.glm5_spec_session_burst_gated(e, sess, target, k, eos, &mut Glm5SpecKnobs::default())
2208 }
2209
2210 /// [`glm5_spec_session_burst`] with GATE INSTRUMENTS (`Glm5SpecKnobs` — never a serving
2211 /// surface; no serving path constructs a non-default value).
2212 pub fn glm5_spec_session_burst_gated(
2213 &self,
2214 e: &Engine,
2215 sess: &mut Glm5SpecSession,
2216 target: usize,
2217 k: usize,
2218 eos: &[u32],
2219 knobs: &mut Glm5SpecKnobs<'_>,
2220 ) -> Res<(Vec<u32>, usize, usize)> {
2221 let cap = Self::hyper_batch_cap();
2222 if k == 0 || k + 1 > cap {
2223 return Err(format!(
2224 "glm5_spec_session_burst: k={k} outside 1..={} (verify rows = k+1 must stay \
2225 inside the decode-exact knee, cap {cap})",
2226 cap - 1
2227 )
2228 .into());
2229 }
2230 if let Glm5DraftState::Dflash2 { .. } = sess.draft {
2231 let b = self
2232 .glm5_dflash
2233 .as_ref()
2234 .ok_or("dflash session on a model with no loaded drafter")?
2235 .draft
2236 .cfg
2237 .block_size;
2238 if k + 1 > b {
2239 return Err(format!(
2240 "glm5_spec_session_burst: k={k} exceeds the DFlash2 drafter's block \
2241 (block_size {b} = anchor + {} drafts, the trained mask pattern) — \
2242 the worker clamps operator K pins to {} for this source; refusing \
2243 loudly rather than drafting an untrained shape",
2244 b - 1,
2245 b - 1
2246 )
2247 .into());
2248 }
2249 }
2250 let d2t = self.glm5_d2t();
2251 if d2t.is_some() && knobs.skip_d2t_remap {
2252 eprintln!("[glm5-spec] d2t REMAP SKIPPED — red-arm instrument, drafts are rank ids");
2253 }
2254 let sp_on: Option<SpecSampling> = sess.sampling.filter(|sp| sp.temp > 0.0);
2255 let mut out: Vec<u32> = Vec::with_capacity(target + k);
2256 let mut drafted = 0usize;
2257 let mut accepted = 0usize;
2258 let mut phase: Option<SpecPhaseNs> =
2259 crate::spec_phase::spec_trace_on().then(SpecPhaseNs::default);
2260 if !sess.anchor_emitted {
2261 // The prime's boundary token: emitted exactly once, by the first burst.
2262 out.push(sess.anchor);
2263 sess.anchor_emitted = true;
2264 if eos.contains(&sess.anchor) {
2265 sess.done = true;
2266 }
2267 }
2268 while out.len() < target && !sess.done {
2269 // Context guard: a round appends up to k+1 trunk rows from `cache.pos` (and the
2270 // draft plane stays <= pos + k), so the next round must fit with one row slack.
2271 if sess.cache.pos + k + 2 > sess.max_ctx {
2272 sess.done = true;
2273 break;
2274 }
2275 let (round_tokens, n_drafted) =
2276 self.glm5_spec_round(e, sess, k, d2t, sp_on.as_ref(), knobs, phase.as_mut())?;
2277 drafted += n_drafted;
2278 accepted += round_tokens.len() - 1; // j accepted drafts + the bonus row
2279 for &tok in &round_tokens {
2280 if eos.contains(&tok) {
2281 sess.done = true;
2282 }
2283 }
2284 out.extend_from_slice(&round_tokens);
2285 sess.rounds += 1;
2286 }
2287 if let Some(ph) = phase.as_ref() {
2288 ph.emit("glm5-phase", "glm5-phase-v", k);
2289 }
2290 Ok((out, drafted, accepted))
2291 }
2292
2293 /// One draft->verify->accept->rollback->re-seed round over the session state. Returns
2294 /// `(round_tokens, n_drafted)`: the round's committed tokens (`j` accepted drafts + the
2295 /// bonus token) and how many drafts actually entered the verify (== `k` today; the
2296 /// confidence gate may truncate it below `k`).
2297 #[allow(clippy::too_many_arguments)]
2298 // allow: the parameter list mirrors the round contract (session + policy + gate knobs +
2299 // the trace accumulator); bundling would hide which inputs are serving vs instrument
2300 fn glm5_spec_round(
2301 &self,
2302 e: &Engine,
2303 sess: &mut Glm5SpecSession,
2304 k: usize,
2305 d2t: Option<&[u32]>,
2306 sp: Option<&SpecSampling>,
2307 knobs: &mut Glm5SpecKnobs<'_>,
2308 mut phase: Option<&mut SpecPhaseNs>,
2309 ) -> Res<(Vec<u32>, usize)> {
2310 let n_vocab = self.output.out_features();
2311 let n_embd = self.cfg.n_embd as usize;
2312 // The MTP block / DFlash2 drafter, the trunk lm head and the verify walk's returned
2313 // rows all live on the LAST stage under a split — every draft-chain and accept-side
2314 // op below runs through the head engine (identity when the door is shut).
2315 let eh = self.glm5_head_engine(e)?;
2316 let mut t_mark = phase.as_ref().map(|_| SpecPhaseNs::clock(e, eh));
2317 // Phase-boundary bump: drain, bucket the elapsed ns, restart the clock. No-op with
2318 // the trace off (t_mark is None and no stream is ever synchronized).
2319 macro_rules! bump {
2320 ($field:ident) => {
2321 if let (Some(ph), Some(t0)) = (phase.as_deref_mut(), t_mark.as_mut()) {
2322 let now = SpecPhaseNs::clock(e, eh);
2323 ph.$field += now.duration_since(*t0).as_nanos() as u64;
2324 *t0 = now;
2325 }
2326 };
2327 }
2328
2329 // CONFIDENCE GATE resolution (loop-port 2): the env pair is the serving surface
2330 // (the step37 family, no new flags); the knobs override is the gate instrument.
2331 let (p_min, pmin0) = knobs
2332 .pmin_override
2333 .unwrap_or_else(|| (glm5_pmin(), glm5_pmin0()));
2334
2335 // ---- 1+2. produce the K drafts (+ the retained q side), SOURCE-KEYED. Everything
2336 // after this point is shared and source-blind — the exactness seam (module doc).
2337 let (drafts, qside, mtp_committed_len) = match sess.draft {
2338 Glm5DraftState::Dflash2 { .. } => {
2339 let (d, q) = self.glm5_dflash_round_drafts(eh, sess, k, sp, knobs, p_min, pmin0)?;
2340 (d, q, 0)
2341 }
2342 Glm5DraftState::NativeMtp => {
2343 let mtp_il = sess.mtp_il.ok_or("native-mtp arm without a plane index")?;
2344 // ---- feed pending committed pairs; the last call yields draft 1 ----
2345 let mut last: Option<(CudaSlice<f32>, CudaSlice<f32>)> = None;
2346 for (tok, h) in sess.pending.drain(..) {
2347 let plane_len = sess.cache.latent[mtp_il]
2348 .as_ref()
2349 .ok_or("MTP plane missing")?
2350 .len;
2351 last = Some(self.mtp_head_forward_mla_cached(
2352 eh,
2353 0,
2354 tok,
2355 &h,
2356 &mut sess.cache,
2357 plane_len,
2358 )?);
2359 }
2360 let (mut d_logits, mut carrier) =
2361 last.ok_or("glm5 spec round started with no pending committed pair")?;
2362 let mtp_committed_len = sess.cache.latent[mtp_il]
2363 .as_ref()
2364 .ok_or("MTP plane missing")?
2365 .len;
2366
2367 // ---- chain K drafts. Greedy route: argmax over the draft head. Sampled
2368 // route: filtered Gumbel draw through the session's device Philox stream
2369 // (`sctr`), with the per-step filtered stats + logits retained — they are
2370 // the q side of the accept walk. Trimmed heads yield RANK ids that remap
2371 // through d2t to true vocab before anything consumes them (chain feed,
2372 // verify, output); the q gather keeps the rank id.
2373 let d_vocab = d2t.map(|m| m.len()).unwrap_or(n_vocab);
2374 let mut drafts: Vec<u32> = Vec::with_capacity(k); // true-vocab tokens
2375 let mut draft_idx: Vec<u32> = Vec::with_capacity(k); // draft-head rank ids
2376 let mut draft_logits: Vec<CudaSlice<f32>> = Vec::new(); // sampled route only
2377 let mut draft_stats: Vec<(f32, f32, f32)> = Vec::new(); // (mx, th, z), sampled only
2378 for ki in 0..k {
2379 let (idx, sampled_stats) = match sp {
2380 Some(sp) => {
2381 let (idx, stats) =
2382 glm5_sampled_draft(eh, &d_logits, d_vocab, sp, &mut sess.sctr)?;
2383 (idx, Some(stats))
2384 }
2385 None => {
2386 // Device argmax + ONE 4-byte readback per draft (loop-port 1)
2387 // — replaces the full d_vocab logits DtoH + host argmax the
2388 // map names at this seam. Same tie-break contract
2389 // (argmax_gate); drafts never decide exactness anyway, the
2390 // verify arbitrates. The #87 sentinel guard mirrors the
2391 // spec.rs graph chain: a device argmax may emit a sentinel
2392 // on a NaN row — refuse loudly, never gather an OOB embed.
2393 let td = eh.argmax_token_device(&d_logits, d_vocab)?;
2394 let idx = crate::spec::guard_vocab_token(
2395 eh.dtoh_u32_one(&td)?,
2396 d_vocab,
2397 &format!(
2398 "glm5 native draft argmax at round {} ki={ki}",
2399 sess.rounds
2400 ),
2401 )?;
2402 (idx, None)
2403 }
2404 };
2405 // P-MIN CONFIDENCE GATE (loop-port 2, the spec.rs chain break): p =
2406 // the head's softmax confidence in its own pick (the `g_p` statistic,
2407 // prob_of_token_device kernels), one 4-byte read — armed rounds only.
2408 // Break BEFORE the pick is drafted or the next full-MoE-layer chain
2409 // forward is paid; a discarded sampled draw's Philox advance stands
2410 // (spec.rs eager parity: "counts the p-min-discarded token too").
2411 if p_min > 0.0 {
2412 let tok_d = eh.htod_u32_v(&[idx])?;
2413 let p_d = eh.prob_of_token_device(&d_logits, &tok_d, d_vocab)?;
2414 let p = eh.dtoh(&p_d)?[0];
2415 if p < p_min && (ki > 0 || pmin0) {
2416 break;
2417 }
2418 }
2419 if let Some(stats) = sampled_stats {
2420 draft_stats.push(stats);
2421 draft_logits.push(eh.clone_dtod(&d_logits)?);
2422 }
2423 let mut d = match d2t {
2424 Some(map) if !knobs.skip_d2t_remap => map[idx as usize],
2425 _ => idx,
2426 };
2427 if let Some(over) = knobs.draft_override.as_mut() {
2428 d = over(sess.rounds, ki, d);
2429 }
2430 drafts.push(d);
2431 draft_idx.push(idx);
2432 if ki + 1 < k {
2433 let plane_len = sess.cache.latent[mtp_il]
2434 .as_ref()
2435 .ok_or("MTP plane missing")?
2436 .len;
2437 let (lg, ca) = self.mtp_head_forward_mla_cached(
2438 eh,
2439 0,
2440 d,
2441 &carrier,
2442 &mut sess.cache,
2443 plane_len,
2444 )?;
2445 d_logits = lg;
2446 carrier = ca;
2447 }
2448 }
2449 let q = match sp {
2450 Some(_) => Glm5DraftQ::Mtp {
2451 draft_idx,
2452 draft_logits,
2453 draft_stats,
2454 },
2455 None => Glm5DraftQ::None,
2456 };
2457 (drafts, q, mtp_committed_len)
2458 }
2459 };
2460 bump!(draft);
2461
2462 // DFlash2 source: arm the verify tap — the walk's rows are next round's drafter
2463 // context features (rows 0..keep survive the accept; the sink is taken in step 7).
2464 // DEVICE-STAGED (loop-port 1): the walk D2Ds each tapped layer's contracted rows
2465 // instead of blocking on five in-walk DtoHs; step 7 drains post-walk.
2466 if let Glm5DraftState::Dflash2 { taps, .. } = &sess.draft {
2467 sess.cache.hc_taps = Some(HcTapSink::new_device_staged(
2468 taps.clone(),
2469 n_embd,
2470 drafts.len() + 1,
2471 ));
2472 }
2473
2474 // ---- 3. verify: one t=K+1 walk over the trunk ----
2475 let mut rows: Vec<u32> = Vec::with_capacity(drafts.len() + 1);
2476 rows.push(sess.anchor);
2477 rows.extend_from_slice(&drafts);
2478 let (vlogits, collapsed, ckpt) = self.glm5_verify_rows(e, &rows, &mut sess.cache)?;
2479 bump!(verify);
2480
2481 // ---- 4. accept ----
2482 // ZERO-DRAFT SAMPLED ROUND (PMIN0): the verify batch is just the anchor row —
2483 // m=1 = a plain decode step, exactly the llama.cpp gating spec.rs vendored. The
2484 // bonus is the full-accept filtered-Gumbel draw from that one row through the
2485 // session's Philox stream (identical in distribution to the plain sampled step
2486 // this round degenerates to). Greedy zero-draft rounds ride the general arm
2487 // below (j=0, bonus = the row-0 device argmax).
2488 let (j, bonus) = if let (true, Some(sp)) = (drafts.is_empty(), sp) {
2489 (
2490 0,
2491 self.glm5_sampled_bonus(eh, sess, sp, &vlogits, 0, n_vocab)?,
2492 )
2493 } else {
2494 match (sp, &qside) {
2495 (None, _) => {
2496 // Greedy longest matching prefix (the DFlash2 probe's rule); bonus = the
2497 // target's own argmax at the first non-accepted slot. Byte-deterministic —
2498 // the instrument the spec-vs-plain identity gates pin.
2499 //
2500 // DEVICE ACCEPT ARGMAXES (loop-port 1, the K=1 flip): per verify row, a
2501 // device argmax into one [t] slot buffer, then ONE tiny u32 readback —
2502 // replacing the (K+1) x n_vocab logits DtoH + (K+1) host argmax scans
2503 // (~2.4 MB + a host walk over 600k floats at K=3 on the real head; the
2504 // 3way arithmetic needs 0.67 ms off the fixed round cost to flip K=1).
2505 // `argmax_token_device_col` carries the host argmax's tie-break contract
2506 // bit for bit (lowest index wins, argmax_gate-validated), so the accept
2507 // walk commits the SAME tokens in the SAME order — the byte-identity
2508 // batteries below stay the proof.
2509 let t = rows.len();
2510 let mut vam_d = eh.alloc_u32_zeroed(t)?;
2511 for r in 0..t {
2512 eh.argmax_token_device_col(&vlogits, r, n_vocab, &mut vam_d, r)?;
2513 }
2514 let vam = eh.dtoh_u32(&vam_d)?;
2515 let mut j = 0usize;
2516 while j < drafts.len() && drafts[j] == vam[j] {
2517 j += 1;
2518 }
2519 if knobs.accept_probe {
2520 self.glm5_accept_probe(eh, sess.rounds, &vlogits, &drafts, &vam, j)?;
2521 }
2522 (j, vam[j])
2523 }
2524 (
2525 Some(sp),
2526 Glm5DraftQ::Mtp {
2527 draft_idx,
2528 draft_logits,
2529 draft_stats,
2530 },
2531 ) => self.glm5_sampled_accept(
2532 eh,
2533 sess,
2534 sp,
2535 &vlogits,
2536 &drafts,
2537 draft_idx,
2538 draft_logits,
2539 draft_stats,
2540 d2t,
2541 drafts.len(),
2542 )?,
2543 (Some(sp), Glm5DraftQ::Selector { prop, dl }) => {
2544 // The q38 serve route's rejection walk, VERBATIM (`dspark_accept_sampled`):
2545 // `rows` = [anchor, drafts..] is its cand contract, verify row j arbitrates
2546 // rows[j+1], the bonus draws from row k on full accept, and the reject-slot
2547 // residual uses the selector's sparse candidate-set q. Philox counters are
2548 // this session's — randomness never repeats across bursts.
2549 let (m, next) = crate::dflash::dspark_accept_sampled(
2550 eh,
2551 &vlogits,
2552 &rows,
2553 rows.len(),
2554 n_vocab,
2555 dl,
2556 prop,
2557 sp,
2558 &[],
2559 &mut sess.sctr,
2560 &mut sess.uctr,
2561 )?;
2562 let next = crate::spec::guard_vocab_token(
2563 next,
2564 n_vocab,
2565 &format!(
2566 "glm5 dflash2 sampled verify bonus at round {} j={m}",
2567 sess.rounds
2568 ),
2569 )?;
2570 (m, next)
2571 }
2572 (Some(_), Glm5DraftQ::None) => {
2573 unreachable!("sampled round without a retained q side")
2574 }
2575 }
2576 };
2577 bump!(accept);
2578
2579 // ---- 5. commit j drafts + the bonus token ----
2580 let mut round_tokens: Vec<u32> = Vec::with_capacity(j + 1);
2581 round_tokens.extend_from_slice(&drafts[..j]);
2582 round_tokens.push(bonus);
2583
2584 // ---- 6. rollback the trunk to the accepted prefix ----
2585 let keep = j + 1;
2586 if knobs.disable_rollback {
2587 // RED-ARM INSTRUMENT: move pos, leave every state plane at post-row-K.
2588 sess.cache.pos = ckpt.pos + keep;
2589 } else {
2590 self.glm5_verify_rollback(e, &mut sess.cache, &ckpt, keep)?;
2591 }
2592 bump!(roll);
2593
2594 // ---- 7+8. draft-source state maintenance, SOURCE-KEYED ----
2595 match &mut sess.draft {
2596 Glm5DraftState::NativeMtp => {
2597 // MTP plane: len reset to the committed boundary (chain rows out), then
2598 // re-seed the pending pairs (token at pos0+i, collapsed row i-1).
2599 self.glm5_mtp_plane_reset(e, &mut sess.cache, mtp_committed_len)?;
2600 for i in 1..=keep {
2601 let tok = round_tokens[i - 1];
2602 let h = self.glm5_seed_row(eh, &collapsed, rows.len(), i - 1)?;
2603 sess.pending.push((tok, h));
2604 }
2605 }
2606 Glm5DraftState::Dflash2 { pending, taps, .. } => {
2607 // The kept verify rows' tap features (rows 0..keep = [anchor, accepted
2608 // drafts]) become next round's drafter context — the probe's
2609 // `F_feat[new_lo:start]` advance. The drafter's own KV block rows were
2610 // transient (forward_round never moves kv.len), so no drafter rollback
2611 // exists to run. The trunk-side MTP plane was never touched.
2612 // Device-staged rows drain HERE — the round's one post-walk sync point
2613 // for tap features (loop-port 1).
2614 let mut sink = sess
2615 .cache
2616 .hc_taps
2617 .take()
2618 .ok_or("glm5 dflash verify tap sink vanished")?;
2619 self.glm5_tap_drain(e, &mut sink)?;
2620 let row_w = taps.len() * n_embd;
2621 pending.extend_from_slice(&sink.rows[..keep * row_w]);
2622 }
2623 }
2624 // Cache-row bookkeeping: the trunk committed rows [anchor, drafts[..j]] (keep =
2625 // j+1), so `committed` gains exactly those tokens; the BONUS is the new live
2626 // anchor — emitted this round, consumed by the trunk as the NEXT round's row 0
2627 // (the dspark `last` convention). Invariant at every round boundary:
2628 // `cache.pos == committed.len()`, token-for-token.
2629 sess.committed.push(sess.anchor);
2630 sess.committed.extend_from_slice(&drafts[..j]);
2631 sess.anchor = bonus;
2632 bump!(maint);
2633 if let Some(ph) = phase {
2634 ph.rounds += 1;
2635 }
2636 Ok((round_tokens, drafts.len()))
2637 }
2638
2639 /// THE ACCEPTANCE-RACE FIX (lane/glm5-accrace 2026-09-01): order the CALLER's stream
2640 /// behind EVERY stage stream — the exit mirror of
2641 /// [`crate::pp::PpNRt::fence_stages_behind`], built from
2642 /// [`crate::pp::PpNRt::publish_all_to`] (event waits, never a device sync, so the stage
2643 /// streams keep running).
2644 ///
2645 /// Call OUTSIDE any `rt.enter` scope: `e.stream()` must resolve to the caller's stream,
2646 /// not a stage's. Door shut or the same-stream seam (`MEMRA_PP_STREAMS=0`): a no-op by
2647 /// construction, so single-device and STREAMS=0 behaviour is untouched.
2648 fn glm5_publish_stages(&self, e: &Engine) -> Res<()> {
2649 if crate::pp::pp_cuts(self.layers.len()).is_some() && !crate::pp::pp2_streams_off() {
2650 let rt = crate::pp::PpNRt::get(e)?;
2651 let dst = e.stream();
2652 rt.publish_all_to(&dst)?;
2653 }
2654 Ok(())
2655 }
2656
2657 /// GATE INSTRUMENT (lane/glm5-accrace; contract in [`Glm5SpecKnobs::accept_probe`]):
2658 /// one stderr line per greedy round pairing the DEVICE accept row against a HOST
2659 /// argmax over the same buffer, plus a per-row (argmax, row hash) census so two runs
2660 /// of the same deterministic fixture can be diffed round-for-round.
2661 fn glm5_accept_probe(
2662 &self,
2663 eh: &Engine,
2664 round: usize,
2665 vlogits: &CudaSlice<f32>,
2666 drafts: &[u32],
2667 vam: &[u32],
2668 j: usize,
2669 ) -> Res<()> {
2670 let n_vocab = self.output.out_features();
2671 let t = vam.len();
2672 let host = eh.dtoh(vlogits)?;
2673 let mut hvam: Vec<u32> = Vec::with_capacity(t);
2674 let mut rows_census: Vec<String> = Vec::with_capacity(t);
2675 for r in 0..t {
2676 let row = &host[r * n_vocab..(r + 1) * n_vocab];
2677 let am = argmax(row) as u32;
2678 hvam.push(am);
2679 // FNV-1a over the row's f32 BITS: a bit-level fingerprint, so a run-to-run
2680 // diff is exact rather than eyeballed at some print precision.
2681 let mut h: u64 = 0xcbf2_9ce4_8422_2325;
2682 for v in row {
2683 for b in v.to_bits().to_le_bytes() {
2684 h ^= u64::from(b);
2685 h = h.wrapping_mul(0x100_0000_01b3);
2686 }
2687 }
2688 rows_census.push(format!("{r}:{am}:{h:016x}"));
2689 }
2690 eprintln!(
2691 "[glm5-accrace] round={round} t={t} j={j} keep={} drafts={drafts:?} \
2692 dev_vam={vam:?} host_vam={hvam:?} agree={} rows=[{}]",
2693 j + 1,
2694 hvam == vam,
2695 rows_census.join(" ")
2696 );
2697 Ok(())
2698 }
2699
2700 /// ONE round's drafts from the DFlash2 source (module doc, DRAFT SOURCE SEAM) — the
2701 /// shipped q38 selector round, re-aimed at glm5's hc-contract features:
2702 ///
2703 /// 1. ingest pending committed feature rows into the drafter's own ctx KV (chunked
2704 /// at 256 rows — the qwen depth-OOM bound; round 1 carries the whole prompt);
2705 /// 2. block forward `[anchor, MASK x b-1]` at absolute positions over the cached ctx
2706 /// (`forward_round` — block K/V transient, exactly the reference crop);
2707 /// 3. draft logits = trunk lm_head over rows 1..b (mask-fill harvest — the DFlash2
2708 /// family census; FR-Spec trim consumed exactly as the dspark serve arm does);
2709 /// 4. selector walk: greedy chain, or the sampled candidate-set walk whose recorded
2710 /// q (`DsparkDraftSample::Selector`) the shared rejection accept consumes.
2711 ///
2712 /// Drafts are truncated to `k` (the chain is sequential, so a prefix is well-formed),
2713 /// then to the CONFIDENCE prefix when `p_min` is armed (loop-port 2, the tau-slot
2714 /// form): the selector's recorded per-slot q — `q_chosen` on the sampled walk, its
2715 /// T=1 twin on the greedy walk — gates each slot through `glm5_conf_keep`, so the
2716 /// low-confidence tail never enters the verify batch (a truncated round rides down
2717 /// the `31.6 + 20.1*K` line; rejection sampling stays exact for any proposal prefix).
2718 /// `knobs.draft_override` applies after — the gate instrument, never serving.
2719 #[allow(clippy::too_many_arguments)]
2720 // allow: the parameter list mirrors the round contract plus the resolved gate pair
2721 fn glm5_dflash_round_drafts(
2722 &self,
2723 eh: &Engine,
2724 sess: &mut Glm5SpecSession,
2725 k: usize,
2726 sp: Option<&SpecSampling>,
2727 knobs: &mut Glm5SpecKnobs<'_>,
2728 p_min: f32,
2729 pmin0: bool,
2730 ) -> Res<(Vec<u32>, Glm5DraftQ)> {
2731 let dr = self
2732 .glm5_dflash
2733 .as_ref()
2734 .ok_or("glm5 dflash draft state without a loaded drafter")?;
2735 let draft = &dr.draft;
2736 let c = &draft.cfg;
2737 let b = c.block_size;
2738 let n_embd = self.cfg.n_embd as usize;
2739 let n_vocab = self.output.out_features();
2740 let Glm5SpecSession {
2741 draft: state,
2742 cache,
2743 anchor,
2744 sctr: _,
2745 uctr,
2746 rounds,
2747 ..
2748 } = sess;
2749 let Glm5DraftState::Dflash2 { kv, pending, taps } = state else {
2750 return Err("glm5_dflash_round_drafts on a native-mtp session".into());
2751 };
2752 let anchor = *anchor;
2753
2754 // ---- 1. ingest pending committed feature rows (positions kv.len..) ----
2755 let row_w = taps.len() * n_embd;
2756 debug_assert_eq!(pending.len() % row_w, 0, "ragged pending feature rows");
2757 let n_new = pending.len() / row_w;
2758 let mut r0 = 0usize;
2759 while r0 < n_new {
2760 let t_c = (n_new - r0).min(256);
2761 let chunk = eh.htod(&pending[r0 * row_w..(r0 + t_c) * row_w])?;
2762 let feats = draft.ctx_features(eh, &chunk, t_c)?;
2763 let pos_c: Vec<i32> = ((kv.len as i32)..(kv.len + t_c) as i32).collect();
2764 draft.ingest_ctx(eh, kv, &feats, &pos_c, t_c)?;
2765 r0 += t_c;
2766 }
2767 pending.clear();
2768 let start = cache.pos;
2769 debug_assert_eq!(
2770 kv.len, start,
2771 "drafter ctx rows must equal committed trunk rows at a round boundary"
2772 );
2773
2774 // ---- 2. block forward over the cached ctx (decode-exact matmul scope: the m=8
2775 // drafter GEMMs otherwise fall into the prefill-GEMM class — the dspark round's
2776 // measured fix; RAII so a `?` exit never latches exact engine-wide) ----
2777 let exact_scope = eh.exact_scope(true);
2778 let mut block: Vec<u32> = vec![c.mask_token_id; b];
2779 block[0] = anchor;
2780 let noise = eh.htod(&self.embd.gather(n_embd, &block))?;
2781 let pos_block: Vec<i32> = ((start as i32)..(start + b) as i32).collect();
2782 let dh = draft.forward_round(eh, kv, &noise, &pos_block)?;
2783
2784 // ---- 3. draft logits over the mask-fill harvest rows 1..b ----
2785 let nd = b - 1;
2786 let mut rows_buf = eh.uninit(nd * n_embd)?;
2787 {
2788 let dv = eh.view(&dh, b * n_embd);
2789 let tail = dv.slice(n_embd..b * n_embd);
2790 eh.copy_view_into(&mut rows_buf, 0, &tail, nd * n_embd)?;
2791 }
2792 // TRIMMED DRAFT HEAD: the dspark serve arm's resolution verbatim — the FR-Spec
2793 // self-trim the load path builds on the MTP struct (gathered rows of the target's
2794 // own head). Available only when the MTP struct loaded; the DFlash2-without-head
2795 // boot (the q38 VRAM pattern) runs the full target head, stated in the boot receipt.
2796 let trim = self
2797 .mtp
2798 .as_ref()
2799 .filter(|m| m.d2t_from_target_head)
2800 .and_then(|m| m.shared_head_head.as_ref().zip(m.d2t.as_ref()))
2801 .filter(|(_, d2t)| !d2t.is_empty());
2802 let (dl_head, dl_vocab) = match trim {
2803 Some((head, d2t)) => (head, d2t.len()),
2804 None => (&self.output, n_vocab),
2805 };
2806 // skip_d2t_remap red arm (the q38 defect made loud): candidates stay RANK ids.
2807 let trim_d2t = trim
2808 .filter(|_| !knobs.skip_d2t_remap)
2809 .map(|(_, d2t)| d2t.as_slice());
2810 let dl = eh.matmul(dl_head, &rows_buf, nd)?;
2811
2812 // ---- 4. selector walk (greedy chain / sampled candidate-set walk) ----
2813 let (mut drafts, slot_q, qside) = match sp {
2814 None => {
2815 let (path, q) = draft
2816 .dflash2_propose_greedy_q(eh, &dl, &rows_buf, nd, dl_vocab, anchor, trim_d2t)?;
2817 (path, q, Glm5DraftQ::None)
2818 }
2819 Some(sp) => {
2820 let (path, q_chosen, cand, q_rows) = draft.dflash2_propose_sampled(
2821 eh, &dl, &rows_buf, nd, dl_vocab, anchor, sp.temp, sp.seed, uctr, trim_d2t,
2822 )?;
2823 let top_k = draft
2824 .dflash2
2825 .as_ref()
2826 .ok_or("glm5 dflash drafter lost its DFlash2 head")?
2827 .top_k;
2828 (
2829 path,
2830 q_chosen.clone(),
2831 Glm5DraftQ::Selector {
2832 prop: DsparkDraftSample::Selector {
2833 cand,
2834 q_rows,
2835 q_chosen,
2836 top_k,
2837 },
2838 dl,
2839 },
2840 )
2841 }
2842 };
2843 drop(exact_scope);
2844 drafts.truncate(k);
2845 // TAU-SLOT CONFIDENCE TRUNCATION (loop-port 2): the low-confidence tail never
2846 // enters verify. Slot-indexed prefix reads keep the retained Selector q side
2847 // consistent (cand/q_rows/q_chosen are per-slot; the accept walk reads slots
2848 // 0..drafts.len()-1 only). p_min unset = today's rounds, untouched.
2849 if p_min > 0.0 {
2850 let kc = glm5_conf_keep(&slot_q[..drafts.len()], p_min, pmin0);
2851 drafts.truncate(kc);
2852 }
2853 if let Some(over) = knobs.draft_override.as_mut() {
2854 for (ki, d) in drafts.iter_mut().enumerate() {
2855 *d = over(*rounds, ki, *d);
2856 }
2857 }
2858 Ok((drafts, qside))
2859 }
2860
2861 /// SAMPLED ACCEPT (module doc): the rejection-sampling walk `u_j * q_j(x_j) < p_j(x_j)`
2862 /// over the verify logit rows — memra's existing sampled spec contract (the
2863 /// MEMRA_SPEC_TEMP route / dspark sampled-admission walk), plugged in at exactly the
2864 /// accept seam; walk and rollback unchanged. p and q take the SAME filter transforms
2865 /// (`filter_stats` + `softmax_gather_filtered`, distribution-exact for the filtered
2866 /// target); the accept-test uniforms come from `spec::host_u01` on the session's `uctr`
2867 /// (tag 0xFFFF_FFFE) and every device draw (draft chain, full-accept bonus, residual
2868 /// resample) advances the session's `sctr` — counters persist on the session so
2869 /// randomness never repeats across bursts. Returns `(j, bonus)`. `e` is the HEAD
2870 /// engine (the round resolves it): the verify rows and retained draft logits live on
2871 /// the last stage under a split.
2872 #[allow(clippy::too_many_arguments)]
2873 // allow: the parameter list mirrors the accept seam's inputs (verify rows + the draft
2874 // chain's retained q side); bundling into a struct would hide the p/q pairing
2875 fn glm5_sampled_accept(
2876 &self,
2877 e: &Engine,
2878 sess: &mut Glm5SpecSession,
2879 sp: &SpecSampling,
2880 vlogits: &CudaSlice<f32>,
2881 drafts: &[u32],
2882 draft_idx: &[u32],
2883 draft_logits: &[CudaSlice<f32>],
2884 draft_stats: &[(f32, f32, f32)],
2885 d2t: Option<&[u32]>,
2886 k: usize,
2887 ) -> Res<(usize, u32)> {
2888 let n_vocab = self.output.out_features();
2889 let d_vocab = d2t.map(|m| m.len()).unwrap_or(n_vocab);
2890 // FILTERED p_j: one batched stats pass over verify rows 0..k-1 (row j is the target
2891 // distribution at draft j's slot), then one batched gather of the drafted tokens.
2892 let rows_i: Vec<i32> = (0..k as i32).collect();
2893 let rows_d = e.htod_i32(&rows_i)?;
2894 let (mut th_d, mut z_d, mut mx_d) = (e.zeros(k)?, e.zeros(k)?, e.zeros(k)?);
2895 e.filter_stats(
2896 vlogits, n_vocab, &rows_d, &mut th_d, &mut z_d, &mut mx_d, n_vocab, k, sp.temp,
2897 sp.top_k, sp.top_p, sp.min_p,
2898 )?;
2899 let ids_d = e.htod_u32_v(drafts)?;
2900 let mut pj_d = e.zeros(k)?;
2901 e.softmax_gather_filtered(
2902 vlogits, n_vocab, &ids_d, &rows_d, &th_d, &z_d, &mut pj_d, n_vocab, k, sp.temp,
2903 )?;
2904 let pj = e.dtoh(&pj_d)?;
2905 let (thv, zv, mxv) = (e.dtoh(&th_d)?, e.dtoh(&z_d)?, e.dtoh(&mx_d)?);
2906
2907 // The walk: FILTERED q_j from the retained draft logits (rank id for trimmed heads),
2908 // host Philox accept test per slot.
2909 let mut j = 0usize;
2910 while j < k {
2911 let (_qmx, qth, qz) = draft_stats[j];
2912 let idsd = e.htod_u32_v(&[draft_idx[j]])?;
2913 let rows0 = e.htod_i32(&[0])?;
2914 let thd = e.htod(&[qth])?;
2915 let zd = e.htod(&[qz])?;
2916 let mut outd = e.zeros(1)?;
2917 e.softmax_gather_filtered(
2918 &draft_logits[j],
2919 d_vocab,
2920 &idsd,
2921 &rows0,
2922 &thd,
2923 &zd,
2924 &mut outd,
2925 d_vocab,
2926 1,
2927 sp.temp,
2928 )?;
2929 let qj = e.dtoh(&outd)?[0];
2930 let u = crate::spec::host_u01(sp.seed, sess.uctr);
2931 sess.uctr = sess.uctr.wrapping_add(1);
2932 if (u as f64) * (qj as f64) < pj[j] as f64 {
2933 j += 1;
2934 } else {
2935 break;
2936 }
2937 }
2938
2939 // Bonus: full accept draws a filtered Gumbel sample from the LAST verify row
2940 // (`glm5_sampled_bonus` — shared with the PMIN0 zero-draft round); rejection at j
2941 // resamples the residual norm(max(0, fp_j - fq_j)) — with a trimmed draft head, q
2942 // scatters back to full vocab first (`scatter_trim_logits`).
2943 if j == k {
2944 return Ok((
2945 j,
2946 self.glm5_sampled_bonus(e, sess, sp, vlogits, k, n_vocab)?,
2947 ));
2948 }
2949 let mut col = e.zeros(n_vocab)?;
2950 let bonus = {
2951 let vv = e.view(vlogits, (k + 1) * n_vocab);
2952 let row = vv.slice(j * n_vocab..(j + 1) * n_vocab);
2953 e.copy_view_into(&mut col, 0, &row, n_vocab)?;
2954 let p_stats = (mxv[j], thv[j], zv[j]);
2955 let q_stats = draft_stats[j];
2956 let sc = sess.sctr;
2957 sess.sctr = sess.sctr.wrapping_add(1);
2958 let mut sample_tok = e.alloc_u32_zeroed(1)?;
2959 match d2t {
2960 Some(map) => {
2961 let map_d = e.htod_u32_v(map)?;
2962 let mut q_full = e.zeros(n_vocab)?;
2963 e.scatter_trim_logits(&draft_logits[j], &map_d, &mut q_full, d_vocab, n_vocab)?;
2964 e.residual_sample_filtered(
2965 &col,
2966 Some(&q_full),
2967 n_vocab,
2968 sp.temp,
2969 sp.seed,
2970 sc,
2971 p_stats,
2972 q_stats,
2973 &mut sample_tok,
2974 )?;
2975 }
2976 None => {
2977 e.residual_sample_filtered(
2978 &col,
2979 Some(&draft_logits[j]),
2980 n_vocab,
2981 sp.temp,
2982 sp.seed,
2983 sc,
2984 p_stats,
2985 q_stats,
2986 &mut sample_tok,
2987 )?;
2988 }
2989 }
2990 e.dtoh_u32(&sample_tok)?[0]
2991 };
2992 let bonus = crate::spec::guard_vocab_token(
2993 bonus,
2994 n_vocab,
2995 &format!("glm5 sampled verify bonus at round {} j={j}", sess.rounds),
2996 )?;
2997 Ok((j, bonus))
2998 }
2999
3000 /// One filtered-Gumbel bonus draw from verify row `row` through the session's device
3001 /// Philox stream — the sampled FULL-ACCEPT bonus, and the entire accept of a PMIN0
3002 /// zero-draft round (whose verify batch is just the anchor row: m=1 = a plain sampled
3003 /// decode step). Advances `sctr` exactly once; byte-for-byte the pre-extraction
3004 /// full-accept arm of `glm5_sampled_accept`.
3005 fn glm5_sampled_bonus(
3006 &self,
3007 e: &Engine,
3008 sess: &mut Glm5SpecSession,
3009 sp: &SpecSampling,
3010 vlogits: &CudaSlice<f32>,
3011 row: usize,
3012 n_vocab: usize,
3013 ) -> Res<u32> {
3014 let mut col = e.zeros(n_vocab)?;
3015 let vv = e.view(vlogits, (row + 1) * n_vocab);
3016 let src = vv.slice(row * n_vocab..(row + 1) * n_vocab);
3017 e.copy_view_into(&mut col, 0, &src, n_vocab)?;
3018 let rows0 = e.htod_i32(&[0])?;
3019 let (mut bth, mut bz, mut bmx) = (e.zeros(1)?, e.zeros(1)?, e.zeros(1)?);
3020 e.filter_stats(
3021 &col, n_vocab, &rows0, &mut bth, &mut bz, &mut bmx, n_vocab, 1, sp.temp, sp.top_k,
3022 sp.top_p, sp.min_p,
3023 )?;
3024 let (th, mx) = (e.dtoh(&bth)?[0], e.dtoh(&bmx)?[0]);
3025 let mut pb = e.zeros(n_vocab)?;
3026 e.gumbel_perturb_filtered(&col, &mut pb, n_vocab, sp.seed, sess.sctr, sp.temp, mx, th)?;
3027 sess.sctr = sess.sctr.wrapping_add(1);
3028 let td = e.argmax_token_device(&pb, n_vocab)?;
3029 crate::spec::guard_vocab_token(
3030 e.dtoh_u32_one(&td)?,
3031 n_vocab,
3032 &format!(
3033 "glm5 sampled verify bonus at round {} (row {row})",
3034 sess.rounds
3035 ),
3036 )
3037 }
3038}
3039
3040/// One filtered Gumbel draw from a draft-head logit row through the session's device Philox
3041/// stream — the sampled route's PROPOSAL. Returns the drawn RANK id and the row's filtered
3042/// stats `(row_max, threshold_e, renorm_mass)`, which the accept walk's q gather and the
3043/// rejection residual both reuse (the q side must be the distribution the draft was actually
3044/// drawn from, or rejection sampling is not exact for the filtered target).
3045fn glm5_sampled_draft(
3046 e: &Engine,
3047 dl: &CudaSlice<f32>,
3048 d_vocab: usize,
3049 sp: &SpecSampling,
3050 sctr: &mut u32,
3051) -> Res<(u32, (f32, f32, f32))> {
3052 let rows0 = e.htod_i32(&[0])?;
3053 let (mut th_d, mut z_d, mut mx_d) = (e.zeros(1)?, e.zeros(1)?, e.zeros(1)?);
3054 e.filter_stats(
3055 dl, d_vocab, &rows0, &mut th_d, &mut z_d, &mut mx_d, d_vocab, 1, sp.temp, sp.top_k,
3056 sp.top_p, sp.min_p,
3057 )?;
3058 let (th, z, mx) = (e.dtoh(&th_d)?[0], e.dtoh(&z_d)?[0], e.dtoh(&mx_d)?[0]);
3059 let mut pb = e.zeros(d_vocab)?;
3060 e.gumbel_perturb_filtered(dl, &mut pb, d_vocab, sp.seed, *sctr, sp.temp, mx, th)?;
3061 *sctr = sctr.wrapping_add(1);
3062 let td = e.argmax_token_device(&pb, d_vocab)?;
3063 let idx =
3064 crate::spec::guard_vocab_token(e.dtoh_u32_one(&td)?, d_vocab, "glm5 sampled draft draw")?;
3065 Ok((idx, (mx, th, z)))
3066}
3067
3068/// glm5_next SERVED speculative session (lane/glm5-spec-routing, 2026-08-30): the state one
3069/// request's spec decoding carries across worker bursts — the dspark/gemma session twins'
3070/// shape. The session OWNS its trunk cache (the worker's `s.cache` stays `None`); at every
3071/// burst boundary the invariant is `cache.pos == committed.len()` with each committed row's
3072/// trunk state exactly what a plain prime of that sequence would hold (the accept walk's
3073/// basis, pinned by the tparallel gate), plus ONE emitted-but-uncommitted `anchor` token
3074/// (the next round's verify row 0 — the dspark `last` convention).
3075pub struct Glm5SpecSession {
3076 cache: Cache,
3077 /// Every token whose trunk state the cache holds, in order (prompt + committed
3078 /// generation). EXCLUDES the live `anchor`.
3079 pub committed: Vec<u32>,
3080 /// The last emitted token, not yet consumed by the trunk — round anchor / verify row 0.
3081 anchor: u32,
3082 /// The prime's boundary token is emitted exactly once, by the first burst.
3083 anchor_emitted: bool,
3084 /// Committed `(token, h_seed)` pairs not yet fed to the MTP draft plane; the last
3085 /// feed's logits double as the next round's first draft (the re-warm contract).
3086 /// NATIVE-MTP arm only; the DFlash2 source keeps its own pending rows in `draft`.
3087 pending: Vec<(u32, CudaSlice<f32>)>,
3088 /// The session's pinned draft source + its state (module doc, DRAFT SOURCE SEAM).
3089 draft: Glm5DraftState,
3090 /// `None` / `temp <= 0` = greedy byte-contract route. Fixed for the session — the
3091 /// worker's admission owns the sampler identity.
3092 sampling: Option<SpecSampling>,
3093 /// Session-continuity Philox counters (never reset across bursts): `sctr` = device
3094 /// sampling events (boundary, draft chain, bonus, residual), `uctr` = host accept-test
3095 /// uniforms (`spec::host_u01`, tag 0xFFFF_FFFE).
3096 sctr: u32,
3097 uctr: u32,
3098 /// Verify rounds completed over the session lifetime (the worker's per-burst
3099 /// rounds-delta receipt, the dspark `rounds` convention).
3100 pub rounds: usize,
3101 done: bool,
3102 max_ctx: usize,
3103 /// MTP draft-plane layer index — `Some` on the native-MTP arm only.
3104 mtp_il: Option<usize>,
3105 /// Prompt-boundary capture for the DEFERRED prefix publication (lane/glm5-prefix-latent2,
3106 /// 2026-09-01; the dspark `prefix_capture` pattern): taken at session creation before any
3107 /// burst mutates the recurrent/tail state, drained by the worker's sweep. `None` when the
3108 /// worker did not request capture or when any boundary invariant refused.
3109 prefix_capture: Option<crate::spec::SpecBoundaryCapture>,
3110}
3111
3112impl Glm5SpecSession {
3113 /// Context capacity of the session's cache (the server's ContextFull guard).
3114 pub fn cache_max_ctx(&self) -> usize {
3115 self.max_ctx
3116 }
3117 /// Drain the prompt-boundary capture (doc on the field; the dspark
3118 /// `take_prefix_capture` twin — the worker publishes it against `cache_ref`).
3119 pub fn take_prefix_capture(&mut self) -> Option<crate::spec::SpecBoundaryCapture> {
3120 self.prefix_capture.take()
3121 }
3122 /// True when the deferred prefix capture can publish NOW: a capture exists AND the
3123 /// DFlash2 drafter KV already covers the boundary. glm5 defers the prompt's feature
3124 /// ingest to round 1 (unlike dspark's at-creation ingest), so a drain that fired before
3125 /// the first burst would export an empty tail and waste the capture — the worker's
3126 /// sweep polls this instead.
3127 pub fn prefix_capture_ready(&self) -> bool {
3128 self.prefix_capture
3129 .as_ref()
3130 .is_some_and(|c| match &self.draft {
3131 Glm5DraftState::Dflash2 { kv, .. } => kv.len >= c.pos,
3132 _ => false,
3133 })
3134 }
3135 /// The drafter's readable KV tail at `upto` rows (the dspark `export_tail` seam) —
3136 /// `None` on the native arm or when the tail cannot cover the drafter window.
3137 pub fn export_draft_tail(
3138 &self,
3139 e: &Engine,
3140 upto: usize,
3141 ) -> Option<crate::dflash::DflashKvTail> {
3142 match &self.draft {
3143 Glm5DraftState::Dflash2 { kv, .. } => kv.export_tail(e, upto),
3144 _ => None,
3145 }
3146 }
3147 /// The session's trunk cache, for the worker's deferred prefix publication (the
3148 /// append-only-below-boundary slices) — never for mutation.
3149 pub fn cache_ref(&self) -> &Cache {
3150 &self.cache
3151 }
3152 /// Trunk rows currently committed (== `committed.len()` at burst boundaries).
3153 pub fn pos(&self) -> usize {
3154 self.cache.pos
3155 }
3156 /// EOS committed or the context guard tripped: the next burst would emit nothing.
3157 pub fn finished(&self) -> bool {
3158 self.done
3159 }
3160 /// True when the session is a legal demotion source (loop-port fold-in, map #8):
3161 /// GREEDY only — a sampled session's committed stream depends on its session-owned
3162 /// Philox counters, and the plain batched sampler is a different random program
3163 /// mid-request (the exact exclusion the MTP and dspark sweeps carry).
3164 pub fn demote_eligible(&self) -> bool {
3165 self.sampling.is_none()
3166 }
3167}
3168
3169impl HybridModel {
3170 /// ONE-WAY DEMOTION HANDOFF for the glm5 session (loop-port fold-in — the map's #8,
3171 /// the `SpecSession::into_demoted` / `DsparkSpecSession::into_demoted` twin): consume
3172 /// the session and hand `(cache, next_pred)` to the plain batched-decode path, so a
3173 /// spec session admitted on a quiet box stops serializing the tick when load arrives
3174 /// (the spec-gate HIGH sweep's ship-safety lever; dspark receipt: "c=8 429.6 = parity
3175 /// (pre-lane -37%)").
3176 ///
3177 /// THE ANCHOR IS THE CARRIED-PENDING SHAPE: glm5 emits each round's bonus immediately
3178 /// (`round_tokens` include it) while the trunk consumes it only as the NEXT round's
3179 /// row 0 — so at every burst boundary the session holds ONE emitted-but-uncommitted
3180 /// token. Handing the cache over as-is would leave it one row short of the public
3181 /// stream, and `device_next` re-emitting the anchor would duplicate a served token.
3182 /// The flush below is `spec_flush_pending`'s exact analogue: ONE plain T=1 decode
3183 /// step commits the anchor (byte-identical to the never-drafted chain — the
3184 /// tparallel gate's accept-j-then-continue identity IS this claim), and its argmax
3185 /// becomes the handoff's `next_pred` — a token the batched path emits and feeds
3186 /// exactly as it would its own. One trunk pass, once per demotion, never per burst.
3187 ///
3188 /// ONE-WAY BY DESIGN: the draft state (MTP pending pairs / DFlash2 drafter KV and
3189 /// feature rows) and the Philox counters are DROPPED, freeing their VRAM; there is
3190 /// no cheap symmetric re-promotion (the spec.rs law, verbatim). Sampled sessions
3191 /// refuse loudly (`demote_eligible`; the worker's sweep excludes them first).
3192 pub fn glm5_spec_into_demoted(
3193 &self,
3194 e: &Engine,
3195 mut sess: Glm5SpecSession,
3196 ) -> Res<(Cache, u32)> {
3197 if !sess.demote_eligible() {
3198 return Err(
3199 "glm5 demote: sampled sessions stay on spec until they end (session-owned \
3200 Philox vs the worker sampler is an unmeasured distributional seam — the \
3201 MTP sweep's exclusion, verbatim)"
3202 .into(),
3203 );
3204 }
3205 if sess.cache.pos + 1 > sess.max_ctx {
3206 return Err(format!(
3207 "glm5 demote: no room to flush the live anchor ({} + 1 > ctx {})",
3208 sess.cache.pos, sess.max_ctx
3209 )
3210 .into());
3211 }
3212 let logits = self.decode_step(e, sess.anchor, &mut sess.cache)?;
3213 sess.committed.push(sess.anchor);
3214 let next = argmax(&logits) as u32;
3215 Ok((sess.cache, next))
3216 }
3217}
3218
3219/// Gate instruments for `generate_spec_glm5_gated`. Documented as instruments: no serving
3220/// path constructs a non-default value.
3221#[derive(Default)]
3222pub struct Glm5SpecKnobs<'a> {
3223 /// `(round, draft_index, greedy_draft) -> draft` — deterministic forced-accept /
3224 /// forced-reject rounds for the end-to-end gate.
3225 pub draft_override: Option<&'a mut dyn FnMut(usize, usize, u32) -> u32>,
3226 /// RED ARM ONLY: skip the state rollback (pos still moves). A corrupted draft must then
3227 /// leave post-row-K KDA state and un-truncated latent rows behind — the end-to-end gate
3228 /// asserts the tape DIVERGES from plain decode (or the kpool residency tripwire fires).
3229 pub disable_rollback: bool,
3230 /// RED ARM ONLY: with an FR-Spec trim loaded, use the draft argmax RANK id as the vocab
3231 /// id (the q38 skipped-remap defect: 0/248 acceptance with every exactness gate green).
3232 /// The gate asserts the drafted sequence diverges from the untrimmed arm's while the
3233 /// output tape STAYS byte-identical to plain decode — the silent failure made loud.
3234 pub skip_d2t_remap: bool,
3235 /// GATE INSTRUMENT for the confidence gate (loop-port 2): `Some((p_min, pmin0))`
3236 /// overrides the `MEMRA_SPEC_PMIN`/`MEMRA_SPEC_PMIN0` env pair for this call — the
3237 /// env statics latch once per process, so the byte-identity gate drives its PMIN
3238 /// arms through here instead of the environment. `None` = the serving resolution.
3239 pub pmin_override: Option<(f32, bool)>,
3240 /// GATE INSTRUMENT (lane/glm5-accrace): trace every GREEDY round's accept decision to
3241 /// stderr as one `[glm5-accrace]` line — round, t, j, the drafts, the DEVICE argmax
3242 /// row (`argmax_token_device_col` + the one u32 readback the accept walk consumes), a
3243 /// HOST argmax over the same `vlogits` buffer, and a per-row (argmax, FNV-1a hash of
3244 /// the row's f32 bits) census.
3245 ///
3246 /// TWO THINGS IT SEPARATES, which is why it exists: (a) `dev != host` means the device
3247 /// accept path published a value the logits buffer does not justify (a readback/scratch
3248 /// race); (b) `dev == host` with a row hash that moves between two runs of the same
3249 /// deterministic fixture means the verify logits themselves were computed over
3250 /// corrupted state (an upstream walk/rollback race). The host read is issued AFTER the
3251 /// device path's own `dtoh_u32` has already synchronized the consuming stream, so the
3252 /// probe can only observe the race, never mask it.
3253 ///
3254 /// Never a serving surface: no serving path constructs a non-default value.
3255 pub accept_probe: bool,
3256}
3257
3258#[cfg(test)]
3259mod conf_keep_tests {
3260 use super::glm5_conf_keep;
3261
3262 /// The spec.rs chain-break semantics, pinned CPU-side (loop-port 2): break at the
3263 /// first sub-threshold slot; slot 0 survives a miss unless PMIN0.
3264 #[test]
3265 fn conf_keep_matches_the_spec_rs_break_semantics() {
3266 // Gate off: everything kept.
3267 assert_eq!(glm5_conf_keep(&[0.1, 0.1], 0.0, true), 2);
3268 // All confident: everything kept.
3269 assert_eq!(glm5_conf_keep(&[0.9, 0.8, 0.7], 0.5, false), 3);
3270 // Break mid-chain at the first miss; the confident tail after it never rides
3271 // (prefix truncation — the accept rule could never commit past the gap anyway).
3272 assert_eq!(glm5_conf_keep(&[0.9, 0.2, 0.9], 0.5, false), 1);
3273 assert_eq!(glm5_conf_keep(&[0.9, 0.2, 0.9], 0.5, true), 1);
3274 // Slot-0 miss: survives without PMIN0 (the j > 0 arm of the break condition), and
3275 // does NOT latch — slot 1 is judged on its own confidence (the spec.rs chain
3276 // evaluates each slot's p independently)...
3277 assert_eq!(glm5_conf_keep(&[0.2, 0.9], 0.5, false), 2);
3278 // ...but a sub-threshold slot past 0 still breaks.
3279 assert_eq!(glm5_conf_keep(&[0.2, 0.2], 0.5, false), 1);
3280 // PMIN0 arms the zero-draft round.
3281 assert_eq!(glm5_conf_keep(&[0.2, 0.9], 0.5, true), 0);
3282 // Boundary: q == p_min is NOT below it (strict <, the spec.rs test).
3283 assert_eq!(glm5_conf_keep(&[0.5, 0.5], 0.5, true), 2);
3284 // Empty chain: nothing to keep.
3285 assert_eq!(glm5_conf_keep(&[], 0.5, true), 0);
3286 }
3287}