Skip to main content

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_TP` (default OFF, lane/glm5-composition 2026-09-01): admit glm5 spec
224/// SESSIONS on a `MEMRA_GLM5_TP`-armed model. DEFAULT OFF BY DESIGN (new-flags law): the
225/// composition's verify/rollback wiring is rig-gated for correctness (per-rank KDA
226/// snapshot/replay, per-replica MLA latent truncation — `glm5-tp-gate` arms S*), but it has
227/// ZERO real-artifact receipts and the TP serving wiring is still the named box increment;
228/// an unmeasured composition does not default ON. `=1` lifts ONLY the session co-refusal —
229/// every other admission law holds (draft source required, batched verify walk required:
230/// the per-row rollback seam carries no TP arm and refuses by name). Read per session
231/// creation. Rollback seam: unset (the co-refusal is restored verbatim).
232pub fn glm5_spec_tp_on() -> bool {
233    std::env::var("MEMRA_GLM5_SPEC_TP").as_deref() == Ok("1")
234}
235
236/// `MEMRA_SPEC_PMIN`, honored by the glm5 loop (loop-port 2 — the step37 shipping family,
237/// `MEMRA_SPEC_PMIN=0.5 MEMRA_SPEC_PMIN0=1` is what step37 serves; NO new flag): stop the
238/// draft chain early when the drafter's confidence in its own pick drops below p_min.
239/// Native chain: p = the head's softmax confidence in its pick (the spec.rs `g_p`
240/// statistic, `prob_of_token_device`). DFlash2: q = the selector's recorded per-slot
241/// candidate-set confidence (`q_chosen`; T=1 twin on the greedy walk) — the owner's
242/// "take only high confidence offers" tau-slot form, truncated PRE-verify.
243/// Unset/0 = OFF (today's rounds, byte-identical). The VALUE is a per-model measurement
244/// (spec.rs bank: q27 PMIN=0.3 was -1.9% on one pack; step37 ships 0.5) — the box-B tau
245/// ladder prices glm5's.
246pub(crate) fn glm5_pmin() -> f32 {
247    use std::sync::OnceLock;
248    static P: OnceLock<f32> = OnceLock::new();
249    *P.get_or_init(|| {
250        std::env::var("MEMRA_SPEC_PMIN")
251            .ok()
252            .and_then(|v| v.parse().ok())
253            .unwrap_or(0.0)
254    })
255}
256
257/// `MEMRA_SPEC_PMIN0=1` (llama.cpp's draft gating, vendored via spec.rs): the p-min gate
258/// applies at slot 0 too, so a low-confidence round drafts NOTHING and the verify batch is
259/// just the anchor row — m=1 = a plain decode step. Always legal for glm5 (the anchor row
260/// exists every round). "llama's 35B win rides exactly this — draft acceptance 76% at mean
261/// len 2.5 because unpredictable stretches never pay draft+verify overhead" (spec.rs).
262pub(crate) fn glm5_pmin0() -> bool {
263    use std::sync::OnceLock;
264    static P: OnceLock<bool> = OnceLock::new();
265    *P.get_or_init(|| std::env::var("MEMRA_SPEC_PMIN0").as_deref() == Ok("1"))
266}
267
268/// MEMRA_SPEC_PMIN break semantics — hoisted to the shared K-policy surface
269/// ([`crate::spec::spec_conf_keep`], lane/glm5-extract-general); re-exported here so the
270/// glm5 gates and call sites keep their name.
271pub use crate::spec::spec_conf_keep as glm5_conf_keep;
272
273/// The loaded glm5 DFlash2 drafter (module doc, DRAFT SOURCE SEAM): the model-level half of
274/// the `Dflash2` draft source — weights loaded ONCE per model on the head engine
275/// (`hybrid.rs`, `MEMRA_GLM5_DFLASH`); per-session state lives in [`Glm5DraftState`].
276pub struct Glm5DflashDrafter {
277    pub draft: DflashDraft,
278    /// First 8 hex of sha256(model.safetensors) — the boot-receipt identity pin
279    /// (`b33c0347` for the probe-pinned incoai/GLM-5.3-Flash-DFlash2 @ dc77ff1c bytes).
280    pub sha8: String,
281}
282
283/// Per-session draft-source state (module doc, DRAFT SOURCE SEAM). Selected at session
284/// creation from the model's loaded sources and pinned for the session's lifetime.
285pub(crate) enum Glm5DraftState {
286    /// Embedded NextN head: state = the MTP latent plane + `Glm5SpecSession::pending`
287    /// (token, h_seed) pairs — the pre-seam program, byte-identical.
288    NativeMtp,
289    /// DFlash2 block-diffusion drafter: state = the drafter's own ctx-feature KV cache
290    /// plus host feature rows not yet ingested. Invariant at every round boundary:
291    /// `kv.len + pending.len()/(taps.len()*n_embd) == committed.len()` — the drafter's
292    /// context is exactly the committed tokens (the probe's `F_feat[new_lo:start]` walk).
293    Dflash2 {
294        kv: DflashKv,
295        /// Committed-position feature rows awaiting ingest, `[n, n_taps*n_embd]` host
296        /// (the prompt's prime taps at session start; each round's kept verify taps after).
297        pending: Vec<f32>,
298        /// Resolved tap layers (drafter config `target_layer_ids`, red-arm shift applied).
299        taps: Vec<usize>,
300    },
301}
302
303/// The retained q side of one round's draft chain — what the sampled accept walk consumes.
304/// Greedy rounds carry `None` (the accept is the byte-deterministic prefix walk).
305enum Glm5DraftQ {
306    None,
307    /// Native MTP chain: per-slot retained draft logits (rank space under a trim) + the
308    /// filtered stats of the distribution each draft was drawn from.
309    Mtp {
310        draft_idx: Vec<u32>,
311        draft_logits: Vec<CudaSlice<f32>>,
312        draft_stats: Vec<(f32, f32, f32)>,
313    },
314    /// DFlash2 selector proposal (the recorded candidate-set q) + the retained draft-logit
315    /// rows `dl` (`dspark_accept_sampled`'s buffer contract; unread on the Selector q path).
316    Selector {
317        prop: DsparkDraftSample,
318        dl: CudaSlice<f32>,
319    },
320}
321
322/// Resolve the drafter's tap layers against the trunk: the drafter config's
323/// `target_layer_ids` are memra PLAN layer indices whose COMPLETED output feeds the fc
324/// (the probe's capture convention: `MEMRA_TRACE_LAYER_ROWS_LAYERS=5,14,24,33,42` == the
325/// drafter's own `target_layer_ids`, asserted 1:1 in `score_dflash2.py`).
326/// `MEMRA_GLM5_DFLASH_GATE_RED=tap-shift` is the RED-ARM INSTRUMENT: every tap moves +1
327/// layer — deliberately wrong features whose acceptance collapse the gate asserts while
328/// the output tape stays byte-identical. Unknown values refuse loudly.
329fn glm5_dflash_tap_layers(draft: &DflashDraft, n_trunk: usize) -> Res<Vec<usize>> {
330    let mut taps = draft.cfg.target_layer_ids.clone();
331    if taps.is_empty() {
332        return Err("glm5 DFlash2 drafter config carries no target_layer_ids".into());
333    }
334    match std::env::var("MEMRA_GLM5_DFLASH_GATE_RED").ok().as_deref() {
335        Some("tap-shift") => {
336            eprintln!(
337                "[glm5-spec] RED-ARM tap-shift: drafter tap layers shifted +1 (gate \
338                 instrument, never a serving flag)"
339            );
340            for t in taps.iter_mut() {
341                *t += 1;
342            }
343        }
344        Some("") | None => {}
345        Some(other) => {
346            return Err(format!(
347                "MEMRA_GLM5_DFLASH_GATE_RED={other:?}: unknown red arm (want tap-shift)"
348            )
349            .into());
350        }
351    }
352    if let Some(&bad) = taps.iter().find(|&&t| t >= n_trunk) {
353        return Err(
354            format!("glm5 DFlash2 tap layer {bad} is outside the {n_trunk}-layer trunk").into(),
355        );
356    }
357    Ok(taps)
358}
359
360/// Pre-round state checkpoint for one glm5 verify round. Captured by `glm5_verify_rows`
361/// BEFORE any row runs; consumed by `glm5_verify_rollback`.
362///
363/// Covers exactly the state planes a glm5_next trunk mutates in a verify round:
364/// - `latent_len`: per-layer MLA latent length at round start (rollback = truncate).
365/// - KDA state (loop-port 3, the module doc's GdnStash/ReplaySSM diet LANDED): the old
366///   per-row (conv, ssm) column clones — 4 MiB x 34 layers per COLUMN, ~0.95 GiB of
367///   transient at K=7 — are replaced by
368///   * `kda_ssm_snap[il]`: ONE recurrent-state clone per layer per round (the state
369///     BEFORE row 0),
370///   * `kda_scan_stash[il][r]`: row r's scan-input buffers, STOLEN from the step (zero
371///     copies, ~160 KB/row/layer — `kda::KdaScanInputs`), rows `0..t-1` except the last
372///     (`keep == t` needs no restore, so row `t-1` is never a replay target),
373///   * `kda_conv_cols[il][r]`: the conv ring stays PER-ROW CLONED (288 KiB, 1.4% of the
374///     ssm plane it rode beside — not worth a replay arm).
375///
376///   Partial-accept rollback REPLAYS rows `0..keep` from the snapshot
377///   (`kda::kda_scan_replay`): each replay is the original t=1 scan launch re-issued over
378///   the very buffers that row consumed, so the rebuilt state is byte-identical to the
379///   clone it replaces by construction. Under a ppN split every clone/stash lives on its
380///   layer's OWNING stage engine; rollback restores through the same per-stage seam.
381/// - `pos`: `cache.pos` at round start.
382///
383/// glm5_next has no Full/Linear trunk mixers (the walk refuses them by name), so `kv`,
384/// `tp_kv` and GDN stashes have no arm here — growing one is a deliberate extension with
385/// its own gate, not a silent default.
386pub struct Glm5VerifyCkpt {
387    pos: usize,
388    latent_len: Vec<Option<usize>>,
389    /// Per-row conv-ring clones, rows `0..t-1` except the last (doc above). PER-ROW walk
390    /// only (`MEMRA_GLM5_VERIFY_BATCH=0`); the batched walk fills `kda_rows` instead.
391    kda_conv_cols: Vec<Option<Vec<CudaSlice<f32>>>>,
392    /// The recurrent state BEFORE row 0, one clone per KDA layer per round (doc above).
393    /// BOTH walks fill this — it is the batched replay's scan base too.
394    kda_ssm_snap: Vec<Option<CudaSlice<f32>>>,
395    /// Stolen per-row scan inputs, rows `0..t-1` except the last (doc above). PER-ROW
396    /// walk only.
397    kda_scan_stash: Vec<Option<Vec<crate::kda::KdaScanInputs>>>,
398    /// BATCHED walk (lane/glm5-verify-batch): one [`crate::kda::KdaRowsStash`] per KDA
399    /// layer per round — ring snapshot + stolen raw conv rows + stolen batched scan
400    /// inputs; rollback re-rolls the ring and replays the scan ONCE at T=keep.
401    kda_rows: Vec<Option<crate::kda::KdaRowsStash>>,
402    /// glm5 TP composition (lane/glm5-composition): per-rank rollback material of each
403    /// SHARDED KDA layer's batched verify call — the rank-indexed twin of
404    /// (`kda_ssm_snap`, `kda_rows`), restored through each rank's own engine. `None` on
405    /// every unsharded layer.
406    kda_tp: Vec<Option<crate::glm5_tp::Glm5TpKdaVerifyStash>>,
407    /// Row count of the walk that filled this ckpt; rollback validates `keep` against it.
408    rows: usize,
409}
410
411impl Glm5VerifyCkpt {
412    /// GATE RECEIPT (wiring anchor, not a serving surface): how many KDA layers filled
413    /// the BATCHED rows stash vs the PER-ROW column stash — the flag A/B gate asserts
414    /// the arm it set actually ran (wiring-assertions-match-prose law: anchor on the
415    /// invocation's artifact, never the log prose).
416    pub fn kda_stash_kinds(&self) -> (usize, usize) {
417        (
418            self.kda_rows.iter().filter(|s| s.is_some()).count(),
419            self.kda_conv_cols.iter().filter(|s| s.is_some()).count(),
420        )
421    }
422}
423
424/// Position buffers for one verify walk range (per stage engine under a split — the
425/// per-stage pos_d law): `all` = the `[t]` vector the BATCHED per-layer mixer calls
426/// consume; `rows` = the per-row single-position buffers of the per-row arm, built only
427/// when that arm can run (flag off) — the batched arm never reads them.
428struct Glm5VerifyPos {
429    pos0: usize,
430    t: usize,
431    all: CudaSlice<i32>,
432    rows: Vec<CudaSlice<i32>>,
433}
434
435impl Glm5VerifyPos {
436    fn new(e: &Engine, pos0: usize, t: usize) -> Res<Self> {
437        let v: Vec<i32> = (0..t as i32).map(|r| pos0 as i32 + r).collect();
438        let all = e.htod_i32(&v)?;
439        let rows = if glm5_verify_batch_on() && t > 1 {
440            Vec::new()
441        } else {
442            (0..t)
443                .map(|r| e.htod_i32(&[(pos0 + r) as i32]))
444                .collect::<Result<_, _>>()?
445        };
446        Ok(Self { pos0, t, all, rows })
447    }
448}
449
450impl HybridModel {
451    /// THE T-PARALLEL VERIFY WALK: score `tokens` (row 0 = the last committed token, rows
452    /// 1..t = the K drafted tokens) in ONE forward over the hc trunk at positions
453    /// `cache.pos .. cache.pos + t`, in the batched-decode kernel classes (module doc).
454    ///
455    /// Returns `(logits [t, n_vocab] device, collapsed [t, n_embd] device, ckpt)`:
456    /// `logits` row r is bit-identical to the plain `decode_step_hyper` logits after
457    /// consuming `tokens[r]` at that position (the gate's bar); `collapsed` row r is the
458    /// pre-output_norm hidden — the MTP `h_seed` for position `cache.pos + r`.
459    ///
460    /// State effects: every trunk MLA plane appends `t` rows; every trunk KDA state
461    /// advances `t` steps (per-step columns stashed in the ckpt); `cache.pos` is NOT moved
462    /// (rollback owns it). The MTP block's plane (il = n_trunk) is untouched.
463    pub fn glm5_verify_rows(
464        &self,
465        e: &Engine,
466        tokens: &[u32],
467        cache: &mut Cache,
468    ) -> Res<(CudaSlice<f32>, CudaSlice<f32>, Glm5VerifyCkpt)> {
469        let topology = *self
470            .hyper
471            .as_ref()
472            .ok_or("glm5_verify_rows on a model with no HyperConnections topology")?;
473        let t = tokens.len();
474        let cap = Self::hyper_batch_cap();
475        if t == 0 {
476            return Err("glm5_verify_rows: empty verify row set".into());
477        }
478        if t > cap {
479            return Err(format!(
480                "glm5_verify_rows: t={t} > cap {cap} — at t >= PRIME_MIN_T (16) the MoE \
481                 shared-expert trio crosses off the decode-exact class (the batched-decode \
482                 gate's measured B=16 knee), so per-row bit-identity vs plain decode breaks. \
483                 K <= cap-1 drafts per round"
484            )
485            .into());
486        }
487        let mut any_sharded = false;
488        for (il, layer) in self.layers.iter().enumerate() {
489            match &layer.mixer {
490                Mixer::Kda(la) => any_sharded |= la.tp.is_some(),
491                Mixer::Mla(mla) => any_sharded |= mla.tp.is_some(),
492                _ => {
493                    return Err(format!(
494                        "glm5_verify_rows: trunk layer {il} is not a KDA or MLA mixer — the \
495                         rollback contract below is built and gated for glm5_next's two state \
496                         classes only; a Full/Linear arm needs its own ckpt plane and gate"
497                    )
498                    .into());
499                }
500            }
501        }
502        // spec x TP composition (lane/glm5-composition): the per-row walk carries no TP
503        // rollback arm — a sharded trunk demands the BATCHED walk at t > 1 (t = 1 rounds
504        // ride the TP decode walk below; full accept is the only legal outcome there).
505        if any_sharded && t > 1 && !glm5_verify_batch_on() {
506            return Err(
507                "glm5_verify_rows: MEMRA_GLM5_TP is armed and MEMRA_GLM5_VERIFY_BATCH=0 — \
508                 the per-row rollback seam carries no TP arm; the spec x TP composition \
509                 requires the batched verify walk (unset MEMRA_GLM5_VERIFY_BATCH or run \
510                 without the TP door)"
511                    .into(),
512            );
513        }
514
515        let n_embd = self.cfg.n_embd as usize;
516        let pos0 = cache.pos;
517
518        // Ckpt BEFORE any state moves.
519        let mut ckpt = Glm5VerifyCkpt {
520            pos: pos0,
521            latent_len: cache
522                .latent
523                .iter()
524                .take(self.layers.len())
525                .map(|plane| plane.as_ref().map(|plane| plane.len))
526                .collect(),
527            kda_conv_cols: (0..self.layers.len()).map(|_| None).collect(),
528            kda_ssm_snap: (0..self.layers.len()).map(|_| None).collect(),
529            kda_scan_stash: (0..self.layers.len()).map(|_| None).collect(),
530            kda_rows: (0..self.layers.len()).map(|_| None).collect(),
531            kda_tp: (0..self.layers.len()).map(|_| None).collect(),
532            rows: t,
533        };
534
535        // ppN door — the verify walk owns its stage split exactly as the batched decode
536        // walk does (`decode_step_batch_hyper_ppn`, decode_batch.rs). Loud refusal on an
537        // unqualified pipeline rewrite, never a single-engine walk over stage-sharded
538        // weights.
539        if let Some(fence) = crate::pp::pp_cuts(self.layers.len()) {
540            if !self.rewrite_allowed(memra_gguf::execution_manifest::RewriteSurface::Pipeline) {
541                return Err("pipeline rewrite is not qualified for this ModelPlan".into());
542            }
543            return self.glm5_verify_rows_ppn(e, tokens, cache, ckpt, &topology, &fence);
544        }
545
546        let pos = Glm5VerifyPos::new(e, pos0, t)?;
547        let embedded = e.htod(&self.embd.gather(n_embd, tokens))?;
548        let x = crate::hyper::expand(e, &topology, &embedded, t, n_embd)?;
549        let x = self.glm5_verify_range(
550            e,
551            &topology,
552            x,
553            0,
554            self.layers.len(),
555            &pos,
556            cache,
557            &mut ckpt,
558        )?;
559        let (logits, collapsed) = self.glm5_verify_head(e, &topology, &x, t)?;
560        Ok((logits, collapsed, ckpt))
561    }
562
563    /// One hc layer RANGE `[lo, hi)` of the verify walk — the body `glm5_verify_rows` ran
564    /// inline before the ppN twin landed, extracted so the unsplit walk and every pipeline
565    /// stage run the SAME code over their own range (the `hyper_range_decode` /
566    /// `decode_batch_layers` precedent: bit-identity between the arms is then structural,
567    /// not a coincidence of two maintained copies). At `lo=0, hi=n_layers` the launch
568    /// sequence is identical to the pre-extraction walk.
569    ///
570    /// KDA ckpt columns are cloned THROUGH `e` — under a split that is the owning stage's
571    /// engine, so each column lives on the device (and is ordered on the stream) that owns
572    /// its layer's state; `glm5_verify_rollback` restores through the same per-stage seam.
573    #[allow(clippy::too_many_arguments)]
574    // allow: the parameter list mirrors the range-walk call contract its siblings share
575    fn glm5_verify_range(
576        &self,
577        e: &Engine,
578        topology: &crate::hyper::HyperTopology,
579        mut x: CudaSlice<f32>,
580        lo: usize,
581        hi: usize,
582        pos: &Glm5VerifyPos,
583        cache: &mut Cache,
584        ckpt: &mut Glm5VerifyCkpt,
585    ) -> Res<CudaSlice<f32>> {
586        let t = pos.t;
587        let n_embd = self.cfg.n_embd as usize;
588        let eps = self.cfg.rms_eps;
589        // THE BATCHED MIXER ARM (lane/glm5-verify-batch, default ON): one t=K+1 call per
590        // layer per class instead of the per-row loop — KDA batches projections/conv/
591        // gates through the decode-exact classes with the recurrence sequential INSIDE
592        // one scan launch; MLA runs the SAME cached core at t rows on the rows-exact
593        // matmul classes (per-query causal selection + attention by construction).
594        // `0` = the per-row walk below, byte-for-byte (the rollback seam). Engagement is
595        // a receipt, announced once per process.
596        let batch = glm5_verify_batch_on() && t > 1;
597        {
598            static SAID: std::sync::Once = std::sync::Once::new();
599            SAID.call_once(|| {
600                if batch {
601                    eprintln!(
602                        "[glm5-spec] verify walk BATCHED per layer: kda=one t-call (scan \
603                         sequential in-kernel), mla=rows-exact t-call, head=rows-exact, \
604                         moe=pairs rows-call where qualified \
605                         (MEMRA_GLM5_VERIFY_BATCH default ON)"
606                    );
607                } else {
608                    eprintln!("[glm5-spec] verify walk PER-ROW (MEMRA_GLM5_VERIFY_BATCH=0 or t=1)");
609                }
610            });
611        }
612        let trace_v = crate::spec_phase::spec_trace_level() >= 2;
613        // Sub-phase clock (trace level 2 only): drain the walking stream so the elapsed
614        // ns lands in the mixer-class bucket — shares, never walls.
615        let vclock = |on: bool| -> Option<std::time::Instant> {
616            on.then(|| {
617                let _ = e.stream().synchronize();
618                std::time::Instant::now()
619            })
620        };
621        for il in lo..hi {
622            let layer = &self.layers[il];
623            let hyper = layer.hyper.as_ref().ok_or_else(|| {
624                format!("layer {il} carries no hyper-connection weights under an hc plan")
625            })?;
626
627            let (y, mix) = crate::hyper::pre_exact(e, topology, &hyper.attn, &x, t, n_embd)?;
628            let mut h = e.uninit(t * n_embd)?;
629            e.rms_norm(&y, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
630            // The batched arm refuses per-layer only for an MLA layer WITHOUT the DSA
631            // indexer: the absorbed t>1 attention arm has no per-row bit-identity claim
632            // at this seam, so it stays on the per-row loop by name (glm5_next always
633            // carries the indexer, so this is a foreign-geometry guard, not a live path).
634            let layer_batched = batch
635                && match &layer.mixer {
636                    Mixer::Kda(_) => true,
637                    Mixer::Mla(mla) => mla.index.is_some(),
638                    Mixer::Full(_) | Mixer::Linear(_) => unreachable!("refused at walk entry"),
639                };
640            let mixed = if layer_batched {
641                match &layer.mixer {
642                    // spec x TP composition: sharded mixers ride the TP verify walks —
643                    // per-rank batched rows calls, column-parallel-over-gather joins on
644                    // the rows-exact classes, per-rank rollback stash into the ckpt.
645                    Mixer::Kda(la) if la.tp.is_some() => {
646                        let t0 = vclock(trace_v);
647                        let (out, stash) =
648                            crate::glm5_tp::kda_tp_verify_rows(e, la, &h, t, eps, cache, il)?;
649                        ckpt.kda_tp[il] = Some(stash);
650                        if let Some(t0) = t0 {
651                            let _ = e.stream().synchronize();
652                            use std::sync::atomic::Ordering;
653                            crate::spec_phase::V_KDA_NS
654                                .fetch_add(t0.elapsed().as_nanos() as u64, Ordering::Relaxed);
655                        }
656                        out
657                    }
658                    Mixer::Mla(mla) if mla.tp.is_some() => {
659                        let t0 = vclock(trace_v);
660                        let out =
661                            self.mla_tp_attn_cached(e, mla, &h, &pos.all, t, il, cache, true)?;
662                        if let Some(t0) = t0 {
663                            let _ = e.stream().synchronize();
664                            use std::sync::atomic::Ordering;
665                            crate::spec_phase::V_MLA_NS
666                                .fetch_add(t0.elapsed().as_nanos() as u64, Ordering::Relaxed);
667                        }
668                        out
669                    }
670                    Mixer::Kda(la) => {
671                        // Pre-round snapshot: ONE ssm clone per layer per round, BEFORE
672                        // the batched call advances the resident state (ckpt doc; also
673                        // the batched rollback's scan-replay base).
674                        {
675                            let rl = cache.recur[il]
676                                .as_ref()
677                                .ok_or("glm5 verify KDA layer has no recurrent state")?;
678                            ckpt.kda_ssm_snap[il] = Some(e.clone_dtod(&rl.ssm_state)?);
679                        }
680                        let t0 = vclock(trace_v);
681                        let mut scan_ns = 0u64;
682                        let (out, stash) = crate::kda::kda_verify_rows_cached(
683                            e,
684                            la,
685                            &h,
686                            t,
687                            eps,
688                            cache,
689                            il,
690                            trace_v.then_some(&mut scan_ns),
691                        )?;
692                        ckpt.kda_rows[il] = Some(stash);
693                        if let Some(t0) = t0 {
694                            let _ = e.stream().synchronize();
695                            use std::sync::atomic::Ordering;
696                            crate::spec_phase::V_KDA_NS
697                                .fetch_add(t0.elapsed().as_nanos() as u64, Ordering::Relaxed);
698                            crate::spec_phase::V_KDA_SCAN_NS.fetch_add(scan_ns, Ordering::Relaxed);
699                        }
700                        out
701                    }
702                    Mixer::Mla(mla) => {
703                        let t0 = vclock(trace_v);
704                        let out =
705                            self.mla_attn_cached_rows_exact(e, mla, &h, &pos.all, t, il, cache)?;
706                        if let Some(t0) = t0 {
707                            let _ = e.stream().synchronize();
708                            use std::sync::atomic::Ordering;
709                            crate::spec_phase::V_MLA_NS
710                                .fetch_add(t0.elapsed().as_nanos() as u64, Ordering::Relaxed);
711                        }
712                        out
713                    }
714                    Mixer::Full(_) | Mixer::Linear(_) => unreachable!("refused at walk entry"),
715                }
716            } else {
717                // ---- PER-ROW mixer walk (the rollback seam; also t=1 rounds and the
718                // no-indexer MLA guard): row r's state input is row r-1's state output
719                // (KDA) / rows 0..pos0+r (MLA latent) — each row the SAME t=1 call its
720                // plain decode step makes.
721                // h_row is hoisted out of the row loop (loop-port 3): the mixer consumes
722                // it in stream order before the next row's overwrite, so ONE buffer per
723                // layer replaces t allocations (stream-ordered pool churn is the dsv4
724                // lesson).
725                let mut mixed = e.uninit(t * n_embd)?;
726                let mut h_row = e.uninit(n_embd)?;
727                #[allow(clippy::needless_range_loop)]
728                // allow: r is the sequential row cursor (slices h, offsets pos); iterating pos buffers would hide the row-chaining contract
729                for r in 0..t {
730                    e.dtod_copy_view(&h.slice(r * n_embd..(r + 1) * n_embd), &mut h_row)?;
731                    // The per-row position buffer: prebuilt when the per-row arm owns the
732                    // walk; built on demand for the rare per-layer refusal under batch.
733                    let pos_row: CudaSlice<i32>;
734                    let pos_r = if let Some(p) = pos.rows.get(r) {
735                        p
736                    } else {
737                        pos_row = e.htod_i32(&[(pos.pos0 + r) as i32])?;
738                        &pos_row
739                    };
740                    let out_row = match &layer.mixer {
741                        // spec x TP composition, t = 1 rounds only (the entry guard
742                        // refuses t > 1 on the per-row arm): one TP decode-walk call —
743                        // keep == rows == 1 is the only legal rollback there, which the
744                        // KDA arm satisfies with the resident state and the MLA arm with
745                        // the len truncation to saved + 1.
746                        Mixer::Kda(la) if la.tp.is_some() => crate::glm5_tp::kda_tp_cached(
747                            e,
748                            la,
749                            &h_row,
750                            1,
751                            eps,
752                            cache,
753                            il,
754                            crate::kda::ConvArm::Decode,
755                        )?,
756                        Mixer::Mla(mla) if mla.tp.is_some() => {
757                            self.mla_tp_attn_cached(e, mla, &h_row, pos_r, 1, il, cache, false)?
758                        }
759                        Mixer::Kda(la) => {
760                            // Pre-round snapshot: ONE ssm clone per layer per round taken
761                            // before row 0 mutates the resident state (loop-port 3).
762                            if r == 0 && t > 1 {
763                                let rl = cache.recur[il]
764                                    .as_ref()
765                                    .ok_or("glm5 verify KDA layer has no recurrent state")?;
766                                ckpt.kda_ssm_snap[il] = Some(e.clone_dtod(&rl.ssm_state)?);
767                            }
768                            if r + 1 < t {
769                                // Steal the row's scan inputs for the replay stash (zero
770                                // copies); clone only the small conv ring per row.
771                                let (out, inputs) = crate::kda::kda_decode_cached_stash(
772                                    e, la, &h_row, eps, cache, il,
773                                )?;
774                                let rl = cache.recur[il]
775                                    .as_ref()
776                                    .ok_or("glm5 verify KDA layer has no recurrent state")?;
777                                ckpt.kda_conv_cols[il]
778                                    .get_or_insert_with(Vec::new)
779                                    .push(e.clone_dtod(&rl.conv_state)?);
780                                ckpt.kda_scan_stash[il]
781                                    .get_or_insert_with(Vec::new)
782                                    .push(inputs);
783                                out
784                            } else {
785                                crate::kda::kda_decode_cached(e, la, &h_row, eps, cache, il)?
786                            }
787                        }
788                        Mixer::Mla(mla) => {
789                            self.mla_attn_cached(e, mla, &h_row, pos_r, 1, il, cache)?
790                        }
791                        // Refused at entry; unreachable keeps the match total without a silent arm.
792                        Mixer::Full(_) | Mixer::Linear(_) => unreachable!("refused at walk entry"),
793                    };
794                    e.copy_into(&mut mixed, r * n_embd, &out_row, n_embd)?;
795                }
796                mixed
797            };
798            x = crate::hyper::post(e, topology, &mixed, &x, &mix, t, n_embd)?;
799
800            let (y, mix) = crate::hyper::pre_exact(e, topology, &hyper.mlp, &x, t, n_embd)?;
801            let mut z = e.uninit(t * n_embd)?;
802            e.rms_norm(
803                &y,
804                layer.post_attn_norm.float_data(),
805                &mut z,
806                n_embd,
807                t,
808                eps,
809            )?;
810            // FFN branch: `batch` arms the pairs-shaped batched MoE across the t rows
811            // (lane/glm5-vrest — fail-closed inside to the byte-identical sequential
812            // loop); the =0 arm keeps the pre-lane per-(token,expert) class. Clocked
813            // into the vffn sub-bucket at trace level 2 (batched arm only, like vkda).
814            let t0 = vclock(trace_v && batch);
815            let ffn_out = self.hyper_ffn_branch_batch(e, layer, &z, t, il, batch)?;
816            if let Some(t0) = t0 {
817                let _ = e.stream().synchronize();
818                use std::sync::atomic::Ordering;
819                crate::spec_phase::V_FFN_NS
820                    .fetch_add(t0.elapsed().as_nanos() as u64, Ordering::Relaxed);
821            }
822            x = crate::hyper::post(e, topology, &ffn_out, &x, &mix, t, n_embd)?;
823            // glm5 DFlash2 feature tap (module doc, DRAFT SOURCE SEAM): the verify rows'
824            // contracted completed-layer outputs are next round's drafter context.
825            self.glm5_hc_tap(e, cache, topology, il, &x, t)?;
826        }
827        Ok(x)
828    }
829
830    /// Write one tapped layer's CONTRACTED completed output into the armed
831    /// [`HcTapSink`] — the glm5 DFlash2 drafter's measured feature contract (stream-mean
832    /// over the hyper streams, the probe's `hc_contract` capture definition). Staged
833    /// through the WALKING engine `e` (the owning stage engine under a ppN split), so the
834    /// sink is placement-invariant. One Option check when unarmed; nothing else pays.
835    ///
836    /// Two staging arms (loop-port 1):
837    ///   * `device_stage` (the verify-round sink): ONE async D2D into the slot's device
838    ///     buffer — the walk never blocks; the round drains all slots post-walk in its
839    ///     single sync point (`glm5_tap_drain`). Kills the five in-walk DtoHs the 3way
840    ///     window priced into the 31.6 ms fixed round cost (map row #17).
841    ///   * host-staged (prime sinks): the pre-port behavior — per-chunk DtoH, amortized
842    ///     over the prime's >= 256-row chunks.
843    pub(crate) fn glm5_hc_tap(
844        &self,
845        e: &Engine,
846        cache: &mut Cache,
847        topology: &crate::hyper::HyperTopology,
848        il: usize,
849        x: &CudaSlice<f32>,
850        t: usize,
851    ) -> Res<()> {
852        let Some(sink) = cache.hc_taps.as_mut() else {
853            return Ok(());
854        };
855        let Some(slot) = sink.layer_ids.iter().position(|&l| l == il) else {
856            return Ok(());
857        };
858        let h = sink.hidden;
859        let n_taps = sink.layer_ids.len();
860        let base = sink.base;
861        debug_assert!(
862            base + t <= sink.t,
863            "hc tap window {base}+{t} exceeds sink {}",
864            sink.t
865        );
866        let contracted = crate::hyper::contract_mean(e, topology, x, t, h)?;
867        if sink.device_stage {
868            // Lazy slot buffer on the WRITING engine (this layer always walks on one
869            // stage, so the buffer's device is stable for the sink's lifetime). Every
870            // walk row writes every tapped layer, so the buffer is fully covered by the
871            // walk that armed the sink.
872            if sink.dev[slot].is_none() {
873                sink.dev[slot] = Some(e.uninit(sink.t * h)?);
874            }
875            let buf = sink.dev[slot].as_mut().expect("just filled");
876            e.copy_into(buf, base * h, &contracted, t * h)?;
877            return Ok(());
878        }
879        let host = e.dtoh(&contracted)?;
880        for r in 0..t {
881            let dst = (base + r) * n_taps * h + slot * h;
882            sink.rows[dst..dst + h].copy_from_slice(&host[r * h..(r + 1) * h]);
883        }
884        Ok(())
885    }
886
887    /// Drain a device-staged tap sink into its host `rows` — the round's ONE post-walk
888    /// sync point for tap features (loop-port 1). Each slot reads back through its
889    /// layer's OWNING engine (the stage engine under a live split, the caller's engine
890    /// otherwise); the verify walk's terminal drain has already retired every stage's
891    /// writes (stream program order: the slot copy precedes its stage's TX, and the
892    /// TX-wait chain covers it transitively — the pp.rs multi-stream law).
893    fn glm5_tap_drain(&self, e: &Engine, sink: &mut HcTapSink) -> Res<()> {
894        if !sink.device_stage {
895            return Ok(());
896        }
897        let h = sink.hidden;
898        let n_taps = sink.layer_ids.len();
899        let split = match crate::pp::pp_cuts(self.layers.len()) {
900            Some(fence) if !crate::pp::pp2_streams_off() => {
901                Some((crate::pp::PpNRt::get(e)?, fence))
902            }
903            _ => None,
904        };
905        for slot in 0..n_taps {
906            let Some(buf) = sink.dev[slot].take() else {
907                continue;
908            };
909            let il = sink.layer_ids[slot];
910            let es = match split.as_ref() {
911                Some((rt, fence)) => {
912                    let stage = fence
913                        .windows(2)
914                        .position(|w| il >= w[0] && il < w[1])
915                        .ok_or_else(|| format!("tap layer {il} outside every stage range"))?;
916                    rt.engine(stage, e)
917                }
918                None => e,
919            };
920            let host = es.dtoh(&buf)?;
921            for r in 0..sink.t {
922                let dst = r * n_taps * h + slot * h;
923                sink.rows[dst..dst + h].copy_from_slice(&host[r * h..(r + 1) * h]);
924            }
925        }
926        Ok(())
927    }
928
929    /// Trunk exit of the verify walk, the batched head's decode-exact form
930    /// (`hyper_batch_head_logits`), with the collapsed pre-output_norm rows kept — they are
931    /// the h_seeds the MTP head re-seeds from (LANE.md §A). Under a split this runs on the
932    /// LAST stage's engine, where the loader put `output_norm` + the lm head.
933    fn glm5_verify_head(
934        &self,
935        e: &Engine,
936        topology: &crate::hyper::HyperTopology,
937        x: &CudaSlice<f32>,
938        t: usize,
939    ) -> Res<(CudaSlice<f32>, CudaSlice<f32>)> {
940        let n_embd = self.cfg.n_embd as usize;
941        let eps = self.cfg.rms_eps;
942        let collapsed =
943            crate::hyper::collapse(e, topology, self.hyper_head.as_ref(), x, t, n_embd)?;
944        let mut hn = e.uninit(t * n_embd)?;
945        e.rms_norm(
946            &collapsed,
947            self.output_norm.float_data(),
948            &mut hn,
949            n_embd,
950            t,
951            eps,
952        )?;
953        // Under the batched walk the lm head rides the rows-exact classes too (the bf16
954        // tcols twin reads the 1.27 GB head ONCE per round instead of once per row);
955        // per-row bits unchanged by contract, the tcols bit-gate holds it.
956        let logits = if glm5_verify_batch_on() && t > 1 {
957            e.matmul_rows_exact(&self.output, &hn, t)?
958        } else {
959            e.matmul_decode_exact(&self.output, &hn, t)?
960        };
961        Ok((logits, collapsed))
962    }
963
964    /// ppN twin of the verify walk (lane/glm5-ppn-verify, 2026-08-30), mirroring
965    /// `decode_step_batch_hyper_ppn` (decode_batch.rs): the t=K+1 rows walk as N stage
966    /// subgraphs — per-stage engine, per-stage pos_rows uploads, ONE `[t, streams, n_embd]`
967    /// boundary payload per fence cut. Row chaining is per-LAYER through the one cache
968    /// (row r+1 at layer il depends only on row r at layer il), so a straight layer-range
969    /// split preserves it exactly; no row ever crosses a boundary individually. Head +
970    /// collapsed rows land on the LAST stage's engine — where the loader put the lm head
971    /// and where `pp::new_cache*` places the MTP plane the re-seed feeds.
972    ///
973    /// DRAIN CONTRACT: this walk returns DEVICE buffers with no terminal dtoh (unlike its
974    /// decode twins, whose epilogue reads back on the last stage's stream), so it owns the
975    /// settle — the per-stage arm synchronizes the LAST stage's stream before returning.
976    /// The TX-wait chain transitively covers every earlier stage (pp.rs multi-stream law),
977    /// so the logits, the collapsed rows AND the ckpt's per-stage KDA columns are all safe
978    /// for consumption from the caller's streams after this returns.
979    #[allow(clippy::too_many_arguments)]
980    // allow: the parameter list mirrors its decode twin's stage-walk contract
981    fn glm5_verify_rows_ppn(
982        &self,
983        e: &Engine,
984        tokens: &[u32],
985        cache: &mut Cache,
986        mut ckpt: Glm5VerifyCkpt,
987        topology: &crate::hyper::HyperTopology,
988        fence: &[usize],
989    ) -> Res<(CudaSlice<f32>, CudaSlice<f32>, Glm5VerifyCkpt)> {
990        let t = tokens.len();
991        let n_embd = self.cfg.n_embd as usize;
992        let pos0 = ckpt.pos;
993        let payload = t * topology.streams * n_embd;
994        // Position buffers through THIS stage's engine (the per-stage pos_d law:
995        // allocated, consumed and freed on one stage's stream).
996        let pos_on = |eng: &Engine| -> Res<Glm5VerifyPos> { Glm5VerifyPos::new(eng, pos0, t) };
997
998        // Same-stream seam (MEMRA_PP_STREAMS=0): one engine, boundary copies between
999        // ranges — the shape every hc ppN walk uses for this knob.
1000        if crate::pp::pp2_streams_off() {
1001            let pos = pos_on(e)?;
1002            let embedded = e.htod(&self.embd.gather(n_embd, tokens))?;
1003            let mut x = crate::hyper::expand(e, topology, &embedded, t, n_embd)?;
1004            x =
1005                self.glm5_verify_range(e, topology, x, fence[0], fence[1], &pos, cache, &mut ckpt)?;
1006            for s in 1..fence.len() - 1 {
1007                let boundary_tx = e.clone_dtod(&x)?;
1008                let boundary_rx = e.clone_dtod(&boundary_tx)?;
1009                x = self.glm5_verify_range(
1010                    e,
1011                    topology,
1012                    boundary_rx,
1013                    fence[s],
1014                    fence[s + 1],
1015                    &pos,
1016                    cache,
1017                    &mut ckpt,
1018                )?;
1019            }
1020            let (logits, collapsed) = self.glm5_verify_head(e, topology, &x, t)?;
1021            return Ok((logits, collapsed, ckpt));
1022        }
1023
1024        let rt = crate::pp::PpNRt::get(e)?;
1025        let n_st = fence.len() - 1;
1026        assert_eq!(
1027            rt.n_stages(),
1028            n_st,
1029            "PpNRt stage count {} != fence stages {n_st}",
1030            rt.n_stages()
1031        );
1032        // #87 reverse publication (see decode_step_batch_ppn): order every stage stream
1033        // behind the caller before this body's first stage allocation.
1034        rt.fence_stages_behind(&e.stream())?;
1035
1036        // ---- STAGE 0: embed + expand (no weights) + layers [0, fence[1]) + TX ----
1037        let mut slot = {
1038            let _st0 = rt.enter(0);
1039            let e0 = rt.engine(0, e);
1040            let pos = pos_on(e0)?;
1041            let embedded = e0.htod(&self.embd.gather(n_embd, tokens))?;
1042            let x = crate::hyper::expand(e0, topology, &embedded, t, n_embd)?;
1043            let x = self
1044                .glm5_verify_range(e0, topology, x, fence[0], fence[1], &pos, cache, &mut ckpt)?;
1045            rt.tx(0, &x, payload)?
1046        };
1047
1048        // ---- MIDDLE STAGES: RX -> range -> TX ----
1049        for s in 1..n_st - 1 {
1050            let _st = rt.enter(s);
1051            let es = rt.engine(s, e);
1052            let pos = pos_on(es)?;
1053            let x = rt.rx(s - 1, slot, payload)?;
1054            let x = self.glm5_verify_range(
1055                es,
1056                topology,
1057                x,
1058                fence[s],
1059                fence[s + 1],
1060                &pos,
1061                cache,
1062                &mut ckpt,
1063            )?;
1064            slot = rt.tx(s, &x, payload)?;
1065        }
1066
1067        // ---- LAST STAGE: RX + final range + collapse/head + the drain (doc above) ----
1068        let _stl = rt.enter(n_st - 1);
1069        let el = rt.engine(n_st - 1, e);
1070        let pos = pos_on(el)?;
1071        let x = rt.rx(n_st - 2, slot, payload)?;
1072        let x = self.glm5_verify_range(
1073            el,
1074            topology,
1075            x,
1076            fence[n_st - 1],
1077            fence[n_st],
1078            &pos,
1079            cache,
1080            &mut ckpt,
1081        )?;
1082        let (logits, collapsed) = self.glm5_verify_head(el, topology, &x, t)?;
1083        // el.stream() under the enter-guard IS the stage stream (memra_runtime ambient
1084        // override) — this drain settles the whole walk transitively.
1085        el.stream().synchronize()?;
1086        Ok((logits, collapsed, ckpt))
1087    }
1088
1089    /// The engine that owns the trunk exit (collapse + output_norm + lm head), the MTP
1090    /// block's weights AND its latent plane under a ppN split: the LAST stage's engine —
1091    /// `hybrid.rs` uploads the head there (`pp::layer_engine(e, n_trunk, n_trunk - 1)`)
1092    /// and `pp::new_cache*` maps trailing MTP/NextN planes to the last stage. Door shut or
1093    /// the same-stream seam: the caller's engine, unchanged (single-device callers pay
1094    /// nothing — `e` is returned by identity).
1095    fn glm5_head_engine<'e>(&self, e: &'e Engine) -> Res<&'e Engine> {
1096        match crate::pp::pp_cuts(self.layers.len()) {
1097            Some(fence) if !crate::pp::pp2_streams_off() => {
1098                let rt = crate::pp::PpNRt::get(e)?;
1099                Ok(rt.engine(fence.len() - 2, e))
1100            }
1101            _ => Ok(e),
1102        }
1103    }
1104
1105    /// Roll the trunk back to exactly `keep` accepted verify rows (1 <= keep <= t; keep =
1106    /// j+1: the always-committed anchor row plus j accepted drafts).
1107    ///
1108    /// - MLA latent planes: `len = snapshot + keep` (truncate; rows are position-addressed
1109    ///   and append-only, so the kept rows ARE what a plain decode chain would have
1110    ///   written — the decode-exact contract), device `len_d` in lock-step, and
1111    ///   `truncate_index_pool_keys(pool)` clamps pool-key finality to what the shortened
1112    ///   `len` still justifies (the tail-ring residency tripwire fires on the next call if
1113    ///   this clamp is ever skipped).
1114    /// - KDA state: restore column keep-1 (state after the last kept row); full accept
1115    ///   (keep == t) keeps the resident state — the columns are clones OF it.
1116    /// - `cache.pos = snapshot + keep`.
1117    ///
1118    /// Under a live ppN split each stage's layers restore THROUGH that stage's engine ON
1119    /// its stream: the state planes and the ckpt columns live on the owning stage's device
1120    /// (per-stage `KvDev` allocation; per-stage clones in the walk), and enqueuing the
1121    /// restores on the same stage streams the walk writes on orders them without any extra
1122    /// fence — the next walk's own entry fence covers the primary-stream seam.
1123    pub fn glm5_verify_rollback(
1124        &self,
1125        e: &Engine,
1126        cache: &mut Cache,
1127        ckpt: &Glm5VerifyCkpt,
1128        keep: usize,
1129    ) -> Res<()> {
1130        if keep == 0 || keep > ckpt.rows {
1131            return Err(format!(
1132                "glm5_verify_rollback: keep={keep} outside 1..={} (the anchor row is always \
1133                 committed; keep = accepted drafts + 1)",
1134                ckpt.rows
1135            )
1136            .into());
1137        }
1138        match crate::pp::pp_cuts(self.layers.len()) {
1139            Some(fence) if !crate::pp::pp2_streams_off() => {
1140                let rt = crate::pp::PpNRt::get(e)?;
1141                for s in 0..fence.len() - 1 {
1142                    let _st = rt.enter(s);
1143                    let es = rt.engine(s, e);
1144                    for il in fence[s]..fence[s + 1] {
1145                        self.glm5_rollback_layer(es, cache, ckpt, keep, il)?;
1146                    }
1147                }
1148            }
1149            _ => {
1150                for il in 0..self.layers.len() {
1151                    self.glm5_rollback_layer(e, cache, ckpt, keep, il)?;
1152                }
1153            }
1154        }
1155        cache.pos = ckpt.pos + keep;
1156        Ok(())
1157    }
1158
1159    /// Restore ONE trunk layer to the ckpt's `keep`-row state (the per-plane contract in
1160    /// [`Self::glm5_verify_rollback`]'s doc). `e` is the layer's OWNING engine — the stage
1161    /// engine under a split, the caller's engine otherwise.
1162    fn glm5_rollback_layer(
1163        &self,
1164        e: &Engine,
1165        cache: &mut Cache,
1166        ckpt: &Glm5VerifyCkpt,
1167        keep: usize,
1168        il: usize,
1169    ) -> Res<()> {
1170        match &self.layers[il].mixer {
1171            Mixer::Mla(mla) => {
1172                let saved = ckpt.latent_len[il].ok_or_else(|| {
1173                    format!("glm5_verify_rollback: MLA layer {il} missing from the ckpt")
1174                })?;
1175                let plane = cache.latent[il].as_mut().ok_or_else(|| {
1176                    format!("glm5_verify_rollback: MLA layer {il} has no latent plane")
1177                })?;
1178                plane.len = saved + keep;
1179                let len_i32 =
1180                    i32::try_from(plane.len).map_err(|_| "latent length exceeds i32 mirror")?;
1181                // Door H (`MEMRA_GLM5_HTOD_DIET`): async `i32_set_k` instead of the synchronizing
1182                // pageable 4-byte copy — 11 of these per round, and unconditional (unlike the
1183                // KDA arm, which short-circuits when `keep == rows`).
1184                e.i32_mirror_store(&mut plane.len_d, len_i32)?;
1185                if let Some(indexer) = mla.index.as_ref() {
1186                    plane.truncate_index_pool_keys(indexer.geom.pool);
1187                }
1188                // spec x TP composition: the PEER latent replicas append in lock-step with
1189                // the canonical plane (the TP walk's construction), so the same truncation
1190                // restores each of them — through its own rank's engine for the device
1191                // `len_d` mirror.
1192                if let Some(tp) = mla.tp.as_ref()
1193                    && let Some(replicas) = cache.glm5_tp_latent_peer[il].as_mut()
1194                {
1195                    for (i, replica) in replicas.iter_mut().enumerate() {
1196                        replica.len = saved + keep;
1197                        tp.rt.peers[i].i32_mirror_store(&mut replica.len_d, len_i32)?;
1198                        if let Some(indexer) = mla.index.as_ref() {
1199                            replica.truncate_index_pool_keys(indexer.geom.pool);
1200                        }
1201                    }
1202                }
1203            }
1204            Mixer::Kda(la) if la.tp.is_some() => {
1205                if keep == ckpt.rows {
1206                    return Ok(()); // resident per-rank states ARE the post-keep states
1207                }
1208                let stash = ckpt.kda_tp[il].as_ref().ok_or_else(|| {
1209                    format!(
1210                        "glm5_verify_rollback: sharded KDA layer {il} has no per-rank stash \
1211                         (the batched TP verify walk fills it; the per-row arm is refused \
1212                         at walk entry)"
1213                    )
1214                })?;
1215                crate::glm5_tp::kda_tp_verify_rollback(e, la, stash, keep, cache, il)?;
1216            }
1217            Mixer::Kda(la) => {
1218                if keep == ckpt.rows {
1219                    return Ok(()); // resident state IS the state after the last kept row
1220                }
1221                // BATCHED-walk stash (lane/glm5-verify-batch): ring restore + re-roll,
1222                // then ONE scan replay at T=keep from the pre-round snapshot.
1223                if let Some(stash) = ckpt.kda_rows[il].as_ref() {
1224                    let snap = ckpt.kda_ssm_snap[il].as_ref().ok_or_else(|| {
1225                        format!("glm5_verify_rollback: KDA layer {il} has no ssm snapshot")
1226                    })?;
1227                    return crate::kda::kda_verify_rollback_rows(
1228                        e, la, snap, stash, keep, cache, il,
1229                    );
1230                }
1231                // Conv ring: restore the cloned column (unchanged — 288 KiB).
1232                let conv_cols = ckpt.kda_conv_cols[il].as_ref().ok_or_else(|| {
1233                    format!("glm5_verify_rollback: KDA layer {il} has no conv columns")
1234                })?;
1235                let conv = &conv_cols[keep - 1];
1236                {
1237                    let rl = cache.recur[il].as_mut().ok_or_else(|| {
1238                        format!("glm5_verify_rollback: KDA layer {il} has no recurrent state")
1239                    })?;
1240                    e.copy_into(&mut rl.conv_state, 0, conv, conv.len())?;
1241                }
1242                // Recurrent state: REPLAY rows 0..keep from the pre-round snapshot
1243                // (loop-port 3; ckpt doc) — each replay re-issues that row's original
1244                // t=1 scan over its stolen inputs, so the rebuilt state is byte-identical
1245                // to the per-row clone this retires.
1246                let snap = ckpt.kda_ssm_snap[il].as_ref().ok_or_else(|| {
1247                    format!("glm5_verify_rollback: KDA layer {il} has no ssm snapshot")
1248                })?;
1249                let stash = ckpt.kda_scan_stash[il].as_ref().ok_or_else(|| {
1250                    format!("glm5_verify_rollback: KDA layer {il} has no scan stash")
1251                })?;
1252                crate::kda::kda_scan_replay(e, la, snap, &stash[..keep], cache, il)?;
1253            }
1254            Mixer::Full(_) | Mixer::Linear(_) => {
1255                return Err(format!(
1256                    "glm5_verify_rollback: layer {il} mixer class was refused at walk \
1257                     entry and cannot appear in a ckpt"
1258                )
1259                .into());
1260            }
1261        }
1262        Ok(())
1263    }
1264
1265    /// Reset the MTP draft plane (il = n_trunk) to `len` rows — the LANE.md rollback
1266    /// contract ("one row per step, rollback = plane len reset") plus the same pool-key
1267    /// clamp every len-shortening path owes the tail ring. The plane lives on the LAST
1268    /// stage under a split (`pp::new_cache*` maps trailing MTP planes there), so the
1269    /// device mirror writes through the head engine.
1270    fn glm5_mtp_plane_reset(&self, e: &Engine, cache: &mut Cache, len: usize) -> Res<()> {
1271        let e = self.glm5_head_engine(e)?;
1272        let mtp = self
1273            .mtp
1274            .as_ref()
1275            .ok_or("glm5_mtp_plane_reset with no MTP head loaded")?;
1276        let il = self
1277            .plan
1278            .mtp_blocks
1279            .first()
1280            .ok_or("ModelPlan declares no MTP block")?
1281            .layer
1282            .index as usize;
1283        let plane = cache
1284            .latent
1285            .get_mut(il)
1286            .and_then(|plane| plane.as_mut())
1287            .ok_or_else(|| format!("MTP block layer {il} has no latent cache plane"))?;
1288        if len > plane.len {
1289            return Err(format!(
1290                "glm5_mtp_plane_reset: target {len} is past the plane's {} rows — a reset \
1291                 only ever shortens",
1292                plane.len
1293            )
1294            .into());
1295        }
1296        plane.len = len;
1297        let len_i32 = i32::try_from(len).map_err(|_| "latent length exceeds i32 mirror")?;
1298        e.stream().memcpy_htod(&[len_i32], &mut plane.len_d)?;
1299        if let Mixer::Mla(mla) = &mtp.mixer
1300            && let Some(indexer) = mla.index.as_ref()
1301        {
1302            plane.truncate_index_pool_keys(indexer.geom.pool);
1303        }
1304        Ok(())
1305    }
1306
1307    /// Single-shot glm5 speculative generation: draft (MTP head, K steps) -> verify (one
1308    /// t=K+1 walk) -> accept j (greedy longest matching prefix) -> rollback -> re-seed.
1309    /// Returns `(tokens, drafted, accepted)` — `generate_spec`'s contract. One-shot form:
1310    /// builds a [`Glm5SpecSession`] over a fresh cache and drives it to `max_new` — the
1311    /// SAME round machinery the serve worker bursts, so the tparallel gate's byte-identity
1312    /// pins cover the served path's rounds too.
1313    pub fn generate_spec_glm5(
1314        &self,
1315        e: &Engine,
1316        prompt: &[u32],
1317        max_new: usize,
1318        k: usize,
1319    ) -> Res<(Vec<u32>, usize, usize)> {
1320        self.generate_spec_glm5_gated(e, prompt, max_new, k, Glm5SpecKnobs::default())
1321    }
1322
1323    /// `generate_spec_glm5` with GATE INSTRUMENTS (never a serving surface): a draft
1324    /// override for deterministic forced-accept / forced-reject rounds, and a
1325    /// rollback-disable arm that red-proves the end-to-end byte-identity gate.
1326    pub fn generate_spec_glm5_gated(
1327        &self,
1328        e: &Engine,
1329        prompt: &[u32],
1330        max_new: usize,
1331        k: usize,
1332        mut knobs: Glm5SpecKnobs<'_>,
1333    ) -> Res<(Vec<u32>, usize, usize)> {
1334        let cap = Self::hyper_batch_cap();
1335        if k == 0 || k + 1 > cap {
1336            return Err(format!(
1337                "generate_spec_glm5: k={k} outside 1..={} (verify rows = k+1 must stay \
1338                 inside the decode-exact knee, cap {cap})",
1339                cap - 1
1340            )
1341            .into());
1342        }
1343        if max_new == 0 {
1344            return Ok((Vec::new(), 0, 0));
1345        }
1346        let max_ctx = prompt.len() + max_new + k + 8;
1347        let mut sess = self.glm5_spec_session_new(e, prompt, max_ctx, None)?;
1348        let mut out: Vec<u32> = Vec::with_capacity(max_new + k);
1349        let mut drafted = 0usize;
1350        let mut accepted = 0usize;
1351        while out.len() < max_new && !sess.finished() {
1352            let (burst, d, a) = self.glm5_spec_session_burst_gated(
1353                e,
1354                &mut sess,
1355                max_new - out.len(),
1356                k,
1357                &[],
1358                &mut knobs,
1359            )?;
1360            if burst.is_empty() {
1361                break; // ctx guard tripped with nothing new — never spin
1362            }
1363            out.extend(burst);
1364            drafted += d;
1365            accepted += a;
1366        }
1367        out.truncate(max_new);
1368        Ok((out, drafted, accepted))
1369    }
1370
1371    /// BATCHED MTP-PLANE WARM (loop-port fold-in; doc at the call site in
1372    /// `glm5_spec_session_new`): fill the NextN block's latent plane with rows for pairs
1373    /// `(tokens_next[i], hiddens row i)`, i in `0..t`, in chunked t-parallel passes —
1374    /// ops 1-7 of `mtp_head_forward_mla_cached` batched over the chunk (embed gather,
1375    /// enorm/hnorm, the eh_proj concat via `place_rows_strided`, attn_norm), then ONE
1376    /// `mla_attn_cached` append per chunk (the prime-class t>1 arm the trunk's own MLA
1377    /// layers warm through; its attention output is discarded — the plane rows are the
1378    /// product). The MoE FFN, final norm and lm-head of the per-token chain are never
1379    /// run: they fed nothing but the (discarded) draft logits of prompt positions.
1380    fn glm5_mtp_plane_fill(
1381        &self,
1382        e: &Engine,
1383        tokens_next: &[u32],
1384        hiddens: &CudaSlice<f32>,
1385        t: usize,
1386        cache: &mut Cache,
1387    ) -> Res<()> {
1388        let mtp = self
1389            .mtp
1390            .as_ref()
1391            .ok_or("glm5_mtp_plane_fill with no MTP head loaded")?;
1392        let il = self
1393            .plan
1394            .mtp_blocks
1395            .first()
1396            .ok_or("ModelPlan declares no MTP block")?
1397            .layer
1398            .index as usize;
1399        let Mixer::Mla(mla) = &mtp.mixer else {
1400            return Err("glm5_mtp_plane_fill serves MLA-mixer MTP blocks only".into());
1401        };
1402        if tokens_next.len() < t {
1403            return Err(format!(
1404                "glm5_mtp_plane_fill: {t} rows requested over {} successor tokens",
1405                tokens_next.len()
1406            )
1407            .into());
1408        }
1409        let n_embd = self.cfg.n_embd as usize;
1410        let eps = self.cfg.rms_eps;
1411        // Chunk bound: the trunk prime's workspace discipline — bounds the t>1 attention
1412        // workspace and the transient buffers below without changing the append semantics
1413        // (`mla_attn_cached` appends at the plane's running length either way).
1414        const CHUNK: usize = 512;
1415        let mut done = 0usize;
1416        while done < t {
1417            let tc = (t - done).min(CHUNK);
1418            let e_emb = e.htod(&self.embd.gather(n_embd, &tokens_next[done..done + tc]))?;
1419            let mut e_norm = e.uninit(tc * n_embd)?;
1420            e.rms_norm(&e_emb, mtp.enorm.float_data(), &mut e_norm, n_embd, tc, eps)?;
1421            // hnorm over the chunk's hidden rows (one contiguous view copy — rms_norm
1422            // takes an owned-slice operand).
1423            let hv = e.view(hiddens, (done + tc) * n_embd);
1424            let mut h_rows = e.uninit(tc * n_embd)?;
1425            e.copy_view_into(
1426                &mut h_rows,
1427                0,
1428                &hv.slice(done * n_embd..(done + tc) * n_embd),
1429                tc * n_embd,
1430            )?;
1431            let mut h_norm = e.uninit(tc * n_embd)?;
1432            e.rms_norm(
1433                &h_rows,
1434                mtp.hnorm.float_data(),
1435                &mut h_norm,
1436                n_embd,
1437                tc,
1438                eps,
1439            )?;
1440            // concat rows [tc, 2*n_embd] = [enorm ; hnorm] — two strided placements.
1441            let mut concat = e.uninit(tc * 2 * n_embd)?;
1442            e.place_rows_strided(&e_norm, &mut concat, n_embd, tc, 2 * n_embd, 0)?;
1443            e.place_rows_strided(&h_norm, &mut concat, n_embd, tc, 2 * n_embd, n_embd)?;
1444            let inp_sa = e.matmul(&mtp.eh_proj, &concat, tc)?;
1445            let mut a_norm = e.uninit(tc * n_embd)?;
1446            e.rms_norm(
1447                &inp_sa,
1448                mtp.attn_norm.float_data(),
1449                &mut a_norm,
1450                n_embd,
1451                tc,
1452                eps,
1453            )?;
1454            let pos: Vec<i32> = (done as i32..(done + tc) as i32).collect();
1455            let pos_d = e.htod_i32(&pos)?;
1456            let _ = self.mla_attn_cached(e, mla, &a_norm, &pos_d, tc, il, cache)?;
1457            done += tc;
1458        }
1459        Ok(())
1460    }
1461
1462    /// Row `row` of a `[rows, n_embd]` device stack, copied into its own `[n_embd]` buffer
1463    /// (the MTP `h_seed` handoff shape).
1464    fn glm5_seed_row(
1465        &self,
1466        e: &Engine,
1467        src: &CudaSlice<f32>,
1468        rows: usize,
1469        row: usize,
1470    ) -> Res<CudaSlice<f32>> {
1471        let n_embd = self.cfg.n_embd as usize;
1472        let stack = e.view(src, rows * n_embd);
1473        let view = stack.slice(row * n_embd..(row + 1) * n_embd);
1474        let mut seed = e.uninit(n_embd)?;
1475        e.copy_view_into(&mut seed, 0, &view, n_embd)?;
1476        Ok(seed)
1477    }
1478
1479    /// SERVED-SESSION ENTRY (lane/glm5-spec-routing): prime the prompt, warm the MTP plane,
1480    /// draw the boundary token, and hand back a [`Glm5SpecSession`] the worker bursts.
1481    ///
1482    /// `sampling`: `None` / `temp <= 0` = the greedy byte-contract route (the instrument);
1483    /// `Some` with `temp > 0` = the sampled route — the boundary token, the draft chain and
1484    /// the accept walk all draw through the session's own Philox counters (`sctr` device
1485    /// events, `uctr` host accept-test uniforms via `spec::host_u01`, tag 0xFFFF_FFFE), so a
1486    /// session's randomness never repeats across bursts (the session-continuity law).
1487    /// PENALIZED sampled requests are refused loudly — the glm5 accept walk has no penalty
1488    /// arm yet; worker admission keeps them on the plain path (same split as dspark's
1489    /// penalized-greedy exclusion).
1490    pub fn glm5_spec_session_new(
1491        &self,
1492        e: &Engine,
1493        prompt: &[u32],
1494        ctx_cap: usize,
1495        sampling: Option<SpecSampling>,
1496    ) -> Res<Glm5SpecSession> {
1497        if self.hyper.is_none() {
1498            return Err("generate_spec_glm5 requires a HyperConnections trunk".into());
1499        }
1500        // Two parallel/spec programs on one model never silently coexist unless the
1501        // composition is EXPLICITLY armed: the spec x TP verify/rollback wiring
1502        // (lane/glm5-composition) exists and is rig-gated, but it has zero real-artifact
1503        // receipts, so sessions on a TP-armed model stay co-refused unless
1504        // MEMRA_GLM5_SPEC_TP=1 lifts the refusal (default OFF by design — the FLAGS row).
1505        if crate::glm5_tp::glm5_tp_armed() {
1506            if !glm5_spec_tp_on() {
1507                return Err("glm5 spec is co-refused while MEMRA_GLM5_TP is armed: set \
1508                     MEMRA_GLM5_SPEC_TP=1 to run the gated spec x TP composition \
1509                     (default OFF — zero real-artifact receipts; every other admission \
1510                     law still holds)"
1511                    .into());
1512            }
1513            if !glm5_verify_batch_on() {
1514                return Err("MEMRA_GLM5_SPEC_TP=1 requires the BATCHED verify walk \
1515                     (MEMRA_GLM5_VERIFY_BATCH must not be 0): the per-row rollback seam \
1516                     carries no TP arm"
1517                    .into());
1518            }
1519            eprintln!(
1520                "[glm5-spec] spec x TP composition ARMED (MEMRA_GLM5_SPEC_TP=1): verify \
1521                 rows ride the TP shards; rollback restores per-rank planes \
1522                 performance_claim=false"
1523            );
1524        }
1525        // DRAFT-SOURCE SELECTION (module doc): DFlash2 when the drafter is loaded
1526        // (MEMRA_GLM5_DFLASH — it wins over a co-loaded MTP head, the boot receipt states
1527        // the selection), native MTP otherwise; neither = the loud refusal below. The
1528        // native MTP head is NOT required for the DFlash2 source (the q38 pattern).
1529        let dflash_src = self.glm5_dflash.as_ref();
1530        if dflash_src.is_none() && self.mtp.is_none() {
1531            return Err(
1532                "generate_spec_glm5 requires a draft source: the embedded MTP head \
1533                 (MEMRA_GLM5_MTP=1; a full MoE layer, unloaded by default) or the DFlash2 \
1534                 drafter (MEMRA_GLM5_DFLASH=<dir-or-hf-spec>)"
1535                    .into(),
1536            );
1537        }
1538        if prompt.len() < 2 {
1539            return Err(
1540                "generate_spec_glm5 needs a prompt of >= 2 tokens (the MTP plane warms on \
1541                 (token[i+1], hidden[i]) pairs)"
1542                    .into(),
1543            );
1544        }
1545        if dflash_src.is_none() && crate::spec::spec_hpost() {
1546            // MTP-carrier-specific refusal: the DFlash2 source consumes tapped trunk
1547            // features, not the h_seed carrier, so the flag has nothing to flip there.
1548            return Err(
1549                "generate_spec_glm5 has no MEMRA_SPEC_HPOST arm: the flag flips the MTP \
1550                 carrier to the post-norm hidden, but this loop seeds every committed pair \
1551                 from the trunk's PRE-output_norm collapsed rows (LANE.md §A). Mixing the \
1552                 two silently degrades drafts; the HPOST twin needs its own gate before it \
1553                 may run"
1554                    .into(),
1555            );
1556        }
1557        // ppN split (lane/glm5-ppn-verify): the verify walk, the rollback and the MTP
1558        // chain all run under the split now — but an UNQUALIFIED pipeline rewrite still
1559        // refuses loudly at the session seam, before any cache is allocated over
1560        // stage-sharded weights (worker admission additionally bounds the stage count to
1561        // the gated set; see glm5_spec_capable).
1562        if crate::pp::pp_cuts(self.layers.len()).is_some()
1563            && !self.rewrite_allowed(memra_gguf::execution_manifest::RewriteSurface::Pipeline)
1564        {
1565            return Err("pipeline rewrite is not qualified for this ModelPlan".into());
1566        }
1567        let sampling = sampling.filter(|sp| sp.temp > 0.0);
1568        if let Some(sp) = sampling.as_ref()
1569            && sp.pen_on()
1570        {
1571            return Err(
1572                "glm5 spec has no penalty arm yet: penalized sampled requests serve on the \
1573                 plain path (worker admission owns the exclusion; silently dropping the \
1574                 request's penalties is the failure class this refusal prevents)"
1575                    .into(),
1576            );
1577        }
1578        // Room for the prompt, the anchor row and at least one verify round.
1579        if prompt.len() + 4 > ctx_cap {
1580            return Err(format!(
1581                "glm5 spec session needs ctx for prompt {} + anchor + one verify round, \
1582                 cap {ctx_cap}",
1583                prompt.len()
1584            )
1585            .into());
1586        }
1587        let n_vocab = self.output.out_features();
1588        // FR-SPEC TRIM (module doc): a loaded `MEMRA_FRSPEC_TRIM` artifact means the draft
1589        // head projects over gathered top-N rows and every draft pick is a RANK id that
1590        // must remap through d2t to the true vocabulary BEFORE it is chained or verified.
1591        // The verify walk stays full-vocab regardless — the invariant under gate.
1592        if let Some(map) = self.glm5_d2t() {
1593            if map.iter().any(|&t| t as usize >= n_vocab) {
1594                return Err(format!(
1595                    "glm5 FR-Spec d2t carries a token id >= n_vocab {n_vocab} — the ranks \
1596                     artifact was minted for a different vocabulary"
1597                )
1598                .into());
1599            }
1600            // Engagement receipt (the dspark trim receipt's shape): the server-log line
1601            // the trim arm's per-session engagement is verified by.
1602            eprintln!(
1603                "[glm5-spec] draft head TRIMMED to {} rows (FR-Spec d2t engaged)",
1604                map.len(),
1605            );
1606        }
1607        // Confidence-gate engagement receipt (loop-port 2; the deploy-gate greps this —
1608        // never-serve-greedy law's receipt discipline): armed iff MEMRA_SPEC_PMIN > 0.
1609        if glm5_pmin() > 0.0 {
1610            eprintln!(
1611                "[glm5-spec] draft confidence gate armed: PMIN={:.3} PMIN0={} (native \
1612                 chain p-of-pick; DFlash2 selector-q tau-slot truncation)",
1613                glm5_pmin(),
1614                glm5_pmin0() as u8,
1615            );
1616        }
1617        // The MTP plane index is a NATIVE-arm need; the DFlash2 source never touches the
1618        // plane (it still allocates below — plan-structural, the named cost in the module
1619        // doc — but nothing reads or resets it).
1620        let mtp_il = match dflash_src {
1621            Some(_) => None,
1622            None => Some(
1623                self.plan
1624                    .mtp_blocks
1625                    .first()
1626                    .ok_or("ModelPlan declares no MTP block")?
1627                    .layer
1628                    .index as usize,
1629            ),
1630        };
1631        // Stage-owned allocation under a split (each layer's planes on its stage's device,
1632        // trailing MTP plane on the last stage); door shut = plain `Cache::new_planned`.
1633        let mut cache = crate::pp::new_cache_planned(e, &self.cfg, &self.plan, ctx_cap)?;
1634
1635        // ---- prime, boundary token, draft-source warm over the prompt ----
1636        // `prime_cache` routes to its own ppN twin under the split; `hiddens` is owned by
1637        // the LAST stage's engine (its published contract) — exactly where the MTP chain
1638        // below runs, so the warm consumes it with no device bounce. DFlash2 source: the
1639        // prime walk fills the armed HcTapSink with every prompt row's contracted tap
1640        // features (the drafter's context; round 1 ingests them into its own KV).
1641        let plen = prompt.len();
1642        let n_embd = self.cfg.n_embd as usize;
1643        let tap_layers = match dflash_src {
1644            Some(dr) => {
1645                let taps = glm5_dflash_tap_layers(&dr.draft, self.layers.len())?;
1646                cache.hc_taps = Some(HcTapSink::new(taps.clone(), n_embd, plen));
1647                Some(taps)
1648            }
1649            None => None,
1650        };
1651        let (logits0, _seed, hiddens) = self.prime_cache(e, prompt, &mut cache, 0)?;
1652        let eh = self.glm5_head_engine(e)?;
1653        let mut sctr = 0u32;
1654        let anchor = match sampling.as_ref() {
1655            Some(sp) => {
1656                crate::spec::sample_boundary_token(eh, &logits0, sp, &[], &mut sctr, "glm5-prime")?
1657            }
1658            None => argmax(&logits0) as u32,
1659        };
1660
1661        let (draft, pending) = match (dflash_src, tap_layers) {
1662            (Some(dr), Some(taps)) => {
1663                let sink = cache
1664                    .hc_taps
1665                    .take()
1666                    .ok_or("glm5 dflash prime tap sink vanished")?;
1667                // Drafter ctx KV on the HEAD engine (where the drafter weights loaded and
1668                // every round's chain runs); prompt feature rows ride `pending` so round 1
1669                // ingests them through the one chunked path.
1670                let kv = DflashKv::new(eh, &dr.draft.cfg, ctx_cap)?;
1671                (
1672                    Glm5DraftState::Dflash2 {
1673                        kv,
1674                        pending: sink.rows,
1675                        taps,
1676                    },
1677                    Vec::new(),
1678                )
1679            }
1680            _ => {
1681                // BATCHED PLANE WARM (loop-port fold-in — the map's #4, the spec.rs
1682                // `mtp_kv_fill_all` pattern re-aimed at the MLA plane): pairs
1683                // (prompt[i+1], h_i) at plane pos i, i in 0..P-1, filled in CHUNKED
1684                // t-parallel passes instead of P-1 sequential full-block forwards. The
1685                // sequential warm ran ~400 tok/s — the measured +2.5 s TTFT per 1k
1686                // prompt tokens, spec-battery flip condition 1 by name. MTP rows are
1687                // INDEPENDENT given the trunk hiddens (no row-to-row recurrence — the
1688                // plane is the only carrier), so the fill is exact in structure; the
1689                // t>1 attention takes the prime-class program, which can only move
1690                // DRAFTS, never output (verify arbitrates; the byte-identity batteries
1691                // stay the proof).
1692                self.glm5_mtp_plane_fill(eh, &prompt[1..], &hiddens, plen - 1, &mut cache)?;
1693                // pending = committed (token, h_seed) pairs not yet fed to the MTP plane.
1694                // The LAST pair's logits are the next round's first draft — the re-warm
1695                // doubles as draft 1.
1696                let pending = vec![(anchor, self.glm5_seed_row(eh, &hiddens, plen, plen - 1)?)];
1697                (Glm5DraftState::NativeMtp, pending)
1698            }
1699        };
1700        Ok(Glm5SpecSession {
1701            cache,
1702            committed: prompt.to_vec(),
1703            anchor,
1704            anchor_emitted: false,
1705            pending,
1706            draft,
1707            sampling,
1708            sctr,
1709            uctr: 0,
1710            rounds: 0,
1711            done: false,
1712            max_ctx: ctx_cap,
1713            mtp_il,
1714        })
1715    }
1716
1717    /// The loaded FR-Spec draft->target map, when a trim artifact actually landed on the
1718    /// embedded head (None = full-vocab head, rank id == token id).
1719    fn glm5_d2t(&self) -> Option<&[u32]> {
1720        self.mtp
1721            .as_ref()
1722            .and_then(|head| head.d2t.as_deref())
1723            .filter(|map| !map.is_empty())
1724    }
1725
1726    /// ONE serve burst (the worker's per-tick call, `step_glm5_spec`): rounds of
1727    /// draft(K) -> `glm5_verify_rows` -> accept -> rollback/commit until `target` new
1728    /// tokens are out, EOS commits, or the context guard trips. Returns
1729    /// `(burst, drafted, accepted)`; the burst may overshoot `target` by up to K (a
1730    /// round commits j+1 tokens atomically — the engine surplus stays committed in the
1731    /// session cache and the WORKER clamps public emission to the request budget, the
1732    /// SpecSession overshoot contract).
1733    pub fn glm5_spec_session_burst(
1734        &self,
1735        e: &Engine,
1736        sess: &mut Glm5SpecSession,
1737        target: usize,
1738        k: usize,
1739        eos: &[u32],
1740    ) -> Res<(Vec<u32>, usize, usize)> {
1741        self.glm5_spec_session_burst_gated(e, sess, target, k, eos, &mut Glm5SpecKnobs::default())
1742    }
1743
1744    /// [`glm5_spec_session_burst`] with GATE INSTRUMENTS (`Glm5SpecKnobs` — never a serving
1745    /// surface; no serving path constructs a non-default value).
1746    pub fn glm5_spec_session_burst_gated(
1747        &self,
1748        e: &Engine,
1749        sess: &mut Glm5SpecSession,
1750        target: usize,
1751        k: usize,
1752        eos: &[u32],
1753        knobs: &mut Glm5SpecKnobs<'_>,
1754    ) -> Res<(Vec<u32>, usize, usize)> {
1755        let cap = Self::hyper_batch_cap();
1756        if k == 0 || k + 1 > cap {
1757            return Err(format!(
1758                "glm5_spec_session_burst: k={k} outside 1..={} (verify rows = k+1 must stay \
1759                 inside the decode-exact knee, cap {cap})",
1760                cap - 1
1761            )
1762            .into());
1763        }
1764        if let Glm5DraftState::Dflash2 { .. } = sess.draft {
1765            let b = self
1766                .glm5_dflash
1767                .as_ref()
1768                .ok_or("dflash session on a model with no loaded drafter")?
1769                .draft
1770                .cfg
1771                .block_size;
1772            if k + 1 > b {
1773                return Err(format!(
1774                    "glm5_spec_session_burst: k={k} exceeds the DFlash2 drafter's block \
1775                     (block_size {b} = anchor + {} drafts, the trained mask pattern) — \
1776                     the worker clamps operator K pins to {} for this source; refusing \
1777                     loudly rather than drafting an untrained shape",
1778                    b - 1,
1779                    b - 1
1780                )
1781                .into());
1782            }
1783        }
1784        let d2t = self.glm5_d2t();
1785        if d2t.is_some() && knobs.skip_d2t_remap {
1786            eprintln!("[glm5-spec] d2t REMAP SKIPPED — red-arm instrument, drafts are rank ids");
1787        }
1788        let sp_on: Option<SpecSampling> = sess.sampling.filter(|sp| sp.temp > 0.0);
1789        let mut out: Vec<u32> = Vec::with_capacity(target + k);
1790        let mut drafted = 0usize;
1791        let mut accepted = 0usize;
1792        let mut phase: Option<SpecPhaseNs> =
1793            crate::spec_phase::spec_trace_on().then(SpecPhaseNs::default);
1794        if !sess.anchor_emitted {
1795            // The prime's boundary token: emitted exactly once, by the first burst.
1796            out.push(sess.anchor);
1797            sess.anchor_emitted = true;
1798            if eos.contains(&sess.anchor) {
1799                sess.done = true;
1800            }
1801        }
1802        while out.len() < target && !sess.done {
1803            // Context guard: a round appends up to k+1 trunk rows from `cache.pos` (and the
1804            // draft plane stays <= pos + k), so the next round must fit with one row slack.
1805            if sess.cache.pos + k + 2 > sess.max_ctx {
1806                sess.done = true;
1807                break;
1808            }
1809            let (round_tokens, n_drafted) =
1810                self.glm5_spec_round(e, sess, k, d2t, sp_on.as_ref(), knobs, phase.as_mut())?;
1811            drafted += n_drafted;
1812            accepted += round_tokens.len() - 1; // j accepted drafts + the bonus row
1813            for &tok in &round_tokens {
1814                if eos.contains(&tok) {
1815                    sess.done = true;
1816                }
1817            }
1818            out.extend_from_slice(&round_tokens);
1819            sess.rounds += 1;
1820        }
1821        if let Some(ph) = phase.as_ref() {
1822            ph.emit("glm5-phase", "glm5-phase-v", k);
1823        }
1824        Ok((out, drafted, accepted))
1825    }
1826
1827    /// One draft->verify->accept->rollback->re-seed round over the session state. Returns
1828    /// `(round_tokens, n_drafted)`: the round's committed tokens (`j` accepted drafts + the
1829    /// bonus token) and how many drafts actually entered the verify (== `k` today; the
1830    /// confidence gate may truncate it below `k`).
1831    #[allow(clippy::too_many_arguments)]
1832    // allow: the parameter list mirrors the round contract (session + policy + gate knobs +
1833    // the trace accumulator); bundling would hide which inputs are serving vs instrument
1834    fn glm5_spec_round(
1835        &self,
1836        e: &Engine,
1837        sess: &mut Glm5SpecSession,
1838        k: usize,
1839        d2t: Option<&[u32]>,
1840        sp: Option<&SpecSampling>,
1841        knobs: &mut Glm5SpecKnobs<'_>,
1842        mut phase: Option<&mut SpecPhaseNs>,
1843    ) -> Res<(Vec<u32>, usize)> {
1844        let n_vocab = self.output.out_features();
1845        let n_embd = self.cfg.n_embd as usize;
1846        // The MTP block / DFlash2 drafter, the trunk lm head and the verify walk's returned
1847        // rows all live on the LAST stage under a split — every draft-chain and accept-side
1848        // op below runs through the head engine (identity when the door is shut).
1849        let eh = self.glm5_head_engine(e)?;
1850        let mut t_mark = phase.as_ref().map(|_| SpecPhaseNs::clock(e, eh));
1851        // Phase-boundary bump: drain, bucket the elapsed ns, restart the clock. No-op with
1852        // the trace off (t_mark is None and no stream is ever synchronized).
1853        macro_rules! bump {
1854            ($field:ident) => {
1855                if let (Some(ph), Some(t0)) = (phase.as_deref_mut(), t_mark.as_mut()) {
1856                    let now = SpecPhaseNs::clock(e, eh);
1857                    ph.$field += now.duration_since(*t0).as_nanos() as u64;
1858                    *t0 = now;
1859                }
1860            };
1861        }
1862
1863        // CONFIDENCE GATE resolution (loop-port 2): the env pair is the serving surface
1864        // (the step37 family, no new flags); the knobs override is the gate instrument.
1865        let (p_min, pmin0) = knobs
1866            .pmin_override
1867            .unwrap_or_else(|| (glm5_pmin(), glm5_pmin0()));
1868
1869        // ---- 1+2. produce the K drafts (+ the retained q side), SOURCE-KEYED. Everything
1870        // after this point is shared and source-blind — the exactness seam (module doc).
1871        let (drafts, qside, mtp_committed_len) = match sess.draft {
1872            Glm5DraftState::Dflash2 { .. } => {
1873                let (d, q) = self.glm5_dflash_round_drafts(eh, sess, k, sp, knobs, p_min, pmin0)?;
1874                (d, q, 0)
1875            }
1876            Glm5DraftState::NativeMtp => {
1877                let mtp_il = sess.mtp_il.ok_or("native-mtp arm without a plane index")?;
1878                // ---- feed pending committed pairs; the last call yields draft 1 ----
1879                let mut last: Option<(CudaSlice<f32>, CudaSlice<f32>)> = None;
1880                for (tok, h) in sess.pending.drain(..) {
1881                    let plane_len = sess.cache.latent[mtp_il]
1882                        .as_ref()
1883                        .ok_or("MTP plane missing")?
1884                        .len;
1885                    last = Some(self.mtp_head_forward_mla_cached(
1886                        eh,
1887                        0,
1888                        tok,
1889                        &h,
1890                        &mut sess.cache,
1891                        plane_len,
1892                    )?);
1893                }
1894                let (mut d_logits, mut carrier) =
1895                    last.ok_or("glm5 spec round started with no pending committed pair")?;
1896                let mtp_committed_len = sess.cache.latent[mtp_il]
1897                    .as_ref()
1898                    .ok_or("MTP plane missing")?
1899                    .len;
1900
1901                // ---- chain K drafts. Greedy route: argmax over the draft head. Sampled
1902                // route: filtered Gumbel draw through the session's device Philox stream
1903                // (`sctr`), with the per-step filtered stats + logits retained — they are
1904                // the q side of the accept walk. Trimmed heads yield RANK ids that remap
1905                // through d2t to true vocab before anything consumes them (chain feed,
1906                // verify, output); the q gather keeps the rank id.
1907                let d_vocab = d2t.map(|m| m.len()).unwrap_or(n_vocab);
1908                let mut drafts: Vec<u32> = Vec::with_capacity(k); // true-vocab tokens
1909                let mut draft_idx: Vec<u32> = Vec::with_capacity(k); // draft-head rank ids
1910                let mut draft_logits: Vec<CudaSlice<f32>> = Vec::new(); // sampled route only
1911                let mut draft_stats: Vec<(f32, f32, f32)> = Vec::new(); // (mx, th, z), sampled only
1912                for ki in 0..k {
1913                    let (idx, sampled_stats) = match sp {
1914                        Some(sp) => {
1915                            let (idx, stats) =
1916                                glm5_sampled_draft(eh, &d_logits, d_vocab, sp, &mut sess.sctr)?;
1917                            (idx, Some(stats))
1918                        }
1919                        None => {
1920                            // Device argmax + ONE 4-byte readback per draft (loop-port 1)
1921                            // — replaces the full d_vocab logits DtoH + host argmax the
1922                            // map names at this seam. Same tie-break contract
1923                            // (argmax_gate); drafts never decide exactness anyway, the
1924                            // verify arbitrates. The #87 sentinel guard mirrors the
1925                            // spec.rs graph chain: a device argmax may emit a sentinel
1926                            // on a NaN row — refuse loudly, never gather an OOB embed.
1927                            let td = eh.argmax_token_device(&d_logits, d_vocab)?;
1928                            let idx = crate::spec::guard_vocab_token(
1929                                eh.dtoh_u32_one(&td)?,
1930                                d_vocab,
1931                                &format!(
1932                                    "glm5 native draft argmax at round {} ki={ki}",
1933                                    sess.rounds
1934                                ),
1935                            )?;
1936                            (idx, None)
1937                        }
1938                    };
1939                    // P-MIN CONFIDENCE GATE (loop-port 2, the spec.rs chain break): p =
1940                    // the head's softmax confidence in its own pick (the `g_p` statistic,
1941                    // prob_of_token_device kernels), one 4-byte read — armed rounds only.
1942                    // Break BEFORE the pick is drafted or the next full-MoE-layer chain
1943                    // forward is paid; a discarded sampled draw's Philox advance stands
1944                    // (spec.rs eager parity: "counts the p-min-discarded token too").
1945                    if p_min > 0.0 {
1946                        let tok_d = eh.htod_u32_v(&[idx])?;
1947                        let p_d = eh.prob_of_token_device(&d_logits, &tok_d, d_vocab)?;
1948                        let p = eh.dtoh(&p_d)?[0];
1949                        if p < p_min && (ki > 0 || pmin0) {
1950                            break;
1951                        }
1952                    }
1953                    if let Some(stats) = sampled_stats {
1954                        draft_stats.push(stats);
1955                        draft_logits.push(eh.clone_dtod(&d_logits)?);
1956                    }
1957                    let mut d = match d2t {
1958                        Some(map) if !knobs.skip_d2t_remap => map[idx as usize],
1959                        _ => idx,
1960                    };
1961                    if let Some(over) = knobs.draft_override.as_mut() {
1962                        d = over(sess.rounds, ki, d);
1963                    }
1964                    drafts.push(d);
1965                    draft_idx.push(idx);
1966                    if ki + 1 < k {
1967                        let plane_len = sess.cache.latent[mtp_il]
1968                            .as_ref()
1969                            .ok_or("MTP plane missing")?
1970                            .len;
1971                        let (lg, ca) = self.mtp_head_forward_mla_cached(
1972                            eh,
1973                            0,
1974                            d,
1975                            &carrier,
1976                            &mut sess.cache,
1977                            plane_len,
1978                        )?;
1979                        d_logits = lg;
1980                        carrier = ca;
1981                    }
1982                }
1983                let q = match sp {
1984                    Some(_) => Glm5DraftQ::Mtp {
1985                        draft_idx,
1986                        draft_logits,
1987                        draft_stats,
1988                    },
1989                    None => Glm5DraftQ::None,
1990                };
1991                (drafts, q, mtp_committed_len)
1992            }
1993        };
1994        bump!(draft);
1995
1996        // DFlash2 source: arm the verify tap — the walk's rows are next round's drafter
1997        // context features (rows 0..keep survive the accept; the sink is taken in step 7).
1998        // DEVICE-STAGED (loop-port 1): the walk D2Ds each tapped layer's contracted rows
1999        // instead of blocking on five in-walk DtoHs; step 7 drains post-walk.
2000        if let Glm5DraftState::Dflash2 { taps, .. } = &sess.draft {
2001            sess.cache.hc_taps = Some(HcTapSink::new_device_staged(
2002                taps.clone(),
2003                n_embd,
2004                drafts.len() + 1,
2005            ));
2006        }
2007
2008        // ---- 3. verify: one t=K+1 walk over the trunk ----
2009        let mut rows: Vec<u32> = Vec::with_capacity(drafts.len() + 1);
2010        rows.push(sess.anchor);
2011        rows.extend_from_slice(&drafts);
2012        let (vlogits, collapsed, ckpt) = self.glm5_verify_rows(e, &rows, &mut sess.cache)?;
2013        bump!(verify);
2014
2015        // ---- 4. accept ----
2016        // ZERO-DRAFT SAMPLED ROUND (PMIN0): the verify batch is just the anchor row —
2017        // m=1 = a plain decode step, exactly the llama.cpp gating spec.rs vendored. The
2018        // bonus is the full-accept filtered-Gumbel draw from that one row through the
2019        // session's Philox stream (identical in distribution to the plain sampled step
2020        // this round degenerates to). Greedy zero-draft rounds ride the general arm
2021        // below (j=0, bonus = the row-0 device argmax).
2022        let (j, bonus) = if let (true, Some(sp)) = (drafts.is_empty(), sp) {
2023            (
2024                0,
2025                self.glm5_sampled_bonus(eh, sess, sp, &vlogits, 0, n_vocab)?,
2026            )
2027        } else {
2028            match (sp, &qside) {
2029                (None, _) => {
2030                    // Greedy longest matching prefix (the DFlash2 probe's rule); bonus = the
2031                    // target's own argmax at the first non-accepted slot. Byte-deterministic —
2032                    // the instrument the spec-vs-plain identity gates pin.
2033                    //
2034                    // DEVICE ACCEPT ARGMAXES (loop-port 1, the K=1 flip): per verify row, a
2035                    // device argmax into one [t] slot buffer, then ONE tiny u32 readback —
2036                    // replacing the (K+1) x n_vocab logits DtoH + (K+1) host argmax scans
2037                    // (~2.4 MB + a host walk over 600k floats at K=3 on the real head; the
2038                    // 3way arithmetic needs 0.67 ms off the fixed round cost to flip K=1).
2039                    // `argmax_token_device_col` carries the host argmax's tie-break contract
2040                    // bit for bit (lowest index wins, argmax_gate-validated), so the accept
2041                    // walk commits the SAME tokens in the SAME order — the byte-identity
2042                    // batteries below stay the proof.
2043                    let t = rows.len();
2044                    let mut vam_d = eh.alloc_u32_zeroed(t)?;
2045                    for r in 0..t {
2046                        eh.argmax_token_device_col(&vlogits, r, n_vocab, &mut vam_d, r)?;
2047                    }
2048                    let vam = eh.dtoh_u32(&vam_d)?;
2049                    let mut j = 0usize;
2050                    while j < drafts.len() && drafts[j] == vam[j] {
2051                        j += 1;
2052                    }
2053                    (j, vam[j])
2054                }
2055                (
2056                    Some(sp),
2057                    Glm5DraftQ::Mtp {
2058                        draft_idx,
2059                        draft_logits,
2060                        draft_stats,
2061                    },
2062                ) => self.glm5_sampled_accept(
2063                    eh,
2064                    sess,
2065                    sp,
2066                    &vlogits,
2067                    &drafts,
2068                    draft_idx,
2069                    draft_logits,
2070                    draft_stats,
2071                    d2t,
2072                    drafts.len(),
2073                )?,
2074                (Some(sp), Glm5DraftQ::Selector { prop, dl }) => {
2075                    // The q38 serve route's rejection walk, VERBATIM (`dspark_accept_sampled`):
2076                    // `rows` = [anchor, drafts..] is its cand contract, verify row j arbitrates
2077                    // rows[j+1], the bonus draws from row k on full accept, and the reject-slot
2078                    // residual uses the selector's sparse candidate-set q. Philox counters are
2079                    // this session's — randomness never repeats across bursts.
2080                    let (m, next) = crate::dflash::dspark_accept_sampled(
2081                        eh,
2082                        &vlogits,
2083                        &rows,
2084                        rows.len(),
2085                        n_vocab,
2086                        dl,
2087                        prop,
2088                        sp,
2089                        &[],
2090                        &mut sess.sctr,
2091                        &mut sess.uctr,
2092                    )?;
2093                    let next = crate::spec::guard_vocab_token(
2094                        next,
2095                        n_vocab,
2096                        &format!(
2097                            "glm5 dflash2 sampled verify bonus at round {} j={m}",
2098                            sess.rounds
2099                        ),
2100                    )?;
2101                    (m, next)
2102                }
2103                (Some(_), Glm5DraftQ::None) => {
2104                    unreachable!("sampled round without a retained q side")
2105                }
2106            }
2107        };
2108        bump!(accept);
2109
2110        // ---- 5. commit j drafts + the bonus token ----
2111        let mut round_tokens: Vec<u32> = Vec::with_capacity(j + 1);
2112        round_tokens.extend_from_slice(&drafts[..j]);
2113        round_tokens.push(bonus);
2114
2115        // ---- 6. rollback the trunk to the accepted prefix ----
2116        let keep = j + 1;
2117        if knobs.disable_rollback {
2118            // RED-ARM INSTRUMENT: move pos, leave every state plane at post-row-K.
2119            sess.cache.pos = ckpt.pos + keep;
2120        } else {
2121            self.glm5_verify_rollback(e, &mut sess.cache, &ckpt, keep)?;
2122        }
2123        bump!(roll);
2124
2125        // ---- 7+8. draft-source state maintenance, SOURCE-KEYED ----
2126        match &mut sess.draft {
2127            Glm5DraftState::NativeMtp => {
2128                // MTP plane: len reset to the committed boundary (chain rows out), then
2129                // re-seed the pending pairs (token at pos0+i, collapsed row i-1).
2130                self.glm5_mtp_plane_reset(e, &mut sess.cache, mtp_committed_len)?;
2131                for i in 1..=keep {
2132                    let tok = round_tokens[i - 1];
2133                    let h = self.glm5_seed_row(eh, &collapsed, rows.len(), i - 1)?;
2134                    sess.pending.push((tok, h));
2135                }
2136            }
2137            Glm5DraftState::Dflash2 { pending, taps, .. } => {
2138                // The kept verify rows' tap features (rows 0..keep = [anchor, accepted
2139                // drafts]) become next round's drafter context — the probe's
2140                // `F_feat[new_lo:start]` advance. The drafter's own KV block rows were
2141                // transient (forward_round never moves kv.len), so no drafter rollback
2142                // exists to run. The trunk-side MTP plane was never touched.
2143                // Device-staged rows drain HERE — the round's one post-walk sync point
2144                // for tap features (loop-port 1).
2145                let mut sink = sess
2146                    .cache
2147                    .hc_taps
2148                    .take()
2149                    .ok_or("glm5 dflash verify tap sink vanished")?;
2150                self.glm5_tap_drain(e, &mut sink)?;
2151                let row_w = taps.len() * n_embd;
2152                pending.extend_from_slice(&sink.rows[..keep * row_w]);
2153            }
2154        }
2155        // Cache-row bookkeeping: the trunk committed rows [anchor, drafts[..j]] (keep =
2156        // j+1), so `committed` gains exactly those tokens; the BONUS is the new live
2157        // anchor — emitted this round, consumed by the trunk as the NEXT round's row 0
2158        // (the dspark `last` convention). Invariant at every round boundary:
2159        // `cache.pos == committed.len()`, token-for-token.
2160        sess.committed.push(sess.anchor);
2161        sess.committed.extend_from_slice(&drafts[..j]);
2162        sess.anchor = bonus;
2163        bump!(maint);
2164        if let Some(ph) = phase {
2165            ph.rounds += 1;
2166        }
2167        Ok((round_tokens, drafts.len()))
2168    }
2169
2170    /// ONE round's drafts from the DFlash2 source (module doc, DRAFT SOURCE SEAM) — the
2171    /// shipped q38 selector round, re-aimed at glm5's hc-contract features:
2172    ///
2173    ///   1. ingest pending committed feature rows into the drafter's own ctx KV (chunked
2174    ///      at 256 rows — the qwen depth-OOM bound; round 1 carries the whole prompt);
2175    ///   2. block forward `[anchor, MASK x b-1]` at absolute positions over the cached ctx
2176    ///      (`forward_round` — block K/V transient, exactly the reference crop);
2177    ///   3. draft logits = trunk lm_head over rows 1..b (mask-fill harvest — the DFlash2
2178    ///      family census; FR-Spec trim consumed exactly as the dspark serve arm does);
2179    ///   4. selector walk: greedy chain, or the sampled candidate-set walk whose recorded
2180    ///      q (`DsparkDraftSample::Selector`) the shared rejection accept consumes.
2181    ///
2182    /// Drafts are truncated to `k` (the chain is sequential, so a prefix is well-formed),
2183    /// then to the CONFIDENCE prefix when `p_min` is armed (loop-port 2, the tau-slot
2184    /// form): the selector's recorded per-slot q — `q_chosen` on the sampled walk, its
2185    /// T=1 twin on the greedy walk — gates each slot through `glm5_conf_keep`, so the
2186    /// low-confidence tail never enters the verify batch (a truncated round rides down
2187    /// the `31.6 + 20.1*K` line; rejection sampling stays exact for any proposal prefix).
2188    /// `knobs.draft_override` applies after — the gate instrument, never serving.
2189    #[allow(clippy::too_many_arguments)]
2190    // allow: the parameter list mirrors the round contract plus the resolved gate pair
2191    fn glm5_dflash_round_drafts(
2192        &self,
2193        eh: &Engine,
2194        sess: &mut Glm5SpecSession,
2195        k: usize,
2196        sp: Option<&SpecSampling>,
2197        knobs: &mut Glm5SpecKnobs<'_>,
2198        p_min: f32,
2199        pmin0: bool,
2200    ) -> Res<(Vec<u32>, Glm5DraftQ)> {
2201        let dr = self
2202            .glm5_dflash
2203            .as_ref()
2204            .ok_or("glm5 dflash draft state without a loaded drafter")?;
2205        let draft = &dr.draft;
2206        let c = &draft.cfg;
2207        let b = c.block_size;
2208        let n_embd = self.cfg.n_embd as usize;
2209        let n_vocab = self.output.out_features();
2210        let Glm5SpecSession {
2211            draft: state,
2212            cache,
2213            anchor,
2214            sctr: _,
2215            uctr,
2216            rounds,
2217            ..
2218        } = sess;
2219        let Glm5DraftState::Dflash2 { kv, pending, taps } = state else {
2220            return Err("glm5_dflash_round_drafts on a native-mtp session".into());
2221        };
2222        let anchor = *anchor;
2223
2224        // ---- 1. ingest pending committed feature rows (positions kv.len..) ----
2225        let row_w = taps.len() * n_embd;
2226        debug_assert_eq!(pending.len() % row_w, 0, "ragged pending feature rows");
2227        let n_new = pending.len() / row_w;
2228        let mut r0 = 0usize;
2229        while r0 < n_new {
2230            let t_c = (n_new - r0).min(256);
2231            let chunk = eh.htod(&pending[r0 * row_w..(r0 + t_c) * row_w])?;
2232            let feats = draft.ctx_features(eh, &chunk, t_c)?;
2233            let pos_c: Vec<i32> = ((kv.len as i32)..(kv.len + t_c) as i32).collect();
2234            draft.ingest_ctx(eh, kv, &feats, &pos_c, t_c)?;
2235            r0 += t_c;
2236        }
2237        pending.clear();
2238        let start = cache.pos;
2239        debug_assert_eq!(
2240            kv.len, start,
2241            "drafter ctx rows must equal committed trunk rows at a round boundary"
2242        );
2243
2244        // ---- 2. block forward over the cached ctx (decode-exact matmul scope: the m=8
2245        // drafter GEMMs otherwise fall into the prefill-GEMM class — the dspark round's
2246        // measured fix; RAII so a `?` exit never latches exact engine-wide) ----
2247        let exact_scope = eh.exact_scope(true);
2248        let mut block: Vec<u32> = vec![c.mask_token_id; b];
2249        block[0] = anchor;
2250        let noise = eh.htod(&self.embd.gather(n_embd, &block))?;
2251        let pos_block: Vec<i32> = ((start as i32)..(start + b) as i32).collect();
2252        let dh = draft.forward_round(eh, kv, &noise, &pos_block)?;
2253
2254        // ---- 3. draft logits over the mask-fill harvest rows 1..b ----
2255        let nd = b - 1;
2256        let mut rows_buf = eh.uninit(nd * n_embd)?;
2257        {
2258            let dv = eh.view(&dh, b * n_embd);
2259            let tail = dv.slice(n_embd..b * n_embd);
2260            eh.copy_view_into(&mut rows_buf, 0, &tail, nd * n_embd)?;
2261        }
2262        // TRIMMED DRAFT HEAD: the dspark serve arm's resolution verbatim — the FR-Spec
2263        // self-trim the load path builds on the MTP struct (gathered rows of the target's
2264        // own head). Available only when the MTP struct loaded; the DFlash2-without-head
2265        // boot (the q38 VRAM pattern) runs the full target head, stated in the boot receipt.
2266        let trim = self
2267            .mtp
2268            .as_ref()
2269            .filter(|m| m.d2t_from_target_head)
2270            .and_then(|m| m.shared_head_head.as_ref().zip(m.d2t.as_ref()))
2271            .filter(|(_, d2t)| !d2t.is_empty());
2272        let (dl_head, dl_vocab) = match trim {
2273            Some((head, d2t)) => (head, d2t.len()),
2274            None => (&self.output, n_vocab),
2275        };
2276        // skip_d2t_remap red arm (the q38 defect made loud): candidates stay RANK ids.
2277        let trim_d2t = trim
2278            .filter(|_| !knobs.skip_d2t_remap)
2279            .map(|(_, d2t)| d2t.as_slice());
2280        let dl = eh.matmul(dl_head, &rows_buf, nd)?;
2281
2282        // ---- 4. selector walk (greedy chain / sampled candidate-set walk) ----
2283        let (mut drafts, slot_q, qside) = match sp {
2284            None => {
2285                let (path, q) = draft
2286                    .dflash2_propose_greedy_q(eh, &dl, &rows_buf, nd, dl_vocab, anchor, trim_d2t)?;
2287                (path, q, Glm5DraftQ::None)
2288            }
2289            Some(sp) => {
2290                let (path, q_chosen, cand, q_rows) = draft.dflash2_propose_sampled(
2291                    eh, &dl, &rows_buf, nd, dl_vocab, anchor, sp.temp, sp.seed, uctr, trim_d2t,
2292                )?;
2293                let top_k = draft
2294                    .dflash2
2295                    .as_ref()
2296                    .ok_or("glm5 dflash drafter lost its DFlash2 head")?
2297                    .top_k;
2298                (
2299                    path,
2300                    q_chosen.clone(),
2301                    Glm5DraftQ::Selector {
2302                        prop: DsparkDraftSample::Selector {
2303                            cand,
2304                            q_rows,
2305                            q_chosen,
2306                            top_k,
2307                        },
2308                        dl,
2309                    },
2310                )
2311            }
2312        };
2313        drop(exact_scope);
2314        drafts.truncate(k);
2315        // TAU-SLOT CONFIDENCE TRUNCATION (loop-port 2): the low-confidence tail never
2316        // enters verify. Slot-indexed prefix reads keep the retained Selector q side
2317        // consistent (cand/q_rows/q_chosen are per-slot; the accept walk reads slots
2318        // 0..drafts.len()-1 only). p_min unset = today's rounds, untouched.
2319        if p_min > 0.0 {
2320            let kc = glm5_conf_keep(&slot_q[..drafts.len()], p_min, pmin0);
2321            drafts.truncate(kc);
2322        }
2323        if let Some(over) = knobs.draft_override.as_mut() {
2324            for (ki, d) in drafts.iter_mut().enumerate() {
2325                *d = over(*rounds, ki, *d);
2326            }
2327        }
2328        Ok((drafts, qside))
2329    }
2330
2331    /// SAMPLED ACCEPT (module doc): the rejection-sampling walk `u_j * q_j(x_j) < p_j(x_j)`
2332    /// over the verify logit rows — memra's existing sampled spec contract (the
2333    /// MEMRA_SPEC_TEMP route / dspark sampled-admission walk), plugged in at exactly the
2334    /// accept seam; walk and rollback unchanged. p and q take the SAME filter transforms
2335    /// (`filter_stats` + `softmax_gather_filtered`, distribution-exact for the filtered
2336    /// target); the accept-test uniforms come from `spec::host_u01` on the session's `uctr`
2337    /// (tag 0xFFFF_FFFE) and every device draw (draft chain, full-accept bonus, residual
2338    /// resample) advances the session's `sctr` — counters persist on the session so
2339    /// randomness never repeats across bursts. Returns `(j, bonus)`. `e` is the HEAD
2340    /// engine (the round resolves it): the verify rows and retained draft logits live on
2341    /// the last stage under a split.
2342    #[allow(clippy::too_many_arguments)]
2343    // allow: the parameter list mirrors the accept seam's inputs (verify rows + the draft
2344    // chain's retained q side); bundling into a struct would hide the p/q pairing
2345    fn glm5_sampled_accept(
2346        &self,
2347        e: &Engine,
2348        sess: &mut Glm5SpecSession,
2349        sp: &SpecSampling,
2350        vlogits: &CudaSlice<f32>,
2351        drafts: &[u32],
2352        draft_idx: &[u32],
2353        draft_logits: &[CudaSlice<f32>],
2354        draft_stats: &[(f32, f32, f32)],
2355        d2t: Option<&[u32]>,
2356        k: usize,
2357    ) -> Res<(usize, u32)> {
2358        let n_vocab = self.output.out_features();
2359        let d_vocab = d2t.map(|m| m.len()).unwrap_or(n_vocab);
2360        // FILTERED p_j: one batched stats pass over verify rows 0..k-1 (row j is the target
2361        // distribution at draft j's slot), then one batched gather of the drafted tokens.
2362        let rows_i: Vec<i32> = (0..k as i32).collect();
2363        let rows_d = e.htod_i32(&rows_i)?;
2364        let (mut th_d, mut z_d, mut mx_d) = (e.zeros(k)?, e.zeros(k)?, e.zeros(k)?);
2365        e.filter_stats(
2366            vlogits, n_vocab, &rows_d, &mut th_d, &mut z_d, &mut mx_d, n_vocab, k, sp.temp,
2367            sp.top_k, sp.top_p, sp.min_p,
2368        )?;
2369        let ids_d = e.htod_u32_v(drafts)?;
2370        let mut pj_d = e.zeros(k)?;
2371        e.softmax_gather_filtered(
2372            vlogits, n_vocab, &ids_d, &rows_d, &th_d, &z_d, &mut pj_d, n_vocab, k, sp.temp,
2373        )?;
2374        let pj = e.dtoh(&pj_d)?;
2375        let (thv, zv, mxv) = (e.dtoh(&th_d)?, e.dtoh(&z_d)?, e.dtoh(&mx_d)?);
2376
2377        // The walk: FILTERED q_j from the retained draft logits (rank id for trimmed heads),
2378        // host Philox accept test per slot.
2379        let mut j = 0usize;
2380        while j < k {
2381            let (_qmx, qth, qz) = draft_stats[j];
2382            let idsd = e.htod_u32_v(&[draft_idx[j]])?;
2383            let rows0 = e.htod_i32(&[0])?;
2384            let thd = e.htod(&[qth])?;
2385            let zd = e.htod(&[qz])?;
2386            let mut outd = e.zeros(1)?;
2387            e.softmax_gather_filtered(
2388                &draft_logits[j],
2389                d_vocab,
2390                &idsd,
2391                &rows0,
2392                &thd,
2393                &zd,
2394                &mut outd,
2395                d_vocab,
2396                1,
2397                sp.temp,
2398            )?;
2399            let qj = e.dtoh(&outd)?[0];
2400            let u = crate::spec::host_u01(sp.seed, sess.uctr);
2401            sess.uctr = sess.uctr.wrapping_add(1);
2402            if (u as f64) * (qj as f64) < pj[j] as f64 {
2403                j += 1;
2404            } else {
2405                break;
2406            }
2407        }
2408
2409        // Bonus: full accept draws a filtered Gumbel sample from the LAST verify row
2410        // (`glm5_sampled_bonus` — shared with the PMIN0 zero-draft round); rejection at j
2411        // resamples the residual norm(max(0, fp_j - fq_j)) — with a trimmed draft head, q
2412        // scatters back to full vocab first (`scatter_trim_logits`).
2413        if j == k {
2414            return Ok((
2415                j,
2416                self.glm5_sampled_bonus(e, sess, sp, vlogits, k, n_vocab)?,
2417            ));
2418        }
2419        let mut col = e.zeros(n_vocab)?;
2420        let bonus = {
2421            let vv = e.view(vlogits, (k + 1) * n_vocab);
2422            let row = vv.slice(j * n_vocab..(j + 1) * n_vocab);
2423            e.copy_view_into(&mut col, 0, &row, n_vocab)?;
2424            let p_stats = (mxv[j], thv[j], zv[j]);
2425            let q_stats = draft_stats[j];
2426            let sc = sess.sctr;
2427            sess.sctr = sess.sctr.wrapping_add(1);
2428            let mut sample_tok = e.alloc_u32_zeroed(1)?;
2429            match d2t {
2430                Some(map) => {
2431                    let map_d = e.htod_u32_v(map)?;
2432                    let mut q_full = e.zeros(n_vocab)?;
2433                    e.scatter_trim_logits(&draft_logits[j], &map_d, &mut q_full, d_vocab, n_vocab)?;
2434                    e.residual_sample_filtered(
2435                        &col,
2436                        Some(&q_full),
2437                        n_vocab,
2438                        sp.temp,
2439                        sp.seed,
2440                        sc,
2441                        p_stats,
2442                        q_stats,
2443                        &mut sample_tok,
2444                    )?;
2445                }
2446                None => {
2447                    e.residual_sample_filtered(
2448                        &col,
2449                        Some(&draft_logits[j]),
2450                        n_vocab,
2451                        sp.temp,
2452                        sp.seed,
2453                        sc,
2454                        p_stats,
2455                        q_stats,
2456                        &mut sample_tok,
2457                    )?;
2458                }
2459            }
2460            e.dtoh_u32(&sample_tok)?[0]
2461        };
2462        let bonus = crate::spec::guard_vocab_token(
2463            bonus,
2464            n_vocab,
2465            &format!("glm5 sampled verify bonus at round {} j={j}", sess.rounds),
2466        )?;
2467        Ok((j, bonus))
2468    }
2469
2470    /// One filtered-Gumbel bonus draw from verify row `row` through the session's device
2471    /// Philox stream — the sampled FULL-ACCEPT bonus, and the entire accept of a PMIN0
2472    /// zero-draft round (whose verify batch is just the anchor row: m=1 = a plain sampled
2473    /// decode step). Advances `sctr` exactly once; byte-for-byte the pre-extraction
2474    /// full-accept arm of `glm5_sampled_accept`.
2475    fn glm5_sampled_bonus(
2476        &self,
2477        e: &Engine,
2478        sess: &mut Glm5SpecSession,
2479        sp: &SpecSampling,
2480        vlogits: &CudaSlice<f32>,
2481        row: usize,
2482        n_vocab: usize,
2483    ) -> Res<u32> {
2484        let mut col = e.zeros(n_vocab)?;
2485        let vv = e.view(vlogits, (row + 1) * n_vocab);
2486        let src = vv.slice(row * n_vocab..(row + 1) * n_vocab);
2487        e.copy_view_into(&mut col, 0, &src, n_vocab)?;
2488        let rows0 = e.htod_i32(&[0])?;
2489        let (mut bth, mut bz, mut bmx) = (e.zeros(1)?, e.zeros(1)?, e.zeros(1)?);
2490        e.filter_stats(
2491            &col, n_vocab, &rows0, &mut bth, &mut bz, &mut bmx, n_vocab, 1, sp.temp, sp.top_k,
2492            sp.top_p, sp.min_p,
2493        )?;
2494        let (th, mx) = (e.dtoh(&bth)?[0], e.dtoh(&bmx)?[0]);
2495        let mut pb = e.zeros(n_vocab)?;
2496        e.gumbel_perturb_filtered(&col, &mut pb, n_vocab, sp.seed, sess.sctr, sp.temp, mx, th)?;
2497        sess.sctr = sess.sctr.wrapping_add(1);
2498        let td = e.argmax_token_device(&pb, n_vocab)?;
2499        crate::spec::guard_vocab_token(
2500            e.dtoh_u32_one(&td)?,
2501            n_vocab,
2502            &format!(
2503                "glm5 sampled verify bonus at round {} (row {row})",
2504                sess.rounds
2505            ),
2506        )
2507    }
2508}
2509
2510/// One filtered Gumbel draw from a draft-head logit row through the session's device Philox
2511/// stream — the sampled route's PROPOSAL. Returns the drawn RANK id and the row's filtered
2512/// stats `(row_max, threshold_e, renorm_mass)`, which the accept walk's q gather and the
2513/// rejection residual both reuse (the q side must be the distribution the draft was actually
2514/// drawn from, or rejection sampling is not exact for the filtered target).
2515fn glm5_sampled_draft(
2516    e: &Engine,
2517    dl: &CudaSlice<f32>,
2518    d_vocab: usize,
2519    sp: &SpecSampling,
2520    sctr: &mut u32,
2521) -> Res<(u32, (f32, f32, f32))> {
2522    let rows0 = e.htod_i32(&[0])?;
2523    let (mut th_d, mut z_d, mut mx_d) = (e.zeros(1)?, e.zeros(1)?, e.zeros(1)?);
2524    e.filter_stats(
2525        dl, d_vocab, &rows0, &mut th_d, &mut z_d, &mut mx_d, d_vocab, 1, sp.temp, sp.top_k,
2526        sp.top_p, sp.min_p,
2527    )?;
2528    let (th, z, mx) = (e.dtoh(&th_d)?[0], e.dtoh(&z_d)?[0], e.dtoh(&mx_d)?[0]);
2529    let mut pb = e.zeros(d_vocab)?;
2530    e.gumbel_perturb_filtered(dl, &mut pb, d_vocab, sp.seed, *sctr, sp.temp, mx, th)?;
2531    *sctr = sctr.wrapping_add(1);
2532    let td = e.argmax_token_device(&pb, d_vocab)?;
2533    let idx =
2534        crate::spec::guard_vocab_token(e.dtoh_u32_one(&td)?, d_vocab, "glm5 sampled draft draw")?;
2535    Ok((idx, (mx, th, z)))
2536}
2537
2538/// glm5_next SERVED speculative session (lane/glm5-spec-routing, 2026-08-30): the state one
2539/// request's spec decoding carries across worker bursts — the dspark/gemma session twins'
2540/// shape. The session OWNS its trunk cache (the worker's `s.cache` stays `None`); at every
2541/// burst boundary the invariant is `cache.pos == committed.len()` with each committed row's
2542/// trunk state exactly what a plain prime of that sequence would hold (the accept walk's
2543/// basis, pinned by the tparallel gate), plus ONE emitted-but-uncommitted `anchor` token
2544/// (the next round's verify row 0 — the dspark `last` convention).
2545pub struct Glm5SpecSession {
2546    cache: Cache,
2547    /// Every token whose trunk state the cache holds, in order (prompt + committed
2548    /// generation). EXCLUDES the live `anchor`.
2549    pub committed: Vec<u32>,
2550    /// The last emitted token, not yet consumed by the trunk — round anchor / verify row 0.
2551    anchor: u32,
2552    /// The prime's boundary token is emitted exactly once, by the first burst.
2553    anchor_emitted: bool,
2554    /// Committed `(token, h_seed)` pairs not yet fed to the MTP draft plane; the last
2555    /// feed's logits double as the next round's first draft (the re-warm contract).
2556    /// NATIVE-MTP arm only; the DFlash2 source keeps its own pending rows in `draft`.
2557    pending: Vec<(u32, CudaSlice<f32>)>,
2558    /// The session's pinned draft source + its state (module doc, DRAFT SOURCE SEAM).
2559    draft: Glm5DraftState,
2560    /// `None` / `temp <= 0` = greedy byte-contract route. Fixed for the session — the
2561    /// worker's admission owns the sampler identity.
2562    sampling: Option<SpecSampling>,
2563    /// Session-continuity Philox counters (never reset across bursts): `sctr` = device
2564    /// sampling events (boundary, draft chain, bonus, residual), `uctr` = host accept-test
2565    /// uniforms (`spec::host_u01`, tag 0xFFFF_FFFE).
2566    sctr: u32,
2567    uctr: u32,
2568    /// Verify rounds completed over the session lifetime (the worker's per-burst
2569    /// rounds-delta receipt, the dspark `rounds` convention).
2570    pub rounds: usize,
2571    done: bool,
2572    max_ctx: usize,
2573    /// MTP draft-plane layer index — `Some` on the native-MTP arm only.
2574    mtp_il: Option<usize>,
2575}
2576
2577impl Glm5SpecSession {
2578    /// Context capacity of the session's cache (the server's ContextFull guard).
2579    pub fn cache_max_ctx(&self) -> usize {
2580        self.max_ctx
2581    }
2582    /// Trunk rows currently committed (== `committed.len()` at burst boundaries).
2583    pub fn pos(&self) -> usize {
2584        self.cache.pos
2585    }
2586    /// EOS committed or the context guard tripped: the next burst would emit nothing.
2587    pub fn finished(&self) -> bool {
2588        self.done
2589    }
2590    /// True when the session is a legal demotion source (loop-port fold-in, map #8):
2591    /// GREEDY only — a sampled session's committed stream depends on its session-owned
2592    /// Philox counters, and the plain batched sampler is a different random program
2593    /// mid-request (the exact exclusion the MTP and dspark sweeps carry).
2594    pub fn demote_eligible(&self) -> bool {
2595        self.sampling.is_none()
2596    }
2597}
2598
2599impl HybridModel {
2600    /// ONE-WAY DEMOTION HANDOFF for the glm5 session (loop-port fold-in — the map's #8,
2601    /// the `SpecSession::into_demoted` / `DsparkSpecSession::into_demoted` twin): consume
2602    /// the session and hand `(cache, next_pred)` to the plain batched-decode path, so a
2603    /// spec session admitted on a quiet box stops serializing the tick when load arrives
2604    /// (the spec-gate HIGH sweep's ship-safety lever; dspark receipt: "c=8 429.6 = parity
2605    /// (pre-lane -37%)").
2606    ///
2607    /// THE ANCHOR IS THE CARRIED-PENDING SHAPE: glm5 emits each round's bonus immediately
2608    /// (`round_tokens` include it) while the trunk consumes it only as the NEXT round's
2609    /// row 0 — so at every burst boundary the session holds ONE emitted-but-uncommitted
2610    /// token. Handing the cache over as-is would leave it one row short of the public
2611    /// stream, and `device_next` re-emitting the anchor would duplicate a served token.
2612    /// The flush below is `spec_flush_pending`'s exact analogue: ONE plain T=1 decode
2613    /// step commits the anchor (byte-identical to the never-drafted chain — the
2614    /// tparallel gate's accept-j-then-continue identity IS this claim), and its argmax
2615    /// becomes the handoff's `next_pred` — a token the batched path emits and feeds
2616    /// exactly as it would its own. One trunk pass, once per demotion, never per burst.
2617    ///
2618    /// ONE-WAY BY DESIGN: the draft state (MTP pending pairs / DFlash2 drafter KV and
2619    /// feature rows) and the Philox counters are DROPPED, freeing their VRAM; there is
2620    /// no cheap symmetric re-promotion (the spec.rs law, verbatim). Sampled sessions
2621    /// refuse loudly (`demote_eligible`; the worker's sweep excludes them first).
2622    pub fn glm5_spec_into_demoted(
2623        &self,
2624        e: &Engine,
2625        mut sess: Glm5SpecSession,
2626    ) -> Res<(Cache, u32)> {
2627        if !sess.demote_eligible() {
2628            return Err(
2629                "glm5 demote: sampled sessions stay on spec until they end (session-owned \
2630                 Philox vs the worker sampler is an unmeasured distributional seam — the \
2631                 MTP sweep's exclusion, verbatim)"
2632                    .into(),
2633            );
2634        }
2635        if sess.cache.pos + 1 > sess.max_ctx {
2636            return Err(format!(
2637                "glm5 demote: no room to flush the live anchor ({} + 1 > ctx {})",
2638                sess.cache.pos, sess.max_ctx
2639            )
2640            .into());
2641        }
2642        let logits = self.decode_step(e, sess.anchor, &mut sess.cache)?;
2643        sess.committed.push(sess.anchor);
2644        let next = argmax(&logits) as u32;
2645        Ok((sess.cache, next))
2646    }
2647}
2648
2649/// Gate instruments for `generate_spec_glm5_gated`. Documented as instruments: no serving
2650/// path constructs a non-default value.
2651#[derive(Default)]
2652pub struct Glm5SpecKnobs<'a> {
2653    /// `(round, draft_index, greedy_draft) -> draft` — deterministic forced-accept /
2654    /// forced-reject rounds for the end-to-end gate.
2655    pub draft_override: Option<&'a mut dyn FnMut(usize, usize, u32) -> u32>,
2656    /// RED ARM ONLY: skip the state rollback (pos still moves). A corrupted draft must then
2657    /// leave post-row-K KDA state and un-truncated latent rows behind — the end-to-end gate
2658    /// asserts the tape DIVERGES from plain decode (or the kpool residency tripwire fires).
2659    pub disable_rollback: bool,
2660    /// RED ARM ONLY: with an FR-Spec trim loaded, use the draft argmax RANK id as the vocab
2661    /// id (the q38 skipped-remap defect: 0/248 acceptance with every exactness gate green).
2662    /// The gate asserts the drafted sequence diverges from the untrimmed arm's while the
2663    /// output tape STAYS byte-identical to plain decode — the silent failure made loud.
2664    pub skip_d2t_remap: bool,
2665    /// GATE INSTRUMENT for the confidence gate (loop-port 2): `Some((p_min, pmin0))`
2666    /// overrides the `MEMRA_SPEC_PMIN`/`MEMRA_SPEC_PMIN0` env pair for this call — the
2667    /// env statics latch once per process, so the byte-identity gate drives its PMIN
2668    /// arms through here instead of the environment. `None` = the serving resolution.
2669    pub pmin_override: Option<(f32, bool)>,
2670}
2671
2672#[cfg(test)]
2673mod conf_keep_tests {
2674    use super::glm5_conf_keep;
2675
2676    /// The spec.rs chain-break semantics, pinned CPU-side (loop-port 2): break at the
2677    /// first sub-threshold slot; slot 0 survives a miss unless PMIN0.
2678    #[test]
2679    fn conf_keep_matches_the_spec_rs_break_semantics() {
2680        // Gate off: everything kept.
2681        assert_eq!(glm5_conf_keep(&[0.1, 0.1], 0.0, true), 2);
2682        // All confident: everything kept.
2683        assert_eq!(glm5_conf_keep(&[0.9, 0.8, 0.7], 0.5, false), 3);
2684        // Break mid-chain at the first miss; the confident tail after it never rides
2685        // (prefix truncation — the accept rule could never commit past the gap anyway).
2686        assert_eq!(glm5_conf_keep(&[0.9, 0.2, 0.9], 0.5, false), 1);
2687        assert_eq!(glm5_conf_keep(&[0.9, 0.2, 0.9], 0.5, true), 1);
2688        // Slot-0 miss: survives without PMIN0 (the j > 0 arm of the break condition), and
2689        // does NOT latch — slot 1 is judged on its own confidence (the spec.rs chain
2690        // evaluates each slot's p independently)...
2691        assert_eq!(glm5_conf_keep(&[0.2, 0.9], 0.5, false), 2);
2692        // ...but a sub-threshold slot past 0 still breaks.
2693        assert_eq!(glm5_conf_keep(&[0.2, 0.2], 0.5, false), 1);
2694        // PMIN0 arms the zero-draft round.
2695        assert_eq!(glm5_conf_keep(&[0.2, 0.9], 0.5, true), 0);
2696        // Boundary: q == p_min is NOT below it (strict <, the spec.rs test).
2697        assert_eq!(glm5_conf_keep(&[0.5, 0.5], 0.5, true), 2);
2698        // Empty chain: nothing to keep.
2699        assert_eq!(glm5_conf_keep(&[], 0.5, true), 0);
2700    }
2701}